1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, EpochGuard, MaintenanceReceipt, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, IndexGeneration, MinHashIndex,
18 SparseIndex,
19};
20use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
21use crate::memtable::{Memtable, Row, Value};
22use crate::mutable_run::MutableRun;
23use crate::row_id_set::RowIdSet;
24use crate::rowid::{RowId, RowIdAllocator};
25use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
26use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
27use crate::txn::{GroupCommit, OwnedRow};
28use crate::wal::{Op, SharedWal, Wal};
29use crate::{MongrelError, Result};
30use arc_swap::ArcSwap;
31use std::cmp::Reverse;
32use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
33use std::path::{Path, PathBuf};
34use std::sync::atomic::AtomicBool;
35use std::sync::Arc;
36use zeroize::Zeroizing;
37
38pub const WAL_DIR: &str = "_wal";
39pub const RUNS_DIR: &str = "_runs";
40pub const CACHE_DIR: &str = "_cache";
41pub const META_DIR: &str = "_meta";
42pub const RCACHE_DIR: &str = "_rcache";
43pub const KEYS_FILENAME: &str = "keys";
44pub const SCHEMA_FILENAME: &str = "schema.json";
45
46fn derive_next_run_id(
47 dir: &Path,
48 runs_root: Option<&crate::durable_file::DurableRoot>,
49 active: &[RunRef],
50 retiring: &[crate::manifest::RetiredRun],
51) -> Result<u64> {
52 let mut maximum = 0_u64;
53 for run_id in active
54 .iter()
55 .map(|run| run.run_id)
56 .chain(retiring.iter().map(|run| run.run_id))
57 {
58 let run_id = u64::try_from(run_id)
59 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
60 maximum = maximum.max(run_id);
61 }
62 let names = match runs_root {
63 Some(root) => root.list_regular_files(".")?,
64 None => std::fs::read_dir(dir.join(RUNS_DIR))?
65 .map(|entry| entry.map(|entry| entry.file_name()))
66 .collect::<std::io::Result<Vec<_>>>()?,
67 };
68 for name in names {
69 let Some(name) = name.to_str() else {
70 continue;
71 };
72 let Some(digits) = name
73 .strip_prefix("r-")
74 .and_then(|name| name.strip_suffix(".sr"))
75 else {
76 continue;
77 };
78 let Ok(run_id) = digits.parse::<u64>() else {
79 continue;
80 };
81 if name == format!("r-{run_id}.sr") {
82 maximum = maximum.max(run_id);
83 }
84 }
85 maximum
86 .checked_add(1)
87 .map(|next| next.max(1))
88 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
89}
90
91enum ControlledVisibleCandidate {
92 Memory(Row),
93 Run(RunVisibleVersion),
94}
95
96impl ControlledVisibleCandidate {
97 fn row_id(&self) -> RowId {
98 match self {
99 Self::Memory(row) => row.row_id,
100 Self::Run(version) => version.row_id,
101 }
102 }
103
104 fn committed_epoch(&self) -> Epoch {
105 match self {
106 Self::Memory(row) => row.committed_epoch,
107 Self::Run(version) => version.committed_epoch,
108 }
109 }
110
111 fn deleted(&self) -> bool {
112 match self {
113 Self::Memory(row) => row.deleted,
114 Self::Run(version) => version.deleted,
115 }
116 }
117
118 fn commit_ts(&self) -> Option<mongreldb_types::hlc::HlcTimestamp> {
119 match self {
120 Self::Memory(row) => row.commit_ts,
121 Self::Run(version) => version.commit_ts,
124 }
125 }
126}
127
128enum ControlledVisibleCursor {
129 Memory(std::vec::IntoIter<Row>),
131 MemoryStreaming {
135 active: std::vec::IntoIter<Row>,
136 rest: std::collections::btree_map::IntoValues<RowId, Row>,
137 batch_cap: usize,
138 #[allow(dead_code)]
140 peak_active: usize,
141 #[allow(dead_code)]
143 total: usize,
144 },
145 Run(Box<RunVisibleVersionCursor>),
146 #[cfg(test)]
147 Synthetic {
148 next: u64,
149 end: u64,
150 },
151}
152
153const CONTROLLED_HOT_BATCH: usize = 256;
155
156struct ControlledVisibleSource {
157 cursor: ControlledVisibleCursor,
158 current: Option<ControlledVisibleCandidate>,
159}
160
161impl ControlledVisibleSource {
162 #[cfg(test)]
164 fn memory(rows: Vec<Row>) -> Self {
165 Self {
166 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
167 current: None,
168 }
169 }
170
171 fn memory_from_map(map: BTreeMap<RowId, Row>) -> Self {
175 let total = map.len();
176 let mut values = map.into_values();
177 if total > CONTROLLED_HOT_BATCH {
178 let active: Vec<Row> = values.by_ref().take(CONTROLLED_HOT_BATCH).collect();
179 let peak = active.len();
180 Self {
181 cursor: ControlledVisibleCursor::MemoryStreaming {
182 active: active.into_iter(),
183 rest: values,
184 batch_cap: CONTROLLED_HOT_BATCH,
185 peak_active: peak,
186 total,
187 },
188 current: None,
189 }
190 } else {
191 let active: Vec<Row> = values.collect();
192 Self {
193 cursor: ControlledVisibleCursor::Memory(active.into_iter()),
194 current: None,
195 }
196 }
197 }
198
199 fn run(cursor: RunVisibleVersionCursor) -> Self {
200 Self {
201 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
202 current: None,
203 }
204 }
205
206 #[cfg(test)]
207 fn synthetic(end: u64) -> Self {
208 Self {
209 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
210 current: None,
211 }
212 }
213
214 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
215 self.current = match &mut self.cursor {
216 ControlledVisibleCursor::Memory(rows) => {
217 rows.next().map(ControlledVisibleCandidate::Memory)
218 }
219 ControlledVisibleCursor::MemoryStreaming {
220 active,
221 rest,
222 batch_cap,
223 peak_active,
224 total: _,
225 } => {
226 if let Some(row) = active.next() {
227 Some(ControlledVisibleCandidate::Memory(row))
228 } else {
229 control.checkpoint()?;
230 let next_batch: Vec<Row> = rest.by_ref().take(*batch_cap).collect();
231 if next_batch.is_empty() {
232 None
233 } else {
234 *peak_active = (*peak_active).max(next_batch.len());
235 *active = next_batch.into_iter();
236 active.next().map(ControlledVisibleCandidate::Memory)
237 }
238 }
239 }
240 ControlledVisibleCursor::Run(cursor) => cursor
241 .next_visible_version(control)?
242 .map(ControlledVisibleCandidate::Run),
243 #[cfg(test)]
244 ControlledVisibleCursor::Synthetic { next, end } => {
245 if *next > *end {
246 None
247 } else {
248 let row = Row::new(RowId(*next), Epoch(1));
249 *next += 1;
250 Some(ControlledVisibleCandidate::Memory(row))
251 }
252 }
253 };
254 Ok(())
255 }
256
257 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate> {
258 let current = self.current.take().ok_or_else(|| {
259 MongrelError::Other("controlled visible source was not primed".into())
260 })?;
261 self.advance(control)?;
262 Ok(current)
263 }
264
265 fn materialize(
266 &mut self,
267 candidate: ControlledVisibleCandidate,
268 control: &crate::ExecutionControl,
269 ) -> Result<Row> {
270 match candidate {
271 ControlledVisibleCandidate::Memory(row) => Ok(row),
272 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
273 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
274 _ => Err(MongrelError::Other(
275 "run candidate escaped its controlled cursor".into(),
276 )),
277 },
278 }
279 }
280
281 #[cfg(test)]
282 fn is_streaming_memory(&self) -> bool {
283 matches!(self.cursor, ControlledVisibleCursor::MemoryStreaming { .. })
284 }
285
286 #[cfg(test)]
287 fn streaming_peak_active(&self) -> Option<usize> {
288 match &self.cursor {
289 ControlledVisibleCursor::MemoryStreaming { peak_active, .. } => Some(*peak_active),
290 _ => None,
291 }
292 }
293
294 #[cfg(test)]
295 fn streaming_total(&self) -> Option<usize> {
296 match &self.cursor {
297 ControlledVisibleCursor::MemoryStreaming { total, .. } => Some(*total),
298 _ => None,
299 }
300 }
301}
302
303fn merge_controlled_visible_sources(
304 sources: &mut [ControlledVisibleSource],
305 control: &crate::ExecutionControl,
306 mut expired: impl FnMut(&Row) -> bool,
307 mut visit: impl FnMut(Row) -> Result<()>,
308) -> Result<()> {
309 let mut heap = BinaryHeap::new();
310 for (source_index, source) in sources.iter_mut().enumerate() {
311 source.advance(control)?;
312 if let Some(candidate) = &source.current {
313 heap.push(Reverse((candidate.row_id(), source_index)));
314 }
315 }
316 let mut merged = 0_usize;
317 while let Some(Reverse((row_id, source_index))) = heap.pop() {
318 if merged.is_multiple_of(256) {
319 control.checkpoint()?;
320 }
321 merged += 1;
322 let mut best_source = source_index;
323 let mut best = sources[source_index].pop(control)?;
324 if let Some(next) = &sources[source_index].current {
325 heap.push(Reverse((next.row_id(), source_index)));
326 }
327 while heap
328 .peek()
329 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
330 {
331 let Some(Reverse((_, source_index))) = heap.pop() else {
332 break;
333 };
334 let candidate = sources[source_index].pop(control)?;
335 if Snapshot::version_is_newer(
339 candidate.committed_epoch(),
340 candidate.commit_ts(),
341 best.committed_epoch(),
342 best.commit_ts(),
343 ) {
344 best = candidate;
345 best_source = source_index;
346 }
347 if let Some(next) = &sources[source_index].current {
348 heap.push(Reverse((next.row_id(), source_index)));
349 }
350 }
351 if best.deleted() {
352 continue;
353 }
354 let row = sources[best_source].materialize(best, control)?;
355 if !expired(&row) {
356 visit(row)?;
357 }
358 }
359 control.checkpoint()
360}
361
362#[cfg(test)]
363mod controlled_visible_cursor_tests {
364 use super::*;
365
366 #[test]
367 fn streams_more_than_one_million_rows_without_a_source_cap() {
368 let control = crate::ExecutionControl::new(None);
369 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
370 let mut count = 0_u64;
371 let mut last = 0_u64;
372 merge_controlled_visible_sources(
373 &mut sources,
374 &control,
375 |_| false,
376 |row| {
377 count += 1;
378 assert!(row.row_id.0 > last);
379 last = row.row_id.0;
380 Ok(())
381 },
382 )
383 .unwrap();
384 assert_eq!(count, 1_000_001);
385 assert_eq!(last, 1_000_001);
386 }
387
388 #[test]
389 fn merge_orders_rows_and_honors_newest_tombstones() {
390 let control = crate::ExecutionControl::new(None);
391 let older = vec![
392 Row::new(RowId(1), Epoch(1)),
393 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
394 Row::new(RowId(4), Epoch(1)),
395 ];
396 let mut deleted = Row::new(RowId(1), Epoch(2));
397 deleted.deleted = true;
398 let newer = vec![
399 deleted,
400 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
401 Row::new(RowId(3), Epoch(2)),
402 ];
403 let mut sources = vec![
404 ControlledVisibleSource::memory(older),
405 ControlledVisibleSource::memory(newer),
406 ];
407 let mut rows = Vec::new();
408 merge_controlled_visible_sources(
409 &mut sources,
410 &control,
411 |_| false,
412 |row| {
413 rows.push(row);
414 Ok(())
415 },
416 )
417 .unwrap();
418 assert_eq!(
419 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
420 vec![2, 3, 4]
421 );
422 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
423 }
424
425 #[test]
426 fn controlled_merge_uses_hlc_authority_when_epochs_inverted() {
427 use mongreldb_types::hlc::HlcTimestamp;
428 let hlc_old = HlcTimestamp {
429 physical_micros: 100,
430 logical: 0,
431 node_tiebreaker: 1,
432 };
433 let hlc_new = HlcTimestamp {
434 physical_micros: 200,
435 logical: 0,
436 node_tiebreaker: 1,
437 };
438 let a = vec![
441 Row::new_with_hlc(RowId(1), Epoch(50), hlc_old)
442 .with_column(1, Value::Int64(999)),
443 ];
444 let b = vec![
445 Row::new_with_hlc(RowId(1), Epoch(1), hlc_new)
446 .with_column(1, Value::Int64(11)),
447 ];
448 let control = crate::ExecutionControl::new(None);
449 let mut sources = vec![
450 ControlledVisibleSource::memory(a),
451 ControlledVisibleSource::memory(b),
452 ];
453 let mut rows = Vec::new();
454 merge_controlled_visible_sources(
455 &mut sources,
456 &control,
457 |_| false,
458 |row| {
459 rows.push(row);
460 Ok(())
461 },
462 )
463 .unwrap();
464 assert_eq!(rows.len(), 1);
465 assert_eq!(rows[0].row_id.0, 1);
466 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(11)));
468 assert_eq!(rows[0].commit_ts, Some(hlc_new));
469 }
470
471 #[test]
472 fn controlled_merge_epoch_wins_when_one_side_lacks_hlc() {
473 use mongreldb_types::hlc::HlcTimestamp;
474 let hlc_old = HlcTimestamp {
475 physical_micros: 100,
476 logical: 0,
477 node_tiebreaker: 1,
478 };
479 let a = vec![
482 Row::new_with_hlc(RowId(1), Epoch(10), hlc_old)
483 .with_column(1, Value::Int64(10)),
484 ];
485 let b = vec![
486 Row::new(RowId(1), Epoch(5)).with_column(1, Value::Int64(5)),
487 ];
488 let control = crate::ExecutionControl::new(None);
489 let mut sources = vec![
490 ControlledVisibleSource::memory(a),
491 ControlledVisibleSource::memory(b),
492 ];
493 let mut rows = Vec::new();
494 merge_controlled_visible_sources(
495 &mut sources,
496 &control,
497 |_| false,
498 |row| {
499 rows.push(row);
500 Ok(())
501 },
502 )
503 .unwrap();
504 assert_eq!(rows.len(), 1);
505 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(10)));
506 }
507
508 #[test]
511 fn controlled_merge_memory_vs_run_uses_hlc_when_run_stamped() {
512 use mongreldb_types::hlc::HlcTimestamp;
513 use crate::sorted_run::{RunReader, RunWriter};
514 use tempfile::tempdir;
515
516 let hlc_old = HlcTimestamp {
517 physical_micros: 100,
518 logical: 0,
519 node_tiebreaker: 1,
520 };
521 let hlc_new = HlcTimestamp {
522 physical_micros: 200,
523 logical: 0,
524 node_tiebreaker: 1,
525 };
526 let schema = Schema {
527 schema_id: 1,
528 columns: vec![ColumnDef {
529 id: 1,
530 name: "v".into(),
531 ty: TypeId::Int64,
532 flags: ColumnFlags::empty(),
533 default_value: None,
534 embedding_source: None,
535 }],
536 indexes: vec![],
537 colocation: vec![],
538 constraints: Default::default(),
539 clustered: false,
540 };
541 let dir = tempdir().unwrap();
542 let path = dir.path().join("r-hlc.sr");
543 let run_rows = vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old)
545 .with_column(1, Value::Int64(999))];
546 RunWriter::new(&schema, 1, Epoch(50), 0)
547 .write(&path, &run_rows)
548 .unwrap();
549 let reader = RunReader::open(&path, schema, None).unwrap();
550 assert!(reader.has_column(crate::sorted_run::SYS_COMMIT_TS));
551
552 let mem = vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new)
554 .with_column(1, Value::Int64(11))];
555 let control = crate::ExecutionControl::new(None);
556 let mut sources = vec![
557 ControlledVisibleSource::memory(mem),
558 ControlledVisibleSource::run(
559 reader
560 .into_visible_version_cursor(Epoch(u64::MAX))
561 .unwrap(),
562 ),
563 ];
564 let mut rows = Vec::new();
565 merge_controlled_visible_sources(
566 &mut sources,
567 &control,
568 |_| false,
569 |row| {
570 rows.push(row);
571 Ok(())
572 },
573 )
574 .unwrap();
575 assert_eq!(rows.len(), 1);
576 assert_eq!(
577 rows[0].columns.get(&1),
578 Some(&Value::Int64(11)),
579 "HLC-newer memory version must beat epoch-newer run"
580 );
581 assert_eq!(rows[0].commit_ts, Some(hlc_new));
582 }
583
584 #[test]
585 fn controlled_memory_source_streams_large_overlays_without_full_active_vec() {
586 let mut map = BTreeMap::new();
587 for i in 1..=500u64 {
588 map.insert(
589 RowId(i),
590 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
591 );
592 }
593 let source = ControlledVisibleSource::memory_from_map(map);
594 assert!(
595 source.is_streaming_memory(),
596 "overlays larger than CONTROLLED_HOT_BATCH must use streaming cursor"
597 );
598 assert_eq!(source.streaming_total(), Some(500));
599 assert!(
601 source.streaming_peak_active().unwrap_or(usize::MAX) <= CONTROLLED_HOT_BATCH,
602 "peak active buffer must be <= CONTROLLED_HOT_BATCH"
603 );
604
605 let control = crate::ExecutionControl::new(None);
607 let mut sources = vec![ControlledVisibleSource::memory_from_map({
608 let mut m = BTreeMap::new();
609 for i in 1..=500u64 {
610 m.insert(
611 RowId(i),
612 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
613 );
614 }
615 m
616 })];
617 let mut n = 0usize;
618 merge_controlled_visible_sources(
619 &mut sources,
620 &control,
621 |_| false,
622 |_| {
623 n += 1;
624 Ok(())
625 },
626 )
627 .unwrap();
628 assert_eq!(n, 500);
629 assert!(
630 sources[0].streaming_peak_active().unwrap_or(0) <= CONTROLLED_HOT_BATCH,
631 "after full drain peak active still <= batch cap"
632 );
633 }
634}
635
636fn iso_now_bytes() -> Vec<u8> {
639 let secs = std::time::SystemTime::now()
640 .duration_since(std::time::UNIX_EPOCH)
641 .map(|d| d.as_secs() as i64)
642 .unwrap_or(0);
643 let days = secs.div_euclid(86_400);
644 let rem = secs.rem_euclid(86_400);
645 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
646 let (year, month, day) = civil_from_days(days);
647 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
648}
649
650pub(crate) fn unix_nanos_now() -> i64 {
651 std::time::SystemTime::now()
652 .duration_since(std::time::UNIX_EPOCH)
653 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
654 .unwrap_or(0)
655}
656
657fn ann_candidate_cap(
658 index_len: usize,
659 context: Option<&crate::query::AiExecutionContext>,
660) -> usize {
661 index_len
662 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
663 .min(context.map_or(
664 crate::query::MAX_RAW_INDEX_CANDIDATES,
665 crate::query::AiExecutionContext::max_fused_candidates,
666 ))
667}
668
669#[cfg(test)]
670mod ann_candidate_cap_tests {
671 use super::*;
672
673 #[test]
674 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
675 assert_eq!(
676 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
677 crate::query::MAX_RAW_INDEX_CANDIDATES,
678 );
679 let context = crate::query::AiExecutionContext::with_limits(
680 std::time::Duration::from_secs(1),
681 usize::MAX,
682 17,
683 );
684 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
685 }
686}
687
688fn civil_from_days(z: i64) -> (i64, u32, u32) {
689 let z = z + 719_468;
690 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
691 let doe = z - era * 146_097;
692 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
693 let y = yoe + era * 400;
694 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
695 let mp = (5 * doy + 2) / 153;
696 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
697 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
698 (if m <= 2 { y + 1 } else { y }, m, d)
699}
700
701pub fn clustered_row_id(primary_key: &Value) -> RowId {
708 let mut hash: u64 = 0xcbf29ce484222325;
709 for byte in primary_key.encode_key() {
710 hash ^= u64::from(byte);
711 hash = hash.wrapping_mul(0x100000001b3);
712 }
713 RowId(hash.max(1))
714}
715
716const 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;
723
724#[derive(Clone, Copy, Debug)]
739struct AutoIncState {
740 column_id: u16,
741 next: i64,
742 seeded: bool,
743}
744
745pub(crate) struct RecoveryMetadataPlan {
746 live_count: u64,
747 auto_inc: Option<AutoIncState>,
748 changed: bool,
749}
750
751type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
752
753fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
756 schema.auto_increment_column().map(|c| AutoIncState {
757 column_id: c.id,
758 next: 0,
759 seeded: false,
760 })
761}
762
763#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
776pub enum IndexBuildPolicy {
777 #[default]
779 Deferred,
780 Eager,
782}
783
784#[derive(Clone)]
785struct ReversePkSegment {
786 values: HashMap<RowId, Vec<u8>>,
787 removed: HashSet<RowId>,
788}
789
790#[derive(Clone)]
791struct ReversePkMap {
792 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
793 active: ReversePkSegment,
794}
795
796impl ReversePkMap {
797 fn new() -> Self {
798 Self {
799 frozen: Arc::new(Vec::new()),
800 active: ReversePkSegment {
801 values: HashMap::new(),
802 removed: HashSet::new(),
803 },
804 }
805 }
806
807 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
808 let mut map = Self::new();
809 map.active.values.extend(entries);
810 map
811 }
812
813 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
814 self.active.removed.remove(&row_id);
815 self.active.values.insert(row_id, key);
816 }
817
818 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
819 if let Some(key) = self.active.values.get(row_id) {
820 return Some(key);
821 }
822 if self.active.removed.contains(row_id) {
823 return None;
824 }
825 for segment in self.frozen.iter().rev() {
826 if let Some(key) = segment.values.get(row_id) {
827 return Some(key);
828 }
829 if segment.removed.contains(row_id) {
830 return None;
831 }
832 }
833 None
834 }
835
836 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
837 let previous = self.get(row_id).cloned();
838 self.active.values.remove(row_id);
839 self.active.removed.insert(*row_id);
840 previous
841 }
842
843 fn clear(&mut self) {
844 *self = Self::new();
845 }
846
847 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
848 let mut entries = HashMap::new();
849 for segment in self
850 .frozen
851 .iter()
852 .map(Arc::as_ref)
853 .chain(std::iter::once(&self.active))
854 {
855 for row_id in &segment.removed {
856 entries.remove(row_id);
857 }
858 entries.extend(
859 segment
860 .values
861 .iter()
862 .map(|(row_id, key)| (*row_id, key.clone())),
863 );
864 }
865 entries
866 }
867
868 fn seal(&mut self) {
869 if self.active.values.is_empty() && self.active.removed.is_empty() {
870 return;
871 }
872 let active = std::mem::replace(
873 &mut self.active,
874 ReversePkSegment {
875 values: HashMap::new(),
876 removed: HashSet::new(),
877 },
878 );
879 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
880 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
881 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
882 values: self.entries(),
883 removed: HashSet::new(),
884 })]);
885 }
886 }
887}
888
889#[derive(Clone)]
902pub struct ReadGeneration {
903 schema: Arc<Schema>,
904 base_runs: Arc<Vec<RunRef>>,
905 deltas: TableDeltas,
906 indexes: Arc<IndexGeneration>,
907 visible_through: Epoch,
908}
909
910pub(crate) enum SecondaryIndexArtifact {
914 Bitmap(u16, BitmapIndex),
915 LearnedRange(u16, ColumnLearnedRange),
916 Fm(u16, Box<FmIndex>),
917 Ann(u16, Box<AnnIndex>),
918 Sparse(u16, SparseIndex),
919 MinHash(u16, MinHashIndex),
920}
921
922#[derive(Clone)]
928pub struct TableDeltas {
929 memtable: Memtable,
930 mutable_run: MutableRun,
931 hot: HotIndex,
932 pk_by_row: ReversePkMap,
933}
934
935impl ReadGeneration {
936 fn empty(schema: &Schema) -> Self {
939 Self {
940 schema: Arc::new(schema.clone()),
941 base_runs: Arc::new(Vec::new()),
942 deltas: TableDeltas {
943 memtable: Memtable::new(),
944 mutable_run: MutableRun::new(),
945 hot: HotIndex::new(),
946 pk_by_row: ReversePkMap::new(),
947 },
948 indexes: Arc::new(IndexGeneration::default()),
949 visible_through: Epoch(0),
950 }
951 }
952
953 pub fn schema(&self) -> &Arc<Schema> {
955 &self.schema
956 }
957
958 pub fn base_runs(&self) -> &[RunRef] {
960 &self.base_runs
961 }
962
963 pub fn indexes(&self) -> &Arc<IndexGeneration> {
965 &self.indexes
966 }
967
968 pub fn visible_through(&self) -> Epoch {
970 self.visible_through
971 }
972
973 pub fn deltas(&self) -> &TableDeltas {
977 &self.deltas
978 }
979}
980
981impl TableDeltas {
982 pub fn approx_bytes(&self) -> u64 {
985 self.memtable
986 .approx_bytes()
987 .saturating_add(self.mutable_run.approx_bytes())
988 }
989
990 pub fn memtable_len(&self) -> usize {
992 self.memtable.len()
993 }
994
995 pub fn mutable_run_len(&self) -> usize {
997 self.mutable_run.len()
998 }
999
1000 pub fn hot_len(&self) -> usize {
1002 self.hot.len()
1003 }
1004
1005 pub fn reverse_pk_len(&self) -> usize {
1007 self.pk_by_row.entries().len()
1008 }
1009}
1010
1011#[derive(Clone)]
1013pub struct Table {
1014 dir: PathBuf,
1015 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1016 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
1017 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
1018 table_id: u64,
1019 name: String,
1023 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1028 read_only: bool,
1031 durable_commit_failed: bool,
1035 wal: WalSink,
1036 memtable: Memtable,
1037 mutable_run: MutableRun,
1042 mutable_run_spill_bytes: u64,
1044 compaction_zstd_level: i32,
1047 allocator: RowIdAllocator,
1048 epoch: Arc<EpochAuthority>,
1049 data_generation: u64,
1052 schema: Schema,
1053 hot: HotIndex,
1054 kek: Option<Arc<Kek>>,
1057 column_keys: HashMap<u16, ([u8; 32], u8)>,
1061 run_refs: Vec<RunRef>,
1062 retiring: Vec<crate::manifest::RetiredRun>,
1065 next_run_id: u64,
1066 sync_byte_threshold: u64,
1067 current_txn_id: u64,
1072 pending_private_mutations: bool,
1076 bitmap: HashMap<u16, BitmapIndex>,
1077 ann: HashMap<u16, AnnIndex>,
1078 fm: HashMap<u16, FmIndex>,
1079 sparse: HashMap<u16, SparseIndex>,
1080 minhash: HashMap<u16, MinHashIndex>,
1081 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
1084 pk_by_row: ReversePkMap,
1086 pinned: BTreeMap<Epoch, usize>,
1089 pub(crate) live_count: u64,
1092 reservoir: crate::reservoir::Reservoir,
1095 reservoir_complete: bool,
1103 had_deletes: bool,
1107 recent_delete_preimages: HashMap<Vec<u8>, Row>,
1112 run_row_id_ranges: HashMap<u128, (u64, u64)>,
1115 agg_cache: Arc<HashMap<u64, CachedAgg>>,
1119 global_idx_epoch: u64,
1123 indexes_complete: bool,
1128 index_build_policy: IndexBuildPolicy,
1130 pk_by_row_complete: bool,
1137 flushed_epoch: u64,
1140 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1143 snapshots: Arc<crate::retention::SnapshotRegistry>,
1146 commit_lock: Arc<parking_lot::Mutex<()>>,
1148 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1151 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
1161 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
1170 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1172 pending_delete_rids: roaring::RoaringBitmap,
1175 pending_put_cols: std::collections::HashSet<u16>,
1178 pending_rows: Vec<Row>,
1184 pending_rows_auto_inc: Vec<bool>,
1185 pending_dels: Vec<RowId>,
1188 pending_truncate: Option<Epoch>,
1192 auto_inc: Option<AutoIncState>,
1195 ttl: Option<TtlPolicy>,
1198 pins: Arc<crate::retention::PinRegistry>,
1205 published: Arc<ArcSwap<ReadGeneration>>,
1210 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
1215 lookup_metrics: LookupMetrics,
1220}
1221
1222#[derive(Debug, Default)]
1223struct LookupMetrics {
1224 hot_lookup_hit: std::sync::atomic::AtomicU64,
1225 hot_lookup_fallback: std::sync::atomic::AtomicU64,
1226 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64,
1227 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64,
1228 result_cache_memory_hit: std::sync::atomic::AtomicU64,
1229 result_cache_disk_hit: std::sync::atomic::AtomicU64,
1230 result_cache_miss: std::sync::atomic::AtomicU64,
1231 result_cache_persistent_write_us: std::sync::atomic::AtomicU64,
1232 get_run_opened: std::sync::atomic::AtomicU64,
1234 get_run_skipped: std::sync::atomic::AtomicU64,
1235}
1236
1237impl Clone for LookupMetrics {
1238 fn clone(&self) -> Self {
1239 Self {
1240 hot_lookup_hit: std::sync::atomic::AtomicU64::new(
1241 self.hot_lookup_hit.load(std::sync::atomic::Ordering::Relaxed),
1242 ),
1243 hot_lookup_fallback: std::sync::atomic::AtomicU64::new(
1244 self.hot_lookup_fallback.load(std::sync::atomic::Ordering::Relaxed),
1245 ),
1246 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64::new(
1247 self.hot_lookup_fallback_overlay_rows
1248 .load(std::sync::atomic::Ordering::Relaxed),
1249 ),
1250 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64::new(
1251 self.hot_lookup_fallback_runs.load(std::sync::atomic::Ordering::Relaxed),
1252 ),
1253 result_cache_memory_hit: std::sync::atomic::AtomicU64::new(
1254 self.result_cache_memory_hit
1255 .load(std::sync::atomic::Ordering::Relaxed),
1256 ),
1257 result_cache_disk_hit: std::sync::atomic::AtomicU64::new(
1258 self.result_cache_disk_hit
1259 .load(std::sync::atomic::Ordering::Relaxed),
1260 ),
1261 result_cache_miss: std::sync::atomic::AtomicU64::new(
1262 self.result_cache_miss.load(std::sync::atomic::Ordering::Relaxed),
1263 ),
1264 result_cache_persistent_write_us: std::sync::atomic::AtomicU64::new(
1265 self.result_cache_persistent_write_us
1266 .load(std::sync::atomic::Ordering::Relaxed),
1267 ),
1268 get_run_opened: std::sync::atomic::AtomicU64::new(
1269 self.get_run_opened
1270 .load(std::sync::atomic::Ordering::Relaxed),
1271 ),
1272 get_run_skipped: std::sync::atomic::AtomicU64::new(
1273 self.get_run_skipped
1274 .load(std::sync::atomic::Ordering::Relaxed),
1275 ),
1276 }
1277 }
1278}
1279
1280impl LookupMetrics {
1281 fn snapshot(&self) -> LookupMetricsSnapshot {
1282 LookupMetricsSnapshot {
1283 hot_lookup_hit: self.hot_lookup_hit.load(std::sync::atomic::Ordering::Relaxed),
1284 hot_lookup_fallback: self.hot_lookup_fallback.load(std::sync::atomic::Ordering::Relaxed),
1285 hot_lookup_fallback_overlay_rows: self
1286 .hot_lookup_fallback_overlay_rows
1287 .load(std::sync::atomic::Ordering::Relaxed),
1288 hot_lookup_fallback_runs: self
1289 .hot_lookup_fallback_runs
1290 .load(std::sync::atomic::Ordering::Relaxed),
1291 result_cache_memory_hit: self
1292 .result_cache_memory_hit
1293 .load(std::sync::atomic::Ordering::Relaxed),
1294 result_cache_disk_hit: self
1295 .result_cache_disk_hit
1296 .load(std::sync::atomic::Ordering::Relaxed),
1297 result_cache_miss: self
1298 .result_cache_miss
1299 .load(std::sync::atomic::Ordering::Relaxed),
1300 result_cache_persistent_write_us: self
1301 .result_cache_persistent_write_us
1302 .load(std::sync::atomic::Ordering::Relaxed),
1303 get_run_opened: self
1304 .get_run_opened
1305 .load(std::sync::atomic::Ordering::Relaxed),
1306 get_run_skipped: self
1307 .get_run_skipped
1308 .load(std::sync::atomic::Ordering::Relaxed),
1309 }
1310 }
1311}
1312
1313#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
1316pub struct LookupMetricsSnapshot {
1317 pub hot_lookup_hit: u64,
1318 pub hot_lookup_fallback: u64,
1319 pub hot_lookup_fallback_overlay_rows: u64,
1320 pub hot_lookup_fallback_runs: u64,
1321 pub result_cache_memory_hit: u64,
1322 pub result_cache_disk_hit: u64,
1323 pub result_cache_miss: u64,
1324 pub result_cache_persistent_write_us: u64,
1325 pub get_run_opened: u64,
1326 pub get_run_skipped: u64,
1327}
1328
1329const _: () = {
1336 const fn assert_sync<T: ?Sized + Sync>() {}
1337 assert_sync::<Table>();
1338};
1339
1340enum CachedData {
1346 Rows(Arc<Vec<Row>>),
1347 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
1348}
1349
1350impl CachedData {
1351 fn approx_bytes(&self) -> u64 {
1352 match self {
1353 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
1354 CachedData::Columns(c) => c
1355 .iter()
1356 .map(|(_, c)| c.approx_bytes())
1357 .sum::<u64>()
1358 .saturating_add(c.len() as u64 * 16),
1359 }
1360 }
1361}
1362
1363struct CachedEntry {
1367 data: CachedData,
1368 footprint: roaring::RoaringBitmap,
1369 condition_cols: Vec<u16>,
1370}
1371
1372struct ResultCache {
1383 entries: std::collections::HashMap<u64, CachedEntry>,
1384 order: BTreeSet<(u64, u64)>,
1385 generations: HashMap<u64, u64>,
1386 next_generation: u64,
1387 condition_index: HashMap<u16, HashSet<u64>>,
1390 empty_footprint_entries: HashSet<u64>,
1393 bytes: u64,
1394 max_bytes: u64,
1395 dir: Option<std::path::PathBuf>,
1396 #[allow(dead_code)]
1397 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1398 memory_hit: std::sync::atomic::AtomicU64,
1402 disk_hit: std::sync::atomic::AtomicU64,
1403 miss: std::sync::atomic::AtomicU64,
1404 persistent_write_us: std::sync::atomic::AtomicU64,
1405 persist_min_bytes: u64,
1410}
1411
1412#[derive(serde::Serialize, serde::Deserialize)]
1414struct SerializedEntry {
1415 condition_cols: Vec<u16>,
1416 footprint_bits: Vec<u32>,
1417 data: SerializedData,
1418}
1419
1420#[derive(serde::Serialize, serde::Deserialize)]
1421enum SerializedData {
1422 Rows(Vec<Row>),
1423 Columns(Vec<(u16, columnar::NativeColumn)>),
1424}
1425
1426#[derive(serde::Serialize)]
1427struct SerializedEntryRef<'a> {
1428 condition_cols: &'a [u16],
1429 footprint_bits: Vec<u32>,
1430 data: SerializedDataRef<'a>,
1431}
1432
1433#[derive(serde::Serialize)]
1434enum SerializedDataRef<'a> {
1435 Rows(&'a [Row]),
1436 Columns(&'a [(u16, columnar::NativeColumn)]),
1437}
1438
1439impl<'a> SerializedEntryRef<'a> {
1440 fn from_entry(entry: &'a CachedEntry) -> Self {
1441 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
1442 let data = match &entry.data {
1443 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
1444 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
1445 };
1446 Self {
1447 condition_cols: &entry.condition_cols,
1448 footprint_bits,
1449 data,
1450 }
1451 }
1452}
1453
1454impl SerializedEntry {
1455 fn into_entry(self) -> Option<CachedEntry> {
1456 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1457 let data = match self.data {
1458 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1459 SerializedData::Columns(c) => {
1460 if !c.iter().all(|(_, col)| col.validate()) {
1463 return None;
1464 }
1465 CachedData::Columns(Arc::new(c))
1466 }
1467 };
1468 Some(CachedEntry {
1469 data,
1470 footprint,
1471 condition_cols: self.condition_cols,
1472 })
1473 }
1474}
1475
1476impl ResultCache {
1477 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1478
1479 fn new() -> Self {
1480 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1481 }
1482
1483 fn with_max_bytes(max_bytes: u64) -> Self {
1484 Self {
1485 entries: std::collections::HashMap::new(),
1486 order: BTreeSet::new(),
1487 generations: HashMap::new(),
1488 next_generation: 0,
1489 condition_index: HashMap::new(),
1490 empty_footprint_entries: HashSet::new(),
1491 bytes: 0,
1492 max_bytes,
1493 dir: None,
1494 cache_dek: None,
1495 memory_hit: std::sync::atomic::AtomicU64::new(0),
1496 disk_hit: std::sync::atomic::AtomicU64::new(0),
1497 miss: std::sync::atomic::AtomicU64::new(0),
1498 persistent_write_us: std::sync::atomic::AtomicU64::new(0),
1499 persist_min_bytes: 4 * 1024,
1502 }
1503 }
1504
1505 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1506 let _ = std::fs::create_dir_all(&dir);
1507 self.dir = Some(dir);
1508 self
1509 }
1510
1511 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1512 self.cache_dek = dek;
1513 self
1514 }
1515
1516 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1517 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1518 }
1519
1520 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1524 let Some(path) = self.disk_path(key) else {
1525 return;
1526 };
1527 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1528 Ok(s) => s,
1529 Err(_) => return,
1530 };
1531 let on_disk = if let Some(dek) = &self.cache_dek {
1533 match self.encrypt_cache(&serialized, dek) {
1534 Some(b) => b,
1535 None => return,
1536 }
1537 } else {
1538 serialized
1539 };
1540 let tmp = path.with_extension("tmp");
1541 use std::io::Write;
1542 let write = || -> std::io::Result<()> {
1543 let mut f = std::fs::File::create(&tmp)?;
1544 f.write_all(&on_disk)?;
1545 f.flush()?;
1546 Ok(())
1547 };
1548 if write().is_err() {
1549 let _ = std::fs::remove_file(&tmp);
1550 return;
1551 }
1552 let _ = std::fs::rename(&tmp, &path);
1553 }
1554
1555 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1557 let path = self.disk_path(key)?;
1558 let bytes = std::fs::read(&path).ok()?;
1559 let plaintext = if let Some(dek) = &self.cache_dek {
1560 self.decrypt_cache(&bytes, dek)?
1561 } else {
1562 bytes
1563 };
1564 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1565 serialized.into_entry()
1566 }
1567
1568 fn remove_from_disk(&self, key: u64) {
1570 if let Some(path) = self.disk_path(key) {
1571 let _ = std::fs::remove_file(&path);
1572 }
1573 }
1574
1575 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1577 use crate::encryption::Cipher;
1578 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1579 let mut nonce = [0u8; 12];
1580 crate::encryption::fill_random(&mut nonce).ok()?;
1581 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1582 let mut out = Vec::with_capacity(12 + ct.len());
1583 out.extend_from_slice(&nonce);
1584 out.extend_from_slice(&ct);
1585 Some(out)
1586 }
1587
1588 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1590 use crate::encryption::Cipher;
1591 if bytes.len() < 28 {
1592 return None;
1593 }
1594 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1595 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1596 let ct = &bytes[12..];
1597 cipher.decrypt_page(&nonce, ct).ok()
1598 }
1599
1600 fn load_persistent(&mut self) {
1603 let Some(dir) = self.dir.as_ref().cloned() else {
1604 return;
1605 };
1606 let entries = match std::fs::read_dir(&dir) {
1607 Ok(e) => e,
1608 Err(_) => return,
1609 };
1610 for entry in entries.flatten() {
1611 let path = entry.path();
1612 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1614 let _ = std::fs::remove_file(&path);
1615 continue;
1616 }
1617 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1618 continue;
1619 }
1620 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1621 Some(s) => s,
1622 None => continue,
1623 };
1624 let key = match u64::from_str_radix(stem, 16) {
1625 Ok(k) => k,
1626 Err(_) => continue,
1627 };
1628 let bytes = match std::fs::read(&path) {
1629 Ok(b) => b,
1630 Err(_) => continue,
1631 };
1632 let plaintext = if let Some(dek) = &self.cache_dek {
1634 match self.decrypt_cache(&bytes, dek) {
1635 Some(p) => p,
1636 None => {
1637 let _ = std::fs::remove_file(&path);
1638 continue;
1639 }
1640 }
1641 } else {
1642 bytes
1643 };
1644 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1645 Ok(serialized) => {
1646 if let Some(entry) = serialized.into_entry() {
1647 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1648 self.index_entry(key, &entry);
1649 self.entries.insert(key, entry);
1650 self.touch(key);
1651 } else {
1652 let _ = std::fs::remove_file(&path);
1653 }
1654 }
1655 Err(_) => {
1656 let _ = std::fs::remove_file(&path);
1657 }
1658 }
1659 }
1660 self.evict();
1661 }
1662
1663 fn set_max_bytes(&mut self, max_bytes: u64) {
1664 self.max_bytes = max_bytes;
1665 self.evict();
1666 }
1667
1668 fn touch(&mut self, key: u64) {
1671 if !self.entries.contains_key(&key) {
1672 return;
1673 }
1674 if let Some(previous) = self.generations.remove(&key) {
1675 self.order.remove(&(previous, key));
1676 }
1677 let generation = self.allocate_generation();
1678 self.generations.insert(key, generation);
1679 self.order.insert((generation, key));
1680 }
1681
1682 fn untrack(&mut self, key: u64) {
1683 if let Some(generation) = self.generations.remove(&key) {
1684 self.order.remove(&(generation, key));
1685 }
1686 }
1687
1688 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
1689 for column in &entry.condition_cols {
1690 self.condition_index.entry(*column).or_default().insert(key);
1691 }
1692 if entry.footprint.is_empty() {
1693 self.empty_footprint_entries.insert(key);
1694 }
1695 }
1696
1697 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
1698 for column in &entry.condition_cols {
1699 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
1700 keys.remove(&key);
1701 keys.is_empty()
1702 });
1703 if remove_column {
1704 self.condition_index.remove(column);
1705 }
1706 }
1707 if entry.footprint.is_empty() {
1708 self.empty_footprint_entries.remove(&key);
1709 }
1710 }
1711
1712 fn allocate_generation(&mut self) -> u64 {
1713 if self.next_generation == u64::MAX {
1714 self.rebase_generations();
1715 }
1716 let generation = self.next_generation;
1717 self.next_generation += 1;
1718 generation
1719 }
1720
1721 fn rebase_generations(&mut self) {
1722 let ordered = std::mem::take(&mut self.order);
1723 self.generations.clear();
1724 for (generation, (_, key)) in ordered.into_iter().enumerate() {
1725 let generation = generation as u64;
1726 self.generations.insert(key, generation);
1727 self.order.insert((generation, key));
1728 }
1729 self.next_generation = self.order.len() as u64;
1730 }
1731
1732 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1733 let res = self.entries.get(&key).and_then(|e| match &e.data {
1734 CachedData::Rows(r) => Some(r.clone()),
1735 CachedData::Columns(_) => None,
1736 });
1737 if res.is_some() {
1738 self.touch(key);
1739 self.memory_hit
1740 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1741 return res;
1742 }
1743 if let Some(entry) = self.load_from_disk(key) {
1745 let res = match &entry.data {
1746 CachedData::Rows(r) => Some(r.clone()),
1747 CachedData::Columns(_) => None,
1748 };
1749 if res.is_some() {
1750 let approx = entry.data.approx_bytes();
1751 self.bytes = self.bytes.saturating_add(approx);
1752 self.index_entry(key, &entry);
1753 self.entries.insert(key, entry);
1754 self.touch(key);
1755 self.evict();
1756 self.disk_hit
1757 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1758 return res;
1759 }
1760 }
1761 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1762 None
1763 }
1764
1765 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1766 let res = self.entries.get(&key).and_then(|e| match &e.data {
1767 CachedData::Columns(c) => Some(c.clone()),
1768 CachedData::Rows(_) => None,
1769 });
1770 if res.is_some() {
1771 self.touch(key);
1772 self.memory_hit
1773 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1774 return res;
1775 }
1776 if let Some(entry) = self.load_from_disk(key) {
1778 let res = match &entry.data {
1779 CachedData::Columns(c) => Some(c.clone()),
1780 CachedData::Rows(_) => None,
1781 };
1782 if res.is_some() {
1783 let approx = entry.data.approx_bytes();
1784 self.bytes = self.bytes.saturating_add(approx);
1785 self.index_entry(key, &entry);
1786 self.entries.insert(key, entry);
1787 self.touch(key);
1788 self.evict();
1789 self.disk_hit
1790 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1791 return res;
1792 }
1793 }
1794 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1795 None
1796 }
1797
1798 fn insert(&mut self, key: u64, entry: CachedEntry) {
1799 let approx = entry.data.approx_bytes();
1800 if let Some(previous) = self.entries.remove(&key) {
1801 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
1802 self.unindex_entry(key, &previous);
1803 self.untrack(key);
1804 }
1805 if self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes) {
1809 let write_start = std::time::Instant::now();
1810 self.store_to_disk(key, &entry);
1811 let write_us = write_start.elapsed().as_micros() as u64;
1812 self.persistent_write_us
1813 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
1814 }
1815 self.bytes = self.bytes.saturating_add(approx);
1816 self.index_entry(key, &entry);
1817 self.entries.insert(key, entry);
1818 self.touch(key);
1819 self.evict();
1820 }
1821
1822 #[cfg(test)]
1823 fn set_persist_min_bytes(&mut self, min: u64) {
1824 self.persist_min_bytes = min;
1825 }
1826
1827 #[cfg(test)]
1829 fn would_persist(&self, approx: u64) -> bool {
1830 self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
1831 }
1832
1833 fn cache_counters(&self) -> (u64, u64, u64, u64) {
1836 (
1837 self.memory_hit
1838 .load(std::sync::atomic::Ordering::Relaxed),
1839 self.disk_hit
1840 .load(std::sync::atomic::Ordering::Relaxed),
1841 self.miss.load(std::sync::atomic::Ordering::Relaxed),
1842 self.persistent_write_us
1843 .load(std::sync::atomic::Ordering::Relaxed),
1844 )
1845 }
1846
1847 fn invalidate(
1856 &mut self,
1857 delete_rids: &roaring::RoaringBitmap,
1858 put_cols: &std::collections::HashSet<u16>,
1859 ) {
1860 if self.entries.is_empty() {
1861 return;
1862 }
1863 let has_deletes = !delete_rids.is_empty();
1864 let mut to_remove = HashSet::new();
1865
1866 for column in put_cols {
1870 if let Some(keys) = self.condition_index.get(column) {
1871 to_remove.extend(keys.iter().copied());
1872 }
1873 }
1874
1875 if has_deletes {
1876 to_remove.extend(self.empty_footprint_entries.iter().copied());
1880 for (&key, entry) in &self.entries {
1881 if !to_remove.contains(&key)
1882 && !entry.footprint.is_empty()
1883 && entry.footprint.intersection_len(delete_rids) > 0
1884 {
1885 to_remove.insert(key);
1886 }
1887 }
1888 }
1889
1890 for key in to_remove {
1891 if let Some(entry) = self.entries.remove(&key) {
1892 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
1893 self.unindex_entry(key, &entry);
1894 }
1895 self.remove_from_disk(key);
1896 self.untrack(key);
1897 }
1898 }
1899
1900 fn clear(&mut self) {
1901 if let Some(dir) = &self.dir {
1903 if let Ok(entries) = std::fs::read_dir(dir) {
1904 for entry in entries.flatten() {
1905 let path = entry.path();
1906 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
1907 let _ = std::fs::remove_file(&path);
1908 }
1909 }
1910 }
1911 }
1912 self.entries.clear();
1913 self.order.clear();
1914 self.generations.clear();
1915 self.next_generation = 0;
1916 self.condition_index.clear();
1917 self.empty_footprint_entries.clear();
1918 self.bytes = 0;
1919 }
1920
1921 fn evict(&mut self) {
1922 while self.bytes > self.max_bytes {
1923 let Some((_, key)) = self.order.pop_first() else {
1924 break;
1925 };
1926 self.generations.remove(&key);
1927 if let Some(entry) = self.entries.remove(&key) {
1928 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
1929 self.unindex_entry(key, &entry);
1930 self.remove_from_disk(key);
1934 }
1935 }
1936 }
1937}
1938
1939#[cfg(test)]
1940mod result_cache_lru_tests {
1941 use super::*;
1942
1943 fn row_entry(row_id: u64) -> CachedEntry {
1944 CachedEntry {
1945 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
1946 footprint: std::iter::once(row_id as u32).collect(),
1947 condition_cols: vec![1],
1948 }
1949 }
1950
1951 #[test]
1952 fn hits_update_recency_without_growing_an_order_queue() {
1953 let mut cache = ResultCache::with_max_bytes(u64::MAX);
1954 cache.insert(1, row_entry(1));
1955 cache.insert(2, row_entry(2));
1956 assert_eq!(cache.order.len(), cache.entries.len());
1957
1958 for _ in 0..1_000 {
1959 assert!(cache.get_rows(1).is_some());
1960 }
1961
1962 assert_eq!(cache.order.len(), cache.entries.len());
1963 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
1964
1965 let one_entry = cache.entries[&1].data.approx_bytes();
1966 cache.set_max_bytes(one_entry);
1967 assert!(cache.entries.contains_key(&1));
1968 assert!(!cache.entries.contains_key(&2));
1969 assert_eq!(cache.order.len(), cache.entries.len());
1970 }
1971
1972 #[test]
1973 fn tiny_entries_skip_persistent_tier_by_default() {
1974 let dir = tempfile::tempdir().unwrap();
1975 let mut cache = ResultCache::with_max_bytes(u64::MAX).with_dir(dir.path().to_path_buf());
1976 let tiny = row_entry(1).data.approx_bytes();
1977 assert!(
1978 tiny < 4 * 1024,
1979 "row_entry fixture must be below default 4KiB threshold"
1980 );
1981 assert!(
1982 !cache.would_persist(tiny),
1983 "tiny entry must not force synchronous disk publish"
1984 );
1985 assert!(
1986 cache.would_persist(8 * 1024),
1987 "large entry must still persist when dir is set"
1988 );
1989 cache.set_persist_min_bytes(0);
1990 assert!(
1991 cache.would_persist(tiny),
1992 "persist_min_bytes=0 restores always-persist policy"
1993 );
1994 }
1995
1996 #[test]
1997 fn insert_invalidation_uses_the_condition_reverse_index() {
1998 let mut cache = ResultCache::with_max_bytes(u64::MAX);
1999 let mut first = row_entry(1);
2000 first.condition_cols = vec![1];
2001 let mut second = row_entry(2);
2002 second.condition_cols = vec![2];
2003 cache.insert(1, first);
2004 cache.insert(2, second);
2005
2006 let put_cols = std::iter::once(1_u16).collect();
2007 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
2008
2009 assert!(!cache.entries.contains_key(&1));
2010 assert!(cache.entries.contains_key(&2));
2011 assert!(!cache
2012 .condition_index
2013 .get(&1)
2014 .is_some_and(|keys| keys.contains(&1)));
2015 assert!(cache
2016 .condition_index
2017 .get(&2)
2018 .is_some_and(|keys| keys.contains(&2)));
2019 assert_eq!(cache.order.len(), cache.entries.len());
2020 }
2021
2022 #[test]
2023 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
2024 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2025 cache.insert(1, row_entry(1));
2026 cache.insert(2, row_entry(2));
2027 let mut unknown = row_entry(3);
2028 unknown.footprint.clear();
2029 cache.insert(3, unknown);
2030
2031 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
2032 cache.invalidate(&delete_rids, &HashSet::new());
2033
2034 assert!(!cache.entries.contains_key(&1));
2035 assert!(cache.entries.contains_key(&2));
2036 assert!(!cache.entries.contains_key(&3));
2037 assert!(cache.empty_footprint_entries.is_empty());
2038 assert_eq!(cache.order.len(), cache.entries.len());
2039 }
2040
2041 #[test]
2042 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
2043 let entry = row_entry(7);
2044 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
2045 let owned = bincode::serialize(&SerializedEntry {
2046 condition_cols: entry.condition_cols.clone(),
2047 footprint_bits: entry.footprint.iter().collect(),
2048 data: match &entry.data {
2049 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
2050 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
2051 },
2052 })
2053 .unwrap();
2054 assert_eq!(borrowed, owned);
2055 }
2056}
2057
2058type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
2065
2066fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
2067 let _ = kek;
2068 {
2069 if let Some(k) = kek {
2070 return (
2071 Some(k.derive_table_wal_key(_table_id)),
2072 Some(k.derive_cache_key()),
2073 );
2074 }
2075 }
2076 (None, None)
2077}
2078
2079fn read_table_encryption_salt_root(
2080 root: &crate::durable_file::DurableRoot,
2081) -> Result<[u8; crate::encryption::SALT_LEN]> {
2082 use std::io::Read;
2083
2084 let mut file = root
2085 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
2086 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
2087 let length = file.metadata()?.len();
2088 if length != crate::encryption::SALT_LEN as u64 {
2089 return Err(MongrelError::InvalidArgument(format!(
2090 "salt file is {length} bytes, expected {}",
2091 crate::encryption::SALT_LEN
2092 )));
2093 }
2094 let mut salt = [0_u8; crate::encryption::SALT_LEN];
2095 file.read_exact(&mut salt)?;
2096 Ok(salt)
2097}
2098
2099fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
2101 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
2102}
2103
2104fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
2105 let Some(kek) = kek else {
2106 return HashMap::new();
2107 };
2108 {
2109 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
2110 schema
2111 .columns
2112 .iter()
2113 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
2114 .map(|c| {
2115 let scheme = if schema
2116 .indexes
2117 .iter()
2118 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
2119 {
2120 SCHEME_OPE_RANGE
2121 } else {
2122 SCHEME_HMAC_EQ
2123 };
2124 let key: [u8; 32] = *kek.derive_column_key(c.id);
2125 (c.id, (key, scheme))
2126 })
2127 .collect()
2128 }
2129}
2130
2131pub(crate) struct SharedCtx {
2136 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
2137 pub epoch: Arc<EpochAuthority>,
2138 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
2139 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
2140 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
2141 pub kek: Option<Arc<Kek>>,
2142 pub commit_lock: Arc<parking_lot::Mutex<()>>,
2148 pub shared: Option<SharedWalCtx>,
2152 pub table_name: Option<String>,
2155 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
2158 pub read_only: bool,
2160}
2161
2162#[derive(Clone)]
2168pub(crate) struct SharedWalCtx {
2169 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
2170 pub group: Arc<GroupCommit>,
2171 pub poisoned: Arc<AtomicBool>,
2172 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
2173 pub change_wake: tokio::sync::broadcast::Sender<()>,
2174 pub lifecycle: Arc<crate::core::LifecycleController>,
2177 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
2179}
2180
2181enum WalSink {
2184 Private(Wal),
2185 Shared(SharedWalCtx),
2186 ReadOnly,
2187}
2188
2189impl Clone for WalSink {
2190 fn clone(&self) -> Self {
2191 match self {
2192 Self::Shared(shared) => Self::Shared(shared.clone()),
2193 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
2194 }
2195 }
2196}
2197
2198impl SharedCtx {
2199 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
2203 let n_shards = if cache_dir.is_some() {
2207 1
2208 } else {
2209 crate::cache::CACHE_SHARDS
2210 };
2211 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
2212 let page_cache = if let Some(d) = cache_dir {
2213 Arc::new(crate::cache::Sharded::new(1, || {
2214 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
2215 }))
2216 } else {
2217 Arc::new(crate::cache::Sharded::new(n_shards, || {
2218 crate::cache::PageCache::new(per_shard)
2219 }))
2220 };
2221 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
2222 let decoded_cache = Arc::new(crate::cache::Sharded::new(
2223 crate::cache::CACHE_SHARDS,
2224 || crate::cache::DecodedPageCache::new(decoded_per_shard),
2225 ));
2226 Self {
2227 root_guard: None,
2228 epoch: Arc::new(EpochAuthority::new(0)),
2229 page_cache,
2230 decoded_cache,
2231 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
2232 kek,
2233 commit_lock: Arc::new(parking_lot::Mutex::new(())),
2234 shared: None,
2235 table_name: None,
2236 auth: None,
2237 read_only: false,
2238 }
2239 }
2240}
2241
2242fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
2246 use crate::query::Condition;
2247 match c {
2248 Condition::Pk(_)
2250 | Condition::BitmapEq { .. }
2251 | Condition::BitmapIn { .. }
2252 | Condition::BytesPrefix { .. }
2253 | Condition::IsNull { .. }
2254 | Condition::IsNotNull { .. } => 0,
2255 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
2257 1
2258 }
2259 Condition::FmContains { .. }
2261 | Condition::FmContainsAll { .. }
2262 | Condition::Ann { .. }
2263 | Condition::SparseMatch { .. } => 2,
2264 }
2265}
2266
2267impl Table {
2268 pub(crate) fn build_secondary_index_artifact<F>(
2274 &self,
2275 definition: &IndexDef,
2276 rows: &[Row],
2277 mut checkpoint: F,
2278 ) -> Result<SecondaryIndexArtifact>
2279 where
2280 F: FnMut(usize, usize) -> Result<()>,
2281 {
2282 let mut build_schema = self.schema.clone();
2283 build_schema.indexes = vec![definition.clone()];
2284 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
2285 let name_to_id: HashMap<&str, u16> = self
2286 .schema
2287 .columns
2288 .iter()
2289 .map(|column| (column.name.as_str(), column.id))
2290 .collect();
2291 let mut hot = HotIndex::new();
2292 let mut accepted = Vec::with_capacity(rows.len());
2293
2294 for (offset, row) in rows.iter().enumerate() {
2295 checkpoint(offset, rows.len())?;
2296 if row.deleted {
2297 continue;
2298 }
2299 if let Some(predicate) = &definition.predicate {
2300 let columns: HashMap<u16, &Value> =
2301 row.columns.iter().map(|(id, value)| (*id, value)).collect();
2302 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
2303 continue;
2304 }
2305 }
2306 accepted.push(row);
2307 if definition.kind == IndexKind::LearnedRange {
2308 continue;
2309 }
2310 let effective = if self.column_keys.is_empty() {
2311 row.clone()
2312 } else {
2313 self.tokenized_for_indexes(row)
2314 };
2315 if definition.kind == IndexKind::Ann {
2316 if let (Some(index), Some(vector)) = (
2317 ann.get_mut(&definition.column_id),
2318 effective
2319 .columns
2320 .get(&definition.column_id)
2321 .and_then(Value::as_embedding),
2322 ) {
2323 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
2324 checkpoint(offset, rows.len())
2325 })?;
2326 }
2327 } else {
2328 index_into_single(
2329 definition,
2330 &build_schema,
2331 &effective,
2332 &mut hot,
2333 &mut bitmap,
2334 &mut ann,
2335 &mut fm,
2336 &mut sparse,
2337 &mut minhash,
2338 );
2339 }
2340 }
2341 checkpoint(rows.len(), rows.len())?;
2342
2343 if let Some(index) = ann.get_mut(&definition.column_id) {
2344 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
2345 }
2346
2347 let column_id = definition.column_id;
2348 match definition.kind {
2349 IndexKind::Bitmap => bitmap
2350 .remove(&column_id)
2351 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
2352 IndexKind::Ann => ann
2353 .remove(&column_id)
2354 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
2355 IndexKind::FmIndex => fm
2356 .remove(&column_id)
2357 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
2358 IndexKind::Sparse => sparse
2359 .remove(&column_id)
2360 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
2361 IndexKind::MinHash => minhash
2362 .remove(&column_id)
2363 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
2364 IndexKind::LearnedRange => {
2365 let epsilon = definition
2366 .options
2367 .learned_range
2368 .as_ref()
2369 .map(|options| options.epsilon)
2370 .unwrap_or(16);
2371 let column = self
2372 .schema
2373 .columns
2374 .iter()
2375 .find(|column| column.id == column_id)
2376 .ok_or_else(|| {
2377 MongrelError::Schema(format!(
2378 "index {} references unknown column {column_id}",
2379 definition.name
2380 ))
2381 })?;
2382 let index = match column.ty {
2383 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2384 let pairs: Vec<_> = accepted
2385 .iter()
2386 .filter_map(|row| match row.columns.get(&column_id) {
2387 Some(Value::Int64(value)) if !row.deleted => {
2388 Some((*value, row.row_id.0))
2389 }
2390 _ => None,
2391 })
2392 .collect();
2393 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
2394 }
2395 TypeId::Float32 | TypeId::Float64 => {
2396 let pairs: Vec<_> = accepted
2397 .iter()
2398 .filter_map(|row| match row.columns.get(&column_id) {
2399 Some(Value::Float64(value)) if !row.deleted => {
2400 Some((*value, row.row_id.0))
2401 }
2402 _ => None,
2403 })
2404 .collect();
2405 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
2406 }
2407 ref ty => {
2408 return Err(MongrelError::Schema(format!(
2409 "LearnedRange index {} does not support {ty:?}",
2410 definition.name
2411 )));
2412 }
2413 };
2414 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
2415 }
2416 }
2417 .ok_or_else(|| {
2418 MongrelError::Other(format!(
2419 "failed to construct hidden index {}",
2420 definition.name
2421 ))
2422 })
2423 }
2424
2425 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
2426 let dir = dir.as_ref().to_path_buf();
2427 std::fs::create_dir_all(&dir)?;
2430 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2431 let pinned = root.io_path()?;
2432 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2433 ctx.root_guard = Some(root.clone());
2434 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2435 root.finalize_deferred_sync_shallow()?;
2441 Ok(table)
2442 }
2443
2444 pub fn create_encrypted(
2455 dir: impl AsRef<Path>,
2456 schema: Schema,
2457 table_id: u64,
2458 passphrase: &str,
2459 ) -> Result<Self> {
2460 let dir = dir.as_ref().to_path_buf();
2461 std::fs::create_dir_all(&dir)?;
2462 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2463 root.create_directory_all(META_DIR)?;
2464 let salt = crate::encryption::random_salt()?;
2465 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
2466 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2467 let pinned = root.io_path()?;
2468 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2469 ctx.root_guard = Some(root.clone());
2470 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2471 root.finalize_deferred_sync()?;
2472 Ok(table)
2473 }
2474
2475 pub fn create_with_key(
2480 dir: impl AsRef<Path>,
2481 schema: Schema,
2482 table_id: u64,
2483 key: &[u8],
2484 ) -> Result<Self> {
2485 let dir = dir.as_ref().to_path_buf();
2486 std::fs::create_dir_all(&dir)?;
2487 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2488 root.create_directory_all(META_DIR)?;
2489 let salt = crate::encryption::random_salt()?;
2490 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
2491 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
2492 let pinned = root.io_path()?;
2493 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2494 ctx.root_guard = Some(root.clone());
2495 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2496 root.finalize_deferred_sync()?;
2497 Ok(table)
2498 }
2499
2500 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
2502 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2503 let salt = read_table_encryption_salt_root(&root)?;
2504 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
2505 let pinned = root.io_path()?;
2506 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2507 ctx.root_guard = Some(root);
2508 Self::open_in(&pinned, ctx)
2509 }
2510
2511 pub(crate) fn create_in(
2512 dir: impl AsRef<Path>,
2513 schema: Schema,
2514 table_id: u64,
2515 ctx: SharedCtx,
2516 ) -> Result<Self> {
2517 schema.validate_auto_increment()?;
2518 schema.validate_defaults()?;
2519 schema.validate_ai()?;
2520 for index in &schema.indexes {
2521 index.validate_options()?;
2522 }
2523 let dir = dir.as_ref().to_path_buf();
2524 let runs_root = match ctx.root_guard.as_ref() {
2525 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
2526 None => {
2527 crate::durable_file::create_directory_all(&dir)?;
2528 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
2529 None
2530 }
2531 };
2532 match ctx.root_guard.as_deref() {
2533 Some(root) => write_schema_durable(root, &schema)?,
2534 None => write_schema(&dir, &schema)?,
2535 }
2536 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
2537 let (wal, current_txn_id) = match ctx.shared.clone() {
2540 Some(s) => (WalSink::Shared(s), 0),
2541 None => {
2542 let pinned_wal_root = match ctx.root_guard.as_deref() {
2543 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
2544 None => None,
2545 };
2546 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
2547 root.io_path()?
2548 } else {
2549 let wal_dir = dir.join(WAL_DIR);
2550 std::fs::create_dir_all(&wal_dir)?;
2551 wal_dir
2552 };
2553 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
2554 (Some(root), Some(dk)) => {
2555 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
2556 }
2557 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
2558 (None, Some(dk)) => Wal::create_with_cipher(
2559 wal_dir.join("seg-000000.wal"),
2560 Epoch(0),
2561 Some(make_cipher(dk)),
2562 0,
2563 )?,
2564 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
2565 };
2566 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2567 (WalSink::Private(w), 1)
2568 }
2569 };
2570 let mut manifest = Manifest::new(table_id, schema.schema_id);
2571 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2576 match ctx.root_guard.as_deref() {
2577 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
2578 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
2579 }
2580 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
2581 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2582 let auto_inc = resolve_auto_inc(&schema);
2583 let rcache_dir = dir.join(RCACHE_DIR);
2584 let initial_view = ReadGeneration::empty(&schema);
2585 Ok(Self {
2586 dir,
2587 _root_guard: ctx.root_guard,
2588 runs_root,
2589 idx_root: None,
2590 table_id,
2591 name: ctx.table_name.unwrap_or_default(),
2592 auth: ctx.auth,
2593 read_only: ctx.read_only,
2594 durable_commit_failed: false,
2595 wal,
2596 memtable: Memtable::new(),
2597 mutable_run: MutableRun::new(),
2598 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2599 compaction_zstd_level: 3,
2600 allocator: RowIdAllocator::new(0),
2601 epoch: ctx.epoch,
2602 data_generation: 0,
2603 schema,
2604 hot: HotIndex::new(),
2605 kek: ctx.kek,
2606 column_keys,
2607 run_refs: Vec::new(),
2608 retiring: Vec::new(),
2609 next_run_id: 1,
2610 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2611 current_txn_id,
2612 pending_private_mutations: false,
2613 bitmap,
2614 ann,
2615 fm,
2616 sparse,
2617 minhash,
2618 learned_range: Arc::new(HashMap::new()),
2619 pk_by_row: ReversePkMap::new(),
2620 pinned: BTreeMap::new(),
2621 live_count: 0,
2622 reservoir: crate::reservoir::Reservoir::default(),
2623 reservoir_complete: true,
2624 had_deletes: false,
2625 recent_delete_preimages: HashMap::new(),
2626 run_row_id_ranges: HashMap::new(),
2627 agg_cache: Arc::new(HashMap::new()),
2628 global_idx_epoch: 0,
2629 indexes_complete: true,
2630 index_build_policy: IndexBuildPolicy::default(),
2631 pk_by_row_complete: false,
2632 flushed_epoch: 0,
2633 page_cache: ctx.page_cache,
2634 decoded_cache: ctx.decoded_cache,
2635 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2636 snapshots: ctx.snapshots,
2637 commit_lock: ctx.commit_lock,
2638 result_cache: Arc::new(parking_lot::Mutex::new(
2639 ResultCache::new()
2640 .with_dir(rcache_dir)
2641 .with_cache_dek(cache_dek.clone()),
2642 )),
2643 pending_delete_rids: roaring::RoaringBitmap::new(),
2644 pending_put_cols: std::collections::HashSet::new(),
2645 pending_rows: Vec::new(),
2646 pending_rows_auto_inc: Vec::new(),
2647 pending_dels: Vec::new(),
2648 pending_truncate: None,
2649 wal_dek,
2650 auto_inc,
2651 ttl: None,
2652 pins: Arc::new(crate::retention::PinRegistry::new()),
2653 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2654 read_generation_pin: None,
2655 lookup_metrics: LookupMetrics::default(),
2656 })
2657 }
2658
2659 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
2663 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2664 let pinned = root.io_path()?;
2665 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2666 ctx.root_guard = Some(root);
2667 Self::open_in(&pinned, ctx)
2668 }
2669
2670 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
2673 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2674 let salt = read_table_encryption_salt_root(&root)?;
2675 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2676 let pinned = root.io_path()?;
2677 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2678 ctx.root_guard = Some(root);
2679 let t = Self::open_in(&pinned, ctx)?;
2680 Ok(t)
2681 }
2682
2683 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
2684 let dir = dir.as_ref().to_path_buf();
2685 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2686 let mut manifest = match ctx.root_guard.as_ref() {
2687 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
2688 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
2689 };
2690 let schema: Schema = match ctx.root_guard.as_ref() {
2691 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
2692 None => read_schema(&dir)?,
2693 };
2694 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
2700 let runs_root = match ctx.root_guard.as_ref() {
2701 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
2702 None => None,
2703 };
2704 let idx_root = match ctx.root_guard.as_ref() {
2705 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
2706 Ok(root) => Some(Arc::new(root)),
2707 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2708 Err(error) => return Err(error.into()),
2709 },
2710 None => None,
2711 };
2712 schema.validate_auto_increment()?;
2713 schema.validate_defaults()?;
2714 schema.validate_ai()?;
2715 for index in &schema.indexes {
2716 index.validate_options()?;
2717 }
2718 let replay_epoch = Epoch(manifest.current_epoch);
2719 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
2720 let private_replayed = if ctx.shared.is_none() {
2721 match latest_wal_segment(&dir.join(WAL_DIR))? {
2722 Some(path) => {
2723 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
2724 crate::wal::replay_with_cipher(path, cipher)?
2725 }
2726 None => Vec::new(),
2727 }
2728 } else {
2729 Vec::new()
2730 };
2731 if ctx.shared.is_none() {
2732 preflight_standalone_open(
2733 &dir,
2734 runs_root.as_deref(),
2735 idx_root.as_deref(),
2736 &manifest,
2737 &schema,
2738 &private_replayed,
2739 ctx.kek.clone(),
2740 )?;
2741 }
2742 let next_run_id = derive_next_run_id(
2743 &dir,
2744 runs_root.as_deref(),
2745 &manifest.runs,
2746 &manifest.retiring,
2747 )?;
2748 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2752 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2753 None => {
2754 let replayed = private_replayed;
2755 let wal_dir = dir.join(WAL_DIR);
2761 crate::durable_file::create_directory_all(&wal_dir)?;
2762 let segment = next_wal_segment(&wal_dir)?;
2763 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2764 let temporary = wal_dir.join(format!(
2765 ".recovery-{}-{}-{segment_no:06}.tmp",
2766 std::process::id(),
2767 std::time::SystemTime::now()
2768 .duration_since(std::time::UNIX_EPOCH)
2769 .unwrap_or_default()
2770 .as_nanos()
2771 ));
2772 let mut w = Wal::create_with_cipher(
2773 &temporary,
2774 replay_epoch,
2775 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2776 segment_no,
2777 )?;
2778 for record in &replayed {
2779 w.append_txn(record.txn_id, record.op.clone())?;
2780 }
2781 let mut w = w.publish_as(segment)?;
2782 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2783 let next_txn_id = replayed
2784 .iter()
2785 .map(|record| record.txn_id)
2786 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
2787 .max()
2788 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
2789 .unwrap_or(1);
2790 (WalSink::Private(w), replayed, next_txn_id)
2791 }
2792 };
2793
2794 let mut memtable = Memtable::new();
2795 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
2796 let persisted_epoch = manifest.current_epoch;
2797 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
2804 s.next = manifest.auto_inc_next;
2805 s.seeded = manifest.auto_inc_next > 0;
2806 s
2807 });
2808
2809 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
2816 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
2817 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
2818 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
2819 std::collections::BTreeMap::new();
2820 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
2821 let mut recovered_epoch = manifest.current_epoch;
2822 let mut recovered_manifest_dirty = schema_manifest_repair;
2823 let mut saw_delete = false;
2824 for record in replayed {
2825 let txn_id = record.txn_id;
2826 match record.op {
2827 Op::Put { rows, .. } => {
2828 let rows: Vec<Row> = bincode::deserialize(&rows)?;
2829 for row in &rows {
2830 allocator.advance_to(row.row_id)?;
2831 if let Some(ai) = auto_inc.as_mut() {
2832 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2833 let next = n.checked_add(1).ok_or_else(|| {
2834 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
2835 })?;
2836 if next > ai.next {
2837 ai.next = next;
2838 }
2839 }
2840 }
2841 }
2842 staged_puts.entry(txn_id).or_default().extend(rows);
2843 }
2844 Op::Delete { row_ids, .. } => {
2845 staged_deletes.entry(txn_id).or_default().extend(row_ids);
2846 }
2847 Op::TxnCommit { epoch, .. } => {
2848 let commit_epoch = Epoch(epoch);
2849 recovered_epoch = recovered_epoch.max(epoch);
2850 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
2851 memtable = Memtable::new();
2852 replayed_puts.clear();
2853 replayed_deletes.clear();
2854 manifest.runs.clear();
2855 manifest.retiring.clear();
2856 manifest.live_count = 0;
2857 manifest.global_idx_epoch = 0;
2858 manifest.current_epoch = manifest.current_epoch.max(epoch);
2859 recovered_manifest_dirty = true;
2860 saw_delete = true;
2861 }
2862 if let Some(puts) = staged_puts.remove(&txn_id) {
2863 if commit_epoch.0 > manifest.flushed_epoch {
2864 for row in &puts {
2865 memtable.upsert(row.clone());
2866 }
2867 replayed_puts.entry(commit_epoch).or_default().extend(puts);
2868 }
2869 }
2870 if let Some(dels) = staged_deletes.remove(&txn_id) {
2871 saw_delete = true;
2872 if commit_epoch.0 > manifest.flushed_epoch {
2873 for rid in dels {
2874 memtable.tombstone(rid, commit_epoch);
2875 replayed_deletes.push((rid, commit_epoch));
2876 }
2877 }
2878 }
2879 }
2880 Op::TxnAbort => {
2881 staged_puts.remove(&txn_id);
2882 staged_deletes.remove(&txn_id);
2883 staged_truncates.remove(&txn_id);
2884 }
2885 Op::TruncateTable { .. } => {
2886 staged_puts.remove(&txn_id);
2887 staged_deletes.remove(&txn_id);
2888 staged_truncates.insert(txn_id);
2889 }
2890 Op::ExternalTableState { .. }
2891 | Op::Flush { .. }
2892 | Op::Ddl(_)
2893 | Op::BeforeImage { .. }
2894 | Op::CommitTimestamp { .. }
2895 | Op::SpilledRows { .. } => {}
2896 }
2897 }
2898
2899 let rcache_dir = dir.join(RCACHE_DIR);
2900 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2901 let initial_view = ReadGeneration::empty(&schema);
2902 let mut db = Self {
2903 dir,
2904 _root_guard: ctx.root_guard,
2905 runs_root,
2906 idx_root,
2907 table_id: manifest.table_id,
2908 name: ctx.table_name.unwrap_or_default(),
2909 auth: ctx.auth,
2910 read_only: ctx.read_only,
2911 durable_commit_failed: false,
2912 wal,
2913 memtable,
2914 mutable_run: MutableRun::new(),
2915 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2916 compaction_zstd_level: 3,
2917 allocator,
2918 epoch: ctx.epoch,
2919 data_generation: persisted_epoch,
2920 schema,
2921 hot: HotIndex::new(),
2922 kek: ctx.kek,
2923 column_keys,
2924 run_refs: manifest.runs.clone(),
2925 retiring: manifest.retiring.clone(),
2926 next_run_id,
2927 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2928 current_txn_id,
2929 pending_private_mutations: false,
2930 bitmap: HashMap::new(),
2931 ann: HashMap::new(),
2932 fm: HashMap::new(),
2933 sparse: HashMap::new(),
2934 minhash: HashMap::new(),
2935 learned_range: Arc::new(HashMap::new()),
2936 pk_by_row: ReversePkMap::new(),
2937 pinned: BTreeMap::new(),
2938 live_count: manifest.live_count,
2939 reservoir: crate::reservoir::Reservoir::default(),
2940 reservoir_complete: false,
2941 had_deletes: saw_delete
2942 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
2943 != manifest.live_count,
2944 recent_delete_preimages: HashMap::new(),
2945 run_row_id_ranges: HashMap::new(),
2946 agg_cache: Arc::new(HashMap::new()),
2947 global_idx_epoch: manifest.global_idx_epoch,
2948 indexes_complete: true,
2949 index_build_policy: IndexBuildPolicy::default(),
2950 pk_by_row_complete: false,
2951 flushed_epoch: manifest.flushed_epoch,
2952 page_cache: ctx.page_cache,
2953 decoded_cache: ctx.decoded_cache,
2954 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2955 snapshots: ctx.snapshots,
2956 commit_lock: ctx.commit_lock,
2957 result_cache: Arc::new(parking_lot::Mutex::new(
2958 ResultCache::new()
2959 .with_dir(rcache_dir)
2960 .with_cache_dek(cache_dek.clone()),
2961 )),
2962 pending_delete_rids: roaring::RoaringBitmap::new(),
2963 pending_put_cols: std::collections::HashSet::new(),
2964 pending_rows: Vec::new(),
2965 pending_rows_auto_inc: Vec::new(),
2966 pending_dels: Vec::new(),
2967 pending_truncate: None,
2968 wal_dek,
2969 auto_inc,
2970 ttl: manifest.ttl,
2971 pins: Arc::new(crate::retention::PinRegistry::new()),
2972 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2973 read_generation_pin: None,
2974 lookup_metrics: LookupMetrics::default(),
2975 };
2976
2977 db.epoch.advance_recovered(Epoch(recovered_epoch));
2980
2981 let checkpoint = match db.idx_root.as_deref() {
2986 Some(root) => {
2987 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
2988 }
2989 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
2990 };
2991 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
2992 c.epoch_built == manifest.global_idx_epoch
2993 && manifest.global_idx_epoch > 0
2994 && manifest
2995 .runs
2996 .iter()
2997 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
2998 });
2999 if let Some(loaded) = checkpoint {
3000 if checkpoint_valid {
3001 db.hot = loaded.hot;
3002 db.bitmap = loaded.bitmap;
3003 db.ann = loaded.ann;
3004 db.fm = loaded.fm;
3005 db.sparse = loaded.sparse;
3006 db.minhash = loaded.minhash;
3007 db.learned_range = Arc::new(loaded.learned_range);
3008 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
3013 for (cid, idx) in bitmap0 {
3014 db.bitmap.entry(cid).or_insert(idx);
3015 }
3016 for (cid, idx) in ann0 {
3017 db.ann.entry(cid).or_insert(idx);
3018 }
3019 for (cid, idx) in fm0 {
3020 db.fm.entry(cid).or_insert(idx);
3021 }
3022 for (cid, idx) in sparse0 {
3023 db.sparse.entry(cid).or_insert(idx);
3024 }
3025 for (cid, idx) in minhash0 {
3026 db.minhash.entry(cid).or_insert(idx);
3027 }
3028 }
3031 }
3032 if !checkpoint_valid {
3033 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
3034 db.bitmap = bitmap;
3035 db.ann = ann;
3036 db.fm = fm;
3037 db.sparse = sparse;
3038 db.minhash = minhash;
3039 db.rebuild_indexes_from_runs()?;
3040 db.build_learned_ranges()?;
3041 }
3042
3043 for (epoch, group) in replayed_puts {
3048 let (losers, winner_pks) = db.partition_pk_winners(&group);
3049 for (key, &row_id) in &winner_pks {
3050 if let Some(old_rid) = db.hot.get(key) {
3051 if old_rid != row_id {
3052 db.tombstone_row(old_rid, epoch, None, false);
3053 }
3054 }
3055 }
3056 for &loser_rid in &losers {
3057 db.tombstone_row(loser_rid, epoch, None, false);
3058 }
3059 for (key, row_id) in winner_pks {
3060 db.insert_hot_pk(key, row_id);
3061 }
3062 if db.schema.primary_key().is_none() {
3063 for r in &group {
3064 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3065 }
3066 }
3067 for r in &group {
3068 if !losers.contains(&r.row_id) {
3069 db.index_row(r);
3070 }
3071 }
3072 }
3073 for (rid, epoch) in &replayed_deletes {
3077 db.remove_hot_for_row(*rid, *epoch);
3078 }
3079
3080 if recovered_manifest_dirty {
3081 let rows = db.visible_rows(Snapshot::unbounded())?;
3082 db.live_count = rows.len() as u64;
3083 db.persist_manifest(Epoch(recovered_epoch))?;
3084 }
3085
3086 db.result_cache.lock().load_persistent();
3093 db.refresh_run_row_id_ranges();
3095 Ok(db)
3096 }
3097
3098 fn ensure_reservoir_complete(&mut self) -> Result<()> {
3104 if self.reservoir_complete {
3105 return Ok(());
3106 }
3107 self.rebuild_reservoir()?;
3108 self.reservoir_complete = true;
3109 Ok(())
3110 }
3111
3112 fn rebuild_reservoir(&mut self) -> Result<()> {
3115 let snap = self.snapshot();
3116 let rows = self.visible_rows(snap)?;
3117 self.reservoir.reset();
3118 for r in rows {
3119 self.reservoir.offer(r.row_id.0);
3120 }
3121 Ok(())
3122 }
3123
3124 pub fn rebuild_indexes(&mut self) -> Result<()> {
3128 self.rebuild_indexes_from_runs_inner(None)
3129 }
3130
3131 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
3132 self.rebuild_indexes_from_runs_inner(None)
3133 }
3134
3135 fn pk_equality_fallback(
3138 &self,
3139 pk_column_id: u16,
3140 lookup: &[u8],
3141 snapshot: Snapshot,
3142 ) -> Result<RowIdSet> {
3143 for row in self.memtable.visible_versions_at(snapshot) {
3145 self.lookup_metrics
3146 .hot_lookup_fallback_overlay_rows
3147 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3148 if row.deleted {
3149 continue;
3150 }
3151 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3152 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3153 return Ok(RowIdSet::one(row.row_id.0));
3154 }
3155 }
3156 }
3157 for row in self.mutable_run.visible_versions_at(snapshot) {
3158 self.lookup_metrics
3159 .hot_lookup_fallback_overlay_rows
3160 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3161 if row.deleted {
3162 continue;
3163 }
3164 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3165 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3166 return Ok(RowIdSet::one(row.row_id.0));
3167 }
3168 }
3169 }
3170 if lookup.len() == 8 {
3173 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
3174 let n = i64::from_be_bytes(arr);
3175 let result = self.range_scan_i64(pk_column_id, n, n, snapshot)?;
3176 self.lookup_metrics
3177 .hot_lookup_fallback_runs
3178 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3179 return Ok(result);
3180 }
3181 }
3182 let mut found = Vec::new();
3185 let overlay = self.overlay_rid_set(snapshot);
3186 for rr in &self.run_refs {
3187 self.lookup_metrics
3188 .hot_lookup_fallback_runs
3189 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3190 let mut reader = self.open_reader(rr.run_id)?;
3191 for row in reader.visible_rows(snapshot.epoch)? {
3192 if overlay.contains(&row.row_id.0) || row.deleted {
3193 continue;
3194 }
3195 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3196 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3197 found.push(row.row_id.0);
3198 }
3199 }
3200 }
3201 }
3202 Ok(RowIdSet::from_unsorted(found))
3203 }
3204
3205 fn rebuild_indexes_from_runs_inner(
3206 &mut self,
3207 control: Option<&crate::ExecutionControl>,
3208 ) -> Result<()> {
3209 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
3212 crate::retention::PinSource::OnlineIndexBuild,
3213 self.current_epoch(),
3214 );
3215 self.hot = HotIndex::new();
3216 self.pk_by_row.clear();
3217 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3218 self.bitmap = bitmap;
3219 self.ann = ann;
3220 self.fm = fm;
3221 self.sparse = sparse;
3222 self.minhash = minhash;
3223 let snapshot = Epoch(u64::MAX);
3224 let ttl_now = unix_nanos_now();
3225 let mut scanned = 0_usize;
3226 for rr in self.run_refs.clone() {
3227 if let Some(control) = control {
3228 control.checkpoint()?;
3229 }
3230 let mut reader = self.open_reader(rr.run_id)?;
3231 for row in reader.visible_rows(snapshot)? {
3232 if scanned.is_multiple_of(256) {
3233 if let Some(control) = control {
3234 control.checkpoint()?;
3235 }
3236 }
3237 scanned += 1;
3238 if self.row_expired_at(&row, ttl_now) {
3239 continue;
3240 }
3241 let tok_row = self.tokenized_for_indexes(&row);
3242 index_into(
3243 &self.schema,
3244 &tok_row,
3245 &mut self.hot,
3246 &mut self.bitmap,
3247 &mut self.ann,
3248 &mut self.fm,
3249 &mut self.sparse,
3250 &mut self.minhash,
3251 );
3252 }
3253 }
3254 for row in self.mutable_run.visible_versions(snapshot) {
3255 if scanned.is_multiple_of(256) {
3256 if let Some(control) = control {
3257 control.checkpoint()?;
3258 }
3259 }
3260 scanned += 1;
3261 if row.deleted {
3262 self.remove_hot_for_row(row.row_id, snapshot);
3263 } else if !self.row_expired_at(&row, ttl_now) {
3264 self.index_row(&row);
3265 }
3266 }
3267 for row in self.memtable.visible_versions(snapshot) {
3268 if scanned.is_multiple_of(256) {
3269 if let Some(control) = control {
3270 control.checkpoint()?;
3271 }
3272 }
3273 scanned += 1;
3274 if row.deleted {
3275 self.remove_hot_for_row(row.row_id, snapshot);
3276 } else if !self.row_expired_at(&row, ttl_now) {
3277 self.index_row(&row);
3278 }
3279 }
3280 let pin_epochs = self.active_pin_epochs_for_rebuild();
3286 if !pin_epochs.is_empty() {
3287 let current_snap = Snapshot::at(Epoch(u64::MAX));
3288 for pin_epoch in pin_epochs {
3289 if let Some(control) = control {
3290 control.checkpoint()?;
3291 }
3292 let pin_snap = Snapshot::at(pin_epoch);
3293 for rr in self.run_refs.clone() {
3294 if let Some(control) = control {
3295 control.checkpoint()?;
3296 }
3297 let mut reader = self.open_reader(rr.run_id)?;
3298 for row in reader.visible_rows(pin_epoch)? {
3299 if row.deleted || self.row_expired_at(&row, ttl_now) {
3300 continue;
3301 }
3302 if self.get(row.row_id, current_snap).is_none() {
3303 self.index_bitmap_membership_only(&row);
3304 }
3305 }
3306 }
3307 for row in self
3308 .mutable_run
3309 .visible_versions_at(pin_snap)
3310 .into_iter()
3311 .chain(self.memtable.visible_versions_at(pin_snap))
3312 {
3313 if row.deleted || self.row_expired_at(&row, ttl_now) {
3314 continue;
3315 }
3316 if self.get(row.row_id, current_snap).is_none() {
3317 self.index_bitmap_membership_only(&row);
3318 }
3319 }
3320 }
3321 }
3322 self.recent_delete_preimages.clear();
3323 self.refresh_pk_by_row_from_hot();
3324 Ok(())
3325 }
3326
3327 fn index_bitmap_membership_only(&mut self, row: &Row) {
3330 if row.deleted {
3331 return;
3332 }
3333 let columns_map: HashMap<u16, &Value> =
3334 row.columns.iter().map(|(k, v)| (*k, v)).collect();
3335 let name_to_id: HashMap<&str, u16> = self
3336 .schema
3337 .columns
3338 .iter()
3339 .map(|c| (c.name.as_str(), c.id))
3340 .collect();
3341 for idef in &self.schema.indexes {
3342 if idef.kind != crate::schema::IndexKind::Bitmap {
3343 continue;
3344 }
3345 if let Some(pred) = &idef.predicate {
3346 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3347 continue;
3348 }
3349 }
3350 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
3351 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3352 b.insert(key, row.row_id);
3353 }
3354 }
3355 }
3356 }
3357
3358 fn refresh_pk_by_row_from_hot(&mut self) {
3359 self.pk_by_row_complete = true;
3360 if self.schema.primary_key().is_none() {
3361 self.pk_by_row.clear();
3362 return;
3363 }
3364 self.pk_by_row = ReversePkMap::from_entries(
3370 self.hot
3371 .entries()
3372 .into_iter()
3373 .map(|(key, row_id)| (row_id, key)),
3374 );
3375 }
3376
3377 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
3378 if self.schema.primary_key().is_some() {
3379 self.pk_by_row.insert(row_id, key.clone());
3380 }
3381 self.hot.insert(key, row_id);
3382 }
3383
3384 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
3388 self.build_learned_ranges_inner(None)
3389 }
3390
3391 fn build_learned_ranges_inner(
3392 &mut self,
3393 control: Option<&crate::ExecutionControl>,
3394 ) -> Result<()> {
3395 self.learned_range = Arc::new(HashMap::new());
3396 if self.run_refs.len() != 1 {
3397 return Ok(());
3398 }
3399 let cols: Vec<(u16, usize)> = self
3400 .schema
3401 .indexes
3402 .iter()
3403 .filter(|i| i.kind == IndexKind::LearnedRange)
3404 .map(|i| {
3405 (
3406 i.column_id,
3407 i.options
3408 .learned_range
3409 .as_ref()
3410 .map(|options| options.epsilon)
3411 .unwrap_or(16),
3412 )
3413 })
3414 .collect();
3415 if cols.is_empty() {
3416 return Ok(());
3417 }
3418 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3419 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
3420 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
3421 _ => return Ok(()),
3422 };
3423 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
3424 if column_index % 256 == 0 {
3425 if let Some(control) = control {
3426 control.checkpoint()?;
3427 }
3428 }
3429 let ty = self
3430 .schema
3431 .columns
3432 .iter()
3433 .find(|c| c.id == cid)
3434 .map(|c| c.ty.clone())
3435 .unwrap_or(TypeId::Int64);
3436 match ty {
3437 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
3438 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
3439 let pairs: Vec<(i64, u64)> = data
3440 .iter()
3441 .zip(row_ids.iter())
3442 .map(|(v, r)| (*v, *r))
3443 .collect();
3444 Arc::make_mut(&mut self.learned_range).insert(
3445 cid,
3446 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
3447 );
3448 }
3449 }
3450 TypeId::Float64 => {
3451 if let columnar::NativeColumn::Float64 { data, .. } =
3452 reader.column_native(cid)?
3453 {
3454 let pairs: Vec<(f64, u64)> = data
3455 .iter()
3456 .zip(row_ids.iter())
3457 .map(|(v, r)| (*v, *r))
3458 .collect();
3459 Arc::make_mut(&mut self.learned_range).insert(
3460 cid,
3461 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
3462 );
3463 }
3464 }
3465 _ => {}
3466 }
3467 }
3468 Ok(())
3469 }
3470
3471 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
3478 if self.indexes_complete {
3479 crate::trace::QueryTrace::record(|t| {
3480 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
3481 });
3482 return Ok(());
3483 }
3484 crate::trace::QueryTrace::record(|t| {
3485 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
3486 });
3487 self.rebuild_indexes_from_runs()?;
3488 self.build_learned_ranges()?;
3489 self.indexes_complete = true;
3490 let epoch = self.current_epoch();
3491 self.checkpoint_indexes(epoch);
3492 Ok(())
3493 }
3494
3495 #[doc(hidden)]
3498 pub fn ensure_indexes_complete_controlled<F>(
3499 &mut self,
3500 control: &crate::ExecutionControl,
3501 before_publish: F,
3502 ) -> Result<bool>
3503 where
3504 F: FnOnce() -> bool,
3505 {
3506 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
3507 .map(|(changed, _)| changed)
3508 }
3509
3510 #[doc(hidden)]
3513 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
3514 &mut self,
3515 control: &crate::ExecutionControl,
3516 before_publish: F,
3517 ) -> Result<(bool, Option<MaintenanceReceipt>)>
3518 where
3519 F: FnOnce() -> bool,
3520 {
3521 if self.indexes_complete {
3522 crate::trace::QueryTrace::record(|trace| {
3523 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
3524 });
3525 return Ok((false, None));
3526 }
3527 crate::trace::QueryTrace::record(|trace| {
3528 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
3529 });
3530 control.checkpoint()?;
3531 let maintenance_epoch = self.current_epoch();
3532 self.rebuild_indexes_from_runs_inner(Some(control))?;
3533 self.build_learned_ranges_inner(Some(control))?;
3534 control.checkpoint()?;
3535 if !before_publish() {
3536 return Err(MongrelError::Cancelled);
3537 }
3538 self.indexes_complete = true;
3539 self.checkpoint_indexes(maintenance_epoch);
3540 Ok((
3541 true,
3542 Some(MaintenanceReceipt {
3543 epoch: maintenance_epoch,
3544 }),
3545 ))
3546 }
3547
3548 fn pending_epoch(&self) -> Epoch {
3549 Epoch(self.epoch.visible().0 + 1)
3550 }
3551
3552 fn is_shared(&self) -> bool {
3555 matches!(self.wal, WalSink::Shared(_))
3556 }
3557
3558 fn ensure_txn_id(&mut self) -> Result<u64> {
3562 if self.current_txn_id == 0 {
3563 let id = match &self.wal {
3564 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
3565 WalSink::Private(_) => {
3566 return Err(MongrelError::Full(
3567 "standalone transaction id namespace exhausted".into(),
3568 ))
3569 }
3570 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3571 };
3572 self.current_txn_id = id;
3573 }
3574 Ok(self.current_txn_id)
3575 }
3576
3577 fn wal_append_data(&mut self, op: Op) -> Result<()> {
3580 self.ensure_writable()?;
3581 let txn_id = self.ensure_txn_id()?;
3582 let table_id = self.table_id;
3583 match &mut self.wal {
3584 WalSink::Private(w) => {
3585 w.append_txn(txn_id, op)?;
3586 self.pending_private_mutations = true;
3587 }
3588 WalSink::Shared(s) => {
3589 s.wal.lock().append(txn_id, table_id, op)?;
3590 }
3591 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3592 }
3593 Ok(())
3594 }
3595
3596 fn ensure_writable(&self) -> Result<()> {
3597 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
3598 return Err(MongrelError::ReadOnlyReplica);
3599 }
3600 if self.durable_commit_failed {
3601 return Err(MongrelError::Other(
3602 "table poisoned by post-commit failure; reopen required".into(),
3603 ));
3604 }
3605 Ok(())
3606 }
3607
3608 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
3619 match &self.auth {
3620 Some(checker) => checker.check(&self.name, perm),
3621 None => Ok(()),
3622 }
3623 }
3624 pub fn require_select(&self) -> Result<()> {
3629 self.require(crate::auth_state::RequiredPermission::Select)
3630 }
3631 fn require_insert(&self) -> Result<()> {
3632 self.require(crate::auth_state::RequiredPermission::Insert)
3633 }
3634 #[allow(dead_code)]
3638 fn require_update(&self) -> Result<()> {
3639 self.require(crate::auth_state::RequiredPermission::Update)
3640 }
3641 fn require_delete(&self) -> Result<()> {
3642 self.require(crate::auth_state::RequiredPermission::Delete)
3643 }
3644
3645 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
3648 self.require_insert()?;
3649 Ok(self.put_returning(columns)?.0)
3650 }
3651
3652 pub fn put_returning(
3657 &mut self,
3658 mut columns: Vec<(u16, Value)>,
3659 ) -> Result<(RowId, Option<i64>)> {
3660 self.require_insert()?;
3661 let assigned = self.fill_auto_inc(&mut columns)?;
3662 self.apply_defaults(&mut columns)?;
3663 self.schema.validate_values(&columns)?;
3664 let row_id = if self.schema.clustered {
3669 self.derive_clustered_row_id(&columns)?
3670 } else {
3671 self.allocator.alloc()?
3672 };
3673 let epoch = self.pending_epoch();
3674 let mut row = Row::new(row_id, epoch);
3675 for (col_id, val) in columns {
3676 row.columns.insert(col_id, val);
3677 }
3678 self.commit_rows(vec![row], assigned.is_some())?;
3679 Ok((row_id, assigned))
3680 }
3681
3682 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
3685 self.require_insert()?;
3686 Ok(self
3687 .put_batch_returning(batch)?
3688 .into_iter()
3689 .map(|(r, _)| r)
3690 .collect())
3691 }
3692
3693 pub fn put_batch_returning(
3696 &mut self,
3697 batch: Vec<Vec<(u16, Value)>>,
3698 ) -> Result<Vec<(RowId, Option<i64>)>> {
3699 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
3700 for mut cols in batch {
3701 let assigned = self.fill_auto_inc(&mut cols)?;
3702 self.apply_defaults(&mut cols)?;
3703 filled.push((cols, assigned));
3704 }
3705 for (cols, _) in &filled {
3706 self.schema.validate_values(cols)?;
3707 }
3708 let epoch = self.pending_epoch();
3709 let mut rows = Vec::with_capacity(filled.len());
3710 let mut ids = Vec::with_capacity(filled.len());
3711 let first_row_id = if self.schema.clustered {
3712 None
3713 } else {
3714 let count = u64::try_from(filled.len())
3715 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
3716 Some(self.allocator.alloc_range(count)?.0)
3717 };
3718 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
3719 let row_id = match first_row_id {
3720 Some(first) => RowId(first + row_index as u64),
3721 None => self.derive_clustered_row_id(&cols)?,
3722 };
3723 let mut row = Row::new(row_id, epoch);
3724 for (c, v) in cols {
3725 row.columns.insert(c, v);
3726 }
3727 ids.push((row_id, assigned));
3728 rows.push(row);
3729 }
3730 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
3731 self.commit_rows(rows, all_auto_generated)?;
3732 Ok(ids)
3733 }
3734
3735 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
3741 self.ensure_writable()?;
3742 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
3743 return Ok(None);
3744 };
3745 let pos = columns.iter().position(|(c, _)| *c == cid);
3746 let assigned = match pos {
3747 Some(i) => match &columns[i].1 {
3748 Value::Null => {
3749 let next = self.alloc_auto_inc_value()?;
3750 columns[i].1 = Value::Int64(next);
3751 Some(next)
3752 }
3753 Value::Int64(n) => {
3754 self.advance_auto_inc_past(*n)?;
3755 None
3756 }
3757 other => {
3758 return Err(MongrelError::InvalidArgument(format!(
3759 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
3760 other
3761 )))
3762 }
3763 },
3764 None => {
3765 let next = self.alloc_auto_inc_value()?;
3766 columns.push((cid, Value::Int64(next)));
3767 Some(next)
3768 }
3769 };
3770 Ok(assigned)
3771 }
3772
3773 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
3779 for col in &self.schema.columns {
3780 let Some(expr) = &col.default_value else {
3781 continue;
3782 };
3783 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
3785 continue;
3786 }
3787 let pos = columns.iter().position(|(c, _)| *c == col.id);
3788 let needs_default = match pos {
3789 None => true,
3790 Some(i) => matches!(columns[i].1, Value::Null),
3791 };
3792 if !needs_default {
3793 continue;
3794 }
3795 let v = match expr {
3796 crate::schema::DefaultExpr::Static(v) => v.clone(),
3797 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
3798 crate::schema::DefaultExpr::Uuid => {
3799 let mut buf = [0u8; 16];
3800 getrandom::getrandom(&mut buf)
3801 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
3802 Value::Uuid(buf)
3803 }
3804 };
3805 match pos {
3806 None => columns.push((col.id, v)),
3807 Some(i) => columns[i].1 = v,
3808 }
3809 }
3810 Ok(())
3811 }
3812
3813 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
3815 self.ensure_auto_inc_seeded()?;
3816 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3818 let v = ai.next;
3819 ai.next = ai
3820 .next
3821 .checked_add(1)
3822 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3823 Ok(v)
3824 }
3825
3826 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
3829 self.ensure_auto_inc_seeded()?;
3830 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3831 let floor = used
3832 .checked_add(1)
3833 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3834 .max(1);
3835 if ai.next < floor {
3836 ai.next = floor;
3837 }
3838 Ok(())
3839 }
3840
3841 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
3846 let needs_seed = match self.auto_inc {
3847 Some(ai) => !ai.seeded,
3848 None => return Ok(()),
3849 };
3850 if !needs_seed {
3851 return Ok(());
3852 }
3853 if self.seed_empty_auto_inc() {
3854 return Ok(());
3855 }
3856 let cid = self
3857 .auto_inc
3858 .as_ref()
3859 .expect("auto-inc column present")
3860 .column_id;
3861 let max = self.scan_max_int64(cid)?;
3862 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3863 let floor = max
3864 .checked_add(1)
3865 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3866 .max(1);
3867 if ai.next < floor {
3868 ai.next = floor;
3869 }
3870 ai.seeded = true;
3871 Ok(())
3872 }
3873
3874 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
3875 if n == 0 || self.auto_inc.is_none() {
3876 return Ok(None);
3877 }
3878 self.ensure_auto_inc_seeded()?;
3879 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
3880 let start = ai.next;
3881 let count = i64::try_from(n)
3882 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
3883 ai.next = ai
3884 .next
3885 .checked_add(count)
3886 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
3887 Ok(Some(start))
3888 }
3889
3890 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
3894 let mut max: i64 = 0;
3895 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
3896 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3897 if *n > max {
3898 max = *n;
3899 }
3900 }
3901 }
3902 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
3903 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
3904 if *n > max {
3905 max = *n;
3906 }
3907 }
3908 }
3909 for rr in self.run_refs.clone() {
3910 let reader = self.open_reader(rr.run_id)?;
3911 if let Some(stats) = reader.column_page_stats(column_id) {
3912 for s in stats {
3913 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
3914 if n > max {
3915 max = n;
3916 }
3917 }
3918 }
3919 } else if reader.has_column(column_id) {
3920 if let columnar::NativeColumn::Int64 { data, validity } =
3921 reader.column_native_shared(column_id)?
3922 {
3923 for (i, n) in data.iter().enumerate() {
3924 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
3925 {
3926 max = *n;
3927 }
3928 }
3929 }
3930 }
3931 }
3932 Ok(max)
3933 }
3934
3935 fn seed_empty_auto_inc(&mut self) -> bool {
3936 let Some(ai) = self.auto_inc.as_mut() else {
3937 return false;
3938 };
3939 if ai.seeded || self.live_count != 0 {
3940 return false;
3941 }
3942 if ai.next < 1 {
3943 ai.next = 1;
3944 }
3945 ai.seeded = true;
3946 true
3947 }
3948
3949 fn advance_auto_inc_from_native_columns(
3950 &mut self,
3951 columns: &[(u16, columnar::NativeColumn)],
3952 n: usize,
3953 live_before: u64,
3954 ) -> Result<()> {
3955 let Some(ai) = self.auto_inc.as_mut() else {
3956 return Ok(());
3957 };
3958 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
3959 return Ok(());
3960 };
3961 let columnar::NativeColumn::Int64 { data, validity } = col else {
3962 return Err(MongrelError::InvalidArgument(format!(
3963 "AUTO_INCREMENT column {} must be Int64",
3964 ai.column_id
3965 )));
3966 };
3967 let max = if native_int64_strictly_increasing(col, n) {
3968 data.get(n.saturating_sub(1)).copied()
3969 } else {
3970 data.iter()
3971 .take(n)
3972 .enumerate()
3973 .filter_map(|(i, v)| {
3974 if validity.is_empty() || columnar::validity_bit(validity, i) {
3975 Some(*v)
3976 } else {
3977 None
3978 }
3979 })
3980 .max()
3981 };
3982 if let Some(max) = max {
3983 let floor = max
3984 .checked_add(1)
3985 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
3986 .max(1);
3987 if ai.next < floor {
3988 ai.next = floor;
3989 }
3990 if ai.seeded || live_before == 0 {
3991 ai.seeded = true;
3992 }
3993 }
3994 Ok(())
3995 }
3996
3997 fn fill_auto_inc_native_columns(
3998 &mut self,
3999 columns: &mut Vec<(u16, columnar::NativeColumn)>,
4000 n: usize,
4001 ) -> Result<()> {
4002 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4003 return Ok(());
4004 };
4005 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
4006 if let Some(start) = self.alloc_auto_inc_range(n)? {
4007 columns.push((
4008 cid,
4009 columnar::NativeColumn::Int64 {
4010 data: (start..start.saturating_add(n as i64)).collect(),
4011 validity: vec![0xFF; n.div_ceil(8)],
4012 },
4013 ));
4014 }
4015 return Ok(());
4016 };
4017
4018 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
4019 return Err(MongrelError::InvalidArgument(format!(
4020 "AUTO_INCREMENT column {cid} must be Int64"
4021 )));
4022 };
4023 if data.len() < n {
4024 return Err(MongrelError::InvalidArgument(format!(
4025 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
4026 data.len()
4027 )));
4028 }
4029 if columnar::all_non_null(validity, n) {
4030 return Ok(());
4031 }
4032 if validity.iter().all(|b| *b == 0) {
4033 if let Some(start) = self.alloc_auto_inc_range(n)? {
4034 for (i, slot) in data.iter_mut().take(n).enumerate() {
4035 *slot = start.saturating_add(i as i64);
4036 }
4037 *validity = vec![0xFF; n.div_ceil(8)];
4038 }
4039 return Ok(());
4040 }
4041
4042 let new_validity = vec![0xFF; data.len().div_ceil(8)];
4043 for (i, slot) in data.iter_mut().enumerate().take(n) {
4044 if columnar::validity_bit(validity, i) {
4045 self.advance_auto_inc_past(*slot)?;
4046 } else {
4047 *slot = self.alloc_auto_inc_value()?;
4048 }
4049 }
4050 *validity = new_validity;
4051 Ok(())
4052 }
4053
4054 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
4068 self.ensure_writable()?;
4069 if self.auto_inc.is_none() {
4070 return Ok(None);
4071 }
4072 Ok(Some(self.alloc_auto_inc_value()?))
4073 }
4074
4075 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
4081 let payload = bincode::serialize(&rows)?;
4082 self.wal_append_data(Op::Put {
4083 table_id: self.table_id,
4084 rows: payload,
4085 })?;
4086 if self.is_shared() {
4087 self.pending_rows_auto_inc
4088 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
4089 self.pending_rows.extend(rows);
4090 } else {
4091 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
4092 }
4093 Ok(())
4094 }
4095
4096 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
4099 self.ensure_indexes_complete()
4100 }
4101
4102 pub(crate) fn prepare_durable_publish_controlled(
4103 &mut self,
4104 control: &crate::ExecutionControl,
4105 ) -> Result<()> {
4106 self.ensure_indexes_complete_controlled(control, || true)?;
4107 Ok(())
4108 }
4109
4110 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
4111 self.apply_put_rows_inner_prepared(rows, true);
4112 }
4113
4114 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
4115 if check_existing_pk {
4116 self.ensure_indexes_complete()?;
4117 }
4118 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
4119 Ok(())
4120 }
4121
4122 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
4126 if rows.len() == 1 {
4130 let row = rows.into_iter().next().expect("len checked");
4131 self.apply_put_row_single(row, check_existing_pk);
4132 return;
4133 }
4134 let pk_id = self.schema.primary_key().map(|c| c.id);
4151 let probe = match pk_id {
4152 Some(pid) => {
4153 check_existing_pk
4154 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
4155 }
4156 None => false,
4157 };
4158 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
4161 for r in rows {
4162 for &cid in r.columns.keys() {
4163 self.pending_put_cols.insert(cid);
4164 }
4165 let mut replaced_image: Option<Row> = None;
4166 match pk_id {
4167 Some(pid) if probe || maintain_pk_by_row => {
4168 if let Some(pk_val) = r.columns.get(&pid) {
4169 let key = self.index_lookup_key(pid, pk_val);
4170 if probe {
4171 if let Some(old_rid) = self.hot.get(&key) {
4172 if old_rid != r.row_id {
4173 replaced_image = self.get(old_rid, self.snapshot());
4174 self.tombstone_row(
4175 old_rid,
4176 r.committed_epoch,
4177 r.commit_ts,
4178 true,
4179 );
4180 }
4181 }
4182 }
4183 if maintain_pk_by_row {
4184 self.pk_by_row.insert(r.row_id, key);
4185 }
4186 }
4187 }
4188 Some(_) => {}
4189 None => {
4190 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4191 }
4192 }
4193 if let Some(old) = replaced_image {
4194 self.maintain_indexes_on_pk_replace(&old, &r);
4195 } else {
4196 self.index_row(&r);
4197 }
4198 self.reservoir.offer(r.row_id.0);
4199 self.memtable.upsert(r);
4200 self.live_count = self.live_count.saturating_add(1);
4203 }
4204 self.data_generation = self.data_generation.wrapping_add(1);
4205 }
4206
4207 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
4217 for &cid in row.columns.keys() {
4218 self.pending_put_cols.insert(cid);
4219 }
4220 let epoch = row.committed_epoch;
4221 let mut replaced_image: Option<Row> = None;
4222 if let Some(pk_col) = self.schema.primary_key() {
4223 let pk_id = pk_col.id;
4224 if let Some(pk_val) = row.columns.get(&pk_id) {
4225 let maintain_pk_by_row = self.pk_by_row_complete;
4229 if check_existing_pk || maintain_pk_by_row {
4230 let key = self.index_lookup_key(pk_id, pk_val);
4231 if check_existing_pk {
4232 if let Some(old_rid) = self.hot.get(&key) {
4233 if old_rid != row.row_id {
4234 replaced_image = self.get(old_rid, self.snapshot());
4238 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
4239 }
4240 } else if let Some(old) = self.recent_delete_preimages.remove(&key) {
4241 replaced_image = Some(old);
4244 }
4245 }
4246 if maintain_pk_by_row {
4247 self.pk_by_row.insert(row.row_id, key);
4248 }
4249 }
4250 }
4251 } else {
4252 self.hot
4253 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
4254 }
4255 if let Some(old) = replaced_image {
4256 self.maintain_indexes_on_pk_replace(&old, &row);
4257 } else {
4258 self.index_row(&row);
4259 }
4260 self.reservoir.offer(row.row_id.0);
4261 self.memtable.upsert(row);
4262 self.live_count = self.live_count.saturating_add(1);
4263 self.data_generation = self.data_generation.wrapping_add(1);
4264 }
4265
4266 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
4269 let has_partial = self
4270 .schema
4271 .indexes
4272 .iter()
4273 .any(|idx| idx.predicate.is_some());
4274 if has_partial {
4275 self.unindex_bitmap_membership(old);
4279 self.index_row(new);
4280 return;
4281 }
4282
4283 crate::index::maintain_bitmap_secondary_on_replace(
4284 &self.schema,
4285 &mut self.bitmap,
4286 old,
4287 new,
4288 );
4289 self.index_row_non_bitmap(new);
4290 }
4291
4292 fn index_row_non_bitmap(&mut self, row: &Row) {
4295 if row.deleted {
4296 return;
4297 }
4298 let effective = if self.column_keys.is_empty() {
4299 None
4300 } else {
4301 Some(self.tokenized_for_indexes(row))
4302 };
4303 let source = effective.as_ref().unwrap_or(row);
4304 for idef in &self.schema.indexes {
4305 if idef.kind == crate::schema::IndexKind::Bitmap {
4306 continue;
4307 }
4308 index_into_single(
4309 idef,
4310 &self.schema,
4311 source,
4312 &mut self.hot,
4313 &mut self.bitmap,
4314 &mut self.ann,
4315 &mut self.fm,
4316 &mut self.sparse,
4317 &mut self.minhash,
4318 );
4319 }
4320 if let Some(pk_col) = self.schema.primary_key() {
4321 if let Some(pk_val) = source.columns.get(&pk_col.id) {
4322 let key = if self.column_keys.is_empty() {
4323 pk_val.encode_key()
4324 } else {
4325 self.index_lookup_key(pk_col.id, pk_val)
4326 };
4327 self.hot.insert(key, source.row_id);
4328 }
4329 }
4330 }
4331
4332 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
4335 self.allocator.alloc()
4336 }
4337
4338 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
4344 let pk = self.schema.primary_key().ok_or_else(|| {
4345 MongrelError::Schema("clustered table requires a single-column primary key".into())
4346 })?;
4347 let pk_val = columns
4348 .iter()
4349 .find(|(id, _)| *id == pk.id)
4350 .map(|(_, v)| v)
4351 .ok_or_else(|| {
4352 MongrelError::Schema(format!(
4353 "clustered table missing primary key column {} ({})",
4354 pk.id, pk.name
4355 ))
4356 })?;
4357 Ok(clustered_row_id(pk_val))
4358 }
4359
4360 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
4368 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
4369 return Err(MongrelError::Full("row-id namespace exhausted".into()));
4370 }
4371 let n = rows.len();
4372 for r in rows {
4373 for &cid in r.columns.keys() {
4374 self.pending_put_cols.insert(cid);
4375 }
4376 }
4377 let (losers, winner_pks) = self.partition_pk_winners(rows);
4378 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
4379 let group_ts = rows.first().and_then(|r| r.commit_ts);
4381 for (key, &row_id) in &winner_pks {
4382 if let Some(old_rid) = self.hot.get(key) {
4383 if old_rid != row_id {
4384 self.tombstone_row(old_rid, epoch, group_ts, true);
4385 }
4386 }
4387 }
4388 for &loser_rid in &losers {
4391 self.tombstone_row(loser_rid, epoch, group_ts, false);
4392 }
4393 for (key, row_id) in winner_pks {
4395 self.insert_hot_pk(key, row_id);
4396 }
4397 if self.schema.primary_key().is_none() {
4398 for r in rows {
4399 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4400 }
4401 }
4402 for r in rows {
4403 self.allocator.advance_to(r.row_id)?;
4404 if !losers.contains(&r.row_id) {
4405 self.index_row(r);
4406 }
4407 }
4408 for r in rows {
4409 if !losers.contains(&r.row_id) {
4410 self.reservoir.offer(r.row_id.0);
4411 }
4412 }
4413 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
4414 self.data_generation = self.data_generation.wrapping_add(1);
4415 Ok(())
4416 }
4417
4418 pub(crate) fn recover_apply(
4423 &mut self,
4424 rows: Vec<Row>,
4425 deletes: Vec<(RowId, Epoch)>,
4426 ) -> Result<()> {
4427 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
4431 std::collections::BTreeMap::new();
4432 for row in rows {
4433 if row.row_id.0 >= u64::MAX - 1 {
4434 return Err(MongrelError::Full("row-id namespace exhausted".into()));
4435 }
4436 self.allocator.advance_to(row.row_id)?;
4437 if let Some(ai) = self.auto_inc.as_mut() {
4442 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
4443 let next = n.checked_add(1).ok_or_else(|| {
4444 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4445 })?;
4446 if next > ai.next {
4447 ai.next = next;
4448 }
4449 }
4450 }
4451 by_epoch.entry(row.committed_epoch).or_default().push(row);
4452 }
4453 for (epoch, group) in by_epoch {
4454 let (losers, winner_pks) = self.partition_pk_winners(&group);
4455 let group_ts = group.first().and_then(|r| r.commit_ts);
4457 for (key, &row_id) in &winner_pks {
4458 if let Some(old_rid) = self.hot.get(key) {
4459 if old_rid != row_id {
4460 self.tombstone_row(old_rid, epoch, group_ts, false);
4461 }
4462 }
4463 }
4464 for (key, row_id) in winner_pks {
4465 self.insert_hot_pk(key, row_id);
4466 }
4467 if self.schema.primary_key().is_none() {
4468 for r in &group {
4469 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4470 }
4471 }
4472 for r in &group {
4473 if !losers.contains(&r.row_id) {
4474 self.memtable.upsert(r.clone());
4475 self.index_row(r);
4476 }
4477 }
4478 }
4479 for (rid, epoch) in deletes {
4480 self.memtable.tombstone(rid, epoch);
4481 self.remove_hot_for_row(rid, epoch);
4482 }
4483 self.reservoir_complete = false;
4486 Ok(())
4487 }
4488
4489 pub(crate) fn flushed_epoch(&self) -> u64 {
4491 self.flushed_epoch
4492 }
4493
4494 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
4495 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
4496 }
4497
4498 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
4500 self.schema.validate_values(cells)
4501 }
4502
4503 fn validate_columns_not_null(
4507 &self,
4508 columns: &[(u16, columnar::NativeColumn)],
4509 n: usize,
4510 ) -> Result<()> {
4511 let by_id: HashMap<u16, &columnar::NativeColumn> =
4512 columns.iter().map(|(id, c)| (*id, c)).collect();
4513 for col in &self.schema.columns {
4514 if !col.flags.contains(ColumnFlags::NULLABLE) {
4515 match by_id.get(&col.id) {
4516 None => {
4517 return Err(MongrelError::InvalidArgument(format!(
4518 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
4519 col.name, col.id
4520 )));
4521 }
4522 Some(c) => {
4523 if c.null_count(n) != 0 {
4524 return Err(MongrelError::InvalidArgument(format!(
4525 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
4526 col.name, col.id
4527 )));
4528 }
4529 }
4530 }
4531 }
4532 if let TypeId::Enum { variants } = &col.ty {
4533 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
4534 if by_id.contains_key(&col.id) {
4535 return Err(MongrelError::InvalidArgument(format!(
4536 "column '{}' ({}) enum requires a bytes column",
4537 col.name, col.id
4538 )));
4539 }
4540 continue;
4541 };
4542 for index in 0..n {
4543 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
4544 continue;
4545 };
4546 if !variants.iter().any(|variant| variant.as_bytes() == value) {
4547 return Err(MongrelError::InvalidArgument(format!(
4548 "column '{}' ({}) enum value {:?} is not one of {:?}",
4549 col.name,
4550 col.id,
4551 String::from_utf8_lossy(value),
4552 variants
4553 )));
4554 }
4555 }
4556 }
4557 }
4558 Ok(())
4559 }
4560
4561 fn bulk_pk_winner_indices(
4566 &self,
4567 columns: &[(u16, columnar::NativeColumn)],
4568 n: usize,
4569 ) -> Option<Vec<usize>> {
4570 let pk_col = self.schema.primary_key()?;
4571 let pk_id = pk_col.id;
4572 let pk_ty = pk_col.ty.clone();
4573 let by_id: HashMap<u16, &columnar::NativeColumn> =
4574 columns.iter().map(|(id, c)| (*id, c)).collect();
4575 let pk_native = by_id.get(&pk_id)?;
4576 if native_int64_strictly_increasing(pk_native, n) {
4577 return None;
4578 }
4579 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
4581 let mut null_pk_rows: Vec<usize> = Vec::new();
4582 for i in 0..n {
4583 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
4584 Some(key) => {
4585 last.insert(key, i);
4586 }
4587 None => null_pk_rows.push(i),
4588 }
4589 }
4590 let mut winners: HashSet<usize> = last.values().copied().collect();
4591 for i in null_pk_rows {
4592 winners.insert(i);
4593 }
4594 Some((0..n).filter(|i| winners.contains(i)).collect())
4595 }
4596
4597 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
4599 self.require_delete()?;
4600 let epoch = self.pending_epoch();
4601 self.wal_append_data(Op::Delete {
4602 table_id: self.table_id,
4603 row_ids: vec![row_id],
4604 })?;
4605 if self.is_shared() {
4606 self.pending_dels.push(row_id);
4607 } else {
4608 self.apply_delete(row_id, epoch);
4609 }
4610 Ok(())
4611 }
4612
4613 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
4614 let pre = self.get(row_id, self.snapshot());
4615 self.delete(row_id)?;
4616 Ok(pre.map(|row| {
4617 let mut columns: Vec<_> = row.columns.into_iter().collect();
4618 columns.sort_by_key(|(id, _)| *id);
4619 OwnedRow { columns }
4620 }))
4621 }
4622
4623 pub fn truncate(&mut self) -> Result<()> {
4625 self.require_delete()?;
4626 let epoch = self.pending_epoch();
4627 self.wal_append_data(Op::TruncateTable {
4628 table_id: self.table_id,
4629 })?;
4630 self.pending_rows.clear();
4631 self.pending_rows_auto_inc.clear();
4632 self.pending_dels.clear();
4633 self.pending_truncate = Some(epoch);
4634 Ok(())
4635 }
4636
4637 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
4639 self.run_refs.clear();
4644 self.retiring.clear();
4645 self.memtable = Memtable::new();
4646 self.mutable_run = MutableRun::new();
4647 self.hot = HotIndex::new();
4648 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
4649 self.bitmap = bitmap;
4650 self.ann = ann;
4651 self.fm = fm;
4652 self.sparse = sparse;
4653 self.minhash = minhash;
4654 self.learned_range = Arc::new(HashMap::new());
4655 self.pk_by_row.clear();
4656 self.pk_by_row_complete = false;
4657 self.live_count = 0;
4658 self.reservoir = crate::reservoir::Reservoir::default();
4659 self.reservoir_complete = true;
4660 self.had_deletes = true;
4661 self.agg_cache = Arc::new(HashMap::new());
4662 self.global_idx_epoch = 0;
4663 self.indexes_complete = true;
4664 self.pending_delete_rids.clear();
4665 self.pending_put_cols.clear();
4666 self.pending_rows.clear();
4667 self.pending_rows_auto_inc.clear();
4668 self.pending_dels.clear();
4669 self.clear_result_cache();
4670 self.invalidate_index_checkpoint();
4671 self.data_generation = self.data_generation.wrapping_add(1);
4672 }
4673
4674 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
4677 self.apply_delete_at(row_id, epoch, None);
4678 }
4679
4680 pub(crate) fn apply_delete_at(
4682 &mut self,
4683 row_id: RowId,
4684 epoch: Epoch,
4685 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4686 ) {
4687 let preimage = self.get(row_id, self.snapshot());
4692 self.remove_hot_for_row(row_id, epoch);
4693 if let Some(row) = preimage {
4694 if let Some(pk_col) = self.schema.primary_key() {
4695 if let Some(pk_val) = row.columns.get(&pk_col.id) {
4696 let key = self.index_lookup_key(pk_col.id, pk_val);
4697 self.recent_delete_preimages.insert(key, row);
4698 }
4699 }
4700 }
4701 self.tombstone_row(row_id, epoch, commit_ts, true);
4702 self.data_generation = self.data_generation.wrapping_add(1);
4703 }
4704
4705 fn unindex_bitmap_membership(&mut self, row: &Row) {
4710 if row.deleted {
4711 return;
4712 }
4713 for idef in &self.schema.indexes {
4714 if idef.kind != crate::schema::IndexKind::Bitmap {
4715 continue;
4716 }
4717 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
4718 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4719 b.remove(&key, row.row_id);
4720 }
4721 }
4722 }
4723 }
4724
4725 fn union_bitmap_point_i64(
4729 &self,
4730 set: &mut RowIdSet,
4731 column_id: u16,
4732 value: i64,
4733 snapshot: Snapshot,
4734 ) {
4735 let Some(b) = self.bitmap.get(&column_id) else {
4736 return;
4737 };
4738 let encoded = Value::Int64(value).encode_key();
4739 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
4740 for rid in b.get(&lookup).iter() {
4741 set.insert(u64::from(rid));
4742 }
4743 set.remove_many(self.overlay_tombstoned_rids(snapshot));
4747 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
4748 }
4749
4750 fn tombstone_row(
4754 &mut self,
4755 row_id: RowId,
4756 epoch: Epoch,
4757 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
4758 adjust_live_count: bool,
4759 ) {
4760 let tombstone = Row {
4761 row_id,
4762 committed_epoch: epoch,
4763 columns: std::collections::HashMap::new(),
4764 deleted: true,
4765 commit_ts,
4766 };
4767 self.memtable.upsert(tombstone);
4768 self.pk_by_row.remove(&row_id);
4769 if adjust_live_count {
4770 self.live_count = self.live_count.saturating_sub(1);
4771 }
4772 self.pending_delete_rids.insert(row_id.0 as u32);
4774 self.had_deletes = true;
4777 self.agg_cache = Arc::new(HashMap::new());
4778 }
4779
4780 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
4784 let Some(pk_col) = self.schema.primary_key() else {
4785 return;
4786 };
4787 if self.pk_by_row_complete {
4790 if let Some(key) = self.pk_by_row.remove(&row_id) {
4791 if self.hot.get(&key) == Some(row_id) {
4792 self.hot.remove(&key);
4793 }
4794 }
4795 return;
4796 }
4797 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
4816 if self.indexes_complete {
4817 let pk_val = self
4818 .memtable
4819 .get_version(row_id, lookup_epoch)
4820 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4821 .or_else(|| {
4822 self.mutable_run
4823 .get_version(row_id, lookup_epoch)
4824 .filter(|(_, r)| !r.deleted)
4825 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
4826 })
4827 .or_else(|| {
4828 self.run_refs.iter().find_map(|rr| {
4829 let mut reader = self.open_reader(rr.run_id).ok()?;
4830 let (_, deleted, val) = reader
4831 .get_version_column(row_id, lookup_epoch, pk_col.id)
4832 .ok()??;
4833 if deleted {
4834 return None;
4835 }
4836 val
4837 })
4838 });
4839 if let Some(pk_val) = pk_val {
4840 let key = self.index_lookup_key(pk_col.id, &pk_val);
4841 if self.hot.get(&key) == Some(row_id) {
4842 self.hot.remove(&key);
4843 }
4844 return;
4845 }
4846 }
4847 self.refresh_pk_by_row_from_hot();
4852 if let Some(key) = self.pk_by_row.remove(&row_id) {
4853 if self.hot.get(&key) == Some(row_id) {
4854 self.hot.remove(&key);
4855 }
4856 }
4857 }
4858
4859 fn partition_pk_winners(
4864 &self,
4865 rows: &[Row],
4866 ) -> (
4867 std::collections::HashSet<RowId>,
4868 std::collections::HashMap<Vec<u8>, RowId>,
4869 ) {
4870 let mut losers = std::collections::HashSet::new();
4871 let Some(pk_col) = self.schema.primary_key() else {
4872 return (losers, std::collections::HashMap::new());
4873 };
4874 let pk_id = pk_col.id;
4875 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
4876 std::collections::HashMap::new();
4877 for r in rows {
4878 let Some(pk_val) = r.columns.get(&pk_id) else {
4879 continue;
4880 };
4881 let key = self.index_lookup_key(pk_id, pk_val);
4882 if let Some(&old_rid) = winners.get(&key) {
4883 losers.insert(old_rid);
4884 }
4885 winners.insert(key, r.row_id);
4886 }
4887 (losers, winners)
4888 }
4889
4890 fn index_row(&mut self, row: &Row) {
4891 if row.deleted {
4892 return;
4893 }
4894 let any_predicate = self
4902 .schema
4903 .indexes
4904 .iter()
4905 .any(|idx| idx.predicate.is_some());
4906 if any_predicate {
4907 let columns_map: HashMap<u16, &Value> =
4908 row.columns.iter().map(|(k, v)| (*k, v)).collect();
4909 let name_to_id: HashMap<&str, u16> = self
4910 .schema
4911 .columns
4912 .iter()
4913 .map(|c| (c.name.as_str(), c.id))
4914 .collect();
4915 for idx in &self.schema.indexes {
4916 if let Some(pred) = &idx.predicate {
4917 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4918 continue; }
4920 }
4921 index_into_single(
4923 idx,
4924 &self.schema,
4925 row,
4926 &mut self.hot,
4927 &mut self.bitmap,
4928 &mut self.ann,
4929 &mut self.fm,
4930 &mut self.sparse,
4931 &mut self.minhash,
4932 );
4933 }
4934 return;
4935 }
4936 if self.column_keys.is_empty() {
4940 index_into(
4941 &self.schema,
4942 row,
4943 &mut self.hot,
4944 &mut self.bitmap,
4945 &mut self.ann,
4946 &mut self.fm,
4947 &mut self.sparse,
4948 &mut self.minhash,
4949 );
4950 return;
4951 }
4952 let effective_row = self.tokenized_for_indexes(row);
4953 index_into(
4954 &self.schema,
4955 &effective_row,
4956 &mut self.hot,
4957 &mut self.bitmap,
4958 &mut self.ann,
4959 &mut self.fm,
4960 &mut self.sparse,
4961 &mut self.minhash,
4962 );
4963 }
4964
4965 fn tokenized_for_indexes(&self, row: &Row) -> Row {
4971 if self.column_keys.is_empty() {
4972 return row.clone();
4973 }
4974 {
4975 use crate::encryption::SCHEME_HMAC_EQ;
4976 let mut tok = row.clone();
4977 for (&cid, &(_, scheme)) in &self.column_keys {
4978 if scheme != SCHEME_HMAC_EQ {
4979 continue;
4980 }
4981 if let Some(v) = tok.columns.get(&cid).cloned() {
4982 if let Some(t) = self.tokenize_value(cid, &v) {
4983 tok.columns.insert(cid, t);
4984 }
4985 }
4986 }
4987 tok
4988 }
4989 }
4990
4991 pub fn commit(&mut self) -> Result<Epoch> {
4996 self.commit_inner(None)
4997 }
4998
4999 #[doc(hidden)]
5002 pub fn commit_controlled<F>(
5003 &mut self,
5004 control: &crate::ExecutionControl,
5005 mut before_commit: F,
5006 ) -> Result<Epoch>
5007 where
5008 F: FnMut() -> Result<()>,
5009 {
5010 self.commit_inner(Some((control, &mut before_commit)))
5011 }
5012
5013 fn commit_inner(
5014 &mut self,
5015 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5016 ) -> Result<Epoch> {
5017 self.ensure_writable()?;
5018 if !self.has_pending_mutations() {
5019 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
5020 return Err(MongrelError::Full(
5021 "standalone transaction id namespace exhausted".into(),
5022 ));
5023 }
5024 return Ok(self.epoch.visible());
5025 }
5026 self.commit_new_epoch_inner(controlled)
5027 }
5028
5029 fn commit_new_epoch(&mut self) -> Result<Epoch> {
5033 self.commit_new_epoch_inner(None)
5034 }
5035
5036 fn commit_new_epoch_inner(
5037 &mut self,
5038 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5039 ) -> Result<Epoch> {
5040 self.ensure_writable()?;
5041 if self.is_shared() {
5042 self.commit_shared(controlled)
5043 } else {
5044 self.commit_private(controlled)
5045 }
5046 }
5047
5048 fn commit_private(
5050 &mut self,
5051 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5052 ) -> Result<Epoch> {
5053 let commit_lock = Arc::clone(&self.commit_lock);
5057 let _g = commit_lock.lock();
5058 let txn_id = self.ensure_txn_id()?;
5061 if let Some((control, before_commit)) = controlled {
5062 control.checkpoint()?;
5063 before_commit()?;
5064 }
5065 let new_epoch = self.epoch.bump_assigned();
5066 let epoch_authority = Arc::clone(&self.epoch);
5067 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5068 let wal_result = match &mut self.wal {
5072 WalSink::Private(w) => w
5073 .append_txn(
5074 txn_id,
5075 Op::TxnCommit {
5076 epoch: new_epoch.0,
5077 added_runs: Vec::new(),
5078 },
5079 )
5080 .and_then(|_| w.sync()),
5081 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
5082 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
5083 };
5084 if let Err(error) = wal_result {
5085 self.durable_commit_failed = true;
5086 return Err(MongrelError::CommitOutcomeUnknown {
5087 epoch: new_epoch.0,
5088 message: error.to_string(),
5089 });
5090 }
5091 if let Some(epoch) = self.pending_truncate.take() {
5094 self.apply_truncate(epoch);
5095 }
5096 self.invalidate_pending_cache();
5097 let publish_result = self.persist_manifest(new_epoch);
5098 self.epoch.publish_in_order(new_epoch);
5102 epoch_guard.disarm();
5103 if let Err(error) = publish_result {
5104 self.durable_commit_failed = true;
5105 return Err(MongrelError::DurableCommit {
5106 epoch: new_epoch.0,
5107 message: error.to_string(),
5108 });
5109 }
5110 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
5111 self.pending_private_mutations = false;
5112 self.data_generation = self.data_generation.wrapping_add(1);
5113 Ok(new_epoch)
5114 }
5115
5116 fn commit_shared(
5122 &mut self,
5123 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5124 ) -> Result<Epoch> {
5125 use std::sync::atomic::Ordering;
5126 let s = match &self.wal {
5127 WalSink::Shared(s) => s.clone(),
5128 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
5129 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
5130 };
5131 if s.poisoned.load(Ordering::Relaxed) {
5132 return Err(MongrelError::Other(
5133 "database poisoned by fsync error".into(),
5134 ));
5135 }
5136 let commit_lock = Arc::clone(&self.commit_lock);
5143 let _g = commit_lock.lock();
5144 if !self.pending_rows.is_empty() {
5145 match controlled.as_ref() {
5146 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
5147 None => self.prepare_durable_publish()?,
5148 }
5149 }
5150 let txn_id = self.ensure_txn_id()?;
5153 let mut wal = s.wal.lock();
5154 if let Some((control, before_commit)) = controlled {
5155 control.checkpoint()?;
5156 before_commit()?;
5157 }
5158 let new_epoch = self.epoch.bump_assigned();
5159 let epoch_authority = Arc::clone(&self.epoch);
5160 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5161 let commit_ts = s.hlc.now().map_err(|skew| {
5164 MongrelError::Other(format!(
5165 "clock skew rejected commit timestamp allocation: {skew}"
5166 ))
5167 })?;
5168 let commit_seq = match wal.append_commit_at(
5169 txn_id,
5170 new_epoch,
5171 &[],
5172 commit_ts.physical_micros.saturating_mul(1_000),
5173 ) {
5174 Ok(commit_seq) => commit_seq,
5175 Err(error) => {
5176 s.poisoned.store(true, Ordering::Relaxed);
5177 s.lifecycle.poison();
5178 return Err(MongrelError::CommitOutcomeUnknown {
5179 epoch: new_epoch.0,
5180 message: error.to_string(),
5181 });
5182 }
5183 };
5184 drop(wal);
5185 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
5186 s.poisoned.store(true, Ordering::Relaxed);
5187 s.lifecycle.poison();
5188 return Err(MongrelError::CommitOutcomeUnknown {
5189 epoch: new_epoch.0,
5190 message: error.to_string(),
5191 });
5192 }
5193
5194 if self.pending_truncate.take().is_some() {
5197 self.apply_truncate(new_epoch);
5198 }
5199 let mut rows = std::mem::take(&mut self.pending_rows);
5200 if !rows.is_empty() {
5201 for r in &mut rows {
5202 r.committed_epoch = new_epoch;
5203 r.commit_ts = Some(commit_ts);
5204 }
5205 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
5206 let all_auto_generated =
5207 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
5208 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
5209 } else {
5210 self.pending_rows_auto_inc.clear();
5211 }
5212 let dels = std::mem::take(&mut self.pending_dels);
5213 for rid in dels {
5214 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
5215 }
5216
5217 self.invalidate_pending_cache();
5218 let publish_result = self.persist_manifest(new_epoch);
5219 self.epoch.publish_in_order(new_epoch);
5220 epoch_guard.disarm();
5221 let _ = s.change_wake.send(());
5222 if let Err(error) = publish_result {
5223 self.durable_commit_failed = true;
5224 s.poisoned.store(true, Ordering::Relaxed);
5225 s.lifecycle.poison();
5226 return Err(MongrelError::DurableCommit {
5227 epoch: new_epoch.0,
5228 message: error.to_string(),
5229 });
5230 }
5231 self.current_txn_id = 0;
5233 self.data_generation = self.data_generation.wrapping_add(1);
5234 Ok(new_epoch)
5235 }
5236
5237 pub fn flush(&mut self) -> Result<Epoch> {
5245 self.flush_with_outcome().map(|(epoch, _)| epoch)
5246 }
5247
5248 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
5250 self.flush_with_outcome_inner(None)
5251 }
5252
5253 #[doc(hidden)]
5256 pub fn flush_with_outcome_controlled<F>(
5257 &mut self,
5258 control: &crate::ExecutionControl,
5259 mut before_commit: F,
5260 ) -> Result<(Epoch, bool)>
5261 where
5262 F: FnMut() -> Result<()>,
5263 {
5264 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
5265 }
5266
5267 fn flush_with_outcome_inner(
5268 &mut self,
5269 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5270 ) -> Result<(Epoch, bool)> {
5271 match controlled.as_ref() {
5272 Some((control, _)) => {
5273 self.ensure_indexes_complete_controlled(control, || true)?;
5274 }
5275 None => self.ensure_indexes_complete()?,
5276 }
5277 let committed = self.has_pending_mutations();
5278 let epoch = self.commit_inner(controlled)?;
5279 let finish: Result<(Epoch, bool)> = (|| {
5280 let rows = self.memtable.drain_sorted();
5281 if !rows.is_empty() {
5282 self.mutable_run.insert_many(rows);
5283 }
5284 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
5285 self.spill_mutable_run(epoch)?;
5286 self.mark_flushed(epoch)?;
5290 self.persist_manifest(epoch)?;
5291 self.build_learned_ranges()?;
5292 self.checkpoint_indexes(epoch);
5295 }
5296 Ok((epoch, committed))
5299 })();
5300 let outcome = match finish {
5301 Err(error) if committed => Err(MongrelError::DurableCommit {
5302 epoch: epoch.0,
5303 message: error.to_string(),
5304 }),
5305 result => result,
5306 };
5307 if outcome.is_ok() {
5308 let _ = self.publish_read_generation();
5314 }
5315 outcome
5316 }
5317
5318 fn has_pending_mutations(&self) -> bool {
5319 self.pending_private_mutations
5320 || !self.pending_rows.is_empty()
5321 || !self.pending_dels.is_empty()
5322 || self.pending_truncate.is_some()
5323 }
5324
5325 pub fn has_pending_writes(&self) -> bool {
5326 self.has_pending_mutations()
5327 }
5328
5329 pub fn force_flush(&mut self) -> Result<Epoch> {
5338 let saved = self.mutable_run_spill_bytes;
5339 self.mutable_run_spill_bytes = 1;
5340 let result = self.flush();
5341 self.mutable_run_spill_bytes = saved;
5342 result
5343 }
5344
5345 pub fn close(&mut self) -> Result<()> {
5352 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
5353 self.force_flush()?;
5354 }
5355 Ok(())
5356 }
5357
5358 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
5365 let op = Op::Flush {
5366 table_id: self.table_id,
5367 flushed_epoch: epoch.0,
5368 };
5369 match &mut self.wal {
5370 WalSink::Private(w) => {
5371 w.append_system(op)?;
5372 w.sync()?;
5373 }
5374 WalSink::Shared(s) => {
5375 s.wal.lock().append_system(op)?;
5380 }
5381 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
5382 }
5383 self.flushed_epoch = epoch.0;
5384 if matches!(self.wal, WalSink::Private(_)) {
5385 self.rotate_wal(epoch)?;
5386 }
5387 Ok(())
5388 }
5389
5390 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
5394 if self.mutable_run.is_empty() {
5395 return Ok(());
5396 }
5397 let run_id = self.alloc_run_id()?;
5398 let rows = self.mutable_run.drain_sorted();
5399 if rows.is_empty() {
5400 return Ok(());
5401 }
5402 let path = self.run_path(run_id);
5403 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
5404 if let Some(kek) = &self.kek {
5405 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5406 }
5407 let header = match self.create_run_file(run_id)? {
5408 Some(file) => writer.write_file(file, &rows)?,
5409 None => writer.write(&path, &rows)?,
5410 };
5411 self.run_refs.push(RunRef {
5412 run_id: run_id as u128,
5413 level: 0,
5414 epoch_created: epoch.0,
5415 row_count: header.row_count,
5416 });
5417 self.run_row_id_ranges
5418 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
5419 Ok(())
5420 }
5421
5422 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
5426 self.mutable_run_spill_bytes = bytes.max(1);
5427 }
5428
5429 pub fn set_compaction_zstd_level(&mut self, level: i32) {
5433 self.compaction_zstd_level = level;
5434 }
5435
5436 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
5440 self.result_cache.lock().set_max_bytes(max_bytes);
5441 }
5442
5443 pub(crate) fn clear_result_cache(&mut self) {
5447 self.result_cache.lock().clear();
5448 }
5449
5450 pub fn mutable_run_len(&self) -> usize {
5452 self.mutable_run.len()
5453 }
5454
5455 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
5458 self.mutable_run.drain_sorted()
5459 }
5460
5461 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
5463 let mut snapshot = self.mutable_run.clone();
5464 snapshot.drain_sorted()
5465 }
5466
5467 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
5472 self.ensure_writable()?;
5473 let n = batch.len();
5474 if n == 0 {
5475 return Ok(self.current_epoch());
5476 }
5477 for row in &batch {
5478 self.schema.validate_values(row)?;
5479 }
5480 let epoch = self.commit_new_epoch()?;
5481 let live_before = self.live_count;
5482 self.spill_mutable_run(epoch)?;
5486 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
5487 && self.indexes_complete
5488 && self.run_refs.is_empty()
5489 && self.memtable.is_empty()
5490 && self.mutable_run.is_empty();
5491 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
5497 use rayon::prelude::*;
5498 self.schema
5499 .columns
5500 .par_iter()
5501 .map(|cdef| {
5502 (
5503 cdef.id,
5504 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
5505 )
5506 })
5507 .collect::<Vec<_>>()
5508 };
5509 drop(batch);
5510 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
5515 self.validate_columns_not_null(&user_columns, n)?;
5516 let winner_idx = self
5517 .bulk_pk_winner_indices(&user_columns, n)
5518 .filter(|idx| idx.len() != n);
5519 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
5520 match winner_idx.as_deref() {
5521 Some(idx) => {
5522 let compacted = user_columns
5523 .iter()
5524 .map(|(id, c)| (*id, c.gather(idx)))
5525 .collect();
5526 (compacted, idx.len())
5527 }
5528 None => (std::mem::take(&mut user_columns), n),
5529 };
5530 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
5531 let first = self.allocator.alloc_range(write_n as u64)?.0;
5532 for rid in first..first + write_n as u64 {
5533 self.reservoir.offer(rid);
5534 }
5535 let run_id = self.alloc_run_id()?;
5536 let path = self.run_path(run_id);
5537 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
5538 .clean(true)
5539 .with_lz4()
5540 .with_native_endian();
5541 if let Some(kek) = &self.kek {
5542 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5543 }
5544 let header = match self.create_run_file(run_id)? {
5545 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
5546 None => writer.write_native(&path, &write_columns, write_n, first)?,
5547 };
5548 self.run_refs.push(RunRef {
5549 run_id: run_id as u128,
5550 level: 0,
5551 epoch_created: epoch.0,
5552 row_count: header.row_count,
5553 });
5554 self.run_row_id_ranges
5555 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
5556 self.live_count = self.live_count.saturating_add(write_n as u64);
5557 if eager_index_build {
5558 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
5559 self.index_columns_bulk(&write_columns, &row_ids);
5560 self.indexes_complete = true;
5561 self.build_learned_ranges()?;
5562 } else {
5563 self.indexes_complete = false;
5564 }
5565 self.mark_flushed(epoch)?;
5566 self.persist_manifest(epoch)?;
5567 if eager_index_build {
5568 self.checkpoint_indexes(epoch);
5569 }
5570 self.clear_result_cache();
5571 Ok(epoch)
5572 }
5573
5574 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
5577 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
5578 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
5579 let segment_no = segment
5582 .file_stem()
5583 .and_then(|s| s.to_str())
5584 .and_then(|s| s.strip_prefix("seg-"))
5585 .and_then(|s| s.parse::<u64>().ok())
5586 .unwrap_or(0);
5587 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
5588 wal.set_sync_byte_threshold(self.sync_byte_threshold);
5589 wal.sync()?;
5590 self.wal = WalSink::Private(wal);
5591 Ok(())
5592 }
5593
5594 pub(crate) fn invalidate_pending_cache(&mut self) {
5599 self.result_cache
5600 .lock()
5601 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
5602 self.pending_delete_rids.clear();
5603 self.pending_put_cols.clear();
5604 }
5605
5606 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
5607 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
5608 m.current_epoch = epoch.0;
5609 m.next_row_id = self.allocator.current().0;
5610 m.runs = self.run_refs.clone();
5611 m.live_count = self.live_count;
5612 m.global_idx_epoch = self.global_idx_epoch;
5613 m.flushed_epoch = self.flushed_epoch;
5614 m.retiring = self.retiring.clone();
5615 m.auto_inc_next = match self.auto_inc {
5619 Some(ai) if ai.seeded => ai.next,
5620 _ => 0,
5621 };
5622 m.ttl = self.ttl;
5623 let meta_dek = self.manifest_meta_dek();
5624 match self._root_guard.as_deref() {
5625 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
5626 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
5627 }
5628 Ok(())
5629 }
5630
5631 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
5632 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
5636 let live_count = u64::try_from(rows.len())
5637 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
5638 let auto_inc = match self.auto_inc {
5639 Some(mut state) => {
5640 let maximum = self.scan_max_int64(state.column_id)?;
5641 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
5642 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
5643 })?;
5644 state.next = state.next.max(after_maximum).max(1);
5645 state.seeded = true;
5646 Some(state)
5647 }
5648 None => None,
5649 };
5650 Ok(RecoveryMetadataPlan {
5651 live_count,
5652 auto_inc,
5653 changed: live_count != self.live_count
5654 || auto_inc.is_some_and(|planned| {
5655 self.auto_inc.is_none_or(|current| {
5656 current.next != planned.next || current.seeded != planned.seeded
5657 })
5658 }),
5659 })
5660 }
5661
5662 pub(crate) fn apply_recovered_metadata(
5663 &mut self,
5664 plan: RecoveryMetadataPlan,
5665 epoch: Epoch,
5666 ) -> Result<()> {
5667 if !plan.changed {
5668 return Ok(());
5669 }
5670 self.live_count = plan.live_count;
5671 self.auto_inc = plan.auto_inc;
5672 self.persist_manifest(epoch)
5673 }
5674
5675 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
5681 if !self.indexes_complete {
5684 return;
5685 }
5686 if crate::catalog::inject_hook("index.publish.before").is_err() {
5689 return;
5690 }
5691 if self.idx_root.is_none() {
5692 if let Some(root) = self._root_guard.as_ref() {
5693 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
5694 return;
5695 };
5696 self.idx_root = Some(Arc::new(idx_root));
5697 }
5698 }
5699 let snap = global_idx::IndexSnapshot {
5700 hot: &self.hot,
5701 bitmap: &self.bitmap,
5702 ann: &self.ann,
5703 fm: &self.fm,
5704 sparse: &self.sparse,
5705 minhash: &self.minhash,
5706 learned_range: &self.learned_range,
5707 };
5708 let idx_dek = self.idx_dek();
5710 let written = match self.idx_root.as_deref() {
5711 Some(root) => global_idx::write_atomic_root(
5712 root,
5713 self.table_id,
5714 epoch.0,
5715 snap,
5716 idx_dek.as_deref(),
5717 ),
5718 None => global_idx::write_atomic(
5719 &self.dir,
5720 self.table_id,
5721 epoch.0,
5722 snap,
5723 idx_dek.as_deref(),
5724 ),
5725 };
5726 if written.is_ok() {
5727 self.global_idx_epoch = epoch.0;
5728 let _ = self.persist_manifest(epoch);
5729 let _ = crate::catalog::inject_hook("index.publish.after");
5731 }
5732 }
5733
5734 pub(crate) fn invalidate_index_checkpoint(&mut self) {
5737 self.global_idx_epoch = 0;
5738 if let Some(root) = self.idx_root.as_deref() {
5739 let _ = root.remove_file(global_idx::IDX_FILENAME);
5740 } else {
5741 global_idx::remove(&self.dir);
5742 }
5743 let _ = self.persist_manifest(self.epoch.visible());
5744 }
5745
5746 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
5751 self.indexes_complete = false;
5752 self.global_idx_epoch = 0;
5753 if let Some(root) = self.idx_root.as_deref() {
5754 let _ = root.remove_file(global_idx::IDX_FILENAME);
5755 } else {
5756 global_idx::remove(&self.dir);
5757 }
5758 }
5759
5760 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
5761 self.indexes_complete = true;
5762 }
5763
5764 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
5770 self.durable_commit_failed = true;
5771 if let WalSink::Shared(shared) = &self.wal {
5772 shared
5773 .poisoned
5774 .store(true, std::sync::atomic::Ordering::Relaxed);
5775 }
5776 }
5777
5778 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
5782 self.durable_commit_failed = true;
5783 }
5784
5785 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
5794 let mut best: Option<Row> = None;
5795 let mut consider = |row: Row| {
5796 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5797 return;
5798 }
5799 if best.as_ref().is_none_or(|current| {
5800 Snapshot::version_is_newer(
5801 row.committed_epoch,
5802 row.commit_ts,
5803 current.committed_epoch,
5804 current.commit_ts,
5805 )
5806 }) {
5807 best = Some(row);
5808 }
5809 };
5810 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
5811 consider(row);
5812 }
5813 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
5814 consider(row);
5815 }
5816 for rr in &self.run_refs {
5817 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
5820 if row_id.0 < min_rid || row_id.0 > max_rid {
5821 self.lookup_metrics
5822 .get_run_skipped
5823 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5824 continue;
5825 }
5826 }
5827 self.lookup_metrics
5828 .get_run_opened
5829 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
5830 let Ok(mut reader) = self.open_reader(rr.run_id) else {
5831 continue;
5832 };
5833 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
5836 continue;
5837 };
5838 consider(row);
5839 }
5840 let now_nanos = unix_nanos_now();
5841 match best {
5842 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
5843 Some(r) => Some(r),
5844 None => None,
5845 }
5846 }
5847
5848 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
5852 self.visible_rows_at_time(snapshot, unix_nanos_now())
5853 }
5854
5855 #[doc(hidden)]
5858 pub fn visible_rows_controlled(
5859 &self,
5860 snapshot: Snapshot,
5861 control: &crate::ExecutionControl,
5862 ) -> Result<Vec<Row>> {
5863 let mut out = Vec::new();
5864 self.for_each_visible_row_controlled(snapshot, control, |row| {
5865 out.push(row);
5866 Ok(())
5867 })?;
5868 Ok(out)
5869 }
5870
5871 #[doc(hidden)]
5874 pub fn for_each_visible_row_controlled<F>(
5875 &self,
5876 snapshot: Snapshot,
5877 control: &crate::ExecutionControl,
5878 mut visit: F,
5879 ) -> Result<()>
5880 where
5881 F: FnMut(Row) -> Result<()>,
5882 {
5883 let mut sources = Vec::with_capacity(self.run_refs.len() + 2);
5884 control.checkpoint()?;
5885 let memtable_map = self.memtable.newest_visible_map(snapshot);
5888 if !memtable_map.is_empty() {
5889 sources.push(ControlledVisibleSource::memory_from_map(memtable_map));
5890 }
5891 control.checkpoint()?;
5892 let mutable_map = self.mutable_run.newest_visible_map(snapshot);
5895 if !mutable_map.is_empty() {
5896 sources.push(ControlledVisibleSource::memory_from_map(mutable_map));
5897 }
5898 for run in &self.run_refs {
5899 control.checkpoint()?;
5900 let reader = self.open_reader(run.run_id)?;
5901 sources.push(ControlledVisibleSource::run(
5905 reader.into_visible_version_cursor(snapshot.epoch)?,
5906 ));
5907 }
5908 let now_nanos = unix_nanos_now();
5909 merge_controlled_visible_sources(
5910 &mut sources,
5911 control,
5912 |row| self.row_expired_at(row, now_nanos),
5913 |row| {
5914 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5917 return Ok(());
5918 }
5919 visit(row)
5920 },
5921 )
5922 }
5923
5924 #[doc(hidden)]
5925 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
5926 let mut best: HashMap<u64, Row> = HashMap::new();
5927 let mut fold = |row: Row| {
5928 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
5929 return;
5930 }
5931 best.entry(row.row_id.0)
5932 .and_modify(|existing| {
5933 if Snapshot::version_is_newer(
5934 row.committed_epoch,
5935 row.commit_ts,
5936 existing.committed_epoch,
5937 existing.commit_ts,
5938 ) {
5939 *existing = row.clone();
5940 }
5941 })
5942 .or_insert(row);
5943 };
5944 for row in self.memtable.visible_versions_at(snapshot) {
5945 fold(row);
5946 }
5947 for row in self.mutable_run.visible_versions_at(snapshot) {
5948 fold(row);
5949 }
5950 for rr in &self.run_refs {
5951 let mut reader = self.open_reader(rr.run_id)?;
5952 for row in reader.visible_versions(snapshot.epoch)? {
5954 fold(row);
5955 }
5956 }
5957 let mut out: Vec<Row> = best
5958 .into_values()
5959 .filter_map(|r| {
5960 if r.deleted || self.row_expired_at(&r, now_nanos) {
5961 None
5962 } else {
5963 Some(r)
5964 }
5965 })
5966 .collect();
5967 out.sort_by_key(|r| r.row_id);
5968 Ok(out)
5969 }
5970
5971 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
5978 if self.ttl.is_none()
5979 && self.memtable.is_empty()
5980 && self.mutable_run.is_empty()
5981 && self.run_refs.len() == 1
5982 {
5983 let rr = self.run_refs[0].clone();
5984 let mut reader = self.open_reader(rr.run_id)?;
5985 let idxs = reader.visible_indices(snapshot.epoch)?;
5986 let mut cols = Vec::with_capacity(self.schema.columns.len());
5987 for cdef in &self.schema.columns {
5988 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
5989 }
5990 return Ok(cols);
5991 }
5992 let rows = self.visible_rows(snapshot)?;
5994 let mut cols: Vec<(u16, Vec<Value>)> = self
5995 .schema
5996 .columns
5997 .iter()
5998 .map(|c| (c.id, Vec::with_capacity(rows.len())))
5999 .collect();
6000 for r in &rows {
6001 for (cid, vec) in cols.iter_mut() {
6002 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
6003 }
6004 }
6005 Ok(cols)
6006 }
6007
6008 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
6010 let row_id = self.hot.get(key)?;
6011 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
6012 Some(row_id)
6013 } else {
6014 None
6015 }
6016 }
6017
6018 pub fn lookup_metrics_snapshot(&self) -> LookupMetricsSnapshot {
6022 let mut snap = self.lookup_metrics.snapshot();
6023 let (mem, disk, miss, write_us) = self.result_cache.lock().cache_counters();
6024 snap.result_cache_memory_hit = mem;
6025 snap.result_cache_disk_hit = disk;
6026 snap.result_cache_miss = miss;
6027 snap.result_cache_persistent_write_us = write_us;
6028 snap
6029 }
6030
6031 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
6036 self.query_at_with_allowed(q, self.snapshot(), None)
6037 }
6038
6039 pub fn query_controlled(
6042 &mut self,
6043 q: &crate::query::Query,
6044 control: &crate::ExecutionControl,
6045 ) -> Result<Vec<Row>> {
6046 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
6047 }
6048
6049 pub fn query_at_with_allowed(
6052 &mut self,
6053 q: &crate::query::Query,
6054 snapshot: Snapshot,
6055 allowed: Option<&std::collections::HashSet<RowId>>,
6056 ) -> Result<Vec<Row>> {
6057 self.query_at_with_allowed_after(q, snapshot, allowed, None)
6058 }
6059
6060 #[doc(hidden)]
6061 pub fn query_at_with_allowed_controlled(
6062 &mut self,
6063 q: &crate::query::Query,
6064 snapshot: Snapshot,
6065 allowed: Option<&std::collections::HashSet<RowId>>,
6066 control: &crate::ExecutionControl,
6067 ) -> Result<Vec<Row>> {
6068 self.require_select()?;
6069 self.ensure_indexes_complete_controlled(control, || true)?;
6070 self.validate_native_query(q)?;
6071 self.query_conditions_at(
6072 &q.conditions,
6073 snapshot,
6074 allowed,
6075 q.limit,
6076 q.offset,
6077 None,
6078 unix_nanos_now(),
6079 Some(control),
6080 )
6081 }
6082
6083 #[doc(hidden)]
6084 pub fn query_at_with_allowed_after(
6085 &mut self,
6086 q: &crate::query::Query,
6087 snapshot: Snapshot,
6088 allowed: Option<&std::collections::HashSet<RowId>>,
6089 after_row_id: Option<RowId>,
6090 ) -> Result<Vec<Row>> {
6091 self.query_at_with_allowed_after_at_time(
6092 q,
6093 snapshot,
6094 allowed,
6095 after_row_id,
6096 unix_nanos_now(),
6097 )
6098 }
6099
6100 #[doc(hidden)]
6101 pub fn query_at_with_allowed_after_at_time(
6102 &mut self,
6103 q: &crate::query::Query,
6104 snapshot: Snapshot,
6105 allowed: Option<&std::collections::HashSet<RowId>>,
6106 after_row_id: Option<RowId>,
6107 query_time_nanos: i64,
6108 ) -> Result<Vec<Row>> {
6109 self.require_select()?;
6110 self.ensure_indexes_complete()?;
6111 self.validate_native_query(q)?;
6112 self.query_conditions_at(
6113 &q.conditions,
6114 snapshot,
6115 allowed,
6116 q.limit,
6117 q.offset,
6118 after_row_id,
6119 query_time_nanos,
6120 None,
6121 )
6122 }
6123
6124 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
6125 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
6126 return Err(MongrelError::InvalidArgument(format!(
6127 "query exceeds {} conditions",
6128 crate::query::MAX_HARD_CONDITIONS
6129 )));
6130 }
6131 if let Some(limit) = q.limit {
6132 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
6133 return Err(MongrelError::InvalidArgument(format!(
6134 "query limit must be between 1 and {}",
6135 crate::query::MAX_FINAL_LIMIT
6136 )));
6137 }
6138 }
6139 if q.offset > crate::query::MAX_QUERY_OFFSET {
6140 return Err(MongrelError::InvalidArgument(format!(
6141 "query offset exceeds {}",
6142 crate::query::MAX_QUERY_OFFSET
6143 )));
6144 }
6145 Ok(())
6146 }
6147
6148 #[doc(hidden)]
6151 pub fn query_all_at(
6152 &mut self,
6153 conditions: &[crate::query::Condition],
6154 snapshot: Snapshot,
6155 ) -> Result<Vec<Row>> {
6156 self.require_select()?;
6157 self.ensure_indexes_complete()?;
6158 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
6159 return Err(MongrelError::InvalidArgument(format!(
6160 "query exceeds {} conditions",
6161 crate::query::MAX_HARD_CONDITIONS
6162 )));
6163 }
6164 self.query_conditions_at(
6165 conditions,
6166 snapshot,
6167 None,
6168 None,
6169 0,
6170 None,
6171 unix_nanos_now(),
6172 None,
6173 )
6174 }
6175
6176 #[allow(clippy::too_many_arguments)]
6177 fn query_conditions_at(
6178 &self,
6179 conditions: &[crate::query::Condition],
6180 snapshot: Snapshot,
6181 allowed: Option<&std::collections::HashSet<RowId>>,
6182 limit: Option<usize>,
6183 offset: usize,
6184 after_row_id: Option<RowId>,
6185 query_time_nanos: i64,
6186 control: Option<&crate::ExecutionControl>,
6187 ) -> Result<Vec<Row>> {
6188 control
6189 .map(crate::ExecutionControl::checkpoint)
6190 .transpose()?;
6191 crate::trace::QueryTrace::record(|t| {
6192 t.run_count = self.run_refs.len();
6193 t.memtable_rows = self.memtable.len();
6194 t.mutable_run_rows = self.mutable_run.len();
6195 });
6196 if conditions.is_empty() {
6200 crate::trace::QueryTrace::record(|t| {
6201 t.scan_mode = crate::trace::ScanMode::Materialized;
6202 t.row_materialized = true;
6203 });
6204 let mut rows = match control {
6205 Some(control) => self.visible_rows_controlled(snapshot, control)?,
6206 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
6207 };
6208 if let Some(allowed) = allowed {
6209 let mut filtered = Vec::with_capacity(rows.len());
6210 for (index, row) in rows.into_iter().enumerate() {
6211 if index & 255 == 0 {
6212 control
6213 .map(crate::ExecutionControl::checkpoint)
6214 .transpose()?;
6215 }
6216 if allowed.contains(&row.row_id) {
6217 filtered.push(row);
6218 }
6219 }
6220 rows = filtered;
6221 }
6222 if let Some(after_row_id) = after_row_id {
6223 rows.retain(|row| row.row_id > after_row_id);
6224 }
6225 rows.drain(..offset.min(rows.len()));
6226 if let Some(limit) = limit {
6227 rows.truncate(limit);
6228 }
6229 return Ok(rows);
6230 }
6231 crate::trace::QueryTrace::record(|t| {
6232 t.conditions_pushed = conditions.len();
6233 t.scan_mode = crate::trace::ScanMode::Materialized;
6234 t.row_materialized = true;
6235 });
6236 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
6243 ordered.sort_by_key(|c| condition_cost_rank(c));
6244 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
6245 for c in &ordered {
6246 control
6247 .map(crate::ExecutionControl::checkpoint)
6248 .transpose()?;
6249 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
6250 let empty = s.is_empty();
6251 sets.push(s);
6252 if empty {
6253 break;
6254 }
6255 }
6256 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6257 if let Some(allowed) = allowed {
6258 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
6259 }
6260 if let Some(after_row_id) = after_row_id {
6261 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
6262 rids.drain(..first);
6263 }
6264 rids.drain(..offset.min(rids.len()));
6265 if let Some(limit) = limit {
6266 rids.truncate(limit);
6267 }
6268 control
6269 .map(crate::ExecutionControl::checkpoint)
6270 .transpose()?;
6271 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
6272 }
6273
6274 pub fn retrieve(
6276 &mut self,
6277 retriever: &crate::query::Retriever,
6278 ) -> Result<Vec<crate::query::RetrieverHit>> {
6279 self.retrieve_with_allowed(retriever, None)
6280 }
6281
6282 pub fn retrieve_at(
6283 &mut self,
6284 retriever: &crate::query::Retriever,
6285 snapshot: Snapshot,
6286 allowed: Option<&std::collections::HashSet<RowId>>,
6287 ) -> Result<Vec<crate::query::RetrieverHit>> {
6288 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
6289 }
6290
6291 pub fn retrieve_with_allowed(
6294 &mut self,
6295 retriever: &crate::query::Retriever,
6296 allowed: Option<&std::collections::HashSet<RowId>>,
6297 ) -> Result<Vec<crate::query::RetrieverHit>> {
6298 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
6299 }
6300
6301 pub fn retrieve_at_with_allowed(
6302 &mut self,
6303 retriever: &crate::query::Retriever,
6304 snapshot: Snapshot,
6305 allowed: Option<&std::collections::HashSet<RowId>>,
6306 ) -> Result<Vec<crate::query::RetrieverHit>> {
6307 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
6308 }
6309
6310 pub fn retrieve_at_with_allowed_and_context(
6311 &mut self,
6312 retriever: &crate::query::Retriever,
6313 snapshot: Snapshot,
6314 allowed: Option<&std::collections::HashSet<RowId>>,
6315 context: Option<&crate::query::AiExecutionContext>,
6316 ) -> Result<Vec<crate::query::RetrieverHit>> {
6317 self.require_select()?;
6318 self.ensure_indexes_complete()?;
6319 self.validate_retriever(retriever)?;
6320 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
6321 }
6322
6323 pub fn retrieve_at_with_candidate_authorization_and_context(
6324 &mut self,
6325 retriever: &crate::query::Retriever,
6326 snapshot: Snapshot,
6327 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6328 context: Option<&crate::query::AiExecutionContext>,
6329 ) -> Result<Vec<crate::query::RetrieverHit>> {
6330 self.require_select()?;
6331 self.ensure_indexes_complete()?;
6332 self.retrieve_at_with_candidate_authorization_on_generation(
6333 retriever,
6334 snapshot,
6335 authorization,
6336 context,
6337 )
6338 }
6339
6340 #[doc(hidden)]
6341 pub fn retrieve_at_with_candidate_authorization_on_generation(
6342 &self,
6343 retriever: &crate::query::Retriever,
6344 snapshot: Snapshot,
6345 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6346 context: Option<&crate::query::AiExecutionContext>,
6347 ) -> Result<Vec<crate::query::RetrieverHit>> {
6348 self.require_select()?;
6349 self.validate_retriever(retriever)?;
6350 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
6351 }
6352
6353 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
6354 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
6355 let (column_id, k) = match retriever {
6356 Retriever::Ann {
6357 column_id,
6358 query,
6359 k,
6360 } => {
6361 let index = self.ann.get(column_id).ok_or_else(|| {
6362 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
6363 })?;
6364 if query.len() != index.dim() {
6365 return Err(MongrelError::InvalidArgument(format!(
6366 "ANN query dimension must be {}, got {}",
6367 index.dim(),
6368 query.len()
6369 )));
6370 }
6371 if query.iter().any(|value| !value.is_finite()) {
6372 return Err(MongrelError::InvalidArgument(
6373 "ANN query values must be finite".into(),
6374 ));
6375 }
6376 (*column_id, *k)
6377 }
6378 Retriever::Sparse {
6379 column_id,
6380 query,
6381 k,
6382 } => {
6383 if !self.sparse.contains_key(column_id) {
6384 return Err(MongrelError::InvalidArgument(format!(
6385 "column {column_id} has no Sparse index"
6386 )));
6387 }
6388 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
6389 return Err(MongrelError::InvalidArgument(
6390 "Sparse query must be non-empty with finite weights".into(),
6391 ));
6392 }
6393 if query.len() > MAX_SPARSE_TERMS {
6394 return Err(MongrelError::InvalidArgument(format!(
6395 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
6396 )));
6397 }
6398 (*column_id, *k)
6399 }
6400 Retriever::MinHash {
6401 column_id,
6402 members,
6403 k,
6404 } => {
6405 if !self.minhash.contains_key(column_id) {
6406 return Err(MongrelError::InvalidArgument(format!(
6407 "column {column_id} has no MinHash index"
6408 )));
6409 }
6410 if members.is_empty() {
6411 return Err(MongrelError::InvalidArgument(
6412 "MinHash members must not be empty".into(),
6413 ));
6414 }
6415 if members.len() > MAX_SET_MEMBERS {
6416 return Err(MongrelError::InvalidArgument(format!(
6417 "MinHash query exceeds {MAX_SET_MEMBERS} members"
6418 )));
6419 }
6420 let mut total_bytes = 0usize;
6421 for member in members {
6422 let bytes = member.encoded_len();
6423 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
6424 return Err(MongrelError::InvalidArgument(format!(
6425 "MinHash member exceeds {} bytes",
6426 crate::query::MAX_SET_MEMBER_BYTES
6427 )));
6428 }
6429 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
6430 MongrelError::InvalidArgument("MinHash input size overflow".into())
6431 })?;
6432 }
6433 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
6434 return Err(MongrelError::InvalidArgument(format!(
6435 "MinHash input exceeds {} bytes",
6436 crate::query::MAX_SET_INPUT_BYTES
6437 )));
6438 }
6439 (*column_id, *k)
6440 }
6441 };
6442 if k == 0 {
6443 return Err(MongrelError::InvalidArgument(
6444 "retriever k must be > 0".into(),
6445 ));
6446 }
6447 if k > MAX_RETRIEVER_K {
6448 return Err(MongrelError::InvalidArgument(format!(
6449 "retriever k exceeds {MAX_RETRIEVER_K}"
6450 )));
6451 }
6452 debug_assert!(self
6453 .schema
6454 .columns
6455 .iter()
6456 .any(|column| column.id == column_id));
6457 Ok(())
6458 }
6459
6460 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
6461 use crate::query::Condition;
6462 match condition {
6463 Condition::Ann {
6464 column_id,
6465 query,
6466 k,
6467 } => self.validate_retriever(&crate::query::Retriever::Ann {
6468 column_id: *column_id,
6469 query: query.clone(),
6470 k: *k,
6471 }),
6472 Condition::SparseMatch {
6473 column_id,
6474 query,
6475 k,
6476 } => self.validate_retriever(&crate::query::Retriever::Sparse {
6477 column_id: *column_id,
6478 query: query.clone(),
6479 k: *k,
6480 }),
6481 Condition::MinHashSimilar {
6482 column_id,
6483 query,
6484 k,
6485 } => {
6486 if !self.minhash.contains_key(column_id) {
6487 return Err(MongrelError::InvalidArgument(format!(
6488 "column {column_id} has no MinHash index"
6489 )));
6490 }
6491 if query.is_empty() || *k == 0 {
6492 return Err(MongrelError::InvalidArgument(
6493 "MinHash query must be non-empty and k must be > 0".into(),
6494 ));
6495 }
6496 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
6497 {
6498 return Err(MongrelError::InvalidArgument(format!(
6499 "MinHash query must have <= {} members and k <= {}",
6500 crate::query::MAX_SET_MEMBERS,
6501 crate::query::MAX_RETRIEVER_K
6502 )));
6503 }
6504 Ok(())
6505 }
6506 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
6507 Err(MongrelError::InvalidArgument(format!(
6508 "bitmap IN exceeds {} values",
6509 crate::query::MAX_SET_MEMBERS
6510 )))
6511 }
6512 Condition::FmContainsAll { patterns, .. }
6513 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
6514 {
6515 Err(MongrelError::InvalidArgument(format!(
6516 "FM query exceeds {} patterns",
6517 crate::query::MAX_HARD_CONDITIONS
6518 )))
6519 }
6520 _ => Ok(()),
6521 }
6522 }
6523
6524 fn retrieve_filtered(
6525 &self,
6526 retriever: &crate::query::Retriever,
6527 snapshot: Snapshot,
6528 hard_filter: Option<&RowIdSet>,
6529 allowed: Option<&std::collections::HashSet<RowId>>,
6530 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6531 context: Option<&crate::query::AiExecutionContext>,
6532 ) -> Result<Vec<crate::query::RetrieverHit>> {
6533 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
6534 let started = std::time::Instant::now();
6535 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
6536 Retriever::Ann {
6537 column_id,
6538 query,
6539 k,
6540 } => {
6541 let Some(index) = self.ann.get(column_id) else {
6542 return Ok(Vec::new());
6543 };
6544 let cap = ann_candidate_cap(index.len(), context);
6545 if cap == 0 {
6546 return Ok(Vec::new());
6547 }
6548 let mut breadth = (*k).max(1).min(cap);
6549 let mut eligibility = std::collections::HashMap::new();
6550 let mut filtered = loop {
6551 let mut seen = std::collections::HashSet::new();
6552 if let Some(context) = context {
6553 context.checkpoint()?;
6554 }
6555 let raw = index.search_with_context(query, breadth, context)?;
6556 let unchecked: Vec<_> = raw
6557 .iter()
6558 .map(|(row_id, _)| *row_id)
6559 .filter(|row_id| !eligibility.contains_key(row_id))
6560 .filter(|row_id| {
6561 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
6562 && allowed.is_none_or(|allowed| allowed.contains(row_id))
6563 })
6564 .collect();
6565 let eligible = self.eligible_and_authorized_candidate_ids(
6566 &unchecked,
6567 *column_id,
6568 snapshot,
6569 candidate_authorization,
6570 context,
6571 )?;
6572 for row_id in unchecked {
6573 eligibility.insert(row_id, eligible.contains(&row_id));
6574 }
6575 let filtered: Vec<_> = raw
6576 .into_iter()
6577 .filter(|(row_id, _)| {
6578 seen.insert(*row_id)
6579 && eligibility.get(row_id).copied().unwrap_or(false)
6580 })
6581 .map(|(row_id, score)| {
6582 let score = match score {
6583 crate::index::AnnDistance::Hamming(d) => {
6584 RetrieverScore::AnnHammingDistance(d)
6585 }
6586 crate::index::AnnDistance::Cosine(d) => {
6587 RetrieverScore::AnnCosineDistance(d)
6588 }
6589 };
6590 (row_id, score)
6591 })
6592 .collect();
6593 if filtered.len() >= *k || breadth >= cap {
6594 if filtered.len() < *k && index.len() > cap && breadth >= cap {
6595 crate::trace::QueryTrace::record(|trace| {
6596 trace.ann_candidate_cap_hit = true;
6597 });
6598 }
6599 break filtered;
6600 }
6601 breadth = breadth.saturating_mul(2).min(cap);
6602 };
6603 filtered.truncate(*k);
6604 filtered
6605 }
6606 Retriever::Sparse {
6607 column_id,
6608 query,
6609 k,
6610 } => self
6611 .sparse
6612 .get(column_id)
6613 .map(|index| -> Result<Vec<_>> {
6614 let mut breadth = (*k).max(1);
6615 let mut eligibility = std::collections::HashMap::new();
6616 loop {
6617 if let Some(context) = context {
6618 context.checkpoint()?;
6619 }
6620 let raw = index.search_with_context(query, breadth, context)?;
6621 let unchecked: Vec<_> = raw
6622 .iter()
6623 .map(|(row_id, _)| *row_id)
6624 .filter(|row_id| !eligibility.contains_key(row_id))
6625 .filter(|row_id| {
6626 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
6627 && allowed.is_none_or(|allowed| allowed.contains(row_id))
6628 })
6629 .collect();
6630 let eligible = self.eligible_and_authorized_candidate_ids(
6631 &unchecked,
6632 *column_id,
6633 snapshot,
6634 candidate_authorization,
6635 context,
6636 )?;
6637 for row_id in unchecked {
6638 eligibility.insert(row_id, eligible.contains(&row_id));
6639 }
6640 let filtered: Vec<_> = raw
6641 .iter()
6642 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
6643 .take(*k)
6644 .map(|(row_id, score)| {
6645 (*row_id, RetrieverScore::SparseDotProduct(*score))
6646 })
6647 .collect();
6648 if filtered.len() >= *k || raw.len() < breadth {
6649 break Ok(filtered);
6650 }
6651 let next = breadth.saturating_mul(2);
6652 if next == breadth {
6653 break Ok(filtered);
6654 }
6655 breadth = next;
6656 }
6657 })
6658 .transpose()?
6659 .unwrap_or_default(),
6660 Retriever::MinHash {
6661 column_id,
6662 members,
6663 k,
6664 } => self
6665 .minhash
6666 .get(column_id)
6667 .map(|index| -> Result<Vec<_>> {
6668 let mut hashes = Vec::with_capacity(members.len());
6669 for member in members {
6670 if let Some(context) = context {
6671 context.consume(crate::query::work_units(
6672 member.encoded_len(),
6673 crate::query::PARSE_WORK_QUANTUM,
6674 ))?;
6675 }
6676 hashes.push(member.hash_v1());
6677 }
6678 let mut breadth = (*k).max(1);
6679 let mut eligibility = std::collections::HashMap::new();
6680 loop {
6681 if let Some(context) = context {
6682 context.checkpoint()?;
6683 }
6684 let raw = index.search_with_context(&hashes, breadth, context)?;
6685 let unchecked: Vec<_> = raw
6686 .iter()
6687 .map(|(row_id, _)| *row_id)
6688 .filter(|row_id| !eligibility.contains_key(row_id))
6689 .filter(|row_id| {
6690 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
6691 && allowed.is_none_or(|allowed| allowed.contains(row_id))
6692 })
6693 .collect();
6694 let eligible = self.eligible_and_authorized_candidate_ids(
6695 &unchecked,
6696 *column_id,
6697 snapshot,
6698 candidate_authorization,
6699 context,
6700 )?;
6701 for row_id in unchecked {
6702 eligibility.insert(row_id, eligible.contains(&row_id));
6703 }
6704 let filtered: Vec<_> = raw
6705 .iter()
6706 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
6707 .take(*k)
6708 .map(|(row_id, score)| {
6709 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
6710 })
6711 .collect();
6712 if filtered.len() >= *k || raw.len() < breadth {
6713 break Ok(filtered);
6714 }
6715 let next = breadth.saturating_mul(2);
6716 if next == breadth {
6717 break Ok(filtered);
6718 }
6719 breadth = next;
6720 }
6721 })
6722 .transpose()?
6723 .unwrap_or_default(),
6724 };
6725 let elapsed = started.elapsed().as_nanos() as u64;
6726 crate::trace::QueryTrace::record(|trace| {
6727 match retriever {
6728 Retriever::Ann { .. } => {
6729 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
6730 if let Retriever::Ann { column_id, .. } = retriever {
6731 if let Some(index) = self.ann.get(column_id) {
6732 trace.ann_algorithm = Some(index.algorithm());
6733 trace.ann_quantization = Some(index.quantization());
6734 trace.ann_backend = Some(index.backend_name());
6735 }
6736 }
6737 }
6738 Retriever::Sparse { .. } => {
6739 trace.sparse_candidate_nanos =
6740 trace.sparse_candidate_nanos.saturating_add(elapsed)
6741 }
6742 Retriever::MinHash { .. } => {
6743 trace.minhash_candidate_nanos =
6744 trace.minhash_candidate_nanos.saturating_add(elapsed)
6745 }
6746 }
6747 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
6748 });
6749 Ok(scored
6750 .into_iter()
6751 .enumerate()
6752 .map(|(rank, (row_id, score))| RetrieverHit {
6753 row_id,
6754 rank: rank + 1,
6755 score,
6756 })
6757 .collect())
6758 }
6759
6760 fn eligible_candidate_ids(
6761 &self,
6762 candidates: &[RowId],
6763 _column_id: u16,
6764 snapshot: Snapshot,
6765 context: Option<&crate::query::AiExecutionContext>,
6766 ) -> Result<std::collections::HashSet<RowId>> {
6767 if !self.had_deletes
6768 && self.ttl.is_none()
6769 && self.pending_put_cols.is_empty()
6770 && snapshot.epoch == self.snapshot().epoch
6771 {
6772 return Ok(candidates.iter().copied().collect());
6773 }
6774 let mut readers: Vec<_> = self
6775 .run_refs
6776 .iter()
6777 .map(|run| self.open_reader(run.run_id))
6778 .collect::<Result<_>>()?;
6779 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6780 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
6781 for &row_id in candidates {
6782 if let Some(context) = context {
6783 context.consume(1)?;
6784 }
6785 let mem = self.memtable.get_version_at(row_id, snapshot);
6786 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
6787 let overlay = match (mem, mutable) {
6788 (Some(left), Some(right)) => Some(
6789 if Snapshot::version_is_newer(
6790 left.1.committed_epoch,
6791 left.1.commit_ts,
6792 right.1.committed_epoch,
6793 right.1.commit_ts,
6794 ) {
6795 left
6796 } else {
6797 right
6798 },
6799 ),
6800 (Some(value), None) | (None, Some(value)) => Some(value),
6801 (None, None) => None,
6802 };
6803 if let Some((_, row)) = overlay {
6804 if !row.deleted && !self.row_expired_at(&row, now) {
6805 eligible.insert(row_id);
6806 }
6807 continue;
6808 }
6809 let mut best: Option<(Epoch, bool, usize)> = None;
6810 for (index, reader) in readers.iter_mut().enumerate() {
6811 if let Some((epoch, deleted)) =
6812 reader.get_version_visibility(row_id, snapshot.epoch)?
6813 {
6814 if best
6815 .as_ref()
6816 .map(|(best_epoch, ..)| epoch > *best_epoch)
6817 .unwrap_or(true)
6818 {
6819 best = Some((epoch, deleted, index));
6820 }
6821 }
6822 }
6823 let Some((_, false, reader_index)) = best else {
6824 continue;
6825 };
6826 if let Some(ttl) = self.ttl {
6827 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
6828 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
6829 {
6830 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
6831 continue;
6832 }
6833 }
6834 }
6835 eligible.insert(row_id);
6836 }
6837 Ok(eligible)
6838 }
6839
6840 fn eligible_and_authorized_candidate_ids(
6841 &self,
6842 candidates: &[RowId],
6843 column_id: u16,
6844 snapshot: Snapshot,
6845 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6846 context: Option<&crate::query::AiExecutionContext>,
6847 ) -> Result<std::collections::HashSet<RowId>> {
6848 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
6849 let Some(authorization) = authorization else {
6850 return Ok(eligible);
6851 };
6852 let candidates: Vec<_> = eligible.into_iter().collect();
6853 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
6854 }
6855
6856 fn policy_allowed_candidate_ids(
6857 &self,
6858 candidates: &[RowId],
6859 snapshot: Snapshot,
6860 authorization: &crate::security::CandidateAuthorization<'_>,
6861 context: Option<&crate::query::AiExecutionContext>,
6862 ) -> Result<std::collections::HashSet<RowId>> {
6863 let started = std::time::Instant::now();
6864 if candidates.is_empty()
6865 || authorization.principal.is_admin
6866 || !authorization.security.rls_enabled(authorization.table)
6867 {
6868 return Ok(candidates.iter().copied().collect());
6869 }
6870 if let Some(context) = context {
6871 context.checkpoint()?;
6872 }
6873 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
6874 let mut rows: std::collections::HashMap<RowId, Row> = candidates
6875 .iter()
6876 .map(|row_id| {
6877 (
6878 *row_id,
6879 Row {
6880 row_id: *row_id,
6881 committed_epoch: snapshot.epoch,
6882 columns: std::collections::HashMap::new(),
6883 deleted: false,
6884 commit_ts: None,
6885 },
6886 )
6887 })
6888 .collect();
6889 let columns = authorization
6890 .security
6891 .select_policy_columns(authorization.table, authorization.principal);
6892 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
6893 let mut decoded = 0usize;
6894 for column_id in &columns {
6895 if let Some(context) = context {
6896 context.checkpoint()?;
6897 }
6898 for (row_id, value) in self.values_for_rids_batch_at_with_context(
6899 &row_ids, *column_id, snapshot, query_now, context,
6900 )? {
6901 if let Some(row) = rows.get_mut(&row_id) {
6902 row.columns.insert(*column_id, value);
6903 decoded = decoded.saturating_add(1);
6904 }
6905 }
6906 }
6907 if let Some(context) = context {
6908 context.consume(candidates.len().saturating_add(decoded))?;
6909 }
6910 let allowed = rows
6911 .into_values()
6912 .filter_map(|row| {
6913 authorization
6914 .security
6915 .row_allowed(
6916 authorization.table,
6917 crate::security::PolicyCommand::Select,
6918 &row,
6919 authorization.principal,
6920 false,
6921 )
6922 .then_some(row.row_id)
6923 })
6924 .collect();
6925 crate::trace::QueryTrace::record(|trace| {
6926 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
6927 trace.rls_policy_columns_decoded =
6928 trace.rls_policy_columns_decoded.saturating_add(decoded);
6929 trace.authorization_nanos = trace
6930 .authorization_nanos
6931 .saturating_add(started.elapsed().as_nanos() as u64);
6932 });
6933 Ok(allowed)
6934 }
6935
6936 pub fn search(
6938 &mut self,
6939 request: &crate::query::SearchRequest,
6940 ) -> Result<Vec<crate::query::SearchHit>> {
6941 self.search_with_allowed(request, None)
6942 }
6943
6944 pub fn search_at(
6945 &mut self,
6946 request: &crate::query::SearchRequest,
6947 snapshot: Snapshot,
6948 authorized: Option<&std::collections::HashSet<RowId>>,
6949 ) -> Result<Vec<crate::query::SearchHit>> {
6950 self.search_at_with_allowed(request, snapshot, authorized)
6951 }
6952
6953 pub fn search_with_allowed(
6954 &mut self,
6955 request: &crate::query::SearchRequest,
6956 authorized: Option<&std::collections::HashSet<RowId>>,
6957 ) -> Result<Vec<crate::query::SearchHit>> {
6958 self.search_at_with_allowed(request, self.snapshot(), authorized)
6959 }
6960
6961 pub fn search_at_with_allowed(
6962 &mut self,
6963 request: &crate::query::SearchRequest,
6964 snapshot: Snapshot,
6965 authorized: Option<&std::collections::HashSet<RowId>>,
6966 ) -> Result<Vec<crate::query::SearchHit>> {
6967 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
6968 }
6969
6970 pub fn search_at_with_allowed_and_context(
6971 &mut self,
6972 request: &crate::query::SearchRequest,
6973 snapshot: Snapshot,
6974 authorized: Option<&std::collections::HashSet<RowId>>,
6975 context: Option<&crate::query::AiExecutionContext>,
6976 ) -> Result<Vec<crate::query::SearchHit>> {
6977 self.ensure_indexes_complete()?;
6978 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
6979 }
6980
6981 pub fn search_at_with_candidate_authorization_and_context(
6982 &mut self,
6983 request: &crate::query::SearchRequest,
6984 snapshot: Snapshot,
6985 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6986 context: Option<&crate::query::AiExecutionContext>,
6987 ) -> Result<Vec<crate::query::SearchHit>> {
6988 self.ensure_indexes_complete()?;
6989 self.search_at_with_filters_and_context(
6990 request,
6991 snapshot,
6992 None,
6993 authorization,
6994 context,
6995 None,
6996 )
6997 }
6998
6999 #[doc(hidden)]
7000 pub fn search_at_with_candidate_authorization_on_generation(
7001 &self,
7002 request: &crate::query::SearchRequest,
7003 snapshot: Snapshot,
7004 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7005 context: Option<&crate::query::AiExecutionContext>,
7006 ) -> Result<Vec<crate::query::SearchHit>> {
7007 self.search_at_with_filters_and_context(
7008 request,
7009 snapshot,
7010 None,
7011 authorization,
7012 context,
7013 None,
7014 )
7015 }
7016
7017 #[doc(hidden)]
7018 pub fn search_at_with_candidate_authorization_on_generation_after(
7019 &self,
7020 request: &crate::query::SearchRequest,
7021 snapshot: Snapshot,
7022 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7023 context: Option<&crate::query::AiExecutionContext>,
7024 after: Option<crate::query::SearchAfter>,
7025 ) -> Result<Vec<crate::query::SearchHit>> {
7026 self.search_at_with_filters_and_context(
7027 request,
7028 snapshot,
7029 None,
7030 authorization,
7031 context,
7032 after,
7033 )
7034 }
7035
7036 fn search_at_with_filters_and_context(
7037 &self,
7038 request: &crate::query::SearchRequest,
7039 snapshot: Snapshot,
7040 authorized: Option<&std::collections::HashSet<RowId>>,
7041 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7042 context: Option<&crate::query::AiExecutionContext>,
7043 after: Option<crate::query::SearchAfter>,
7044 ) -> Result<Vec<crate::query::SearchHit>> {
7045 use crate::query::{
7046 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
7047 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
7048 };
7049 let total_started = std::time::Instant::now();
7050 let rank_offset = after.map_or(0, |after| after.returned_count);
7051 self.require_select()?;
7052 if request.limit == 0 {
7053 return Err(MongrelError::InvalidArgument(
7054 "search limit must be > 0".into(),
7055 ));
7056 }
7057 if request.limit > MAX_FINAL_LIMIT {
7058 return Err(MongrelError::InvalidArgument(format!(
7059 "search limit exceeds {MAX_FINAL_LIMIT}"
7060 )));
7061 }
7062 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
7063 return Err(MongrelError::InvalidArgument(
7064 "search-after score must be finite".into(),
7065 ));
7066 }
7067 if request.retrievers.is_empty() {
7068 return Err(MongrelError::InvalidArgument(
7069 "search requires at least one retriever".into(),
7070 ));
7071 }
7072 if request.retrievers.len() > MAX_RETRIEVERS {
7073 return Err(MongrelError::InvalidArgument(format!(
7074 "search exceeds {MAX_RETRIEVERS} retrievers"
7075 )));
7076 }
7077 if request.must.len() > MAX_HARD_CONDITIONS {
7078 return Err(MongrelError::InvalidArgument(format!(
7079 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
7080 )));
7081 }
7082 for condition in &request.must {
7083 self.validate_condition(condition)?;
7084 }
7085 if request.must.iter().any(|condition| {
7086 matches!(
7087 condition,
7088 Condition::Ann { .. }
7089 | Condition::SparseMatch { .. }
7090 | Condition::MinHashSimilar { .. }
7091 )
7092 }) {
7093 return Err(MongrelError::InvalidArgument(
7094 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
7095 .into(),
7096 ));
7097 }
7098 let mut names = std::collections::HashSet::new();
7099 for named in &request.retrievers {
7100 if named.name.is_empty()
7101 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
7102 || !names.insert(named.name.as_str())
7103 {
7104 return Err(MongrelError::InvalidArgument(format!(
7105 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
7106 crate::query::MAX_RETRIEVER_NAME_BYTES
7107 )));
7108 }
7109 if !named.weight.is_finite()
7110 || named.weight < 0.0
7111 || named.weight > MAX_RETRIEVER_WEIGHT
7112 {
7113 return Err(MongrelError::InvalidArgument(format!(
7114 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
7115 )));
7116 }
7117 self.validate_retriever(&named.retriever)?;
7118 }
7119 let projection = request
7120 .projection
7121 .clone()
7122 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
7123 if projection.len() > MAX_PROJECTION_COLUMNS {
7124 return Err(MongrelError::InvalidArgument(format!(
7125 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
7126 )));
7127 }
7128 for column_id in &projection {
7129 if !self
7130 .schema
7131 .columns
7132 .iter()
7133 .any(|column| column.id == *column_id)
7134 {
7135 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
7136 }
7137 }
7138 if let Some(crate::query::Rerank::ExactVector {
7139 embedding_column,
7140 query,
7141 candidate_limit,
7142 weight,
7143 ..
7144 }) = &request.rerank
7145 {
7146 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
7147 {
7148 return Err(MongrelError::InvalidArgument(format!(
7149 "rerank candidate_limit must be between search limit and {}",
7150 crate::query::MAX_RETRIEVER_K
7151 )));
7152 }
7153 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
7154 return Err(MongrelError::InvalidArgument(format!(
7155 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
7156 )));
7157 }
7158 let column = self
7159 .schema
7160 .columns
7161 .iter()
7162 .find(|column| column.id == *embedding_column)
7163 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
7164 let crate::schema::TypeId::Embedding { dim } = column.ty else {
7165 return Err(MongrelError::InvalidArgument(format!(
7166 "rerank column {embedding_column} is not an embedding"
7167 )));
7168 };
7169 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
7170 return Err(MongrelError::InvalidArgument(format!(
7171 "rerank query must contain {dim} finite values"
7172 )));
7173 }
7174 }
7175
7176 let hard_filter_started = std::time::Instant::now();
7177 let hard_filter = if request.must.is_empty() {
7178 None
7179 } else {
7180 let mut sets = Vec::with_capacity(request.must.len());
7181 for condition in &request.must {
7182 if let Some(context) = context {
7183 context.checkpoint()?;
7184 }
7185 sets.push(self.resolve_condition(condition, snapshot)?);
7186 }
7187 Some(RowIdSet::intersect_many(sets))
7188 };
7189 crate::trace::QueryTrace::record(|trace| {
7190 trace.hard_filter_nanos = trace
7191 .hard_filter_nanos
7192 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
7193 });
7194 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
7195 return Ok(Vec::new());
7196 }
7197
7198 let constant = match request.fusion {
7199 Fusion::ReciprocalRank { constant } => constant,
7200 };
7201 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
7202 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
7203 let mut fusion_nanos = 0u64;
7204 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
7205 std::collections::HashMap::new();
7206 for named in retrievers {
7207 if named.weight == 0.0 {
7208 continue;
7209 }
7210 if let Some(context) = context {
7211 context.checkpoint()?;
7212 }
7213 let hits = self.retrieve_filtered(
7214 &named.retriever,
7215 snapshot,
7216 hard_filter.as_ref(),
7217 authorized,
7218 candidate_authorization,
7219 context,
7220 )?;
7221 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
7222 let fusion_started = std::time::Instant::now();
7223 for hit in hits {
7224 if let Some(context) = context {
7225 context.consume(1)?;
7226 }
7227 let contribution = named.weight / (constant as f64 + hit.rank as f64);
7228 if !contribution.is_finite() {
7229 return Err(MongrelError::InvalidArgument(
7230 "retriever contribution must be finite".into(),
7231 ));
7232 }
7233 let max_fused_candidates = context.map_or(
7234 crate::query::MAX_FUSED_CANDIDATES,
7235 crate::query::AiExecutionContext::max_fused_candidates,
7236 );
7237 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
7238 return Err(MongrelError::WorkBudgetExceeded);
7239 }
7240 let entry = fused.entry(hit.row_id).or_default();
7241 entry.0 += contribution;
7242 if !entry.0.is_finite() {
7243 return Err(MongrelError::InvalidArgument(
7244 "fused score must be finite".into(),
7245 ));
7246 }
7247 entry.1.push(ComponentScore {
7248 retriever_name: retriever_name.clone(),
7249 rank: hit.rank,
7250 raw_score: hit.score,
7251 contribution,
7252 });
7253 }
7254 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
7255 }
7256 let union_size = fused.len();
7257 let mut ranked: Vec<_> = fused
7258 .into_iter()
7259 .map(|(row_id, (fused_score, components))| {
7260 (row_id, fused_score, components, None, fused_score)
7261 })
7262 .collect();
7263 let order = |(a_row, _, _, _, a_score): &(
7264 RowId,
7265 f64,
7266 Vec<ComponentScore>,
7267 Option<f32>,
7268 f64,
7269 ),
7270 (b_row, _, _, _, b_score): &(
7271 RowId,
7272 f64,
7273 Vec<ComponentScore>,
7274 Option<f32>,
7275 f64,
7276 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
7277 if let Some(crate::query::Rerank::ExactVector {
7278 embedding_column,
7279 query,
7280 metric,
7281 candidate_limit,
7282 weight,
7283 }) = &request.rerank
7284 {
7285 let fused_order = |(a_row, a_score, ..): &(
7286 RowId,
7287 f64,
7288 Vec<ComponentScore>,
7289 Option<f32>,
7290 f64,
7291 ),
7292 (b_row, b_score, ..): &(
7293 RowId,
7294 f64,
7295 Vec<ComponentScore>,
7296 Option<f32>,
7297 f64,
7298 )| {
7299 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
7300 };
7301 let selection_started = std::time::Instant::now();
7302 if let Some(context) = context {
7303 context.consume(ranked.len())?;
7304 }
7305 if ranked.len() > *candidate_limit {
7306 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
7307 ranked.truncate(*candidate_limit);
7308 }
7309 ranked.sort_by(fused_order);
7310 fusion_nanos =
7311 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
7312 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
7313 if let Some(context) = context {
7314 context.consume(row_ids.len())?;
7315 }
7316 let query_now =
7317 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7318 let gather_started = std::time::Instant::now();
7319 let vectors = self.values_for_rids_batch_at_with_context(
7320 &row_ids,
7321 *embedding_column,
7322 snapshot,
7323 query_now,
7324 context,
7325 )?;
7326 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
7327 let vector_work =
7328 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
7329 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
7330 if let Some(context) = context {
7331 context.consume(vector_work)?;
7332 }
7333 query
7334 .iter()
7335 .map(|value| f64::from(*value).powi(2))
7336 .sum::<f64>()
7337 .sqrt()
7338 } else {
7339 0.0
7340 };
7341 let score_started = std::time::Instant::now();
7342 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
7343 for (row_id, value) in vectors {
7344 if let Some(meta) = value.generated_embedding_metadata() {
7345 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
7346 continue;
7347 }
7348 }
7349 let Some(vector) = value.as_embedding() else {
7350 continue;
7351 };
7352 let score = match metric {
7353 crate::query::VectorMetric::DotProduct => {
7354 if let Some(context) = context {
7355 context.consume(vector_work)?;
7356 }
7357 query
7358 .iter()
7359 .zip(vector)
7360 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7361 .sum::<f64>()
7362 }
7363 crate::query::VectorMetric::Cosine => {
7364 if let Some(context) = context {
7365 context.consume(vector_work.saturating_mul(2))?;
7366 }
7367 let dot = query
7368 .iter()
7369 .zip(vector)
7370 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7371 .sum::<f64>();
7372 let norm = vector
7373 .iter()
7374 .map(|value| f64::from(*value).powi(2))
7375 .sum::<f64>()
7376 .sqrt();
7377 if query_norm == 0.0 || norm == 0.0 {
7378 0.0
7379 } else {
7380 dot / (query_norm * norm)
7381 }
7382 }
7383 crate::query::VectorMetric::Euclidean => {
7384 if let Some(context) = context {
7385 context.consume(vector_work)?;
7386 }
7387 query
7388 .iter()
7389 .zip(vector)
7390 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
7391 .sum::<f64>()
7392 .sqrt()
7393 }
7394 };
7395 if !score.is_finite() {
7396 return Err(MongrelError::InvalidArgument(
7397 "exact rerank score must be finite".into(),
7398 ));
7399 }
7400 scores.insert(row_id, score as f32);
7401 }
7402 let mut reranked = Vec::with_capacity(ranked.len());
7403 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
7404 let Some(score) = scores.get(&row_id).copied() else {
7405 continue;
7406 };
7407 let ordering_score = match metric {
7408 crate::query::VectorMetric::Euclidean => -f64::from(score),
7409 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
7410 f64::from(score)
7411 }
7412 };
7413 let final_score = fused_score + *weight * ordering_score;
7414 if !final_score.is_finite() {
7415 return Err(MongrelError::InvalidArgument(
7416 "final rerank score must be finite".into(),
7417 ));
7418 }
7419 reranked.push((row_id, fused_score, components, Some(score), final_score));
7420 }
7421 ranked = reranked;
7422 ranked.sort_by(order);
7423 crate::trace::QueryTrace::record(|trace| {
7424 trace.exact_vector_gather_nanos =
7425 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
7426 trace.exact_vector_score_nanos = trace
7427 .exact_vector_score_nanos
7428 .saturating_add(score_started.elapsed().as_nanos() as u64);
7429 });
7430 }
7431 if let Some(after) = after {
7432 ranked.retain(|(row_id, _, _, _, final_score)| {
7433 final_score.total_cmp(&after.final_score).is_lt()
7434 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
7435 });
7436 }
7437 let projection_started = std::time::Instant::now();
7438 let sentinel = projection
7439 .first()
7440 .copied()
7441 .or_else(|| self.schema.columns.first().map(|column| column.id));
7442 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7443 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
7444 let mut projection_rows = 0usize;
7445 let mut projection_cells = 0usize;
7446 while out.len() < request.limit && !ranked.is_empty() {
7447 if let Some(context) = context {
7448 context.checkpoint()?;
7449 context.consume(ranked.len())?;
7450 }
7451 let needed = request.limit - out.len();
7452 let window_size = ranked
7453 .len()
7454 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
7455 let selection_started = std::time::Instant::now();
7456 let mut remainder = if ranked.len() > window_size {
7457 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
7458 ranked.split_off(window_size)
7459 } else {
7460 Vec::new()
7461 };
7462 ranked.sort_by(order);
7463 fusion_nanos =
7464 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
7465 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
7466 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
7467 if let Some(context) = context {
7468 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
7469 }
7470 projection_rows = projection_rows.saturating_add(row_ids.len());
7471 projection_cells =
7472 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
7473 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
7474 std::collections::HashMap::new();
7475 if let Some(column_id) = sentinel {
7476 for (row_id, value) in self.values_for_rids_batch_at_with_context(
7477 &row_ids, column_id, snapshot, query_now, context,
7478 )? {
7479 cells.entry(row_id).or_default().insert(column_id, value);
7480 }
7481 }
7482 for &column_id in &projection {
7483 if Some(column_id) == sentinel {
7484 continue;
7485 }
7486 for (row_id, value) in self.values_for_rids_batch_at_with_context(
7487 &row_ids, column_id, snapshot, query_now, context,
7488 )? {
7489 cells.entry(row_id).or_default().insert(column_id, value);
7490 }
7491 }
7492 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
7493 ranked.drain(..)
7494 {
7495 let Some(row_cells) = cells.remove(&row_id) else {
7496 continue;
7497 };
7498 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
7499 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
7500 out.push(SearchHit {
7501 row_id,
7502 cells: projection
7503 .iter()
7504 .filter_map(|column_id| {
7505 row_cells
7506 .get(column_id)
7507 .cloned()
7508 .map(|value| (*column_id, value))
7509 })
7510 .collect(),
7511 components,
7512 fused_score,
7513 exact_rerank_score,
7514 final_score,
7515 final_rank,
7516 });
7517 if out.len() == request.limit {
7518 break;
7519 }
7520 }
7521 ranked.append(&mut remainder);
7522 }
7523 crate::trace::QueryTrace::record(|trace| {
7524 trace.union_size = union_size;
7525 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
7526 trace.projection_nanos = trace
7527 .projection_nanos
7528 .saturating_add(projection_started.elapsed().as_nanos() as u64);
7529 trace.total_nanos = trace
7530 .total_nanos
7531 .saturating_add(total_started.elapsed().as_nanos() as u64);
7532 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
7533 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
7534 if let Some(context) = context {
7535 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
7536 }
7537 });
7538 Ok(out)
7539 }
7540
7541 pub fn set_similarity(
7544 &mut self,
7545 request: &crate::query::SetSimilarityRequest,
7546 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7547 self.set_similarity_with_allowed(request, None)
7548 }
7549
7550 pub fn set_similarity_at(
7551 &mut self,
7552 request: &crate::query::SetSimilarityRequest,
7553 snapshot: Snapshot,
7554 allowed: Option<&std::collections::HashSet<RowId>>,
7555 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7556 self.set_similarity_explained_at(request, snapshot, allowed)
7557 .map(|(hits, _)| hits)
7558 }
7559
7560 pub fn ann_rerank(
7562 &mut self,
7563 request: &crate::query::AnnRerankRequest,
7564 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7565 self.ann_rerank_with_allowed(request, None)
7566 }
7567
7568 pub fn ann_rerank_with_allowed(
7569 &mut self,
7570 request: &crate::query::AnnRerankRequest,
7571 allowed: Option<&std::collections::HashSet<RowId>>,
7572 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7573 self.ann_rerank_at(request, self.snapshot(), allowed)
7574 }
7575
7576 pub fn ann_rerank_at(
7577 &mut self,
7578 request: &crate::query::AnnRerankRequest,
7579 snapshot: Snapshot,
7580 allowed: Option<&std::collections::HashSet<RowId>>,
7581 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7582 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
7583 }
7584
7585 pub fn ann_rerank_at_with_context(
7586 &mut self,
7587 request: &crate::query::AnnRerankRequest,
7588 snapshot: Snapshot,
7589 allowed: Option<&std::collections::HashSet<RowId>>,
7590 context: Option<&crate::query::AiExecutionContext>,
7591 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7592 self.ensure_indexes_complete()?;
7593 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
7594 }
7595
7596 pub fn ann_rerank_at_with_candidate_authorization_and_context(
7597 &mut self,
7598 request: &crate::query::AnnRerankRequest,
7599 snapshot: Snapshot,
7600 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7601 context: Option<&crate::query::AiExecutionContext>,
7602 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7603 self.ensure_indexes_complete()?;
7604 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
7605 }
7606
7607 #[doc(hidden)]
7608 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
7609 &self,
7610 request: &crate::query::AnnRerankRequest,
7611 snapshot: Snapshot,
7612 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7613 context: Option<&crate::query::AiExecutionContext>,
7614 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7615 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
7616 }
7617
7618 fn ann_rerank_at_with_filters_and_context(
7619 &self,
7620 request: &crate::query::AnnRerankRequest,
7621 snapshot: Snapshot,
7622 allowed: Option<&std::collections::HashSet<RowId>>,
7623 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7624 context: Option<&crate::query::AiExecutionContext>,
7625 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7626 use crate::query::{
7627 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
7628 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
7629 };
7630 if request.candidate_k == 0 || request.limit == 0 {
7631 return Err(MongrelError::InvalidArgument(
7632 "candidate_k and limit must be > 0".into(),
7633 ));
7634 }
7635 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
7636 return Err(MongrelError::InvalidArgument(format!(
7637 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
7638 )));
7639 }
7640 let retriever = Retriever::Ann {
7641 column_id: request.column_id,
7642 query: request.query.clone(),
7643 k: request.candidate_k,
7644 };
7645 self.require_select()?;
7646 self.validate_retriever(&retriever)?;
7647 let hits = self.retrieve_filtered(
7648 &retriever,
7649 snapshot,
7650 None,
7651 allowed,
7652 candidate_authorization,
7653 context,
7654 )?;
7655 let distances: std::collections::HashMap<_, _> = hits
7656 .iter()
7657 .filter_map(|hit| match hit.score {
7658 RetrieverScore::AnnHammingDistance(distance) => {
7659 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
7660 }
7661 RetrieverScore::AnnCosineDistance(distance) => {
7662 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
7663 }
7664 _ => None,
7665 })
7666 .collect();
7667 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7668 if let Some(context) = context {
7669 context.consume(row_ids.len())?;
7670 }
7671 let gather_started = std::time::Instant::now();
7672 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7673 let values = self.values_for_rids_batch_at_with_context(
7674 &row_ids,
7675 request.column_id,
7676 snapshot,
7677 query_now,
7678 context,
7679 )?;
7680 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
7681 let score_started = std::time::Instant::now();
7682 let vector_work =
7683 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
7684 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
7685 if let Some(context) = context {
7686 context.consume(vector_work)?;
7687 }
7688 request
7689 .query
7690 .iter()
7691 .map(|value| f64::from(*value).powi(2))
7692 .sum::<f64>()
7693 .sqrt()
7694 } else {
7695 0.0
7696 };
7697 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
7698 for (row_id, value) in values {
7699 if let Some(meta) = value.generated_embedding_metadata() {
7700 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
7701 continue;
7702 }
7703 }
7704 let Some(vector) = value.as_embedding() else {
7705 continue;
7706 };
7707 let exact_score = match request.metric {
7708 VectorMetric::DotProduct => {
7709 if let Some(context) = context {
7710 context.consume(vector_work)?;
7711 }
7712 request
7713 .query
7714 .iter()
7715 .zip(vector)
7716 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7717 .sum::<f64>()
7718 }
7719 VectorMetric::Cosine => {
7720 if let Some(context) = context {
7721 context.consume(vector_work.saturating_mul(2))?;
7722 }
7723 let dot = request
7724 .query
7725 .iter()
7726 .zip(vector)
7727 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7728 .sum::<f64>();
7729 let norm = vector
7730 .iter()
7731 .map(|value| f64::from(*value).powi(2))
7732 .sum::<f64>()
7733 .sqrt();
7734 if query_norm == 0.0 || norm == 0.0 {
7735 0.0
7736 } else {
7737 dot / (query_norm * norm)
7738 }
7739 }
7740 VectorMetric::Euclidean => {
7741 if let Some(context) = context {
7742 context.consume(vector_work)?;
7743 }
7744 request
7745 .query
7746 .iter()
7747 .zip(vector)
7748 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
7749 .sum::<f64>()
7750 .sqrt()
7751 }
7752 };
7753 let exact_score = exact_score as f32;
7754 if !exact_score.is_finite() {
7755 return Err(MongrelError::InvalidArgument(
7756 "exact ANN score must be finite".into(),
7757 ));
7758 }
7759 let Some(candidate_distance) = distances.get(&row_id).copied() else {
7760 continue;
7761 };
7762 reranked.push(AnnRerankHit {
7763 row_id,
7764 candidate_distance,
7765 exact_score,
7766 });
7767 }
7768 reranked.sort_by(|left, right| {
7769 let score = match request.metric {
7770 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
7771 VectorMetric::Cosine | VectorMetric::DotProduct => {
7772 right.exact_score.total_cmp(&left.exact_score)
7773 }
7774 };
7775 score.then_with(|| left.row_id.cmp(&right.row_id))
7776 });
7777 reranked.truncate(request.limit);
7778 crate::trace::QueryTrace::record(|trace| {
7779 trace.exact_vector_gather_nanos =
7780 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
7781 trace.exact_vector_score_nanos = trace
7782 .exact_vector_score_nanos
7783 .saturating_add(score_started.elapsed().as_nanos() as u64);
7784 });
7785 Ok(reranked)
7786 }
7787
7788 pub fn set_similarity_with_allowed(
7789 &mut self,
7790 request: &crate::query::SetSimilarityRequest,
7791 allowed: Option<&std::collections::HashSet<RowId>>,
7792 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7793 self.set_similarity_explained_at(request, self.snapshot(), allowed)
7794 .map(|(hits, _)| hits)
7795 }
7796
7797 pub fn set_similarity_explained(
7798 &mut self,
7799 request: &crate::query::SetSimilarityRequest,
7800 ) -> Result<(
7801 Vec<crate::query::SetSimilarityHit>,
7802 crate::query::SetSimilarityTrace,
7803 )> {
7804 self.set_similarity_explained_at(request, self.snapshot(), None)
7805 }
7806
7807 fn set_similarity_explained_at(
7808 &mut self,
7809 request: &crate::query::SetSimilarityRequest,
7810 snapshot: Snapshot,
7811 allowed: Option<&std::collections::HashSet<RowId>>,
7812 ) -> Result<(
7813 Vec<crate::query::SetSimilarityHit>,
7814 crate::query::SetSimilarityTrace,
7815 )> {
7816 self.ensure_indexes_complete()?;
7817 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
7818 }
7819
7820 pub fn set_similarity_at_with_context(
7821 &mut self,
7822 request: &crate::query::SetSimilarityRequest,
7823 snapshot: Snapshot,
7824 allowed: Option<&std::collections::HashSet<RowId>>,
7825 context: Option<&crate::query::AiExecutionContext>,
7826 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7827 self.ensure_indexes_complete()?;
7828 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
7829 .map(|(hits, _)| hits)
7830 }
7831
7832 pub fn set_similarity_at_with_candidate_authorization_and_context(
7833 &mut self,
7834 request: &crate::query::SetSimilarityRequest,
7835 snapshot: Snapshot,
7836 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7837 context: Option<&crate::query::AiExecutionContext>,
7838 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7839 self.ensure_indexes_complete()?;
7840 self.set_similarity_explained_at_with_context(
7841 request,
7842 snapshot,
7843 None,
7844 authorization,
7845 context,
7846 )
7847 .map(|(hits, _)| hits)
7848 }
7849
7850 #[doc(hidden)]
7851 pub fn set_similarity_at_with_candidate_authorization_on_generation(
7852 &self,
7853 request: &crate::query::SetSimilarityRequest,
7854 snapshot: Snapshot,
7855 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7856 context: Option<&crate::query::AiExecutionContext>,
7857 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7858 self.set_similarity_explained_at_with_context(
7859 request,
7860 snapshot,
7861 None,
7862 authorization,
7863 context,
7864 )
7865 .map(|(hits, _)| hits)
7866 }
7867
7868 fn set_similarity_explained_at_with_context(
7869 &self,
7870 request: &crate::query::SetSimilarityRequest,
7871 snapshot: Snapshot,
7872 allowed: Option<&std::collections::HashSet<RowId>>,
7873 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7874 context: Option<&crate::query::AiExecutionContext>,
7875 ) -> Result<(
7876 Vec<crate::query::SetSimilarityHit>,
7877 crate::query::SetSimilarityTrace,
7878 )> {
7879 use crate::query::{
7880 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
7881 MAX_SET_MEMBERS,
7882 };
7883 let mut trace = crate::query::SetSimilarityTrace::default();
7884 if request.members.is_empty() {
7885 return Ok((Vec::new(), trace));
7886 }
7887 if request.candidate_k == 0 || request.limit == 0 {
7888 return Err(MongrelError::InvalidArgument(
7889 "candidate_k and limit must be > 0".into(),
7890 ));
7891 }
7892 if request.candidate_k > MAX_RETRIEVER_K
7893 || request.limit > MAX_FINAL_LIMIT
7894 || request.members.len() > MAX_SET_MEMBERS
7895 {
7896 return Err(MongrelError::InvalidArgument(format!(
7897 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
7898 )));
7899 }
7900 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
7901 return Err(MongrelError::InvalidArgument(
7902 "min_jaccard must be finite and between 0 and 1".into(),
7903 ));
7904 }
7905 let started = std::time::Instant::now();
7906 let retriever = Retriever::MinHash {
7907 column_id: request.column_id,
7908 members: request.members.clone(),
7909 k: request.candidate_k,
7910 };
7911 self.require_select()?;
7912 self.validate_retriever(&retriever)?;
7913 let hits = self.retrieve_filtered(
7914 &retriever,
7915 snapshot,
7916 None,
7917 allowed,
7918 candidate_authorization,
7919 context,
7920 )?;
7921 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
7922 trace.candidate_count = hits.len();
7923 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
7924 if let Some(context) = context {
7925 context.consume(row_ids.len())?;
7926 }
7927 let started = std::time::Instant::now();
7928 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7929 let values = self.values_for_rids_batch_at_with_context(
7930 &row_ids,
7931 request.column_id,
7932 snapshot,
7933 query_now,
7934 context,
7935 )?;
7936 trace.gather_us = started.elapsed().as_micros() as u64;
7937 if let Some(context) = context {
7938 context.consume(request.members.len())?;
7939 }
7940 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
7941 let estimates: std::collections::HashMap<_, _> = hits
7942 .into_iter()
7943 .filter_map(|hit| match hit.score {
7944 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
7945 _ => None,
7946 })
7947 .collect();
7948 let started = std::time::Instant::now();
7949 let mut parsed = Vec::with_capacity(values.len());
7950 for (row_id, value) in values {
7951 let Value::Bytes(bytes) = value else {
7952 continue;
7953 };
7954 if let Some(context) = context {
7955 context.consume(crate::query::work_units(
7956 bytes.len(),
7957 crate::query::PARSE_WORK_QUANTUM,
7958 ))?;
7959 }
7960 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
7961 continue;
7962 };
7963 if let Some(context) = context {
7964 context.consume(members.len())?;
7965 }
7966 let stored = members
7967 .into_iter()
7968 .filter_map(|member| match member {
7969 serde_json::Value::String(value) => {
7970 Some(crate::query::SetMember::String(value))
7971 }
7972 serde_json::Value::Number(value) => {
7973 Some(crate::query::SetMember::Number(value))
7974 }
7975 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
7976 _ => None,
7977 })
7978 .collect::<std::collections::HashSet<_>>();
7979 parsed.push((row_id, stored));
7980 }
7981 trace.parse_us = started.elapsed().as_micros() as u64;
7982 trace.verified_count = parsed.len();
7983 let started = std::time::Instant::now();
7984 let mut exact = Vec::new();
7985 for (row_id, stored) in parsed {
7986 if let Some(context) = context {
7987 context.consume(query.len().saturating_add(stored.len()))?;
7988 }
7989 let union = query.union(&stored).count();
7990 let score = if union == 0 {
7991 1.0
7992 } else {
7993 query.intersection(&stored).count() as f32 / union as f32
7994 };
7995 if score >= request.min_jaccard {
7996 exact.push(SetSimilarityHit {
7997 row_id,
7998 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
7999 exact_jaccard: score,
8000 });
8001 }
8002 }
8003 exact.sort_by(|a, b| {
8004 b.exact_jaccard
8005 .total_cmp(&a.exact_jaccard)
8006 .then_with(|| a.row_id.cmp(&b.row_id))
8007 });
8008 exact.truncate(request.limit);
8009 trace.score_us = started.elapsed().as_micros() as u64;
8010 crate::trace::QueryTrace::record(|query_trace| {
8011 query_trace.exact_set_gather_nanos = query_trace
8012 .exact_set_gather_nanos
8013 .saturating_add(trace.gather_us.saturating_mul(1_000));
8014 query_trace.exact_set_parse_nanos = query_trace
8015 .exact_set_parse_nanos
8016 .saturating_add(trace.parse_us.saturating_mul(1_000));
8017 query_trace.exact_set_score_nanos = query_trace
8018 .exact_set_score_nanos
8019 .saturating_add(trace.score_us.saturating_mul(1_000));
8020 });
8021 Ok((exact, trace))
8022 }
8023
8024 fn values_for_rids_batch_at(
8026 &self,
8027 row_ids: &[u64],
8028 column_id: u16,
8029 snapshot: Snapshot,
8030 now: i64,
8031 ) -> Result<Vec<(RowId, Value)>> {
8032 if self.ttl.is_none()
8033 && self.memtable.is_empty()
8034 && self.mutable_run.is_empty()
8035 && self.run_refs.len() == 1
8036 {
8037 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8038 if row_ids.len().saturating_mul(24) < reader.row_count() {
8043 let mut values = Vec::with_capacity(row_ids.len());
8044 for &raw_row_id in row_ids {
8045 let row_id = RowId(raw_row_id);
8046 if let Some((_, false, Some(value))) =
8047 reader.get_version_column(row_id, snapshot.epoch, column_id)?
8048 {
8049 values.push((row_id, value));
8050 }
8051 }
8052 return Ok(values);
8053 }
8054 let (positions, visible_row_ids) =
8055 reader.visible_positions_with_rids(snapshot.epoch)?;
8056 let requested: Vec<(RowId, usize)> = row_ids
8057 .iter()
8058 .filter_map(|raw| {
8059 visible_row_ids
8060 .binary_search(&(*raw as i64))
8061 .ok()
8062 .map(|index| (RowId(*raw), positions[index]))
8063 })
8064 .collect();
8065 let values = reader.gather_column(
8066 column_id,
8067 &requested
8068 .iter()
8069 .map(|(_, position)| *position)
8070 .collect::<Vec<_>>(),
8071 )?;
8072 return Ok(requested
8073 .into_iter()
8074 .zip(values)
8075 .map(|((row_id, _), value)| (row_id, value))
8076 .collect());
8077 }
8078 self.values_for_rids_at(row_ids, column_id, snapshot, now)
8079 }
8080
8081 fn values_for_rids_batch_at_with_context(
8082 &self,
8083 row_ids: &[u64],
8084 column_id: u16,
8085 snapshot: Snapshot,
8086 now: i64,
8087 context: Option<&crate::query::AiExecutionContext>,
8088 ) -> Result<Vec<(RowId, Value)>> {
8089 let Some(context) = context else {
8090 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
8091 };
8092 let mut values = Vec::with_capacity(row_ids.len());
8093 for chunk in row_ids.chunks(256) {
8094 context.checkpoint()?;
8095 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
8096 }
8097 Ok(values)
8098 }
8099
8100 fn values_for_rids_at(
8102 &self,
8103 row_ids: &[u64],
8104 column_id: u16,
8105 snapshot: Snapshot,
8106 now: i64,
8107 ) -> Result<Vec<(RowId, Value)>> {
8108 let mut readers: Vec<_> = self
8109 .run_refs
8110 .iter()
8111 .map(|run| self.open_reader(run.run_id))
8112 .collect::<Result<_>>()?;
8113 let mut out = Vec::with_capacity(row_ids.len());
8114 for &raw_row_id in row_ids {
8115 let row_id = RowId(raw_row_id);
8116 let mem = self.memtable.get_version_at(row_id, snapshot);
8117 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
8118 let overlay = match (mem, mutable) {
8119 (Some((_, a)), Some((_, b))) => Some(
8120 if Snapshot::version_is_newer(
8121 a.committed_epoch,
8122 a.commit_ts,
8123 b.committed_epoch,
8124 b.commit_ts,
8125 ) {
8126 a
8127 } else {
8128 b
8129 },
8130 ),
8131 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
8132 (None, None) => None,
8133 };
8134 if let Some(row) = overlay {
8135 if !row.deleted && !self.row_expired_at(&row, now) {
8136 if let Some(value) = row.columns.get(&column_id) {
8137 out.push((row_id, value.clone()));
8138 }
8139 }
8140 continue;
8141 }
8142
8143 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
8144 for (index, reader) in readers.iter_mut().enumerate() {
8145 if let Some((epoch, deleted, value)) =
8146 reader.get_version_column(row_id, snapshot.epoch, column_id)?
8147 {
8148 if best
8149 .as_ref()
8150 .map(|(best_epoch, ..)| epoch > *best_epoch)
8151 .unwrap_or(true)
8152 {
8153 best = Some((epoch, deleted, value, index));
8154 }
8155 }
8156 }
8157 let Some((_, false, Some(value), reader_index)) = best else {
8158 continue;
8159 };
8160 if let Some(ttl) = self.ttl {
8161 if ttl.column_id != column_id {
8162 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
8163 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
8164 {
8165 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8166 continue;
8167 }
8168 }
8169 } else if let Value::Int64(timestamp) = value {
8170 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8171 continue;
8172 }
8173 }
8174 }
8175 out.push((row_id, value));
8176 }
8177 Ok(out)
8178 }
8179
8180 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
8185 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
8186 }
8187
8188 pub fn rows_for_rids_with_context(
8189 &self,
8190 rids: &[u64],
8191 snapshot: Snapshot,
8192 context: &crate::query::AiExecutionContext,
8193 ) -> Result<Vec<Row>> {
8194 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
8195 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
8196 }
8197
8198 fn rows_for_rids_at_time(
8199 &self,
8200 rids: &[u64],
8201 snapshot: Snapshot,
8202 ttl_now: i64,
8203 control: Option<&crate::ExecutionControl>,
8204 ) -> Result<Vec<Row>> {
8205 use std::collections::HashMap;
8206 let mut rows = Vec::with_capacity(rids.len());
8207 let tier_size = self.memtable.len() + self.mutable_run.len();
8222 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
8223 if rids.len().saturating_mul(24) < tier_size {
8224 for &rid in rids {
8225 if overlay.len() & 255 == 0 {
8226 control
8227 .map(crate::ExecutionControl::checkpoint)
8228 .transpose()?;
8229 }
8230 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
8231 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
8232 let newest = match (mem, mrun) {
8233 (Some((_, mr)), Some((_, rr))) => Some(
8234 if Snapshot::version_is_newer(
8235 mr.committed_epoch,
8236 mr.commit_ts,
8237 rr.committed_epoch,
8238 rr.commit_ts,
8239 ) {
8240 mr
8241 } else {
8242 rr
8243 },
8244 ),
8245 (Some((_, mr)), None) => Some(mr),
8246 (None, Some((_, rr))) => Some(rr),
8247 (None, None) => None,
8248 };
8249 if let Some(row) = newest {
8250 overlay.insert(rid, row);
8251 }
8252 }
8253 } else {
8254 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
8255 overlay
8256 .entry(row.row_id.0)
8257 .and_modify(|e| {
8258 if Snapshot::version_is_newer(
8259 row.committed_epoch,
8260 row.commit_ts,
8261 e.committed_epoch,
8262 e.commit_ts,
8263 ) {
8264 *e = row.clone();
8265 }
8266 })
8267 .or_insert(row);
8268 };
8269 for (index, row) in self
8270 .memtable
8271 .visible_versions_at(snapshot)
8272 .into_iter()
8273 .enumerate()
8274 {
8275 if index & 255 == 0 {
8276 control
8277 .map(crate::ExecutionControl::checkpoint)
8278 .transpose()?;
8279 }
8280 fold_newest(row, &mut overlay);
8281 }
8282 for (index, row) in self
8283 .mutable_run
8284 .visible_versions_at(snapshot)
8285 .into_iter()
8286 .enumerate()
8287 {
8288 if index & 255 == 0 {
8289 control
8290 .map(crate::ExecutionControl::checkpoint)
8291 .transpose()?;
8292 }
8293 fold_newest(row, &mut overlay);
8294 }
8295 }
8296 if self.run_refs.len() == 1 {
8297 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8298 if rids.len().saturating_mul(24) < reader.row_count() {
8306 for (index, &rid) in rids.iter().enumerate() {
8307 if index & 255 == 0 {
8308 control
8309 .map(crate::ExecutionControl::checkpoint)
8310 .transpose()?;
8311 }
8312 if let Some(r) = overlay.get(&rid) {
8313 if !r.deleted {
8314 rows.push(r.clone());
8315 }
8316 continue;
8317 }
8318 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
8319 if !row.deleted {
8320 rows.push(row);
8321 }
8322 }
8323 }
8324 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8325 return Ok(rows);
8326 }
8327 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8336 enum Src {
8339 Overlay,
8340 Run,
8341 }
8342 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
8343 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
8344 for (index, rid) in rids.iter().enumerate() {
8345 if index & 255 == 0 {
8346 control
8347 .map(crate::ExecutionControl::checkpoint)
8348 .transpose()?;
8349 }
8350 if overlay.contains_key(rid) {
8351 plan.push(Src::Overlay);
8352 continue;
8353 }
8354 match vis_rids.binary_search(&(*rid as i64)) {
8355 Ok(i) => {
8356 plan.push(Src::Run);
8357 fetch.push(positions[i]);
8358 }
8359 Err(_) => { }
8360 }
8361 }
8362 let fetched = reader.materialize_batch(&fetch)?;
8363 let mut fetched_iter = fetched.into_iter();
8364 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
8365 if index & 255 == 0 {
8366 control
8367 .map(crate::ExecutionControl::checkpoint)
8368 .transpose()?;
8369 }
8370 match src {
8371 Src::Overlay => {
8372 if let Some(r) = overlay.get(rid) {
8373 if !r.deleted {
8374 rows.push(r.clone());
8375 }
8376 }
8377 }
8378 Src::Run => {
8379 if let Some(row) = fetched_iter.next() {
8380 if !row.deleted {
8381 rows.push(row);
8382 }
8383 }
8384 }
8385 }
8386 }
8387 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8388 return Ok(rows);
8389 }
8390 let mut readers: Vec<_> = self
8394 .run_refs
8395 .iter()
8396 .map(|rr| self.open_reader(rr.run_id))
8397 .collect::<Result<Vec<_>>>()?;
8398 for (index, rid) in rids.iter().enumerate() {
8399 if index & 255 == 0 {
8400 control
8401 .map(crate::ExecutionControl::checkpoint)
8402 .transpose()?;
8403 }
8404 if let Some(r) = overlay.get(rid) {
8405 if !r.deleted {
8406 rows.push(r.clone());
8407 }
8408 continue;
8409 }
8410 let mut best: Option<(Epoch, Row)> = None;
8411 for reader in readers.iter_mut() {
8412 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
8413 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
8414 best = Some((epoch, row));
8415 }
8416 }
8417 }
8418 if let Some((_, r)) = best {
8419 if !r.deleted {
8420 rows.push(r);
8421 }
8422 }
8423 }
8424 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8425 Ok(rows)
8426 }
8427
8428 pub fn indexes_complete(&self) -> bool {
8438 self.indexes_complete
8439 }
8440
8441 pub fn index_build_policy(&self) -> IndexBuildPolicy {
8443 self.index_build_policy
8444 }
8445
8446 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
8450 self.index_build_policy = policy;
8451 }
8452
8453 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
8458 if !self.indexes_complete {
8462 return None;
8463 }
8464 let b = self.bitmap.get(&column_id)?;
8465 let result: Vec<Vec<u8>> = b
8466 .keys()
8467 .into_iter()
8468 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
8469 .collect();
8470 Some(result)
8471 }
8472
8473 pub fn fk_join_row_ids(
8474 &self,
8475 fk_column_id: u16,
8476 pk_values: &[Vec<u8>],
8477 fk_conditions: &[crate::query::Condition],
8478 snapshot: Snapshot,
8479 ) -> Result<Vec<u64>> {
8480 let Some(b) = self.bitmap.get(&fk_column_id) else {
8481 return Ok(Vec::new());
8482 };
8483 let mut join_set = {
8484 let mut acc = roaring::RoaringBitmap::new();
8485 for v in pk_values {
8486 acc |= b.get(v);
8487 }
8488 RowIdSet::from_roaring(acc)
8489 };
8490 if !fk_conditions.is_empty() {
8491 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
8492 sets.push(join_set);
8493 for c in fk_conditions {
8494 sets.push(self.resolve_condition(c, snapshot)?);
8495 }
8496 join_set = RowIdSet::intersect_many(sets);
8497 }
8498 Ok(join_set.into_sorted_vec())
8499 }
8500
8501 pub fn fk_join_count(
8507 &self,
8508 fk_column_id: u16,
8509 pk_values: &[Vec<u8>],
8510 fk_conditions: &[crate::query::Condition],
8511 snapshot: Snapshot,
8512 ) -> Result<u64> {
8513 let Some(b) = self.bitmap.get(&fk_column_id) else {
8514 return Ok(0);
8515 };
8516 let mut acc = roaring::RoaringBitmap::new();
8517 for v in pk_values {
8518 acc |= b.get(v);
8519 }
8520 if fk_conditions.is_empty() {
8521 return Ok(acc.len());
8522 }
8523 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
8524 sets.push(RowIdSet::from_roaring(acc));
8525 for c in fk_conditions {
8526 sets.push(self.resolve_condition(c, snapshot)?);
8527 }
8528 Ok(RowIdSet::intersect_many(sets).len() as u64)
8529 }
8530
8531 fn resolve_condition(
8536 &self,
8537 c: &crate::query::Condition,
8538 snapshot: Snapshot,
8539 ) -> Result<RowIdSet> {
8540 self.resolve_condition_with_allowed(c, snapshot, None)
8541 }
8542
8543 fn resolve_condition_with_allowed(
8544 &self,
8545 c: &crate::query::Condition,
8546 snapshot: Snapshot,
8547 allowed: Option<&std::collections::HashSet<RowId>>,
8548 ) -> Result<RowIdSet> {
8549 use crate::query::Condition;
8550 self.validate_condition(c)?;
8551 Ok(match c {
8552 Condition::Pk(key) => {
8553 let lookup = self
8554 .schema
8555 .primary_key()
8556 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
8557 .unwrap_or_else(|| key.clone());
8558 if let Some(r) = self.hot.get(&lookup) {
8559 self.lookup_metrics
8560 .hot_lookup_hit
8561 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8562 RowIdSet::one(r.0)
8563 } else if let Some(pk_col) = self.schema.primary_key() {
8564 self.lookup_metrics
8570 .hot_lookup_fallback
8571 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
8572 self.pk_equality_fallback(pk_col.id, &lookup, snapshot)?
8573 } else {
8574 RowIdSet::empty()
8575 }
8576 }
8577 Condition::BitmapEq { column_id, value } => {
8578 let lookup = self.index_lookup_key_bytes(*column_id, value);
8579 self.bitmap
8580 .get(column_id)
8581 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
8582 .unwrap_or_else(RowIdSet::empty)
8583 }
8584 Condition::BitmapIn { column_id, values } => {
8585 let bm = self.bitmap.get(column_id);
8586 let mut acc = roaring::RoaringBitmap::new();
8587 if let Some(b) = bm {
8588 for v in values {
8589 let lookup = self.index_lookup_key_bytes(*column_id, v);
8590 acc |= b.get(&lookup);
8591 }
8592 }
8593 RowIdSet::from_roaring(acc)
8594 }
8595 Condition::BytesPrefix { column_id, prefix } => {
8596 if let Some(b) = self.bitmap.get(column_id) {
8601 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
8602 let mut acc = roaring::RoaringBitmap::new();
8603 for key in b.keys() {
8604 if key.starts_with(&lookup_prefix) {
8605 acc |= b.get(&key);
8606 }
8607 }
8608 RowIdSet::from_roaring(acc)
8609 } else {
8610 RowIdSet::empty()
8611 }
8612 }
8613 Condition::FmContains { column_id, pattern } => self
8614 .fm
8615 .get(column_id)
8616 .map(|f| {
8617 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
8618 })
8619 .unwrap_or_else(RowIdSet::empty),
8620 Condition::FmContainsAll {
8621 column_id,
8622 patterns,
8623 } => {
8624 if let Some(f) = self.fm.get(column_id) {
8627 let sets: Vec<RowIdSet> = patterns
8628 .iter()
8629 .map(|pat| {
8630 RowIdSet::from_unsorted(
8631 f.locate(pat).into_iter().map(|r| r.0).collect(),
8632 )
8633 })
8634 .collect();
8635 RowIdSet::intersect_many(sets)
8636 } else {
8637 RowIdSet::empty()
8638 }
8639 }
8640 Condition::Ann {
8641 column_id,
8642 query,
8643 k,
8644 } => RowIdSet::from_unsorted(
8645 self.retrieve_filtered(
8646 &crate::query::Retriever::Ann {
8647 column_id: *column_id,
8648 query: query.clone(),
8649 k: *k,
8650 },
8651 snapshot,
8652 None,
8653 allowed,
8654 None,
8655 None,
8656 )?
8657 .into_iter()
8658 .map(|hit| hit.row_id.0)
8659 .collect(),
8660 ),
8661 Condition::SparseMatch {
8662 column_id,
8663 query,
8664 k,
8665 } => RowIdSet::from_unsorted(
8666 self.retrieve_filtered(
8667 &crate::query::Retriever::Sparse {
8668 column_id: *column_id,
8669 query: query.clone(),
8670 k: *k,
8671 },
8672 snapshot,
8673 None,
8674 allowed,
8675 None,
8676 None,
8677 )?
8678 .into_iter()
8679 .map(|hit| hit.row_id.0)
8680 .collect(),
8681 ),
8682 Condition::MinHashSimilar {
8683 column_id,
8684 query,
8685 k,
8686 } => match self.minhash.get(column_id) {
8687 Some(index) => {
8688 let candidates = index.candidate_row_ids(query);
8689 let eligible =
8690 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
8691 RowIdSet::from_unsorted(
8692 index
8693 .search_filtered(query, *k, |row_id| {
8694 eligible.contains(&row_id)
8695 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
8696 })
8697 .into_iter()
8698 .map(|(row_id, _)| row_id.0)
8699 .collect(),
8700 )
8701 }
8702 None => RowIdSet::empty(),
8703 },
8704 Condition::Range { column_id, lo, hi } => {
8705 let mut set = if let Some(li) = self.learned_range.get(column_id) {
8724 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
8725 } else if self.run_refs.len() == 1 {
8726 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8727 r.range_row_id_set_i64(*column_id, *lo, *hi)?
8728 } else {
8729 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
8732 if lo == hi {
8733 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
8734 }
8735 return Ok(multi);
8736 };
8737 set.remove_many(self.overlay_rid_set(snapshot));
8738 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
8739 if lo == hi {
8740 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
8741 }
8742 set
8743 }
8744 Condition::RangeF64 {
8745 column_id,
8746 lo,
8747 lo_inclusive,
8748 hi,
8749 hi_inclusive,
8750 } => {
8751 let mut set = if let Some(li) = self.learned_range.get(column_id) {
8754 RowIdSet::from_unsorted(
8755 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
8756 .into_iter()
8757 .collect(),
8758 )
8759 } else if self.run_refs.len() == 1 {
8760 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8761 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
8762 } else {
8763 return self.range_scan_f64(
8764 *column_id,
8765 *lo,
8766 *lo_inclusive,
8767 *hi,
8768 *hi_inclusive,
8769 snapshot,
8770 );
8771 };
8772 set.remove_many(self.overlay_rid_set(snapshot));
8773 self.range_scan_overlay_f64(
8774 &mut set,
8775 *column_id,
8776 *lo,
8777 *lo_inclusive,
8778 *hi,
8779 *hi_inclusive,
8780 snapshot,
8781 );
8782 set
8783 }
8784 Condition::IsNull { column_id } => {
8785 let mut set = if self.run_refs.len() == 1 {
8786 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8787 r.null_row_id_set(*column_id, true)?
8788 } else {
8789 return self.null_scan(*column_id, true, snapshot);
8790 };
8791 set.remove_many(self.overlay_rid_set(snapshot));
8792 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
8793 set
8794 }
8795 Condition::IsNotNull { column_id } => {
8796 let mut set = if self.run_refs.len() == 1 {
8797 let mut r = self.open_reader(self.run_refs[0].run_id)?;
8798 r.null_row_id_set(*column_id, false)?
8799 } else {
8800 return self.null_scan(*column_id, false, snapshot);
8801 };
8802 set.remove_many(self.overlay_rid_set(snapshot));
8803 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
8804 set
8805 }
8806 })
8807 }
8808
8809 fn range_scan_i64(
8817 &self,
8818 column_id: u16,
8819 lo: i64,
8820 hi: i64,
8821 snapshot: Snapshot,
8822 ) -> Result<RowIdSet> {
8823 let mut row_ids = Vec::new();
8824 let overlay_rids = self.overlay_rid_set(snapshot);
8825 for rr in &self.run_refs {
8826 let mut reader = self.open_reader(rr.run_id)?;
8827 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
8828 for rid in matched {
8829 if !overlay_rids.contains(&rid) {
8830 row_ids.push(rid);
8831 }
8832 }
8833 }
8834 let mut s = RowIdSet::from_unsorted(row_ids);
8835 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
8836 Ok(s)
8837 }
8838
8839 fn range_scan_f64(
8842 &self,
8843 column_id: u16,
8844 lo: f64,
8845 lo_inclusive: bool,
8846 hi: f64,
8847 hi_inclusive: bool,
8848 snapshot: Snapshot,
8849 ) -> Result<RowIdSet> {
8850 let mut row_ids = Vec::new();
8851 let overlay_rids = self.overlay_rid_set(snapshot);
8852 for rr in &self.run_refs {
8853 let mut reader = self.open_reader(rr.run_id)?;
8854 let matched = reader.range_row_ids_visible_f64(
8855 column_id,
8856 lo,
8857 lo_inclusive,
8858 hi,
8859 hi_inclusive,
8860 snapshot.epoch,
8861 )?;
8862 for rid in matched {
8863 if !overlay_rids.contains(&rid) {
8864 row_ids.push(rid);
8865 }
8866 }
8867 }
8868 let mut s = RowIdSet::from_unsorted(row_ids);
8869 self.range_scan_overlay_f64(
8870 &mut s,
8871 column_id,
8872 lo,
8873 lo_inclusive,
8874 hi,
8875 hi_inclusive,
8876 snapshot,
8877 );
8878 Ok(s)
8879 }
8880
8881 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
8883 let mut s = HashSet::new();
8884 for row in self.memtable.visible_versions_at(snapshot) {
8885 s.insert(row.row_id.0);
8886 }
8887 for row in self.mutable_run.visible_versions_at(snapshot) {
8888 s.insert(row.row_id.0);
8889 }
8890 s
8891 }
8892
8893 fn range_scan_overlay_i64(
8894 &self,
8895 s: &mut RowIdSet,
8896 column_id: u16,
8897 lo: i64,
8898 hi: i64,
8899 snapshot: Snapshot,
8900 ) {
8901 let mut newest: HashMap<u64, Row> = HashMap::new();
8908 for r in self.mutable_run.visible_versions_at(snapshot) {
8909 newest.insert(r.row_id.0, r);
8910 }
8911 for r in self.memtable.visible_versions_at(snapshot) {
8912 newest
8913 .entry(r.row_id.0)
8914 .and_modify(|cur| {
8915 if Snapshot::version_is_newer(
8916 r.committed_epoch,
8917 r.commit_ts,
8918 cur.committed_epoch,
8919 cur.commit_ts,
8920 ) {
8921 *cur = r.clone();
8922 }
8923 })
8924 .or_insert(r);
8925 }
8926 for row in newest.values() {
8927 if !row.deleted {
8928 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
8929 if *v >= lo && *v <= hi {
8930 s.insert(row.row_id.0);
8931 }
8932 }
8933 }
8934 }
8935 }
8936
8937 #[allow(clippy::too_many_arguments)]
8938 fn range_scan_overlay_f64(
8939 &self,
8940 s: &mut RowIdSet,
8941 column_id: u16,
8942 lo: f64,
8943 lo_inclusive: bool,
8944 hi: f64,
8945 hi_inclusive: bool,
8946 snapshot: Snapshot,
8947 ) {
8948 let mut newest: HashMap<u64, Row> = HashMap::new();
8951 for r in self.mutable_run.visible_versions_at(snapshot) {
8952 newest.insert(r.row_id.0, r);
8953 }
8954 for r in self.memtable.visible_versions_at(snapshot) {
8955 newest
8956 .entry(r.row_id.0)
8957 .and_modify(|cur| {
8958 if Snapshot::version_is_newer(
8959 r.committed_epoch,
8960 r.commit_ts,
8961 cur.committed_epoch,
8962 cur.commit_ts,
8963 ) {
8964 *cur = r.clone();
8965 }
8966 })
8967 .or_insert(r);
8968 }
8969 for row in newest.values() {
8970 if !row.deleted {
8971 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
8972 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
8973 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
8974 if ok_lo && ok_hi {
8975 s.insert(row.row_id.0);
8976 }
8977 }
8978 }
8979 }
8980 }
8981
8982 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
8985 let mut row_ids = Vec::new();
8986 let overlay_rids = self.overlay_rid_set(snapshot);
8987 for rr in &self.run_refs {
8988 let mut reader = self.open_reader(rr.run_id)?;
8989 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
8990 for rid in matched {
8991 if !overlay_rids.contains(&rid) {
8992 row_ids.push(rid);
8993 }
8994 }
8995 }
8996 let mut s = RowIdSet::from_unsorted(row_ids);
8997 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
8998 Ok(s)
8999 }
9000
9001 fn null_scan_overlay(
9005 &self,
9006 s: &mut RowIdSet,
9007 column_id: u16,
9008 want_nulls: bool,
9009 snapshot: Snapshot,
9010 ) {
9011 let mut newest: HashMap<u64, Row> = HashMap::new();
9012 for r in self.mutable_run.visible_versions_at(snapshot) {
9013 newest.insert(r.row_id.0, r);
9014 }
9015 for r in self.memtable.visible_versions_at(snapshot) {
9016 newest
9017 .entry(r.row_id.0)
9018 .and_modify(|cur| {
9019 if Snapshot::version_is_newer(
9020 r.committed_epoch,
9021 r.commit_ts,
9022 cur.committed_epoch,
9023 cur.commit_ts,
9024 ) {
9025 *cur = r.clone();
9026 }
9027 })
9028 .or_insert(r);
9029 }
9030 for row in newest.values() {
9031 if row.deleted {
9032 continue;
9033 }
9034 let is_null = !row.columns.contains_key(&column_id)
9035 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
9036 if is_null == want_nulls {
9037 s.insert(row.row_id.0);
9038 }
9039 }
9040 }
9041
9042 pub fn snapshot(&self) -> Snapshot {
9043 let epoch = self.epoch.visible();
9044 match &self.wal {
9048 WalSink::Shared(shared) => match shared.hlc.now() {
9049 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
9050 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
9052 },
9053 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
9054 }
9055 }
9056
9057 pub fn data_generation(&self) -> u64 {
9059 self.data_generation
9060 }
9061
9062 pub(crate) fn bump_data_generation(&mut self) {
9063 self.data_generation = self.data_generation.wrapping_add(1);
9064 }
9065
9066 pub fn table_id(&self) -> u64 {
9068 self.table_id
9069 }
9070
9071 fn seal_generations(&mut self) {
9077 self.memtable.seal();
9078 self.mutable_run.seal();
9079 self.hot.seal();
9080 for index in self.bitmap.values_mut() {
9081 index.seal();
9082 }
9083 for index in self.ann.values_mut() {
9084 index.seal();
9085 }
9086 for index in self.fm.values_mut() {
9087 index.seal();
9088 }
9089 for index in self.sparse.values_mut() {
9090 index.seal();
9091 }
9092 for index in self.minhash.values_mut() {
9093 index.seal();
9094 }
9095 self.pk_by_row.seal();
9096 }
9097
9098 fn capture_read_generation(&self) -> ReadGeneration {
9103 let visible_through = self.current_epoch();
9104 ReadGeneration {
9105 schema: Arc::new(self.schema.clone()),
9106 base_runs: Arc::new(self.run_refs.clone()),
9107 deltas: TableDeltas {
9108 memtable: self.memtable.clone(),
9109 mutable_run: self.mutable_run.clone(),
9110 hot: self.hot.clone(),
9111 pk_by_row: self.pk_by_row.clone(),
9112 },
9113 indexes: Arc::new(IndexGeneration::capture(
9114 &self.bitmap,
9115 &self.learned_range,
9116 &self.fm,
9117 &self.ann,
9118 &self.sparse,
9119 &self.minhash,
9120 visible_through,
9121 match &self.wal {
9123 WalSink::Shared(shared) => shared
9124 .hlc
9125 .now()
9126 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
9127 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
9128 },
9129 )),
9130 visible_through,
9131 }
9132 }
9133
9134 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
9139 self.ensure_indexes_complete()?;
9140 self.seal_generations();
9141 let view = Arc::new(self.capture_read_generation());
9142 self.published.store(Arc::clone(&view));
9143 Ok(view)
9144 }
9145
9146 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
9152 self.published.load_full()
9153 }
9154
9155 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
9157 &self.pins
9158 }
9159
9160 pub fn version_gc_floor(&self) -> Epoch {
9165 self.min_active_snapshot()
9166 .unwrap_or_else(|| self.current_epoch())
9167 }
9168
9169 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
9176 let mut report = self.pins.report();
9177 let transaction_floor = [
9178 self.pinned.keys().next().copied(),
9179 self.snapshots.min_pinned(),
9180 ]
9181 .into_iter()
9182 .flatten()
9183 .min();
9184 if let Some(epoch) = transaction_floor {
9185 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
9186 }
9187 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
9188 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
9189 }
9190 report
9191 }
9192
9193 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
9194 self.publish_read_generation()?;
9195 let mut generation = self.clone();
9196 generation.read_only = true;
9197 generation.wal = WalSink::ReadOnly;
9198 generation.pending_delete_rids.clear();
9199 generation.pending_put_cols.clear();
9200 generation.pending_rows.clear();
9201 generation.pending_rows_auto_inc.clear();
9202 generation.pending_dels.clear();
9203 generation.pending_truncate = None;
9204 generation.agg_cache = Arc::new(HashMap::new());
9205 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
9208 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
9211 crate::retention::PinSource::ReadGeneration,
9212 self.current_epoch(),
9213 )));
9214 Ok(generation)
9215 }
9216
9217 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
9218 (std::mem::size_of::<Self>() as u64)
9219 .saturating_add(self.memtable.approx_bytes())
9220 .saturating_add(self.mutable_run.approx_bytes())
9221 .saturating_add(self.live_count.saturating_mul(64))
9222 }
9223
9224 pub fn pin_snapshot(&mut self) -> Snapshot {
9231 let snap = self.snapshot();
9232 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
9233 snap
9234 }
9235
9236 fn active_pin_epochs_for_rebuild(&self) -> Vec<Epoch> {
9248 let mut set = BTreeSet::new();
9249 for epoch in self.pinned.keys().copied() {
9250 set.insert(epoch);
9251 }
9252 for epoch in self.snapshots.live_pinned_epochs() {
9253 set.insert(epoch);
9254 }
9255 for epoch in self.pins.live_pin_epochs() {
9256 set.insert(epoch);
9257 }
9258 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
9259 set.insert(floor);
9260 }
9261 set.into_iter().collect()
9262 }
9263
9264 pub fn hlc_gc_floor(
9274 &self,
9275 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
9276 ) -> crate::epoch::GcFloor {
9277 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
9278 epoch
9279 .and_then(&mut project_epoch)
9280 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
9281 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
9282 };
9283 let transaction = [
9284 self.pinned.keys().next().copied(),
9285 self.snapshots.min_pinned(),
9286 ]
9287 .into_iter()
9288 .flatten()
9289 .min();
9290 let history = self.snapshots.history_floor(self.current_epoch());
9291 crate::epoch::GcFloor {
9292 transaction_snapshot: project(transaction),
9293 history_retention: project(history),
9294 backup_pitr: project(
9295 self.pins
9296 .oldest_for(crate::retention::PinSource::BackupPitr),
9297 ),
9298 replication: project(
9299 self.pins
9300 .oldest_for(crate::retention::PinSource::Replication),
9301 ),
9302 read_generation: project(
9303 self.pins
9304 .oldest_for(crate::retention::PinSource::ReadGeneration),
9305 ),
9306 online_index_build: project(
9307 self.pins
9308 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
9309 ),
9310 }
9311 }
9312
9313 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
9315 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
9316 *count -= 1;
9317 if *count == 0 {
9318 self.pinned.remove(&snap.epoch);
9319 }
9320 }
9321 }
9322
9323 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
9335 let local = self.pinned.keys().next().copied();
9336 let global = self.snapshots.min_pinned();
9337 let history = self.snapshots.history_floor(self.current_epoch());
9338 let pinned = self.pins.oldest_pinned();
9339 [local, global, history, pinned].into_iter().flatten().min()
9340 }
9341
9342 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
9346 self.ensure_writable()?;
9347 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
9348 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
9349 }
9350
9351 pub fn clear_ttl(&mut self) -> Result<()> {
9352 self.ensure_writable()?;
9353 self.apply_ttl_policy_at(None, self.current_epoch())
9354 }
9355
9356 pub fn ttl(&self) -> Option<TtlPolicy> {
9357 self.ttl
9358 }
9359
9360 pub(crate) fn prepare_ttl_policy(
9361 &self,
9362 column_name: &str,
9363 duration_nanos: u64,
9364 ) -> Result<TtlPolicy> {
9365 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
9366 return Err(MongrelError::InvalidArgument(
9367 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
9368 ));
9369 }
9370 let column = self
9371 .schema
9372 .columns
9373 .iter()
9374 .find(|column| column.name == column_name)
9375 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
9376 if column.ty != TypeId::TimestampNanos {
9377 return Err(MongrelError::Schema(format!(
9378 "TTL column {column_name} must be TimestampNanos, is {:?}",
9379 column.ty
9380 )));
9381 }
9382 Ok(TtlPolicy {
9383 column_id: column.id,
9384 duration_nanos,
9385 })
9386 }
9387
9388 pub(crate) fn apply_ttl_policy_at(
9389 &mut self,
9390 policy: Option<TtlPolicy>,
9391 epoch: Epoch,
9392 ) -> Result<()> {
9393 if let Some(policy) = policy {
9394 let column = self
9395 .schema
9396 .columns
9397 .iter()
9398 .find(|column| column.id == policy.column_id)
9399 .ok_or_else(|| {
9400 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
9401 })?;
9402 if column.ty != TypeId::TimestampNanos
9403 || policy.duration_nanos == 0
9404 || policy.duration_nanos > i64::MAX as u64
9405 {
9406 return Err(MongrelError::Schema("invalid TTL policy".into()));
9407 }
9408 }
9409 self.ttl = policy;
9410 self.agg_cache = Arc::new(HashMap::new());
9411 self.clear_result_cache();
9412 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
9413 self.persist_manifest(epoch)
9414 }
9415
9416 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
9417 let Some(policy) = self.ttl else {
9418 return false;
9419 };
9420 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
9421 return false;
9422 };
9423 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
9424 }
9425
9426 pub fn current_epoch(&self) -> Epoch {
9427 self.epoch.visible()
9428 }
9429
9430 pub fn memtable_len(&self) -> usize {
9431 self.memtable.len()
9432 }
9433
9434 pub fn count(&self) -> u64 {
9437 if self.ttl.is_none()
9438 && self.pending_put_cols.is_empty()
9439 && self.pending_delete_rids.is_empty()
9440 && self.pending_rows.is_empty()
9441 && self.pending_dels.is_empty()
9442 && self.pending_truncate.is_none()
9443 {
9444 self.live_count
9445 } else {
9446 self.visible_rows(self.snapshot())
9447 .map(|rows| rows.len() as u64)
9448 .unwrap_or(self.live_count)
9449 }
9450 }
9451
9452 pub fn count_conditions(
9456 &mut self,
9457 conditions: &[crate::query::Condition],
9458 snapshot: Snapshot,
9459 ) -> Result<Option<u64>> {
9460 use crate::query::Condition;
9461 if self.ttl.is_some() {
9462 if conditions.is_empty() {
9463 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
9464 }
9465 let mut sets = Vec::with_capacity(conditions.len());
9466 for condition in conditions {
9467 sets.push(self.resolve_condition(condition, snapshot)?);
9468 }
9469 let survivors = RowIdSet::intersect_many(sets);
9470 let rows = self.visible_rows(snapshot)?;
9471 return Ok(Some(
9472 rows.into_iter()
9473 .filter(|row| survivors.contains(row.row_id.0))
9474 .count() as u64,
9475 ));
9476 }
9477 if conditions.is_empty() {
9478 return Ok(Some(self.count()));
9479 }
9480 let served = |c: &Condition| {
9481 matches!(
9482 c,
9483 Condition::Pk(_)
9484 | Condition::BitmapEq { .. }
9485 | Condition::BitmapIn { .. }
9486 | Condition::BytesPrefix { .. }
9487 | Condition::FmContains { .. }
9488 | Condition::FmContainsAll { .. }
9489 | Condition::Ann { .. }
9490 | Condition::Range { .. }
9491 | Condition::RangeF64 { .. }
9492 | Condition::SparseMatch { .. }
9493 | Condition::MinHashSimilar { .. }
9494 | Condition::IsNull { .. }
9495 | Condition::IsNotNull { .. }
9496 )
9497 };
9498 if !conditions.iter().all(served) {
9499 return Ok(None);
9500 }
9501 self.ensure_indexes_complete()?;
9502 if !self.pending_put_cols.is_empty()
9503 || !self.pending_delete_rids.is_empty()
9504 || !self.pending_rows.is_empty()
9505 || !self.pending_dels.is_empty()
9506 || self.pending_truncate.is_some()
9507 {
9508 let mut sets = Vec::with_capacity(conditions.len());
9509 for condition in conditions {
9510 sets.push(self.resolve_condition(condition, snapshot)?);
9511 }
9512 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
9513 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
9514 }
9515 let mut sets = Vec::with_capacity(conditions.len());
9516 for condition in conditions {
9517 sets.push(self.resolve_condition(condition, snapshot)?);
9518 }
9519 let rids = RowIdSet::intersect_many(sets);
9520 if self.had_deletes
9528 || !self.memtable.is_empty()
9529 || !self.mutable_run.is_empty()
9530 {
9531 let sorted = rids.into_sorted_vec();
9532 let count = self.rows_for_rids(&sorted, snapshot)?.len() as u64;
9533 crate::trace::QueryTrace::record(|t| {
9534 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
9535 t.survivor_count = Some(count as usize);
9536 t.conditions_pushed = conditions.len();
9537 });
9538 return Ok(Some(count));
9539 }
9540 let count = rids.len() as u64;
9541 crate::trace::QueryTrace::record(|t| {
9542 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
9543 t.survivor_count = Some(count as usize);
9544 t.conditions_pushed = conditions.len();
9545 });
9546 Ok(Some(count))
9547 }
9548
9549 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
9555 let mut out = Vec::new();
9556 for row in self.memtable.visible_versions(snapshot.epoch) {
9557 if row.deleted {
9558 out.push(row.row_id.0);
9559 }
9560 }
9561 for row in self.mutable_run.visible_versions(snapshot.epoch) {
9562 if row.deleted {
9563 out.push(row.row_id.0);
9564 }
9565 }
9566 out
9567 }
9568
9569 pub fn bulk_load_columns(
9578 &mut self,
9579 user_columns: Vec<(u16, columnar::NativeColumn)>,
9580 ) -> Result<Epoch> {
9581 self.bulk_load_columns_with(user_columns, 3, false, true)
9582 }
9583
9584 pub fn bulk_load_fast(
9591 &mut self,
9592 user_columns: Vec<(u16, columnar::NativeColumn)>,
9593 ) -> Result<Epoch> {
9594 self.bulk_load_columns_with(user_columns, -1, true, false)
9595 }
9596
9597 fn bulk_load_columns_with(
9598 &mut self,
9599 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
9600 zstd_level: i32,
9601 force_plain: bool,
9602 lz4: bool,
9603 ) -> Result<Epoch> {
9604 self.ensure_writable()?;
9605 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
9606 if n == 0 {
9607 return Ok(self.current_epoch());
9608 }
9609 let epoch = self.commit_new_epoch()?;
9610 let live_before = self.live_count;
9611 self.spill_mutable_run(epoch)?;
9613 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
9614 && self.indexes_complete
9615 && self.run_refs.is_empty()
9616 && self.memtable.is_empty()
9617 && self.mutable_run.is_empty();
9618 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
9621 self.validate_columns_not_null(&user_columns, n)?;
9622 let winner_idx = self
9623 .bulk_pk_winner_indices(&user_columns, n)
9624 .filter(|idx| idx.len() != n);
9625 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
9626 match winner_idx.as_deref() {
9627 Some(idx) => {
9628 let compacted = user_columns
9629 .iter()
9630 .map(|(id, c)| (*id, c.gather(idx)))
9631 .collect();
9632 (compacted, idx.len())
9633 }
9634 None => (user_columns, n),
9635 };
9636 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
9637 let first = self.allocator.alloc_range(write_n as u64)?.0;
9638 for rid in first..first + write_n as u64 {
9639 self.reservoir.offer(rid);
9640 }
9641 let run_id = self.alloc_run_id()?;
9642 let path = self.run_path(run_id);
9643 let mut writer =
9644 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
9645 if force_plain {
9646 writer = writer.with_plain();
9647 } else if lz4 {
9648 writer = writer.with_lz4();
9651 } else {
9652 writer = writer.with_zstd_level(zstd_level);
9653 }
9654 if let Some(kek) = &self.kek {
9655 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
9656 }
9657 let header = match self.create_run_file(run_id)? {
9658 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
9659 None => writer.write_native(&path, &write_columns, write_n, first)?,
9660 };
9661 self.run_refs.push(RunRef {
9662 run_id: run_id as u128,
9663 level: 0,
9664 epoch_created: epoch.0,
9665 row_count: header.row_count,
9666 });
9667 self.run_row_id_ranges
9668 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
9669 self.live_count = self.live_count.saturating_add(write_n as u64);
9670 if eager_index_build {
9671 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
9672 self.index_columns_bulk(&write_columns, &row_ids);
9673 self.indexes_complete = true;
9674 self.build_learned_ranges()?;
9675 } else {
9676 self.indexes_complete = false;
9680 }
9681 self.mark_flushed(epoch)?;
9682 self.persist_manifest(epoch)?;
9683 if eager_index_build {
9684 self.checkpoint_indexes(epoch);
9685 }
9686 self.clear_result_cache();
9687 self.data_generation = self.data_generation.wrapping_add(1);
9688 Ok(epoch)
9689 }
9690
9691 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
9709 let n = row_ids.len();
9710 if n == 0 {
9711 return;
9712 }
9713 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
9714 columns.iter().map(|(id, c)| (*id, c)).collect();
9715 let ty_of: std::collections::HashMap<u16, TypeId> = self
9716 .schema
9717 .columns
9718 .iter()
9719 .map(|c| (c.id, c.ty.clone()))
9720 .collect();
9721 let pk_id = self.schema.primary_key().map(|c| c.id);
9722
9723 for (i, &rid) in row_ids.iter().enumerate() {
9724 let row_id = RowId(rid);
9725 if let Some(pid) = pk_id {
9726 if let Some(col) = by_id.get(&pid) {
9727 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
9728 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
9729 self.insert_hot_pk(key, row_id);
9730 }
9731 }
9732 }
9733 for idef in &self.schema.indexes {
9734 let Some(col) = by_id.get(&idef.column_id) else {
9735 continue;
9736 };
9737 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
9738 match idef.kind {
9739 IndexKind::Bitmap => {
9740 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
9741 if let Some(key) =
9742 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
9743 {
9744 b.insert(key, row_id);
9745 }
9746 }
9747 }
9748 IndexKind::FmIndex => {
9749 if let Some(f) = self.fm.get_mut(&idef.column_id) {
9750 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9751 f.insert(bytes.to_vec(), row_id);
9752 }
9753 }
9754 }
9755 IndexKind::Sparse => {
9756 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
9757 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9758 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
9759 s.insert(&terms, row_id);
9760 }
9761 }
9762 }
9763 }
9764 IndexKind::MinHash => {
9765 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
9766 if let Some(bytes) = columnar::native_bytes_at(col, i) {
9767 let tokens = crate::index::token_hashes_from_bytes(bytes);
9768 mh.insert(&tokens, row_id);
9769 }
9770 }
9771 }
9772 _ => {}
9773 }
9774 }
9775 }
9776 }
9777
9778 pub fn visible_columns_native(
9783 &self,
9784 snapshot: Snapshot,
9785 projection: Option<&[u16]>,
9786 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9787 self.visible_columns_native_inner(snapshot, projection, None)
9788 }
9789
9790 pub fn visible_columns_native_with_control(
9791 &self,
9792 snapshot: Snapshot,
9793 projection: Option<&[u16]>,
9794 control: &crate::ExecutionControl,
9795 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9796 self.visible_columns_native_inner(snapshot, projection, Some(control))
9797 }
9798
9799 fn visible_columns_native_inner(
9800 &self,
9801 snapshot: Snapshot,
9802 projection: Option<&[u16]>,
9803 control: Option<&crate::ExecutionControl>,
9804 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
9805 execution_checkpoint(control, 0)?;
9806 let wanted: Vec<u16> = match projection {
9807 Some(p) => p.to_vec(),
9808 None => self.schema.columns.iter().map(|c| c.id).collect(),
9809 };
9810 if self.ttl.is_none()
9811 && self.memtable.is_empty()
9812 && self.mutable_run.is_empty()
9813 && self.run_refs.len() == 1
9814 {
9815 let rr = self.run_refs[0].clone();
9816 let mut reader = self.open_reader(rr.run_id)?;
9817 let idxs = reader.visible_indices_native(snapshot.epoch)?;
9818 execution_checkpoint(control, 0)?;
9819 let all_visible = idxs.len() == reader.row_count();
9820 if reader.has_mmap() && control.is_none() {
9826 use rayon::prelude::*;
9827 let valid: Vec<u16> = wanted
9830 .iter()
9831 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
9832 .copied()
9833 .collect();
9834 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
9836 .par_iter()
9837 .filter_map(|cid| {
9838 reader
9839 .column_native_shared(*cid)
9840 .ok()
9841 .map(|col| (*cid, col))
9842 })
9843 .collect();
9844 let cols = decoded
9845 .into_iter()
9846 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
9847 .collect();
9848 return Ok(cols);
9849 }
9850 let mut cols = Vec::with_capacity(wanted.len());
9851 for (index, cid) in wanted.iter().enumerate() {
9852 execution_checkpoint(control, index)?;
9853 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
9854 Some(c) => c,
9855 None => continue,
9856 };
9857 let col = reader.column_native(cdef.id)?;
9858 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
9859 }
9860 return Ok(cols);
9861 }
9862 let vcols = self.visible_columns(snapshot)?;
9863 execution_checkpoint(control, 0)?;
9864 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
9865 let out: Vec<(u16, columnar::NativeColumn)> = vcols
9866 .into_iter()
9867 .filter(|(id, _)| want_set.contains(id))
9868 .map(|(id, vals)| {
9869 let ty = self
9870 .schema
9871 .columns
9872 .iter()
9873 .find(|c| c.id == id)
9874 .map(|c| c.ty.clone())
9875 .unwrap_or(TypeId::Bytes);
9876 (id, columnar::values_to_native(ty, &vals))
9877 })
9878 .collect();
9879 Ok(out)
9880 }
9881
9882 pub fn run_count(&self) -> usize {
9883 self.run_refs.len()
9884 }
9885
9886 pub fn memtable_is_empty(&self) -> bool {
9888 self.memtable.is_empty()
9889 }
9890
9891 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
9895 self.page_cache.stats()
9896 }
9897
9898 pub fn reset_page_cache_stats(&self) {
9900 self.page_cache.reset_stats();
9901 }
9902
9903 pub fn run_ids(&self) -> Vec<u128> {
9906 self.run_refs.iter().map(|r| r.run_id).collect()
9907 }
9908
9909 pub fn single_run_is_clean(&self) -> bool {
9913 if self.ttl.is_some() || self.run_refs.len() != 1 {
9914 return false;
9915 }
9916 self.open_reader(self.run_refs[0].run_id)
9917 .map(|r| r.is_clean())
9918 .unwrap_or(false)
9919 }
9920
9921 fn resolve_footprint(
9928 &self,
9929 conditions: &[crate::query::Condition],
9930 snapshot: Snapshot,
9931 ) -> roaring::RoaringBitmap {
9932 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
9933 return roaring::RoaringBitmap::new();
9934 }
9935 if self.run_refs.is_empty() {
9936 return roaring::RoaringBitmap::new();
9937 }
9938 if self.run_refs.len() == 1 {
9940 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
9941 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
9942 return rids.to_roaring_lossy();
9943 }
9944 }
9945 }
9946 roaring::RoaringBitmap::new()
9947 }
9948
9949 pub fn query_columns_native_cached(
9960 &mut self,
9961 conditions: &[crate::query::Condition],
9962 projection: Option<&[u16]>,
9963 snapshot: Snapshot,
9964 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9965 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
9966 }
9967
9968 pub fn query_columns_native_cached_with_control(
9969 &mut self,
9970 conditions: &[crate::query::Condition],
9971 projection: Option<&[u16]>,
9972 snapshot: Snapshot,
9973 control: &crate::ExecutionControl,
9974 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9975 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
9976 }
9977
9978 fn query_columns_native_cached_inner(
9979 &mut self,
9980 conditions: &[crate::query::Condition],
9981 projection: Option<&[u16]>,
9982 snapshot: Snapshot,
9983 control: Option<&crate::ExecutionControl>,
9984 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
9985 execution_checkpoint(control, 0)?;
9986 if self.ttl.is_some() {
9989 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9990 }
9991 if conditions.is_empty() {
9992 return self.query_columns_native_inner(conditions, projection, snapshot, control);
9993 }
9994 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
9998 if let Some(hit) = self.result_cache.lock().get_columns(key) {
9999 crate::trace::QueryTrace::record(|t| {
10000 t.result_cache_hit = true;
10001 t.scan_mode = crate::trace::ScanMode::NativePushdown;
10002 });
10003 return Ok(Some((*hit).clone()));
10004 }
10005 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
10006 execution_checkpoint(control, 0)?;
10007 if let Some(cols) = &res {
10008 let footprint = self.resolve_footprint(conditions, snapshot);
10009 let condition_cols = crate::query::condition_columns(conditions);
10010 execution_checkpoint(control, 0)?;
10011 self.result_cache.lock().insert(
10012 key,
10013 CachedEntry {
10014 data: CachedData::Columns(Arc::new(cols.clone())),
10015 footprint,
10016 condition_cols,
10017 },
10018 );
10019 }
10020 Ok(res)
10021 }
10022
10023 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
10028 if self.ttl.is_some() {
10029 return self.query(q);
10030 }
10031 if q.conditions.is_empty() {
10032 return self.query(q);
10033 }
10034 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
10035 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
10036 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
10037 if let Some(hit) = self.result_cache.lock().get_rows(key) {
10038 crate::trace::QueryTrace::record(|t| {
10039 t.result_cache_hit = true;
10040 t.scan_mode = crate::trace::ScanMode::Materialized;
10041 });
10042 return Ok((*hit).clone());
10043 }
10044 let rows = self.query(q)?;
10045 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
10046 let condition_cols = crate::query::condition_columns(&q.conditions);
10047 self.result_cache.lock().insert(
10048 key,
10049 CachedEntry {
10050 data: CachedData::Rows(Arc::new(rows.clone())),
10051 footprint,
10052 condition_cols,
10053 },
10054 );
10055 Ok(rows)
10056 }
10057
10058 #[allow(clippy::type_complexity)]
10073 pub fn query_columns_native_traced(
10074 &mut self,
10075 conditions: &[crate::query::Condition],
10076 projection: Option<&[u16]>,
10077 snapshot: Snapshot,
10078 ) -> Result<(
10079 Option<Vec<(u16, columnar::NativeColumn)>>,
10080 crate::trace::QueryTrace,
10081 )> {
10082 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10083 self.query_columns_native(conditions, projection, snapshot)
10084 });
10085 Ok((result?, trace))
10086 }
10087
10088 #[allow(clippy::type_complexity)]
10091 pub fn query_columns_native_cached_traced(
10092 &mut self,
10093 conditions: &[crate::query::Condition],
10094 projection: Option<&[u16]>,
10095 snapshot: Snapshot,
10096 ) -> Result<(
10097 Option<Vec<(u16, columnar::NativeColumn)>>,
10098 crate::trace::QueryTrace,
10099 )> {
10100 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10101 self.query_columns_native_cached(conditions, projection, snapshot)
10102 });
10103 Ok((result?, trace))
10104 }
10105
10106 pub fn native_page_cursor_traced(
10108 &self,
10109 snapshot: Snapshot,
10110 projection: Vec<(u16, TypeId)>,
10111 conditions: &[crate::query::Condition],
10112 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
10113 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10114 self.native_page_cursor(snapshot, projection, conditions)
10115 });
10116 Ok((result?, trace))
10117 }
10118
10119 pub fn native_multi_run_cursor_traced(
10121 &self,
10122 snapshot: Snapshot,
10123 projection: Vec<(u16, TypeId)>,
10124 conditions: &[crate::query::Condition],
10125 ) -> Result<(
10126 Option<crate::cursor::MultiRunCursor>,
10127 crate::trace::QueryTrace,
10128 )> {
10129 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10130 self.native_multi_run_cursor(snapshot, projection, conditions)
10131 });
10132 Ok((result?, trace))
10133 }
10134
10135 pub fn count_conditions_traced(
10137 &mut self,
10138 conditions: &[crate::query::Condition],
10139 snapshot: Snapshot,
10140 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
10141 let (result, trace) =
10142 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
10143 Ok((result?, trace))
10144 }
10145
10146 pub fn query_traced(
10148 &mut self,
10149 q: &crate::query::Query,
10150 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
10151 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
10152 Ok((result?, trace))
10153 }
10154
10155 pub fn query_columns_native(
10160 &mut self,
10161 conditions: &[crate::query::Condition],
10162 projection: Option<&[u16]>,
10163 snapshot: Snapshot,
10164 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10165 self.query_columns_native_inner(conditions, projection, snapshot, None)
10166 }
10167
10168 pub fn query_columns_native_with_control(
10169 &mut self,
10170 conditions: &[crate::query::Condition],
10171 projection: Option<&[u16]>,
10172 snapshot: Snapshot,
10173 control: &crate::ExecutionControl,
10174 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10175 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
10176 }
10177
10178 fn query_columns_native_inner(
10179 &mut self,
10180 conditions: &[crate::query::Condition],
10181 projection: Option<&[u16]>,
10182 snapshot: Snapshot,
10183 control: Option<&crate::ExecutionControl>,
10184 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10185 use crate::query::Condition;
10186 execution_checkpoint(control, 0)?;
10187 if self.ttl.is_some() {
10190 return Ok(None);
10191 }
10192 if conditions.is_empty() {
10193 return Ok(None);
10194 }
10195 self.ensure_indexes_complete()?;
10196
10197 let served = |c: &Condition| {
10202 matches!(
10203 c,
10204 Condition::Pk(_)
10205 | Condition::BitmapEq { .. }
10206 | Condition::BitmapIn { .. }
10207 | Condition::BytesPrefix { .. }
10208 | Condition::FmContains { .. }
10209 | Condition::FmContainsAll { .. }
10210 | Condition::Ann { .. }
10211 | Condition::Range { .. }
10212 | Condition::RangeF64 { .. }
10213 | Condition::SparseMatch { .. }
10214 | Condition::MinHashSimilar { .. }
10215 | Condition::IsNull { .. }
10216 | Condition::IsNotNull { .. }
10217 )
10218 };
10219 if !conditions.iter().all(served) {
10220 return Ok(None);
10221 }
10222 let fast_path =
10223 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10224 crate::trace::QueryTrace::record(|t| {
10225 t.run_count = self.run_refs.len();
10226 t.memtable_rows = self.memtable.len();
10227 t.mutable_run_rows = self.mutable_run.len();
10228 t.conditions_pushed = conditions.len();
10229 t.learned_range_used = conditions.iter().any(|c| match c {
10230 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
10231 self.learned_range.contains_key(column_id)
10232 }
10233 _ => false,
10234 });
10235 });
10236 let col_ids: Vec<u16> = projection
10238 .map(|p| p.to_vec())
10239 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
10240 let proj_pairs: Vec<(u16, TypeId)> = col_ids
10241 .iter()
10242 .map(|&cid| {
10243 let ty = self
10244 .schema
10245 .columns
10246 .iter()
10247 .find(|c| c.id == cid)
10248 .map(|c| c.ty.clone())
10249 .unwrap_or(TypeId::Bytes);
10250 (cid, ty)
10251 })
10252 .collect();
10253
10254 if fast_path {
10260 let needs_column = conditions.iter().any(|c| match c {
10263 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
10264 Condition::RangeF64 { column_id, .. } => {
10265 !self.learned_range.contains_key(column_id)
10266 }
10267 _ => false,
10268 });
10269 let mut reader_opt: Option<RunReader> = if needs_column {
10270 Some(self.open_reader(self.run_refs[0].run_id)?)
10271 } else {
10272 None
10273 };
10274 let mut sets: Vec<RowIdSet> = Vec::new();
10275 for (index, c) in conditions.iter().enumerate() {
10276 execution_checkpoint(control, index)?;
10277 let s = match c {
10278 Condition::Range { column_id, lo, hi }
10279 if !self.learned_range.contains_key(column_id) =>
10280 {
10281 if reader_opt.is_none() {
10282 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
10283 }
10284 reader_opt
10285 .as_mut()
10286 .expect("reader opened for range")
10287 .range_row_id_set_i64(*column_id, *lo, *hi)?
10288 }
10289 Condition::RangeF64 {
10290 column_id,
10291 lo,
10292 lo_inclusive,
10293 hi,
10294 hi_inclusive,
10295 } if !self.learned_range.contains_key(column_id) => {
10296 if reader_opt.is_none() {
10297 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
10298 }
10299 reader_opt
10300 .as_mut()
10301 .expect("reader opened for range")
10302 .range_row_id_set_f64(
10303 *column_id,
10304 *lo,
10305 *lo_inclusive,
10306 *hi,
10307 *hi_inclusive,
10308 )?
10309 }
10310 _ => self.resolve_condition(c, snapshot)?,
10311 };
10312 sets.push(s);
10313 }
10314 let candidates = RowIdSet::intersect_many(sets);
10315 crate::trace::QueryTrace::record(|t| {
10316 t.survivor_count = Some(candidates.len());
10317 });
10318 if candidates.is_empty() {
10319 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
10320 .iter()
10321 .map(|&id| {
10322 (
10323 id,
10324 columnar::null_native(
10325 proj_pairs
10326 .iter()
10327 .find(|(c, _)| c == &id)
10328 .map(|(_, t)| t.clone())
10329 .unwrap_or(TypeId::Bytes),
10330 0,
10331 ),
10332 )
10333 })
10334 .collect();
10335 return Ok(Some(cols));
10336 }
10337 let mut reader = match reader_opt.take() {
10338 Some(r) => r,
10339 None => self.open_reader(self.run_refs[0].run_id)?,
10340 };
10341 let candidate_ids = candidates.into_sorted_vec();
10342 let (positions, fast_rid) = if let Some(positions) =
10343 reader.positions_for_row_ids_fast(&candidate_ids)
10344 {
10345 (positions, true)
10346 } else {
10347 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
10348 match col {
10349 columnar::NativeColumn::Int64 { data, .. } => {
10350 let mut p = Vec::with_capacity(candidate_ids.len());
10351 for (index, rid) in candidate_ids.iter().enumerate() {
10352 execution_checkpoint(control, index)?;
10353 if let Ok(position) = data.binary_search(&(*rid as i64)) {
10354 p.push(position);
10355 }
10356 }
10357 p.sort_unstable();
10358 (p, false)
10359 }
10360 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
10361 }
10362 };
10363 crate::trace::QueryTrace::record(|t| {
10364 t.scan_mode = crate::trace::ScanMode::NativePushdown;
10365 t.fast_row_id_map = fast_rid;
10366 });
10367 let mut cols = Vec::with_capacity(col_ids.len());
10368 for (index, cid) in col_ids.iter().enumerate() {
10369 execution_checkpoint(control, index)?;
10370 let col = reader.column_native(*cid)?;
10371 cols.push((*cid, col.gather(&positions)));
10372 }
10373 return Ok(Some(cols));
10374 }
10375
10376 if !self.run_refs.is_empty() {
10389 use crate::cursor::{
10390 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
10391 };
10392 let remaining: usize;
10393 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
10394 let c = self
10395 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
10396 .expect("single-run cursor should build when run_refs.len() == 1");
10397 remaining = c.remaining_rows();
10398 Box::new(c)
10399 } else {
10400 let c = self
10401 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
10402 .expect("multi-run cursor should build when run_refs.len() >= 1");
10403 remaining = c.remaining_rows();
10404 Box::new(c)
10405 };
10406 crate::trace::QueryTrace::record(|t| {
10407 if t.survivor_count.is_none() {
10408 t.survivor_count = Some(remaining);
10409 }
10410 });
10411 let cols = match control {
10412 Some(control) => {
10413 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
10414 }
10415 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
10416 };
10417 return Ok(Some(cols));
10418 }
10419
10420 crate::trace::QueryTrace::record(|t| {
10425 t.scan_mode = crate::trace::ScanMode::Materialized;
10426 t.row_materialized = true;
10427 });
10428 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10429 for (index, c) in conditions.iter().enumerate() {
10430 execution_checkpoint(control, index)?;
10431 sets.push(self.resolve_condition(c, snapshot)?);
10432 }
10433 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
10434 let rows = self.rows_for_rids(&rids, snapshot)?;
10435 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
10436 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
10437 execution_checkpoint(control, index)?;
10438 let vals: Vec<Value> = rows
10439 .iter()
10440 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
10441 .collect();
10442 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
10443 }
10444 Ok(Some(cols))
10445 }
10446
10447 pub fn native_page_cursor(
10462 &self,
10463 snapshot: Snapshot,
10464 projection: Vec<(u16, TypeId)>,
10465 conditions: &[crate::query::Condition],
10466 ) -> Result<Option<NativePageCursor>> {
10467 use crate::cursor::build_page_plans;
10468 if self.ttl.is_some() {
10469 return Ok(None);
10470 }
10471 if !conditions.is_empty() && !self.indexes_complete {
10474 return Ok(None);
10475 }
10476 if self.run_refs.len() != 1 {
10477 return Ok(None);
10478 }
10479 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
10480 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
10481
10482 let overlay_rids: HashSet<u64> = {
10485 let mut s = HashSet::new();
10486 for row in self.memtable.visible_versions(snapshot.epoch) {
10487 s.insert(row.row_id.0);
10488 }
10489 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10490 s.insert(row.row_id.0);
10491 }
10492 s
10493 };
10494
10495 let survivors = if conditions.is_empty() {
10499 None
10500 } else {
10501 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
10502 };
10503
10504 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
10511 survivors.clone()
10512 } else if let Some(s) = &survivors {
10513 let mut run_set = s.clone();
10514 run_set.remove_many(overlay_rids.iter().copied());
10515 Some(run_set)
10516 } else {
10517 Some(RowIdSet::from_unsorted(
10518 rids.iter()
10519 .map(|&r| r as u64)
10520 .filter(|r| !overlay_rids.contains(r))
10521 .collect(),
10522 ))
10523 };
10524
10525 let overlay_rows = if overlay_rids.is_empty() {
10526 Vec::new()
10527 } else {
10528 let bound = Self::overlay_materialization_bound(conditions, &survivors);
10529 self.overlay_visible_rows(snapshot, bound)
10530 };
10531
10532 let plans = if positions.is_empty() {
10534 Vec::new()
10535 } else {
10536 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
10537 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
10538 };
10539
10540 let overlay = if overlay_rows.is_empty() {
10542 None
10543 } else {
10544 let filtered =
10545 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
10546 if filtered.is_empty() {
10547 None
10548 } else {
10549 Some(self.materialize_overlay(&filtered, &projection))
10550 }
10551 };
10552
10553 let overlay_row_count = overlay
10554 .as_ref()
10555 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
10556 .unwrap_or(0);
10557 crate::trace::QueryTrace::record(|t| {
10558 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
10559 t.run_count = self.run_refs.len();
10560 t.memtable_rows = self.memtable.len();
10561 t.mutable_run_rows = self.mutable_run.len();
10562 t.overlay_rows = overlay_row_count;
10563 t.conditions_pushed = conditions.len();
10564 t.pages_decoded = plans
10565 .iter()
10566 .map(|p| p.positions.len())
10567 .sum::<usize>()
10568 .min(1);
10569 });
10570
10571 Ok(Some(NativePageCursor::new_with_overlay(
10572 reader, projection, plans, overlay,
10573 )))
10574 }
10575 #[allow(clippy::type_complexity)]
10585 pub fn native_multi_run_cursor(
10586 &self,
10587 snapshot: Snapshot,
10588 projection: Vec<(u16, TypeId)>,
10589 conditions: &[crate::query::Condition],
10590 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
10591 use crate::cursor::{MultiRunCursor, RunStream};
10592 use crate::sorted_run::SYS_ROW_ID;
10593 use std::collections::{BinaryHeap, HashMap, HashSet};
10594 if self.ttl.is_some() {
10595 return Ok(None);
10596 }
10597 if !conditions.is_empty() && !self.indexes_complete {
10600 return Ok(None);
10601 }
10602 if self.run_refs.is_empty() {
10603 return Ok(None);
10604 }
10605
10606 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
10608 Vec::with_capacity(self.run_refs.len());
10609 for rr in &self.run_refs {
10610 let mut reader = self.open_reader(rr.run_id)?;
10611 let (rids, eps, del) = reader.system_columns_native()?;
10612 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
10613 run_meta.push((reader, rids, eps, del, page_rows));
10614 }
10615
10616 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
10620 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
10621 for i in 0..rids.len() {
10622 let rid = rids[i] as u64;
10623 let e = eps[i] as u64;
10624 if e > snapshot.epoch.0 {
10625 continue;
10626 }
10627 let is_del = del[i] != 0;
10628 best.entry(rid)
10629 .and_modify(|cur| {
10630 if e > cur.0 {
10631 *cur = (e, run_idx, i, is_del);
10632 }
10633 })
10634 .or_insert((e, run_idx, i, is_del));
10635 }
10636 }
10637
10638 let overlay_rids: HashSet<u64> = {
10640 let mut s = HashSet::new();
10641 for row in self.memtable.visible_versions(snapshot.epoch) {
10642 s.insert(row.row_id.0);
10643 }
10644 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10645 s.insert(row.row_id.0);
10646 }
10647 s
10648 };
10649
10650 let survivors: Option<RowIdSet> = if conditions.is_empty() {
10652 None
10653 } else {
10654 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10655 for c in conditions {
10656 sets.push(self.resolve_condition(c, snapshot)?);
10657 }
10658 Some(RowIdSet::intersect_many(sets))
10659 };
10660
10661 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
10665 for (rid, (_, run_idx, pos, deleted)) in &best {
10666 if *deleted {
10667 continue;
10668 }
10669 if overlay_rids.contains(rid) {
10670 continue;
10671 }
10672 if let Some(s) = &survivors {
10673 if !s.contains(*rid) {
10674 continue;
10675 }
10676 }
10677 per_run[*run_idx].push((*rid, *pos));
10678 }
10679 for v in per_run.iter_mut() {
10680 v.sort_unstable_by_key(|&(rid, _)| rid);
10681 }
10682
10683 let mut streams = Vec::with_capacity(run_meta.len());
10685 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
10686 let mut total = 0usize;
10687 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
10688 let mut starts = Vec::with_capacity(page_rows.len());
10689 let mut acc = 0usize;
10690 for &r in &page_rows {
10691 starts.push(acc);
10692 acc += r;
10693 }
10694 let mut survivors_vec: Vec<(u64, usize, usize)> =
10695 Vec::with_capacity(per_run[run_idx].len());
10696 for &(rid, pos) in &per_run[run_idx] {
10697 let page_seq = match starts.partition_point(|&s| s <= pos) {
10698 0 => continue,
10699 p => p - 1,
10700 };
10701 let within = pos - starts[page_seq];
10702 survivors_vec.push((rid, page_seq, within));
10703 }
10704 total += survivors_vec.len();
10705 if let Some(&(rid, _, _)) = survivors_vec.first() {
10706 heap.push(std::cmp::Reverse((rid, run_idx)));
10707 }
10708 streams.push(RunStream::new(reader, survivors_vec, page_rows));
10709 }
10710
10711 let overlay_rows = if overlay_rids.is_empty() {
10713 Vec::new()
10714 } else {
10715 let bound = Self::overlay_materialization_bound(conditions, &survivors);
10716 self.overlay_visible_rows(snapshot, bound)
10717 };
10718 let overlay = if overlay_rows.is_empty() {
10719 None
10720 } else {
10721 let filtered =
10722 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
10723 if filtered.is_empty() {
10724 None
10725 } else {
10726 Some(self.materialize_overlay(&filtered, &projection))
10727 }
10728 };
10729
10730 let overlay_row_count = overlay
10731 .as_ref()
10732 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
10733 .unwrap_or(0);
10734 crate::trace::QueryTrace::record(|t| {
10735 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
10736 t.run_count = self.run_refs.len();
10737 t.memtable_rows = self.memtable.len();
10738 t.mutable_run_rows = self.mutable_run.len();
10739 t.overlay_rows = overlay_row_count;
10740 t.conditions_pushed = conditions.len();
10741 t.survivor_count = Some(total);
10742 });
10743
10744 Ok(Some(MultiRunCursor::new(
10745 streams, projection, heap, total, overlay,
10746 )))
10747 }
10748
10749 fn overlay_materialization_bound<'a>(
10761 conditions: &[crate::query::Condition],
10762 survivors: &'a Option<RowIdSet>,
10763 ) -> Option<&'a RowIdSet> {
10764 use crate::query::Condition;
10765 let has_range = conditions
10766 .iter()
10767 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
10768 if has_range {
10769 None
10770 } else {
10771 survivors.as_ref()
10772 }
10773 }
10774
10775 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
10787 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
10788 let mut fold = |row: Row| {
10789 if let Some(b) = bound {
10790 if !b.contains(row.row_id.0) {
10791 return;
10792 }
10793 }
10794 best.entry(row.row_id.0)
10795 .and_modify(|(be, br)| {
10796 if row.committed_epoch > *be {
10797 *be = row.committed_epoch;
10798 *br = row.clone();
10799 }
10800 })
10801 .or_insert_with(|| (row.committed_epoch, row));
10802 };
10803 for row in self.memtable.visible_versions(snapshot.epoch) {
10804 fold(row);
10805 }
10806 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10807 fold(row);
10808 }
10809 let mut out: Vec<Row> = best
10810 .into_values()
10811 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
10812 .collect();
10813 out.sort_by_key(|r| r.row_id);
10814 out
10815 }
10816
10817 fn filter_overlay_rows(
10825 &self,
10826 rows: Vec<Row>,
10827 conditions: &[crate::query::Condition],
10828 survivors: Option<&RowIdSet>,
10829 snapshot: Snapshot,
10830 ) -> Result<Vec<Row>> {
10831 if conditions.is_empty() {
10832 return Ok(rows);
10833 }
10834 use crate::query::Condition;
10835 let all_index_served = !conditions
10839 .iter()
10840 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
10841 if all_index_served {
10842 return Ok(rows
10843 .into_iter()
10844 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
10845 .collect());
10846 }
10847 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10850 for c in conditions {
10851 let s = match c {
10852 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
10853 _ => self.resolve_condition(c, snapshot)?,
10854 };
10855 per_cond_sets.push(s);
10856 }
10857 Ok(rows
10858 .into_iter()
10859 .filter(|row| {
10860 conditions.iter().enumerate().all(|(i, c)| match c {
10861 Condition::Range { column_id, lo, hi } => {
10862 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
10863 }
10864 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
10865 match row.columns.get(column_id) {
10866 Some(Value::Float64(v)) => {
10867 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
10868 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
10869 lo_ok && hi_ok
10870 }
10871 _ => false,
10872 }
10873 }
10874 _ => per_cond_sets[i].contains(row.row_id.0),
10875 })
10876 })
10877 .collect())
10878 }
10879
10880 fn materialize_overlay(
10883 &self,
10884 rows: &[Row],
10885 projection: &[(u16, TypeId)],
10886 ) -> Vec<columnar::NativeColumn> {
10887 if projection.is_empty() {
10888 return vec![columnar::null_native(TypeId::Int64, rows.len())];
10889 }
10890 let mut cols = Vec::with_capacity(projection.len());
10891 for (cid, ty) in projection {
10892 let vals: Vec<Value> = rows
10893 .iter()
10894 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
10895 .collect();
10896 cols.push(columnar::values_to_native(ty.clone(), &vals));
10897 }
10898 cols
10899 }
10900
10901 fn resolve_survivor_rids(
10906 &self,
10907 conditions: &[crate::query::Condition],
10908 reader: &mut RunReader,
10909 snapshot: Snapshot,
10910 ) -> Result<RowIdSet> {
10911 use crate::query::Condition;
10912 let mut sets: Vec<RowIdSet> = Vec::new();
10913 for c in conditions {
10914 self.validate_condition(c)?;
10915 let s: RowIdSet = match c {
10916 Condition::Pk(key) => {
10917 let lookup = self
10918 .schema
10919 .primary_key()
10920 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
10921 .unwrap_or_else(|| key.clone());
10922 self.hot
10923 .get(&lookup)
10924 .map(|r| RowIdSet::one(r.0))
10925 .unwrap_or_else(RowIdSet::empty)
10926 }
10927 Condition::BitmapEq { column_id, value } => {
10928 let lookup = self.index_lookup_key_bytes(*column_id, value);
10929 self.bitmap
10930 .get(column_id)
10931 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
10932 .unwrap_or_else(RowIdSet::empty)
10933 }
10934 Condition::BitmapIn { column_id, values } => {
10935 let bm = self.bitmap.get(column_id);
10936 let mut acc = roaring::RoaringBitmap::new();
10937 if let Some(b) = bm {
10938 for v in values {
10939 let lookup = self.index_lookup_key_bytes(*column_id, v);
10940 acc |= b.get(&lookup);
10941 }
10942 }
10943 RowIdSet::from_roaring(acc)
10944 }
10945 Condition::BytesPrefix { column_id, prefix } => {
10946 if let Some(b) = self.bitmap.get(column_id) {
10947 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
10948 let mut acc = roaring::RoaringBitmap::new();
10949 for key in b.keys() {
10950 if key.starts_with(&lookup_prefix) {
10951 acc |= b.get(&key);
10952 }
10953 }
10954 RowIdSet::from_roaring(acc)
10955 } else {
10956 RowIdSet::empty()
10957 }
10958 }
10959 Condition::FmContains { column_id, pattern } => self
10960 .fm
10961 .get(column_id)
10962 .map(|f| {
10963 RowIdSet::from_unsorted(
10964 f.locate(pattern).into_iter().map(|r| r.0).collect(),
10965 )
10966 })
10967 .unwrap_or_else(RowIdSet::empty),
10968 Condition::FmContainsAll {
10969 column_id,
10970 patterns,
10971 } => {
10972 if let Some(f) = self.fm.get(column_id) {
10973 let sets: Vec<RowIdSet> = patterns
10974 .iter()
10975 .map(|pat| {
10976 RowIdSet::from_unsorted(
10977 f.locate(pat).into_iter().map(|r| r.0).collect(),
10978 )
10979 })
10980 .collect();
10981 RowIdSet::intersect_many(sets)
10982 } else {
10983 RowIdSet::empty()
10984 }
10985 }
10986 Condition::Ann {
10987 column_id,
10988 query,
10989 k,
10990 } => RowIdSet::from_unsorted(
10991 self.retrieve_filtered(
10992 &crate::query::Retriever::Ann {
10993 column_id: *column_id,
10994 query: query.clone(),
10995 k: *k,
10996 },
10997 snapshot,
10998 None,
10999 None,
11000 None,
11001 None,
11002 )?
11003 .into_iter()
11004 .map(|hit| hit.row_id.0)
11005 .collect(),
11006 ),
11007 Condition::SparseMatch {
11008 column_id,
11009 query,
11010 k,
11011 } => RowIdSet::from_unsorted(
11012 self.retrieve_filtered(
11013 &crate::query::Retriever::Sparse {
11014 column_id: *column_id,
11015 query: query.clone(),
11016 k: *k,
11017 },
11018 snapshot,
11019 None,
11020 None,
11021 None,
11022 None,
11023 )?
11024 .into_iter()
11025 .map(|hit| hit.row_id.0)
11026 .collect(),
11027 ),
11028 Condition::MinHashSimilar {
11029 column_id,
11030 query,
11031 k,
11032 } => match self.minhash.get(column_id) {
11033 Some(index) => {
11034 let candidates = index.candidate_row_ids(query);
11035 let eligible =
11036 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
11037 RowIdSet::from_unsorted(
11038 index
11039 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
11040 .into_iter()
11041 .map(|(row_id, _)| row_id.0)
11042 .collect(),
11043 )
11044 }
11045 None => RowIdSet::empty(),
11046 },
11047 Condition::Range { column_id, lo, hi } => {
11048 if let Some(li) = self.learned_range.get(column_id) {
11049 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
11050 } else {
11051 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
11052 }
11053 }
11054 Condition::RangeF64 {
11055 column_id,
11056 lo,
11057 lo_inclusive,
11058 hi,
11059 hi_inclusive,
11060 } => {
11061 if let Some(li) = self.learned_range.get(column_id) {
11062 RowIdSet::from_unsorted(
11063 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
11064 .into_iter()
11065 .collect(),
11066 )
11067 } else {
11068 reader.range_row_id_set_f64(
11069 *column_id,
11070 *lo,
11071 *lo_inclusive,
11072 *hi,
11073 *hi_inclusive,
11074 )?
11075 }
11076 }
11077 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
11078 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
11079 };
11080 sets.push(s);
11081 }
11082 Ok(RowIdSet::intersect_many(sets))
11083 }
11084
11085 pub fn scan_cursor(
11106 &self,
11107 snapshot: Snapshot,
11108 projection: Vec<(u16, TypeId)>,
11109 conditions: &[crate::query::Condition],
11110 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
11111 if self.ttl.is_some() {
11112 return Ok(None);
11113 }
11114 if !conditions.is_empty() && !self.indexes_complete {
11120 return Ok(None);
11121 }
11122 if self.run_refs.len() == 1 {
11123 Ok(self
11124 .native_page_cursor(snapshot, projection, conditions)?
11125 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
11126 } else {
11127 Ok(self
11128 .native_multi_run_cursor(snapshot, projection, conditions)?
11129 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
11130 }
11131 }
11132
11133 pub fn aggregate_native(
11147 &self,
11148 snapshot: Snapshot,
11149 column: Option<u16>,
11150 conditions: &[crate::query::Condition],
11151 agg: NativeAgg,
11152 ) -> Result<Option<NativeAggResult>> {
11153 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
11154 }
11155
11156 pub fn aggregate_native_with_control(
11157 &self,
11158 snapshot: Snapshot,
11159 column: Option<u16>,
11160 conditions: &[crate::query::Condition],
11161 agg: NativeAgg,
11162 control: &crate::ExecutionControl,
11163 ) -> Result<Option<NativeAggResult>> {
11164 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
11165 }
11166
11167 fn aggregate_native_inner(
11168 &self,
11169 snapshot: Snapshot,
11170 column: Option<u16>,
11171 conditions: &[crate::query::Condition],
11172 agg: NativeAgg,
11173 control: Option<&crate::ExecutionControl>,
11174 ) -> Result<Option<NativeAggResult>> {
11175 execution_checkpoint(control, 0)?;
11176 if self.ttl.is_some() {
11177 return Ok(None);
11178 }
11179 if self.run_refs.len() == 1 && conditions.is_empty() {
11181 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
11182 return Ok(Some(res));
11183 }
11184 }
11185 if matches!(agg, NativeAgg::Count) && column.is_none() {
11189 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
11190 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
11191 }
11192 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
11193 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
11194 }
11195 let cid = match column {
11198 Some(c) => c,
11199 None => return Ok(None),
11200 };
11201 let ty = self.column_type(cid);
11202 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
11203 execution_checkpoint(control, 0)?;
11204 return match ty {
11205 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11206 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
11207 Ok(Some(pack_int(agg, count, sum, mn, mx)))
11208 }
11209 TypeId::Float64 => {
11210 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
11211 Ok(Some(pack_float(agg, count, sum, mn, mx)))
11212 }
11213 _ => Ok(None),
11214 };
11215 }
11216 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
11218 execution_checkpoint(control, 0)?;
11219 match ty {
11220 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11221 let mut count = 0u64;
11222 let mut sum = 0i128;
11223 let mut mn = i64::MAX;
11224 let mut mx = i64::MIN;
11225 for row in &rows {
11226 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
11227 count += 1;
11228 sum += i128::from(*v);
11229 mn = mn.min(*v);
11230 mx = mx.max(*v);
11231 }
11232 }
11233 Ok(Some(pack_int(agg, count, sum, mn, mx)))
11234 }
11235 TypeId::Float64 => {
11236 let mut count = 0u64;
11237 let mut sum = 0.0f64;
11238 let mut mn = f64::INFINITY;
11239 let mut mx = f64::NEG_INFINITY;
11240 for row in &rows {
11241 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
11242 count += 1;
11243 sum += *v;
11244 mn = mn.min(*v);
11245 mx = mx.max(*v);
11246 }
11247 }
11248 Ok(Some(pack_float(agg, count, sum, mn, mx)))
11249 }
11250 _ => Ok(None),
11251 }
11252 }
11253
11254 fn visible_rows_filtered(
11256 &self,
11257 snapshot: Snapshot,
11258 conditions: &[crate::query::Condition],
11259 control: Option<&crate::ExecutionControl>,
11260 ) -> Result<Vec<Row>> {
11261 let rows = if let Some(control) = control {
11262 self.visible_rows_controlled(snapshot, control)?
11263 } else {
11264 self.visible_rows(snapshot)?
11265 };
11266 if conditions.is_empty() {
11267 return Ok(rows);
11268 }
11269 Ok(rows
11270 .into_iter()
11271 .filter(|row| {
11272 conditions
11273 .iter()
11274 .all(|cond| condition_matches_row(cond, row, &self.schema))
11275 })
11276 .collect())
11277 }
11278
11279 fn aggregate_from_stats(
11287 &self,
11288 snapshot: Snapshot,
11289 column: Option<u16>,
11290 agg: NativeAgg,
11291 ) -> Result<Option<NativeAggResult>> {
11292 let cid = match (agg, column) {
11293 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
11294 _ => return Ok(None), };
11296 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
11297 return Ok(None);
11298 };
11299 let Some(cs) = stats.get(&cid) else {
11300 return Ok(None);
11301 };
11302 match agg {
11303 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
11305 self.live_count.saturating_sub(cs.null_count),
11306 ))),
11307 NativeAgg::Min | NativeAgg::Max => {
11308 let bound = if agg == NativeAgg::Min {
11309 &cs.min
11310 } else {
11311 &cs.max
11312 };
11313 match bound {
11314 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
11315 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
11316 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
11321 None => Ok(None),
11322 }
11323 }
11324 _ => Ok(None),
11325 }
11326 }
11327
11328 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
11337 if self.ttl.is_some() {
11338 return Ok(None);
11339 }
11340 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
11341 return Ok(None);
11342 }
11343 self.ensure_indexes_complete()?;
11346 let reader = self.open_reader(self.run_refs[0].run_id)?;
11347 if self.live_count != reader.row_count() as u64 {
11348 return Ok(None);
11349 }
11350 let Some(bm) = self.bitmap.get(&column_id) else {
11351 return Ok(None); };
11353 let mut distinct = bm.value_count() as u64;
11354 if !bm.get(&Value::Null.encode_key()).is_empty() {
11357 distinct = distinct.saturating_sub(1);
11358 }
11359 Ok(Some(distinct))
11360 }
11361
11362 pub fn aggregate_incremental(
11374 &mut self,
11375 cache_key: u64,
11376 conditions: &[crate::query::Condition],
11377 column: Option<u16>,
11378 agg: NativeAgg,
11379 ) -> Result<IncrementalAggResult> {
11380 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
11381 }
11382
11383 pub fn aggregate_incremental_with_control(
11384 &mut self,
11385 cache_key: u64,
11386 conditions: &[crate::query::Condition],
11387 column: Option<u16>,
11388 agg: NativeAgg,
11389 control: &crate::ExecutionControl,
11390 ) -> Result<IncrementalAggResult> {
11391 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
11392 }
11393
11394 fn aggregate_incremental_inner(
11395 &mut self,
11396 cache_key: u64,
11397 conditions: &[crate::query::Condition],
11398 column: Option<u16>,
11399 agg: NativeAgg,
11400 control: Option<&crate::ExecutionControl>,
11401 ) -> Result<IncrementalAggResult> {
11402 execution_checkpoint(control, 0)?;
11403 let snap = self.snapshot();
11404 let cur_wm = self.allocator.current().0;
11405 let cur_epoch = snap.epoch.0;
11406 let incremental_ok = self.ttl.is_none()
11413 && !self.had_deletes
11414 && self.memtable.is_empty()
11415 && self.mutable_run.is_empty();
11416
11417 if incremental_ok {
11420 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
11421 if cached.epoch == cur_epoch {
11422 return Ok(IncrementalAggResult {
11423 state: cached.state,
11424 incremental: true,
11425 delta_rows: 0,
11426 });
11427 }
11428 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
11429 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
11430 let mut delta_rids = Vec::with_capacity(delta_len);
11431 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
11432 execution_checkpoint(control, index)?;
11433 delta_rids.push(row_id);
11434 }
11435 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
11436 execution_checkpoint(control, 0)?;
11437 let index_sets = self.resolve_index_conditions(conditions, snap)?;
11438 let delta_state = agg_state_from_rows(
11439 &delta_rows,
11440 conditions,
11441 &index_sets,
11442 column,
11443 agg,
11444 &self.schema,
11445 control,
11446 )?;
11447 let merged = cached.state.merge(delta_state);
11448 let delta_n = delta_rids.len() as u64;
11449 Arc::make_mut(&mut self.agg_cache).insert(
11450 cache_key,
11451 CachedAgg {
11452 state: merged.clone(),
11453 watermark: cur_wm,
11454 epoch: cur_epoch,
11455 },
11456 );
11457 return Ok(IncrementalAggResult {
11458 state: merged,
11459 incremental: true,
11460 delta_rows: delta_n,
11461 });
11462 }
11463 }
11464 }
11465
11466 let cursor_ok =
11471 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
11472 let state = if cursor_ok && agg != NativeAgg::Avg {
11473 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
11474 Some(result) => {
11475 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
11476 }
11477 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
11478 }
11479 } else {
11480 self.agg_state_full_scan(conditions, column, agg, snap, control)?
11481 };
11482 if incremental_ok {
11484 Arc::make_mut(&mut self.agg_cache).insert(
11485 cache_key,
11486 CachedAgg {
11487 state: state.clone(),
11488 watermark: cur_wm,
11489 epoch: cur_epoch,
11490 },
11491 );
11492 }
11493 Ok(IncrementalAggResult {
11494 state,
11495 incremental: false,
11496 delta_rows: 0,
11497 })
11498 }
11499
11500 fn agg_state_full_scan(
11503 &self,
11504 conditions: &[crate::query::Condition],
11505 column: Option<u16>,
11506 agg: NativeAgg,
11507 snap: Snapshot,
11508 control: Option<&crate::ExecutionControl>,
11509 ) -> Result<AggState> {
11510 execution_checkpoint(control, 0)?;
11511 let rows = self.visible_rows(snap)?;
11512 execution_checkpoint(control, 0)?;
11513 let index_sets = self.resolve_index_conditions(conditions, snap)?;
11514 agg_state_from_rows(
11515 &rows,
11516 conditions,
11517 &index_sets,
11518 column,
11519 agg,
11520 &self.schema,
11521 control,
11522 )
11523 }
11524
11525 fn resolve_index_conditions(
11528 &self,
11529 conditions: &[crate::query::Condition],
11530 snapshot: Snapshot,
11531 ) -> Result<Vec<RowIdSet>> {
11532 use crate::query::Condition;
11533 let mut sets = Vec::new();
11534 for c in conditions {
11535 if matches!(
11536 c,
11537 Condition::Ann { .. }
11538 | Condition::SparseMatch { .. }
11539 | Condition::MinHashSimilar { .. }
11540 ) {
11541 sets.push(self.resolve_condition(c, snapshot)?);
11542 }
11543 }
11544 Ok(sets)
11545 }
11546
11547 fn column_type(&self, cid: u16) -> TypeId {
11548 self.schema
11549 .columns
11550 .iter()
11551 .find(|c| c.id == cid)
11552 .map(|c| c.ty.clone())
11553 .unwrap_or(TypeId::Bytes)
11554 }
11555
11556 pub fn approx_aggregate(
11565 &mut self,
11566 conditions: &[crate::query::Condition],
11567 column: Option<u16>,
11568 agg: ApproxAgg,
11569 z: f64,
11570 ) -> Result<Option<ApproxResult>> {
11571 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
11572 }
11573
11574 pub fn approx_aggregate_with_candidate_authorization(
11577 &mut self,
11578 conditions: &[crate::query::Condition],
11579 column: Option<u16>,
11580 agg: ApproxAgg,
11581 z: f64,
11582 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
11583 ) -> Result<Option<ApproxResult>> {
11584 use crate::query::Condition;
11585 self.ensure_reservoir_complete()?;
11586 let mut snapshot = self.snapshot();
11591 let n_pop = self.count();
11592 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
11593 if sample_rids.is_empty() {
11594 return Ok(None);
11595 }
11596 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
11598 if live_sample.is_empty() {
11599 snapshot = Snapshot::unbounded();
11600 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
11601 }
11602 let s = live_sample.len();
11603 if s == 0 {
11604 return Ok(None);
11605 }
11606 let authorized = authorization
11607 .map(|authorization| {
11608 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
11609 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
11610 })
11611 .transpose()?;
11612
11613 let mut index_sets: Vec<RowIdSet> = Vec::new();
11616 for c in conditions {
11617 if matches!(
11618 c,
11619 Condition::Ann { .. }
11620 | Condition::SparseMatch { .. }
11621 | Condition::MinHashSimilar { .. }
11622 ) {
11623 index_sets.push(self.resolve_condition(c, snapshot)?);
11624 }
11625 }
11626
11627 let cid = match (agg, column) {
11629 (ApproxAgg::Count, _) => None,
11630 (_, Some(c)) => Some(c),
11631 _ => return Ok(None),
11632 };
11633 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
11634 for r in &live_sample {
11635 if authorized
11636 .as_ref()
11637 .is_some_and(|authorized| !authorized.contains(&r.row_id))
11638 {
11639 continue;
11640 }
11641 if !conditions
11643 .iter()
11644 .all(|c| condition_matches_row(c, r, &self.schema))
11645 {
11646 continue;
11647 }
11648 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
11650 continue;
11651 }
11652 if let Some(cid) = cid {
11653 let mut cells = r
11654 .columns
11655 .get(&cid)
11656 .cloned()
11657 .map(|value| vec![(cid, value)])
11658 .unwrap_or_default();
11659 if let Some(authorization) = authorization {
11660 authorization.security.apply_masks_to_cells(
11661 authorization.table,
11662 &mut cells,
11663 authorization.principal,
11664 );
11665 }
11666 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
11667 passing_vals.push(v);
11668 } } else {
11670 passing_vals.push(0.0); }
11672 }
11673 let m = passing_vals.len();
11674
11675 let (point, half) = match agg {
11676 ApproxAgg::Count => {
11677 let p = m as f64 / s as f64;
11679 let point = n_pop as f64 * p;
11680 let var = if s > 1 {
11681 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
11682 * (1.0 - s as f64 / n_pop as f64).max(0.0)
11683 } else {
11684 0.0
11685 };
11686 (point, z * var.sqrt())
11687 }
11688 ApproxAgg::Sum => {
11689 let y: Vec<f64> = live_sample
11691 .iter()
11692 .map(|r| {
11693 let passes_row = authorized
11694 .as_ref()
11695 .is_none_or(|authorized| authorized.contains(&r.row_id))
11696 && conditions
11697 .iter()
11698 .all(|c| condition_matches_row(c, r, &self.schema))
11699 && index_sets.iter().all(|set| set.contains(r.row_id.0));
11700 if passes_row {
11701 cid.and_then(|cid| {
11702 let mut cells = r
11703 .columns
11704 .get(&cid)
11705 .cloned()
11706 .map(|value| vec![(cid, value)])
11707 .unwrap_or_default();
11708 if let Some(authorization) = authorization {
11709 authorization.security.apply_masks_to_cells(
11710 authorization.table,
11711 &mut cells,
11712 authorization.principal,
11713 );
11714 }
11715 as_f64(cells.first().map(|(_, value)| value))
11716 })
11717 .unwrap_or(0.0)
11718 } else {
11719 0.0
11720 }
11721 })
11722 .collect();
11723 let mean_y = y.iter().sum::<f64>() / s as f64;
11724 let point = n_pop as f64 * mean_y;
11725 let var = if s > 1 {
11726 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
11727 let var_y = ss / (s - 1) as f64;
11728 n_pop as f64 * n_pop as f64 * var_y / s as f64
11729 * (1.0 - s as f64 / n_pop as f64).max(0.0)
11730 } else {
11731 0.0
11732 };
11733 (point, z * var.sqrt())
11734 }
11735 ApproxAgg::Avg => {
11736 if m == 0 {
11737 return Ok(Some(ApproxResult {
11738 point: 0.0,
11739 ci_low: 0.0,
11740 ci_high: 0.0,
11741 n_population: n_pop,
11742 n_sample_live: s,
11743 n_passing: 0,
11744 }));
11745 }
11746 let mean = passing_vals.iter().sum::<f64>() / m as f64;
11747 let half = if m > 1 {
11748 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
11749 let sd = (ss / (m - 1) as f64).sqrt();
11750 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
11751 z * sd / (m as f64).sqrt() * fpc.sqrt()
11752 } else {
11753 0.0
11754 };
11755 (mean, half)
11756 }
11757 };
11758
11759 Ok(Some(ApproxResult {
11760 point,
11761 ci_low: point - half,
11762 ci_high: point + half,
11763 n_population: n_pop,
11764 n_sample_live: s,
11765 n_passing: m,
11766 }))
11767 }
11768
11769 pub fn exact_column_stats(
11777 &self,
11778 _snapshot: Snapshot,
11779 projection: &[u16],
11780 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
11781 if self.ttl.is_some()
11782 || !(self.memtable.is_empty()
11783 && self.mutable_run.is_empty()
11784 && self.run_refs.len() == 1)
11785 {
11786 return Ok(None);
11787 }
11788 let reader = self.open_reader(self.run_refs[0].run_id)?;
11789 if self.live_count != reader.row_count() as u64 {
11790 return Ok(None);
11791 }
11792 let mut out = HashMap::new();
11793 for &cid in projection {
11794 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
11795 Some(c) => c,
11796 None => continue,
11797 };
11798 let Some(stats) = reader.column_page_stats(cid) else {
11800 out.insert(
11801 cid,
11802 ColumnStat {
11803 min: None,
11804 max: None,
11805 null_count: self.live_count,
11806 },
11807 );
11808 continue;
11809 };
11810 let stat = match cdef.ty {
11811 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11812 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
11813 min: mn.map(Value::Int64),
11814 max: mx.map(Value::Int64),
11815 null_count: n,
11816 })
11817 }
11818 TypeId::Float64 => {
11819 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
11820 min: mn.map(Value::Float64),
11821 max: mx.map(Value::Float64),
11822 null_count: n,
11823 })
11824 }
11825 _ => None,
11826 };
11827 if let Some(s) = stat {
11828 out.insert(cid, s);
11829 }
11830 }
11831 Ok(Some(out))
11832 }
11833
11834 pub fn dir(&self) -> &Path {
11835 &self.dir
11836 }
11837
11838 pub fn schema(&self) -> &Schema {
11839 &self.schema
11840 }
11841
11842 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
11843 self.ann.get(&column_id)
11844 }
11845
11846 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
11847 self.ann.get_mut(&column_id)
11848 }
11849
11850 pub(crate) fn set_catalog_name(&mut self, name: String) {
11851 self.name = name;
11852 }
11853
11854 pub(crate) fn prepare_alter_column(
11855 &mut self,
11856 column_name: &str,
11857 change: &AlterColumn,
11858 ) -> Result<(ColumnDef, Option<Schema>)> {
11859 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
11860 return Err(MongrelError::InvalidArgument(
11861 "ALTER COLUMN requires committing staged writes first".into(),
11862 ));
11863 }
11864 let old = self
11865 .schema
11866 .columns
11867 .iter()
11868 .find(|c| c.name == column_name)
11869 .cloned()
11870 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
11871 let mut next = old.clone();
11872
11873 if let Some(name) = &change.name {
11874 let trimmed = name.trim();
11875 if trimmed.is_empty() {
11876 return Err(MongrelError::InvalidArgument(
11877 "ALTER COLUMN name must not be empty".into(),
11878 ));
11879 }
11880 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
11881 return Err(MongrelError::Schema(format!(
11882 "column {trimmed} already exists"
11883 )));
11884 }
11885 next.name = trimmed.to_string();
11886 }
11887
11888 if let Some(ty) = &change.ty {
11889 next.ty = ty.clone();
11890 }
11891 if let Some(flags) = change.flags {
11892 validate_alter_column_flags(old.flags, flags)?;
11893 next.flags = flags;
11894 }
11895
11896 if let Some(default_change) = &change.default_value {
11897 next.default_value = default_change.clone();
11898 }
11899 if let Some(source_change) = &change.embedding_source {
11900 next.embedding_source = source_change.clone();
11901 }
11902
11903 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
11904 if old.flags.contains(ColumnFlags::NULLABLE)
11905 && !next.flags.contains(ColumnFlags::NULLABLE)
11906 && self.column_has_nulls(old.id)?
11907 {
11908 return Err(MongrelError::InvalidArgument(format!(
11909 "column '{}' contains NULL values",
11910 old.name
11911 )));
11912 }
11913 if next == old {
11914 return Ok((next, None));
11915 }
11916 let mut schema = self.schema.clone();
11917 let index = schema
11918 .columns
11919 .iter()
11920 .position(|column| column.id == next.id)
11921 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
11922 schema.columns[index] = next.clone();
11923 schema.schema_id = schema
11924 .schema_id
11925 .checked_add(1)
11926 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
11927 schema.validate_auto_increment()?;
11928 schema.validate_defaults()?;
11929 Ok((next, Some(schema)))
11930 }
11931
11932 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
11933 self.schema = schema;
11934 self.auto_inc = resolve_auto_inc(&self.schema);
11935 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
11936 self.clear_result_cache();
11937 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
11938 }
11939
11940 pub(crate) fn publish_index_schema_change(
11944 &mut self,
11945 schema: Schema,
11946 artifact: SecondaryIndexArtifact,
11947 ) -> Result<()> {
11948 self.apply_altered_schema_prepared(schema);
11949 self.prune_index_maps_to_schema();
11950 match artifact {
11951 SecondaryIndexArtifact::Bitmap(column_id, index) => {
11952 self.bitmap.insert(column_id, index);
11953 }
11954 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
11955 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
11956 }
11957 SecondaryIndexArtifact::Fm(column_id, index) => {
11958 self.fm.insert(column_id, *index);
11959 }
11960 SecondaryIndexArtifact::Ann(column_id, index) => {
11961 self.ann.insert(column_id, *index);
11962 }
11963 SecondaryIndexArtifact::Sparse(column_id, index) => {
11964 self.sparse.insert(column_id, index);
11965 }
11966 SecondaryIndexArtifact::MinHash(column_id, index) => {
11967 self.minhash.insert(column_id, index);
11968 }
11969 }
11970 self.indexes_complete = true;
11971 self.seal_generations();
11972 let view = Arc::new(self.capture_read_generation());
11973 self.published.store(Arc::clone(&view));
11974 checkpoint_current_schema(self)?;
11975 self.invalidate_index_checkpoint();
11978 Ok(())
11979 }
11980
11981 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
11987 self.apply_altered_schema_prepared(schema);
11988 self.prune_index_maps_to_schema();
11989 self.indexes_complete = true;
11990 self.seal_generations();
11991 let view = Arc::new(self.capture_read_generation());
11992 self.published.store(Arc::clone(&view));
11993 checkpoint_current_schema(self)?;
11994 self.invalidate_index_checkpoint();
11996 Ok(())
11997 }
11998
11999 fn prune_index_maps_to_schema(&mut self) {
12000 let keep_bitmap: std::collections::HashSet<u16> = self
12001 .schema
12002 .indexes
12003 .iter()
12004 .filter(|index| index.kind == IndexKind::Bitmap)
12005 .map(|index| index.column_id)
12006 .collect();
12007 let keep_ann: std::collections::HashSet<u16> = self
12008 .schema
12009 .indexes
12010 .iter()
12011 .filter(|index| index.kind == IndexKind::Ann)
12012 .map(|index| index.column_id)
12013 .collect();
12014 let keep_fm: std::collections::HashSet<u16> = self
12015 .schema
12016 .indexes
12017 .iter()
12018 .filter(|index| index.kind == IndexKind::FmIndex)
12019 .map(|index| index.column_id)
12020 .collect();
12021 let keep_sparse: std::collections::HashSet<u16> = self
12022 .schema
12023 .indexes
12024 .iter()
12025 .filter(|index| index.kind == IndexKind::Sparse)
12026 .map(|index| index.column_id)
12027 .collect();
12028 let keep_minhash: std::collections::HashSet<u16> = self
12029 .schema
12030 .indexes
12031 .iter()
12032 .filter(|index| index.kind == IndexKind::MinHash)
12033 .map(|index| index.column_id)
12034 .collect();
12035 let keep_learned: std::collections::HashSet<u16> = self
12036 .schema
12037 .indexes
12038 .iter()
12039 .filter(|index| index.kind == IndexKind::LearnedRange)
12040 .map(|index| index.column_id)
12041 .collect();
12042 self.bitmap
12043 .retain(|column_id, _| keep_bitmap.contains(column_id));
12044 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
12045 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
12046 self.sparse
12047 .retain(|column_id, _| keep_sparse.contains(column_id));
12048 self.minhash
12049 .retain(|column_id, _| keep_minhash.contains(column_id));
12050 {
12051 let learned = Arc::make_mut(&mut self.learned_range);
12052 learned.retain(|column_id, _| keep_learned.contains(column_id));
12053 }
12054 }
12055
12056 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
12057 checkpoint_current_schema(self)
12058 }
12059
12060 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
12061 self.ensure_writable()?;
12062 let previous_schema = self.schema.clone();
12063 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
12064 if let Some(schema) = schema {
12065 self.apply_altered_schema_prepared(schema);
12066 self.checkpoint_standalone_schema_change(previous_schema)?;
12067 }
12068 Ok(column)
12069 }
12070
12071 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
12072 if self.live_count == 0 {
12073 return Ok(false);
12074 }
12075 let snap = self.snapshot();
12076 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
12077 Ok(columns
12078 .first()
12079 .map(|(_, col)| col.null_count(col.len()) != 0)
12080 .unwrap_or(true))
12081 }
12082
12083 fn has_stored_versions(&self) -> bool {
12084 !self.memtable.is_empty()
12085 || !self.mutable_run.is_empty()
12086 || self.run_refs.iter().any(|r| r.row_count > 0)
12087 || !self.retiring.is_empty()
12088 }
12089
12090 pub fn add_column(
12095 &mut self,
12096 name: &str,
12097 ty: TypeId,
12098 flags: ColumnFlags,
12099 default_value: Option<crate::schema::DefaultExpr>,
12100 ) -> Result<u16> {
12101 self.add_column_with_id(name, ty, flags, default_value, None)
12102 }
12103
12104 pub fn add_column_with_id(
12105 &mut self,
12106 name: &str,
12107 ty: TypeId,
12108 flags: ColumnFlags,
12109 default_value: Option<crate::schema::DefaultExpr>,
12110 requested_id: Option<u16>,
12111 ) -> Result<u16> {
12112 self.ensure_writable()?;
12113 let previous_schema = self.schema.clone();
12114 let (column, schema) =
12115 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
12116 self.apply_altered_schema_prepared(schema);
12117 self.checkpoint_standalone_schema_change(previous_schema)?;
12118 Ok(column.id)
12119 }
12120
12121 pub(crate) fn prepare_add_column(
12122 &mut self,
12123 name: &str,
12124 ty: TypeId,
12125 flags: ColumnFlags,
12126 default_value: Option<crate::schema::DefaultExpr>,
12127 requested_id: Option<u16>,
12128 ) -> Result<(ColumnDef, Schema)> {
12129 self.ensure_writable()?;
12130 if self.schema.columns.iter().any(|c| c.name == name) {
12131 return Err(MongrelError::Schema(format!(
12132 "column {name} already exists"
12133 )));
12134 }
12135 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
12136 if self.schema.columns.iter().any(|c| c.id == id) {
12137 return Err(MongrelError::Schema(format!(
12138 "column id {id} already exists"
12139 )));
12140 }
12141 id
12142 } else {
12143 self.schema
12144 .columns
12145 .iter()
12146 .map(|c| c.id)
12147 .max()
12148 .unwrap_or(0)
12149 .checked_add(1)
12150 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
12151 };
12152 let column = ColumnDef {
12153 id,
12154 name: name.to_string(),
12155 ty,
12156 flags,
12157 default_value,
12158 embedding_source: None,
12159 };
12160 let mut next_schema = self.schema.clone();
12161 next_schema.columns.push(column.clone());
12162 next_schema.schema_id = next_schema
12163 .schema_id
12164 .checked_add(1)
12165 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
12166 next_schema.validate_auto_increment()?;
12167 next_schema.validate_defaults()?;
12168 Ok((column, next_schema))
12169 }
12170
12171 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
12180 self.ensure_writable()?;
12181 let cid = self
12182 .schema
12183 .columns
12184 .iter()
12185 .find(|c| c.name == column_name)
12186 .map(|c| c.id)
12187 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
12188 let ty = self
12189 .schema
12190 .columns
12191 .iter()
12192 .find(|c| c.id == cid)
12193 .map(|c| c.ty.clone())
12194 .unwrap_or(TypeId::Int64);
12195 if !matches!(
12196 ty,
12197 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
12198 ) {
12199 return Err(MongrelError::Schema(format!(
12200 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
12201 )));
12202 }
12203 if self
12204 .schema
12205 .indexes
12206 .iter()
12207 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
12208 {
12209 return Ok(()); }
12211 let previous_schema = self.schema.clone();
12212 let previous_learned_range = Arc::clone(&self.learned_range);
12213 let mut next_schema = previous_schema.clone();
12214 next_schema.indexes.push(IndexDef {
12215 name: format!("{}_learned_range", column_name),
12216 column_id: cid,
12217 kind: IndexKind::LearnedRange,
12218 predicate: None,
12219 options: Default::default(),
12220 });
12221 next_schema.schema_id = next_schema
12222 .schema_id
12223 .checked_add(1)
12224 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
12225 self.apply_altered_schema_prepared(next_schema);
12226 if let Err(error) = self.build_learned_ranges() {
12227 self.apply_altered_schema_prepared(previous_schema);
12228 self.learned_range = previous_learned_range;
12229 return Err(error);
12230 }
12231 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
12232 if !matches!(
12233 &error,
12234 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
12235 ) {
12236 self.learned_range = previous_learned_range;
12237 }
12238 return Err(error);
12239 }
12240 Ok(())
12241 }
12242
12243 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
12244 let mut schema_published = false;
12245 let schema_result = match self._root_guard.as_deref() {
12246 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
12247 schema_published = true;
12248 }),
12249 None => write_schema_with_after(&self.dir, &self.schema, || {
12250 schema_published = true;
12251 }),
12252 };
12253 if schema_result.is_err() && !schema_published {
12254 self.apply_altered_schema_prepared(previous_schema);
12255 return schema_result;
12256 }
12257
12258 let manifest_result = self.persist_manifest(self.current_epoch());
12259 match (schema_result, manifest_result) {
12260 (_, Ok(())) => Ok(()),
12261 (Ok(()), Err(error)) => {
12262 self.poison_after_maintenance_publish_failure();
12263 Err(MongrelError::DurableCommit {
12264 epoch: self.current_epoch().0,
12265 message: format!(
12266 "schema is durable but matching manifest publication failed: {error}"
12267 ),
12268 })
12269 }
12270 (Err(schema_error), Err(manifest_error)) => {
12271 self.poison_after_maintenance_publish_failure();
12272 Err(MongrelError::CommitOutcomeUnknown {
12273 epoch: self.current_epoch().0,
12274 message: format!(
12275 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12276 ),
12277 })
12278 }
12279 }
12280 }
12281
12282 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
12285 self.sync_byte_threshold = threshold;
12286 if let WalSink::Private(w) = &mut self.wal {
12287 w.set_sync_byte_threshold(threshold);
12288 }
12289 }
12290
12291 pub fn page_cache_flush(&self) {
12295 self.page_cache.flush_to_disk();
12296 }
12297
12298 pub fn page_cache_len(&self) -> usize {
12300 self.page_cache.len()
12301 }
12302
12303 pub fn decoded_cache_len(&self) -> usize {
12306 self.decoded_cache.len()
12307 }
12308
12309 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
12312 self.memtable.drain_sorted()
12313 }
12314
12315 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
12316 self.runs_dir().join(format!("r-{run_id}.sr"))
12317 }
12318
12319 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
12320 match self.runs_root.as_deref() {
12321 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
12322 None => Ok(None),
12323 }
12324 }
12325
12326 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
12327 match self.runs_root.as_deref() {
12328 Some(root) => Ok(Some(root.create_regular_new(name)?)),
12329 None => Ok(None),
12330 }
12331 }
12332
12333 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
12334 match self.runs_root.as_deref() {
12335 Some(root) => match root.remove_file(name) {
12336 Ok(()) => Ok(()),
12337 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
12338 Err(error) => Err(error.into()),
12339 },
12340 None => match std::fs::remove_file(self.runs_dir().join(name)) {
12341 Ok(()) => Ok(()),
12342 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
12343 Err(error) => Err(error.into()),
12344 },
12345 }
12346 }
12347
12348 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
12349 match self.runs_root.as_deref() {
12350 Some(root) => root
12351 .rename_file_new(source, destination)
12352 .map_err(Into::into),
12353 None => crate::durable_file::rename(
12354 &self.runs_dir().join(source),
12355 &self.runs_dir().join(destination),
12356 )
12357 .map_err(Into::into),
12358 }
12359 }
12360
12361 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
12362 self.run_refs.iter().map(|run| run.run_id)
12363 }
12364
12365 pub(crate) fn table_dir(&self) -> &Path {
12366 &self.dir
12367 }
12368
12369 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
12370 &self.schema
12371 }
12372
12373 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
12374 let id = self.next_run_id;
12375 self.next_run_id = self
12376 .next_run_id
12377 .checked_add(1)
12378 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
12379 Ok(id)
12380 }
12381
12382 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
12383 self.run_refs.push(run_ref);
12384 }
12385
12386 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
12396 self.retiring.push(crate::manifest::RetiredRun {
12397 run_id,
12398 retire_epoch,
12399 });
12400 }
12401
12402 pub(crate) fn reap_retiring(
12406 &mut self,
12407 min_active: Epoch,
12408 backup_pinned: &std::collections::HashSet<u128>,
12409 ) -> Result<usize> {
12410 if self.retiring.is_empty() {
12411 return Ok(0);
12412 }
12413 let mut reaped = 0;
12414 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
12415 for r in std::mem::take(&mut self.retiring) {
12421 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
12422 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
12423 reaped += 1;
12424 } else {
12425 kept.push(r);
12426 }
12427 }
12428 self.retiring = kept;
12429 if reaped > 0 {
12430 self.persist_manifest(self.current_epoch())?;
12431 }
12432 Ok(reaped)
12433 }
12434
12435 pub(crate) fn has_reapable_retiring(
12436 &self,
12437 min_active: Epoch,
12438 backup_pinned: &std::collections::HashSet<u128>,
12439 ) -> bool {
12440 self.retiring
12441 .iter()
12442 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
12443 }
12444
12445 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
12446 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
12447 return false;
12448 }
12449 self.live_count = self.live_count.saturating_add(run_ref.row_count);
12450 self.run_refs.push(run_ref);
12451 self.indexes_complete = false;
12452 true
12453 }
12454
12455 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
12456 self.kek.as_ref()
12457 }
12458
12459 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
12460 let mut reader = match self.runs_root.as_deref() {
12461 Some(root) => RunReader::open_file_with_cache(
12462 root.open_regular(format!("r-{run_id}.sr"))?,
12463 self.schema.clone(),
12464 self.kek.clone(),
12465 Some(self.page_cache.clone()),
12466 Some(self.decoded_cache.clone()),
12467 self.table_id,
12468 Some(&self.verified_runs),
12469 None,
12470 )?,
12471 None => RunReader::open_with_cache(
12472 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
12473 self.schema.clone(),
12474 self.kek.clone(),
12475 Some(self.page_cache.clone()),
12476 Some(self.decoded_cache.clone()),
12477 self.table_id,
12478 Some(&self.verified_runs),
12479 )?,
12480 };
12481 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
12485 reader.set_uniform_epoch(Epoch(rr.epoch_created));
12486 }
12487 Ok(reader)
12488 }
12489
12490 pub(crate) fn run_refs(&self) -> &[RunRef] {
12491 &self.run_refs
12492 }
12493
12494 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
12495 self.retiring.iter().map(|run| run.run_id)
12496 }
12497
12498 pub(crate) fn runs_dir(&self) -> PathBuf {
12499 self.runs_root
12500 .as_deref()
12501 .and_then(|root| root.io_path().ok())
12502 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
12503 }
12504
12505 pub(crate) fn wal_dir(&self) -> PathBuf {
12506 self.dir.join(WAL_DIR)
12507 }
12508
12509 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
12510 self.run_refs = refs;
12511 self.refresh_run_row_id_ranges();
12512 }
12513
12514 pub(crate) fn refresh_run_row_id_ranges(&mut self) {
12517 let mut ranges = HashMap::with_capacity(self.run_refs.len());
12518 for rr in &self.run_refs {
12519 if let Ok(reader) = self.open_reader(rr.run_id) {
12520 let h = reader.header();
12521 ranges.insert(rr.run_id, (h.min_row_id, h.max_row_id));
12522 }
12523 }
12524 self.run_row_id_ranges = ranges;
12525 }
12526
12527 pub(crate) fn compaction_zstd_level(&self) -> i32 {
12528 self.compaction_zstd_level
12529 }
12530
12531 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
12532 self.kek.clone()
12533 }
12534
12535 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
12539 self.kek.as_ref().map(|k| k.derive_idx_key())
12540 }
12541
12542 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
12546 self.kek.as_ref().map(|k| *k.derive_meta_key())
12547 }
12548
12549 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
12552 self.column_keys
12553 .iter()
12554 .map(|(&id, &(_, scheme))| (id, scheme))
12555 .collect()
12556 }
12557
12558 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
12563 self.tokenize_value_enc(column_id, v)
12564 }
12565
12566 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
12567 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
12568 let (key, scheme) = self.column_keys.get(&column_id)?;
12569 let token: Vec<u8> = match (*scheme, v) {
12570 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
12571 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
12572 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
12573 _ => hmac_token(key, &v.encode_key()).to_vec(),
12574 };
12575 Some(Value::Bytes(token))
12576 }
12577
12578 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
12580 self.index_lookup_key_bytes(column_id, &v.encode_key())
12581 }
12582
12583 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
12586 {
12587 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
12588 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
12589 if *scheme == SCHEME_HMAC_EQ {
12590 return hmac_token(key, encoded).to_vec();
12591 }
12592 }
12593 }
12594 let _ = column_id;
12595 encoded.to_vec()
12596 }
12597}
12598
12599fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
12600 let columnar::NativeColumn::Int64 { data, validity } = col else {
12601 return false;
12602 };
12603 if data.len() < n || !columnar::all_non_null(validity, n) {
12604 return false;
12605 }
12606 data.iter()
12607 .take(n)
12608 .zip(data.iter().skip(1))
12609 .all(|(a, b)| a < b)
12610}
12611
12612#[derive(Debug, Clone)]
12616pub struct ColumnStat {
12617 pub min: Option<Value>,
12618 pub max: Option<Value>,
12619 pub null_count: u64,
12620}
12621
12622#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12624pub enum NativeAgg {
12625 Count,
12626 Sum,
12627 Min,
12628 Max,
12629 Avg,
12630}
12631
12632#[derive(Debug, Clone, PartialEq)]
12634pub enum NativeAggResult {
12635 Count(u64),
12636 Int(i64),
12637 Float(f64),
12638 Null,
12640}
12641
12642#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12644pub enum ApproxAgg {
12645 Count,
12646 Sum,
12647 Avg,
12648}
12649
12650#[derive(Debug, Clone)]
12654pub struct ApproxResult {
12655 pub point: f64,
12657 pub ci_low: f64,
12659 pub ci_high: f64,
12661 pub n_population: u64,
12663 pub n_sample_live: usize,
12665 pub n_passing: usize,
12667}
12668
12669#[derive(Debug, Clone, PartialEq)]
12674pub enum AggState {
12675 Count(u64),
12677 SumI {
12679 sum: i128,
12680 count: u64,
12681 },
12682 SumF {
12684 sum: f64,
12685 count: u64,
12686 },
12687 AvgI {
12689 sum: i128,
12690 count: u64,
12691 },
12692 AvgF {
12694 sum: f64,
12695 count: u64,
12696 },
12697 MinI(i64),
12699 MaxI(i64),
12700 MinF(f64),
12702 MaxF(f64),
12703 Empty,
12705}
12706
12707impl AggState {
12708 pub fn merge(self, other: AggState) -> AggState {
12710 use AggState::*;
12711 match (self, other) {
12712 (Empty, x) | (x, Empty) => x,
12713 (Count(a), Count(b)) => Count(a + b),
12714 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
12715 sum: sa + sb,
12716 count: ca + cb,
12717 },
12718 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
12719 sum: sa + sb,
12720 count: ca + cb,
12721 },
12722 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
12723 sum: sa + sb,
12724 count: ca + cb,
12725 },
12726 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
12727 sum: sa + sb,
12728 count: ca + cb,
12729 },
12730 (MinI(a), MinI(b)) => MinI(a.min(b)),
12731 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
12732 (MinF(a), MinF(b)) => MinF(a.min(b)),
12733 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
12734 _ => Empty, }
12736 }
12737
12738 pub fn point(&self) -> Option<f64> {
12740 match self {
12741 AggState::Count(n) => Some(*n as f64),
12742 AggState::SumI { sum, .. } => Some(*sum as f64),
12743 AggState::SumF { sum, .. } => Some(*sum),
12744 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
12745 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
12746 AggState::MinI(n) => Some(*n as f64),
12747 AggState::MaxI(n) => Some(*n as f64),
12748 AggState::MinF(n) => Some(*n),
12749 AggState::MaxF(n) => Some(*n),
12750 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
12751 }
12752 }
12753
12754 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
12758 let is_float = matches!(ty, Some(TypeId::Float64));
12759 match (agg, result) {
12760 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
12761 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
12762 sum: x as i128,
12763 count: 1, },
12765 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
12766 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
12767 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
12768 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
12769 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
12770 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
12771 (NativeAgg::Count, _) => AggState::Empty,
12772 (_, NativeAggResult::Null) => AggState::Empty,
12773 _ => {
12774 let _ = is_float;
12775 AggState::Empty
12776 }
12777 }
12778 }
12779}
12780
12781#[derive(Debug, Clone)]
12784pub struct CachedAgg {
12785 pub state: AggState,
12786 pub watermark: u64,
12787 pub epoch: u64,
12788}
12789
12790#[derive(Debug, Clone)]
12792pub struct IncrementalAggResult {
12793 pub state: AggState,
12795 pub incremental: bool,
12798 pub delta_rows: u64,
12800}
12801
12802fn agg_state_from_rows(
12806 rows: &[Row],
12807 conditions: &[crate::query::Condition],
12808 index_sets: &[RowIdSet],
12809 column: Option<u16>,
12810 agg: NativeAgg,
12811 schema: &Schema,
12812 control: Option<&crate::ExecutionControl>,
12813) -> Result<AggState> {
12814 let mut count: u64 = 0;
12815 let mut sum_i: i128 = 0;
12816 let mut sum_f: f64 = 0.0;
12817 let mut mn_i: i64 = i64::MAX;
12818 let mut mx_i: i64 = i64::MIN;
12819 let mut mn_f: f64 = f64::INFINITY;
12820 let mut mx_f: f64 = f64::NEG_INFINITY;
12821 let mut saw_int = false;
12822 let mut saw_float = false;
12823 for (index, r) in rows.iter().enumerate() {
12824 execution_checkpoint(control, index)?;
12825 if !conditions
12826 .iter()
12827 .all(|c| condition_matches_row(c, r, schema))
12828 {
12829 continue;
12830 }
12831 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
12832 continue;
12833 }
12834 match agg {
12835 NativeAgg::Count => match column {
12836 None => count += 1,
12838 Some(cid) => match r.columns.get(&cid) {
12841 None | Some(Value::Null) => {}
12842 Some(_) => count += 1,
12843 },
12844 },
12845 _ => match column.and_then(|cid| r.columns.get(&cid)) {
12846 Some(Value::Int64(n)) => {
12847 count += 1;
12848 sum_i += *n as i128;
12849 mn_i = mn_i.min(*n);
12850 mx_i = mx_i.max(*n);
12851 saw_int = true;
12852 }
12853 Some(Value::Float64(f)) => {
12854 count += 1;
12855 sum_f += f;
12856 mn_f = mn_f.min(*f);
12857 mx_f = mx_f.max(*f);
12858 saw_float = true;
12859 }
12860 _ => {}
12861 },
12862 }
12863 }
12864 Ok(match agg {
12865 NativeAgg::Count => {
12866 if count == 0 {
12867 AggState::Empty
12868 } else {
12869 AggState::Count(count)
12870 }
12871 }
12872 NativeAgg::Sum => {
12873 if count == 0 {
12874 AggState::Empty
12875 } else if saw_int {
12876 AggState::SumI { sum: sum_i, count }
12877 } else {
12878 AggState::SumF { sum: sum_f, count }
12879 }
12880 }
12881 NativeAgg::Avg => {
12882 if count == 0 {
12883 AggState::Empty
12884 } else if saw_int {
12885 AggState::AvgI { sum: sum_i, count }
12886 } else {
12887 AggState::AvgF { sum: sum_f, count }
12888 }
12889 }
12890 NativeAgg::Min => {
12891 if !saw_int && !saw_float {
12892 AggState::Empty
12893 } else if saw_int {
12894 AggState::MinI(mn_i)
12895 } else {
12896 AggState::MinF(mn_f)
12897 }
12898 }
12899 NativeAgg::Max => {
12900 if !saw_int && !saw_float {
12901 AggState::Empty
12902 } else if saw_int {
12903 AggState::MaxI(mx_i)
12904 } else {
12905 AggState::MaxF(mx_f)
12906 }
12907 }
12908 })
12909}
12910
12911fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
12915 use crate::query::Condition;
12916 match c {
12917 Condition::Pk(key) => match schema.primary_key() {
12918 Some(pk) => row
12919 .columns
12920 .get(&pk.id)
12921 .map(|v| v.encode_key() == *key)
12922 .unwrap_or(false),
12923 None => false,
12924 },
12925 Condition::BitmapEq { column_id, value } => row
12926 .columns
12927 .get(column_id)
12928 .map(|v| v.encode_key() == *value)
12929 .unwrap_or(false),
12930 Condition::BitmapIn { column_id, values } => {
12931 let key = row.columns.get(column_id).map(|v| v.encode_key());
12932 match key {
12933 Some(k) => values.contains(&k),
12934 None => false,
12935 }
12936 }
12937 Condition::BytesPrefix { column_id, prefix } => row
12938 .columns
12939 .get(column_id)
12940 .map(|v| v.encode_key().starts_with(prefix))
12941 .unwrap_or(false),
12942 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
12943 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
12944 _ => false,
12945 },
12946 Condition::RangeF64 {
12947 column_id,
12948 lo,
12949 lo_inclusive,
12950 hi,
12951 hi_inclusive,
12952 } => match row.columns.get(column_id) {
12953 Some(Value::Float64(n)) => {
12954 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
12955 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
12956 lo_ok && hi_ok
12957 }
12958 _ => false,
12959 },
12960 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
12961 Some(Value::Bytes(b)) => {
12962 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
12963 }
12964 _ => false,
12965 },
12966 Condition::FmContainsAll {
12967 column_id,
12968 patterns,
12969 } => match row.columns.get(column_id) {
12970 Some(Value::Bytes(b)) => patterns
12971 .iter()
12972 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
12973 _ => false,
12974 },
12975 Condition::Ann { .. }
12976 | Condition::SparseMatch { .. }
12977 | Condition::MinHashSimilar { .. } => true,
12978 Condition::IsNull { column_id } => {
12979 matches!(row.columns.get(column_id), Some(Value::Null) | None)
12980 }
12981 Condition::IsNotNull { column_id } => {
12982 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
12983 }
12984 }
12985}
12986
12987fn as_f64(v: Option<&Value>) -> Option<f64> {
12989 match v {
12990 Some(Value::Int64(n)) => Some(*n as f64),
12991 Some(Value::Float64(f)) => Some(*f),
12992 _ => None,
12993 }
12994}
12995
12996fn accumulate_int(
13000 cursor: &mut dyn crate::cursor::Cursor,
13001 control: Option<&crate::ExecutionControl>,
13002) -> Result<(u64, i128, i64, i64)> {
13003 let mut count: u64 = 0;
13004 let mut sum: i128 = 0;
13005 let mut mn: i64 = i64::MAX;
13006 let mut mx: i64 = i64::MIN;
13007 while let Some(cols) = cursor.next_batch()? {
13008 execution_checkpoint(control, 0)?;
13009 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
13010 if crate::columnar::all_non_null(validity, data.len()) {
13011 count += data.len() as u64;
13013 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
13014 execution_checkpoint(control, chunk_index * 1024)?;
13015 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
13016 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
13017 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
13018 }
13019 } else {
13020 for (i, &v) in data.iter().enumerate() {
13021 execution_checkpoint(control, i)?;
13022 if crate::columnar::validity_bit(validity, i) {
13023 count += 1;
13024 sum += v as i128;
13025 mn = mn.min(v);
13026 mx = mx.max(v);
13027 }
13028 }
13029 }
13030 }
13031 }
13032 Ok((count, sum, mn, mx))
13033}
13034
13035fn accumulate_float(
13037 cursor: &mut dyn crate::cursor::Cursor,
13038 control: Option<&crate::ExecutionControl>,
13039) -> Result<(u64, f64, f64, f64)> {
13040 let mut count: u64 = 0;
13041 let mut sum: f64 = 0.0;
13042 let mut mn: f64 = f64::INFINITY;
13043 let mut mx: f64 = f64::NEG_INFINITY;
13044 while let Some(cols) = cursor.next_batch()? {
13045 execution_checkpoint(control, 0)?;
13046 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
13047 if crate::columnar::all_non_null(validity, data.len()) {
13048 count += data.len() as u64;
13049 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
13050 execution_checkpoint(control, chunk_index * 1024)?;
13051 sum += chunk.iter().sum::<f64>();
13052 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
13053 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
13054 }
13055 } else {
13056 for (i, &v) in data.iter().enumerate() {
13057 execution_checkpoint(control, i)?;
13058 if crate::columnar::validity_bit(validity, i) {
13059 count += 1;
13060 sum += v;
13061 mn = mn.min(v);
13062 mx = mx.max(v);
13063 }
13064 }
13065 }
13066 }
13067 }
13068 Ok((count, sum, mn, mx))
13069}
13070
13071#[inline]
13072fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
13073 if index.is_multiple_of(256) {
13074 control
13075 .map(crate::ExecutionControl::checkpoint)
13076 .transpose()?;
13077 }
13078 Ok(())
13079}
13080
13081fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
13082 if count == 0 && !matches!(agg, NativeAgg::Count) {
13083 return NativeAggResult::Null;
13084 }
13085 match agg {
13086 NativeAgg::Count => NativeAggResult::Count(count),
13087 NativeAgg::Sum => match sum.try_into() {
13090 Ok(v) => NativeAggResult::Int(v),
13091 Err(_) => NativeAggResult::Null,
13092 },
13093 NativeAgg::Min => NativeAggResult::Int(mn),
13094 NativeAgg::Max => NativeAggResult::Int(mx),
13095 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
13096 }
13097}
13098
13099fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
13100 if count == 0 && !matches!(agg, NativeAgg::Count) {
13101 return NativeAggResult::Null;
13102 }
13103 match agg {
13104 NativeAgg::Count => NativeAggResult::Count(count),
13105 NativeAgg::Sum => NativeAggResult::Float(sum),
13106 NativeAgg::Min => NativeAggResult::Float(mn),
13107 NativeAgg::Max => NativeAggResult::Float(mx),
13108 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
13109 }
13110}
13111
13112fn agg_int(
13115 stats: &[crate::page::PageStat],
13116 decode: fn(Option<&[u8]>) -> Option<i64>,
13117) -> Option<(Option<i64>, Option<i64>, u64)> {
13118 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
13119 let mut any = false;
13120 for s in stats {
13121 if let Some(v) = decode(s.min.as_deref()) {
13122 mn = mn.min(v);
13123 any = true;
13124 }
13125 if let Some(v) = decode(s.max.as_deref()) {
13126 mx = mx.max(v);
13127 any = true;
13128 }
13129 nulls += s.null_count;
13130 }
13131 any.then_some((Some(mn), Some(mx), nulls))
13132}
13133
13134fn agg_float(
13136 stats: &[crate::page::PageStat],
13137 decode: fn(Option<&[u8]>) -> Option<f64>,
13138) -> Option<(Option<f64>, Option<f64>, u64)> {
13139 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
13140 let mut any = false;
13141 for s in stats {
13142 if let Some(v) = decode(s.min.as_deref()) {
13143 mn = mn.min(v);
13144 any = true;
13145 }
13146 if let Some(v) = decode(s.max.as_deref()) {
13147 mx = mx.max(v);
13148 any = true;
13149 }
13150 nulls += s.null_count;
13151 }
13152 any.then_some((Some(mn), Some(mx), nulls))
13153}
13154
13155type SecondaryIndexes = (
13157 HashMap<u16, BitmapIndex>,
13158 HashMap<u16, AnnIndex>,
13159 HashMap<u16, FmIndex>,
13160 HashMap<u16, SparseIndex>,
13161 HashMap<u16, MinHashIndex>,
13162);
13163
13164fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
13165 let mut bitmap = HashMap::new();
13166 let mut ann = HashMap::new();
13167 let mut fm = HashMap::new();
13168 let mut sparse = HashMap::new();
13169 let mut minhash = HashMap::new();
13170 for idef in &schema.indexes {
13171 match idef.kind {
13172 IndexKind::Bitmap => {
13173 bitmap.insert(idef.column_id, BitmapIndex::new());
13174 }
13175 IndexKind::Ann => {
13176 let dim = schema
13177 .columns
13178 .iter()
13179 .find(|c| c.id == idef.column_id)
13180 .and_then(|c| match c.ty {
13181 TypeId::Embedding { dim } => Some(dim as usize),
13182 _ => None,
13183 })
13184 .unwrap_or(0);
13185 let options = idef.options.ann.clone().unwrap_or_default();
13186 ann.insert(
13187 idef.column_id,
13188 AnnIndex::with_full_options(
13189 dim,
13190 options.m,
13191 options.ef_construction,
13192 options.ef_search,
13193 &options,
13194 ),
13195 );
13196 }
13197 IndexKind::FmIndex => {
13198 fm.insert(idef.column_id, FmIndex::new());
13199 }
13200 IndexKind::Sparse => {
13201 sparse.insert(idef.column_id, SparseIndex::new());
13202 }
13203 IndexKind::MinHash => {
13204 let options = idef.options.minhash.clone().unwrap_or_default();
13205 minhash.insert(
13206 idef.column_id,
13207 MinHashIndex::with_options(options.permutations, options.bands),
13208 );
13209 }
13210 _ => {}
13211 }
13212 }
13213 (bitmap, ann, fm, sparse, minhash)
13214}
13215
13216const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
13217 | ColumnFlags::AUTO_INCREMENT
13218 | ColumnFlags::ENCRYPTED
13219 | ColumnFlags::ENCRYPTED_INDEXABLE
13220 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
13221
13222fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
13223 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
13224 return Err(MongrelError::Schema(
13225 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
13226 ));
13227 }
13228 Ok(())
13229}
13230
13231fn validate_alter_column_type(
13232 schema: &Schema,
13233 old: &ColumnDef,
13234 next: &ColumnDef,
13235 has_stored_versions: bool,
13236) -> Result<()> {
13237 if old.ty == next.ty {
13238 return Ok(());
13239 }
13240 if schema.indexes.iter().any(|i| i.column_id == old.id) {
13241 return Err(MongrelError::Schema(format!(
13242 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
13243 old.name
13244 )));
13245 }
13246 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
13247 return Ok(());
13248 }
13249 Err(MongrelError::Schema(format!(
13250 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
13251 old.ty, next.ty
13252 )))
13253}
13254
13255fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
13256 matches!(
13257 (old, new),
13258 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
13259 )
13260}
13261
13262fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
13268 let mut prev: Option<i64> = None;
13269 for r in rows {
13270 match r.columns.get(&pk_id) {
13271 Some(Value::Int64(v)) => {
13272 if prev.is_some_and(|p| p >= *v) {
13273 return false;
13274 }
13275 prev = Some(*v);
13276 }
13277 _ => return false,
13278 }
13279 }
13280 true
13281}
13282
13283#[allow(clippy::too_many_arguments)]
13284fn index_into(
13285 schema: &Schema,
13286 row: &Row,
13287 hot: &mut HotIndex,
13288 bitmap: &mut HashMap<u16, BitmapIndex>,
13289 ann: &mut HashMap<u16, AnnIndex>,
13290 fm: &mut HashMap<u16, FmIndex>,
13291 sparse: &mut HashMap<u16, SparseIndex>,
13292 minhash: &mut HashMap<u16, MinHashIndex>,
13293) {
13294 for idef in &schema.indexes {
13295 let Some(val) = row.columns.get(&idef.column_id) else {
13296 continue;
13297 };
13298 match idef.kind {
13299 IndexKind::Bitmap => {
13300 if let Some(b) = bitmap.get_mut(&idef.column_id) {
13301 b.insert(val.encode_key(), row.row_id);
13302 }
13303 }
13304 IndexKind::Ann => {
13305 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
13306 if let Some(meta) = val.generated_embedding_metadata() {
13307 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
13309 continue;
13310 }
13311 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
13312 .is_err()
13313 {
13314 continue;
13315 }
13316 }
13317 a.insert_validated(v, row.row_id);
13318 }
13319 }
13320 IndexKind::FmIndex => {
13321 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
13322 f.insert(b.clone(), row.row_id);
13323 }
13324 }
13325 IndexKind::Sparse => {
13326 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
13327 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
13330 s.insert(&terms, row.row_id);
13331 }
13332 }
13333 }
13334 IndexKind::MinHash => {
13335 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
13336 let tokens = crate::index::token_hashes_from_bytes(b);
13339 mh.insert(&tokens, row.row_id);
13340 }
13341 }
13342 _ => {}
13343 }
13344 }
13345 if let Some(pk_col) = schema.primary_key() {
13346 if let Some(pk_val) = row.columns.get(&pk_col.id) {
13347 hot.insert(pk_val.encode_key(), row.row_id);
13348 }
13349 }
13350}
13351
13352#[allow(clippy::too_many_arguments)]
13355fn index_into_single(
13356 idef: &IndexDef,
13357 _schema: &Schema,
13358 row: &Row,
13359 _hot: &mut HotIndex,
13360 bitmap: &mut HashMap<u16, BitmapIndex>,
13361 ann: &mut HashMap<u16, AnnIndex>,
13362 fm: &mut HashMap<u16, FmIndex>,
13363 sparse: &mut HashMap<u16, SparseIndex>,
13364 minhash: &mut HashMap<u16, MinHashIndex>,
13365) {
13366 let Some(val) = row.columns.get(&idef.column_id) else {
13367 return;
13368 };
13369 match idef.kind {
13370 IndexKind::Bitmap => {
13371 if let Some(b) = bitmap.get_mut(&idef.column_id) {
13372 b.insert(val.encode_key(), row.row_id);
13373 }
13374 }
13375 IndexKind::Ann => {
13376 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
13377 if let Some(meta) = val.generated_embedding_metadata() {
13378 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
13380 return;
13381 }
13382 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
13383 .is_err()
13384 {
13385 return;
13386 }
13387 }
13388 a.insert_validated(v, row.row_id);
13389 }
13390 }
13391 IndexKind::FmIndex => {
13392 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
13393 f.insert(b.clone(), row.row_id);
13394 }
13395 }
13396 IndexKind::Sparse => {
13397 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
13398 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
13399 s.insert(&terms, row.row_id);
13400 }
13401 }
13402 }
13403 IndexKind::MinHash => {
13404 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
13405 let tokens = crate::index::token_hashes_from_bytes(b);
13406 mh.insert(&tokens, row.row_id);
13407 }
13408 }
13409 _ => {}
13410 }
13411}
13412
13413fn eval_partial_predicate(
13419 pred: &str,
13420 columns_map: &HashMap<u16, &Value>,
13421 name_to_id: &HashMap<&str, u16>,
13422) -> bool {
13423 let lower = pred.trim().to_ascii_lowercase();
13424 if let Some(rest) = lower.strip_suffix(" is not null") {
13426 let col_name = rest.trim();
13427 if let Some(col_id) = name_to_id.get(col_name) {
13428 return columns_map
13429 .get(col_id)
13430 .is_some_and(|v| !matches!(v, Value::Null));
13431 }
13432 }
13433 if let Some(rest) = lower.strip_suffix(" is null") {
13435 let col_name = rest.trim();
13436 if let Some(col_id) = name_to_id.get(col_name) {
13437 return columns_map
13438 .get(col_id)
13439 .is_none_or(|v| matches!(v, Value::Null));
13440 }
13441 }
13442 true
13445}
13446
13447#[allow(dead_code)]
13453fn bulk_index_key(
13454 column_keys: &HashMap<u16, ([u8; 32], u8)>,
13455 column_id: u16,
13456 ty: TypeId,
13457 col: &columnar::NativeColumn,
13458 i: usize,
13459) -> Option<Vec<u8>> {
13460 let encoded = columnar::encode_key_native(ty, col, i)?;
13461 {
13462 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
13463 if let Some((key, scheme)) = column_keys.get(&column_id) {
13464 return Some(match (*scheme, col) {
13465 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
13466 (_, columnar::NativeColumn::Int64 { data, .. }) => {
13467 ope_token_i64(key, data[i]).to_vec()
13468 }
13469 (_, columnar::NativeColumn::Float64 { data, .. }) => {
13470 ope_token_f64(key, data[i]).to_vec()
13471 }
13472 _ => hmac_token(key, &encoded).to_vec(),
13473 });
13474 }
13475 }
13476 Some(encoded)
13477}
13478
13479pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
13480 write_schema_with_after(dir, schema, || {})
13481}
13482
13483pub(crate) fn write_schema_durable(
13484 root: &crate::durable_file::DurableRoot,
13485 schema: &Schema,
13486) -> Result<()> {
13487 write_schema_durable_with_after(root, schema, || {})
13488}
13489
13490fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
13491where
13492 F: FnOnce(),
13493{
13494 let json = serde_json::to_string_pretty(schema)
13495 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
13496 crate::durable_file::write_atomic_with_after(
13497 &dir.join(SCHEMA_FILENAME),
13498 json.as_bytes(),
13499 after_publish,
13500 )?;
13501 Ok(())
13502}
13503
13504fn write_schema_durable_with_after<F>(
13505 root: &crate::durable_file::DurableRoot,
13506 schema: &Schema,
13507 after_publish: F,
13508) -> Result<()>
13509where
13510 F: FnOnce(),
13511{
13512 let json = serde_json::to_string_pretty(schema)
13513 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
13514 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
13515 Ok(())
13516}
13517
13518fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
13519 let mut schema_published = false;
13520 let schema_result = match table._root_guard.as_deref() {
13521 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
13522 schema_published = true;
13523 }),
13524 None => write_schema_with_after(&table.dir, &table.schema, || {
13525 schema_published = true;
13526 }),
13527 };
13528 if schema_result.is_err() && !schema_published {
13529 return schema_result;
13530 }
13531 match table.persist_manifest(table.current_epoch()) {
13532 Ok(()) => Ok(()),
13533 Err(manifest_error) => Err(match schema_result {
13534 Ok(()) => manifest_error,
13535 Err(schema_error) => MongrelError::Other(format!(
13536 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
13537 )),
13538 }),
13539 }
13540}
13541
13542fn read_schema(dir: &Path) -> Result<Schema> {
13543 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
13544 read_schema_file(file)
13545}
13546
13547fn read_schema_file(file: std::fs::File) -> Result<Schema> {
13548 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
13549 use std::io::Read;
13550
13551 let length = file.metadata()?.len();
13552 if length > MAX_SCHEMA_BYTES {
13553 return Err(MongrelError::ResourceLimitExceeded {
13554 resource: "schema bytes",
13555 requested: usize::try_from(length).unwrap_or(usize::MAX),
13556 limit: MAX_SCHEMA_BYTES as usize,
13557 });
13558 }
13559 let mut bytes = Vec::with_capacity(length as usize);
13560 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
13561 if bytes.len() as u64 != length {
13562 return Err(MongrelError::Schema(
13563 "schema length changed while reading".into(),
13564 ));
13565 }
13566 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
13567}
13568
13569fn preflight_standalone_open(
13570 dir: &Path,
13571 runs_root: Option<&crate::durable_file::DurableRoot>,
13572 idx_root: Option<&crate::durable_file::DurableRoot>,
13573 manifest: &Manifest,
13574 schema: &Schema,
13575 records: &[crate::wal::Record],
13576 kek: Option<Arc<Kek>>,
13577) -> Result<()> {
13578 crate::wal::validate_shared_transaction_framing(records)?;
13579 if manifest.schema_id > schema.schema_id
13580 || manifest.flushed_epoch > manifest.current_epoch
13581 || manifest.global_idx_epoch > manifest.current_epoch
13582 || manifest.next_row_id == u64::MAX
13583 || manifest.auto_inc_next < 0
13584 || manifest.auto_inc_next == i64::MAX
13585 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
13586 {
13587 return Err(MongrelError::InvalidArgument(
13588 "manifest counters or schema identity are invalid".into(),
13589 ));
13590 }
13591 let mut run_ids = HashSet::new();
13592 let mut maximum_row_id = None::<u64>;
13593 for run in &manifest.runs {
13594 if run.run_id >= u64::MAX as u128
13595 || !run_ids.insert(run.run_id)
13596 || run.epoch_created > manifest.current_epoch
13597 {
13598 return Err(MongrelError::InvalidArgument(
13599 "manifest contains an invalid or duplicate active run".into(),
13600 ));
13601 }
13602 let mut reader = match runs_root {
13603 Some(root) => RunReader::open_file(
13604 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
13605 schema.clone(),
13606 kek.clone(),
13607 )?,
13608 None => RunReader::open(
13609 dir.join(RUNS_DIR)
13610 .join(format!("r-{}.sr", run.run_id as u64)),
13611 schema.clone(),
13612 kek.clone(),
13613 )?,
13614 };
13615 let header = reader.header();
13616 if header.run_id != run.run_id
13617 || header.level != run.level
13618 || header.row_count != run.row_count
13619 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
13620 || header.is_uniform_epoch() && header.epoch_created != 0
13621 || header.schema_id > schema.schema_id
13622 {
13623 return Err(MongrelError::InvalidArgument(format!(
13624 "run {} differs from its manifest",
13625 run.run_id
13626 )));
13627 }
13628 if header.row_count != 0 {
13629 maximum_row_id = Some(
13630 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
13631 );
13632 }
13633 reader.validate_all_pages()?;
13634 }
13635 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
13636 return Err(MongrelError::InvalidArgument(
13637 "manifest next_row_id does not advance beyond persisted rows".into(),
13638 ));
13639 }
13640 for run in &manifest.retiring {
13641 if run.run_id >= u64::MAX as u128
13642 || run.retire_epoch > manifest.current_epoch
13643 || !run_ids.insert(run.run_id)
13644 {
13645 return Err(MongrelError::InvalidArgument(
13646 "manifest contains an invalid or duplicate retired run".into(),
13647 ));
13648 }
13649 }
13650 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
13651 match idx_root {
13652 Some(root) => {
13653 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
13654 }
13655 None => {
13656 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
13657 }
13658 }
13659
13660 let committed = records
13661 .iter()
13662 .filter_map(|record| match record.op {
13663 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
13664 _ => None,
13665 })
13666 .collect::<HashMap<_, _>>();
13667 for record in records {
13668 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
13669 continue;
13670 };
13671 match &record.op {
13672 Op::Put { table_id, rows } => {
13673 if *table_id != manifest.table_id {
13674 return Err(MongrelError::CorruptWal {
13675 offset: record.seq.0,
13676 reason: format!(
13677 "private WAL record references table {table_id}, expected {}",
13678 manifest.table_id
13679 ),
13680 });
13681 }
13682 let rows: Vec<Row> =
13683 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
13684 offset: record.seq.0,
13685 reason: format!("committed Put payload could not be decoded: {error}"),
13686 })?;
13687 for row in rows {
13688 if row.deleted || row.row_id.0 == u64::MAX {
13689 return Err(MongrelError::CorruptWal {
13690 offset: record.seq.0,
13691 reason: "committed Put contains an invalid row identity".into(),
13692 });
13693 }
13694 let cells = row.columns.into_iter().collect::<Vec<_>>();
13695 schema
13696 .validate_values(&cells)
13697 .map_err(|error| MongrelError::CorruptWal {
13698 offset: record.seq.0,
13699 reason: format!("committed Put violates table schema: {error}"),
13700 })?;
13701 if schema.auto_increment_column().is_some_and(|column| {
13702 matches!(
13703 cells.iter().find(|(id, _)| *id == column.id),
13704 Some((_, Value::Int64(value))) if *value == i64::MAX
13705 )
13706 }) {
13707 return Err(MongrelError::CorruptWal {
13708 offset: record.seq.0,
13709 reason: "committed Put exhausts AUTO_INCREMENT".into(),
13710 });
13711 }
13712 }
13713 }
13714 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
13715 if *table_id != manifest.table_id =>
13716 {
13717 return Err(MongrelError::CorruptWal {
13718 offset: record.seq.0,
13719 reason: format!(
13720 "private WAL record references table {table_id}, expected {}",
13721 manifest.table_id
13722 ),
13723 });
13724 }
13725 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
13726 return Err(MongrelError::CorruptWal {
13727 offset: record.seq.0,
13728 reason: "private WAL contains shared spilled-run metadata".into(),
13729 });
13730 }
13731 _ => {}
13732 }
13733 }
13734 Ok(())
13735}
13736
13737fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
13738 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
13739}
13740
13741fn wal_segment_number(path: &Path) -> Option<u64> {
13742 path.file_stem()
13743 .and_then(|stem| stem.to_str())
13744 .and_then(|stem| stem.strip_prefix("seg-"))
13745 .and_then(|number| number.parse().ok())
13746}
13747
13748fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
13749 let n = list_wal_numbers(wal_dir)?;
13750 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
13751}
13752
13753fn next_wal_number(wal_dir: &Path) -> Result<u32> {
13754 list_wal_numbers(wal_dir)?
13755 .map(|maximum| {
13756 maximum
13757 .checked_add(1)
13758 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
13759 })
13760 .unwrap_or(Ok(0))
13761}
13762
13763fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
13764 let mut max_n = None;
13765 let entries = match std::fs::read_dir(wal_dir) {
13766 Ok(entries) => entries,
13767 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
13768 Err(error) => return Err(error.into()),
13769 };
13770 for entry in entries {
13771 let entry = entry?;
13772 let fname = entry.file_name();
13773 let Some(s) = fname.to_str() else {
13774 continue;
13775 };
13776 let Some(stripped) = s.strip_prefix("seg-") else {
13777 continue;
13778 };
13779 let Some(number) = stripped.strip_suffix(".wal") else {
13780 return Err(MongrelError::CorruptWal {
13781 offset: 0,
13782 reason: format!("malformed WAL segment name {s:?}"),
13783 });
13784 };
13785 let n = number
13786 .parse::<u32>()
13787 .map_err(|_| MongrelError::CorruptWal {
13788 offset: 0,
13789 reason: format!("malformed WAL segment name {s:?}"),
13790 })?;
13791 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
13792 return Err(MongrelError::CorruptWal {
13793 offset: n as u64,
13794 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
13795 });
13796 }
13797 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
13798 }
13799 Ok(max_n)
13800}