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, MemtableVisibleVersionCursor, Row, Value};
22use crate::mutable_run::{MutableRun, MutableRunVisibleVersionCursor};
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<'a> {
92 Memory(Row),
93 Memtable(RowId, Epoch, &'a Row),
95 MutableRun(RowId, Epoch, &'a Row),
97 Run(RunVisibleVersion),
98}
99
100impl<'a> ControlledVisibleCandidate<'a> {
101 fn row_id(&self) -> RowId {
102 match self {
103 Self::Memory(row) => row.row_id,
104 Self::Memtable(rid, _, _) => *rid,
105 Self::MutableRun(rid, _, _) => *rid,
106 Self::Run(version) => version.row_id,
107 }
108 }
109
110 fn committed_epoch(&self) -> Epoch {
111 match self {
112 Self::Memory(row) => row.committed_epoch,
113 Self::Memtable(_, epoch, _) => *epoch,
114 Self::MutableRun(_, epoch, _) => *epoch,
115 Self::Run(version) => version.committed_epoch,
116 }
117 }
118
119 fn deleted(&self) -> bool {
120 match self {
121 Self::Memory(row) => row.deleted,
122 Self::Memtable(_, _, row) => row.deleted,
123 Self::MutableRun(_, _, row) => row.deleted,
124 Self::Run(version) => version.deleted,
125 }
126 }
127
128 fn commit_ts(&self) -> Option<mongreldb_types::hlc::HlcTimestamp> {
129 match self {
130 Self::Memory(row) => row.commit_ts,
131 Self::Memtable(_, _, row) => row.commit_ts,
132 Self::MutableRun(_, _, row) => row.commit_ts,
133 Self::Run(version) => version.commit_ts,
136 }
137 }
138}
139
140enum ControlledVisibleCursor<'a> {
141 Memory(std::vec::IntoIter<Row>),
143 MemoryStreaming {
147 active: std::vec::IntoIter<Row>,
148 rest: std::collections::btree_map::IntoValues<RowId, Row>,
149 batch_cap: usize,
150 #[allow(dead_code)]
152 peak_active: usize,
153 #[allow(dead_code)]
155 total: usize,
156 },
157 Memtable(MemtableVisibleVersionCursor<'a>),
161 MutableRun(MutableRunVisibleVersionCursor<'a>),
163 Run(Box<RunVisibleVersionCursor>),
164 #[cfg(test)]
165 Synthetic {
166 next: u64,
167 end: u64,
168 },
169}
170
171#[allow(dead_code)] const CONTROLLED_HOT_BATCH: usize = 256;
174
175struct ControlledVisibleSource<'a> {
176 cursor: ControlledVisibleCursor<'a>,
177 current: Option<ControlledVisibleCandidate<'a>>,
178}
179
180impl<'a> ControlledVisibleSource<'a> {
181 #[cfg(test)]
183 fn memory(rows: Vec<Row>) -> Self {
184 Self {
185 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
186 current: None,
187 }
188 }
189
190 #[allow(dead_code)] fn memory_from_map(map: BTreeMap<RowId, Row>) -> Self {
195 let total = map.len();
196 let mut values = map.into_values();
197 if total > CONTROLLED_HOT_BATCH {
198 let active: Vec<Row> = values.by_ref().take(CONTROLLED_HOT_BATCH).collect();
199 let peak = active.len();
200 Self {
201 cursor: ControlledVisibleCursor::MemoryStreaming {
202 active: active.into_iter(),
203 rest: values,
204 batch_cap: CONTROLLED_HOT_BATCH,
205 peak_active: peak,
206 total,
207 },
208 current: None,
209 }
210 } else {
211 let active: Vec<Row> = values.collect();
212 Self {
213 cursor: ControlledVisibleCursor::Memory(active.into_iter()),
214 current: None,
215 }
216 }
217 }
218
219 fn memtable_cursor(cursor: MemtableVisibleVersionCursor<'a>) -> Self {
223 Self {
224 cursor: ControlledVisibleCursor::Memtable(cursor),
225 current: None,
226 }
227 }
228
229 fn mutable_run_cursor(cursor: MutableRunVisibleVersionCursor<'a>) -> Self {
233 Self {
234 cursor: ControlledVisibleCursor::MutableRun(cursor),
235 current: None,
236 }
237 }
238
239 fn run(cursor: RunVisibleVersionCursor) -> Self {
240 Self {
241 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
242 current: None,
243 }
244 }
245
246 #[cfg(test)]
247 fn synthetic(end: u64) -> Self {
248 Self {
249 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
250 current: None,
251 }
252 }
253
254 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
255 self.current = match &mut self.cursor {
256 ControlledVisibleCursor::Memory(rows) => {
257 rows.next().map(ControlledVisibleCandidate::Memory)
258 }
259 ControlledVisibleCursor::MemoryStreaming {
260 active,
261 rest,
262 batch_cap,
263 peak_active,
264 total: _,
265 } => {
266 if let Some(row) = active.next() {
267 Some(ControlledVisibleCandidate::Memory(row))
268 } else {
269 control.checkpoint()?;
270 let next_batch: Vec<Row> = rest.by_ref().take(*batch_cap).collect();
271 if next_batch.is_empty() {
272 None
273 } else {
274 *peak_active = (*peak_active).max(next_batch.len());
275 *active = next_batch.into_iter();
276 active.next().map(ControlledVisibleCandidate::Memory)
277 }
278 }
279 }
280 ControlledVisibleCursor::Memtable(iter) => iter
281 .next()
282 .map(|(rid, epoch, row)| ControlledVisibleCandidate::Memtable(rid, epoch, row)),
283 ControlledVisibleCursor::MutableRun(iter) => iter
284 .next()
285 .map(|(rid, epoch, row)| ControlledVisibleCandidate::MutableRun(rid, epoch, row)),
286 ControlledVisibleCursor::Run(cursor) => cursor
287 .next_visible_version(control)?
288 .map(ControlledVisibleCandidate::Run),
289 #[cfg(test)]
290 ControlledVisibleCursor::Synthetic { next, end } => {
291 if *next > *end {
292 None
293 } else {
294 let row = Row::new(RowId(*next), Epoch(1));
295 *next += 1;
296 Some(ControlledVisibleCandidate::Memory(row))
297 }
298 }
299 };
300 Ok(())
301 }
302
303 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate<'a>> {
304 let current = self.current.take().ok_or_else(|| {
305 MongrelError::Other("controlled visible source was not primed".into())
306 })?;
307 self.advance(control)?;
308 Ok(current)
309 }
310
311 fn materialize(
312 &mut self,
313 candidate: ControlledVisibleCandidate<'a>,
314 control: &crate::ExecutionControl,
315 ) -> Result<Row> {
316 match candidate {
317 ControlledVisibleCandidate::Memory(row) => Ok(row),
318 ControlledVisibleCandidate::Memtable(_, _, row) => Ok(row.clone()),
319 ControlledVisibleCandidate::MutableRun(_, _, row) => Ok(row.clone()),
320 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
321 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
322 _ => Err(MongrelError::Other(
323 "run candidate escaped its controlled cursor".into(),
324 )),
325 },
326 }
327 }
328
329 #[cfg(test)]
330 fn is_streaming_memory(&self) -> bool {
331 matches!(self.cursor, ControlledVisibleCursor::MemoryStreaming { .. })
332 }
333
334 #[cfg(test)]
335 fn streaming_peak_active(&self) -> Option<usize> {
336 match &self.cursor {
337 ControlledVisibleCursor::MemoryStreaming { peak_active, .. } => Some(*peak_active),
338 _ => None,
339 }
340 }
341
342 #[cfg(test)]
343 fn streaming_total(&self) -> Option<usize> {
344 match &self.cursor {
345 ControlledVisibleCursor::MemoryStreaming { total, .. } => Some(*total),
346 _ => None,
347 }
348 }
349}
350
351fn merge_controlled_visible_sources<'a>(
352 sources: &mut [ControlledVisibleSource<'a>],
353 control: &crate::ExecutionControl,
354 mut expired: impl FnMut(&Row) -> bool,
355 mut visit: impl FnMut(Row) -> Result<()>,
356) -> Result<()> {
357 let mut heap = BinaryHeap::new();
358 for (source_index, source) in sources.iter_mut().enumerate() {
359 source.advance(control)?;
360 if let Some(candidate) = &source.current {
361 heap.push(Reverse((candidate.row_id(), source_index)));
362 }
363 }
364 let mut merged = 0_usize;
365 while let Some(Reverse((row_id, source_index))) = heap.pop() {
366 if merged.is_multiple_of(256) {
367 control.checkpoint()?;
368 }
369 merged += 1;
370 let mut best_source = source_index;
371 let mut best = sources[source_index].pop(control)?;
372 if let Some(next) = &sources[source_index].current {
373 heap.push(Reverse((next.row_id(), source_index)));
374 }
375 while heap
376 .peek()
377 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
378 {
379 let Some(Reverse((_, source_index))) = heap.pop() else {
380 break;
381 };
382 let candidate = sources[source_index].pop(control)?;
383 if Snapshot::version_is_newer(
387 candidate.committed_epoch(),
388 candidate.commit_ts(),
389 best.committed_epoch(),
390 best.commit_ts(),
391 ) {
392 best = candidate;
393 best_source = source_index;
394 }
395 if let Some(next) = &sources[source_index].current {
396 heap.push(Reverse((next.row_id(), source_index)));
397 }
398 }
399 if best.deleted() {
400 continue;
401 }
402 let row = sources[best_source].materialize(best, control)?;
403 if !expired(&row) {
404 visit(row)?;
405 }
406 }
407 control.checkpoint()
408}
409
410#[cfg(test)]
411mod controlled_visible_cursor_tests {
412 use super::*;
413
414 #[test]
415 fn streams_more_than_one_million_rows_without_a_source_cap() {
416 let control = crate::ExecutionControl::new(None);
417 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
418 let mut count = 0_u64;
419 let mut last = 0_u64;
420 merge_controlled_visible_sources(
421 &mut sources,
422 &control,
423 |_| false,
424 |row| {
425 count += 1;
426 assert!(row.row_id.0 > last);
427 last = row.row_id.0;
428 Ok(())
429 },
430 )
431 .unwrap();
432 assert_eq!(count, 1_000_001);
433 assert_eq!(last, 1_000_001);
434 }
435
436 #[test]
437 fn merge_orders_rows_and_honors_newest_tombstones() {
438 let control = crate::ExecutionControl::new(None);
439 let older = vec![
440 Row::new(RowId(1), Epoch(1)),
441 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
442 Row::new(RowId(4), Epoch(1)),
443 ];
444 let mut deleted = Row::new(RowId(1), Epoch(2));
445 deleted.deleted = true;
446 let newer = vec![
447 deleted,
448 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
449 Row::new(RowId(3), Epoch(2)),
450 ];
451 let mut sources = vec![
452 ControlledVisibleSource::memory(older),
453 ControlledVisibleSource::memory(newer),
454 ];
455 let mut rows = Vec::new();
456 merge_controlled_visible_sources(
457 &mut sources,
458 &control,
459 |_| false,
460 |row| {
461 rows.push(row);
462 Ok(())
463 },
464 )
465 .unwrap();
466 assert_eq!(
467 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
468 vec![2, 3, 4]
469 );
470 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
471 }
472
473 #[test]
474 fn controlled_merge_uses_hlc_authority_when_epochs_inverted() {
475 use mongreldb_types::hlc::HlcTimestamp;
476 let hlc_old = HlcTimestamp {
477 physical_micros: 100,
478 logical: 0,
479 node_tiebreaker: 1,
480 };
481 let hlc_new = HlcTimestamp {
482 physical_micros: 200,
483 logical: 0,
484 node_tiebreaker: 1,
485 };
486 let a =
489 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
490 let b =
491 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
492 let control = crate::ExecutionControl::new(None);
493 let mut sources = vec![
494 ControlledVisibleSource::memory(a),
495 ControlledVisibleSource::memory(b),
496 ];
497 let mut rows = Vec::new();
498 merge_controlled_visible_sources(
499 &mut sources,
500 &control,
501 |_| false,
502 |row| {
503 rows.push(row);
504 Ok(())
505 },
506 )
507 .unwrap();
508 assert_eq!(rows.len(), 1);
509 assert_eq!(rows[0].row_id.0, 1);
510 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(11)));
512 assert_eq!(rows[0].commit_ts, Some(hlc_new));
513 }
514
515 #[test]
516 fn controlled_merge_epoch_wins_when_one_side_lacks_hlc() {
517 use mongreldb_types::hlc::HlcTimestamp;
518 let hlc_old = HlcTimestamp {
519 physical_micros: 100,
520 logical: 0,
521 node_tiebreaker: 1,
522 };
523 let a =
526 vec![Row::new_with_hlc(RowId(1), Epoch(10), hlc_old).with_column(1, Value::Int64(10))];
527 let b = vec![Row::new(RowId(1), Epoch(5)).with_column(1, Value::Int64(5))];
528 let control = crate::ExecutionControl::new(None);
529 let mut sources = vec![
530 ControlledVisibleSource::memory(a),
531 ControlledVisibleSource::memory(b),
532 ];
533 let mut rows = Vec::new();
534 merge_controlled_visible_sources(
535 &mut sources,
536 &control,
537 |_| false,
538 |row| {
539 rows.push(row);
540 Ok(())
541 },
542 )
543 .unwrap();
544 assert_eq!(rows.len(), 1);
545 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(10)));
546 }
547
548 #[test]
551 fn controlled_merge_memory_vs_run_uses_hlc_when_run_stamped() {
552 use crate::sorted_run::{RunReader, RunWriter};
553 use mongreldb_types::hlc::HlcTimestamp;
554 use tempfile::tempdir;
555
556 let hlc_old = HlcTimestamp {
557 physical_micros: 100,
558 logical: 0,
559 node_tiebreaker: 1,
560 };
561 let hlc_new = HlcTimestamp {
562 physical_micros: 200,
563 logical: 0,
564 node_tiebreaker: 1,
565 };
566 let schema = Schema {
567 schema_id: 1,
568 columns: vec![ColumnDef {
569 id: 1,
570 name: "v".into(),
571 ty: TypeId::Int64,
572 flags: ColumnFlags::empty(),
573 default_value: None,
574 embedding_source: None,
575 }],
576 indexes: vec![],
577 colocation: vec![],
578 constraints: Default::default(),
579 clustered: false,
580 };
581 let dir = tempdir().unwrap();
582 let path = dir.path().join("r-hlc.sr");
583 let run_rows =
585 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
586 RunWriter::new(&schema, 1, Epoch(50), 0)
587 .write(&path, &run_rows)
588 .unwrap();
589 let reader = RunReader::open(&path, schema, None).unwrap();
590 assert!(reader.has_column(crate::sorted_run::SYS_COMMIT_TS));
591
592 let mem =
594 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
595 let control = crate::ExecutionControl::new(None);
596 let mut sources = vec![
597 ControlledVisibleSource::memory(mem),
598 ControlledVisibleSource::run(
599 reader.into_visible_version_cursor(Epoch(u64::MAX)).unwrap(),
600 ),
601 ];
602 let mut rows = Vec::new();
603 merge_controlled_visible_sources(
604 &mut sources,
605 &control,
606 |_| false,
607 |row| {
608 rows.push(row);
609 Ok(())
610 },
611 )
612 .unwrap();
613 assert_eq!(rows.len(), 1);
614 assert_eq!(
615 rows[0].columns.get(&1),
616 Some(&Value::Int64(11)),
617 "HLC-newer memory version must beat epoch-newer run"
618 );
619 assert_eq!(rows[0].commit_ts, Some(hlc_new));
620 }
621
622 #[test]
623 fn controlled_memory_source_streams_large_overlays_without_full_active_vec() {
624 let mut map = BTreeMap::new();
625 for i in 1..=500u64 {
626 map.insert(
627 RowId(i),
628 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
629 );
630 }
631 let source = ControlledVisibleSource::memory_from_map(map);
632 assert!(
633 source.is_streaming_memory(),
634 "overlays larger than CONTROLLED_HOT_BATCH must use streaming cursor"
635 );
636 assert_eq!(source.streaming_total(), Some(500));
637 assert!(
639 source.streaming_peak_active().unwrap_or(usize::MAX) <= CONTROLLED_HOT_BATCH,
640 "peak active buffer must be <= CONTROLLED_HOT_BATCH"
641 );
642
643 let control = crate::ExecutionControl::new(None);
645 let mut sources = vec![ControlledVisibleSource::memory_from_map({
646 let mut m = BTreeMap::new();
647 for i in 1..=500u64 {
648 m.insert(
649 RowId(i),
650 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
651 );
652 }
653 m
654 })];
655 let mut n = 0usize;
656 merge_controlled_visible_sources(
657 &mut sources,
658 &control,
659 |_| false,
660 |_| {
661 n += 1;
662 Ok(())
663 },
664 )
665 .unwrap();
666 assert_eq!(n, 500);
667 assert!(
668 sources[0].streaming_peak_active().unwrap_or(0) <= CONTROLLED_HOT_BATCH,
669 "after full drain peak active still <= batch cap"
670 );
671 }
672}
673
674fn iso_now_bytes() -> Vec<u8> {
677 let secs = std::time::SystemTime::now()
678 .duration_since(std::time::UNIX_EPOCH)
679 .map(|d| d.as_secs() as i64)
680 .unwrap_or(0);
681 let days = secs.div_euclid(86_400);
682 let rem = secs.rem_euclid(86_400);
683 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
684 let (year, month, day) = civil_from_days(days);
685 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
686}
687
688pub(crate) fn unix_nanos_now() -> i64 {
689 std::time::SystemTime::now()
690 .duration_since(std::time::UNIX_EPOCH)
691 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
692 .unwrap_or(0)
693}
694
695fn ann_candidate_cap(
696 index_len: usize,
697 context: Option<&crate::query::AiExecutionContext>,
698) -> usize {
699 index_len
700 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
701 .min(context.map_or(
702 crate::query::MAX_RAW_INDEX_CANDIDATES,
703 crate::query::AiExecutionContext::max_fused_candidates,
704 ))
705}
706
707#[cfg(test)]
708mod ann_candidate_cap_tests {
709 use super::*;
710
711 #[test]
712 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
713 assert_eq!(
714 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
715 crate::query::MAX_RAW_INDEX_CANDIDATES,
716 );
717 let context = crate::query::AiExecutionContext::with_limits(
718 std::time::Duration::from_secs(1),
719 usize::MAX,
720 17,
721 );
722 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
723 }
724}
725
726fn civil_from_days(z: i64) -> (i64, u32, u32) {
727 let z = z + 719_468;
728 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
729 let doe = z - era * 146_097;
730 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
731 let y = yoe + era * 400;
732 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
733 let mp = (5 * doy + 2) / 153;
734 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
735 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
736 (if m <= 2 { y + 1 } else { y }, m, d)
737}
738
739pub fn clustered_row_id(primary_key: &Value) -> RowId {
746 let mut hash: u64 = 0xcbf29ce484222325;
747 for byte in primary_key.encode_key() {
748 hash ^= u64::from(byte);
749 hash = hash.wrapping_mul(0x100000001b3);
750 }
751 RowId(hash.max(1))
752}
753
754const 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;
761
762#[derive(Clone, Copy, Debug)]
777struct AutoIncState {
778 column_id: u16,
779 next: i64,
780 seeded: bool,
781}
782
783pub(crate) struct RecoveryMetadataPlan {
784 live_count: u64,
785 auto_inc: Option<AutoIncState>,
786 changed: bool,
787}
788
789type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
790
791fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
794 schema.auto_increment_column().map(|c| AutoIncState {
795 column_id: c.id,
796 next: 0,
797 seeded: false,
798 })
799}
800
801#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
814pub enum IndexBuildPolicy {
815 #[default]
817 Deferred,
818 Eager,
820}
821
822#[derive(Clone)]
823struct ReversePkSegment {
824 values: HashMap<RowId, Vec<u8>>,
825 removed: HashSet<RowId>,
826}
827
828#[derive(Clone)]
829struct ReversePkMap {
830 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
831 active: ReversePkSegment,
832}
833
834impl ReversePkMap {
835 fn new() -> Self {
836 Self {
837 frozen: Arc::new(Vec::new()),
838 active: ReversePkSegment {
839 values: HashMap::new(),
840 removed: HashSet::new(),
841 },
842 }
843 }
844
845 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
846 let mut map = Self::new();
847 map.active.values.extend(entries);
848 map
849 }
850
851 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
852 self.active.removed.remove(&row_id);
853 self.active.values.insert(row_id, key);
854 }
855
856 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
857 if let Some(key) = self.active.values.get(row_id) {
858 return Some(key);
859 }
860 if self.active.removed.contains(row_id) {
861 return None;
862 }
863 for segment in self.frozen.iter().rev() {
864 if let Some(key) = segment.values.get(row_id) {
865 return Some(key);
866 }
867 if segment.removed.contains(row_id) {
868 return None;
869 }
870 }
871 None
872 }
873
874 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
875 let previous = self.get(row_id).cloned();
876 self.active.values.remove(row_id);
877 self.active.removed.insert(*row_id);
878 previous
879 }
880
881 fn clear(&mut self) {
882 *self = Self::new();
883 }
884
885 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
886 let mut entries = HashMap::new();
887 for segment in self
888 .frozen
889 .iter()
890 .map(Arc::as_ref)
891 .chain(std::iter::once(&self.active))
892 {
893 for row_id in &segment.removed {
894 entries.remove(row_id);
895 }
896 entries.extend(
897 segment
898 .values
899 .iter()
900 .map(|(row_id, key)| (*row_id, key.clone())),
901 );
902 }
903 entries
904 }
905
906 fn seal(&mut self) {
907 if self.active.values.is_empty() && self.active.removed.is_empty() {
908 return;
909 }
910 let active = std::mem::replace(
911 &mut self.active,
912 ReversePkSegment {
913 values: HashMap::new(),
914 removed: HashSet::new(),
915 },
916 );
917 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
918 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
919 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
920 values: self.entries(),
921 removed: HashSet::new(),
922 })]);
923 }
924 }
925}
926
927#[derive(Clone)]
940pub struct ReadGeneration {
941 schema: Arc<Schema>,
942 base_runs: Arc<Vec<RunRef>>,
943 deltas: TableDeltas,
944 indexes: Arc<IndexGeneration>,
945 visible_through: Epoch,
946}
947
948pub(crate) enum SecondaryIndexArtifact {
952 Bitmap(u16, BitmapIndex),
953 LearnedRange(u16, ColumnLearnedRange),
954 Fm(u16, Box<FmIndex>),
955 Ann(u16, Box<AnnIndex>),
956 Sparse(u16, SparseIndex),
957 MinHash(u16, MinHashIndex),
958}
959
960#[derive(Clone)]
966pub struct TableDeltas {
967 memtable: Memtable,
968 mutable_run: MutableRun,
969 hot: HotIndex,
970 pk_by_row: ReversePkMap,
971}
972
973impl ReadGeneration {
974 fn empty(schema: &Schema) -> Self {
977 Self {
978 schema: Arc::new(schema.clone()),
979 base_runs: Arc::new(Vec::new()),
980 deltas: TableDeltas {
981 memtable: Memtable::new(),
982 mutable_run: MutableRun::new(),
983 hot: HotIndex::new(),
984 pk_by_row: ReversePkMap::new(),
985 },
986 indexes: Arc::new(IndexGeneration::default()),
987 visible_through: Epoch(0),
988 }
989 }
990
991 pub fn schema(&self) -> &Arc<Schema> {
993 &self.schema
994 }
995
996 pub fn base_runs(&self) -> &[RunRef] {
998 &self.base_runs
999 }
1000
1001 pub fn indexes(&self) -> &Arc<IndexGeneration> {
1003 &self.indexes
1004 }
1005
1006 pub fn visible_through(&self) -> Epoch {
1008 self.visible_through
1009 }
1010
1011 pub fn deltas(&self) -> &TableDeltas {
1015 &self.deltas
1016 }
1017}
1018
1019impl TableDeltas {
1020 pub fn approx_bytes(&self) -> u64 {
1023 self.memtable
1024 .approx_bytes()
1025 .saturating_add(self.mutable_run.approx_bytes())
1026 }
1027
1028 pub fn memtable_len(&self) -> usize {
1030 self.memtable.len()
1031 }
1032
1033 pub fn mutable_run_len(&self) -> usize {
1035 self.mutable_run.len()
1036 }
1037
1038 pub fn hot_len(&self) -> usize {
1040 self.hot.len()
1041 }
1042
1043 pub fn reverse_pk_len(&self) -> usize {
1045 self.pk_by_row.entries().len()
1046 }
1047}
1048
1049#[derive(Clone)]
1051pub struct Table {
1052 dir: PathBuf,
1053 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1054 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
1055 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
1056 table_id: u64,
1057 name: String,
1061 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1066 read_only: bool,
1069 durable_commit_failed: bool,
1073 wal: WalSink,
1074 memtable: Memtable,
1075 mutable_run: MutableRun,
1080 mutable_run_spill_bytes: u64,
1082 compaction_zstd_level: i32,
1085 allocator: RowIdAllocator,
1086 epoch: Arc<EpochAuthority>,
1087 data_generation: u64,
1090 schema: Schema,
1091 hot: HotIndex,
1092 kek: Option<Arc<Kek>>,
1095 column_keys: HashMap<u16, ([u8; 32], u8)>,
1099 run_refs: Vec<RunRef>,
1100 retiring: Vec<crate::manifest::RetiredRun>,
1103 next_run_id: u64,
1104 sync_byte_threshold: u64,
1105 current_txn_id: u64,
1110 pending_private_mutations: bool,
1114 bitmap: HashMap<u16, BitmapIndex>,
1115 ann: HashMap<u16, AnnIndex>,
1116 fm: HashMap<u16, FmIndex>,
1117 sparse: HashMap<u16, SparseIndex>,
1118 minhash: HashMap<u16, MinHashIndex>,
1119 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
1122 pk_by_row: ReversePkMap,
1124 pinned: BTreeMap<Epoch, usize>,
1127 pub(crate) live_count: u64,
1130 reservoir: crate::reservoir::Reservoir,
1133 reservoir_complete: bool,
1141 had_deletes: bool,
1145 recent_delete_preimages: HashMap<Vec<u8>, Row>,
1150 run_row_id_ranges: HashMap<u128, (u64, u64)>,
1153 agg_cache: Arc<HashMap<u64, CachedAgg>>,
1157 global_idx_epoch: u64,
1161 indexes_complete: bool,
1166 index_build_policy: IndexBuildPolicy,
1168 pk_by_row_complete: bool,
1175 flushed_epoch: u64,
1178 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1181 snapshots: Arc<crate::retention::SnapshotRegistry>,
1184 commit_lock: Arc<parking_lot::Mutex<()>>,
1186 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1189 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
1199 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
1208 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1210 pending_delete_rids: roaring::RoaringBitmap,
1213 pending_put_cols: std::collections::HashSet<u16>,
1216 pending_rows: Vec<Row>,
1222 pending_rows_auto_inc: Vec<bool>,
1223 pending_dels: Vec<RowId>,
1226 pending_truncate: Option<Epoch>,
1230 auto_inc: Option<AutoIncState>,
1233 ttl: Option<TtlPolicy>,
1236 pins: Arc<crate::retention::PinRegistry>,
1243 published: Arc<ArcSwap<ReadGeneration>>,
1248 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
1253 lookup_metrics: LookupMetrics,
1258}
1259
1260#[derive(Debug, Default)]
1261pub struct LookupMetrics {
1262 hot_lookup_hit: std::sync::atomic::AtomicU64,
1263 hot_lookup_fallback: std::sync::atomic::AtomicU64,
1264 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64,
1265 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64,
1266 result_cache_memory_hit: std::sync::atomic::AtomicU64,
1267 result_cache_disk_hit: std::sync::atomic::AtomicU64,
1268 result_cache_miss: std::sync::atomic::AtomicU64,
1269 result_cache_persistent_write_us: std::sync::atomic::AtomicU64,
1270 get_run_opened: std::sync::atomic::AtomicU64,
1272 get_run_skipped: std::sync::atomic::AtomicU64,
1273 pub(crate) directory_lookup_hit: std::sync::atomic::AtomicU64,
1275 pub(crate) directory_lookup_fallback: std::sync::atomic::AtomicU64,
1276 pub(crate) directory_incomplete: std::sync::atomic::AtomicU64,
1277 pub(crate) directory_run_readers_opened: std::sync::atomic::AtomicU64,
1278 pub(crate) directory_early_stop_total: std::sync::atomic::AtomicU64,
1279 pub(crate) result_cache_persist_enqueued_total: std::sync::atomic::AtomicU64,
1281 pub(crate) result_cache_persist_coalesced_total: std::sync::atomic::AtomicU64,
1282 pub(crate) result_cache_persist_dropped_store_total: std::sync::atomic::AtomicU64,
1283 pub(crate) result_cache_persist_remove_total: std::sync::atomic::AtomicU64,
1284 pub(crate) result_cache_persist_stale_store_skipped_total: std::sync::atomic::AtomicU64,
1285 pub(crate) result_cache_persist_errors_total: std::sync::atomic::AtomicU64,
1286 pub(crate) result_cache_persist_shutdown_abandoned_total: std::sync::atomic::AtomicU64,
1287 pub(crate) result_cache_persist_queue_depth: std::sync::atomic::AtomicU64,
1288 pub(crate) hot_fallback_reasons: [std::sync::atomic::AtomicU64; 9],
1290 pub(crate) hot_fallback_overlay_versions_total: std::sync::atomic::AtomicU64,
1291 pub(crate) hot_fallback_runs_considered_total: std::sync::atomic::AtomicU64,
1292 pub(crate) hot_fallback_runs_opened_total: std::sync::atomic::AtomicU64,
1293 pub(crate) hot_fallback_pages_decoded_total: std::sync::atomic::AtomicU64,
1294 pub(crate) hot_fallback_rows_materialized_total: std::sync::atomic::AtomicU64,
1295 pub(crate) hot_lookup_duration_nanos: std::sync::atomic::AtomicU64,
1296 pub(crate) hot_fallback_duration_nanos: std::sync::atomic::AtomicU64,
1297 pub(crate) hot_mapping_rebuild_total: std::sync::atomic::AtomicU64,
1298 pub(crate) hot_checkpoint_rejected_total: std::sync::atomic::AtomicU64,
1299}
1300
1301impl Clone for LookupMetrics {
1302 fn clone(&self) -> Self {
1303 let mut hot_fallback_reasons = [
1304 std::sync::atomic::AtomicU64::new(0),
1305 std::sync::atomic::AtomicU64::new(0),
1306 std::sync::atomic::AtomicU64::new(0),
1307 std::sync::atomic::AtomicU64::new(0),
1308 std::sync::atomic::AtomicU64::new(0),
1309 std::sync::atomic::AtomicU64::new(0),
1310 std::sync::atomic::AtomicU64::new(0),
1311 std::sync::atomic::AtomicU64::new(0),
1312 std::sync::atomic::AtomicU64::new(0),
1313 ];
1314 for (dst, src) in hot_fallback_reasons
1315 .iter_mut()
1316 .zip(self.hot_fallback_reasons.iter())
1317 {
1318 dst.store(
1319 src.load(std::sync::atomic::Ordering::Relaxed),
1320 std::sync::atomic::Ordering::Relaxed,
1321 );
1322 }
1323 let copy_atomic = |src: &std::sync::atomic::AtomicU64| {
1324 std::sync::atomic::AtomicU64::new(src.load(std::sync::atomic::Ordering::Relaxed))
1325 };
1326 Self {
1327 hot_lookup_hit: copy_atomic(&self.hot_lookup_hit),
1328 hot_lookup_fallback: copy_atomic(&self.hot_lookup_fallback),
1329 hot_lookup_fallback_overlay_rows: copy_atomic(&self.hot_lookup_fallback_overlay_rows),
1330 hot_lookup_fallback_runs: copy_atomic(&self.hot_lookup_fallback_runs),
1331 result_cache_memory_hit: copy_atomic(&self.result_cache_memory_hit),
1332 result_cache_disk_hit: copy_atomic(&self.result_cache_disk_hit),
1333 result_cache_miss: copy_atomic(&self.result_cache_miss),
1334 result_cache_persistent_write_us: copy_atomic(&self.result_cache_persistent_write_us),
1335 get_run_opened: copy_atomic(&self.get_run_opened),
1336 get_run_skipped: copy_atomic(&self.get_run_skipped),
1337 directory_lookup_hit: copy_atomic(&self.directory_lookup_hit),
1338 directory_lookup_fallback: copy_atomic(&self.directory_lookup_fallback),
1339 directory_incomplete: copy_atomic(&self.directory_incomplete),
1340 directory_run_readers_opened: copy_atomic(&self.directory_run_readers_opened),
1341 directory_early_stop_total: copy_atomic(&self.directory_early_stop_total),
1342 result_cache_persist_enqueued_total: copy_atomic(
1343 &self.result_cache_persist_enqueued_total,
1344 ),
1345 result_cache_persist_coalesced_total: copy_atomic(
1346 &self.result_cache_persist_coalesced_total,
1347 ),
1348 result_cache_persist_dropped_store_total: copy_atomic(
1349 &self.result_cache_persist_dropped_store_total,
1350 ),
1351 result_cache_persist_remove_total: copy_atomic(&self.result_cache_persist_remove_total),
1352 result_cache_persist_stale_store_skipped_total: copy_atomic(
1353 &self.result_cache_persist_stale_store_skipped_total,
1354 ),
1355 result_cache_persist_errors_total: copy_atomic(&self.result_cache_persist_errors_total),
1356 result_cache_persist_shutdown_abandoned_total: copy_atomic(
1357 &self.result_cache_persist_shutdown_abandoned_total,
1358 ),
1359 result_cache_persist_queue_depth: copy_atomic(&self.result_cache_persist_queue_depth),
1360 hot_fallback_reasons,
1361 hot_fallback_overlay_versions_total: copy_atomic(
1362 &self.hot_fallback_overlay_versions_total,
1363 ),
1364 hot_fallback_runs_considered_total: copy_atomic(
1365 &self.hot_fallback_runs_considered_total,
1366 ),
1367 hot_fallback_runs_opened_total: copy_atomic(&self.hot_fallback_runs_opened_total),
1368 hot_fallback_pages_decoded_total: copy_atomic(&self.hot_fallback_pages_decoded_total),
1369 hot_fallback_rows_materialized_total: copy_atomic(
1370 &self.hot_fallback_rows_materialized_total,
1371 ),
1372 hot_lookup_duration_nanos: copy_atomic(&self.hot_lookup_duration_nanos),
1373 hot_fallback_duration_nanos: copy_atomic(&self.hot_fallback_duration_nanos),
1374 hot_mapping_rebuild_total: copy_atomic(&self.hot_mapping_rebuild_total),
1375 hot_checkpoint_rejected_total: copy_atomic(&self.hot_checkpoint_rejected_total),
1376 }
1377 }
1378}
1379
1380impl LookupMetrics {
1381 fn snapshot(&self) -> LookupMetricsSnapshot {
1382 LookupMetricsSnapshot {
1383 hot_lookup_hit: self
1384 .hot_lookup_hit
1385 .load(std::sync::atomic::Ordering::Relaxed),
1386 hot_lookup_fallback: self
1387 .hot_lookup_fallback
1388 .load(std::sync::atomic::Ordering::Relaxed),
1389 hot_lookup_fallback_overlay_rows: self
1390 .hot_lookup_fallback_overlay_rows
1391 .load(std::sync::atomic::Ordering::Relaxed),
1392 hot_lookup_fallback_runs: self
1393 .hot_lookup_fallback_runs
1394 .load(std::sync::atomic::Ordering::Relaxed),
1395 result_cache_memory_hit: self
1396 .result_cache_memory_hit
1397 .load(std::sync::atomic::Ordering::Relaxed),
1398 result_cache_disk_hit: self
1399 .result_cache_disk_hit
1400 .load(std::sync::atomic::Ordering::Relaxed),
1401 result_cache_miss: self
1402 .result_cache_miss
1403 .load(std::sync::atomic::Ordering::Relaxed),
1404 result_cache_persistent_write_us: self
1405 .result_cache_persistent_write_us
1406 .load(std::sync::atomic::Ordering::Relaxed),
1407 get_run_opened: self
1408 .get_run_opened
1409 .load(std::sync::atomic::Ordering::Relaxed),
1410 get_run_skipped: self
1411 .get_run_skipped
1412 .load(std::sync::atomic::Ordering::Relaxed),
1413 directory_lookup_hit: self
1414 .directory_lookup_hit
1415 .load(std::sync::atomic::Ordering::Relaxed),
1416 directory_lookup_fallback: self
1417 .directory_lookup_fallback
1418 .load(std::sync::atomic::Ordering::Relaxed),
1419 directory_incomplete: self
1420 .directory_incomplete
1421 .load(std::sync::atomic::Ordering::Relaxed),
1422 directory_run_readers_opened: self
1423 .directory_run_readers_opened
1424 .load(std::sync::atomic::Ordering::Relaxed),
1425 directory_early_stop_total: self
1426 .directory_early_stop_total
1427 .load(std::sync::atomic::Ordering::Relaxed),
1428 result_cache_persist_enqueued_total: self
1429 .result_cache_persist_enqueued_total
1430 .load(std::sync::atomic::Ordering::Relaxed),
1431 result_cache_persist_coalesced_total: self
1432 .result_cache_persist_coalesced_total
1433 .load(std::sync::atomic::Ordering::Relaxed),
1434 result_cache_persist_dropped_store_total: self
1435 .result_cache_persist_dropped_store_total
1436 .load(std::sync::atomic::Ordering::Relaxed),
1437 result_cache_persist_remove_total: self
1438 .result_cache_persist_remove_total
1439 .load(std::sync::atomic::Ordering::Relaxed),
1440 result_cache_persist_stale_store_skipped_total: self
1441 .result_cache_persist_stale_store_skipped_total
1442 .load(std::sync::atomic::Ordering::Relaxed),
1443 result_cache_persist_errors_total: self
1444 .result_cache_persist_errors_total
1445 .load(std::sync::atomic::Ordering::Relaxed),
1446 result_cache_persist_shutdown_abandoned_total: self
1447 .result_cache_persist_shutdown_abandoned_total
1448 .load(std::sync::atomic::Ordering::Relaxed),
1449 result_cache_persist_queue_depth: self
1450 .result_cache_persist_queue_depth
1451 .load(std::sync::atomic::Ordering::Relaxed),
1452 hot_fallback_reasons: [
1453 self.hot_fallback_reasons[0].load(std::sync::atomic::Ordering::Relaxed),
1454 self.hot_fallback_reasons[1].load(std::sync::atomic::Ordering::Relaxed),
1455 self.hot_fallback_reasons[2].load(std::sync::atomic::Ordering::Relaxed),
1456 self.hot_fallback_reasons[3].load(std::sync::atomic::Ordering::Relaxed),
1457 self.hot_fallback_reasons[4].load(std::sync::atomic::Ordering::Relaxed),
1458 self.hot_fallback_reasons[5].load(std::sync::atomic::Ordering::Relaxed),
1459 self.hot_fallback_reasons[6].load(std::sync::atomic::Ordering::Relaxed),
1460 self.hot_fallback_reasons[7].load(std::sync::atomic::Ordering::Relaxed),
1461 self.hot_fallback_reasons[8].load(std::sync::atomic::Ordering::Relaxed),
1462 ],
1463 hot_fallback_overlay_versions_total: self
1464 .hot_fallback_overlay_versions_total
1465 .load(std::sync::atomic::Ordering::Relaxed),
1466 hot_fallback_runs_considered_total: self
1467 .hot_fallback_runs_considered_total
1468 .load(std::sync::atomic::Ordering::Relaxed),
1469 hot_fallback_runs_opened_total: self
1470 .hot_fallback_runs_opened_total
1471 .load(std::sync::atomic::Ordering::Relaxed),
1472 hot_fallback_pages_decoded_total: self
1473 .hot_fallback_pages_decoded_total
1474 .load(std::sync::atomic::Ordering::Relaxed),
1475 hot_fallback_rows_materialized_total: self
1476 .hot_fallback_rows_materialized_total
1477 .load(std::sync::atomic::Ordering::Relaxed),
1478 hot_lookup_duration_nanos: self
1479 .hot_lookup_duration_nanos
1480 .load(std::sync::atomic::Ordering::Relaxed),
1481 hot_fallback_duration_nanos: self
1482 .hot_fallback_duration_nanos
1483 .load(std::sync::atomic::Ordering::Relaxed),
1484 hot_mapping_rebuild_total: self
1485 .hot_mapping_rebuild_total
1486 .load(std::sync::atomic::Ordering::Relaxed),
1487 hot_checkpoint_rejected_total: self
1488 .hot_checkpoint_rejected_total
1489 .load(std::sync::atomic::Ordering::Relaxed),
1490 }
1491 }
1492}
1493
1494#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
1497pub struct LookupMetricsSnapshot {
1498 pub hot_lookup_hit: u64,
1499 pub hot_lookup_fallback: u64,
1500 pub hot_lookup_fallback_overlay_rows: u64,
1501 pub hot_lookup_fallback_runs: u64,
1502 pub result_cache_memory_hit: u64,
1503 pub result_cache_disk_hit: u64,
1504 pub result_cache_miss: u64,
1505 pub result_cache_persistent_write_us: u64,
1506 pub get_run_opened: u64,
1507 pub get_run_skipped: u64,
1508 pub directory_lookup_hit: u64,
1510 pub directory_lookup_fallback: u64,
1511 pub directory_incomplete: u64,
1512 pub directory_run_readers_opened: u64,
1513 pub directory_early_stop_total: u64,
1514 pub result_cache_persist_enqueued_total: u64,
1516 pub result_cache_persist_coalesced_total: u64,
1517 pub result_cache_persist_dropped_store_total: u64,
1518 pub result_cache_persist_remove_total: u64,
1519 pub result_cache_persist_stale_store_skipped_total: u64,
1520 pub result_cache_persist_errors_total: u64,
1521 pub result_cache_persist_shutdown_abandoned_total: u64,
1522 pub result_cache_persist_queue_depth: u64,
1523 pub hot_fallback_reasons: [u64; 9],
1525 pub hot_fallback_overlay_versions_total: u64,
1526 pub hot_fallback_runs_considered_total: u64,
1527 pub hot_fallback_runs_opened_total: u64,
1528 pub hot_fallback_pages_decoded_total: u64,
1529 pub hot_fallback_rows_materialized_total: u64,
1530 pub hot_lookup_duration_nanos: u64,
1531 pub hot_fallback_duration_nanos: u64,
1532 pub hot_mapping_rebuild_total: u64,
1533 pub hot_checkpoint_rejected_total: u64,
1534}
1535
1536pub fn hot_fallback_reason_index(r: crate::trace::HotFallbackReason) -> usize {
1539 match r {
1540 crate::trace::HotFallbackReason::MissingMapping => 0,
1541 crate::trace::HotFallbackReason::StaleRowId => 1,
1542 crate::trace::HotFallbackReason::InvisibleAtSnapshot => 2,
1543 crate::trace::HotFallbackReason::HistoricalSnapshot => 3,
1544 crate::trace::HotFallbackReason::Tombstone => 4,
1545 crate::trace::HotFallbackReason::TtlExpired => 5,
1546 crate::trace::HotFallbackReason::PrimaryKeyMismatch => 6,
1547 crate::trace::HotFallbackReason::IndexIncomplete => 7,
1548 crate::trace::HotFallbackReason::CheckpointRejected => 8,
1549 }
1550}
1551
1552const _: () = {
1559 const fn assert_sync<T: ?Sized + Sync>() {}
1560 assert_sync::<Table>();
1561};
1562
1563enum CachedData {
1569 Rows(Arc<Vec<Row>>),
1570 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
1571}
1572
1573impl CachedData {
1574 fn approx_bytes(&self) -> u64 {
1575 match self {
1576 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
1577 CachedData::Columns(c) => c
1578 .iter()
1579 .map(|(_, c)| c.approx_bytes())
1580 .sum::<u64>()
1581 .saturating_add(c.len() as u64 * 16),
1582 }
1583 }
1584}
1585
1586struct CachedEntry {
1590 data: CachedData,
1591 footprint: roaring::RoaringBitmap,
1592 condition_cols: Vec<u16>,
1593}
1594
1595struct ResultCache {
1606 entries: std::collections::HashMap<u64, CachedEntry>,
1607 order: BTreeSet<(u64, u64)>,
1608 generations: HashMap<u64, u64>,
1609 next_generation: u64,
1610 condition_index: HashMap<u16, HashSet<u64>>,
1613 empty_footprint_entries: HashSet<u64>,
1616 bytes: u64,
1617 max_bytes: u64,
1618 dir: Option<std::path::PathBuf>,
1619 #[allow(dead_code)]
1620 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1621 memory_hit: std::sync::atomic::AtomicU64,
1625 disk_hit: std::sync::atomic::AtomicU64,
1626 miss: std::sync::atomic::AtomicU64,
1627 persistent_write_us: std::sync::atomic::AtomicU64,
1628 persist_min_bytes: u64,
1633}
1634
1635#[derive(serde::Serialize, serde::Deserialize)]
1637struct SerializedEntry {
1638 condition_cols: Vec<u16>,
1639 footprint_bits: Vec<u32>,
1640 data: SerializedData,
1641}
1642
1643#[derive(serde::Serialize, serde::Deserialize)]
1644enum SerializedData {
1645 Rows(Vec<Row>),
1646 Columns(Vec<(u16, columnar::NativeColumn)>),
1647}
1648
1649#[derive(serde::Serialize)]
1650struct SerializedEntryRef<'a> {
1651 condition_cols: &'a [u16],
1652 footprint_bits: Vec<u32>,
1653 data: SerializedDataRef<'a>,
1654}
1655
1656#[derive(serde::Serialize)]
1657enum SerializedDataRef<'a> {
1658 Rows(&'a [Row]),
1659 Columns(&'a [(u16, columnar::NativeColumn)]),
1660}
1661
1662impl<'a> SerializedEntryRef<'a> {
1663 fn from_entry(entry: &'a CachedEntry) -> Self {
1664 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
1665 let data = match &entry.data {
1666 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
1667 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
1668 };
1669 Self {
1670 condition_cols: &entry.condition_cols,
1671 footprint_bits,
1672 data,
1673 }
1674 }
1675}
1676
1677impl SerializedEntry {
1678 fn into_entry(self) -> Option<CachedEntry> {
1679 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1680 let data = match self.data {
1681 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1682 SerializedData::Columns(c) => {
1683 if !c.iter().all(|(_, col)| col.validate()) {
1686 return None;
1687 }
1688 CachedData::Columns(Arc::new(c))
1689 }
1690 };
1691 Some(CachedEntry {
1692 data,
1693 footprint,
1694 condition_cols: self.condition_cols,
1695 })
1696 }
1697}
1698
1699impl ResultCache {
1700 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1701
1702 fn new() -> Self {
1703 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1704 }
1705
1706 fn with_max_bytes(max_bytes: u64) -> Self {
1707 Self {
1708 entries: std::collections::HashMap::new(),
1709 order: BTreeSet::new(),
1710 generations: HashMap::new(),
1711 next_generation: 0,
1712 condition_index: HashMap::new(),
1713 empty_footprint_entries: HashSet::new(),
1714 bytes: 0,
1715 max_bytes,
1716 dir: None,
1717 cache_dek: None,
1718 memory_hit: std::sync::atomic::AtomicU64::new(0),
1719 disk_hit: std::sync::atomic::AtomicU64::new(0),
1720 miss: std::sync::atomic::AtomicU64::new(0),
1721 persistent_write_us: std::sync::atomic::AtomicU64::new(0),
1722 persist_min_bytes: 4 * 1024,
1725 }
1726 }
1727
1728 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1729 let _ = std::fs::create_dir_all(&dir);
1730 self.dir = Some(dir);
1731 self
1732 }
1733
1734 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1735 self.cache_dek = dek;
1736 self
1737 }
1738
1739 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1740 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1741 }
1742
1743 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1747 let Some(path) = self.disk_path(key) else {
1748 return;
1749 };
1750 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1751 Ok(s) => s,
1752 Err(_) => return,
1753 };
1754 let on_disk = if let Some(dek) = &self.cache_dek {
1756 match self.encrypt_cache(&serialized, dek) {
1757 Some(b) => b,
1758 None => return,
1759 }
1760 } else {
1761 serialized
1762 };
1763 let tmp = path.with_extension("tmp");
1764 use std::io::Write;
1765 let write = || -> std::io::Result<()> {
1766 let mut f = std::fs::File::create(&tmp)?;
1767 f.write_all(&on_disk)?;
1768 f.flush()?;
1769 Ok(())
1770 };
1771 if write().is_err() {
1772 let _ = std::fs::remove_file(&tmp);
1773 return;
1774 }
1775 let _ = std::fs::rename(&tmp, &path);
1776 }
1777
1778 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1780 let path = self.disk_path(key)?;
1781 let bytes = std::fs::read(&path).ok()?;
1782 let plaintext = if let Some(dek) = &self.cache_dek {
1783 self.decrypt_cache(&bytes, dek)?
1784 } else {
1785 bytes
1786 };
1787 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1788 serialized.into_entry()
1789 }
1790
1791 fn remove_from_disk(&self, key: u64) {
1793 if let Some(path) = self.disk_path(key) {
1794 let _ = std::fs::remove_file(&path);
1795 }
1796 }
1797
1798 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1800 use crate::encryption::Cipher;
1801 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1802 let mut nonce = [0u8; 12];
1803 crate::encryption::fill_random(&mut nonce).ok()?;
1804 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1805 let mut out = Vec::with_capacity(12 + ct.len());
1806 out.extend_from_slice(&nonce);
1807 out.extend_from_slice(&ct);
1808 Some(out)
1809 }
1810
1811 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1813 use crate::encryption::Cipher;
1814 if bytes.len() < 28 {
1815 return None;
1816 }
1817 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1818 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1819 let ct = &bytes[12..];
1820 cipher.decrypt_page(&nonce, ct).ok()
1821 }
1822
1823 fn load_persistent(&mut self) {
1826 let Some(dir) = self.dir.as_ref().cloned() else {
1827 return;
1828 };
1829 let entries = match std::fs::read_dir(&dir) {
1830 Ok(e) => e,
1831 Err(_) => return,
1832 };
1833 for entry in entries.flatten() {
1834 let path = entry.path();
1835 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1837 let _ = std::fs::remove_file(&path);
1838 continue;
1839 }
1840 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1841 continue;
1842 }
1843 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1844 Some(s) => s,
1845 None => continue,
1846 };
1847 let key = match u64::from_str_radix(stem, 16) {
1848 Ok(k) => k,
1849 Err(_) => continue,
1850 };
1851 let bytes = match std::fs::read(&path) {
1852 Ok(b) => b,
1853 Err(_) => continue,
1854 };
1855 let plaintext = if let Some(dek) = &self.cache_dek {
1857 match self.decrypt_cache(&bytes, dek) {
1858 Some(p) => p,
1859 None => {
1860 let _ = std::fs::remove_file(&path);
1861 continue;
1862 }
1863 }
1864 } else {
1865 bytes
1866 };
1867 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1868 Ok(serialized) => {
1869 if let Some(entry) = serialized.into_entry() {
1870 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1871 self.index_entry(key, &entry);
1872 self.entries.insert(key, entry);
1873 self.touch(key);
1874 } else {
1875 let _ = std::fs::remove_file(&path);
1876 }
1877 }
1878 Err(_) => {
1879 let _ = std::fs::remove_file(&path);
1880 }
1881 }
1882 }
1883 self.evict();
1884 }
1885
1886 fn set_max_bytes(&mut self, max_bytes: u64) {
1887 self.max_bytes = max_bytes;
1888 self.evict();
1889 }
1890
1891 fn touch(&mut self, key: u64) {
1894 if !self.entries.contains_key(&key) {
1895 return;
1896 }
1897 if let Some(previous) = self.generations.remove(&key) {
1898 self.order.remove(&(previous, key));
1899 }
1900 let generation = self.allocate_generation();
1901 self.generations.insert(key, generation);
1902 self.order.insert((generation, key));
1903 }
1904
1905 fn untrack(&mut self, key: u64) {
1906 if let Some(generation) = self.generations.remove(&key) {
1907 self.order.remove(&(generation, key));
1908 }
1909 }
1910
1911 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
1912 for column in &entry.condition_cols {
1913 self.condition_index.entry(*column).or_default().insert(key);
1914 }
1915 if entry.footprint.is_empty() {
1916 self.empty_footprint_entries.insert(key);
1917 }
1918 }
1919
1920 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
1921 for column in &entry.condition_cols {
1922 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
1923 keys.remove(&key);
1924 keys.is_empty()
1925 });
1926 if remove_column {
1927 self.condition_index.remove(column);
1928 }
1929 }
1930 if entry.footprint.is_empty() {
1931 self.empty_footprint_entries.remove(&key);
1932 }
1933 }
1934
1935 fn allocate_generation(&mut self) -> u64 {
1936 if self.next_generation == u64::MAX {
1937 self.rebase_generations();
1938 }
1939 let generation = self.next_generation;
1940 self.next_generation += 1;
1941 generation
1942 }
1943
1944 fn rebase_generations(&mut self) {
1945 let ordered = std::mem::take(&mut self.order);
1946 self.generations.clear();
1947 for (generation, (_, key)) in ordered.into_iter().enumerate() {
1948 let generation = generation as u64;
1949 self.generations.insert(key, generation);
1950 self.order.insert((generation, key));
1951 }
1952 self.next_generation = self.order.len() as u64;
1953 }
1954
1955 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
1956 let res = self.entries.get(&key).and_then(|e| match &e.data {
1957 CachedData::Rows(r) => Some(r.clone()),
1958 CachedData::Columns(_) => None,
1959 });
1960 if res.is_some() {
1961 self.touch(key);
1962 self.memory_hit
1963 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1964 return res;
1965 }
1966 if let Some(entry) = self.load_from_disk(key) {
1968 let res = match &entry.data {
1969 CachedData::Rows(r) => Some(r.clone()),
1970 CachedData::Columns(_) => None,
1971 };
1972 if res.is_some() {
1973 let approx = entry.data.approx_bytes();
1974 self.bytes = self.bytes.saturating_add(approx);
1975 self.index_entry(key, &entry);
1976 self.entries.insert(key, entry);
1977 self.touch(key);
1978 self.evict();
1979 self.disk_hit
1980 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1981 return res;
1982 }
1983 }
1984 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1985 None
1986 }
1987
1988 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
1989 let res = self.entries.get(&key).and_then(|e| match &e.data {
1990 CachedData::Columns(c) => Some(c.clone()),
1991 CachedData::Rows(_) => None,
1992 });
1993 if res.is_some() {
1994 self.touch(key);
1995 self.memory_hit
1996 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1997 return res;
1998 }
1999 if let Some(entry) = self.load_from_disk(key) {
2001 let res = match &entry.data {
2002 CachedData::Columns(c) => Some(c.clone()),
2003 CachedData::Rows(_) => None,
2004 };
2005 if res.is_some() {
2006 let approx = entry.data.approx_bytes();
2007 self.bytes = self.bytes.saturating_add(approx);
2008 self.index_entry(key, &entry);
2009 self.entries.insert(key, entry);
2010 self.touch(key);
2011 self.evict();
2012 self.disk_hit
2013 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2014 return res;
2015 }
2016 }
2017 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2018 None
2019 }
2020
2021 fn insert(&mut self, key: u64, entry: CachedEntry) {
2022 let approx = entry.data.approx_bytes();
2023 if let Some(previous) = self.entries.remove(&key) {
2024 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
2025 self.unindex_entry(key, &previous);
2026 self.untrack(key);
2027 }
2028 if self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes) {
2032 let write_start = std::time::Instant::now();
2033 self.store_to_disk(key, &entry);
2034 let write_us = write_start.elapsed().as_micros() as u64;
2035 self.persistent_write_us
2036 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2037 }
2038 self.bytes = self.bytes.saturating_add(approx);
2039 self.index_entry(key, &entry);
2040 self.entries.insert(key, entry);
2041 self.touch(key);
2042 self.evict();
2043 }
2044
2045 #[cfg(test)]
2046 fn set_persist_min_bytes(&mut self, min: u64) {
2047 self.persist_min_bytes = min;
2048 }
2049
2050 #[cfg(test)]
2052 fn would_persist(&self, approx: u64) -> bool {
2053 self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
2054 }
2055
2056 fn cache_counters(&self) -> (u64, u64, u64, u64) {
2059 (
2060 self.memory_hit.load(std::sync::atomic::Ordering::Relaxed),
2061 self.disk_hit.load(std::sync::atomic::Ordering::Relaxed),
2062 self.miss.load(std::sync::atomic::Ordering::Relaxed),
2063 self.persistent_write_us
2064 .load(std::sync::atomic::Ordering::Relaxed),
2065 )
2066 }
2067
2068 fn invalidate(
2077 &mut self,
2078 delete_rids: &roaring::RoaringBitmap,
2079 put_cols: &std::collections::HashSet<u16>,
2080 ) {
2081 if self.entries.is_empty() {
2082 return;
2083 }
2084 let has_deletes = !delete_rids.is_empty();
2085 let mut to_remove = HashSet::new();
2086
2087 for column in put_cols {
2091 if let Some(keys) = self.condition_index.get(column) {
2092 to_remove.extend(keys.iter().copied());
2093 }
2094 }
2095
2096 if has_deletes {
2097 to_remove.extend(self.empty_footprint_entries.iter().copied());
2101 for (&key, entry) in &self.entries {
2102 if !to_remove.contains(&key)
2103 && !entry.footprint.is_empty()
2104 && entry.footprint.intersection_len(delete_rids) > 0
2105 {
2106 to_remove.insert(key);
2107 }
2108 }
2109 }
2110
2111 for key in to_remove {
2112 if let Some(entry) = self.entries.remove(&key) {
2113 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2114 self.unindex_entry(key, &entry);
2115 }
2116 self.remove_from_disk(key);
2117 self.untrack(key);
2118 }
2119 }
2120
2121 fn clear(&mut self) {
2122 if let Some(dir) = &self.dir {
2124 if let Ok(entries) = std::fs::read_dir(dir) {
2125 for entry in entries.flatten() {
2126 let path = entry.path();
2127 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
2128 let _ = std::fs::remove_file(&path);
2129 }
2130 }
2131 }
2132 }
2133 self.entries.clear();
2134 self.order.clear();
2135 self.generations.clear();
2136 self.next_generation = 0;
2137 self.condition_index.clear();
2138 self.empty_footprint_entries.clear();
2139 self.bytes = 0;
2140 }
2141
2142 fn evict(&mut self) {
2143 while self.bytes > self.max_bytes {
2144 let Some((_, key)) = self.order.pop_first() else {
2145 break;
2146 };
2147 self.generations.remove(&key);
2148 if let Some(entry) = self.entries.remove(&key) {
2149 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2150 self.unindex_entry(key, &entry);
2151 self.remove_from_disk(key);
2155 }
2156 }
2157 }
2158}
2159
2160#[cfg(test)]
2161mod result_cache_lru_tests {
2162 use super::*;
2163
2164 fn row_entry(row_id: u64) -> CachedEntry {
2165 CachedEntry {
2166 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
2167 footprint: std::iter::once(row_id as u32).collect(),
2168 condition_cols: vec![1],
2169 }
2170 }
2171
2172 #[test]
2173 fn hits_update_recency_without_growing_an_order_queue() {
2174 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2175 cache.insert(1, row_entry(1));
2176 cache.insert(2, row_entry(2));
2177 assert_eq!(cache.order.len(), cache.entries.len());
2178
2179 for _ in 0..1_000 {
2180 assert!(cache.get_rows(1).is_some());
2181 }
2182
2183 assert_eq!(cache.order.len(), cache.entries.len());
2184 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
2185
2186 let one_entry = cache.entries[&1].data.approx_bytes();
2187 cache.set_max_bytes(one_entry);
2188 assert!(cache.entries.contains_key(&1));
2189 assert!(!cache.entries.contains_key(&2));
2190 assert_eq!(cache.order.len(), cache.entries.len());
2191 }
2192
2193 #[test]
2194 fn tiny_entries_skip_persistent_tier_by_default() {
2195 let dir = tempfile::tempdir().unwrap();
2196 let mut cache = ResultCache::with_max_bytes(u64::MAX).with_dir(dir.path().to_path_buf());
2197 let tiny = row_entry(1).data.approx_bytes();
2198 assert!(
2199 tiny < 4 * 1024,
2200 "row_entry fixture must be below default 4KiB threshold"
2201 );
2202 assert!(
2203 !cache.would_persist(tiny),
2204 "tiny entry must not force synchronous disk publish"
2205 );
2206 assert!(
2207 cache.would_persist(8 * 1024),
2208 "large entry must still persist when dir is set"
2209 );
2210 cache.set_persist_min_bytes(0);
2211 assert!(
2212 cache.would_persist(tiny),
2213 "persist_min_bytes=0 restores always-persist policy"
2214 );
2215 }
2216
2217 #[test]
2218 fn insert_invalidation_uses_the_condition_reverse_index() {
2219 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2220 let mut first = row_entry(1);
2221 first.condition_cols = vec![1];
2222 let mut second = row_entry(2);
2223 second.condition_cols = vec![2];
2224 cache.insert(1, first);
2225 cache.insert(2, second);
2226
2227 let put_cols = std::iter::once(1_u16).collect();
2228 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
2229
2230 assert!(!cache.entries.contains_key(&1));
2231 assert!(cache.entries.contains_key(&2));
2232 assert!(!cache
2233 .condition_index
2234 .get(&1)
2235 .is_some_and(|keys| keys.contains(&1)));
2236 assert!(cache
2237 .condition_index
2238 .get(&2)
2239 .is_some_and(|keys| keys.contains(&2)));
2240 assert_eq!(cache.order.len(), cache.entries.len());
2241 }
2242
2243 #[test]
2244 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
2245 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2246 cache.insert(1, row_entry(1));
2247 cache.insert(2, row_entry(2));
2248 let mut unknown = row_entry(3);
2249 unknown.footprint.clear();
2250 cache.insert(3, unknown);
2251
2252 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
2253 cache.invalidate(&delete_rids, &HashSet::new());
2254
2255 assert!(!cache.entries.contains_key(&1));
2256 assert!(cache.entries.contains_key(&2));
2257 assert!(!cache.entries.contains_key(&3));
2258 assert!(cache.empty_footprint_entries.is_empty());
2259 assert_eq!(cache.order.len(), cache.entries.len());
2260 }
2261
2262 #[test]
2263 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
2264 let entry = row_entry(7);
2265 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
2266 let owned = bincode::serialize(&SerializedEntry {
2267 condition_cols: entry.condition_cols.clone(),
2268 footprint_bits: entry.footprint.iter().collect(),
2269 data: match &entry.data {
2270 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
2271 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
2272 },
2273 })
2274 .unwrap();
2275 assert_eq!(borrowed, owned);
2276 }
2277}
2278
2279type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
2286
2287fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
2288 let _ = kek;
2289 {
2290 if let Some(k) = kek {
2291 return (
2292 Some(k.derive_table_wal_key(_table_id)),
2293 Some(k.derive_cache_key()),
2294 );
2295 }
2296 }
2297 (None, None)
2298}
2299
2300fn read_table_encryption_salt_root(
2301 root: &crate::durable_file::DurableRoot,
2302) -> Result<[u8; crate::encryption::SALT_LEN]> {
2303 use std::io::Read;
2304
2305 let mut file = root
2306 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
2307 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
2308 let length = file.metadata()?.len();
2309 if length != crate::encryption::SALT_LEN as u64 {
2310 return Err(MongrelError::InvalidArgument(format!(
2311 "salt file is {length} bytes, expected {}",
2312 crate::encryption::SALT_LEN
2313 )));
2314 }
2315 let mut salt = [0_u8; crate::encryption::SALT_LEN];
2316 file.read_exact(&mut salt)?;
2317 Ok(salt)
2318}
2319
2320fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
2322 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
2323}
2324
2325fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
2326 let Some(kek) = kek else {
2327 return HashMap::new();
2328 };
2329 {
2330 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
2331 schema
2332 .columns
2333 .iter()
2334 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
2335 .map(|c| {
2336 let scheme = if schema
2337 .indexes
2338 .iter()
2339 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
2340 {
2341 SCHEME_OPE_RANGE
2342 } else {
2343 SCHEME_HMAC_EQ
2344 };
2345 let key: [u8; 32] = *kek.derive_column_key(c.id);
2346 (c.id, (key, scheme))
2347 })
2348 .collect()
2349 }
2350}
2351
2352pub(crate) struct SharedCtx {
2357 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
2358 pub epoch: Arc<EpochAuthority>,
2359 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
2360 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
2361 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
2362 pub kek: Option<Arc<Kek>>,
2363 pub commit_lock: Arc<parking_lot::Mutex<()>>,
2369 pub shared: Option<SharedWalCtx>,
2373 pub table_name: Option<String>,
2376 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
2379 pub read_only: bool,
2381}
2382
2383#[derive(Clone)]
2389pub(crate) struct SharedWalCtx {
2390 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
2391 pub group: Arc<GroupCommit>,
2392 pub poisoned: Arc<AtomicBool>,
2393 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
2394 pub change_wake: tokio::sync::broadcast::Sender<()>,
2395 pub lifecycle: Arc<crate::core::LifecycleController>,
2398 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
2400}
2401
2402enum WalSink {
2405 Private(Wal),
2406 Shared(SharedWalCtx),
2407 ReadOnly,
2408}
2409
2410impl Clone for WalSink {
2411 fn clone(&self) -> Self {
2412 match self {
2413 Self::Shared(shared) => Self::Shared(shared.clone()),
2414 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
2415 }
2416 }
2417}
2418
2419impl SharedCtx {
2420 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
2424 let n_shards = if cache_dir.is_some() {
2428 1
2429 } else {
2430 crate::cache::CACHE_SHARDS
2431 };
2432 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
2433 let page_cache = if let Some(d) = cache_dir {
2434 Arc::new(crate::cache::Sharded::new(1, || {
2435 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
2436 }))
2437 } else {
2438 Arc::new(crate::cache::Sharded::new(n_shards, || {
2439 crate::cache::PageCache::new(per_shard)
2440 }))
2441 };
2442 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
2443 let decoded_cache = Arc::new(crate::cache::Sharded::new(
2444 crate::cache::CACHE_SHARDS,
2445 || crate::cache::DecodedPageCache::new(decoded_per_shard),
2446 ));
2447 Self {
2448 root_guard: None,
2449 epoch: Arc::new(EpochAuthority::new(0)),
2450 page_cache,
2451 decoded_cache,
2452 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
2453 kek,
2454 commit_lock: Arc::new(parking_lot::Mutex::new(())),
2455 shared: None,
2456 table_name: None,
2457 auth: None,
2458 read_only: false,
2459 }
2460 }
2461}
2462
2463fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
2467 use crate::query::Condition;
2468 match c {
2469 Condition::Pk(_)
2471 | Condition::BitmapEq { .. }
2472 | Condition::BitmapIn { .. }
2473 | Condition::BytesPrefix { .. }
2474 | Condition::IsNull { .. }
2475 | Condition::IsNotNull { .. } => 0,
2476 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
2478 1
2479 }
2480 Condition::FmContains { .. }
2482 | Condition::FmContainsAll { .. }
2483 | Condition::Ann { .. }
2484 | Condition::SparseMatch { .. } => 2,
2485 }
2486}
2487
2488impl Table {
2489 pub(crate) fn build_secondary_index_artifact<F>(
2495 &self,
2496 definition: &IndexDef,
2497 rows: &[Row],
2498 mut checkpoint: F,
2499 ) -> Result<SecondaryIndexArtifact>
2500 where
2501 F: FnMut(usize, usize) -> Result<()>,
2502 {
2503 let mut build_schema = self.schema.clone();
2504 build_schema.indexes = vec![definition.clone()];
2505 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
2506 let name_to_id: HashMap<&str, u16> = self
2507 .schema
2508 .columns
2509 .iter()
2510 .map(|column| (column.name.as_str(), column.id))
2511 .collect();
2512 let mut hot = HotIndex::new();
2513 let mut accepted = Vec::with_capacity(rows.len());
2514
2515 for (offset, row) in rows.iter().enumerate() {
2516 checkpoint(offset, rows.len())?;
2517 if row.deleted {
2518 continue;
2519 }
2520 if let Some(predicate) = &definition.predicate {
2521 let columns: HashMap<u16, &Value> =
2522 row.columns.iter().map(|(id, value)| (*id, value)).collect();
2523 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
2524 continue;
2525 }
2526 }
2527 accepted.push(row);
2528 if definition.kind == IndexKind::LearnedRange {
2529 continue;
2530 }
2531 let effective = if self.column_keys.is_empty() {
2532 row.clone()
2533 } else {
2534 self.tokenized_for_indexes(row)
2535 };
2536 if definition.kind == IndexKind::Ann {
2537 if let (Some(index), Some(vector)) = (
2538 ann.get_mut(&definition.column_id),
2539 effective
2540 .columns
2541 .get(&definition.column_id)
2542 .and_then(Value::as_embedding),
2543 ) {
2544 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
2545 checkpoint(offset, rows.len())
2546 })?;
2547 }
2548 } else {
2549 index_into_single(
2550 definition,
2551 &build_schema,
2552 &effective,
2553 &mut hot,
2554 &mut bitmap,
2555 &mut ann,
2556 &mut fm,
2557 &mut sparse,
2558 &mut minhash,
2559 );
2560 }
2561 }
2562 checkpoint(rows.len(), rows.len())?;
2563
2564 if let Some(index) = ann.get_mut(&definition.column_id) {
2565 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
2566 }
2567
2568 let column_id = definition.column_id;
2569 match definition.kind {
2570 IndexKind::Bitmap => bitmap
2571 .remove(&column_id)
2572 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
2573 IndexKind::Ann => ann
2574 .remove(&column_id)
2575 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
2576 IndexKind::FmIndex => fm
2577 .remove(&column_id)
2578 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
2579 IndexKind::Sparse => sparse
2580 .remove(&column_id)
2581 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
2582 IndexKind::MinHash => minhash
2583 .remove(&column_id)
2584 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
2585 IndexKind::LearnedRange => {
2586 let epsilon = definition
2587 .options
2588 .learned_range
2589 .as_ref()
2590 .map(|options| options.epsilon)
2591 .unwrap_or(16);
2592 let column = self
2593 .schema
2594 .columns
2595 .iter()
2596 .find(|column| column.id == column_id)
2597 .ok_or_else(|| {
2598 MongrelError::Schema(format!(
2599 "index {} references unknown column {column_id}",
2600 definition.name
2601 ))
2602 })?;
2603 let index = match column.ty {
2604 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2605 let pairs: Vec<_> = accepted
2606 .iter()
2607 .filter_map(|row| match row.columns.get(&column_id) {
2608 Some(Value::Int64(value)) if !row.deleted => {
2609 Some((*value, row.row_id.0))
2610 }
2611 _ => None,
2612 })
2613 .collect();
2614 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
2615 }
2616 TypeId::Float32 | TypeId::Float64 => {
2617 let pairs: Vec<_> = accepted
2618 .iter()
2619 .filter_map(|row| match row.columns.get(&column_id) {
2620 Some(Value::Float64(value)) if !row.deleted => {
2621 Some((*value, row.row_id.0))
2622 }
2623 _ => None,
2624 })
2625 .collect();
2626 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
2627 }
2628 ref ty => {
2629 return Err(MongrelError::Schema(format!(
2630 "LearnedRange index {} does not support {ty:?}",
2631 definition.name
2632 )));
2633 }
2634 };
2635 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
2636 }
2637 }
2638 .ok_or_else(|| {
2639 MongrelError::Other(format!(
2640 "failed to construct hidden index {}",
2641 definition.name
2642 ))
2643 })
2644 }
2645
2646 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
2647 let dir = dir.as_ref().to_path_buf();
2648 std::fs::create_dir_all(&dir)?;
2651 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2652 let pinned = root.io_path()?;
2653 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2654 ctx.root_guard = Some(root.clone());
2655 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2656 root.finalize_deferred_sync_shallow()?;
2662 Ok(table)
2663 }
2664
2665 pub fn create_encrypted(
2676 dir: impl AsRef<Path>,
2677 schema: Schema,
2678 table_id: u64,
2679 passphrase: &str,
2680 ) -> Result<Self> {
2681 let dir = dir.as_ref().to_path_buf();
2682 std::fs::create_dir_all(&dir)?;
2683 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2684 root.create_directory_all(META_DIR)?;
2685 let salt = crate::encryption::random_salt()?;
2686 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
2687 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2688 let pinned = root.io_path()?;
2689 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2690 ctx.root_guard = Some(root.clone());
2691 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2692 root.finalize_deferred_sync()?;
2693 Ok(table)
2694 }
2695
2696 pub fn create_with_key(
2701 dir: impl AsRef<Path>,
2702 schema: Schema,
2703 table_id: u64,
2704 key: &[u8],
2705 ) -> Result<Self> {
2706 let dir = dir.as_ref().to_path_buf();
2707 std::fs::create_dir_all(&dir)?;
2708 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2709 root.create_directory_all(META_DIR)?;
2710 let salt = crate::encryption::random_salt()?;
2711 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
2712 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
2713 let pinned = root.io_path()?;
2714 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2715 ctx.root_guard = Some(root.clone());
2716 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2717 root.finalize_deferred_sync()?;
2718 Ok(table)
2719 }
2720
2721 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
2723 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2724 let salt = read_table_encryption_salt_root(&root)?;
2725 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
2726 let pinned = root.io_path()?;
2727 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2728 ctx.root_guard = Some(root);
2729 Self::open_in(&pinned, ctx)
2730 }
2731
2732 pub(crate) fn create_in(
2733 dir: impl AsRef<Path>,
2734 schema: Schema,
2735 table_id: u64,
2736 ctx: SharedCtx,
2737 ) -> Result<Self> {
2738 schema.validate_auto_increment()?;
2739 schema.validate_defaults()?;
2740 schema.validate_ai()?;
2741 for index in &schema.indexes {
2742 index.validate_options()?;
2743 }
2744 let dir = dir.as_ref().to_path_buf();
2745 let runs_root = match ctx.root_guard.as_ref() {
2746 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
2747 None => {
2748 crate::durable_file::create_directory_all(&dir)?;
2749 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
2750 None
2751 }
2752 };
2753 match ctx.root_guard.as_deref() {
2754 Some(root) => write_schema_durable(root, &schema)?,
2755 None => write_schema(&dir, &schema)?,
2756 }
2757 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
2758 let (wal, current_txn_id) = match ctx.shared.clone() {
2761 Some(s) => (WalSink::Shared(s), 0),
2762 None => {
2763 let pinned_wal_root = match ctx.root_guard.as_deref() {
2764 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
2765 None => None,
2766 };
2767 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
2768 root.io_path()?
2769 } else {
2770 let wal_dir = dir.join(WAL_DIR);
2771 std::fs::create_dir_all(&wal_dir)?;
2772 wal_dir
2773 };
2774 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
2775 (Some(root), Some(dk)) => {
2776 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
2777 }
2778 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
2779 (None, Some(dk)) => Wal::create_with_cipher(
2780 wal_dir.join("seg-000000.wal"),
2781 Epoch(0),
2782 Some(make_cipher(dk)),
2783 0,
2784 )?,
2785 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
2786 };
2787 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
2788 (WalSink::Private(w), 1)
2789 }
2790 };
2791 let mut manifest = Manifest::new(table_id, schema.schema_id);
2792 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2797 match ctx.root_guard.as_deref() {
2798 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
2799 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
2800 }
2801 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
2802 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
2803 let auto_inc = resolve_auto_inc(&schema);
2804 let rcache_dir = dir.join(RCACHE_DIR);
2805 let initial_view = ReadGeneration::empty(&schema);
2806 Ok(Self {
2807 dir,
2808 _root_guard: ctx.root_guard,
2809 runs_root,
2810 idx_root: None,
2811 table_id,
2812 name: ctx.table_name.unwrap_or_default(),
2813 auth: ctx.auth,
2814 read_only: ctx.read_only,
2815 durable_commit_failed: false,
2816 wal,
2817 memtable: Memtable::new(),
2818 mutable_run: MutableRun::new(),
2819 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
2820 compaction_zstd_level: 3,
2821 allocator: RowIdAllocator::new(0),
2822 epoch: ctx.epoch,
2823 data_generation: 0,
2824 schema,
2825 hot: HotIndex::new(),
2826 kek: ctx.kek,
2827 column_keys,
2828 run_refs: Vec::new(),
2829 retiring: Vec::new(),
2830 next_run_id: 1,
2831 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
2832 current_txn_id,
2833 pending_private_mutations: false,
2834 bitmap,
2835 ann,
2836 fm,
2837 sparse,
2838 minhash,
2839 learned_range: Arc::new(HashMap::new()),
2840 pk_by_row: ReversePkMap::new(),
2841 pinned: BTreeMap::new(),
2842 live_count: 0,
2843 reservoir: crate::reservoir::Reservoir::default(),
2844 reservoir_complete: true,
2845 had_deletes: false,
2846 recent_delete_preimages: HashMap::new(),
2847 run_row_id_ranges: HashMap::new(),
2848 agg_cache: Arc::new(HashMap::new()),
2849 global_idx_epoch: 0,
2850 indexes_complete: true,
2851 index_build_policy: IndexBuildPolicy::default(),
2852 pk_by_row_complete: false,
2853 flushed_epoch: 0,
2854 page_cache: ctx.page_cache,
2855 decoded_cache: ctx.decoded_cache,
2856 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
2857 snapshots: ctx.snapshots,
2858 commit_lock: ctx.commit_lock,
2859 result_cache: Arc::new(parking_lot::Mutex::new(
2860 ResultCache::new()
2861 .with_dir(rcache_dir)
2862 .with_cache_dek(cache_dek.clone()),
2863 )),
2864 pending_delete_rids: roaring::RoaringBitmap::new(),
2865 pending_put_cols: std::collections::HashSet::new(),
2866 pending_rows: Vec::new(),
2867 pending_rows_auto_inc: Vec::new(),
2868 pending_dels: Vec::new(),
2869 pending_truncate: None,
2870 wal_dek,
2871 auto_inc,
2872 ttl: None,
2873 pins: Arc::new(crate::retention::PinRegistry::new()),
2874 published: Arc::new(ArcSwap::from_pointee(initial_view)),
2875 read_generation_pin: None,
2876 lookup_metrics: LookupMetrics::default(),
2877 })
2878 }
2879
2880 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
2884 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2885 let pinned = root.io_path()?;
2886 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2887 ctx.root_guard = Some(root);
2888 Self::open_in(&pinned, ctx)
2889 }
2890
2891 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
2894 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
2895 let salt = read_table_encryption_salt_root(&root)?;
2896 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2897 let pinned = root.io_path()?;
2898 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2899 ctx.root_guard = Some(root);
2900 let t = Self::open_in(&pinned, ctx)?;
2901 Ok(t)
2902 }
2903
2904 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
2905 let dir = dir.as_ref().to_path_buf();
2906 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
2907 let mut manifest = match ctx.root_guard.as_ref() {
2908 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
2909 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
2910 };
2911 let schema: Schema = match ctx.root_guard.as_ref() {
2912 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
2913 None => read_schema(&dir)?,
2914 };
2915 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
2921 let runs_root = match ctx.root_guard.as_ref() {
2922 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
2923 None => None,
2924 };
2925 let idx_root = match ctx.root_guard.as_ref() {
2926 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
2927 Ok(root) => Some(Arc::new(root)),
2928 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
2929 Err(error) => return Err(error.into()),
2930 },
2931 None => None,
2932 };
2933 schema.validate_auto_increment()?;
2934 schema.validate_defaults()?;
2935 schema.validate_ai()?;
2936 for index in &schema.indexes {
2937 index.validate_options()?;
2938 }
2939 let replay_epoch = Epoch(manifest.current_epoch);
2940 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
2941 let private_replayed = if ctx.shared.is_none() {
2942 match latest_wal_segment(&dir.join(WAL_DIR))? {
2943 Some(path) => {
2944 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
2945 crate::wal::replay_with_cipher(path, cipher)?
2946 }
2947 None => Vec::new(),
2948 }
2949 } else {
2950 Vec::new()
2951 };
2952 if ctx.shared.is_none() {
2953 preflight_standalone_open(
2954 &dir,
2955 runs_root.as_deref(),
2956 idx_root.as_deref(),
2957 &manifest,
2958 &schema,
2959 &private_replayed,
2960 ctx.kek.clone(),
2961 )?;
2962 }
2963 let next_run_id = derive_next_run_id(
2964 &dir,
2965 runs_root.as_deref(),
2966 &manifest.runs,
2967 &manifest.retiring,
2968 )?;
2969 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
2973 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
2974 None => {
2975 let replayed = private_replayed;
2976 let wal_dir = dir.join(WAL_DIR);
2982 crate::durable_file::create_directory_all(&wal_dir)?;
2983 let segment = next_wal_segment(&wal_dir)?;
2984 let segment_no = wal_segment_number(&segment).unwrap_or(0);
2985 let temporary = wal_dir.join(format!(
2986 ".recovery-{}-{}-{segment_no:06}.tmp",
2987 std::process::id(),
2988 std::time::SystemTime::now()
2989 .duration_since(std::time::UNIX_EPOCH)
2990 .unwrap_or_default()
2991 .as_nanos()
2992 ));
2993 let mut w = Wal::create_with_cipher(
2994 &temporary,
2995 replay_epoch,
2996 wal_dek.as_ref().map(|dk| make_cipher(dk)),
2997 segment_no,
2998 )?;
2999 for record in &replayed {
3000 w.append_txn(record.txn_id, record.op.clone())?;
3001 }
3002 let mut w = w.publish_as(segment)?;
3003 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3004 let next_txn_id = replayed
3005 .iter()
3006 .map(|record| record.txn_id)
3007 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
3008 .max()
3009 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
3010 .unwrap_or(1);
3011 (WalSink::Private(w), replayed, next_txn_id)
3012 }
3013 };
3014
3015 let mut memtable = Memtable::new();
3016 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
3017 let persisted_epoch = manifest.current_epoch;
3018 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
3025 s.next = manifest.auto_inc_next;
3026 s.seeded = manifest.auto_inc_next > 0;
3027 s
3028 });
3029
3030 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
3037 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
3038 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
3039 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
3040 std::collections::BTreeMap::new();
3041 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
3042 let mut recovered_epoch = manifest.current_epoch;
3043 let mut recovered_manifest_dirty = schema_manifest_repair;
3044 let mut saw_delete = false;
3045 for record in replayed {
3046 let txn_id = record.txn_id;
3047 match record.op {
3048 Op::Put { rows, .. } => {
3049 let rows: Vec<Row> = bincode::deserialize(&rows)?;
3050 for row in &rows {
3051 allocator.advance_to(row.row_id)?;
3052 if let Some(ai) = auto_inc.as_mut() {
3053 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3054 let next = n.checked_add(1).ok_or_else(|| {
3055 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3056 })?;
3057 if next > ai.next {
3058 ai.next = next;
3059 }
3060 }
3061 }
3062 }
3063 staged_puts.entry(txn_id).or_default().extend(rows);
3064 }
3065 Op::Delete { row_ids, .. } => {
3066 staged_deletes.entry(txn_id).or_default().extend(row_ids);
3067 }
3068 Op::TxnCommit { epoch, .. } => {
3069 let commit_epoch = Epoch(epoch);
3070 recovered_epoch = recovered_epoch.max(epoch);
3071 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
3072 memtable = Memtable::new();
3073 replayed_puts.clear();
3074 replayed_deletes.clear();
3075 manifest.runs.clear();
3076 manifest.retiring.clear();
3077 manifest.live_count = 0;
3078 manifest.global_idx_epoch = 0;
3079 manifest.current_epoch = manifest.current_epoch.max(epoch);
3080 recovered_manifest_dirty = true;
3081 saw_delete = true;
3082 }
3083 if let Some(puts) = staged_puts.remove(&txn_id) {
3084 if commit_epoch.0 > manifest.flushed_epoch {
3085 for row in &puts {
3086 memtable.upsert(row.clone());
3087 }
3088 replayed_puts.entry(commit_epoch).or_default().extend(puts);
3089 }
3090 }
3091 if let Some(dels) = staged_deletes.remove(&txn_id) {
3092 saw_delete = true;
3093 if commit_epoch.0 > manifest.flushed_epoch {
3094 for rid in dels {
3095 memtable.tombstone(rid, commit_epoch);
3096 replayed_deletes.push((rid, commit_epoch));
3097 }
3098 }
3099 }
3100 }
3101 Op::TxnAbort => {
3102 staged_puts.remove(&txn_id);
3103 staged_deletes.remove(&txn_id);
3104 staged_truncates.remove(&txn_id);
3105 }
3106 Op::TruncateTable { .. } => {
3107 staged_puts.remove(&txn_id);
3108 staged_deletes.remove(&txn_id);
3109 staged_truncates.insert(txn_id);
3110 }
3111 Op::ExternalTableState { .. }
3112 | Op::Flush { .. }
3113 | Op::Ddl(_)
3114 | Op::BeforeImage { .. }
3115 | Op::CommitTimestamp { .. }
3116 | Op::SpilledRows { .. } => {}
3117 }
3118 }
3119
3120 let rcache_dir = dir.join(RCACHE_DIR);
3121 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3122 let initial_view = ReadGeneration::empty(&schema);
3123 let mut db = Self {
3124 dir,
3125 _root_guard: ctx.root_guard,
3126 runs_root,
3127 idx_root,
3128 table_id: manifest.table_id,
3129 name: ctx.table_name.unwrap_or_default(),
3130 auth: ctx.auth,
3131 read_only: ctx.read_only,
3132 durable_commit_failed: false,
3133 wal,
3134 memtable,
3135 mutable_run: MutableRun::new(),
3136 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3137 compaction_zstd_level: 3,
3138 allocator,
3139 epoch: ctx.epoch,
3140 data_generation: persisted_epoch,
3141 schema,
3142 hot: HotIndex::new(),
3143 kek: ctx.kek,
3144 column_keys,
3145 run_refs: manifest.runs.clone(),
3146 retiring: manifest.retiring.clone(),
3147 next_run_id,
3148 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3149 current_txn_id,
3150 pending_private_mutations: false,
3151 bitmap: HashMap::new(),
3152 ann: HashMap::new(),
3153 fm: HashMap::new(),
3154 sparse: HashMap::new(),
3155 minhash: HashMap::new(),
3156 learned_range: Arc::new(HashMap::new()),
3157 pk_by_row: ReversePkMap::new(),
3158 pinned: BTreeMap::new(),
3159 live_count: manifest.live_count,
3160 reservoir: crate::reservoir::Reservoir::default(),
3161 reservoir_complete: false,
3162 had_deletes: saw_delete
3163 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
3164 != manifest.live_count,
3165 recent_delete_preimages: HashMap::new(),
3166 run_row_id_ranges: HashMap::new(),
3167 agg_cache: Arc::new(HashMap::new()),
3168 global_idx_epoch: manifest.global_idx_epoch,
3169 indexes_complete: true,
3170 index_build_policy: IndexBuildPolicy::default(),
3171 pk_by_row_complete: false,
3172 flushed_epoch: manifest.flushed_epoch,
3173 page_cache: ctx.page_cache,
3174 decoded_cache: ctx.decoded_cache,
3175 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3176 snapshots: ctx.snapshots,
3177 commit_lock: ctx.commit_lock,
3178 result_cache: Arc::new(parking_lot::Mutex::new(
3179 ResultCache::new()
3180 .with_dir(rcache_dir)
3181 .with_cache_dek(cache_dek.clone()),
3182 )),
3183 pending_delete_rids: roaring::RoaringBitmap::new(),
3184 pending_put_cols: std::collections::HashSet::new(),
3185 pending_rows: Vec::new(),
3186 pending_rows_auto_inc: Vec::new(),
3187 pending_dels: Vec::new(),
3188 pending_truncate: None,
3189 wal_dek,
3190 auto_inc,
3191 ttl: manifest.ttl,
3192 pins: Arc::new(crate::retention::PinRegistry::new()),
3193 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3194 read_generation_pin: None,
3195 lookup_metrics: LookupMetrics::default(),
3196 };
3197
3198 db.epoch.advance_recovered(Epoch(recovered_epoch));
3201
3202 let checkpoint = match db.idx_root.as_deref() {
3207 Some(root) => {
3208 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
3209 }
3210 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
3211 };
3212 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
3213 c.epoch_built == manifest.global_idx_epoch
3214 && manifest.global_idx_epoch > 0
3215 && manifest
3216 .runs
3217 .iter()
3218 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
3219 });
3220 if let Some(loaded) = checkpoint {
3221 if checkpoint_valid {
3222 db.hot = loaded.hot;
3223 db.bitmap = loaded.bitmap;
3224 db.ann = loaded.ann;
3225 db.fm = loaded.fm;
3226 db.sparse = loaded.sparse;
3227 db.minhash = loaded.minhash;
3228 db.learned_range = Arc::new(loaded.learned_range);
3229 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
3234 for (cid, idx) in bitmap0 {
3235 db.bitmap.entry(cid).or_insert(idx);
3236 }
3237 for (cid, idx) in ann0 {
3238 db.ann.entry(cid).or_insert(idx);
3239 }
3240 for (cid, idx) in fm0 {
3241 db.fm.entry(cid).or_insert(idx);
3242 }
3243 for (cid, idx) in sparse0 {
3244 db.sparse.entry(cid).or_insert(idx);
3245 }
3246 for (cid, idx) in minhash0 {
3247 db.minhash.entry(cid).or_insert(idx);
3248 }
3249 }
3252 }
3253 if !checkpoint_valid {
3254 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
3255 db.bitmap = bitmap;
3256 db.ann = ann;
3257 db.fm = fm;
3258 db.sparse = sparse;
3259 db.minhash = minhash;
3260 db.rebuild_indexes_from_runs()?;
3261 db.build_learned_ranges()?;
3262 }
3263
3264 for (epoch, group) in replayed_puts {
3269 let (losers, winner_pks) = db.partition_pk_winners(&group);
3270 for (key, &row_id) in &winner_pks {
3271 if let Some(old_rid) = db.hot.get(key) {
3272 if old_rid != row_id {
3273 db.tombstone_row(old_rid, epoch, None, false);
3274 }
3275 }
3276 }
3277 for &loser_rid in &losers {
3278 db.tombstone_row(loser_rid, epoch, None, false);
3279 }
3280 for (key, row_id) in winner_pks {
3281 db.insert_hot_pk(key, row_id);
3282 }
3283 if db.schema.primary_key().is_none() {
3284 for r in &group {
3285 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3286 }
3287 }
3288 for r in &group {
3289 if !losers.contains(&r.row_id) {
3290 db.index_row(r);
3291 }
3292 }
3293 }
3294 for (rid, epoch) in &replayed_deletes {
3298 db.remove_hot_for_row(*rid, *epoch);
3299 }
3300
3301 if recovered_manifest_dirty {
3302 let rows = db.visible_rows(Snapshot::unbounded())?;
3303 db.live_count = rows.len() as u64;
3304 db.persist_manifest(Epoch(recovered_epoch))?;
3305 }
3306
3307 db.result_cache.lock().load_persistent();
3314 db.refresh_run_row_id_ranges();
3316 Ok(db)
3317 }
3318
3319 fn ensure_reservoir_complete(&mut self) -> Result<()> {
3325 if self.reservoir_complete {
3326 return Ok(());
3327 }
3328 self.rebuild_reservoir()?;
3329 self.reservoir_complete = true;
3330 Ok(())
3331 }
3332
3333 fn rebuild_reservoir(&mut self) -> Result<()> {
3336 let snap = self.snapshot();
3337 let rows = self.visible_rows(snap)?;
3338 self.reservoir.reset();
3339 for r in rows {
3340 self.reservoir.offer(r.row_id.0);
3341 }
3342 Ok(())
3343 }
3344
3345 pub fn rebuild_indexes(&mut self) -> Result<()> {
3349 self.rebuild_indexes_from_runs_inner(None)
3350 }
3351
3352 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
3353 self.rebuild_indexes_from_runs_inner(None)
3354 }
3355
3356 fn pk_equality_fallback(
3369 &self,
3370 pk_column_id: u16,
3371 lookup: &[u8],
3372 snapshot: Snapshot,
3373 ) -> Result<(RowIdSet, crate::trace::HotFallbackReason)> {
3374 let mut tombstone_hit = false;
3375 let mut overlay_versions = 0u64;
3376 for row in self.memtable.visible_versions_at(snapshot) {
3378 overlay_versions += 1;
3379 self.lookup_metrics
3380 .hot_lookup_fallback_overlay_rows
3381 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3382 if row.deleted {
3383 tombstone_hit = true;
3384 continue;
3385 }
3386 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3387 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3388 self.lookup_metrics
3389 .hot_fallback_overlay_versions_total
3390 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3391 return Ok((
3392 RowIdSet::one(row.row_id.0),
3393 crate::trace::HotFallbackReason::MissingMapping,
3394 ));
3395 }
3396 }
3397 }
3398 for row in self.mutable_run.visible_versions_at(snapshot) {
3399 overlay_versions += 1;
3400 self.lookup_metrics
3401 .hot_lookup_fallback_overlay_rows
3402 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3403 if row.deleted {
3404 tombstone_hit = true;
3405 continue;
3406 }
3407 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3408 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3409 self.lookup_metrics
3410 .hot_fallback_overlay_versions_total
3411 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3412 return Ok((
3413 RowIdSet::one(row.row_id.0),
3414 crate::trace::HotFallbackReason::MissingMapping,
3415 ));
3416 }
3417 }
3418 }
3419 if lookup.len() == 8 {
3422 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
3423 let n = i64::from_be_bytes(arr);
3424 let result = self.range_scan_i64(pk_column_id, n, n, snapshot)?;
3425 self.lookup_metrics
3426 .hot_lookup_fallback_runs
3427 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3428 self.lookup_metrics
3432 .hot_fallback_runs_considered_total
3433 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3434 self.lookup_metrics
3435 .hot_fallback_overlay_versions_total
3436 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3437 let reason = if result.is_empty() {
3438 if tombstone_hit {
3439 crate::trace::HotFallbackReason::Tombstone
3440 } else {
3441 crate::trace::HotFallbackReason::MissingMapping
3442 }
3443 } else {
3444 crate::trace::HotFallbackReason::MissingMapping
3445 };
3446 return Ok((result, reason));
3447 }
3448 }
3449 let mut found = Vec::new();
3452 let overlay = self.overlay_rid_set(snapshot);
3453 let mut runs_considered = 0u64;
3454 for rr in &self.run_refs {
3455 runs_considered += 1;
3456 self.lookup_metrics
3457 .hot_lookup_fallback_runs
3458 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3459 let mut reader = self.open_reader(rr.run_id)?;
3460 for row in reader.visible_rows(snapshot.epoch)? {
3461 if overlay.contains(&row.row_id.0) || row.deleted {
3462 if row.deleted {
3463 tombstone_hit = true;
3464 }
3465 continue;
3466 }
3467 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3468 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3469 found.push(row.row_id.0);
3470 }
3471 }
3472 }
3473 }
3474 let reason = if tombstone_hit {
3475 crate::trace::HotFallbackReason::Tombstone
3476 } else {
3477 crate::trace::HotFallbackReason::MissingMapping
3480 };
3481 self.lookup_metrics
3482 .hot_fallback_overlay_versions_total
3483 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3484 self.lookup_metrics
3485 .hot_fallback_runs_considered_total
3486 .fetch_add(runs_considered, std::sync::atomic::Ordering::Relaxed);
3487 Ok((RowIdSet::from_unsorted(found), reason))
3488 }
3489
3490 fn record_hot_fallback_reason(&self, reason: crate::trace::HotFallbackReason) {
3494 let idx = hot_fallback_reason_index(reason);
3495 self.lookup_metrics.hot_fallback_reasons[idx]
3496 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3497 self.lookup_metrics
3498 .hot_lookup_fallback
3499 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3500 crate::trace::QueryTrace::record(|t| {
3501 t.hot_lookup_attempted = true;
3502 t.hot_lookup_hit = false;
3503 t.hot_fallback_reason = Some(reason.as_str());
3504 });
3505 }
3506
3507 fn rebuild_indexes_from_runs_inner(
3508 &mut self,
3509 control: Option<&crate::ExecutionControl>,
3510 ) -> Result<()> {
3511 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
3514 crate::retention::PinSource::OnlineIndexBuild,
3515 self.current_epoch(),
3516 );
3517 self.hot = HotIndex::new();
3518 self.pk_by_row.clear();
3519 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3520 self.bitmap = bitmap;
3521 self.ann = ann;
3522 self.fm = fm;
3523 self.sparse = sparse;
3524 self.minhash = minhash;
3525 let snapshot = Epoch(u64::MAX);
3526 let ttl_now = unix_nanos_now();
3527 let mut scanned = 0_usize;
3528 for rr in self.run_refs.clone() {
3529 if let Some(control) = control {
3530 control.checkpoint()?;
3531 }
3532 let mut reader = self.open_reader(rr.run_id)?;
3533 for row in reader.visible_rows(snapshot)? {
3534 if scanned.is_multiple_of(256) {
3535 if let Some(control) = control {
3536 control.checkpoint()?;
3537 }
3538 }
3539 scanned += 1;
3540 if self.row_expired_at(&row, ttl_now) {
3541 continue;
3542 }
3543 let tok_row = self.tokenized_for_indexes(&row);
3544 index_into(
3545 &self.schema,
3546 &tok_row,
3547 &mut self.hot,
3548 &mut self.bitmap,
3549 &mut self.ann,
3550 &mut self.fm,
3551 &mut self.sparse,
3552 &mut self.minhash,
3553 );
3554 }
3555 }
3556 for row in self.mutable_run.visible_versions(snapshot) {
3557 if scanned.is_multiple_of(256) {
3558 if let Some(control) = control {
3559 control.checkpoint()?;
3560 }
3561 }
3562 scanned += 1;
3563 if row.deleted {
3564 self.remove_hot_for_row(row.row_id, snapshot);
3565 } else if !self.row_expired_at(&row, ttl_now) {
3566 self.index_row(&row);
3567 }
3568 }
3569 for row in self.memtable.visible_versions(snapshot) {
3570 if scanned.is_multiple_of(256) {
3571 if let Some(control) = control {
3572 control.checkpoint()?;
3573 }
3574 }
3575 scanned += 1;
3576 if row.deleted {
3577 self.remove_hot_for_row(row.row_id, snapshot);
3578 } else if !self.row_expired_at(&row, ttl_now) {
3579 self.index_row(&row);
3580 }
3581 }
3582 let pin_epochs = self.active_pin_epochs_for_rebuild();
3588 if !pin_epochs.is_empty() {
3589 let current_snap = Snapshot::at(Epoch(u64::MAX));
3590 for pin_epoch in pin_epochs {
3591 if let Some(control) = control {
3592 control.checkpoint()?;
3593 }
3594 let pin_snap = Snapshot::at(pin_epoch);
3595 for rr in self.run_refs.clone() {
3596 if let Some(control) = control {
3597 control.checkpoint()?;
3598 }
3599 let mut reader = self.open_reader(rr.run_id)?;
3600 for row in reader.visible_rows(pin_epoch)? {
3601 if row.deleted || self.row_expired_at(&row, ttl_now) {
3602 continue;
3603 }
3604 if self.get(row.row_id, current_snap).is_none() {
3605 self.index_bitmap_membership_only(&row);
3606 }
3607 }
3608 }
3609 for row in self
3610 .mutable_run
3611 .visible_versions_at(pin_snap)
3612 .into_iter()
3613 .chain(self.memtable.visible_versions_at(pin_snap))
3614 {
3615 if row.deleted || self.row_expired_at(&row, ttl_now) {
3616 continue;
3617 }
3618 if self.get(row.row_id, current_snap).is_none() {
3619 self.index_bitmap_membership_only(&row);
3620 }
3621 }
3622 }
3623 }
3624 self.recent_delete_preimages.clear();
3625 self.refresh_pk_by_row_from_hot();
3626 Ok(())
3627 }
3628
3629 fn index_bitmap_membership_only(&mut self, row: &Row) {
3632 if row.deleted {
3633 return;
3634 }
3635 let columns_map: HashMap<u16, &Value> = row.columns.iter().map(|(k, v)| (*k, v)).collect();
3636 let name_to_id: HashMap<&str, u16> = self
3637 .schema
3638 .columns
3639 .iter()
3640 .map(|c| (c.name.as_str(), c.id))
3641 .collect();
3642 for idef in &self.schema.indexes {
3643 if idef.kind != crate::schema::IndexKind::Bitmap {
3644 continue;
3645 }
3646 if let Some(pred) = &idef.predicate {
3647 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
3648 continue;
3649 }
3650 }
3651 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
3652 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3653 b.insert(key, row.row_id);
3654 }
3655 }
3656 }
3657 }
3658
3659 fn refresh_pk_by_row_from_hot(&mut self) {
3660 self.pk_by_row_complete = true;
3661 if self.schema.primary_key().is_none() {
3662 self.pk_by_row.clear();
3663 return;
3664 }
3665 self.pk_by_row = ReversePkMap::from_entries(
3671 self.hot
3672 .entries()
3673 .into_iter()
3674 .map(|(key, row_id)| (row_id, key)),
3675 );
3676 }
3677
3678 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
3679 if self.schema.primary_key().is_some() {
3680 self.pk_by_row.insert(row_id, key.clone());
3681 }
3682 self.hot.insert(key, row_id);
3683 }
3684
3685 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
3689 self.build_learned_ranges_inner(None)
3690 }
3691
3692 fn build_learned_ranges_inner(
3693 &mut self,
3694 control: Option<&crate::ExecutionControl>,
3695 ) -> Result<()> {
3696 self.learned_range = Arc::new(HashMap::new());
3697 if self.run_refs.len() != 1 {
3698 return Ok(());
3699 }
3700 let cols: Vec<(u16, usize)> = self
3701 .schema
3702 .indexes
3703 .iter()
3704 .filter(|i| i.kind == IndexKind::LearnedRange)
3705 .map(|i| {
3706 (
3707 i.column_id,
3708 i.options
3709 .learned_range
3710 .as_ref()
3711 .map(|options| options.epsilon)
3712 .unwrap_or(16),
3713 )
3714 })
3715 .collect();
3716 if cols.is_empty() {
3717 return Ok(());
3718 }
3719 let snapshot_epoch = self.current_epoch();
3726 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3727 let (visible_positions, visible_rids) =
3728 reader.visible_positions_with_rids(snapshot_epoch)?;
3729 let row_ids: Vec<u64> = visible_rids.iter().map(|r| *r as u64).collect();
3730 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
3731 if column_index % 256 == 0 {
3732 if let Some(control) = control {
3733 control.checkpoint()?;
3734 }
3735 }
3736 let ty = self
3737 .schema
3738 .columns
3739 .iter()
3740 .find(|c| c.id == cid)
3741 .map(|c| c.ty.clone())
3742 .unwrap_or(TypeId::Int64);
3743 match ty {
3744 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
3745 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
3746 let pairs: Vec<(i64, u64)> = visible_positions
3747 .iter()
3748 .zip(row_ids.iter())
3749 .map(|(&p, &r)| (data[p], r))
3750 .collect();
3751 Arc::make_mut(&mut self.learned_range).insert(
3752 cid,
3753 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
3754 );
3755 }
3756 }
3757 TypeId::Float64 => {
3758 if let columnar::NativeColumn::Float64 { data, .. } =
3759 reader.column_native(cid)?
3760 {
3761 let pairs: Vec<(f64, u64)> = visible_positions
3762 .iter()
3763 .zip(row_ids.iter())
3764 .map(|(&p, &r)| (data[p], r))
3765 .collect();
3766 Arc::make_mut(&mut self.learned_range).insert(
3767 cid,
3768 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
3769 );
3770 }
3771 }
3772 _ => {}
3773 }
3774 }
3775 Ok(())
3776 }
3777
3778 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
3785 if self.indexes_complete {
3786 crate::trace::QueryTrace::record(|t| {
3787 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
3788 });
3789 return Ok(());
3790 }
3791 crate::trace::QueryTrace::record(|t| {
3792 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
3793 });
3794 self.rebuild_indexes_from_runs()?;
3795 self.build_learned_ranges()?;
3796 self.indexes_complete = true;
3797 let epoch = self.current_epoch();
3798 self.checkpoint_indexes(epoch);
3799 Ok(())
3800 }
3801
3802 #[doc(hidden)]
3805 pub fn ensure_indexes_complete_controlled<F>(
3806 &mut self,
3807 control: &crate::ExecutionControl,
3808 before_publish: F,
3809 ) -> Result<bool>
3810 where
3811 F: FnOnce() -> bool,
3812 {
3813 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
3814 .map(|(changed, _)| changed)
3815 }
3816
3817 #[doc(hidden)]
3820 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
3821 &mut self,
3822 control: &crate::ExecutionControl,
3823 before_publish: F,
3824 ) -> Result<(bool, Option<MaintenanceReceipt>)>
3825 where
3826 F: FnOnce() -> bool,
3827 {
3828 if self.indexes_complete {
3829 crate::trace::QueryTrace::record(|trace| {
3830 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
3831 });
3832 return Ok((false, None));
3833 }
3834 crate::trace::QueryTrace::record(|trace| {
3835 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
3836 });
3837 control.checkpoint()?;
3838 let maintenance_epoch = self.current_epoch();
3839 self.rebuild_indexes_from_runs_inner(Some(control))?;
3840 self.build_learned_ranges_inner(Some(control))?;
3841 control.checkpoint()?;
3842 if !before_publish() {
3843 return Err(MongrelError::Cancelled);
3844 }
3845 self.indexes_complete = true;
3846 self.checkpoint_indexes(maintenance_epoch);
3847 Ok((
3848 true,
3849 Some(MaintenanceReceipt {
3850 epoch: maintenance_epoch,
3851 }),
3852 ))
3853 }
3854
3855 fn pending_epoch(&self) -> Epoch {
3856 Epoch(self.epoch.visible().0 + 1)
3857 }
3858
3859 fn is_shared(&self) -> bool {
3862 matches!(self.wal, WalSink::Shared(_))
3863 }
3864
3865 fn ensure_txn_id(&mut self) -> Result<u64> {
3869 if self.current_txn_id == 0 {
3870 let id = match &self.wal {
3871 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
3872 WalSink::Private(_) => {
3873 return Err(MongrelError::Full(
3874 "standalone transaction id namespace exhausted".into(),
3875 ))
3876 }
3877 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3878 };
3879 self.current_txn_id = id;
3880 }
3881 Ok(self.current_txn_id)
3882 }
3883
3884 fn wal_append_data(&mut self, op: Op) -> Result<()> {
3887 self.ensure_writable()?;
3888 let txn_id = self.ensure_txn_id()?;
3889 let table_id = self.table_id;
3890 match &mut self.wal {
3891 WalSink::Private(w) => {
3892 w.append_txn(txn_id, op)?;
3893 self.pending_private_mutations = true;
3894 }
3895 WalSink::Shared(s) => {
3896 s.wal.lock().append(txn_id, table_id, op)?;
3897 }
3898 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
3899 }
3900 Ok(())
3901 }
3902
3903 fn ensure_writable(&self) -> Result<()> {
3904 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
3905 return Err(MongrelError::ReadOnlyReplica);
3906 }
3907 if self.durable_commit_failed {
3908 return Err(MongrelError::Other(
3909 "table poisoned by post-commit failure; reopen required".into(),
3910 ));
3911 }
3912 Ok(())
3913 }
3914
3915 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
3926 match &self.auth {
3927 Some(checker) => checker.check(&self.name, perm),
3928 None => Ok(()),
3929 }
3930 }
3931 pub fn require_select(&self) -> Result<()> {
3936 self.require(crate::auth_state::RequiredPermission::Select)
3937 }
3938 fn require_insert(&self) -> Result<()> {
3939 self.require(crate::auth_state::RequiredPermission::Insert)
3940 }
3941 #[allow(dead_code)]
3945 fn require_update(&self) -> Result<()> {
3946 self.require(crate::auth_state::RequiredPermission::Update)
3947 }
3948 fn require_delete(&self) -> Result<()> {
3949 self.require(crate::auth_state::RequiredPermission::Delete)
3950 }
3951
3952 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
3955 self.require_insert()?;
3956 Ok(self.put_returning(columns)?.0)
3957 }
3958
3959 pub fn put_returning(
3964 &mut self,
3965 mut columns: Vec<(u16, Value)>,
3966 ) -> Result<(RowId, Option<i64>)> {
3967 self.require_insert()?;
3968 let assigned = self.fill_auto_inc(&mut columns)?;
3969 self.apply_defaults(&mut columns)?;
3970 self.schema.validate_values(&columns)?;
3971 let row_id = if self.schema.clustered {
3976 self.derive_clustered_row_id(&columns)?
3977 } else {
3978 self.allocator.alloc()?
3979 };
3980 let epoch = self.pending_epoch();
3981 let mut row = Row::new(row_id, epoch);
3982 for (col_id, val) in columns {
3983 row.columns.insert(col_id, val);
3984 }
3985 self.commit_rows(vec![row], assigned.is_some())?;
3986 Ok((row_id, assigned))
3987 }
3988
3989 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
3992 self.require_insert()?;
3993 Ok(self
3994 .put_batch_returning(batch)?
3995 .into_iter()
3996 .map(|(r, _)| r)
3997 .collect())
3998 }
3999
4000 pub fn put_batch_returning(
4003 &mut self,
4004 batch: Vec<Vec<(u16, Value)>>,
4005 ) -> Result<Vec<(RowId, Option<i64>)>> {
4006 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
4007 for mut cols in batch {
4008 let assigned = self.fill_auto_inc(&mut cols)?;
4009 self.apply_defaults(&mut cols)?;
4010 filled.push((cols, assigned));
4011 }
4012 for (cols, _) in &filled {
4013 self.schema.validate_values(cols)?;
4014 }
4015 let epoch = self.pending_epoch();
4016 let mut rows = Vec::with_capacity(filled.len());
4017 let mut ids = Vec::with_capacity(filled.len());
4018 let first_row_id = if self.schema.clustered {
4019 None
4020 } else {
4021 let count = u64::try_from(filled.len())
4022 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
4023 Some(self.allocator.alloc_range(count)?.0)
4024 };
4025 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
4026 let row_id = match first_row_id {
4027 Some(first) => RowId(first + row_index as u64),
4028 None => self.derive_clustered_row_id(&cols)?,
4029 };
4030 let mut row = Row::new(row_id, epoch);
4031 for (c, v) in cols {
4032 row.columns.insert(c, v);
4033 }
4034 ids.push((row_id, assigned));
4035 rows.push(row);
4036 }
4037 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
4038 self.commit_rows(rows, all_auto_generated)?;
4039 Ok(ids)
4040 }
4041
4042 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
4048 self.ensure_writable()?;
4049 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4050 return Ok(None);
4051 };
4052 let pos = columns.iter().position(|(c, _)| *c == cid);
4053 let assigned = match pos {
4054 Some(i) => match &columns[i].1 {
4055 Value::Null => {
4056 let next = self.alloc_auto_inc_value()?;
4057 columns[i].1 = Value::Int64(next);
4058 Some(next)
4059 }
4060 Value::Int64(n) => {
4061 self.advance_auto_inc_past(*n)?;
4062 None
4063 }
4064 other => {
4065 return Err(MongrelError::InvalidArgument(format!(
4066 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
4067 other
4068 )))
4069 }
4070 },
4071 None => {
4072 let next = self.alloc_auto_inc_value()?;
4073 columns.push((cid, Value::Int64(next)));
4074 Some(next)
4075 }
4076 };
4077 Ok(assigned)
4078 }
4079
4080 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
4086 for col in &self.schema.columns {
4087 let Some(expr) = &col.default_value else {
4088 continue;
4089 };
4090 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
4092 continue;
4093 }
4094 let pos = columns.iter().position(|(c, _)| *c == col.id);
4095 let needs_default = match pos {
4096 None => true,
4097 Some(i) => matches!(columns[i].1, Value::Null),
4098 };
4099 if !needs_default {
4100 continue;
4101 }
4102 let v = match expr {
4103 crate::schema::DefaultExpr::Static(v) => v.clone(),
4104 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
4105 crate::schema::DefaultExpr::Uuid => {
4106 let mut buf = [0u8; 16];
4107 getrandom::getrandom(&mut buf)
4108 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
4109 Value::Uuid(buf)
4110 }
4111 };
4112 match pos {
4113 None => columns.push((col.id, v)),
4114 Some(i) => columns[i].1 = v,
4115 }
4116 }
4117 Ok(())
4118 }
4119
4120 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
4122 self.ensure_auto_inc_seeded()?;
4123 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4125 let v = ai.next;
4126 ai.next = ai
4127 .next
4128 .checked_add(1)
4129 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4130 Ok(v)
4131 }
4132
4133 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
4136 self.ensure_auto_inc_seeded()?;
4137 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4138 let floor = used
4139 .checked_add(1)
4140 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4141 .max(1);
4142 if ai.next < floor {
4143 ai.next = floor;
4144 }
4145 Ok(())
4146 }
4147
4148 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
4153 let needs_seed = match self.auto_inc {
4154 Some(ai) => !ai.seeded,
4155 None => return Ok(()),
4156 };
4157 if !needs_seed {
4158 return Ok(());
4159 }
4160 if self.seed_empty_auto_inc() {
4161 return Ok(());
4162 }
4163 let cid = self
4164 .auto_inc
4165 .as_ref()
4166 .expect("auto-inc column present")
4167 .column_id;
4168 let max = self.scan_max_int64(cid)?;
4169 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4170 let floor = max
4171 .checked_add(1)
4172 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4173 .max(1);
4174 if ai.next < floor {
4175 ai.next = floor;
4176 }
4177 ai.seeded = true;
4178 Ok(())
4179 }
4180
4181 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
4182 if n == 0 || self.auto_inc.is_none() {
4183 return Ok(None);
4184 }
4185 self.ensure_auto_inc_seeded()?;
4186 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4187 let start = ai.next;
4188 let count = i64::try_from(n)
4189 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
4190 ai.next = ai
4191 .next
4192 .checked_add(count)
4193 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4194 Ok(Some(start))
4195 }
4196
4197 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
4201 let mut max: i64 = 0;
4202 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
4203 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4204 if *n > max {
4205 max = *n;
4206 }
4207 }
4208 }
4209 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
4210 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4211 if *n > max {
4212 max = *n;
4213 }
4214 }
4215 }
4216 for rr in self.run_refs.clone() {
4217 let reader = self.open_reader(rr.run_id)?;
4218 if let Some(stats) = reader.column_page_stats(column_id) {
4219 for s in stats {
4220 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
4221 if n > max {
4222 max = n;
4223 }
4224 }
4225 }
4226 } else if reader.has_column(column_id) {
4227 if let columnar::NativeColumn::Int64 { data, validity } =
4228 reader.column_native_shared(column_id)?
4229 {
4230 for (i, n) in data.iter().enumerate() {
4231 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
4232 {
4233 max = *n;
4234 }
4235 }
4236 }
4237 }
4238 }
4239 Ok(max)
4240 }
4241
4242 fn seed_empty_auto_inc(&mut self) -> bool {
4243 let Some(ai) = self.auto_inc.as_mut() else {
4244 return false;
4245 };
4246 if ai.seeded || self.live_count != 0 {
4247 return false;
4248 }
4249 if ai.next < 1 {
4250 ai.next = 1;
4251 }
4252 ai.seeded = true;
4253 true
4254 }
4255
4256 fn advance_auto_inc_from_native_columns(
4257 &mut self,
4258 columns: &[(u16, columnar::NativeColumn)],
4259 n: usize,
4260 live_before: u64,
4261 ) -> Result<()> {
4262 let Some(ai) = self.auto_inc.as_mut() else {
4263 return Ok(());
4264 };
4265 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
4266 return Ok(());
4267 };
4268 let columnar::NativeColumn::Int64 { data, validity } = col else {
4269 return Err(MongrelError::InvalidArgument(format!(
4270 "AUTO_INCREMENT column {} must be Int64",
4271 ai.column_id
4272 )));
4273 };
4274 let max = if native_int64_strictly_increasing(col, n) {
4275 data.get(n.saturating_sub(1)).copied()
4276 } else {
4277 data.iter()
4278 .take(n)
4279 .enumerate()
4280 .filter_map(|(i, v)| {
4281 if validity.is_empty() || columnar::validity_bit(validity, i) {
4282 Some(*v)
4283 } else {
4284 None
4285 }
4286 })
4287 .max()
4288 };
4289 if let Some(max) = max {
4290 let floor = max
4291 .checked_add(1)
4292 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4293 .max(1);
4294 if ai.next < floor {
4295 ai.next = floor;
4296 }
4297 if ai.seeded || live_before == 0 {
4298 ai.seeded = true;
4299 }
4300 }
4301 Ok(())
4302 }
4303
4304 fn fill_auto_inc_native_columns(
4305 &mut self,
4306 columns: &mut Vec<(u16, columnar::NativeColumn)>,
4307 n: usize,
4308 ) -> Result<()> {
4309 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4310 return Ok(());
4311 };
4312 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
4313 if let Some(start) = self.alloc_auto_inc_range(n)? {
4314 columns.push((
4315 cid,
4316 columnar::NativeColumn::Int64 {
4317 data: (start..start.saturating_add(n as i64)).collect(),
4318 validity: vec![0xFF; n.div_ceil(8)],
4319 },
4320 ));
4321 }
4322 return Ok(());
4323 };
4324
4325 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
4326 return Err(MongrelError::InvalidArgument(format!(
4327 "AUTO_INCREMENT column {cid} must be Int64"
4328 )));
4329 };
4330 if data.len() < n {
4331 return Err(MongrelError::InvalidArgument(format!(
4332 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
4333 data.len()
4334 )));
4335 }
4336 if columnar::all_non_null(validity, n) {
4337 return Ok(());
4338 }
4339 if validity.iter().all(|b| *b == 0) {
4340 if let Some(start) = self.alloc_auto_inc_range(n)? {
4341 for (i, slot) in data.iter_mut().take(n).enumerate() {
4342 *slot = start.saturating_add(i as i64);
4343 }
4344 *validity = vec![0xFF; n.div_ceil(8)];
4345 }
4346 return Ok(());
4347 }
4348
4349 let new_validity = vec![0xFF; data.len().div_ceil(8)];
4350 for (i, slot) in data.iter_mut().enumerate().take(n) {
4351 if columnar::validity_bit(validity, i) {
4352 self.advance_auto_inc_past(*slot)?;
4353 } else {
4354 *slot = self.alloc_auto_inc_value()?;
4355 }
4356 }
4357 *validity = new_validity;
4358 Ok(())
4359 }
4360
4361 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
4375 self.ensure_writable()?;
4376 if self.auto_inc.is_none() {
4377 return Ok(None);
4378 }
4379 Ok(Some(self.alloc_auto_inc_value()?))
4380 }
4381
4382 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
4388 let payload = bincode::serialize(&rows)?;
4389 self.wal_append_data(Op::Put {
4390 table_id: self.table_id,
4391 rows: payload,
4392 })?;
4393 if self.is_shared() {
4394 self.pending_rows_auto_inc
4395 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
4396 self.pending_rows.extend(rows);
4397 } else {
4398 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
4399 }
4400 Ok(())
4401 }
4402
4403 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
4406 self.ensure_indexes_complete()
4407 }
4408
4409 pub(crate) fn prepare_durable_publish_controlled(
4410 &mut self,
4411 control: &crate::ExecutionControl,
4412 ) -> Result<()> {
4413 self.ensure_indexes_complete_controlled(control, || true)?;
4414 Ok(())
4415 }
4416
4417 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
4418 self.apply_put_rows_inner_prepared(rows, true);
4419 }
4420
4421 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
4422 if check_existing_pk {
4423 self.ensure_indexes_complete()?;
4424 }
4425 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
4426 Ok(())
4427 }
4428
4429 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
4433 if rows.len() == 1 {
4437 let row = rows.into_iter().next().expect("len checked");
4438 self.apply_put_row_single(row, check_existing_pk);
4439 return;
4440 }
4441 let pk_id = self.schema.primary_key().map(|c| c.id);
4458 let probe = match pk_id {
4459 Some(pid) => {
4460 check_existing_pk
4461 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
4462 }
4463 None => false,
4464 };
4465 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
4468 for r in rows {
4469 for &cid in r.columns.keys() {
4470 self.pending_put_cols.insert(cid);
4471 }
4472 let mut replaced_image: Option<Row> = None;
4473 match pk_id {
4474 Some(pid) if probe || maintain_pk_by_row => {
4475 if let Some(pk_val) = r.columns.get(&pid) {
4476 let key = self.index_lookup_key(pid, pk_val);
4477 if probe {
4478 if let Some(old) = self.recent_delete_preimages.remove(&key) {
4487 replaced_image = Some(old);
4488 if let Some(old_rid) = self.hot.get(&key) {
4489 if old_rid != r.row_id {
4490 self.tombstone_row(
4491 old_rid,
4492 r.committed_epoch,
4493 r.commit_ts,
4494 true,
4495 );
4496 }
4497 }
4498 } else if let Some(old_rid) = self.hot.get(&key) {
4499 if old_rid != r.row_id {
4500 replaced_image = self.get(old_rid, self.snapshot());
4501 self.tombstone_row(
4502 old_rid,
4503 r.committed_epoch,
4504 r.commit_ts,
4505 true,
4506 );
4507 }
4508 }
4509 }
4510 if maintain_pk_by_row {
4511 self.pk_by_row.insert(r.row_id, key);
4512 }
4513 }
4514 }
4515 Some(_) => {}
4516 None => {
4517 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4518 }
4519 }
4520 if let Some(old) = replaced_image {
4521 self.maintain_indexes_on_pk_replace(&old, &r);
4522 } else {
4523 self.index_row(&r);
4524 }
4525 self.reservoir.offer(r.row_id.0);
4526 self.memtable.upsert(r);
4527 self.live_count = self.live_count.saturating_add(1);
4530 }
4531 self.data_generation = self.data_generation.wrapping_add(1);
4532 }
4533
4534 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
4544 for &cid in row.columns.keys() {
4545 self.pending_put_cols.insert(cid);
4546 }
4547 let epoch = row.committed_epoch;
4548 let mut replaced_image: Option<Row> = None;
4549 if let Some(pk_col) = self.schema.primary_key() {
4550 let pk_id = pk_col.id;
4551 if let Some(pk_val) = row.columns.get(&pk_id) {
4552 let maintain_pk_by_row = self.pk_by_row_complete;
4556 if check_existing_pk || maintain_pk_by_row {
4557 let key = self.index_lookup_key(pk_id, pk_val);
4558 if check_existing_pk {
4559 if let Some(old) = self.recent_delete_preimages.remove(&key) {
4568 replaced_image = Some(old);
4569 if let Some(old_rid) = self.hot.get(&key) {
4570 if old_rid != row.row_id {
4571 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
4572 }
4573 }
4574 } else if let Some(old_rid) = self.hot.get(&key) {
4575 if old_rid != row.row_id {
4576 replaced_image = self.get(old_rid, self.snapshot());
4580 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
4581 }
4582 }
4583 }
4584 if maintain_pk_by_row {
4585 self.pk_by_row.insert(row.row_id, key);
4586 }
4587 }
4588 }
4589 } else {
4590 self.hot
4591 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
4592 }
4593 if let Some(old) = replaced_image {
4594 self.maintain_indexes_on_pk_replace(&old, &row);
4595 } else {
4596 self.index_row(&row);
4597 }
4598 self.reservoir.offer(row.row_id.0);
4599 self.memtable.upsert(row);
4600 self.live_count = self.live_count.saturating_add(1);
4601 self.data_generation = self.data_generation.wrapping_add(1);
4602 }
4603
4604 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
4607 let has_partial = self
4608 .schema
4609 .indexes
4610 .iter()
4611 .any(|idx| idx.predicate.is_some());
4612 if has_partial {
4613 self.unindex_bitmap_membership(old);
4617 self.index_row(new);
4618 return;
4619 }
4620
4621 crate::index::maintain_bitmap_secondary_on_replace(
4622 &self.schema,
4623 &mut self.bitmap,
4624 old,
4625 new,
4626 );
4627 self.index_row_non_bitmap(new);
4628 }
4629
4630 fn index_row_non_bitmap(&mut self, row: &Row) {
4633 if row.deleted {
4634 return;
4635 }
4636 let effective = if self.column_keys.is_empty() {
4637 None
4638 } else {
4639 Some(self.tokenized_for_indexes(row))
4640 };
4641 let source = effective.as_ref().unwrap_or(row);
4642 for idef in &self.schema.indexes {
4643 if idef.kind == crate::schema::IndexKind::Bitmap {
4644 continue;
4645 }
4646 index_into_single(
4647 idef,
4648 &self.schema,
4649 source,
4650 &mut self.hot,
4651 &mut self.bitmap,
4652 &mut self.ann,
4653 &mut self.fm,
4654 &mut self.sparse,
4655 &mut self.minhash,
4656 );
4657 }
4658 if let Some(pk_col) = self.schema.primary_key() {
4659 if let Some(pk_val) = source.columns.get(&pk_col.id) {
4660 let key = if self.column_keys.is_empty() {
4661 pk_val.encode_key()
4662 } else {
4663 self.index_lookup_key(pk_col.id, pk_val)
4664 };
4665 self.hot.insert(key, source.row_id);
4666 }
4667 }
4668 }
4669
4670 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
4673 self.allocator.alloc()
4674 }
4675
4676 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
4682 let pk = self.schema.primary_key().ok_or_else(|| {
4683 MongrelError::Schema("clustered table requires a single-column primary key".into())
4684 })?;
4685 let pk_val = columns
4686 .iter()
4687 .find(|(id, _)| *id == pk.id)
4688 .map(|(_, v)| v)
4689 .ok_or_else(|| {
4690 MongrelError::Schema(format!(
4691 "clustered table missing primary key column {} ({})",
4692 pk.id, pk.name
4693 ))
4694 })?;
4695 Ok(clustered_row_id(pk_val))
4696 }
4697
4698 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
4706 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
4707 return Err(MongrelError::Full("row-id namespace exhausted".into()));
4708 }
4709 let n = rows.len();
4710 for r in rows {
4711 for &cid in r.columns.keys() {
4712 self.pending_put_cols.insert(cid);
4713 }
4714 }
4715 let (losers, winner_pks) = self.partition_pk_winners(rows);
4716 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
4717 let group_ts = rows.first().and_then(|r| r.commit_ts);
4719 for (key, &row_id) in &winner_pks {
4720 if let Some(old_rid) = self.hot.get(key) {
4721 if old_rid != row_id {
4722 self.tombstone_row(old_rid, epoch, group_ts, true);
4723 }
4724 }
4725 }
4726 for &loser_rid in &losers {
4729 self.tombstone_row(loser_rid, epoch, group_ts, false);
4730 }
4731 for (key, row_id) in winner_pks {
4733 self.insert_hot_pk(key, row_id);
4734 }
4735 if self.schema.primary_key().is_none() {
4736 for r in rows {
4737 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4738 }
4739 }
4740 for r in rows {
4741 self.allocator.advance_to(r.row_id)?;
4742 if !losers.contains(&r.row_id) {
4743 self.index_row(r);
4744 }
4745 }
4746 for r in rows {
4747 if !losers.contains(&r.row_id) {
4748 self.reservoir.offer(r.row_id.0);
4749 }
4750 }
4751 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
4752 self.data_generation = self.data_generation.wrapping_add(1);
4753 Ok(())
4754 }
4755
4756 pub(crate) fn recover_apply(
4761 &mut self,
4762 rows: Vec<Row>,
4763 deletes: Vec<(RowId, Epoch)>,
4764 ) -> Result<()> {
4765 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
4769 std::collections::BTreeMap::new();
4770 for row in rows {
4771 if row.row_id.0 >= u64::MAX - 1 {
4772 return Err(MongrelError::Full("row-id namespace exhausted".into()));
4773 }
4774 self.allocator.advance_to(row.row_id)?;
4775 if let Some(ai) = self.auto_inc.as_mut() {
4780 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
4781 let next = n.checked_add(1).ok_or_else(|| {
4782 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
4783 })?;
4784 if next > ai.next {
4785 ai.next = next;
4786 }
4787 }
4788 }
4789 by_epoch.entry(row.committed_epoch).or_default().push(row);
4790 }
4791 for (epoch, group) in by_epoch {
4792 let (losers, winner_pks) = self.partition_pk_winners(&group);
4793 let group_ts = group.first().and_then(|r| r.commit_ts);
4795 for (key, &row_id) in &winner_pks {
4796 if let Some(old_rid) = self.hot.get(key) {
4797 if old_rid != row_id {
4798 self.tombstone_row(old_rid, epoch, group_ts, false);
4799 }
4800 }
4801 }
4802 for (key, row_id) in winner_pks {
4803 self.insert_hot_pk(key, row_id);
4804 }
4805 if self.schema.primary_key().is_none() {
4806 for r in &group {
4807 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4808 }
4809 }
4810 for r in &group {
4811 if !losers.contains(&r.row_id) {
4812 self.memtable.upsert(r.clone());
4813 self.index_row(r);
4814 }
4815 }
4816 }
4817 for (rid, epoch) in deletes {
4818 self.memtable.tombstone(rid, epoch);
4819 self.remove_hot_for_row(rid, epoch);
4820 }
4821 self.reservoir_complete = false;
4824 Ok(())
4825 }
4826
4827 pub(crate) fn flushed_epoch(&self) -> u64 {
4829 self.flushed_epoch
4830 }
4831
4832 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
4833 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
4834 }
4835
4836 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
4838 self.schema.validate_values(cells)
4839 }
4840
4841 fn validate_columns_not_null(
4845 &self,
4846 columns: &[(u16, columnar::NativeColumn)],
4847 n: usize,
4848 ) -> Result<()> {
4849 let by_id: HashMap<u16, &columnar::NativeColumn> =
4850 columns.iter().map(|(id, c)| (*id, c)).collect();
4851 for col in &self.schema.columns {
4852 if !col.flags.contains(ColumnFlags::NULLABLE) {
4853 match by_id.get(&col.id) {
4854 None => {
4855 return Err(MongrelError::InvalidArgument(format!(
4856 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
4857 col.name, col.id
4858 )));
4859 }
4860 Some(c) => {
4861 if c.null_count(n) != 0 {
4862 return Err(MongrelError::InvalidArgument(format!(
4863 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
4864 col.name, col.id
4865 )));
4866 }
4867 }
4868 }
4869 }
4870 if let TypeId::Enum { variants } = &col.ty {
4871 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
4872 if by_id.contains_key(&col.id) {
4873 return Err(MongrelError::InvalidArgument(format!(
4874 "column '{}' ({}) enum requires a bytes column",
4875 col.name, col.id
4876 )));
4877 }
4878 continue;
4879 };
4880 for index in 0..n {
4881 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
4882 continue;
4883 };
4884 if !variants.iter().any(|variant| variant.as_bytes() == value) {
4885 return Err(MongrelError::InvalidArgument(format!(
4886 "column '{}' ({}) enum value {:?} is not one of {:?}",
4887 col.name,
4888 col.id,
4889 String::from_utf8_lossy(value),
4890 variants
4891 )));
4892 }
4893 }
4894 }
4895 }
4896 Ok(())
4897 }
4898
4899 fn bulk_pk_winner_indices(
4904 &self,
4905 columns: &[(u16, columnar::NativeColumn)],
4906 n: usize,
4907 ) -> Option<Vec<usize>> {
4908 let pk_col = self.schema.primary_key()?;
4909 let pk_id = pk_col.id;
4910 let pk_ty = pk_col.ty.clone();
4911 let by_id: HashMap<u16, &columnar::NativeColumn> =
4912 columns.iter().map(|(id, c)| (*id, c)).collect();
4913 let pk_native = by_id.get(&pk_id)?;
4914 if native_int64_strictly_increasing(pk_native, n) {
4915 return None;
4916 }
4917 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
4919 let mut null_pk_rows: Vec<usize> = Vec::new();
4920 for i in 0..n {
4921 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
4922 Some(key) => {
4923 last.insert(key, i);
4924 }
4925 None => null_pk_rows.push(i),
4926 }
4927 }
4928 let mut winners: HashSet<usize> = last.values().copied().collect();
4929 for i in null_pk_rows {
4930 winners.insert(i);
4931 }
4932 Some((0..n).filter(|i| winners.contains(i)).collect())
4933 }
4934
4935 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
4937 self.require_delete()?;
4938 let epoch = self.pending_epoch();
4939 self.wal_append_data(Op::Delete {
4940 table_id: self.table_id,
4941 row_ids: vec![row_id],
4942 })?;
4943 if self.is_shared() {
4944 self.pending_dels.push(row_id);
4945 } else {
4946 self.apply_delete(row_id, epoch);
4947 }
4948 Ok(())
4949 }
4950
4951 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
4952 let pre = self.get(row_id, self.snapshot());
4953 self.delete(row_id)?;
4954 Ok(pre.map(|row| {
4955 let mut columns: Vec<_> = row.columns.into_iter().collect();
4956 columns.sort_by_key(|(id, _)| *id);
4957 OwnedRow { columns }
4958 }))
4959 }
4960
4961 pub fn truncate(&mut self) -> Result<()> {
4963 self.require_delete()?;
4964 let epoch = self.pending_epoch();
4965 self.wal_append_data(Op::TruncateTable {
4966 table_id: self.table_id,
4967 })?;
4968 self.pending_rows.clear();
4969 self.pending_rows_auto_inc.clear();
4970 self.pending_dels.clear();
4971 self.pending_truncate = Some(epoch);
4972 Ok(())
4973 }
4974
4975 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
4977 self.run_refs.clear();
4982 self.retiring.clear();
4983 self.memtable = Memtable::new();
4984 self.mutable_run = MutableRun::new();
4985 self.hot = HotIndex::new();
4986 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
4987 self.bitmap = bitmap;
4988 self.ann = ann;
4989 self.fm = fm;
4990 self.sparse = sparse;
4991 self.minhash = minhash;
4992 self.learned_range = Arc::new(HashMap::new());
4993 self.pk_by_row.clear();
4994 self.pk_by_row_complete = false;
4995 self.live_count = 0;
4996 self.reservoir = crate::reservoir::Reservoir::default();
4997 self.reservoir_complete = true;
4998 self.had_deletes = true;
4999 self.agg_cache = Arc::new(HashMap::new());
5000 self.global_idx_epoch = 0;
5001 self.indexes_complete = true;
5002 self.pending_delete_rids.clear();
5003 self.pending_put_cols.clear();
5004 self.pending_rows.clear();
5005 self.pending_rows_auto_inc.clear();
5006 self.pending_dels.clear();
5007 self.clear_result_cache();
5008 self.invalidate_index_checkpoint();
5009 self.data_generation = self.data_generation.wrapping_add(1);
5010 }
5011
5012 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
5015 self.apply_delete_at(row_id, epoch, None);
5016 }
5017
5018 pub(crate) fn apply_delete_at(
5020 &mut self,
5021 row_id: RowId,
5022 epoch: Epoch,
5023 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5024 ) {
5025 let preimage = self.get(row_id, self.snapshot());
5035 if let Some(row) = preimage {
5036 if let Some(pk_col) = self.schema.primary_key() {
5037 if let Some(pk_val) = row.columns.get(&pk_col.id) {
5038 let key = self.index_lookup_key(pk_col.id, pk_val);
5039 self.recent_delete_preimages.insert(key, row);
5040 }
5041 }
5042 }
5043 self.tombstone_row(row_id, epoch, commit_ts, true);
5044 self.data_generation = self.data_generation.wrapping_add(1);
5045 }
5046
5047 fn unindex_bitmap_membership(&mut self, row: &Row) {
5052 if row.deleted {
5053 return;
5054 }
5055 for idef in &self.schema.indexes {
5056 if idef.kind != crate::schema::IndexKind::Bitmap {
5057 continue;
5058 }
5059 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
5060 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5061 b.remove(&key, row.row_id);
5062 }
5063 }
5064 }
5065 }
5066
5067 fn union_bitmap_point_i64(
5071 &self,
5072 set: &mut RowIdSet,
5073 column_id: u16,
5074 value: i64,
5075 snapshot: Snapshot,
5076 ) {
5077 let Some(b) = self.bitmap.get(&column_id) else {
5078 return;
5079 };
5080 let encoded = Value::Int64(value).encode_key();
5081 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
5082 for rid in b.get(&lookup).iter() {
5083 set.insert(u64::from(rid));
5084 }
5085 set.remove_many(self.overlay_tombstoned_rids(snapshot));
5089 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
5090 }
5091
5092 fn tombstone_row(
5106 &mut self,
5107 row_id: RowId,
5108 epoch: Epoch,
5109 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5110 adjust_live_count: bool,
5111 ) {
5112 let prev_was_live = if adjust_live_count {
5113 match self.memtable.get_version(row_id, epoch) {
5114 Some((_, prev)) => !prev.deleted,
5115 None => true,
5116 }
5117 } else {
5118 false
5119 };
5120 let tombstone = Row {
5121 row_id,
5122 committed_epoch: epoch,
5123 columns: std::collections::HashMap::new(),
5124 deleted: true,
5125 commit_ts,
5126 };
5127 self.memtable.upsert(tombstone);
5128 self.pk_by_row.remove(&row_id);
5129 if prev_was_live {
5130 self.live_count = self.live_count.saturating_sub(1);
5131 }
5132 self.pending_delete_rids.insert(row_id.0 as u32);
5134 self.had_deletes = true;
5137 self.agg_cache = Arc::new(HashMap::new());
5138 }
5139
5140 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
5144 let Some(pk_col) = self.schema.primary_key() else {
5145 return;
5146 };
5147 if self.pk_by_row_complete {
5150 if let Some(key) = self.pk_by_row.remove(&row_id) {
5151 if self.hot.get(&key) == Some(row_id) {
5152 self.hot.remove(&key);
5153 }
5154 }
5155 return;
5156 }
5157 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
5176 if self.indexes_complete {
5177 let pk_val = self
5178 .memtable
5179 .get_version(row_id, lookup_epoch)
5180 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5181 .or_else(|| {
5182 self.mutable_run
5183 .get_version(row_id, lookup_epoch)
5184 .filter(|(_, r)| !r.deleted)
5185 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5186 })
5187 .or_else(|| {
5188 self.run_refs.iter().find_map(|rr| {
5189 let mut reader = self.open_reader(rr.run_id).ok()?;
5190 let (_, deleted, val) = reader
5191 .get_version_column(row_id, lookup_epoch, pk_col.id)
5192 .ok()??;
5193 if deleted {
5194 return None;
5195 }
5196 val
5197 })
5198 });
5199 if let Some(pk_val) = pk_val {
5200 let key = self.index_lookup_key(pk_col.id, &pk_val);
5201 if self.hot.get(&key) == Some(row_id) {
5202 self.hot.remove(&key);
5203 }
5204 return;
5205 }
5206 }
5207 self.refresh_pk_by_row_from_hot();
5212 if let Some(key) = self.pk_by_row.remove(&row_id) {
5213 if self.hot.get(&key) == Some(row_id) {
5214 self.hot.remove(&key);
5215 }
5216 }
5217 }
5218
5219 fn partition_pk_winners(
5224 &self,
5225 rows: &[Row],
5226 ) -> (
5227 std::collections::HashSet<RowId>,
5228 std::collections::HashMap<Vec<u8>, RowId>,
5229 ) {
5230 let mut losers = std::collections::HashSet::new();
5231 let Some(pk_col) = self.schema.primary_key() else {
5232 return (losers, std::collections::HashMap::new());
5233 };
5234 let pk_id = pk_col.id;
5235 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
5236 std::collections::HashMap::new();
5237 for r in rows {
5238 let Some(pk_val) = r.columns.get(&pk_id) else {
5239 continue;
5240 };
5241 let key = self.index_lookup_key(pk_id, pk_val);
5242 if let Some(&old_rid) = winners.get(&key) {
5243 losers.insert(old_rid);
5244 }
5245 winners.insert(key, r.row_id);
5246 }
5247 (losers, winners)
5248 }
5249
5250 fn index_row(&mut self, row: &Row) {
5251 if row.deleted {
5252 return;
5253 }
5254 let any_predicate = self
5262 .schema
5263 .indexes
5264 .iter()
5265 .any(|idx| idx.predicate.is_some());
5266 if any_predicate {
5267 let columns_map: HashMap<u16, &Value> =
5268 row.columns.iter().map(|(k, v)| (*k, v)).collect();
5269 let name_to_id: HashMap<&str, u16> = self
5270 .schema
5271 .columns
5272 .iter()
5273 .map(|c| (c.name.as_str(), c.id))
5274 .collect();
5275 for idx in &self.schema.indexes {
5276 if let Some(pred) = &idx.predicate {
5277 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
5278 continue; }
5280 }
5281 index_into_single(
5283 idx,
5284 &self.schema,
5285 row,
5286 &mut self.hot,
5287 &mut self.bitmap,
5288 &mut self.ann,
5289 &mut self.fm,
5290 &mut self.sparse,
5291 &mut self.minhash,
5292 );
5293 }
5294 return;
5295 }
5296 if self.column_keys.is_empty() {
5300 index_into(
5301 &self.schema,
5302 row,
5303 &mut self.hot,
5304 &mut self.bitmap,
5305 &mut self.ann,
5306 &mut self.fm,
5307 &mut self.sparse,
5308 &mut self.minhash,
5309 );
5310 return;
5311 }
5312 let effective_row = self.tokenized_for_indexes(row);
5313 index_into(
5314 &self.schema,
5315 &effective_row,
5316 &mut self.hot,
5317 &mut self.bitmap,
5318 &mut self.ann,
5319 &mut self.fm,
5320 &mut self.sparse,
5321 &mut self.minhash,
5322 );
5323 }
5324
5325 fn tokenized_for_indexes(&self, row: &Row) -> Row {
5331 if self.column_keys.is_empty() {
5332 return row.clone();
5333 }
5334 {
5335 use crate::encryption::SCHEME_HMAC_EQ;
5336 let mut tok = row.clone();
5337 for (&cid, &(_, scheme)) in &self.column_keys {
5338 if scheme != SCHEME_HMAC_EQ {
5339 continue;
5340 }
5341 if let Some(v) = tok.columns.get(&cid).cloned() {
5342 if let Some(t) = self.tokenize_value(cid, &v) {
5343 tok.columns.insert(cid, t);
5344 }
5345 }
5346 }
5347 tok
5348 }
5349 }
5350
5351 pub fn commit(&mut self) -> Result<Epoch> {
5356 self.commit_inner(None)
5357 }
5358
5359 #[doc(hidden)]
5362 pub fn commit_controlled<F>(
5363 &mut self,
5364 control: &crate::ExecutionControl,
5365 mut before_commit: F,
5366 ) -> Result<Epoch>
5367 where
5368 F: FnMut() -> Result<()>,
5369 {
5370 self.commit_inner(Some((control, &mut before_commit)))
5371 }
5372
5373 fn commit_inner(
5374 &mut self,
5375 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5376 ) -> Result<Epoch> {
5377 self.ensure_writable()?;
5378 if !self.has_pending_mutations() {
5379 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
5380 return Err(MongrelError::Full(
5381 "standalone transaction id namespace exhausted".into(),
5382 ));
5383 }
5384 return Ok(self.epoch.visible());
5385 }
5386 self.commit_new_epoch_inner(controlled)
5387 }
5388
5389 fn commit_new_epoch(&mut self) -> Result<Epoch> {
5393 self.commit_new_epoch_inner(None)
5394 }
5395
5396 fn commit_new_epoch_inner(
5397 &mut self,
5398 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5399 ) -> Result<Epoch> {
5400 self.ensure_writable()?;
5401 if self.is_shared() {
5402 self.commit_shared(controlled)
5403 } else {
5404 self.commit_private(controlled)
5405 }
5406 }
5407
5408 fn commit_private(
5410 &mut self,
5411 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5412 ) -> Result<Epoch> {
5413 let commit_lock = Arc::clone(&self.commit_lock);
5417 let _g = commit_lock.lock();
5418 let txn_id = self.ensure_txn_id()?;
5421 if let Some((control, before_commit)) = controlled {
5422 control.checkpoint()?;
5423 before_commit()?;
5424 }
5425 let new_epoch = self.epoch.bump_assigned();
5426 let epoch_authority = Arc::clone(&self.epoch);
5427 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5428 let wal_result = match &mut self.wal {
5432 WalSink::Private(w) => w
5433 .append_txn(
5434 txn_id,
5435 Op::TxnCommit {
5436 epoch: new_epoch.0,
5437 added_runs: Vec::new(),
5438 },
5439 )
5440 .and_then(|_| w.sync()),
5441 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
5442 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
5443 };
5444 if let Err(error) = wal_result {
5445 self.durable_commit_failed = true;
5446 return Err(MongrelError::CommitOutcomeUnknown {
5447 epoch: new_epoch.0,
5448 message: error.to_string(),
5449 });
5450 }
5451 if let Some(epoch) = self.pending_truncate.take() {
5454 self.apply_truncate(epoch);
5455 }
5456 self.invalidate_pending_cache();
5457 let publish_result = self.persist_manifest(new_epoch);
5458 self.epoch.publish_in_order(new_epoch);
5462 epoch_guard.disarm();
5463 if let Err(error) = publish_result {
5464 self.durable_commit_failed = true;
5465 return Err(MongrelError::DurableCommit {
5466 epoch: new_epoch.0,
5467 message: error.to_string(),
5468 });
5469 }
5470 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
5471 self.pending_private_mutations = false;
5472 self.data_generation = self.data_generation.wrapping_add(1);
5473 Ok(new_epoch)
5474 }
5475
5476 fn commit_shared(
5482 &mut self,
5483 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5484 ) -> Result<Epoch> {
5485 use std::sync::atomic::Ordering;
5486 let s = match &self.wal {
5487 WalSink::Shared(s) => s.clone(),
5488 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
5489 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
5490 };
5491 if s.poisoned.load(Ordering::Relaxed) {
5492 return Err(MongrelError::Other(
5493 "database poisoned by fsync error".into(),
5494 ));
5495 }
5496 let commit_lock = Arc::clone(&self.commit_lock);
5503 let _g = commit_lock.lock();
5504 if !self.pending_rows.is_empty() {
5505 match controlled.as_ref() {
5506 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
5507 None => self.prepare_durable_publish()?,
5508 }
5509 }
5510 let txn_id = self.ensure_txn_id()?;
5513 let mut wal = s.wal.lock();
5514 if let Some((control, before_commit)) = controlled {
5515 control.checkpoint()?;
5516 before_commit()?;
5517 }
5518 let new_epoch = self.epoch.bump_assigned();
5519 let epoch_authority = Arc::clone(&self.epoch);
5520 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5521 let commit_ts = s.hlc.now().map_err(|skew| {
5524 MongrelError::Other(format!(
5525 "clock skew rejected commit timestamp allocation: {skew}"
5526 ))
5527 })?;
5528 let commit_seq = match wal.append_commit_at(
5529 txn_id,
5530 new_epoch,
5531 &[],
5532 commit_ts.physical_micros.saturating_mul(1_000),
5533 ) {
5534 Ok(commit_seq) => commit_seq,
5535 Err(error) => {
5536 s.poisoned.store(true, Ordering::Relaxed);
5537 s.lifecycle.poison();
5538 return Err(MongrelError::CommitOutcomeUnknown {
5539 epoch: new_epoch.0,
5540 message: error.to_string(),
5541 });
5542 }
5543 };
5544 drop(wal);
5545 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
5546 s.poisoned.store(true, Ordering::Relaxed);
5547 s.lifecycle.poison();
5548 return Err(MongrelError::CommitOutcomeUnknown {
5549 epoch: new_epoch.0,
5550 message: error.to_string(),
5551 });
5552 }
5553
5554 if self.pending_truncate.take().is_some() {
5557 self.apply_truncate(new_epoch);
5558 }
5559 let mut rows = std::mem::take(&mut self.pending_rows);
5560 if !rows.is_empty() {
5561 for r in &mut rows {
5562 r.committed_epoch = new_epoch;
5563 r.commit_ts = Some(commit_ts);
5564 }
5565 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
5566 let all_auto_generated =
5567 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
5568 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
5569 } else {
5570 self.pending_rows_auto_inc.clear();
5571 }
5572 let dels = std::mem::take(&mut self.pending_dels);
5573 for rid in dels {
5574 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
5575 }
5576
5577 self.invalidate_pending_cache();
5578 let publish_result = self.persist_manifest(new_epoch);
5579 self.epoch.publish_in_order(new_epoch);
5580 epoch_guard.disarm();
5581 let _ = s.change_wake.send(());
5582 if let Err(error) = publish_result {
5583 self.durable_commit_failed = true;
5584 s.poisoned.store(true, Ordering::Relaxed);
5585 s.lifecycle.poison();
5586 return Err(MongrelError::DurableCommit {
5587 epoch: new_epoch.0,
5588 message: error.to_string(),
5589 });
5590 }
5591 self.current_txn_id = 0;
5593 self.data_generation = self.data_generation.wrapping_add(1);
5594 Ok(new_epoch)
5595 }
5596
5597 pub fn flush(&mut self) -> Result<Epoch> {
5605 self.flush_with_outcome().map(|(epoch, _)| epoch)
5606 }
5607
5608 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
5610 self.flush_with_outcome_inner(None)
5611 }
5612
5613 #[doc(hidden)]
5616 pub fn flush_with_outcome_controlled<F>(
5617 &mut self,
5618 control: &crate::ExecutionControl,
5619 mut before_commit: F,
5620 ) -> Result<(Epoch, bool)>
5621 where
5622 F: FnMut() -> Result<()>,
5623 {
5624 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
5625 }
5626
5627 fn flush_with_outcome_inner(
5628 &mut self,
5629 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5630 ) -> Result<(Epoch, bool)> {
5631 match controlled.as_ref() {
5632 Some((control, _)) => {
5633 self.ensure_indexes_complete_controlled(control, || true)?;
5634 }
5635 None => self.ensure_indexes_complete()?,
5636 }
5637 let committed = self.has_pending_mutations();
5638 let epoch = self.commit_inner(controlled)?;
5639 let finish: Result<(Epoch, bool)> = (|| {
5640 let rows = self.memtable.drain_sorted();
5641 if !rows.is_empty() {
5642 self.mutable_run.insert_many(rows);
5643 }
5644 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
5645 self.spill_mutable_run(epoch)?;
5646 self.mark_flushed(epoch)?;
5650 self.persist_manifest(epoch)?;
5651 self.build_learned_ranges()?;
5652 self.checkpoint_indexes(epoch);
5655 }
5656 Ok((epoch, committed))
5659 })();
5660 let outcome = match finish {
5661 Err(error) if committed => Err(MongrelError::DurableCommit {
5662 epoch: epoch.0,
5663 message: error.to_string(),
5664 }),
5665 result => result,
5666 };
5667 if outcome.is_ok() {
5668 let _ = self.publish_read_generation();
5674 }
5675 outcome
5676 }
5677
5678 fn has_pending_mutations(&self) -> bool {
5679 self.pending_private_mutations
5680 || !self.pending_rows.is_empty()
5681 || !self.pending_dels.is_empty()
5682 || self.pending_truncate.is_some()
5683 }
5684
5685 pub fn has_pending_writes(&self) -> bool {
5686 self.has_pending_mutations()
5687 }
5688
5689 pub fn force_flush(&mut self) -> Result<Epoch> {
5698 let saved = self.mutable_run_spill_bytes;
5699 self.mutable_run_spill_bytes = 1;
5700 let result = self.flush();
5701 self.mutable_run_spill_bytes = saved;
5702 result
5703 }
5704
5705 pub fn close(&mut self) -> Result<()> {
5712 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
5713 self.force_flush()?;
5714 }
5715 Ok(())
5716 }
5717
5718 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
5725 let op = Op::Flush {
5726 table_id: self.table_id,
5727 flushed_epoch: epoch.0,
5728 };
5729 match &mut self.wal {
5730 WalSink::Private(w) => {
5731 w.append_system(op)?;
5732 w.sync()?;
5733 }
5734 WalSink::Shared(s) => {
5735 s.wal.lock().append_system(op)?;
5740 }
5741 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
5742 }
5743 self.flushed_epoch = epoch.0;
5744 if matches!(self.wal, WalSink::Private(_)) {
5745 self.rotate_wal(epoch)?;
5746 }
5747 Ok(())
5748 }
5749
5750 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
5754 if self.mutable_run.is_empty() {
5755 return Ok(());
5756 }
5757 let run_id = self.alloc_run_id()?;
5758 let rows = self.mutable_run.drain_sorted();
5759 if rows.is_empty() {
5760 return Ok(());
5761 }
5762 let path = self.run_path(run_id);
5763 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
5764 if let Some(kek) = &self.kek {
5765 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5766 }
5767 let header = match self.create_run_file(run_id)? {
5768 Some(file) => writer.write_file(file, &rows)?,
5769 None => writer.write(&path, &rows)?,
5770 };
5771 self.run_refs.push(RunRef {
5772 run_id: run_id as u128,
5773 level: 0,
5774 epoch_created: epoch.0,
5775 row_count: header.row_count,
5776 });
5777 self.run_row_id_ranges
5778 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
5779 Ok(())
5780 }
5781
5782 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
5786 self.mutable_run_spill_bytes = bytes.max(1);
5787 }
5788
5789 pub fn set_compaction_zstd_level(&mut self, level: i32) {
5793 self.compaction_zstd_level = level;
5794 }
5795
5796 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
5800 self.result_cache.lock().set_max_bytes(max_bytes);
5801 }
5802
5803 pub(crate) fn clear_result_cache(&mut self) {
5807 self.result_cache.lock().clear();
5808 }
5809
5810 pub fn mutable_run_len(&self) -> usize {
5812 self.mutable_run.len()
5813 }
5814
5815 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
5818 self.mutable_run.drain_sorted()
5819 }
5820
5821 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
5823 let mut snapshot = self.mutable_run.clone();
5824 snapshot.drain_sorted()
5825 }
5826
5827 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
5832 self.ensure_writable()?;
5833 let n = batch.len();
5834 if n == 0 {
5835 return Ok(self.current_epoch());
5836 }
5837 for row in &batch {
5838 self.schema.validate_values(row)?;
5839 }
5840 let epoch = self.commit_new_epoch()?;
5841 let live_before = self.live_count;
5842 self.spill_mutable_run(epoch)?;
5846 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
5847 && self.indexes_complete
5848 && self.run_refs.is_empty()
5849 && self.memtable.is_empty()
5850 && self.mutable_run.is_empty();
5851 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
5857 use rayon::prelude::*;
5858 self.schema
5859 .columns
5860 .par_iter()
5861 .map(|cdef| {
5862 (
5863 cdef.id,
5864 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
5865 )
5866 })
5867 .collect::<Vec<_>>()
5868 };
5869 drop(batch);
5870 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
5875 self.validate_columns_not_null(&user_columns, n)?;
5876 let winner_idx = self
5877 .bulk_pk_winner_indices(&user_columns, n)
5878 .filter(|idx| idx.len() != n);
5879 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
5880 match winner_idx.as_deref() {
5881 Some(idx) => {
5882 let compacted = user_columns
5883 .iter()
5884 .map(|(id, c)| (*id, c.gather(idx)))
5885 .collect();
5886 (compacted, idx.len())
5887 }
5888 None => (std::mem::take(&mut user_columns), n),
5889 };
5890 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
5891 let first = self.allocator.alloc_range(write_n as u64)?.0;
5892 for rid in first..first + write_n as u64 {
5893 self.reservoir.offer(rid);
5894 }
5895 let run_id = self.alloc_run_id()?;
5896 let path = self.run_path(run_id);
5897 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
5898 .clean(true)
5899 .with_lz4()
5900 .with_native_endian();
5901 if let Some(kek) = &self.kek {
5902 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
5903 }
5904 let header = match self.create_run_file(run_id)? {
5905 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
5906 None => writer.write_native(&path, &write_columns, write_n, first)?,
5907 };
5908 self.run_refs.push(RunRef {
5909 run_id: run_id as u128,
5910 level: 0,
5911 epoch_created: epoch.0,
5912 row_count: header.row_count,
5913 });
5914 self.run_row_id_ranges
5915 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
5916 self.live_count = self.live_count.saturating_add(write_n as u64);
5917 if eager_index_build {
5918 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
5919 self.index_columns_bulk(&write_columns, &row_ids);
5920 self.indexes_complete = true;
5921 self.build_learned_ranges()?;
5922 } else {
5923 self.indexes_complete = false;
5924 }
5925 self.mark_flushed(epoch)?;
5926 self.persist_manifest(epoch)?;
5927 if eager_index_build {
5928 self.checkpoint_indexes(epoch);
5929 }
5930 self.clear_result_cache();
5931 Ok(epoch)
5932 }
5933
5934 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
5937 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
5938 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
5939 let segment_no = segment
5942 .file_stem()
5943 .and_then(|s| s.to_str())
5944 .and_then(|s| s.strip_prefix("seg-"))
5945 .and_then(|s| s.parse::<u64>().ok())
5946 .unwrap_or(0);
5947 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
5948 wal.set_sync_byte_threshold(self.sync_byte_threshold);
5949 wal.sync()?;
5950 self.wal = WalSink::Private(wal);
5951 Ok(())
5952 }
5953
5954 pub(crate) fn invalidate_pending_cache(&mut self) {
5959 self.result_cache
5960 .lock()
5961 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
5962 self.pending_delete_rids.clear();
5963 self.pending_put_cols.clear();
5964 }
5965
5966 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
5967 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
5968 m.current_epoch = epoch.0;
5969 m.next_row_id = self.allocator.current().0;
5970 m.runs = self.run_refs.clone();
5971 m.live_count = self.live_count;
5972 m.global_idx_epoch = self.global_idx_epoch;
5973 m.flushed_epoch = self.flushed_epoch;
5974 m.retiring = self.retiring.clone();
5975 m.auto_inc_next = match self.auto_inc {
5979 Some(ai) if ai.seeded => ai.next,
5980 _ => 0,
5981 };
5982 m.ttl = self.ttl;
5983 let meta_dek = self.manifest_meta_dek();
5984 match self._root_guard.as_deref() {
5985 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
5986 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
5987 }
5988 Ok(())
5989 }
5990
5991 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
5992 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
5996 let live_count = u64::try_from(rows.len())
5997 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
5998 let auto_inc = match self.auto_inc {
5999 Some(mut state) => {
6000 let maximum = self.scan_max_int64(state.column_id)?;
6001 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
6002 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
6003 })?;
6004 state.next = state.next.max(after_maximum).max(1);
6005 state.seeded = true;
6006 Some(state)
6007 }
6008 None => None,
6009 };
6010 Ok(RecoveryMetadataPlan {
6011 live_count,
6012 auto_inc,
6013 changed: live_count != self.live_count
6014 || auto_inc.is_some_and(|planned| {
6015 self.auto_inc.is_none_or(|current| {
6016 current.next != planned.next || current.seeded != planned.seeded
6017 })
6018 }),
6019 })
6020 }
6021
6022 pub(crate) fn apply_recovered_metadata(
6023 &mut self,
6024 plan: RecoveryMetadataPlan,
6025 epoch: Epoch,
6026 ) -> Result<()> {
6027 if !plan.changed {
6028 return Ok(());
6029 }
6030 self.live_count = plan.live_count;
6031 self.auto_inc = plan.auto_inc;
6032 self.persist_manifest(epoch)
6033 }
6034
6035 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
6041 if !self.indexes_complete {
6044 return;
6045 }
6046 if crate::catalog::inject_hook("index.publish.before").is_err() {
6049 return;
6050 }
6051 if self.idx_root.is_none() {
6052 if let Some(root) = self._root_guard.as_ref() {
6053 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
6054 return;
6055 };
6056 self.idx_root = Some(Arc::new(idx_root));
6057 }
6058 }
6059 let snap = global_idx::IndexSnapshot {
6060 hot: &self.hot,
6061 bitmap: &self.bitmap,
6062 ann: &self.ann,
6063 fm: &self.fm,
6064 sparse: &self.sparse,
6065 minhash: &self.minhash,
6066 learned_range: &self.learned_range,
6067 };
6068 let idx_dek = self.idx_dek();
6070 let written = match self.idx_root.as_deref() {
6071 Some(root) => global_idx::write_atomic_root(
6072 root,
6073 self.table_id,
6074 epoch.0,
6075 snap,
6076 idx_dek.as_deref(),
6077 ),
6078 None => global_idx::write_atomic(
6079 &self.dir,
6080 self.table_id,
6081 epoch.0,
6082 snap,
6083 idx_dek.as_deref(),
6084 ),
6085 };
6086 if written.is_ok() {
6087 self.global_idx_epoch = epoch.0;
6088 let _ = self.persist_manifest(epoch);
6089 let _ = crate::catalog::inject_hook("index.publish.after");
6091 }
6092 }
6093
6094 pub(crate) fn invalidate_index_checkpoint(&mut self) {
6097 self.global_idx_epoch = 0;
6098 if let Some(root) = self.idx_root.as_deref() {
6099 let _ = root.remove_file(global_idx::IDX_FILENAME);
6100 } else {
6101 global_idx::remove(&self.dir);
6102 }
6103 let _ = self.persist_manifest(self.epoch.visible());
6104 }
6105
6106 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
6111 self.indexes_complete = false;
6112 self.global_idx_epoch = 0;
6113 if let Some(root) = self.idx_root.as_deref() {
6114 let _ = root.remove_file(global_idx::IDX_FILENAME);
6115 } else {
6116 global_idx::remove(&self.dir);
6117 }
6118 }
6119
6120 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
6121 self.indexes_complete = true;
6122 }
6123
6124 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
6130 self.durable_commit_failed = true;
6131 if let WalSink::Shared(shared) = &self.wal {
6132 shared
6133 .poisoned
6134 .store(true, std::sync::atomic::Ordering::Relaxed);
6135 }
6136 }
6137
6138 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
6142 self.durable_commit_failed = true;
6143 }
6144
6145 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
6154 let mut best: Option<Row> = None;
6155 let mut consider = |row: Row| {
6156 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6157 return;
6158 }
6159 if best.as_ref().is_none_or(|current| {
6160 Snapshot::version_is_newer(
6161 row.committed_epoch,
6162 row.commit_ts,
6163 current.committed_epoch,
6164 current.commit_ts,
6165 )
6166 }) {
6167 best = Some(row);
6168 }
6169 };
6170 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
6171 consider(row);
6172 }
6173 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
6174 consider(row);
6175 }
6176 for rr in &self.run_refs {
6177 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
6180 if row_id.0 < min_rid || row_id.0 > max_rid {
6181 self.lookup_metrics
6182 .get_run_skipped
6183 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6184 continue;
6185 }
6186 }
6187 self.lookup_metrics
6188 .get_run_opened
6189 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6190 let Ok(mut reader) = self.open_reader(rr.run_id) else {
6191 continue;
6192 };
6193 let Ok(Some((_, row))) = reader.get_version(row_id, snapshot.epoch) else {
6196 continue;
6197 };
6198 consider(row);
6199 }
6200 let now_nanos = unix_nanos_now();
6201 match best {
6202 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
6203 Some(r) => Some(r),
6204 None => None,
6205 }
6206 }
6207
6208 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
6212 self.visible_rows_at_time(snapshot, unix_nanos_now())
6213 }
6214
6215 #[doc(hidden)]
6218 pub fn visible_rows_controlled(
6219 &self,
6220 snapshot: Snapshot,
6221 control: &crate::ExecutionControl,
6222 ) -> Result<Vec<Row>> {
6223 let mut out = Vec::new();
6224 self.for_each_visible_row_controlled(snapshot, control, |row| {
6225 out.push(row);
6226 Ok(())
6227 })?;
6228 Ok(out)
6229 }
6230
6231 #[doc(hidden)]
6234 pub fn for_each_visible_row_controlled<F>(
6235 &self,
6236 snapshot: Snapshot,
6237 control: &crate::ExecutionControl,
6238 mut visit: F,
6239 ) -> Result<()>
6240 where
6241 F: FnMut(Row) -> Result<()>,
6242 {
6243 let setup_start = std::time::Instant::now();
6250 let mut versions_examined: u64 = 0;
6251 let mut rows_emitted: u64 = 0;
6252 let mut checkpoints: u64 = 0;
6253 let mut first_row_recorded = false;
6254 let mut sources: Vec<ControlledVisibleSource<'_>> =
6255 Vec::with_capacity(self.run_refs.len() + 2);
6256 control.checkpoint()?;
6257 checkpoints += 1;
6258 let memtable_map = self.memtable.newest_visible_map(snapshot);
6264 if !memtable_map.is_empty() {
6265 sources.push(ControlledVisibleSource::memory_from_map(memtable_map));
6266 }
6267 control.checkpoint()?;
6268 checkpoints += 1;
6269 let mutable_map = self.mutable_run.newest_visible_map(snapshot);
6270 if !mutable_map.is_empty() {
6271 sources.push(ControlledVisibleSource::memory_from_map(mutable_map));
6272 }
6273 for run in &self.run_refs {
6274 control.checkpoint()?;
6275 checkpoints += 1;
6276 let reader = self.open_reader(run.run_id)?;
6277 sources.push(ControlledVisibleSource::run(
6278 reader.into_visible_version_cursor(snapshot.epoch)?,
6279 ));
6280 }
6281 let start = std::time::Instant::now();
6286 let _setup_us = start.duration_since(setup_start).as_micros() as u64;
6287 let now_nanos = unix_nanos_now();
6288 let mut first_row_us: u64 = 0;
6289 let result = merge_controlled_visible_sources(
6290 &mut sources,
6291 control,
6292 |row| self.row_expired_at(row, now_nanos),
6293 |row| {
6294 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6295 return Ok(());
6296 }
6297 versions_examined += 1;
6298 if !first_row_recorded {
6299 first_row_us = start.elapsed().as_micros() as u64;
6300 first_row_recorded = true;
6301 }
6302 rows_emitted += 1;
6303 visit(row)
6304 },
6305 );
6306 let time_to_first_row_us = if first_row_recorded { first_row_us } else { 0 };
6307 crate::trace::QueryTrace::record(|t| {
6308 t.controlled_scan_versions_examined = t
6309 .controlled_scan_versions_examined
6310 .saturating_add(versions_examined as usize);
6311 t.controlled_scan_rows_emitted = t
6312 .controlled_scan_rows_emitted
6313 .saturating_add(rows_emitted as usize);
6314 t.controlled_scan_checkpoints = t
6315 .controlled_scan_checkpoints
6316 .saturating_add(checkpoints as usize);
6317 t.controlled_scan_time_to_first_row_us = t
6318 .controlled_scan_time_to_first_row_us
6319 .saturating_add(time_to_first_row_us);
6320 });
6321 result
6322 }
6323
6324 #[doc(hidden)]
6325 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
6326 let mut best: HashMap<u64, Row> = HashMap::new();
6327 let mut fold = |row: Row| {
6328 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6329 return;
6330 }
6331 best.entry(row.row_id.0)
6332 .and_modify(|existing| {
6333 if Snapshot::version_is_newer(
6334 row.committed_epoch,
6335 row.commit_ts,
6336 existing.committed_epoch,
6337 existing.commit_ts,
6338 ) {
6339 *existing = row.clone();
6340 }
6341 })
6342 .or_insert(row);
6343 };
6344 for row in self.memtable.visible_versions_at(snapshot) {
6345 fold(row);
6346 }
6347 for row in self.mutable_run.visible_versions_at(snapshot) {
6348 fold(row);
6349 }
6350 for rr in &self.run_refs {
6351 let mut reader = self.open_reader(rr.run_id)?;
6352 for row in reader.visible_versions(snapshot.epoch)? {
6354 fold(row);
6355 }
6356 }
6357 let mut out: Vec<Row> = best
6358 .into_values()
6359 .filter_map(|r| {
6360 if r.deleted || self.row_expired_at(&r, now_nanos) {
6361 None
6362 } else {
6363 Some(r)
6364 }
6365 })
6366 .collect();
6367 out.sort_by_key(|r| r.row_id);
6368 Ok(out)
6369 }
6370
6371 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
6378 if self.ttl.is_none()
6379 && self.memtable.is_empty()
6380 && self.mutable_run.is_empty()
6381 && self.run_refs.len() == 1
6382 {
6383 let rr = self.run_refs[0].clone();
6384 let mut reader = self.open_reader(rr.run_id)?;
6385 let idxs = reader.visible_indices(snapshot.epoch)?;
6386 let mut cols = Vec::with_capacity(self.schema.columns.len());
6387 for cdef in &self.schema.columns {
6388 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
6389 }
6390 return Ok(cols);
6391 }
6392 let rows = self.visible_rows(snapshot)?;
6394 let mut cols: Vec<(u16, Vec<Value>)> = self
6395 .schema
6396 .columns
6397 .iter()
6398 .map(|c| (c.id, Vec::with_capacity(rows.len())))
6399 .collect();
6400 for r in &rows {
6401 for (cid, vec) in cols.iter_mut() {
6402 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
6403 }
6404 }
6405 Ok(cols)
6406 }
6407
6408 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
6410 let row_id = self.hot.get(key)?;
6411 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
6412 Some(row_id)
6413 } else {
6414 None
6415 }
6416 }
6417
6418 pub fn lookup_metrics_snapshot(&self) -> LookupMetricsSnapshot {
6422 let mut snap = self.lookup_metrics.snapshot();
6423 let (mem, disk, miss, write_us) = self.result_cache.lock().cache_counters();
6424 snap.result_cache_memory_hit = mem;
6425 snap.result_cache_disk_hit = disk;
6426 snap.result_cache_miss = miss;
6427 snap.result_cache_persistent_write_us = write_us;
6428 snap
6429 }
6430
6431 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
6436 self.query_at_with_allowed(q, self.snapshot(), None)
6437 }
6438
6439 pub fn query_controlled(
6442 &mut self,
6443 q: &crate::query::Query,
6444 control: &crate::ExecutionControl,
6445 ) -> Result<Vec<Row>> {
6446 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
6447 }
6448
6449 pub fn query_at_with_allowed(
6452 &mut self,
6453 q: &crate::query::Query,
6454 snapshot: Snapshot,
6455 allowed: Option<&std::collections::HashSet<RowId>>,
6456 ) -> Result<Vec<Row>> {
6457 self.query_at_with_allowed_after(q, snapshot, allowed, None)
6458 }
6459
6460 #[doc(hidden)]
6461 pub fn query_at_with_allowed_controlled(
6462 &mut self,
6463 q: &crate::query::Query,
6464 snapshot: Snapshot,
6465 allowed: Option<&std::collections::HashSet<RowId>>,
6466 control: &crate::ExecutionControl,
6467 ) -> Result<Vec<Row>> {
6468 self.require_select()?;
6469 self.ensure_indexes_complete_controlled(control, || true)?;
6470 self.validate_native_query(q)?;
6471 self.query_conditions_at(
6472 &q.conditions,
6473 snapshot,
6474 allowed,
6475 q.limit,
6476 q.offset,
6477 None,
6478 unix_nanos_now(),
6479 Some(control),
6480 )
6481 }
6482
6483 #[doc(hidden)]
6484 pub fn query_at_with_allowed_after(
6485 &mut self,
6486 q: &crate::query::Query,
6487 snapshot: Snapshot,
6488 allowed: Option<&std::collections::HashSet<RowId>>,
6489 after_row_id: Option<RowId>,
6490 ) -> Result<Vec<Row>> {
6491 self.query_at_with_allowed_after_at_time(
6492 q,
6493 snapshot,
6494 allowed,
6495 after_row_id,
6496 unix_nanos_now(),
6497 )
6498 }
6499
6500 #[doc(hidden)]
6501 pub fn query_at_with_allowed_after_at_time(
6502 &mut self,
6503 q: &crate::query::Query,
6504 snapshot: Snapshot,
6505 allowed: Option<&std::collections::HashSet<RowId>>,
6506 after_row_id: Option<RowId>,
6507 query_time_nanos: i64,
6508 ) -> Result<Vec<Row>> {
6509 self.require_select()?;
6510 self.ensure_indexes_complete()?;
6511 self.validate_native_query(q)?;
6512 self.query_conditions_at(
6513 &q.conditions,
6514 snapshot,
6515 allowed,
6516 q.limit,
6517 q.offset,
6518 after_row_id,
6519 query_time_nanos,
6520 None,
6521 )
6522 }
6523
6524 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
6525 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
6526 return Err(MongrelError::InvalidArgument(format!(
6527 "query exceeds {} conditions",
6528 crate::query::MAX_HARD_CONDITIONS
6529 )));
6530 }
6531 if let Some(limit) = q.limit {
6532 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
6533 return Err(MongrelError::InvalidArgument(format!(
6534 "query limit must be between 1 and {}",
6535 crate::query::MAX_FINAL_LIMIT
6536 )));
6537 }
6538 }
6539 if q.offset > crate::query::MAX_QUERY_OFFSET {
6540 return Err(MongrelError::InvalidArgument(format!(
6541 "query offset exceeds {}",
6542 crate::query::MAX_QUERY_OFFSET
6543 )));
6544 }
6545 Ok(())
6546 }
6547
6548 #[doc(hidden)]
6551 pub fn query_all_at(
6552 &mut self,
6553 conditions: &[crate::query::Condition],
6554 snapshot: Snapshot,
6555 ) -> Result<Vec<Row>> {
6556 self.require_select()?;
6557 self.ensure_indexes_complete()?;
6558 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
6559 return Err(MongrelError::InvalidArgument(format!(
6560 "query exceeds {} conditions",
6561 crate::query::MAX_HARD_CONDITIONS
6562 )));
6563 }
6564 self.query_conditions_at(
6565 conditions,
6566 snapshot,
6567 None,
6568 None,
6569 0,
6570 None,
6571 unix_nanos_now(),
6572 None,
6573 )
6574 }
6575
6576 #[allow(clippy::too_many_arguments)]
6577 fn query_conditions_at(
6578 &self,
6579 conditions: &[crate::query::Condition],
6580 snapshot: Snapshot,
6581 allowed: Option<&std::collections::HashSet<RowId>>,
6582 limit: Option<usize>,
6583 offset: usize,
6584 after_row_id: Option<RowId>,
6585 query_time_nanos: i64,
6586 control: Option<&crate::ExecutionControl>,
6587 ) -> Result<Vec<Row>> {
6588 control
6589 .map(crate::ExecutionControl::checkpoint)
6590 .transpose()?;
6591 crate::trace::QueryTrace::record(|t| {
6592 t.run_count = self.run_refs.len();
6593 t.memtable_rows = self.memtable.len();
6594 t.mutable_run_rows = self.mutable_run.len();
6595 });
6596 if conditions.is_empty() {
6600 crate::trace::QueryTrace::record(|t| {
6601 t.scan_mode = crate::trace::ScanMode::Materialized;
6602 t.row_materialized = true;
6603 });
6604 let mut rows = match control {
6605 Some(control) => self.visible_rows_controlled(snapshot, control)?,
6606 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
6607 };
6608 if let Some(allowed) = allowed {
6609 let mut filtered = Vec::with_capacity(rows.len());
6610 for (index, row) in rows.into_iter().enumerate() {
6611 if index & 255 == 0 {
6612 control
6613 .map(crate::ExecutionControl::checkpoint)
6614 .transpose()?;
6615 }
6616 if allowed.contains(&row.row_id) {
6617 filtered.push(row);
6618 }
6619 }
6620 rows = filtered;
6621 }
6622 if let Some(after_row_id) = after_row_id {
6623 rows.retain(|row| row.row_id > after_row_id);
6624 }
6625 rows.drain(..offset.min(rows.len()));
6626 if let Some(limit) = limit {
6627 rows.truncate(limit);
6628 }
6629 return Ok(rows);
6630 }
6631 crate::trace::QueryTrace::record(|t| {
6632 t.conditions_pushed = conditions.len();
6633 t.scan_mode = crate::trace::ScanMode::Materialized;
6634 t.row_materialized = true;
6635 });
6636 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
6643 ordered.sort_by_key(|c| condition_cost_rank(c));
6644 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
6645 for c in &ordered {
6646 control
6647 .map(crate::ExecutionControl::checkpoint)
6648 .transpose()?;
6649 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
6650 let empty = s.is_empty();
6651 sets.push(s);
6652 if empty {
6653 break;
6654 }
6655 }
6656 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
6657 if let Some(allowed) = allowed {
6658 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
6659 }
6660 if let Some(after_row_id) = after_row_id {
6661 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
6662 rids.drain(..first);
6663 }
6664 rids.drain(..offset.min(rids.len()));
6665 if let Some(limit) = limit {
6666 rids.truncate(limit);
6667 }
6668 control
6669 .map(crate::ExecutionControl::checkpoint)
6670 .transpose()?;
6671 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
6672 }
6673
6674 pub fn retrieve(
6676 &mut self,
6677 retriever: &crate::query::Retriever,
6678 ) -> Result<Vec<crate::query::RetrieverHit>> {
6679 self.retrieve_with_allowed(retriever, None)
6680 }
6681
6682 pub fn retrieve_at(
6683 &mut self,
6684 retriever: &crate::query::Retriever,
6685 snapshot: Snapshot,
6686 allowed: Option<&std::collections::HashSet<RowId>>,
6687 ) -> Result<Vec<crate::query::RetrieverHit>> {
6688 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
6689 }
6690
6691 pub fn retrieve_with_allowed(
6694 &mut self,
6695 retriever: &crate::query::Retriever,
6696 allowed: Option<&std::collections::HashSet<RowId>>,
6697 ) -> Result<Vec<crate::query::RetrieverHit>> {
6698 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
6699 }
6700
6701 pub fn retrieve_at_with_allowed(
6702 &mut self,
6703 retriever: &crate::query::Retriever,
6704 snapshot: Snapshot,
6705 allowed: Option<&std::collections::HashSet<RowId>>,
6706 ) -> Result<Vec<crate::query::RetrieverHit>> {
6707 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
6708 }
6709
6710 pub fn retrieve_at_with_allowed_and_context(
6711 &mut self,
6712 retriever: &crate::query::Retriever,
6713 snapshot: Snapshot,
6714 allowed: Option<&std::collections::HashSet<RowId>>,
6715 context: Option<&crate::query::AiExecutionContext>,
6716 ) -> Result<Vec<crate::query::RetrieverHit>> {
6717 self.require_select()?;
6718 self.ensure_indexes_complete()?;
6719 self.validate_retriever(retriever)?;
6720 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
6721 }
6722
6723 pub fn retrieve_at_with_candidate_authorization_and_context(
6724 &mut self,
6725 retriever: &crate::query::Retriever,
6726 snapshot: Snapshot,
6727 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6728 context: Option<&crate::query::AiExecutionContext>,
6729 ) -> Result<Vec<crate::query::RetrieverHit>> {
6730 self.require_select()?;
6731 self.ensure_indexes_complete()?;
6732 self.retrieve_at_with_candidate_authorization_on_generation(
6733 retriever,
6734 snapshot,
6735 authorization,
6736 context,
6737 )
6738 }
6739
6740 #[doc(hidden)]
6741 pub fn retrieve_at_with_candidate_authorization_on_generation(
6742 &self,
6743 retriever: &crate::query::Retriever,
6744 snapshot: Snapshot,
6745 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6746 context: Option<&crate::query::AiExecutionContext>,
6747 ) -> Result<Vec<crate::query::RetrieverHit>> {
6748 self.require_select()?;
6749 self.validate_retriever(retriever)?;
6750 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
6751 }
6752
6753 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
6754 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
6755 let (column_id, k) = match retriever {
6756 Retriever::Ann {
6757 column_id,
6758 query,
6759 k,
6760 } => {
6761 let index = self.ann.get(column_id).ok_or_else(|| {
6762 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
6763 })?;
6764 if query.len() != index.dim() {
6765 return Err(MongrelError::InvalidArgument(format!(
6766 "ANN query dimension must be {}, got {}",
6767 index.dim(),
6768 query.len()
6769 )));
6770 }
6771 if query.iter().any(|value| !value.is_finite()) {
6772 return Err(MongrelError::InvalidArgument(
6773 "ANN query values must be finite".into(),
6774 ));
6775 }
6776 (*column_id, *k)
6777 }
6778 Retriever::Sparse {
6779 column_id,
6780 query,
6781 k,
6782 } => {
6783 if !self.sparse.contains_key(column_id) {
6784 return Err(MongrelError::InvalidArgument(format!(
6785 "column {column_id} has no Sparse index"
6786 )));
6787 }
6788 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
6789 return Err(MongrelError::InvalidArgument(
6790 "Sparse query must be non-empty with finite weights".into(),
6791 ));
6792 }
6793 if query.len() > MAX_SPARSE_TERMS {
6794 return Err(MongrelError::InvalidArgument(format!(
6795 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
6796 )));
6797 }
6798 (*column_id, *k)
6799 }
6800 Retriever::MinHash {
6801 column_id,
6802 members,
6803 k,
6804 } => {
6805 if !self.minhash.contains_key(column_id) {
6806 return Err(MongrelError::InvalidArgument(format!(
6807 "column {column_id} has no MinHash index"
6808 )));
6809 }
6810 if members.is_empty() {
6811 return Err(MongrelError::InvalidArgument(
6812 "MinHash members must not be empty".into(),
6813 ));
6814 }
6815 if members.len() > MAX_SET_MEMBERS {
6816 return Err(MongrelError::InvalidArgument(format!(
6817 "MinHash query exceeds {MAX_SET_MEMBERS} members"
6818 )));
6819 }
6820 let mut total_bytes = 0usize;
6821 for member in members {
6822 let bytes = member.encoded_len();
6823 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
6824 return Err(MongrelError::InvalidArgument(format!(
6825 "MinHash member exceeds {} bytes",
6826 crate::query::MAX_SET_MEMBER_BYTES
6827 )));
6828 }
6829 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
6830 MongrelError::InvalidArgument("MinHash input size overflow".into())
6831 })?;
6832 }
6833 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
6834 return Err(MongrelError::InvalidArgument(format!(
6835 "MinHash input exceeds {} bytes",
6836 crate::query::MAX_SET_INPUT_BYTES
6837 )));
6838 }
6839 (*column_id, *k)
6840 }
6841 };
6842 if k == 0 {
6843 return Err(MongrelError::InvalidArgument(
6844 "retriever k must be > 0".into(),
6845 ));
6846 }
6847 if k > MAX_RETRIEVER_K {
6848 return Err(MongrelError::InvalidArgument(format!(
6849 "retriever k exceeds {MAX_RETRIEVER_K}"
6850 )));
6851 }
6852 debug_assert!(self
6853 .schema
6854 .columns
6855 .iter()
6856 .any(|column| column.id == column_id));
6857 Ok(())
6858 }
6859
6860 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
6861 use crate::query::Condition;
6862 match condition {
6863 Condition::Ann {
6864 column_id,
6865 query,
6866 k,
6867 } => self.validate_retriever(&crate::query::Retriever::Ann {
6868 column_id: *column_id,
6869 query: query.clone(),
6870 k: *k,
6871 }),
6872 Condition::SparseMatch {
6873 column_id,
6874 query,
6875 k,
6876 } => self.validate_retriever(&crate::query::Retriever::Sparse {
6877 column_id: *column_id,
6878 query: query.clone(),
6879 k: *k,
6880 }),
6881 Condition::MinHashSimilar {
6882 column_id,
6883 query,
6884 k,
6885 } => {
6886 if !self.minhash.contains_key(column_id) {
6887 return Err(MongrelError::InvalidArgument(format!(
6888 "column {column_id} has no MinHash index"
6889 )));
6890 }
6891 if query.is_empty() || *k == 0 {
6892 return Err(MongrelError::InvalidArgument(
6893 "MinHash query must be non-empty and k must be > 0".into(),
6894 ));
6895 }
6896 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
6897 {
6898 return Err(MongrelError::InvalidArgument(format!(
6899 "MinHash query must have <= {} members and k <= {}",
6900 crate::query::MAX_SET_MEMBERS,
6901 crate::query::MAX_RETRIEVER_K
6902 )));
6903 }
6904 Ok(())
6905 }
6906 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
6907 Err(MongrelError::InvalidArgument(format!(
6908 "bitmap IN exceeds {} values",
6909 crate::query::MAX_SET_MEMBERS
6910 )))
6911 }
6912 Condition::FmContainsAll { patterns, .. }
6913 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
6914 {
6915 Err(MongrelError::InvalidArgument(format!(
6916 "FM query exceeds {} patterns",
6917 crate::query::MAX_HARD_CONDITIONS
6918 )))
6919 }
6920 _ => Ok(()),
6921 }
6922 }
6923
6924 fn retrieve_filtered(
6925 &self,
6926 retriever: &crate::query::Retriever,
6927 snapshot: Snapshot,
6928 hard_filter: Option<&RowIdSet>,
6929 allowed: Option<&std::collections::HashSet<RowId>>,
6930 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
6931 context: Option<&crate::query::AiExecutionContext>,
6932 ) -> Result<Vec<crate::query::RetrieverHit>> {
6933 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
6934 let started = std::time::Instant::now();
6935 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
6936 Retriever::Ann {
6937 column_id,
6938 query,
6939 k,
6940 } => {
6941 let Some(index) = self.ann.get(column_id) else {
6942 return Ok(Vec::new());
6943 };
6944 let cap = ann_candidate_cap(index.len(), context);
6945 if cap == 0 {
6946 return Ok(Vec::new());
6947 }
6948 let mut breadth = (*k).max(1).min(cap);
6949 let mut eligibility = std::collections::HashMap::new();
6950 let mut filtered = loop {
6951 let mut seen = std::collections::HashSet::new();
6952 if let Some(context) = context {
6953 context.checkpoint()?;
6954 }
6955 let raw = index.search_with_context(query, breadth, context)?;
6956 let unchecked: Vec<_> = raw
6957 .iter()
6958 .map(|(row_id, _)| *row_id)
6959 .filter(|row_id| !eligibility.contains_key(row_id))
6960 .filter(|row_id| {
6961 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
6962 && allowed.is_none_or(|allowed| allowed.contains(row_id))
6963 })
6964 .collect();
6965 let eligible = self.eligible_and_authorized_candidate_ids(
6966 &unchecked,
6967 *column_id,
6968 snapshot,
6969 candidate_authorization,
6970 context,
6971 )?;
6972 for row_id in unchecked {
6973 eligibility.insert(row_id, eligible.contains(&row_id));
6974 }
6975 let filtered: Vec<_> = raw
6976 .into_iter()
6977 .filter(|(row_id, _)| {
6978 seen.insert(*row_id)
6979 && eligibility.get(row_id).copied().unwrap_or(false)
6980 })
6981 .map(|(row_id, score)| {
6982 let score = match score {
6983 crate::index::AnnDistance::Hamming(d) => {
6984 RetrieverScore::AnnHammingDistance(d)
6985 }
6986 crate::index::AnnDistance::Cosine(d) => {
6987 RetrieverScore::AnnCosineDistance(d)
6988 }
6989 };
6990 (row_id, score)
6991 })
6992 .collect();
6993 if filtered.len() >= *k || breadth >= cap {
6994 if filtered.len() < *k && index.len() > cap && breadth >= cap {
6995 crate::trace::QueryTrace::record(|trace| {
6996 trace.ann_candidate_cap_hit = true;
6997 });
6998 }
6999 break filtered;
7000 }
7001 breadth = breadth.saturating_mul(2).min(cap);
7002 };
7003 filtered.truncate(*k);
7004 filtered
7005 }
7006 Retriever::Sparse {
7007 column_id,
7008 query,
7009 k,
7010 } => self
7011 .sparse
7012 .get(column_id)
7013 .map(|index| -> Result<Vec<_>> {
7014 let mut breadth = (*k).max(1);
7015 let mut eligibility = std::collections::HashMap::new();
7016 loop {
7017 if let Some(context) = context {
7018 context.checkpoint()?;
7019 }
7020 let raw = index.search_with_context(query, breadth, context)?;
7021 let unchecked: Vec<_> = raw
7022 .iter()
7023 .map(|(row_id, _)| *row_id)
7024 .filter(|row_id| !eligibility.contains_key(row_id))
7025 .filter(|row_id| {
7026 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7027 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7028 })
7029 .collect();
7030 let eligible = self.eligible_and_authorized_candidate_ids(
7031 &unchecked,
7032 *column_id,
7033 snapshot,
7034 candidate_authorization,
7035 context,
7036 )?;
7037 for row_id in unchecked {
7038 eligibility.insert(row_id, eligible.contains(&row_id));
7039 }
7040 let filtered: Vec<_> = raw
7041 .iter()
7042 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
7043 .take(*k)
7044 .map(|(row_id, score)| {
7045 (*row_id, RetrieverScore::SparseDotProduct(*score))
7046 })
7047 .collect();
7048 if filtered.len() >= *k || raw.len() < breadth {
7049 break Ok(filtered);
7050 }
7051 let next = breadth.saturating_mul(2);
7052 if next == breadth {
7053 break Ok(filtered);
7054 }
7055 breadth = next;
7056 }
7057 })
7058 .transpose()?
7059 .unwrap_or_default(),
7060 Retriever::MinHash {
7061 column_id,
7062 members,
7063 k,
7064 } => self
7065 .minhash
7066 .get(column_id)
7067 .map(|index| -> Result<Vec<_>> {
7068 let mut hashes = Vec::with_capacity(members.len());
7069 for member in members {
7070 if let Some(context) = context {
7071 context.consume(crate::query::work_units(
7072 member.encoded_len(),
7073 crate::query::PARSE_WORK_QUANTUM,
7074 ))?;
7075 }
7076 hashes.push(member.hash_v1());
7077 }
7078 let mut breadth = (*k).max(1);
7079 let mut eligibility = std::collections::HashMap::new();
7080 loop {
7081 if let Some(context) = context {
7082 context.checkpoint()?;
7083 }
7084 let raw = index.search_with_context(&hashes, breadth, context)?;
7085 let unchecked: Vec<_> = raw
7086 .iter()
7087 .map(|(row_id, _)| *row_id)
7088 .filter(|row_id| !eligibility.contains_key(row_id))
7089 .filter(|row_id| {
7090 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7091 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7092 })
7093 .collect();
7094 let eligible = self.eligible_and_authorized_candidate_ids(
7095 &unchecked,
7096 *column_id,
7097 snapshot,
7098 candidate_authorization,
7099 context,
7100 )?;
7101 for row_id in unchecked {
7102 eligibility.insert(row_id, eligible.contains(&row_id));
7103 }
7104 let filtered: Vec<_> = raw
7105 .iter()
7106 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
7107 .take(*k)
7108 .map(|(row_id, score)| {
7109 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
7110 })
7111 .collect();
7112 if filtered.len() >= *k || raw.len() < breadth {
7113 break Ok(filtered);
7114 }
7115 let next = breadth.saturating_mul(2);
7116 if next == breadth {
7117 break Ok(filtered);
7118 }
7119 breadth = next;
7120 }
7121 })
7122 .transpose()?
7123 .unwrap_or_default(),
7124 };
7125 let elapsed = started.elapsed().as_nanos() as u64;
7126 crate::trace::QueryTrace::record(|trace| {
7127 match retriever {
7128 Retriever::Ann { .. } => {
7129 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
7130 if let Retriever::Ann { column_id, .. } = retriever {
7131 if let Some(index) = self.ann.get(column_id) {
7132 trace.ann_algorithm = Some(index.algorithm());
7133 trace.ann_quantization = Some(index.quantization());
7134 trace.ann_backend = Some(index.backend_name());
7135 }
7136 }
7137 }
7138 Retriever::Sparse { .. } => {
7139 trace.sparse_candidate_nanos =
7140 trace.sparse_candidate_nanos.saturating_add(elapsed)
7141 }
7142 Retriever::MinHash { .. } => {
7143 trace.minhash_candidate_nanos =
7144 trace.minhash_candidate_nanos.saturating_add(elapsed)
7145 }
7146 }
7147 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
7148 });
7149 Ok(scored
7150 .into_iter()
7151 .enumerate()
7152 .map(|(rank, (row_id, score))| RetrieverHit {
7153 row_id,
7154 rank: rank + 1,
7155 score,
7156 })
7157 .collect())
7158 }
7159
7160 fn eligible_candidate_ids(
7161 &self,
7162 candidates: &[RowId],
7163 _column_id: u16,
7164 snapshot: Snapshot,
7165 context: Option<&crate::query::AiExecutionContext>,
7166 ) -> Result<std::collections::HashSet<RowId>> {
7167 let lookup_snapshot = if self.is_shared()
7177 || (self.pending_put_cols.is_empty() && self.pending_delete_rids.is_empty())
7178 || snapshot.epoch.0 >= self.pending_epoch().0
7179 {
7180 snapshot
7181 } else {
7182 Snapshot::at(self.pending_epoch())
7183 };
7184 if !self.had_deletes && self.ttl.is_none() && lookup_snapshot.epoch == self.snapshot().epoch
7185 {
7186 return Ok(candidates.iter().copied().collect());
7187 }
7188 let mut readers: Vec<_> = self
7189 .run_refs
7190 .iter()
7191 .map(|run| self.open_reader(run.run_id))
7192 .collect::<Result<_>>()?;
7193 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7194 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
7195 for &row_id in candidates {
7196 if let Some(context) = context {
7197 context.consume(1)?;
7198 }
7199 let mem = self.memtable.get_version_at(row_id, lookup_snapshot);
7200 let mutable = self.mutable_run.get_version_at(row_id, lookup_snapshot);
7201 let overlay = match (mem, mutable) {
7202 (Some(left), Some(right)) => Some(
7203 if Snapshot::version_is_newer(
7204 left.1.committed_epoch,
7205 left.1.commit_ts,
7206 right.1.committed_epoch,
7207 right.1.commit_ts,
7208 ) {
7209 left
7210 } else {
7211 right
7212 },
7213 ),
7214 (Some(value), None) | (None, Some(value)) => Some(value),
7215 (None, None) => None,
7216 };
7217 if let Some((_, row)) = overlay {
7218 if !row.deleted && !self.row_expired_at(&row, now) {
7219 eligible.insert(row_id);
7220 }
7221 continue;
7222 }
7223 let mut best: Option<(Epoch, bool, usize)> = None;
7224 for (index, reader) in readers.iter_mut().enumerate() {
7225 if let Some((epoch, deleted)) =
7226 reader.get_version_visibility(row_id, snapshot.epoch)?
7227 {
7228 if best
7229 .as_ref()
7230 .map(|(best_epoch, ..)| epoch > *best_epoch)
7231 .unwrap_or(true)
7232 {
7233 best = Some((epoch, deleted, index));
7234 }
7235 }
7236 }
7237 let Some((_, false, reader_index)) = best else {
7238 continue;
7239 };
7240 if let Some(ttl) = self.ttl {
7241 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7242 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
7243 {
7244 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7245 continue;
7246 }
7247 }
7248 }
7249 eligible.insert(row_id);
7250 }
7251 Ok(eligible)
7252 }
7253
7254 fn eligible_and_authorized_candidate_ids(
7255 &self,
7256 candidates: &[RowId],
7257 column_id: u16,
7258 snapshot: Snapshot,
7259 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7260 context: Option<&crate::query::AiExecutionContext>,
7261 ) -> Result<std::collections::HashSet<RowId>> {
7262 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
7263 let Some(authorization) = authorization else {
7264 return Ok(eligible);
7265 };
7266 let candidates: Vec<_> = eligible.into_iter().collect();
7267 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
7268 }
7269
7270 fn policy_allowed_candidate_ids(
7271 &self,
7272 candidates: &[RowId],
7273 snapshot: Snapshot,
7274 authorization: &crate::security::CandidateAuthorization<'_>,
7275 context: Option<&crate::query::AiExecutionContext>,
7276 ) -> Result<std::collections::HashSet<RowId>> {
7277 let started = std::time::Instant::now();
7278 if candidates.is_empty()
7279 || authorization.principal.is_admin
7280 || !authorization.security.rls_enabled(authorization.table)
7281 {
7282 return Ok(candidates.iter().copied().collect());
7283 }
7284 if let Some(context) = context {
7285 context.checkpoint()?;
7286 }
7287 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
7288 let mut rows: std::collections::HashMap<RowId, Row> = candidates
7289 .iter()
7290 .map(|row_id| {
7291 (
7292 *row_id,
7293 Row {
7294 row_id: *row_id,
7295 committed_epoch: snapshot.epoch,
7296 columns: std::collections::HashMap::new(),
7297 deleted: false,
7298 commit_ts: None,
7299 },
7300 )
7301 })
7302 .collect();
7303 let columns = authorization
7304 .security
7305 .select_policy_columns(authorization.table, authorization.principal);
7306 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7307 let mut decoded = 0usize;
7308 for column_id in &columns {
7309 if let Some(context) = context {
7310 context.checkpoint()?;
7311 }
7312 for (row_id, value) in self.values_for_rids_batch_at_with_context(
7313 &row_ids, *column_id, snapshot, query_now, context,
7314 )? {
7315 if let Some(row) = rows.get_mut(&row_id) {
7316 row.columns.insert(*column_id, value);
7317 decoded = decoded.saturating_add(1);
7318 }
7319 }
7320 }
7321 if let Some(context) = context {
7322 context.consume(candidates.len().saturating_add(decoded))?;
7323 }
7324 let allowed = rows
7325 .into_values()
7326 .filter_map(|row| {
7327 authorization
7328 .security
7329 .row_allowed(
7330 authorization.table,
7331 crate::security::PolicyCommand::Select,
7332 &row,
7333 authorization.principal,
7334 false,
7335 )
7336 .then_some(row.row_id)
7337 })
7338 .collect();
7339 crate::trace::QueryTrace::record(|trace| {
7340 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
7341 trace.rls_policy_columns_decoded =
7342 trace.rls_policy_columns_decoded.saturating_add(decoded);
7343 trace.authorization_nanos = trace
7344 .authorization_nanos
7345 .saturating_add(started.elapsed().as_nanos() as u64);
7346 });
7347 Ok(allowed)
7348 }
7349
7350 pub fn search(
7352 &mut self,
7353 request: &crate::query::SearchRequest,
7354 ) -> Result<Vec<crate::query::SearchHit>> {
7355 self.search_with_allowed(request, None)
7356 }
7357
7358 pub fn search_at(
7359 &mut self,
7360 request: &crate::query::SearchRequest,
7361 snapshot: Snapshot,
7362 authorized: Option<&std::collections::HashSet<RowId>>,
7363 ) -> Result<Vec<crate::query::SearchHit>> {
7364 self.search_at_with_allowed(request, snapshot, authorized)
7365 }
7366
7367 pub fn search_with_allowed(
7368 &mut self,
7369 request: &crate::query::SearchRequest,
7370 authorized: Option<&std::collections::HashSet<RowId>>,
7371 ) -> Result<Vec<crate::query::SearchHit>> {
7372 self.search_at_with_allowed(request, self.snapshot(), authorized)
7373 }
7374
7375 pub fn search_at_with_allowed(
7376 &mut self,
7377 request: &crate::query::SearchRequest,
7378 snapshot: Snapshot,
7379 authorized: Option<&std::collections::HashSet<RowId>>,
7380 ) -> Result<Vec<crate::query::SearchHit>> {
7381 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
7382 }
7383
7384 pub fn search_at_with_allowed_and_context(
7385 &mut self,
7386 request: &crate::query::SearchRequest,
7387 snapshot: Snapshot,
7388 authorized: Option<&std::collections::HashSet<RowId>>,
7389 context: Option<&crate::query::AiExecutionContext>,
7390 ) -> Result<Vec<crate::query::SearchHit>> {
7391 self.ensure_indexes_complete()?;
7392 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
7393 }
7394
7395 pub fn search_at_with_candidate_authorization_and_context(
7396 &mut self,
7397 request: &crate::query::SearchRequest,
7398 snapshot: Snapshot,
7399 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7400 context: Option<&crate::query::AiExecutionContext>,
7401 ) -> Result<Vec<crate::query::SearchHit>> {
7402 self.ensure_indexes_complete()?;
7403 self.search_at_with_filters_and_context(
7404 request,
7405 snapshot,
7406 None,
7407 authorization,
7408 context,
7409 None,
7410 )
7411 }
7412
7413 #[doc(hidden)]
7414 pub fn search_at_with_candidate_authorization_on_generation(
7415 &self,
7416 request: &crate::query::SearchRequest,
7417 snapshot: Snapshot,
7418 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7419 context: Option<&crate::query::AiExecutionContext>,
7420 ) -> Result<Vec<crate::query::SearchHit>> {
7421 self.search_at_with_filters_and_context(
7422 request,
7423 snapshot,
7424 None,
7425 authorization,
7426 context,
7427 None,
7428 )
7429 }
7430
7431 #[doc(hidden)]
7432 pub fn search_at_with_candidate_authorization_on_generation_after(
7433 &self,
7434 request: &crate::query::SearchRequest,
7435 snapshot: Snapshot,
7436 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7437 context: Option<&crate::query::AiExecutionContext>,
7438 after: Option<crate::query::SearchAfter>,
7439 ) -> Result<Vec<crate::query::SearchHit>> {
7440 self.search_at_with_filters_and_context(
7441 request,
7442 snapshot,
7443 None,
7444 authorization,
7445 context,
7446 after,
7447 )
7448 }
7449
7450 fn search_at_with_filters_and_context(
7451 &self,
7452 request: &crate::query::SearchRequest,
7453 snapshot: Snapshot,
7454 authorized: Option<&std::collections::HashSet<RowId>>,
7455 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7456 context: Option<&crate::query::AiExecutionContext>,
7457 after: Option<crate::query::SearchAfter>,
7458 ) -> Result<Vec<crate::query::SearchHit>> {
7459 use crate::query::{
7460 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
7461 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
7462 };
7463 let total_started = std::time::Instant::now();
7464 let rank_offset = after.map_or(0, |after| after.returned_count);
7465 self.require_select()?;
7466 if request.limit == 0 {
7467 return Err(MongrelError::InvalidArgument(
7468 "search limit must be > 0".into(),
7469 ));
7470 }
7471 if request.limit > MAX_FINAL_LIMIT {
7472 return Err(MongrelError::InvalidArgument(format!(
7473 "search limit exceeds {MAX_FINAL_LIMIT}"
7474 )));
7475 }
7476 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
7477 return Err(MongrelError::InvalidArgument(
7478 "search-after score must be finite".into(),
7479 ));
7480 }
7481 if request.retrievers.is_empty() {
7482 return Err(MongrelError::InvalidArgument(
7483 "search requires at least one retriever".into(),
7484 ));
7485 }
7486 if request.retrievers.len() > MAX_RETRIEVERS {
7487 return Err(MongrelError::InvalidArgument(format!(
7488 "search exceeds {MAX_RETRIEVERS} retrievers"
7489 )));
7490 }
7491 if request.must.len() > MAX_HARD_CONDITIONS {
7492 return Err(MongrelError::InvalidArgument(format!(
7493 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
7494 )));
7495 }
7496 for condition in &request.must {
7497 self.validate_condition(condition)?;
7498 }
7499 if request.must.iter().any(|condition| {
7500 matches!(
7501 condition,
7502 Condition::Ann { .. }
7503 | Condition::SparseMatch { .. }
7504 | Condition::MinHashSimilar { .. }
7505 )
7506 }) {
7507 return Err(MongrelError::InvalidArgument(
7508 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
7509 .into(),
7510 ));
7511 }
7512 let mut names = std::collections::HashSet::new();
7513 for named in &request.retrievers {
7514 if named.name.is_empty()
7515 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
7516 || !names.insert(named.name.as_str())
7517 {
7518 return Err(MongrelError::InvalidArgument(format!(
7519 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
7520 crate::query::MAX_RETRIEVER_NAME_BYTES
7521 )));
7522 }
7523 if !named.weight.is_finite()
7524 || named.weight < 0.0
7525 || named.weight > MAX_RETRIEVER_WEIGHT
7526 {
7527 return Err(MongrelError::InvalidArgument(format!(
7528 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
7529 )));
7530 }
7531 self.validate_retriever(&named.retriever)?;
7532 }
7533 let projection = request
7534 .projection
7535 .clone()
7536 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
7537 if projection.len() > MAX_PROJECTION_COLUMNS {
7538 return Err(MongrelError::InvalidArgument(format!(
7539 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
7540 )));
7541 }
7542 for column_id in &projection {
7543 if !self
7544 .schema
7545 .columns
7546 .iter()
7547 .any(|column| column.id == *column_id)
7548 {
7549 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
7550 }
7551 }
7552 if let Some(crate::query::Rerank::ExactVector {
7553 embedding_column,
7554 query,
7555 candidate_limit,
7556 weight,
7557 ..
7558 }) = &request.rerank
7559 {
7560 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
7561 {
7562 return Err(MongrelError::InvalidArgument(format!(
7563 "rerank candidate_limit must be between search limit and {}",
7564 crate::query::MAX_RETRIEVER_K
7565 )));
7566 }
7567 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
7568 return Err(MongrelError::InvalidArgument(format!(
7569 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
7570 )));
7571 }
7572 let column = self
7573 .schema
7574 .columns
7575 .iter()
7576 .find(|column| column.id == *embedding_column)
7577 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
7578 let crate::schema::TypeId::Embedding { dim } = column.ty else {
7579 return Err(MongrelError::InvalidArgument(format!(
7580 "rerank column {embedding_column} is not an embedding"
7581 )));
7582 };
7583 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
7584 return Err(MongrelError::InvalidArgument(format!(
7585 "rerank query must contain {dim} finite values"
7586 )));
7587 }
7588 }
7589
7590 let hard_filter_started = std::time::Instant::now();
7591 let hard_filter = if request.must.is_empty() {
7592 None
7593 } else {
7594 let mut sets = Vec::with_capacity(request.must.len());
7595 for condition in &request.must {
7596 if let Some(context) = context {
7597 context.checkpoint()?;
7598 }
7599 sets.push(self.resolve_condition(condition, snapshot)?);
7600 }
7601 Some(RowIdSet::intersect_many(sets))
7602 };
7603 crate::trace::QueryTrace::record(|trace| {
7604 trace.hard_filter_nanos = trace
7605 .hard_filter_nanos
7606 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
7607 });
7608 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
7609 return Ok(Vec::new());
7610 }
7611
7612 let constant = match request.fusion {
7613 Fusion::ReciprocalRank { constant } => constant,
7614 };
7615 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
7616 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
7617 let mut fusion_nanos = 0u64;
7618 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
7619 std::collections::HashMap::new();
7620 for named in retrievers {
7621 if named.weight == 0.0 {
7622 continue;
7623 }
7624 if let Some(context) = context {
7625 context.checkpoint()?;
7626 }
7627 let hits = self.retrieve_filtered(
7628 &named.retriever,
7629 snapshot,
7630 hard_filter.as_ref(),
7631 authorized,
7632 candidate_authorization,
7633 context,
7634 )?;
7635 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
7636 let fusion_started = std::time::Instant::now();
7637 for hit in hits {
7638 if let Some(context) = context {
7639 context.consume(1)?;
7640 }
7641 let contribution = named.weight / (constant as f64 + hit.rank as f64);
7642 if !contribution.is_finite() {
7643 return Err(MongrelError::InvalidArgument(
7644 "retriever contribution must be finite".into(),
7645 ));
7646 }
7647 let max_fused_candidates = context.map_or(
7648 crate::query::MAX_FUSED_CANDIDATES,
7649 crate::query::AiExecutionContext::max_fused_candidates,
7650 );
7651 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
7652 return Err(MongrelError::WorkBudgetExceeded);
7653 }
7654 let entry = fused.entry(hit.row_id).or_default();
7655 entry.0 += contribution;
7656 if !entry.0.is_finite() {
7657 return Err(MongrelError::InvalidArgument(
7658 "fused score must be finite".into(),
7659 ));
7660 }
7661 entry.1.push(ComponentScore {
7662 retriever_name: retriever_name.clone(),
7663 rank: hit.rank,
7664 raw_score: hit.score,
7665 contribution,
7666 });
7667 }
7668 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
7669 }
7670 let union_size = fused.len();
7671 let mut ranked: Vec<_> = fused
7672 .into_iter()
7673 .map(|(row_id, (fused_score, components))| {
7674 (row_id, fused_score, components, None, fused_score)
7675 })
7676 .collect();
7677 let order = |(a_row, _, _, _, a_score): &(
7678 RowId,
7679 f64,
7680 Vec<ComponentScore>,
7681 Option<f32>,
7682 f64,
7683 ),
7684 (b_row, _, _, _, b_score): &(
7685 RowId,
7686 f64,
7687 Vec<ComponentScore>,
7688 Option<f32>,
7689 f64,
7690 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
7691 if let Some(crate::query::Rerank::ExactVector {
7692 embedding_column,
7693 query,
7694 metric,
7695 candidate_limit,
7696 weight,
7697 }) = &request.rerank
7698 {
7699 let fused_order = |(a_row, a_score, ..): &(
7700 RowId,
7701 f64,
7702 Vec<ComponentScore>,
7703 Option<f32>,
7704 f64,
7705 ),
7706 (b_row, b_score, ..): &(
7707 RowId,
7708 f64,
7709 Vec<ComponentScore>,
7710 Option<f32>,
7711 f64,
7712 )| {
7713 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
7714 };
7715 let selection_started = std::time::Instant::now();
7716 if let Some(context) = context {
7717 context.consume(ranked.len())?;
7718 }
7719 if ranked.len() > *candidate_limit {
7720 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
7721 ranked.truncate(*candidate_limit);
7722 }
7723 ranked.sort_by(fused_order);
7724 fusion_nanos =
7725 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
7726 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
7727 if let Some(context) = context {
7728 context.consume(row_ids.len())?;
7729 }
7730 let query_now =
7731 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7732 let gather_started = std::time::Instant::now();
7733 let vectors = self.values_for_rids_batch_at_with_context(
7734 &row_ids,
7735 *embedding_column,
7736 snapshot,
7737 query_now,
7738 context,
7739 )?;
7740 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
7741 let vector_work =
7742 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
7743 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
7744 if let Some(context) = context {
7745 context.consume(vector_work)?;
7746 }
7747 query
7748 .iter()
7749 .map(|value| f64::from(*value).powi(2))
7750 .sum::<f64>()
7751 .sqrt()
7752 } else {
7753 0.0
7754 };
7755 let score_started = std::time::Instant::now();
7756 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
7757 for (row_id, value) in vectors {
7758 if let Some(meta) = value.generated_embedding_metadata() {
7759 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
7760 continue;
7761 }
7762 }
7763 let Some(vector) = value.as_embedding() else {
7764 continue;
7765 };
7766 let score = match metric {
7767 crate::query::VectorMetric::DotProduct => {
7768 if let Some(context) = context {
7769 context.consume(vector_work)?;
7770 }
7771 query
7772 .iter()
7773 .zip(vector)
7774 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7775 .sum::<f64>()
7776 }
7777 crate::query::VectorMetric::Cosine => {
7778 if let Some(context) = context {
7779 context.consume(vector_work.saturating_mul(2))?;
7780 }
7781 let dot = query
7782 .iter()
7783 .zip(vector)
7784 .map(|(left, right)| f64::from(*left) * f64::from(*right))
7785 .sum::<f64>();
7786 let norm = vector
7787 .iter()
7788 .map(|value| f64::from(*value).powi(2))
7789 .sum::<f64>()
7790 .sqrt();
7791 if query_norm == 0.0 || norm == 0.0 {
7792 0.0
7793 } else {
7794 dot / (query_norm * norm)
7795 }
7796 }
7797 crate::query::VectorMetric::Euclidean => {
7798 if let Some(context) = context {
7799 context.consume(vector_work)?;
7800 }
7801 query
7802 .iter()
7803 .zip(vector)
7804 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
7805 .sum::<f64>()
7806 .sqrt()
7807 }
7808 };
7809 if !score.is_finite() {
7810 return Err(MongrelError::InvalidArgument(
7811 "exact rerank score must be finite".into(),
7812 ));
7813 }
7814 scores.insert(row_id, score as f32);
7815 }
7816 let mut reranked = Vec::with_capacity(ranked.len());
7817 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
7818 let Some(score) = scores.get(&row_id).copied() else {
7819 continue;
7820 };
7821 let ordering_score = match metric {
7822 crate::query::VectorMetric::Euclidean => -f64::from(score),
7823 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
7824 f64::from(score)
7825 }
7826 };
7827 let final_score = fused_score + *weight * ordering_score;
7828 if !final_score.is_finite() {
7829 return Err(MongrelError::InvalidArgument(
7830 "final rerank score must be finite".into(),
7831 ));
7832 }
7833 reranked.push((row_id, fused_score, components, Some(score), final_score));
7834 }
7835 ranked = reranked;
7836 ranked.sort_by(order);
7837 crate::trace::QueryTrace::record(|trace| {
7838 trace.exact_vector_gather_nanos =
7839 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
7840 trace.exact_vector_score_nanos = trace
7841 .exact_vector_score_nanos
7842 .saturating_add(score_started.elapsed().as_nanos() as u64);
7843 });
7844 }
7845 if let Some(after) = after {
7846 ranked.retain(|(row_id, _, _, _, final_score)| {
7847 final_score.total_cmp(&after.final_score).is_lt()
7848 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
7849 });
7850 }
7851 let projection_started = std::time::Instant::now();
7852 let sentinel = projection
7853 .first()
7854 .copied()
7855 .or_else(|| self.schema.columns.first().map(|column| column.id));
7856 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7857 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
7858 let mut projection_rows = 0usize;
7859 let mut projection_cells = 0usize;
7860 while out.len() < request.limit && !ranked.is_empty() {
7861 if let Some(context) = context {
7862 context.checkpoint()?;
7863 context.consume(ranked.len())?;
7864 }
7865 let needed = request.limit - out.len();
7866 let window_size = ranked
7867 .len()
7868 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
7869 let selection_started = std::time::Instant::now();
7870 let mut remainder = if ranked.len() > window_size {
7871 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
7872 ranked.split_off(window_size)
7873 } else {
7874 Vec::new()
7875 };
7876 ranked.sort_by(order);
7877 fusion_nanos =
7878 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
7879 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
7880 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
7881 if let Some(context) = context {
7882 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
7883 }
7884 projection_rows = projection_rows.saturating_add(row_ids.len());
7885 projection_cells =
7886 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
7887 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
7888 std::collections::HashMap::new();
7889 if let Some(column_id) = sentinel {
7890 for (row_id, value) in self.values_for_rids_batch_at_with_context(
7891 &row_ids, column_id, snapshot, query_now, context,
7892 )? {
7893 cells.entry(row_id).or_default().insert(column_id, value);
7894 }
7895 }
7896 for &column_id in &projection {
7897 if Some(column_id) == sentinel {
7898 continue;
7899 }
7900 for (row_id, value) in self.values_for_rids_batch_at_with_context(
7901 &row_ids, column_id, snapshot, query_now, context,
7902 )? {
7903 cells.entry(row_id).or_default().insert(column_id, value);
7904 }
7905 }
7906 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
7907 ranked.drain(..)
7908 {
7909 let Some(row_cells) = cells.remove(&row_id) else {
7910 continue;
7911 };
7912 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
7913 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
7914 out.push(SearchHit {
7915 row_id,
7916 cells: projection
7917 .iter()
7918 .filter_map(|column_id| {
7919 row_cells
7920 .get(column_id)
7921 .cloned()
7922 .map(|value| (*column_id, value))
7923 })
7924 .collect(),
7925 components,
7926 fused_score,
7927 exact_rerank_score,
7928 final_score,
7929 final_rank,
7930 });
7931 if out.len() == request.limit {
7932 break;
7933 }
7934 }
7935 ranked.append(&mut remainder);
7936 }
7937 crate::trace::QueryTrace::record(|trace| {
7938 trace.union_size = union_size;
7939 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
7940 trace.projection_nanos = trace
7941 .projection_nanos
7942 .saturating_add(projection_started.elapsed().as_nanos() as u64);
7943 trace.total_nanos = trace
7944 .total_nanos
7945 .saturating_add(total_started.elapsed().as_nanos() as u64);
7946 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
7947 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
7948 if let Some(context) = context {
7949 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
7950 }
7951 });
7952 Ok(out)
7953 }
7954
7955 pub fn set_similarity(
7958 &mut self,
7959 request: &crate::query::SetSimilarityRequest,
7960 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7961 self.set_similarity_with_allowed(request, None)
7962 }
7963
7964 pub fn set_similarity_at(
7965 &mut self,
7966 request: &crate::query::SetSimilarityRequest,
7967 snapshot: Snapshot,
7968 allowed: Option<&std::collections::HashSet<RowId>>,
7969 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
7970 self.set_similarity_explained_at(request, snapshot, allowed)
7971 .map(|(hits, _)| hits)
7972 }
7973
7974 pub fn ann_rerank(
7976 &mut self,
7977 request: &crate::query::AnnRerankRequest,
7978 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7979 self.ann_rerank_with_allowed(request, None)
7980 }
7981
7982 pub fn ann_rerank_with_allowed(
7983 &mut self,
7984 request: &crate::query::AnnRerankRequest,
7985 allowed: Option<&std::collections::HashSet<RowId>>,
7986 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7987 self.ann_rerank_at(request, self.snapshot(), allowed)
7988 }
7989
7990 pub fn ann_rerank_at(
7991 &mut self,
7992 request: &crate::query::AnnRerankRequest,
7993 snapshot: Snapshot,
7994 allowed: Option<&std::collections::HashSet<RowId>>,
7995 ) -> Result<Vec<crate::query::AnnRerankHit>> {
7996 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
7997 }
7998
7999 pub fn ann_rerank_at_with_context(
8000 &mut self,
8001 request: &crate::query::AnnRerankRequest,
8002 snapshot: Snapshot,
8003 allowed: Option<&std::collections::HashSet<RowId>>,
8004 context: Option<&crate::query::AiExecutionContext>,
8005 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8006 self.ensure_indexes_complete()?;
8007 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
8008 }
8009
8010 pub fn ann_rerank_at_with_candidate_authorization_and_context(
8011 &mut self,
8012 request: &crate::query::AnnRerankRequest,
8013 snapshot: Snapshot,
8014 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8015 context: Option<&crate::query::AiExecutionContext>,
8016 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8017 self.ensure_indexes_complete()?;
8018 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
8019 }
8020
8021 #[doc(hidden)]
8022 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
8023 &self,
8024 request: &crate::query::AnnRerankRequest,
8025 snapshot: Snapshot,
8026 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8027 context: Option<&crate::query::AiExecutionContext>,
8028 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8029 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
8030 }
8031
8032 fn ann_rerank_at_with_filters_and_context(
8033 &self,
8034 request: &crate::query::AnnRerankRequest,
8035 snapshot: Snapshot,
8036 allowed: Option<&std::collections::HashSet<RowId>>,
8037 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8038 context: Option<&crate::query::AiExecutionContext>,
8039 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8040 use crate::query::{
8041 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
8042 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
8043 };
8044 if request.candidate_k == 0 || request.limit == 0 {
8045 return Err(MongrelError::InvalidArgument(
8046 "candidate_k and limit must be > 0".into(),
8047 ));
8048 }
8049 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
8050 return Err(MongrelError::InvalidArgument(format!(
8051 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
8052 )));
8053 }
8054 let retriever = Retriever::Ann {
8055 column_id: request.column_id,
8056 query: request.query.clone(),
8057 k: request.candidate_k,
8058 };
8059 self.require_select()?;
8060 self.validate_retriever(&retriever)?;
8061 let hits = self.retrieve_filtered(
8062 &retriever,
8063 snapshot,
8064 None,
8065 allowed,
8066 candidate_authorization,
8067 context,
8068 )?;
8069 let distances: std::collections::HashMap<_, _> = hits
8070 .iter()
8071 .filter_map(|hit| match hit.score {
8072 RetrieverScore::AnnHammingDistance(distance) => {
8073 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
8074 }
8075 RetrieverScore::AnnCosineDistance(distance) => {
8076 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
8077 }
8078 _ => None,
8079 })
8080 .collect();
8081 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
8082 if let Some(context) = context {
8083 context.consume(row_ids.len())?;
8084 }
8085 let gather_started = std::time::Instant::now();
8086 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8087 let values = self.values_for_rids_batch_at_with_context(
8088 &row_ids,
8089 request.column_id,
8090 snapshot,
8091 query_now,
8092 context,
8093 )?;
8094 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
8095 let score_started = std::time::Instant::now();
8096 let vector_work =
8097 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
8098 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
8099 if let Some(context) = context {
8100 context.consume(vector_work)?;
8101 }
8102 request
8103 .query
8104 .iter()
8105 .map(|value| f64::from(*value).powi(2))
8106 .sum::<f64>()
8107 .sqrt()
8108 } else {
8109 0.0
8110 };
8111 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
8112 for (row_id, value) in values {
8113 if let Some(meta) = value.generated_embedding_metadata() {
8114 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
8115 continue;
8116 }
8117 }
8118 let Some(vector) = value.as_embedding() else {
8119 continue;
8120 };
8121 let exact_score = match request.metric {
8122 VectorMetric::DotProduct => {
8123 if let Some(context) = context {
8124 context.consume(vector_work)?;
8125 }
8126 request
8127 .query
8128 .iter()
8129 .zip(vector)
8130 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8131 .sum::<f64>()
8132 }
8133 VectorMetric::Cosine => {
8134 if let Some(context) = context {
8135 context.consume(vector_work.saturating_mul(2))?;
8136 }
8137 let dot = request
8138 .query
8139 .iter()
8140 .zip(vector)
8141 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8142 .sum::<f64>();
8143 let norm = vector
8144 .iter()
8145 .map(|value| f64::from(*value).powi(2))
8146 .sum::<f64>()
8147 .sqrt();
8148 if query_norm == 0.0 || norm == 0.0 {
8149 0.0
8150 } else {
8151 dot / (query_norm * norm)
8152 }
8153 }
8154 VectorMetric::Euclidean => {
8155 if let Some(context) = context {
8156 context.consume(vector_work)?;
8157 }
8158 request
8159 .query
8160 .iter()
8161 .zip(vector)
8162 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
8163 .sum::<f64>()
8164 .sqrt()
8165 }
8166 };
8167 let exact_score = exact_score as f32;
8168 if !exact_score.is_finite() {
8169 return Err(MongrelError::InvalidArgument(
8170 "exact ANN score must be finite".into(),
8171 ));
8172 }
8173 let Some(candidate_distance) = distances.get(&row_id).copied() else {
8174 continue;
8175 };
8176 reranked.push(AnnRerankHit {
8177 row_id,
8178 candidate_distance,
8179 exact_score,
8180 });
8181 }
8182 reranked.sort_by(|left, right| {
8183 let score = match request.metric {
8184 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
8185 VectorMetric::Cosine | VectorMetric::DotProduct => {
8186 right.exact_score.total_cmp(&left.exact_score)
8187 }
8188 };
8189 score.then_with(|| left.row_id.cmp(&right.row_id))
8190 });
8191 reranked.truncate(request.limit);
8192 crate::trace::QueryTrace::record(|trace| {
8193 trace.exact_vector_gather_nanos =
8194 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
8195 trace.exact_vector_score_nanos = trace
8196 .exact_vector_score_nanos
8197 .saturating_add(score_started.elapsed().as_nanos() as u64);
8198 });
8199 Ok(reranked)
8200 }
8201
8202 pub fn set_similarity_with_allowed(
8203 &mut self,
8204 request: &crate::query::SetSimilarityRequest,
8205 allowed: Option<&std::collections::HashSet<RowId>>,
8206 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8207 self.set_similarity_explained_at(request, self.snapshot(), allowed)
8208 .map(|(hits, _)| hits)
8209 }
8210
8211 pub fn set_similarity_explained(
8212 &mut self,
8213 request: &crate::query::SetSimilarityRequest,
8214 ) -> Result<(
8215 Vec<crate::query::SetSimilarityHit>,
8216 crate::query::SetSimilarityTrace,
8217 )> {
8218 self.set_similarity_explained_at(request, self.snapshot(), None)
8219 }
8220
8221 fn set_similarity_explained_at(
8222 &mut self,
8223 request: &crate::query::SetSimilarityRequest,
8224 snapshot: Snapshot,
8225 allowed: Option<&std::collections::HashSet<RowId>>,
8226 ) -> Result<(
8227 Vec<crate::query::SetSimilarityHit>,
8228 crate::query::SetSimilarityTrace,
8229 )> {
8230 self.ensure_indexes_complete()?;
8231 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
8232 }
8233
8234 pub fn set_similarity_at_with_context(
8235 &mut self,
8236 request: &crate::query::SetSimilarityRequest,
8237 snapshot: Snapshot,
8238 allowed: Option<&std::collections::HashSet<RowId>>,
8239 context: Option<&crate::query::AiExecutionContext>,
8240 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8241 self.ensure_indexes_complete()?;
8242 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
8243 .map(|(hits, _)| hits)
8244 }
8245
8246 pub fn set_similarity_at_with_candidate_authorization_and_context(
8247 &mut self,
8248 request: &crate::query::SetSimilarityRequest,
8249 snapshot: Snapshot,
8250 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8251 context: Option<&crate::query::AiExecutionContext>,
8252 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8253 self.ensure_indexes_complete()?;
8254 self.set_similarity_explained_at_with_context(
8255 request,
8256 snapshot,
8257 None,
8258 authorization,
8259 context,
8260 )
8261 .map(|(hits, _)| hits)
8262 }
8263
8264 #[doc(hidden)]
8265 pub fn set_similarity_at_with_candidate_authorization_on_generation(
8266 &self,
8267 request: &crate::query::SetSimilarityRequest,
8268 snapshot: Snapshot,
8269 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8270 context: Option<&crate::query::AiExecutionContext>,
8271 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8272 self.set_similarity_explained_at_with_context(
8273 request,
8274 snapshot,
8275 None,
8276 authorization,
8277 context,
8278 )
8279 .map(|(hits, _)| hits)
8280 }
8281
8282 fn set_similarity_explained_at_with_context(
8283 &self,
8284 request: &crate::query::SetSimilarityRequest,
8285 snapshot: Snapshot,
8286 allowed: Option<&std::collections::HashSet<RowId>>,
8287 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8288 context: Option<&crate::query::AiExecutionContext>,
8289 ) -> Result<(
8290 Vec<crate::query::SetSimilarityHit>,
8291 crate::query::SetSimilarityTrace,
8292 )> {
8293 use crate::query::{
8294 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
8295 MAX_SET_MEMBERS,
8296 };
8297 let mut trace = crate::query::SetSimilarityTrace::default();
8298 if request.members.is_empty() {
8299 return Ok((Vec::new(), trace));
8300 }
8301 if request.candidate_k == 0 || request.limit == 0 {
8302 return Err(MongrelError::InvalidArgument(
8303 "candidate_k and limit must be > 0".into(),
8304 ));
8305 }
8306 if request.candidate_k > MAX_RETRIEVER_K
8307 || request.limit > MAX_FINAL_LIMIT
8308 || request.members.len() > MAX_SET_MEMBERS
8309 {
8310 return Err(MongrelError::InvalidArgument(format!(
8311 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
8312 )));
8313 }
8314 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
8315 return Err(MongrelError::InvalidArgument(
8316 "min_jaccard must be finite and between 0 and 1".into(),
8317 ));
8318 }
8319 let started = std::time::Instant::now();
8320 let retriever = Retriever::MinHash {
8321 column_id: request.column_id,
8322 members: request.members.clone(),
8323 k: request.candidate_k,
8324 };
8325 self.require_select()?;
8326 self.validate_retriever(&retriever)?;
8327 let hits = self.retrieve_filtered(
8328 &retriever,
8329 snapshot,
8330 None,
8331 allowed,
8332 candidate_authorization,
8333 context,
8334 )?;
8335 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
8336 trace.candidate_count = hits.len();
8337 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
8338 if let Some(context) = context {
8339 context.consume(row_ids.len())?;
8340 }
8341 let started = std::time::Instant::now();
8342 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8343 let values = self.values_for_rids_batch_at_with_context(
8344 &row_ids,
8345 request.column_id,
8346 snapshot,
8347 query_now,
8348 context,
8349 )?;
8350 trace.gather_us = started.elapsed().as_micros() as u64;
8351 if let Some(context) = context {
8352 context.consume(request.members.len())?;
8353 }
8354 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
8355 let estimates: std::collections::HashMap<_, _> = hits
8356 .into_iter()
8357 .filter_map(|hit| match hit.score {
8358 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
8359 _ => None,
8360 })
8361 .collect();
8362 let started = std::time::Instant::now();
8363 let mut parsed = Vec::with_capacity(values.len());
8364 for (row_id, value) in values {
8365 let Value::Bytes(bytes) = value else {
8366 continue;
8367 };
8368 if let Some(context) = context {
8369 context.consume(crate::query::work_units(
8370 bytes.len(),
8371 crate::query::PARSE_WORK_QUANTUM,
8372 ))?;
8373 }
8374 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
8375 continue;
8376 };
8377 if let Some(context) = context {
8378 context.consume(members.len())?;
8379 }
8380 let stored = members
8381 .into_iter()
8382 .filter_map(|member| match member {
8383 serde_json::Value::String(value) => {
8384 Some(crate::query::SetMember::String(value))
8385 }
8386 serde_json::Value::Number(value) => {
8387 Some(crate::query::SetMember::Number(value))
8388 }
8389 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
8390 _ => None,
8391 })
8392 .collect::<std::collections::HashSet<_>>();
8393 parsed.push((row_id, stored));
8394 }
8395 trace.parse_us = started.elapsed().as_micros() as u64;
8396 trace.verified_count = parsed.len();
8397 let started = std::time::Instant::now();
8398 let mut exact = Vec::new();
8399 for (row_id, stored) in parsed {
8400 if let Some(context) = context {
8401 context.consume(query.len().saturating_add(stored.len()))?;
8402 }
8403 let union = query.union(&stored).count();
8404 let score = if union == 0 {
8405 1.0
8406 } else {
8407 query.intersection(&stored).count() as f32 / union as f32
8408 };
8409 if score >= request.min_jaccard {
8410 exact.push(SetSimilarityHit {
8411 row_id,
8412 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
8413 exact_jaccard: score,
8414 });
8415 }
8416 }
8417 exact.sort_by(|a, b| {
8418 b.exact_jaccard
8419 .total_cmp(&a.exact_jaccard)
8420 .then_with(|| a.row_id.cmp(&b.row_id))
8421 });
8422 exact.truncate(request.limit);
8423 trace.score_us = started.elapsed().as_micros() as u64;
8424 crate::trace::QueryTrace::record(|query_trace| {
8425 query_trace.exact_set_gather_nanos = query_trace
8426 .exact_set_gather_nanos
8427 .saturating_add(trace.gather_us.saturating_mul(1_000));
8428 query_trace.exact_set_parse_nanos = query_trace
8429 .exact_set_parse_nanos
8430 .saturating_add(trace.parse_us.saturating_mul(1_000));
8431 query_trace.exact_set_score_nanos = query_trace
8432 .exact_set_score_nanos
8433 .saturating_add(trace.score_us.saturating_mul(1_000));
8434 });
8435 Ok((exact, trace))
8436 }
8437
8438 fn values_for_rids_batch_at(
8440 &self,
8441 row_ids: &[u64],
8442 column_id: u16,
8443 snapshot: Snapshot,
8444 now: i64,
8445 ) -> Result<Vec<(RowId, Value)>> {
8446 if self.ttl.is_none()
8447 && self.memtable.is_empty()
8448 && self.mutable_run.is_empty()
8449 && self.run_refs.len() == 1
8450 {
8451 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8452 if row_ids.len().saturating_mul(24) < reader.row_count() {
8457 let mut values = Vec::with_capacity(row_ids.len());
8458 for &raw_row_id in row_ids {
8459 let row_id = RowId(raw_row_id);
8460 if let Some((_, false, Some(value))) =
8461 reader.get_version_column(row_id, snapshot.epoch, column_id)?
8462 {
8463 values.push((row_id, value));
8464 }
8465 }
8466 return Ok(values);
8467 }
8468 let (positions, visible_row_ids) =
8469 reader.visible_positions_with_rids(snapshot.epoch)?;
8470 let requested: Vec<(RowId, usize)> = row_ids
8471 .iter()
8472 .filter_map(|raw| {
8473 visible_row_ids
8474 .binary_search(&(*raw as i64))
8475 .ok()
8476 .map(|index| (RowId(*raw), positions[index]))
8477 })
8478 .collect();
8479 let values = reader.gather_column(
8480 column_id,
8481 &requested
8482 .iter()
8483 .map(|(_, position)| *position)
8484 .collect::<Vec<_>>(),
8485 )?;
8486 return Ok(requested
8487 .into_iter()
8488 .zip(values)
8489 .map(|((row_id, _), value)| (row_id, value))
8490 .collect());
8491 }
8492 self.values_for_rids_at(row_ids, column_id, snapshot, now)
8493 }
8494
8495 fn values_for_rids_batch_at_with_context(
8496 &self,
8497 row_ids: &[u64],
8498 column_id: u16,
8499 snapshot: Snapshot,
8500 now: i64,
8501 context: Option<&crate::query::AiExecutionContext>,
8502 ) -> Result<Vec<(RowId, Value)>> {
8503 let Some(context) = context else {
8504 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
8505 };
8506 let mut values = Vec::with_capacity(row_ids.len());
8507 for chunk in row_ids.chunks(256) {
8508 context.checkpoint()?;
8509 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
8510 }
8511 Ok(values)
8512 }
8513
8514 fn values_for_rids_at(
8516 &self,
8517 row_ids: &[u64],
8518 column_id: u16,
8519 snapshot: Snapshot,
8520 now: i64,
8521 ) -> Result<Vec<(RowId, Value)>> {
8522 let mut readers: Vec<_> = self
8523 .run_refs
8524 .iter()
8525 .map(|run| self.open_reader(run.run_id))
8526 .collect::<Result<_>>()?;
8527 let mut out = Vec::with_capacity(row_ids.len());
8528 for &raw_row_id in row_ids {
8529 let row_id = RowId(raw_row_id);
8530 let mem = self.memtable.get_version_at(row_id, snapshot);
8531 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
8532 let overlay = match (mem, mutable) {
8533 (Some((_, a)), Some((_, b))) => Some(
8534 if Snapshot::version_is_newer(
8535 a.committed_epoch,
8536 a.commit_ts,
8537 b.committed_epoch,
8538 b.commit_ts,
8539 ) {
8540 a
8541 } else {
8542 b
8543 },
8544 ),
8545 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
8546 (None, None) => None,
8547 };
8548 if let Some(row) = overlay {
8549 if !row.deleted && !self.row_expired_at(&row, now) {
8550 if let Some(value) = row.columns.get(&column_id) {
8551 out.push((row_id, value.clone()));
8552 }
8553 }
8554 continue;
8555 }
8556
8557 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
8558 for (index, reader) in readers.iter_mut().enumerate() {
8559 if let Some((epoch, deleted, value)) =
8560 reader.get_version_column(row_id, snapshot.epoch, column_id)?
8561 {
8562 if best
8563 .as_ref()
8564 .map(|(best_epoch, ..)| epoch > *best_epoch)
8565 .unwrap_or(true)
8566 {
8567 best = Some((epoch, deleted, value, index));
8568 }
8569 }
8570 }
8571 let Some((_, false, Some(value), reader_index)) = best else {
8572 continue;
8573 };
8574 if let Some(ttl) = self.ttl {
8575 if ttl.column_id != column_id {
8576 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
8577 .get_version_column(row_id, snapshot.epoch, ttl.column_id)?
8578 {
8579 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8580 continue;
8581 }
8582 }
8583 } else if let Value::Int64(timestamp) = value {
8584 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8585 continue;
8586 }
8587 }
8588 }
8589 out.push((row_id, value));
8590 }
8591 Ok(out)
8592 }
8593
8594 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
8599 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
8600 }
8601
8602 pub fn rows_for_rids_with_context(
8603 &self,
8604 rids: &[u64],
8605 snapshot: Snapshot,
8606 context: &crate::query::AiExecutionContext,
8607 ) -> Result<Vec<Row>> {
8608 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
8609 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
8610 }
8611
8612 fn rows_for_rids_at_time(
8613 &self,
8614 rids: &[u64],
8615 snapshot: Snapshot,
8616 ttl_now: i64,
8617 control: Option<&crate::ExecutionControl>,
8618 ) -> Result<Vec<Row>> {
8619 use std::collections::HashMap;
8620 let mut rows = Vec::with_capacity(rids.len());
8621 let tier_size = self.memtable.len() + self.mutable_run.len();
8636 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
8637 if rids.len().saturating_mul(24) < tier_size {
8638 for &rid in rids {
8639 if overlay.len() & 255 == 0 {
8640 control
8641 .map(crate::ExecutionControl::checkpoint)
8642 .transpose()?;
8643 }
8644 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
8645 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
8646 let newest = match (mem, mrun) {
8647 (Some((_, mr)), Some((_, rr))) => Some(
8648 if Snapshot::version_is_newer(
8649 mr.committed_epoch,
8650 mr.commit_ts,
8651 rr.committed_epoch,
8652 rr.commit_ts,
8653 ) {
8654 mr
8655 } else {
8656 rr
8657 },
8658 ),
8659 (Some((_, mr)), None) => Some(mr),
8660 (None, Some((_, rr))) => Some(rr),
8661 (None, None) => None,
8662 };
8663 if let Some(row) = newest {
8664 overlay.insert(rid, row);
8665 }
8666 }
8667 } else {
8668 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
8669 overlay
8670 .entry(row.row_id.0)
8671 .and_modify(|e| {
8672 if Snapshot::version_is_newer(
8673 row.committed_epoch,
8674 row.commit_ts,
8675 e.committed_epoch,
8676 e.commit_ts,
8677 ) {
8678 *e = row.clone();
8679 }
8680 })
8681 .or_insert(row);
8682 };
8683 for (index, row) in self
8684 .memtable
8685 .visible_versions_at(snapshot)
8686 .into_iter()
8687 .enumerate()
8688 {
8689 if index & 255 == 0 {
8690 control
8691 .map(crate::ExecutionControl::checkpoint)
8692 .transpose()?;
8693 }
8694 fold_newest(row, &mut overlay);
8695 }
8696 for (index, row) in self
8697 .mutable_run
8698 .visible_versions_at(snapshot)
8699 .into_iter()
8700 .enumerate()
8701 {
8702 if index & 255 == 0 {
8703 control
8704 .map(crate::ExecutionControl::checkpoint)
8705 .transpose()?;
8706 }
8707 fold_newest(row, &mut overlay);
8708 }
8709 }
8710 if self.run_refs.len() == 1 {
8711 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
8712 if rids.len().saturating_mul(24) < reader.row_count() {
8720 for (index, &rid) in rids.iter().enumerate() {
8721 if index & 255 == 0 {
8722 control
8723 .map(crate::ExecutionControl::checkpoint)
8724 .transpose()?;
8725 }
8726 if let Some(r) = overlay.get(&rid) {
8727 if !r.deleted {
8728 rows.push(r.clone());
8729 }
8730 continue;
8731 }
8732 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
8733 if !row.deleted {
8734 rows.push(row);
8735 }
8736 }
8737 }
8738 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8739 return Ok(rows);
8740 }
8741 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
8750 enum Src {
8753 Overlay,
8754 Run,
8755 }
8756 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
8757 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
8758 for (index, rid) in rids.iter().enumerate() {
8759 if index & 255 == 0 {
8760 control
8761 .map(crate::ExecutionControl::checkpoint)
8762 .transpose()?;
8763 }
8764 if overlay.contains_key(rid) {
8765 plan.push(Src::Overlay);
8766 continue;
8767 }
8768 match vis_rids.binary_search(&(*rid as i64)) {
8769 Ok(i) => {
8770 plan.push(Src::Run);
8771 fetch.push(positions[i]);
8772 }
8773 Err(_) => { }
8774 }
8775 }
8776 let fetched = reader.materialize_batch(&fetch)?;
8777 let mut fetched_iter = fetched.into_iter();
8778 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
8779 if index & 255 == 0 {
8780 control
8781 .map(crate::ExecutionControl::checkpoint)
8782 .transpose()?;
8783 }
8784 match src {
8785 Src::Overlay => {
8786 if let Some(r) = overlay.get(rid) {
8787 if !r.deleted {
8788 rows.push(r.clone());
8789 }
8790 }
8791 }
8792 Src::Run => {
8793 if let Some(row) = fetched_iter.next() {
8794 if !row.deleted {
8795 rows.push(row);
8796 }
8797 }
8798 }
8799 }
8800 }
8801 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8802 return Ok(rows);
8803 }
8804 let mut readers: Vec<_> = self
8808 .run_refs
8809 .iter()
8810 .map(|rr| self.open_reader(rr.run_id))
8811 .collect::<Result<Vec<_>>>()?;
8812 for (index, rid) in rids.iter().enumerate() {
8813 if index & 255 == 0 {
8814 control
8815 .map(crate::ExecutionControl::checkpoint)
8816 .transpose()?;
8817 }
8818 if let Some(r) = overlay.get(rid) {
8819 if !r.deleted {
8820 rows.push(r.clone());
8821 }
8822 continue;
8823 }
8824 let mut best: Option<(Epoch, Row)> = None;
8825 for reader in readers.iter_mut() {
8826 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
8827 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
8828 best = Some((epoch, row));
8829 }
8830 }
8831 }
8832 if let Some((_, r)) = best {
8833 if !r.deleted {
8834 rows.push(r);
8835 }
8836 }
8837 }
8838 rows.retain(|row| !self.row_expired_at(row, ttl_now));
8839 Ok(rows)
8840 }
8841
8842 pub fn indexes_complete(&self) -> bool {
8852 self.indexes_complete
8853 }
8854
8855 pub fn index_build_policy(&self) -> IndexBuildPolicy {
8857 self.index_build_policy
8858 }
8859
8860 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
8864 self.index_build_policy = policy;
8865 }
8866
8867 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
8872 if !self.indexes_complete {
8876 return None;
8877 }
8878 let b = self.bitmap.get(&column_id)?;
8879 let result: Vec<Vec<u8>> = b
8880 .keys()
8881 .into_iter()
8882 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
8883 .collect();
8884 Some(result)
8885 }
8886
8887 pub fn fk_join_row_ids(
8888 &self,
8889 fk_column_id: u16,
8890 pk_values: &[Vec<u8>],
8891 fk_conditions: &[crate::query::Condition],
8892 snapshot: Snapshot,
8893 ) -> Result<Vec<u64>> {
8894 let Some(b) = self.bitmap.get(&fk_column_id) else {
8895 return Ok(Vec::new());
8896 };
8897 let mut join_set = {
8898 let mut acc = roaring::RoaringBitmap::new();
8899 for v in pk_values {
8900 acc |= b.get(v);
8901 }
8902 RowIdSet::from_roaring(acc)
8903 };
8904 if !fk_conditions.is_empty() {
8905 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
8906 sets.push(join_set);
8907 for c in fk_conditions {
8908 sets.push(self.resolve_condition(c, snapshot)?);
8909 }
8910 join_set = RowIdSet::intersect_many(sets);
8911 }
8912 Ok(join_set.into_sorted_vec())
8913 }
8914
8915 pub fn fk_join_count(
8921 &self,
8922 fk_column_id: u16,
8923 pk_values: &[Vec<u8>],
8924 fk_conditions: &[crate::query::Condition],
8925 snapshot: Snapshot,
8926 ) -> Result<u64> {
8927 let Some(b) = self.bitmap.get(&fk_column_id) else {
8928 return Ok(0);
8929 };
8930 let mut acc = roaring::RoaringBitmap::new();
8931 for v in pk_values {
8932 acc |= b.get(v);
8933 }
8934 if fk_conditions.is_empty() {
8935 return Ok(acc.len());
8936 }
8937 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
8938 sets.push(RowIdSet::from_roaring(acc));
8939 for c in fk_conditions {
8940 sets.push(self.resolve_condition(c, snapshot)?);
8941 }
8942 Ok(RowIdSet::intersect_many(sets).len() as u64)
8943 }
8944
8945 fn resolve_condition(
8950 &self,
8951 c: &crate::query::Condition,
8952 snapshot: Snapshot,
8953 ) -> Result<RowIdSet> {
8954 self.resolve_condition_with_allowed(c, snapshot, None)
8955 }
8956
8957 fn resolve_condition_with_allowed(
8958 &self,
8959 c: &crate::query::Condition,
8960 snapshot: Snapshot,
8961 allowed: Option<&std::collections::HashSet<RowId>>,
8962 ) -> Result<RowIdSet> {
8963 use crate::query::Condition;
8964 self.validate_condition(c)?;
8965 Ok(match c {
8966 Condition::Pk(key) => {
8967 let lookup = self
8968 .schema
8969 .primary_key()
8970 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
8971 .unwrap_or_else(|| key.clone());
8972 if let Some(r) = self.hot.get(&lookup) {
8973 if snapshot.epoch < self.current_epoch() {
8979 self.record_hot_fallback_reason(
8987 crate::trace::HotFallbackReason::HistoricalSnapshot,
8988 );
8989 if let Some(pk_col) = self.schema.primary_key() {
8990 let (result, _inner) =
8991 self.pk_equality_fallback(pk_col.id, &lookup, snapshot)?;
8992 return Ok(result);
8993 }
8994 return Ok(RowIdSet::empty());
8995 }
8996 if let Some(row) = self.get(crate::rowid::RowId(r.0), snapshot) {
8997 let pk_ok = self
8998 .schema
8999 .primary_key()
9000 .map(|pk| {
9001 row.columns
9002 .get(&pk.id)
9003 .map(|v| self.index_lookup_key(pk.id, v) == lookup)
9004 .unwrap_or(false)
9005 })
9006 .unwrap_or(true);
9007 if pk_ok {
9008 self.lookup_metrics
9009 .hot_lookup_hit
9010 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9011 crate::trace::QueryTrace::record(|t| {
9012 t.hot_lookup_attempted = true;
9013 t.hot_lookup_hit = true;
9014 });
9015 RowIdSet::one(r.0)
9016 } else {
9017 self.record_hot_fallback_reason(
9018 crate::trace::HotFallbackReason::PrimaryKeyMismatch,
9019 );
9020 RowIdSet::one(r.0)
9021 }
9022 } else {
9023 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::Tombstone);
9029 self.lookup_metrics
9033 .hot_fallback_runs_considered_total
9034 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9035 RowIdSet::empty()
9036 }
9037 } else if let Some(pk_col) = self.schema.primary_key() {
9038 let (result, reason) =
9046 self.pk_equality_fallback(pk_col.id, &lookup, snapshot)?;
9047 self.record_hot_fallback_reason(reason);
9048 return Ok(result);
9049 } else {
9050 RowIdSet::empty()
9051 }
9052 }
9053 Condition::BitmapEq { column_id, value } => {
9054 let lookup = self.index_lookup_key_bytes(*column_id, value);
9055 self.bitmap
9056 .get(column_id)
9057 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9058 .unwrap_or_else(RowIdSet::empty)
9059 }
9060 Condition::BitmapIn { column_id, values } => {
9061 let bm = self.bitmap.get(column_id);
9062 let mut acc = roaring::RoaringBitmap::new();
9063 if let Some(b) = bm {
9064 for v in values {
9065 let lookup = self.index_lookup_key_bytes(*column_id, v);
9066 acc |= b.get(&lookup);
9067 }
9068 }
9069 RowIdSet::from_roaring(acc)
9070 }
9071 Condition::BytesPrefix { column_id, prefix } => {
9072 if let Some(b) = self.bitmap.get(column_id) {
9077 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9078 let mut acc = roaring::RoaringBitmap::new();
9079 for key in b.keys() {
9080 if key.starts_with(&lookup_prefix) {
9081 acc |= b.get(&key);
9082 }
9083 }
9084 RowIdSet::from_roaring(acc)
9085 } else {
9086 RowIdSet::empty()
9087 }
9088 }
9089 Condition::FmContains { column_id, pattern } => self
9090 .fm
9091 .get(column_id)
9092 .map(|f| {
9093 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
9094 })
9095 .unwrap_or_else(RowIdSet::empty),
9096 Condition::FmContainsAll {
9097 column_id,
9098 patterns,
9099 } => {
9100 if let Some(f) = self.fm.get(column_id) {
9103 let sets: Vec<RowIdSet> = patterns
9104 .iter()
9105 .map(|pat| {
9106 RowIdSet::from_unsorted(
9107 f.locate(pat).into_iter().map(|r| r.0).collect(),
9108 )
9109 })
9110 .collect();
9111 RowIdSet::intersect_many(sets)
9112 } else {
9113 RowIdSet::empty()
9114 }
9115 }
9116 Condition::Ann {
9117 column_id,
9118 query,
9119 k,
9120 } => RowIdSet::from_unsorted(
9121 self.retrieve_filtered(
9122 &crate::query::Retriever::Ann {
9123 column_id: *column_id,
9124 query: query.clone(),
9125 k: *k,
9126 },
9127 snapshot,
9128 None,
9129 allowed,
9130 None,
9131 None,
9132 )?
9133 .into_iter()
9134 .map(|hit| hit.row_id.0)
9135 .collect(),
9136 ),
9137 Condition::SparseMatch {
9138 column_id,
9139 query,
9140 k,
9141 } => RowIdSet::from_unsorted(
9142 self.retrieve_filtered(
9143 &crate::query::Retriever::Sparse {
9144 column_id: *column_id,
9145 query: query.clone(),
9146 k: *k,
9147 },
9148 snapshot,
9149 None,
9150 allowed,
9151 None,
9152 None,
9153 )?
9154 .into_iter()
9155 .map(|hit| hit.row_id.0)
9156 .collect(),
9157 ),
9158 Condition::MinHashSimilar {
9159 column_id,
9160 query,
9161 k,
9162 } => match self.minhash.get(column_id) {
9163 Some(index) => {
9164 let candidates = index.candidate_row_ids(query);
9165 let eligible =
9166 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9167 RowIdSet::from_unsorted(
9168 index
9169 .search_filtered(query, *k, |row_id| {
9170 eligible.contains(&row_id)
9171 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
9172 })
9173 .into_iter()
9174 .map(|(row_id, _)| row_id.0)
9175 .collect(),
9176 )
9177 }
9178 None => RowIdSet::empty(),
9179 },
9180 Condition::Range { column_id, lo, hi } => {
9181 let mut set = if let Some(li) = self.learned_range.get(column_id) {
9200 if self.run_refs.len() == 1 {
9201 RowIdSet::from_unsorted(
9204 li.range(*lo, *hi).into_iter().collect(),
9205 )
9206 } else {
9207 let mut multi =
9215 self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
9216 if lo == hi {
9217 self.union_bitmap_point_i64(
9218 &mut multi,
9219 *column_id,
9220 *lo,
9221 snapshot,
9222 );
9223 }
9224 return Ok(multi);
9225 }
9226 } else if self.run_refs.len() == 1 {
9227 let mut r = self.open_reader(self.run_refs[0].run_id)?;
9228 r.range_row_id_set_i64(*column_id, *lo, *hi)?
9229 } else {
9230 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
9233 if lo == hi {
9234 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
9235 }
9236 return Ok(multi);
9237 };
9238 set.remove_many(self.overlay_rid_set(snapshot));
9239 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
9240 if lo == hi {
9241 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
9242 }
9243 set
9244 }
9245 Condition::RangeF64 {
9246 column_id,
9247 lo,
9248 lo_inclusive,
9249 hi,
9250 hi_inclusive,
9251 } => {
9252 let mut set = if let Some(li) = self.learned_range.get(column_id) {
9255 RowIdSet::from_unsorted(
9256 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
9257 .into_iter()
9258 .collect(),
9259 )
9260 } else if self.run_refs.len() == 1 {
9261 let mut r = self.open_reader(self.run_refs[0].run_id)?;
9262 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
9263 } else {
9264 return self.range_scan_f64(
9265 *column_id,
9266 *lo,
9267 *lo_inclusive,
9268 *hi,
9269 *hi_inclusive,
9270 snapshot,
9271 );
9272 };
9273 set.remove_many(self.overlay_rid_set(snapshot));
9274 self.range_scan_overlay_f64(
9275 &mut set,
9276 *column_id,
9277 *lo,
9278 *lo_inclusive,
9279 *hi,
9280 *hi_inclusive,
9281 snapshot,
9282 );
9283 set
9284 }
9285 Condition::IsNull { column_id } => {
9286 let mut set = if self.run_refs.len() == 1 {
9287 let mut r = self.open_reader(self.run_refs[0].run_id)?;
9288 r.null_row_id_set(*column_id, true)?
9289 } else {
9290 return self.null_scan(*column_id, true, snapshot);
9291 };
9292 set.remove_many(self.overlay_rid_set(snapshot));
9293 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
9294 set
9295 }
9296 Condition::IsNotNull { column_id } => {
9297 let mut set = if self.run_refs.len() == 1 {
9298 let mut r = self.open_reader(self.run_refs[0].run_id)?;
9299 r.null_row_id_set(*column_id, false)?
9300 } else {
9301 return self.null_scan(*column_id, false, snapshot);
9302 };
9303 set.remove_many(self.overlay_rid_set(snapshot));
9304 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
9305 set
9306 }
9307 })
9308 }
9309
9310 fn range_scan_i64(
9318 &self,
9319 column_id: u16,
9320 lo: i64,
9321 hi: i64,
9322 snapshot: Snapshot,
9323 ) -> Result<RowIdSet> {
9324 let mut row_ids = Vec::new();
9325 let mut tomb_rids: HashSet<u64> = HashSet::new();
9332 let overlay_rids = self.overlay_rid_set(snapshot);
9333 for rr in &self.run_refs {
9334 let mut reader = self.open_reader(rr.run_id)?;
9335 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
9336 for rid in matched {
9337 if !overlay_rids.contains(&rid) {
9338 row_ids.push(rid);
9339 }
9340 }
9341 for rid in reader.tombstoned_row_ids(snapshot.epoch)? {
9342 tomb_rids.insert(rid);
9343 }
9344 }
9345 let mut s = RowIdSet::from_unsorted(row_ids);
9346 s.remove_many(tomb_rids);
9347 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
9348 Ok(s)
9349 }
9350
9351 fn range_scan_f64(
9354 &self,
9355 column_id: u16,
9356 lo: f64,
9357 lo_inclusive: bool,
9358 hi: f64,
9359 hi_inclusive: bool,
9360 snapshot: Snapshot,
9361 ) -> Result<RowIdSet> {
9362 let mut row_ids = Vec::new();
9363 let overlay_rids = self.overlay_rid_set(snapshot);
9364 for rr in &self.run_refs {
9365 let mut reader = self.open_reader(rr.run_id)?;
9366 let matched = reader.range_row_ids_visible_f64(
9367 column_id,
9368 lo,
9369 lo_inclusive,
9370 hi,
9371 hi_inclusive,
9372 snapshot.epoch,
9373 )?;
9374 for rid in matched {
9375 if !overlay_rids.contains(&rid) {
9376 row_ids.push(rid);
9377 }
9378 }
9379 }
9380 let mut s = RowIdSet::from_unsorted(row_ids);
9381 self.range_scan_overlay_f64(
9382 &mut s,
9383 column_id,
9384 lo,
9385 lo_inclusive,
9386 hi,
9387 hi_inclusive,
9388 snapshot,
9389 );
9390 Ok(s)
9391 }
9392
9393 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
9395 let mut s = HashSet::new();
9396 for row in self.memtable.visible_versions_at(snapshot) {
9397 s.insert(row.row_id.0);
9398 }
9399 for row in self.mutable_run.visible_versions_at(snapshot) {
9400 s.insert(row.row_id.0);
9401 }
9402 s
9403 }
9404
9405 fn range_scan_overlay_i64(
9406 &self,
9407 s: &mut RowIdSet,
9408 column_id: u16,
9409 lo: i64,
9410 hi: i64,
9411 snapshot: Snapshot,
9412 ) {
9413 let mut newest: HashMap<u64, Row> = HashMap::new();
9420 for r in self.mutable_run.visible_versions_at(snapshot) {
9421 newest.insert(r.row_id.0, r);
9422 }
9423 for r in self.memtable.visible_versions_at(snapshot) {
9424 newest
9425 .entry(r.row_id.0)
9426 .and_modify(|cur| {
9427 if Snapshot::version_is_newer(
9428 r.committed_epoch,
9429 r.commit_ts,
9430 cur.committed_epoch,
9431 cur.commit_ts,
9432 ) {
9433 *cur = r.clone();
9434 }
9435 })
9436 .or_insert(r);
9437 }
9438 for row in newest.values() {
9439 if !row.deleted {
9440 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
9441 if *v >= lo && *v <= hi {
9442 s.insert(row.row_id.0);
9443 }
9444 }
9445 }
9446 }
9447 }
9448
9449 #[allow(clippy::too_many_arguments)]
9450 fn range_scan_overlay_f64(
9451 &self,
9452 s: &mut RowIdSet,
9453 column_id: u16,
9454 lo: f64,
9455 lo_inclusive: bool,
9456 hi: f64,
9457 hi_inclusive: bool,
9458 snapshot: Snapshot,
9459 ) {
9460 let mut newest: HashMap<u64, Row> = HashMap::new();
9463 for r in self.mutable_run.visible_versions_at(snapshot) {
9464 newest.insert(r.row_id.0, r);
9465 }
9466 for r in self.memtable.visible_versions_at(snapshot) {
9467 newest
9468 .entry(r.row_id.0)
9469 .and_modify(|cur| {
9470 if Snapshot::version_is_newer(
9471 r.committed_epoch,
9472 r.commit_ts,
9473 cur.committed_epoch,
9474 cur.commit_ts,
9475 ) {
9476 *cur = r.clone();
9477 }
9478 })
9479 .or_insert(r);
9480 }
9481 for row in newest.values() {
9482 if !row.deleted {
9483 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
9484 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
9485 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
9486 if ok_lo && ok_hi {
9487 s.insert(row.row_id.0);
9488 }
9489 }
9490 }
9491 }
9492 }
9493
9494 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
9497 let mut row_ids = Vec::new();
9498 let overlay_rids = self.overlay_rid_set(snapshot);
9499 for rr in &self.run_refs {
9500 let mut reader = self.open_reader(rr.run_id)?;
9501 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
9502 for rid in matched {
9503 if !overlay_rids.contains(&rid) {
9504 row_ids.push(rid);
9505 }
9506 }
9507 }
9508 let mut s = RowIdSet::from_unsorted(row_ids);
9509 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
9510 Ok(s)
9511 }
9512
9513 fn null_scan_overlay(
9517 &self,
9518 s: &mut RowIdSet,
9519 column_id: u16,
9520 want_nulls: bool,
9521 snapshot: Snapshot,
9522 ) {
9523 let mut newest: HashMap<u64, Row> = HashMap::new();
9524 for r in self.mutable_run.visible_versions_at(snapshot) {
9525 newest.insert(r.row_id.0, r);
9526 }
9527 for r in self.memtable.visible_versions_at(snapshot) {
9528 newest
9529 .entry(r.row_id.0)
9530 .and_modify(|cur| {
9531 if Snapshot::version_is_newer(
9532 r.committed_epoch,
9533 r.commit_ts,
9534 cur.committed_epoch,
9535 cur.commit_ts,
9536 ) {
9537 *cur = r.clone();
9538 }
9539 })
9540 .or_insert(r);
9541 }
9542 for row in newest.values() {
9543 if row.deleted {
9544 continue;
9545 }
9546 let is_null = !row.columns.contains_key(&column_id)
9547 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
9548 if is_null == want_nulls {
9549 s.insert(row.row_id.0);
9550 }
9551 }
9552 }
9553
9554 pub fn snapshot(&self) -> Snapshot {
9555 let epoch = self.epoch.visible();
9556 match &self.wal {
9560 WalSink::Shared(shared) => match shared.hlc.now() {
9561 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
9562 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
9564 },
9565 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
9566 }
9567 }
9568
9569 pub fn data_generation(&self) -> u64 {
9571 self.data_generation
9572 }
9573
9574 pub(crate) fn bump_data_generation(&mut self) {
9575 self.data_generation = self.data_generation.wrapping_add(1);
9576 }
9577
9578 pub fn table_id(&self) -> u64 {
9580 self.table_id
9581 }
9582
9583 fn seal_generations(&mut self) {
9589 self.memtable.seal();
9590 self.mutable_run.seal();
9591 self.hot.seal();
9592 for index in self.bitmap.values_mut() {
9593 index.seal();
9594 }
9595 for index in self.ann.values_mut() {
9596 index.seal();
9597 }
9598 for index in self.fm.values_mut() {
9599 index.seal();
9600 }
9601 for index in self.sparse.values_mut() {
9602 index.seal();
9603 }
9604 for index in self.minhash.values_mut() {
9605 index.seal();
9606 }
9607 self.pk_by_row.seal();
9608 }
9609
9610 fn capture_read_generation(&self) -> ReadGeneration {
9615 let visible_through = self.current_epoch();
9616 ReadGeneration {
9617 schema: Arc::new(self.schema.clone()),
9618 base_runs: Arc::new(self.run_refs.clone()),
9619 deltas: TableDeltas {
9620 memtable: self.memtable.clone(),
9621 mutable_run: self.mutable_run.clone(),
9622 hot: self.hot.clone(),
9623 pk_by_row: self.pk_by_row.clone(),
9624 },
9625 indexes: Arc::new(IndexGeneration::capture(
9626 &self.bitmap,
9627 &self.learned_range,
9628 &self.fm,
9629 &self.ann,
9630 &self.sparse,
9631 &self.minhash,
9632 visible_through,
9633 match &self.wal {
9635 WalSink::Shared(shared) => shared
9636 .hlc
9637 .now()
9638 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
9639 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
9640 },
9641 )),
9642 visible_through,
9643 }
9644 }
9645
9646 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
9651 self.ensure_indexes_complete()?;
9652 self.seal_generations();
9653 let view = Arc::new(self.capture_read_generation());
9654 self.published.store(Arc::clone(&view));
9655 Ok(view)
9656 }
9657
9658 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
9664 self.published.load_full()
9665 }
9666
9667 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
9669 &self.pins
9670 }
9671
9672 pub fn version_gc_floor(&self) -> Epoch {
9677 self.min_active_snapshot()
9678 .unwrap_or_else(|| self.current_epoch())
9679 }
9680
9681 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
9688 let mut report = self.pins.report();
9689 let transaction_floor = [
9690 self.pinned.keys().next().copied(),
9691 self.snapshots.min_pinned(),
9692 ]
9693 .into_iter()
9694 .flatten()
9695 .min();
9696 if let Some(epoch) = transaction_floor {
9697 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
9698 }
9699 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
9700 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
9701 }
9702 report
9703 }
9704
9705 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
9706 self.publish_read_generation()?;
9707 let mut generation = self.clone();
9708 generation.read_only = true;
9709 generation.wal = WalSink::ReadOnly;
9710 generation.pending_delete_rids.clear();
9711 generation.pending_put_cols.clear();
9712 generation.pending_rows.clear();
9713 generation.pending_rows_auto_inc.clear();
9714 generation.pending_dels.clear();
9715 generation.pending_truncate = None;
9716 generation.agg_cache = Arc::new(HashMap::new());
9717 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
9720 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
9723 crate::retention::PinSource::ReadGeneration,
9724 self.current_epoch(),
9725 )));
9726 Ok(generation)
9727 }
9728
9729 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
9730 (std::mem::size_of::<Self>() as u64)
9731 .saturating_add(self.memtable.approx_bytes())
9732 .saturating_add(self.mutable_run.approx_bytes())
9733 .saturating_add(self.live_count.saturating_mul(64))
9734 }
9735
9736 pub fn pin_snapshot(&mut self) -> Snapshot {
9743 let snap = self.snapshot();
9744 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
9745 snap
9746 }
9747
9748 fn active_pin_epochs_for_rebuild(&self) -> Vec<Epoch> {
9760 let mut set = BTreeSet::new();
9761 for epoch in self.pinned.keys().copied() {
9762 set.insert(epoch);
9763 }
9764 for epoch in self.snapshots.live_pinned_epochs() {
9765 set.insert(epoch);
9766 }
9767 for epoch in self.pins.live_pin_epochs() {
9768 set.insert(epoch);
9769 }
9770 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
9771 set.insert(floor);
9772 }
9773 set.into_iter().collect()
9774 }
9775
9776 pub fn hlc_gc_floor(
9786 &self,
9787 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
9788 ) -> crate::epoch::GcFloor {
9789 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
9790 epoch
9791 .and_then(&mut project_epoch)
9792 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
9793 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
9794 };
9795 let transaction = [
9796 self.pinned.keys().next().copied(),
9797 self.snapshots.min_pinned(),
9798 ]
9799 .into_iter()
9800 .flatten()
9801 .min();
9802 let history = self.snapshots.history_floor(self.current_epoch());
9803 crate::epoch::GcFloor {
9804 transaction_snapshot: project(transaction),
9805 history_retention: project(history),
9806 backup_pitr: project(
9807 self.pins
9808 .oldest_for(crate::retention::PinSource::BackupPitr),
9809 ),
9810 replication: project(
9811 self.pins
9812 .oldest_for(crate::retention::PinSource::Replication),
9813 ),
9814 read_generation: project(
9815 self.pins
9816 .oldest_for(crate::retention::PinSource::ReadGeneration),
9817 ),
9818 online_index_build: project(
9819 self.pins
9820 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
9821 ),
9822 }
9823 }
9824
9825 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
9827 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
9828 *count -= 1;
9829 if *count == 0 {
9830 self.pinned.remove(&snap.epoch);
9831 }
9832 }
9833 }
9834
9835 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
9847 let local = self.pinned.keys().next().copied();
9848 let global = self.snapshots.min_pinned();
9849 let history = self.snapshots.history_floor(self.current_epoch());
9850 let pinned = self.pins.oldest_pinned();
9851 [local, global, history, pinned].into_iter().flatten().min()
9852 }
9853
9854 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
9858 self.ensure_writable()?;
9859 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
9860 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
9861 }
9862
9863 pub fn clear_ttl(&mut self) -> Result<()> {
9864 self.ensure_writable()?;
9865 self.apply_ttl_policy_at(None, self.current_epoch())
9866 }
9867
9868 pub fn ttl(&self) -> Option<TtlPolicy> {
9869 self.ttl
9870 }
9871
9872 pub(crate) fn prepare_ttl_policy(
9873 &self,
9874 column_name: &str,
9875 duration_nanos: u64,
9876 ) -> Result<TtlPolicy> {
9877 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
9878 return Err(MongrelError::InvalidArgument(
9879 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
9880 ));
9881 }
9882 let column = self
9883 .schema
9884 .columns
9885 .iter()
9886 .find(|column| column.name == column_name)
9887 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
9888 if column.ty != TypeId::TimestampNanos {
9889 return Err(MongrelError::Schema(format!(
9890 "TTL column {column_name} must be TimestampNanos, is {:?}",
9891 column.ty
9892 )));
9893 }
9894 Ok(TtlPolicy {
9895 column_id: column.id,
9896 duration_nanos,
9897 })
9898 }
9899
9900 pub(crate) fn apply_ttl_policy_at(
9901 &mut self,
9902 policy: Option<TtlPolicy>,
9903 epoch: Epoch,
9904 ) -> Result<()> {
9905 if let Some(policy) = policy {
9906 let column = self
9907 .schema
9908 .columns
9909 .iter()
9910 .find(|column| column.id == policy.column_id)
9911 .ok_or_else(|| {
9912 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
9913 })?;
9914 if column.ty != TypeId::TimestampNanos
9915 || policy.duration_nanos == 0
9916 || policy.duration_nanos > i64::MAX as u64
9917 {
9918 return Err(MongrelError::Schema("invalid TTL policy".into()));
9919 }
9920 }
9921 self.ttl = policy;
9922 self.agg_cache = Arc::new(HashMap::new());
9923 self.clear_result_cache();
9924 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
9925 self.persist_manifest(epoch)
9926 }
9927
9928 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
9929 let Some(policy) = self.ttl else {
9930 return false;
9931 };
9932 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
9933 return false;
9934 };
9935 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
9936 }
9937
9938 pub fn current_epoch(&self) -> Epoch {
9939 self.epoch.visible()
9940 }
9941
9942 pub fn memtable_len(&self) -> usize {
9943 self.memtable.len()
9944 }
9945
9946 pub fn count(&self) -> u64 {
9949 if self.ttl.is_none()
9950 && self.pending_put_cols.is_empty()
9951 && self.pending_delete_rids.is_empty()
9952 && self.pending_rows.is_empty()
9953 && self.pending_dels.is_empty()
9954 && self.pending_truncate.is_none()
9955 {
9956 self.live_count
9957 } else {
9958 self.visible_rows(self.snapshot())
9959 .map(|rows| rows.len() as u64)
9960 .unwrap_or(self.live_count)
9961 }
9962 }
9963
9964 pub fn count_conditions(
9968 &mut self,
9969 conditions: &[crate::query::Condition],
9970 snapshot: Snapshot,
9971 ) -> Result<Option<u64>> {
9972 use crate::query::Condition;
9973 if self.ttl.is_some() {
9974 if conditions.is_empty() {
9975 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
9976 }
9977 let mut sets = Vec::with_capacity(conditions.len());
9978 for condition in conditions {
9979 sets.push(self.resolve_condition(condition, snapshot)?);
9980 }
9981 let survivors = RowIdSet::intersect_many(sets);
9982 let rows = self.visible_rows(snapshot)?;
9983 return Ok(Some(
9984 rows.into_iter()
9985 .filter(|row| survivors.contains(row.row_id.0))
9986 .count() as u64,
9987 ));
9988 }
9989 if conditions.is_empty() {
9990 return Ok(Some(self.count()));
9991 }
9992 let served = |c: &Condition| {
9993 matches!(
9994 c,
9995 Condition::Pk(_)
9996 | Condition::BitmapEq { .. }
9997 | Condition::BitmapIn { .. }
9998 | Condition::BytesPrefix { .. }
9999 | Condition::FmContains { .. }
10000 | Condition::FmContainsAll { .. }
10001 | Condition::Ann { .. }
10002 | Condition::Range { .. }
10003 | Condition::RangeF64 { .. }
10004 | Condition::SparseMatch { .. }
10005 | Condition::MinHashSimilar { .. }
10006 | Condition::IsNull { .. }
10007 | Condition::IsNotNull { .. }
10008 )
10009 };
10010 if !conditions.iter().all(served) {
10011 return Ok(None);
10012 }
10013 self.ensure_indexes_complete()?;
10014 if !self.pending_put_cols.is_empty()
10015 || !self.pending_delete_rids.is_empty()
10016 || !self.pending_rows.is_empty()
10017 || !self.pending_dels.is_empty()
10018 || self.pending_truncate.is_some()
10019 {
10020 let mut sets = Vec::with_capacity(conditions.len());
10021 for condition in conditions {
10022 sets.push(self.resolve_condition(condition, snapshot)?);
10023 }
10024 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
10025 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
10026 }
10027 let mut sets = Vec::with_capacity(conditions.len());
10028 for condition in conditions {
10029 sets.push(self.resolve_condition(condition, snapshot)?);
10030 }
10031 let rids = RowIdSet::intersect_many(sets);
10032 if self.had_deletes || !self.memtable.is_empty() || !self.mutable_run.is_empty() {
10040 let sorted = rids.into_sorted_vec();
10041 let count = self.rows_for_rids(&sorted, snapshot)?.len() as u64;
10042 crate::trace::QueryTrace::record(|t| {
10043 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
10044 t.survivor_count = Some(count as usize);
10045 t.conditions_pushed = conditions.len();
10046 });
10047 return Ok(Some(count));
10048 }
10049 let count = rids.len() as u64;
10050 crate::trace::QueryTrace::record(|t| {
10051 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
10052 t.survivor_count = Some(count as usize);
10053 t.conditions_pushed = conditions.len();
10054 });
10055 Ok(Some(count))
10056 }
10057
10058 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
10064 let mut out = Vec::new();
10065 for row in self.memtable.visible_versions(snapshot.epoch) {
10066 if row.deleted {
10067 out.push(row.row_id.0);
10068 }
10069 }
10070 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10071 if row.deleted {
10072 out.push(row.row_id.0);
10073 }
10074 }
10075 out
10076 }
10077
10078 pub fn bulk_load_columns(
10087 &mut self,
10088 user_columns: Vec<(u16, columnar::NativeColumn)>,
10089 ) -> Result<Epoch> {
10090 self.bulk_load_columns_with(user_columns, 3, false, true)
10091 }
10092
10093 pub fn bulk_load_fast(
10100 &mut self,
10101 user_columns: Vec<(u16, columnar::NativeColumn)>,
10102 ) -> Result<Epoch> {
10103 self.bulk_load_columns_with(user_columns, -1, true, false)
10104 }
10105
10106 fn bulk_load_columns_with(
10107 &mut self,
10108 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
10109 zstd_level: i32,
10110 force_plain: bool,
10111 lz4: bool,
10112 ) -> Result<Epoch> {
10113 self.ensure_writable()?;
10114 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
10115 if n == 0 {
10116 return Ok(self.current_epoch());
10117 }
10118 let epoch = self.commit_new_epoch()?;
10119 let live_before = self.live_count;
10120 self.spill_mutable_run(epoch)?;
10122 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
10123 && self.indexes_complete
10124 && self.run_refs.is_empty()
10125 && self.memtable.is_empty()
10126 && self.mutable_run.is_empty();
10127 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
10130 self.validate_columns_not_null(&user_columns, n)?;
10131 let winner_idx = self
10132 .bulk_pk_winner_indices(&user_columns, n)
10133 .filter(|idx| idx.len() != n);
10134 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
10135 match winner_idx.as_deref() {
10136 Some(idx) => {
10137 let compacted = user_columns
10138 .iter()
10139 .map(|(id, c)| (*id, c.gather(idx)))
10140 .collect();
10141 (compacted, idx.len())
10142 }
10143 None => (user_columns, n),
10144 };
10145 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
10146 let first = self.allocator.alloc_range(write_n as u64)?.0;
10147 for rid in first..first + write_n as u64 {
10148 self.reservoir.offer(rid);
10149 }
10150 let run_id = self.alloc_run_id()?;
10151 let path = self.run_path(run_id);
10152 let mut writer =
10153 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
10154 if force_plain {
10155 writer = writer.with_plain();
10156 } else if lz4 {
10157 writer = writer.with_lz4();
10160 } else {
10161 writer = writer.with_zstd_level(zstd_level);
10162 }
10163 if let Some(kek) = &self.kek {
10164 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
10165 }
10166 let header = match self.create_run_file(run_id)? {
10167 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
10168 None => writer.write_native(&path, &write_columns, write_n, first)?,
10169 };
10170 self.run_refs.push(RunRef {
10171 run_id: run_id as u128,
10172 level: 0,
10173 epoch_created: epoch.0,
10174 row_count: header.row_count,
10175 });
10176 self.run_row_id_ranges
10177 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
10178 self.live_count = self.live_count.saturating_add(write_n as u64);
10179 if eager_index_build {
10180 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
10181 self.index_columns_bulk(&write_columns, &row_ids);
10182 self.indexes_complete = true;
10183 self.build_learned_ranges()?;
10184 } else {
10185 self.indexes_complete = false;
10189 }
10190 self.mark_flushed(epoch)?;
10191 self.persist_manifest(epoch)?;
10192 if eager_index_build {
10193 self.checkpoint_indexes(epoch);
10194 }
10195 self.clear_result_cache();
10196 self.data_generation = self.data_generation.wrapping_add(1);
10197 Ok(epoch)
10198 }
10199
10200 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
10218 let n = row_ids.len();
10219 if n == 0 {
10220 return;
10221 }
10222 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
10223 columns.iter().map(|(id, c)| (*id, c)).collect();
10224 let ty_of: std::collections::HashMap<u16, TypeId> = self
10225 .schema
10226 .columns
10227 .iter()
10228 .map(|c| (c.id, c.ty.clone()))
10229 .collect();
10230 let pk_id = self.schema.primary_key().map(|c| c.id);
10231
10232 for (i, &rid) in row_ids.iter().enumerate() {
10233 let row_id = RowId(rid);
10234 if let Some(pid) = pk_id {
10235 if let Some(col) = by_id.get(&pid) {
10236 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
10237 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
10238 self.insert_hot_pk(key, row_id);
10239 }
10240 }
10241 }
10242 for idef in &self.schema.indexes {
10243 let Some(col) = by_id.get(&idef.column_id) else {
10244 continue;
10245 };
10246 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
10247 match idef.kind {
10248 IndexKind::Bitmap => {
10249 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
10250 if let Some(key) =
10251 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
10252 {
10253 b.insert(key, row_id);
10254 }
10255 }
10256 }
10257 IndexKind::FmIndex => {
10258 if let Some(f) = self.fm.get_mut(&idef.column_id) {
10259 if let Some(bytes) = columnar::native_bytes_at(col, i) {
10260 f.insert(bytes.to_vec(), row_id);
10261 }
10262 }
10263 }
10264 IndexKind::Sparse => {
10265 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
10266 if let Some(bytes) = columnar::native_bytes_at(col, i) {
10267 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
10268 s.insert(&terms, row_id);
10269 }
10270 }
10271 }
10272 }
10273 IndexKind::MinHash => {
10274 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
10275 if let Some(bytes) = columnar::native_bytes_at(col, i) {
10276 let tokens = crate::index::token_hashes_from_bytes(bytes);
10277 mh.insert(&tokens, row_id);
10278 }
10279 }
10280 }
10281 _ => {}
10282 }
10283 }
10284 }
10285 }
10286
10287 pub fn visible_columns_native(
10292 &self,
10293 snapshot: Snapshot,
10294 projection: Option<&[u16]>,
10295 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
10296 self.visible_columns_native_inner(snapshot, projection, None)
10297 }
10298
10299 pub fn visible_columns_native_with_control(
10300 &self,
10301 snapshot: Snapshot,
10302 projection: Option<&[u16]>,
10303 control: &crate::ExecutionControl,
10304 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
10305 self.visible_columns_native_inner(snapshot, projection, Some(control))
10306 }
10307
10308 fn visible_columns_native_inner(
10309 &self,
10310 snapshot: Snapshot,
10311 projection: Option<&[u16]>,
10312 control: Option<&crate::ExecutionControl>,
10313 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
10314 execution_checkpoint(control, 0)?;
10315 let wanted: Vec<u16> = match projection {
10316 Some(p) => p.to_vec(),
10317 None => self.schema.columns.iter().map(|c| c.id).collect(),
10318 };
10319 if self.ttl.is_none()
10320 && self.memtable.is_empty()
10321 && self.mutable_run.is_empty()
10322 && self.run_refs.len() == 1
10323 {
10324 let rr = self.run_refs[0].clone();
10325 let mut reader = self.open_reader(rr.run_id)?;
10326 let idxs = reader.visible_indices_native(snapshot.epoch)?;
10327 execution_checkpoint(control, 0)?;
10328 let all_visible = idxs.len() == reader.row_count();
10329 if reader.has_mmap() && control.is_none() {
10335 use rayon::prelude::*;
10336 let valid: Vec<u16> = wanted
10339 .iter()
10340 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
10341 .copied()
10342 .collect();
10343 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
10345 .par_iter()
10346 .filter_map(|cid| {
10347 reader
10348 .column_native_shared(*cid)
10349 .ok()
10350 .map(|col| (*cid, col))
10351 })
10352 .collect();
10353 let cols = decoded
10354 .into_iter()
10355 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
10356 .collect();
10357 return Ok(cols);
10358 }
10359 let mut cols = Vec::with_capacity(wanted.len());
10360 for (index, cid) in wanted.iter().enumerate() {
10361 execution_checkpoint(control, index)?;
10362 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
10363 Some(c) => c,
10364 None => continue,
10365 };
10366 let col = reader.column_native(cdef.id)?;
10367 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
10368 }
10369 return Ok(cols);
10370 }
10371 let vcols = self.visible_columns(snapshot)?;
10372 execution_checkpoint(control, 0)?;
10373 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
10374 let out: Vec<(u16, columnar::NativeColumn)> = vcols
10375 .into_iter()
10376 .filter(|(id, _)| want_set.contains(id))
10377 .map(|(id, vals)| {
10378 let ty = self
10379 .schema
10380 .columns
10381 .iter()
10382 .find(|c| c.id == id)
10383 .map(|c| c.ty.clone())
10384 .unwrap_or(TypeId::Bytes);
10385 (id, columnar::values_to_native(ty, &vals))
10386 })
10387 .collect();
10388 Ok(out)
10389 }
10390
10391 pub fn run_count(&self) -> usize {
10392 self.run_refs.len()
10393 }
10394
10395 pub fn memtable_is_empty(&self) -> bool {
10397 self.memtable.is_empty()
10398 }
10399
10400 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
10404 self.page_cache.stats()
10405 }
10406
10407 pub fn reset_page_cache_stats(&self) {
10409 self.page_cache.reset_stats();
10410 }
10411
10412 pub fn run_ids(&self) -> Vec<u128> {
10415 self.run_refs.iter().map(|r| r.run_id).collect()
10416 }
10417
10418 pub fn single_run_is_clean(&self) -> bool {
10422 if self.ttl.is_some() || self.run_refs.len() != 1 {
10423 return false;
10424 }
10425 self.open_reader(self.run_refs[0].run_id)
10426 .map(|r| r.is_clean())
10427 .unwrap_or(false)
10428 }
10429
10430 fn resolve_footprint(
10437 &self,
10438 conditions: &[crate::query::Condition],
10439 snapshot: Snapshot,
10440 ) -> roaring::RoaringBitmap {
10441 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
10442 return roaring::RoaringBitmap::new();
10443 }
10444 if self.run_refs.is_empty() {
10445 return roaring::RoaringBitmap::new();
10446 }
10447 if self.run_refs.len() == 1 {
10449 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
10450 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
10451 return rids.to_roaring_lossy();
10452 }
10453 }
10454 }
10455 roaring::RoaringBitmap::new()
10456 }
10457
10458 pub fn query_columns_native_cached(
10469 &mut self,
10470 conditions: &[crate::query::Condition],
10471 projection: Option<&[u16]>,
10472 snapshot: Snapshot,
10473 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10474 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
10475 }
10476
10477 pub fn query_columns_native_cached_with_control(
10478 &mut self,
10479 conditions: &[crate::query::Condition],
10480 projection: Option<&[u16]>,
10481 snapshot: Snapshot,
10482 control: &crate::ExecutionControl,
10483 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10484 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
10485 }
10486
10487 fn query_columns_native_cached_inner(
10488 &mut self,
10489 conditions: &[crate::query::Condition],
10490 projection: Option<&[u16]>,
10491 snapshot: Snapshot,
10492 control: Option<&crate::ExecutionControl>,
10493 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10494 execution_checkpoint(control, 0)?;
10495 if self.ttl.is_some() {
10498 return self.query_columns_native_inner(conditions, projection, snapshot, control);
10499 }
10500 if conditions.is_empty() {
10501 return self.query_columns_native_inner(conditions, projection, snapshot, control);
10502 }
10503 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
10507 if let Some(hit) = self.result_cache.lock().get_columns(key) {
10508 crate::trace::QueryTrace::record(|t| {
10509 t.result_cache_hit = true;
10510 t.scan_mode = crate::trace::ScanMode::NativePushdown;
10511 });
10512 return Ok(Some((*hit).clone()));
10513 }
10514 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
10515 execution_checkpoint(control, 0)?;
10516 if let Some(cols) = &res {
10517 let footprint = self.resolve_footprint(conditions, snapshot);
10518 let condition_cols = crate::query::condition_columns(conditions);
10519 execution_checkpoint(control, 0)?;
10520 self.result_cache.lock().insert(
10521 key,
10522 CachedEntry {
10523 data: CachedData::Columns(Arc::new(cols.clone())),
10524 footprint,
10525 condition_cols,
10526 },
10527 );
10528 }
10529 Ok(res)
10530 }
10531
10532 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
10537 if self.ttl.is_some() {
10538 return self.query(q);
10539 }
10540 if q.conditions.is_empty() {
10541 return self.query(q);
10542 }
10543 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
10544 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
10545 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
10546 if let Some(hit) = self.result_cache.lock().get_rows(key) {
10547 crate::trace::QueryTrace::record(|t| {
10548 t.result_cache_hit = true;
10549 t.scan_mode = crate::trace::ScanMode::Materialized;
10550 });
10551 return Ok((*hit).clone());
10552 }
10553 let rows = self.query(q)?;
10554 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
10555 let condition_cols = crate::query::condition_columns(&q.conditions);
10556 self.result_cache.lock().insert(
10557 key,
10558 CachedEntry {
10559 data: CachedData::Rows(Arc::new(rows.clone())),
10560 footprint,
10561 condition_cols,
10562 },
10563 );
10564 Ok(rows)
10565 }
10566
10567 #[allow(clippy::type_complexity)]
10582 pub fn query_columns_native_traced(
10583 &mut self,
10584 conditions: &[crate::query::Condition],
10585 projection: Option<&[u16]>,
10586 snapshot: Snapshot,
10587 ) -> Result<(
10588 Option<Vec<(u16, columnar::NativeColumn)>>,
10589 crate::trace::QueryTrace,
10590 )> {
10591 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10592 self.query_columns_native(conditions, projection, snapshot)
10593 });
10594 Ok((result?, trace))
10595 }
10596
10597 #[allow(clippy::type_complexity)]
10600 pub fn query_columns_native_cached_traced(
10601 &mut self,
10602 conditions: &[crate::query::Condition],
10603 projection: Option<&[u16]>,
10604 snapshot: Snapshot,
10605 ) -> Result<(
10606 Option<Vec<(u16, columnar::NativeColumn)>>,
10607 crate::trace::QueryTrace,
10608 )> {
10609 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10610 self.query_columns_native_cached(conditions, projection, snapshot)
10611 });
10612 Ok((result?, trace))
10613 }
10614
10615 pub fn native_page_cursor_traced(
10617 &self,
10618 snapshot: Snapshot,
10619 projection: Vec<(u16, TypeId)>,
10620 conditions: &[crate::query::Condition],
10621 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
10622 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10623 self.native_page_cursor(snapshot, projection, conditions)
10624 });
10625 Ok((result?, trace))
10626 }
10627
10628 pub fn native_multi_run_cursor_traced(
10630 &self,
10631 snapshot: Snapshot,
10632 projection: Vec<(u16, TypeId)>,
10633 conditions: &[crate::query::Condition],
10634 ) -> Result<(
10635 Option<crate::cursor::MultiRunCursor>,
10636 crate::trace::QueryTrace,
10637 )> {
10638 let (result, trace) = crate::trace::QueryTrace::capture(|| {
10639 self.native_multi_run_cursor(snapshot, projection, conditions)
10640 });
10641 Ok((result?, trace))
10642 }
10643
10644 pub fn count_conditions_traced(
10646 &mut self,
10647 conditions: &[crate::query::Condition],
10648 snapshot: Snapshot,
10649 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
10650 let (result, trace) =
10651 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
10652 Ok((result?, trace))
10653 }
10654
10655 pub fn query_traced(
10657 &mut self,
10658 q: &crate::query::Query,
10659 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
10660 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
10661 Ok((result?, trace))
10662 }
10663
10664 pub fn query_columns_native(
10669 &mut self,
10670 conditions: &[crate::query::Condition],
10671 projection: Option<&[u16]>,
10672 snapshot: Snapshot,
10673 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10674 self.query_columns_native_inner(conditions, projection, snapshot, None)
10675 }
10676
10677 pub fn query_columns_native_with_control(
10678 &mut self,
10679 conditions: &[crate::query::Condition],
10680 projection: Option<&[u16]>,
10681 snapshot: Snapshot,
10682 control: &crate::ExecutionControl,
10683 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10684 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
10685 }
10686
10687 fn query_columns_native_inner(
10688 &mut self,
10689 conditions: &[crate::query::Condition],
10690 projection: Option<&[u16]>,
10691 snapshot: Snapshot,
10692 control: Option<&crate::ExecutionControl>,
10693 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
10694 use crate::query::Condition;
10695 execution_checkpoint(control, 0)?;
10696 if self.ttl.is_some() {
10699 return Ok(None);
10700 }
10701 if conditions.is_empty() {
10702 return Ok(None);
10703 }
10704 self.ensure_indexes_complete()?;
10705
10706 let served = |c: &Condition| {
10711 matches!(
10712 c,
10713 Condition::Pk(_)
10714 | Condition::BitmapEq { .. }
10715 | Condition::BitmapIn { .. }
10716 | Condition::BytesPrefix { .. }
10717 | Condition::FmContains { .. }
10718 | Condition::FmContainsAll { .. }
10719 | Condition::Ann { .. }
10720 | Condition::Range { .. }
10721 | Condition::RangeF64 { .. }
10722 | Condition::SparseMatch { .. }
10723 | Condition::MinHashSimilar { .. }
10724 | Condition::IsNull { .. }
10725 | Condition::IsNotNull { .. }
10726 )
10727 };
10728 if !conditions.iter().all(served) {
10729 return Ok(None);
10730 }
10731 let fast_path =
10732 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
10733 crate::trace::QueryTrace::record(|t| {
10734 t.run_count = self.run_refs.len();
10735 t.memtable_rows = self.memtable.len();
10736 t.mutable_run_rows = self.mutable_run.len();
10737 t.conditions_pushed = conditions.len();
10738 t.learned_range_used = conditions.iter().any(|c| match c {
10739 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
10740 self.learned_range.contains_key(column_id)
10741 }
10742 _ => false,
10743 });
10744 });
10745 let col_ids: Vec<u16> = projection
10747 .map(|p| p.to_vec())
10748 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
10749 let proj_pairs: Vec<(u16, TypeId)> = col_ids
10750 .iter()
10751 .map(|&cid| {
10752 let ty = self
10753 .schema
10754 .columns
10755 .iter()
10756 .find(|c| c.id == cid)
10757 .map(|c| c.ty.clone())
10758 .unwrap_or(TypeId::Bytes);
10759 (cid, ty)
10760 })
10761 .collect();
10762
10763 if fast_path {
10769 let needs_column = conditions.iter().any(|c| match c {
10772 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
10773 Condition::RangeF64 { column_id, .. } => {
10774 !self.learned_range.contains_key(column_id)
10775 }
10776 _ => false,
10777 });
10778 let mut reader_opt: Option<RunReader> = if needs_column {
10779 Some(self.open_reader(self.run_refs[0].run_id)?)
10780 } else {
10781 None
10782 };
10783 let mut sets: Vec<RowIdSet> = Vec::new();
10784 for (index, c) in conditions.iter().enumerate() {
10785 execution_checkpoint(control, index)?;
10786 let s = match c {
10787 Condition::Range { column_id, lo, hi }
10788 if !self.learned_range.contains_key(column_id) =>
10789 {
10790 if reader_opt.is_none() {
10791 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
10792 }
10793 reader_opt
10794 .as_mut()
10795 .expect("reader opened for range")
10796 .range_row_id_set_i64(*column_id, *lo, *hi)?
10797 }
10798 Condition::RangeF64 {
10799 column_id,
10800 lo,
10801 lo_inclusive,
10802 hi,
10803 hi_inclusive,
10804 } if !self.learned_range.contains_key(column_id) => {
10805 if reader_opt.is_none() {
10806 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
10807 }
10808 reader_opt
10809 .as_mut()
10810 .expect("reader opened for range")
10811 .range_row_id_set_f64(
10812 *column_id,
10813 *lo,
10814 *lo_inclusive,
10815 *hi,
10816 *hi_inclusive,
10817 )?
10818 }
10819 _ => self.resolve_condition(c, snapshot)?,
10820 };
10821 sets.push(s);
10822 }
10823 let candidates = RowIdSet::intersect_many(sets);
10824 crate::trace::QueryTrace::record(|t| {
10825 t.survivor_count = Some(candidates.len());
10826 });
10827 if candidates.is_empty() {
10828 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
10829 .iter()
10830 .map(|&id| {
10831 (
10832 id,
10833 columnar::null_native(
10834 proj_pairs
10835 .iter()
10836 .find(|(c, _)| c == &id)
10837 .map(|(_, t)| t.clone())
10838 .unwrap_or(TypeId::Bytes),
10839 0,
10840 ),
10841 )
10842 })
10843 .collect();
10844 return Ok(Some(cols));
10845 }
10846 let mut reader = match reader_opt.take() {
10847 Some(r) => r,
10848 None => self.open_reader(self.run_refs[0].run_id)?,
10849 };
10850 let candidate_ids = candidates.into_sorted_vec();
10851 let (positions, fast_rid) = if let Some(positions) =
10852 reader.positions_for_row_ids_fast(&candidate_ids)
10853 {
10854 (positions, true)
10855 } else {
10856 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
10857 match col {
10858 columnar::NativeColumn::Int64 { data, .. } => {
10859 let mut p = Vec::with_capacity(candidate_ids.len());
10860 for (index, rid) in candidate_ids.iter().enumerate() {
10861 execution_checkpoint(control, index)?;
10862 if let Ok(position) = data.binary_search(&(*rid as i64)) {
10863 p.push(position);
10864 }
10865 }
10866 p.sort_unstable();
10867 (p, false)
10868 }
10869 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
10870 }
10871 };
10872 crate::trace::QueryTrace::record(|t| {
10873 t.scan_mode = crate::trace::ScanMode::NativePushdown;
10874 t.fast_row_id_map = fast_rid;
10875 });
10876 let mut cols = Vec::with_capacity(col_ids.len());
10877 for (index, cid) in col_ids.iter().enumerate() {
10878 execution_checkpoint(control, index)?;
10879 let col = reader.column_native(*cid)?;
10880 cols.push((*cid, col.gather(&positions)));
10881 }
10882 return Ok(Some(cols));
10883 }
10884
10885 if !self.run_refs.is_empty() {
10898 use crate::cursor::{
10899 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
10900 };
10901 let remaining: usize;
10902 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
10903 let c = self
10904 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
10905 .expect("single-run cursor should build when run_refs.len() == 1");
10906 remaining = c.remaining_rows();
10907 Box::new(c)
10908 } else {
10909 let c = self
10910 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
10911 .expect("multi-run cursor should build when run_refs.len() >= 1");
10912 remaining = c.remaining_rows();
10913 Box::new(c)
10914 };
10915 crate::trace::QueryTrace::record(|t| {
10916 if t.survivor_count.is_none() {
10917 t.survivor_count = Some(remaining);
10918 }
10919 });
10920 let cols = match control {
10921 Some(control) => {
10922 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
10923 }
10924 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
10925 };
10926 return Ok(Some(cols));
10927 }
10928
10929 crate::trace::QueryTrace::record(|t| {
10934 t.scan_mode = crate::trace::ScanMode::Materialized;
10935 t.row_materialized = true;
10936 });
10937 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
10938 for (index, c) in conditions.iter().enumerate() {
10939 execution_checkpoint(control, index)?;
10940 sets.push(self.resolve_condition(c, snapshot)?);
10941 }
10942 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
10943 let rows = self.rows_for_rids(&rids, snapshot)?;
10944 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
10945 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
10946 execution_checkpoint(control, index)?;
10947 let vals: Vec<Value> = rows
10948 .iter()
10949 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
10950 .collect();
10951 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
10952 }
10953 Ok(Some(cols))
10954 }
10955
10956 pub fn native_page_cursor(
10971 &self,
10972 snapshot: Snapshot,
10973 projection: Vec<(u16, TypeId)>,
10974 conditions: &[crate::query::Condition],
10975 ) -> Result<Option<NativePageCursor>> {
10976 use crate::cursor::build_page_plans;
10977 if self.ttl.is_some() {
10978 return Ok(None);
10979 }
10980 if !conditions.is_empty() && !self.indexes_complete {
10983 return Ok(None);
10984 }
10985 if self.run_refs.len() != 1 {
10986 return Ok(None);
10987 }
10988 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
10989 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
10990
10991 let overlay_rids: HashSet<u64> = {
10994 let mut s = HashSet::new();
10995 for row in self.memtable.visible_versions(snapshot.epoch) {
10996 s.insert(row.row_id.0);
10997 }
10998 for row in self.mutable_run.visible_versions(snapshot.epoch) {
10999 s.insert(row.row_id.0);
11000 }
11001 s
11002 };
11003
11004 let survivors = if conditions.is_empty() {
11008 None
11009 } else {
11010 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
11011 };
11012
11013 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
11020 survivors.clone()
11021 } else if let Some(s) = &survivors {
11022 let mut run_set = s.clone();
11023 run_set.remove_many(overlay_rids.iter().copied());
11024 Some(run_set)
11025 } else {
11026 Some(RowIdSet::from_unsorted(
11027 rids.iter()
11028 .map(|&r| r as u64)
11029 .filter(|r| !overlay_rids.contains(r))
11030 .collect(),
11031 ))
11032 };
11033
11034 let overlay_rows = if overlay_rids.is_empty() {
11035 Vec::new()
11036 } else {
11037 let bound = Self::overlay_materialization_bound(conditions, &survivors);
11038 self.overlay_visible_rows(snapshot, bound)
11039 };
11040
11041 let plans = if positions.is_empty() {
11043 Vec::new()
11044 } else {
11045 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
11046 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
11047 };
11048
11049 let overlay = if overlay_rows.is_empty() {
11051 None
11052 } else {
11053 let filtered =
11054 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
11055 if filtered.is_empty() {
11056 None
11057 } else {
11058 Some(self.materialize_overlay(&filtered, &projection))
11059 }
11060 };
11061
11062 let overlay_row_count = overlay
11063 .as_ref()
11064 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
11065 .unwrap_or(0);
11066 crate::trace::QueryTrace::record(|t| {
11067 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
11068 t.run_count = self.run_refs.len();
11069 t.memtable_rows = self.memtable.len();
11070 t.mutable_run_rows = self.mutable_run.len();
11071 t.overlay_rows = overlay_row_count;
11072 t.conditions_pushed = conditions.len();
11073 t.pages_decoded = plans
11074 .iter()
11075 .map(|p| p.positions.len())
11076 .sum::<usize>()
11077 .min(1);
11078 });
11079
11080 Ok(Some(NativePageCursor::new_with_overlay(
11081 reader, projection, plans, overlay,
11082 )))
11083 }
11084 #[allow(clippy::type_complexity)]
11094 pub fn native_multi_run_cursor(
11095 &self,
11096 snapshot: Snapshot,
11097 projection: Vec<(u16, TypeId)>,
11098 conditions: &[crate::query::Condition],
11099 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
11100 use crate::cursor::{MultiRunCursor, RunStream};
11101 use crate::sorted_run::SYS_ROW_ID;
11102 use std::collections::{BinaryHeap, HashMap, HashSet};
11103 if self.ttl.is_some() {
11104 return Ok(None);
11105 }
11106 if !conditions.is_empty() && !self.indexes_complete {
11109 return Ok(None);
11110 }
11111 if self.run_refs.is_empty() {
11112 return Ok(None);
11113 }
11114
11115 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
11117 Vec::with_capacity(self.run_refs.len());
11118 for rr in &self.run_refs {
11119 let mut reader = self.open_reader(rr.run_id)?;
11120 let (rids, eps, del) = reader.system_columns_native()?;
11121 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
11122 run_meta.push((reader, rids, eps, del, page_rows));
11123 }
11124
11125 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
11129 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
11130 for i in 0..rids.len() {
11131 let rid = rids[i] as u64;
11132 let e = eps[i] as u64;
11133 if e > snapshot.epoch.0 {
11134 continue;
11135 }
11136 let is_del = del[i] != 0;
11137 best.entry(rid)
11138 .and_modify(|cur| {
11139 if e > cur.0 {
11140 *cur = (e, run_idx, i, is_del);
11141 }
11142 })
11143 .or_insert((e, run_idx, i, is_del));
11144 }
11145 }
11146
11147 let overlay_rids: HashSet<u64> = {
11149 let mut s = HashSet::new();
11150 for row in self.memtable.visible_versions(snapshot.epoch) {
11151 s.insert(row.row_id.0);
11152 }
11153 for row in self.mutable_run.visible_versions(snapshot.epoch) {
11154 s.insert(row.row_id.0);
11155 }
11156 s
11157 };
11158
11159 let survivors: Option<RowIdSet> = if conditions.is_empty() {
11161 None
11162 } else {
11163 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
11164 for c in conditions {
11165 sets.push(self.resolve_condition(c, snapshot)?);
11166 }
11167 Some(RowIdSet::intersect_many(sets))
11168 };
11169
11170 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
11174 for (rid, (_, run_idx, pos, deleted)) in &best {
11175 if *deleted {
11176 continue;
11177 }
11178 if overlay_rids.contains(rid) {
11179 continue;
11180 }
11181 if let Some(s) = &survivors {
11182 if !s.contains(*rid) {
11183 continue;
11184 }
11185 }
11186 per_run[*run_idx].push((*rid, *pos));
11187 }
11188 for v in per_run.iter_mut() {
11189 v.sort_unstable_by_key(|&(rid, _)| rid);
11190 }
11191
11192 let mut streams = Vec::with_capacity(run_meta.len());
11194 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
11195 let mut total = 0usize;
11196 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
11197 let mut starts = Vec::with_capacity(page_rows.len());
11198 let mut acc = 0usize;
11199 for &r in &page_rows {
11200 starts.push(acc);
11201 acc += r;
11202 }
11203 let mut survivors_vec: Vec<(u64, usize, usize)> =
11204 Vec::with_capacity(per_run[run_idx].len());
11205 for &(rid, pos) in &per_run[run_idx] {
11206 let page_seq = match starts.partition_point(|&s| s <= pos) {
11207 0 => continue,
11208 p => p - 1,
11209 };
11210 let within = pos - starts[page_seq];
11211 survivors_vec.push((rid, page_seq, within));
11212 }
11213 total += survivors_vec.len();
11214 if let Some(&(rid, _, _)) = survivors_vec.first() {
11215 heap.push(std::cmp::Reverse((rid, run_idx)));
11216 }
11217 streams.push(RunStream::new(reader, survivors_vec, page_rows));
11218 }
11219
11220 let overlay_rows = if overlay_rids.is_empty() {
11222 Vec::new()
11223 } else {
11224 let bound = Self::overlay_materialization_bound(conditions, &survivors);
11225 self.overlay_visible_rows(snapshot, bound)
11226 };
11227 let overlay = if overlay_rows.is_empty() {
11228 None
11229 } else {
11230 let filtered =
11231 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
11232 if filtered.is_empty() {
11233 None
11234 } else {
11235 Some(self.materialize_overlay(&filtered, &projection))
11236 }
11237 };
11238
11239 let overlay_row_count = overlay
11240 .as_ref()
11241 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
11242 .unwrap_or(0);
11243 crate::trace::QueryTrace::record(|t| {
11244 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
11245 t.run_count = self.run_refs.len();
11246 t.memtable_rows = self.memtable.len();
11247 t.mutable_run_rows = self.mutable_run.len();
11248 t.overlay_rows = overlay_row_count;
11249 t.conditions_pushed = conditions.len();
11250 t.survivor_count = Some(total);
11251 });
11252
11253 Ok(Some(MultiRunCursor::new(
11254 streams, projection, heap, total, overlay,
11255 )))
11256 }
11257
11258 fn overlay_materialization_bound<'a>(
11270 conditions: &[crate::query::Condition],
11271 survivors: &'a Option<RowIdSet>,
11272 ) -> Option<&'a RowIdSet> {
11273 use crate::query::Condition;
11274 let has_range = conditions
11275 .iter()
11276 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
11277 if has_range {
11278 None
11279 } else {
11280 survivors.as_ref()
11281 }
11282 }
11283
11284 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
11296 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
11297 let mut fold = |row: Row| {
11298 if let Some(b) = bound {
11299 if !b.contains(row.row_id.0) {
11300 return;
11301 }
11302 }
11303 best.entry(row.row_id.0)
11304 .and_modify(|(be, br)| {
11305 if row.committed_epoch > *be {
11306 *be = row.committed_epoch;
11307 *br = row.clone();
11308 }
11309 })
11310 .or_insert_with(|| (row.committed_epoch, row));
11311 };
11312 for row in self.memtable.visible_versions(snapshot.epoch) {
11313 fold(row);
11314 }
11315 for row in self.mutable_run.visible_versions(snapshot.epoch) {
11316 fold(row);
11317 }
11318 let mut out: Vec<Row> = best
11319 .into_values()
11320 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
11321 .collect();
11322 out.sort_by_key(|r| r.row_id);
11323 out
11324 }
11325
11326 fn filter_overlay_rows(
11334 &self,
11335 rows: Vec<Row>,
11336 conditions: &[crate::query::Condition],
11337 survivors: Option<&RowIdSet>,
11338 snapshot: Snapshot,
11339 ) -> Result<Vec<Row>> {
11340 if conditions.is_empty() {
11341 return Ok(rows);
11342 }
11343 use crate::query::Condition;
11344 let all_index_served = !conditions
11348 .iter()
11349 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
11350 if all_index_served {
11351 return Ok(rows
11352 .into_iter()
11353 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
11354 .collect());
11355 }
11356 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
11359 for c in conditions {
11360 let s = match c {
11361 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
11362 _ => self.resolve_condition(c, snapshot)?,
11363 };
11364 per_cond_sets.push(s);
11365 }
11366 Ok(rows
11367 .into_iter()
11368 .filter(|row| {
11369 conditions.iter().enumerate().all(|(i, c)| match c {
11370 Condition::Range { column_id, lo, hi } => {
11371 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
11372 }
11373 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
11374 match row.columns.get(column_id) {
11375 Some(Value::Float64(v)) => {
11376 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
11377 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
11378 lo_ok && hi_ok
11379 }
11380 _ => false,
11381 }
11382 }
11383 _ => per_cond_sets[i].contains(row.row_id.0),
11384 })
11385 })
11386 .collect())
11387 }
11388
11389 fn materialize_overlay(
11392 &self,
11393 rows: &[Row],
11394 projection: &[(u16, TypeId)],
11395 ) -> Vec<columnar::NativeColumn> {
11396 if projection.is_empty() {
11397 return vec![columnar::null_native(TypeId::Int64, rows.len())];
11398 }
11399 let mut cols = Vec::with_capacity(projection.len());
11400 for (cid, ty) in projection {
11401 let vals: Vec<Value> = rows
11402 .iter()
11403 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
11404 .collect();
11405 cols.push(columnar::values_to_native(ty.clone(), &vals));
11406 }
11407 cols
11408 }
11409
11410 fn resolve_survivor_rids(
11415 &self,
11416 conditions: &[crate::query::Condition],
11417 reader: &mut RunReader,
11418 snapshot: Snapshot,
11419 ) -> Result<RowIdSet> {
11420 use crate::query::Condition;
11421 let mut sets: Vec<RowIdSet> = Vec::new();
11422 for c in conditions {
11423 self.validate_condition(c)?;
11424 let s: RowIdSet = match c {
11425 Condition::Pk(key) => {
11426 let lookup = self
11427 .schema
11428 .primary_key()
11429 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
11430 .unwrap_or_else(|| key.clone());
11431 self.hot
11432 .get(&lookup)
11433 .map(|r| RowIdSet::one(r.0))
11434 .unwrap_or_else(RowIdSet::empty)
11435 }
11436 Condition::BitmapEq { column_id, value } => {
11437 let lookup = self.index_lookup_key_bytes(*column_id, value);
11438 self.bitmap
11439 .get(column_id)
11440 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
11441 .unwrap_or_else(RowIdSet::empty)
11442 }
11443 Condition::BitmapIn { column_id, values } => {
11444 let bm = self.bitmap.get(column_id);
11445 let mut acc = roaring::RoaringBitmap::new();
11446 if let Some(b) = bm {
11447 for v in values {
11448 let lookup = self.index_lookup_key_bytes(*column_id, v);
11449 acc |= b.get(&lookup);
11450 }
11451 }
11452 RowIdSet::from_roaring(acc)
11453 }
11454 Condition::BytesPrefix { column_id, prefix } => {
11455 if let Some(b) = self.bitmap.get(column_id) {
11456 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
11457 let mut acc = roaring::RoaringBitmap::new();
11458 for key in b.keys() {
11459 if key.starts_with(&lookup_prefix) {
11460 acc |= b.get(&key);
11461 }
11462 }
11463 RowIdSet::from_roaring(acc)
11464 } else {
11465 RowIdSet::empty()
11466 }
11467 }
11468 Condition::FmContains { column_id, pattern } => self
11469 .fm
11470 .get(column_id)
11471 .map(|f| {
11472 RowIdSet::from_unsorted(
11473 f.locate(pattern).into_iter().map(|r| r.0).collect(),
11474 )
11475 })
11476 .unwrap_or_else(RowIdSet::empty),
11477 Condition::FmContainsAll {
11478 column_id,
11479 patterns,
11480 } => {
11481 if let Some(f) = self.fm.get(column_id) {
11482 let sets: Vec<RowIdSet> = patterns
11483 .iter()
11484 .map(|pat| {
11485 RowIdSet::from_unsorted(
11486 f.locate(pat).into_iter().map(|r| r.0).collect(),
11487 )
11488 })
11489 .collect();
11490 RowIdSet::intersect_many(sets)
11491 } else {
11492 RowIdSet::empty()
11493 }
11494 }
11495 Condition::Ann {
11496 column_id,
11497 query,
11498 k,
11499 } => RowIdSet::from_unsorted(
11500 self.retrieve_filtered(
11501 &crate::query::Retriever::Ann {
11502 column_id: *column_id,
11503 query: query.clone(),
11504 k: *k,
11505 },
11506 snapshot,
11507 None,
11508 None,
11509 None,
11510 None,
11511 )?
11512 .into_iter()
11513 .map(|hit| hit.row_id.0)
11514 .collect(),
11515 ),
11516 Condition::SparseMatch {
11517 column_id,
11518 query,
11519 k,
11520 } => RowIdSet::from_unsorted(
11521 self.retrieve_filtered(
11522 &crate::query::Retriever::Sparse {
11523 column_id: *column_id,
11524 query: query.clone(),
11525 k: *k,
11526 },
11527 snapshot,
11528 None,
11529 None,
11530 None,
11531 None,
11532 )?
11533 .into_iter()
11534 .map(|hit| hit.row_id.0)
11535 .collect(),
11536 ),
11537 Condition::MinHashSimilar {
11538 column_id,
11539 query,
11540 k,
11541 } => match self.minhash.get(column_id) {
11542 Some(index) => {
11543 let candidates = index.candidate_row_ids(query);
11544 let eligible =
11545 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
11546 RowIdSet::from_unsorted(
11547 index
11548 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
11549 .into_iter()
11550 .map(|(row_id, _)| row_id.0)
11551 .collect(),
11552 )
11553 }
11554 None => RowIdSet::empty(),
11555 },
11556 Condition::Range { column_id, lo, hi } => {
11557 if let Some(li) = self.learned_range.get(column_id) {
11558 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
11559 } else {
11560 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
11561 }
11562 }
11563 Condition::RangeF64 {
11564 column_id,
11565 lo,
11566 lo_inclusive,
11567 hi,
11568 hi_inclusive,
11569 } => {
11570 if let Some(li) = self.learned_range.get(column_id) {
11571 RowIdSet::from_unsorted(
11572 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
11573 .into_iter()
11574 .collect(),
11575 )
11576 } else {
11577 reader.range_row_id_set_f64(
11578 *column_id,
11579 *lo,
11580 *lo_inclusive,
11581 *hi,
11582 *hi_inclusive,
11583 )?
11584 }
11585 }
11586 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
11587 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
11588 };
11589 sets.push(s);
11590 }
11591 Ok(RowIdSet::intersect_many(sets))
11592 }
11593
11594 pub fn scan_cursor(
11615 &self,
11616 snapshot: Snapshot,
11617 projection: Vec<(u16, TypeId)>,
11618 conditions: &[crate::query::Condition],
11619 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
11620 if self.ttl.is_some() {
11621 return Ok(None);
11622 }
11623 if !conditions.is_empty() && !self.indexes_complete {
11629 return Ok(None);
11630 }
11631 if self.run_refs.len() == 1 {
11632 Ok(self
11633 .native_page_cursor(snapshot, projection, conditions)?
11634 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
11635 } else {
11636 Ok(self
11637 .native_multi_run_cursor(snapshot, projection, conditions)?
11638 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
11639 }
11640 }
11641
11642 pub fn aggregate_native(
11656 &self,
11657 snapshot: Snapshot,
11658 column: Option<u16>,
11659 conditions: &[crate::query::Condition],
11660 agg: NativeAgg,
11661 ) -> Result<Option<NativeAggResult>> {
11662 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
11663 }
11664
11665 pub fn aggregate_native_with_control(
11666 &self,
11667 snapshot: Snapshot,
11668 column: Option<u16>,
11669 conditions: &[crate::query::Condition],
11670 agg: NativeAgg,
11671 control: &crate::ExecutionControl,
11672 ) -> Result<Option<NativeAggResult>> {
11673 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
11674 }
11675
11676 fn aggregate_native_inner(
11677 &self,
11678 snapshot: Snapshot,
11679 column: Option<u16>,
11680 conditions: &[crate::query::Condition],
11681 agg: NativeAgg,
11682 control: Option<&crate::ExecutionControl>,
11683 ) -> Result<Option<NativeAggResult>> {
11684 execution_checkpoint(control, 0)?;
11685 if self.ttl.is_some() {
11686 return Ok(None);
11687 }
11688 if self.run_refs.len() == 1 && conditions.is_empty() {
11690 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
11691 return Ok(Some(res));
11692 }
11693 }
11694 if matches!(agg, NativeAgg::Count) && column.is_none() {
11698 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
11699 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
11700 }
11701 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
11702 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
11703 }
11704 let cid = match column {
11707 Some(c) => c,
11708 None => return Ok(None),
11709 };
11710 let ty = self.column_type(cid);
11711 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
11712 execution_checkpoint(control, 0)?;
11713 return match ty {
11714 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11715 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
11716 Ok(Some(pack_int(agg, count, sum, mn, mx)))
11717 }
11718 TypeId::Float64 => {
11719 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
11720 Ok(Some(pack_float(agg, count, sum, mn, mx)))
11721 }
11722 _ => Ok(None),
11723 };
11724 }
11725 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
11727 execution_checkpoint(control, 0)?;
11728 match ty {
11729 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
11730 let mut count = 0u64;
11731 let mut sum = 0i128;
11732 let mut mn = i64::MAX;
11733 let mut mx = i64::MIN;
11734 for row in &rows {
11735 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
11736 count += 1;
11737 sum += i128::from(*v);
11738 mn = mn.min(*v);
11739 mx = mx.max(*v);
11740 }
11741 }
11742 Ok(Some(pack_int(agg, count, sum, mn, mx)))
11743 }
11744 TypeId::Float64 => {
11745 let mut count = 0u64;
11746 let mut sum = 0.0f64;
11747 let mut mn = f64::INFINITY;
11748 let mut mx = f64::NEG_INFINITY;
11749 for row in &rows {
11750 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
11751 count += 1;
11752 sum += *v;
11753 mn = mn.min(*v);
11754 mx = mx.max(*v);
11755 }
11756 }
11757 Ok(Some(pack_float(agg, count, sum, mn, mx)))
11758 }
11759 _ => Ok(None),
11760 }
11761 }
11762
11763 fn visible_rows_filtered(
11765 &self,
11766 snapshot: Snapshot,
11767 conditions: &[crate::query::Condition],
11768 control: Option<&crate::ExecutionControl>,
11769 ) -> Result<Vec<Row>> {
11770 let rows = if let Some(control) = control {
11771 self.visible_rows_controlled(snapshot, control)?
11772 } else {
11773 self.visible_rows(snapshot)?
11774 };
11775 if conditions.is_empty() {
11776 return Ok(rows);
11777 }
11778 Ok(rows
11779 .into_iter()
11780 .filter(|row| {
11781 conditions
11782 .iter()
11783 .all(|cond| condition_matches_row(cond, row, &self.schema))
11784 })
11785 .collect())
11786 }
11787
11788 fn aggregate_from_stats(
11796 &self,
11797 snapshot: Snapshot,
11798 column: Option<u16>,
11799 agg: NativeAgg,
11800 ) -> Result<Option<NativeAggResult>> {
11801 let cid = match (agg, column) {
11802 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
11803 _ => return Ok(None), };
11805 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
11806 return Ok(None);
11807 };
11808 let Some(cs) = stats.get(&cid) else {
11809 return Ok(None);
11810 };
11811 match agg {
11812 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
11814 self.live_count.saturating_sub(cs.null_count),
11815 ))),
11816 NativeAgg::Min | NativeAgg::Max => {
11817 let bound = if agg == NativeAgg::Min {
11818 &cs.min
11819 } else {
11820 &cs.max
11821 };
11822 match bound {
11823 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
11824 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
11825 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
11830 None => Ok(None),
11831 }
11832 }
11833 _ => Ok(None),
11834 }
11835 }
11836
11837 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
11846 if self.ttl.is_some() {
11847 return Ok(None);
11848 }
11849 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
11850 return Ok(None);
11851 }
11852 self.ensure_indexes_complete()?;
11855 let reader = self.open_reader(self.run_refs[0].run_id)?;
11856 if self.live_count != reader.row_count() as u64 {
11857 return Ok(None);
11858 }
11859 let Some(bm) = self.bitmap.get(&column_id) else {
11860 return Ok(None); };
11862 let mut distinct = bm.value_count() as u64;
11863 if !bm.get(&Value::Null.encode_key()).is_empty() {
11866 distinct = distinct.saturating_sub(1);
11867 }
11868 Ok(Some(distinct))
11869 }
11870
11871 pub fn aggregate_incremental(
11883 &mut self,
11884 cache_key: u64,
11885 conditions: &[crate::query::Condition],
11886 column: Option<u16>,
11887 agg: NativeAgg,
11888 ) -> Result<IncrementalAggResult> {
11889 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
11890 }
11891
11892 pub fn aggregate_incremental_with_control(
11893 &mut self,
11894 cache_key: u64,
11895 conditions: &[crate::query::Condition],
11896 column: Option<u16>,
11897 agg: NativeAgg,
11898 control: &crate::ExecutionControl,
11899 ) -> Result<IncrementalAggResult> {
11900 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
11901 }
11902
11903 fn aggregate_incremental_inner(
11904 &mut self,
11905 cache_key: u64,
11906 conditions: &[crate::query::Condition],
11907 column: Option<u16>,
11908 agg: NativeAgg,
11909 control: Option<&crate::ExecutionControl>,
11910 ) -> Result<IncrementalAggResult> {
11911 execution_checkpoint(control, 0)?;
11912 let snap = self.snapshot();
11913 let cur_wm = self.allocator.current().0;
11914 let cur_epoch = snap.epoch.0;
11915 let incremental_ok = self.ttl.is_none()
11922 && !self.had_deletes
11923 && self.memtable.is_empty()
11924 && self.mutable_run.is_empty();
11925
11926 if incremental_ok {
11929 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
11930 if cached.epoch == cur_epoch {
11931 return Ok(IncrementalAggResult {
11932 state: cached.state,
11933 incremental: true,
11934 delta_rows: 0,
11935 });
11936 }
11937 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
11938 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
11939 let mut delta_rids = Vec::with_capacity(delta_len);
11940 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
11941 execution_checkpoint(control, index)?;
11942 delta_rids.push(row_id);
11943 }
11944 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
11945 execution_checkpoint(control, 0)?;
11946 let index_sets = self.resolve_index_conditions(conditions, snap)?;
11947 let delta_state = agg_state_from_rows(
11948 &delta_rows,
11949 conditions,
11950 &index_sets,
11951 column,
11952 agg,
11953 &self.schema,
11954 control,
11955 )?;
11956 let merged = cached.state.merge(delta_state);
11957 let delta_n = delta_rids.len() as u64;
11958 Arc::make_mut(&mut self.agg_cache).insert(
11959 cache_key,
11960 CachedAgg {
11961 state: merged.clone(),
11962 watermark: cur_wm,
11963 epoch: cur_epoch,
11964 },
11965 );
11966 return Ok(IncrementalAggResult {
11967 state: merged,
11968 incremental: true,
11969 delta_rows: delta_n,
11970 });
11971 }
11972 }
11973 }
11974
11975 let cursor_ok =
11980 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
11981 let state = if cursor_ok && agg != NativeAgg::Avg {
11982 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
11983 Some(result) => {
11984 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
11985 }
11986 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
11987 }
11988 } else {
11989 self.agg_state_full_scan(conditions, column, agg, snap, control)?
11990 };
11991 if incremental_ok {
11993 Arc::make_mut(&mut self.agg_cache).insert(
11994 cache_key,
11995 CachedAgg {
11996 state: state.clone(),
11997 watermark: cur_wm,
11998 epoch: cur_epoch,
11999 },
12000 );
12001 }
12002 Ok(IncrementalAggResult {
12003 state,
12004 incremental: false,
12005 delta_rows: 0,
12006 })
12007 }
12008
12009 fn agg_state_full_scan(
12012 &self,
12013 conditions: &[crate::query::Condition],
12014 column: Option<u16>,
12015 agg: NativeAgg,
12016 snap: Snapshot,
12017 control: Option<&crate::ExecutionControl>,
12018 ) -> Result<AggState> {
12019 execution_checkpoint(control, 0)?;
12020 let rows = self.visible_rows(snap)?;
12021 execution_checkpoint(control, 0)?;
12022 let index_sets = self.resolve_index_conditions(conditions, snap)?;
12023 agg_state_from_rows(
12024 &rows,
12025 conditions,
12026 &index_sets,
12027 column,
12028 agg,
12029 &self.schema,
12030 control,
12031 )
12032 }
12033
12034 fn resolve_index_conditions(
12037 &self,
12038 conditions: &[crate::query::Condition],
12039 snapshot: Snapshot,
12040 ) -> Result<Vec<RowIdSet>> {
12041 use crate::query::Condition;
12042 let mut sets = Vec::new();
12043 for c in conditions {
12044 if matches!(
12045 c,
12046 Condition::Ann { .. }
12047 | Condition::SparseMatch { .. }
12048 | Condition::MinHashSimilar { .. }
12049 ) {
12050 sets.push(self.resolve_condition(c, snapshot)?);
12051 }
12052 }
12053 Ok(sets)
12054 }
12055
12056 fn column_type(&self, cid: u16) -> TypeId {
12057 self.schema
12058 .columns
12059 .iter()
12060 .find(|c| c.id == cid)
12061 .map(|c| c.ty.clone())
12062 .unwrap_or(TypeId::Bytes)
12063 }
12064
12065 pub fn approx_aggregate(
12074 &mut self,
12075 conditions: &[crate::query::Condition],
12076 column: Option<u16>,
12077 agg: ApproxAgg,
12078 z: f64,
12079 ) -> Result<Option<ApproxResult>> {
12080 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
12081 }
12082
12083 pub fn approx_aggregate_with_candidate_authorization(
12086 &mut self,
12087 conditions: &[crate::query::Condition],
12088 column: Option<u16>,
12089 agg: ApproxAgg,
12090 z: f64,
12091 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
12092 ) -> Result<Option<ApproxResult>> {
12093 use crate::query::Condition;
12094 self.ensure_reservoir_complete()?;
12095 let mut snapshot = self.snapshot();
12100 let n_pop = self.count();
12101 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
12102 if sample_rids.is_empty() {
12103 return Ok(None);
12104 }
12105 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
12107 if live_sample.is_empty() {
12108 snapshot = Snapshot::unbounded();
12109 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
12110 }
12111 let s = live_sample.len();
12112 if s == 0 {
12113 return Ok(None);
12114 }
12115 let authorized = authorization
12116 .map(|authorization| {
12117 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
12118 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
12119 })
12120 .transpose()?;
12121
12122 let mut index_sets: Vec<RowIdSet> = Vec::new();
12125 for c in conditions {
12126 if matches!(
12127 c,
12128 Condition::Ann { .. }
12129 | Condition::SparseMatch { .. }
12130 | Condition::MinHashSimilar { .. }
12131 ) {
12132 index_sets.push(self.resolve_condition(c, snapshot)?);
12133 }
12134 }
12135
12136 let cid = match (agg, column) {
12138 (ApproxAgg::Count, _) => None,
12139 (_, Some(c)) => Some(c),
12140 _ => return Ok(None),
12141 };
12142 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
12143 for r in &live_sample {
12144 if authorized
12145 .as_ref()
12146 .is_some_and(|authorized| !authorized.contains(&r.row_id))
12147 {
12148 continue;
12149 }
12150 if !conditions
12152 .iter()
12153 .all(|c| condition_matches_row(c, r, &self.schema))
12154 {
12155 continue;
12156 }
12157 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
12159 continue;
12160 }
12161 if let Some(cid) = cid {
12162 let mut cells = r
12163 .columns
12164 .get(&cid)
12165 .cloned()
12166 .map(|value| vec![(cid, value)])
12167 .unwrap_or_default();
12168 if let Some(authorization) = authorization {
12169 authorization.security.apply_masks_to_cells(
12170 authorization.table,
12171 &mut cells,
12172 authorization.principal,
12173 );
12174 }
12175 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
12176 passing_vals.push(v);
12177 } } else {
12179 passing_vals.push(0.0); }
12181 }
12182 let m = passing_vals.len();
12183
12184 let (point, half) = match agg {
12185 ApproxAgg::Count => {
12186 let p = m as f64 / s as f64;
12188 let point = n_pop as f64 * p;
12189 let var = if s > 1 {
12190 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
12191 * (1.0 - s as f64 / n_pop as f64).max(0.0)
12192 } else {
12193 0.0
12194 };
12195 (point, z * var.sqrt())
12196 }
12197 ApproxAgg::Sum => {
12198 let y: Vec<f64> = live_sample
12200 .iter()
12201 .map(|r| {
12202 let passes_row = authorized
12203 .as_ref()
12204 .is_none_or(|authorized| authorized.contains(&r.row_id))
12205 && conditions
12206 .iter()
12207 .all(|c| condition_matches_row(c, r, &self.schema))
12208 && index_sets.iter().all(|set| set.contains(r.row_id.0));
12209 if passes_row {
12210 cid.and_then(|cid| {
12211 let mut cells = r
12212 .columns
12213 .get(&cid)
12214 .cloned()
12215 .map(|value| vec![(cid, value)])
12216 .unwrap_or_default();
12217 if let Some(authorization) = authorization {
12218 authorization.security.apply_masks_to_cells(
12219 authorization.table,
12220 &mut cells,
12221 authorization.principal,
12222 );
12223 }
12224 as_f64(cells.first().map(|(_, value)| value))
12225 })
12226 .unwrap_or(0.0)
12227 } else {
12228 0.0
12229 }
12230 })
12231 .collect();
12232 let mean_y = y.iter().sum::<f64>() / s as f64;
12233 let point = n_pop as f64 * mean_y;
12234 let var = if s > 1 {
12235 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
12236 let var_y = ss / (s - 1) as f64;
12237 n_pop as f64 * n_pop as f64 * var_y / s as f64
12238 * (1.0 - s as f64 / n_pop as f64).max(0.0)
12239 } else {
12240 0.0
12241 };
12242 (point, z * var.sqrt())
12243 }
12244 ApproxAgg::Avg => {
12245 if m == 0 {
12246 return Ok(Some(ApproxResult {
12247 point: 0.0,
12248 ci_low: 0.0,
12249 ci_high: 0.0,
12250 n_population: n_pop,
12251 n_sample_live: s,
12252 n_passing: 0,
12253 }));
12254 }
12255 let mean = passing_vals.iter().sum::<f64>() / m as f64;
12256 let half = if m > 1 {
12257 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
12258 let sd = (ss / (m - 1) as f64).sqrt();
12259 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
12260 z * sd / (m as f64).sqrt() * fpc.sqrt()
12261 } else {
12262 0.0
12263 };
12264 (mean, half)
12265 }
12266 };
12267
12268 Ok(Some(ApproxResult {
12269 point,
12270 ci_low: point - half,
12271 ci_high: point + half,
12272 n_population: n_pop,
12273 n_sample_live: s,
12274 n_passing: m,
12275 }))
12276 }
12277
12278 pub fn exact_column_stats(
12286 &self,
12287 _snapshot: Snapshot,
12288 projection: &[u16],
12289 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
12290 if self.ttl.is_some()
12291 || !(self.memtable.is_empty()
12292 && self.mutable_run.is_empty()
12293 && self.run_refs.len() == 1)
12294 {
12295 return Ok(None);
12296 }
12297 let reader = self.open_reader(self.run_refs[0].run_id)?;
12298 if self.live_count != reader.row_count() as u64 {
12299 return Ok(None);
12300 }
12301 let mut out = HashMap::new();
12302 for &cid in projection {
12303 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
12304 Some(c) => c,
12305 None => continue,
12306 };
12307 let Some(stats) = reader.column_page_stats(cid) else {
12309 out.insert(
12310 cid,
12311 ColumnStat {
12312 min: None,
12313 max: None,
12314 null_count: self.live_count,
12315 },
12316 );
12317 continue;
12318 };
12319 let stat = match cdef.ty {
12320 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
12321 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
12322 min: mn.map(Value::Int64),
12323 max: mx.map(Value::Int64),
12324 null_count: n,
12325 })
12326 }
12327 TypeId::Float64 => {
12328 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
12329 min: mn.map(Value::Float64),
12330 max: mx.map(Value::Float64),
12331 null_count: n,
12332 })
12333 }
12334 _ => None,
12335 };
12336 if let Some(s) = stat {
12337 out.insert(cid, s);
12338 }
12339 }
12340 Ok(Some(out))
12341 }
12342
12343 pub fn dir(&self) -> &Path {
12344 &self.dir
12345 }
12346
12347 pub fn schema(&self) -> &Schema {
12348 &self.schema
12349 }
12350
12351 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
12352 self.ann.get(&column_id)
12353 }
12354
12355 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
12356 self.ann.get_mut(&column_id)
12357 }
12358
12359 pub(crate) fn set_catalog_name(&mut self, name: String) {
12360 self.name = name;
12361 }
12362
12363 pub(crate) fn prepare_alter_column(
12364 &mut self,
12365 column_name: &str,
12366 change: &AlterColumn,
12367 ) -> Result<(ColumnDef, Option<Schema>)> {
12368 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
12369 return Err(MongrelError::InvalidArgument(
12370 "ALTER COLUMN requires committing staged writes first".into(),
12371 ));
12372 }
12373 let old = self
12374 .schema
12375 .columns
12376 .iter()
12377 .find(|c| c.name == column_name)
12378 .cloned()
12379 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
12380 let mut next = old.clone();
12381
12382 if let Some(name) = &change.name {
12383 let trimmed = name.trim();
12384 if trimmed.is_empty() {
12385 return Err(MongrelError::InvalidArgument(
12386 "ALTER COLUMN name must not be empty".into(),
12387 ));
12388 }
12389 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
12390 return Err(MongrelError::Schema(format!(
12391 "column {trimmed} already exists"
12392 )));
12393 }
12394 next.name = trimmed.to_string();
12395 }
12396
12397 if let Some(ty) = &change.ty {
12398 next.ty = ty.clone();
12399 }
12400 if let Some(flags) = change.flags {
12401 validate_alter_column_flags(old.flags, flags)?;
12402 next.flags = flags;
12403 }
12404
12405 if let Some(default_change) = &change.default_value {
12406 next.default_value = default_change.clone();
12407 }
12408 if let Some(source_change) = &change.embedding_source {
12409 next.embedding_source = source_change.clone();
12410 }
12411
12412 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
12413 if old.flags.contains(ColumnFlags::NULLABLE)
12414 && !next.flags.contains(ColumnFlags::NULLABLE)
12415 && self.column_has_nulls(old.id)?
12416 {
12417 return Err(MongrelError::InvalidArgument(format!(
12418 "column '{}' contains NULL values",
12419 old.name
12420 )));
12421 }
12422 if next == old {
12423 return Ok((next, None));
12424 }
12425 let mut schema = self.schema.clone();
12426 let index = schema
12427 .columns
12428 .iter()
12429 .position(|column| column.id == next.id)
12430 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
12431 schema.columns[index] = next.clone();
12432 schema.schema_id = schema
12433 .schema_id
12434 .checked_add(1)
12435 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
12436 schema.validate_auto_increment()?;
12437 schema.validate_defaults()?;
12438 Ok((next, Some(schema)))
12439 }
12440
12441 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
12442 self.schema = schema;
12443 self.auto_inc = resolve_auto_inc(&self.schema);
12444 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
12445 self.clear_result_cache();
12446 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
12447 }
12448
12449 pub(crate) fn publish_index_schema_change(
12453 &mut self,
12454 schema: Schema,
12455 artifact: SecondaryIndexArtifact,
12456 ) -> Result<()> {
12457 self.apply_altered_schema_prepared(schema);
12458 self.prune_index_maps_to_schema();
12459 match artifact {
12460 SecondaryIndexArtifact::Bitmap(column_id, index) => {
12461 self.bitmap.insert(column_id, index);
12462 }
12463 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
12464 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
12465 }
12466 SecondaryIndexArtifact::Fm(column_id, index) => {
12467 self.fm.insert(column_id, *index);
12468 }
12469 SecondaryIndexArtifact::Ann(column_id, index) => {
12470 self.ann.insert(column_id, *index);
12471 }
12472 SecondaryIndexArtifact::Sparse(column_id, index) => {
12473 self.sparse.insert(column_id, index);
12474 }
12475 SecondaryIndexArtifact::MinHash(column_id, index) => {
12476 self.minhash.insert(column_id, index);
12477 }
12478 }
12479 self.indexes_complete = true;
12480 self.seal_generations();
12481 let view = Arc::new(self.capture_read_generation());
12482 self.published.store(Arc::clone(&view));
12483 checkpoint_current_schema(self)?;
12484 self.invalidate_index_checkpoint();
12487 Ok(())
12488 }
12489
12490 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
12496 self.apply_altered_schema_prepared(schema);
12497 self.prune_index_maps_to_schema();
12498 self.indexes_complete = true;
12499 self.seal_generations();
12500 let view = Arc::new(self.capture_read_generation());
12501 self.published.store(Arc::clone(&view));
12502 checkpoint_current_schema(self)?;
12503 self.invalidate_index_checkpoint();
12505 Ok(())
12506 }
12507
12508 fn prune_index_maps_to_schema(&mut self) {
12509 let keep_bitmap: std::collections::HashSet<u16> = self
12510 .schema
12511 .indexes
12512 .iter()
12513 .filter(|index| index.kind == IndexKind::Bitmap)
12514 .map(|index| index.column_id)
12515 .collect();
12516 let keep_ann: std::collections::HashSet<u16> = self
12517 .schema
12518 .indexes
12519 .iter()
12520 .filter(|index| index.kind == IndexKind::Ann)
12521 .map(|index| index.column_id)
12522 .collect();
12523 let keep_fm: std::collections::HashSet<u16> = self
12524 .schema
12525 .indexes
12526 .iter()
12527 .filter(|index| index.kind == IndexKind::FmIndex)
12528 .map(|index| index.column_id)
12529 .collect();
12530 let keep_sparse: std::collections::HashSet<u16> = self
12531 .schema
12532 .indexes
12533 .iter()
12534 .filter(|index| index.kind == IndexKind::Sparse)
12535 .map(|index| index.column_id)
12536 .collect();
12537 let keep_minhash: std::collections::HashSet<u16> = self
12538 .schema
12539 .indexes
12540 .iter()
12541 .filter(|index| index.kind == IndexKind::MinHash)
12542 .map(|index| index.column_id)
12543 .collect();
12544 let keep_learned: std::collections::HashSet<u16> = self
12545 .schema
12546 .indexes
12547 .iter()
12548 .filter(|index| index.kind == IndexKind::LearnedRange)
12549 .map(|index| index.column_id)
12550 .collect();
12551 self.bitmap
12552 .retain(|column_id, _| keep_bitmap.contains(column_id));
12553 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
12554 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
12555 self.sparse
12556 .retain(|column_id, _| keep_sparse.contains(column_id));
12557 self.minhash
12558 .retain(|column_id, _| keep_minhash.contains(column_id));
12559 {
12560 let learned = Arc::make_mut(&mut self.learned_range);
12561 learned.retain(|column_id, _| keep_learned.contains(column_id));
12562 }
12563 }
12564
12565 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
12566 checkpoint_current_schema(self)
12567 }
12568
12569 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
12570 self.ensure_writable()?;
12571 let previous_schema = self.schema.clone();
12572 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
12573 if let Some(schema) = schema {
12574 self.apply_altered_schema_prepared(schema);
12575 self.checkpoint_standalone_schema_change(previous_schema)?;
12576 }
12577 Ok(column)
12578 }
12579
12580 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
12581 if self.live_count == 0 {
12582 return Ok(false);
12583 }
12584 let snap = self.snapshot();
12585 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
12586 Ok(columns
12587 .first()
12588 .map(|(_, col)| col.null_count(col.len()) != 0)
12589 .unwrap_or(true))
12590 }
12591
12592 fn has_stored_versions(&self) -> bool {
12593 !self.memtable.is_empty()
12594 || !self.mutable_run.is_empty()
12595 || self.run_refs.iter().any(|r| r.row_count > 0)
12596 || !self.retiring.is_empty()
12597 }
12598
12599 pub fn add_column(
12604 &mut self,
12605 name: &str,
12606 ty: TypeId,
12607 flags: ColumnFlags,
12608 default_value: Option<crate::schema::DefaultExpr>,
12609 ) -> Result<u16> {
12610 self.add_column_with_id(name, ty, flags, default_value, None)
12611 }
12612
12613 pub fn add_column_with_id(
12614 &mut self,
12615 name: &str,
12616 ty: TypeId,
12617 flags: ColumnFlags,
12618 default_value: Option<crate::schema::DefaultExpr>,
12619 requested_id: Option<u16>,
12620 ) -> Result<u16> {
12621 self.ensure_writable()?;
12622 let previous_schema = self.schema.clone();
12623 let (column, schema) =
12624 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
12625 self.apply_altered_schema_prepared(schema);
12626 self.checkpoint_standalone_schema_change(previous_schema)?;
12627 Ok(column.id)
12628 }
12629
12630 pub(crate) fn prepare_add_column(
12631 &mut self,
12632 name: &str,
12633 ty: TypeId,
12634 flags: ColumnFlags,
12635 default_value: Option<crate::schema::DefaultExpr>,
12636 requested_id: Option<u16>,
12637 ) -> Result<(ColumnDef, Schema)> {
12638 self.ensure_writable()?;
12639 if self.schema.columns.iter().any(|c| c.name == name) {
12640 return Err(MongrelError::Schema(format!(
12641 "column {name} already exists"
12642 )));
12643 }
12644 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
12645 if self.schema.columns.iter().any(|c| c.id == id) {
12646 return Err(MongrelError::Schema(format!(
12647 "column id {id} already exists"
12648 )));
12649 }
12650 id
12651 } else {
12652 self.schema
12653 .columns
12654 .iter()
12655 .map(|c| c.id)
12656 .max()
12657 .unwrap_or(0)
12658 .checked_add(1)
12659 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
12660 };
12661 let column = ColumnDef {
12662 id,
12663 name: name.to_string(),
12664 ty,
12665 flags,
12666 default_value,
12667 embedding_source: None,
12668 };
12669 let mut next_schema = self.schema.clone();
12670 next_schema.columns.push(column.clone());
12671 next_schema.schema_id = next_schema
12672 .schema_id
12673 .checked_add(1)
12674 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
12675 next_schema.validate_auto_increment()?;
12676 next_schema.validate_defaults()?;
12677 Ok((column, next_schema))
12678 }
12679
12680 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
12689 self.ensure_writable()?;
12690 let cid = self
12691 .schema
12692 .columns
12693 .iter()
12694 .find(|c| c.name == column_name)
12695 .map(|c| c.id)
12696 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
12697 let ty = self
12698 .schema
12699 .columns
12700 .iter()
12701 .find(|c| c.id == cid)
12702 .map(|c| c.ty.clone())
12703 .unwrap_or(TypeId::Int64);
12704 if !matches!(
12705 ty,
12706 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
12707 ) {
12708 return Err(MongrelError::Schema(format!(
12709 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
12710 )));
12711 }
12712 if self
12713 .schema
12714 .indexes
12715 .iter()
12716 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
12717 {
12718 return Ok(()); }
12720 let previous_schema = self.schema.clone();
12721 let previous_learned_range = Arc::clone(&self.learned_range);
12722 let mut next_schema = previous_schema.clone();
12723 next_schema.indexes.push(IndexDef {
12724 name: format!("{}_learned_range", column_name),
12725 column_id: cid,
12726 kind: IndexKind::LearnedRange,
12727 predicate: None,
12728 options: Default::default(),
12729 });
12730 next_schema.schema_id = next_schema
12731 .schema_id
12732 .checked_add(1)
12733 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
12734 self.apply_altered_schema_prepared(next_schema);
12735 if let Err(error) = self.build_learned_ranges() {
12736 self.apply_altered_schema_prepared(previous_schema);
12737 self.learned_range = previous_learned_range;
12738 return Err(error);
12739 }
12740 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
12741 if !matches!(
12742 &error,
12743 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
12744 ) {
12745 self.learned_range = previous_learned_range;
12746 }
12747 return Err(error);
12748 }
12749 Ok(())
12750 }
12751
12752 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
12753 let mut schema_published = false;
12754 let schema_result = match self._root_guard.as_deref() {
12755 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
12756 schema_published = true;
12757 }),
12758 None => write_schema_with_after(&self.dir, &self.schema, || {
12759 schema_published = true;
12760 }),
12761 };
12762 if schema_result.is_err() && !schema_published {
12763 self.apply_altered_schema_prepared(previous_schema);
12764 return schema_result;
12765 }
12766
12767 let manifest_result = self.persist_manifest(self.current_epoch());
12768 match (schema_result, manifest_result) {
12769 (_, Ok(())) => Ok(()),
12770 (Ok(()), Err(error)) => {
12771 self.poison_after_maintenance_publish_failure();
12772 Err(MongrelError::DurableCommit {
12773 epoch: self.current_epoch().0,
12774 message: format!(
12775 "schema is durable but matching manifest publication failed: {error}"
12776 ),
12777 })
12778 }
12779 (Err(schema_error), Err(manifest_error)) => {
12780 self.poison_after_maintenance_publish_failure();
12781 Err(MongrelError::CommitOutcomeUnknown {
12782 epoch: self.current_epoch().0,
12783 message: format!(
12784 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
12785 ),
12786 })
12787 }
12788 }
12789 }
12790
12791 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
12794 self.sync_byte_threshold = threshold;
12795 if let WalSink::Private(w) = &mut self.wal {
12796 w.set_sync_byte_threshold(threshold);
12797 }
12798 }
12799
12800 pub fn page_cache_flush(&self) {
12804 self.page_cache.flush_to_disk();
12805 }
12806
12807 pub fn page_cache_len(&self) -> usize {
12809 self.page_cache.len()
12810 }
12811
12812 pub fn decoded_cache_len(&self) -> usize {
12815 self.decoded_cache.len()
12816 }
12817
12818 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
12821 self.memtable.drain_sorted()
12822 }
12823
12824 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
12825 self.runs_dir().join(format!("r-{run_id}.sr"))
12826 }
12827
12828 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
12829 match self.runs_root.as_deref() {
12830 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
12831 None => Ok(None),
12832 }
12833 }
12834
12835 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
12836 match self.runs_root.as_deref() {
12837 Some(root) => Ok(Some(root.create_regular_new(name)?)),
12838 None => Ok(None),
12839 }
12840 }
12841
12842 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
12843 match self.runs_root.as_deref() {
12844 Some(root) => match root.remove_file(name) {
12845 Ok(()) => Ok(()),
12846 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
12847 Err(error) => Err(error.into()),
12848 },
12849 None => match std::fs::remove_file(self.runs_dir().join(name)) {
12850 Ok(()) => Ok(()),
12851 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
12852 Err(error) => Err(error.into()),
12853 },
12854 }
12855 }
12856
12857 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
12858 match self.runs_root.as_deref() {
12859 Some(root) => root
12860 .rename_file_new(source, destination)
12861 .map_err(Into::into),
12862 None => crate::durable_file::rename(
12863 &self.runs_dir().join(source),
12864 &self.runs_dir().join(destination),
12865 )
12866 .map_err(Into::into),
12867 }
12868 }
12869
12870 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
12871 self.run_refs.iter().map(|run| run.run_id)
12872 }
12873
12874 pub(crate) fn table_dir(&self) -> &Path {
12875 &self.dir
12876 }
12877
12878 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
12879 &self.schema
12880 }
12881
12882 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
12883 let id = self.next_run_id;
12884 self.next_run_id = self
12885 .next_run_id
12886 .checked_add(1)
12887 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
12888 Ok(id)
12889 }
12890
12891 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
12892 self.run_refs.push(run_ref);
12893 }
12894
12895 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
12905 self.retiring.push(crate::manifest::RetiredRun {
12906 run_id,
12907 retire_epoch,
12908 });
12909 }
12910
12911 pub(crate) fn reap_retiring(
12915 &mut self,
12916 min_active: Epoch,
12917 backup_pinned: &std::collections::HashSet<u128>,
12918 ) -> Result<usize> {
12919 if self.retiring.is_empty() {
12920 return Ok(0);
12921 }
12922 let mut reaped = 0;
12923 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
12924 for r in std::mem::take(&mut self.retiring) {
12930 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
12931 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
12932 reaped += 1;
12933 } else {
12934 kept.push(r);
12935 }
12936 }
12937 self.retiring = kept;
12938 if reaped > 0 {
12939 self.persist_manifest(self.current_epoch())?;
12940 }
12941 Ok(reaped)
12942 }
12943
12944 pub(crate) fn has_reapable_retiring(
12945 &self,
12946 min_active: Epoch,
12947 backup_pinned: &std::collections::HashSet<u128>,
12948 ) -> bool {
12949 self.retiring
12950 .iter()
12951 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
12952 }
12953
12954 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
12955 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
12956 return false;
12957 }
12958 self.live_count = self.live_count.saturating_add(run_ref.row_count);
12959 self.run_refs.push(run_ref);
12960 self.indexes_complete = false;
12961 true
12962 }
12963
12964 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
12965 self.kek.as_ref()
12966 }
12967
12968 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
12969 let mut reader = match self.runs_root.as_deref() {
12970 Some(root) => RunReader::open_file_with_cache(
12971 root.open_regular(format!("r-{run_id}.sr"))?,
12972 self.schema.clone(),
12973 self.kek.clone(),
12974 Some(self.page_cache.clone()),
12975 Some(self.decoded_cache.clone()),
12976 self.table_id,
12977 Some(&self.verified_runs),
12978 None,
12979 )?,
12980 None => RunReader::open_with_cache(
12981 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
12982 self.schema.clone(),
12983 self.kek.clone(),
12984 Some(self.page_cache.clone()),
12985 Some(self.decoded_cache.clone()),
12986 self.table_id,
12987 Some(&self.verified_runs),
12988 )?,
12989 };
12990 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
12994 reader.set_uniform_epoch(Epoch(rr.epoch_created));
12995 }
12996 Ok(reader)
12997 }
12998
12999 pub(crate) fn run_refs(&self) -> &[RunRef] {
13000 &self.run_refs
13001 }
13002
13003 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
13004 self.retiring.iter().map(|run| run.run_id)
13005 }
13006
13007 pub(crate) fn runs_dir(&self) -> PathBuf {
13008 self.runs_root
13009 .as_deref()
13010 .and_then(|root| root.io_path().ok())
13011 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
13012 }
13013
13014 pub(crate) fn wal_dir(&self) -> PathBuf {
13015 self.dir.join(WAL_DIR)
13016 }
13017
13018 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
13019 self.run_refs = refs;
13020 self.refresh_run_row_id_ranges();
13021 }
13022
13023 pub(crate) fn run_row_id_range(&self, run_id: u128) -> Option<(u64, u64)> {
13027 self.run_row_id_ranges.get(&run_id).copied()
13028 }
13029
13030 pub(crate) fn refresh_run_row_id_ranges(&mut self) {
13033 let mut ranges = HashMap::with_capacity(self.run_refs.len());
13034 for rr in &self.run_refs {
13035 if let Ok(reader) = self.open_reader(rr.run_id) {
13036 let h = reader.header();
13037 ranges.insert(rr.run_id, (h.min_row_id, h.max_row_id));
13038 }
13039 }
13040 self.run_row_id_ranges = ranges;
13041 }
13042
13043 pub(crate) fn compaction_zstd_level(&self) -> i32 {
13044 self.compaction_zstd_level
13045 }
13046
13047 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
13048 self.kek.clone()
13049 }
13050
13051 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
13055 self.kek.as_ref().map(|k| k.derive_idx_key())
13056 }
13057
13058 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
13062 self.kek.as_ref().map(|k| *k.derive_meta_key())
13063 }
13064
13065 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
13068 self.column_keys
13069 .iter()
13070 .map(|(&id, &(_, scheme))| (id, scheme))
13071 .collect()
13072 }
13073
13074 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
13079 self.tokenize_value_enc(column_id, v)
13080 }
13081
13082 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
13083 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
13084 let (key, scheme) = self.column_keys.get(&column_id)?;
13085 let token: Vec<u8> = match (*scheme, v) {
13086 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
13087 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
13088 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
13089 _ => hmac_token(key, &v.encode_key()).to_vec(),
13090 };
13091 Some(Value::Bytes(token))
13092 }
13093
13094 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
13096 self.index_lookup_key_bytes(column_id, &v.encode_key())
13097 }
13098
13099 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
13102 {
13103 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
13104 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
13105 if *scheme == SCHEME_HMAC_EQ {
13106 return hmac_token(key, encoded).to_vec();
13107 }
13108 }
13109 }
13110 let _ = column_id;
13111 encoded.to_vec()
13112 }
13113}
13114
13115fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
13116 let columnar::NativeColumn::Int64 { data, validity } = col else {
13117 return false;
13118 };
13119 if data.len() < n || !columnar::all_non_null(validity, n) {
13120 return false;
13121 }
13122 data.iter()
13123 .take(n)
13124 .zip(data.iter().skip(1))
13125 .all(|(a, b)| a < b)
13126}
13127
13128#[derive(Debug, Clone)]
13132pub struct ColumnStat {
13133 pub min: Option<Value>,
13134 pub max: Option<Value>,
13135 pub null_count: u64,
13136}
13137
13138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13140pub enum NativeAgg {
13141 Count,
13142 Sum,
13143 Min,
13144 Max,
13145 Avg,
13146}
13147
13148#[derive(Debug, Clone, PartialEq)]
13150pub enum NativeAggResult {
13151 Count(u64),
13152 Int(i64),
13153 Float(f64),
13154 Null,
13156}
13157
13158#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13160pub enum ApproxAgg {
13161 Count,
13162 Sum,
13163 Avg,
13164}
13165
13166#[derive(Debug, Clone)]
13170pub struct ApproxResult {
13171 pub point: f64,
13173 pub ci_low: f64,
13175 pub ci_high: f64,
13177 pub n_population: u64,
13179 pub n_sample_live: usize,
13181 pub n_passing: usize,
13183}
13184
13185#[derive(Debug, Clone, PartialEq)]
13190pub enum AggState {
13191 Count(u64),
13193 SumI {
13195 sum: i128,
13196 count: u64,
13197 },
13198 SumF {
13200 sum: f64,
13201 count: u64,
13202 },
13203 AvgI {
13205 sum: i128,
13206 count: u64,
13207 },
13208 AvgF {
13210 sum: f64,
13211 count: u64,
13212 },
13213 MinI(i64),
13215 MaxI(i64),
13216 MinF(f64),
13218 MaxF(f64),
13219 Empty,
13221}
13222
13223impl AggState {
13224 pub fn merge(self, other: AggState) -> AggState {
13226 use AggState::*;
13227 match (self, other) {
13228 (Empty, x) | (x, Empty) => x,
13229 (Count(a), Count(b)) => Count(a + b),
13230 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
13231 sum: sa + sb,
13232 count: ca + cb,
13233 },
13234 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
13235 sum: sa + sb,
13236 count: ca + cb,
13237 },
13238 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
13239 sum: sa + sb,
13240 count: ca + cb,
13241 },
13242 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
13243 sum: sa + sb,
13244 count: ca + cb,
13245 },
13246 (MinI(a), MinI(b)) => MinI(a.min(b)),
13247 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
13248 (MinF(a), MinF(b)) => MinF(a.min(b)),
13249 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
13250 _ => Empty, }
13252 }
13253
13254 pub fn point(&self) -> Option<f64> {
13256 match self {
13257 AggState::Count(n) => Some(*n as f64),
13258 AggState::SumI { sum, .. } => Some(*sum as f64),
13259 AggState::SumF { sum, .. } => Some(*sum),
13260 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
13261 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
13262 AggState::MinI(n) => Some(*n as f64),
13263 AggState::MaxI(n) => Some(*n as f64),
13264 AggState::MinF(n) => Some(*n),
13265 AggState::MaxF(n) => Some(*n),
13266 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
13267 }
13268 }
13269
13270 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
13274 let is_float = matches!(ty, Some(TypeId::Float64));
13275 match (agg, result) {
13276 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
13277 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
13278 sum: x as i128,
13279 count: 1, },
13281 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
13282 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
13283 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
13284 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
13285 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
13286 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
13287 (NativeAgg::Count, _) => AggState::Empty,
13288 (_, NativeAggResult::Null) => AggState::Empty,
13289 _ => {
13290 let _ = is_float;
13291 AggState::Empty
13292 }
13293 }
13294 }
13295}
13296
13297#[derive(Debug, Clone)]
13300pub struct CachedAgg {
13301 pub state: AggState,
13302 pub watermark: u64,
13303 pub epoch: u64,
13304}
13305
13306#[derive(Debug, Clone)]
13308pub struct IncrementalAggResult {
13309 pub state: AggState,
13311 pub incremental: bool,
13314 pub delta_rows: u64,
13316}
13317
13318fn agg_state_from_rows(
13322 rows: &[Row],
13323 conditions: &[crate::query::Condition],
13324 index_sets: &[RowIdSet],
13325 column: Option<u16>,
13326 agg: NativeAgg,
13327 schema: &Schema,
13328 control: Option<&crate::ExecutionControl>,
13329) -> Result<AggState> {
13330 let mut count: u64 = 0;
13331 let mut sum_i: i128 = 0;
13332 let mut sum_f: f64 = 0.0;
13333 let mut mn_i: i64 = i64::MAX;
13334 let mut mx_i: i64 = i64::MIN;
13335 let mut mn_f: f64 = f64::INFINITY;
13336 let mut mx_f: f64 = f64::NEG_INFINITY;
13337 let mut saw_int = false;
13338 let mut saw_float = false;
13339 for (index, r) in rows.iter().enumerate() {
13340 execution_checkpoint(control, index)?;
13341 if !conditions
13342 .iter()
13343 .all(|c| condition_matches_row(c, r, schema))
13344 {
13345 continue;
13346 }
13347 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
13348 continue;
13349 }
13350 match agg {
13351 NativeAgg::Count => match column {
13352 None => count += 1,
13354 Some(cid) => match r.columns.get(&cid) {
13357 None | Some(Value::Null) => {}
13358 Some(_) => count += 1,
13359 },
13360 },
13361 _ => match column.and_then(|cid| r.columns.get(&cid)) {
13362 Some(Value::Int64(n)) => {
13363 count += 1;
13364 sum_i += *n as i128;
13365 mn_i = mn_i.min(*n);
13366 mx_i = mx_i.max(*n);
13367 saw_int = true;
13368 }
13369 Some(Value::Float64(f)) => {
13370 count += 1;
13371 sum_f += f;
13372 mn_f = mn_f.min(*f);
13373 mx_f = mx_f.max(*f);
13374 saw_float = true;
13375 }
13376 _ => {}
13377 },
13378 }
13379 }
13380 Ok(match agg {
13381 NativeAgg::Count => {
13382 if count == 0 {
13383 AggState::Empty
13384 } else {
13385 AggState::Count(count)
13386 }
13387 }
13388 NativeAgg::Sum => {
13389 if count == 0 {
13390 AggState::Empty
13391 } else if saw_int {
13392 AggState::SumI { sum: sum_i, count }
13393 } else {
13394 AggState::SumF { sum: sum_f, count }
13395 }
13396 }
13397 NativeAgg::Avg => {
13398 if count == 0 {
13399 AggState::Empty
13400 } else if saw_int {
13401 AggState::AvgI { sum: sum_i, count }
13402 } else {
13403 AggState::AvgF { sum: sum_f, count }
13404 }
13405 }
13406 NativeAgg::Min => {
13407 if !saw_int && !saw_float {
13408 AggState::Empty
13409 } else if saw_int {
13410 AggState::MinI(mn_i)
13411 } else {
13412 AggState::MinF(mn_f)
13413 }
13414 }
13415 NativeAgg::Max => {
13416 if !saw_int && !saw_float {
13417 AggState::Empty
13418 } else if saw_int {
13419 AggState::MaxI(mx_i)
13420 } else {
13421 AggState::MaxF(mx_f)
13422 }
13423 }
13424 })
13425}
13426
13427fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
13431 use crate::query::Condition;
13432 match c {
13433 Condition::Pk(key) => match schema.primary_key() {
13434 Some(pk) => row
13435 .columns
13436 .get(&pk.id)
13437 .map(|v| v.encode_key() == *key)
13438 .unwrap_or(false),
13439 None => false,
13440 },
13441 Condition::BitmapEq { column_id, value } => row
13442 .columns
13443 .get(column_id)
13444 .map(|v| v.encode_key() == *value)
13445 .unwrap_or(false),
13446 Condition::BitmapIn { column_id, values } => {
13447 let key = row.columns.get(column_id).map(|v| v.encode_key());
13448 match key {
13449 Some(k) => values.contains(&k),
13450 None => false,
13451 }
13452 }
13453 Condition::BytesPrefix { column_id, prefix } => row
13454 .columns
13455 .get(column_id)
13456 .map(|v| v.encode_key().starts_with(prefix))
13457 .unwrap_or(false),
13458 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
13459 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
13460 _ => false,
13461 },
13462 Condition::RangeF64 {
13463 column_id,
13464 lo,
13465 lo_inclusive,
13466 hi,
13467 hi_inclusive,
13468 } => match row.columns.get(column_id) {
13469 Some(Value::Float64(n)) => {
13470 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
13471 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
13472 lo_ok && hi_ok
13473 }
13474 _ => false,
13475 },
13476 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
13477 Some(Value::Bytes(b)) => {
13478 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
13479 }
13480 _ => false,
13481 },
13482 Condition::FmContainsAll {
13483 column_id,
13484 patterns,
13485 } => match row.columns.get(column_id) {
13486 Some(Value::Bytes(b)) => patterns
13487 .iter()
13488 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
13489 _ => false,
13490 },
13491 Condition::Ann { .. }
13492 | Condition::SparseMatch { .. }
13493 | Condition::MinHashSimilar { .. } => true,
13494 Condition::IsNull { column_id } => {
13495 matches!(row.columns.get(column_id), Some(Value::Null) | None)
13496 }
13497 Condition::IsNotNull { column_id } => {
13498 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
13499 }
13500 }
13501}
13502
13503fn as_f64(v: Option<&Value>) -> Option<f64> {
13505 match v {
13506 Some(Value::Int64(n)) => Some(*n as f64),
13507 Some(Value::Float64(f)) => Some(*f),
13508 _ => None,
13509 }
13510}
13511
13512fn accumulate_int(
13516 cursor: &mut dyn crate::cursor::Cursor,
13517 control: Option<&crate::ExecutionControl>,
13518) -> Result<(u64, i128, i64, i64)> {
13519 let mut count: u64 = 0;
13520 let mut sum: i128 = 0;
13521 let mut mn: i64 = i64::MAX;
13522 let mut mx: i64 = i64::MIN;
13523 while let Some(cols) = cursor.next_batch()? {
13524 execution_checkpoint(control, 0)?;
13525 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
13526 if crate::columnar::all_non_null(validity, data.len()) {
13527 count += data.len() as u64;
13529 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
13530 execution_checkpoint(control, chunk_index * 1024)?;
13531 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
13532 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
13533 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
13534 }
13535 } else {
13536 for (i, &v) in data.iter().enumerate() {
13537 execution_checkpoint(control, i)?;
13538 if crate::columnar::validity_bit(validity, i) {
13539 count += 1;
13540 sum += v as i128;
13541 mn = mn.min(v);
13542 mx = mx.max(v);
13543 }
13544 }
13545 }
13546 }
13547 }
13548 Ok((count, sum, mn, mx))
13549}
13550
13551fn accumulate_float(
13553 cursor: &mut dyn crate::cursor::Cursor,
13554 control: Option<&crate::ExecutionControl>,
13555) -> Result<(u64, f64, f64, f64)> {
13556 let mut count: u64 = 0;
13557 let mut sum: f64 = 0.0;
13558 let mut mn: f64 = f64::INFINITY;
13559 let mut mx: f64 = f64::NEG_INFINITY;
13560 while let Some(cols) = cursor.next_batch()? {
13561 execution_checkpoint(control, 0)?;
13562 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
13563 if crate::columnar::all_non_null(validity, data.len()) {
13564 count += data.len() as u64;
13565 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
13566 execution_checkpoint(control, chunk_index * 1024)?;
13567 sum += chunk.iter().sum::<f64>();
13568 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
13569 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
13570 }
13571 } else {
13572 for (i, &v) in data.iter().enumerate() {
13573 execution_checkpoint(control, i)?;
13574 if crate::columnar::validity_bit(validity, i) {
13575 count += 1;
13576 sum += v;
13577 mn = mn.min(v);
13578 mx = mx.max(v);
13579 }
13580 }
13581 }
13582 }
13583 }
13584 Ok((count, sum, mn, mx))
13585}
13586
13587#[inline]
13588fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
13589 if index.is_multiple_of(256) {
13590 control
13591 .map(crate::ExecutionControl::checkpoint)
13592 .transpose()?;
13593 }
13594 Ok(())
13595}
13596
13597fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
13598 if count == 0 && !matches!(agg, NativeAgg::Count) {
13599 return NativeAggResult::Null;
13600 }
13601 match agg {
13602 NativeAgg::Count => NativeAggResult::Count(count),
13603 NativeAgg::Sum => match sum.try_into() {
13606 Ok(v) => NativeAggResult::Int(v),
13607 Err(_) => NativeAggResult::Null,
13608 },
13609 NativeAgg::Min => NativeAggResult::Int(mn),
13610 NativeAgg::Max => NativeAggResult::Int(mx),
13611 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
13612 }
13613}
13614
13615fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
13616 if count == 0 && !matches!(agg, NativeAgg::Count) {
13617 return NativeAggResult::Null;
13618 }
13619 match agg {
13620 NativeAgg::Count => NativeAggResult::Count(count),
13621 NativeAgg::Sum => NativeAggResult::Float(sum),
13622 NativeAgg::Min => NativeAggResult::Float(mn),
13623 NativeAgg::Max => NativeAggResult::Float(mx),
13624 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
13625 }
13626}
13627
13628fn agg_int(
13631 stats: &[crate::page::PageStat],
13632 decode: fn(Option<&[u8]>) -> Option<i64>,
13633) -> Option<(Option<i64>, Option<i64>, u64)> {
13634 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
13635 let mut any = false;
13636 for s in stats {
13637 if let Some(v) = decode(s.min.as_deref()) {
13638 mn = mn.min(v);
13639 any = true;
13640 }
13641 if let Some(v) = decode(s.max.as_deref()) {
13642 mx = mx.max(v);
13643 any = true;
13644 }
13645 nulls += s.null_count;
13646 }
13647 any.then_some((Some(mn), Some(mx), nulls))
13648}
13649
13650fn agg_float(
13652 stats: &[crate::page::PageStat],
13653 decode: fn(Option<&[u8]>) -> Option<f64>,
13654) -> Option<(Option<f64>, Option<f64>, u64)> {
13655 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
13656 let mut any = false;
13657 for s in stats {
13658 if let Some(v) = decode(s.min.as_deref()) {
13659 mn = mn.min(v);
13660 any = true;
13661 }
13662 if let Some(v) = decode(s.max.as_deref()) {
13663 mx = mx.max(v);
13664 any = true;
13665 }
13666 nulls += s.null_count;
13667 }
13668 any.then_some((Some(mn), Some(mx), nulls))
13669}
13670
13671type SecondaryIndexes = (
13673 HashMap<u16, BitmapIndex>,
13674 HashMap<u16, AnnIndex>,
13675 HashMap<u16, FmIndex>,
13676 HashMap<u16, SparseIndex>,
13677 HashMap<u16, MinHashIndex>,
13678);
13679
13680fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
13681 let mut bitmap = HashMap::new();
13682 let mut ann = HashMap::new();
13683 let mut fm = HashMap::new();
13684 let mut sparse = HashMap::new();
13685 let mut minhash = HashMap::new();
13686 for idef in &schema.indexes {
13687 match idef.kind {
13688 IndexKind::Bitmap => {
13689 bitmap.insert(idef.column_id, BitmapIndex::new());
13690 }
13691 IndexKind::Ann => {
13692 let dim = schema
13693 .columns
13694 .iter()
13695 .find(|c| c.id == idef.column_id)
13696 .and_then(|c| match c.ty {
13697 TypeId::Embedding { dim } => Some(dim as usize),
13698 _ => None,
13699 })
13700 .unwrap_or(0);
13701 let options = idef.options.ann.clone().unwrap_or_default();
13702 ann.insert(
13703 idef.column_id,
13704 AnnIndex::with_full_options(
13705 dim,
13706 options.m,
13707 options.ef_construction,
13708 options.ef_search,
13709 &options,
13710 ),
13711 );
13712 }
13713 IndexKind::FmIndex => {
13714 fm.insert(idef.column_id, FmIndex::new());
13715 }
13716 IndexKind::Sparse => {
13717 sparse.insert(idef.column_id, SparseIndex::new());
13718 }
13719 IndexKind::MinHash => {
13720 let options = idef.options.minhash.clone().unwrap_or_default();
13721 minhash.insert(
13722 idef.column_id,
13723 MinHashIndex::with_options(options.permutations, options.bands),
13724 );
13725 }
13726 _ => {}
13727 }
13728 }
13729 (bitmap, ann, fm, sparse, minhash)
13730}
13731
13732const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
13733 | ColumnFlags::AUTO_INCREMENT
13734 | ColumnFlags::ENCRYPTED
13735 | ColumnFlags::ENCRYPTED_INDEXABLE
13736 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
13737
13738fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
13739 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
13740 return Err(MongrelError::Schema(
13741 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
13742 ));
13743 }
13744 Ok(())
13745}
13746
13747fn validate_alter_column_type(
13748 schema: &Schema,
13749 old: &ColumnDef,
13750 next: &ColumnDef,
13751 has_stored_versions: bool,
13752) -> Result<()> {
13753 if old.ty == next.ty {
13754 return Ok(());
13755 }
13756 if schema.indexes.iter().any(|i| i.column_id == old.id) {
13757 return Err(MongrelError::Schema(format!(
13758 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
13759 old.name
13760 )));
13761 }
13762 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
13763 return Ok(());
13764 }
13765 Err(MongrelError::Schema(format!(
13766 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
13767 old.ty, next.ty
13768 )))
13769}
13770
13771fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
13772 matches!(
13773 (old, new),
13774 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
13775 )
13776}
13777
13778fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
13784 let mut prev: Option<i64> = None;
13785 for r in rows {
13786 match r.columns.get(&pk_id) {
13787 Some(Value::Int64(v)) => {
13788 if prev.is_some_and(|p| p >= *v) {
13789 return false;
13790 }
13791 prev = Some(*v);
13792 }
13793 _ => return false,
13794 }
13795 }
13796 true
13797}
13798
13799#[allow(clippy::too_many_arguments)]
13800fn index_into(
13801 schema: &Schema,
13802 row: &Row,
13803 hot: &mut HotIndex,
13804 bitmap: &mut HashMap<u16, BitmapIndex>,
13805 ann: &mut HashMap<u16, AnnIndex>,
13806 fm: &mut HashMap<u16, FmIndex>,
13807 sparse: &mut HashMap<u16, SparseIndex>,
13808 minhash: &mut HashMap<u16, MinHashIndex>,
13809) {
13810 for idef in &schema.indexes {
13811 let Some(val) = row.columns.get(&idef.column_id) else {
13812 continue;
13813 };
13814 match idef.kind {
13815 IndexKind::Bitmap => {
13816 if let Some(b) = bitmap.get_mut(&idef.column_id) {
13817 b.insert(val.encode_key(), row.row_id);
13818 }
13819 }
13820 IndexKind::Ann => {
13821 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
13822 if let Some(meta) = val.generated_embedding_metadata() {
13823 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
13825 continue;
13826 }
13827 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
13828 .is_err()
13829 {
13830 continue;
13831 }
13832 }
13833 a.insert_validated(v, row.row_id);
13834 }
13835 }
13836 IndexKind::FmIndex => {
13837 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
13838 f.insert(b.clone(), row.row_id);
13839 }
13840 }
13841 IndexKind::Sparse => {
13842 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
13843 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
13846 s.insert(&terms, row.row_id);
13847 }
13848 }
13849 }
13850 IndexKind::MinHash => {
13851 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
13852 let tokens = crate::index::token_hashes_from_bytes(b);
13855 mh.insert(&tokens, row.row_id);
13856 }
13857 }
13858 _ => {}
13859 }
13860 }
13861 if let Some(pk_col) = schema.primary_key() {
13862 if let Some(pk_val) = row.columns.get(&pk_col.id) {
13863 hot.insert(pk_val.encode_key(), row.row_id);
13864 }
13865 }
13866}
13867
13868#[allow(clippy::too_many_arguments)]
13871fn index_into_single(
13872 idef: &IndexDef,
13873 _schema: &Schema,
13874 row: &Row,
13875 _hot: &mut HotIndex,
13876 bitmap: &mut HashMap<u16, BitmapIndex>,
13877 ann: &mut HashMap<u16, AnnIndex>,
13878 fm: &mut HashMap<u16, FmIndex>,
13879 sparse: &mut HashMap<u16, SparseIndex>,
13880 minhash: &mut HashMap<u16, MinHashIndex>,
13881) {
13882 let Some(val) = row.columns.get(&idef.column_id) else {
13883 return;
13884 };
13885 match idef.kind {
13886 IndexKind::Bitmap => {
13887 if let Some(b) = bitmap.get_mut(&idef.column_id) {
13888 b.insert(val.encode_key(), row.row_id);
13889 }
13890 }
13891 IndexKind::Ann => {
13892 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
13893 if let Some(meta) = val.generated_embedding_metadata() {
13894 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
13896 return;
13897 }
13898 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
13899 .is_err()
13900 {
13901 return;
13902 }
13903 }
13904 a.insert_validated(v, row.row_id);
13905 }
13906 }
13907 IndexKind::FmIndex => {
13908 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
13909 f.insert(b.clone(), row.row_id);
13910 }
13911 }
13912 IndexKind::Sparse => {
13913 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
13914 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
13915 s.insert(&terms, row.row_id);
13916 }
13917 }
13918 }
13919 IndexKind::MinHash => {
13920 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
13921 let tokens = crate::index::token_hashes_from_bytes(b);
13922 mh.insert(&tokens, row.row_id);
13923 }
13924 }
13925 _ => {}
13926 }
13927}
13928
13929fn eval_partial_predicate(
13935 pred: &str,
13936 columns_map: &HashMap<u16, &Value>,
13937 name_to_id: &HashMap<&str, u16>,
13938) -> bool {
13939 let lower = pred.trim().to_ascii_lowercase();
13940 if let Some(rest) = lower.strip_suffix(" is not null") {
13942 let col_name = rest.trim();
13943 if let Some(col_id) = name_to_id.get(col_name) {
13944 return columns_map
13945 .get(col_id)
13946 .is_some_and(|v| !matches!(v, Value::Null));
13947 }
13948 }
13949 if let Some(rest) = lower.strip_suffix(" is null") {
13951 let col_name = rest.trim();
13952 if let Some(col_id) = name_to_id.get(col_name) {
13953 return columns_map
13954 .get(col_id)
13955 .is_none_or(|v| matches!(v, Value::Null));
13956 }
13957 }
13958 true
13961}
13962
13963#[allow(dead_code)]
13969fn bulk_index_key(
13970 column_keys: &HashMap<u16, ([u8; 32], u8)>,
13971 column_id: u16,
13972 ty: TypeId,
13973 col: &columnar::NativeColumn,
13974 i: usize,
13975) -> Option<Vec<u8>> {
13976 let encoded = columnar::encode_key_native(ty, col, i)?;
13977 {
13978 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
13979 if let Some((key, scheme)) = column_keys.get(&column_id) {
13980 return Some(match (*scheme, col) {
13981 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
13982 (_, columnar::NativeColumn::Int64 { data, .. }) => {
13983 ope_token_i64(key, data[i]).to_vec()
13984 }
13985 (_, columnar::NativeColumn::Float64 { data, .. }) => {
13986 ope_token_f64(key, data[i]).to_vec()
13987 }
13988 _ => hmac_token(key, &encoded).to_vec(),
13989 });
13990 }
13991 }
13992 Some(encoded)
13993}
13994
13995pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
13996 write_schema_with_after(dir, schema, || {})
13997}
13998
13999pub(crate) fn write_schema_durable(
14000 root: &crate::durable_file::DurableRoot,
14001 schema: &Schema,
14002) -> Result<()> {
14003 write_schema_durable_with_after(root, schema, || {})
14004}
14005
14006fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
14007where
14008 F: FnOnce(),
14009{
14010 let json = serde_json::to_string_pretty(schema)
14011 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
14012 crate::durable_file::write_atomic_with_after(
14013 &dir.join(SCHEMA_FILENAME),
14014 json.as_bytes(),
14015 after_publish,
14016 )?;
14017 Ok(())
14018}
14019
14020fn write_schema_durable_with_after<F>(
14021 root: &crate::durable_file::DurableRoot,
14022 schema: &Schema,
14023 after_publish: F,
14024) -> Result<()>
14025where
14026 F: FnOnce(),
14027{
14028 let json = serde_json::to_string_pretty(schema)
14029 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
14030 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
14031 Ok(())
14032}
14033
14034fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
14035 let mut schema_published = false;
14036 let schema_result = match table._root_guard.as_deref() {
14037 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
14038 schema_published = true;
14039 }),
14040 None => write_schema_with_after(&table.dir, &table.schema, || {
14041 schema_published = true;
14042 }),
14043 };
14044 if schema_result.is_err() && !schema_published {
14045 return schema_result;
14046 }
14047 match table.persist_manifest(table.current_epoch()) {
14048 Ok(()) => Ok(()),
14049 Err(manifest_error) => Err(match schema_result {
14050 Ok(()) => manifest_error,
14051 Err(schema_error) => MongrelError::Other(format!(
14052 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
14053 )),
14054 }),
14055 }
14056}
14057
14058fn read_schema(dir: &Path) -> Result<Schema> {
14059 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
14060 read_schema_file(file)
14061}
14062
14063fn read_schema_file(file: std::fs::File) -> Result<Schema> {
14064 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
14065 use std::io::Read;
14066
14067 let length = file.metadata()?.len();
14068 if length > MAX_SCHEMA_BYTES {
14069 return Err(MongrelError::ResourceLimitExceeded {
14070 resource: "schema bytes",
14071 requested: usize::try_from(length).unwrap_or(usize::MAX),
14072 limit: MAX_SCHEMA_BYTES as usize,
14073 });
14074 }
14075 let mut bytes = Vec::with_capacity(length as usize);
14076 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
14077 if bytes.len() as u64 != length {
14078 return Err(MongrelError::Schema(
14079 "schema length changed while reading".into(),
14080 ));
14081 }
14082 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
14083}
14084
14085fn preflight_standalone_open(
14086 dir: &Path,
14087 runs_root: Option<&crate::durable_file::DurableRoot>,
14088 idx_root: Option<&crate::durable_file::DurableRoot>,
14089 manifest: &Manifest,
14090 schema: &Schema,
14091 records: &[crate::wal::Record],
14092 kek: Option<Arc<Kek>>,
14093) -> Result<()> {
14094 crate::wal::validate_shared_transaction_framing(records)?;
14095 if manifest.schema_id > schema.schema_id
14096 || manifest.flushed_epoch > manifest.current_epoch
14097 || manifest.global_idx_epoch > manifest.current_epoch
14098 || manifest.next_row_id == u64::MAX
14099 || manifest.auto_inc_next < 0
14100 || manifest.auto_inc_next == i64::MAX
14101 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
14102 {
14103 return Err(MongrelError::InvalidArgument(
14104 "manifest counters or schema identity are invalid".into(),
14105 ));
14106 }
14107 let mut run_ids = HashSet::new();
14108 let mut maximum_row_id = None::<u64>;
14109 for run in &manifest.runs {
14110 if run.run_id >= u64::MAX as u128
14111 || !run_ids.insert(run.run_id)
14112 || run.epoch_created > manifest.current_epoch
14113 {
14114 return Err(MongrelError::InvalidArgument(
14115 "manifest contains an invalid or duplicate active run".into(),
14116 ));
14117 }
14118 let mut reader = match runs_root {
14119 Some(root) => RunReader::open_file(
14120 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
14121 schema.clone(),
14122 kek.clone(),
14123 )?,
14124 None => RunReader::open(
14125 dir.join(RUNS_DIR)
14126 .join(format!("r-{}.sr", run.run_id as u64)),
14127 schema.clone(),
14128 kek.clone(),
14129 )?,
14130 };
14131 let header = reader.header();
14132 if header.run_id != run.run_id
14133 || header.level != run.level
14134 || header.row_count != run.row_count
14135 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
14136 || header.is_uniform_epoch() && header.epoch_created != 0
14137 || header.schema_id > schema.schema_id
14138 {
14139 return Err(MongrelError::InvalidArgument(format!(
14140 "run {} differs from its manifest",
14141 run.run_id
14142 )));
14143 }
14144 if header.row_count != 0 {
14145 maximum_row_id = Some(
14146 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
14147 );
14148 }
14149 reader.validate_all_pages()?;
14150 }
14151 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
14152 return Err(MongrelError::InvalidArgument(
14153 "manifest next_row_id does not advance beyond persisted rows".into(),
14154 ));
14155 }
14156 for run in &manifest.retiring {
14157 if run.run_id >= u64::MAX as u128
14158 || run.retire_epoch > manifest.current_epoch
14159 || !run_ids.insert(run.run_id)
14160 {
14161 return Err(MongrelError::InvalidArgument(
14162 "manifest contains an invalid or duplicate retired run".into(),
14163 ));
14164 }
14165 }
14166 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
14167 match idx_root {
14168 Some(root) => {
14169 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
14170 }
14171 None => {
14172 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
14173 }
14174 }
14175
14176 let committed = records
14177 .iter()
14178 .filter_map(|record| match record.op {
14179 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
14180 _ => None,
14181 })
14182 .collect::<HashMap<_, _>>();
14183 for record in records {
14184 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
14185 continue;
14186 };
14187 match &record.op {
14188 Op::Put { table_id, rows } => {
14189 if *table_id != manifest.table_id {
14190 return Err(MongrelError::CorruptWal {
14191 offset: record.seq.0,
14192 reason: format!(
14193 "private WAL record references table {table_id}, expected {}",
14194 manifest.table_id
14195 ),
14196 });
14197 }
14198 let rows: Vec<Row> =
14199 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
14200 offset: record.seq.0,
14201 reason: format!("committed Put payload could not be decoded: {error}"),
14202 })?;
14203 for row in rows {
14204 if row.deleted || row.row_id.0 == u64::MAX {
14205 return Err(MongrelError::CorruptWal {
14206 offset: record.seq.0,
14207 reason: "committed Put contains an invalid row identity".into(),
14208 });
14209 }
14210 let cells = row.columns.into_iter().collect::<Vec<_>>();
14211 schema
14212 .validate_values(&cells)
14213 .map_err(|error| MongrelError::CorruptWal {
14214 offset: record.seq.0,
14215 reason: format!("committed Put violates table schema: {error}"),
14216 })?;
14217 if schema.auto_increment_column().is_some_and(|column| {
14218 matches!(
14219 cells.iter().find(|(id, _)| *id == column.id),
14220 Some((_, Value::Int64(value))) if *value == i64::MAX
14221 )
14222 }) {
14223 return Err(MongrelError::CorruptWal {
14224 offset: record.seq.0,
14225 reason: "committed Put exhausts AUTO_INCREMENT".into(),
14226 });
14227 }
14228 }
14229 }
14230 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
14231 if *table_id != manifest.table_id =>
14232 {
14233 return Err(MongrelError::CorruptWal {
14234 offset: record.seq.0,
14235 reason: format!(
14236 "private WAL record references table {table_id}, expected {}",
14237 manifest.table_id
14238 ),
14239 });
14240 }
14241 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
14242 return Err(MongrelError::CorruptWal {
14243 offset: record.seq.0,
14244 reason: "private WAL contains shared spilled-run metadata".into(),
14245 });
14246 }
14247 _ => {}
14248 }
14249 }
14250 Ok(())
14251}
14252
14253fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
14254 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
14255}
14256
14257fn wal_segment_number(path: &Path) -> Option<u64> {
14258 path.file_stem()
14259 .and_then(|stem| stem.to_str())
14260 .and_then(|stem| stem.strip_prefix("seg-"))
14261 .and_then(|number| number.parse().ok())
14262}
14263
14264fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
14265 let n = list_wal_numbers(wal_dir)?;
14266 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
14267}
14268
14269fn next_wal_number(wal_dir: &Path) -> Result<u32> {
14270 list_wal_numbers(wal_dir)?
14271 .map(|maximum| {
14272 maximum
14273 .checked_add(1)
14274 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
14275 })
14276 .unwrap_or(Ok(0))
14277}
14278
14279fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
14280 let mut max_n = None;
14281 let entries = match std::fs::read_dir(wal_dir) {
14282 Ok(entries) => entries,
14283 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
14284 Err(error) => return Err(error.into()),
14285 };
14286 for entry in entries {
14287 let entry = entry?;
14288 let fname = entry.file_name();
14289 let Some(s) = fname.to_str() else {
14290 continue;
14291 };
14292 let Some(stripped) = s.strip_prefix("seg-") else {
14293 continue;
14294 };
14295 let Some(number) = stripped.strip_suffix(".wal") else {
14296 return Err(MongrelError::CorruptWal {
14297 offset: 0,
14298 reason: format!("malformed WAL segment name {s:?}"),
14299 });
14300 };
14301 let n = number
14302 .parse::<u32>()
14303 .map_err(|_| MongrelError::CorruptWal {
14304 offset: 0,
14305 reason: format!("malformed WAL segment name {s:?}"),
14306 })?;
14307 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
14308 return Err(MongrelError::CorruptWal {
14309 offset: n as u64,
14310 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
14311 });
14312 }
14313 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
14314 }
14315 Ok(max_n)
14316}