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, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::cmp::Reverse;
30use std::collections::{BTreeMap, BinaryHeap, HashMap, HashSet};
31use std::path::{Path, PathBuf};
32use std::sync::atomic::AtomicBool;
33use std::sync::Arc;
34use zeroize::Zeroizing;
35
36pub const WAL_DIR: &str = "_wal";
37pub const RUNS_DIR: &str = "_runs";
38pub const CACHE_DIR: &str = "_cache";
39pub const META_DIR: &str = "_meta";
40pub const RCACHE_DIR: &str = "_rcache";
41pub const KEYS_FILENAME: &str = "keys";
42pub const SCHEMA_FILENAME: &str = "schema.json";
43
44fn derive_next_run_id(
45 dir: &Path,
46 runs_root: Option<&crate::durable_file::DurableRoot>,
47 active: &[RunRef],
48 retiring: &[crate::manifest::RetiredRun],
49) -> Result<u64> {
50 let mut maximum = 0_u64;
51 for run_id in active
52 .iter()
53 .map(|run| run.run_id)
54 .chain(retiring.iter().map(|run| run.run_id))
55 {
56 let run_id = u64::try_from(run_id)
57 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
58 maximum = maximum.max(run_id);
59 }
60 let names = match runs_root {
61 Some(root) => root.list_regular_files(".")?,
62 None => std::fs::read_dir(dir.join(RUNS_DIR))?
63 .map(|entry| entry.map(|entry| entry.file_name()))
64 .collect::<std::io::Result<Vec<_>>>()?,
65 };
66 for name in names {
67 let Some(name) = name.to_str() else {
68 continue;
69 };
70 let Some(digits) = name
71 .strip_prefix("r-")
72 .and_then(|name| name.strip_suffix(".sr"))
73 else {
74 continue;
75 };
76 let Ok(run_id) = digits.parse::<u64>() else {
77 continue;
78 };
79 if name == format!("r-{run_id}.sr") {
80 maximum = maximum.max(run_id);
81 }
82 }
83 maximum
84 .checked_add(1)
85 .map(|next| next.max(1))
86 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
87}
88
89enum ControlledVisibleCandidate {
90 Memory(Row),
91 Run(RunVisibleVersion),
92}
93
94impl ControlledVisibleCandidate {
95 fn row_id(&self) -> RowId {
96 match self {
97 Self::Memory(row) => row.row_id,
98 Self::Run(version) => version.row_id,
99 }
100 }
101
102 fn committed_epoch(&self) -> Epoch {
103 match self {
104 Self::Memory(row) => row.committed_epoch,
105 Self::Run(version) => version.committed_epoch,
106 }
107 }
108
109 fn deleted(&self) -> bool {
110 match self {
111 Self::Memory(row) => row.deleted,
112 Self::Run(version) => version.deleted,
113 }
114 }
115}
116
117enum ControlledVisibleCursor {
118 Memory(std::vec::IntoIter<Row>),
119 Run(Box<RunVisibleVersionCursor>),
120 #[cfg(test)]
121 Synthetic {
122 next: u64,
123 end: u64,
124 },
125}
126
127struct ControlledVisibleSource {
128 cursor: ControlledVisibleCursor,
129 current: Option<ControlledVisibleCandidate>,
130}
131
132impl ControlledVisibleSource {
133 fn memory(rows: Vec<Row>) -> Self {
134 Self {
135 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
136 current: None,
137 }
138 }
139
140 fn run(cursor: RunVisibleVersionCursor) -> Self {
141 Self {
142 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
143 current: None,
144 }
145 }
146
147 #[cfg(test)]
148 fn synthetic(end: u64) -> Self {
149 Self {
150 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
151 current: None,
152 }
153 }
154
155 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
156 self.current = match &mut self.cursor {
157 ControlledVisibleCursor::Memory(rows) => {
158 rows.next().map(ControlledVisibleCandidate::Memory)
159 }
160 ControlledVisibleCursor::Run(cursor) => cursor
161 .next_visible_version(control)?
162 .map(ControlledVisibleCandidate::Run),
163 #[cfg(test)]
164 ControlledVisibleCursor::Synthetic { next, end } => {
165 if *next > *end {
166 None
167 } else {
168 let row = Row::new(RowId(*next), Epoch(1));
169 *next += 1;
170 Some(ControlledVisibleCandidate::Memory(row))
171 }
172 }
173 };
174 Ok(())
175 }
176
177 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate> {
178 let current = self.current.take().ok_or_else(|| {
179 MongrelError::Other("controlled visible source was not primed".into())
180 })?;
181 self.advance(control)?;
182 Ok(current)
183 }
184
185 fn materialize(
186 &mut self,
187 candidate: ControlledVisibleCandidate,
188 control: &crate::ExecutionControl,
189 ) -> Result<Row> {
190 match candidate {
191 ControlledVisibleCandidate::Memory(row) => Ok(row),
192 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
193 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
194 _ => Err(MongrelError::Other(
195 "run candidate escaped its controlled cursor".into(),
196 )),
197 },
198 }
199 }
200}
201
202fn merge_controlled_visible_sources(
203 sources: &mut [ControlledVisibleSource],
204 control: &crate::ExecutionControl,
205 mut expired: impl FnMut(&Row) -> bool,
206 mut visit: impl FnMut(Row) -> Result<()>,
207) -> Result<()> {
208 let mut heap = BinaryHeap::new();
209 for (source_index, source) in sources.iter_mut().enumerate() {
210 source.advance(control)?;
211 if let Some(candidate) = &source.current {
212 heap.push(Reverse((candidate.row_id(), source_index)));
213 }
214 }
215 let mut merged = 0_usize;
216 while let Some(Reverse((row_id, source_index))) = heap.pop() {
217 if merged.is_multiple_of(256) {
218 control.checkpoint()?;
219 }
220 merged += 1;
221 let mut best_source = source_index;
222 let mut best = sources[source_index].pop(control)?;
223 if let Some(next) = &sources[source_index].current {
224 heap.push(Reverse((next.row_id(), source_index)));
225 }
226 while heap
227 .peek()
228 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
229 {
230 let Some(Reverse((_, source_index))) = heap.pop() else {
231 break;
232 };
233 let candidate = sources[source_index].pop(control)?;
234 if candidate.committed_epoch() > best.committed_epoch() {
235 best = candidate;
236 best_source = source_index;
237 }
238 if let Some(next) = &sources[source_index].current {
239 heap.push(Reverse((next.row_id(), source_index)));
240 }
241 }
242 if best.deleted() {
243 continue;
244 }
245 let row = sources[best_source].materialize(best, control)?;
246 if !expired(&row) {
247 visit(row)?;
248 }
249 }
250 control.checkpoint()
251}
252
253#[cfg(test)]
254mod controlled_visible_cursor_tests {
255 use super::*;
256
257 #[test]
258 fn streams_more_than_one_million_rows_without_a_source_cap() {
259 let control = crate::ExecutionControl::new(None);
260 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
261 let mut count = 0_u64;
262 let mut last = 0_u64;
263 merge_controlled_visible_sources(
264 &mut sources,
265 &control,
266 |_| false,
267 |row| {
268 count += 1;
269 assert!(row.row_id.0 > last);
270 last = row.row_id.0;
271 Ok(())
272 },
273 )
274 .unwrap();
275 assert_eq!(count, 1_000_001);
276 assert_eq!(last, 1_000_001);
277 }
278
279 #[test]
280 fn merge_orders_rows_and_honors_newest_tombstones() {
281 let control = crate::ExecutionControl::new(None);
282 let older = vec![
283 Row::new(RowId(1), Epoch(1)),
284 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
285 Row::new(RowId(4), Epoch(1)),
286 ];
287 let mut deleted = Row::new(RowId(1), Epoch(2));
288 deleted.deleted = true;
289 let newer = vec![
290 deleted,
291 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
292 Row::new(RowId(3), Epoch(2)),
293 ];
294 let mut sources = vec![
295 ControlledVisibleSource::memory(older),
296 ControlledVisibleSource::memory(newer),
297 ];
298 let mut rows = Vec::new();
299 merge_controlled_visible_sources(
300 &mut sources,
301 &control,
302 |_| false,
303 |row| {
304 rows.push(row);
305 Ok(())
306 },
307 )
308 .unwrap();
309 assert_eq!(
310 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
311 vec![2, 3, 4]
312 );
313 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
314 }
315}
316
317fn iso_now_bytes() -> Vec<u8> {
320 let secs = std::time::SystemTime::now()
321 .duration_since(std::time::UNIX_EPOCH)
322 .map(|d| d.as_secs() as i64)
323 .unwrap_or(0);
324 let days = secs.div_euclid(86_400);
325 let rem = secs.rem_euclid(86_400);
326 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
327 let (year, month, day) = civil_from_days(days);
328 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
329}
330
331pub(crate) fn unix_nanos_now() -> i64 {
332 std::time::SystemTime::now()
333 .duration_since(std::time::UNIX_EPOCH)
334 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
335 .unwrap_or(0)
336}
337
338fn ann_candidate_cap(
339 index_len: usize,
340 context: Option<&crate::query::AiExecutionContext>,
341) -> usize {
342 index_len
343 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
344 .min(context.map_or(
345 crate::query::MAX_RAW_INDEX_CANDIDATES,
346 crate::query::AiExecutionContext::max_fused_candidates,
347 ))
348}
349
350#[cfg(test)]
351mod ann_candidate_cap_tests {
352 use super::*;
353
354 #[test]
355 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
356 assert_eq!(
357 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
358 crate::query::MAX_RAW_INDEX_CANDIDATES,
359 );
360 let context = crate::query::AiExecutionContext::with_limits(
361 std::time::Duration::from_secs(1),
362 usize::MAX,
363 17,
364 );
365 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
366 }
367}
368
369fn civil_from_days(z: i64) -> (i64, u32, u32) {
370 let z = z + 719_468;
371 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
372 let doe = z - era * 146_097;
373 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
374 let y = yoe + era * 400;
375 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
376 let mp = (5 * doy + 2) / 153;
377 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
378 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
379 (if m <= 2 { y + 1 } else { y }, m, d)
380}
381
382const 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;
389
390#[derive(Clone, Copy, Debug)]
405struct AutoIncState {
406 column_id: u16,
407 next: i64,
408 seeded: bool,
409}
410
411pub(crate) struct RecoveryMetadataPlan {
412 live_count: u64,
413 auto_inc: Option<AutoIncState>,
414 changed: bool,
415}
416
417type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
418
419fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
422 schema.auto_increment_column().map(|c| AutoIncState {
423 column_id: c.id,
424 next: 0,
425 seeded: false,
426 })
427}
428
429#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
442pub enum IndexBuildPolicy {
443 #[default]
445 Deferred,
446 Eager,
448}
449
450#[derive(Clone)]
451struct ReversePkSegment {
452 values: HashMap<RowId, Vec<u8>>,
453 removed: HashSet<RowId>,
454}
455
456#[derive(Clone)]
457struct ReversePkMap {
458 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
459 active: ReversePkSegment,
460}
461
462impl ReversePkMap {
463 fn new() -> Self {
464 Self {
465 frozen: Arc::new(Vec::new()),
466 active: ReversePkSegment {
467 values: HashMap::new(),
468 removed: HashSet::new(),
469 },
470 }
471 }
472
473 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
474 let mut map = Self::new();
475 map.active.values.extend(entries);
476 map
477 }
478
479 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
480 self.active.removed.remove(&row_id);
481 self.active.values.insert(row_id, key);
482 }
483
484 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
485 if let Some(key) = self.active.values.get(row_id) {
486 return Some(key);
487 }
488 if self.active.removed.contains(row_id) {
489 return None;
490 }
491 for segment in self.frozen.iter().rev() {
492 if let Some(key) = segment.values.get(row_id) {
493 return Some(key);
494 }
495 if segment.removed.contains(row_id) {
496 return None;
497 }
498 }
499 None
500 }
501
502 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
503 let previous = self.get(row_id).cloned();
504 self.active.values.remove(row_id);
505 self.active.removed.insert(*row_id);
506 previous
507 }
508
509 fn clear(&mut self) {
510 *self = Self::new();
511 }
512
513 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
514 let mut entries = HashMap::new();
515 for segment in self
516 .frozen
517 .iter()
518 .map(Arc::as_ref)
519 .chain(std::iter::once(&self.active))
520 {
521 for row_id in &segment.removed {
522 entries.remove(row_id);
523 }
524 entries.extend(
525 segment
526 .values
527 .iter()
528 .map(|(row_id, key)| (*row_id, key.clone())),
529 );
530 }
531 entries
532 }
533
534 fn seal(&mut self) {
535 if self.active.values.is_empty() && self.active.removed.is_empty() {
536 return;
537 }
538 let active = std::mem::replace(
539 &mut self.active,
540 ReversePkSegment {
541 values: HashMap::new(),
542 removed: HashSet::new(),
543 },
544 );
545 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
546 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
547 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
548 values: self.entries(),
549 removed: HashSet::new(),
550 })]);
551 }
552 }
553}
554
555#[derive(Clone)]
557pub struct Table {
558 dir: PathBuf,
559 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
560 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
561 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
562 table_id: u64,
563 name: String,
567 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
572 read_only: bool,
575 durable_commit_failed: bool,
579 wal: WalSink,
580 memtable: Memtable,
581 mutable_run: MutableRun,
586 mutable_run_spill_bytes: u64,
588 compaction_zstd_level: i32,
591 allocator: RowIdAllocator,
592 epoch: Arc<EpochAuthority>,
593 data_generation: u64,
596 schema: Schema,
597 hot: HotIndex,
598 kek: Option<Arc<Kek>>,
601 column_keys: HashMap<u16, ([u8; 32], u8)>,
605 run_refs: Vec<RunRef>,
606 retiring: Vec<crate::manifest::RetiredRun>,
609 next_run_id: u64,
610 sync_byte_threshold: u64,
611 current_txn_id: u64,
616 pending_private_mutations: bool,
620 bitmap: HashMap<u16, BitmapIndex>,
621 ann: HashMap<u16, AnnIndex>,
622 fm: HashMap<u16, FmIndex>,
623 sparse: HashMap<u16, SparseIndex>,
624 minhash: HashMap<u16, MinHashIndex>,
625 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
628 pk_by_row: ReversePkMap,
630 pinned: BTreeMap<Epoch, usize>,
633 pub(crate) live_count: u64,
636 reservoir: crate::reservoir::Reservoir,
639 reservoir_complete: bool,
647 had_deletes: bool,
651 agg_cache: Arc<HashMap<u64, CachedAgg>>,
655 global_idx_epoch: u64,
659 indexes_complete: bool,
664 index_build_policy: IndexBuildPolicy,
666 pk_by_row_complete: bool,
673 flushed_epoch: u64,
676 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
679 snapshots: Arc<crate::retention::SnapshotRegistry>,
682 commit_lock: Arc<parking_lot::Mutex<()>>,
684 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
687 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
697 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
706 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
708 pending_delete_rids: roaring::RoaringBitmap,
711 pending_put_cols: std::collections::HashSet<u16>,
714 pending_rows: Vec<Row>,
720 pending_rows_auto_inc: Vec<bool>,
721 pending_dels: Vec<RowId>,
724 pending_truncate: Option<Epoch>,
728 auto_inc: Option<AutoIncState>,
731 ttl: Option<TtlPolicy>,
734}
735
736const _: () = {
743 const fn assert_sync<T: ?Sized + Sync>() {}
744 assert_sync::<Table>();
745};
746
747enum CachedData {
753 Rows(Arc<Vec<Row>>),
754 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
755}
756
757impl CachedData {
758 fn approx_bytes(&self) -> u64 {
759 match self {
760 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
761 CachedData::Columns(c) => c
762 .iter()
763 .map(|(_, c)| c.approx_bytes())
764 .sum::<u64>()
765 .saturating_add(c.len() as u64 * 16),
766 }
767 }
768}
769
770struct CachedEntry {
774 data: CachedData,
775 footprint: roaring::RoaringBitmap,
776 condition_cols: Vec<u16>,
777}
778
779struct ResultCache {
790 entries: std::collections::HashMap<u64, CachedEntry>,
791 order: std::collections::VecDeque<u64>,
792 bytes: u64,
793 max_bytes: u64,
794 dir: Option<std::path::PathBuf>,
795 #[allow(dead_code)]
796 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
797}
798
799#[derive(serde::Serialize, serde::Deserialize)]
801struct SerializedEntry {
802 condition_cols: Vec<u16>,
803 footprint_bits: Vec<u32>,
804 data: SerializedData,
805}
806
807#[derive(serde::Serialize, serde::Deserialize)]
808enum SerializedData {
809 Rows(Vec<Row>),
810 Columns(Vec<(u16, columnar::NativeColumn)>),
811}
812
813impl SerializedEntry {
814 fn from_entry(entry: &CachedEntry) -> Self {
815 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
816 let data = match &entry.data {
817 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
818 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
819 };
820 Self {
821 condition_cols: entry.condition_cols.clone(),
822 footprint_bits,
823 data,
824 }
825 }
826
827 fn into_entry(self) -> Option<CachedEntry> {
828 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
829 let data = match self.data {
830 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
831 SerializedData::Columns(c) => {
832 if !c.iter().all(|(_, col)| col.validate()) {
835 return None;
836 }
837 CachedData::Columns(Arc::new(c))
838 }
839 };
840 Some(CachedEntry {
841 data,
842 footprint,
843 condition_cols: self.condition_cols,
844 })
845 }
846}
847
848impl ResultCache {
849 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
850
851 fn new() -> Self {
852 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
853 }
854
855 fn with_max_bytes(max_bytes: u64) -> Self {
856 Self {
857 entries: std::collections::HashMap::new(),
858 order: std::collections::VecDeque::new(),
859 bytes: 0,
860 max_bytes,
861 dir: None,
862 cache_dek: None,
863 }
864 }
865
866 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
867 let _ = std::fs::create_dir_all(&dir);
868 self.dir = Some(dir);
869 self
870 }
871
872 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
873 self.cache_dek = dek;
874 self
875 }
876
877 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
878 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
879 }
880
881 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
885 let Some(path) = self.disk_path(key) else {
886 return;
887 };
888 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
889 Ok(s) => s,
890 Err(_) => return,
891 };
892 let on_disk = if let Some(dek) = &self.cache_dek {
894 match self.encrypt_cache(&serialized, dek) {
895 Some(b) => b,
896 None => return,
897 }
898 } else {
899 serialized
900 };
901 let tmp = path.with_extension("tmp");
902 use std::io::Write;
903 let write = || -> std::io::Result<()> {
904 let mut f = std::fs::File::create(&tmp)?;
905 f.write_all(&on_disk)?;
906 f.flush()?;
907 Ok(())
908 };
909 if write().is_err() {
910 let _ = std::fs::remove_file(&tmp);
911 return;
912 }
913 let _ = std::fs::rename(&tmp, &path);
914 }
915
916 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
918 let path = self.disk_path(key)?;
919 let bytes = std::fs::read(&path).ok()?;
920 let plaintext = if let Some(dek) = &self.cache_dek {
921 self.decrypt_cache(&bytes, dek)?
922 } else {
923 bytes
924 };
925 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
926 serialized.into_entry()
927 }
928
929 fn remove_from_disk(&self, key: u64) {
931 if let Some(path) = self.disk_path(key) {
932 let _ = std::fs::remove_file(&path);
933 }
934 }
935
936 #[cfg(feature = "encryption")]
938 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
939 use crate::encryption::Cipher;
940 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
941 let mut nonce = [0u8; 12];
942 crate::encryption::fill_random(&mut nonce).ok()?;
943 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
944 let mut out = Vec::with_capacity(12 + ct.len());
945 out.extend_from_slice(&nonce);
946 out.extend_from_slice(&ct);
947 Some(out)
948 }
949
950 #[cfg(not(feature = "encryption"))]
951 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
952 None
953 }
954
955 #[cfg(feature = "encryption")]
957 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
958 use crate::encryption::Cipher;
959 if bytes.len() < 28 {
960 return None;
961 }
962 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
963 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
964 let ct = &bytes[12..];
965 cipher.decrypt_page(&nonce, ct).ok()
966 }
967
968 #[cfg(not(feature = "encryption"))]
969 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
970 None
971 }
972
973 fn load_persistent(&mut self) {
976 let Some(dir) = self.dir.as_ref().cloned() else {
977 return;
978 };
979 let entries = match std::fs::read_dir(&dir) {
980 Ok(e) => e,
981 Err(_) => return,
982 };
983 for entry in entries.flatten() {
984 let path = entry.path();
985 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
987 let _ = std::fs::remove_file(&path);
988 continue;
989 }
990 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
991 continue;
992 }
993 let stem = match path.file_stem().and_then(|s| s.to_str()) {
994 Some(s) => s,
995 None => continue,
996 };
997 let key = match u64::from_str_radix(stem, 16) {
998 Ok(k) => k,
999 Err(_) => continue,
1000 };
1001 let bytes = match std::fs::read(&path) {
1002 Ok(b) => b,
1003 Err(_) => continue,
1004 };
1005 let plaintext = if let Some(dek) = &self.cache_dek {
1007 match self.decrypt_cache(&bytes, dek) {
1008 Some(p) => p,
1009 None => {
1010 let _ = std::fs::remove_file(&path);
1011 continue;
1012 }
1013 }
1014 } else {
1015 bytes
1016 };
1017 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1018 Ok(serialized) => {
1019 if let Some(entry) = serialized.into_entry() {
1020 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1021 self.entries.insert(key, entry);
1022 self.order.push_back(key);
1023 } else {
1024 let _ = std::fs::remove_file(&path);
1025 }
1026 }
1027 Err(_) => {
1028 let _ = std::fs::remove_file(&path);
1029 }
1030 }
1031 }
1032 self.evict();
1033 }
1034
1035 fn set_max_bytes(&mut self, max_bytes: u64) {
1036 self.max_bytes = max_bytes;
1037 self.evict();
1038 }
1039
1040 fn touch(&mut self, key: u64) {
1042 self.order.retain(|k| *k != key);
1043 self.order.push_back(key);
1044 }
1045
1046 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1047 let res = self.entries.get(&key).and_then(|e| match &e.data {
1048 CachedData::Rows(r) => Some(r.clone()),
1049 CachedData::Columns(_) => None,
1050 });
1051 if res.is_some() {
1052 self.touch(key);
1053 return res;
1054 }
1055 if let Some(entry) = self.load_from_disk(key) {
1057 let res = match &entry.data {
1058 CachedData::Rows(r) => Some(r.clone()),
1059 CachedData::Columns(_) => None,
1060 };
1061 if res.is_some() {
1062 let approx = entry.data.approx_bytes();
1063 self.bytes = self.bytes.saturating_add(approx);
1064 self.entries.insert(key, entry);
1065 self.order.push_back(key);
1066 self.evict();
1067 return res;
1068 }
1069 }
1070 None
1071 }
1072
1073 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1074 let res = self.entries.get(&key).and_then(|e| match &e.data {
1075 CachedData::Columns(c) => Some(c.clone()),
1076 CachedData::Rows(_) => None,
1077 });
1078 if res.is_some() {
1079 self.touch(key);
1080 return res;
1081 }
1082 if let Some(entry) = self.load_from_disk(key) {
1084 let res = match &entry.data {
1085 CachedData::Columns(c) => Some(c.clone()),
1086 CachedData::Rows(_) => None,
1087 };
1088 if res.is_some() {
1089 let approx = entry.data.approx_bytes();
1090 self.bytes = self.bytes.saturating_add(approx);
1091 self.entries.insert(key, entry);
1092 self.order.push_back(key);
1093 self.evict();
1094 return res;
1095 }
1096 }
1097 None
1098 }
1099
1100 fn insert(&mut self, key: u64, entry: CachedEntry) {
1101 let approx = entry.data.approx_bytes();
1102 if self.entries.remove(&key).is_some() {
1103 self.order.retain(|k| *k != key);
1104 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
1105 }
1106 self.store_to_disk(key, &entry);
1108 self.bytes = self.bytes.saturating_add(approx);
1109 self.entries.insert(key, entry);
1110 self.order.push_back(key);
1111 self.evict();
1112 }
1113
1114 fn invalidate(
1123 &mut self,
1124 delete_rids: &roaring::RoaringBitmap,
1125 put_cols: &std::collections::HashSet<u16>,
1126 ) {
1127 if self.entries.is_empty() {
1128 return;
1129 }
1130 let has_deletes = !delete_rids.is_empty();
1131 let to_remove: std::collections::HashSet<u64> = self
1132 .entries
1133 .iter()
1134 .filter(|(_, e)| {
1135 let delete_hit = if e.footprint.is_empty() {
1136 has_deletes
1137 } else {
1138 e.footprint.intersection_len(delete_rids) > 0
1139 };
1140 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
1141 delete_hit || col_hit
1142 })
1143 .map(|(&k, _)| k)
1144 .collect();
1145 for key in &to_remove {
1146 if let Some(e) = self.entries.remove(key) {
1147 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1148 }
1149 self.remove_from_disk(*key);
1150 }
1151 if !to_remove.is_empty() {
1152 self.order.retain(|k| !to_remove.contains(k));
1153 }
1154 }
1155
1156 fn clear(&mut self) {
1157 if let Some(dir) = &self.dir {
1159 if let Ok(entries) = std::fs::read_dir(dir) {
1160 for entry in entries.flatten() {
1161 let path = entry.path();
1162 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1163 let _ = std::fs::remove_file(&path);
1164 }
1165 }
1166 }
1167 }
1168 self.entries.clear();
1169 self.order.clear();
1170 self.bytes = 0;
1171 }
1172
1173 fn evict(&mut self) {
1174 while self.bytes > self.max_bytes {
1175 let Some(k) = self.order.pop_front() else {
1176 break;
1177 };
1178 if let Some(e) = self.entries.remove(&k) {
1179 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
1180 self.remove_from_disk(k);
1184 }
1185 }
1186 }
1187}
1188
1189type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
1196
1197fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
1198 let _ = kek;
1199 #[cfg(feature = "encryption")]
1200 {
1201 if let Some(k) = kek {
1202 return (
1203 Some(k.derive_table_wal_key(_table_id)),
1204 Some(k.derive_cache_key()),
1205 );
1206 }
1207 }
1208 (None, None)
1209}
1210
1211#[cfg(feature = "encryption")]
1212fn read_table_encryption_salt_root(
1213 root: &crate::durable_file::DurableRoot,
1214) -> Result<[u8; crate::encryption::SALT_LEN]> {
1215 use std::io::Read;
1216
1217 let mut file = root
1218 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
1219 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
1220 let length = file.metadata()?.len();
1221 if length != crate::encryption::SALT_LEN as u64 {
1222 return Err(MongrelError::InvalidArgument(format!(
1223 "salt file is {length} bytes, expected {}",
1224 crate::encryption::SALT_LEN
1225 )));
1226 }
1227 let mut salt = [0_u8; crate::encryption::SALT_LEN];
1228 file.read_exact(&mut salt)?;
1229 Ok(salt)
1230}
1231
1232#[cfg(feature = "encryption")]
1234fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1235 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
1236}
1237
1238#[cfg(not(feature = "encryption"))]
1239fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
1240 Box::new(crate::encryption::PlaintextCipher)
1241}
1242
1243fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
1244 let Some(kek) = kek else {
1245 return HashMap::new();
1246 };
1247 #[cfg(feature = "encryption")]
1248 {
1249 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
1250 schema
1251 .columns
1252 .iter()
1253 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
1254 .map(|c| {
1255 let scheme = if schema
1256 .indexes
1257 .iter()
1258 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
1259 {
1260 SCHEME_OPE_RANGE
1261 } else {
1262 SCHEME_HMAC_EQ
1263 };
1264 let key: [u8; 32] = *kek.derive_column_key(c.id);
1265 (c.id, (key, scheme))
1266 })
1267 .collect()
1268 }
1269 #[cfg(not(feature = "encryption"))]
1270 {
1271 let _ = (kek, schema);
1272 HashMap::new()
1273 }
1274}
1275
1276pub(crate) struct SharedCtx {
1281 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1282 pub epoch: Arc<EpochAuthority>,
1283 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1284 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1285 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
1286 pub kek: Option<Arc<Kek>>,
1287 pub commit_lock: Arc<parking_lot::Mutex<()>>,
1293 pub shared: Option<SharedWalCtx>,
1297 pub table_name: Option<String>,
1300 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1303 pub read_only: bool,
1305}
1306
1307#[derive(Clone)]
1313pub(crate) struct SharedWalCtx {
1314 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
1315 pub group: Arc<GroupCommit>,
1316 pub poisoned: Arc<AtomicBool>,
1317 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
1318 pub change_wake: tokio::sync::broadcast::Sender<()>,
1319}
1320
1321enum WalSink {
1324 Private(Wal),
1325 Shared(SharedWalCtx),
1326 ReadOnly,
1327}
1328
1329impl Clone for WalSink {
1330 fn clone(&self) -> Self {
1331 match self {
1332 Self::Shared(shared) => Self::Shared(shared.clone()),
1333 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
1334 }
1335 }
1336}
1337
1338impl SharedCtx {
1339 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
1343 let n_shards = if cache_dir.is_some() {
1347 1
1348 } else {
1349 crate::cache::CACHE_SHARDS
1350 };
1351 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
1352 let page_cache = if let Some(d) = cache_dir {
1353 Arc::new(crate::cache::Sharded::new(1, || {
1354 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
1355 }))
1356 } else {
1357 Arc::new(crate::cache::Sharded::new(n_shards, || {
1358 crate::cache::PageCache::new(per_shard)
1359 }))
1360 };
1361 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
1362 let decoded_cache = Arc::new(crate::cache::Sharded::new(
1363 crate::cache::CACHE_SHARDS,
1364 || crate::cache::DecodedPageCache::new(decoded_per_shard),
1365 ));
1366 Self {
1367 root_guard: None,
1368 epoch: Arc::new(EpochAuthority::new(0)),
1369 page_cache,
1370 decoded_cache,
1371 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
1372 kek,
1373 commit_lock: Arc::new(parking_lot::Mutex::new(())),
1374 shared: None,
1375 table_name: None,
1376 auth: None,
1377 read_only: false,
1378 }
1379 }
1380}
1381
1382fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
1386 use crate::query::Condition;
1387 match c {
1388 Condition::Pk(_)
1390 | Condition::BitmapEq { .. }
1391 | Condition::BitmapIn { .. }
1392 | Condition::BytesPrefix { .. }
1393 | Condition::IsNull { .. }
1394 | Condition::IsNotNull { .. } => 0,
1395 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
1397 1
1398 }
1399 Condition::FmContains { .. }
1401 | Condition::FmContainsAll { .. }
1402 | Condition::Ann { .. }
1403 | Condition::SparseMatch { .. } => 2,
1404 }
1405}
1406
1407impl Table {
1408 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
1409 let dir = dir.as_ref().to_path_buf();
1410 crate::durable_file::create_directory_all(&dir)?;
1411 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1412 let pinned = root.io_path()?;
1413 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1414 ctx.root_guard = Some(root);
1415 Self::create_in(&pinned, schema, table_id, ctx)
1416 }
1417
1418 #[cfg(feature = "encryption")]
1429 pub fn create_encrypted(
1430 dir: impl AsRef<Path>,
1431 schema: Schema,
1432 table_id: u64,
1433 passphrase: &str,
1434 ) -> Result<Self> {
1435 let dir = dir.as_ref().to_path_buf();
1436 crate::durable_file::create_directory_all(&dir)?;
1437 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1438 root.create_directory_all(META_DIR)?;
1439 let salt = crate::encryption::random_salt()?;
1440 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1441 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1442 let pinned = root.io_path()?;
1443 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1444 ctx.root_guard = Some(root);
1445 Self::create_in(&pinned, schema, table_id, ctx)
1446 }
1447
1448 #[cfg(feature = "encryption")]
1453 pub fn create_with_key(
1454 dir: impl AsRef<Path>,
1455 schema: Schema,
1456 table_id: u64,
1457 key: &[u8],
1458 ) -> Result<Self> {
1459 let dir = dir.as_ref().to_path_buf();
1460 crate::durable_file::create_directory_all(&dir)?;
1461 let root = Arc::new(crate::durable_file::DurableRoot::open(&dir)?);
1462 root.create_directory_all(META_DIR)?;
1463 let salt = crate::encryption::random_salt()?;
1464 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
1465 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
1466 let pinned = root.io_path()?;
1467 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1468 ctx.root_guard = Some(root);
1469 Self::create_in(&pinned, schema, table_id, ctx)
1470 }
1471
1472 #[cfg(feature = "encryption")]
1474 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
1475 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1476 let salt = read_table_encryption_salt_root(&root)?;
1477 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1478 let pinned = root.io_path()?;
1479 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1480 ctx.root_guard = Some(root);
1481 Self::open_in(&pinned, ctx)
1482 }
1483
1484 pub(crate) fn create_in(
1485 dir: impl AsRef<Path>,
1486 schema: Schema,
1487 table_id: u64,
1488 ctx: SharedCtx,
1489 ) -> Result<Self> {
1490 schema.validate_auto_increment()?;
1491 schema.validate_defaults()?;
1492 schema.validate_ai()?;
1493 for index in &schema.indexes {
1494 index.validate_options()?;
1495 }
1496 let dir = dir.as_ref().to_path_buf();
1497 let runs_root = match ctx.root_guard.as_ref() {
1498 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
1499 None => {
1500 crate::durable_file::create_directory_all(&dir)?;
1501 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
1502 None
1503 }
1504 };
1505 match ctx.root_guard.as_deref() {
1506 Some(root) => write_schema_durable(root, &schema)?,
1507 None => write_schema(&dir, &schema)?,
1508 }
1509 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1510 let (wal, current_txn_id) = match ctx.shared.clone() {
1513 Some(s) => (WalSink::Shared(s), 0),
1514 None => {
1515 let pinned_wal_root = match ctx.root_guard.as_deref() {
1516 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
1517 None => None,
1518 };
1519 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
1520 root.io_path()?
1521 } else {
1522 let wal_dir = dir.join(WAL_DIR);
1523 crate::durable_file::create_directory_all(&wal_dir)?;
1524 wal_dir
1525 };
1526 let mut w = if let Some(ref dk) = wal_dek {
1527 Wal::create_with_cipher(
1528 wal_dir.join("seg-000000.wal"),
1529 Epoch(0),
1530 Some(make_cipher(dk)),
1531 0,
1532 )?
1533 } else {
1534 Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?
1535 };
1536 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1537 (WalSink::Private(w), 1)
1538 }
1539 };
1540 let mut manifest = Manifest::new(table_id, schema.schema_id);
1541 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1546 match ctx.root_guard.as_deref() {
1547 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
1548 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
1549 }
1550 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1551 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1552 let auto_inc = resolve_auto_inc(&schema);
1553 let rcache_dir = dir.join(RCACHE_DIR);
1554 Ok(Self {
1555 dir,
1556 _root_guard: ctx.root_guard,
1557 runs_root,
1558 idx_root: None,
1559 table_id,
1560 name: ctx.table_name.unwrap_or_default(),
1561 auth: ctx.auth,
1562 read_only: ctx.read_only,
1563 durable_commit_failed: false,
1564 wal,
1565 memtable: Memtable::new(),
1566 mutable_run: MutableRun::new(),
1567 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1568 compaction_zstd_level: 3,
1569 allocator: RowIdAllocator::new(0),
1570 epoch: ctx.epoch,
1571 data_generation: 0,
1572 schema,
1573 hot: HotIndex::new(),
1574 kek: ctx.kek,
1575 column_keys,
1576 run_refs: Vec::new(),
1577 retiring: Vec::new(),
1578 next_run_id: 1,
1579 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1580 current_txn_id,
1581 pending_private_mutations: false,
1582 bitmap,
1583 ann,
1584 fm,
1585 sparse,
1586 minhash,
1587 learned_range: Arc::new(HashMap::new()),
1588 pk_by_row: ReversePkMap::new(),
1589 pinned: BTreeMap::new(),
1590 live_count: 0,
1591 reservoir: crate::reservoir::Reservoir::default(),
1592 reservoir_complete: true,
1593 had_deletes: false,
1594 agg_cache: Arc::new(HashMap::new()),
1595 global_idx_epoch: 0,
1596 indexes_complete: true,
1597 index_build_policy: IndexBuildPolicy::default(),
1598 pk_by_row_complete: false,
1599 flushed_epoch: 0,
1600 page_cache: ctx.page_cache,
1601 decoded_cache: ctx.decoded_cache,
1602 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1603 snapshots: ctx.snapshots,
1604 commit_lock: ctx.commit_lock,
1605 result_cache: Arc::new(parking_lot::Mutex::new(
1606 ResultCache::new()
1607 .with_dir(rcache_dir)
1608 .with_cache_dek(cache_dek.clone()),
1609 )),
1610 pending_delete_rids: roaring::RoaringBitmap::new(),
1611 pending_put_cols: std::collections::HashSet::new(),
1612 pending_rows: Vec::new(),
1613 pending_rows_auto_inc: Vec::new(),
1614 pending_dels: Vec::new(),
1615 pending_truncate: None,
1616 wal_dek,
1617 auto_inc,
1618 ttl: None,
1619 })
1620 }
1621
1622 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1626 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1627 let pinned = root.io_path()?;
1628 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
1629 ctx.root_guard = Some(root);
1630 Self::open_in(&pinned, ctx)
1631 }
1632
1633 #[cfg(feature = "encryption")]
1636 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1637 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
1638 let salt = read_table_encryption_salt_root(&root)?;
1639 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1640 let pinned = root.io_path()?;
1641 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
1642 ctx.root_guard = Some(root);
1643 let t = Self::open_in(&pinned, ctx)?;
1644 Ok(t)
1645 }
1646
1647 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1648 let dir = dir.as_ref().to_path_buf();
1649 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1650 let mut manifest = match ctx.root_guard.as_ref() {
1651 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
1652 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
1653 };
1654 let schema: Schema = match ctx.root_guard.as_ref() {
1655 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
1656 None => read_schema(&dir)?,
1657 };
1658 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
1664 let runs_root = match ctx.root_guard.as_ref() {
1665 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
1666 None => None,
1667 };
1668 let idx_root = match ctx.root_guard.as_ref() {
1669 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
1670 Ok(root) => Some(Arc::new(root)),
1671 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
1672 Err(error) => return Err(error.into()),
1673 },
1674 None => None,
1675 };
1676 schema.validate_auto_increment()?;
1677 schema.validate_defaults()?;
1678 schema.validate_ai()?;
1679 for index in &schema.indexes {
1680 index.validate_options()?;
1681 }
1682 let replay_epoch = Epoch(manifest.current_epoch);
1683 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1684 let private_replayed = if ctx.shared.is_none() {
1685 match latest_wal_segment(&dir.join(WAL_DIR))? {
1686 Some(path) => {
1687 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1688 crate::wal::replay_with_cipher(path, cipher)?
1689 }
1690 None => Vec::new(),
1691 }
1692 } else {
1693 Vec::new()
1694 };
1695 if ctx.shared.is_none() {
1696 preflight_standalone_open(
1697 &dir,
1698 runs_root.as_deref(),
1699 idx_root.as_deref(),
1700 &manifest,
1701 &schema,
1702 &private_replayed,
1703 ctx.kek.clone(),
1704 )?;
1705 }
1706 let next_run_id = derive_next_run_id(
1707 &dir,
1708 runs_root.as_deref(),
1709 &manifest.runs,
1710 &manifest.retiring,
1711 )?;
1712 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1716 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1717 None => {
1718 let replayed = private_replayed;
1719 let wal_dir = dir.join(WAL_DIR);
1725 crate::durable_file::create_directory_all(&wal_dir)?;
1726 let segment = next_wal_segment(&wal_dir)?;
1727 let segment_no = wal_segment_number(&segment).unwrap_or(0);
1728 let temporary = wal_dir.join(format!(
1729 ".recovery-{}-{}-{segment_no:06}.tmp",
1730 std::process::id(),
1731 std::time::SystemTime::now()
1732 .duration_since(std::time::UNIX_EPOCH)
1733 .unwrap_or_default()
1734 .as_nanos()
1735 ));
1736 let mut w = Wal::create_with_cipher(
1737 &temporary,
1738 replay_epoch,
1739 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1740 segment_no,
1741 )?;
1742 for record in &replayed {
1743 w.append_txn(record.txn_id, record.op.clone())?;
1744 }
1745 let mut w = w.publish_as(segment)?;
1746 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1747 let next_txn_id = replayed
1748 .iter()
1749 .map(|record| record.txn_id)
1750 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
1751 .max()
1752 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
1753 .unwrap_or(1);
1754 (WalSink::Private(w), replayed, next_txn_id)
1755 }
1756 };
1757
1758 let mut memtable = Memtable::new();
1759 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1760 let persisted_epoch = manifest.current_epoch;
1761 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1768 s.next = manifest.auto_inc_next;
1769 s.seeded = manifest.auto_inc_next > 0;
1770 s
1771 });
1772
1773 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1780 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1781 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
1782 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1783 std::collections::BTreeMap::new();
1784 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1785 let mut recovered_epoch = manifest.current_epoch;
1786 let mut recovered_manifest_dirty = schema_manifest_repair;
1787 let mut saw_delete = false;
1788 for record in replayed {
1789 let txn_id = record.txn_id;
1790 match record.op {
1791 Op::Put { rows, .. } => {
1792 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1793 for row in &rows {
1794 allocator.advance_to(row.row_id)?;
1795 if let Some(ai) = auto_inc.as_mut() {
1796 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1797 let next = n.checked_add(1).ok_or_else(|| {
1798 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
1799 })?;
1800 if next > ai.next {
1801 ai.next = next;
1802 }
1803 }
1804 }
1805 }
1806 staged_puts.entry(txn_id).or_default().extend(rows);
1807 }
1808 Op::Delete { row_ids, .. } => {
1809 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1810 }
1811 Op::TxnCommit { epoch, .. } => {
1812 let commit_epoch = Epoch(epoch);
1813 recovered_epoch = recovered_epoch.max(epoch);
1814 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
1815 memtable = Memtable::new();
1816 replayed_puts.clear();
1817 replayed_deletes.clear();
1818 manifest.runs.clear();
1819 manifest.retiring.clear();
1820 manifest.live_count = 0;
1821 manifest.global_idx_epoch = 0;
1822 manifest.current_epoch = manifest.current_epoch.max(epoch);
1823 recovered_manifest_dirty = true;
1824 saw_delete = true;
1825 }
1826 if let Some(puts) = staged_puts.remove(&txn_id) {
1827 if commit_epoch.0 > manifest.flushed_epoch {
1828 for row in &puts {
1829 memtable.upsert(row.clone());
1830 }
1831 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1832 }
1833 }
1834 if let Some(dels) = staged_deletes.remove(&txn_id) {
1835 saw_delete = true;
1836 if commit_epoch.0 > manifest.flushed_epoch {
1837 for rid in dels {
1838 memtable.tombstone(rid, commit_epoch);
1839 replayed_deletes.push((rid, commit_epoch));
1840 }
1841 }
1842 }
1843 }
1844 Op::TxnAbort => {
1845 staged_puts.remove(&txn_id);
1846 staged_deletes.remove(&txn_id);
1847 staged_truncates.remove(&txn_id);
1848 }
1849 Op::TruncateTable { .. } => {
1850 staged_puts.remove(&txn_id);
1851 staged_deletes.remove(&txn_id);
1852 staged_truncates.insert(txn_id);
1853 }
1854 Op::ExternalTableState { .. }
1855 | Op::Flush { .. }
1856 | Op::Ddl(_)
1857 | Op::BeforeImage { .. }
1858 | Op::CommitTimestamp { .. }
1859 | Op::SpilledRows { .. } => {}
1860 }
1861 }
1862
1863 let rcache_dir = dir.join(RCACHE_DIR);
1864 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1865 let mut db = Self {
1866 dir,
1867 _root_guard: ctx.root_guard,
1868 runs_root,
1869 idx_root,
1870 table_id: manifest.table_id,
1871 name: ctx.table_name.unwrap_or_default(),
1872 auth: ctx.auth,
1873 read_only: ctx.read_only,
1874 durable_commit_failed: false,
1875 wal,
1876 memtable,
1877 mutable_run: MutableRun::new(),
1878 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1879 compaction_zstd_level: 3,
1880 allocator,
1881 epoch: ctx.epoch,
1882 data_generation: persisted_epoch,
1883 schema,
1884 hot: HotIndex::new(),
1885 kek: ctx.kek,
1886 column_keys,
1887 run_refs: manifest.runs.clone(),
1888 retiring: manifest.retiring.clone(),
1889 next_run_id,
1890 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1891 current_txn_id,
1892 pending_private_mutations: false,
1893 bitmap: HashMap::new(),
1894 ann: HashMap::new(),
1895 fm: HashMap::new(),
1896 sparse: HashMap::new(),
1897 minhash: HashMap::new(),
1898 learned_range: Arc::new(HashMap::new()),
1899 pk_by_row: ReversePkMap::new(),
1900 pinned: BTreeMap::new(),
1901 live_count: manifest.live_count,
1902 reservoir: crate::reservoir::Reservoir::default(),
1903 reservoir_complete: false,
1904 had_deletes: saw_delete
1905 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
1906 != manifest.live_count,
1907 agg_cache: Arc::new(HashMap::new()),
1908 global_idx_epoch: manifest.global_idx_epoch,
1909 indexes_complete: true,
1910 index_build_policy: IndexBuildPolicy::default(),
1911 pk_by_row_complete: false,
1912 flushed_epoch: manifest.flushed_epoch,
1913 page_cache: ctx.page_cache,
1914 decoded_cache: ctx.decoded_cache,
1915 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1916 snapshots: ctx.snapshots,
1917 commit_lock: ctx.commit_lock,
1918 result_cache: Arc::new(parking_lot::Mutex::new(
1919 ResultCache::new()
1920 .with_dir(rcache_dir)
1921 .with_cache_dek(cache_dek.clone()),
1922 )),
1923 pending_delete_rids: roaring::RoaringBitmap::new(),
1924 pending_put_cols: std::collections::HashSet::new(),
1925 pending_rows: Vec::new(),
1926 pending_rows_auto_inc: Vec::new(),
1927 pending_dels: Vec::new(),
1928 pending_truncate: None,
1929 wal_dek,
1930 auto_inc,
1931 ttl: manifest.ttl,
1932 };
1933
1934 db.epoch.advance_recovered(Epoch(recovered_epoch));
1937
1938 let checkpoint = match db.idx_root.as_deref() {
1943 Some(root) => {
1944 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
1945 }
1946 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
1947 };
1948 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1949 c.epoch_built == manifest.global_idx_epoch
1950 && manifest.global_idx_epoch > 0
1951 && manifest
1952 .runs
1953 .iter()
1954 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1955 });
1956 if let Some(loaded) = checkpoint {
1957 if checkpoint_valid {
1958 db.hot = loaded.hot;
1959 db.bitmap = loaded.bitmap;
1960 db.ann = loaded.ann;
1961 db.fm = loaded.fm;
1962 db.sparse = loaded.sparse;
1963 db.minhash = loaded.minhash;
1964 db.learned_range = Arc::new(loaded.learned_range);
1965 }
1968 }
1969 if !checkpoint_valid {
1970 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1971 db.bitmap = bitmap;
1972 db.ann = ann;
1973 db.fm = fm;
1974 db.sparse = sparse;
1975 db.minhash = minhash;
1976 db.rebuild_indexes_from_runs()?;
1977 db.build_learned_ranges()?;
1978 }
1979
1980 for (epoch, group) in replayed_puts {
1985 let (losers, winner_pks) = db.partition_pk_winners(&group);
1986 for (key, &row_id) in &winner_pks {
1987 if let Some(old_rid) = db.hot.get(key) {
1988 if old_rid != row_id {
1989 db.tombstone_row(old_rid, epoch, false);
1990 }
1991 }
1992 }
1993 for &loser_rid in &losers {
1994 db.tombstone_row(loser_rid, epoch, false);
1995 }
1996 for (key, row_id) in winner_pks {
1997 db.insert_hot_pk(key, row_id);
1998 }
1999 if db.schema.primary_key().is_none() {
2000 for r in &group {
2001 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2002 }
2003 }
2004 for r in &group {
2005 if !losers.contains(&r.row_id) {
2006 db.index_row(r);
2007 }
2008 }
2009 }
2010 for (rid, epoch) in &replayed_deletes {
2014 db.remove_hot_for_row(*rid, *epoch);
2015 }
2016
2017 if recovered_manifest_dirty {
2018 let rows = db.visible_rows(Snapshot::at(Epoch(u64::MAX)))?;
2019 db.live_count = rows.len() as u64;
2020 db.persist_manifest(Epoch(recovered_epoch))?;
2021 }
2022
2023 db.result_cache.lock().load_persistent();
2030 Ok(db)
2031 }
2032
2033 fn ensure_reservoir_complete(&mut self) -> Result<()> {
2039 if self.reservoir_complete {
2040 return Ok(());
2041 }
2042 self.rebuild_reservoir()?;
2043 self.reservoir_complete = true;
2044 Ok(())
2045 }
2046
2047 fn rebuild_reservoir(&mut self) -> Result<()> {
2050 let snap = self.snapshot();
2051 let rows = self.visible_rows(snap)?;
2052 self.reservoir.reset();
2053 for r in rows {
2054 self.reservoir.offer(r.row_id.0);
2055 }
2056 Ok(())
2057 }
2058
2059 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
2060 self.rebuild_indexes_from_runs_inner(None)
2061 }
2062
2063 fn rebuild_indexes_from_runs_inner(
2064 &mut self,
2065 control: Option<&crate::ExecutionControl>,
2066 ) -> Result<()> {
2067 self.hot = HotIndex::new();
2068 self.pk_by_row.clear();
2069 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2070 self.bitmap = bitmap;
2071 self.ann = ann;
2072 self.fm = fm;
2073 self.sparse = sparse;
2074 self.minhash = minhash;
2075 let snapshot = Epoch(u64::MAX);
2076 let ttl_now = unix_nanos_now();
2077 let mut scanned = 0_usize;
2078 for rr in self.run_refs.clone() {
2079 if let Some(control) = control {
2080 control.checkpoint()?;
2081 }
2082 let mut reader = self.open_reader(rr.run_id)?;
2083 for row in reader.visible_rows(snapshot)? {
2084 if scanned.is_multiple_of(256) {
2085 if let Some(control) = control {
2086 control.checkpoint()?;
2087 }
2088 }
2089 scanned += 1;
2090 if self.row_expired_at(&row, ttl_now) {
2091 continue;
2092 }
2093 let tok_row = self.tokenized_for_indexes(&row);
2094 index_into(
2095 &self.schema,
2096 &tok_row,
2097 &mut self.hot,
2098 &mut self.bitmap,
2099 &mut self.ann,
2100 &mut self.fm,
2101 &mut self.sparse,
2102 &mut self.minhash,
2103 );
2104 }
2105 }
2106 for row in self.mutable_run.visible_versions(snapshot) {
2107 if scanned.is_multiple_of(256) {
2108 if let Some(control) = control {
2109 control.checkpoint()?;
2110 }
2111 }
2112 scanned += 1;
2113 if row.deleted {
2114 self.remove_hot_for_row(row.row_id, snapshot);
2115 } else if !self.row_expired_at(&row, ttl_now) {
2116 self.index_row(&row);
2117 }
2118 }
2119 for row in self.memtable.visible_versions(snapshot) {
2120 if scanned.is_multiple_of(256) {
2121 if let Some(control) = control {
2122 control.checkpoint()?;
2123 }
2124 }
2125 scanned += 1;
2126 if row.deleted {
2127 self.remove_hot_for_row(row.row_id, snapshot);
2128 } else if !self.row_expired_at(&row, ttl_now) {
2129 self.index_row(&row);
2130 }
2131 }
2132 self.refresh_pk_by_row_from_hot();
2133 Ok(())
2134 }
2135
2136 fn refresh_pk_by_row_from_hot(&mut self) {
2137 self.pk_by_row_complete = true;
2138 if self.schema.primary_key().is_none() {
2139 self.pk_by_row.clear();
2140 return;
2141 }
2142 self.pk_by_row = ReversePkMap::from_entries(
2148 self.hot
2149 .entries()
2150 .into_iter()
2151 .map(|(key, row_id)| (row_id, key)),
2152 );
2153 }
2154
2155 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
2156 if self.schema.primary_key().is_some() {
2157 self.pk_by_row.insert(row_id, key.clone());
2158 }
2159 self.hot.insert(key, row_id);
2160 }
2161
2162 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
2166 self.build_learned_ranges_inner(None)
2167 }
2168
2169 fn build_learned_ranges_inner(
2170 &mut self,
2171 control: Option<&crate::ExecutionControl>,
2172 ) -> Result<()> {
2173 self.learned_range = Arc::new(HashMap::new());
2174 if self.run_refs.len() != 1 {
2175 return Ok(());
2176 }
2177 let cols: Vec<(u16, usize)> = self
2178 .schema
2179 .indexes
2180 .iter()
2181 .filter(|i| i.kind == IndexKind::LearnedRange)
2182 .map(|i| {
2183 (
2184 i.column_id,
2185 i.options
2186 .learned_range
2187 .as_ref()
2188 .map(|options| options.epsilon)
2189 .unwrap_or(16),
2190 )
2191 })
2192 .collect();
2193 if cols.is_empty() {
2194 return Ok(());
2195 }
2196 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2197 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
2198 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
2199 _ => return Ok(()),
2200 };
2201 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
2202 if column_index % 256 == 0 {
2203 if let Some(control) = control {
2204 control.checkpoint()?;
2205 }
2206 }
2207 let ty = self
2208 .schema
2209 .columns
2210 .iter()
2211 .find(|c| c.id == cid)
2212 .map(|c| c.ty.clone())
2213 .unwrap_or(TypeId::Int64);
2214 match ty {
2215 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2216 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
2217 let pairs: Vec<(i64, u64)> = data
2218 .iter()
2219 .zip(row_ids.iter())
2220 .map(|(v, r)| (*v, *r))
2221 .collect();
2222 Arc::make_mut(&mut self.learned_range).insert(
2223 cid,
2224 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
2225 );
2226 }
2227 }
2228 TypeId::Float64 => {
2229 if let columnar::NativeColumn::Float64 { data, .. } =
2230 reader.column_native(cid)?
2231 {
2232 let pairs: Vec<(f64, u64)> = data
2233 .iter()
2234 .zip(row_ids.iter())
2235 .map(|(v, r)| (*v, *r))
2236 .collect();
2237 Arc::make_mut(&mut self.learned_range).insert(
2238 cid,
2239 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
2240 );
2241 }
2242 }
2243 _ => {}
2244 }
2245 }
2246 Ok(())
2247 }
2248
2249 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
2256 if self.indexes_complete {
2257 crate::trace::QueryTrace::record(|t| {
2258 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2259 });
2260 return Ok(());
2261 }
2262 crate::trace::QueryTrace::record(|t| {
2263 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2264 });
2265 self.rebuild_indexes_from_runs()?;
2266 self.build_learned_ranges()?;
2267 self.indexes_complete = true;
2268 let epoch = self.current_epoch();
2269 self.checkpoint_indexes(epoch);
2270 Ok(())
2271 }
2272
2273 #[doc(hidden)]
2276 pub fn ensure_indexes_complete_controlled<F>(
2277 &mut self,
2278 control: &crate::ExecutionControl,
2279 before_publish: F,
2280 ) -> Result<bool>
2281 where
2282 F: FnOnce() -> bool,
2283 {
2284 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
2285 .map(|(changed, _)| changed)
2286 }
2287
2288 #[doc(hidden)]
2291 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
2292 &mut self,
2293 control: &crate::ExecutionControl,
2294 before_publish: F,
2295 ) -> Result<(bool, Option<MaintenanceReceipt>)>
2296 where
2297 F: FnOnce() -> bool,
2298 {
2299 if self.indexes_complete {
2300 crate::trace::QueryTrace::record(|trace| {
2301 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
2302 });
2303 return Ok((false, None));
2304 }
2305 crate::trace::QueryTrace::record(|trace| {
2306 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
2307 });
2308 control.checkpoint()?;
2309 let maintenance_epoch = self.current_epoch();
2310 self.rebuild_indexes_from_runs_inner(Some(control))?;
2311 self.build_learned_ranges_inner(Some(control))?;
2312 control.checkpoint()?;
2313 if !before_publish() {
2314 return Err(MongrelError::Cancelled);
2315 }
2316 self.indexes_complete = true;
2317 self.checkpoint_indexes(maintenance_epoch);
2318 Ok((
2319 true,
2320 Some(MaintenanceReceipt {
2321 epoch: maintenance_epoch,
2322 }),
2323 ))
2324 }
2325
2326 fn pending_epoch(&self) -> Epoch {
2327 Epoch(self.epoch.visible().0 + 1)
2328 }
2329
2330 fn is_shared(&self) -> bool {
2333 matches!(self.wal, WalSink::Shared(_))
2334 }
2335
2336 fn ensure_txn_id(&mut self) -> Result<u64> {
2340 if self.current_txn_id == 0 {
2341 let id = match &self.wal {
2342 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
2343 WalSink::Private(_) => {
2344 return Err(MongrelError::Full(
2345 "standalone transaction id namespace exhausted".into(),
2346 ))
2347 }
2348 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2349 };
2350 self.current_txn_id = id;
2351 }
2352 Ok(self.current_txn_id)
2353 }
2354
2355 fn wal_append_data(&mut self, op: Op) -> Result<()> {
2358 self.ensure_writable()?;
2359 let txn_id = self.ensure_txn_id()?;
2360 let table_id = self.table_id;
2361 match &mut self.wal {
2362 WalSink::Private(w) => {
2363 w.append_txn(txn_id, op)?;
2364 self.pending_private_mutations = true;
2365 }
2366 WalSink::Shared(s) => {
2367 s.wal.lock().append(txn_id, table_id, op)?;
2368 }
2369 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
2370 }
2371 Ok(())
2372 }
2373
2374 fn ensure_writable(&self) -> Result<()> {
2375 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
2376 return Err(MongrelError::ReadOnlyReplica);
2377 }
2378 if self.durable_commit_failed {
2379 return Err(MongrelError::Other(
2380 "table poisoned by post-commit failure; reopen required".into(),
2381 ));
2382 }
2383 Ok(())
2384 }
2385
2386 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
2397 match &self.auth {
2398 Some(checker) => checker.check(&self.name, perm),
2399 None => Ok(()),
2400 }
2401 }
2402 pub fn require_select(&self) -> Result<()> {
2407 self.require(crate::auth_state::RequiredPermission::Select)
2408 }
2409 fn require_insert(&self) -> Result<()> {
2410 self.require(crate::auth_state::RequiredPermission::Insert)
2411 }
2412 #[allow(dead_code)]
2416 fn require_update(&self) -> Result<()> {
2417 self.require(crate::auth_state::RequiredPermission::Update)
2418 }
2419 fn require_delete(&self) -> Result<()> {
2420 self.require(crate::auth_state::RequiredPermission::Delete)
2421 }
2422
2423 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
2426 self.require_insert()?;
2427 Ok(self.put_returning(columns)?.0)
2428 }
2429
2430 pub fn put_returning(
2435 &mut self,
2436 mut columns: Vec<(u16, Value)>,
2437 ) -> Result<(RowId, Option<i64>)> {
2438 self.require_insert()?;
2439 let assigned = self.fill_auto_inc(&mut columns)?;
2440 self.apply_defaults(&mut columns)?;
2441 self.schema.validate_values(&columns)?;
2442 let row_id = if self.schema.clustered {
2447 self.derive_clustered_row_id(&columns)?
2448 } else {
2449 self.allocator.alloc()?
2450 };
2451 let epoch = self.pending_epoch();
2452 let mut row = Row::new(row_id, epoch);
2453 for (col_id, val) in columns {
2454 row.columns.insert(col_id, val);
2455 }
2456 self.commit_rows(vec![row], assigned.is_some())?;
2457 Ok((row_id, assigned))
2458 }
2459
2460 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
2463 self.require_insert()?;
2464 Ok(self
2465 .put_batch_returning(batch)?
2466 .into_iter()
2467 .map(|(r, _)| r)
2468 .collect())
2469 }
2470
2471 pub fn put_batch_returning(
2474 &mut self,
2475 batch: Vec<Vec<(u16, Value)>>,
2476 ) -> Result<Vec<(RowId, Option<i64>)>> {
2477 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
2478 for mut cols in batch {
2479 let assigned = self.fill_auto_inc(&mut cols)?;
2480 self.apply_defaults(&mut cols)?;
2481 filled.push((cols, assigned));
2482 }
2483 for (cols, _) in &filled {
2484 self.schema.validate_values(cols)?;
2485 }
2486 let epoch = self.pending_epoch();
2487 let mut rows = Vec::with_capacity(filled.len());
2488 let mut ids = Vec::with_capacity(filled.len());
2489 let first_row_id = if self.schema.clustered {
2490 None
2491 } else {
2492 let count = u64::try_from(filled.len())
2493 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
2494 Some(self.allocator.alloc_range(count)?.0)
2495 };
2496 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
2497 let row_id = match first_row_id {
2498 Some(first) => RowId(first + row_index as u64),
2499 None => self.derive_clustered_row_id(&cols)?,
2500 };
2501 let mut row = Row::new(row_id, epoch);
2502 for (c, v) in cols {
2503 row.columns.insert(c, v);
2504 }
2505 ids.push((row_id, assigned));
2506 rows.push(row);
2507 }
2508 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
2509 self.commit_rows(rows, all_auto_generated)?;
2510 Ok(ids)
2511 }
2512
2513 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
2519 self.ensure_writable()?;
2520 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2521 return Ok(None);
2522 };
2523 let pos = columns.iter().position(|(c, _)| *c == cid);
2524 let assigned = match pos {
2525 Some(i) => match &columns[i].1 {
2526 Value::Null => {
2527 let next = self.alloc_auto_inc_value()?;
2528 columns[i].1 = Value::Int64(next);
2529 Some(next)
2530 }
2531 Value::Int64(n) => {
2532 self.advance_auto_inc_past(*n)?;
2533 None
2534 }
2535 other => {
2536 return Err(MongrelError::InvalidArgument(format!(
2537 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
2538 other
2539 )))
2540 }
2541 },
2542 None => {
2543 let next = self.alloc_auto_inc_value()?;
2544 columns.push((cid, Value::Int64(next)));
2545 Some(next)
2546 }
2547 };
2548 Ok(assigned)
2549 }
2550
2551 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
2557 for col in &self.schema.columns {
2558 let Some(expr) = &col.default_value else {
2559 continue;
2560 };
2561 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
2563 continue;
2564 }
2565 let pos = columns.iter().position(|(c, _)| *c == col.id);
2566 let needs_default = match pos {
2567 None => true,
2568 Some(i) => matches!(columns[i].1, Value::Null),
2569 };
2570 if !needs_default {
2571 continue;
2572 }
2573 let v = match expr {
2574 crate::schema::DefaultExpr::Static(v) => v.clone(),
2575 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
2576 crate::schema::DefaultExpr::Uuid => {
2577 let mut buf = [0u8; 16];
2578 getrandom::getrandom(&mut buf)
2579 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
2580 Value::Uuid(buf)
2581 }
2582 };
2583 match pos {
2584 None => columns.push((col.id, v)),
2585 Some(i) => columns[i].1 = v,
2586 }
2587 }
2588 Ok(())
2589 }
2590
2591 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
2593 self.ensure_auto_inc_seeded()?;
2594 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2596 let v = ai.next;
2597 ai.next = ai
2598 .next
2599 .checked_add(1)
2600 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2601 Ok(v)
2602 }
2603
2604 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
2607 self.ensure_auto_inc_seeded()?;
2608 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2609 let floor = used
2610 .checked_add(1)
2611 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2612 .max(1);
2613 if ai.next < floor {
2614 ai.next = floor;
2615 }
2616 Ok(())
2617 }
2618
2619 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
2624 let needs_seed = match self.auto_inc {
2625 Some(ai) => !ai.seeded,
2626 None => return Ok(()),
2627 };
2628 if !needs_seed {
2629 return Ok(());
2630 }
2631 if self.seed_empty_auto_inc() {
2632 return Ok(());
2633 }
2634 let cid = self
2635 .auto_inc
2636 .as_ref()
2637 .expect("auto-inc column present")
2638 .column_id;
2639 let max = self.scan_max_int64(cid)?;
2640 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2641 let floor = max
2642 .checked_add(1)
2643 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2644 .max(1);
2645 if ai.next < floor {
2646 ai.next = floor;
2647 }
2648 ai.seeded = true;
2649 Ok(())
2650 }
2651
2652 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
2653 if n == 0 || self.auto_inc.is_none() {
2654 return Ok(None);
2655 }
2656 self.ensure_auto_inc_seeded()?;
2657 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
2658 let start = ai.next;
2659 let count = i64::try_from(n)
2660 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
2661 ai.next = ai
2662 .next
2663 .checked_add(count)
2664 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
2665 Ok(Some(start))
2666 }
2667
2668 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
2672 let mut max: i64 = 0;
2673 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
2674 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2675 if *n > max {
2676 max = *n;
2677 }
2678 }
2679 }
2680 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
2681 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
2682 if *n > max {
2683 max = *n;
2684 }
2685 }
2686 }
2687 for rr in self.run_refs.clone() {
2688 let reader = self.open_reader(rr.run_id)?;
2689 if let Some(stats) = reader.column_page_stats(column_id) {
2690 for s in stats {
2691 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
2692 if n > max {
2693 max = n;
2694 }
2695 }
2696 }
2697 } else if reader.has_column(column_id) {
2698 if let columnar::NativeColumn::Int64 { data, validity } =
2699 reader.column_native_shared(column_id)?
2700 {
2701 for (i, n) in data.iter().enumerate() {
2702 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
2703 {
2704 max = *n;
2705 }
2706 }
2707 }
2708 }
2709 }
2710 Ok(max)
2711 }
2712
2713 fn seed_empty_auto_inc(&mut self) -> bool {
2714 let Some(ai) = self.auto_inc.as_mut() else {
2715 return false;
2716 };
2717 if ai.seeded || self.live_count != 0 {
2718 return false;
2719 }
2720 if ai.next < 1 {
2721 ai.next = 1;
2722 }
2723 ai.seeded = true;
2724 true
2725 }
2726
2727 fn advance_auto_inc_from_native_columns(
2728 &mut self,
2729 columns: &[(u16, columnar::NativeColumn)],
2730 n: usize,
2731 live_before: u64,
2732 ) -> Result<()> {
2733 let Some(ai) = self.auto_inc.as_mut() else {
2734 return Ok(());
2735 };
2736 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
2737 return Ok(());
2738 };
2739 let columnar::NativeColumn::Int64 { data, validity } = col else {
2740 return Err(MongrelError::InvalidArgument(format!(
2741 "AUTO_INCREMENT column {} must be Int64",
2742 ai.column_id
2743 )));
2744 };
2745 let max = if native_int64_strictly_increasing(col, n) {
2746 data.get(n.saturating_sub(1)).copied()
2747 } else {
2748 data.iter()
2749 .take(n)
2750 .enumerate()
2751 .filter_map(|(i, v)| {
2752 if validity.is_empty() || columnar::validity_bit(validity, i) {
2753 Some(*v)
2754 } else {
2755 None
2756 }
2757 })
2758 .max()
2759 };
2760 if let Some(max) = max {
2761 let floor = max
2762 .checked_add(1)
2763 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
2764 .max(1);
2765 if ai.next < floor {
2766 ai.next = floor;
2767 }
2768 if ai.seeded || live_before == 0 {
2769 ai.seeded = true;
2770 }
2771 }
2772 Ok(())
2773 }
2774
2775 fn fill_auto_inc_native_columns(
2776 &mut self,
2777 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2778 n: usize,
2779 ) -> Result<()> {
2780 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2781 return Ok(());
2782 };
2783 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2784 if let Some(start) = self.alloc_auto_inc_range(n)? {
2785 columns.push((
2786 cid,
2787 columnar::NativeColumn::Int64 {
2788 data: (start..start.saturating_add(n as i64)).collect(),
2789 validity: vec![0xFF; n.div_ceil(8)],
2790 },
2791 ));
2792 }
2793 return Ok(());
2794 };
2795
2796 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2797 return Err(MongrelError::InvalidArgument(format!(
2798 "AUTO_INCREMENT column {cid} must be Int64"
2799 )));
2800 };
2801 if data.len() < n {
2802 return Err(MongrelError::InvalidArgument(format!(
2803 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2804 data.len()
2805 )));
2806 }
2807 if columnar::all_non_null(validity, n) {
2808 return Ok(());
2809 }
2810 if validity.iter().all(|b| *b == 0) {
2811 if let Some(start) = self.alloc_auto_inc_range(n)? {
2812 for (i, slot) in data.iter_mut().take(n).enumerate() {
2813 *slot = start.saturating_add(i as i64);
2814 }
2815 *validity = vec![0xFF; n.div_ceil(8)];
2816 }
2817 return Ok(());
2818 }
2819
2820 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2821 for (i, slot) in data.iter_mut().enumerate().take(n) {
2822 if columnar::validity_bit(validity, i) {
2823 self.advance_auto_inc_past(*slot)?;
2824 } else {
2825 *slot = self.alloc_auto_inc_value()?;
2826 }
2827 }
2828 *validity = new_validity;
2829 Ok(())
2830 }
2831
2832 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2846 self.ensure_writable()?;
2847 if self.auto_inc.is_none() {
2848 return Ok(None);
2849 }
2850 Ok(Some(self.alloc_auto_inc_value()?))
2851 }
2852
2853 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2859 let payload = bincode::serialize(&rows)?;
2860 self.wal_append_data(Op::Put {
2861 table_id: self.table_id,
2862 rows: payload,
2863 })?;
2864 if self.is_shared() {
2865 self.pending_rows_auto_inc
2866 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
2867 self.pending_rows.extend(rows);
2868 } else {
2869 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2870 }
2871 Ok(())
2872 }
2873
2874 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
2877 self.ensure_indexes_complete()
2878 }
2879
2880 pub(crate) fn prepare_durable_publish_controlled(
2881 &mut self,
2882 control: &crate::ExecutionControl,
2883 ) -> Result<()> {
2884 self.ensure_indexes_complete_controlled(control, || true)?;
2885 Ok(())
2886 }
2887
2888 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
2889 self.apply_put_rows_inner_prepared(rows, true);
2890 }
2891
2892 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2893 if check_existing_pk {
2894 self.ensure_indexes_complete()?;
2895 }
2896 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
2897 Ok(())
2898 }
2899
2900 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
2904 if rows.len() == 1 {
2908 let row = rows.into_iter().next().expect("len checked");
2909 self.apply_put_row_single(row, check_existing_pk);
2910 return;
2911 }
2912 let pk_id = self.schema.primary_key().map(|c| c.id);
2929 let probe = match pk_id {
2930 Some(pid) => {
2931 check_existing_pk
2932 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2933 }
2934 None => false,
2935 };
2936 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2939 for r in rows {
2940 for &cid in r.columns.keys() {
2941 self.pending_put_cols.insert(cid);
2942 }
2943 match pk_id {
2944 Some(pid) if probe || maintain_pk_by_row => {
2945 if let Some(pk_val) = r.columns.get(&pid) {
2946 let key = self.index_lookup_key(pid, pk_val);
2947 if probe {
2948 if let Some(old_rid) = self.hot.get(&key) {
2949 if old_rid != r.row_id {
2950 self.tombstone_row(old_rid, r.committed_epoch, true);
2951 }
2952 }
2953 }
2954 if maintain_pk_by_row {
2955 self.pk_by_row.insert(r.row_id, key);
2956 }
2957 }
2958 }
2959 Some(_) => {}
2960 None => {
2961 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2962 }
2963 }
2964 self.index_row(&r);
2965 self.reservoir.offer(r.row_id.0);
2966 self.memtable.upsert(r);
2967 self.live_count = self.live_count.saturating_add(1);
2970 }
2971 self.data_generation = self.data_generation.wrapping_add(1);
2972 }
2973
2974 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
2978 for &cid in row.columns.keys() {
2979 self.pending_put_cols.insert(cid);
2980 }
2981 let epoch = row.committed_epoch;
2982 if let Some(pk_col) = self.schema.primary_key() {
2983 let pk_id = pk_col.id;
2984 if let Some(pk_val) = row.columns.get(&pk_id) {
2985 let maintain_pk_by_row = self.pk_by_row_complete;
2989 if check_existing_pk || maintain_pk_by_row {
2990 let key = self.index_lookup_key(pk_id, pk_val);
2991 if check_existing_pk {
2992 if let Some(old_rid) = self.hot.get(&key) {
2993 if old_rid != row.row_id {
2994 self.tombstone_row(old_rid, epoch, true);
2995 }
2996 }
2997 }
2998 if maintain_pk_by_row {
2999 self.pk_by_row.insert(row.row_id, key);
3000 }
3001 }
3002 }
3003 } else {
3004 self.hot
3005 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
3006 }
3007 self.index_row(&row);
3008 self.reservoir.offer(row.row_id.0);
3009 self.memtable.upsert(row);
3010 self.live_count = self.live_count.saturating_add(1);
3011 self.data_generation = self.data_generation.wrapping_add(1);
3012 }
3013
3014 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
3017 self.allocator.alloc()
3018 }
3019
3020 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
3026 let pk = self.schema.primary_key().ok_or_else(|| {
3027 MongrelError::Schema("clustered table requires a single-column primary key".into())
3028 })?;
3029 let pk_val = columns
3030 .iter()
3031 .find(|(id, _)| *id == pk.id)
3032 .map(|(_, v)| v)
3033 .ok_or_else(|| {
3034 MongrelError::Schema(format!(
3035 "clustered table missing primary key column {} ({})",
3036 pk.id, pk.name
3037 ))
3038 })?;
3039 let key_bytes = pk_val.encode_key();
3040 let mut hash: u64 = 0xcbf29ce484222325;
3042 for &b in &key_bytes {
3043 hash ^= b as u64;
3044 hash = hash.wrapping_mul(0x100000001b3);
3045 }
3046 Ok(RowId(hash.max(1)))
3049 }
3050
3051 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
3059 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
3060 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3061 }
3062 let n = rows.len();
3063 for r in rows {
3064 for &cid in r.columns.keys() {
3065 self.pending_put_cols.insert(cid);
3066 }
3067 }
3068 let (losers, winner_pks) = self.partition_pk_winners(rows);
3069 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
3070 for (key, &row_id) in &winner_pks {
3072 if let Some(old_rid) = self.hot.get(key) {
3073 if old_rid != row_id {
3074 self.tombstone_row(old_rid, epoch, true);
3075 }
3076 }
3077 }
3078 for &loser_rid in &losers {
3081 self.tombstone_row(loser_rid, epoch, false);
3082 }
3083 for (key, row_id) in winner_pks {
3085 self.insert_hot_pk(key, row_id);
3086 }
3087 if self.schema.primary_key().is_none() {
3088 for r in rows {
3089 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3090 }
3091 }
3092 for r in rows {
3093 self.allocator.advance_to(r.row_id)?;
3094 if !losers.contains(&r.row_id) {
3095 self.index_row(r);
3096 }
3097 }
3098 for r in rows {
3099 if !losers.contains(&r.row_id) {
3100 self.reservoir.offer(r.row_id.0);
3101 }
3102 }
3103 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
3104 self.data_generation = self.data_generation.wrapping_add(1);
3105 Ok(())
3106 }
3107
3108 pub(crate) fn recover_apply(
3113 &mut self,
3114 rows: Vec<Row>,
3115 deletes: Vec<(RowId, Epoch)>,
3116 ) -> Result<()> {
3117 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
3121 std::collections::BTreeMap::new();
3122 for row in rows {
3123 if row.row_id.0 >= u64::MAX - 1 {
3124 return Err(MongrelError::Full("row-id namespace exhausted".into()));
3125 }
3126 self.allocator.advance_to(row.row_id)?;
3127 if let Some(ai) = self.auto_inc.as_mut() {
3132 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3133 let next = n.checked_add(1).ok_or_else(|| {
3134 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3135 })?;
3136 if next > ai.next {
3137 ai.next = next;
3138 }
3139 }
3140 }
3141 by_epoch.entry(row.committed_epoch).or_default().push(row);
3142 }
3143 for (epoch, group) in by_epoch {
3144 let (losers, winner_pks) = self.partition_pk_winners(&group);
3145 for (key, &row_id) in &winner_pks {
3147 if let Some(old_rid) = self.hot.get(key) {
3148 if old_rid != row_id {
3149 self.tombstone_row(old_rid, epoch, false);
3150 }
3151 }
3152 }
3153 for (key, row_id) in winner_pks {
3154 self.insert_hot_pk(key, row_id);
3155 }
3156 if self.schema.primary_key().is_none() {
3157 for r in &group {
3158 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3159 }
3160 }
3161 for r in &group {
3162 if !losers.contains(&r.row_id) {
3163 self.memtable.upsert(r.clone());
3164 self.index_row(r);
3165 }
3166 }
3167 }
3168 for (rid, epoch) in deletes {
3169 self.memtable.tombstone(rid, epoch);
3170 self.remove_hot_for_row(rid, epoch);
3171 }
3172 self.reservoir_complete = false;
3175 Ok(())
3176 }
3177
3178 pub(crate) fn flushed_epoch(&self) -> u64 {
3180 self.flushed_epoch
3181 }
3182
3183 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
3184 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
3185 }
3186
3187 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
3189 self.schema.validate_values(cells)
3190 }
3191
3192 fn validate_columns_not_null(
3196 &self,
3197 columns: &[(u16, columnar::NativeColumn)],
3198 n: usize,
3199 ) -> Result<()> {
3200 let by_id: HashMap<u16, &columnar::NativeColumn> =
3201 columns.iter().map(|(id, c)| (*id, c)).collect();
3202 for col in &self.schema.columns {
3203 if !col.flags.contains(ColumnFlags::NULLABLE) {
3204 match by_id.get(&col.id) {
3205 None => {
3206 return Err(MongrelError::InvalidArgument(format!(
3207 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
3208 col.name, col.id
3209 )));
3210 }
3211 Some(c) => {
3212 if c.null_count(n) != 0 {
3213 return Err(MongrelError::InvalidArgument(format!(
3214 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
3215 col.name, col.id
3216 )));
3217 }
3218 }
3219 }
3220 }
3221 if let TypeId::Enum { variants } = &col.ty {
3222 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
3223 if by_id.contains_key(&col.id) {
3224 return Err(MongrelError::InvalidArgument(format!(
3225 "column '{}' ({}) enum requires a bytes column",
3226 col.name, col.id
3227 )));
3228 }
3229 continue;
3230 };
3231 for index in 0..n {
3232 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
3233 continue;
3234 };
3235 if !variants.iter().any(|variant| variant.as_bytes() == value) {
3236 return Err(MongrelError::InvalidArgument(format!(
3237 "column '{}' ({}) enum value {:?} is not one of {:?}",
3238 col.name,
3239 col.id,
3240 String::from_utf8_lossy(value),
3241 variants
3242 )));
3243 }
3244 }
3245 }
3246 }
3247 Ok(())
3248 }
3249
3250 fn bulk_pk_winner_indices(
3255 &self,
3256 columns: &[(u16, columnar::NativeColumn)],
3257 n: usize,
3258 ) -> Option<Vec<usize>> {
3259 let pk_col = self.schema.primary_key()?;
3260 let pk_id = pk_col.id;
3261 let pk_ty = pk_col.ty.clone();
3262 let by_id: HashMap<u16, &columnar::NativeColumn> =
3263 columns.iter().map(|(id, c)| (*id, c)).collect();
3264 let pk_native = by_id.get(&pk_id)?;
3265 if native_int64_strictly_increasing(pk_native, n) {
3266 return None;
3267 }
3268 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
3270 let mut null_pk_rows: Vec<usize> = Vec::new();
3271 for i in 0..n {
3272 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
3273 Some(key) => {
3274 last.insert(key, i);
3275 }
3276 None => null_pk_rows.push(i),
3277 }
3278 }
3279 let mut winners: HashSet<usize> = last.values().copied().collect();
3280 for i in null_pk_rows {
3281 winners.insert(i);
3282 }
3283 Some((0..n).filter(|i| winners.contains(i)).collect())
3284 }
3285
3286 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
3288 self.require_delete()?;
3289 let epoch = self.pending_epoch();
3290 self.wal_append_data(Op::Delete {
3291 table_id: self.table_id,
3292 row_ids: vec![row_id],
3293 })?;
3294 if self.is_shared() {
3295 self.pending_dels.push(row_id);
3296 } else {
3297 self.apply_delete(row_id, epoch);
3298 }
3299 Ok(())
3300 }
3301
3302 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
3303 let pre = self.get(row_id, self.snapshot());
3304 self.delete(row_id)?;
3305 Ok(pre.map(|row| {
3306 let mut columns: Vec<_> = row.columns.into_iter().collect();
3307 columns.sort_by_key(|(id, _)| *id);
3308 OwnedRow { columns }
3309 }))
3310 }
3311
3312 pub fn truncate(&mut self) -> Result<()> {
3314 self.require_delete()?;
3315 let epoch = self.pending_epoch();
3316 self.wal_append_data(Op::TruncateTable {
3317 table_id: self.table_id,
3318 })?;
3319 self.pending_rows.clear();
3320 self.pending_rows_auto_inc.clear();
3321 self.pending_dels.clear();
3322 self.pending_truncate = Some(epoch);
3323 Ok(())
3324 }
3325
3326 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
3328 self.run_refs.clear();
3333 self.retiring.clear();
3334 self.memtable = Memtable::new();
3335 self.mutable_run = MutableRun::new();
3336 self.hot = HotIndex::new();
3337 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3338 self.bitmap = bitmap;
3339 self.ann = ann;
3340 self.fm = fm;
3341 self.sparse = sparse;
3342 self.minhash = minhash;
3343 self.learned_range = Arc::new(HashMap::new());
3344 self.pk_by_row.clear();
3345 self.pk_by_row_complete = false;
3346 self.live_count = 0;
3347 self.reservoir = crate::reservoir::Reservoir::default();
3348 self.reservoir_complete = true;
3349 self.had_deletes = true;
3350 self.agg_cache = Arc::new(HashMap::new());
3351 self.global_idx_epoch = 0;
3352 self.indexes_complete = true;
3353 self.pending_delete_rids.clear();
3354 self.pending_put_cols.clear();
3355 self.pending_rows.clear();
3356 self.pending_rows_auto_inc.clear();
3357 self.pending_dels.clear();
3358 self.clear_result_cache();
3359 self.invalidate_index_checkpoint();
3360 self.data_generation = self.data_generation.wrapping_add(1);
3361 }
3362
3363 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
3366 self.remove_hot_for_row(row_id, epoch);
3367 self.tombstone_row(row_id, epoch, true);
3368 self.data_generation = self.data_generation.wrapping_add(1);
3369 }
3370
3371 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
3375 let tombstone = Row {
3376 row_id,
3377 committed_epoch: epoch,
3378 columns: std::collections::HashMap::new(),
3379 deleted: true,
3380 };
3381 self.memtable.upsert(tombstone);
3382 self.pk_by_row.remove(&row_id);
3383 if adjust_live_count {
3384 self.live_count = self.live_count.saturating_sub(1);
3385 }
3386 self.pending_delete_rids.insert(row_id.0 as u32);
3388 self.had_deletes = true;
3391 self.agg_cache = Arc::new(HashMap::new());
3392 }
3393
3394 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
3398 let Some(pk_col) = self.schema.primary_key() else {
3399 return;
3400 };
3401 if self.pk_by_row_complete {
3404 if let Some(key) = self.pk_by_row.remove(&row_id) {
3405 if self.hot.get(&key) == Some(row_id) {
3406 self.hot.remove(&key);
3407 }
3408 }
3409 return;
3410 }
3411 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
3430 if self.indexes_complete {
3431 let pk_val = self
3432 .memtable
3433 .get_version(row_id, lookup_epoch)
3434 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3435 .or_else(|| {
3436 self.mutable_run
3437 .get_version(row_id, lookup_epoch)
3438 .filter(|(_, r)| !r.deleted)
3439 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
3440 })
3441 .or_else(|| {
3442 self.run_refs.iter().find_map(|rr| {
3443 let mut reader = self.open_reader(rr.run_id).ok()?;
3444 let (_, deleted, val) = reader
3445 .get_version_column(row_id, lookup_epoch, pk_col.id)
3446 .ok()??;
3447 if deleted {
3448 return None;
3449 }
3450 val
3451 })
3452 });
3453 if let Some(pk_val) = pk_val {
3454 let key = self.index_lookup_key(pk_col.id, &pk_val);
3455 if self.hot.get(&key) == Some(row_id) {
3456 self.hot.remove(&key);
3457 }
3458 return;
3459 }
3460 }
3461 self.refresh_pk_by_row_from_hot();
3466 if let Some(key) = self.pk_by_row.remove(&row_id) {
3467 if self.hot.get(&key) == Some(row_id) {
3468 self.hot.remove(&key);
3469 }
3470 }
3471 }
3472
3473 fn partition_pk_winners(
3478 &self,
3479 rows: &[Row],
3480 ) -> (
3481 std::collections::HashSet<RowId>,
3482 std::collections::HashMap<Vec<u8>, RowId>,
3483 ) {
3484 let mut losers = std::collections::HashSet::new();
3485 let Some(pk_col) = self.schema.primary_key() else {
3486 return (losers, std::collections::HashMap::new());
3487 };
3488 let pk_id = pk_col.id;
3489 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
3490 std::collections::HashMap::new();
3491 for r in rows {
3492 let Some(pk_val) = r.columns.get(&pk_id) else {
3493 continue;
3494 };
3495 let key = self.index_lookup_key(pk_id, pk_val);
3496 if let Some(&old_rid) = winners.get(&key) {
3497 losers.insert(old_rid);
3498 }
3499 winners.insert(key, r.row_id);
3500 }
3501 (losers, winners)
3502 }
3503
3504 fn index_row(&mut self, row: &Row) {
3505 if row.deleted {
3506 return;
3507 }
3508 let any_predicate = self
3516 .schema
3517 .indexes
3518 .iter()
3519 .any(|idx| idx.predicate.is_some());
3520 if any_predicate {
3521 let columns_map: HashMap<u16, &Value> =
3522 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3523 let name_to_id: HashMap<&str, u16> = self
3524 .schema
3525 .columns
3526 .iter()
3527 .map(|c| (c.name.as_str(), c.id))
3528 .collect();
3529 for idx in &self.schema.indexes {
3530 if let Some(pred) = &idx.predicate {
3531 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3532 continue; }
3534 }
3535 index_into_single(
3537 idx,
3538 &self.schema,
3539 row,
3540 &mut self.hot,
3541 &mut self.bitmap,
3542 &mut self.ann,
3543 &mut self.fm,
3544 &mut self.sparse,
3545 &mut self.minhash,
3546 );
3547 }
3548 return;
3549 }
3550 if self.column_keys.is_empty() {
3554 index_into(
3555 &self.schema,
3556 row,
3557 &mut self.hot,
3558 &mut self.bitmap,
3559 &mut self.ann,
3560 &mut self.fm,
3561 &mut self.sparse,
3562 &mut self.minhash,
3563 );
3564 return;
3565 }
3566 let effective_row = self.tokenized_for_indexes(row);
3567 index_into(
3568 &self.schema,
3569 &effective_row,
3570 &mut self.hot,
3571 &mut self.bitmap,
3572 &mut self.ann,
3573 &mut self.fm,
3574 &mut self.sparse,
3575 &mut self.minhash,
3576 );
3577 }
3578
3579 fn tokenized_for_indexes(&self, row: &Row) -> Row {
3585 if self.column_keys.is_empty() {
3586 return row.clone();
3587 }
3588 #[cfg(feature = "encryption")]
3589 {
3590 use crate::encryption::SCHEME_HMAC_EQ;
3591 let mut tok = row.clone();
3592 for (&cid, &(_, scheme)) in &self.column_keys {
3593 if scheme != SCHEME_HMAC_EQ {
3594 continue;
3595 }
3596 if let Some(v) = tok.columns.get(&cid).cloned() {
3597 if let Some(t) = self.tokenize_value(cid, &v) {
3598 tok.columns.insert(cid, t);
3599 }
3600 }
3601 }
3602 tok
3603 }
3604 #[cfg(not(feature = "encryption"))]
3605 {
3606 row.clone()
3607 }
3608 }
3609
3610 pub fn commit(&mut self) -> Result<Epoch> {
3615 self.commit_inner(None)
3616 }
3617
3618 #[doc(hidden)]
3621 pub fn commit_controlled<F>(
3622 &mut self,
3623 control: &crate::ExecutionControl,
3624 mut before_commit: F,
3625 ) -> Result<Epoch>
3626 where
3627 F: FnMut() -> Result<()>,
3628 {
3629 self.commit_inner(Some((control, &mut before_commit)))
3630 }
3631
3632 fn commit_inner(
3633 &mut self,
3634 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3635 ) -> Result<Epoch> {
3636 self.ensure_writable()?;
3637 if !self.has_pending_mutations() {
3638 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
3639 return Err(MongrelError::Full(
3640 "standalone transaction id namespace exhausted".into(),
3641 ));
3642 }
3643 return Ok(self.epoch.visible());
3644 }
3645 self.commit_new_epoch_inner(controlled)
3646 }
3647
3648 fn commit_new_epoch(&mut self) -> Result<Epoch> {
3652 self.commit_new_epoch_inner(None)
3653 }
3654
3655 fn commit_new_epoch_inner(
3656 &mut self,
3657 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3658 ) -> Result<Epoch> {
3659 self.ensure_writable()?;
3660 if self.is_shared() {
3661 self.commit_shared(controlled)
3662 } else {
3663 self.commit_private(controlled)
3664 }
3665 }
3666
3667 fn commit_private(
3669 &mut self,
3670 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3671 ) -> Result<Epoch> {
3672 let commit_lock = Arc::clone(&self.commit_lock);
3676 let _g = commit_lock.lock();
3677 let txn_id = self.ensure_txn_id()?;
3680 if let Some((control, before_commit)) = controlled {
3681 control.checkpoint()?;
3682 before_commit()?;
3683 }
3684 let new_epoch = self.epoch.bump_assigned();
3685 let epoch_authority = Arc::clone(&self.epoch);
3686 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3687 let wal_result = match &mut self.wal {
3691 WalSink::Private(w) => w
3692 .append_txn(
3693 txn_id,
3694 Op::TxnCommit {
3695 epoch: new_epoch.0,
3696 added_runs: Vec::new(),
3697 },
3698 )
3699 .and_then(|_| w.sync()),
3700 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
3701 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
3702 };
3703 if let Err(error) = wal_result {
3704 self.durable_commit_failed = true;
3705 return Err(MongrelError::CommitOutcomeUnknown {
3706 epoch: new_epoch.0,
3707 message: error.to_string(),
3708 });
3709 }
3710 if let Some(epoch) = self.pending_truncate.take() {
3713 self.apply_truncate(epoch);
3714 }
3715 self.invalidate_pending_cache();
3716 let publish_result = self.persist_manifest(new_epoch);
3717 self.epoch.publish_in_order(new_epoch);
3721 epoch_guard.disarm();
3722 if let Err(error) = publish_result {
3723 self.durable_commit_failed = true;
3724 return Err(MongrelError::DurableCommit {
3725 epoch: new_epoch.0,
3726 message: error.to_string(),
3727 });
3728 }
3729 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
3730 self.pending_private_mutations = false;
3731 self.data_generation = self.data_generation.wrapping_add(1);
3732 Ok(new_epoch)
3733 }
3734
3735 fn commit_shared(
3741 &mut self,
3742 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3743 ) -> Result<Epoch> {
3744 use std::sync::atomic::Ordering;
3745 let s = match &self.wal {
3746 WalSink::Shared(s) => s.clone(),
3747 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
3748 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3749 };
3750 if s.poisoned.load(Ordering::Relaxed) {
3751 return Err(MongrelError::Other(
3752 "database poisoned by fsync error".into(),
3753 ));
3754 }
3755 let commit_lock = Arc::clone(&self.commit_lock);
3762 let _g = commit_lock.lock();
3763 if !self.pending_rows.is_empty() {
3764 match controlled.as_ref() {
3765 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
3766 None => self.prepare_durable_publish()?,
3767 }
3768 }
3769 let txn_id = self.ensure_txn_id()?;
3772 let mut wal = s.wal.lock();
3773 if let Some((control, before_commit)) = controlled {
3774 control.checkpoint()?;
3775 before_commit()?;
3776 }
3777 let new_epoch = self.epoch.bump_assigned();
3778 let epoch_authority = Arc::clone(&self.epoch);
3779 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
3780 let commit_seq = match wal.append_commit(txn_id, new_epoch, &[]) {
3781 Ok(commit_seq) => commit_seq,
3782 Err(error) => {
3783 s.poisoned.store(true, Ordering::Relaxed);
3784 return Err(MongrelError::CommitOutcomeUnknown {
3785 epoch: new_epoch.0,
3786 message: error.to_string(),
3787 });
3788 }
3789 };
3790 drop(wal);
3791 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
3792 s.poisoned.store(true, Ordering::Relaxed);
3793 return Err(MongrelError::CommitOutcomeUnknown {
3794 epoch: new_epoch.0,
3795 message: error.to_string(),
3796 });
3797 }
3798
3799 if self.pending_truncate.take().is_some() {
3802 self.apply_truncate(new_epoch);
3803 }
3804 let mut rows = std::mem::take(&mut self.pending_rows);
3805 if !rows.is_empty() {
3806 for r in &mut rows {
3807 r.committed_epoch = new_epoch;
3808 }
3809 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
3810 let all_auto_generated =
3811 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
3812 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
3813 } else {
3814 self.pending_rows_auto_inc.clear();
3815 }
3816 let dels = std::mem::take(&mut self.pending_dels);
3817 for rid in dels {
3818 self.apply_delete(rid, new_epoch);
3819 }
3820
3821 self.invalidate_pending_cache();
3822 let publish_result = self.persist_manifest(new_epoch);
3823 self.epoch.publish_in_order(new_epoch);
3824 epoch_guard.disarm();
3825 let _ = s.change_wake.send(());
3826 if let Err(error) = publish_result {
3827 self.durable_commit_failed = true;
3828 s.poisoned.store(true, Ordering::Relaxed);
3829 return Err(MongrelError::DurableCommit {
3830 epoch: new_epoch.0,
3831 message: error.to_string(),
3832 });
3833 }
3834 self.current_txn_id = 0;
3836 self.data_generation = self.data_generation.wrapping_add(1);
3837 Ok(new_epoch)
3838 }
3839
3840 pub fn flush(&mut self) -> Result<Epoch> {
3848 self.flush_with_outcome().map(|(epoch, _)| epoch)
3849 }
3850
3851 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
3853 self.flush_with_outcome_inner(None)
3854 }
3855
3856 #[doc(hidden)]
3859 pub fn flush_with_outcome_controlled<F>(
3860 &mut self,
3861 control: &crate::ExecutionControl,
3862 mut before_commit: F,
3863 ) -> Result<(Epoch, bool)>
3864 where
3865 F: FnMut() -> Result<()>,
3866 {
3867 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
3868 }
3869
3870 fn flush_with_outcome_inner(
3871 &mut self,
3872 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
3873 ) -> Result<(Epoch, bool)> {
3874 match controlled.as_ref() {
3875 Some((control, _)) => {
3876 self.ensure_indexes_complete_controlled(control, || true)?;
3877 }
3878 None => self.ensure_indexes_complete()?,
3879 }
3880 let committed = self.has_pending_mutations();
3881 let epoch = self.commit_inner(controlled)?;
3882 let finish: Result<(Epoch, bool)> = (|| {
3883 let rows = self.memtable.drain_sorted();
3884 if !rows.is_empty() {
3885 self.mutable_run.insert_many(rows);
3886 }
3887 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
3888 self.spill_mutable_run(epoch)?;
3889 self.mark_flushed(epoch)?;
3893 self.persist_manifest(epoch)?;
3894 self.build_learned_ranges()?;
3895 self.checkpoint_indexes(epoch);
3898 }
3899 Ok((epoch, committed))
3902 })();
3903 match finish {
3904 Err(error) if committed => Err(MongrelError::DurableCommit {
3905 epoch: epoch.0,
3906 message: error.to_string(),
3907 }),
3908 result => result,
3909 }
3910 }
3911
3912 fn has_pending_mutations(&self) -> bool {
3913 self.pending_private_mutations
3914 || !self.pending_rows.is_empty()
3915 || !self.pending_dels.is_empty()
3916 || self.pending_truncate.is_some()
3917 }
3918
3919 pub fn has_pending_writes(&self) -> bool {
3920 self.has_pending_mutations()
3921 }
3922
3923 pub fn force_flush(&mut self) -> Result<Epoch> {
3932 let saved = self.mutable_run_spill_bytes;
3933 self.mutable_run_spill_bytes = 1;
3934 let result = self.flush();
3935 self.mutable_run_spill_bytes = saved;
3936 result
3937 }
3938
3939 pub fn close(&mut self) -> Result<()> {
3946 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
3947 self.force_flush()?;
3948 }
3949 Ok(())
3950 }
3951
3952 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
3959 let op = Op::Flush {
3960 table_id: self.table_id,
3961 flushed_epoch: epoch.0,
3962 };
3963 match &mut self.wal {
3964 WalSink::Private(w) => {
3965 w.append_system(op)?;
3966 w.sync()?;
3967 }
3968 WalSink::Shared(s) => {
3969 s.wal.lock().append_system(op)?;
3974 }
3975 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3976 }
3977 self.flushed_epoch = epoch.0;
3978 if matches!(self.wal, WalSink::Private(_)) {
3979 self.rotate_wal(epoch)?;
3980 }
3981 Ok(())
3982 }
3983
3984 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
3988 if self.mutable_run.is_empty() {
3989 return Ok(());
3990 }
3991 let run_id = self.alloc_run_id()?;
3992 let rows = self.mutable_run.drain_sorted();
3993 if rows.is_empty() {
3994 return Ok(());
3995 }
3996 let path = self.run_path(run_id);
3997 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3998 if let Some(kek) = &self.kek {
3999 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4000 }
4001 let header = match self.create_run_file(run_id)? {
4002 Some(file) => writer.write_file(file, &rows)?,
4003 None => writer.write(&path, &rows)?,
4004 };
4005 self.run_refs.push(RunRef {
4006 run_id: run_id as u128,
4007 level: 0,
4008 epoch_created: epoch.0,
4009 row_count: header.row_count,
4010 });
4011 Ok(())
4012 }
4013
4014 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
4018 self.mutable_run_spill_bytes = bytes.max(1);
4019 }
4020
4021 pub fn set_compaction_zstd_level(&mut self, level: i32) {
4025 self.compaction_zstd_level = level;
4026 }
4027
4028 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
4032 self.result_cache.lock().set_max_bytes(max_bytes);
4033 }
4034
4035 pub(crate) fn clear_result_cache(&mut self) {
4039 self.result_cache.lock().clear();
4040 }
4041
4042 pub fn mutable_run_len(&self) -> usize {
4044 self.mutable_run.len()
4045 }
4046
4047 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
4050 self.mutable_run.drain_sorted()
4051 }
4052
4053 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
4055 let mut snapshot = self.mutable_run.clone();
4056 snapshot.drain_sorted()
4057 }
4058
4059 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
4064 self.ensure_writable()?;
4065 let n = batch.len();
4066 if n == 0 {
4067 return Ok(self.current_epoch());
4068 }
4069 for row in &batch {
4070 self.schema.validate_values(row)?;
4071 }
4072 let epoch = self.commit_new_epoch()?;
4073 let live_before = self.live_count;
4074 self.spill_mutable_run(epoch)?;
4078 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4079 && self.indexes_complete
4080 && self.run_refs.is_empty()
4081 && self.memtable.is_empty()
4082 && self.mutable_run.is_empty();
4083 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
4089 use rayon::prelude::*;
4090 self.schema
4091 .columns
4092 .par_iter()
4093 .map(|cdef| {
4094 (
4095 cdef.id,
4096 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
4097 )
4098 })
4099 .collect::<Vec<_>>()
4100 };
4101 drop(batch);
4102 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4107 self.validate_columns_not_null(&user_columns, n)?;
4108 let winner_idx = self
4109 .bulk_pk_winner_indices(&user_columns, n)
4110 .filter(|idx| idx.len() != n);
4111 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4112 match winner_idx.as_deref() {
4113 Some(idx) => {
4114 let compacted = user_columns
4115 .iter()
4116 .map(|(id, c)| (*id, c.gather(idx)))
4117 .collect();
4118 (compacted, idx.len())
4119 }
4120 None => (std::mem::take(&mut user_columns), n),
4121 };
4122 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4123 let first = self.allocator.alloc_range(write_n as u64)?.0;
4124 for rid in first..first + write_n as u64 {
4125 self.reservoir.offer(rid);
4126 }
4127 let run_id = self.alloc_run_id()?;
4128 let path = self.run_path(run_id);
4129 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
4130 .clean(true)
4131 .with_lz4()
4132 .with_native_endian();
4133 if let Some(kek) = &self.kek {
4134 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4135 }
4136 let header = match self.create_run_file(run_id)? {
4137 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
4138 None => writer.write_native(&path, &write_columns, write_n, first)?,
4139 };
4140 self.run_refs.push(RunRef {
4141 run_id: run_id as u128,
4142 level: 0,
4143 epoch_created: epoch.0,
4144 row_count: header.row_count,
4145 });
4146 self.live_count = self.live_count.saturating_add(write_n as u64);
4147 if eager_index_build {
4148 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4149 self.index_columns_bulk(&write_columns, &row_ids);
4150 self.indexes_complete = true;
4151 self.build_learned_ranges()?;
4152 } else {
4153 self.indexes_complete = false;
4154 }
4155 self.mark_flushed(epoch)?;
4156 self.persist_manifest(epoch)?;
4157 if eager_index_build {
4158 self.checkpoint_indexes(epoch);
4159 }
4160 self.clear_result_cache();
4161 Ok(epoch)
4162 }
4163
4164 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
4167 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
4168 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
4169 let segment_no = segment
4172 .file_stem()
4173 .and_then(|s| s.to_str())
4174 .and_then(|s| s.strip_prefix("seg-"))
4175 .and_then(|s| s.parse::<u64>().ok())
4176 .unwrap_or(0);
4177 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
4178 wal.set_sync_byte_threshold(self.sync_byte_threshold);
4179 wal.sync()?;
4180 self.wal = WalSink::Private(wal);
4181 Ok(())
4182 }
4183
4184 pub(crate) fn invalidate_pending_cache(&mut self) {
4189 self.result_cache
4190 .lock()
4191 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
4192 self.pending_delete_rids.clear();
4193 self.pending_put_cols.clear();
4194 }
4195
4196 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
4197 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
4198 m.current_epoch = epoch.0;
4199 m.next_row_id = self.allocator.current().0;
4200 m.runs = self.run_refs.clone();
4201 m.live_count = self.live_count;
4202 m.global_idx_epoch = self.global_idx_epoch;
4203 m.flushed_epoch = self.flushed_epoch;
4204 m.retiring = self.retiring.clone();
4205 m.auto_inc_next = match self.auto_inc {
4209 Some(ai) if ai.seeded => ai.next,
4210 _ => 0,
4211 };
4212 m.ttl = self.ttl;
4213 let meta_dek = self.manifest_meta_dek();
4214 match self._root_guard.as_deref() {
4215 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
4216 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
4217 }
4218 Ok(())
4219 }
4220
4221 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
4222 let rows = self.visible_rows_at_time(Snapshot::at(Epoch(u64::MAX)), i64::MIN)?;
4226 let live_count = u64::try_from(rows.len())
4227 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
4228 let auto_inc = match self.auto_inc {
4229 Some(mut state) => {
4230 let maximum = self.scan_max_int64(state.column_id)?;
4231 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
4232 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4233 })?;
4234 state.next = state.next.max(after_maximum).max(1);
4235 state.seeded = true;
4236 Some(state)
4237 }
4238 None => None,
4239 };
4240 Ok(RecoveryMetadataPlan {
4241 live_count,
4242 auto_inc,
4243 changed: live_count != self.live_count
4244 || auto_inc.is_some_and(|planned| {
4245 self.auto_inc.is_none_or(|current| {
4246 current.next != planned.next || current.seeded != planned.seeded
4247 })
4248 }),
4249 })
4250 }
4251
4252 pub(crate) fn apply_recovered_metadata(
4253 &mut self,
4254 plan: RecoveryMetadataPlan,
4255 epoch: Epoch,
4256 ) -> Result<()> {
4257 if !plan.changed {
4258 return Ok(());
4259 }
4260 self.live_count = plan.live_count;
4261 self.auto_inc = plan.auto_inc;
4262 self.persist_manifest(epoch)
4263 }
4264
4265 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
4271 if !self.indexes_complete {
4274 return;
4275 }
4276 if self.idx_root.is_none() {
4277 if let Some(root) = self._root_guard.as_ref() {
4278 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
4279 return;
4280 };
4281 self.idx_root = Some(Arc::new(idx_root));
4282 }
4283 }
4284 let snap = global_idx::IndexSnapshot {
4285 hot: &self.hot,
4286 bitmap: &self.bitmap,
4287 ann: &self.ann,
4288 fm: &self.fm,
4289 sparse: &self.sparse,
4290 minhash: &self.minhash,
4291 learned_range: &self.learned_range,
4292 };
4293 let idx_dek = self.idx_dek();
4295 let written = match self.idx_root.as_deref() {
4296 Some(root) => global_idx::write_atomic_root(
4297 root,
4298 self.table_id,
4299 epoch.0,
4300 snap,
4301 idx_dek.as_deref(),
4302 ),
4303 None => global_idx::write_atomic(
4304 &self.dir,
4305 self.table_id,
4306 epoch.0,
4307 snap,
4308 idx_dek.as_deref(),
4309 ),
4310 };
4311 if written.is_ok() {
4312 self.global_idx_epoch = epoch.0;
4313 let _ = self.persist_manifest(epoch);
4314 }
4315 }
4316
4317 pub(crate) fn invalidate_index_checkpoint(&mut self) {
4320 self.global_idx_epoch = 0;
4321 if let Some(root) = self.idx_root.as_deref() {
4322 let _ = root.remove_file(global_idx::IDX_FILENAME);
4323 } else {
4324 global_idx::remove(&self.dir);
4325 }
4326 let _ = self.persist_manifest(self.epoch.visible());
4327 }
4328
4329 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
4334 self.indexes_complete = false;
4335 self.global_idx_epoch = 0;
4336 if let Some(root) = self.idx_root.as_deref() {
4337 let _ = root.remove_file(global_idx::IDX_FILENAME);
4338 } else {
4339 global_idx::remove(&self.dir);
4340 }
4341 }
4342
4343 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
4344 self.indexes_complete = true;
4345 }
4346
4347 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
4353 self.durable_commit_failed = true;
4354 if let WalSink::Shared(shared) = &self.wal {
4355 shared
4356 .poisoned
4357 .store(true, std::sync::atomic::Ordering::Relaxed);
4358 }
4359 }
4360
4361 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
4365 self.durable_commit_failed = true;
4366 }
4367
4368 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
4371 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
4372 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
4373 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4374 best = Some((epoch, row));
4375 }
4376 }
4377 for rr in &self.run_refs {
4378 let Ok(mut reader) = self.open_reader(rr.run_id) else {
4379 continue;
4380 };
4381 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
4382 continue;
4383 };
4384 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
4385 best = Some((epoch, row));
4386 }
4387 }
4388 let now_nanos = unix_nanos_now();
4389 match best {
4390 Some((_, r)) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
4391 Some((_, r)) => Some(r),
4392 None => None,
4393 }
4394 }
4395
4396 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
4400 self.visible_rows_at_time(snapshot, unix_nanos_now())
4401 }
4402
4403 #[doc(hidden)]
4406 pub fn visible_rows_controlled(
4407 &self,
4408 snapshot: Snapshot,
4409 control: &crate::ExecutionControl,
4410 ) -> Result<Vec<Row>> {
4411 let mut out = Vec::new();
4412 self.for_each_visible_row_controlled(snapshot, control, |row| {
4413 out.push(row);
4414 Ok(())
4415 })?;
4416 Ok(out)
4417 }
4418
4419 #[doc(hidden)]
4422 pub fn for_each_visible_row_controlled<F>(
4423 &self,
4424 snapshot: Snapshot,
4425 control: &crate::ExecutionControl,
4426 visit: F,
4427 ) -> Result<()>
4428 where
4429 F: FnMut(Row) -> Result<()>,
4430 {
4431 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
4432 control.checkpoint()?;
4433 let memtable = self.memtable.visible_versions(snapshot.epoch);
4434 if !memtable.is_empty() {
4435 sources.push(ControlledVisibleSource::memory(memtable));
4436 }
4437 control.checkpoint()?;
4438 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4439 if !mutable.is_empty() {
4440 sources.push(ControlledVisibleSource::memory(mutable));
4441 }
4442 for run in &self.run_refs {
4443 control.checkpoint()?;
4444 let reader = self.open_reader(run.run_id)?;
4445 sources.push(ControlledVisibleSource::run(
4446 reader.into_visible_version_cursor(snapshot.epoch)?,
4447 ));
4448 }
4449 let now_nanos = unix_nanos_now();
4450 merge_controlled_visible_sources(
4451 &mut sources,
4452 control,
4453 |row| self.row_expired_at(row, now_nanos),
4454 visit,
4455 )
4456 }
4457
4458 #[doc(hidden)]
4459 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
4460 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4461 let mut fold = |row: Row| {
4462 best.entry(row.row_id.0)
4463 .and_modify(|e| {
4464 if row.committed_epoch > e.0 {
4465 *e = (row.committed_epoch, row.clone());
4466 }
4467 })
4468 .or_insert_with(|| (row.committed_epoch, row));
4469 };
4470 for row in self.memtable.visible_versions(snapshot.epoch) {
4471 fold(row);
4472 }
4473 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4474 fold(row);
4475 }
4476 for rr in &self.run_refs {
4477 let mut reader = self.open_reader(rr.run_id)?;
4478 for row in reader.visible_versions(snapshot.epoch)? {
4479 fold(row);
4480 }
4481 }
4482 let mut out: Vec<Row> = best
4483 .into_values()
4484 .filter_map(|(_, r)| {
4485 if r.deleted || self.row_expired_at(&r, now_nanos) {
4486 None
4487 } else {
4488 Some(r)
4489 }
4490 })
4491 .collect();
4492 out.sort_by_key(|r| r.row_id);
4493 Ok(out)
4494 }
4495
4496 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
4503 if self.ttl.is_none()
4504 && self.memtable.is_empty()
4505 && self.mutable_run.is_empty()
4506 && self.run_refs.len() == 1
4507 {
4508 let rr = self.run_refs[0].clone();
4509 let mut reader = self.open_reader(rr.run_id)?;
4510 let idxs = reader.visible_indices(snapshot.epoch)?;
4511 let mut cols = Vec::with_capacity(self.schema.columns.len());
4512 for cdef in &self.schema.columns {
4513 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
4514 }
4515 return Ok(cols);
4516 }
4517 let rows = self.visible_rows(snapshot)?;
4519 let mut cols: Vec<(u16, Vec<Value>)> = self
4520 .schema
4521 .columns
4522 .iter()
4523 .map(|c| (c.id, Vec::with_capacity(rows.len())))
4524 .collect();
4525 for r in &rows {
4526 for (cid, vec) in cols.iter_mut() {
4527 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
4528 }
4529 }
4530 Ok(cols)
4531 }
4532
4533 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
4535 let row_id = self.hot.get(key)?;
4536 if self.ttl.is_none() || self.get(row_id, Snapshot::at(Epoch(u64::MAX))).is_some() {
4537 Some(row_id)
4538 } else {
4539 None
4540 }
4541 }
4542
4543 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4548 self.query_at_with_allowed(q, self.snapshot(), None)
4549 }
4550
4551 pub fn query_at_with_allowed(
4554 &mut self,
4555 q: &crate::query::Query,
4556 snapshot: Snapshot,
4557 allowed: Option<&std::collections::HashSet<RowId>>,
4558 ) -> Result<Vec<Row>> {
4559 self.query_at_with_allowed_after(q, snapshot, allowed, None)
4560 }
4561
4562 #[doc(hidden)]
4563 pub fn query_at_with_allowed_after(
4564 &mut self,
4565 q: &crate::query::Query,
4566 snapshot: Snapshot,
4567 allowed: Option<&std::collections::HashSet<RowId>>,
4568 after_row_id: Option<RowId>,
4569 ) -> Result<Vec<Row>> {
4570 self.query_at_with_allowed_after_at_time(
4571 q,
4572 snapshot,
4573 allowed,
4574 after_row_id,
4575 unix_nanos_now(),
4576 )
4577 }
4578
4579 #[doc(hidden)]
4580 pub fn query_at_with_allowed_after_at_time(
4581 &mut self,
4582 q: &crate::query::Query,
4583 snapshot: Snapshot,
4584 allowed: Option<&std::collections::HashSet<RowId>>,
4585 after_row_id: Option<RowId>,
4586 query_time_nanos: i64,
4587 ) -> Result<Vec<Row>> {
4588 self.require_select()?;
4589 self.ensure_indexes_complete()?;
4590 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4591 return Err(MongrelError::InvalidArgument(format!(
4592 "query exceeds {} conditions",
4593 crate::query::MAX_HARD_CONDITIONS
4594 )));
4595 }
4596 if let Some(limit) = q.limit {
4597 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
4598 return Err(MongrelError::InvalidArgument(format!(
4599 "query limit must be between 1 and {}",
4600 crate::query::MAX_FINAL_LIMIT
4601 )));
4602 }
4603 }
4604 if q.offset > crate::query::MAX_QUERY_OFFSET {
4605 return Err(MongrelError::InvalidArgument(format!(
4606 "query offset exceeds {}",
4607 crate::query::MAX_QUERY_OFFSET
4608 )));
4609 }
4610 self.query_conditions_at(
4611 &q.conditions,
4612 snapshot,
4613 allowed,
4614 q.limit,
4615 q.offset,
4616 after_row_id,
4617 query_time_nanos,
4618 )
4619 }
4620
4621 #[doc(hidden)]
4624 pub fn query_all_at(
4625 &mut self,
4626 conditions: &[crate::query::Condition],
4627 snapshot: Snapshot,
4628 ) -> Result<Vec<Row>> {
4629 self.require_select()?;
4630 self.ensure_indexes_complete()?;
4631 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
4632 return Err(MongrelError::InvalidArgument(format!(
4633 "query exceeds {} conditions",
4634 crate::query::MAX_HARD_CONDITIONS
4635 )));
4636 }
4637 self.query_conditions_at(conditions, snapshot, None, None, 0, None, unix_nanos_now())
4638 }
4639
4640 #[allow(clippy::too_many_arguments)]
4641 fn query_conditions_at(
4642 &self,
4643 conditions: &[crate::query::Condition],
4644 snapshot: Snapshot,
4645 allowed: Option<&std::collections::HashSet<RowId>>,
4646 limit: Option<usize>,
4647 offset: usize,
4648 after_row_id: Option<RowId>,
4649 query_time_nanos: i64,
4650 ) -> Result<Vec<Row>> {
4651 crate::trace::QueryTrace::record(|t| {
4652 t.run_count = self.run_refs.len();
4653 t.memtable_rows = self.memtable.len();
4654 t.mutable_run_rows = self.mutable_run.len();
4655 });
4656 if conditions.is_empty() {
4660 crate::trace::QueryTrace::record(|t| {
4661 t.scan_mode = crate::trace::ScanMode::Materialized;
4662 t.row_materialized = true;
4663 });
4664 let mut rows = self.visible_rows_at_time(snapshot, query_time_nanos)?;
4665 if let Some(allowed) = allowed {
4666 rows.retain(|row| allowed.contains(&row.row_id));
4667 }
4668 if let Some(after_row_id) = after_row_id {
4669 rows.retain(|row| row.row_id > after_row_id);
4670 }
4671 rows.drain(..offset.min(rows.len()));
4672 if let Some(limit) = limit {
4673 rows.truncate(limit);
4674 }
4675 return Ok(rows);
4676 }
4677 crate::trace::QueryTrace::record(|t| {
4678 t.conditions_pushed = conditions.len();
4679 t.scan_mode = crate::trace::ScanMode::Materialized;
4680 t.row_materialized = true;
4681 });
4682 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
4689 ordered.sort_by_key(|c| condition_cost_rank(c));
4690 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
4691 for c in &ordered {
4692 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
4693 let empty = s.is_empty();
4694 sets.push(s);
4695 if empty {
4696 break;
4697 }
4698 }
4699 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4700 if let Some(allowed) = allowed {
4701 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
4702 }
4703 if let Some(after_row_id) = after_row_id {
4704 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
4705 rids.drain(..first);
4706 }
4707 rids.drain(..offset.min(rids.len()));
4708 if let Some(limit) = limit {
4709 rids.truncate(limit);
4710 }
4711 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos)
4712 }
4713
4714 pub fn retrieve(
4716 &mut self,
4717 retriever: &crate::query::Retriever,
4718 ) -> Result<Vec<crate::query::RetrieverHit>> {
4719 self.retrieve_with_allowed(retriever, None)
4720 }
4721
4722 pub fn retrieve_at(
4723 &mut self,
4724 retriever: &crate::query::Retriever,
4725 snapshot: Snapshot,
4726 allowed: Option<&std::collections::HashSet<RowId>>,
4727 ) -> Result<Vec<crate::query::RetrieverHit>> {
4728 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
4729 }
4730
4731 pub fn retrieve_with_allowed(
4734 &mut self,
4735 retriever: &crate::query::Retriever,
4736 allowed: Option<&std::collections::HashSet<RowId>>,
4737 ) -> Result<Vec<crate::query::RetrieverHit>> {
4738 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
4739 }
4740
4741 pub fn retrieve_at_with_allowed(
4742 &mut self,
4743 retriever: &crate::query::Retriever,
4744 snapshot: Snapshot,
4745 allowed: Option<&std::collections::HashSet<RowId>>,
4746 ) -> Result<Vec<crate::query::RetrieverHit>> {
4747 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
4748 }
4749
4750 pub fn retrieve_at_with_allowed_and_context(
4751 &mut self,
4752 retriever: &crate::query::Retriever,
4753 snapshot: Snapshot,
4754 allowed: Option<&std::collections::HashSet<RowId>>,
4755 context: Option<&crate::query::AiExecutionContext>,
4756 ) -> Result<Vec<crate::query::RetrieverHit>> {
4757 self.require_select()?;
4758 self.ensure_indexes_complete()?;
4759 self.validate_retriever(retriever)?;
4760 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
4761 }
4762
4763 pub fn retrieve_at_with_candidate_authorization_and_context(
4764 &mut self,
4765 retriever: &crate::query::Retriever,
4766 snapshot: Snapshot,
4767 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4768 context: Option<&crate::query::AiExecutionContext>,
4769 ) -> Result<Vec<crate::query::RetrieverHit>> {
4770 self.require_select()?;
4771 self.ensure_indexes_complete()?;
4772 self.retrieve_at_with_candidate_authorization_on_generation(
4773 retriever,
4774 snapshot,
4775 authorization,
4776 context,
4777 )
4778 }
4779
4780 #[doc(hidden)]
4781 pub fn retrieve_at_with_candidate_authorization_on_generation(
4782 &self,
4783 retriever: &crate::query::Retriever,
4784 snapshot: Snapshot,
4785 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4786 context: Option<&crate::query::AiExecutionContext>,
4787 ) -> Result<Vec<crate::query::RetrieverHit>> {
4788 self.require_select()?;
4789 self.validate_retriever(retriever)?;
4790 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
4791 }
4792
4793 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
4794 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
4795 let (column_id, k) = match retriever {
4796 Retriever::Ann {
4797 column_id,
4798 query,
4799 k,
4800 } => {
4801 let index = self.ann.get(column_id).ok_or_else(|| {
4802 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
4803 })?;
4804 if query.len() != index.dim() {
4805 return Err(MongrelError::InvalidArgument(format!(
4806 "ANN query dimension must be {}, got {}",
4807 index.dim(),
4808 query.len()
4809 )));
4810 }
4811 if query.iter().any(|value| !value.is_finite()) {
4812 return Err(MongrelError::InvalidArgument(
4813 "ANN query values must be finite".into(),
4814 ));
4815 }
4816 (*column_id, *k)
4817 }
4818 Retriever::Sparse {
4819 column_id,
4820 query,
4821 k,
4822 } => {
4823 if !self.sparse.contains_key(column_id) {
4824 return Err(MongrelError::InvalidArgument(format!(
4825 "column {column_id} has no Sparse index"
4826 )));
4827 }
4828 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
4829 return Err(MongrelError::InvalidArgument(
4830 "Sparse query must be non-empty with finite weights".into(),
4831 ));
4832 }
4833 if query.len() > MAX_SPARSE_TERMS {
4834 return Err(MongrelError::InvalidArgument(format!(
4835 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
4836 )));
4837 }
4838 (*column_id, *k)
4839 }
4840 Retriever::MinHash {
4841 column_id,
4842 members,
4843 k,
4844 } => {
4845 if !self.minhash.contains_key(column_id) {
4846 return Err(MongrelError::InvalidArgument(format!(
4847 "column {column_id} has no MinHash index"
4848 )));
4849 }
4850 if members.is_empty() {
4851 return Err(MongrelError::InvalidArgument(
4852 "MinHash members must not be empty".into(),
4853 ));
4854 }
4855 if members.len() > MAX_SET_MEMBERS {
4856 return Err(MongrelError::InvalidArgument(format!(
4857 "MinHash query exceeds {MAX_SET_MEMBERS} members"
4858 )));
4859 }
4860 let mut total_bytes = 0usize;
4861 for member in members {
4862 let bytes = member.encoded_len();
4863 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
4864 return Err(MongrelError::InvalidArgument(format!(
4865 "MinHash member exceeds {} bytes",
4866 crate::query::MAX_SET_MEMBER_BYTES
4867 )));
4868 }
4869 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
4870 MongrelError::InvalidArgument("MinHash input size overflow".into())
4871 })?;
4872 }
4873 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
4874 return Err(MongrelError::InvalidArgument(format!(
4875 "MinHash input exceeds {} bytes",
4876 crate::query::MAX_SET_INPUT_BYTES
4877 )));
4878 }
4879 (*column_id, *k)
4880 }
4881 };
4882 if k == 0 {
4883 return Err(MongrelError::InvalidArgument(
4884 "retriever k must be > 0".into(),
4885 ));
4886 }
4887 if k > MAX_RETRIEVER_K {
4888 return Err(MongrelError::InvalidArgument(format!(
4889 "retriever k exceeds {MAX_RETRIEVER_K}"
4890 )));
4891 }
4892 debug_assert!(self
4893 .schema
4894 .columns
4895 .iter()
4896 .any(|column| column.id == column_id));
4897 Ok(())
4898 }
4899
4900 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
4901 use crate::query::Condition;
4902 match condition {
4903 Condition::Ann {
4904 column_id,
4905 query,
4906 k,
4907 } => self.validate_retriever(&crate::query::Retriever::Ann {
4908 column_id: *column_id,
4909 query: query.clone(),
4910 k: *k,
4911 }),
4912 Condition::SparseMatch {
4913 column_id,
4914 query,
4915 k,
4916 } => self.validate_retriever(&crate::query::Retriever::Sparse {
4917 column_id: *column_id,
4918 query: query.clone(),
4919 k: *k,
4920 }),
4921 Condition::MinHashSimilar {
4922 column_id,
4923 query,
4924 k,
4925 } => {
4926 if !self.minhash.contains_key(column_id) {
4927 return Err(MongrelError::InvalidArgument(format!(
4928 "column {column_id} has no MinHash index"
4929 )));
4930 }
4931 if query.is_empty() || *k == 0 {
4932 return Err(MongrelError::InvalidArgument(
4933 "MinHash query must be non-empty and k must be > 0".into(),
4934 ));
4935 }
4936 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
4937 {
4938 return Err(MongrelError::InvalidArgument(format!(
4939 "MinHash query must have <= {} members and k <= {}",
4940 crate::query::MAX_SET_MEMBERS,
4941 crate::query::MAX_RETRIEVER_K
4942 )));
4943 }
4944 Ok(())
4945 }
4946 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
4947 Err(MongrelError::InvalidArgument(format!(
4948 "bitmap IN exceeds {} values",
4949 crate::query::MAX_SET_MEMBERS
4950 )))
4951 }
4952 Condition::FmContainsAll { patterns, .. }
4953 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
4954 {
4955 Err(MongrelError::InvalidArgument(format!(
4956 "FM query exceeds {} patterns",
4957 crate::query::MAX_HARD_CONDITIONS
4958 )))
4959 }
4960 _ => Ok(()),
4961 }
4962 }
4963
4964 fn retrieve_filtered(
4965 &self,
4966 retriever: &crate::query::Retriever,
4967 snapshot: Snapshot,
4968 hard_filter: Option<&RowIdSet>,
4969 allowed: Option<&std::collections::HashSet<RowId>>,
4970 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
4971 context: Option<&crate::query::AiExecutionContext>,
4972 ) -> Result<Vec<crate::query::RetrieverHit>> {
4973 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
4974 let started = std::time::Instant::now();
4975 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
4976 Retriever::Ann {
4977 column_id,
4978 query,
4979 k,
4980 } => {
4981 let Some(index) = self.ann.get(column_id) else {
4982 return Ok(Vec::new());
4983 };
4984 let cap = ann_candidate_cap(index.len(), context);
4985 if cap == 0 {
4986 return Ok(Vec::new());
4987 }
4988 let mut breadth = (*k).max(1).min(cap);
4989 let mut eligibility = std::collections::HashMap::new();
4990 let mut filtered = loop {
4991 let mut seen = std::collections::HashSet::new();
4992 if let Some(context) = context {
4993 context.checkpoint()?;
4994 }
4995 let raw = index.search_with_context(query, breadth, context)?;
4996 let unchecked: Vec<_> = raw
4997 .iter()
4998 .map(|(row_id, _)| *row_id)
4999 .filter(|row_id| !eligibility.contains_key(row_id))
5000 .filter(|row_id| {
5001 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5002 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5003 })
5004 .collect();
5005 let eligible = self.eligible_and_authorized_candidate_ids(
5006 &unchecked,
5007 *column_id,
5008 snapshot,
5009 candidate_authorization,
5010 context,
5011 )?;
5012 for row_id in unchecked {
5013 eligibility.insert(row_id, eligible.contains(&row_id));
5014 }
5015 let filtered: Vec<_> = raw
5016 .into_iter()
5017 .filter(|(row_id, _)| {
5018 seen.insert(*row_id)
5019 && eligibility.get(row_id).copied().unwrap_or(false)
5020 })
5021 .map(|(row_id, score)| (row_id, RetrieverScore::AnnHammingDistance(score)))
5022 .collect();
5023 if filtered.len() >= *k || breadth >= cap {
5024 if filtered.len() < *k && index.len() > cap && breadth >= cap {
5025 crate::trace::QueryTrace::record(|trace| {
5026 trace.ann_candidate_cap_hit = true;
5027 });
5028 }
5029 break filtered;
5030 }
5031 breadth = breadth.saturating_mul(2).min(cap);
5032 };
5033 filtered.truncate(*k);
5034 filtered
5035 }
5036 Retriever::Sparse {
5037 column_id,
5038 query,
5039 k,
5040 } => self
5041 .sparse
5042 .get(column_id)
5043 .map(|index| -> Result<Vec<_>> {
5044 let mut breadth = (*k).max(1);
5045 let mut eligibility = std::collections::HashMap::new();
5046 loop {
5047 if let Some(context) = context {
5048 context.checkpoint()?;
5049 }
5050 let raw = index.search_with_context(query, breadth, context)?;
5051 let unchecked: Vec<_> = raw
5052 .iter()
5053 .map(|(row_id, _)| *row_id)
5054 .filter(|row_id| !eligibility.contains_key(row_id))
5055 .filter(|row_id| {
5056 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5057 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5058 })
5059 .collect();
5060 let eligible = self.eligible_and_authorized_candidate_ids(
5061 &unchecked,
5062 *column_id,
5063 snapshot,
5064 candidate_authorization,
5065 context,
5066 )?;
5067 for row_id in unchecked {
5068 eligibility.insert(row_id, eligible.contains(&row_id));
5069 }
5070 let filtered: Vec<_> = raw
5071 .iter()
5072 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5073 .take(*k)
5074 .map(|(row_id, score)| {
5075 (*row_id, RetrieverScore::SparseDotProduct(*score))
5076 })
5077 .collect();
5078 if filtered.len() >= *k || raw.len() < breadth {
5079 break Ok(filtered);
5080 }
5081 let next = breadth.saturating_mul(2);
5082 if next == breadth {
5083 break Ok(filtered);
5084 }
5085 breadth = next;
5086 }
5087 })
5088 .transpose()?
5089 .unwrap_or_default(),
5090 Retriever::MinHash {
5091 column_id,
5092 members,
5093 k,
5094 } => self
5095 .minhash
5096 .get(column_id)
5097 .map(|index| -> Result<Vec<_>> {
5098 let mut hashes = Vec::with_capacity(members.len());
5099 for member in members {
5100 if let Some(context) = context {
5101 context.consume(crate::query::work_units(
5102 member.encoded_len(),
5103 crate::query::PARSE_WORK_QUANTUM,
5104 ))?;
5105 }
5106 hashes.push(member.hash_v1());
5107 }
5108 let mut breadth = (*k).max(1);
5109 let mut eligibility = std::collections::HashMap::new();
5110 loop {
5111 if let Some(context) = context {
5112 context.checkpoint()?;
5113 }
5114 let raw = index.search_with_context(&hashes, breadth, context)?;
5115 let unchecked: Vec<_> = raw
5116 .iter()
5117 .map(|(row_id, _)| *row_id)
5118 .filter(|row_id| !eligibility.contains_key(row_id))
5119 .filter(|row_id| {
5120 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
5121 && allowed.is_none_or(|allowed| allowed.contains(row_id))
5122 })
5123 .collect();
5124 let eligible = self.eligible_and_authorized_candidate_ids(
5125 &unchecked,
5126 *column_id,
5127 snapshot,
5128 candidate_authorization,
5129 context,
5130 )?;
5131 for row_id in unchecked {
5132 eligibility.insert(row_id, eligible.contains(&row_id));
5133 }
5134 let filtered: Vec<_> = raw
5135 .iter()
5136 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
5137 .take(*k)
5138 .map(|(row_id, score)| {
5139 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
5140 })
5141 .collect();
5142 if filtered.len() >= *k || raw.len() < breadth {
5143 break Ok(filtered);
5144 }
5145 let next = breadth.saturating_mul(2);
5146 if next == breadth {
5147 break Ok(filtered);
5148 }
5149 breadth = next;
5150 }
5151 })
5152 .transpose()?
5153 .unwrap_or_default(),
5154 };
5155 let elapsed = started.elapsed().as_nanos() as u64;
5156 crate::trace::QueryTrace::record(|trace| {
5157 match retriever {
5158 Retriever::Ann { .. } => {
5159 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed)
5160 }
5161 Retriever::Sparse { .. } => {
5162 trace.sparse_candidate_nanos =
5163 trace.sparse_candidate_nanos.saturating_add(elapsed)
5164 }
5165 Retriever::MinHash { .. } => {
5166 trace.minhash_candidate_nanos =
5167 trace.minhash_candidate_nanos.saturating_add(elapsed)
5168 }
5169 }
5170 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
5171 });
5172 Ok(scored
5173 .into_iter()
5174 .enumerate()
5175 .map(|(rank, (row_id, score))| RetrieverHit {
5176 row_id,
5177 rank: rank + 1,
5178 score,
5179 })
5180 .collect())
5181 }
5182
5183 fn eligible_candidate_ids(
5184 &self,
5185 candidates: &[RowId],
5186 _column_id: u16,
5187 snapshot: Snapshot,
5188 context: Option<&crate::query::AiExecutionContext>,
5189 ) -> Result<std::collections::HashSet<RowId>> {
5190 if !self.had_deletes
5191 && self.ttl.is_none()
5192 && self.pending_put_cols.is_empty()
5193 && snapshot.epoch == self.snapshot().epoch
5194 {
5195 return Ok(candidates.iter().copied().collect());
5196 }
5197 let mut readers: Vec<_> = self
5198 .run_refs
5199 .iter()
5200 .map(|run| self.open_reader(run.run_id))
5201 .collect::<Result<_>>()?;
5202 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5203 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
5204 for &row_id in candidates {
5205 if let Some(context) = context {
5206 context.consume(1)?;
5207 }
5208 let mem = self.memtable.get_version(row_id, snapshot.epoch);
5209 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
5210 let overlay = match (mem, mutable) {
5211 (Some(left), Some(right)) => Some(if left.0 >= right.0 { left } else { right }),
5212 (Some(value), None) | (None, Some(value)) => Some(value),
5213 (None, None) => None,
5214 };
5215 if let Some((_, row)) = overlay {
5216 if !row.deleted && !self.row_expired_at(&row, now) {
5217 eligible.insert(row_id);
5218 }
5219 continue;
5220 }
5221 let mut best: Option<(Epoch, bool, usize)> = None;
5222 for (index, reader) in readers.iter_mut().enumerate() {
5223 if let Some((epoch, deleted)) =
5224 reader.get_version_visibility(row_id, snapshot.epoch)?
5225 {
5226 if best
5227 .as_ref()
5228 .map(|(best_epoch, ..)| epoch > *best_epoch)
5229 .unwrap_or(true)
5230 {
5231 best = Some((epoch, deleted, index));
5232 }
5233 }
5234 }
5235 let Some((_, false, reader_index)) = best else {
5236 continue;
5237 };
5238 if let Some(ttl) = self.ttl {
5239 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
5240 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
5241 {
5242 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
5243 continue;
5244 }
5245 }
5246 }
5247 eligible.insert(row_id);
5248 }
5249 Ok(eligible)
5250 }
5251
5252 fn eligible_and_authorized_candidate_ids(
5253 &self,
5254 candidates: &[RowId],
5255 column_id: u16,
5256 snapshot: Snapshot,
5257 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5258 context: Option<&crate::query::AiExecutionContext>,
5259 ) -> Result<std::collections::HashSet<RowId>> {
5260 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
5261 let Some(authorization) = authorization else {
5262 return Ok(eligible);
5263 };
5264 let candidates: Vec<_> = eligible.into_iter().collect();
5265 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
5266 }
5267
5268 fn policy_allowed_candidate_ids(
5269 &self,
5270 candidates: &[RowId],
5271 snapshot: Snapshot,
5272 authorization: &crate::security::CandidateAuthorization<'_>,
5273 context: Option<&crate::query::AiExecutionContext>,
5274 ) -> Result<std::collections::HashSet<RowId>> {
5275 let started = std::time::Instant::now();
5276 if candidates.is_empty()
5277 || authorization.principal.is_admin
5278 || !authorization.security.rls_enabled(authorization.table)
5279 {
5280 return Ok(candidates.iter().copied().collect());
5281 }
5282 if let Some(context) = context {
5283 context.checkpoint()?;
5284 }
5285 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
5286 let mut rows: std::collections::HashMap<RowId, Row> = candidates
5287 .iter()
5288 .map(|row_id| {
5289 (
5290 *row_id,
5291 Row {
5292 row_id: *row_id,
5293 committed_epoch: snapshot.epoch,
5294 columns: std::collections::HashMap::new(),
5295 deleted: false,
5296 },
5297 )
5298 })
5299 .collect();
5300 let columns = authorization
5301 .security
5302 .select_policy_columns(authorization.table, authorization.principal);
5303 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5304 let mut decoded = 0usize;
5305 for column_id in &columns {
5306 if let Some(context) = context {
5307 context.checkpoint()?;
5308 }
5309 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5310 &row_ids, *column_id, snapshot, query_now, context,
5311 )? {
5312 if let Some(row) = rows.get_mut(&row_id) {
5313 row.columns.insert(*column_id, value);
5314 decoded = decoded.saturating_add(1);
5315 }
5316 }
5317 }
5318 if let Some(context) = context {
5319 context.consume(candidates.len().saturating_add(decoded))?;
5320 }
5321 let allowed = rows
5322 .into_values()
5323 .filter_map(|row| {
5324 authorization
5325 .security
5326 .row_allowed(
5327 authorization.table,
5328 crate::security::PolicyCommand::Select,
5329 &row,
5330 authorization.principal,
5331 false,
5332 )
5333 .then_some(row.row_id)
5334 })
5335 .collect();
5336 crate::trace::QueryTrace::record(|trace| {
5337 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
5338 trace.rls_policy_columns_decoded =
5339 trace.rls_policy_columns_decoded.saturating_add(decoded);
5340 trace.authorization_nanos = trace
5341 .authorization_nanos
5342 .saturating_add(started.elapsed().as_nanos() as u64);
5343 });
5344 Ok(allowed)
5345 }
5346
5347 pub fn search(
5349 &mut self,
5350 request: &crate::query::SearchRequest,
5351 ) -> Result<Vec<crate::query::SearchHit>> {
5352 self.search_with_allowed(request, None)
5353 }
5354
5355 pub fn search_at(
5356 &mut self,
5357 request: &crate::query::SearchRequest,
5358 snapshot: Snapshot,
5359 authorized: Option<&std::collections::HashSet<RowId>>,
5360 ) -> Result<Vec<crate::query::SearchHit>> {
5361 self.search_at_with_allowed(request, snapshot, authorized)
5362 }
5363
5364 pub fn search_with_allowed(
5365 &mut self,
5366 request: &crate::query::SearchRequest,
5367 authorized: Option<&std::collections::HashSet<RowId>>,
5368 ) -> Result<Vec<crate::query::SearchHit>> {
5369 self.search_at_with_allowed(request, self.snapshot(), authorized)
5370 }
5371
5372 pub fn search_at_with_allowed(
5373 &mut self,
5374 request: &crate::query::SearchRequest,
5375 snapshot: Snapshot,
5376 authorized: Option<&std::collections::HashSet<RowId>>,
5377 ) -> Result<Vec<crate::query::SearchHit>> {
5378 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
5379 }
5380
5381 pub fn search_at_with_allowed_and_context(
5382 &mut self,
5383 request: &crate::query::SearchRequest,
5384 snapshot: Snapshot,
5385 authorized: Option<&std::collections::HashSet<RowId>>,
5386 context: Option<&crate::query::AiExecutionContext>,
5387 ) -> Result<Vec<crate::query::SearchHit>> {
5388 self.ensure_indexes_complete()?;
5389 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
5390 }
5391
5392 pub fn search_at_with_candidate_authorization_and_context(
5393 &mut self,
5394 request: &crate::query::SearchRequest,
5395 snapshot: Snapshot,
5396 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5397 context: Option<&crate::query::AiExecutionContext>,
5398 ) -> Result<Vec<crate::query::SearchHit>> {
5399 self.ensure_indexes_complete()?;
5400 self.search_at_with_filters_and_context(
5401 request,
5402 snapshot,
5403 None,
5404 authorization,
5405 context,
5406 None,
5407 )
5408 }
5409
5410 #[doc(hidden)]
5411 pub fn search_at_with_candidate_authorization_on_generation(
5412 &self,
5413 request: &crate::query::SearchRequest,
5414 snapshot: Snapshot,
5415 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5416 context: Option<&crate::query::AiExecutionContext>,
5417 ) -> Result<Vec<crate::query::SearchHit>> {
5418 self.search_at_with_filters_and_context(
5419 request,
5420 snapshot,
5421 None,
5422 authorization,
5423 context,
5424 None,
5425 )
5426 }
5427
5428 #[doc(hidden)]
5429 pub fn search_at_with_candidate_authorization_on_generation_after(
5430 &self,
5431 request: &crate::query::SearchRequest,
5432 snapshot: Snapshot,
5433 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5434 context: Option<&crate::query::AiExecutionContext>,
5435 after: Option<crate::query::SearchAfter>,
5436 ) -> Result<Vec<crate::query::SearchHit>> {
5437 self.search_at_with_filters_and_context(
5438 request,
5439 snapshot,
5440 None,
5441 authorization,
5442 context,
5443 after,
5444 )
5445 }
5446
5447 fn search_at_with_filters_and_context(
5448 &self,
5449 request: &crate::query::SearchRequest,
5450 snapshot: Snapshot,
5451 authorized: Option<&std::collections::HashSet<RowId>>,
5452 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
5453 context: Option<&crate::query::AiExecutionContext>,
5454 after: Option<crate::query::SearchAfter>,
5455 ) -> Result<Vec<crate::query::SearchHit>> {
5456 use crate::query::{
5457 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
5458 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
5459 };
5460 let total_started = std::time::Instant::now();
5461 let rank_offset = after.map_or(0, |after| after.returned_count);
5462 self.require_select()?;
5463 if request.limit == 0 {
5464 return Err(MongrelError::InvalidArgument(
5465 "search limit must be > 0".into(),
5466 ));
5467 }
5468 if request.limit > MAX_FINAL_LIMIT {
5469 return Err(MongrelError::InvalidArgument(format!(
5470 "search limit exceeds {MAX_FINAL_LIMIT}"
5471 )));
5472 }
5473 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
5474 return Err(MongrelError::InvalidArgument(
5475 "search-after score must be finite".into(),
5476 ));
5477 }
5478 if request.retrievers.is_empty() {
5479 return Err(MongrelError::InvalidArgument(
5480 "search requires at least one retriever".into(),
5481 ));
5482 }
5483 if request.retrievers.len() > MAX_RETRIEVERS {
5484 return Err(MongrelError::InvalidArgument(format!(
5485 "search exceeds {MAX_RETRIEVERS} retrievers"
5486 )));
5487 }
5488 if request.must.len() > MAX_HARD_CONDITIONS {
5489 return Err(MongrelError::InvalidArgument(format!(
5490 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
5491 )));
5492 }
5493 for condition in &request.must {
5494 self.validate_condition(condition)?;
5495 }
5496 if request.must.iter().any(|condition| {
5497 matches!(
5498 condition,
5499 Condition::Ann { .. }
5500 | Condition::SparseMatch { .. }
5501 | Condition::MinHashSimilar { .. }
5502 )
5503 }) {
5504 return Err(MongrelError::InvalidArgument(
5505 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
5506 .into(),
5507 ));
5508 }
5509 let mut names = std::collections::HashSet::new();
5510 for named in &request.retrievers {
5511 if named.name.is_empty()
5512 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
5513 || !names.insert(named.name.as_str())
5514 {
5515 return Err(MongrelError::InvalidArgument(format!(
5516 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
5517 crate::query::MAX_RETRIEVER_NAME_BYTES
5518 )));
5519 }
5520 if !named.weight.is_finite()
5521 || named.weight < 0.0
5522 || named.weight > MAX_RETRIEVER_WEIGHT
5523 {
5524 return Err(MongrelError::InvalidArgument(format!(
5525 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5526 )));
5527 }
5528 self.validate_retriever(&named.retriever)?;
5529 }
5530 let projection = request
5531 .projection
5532 .clone()
5533 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
5534 if projection.len() > MAX_PROJECTION_COLUMNS {
5535 return Err(MongrelError::InvalidArgument(format!(
5536 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
5537 )));
5538 }
5539 for column_id in &projection {
5540 if !self
5541 .schema
5542 .columns
5543 .iter()
5544 .any(|column| column.id == *column_id)
5545 {
5546 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
5547 }
5548 }
5549 if let Some(crate::query::Rerank::ExactVector {
5550 embedding_column,
5551 query,
5552 candidate_limit,
5553 weight,
5554 ..
5555 }) = &request.rerank
5556 {
5557 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
5558 {
5559 return Err(MongrelError::InvalidArgument(format!(
5560 "rerank candidate_limit must be between search limit and {}",
5561 crate::query::MAX_RETRIEVER_K
5562 )));
5563 }
5564 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
5565 return Err(MongrelError::InvalidArgument(format!(
5566 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
5567 )));
5568 }
5569 let column = self
5570 .schema
5571 .columns
5572 .iter()
5573 .find(|column| column.id == *embedding_column)
5574 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
5575 let crate::schema::TypeId::Embedding { dim } = column.ty else {
5576 return Err(MongrelError::InvalidArgument(format!(
5577 "rerank column {embedding_column} is not an embedding"
5578 )));
5579 };
5580 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
5581 return Err(MongrelError::InvalidArgument(format!(
5582 "rerank query must contain {dim} finite values"
5583 )));
5584 }
5585 }
5586
5587 let hard_filter_started = std::time::Instant::now();
5588 let hard_filter = if request.must.is_empty() {
5589 None
5590 } else {
5591 let mut sets = Vec::with_capacity(request.must.len());
5592 for condition in &request.must {
5593 if let Some(context) = context {
5594 context.checkpoint()?;
5595 }
5596 sets.push(self.resolve_condition(condition, snapshot)?);
5597 }
5598 Some(RowIdSet::intersect_many(sets))
5599 };
5600 crate::trace::QueryTrace::record(|trace| {
5601 trace.hard_filter_nanos = trace
5602 .hard_filter_nanos
5603 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
5604 });
5605 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
5606 return Ok(Vec::new());
5607 }
5608
5609 let constant = match request.fusion {
5610 Fusion::ReciprocalRank { constant } => constant,
5611 };
5612 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
5613 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
5614 let mut fusion_nanos = 0u64;
5615 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
5616 std::collections::HashMap::new();
5617 for named in retrievers {
5618 if named.weight == 0.0 {
5619 continue;
5620 }
5621 if let Some(context) = context {
5622 context.checkpoint()?;
5623 }
5624 let hits = self.retrieve_filtered(
5625 &named.retriever,
5626 snapshot,
5627 hard_filter.as_ref(),
5628 authorized,
5629 candidate_authorization,
5630 context,
5631 )?;
5632 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
5633 let fusion_started = std::time::Instant::now();
5634 for hit in hits {
5635 if let Some(context) = context {
5636 context.consume(1)?;
5637 }
5638 let contribution = named.weight / (constant as f64 + hit.rank as f64);
5639 if !contribution.is_finite() {
5640 return Err(MongrelError::InvalidArgument(
5641 "retriever contribution must be finite".into(),
5642 ));
5643 }
5644 let max_fused_candidates = context.map_or(
5645 crate::query::MAX_FUSED_CANDIDATES,
5646 crate::query::AiExecutionContext::max_fused_candidates,
5647 );
5648 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
5649 return Err(MongrelError::WorkBudgetExceeded);
5650 }
5651 let entry = fused.entry(hit.row_id).or_default();
5652 entry.0 += contribution;
5653 if !entry.0.is_finite() {
5654 return Err(MongrelError::InvalidArgument(
5655 "fused score must be finite".into(),
5656 ));
5657 }
5658 entry.1.push(ComponentScore {
5659 retriever_name: retriever_name.clone(),
5660 rank: hit.rank,
5661 raw_score: hit.score,
5662 contribution,
5663 });
5664 }
5665 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
5666 }
5667 let union_size = fused.len();
5668 let mut ranked: Vec<_> = fused
5669 .into_iter()
5670 .map(|(row_id, (fused_score, components))| {
5671 (row_id, fused_score, components, None, fused_score)
5672 })
5673 .collect();
5674 let order = |(a_row, _, _, _, a_score): &(
5675 RowId,
5676 f64,
5677 Vec<ComponentScore>,
5678 Option<f32>,
5679 f64,
5680 ),
5681 (b_row, _, _, _, b_score): &(
5682 RowId,
5683 f64,
5684 Vec<ComponentScore>,
5685 Option<f32>,
5686 f64,
5687 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
5688 if let Some(crate::query::Rerank::ExactVector {
5689 embedding_column,
5690 query,
5691 metric,
5692 candidate_limit,
5693 weight,
5694 }) = &request.rerank
5695 {
5696 let fused_order = |(a_row, a_score, ..): &(
5697 RowId,
5698 f64,
5699 Vec<ComponentScore>,
5700 Option<f32>,
5701 f64,
5702 ),
5703 (b_row, b_score, ..): &(
5704 RowId,
5705 f64,
5706 Vec<ComponentScore>,
5707 Option<f32>,
5708 f64,
5709 )| {
5710 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
5711 };
5712 let selection_started = std::time::Instant::now();
5713 if let Some(context) = context {
5714 context.consume(ranked.len())?;
5715 }
5716 if ranked.len() > *candidate_limit {
5717 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
5718 ranked.truncate(*candidate_limit);
5719 }
5720 ranked.sort_by(fused_order);
5721 fusion_nanos =
5722 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5723 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5724 if let Some(context) = context {
5725 context.consume(row_ids.len())?;
5726 }
5727 let query_now =
5728 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5729 let gather_started = std::time::Instant::now();
5730 let vectors = self.values_for_rids_batch_at_with_context(
5731 &row_ids,
5732 *embedding_column,
5733 snapshot,
5734 query_now,
5735 context,
5736 )?;
5737 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
5738 let vector_work =
5739 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
5740 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
5741 if let Some(context) = context {
5742 context.consume(vector_work)?;
5743 }
5744 query
5745 .iter()
5746 .map(|value| f64::from(*value).powi(2))
5747 .sum::<f64>()
5748 .sqrt()
5749 } else {
5750 0.0
5751 };
5752 let score_started = std::time::Instant::now();
5753 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
5754 for (row_id, value) in vectors {
5755 let Value::Embedding(vector) = value else {
5756 continue;
5757 };
5758 let score = match metric {
5759 crate::query::VectorMetric::DotProduct => {
5760 if let Some(context) = context {
5761 context.consume(vector_work)?;
5762 }
5763 query
5764 .iter()
5765 .zip(&vector)
5766 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5767 .sum::<f64>()
5768 }
5769 crate::query::VectorMetric::Cosine => {
5770 if let Some(context) = context {
5771 context.consume(vector_work.saturating_mul(2))?;
5772 }
5773 let dot = query
5774 .iter()
5775 .zip(&vector)
5776 .map(|(left, right)| f64::from(*left) * f64::from(*right))
5777 .sum::<f64>();
5778 let norm = vector
5779 .iter()
5780 .map(|value| f64::from(*value).powi(2))
5781 .sum::<f64>()
5782 .sqrt();
5783 if query_norm == 0.0 || norm == 0.0 {
5784 0.0
5785 } else {
5786 dot / (query_norm * norm)
5787 }
5788 }
5789 crate::query::VectorMetric::Euclidean => {
5790 if let Some(context) = context {
5791 context.consume(vector_work)?;
5792 }
5793 query
5794 .iter()
5795 .zip(&vector)
5796 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
5797 .sum::<f64>()
5798 .sqrt()
5799 }
5800 };
5801 if !score.is_finite() {
5802 return Err(MongrelError::InvalidArgument(
5803 "exact rerank score must be finite".into(),
5804 ));
5805 }
5806 scores.insert(row_id, score as f32);
5807 }
5808 let mut reranked = Vec::with_capacity(ranked.len());
5809 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
5810 let Some(score) = scores.get(&row_id).copied() else {
5811 continue;
5812 };
5813 let ordering_score = match metric {
5814 crate::query::VectorMetric::Euclidean => -f64::from(score),
5815 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
5816 f64::from(score)
5817 }
5818 };
5819 let final_score = fused_score + *weight * ordering_score;
5820 if !final_score.is_finite() {
5821 return Err(MongrelError::InvalidArgument(
5822 "final rerank score must be finite".into(),
5823 ));
5824 }
5825 reranked.push((row_id, fused_score, components, Some(score), final_score));
5826 }
5827 ranked = reranked;
5828 ranked.sort_by(order);
5829 crate::trace::QueryTrace::record(|trace| {
5830 trace.exact_vector_gather_nanos =
5831 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
5832 trace.exact_vector_score_nanos = trace
5833 .exact_vector_score_nanos
5834 .saturating_add(score_started.elapsed().as_nanos() as u64);
5835 });
5836 }
5837 if let Some(after) = after {
5838 ranked.retain(|(row_id, _, _, _, final_score)| {
5839 final_score.total_cmp(&after.final_score).is_lt()
5840 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
5841 });
5842 }
5843 let projection_started = std::time::Instant::now();
5844 let sentinel = projection
5845 .first()
5846 .copied()
5847 .or_else(|| self.schema.columns.first().map(|column| column.id));
5848 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
5849 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
5850 let mut projection_rows = 0usize;
5851 let mut projection_cells = 0usize;
5852 while out.len() < request.limit && !ranked.is_empty() {
5853 if let Some(context) = context {
5854 context.checkpoint()?;
5855 context.consume(ranked.len())?;
5856 }
5857 let needed = request.limit - out.len();
5858 let window_size = ranked
5859 .len()
5860 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
5861 let selection_started = std::time::Instant::now();
5862 let mut remainder = if ranked.len() > window_size {
5863 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
5864 ranked.split_off(window_size)
5865 } else {
5866 Vec::new()
5867 };
5868 ranked.sort_by(order);
5869 fusion_nanos =
5870 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
5871 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
5872 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
5873 if let Some(context) = context {
5874 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
5875 }
5876 projection_rows = projection_rows.saturating_add(row_ids.len());
5877 projection_cells =
5878 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
5879 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
5880 std::collections::HashMap::new();
5881 if let Some(column_id) = sentinel {
5882 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5883 &row_ids, column_id, snapshot, query_now, context,
5884 )? {
5885 cells.entry(row_id).or_default().insert(column_id, value);
5886 }
5887 }
5888 for &column_id in &projection {
5889 if Some(column_id) == sentinel {
5890 continue;
5891 }
5892 for (row_id, value) in self.values_for_rids_batch_at_with_context(
5893 &row_ids, column_id, snapshot, query_now, context,
5894 )? {
5895 cells.entry(row_id).or_default().insert(column_id, value);
5896 }
5897 }
5898 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
5899 ranked.drain(..)
5900 {
5901 let Some(row_cells) = cells.remove(&row_id) else {
5902 continue;
5903 };
5904 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
5905 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
5906 out.push(SearchHit {
5907 row_id,
5908 cells: projection
5909 .iter()
5910 .filter_map(|column_id| {
5911 row_cells
5912 .get(column_id)
5913 .cloned()
5914 .map(|value| (*column_id, value))
5915 })
5916 .collect(),
5917 components,
5918 fused_score,
5919 exact_rerank_score,
5920 final_score,
5921 final_rank,
5922 });
5923 if out.len() == request.limit {
5924 break;
5925 }
5926 }
5927 ranked.append(&mut remainder);
5928 }
5929 crate::trace::QueryTrace::record(|trace| {
5930 trace.union_size = union_size;
5931 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
5932 trace.projection_nanos = trace
5933 .projection_nanos
5934 .saturating_add(projection_started.elapsed().as_nanos() as u64);
5935 trace.total_nanos = trace
5936 .total_nanos
5937 .saturating_add(total_started.elapsed().as_nanos() as u64);
5938 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
5939 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
5940 if let Some(context) = context {
5941 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
5942 }
5943 });
5944 Ok(out)
5945 }
5946
5947 pub fn set_similarity(
5950 &mut self,
5951 request: &crate::query::SetSimilarityRequest,
5952 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5953 self.set_similarity_with_allowed(request, None)
5954 }
5955
5956 pub fn set_similarity_at(
5957 &mut self,
5958 request: &crate::query::SetSimilarityRequest,
5959 snapshot: Snapshot,
5960 allowed: Option<&std::collections::HashSet<RowId>>,
5961 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
5962 self.set_similarity_explained_at(request, snapshot, allowed)
5963 .map(|(hits, _)| hits)
5964 }
5965
5966 pub fn ann_rerank(
5968 &mut self,
5969 request: &crate::query::AnnRerankRequest,
5970 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5971 self.ann_rerank_with_allowed(request, None)
5972 }
5973
5974 pub fn ann_rerank_with_allowed(
5975 &mut self,
5976 request: &crate::query::AnnRerankRequest,
5977 allowed: Option<&std::collections::HashSet<RowId>>,
5978 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5979 self.ann_rerank_at(request, self.snapshot(), allowed)
5980 }
5981
5982 pub fn ann_rerank_at(
5983 &mut self,
5984 request: &crate::query::AnnRerankRequest,
5985 snapshot: Snapshot,
5986 allowed: Option<&std::collections::HashSet<RowId>>,
5987 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5988 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
5989 }
5990
5991 pub fn ann_rerank_at_with_context(
5992 &mut self,
5993 request: &crate::query::AnnRerankRequest,
5994 snapshot: Snapshot,
5995 allowed: Option<&std::collections::HashSet<RowId>>,
5996 context: Option<&crate::query::AiExecutionContext>,
5997 ) -> Result<Vec<crate::query::AnnRerankHit>> {
5998 self.ensure_indexes_complete()?;
5999 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
6000 }
6001
6002 pub fn ann_rerank_at_with_candidate_authorization_and_context(
6003 &mut self,
6004 request: &crate::query::AnnRerankRequest,
6005 snapshot: Snapshot,
6006 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6007 context: Option<&crate::query::AiExecutionContext>,
6008 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6009 self.ensure_indexes_complete()?;
6010 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6011 }
6012
6013 #[doc(hidden)]
6014 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
6015 &self,
6016 request: &crate::query::AnnRerankRequest,
6017 snapshot: Snapshot,
6018 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6019 context: Option<&crate::query::AiExecutionContext>,
6020 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6021 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
6022 }
6023
6024 fn ann_rerank_at_with_filters_and_context(
6025 &self,
6026 request: &crate::query::AnnRerankRequest,
6027 snapshot: Snapshot,
6028 allowed: Option<&std::collections::HashSet<RowId>>,
6029 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6030 context: Option<&crate::query::AiExecutionContext>,
6031 ) -> Result<Vec<crate::query::AnnRerankHit>> {
6032 use crate::query::{
6033 AnnRerankHit, Retriever, RetrieverScore, VectorMetric, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6034 };
6035 if request.candidate_k == 0 || request.limit == 0 {
6036 return Err(MongrelError::InvalidArgument(
6037 "candidate_k and limit must be > 0".into(),
6038 ));
6039 }
6040 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
6041 return Err(MongrelError::InvalidArgument(format!(
6042 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
6043 )));
6044 }
6045 let retriever = Retriever::Ann {
6046 column_id: request.column_id,
6047 query: request.query.clone(),
6048 k: request.candidate_k,
6049 };
6050 self.require_select()?;
6051 self.validate_retriever(&retriever)?;
6052 let hits = self.retrieve_filtered(
6053 &retriever,
6054 snapshot,
6055 None,
6056 allowed,
6057 candidate_authorization,
6058 context,
6059 )?;
6060 let distances: std::collections::HashMap<_, _> = hits
6061 .iter()
6062 .filter_map(|hit| match hit.score {
6063 RetrieverScore::AnnHammingDistance(distance) => Some((hit.row_id, distance)),
6064 _ => None,
6065 })
6066 .collect();
6067 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6068 if let Some(context) = context {
6069 context.consume(row_ids.len())?;
6070 }
6071 let gather_started = std::time::Instant::now();
6072 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6073 let values = self.values_for_rids_batch_at_with_context(
6074 &row_ids,
6075 request.column_id,
6076 snapshot,
6077 query_now,
6078 context,
6079 )?;
6080 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
6081 let score_started = std::time::Instant::now();
6082 let vector_work =
6083 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
6084 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
6085 if let Some(context) = context {
6086 context.consume(vector_work)?;
6087 }
6088 request
6089 .query
6090 .iter()
6091 .map(|value| f64::from(*value).powi(2))
6092 .sum::<f64>()
6093 .sqrt()
6094 } else {
6095 0.0
6096 };
6097 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
6098 for (row_id, value) in values {
6099 let Value::Embedding(vector) = value else {
6100 continue;
6101 };
6102 let exact_score = match request.metric {
6103 VectorMetric::DotProduct => {
6104 if let Some(context) = context {
6105 context.consume(vector_work)?;
6106 }
6107 request
6108 .query
6109 .iter()
6110 .zip(&vector)
6111 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6112 .sum::<f64>()
6113 }
6114 VectorMetric::Cosine => {
6115 if let Some(context) = context {
6116 context.consume(vector_work.saturating_mul(2))?;
6117 }
6118 let dot = request
6119 .query
6120 .iter()
6121 .zip(&vector)
6122 .map(|(left, right)| f64::from(*left) * f64::from(*right))
6123 .sum::<f64>();
6124 let norm = vector
6125 .iter()
6126 .map(|value| f64::from(*value).powi(2))
6127 .sum::<f64>()
6128 .sqrt();
6129 if query_norm == 0.0 || norm == 0.0 {
6130 0.0
6131 } else {
6132 dot / (query_norm * norm)
6133 }
6134 }
6135 VectorMetric::Euclidean => {
6136 if let Some(context) = context {
6137 context.consume(vector_work)?;
6138 }
6139 request
6140 .query
6141 .iter()
6142 .zip(&vector)
6143 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
6144 .sum::<f64>()
6145 .sqrt()
6146 }
6147 };
6148 let exact_score = exact_score as f32;
6149 if !exact_score.is_finite() {
6150 return Err(MongrelError::InvalidArgument(
6151 "exact ANN score must be finite".into(),
6152 ));
6153 }
6154 reranked.push(AnnRerankHit {
6155 row_id,
6156 hamming_distance: distances.get(&row_id).copied().unwrap_or_default(),
6157 exact_score,
6158 });
6159 }
6160 reranked.sort_by(|left, right| {
6161 let score = match request.metric {
6162 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
6163 VectorMetric::Cosine | VectorMetric::DotProduct => {
6164 right.exact_score.total_cmp(&left.exact_score)
6165 }
6166 };
6167 score.then_with(|| left.row_id.cmp(&right.row_id))
6168 });
6169 reranked.truncate(request.limit);
6170 crate::trace::QueryTrace::record(|trace| {
6171 trace.exact_vector_gather_nanos =
6172 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
6173 trace.exact_vector_score_nanos = trace
6174 .exact_vector_score_nanos
6175 .saturating_add(score_started.elapsed().as_nanos() as u64);
6176 });
6177 Ok(reranked)
6178 }
6179
6180 pub fn set_similarity_with_allowed(
6181 &mut self,
6182 request: &crate::query::SetSimilarityRequest,
6183 allowed: Option<&std::collections::HashSet<RowId>>,
6184 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6185 self.set_similarity_explained_at(request, self.snapshot(), allowed)
6186 .map(|(hits, _)| hits)
6187 }
6188
6189 pub fn set_similarity_explained(
6190 &mut self,
6191 request: &crate::query::SetSimilarityRequest,
6192 ) -> Result<(
6193 Vec<crate::query::SetSimilarityHit>,
6194 crate::query::SetSimilarityTrace,
6195 )> {
6196 self.set_similarity_explained_at(request, self.snapshot(), None)
6197 }
6198
6199 fn set_similarity_explained_at(
6200 &mut self,
6201 request: &crate::query::SetSimilarityRequest,
6202 snapshot: Snapshot,
6203 allowed: Option<&std::collections::HashSet<RowId>>,
6204 ) -> Result<(
6205 Vec<crate::query::SetSimilarityHit>,
6206 crate::query::SetSimilarityTrace,
6207 )> {
6208 self.ensure_indexes_complete()?;
6209 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
6210 }
6211
6212 pub fn set_similarity_at_with_context(
6213 &mut self,
6214 request: &crate::query::SetSimilarityRequest,
6215 snapshot: Snapshot,
6216 allowed: Option<&std::collections::HashSet<RowId>>,
6217 context: Option<&crate::query::AiExecutionContext>,
6218 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6219 self.ensure_indexes_complete()?;
6220 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
6221 .map(|(hits, _)| hits)
6222 }
6223
6224 pub fn set_similarity_at_with_candidate_authorization_and_context(
6225 &mut self,
6226 request: &crate::query::SetSimilarityRequest,
6227 snapshot: Snapshot,
6228 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6229 context: Option<&crate::query::AiExecutionContext>,
6230 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6231 self.ensure_indexes_complete()?;
6232 self.set_similarity_explained_at_with_context(
6233 request,
6234 snapshot,
6235 None,
6236 authorization,
6237 context,
6238 )
6239 .map(|(hits, _)| hits)
6240 }
6241
6242 #[doc(hidden)]
6243 pub fn set_similarity_at_with_candidate_authorization_on_generation(
6244 &self,
6245 request: &crate::query::SetSimilarityRequest,
6246 snapshot: Snapshot,
6247 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6248 context: Option<&crate::query::AiExecutionContext>,
6249 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
6250 self.set_similarity_explained_at_with_context(
6251 request,
6252 snapshot,
6253 None,
6254 authorization,
6255 context,
6256 )
6257 .map(|(hits, _)| hits)
6258 }
6259
6260 fn set_similarity_explained_at_with_context(
6261 &self,
6262 request: &crate::query::SetSimilarityRequest,
6263 snapshot: Snapshot,
6264 allowed: Option<&std::collections::HashSet<RowId>>,
6265 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6266 context: Option<&crate::query::AiExecutionContext>,
6267 ) -> Result<(
6268 Vec<crate::query::SetSimilarityHit>,
6269 crate::query::SetSimilarityTrace,
6270 )> {
6271 use crate::query::{
6272 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
6273 MAX_SET_MEMBERS,
6274 };
6275 let mut trace = crate::query::SetSimilarityTrace::default();
6276 if request.members.is_empty() {
6277 return Ok((Vec::new(), trace));
6278 }
6279 if request.candidate_k == 0 || request.limit == 0 {
6280 return Err(MongrelError::InvalidArgument(
6281 "candidate_k and limit must be > 0".into(),
6282 ));
6283 }
6284 if request.candidate_k > MAX_RETRIEVER_K
6285 || request.limit > MAX_FINAL_LIMIT
6286 || request.members.len() > MAX_SET_MEMBERS
6287 {
6288 return Err(MongrelError::InvalidArgument(format!(
6289 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
6290 )));
6291 }
6292 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
6293 return Err(MongrelError::InvalidArgument(
6294 "min_jaccard must be finite and between 0 and 1".into(),
6295 ));
6296 }
6297 let started = std::time::Instant::now();
6298 let retriever = Retriever::MinHash {
6299 column_id: request.column_id,
6300 members: request.members.clone(),
6301 k: request.candidate_k,
6302 };
6303 self.require_select()?;
6304 self.validate_retriever(&retriever)?;
6305 let hits = self.retrieve_filtered(
6306 &retriever,
6307 snapshot,
6308 None,
6309 allowed,
6310 candidate_authorization,
6311 context,
6312 )?;
6313 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
6314 trace.candidate_count = hits.len();
6315 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
6316 if let Some(context) = context {
6317 context.consume(row_ids.len())?;
6318 }
6319 let started = std::time::Instant::now();
6320 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6321 let values = self.values_for_rids_batch_at_with_context(
6322 &row_ids,
6323 request.column_id,
6324 snapshot,
6325 query_now,
6326 context,
6327 )?;
6328 trace.gather_us = started.elapsed().as_micros() as u64;
6329 if let Some(context) = context {
6330 context.consume(request.members.len())?;
6331 }
6332 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
6333 let estimates: std::collections::HashMap<_, _> = hits
6334 .into_iter()
6335 .filter_map(|hit| match hit.score {
6336 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
6337 _ => None,
6338 })
6339 .collect();
6340 let started = std::time::Instant::now();
6341 let mut parsed = Vec::with_capacity(values.len());
6342 for (row_id, value) in values {
6343 let Value::Bytes(bytes) = value else {
6344 continue;
6345 };
6346 if let Some(context) = context {
6347 context.consume(crate::query::work_units(
6348 bytes.len(),
6349 crate::query::PARSE_WORK_QUANTUM,
6350 ))?;
6351 }
6352 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
6353 continue;
6354 };
6355 if let Some(context) = context {
6356 context.consume(members.len())?;
6357 }
6358 let stored = members
6359 .into_iter()
6360 .filter_map(|member| match member {
6361 serde_json::Value::String(value) => {
6362 Some(crate::query::SetMember::String(value))
6363 }
6364 serde_json::Value::Number(value) => {
6365 Some(crate::query::SetMember::Number(value))
6366 }
6367 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
6368 _ => None,
6369 })
6370 .collect::<std::collections::HashSet<_>>();
6371 parsed.push((row_id, stored));
6372 }
6373 trace.parse_us = started.elapsed().as_micros() as u64;
6374 trace.verified_count = parsed.len();
6375 let started = std::time::Instant::now();
6376 let mut exact = Vec::new();
6377 for (row_id, stored) in parsed {
6378 if let Some(context) = context {
6379 context.consume(query.len().saturating_add(stored.len()))?;
6380 }
6381 let union = query.union(&stored).count();
6382 let score = if union == 0 {
6383 1.0
6384 } else {
6385 query.intersection(&stored).count() as f32 / union as f32
6386 };
6387 if score >= request.min_jaccard {
6388 exact.push(SetSimilarityHit {
6389 row_id,
6390 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
6391 exact_jaccard: score,
6392 });
6393 }
6394 }
6395 exact.sort_by(|a, b| {
6396 b.exact_jaccard
6397 .total_cmp(&a.exact_jaccard)
6398 .then_with(|| a.row_id.cmp(&b.row_id))
6399 });
6400 exact.truncate(request.limit);
6401 trace.score_us = started.elapsed().as_micros() as u64;
6402 crate::trace::QueryTrace::record(|query_trace| {
6403 query_trace.exact_set_gather_nanos = query_trace
6404 .exact_set_gather_nanos
6405 .saturating_add(trace.gather_us.saturating_mul(1_000));
6406 query_trace.exact_set_parse_nanos = query_trace
6407 .exact_set_parse_nanos
6408 .saturating_add(trace.parse_us.saturating_mul(1_000));
6409 query_trace.exact_set_score_nanos = query_trace
6410 .exact_set_score_nanos
6411 .saturating_add(trace.score_us.saturating_mul(1_000));
6412 });
6413 Ok((exact, trace))
6414 }
6415
6416 fn values_for_rids_batch_at(
6418 &self,
6419 row_ids: &[u64],
6420 column_id: u16,
6421 snapshot: Snapshot,
6422 now: i64,
6423 ) -> Result<Vec<(RowId, Value)>> {
6424 if self.ttl.is_none()
6425 && self.memtable.is_empty()
6426 && self.mutable_run.is_empty()
6427 && self.run_refs.len() == 1
6428 {
6429 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6430 if row_ids.len().saturating_mul(24) < reader.row_count() {
6435 let mut values = Vec::with_capacity(row_ids.len());
6436 for &raw_row_id in row_ids {
6437 let row_id = RowId(raw_row_id);
6438 if let Some((_, false, Some(value))) =
6439 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6440 {
6441 values.push((row_id, value));
6442 }
6443 }
6444 return Ok(values);
6445 }
6446 let (positions, visible_row_ids) =
6447 reader.visible_positions_with_rids(snapshot.epoch)?;
6448 let requested: Vec<(RowId, usize)> = row_ids
6449 .iter()
6450 .filter_map(|raw| {
6451 visible_row_ids
6452 .binary_search(&(*raw as i64))
6453 .ok()
6454 .map(|index| (RowId(*raw), positions[index]))
6455 })
6456 .collect();
6457 let values = reader.gather_column(
6458 column_id,
6459 &requested
6460 .iter()
6461 .map(|(_, position)| *position)
6462 .collect::<Vec<_>>(),
6463 )?;
6464 return Ok(requested
6465 .into_iter()
6466 .zip(values)
6467 .map(|((row_id, _), value)| (row_id, value))
6468 .collect());
6469 }
6470 self.values_for_rids_at(row_ids, column_id, snapshot, now)
6471 }
6472
6473 fn values_for_rids_batch_at_with_context(
6474 &self,
6475 row_ids: &[u64],
6476 column_id: u16,
6477 snapshot: Snapshot,
6478 now: i64,
6479 context: Option<&crate::query::AiExecutionContext>,
6480 ) -> Result<Vec<(RowId, Value)>> {
6481 let Some(context) = context else {
6482 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
6483 };
6484 let mut values = Vec::with_capacity(row_ids.len());
6485 for chunk in row_ids.chunks(256) {
6486 context.checkpoint()?;
6487 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
6488 }
6489 Ok(values)
6490 }
6491
6492 fn values_for_rids_at(
6494 &self,
6495 row_ids: &[u64],
6496 column_id: u16,
6497 snapshot: Snapshot,
6498 now: i64,
6499 ) -> Result<Vec<(RowId, Value)>> {
6500 let mut readers: Vec<_> = self
6501 .run_refs
6502 .iter()
6503 .map(|run| self.open_reader(run.run_id))
6504 .collect::<Result<_>>()?;
6505 let mut out = Vec::with_capacity(row_ids.len());
6506 for &raw_row_id in row_ids {
6507 let row_id = RowId(raw_row_id);
6508 let mem = self.memtable.get_version(row_id, snapshot.epoch);
6509 let mutable = self.mutable_run.get_version(row_id, snapshot.epoch);
6510 let overlay = match (mem, mutable) {
6511 (Some((a_epoch, a)), Some((b_epoch, b))) => Some(if a_epoch >= b_epoch {
6512 (a_epoch, a)
6513 } else {
6514 (b_epoch, b)
6515 }),
6516 (Some(value), None) | (None, Some(value)) => Some(value),
6517 (None, None) => None,
6518 };
6519 if let Some((_, row)) = overlay {
6520 if !row.deleted && !self.row_expired_at(&row, now) {
6521 if let Some(value) = row.columns.get(&column_id) {
6522 out.push((row_id, value.clone()));
6523 }
6524 }
6525 continue;
6526 }
6527
6528 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
6529 for (index, reader) in readers.iter_mut().enumerate() {
6530 if let Some((epoch, deleted, value)) =
6531 reader.get_version_column(row_id, snapshot.epoch, column_id)?
6532 {
6533 if best
6534 .as_ref()
6535 .map(|(best_epoch, ..)| epoch > *best_epoch)
6536 .unwrap_or(true)
6537 {
6538 best = Some((epoch, deleted, value, index));
6539 }
6540 }
6541 }
6542 let Some((_, false, Some(value), reader_index)) = best else {
6543 continue;
6544 };
6545 if let Some(ttl) = self.ttl {
6546 if ttl.column_id != column_id {
6547 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6548 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6549 {
6550 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6551 continue;
6552 }
6553 }
6554 } else if let Value::Int64(timestamp) = value {
6555 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6556 continue;
6557 }
6558 }
6559 }
6560 out.push((row_id, value));
6561 }
6562 Ok(out)
6563 }
6564
6565 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
6570 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now())
6571 }
6572
6573 pub fn rows_for_rids_with_context(
6574 &self,
6575 rids: &[u64],
6576 snapshot: Snapshot,
6577 context: &crate::query::AiExecutionContext,
6578 ) -> Result<Vec<Row>> {
6579 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
6580 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos())
6581 }
6582
6583 fn rows_for_rids_at_time(
6584 &self,
6585 rids: &[u64],
6586 snapshot: Snapshot,
6587 ttl_now: i64,
6588 ) -> Result<Vec<Row>> {
6589 use std::collections::HashMap;
6590 let mut rows = Vec::with_capacity(rids.len());
6591 let tier_size = self.memtable.len() + self.mutable_run.len();
6608 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
6609 if rids.len().saturating_mul(24) < tier_size {
6610 for &rid in rids {
6611 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
6612 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
6613 let newest = match (mem, mrun) {
6614 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
6615 (Some((_, mr)), None) => Some(mr),
6616 (None, Some((_, rr))) => Some(rr),
6617 (None, None) => None,
6618 };
6619 if let Some(row) = newest {
6620 overlay.insert(rid, row);
6621 }
6622 }
6623 } else {
6624 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
6625 overlay
6626 .entry(row.row_id.0)
6627 .and_modify(|e| {
6628 if row.committed_epoch > e.committed_epoch {
6629 *e = row.clone();
6630 }
6631 })
6632 .or_insert(row);
6633 };
6634 for row in self.memtable.visible_versions(snapshot.epoch) {
6635 fold_newest(row, &mut overlay);
6636 }
6637 for row in self.mutable_run.visible_versions(snapshot.epoch) {
6638 fold_newest(row, &mut overlay);
6639 }
6640 }
6641 if self.run_refs.len() == 1 {
6642 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
6643 if rids.len().saturating_mul(24) < reader.row_count() {
6651 for &rid in rids {
6652 if let Some(r) = overlay.get(&rid) {
6653 if !r.deleted {
6654 rows.push(r.clone());
6655 }
6656 continue;
6657 }
6658 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
6659 if !row.deleted {
6660 rows.push(row);
6661 }
6662 }
6663 }
6664 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6665 return Ok(rows);
6666 }
6667 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
6676 enum Src {
6679 Overlay,
6680 Run,
6681 }
6682 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
6683 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
6684 for rid in rids {
6685 if overlay.contains_key(rid) {
6686 plan.push(Src::Overlay);
6687 continue;
6688 }
6689 match vis_rids.binary_search(&(*rid as i64)) {
6690 Ok(i) => {
6691 plan.push(Src::Run);
6692 fetch.push(positions[i]);
6693 }
6694 Err(_) => { }
6695 }
6696 }
6697 let fetched = reader.materialize_batch(&fetch)?;
6698 let mut fetched_iter = fetched.into_iter();
6699 for (rid, src) in rids.iter().zip(plan) {
6700 match src {
6701 Src::Overlay => {
6702 if let Some(r) = overlay.get(rid) {
6703 if !r.deleted {
6704 rows.push(r.clone());
6705 }
6706 }
6707 }
6708 Src::Run => {
6709 if let Some(row) = fetched_iter.next() {
6710 if !row.deleted {
6711 rows.push(row);
6712 }
6713 }
6714 }
6715 }
6716 }
6717 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6718 return Ok(rows);
6719 }
6720 let mut readers: Vec<_> = self
6724 .run_refs
6725 .iter()
6726 .map(|rr| self.open_reader(rr.run_id))
6727 .collect::<Result<Vec<_>>>()?;
6728 for rid in rids {
6729 if let Some(r) = overlay.get(rid) {
6730 if !r.deleted {
6731 rows.push(r.clone());
6732 }
6733 continue;
6734 }
6735 let mut best: Option<(Epoch, Row)> = None;
6736 for reader in readers.iter_mut() {
6737 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
6738 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
6739 best = Some((epoch, row));
6740 }
6741 }
6742 }
6743 if let Some((_, r)) = best {
6744 if !r.deleted {
6745 rows.push(r);
6746 }
6747 }
6748 }
6749 rows.retain(|row| !self.row_expired_at(row, ttl_now));
6750 Ok(rows)
6751 }
6752
6753 pub fn indexes_complete(&self) -> bool {
6763 self.indexes_complete
6764 }
6765
6766 pub fn index_build_policy(&self) -> IndexBuildPolicy {
6768 self.index_build_policy
6769 }
6770
6771 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
6775 self.index_build_policy = policy;
6776 }
6777
6778 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
6783 if !self.indexes_complete {
6787 return None;
6788 }
6789 let b = self.bitmap.get(&column_id)?;
6790 let result: Vec<Vec<u8>> = b
6791 .keys()
6792 .into_iter()
6793 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
6794 .collect();
6795 Some(result)
6796 }
6797
6798 pub fn fk_join_row_ids(
6799 &self,
6800 fk_column_id: u16,
6801 pk_values: &[Vec<u8>],
6802 fk_conditions: &[crate::query::Condition],
6803 snapshot: Snapshot,
6804 ) -> Result<Vec<u64>> {
6805 let Some(b) = self.bitmap.get(&fk_column_id) else {
6806 return Ok(Vec::new());
6807 };
6808 let mut join_set = {
6809 let mut acc = roaring::RoaringBitmap::new();
6810 for v in pk_values {
6811 acc |= b.get(v);
6812 }
6813 RowIdSet::from_roaring(acc)
6814 };
6815 if !fk_conditions.is_empty() {
6816 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
6817 sets.push(join_set);
6818 for c in fk_conditions {
6819 sets.push(self.resolve_condition(c, snapshot)?);
6820 }
6821 join_set = RowIdSet::intersect_many(sets);
6822 }
6823 Ok(join_set.into_sorted_vec())
6824 }
6825
6826 pub fn fk_join_count(
6832 &self,
6833 fk_column_id: u16,
6834 pk_values: &[Vec<u8>],
6835 fk_conditions: &[crate::query::Condition],
6836 snapshot: Snapshot,
6837 ) -> Result<u64> {
6838 let Some(b) = self.bitmap.get(&fk_column_id) else {
6839 return Ok(0);
6840 };
6841 let mut acc = roaring::RoaringBitmap::new();
6842 for v in pk_values {
6843 acc |= b.get(v);
6844 }
6845 if fk_conditions.is_empty() {
6846 return Ok(acc.len());
6847 }
6848 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
6849 sets.push(RowIdSet::from_roaring(acc));
6850 for c in fk_conditions {
6851 sets.push(self.resolve_condition(c, snapshot)?);
6852 }
6853 Ok(RowIdSet::intersect_many(sets).len() as u64)
6854 }
6855
6856 fn resolve_condition(
6861 &self,
6862 c: &crate::query::Condition,
6863 snapshot: Snapshot,
6864 ) -> Result<RowIdSet> {
6865 self.resolve_condition_with_allowed(c, snapshot, None)
6866 }
6867
6868 fn resolve_condition_with_allowed(
6869 &self,
6870 c: &crate::query::Condition,
6871 snapshot: Snapshot,
6872 allowed: Option<&std::collections::HashSet<RowId>>,
6873 ) -> Result<RowIdSet> {
6874 use crate::query::Condition;
6875 self.validate_condition(c)?;
6876 Ok(match c {
6877 Condition::Pk(key) => {
6878 let lookup = self
6879 .schema
6880 .primary_key()
6881 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
6882 .unwrap_or_else(|| key.clone());
6883 self.hot
6884 .get(&lookup)
6885 .map(|r| RowIdSet::one(r.0))
6886 .unwrap_or_else(RowIdSet::empty)
6887 }
6888 Condition::BitmapEq { column_id, value } => {
6889 let lookup = self.index_lookup_key_bytes(*column_id, value);
6890 self.bitmap
6891 .get(column_id)
6892 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
6893 .unwrap_or_else(RowIdSet::empty)
6894 }
6895 Condition::BitmapIn { column_id, values } => {
6896 let bm = self.bitmap.get(column_id);
6897 let mut acc = roaring::RoaringBitmap::new();
6898 if let Some(b) = bm {
6899 for v in values {
6900 let lookup = self.index_lookup_key_bytes(*column_id, v);
6901 acc |= b.get(&lookup);
6902 }
6903 }
6904 RowIdSet::from_roaring(acc)
6905 }
6906 Condition::BytesPrefix { column_id, prefix } => {
6907 if let Some(b) = self.bitmap.get(column_id) {
6912 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
6913 let mut acc = roaring::RoaringBitmap::new();
6914 for key in b.keys() {
6915 if key.starts_with(&lookup_prefix) {
6916 acc |= b.get(&key);
6917 }
6918 }
6919 RowIdSet::from_roaring(acc)
6920 } else {
6921 RowIdSet::empty()
6922 }
6923 }
6924 Condition::FmContains { column_id, pattern } => self
6925 .fm
6926 .get(column_id)
6927 .map(|f| {
6928 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
6929 })
6930 .unwrap_or_else(RowIdSet::empty),
6931 Condition::FmContainsAll {
6932 column_id,
6933 patterns,
6934 } => {
6935 if let Some(f) = self.fm.get(column_id) {
6938 let sets: Vec<RowIdSet> = patterns
6939 .iter()
6940 .map(|pat| {
6941 RowIdSet::from_unsorted(
6942 f.locate(pat).into_iter().map(|r| r.0).collect(),
6943 )
6944 })
6945 .collect();
6946 RowIdSet::intersect_many(sets)
6947 } else {
6948 RowIdSet::empty()
6949 }
6950 }
6951 Condition::Ann {
6952 column_id,
6953 query,
6954 k,
6955 } => RowIdSet::from_unsorted(
6956 self.retrieve_filtered(
6957 &crate::query::Retriever::Ann {
6958 column_id: *column_id,
6959 query: query.clone(),
6960 k: *k,
6961 },
6962 snapshot,
6963 None,
6964 allowed,
6965 None,
6966 None,
6967 )?
6968 .into_iter()
6969 .map(|hit| hit.row_id.0)
6970 .collect(),
6971 ),
6972 Condition::SparseMatch {
6973 column_id,
6974 query,
6975 k,
6976 } => RowIdSet::from_unsorted(
6977 self.retrieve_filtered(
6978 &crate::query::Retriever::Sparse {
6979 column_id: *column_id,
6980 query: query.clone(),
6981 k: *k,
6982 },
6983 snapshot,
6984 None,
6985 allowed,
6986 None,
6987 None,
6988 )?
6989 .into_iter()
6990 .map(|hit| hit.row_id.0)
6991 .collect(),
6992 ),
6993 Condition::MinHashSimilar {
6994 column_id,
6995 query,
6996 k,
6997 } => match self.minhash.get(column_id) {
6998 Some(index) => {
6999 let candidates = index.candidate_row_ids(query);
7000 let eligible =
7001 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
7002 RowIdSet::from_unsorted(
7003 index
7004 .search_filtered(query, *k, |row_id| {
7005 eligible.contains(&row_id)
7006 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
7007 })
7008 .into_iter()
7009 .map(|(row_id, _)| row_id.0)
7010 .collect(),
7011 )
7012 }
7013 None => RowIdSet::empty(),
7014 },
7015 Condition::Range { column_id, lo, hi } => {
7016 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7025 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
7026 } else if self.run_refs.len() == 1 {
7027 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7028 r.range_row_id_set_i64(*column_id, *lo, *hi)?
7029 } else {
7030 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
7031 };
7032 set.remove_many(self.overlay_rid_set(snapshot));
7033 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
7034 set
7035 }
7036 Condition::RangeF64 {
7037 column_id,
7038 lo,
7039 lo_inclusive,
7040 hi,
7041 hi_inclusive,
7042 } => {
7043 let mut set = if let Some(li) = self.learned_range.get(column_id) {
7046 RowIdSet::from_unsorted(
7047 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
7048 .into_iter()
7049 .collect(),
7050 )
7051 } else if self.run_refs.len() == 1 {
7052 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7053 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
7054 } else {
7055 return self.range_scan_f64(
7056 *column_id,
7057 *lo,
7058 *lo_inclusive,
7059 *hi,
7060 *hi_inclusive,
7061 snapshot,
7062 );
7063 };
7064 set.remove_many(self.overlay_rid_set(snapshot));
7065 self.range_scan_overlay_f64(
7066 &mut set,
7067 *column_id,
7068 *lo,
7069 *lo_inclusive,
7070 *hi,
7071 *hi_inclusive,
7072 snapshot,
7073 );
7074 set
7075 }
7076 Condition::IsNull { column_id } => {
7077 let mut set = if self.run_refs.len() == 1 {
7078 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7079 r.null_row_id_set(*column_id, true)?
7080 } else {
7081 return self.null_scan(*column_id, true, snapshot);
7082 };
7083 set.remove_many(self.overlay_rid_set(snapshot));
7084 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
7085 set
7086 }
7087 Condition::IsNotNull { column_id } => {
7088 let mut set = if self.run_refs.len() == 1 {
7089 let mut r = self.open_reader(self.run_refs[0].run_id)?;
7090 r.null_row_id_set(*column_id, false)?
7091 } else {
7092 return self.null_scan(*column_id, false, snapshot);
7093 };
7094 set.remove_many(self.overlay_rid_set(snapshot));
7095 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
7096 set
7097 }
7098 })
7099 }
7100
7101 fn range_scan_i64(
7109 &self,
7110 column_id: u16,
7111 lo: i64,
7112 hi: i64,
7113 snapshot: Snapshot,
7114 ) -> Result<RowIdSet> {
7115 let mut row_ids = Vec::new();
7116 let overlay_rids = self.overlay_rid_set(snapshot);
7117 for rr in &self.run_refs {
7118 let mut reader = self.open_reader(rr.run_id)?;
7119 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
7120 for rid in matched {
7121 if !overlay_rids.contains(&rid) {
7122 row_ids.push(rid);
7123 }
7124 }
7125 }
7126 let mut s = RowIdSet::from_unsorted(row_ids);
7127 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
7128 Ok(s)
7129 }
7130
7131 fn range_scan_f64(
7134 &self,
7135 column_id: u16,
7136 lo: f64,
7137 lo_inclusive: bool,
7138 hi: f64,
7139 hi_inclusive: bool,
7140 snapshot: Snapshot,
7141 ) -> Result<RowIdSet> {
7142 let mut row_ids = Vec::new();
7143 let overlay_rids = self.overlay_rid_set(snapshot);
7144 for rr in &self.run_refs {
7145 let mut reader = self.open_reader(rr.run_id)?;
7146 let matched = reader.range_row_ids_visible_f64(
7147 column_id,
7148 lo,
7149 lo_inclusive,
7150 hi,
7151 hi_inclusive,
7152 snapshot.epoch,
7153 )?;
7154 for rid in matched {
7155 if !overlay_rids.contains(&rid) {
7156 row_ids.push(rid);
7157 }
7158 }
7159 }
7160 let mut s = RowIdSet::from_unsorted(row_ids);
7161 self.range_scan_overlay_f64(
7162 &mut s,
7163 column_id,
7164 lo,
7165 lo_inclusive,
7166 hi,
7167 hi_inclusive,
7168 snapshot,
7169 );
7170 Ok(s)
7171 }
7172
7173 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
7175 let mut s = HashSet::new();
7176 for row in self.memtable.visible_versions(snapshot.epoch) {
7177 s.insert(row.row_id.0);
7178 }
7179 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7180 s.insert(row.row_id.0);
7181 }
7182 s
7183 }
7184
7185 fn range_scan_overlay_i64(
7186 &self,
7187 s: &mut RowIdSet,
7188 column_id: u16,
7189 lo: i64,
7190 hi: i64,
7191 snapshot: Snapshot,
7192 ) {
7193 let mut newest: HashMap<u64, &Row> = HashMap::new();
7198 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7199 let memtable = self.memtable.visible_versions(snapshot.epoch);
7200 for r in &mutable {
7201 newest.entry(r.row_id.0).or_insert(r);
7202 }
7203 for r in &memtable {
7204 newest.insert(r.row_id.0, r);
7205 }
7206 for row in newest.values() {
7207 if !row.deleted {
7208 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
7209 if *v >= lo && *v <= hi {
7210 s.insert(row.row_id.0);
7211 }
7212 }
7213 }
7214 }
7215 }
7216
7217 #[allow(clippy::too_many_arguments)]
7218 fn range_scan_overlay_f64(
7219 &self,
7220 s: &mut RowIdSet,
7221 column_id: u16,
7222 lo: f64,
7223 lo_inclusive: bool,
7224 hi: f64,
7225 hi_inclusive: bool,
7226 snapshot: Snapshot,
7227 ) {
7228 let mut newest: HashMap<u64, &Row> = HashMap::new();
7231 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7232 let memtable = self.memtable.visible_versions(snapshot.epoch);
7233 for r in &mutable {
7234 newest.entry(r.row_id.0).or_insert(r);
7235 }
7236 for r in &memtable {
7237 newest.insert(r.row_id.0, r);
7238 }
7239 for row in newest.values() {
7240 if !row.deleted {
7241 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
7242 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
7243 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
7244 if ok_lo && ok_hi {
7245 s.insert(row.row_id.0);
7246 }
7247 }
7248 }
7249 }
7250 }
7251
7252 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
7255 let mut row_ids = Vec::new();
7256 let overlay_rids = self.overlay_rid_set(snapshot);
7257 for rr in &self.run_refs {
7258 let mut reader = self.open_reader(rr.run_id)?;
7259 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
7260 for rid in matched {
7261 if !overlay_rids.contains(&rid) {
7262 row_ids.push(rid);
7263 }
7264 }
7265 }
7266 let mut s = RowIdSet::from_unsorted(row_ids);
7267 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
7268 Ok(s)
7269 }
7270
7271 fn null_scan_overlay(
7275 &self,
7276 s: &mut RowIdSet,
7277 column_id: u16,
7278 want_nulls: bool,
7279 snapshot: Snapshot,
7280 ) {
7281 let mut newest: HashMap<u64, &Row> = HashMap::new();
7282 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
7283 let memtable = self.memtable.visible_versions(snapshot.epoch);
7284 for r in &mutable {
7285 newest.entry(r.row_id.0).or_insert(r);
7286 }
7287 for r in &memtable {
7288 newest.insert(r.row_id.0, r);
7289 }
7290 for row in newest.values() {
7291 if row.deleted {
7292 continue;
7293 }
7294 let is_null = !row.columns.contains_key(&column_id)
7295 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
7296 if is_null == want_nulls {
7297 s.insert(row.row_id.0);
7298 }
7299 }
7300 }
7301
7302 pub fn snapshot(&self) -> Snapshot {
7303 Snapshot::at(self.epoch.visible())
7304 }
7305
7306 pub fn data_generation(&self) -> u64 {
7308 self.data_generation
7309 }
7310
7311 pub(crate) fn bump_data_generation(&mut self) {
7312 self.data_generation = self.data_generation.wrapping_add(1);
7313 }
7314
7315 pub(crate) fn table_id(&self) -> u64 {
7316 self.table_id
7317 }
7318
7319 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
7320 self.ensure_indexes_complete()?;
7321 self.memtable.seal();
7322 self.mutable_run.seal();
7323 self.hot.seal();
7324 for index in self.bitmap.values_mut() {
7325 index.seal();
7326 }
7327 for index in self.ann.values_mut() {
7328 index.seal();
7329 }
7330 for index in self.fm.values_mut() {
7331 index.seal();
7332 }
7333 for index in self.sparse.values_mut() {
7334 index.seal();
7335 }
7336 for index in self.minhash.values_mut() {
7337 index.seal();
7338 }
7339 self.pk_by_row.seal();
7340 let mut generation = self.clone();
7341 generation.read_only = true;
7342 generation.wal = WalSink::ReadOnly;
7343 generation.pending_delete_rids.clear();
7344 generation.pending_put_cols.clear();
7345 generation.pending_rows.clear();
7346 generation.pending_rows_auto_inc.clear();
7347 generation.pending_dels.clear();
7348 generation.pending_truncate = None;
7349 generation.agg_cache = Arc::new(HashMap::new());
7350 Ok(generation)
7351 }
7352
7353 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
7354 (std::mem::size_of::<Self>() as u64)
7355 .saturating_add(self.memtable.approx_bytes())
7356 .saturating_add(self.mutable_run.approx_bytes())
7357 .saturating_add(self.live_count.saturating_mul(64))
7358 }
7359
7360 pub fn pin_snapshot(&mut self) -> Snapshot {
7363 let e = self.epoch.visible();
7364 *self.pinned.entry(e).or_insert(0) += 1;
7365 Snapshot::at(e)
7366 }
7367
7368 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
7370 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
7371 *count -= 1;
7372 if *count == 0 {
7373 self.pinned.remove(&snap.epoch);
7374 }
7375 }
7376 }
7377
7378 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
7388 let local = self.pinned.keys().next().copied();
7389 let global = self.snapshots.min_pinned();
7390 let history = self.snapshots.history_floor(self.current_epoch());
7391 [local, global, history].into_iter().flatten().min()
7392 }
7393
7394 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
7398 self.ensure_writable()?;
7399 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
7400 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
7401 }
7402
7403 pub fn clear_ttl(&mut self) -> Result<()> {
7404 self.ensure_writable()?;
7405 self.apply_ttl_policy_at(None, self.current_epoch())
7406 }
7407
7408 pub fn ttl(&self) -> Option<TtlPolicy> {
7409 self.ttl
7410 }
7411
7412 pub(crate) fn prepare_ttl_policy(
7413 &self,
7414 column_name: &str,
7415 duration_nanos: u64,
7416 ) -> Result<TtlPolicy> {
7417 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
7418 return Err(MongrelError::InvalidArgument(
7419 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
7420 ));
7421 }
7422 let column = self
7423 .schema
7424 .columns
7425 .iter()
7426 .find(|column| column.name == column_name)
7427 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
7428 if column.ty != TypeId::TimestampNanos {
7429 return Err(MongrelError::Schema(format!(
7430 "TTL column {column_name} must be TimestampNanos, is {:?}",
7431 column.ty
7432 )));
7433 }
7434 Ok(TtlPolicy {
7435 column_id: column.id,
7436 duration_nanos,
7437 })
7438 }
7439
7440 pub(crate) fn apply_ttl_policy_at(
7441 &mut self,
7442 policy: Option<TtlPolicy>,
7443 epoch: Epoch,
7444 ) -> Result<()> {
7445 if let Some(policy) = policy {
7446 let column = self
7447 .schema
7448 .columns
7449 .iter()
7450 .find(|column| column.id == policy.column_id)
7451 .ok_or_else(|| {
7452 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
7453 })?;
7454 if column.ty != TypeId::TimestampNanos
7455 || policy.duration_nanos == 0
7456 || policy.duration_nanos > i64::MAX as u64
7457 {
7458 return Err(MongrelError::Schema("invalid TTL policy".into()));
7459 }
7460 }
7461 self.ttl = policy;
7462 self.agg_cache = Arc::new(HashMap::new());
7463 self.clear_result_cache();
7464 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
7465 self.persist_manifest(epoch)
7466 }
7467
7468 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
7469 let Some(policy) = self.ttl else {
7470 return false;
7471 };
7472 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
7473 return false;
7474 };
7475 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
7476 }
7477
7478 pub fn current_epoch(&self) -> Epoch {
7479 self.epoch.visible()
7480 }
7481
7482 pub fn memtable_len(&self) -> usize {
7483 self.memtable.len()
7484 }
7485
7486 pub fn count(&self) -> u64 {
7489 if self.ttl.is_none()
7490 && self.pending_put_cols.is_empty()
7491 && self.pending_delete_rids.is_empty()
7492 && self.pending_rows.is_empty()
7493 && self.pending_dels.is_empty()
7494 && self.pending_truncate.is_none()
7495 {
7496 self.live_count
7497 } else {
7498 self.visible_rows(self.snapshot())
7499 .map(|rows| rows.len() as u64)
7500 .unwrap_or(self.live_count)
7501 }
7502 }
7503
7504 pub fn count_conditions(
7508 &mut self,
7509 conditions: &[crate::query::Condition],
7510 snapshot: Snapshot,
7511 ) -> Result<Option<u64>> {
7512 use crate::query::Condition;
7513 if self.ttl.is_some() {
7514 if conditions.is_empty() {
7515 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
7516 }
7517 let mut sets = Vec::with_capacity(conditions.len());
7518 for condition in conditions {
7519 sets.push(self.resolve_condition(condition, snapshot)?);
7520 }
7521 let survivors = RowIdSet::intersect_many(sets);
7522 let rows = self.visible_rows(snapshot)?;
7523 return Ok(Some(
7524 rows.into_iter()
7525 .filter(|row| survivors.contains(row.row_id.0))
7526 .count() as u64,
7527 ));
7528 }
7529 if conditions.is_empty() {
7530 return Ok(Some(self.count()));
7531 }
7532 let served = |c: &Condition| {
7533 matches!(
7534 c,
7535 Condition::Pk(_)
7536 | Condition::BitmapEq { .. }
7537 | Condition::BitmapIn { .. }
7538 | Condition::BytesPrefix { .. }
7539 | Condition::FmContains { .. }
7540 | Condition::FmContainsAll { .. }
7541 | Condition::Ann { .. }
7542 | Condition::Range { .. }
7543 | Condition::RangeF64 { .. }
7544 | Condition::SparseMatch { .. }
7545 | Condition::MinHashSimilar { .. }
7546 | Condition::IsNull { .. }
7547 | Condition::IsNotNull { .. }
7548 )
7549 };
7550 if !conditions.iter().all(served) {
7551 return Ok(None);
7552 }
7553 self.ensure_indexes_complete()?;
7554 if !self.pending_put_cols.is_empty()
7555 || !self.pending_delete_rids.is_empty()
7556 || !self.pending_rows.is_empty()
7557 || !self.pending_dels.is_empty()
7558 || self.pending_truncate.is_some()
7559 {
7560 let mut sets = Vec::with_capacity(conditions.len());
7561 for condition in conditions {
7562 sets.push(self.resolve_condition(condition, snapshot)?);
7563 }
7564 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7565 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
7566 }
7567 let mut sets = Vec::with_capacity(conditions.len());
7568 for condition in conditions {
7569 sets.push(self.resolve_condition(condition, snapshot)?);
7570 }
7571 let mut rids = RowIdSet::intersect_many(sets);
7572 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7582 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
7583 }
7584 let count = rids.len() as u64;
7585 crate::trace::QueryTrace::record(|t| {
7586 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
7587 t.survivor_count = Some(count as usize);
7588 t.conditions_pushed = conditions.len();
7589 });
7590 Ok(Some(count))
7591 }
7592
7593 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
7598 let mut out = Vec::new();
7599 for row in self.memtable.visible_versions(snapshot.epoch) {
7600 if row.deleted {
7601 out.push(row.row_id.0);
7602 }
7603 }
7604 for row in self.mutable_run.visible_versions(snapshot.epoch) {
7605 if row.deleted {
7606 out.push(row.row_id.0);
7607 }
7608 }
7609 out
7610 }
7611
7612 pub fn bulk_load_columns(
7621 &mut self,
7622 user_columns: Vec<(u16, columnar::NativeColumn)>,
7623 ) -> Result<Epoch> {
7624 self.bulk_load_columns_with(user_columns, 3, false, true)
7625 }
7626
7627 pub fn bulk_load_fast(
7634 &mut self,
7635 user_columns: Vec<(u16, columnar::NativeColumn)>,
7636 ) -> Result<Epoch> {
7637 self.bulk_load_columns_with(user_columns, -1, true, false)
7638 }
7639
7640 fn bulk_load_columns_with(
7641 &mut self,
7642 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
7643 zstd_level: i32,
7644 force_plain: bool,
7645 lz4: bool,
7646 ) -> Result<Epoch> {
7647 self.ensure_writable()?;
7648 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
7649 if n == 0 {
7650 return Ok(self.current_epoch());
7651 }
7652 let epoch = self.commit_new_epoch()?;
7653 let live_before = self.live_count;
7654 self.spill_mutable_run(epoch)?;
7656 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
7657 && self.indexes_complete
7658 && self.run_refs.is_empty()
7659 && self.memtable.is_empty()
7660 && self.mutable_run.is_empty();
7661 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
7664 self.validate_columns_not_null(&user_columns, n)?;
7665 let winner_idx = self
7666 .bulk_pk_winner_indices(&user_columns, n)
7667 .filter(|idx| idx.len() != n);
7668 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
7669 match winner_idx.as_deref() {
7670 Some(idx) => {
7671 let compacted = user_columns
7672 .iter()
7673 .map(|(id, c)| (*id, c.gather(idx)))
7674 .collect();
7675 (compacted, idx.len())
7676 }
7677 None => (user_columns, n),
7678 };
7679 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
7680 let first = self.allocator.alloc_range(write_n as u64)?.0;
7681 for rid in first..first + write_n as u64 {
7682 self.reservoir.offer(rid);
7683 }
7684 let run_id = self.alloc_run_id()?;
7685 let path = self.run_path(run_id);
7686 let mut writer =
7687 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
7688 if force_plain {
7689 writer = writer.with_plain();
7690 } else if lz4 {
7691 writer = writer.with_lz4();
7694 } else {
7695 writer = writer.with_zstd_level(zstd_level);
7696 }
7697 if let Some(kek) = &self.kek {
7698 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
7699 }
7700 let header = match self.create_run_file(run_id)? {
7701 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
7702 None => writer.write_native(&path, &write_columns, write_n, first)?,
7703 };
7704 self.run_refs.push(RunRef {
7705 run_id: run_id as u128,
7706 level: 0,
7707 epoch_created: epoch.0,
7708 row_count: header.row_count,
7709 });
7710 self.live_count = self.live_count.saturating_add(write_n as u64);
7711 if eager_index_build {
7712 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
7713 self.index_columns_bulk(&write_columns, &row_ids);
7714 self.indexes_complete = true;
7715 self.build_learned_ranges()?;
7716 } else {
7717 self.indexes_complete = false;
7721 }
7722 self.mark_flushed(epoch)?;
7723 self.persist_manifest(epoch)?;
7724 if eager_index_build {
7725 self.checkpoint_indexes(epoch);
7726 }
7727 self.clear_result_cache();
7728 self.data_generation = self.data_generation.wrapping_add(1);
7729 Ok(epoch)
7730 }
7731
7732 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
7750 let n = row_ids.len();
7751 if n == 0 {
7752 return;
7753 }
7754 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
7755 columns.iter().map(|(id, c)| (*id, c)).collect();
7756 let ty_of: std::collections::HashMap<u16, TypeId> = self
7757 .schema
7758 .columns
7759 .iter()
7760 .map(|c| (c.id, c.ty.clone()))
7761 .collect();
7762 let pk_id = self.schema.primary_key().map(|c| c.id);
7763
7764 for (i, &rid) in row_ids.iter().enumerate() {
7765 let row_id = RowId(rid);
7766 if let Some(pid) = pk_id {
7767 if let Some(col) = by_id.get(&pid) {
7768 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
7769 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
7770 self.insert_hot_pk(key, row_id);
7771 }
7772 }
7773 }
7774 for idef in &self.schema.indexes {
7775 let Some(col) = by_id.get(&idef.column_id) else {
7776 continue;
7777 };
7778 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
7779 match idef.kind {
7780 IndexKind::Bitmap => {
7781 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
7782 if let Some(key) =
7783 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
7784 {
7785 b.insert(key, row_id);
7786 }
7787 }
7788 }
7789 IndexKind::FmIndex => {
7790 if let Some(f) = self.fm.get_mut(&idef.column_id) {
7791 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7792 f.insert(bytes.to_vec(), row_id);
7793 }
7794 }
7795 }
7796 IndexKind::Sparse => {
7797 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
7798 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7799 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
7800 s.insert(&terms, row_id);
7801 }
7802 }
7803 }
7804 }
7805 IndexKind::MinHash => {
7806 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
7807 if let Some(bytes) = columnar::native_bytes_at(col, i) {
7808 let tokens = crate::index::token_hashes_from_bytes(bytes);
7809 mh.insert(&tokens, row_id);
7810 }
7811 }
7812 }
7813 _ => {}
7814 }
7815 }
7816 }
7817 }
7818
7819 pub fn visible_columns_native(
7824 &self,
7825 snapshot: Snapshot,
7826 projection: Option<&[u16]>,
7827 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7828 self.visible_columns_native_inner(snapshot, projection, None)
7829 }
7830
7831 pub fn visible_columns_native_with_control(
7832 &self,
7833 snapshot: Snapshot,
7834 projection: Option<&[u16]>,
7835 control: &crate::ExecutionControl,
7836 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7837 self.visible_columns_native_inner(snapshot, projection, Some(control))
7838 }
7839
7840 fn visible_columns_native_inner(
7841 &self,
7842 snapshot: Snapshot,
7843 projection: Option<&[u16]>,
7844 control: Option<&crate::ExecutionControl>,
7845 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
7846 execution_checkpoint(control, 0)?;
7847 let wanted: Vec<u16> = match projection {
7848 Some(p) => p.to_vec(),
7849 None => self.schema.columns.iter().map(|c| c.id).collect(),
7850 };
7851 if self.ttl.is_none()
7852 && self.memtable.is_empty()
7853 && self.mutable_run.is_empty()
7854 && self.run_refs.len() == 1
7855 {
7856 let rr = self.run_refs[0].clone();
7857 let mut reader = self.open_reader(rr.run_id)?;
7858 let idxs = reader.visible_indices_native(snapshot.epoch)?;
7859 execution_checkpoint(control, 0)?;
7860 let all_visible = idxs.len() == reader.row_count();
7861 if reader.has_mmap() && control.is_none() {
7867 use rayon::prelude::*;
7868 let valid: Vec<u16> = wanted
7871 .iter()
7872 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
7873 .copied()
7874 .collect();
7875 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
7877 .par_iter()
7878 .filter_map(|cid| {
7879 reader
7880 .column_native_shared(*cid)
7881 .ok()
7882 .map(|col| (*cid, col))
7883 })
7884 .collect();
7885 let cols = decoded
7886 .into_iter()
7887 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
7888 .collect();
7889 return Ok(cols);
7890 }
7891 let mut cols = Vec::with_capacity(wanted.len());
7892 for (index, cid) in wanted.iter().enumerate() {
7893 execution_checkpoint(control, index)?;
7894 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
7895 Some(c) => c,
7896 None => continue,
7897 };
7898 let col = reader.column_native(cdef.id)?;
7899 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
7900 }
7901 return Ok(cols);
7902 }
7903 let vcols = self.visible_columns(snapshot)?;
7904 execution_checkpoint(control, 0)?;
7905 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
7906 let out: Vec<(u16, columnar::NativeColumn)> = vcols
7907 .into_iter()
7908 .filter(|(id, _)| want_set.contains(id))
7909 .map(|(id, vals)| {
7910 let ty = self
7911 .schema
7912 .columns
7913 .iter()
7914 .find(|c| c.id == id)
7915 .map(|c| c.ty.clone())
7916 .unwrap_or(TypeId::Bytes);
7917 (id, columnar::values_to_native(ty, &vals))
7918 })
7919 .collect();
7920 Ok(out)
7921 }
7922
7923 pub fn run_count(&self) -> usize {
7924 self.run_refs.len()
7925 }
7926
7927 pub fn memtable_is_empty(&self) -> bool {
7929 self.memtable.is_empty()
7930 }
7931
7932 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
7936 self.page_cache.stats()
7937 }
7938
7939 pub fn reset_page_cache_stats(&self) {
7941 self.page_cache.reset_stats();
7942 }
7943
7944 pub fn run_ids(&self) -> Vec<u128> {
7947 self.run_refs.iter().map(|r| r.run_id).collect()
7948 }
7949
7950 pub fn single_run_is_clean(&self) -> bool {
7954 if self.ttl.is_some() || self.run_refs.len() != 1 {
7955 return false;
7956 }
7957 self.open_reader(self.run_refs[0].run_id)
7958 .map(|r| r.is_clean())
7959 .unwrap_or(false)
7960 }
7961
7962 fn resolve_footprint(
7969 &self,
7970 conditions: &[crate::query::Condition],
7971 snapshot: Snapshot,
7972 ) -> roaring::RoaringBitmap {
7973 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
7974 return roaring::RoaringBitmap::new();
7975 }
7976 if self.run_refs.is_empty() {
7977 return roaring::RoaringBitmap::new();
7978 }
7979 if self.run_refs.len() == 1 {
7981 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
7982 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
7983 return rids.to_roaring_lossy();
7984 }
7985 }
7986 }
7987 roaring::RoaringBitmap::new()
7988 }
7989
7990 pub fn query_columns_native_cached(
8001 &mut self,
8002 conditions: &[crate::query::Condition],
8003 projection: Option<&[u16]>,
8004 snapshot: Snapshot,
8005 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8006 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
8007 }
8008
8009 pub fn query_columns_native_cached_with_control(
8010 &mut self,
8011 conditions: &[crate::query::Condition],
8012 projection: Option<&[u16]>,
8013 snapshot: Snapshot,
8014 control: &crate::ExecutionControl,
8015 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8016 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
8017 }
8018
8019 fn query_columns_native_cached_inner(
8020 &mut self,
8021 conditions: &[crate::query::Condition],
8022 projection: Option<&[u16]>,
8023 snapshot: Snapshot,
8024 control: Option<&crate::ExecutionControl>,
8025 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8026 execution_checkpoint(control, 0)?;
8027 if self.ttl.is_some() {
8030 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8031 }
8032 if conditions.is_empty() {
8033 return self.query_columns_native_inner(conditions, projection, snapshot, control);
8034 }
8035 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
8039 if let Some(hit) = self.result_cache.lock().get_columns(key) {
8040 crate::trace::QueryTrace::record(|t| {
8041 t.result_cache_hit = true;
8042 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8043 });
8044 return Ok(Some((*hit).clone()));
8045 }
8046 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
8047 execution_checkpoint(control, 0)?;
8048 if let Some(cols) = &res {
8049 let footprint = self.resolve_footprint(conditions, snapshot);
8050 let condition_cols = crate::query::condition_columns(conditions);
8051 execution_checkpoint(control, 0)?;
8052 self.result_cache.lock().insert(
8053 key,
8054 CachedEntry {
8055 data: CachedData::Columns(Arc::new(cols.clone())),
8056 footprint,
8057 condition_cols,
8058 },
8059 );
8060 }
8061 Ok(res)
8062 }
8063
8064 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
8069 if self.ttl.is_some() {
8070 return self.query(q);
8071 }
8072 if q.conditions.is_empty() {
8073 return self.query(q);
8074 }
8075 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
8076 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
8077 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
8078 if let Some(hit) = self.result_cache.lock().get_rows(key) {
8079 crate::trace::QueryTrace::record(|t| {
8080 t.result_cache_hit = true;
8081 t.scan_mode = crate::trace::ScanMode::Materialized;
8082 });
8083 return Ok((*hit).clone());
8084 }
8085 let rows = self.query(q)?;
8086 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
8087 let condition_cols = crate::query::condition_columns(&q.conditions);
8088 self.result_cache.lock().insert(
8089 key,
8090 CachedEntry {
8091 data: CachedData::Rows(Arc::new(rows.clone())),
8092 footprint,
8093 condition_cols,
8094 },
8095 );
8096 Ok(rows)
8097 }
8098
8099 #[allow(clippy::type_complexity)]
8114 pub fn query_columns_native_traced(
8115 &mut self,
8116 conditions: &[crate::query::Condition],
8117 projection: Option<&[u16]>,
8118 snapshot: Snapshot,
8119 ) -> Result<(
8120 Option<Vec<(u16, columnar::NativeColumn)>>,
8121 crate::trace::QueryTrace,
8122 )> {
8123 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8124 self.query_columns_native(conditions, projection, snapshot)
8125 });
8126 Ok((result?, trace))
8127 }
8128
8129 #[allow(clippy::type_complexity)]
8132 pub fn query_columns_native_cached_traced(
8133 &mut self,
8134 conditions: &[crate::query::Condition],
8135 projection: Option<&[u16]>,
8136 snapshot: Snapshot,
8137 ) -> Result<(
8138 Option<Vec<(u16, columnar::NativeColumn)>>,
8139 crate::trace::QueryTrace,
8140 )> {
8141 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8142 self.query_columns_native_cached(conditions, projection, snapshot)
8143 });
8144 Ok((result?, trace))
8145 }
8146
8147 pub fn native_page_cursor_traced(
8149 &self,
8150 snapshot: Snapshot,
8151 projection: Vec<(u16, TypeId)>,
8152 conditions: &[crate::query::Condition],
8153 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
8154 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8155 self.native_page_cursor(snapshot, projection, conditions)
8156 });
8157 Ok((result?, trace))
8158 }
8159
8160 pub fn native_multi_run_cursor_traced(
8162 &self,
8163 snapshot: Snapshot,
8164 projection: Vec<(u16, TypeId)>,
8165 conditions: &[crate::query::Condition],
8166 ) -> Result<(
8167 Option<crate::cursor::MultiRunCursor>,
8168 crate::trace::QueryTrace,
8169 )> {
8170 let (result, trace) = crate::trace::QueryTrace::capture(|| {
8171 self.native_multi_run_cursor(snapshot, projection, conditions)
8172 });
8173 Ok((result?, trace))
8174 }
8175
8176 pub fn count_conditions_traced(
8178 &mut self,
8179 conditions: &[crate::query::Condition],
8180 snapshot: Snapshot,
8181 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
8182 let (result, trace) =
8183 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
8184 Ok((result?, trace))
8185 }
8186
8187 pub fn query_traced(
8189 &mut self,
8190 q: &crate::query::Query,
8191 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
8192 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
8193 Ok((result?, trace))
8194 }
8195
8196 pub fn query_columns_native(
8201 &mut self,
8202 conditions: &[crate::query::Condition],
8203 projection: Option<&[u16]>,
8204 snapshot: Snapshot,
8205 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8206 self.query_columns_native_inner(conditions, projection, snapshot, None)
8207 }
8208
8209 pub fn query_columns_native_with_control(
8210 &mut self,
8211 conditions: &[crate::query::Condition],
8212 projection: Option<&[u16]>,
8213 snapshot: Snapshot,
8214 control: &crate::ExecutionControl,
8215 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8216 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
8217 }
8218
8219 fn query_columns_native_inner(
8220 &mut self,
8221 conditions: &[crate::query::Condition],
8222 projection: Option<&[u16]>,
8223 snapshot: Snapshot,
8224 control: Option<&crate::ExecutionControl>,
8225 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
8226 use crate::query::Condition;
8227 execution_checkpoint(control, 0)?;
8228 if self.ttl.is_some() {
8231 return Ok(None);
8232 }
8233 if conditions.is_empty() {
8234 return Ok(None);
8235 }
8236 self.ensure_indexes_complete()?;
8237
8238 let served = |c: &Condition| {
8243 matches!(
8244 c,
8245 Condition::Pk(_)
8246 | Condition::BitmapEq { .. }
8247 | Condition::BitmapIn { .. }
8248 | Condition::BytesPrefix { .. }
8249 | Condition::FmContains { .. }
8250 | Condition::FmContainsAll { .. }
8251 | Condition::Ann { .. }
8252 | Condition::Range { .. }
8253 | Condition::RangeF64 { .. }
8254 | Condition::SparseMatch { .. }
8255 | Condition::MinHashSimilar { .. }
8256 | Condition::IsNull { .. }
8257 | Condition::IsNotNull { .. }
8258 )
8259 };
8260 if !conditions.iter().all(served) {
8261 return Ok(None);
8262 }
8263 let fast_path =
8264 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
8265 crate::trace::QueryTrace::record(|t| {
8266 t.run_count = self.run_refs.len();
8267 t.memtable_rows = self.memtable.len();
8268 t.mutable_run_rows = self.mutable_run.len();
8269 t.conditions_pushed = conditions.len();
8270 t.learned_range_used = conditions.iter().any(|c| match c {
8271 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
8272 self.learned_range.contains_key(column_id)
8273 }
8274 _ => false,
8275 });
8276 });
8277 let col_ids: Vec<u16> = projection
8279 .map(|p| p.to_vec())
8280 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
8281 let proj_pairs: Vec<(u16, TypeId)> = col_ids
8282 .iter()
8283 .map(|&cid| {
8284 let ty = self
8285 .schema
8286 .columns
8287 .iter()
8288 .find(|c| c.id == cid)
8289 .map(|c| c.ty.clone())
8290 .unwrap_or(TypeId::Bytes);
8291 (cid, ty)
8292 })
8293 .collect();
8294
8295 if fast_path {
8301 let needs_column = conditions.iter().any(|c| match c {
8304 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
8305 Condition::RangeF64 { column_id, .. } => {
8306 !self.learned_range.contains_key(column_id)
8307 }
8308 _ => false,
8309 });
8310 let mut reader_opt: Option<RunReader> = if needs_column {
8311 Some(self.open_reader(self.run_refs[0].run_id)?)
8312 } else {
8313 None
8314 };
8315 let mut sets: Vec<RowIdSet> = Vec::new();
8316 for (index, c) in conditions.iter().enumerate() {
8317 execution_checkpoint(control, index)?;
8318 let s = match c {
8319 Condition::Range { column_id, lo, hi }
8320 if !self.learned_range.contains_key(column_id) =>
8321 {
8322 if reader_opt.is_none() {
8323 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8324 }
8325 reader_opt
8326 .as_mut()
8327 .expect("reader opened for range")
8328 .range_row_id_set_i64(*column_id, *lo, *hi)?
8329 }
8330 Condition::RangeF64 {
8331 column_id,
8332 lo,
8333 lo_inclusive,
8334 hi,
8335 hi_inclusive,
8336 } if !self.learned_range.contains_key(column_id) => {
8337 if reader_opt.is_none() {
8338 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
8339 }
8340 reader_opt
8341 .as_mut()
8342 .expect("reader opened for range")
8343 .range_row_id_set_f64(
8344 *column_id,
8345 *lo,
8346 *lo_inclusive,
8347 *hi,
8348 *hi_inclusive,
8349 )?
8350 }
8351 _ => self.resolve_condition(c, snapshot)?,
8352 };
8353 sets.push(s);
8354 }
8355 let candidates = RowIdSet::intersect_many(sets);
8356 crate::trace::QueryTrace::record(|t| {
8357 t.survivor_count = Some(candidates.len());
8358 });
8359 if candidates.is_empty() {
8360 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
8361 .iter()
8362 .map(|&id| {
8363 (
8364 id,
8365 columnar::null_native(
8366 proj_pairs
8367 .iter()
8368 .find(|(c, _)| c == &id)
8369 .map(|(_, t)| t.clone())
8370 .unwrap_or(TypeId::Bytes),
8371 0,
8372 ),
8373 )
8374 })
8375 .collect();
8376 return Ok(Some(cols));
8377 }
8378 let mut reader = match reader_opt.take() {
8379 Some(r) => r,
8380 None => self.open_reader(self.run_refs[0].run_id)?,
8381 };
8382 let candidate_ids = candidates.into_sorted_vec();
8383 let (positions, fast_rid) = if let Some(positions) =
8384 reader.positions_for_row_ids_fast(&candidate_ids)
8385 {
8386 (positions, true)
8387 } else {
8388 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
8389 match col {
8390 columnar::NativeColumn::Int64 { data, .. } => {
8391 let mut p = Vec::with_capacity(candidate_ids.len());
8392 for (index, rid) in candidate_ids.iter().enumerate() {
8393 execution_checkpoint(control, index)?;
8394 if let Ok(position) = data.binary_search(&(*rid as i64)) {
8395 p.push(position);
8396 }
8397 }
8398 p.sort_unstable();
8399 (p, false)
8400 }
8401 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
8402 }
8403 };
8404 crate::trace::QueryTrace::record(|t| {
8405 t.scan_mode = crate::trace::ScanMode::NativePushdown;
8406 t.fast_row_id_map = fast_rid;
8407 });
8408 let mut cols = Vec::with_capacity(col_ids.len());
8409 for (index, cid) in col_ids.iter().enumerate() {
8410 execution_checkpoint(control, index)?;
8411 let col = reader.column_native(*cid)?;
8412 cols.push((*cid, col.gather(&positions)));
8413 }
8414 return Ok(Some(cols));
8415 }
8416
8417 if !self.run_refs.is_empty() {
8430 use crate::cursor::{
8431 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
8432 };
8433 let remaining: usize;
8434 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
8435 let c = self
8436 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
8437 .expect("single-run cursor should build when run_refs.len() == 1");
8438 remaining = c.remaining_rows();
8439 Box::new(c)
8440 } else {
8441 let c = self
8442 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
8443 .expect("multi-run cursor should build when run_refs.len() >= 1");
8444 remaining = c.remaining_rows();
8445 Box::new(c)
8446 };
8447 crate::trace::QueryTrace::record(|t| {
8448 if t.survivor_count.is_none() {
8449 t.survivor_count = Some(remaining);
8450 }
8451 });
8452 let cols = match control {
8453 Some(control) => {
8454 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
8455 }
8456 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
8457 };
8458 return Ok(Some(cols));
8459 }
8460
8461 crate::trace::QueryTrace::record(|t| {
8466 t.scan_mode = crate::trace::ScanMode::Materialized;
8467 t.row_materialized = true;
8468 });
8469 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8470 for (index, c) in conditions.iter().enumerate() {
8471 execution_checkpoint(control, index)?;
8472 sets.push(self.resolve_condition(c, snapshot)?);
8473 }
8474 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
8475 let rows = self.rows_for_rids(&rids, snapshot)?;
8476 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
8477 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
8478 execution_checkpoint(control, index)?;
8479 let vals: Vec<Value> = rows
8480 .iter()
8481 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8482 .collect();
8483 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
8484 }
8485 Ok(Some(cols))
8486 }
8487
8488 pub fn native_page_cursor(
8503 &self,
8504 snapshot: Snapshot,
8505 projection: Vec<(u16, TypeId)>,
8506 conditions: &[crate::query::Condition],
8507 ) -> Result<Option<NativePageCursor>> {
8508 use crate::cursor::build_page_plans;
8509 if self.ttl.is_some() {
8510 return Ok(None);
8511 }
8512 if !conditions.is_empty() && !self.indexes_complete {
8515 return Ok(None);
8516 }
8517 if self.run_refs.len() != 1 {
8518 return Ok(None);
8519 }
8520 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8521 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8522
8523 let overlay_rids: HashSet<u64> = {
8526 let mut s = HashSet::new();
8527 for row in self.memtable.visible_versions(snapshot.epoch) {
8528 s.insert(row.row_id.0);
8529 }
8530 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8531 s.insert(row.row_id.0);
8532 }
8533 s
8534 };
8535
8536 let survivors = if conditions.is_empty() {
8540 None
8541 } else {
8542 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
8543 };
8544
8545 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
8552 survivors.clone()
8553 } else if let Some(s) = &survivors {
8554 let mut run_set = s.clone();
8555 run_set.remove_many(overlay_rids.iter().copied());
8556 Some(run_set)
8557 } else {
8558 Some(RowIdSet::from_unsorted(
8559 rids.iter()
8560 .map(|&r| r as u64)
8561 .filter(|r| !overlay_rids.contains(r))
8562 .collect(),
8563 ))
8564 };
8565
8566 let overlay_rows = if overlay_rids.is_empty() {
8567 Vec::new()
8568 } else {
8569 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8570 self.overlay_visible_rows(snapshot, bound)
8571 };
8572
8573 let plans = if positions.is_empty() {
8575 Vec::new()
8576 } else {
8577 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
8578 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
8579 };
8580
8581 let overlay = if overlay_rows.is_empty() {
8583 None
8584 } else {
8585 let filtered =
8586 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8587 if filtered.is_empty() {
8588 None
8589 } else {
8590 Some(self.materialize_overlay(&filtered, &projection))
8591 }
8592 };
8593
8594 let overlay_row_count = overlay
8595 .as_ref()
8596 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8597 .unwrap_or(0);
8598 crate::trace::QueryTrace::record(|t| {
8599 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
8600 t.run_count = self.run_refs.len();
8601 t.memtable_rows = self.memtable.len();
8602 t.mutable_run_rows = self.mutable_run.len();
8603 t.overlay_rows = overlay_row_count;
8604 t.conditions_pushed = conditions.len();
8605 t.pages_decoded = plans
8606 .iter()
8607 .map(|p| p.positions.len())
8608 .sum::<usize>()
8609 .min(1);
8610 });
8611
8612 Ok(Some(NativePageCursor::new_with_overlay(
8613 reader, projection, plans, overlay,
8614 )))
8615 }
8616 #[allow(clippy::type_complexity)]
8626 pub fn native_multi_run_cursor(
8627 &self,
8628 snapshot: Snapshot,
8629 projection: Vec<(u16, TypeId)>,
8630 conditions: &[crate::query::Condition],
8631 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
8632 use crate::cursor::{MultiRunCursor, RunStream};
8633 use crate::sorted_run::SYS_ROW_ID;
8634 use std::collections::{BinaryHeap, HashMap, HashSet};
8635 if self.ttl.is_some() {
8636 return Ok(None);
8637 }
8638 if !conditions.is_empty() && !self.indexes_complete {
8641 return Ok(None);
8642 }
8643 if self.run_refs.is_empty() {
8644 return Ok(None);
8645 }
8646
8647 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
8649 Vec::with_capacity(self.run_refs.len());
8650 for rr in &self.run_refs {
8651 let mut reader = self.open_reader(rr.run_id)?;
8652 let (rids, eps, del) = reader.system_columns_native()?;
8653 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
8654 run_meta.push((reader, rids, eps, del, page_rows));
8655 }
8656
8657 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
8661 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
8662 for i in 0..rids.len() {
8663 let rid = rids[i] as u64;
8664 let e = eps[i] as u64;
8665 if e > snapshot.epoch.0 {
8666 continue;
8667 }
8668 let is_del = del[i] != 0;
8669 best.entry(rid)
8670 .and_modify(|cur| {
8671 if e > cur.0 {
8672 *cur = (e, run_idx, i, is_del);
8673 }
8674 })
8675 .or_insert((e, run_idx, i, is_del));
8676 }
8677 }
8678
8679 let overlay_rids: HashSet<u64> = {
8681 let mut s = HashSet::new();
8682 for row in self.memtable.visible_versions(snapshot.epoch) {
8683 s.insert(row.row_id.0);
8684 }
8685 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8686 s.insert(row.row_id.0);
8687 }
8688 s
8689 };
8690
8691 let survivors: Option<RowIdSet> = if conditions.is_empty() {
8693 None
8694 } else {
8695 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8696 for c in conditions {
8697 sets.push(self.resolve_condition(c, snapshot)?);
8698 }
8699 Some(RowIdSet::intersect_many(sets))
8700 };
8701
8702 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
8706 for (rid, (_, run_idx, pos, deleted)) in &best {
8707 if *deleted {
8708 continue;
8709 }
8710 if overlay_rids.contains(rid) {
8711 continue;
8712 }
8713 if let Some(s) = &survivors {
8714 if !s.contains(*rid) {
8715 continue;
8716 }
8717 }
8718 per_run[*run_idx].push((*rid, *pos));
8719 }
8720 for v in per_run.iter_mut() {
8721 v.sort_unstable_by_key(|&(rid, _)| rid);
8722 }
8723
8724 let mut streams = Vec::with_capacity(run_meta.len());
8726 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
8727 let mut total = 0usize;
8728 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
8729 let mut starts = Vec::with_capacity(page_rows.len());
8730 let mut acc = 0usize;
8731 for &r in &page_rows {
8732 starts.push(acc);
8733 acc += r;
8734 }
8735 let mut survivors_vec: Vec<(u64, usize, usize)> =
8736 Vec::with_capacity(per_run[run_idx].len());
8737 for &(rid, pos) in &per_run[run_idx] {
8738 let page_seq = match starts.partition_point(|&s| s <= pos) {
8739 0 => continue,
8740 p => p - 1,
8741 };
8742 let within = pos - starts[page_seq];
8743 survivors_vec.push((rid, page_seq, within));
8744 }
8745 total += survivors_vec.len();
8746 if let Some(&(rid, _, _)) = survivors_vec.first() {
8747 heap.push(std::cmp::Reverse((rid, run_idx)));
8748 }
8749 streams.push(RunStream::new(reader, survivors_vec, page_rows));
8750 }
8751
8752 let overlay_rows = if overlay_rids.is_empty() {
8754 Vec::new()
8755 } else {
8756 let bound = Self::overlay_materialization_bound(conditions, &survivors);
8757 self.overlay_visible_rows(snapshot, bound)
8758 };
8759 let overlay = if overlay_rows.is_empty() {
8760 None
8761 } else {
8762 let filtered =
8763 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
8764 if filtered.is_empty() {
8765 None
8766 } else {
8767 Some(self.materialize_overlay(&filtered, &projection))
8768 }
8769 };
8770
8771 let overlay_row_count = overlay
8772 .as_ref()
8773 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
8774 .unwrap_or(0);
8775 crate::trace::QueryTrace::record(|t| {
8776 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
8777 t.run_count = self.run_refs.len();
8778 t.memtable_rows = self.memtable.len();
8779 t.mutable_run_rows = self.mutable_run.len();
8780 t.overlay_rows = overlay_row_count;
8781 t.conditions_pushed = conditions.len();
8782 t.survivor_count = Some(total);
8783 });
8784
8785 Ok(Some(MultiRunCursor::new(
8786 streams, projection, heap, total, overlay,
8787 )))
8788 }
8789
8790 fn overlay_materialization_bound<'a>(
8802 conditions: &[crate::query::Condition],
8803 survivors: &'a Option<RowIdSet>,
8804 ) -> Option<&'a RowIdSet> {
8805 use crate::query::Condition;
8806 let has_range = conditions
8807 .iter()
8808 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
8809 if has_range {
8810 None
8811 } else {
8812 survivors.as_ref()
8813 }
8814 }
8815
8816 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
8828 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
8829 let mut fold = |row: Row| {
8830 if let Some(b) = bound {
8831 if !b.contains(row.row_id.0) {
8832 return;
8833 }
8834 }
8835 best.entry(row.row_id.0)
8836 .and_modify(|(be, br)| {
8837 if row.committed_epoch > *be {
8838 *be = row.committed_epoch;
8839 *br = row.clone();
8840 }
8841 })
8842 .or_insert_with(|| (row.committed_epoch, row));
8843 };
8844 for row in self.memtable.visible_versions(snapshot.epoch) {
8845 fold(row);
8846 }
8847 for row in self.mutable_run.visible_versions(snapshot.epoch) {
8848 fold(row);
8849 }
8850 let mut out: Vec<Row> = best
8851 .into_values()
8852 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
8853 .collect();
8854 out.sort_by_key(|r| r.row_id);
8855 out
8856 }
8857
8858 fn filter_overlay_rows(
8866 &self,
8867 rows: Vec<Row>,
8868 conditions: &[crate::query::Condition],
8869 survivors: Option<&RowIdSet>,
8870 snapshot: Snapshot,
8871 ) -> Result<Vec<Row>> {
8872 if conditions.is_empty() {
8873 return Ok(rows);
8874 }
8875 use crate::query::Condition;
8876 let all_index_served = !conditions
8880 .iter()
8881 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
8882 if all_index_served {
8883 return Ok(rows
8884 .into_iter()
8885 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
8886 .collect());
8887 }
8888 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
8891 for c in conditions {
8892 let s = match c {
8893 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
8894 _ => self.resolve_condition(c, snapshot)?,
8895 };
8896 per_cond_sets.push(s);
8897 }
8898 Ok(rows
8899 .into_iter()
8900 .filter(|row| {
8901 conditions.iter().enumerate().all(|(i, c)| match c {
8902 Condition::Range { column_id, lo, hi } => {
8903 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
8904 }
8905 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
8906 match row.columns.get(column_id) {
8907 Some(Value::Float64(v)) => {
8908 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
8909 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
8910 lo_ok && hi_ok
8911 }
8912 _ => false,
8913 }
8914 }
8915 _ => per_cond_sets[i].contains(row.row_id.0),
8916 })
8917 })
8918 .collect())
8919 }
8920
8921 fn materialize_overlay(
8924 &self,
8925 rows: &[Row],
8926 projection: &[(u16, TypeId)],
8927 ) -> Vec<columnar::NativeColumn> {
8928 if projection.is_empty() {
8929 return vec![columnar::null_native(TypeId::Int64, rows.len())];
8930 }
8931 let mut cols = Vec::with_capacity(projection.len());
8932 for (cid, ty) in projection {
8933 let vals: Vec<Value> = rows
8934 .iter()
8935 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
8936 .collect();
8937 cols.push(columnar::values_to_native(ty.clone(), &vals));
8938 }
8939 cols
8940 }
8941
8942 fn resolve_survivor_rids(
8947 &self,
8948 conditions: &[crate::query::Condition],
8949 reader: &mut RunReader,
8950 snapshot: Snapshot,
8951 ) -> Result<RowIdSet> {
8952 use crate::query::Condition;
8953 let mut sets: Vec<RowIdSet> = Vec::new();
8954 for c in conditions {
8955 self.validate_condition(c)?;
8956 let s: RowIdSet = match c {
8957 Condition::Pk(key) => {
8958 let lookup = self
8959 .schema
8960 .primary_key()
8961 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
8962 .unwrap_or_else(|| key.clone());
8963 self.hot
8964 .get(&lookup)
8965 .map(|r| RowIdSet::one(r.0))
8966 .unwrap_or_else(RowIdSet::empty)
8967 }
8968 Condition::BitmapEq { column_id, value } => {
8969 let lookup = self.index_lookup_key_bytes(*column_id, value);
8970 self.bitmap
8971 .get(column_id)
8972 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
8973 .unwrap_or_else(RowIdSet::empty)
8974 }
8975 Condition::BitmapIn { column_id, values } => {
8976 let bm = self.bitmap.get(column_id);
8977 let mut acc = roaring::RoaringBitmap::new();
8978 if let Some(b) = bm {
8979 for v in values {
8980 let lookup = self.index_lookup_key_bytes(*column_id, v);
8981 acc |= b.get(&lookup);
8982 }
8983 }
8984 RowIdSet::from_roaring(acc)
8985 }
8986 Condition::BytesPrefix { column_id, prefix } => {
8987 if let Some(b) = self.bitmap.get(column_id) {
8988 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
8989 let mut acc = roaring::RoaringBitmap::new();
8990 for key in b.keys() {
8991 if key.starts_with(&lookup_prefix) {
8992 acc |= b.get(&key);
8993 }
8994 }
8995 RowIdSet::from_roaring(acc)
8996 } else {
8997 RowIdSet::empty()
8998 }
8999 }
9000 Condition::FmContains { column_id, pattern } => self
9001 .fm
9002 .get(column_id)
9003 .map(|f| {
9004 RowIdSet::from_unsorted(
9005 f.locate(pattern).into_iter().map(|r| r.0).collect(),
9006 )
9007 })
9008 .unwrap_or_else(RowIdSet::empty),
9009 Condition::FmContainsAll {
9010 column_id,
9011 patterns,
9012 } => {
9013 if let Some(f) = self.fm.get(column_id) {
9014 let sets: Vec<RowIdSet> = patterns
9015 .iter()
9016 .map(|pat| {
9017 RowIdSet::from_unsorted(
9018 f.locate(pat).into_iter().map(|r| r.0).collect(),
9019 )
9020 })
9021 .collect();
9022 RowIdSet::intersect_many(sets)
9023 } else {
9024 RowIdSet::empty()
9025 }
9026 }
9027 Condition::Ann {
9028 column_id,
9029 query,
9030 k,
9031 } => RowIdSet::from_unsorted(
9032 self.retrieve_filtered(
9033 &crate::query::Retriever::Ann {
9034 column_id: *column_id,
9035 query: query.clone(),
9036 k: *k,
9037 },
9038 snapshot,
9039 None,
9040 None,
9041 None,
9042 None,
9043 )?
9044 .into_iter()
9045 .map(|hit| hit.row_id.0)
9046 .collect(),
9047 ),
9048 Condition::SparseMatch {
9049 column_id,
9050 query,
9051 k,
9052 } => RowIdSet::from_unsorted(
9053 self.retrieve_filtered(
9054 &crate::query::Retriever::Sparse {
9055 column_id: *column_id,
9056 query: query.clone(),
9057 k: *k,
9058 },
9059 snapshot,
9060 None,
9061 None,
9062 None,
9063 None,
9064 )?
9065 .into_iter()
9066 .map(|hit| hit.row_id.0)
9067 .collect(),
9068 ),
9069 Condition::MinHashSimilar {
9070 column_id,
9071 query,
9072 k,
9073 } => match self.minhash.get(column_id) {
9074 Some(index) => {
9075 let candidates = index.candidate_row_ids(query);
9076 let eligible =
9077 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9078 RowIdSet::from_unsorted(
9079 index
9080 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
9081 .into_iter()
9082 .map(|(row_id, _)| row_id.0)
9083 .collect(),
9084 )
9085 }
9086 None => RowIdSet::empty(),
9087 },
9088 Condition::Range { column_id, lo, hi } => {
9089 if let Some(li) = self.learned_range.get(column_id) {
9090 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
9091 } else {
9092 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
9093 }
9094 }
9095 Condition::RangeF64 {
9096 column_id,
9097 lo,
9098 lo_inclusive,
9099 hi,
9100 hi_inclusive,
9101 } => {
9102 if let Some(li) = self.learned_range.get(column_id) {
9103 RowIdSet::from_unsorted(
9104 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9105 .into_iter()
9106 .collect(),
9107 )
9108 } else {
9109 reader.range_row_id_set_f64(
9110 *column_id,
9111 *lo,
9112 *lo_inclusive,
9113 *hi,
9114 *hi_inclusive,
9115 )?
9116 }
9117 }
9118 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
9119 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
9120 };
9121 sets.push(s);
9122 }
9123 Ok(RowIdSet::intersect_many(sets))
9124 }
9125
9126 pub fn scan_cursor(
9147 &self,
9148 snapshot: Snapshot,
9149 projection: Vec<(u16, TypeId)>,
9150 conditions: &[crate::query::Condition],
9151 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
9152 if self.ttl.is_some() {
9153 return Ok(None);
9154 }
9155 if !conditions.is_empty() && !self.indexes_complete {
9161 return Ok(None);
9162 }
9163 if self.run_refs.len() == 1 {
9164 Ok(self
9165 .native_page_cursor(snapshot, projection, conditions)?
9166 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9167 } else {
9168 Ok(self
9169 .native_multi_run_cursor(snapshot, projection, conditions)?
9170 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
9171 }
9172 }
9173
9174 pub fn aggregate_native(
9188 &self,
9189 snapshot: Snapshot,
9190 column: Option<u16>,
9191 conditions: &[crate::query::Condition],
9192 agg: NativeAgg,
9193 ) -> Result<Option<NativeAggResult>> {
9194 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
9195 }
9196
9197 pub fn aggregate_native_with_control(
9198 &self,
9199 snapshot: Snapshot,
9200 column: Option<u16>,
9201 conditions: &[crate::query::Condition],
9202 agg: NativeAgg,
9203 control: &crate::ExecutionControl,
9204 ) -> Result<Option<NativeAggResult>> {
9205 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
9206 }
9207
9208 fn aggregate_native_inner(
9209 &self,
9210 snapshot: Snapshot,
9211 column: Option<u16>,
9212 conditions: &[crate::query::Condition],
9213 agg: NativeAgg,
9214 control: Option<&crate::ExecutionControl>,
9215 ) -> Result<Option<NativeAggResult>> {
9216 execution_checkpoint(control, 0)?;
9217 if self.ttl.is_some() {
9218 return Ok(None);
9219 }
9220 if self.run_refs.len() == 1 && conditions.is_empty() {
9222 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
9223 return Ok(Some(res));
9224 }
9225 }
9226 if matches!(agg, NativeAgg::Count) && column.is_none() {
9228 return Ok(self
9229 .scan_cursor(snapshot, Vec::new(), conditions)?
9230 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
9231 }
9232 let cid = match column {
9235 Some(c) => c,
9236 None => return Ok(None),
9237 };
9238 let ty = self.column_type(cid);
9239 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
9240 else {
9241 return Ok(None);
9242 };
9243 execution_checkpoint(control, 0)?;
9244 match ty {
9245 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9246 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
9247 Ok(Some(pack_int(agg, count, sum, mn, mx)))
9248 }
9249 TypeId::Float64 => {
9250 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
9251 Ok(Some(pack_float(agg, count, sum, mn, mx)))
9252 }
9253 _ => Ok(None),
9254 }
9255 }
9256
9257 fn aggregate_from_stats(
9265 &self,
9266 snapshot: Snapshot,
9267 column: Option<u16>,
9268 agg: NativeAgg,
9269 ) -> Result<Option<NativeAggResult>> {
9270 let cid = match (agg, column) {
9271 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
9272 _ => return Ok(None), };
9274 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
9275 return Ok(None);
9276 };
9277 let Some(cs) = stats.get(&cid) else {
9278 return Ok(None);
9279 };
9280 match agg {
9281 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
9283 self.live_count.saturating_sub(cs.null_count),
9284 ))),
9285 NativeAgg::Min | NativeAgg::Max => {
9286 let bound = if agg == NativeAgg::Min {
9287 &cs.min
9288 } else {
9289 &cs.max
9290 };
9291 match bound {
9292 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
9293 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
9294 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
9299 None => Ok(None),
9300 }
9301 }
9302 _ => Ok(None),
9303 }
9304 }
9305
9306 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
9315 if self.ttl.is_some() {
9316 return Ok(None);
9317 }
9318 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
9319 return Ok(None);
9320 }
9321 self.ensure_indexes_complete()?;
9324 let reader = self.open_reader(self.run_refs[0].run_id)?;
9325 if self.live_count != reader.row_count() as u64 {
9326 return Ok(None);
9327 }
9328 let Some(bm) = self.bitmap.get(&column_id) else {
9329 return Ok(None); };
9331 let mut distinct = bm.value_count() as u64;
9332 if !bm.get(&Value::Null.encode_key()).is_empty() {
9335 distinct = distinct.saturating_sub(1);
9336 }
9337 Ok(Some(distinct))
9338 }
9339
9340 pub fn aggregate_incremental(
9352 &mut self,
9353 cache_key: u64,
9354 conditions: &[crate::query::Condition],
9355 column: Option<u16>,
9356 agg: NativeAgg,
9357 ) -> Result<IncrementalAggResult> {
9358 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
9359 }
9360
9361 pub fn aggregate_incremental_with_control(
9362 &mut self,
9363 cache_key: u64,
9364 conditions: &[crate::query::Condition],
9365 column: Option<u16>,
9366 agg: NativeAgg,
9367 control: &crate::ExecutionControl,
9368 ) -> Result<IncrementalAggResult> {
9369 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
9370 }
9371
9372 fn aggregate_incremental_inner(
9373 &mut self,
9374 cache_key: u64,
9375 conditions: &[crate::query::Condition],
9376 column: Option<u16>,
9377 agg: NativeAgg,
9378 control: Option<&crate::ExecutionControl>,
9379 ) -> Result<IncrementalAggResult> {
9380 execution_checkpoint(control, 0)?;
9381 let snap = self.snapshot();
9382 let cur_wm = self.allocator.current().0;
9383 let cur_epoch = snap.epoch.0;
9384 let incremental_ok = self.ttl.is_none()
9391 && !self.had_deletes
9392 && self.memtable.is_empty()
9393 && self.mutable_run.is_empty();
9394
9395 if incremental_ok {
9398 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
9399 if cached.epoch == cur_epoch {
9400 return Ok(IncrementalAggResult {
9401 state: cached.state,
9402 incremental: true,
9403 delta_rows: 0,
9404 });
9405 }
9406 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
9407 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
9408 let mut delta_rids = Vec::with_capacity(delta_len);
9409 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
9410 execution_checkpoint(control, index)?;
9411 delta_rids.push(row_id);
9412 }
9413 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
9414 execution_checkpoint(control, 0)?;
9415 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9416 let delta_state = agg_state_from_rows(
9417 &delta_rows,
9418 conditions,
9419 &index_sets,
9420 column,
9421 agg,
9422 &self.schema,
9423 control,
9424 )?;
9425 let merged = cached.state.merge(delta_state);
9426 let delta_n = delta_rids.len() as u64;
9427 Arc::make_mut(&mut self.agg_cache).insert(
9428 cache_key,
9429 CachedAgg {
9430 state: merged.clone(),
9431 watermark: cur_wm,
9432 epoch: cur_epoch,
9433 },
9434 );
9435 return Ok(IncrementalAggResult {
9436 state: merged,
9437 incremental: true,
9438 delta_rows: delta_n,
9439 });
9440 }
9441 }
9442 }
9443
9444 let cursor_ok =
9449 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
9450 let state = if cursor_ok && agg != NativeAgg::Avg {
9451 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
9452 Some(result) => {
9453 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
9454 }
9455 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
9456 }
9457 } else {
9458 self.agg_state_full_scan(conditions, column, agg, snap, control)?
9459 };
9460 if incremental_ok {
9462 Arc::make_mut(&mut self.agg_cache).insert(
9463 cache_key,
9464 CachedAgg {
9465 state: state.clone(),
9466 watermark: cur_wm,
9467 epoch: cur_epoch,
9468 },
9469 );
9470 }
9471 Ok(IncrementalAggResult {
9472 state,
9473 incremental: false,
9474 delta_rows: 0,
9475 })
9476 }
9477
9478 fn agg_state_full_scan(
9481 &self,
9482 conditions: &[crate::query::Condition],
9483 column: Option<u16>,
9484 agg: NativeAgg,
9485 snap: Snapshot,
9486 control: Option<&crate::ExecutionControl>,
9487 ) -> Result<AggState> {
9488 execution_checkpoint(control, 0)?;
9489 let rows = self.visible_rows(snap)?;
9490 execution_checkpoint(control, 0)?;
9491 let index_sets = self.resolve_index_conditions(conditions, snap)?;
9492 agg_state_from_rows(
9493 &rows,
9494 conditions,
9495 &index_sets,
9496 column,
9497 agg,
9498 &self.schema,
9499 control,
9500 )
9501 }
9502
9503 fn resolve_index_conditions(
9506 &self,
9507 conditions: &[crate::query::Condition],
9508 snapshot: Snapshot,
9509 ) -> Result<Vec<RowIdSet>> {
9510 use crate::query::Condition;
9511 let mut sets = Vec::new();
9512 for c in conditions {
9513 if matches!(
9514 c,
9515 Condition::Ann { .. }
9516 | Condition::SparseMatch { .. }
9517 | Condition::MinHashSimilar { .. }
9518 ) {
9519 sets.push(self.resolve_condition(c, snapshot)?);
9520 }
9521 }
9522 Ok(sets)
9523 }
9524
9525 fn column_type(&self, cid: u16) -> TypeId {
9526 self.schema
9527 .columns
9528 .iter()
9529 .find(|c| c.id == cid)
9530 .map(|c| c.ty.clone())
9531 .unwrap_or(TypeId::Bytes)
9532 }
9533
9534 pub fn approx_aggregate(
9543 &mut self,
9544 conditions: &[crate::query::Condition],
9545 column: Option<u16>,
9546 agg: ApproxAgg,
9547 z: f64,
9548 ) -> Result<Option<ApproxResult>> {
9549 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
9550 }
9551
9552 pub fn approx_aggregate_with_candidate_authorization(
9555 &mut self,
9556 conditions: &[crate::query::Condition],
9557 column: Option<u16>,
9558 agg: ApproxAgg,
9559 z: f64,
9560 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9561 ) -> Result<Option<ApproxResult>> {
9562 use crate::query::Condition;
9563 self.ensure_reservoir_complete()?;
9564 let snapshot = self.snapshot();
9565 let n_pop = self.count();
9566 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
9567 if sample_rids.is_empty() {
9568 return Ok(None);
9569 }
9570 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
9572 let s = live_sample.len();
9573 if s == 0 {
9574 return Ok(None);
9575 }
9576 let authorized = authorization
9577 .map(|authorization| {
9578 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
9579 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
9580 })
9581 .transpose()?;
9582
9583 let mut index_sets: Vec<RowIdSet> = Vec::new();
9586 for c in conditions {
9587 if matches!(
9588 c,
9589 Condition::Ann { .. }
9590 | Condition::SparseMatch { .. }
9591 | Condition::MinHashSimilar { .. }
9592 ) {
9593 index_sets.push(self.resolve_condition(c, snapshot)?);
9594 }
9595 }
9596
9597 let cid = match (agg, column) {
9599 (ApproxAgg::Count, _) => None,
9600 (_, Some(c)) => Some(c),
9601 _ => return Ok(None),
9602 };
9603 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
9604 for r in &live_sample {
9605 if authorized
9606 .as_ref()
9607 .is_some_and(|authorized| !authorized.contains(&r.row_id))
9608 {
9609 continue;
9610 }
9611 if !conditions
9613 .iter()
9614 .all(|c| condition_matches_row(c, r, &self.schema))
9615 {
9616 continue;
9617 }
9618 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
9620 continue;
9621 }
9622 if let Some(cid) = cid {
9623 let mut cells = r
9624 .columns
9625 .get(&cid)
9626 .cloned()
9627 .map(|value| vec![(cid, value)])
9628 .unwrap_or_default();
9629 if let Some(authorization) = authorization {
9630 authorization.security.apply_masks_to_cells(
9631 authorization.table,
9632 &mut cells,
9633 authorization.principal,
9634 );
9635 }
9636 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
9637 passing_vals.push(v);
9638 } } else {
9640 passing_vals.push(0.0); }
9642 }
9643 let m = passing_vals.len();
9644
9645 let (point, half) = match agg {
9646 ApproxAgg::Count => {
9647 let p = m as f64 / s as f64;
9649 let point = n_pop as f64 * p;
9650 let var = if s > 1 {
9651 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
9652 * (1.0 - s as f64 / n_pop as f64).max(0.0)
9653 } else {
9654 0.0
9655 };
9656 (point, z * var.sqrt())
9657 }
9658 ApproxAgg::Sum => {
9659 let y: Vec<f64> = live_sample
9661 .iter()
9662 .map(|r| {
9663 let passes_row = authorized
9664 .as_ref()
9665 .is_none_or(|authorized| authorized.contains(&r.row_id))
9666 && conditions
9667 .iter()
9668 .all(|c| condition_matches_row(c, r, &self.schema))
9669 && index_sets.iter().all(|set| set.contains(r.row_id.0));
9670 if passes_row {
9671 cid.and_then(|cid| {
9672 let mut cells = r
9673 .columns
9674 .get(&cid)
9675 .cloned()
9676 .map(|value| vec![(cid, value)])
9677 .unwrap_or_default();
9678 if let Some(authorization) = authorization {
9679 authorization.security.apply_masks_to_cells(
9680 authorization.table,
9681 &mut cells,
9682 authorization.principal,
9683 );
9684 }
9685 as_f64(cells.first().map(|(_, value)| value))
9686 })
9687 .unwrap_or(0.0)
9688 } else {
9689 0.0
9690 }
9691 })
9692 .collect();
9693 let mean_y = y.iter().sum::<f64>() / s as f64;
9694 let point = n_pop as f64 * mean_y;
9695 let var = if s > 1 {
9696 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
9697 let var_y = ss / (s - 1) as f64;
9698 n_pop as f64 * n_pop as f64 * var_y / s as f64
9699 * (1.0 - s as f64 / n_pop as f64).max(0.0)
9700 } else {
9701 0.0
9702 };
9703 (point, z * var.sqrt())
9704 }
9705 ApproxAgg::Avg => {
9706 if m == 0 {
9707 return Ok(Some(ApproxResult {
9708 point: 0.0,
9709 ci_low: 0.0,
9710 ci_high: 0.0,
9711 n_population: n_pop,
9712 n_sample_live: s,
9713 n_passing: 0,
9714 }));
9715 }
9716 let mean = passing_vals.iter().sum::<f64>() / m as f64;
9717 let half = if m > 1 {
9718 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
9719 let sd = (ss / (m - 1) as f64).sqrt();
9720 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
9721 z * sd / (m as f64).sqrt() * fpc.sqrt()
9722 } else {
9723 0.0
9724 };
9725 (mean, half)
9726 }
9727 };
9728
9729 Ok(Some(ApproxResult {
9730 point,
9731 ci_low: point - half,
9732 ci_high: point + half,
9733 n_population: n_pop,
9734 n_sample_live: s,
9735 n_passing: m,
9736 }))
9737 }
9738
9739 pub fn exact_column_stats(
9747 &self,
9748 _snapshot: Snapshot,
9749 projection: &[u16],
9750 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
9751 if self.ttl.is_some()
9752 || !(self.memtable.is_empty()
9753 && self.mutable_run.is_empty()
9754 && self.run_refs.len() == 1)
9755 {
9756 return Ok(None);
9757 }
9758 let reader = self.open_reader(self.run_refs[0].run_id)?;
9759 if self.live_count != reader.row_count() as u64 {
9760 return Ok(None);
9761 }
9762 let mut out = HashMap::new();
9763 for &cid in projection {
9764 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
9765 Some(c) => c,
9766 None => continue,
9767 };
9768 let Some(stats) = reader.column_page_stats(cid) else {
9770 out.insert(
9771 cid,
9772 ColumnStat {
9773 min: None,
9774 max: None,
9775 null_count: self.live_count,
9776 },
9777 );
9778 continue;
9779 };
9780 let stat = match cdef.ty {
9781 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
9782 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
9783 min: mn.map(Value::Int64),
9784 max: mx.map(Value::Int64),
9785 null_count: n,
9786 })
9787 }
9788 TypeId::Float64 => {
9789 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
9790 min: mn.map(Value::Float64),
9791 max: mx.map(Value::Float64),
9792 null_count: n,
9793 })
9794 }
9795 _ => None,
9796 };
9797 if let Some(s) = stat {
9798 out.insert(cid, s);
9799 }
9800 }
9801 Ok(Some(out))
9802 }
9803
9804 pub fn dir(&self) -> &Path {
9805 &self.dir
9806 }
9807
9808 pub fn schema(&self) -> &Schema {
9809 &self.schema
9810 }
9811
9812 pub(crate) fn set_catalog_name(&mut self, name: String) {
9813 self.name = name;
9814 }
9815
9816 pub(crate) fn prepare_alter_column(
9817 &mut self,
9818 column_name: &str,
9819 change: &AlterColumn,
9820 ) -> Result<(ColumnDef, Option<Schema>)> {
9821 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
9822 return Err(MongrelError::InvalidArgument(
9823 "ALTER COLUMN requires committing staged writes first".into(),
9824 ));
9825 }
9826 let old = self
9827 .schema
9828 .columns
9829 .iter()
9830 .find(|c| c.name == column_name)
9831 .cloned()
9832 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
9833 let mut next = old.clone();
9834
9835 if let Some(name) = &change.name {
9836 let trimmed = name.trim();
9837 if trimmed.is_empty() {
9838 return Err(MongrelError::InvalidArgument(
9839 "ALTER COLUMN name must not be empty".into(),
9840 ));
9841 }
9842 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
9843 return Err(MongrelError::Schema(format!(
9844 "column {trimmed} already exists"
9845 )));
9846 }
9847 next.name = trimmed.to_string();
9848 }
9849
9850 if let Some(ty) = &change.ty {
9851 next.ty = ty.clone();
9852 }
9853 if let Some(flags) = change.flags {
9854 validate_alter_column_flags(old.flags, flags)?;
9855 next.flags = flags;
9856 }
9857
9858 if let Some(default_change) = &change.default_value {
9859 next.default_value = default_change.clone();
9860 }
9861
9862 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
9863 if old.flags.contains(ColumnFlags::NULLABLE)
9864 && !next.flags.contains(ColumnFlags::NULLABLE)
9865 && self.column_has_nulls(old.id)?
9866 {
9867 return Err(MongrelError::InvalidArgument(format!(
9868 "column '{}' contains NULL values",
9869 old.name
9870 )));
9871 }
9872 if next == old {
9873 return Ok((next, None));
9874 }
9875 let mut schema = self.schema.clone();
9876 let index = schema
9877 .columns
9878 .iter()
9879 .position(|column| column.id == next.id)
9880 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
9881 schema.columns[index] = next.clone();
9882 schema.schema_id = schema
9883 .schema_id
9884 .checked_add(1)
9885 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
9886 schema.validate_auto_increment()?;
9887 schema.validate_defaults()?;
9888 Ok((next, Some(schema)))
9889 }
9890
9891 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
9892 self.schema = schema;
9893 self.auto_inc = resolve_auto_inc(&self.schema);
9894 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
9895 self.clear_result_cache();
9896 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
9897 }
9898
9899 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
9900 checkpoint_current_schema(self)
9901 }
9902
9903 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
9904 self.ensure_writable()?;
9905 let previous_schema = self.schema.clone();
9906 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
9907 if let Some(schema) = schema {
9908 self.apply_altered_schema_prepared(schema);
9909 self.checkpoint_standalone_schema_change(previous_schema)?;
9910 }
9911 Ok(column)
9912 }
9913
9914 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
9915 if self.live_count == 0 {
9916 return Ok(false);
9917 }
9918 let snap = self.snapshot();
9919 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
9920 Ok(columns
9921 .first()
9922 .map(|(_, col)| col.null_count(col.len()) != 0)
9923 .unwrap_or(true))
9924 }
9925
9926 fn has_stored_versions(&self) -> bool {
9927 !self.memtable.is_empty()
9928 || !self.mutable_run.is_empty()
9929 || self.run_refs.iter().any(|r| r.row_count > 0)
9930 || !self.retiring.is_empty()
9931 }
9932
9933 pub fn add_column(
9938 &mut self,
9939 name: &str,
9940 ty: TypeId,
9941 flags: ColumnFlags,
9942 default_value: Option<crate::schema::DefaultExpr>,
9943 ) -> Result<u16> {
9944 self.add_column_with_id(name, ty, flags, default_value, None)
9945 }
9946
9947 pub fn add_column_with_id(
9948 &mut self,
9949 name: &str,
9950 ty: TypeId,
9951 flags: ColumnFlags,
9952 default_value: Option<crate::schema::DefaultExpr>,
9953 requested_id: Option<u16>,
9954 ) -> Result<u16> {
9955 self.ensure_writable()?;
9956 if self.schema.columns.iter().any(|c| c.name == name) {
9957 return Err(MongrelError::Schema(format!(
9958 "column {name} already exists"
9959 )));
9960 }
9961 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
9962 if self.schema.columns.iter().any(|c| c.id == id) {
9963 return Err(MongrelError::Schema(format!(
9964 "column id {id} already exists"
9965 )));
9966 }
9967 id
9968 } else {
9969 self.schema
9970 .columns
9971 .iter()
9972 .map(|c| c.id)
9973 .max()
9974 .unwrap_or(0)
9975 .checked_add(1)
9976 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
9977 };
9978 let previous_schema = self.schema.clone();
9979 let mut next_schema = previous_schema.clone();
9980 next_schema.columns.push(ColumnDef {
9981 id,
9982 name: name.to_string(),
9983 ty,
9984 flags,
9985 default_value,
9986 });
9987 next_schema.schema_id = next_schema
9988 .schema_id
9989 .checked_add(1)
9990 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
9991 next_schema.validate_auto_increment()?;
9992 next_schema.validate_defaults()?;
9993 self.apply_altered_schema_prepared(next_schema);
9994 self.checkpoint_standalone_schema_change(previous_schema)?;
9995 Ok(id)
9996 }
9997
9998 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
10007 self.ensure_writable()?;
10008 let cid = self
10009 .schema
10010 .columns
10011 .iter()
10012 .find(|c| c.name == column_name)
10013 .map(|c| c.id)
10014 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
10015 let ty = self
10016 .schema
10017 .columns
10018 .iter()
10019 .find(|c| c.id == cid)
10020 .map(|c| c.ty.clone())
10021 .unwrap_or(TypeId::Int64);
10022 if !matches!(
10023 ty,
10024 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
10025 ) {
10026 return Err(MongrelError::Schema(format!(
10027 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
10028 )));
10029 }
10030 if self
10031 .schema
10032 .indexes
10033 .iter()
10034 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
10035 {
10036 return Ok(()); }
10038 let previous_schema = self.schema.clone();
10039 let previous_learned_range = Arc::clone(&self.learned_range);
10040 let mut next_schema = previous_schema.clone();
10041 next_schema.indexes.push(IndexDef {
10042 name: format!("{}_learned_range", column_name),
10043 column_id: cid,
10044 kind: IndexKind::LearnedRange,
10045 predicate: None,
10046 options: Default::default(),
10047 });
10048 next_schema.schema_id = next_schema
10049 .schema_id
10050 .checked_add(1)
10051 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
10052 self.apply_altered_schema_prepared(next_schema);
10053 if let Err(error) = self.build_learned_ranges() {
10054 self.apply_altered_schema_prepared(previous_schema);
10055 self.learned_range = previous_learned_range;
10056 return Err(error);
10057 }
10058 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
10059 if !matches!(
10060 &error,
10061 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
10062 ) {
10063 self.learned_range = previous_learned_range;
10064 }
10065 return Err(error);
10066 }
10067 Ok(())
10068 }
10069
10070 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
10071 let mut schema_published = false;
10072 let schema_result = match self._root_guard.as_deref() {
10073 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
10074 schema_published = true;
10075 }),
10076 None => write_schema_with_after(&self.dir, &self.schema, || {
10077 schema_published = true;
10078 }),
10079 };
10080 if schema_result.is_err() && !schema_published {
10081 self.apply_altered_schema_prepared(previous_schema);
10082 return schema_result;
10083 }
10084
10085 let manifest_result = self.persist_manifest(self.current_epoch());
10086 match (schema_result, manifest_result) {
10087 (_, Ok(())) => Ok(()),
10088 (Ok(()), Err(error)) => {
10089 self.poison_after_maintenance_publish_failure();
10090 Err(MongrelError::DurableCommit {
10091 epoch: self.current_epoch().0,
10092 message: format!(
10093 "schema is durable but matching manifest publication failed: {error}"
10094 ),
10095 })
10096 }
10097 (Err(schema_error), Err(manifest_error)) => {
10098 self.poison_after_maintenance_publish_failure();
10099 Err(MongrelError::CommitOutcomeUnknown {
10100 epoch: self.current_epoch().0,
10101 message: format!(
10102 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
10103 ),
10104 })
10105 }
10106 }
10107 }
10108
10109 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
10112 self.sync_byte_threshold = threshold;
10113 if let WalSink::Private(w) = &mut self.wal {
10114 w.set_sync_byte_threshold(threshold);
10115 }
10116 }
10117
10118 pub fn page_cache_flush(&self) {
10122 self.page_cache.flush_to_disk();
10123 }
10124
10125 pub fn page_cache_len(&self) -> usize {
10127 self.page_cache.len()
10128 }
10129
10130 pub fn decoded_cache_len(&self) -> usize {
10133 self.decoded_cache.len()
10134 }
10135
10136 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
10139 self.memtable.drain_sorted()
10140 }
10141
10142 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
10143 self.runs_dir().join(format!("r-{run_id}.sr"))
10144 }
10145
10146 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
10147 match self.runs_root.as_deref() {
10148 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
10149 None => Ok(None),
10150 }
10151 }
10152
10153 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
10154 match self.runs_root.as_deref() {
10155 Some(root) => Ok(Some(root.create_regular_new(name)?)),
10156 None => Ok(None),
10157 }
10158 }
10159
10160 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
10161 match self.runs_root.as_deref() {
10162 Some(root) => match root.remove_file(name) {
10163 Ok(()) => Ok(()),
10164 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10165 Err(error) => Err(error.into()),
10166 },
10167 None => match std::fs::remove_file(self.runs_dir().join(name)) {
10168 Ok(()) => Ok(()),
10169 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
10170 Err(error) => Err(error.into()),
10171 },
10172 }
10173 }
10174
10175 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
10176 match self.runs_root.as_deref() {
10177 Some(root) => root
10178 .rename_file_new(source, destination)
10179 .map_err(Into::into),
10180 None => crate::durable_file::rename(
10181 &self.runs_dir().join(source),
10182 &self.runs_dir().join(destination),
10183 )
10184 .map_err(Into::into),
10185 }
10186 }
10187
10188 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10189 self.run_refs.iter().map(|run| run.run_id)
10190 }
10191
10192 pub(crate) fn table_dir(&self) -> &Path {
10193 &self.dir
10194 }
10195
10196 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
10197 &self.schema
10198 }
10199
10200 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
10201 let id = self.next_run_id;
10202 self.next_run_id = self
10203 .next_run_id
10204 .checked_add(1)
10205 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
10206 Ok(id)
10207 }
10208
10209 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
10210 self.run_refs.push(run_ref);
10211 }
10212
10213 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
10223 self.retiring.push(crate::manifest::RetiredRun {
10224 run_id,
10225 retire_epoch,
10226 });
10227 }
10228
10229 pub(crate) fn reap_retiring(
10233 &mut self,
10234 min_active: Epoch,
10235 backup_pinned: &std::collections::HashSet<u128>,
10236 ) -> Result<usize> {
10237 if self.retiring.is_empty() {
10238 return Ok(0);
10239 }
10240 let mut reaped = 0;
10241 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
10242 for r in std::mem::take(&mut self.retiring) {
10248 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
10249 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
10250 reaped += 1;
10251 } else {
10252 kept.push(r);
10253 }
10254 }
10255 self.retiring = kept;
10256 if reaped > 0 {
10257 self.persist_manifest(self.current_epoch())?;
10258 }
10259 Ok(reaped)
10260 }
10261
10262 pub(crate) fn has_reapable_retiring(
10263 &self,
10264 min_active: Epoch,
10265 backup_pinned: &std::collections::HashSet<u128>,
10266 ) -> bool {
10267 self.retiring
10268 .iter()
10269 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
10270 }
10271
10272 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
10273 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
10274 return false;
10275 }
10276 self.live_count = self.live_count.saturating_add(run_ref.row_count);
10277 self.run_refs.push(run_ref);
10278 self.indexes_complete = false;
10279 true
10280 }
10281
10282 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
10283 self.kek.as_ref()
10284 }
10285
10286 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
10287 let mut reader = match self.runs_root.as_deref() {
10288 Some(root) => RunReader::open_file_with_cache(
10289 root.open_regular(format!("r-{run_id}.sr"))?,
10290 self.schema.clone(),
10291 self.kek.clone(),
10292 Some(self.page_cache.clone()),
10293 Some(self.decoded_cache.clone()),
10294 self.table_id,
10295 Some(&self.verified_runs),
10296 None,
10297 )?,
10298 None => RunReader::open_with_cache(
10299 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
10300 self.schema.clone(),
10301 self.kek.clone(),
10302 Some(self.page_cache.clone()),
10303 Some(self.decoded_cache.clone()),
10304 self.table_id,
10305 Some(&self.verified_runs),
10306 )?,
10307 };
10308 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
10312 reader.set_uniform_epoch(Epoch(rr.epoch_created));
10313 }
10314 Ok(reader)
10315 }
10316
10317 pub(crate) fn run_refs(&self) -> &[RunRef] {
10318 &self.run_refs
10319 }
10320
10321 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
10322 self.retiring.iter().map(|run| run.run_id)
10323 }
10324
10325 pub(crate) fn runs_dir(&self) -> PathBuf {
10326 self.runs_root
10327 .as_deref()
10328 .and_then(|root| root.io_path().ok())
10329 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
10330 }
10331
10332 pub(crate) fn wal_dir(&self) -> PathBuf {
10333 self.dir.join(WAL_DIR)
10334 }
10335
10336 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
10337 self.run_refs = refs;
10338 }
10339
10340 pub(crate) fn compaction_zstd_level(&self) -> i32 {
10341 self.compaction_zstd_level
10342 }
10343
10344 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
10345 self.kek.clone()
10346 }
10347
10348 #[cfg(feature = "encryption")]
10352 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10353 self.kek.as_ref().map(|k| k.derive_idx_key())
10354 }
10355
10356 #[cfg(not(feature = "encryption"))]
10357 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
10358 None
10359 }
10360
10361 #[cfg(feature = "encryption")]
10365 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10366 self.kek.as_ref().map(|k| *k.derive_meta_key())
10367 }
10368
10369 #[cfg(not(feature = "encryption"))]
10370 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
10371 None
10372 }
10373
10374 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
10377 self.column_keys
10378 .iter()
10379 .map(|(&id, &(_, scheme))| (id, scheme))
10380 .collect()
10381 }
10382
10383 #[cfg(feature = "encryption")]
10388 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
10389 self.tokenize_value_enc(column_id, v)
10390 }
10391
10392 #[cfg(feature = "encryption")]
10393 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
10394 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
10395 let (key, scheme) = self.column_keys.get(&column_id)?;
10396 let token: Vec<u8> = match (*scheme, v) {
10397 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
10398 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
10399 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
10400 _ => hmac_token(key, &v.encode_key()).to_vec(),
10401 };
10402 Some(Value::Bytes(token))
10403 }
10404
10405 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
10407 self.index_lookup_key_bytes(column_id, &v.encode_key())
10408 }
10409
10410 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
10413 #[cfg(feature = "encryption")]
10414 {
10415 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
10416 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
10417 if *scheme == SCHEME_HMAC_EQ {
10418 return hmac_token(key, encoded).to_vec();
10419 }
10420 }
10421 }
10422 let _ = column_id;
10423 encoded.to_vec()
10424 }
10425}
10426
10427fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
10428 let columnar::NativeColumn::Int64 { data, validity } = col else {
10429 return false;
10430 };
10431 if data.len() < n || !columnar::all_non_null(validity, n) {
10432 return false;
10433 }
10434 data.iter()
10435 .take(n)
10436 .zip(data.iter().skip(1))
10437 .all(|(a, b)| a < b)
10438}
10439
10440#[derive(Debug, Clone)]
10444pub struct ColumnStat {
10445 pub min: Option<Value>,
10446 pub max: Option<Value>,
10447 pub null_count: u64,
10448}
10449
10450#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10452pub enum NativeAgg {
10453 Count,
10454 Sum,
10455 Min,
10456 Max,
10457 Avg,
10458}
10459
10460#[derive(Debug, Clone, PartialEq)]
10462pub enum NativeAggResult {
10463 Count(u64),
10464 Int(i64),
10465 Float(f64),
10466 Null,
10468}
10469
10470#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10472pub enum ApproxAgg {
10473 Count,
10474 Sum,
10475 Avg,
10476}
10477
10478#[derive(Debug, Clone)]
10482pub struct ApproxResult {
10483 pub point: f64,
10485 pub ci_low: f64,
10487 pub ci_high: f64,
10489 pub n_population: u64,
10491 pub n_sample_live: usize,
10493 pub n_passing: usize,
10495}
10496
10497#[derive(Debug, Clone, PartialEq)]
10502pub enum AggState {
10503 Count(u64),
10505 SumI {
10507 sum: i128,
10508 count: u64,
10509 },
10510 SumF {
10512 sum: f64,
10513 count: u64,
10514 },
10515 AvgI {
10517 sum: i128,
10518 count: u64,
10519 },
10520 AvgF {
10522 sum: f64,
10523 count: u64,
10524 },
10525 MinI(i64),
10527 MaxI(i64),
10528 MinF(f64),
10530 MaxF(f64),
10531 Empty,
10533}
10534
10535impl AggState {
10536 pub fn merge(self, other: AggState) -> AggState {
10538 use AggState::*;
10539 match (self, other) {
10540 (Empty, x) | (x, Empty) => x,
10541 (Count(a), Count(b)) => Count(a + b),
10542 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
10543 sum: sa + sb,
10544 count: ca + cb,
10545 },
10546 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
10547 sum: sa + sb,
10548 count: ca + cb,
10549 },
10550 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
10551 sum: sa + sb,
10552 count: ca + cb,
10553 },
10554 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
10555 sum: sa + sb,
10556 count: ca + cb,
10557 },
10558 (MinI(a), MinI(b)) => MinI(a.min(b)),
10559 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
10560 (MinF(a), MinF(b)) => MinF(a.min(b)),
10561 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
10562 _ => Empty, }
10564 }
10565
10566 pub fn point(&self) -> Option<f64> {
10568 match self {
10569 AggState::Count(n) => Some(*n as f64),
10570 AggState::SumI { sum, .. } => Some(*sum as f64),
10571 AggState::SumF { sum, .. } => Some(*sum),
10572 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
10573 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
10574 AggState::MinI(n) => Some(*n as f64),
10575 AggState::MaxI(n) => Some(*n as f64),
10576 AggState::MinF(n) => Some(*n),
10577 AggState::MaxF(n) => Some(*n),
10578 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
10579 }
10580 }
10581
10582 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
10586 let is_float = matches!(ty, Some(TypeId::Float64));
10587 match (agg, result) {
10588 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
10589 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
10590 sum: x as i128,
10591 count: 1, },
10593 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
10594 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
10595 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
10596 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
10597 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
10598 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
10599 (NativeAgg::Count, _) => AggState::Empty,
10600 (_, NativeAggResult::Null) => AggState::Empty,
10601 _ => {
10602 let _ = is_float;
10603 AggState::Empty
10604 }
10605 }
10606 }
10607}
10608
10609#[derive(Debug, Clone)]
10612pub struct CachedAgg {
10613 pub state: AggState,
10614 pub watermark: u64,
10615 pub epoch: u64,
10616}
10617
10618#[derive(Debug, Clone)]
10620pub struct IncrementalAggResult {
10621 pub state: AggState,
10623 pub incremental: bool,
10626 pub delta_rows: u64,
10628}
10629
10630fn agg_state_from_rows(
10634 rows: &[Row],
10635 conditions: &[crate::query::Condition],
10636 index_sets: &[RowIdSet],
10637 column: Option<u16>,
10638 agg: NativeAgg,
10639 schema: &Schema,
10640 control: Option<&crate::ExecutionControl>,
10641) -> Result<AggState> {
10642 let mut count: u64 = 0;
10643 let mut sum_i: i128 = 0;
10644 let mut sum_f: f64 = 0.0;
10645 let mut mn_i: i64 = i64::MAX;
10646 let mut mx_i: i64 = i64::MIN;
10647 let mut mn_f: f64 = f64::INFINITY;
10648 let mut mx_f: f64 = f64::NEG_INFINITY;
10649 let mut saw_int = false;
10650 let mut saw_float = false;
10651 for (index, r) in rows.iter().enumerate() {
10652 execution_checkpoint(control, index)?;
10653 if !conditions
10654 .iter()
10655 .all(|c| condition_matches_row(c, r, schema))
10656 {
10657 continue;
10658 }
10659 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
10660 continue;
10661 }
10662 match agg {
10663 NativeAgg::Count => match column {
10664 None => count += 1,
10666 Some(cid) => match r.columns.get(&cid) {
10669 None | Some(Value::Null) => {}
10670 Some(_) => count += 1,
10671 },
10672 },
10673 _ => match column.and_then(|cid| r.columns.get(&cid)) {
10674 Some(Value::Int64(n)) => {
10675 count += 1;
10676 sum_i += *n as i128;
10677 mn_i = mn_i.min(*n);
10678 mx_i = mx_i.max(*n);
10679 saw_int = true;
10680 }
10681 Some(Value::Float64(f)) => {
10682 count += 1;
10683 sum_f += f;
10684 mn_f = mn_f.min(*f);
10685 mx_f = mx_f.max(*f);
10686 saw_float = true;
10687 }
10688 _ => {}
10689 },
10690 }
10691 }
10692 Ok(match agg {
10693 NativeAgg::Count => {
10694 if count == 0 {
10695 AggState::Empty
10696 } else {
10697 AggState::Count(count)
10698 }
10699 }
10700 NativeAgg::Sum => {
10701 if count == 0 {
10702 AggState::Empty
10703 } else if saw_int {
10704 AggState::SumI { sum: sum_i, count }
10705 } else {
10706 AggState::SumF { sum: sum_f, count }
10707 }
10708 }
10709 NativeAgg::Avg => {
10710 if count == 0 {
10711 AggState::Empty
10712 } else if saw_int {
10713 AggState::AvgI { sum: sum_i, count }
10714 } else {
10715 AggState::AvgF { sum: sum_f, count }
10716 }
10717 }
10718 NativeAgg::Min => {
10719 if !saw_int && !saw_float {
10720 AggState::Empty
10721 } else if saw_int {
10722 AggState::MinI(mn_i)
10723 } else {
10724 AggState::MinF(mn_f)
10725 }
10726 }
10727 NativeAgg::Max => {
10728 if !saw_int && !saw_float {
10729 AggState::Empty
10730 } else if saw_int {
10731 AggState::MaxI(mx_i)
10732 } else {
10733 AggState::MaxF(mx_f)
10734 }
10735 }
10736 })
10737}
10738
10739fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
10743 use crate::query::Condition;
10744 match c {
10745 Condition::Pk(key) => match schema.primary_key() {
10746 Some(pk) => row
10747 .columns
10748 .get(&pk.id)
10749 .map(|v| v.encode_key() == *key)
10750 .unwrap_or(false),
10751 None => false,
10752 },
10753 Condition::BitmapEq { column_id, value } => row
10754 .columns
10755 .get(column_id)
10756 .map(|v| v.encode_key() == *value)
10757 .unwrap_or(false),
10758 Condition::BitmapIn { column_id, values } => {
10759 let key = row.columns.get(column_id).map(|v| v.encode_key());
10760 match key {
10761 Some(k) => values.contains(&k),
10762 None => false,
10763 }
10764 }
10765 Condition::BytesPrefix { column_id, prefix } => row
10766 .columns
10767 .get(column_id)
10768 .map(|v| v.encode_key().starts_with(prefix))
10769 .unwrap_or(false),
10770 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
10771 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
10772 _ => false,
10773 },
10774 Condition::RangeF64 {
10775 column_id,
10776 lo,
10777 lo_inclusive,
10778 hi,
10779 hi_inclusive,
10780 } => match row.columns.get(column_id) {
10781 Some(Value::Float64(n)) => {
10782 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
10783 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
10784 lo_ok && hi_ok
10785 }
10786 _ => false,
10787 },
10788 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
10789 Some(Value::Bytes(b)) => {
10790 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
10791 }
10792 _ => false,
10793 },
10794 Condition::FmContainsAll {
10795 column_id,
10796 patterns,
10797 } => match row.columns.get(column_id) {
10798 Some(Value::Bytes(b)) => patterns
10799 .iter()
10800 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
10801 _ => false,
10802 },
10803 Condition::Ann { .. }
10804 | Condition::SparseMatch { .. }
10805 | Condition::MinHashSimilar { .. } => true,
10806 Condition::IsNull { column_id } => {
10807 matches!(row.columns.get(column_id), Some(Value::Null) | None)
10808 }
10809 Condition::IsNotNull { column_id } => {
10810 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
10811 }
10812 }
10813}
10814
10815fn as_f64(v: Option<&Value>) -> Option<f64> {
10817 match v {
10818 Some(Value::Int64(n)) => Some(*n as f64),
10819 Some(Value::Float64(f)) => Some(*f),
10820 _ => None,
10821 }
10822}
10823
10824fn accumulate_int(
10828 cursor: &mut dyn crate::cursor::Cursor,
10829 control: Option<&crate::ExecutionControl>,
10830) -> Result<(u64, i128, i64, i64)> {
10831 let mut count: u64 = 0;
10832 let mut sum: i128 = 0;
10833 let mut mn: i64 = i64::MAX;
10834 let mut mx: i64 = i64::MIN;
10835 while let Some(cols) = cursor.next_batch()? {
10836 execution_checkpoint(control, 0)?;
10837 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
10838 if crate::columnar::all_non_null(validity, data.len()) {
10839 count += data.len() as u64;
10841 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
10842 execution_checkpoint(control, chunk_index * 1024)?;
10843 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
10844 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
10845 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
10846 }
10847 } else {
10848 for (i, &v) in data.iter().enumerate() {
10849 execution_checkpoint(control, i)?;
10850 if crate::columnar::validity_bit(validity, i) {
10851 count += 1;
10852 sum += v as i128;
10853 mn = mn.min(v);
10854 mx = mx.max(v);
10855 }
10856 }
10857 }
10858 }
10859 }
10860 Ok((count, sum, mn, mx))
10861}
10862
10863fn accumulate_float(
10865 cursor: &mut dyn crate::cursor::Cursor,
10866 control: Option<&crate::ExecutionControl>,
10867) -> Result<(u64, f64, f64, f64)> {
10868 let mut count: u64 = 0;
10869 let mut sum: f64 = 0.0;
10870 let mut mn: f64 = f64::INFINITY;
10871 let mut mx: f64 = f64::NEG_INFINITY;
10872 while let Some(cols) = cursor.next_batch()? {
10873 execution_checkpoint(control, 0)?;
10874 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
10875 if crate::columnar::all_non_null(validity, data.len()) {
10876 count += data.len() as u64;
10877 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
10878 execution_checkpoint(control, chunk_index * 1024)?;
10879 sum += chunk.iter().sum::<f64>();
10880 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
10881 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
10882 }
10883 } else {
10884 for (i, &v) in data.iter().enumerate() {
10885 execution_checkpoint(control, i)?;
10886 if crate::columnar::validity_bit(validity, i) {
10887 count += 1;
10888 sum += v;
10889 mn = mn.min(v);
10890 mx = mx.max(v);
10891 }
10892 }
10893 }
10894 }
10895 }
10896 Ok((count, sum, mn, mx))
10897}
10898
10899#[inline]
10900fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
10901 if index.is_multiple_of(256) {
10902 control
10903 .map(crate::ExecutionControl::checkpoint)
10904 .transpose()?;
10905 }
10906 Ok(())
10907}
10908
10909fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
10910 if count == 0 && !matches!(agg, NativeAgg::Count) {
10911 return NativeAggResult::Null;
10912 }
10913 match agg {
10914 NativeAgg::Count => NativeAggResult::Count(count),
10915 NativeAgg::Sum => match sum.try_into() {
10918 Ok(v) => NativeAggResult::Int(v),
10919 Err(_) => NativeAggResult::Null,
10920 },
10921 NativeAgg::Min => NativeAggResult::Int(mn),
10922 NativeAgg::Max => NativeAggResult::Int(mx),
10923 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
10924 }
10925}
10926
10927fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
10928 if count == 0 && !matches!(agg, NativeAgg::Count) {
10929 return NativeAggResult::Null;
10930 }
10931 match agg {
10932 NativeAgg::Count => NativeAggResult::Count(count),
10933 NativeAgg::Sum => NativeAggResult::Float(sum),
10934 NativeAgg::Min => NativeAggResult::Float(mn),
10935 NativeAgg::Max => NativeAggResult::Float(mx),
10936 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
10937 }
10938}
10939
10940fn agg_int(
10943 stats: &[crate::page::PageStat],
10944 decode: fn(Option<&[u8]>) -> Option<i64>,
10945) -> Option<(Option<i64>, Option<i64>, u64)> {
10946 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
10947 let mut any = false;
10948 for s in stats {
10949 if let Some(v) = decode(s.min.as_deref()) {
10950 mn = mn.min(v);
10951 any = true;
10952 }
10953 if let Some(v) = decode(s.max.as_deref()) {
10954 mx = mx.max(v);
10955 any = true;
10956 }
10957 nulls += s.null_count;
10958 }
10959 any.then_some((Some(mn), Some(mx), nulls))
10960}
10961
10962fn agg_float(
10964 stats: &[crate::page::PageStat],
10965 decode: fn(Option<&[u8]>) -> Option<f64>,
10966) -> Option<(Option<f64>, Option<f64>, u64)> {
10967 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
10968 let mut any = false;
10969 for s in stats {
10970 if let Some(v) = decode(s.min.as_deref()) {
10971 mn = mn.min(v);
10972 any = true;
10973 }
10974 if let Some(v) = decode(s.max.as_deref()) {
10975 mx = mx.max(v);
10976 any = true;
10977 }
10978 nulls += s.null_count;
10979 }
10980 any.then_some((Some(mn), Some(mx), nulls))
10981}
10982
10983type SecondaryIndexes = (
10985 HashMap<u16, BitmapIndex>,
10986 HashMap<u16, AnnIndex>,
10987 HashMap<u16, FmIndex>,
10988 HashMap<u16, SparseIndex>,
10989 HashMap<u16, MinHashIndex>,
10990);
10991
10992fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
10993 let mut bitmap = HashMap::new();
10994 let mut ann = HashMap::new();
10995 let mut fm = HashMap::new();
10996 let mut sparse = HashMap::new();
10997 let mut minhash = HashMap::new();
10998 for idef in &schema.indexes {
10999 match idef.kind {
11000 IndexKind::Bitmap => {
11001 bitmap.insert(idef.column_id, BitmapIndex::new());
11002 }
11003 IndexKind::Ann => {
11004 let dim = schema
11005 .columns
11006 .iter()
11007 .find(|c| c.id == idef.column_id)
11008 .and_then(|c| match c.ty {
11009 TypeId::Embedding { dim } => Some(dim as usize),
11010 _ => None,
11011 })
11012 .unwrap_or(0);
11013 let options = idef.options.ann.clone().unwrap_or_default();
11014 ann.insert(
11015 idef.column_id,
11016 AnnIndex::with_options(
11017 dim,
11018 options.m,
11019 options.ef_construction,
11020 options.ef_search,
11021 ),
11022 );
11023 }
11024 IndexKind::FmIndex => {
11025 fm.insert(idef.column_id, FmIndex::new());
11026 }
11027 IndexKind::Sparse => {
11028 sparse.insert(idef.column_id, SparseIndex::new());
11029 }
11030 IndexKind::MinHash => {
11031 let options = idef.options.minhash.clone().unwrap_or_default();
11032 minhash.insert(
11033 idef.column_id,
11034 MinHashIndex::with_options(options.permutations, options.bands),
11035 );
11036 }
11037 _ => {}
11038 }
11039 }
11040 (bitmap, ann, fm, sparse, minhash)
11041}
11042
11043const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
11044 | ColumnFlags::AUTO_INCREMENT
11045 | ColumnFlags::ENCRYPTED
11046 | ColumnFlags::ENCRYPTED_INDEXABLE
11047 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
11048
11049fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
11050 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
11051 return Err(MongrelError::Schema(
11052 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
11053 ));
11054 }
11055 Ok(())
11056}
11057
11058fn validate_alter_column_type(
11059 schema: &Schema,
11060 old: &ColumnDef,
11061 next: &ColumnDef,
11062 has_stored_versions: bool,
11063) -> Result<()> {
11064 if old.ty == next.ty {
11065 return Ok(());
11066 }
11067 if schema.indexes.iter().any(|i| i.column_id == old.id) {
11068 return Err(MongrelError::Schema(format!(
11069 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
11070 old.name
11071 )));
11072 }
11073 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
11074 return Ok(());
11075 }
11076 Err(MongrelError::Schema(format!(
11077 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
11078 old.ty, next.ty
11079 )))
11080}
11081
11082fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
11083 matches!(
11084 (old, new),
11085 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
11086 )
11087}
11088
11089fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
11095 let mut prev: Option<i64> = None;
11096 for r in rows {
11097 match r.columns.get(&pk_id) {
11098 Some(Value::Int64(v)) => {
11099 if prev.is_some_and(|p| p >= *v) {
11100 return false;
11101 }
11102 prev = Some(*v);
11103 }
11104 _ => return false,
11105 }
11106 }
11107 true
11108}
11109
11110#[allow(clippy::too_many_arguments)]
11111fn index_into(
11112 schema: &Schema,
11113 row: &Row,
11114 hot: &mut HotIndex,
11115 bitmap: &mut HashMap<u16, BitmapIndex>,
11116 ann: &mut HashMap<u16, AnnIndex>,
11117 fm: &mut HashMap<u16, FmIndex>,
11118 sparse: &mut HashMap<u16, SparseIndex>,
11119 minhash: &mut HashMap<u16, MinHashIndex>,
11120) {
11121 for idef in &schema.indexes {
11122 let Some(val) = row.columns.get(&idef.column_id) else {
11123 continue;
11124 };
11125 match idef.kind {
11126 IndexKind::Bitmap => {
11127 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11128 b.insert(val.encode_key(), row.row_id);
11129 }
11130 }
11131 IndexKind::Ann => {
11132 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11133 a.insert_validated(v, row.row_id);
11134 }
11135 }
11136 IndexKind::FmIndex => {
11137 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11138 f.insert(b.clone(), row.row_id);
11139 }
11140 }
11141 IndexKind::Sparse => {
11142 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11143 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11146 s.insert(&terms, row.row_id);
11147 }
11148 }
11149 }
11150 IndexKind::MinHash => {
11151 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11152 let tokens = crate::index::token_hashes_from_bytes(b);
11155 mh.insert(&tokens, row.row_id);
11156 }
11157 }
11158 _ => {}
11159 }
11160 }
11161 if let Some(pk_col) = schema.primary_key() {
11162 if let Some(pk_val) = row.columns.get(&pk_col.id) {
11163 hot.insert(pk_val.encode_key(), row.row_id);
11164 }
11165 }
11166}
11167
11168#[allow(clippy::too_many_arguments)]
11171fn index_into_single(
11172 idef: &IndexDef,
11173 _schema: &Schema,
11174 row: &Row,
11175 _hot: &mut HotIndex,
11176 bitmap: &mut HashMap<u16, BitmapIndex>,
11177 ann: &mut HashMap<u16, AnnIndex>,
11178 fm: &mut HashMap<u16, FmIndex>,
11179 sparse: &mut HashMap<u16, SparseIndex>,
11180 minhash: &mut HashMap<u16, MinHashIndex>,
11181) {
11182 let Some(val) = row.columns.get(&idef.column_id) else {
11183 return;
11184 };
11185 match idef.kind {
11186 IndexKind::Bitmap => {
11187 if let Some(b) = bitmap.get_mut(&idef.column_id) {
11188 b.insert(val.encode_key(), row.row_id);
11189 }
11190 }
11191 IndexKind::Ann => {
11192 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
11193 a.insert_validated(v, row.row_id);
11194 }
11195 }
11196 IndexKind::FmIndex => {
11197 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
11198 f.insert(b.clone(), row.row_id);
11199 }
11200 }
11201 IndexKind::Sparse => {
11202 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
11203 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
11204 s.insert(&terms, row.row_id);
11205 }
11206 }
11207 }
11208 IndexKind::MinHash => {
11209 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
11210 let tokens = crate::index::token_hashes_from_bytes(b);
11211 mh.insert(&tokens, row.row_id);
11212 }
11213 }
11214 _ => {}
11215 }
11216}
11217
11218fn eval_partial_predicate(
11224 pred: &str,
11225 columns_map: &HashMap<u16, &Value>,
11226 name_to_id: &HashMap<&str, u16>,
11227) -> bool {
11228 let lower = pred.trim().to_ascii_lowercase();
11229 if let Some(rest) = lower.strip_suffix(" is not null") {
11231 let col_name = rest.trim();
11232 if let Some(col_id) = name_to_id.get(col_name) {
11233 return columns_map
11234 .get(col_id)
11235 .is_some_and(|v| !matches!(v, Value::Null));
11236 }
11237 }
11238 if let Some(rest) = lower.strip_suffix(" is null") {
11240 let col_name = rest.trim();
11241 if let Some(col_id) = name_to_id.get(col_name) {
11242 return columns_map
11243 .get(col_id)
11244 .is_none_or(|v| matches!(v, Value::Null));
11245 }
11246 }
11247 true
11250}
11251
11252#[allow(dead_code)]
11258fn bulk_index_key(
11259 column_keys: &HashMap<u16, ([u8; 32], u8)>,
11260 column_id: u16,
11261 ty: TypeId,
11262 col: &columnar::NativeColumn,
11263 i: usize,
11264) -> Option<Vec<u8>> {
11265 let encoded = columnar::encode_key_native(ty, col, i)?;
11266 #[cfg(feature = "encryption")]
11267 {
11268 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
11269 if let Some((key, scheme)) = column_keys.get(&column_id) {
11270 return Some(match (*scheme, col) {
11271 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
11272 (_, columnar::NativeColumn::Int64 { data, .. }) => {
11273 ope_token_i64(key, data[i]).to_vec()
11274 }
11275 (_, columnar::NativeColumn::Float64 { data, .. }) => {
11276 ope_token_f64(key, data[i]).to_vec()
11277 }
11278 _ => hmac_token(key, &encoded).to_vec(),
11279 });
11280 }
11281 }
11282 #[cfg(not(feature = "encryption"))]
11283 {
11284 let _ = (column_id, column_keys, col);
11285 }
11286 Some(encoded)
11287}
11288
11289pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
11290 write_schema_with_after(dir, schema, || {})
11291}
11292
11293pub(crate) fn write_schema_durable(
11294 root: &crate::durable_file::DurableRoot,
11295 schema: &Schema,
11296) -> Result<()> {
11297 write_schema_durable_with_after(root, schema, || {})
11298}
11299
11300fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
11301where
11302 F: FnOnce(),
11303{
11304 let json = serde_json::to_string_pretty(schema)
11305 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
11306 crate::durable_file::write_atomic_with_after(
11307 &dir.join(SCHEMA_FILENAME),
11308 json.as_bytes(),
11309 after_publish,
11310 )?;
11311 Ok(())
11312}
11313
11314fn write_schema_durable_with_after<F>(
11315 root: &crate::durable_file::DurableRoot,
11316 schema: &Schema,
11317 after_publish: F,
11318) -> Result<()>
11319where
11320 F: FnOnce(),
11321{
11322 let json = serde_json::to_string_pretty(schema)
11323 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
11324 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
11325 Ok(())
11326}
11327
11328fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
11329 let mut schema_published = false;
11330 let schema_result = match table._root_guard.as_deref() {
11331 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
11332 schema_published = true;
11333 }),
11334 None => write_schema_with_after(&table.dir, &table.schema, || {
11335 schema_published = true;
11336 }),
11337 };
11338 if schema_result.is_err() && !schema_published {
11339 return schema_result;
11340 }
11341 match table.persist_manifest(table.current_epoch()) {
11342 Ok(()) => Ok(()),
11343 Err(manifest_error) => Err(match schema_result {
11344 Ok(()) => manifest_error,
11345 Err(schema_error) => MongrelError::Other(format!(
11346 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
11347 )),
11348 }),
11349 }
11350}
11351
11352fn read_schema(dir: &Path) -> Result<Schema> {
11353 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
11354 read_schema_file(file)
11355}
11356
11357fn read_schema_file(file: std::fs::File) -> Result<Schema> {
11358 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
11359 use std::io::Read;
11360
11361 let length = file.metadata()?.len();
11362 if length > MAX_SCHEMA_BYTES {
11363 return Err(MongrelError::ResourceLimitExceeded {
11364 resource: "schema bytes",
11365 requested: usize::try_from(length).unwrap_or(usize::MAX),
11366 limit: MAX_SCHEMA_BYTES as usize,
11367 });
11368 }
11369 let mut bytes = Vec::with_capacity(length as usize);
11370 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
11371 if bytes.len() as u64 != length {
11372 return Err(MongrelError::Schema(
11373 "schema length changed while reading".into(),
11374 ));
11375 }
11376 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
11377}
11378
11379fn preflight_standalone_open(
11380 dir: &Path,
11381 runs_root: Option<&crate::durable_file::DurableRoot>,
11382 idx_root: Option<&crate::durable_file::DurableRoot>,
11383 manifest: &Manifest,
11384 schema: &Schema,
11385 records: &[crate::wal::Record],
11386 kek: Option<Arc<Kek>>,
11387) -> Result<()> {
11388 crate::wal::validate_shared_transaction_framing(records)?;
11389 if manifest.schema_id > schema.schema_id
11390 || manifest.flushed_epoch > manifest.current_epoch
11391 || manifest.global_idx_epoch > manifest.current_epoch
11392 || manifest.next_row_id == u64::MAX
11393 || manifest.auto_inc_next < 0
11394 || manifest.auto_inc_next == i64::MAX
11395 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
11396 {
11397 return Err(MongrelError::InvalidArgument(
11398 "manifest counters or schema identity are invalid".into(),
11399 ));
11400 }
11401 let mut run_ids = HashSet::new();
11402 let mut maximum_row_id = None::<u64>;
11403 for run in &manifest.runs {
11404 if run.run_id >= u64::MAX as u128
11405 || !run_ids.insert(run.run_id)
11406 || run.epoch_created > manifest.current_epoch
11407 {
11408 return Err(MongrelError::InvalidArgument(
11409 "manifest contains an invalid or duplicate active run".into(),
11410 ));
11411 }
11412 let mut reader = match runs_root {
11413 Some(root) => RunReader::open_file(
11414 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
11415 schema.clone(),
11416 kek.clone(),
11417 )?,
11418 None => RunReader::open(
11419 dir.join(RUNS_DIR)
11420 .join(format!("r-{}.sr", run.run_id as u64)),
11421 schema.clone(),
11422 kek.clone(),
11423 )?,
11424 };
11425 let header = reader.header();
11426 if header.run_id != run.run_id
11427 || header.level != run.level
11428 || header.row_count != run.row_count
11429 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
11430 || header.is_uniform_epoch() && header.epoch_created != 0
11431 || header.schema_id > schema.schema_id
11432 {
11433 return Err(MongrelError::InvalidArgument(format!(
11434 "run {} differs from its manifest",
11435 run.run_id
11436 )));
11437 }
11438 if header.row_count != 0 {
11439 maximum_row_id = Some(
11440 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
11441 );
11442 }
11443 reader.validate_all_pages()?;
11444 }
11445 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
11446 return Err(MongrelError::InvalidArgument(
11447 "manifest next_row_id does not advance beyond persisted rows".into(),
11448 ));
11449 }
11450 for run in &manifest.retiring {
11451 if run.run_id >= u64::MAX as u128
11452 || run.retire_epoch > manifest.current_epoch
11453 || !run_ids.insert(run.run_id)
11454 {
11455 return Err(MongrelError::InvalidArgument(
11456 "manifest contains an invalid or duplicate retired run".into(),
11457 ));
11458 }
11459 }
11460 #[cfg(feature = "encryption")]
11461 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
11462 #[cfg(not(feature = "encryption"))]
11463 let idx_dek: Option<Zeroizing<[u8; DEK_LEN]>> = None;
11464 match idx_root {
11465 Some(root) => {
11466 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
11467 }
11468 None => {
11469 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
11470 }
11471 }
11472
11473 let committed = records
11474 .iter()
11475 .filter_map(|record| match record.op {
11476 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
11477 _ => None,
11478 })
11479 .collect::<HashMap<_, _>>();
11480 for record in records {
11481 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
11482 continue;
11483 };
11484 match &record.op {
11485 Op::Put { table_id, rows } => {
11486 if *table_id != manifest.table_id {
11487 return Err(MongrelError::CorruptWal {
11488 offset: record.seq.0,
11489 reason: format!(
11490 "private WAL record references table {table_id}, expected {}",
11491 manifest.table_id
11492 ),
11493 });
11494 }
11495 let rows: Vec<Row> =
11496 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
11497 offset: record.seq.0,
11498 reason: format!("committed Put payload could not be decoded: {error}"),
11499 })?;
11500 for row in rows {
11501 if row.deleted || row.row_id.0 == u64::MAX {
11502 return Err(MongrelError::CorruptWal {
11503 offset: record.seq.0,
11504 reason: "committed Put contains an invalid row identity".into(),
11505 });
11506 }
11507 let cells = row.columns.into_iter().collect::<Vec<_>>();
11508 schema
11509 .validate_values(&cells)
11510 .map_err(|error| MongrelError::CorruptWal {
11511 offset: record.seq.0,
11512 reason: format!("committed Put violates table schema: {error}"),
11513 })?;
11514 if schema.auto_increment_column().is_some_and(|column| {
11515 matches!(
11516 cells.iter().find(|(id, _)| *id == column.id),
11517 Some((_, Value::Int64(value))) if *value == i64::MAX
11518 )
11519 }) {
11520 return Err(MongrelError::CorruptWal {
11521 offset: record.seq.0,
11522 reason: "committed Put exhausts AUTO_INCREMENT".into(),
11523 });
11524 }
11525 }
11526 }
11527 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
11528 if *table_id != manifest.table_id =>
11529 {
11530 return Err(MongrelError::CorruptWal {
11531 offset: record.seq.0,
11532 reason: format!(
11533 "private WAL record references table {table_id}, expected {}",
11534 manifest.table_id
11535 ),
11536 });
11537 }
11538 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
11539 return Err(MongrelError::CorruptWal {
11540 offset: record.seq.0,
11541 reason: "private WAL contains shared spilled-run metadata".into(),
11542 });
11543 }
11544 _ => {}
11545 }
11546 }
11547 Ok(())
11548}
11549
11550fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
11551 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
11552}
11553
11554fn wal_segment_number(path: &Path) -> Option<u64> {
11555 path.file_stem()
11556 .and_then(|stem| stem.to_str())
11557 .and_then(|stem| stem.strip_prefix("seg-"))
11558 .and_then(|number| number.parse().ok())
11559}
11560
11561fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
11562 let n = list_wal_numbers(wal_dir)?;
11563 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
11564}
11565
11566fn next_wal_number(wal_dir: &Path) -> Result<u32> {
11567 list_wal_numbers(wal_dir)?
11568 .map(|maximum| {
11569 maximum
11570 .checked_add(1)
11571 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
11572 })
11573 .unwrap_or(Ok(0))
11574}
11575
11576fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
11577 let mut max_n = None;
11578 let entries = match std::fs::read_dir(wal_dir) {
11579 Ok(entries) => entries,
11580 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
11581 Err(error) => return Err(error.into()),
11582 };
11583 for entry in entries {
11584 let entry = entry?;
11585 let fname = entry.file_name();
11586 let Some(s) = fname.to_str() else {
11587 continue;
11588 };
11589 let Some(stripped) = s.strip_prefix("seg-") else {
11590 continue;
11591 };
11592 let Some(number) = stripped.strip_suffix(".wal") else {
11593 return Err(MongrelError::CorruptWal {
11594 offset: 0,
11595 reason: format!("malformed WAL segment name {s:?}"),
11596 });
11597 };
11598 let n = number
11599 .parse::<u32>()
11600 .map_err(|_| MongrelError::CorruptWal {
11601 offset: 0,
11602 reason: format!("malformed WAL segment name {s:?}"),
11603 })?;
11604 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
11605 return Err(MongrelError::CorruptWal {
11606 offset: n as u64,
11607 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
11608 });
11609 }
11610 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
11611 }
11612 Ok(max_n)
11613}