1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, EpochGuard, MaintenanceReceipt, Snapshot, VersionStamp};
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::result_cache::{
24 DrainOutcome, PersistContext, PersistenceDisabledReason, PersistentCacheIdentity,
25 PersistentPublicationState, FRAME_MAGIC,
26};
27use crate::row_id_set::RowIdSet;
28use crate::rowid::{RowId, RowIdAllocator};
29use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
30use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
31use crate::txn::{GroupCommit, OwnedRow};
32use crate::wal::{Op, SharedWal, Wal};
33use crate::{MongrelError, Result};
34use arc_swap::ArcSwap;
35use mongreldb_types::hlc::HlcTimestamp;
36use std::borrow::Cow;
37use std::cmp::Reverse;
38use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
39use std::path::{Path, PathBuf};
40use std::sync::atomic::AtomicBool;
41use std::sync::Arc;
42use std::time::Instant;
43use zeroize::Zeroizing;
44
45pub const WAL_DIR: &str = "_wal";
46pub const RUNS_DIR: &str = "_runs";
47pub const CACHE_DIR: &str = "_cache";
48pub const META_DIR: &str = "_meta";
49pub const RCACHE_DIR: &str = "_rcache";
50pub const KEYS_FILENAME: &str = "keys";
51pub const SCHEMA_FILENAME: &str = "schema.json";
52
53fn derive_next_run_id(
54 dir: &Path,
55 runs_root: Option<&crate::durable_file::DurableRoot>,
56 active: &[RunRef],
57 retiring: &[crate::manifest::RetiredRun],
58) -> Result<u64> {
59 let mut maximum = 0_u64;
60 for run_id in active
61 .iter()
62 .map(|run| run.run_id)
63 .chain(retiring.iter().map(|run| run.run_id))
64 {
65 let run_id = u64::try_from(run_id)
66 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
67 maximum = maximum.max(run_id);
68 }
69 let names = match runs_root {
70 Some(root) => root.list_regular_files(".")?,
71 None => std::fs::read_dir(dir.join(RUNS_DIR))?
72 .map(|entry| entry.map(|entry| entry.file_name()))
73 .collect::<std::io::Result<Vec<_>>>()?,
74 };
75 for name in names {
76 let Some(name) = name.to_str() else {
77 continue;
78 };
79 let Some(digits) = name
80 .strip_prefix("r-")
81 .and_then(|name| name.strip_suffix(".sr"))
82 else {
83 continue;
84 };
85 let Ok(run_id) = digits.parse::<u64>() else {
86 continue;
87 };
88 if name == format!("r-{run_id}.sr") {
89 maximum = maximum.max(run_id);
90 }
91 }
92 maximum
93 .checked_add(1)
94 .map(|next| next.max(1))
95 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
96}
97
98enum ControlledVisibleCandidate<'a> {
99 Memory(Row),
100 Memtable(RowId, Epoch, Cow<'a, Row>),
106 MutableRun(RowId, Epoch, &'a Row),
108 Run(RunVisibleVersion),
109}
110
111impl<'a> ControlledVisibleCandidate<'a> {
112 fn row_id(&self) -> RowId {
113 match self {
114 Self::Memory(row) => row.row_id,
115 Self::Memtable(rid, _, _) => *rid,
116 Self::MutableRun(rid, _, _) => *rid,
117 Self::Run(version) => version.row_id,
118 }
119 }
120
121 fn committed_epoch(&self) -> Epoch {
122 match self {
123 Self::Memory(row) => row.committed_epoch,
124 Self::Memtable(_, epoch, _) => *epoch,
125 Self::MutableRun(_, epoch, _) => *epoch,
126 Self::Run(version) => version.committed_epoch,
127 }
128 }
129
130 fn deleted(&self) -> bool {
131 match self {
132 Self::Memory(row) => row.deleted,
133 Self::Memtable(_, _, row) => row.deleted,
134 Self::MutableRun(_, _, row) => row.deleted,
135 Self::Run(version) => version.deleted,
136 }
137 }
138
139 fn commit_ts(&self) -> Option<mongreldb_types::hlc::HlcTimestamp> {
140 match self {
141 Self::Memory(row) => row.commit_ts,
142 Self::Memtable(_, _, row) => row.commit_ts,
143 Self::MutableRun(_, _, row) => row.commit_ts,
144 Self::Run(version) => version.commit_ts,
147 }
148 }
149}
150
151enum ControlledVisibleCursor<'a> {
152 Memory(std::vec::IntoIter<Row>),
154 MemoryStreaming {
158 active: std::vec::IntoIter<Row>,
159 rest: std::collections::btree_map::IntoValues<RowId, Row>,
160 batch_cap: usize,
161 #[allow(dead_code)]
163 peak_active: usize,
164 #[allow(dead_code)]
166 total: usize,
167 },
168 Memtable(MemtableVisibleVersionCursor<'a>),
172 MutableRun(MutableRunVisibleVersionCursor<'a>),
174 Run(Box<RunVisibleVersionCursor>),
175 #[cfg(test)]
176 Synthetic {
177 next: u64,
178 end: u64,
179 },
180}
181
182#[allow(dead_code)] const CONTROLLED_HOT_BATCH: usize = 256;
185
186struct ControlledVisibleSource<'a> {
187 cursor: ControlledVisibleCursor<'a>,
188 current: Option<ControlledVisibleCandidate<'a>>,
189}
190
191impl<'a> ControlledVisibleSource<'a> {
192 #[cfg(test)]
194 fn memory(rows: Vec<Row>) -> Self {
195 Self {
196 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
197 current: None,
198 }
199 }
200
201 #[allow(dead_code)] fn memory_from_map(map: BTreeMap<RowId, Row>) -> Self {
206 let total = map.len();
207 let mut values = map.into_values();
208 if total > CONTROLLED_HOT_BATCH {
209 let active: Vec<Row> = values.by_ref().take(CONTROLLED_HOT_BATCH).collect();
210 let peak = active.len();
211 Self {
212 cursor: ControlledVisibleCursor::MemoryStreaming {
213 active: active.into_iter(),
214 rest: values,
215 batch_cap: CONTROLLED_HOT_BATCH,
216 peak_active: peak,
217 total,
218 },
219 current: None,
220 }
221 } else {
222 let active: Vec<Row> = values.collect();
223 Self {
224 cursor: ControlledVisibleCursor::Memory(active.into_iter()),
225 current: None,
226 }
227 }
228 }
229
230 #[allow(dead_code)] fn memtable_cursor(cursor: MemtableVisibleVersionCursor<'a>) -> Self {
235 Self {
236 cursor: ControlledVisibleCursor::Memtable(cursor),
237 current: None,
238 }
239 }
240
241 #[allow(dead_code)] fn mutable_run_cursor(cursor: MutableRunVisibleVersionCursor<'a>) -> Self {
246 Self {
247 cursor: ControlledVisibleCursor::MutableRun(cursor),
248 current: None,
249 }
250 }
251
252 fn run(cursor: RunVisibleVersionCursor) -> Self {
253 Self {
254 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
255 current: None,
256 }
257 }
258
259 #[cfg(test)]
260 fn synthetic(end: u64) -> Self {
261 Self {
262 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
263 current: None,
264 }
265 }
266
267 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
268 self.current = match &mut self.cursor {
269 ControlledVisibleCursor::Memory(rows) => {
270 rows.next().map(ControlledVisibleCandidate::Memory)
271 }
272 ControlledVisibleCursor::MemoryStreaming {
273 active,
274 rest,
275 batch_cap,
276 peak_active,
277 total: _,
278 } => {
279 if let Some(row) = active.next() {
280 Some(ControlledVisibleCandidate::Memory(row))
281 } else {
282 control.checkpoint()?;
283 let next_batch: Vec<Row> = rest.by_ref().take(*batch_cap).collect();
284 if next_batch.is_empty() {
285 None
286 } else {
287 *peak_active = (*peak_active).max(next_batch.len());
288 *active = next_batch.into_iter();
289 active.next().map(ControlledVisibleCandidate::Memory)
290 }
291 }
292 }
293 ControlledVisibleCursor::Memtable(iter) => {
294 let prev_peak = iter.peak_examined;
295 let next = iter
299 .next_controlled(control)?
300 .map(|(rid, epoch, row)| ControlledVisibleCandidate::Memtable(rid, epoch, row));
301 if let Some(peak) = iter.peak_examined.checked_sub(prev_peak) {
302 crate::trace::QueryTrace::record(|t| {
303 t.controlled_scan_peak_same_row_versions = t
304 .controlled_scan_peak_same_row_versions
305 .saturating_add(peak);
306 });
307 }
308 next
309 }
310 ControlledVisibleCursor::MutableRun(iter) => {
311 let prev_peak = iter.peak_examined;
312 let next = iter.next().map(|(rid, epoch, row)| {
313 ControlledVisibleCandidate::MutableRun(rid, epoch, row)
314 });
315 if let Some(peak) = iter.peak_examined.checked_sub(prev_peak) {
316 crate::trace::QueryTrace::record(|t| {
317 t.controlled_scan_peak_same_row_versions = t
318 .controlled_scan_peak_same_row_versions
319 .saturating_add(peak);
320 });
321 }
322 next
323 }
324 ControlledVisibleCursor::Run(cursor) => cursor
325 .next_visible_version(control)?
326 .map(ControlledVisibleCandidate::Run),
327 #[cfg(test)]
328 ControlledVisibleCursor::Synthetic { next, end } => {
329 if *next > *end {
330 None
331 } else {
332 let row = Row::new(RowId(*next), Epoch(1));
333 *next += 1;
334 Some(ControlledVisibleCandidate::Memory(row))
335 }
336 }
337 };
338 Ok(())
339 }
340
341 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate<'a>> {
342 let current = self.current.take().ok_or_else(|| {
343 MongrelError::Other("controlled visible source was not primed".into())
344 })?;
345 self.advance(control)?;
346 Ok(current)
347 }
348
349 fn materialize(
350 &mut self,
351 candidate: ControlledVisibleCandidate<'a>,
352 control: &crate::ExecutionControl,
353 ) -> Result<Row> {
354 match candidate {
355 ControlledVisibleCandidate::Memory(row) => Ok(row),
356 ControlledVisibleCandidate::Memtable(_, _, row) => Ok(row.into_owned()),
357 ControlledVisibleCandidate::MutableRun(_, _, row) => Ok(row.clone()),
358 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
359 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
360 _ => Err(MongrelError::Other(
361 "run candidate escaped its controlled cursor".into(),
362 )),
363 },
364 }
365 }
366
367 #[cfg(test)]
368 fn is_streaming_memory(&self) -> bool {
369 matches!(self.cursor, ControlledVisibleCursor::MemoryStreaming { .. })
370 }
371
372 #[cfg(test)]
373 fn streaming_peak_active(&self) -> Option<usize> {
374 match &self.cursor {
375 ControlledVisibleCursor::MemoryStreaming { peak_active, .. } => Some(*peak_active),
376 _ => None,
377 }
378 }
379
380 #[cfg(test)]
381 fn streaming_total(&self) -> Option<usize> {
382 match &self.cursor {
383 ControlledVisibleCursor::MemoryStreaming { total, .. } => Some(*total),
384 _ => None,
385 }
386 }
387}
388
389fn merge_controlled_visible_sources<'a>(
390 sources: &mut [ControlledVisibleSource<'a>],
391 control: &crate::ExecutionControl,
392 mut expired: impl FnMut(&Row) -> bool,
393 mut visit: impl FnMut(Row) -> Result<()>,
394) -> Result<()> {
395 let mut heap = BinaryHeap::new();
396 for (source_index, source) in sources.iter_mut().enumerate() {
397 source.advance(control)?;
398 if let Some(candidate) = &source.current {
399 heap.push(Reverse((candidate.row_id(), source_index)));
400 }
401 }
402 let mut merged = 0_usize;
403 let mut peak_buffer_rows: usize = heap.len();
404 let mut peak_same_row_versions: usize = 0;
405 let mut cancel_started: Option<std::time::Instant> = None;
406 while let Some(Reverse((row_id, source_index))) = heap.pop() {
407 if merged.is_multiple_of(256) {
408 control.checkpoint()?;
409 }
410 merged += 1;
411 let mut best_source = source_index;
412 let mut best = sources[source_index].pop(control)?;
413 if let Some(next) = &sources[source_index].current {
414 heap.push(Reverse((next.row_id(), source_index)));
415 }
416 let mut same_row_versions: usize = 1;
417 while heap
418 .peek()
419 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
420 {
421 control.checkpoint()?;
422 let Some(Reverse((_, source_index))) = heap.pop() else {
423 break;
424 };
425 let candidate = sources[source_index].pop(control)?;
426 same_row_versions += 1;
427 if Snapshot::version_is_newer(
431 candidate.committed_epoch(),
432 candidate.commit_ts(),
433 best.committed_epoch(),
434 best.commit_ts(),
435 ) {
436 best = candidate;
437 best_source = source_index;
438 }
439 if let Some(next) = &sources[source_index].current {
440 heap.push(Reverse((next.row_id(), source_index)));
441 }
442 }
443 peak_same_row_versions = peak_same_row_versions.max(same_row_versions);
444 peak_buffer_rows = peak_buffer_rows.max(heap.len());
445 if best.deleted() {
446 continue;
447 }
448 let row = sources[best_source].materialize(best, control)?;
449 if !expired(&row) {
450 visit(row)?;
451 }
452 if cancel_started.is_none() && control.is_cancelled() {
453 cancel_started = Some(std::time::Instant::now());
454 }
455 }
456 crate::trace::QueryTrace::record(|t| {
457 t.controlled_scan_source_refills = t.controlled_scan_source_refills.saturating_add(0); t.controlled_scan_peak_source_buffer_rows = t
459 .controlled_scan_peak_source_buffer_rows
460 .saturating_add(peak_buffer_rows);
461 t.controlled_scan_peak_same_row_versions = t
462 .controlled_scan_peak_same_row_versions
463 .saturating_add(peak_same_row_versions);
464 });
465 control.checkpoint()?;
466 Ok(())
467}
468
469#[cfg(test)]
470mod controlled_visible_cursor_tests {
471 use super::*;
472
473 #[test]
474 fn streams_more_than_one_million_rows_without_a_source_cap() {
475 let control = crate::ExecutionControl::new(None);
476 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
477 let mut count = 0_u64;
478 let mut last = 0_u64;
479 merge_controlled_visible_sources(
480 &mut sources,
481 &control,
482 |_| false,
483 |row| {
484 count += 1;
485 assert!(row.row_id.0 > last);
486 last = row.row_id.0;
487 Ok(())
488 },
489 )
490 .unwrap();
491 assert_eq!(count, 1_000_001);
492 assert_eq!(last, 1_000_001);
493 }
494
495 #[test]
496 fn merge_orders_rows_and_honors_newest_tombstones() {
497 let control = crate::ExecutionControl::new(None);
498 let older = vec![
499 Row::new(RowId(1), Epoch(1)),
500 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
501 Row::new(RowId(4), Epoch(1)),
502 ];
503 let mut deleted = Row::new(RowId(1), Epoch(2));
504 deleted.deleted = true;
505 let newer = vec![
506 deleted,
507 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
508 Row::new(RowId(3), Epoch(2)),
509 ];
510 let mut sources = vec![
511 ControlledVisibleSource::memory(older),
512 ControlledVisibleSource::memory(newer),
513 ];
514 let mut rows = Vec::new();
515 merge_controlled_visible_sources(
516 &mut sources,
517 &control,
518 |_| false,
519 |row| {
520 rows.push(row);
521 Ok(())
522 },
523 )
524 .unwrap();
525 assert_eq!(
526 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
527 vec![2, 3, 4]
528 );
529 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
530 }
531
532 #[test]
533 fn controlled_merge_uses_hlc_authority_when_epochs_inverted() {
534 use mongreldb_types::hlc::HlcTimestamp;
535 let hlc_old = HlcTimestamp {
536 physical_micros: 100,
537 logical: 0,
538 node_tiebreaker: 1,
539 };
540 let hlc_new = HlcTimestamp {
541 physical_micros: 200,
542 logical: 0,
543 node_tiebreaker: 1,
544 };
545 let a =
548 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
549 let b =
550 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
551 let control = crate::ExecutionControl::new(None);
552 let mut sources = vec![
553 ControlledVisibleSource::memory(a),
554 ControlledVisibleSource::memory(b),
555 ];
556 let mut rows = Vec::new();
557 merge_controlled_visible_sources(
558 &mut sources,
559 &control,
560 |_| false,
561 |row| {
562 rows.push(row);
563 Ok(())
564 },
565 )
566 .unwrap();
567 assert_eq!(rows.len(), 1);
568 assert_eq!(rows[0].row_id.0, 1);
569 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(11)));
571 assert_eq!(rows[0].commit_ts, Some(hlc_new));
572 }
573
574 #[test]
575 fn controlled_merge_epoch_wins_when_one_side_lacks_hlc() {
576 use mongreldb_types::hlc::HlcTimestamp;
577 let hlc_old = HlcTimestamp {
578 physical_micros: 100,
579 logical: 0,
580 node_tiebreaker: 1,
581 };
582 let a =
585 vec![Row::new_with_hlc(RowId(1), Epoch(10), hlc_old).with_column(1, Value::Int64(10))];
586 let b = vec![Row::new(RowId(1), Epoch(5)).with_column(1, Value::Int64(5))];
587 let control = crate::ExecutionControl::new(None);
588 let mut sources = vec![
589 ControlledVisibleSource::memory(a),
590 ControlledVisibleSource::memory(b),
591 ];
592 let mut rows = Vec::new();
593 merge_controlled_visible_sources(
594 &mut sources,
595 &control,
596 |_| false,
597 |row| {
598 rows.push(row);
599 Ok(())
600 },
601 )
602 .unwrap();
603 assert_eq!(rows.len(), 1);
604 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(10)));
605 }
606
607 #[test]
610 fn controlled_merge_memory_vs_run_uses_hlc_when_run_stamped() {
611 use crate::sorted_run::{RunReader, RunWriter};
612 use mongreldb_types::hlc::HlcTimestamp;
613 use tempfile::tempdir;
614
615 let hlc_old = HlcTimestamp {
616 physical_micros: 100,
617 logical: 0,
618 node_tiebreaker: 1,
619 };
620 let hlc_new = HlcTimestamp {
621 physical_micros: 200,
622 logical: 0,
623 node_tiebreaker: 1,
624 };
625 let schema = Schema {
626 schema_id: 1,
627 columns: vec![ColumnDef {
628 id: 1,
629 name: "v".into(),
630 ty: TypeId::Int64,
631 flags: ColumnFlags::empty(),
632 default_value: None,
633 embedding_source: None,
634 }],
635 indexes: vec![],
636 colocation: vec![],
637 constraints: Default::default(),
638 clustered: false,
639 };
640 let dir = tempdir().unwrap();
641 let path = dir.path().join("r-hlc.sr");
642 let run_rows =
644 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
645 RunWriter::new(&schema, 1, Epoch(50), 0)
646 .write(&path, &run_rows)
647 .unwrap();
648 let reader = RunReader::open(&path, schema, None).unwrap();
649 assert!(reader.has_column(crate::sorted_run::SYS_COMMIT_TS));
650
651 let mem =
653 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
654 let control = crate::ExecutionControl::new(None);
655 let mut sources = vec![
656 ControlledVisibleSource::memory(mem),
657 ControlledVisibleSource::run(
658 reader.into_visible_version_cursor(Epoch(u64::MAX)).unwrap(),
659 ),
660 ];
661 let mut rows = Vec::new();
662 merge_controlled_visible_sources(
663 &mut sources,
664 &control,
665 |_| false,
666 |row| {
667 rows.push(row);
668 Ok(())
669 },
670 )
671 .unwrap();
672 assert_eq!(rows.len(), 1);
673 assert_eq!(
674 rows[0].columns.get(&1),
675 Some(&Value::Int64(11)),
676 "HLC-newer memory version must beat epoch-newer run"
677 );
678 assert_eq!(rows[0].commit_ts, Some(hlc_new));
679 }
680
681 #[test]
682 fn controlled_memory_source_streams_large_overlays_without_full_active_vec() {
683 let mut map = BTreeMap::new();
684 for i in 1..=500u64 {
685 map.insert(
686 RowId(i),
687 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
688 );
689 }
690 let source = ControlledVisibleSource::memory_from_map(map);
691 assert!(
692 source.is_streaming_memory(),
693 "overlays larger than CONTROLLED_HOT_BATCH must use streaming cursor"
694 );
695 assert_eq!(source.streaming_total(), Some(500));
696 assert!(
698 source.streaming_peak_active().unwrap_or(usize::MAX) <= CONTROLLED_HOT_BATCH,
699 "peak active buffer must be <= CONTROLLED_HOT_BATCH"
700 );
701
702 let control = crate::ExecutionControl::new(None);
704 let mut sources = vec![ControlledVisibleSource::memory_from_map({
705 let mut m = BTreeMap::new();
706 for i in 1..=500u64 {
707 m.insert(
708 RowId(i),
709 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
710 );
711 }
712 m
713 })];
714 let mut n = 0usize;
715 merge_controlled_visible_sources(
716 &mut sources,
717 &control,
718 |_| false,
719 |_| {
720 n += 1;
721 Ok(())
722 },
723 )
724 .unwrap();
725 assert_eq!(n, 500);
726 assert!(
727 sources[0].streaming_peak_active().unwrap_or(0) <= CONTROLLED_HOT_BATCH,
728 "after full drain peak active still <= batch cap"
729 );
730 }
731}
732
733fn iso_now_bytes() -> Vec<u8> {
736 let secs = std::time::SystemTime::now()
737 .duration_since(std::time::UNIX_EPOCH)
738 .map(|d| d.as_secs() as i64)
739 .unwrap_or(0);
740 let days = secs.div_euclid(86_400);
741 let rem = secs.rem_euclid(86_400);
742 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
743 let (year, month, day) = civil_from_days(days);
744 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
745}
746
747pub(crate) fn unix_nanos_now() -> i64 {
748 std::time::SystemTime::now()
749 .duration_since(std::time::UNIX_EPOCH)
750 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
751 .unwrap_or(0)
752}
753
754fn ann_candidate_cap(
755 index_len: usize,
756 context: Option<&crate::query::AiExecutionContext>,
757) -> usize {
758 index_len
759 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
760 .min(context.map_or(
761 crate::query::MAX_RAW_INDEX_CANDIDATES,
762 crate::query::AiExecutionContext::max_fused_candidates,
763 ))
764}
765
766#[cfg(test)]
767mod ann_candidate_cap_tests {
768 use super::*;
769
770 #[test]
771 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
772 assert_eq!(
773 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
774 crate::query::MAX_RAW_INDEX_CANDIDATES,
775 );
776 let context = crate::query::AiExecutionContext::with_limits(
777 std::time::Duration::from_secs(1),
778 usize::MAX,
779 17,
780 );
781 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
782 }
783}
784
785fn civil_from_days(z: i64) -> (i64, u32, u32) {
786 let z = z + 719_468;
787 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
788 let doe = z - era * 146_097;
789 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
790 let y = yoe + era * 400;
791 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
792 let mp = (5 * doy + 2) / 153;
793 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
794 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
795 (if m <= 2 { y + 1 } else { y }, m, d)
796}
797
798pub fn clustered_row_id(primary_key: &Value) -> RowId {
805 let mut hash: u64 = 0xcbf29ce484222325;
806 for byte in primary_key.encode_key() {
807 hash ^= u64::from(byte);
808 hash = hash.wrapping_mul(0x100000001b3);
809 }
810 RowId(hash.max(1))
811}
812
813const 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;
820
821#[derive(Clone, Copy, Debug)]
836struct AutoIncState {
837 column_id: u16,
838 next: i64,
839 seeded: bool,
840}
841
842pub(crate) struct RecoveryMetadataPlan {
843 live_count: u64,
844 auto_inc: Option<AutoIncState>,
845 changed: bool,
846}
847
848type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
849
850fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
853 schema.auto_increment_column().map(|c| AutoIncState {
854 column_id: c.id,
855 next: 0,
856 seeded: false,
857 })
858}
859
860#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
873pub enum IndexBuildPolicy {
874 #[default]
876 Deferred,
877 Eager,
879}
880
881#[derive(Clone)]
882struct ReversePkSegment {
883 values: HashMap<RowId, Vec<u8>>,
884 removed: HashSet<RowId>,
885}
886
887#[derive(Clone)]
888struct ReversePkMap {
889 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
890 active: ReversePkSegment,
891}
892
893impl ReversePkMap {
894 fn new() -> Self {
895 Self {
896 frozen: Arc::new(Vec::new()),
897 active: ReversePkSegment {
898 values: HashMap::new(),
899 removed: HashSet::new(),
900 },
901 }
902 }
903
904 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
905 let mut map = Self::new();
906 map.active.values.extend(entries);
907 map
908 }
909
910 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
911 self.active.removed.remove(&row_id);
912 self.active.values.insert(row_id, key);
913 }
914
915 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
916 if let Some(key) = self.active.values.get(row_id) {
917 return Some(key);
918 }
919 if self.active.removed.contains(row_id) {
920 return None;
921 }
922 for segment in self.frozen.iter().rev() {
923 if let Some(key) = segment.values.get(row_id) {
924 return Some(key);
925 }
926 if segment.removed.contains(row_id) {
927 return None;
928 }
929 }
930 None
931 }
932
933 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
934 let previous = self.get(row_id).cloned();
935 self.active.values.remove(row_id);
936 self.active.removed.insert(*row_id);
937 previous
938 }
939
940 fn clear(&mut self) {
941 *self = Self::new();
942 }
943
944 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
945 let mut entries = HashMap::new();
946 for segment in self
947 .frozen
948 .iter()
949 .map(Arc::as_ref)
950 .chain(std::iter::once(&self.active))
951 {
952 for row_id in &segment.removed {
953 entries.remove(row_id);
954 }
955 entries.extend(
956 segment
957 .values
958 .iter()
959 .map(|(row_id, key)| (*row_id, key.clone())),
960 );
961 }
962 entries
963 }
964
965 fn seal(&mut self) {
966 if self.active.values.is_empty() && self.active.removed.is_empty() {
967 return;
968 }
969 let active = std::mem::replace(
970 &mut self.active,
971 ReversePkSegment {
972 values: HashMap::new(),
973 removed: HashSet::new(),
974 },
975 );
976 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
977 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
978 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
979 values: self.entries(),
980 removed: HashSet::new(),
981 })]);
982 }
983 }
984}
985
986#[derive(Clone)]
999pub struct ReadGeneration {
1000 schema: Arc<Schema>,
1001 base_runs: Arc<Vec<RunRef>>,
1002 deltas: TableDeltas,
1003 indexes: Arc<IndexGeneration>,
1004 visible_through: Epoch,
1005}
1006
1007pub(crate) enum SecondaryIndexArtifact {
1011 Bitmap(u16, BitmapIndex),
1012 LearnedRange(u16, ColumnLearnedRange),
1013 Fm(u16, Box<FmIndex>),
1014 Ann(u16, Box<AnnIndex>),
1015 Sparse(u16, SparseIndex),
1016 MinHash(u16, MinHashIndex),
1017}
1018
1019#[derive(Clone)]
1025pub struct TableDeltas {
1026 memtable: Memtable,
1027 mutable_run: MutableRun,
1028 hot: HotIndex,
1029 pk_by_row: ReversePkMap,
1030}
1031
1032impl ReadGeneration {
1033 fn empty(schema: &Schema) -> Self {
1036 Self {
1037 schema: Arc::new(schema.clone()),
1038 base_runs: Arc::new(Vec::new()),
1039 deltas: TableDeltas {
1040 memtable: Memtable::new(),
1041 mutable_run: MutableRun::new(),
1042 hot: HotIndex::new(),
1043 pk_by_row: ReversePkMap::new(),
1044 },
1045 indexes: Arc::new(IndexGeneration::default()),
1046 visible_through: Epoch(0),
1047 }
1048 }
1049
1050 pub fn schema(&self) -> &Arc<Schema> {
1052 &self.schema
1053 }
1054
1055 pub fn base_runs(&self) -> &[RunRef] {
1057 &self.base_runs
1058 }
1059
1060 pub fn indexes(&self) -> &Arc<IndexGeneration> {
1062 &self.indexes
1063 }
1064
1065 pub fn visible_through(&self) -> Epoch {
1067 self.visible_through
1068 }
1069
1070 pub fn deltas(&self) -> &TableDeltas {
1074 &self.deltas
1075 }
1076}
1077
1078impl TableDeltas {
1079 pub fn approx_bytes(&self) -> u64 {
1082 self.memtable
1083 .approx_bytes()
1084 .saturating_add(self.mutable_run.approx_bytes())
1085 }
1086
1087 pub fn memtable_len(&self) -> usize {
1089 self.memtable.len()
1090 }
1091
1092 pub fn mutable_run_len(&self) -> usize {
1094 self.mutable_run.len()
1095 }
1096
1097 pub fn hot_len(&self) -> usize {
1099 self.hot.len()
1100 }
1101
1102 pub fn reverse_pk_len(&self) -> usize {
1104 self.pk_by_row.entries().len()
1105 }
1106}
1107
1108#[derive(Clone)]
1110pub struct Table {
1111 dir: PathBuf,
1112 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1113 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
1114 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
1115 table_id: u64,
1116 name: String,
1120 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1125 read_only: bool,
1128 durable_commit_failed: bool,
1132 wal: WalSink,
1133 memtable: Memtable,
1134 mutable_run: MutableRun,
1139 mutable_run_spill_bytes: u64,
1141 compaction_zstd_level: i32,
1144 allocator: RowIdAllocator,
1145 epoch: Arc<EpochAuthority>,
1146 data_generation: u64,
1149 schema: Schema,
1150 hot: HotIndex,
1151 kek: Option<Arc<Kek>>,
1154 column_keys: HashMap<u16, ([u8; 32], u8)>,
1158 run_refs: Vec<RunRef>,
1159 retiring: Vec<crate::manifest::RetiredRun>,
1162 next_run_id: u64,
1163 sync_byte_threshold: u64,
1164 current_txn_id: u64,
1169 pending_private_mutations: bool,
1173 bitmap: HashMap<u16, BitmapIndex>,
1174 ann: HashMap<u16, AnnIndex>,
1175 fm: HashMap<u16, FmIndex>,
1176 sparse: HashMap<u16, SparseIndex>,
1177 minhash: HashMap<u16, MinHashIndex>,
1178 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
1181 pk_by_row: ReversePkMap,
1183 pinned: BTreeMap<Epoch, usize>,
1186 pub(crate) live_count: u64,
1189 reservoir: crate::reservoir::Reservoir,
1192 reservoir_complete: bool,
1200 had_deletes: bool,
1204 recent_delete_preimages: HashMap<Vec<u8>, Row>,
1209 run_row_id_ranges: HashMap<u128, (u64, u64)>,
1212 agg_cache: Arc<HashMap<u64, CachedAgg>>,
1216 global_idx_epoch: u64,
1220 indexes_complete: bool,
1225 index_build_policy: IndexBuildPolicy,
1227 pk_by_row_complete: bool,
1234 flushed_epoch: u64,
1237 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1240 snapshots: Arc<crate::retention::SnapshotRegistry>,
1243 commit_lock: Arc<parking_lot::Mutex<()>>,
1245 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1248 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
1258 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
1267 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1269 pending_delete_rids: roaring::RoaringBitmap,
1272 pending_put_cols: std::collections::HashSet<u16>,
1275 pending_rows: Vec<Row>,
1281 pending_rows_auto_inc: Vec<bool>,
1282 pending_dels: Vec<RowId>,
1285 pending_truncate: Option<Epoch>,
1289 auto_inc: Option<AutoIncState>,
1292 ttl: Option<TtlPolicy>,
1295 pins: Arc<crate::retention::PinRegistry>,
1302 published: Arc<ArcSwap<ReadGeneration>>,
1307 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
1312 pub(crate) lookup_metrics: LookupMetrics,
1317 run_lookup: RunLookupState,
1321}
1322
1323#[derive(Debug, Default, Clone)]
1331struct RunLookupState {
1332 directory: Option<std::sync::Arc<crate::run_lookup::RunLookupDirectory>>,
1333 fingerprint: u64,
1334 complete: bool,
1335}
1336
1337#[derive(Debug, Default)]
1338pub struct LookupMetrics {
1339 hot_lookup_hit: std::sync::atomic::AtomicU64,
1340 hot_lookup_fallback: std::sync::atomic::AtomicU64,
1341 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64,
1342 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64,
1343 result_cache_memory_hit: std::sync::atomic::AtomicU64,
1344 result_cache_disk_hit: std::sync::atomic::AtomicU64,
1345 result_cache_miss: std::sync::atomic::AtomicU64,
1346 result_cache_persistent_write_us: std::sync::atomic::AtomicU64,
1347 get_run_opened: std::sync::atomic::AtomicU64,
1349 get_run_skipped: std::sync::atomic::AtomicU64,
1350 pub(crate) directory_lookup_hit: std::sync::atomic::AtomicU64,
1352 pub(crate) directory_lookup_fallback: std::sync::atomic::AtomicU64,
1353 pub(crate) directory_incomplete: std::sync::atomic::AtomicU64,
1354 pub(crate) directory_run_readers_opened: std::sync::atomic::AtomicU64,
1355 pub(crate) directory_early_stop_total: std::sync::atomic::AtomicU64,
1356 pub(crate) directory_complete_miss_total: std::sync::atomic::AtomicU64,
1362 pub(crate) result_cache_persist_enqueued_total: std::sync::atomic::AtomicU64,
1364 pub(crate) result_cache_persist_coalesced_total: std::sync::atomic::AtomicU64,
1365 pub(crate) result_cache_persist_dropped_store_total: std::sync::atomic::AtomicU64,
1366 pub(crate) result_cache_persist_remove_total: std::sync::atomic::AtomicU64,
1367 pub(crate) result_cache_persist_stale_store_skipped_total: std::sync::atomic::AtomicU64,
1368 pub(crate) result_cache_persist_errors_total: std::sync::atomic::AtomicU64,
1369 pub(crate) result_cache_persist_shutdown_abandoned_total: std::sync::atomic::AtomicU64,
1370 pub(crate) result_cache_persist_queue_depth: std::sync::atomic::AtomicU64,
1371 pub(crate) hot_fallback_reasons: [std::sync::atomic::AtomicU64; 9],
1373 pub(crate) hot_fallback_overlay_versions_total: std::sync::atomic::AtomicU64,
1374 pub(crate) hot_fallback_runs_considered_total: std::sync::atomic::AtomicU64,
1375 pub(crate) hot_fallback_runs_opened_total: std::sync::atomic::AtomicU64,
1376 pub(crate) hot_fallback_pages_decoded_total: std::sync::atomic::AtomicU64,
1377 pub(crate) hot_fallback_rows_materialized_total: std::sync::atomic::AtomicU64,
1378 pub(crate) hot_lookup_duration_nanos: std::sync::atomic::AtomicU64,
1379 pub(crate) hot_fallback_duration_nanos: std::sync::atomic::AtomicU64,
1380 pub(crate) hot_mapping_rebuild_total: std::sync::atomic::AtomicU64,
1381 pub(crate) hot_checkpoint_rejected_total: std::sync::atomic::AtomicU64,
1382}
1383
1384impl Clone for LookupMetrics {
1385 fn clone(&self) -> Self {
1386 let mut hot_fallback_reasons = [
1387 std::sync::atomic::AtomicU64::new(0),
1388 std::sync::atomic::AtomicU64::new(0),
1389 std::sync::atomic::AtomicU64::new(0),
1390 std::sync::atomic::AtomicU64::new(0),
1391 std::sync::atomic::AtomicU64::new(0),
1392 std::sync::atomic::AtomicU64::new(0),
1393 std::sync::atomic::AtomicU64::new(0),
1394 std::sync::atomic::AtomicU64::new(0),
1395 std::sync::atomic::AtomicU64::new(0),
1396 ];
1397 for (dst, src) in hot_fallback_reasons
1398 .iter_mut()
1399 .zip(self.hot_fallback_reasons.iter())
1400 {
1401 dst.store(
1402 src.load(std::sync::atomic::Ordering::Relaxed),
1403 std::sync::atomic::Ordering::Relaxed,
1404 );
1405 }
1406 let copy_atomic = |src: &std::sync::atomic::AtomicU64| {
1407 std::sync::atomic::AtomicU64::new(src.load(std::sync::atomic::Ordering::Relaxed))
1408 };
1409 Self {
1410 hot_lookup_hit: copy_atomic(&self.hot_lookup_hit),
1411 hot_lookup_fallback: copy_atomic(&self.hot_lookup_fallback),
1412 hot_lookup_fallback_overlay_rows: copy_atomic(&self.hot_lookup_fallback_overlay_rows),
1413 hot_lookup_fallback_runs: copy_atomic(&self.hot_lookup_fallback_runs),
1414 result_cache_memory_hit: copy_atomic(&self.result_cache_memory_hit),
1415 result_cache_disk_hit: copy_atomic(&self.result_cache_disk_hit),
1416 result_cache_miss: copy_atomic(&self.result_cache_miss),
1417 result_cache_persistent_write_us: copy_atomic(&self.result_cache_persistent_write_us),
1418 get_run_opened: copy_atomic(&self.get_run_opened),
1419 get_run_skipped: copy_atomic(&self.get_run_skipped),
1420 directory_lookup_hit: copy_atomic(&self.directory_lookup_hit),
1421 directory_lookup_fallback: copy_atomic(&self.directory_lookup_fallback),
1422 directory_incomplete: copy_atomic(&self.directory_incomplete),
1423 directory_run_readers_opened: copy_atomic(&self.directory_run_readers_opened),
1424 directory_early_stop_total: copy_atomic(&self.directory_early_stop_total),
1425 directory_complete_miss_total: copy_atomic(&self.directory_complete_miss_total),
1426 result_cache_persist_enqueued_total: copy_atomic(
1427 &self.result_cache_persist_enqueued_total,
1428 ),
1429 result_cache_persist_coalesced_total: copy_atomic(
1430 &self.result_cache_persist_coalesced_total,
1431 ),
1432 result_cache_persist_dropped_store_total: copy_atomic(
1433 &self.result_cache_persist_dropped_store_total,
1434 ),
1435 result_cache_persist_remove_total: copy_atomic(&self.result_cache_persist_remove_total),
1436 result_cache_persist_stale_store_skipped_total: copy_atomic(
1437 &self.result_cache_persist_stale_store_skipped_total,
1438 ),
1439 result_cache_persist_errors_total: copy_atomic(&self.result_cache_persist_errors_total),
1440 result_cache_persist_shutdown_abandoned_total: copy_atomic(
1441 &self.result_cache_persist_shutdown_abandoned_total,
1442 ),
1443 result_cache_persist_queue_depth: copy_atomic(&self.result_cache_persist_queue_depth),
1444 hot_fallback_reasons,
1445 hot_fallback_overlay_versions_total: copy_atomic(
1446 &self.hot_fallback_overlay_versions_total,
1447 ),
1448 hot_fallback_runs_considered_total: copy_atomic(
1449 &self.hot_fallback_runs_considered_total,
1450 ),
1451 hot_fallback_runs_opened_total: copy_atomic(&self.hot_fallback_runs_opened_total),
1452 hot_fallback_pages_decoded_total: copy_atomic(&self.hot_fallback_pages_decoded_total),
1453 hot_fallback_rows_materialized_total: copy_atomic(
1454 &self.hot_fallback_rows_materialized_total,
1455 ),
1456 hot_lookup_duration_nanos: copy_atomic(&self.hot_lookup_duration_nanos),
1457 hot_fallback_duration_nanos: copy_atomic(&self.hot_fallback_duration_nanos),
1458 hot_mapping_rebuild_total: copy_atomic(&self.hot_mapping_rebuild_total),
1459 hot_checkpoint_rejected_total: copy_atomic(&self.hot_checkpoint_rejected_total),
1460 }
1461 }
1462}
1463
1464impl LookupMetrics {
1465 fn snapshot(&self) -> LookupMetricsSnapshot {
1466 LookupMetricsSnapshot {
1467 hot_lookup_hit: self
1468 .hot_lookup_hit
1469 .load(std::sync::atomic::Ordering::Relaxed),
1470 hot_lookup_fallback: self
1471 .hot_lookup_fallback
1472 .load(std::sync::atomic::Ordering::Relaxed),
1473 hot_lookup_fallback_overlay_rows: self
1474 .hot_lookup_fallback_overlay_rows
1475 .load(std::sync::atomic::Ordering::Relaxed),
1476 hot_lookup_fallback_runs: self
1477 .hot_lookup_fallback_runs
1478 .load(std::sync::atomic::Ordering::Relaxed),
1479 result_cache_memory_hit: self
1480 .result_cache_memory_hit
1481 .load(std::sync::atomic::Ordering::Relaxed),
1482 result_cache_disk_hit: self
1483 .result_cache_disk_hit
1484 .load(std::sync::atomic::Ordering::Relaxed),
1485 result_cache_miss: self
1486 .result_cache_miss
1487 .load(std::sync::atomic::Ordering::Relaxed),
1488 result_cache_persistent_write_us: self
1489 .result_cache_persistent_write_us
1490 .load(std::sync::atomic::Ordering::Relaxed),
1491 get_run_opened: self
1492 .get_run_opened
1493 .load(std::sync::atomic::Ordering::Relaxed),
1494 get_run_skipped: self
1495 .get_run_skipped
1496 .load(std::sync::atomic::Ordering::Relaxed),
1497 directory_lookup_hit: self
1498 .directory_lookup_hit
1499 .load(std::sync::atomic::Ordering::Relaxed),
1500 directory_lookup_fallback: self
1501 .directory_lookup_fallback
1502 .load(std::sync::atomic::Ordering::Relaxed),
1503 directory_incomplete: self
1504 .directory_incomplete
1505 .load(std::sync::atomic::Ordering::Relaxed),
1506 directory_run_readers_opened: self
1507 .directory_run_readers_opened
1508 .load(std::sync::atomic::Ordering::Relaxed),
1509 directory_early_stop_total: self
1510 .directory_early_stop_total
1511 .load(std::sync::atomic::Ordering::Relaxed),
1512 directory_complete_miss_total: self
1513 .directory_complete_miss_total
1514 .load(std::sync::atomic::Ordering::Relaxed),
1515 result_cache_persist_enqueued_total: self
1516 .result_cache_persist_enqueued_total
1517 .load(std::sync::atomic::Ordering::Relaxed),
1518 result_cache_persist_coalesced_total: self
1519 .result_cache_persist_coalesced_total
1520 .load(std::sync::atomic::Ordering::Relaxed),
1521 result_cache_persist_dropped_store_total: self
1522 .result_cache_persist_dropped_store_total
1523 .load(std::sync::atomic::Ordering::Relaxed),
1524 result_cache_persist_remove_total: self
1525 .result_cache_persist_remove_total
1526 .load(std::sync::atomic::Ordering::Relaxed),
1527 result_cache_persist_stale_store_skipped_total: self
1528 .result_cache_persist_stale_store_skipped_total
1529 .load(std::sync::atomic::Ordering::Relaxed),
1530 result_cache_persist_errors_total: self
1531 .result_cache_persist_errors_total
1532 .load(std::sync::atomic::Ordering::Relaxed),
1533 result_cache_persist_shutdown_abandoned_total: self
1534 .result_cache_persist_shutdown_abandoned_total
1535 .load(std::sync::atomic::Ordering::Relaxed),
1536 result_cache_persist_queue_depth: self
1537 .result_cache_persist_queue_depth
1538 .load(std::sync::atomic::Ordering::Relaxed),
1539 result_cache_persist_unavailable_total: 0,
1543 result_cache_persist_skipped_total: 0,
1544 result_cache_worker_spawn_failures_total: 0,
1545 result_cache_worker_shutdown_total: 0,
1546 hot_fallback_reasons: [
1547 self.hot_fallback_reasons[0].load(std::sync::atomic::Ordering::Relaxed),
1548 self.hot_fallback_reasons[1].load(std::sync::atomic::Ordering::Relaxed),
1549 self.hot_fallback_reasons[2].load(std::sync::atomic::Ordering::Relaxed),
1550 self.hot_fallback_reasons[3].load(std::sync::atomic::Ordering::Relaxed),
1551 self.hot_fallback_reasons[4].load(std::sync::atomic::Ordering::Relaxed),
1552 self.hot_fallback_reasons[5].load(std::sync::atomic::Ordering::Relaxed),
1553 self.hot_fallback_reasons[6].load(std::sync::atomic::Ordering::Relaxed),
1554 self.hot_fallback_reasons[7].load(std::sync::atomic::Ordering::Relaxed),
1555 self.hot_fallback_reasons[8].load(std::sync::atomic::Ordering::Relaxed),
1556 ],
1557 hot_fallback_overlay_versions_total: self
1558 .hot_fallback_overlay_versions_total
1559 .load(std::sync::atomic::Ordering::Relaxed),
1560 hot_fallback_runs_considered_total: self
1561 .hot_fallback_runs_considered_total
1562 .load(std::sync::atomic::Ordering::Relaxed),
1563 hot_fallback_runs_opened_total: self
1564 .hot_fallback_runs_opened_total
1565 .load(std::sync::atomic::Ordering::Relaxed),
1566 hot_fallback_pages_decoded_total: self
1567 .hot_fallback_pages_decoded_total
1568 .load(std::sync::atomic::Ordering::Relaxed),
1569 hot_fallback_rows_materialized_total: self
1570 .hot_fallback_rows_materialized_total
1571 .load(std::sync::atomic::Ordering::Relaxed),
1572 hot_lookup_duration_nanos: self
1573 .hot_lookup_duration_nanos
1574 .load(std::sync::atomic::Ordering::Relaxed),
1575 hot_fallback_duration_nanos: self
1576 .hot_fallback_duration_nanos
1577 .load(std::sync::atomic::Ordering::Relaxed),
1578 hot_mapping_rebuild_total: self
1579 .hot_mapping_rebuild_total
1580 .load(std::sync::atomic::Ordering::Relaxed),
1581 hot_checkpoint_rejected_total: self
1582 .hot_checkpoint_rejected_total
1583 .load(std::sync::atomic::Ordering::Relaxed),
1584 }
1585 }
1586}
1587
1588#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
1591pub struct LookupMetricsSnapshot {
1592 pub hot_lookup_hit: u64,
1593 pub hot_lookup_fallback: u64,
1594 pub hot_lookup_fallback_overlay_rows: u64,
1595 pub hot_lookup_fallback_runs: u64,
1596 pub result_cache_memory_hit: u64,
1597 pub result_cache_disk_hit: u64,
1598 pub result_cache_miss: u64,
1599 pub result_cache_persistent_write_us: u64,
1600 pub get_run_opened: u64,
1601 pub get_run_skipped: u64,
1602 pub directory_lookup_hit: u64,
1604 pub directory_lookup_fallback: u64,
1605 pub directory_incomplete: u64,
1606 pub directory_run_readers_opened: u64,
1607 pub directory_early_stop_total: u64,
1608 pub directory_complete_miss_total: u64,
1609 pub result_cache_persist_enqueued_total: u64,
1611 pub result_cache_persist_coalesced_total: u64,
1612 pub result_cache_persist_dropped_store_total: u64,
1613 pub result_cache_persist_remove_total: u64,
1614 pub result_cache_persist_stale_store_skipped_total: u64,
1615 pub result_cache_persist_errors_total: u64,
1616 pub result_cache_persist_shutdown_abandoned_total: u64,
1617 pub result_cache_persist_queue_depth: u64,
1618 pub result_cache_persist_unavailable_total: u64,
1622 pub result_cache_persist_skipped_total: u64,
1625 pub result_cache_worker_spawn_failures_total: u64,
1627 pub result_cache_worker_shutdown_total: u64,
1629 pub hot_fallback_reasons: [u64; 9],
1631 pub hot_fallback_overlay_versions_total: u64,
1632 pub hot_fallback_runs_considered_total: u64,
1633 pub hot_fallback_runs_opened_total: u64,
1634 pub hot_fallback_pages_decoded_total: u64,
1635 pub hot_fallback_rows_materialized_total: u64,
1636 pub hot_lookup_duration_nanos: u64,
1637 pub hot_fallback_duration_nanos: u64,
1638 pub hot_mapping_rebuild_total: u64,
1639 pub hot_checkpoint_rejected_total: u64,
1640}
1641
1642pub fn hot_fallback_reason_index(r: crate::trace::HotFallbackReason) -> usize {
1645 match r {
1646 crate::trace::HotFallbackReason::MissingMapping => 0,
1647 crate::trace::HotFallbackReason::StaleRowId => 1,
1648 crate::trace::HotFallbackReason::InvisibleAtSnapshot => 2,
1649 crate::trace::HotFallbackReason::HistoricalSnapshot => 3,
1650 crate::trace::HotFallbackReason::Tombstone => 4,
1651 crate::trace::HotFallbackReason::TtlExpired => 5,
1652 crate::trace::HotFallbackReason::PrimaryKeyMismatch => 6,
1653 crate::trace::HotFallbackReason::IndexIncomplete => 7,
1654 crate::trace::HotFallbackReason::CheckpointRejected => 8,
1655 }
1656}
1657
1658const _: () = {
1665 const fn assert_sync<T: ?Sized + Sync>() {}
1666 assert_sync::<Table>();
1667};
1668
1669#[derive(Clone)]
1675enum CachedData {
1676 Rows(Arc<Vec<Row>>),
1677 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
1678}
1679
1680impl CachedData {
1681 fn approx_bytes(&self) -> u64 {
1682 match self {
1683 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
1684 CachedData::Columns(c) => c
1685 .iter()
1686 .map(|(_, c)| c.approx_bytes())
1687 .sum::<u64>()
1688 .saturating_add(c.len() as u64 * 16),
1689 }
1690 }
1691}
1692
1693struct CachedEntry {
1697 data: CachedData,
1698 footprint: roaring::RoaringBitmap,
1699 condition_cols: Vec<u16>,
1700}
1701
1702impl CachedEntry {
1703 fn clone_for_persist(&self) -> Self {
1709 Self {
1710 data: self.data.clone(),
1711 footprint: self.footprint.clone(),
1712 condition_cols: self.condition_cols.clone(),
1713 }
1714 }
1715}
1716
1717struct ResultCache {
1730 entries: std::collections::HashMap<u64, CachedEntry>,
1731 order: BTreeSet<(u64, u64)>,
1732 generations: HashMap<u64, u64>,
1733 next_generation: u64,
1734 condition_index: HashMap<u16, HashSet<u64>>,
1737 empty_footprint_entries: HashSet<u64>,
1740 bytes: u64,
1741 max_bytes: u64,
1742 dir: Option<std::path::PathBuf>,
1743 #[allow(dead_code)]
1744 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1745 memory_hit: std::sync::atomic::AtomicU64,
1749 disk_hit: std::sync::atomic::AtomicU64,
1750 miss: std::sync::atomic::AtomicU64,
1751 persistent_write_us: std::sync::atomic::AtomicU64,
1752 persist_min_bytes: u64,
1757 persistence: PersistentPublicationState,
1764 worker_handle: Option<std::thread::JoinHandle<()>>,
1767 completion_rx: Option<std::sync::mpsc::Receiver<()>>,
1772 cache_payload_cipher: Option<std::sync::Arc<crate::encryption::AesCipher>>,
1776 persist_unavailable_total: std::sync::atomic::AtomicU64,
1780 persist_skipped_total: std::sync::atomic::AtomicU64,
1782 worker_spawn_failures_total: std::sync::atomic::AtomicU64,
1784 worker_shutdown_total: std::sync::atomic::AtomicU64,
1786}
1787
1788#[derive(serde::Serialize, serde::Deserialize)]
1790struct SerializedEntry {
1791 condition_cols: Vec<u16>,
1792 footprint_bits: Vec<u32>,
1793 data: SerializedData,
1794}
1795
1796#[derive(serde::Serialize, serde::Deserialize)]
1797enum SerializedData {
1798 Rows(Vec<Row>),
1799 Columns(Vec<(u16, columnar::NativeColumn)>),
1800}
1801
1802#[derive(serde::Serialize)]
1803struct SerializedEntryRef<'a> {
1804 condition_cols: &'a [u16],
1805 footprint_bits: Vec<u32>,
1806 data: SerializedDataRef<'a>,
1807}
1808
1809#[derive(serde::Serialize)]
1810enum SerializedDataRef<'a> {
1811 Rows(&'a [Row]),
1812 Columns(&'a [(u16, columnar::NativeColumn)]),
1813}
1814
1815impl<'a> SerializedEntryRef<'a> {
1816 fn from_entry(entry: &'a CachedEntry) -> Self {
1817 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
1818 let data = match &entry.data {
1819 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
1820 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
1821 };
1822 Self {
1823 condition_cols: &entry.condition_cols,
1824 footprint_bits,
1825 data,
1826 }
1827 }
1828}
1829
1830impl SerializedEntry {
1831 fn into_entry(self) -> Option<CachedEntry> {
1832 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1833 let data = match self.data {
1834 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1835 SerializedData::Columns(c) => {
1836 if !c.iter().all(|(_, col)| col.validate()) {
1839 return None;
1840 }
1841 CachedData::Columns(Arc::new(c))
1842 }
1843 };
1844 Some(CachedEntry {
1845 data,
1846 footprint,
1847 condition_cols: self.condition_cols,
1848 })
1849 }
1850}
1851
1852impl ResultCache {
1853 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1854
1855 fn new() -> Self {
1856 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1857 }
1858
1859 fn with_max_bytes(max_bytes: u64) -> Self {
1860 Self {
1861 entries: std::collections::HashMap::new(),
1862 order: BTreeSet::new(),
1863 generations: HashMap::new(),
1864 next_generation: 0,
1865 condition_index: HashMap::new(),
1866 empty_footprint_entries: HashSet::new(),
1867 bytes: 0,
1868 max_bytes,
1869 dir: None,
1870 cache_dek: None,
1871 memory_hit: std::sync::atomic::AtomicU64::new(0),
1872 disk_hit: std::sync::atomic::AtomicU64::new(0),
1873 miss: std::sync::atomic::AtomicU64::new(0),
1874 persistent_write_us: std::sync::atomic::AtomicU64::new(0),
1875 persist_min_bytes: 4 * 1024,
1878 persistence: PersistentPublicationState::Disabled(
1879 PersistenceDisabledReason::NoDirectory,
1880 ),
1881 worker_handle: None,
1882 completion_rx: None,
1883 cache_payload_cipher: None,
1884 persist_unavailable_total: std::sync::atomic::AtomicU64::new(0),
1885 persist_skipped_total: std::sync::atomic::AtomicU64::new(0),
1886 worker_spawn_failures_total: std::sync::atomic::AtomicU64::new(0),
1887 worker_shutdown_total: std::sync::atomic::AtomicU64::new(0),
1888 }
1889 }
1890
1891 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1892 let _ = std::fs::create_dir_all(&dir);
1893 self.dir = Some(dir);
1894 if matches!(
1895 self.persistence,
1896 PersistentPublicationState::Disabled(PersistenceDisabledReason::NoDirectory)
1897 ) {
1898 self.persistence =
1899 PersistentPublicationState::Disabled(PersistenceDisabledReason::QueueUnavailable);
1900 }
1901 self
1902 }
1903
1904 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1905 self.cache_dek = dek.clone();
1911 if let Some(dek) = dek {
1912 if let Ok(cipher) = crate::encryption::AesCipher::new(&dek[..]) {
1913 self.cache_payload_cipher = Some(std::sync::Arc::new(cipher));
1914 }
1915 }
1916 self
1917 }
1918
1919 fn install_persistent_writer(
1924 &mut self,
1925 writer: std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
1926 worker_handle: std::thread::JoinHandle<()>,
1927 completion_rx: std::sync::mpsc::Receiver<()>,
1928 ) {
1929 self.persistence = PersistentPublicationState::Async(writer);
1930 self.worker_handle = Some(worker_handle);
1931 self.completion_rx = Some(completion_rx);
1932 }
1933
1934 fn note_worker_spawn_failure(&mut self) {
1939 self.worker_spawn_failures_total
1940 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1941 self.persist_unavailable_total
1942 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1943 self.persistence =
1944 PersistentPublicationState::Disabled(PersistenceDisabledReason::WorkerSpawnFailed);
1945 }
1946
1947 fn take_persistent_worker(&mut self) -> Option<std::thread::JoinHandle<()>> {
1950 self.worker_handle.take()
1951 }
1952
1953 fn take_persistent_completion(&mut self) -> Option<std::sync::mpsc::Receiver<()>> {
1956 self.completion_rx.take()
1957 }
1958
1959 fn persistent_writer(
1962 &self,
1963 ) -> Option<&std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>> {
1964 match &self.persistence {
1965 PersistentPublicationState::Async(writer) => Some(writer),
1966 PersistentPublicationState::Disabled(_) => None,
1967 }
1968 }
1969
1970 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1971 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1972 }
1973
1974 fn persist_entry_synchronously_for_maintenance(
1986 &self,
1987 context: PersistContext,
1988 entry: &CachedEntry,
1989 ) {
1990 let Some(path) = self.disk_path(context.key) else {
1991 return;
1992 };
1993 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1994 Ok(s) => s,
1995 Err(_) => return,
1996 };
1997 let bytes = match crate::result_cache::encode_persisted_entry(
1998 context,
1999 &serialized,
2000 self.cache_payload_cipher.as_deref(),
2001 ) {
2002 Ok(b) => b,
2003 Err(_) => return,
2004 };
2005 let tmp = path.with_extension("tmp");
2006 use std::io::Write;
2007 let write = || -> std::io::Result<()> {
2008 let mut f = std::fs::File::create(&tmp)?;
2009 f.write_all(&bytes)?;
2010 f.flush()?;
2011 f.sync_all()?;
2012 Ok(())
2013 };
2014 if write().is_err() {
2015 let _ = std::fs::remove_file(&tmp);
2016 return;
2017 }
2018 if std::fs::rename(&tmp, &path).is_err() {
2019 let _ = std::fs::remove_file(&tmp);
2020 return;
2021 }
2022 if let Ok(dir) = std::fs::File::open(self.dir.as_ref().expect("dir set").clone()) {
2023 let _ = dir.sync_all();
2024 }
2025 }
2026
2027 fn load_from_disk(&self, key: u64, identity: PersistentCacheIdentity) -> Option<CachedEntry> {
2033 let path = self.disk_path(key)?;
2034 let bytes = std::fs::read(&path).ok()?;
2035 if bytes.len() < FRAME_MAGIC.len() || bytes[..FRAME_MAGIC.len()] != FRAME_MAGIC {
2039 let _ = std::fs::remove_file(&path);
2040 return None;
2041 }
2042 let plaintext = match crate::result_cache::decode_persisted_entry(
2043 identity,
2044 key,
2045 &bytes,
2046 self.cache_payload_cipher.as_deref(),
2047 ) {
2048 Ok(p) => p,
2049 Err(_) => {
2050 let _ = std::fs::remove_file(&path);
2051 return None;
2052 }
2053 };
2054 let serialized: SerializedEntry = match bincode::deserialize(&plaintext) {
2055 Ok(s) => s,
2056 Err(_) => {
2057 let _ = std::fs::remove_file(&path);
2058 return None;
2059 }
2060 };
2061 serialized.into_entry()
2062 }
2063
2064 fn remove_from_disk(&self, key: u64) {
2066 if let Some(path) = self.disk_path(key) {
2067 let _ = std::fs::remove_file(&path);
2068 }
2069 }
2070
2071 fn load_persistent(&mut self, identity: PersistentCacheIdentity) {
2075 let Some(dir) = self.dir.as_ref().cloned() else {
2076 return;
2077 };
2078 let entries = match std::fs::read_dir(&dir) {
2079 Ok(e) => e,
2080 Err(_) => return,
2081 };
2082 for entry in entries.flatten() {
2083 let path = entry.path();
2084 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
2086 let _ = std::fs::remove_file(&path);
2087 continue;
2088 }
2089 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
2090 continue;
2091 }
2092 let stem = match path.file_stem().and_then(|s| s.to_str()) {
2093 Some(s) => s,
2094 None => continue,
2095 };
2096 let key = match u64::from_str_radix(stem, 16) {
2097 Ok(k) => k,
2098 Err(_) => continue,
2099 };
2100 if let Some(entry) = self.load_from_disk(key, identity) {
2104 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
2105 self.index_entry(key, &entry);
2106 self.entries.insert(key, entry);
2107 self.touch(key);
2108 }
2109 }
2110 self.evict();
2111 }
2112
2113 fn set_max_bytes(&mut self, max_bytes: u64) {
2114 self.max_bytes = max_bytes;
2115 self.evict();
2116 }
2117
2118 fn touch(&mut self, key: u64) {
2121 if !self.entries.contains_key(&key) {
2122 return;
2123 }
2124 if let Some(previous) = self.generations.remove(&key) {
2125 self.order.remove(&(previous, key));
2126 }
2127 let generation = self.allocate_generation();
2128 self.generations.insert(key, generation);
2129 self.order.insert((generation, key));
2130 }
2131
2132 fn untrack(&mut self, key: u64) {
2133 if let Some(generation) = self.generations.remove(&key) {
2134 self.order.remove(&(generation, key));
2135 }
2136 }
2137
2138 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
2139 for column in &entry.condition_cols {
2140 self.condition_index.entry(*column).or_default().insert(key);
2141 }
2142 if entry.footprint.is_empty() {
2143 self.empty_footprint_entries.insert(key);
2144 }
2145 }
2146
2147 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
2148 for column in &entry.condition_cols {
2149 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
2150 keys.remove(&key);
2151 keys.is_empty()
2152 });
2153 if remove_column {
2154 self.condition_index.remove(column);
2155 }
2156 }
2157 if entry.footprint.is_empty() {
2158 self.empty_footprint_entries.remove(&key);
2159 }
2160 }
2161
2162 fn allocate_generation(&mut self) -> u64 {
2163 if self.next_generation == u64::MAX {
2164 self.rebase_generations();
2165 }
2166 let generation = self.next_generation;
2167 self.next_generation += 1;
2168 generation
2169 }
2170
2171 fn rebase_generations(&mut self) {
2172 let ordered = std::mem::take(&mut self.order);
2173 self.generations.clear();
2174 for (generation, (_, key)) in ordered.into_iter().enumerate() {
2175 let generation = generation as u64;
2176 self.generations.insert(key, generation);
2177 self.order.insert((generation, key));
2178 }
2179 self.next_generation = self.order.len() as u64;
2180 }
2181
2182 fn get_rows(&mut self, key: u64, identity: PersistentCacheIdentity) -> Option<Arc<Vec<Row>>> {
2183 let res = self.entries.get(&key).and_then(|e| match &e.data {
2184 CachedData::Rows(r) => Some(r.clone()),
2185 CachedData::Columns(_) => None,
2186 });
2187 if res.is_some() {
2188 self.touch(key);
2189 self.memory_hit
2190 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2191 return res;
2192 }
2193 if let Some(entry) = self.load_from_disk(key, identity) {
2195 let res = match &entry.data {
2196 CachedData::Rows(r) => Some(r.clone()),
2197 CachedData::Columns(_) => None,
2198 };
2199 if res.is_some() {
2200 let approx = entry.data.approx_bytes();
2201 self.bytes = self.bytes.saturating_add(approx);
2202 self.index_entry(key, &entry);
2203 self.entries.insert(key, entry);
2204 self.touch(key);
2205 self.evict();
2206 self.disk_hit
2207 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2208 return res;
2209 }
2210 }
2211 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2212 None
2213 }
2214
2215 fn get_columns(
2216 &mut self,
2217 key: u64,
2218 identity: PersistentCacheIdentity,
2219 ) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
2220 let res = self.entries.get(&key).and_then(|e| match &e.data {
2221 CachedData::Columns(c) => Some(c.clone()),
2222 CachedData::Rows(_) => None,
2223 });
2224 if res.is_some() {
2225 self.touch(key);
2226 self.memory_hit
2227 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2228 return res;
2229 }
2230 if let Some(entry) = self.load_from_disk(key, identity) {
2232 let res = match &entry.data {
2233 CachedData::Columns(c) => Some(c.clone()),
2234 CachedData::Rows(_) => None,
2235 };
2236 if res.is_some() {
2237 let approx = entry.data.approx_bytes();
2238 self.bytes = self.bytes.saturating_add(approx);
2239 self.index_entry(key, &entry);
2240 self.entries.insert(key, entry);
2241 self.touch(key);
2242 self.evict();
2243 self.disk_hit
2244 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2245 return res;
2246 }
2247 }
2248 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2249 None
2250 }
2251
2252 fn insert(&mut self, key: u64, entry: CachedEntry) {
2257 let approx = entry.data.approx_bytes();
2258 if let Some(previous) = self.entries.remove(&key) {
2259 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
2260 self.unindex_entry(key, &previous);
2261 self.untrack(key);
2262 }
2263 self.bytes = self.bytes.saturating_add(approx);
2264 self.index_entry(key, &entry);
2265 self.entries.insert(key, entry);
2266 self.touch(key);
2267 self.evict();
2268 }
2269
2270 fn persist_entry(&self, context: PersistContext, entry: &CachedEntry) {
2276 let approx = entry.data.approx_bytes();
2277 if self.dir.is_none() {
2278 return;
2279 }
2280 if self.persist_min_bytes > 0 && approx < self.persist_min_bytes {
2281 return;
2282 }
2283 match &self.persistence {
2284 PersistentPublicationState::Async(writer) => {
2285 let entry_clone = entry.clone_for_persist();
2289 let payload_factory: Box<dyn FnOnce() -> Option<Vec<u8>> + Send + 'static> =
2290 Box::new(move || {
2291 bincode::serialize(&SerializedEntryRef::from_entry(&entry_clone)).ok()
2292 });
2293 let persistable = crate::result_cache::PersistableEntry {
2294 key: context.key,
2295 table_id: context.identity.table_id,
2296 schema_id: context.identity.schema_id,
2297 run_generation: context.identity.logical_generation,
2298 entry_generation: context.entry_generation,
2299 bytes: approx as usize,
2300 payload_factory,
2301 };
2302 let write_start = std::time::Instant::now();
2303 writer.enqueue_store(persistable);
2304 let write_us = write_start.elapsed().as_micros() as u64;
2305 self.persistent_write_us
2306 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2307 }
2308 PersistentPublicationState::Disabled(reason) => {
2309 self.persist_skipped_total
2314 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2315 crate::trace::QueryTrace::record(|trace| {
2316 trace.result_cache_persist_skipped = true;
2317 trace.result_cache_persist_skip_reason = Some(reason.label());
2318 });
2319 }
2320 }
2321 }
2322
2323 fn allocate_persist_generation(&mut self, key: u64) -> u64 {
2328 if let Some(writer) = self.persistent_writer() {
2329 let _ = key; writer.bump_persist_generation(key)
2333 } else {
2334 0
2335 }
2336 }
2337
2338 fn enqueue_persist_remove(&self, key: u64) {
2341 if let Some(writer) = self.persistent_writer() {
2342 writer.enqueue_remove(key);
2343 }
2344 }
2345
2346 fn enqueue_persist_clear(&self) {
2353 if let Some(writer) = self.persistent_writer() {
2354 writer.enqueue_clear();
2355 }
2356 }
2357
2358 fn flush_persistent_cache(&self, deadline: std::time::Duration) -> u64 {
2364 let Some(writer) = self.persistent_writer() else {
2365 return 0;
2366 };
2367 let start = std::time::Instant::now();
2368 loop {
2369 let depth = writer.queue_depth() as u64;
2370 let in_flight = writer.writes_in_flight();
2376 if depth == 0 && in_flight == 0 {
2377 return 0;
2378 }
2379 if start.elapsed() >= deadline {
2380 return depth;
2381 }
2382 std::thread::sleep(std::time::Duration::from_millis(2));
2383 }
2384 }
2385
2386 fn shutdown_persistent_cache(&mut self, deadline: std::time::Duration) {
2396 let start = Instant::now();
2397 if let Some(writer) = self.persistent_writer() {
2399 loop {
2400 if writer.queue_depth() == 0 {
2401 break;
2402 }
2403 if start.elapsed() >= deadline {
2404 break;
2405 }
2406 std::thread::sleep(std::time::Duration::from_millis(2));
2407 }
2408 }
2409 let writer = self.persistent_writer().cloned();
2413 if let Some(writer) = writer.as_ref() {
2414 writer.shutdown();
2415 let _ = writer.drain_all_as_abandoned();
2416 }
2417 let mut completion = self.take_persistent_completion();
2424 let completed = if let (Some(rx), Some(_writer)) = (completion.as_mut(), writer.as_ref()) {
2425 let elapsed = start.elapsed();
2426 let remaining = deadline.saturating_sub(elapsed);
2427 crate::result_cache::recv_completion_with_timeout(rx, remaining)
2428 } else {
2429 true
2430 };
2431 if completed {
2432 if let Some(handle) = self.take_persistent_worker() {
2433 let _ = handle.join();
2434 }
2435 drop(completion.take());
2437 } else {
2438 if let Some(writer) = writer.as_ref() {
2442 writer.record_outcome(DrainOutcome::Abandoned);
2443 }
2444 self.take_persistent_worker(); }
2446 if matches!(self.persistence, PersistentPublicationState::Async(_)) {
2450 self.worker_shutdown_total
2451 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2452 self.persist_unavailable_total
2453 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2454 self.persistence =
2455 PersistentPublicationState::Disabled(PersistenceDisabledReason::WorkerShutdown);
2456 }
2457 }
2458
2459 pub(crate) fn set_persist_min_bytes(&mut self, min: u64) {
2463 self.persist_min_bytes = min;
2464 }
2465
2466 #[cfg(test)]
2468 fn would_persist(&self, approx: u64) -> bool {
2469 self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
2470 }
2471
2472 fn cache_counters(&self) -> (u64, u64, u64, u64) {
2475 (
2476 self.memory_hit.load(std::sync::atomic::Ordering::Relaxed),
2477 self.disk_hit.load(std::sync::atomic::Ordering::Relaxed),
2478 self.miss.load(std::sync::atomic::Ordering::Relaxed),
2479 self.persistent_write_us
2480 .load(std::sync::atomic::Ordering::Relaxed),
2481 )
2482 }
2483
2484 fn publication_counters(&self) -> (u64, u64, u64, u64) {
2488 (
2489 self.persist_unavailable_total
2490 .load(std::sync::atomic::Ordering::Relaxed),
2491 self.persist_skipped_total
2492 .load(std::sync::atomic::Ordering::Relaxed),
2493 self.worker_spawn_failures_total
2494 .load(std::sync::atomic::Ordering::Relaxed),
2495 self.worker_shutdown_total
2496 .load(std::sync::atomic::Ordering::Relaxed),
2497 )
2498 }
2499
2500 fn persist_snapshot(&self) -> Option<LookupMetricsSnapshot> {
2503 self.persistent_writer().map(|w| w.persist_snapshot())
2504 }
2505
2506 fn invalidate(
2515 &mut self,
2516 delete_rids: &roaring::RoaringBitmap,
2517 put_cols: &std::collections::HashSet<u16>,
2518 ) {
2519 if self.entries.is_empty() {
2520 return;
2521 }
2522 let has_deletes = !delete_rids.is_empty();
2523 let mut to_remove = HashSet::new();
2524
2525 for column in put_cols {
2529 if let Some(keys) = self.condition_index.get(column) {
2530 to_remove.extend(keys.iter().copied());
2531 }
2532 }
2533
2534 if has_deletes {
2535 to_remove.extend(self.empty_footprint_entries.iter().copied());
2539 for (&key, entry) in &self.entries {
2540 if !to_remove.contains(&key)
2541 && !entry.footprint.is_empty()
2542 && entry.footprint.intersection_len(delete_rids) > 0
2543 {
2544 to_remove.insert(key);
2545 }
2546 }
2547 }
2548
2549 for key in to_remove {
2550 if let Some(entry) = self.entries.remove(&key) {
2551 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2552 self.unindex_entry(key, &entry);
2553 }
2554 if self.persistent_writer().is_some() {
2559 self.enqueue_persist_remove(key);
2560 } else {
2561 self.remove_from_disk(key);
2562 }
2563 self.untrack(key);
2564 }
2565 }
2566
2567 fn clear(&mut self) {
2568 if self.dir.is_some() {
2572 if self.persistent_writer().is_some() {
2573 self.enqueue_persist_clear();
2574 } else if let Some(dir) = &self.dir {
2575 if let Ok(entries) = std::fs::read_dir(dir) {
2576 for entry in entries.flatten() {
2577 let path = entry.path();
2578 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
2579 let _ = std::fs::remove_file(&path);
2580 }
2581 }
2582 }
2583 }
2584 }
2585 self.entries.clear();
2586 self.order.clear();
2587 self.generations.clear();
2588 self.next_generation = 0;
2589 self.condition_index.clear();
2590 self.empty_footprint_entries.clear();
2591 self.bytes = 0;
2592 }
2593
2594 fn evict(&mut self) {
2595 while self.bytes > self.max_bytes {
2596 let Some((_, key)) = self.order.pop_first() else {
2597 break;
2598 };
2599 self.generations.remove(&key);
2600 if let Some(entry) = self.entries.remove(&key) {
2601 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2602 self.unindex_entry(key, &entry);
2603 }
2611 }
2612 }
2613}
2614
2615fn spawn_persistent_cache_worker(
2623 dir: std::path::PathBuf,
2624 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
2625 metrics: LookupMetrics,
2626) -> std::io::Result<(
2627 std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
2628 std::thread::JoinHandle<()>,
2629 std::sync::mpsc::Receiver<()>,
2630)> {
2631 use crate::result_cache::{
2632 spawn_persistent_cache_worker as spawn, RealPersistentCacheIo, StalenessGuard,
2633 WorkerConfig, WriterStalenessGuard,
2634 };
2635 if PERSISTENT_WORKER_SPAWN_FAILURE.with(|flag| flag.get()) {
2639 return Err(std::io::Error::other(
2640 "persistent-cache worker spawn failure forced by test seam",
2641 ));
2642 }
2643 let io: std::sync::Arc<dyn crate::result_cache::PersistentCacheIo> =
2644 std::sync::Arc::new(RealPersistentCacheIo::new(dir)?);
2645 let cipher = match cache_dek {
2646 Some(dek) => {
2647 let cipher = crate::encryption::AesCipher::new(&dek[..])
2648 .map_err(|e| std::io::Error::other(format!("aes: {e}")))?;
2649 Some(std::sync::Arc::new(cipher))
2650 }
2651 None => None,
2652 };
2653 let writer = std::sync::Arc::new(crate::result_cache::PersistentResultCacheWriter::new(
2654 metrics,
2655 crate::result_cache::WriterLimits::default(),
2656 ));
2657 let staleness: std::sync::Arc<dyn StalenessGuard> =
2658 std::sync::Arc::new(WriterStalenessGuard::new(writer.clone()));
2659 let (completion_tx, completion_rx) = std::sync::mpsc::channel();
2660 let config = WorkerConfig {
2661 writer: writer.clone(),
2662 io,
2663 cipher,
2664 staleness,
2665 max_staleness_retries: 8,
2666 completion: Some(completion_tx),
2667 };
2668 let handle = spawn(config);
2669 Ok((writer, handle, completion_rx))
2670}
2671
2672thread_local! {
2673 static PERSISTENT_WORKER_SPAWN_FAILURE: std::cell::Cell<bool> =
2680 const { std::cell::Cell::new(false) };
2681}
2682
2683#[cfg(test)]
2684mod result_cache_lru_tests {
2685 use super::*;
2686
2687 fn row_entry(row_id: u64) -> CachedEntry {
2688 CachedEntry {
2689 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
2690 footprint: std::iter::once(row_id as u32).collect(),
2691 condition_cols: vec![1],
2692 }
2693 }
2694
2695 #[test]
2696 fn hits_update_recency_without_growing_an_order_queue() {
2697 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2698 let identity = PersistentCacheIdentity {
2699 table_id: 0,
2700 schema_id: 0,
2701 logical_generation: 0,
2702 };
2703 cache.insert(1, row_entry(1));
2704 cache.insert(2, row_entry(2));
2705 assert_eq!(cache.order.len(), cache.entries.len());
2706
2707 for _ in 0..1_000 {
2708 assert!(cache.get_rows(1, identity).is_some());
2709 }
2710
2711 assert_eq!(cache.order.len(), cache.entries.len());
2712 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
2713
2714 let one_entry = cache.entries[&1].data.approx_bytes();
2715 cache.set_max_bytes(one_entry);
2716 assert!(cache.entries.contains_key(&1));
2717 assert!(!cache.entries.contains_key(&2));
2718 assert_eq!(cache.order.len(), cache.entries.len());
2719 }
2720
2721 #[test]
2722 fn tiny_entries_skip_persistent_tier_by_default() {
2723 let dir = tempfile::tempdir().unwrap();
2724 let mut cache = ResultCache::with_max_bytes(u64::MAX).with_dir(dir.path().to_path_buf());
2725 let tiny = row_entry(1).data.approx_bytes();
2726 assert!(
2727 tiny < 4 * 1024,
2728 "row_entry fixture must be below default 4KiB threshold"
2729 );
2730 assert!(
2731 !cache.would_persist(tiny),
2732 "tiny entry must not force synchronous disk publish"
2733 );
2734 assert!(
2735 cache.would_persist(8 * 1024),
2736 "large entry must still persist when dir is set"
2737 );
2738 cache.set_persist_min_bytes(0);
2739 assert!(
2740 cache.would_persist(tiny),
2741 "persist_min_bytes=0 restores always-persist policy"
2742 );
2743 }
2744
2745 #[test]
2746 fn insert_invalidation_uses_the_condition_reverse_index() {
2747 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2748 let mut first = row_entry(1);
2749 first.condition_cols = vec![1];
2750 let mut second = row_entry(2);
2751 second.condition_cols = vec![2];
2752 cache.insert(1, first);
2753 cache.insert(2, second);
2754
2755 let put_cols = std::iter::once(1_u16).collect();
2756 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
2757
2758 assert!(!cache.entries.contains_key(&1));
2759 assert!(cache.entries.contains_key(&2));
2760 assert!(!cache
2761 .condition_index
2762 .get(&1)
2763 .is_some_and(|keys| keys.contains(&1)));
2764 assert!(cache
2765 .condition_index
2766 .get(&2)
2767 .is_some_and(|keys| keys.contains(&2)));
2768 assert_eq!(cache.order.len(), cache.entries.len());
2769 }
2770
2771 #[test]
2772 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
2773 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2774 cache.insert(1, row_entry(1));
2775 cache.insert(2, row_entry(2));
2776 let mut unknown = row_entry(3);
2777 unknown.footprint.clear();
2778 cache.insert(3, unknown);
2779
2780 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
2781 cache.invalidate(&delete_rids, &HashSet::new());
2782
2783 assert!(!cache.entries.contains_key(&1));
2784 assert!(cache.entries.contains_key(&2));
2785 assert!(!cache.entries.contains_key(&3));
2786 assert!(cache.empty_footprint_entries.is_empty());
2787 assert_eq!(cache.order.len(), cache.entries.len());
2788 }
2789
2790 #[test]
2791 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
2792 let entry = row_entry(7);
2793 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
2794 let owned = bincode::serialize(&SerializedEntry {
2795 condition_cols: entry.condition_cols.clone(),
2796 footprint_bits: entry.footprint.iter().collect(),
2797 data: match &entry.data {
2798 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
2799 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
2800 },
2801 })
2802 .unwrap();
2803 assert_eq!(borrowed, owned);
2804 }
2805}
2806
2807type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
2814
2815fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
2816 let _ = kek;
2817 {
2818 if let Some(k) = kek {
2819 return (
2820 Some(k.derive_table_wal_key(_table_id)),
2821 Some(k.derive_cache_key()),
2822 );
2823 }
2824 }
2825 (None, None)
2826}
2827
2828fn read_table_encryption_salt_root(
2829 root: &crate::durable_file::DurableRoot,
2830) -> Result<[u8; crate::encryption::SALT_LEN]> {
2831 use std::io::Read;
2832
2833 let mut file = root
2834 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
2835 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
2836 let length = file.metadata()?.len();
2837 if length != crate::encryption::SALT_LEN as u64 {
2838 return Err(MongrelError::InvalidArgument(format!(
2839 "salt file is {length} bytes, expected {}",
2840 crate::encryption::SALT_LEN
2841 )));
2842 }
2843 let mut salt = [0_u8; crate::encryption::SALT_LEN];
2844 file.read_exact(&mut salt)?;
2845 Ok(salt)
2846}
2847
2848fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
2850 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
2851}
2852
2853fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
2854 let Some(kek) = kek else {
2855 return HashMap::new();
2856 };
2857 {
2858 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
2859 schema
2860 .columns
2861 .iter()
2862 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
2863 .map(|c| {
2864 let scheme = if schema
2865 .indexes
2866 .iter()
2867 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
2868 {
2869 SCHEME_OPE_RANGE
2870 } else {
2871 SCHEME_HMAC_EQ
2872 };
2873 let key: [u8; 32] = *kek.derive_column_key(c.id);
2874 (c.id, (key, scheme))
2875 })
2876 .collect()
2877 }
2878}
2879
2880pub(crate) struct SharedCtx {
2885 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
2886 pub epoch: Arc<EpochAuthority>,
2887 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
2888 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
2889 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
2890 pub kek: Option<Arc<Kek>>,
2891 pub commit_lock: Arc<parking_lot::Mutex<()>>,
2897 pub shared: Option<SharedWalCtx>,
2901 pub table_name: Option<String>,
2904 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
2907 pub read_only: bool,
2909}
2910
2911#[derive(Clone)]
2917pub(crate) struct SharedWalCtx {
2918 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
2919 pub group: Arc<GroupCommit>,
2920 pub poisoned: Arc<AtomicBool>,
2921 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
2922 pub change_wake: tokio::sync::broadcast::Sender<()>,
2923 pub lifecycle: Arc<crate::core::LifecycleController>,
2926 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
2928}
2929
2930enum WalSink {
2933 Private(Wal),
2934 Shared(SharedWalCtx),
2935 ReadOnly,
2936}
2937
2938impl Clone for WalSink {
2939 fn clone(&self) -> Self {
2940 match self {
2941 Self::Shared(shared) => Self::Shared(shared.clone()),
2942 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
2943 }
2944 }
2945}
2946
2947impl SharedCtx {
2948 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
2952 let n_shards = if cache_dir.is_some() {
2956 1
2957 } else {
2958 crate::cache::CACHE_SHARDS
2959 };
2960 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
2961 let page_cache = if let Some(d) = cache_dir {
2962 Arc::new(crate::cache::Sharded::new(1, || {
2963 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
2964 }))
2965 } else {
2966 Arc::new(crate::cache::Sharded::new(n_shards, || {
2967 crate::cache::PageCache::new(per_shard)
2968 }))
2969 };
2970 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
2971 let decoded_cache = Arc::new(crate::cache::Sharded::new(
2972 crate::cache::CACHE_SHARDS,
2973 || crate::cache::DecodedPageCache::new(decoded_per_shard),
2974 ));
2975 Self {
2976 root_guard: None,
2977 epoch: Arc::new(EpochAuthority::new(0)),
2978 page_cache,
2979 decoded_cache,
2980 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
2981 kek,
2982 commit_lock: Arc::new(parking_lot::Mutex::new(())),
2983 shared: None,
2984 table_name: None,
2985 auth: None,
2986 read_only: false,
2987 }
2988 }
2989}
2990
2991fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
2995 use crate::query::Condition;
2996 match c {
2997 Condition::Pk(_)
2999 | Condition::BitmapEq { .. }
3000 | Condition::BitmapIn { .. }
3001 | Condition::BytesPrefix { .. }
3002 | Condition::IsNull { .. }
3003 | Condition::IsNotNull { .. } => 0,
3004 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
3006 1
3007 }
3008 Condition::FmContains { .. }
3010 | Condition::FmContainsAll { .. }
3011 | Condition::Ann { .. }
3012 | Condition::SparseMatch { .. } => 2,
3013 }
3014}
3015
3016impl Table {
3017 pub(crate) fn build_secondary_index_artifact<F>(
3023 &self,
3024 definition: &IndexDef,
3025 rows: &[Row],
3026 mut checkpoint: F,
3027 ) -> Result<SecondaryIndexArtifact>
3028 where
3029 F: FnMut(usize, usize) -> Result<()>,
3030 {
3031 let mut build_schema = self.schema.clone();
3032 build_schema.indexes = vec![definition.clone()];
3033 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
3034 let name_to_id: HashMap<&str, u16> = self
3035 .schema
3036 .columns
3037 .iter()
3038 .map(|column| (column.name.as_str(), column.id))
3039 .collect();
3040 let mut hot = HotIndex::new();
3041 let mut accepted = Vec::with_capacity(rows.len());
3042
3043 for (offset, row) in rows.iter().enumerate() {
3044 checkpoint(offset, rows.len())?;
3045 if row.deleted {
3046 continue;
3047 }
3048 if let Some(predicate) = &definition.predicate {
3049 let columns: HashMap<u16, &Value> =
3050 row.columns.iter().map(|(id, value)| (*id, value)).collect();
3051 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
3052 continue;
3053 }
3054 }
3055 accepted.push(row);
3056 if definition.kind == IndexKind::LearnedRange {
3057 continue;
3058 }
3059 let effective = if self.column_keys.is_empty() {
3060 row.clone()
3061 } else {
3062 self.tokenized_for_indexes(row)
3063 };
3064 if definition.kind == IndexKind::Ann {
3065 if let (Some(index), Some(vector)) = (
3066 ann.get_mut(&definition.column_id),
3067 effective
3068 .columns
3069 .get(&definition.column_id)
3070 .and_then(Value::as_embedding),
3071 ) {
3072 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
3073 checkpoint(offset, rows.len())
3074 })?;
3075 }
3076 } else {
3077 index_into_single(
3078 definition,
3079 &build_schema,
3080 &effective,
3081 &mut hot,
3082 &mut bitmap,
3083 &mut ann,
3084 &mut fm,
3085 &mut sparse,
3086 &mut minhash,
3087 );
3088 }
3089 }
3090 checkpoint(rows.len(), rows.len())?;
3091
3092 if let Some(index) = ann.get_mut(&definition.column_id) {
3093 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
3094 }
3095
3096 let column_id = definition.column_id;
3097 match definition.kind {
3098 IndexKind::Bitmap => bitmap
3099 .remove(&column_id)
3100 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
3101 IndexKind::Ann => ann
3102 .remove(&column_id)
3103 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
3104 IndexKind::FmIndex => fm
3105 .remove(&column_id)
3106 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
3107 IndexKind::Sparse => sparse
3108 .remove(&column_id)
3109 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
3110 IndexKind::MinHash => minhash
3111 .remove(&column_id)
3112 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
3113 IndexKind::LearnedRange => {
3114 let epsilon = definition
3115 .options
3116 .learned_range
3117 .as_ref()
3118 .map(|options| options.epsilon)
3119 .unwrap_or(16);
3120 let column = self
3121 .schema
3122 .columns
3123 .iter()
3124 .find(|column| column.id == column_id)
3125 .ok_or_else(|| {
3126 MongrelError::Schema(format!(
3127 "index {} references unknown column {column_id}",
3128 definition.name
3129 ))
3130 })?;
3131 let index = match column.ty {
3132 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
3133 let pairs: Vec<_> = accepted
3134 .iter()
3135 .filter_map(|row| match row.columns.get(&column_id) {
3136 Some(Value::Int64(value)) if !row.deleted => {
3137 Some((*value, row.row_id.0))
3138 }
3139 _ => None,
3140 })
3141 .collect();
3142 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
3143 }
3144 TypeId::Float32 | TypeId::Float64 => {
3145 let pairs: Vec<_> = accepted
3146 .iter()
3147 .filter_map(|row| match row.columns.get(&column_id) {
3148 Some(Value::Float64(value)) if !row.deleted => {
3149 Some((*value, row.row_id.0))
3150 }
3151 _ => None,
3152 })
3153 .collect();
3154 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
3155 }
3156 ref ty => {
3157 return Err(MongrelError::Schema(format!(
3158 "LearnedRange index {} does not support {ty:?}",
3159 definition.name
3160 )));
3161 }
3162 };
3163 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
3164 }
3165 }
3166 .ok_or_else(|| {
3167 MongrelError::Other(format!(
3168 "failed to construct hidden index {}",
3169 definition.name
3170 ))
3171 })
3172 }
3173
3174 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
3175 let dir = dir.as_ref().to_path_buf();
3176 std::fs::create_dir_all(&dir)?;
3179 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3180 let pinned = root.io_path()?;
3181 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
3182 ctx.root_guard = Some(root.clone());
3183 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3184 root.finalize_deferred_sync_shallow()?;
3190 Ok(table)
3191 }
3192
3193 pub fn create_encrypted(
3204 dir: impl AsRef<Path>,
3205 schema: Schema,
3206 table_id: u64,
3207 passphrase: &str,
3208 ) -> Result<Self> {
3209 let dir = dir.as_ref().to_path_buf();
3210 std::fs::create_dir_all(&dir)?;
3211 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3212 root.create_directory_all(META_DIR)?;
3213 let salt = crate::encryption::random_salt()?;
3214 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
3215 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
3216 let pinned = root.io_path()?;
3217 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3218 ctx.root_guard = Some(root.clone());
3219 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3220 root.finalize_deferred_sync()?;
3221 Ok(table)
3222 }
3223
3224 pub fn create_with_key(
3229 dir: impl AsRef<Path>,
3230 schema: Schema,
3231 table_id: u64,
3232 key: &[u8],
3233 ) -> Result<Self> {
3234 let dir = dir.as_ref().to_path_buf();
3235 std::fs::create_dir_all(&dir)?;
3236 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3237 root.create_directory_all(META_DIR)?;
3238 let salt = crate::encryption::random_salt()?;
3239 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
3240 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
3241 let pinned = root.io_path()?;
3242 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3243 ctx.root_guard = Some(root.clone());
3244 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3245 root.finalize_deferred_sync()?;
3246 Ok(table)
3247 }
3248
3249 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
3251 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3252 let salt = read_table_encryption_salt_root(&root)?;
3253 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
3254 let pinned = root.io_path()?;
3255 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3256 ctx.root_guard = Some(root);
3257 Self::open_in(&pinned, ctx)
3258 }
3259
3260 pub(crate) fn create_in(
3261 dir: impl AsRef<Path>,
3262 schema: Schema,
3263 table_id: u64,
3264 ctx: SharedCtx,
3265 ) -> Result<Self> {
3266 schema.validate_auto_increment()?;
3267 schema.validate_defaults()?;
3268 schema.validate_ai()?;
3269 for index in &schema.indexes {
3270 index.validate_options()?;
3271 }
3272 let dir = dir.as_ref().to_path_buf();
3273 let runs_root = match ctx.root_guard.as_ref() {
3274 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
3275 None => {
3276 crate::durable_file::create_directory_all(&dir)?;
3277 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
3278 None
3279 }
3280 };
3281 match ctx.root_guard.as_deref() {
3282 Some(root) => write_schema_durable(root, &schema)?,
3283 None => write_schema(&dir, &schema)?,
3284 }
3285 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
3286 let (wal, current_txn_id) = match ctx.shared.clone() {
3289 Some(s) => (WalSink::Shared(s), 0),
3290 None => {
3291 let pinned_wal_root = match ctx.root_guard.as_deref() {
3292 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
3293 None => None,
3294 };
3295 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
3296 root.io_path()?
3297 } else {
3298 let wal_dir = dir.join(WAL_DIR);
3299 std::fs::create_dir_all(&wal_dir)?;
3300 wal_dir
3301 };
3302 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
3303 (Some(root), Some(dk)) => {
3304 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
3305 }
3306 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
3307 (None, Some(dk)) => Wal::create_with_cipher(
3308 wal_dir.join("seg-000000.wal"),
3309 Epoch(0),
3310 Some(make_cipher(dk)),
3311 0,
3312 )?,
3313 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
3314 };
3315 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3316 (WalSink::Private(w), 1)
3317 }
3318 };
3319 let mut manifest = Manifest::new(table_id, schema.schema_id);
3320 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3325 match ctx.root_guard.as_deref() {
3326 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
3327 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
3328 }
3329 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
3330 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3331 let auto_inc = resolve_auto_inc(&schema);
3332 let rcache_dir = dir.join(RCACHE_DIR);
3333 let initial_view = ReadGeneration::empty(&schema);
3334 Ok(Self {
3335 dir,
3336 _root_guard: ctx.root_guard,
3337 runs_root,
3338 idx_root: None,
3339 table_id,
3340 name: ctx.table_name.unwrap_or_default(),
3341 auth: ctx.auth,
3342 read_only: ctx.read_only,
3343 durable_commit_failed: false,
3344 wal,
3345 memtable: Memtable::new(),
3346 mutable_run: MutableRun::new(),
3347 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3348 compaction_zstd_level: 3,
3349 allocator: RowIdAllocator::new(0),
3350 epoch: ctx.epoch,
3351 data_generation: 0,
3352 schema,
3353 hot: HotIndex::new(),
3354 kek: ctx.kek,
3355 column_keys,
3356 run_refs: Vec::new(),
3357 retiring: Vec::new(),
3358 next_run_id: 1,
3359 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3360 current_txn_id,
3361 pending_private_mutations: false,
3362 bitmap,
3363 ann,
3364 fm,
3365 sparse,
3366 minhash,
3367 learned_range: Arc::new(HashMap::new()),
3368 pk_by_row: ReversePkMap::new(),
3369 pinned: BTreeMap::new(),
3370 live_count: 0,
3371 reservoir: crate::reservoir::Reservoir::default(),
3372 reservoir_complete: true,
3373 had_deletes: false,
3374 recent_delete_preimages: HashMap::new(),
3375 run_row_id_ranges: HashMap::new(),
3376 agg_cache: Arc::new(HashMap::new()),
3377 global_idx_epoch: 0,
3378 indexes_complete: true,
3379 index_build_policy: IndexBuildPolicy::default(),
3380 pk_by_row_complete: false,
3381 flushed_epoch: 0,
3382 page_cache: ctx.page_cache,
3383 decoded_cache: ctx.decoded_cache,
3384 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3385 snapshots: ctx.snapshots,
3386 commit_lock: ctx.commit_lock,
3387 result_cache: {
3394 let cache = ResultCache::new()
3395 .with_dir(rcache_dir.clone())
3396 .with_cache_dek(cache_dek.clone());
3397 let metrics = Arc::new(LookupMetrics::default());
3401 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3402 match spawn_persistent_cache_worker(
3403 rcache_dir.clone(),
3404 cache_dek.clone(),
3405 (*metrics).clone(),
3406 ) {
3407 Ok((writer, handle, completion_rx)) => {
3408 cache_arc
3409 .lock()
3410 .install_persistent_writer(writer, handle, completion_rx);
3411 }
3412 Err(_) => cache_arc.lock().note_worker_spawn_failure(),
3413 }
3414 cache_arc
3415 },
3416 pending_delete_rids: roaring::RoaringBitmap::new(),
3417 pending_put_cols: std::collections::HashSet::new(),
3418 pending_rows: Vec::new(),
3419 pending_rows_auto_inc: Vec::new(),
3420 pending_dels: Vec::new(),
3421 pending_truncate: None,
3422 wal_dek,
3423 auto_inc,
3424 ttl: None,
3425 pins: Arc::new(crate::retention::PinRegistry::new()),
3426 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3427 read_generation_pin: None,
3428 lookup_metrics: LookupMetrics::default(),
3429 run_lookup: RunLookupState::default(),
3430 })
3431 }
3432
3433 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
3437 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3438 let pinned = root.io_path()?;
3439 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
3440 ctx.root_guard = Some(root);
3441 Self::open_in(&pinned, ctx)
3442 }
3443
3444 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
3447 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3448 let salt = read_table_encryption_salt_root(&root)?;
3449 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
3450 let pinned = root.io_path()?;
3451 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3452 ctx.root_guard = Some(root);
3453 let t = Self::open_in(&pinned, ctx)?;
3454 Ok(t)
3455 }
3456
3457 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
3458 let dir = dir.as_ref().to_path_buf();
3459 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3460 let mut manifest = match ctx.root_guard.as_ref() {
3461 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
3462 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
3463 };
3464 let schema: Schema = match ctx.root_guard.as_ref() {
3465 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
3466 None => read_schema(&dir)?,
3467 };
3468 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
3474 let runs_root = match ctx.root_guard.as_ref() {
3475 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
3476 None => None,
3477 };
3478 let idx_root = match ctx.root_guard.as_ref() {
3479 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
3480 Ok(root) => Some(Arc::new(root)),
3481 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
3482 Err(error) => return Err(error.into()),
3483 },
3484 None => None,
3485 };
3486 schema.validate_auto_increment()?;
3487 schema.validate_defaults()?;
3488 schema.validate_ai()?;
3489 for index in &schema.indexes {
3490 index.validate_options()?;
3491 }
3492 let replay_epoch = Epoch(manifest.current_epoch);
3493 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
3494 let private_replayed = if ctx.shared.is_none() {
3495 match latest_wal_segment(&dir.join(WAL_DIR))? {
3496 Some(path) => {
3497 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
3498 crate::wal::replay_with_cipher(path, cipher)?
3499 }
3500 None => Vec::new(),
3501 }
3502 } else {
3503 Vec::new()
3504 };
3505 if ctx.shared.is_none() {
3506 preflight_standalone_open(
3507 &dir,
3508 runs_root.as_deref(),
3509 idx_root.as_deref(),
3510 &manifest,
3511 &schema,
3512 &private_replayed,
3513 ctx.kek.clone(),
3514 )?;
3515 }
3516 let next_run_id = derive_next_run_id(
3517 &dir,
3518 runs_root.as_deref(),
3519 &manifest.runs,
3520 &manifest.retiring,
3521 )?;
3522 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
3526 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
3527 None => {
3528 let replayed = private_replayed;
3529 let wal_dir = dir.join(WAL_DIR);
3535 crate::durable_file::create_directory_all(&wal_dir)?;
3536 let segment = next_wal_segment(&wal_dir)?;
3537 let segment_no = wal_segment_number(&segment).unwrap_or(0);
3538 let temporary = wal_dir.join(format!(
3539 ".recovery-{}-{}-{segment_no:06}.tmp",
3540 std::process::id(),
3541 std::time::SystemTime::now()
3542 .duration_since(std::time::UNIX_EPOCH)
3543 .unwrap_or_default()
3544 .as_nanos()
3545 ));
3546 let mut w = Wal::create_with_cipher(
3547 &temporary,
3548 replay_epoch,
3549 wal_dek.as_ref().map(|dk| make_cipher(dk)),
3550 segment_no,
3551 )?;
3552 for record in &replayed {
3553 w.append_txn(record.txn_id, record.op.clone())?;
3554 }
3555 let mut w = w.publish_as(segment)?;
3556 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3557 let next_txn_id = replayed
3558 .iter()
3559 .map(|record| record.txn_id)
3560 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
3561 .max()
3562 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
3563 .unwrap_or(1);
3564 (WalSink::Private(w), replayed, next_txn_id)
3565 }
3566 };
3567
3568 let mut memtable = Memtable::new();
3569 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
3570 let persisted_epoch = manifest.current_epoch;
3571 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
3578 s.next = manifest.auto_inc_next;
3579 s.seeded = manifest.auto_inc_next > 0;
3580 s
3581 });
3582
3583 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
3590 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
3591 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
3592 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
3593 std::collections::BTreeMap::new();
3594 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
3595 let mut recovered_epoch = manifest.current_epoch;
3596 let mut recovered_manifest_dirty = schema_manifest_repair;
3597 let mut saw_delete = false;
3598 for record in replayed {
3599 let txn_id = record.txn_id;
3600 match record.op {
3601 Op::Put { rows, .. } => {
3602 let rows: Vec<Row> = bincode::deserialize(&rows)?;
3603 for row in &rows {
3604 allocator.advance_to(row.row_id)?;
3605 if let Some(ai) = auto_inc.as_mut() {
3606 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3607 let next = n.checked_add(1).ok_or_else(|| {
3608 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3609 })?;
3610 if next > ai.next {
3611 ai.next = next;
3612 }
3613 }
3614 }
3615 }
3616 staged_puts.entry(txn_id).or_default().extend(rows);
3617 }
3618 Op::Delete { row_ids, .. } => {
3619 staged_deletes.entry(txn_id).or_default().extend(row_ids);
3620 }
3621 Op::TxnCommit { epoch, .. } => {
3622 let commit_epoch = Epoch(epoch);
3623 recovered_epoch = recovered_epoch.max(epoch);
3624 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
3625 memtable = Memtable::new();
3626 replayed_puts.clear();
3627 replayed_deletes.clear();
3628 manifest.runs.clear();
3629 manifest.retiring.clear();
3630 manifest.live_count = 0;
3631 manifest.global_idx_epoch = 0;
3632 manifest.current_epoch = manifest.current_epoch.max(epoch);
3633 recovered_manifest_dirty = true;
3634 saw_delete = true;
3635 }
3636 if let Some(puts) = staged_puts.remove(&txn_id) {
3637 if commit_epoch.0 > manifest.flushed_epoch {
3638 for row in &puts {
3639 memtable.upsert(row.clone());
3640 }
3641 replayed_puts.entry(commit_epoch).or_default().extend(puts);
3642 }
3643 }
3644 if let Some(dels) = staged_deletes.remove(&txn_id) {
3645 saw_delete = true;
3646 if commit_epoch.0 > manifest.flushed_epoch {
3647 for rid in dels {
3648 memtable.tombstone(rid, commit_epoch);
3649 replayed_deletes.push((rid, commit_epoch));
3650 }
3651 }
3652 }
3653 }
3654 Op::TxnAbort => {
3655 staged_puts.remove(&txn_id);
3656 staged_deletes.remove(&txn_id);
3657 staged_truncates.remove(&txn_id);
3658 }
3659 Op::TruncateTable { .. } => {
3660 staged_puts.remove(&txn_id);
3661 staged_deletes.remove(&txn_id);
3662 staged_truncates.insert(txn_id);
3663 }
3664 Op::ExternalTableState { .. }
3665 | Op::Flush { .. }
3666 | Op::Ddl(_)
3667 | Op::BeforeImage { .. }
3668 | Op::CommitTimestamp { .. }
3669 | Op::SpilledRows { .. } => {}
3670 }
3671 }
3672
3673 let rcache_dir = dir.join(RCACHE_DIR);
3674 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3675 let initial_view = ReadGeneration::empty(&schema);
3676 let mut db = Self {
3677 dir,
3678 _root_guard: ctx.root_guard,
3679 runs_root,
3680 idx_root,
3681 table_id: manifest.table_id,
3682 name: ctx.table_name.unwrap_or_default(),
3683 auth: ctx.auth,
3684 read_only: ctx.read_only,
3685 durable_commit_failed: false,
3686 wal,
3687 memtable,
3688 mutable_run: MutableRun::new(),
3689 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3690 compaction_zstd_level: 3,
3691 allocator,
3692 epoch: ctx.epoch,
3693 data_generation: persisted_epoch,
3694 schema,
3695 hot: HotIndex::new(),
3696 kek: ctx.kek,
3697 column_keys,
3698 run_refs: manifest.runs.clone(),
3699 retiring: manifest.retiring.clone(),
3700 next_run_id,
3701 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3702 current_txn_id,
3703 pending_private_mutations: false,
3704 bitmap: HashMap::new(),
3705 ann: HashMap::new(),
3706 fm: HashMap::new(),
3707 sparse: HashMap::new(),
3708 minhash: HashMap::new(),
3709 learned_range: Arc::new(HashMap::new()),
3710 pk_by_row: ReversePkMap::new(),
3711 pinned: BTreeMap::new(),
3712 live_count: manifest.live_count,
3713 reservoir: crate::reservoir::Reservoir::default(),
3714 reservoir_complete: false,
3715 had_deletes: saw_delete
3716 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
3717 != manifest.live_count,
3718 recent_delete_preimages: HashMap::new(),
3719 run_row_id_ranges: HashMap::new(),
3720 agg_cache: Arc::new(HashMap::new()),
3721 global_idx_epoch: manifest.global_idx_epoch,
3722 indexes_complete: true,
3723 index_build_policy: IndexBuildPolicy::default(),
3724 pk_by_row_complete: false,
3725 flushed_epoch: manifest.flushed_epoch,
3726 page_cache: ctx.page_cache,
3727 decoded_cache: ctx.decoded_cache,
3728 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3729 snapshots: ctx.snapshots,
3730 commit_lock: ctx.commit_lock,
3731 result_cache: {
3732 let cache = ResultCache::new()
3733 .with_dir(rcache_dir.clone())
3734 .with_cache_dek(cache_dek.clone());
3735 let metrics = LookupMetrics::default();
3736 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3737 match spawn_persistent_cache_worker(rcache_dir.clone(), cache_dek.clone(), metrics)
3738 {
3739 Ok((writer, handle, completion_rx)) => {
3740 cache_arc
3741 .lock()
3742 .install_persistent_writer(writer, handle, completion_rx);
3743 }
3744 Err(_) => cache_arc.lock().note_worker_spawn_failure(),
3748 }
3749 cache_arc
3750 },
3751 pending_delete_rids: roaring::RoaringBitmap::new(),
3752 pending_put_cols: std::collections::HashSet::new(),
3753 pending_rows: Vec::new(),
3754 pending_rows_auto_inc: Vec::new(),
3755 pending_dels: Vec::new(),
3756 pending_truncate: None,
3757 wal_dek,
3758 auto_inc,
3759 ttl: manifest.ttl,
3760 pins: Arc::new(crate::retention::PinRegistry::new()),
3761 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3762 read_generation_pin: None,
3763 lookup_metrics: LookupMetrics::default(),
3764 run_lookup: RunLookupState::default(),
3765 };
3766
3767 db.epoch.advance_recovered(Epoch(recovered_epoch));
3770
3771 let checkpoint = match db.idx_root.as_deref() {
3776 Some(root) => {
3777 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
3778 }
3779 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
3780 };
3781 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
3782 c.epoch_built == manifest.global_idx_epoch
3783 && manifest.global_idx_epoch > 0
3784 && manifest
3785 .runs
3786 .iter()
3787 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
3788 });
3789 if let Some(loaded) = checkpoint {
3790 if checkpoint_valid {
3791 db.hot = loaded.hot;
3792 db.bitmap = loaded.bitmap;
3793 db.ann = loaded.ann;
3794 db.fm = loaded.fm;
3795 db.sparse = loaded.sparse;
3796 db.minhash = loaded.minhash;
3797 db.learned_range = Arc::new(loaded.learned_range);
3798 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
3803 for (cid, idx) in bitmap0 {
3804 db.bitmap.entry(cid).or_insert(idx);
3805 }
3806 for (cid, idx) in ann0 {
3807 db.ann.entry(cid).or_insert(idx);
3808 }
3809 for (cid, idx) in fm0 {
3810 db.fm.entry(cid).or_insert(idx);
3811 }
3812 for (cid, idx) in sparse0 {
3813 db.sparse.entry(cid).or_insert(idx);
3814 }
3815 for (cid, idx) in minhash0 {
3816 db.minhash.entry(cid).or_insert(idx);
3817 }
3818 }
3821 }
3822 if !checkpoint_valid {
3823 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
3824 db.bitmap = bitmap;
3825 db.ann = ann;
3826 db.fm = fm;
3827 db.sparse = sparse;
3828 db.minhash = minhash;
3829 db.rebuild_indexes_from_runs()?;
3830 db.build_learned_ranges()?;
3831 }
3832
3833 for (epoch, group) in replayed_puts {
3838 let (losers, winner_pks) = db.partition_pk_winners(&group);
3839 for (key, &row_id) in &winner_pks {
3840 if let Some(old_rid) = db.hot.get(key) {
3841 if old_rid != row_id {
3842 db.tombstone_row(old_rid, epoch, None, false);
3843 }
3844 }
3845 }
3846 for &loser_rid in &losers {
3847 db.tombstone_row(loser_rid, epoch, None, false);
3848 }
3849 for (key, row_id) in winner_pks {
3850 db.insert_hot_pk(key, row_id);
3851 }
3852 if db.schema.primary_key().is_none() {
3853 for r in &group {
3854 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3855 }
3856 }
3857 for r in &group {
3858 if !losers.contains(&r.row_id) {
3859 db.index_row(r);
3860 }
3861 }
3862 }
3863 for (rid, epoch) in &replayed_deletes {
3867 db.remove_hot_for_row(*rid, *epoch);
3868 }
3869
3870 if recovered_manifest_dirty {
3871 let rows = db.visible_rows(Snapshot::unbounded())?;
3872 db.live_count = rows.len() as u64;
3873 db.persist_manifest(Epoch(recovered_epoch))?;
3874 }
3875
3876 let identity = crate::result_cache::PersistentCacheIdentity {
3883 table_id: db.table_id,
3884 schema_id: db.schema.schema_id,
3885 logical_generation: db.epoch.visible().0,
3886 };
3887 db.result_cache.lock().load_persistent(identity);
3888 db.refresh_run_row_id_ranges();
3890 db.load_run_lookup_directory();
3895 Ok(db)
3896 }
3897
3898 fn load_run_lookup_directory(&mut self) {
3903 let active_runs = self.run_refs.clone();
3904 let schema_id = self.schema.schema_id;
3905 let run_generation = active_runs.len() as u64;
3906 let index_generation = self.global_idx_epoch;
3907 let fingerprint = crate::run_lookup::compute_fingerprint(
3908 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
3909 schema_id,
3910 index_generation,
3911 run_generation,
3912 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
3913 );
3914 let runs_dir = match self.runs_root.as_deref() {
3916 Some(root) => match root.io_path() {
3917 Ok(p) => p,
3918 Err(_) => self.dir.join(RUNS_DIR),
3919 },
3920 None => self.dir.join(RUNS_DIR),
3921 };
3922 match crate::run_lookup::RunLookupDirectory::read_checkpoint_sharded(&runs_dir, fingerprint)
3923 {
3924 Ok(dir) => {
3925 self.run_lookup = RunLookupState {
3926 directory: Some(std::sync::Arc::new(dir)),
3927 fingerprint,
3928 complete: true,
3929 };
3930 }
3931 Err(error) => {
3932 self.lookup_metrics
3933 .directory_incomplete
3934 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3935 self.lookup_metrics
3936 .directory_lookup_fallback
3937 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3938 self.run_lookup = RunLookupState {
3939 directory: None,
3940 fingerprint,
3941 complete: false,
3942 };
3943 let _ = error;
3947 }
3948 }
3949 }
3950
3951 fn ensure_reservoir_complete(&mut self) -> Result<()> {
3957 if self.reservoir_complete {
3958 return Ok(());
3959 }
3960 self.rebuild_reservoir()?;
3961 self.reservoir_complete = true;
3962 Ok(())
3963 }
3964
3965 fn rebuild_reservoir(&mut self) -> Result<()> {
3968 let snap = self.snapshot();
3969 let rows = self.visible_rows(snap)?;
3970 self.reservoir.reset();
3971 for r in rows {
3972 self.reservoir.offer(r.row_id.0);
3973 }
3974 Ok(())
3975 }
3976
3977 pub fn rebuild_indexes(&mut self) -> Result<()> {
3981 self.rebuild_indexes_from_runs_inner(None)
3982 }
3983
3984 pub fn __force_hot_map_for_test(&mut self, pk_bytes: &[u8], row_id: RowId) {
3995 self.hot.insert(pk_bytes.to_vec(), row_id);
3996 }
3997
3998 pub fn hot_for_test_remove(&mut self, pk_bytes: &[u8]) {
4001 self.hot.remove(pk_bytes);
4002 }
4003
4004 pub fn set_indexes_incomplete_for_test(&mut self) {
4007 self.indexes_complete = false;
4008 }
4009
4010 pub fn bump_hot_checkpoint_rejected_for_test(&self) {
4014 self.lookup_metrics
4015 .hot_checkpoint_rejected_total
4016 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4017 }
4018
4019 pub fn __install_run_for_test(&mut self, rows: &[Row]) -> Result<u128> {
4026 let mut rows = rows.to_vec();
4027 rows.sort_by_key(|r| (r.row_id, r.committed_epoch));
4028 let epoch = rows
4029 .iter()
4030 .map(|r| r.committed_epoch)
4031 .max()
4032 .unwrap_or(Epoch::ZERO);
4033 let run_id = self.alloc_run_id()?;
4034 let path = self.run_path(run_id);
4035 if let Some(parent) = path.parent() {
4036 std::fs::create_dir_all(parent)
4037 .map_err(|e| MongrelError::Other(format!("create runs dir: {e}")))?;
4038 }
4039 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
4040 if let Some(kek) = &self.kek {
4041 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4042 }
4043 let header = match self.create_run_file(run_id)? {
4044 Some(file) => writer.write_file(file, &rows)?,
4045 None => writer.write(&path, &rows)?,
4046 };
4047 self.run_refs.push(RunRef {
4048 run_id: run_id as u128,
4049 level: 0,
4050 epoch_created: epoch.0,
4051 row_count: header.row_count,
4052 });
4053 self.run_row_id_ranges
4054 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
4055 Ok(run_id as u128)
4056 }
4057
4058 pub fn __rebuild_run_lookup_for_test(&mut self) -> Result<()> {
4063 self.publish_run_lookup_directory()
4064 }
4065
4066 pub fn __set_run_lookup_complete_for_test(&mut self, complete: bool) {
4073 self.run_lookup.complete = complete;
4074 }
4075
4076 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
4077 self.rebuild_indexes_from_runs_inner(None)
4078 }
4079
4080 fn pk_equality_fallback(
4107 &self,
4108 pk_column_id: u16,
4109 lookup: &[u8],
4110 snapshot: Snapshot,
4111 ) -> Result<(RowIdSet, crate::trace::HotFallbackReason)> {
4112 let mut tombstone_hit = false;
4113 let mut overlay_versions = 0u64;
4114 let now_nanos = unix_nanos_now();
4115 for row in self.memtable.visible_versions_at(snapshot) {
4117 overlay_versions += 1;
4118 self.lookup_metrics
4119 .hot_lookup_fallback_overlay_rows
4120 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4121 if row.deleted {
4122 tombstone_hit = true;
4123 continue;
4124 }
4125 if self.row_expired_at(&row, now_nanos) {
4126 continue;
4127 }
4128 if let Some(pk_val) = row.columns.get(&pk_column_id) {
4129 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
4130 self.lookup_metrics
4131 .hot_fallback_overlay_versions_total
4132 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
4133 return Ok((
4134 RowIdSet::one(row.row_id.0),
4135 crate::trace::HotFallbackReason::MissingMapping,
4136 ));
4137 }
4138 }
4139 }
4140 for row in self.mutable_run.visible_versions_at(snapshot) {
4141 overlay_versions += 1;
4142 self.lookup_metrics
4143 .hot_lookup_fallback_overlay_rows
4144 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4145 if row.deleted {
4146 tombstone_hit = true;
4147 continue;
4148 }
4149 if self.row_expired_at(&row, now_nanos) {
4150 continue;
4151 }
4152 if let Some(pk_val) = row.columns.get(&pk_column_id) {
4153 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
4154 self.lookup_metrics
4155 .hot_fallback_overlay_versions_total
4156 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
4157 return Ok((
4158 RowIdSet::one(row.row_id.0),
4159 crate::trace::HotFallbackReason::MissingMapping,
4160 ));
4161 }
4162 }
4163 }
4164 if lookup.len() == 8 {
4167 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
4168 let n = i64::from_be_bytes(arr);
4169 let result = self.range_scan_i64(pk_column_id, n, n, snapshot)?;
4170 self.lookup_metrics
4171 .hot_lookup_fallback_runs
4172 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4173 self.lookup_metrics
4177 .hot_fallback_runs_considered_total
4178 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4179 self.lookup_metrics
4180 .hot_fallback_overlay_versions_total
4181 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
4182 let reason = if result.is_empty() {
4183 if tombstone_hit {
4184 crate::trace::HotFallbackReason::Tombstone
4185 } else {
4186 crate::trace::HotFallbackReason::MissingMapping
4187 }
4188 } else {
4189 crate::trace::HotFallbackReason::MissingMapping
4190 };
4191 return Ok((result, reason));
4192 }
4193 }
4194 let mut found: std::collections::BTreeSet<u64> = std::collections::BTreeSet::new();
4197 let overlay = self.overlay_rid_set(snapshot);
4198 let mut runs_considered = 0u64;
4199 for rr in &self.run_refs {
4200 runs_considered += 1;
4201 self.lookup_metrics
4202 .hot_lookup_fallback_runs
4203 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4204 let mut reader = self.open_reader(rr.run_id)?;
4205 for row in reader.visible_versions_at(snapshot)? {
4206 if overlay.contains(&row.row_id.0) || row.deleted {
4207 if row.deleted {
4208 tombstone_hit = true;
4209 }
4210 continue;
4211 }
4212 if self.row_expired_at(&row, now_nanos) {
4213 continue;
4214 }
4215 if let Some(pk_val) = row.columns.get(&pk_column_id) {
4216 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
4217 found.insert(row.row_id.0);
4218 }
4219 }
4220 }
4221 }
4222 let reason = if tombstone_hit {
4223 crate::trace::HotFallbackReason::Tombstone
4224 } else {
4225 crate::trace::HotFallbackReason::MissingMapping
4228 };
4229 self.lookup_metrics
4230 .hot_fallback_overlay_versions_total
4231 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
4232 self.lookup_metrics
4233 .hot_fallback_runs_considered_total
4234 .fetch_add(runs_considered, std::sync::atomic::Ordering::Relaxed);
4235 Ok((RowIdSet::from_unsorted(found.into_iter().collect()), reason))
4236 }
4237
4238 fn record_hot_fallback_reason(&self, reason: crate::trace::HotFallbackReason) {
4242 let idx = hot_fallback_reason_index(reason);
4243 self.lookup_metrics.hot_fallback_reasons[idx]
4244 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4245 self.lookup_metrics
4246 .hot_lookup_fallback
4247 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4248 crate::trace::QueryTrace::record(|t| {
4249 t.hot_lookup_attempted = true;
4250 t.hot_lookup_hit = false;
4251 t.hot_fallback_reason = Some(reason.as_str());
4252 });
4253 }
4254
4255 fn rebuild_indexes_from_runs_inner(
4256 &mut self,
4257 control: Option<&crate::ExecutionControl>,
4258 ) -> Result<()> {
4259 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
4262 crate::retention::PinSource::OnlineIndexBuild,
4263 self.current_epoch(),
4264 );
4265 self.hot = HotIndex::new();
4266 self.pk_by_row.clear();
4267 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
4268 self.bitmap = bitmap;
4269 self.ann = ann;
4270 self.fm = fm;
4271 self.sparse = sparse;
4272 self.minhash = minhash;
4273 let snapshot = Snapshot::unbounded();
4296 let ttl_now = unix_nanos_now();
4297 let mut scanned = 0_usize;
4298 let mut winners: HashMap<RowId, (VersionStamp, Row)> = HashMap::new();
4299 let fold =
4300 |row: Row, winners: &mut HashMap<RowId, (VersionStamp, Row)>, scanned: &mut usize| {
4301 *scanned += 1;
4302 let stamp = VersionStamp {
4303 epoch: row.committed_epoch,
4304 commit_ts: row.commit_ts,
4305 };
4306 winners
4307 .entry(row.row_id)
4308 .and_modify(|entry| {
4309 if stamp.is_newer_than(entry.0) {
4310 *entry = (stamp, row.clone());
4311 }
4312 })
4313 .or_insert((stamp, row));
4314 };
4315 for rr in self.run_refs.clone() {
4316 if let Some(control) = control {
4317 control.checkpoint()?;
4318 }
4319 let mut reader = self.open_reader(rr.run_id)?;
4320 for row in reader.visible_versions_at(snapshot)? {
4321 if scanned.is_multiple_of(256) {
4322 if let Some(control) = control {
4323 control.checkpoint()?;
4324 }
4325 }
4326 fold(row, &mut winners, &mut scanned);
4330 }
4331 }
4332 for (index, row) in self
4333 .memtable
4334 .visible_versions_at(snapshot)
4335 .into_iter()
4336 .enumerate()
4337 {
4338 if index & 255 == 0 {
4339 if let Some(control) = control {
4340 control.checkpoint()?;
4341 }
4342 }
4343 fold(row, &mut winners, &mut scanned);
4344 }
4345 for (index, row) in self
4346 .mutable_run
4347 .visible_versions_at(snapshot)
4348 .into_iter()
4349 .enumerate()
4350 {
4351 if index & 255 == 0 {
4352 if let Some(control) = control {
4353 control.checkpoint()?;
4354 }
4355 }
4356 fold(row, &mut winners, &mut scanned);
4357 }
4358 let any_predicate = self
4363 .schema
4364 .indexes
4365 .iter()
4366 .any(|idx| idx.predicate.is_some());
4367 let name_to_id: HashMap<&str, u16> = self
4368 .schema
4369 .columns
4370 .iter()
4371 .map(|c| (c.name.as_str(), c.id))
4372 .collect();
4373 for (_rid, (stamp, row)) in winners.drain() {
4374 if row.deleted {
4375 let _ = stamp;
4378 continue;
4379 }
4380 let tok_row = self.tokenized_for_indexes(&row);
4383 if !any_predicate {
4384 index_into(
4385 &self.schema,
4386 &tok_row,
4387 &mut self.hot,
4388 &mut self.bitmap,
4389 &mut self.ann,
4390 &mut self.fm,
4391 &mut self.sparse,
4392 &mut self.minhash,
4393 );
4394 continue;
4395 }
4396 if let Some(pk_col) = self.schema.primary_key() {
4397 if let Some(pk_val) = tok_row.columns.get(&pk_col.id) {
4398 self.hot.insert(pk_val.encode_key(), tok_row.row_id);
4399 }
4400 }
4401 let columns_map: HashMap<u16, &Value> =
4402 row.columns.iter().map(|(k, v)| (*k, v)).collect();
4403 for idef in &self.schema.indexes {
4404 if let Some(pred) = &idef.predicate {
4405 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4406 continue;
4407 }
4408 }
4409 index_into_single(
4410 idef,
4411 &self.schema,
4412 &tok_row,
4413 &mut self.hot,
4414 &mut self.bitmap,
4415 &mut self.ann,
4416 &mut self.fm,
4417 &mut self.sparse,
4418 &mut self.minhash,
4419 );
4420 }
4421 }
4422 let pin_epochs = self.active_pin_epochs_for_rebuild();
4428 if !pin_epochs.is_empty() {
4429 let current_snap = Snapshot::at(Epoch(u64::MAX));
4430 for pin_epoch in pin_epochs {
4431 if let Some(control) = control {
4432 control.checkpoint()?;
4433 }
4434 let pin_snap = Snapshot::at(pin_epoch);
4435 for rr in self.run_refs.clone() {
4436 if let Some(control) = control {
4437 control.checkpoint()?;
4438 }
4439 let mut reader = self.open_reader(rr.run_id)?;
4440 for row in reader.visible_rows(pin_epoch)? {
4441 if row.deleted || self.row_expired_at(&row, ttl_now) {
4442 continue;
4443 }
4444 if self.get(row.row_id, current_snap).is_none() {
4445 self.index_bitmap_membership_only(&row);
4446 }
4447 }
4448 }
4449 for row in self
4450 .mutable_run
4451 .visible_versions_at(pin_snap)
4452 .into_iter()
4453 .chain(self.memtable.visible_versions_at(pin_snap))
4454 {
4455 if row.deleted || self.row_expired_at(&row, ttl_now) {
4456 continue;
4457 }
4458 if self.get(row.row_id, current_snap).is_none() {
4459 self.index_bitmap_membership_only(&row);
4460 }
4461 }
4462 }
4463 }
4464 self.recent_delete_preimages.clear();
4465 self.refresh_pk_by_row_from_hot();
4466 Ok(())
4467 }
4468
4469 fn index_bitmap_membership_only(&mut self, row: &Row) {
4472 if row.deleted {
4473 return;
4474 }
4475 let columns_map: HashMap<u16, &Value> = row.columns.iter().map(|(k, v)| (*k, v)).collect();
4476 let name_to_id: HashMap<&str, u16> = self
4477 .schema
4478 .columns
4479 .iter()
4480 .map(|c| (c.name.as_str(), c.id))
4481 .collect();
4482 for idef in &self.schema.indexes {
4483 if idef.kind != crate::schema::IndexKind::Bitmap {
4484 continue;
4485 }
4486 if let Some(pred) = &idef.predicate {
4487 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4488 continue;
4489 }
4490 }
4491 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
4492 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4493 b.insert(key, row.row_id);
4494 }
4495 }
4496 }
4497 }
4498
4499 fn refresh_pk_by_row_from_hot(&mut self) {
4500 self.pk_by_row_complete = true;
4501 if self.schema.primary_key().is_none() {
4502 self.pk_by_row.clear();
4503 return;
4504 }
4505 self.pk_by_row = ReversePkMap::from_entries(
4511 self.hot
4512 .entries()
4513 .into_iter()
4514 .map(|(key, row_id)| (row_id, key)),
4515 );
4516 }
4517
4518 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
4519 if self.schema.primary_key().is_some() {
4520 self.pk_by_row.insert(row_id, key.clone());
4521 }
4522 self.hot.insert(key, row_id);
4523 }
4524
4525 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
4529 self.build_learned_ranges_inner(None)
4530 }
4531
4532 fn build_learned_ranges_inner(
4533 &mut self,
4534 control: Option<&crate::ExecutionControl>,
4535 ) -> Result<()> {
4536 self.learned_range = Arc::new(HashMap::new());
4537 if self.run_refs.len() != 1 {
4538 return Ok(());
4539 }
4540 let cols: Vec<(u16, usize)> = self
4541 .schema
4542 .indexes
4543 .iter()
4544 .filter(|i| i.kind == IndexKind::LearnedRange)
4545 .map(|i| {
4546 (
4547 i.column_id,
4548 i.options
4549 .learned_range
4550 .as_ref()
4551 .map(|options| options.epsilon)
4552 .unwrap_or(16),
4553 )
4554 })
4555 .collect();
4556 if cols.is_empty() {
4557 return Ok(());
4558 }
4559 let snapshot = Snapshot::unbounded();
4570 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4571 let (visible_positions, visible_rids) = reader.visible_positions_with_rids_at(snapshot)?;
4572 let row_ids: Vec<u64> = visible_rids.iter().map(|r| *r as u64).collect();
4573 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
4574 if column_index % 256 == 0 {
4575 if let Some(control) = control {
4576 control.checkpoint()?;
4577 }
4578 }
4579 let ty = self
4580 .schema
4581 .columns
4582 .iter()
4583 .find(|c| c.id == cid)
4584 .map(|c| c.ty.clone())
4585 .unwrap_or(TypeId::Int64);
4586 match ty {
4587 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
4588 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
4589 let pairs: Vec<(i64, u64)> = visible_positions
4590 .iter()
4591 .zip(row_ids.iter())
4592 .map(|(&p, &r)| (data[p], r))
4593 .collect();
4594 Arc::make_mut(&mut self.learned_range).insert(
4595 cid,
4596 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
4597 );
4598 }
4599 }
4600 TypeId::Float64 => {
4601 if let columnar::NativeColumn::Float64 { data, .. } =
4602 reader.column_native(cid)?
4603 {
4604 let pairs: Vec<(f64, u64)> = visible_positions
4605 .iter()
4606 .zip(row_ids.iter())
4607 .map(|(&p, &r)| (data[p], r))
4608 .collect();
4609 Arc::make_mut(&mut self.learned_range).insert(
4610 cid,
4611 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
4612 );
4613 }
4614 }
4615 _ => {}
4616 }
4617 }
4618 Ok(())
4619 }
4620
4621 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
4628 if self.indexes_complete {
4629 crate::trace::QueryTrace::record(|t| {
4630 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4631 });
4632 return Ok(());
4633 }
4634 crate::trace::QueryTrace::record(|t| {
4635 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4636 });
4637 self.rebuild_indexes_from_runs()?;
4638 self.build_learned_ranges()?;
4639 self.indexes_complete = true;
4640 let epoch = self.current_epoch();
4641 self.checkpoint_indexes(epoch);
4642 Ok(())
4643 }
4644
4645 #[doc(hidden)]
4648 pub fn ensure_indexes_complete_controlled<F>(
4649 &mut self,
4650 control: &crate::ExecutionControl,
4651 before_publish: F,
4652 ) -> Result<bool>
4653 where
4654 F: FnOnce() -> bool,
4655 {
4656 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
4657 .map(|(changed, _)| changed)
4658 }
4659
4660 #[doc(hidden)]
4663 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
4664 &mut self,
4665 control: &crate::ExecutionControl,
4666 before_publish: F,
4667 ) -> Result<(bool, Option<MaintenanceReceipt>)>
4668 where
4669 F: FnOnce() -> bool,
4670 {
4671 if self.indexes_complete {
4672 crate::trace::QueryTrace::record(|trace| {
4673 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4674 });
4675 return Ok((false, None));
4676 }
4677 crate::trace::QueryTrace::record(|trace| {
4678 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4679 });
4680 control.checkpoint()?;
4681 let maintenance_epoch = self.current_epoch();
4682 self.rebuild_indexes_from_runs_inner(Some(control))?;
4683 self.build_learned_ranges_inner(Some(control))?;
4684 control.checkpoint()?;
4685 if !before_publish() {
4686 return Err(MongrelError::Cancelled);
4687 }
4688 self.indexes_complete = true;
4689 self.checkpoint_indexes(maintenance_epoch);
4690 Ok((
4691 true,
4692 Some(MaintenanceReceipt {
4693 epoch: maintenance_epoch,
4694 }),
4695 ))
4696 }
4697
4698 fn pending_epoch(&self) -> Epoch {
4699 Epoch(self.epoch.visible().0 + 1)
4700 }
4701
4702 fn is_shared(&self) -> bool {
4705 matches!(self.wal, WalSink::Shared(_))
4706 }
4707
4708 fn ensure_txn_id(&mut self) -> Result<u64> {
4712 if self.current_txn_id == 0 {
4713 let id = match &self.wal {
4714 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
4715 WalSink::Private(_) => {
4716 return Err(MongrelError::Full(
4717 "standalone transaction id namespace exhausted".into(),
4718 ))
4719 }
4720 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4721 };
4722 self.current_txn_id = id;
4723 }
4724 Ok(self.current_txn_id)
4725 }
4726
4727 fn wal_append_data(&mut self, op: Op) -> Result<()> {
4730 self.ensure_writable()?;
4731 let txn_id = self.ensure_txn_id()?;
4732 let table_id = self.table_id;
4733 match &mut self.wal {
4734 WalSink::Private(w) => {
4735 w.append_txn(txn_id, op)?;
4736 self.pending_private_mutations = true;
4737 }
4738 WalSink::Shared(s) => {
4739 s.wal.lock().append(txn_id, table_id, op)?;
4740 }
4741 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4742 }
4743 Ok(())
4744 }
4745
4746 fn ensure_writable(&self) -> Result<()> {
4747 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
4748 return Err(MongrelError::ReadOnlyReplica);
4749 }
4750 if self.durable_commit_failed {
4751 return Err(MongrelError::Other(
4752 "table poisoned by post-commit failure; reopen required".into(),
4753 ));
4754 }
4755 Ok(())
4756 }
4757
4758 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
4769 match &self.auth {
4770 Some(checker) => checker.check(&self.name, perm),
4771 None => Ok(()),
4772 }
4773 }
4774 pub fn require_select(&self) -> Result<()> {
4779 self.require(crate::auth_state::RequiredPermission::Select)
4780 }
4781 fn require_insert(&self) -> Result<()> {
4782 self.require(crate::auth_state::RequiredPermission::Insert)
4783 }
4784 #[allow(dead_code)]
4788 fn require_update(&self) -> Result<()> {
4789 self.require(crate::auth_state::RequiredPermission::Update)
4790 }
4791 fn require_delete(&self) -> Result<()> {
4792 self.require(crate::auth_state::RequiredPermission::Delete)
4793 }
4794
4795 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
4798 self.require_insert()?;
4799 Ok(self.put_returning(columns)?.0)
4800 }
4801
4802 pub fn put_returning(
4807 &mut self,
4808 mut columns: Vec<(u16, Value)>,
4809 ) -> Result<(RowId, Option<i64>)> {
4810 self.require_insert()?;
4811 let assigned = self.fill_auto_inc(&mut columns)?;
4812 self.apply_defaults(&mut columns)?;
4813 self.schema.validate_values(&columns)?;
4814 let row_id = if self.schema.clustered {
4819 self.derive_clustered_row_id(&columns)?
4820 } else {
4821 self.allocator.alloc()?
4822 };
4823 let epoch = self.pending_epoch();
4824 let mut row = Row::new(row_id, epoch);
4825 for (col_id, val) in columns {
4826 row.columns.insert(col_id, val);
4827 }
4828 self.commit_rows(vec![row], assigned.is_some())?;
4829 Ok((row_id, assigned))
4830 }
4831
4832 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
4835 self.require_insert()?;
4836 Ok(self
4837 .put_batch_returning(batch)?
4838 .into_iter()
4839 .map(|(r, _)| r)
4840 .collect())
4841 }
4842
4843 pub fn put_batch_returning(
4846 &mut self,
4847 batch: Vec<Vec<(u16, Value)>>,
4848 ) -> Result<Vec<(RowId, Option<i64>)>> {
4849 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
4850 for mut cols in batch {
4851 let assigned = self.fill_auto_inc(&mut cols)?;
4852 self.apply_defaults(&mut cols)?;
4853 filled.push((cols, assigned));
4854 }
4855 for (cols, _) in &filled {
4856 self.schema.validate_values(cols)?;
4857 }
4858 let epoch = self.pending_epoch();
4859 let mut rows = Vec::with_capacity(filled.len());
4860 let mut ids = Vec::with_capacity(filled.len());
4861 let first_row_id = if self.schema.clustered {
4862 None
4863 } else {
4864 let count = u64::try_from(filled.len())
4865 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
4866 Some(self.allocator.alloc_range(count)?.0)
4867 };
4868 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
4869 let row_id = match first_row_id {
4870 Some(first) => RowId(first + row_index as u64),
4871 None => self.derive_clustered_row_id(&cols)?,
4872 };
4873 let mut row = Row::new(row_id, epoch);
4874 for (c, v) in cols {
4875 row.columns.insert(c, v);
4876 }
4877 ids.push((row_id, assigned));
4878 rows.push(row);
4879 }
4880 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
4881 self.commit_rows(rows, all_auto_generated)?;
4882 Ok(ids)
4883 }
4884
4885 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
4891 self.ensure_writable()?;
4892 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4893 return Ok(None);
4894 };
4895 let pos = columns.iter().position(|(c, _)| *c == cid);
4896 let assigned = match pos {
4897 Some(i) => match &columns[i].1 {
4898 Value::Null => {
4899 let next = self.alloc_auto_inc_value()?;
4900 columns[i].1 = Value::Int64(next);
4901 Some(next)
4902 }
4903 Value::Int64(n) => {
4904 self.advance_auto_inc_past(*n)?;
4905 None
4906 }
4907 other => {
4908 return Err(MongrelError::InvalidArgument(format!(
4909 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
4910 other
4911 )))
4912 }
4913 },
4914 None => {
4915 let next = self.alloc_auto_inc_value()?;
4916 columns.push((cid, Value::Int64(next)));
4917 Some(next)
4918 }
4919 };
4920 Ok(assigned)
4921 }
4922
4923 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
4929 for col in &self.schema.columns {
4930 let Some(expr) = &col.default_value else {
4931 continue;
4932 };
4933 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
4935 continue;
4936 }
4937 let pos = columns.iter().position(|(c, _)| *c == col.id);
4938 let needs_default = match pos {
4939 None => true,
4940 Some(i) => matches!(columns[i].1, Value::Null),
4941 };
4942 if !needs_default {
4943 continue;
4944 }
4945 let v = match expr {
4946 crate::schema::DefaultExpr::Static(v) => v.clone(),
4947 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
4948 crate::schema::DefaultExpr::Uuid => {
4949 let mut buf = [0u8; 16];
4950 getrandom::getrandom(&mut buf)
4951 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
4952 Value::Uuid(buf)
4953 }
4954 };
4955 match pos {
4956 None => columns.push((col.id, v)),
4957 Some(i) => columns[i].1 = v,
4958 }
4959 }
4960 Ok(())
4961 }
4962
4963 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
4965 self.ensure_auto_inc_seeded()?;
4966 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4968 let v = ai.next;
4969 ai.next = ai
4970 .next
4971 .checked_add(1)
4972 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4973 Ok(v)
4974 }
4975
4976 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
4979 self.ensure_auto_inc_seeded()?;
4980 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4981 let floor = used
4982 .checked_add(1)
4983 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4984 .max(1);
4985 if ai.next < floor {
4986 ai.next = floor;
4987 }
4988 Ok(())
4989 }
4990
4991 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
4996 let needs_seed = match self.auto_inc {
4997 Some(ai) => !ai.seeded,
4998 None => return Ok(()),
4999 };
5000 if !needs_seed {
5001 return Ok(());
5002 }
5003 if self.seed_empty_auto_inc() {
5004 return Ok(());
5005 }
5006 let cid = self
5007 .auto_inc
5008 .as_ref()
5009 .expect("auto-inc column present")
5010 .column_id;
5011 let max = self.scan_max_int64(cid)?;
5012 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
5013 let floor = max
5014 .checked_add(1)
5015 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
5016 .max(1);
5017 if ai.next < floor {
5018 ai.next = floor;
5019 }
5020 ai.seeded = true;
5021 Ok(())
5022 }
5023
5024 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
5025 if n == 0 || self.auto_inc.is_none() {
5026 return Ok(None);
5027 }
5028 self.ensure_auto_inc_seeded()?;
5029 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
5030 let start = ai.next;
5031 let count = i64::try_from(n)
5032 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
5033 ai.next = ai
5034 .next
5035 .checked_add(count)
5036 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
5037 Ok(Some(start))
5038 }
5039
5040 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
5044 let mut max: i64 = 0;
5045 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
5046 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
5047 if *n > max {
5048 max = *n;
5049 }
5050 }
5051 }
5052 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
5053 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
5054 if *n > max {
5055 max = *n;
5056 }
5057 }
5058 }
5059 for rr in self.run_refs.clone() {
5060 let reader = self.open_reader(rr.run_id)?;
5061 if let Some(stats) = reader.column_page_stats(column_id) {
5062 for s in stats {
5063 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
5064 if n > max {
5065 max = n;
5066 }
5067 }
5068 }
5069 } else if reader.has_column(column_id) {
5070 if let columnar::NativeColumn::Int64 { data, validity } =
5071 reader.column_native_shared(column_id)?
5072 {
5073 for (i, n) in data.iter().enumerate() {
5074 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
5075 {
5076 max = *n;
5077 }
5078 }
5079 }
5080 }
5081 }
5082 Ok(max)
5083 }
5084
5085 fn seed_empty_auto_inc(&mut self) -> bool {
5086 let Some(ai) = self.auto_inc.as_mut() else {
5087 return false;
5088 };
5089 if ai.seeded || self.live_count != 0 {
5090 return false;
5091 }
5092 if ai.next < 1 {
5093 ai.next = 1;
5094 }
5095 ai.seeded = true;
5096 true
5097 }
5098
5099 fn advance_auto_inc_from_native_columns(
5100 &mut self,
5101 columns: &[(u16, columnar::NativeColumn)],
5102 n: usize,
5103 live_before: u64,
5104 ) -> Result<()> {
5105 let Some(ai) = self.auto_inc.as_mut() else {
5106 return Ok(());
5107 };
5108 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
5109 return Ok(());
5110 };
5111 let columnar::NativeColumn::Int64 { data, validity } = col else {
5112 return Err(MongrelError::InvalidArgument(format!(
5113 "AUTO_INCREMENT column {} must be Int64",
5114 ai.column_id
5115 )));
5116 };
5117 let max = if native_int64_strictly_increasing(col, n) {
5118 data.get(n.saturating_sub(1)).copied()
5119 } else {
5120 data.iter()
5121 .take(n)
5122 .enumerate()
5123 .filter_map(|(i, v)| {
5124 if validity.is_empty() || columnar::validity_bit(validity, i) {
5125 Some(*v)
5126 } else {
5127 None
5128 }
5129 })
5130 .max()
5131 };
5132 if let Some(max) = max {
5133 let floor = max
5134 .checked_add(1)
5135 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
5136 .max(1);
5137 if ai.next < floor {
5138 ai.next = floor;
5139 }
5140 if ai.seeded || live_before == 0 {
5141 ai.seeded = true;
5142 }
5143 }
5144 Ok(())
5145 }
5146
5147 fn fill_auto_inc_native_columns(
5148 &mut self,
5149 columns: &mut Vec<(u16, columnar::NativeColumn)>,
5150 n: usize,
5151 ) -> Result<()> {
5152 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
5153 return Ok(());
5154 };
5155 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
5156 if let Some(start) = self.alloc_auto_inc_range(n)? {
5157 columns.push((
5158 cid,
5159 columnar::NativeColumn::Int64 {
5160 data: (start..start.saturating_add(n as i64)).collect(),
5161 validity: vec![0xFF; n.div_ceil(8)],
5162 },
5163 ));
5164 }
5165 return Ok(());
5166 };
5167
5168 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
5169 return Err(MongrelError::InvalidArgument(format!(
5170 "AUTO_INCREMENT column {cid} must be Int64"
5171 )));
5172 };
5173 if data.len() < n {
5174 return Err(MongrelError::InvalidArgument(format!(
5175 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
5176 data.len()
5177 )));
5178 }
5179 if columnar::all_non_null(validity, n) {
5180 return Ok(());
5181 }
5182 if validity.iter().all(|b| *b == 0) {
5183 if let Some(start) = self.alloc_auto_inc_range(n)? {
5184 for (i, slot) in data.iter_mut().take(n).enumerate() {
5185 *slot = start.saturating_add(i as i64);
5186 }
5187 *validity = vec![0xFF; n.div_ceil(8)];
5188 }
5189 return Ok(());
5190 }
5191
5192 let new_validity = vec![0xFF; data.len().div_ceil(8)];
5193 for (i, slot) in data.iter_mut().enumerate().take(n) {
5194 if columnar::validity_bit(validity, i) {
5195 self.advance_auto_inc_past(*slot)?;
5196 } else {
5197 *slot = self.alloc_auto_inc_value()?;
5198 }
5199 }
5200 *validity = new_validity;
5201 Ok(())
5202 }
5203
5204 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
5218 self.ensure_writable()?;
5219 if self.auto_inc.is_none() {
5220 return Ok(None);
5221 }
5222 Ok(Some(self.alloc_auto_inc_value()?))
5223 }
5224
5225 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
5231 let payload = bincode::serialize(&rows)?;
5232 self.wal_append_data(Op::Put {
5233 table_id: self.table_id,
5234 rows: payload,
5235 })?;
5236 if self.is_shared() {
5237 self.pending_rows_auto_inc
5238 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
5239 self.pending_rows.extend(rows);
5240 } else {
5241 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
5242 }
5243 Ok(())
5244 }
5245
5246 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
5249 self.ensure_indexes_complete()
5250 }
5251
5252 pub(crate) fn prepare_durable_publish_controlled(
5253 &mut self,
5254 control: &crate::ExecutionControl,
5255 ) -> Result<()> {
5256 self.ensure_indexes_complete_controlled(control, || true)?;
5257 Ok(())
5258 }
5259
5260 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
5261 self.apply_put_rows_inner_prepared(rows, true);
5262 }
5263
5264 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
5265 if check_existing_pk {
5266 self.ensure_indexes_complete()?;
5267 }
5268 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
5269 Ok(())
5270 }
5271
5272 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
5276 if rows.len() == 1 {
5280 let row = rows.into_iter().next().expect("len checked");
5281 self.apply_put_row_single(row, check_existing_pk);
5282 return;
5283 }
5284 let pk_id = self.schema.primary_key().map(|c| c.id);
5301 let probe = match pk_id {
5302 Some(pid) => {
5303 check_existing_pk
5304 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
5305 }
5306 None => false,
5307 };
5308 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
5311 for r in rows {
5312 for &cid in r.columns.keys() {
5313 self.pending_put_cols.insert(cid);
5314 }
5315 let mut replaced_image: Option<Row> = None;
5316 match pk_id {
5317 Some(pid) if probe || maintain_pk_by_row => {
5318 if let Some(pk_val) = r.columns.get(&pid) {
5319 let key = self.index_lookup_key(pid, pk_val);
5320 if probe {
5321 if let Some(old) = self.recent_delete_preimages.remove(&key) {
5330 replaced_image = Some(old);
5331 if let Some(old_rid) = self.hot.get(&key) {
5332 if old_rid != r.row_id {
5333 self.tombstone_row(
5334 old_rid,
5335 r.committed_epoch,
5336 r.commit_ts,
5337 true,
5338 );
5339 }
5340 }
5341 } else if let Some(old_rid) = self.hot.get(&key) {
5342 if old_rid != r.row_id {
5343 replaced_image = self.get(old_rid, self.snapshot());
5344 self.tombstone_row(
5345 old_rid,
5346 r.committed_epoch,
5347 r.commit_ts,
5348 true,
5349 );
5350 }
5351 }
5352 }
5353 if maintain_pk_by_row {
5354 self.pk_by_row.insert(r.row_id, key);
5355 }
5356 }
5357 }
5358 Some(_) => {}
5359 None => {
5360 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5361 }
5362 }
5363 if let Some(old) = replaced_image {
5364 self.maintain_indexes_on_pk_replace(&old, &r);
5365 } else {
5366 self.index_row(&r);
5367 }
5368 self.reservoir.offer(r.row_id.0);
5369 self.memtable.upsert(r);
5370 self.live_count = self.live_count.saturating_add(1);
5373 }
5374 self.data_generation = self.data_generation.wrapping_add(1);
5375 }
5376
5377 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
5387 for &cid in row.columns.keys() {
5388 self.pending_put_cols.insert(cid);
5389 }
5390 let epoch = row.committed_epoch;
5391 let mut replaced_image: Option<Row> = None;
5392 if let Some(pk_col) = self.schema.primary_key() {
5393 let pk_id = pk_col.id;
5394 if let Some(pk_val) = row.columns.get(&pk_id) {
5395 let maintain_pk_by_row = self.pk_by_row_complete;
5399 if check_existing_pk || maintain_pk_by_row {
5400 let key = self.index_lookup_key(pk_id, pk_val);
5401 if check_existing_pk {
5402 if let Some(old) = self.recent_delete_preimages.remove(&key) {
5411 replaced_image = Some(old);
5412 if let Some(old_rid) = self.hot.get(&key) {
5413 if old_rid != row.row_id {
5414 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5415 }
5416 }
5417 } else if let Some(old_rid) = self.hot.get(&key) {
5418 if old_rid != row.row_id {
5419 replaced_image = self.get(old_rid, self.snapshot());
5423 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5424 }
5425 }
5426 }
5427 if maintain_pk_by_row {
5428 self.pk_by_row.insert(row.row_id, key);
5429 }
5430 }
5431 }
5432 } else {
5433 self.hot
5434 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
5435 }
5436 if let Some(old) = replaced_image {
5437 self.maintain_indexes_on_pk_replace(&old, &row);
5438 } else {
5439 self.index_row(&row);
5440 }
5441 self.reservoir.offer(row.row_id.0);
5442 self.memtable.upsert(row);
5443 self.live_count = self.live_count.saturating_add(1);
5444 self.data_generation = self.data_generation.wrapping_add(1);
5445 }
5446
5447 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
5450 let has_partial = self
5451 .schema
5452 .indexes
5453 .iter()
5454 .any(|idx| idx.predicate.is_some());
5455 if has_partial {
5456 self.unindex_bitmap_membership(old);
5460 self.index_row(new);
5461 return;
5462 }
5463
5464 crate::index::maintain_bitmap_secondary_on_replace(
5465 &self.schema,
5466 &mut self.bitmap,
5467 old,
5468 new,
5469 );
5470 self.index_row_non_bitmap(new);
5471 }
5472
5473 fn index_row_non_bitmap(&mut self, row: &Row) {
5476 if row.deleted {
5477 return;
5478 }
5479 let effective = if self.column_keys.is_empty() {
5480 None
5481 } else {
5482 Some(self.tokenized_for_indexes(row))
5483 };
5484 let source = effective.as_ref().unwrap_or(row);
5485 for idef in &self.schema.indexes {
5486 if idef.kind == crate::schema::IndexKind::Bitmap {
5487 continue;
5488 }
5489 index_into_single(
5490 idef,
5491 &self.schema,
5492 source,
5493 &mut self.hot,
5494 &mut self.bitmap,
5495 &mut self.ann,
5496 &mut self.fm,
5497 &mut self.sparse,
5498 &mut self.minhash,
5499 );
5500 }
5501 if let Some(pk_col) = self.schema.primary_key() {
5502 if let Some(pk_val) = source.columns.get(&pk_col.id) {
5503 let key = if self.column_keys.is_empty() {
5504 pk_val.encode_key()
5505 } else {
5506 self.index_lookup_key(pk_col.id, pk_val)
5507 };
5508 self.hot.insert(key, source.row_id);
5509 }
5510 }
5511 }
5512
5513 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
5516 self.allocator.alloc()
5517 }
5518
5519 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
5525 let pk = self.schema.primary_key().ok_or_else(|| {
5526 MongrelError::Schema("clustered table requires a single-column primary key".into())
5527 })?;
5528 let pk_val = columns
5529 .iter()
5530 .find(|(id, _)| *id == pk.id)
5531 .map(|(_, v)| v)
5532 .ok_or_else(|| {
5533 MongrelError::Schema(format!(
5534 "clustered table missing primary key column {} ({})",
5535 pk.id, pk.name
5536 ))
5537 })?;
5538 Ok(clustered_row_id(pk_val))
5539 }
5540
5541 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
5549 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
5550 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5551 }
5552 let n = rows.len();
5553 for r in rows {
5554 for &cid in r.columns.keys() {
5555 self.pending_put_cols.insert(cid);
5556 }
5557 }
5558 let (losers, winner_pks) = self.partition_pk_winners(rows);
5559 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
5560 let group_ts = rows.first().and_then(|r| r.commit_ts);
5562 for (key, &row_id) in &winner_pks {
5563 if let Some(old_rid) = self.hot.get(key) {
5564 if old_rid != row_id {
5565 self.tombstone_row(old_rid, epoch, group_ts, true);
5566 }
5567 }
5568 }
5569 for &loser_rid in &losers {
5572 self.tombstone_row(loser_rid, epoch, group_ts, false);
5573 }
5574 for (key, row_id) in winner_pks {
5576 self.insert_hot_pk(key, row_id);
5577 }
5578 if self.schema.primary_key().is_none() {
5579 for r in rows {
5580 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5581 }
5582 }
5583 for r in rows {
5584 self.allocator.advance_to(r.row_id)?;
5585 if !losers.contains(&r.row_id) {
5586 self.index_row(r);
5587 }
5588 }
5589 for r in rows {
5590 if !losers.contains(&r.row_id) {
5591 self.reservoir.offer(r.row_id.0);
5592 }
5593 }
5594 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
5595 self.data_generation = self.data_generation.wrapping_add(1);
5596 Ok(())
5597 }
5598
5599 pub(crate) fn recover_apply(
5604 &mut self,
5605 rows: Vec<Row>,
5606 deletes: Vec<(RowId, Epoch)>,
5607 ) -> Result<()> {
5608 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
5612 std::collections::BTreeMap::new();
5613 for row in rows {
5614 if row.row_id.0 >= u64::MAX - 1 {
5615 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5616 }
5617 self.allocator.advance_to(row.row_id)?;
5618 if let Some(ai) = self.auto_inc.as_mut() {
5623 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
5624 let next = n.checked_add(1).ok_or_else(|| {
5625 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
5626 })?;
5627 if next > ai.next {
5628 ai.next = next;
5629 }
5630 }
5631 }
5632 by_epoch.entry(row.committed_epoch).or_default().push(row);
5633 }
5634 for (epoch, group) in by_epoch {
5635 let (losers, winner_pks) = self.partition_pk_winners(&group);
5636 let group_ts = group.first().and_then(|r| r.commit_ts);
5638 for (key, &row_id) in &winner_pks {
5639 if let Some(old_rid) = self.hot.get(key) {
5640 if old_rid != row_id {
5641 self.tombstone_row(old_rid, epoch, group_ts, false);
5642 }
5643 }
5644 }
5645 for (key, row_id) in winner_pks {
5646 self.insert_hot_pk(key, row_id);
5647 }
5648 if self.schema.primary_key().is_none() {
5649 for r in &group {
5650 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5651 }
5652 }
5653 for r in &group {
5654 if !losers.contains(&r.row_id) {
5655 self.memtable.upsert(r.clone());
5656 self.index_row(r);
5657 }
5658 }
5659 }
5660 for (rid, epoch) in deletes {
5661 self.memtable.tombstone(rid, epoch);
5662 self.remove_hot_for_row(rid, epoch);
5663 }
5664 self.reservoir_complete = false;
5667 Ok(())
5668 }
5669
5670 pub(crate) fn flushed_epoch(&self) -> u64 {
5672 self.flushed_epoch
5673 }
5674
5675 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
5676 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
5677 }
5678
5679 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
5681 self.schema.validate_values(cells)
5682 }
5683
5684 fn validate_columns_not_null(
5688 &self,
5689 columns: &[(u16, columnar::NativeColumn)],
5690 n: usize,
5691 ) -> Result<()> {
5692 let by_id: HashMap<u16, &columnar::NativeColumn> =
5693 columns.iter().map(|(id, c)| (*id, c)).collect();
5694 for col in &self.schema.columns {
5695 if !col.flags.contains(ColumnFlags::NULLABLE) {
5696 match by_id.get(&col.id) {
5697 None => {
5698 return Err(MongrelError::InvalidArgument(format!(
5699 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
5700 col.name, col.id
5701 )));
5702 }
5703 Some(c) => {
5704 if c.null_count(n) != 0 {
5705 return Err(MongrelError::InvalidArgument(format!(
5706 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
5707 col.name, col.id
5708 )));
5709 }
5710 }
5711 }
5712 }
5713 if let TypeId::Enum { variants } = &col.ty {
5714 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
5715 if by_id.contains_key(&col.id) {
5716 return Err(MongrelError::InvalidArgument(format!(
5717 "column '{}' ({}) enum requires a bytes column",
5718 col.name, col.id
5719 )));
5720 }
5721 continue;
5722 };
5723 for index in 0..n {
5724 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
5725 continue;
5726 };
5727 if !variants.iter().any(|variant| variant.as_bytes() == value) {
5728 return Err(MongrelError::InvalidArgument(format!(
5729 "column '{}' ({}) enum value {:?} is not one of {:?}",
5730 col.name,
5731 col.id,
5732 String::from_utf8_lossy(value),
5733 variants
5734 )));
5735 }
5736 }
5737 }
5738 }
5739 Ok(())
5740 }
5741
5742 fn bulk_pk_winner_indices(
5747 &self,
5748 columns: &[(u16, columnar::NativeColumn)],
5749 n: usize,
5750 ) -> Option<Vec<usize>> {
5751 let pk_col = self.schema.primary_key()?;
5752 let pk_id = pk_col.id;
5753 let pk_ty = pk_col.ty.clone();
5754 let by_id: HashMap<u16, &columnar::NativeColumn> =
5755 columns.iter().map(|(id, c)| (*id, c)).collect();
5756 let pk_native = by_id.get(&pk_id)?;
5757 if native_int64_strictly_increasing(pk_native, n) {
5758 return None;
5759 }
5760 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
5762 let mut null_pk_rows: Vec<usize> = Vec::new();
5763 for i in 0..n {
5764 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
5765 Some(key) => {
5766 last.insert(key, i);
5767 }
5768 None => null_pk_rows.push(i),
5769 }
5770 }
5771 let mut winners: HashSet<usize> = last.values().copied().collect();
5772 for i in null_pk_rows {
5773 winners.insert(i);
5774 }
5775 Some((0..n).filter(|i| winners.contains(i)).collect())
5776 }
5777
5778 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
5780 self.require_delete()?;
5781 let epoch = self.pending_epoch();
5782 self.wal_append_data(Op::Delete {
5783 table_id: self.table_id,
5784 row_ids: vec![row_id],
5785 })?;
5786 if self.is_shared() {
5787 self.pending_dels.push(row_id);
5788 } else {
5789 self.apply_delete(row_id, epoch);
5790 }
5791 Ok(())
5792 }
5793
5794 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
5795 let pre = self.get(row_id, self.snapshot());
5796 self.delete(row_id)?;
5797 Ok(pre.map(|row| {
5798 let mut columns: Vec<_> = row.columns.into_iter().collect();
5799 columns.sort_by_key(|(id, _)| *id);
5800 OwnedRow { columns }
5801 }))
5802 }
5803
5804 pub fn truncate(&mut self) -> Result<()> {
5806 self.require_delete()?;
5807 let epoch = self.pending_epoch();
5808 self.wal_append_data(Op::TruncateTable {
5809 table_id: self.table_id,
5810 })?;
5811 self.pending_rows.clear();
5812 self.pending_rows_auto_inc.clear();
5813 self.pending_dels.clear();
5814 self.pending_truncate = Some(epoch);
5815 Ok(())
5816 }
5817
5818 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
5820 self.run_refs.clear();
5825 self.retiring.clear();
5826 self.memtable = Memtable::new();
5827 self.mutable_run = MutableRun::new();
5828 self.hot = HotIndex::new();
5829 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
5830 self.bitmap = bitmap;
5831 self.ann = ann;
5832 self.fm = fm;
5833 self.sparse = sparse;
5834 self.minhash = minhash;
5835 self.learned_range = Arc::new(HashMap::new());
5836 self.pk_by_row.clear();
5837 self.pk_by_row_complete = false;
5838 self.live_count = 0;
5839 self.reservoir = crate::reservoir::Reservoir::default();
5840 self.reservoir_complete = true;
5841 self.had_deletes = true;
5842 self.agg_cache = Arc::new(HashMap::new());
5843 self.global_idx_epoch = 0;
5844 self.indexes_complete = true;
5845 self.pending_delete_rids.clear();
5846 self.pending_put_cols.clear();
5847 self.pending_rows.clear();
5848 self.pending_rows_auto_inc.clear();
5849 self.pending_dels.clear();
5850 self.clear_result_cache();
5851 self.invalidate_index_checkpoint();
5852 self.data_generation = self.data_generation.wrapping_add(1);
5853 }
5854
5855 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
5858 self.apply_delete_at(row_id, epoch, None);
5859 }
5860
5861 pub(crate) fn apply_delete_at(
5863 &mut self,
5864 row_id: RowId,
5865 epoch: Epoch,
5866 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5867 ) {
5868 let preimage = self.get(row_id, self.snapshot());
5878 if let Some(row) = preimage {
5879 if let Some(pk_col) = self.schema.primary_key() {
5880 if let Some(pk_val) = row.columns.get(&pk_col.id) {
5881 let key = self.index_lookup_key(pk_col.id, pk_val);
5882 self.recent_delete_preimages.insert(key, row);
5883 }
5884 }
5885 }
5886 self.tombstone_row(row_id, epoch, commit_ts, true);
5887 self.data_generation = self.data_generation.wrapping_add(1);
5888 }
5889
5890 fn unindex_bitmap_membership(&mut self, row: &Row) {
5895 if row.deleted {
5896 return;
5897 }
5898 for idef in &self.schema.indexes {
5899 if idef.kind != crate::schema::IndexKind::Bitmap {
5900 continue;
5901 }
5902 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
5903 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5904 b.remove(&key, row.row_id);
5905 }
5906 }
5907 }
5908 }
5909
5910 fn union_bitmap_point_i64(
5914 &self,
5915 set: &mut RowIdSet,
5916 column_id: u16,
5917 value: i64,
5918 snapshot: Snapshot,
5919 ) {
5920 let Some(b) = self.bitmap.get(&column_id) else {
5921 return;
5922 };
5923 let encoded = Value::Int64(value).encode_key();
5924 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
5925 for rid in b.get(&lookup).iter() {
5926 set.insert(u64::from(rid));
5927 }
5928 set.remove_many(self.overlay_tombstoned_rids(snapshot));
5932 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
5933 }
5934
5935 fn tombstone_row(
5949 &mut self,
5950 row_id: RowId,
5951 epoch: Epoch,
5952 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5953 adjust_live_count: bool,
5954 ) {
5955 let prev_was_live = if adjust_live_count {
5956 match self.memtable.get_version(row_id, epoch) {
5957 Some((_, prev)) => !prev.deleted,
5958 None => true,
5959 }
5960 } else {
5961 false
5962 };
5963 let tombstone = Row {
5964 row_id,
5965 committed_epoch: epoch,
5966 columns: std::collections::HashMap::new(),
5967 deleted: true,
5968 commit_ts,
5969 };
5970 self.memtable.upsert(tombstone);
5971 self.pk_by_row.remove(&row_id);
5972 if prev_was_live {
5973 self.live_count = self.live_count.saturating_sub(1);
5974 }
5975 self.pending_delete_rids.insert(row_id.0 as u32);
5977 self.had_deletes = true;
5980 self.agg_cache = Arc::new(HashMap::new());
5981 }
5982
5983 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
5987 let Some(pk_col) = self.schema.primary_key() else {
5988 return;
5989 };
5990 if self.pk_by_row_complete {
5993 if let Some(key) = self.pk_by_row.remove(&row_id) {
5994 if self.hot.get(&key) == Some(row_id) {
5995 self.hot.remove(&key);
5996 }
5997 }
5998 return;
5999 }
6000 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
6019 if self.indexes_complete {
6020 let pk_val = self
6021 .memtable
6022 .get_version(row_id, lookup_epoch)
6023 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
6024 .or_else(|| {
6025 self.mutable_run
6026 .get_version(row_id, lookup_epoch)
6027 .filter(|(_, r)| !r.deleted)
6028 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
6029 })
6030 .or_else(|| {
6031 self.run_refs.iter().find_map(|rr| {
6032 let mut reader = self.open_reader(rr.run_id).ok()?;
6033 let (_, deleted, val) = reader
6034 .get_version_column(row_id, lookup_epoch, pk_col.id)
6035 .ok()??;
6036 if deleted {
6037 return None;
6038 }
6039 val
6040 })
6041 });
6042 if let Some(pk_val) = pk_val {
6043 let key = self.index_lookup_key(pk_col.id, &pk_val);
6044 if self.hot.get(&key) == Some(row_id) {
6045 self.hot.remove(&key);
6046 }
6047 return;
6048 }
6049 }
6050 self.refresh_pk_by_row_from_hot();
6055 if let Some(key) = self.pk_by_row.remove(&row_id) {
6056 if self.hot.get(&key) == Some(row_id) {
6057 self.hot.remove(&key);
6058 }
6059 }
6060 }
6061
6062 fn partition_pk_winners(
6067 &self,
6068 rows: &[Row],
6069 ) -> (
6070 std::collections::HashSet<RowId>,
6071 std::collections::HashMap<Vec<u8>, RowId>,
6072 ) {
6073 let mut losers = std::collections::HashSet::new();
6074 let Some(pk_col) = self.schema.primary_key() else {
6075 return (losers, std::collections::HashMap::new());
6076 };
6077 let pk_id = pk_col.id;
6078 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
6079 std::collections::HashMap::new();
6080 for r in rows {
6081 let Some(pk_val) = r.columns.get(&pk_id) else {
6082 continue;
6083 };
6084 let key = self.index_lookup_key(pk_id, pk_val);
6085 if let Some(&old_rid) = winners.get(&key) {
6086 losers.insert(old_rid);
6087 }
6088 winners.insert(key, r.row_id);
6089 }
6090 (losers, winners)
6091 }
6092
6093 fn index_row(&mut self, row: &Row) {
6094 if row.deleted {
6095 return;
6096 }
6097 let any_predicate = self
6105 .schema
6106 .indexes
6107 .iter()
6108 .any(|idx| idx.predicate.is_some());
6109 if any_predicate {
6110 let columns_map: HashMap<u16, &Value> =
6111 row.columns.iter().map(|(k, v)| (*k, v)).collect();
6112 let name_to_id: HashMap<&str, u16> = self
6113 .schema
6114 .columns
6115 .iter()
6116 .map(|c| (c.name.as_str(), c.id))
6117 .collect();
6118 for idx in &self.schema.indexes {
6119 if let Some(pred) = &idx.predicate {
6120 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
6121 continue; }
6123 }
6124 index_into_single(
6126 idx,
6127 &self.schema,
6128 row,
6129 &mut self.hot,
6130 &mut self.bitmap,
6131 &mut self.ann,
6132 &mut self.fm,
6133 &mut self.sparse,
6134 &mut self.minhash,
6135 );
6136 }
6137 return;
6138 }
6139 if self.column_keys.is_empty() {
6143 index_into(
6144 &self.schema,
6145 row,
6146 &mut self.hot,
6147 &mut self.bitmap,
6148 &mut self.ann,
6149 &mut self.fm,
6150 &mut self.sparse,
6151 &mut self.minhash,
6152 );
6153 return;
6154 }
6155 let effective_row = self.tokenized_for_indexes(row);
6156 index_into(
6157 &self.schema,
6158 &effective_row,
6159 &mut self.hot,
6160 &mut self.bitmap,
6161 &mut self.ann,
6162 &mut self.fm,
6163 &mut self.sparse,
6164 &mut self.minhash,
6165 );
6166 }
6167
6168 fn tokenized_for_indexes(&self, row: &Row) -> Row {
6174 if self.column_keys.is_empty() {
6175 return row.clone();
6176 }
6177 {
6178 use crate::encryption::SCHEME_HMAC_EQ;
6179 let mut tok = row.clone();
6180 for (&cid, &(_, scheme)) in &self.column_keys {
6181 if scheme != SCHEME_HMAC_EQ {
6182 continue;
6183 }
6184 if let Some(v) = tok.columns.get(&cid).cloned() {
6185 if let Some(t) = self.tokenize_value(cid, &v) {
6186 tok.columns.insert(cid, t);
6187 }
6188 }
6189 }
6190 tok
6191 }
6192 }
6193
6194 pub fn commit(&mut self) -> Result<Epoch> {
6199 self.commit_inner(None)
6200 }
6201
6202 #[doc(hidden)]
6205 pub fn commit_controlled<F>(
6206 &mut self,
6207 control: &crate::ExecutionControl,
6208 mut before_commit: F,
6209 ) -> Result<Epoch>
6210 where
6211 F: FnMut() -> Result<()>,
6212 {
6213 self.commit_inner(Some((control, &mut before_commit)))
6214 }
6215
6216 fn commit_inner(
6217 &mut self,
6218 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6219 ) -> Result<Epoch> {
6220 self.ensure_writable()?;
6221 if !self.has_pending_mutations() {
6222 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
6223 return Err(MongrelError::Full(
6224 "standalone transaction id namespace exhausted".into(),
6225 ));
6226 }
6227 return Ok(self.epoch.visible());
6228 }
6229 self.commit_new_epoch_inner(controlled)
6230 }
6231
6232 fn commit_new_epoch(&mut self) -> Result<Epoch> {
6236 self.commit_new_epoch_inner(None)
6237 }
6238
6239 fn commit_new_epoch_inner(
6240 &mut self,
6241 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6242 ) -> Result<Epoch> {
6243 self.ensure_writable()?;
6244 if self.is_shared() {
6245 self.commit_shared(controlled)
6246 } else {
6247 self.commit_private(controlled)
6248 }
6249 }
6250
6251 fn commit_private(
6253 &mut self,
6254 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6255 ) -> Result<Epoch> {
6256 let commit_lock = Arc::clone(&self.commit_lock);
6260 let _g = commit_lock.lock();
6261 let txn_id = self.ensure_txn_id()?;
6264 if let Some((control, before_commit)) = controlled {
6265 control.checkpoint()?;
6266 before_commit()?;
6267 }
6268 let new_epoch = self.epoch.bump_assigned();
6269 let epoch_authority = Arc::clone(&self.epoch);
6270 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
6271 let wal_result = match &mut self.wal {
6275 WalSink::Private(w) => w
6276 .append_txn(
6277 txn_id,
6278 Op::TxnCommit {
6279 epoch: new_epoch.0,
6280 added_runs: Vec::new(),
6281 },
6282 )
6283 .and_then(|_| w.sync()),
6284 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
6285 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
6286 };
6287 if let Err(error) = wal_result {
6288 self.durable_commit_failed = true;
6289 return Err(MongrelError::CommitOutcomeUnknown {
6290 epoch: new_epoch.0,
6291 message: error.to_string(),
6292 });
6293 }
6294 if let Some(epoch) = self.pending_truncate.take() {
6297 self.apply_truncate(epoch);
6298 }
6299 self.invalidate_pending_cache();
6300 let publish_result = self.persist_manifest(new_epoch);
6301 self.epoch.publish_in_order(new_epoch);
6305 epoch_guard.disarm();
6306 if let Err(error) = publish_result {
6307 self.durable_commit_failed = true;
6308 return Err(MongrelError::DurableCommit {
6309 epoch: new_epoch.0,
6310 message: error.to_string(),
6311 });
6312 }
6313 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
6314 self.pending_private_mutations = false;
6315 self.data_generation = self.data_generation.wrapping_add(1);
6316 Ok(new_epoch)
6317 }
6318
6319 fn commit_shared(
6325 &mut self,
6326 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6327 ) -> Result<Epoch> {
6328 use std::sync::atomic::Ordering;
6329 let s = match &self.wal {
6330 WalSink::Shared(s) => s.clone(),
6331 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
6332 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
6333 };
6334 if s.poisoned.load(Ordering::Relaxed) {
6335 return Err(MongrelError::Other(
6336 "database poisoned by fsync error".into(),
6337 ));
6338 }
6339 let commit_lock = Arc::clone(&self.commit_lock);
6346 let _g = commit_lock.lock();
6347 if !self.pending_rows.is_empty() {
6348 match controlled.as_ref() {
6349 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
6350 None => self.prepare_durable_publish()?,
6351 }
6352 }
6353 let txn_id = self.ensure_txn_id()?;
6356 let mut wal = s.wal.lock();
6357 if let Some((control, before_commit)) = controlled {
6358 control.checkpoint()?;
6359 before_commit()?;
6360 }
6361 let new_epoch = self.epoch.bump_assigned();
6362 let epoch_authority = Arc::clone(&self.epoch);
6363 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
6364 let commit_ts = s.hlc.now().map_err(|skew| {
6367 MongrelError::Other(format!(
6368 "clock skew rejected commit timestamp allocation: {skew}"
6369 ))
6370 })?;
6371 let commit_seq = match wal.append_commit_at(
6372 txn_id,
6373 new_epoch,
6374 &[],
6375 commit_ts.physical_micros.saturating_mul(1_000),
6376 ) {
6377 Ok(commit_seq) => commit_seq,
6378 Err(error) => {
6379 s.poisoned.store(true, Ordering::Relaxed);
6380 s.lifecycle.poison();
6381 return Err(MongrelError::CommitOutcomeUnknown {
6382 epoch: new_epoch.0,
6383 message: error.to_string(),
6384 });
6385 }
6386 };
6387 drop(wal);
6388 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
6389 s.poisoned.store(true, Ordering::Relaxed);
6390 s.lifecycle.poison();
6391 return Err(MongrelError::CommitOutcomeUnknown {
6392 epoch: new_epoch.0,
6393 message: error.to_string(),
6394 });
6395 }
6396
6397 if self.pending_truncate.take().is_some() {
6400 self.apply_truncate(new_epoch);
6401 }
6402 let mut rows = std::mem::take(&mut self.pending_rows);
6403 if !rows.is_empty() {
6404 for r in &mut rows {
6405 r.committed_epoch = new_epoch;
6406 r.commit_ts = Some(commit_ts);
6407 }
6408 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
6409 let all_auto_generated =
6410 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
6411 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
6412 } else {
6413 self.pending_rows_auto_inc.clear();
6414 }
6415 let dels = std::mem::take(&mut self.pending_dels);
6416 for rid in dels {
6417 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
6418 }
6419
6420 self.invalidate_pending_cache();
6421 let publish_result = self.persist_manifest(new_epoch);
6422 self.epoch.publish_in_order(new_epoch);
6423 epoch_guard.disarm();
6424 let _ = s.change_wake.send(());
6425 if let Err(error) = publish_result {
6426 self.durable_commit_failed = true;
6427 s.poisoned.store(true, Ordering::Relaxed);
6428 s.lifecycle.poison();
6429 return Err(MongrelError::DurableCommit {
6430 epoch: new_epoch.0,
6431 message: error.to_string(),
6432 });
6433 }
6434 self.current_txn_id = 0;
6436 self.data_generation = self.data_generation.wrapping_add(1);
6437 Ok(new_epoch)
6438 }
6439
6440 pub fn flush(&mut self) -> Result<Epoch> {
6448 self.flush_with_outcome().map(|(epoch, _)| epoch)
6449 }
6450
6451 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
6453 self.flush_with_outcome_inner(None)
6454 }
6455
6456 #[doc(hidden)]
6459 pub fn flush_with_outcome_controlled<F>(
6460 &mut self,
6461 control: &crate::ExecutionControl,
6462 mut before_commit: F,
6463 ) -> Result<(Epoch, bool)>
6464 where
6465 F: FnMut() -> Result<()>,
6466 {
6467 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
6468 }
6469
6470 fn flush_with_outcome_inner(
6471 &mut self,
6472 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6473 ) -> Result<(Epoch, bool)> {
6474 match controlled.as_ref() {
6475 Some((control, _)) => {
6476 self.ensure_indexes_complete_controlled(control, || true)?;
6477 }
6478 None => self.ensure_indexes_complete()?,
6479 }
6480 let committed = self.has_pending_mutations();
6481 let epoch = self.commit_inner(controlled)?;
6482 let finish: Result<(Epoch, bool)> = (|| {
6483 let rows = self.memtable.drain_sorted();
6484 if !rows.is_empty() {
6485 self.mutable_run.insert_many(rows);
6486 }
6487 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
6488 self.spill_mutable_run(epoch)?;
6489 self.mark_flushed(epoch)?;
6493 self.persist_manifest(epoch)?;
6494 self.build_learned_ranges()?;
6495 self.checkpoint_indexes(epoch);
6498 let _ = self.publish_run_lookup_directory();
6503 }
6504 Ok((epoch, committed))
6507 })();
6508 let outcome = match finish {
6509 Err(error) if committed => Err(MongrelError::DurableCommit {
6510 epoch: epoch.0,
6511 message: error.to_string(),
6512 }),
6513 result => result,
6514 };
6515 if outcome.is_ok() {
6516 let _ = self.publish_read_generation();
6522 }
6523 outcome
6524 }
6525
6526 fn has_pending_mutations(&self) -> bool {
6527 self.pending_private_mutations
6528 || !self.pending_rows.is_empty()
6529 || !self.pending_dels.is_empty()
6530 || self.pending_truncate.is_some()
6531 }
6532
6533 pub fn has_pending_writes(&self) -> bool {
6534 self.has_pending_mutations()
6535 }
6536
6537 pub fn force_flush(&mut self) -> Result<Epoch> {
6546 let saved = self.mutable_run_spill_bytes;
6547 self.mutable_run_spill_bytes = 1;
6548 let result = self.flush();
6549 self.mutable_run_spill_bytes = saved;
6550 result
6551 }
6552
6553 pub fn close(&mut self) -> Result<()> {
6560 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
6561 self.force_flush()?;
6562 }
6563 self.result_cache
6569 .lock()
6570 .shutdown_persistent_cache(std::time::Duration::from_secs(5));
6571 Ok(())
6572 }
6573
6574 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
6581 let op = Op::Flush {
6582 table_id: self.table_id,
6583 flushed_epoch: epoch.0,
6584 };
6585 match &mut self.wal {
6586 WalSink::Private(w) => {
6587 w.append_system(op)?;
6588 w.sync()?;
6589 }
6590 WalSink::Shared(s) => {
6591 s.wal.lock().append_system(op)?;
6596 }
6597 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
6598 }
6599 self.flushed_epoch = epoch.0;
6600 if matches!(self.wal, WalSink::Private(_)) {
6601 self.rotate_wal(epoch)?;
6602 }
6603 Ok(())
6604 }
6605
6606 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
6610 if self.mutable_run.is_empty() {
6611 return Ok(());
6612 }
6613 let run_id = self.alloc_run_id()?;
6614 let rows = self.mutable_run.drain_sorted();
6615 if rows.is_empty() {
6616 return Ok(());
6617 }
6618 let path = self.run_path(run_id);
6619 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
6620 if let Some(kek) = &self.kek {
6621 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6622 }
6623 let header = match self.create_run_file(run_id)? {
6624 Some(file) => writer.write_file(file, &rows)?,
6625 None => writer.write(&path, &rows)?,
6626 };
6627 self.run_refs.push(RunRef {
6628 run_id: run_id as u128,
6629 level: 0,
6630 epoch_created: epoch.0,
6631 row_count: header.row_count,
6632 });
6633 self.run_row_id_ranges
6634 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6635 Ok(())
6636 }
6637
6638 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
6642 self.mutable_run_spill_bytes = bytes.max(1);
6643 }
6644
6645 pub fn set_compaction_zstd_level(&mut self, level: i32) {
6649 self.compaction_zstd_level = level;
6650 }
6651
6652 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
6656 self.result_cache.lock().set_max_bytes(max_bytes);
6657 }
6658
6659 pub fn flush_persistent_cache(&self, deadline_ms: u64) -> u64 {
6664 self.result_cache
6665 .lock()
6666 .flush_persistent_cache(std::time::Duration::from_millis(deadline_ms))
6667 }
6668
6669 pub fn shutdown_persistent_cache(&mut self, deadline_ms: u64) {
6674 self.result_cache
6675 .lock()
6676 .shutdown_persistent_cache(std::time::Duration::from_millis(deadline_ms));
6677 }
6678
6679 #[doc(hidden)]
6684 pub fn _set_persist_min_bytes_for_test(&mut self, min: u64) {
6685 self.result_cache.lock().set_persist_min_bytes(min);
6686 }
6687
6688 #[doc(hidden)]
6695 pub fn _force_persistent_worker_spawn_failure_for_test(enabled: bool) {
6696 PERSISTENT_WORKER_SPAWN_FAILURE.with(|flag| flag.set(enabled));
6697 }
6698
6699 #[doc(hidden)]
6706 pub fn _persist_cached_entry_synchronously_for_maintenance(
6707 &mut self,
6708 q: &crate::query::Query,
6709 ) -> bool {
6710 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
6711 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
6712 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
6713 let identity = PersistentCacheIdentity {
6714 table_id: self.table_id(),
6715 schema_id: self.schema.schema_id,
6716 logical_generation: self.current_epoch().0,
6717 };
6718 let mut cache = self.result_cache.lock();
6719 let Some(entry) = cache.entries.get(&key).map(|e| e.clone_for_persist()) else {
6720 return false;
6721 };
6722 let entry_generation = cache.allocate_persist_generation(key);
6723 let context = PersistContext {
6724 identity,
6725 key,
6726 entry_generation,
6727 };
6728 cache.persist_entry_synchronously_for_maintenance(context, &entry);
6729 true
6730 }
6731
6732 pub(crate) fn clear_result_cache(&mut self) {
6736 self.result_cache.lock().clear();
6737 }
6738
6739 pub fn mutable_run_len(&self) -> usize {
6741 self.mutable_run.len()
6742 }
6743
6744 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
6747 self.mutable_run.drain_sorted()
6748 }
6749
6750 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
6752 let mut snapshot = self.mutable_run.clone();
6753 snapshot.drain_sorted()
6754 }
6755
6756 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
6761 self.ensure_writable()?;
6762 let n = batch.len();
6763 if n == 0 {
6764 return Ok(self.current_epoch());
6765 }
6766 for row in &batch {
6767 self.schema.validate_values(row)?;
6768 }
6769 let epoch = self.commit_new_epoch()?;
6770 let live_before = self.live_count;
6771 self.spill_mutable_run(epoch)?;
6775 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6776 && self.indexes_complete
6777 && self.run_refs.is_empty()
6778 && self.memtable.is_empty()
6779 && self.mutable_run.is_empty();
6780 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
6786 use rayon::prelude::*;
6787 self.schema
6788 .columns
6789 .par_iter()
6790 .map(|cdef| {
6791 (
6792 cdef.id,
6793 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
6794 )
6795 })
6796 .collect::<Vec<_>>()
6797 };
6798 drop(batch);
6799 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6804 self.validate_columns_not_null(&user_columns, n)?;
6805 let winner_idx = self
6806 .bulk_pk_winner_indices(&user_columns, n)
6807 .filter(|idx| idx.len() != n);
6808 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6809 match winner_idx.as_deref() {
6810 Some(idx) => {
6811 let compacted = user_columns
6812 .iter()
6813 .map(|(id, c)| (*id, c.gather(idx)))
6814 .collect();
6815 (compacted, idx.len())
6816 }
6817 None => (std::mem::take(&mut user_columns), n),
6818 };
6819 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6820 let first = self.allocator.alloc_range(write_n as u64)?.0;
6821 for rid in first..first + write_n as u64 {
6822 self.reservoir.offer(rid);
6823 }
6824 let run_id = self.alloc_run_id()?;
6825 let path = self.run_path(run_id);
6826 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
6827 .clean(true)
6828 .with_lz4()
6829 .with_native_endian();
6830 if let Some(kek) = &self.kek {
6831 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6832 }
6833 let header = match self.create_run_file(run_id)? {
6834 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
6835 None => writer.write_native(&path, &write_columns, write_n, first)?,
6836 };
6837 self.run_refs.push(RunRef {
6838 run_id: run_id as u128,
6839 level: 0,
6840 epoch_created: epoch.0,
6841 row_count: header.row_count,
6842 });
6843 self.run_row_id_ranges
6844 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6845 self.live_count = self.live_count.saturating_add(write_n as u64);
6846 if eager_index_build {
6847 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6848 self.index_columns_bulk(&write_columns, &row_ids);
6849 self.indexes_complete = true;
6850 self.build_learned_ranges()?;
6851 } else {
6852 self.indexes_complete = false;
6853 }
6854 self.mark_flushed(epoch)?;
6855 self.persist_manifest(epoch)?;
6856 if eager_index_build {
6857 self.checkpoint_indexes(epoch);
6858 }
6859 self.clear_result_cache();
6860 Ok(epoch)
6861 }
6862
6863 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
6866 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
6867 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
6868 let segment_no = segment
6871 .file_stem()
6872 .and_then(|s| s.to_str())
6873 .and_then(|s| s.strip_prefix("seg-"))
6874 .and_then(|s| s.parse::<u64>().ok())
6875 .unwrap_or(0);
6876 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
6877 wal.set_sync_byte_threshold(self.sync_byte_threshold);
6878 wal.sync()?;
6879 self.wal = WalSink::Private(wal);
6880 Ok(())
6881 }
6882
6883 pub(crate) fn invalidate_pending_cache(&mut self) {
6888 self.result_cache
6889 .lock()
6890 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
6891 self.pending_delete_rids.clear();
6892 self.pending_put_cols.clear();
6893 }
6894
6895 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
6896 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
6897 m.current_epoch = epoch.0;
6898 m.next_row_id = self.allocator.current().0;
6899 m.runs = self.run_refs.clone();
6900 m.live_count = self.live_count;
6901 m.global_idx_epoch = self.global_idx_epoch;
6902 m.flushed_epoch = self.flushed_epoch;
6903 m.retiring = self.retiring.clone();
6904 m.auto_inc_next = match self.auto_inc {
6908 Some(ai) if ai.seeded => ai.next,
6909 _ => 0,
6910 };
6911 m.ttl = self.ttl;
6912 let meta_dek = self.manifest_meta_dek();
6913 match self._root_guard.as_deref() {
6914 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
6915 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
6916 }
6917 Ok(())
6918 }
6919
6920 pub(crate) fn publish_run_lookup_directory(&mut self) -> Result<()> {
6927 mongreldb_fault::inject("manifest.publication")
6931 .map_err(|e| MongrelError::Other(format!("manifest.publication fault: {e}")))?;
6932 let active_runs = self.run_refs.clone();
6933 let schema_id = self.schema.schema_id;
6934 let index_generation = self.global_idx_epoch;
6935 let run_generation = active_runs.len() as u64;
6936 let fingerprint = crate::run_lookup::compute_fingerprint(
6937 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
6938 schema_id,
6939 index_generation,
6940 run_generation,
6941 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
6942 );
6943 let dir = match crate::run_lookup::RunLookupDirectory::rebuild_from_runs(&active_runs, self)
6944 {
6945 Ok(d) => d,
6946 Err(error) => {
6947 self.run_lookup.complete = false;
6950 self.lookup_metrics
6951 .directory_incomplete
6952 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6953 return Err(error);
6954 }
6955 };
6956 let mut dir = dir;
6960 dir.set_fingerprint(fingerprint);
6961 let runs_dir = match self.runs_root.as_deref() {
6962 Some(root) => match root.io_path() {
6963 Ok(p) => p,
6964 Err(_) => self.dir.join(RUNS_DIR),
6965 },
6966 None => self.dir.join(RUNS_DIR),
6967 };
6968 mongreldb_fault::inject("directory.temp_write")
6971 .map_err(|e| MongrelError::Other(format!("directory.temp_write fault: {e}")))?;
6972 let publish =
6973 dir.write_checkpoint_sharded(&runs_dir, crate::run_lookup::DEFAULT_SHARD_BYTES);
6974 mongreldb_fault::inject("directory.sync")
6975 .map_err(|e| MongrelError::Other(format!("directory.sync fault: {e}")))?;
6976 mongreldb_fault::inject("directory.rename")
6977 .map_err(|e| MongrelError::Other(format!("directory.rename fault: {e}")))?;
6978 if let Err(error) = publish {
6979 self.run_lookup = RunLookupState {
6982 directory: None,
6983 fingerprint,
6984 complete: false,
6985 };
6986 self.lookup_metrics
6987 .directory_incomplete
6988 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6989 self.lookup_metrics
6990 .directory_lookup_fallback
6991 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6992 return Err(error.into());
6993 }
6994 self.run_lookup = RunLookupState {
6995 directory: Some(std::sync::Arc::new(dir)),
6996 fingerprint,
6997 complete: true,
6998 };
6999 Ok(())
7000 }
7001
7002 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
7003 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
7007 let live_count = u64::try_from(rows.len())
7008 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
7009 let auto_inc = match self.auto_inc {
7010 Some(mut state) => {
7011 let maximum = self.scan_max_int64(state.column_id)?;
7012 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
7013 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
7014 })?;
7015 state.next = state.next.max(after_maximum).max(1);
7016 state.seeded = true;
7017 Some(state)
7018 }
7019 None => None,
7020 };
7021 Ok(RecoveryMetadataPlan {
7022 live_count,
7023 auto_inc,
7024 changed: live_count != self.live_count
7025 || auto_inc.is_some_and(|planned| {
7026 self.auto_inc.is_none_or(|current| {
7027 current.next != planned.next || current.seeded != planned.seeded
7028 })
7029 }),
7030 })
7031 }
7032
7033 pub(crate) fn apply_recovered_metadata(
7034 &mut self,
7035 plan: RecoveryMetadataPlan,
7036 epoch: Epoch,
7037 ) -> Result<()> {
7038 if !plan.changed {
7039 return Ok(());
7040 }
7041 self.live_count = plan.live_count;
7042 self.auto_inc = plan.auto_inc;
7043 self.persist_manifest(epoch)
7044 }
7045
7046 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
7052 if !self.indexes_complete {
7055 return;
7056 }
7057 if crate::catalog::inject_hook("index.publish.before").is_err() {
7060 return;
7061 }
7062 if self.idx_root.is_none() {
7063 if let Some(root) = self._root_guard.as_ref() {
7064 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
7065 return;
7066 };
7067 self.idx_root = Some(Arc::new(idx_root));
7068 }
7069 }
7070 let snap = global_idx::IndexSnapshot {
7071 hot: &self.hot,
7072 bitmap: &self.bitmap,
7073 ann: &self.ann,
7074 fm: &self.fm,
7075 sparse: &self.sparse,
7076 minhash: &self.minhash,
7077 learned_range: &self.learned_range,
7078 };
7079 let idx_dek = self.idx_dek();
7081 let written = match self.idx_root.as_deref() {
7082 Some(root) => global_idx::write_atomic_root(
7083 root,
7084 self.table_id,
7085 epoch.0,
7086 snap,
7087 idx_dek.as_deref(),
7088 ),
7089 None => global_idx::write_atomic(
7090 &self.dir,
7091 self.table_id,
7092 epoch.0,
7093 snap,
7094 idx_dek.as_deref(),
7095 ),
7096 };
7097 if written.is_ok() {
7098 self.global_idx_epoch = epoch.0;
7099 let _ = self.persist_manifest(epoch);
7100 let _ = crate::catalog::inject_hook("index.publish.after");
7102 }
7103 }
7104
7105 pub(crate) fn invalidate_index_checkpoint(&mut self) {
7108 self.global_idx_epoch = 0;
7109 if let Some(root) = self.idx_root.as_deref() {
7110 let _ = root.remove_file(global_idx::IDX_FILENAME);
7111 } else {
7112 global_idx::remove(&self.dir);
7113 }
7114 let _ = self.persist_manifest(self.epoch.visible());
7115 }
7116
7117 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
7122 self.indexes_complete = false;
7123 self.global_idx_epoch = 0;
7124 if let Some(root) = self.idx_root.as_deref() {
7125 let _ = root.remove_file(global_idx::IDX_FILENAME);
7126 } else {
7127 global_idx::remove(&self.dir);
7128 }
7129 }
7130
7131 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
7132 self.indexes_complete = true;
7133 }
7134
7135 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
7141 self.durable_commit_failed = true;
7142 if let WalSink::Shared(shared) = &self.wal {
7143 shared
7144 .poisoned
7145 .store(true, std::sync::atomic::Ordering::Relaxed);
7146 }
7147 }
7148
7149 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
7153 self.durable_commit_failed = true;
7154 }
7155
7156 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
7181 let mut best: Option<Row> = None;
7182 fn consider(best: &mut Option<Row>, row: Row, snapshot: Snapshot) {
7185 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
7186 return;
7187 }
7188 if best.as_ref().is_none_or(|current| {
7189 Snapshot::version_is_newer(
7190 row.committed_epoch,
7191 row.commit_ts,
7192 current.committed_epoch,
7193 current.commit_ts,
7194 )
7195 }) {
7196 *best = Some(row);
7197 }
7198 }
7199 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
7200 consider(&mut best, row, snapshot);
7201 }
7202 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
7203 consider(&mut best, row, snapshot);
7204 }
7205 let decision = if self.run_lookup.complete {
7209 let active_ids: Vec<u128> = self.run_refs.iter().map(|r| r.run_id).collect();
7214 let expected_fp = crate::run_lookup::compute_fingerprint(
7215 &active_ids,
7216 self.schema.schema_id,
7217 self.global_idx_epoch,
7218 active_ids.len() as u64,
7219 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
7220 );
7221 if self.run_lookup.fingerprint == expected_fp {
7222 match self.run_lookup.directory.as_ref() {
7223 Some(dir) => dir.decide(row_id),
7224 None => crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale,
7225 }
7226 } else {
7227 self.lookup_metrics
7228 .directory_lookup_fallback
7229 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7230 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale
7231 }
7232 } else {
7233 self.lookup_metrics
7234 .directory_lookup_fallback
7235 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7236 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale
7237 };
7238 match decision {
7239 crate::run_lookup::DirectoryLookupDecision::CompleteMiss => {
7240 self.lookup_metrics
7245 .directory_lookup_hit
7246 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7247 self.lookup_metrics
7248 .directory_complete_miss_total
7249 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7250 }
7251 crate::run_lookup::DirectoryLookupDecision::Candidates(mut locators) => {
7252 locators.sort_by(|a, b| {
7259 if snapshot.uses_hlc_authority() {
7260 b.max_hlc
7261 .cmp(&a.max_hlc)
7262 .then_with(|| b.max_epoch.cmp(&a.max_epoch))
7263 } else {
7264 b.max_epoch
7265 .cmp(&a.max_epoch)
7266 .then_with(|| b.max_hlc.cmp(&a.max_hlc))
7267 }
7268 });
7269 let mut opened: std::collections::HashSet<u128> = std::collections::HashSet::new();
7275 let mut early_stopped: u64 = 0;
7276 for locator in &locators {
7277 let stamp = best
7280 .as_ref()
7281 .map(|current| crate::run_lookup::VersionStamp {
7282 epoch: current.committed_epoch,
7283 hlc: current.commit_ts,
7284 });
7285 if let Some(s) = stamp {
7286 if !locator.can_contain_version_newer_than(s, snapshot) {
7287 early_stopped += 1;
7288 continue;
7289 }
7290 }
7291 if locator.is_impossible_for(snapshot) {
7292 continue;
7293 }
7294 if !opened.insert(locator.run_id) {
7295 continue;
7296 }
7297 self.lookup_metrics
7298 .directory_run_readers_opened
7299 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7300 self.lookup_metrics
7301 .get_run_opened
7302 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7303 let Ok(mut reader) = self.open_reader(locator.run_id) else {
7304 continue;
7305 };
7306 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
7309 continue;
7310 };
7311 consider(&mut best, row, snapshot);
7312 }
7313 if early_stopped > 0 {
7314 self.lookup_metrics
7315 .directory_early_stop_total
7316 .fetch_add(early_stopped, std::sync::atomic::Ordering::Relaxed);
7317 }
7318 self.lookup_metrics
7319 .directory_lookup_hit
7320 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7321 }
7322 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale => {
7323 for rr in &self.run_refs {
7325 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
7329 if row_id.0 < min_rid || row_id.0 > max_rid {
7330 self.lookup_metrics
7331 .get_run_skipped
7332 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7333 continue;
7334 }
7335 }
7336 self.lookup_metrics
7337 .get_run_opened
7338 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7339 let Ok(mut reader) = self.open_reader(rr.run_id) else {
7340 continue;
7341 };
7342 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
7345 continue;
7346 };
7347 consider(&mut best, row, snapshot);
7348 }
7349 }
7350 }
7351 let now_nanos = unix_nanos_now();
7352 match best {
7353 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
7354 Some(r) => Some(r),
7355 None => None,
7356 }
7357 }
7358
7359 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
7363 self.visible_rows_at_time(snapshot, unix_nanos_now())
7364 }
7365
7366 #[doc(hidden)]
7369 pub fn visible_rows_controlled(
7370 &self,
7371 snapshot: Snapshot,
7372 control: &crate::ExecutionControl,
7373 ) -> Result<Vec<Row>> {
7374 let mut out = Vec::new();
7375 self.for_each_visible_row_controlled(snapshot, control, |row| {
7376 out.push(row);
7377 Ok(())
7378 })?;
7379 Ok(out)
7380 }
7381
7382 #[doc(hidden)]
7385 pub fn for_each_visible_row_controlled<F>(
7386 &self,
7387 snapshot: Snapshot,
7388 control: &crate::ExecutionControl,
7389 mut visit: F,
7390 ) -> Result<()>
7391 where
7392 F: FnMut(Row) -> Result<()>,
7393 {
7394 let setup_start = std::time::Instant::now();
7401 let mut versions_examined: u64 = 0;
7402 let mut rows_emitted: u64 = 0;
7403 let mut checkpoints: u64 = 0;
7404 let mut first_row_recorded = false;
7405 let mut sources: Vec<ControlledVisibleSource<'_>> =
7406 Vec::with_capacity(self.run_refs.len() + 2);
7407 control.checkpoint()?;
7408 checkpoints += 1;
7409 if !self.memtable.is_empty() {
7412 sources.push(ControlledVisibleSource::memtable_cursor(
7413 self.memtable.newest_visible_iter(&snapshot),
7414 ));
7415 }
7416 control.checkpoint()?;
7417 checkpoints += 1;
7418 if !self.mutable_run.is_empty() {
7419 sources.push(ControlledVisibleSource::mutable_run_cursor(
7420 self.mutable_run.newest_visible_iter(&snapshot),
7421 ));
7422 }
7423 for run in &self.run_refs {
7424 control.checkpoint()?;
7425 checkpoints += 1;
7426 let reader = self.open_reader(run.run_id)?;
7427 sources.push(ControlledVisibleSource::run(
7428 reader.into_visible_version_cursor_at(snapshot)?,
7429 ));
7430 }
7431 let first_row_start = setup_start;
7434 let _setup_us = setup_start.elapsed().as_micros() as u64;
7435 let now_nanos = unix_nanos_now();
7436 let mut first_row_us: u64 = 0;
7437 let result = merge_controlled_visible_sources(
7438 &mut sources,
7439 control,
7440 |row| self.row_expired_at(row, now_nanos),
7441 |row| {
7442 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
7443 return Ok(());
7444 }
7445 versions_examined += 1;
7446 if !first_row_recorded {
7447 first_row_us = first_row_start.elapsed().as_micros() as u64;
7448 first_row_recorded = true;
7449 }
7450 rows_emitted += 1;
7451 visit(row)
7452 },
7453 );
7454 let time_to_first_row_us = if first_row_recorded { first_row_us } else { 0 };
7455 crate::trace::QueryTrace::record(|t| {
7456 t.controlled_scan_versions_examined = t
7457 .controlled_scan_versions_examined
7458 .saturating_add(versions_examined as usize);
7459 t.controlled_scan_rows_emitted = t
7460 .controlled_scan_rows_emitted
7461 .saturating_add(rows_emitted as usize);
7462 t.controlled_scan_checkpoints = t
7463 .controlled_scan_checkpoints
7464 .saturating_add(checkpoints as usize);
7465 t.controlled_scan_time_to_first_row_us = t
7466 .controlled_scan_time_to_first_row_us
7467 .saturating_add(time_to_first_row_us);
7468 t.controlled_scan_setup_time_us =
7469 t.controlled_scan_setup_time_us.saturating_add(_setup_us);
7470 });
7471 result
7472 }
7473
7474 #[doc(hidden)]
7475 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
7476 let mut best: HashMap<u64, Row> = HashMap::new();
7477 let mut fold = |row: Row| {
7478 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
7479 return;
7480 }
7481 best.entry(row.row_id.0)
7482 .and_modify(|existing| {
7483 if Snapshot::version_is_newer(
7484 row.committed_epoch,
7485 row.commit_ts,
7486 existing.committed_epoch,
7487 existing.commit_ts,
7488 ) {
7489 *existing = row.clone();
7490 }
7491 })
7492 .or_insert(row);
7493 };
7494 for row in self.memtable.visible_versions_at(snapshot) {
7495 fold(row);
7496 }
7497 for row in self.mutable_run.visible_versions_at(snapshot) {
7498 fold(row);
7499 }
7500 for rr in &self.run_refs {
7501 let mut reader = self.open_reader(rr.run_id)?;
7502 for row in reader.visible_versions_at(snapshot)? {
7504 fold(row);
7505 }
7506 }
7507 let mut out: Vec<Row> = best
7508 .into_values()
7509 .filter_map(|r| {
7510 if r.deleted || self.row_expired_at(&r, now_nanos) {
7511 None
7512 } else {
7513 Some(r)
7514 }
7515 })
7516 .collect();
7517 out.sort_by_key(|r| r.row_id);
7518 Ok(out)
7519 }
7520
7521 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
7528 if self.ttl.is_none()
7529 && self.memtable.is_empty()
7530 && self.mutable_run.is_empty()
7531 && self.run_refs.len() == 1
7532 {
7533 let rr = self.run_refs[0].clone();
7534 let mut reader = self.open_reader(rr.run_id)?;
7535 let idxs = reader.visible_indices_at(snapshot)?;
7536 let mut cols = Vec::with_capacity(self.schema.columns.len());
7537 for cdef in &self.schema.columns {
7538 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
7539 }
7540 return Ok(cols);
7541 }
7542 let rows = self.visible_rows(snapshot)?;
7544 let mut cols: Vec<(u16, Vec<Value>)> = self
7545 .schema
7546 .columns
7547 .iter()
7548 .map(|c| (c.id, Vec::with_capacity(rows.len())))
7549 .collect();
7550 for r in &rows {
7551 for (cid, vec) in cols.iter_mut() {
7552 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
7553 }
7554 }
7555 Ok(cols)
7556 }
7557
7558 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
7560 let row_id = self.hot.get(key)?;
7561 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
7562 Some(row_id)
7563 } else {
7564 None
7565 }
7566 }
7567
7568 pub fn lookup_metrics_snapshot(&self) -> LookupMetricsSnapshot {
7572 let mut snap = self.lookup_metrics.snapshot();
7573 let cache = self.result_cache.lock();
7574 let (mem, disk, miss, write_us) = cache.cache_counters();
7575 snap.result_cache_memory_hit = mem;
7576 snap.result_cache_disk_hit = disk;
7577 snap.result_cache_miss = miss;
7578 snap.result_cache_persistent_write_us = write_us;
7579 let (unavailable, skipped, spawn_failures, worker_shutdowns) = cache.publication_counters();
7581 snap.result_cache_persist_unavailable_total = unavailable;
7582 snap.result_cache_persist_skipped_total = skipped;
7583 snap.result_cache_worker_spawn_failures_total = spawn_failures;
7584 snap.result_cache_worker_shutdown_total = worker_shutdowns;
7585 if let Some(writer_snap) = cache.persist_snapshot() {
7588 snap.result_cache_persist_enqueued_total =
7589 writer_snap.result_cache_persist_enqueued_total;
7590 snap.result_cache_persist_coalesced_total =
7591 writer_snap.result_cache_persist_coalesced_total;
7592 snap.result_cache_persist_dropped_store_total =
7593 writer_snap.result_cache_persist_dropped_store_total;
7594 snap.result_cache_persist_remove_total = writer_snap.result_cache_persist_remove_total;
7595 snap.result_cache_persist_stale_store_skipped_total =
7596 writer_snap.result_cache_persist_stale_store_skipped_total;
7597 snap.result_cache_persist_errors_total = writer_snap.result_cache_persist_errors_total;
7598 snap.result_cache_persist_shutdown_abandoned_total =
7599 writer_snap.result_cache_persist_shutdown_abandoned_total;
7600 snap.result_cache_persist_queue_depth = writer_snap.result_cache_persist_queue_depth;
7601 }
7602 snap
7603 }
7604
7605 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
7610 self.query_at_with_allowed(q, self.snapshot(), None)
7611 }
7612
7613 pub fn query_controlled(
7616 &mut self,
7617 q: &crate::query::Query,
7618 control: &crate::ExecutionControl,
7619 ) -> Result<Vec<Row>> {
7620 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
7621 }
7622
7623 pub fn query_at_with_allowed(
7626 &mut self,
7627 q: &crate::query::Query,
7628 snapshot: Snapshot,
7629 allowed: Option<&std::collections::HashSet<RowId>>,
7630 ) -> Result<Vec<Row>> {
7631 self.query_at_with_allowed_after(q, snapshot, allowed, None)
7632 }
7633
7634 #[doc(hidden)]
7635 pub fn query_at_with_allowed_controlled(
7636 &mut self,
7637 q: &crate::query::Query,
7638 snapshot: Snapshot,
7639 allowed: Option<&std::collections::HashSet<RowId>>,
7640 control: &crate::ExecutionControl,
7641 ) -> Result<Vec<Row>> {
7642 self.require_select()?;
7643 self.ensure_indexes_complete_controlled(control, || true)?;
7644 self.validate_native_query(q)?;
7645 self.query_conditions_at(
7646 &q.conditions,
7647 snapshot,
7648 allowed,
7649 q.limit,
7650 q.offset,
7651 None,
7652 unix_nanos_now(),
7653 Some(control),
7654 )
7655 }
7656
7657 #[doc(hidden)]
7658 pub fn query_at_with_allowed_after(
7659 &mut self,
7660 q: &crate::query::Query,
7661 snapshot: Snapshot,
7662 allowed: Option<&std::collections::HashSet<RowId>>,
7663 after_row_id: Option<RowId>,
7664 ) -> Result<Vec<Row>> {
7665 self.query_at_with_allowed_after_at_time(
7666 q,
7667 snapshot,
7668 allowed,
7669 after_row_id,
7670 unix_nanos_now(),
7671 )
7672 }
7673
7674 #[doc(hidden)]
7675 pub fn query_at_with_allowed_after_at_time(
7676 &mut self,
7677 q: &crate::query::Query,
7678 snapshot: Snapshot,
7679 allowed: Option<&std::collections::HashSet<RowId>>,
7680 after_row_id: Option<RowId>,
7681 query_time_nanos: i64,
7682 ) -> Result<Vec<Row>> {
7683 self.require_select()?;
7684 self.ensure_indexes_complete()?;
7685 self.validate_native_query(q)?;
7686 self.query_conditions_at(
7687 &q.conditions,
7688 snapshot,
7689 allowed,
7690 q.limit,
7691 q.offset,
7692 after_row_id,
7693 query_time_nanos,
7694 None,
7695 )
7696 }
7697
7698 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
7699 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7700 return Err(MongrelError::InvalidArgument(format!(
7701 "query exceeds {} conditions",
7702 crate::query::MAX_HARD_CONDITIONS
7703 )));
7704 }
7705 if let Some(limit) = q.limit {
7706 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
7707 return Err(MongrelError::InvalidArgument(format!(
7708 "query limit must be between 1 and {}",
7709 crate::query::MAX_FINAL_LIMIT
7710 )));
7711 }
7712 }
7713 if q.offset > crate::query::MAX_QUERY_OFFSET {
7714 return Err(MongrelError::InvalidArgument(format!(
7715 "query offset exceeds {}",
7716 crate::query::MAX_QUERY_OFFSET
7717 )));
7718 }
7719 Ok(())
7720 }
7721
7722 #[doc(hidden)]
7725 pub fn query_all_at(
7726 &mut self,
7727 conditions: &[crate::query::Condition],
7728 snapshot: Snapshot,
7729 ) -> Result<Vec<Row>> {
7730 self.require_select()?;
7731 self.ensure_indexes_complete()?;
7732 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7733 return Err(MongrelError::InvalidArgument(format!(
7734 "query exceeds {} conditions",
7735 crate::query::MAX_HARD_CONDITIONS
7736 )));
7737 }
7738 self.query_conditions_at(
7739 conditions,
7740 snapshot,
7741 None,
7742 None,
7743 0,
7744 None,
7745 unix_nanos_now(),
7746 None,
7747 )
7748 }
7749
7750 #[allow(clippy::too_many_arguments)]
7751 fn query_conditions_at(
7752 &self,
7753 conditions: &[crate::query::Condition],
7754 snapshot: Snapshot,
7755 allowed: Option<&std::collections::HashSet<RowId>>,
7756 limit: Option<usize>,
7757 offset: usize,
7758 after_row_id: Option<RowId>,
7759 query_time_nanos: i64,
7760 control: Option<&crate::ExecutionControl>,
7761 ) -> Result<Vec<Row>> {
7762 control
7763 .map(crate::ExecutionControl::checkpoint)
7764 .transpose()?;
7765 crate::trace::QueryTrace::record(|t| {
7766 t.run_count = self.run_refs.len();
7767 t.memtable_rows = self.memtable.len();
7768 t.mutable_run_rows = self.mutable_run.len();
7769 });
7770 if conditions.is_empty() {
7774 crate::trace::QueryTrace::record(|t| {
7775 t.scan_mode = crate::trace::ScanMode::Materialized;
7776 t.row_materialized = true;
7777 });
7778 let mut rows = match control {
7779 Some(control) => self.visible_rows_controlled(snapshot, control)?,
7780 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
7781 };
7782 if let Some(allowed) = allowed {
7783 let mut filtered = Vec::with_capacity(rows.len());
7784 for (index, row) in rows.into_iter().enumerate() {
7785 if index & 255 == 0 {
7786 control
7787 .map(crate::ExecutionControl::checkpoint)
7788 .transpose()?;
7789 }
7790 if allowed.contains(&row.row_id) {
7791 filtered.push(row);
7792 }
7793 }
7794 rows = filtered;
7795 }
7796 if let Some(after_row_id) = after_row_id {
7797 rows.retain(|row| row.row_id > after_row_id);
7798 }
7799 rows.drain(..offset.min(rows.len()));
7800 if let Some(limit) = limit {
7801 rows.truncate(limit);
7802 }
7803 return Ok(rows);
7804 }
7805 crate::trace::QueryTrace::record(|t| {
7806 t.conditions_pushed = conditions.len();
7807 t.scan_mode = crate::trace::ScanMode::Materialized;
7808 t.row_materialized = true;
7809 });
7810 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
7817 ordered.sort_by_key(|c| condition_cost_rank(c));
7818 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
7819 for c in &ordered {
7820 control
7821 .map(crate::ExecutionControl::checkpoint)
7822 .transpose()?;
7823 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
7824 let empty = s.is_empty();
7825 sets.push(s);
7826 if empty {
7827 break;
7828 }
7829 }
7830 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7831 if let Some(allowed) = allowed {
7832 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
7833 }
7834 if let Some(after_row_id) = after_row_id {
7835 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
7836 rids.drain(..first);
7837 }
7838 rids.drain(..offset.min(rids.len()));
7839 if let Some(limit) = limit {
7840 rids.truncate(limit);
7841 }
7842 control
7843 .map(crate::ExecutionControl::checkpoint)
7844 .transpose()?;
7845 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
7846 }
7847
7848 pub fn retrieve(
7850 &mut self,
7851 retriever: &crate::query::Retriever,
7852 ) -> Result<Vec<crate::query::RetrieverHit>> {
7853 self.retrieve_with_allowed(retriever, None)
7854 }
7855
7856 pub fn retrieve_at(
7857 &mut self,
7858 retriever: &crate::query::Retriever,
7859 snapshot: Snapshot,
7860 allowed: Option<&std::collections::HashSet<RowId>>,
7861 ) -> Result<Vec<crate::query::RetrieverHit>> {
7862 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
7863 }
7864
7865 pub fn retrieve_with_allowed(
7868 &mut self,
7869 retriever: &crate::query::Retriever,
7870 allowed: Option<&std::collections::HashSet<RowId>>,
7871 ) -> Result<Vec<crate::query::RetrieverHit>> {
7872 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
7873 }
7874
7875 pub fn retrieve_at_with_allowed(
7876 &mut self,
7877 retriever: &crate::query::Retriever,
7878 snapshot: Snapshot,
7879 allowed: Option<&std::collections::HashSet<RowId>>,
7880 ) -> Result<Vec<crate::query::RetrieverHit>> {
7881 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
7882 }
7883
7884 pub fn retrieve_at_with_allowed_and_context(
7885 &mut self,
7886 retriever: &crate::query::Retriever,
7887 snapshot: Snapshot,
7888 allowed: Option<&std::collections::HashSet<RowId>>,
7889 context: Option<&crate::query::AiExecutionContext>,
7890 ) -> Result<Vec<crate::query::RetrieverHit>> {
7891 self.require_select()?;
7892 self.ensure_indexes_complete()?;
7893 self.validate_retriever(retriever)?;
7894 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
7895 }
7896
7897 pub fn retrieve_at_with_candidate_authorization_and_context(
7898 &mut self,
7899 retriever: &crate::query::Retriever,
7900 snapshot: Snapshot,
7901 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7902 context: Option<&crate::query::AiExecutionContext>,
7903 ) -> Result<Vec<crate::query::RetrieverHit>> {
7904 self.require_select()?;
7905 self.ensure_indexes_complete()?;
7906 self.retrieve_at_with_candidate_authorization_on_generation(
7907 retriever,
7908 snapshot,
7909 authorization,
7910 context,
7911 )
7912 }
7913
7914 #[doc(hidden)]
7915 pub fn retrieve_at_with_candidate_authorization_on_generation(
7916 &self,
7917 retriever: &crate::query::Retriever,
7918 snapshot: Snapshot,
7919 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7920 context: Option<&crate::query::AiExecutionContext>,
7921 ) -> Result<Vec<crate::query::RetrieverHit>> {
7922 self.require_select()?;
7923 self.validate_retriever(retriever)?;
7924 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
7925 }
7926
7927 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
7928 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
7929 let (column_id, k) = match retriever {
7930 Retriever::Ann {
7931 column_id,
7932 query,
7933 k,
7934 } => {
7935 let index = self.ann.get(column_id).ok_or_else(|| {
7936 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
7937 })?;
7938 if query.len() != index.dim() {
7939 return Err(MongrelError::InvalidArgument(format!(
7940 "ANN query dimension must be {}, got {}",
7941 index.dim(),
7942 query.len()
7943 )));
7944 }
7945 if query.iter().any(|value| !value.is_finite()) {
7946 return Err(MongrelError::InvalidArgument(
7947 "ANN query values must be finite".into(),
7948 ));
7949 }
7950 (*column_id, *k)
7951 }
7952 Retriever::Sparse {
7953 column_id,
7954 query,
7955 k,
7956 } => {
7957 if !self.sparse.contains_key(column_id) {
7958 return Err(MongrelError::InvalidArgument(format!(
7959 "column {column_id} has no Sparse index"
7960 )));
7961 }
7962 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
7963 return Err(MongrelError::InvalidArgument(
7964 "Sparse query must be non-empty with finite weights".into(),
7965 ));
7966 }
7967 if query.len() > MAX_SPARSE_TERMS {
7968 return Err(MongrelError::InvalidArgument(format!(
7969 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
7970 )));
7971 }
7972 (*column_id, *k)
7973 }
7974 Retriever::MinHash {
7975 column_id,
7976 members,
7977 k,
7978 } => {
7979 if !self.minhash.contains_key(column_id) {
7980 return Err(MongrelError::InvalidArgument(format!(
7981 "column {column_id} has no MinHash index"
7982 )));
7983 }
7984 if members.is_empty() {
7985 return Err(MongrelError::InvalidArgument(
7986 "MinHash members must not be empty".into(),
7987 ));
7988 }
7989 if members.len() > MAX_SET_MEMBERS {
7990 return Err(MongrelError::InvalidArgument(format!(
7991 "MinHash query exceeds {MAX_SET_MEMBERS} members"
7992 )));
7993 }
7994 let mut total_bytes = 0usize;
7995 for member in members {
7996 let bytes = member.encoded_len();
7997 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
7998 return Err(MongrelError::InvalidArgument(format!(
7999 "MinHash member exceeds {} bytes",
8000 crate::query::MAX_SET_MEMBER_BYTES
8001 )));
8002 }
8003 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
8004 MongrelError::InvalidArgument("MinHash input size overflow".into())
8005 })?;
8006 }
8007 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
8008 return Err(MongrelError::InvalidArgument(format!(
8009 "MinHash input exceeds {} bytes",
8010 crate::query::MAX_SET_INPUT_BYTES
8011 )));
8012 }
8013 (*column_id, *k)
8014 }
8015 };
8016 if k == 0 {
8017 return Err(MongrelError::InvalidArgument(
8018 "retriever k must be > 0".into(),
8019 ));
8020 }
8021 if k > MAX_RETRIEVER_K {
8022 return Err(MongrelError::InvalidArgument(format!(
8023 "retriever k exceeds {MAX_RETRIEVER_K}"
8024 )));
8025 }
8026 debug_assert!(self
8027 .schema
8028 .columns
8029 .iter()
8030 .any(|column| column.id == column_id));
8031 Ok(())
8032 }
8033
8034 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
8035 use crate::query::Condition;
8036 match condition {
8037 Condition::Ann {
8038 column_id,
8039 query,
8040 k,
8041 } => self.validate_retriever(&crate::query::Retriever::Ann {
8042 column_id: *column_id,
8043 query: query.clone(),
8044 k: *k,
8045 }),
8046 Condition::SparseMatch {
8047 column_id,
8048 query,
8049 k,
8050 } => self.validate_retriever(&crate::query::Retriever::Sparse {
8051 column_id: *column_id,
8052 query: query.clone(),
8053 k: *k,
8054 }),
8055 Condition::MinHashSimilar {
8056 column_id,
8057 query,
8058 k,
8059 } => {
8060 if !self.minhash.contains_key(column_id) {
8061 return Err(MongrelError::InvalidArgument(format!(
8062 "column {column_id} has no MinHash index"
8063 )));
8064 }
8065 if query.is_empty() || *k == 0 {
8066 return Err(MongrelError::InvalidArgument(
8067 "MinHash query must be non-empty and k must be > 0".into(),
8068 ));
8069 }
8070 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
8071 {
8072 return Err(MongrelError::InvalidArgument(format!(
8073 "MinHash query must have <= {} members and k <= {}",
8074 crate::query::MAX_SET_MEMBERS,
8075 crate::query::MAX_RETRIEVER_K
8076 )));
8077 }
8078 Ok(())
8079 }
8080 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
8081 Err(MongrelError::InvalidArgument(format!(
8082 "bitmap IN exceeds {} values",
8083 crate::query::MAX_SET_MEMBERS
8084 )))
8085 }
8086 Condition::FmContainsAll { patterns, .. }
8087 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
8088 {
8089 Err(MongrelError::InvalidArgument(format!(
8090 "FM query exceeds {} patterns",
8091 crate::query::MAX_HARD_CONDITIONS
8092 )))
8093 }
8094 _ => Ok(()),
8095 }
8096 }
8097
8098 fn retrieve_filtered(
8099 &self,
8100 retriever: &crate::query::Retriever,
8101 snapshot: Snapshot,
8102 hard_filter: Option<&RowIdSet>,
8103 allowed: Option<&std::collections::HashSet<RowId>>,
8104 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8105 context: Option<&crate::query::AiExecutionContext>,
8106 ) -> Result<Vec<crate::query::RetrieverHit>> {
8107 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
8108 let started = std::time::Instant::now();
8109 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
8110 Retriever::Ann {
8111 column_id,
8112 query,
8113 k,
8114 } => {
8115 let Some(index) = self.ann.get(column_id) else {
8116 return Ok(Vec::new());
8117 };
8118 let cap = ann_candidate_cap(index.len(), context);
8119 crate::trace::QueryTrace::record(|trace| trace.candidate_cap = cap);
8120 if cap == 0 {
8121 return Ok(Vec::new());
8122 }
8123 let mut breadth = (*k).max(1).min(cap);
8124 let mut eligibility = std::collections::HashMap::new();
8125 let mut raw_candidates_total: usize = 0;
8130 let mut visibility_rejected_total: usize = 0;
8131 let mut authorization_rejected_total: usize = 0;
8132 let mut candidate_cap_hit_final = false;
8133 let mut filtered = loop {
8134 let mut seen = std::collections::HashSet::new();
8135 if let Some(context) = context {
8136 context.checkpoint()?;
8137 }
8138 let raw = index.search_with_context(query, breadth, context)?;
8139 crate::trace::QueryTrace::record(|trace| {
8140 trace.raw_candidates = raw.len();
8141 let unique = raw
8142 .iter()
8143 .map(|(row_id, _)| *row_id)
8144 .collect::<std::collections::HashSet<_>>()
8145 .len();
8146 trace.unique_candidates = unique;
8147 trace.duplicate_candidates = raw.len().saturating_sub(unique);
8148 trace.authorization_rejected = raw
8149 .iter()
8150 .filter(|(row_id, _)| {
8151 allowed.is_some_and(|allowed| !allowed.contains(row_id))
8152 })
8153 .count();
8154 trace.hard_filter_rejected = raw
8155 .iter()
8156 .filter(|(row_id, _)| {
8157 hard_filter.is_some_and(|filter| !filter.contains(row_id.0))
8158 })
8159 .count();
8160 });
8161 raw_candidates_total = raw_candidates_total.saturating_add(raw.len());
8162 let unchecked: Vec<_> = raw
8163 .iter()
8164 .map(|(row_id, _)| *row_id)
8165 .filter(|row_id| !eligibility.contains_key(row_id))
8166 .filter(|row_id| {
8167 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
8168 && allowed.is_none_or(|allowed| allowed.contains(row_id))
8169 })
8170 .collect();
8171 let eligible = self.eligible_and_authorized_candidate_ids(
8172 &unchecked,
8173 *column_id,
8174 snapshot,
8175 candidate_authorization,
8176 context,
8177 )?;
8178 for row_id in &unchecked {
8179 eligibility.insert(*row_id, eligible.contains(row_id));
8180 }
8181 let new_visibility_or_auth_rejected = unchecked
8184 .iter()
8185 .filter(|row_id| !eligible.contains(row_id))
8186 .count();
8187 let mut new_auth_rejected = 0usize;
8192 if let Some(authorization) = candidate_authorization {
8193 if authorization.security.rls_enabled(authorization.table)
8194 && !authorization.principal.is_admin
8195 {
8196 new_auth_rejected = new_visibility_or_auth_rejected;
8202 }
8203 }
8204 let new_visibility_rejected =
8205 new_visibility_or_auth_rejected.saturating_sub(new_auth_rejected);
8206 visibility_rejected_total =
8207 visibility_rejected_total.saturating_add(new_visibility_rejected);
8208 authorization_rejected_total =
8209 authorization_rejected_total.saturating_add(new_auth_rejected);
8210 let filtered: Vec<_> = raw
8211 .into_iter()
8212 .filter(|(row_id, _)| {
8213 seen.insert(*row_id)
8214 && eligibility.get(row_id).copied().unwrap_or(false)
8215 })
8216 .map(|(row_id, score)| {
8217 let score = match score {
8218 crate::index::AnnDistance::Hamming(d) => {
8219 RetrieverScore::AnnHammingDistance(d)
8220 }
8221 crate::index::AnnDistance::Cosine(d) => {
8222 RetrieverScore::AnnCosineDistance(d)
8223 }
8224 };
8225 (row_id, score)
8226 })
8227 .collect();
8228 if filtered.len() >= *k || breadth >= cap {
8229 if filtered.len() < *k && index.len() > cap && breadth >= cap {
8230 crate::trace::QueryTrace::record(|trace| {
8231 trace.ann_candidate_cap_hit = true;
8232 trace.candidate_cap_hit = true;
8233 });
8234 candidate_cap_hit_final = true;
8235 }
8236 break filtered;
8237 }
8238 breadth = breadth.saturating_mul(2).min(cap);
8239 };
8240 let final_hits = filtered.len();
8243 let index_len = index.len();
8244 crate::trace::QueryTrace::record(|trace| {
8245 trace.candidate_cap = cap;
8246 trace.candidate_cap_hit = candidate_cap_hit_final;
8247 trace.ann_candidate_cap_hit = candidate_cap_hit_final;
8248 trace.raw_candidates =
8249 trace.raw_candidates.saturating_add(raw_candidates_total);
8250 trace.visibility_rejected = trace
8251 .visibility_rejected
8252 .saturating_add(visibility_rejected_total);
8253 trace.authorization_rejected = trace
8254 .authorization_rejected
8255 .saturating_add(authorization_rejected_total);
8256 trace.final_hits = trace.final_hits.saturating_add(final_hits);
8257 let _ = index_len;
8260 });
8261 filtered.truncate(*k);
8262 filtered
8263 }
8264 Retriever::Sparse {
8265 column_id,
8266 query,
8267 k,
8268 } => self
8269 .sparse
8270 .get(column_id)
8271 .map(|index| -> Result<Vec<_>> {
8272 let mut breadth = (*k).max(1);
8273 let mut eligibility = std::collections::HashMap::new();
8274 let mut raw_candidates_total: usize = 0;
8277 let mut visibility_rejected_total: usize = 0;
8278 let mut authorization_rejected_total: usize = 0;
8279 let filtered = loop {
8280 if let Some(context) = context {
8281 context.checkpoint()?;
8282 }
8283 let raw = index.search_with_context(query, breadth, context)?;
8284 raw_candidates_total = raw_candidates_total.saturating_add(raw.len());
8285 authorization_rejected_total = authorization_rejected_total.saturating_add(
8286 raw.iter()
8287 .filter(|(row_id, _)| {
8288 allowed.is_some_and(|allowed| !allowed.contains(row_id))
8289 })
8290 .count(),
8291 );
8292 let unchecked: Vec<_> = raw
8293 .iter()
8294 .map(|(row_id, _)| *row_id)
8295 .filter(|row_id| !eligibility.contains_key(row_id))
8296 .filter(|row_id| {
8297 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
8298 && allowed.is_none_or(|allowed| allowed.contains(row_id))
8299 })
8300 .collect();
8301 let eligible = self.eligible_and_authorized_candidate_ids(
8302 &unchecked,
8303 *column_id,
8304 snapshot,
8305 candidate_authorization,
8306 context,
8307 )?;
8308 visibility_rejected_total = visibility_rejected_total.saturating_add(
8309 unchecked
8310 .iter()
8311 .filter(|row_id| !eligible.contains(row_id))
8312 .count(),
8313 );
8314 for row_id in unchecked {
8315 eligibility.insert(row_id, eligible.contains(&row_id));
8316 }
8317 let filtered: Vec<_> = raw
8318 .iter()
8319 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
8320 .take(*k)
8321 .map(|(row_id, score)| {
8322 (*row_id, RetrieverScore::SparseDotProduct(*score))
8323 })
8324 .collect();
8325 if filtered.len() >= *k || raw.len() < breadth {
8326 break filtered;
8327 }
8328 let next = breadth.saturating_mul(2);
8329 if next == breadth {
8330 break filtered;
8331 }
8332 breadth = next;
8333 };
8334 crate::trace::QueryTrace::record(|trace| {
8335 trace.raw_candidates =
8336 trace.raw_candidates.saturating_add(raw_candidates_total);
8337 trace.visibility_rejected = trace
8338 .visibility_rejected
8339 .saturating_add(visibility_rejected_total);
8340 trace.authorization_rejected = trace
8341 .authorization_rejected
8342 .saturating_add(authorization_rejected_total);
8343 });
8344 Ok(filtered)
8345 })
8346 .transpose()?
8347 .unwrap_or_default(),
8348 Retriever::MinHash {
8349 column_id,
8350 members,
8351 k,
8352 } => self
8353 .minhash
8354 .get(column_id)
8355 .map(|index| -> Result<Vec<_>> {
8356 let mut hashes = Vec::with_capacity(members.len());
8357 for member in members {
8358 if let Some(context) = context {
8359 context.consume(crate::query::work_units(
8360 member.encoded_len(),
8361 crate::query::PARSE_WORK_QUANTUM,
8362 ))?;
8363 }
8364 hashes.push(member.hash_v1());
8365 }
8366 let mut breadth = (*k).max(1);
8367 let mut eligibility = std::collections::HashMap::new();
8368 let mut raw_candidates_total: usize = 0;
8369 let mut visibility_rejected_total: usize = 0;
8370 let mut authorization_rejected_total: usize = 0;
8371 let filtered = loop {
8372 if let Some(context) = context {
8373 context.checkpoint()?;
8374 }
8375 let raw = index.search_with_context(&hashes, breadth, context)?;
8376 raw_candidates_total = raw_candidates_total.saturating_add(raw.len());
8377 authorization_rejected_total = authorization_rejected_total.saturating_add(
8378 raw.iter()
8379 .filter(|(row_id, _)| {
8380 allowed.is_some_and(|allowed| !allowed.contains(row_id))
8381 })
8382 .count(),
8383 );
8384 let unchecked: Vec<_> = raw
8385 .iter()
8386 .map(|(row_id, _)| *row_id)
8387 .filter(|row_id| !eligibility.contains_key(row_id))
8388 .filter(|row_id| {
8389 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
8390 && allowed.is_none_or(|allowed| allowed.contains(row_id))
8391 })
8392 .collect();
8393 let eligible = self.eligible_and_authorized_candidate_ids(
8394 &unchecked,
8395 *column_id,
8396 snapshot,
8397 candidate_authorization,
8398 context,
8399 )?;
8400 visibility_rejected_total = visibility_rejected_total.saturating_add(
8401 unchecked
8402 .iter()
8403 .filter(|row_id| !eligible.contains(row_id))
8404 .count(),
8405 );
8406 for row_id in unchecked {
8407 eligibility.insert(row_id, eligible.contains(&row_id));
8408 }
8409 let filtered: Vec<_> = raw
8410 .iter()
8411 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
8412 .take(*k)
8413 .map(|(row_id, score)| {
8414 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
8415 })
8416 .collect();
8417 if filtered.len() >= *k || raw.len() < breadth {
8418 break filtered;
8419 }
8420 let next = breadth.saturating_mul(2);
8421 if next == breadth {
8422 break filtered;
8423 }
8424 breadth = next;
8425 };
8426 crate::trace::QueryTrace::record(|trace| {
8427 trace.raw_candidates =
8428 trace.raw_candidates.saturating_add(raw_candidates_total);
8429 trace.visibility_rejected = trace
8430 .visibility_rejected
8431 .saturating_add(visibility_rejected_total);
8432 trace.authorization_rejected = trace
8433 .authorization_rejected
8434 .saturating_add(authorization_rejected_total);
8435 });
8436 Ok(filtered)
8437 })
8438 .transpose()?
8439 .unwrap_or_default(),
8440 };
8441 let requested_k = match retriever {
8442 Retriever::Ann { k, .. }
8443 | Retriever::Sparse { k, .. }
8444 | Retriever::MinHash { k, .. } => *k,
8445 };
8446 crate::trace::QueryTrace::record(|trace| {
8447 trace.final_hits = scored.len();
8448 if scored.len() < requested_k {
8449 trace.underfill_reason = Some(if trace.candidate_cap_hit {
8450 "candidate_cap"
8451 } else {
8452 "eligible_candidates_exhausted"
8453 });
8454 }
8455 });
8456 let elapsed = started.elapsed().as_nanos() as u64;
8457 crate::trace::QueryTrace::record(|trace| {
8458 match retriever {
8459 Retriever::Ann { .. } => {
8460 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
8461 if let Retriever::Ann { column_id, .. } = retriever {
8462 if let Some(index) = self.ann.get(column_id) {
8463 trace.ann_algorithm = Some(index.algorithm());
8464 trace.ann_quantization = Some(index.quantization());
8465 trace.ann_backend = Some(index.backend_name());
8466 }
8467 }
8468 }
8469 Retriever::Sparse { .. } => {
8470 trace.sparse_candidate_nanos =
8471 trace.sparse_candidate_nanos.saturating_add(elapsed)
8472 }
8473 Retriever::MinHash { .. } => {
8474 trace.minhash_candidate_nanos =
8475 trace.minhash_candidate_nanos.saturating_add(elapsed)
8476 }
8477 }
8478 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
8479 });
8480 Ok(scored
8481 .into_iter()
8482 .enumerate()
8483 .map(|(rank, (row_id, score))| RetrieverHit {
8484 row_id,
8485 rank: rank + 1,
8486 score,
8487 })
8488 .collect())
8489 }
8490
8491 fn eligible_candidate_ids(
8492 &self,
8493 candidates: &[RowId],
8494 _column_id: u16,
8495 snapshot: Snapshot,
8496 context: Option<&crate::query::AiExecutionContext>,
8497 ) -> Result<std::collections::HashSet<RowId>> {
8498 let lookup_snapshot = snapshot;
8510 let mut readers: Vec<_> = self
8511 .run_refs
8512 .iter()
8513 .map(|run| self.open_reader(run.run_id))
8514 .collect::<Result<_>>()?;
8515 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8516 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
8517 for &row_id in candidates {
8518 if let Some(context) = context {
8519 context.consume(1)?;
8520 }
8521 let mem = self.memtable.get_version_at(row_id, lookup_snapshot);
8522 let mutable = self.mutable_run.get_version_at(row_id, lookup_snapshot);
8523 let overlay = match (mem, mutable) {
8524 (Some(left), Some(right)) => Some(
8525 if Snapshot::version_is_newer(
8526 left.1.committed_epoch,
8527 left.1.commit_ts,
8528 right.1.committed_epoch,
8529 right.1.commit_ts,
8530 ) {
8531 left
8532 } else {
8533 right
8534 },
8535 ),
8536 (Some(value), None) | (None, Some(value)) => Some(value),
8537 (None, None) => None,
8538 };
8539 if let Some((_, row)) = overlay {
8540 if !row.deleted && !self.row_expired_at(&row, now) {
8541 eligible.insert(row_id);
8542 }
8543 continue;
8544 }
8545 let mut best: Option<(crate::sorted_run::VersionVisibility, usize)> = None;
8551 for (index, reader) in readers.iter_mut().enumerate() {
8552 if let Some(visibility) =
8553 reader.get_version_visibility_full_at(row_id, lookup_snapshot)?
8554 {
8555 if best
8556 .as_ref()
8557 .map(|(current, _)| visibility.stamp.is_newer_than(current.stamp))
8558 .unwrap_or(true)
8559 {
8560 best = Some((visibility, index));
8561 }
8562 }
8563 }
8564 let Some((visibility, reader_index)) = best else {
8565 continue;
8566 };
8567 if visibility.deleted {
8568 continue;
8569 }
8570 if let Some(ttl) = self.ttl {
8571 if let Some(ttl_value) = readers[reader_index].get_version_column_full_at(
8572 row_id,
8573 lookup_snapshot,
8574 ttl.column_id,
8575 )? {
8576 if let Some(Value::Int64(timestamp)) = ttl_value.value {
8577 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8578 continue;
8579 }
8580 }
8581 }
8582 }
8583 eligible.insert(row_id);
8584 }
8585 Ok(eligible)
8586 }
8587
8588 fn eligible_and_authorized_candidate_ids(
8589 &self,
8590 candidates: &[RowId],
8591 column_id: u16,
8592 snapshot: Snapshot,
8593 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8594 context: Option<&crate::query::AiExecutionContext>,
8595 ) -> Result<std::collections::HashSet<RowId>> {
8596 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
8597 let Some(authorization) = authorization else {
8598 return Ok(eligible);
8599 };
8600 let candidates: Vec<_> = eligible.into_iter().collect();
8601 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
8602 }
8603
8604 fn policy_allowed_candidate_ids(
8605 &self,
8606 candidates: &[RowId],
8607 snapshot: Snapshot,
8608 authorization: &crate::security::CandidateAuthorization<'_>,
8609 context: Option<&crate::query::AiExecutionContext>,
8610 ) -> Result<std::collections::HashSet<RowId>> {
8611 let started = std::time::Instant::now();
8612 if candidates.is_empty()
8613 || authorization.principal.is_admin
8614 || !authorization.security.rls_enabled(authorization.table)
8615 {
8616 return Ok(candidates.iter().copied().collect());
8617 }
8618 if let Some(context) = context {
8619 context.checkpoint()?;
8620 }
8621 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
8622 let mut rows: std::collections::HashMap<RowId, Row> = candidates
8623 .iter()
8624 .map(|row_id| {
8625 (
8626 *row_id,
8627 Row {
8628 row_id: *row_id,
8629 committed_epoch: snapshot.epoch,
8630 columns: std::collections::HashMap::new(),
8631 deleted: false,
8632 commit_ts: None,
8633 },
8634 )
8635 })
8636 .collect();
8637 let columns = authorization
8638 .security
8639 .select_policy_columns(authorization.table, authorization.principal);
8640 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8641 let mut decoded = 0usize;
8642 for column_id in &columns {
8643 if let Some(context) = context {
8644 context.checkpoint()?;
8645 }
8646 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8647 &row_ids, *column_id, snapshot, query_now, context,
8648 )? {
8649 if let Some(row) = rows.get_mut(&row_id) {
8650 row.columns.insert(*column_id, value);
8651 decoded = decoded.saturating_add(1);
8652 }
8653 }
8654 }
8655 if let Some(context) = context {
8656 context.consume(candidates.len().saturating_add(decoded))?;
8657 }
8658 let allowed = rows
8659 .into_values()
8660 .filter_map(|row| {
8661 authorization
8662 .security
8663 .row_allowed(
8664 authorization.table,
8665 crate::security::PolicyCommand::Select,
8666 &row,
8667 authorization.principal,
8668 false,
8669 )
8670 .then_some(row.row_id)
8671 })
8672 .collect();
8673 crate::trace::QueryTrace::record(|trace| {
8674 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
8675 trace.rls_policy_columns_decoded =
8676 trace.rls_policy_columns_decoded.saturating_add(decoded);
8677 trace.authorization_nanos = trace
8678 .authorization_nanos
8679 .saturating_add(started.elapsed().as_nanos() as u64);
8680 });
8681 Ok(allowed)
8682 }
8683
8684 pub fn search(
8686 &mut self,
8687 request: &crate::query::SearchRequest,
8688 ) -> Result<Vec<crate::query::SearchHit>> {
8689 self.search_with_allowed(request, None)
8690 }
8691
8692 pub fn search_at(
8693 &mut self,
8694 request: &crate::query::SearchRequest,
8695 snapshot: Snapshot,
8696 authorized: Option<&std::collections::HashSet<RowId>>,
8697 ) -> Result<Vec<crate::query::SearchHit>> {
8698 self.search_at_with_allowed(request, snapshot, authorized)
8699 }
8700
8701 pub fn search_with_allowed(
8702 &mut self,
8703 request: &crate::query::SearchRequest,
8704 authorized: Option<&std::collections::HashSet<RowId>>,
8705 ) -> Result<Vec<crate::query::SearchHit>> {
8706 self.search_at_with_allowed(request, self.snapshot(), authorized)
8707 }
8708
8709 pub fn search_at_with_allowed(
8710 &mut self,
8711 request: &crate::query::SearchRequest,
8712 snapshot: Snapshot,
8713 authorized: Option<&std::collections::HashSet<RowId>>,
8714 ) -> Result<Vec<crate::query::SearchHit>> {
8715 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
8716 }
8717
8718 pub fn search_at_with_allowed_and_context(
8719 &mut self,
8720 request: &crate::query::SearchRequest,
8721 snapshot: Snapshot,
8722 authorized: Option<&std::collections::HashSet<RowId>>,
8723 context: Option<&crate::query::AiExecutionContext>,
8724 ) -> Result<Vec<crate::query::SearchHit>> {
8725 self.ensure_indexes_complete()?;
8726 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
8727 }
8728
8729 pub fn search_at_with_candidate_authorization_and_context(
8730 &mut self,
8731 request: &crate::query::SearchRequest,
8732 snapshot: Snapshot,
8733 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8734 context: Option<&crate::query::AiExecutionContext>,
8735 ) -> Result<Vec<crate::query::SearchHit>> {
8736 self.ensure_indexes_complete()?;
8737 self.search_at_with_filters_and_context(
8738 request,
8739 snapshot,
8740 None,
8741 authorization,
8742 context,
8743 None,
8744 )
8745 }
8746
8747 #[doc(hidden)]
8748 pub fn search_at_with_candidate_authorization_on_generation(
8749 &self,
8750 request: &crate::query::SearchRequest,
8751 snapshot: Snapshot,
8752 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8753 context: Option<&crate::query::AiExecutionContext>,
8754 ) -> Result<Vec<crate::query::SearchHit>> {
8755 self.search_at_with_filters_and_context(
8756 request,
8757 snapshot,
8758 None,
8759 authorization,
8760 context,
8761 None,
8762 )
8763 }
8764
8765 #[doc(hidden)]
8766 pub fn search_at_with_candidate_authorization_on_generation_after(
8767 &self,
8768 request: &crate::query::SearchRequest,
8769 snapshot: Snapshot,
8770 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8771 context: Option<&crate::query::AiExecutionContext>,
8772 after: Option<crate::query::SearchAfter>,
8773 ) -> Result<Vec<crate::query::SearchHit>> {
8774 self.search_at_with_filters_and_context(
8775 request,
8776 snapshot,
8777 None,
8778 authorization,
8779 context,
8780 after,
8781 )
8782 }
8783
8784 fn search_at_with_filters_and_context(
8785 &self,
8786 request: &crate::query::SearchRequest,
8787 snapshot: Snapshot,
8788 authorized: Option<&std::collections::HashSet<RowId>>,
8789 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8790 context: Option<&crate::query::AiExecutionContext>,
8791 after: Option<crate::query::SearchAfter>,
8792 ) -> Result<Vec<crate::query::SearchHit>> {
8793 use crate::query::{
8794 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
8795 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
8796 };
8797 let total_started = std::time::Instant::now();
8798 let rank_offset = after.map_or(0, |after| after.returned_count);
8799 self.require_select()?;
8800 if request.limit == 0 {
8801 return Err(MongrelError::InvalidArgument(
8802 "search limit must be > 0".into(),
8803 ));
8804 }
8805 if request.limit > MAX_FINAL_LIMIT {
8806 return Err(MongrelError::InvalidArgument(format!(
8807 "search limit exceeds {MAX_FINAL_LIMIT}"
8808 )));
8809 }
8810 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
8811 return Err(MongrelError::InvalidArgument(
8812 "search-after score must be finite".into(),
8813 ));
8814 }
8815 if request.retrievers.is_empty() {
8816 return Err(MongrelError::InvalidArgument(
8817 "search requires at least one retriever".into(),
8818 ));
8819 }
8820 if request.retrievers.len() > MAX_RETRIEVERS {
8821 return Err(MongrelError::InvalidArgument(format!(
8822 "search exceeds {MAX_RETRIEVERS} retrievers"
8823 )));
8824 }
8825 if request.must.len() > MAX_HARD_CONDITIONS {
8826 return Err(MongrelError::InvalidArgument(format!(
8827 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
8828 )));
8829 }
8830 for condition in &request.must {
8831 self.validate_condition(condition)?;
8832 }
8833 if request.must.iter().any(|condition| {
8834 matches!(
8835 condition,
8836 Condition::Ann { .. }
8837 | Condition::SparseMatch { .. }
8838 | Condition::MinHashSimilar { .. }
8839 )
8840 }) {
8841 return Err(MongrelError::InvalidArgument(
8842 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
8843 .into(),
8844 ));
8845 }
8846 let mut names = std::collections::HashSet::new();
8847 for named in &request.retrievers {
8848 if named.name.is_empty()
8849 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
8850 || !names.insert(named.name.as_str())
8851 {
8852 return Err(MongrelError::InvalidArgument(format!(
8853 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
8854 crate::query::MAX_RETRIEVER_NAME_BYTES
8855 )));
8856 }
8857 if !named.weight.is_finite()
8858 || named.weight < 0.0
8859 || named.weight > MAX_RETRIEVER_WEIGHT
8860 {
8861 return Err(MongrelError::InvalidArgument(format!(
8862 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8863 )));
8864 }
8865 self.validate_retriever(&named.retriever)?;
8866 }
8867 let projection = request
8868 .projection
8869 .clone()
8870 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
8871 if projection.len() > MAX_PROJECTION_COLUMNS {
8872 return Err(MongrelError::InvalidArgument(format!(
8873 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
8874 )));
8875 }
8876 for column_id in &projection {
8877 if !self
8878 .schema
8879 .columns
8880 .iter()
8881 .any(|column| column.id == *column_id)
8882 {
8883 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
8884 }
8885 }
8886 if let Some(crate::query::Rerank::ExactVector {
8887 embedding_column,
8888 query,
8889 candidate_limit,
8890 weight,
8891 ..
8892 }) = &request.rerank
8893 {
8894 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
8895 {
8896 return Err(MongrelError::InvalidArgument(format!(
8897 "rerank candidate_limit must be between search limit and {}",
8898 crate::query::MAX_RETRIEVER_K
8899 )));
8900 }
8901 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
8902 return Err(MongrelError::InvalidArgument(format!(
8903 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8904 )));
8905 }
8906 let column = self
8907 .schema
8908 .columns
8909 .iter()
8910 .find(|column| column.id == *embedding_column)
8911 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
8912 let crate::schema::TypeId::Embedding { dim } = column.ty else {
8913 return Err(MongrelError::InvalidArgument(format!(
8914 "rerank column {embedding_column} is not an embedding"
8915 )));
8916 };
8917 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
8918 return Err(MongrelError::InvalidArgument(format!(
8919 "rerank query must contain {dim} finite values"
8920 )));
8921 }
8922 }
8923
8924 let hard_filter_started = std::time::Instant::now();
8925 let hard_filter = if request.must.is_empty() {
8926 None
8927 } else {
8928 let mut sets = Vec::with_capacity(request.must.len());
8929 for condition in &request.must {
8930 if let Some(context) = context {
8931 context.checkpoint()?;
8932 }
8933 sets.push(self.resolve_condition(condition, snapshot)?);
8934 }
8935 Some(RowIdSet::intersect_many(sets))
8936 };
8937 crate::trace::QueryTrace::record(|trace| {
8938 trace.hard_filter_nanos = trace
8939 .hard_filter_nanos
8940 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
8941 });
8942 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
8943 return Ok(Vec::new());
8944 }
8945
8946 let constant = match request.fusion {
8947 Fusion::ReciprocalRank { constant } => constant,
8948 };
8949 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
8950 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
8951 let mut fusion_nanos = 0u64;
8952 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
8953 std::collections::HashMap::new();
8954 for named in retrievers {
8955 if named.weight == 0.0 {
8956 continue;
8957 }
8958 if let Some(context) = context {
8959 context.checkpoint()?;
8960 }
8961 let hits = self.retrieve_filtered(
8962 &named.retriever,
8963 snapshot,
8964 hard_filter.as_ref(),
8965 authorized,
8966 candidate_authorization,
8967 context,
8968 )?;
8969 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
8970 let fusion_started = std::time::Instant::now();
8971 for hit in hits {
8972 if let Some(context) = context {
8973 context.consume(1)?;
8974 }
8975 let contribution = named.weight / (constant as f64 + hit.rank as f64);
8976 if !contribution.is_finite() {
8977 return Err(MongrelError::InvalidArgument(
8978 "retriever contribution must be finite".into(),
8979 ));
8980 }
8981 let max_fused_candidates = context.map_or(
8982 crate::query::MAX_FUSED_CANDIDATES,
8983 crate::query::AiExecutionContext::max_fused_candidates,
8984 );
8985 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
8986 return Err(MongrelError::WorkBudgetExceeded);
8987 }
8988 let entry = fused.entry(hit.row_id).or_default();
8989 entry.0 += contribution;
8990 if !entry.0.is_finite() {
8991 return Err(MongrelError::InvalidArgument(
8992 "fused score must be finite".into(),
8993 ));
8994 }
8995 entry.1.push(ComponentScore {
8996 retriever_name: retriever_name.clone(),
8997 rank: hit.rank,
8998 raw_score: hit.score,
8999 contribution,
9000 });
9001 }
9002 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
9003 }
9004 let union_size = fused.len();
9005 let mut ranked: Vec<_> = fused
9006 .into_iter()
9007 .map(|(row_id, (fused_score, components))| {
9008 (row_id, fused_score, components, None, fused_score)
9009 })
9010 .collect();
9011 let order = |(a_row, _, _, _, a_score): &(
9012 RowId,
9013 f64,
9014 Vec<ComponentScore>,
9015 Option<f32>,
9016 f64,
9017 ),
9018 (b_row, _, _, _, b_score): &(
9019 RowId,
9020 f64,
9021 Vec<ComponentScore>,
9022 Option<f32>,
9023 f64,
9024 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
9025 if let Some(crate::query::Rerank::ExactVector {
9026 embedding_column,
9027 query,
9028 metric,
9029 candidate_limit,
9030 weight,
9031 }) = &request.rerank
9032 {
9033 let fused_order = |(a_row, a_score, ..): &(
9034 RowId,
9035 f64,
9036 Vec<ComponentScore>,
9037 Option<f32>,
9038 f64,
9039 ),
9040 (b_row, b_score, ..): &(
9041 RowId,
9042 f64,
9043 Vec<ComponentScore>,
9044 Option<f32>,
9045 f64,
9046 )| {
9047 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
9048 };
9049 let selection_started = std::time::Instant::now();
9050 if let Some(context) = context {
9051 context.consume(ranked.len())?;
9052 }
9053 if ranked.len() > *candidate_limit {
9054 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
9055 ranked.truncate(*candidate_limit);
9056 }
9057 ranked.sort_by(fused_order);
9058 fusion_nanos =
9059 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
9060 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
9061 if let Some(context) = context {
9062 context.consume(row_ids.len())?;
9063 }
9064 let query_now =
9065 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9066 let gather_started = std::time::Instant::now();
9067 let vectors = self.values_for_rids_batch_at_with_context(
9068 &row_ids,
9069 *embedding_column,
9070 snapshot,
9071 query_now,
9072 context,
9073 )?;
9074 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
9075 let vector_work =
9076 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
9077 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
9078 if let Some(context) = context {
9079 context.consume(vector_work)?;
9080 }
9081 query
9082 .iter()
9083 .map(|value| f64::from(*value).powi(2))
9084 .sum::<f64>()
9085 .sqrt()
9086 } else {
9087 0.0
9088 };
9089 let score_started = std::time::Instant::now();
9090 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
9091 for (row_id, value) in vectors {
9092 if let Some(meta) = value.generated_embedding_metadata() {
9093 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
9094 continue;
9095 }
9096 }
9097 let Some(vector) = value.as_embedding() else {
9098 continue;
9099 };
9100 let score = match metric {
9101 crate::query::VectorMetric::DotProduct => {
9102 if let Some(context) = context {
9103 context.consume(vector_work)?;
9104 }
9105 query
9106 .iter()
9107 .zip(vector)
9108 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9109 .sum::<f64>()
9110 }
9111 crate::query::VectorMetric::Cosine => {
9112 if let Some(context) = context {
9113 context.consume(vector_work.saturating_mul(2))?;
9114 }
9115 let dot = query
9116 .iter()
9117 .zip(vector)
9118 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9119 .sum::<f64>();
9120 let norm = vector
9121 .iter()
9122 .map(|value| f64::from(*value).powi(2))
9123 .sum::<f64>()
9124 .sqrt();
9125 if query_norm == 0.0 || norm == 0.0 {
9126 0.0
9127 } else {
9128 dot / (query_norm * norm)
9129 }
9130 }
9131 crate::query::VectorMetric::Euclidean => {
9132 if let Some(context) = context {
9133 context.consume(vector_work)?;
9134 }
9135 query
9136 .iter()
9137 .zip(vector)
9138 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
9139 .sum::<f64>()
9140 .sqrt()
9141 }
9142 };
9143 if !score.is_finite() {
9144 return Err(MongrelError::InvalidArgument(
9145 "exact rerank score must be finite".into(),
9146 ));
9147 }
9148 scores.insert(row_id, score as f32);
9149 }
9150 let mut reranked = Vec::with_capacity(ranked.len());
9151 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
9152 let Some(score) = scores.get(&row_id).copied() else {
9153 continue;
9154 };
9155 let ordering_score = match metric {
9156 crate::query::VectorMetric::Euclidean => -f64::from(score),
9157 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
9158 f64::from(score)
9159 }
9160 };
9161 let final_score = fused_score + *weight * ordering_score;
9162 if !final_score.is_finite() {
9163 return Err(MongrelError::InvalidArgument(
9164 "final rerank score must be finite".into(),
9165 ));
9166 }
9167 reranked.push((row_id, fused_score, components, Some(score), final_score));
9168 }
9169 ranked = reranked;
9170 ranked.sort_by(order);
9171 crate::trace::QueryTrace::record(|trace| {
9172 trace.exact_vector_gather_nanos =
9173 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
9174 trace.exact_vector_score_nanos = trace
9175 .exact_vector_score_nanos
9176 .saturating_add(score_started.elapsed().as_nanos() as u64);
9177 });
9178 }
9179 if let Some(after) = after {
9180 ranked.retain(|(row_id, _, _, _, final_score)| {
9181 final_score.total_cmp(&after.final_score).is_lt()
9182 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
9183 });
9184 }
9185 let projection_started = std::time::Instant::now();
9186 let sentinel = projection
9187 .first()
9188 .copied()
9189 .or_else(|| self.schema.columns.first().map(|column| column.id));
9190 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9191 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
9192 let mut projection_rows = 0usize;
9193 let mut projection_cells = 0usize;
9194 while out.len() < request.limit && !ranked.is_empty() {
9195 if let Some(context) = context {
9196 context.checkpoint()?;
9197 context.consume(ranked.len())?;
9198 }
9199 let needed = request.limit - out.len();
9200 let window_size = ranked
9201 .len()
9202 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
9203 let selection_started = std::time::Instant::now();
9204 let mut remainder = if ranked.len() > window_size {
9205 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
9206 ranked.split_off(window_size)
9207 } else {
9208 Vec::new()
9209 };
9210 ranked.sort_by(order);
9211 fusion_nanos =
9212 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
9213 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
9214 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
9215 if let Some(context) = context {
9216 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
9217 }
9218 projection_rows = projection_rows.saturating_add(row_ids.len());
9219 projection_cells =
9220 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
9221 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
9222 std::collections::HashMap::new();
9223 if let Some(column_id) = sentinel {
9224 for (row_id, value) in self.values_for_rids_batch_at_with_context(
9225 &row_ids, column_id, snapshot, query_now, context,
9226 )? {
9227 cells.entry(row_id).or_default().insert(column_id, value);
9228 }
9229 }
9230 for &column_id in &projection {
9231 if Some(column_id) == sentinel {
9232 continue;
9233 }
9234 for (row_id, value) in self.values_for_rids_batch_at_with_context(
9235 &row_ids, column_id, snapshot, query_now, context,
9236 )? {
9237 cells.entry(row_id).or_default().insert(column_id, value);
9238 }
9239 }
9240 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
9241 ranked.drain(..)
9242 {
9243 let Some(row_cells) = cells.remove(&row_id) else {
9244 continue;
9245 };
9246 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
9247 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
9248 out.push(SearchHit {
9249 row_id,
9250 cells: projection
9251 .iter()
9252 .filter_map(|column_id| {
9253 row_cells
9254 .get(column_id)
9255 .cloned()
9256 .map(|value| (*column_id, value))
9257 })
9258 .collect(),
9259 components,
9260 fused_score,
9261 exact_rerank_score,
9262 final_score,
9263 final_rank,
9264 });
9265 if out.len() == request.limit {
9266 break;
9267 }
9268 }
9269 ranked.append(&mut remainder);
9270 }
9271 crate::trace::QueryTrace::record(|trace| {
9272 trace.union_size = union_size;
9273 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
9274 trace.projection_nanos = trace
9275 .projection_nanos
9276 .saturating_add(projection_started.elapsed().as_nanos() as u64);
9277 trace.total_nanos = trace
9278 .total_nanos
9279 .saturating_add(total_started.elapsed().as_nanos() as u64);
9280 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
9281 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
9282 if let Some(context) = context {
9283 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
9284 }
9285 });
9286 Ok(out)
9287 }
9288
9289 pub fn set_similarity(
9292 &mut self,
9293 request: &crate::query::SetSimilarityRequest,
9294 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9295 self.set_similarity_with_allowed(request, None)
9296 }
9297
9298 pub fn set_similarity_at(
9299 &mut self,
9300 request: &crate::query::SetSimilarityRequest,
9301 snapshot: Snapshot,
9302 allowed: Option<&std::collections::HashSet<RowId>>,
9303 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9304 self.set_similarity_explained_at(request, snapshot, allowed)
9305 .map(|(hits, _)| hits)
9306 }
9307
9308 pub fn ann_rerank(
9310 &mut self,
9311 request: &crate::query::AnnRerankRequest,
9312 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9313 self.ann_rerank_with_allowed(request, None)
9314 }
9315
9316 pub fn ann_rerank_with_allowed(
9317 &mut self,
9318 request: &crate::query::AnnRerankRequest,
9319 allowed: Option<&std::collections::HashSet<RowId>>,
9320 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9321 self.ann_rerank_at(request, self.snapshot(), allowed)
9322 }
9323
9324 pub fn ann_rerank_at(
9325 &mut self,
9326 request: &crate::query::AnnRerankRequest,
9327 snapshot: Snapshot,
9328 allowed: Option<&std::collections::HashSet<RowId>>,
9329 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9330 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
9331 }
9332
9333 pub fn ann_rerank_at_with_context(
9334 &mut self,
9335 request: &crate::query::AnnRerankRequest,
9336 snapshot: Snapshot,
9337 allowed: Option<&std::collections::HashSet<RowId>>,
9338 context: Option<&crate::query::AiExecutionContext>,
9339 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9340 self.ensure_indexes_complete()?;
9341 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
9342 }
9343
9344 pub fn ann_rerank_at_with_candidate_authorization_and_context(
9345 &mut self,
9346 request: &crate::query::AnnRerankRequest,
9347 snapshot: Snapshot,
9348 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9349 context: Option<&crate::query::AiExecutionContext>,
9350 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9351 self.ensure_indexes_complete()?;
9352 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
9353 }
9354
9355 #[doc(hidden)]
9356 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
9357 &self,
9358 request: &crate::query::AnnRerankRequest,
9359 snapshot: Snapshot,
9360 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9361 context: Option<&crate::query::AiExecutionContext>,
9362 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9363 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
9364 }
9365
9366 fn ann_rerank_at_with_filters_and_context(
9367 &self,
9368 request: &crate::query::AnnRerankRequest,
9369 snapshot: Snapshot,
9370 allowed: Option<&std::collections::HashSet<RowId>>,
9371 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9372 context: Option<&crate::query::AiExecutionContext>,
9373 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9374 use crate::query::{
9375 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
9376 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
9377 };
9378 if request.candidate_k == 0 || request.limit == 0 {
9379 return Err(MongrelError::InvalidArgument(
9380 "candidate_k and limit must be > 0".into(),
9381 ));
9382 }
9383 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
9384 return Err(MongrelError::InvalidArgument(format!(
9385 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
9386 )));
9387 }
9388 let retriever = Retriever::Ann {
9389 column_id: request.column_id,
9390 query: request.query.clone(),
9391 k: request.candidate_k,
9392 };
9393 self.require_select()?;
9394 self.validate_retriever(&retriever)?;
9395 let hits = self.retrieve_filtered(
9396 &retriever,
9397 snapshot,
9398 None,
9399 allowed,
9400 candidate_authorization,
9401 context,
9402 )?;
9403 let distances: std::collections::HashMap<_, _> = hits
9404 .iter()
9405 .filter_map(|hit| match hit.score {
9406 RetrieverScore::AnnHammingDistance(distance) => {
9407 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
9408 }
9409 RetrieverScore::AnnCosineDistance(distance) => {
9410 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
9411 }
9412 _ => None,
9413 })
9414 .collect();
9415 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
9416 if let Some(context) = context {
9417 context.consume(row_ids.len())?;
9418 }
9419 let gather_started = std::time::Instant::now();
9420 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9421 let values = self.values_for_rids_batch_at_with_context(
9422 &row_ids,
9423 request.column_id,
9424 snapshot,
9425 query_now,
9426 context,
9427 )?;
9428 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
9429 let score_started = std::time::Instant::now();
9430 let vector_work =
9431 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
9432 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
9433 if let Some(context) = context {
9434 context.consume(vector_work)?;
9435 }
9436 request
9437 .query
9438 .iter()
9439 .map(|value| f64::from(*value).powi(2))
9440 .sum::<f64>()
9441 .sqrt()
9442 } else {
9443 0.0
9444 };
9445 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
9446 for (row_id, value) in values {
9447 if let Some(meta) = value.generated_embedding_metadata() {
9448 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
9449 continue;
9450 }
9451 }
9452 let Some(vector) = value.as_embedding() else {
9453 continue;
9454 };
9455 let exact_score = match request.metric {
9456 VectorMetric::DotProduct => {
9457 if let Some(context) = context {
9458 context.consume(vector_work)?;
9459 }
9460 request
9461 .query
9462 .iter()
9463 .zip(vector)
9464 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9465 .sum::<f64>()
9466 }
9467 VectorMetric::Cosine => {
9468 if let Some(context) = context {
9469 context.consume(vector_work.saturating_mul(2))?;
9470 }
9471 let dot = request
9472 .query
9473 .iter()
9474 .zip(vector)
9475 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9476 .sum::<f64>();
9477 let norm = vector
9478 .iter()
9479 .map(|value| f64::from(*value).powi(2))
9480 .sum::<f64>()
9481 .sqrt();
9482 if query_norm == 0.0 || norm == 0.0 {
9483 0.0
9484 } else {
9485 dot / (query_norm * norm)
9486 }
9487 }
9488 VectorMetric::Euclidean => {
9489 if let Some(context) = context {
9490 context.consume(vector_work)?;
9491 }
9492 request
9493 .query
9494 .iter()
9495 .zip(vector)
9496 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
9497 .sum::<f64>()
9498 .sqrt()
9499 }
9500 };
9501 let exact_score = exact_score as f32;
9502 if !exact_score.is_finite() {
9503 return Err(MongrelError::InvalidArgument(
9504 "exact ANN score must be finite".into(),
9505 ));
9506 }
9507 let Some(candidate_distance) = distances.get(&row_id).copied() else {
9508 continue;
9509 };
9510 reranked.push(AnnRerankHit {
9511 row_id,
9512 candidate_distance,
9513 exact_score,
9514 });
9515 }
9516 reranked.sort_by(|left, right| {
9517 let score = match request.metric {
9518 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
9519 VectorMetric::Cosine | VectorMetric::DotProduct => {
9520 right.exact_score.total_cmp(&left.exact_score)
9521 }
9522 };
9523 score.then_with(|| left.row_id.cmp(&right.row_id))
9524 });
9525 reranked.truncate(request.limit);
9526 crate::trace::QueryTrace::record(|trace| {
9527 trace.exact_vector_gather_nanos =
9528 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
9529 trace.exact_vector_score_nanos = trace
9530 .exact_vector_score_nanos
9531 .saturating_add(score_started.elapsed().as_nanos() as u64);
9532 });
9533 Ok(reranked)
9534 }
9535
9536 pub fn set_similarity_with_allowed(
9537 &mut self,
9538 request: &crate::query::SetSimilarityRequest,
9539 allowed: Option<&std::collections::HashSet<RowId>>,
9540 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9541 self.set_similarity_explained_at(request, self.snapshot(), allowed)
9542 .map(|(hits, _)| hits)
9543 }
9544
9545 pub fn set_similarity_explained(
9546 &mut self,
9547 request: &crate::query::SetSimilarityRequest,
9548 ) -> Result<(
9549 Vec<crate::query::SetSimilarityHit>,
9550 crate::query::SetSimilarityTrace,
9551 )> {
9552 self.set_similarity_explained_at(request, self.snapshot(), None)
9553 }
9554
9555 fn set_similarity_explained_at(
9556 &mut self,
9557 request: &crate::query::SetSimilarityRequest,
9558 snapshot: Snapshot,
9559 allowed: Option<&std::collections::HashSet<RowId>>,
9560 ) -> Result<(
9561 Vec<crate::query::SetSimilarityHit>,
9562 crate::query::SetSimilarityTrace,
9563 )> {
9564 self.ensure_indexes_complete()?;
9565 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
9566 }
9567
9568 pub fn set_similarity_at_with_context(
9569 &mut self,
9570 request: &crate::query::SetSimilarityRequest,
9571 snapshot: Snapshot,
9572 allowed: Option<&std::collections::HashSet<RowId>>,
9573 context: Option<&crate::query::AiExecutionContext>,
9574 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9575 self.ensure_indexes_complete()?;
9576 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
9577 .map(|(hits, _)| hits)
9578 }
9579
9580 pub fn set_similarity_at_with_candidate_authorization_and_context(
9581 &mut self,
9582 request: &crate::query::SetSimilarityRequest,
9583 snapshot: Snapshot,
9584 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9585 context: Option<&crate::query::AiExecutionContext>,
9586 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9587 self.ensure_indexes_complete()?;
9588 self.set_similarity_explained_at_with_context(
9589 request,
9590 snapshot,
9591 None,
9592 authorization,
9593 context,
9594 )
9595 .map(|(hits, _)| hits)
9596 }
9597
9598 #[doc(hidden)]
9599 pub fn set_similarity_at_with_candidate_authorization_on_generation(
9600 &self,
9601 request: &crate::query::SetSimilarityRequest,
9602 snapshot: Snapshot,
9603 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9604 context: Option<&crate::query::AiExecutionContext>,
9605 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9606 self.set_similarity_explained_at_with_context(
9607 request,
9608 snapshot,
9609 None,
9610 authorization,
9611 context,
9612 )
9613 .map(|(hits, _)| hits)
9614 }
9615
9616 fn set_similarity_explained_at_with_context(
9617 &self,
9618 request: &crate::query::SetSimilarityRequest,
9619 snapshot: Snapshot,
9620 allowed: Option<&std::collections::HashSet<RowId>>,
9621 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9622 context: Option<&crate::query::AiExecutionContext>,
9623 ) -> Result<(
9624 Vec<crate::query::SetSimilarityHit>,
9625 crate::query::SetSimilarityTrace,
9626 )> {
9627 use crate::query::{
9628 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
9629 MAX_SET_MEMBERS,
9630 };
9631 let mut trace = crate::query::SetSimilarityTrace::default();
9632 if request.members.is_empty() {
9633 return Ok((Vec::new(), trace));
9634 }
9635 if request.candidate_k == 0 || request.limit == 0 {
9636 return Err(MongrelError::InvalidArgument(
9637 "candidate_k and limit must be > 0".into(),
9638 ));
9639 }
9640 if request.candidate_k > MAX_RETRIEVER_K
9641 || request.limit > MAX_FINAL_LIMIT
9642 || request.members.len() > MAX_SET_MEMBERS
9643 {
9644 return Err(MongrelError::InvalidArgument(format!(
9645 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
9646 )));
9647 }
9648 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
9649 return Err(MongrelError::InvalidArgument(
9650 "min_jaccard must be finite and between 0 and 1".into(),
9651 ));
9652 }
9653 let started = std::time::Instant::now();
9654 let retriever = Retriever::MinHash {
9655 column_id: request.column_id,
9656 members: request.members.clone(),
9657 k: request.candidate_k,
9658 };
9659 self.require_select()?;
9660 self.validate_retriever(&retriever)?;
9661 let hits = self.retrieve_filtered(
9662 &retriever,
9663 snapshot,
9664 None,
9665 allowed,
9666 candidate_authorization,
9667 context,
9668 )?;
9669 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
9670 trace.candidate_count = hits.len();
9671 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
9672 if let Some(context) = context {
9673 context.consume(row_ids.len())?;
9674 }
9675 let started = std::time::Instant::now();
9676 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9677 let values = self.values_for_rids_batch_at_with_context(
9678 &row_ids,
9679 request.column_id,
9680 snapshot,
9681 query_now,
9682 context,
9683 )?;
9684 trace.gather_us = started.elapsed().as_micros() as u64;
9685 if let Some(context) = context {
9686 context.consume(request.members.len())?;
9687 }
9688 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
9689 let estimates: std::collections::HashMap<_, _> = hits
9690 .into_iter()
9691 .filter_map(|hit| match hit.score {
9692 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
9693 _ => None,
9694 })
9695 .collect();
9696 let started = std::time::Instant::now();
9697 let mut parsed = Vec::with_capacity(values.len());
9698 for (row_id, value) in values {
9699 let Value::Bytes(bytes) = value else {
9700 continue;
9701 };
9702 if let Some(context) = context {
9703 context.consume(crate::query::work_units(
9704 bytes.len(),
9705 crate::query::PARSE_WORK_QUANTUM,
9706 ))?;
9707 }
9708 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
9709 continue;
9710 };
9711 if let Some(context) = context {
9712 context.consume(members.len())?;
9713 }
9714 let stored = members
9715 .into_iter()
9716 .filter_map(|member| match member {
9717 serde_json::Value::String(value) => {
9718 Some(crate::query::SetMember::String(value))
9719 }
9720 serde_json::Value::Number(value) => {
9721 Some(crate::query::SetMember::Number(value))
9722 }
9723 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
9724 _ => None,
9725 })
9726 .collect::<std::collections::HashSet<_>>();
9727 parsed.push((row_id, stored));
9728 }
9729 trace.parse_us = started.elapsed().as_micros() as u64;
9730 trace.verified_count = parsed.len();
9731 let started = std::time::Instant::now();
9732 let mut exact = Vec::new();
9733 for (row_id, stored) in parsed {
9734 if let Some(context) = context {
9735 context.consume(query.len().saturating_add(stored.len()))?;
9736 }
9737 let union = query.union(&stored).count();
9738 let score = if union == 0 {
9739 1.0
9740 } else {
9741 query.intersection(&stored).count() as f32 / union as f32
9742 };
9743 if score >= request.min_jaccard {
9744 exact.push(SetSimilarityHit {
9745 row_id,
9746 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
9747 exact_jaccard: score,
9748 });
9749 }
9750 }
9751 exact.sort_by(|a, b| {
9752 b.exact_jaccard
9753 .total_cmp(&a.exact_jaccard)
9754 .then_with(|| a.row_id.cmp(&b.row_id))
9755 });
9756 exact.truncate(request.limit);
9757 trace.score_us = started.elapsed().as_micros() as u64;
9758 crate::trace::QueryTrace::record(|query_trace| {
9759 query_trace.exact_set_gather_nanos = query_trace
9760 .exact_set_gather_nanos
9761 .saturating_add(trace.gather_us.saturating_mul(1_000));
9762 query_trace.exact_set_parse_nanos = query_trace
9763 .exact_set_parse_nanos
9764 .saturating_add(trace.parse_us.saturating_mul(1_000));
9765 query_trace.exact_set_score_nanos = query_trace
9766 .exact_set_score_nanos
9767 .saturating_add(trace.score_us.saturating_mul(1_000));
9768 });
9769 Ok((exact, trace))
9770 }
9771
9772 fn values_for_rids_batch_at(
9774 &self,
9775 row_ids: &[u64],
9776 column_id: u16,
9777 snapshot: Snapshot,
9778 now: i64,
9779 ) -> Result<Vec<(RowId, Value)>> {
9780 if self.ttl.is_none()
9781 && self.memtable.is_empty()
9782 && self.mutable_run.is_empty()
9783 && self.run_refs.len() == 1
9784 {
9785 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9786 if row_ids.len().saturating_mul(24) < reader.row_count() {
9791 let mut values = Vec::with_capacity(row_ids.len());
9792 for &raw_row_id in row_ids {
9793 let row_id = RowId(raw_row_id);
9794 if let Some((_, false, Some(value))) =
9795 reader.get_version_column_at(row_id, snapshot, column_id)?
9796 {
9797 values.push((row_id, value));
9798 }
9799 }
9800 return Ok(values);
9801 }
9802 let (positions, visible_row_ids) = reader.visible_positions_with_rids_at(snapshot)?;
9803 let requested: Vec<(RowId, usize)> = row_ids
9804 .iter()
9805 .filter_map(|raw| {
9806 visible_row_ids
9807 .binary_search(&(*raw as i64))
9808 .ok()
9809 .map(|index| (RowId(*raw), positions[index]))
9810 })
9811 .collect();
9812 let values = reader.gather_column(
9813 column_id,
9814 &requested
9815 .iter()
9816 .map(|(_, position)| *position)
9817 .collect::<Vec<_>>(),
9818 )?;
9819 return Ok(requested
9820 .into_iter()
9821 .zip(values)
9822 .map(|((row_id, _), value)| (row_id, value))
9823 .collect());
9824 }
9825 self.values_for_rids_at(row_ids, column_id, snapshot, now)
9826 }
9827
9828 fn values_for_rids_batch_at_with_context(
9829 &self,
9830 row_ids: &[u64],
9831 column_id: u16,
9832 snapshot: Snapshot,
9833 now: i64,
9834 context: Option<&crate::query::AiExecutionContext>,
9835 ) -> Result<Vec<(RowId, Value)>> {
9836 let Some(context) = context else {
9837 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
9838 };
9839 let mut values = Vec::with_capacity(row_ids.len());
9840 for chunk in row_ids.chunks(256) {
9841 context.checkpoint()?;
9842 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
9843 }
9844 Ok(values)
9845 }
9846
9847 fn values_for_rids_at(
9849 &self,
9850 row_ids: &[u64],
9851 column_id: u16,
9852 snapshot: Snapshot,
9853 now: i64,
9854 ) -> Result<Vec<(RowId, Value)>> {
9855 let mut readers: Vec<_> = self
9856 .run_refs
9857 .iter()
9858 .map(|run| self.open_reader(run.run_id))
9859 .collect::<Result<_>>()?;
9860 let mut out = Vec::with_capacity(row_ids.len());
9861 for &raw_row_id in row_ids {
9862 let row_id = RowId(raw_row_id);
9863 let mem = self.memtable.get_version_at(row_id, snapshot);
9864 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
9865 let overlay = match (mem, mutable) {
9866 (Some((_, a)), Some((_, b))) => Some(
9867 if Snapshot::version_is_newer(
9868 a.committed_epoch,
9869 a.commit_ts,
9870 b.committed_epoch,
9871 b.commit_ts,
9872 ) {
9873 a
9874 } else {
9875 b
9876 },
9877 ),
9878 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
9879 (None, None) => None,
9880 };
9881 if let Some(row) = overlay {
9882 if !row.deleted && !self.row_expired_at(&row, now) {
9883 if let Some(value) = row.columns.get(&column_id) {
9884 out.push((row_id, value.clone()));
9885 }
9886 }
9887 continue;
9888 }
9889
9890 let mut best: Option<(crate::sorted_run::VersionedColumnValue, usize)> = None;
9895 for (index, reader) in readers.iter_mut().enumerate() {
9896 if let Some(versioned) =
9897 reader.get_version_column_full_at(row_id, snapshot, column_id)?
9898 {
9899 if best
9900 .as_ref()
9901 .map(|(current, _)| versioned.stamp.is_newer_than(current.stamp))
9902 .unwrap_or(true)
9903 {
9904 best = Some((versioned, index));
9905 }
9906 }
9907 }
9908 let Some((candidate, reader_index)) = best else {
9909 continue;
9910 };
9911 if candidate.deleted {
9912 continue;
9913 }
9914 let Some(value) = candidate.value.clone() else {
9915 continue;
9916 };
9917 if let Some(ttl) = self.ttl {
9918 if ttl.column_id != column_id {
9919 if let Some(ttl_value) = readers[reader_index].get_version_column_full_at(
9920 row_id,
9921 snapshot,
9922 ttl.column_id,
9923 )? {
9924 if let Some(Value::Int64(timestamp)) = ttl_value.value {
9925 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9926 continue;
9927 }
9928 }
9929 }
9930 } else if let Value::Int64(timestamp) = &value {
9931 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9932 continue;
9933 }
9934 }
9935 }
9936 out.push((row_id, value));
9937 }
9938 Ok(out)
9939 }
9940
9941 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
9946 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
9947 }
9948
9949 pub fn rows_for_rids_with_context(
9950 &self,
9951 rids: &[u64],
9952 snapshot: Snapshot,
9953 context: &crate::query::AiExecutionContext,
9954 ) -> Result<Vec<Row>> {
9955 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
9956 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
9957 }
9958
9959 fn rows_for_rids_at_time(
9960 &self,
9961 rids: &[u64],
9962 snapshot: Snapshot,
9963 ttl_now: i64,
9964 control: Option<&crate::ExecutionControl>,
9965 ) -> Result<Vec<Row>> {
9966 use std::collections::HashMap;
9967 let mut rows = Vec::with_capacity(rids.len());
9968 let tier_size = self.memtable.len() + self.mutable_run.len();
9983 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
9984 if rids.len().saturating_mul(24) < tier_size {
9985 for &rid in rids {
9986 if overlay.len() & 255 == 0 {
9987 control
9988 .map(crate::ExecutionControl::checkpoint)
9989 .transpose()?;
9990 }
9991 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
9992 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
9993 let newest = match (mem, mrun) {
9994 (Some((_, mr)), Some((_, rr))) => Some(
9995 if Snapshot::version_is_newer(
9996 mr.committed_epoch,
9997 mr.commit_ts,
9998 rr.committed_epoch,
9999 rr.commit_ts,
10000 ) {
10001 mr
10002 } else {
10003 rr
10004 },
10005 ),
10006 (Some((_, mr)), None) => Some(mr),
10007 (None, Some((_, rr))) => Some(rr),
10008 (None, None) => None,
10009 };
10010 if let Some(row) = newest {
10011 overlay.insert(rid, row);
10012 }
10013 }
10014 } else {
10015 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
10016 overlay
10017 .entry(row.row_id.0)
10018 .and_modify(|e| {
10019 if Snapshot::version_is_newer(
10020 row.committed_epoch,
10021 row.commit_ts,
10022 e.committed_epoch,
10023 e.commit_ts,
10024 ) {
10025 *e = row.clone();
10026 }
10027 })
10028 .or_insert(row);
10029 };
10030 for (index, row) in self
10031 .memtable
10032 .visible_versions_at(snapshot)
10033 .into_iter()
10034 .enumerate()
10035 {
10036 if index & 255 == 0 {
10037 control
10038 .map(crate::ExecutionControl::checkpoint)
10039 .transpose()?;
10040 }
10041 fold_newest(row, &mut overlay);
10042 }
10043 for (index, row) in self
10044 .mutable_run
10045 .visible_versions_at(snapshot)
10046 .into_iter()
10047 .enumerate()
10048 {
10049 if index & 255 == 0 {
10050 control
10051 .map(crate::ExecutionControl::checkpoint)
10052 .transpose()?;
10053 }
10054 fold_newest(row, &mut overlay);
10055 }
10056 }
10057 if self.run_refs.len() == 1 {
10058 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
10059 if rids.len().saturating_mul(24) < reader.row_count() {
10067 for (index, &rid) in rids.iter().enumerate() {
10068 if index & 255 == 0 {
10069 control
10070 .map(crate::ExecutionControl::checkpoint)
10071 .transpose()?;
10072 }
10073 if let Some(r) = overlay.get(&rid) {
10074 if !r.deleted {
10075 rows.push(r.clone());
10076 }
10077 continue;
10078 }
10079 if let Some((_, row)) = reader.get_version_at(RowId(rid), snapshot)? {
10080 if !row.deleted {
10081 rows.push(row);
10082 }
10083 }
10084 }
10085 rows.retain(|row| !self.row_expired_at(row, ttl_now));
10086 return Ok(rows);
10087 }
10088 let (positions, vis_rids) = reader.visible_positions_with_rids_at(snapshot)?;
10097 enum Src {
10100 Overlay,
10101 Run,
10102 }
10103 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
10104 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
10105 for (index, rid) in rids.iter().enumerate() {
10106 if index & 255 == 0 {
10107 control
10108 .map(crate::ExecutionControl::checkpoint)
10109 .transpose()?;
10110 }
10111 if overlay.contains_key(rid) {
10112 plan.push(Src::Overlay);
10113 continue;
10114 }
10115 match vis_rids.binary_search(&(*rid as i64)) {
10116 Ok(i) => {
10117 plan.push(Src::Run);
10118 fetch.push(positions[i]);
10119 }
10120 Err(_) => { }
10121 }
10122 }
10123 let fetched = reader.materialize_batch(&fetch)?;
10124 let mut fetched_iter = fetched.into_iter();
10125 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
10126 if index & 255 == 0 {
10127 control
10128 .map(crate::ExecutionControl::checkpoint)
10129 .transpose()?;
10130 }
10131 match src {
10132 Src::Overlay => {
10133 if let Some(r) = overlay.get(rid) {
10134 if !r.deleted {
10135 rows.push(r.clone());
10136 }
10137 }
10138 }
10139 Src::Run => {
10140 if let Some(row) = fetched_iter.next() {
10141 if !row.deleted {
10142 rows.push(row);
10143 }
10144 }
10145 }
10146 }
10147 }
10148 rows.retain(|row| !self.row_expired_at(row, ttl_now));
10149 return Ok(rows);
10150 }
10151 let mut readers: Vec<_> = self
10155 .run_refs
10156 .iter()
10157 .map(|rr| self.open_reader(rr.run_id))
10158 .collect::<Result<Vec<_>>>()?;
10159 for (index, rid) in rids.iter().enumerate() {
10160 if index & 255 == 0 {
10161 control
10162 .map(crate::ExecutionControl::checkpoint)
10163 .transpose()?;
10164 }
10165 if let Some(r) = overlay.get(rid) {
10166 if !r.deleted {
10167 rows.push(r.clone());
10168 }
10169 continue;
10170 }
10171 let mut best: Option<(VersionStamp, Row)> = None;
10176 for reader in readers.iter_mut() {
10177 if let Ok(Some((stamp, row))) = reader.get_version_stamp_at(RowId(*rid), snapshot) {
10178 if best
10179 .as_ref()
10180 .map(|(current, _)| stamp.is_newer_than(*current))
10181 .unwrap_or(true)
10182 {
10183 best = Some((stamp, row));
10184 }
10185 }
10186 }
10187 if let Some((_, r)) = best {
10188 if !r.deleted {
10189 rows.push(r);
10190 }
10191 }
10192 }
10193 rows.retain(|row| !self.row_expired_at(row, ttl_now));
10194 Ok(rows)
10195 }
10196
10197 pub fn indexes_complete(&self) -> bool {
10207 self.indexes_complete
10208 }
10209
10210 pub fn index_build_policy(&self) -> IndexBuildPolicy {
10212 self.index_build_policy
10213 }
10214
10215 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
10219 self.index_build_policy = policy;
10220 }
10221
10222 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
10227 if !self.indexes_complete {
10231 return None;
10232 }
10233 let b = self.bitmap.get(&column_id)?;
10234 let result: Vec<Vec<u8>> = b
10235 .keys()
10236 .into_iter()
10237 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
10238 .collect();
10239 Some(result)
10240 }
10241
10242 pub fn fk_join_row_ids(
10243 &self,
10244 fk_column_id: u16,
10245 pk_values: &[Vec<u8>],
10246 fk_conditions: &[crate::query::Condition],
10247 snapshot: Snapshot,
10248 ) -> Result<Vec<u64>> {
10249 let Some(b) = self.bitmap.get(&fk_column_id) else {
10250 return Ok(Vec::new());
10251 };
10252 let mut join_set = {
10253 let mut acc = roaring::RoaringBitmap::new();
10254 for v in pk_values {
10255 acc |= b.get(v);
10256 }
10257 RowIdSet::from_roaring(acc)
10258 };
10259 if !fk_conditions.is_empty() {
10260 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
10261 sets.push(join_set);
10262 for c in fk_conditions {
10263 sets.push(self.resolve_condition(c, snapshot)?);
10264 }
10265 join_set = RowIdSet::intersect_many(sets);
10266 }
10267 Ok(join_set.into_sorted_vec())
10268 }
10269
10270 pub fn fk_join_count(
10276 &self,
10277 fk_column_id: u16,
10278 pk_values: &[Vec<u8>],
10279 fk_conditions: &[crate::query::Condition],
10280 snapshot: Snapshot,
10281 ) -> Result<u64> {
10282 let Some(b) = self.bitmap.get(&fk_column_id) else {
10283 return Ok(0);
10284 };
10285 let mut acc = roaring::RoaringBitmap::new();
10286 for v in pk_values {
10287 acc |= b.get(v);
10288 }
10289 if fk_conditions.is_empty() {
10290 return Ok(acc.len());
10291 }
10292 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
10293 sets.push(RowIdSet::from_roaring(acc));
10294 for c in fk_conditions {
10295 sets.push(self.resolve_condition(c, snapshot)?);
10296 }
10297 Ok(RowIdSet::intersect_many(sets).len() as u64)
10298 }
10299
10300 fn inspect_row_at(&self, rid: RowId, snapshot: Snapshot) -> Option<crate::memtable::Row> {
10305 let mut best: Option<crate::memtable::Row> = None;
10306 let mut consider = |row: crate::memtable::Row| {
10307 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
10308 return;
10309 }
10310 if best.as_ref().is_none_or(|current| {
10311 Snapshot::version_is_newer(
10312 row.committed_epoch,
10313 row.commit_ts,
10314 current.committed_epoch,
10315 current.commit_ts,
10316 )
10317 }) {
10318 best = Some(row);
10319 }
10320 };
10321 if let Some((_, row)) = self.memtable.get_version_at(rid, snapshot) {
10322 consider(row);
10323 }
10324 if let Some((_, row)) = self.mutable_run.get_version_at(rid, snapshot) {
10325 consider(row);
10326 }
10327 for rr in &self.run_refs {
10328 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
10329 if rid.0 < min_rid || rid.0 > max_rid {
10330 continue;
10331 }
10332 }
10333 let Ok(mut reader) = self.open_reader(rr.run_id) else {
10334 continue;
10335 };
10336 let Ok(Some((_stamp, row))) = reader.get_version_stamp_at(rid, snapshot) else {
10340 continue;
10341 };
10342 consider(row);
10343 }
10344 best
10345 }
10346
10347 fn resolve_pk_with_hot_fallback(
10348 &self,
10349 pk_column_id: u16,
10350 lookup: &[u8],
10351 snapshot: Snapshot,
10352 ) -> Result<RowIdSet> {
10353 let snapshot = if self.is_shared()
10366 || (self.pending_put_cols.is_empty() && self.pending_delete_rids.is_empty())
10367 || snapshot.epoch.0 != self.epoch.visible().0
10368 {
10369 snapshot
10370 } else {
10371 Snapshot::at(self.pending_epoch())
10372 };
10373 if !self.indexes_complete {
10377 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10378 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::IndexIncomplete);
10379 return Ok(result);
10380 }
10381 let mapped = self.hot.get(lookup);
10384 let Some(r) = mapped else {
10385 let (result, reason) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10386 self.record_hot_fallback_reason(reason);
10387 return Ok(result);
10388 };
10389 if snapshot.epoch < self.current_epoch() {
10394 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::HistoricalSnapshot);
10395 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10396 return Ok(result);
10397 }
10398 let materialized = self.inspect_row_at(r, snapshot);
10402 let now_nanos = unix_nanos_now();
10403 let inspection = crate::trace::inspect_hot_candidate(
10404 materialized.as_ref(),
10405 snapshot,
10406 self.ttl,
10407 now_nanos,
10408 pk_column_id,
10409 lookup,
10410 |row| {
10411 let pk_value = row.columns.get(&pk_column_id);
10412 pk_value
10413 .map(|v| self.index_lookup_key(pk_column_id, v))
10414 .unwrap_or_default()
10415 },
10416 );
10417 match inspection {
10418 crate::trace::HotCandidateInspection::Hit(row) => {
10419 self.lookup_metrics
10420 .hot_lookup_hit
10421 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10422 crate::trace::QueryTrace::record(|t| {
10423 t.hot_lookup_attempted = true;
10424 t.hot_lookup_hit = true;
10425 });
10426 Ok(RowIdSet::one(row.row_id.0))
10427 }
10428 failure => {
10429 let reason = match &failure {
10430 crate::trace::HotCandidateInspection::Hit(_) => unreachable!(),
10431 crate::trace::HotCandidateInspection::MissingRow => {
10432 crate::trace::HotFallbackReason::Tombstone
10437 }
10438 crate::trace::HotCandidateInspection::Invisible => {
10439 crate::trace::HotFallbackReason::InvisibleAtSnapshot
10440 }
10441 crate::trace::HotCandidateInspection::Tombstone => {
10442 crate::trace::HotFallbackReason::Tombstone
10443 }
10444 crate::trace::HotCandidateInspection::TtlExpired => {
10445 crate::trace::HotFallbackReason::TtlExpired
10446 }
10447 crate::trace::HotCandidateInspection::PrimaryKeyMismatch { .. } => {
10448 crate::trace::HotFallbackReason::PrimaryKeyMismatch
10449 }
10450 };
10451 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10457 self.record_hot_fallback_reason(reason);
10458 Ok(result)
10459 }
10460 }
10461 }
10462
10463 fn resolve_condition(
10468 &self,
10469 c: &crate::query::Condition,
10470 snapshot: Snapshot,
10471 ) -> Result<RowIdSet> {
10472 self.resolve_condition_with_allowed(c, snapshot, None)
10473 }
10474
10475 fn resolve_condition_with_allowed(
10476 &self,
10477 c: &crate::query::Condition,
10478 snapshot: Snapshot,
10479 allowed: Option<&std::collections::HashSet<RowId>>,
10480 ) -> Result<RowIdSet> {
10481 use crate::query::Condition;
10482 self.validate_condition(c)?;
10483 Ok(match c {
10484 Condition::Pk(key) => {
10485 let lookup = self
10486 .schema
10487 .primary_key()
10488 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
10489 .unwrap_or_else(|| key.clone());
10490 if let Some(pk_col) = self.schema.primary_key() {
10491 return self.resolve_pk_with_hot_fallback(pk_col.id, &lookup, snapshot);
10492 }
10493 RowIdSet::empty()
10497 }
10498 Condition::BitmapEq { column_id, value } => {
10499 let lookup = self.index_lookup_key_bytes(*column_id, value);
10500 self.bitmap
10501 .get(column_id)
10502 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
10503 .unwrap_or_else(RowIdSet::empty)
10504 }
10505 Condition::BitmapIn { column_id, values } => {
10506 let bm = self.bitmap.get(column_id);
10507 let mut acc = roaring::RoaringBitmap::new();
10508 if let Some(b) = bm {
10509 for v in values {
10510 let lookup = self.index_lookup_key_bytes(*column_id, v);
10511 acc |= b.get(&lookup);
10512 }
10513 }
10514 RowIdSet::from_roaring(acc)
10515 }
10516 Condition::BytesPrefix { column_id, prefix } => {
10517 if let Some(b) = self.bitmap.get(column_id) {
10522 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
10523 let mut acc = roaring::RoaringBitmap::new();
10524 for key in b.keys() {
10525 if key.starts_with(&lookup_prefix) {
10526 acc |= b.get(&key);
10527 }
10528 }
10529 RowIdSet::from_roaring(acc)
10530 } else {
10531 RowIdSet::empty()
10532 }
10533 }
10534 Condition::FmContains { column_id, pattern } => self
10535 .fm
10536 .get(column_id)
10537 .map(|f| {
10538 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
10539 })
10540 .unwrap_or_else(RowIdSet::empty),
10541 Condition::FmContainsAll {
10542 column_id,
10543 patterns,
10544 } => {
10545 if let Some(f) = self.fm.get(column_id) {
10548 let sets: Vec<RowIdSet> = patterns
10549 .iter()
10550 .map(|pat| {
10551 RowIdSet::from_unsorted(
10552 f.locate(pat).into_iter().map(|r| r.0).collect(),
10553 )
10554 })
10555 .collect();
10556 RowIdSet::intersect_many(sets)
10557 } else {
10558 RowIdSet::empty()
10559 }
10560 }
10561 Condition::Ann {
10562 column_id,
10563 query,
10564 k,
10565 } => RowIdSet::from_unsorted(
10566 self.retrieve_filtered(
10567 &crate::query::Retriever::Ann {
10568 column_id: *column_id,
10569 query: query.clone(),
10570 k: *k,
10571 },
10572 snapshot,
10573 None,
10574 allowed,
10575 None,
10576 None,
10577 )?
10578 .into_iter()
10579 .map(|hit| hit.row_id.0)
10580 .collect(),
10581 ),
10582 Condition::SparseMatch {
10583 column_id,
10584 query,
10585 k,
10586 } => RowIdSet::from_unsorted(
10587 self.retrieve_filtered(
10588 &crate::query::Retriever::Sparse {
10589 column_id: *column_id,
10590 query: query.clone(),
10591 k: *k,
10592 },
10593 snapshot,
10594 None,
10595 allowed,
10596 None,
10597 None,
10598 )?
10599 .into_iter()
10600 .map(|hit| hit.row_id.0)
10601 .collect(),
10602 ),
10603 Condition::MinHashSimilar {
10604 column_id,
10605 query,
10606 k,
10607 } => match self.minhash.get(column_id) {
10608 Some(index) => {
10609 let candidates = index.candidate_row_ids(query);
10610 let eligible =
10611 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
10612 RowIdSet::from_unsorted(
10613 index
10614 .search_filtered(query, *k, |row_id| {
10615 eligible.contains(&row_id)
10616 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
10617 })
10618 .into_iter()
10619 .map(|(row_id, _)| row_id.0)
10620 .collect(),
10621 )
10622 }
10623 None => RowIdSet::empty(),
10624 },
10625 Condition::Range { column_id, lo, hi } => {
10626 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10645 if self.run_refs.len() == 1 {
10646 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10649 } else {
10650 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10658 if lo == hi {
10659 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10660 }
10661 return Ok(multi);
10662 }
10663 } else if self.run_refs.len() == 1 {
10664 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10665 r.range_row_id_set_i64(*column_id, *lo, *hi)?
10666 } else {
10667 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10670 if lo == hi {
10671 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10672 }
10673 return Ok(multi);
10674 };
10675 if self.learned_range.get(column_id).is_none() && self.run_refs.len() == 1 {
10676 let candidates: Vec<RowId> =
10677 set.into_sorted_vec().into_iter().map(RowId).collect();
10678 set = RowIdSet::from_unsorted(
10679 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?
10680 .into_iter()
10681 .map(|row_id| row_id.0)
10682 .collect(),
10683 );
10684 }
10685 set.remove_many(self.overlay_rid_set(snapshot));
10686 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
10687 if lo == hi {
10688 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
10689 }
10690 RowIdSet::from_unsorted(
10691 set.into_sorted_vec()
10692 .into_iter()
10693 .filter(|row_id| self.get(RowId(*row_id), snapshot).is_some())
10694 .collect(),
10695 )
10696 }
10697 Condition::RangeF64 {
10698 column_id,
10699 lo,
10700 lo_inclusive,
10701 hi,
10702 hi_inclusive,
10703 } => {
10704 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10707 RowIdSet::from_unsorted(
10708 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10709 .into_iter()
10710 .collect(),
10711 )
10712 } else if self.run_refs.len() == 1 {
10713 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10714 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
10715 } else {
10716 return self.range_scan_f64(
10717 *column_id,
10718 *lo,
10719 *lo_inclusive,
10720 *hi,
10721 *hi_inclusive,
10722 snapshot,
10723 );
10724 };
10725 set.remove_many(self.overlay_rid_set(snapshot));
10726 self.range_scan_overlay_f64(
10727 &mut set,
10728 *column_id,
10729 *lo,
10730 *lo_inclusive,
10731 *hi,
10732 *hi_inclusive,
10733 snapshot,
10734 );
10735 set
10736 }
10737 Condition::IsNull { column_id } => {
10738 let mut set = if self.run_refs.len() == 1 {
10739 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10740 r.null_row_id_set(*column_id, true)?
10741 } else {
10742 return self.null_scan(*column_id, true, snapshot);
10743 };
10744 set.remove_many(self.overlay_rid_set(snapshot));
10745 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
10746 set
10747 }
10748 Condition::IsNotNull { column_id } => {
10749 let mut set = if self.run_refs.len() == 1 {
10750 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10751 r.null_row_id_set(*column_id, false)?
10752 } else {
10753 return self.null_scan(*column_id, false, snapshot);
10754 };
10755 set.remove_many(self.overlay_rid_set(snapshot));
10756 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
10757 set
10758 }
10759 })
10760 }
10761
10762 fn range_scan_i64(
10770 &self,
10771 column_id: u16,
10772 lo: i64,
10773 hi: i64,
10774 snapshot: Snapshot,
10775 ) -> Result<RowIdSet> {
10776 let mut row_ids = Vec::new();
10777 let mut tomb_rids: HashSet<u64> = HashSet::new();
10784 let overlay_rids = self.overlay_rid_set(snapshot);
10785 for rr in &self.run_refs {
10786 let mut reader = self.open_reader(rr.run_id)?;
10787 let matched = reader.range_row_ids_visible_i64_at(column_id, lo, hi, snapshot)?;
10788 for rid in matched {
10789 if !overlay_rids.contains(&rid) {
10790 row_ids.push(rid);
10791 }
10792 }
10793 for rid in reader.tombstoned_row_ids_at(snapshot)? {
10794 tomb_rids.insert(rid);
10795 }
10796 }
10797 let overlay_tombstone_snapshot = if self.is_shared()
10804 || (self.pending_put_cols.is_empty() && self.pending_delete_rids.is_empty())
10805 || snapshot.epoch.0 != self.epoch.visible().0
10806 {
10807 snapshot
10808 } else {
10809 Snapshot::at(self.pending_epoch())
10810 };
10811 for row in self
10812 .memtable
10813 .visible_versions_at(overlay_tombstone_snapshot)
10814 .into_iter()
10815 .chain(
10816 self.mutable_run
10817 .visible_versions_at(overlay_tombstone_snapshot),
10818 )
10819 {
10820 if row.deleted {
10821 tomb_rids.insert(row.row_id.0);
10822 }
10823 }
10824 let mut s = RowIdSet::from_unsorted(row_ids);
10825 s.remove_many(tomb_rids);
10826 let candidates: Vec<RowId> = s.into_sorted_vec().into_iter().map(RowId).collect();
10827 let eligible = self.eligible_candidate_ids(&candidates, column_id, snapshot, None)?;
10828 let mut s = RowIdSet::from_unsorted(eligible.into_iter().map(|row_id| row_id.0).collect());
10829 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
10830 Ok(s)
10831 }
10832
10833 fn range_scan_f64(
10836 &self,
10837 column_id: u16,
10838 lo: f64,
10839 lo_inclusive: bool,
10840 hi: f64,
10841 hi_inclusive: bool,
10842 snapshot: Snapshot,
10843 ) -> Result<RowIdSet> {
10844 let mut row_ids = Vec::new();
10845 let overlay_rids = self.overlay_rid_set(snapshot);
10846 for rr in &self.run_refs {
10847 let mut reader = self.open_reader(rr.run_id)?;
10848 let matched = reader.range_row_ids_visible_f64_at(
10849 column_id,
10850 lo,
10851 lo_inclusive,
10852 hi,
10853 hi_inclusive,
10854 snapshot,
10855 )?;
10856 for rid in matched {
10857 if !overlay_rids.contains(&rid) {
10858 row_ids.push(rid);
10859 }
10860 }
10861 }
10862 let mut s = RowIdSet::from_unsorted(row_ids);
10863 self.range_scan_overlay_f64(
10864 &mut s,
10865 column_id,
10866 lo,
10867 lo_inclusive,
10868 hi,
10869 hi_inclusive,
10870 snapshot,
10871 );
10872 Ok(s)
10873 }
10874
10875 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
10877 let mut s = HashSet::new();
10878 for row in self.memtable.visible_versions_at(snapshot) {
10879 s.insert(row.row_id.0);
10880 }
10881 for row in self.mutable_run.visible_versions_at(snapshot) {
10882 s.insert(row.row_id.0);
10883 }
10884 s
10885 }
10886
10887 fn range_scan_overlay_i64(
10888 &self,
10889 s: &mut RowIdSet,
10890 column_id: u16,
10891 lo: i64,
10892 hi: i64,
10893 snapshot: Snapshot,
10894 ) {
10895 let mut newest: HashMap<u64, Row> = HashMap::new();
10902 for r in self.mutable_run.visible_versions_at(snapshot) {
10903 newest.insert(r.row_id.0, r);
10904 }
10905 for r in self.memtable.visible_versions_at(snapshot) {
10906 newest
10907 .entry(r.row_id.0)
10908 .and_modify(|cur| {
10909 if Snapshot::version_is_newer(
10910 r.committed_epoch,
10911 r.commit_ts,
10912 cur.committed_epoch,
10913 cur.commit_ts,
10914 ) {
10915 *cur = r.clone();
10916 }
10917 })
10918 .or_insert(r);
10919 }
10920 for row in newest.values() {
10921 if !row.deleted {
10922 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
10923 if *v >= lo && *v <= hi && !self.row_expired_at(row, unix_nanos_now()) {
10924 s.insert(row.row_id.0);
10925 }
10926 }
10927 }
10928 }
10929 }
10930
10931 #[allow(clippy::too_many_arguments)]
10932 fn range_scan_overlay_f64(
10933 &self,
10934 s: &mut RowIdSet,
10935 column_id: u16,
10936 lo: f64,
10937 lo_inclusive: bool,
10938 hi: f64,
10939 hi_inclusive: bool,
10940 snapshot: Snapshot,
10941 ) {
10942 let mut newest: HashMap<u64, Row> = HashMap::new();
10945 for r in self.mutable_run.visible_versions_at(snapshot) {
10946 newest.insert(r.row_id.0, r);
10947 }
10948 for r in self.memtable.visible_versions_at(snapshot) {
10949 newest
10950 .entry(r.row_id.0)
10951 .and_modify(|cur| {
10952 if Snapshot::version_is_newer(
10953 r.committed_epoch,
10954 r.commit_ts,
10955 cur.committed_epoch,
10956 cur.commit_ts,
10957 ) {
10958 *cur = r.clone();
10959 }
10960 })
10961 .or_insert(r);
10962 }
10963 for row in newest.values() {
10964 if !row.deleted {
10965 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
10966 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
10967 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
10968 if ok_lo && ok_hi && !self.row_expired_at(row, unix_nanos_now()) {
10969 s.insert(row.row_id.0);
10970 }
10971 }
10972 }
10973 }
10974 }
10975
10976 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
10979 let mut row_ids = Vec::new();
10980 let overlay_rids = self.overlay_rid_set(snapshot);
10981 for rr in &self.run_refs {
10982 let mut reader = self.open_reader(rr.run_id)?;
10983 let matched = reader.null_row_ids_visible_at(column_id, want_nulls, snapshot)?;
10984 for rid in matched {
10985 if !overlay_rids.contains(&rid) {
10986 row_ids.push(rid);
10987 }
10988 }
10989 }
10990 let mut s = RowIdSet::from_unsorted(row_ids);
10991 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
10992 Ok(s)
10993 }
10994
10995 fn null_scan_overlay(
10999 &self,
11000 s: &mut RowIdSet,
11001 column_id: u16,
11002 want_nulls: bool,
11003 snapshot: Snapshot,
11004 ) {
11005 let mut newest: HashMap<u64, Row> = HashMap::new();
11006 for r in self.mutable_run.visible_versions_at(snapshot) {
11007 newest.insert(r.row_id.0, r);
11008 }
11009 for r in self.memtable.visible_versions_at(snapshot) {
11010 newest
11011 .entry(r.row_id.0)
11012 .and_modify(|cur| {
11013 if Snapshot::version_is_newer(
11014 r.committed_epoch,
11015 r.commit_ts,
11016 cur.committed_epoch,
11017 cur.commit_ts,
11018 ) {
11019 *cur = r.clone();
11020 }
11021 })
11022 .or_insert(r);
11023 }
11024 for row in newest.values() {
11025 if row.deleted {
11026 continue;
11027 }
11028 let is_null = !row.columns.contains_key(&column_id)
11029 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
11030 if is_null == want_nulls {
11031 s.insert(row.row_id.0);
11032 }
11033 }
11034 }
11035
11036 pub fn snapshot(&self) -> Snapshot {
11037 let epoch = self.epoch.visible();
11038 match &self.wal {
11042 WalSink::Shared(shared) => match shared.hlc.now() {
11043 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
11044 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
11046 },
11047 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
11048 }
11049 }
11050
11051 pub fn data_generation(&self) -> u64 {
11053 self.data_generation
11054 }
11055
11056 pub(crate) fn bump_data_generation(&mut self) {
11057 self.data_generation = self.data_generation.wrapping_add(1);
11058 }
11059
11060 pub fn table_id(&self) -> u64 {
11062 self.table_id
11063 }
11064
11065 fn seal_generations(&mut self) {
11071 self.memtable.seal();
11072 self.mutable_run.seal();
11073 self.hot.seal();
11074 for index in self.bitmap.values_mut() {
11075 index.seal();
11076 }
11077 for index in self.ann.values_mut() {
11078 index.seal();
11079 }
11080 for index in self.fm.values_mut() {
11081 index.seal();
11082 }
11083 for index in self.sparse.values_mut() {
11084 index.seal();
11085 }
11086 for index in self.minhash.values_mut() {
11087 index.seal();
11088 }
11089 self.pk_by_row.seal();
11090 }
11091
11092 fn capture_read_generation(&self) -> ReadGeneration {
11097 let visible_through = self.current_epoch();
11098 ReadGeneration {
11099 schema: Arc::new(self.schema.clone()),
11100 base_runs: Arc::new(self.run_refs.clone()),
11101 deltas: TableDeltas {
11102 memtable: self.memtable.clone(),
11103 mutable_run: self.mutable_run.clone(),
11104 hot: self.hot.clone(),
11105 pk_by_row: self.pk_by_row.clone(),
11106 },
11107 indexes: Arc::new(IndexGeneration::capture(
11108 &self.bitmap,
11109 &self.learned_range,
11110 &self.fm,
11111 &self.ann,
11112 &self.sparse,
11113 &self.minhash,
11114 visible_through,
11115 match &self.wal {
11117 WalSink::Shared(shared) => shared
11118 .hlc
11119 .now()
11120 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
11121 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
11122 },
11123 )),
11124 visible_through,
11125 }
11126 }
11127
11128 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
11133 self.ensure_indexes_complete()?;
11134 self.seal_generations();
11135 let view = Arc::new(self.capture_read_generation());
11136 self.published.store(Arc::clone(&view));
11137 Ok(view)
11138 }
11139
11140 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
11146 self.published.load_full()
11147 }
11148
11149 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
11151 &self.pins
11152 }
11153
11154 pub fn version_gc_floor(&self) -> Epoch {
11159 self.min_active_snapshot()
11160 .unwrap_or_else(|| self.current_epoch())
11161 }
11162
11163 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
11170 let mut report = self.pins.report();
11171 let transaction_floor = [
11172 self.pinned.keys().next().copied(),
11173 self.snapshots.min_pinned(),
11174 ]
11175 .into_iter()
11176 .flatten()
11177 .min();
11178 if let Some(epoch) = transaction_floor {
11179 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
11180 }
11181 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
11182 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
11183 }
11184 report
11185 }
11186
11187 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
11188 self.publish_read_generation()?;
11189 let mut generation = self.clone();
11190 generation.read_only = true;
11191 generation.wal = WalSink::ReadOnly;
11192 generation.pending_delete_rids.clear();
11193 generation.pending_put_cols.clear();
11194 generation.pending_rows.clear();
11195 generation.pending_rows_auto_inc.clear();
11196 generation.pending_dels.clear();
11197 generation.pending_truncate = None;
11198 generation.agg_cache = Arc::new(HashMap::new());
11199 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
11202 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
11205 crate::retention::PinSource::ReadGeneration,
11206 self.current_epoch(),
11207 )));
11208 Ok(generation)
11209 }
11210
11211 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
11212 (std::mem::size_of::<Self>() as u64)
11213 .saturating_add(self.memtable.approx_bytes())
11214 .saturating_add(self.mutable_run.approx_bytes())
11215 .saturating_add(self.live_count.saturating_mul(64))
11216 }
11217
11218 pub fn pin_snapshot(&mut self) -> Snapshot {
11225 let snap = self.snapshot();
11226 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
11227 snap
11228 }
11229
11230 fn active_pin_epochs_for_rebuild(&self) -> Vec<Epoch> {
11242 let mut set = BTreeSet::new();
11243 for epoch in self.pinned.keys().copied() {
11244 set.insert(epoch);
11245 }
11246 for epoch in self.snapshots.live_pinned_epochs() {
11247 set.insert(epoch);
11248 }
11249 for epoch in self.pins.live_pin_epochs() {
11250 set.insert(epoch);
11251 }
11252 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
11253 set.insert(floor);
11254 }
11255 set.into_iter().collect()
11256 }
11257
11258 pub fn hlc_gc_floor(
11268 &self,
11269 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
11270 ) -> crate::epoch::GcFloor {
11271 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
11272 epoch
11273 .and_then(&mut project_epoch)
11274 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
11275 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
11276 };
11277 let transaction = [
11278 self.pinned.keys().next().copied(),
11279 self.snapshots.min_pinned(),
11280 ]
11281 .into_iter()
11282 .flatten()
11283 .min();
11284 let history = self.snapshots.history_floor(self.current_epoch());
11285 crate::epoch::GcFloor {
11286 transaction_snapshot: project(transaction),
11287 history_retention: project(history),
11288 backup_pitr: project(
11289 self.pins
11290 .oldest_for(crate::retention::PinSource::BackupPitr),
11291 ),
11292 replication: project(
11293 self.pins
11294 .oldest_for(crate::retention::PinSource::Replication),
11295 ),
11296 read_generation: project(
11297 self.pins
11298 .oldest_for(crate::retention::PinSource::ReadGeneration),
11299 ),
11300 online_index_build: project(
11301 self.pins
11302 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
11303 ),
11304 }
11305 }
11306
11307 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
11309 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
11310 *count -= 1;
11311 if *count == 0 {
11312 self.pinned.remove(&snap.epoch);
11313 }
11314 }
11315 }
11316
11317 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
11329 let local = self.pinned.keys().next().copied();
11330 let global = self.snapshots.min_pinned();
11331 let history = self.snapshots.history_floor(self.current_epoch());
11332 let pinned = self.pins.oldest_pinned();
11333 [local, global, history, pinned].into_iter().flatten().min()
11334 }
11335
11336 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
11340 self.ensure_writable()?;
11341 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
11342 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
11343 }
11344
11345 pub fn clear_ttl(&mut self) -> Result<()> {
11346 self.ensure_writable()?;
11347 self.apply_ttl_policy_at(None, self.current_epoch())
11348 }
11349
11350 pub fn ttl(&self) -> Option<TtlPolicy> {
11351 self.ttl
11352 }
11353
11354 pub(crate) fn prepare_ttl_policy(
11355 &self,
11356 column_name: &str,
11357 duration_nanos: u64,
11358 ) -> Result<TtlPolicy> {
11359 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
11360 return Err(MongrelError::InvalidArgument(
11361 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
11362 ));
11363 }
11364 let column = self
11365 .schema
11366 .columns
11367 .iter()
11368 .find(|column| column.name == column_name)
11369 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
11370 if column.ty != TypeId::TimestampNanos {
11371 return Err(MongrelError::Schema(format!(
11372 "TTL column {column_name} must be TimestampNanos, is {:?}",
11373 column.ty
11374 )));
11375 }
11376 Ok(TtlPolicy {
11377 column_id: column.id,
11378 duration_nanos,
11379 })
11380 }
11381
11382 pub(crate) fn apply_ttl_policy_at(
11383 &mut self,
11384 policy: Option<TtlPolicy>,
11385 epoch: Epoch,
11386 ) -> Result<()> {
11387 if let Some(policy) = policy {
11388 let column = self
11389 .schema
11390 .columns
11391 .iter()
11392 .find(|column| column.id == policy.column_id)
11393 .ok_or_else(|| {
11394 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
11395 })?;
11396 if column.ty != TypeId::TimestampNanos
11397 || policy.duration_nanos == 0
11398 || policy.duration_nanos > i64::MAX as u64
11399 {
11400 return Err(MongrelError::Schema("invalid TTL policy".into()));
11401 }
11402 }
11403 self.ttl = policy;
11404 self.agg_cache = Arc::new(HashMap::new());
11405 self.clear_result_cache();
11406 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
11407 self.persist_manifest(epoch)
11408 }
11409
11410 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
11411 let Some(policy) = self.ttl else {
11412 return false;
11413 };
11414 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
11415 return false;
11416 };
11417 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
11418 }
11419
11420 pub fn current_epoch(&self) -> Epoch {
11421 self.epoch.visible()
11422 }
11423
11424 pub fn memtable_len(&self) -> usize {
11425 self.memtable.len()
11426 }
11427
11428 pub fn count(&self) -> u64 {
11431 if self.ttl.is_none()
11432 && self.pending_put_cols.is_empty()
11433 && self.pending_delete_rids.is_empty()
11434 && self.pending_rows.is_empty()
11435 && self.pending_dels.is_empty()
11436 && self.pending_truncate.is_none()
11437 {
11438 self.live_count
11439 } else {
11440 self.visible_rows(self.snapshot())
11441 .map(|rows| rows.len() as u64)
11442 .unwrap_or(self.live_count)
11443 }
11444 }
11445
11446 pub fn count_conditions(
11450 &mut self,
11451 conditions: &[crate::query::Condition],
11452 snapshot: Snapshot,
11453 ) -> Result<Option<u64>> {
11454 use crate::query::Condition;
11455 if self.ttl.is_some() {
11456 if conditions.is_empty() {
11457 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
11458 }
11459 let mut sets = Vec::with_capacity(conditions.len());
11460 for condition in conditions {
11461 sets.push(self.resolve_condition(condition, snapshot)?);
11462 }
11463 let survivors = RowIdSet::intersect_many(sets);
11464 let rows = self.visible_rows(snapshot)?;
11465 return Ok(Some(
11466 rows.into_iter()
11467 .filter(|row| survivors.contains(row.row_id.0))
11468 .count() as u64,
11469 ));
11470 }
11471 if conditions.is_empty() {
11472 return Ok(Some(self.count()));
11473 }
11474 let served = |c: &Condition| {
11475 matches!(
11476 c,
11477 Condition::Pk(_)
11478 | Condition::BitmapEq { .. }
11479 | Condition::BitmapIn { .. }
11480 | Condition::BytesPrefix { .. }
11481 | Condition::FmContains { .. }
11482 | Condition::FmContainsAll { .. }
11483 | Condition::Ann { .. }
11484 | Condition::Range { .. }
11485 | Condition::RangeF64 { .. }
11486 | Condition::SparseMatch { .. }
11487 | Condition::MinHashSimilar { .. }
11488 | Condition::IsNull { .. }
11489 | Condition::IsNotNull { .. }
11490 )
11491 };
11492 if !conditions.iter().all(served) {
11493 return Ok(None);
11494 }
11495 self.ensure_indexes_complete()?;
11496 if !self.pending_put_cols.is_empty()
11497 || !self.pending_delete_rids.is_empty()
11498 || !self.pending_rows.is_empty()
11499 || !self.pending_dels.is_empty()
11500 || self.pending_truncate.is_some()
11501 {
11502 let mut sets = Vec::with_capacity(conditions.len());
11503 for condition in conditions {
11504 sets.push(self.resolve_condition(condition, snapshot)?);
11505 }
11506 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
11507 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
11508 }
11509 let mut sets = Vec::with_capacity(conditions.len());
11510 for condition in conditions {
11511 sets.push(self.resolve_condition(condition, snapshot)?);
11512 }
11513 let rids = RowIdSet::intersect_many(sets);
11514 if self.had_deletes || !self.memtable.is_empty() || !self.mutable_run.is_empty() {
11522 let sorted = rids.into_sorted_vec();
11523 let count = self.rows_for_rids(&sorted, snapshot)?.len() as u64;
11524 crate::trace::QueryTrace::record(|t| {
11525 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
11526 t.survivor_count = Some(count as usize);
11527 t.conditions_pushed = conditions.len();
11528 });
11529 return Ok(Some(count));
11530 }
11531 let count = rids.len() as u64;
11532 crate::trace::QueryTrace::record(|t| {
11533 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
11534 t.survivor_count = Some(count as usize);
11535 t.conditions_pushed = conditions.len();
11536 });
11537 Ok(Some(count))
11538 }
11539
11540 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
11546 let mut out = Vec::new();
11547 for row in self.memtable.visible_versions_at(snapshot) {
11548 if row.deleted {
11549 out.push(row.row_id.0);
11550 }
11551 }
11552 for row in self.mutable_run.visible_versions_at(snapshot) {
11553 if row.deleted {
11554 out.push(row.row_id.0);
11555 }
11556 }
11557 out
11558 }
11559
11560 pub fn bulk_load_columns(
11569 &mut self,
11570 user_columns: Vec<(u16, columnar::NativeColumn)>,
11571 ) -> Result<Epoch> {
11572 self.bulk_load_columns_with(user_columns, 3, false, true)
11573 }
11574
11575 pub fn bulk_load_fast(
11582 &mut self,
11583 user_columns: Vec<(u16, columnar::NativeColumn)>,
11584 ) -> Result<Epoch> {
11585 self.bulk_load_columns_with(user_columns, -1, true, false)
11586 }
11587
11588 fn bulk_load_columns_with(
11589 &mut self,
11590 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
11591 zstd_level: i32,
11592 force_plain: bool,
11593 lz4: bool,
11594 ) -> Result<Epoch> {
11595 self.ensure_writable()?;
11596 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
11597 if n == 0 {
11598 return Ok(self.current_epoch());
11599 }
11600 let epoch = self.commit_new_epoch()?;
11601 let live_before = self.live_count;
11602 self.spill_mutable_run(epoch)?;
11604 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
11605 && self.indexes_complete
11606 && self.run_refs.is_empty()
11607 && self.memtable.is_empty()
11608 && self.mutable_run.is_empty();
11609 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
11612 self.validate_columns_not_null(&user_columns, n)?;
11613 let winner_idx = self
11614 .bulk_pk_winner_indices(&user_columns, n)
11615 .filter(|idx| idx.len() != n);
11616 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
11617 match winner_idx.as_deref() {
11618 Some(idx) => {
11619 let compacted = user_columns
11620 .iter()
11621 .map(|(id, c)| (*id, c.gather(idx)))
11622 .collect();
11623 (compacted, idx.len())
11624 }
11625 None => (user_columns, n),
11626 };
11627 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
11628 let first = self.allocator.alloc_range(write_n as u64)?.0;
11629 for rid in first..first + write_n as u64 {
11630 self.reservoir.offer(rid);
11631 }
11632 let run_id = self.alloc_run_id()?;
11633 let path = self.run_path(run_id);
11634 let mut writer =
11635 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
11636 if force_plain {
11637 writer = writer.with_plain();
11638 } else if lz4 {
11639 writer = writer.with_lz4();
11642 } else {
11643 writer = writer.with_zstd_level(zstd_level);
11644 }
11645 if let Some(kek) = &self.kek {
11646 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
11647 }
11648 let header = match self.create_run_file(run_id)? {
11649 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
11650 None => writer.write_native(&path, &write_columns, write_n, first)?,
11651 };
11652 self.run_refs.push(RunRef {
11653 run_id: run_id as u128,
11654 level: 0,
11655 epoch_created: epoch.0,
11656 row_count: header.row_count,
11657 });
11658 self.run_row_id_ranges
11659 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
11660 self.live_count = self.live_count.saturating_add(write_n as u64);
11661 if eager_index_build {
11662 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
11663 self.index_columns_bulk(&write_columns, &row_ids);
11664 self.indexes_complete = true;
11665 self.build_learned_ranges()?;
11666 } else {
11667 self.indexes_complete = false;
11671 }
11672 self.mark_flushed(epoch)?;
11673 self.persist_manifest(epoch)?;
11674 if eager_index_build {
11675 self.checkpoint_indexes(epoch);
11676 }
11677 self.clear_result_cache();
11678 self.data_generation = self.data_generation.wrapping_add(1);
11679 Ok(epoch)
11680 }
11681
11682 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
11700 let n = row_ids.len();
11701 if n == 0 {
11702 return;
11703 }
11704 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
11705 columns.iter().map(|(id, c)| (*id, c)).collect();
11706 let ty_of: std::collections::HashMap<u16, TypeId> = self
11707 .schema
11708 .columns
11709 .iter()
11710 .map(|c| (c.id, c.ty.clone()))
11711 .collect();
11712 let pk_id = self.schema.primary_key().map(|c| c.id);
11713
11714 for (i, &rid) in row_ids.iter().enumerate() {
11715 let row_id = RowId(rid);
11716 if let Some(pid) = pk_id {
11717 if let Some(col) = by_id.get(&pid) {
11718 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
11719 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
11720 self.insert_hot_pk(key, row_id);
11721 }
11722 }
11723 }
11724 for idef in &self.schema.indexes {
11725 let Some(col) = by_id.get(&idef.column_id) else {
11726 continue;
11727 };
11728 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
11729 match idef.kind {
11730 IndexKind::Bitmap => {
11731 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
11732 if let Some(key) =
11733 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
11734 {
11735 b.insert(key, row_id);
11736 }
11737 }
11738 }
11739 IndexKind::FmIndex => {
11740 if let Some(f) = self.fm.get_mut(&idef.column_id) {
11741 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11742 f.insert(bytes.to_vec(), row_id);
11743 }
11744 }
11745 }
11746 IndexKind::Sparse => {
11747 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
11748 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11749 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
11750 s.insert(&terms, row_id);
11751 }
11752 }
11753 }
11754 }
11755 IndexKind::MinHash => {
11756 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
11757 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11758 let tokens = crate::index::token_hashes_from_bytes(bytes);
11759 mh.insert(&tokens, row_id);
11760 }
11761 }
11762 }
11763 _ => {}
11764 }
11765 }
11766 }
11767 }
11768
11769 pub fn visible_columns_native(
11774 &self,
11775 snapshot: Snapshot,
11776 projection: Option<&[u16]>,
11777 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11778 self.visible_columns_native_inner(snapshot, projection, None)
11779 }
11780
11781 pub fn visible_columns_native_with_control(
11782 &self,
11783 snapshot: Snapshot,
11784 projection: Option<&[u16]>,
11785 control: &crate::ExecutionControl,
11786 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11787 self.visible_columns_native_inner(snapshot, projection, Some(control))
11788 }
11789
11790 fn visible_columns_native_inner(
11791 &self,
11792 snapshot: Snapshot,
11793 projection: Option<&[u16]>,
11794 control: Option<&crate::ExecutionControl>,
11795 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11796 execution_checkpoint(control, 0)?;
11797 let wanted: Vec<u16> = match projection {
11798 Some(p) => p.to_vec(),
11799 None => self.schema.columns.iter().map(|c| c.id).collect(),
11800 };
11801 if self.ttl.is_none()
11802 && self.memtable.is_empty()
11803 && self.mutable_run.is_empty()
11804 && self.run_refs.len() == 1
11805 {
11806 let rr = self.run_refs[0].clone();
11807 let mut reader = self.open_reader(rr.run_id)?;
11808 let idxs = reader.visible_indices_native_at(snapshot)?;
11809 execution_checkpoint(control, 0)?;
11810 let all_visible = idxs.len() == reader.row_count();
11811 if reader.has_mmap() && control.is_none() {
11817 use rayon::prelude::*;
11818 let valid: Vec<u16> = wanted
11821 .iter()
11822 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
11823 .copied()
11824 .collect();
11825 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
11827 .par_iter()
11828 .filter_map(|cid| {
11829 reader
11830 .column_native_shared(*cid)
11831 .ok()
11832 .map(|col| (*cid, col))
11833 })
11834 .collect();
11835 let cols = decoded
11836 .into_iter()
11837 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
11838 .collect();
11839 return Ok(cols);
11840 }
11841 let mut cols = Vec::with_capacity(wanted.len());
11842 for (index, cid) in wanted.iter().enumerate() {
11843 execution_checkpoint(control, index)?;
11844 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
11845 Some(c) => c,
11846 None => continue,
11847 };
11848 let col = reader.column_native(cdef.id)?;
11849 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
11850 }
11851 return Ok(cols);
11852 }
11853 let vcols = self.visible_columns(snapshot)?;
11854 execution_checkpoint(control, 0)?;
11855 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
11856 let out: Vec<(u16, columnar::NativeColumn)> = vcols
11857 .into_iter()
11858 .filter(|(id, _)| want_set.contains(id))
11859 .map(|(id, vals)| {
11860 let ty = self
11861 .schema
11862 .columns
11863 .iter()
11864 .find(|c| c.id == id)
11865 .map(|c| c.ty.clone())
11866 .unwrap_or(TypeId::Bytes);
11867 (id, columnar::values_to_native(ty, &vals))
11868 })
11869 .collect();
11870 Ok(out)
11871 }
11872
11873 pub fn run_count(&self) -> usize {
11874 self.run_refs.len()
11875 }
11876
11877 pub fn memtable_is_empty(&self) -> bool {
11879 self.memtable.is_empty()
11880 }
11881
11882 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
11886 self.page_cache.stats()
11887 }
11888
11889 pub fn reset_page_cache_stats(&self) {
11891 self.page_cache.reset_stats();
11892 }
11893
11894 pub fn run_ids(&self) -> Vec<u128> {
11897 self.run_refs.iter().map(|r| r.run_id).collect()
11898 }
11899
11900 pub fn single_run_is_clean(&self) -> bool {
11904 if self.ttl.is_some() || self.run_refs.len() != 1 {
11905 return false;
11906 }
11907 self.open_reader(self.run_refs[0].run_id)
11908 .map(|r| r.is_clean())
11909 .unwrap_or(false)
11910 }
11911
11912 fn resolve_footprint(
11919 &self,
11920 conditions: &[crate::query::Condition],
11921 snapshot: Snapshot,
11922 ) -> roaring::RoaringBitmap {
11923 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
11924 return roaring::RoaringBitmap::new();
11925 }
11926 if self.run_refs.is_empty() {
11927 return roaring::RoaringBitmap::new();
11928 }
11929 if self.run_refs.len() == 1 {
11931 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
11932 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
11933 return rids.to_roaring_lossy();
11934 }
11935 }
11936 }
11937 roaring::RoaringBitmap::new()
11938 }
11939
11940 pub fn query_columns_native_cached(
11951 &mut self,
11952 conditions: &[crate::query::Condition],
11953 projection: Option<&[u16]>,
11954 snapshot: Snapshot,
11955 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11956 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
11957 }
11958
11959 pub fn query_columns_native_cached_with_control(
11960 &mut self,
11961 conditions: &[crate::query::Condition],
11962 projection: Option<&[u16]>,
11963 snapshot: Snapshot,
11964 control: &crate::ExecutionControl,
11965 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11966 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
11967 }
11968
11969 fn query_columns_native_cached_inner(
11970 &mut self,
11971 conditions: &[crate::query::Condition],
11972 projection: Option<&[u16]>,
11973 snapshot: Snapshot,
11974 control: Option<&crate::ExecutionControl>,
11975 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11976 execution_checkpoint(control, 0)?;
11977 if self.ttl.is_some() {
11980 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11981 }
11982 if conditions.is_empty() {
11983 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11984 }
11985 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
11989 let identity = PersistentCacheIdentity {
11990 table_id: self.table_id(),
11991 schema_id: self.schema.schema_id,
11992 logical_generation: self.current_epoch().0,
11993 };
11994 if let Some(hit) = self.result_cache.lock().get_columns(key, identity) {
11995 crate::trace::QueryTrace::record(|t| {
11996 t.result_cache_hit = true;
11997 t.scan_mode = crate::trace::ScanMode::NativePushdown;
11998 });
11999 return Ok(Some((*hit).clone()));
12000 }
12001 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
12002 execution_checkpoint(control, 0)?;
12003 if let Some(cols) = &res {
12004 let footprint = self.resolve_footprint(conditions, snapshot);
12005 let condition_cols = crate::query::condition_columns(conditions);
12006 execution_checkpoint(control, 0)?;
12007 let entry = CachedEntry {
12008 data: CachedData::Columns(Arc::new(cols.clone())),
12009 footprint,
12010 condition_cols,
12011 };
12012 let mut cache = self.result_cache.lock();
12018 let entry_generation = cache.allocate_persist_generation(key);
12019 let context = PersistContext {
12020 identity,
12021 key,
12022 entry_generation,
12023 };
12024 cache.persist_entry(context, &entry);
12025 cache.insert(key, entry);
12026 }
12027 Ok(res)
12028 }
12029
12030 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
12035 if self.ttl.is_some() {
12036 return self.query(q);
12037 }
12038 if q.conditions.is_empty() {
12039 return self.query(q);
12040 }
12041 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
12042 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
12043 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
12044 let identity = PersistentCacheIdentity {
12045 table_id: self.table_id(),
12046 schema_id: self.schema.schema_id,
12047 logical_generation: self.current_epoch().0,
12048 };
12049 if let Some(hit) = self.result_cache.lock().get_rows(key, identity) {
12050 crate::trace::QueryTrace::record(|t| {
12051 t.result_cache_hit = true;
12052 t.scan_mode = crate::trace::ScanMode::Materialized;
12053 });
12054 return Ok((*hit).clone());
12055 }
12056 let rows = self.query(q)?;
12057 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
12058 let condition_cols = crate::query::condition_columns(&q.conditions);
12059 let entry = CachedEntry {
12060 data: CachedData::Rows(Arc::new(rows.clone())),
12061 footprint,
12062 condition_cols,
12063 };
12064 let mut cache = self.result_cache.lock();
12069 let entry_generation = cache.allocate_persist_generation(key);
12070 let context = PersistContext {
12071 identity,
12072 key,
12073 entry_generation,
12074 };
12075 cache.persist_entry(context, &entry);
12076 cache.insert(key, entry);
12077 Ok(rows)
12078 }
12079
12080 #[allow(clippy::type_complexity)]
12095 pub fn query_columns_native_traced(
12096 &mut self,
12097 conditions: &[crate::query::Condition],
12098 projection: Option<&[u16]>,
12099 snapshot: Snapshot,
12100 ) -> Result<(
12101 Option<Vec<(u16, columnar::NativeColumn)>>,
12102 crate::trace::QueryTrace,
12103 )> {
12104 let (result, trace) = crate::trace::QueryTrace::capture(|| {
12105 self.query_columns_native(conditions, projection, snapshot)
12106 });
12107 Ok((result?, trace))
12108 }
12109
12110 #[allow(clippy::type_complexity)]
12113 pub fn query_columns_native_cached_traced(
12114 &mut self,
12115 conditions: &[crate::query::Condition],
12116 projection: Option<&[u16]>,
12117 snapshot: Snapshot,
12118 ) -> Result<(
12119 Option<Vec<(u16, columnar::NativeColumn)>>,
12120 crate::trace::QueryTrace,
12121 )> {
12122 let (result, trace) = crate::trace::QueryTrace::capture(|| {
12123 self.query_columns_native_cached(conditions, projection, snapshot)
12124 });
12125 Ok((result?, trace))
12126 }
12127
12128 pub fn native_page_cursor_traced(
12130 &self,
12131 snapshot: Snapshot,
12132 projection: Vec<(u16, TypeId)>,
12133 conditions: &[crate::query::Condition],
12134 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
12135 let (result, trace) = crate::trace::QueryTrace::capture(|| {
12136 self.native_page_cursor(snapshot, projection, conditions)
12137 });
12138 Ok((result?, trace))
12139 }
12140
12141 pub fn native_multi_run_cursor_traced(
12143 &self,
12144 snapshot: Snapshot,
12145 projection: Vec<(u16, TypeId)>,
12146 conditions: &[crate::query::Condition],
12147 ) -> Result<(
12148 Option<crate::cursor::MultiRunCursor>,
12149 crate::trace::QueryTrace,
12150 )> {
12151 let (result, trace) = crate::trace::QueryTrace::capture(|| {
12152 self.native_multi_run_cursor(snapshot, projection, conditions)
12153 });
12154 Ok((result?, trace))
12155 }
12156
12157 pub fn count_conditions_traced(
12159 &mut self,
12160 conditions: &[crate::query::Condition],
12161 snapshot: Snapshot,
12162 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
12163 let (result, trace) =
12164 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
12165 Ok((result?, trace))
12166 }
12167
12168 pub fn query_traced(
12170 &mut self,
12171 q: &crate::query::Query,
12172 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
12173 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
12174 Ok((result?, trace))
12175 }
12176
12177 pub fn query_columns_native(
12182 &mut self,
12183 conditions: &[crate::query::Condition],
12184 projection: Option<&[u16]>,
12185 snapshot: Snapshot,
12186 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
12187 self.query_columns_native_inner(conditions, projection, snapshot, None)
12188 }
12189
12190 pub fn query_columns_native_with_control(
12191 &mut self,
12192 conditions: &[crate::query::Condition],
12193 projection: Option<&[u16]>,
12194 snapshot: Snapshot,
12195 control: &crate::ExecutionControl,
12196 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
12197 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
12198 }
12199
12200 fn query_columns_native_inner(
12201 &mut self,
12202 conditions: &[crate::query::Condition],
12203 projection: Option<&[u16]>,
12204 snapshot: Snapshot,
12205 control: Option<&crate::ExecutionControl>,
12206 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
12207 use crate::query::Condition;
12208 execution_checkpoint(control, 0)?;
12209 if self.ttl.is_some() {
12212 return Ok(None);
12213 }
12214 if conditions.is_empty() {
12215 return Ok(None);
12216 }
12217 self.ensure_indexes_complete()?;
12218
12219 let served = |c: &Condition| {
12224 matches!(
12225 c,
12226 Condition::Pk(_)
12227 | Condition::BitmapEq { .. }
12228 | Condition::BitmapIn { .. }
12229 | Condition::BytesPrefix { .. }
12230 | Condition::FmContains { .. }
12231 | Condition::FmContainsAll { .. }
12232 | Condition::Ann { .. }
12233 | Condition::Range { .. }
12234 | Condition::RangeF64 { .. }
12235 | Condition::SparseMatch { .. }
12236 | Condition::MinHashSimilar { .. }
12237 | Condition::IsNull { .. }
12238 | Condition::IsNotNull { .. }
12239 )
12240 };
12241 if !conditions.iter().all(served) {
12242 return Ok(None);
12243 }
12244 let fast_path =
12245 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
12246 crate::trace::QueryTrace::record(|t| {
12247 t.run_count = self.run_refs.len();
12248 t.memtable_rows = self.memtable.len();
12249 t.mutable_run_rows = self.mutable_run.len();
12250 t.conditions_pushed = conditions.len();
12251 t.learned_range_used = conditions.iter().any(|c| match c {
12252 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
12253 self.learned_range.contains_key(column_id)
12254 }
12255 _ => false,
12256 });
12257 });
12258 let col_ids: Vec<u16> = projection
12260 .map(|p| p.to_vec())
12261 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
12262 let proj_pairs: Vec<(u16, TypeId)> = col_ids
12263 .iter()
12264 .map(|&cid| {
12265 let ty = self
12266 .schema
12267 .columns
12268 .iter()
12269 .find(|c| c.id == cid)
12270 .map(|c| c.ty.clone())
12271 .unwrap_or(TypeId::Bytes);
12272 (cid, ty)
12273 })
12274 .collect();
12275
12276 if fast_path {
12282 let needs_column = conditions.iter().any(|c| match c {
12285 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
12286 Condition::RangeF64 { column_id, .. } => {
12287 !self.learned_range.contains_key(column_id)
12288 }
12289 _ => false,
12290 });
12291 let mut reader_opt: Option<RunReader> = if needs_column {
12292 Some(self.open_reader(self.run_refs[0].run_id)?)
12293 } else {
12294 None
12295 };
12296 let mut sets: Vec<RowIdSet> = Vec::new();
12297 for (index, c) in conditions.iter().enumerate() {
12298 execution_checkpoint(control, index)?;
12299 let s = match c {
12300 Condition::Range { column_id, lo, hi }
12301 if !self.learned_range.contains_key(column_id) =>
12302 {
12303 if reader_opt.is_none() {
12304 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
12305 }
12306 reader_opt
12307 .as_mut()
12308 .expect("reader opened for range")
12309 .range_row_id_set_i64(*column_id, *lo, *hi)?
12310 }
12311 Condition::RangeF64 {
12312 column_id,
12313 lo,
12314 lo_inclusive,
12315 hi,
12316 hi_inclusive,
12317 } if !self.learned_range.contains_key(column_id) => {
12318 if reader_opt.is_none() {
12319 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
12320 }
12321 reader_opt
12322 .as_mut()
12323 .expect("reader opened for range")
12324 .range_row_id_set_f64(
12325 *column_id,
12326 *lo,
12327 *lo_inclusive,
12328 *hi,
12329 *hi_inclusive,
12330 )?
12331 }
12332 _ => self.resolve_condition(c, snapshot)?,
12333 };
12334 sets.push(s);
12335 }
12336 let candidates = RowIdSet::intersect_many(sets);
12337 crate::trace::QueryTrace::record(|t| {
12338 t.survivor_count = Some(candidates.len());
12339 });
12340 if candidates.is_empty() {
12341 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
12342 .iter()
12343 .map(|&id| {
12344 (
12345 id,
12346 columnar::null_native(
12347 proj_pairs
12348 .iter()
12349 .find(|(c, _)| c == &id)
12350 .map(|(_, t)| t.clone())
12351 .unwrap_or(TypeId::Bytes),
12352 0,
12353 ),
12354 )
12355 })
12356 .collect();
12357 return Ok(Some(cols));
12358 }
12359 let mut reader = match reader_opt.take() {
12360 Some(r) => r,
12361 None => self.open_reader(self.run_refs[0].run_id)?,
12362 };
12363 let candidate_ids = candidates.into_sorted_vec();
12364 let (positions, fast_rid) = if let Some(positions) =
12365 reader.positions_for_row_ids_fast(&candidate_ids)
12366 {
12367 (positions, true)
12368 } else {
12369 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
12370 match col {
12371 columnar::NativeColumn::Int64 { data, .. } => {
12372 let mut p = Vec::with_capacity(candidate_ids.len());
12373 for (index, rid) in candidate_ids.iter().enumerate() {
12374 execution_checkpoint(control, index)?;
12375 if let Ok(position) = data.binary_search(&(*rid as i64)) {
12376 p.push(position);
12377 }
12378 }
12379 p.sort_unstable();
12380 (p, false)
12381 }
12382 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
12383 }
12384 };
12385 crate::trace::QueryTrace::record(|t| {
12386 t.scan_mode = crate::trace::ScanMode::NativePushdown;
12387 t.fast_row_id_map = fast_rid;
12388 });
12389 let mut cols = Vec::with_capacity(col_ids.len());
12390 for (index, cid) in col_ids.iter().enumerate() {
12391 execution_checkpoint(control, index)?;
12392 let col = reader.column_native(*cid)?;
12393 cols.push((*cid, col.gather(&positions)));
12394 }
12395 return Ok(Some(cols));
12396 }
12397
12398 if !self.run_refs.is_empty() {
12411 use crate::cursor::{
12412 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
12413 };
12414 let remaining: usize;
12415 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
12416 let c = self
12417 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
12418 .expect("single-run cursor should build when run_refs.len() == 1");
12419 remaining = c.remaining_rows();
12420 Box::new(c)
12421 } else {
12422 let c = self
12423 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
12424 .expect("multi-run cursor should build when run_refs.len() >= 1");
12425 remaining = c.remaining_rows();
12426 Box::new(c)
12427 };
12428 crate::trace::QueryTrace::record(|t| {
12429 if t.survivor_count.is_none() {
12430 t.survivor_count = Some(remaining);
12431 }
12432 });
12433 let cols = match control {
12434 Some(control) => {
12435 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
12436 }
12437 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
12438 };
12439 return Ok(Some(cols));
12440 }
12441
12442 crate::trace::QueryTrace::record(|t| {
12447 t.scan_mode = crate::trace::ScanMode::Materialized;
12448 t.row_materialized = true;
12449 });
12450 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12451 for (index, c) in conditions.iter().enumerate() {
12452 execution_checkpoint(control, index)?;
12453 sets.push(self.resolve_condition(c, snapshot)?);
12454 }
12455 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
12456 let rows = self.rows_for_rids(&rids, snapshot)?;
12457 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
12458 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
12459 execution_checkpoint(control, index)?;
12460 let vals: Vec<Value> = rows
12461 .iter()
12462 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
12463 .collect();
12464 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
12465 }
12466 Ok(Some(cols))
12467 }
12468
12469 pub fn native_page_cursor(
12484 &self,
12485 snapshot: Snapshot,
12486 projection: Vec<(u16, TypeId)>,
12487 conditions: &[crate::query::Condition],
12488 ) -> Result<Option<NativePageCursor>> {
12489 use crate::cursor::build_page_plans;
12490 if self.ttl.is_some() {
12491 return Ok(None);
12492 }
12493 if !conditions.is_empty() && !self.indexes_complete {
12496 return Ok(None);
12497 }
12498 if self.run_refs.len() != 1 {
12499 return Ok(None);
12500 }
12501 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
12502 let (positions, rids) = reader.visible_positions_with_rids_at(snapshot)?;
12503
12504 let overlay_rids: HashSet<u64> = {
12507 let mut s = HashSet::new();
12508 for row in self.memtable.visible_versions_at(snapshot) {
12509 s.insert(row.row_id.0);
12510 }
12511 for row in self.mutable_run.visible_versions_at(snapshot) {
12512 s.insert(row.row_id.0);
12513 }
12514 s
12515 };
12516
12517 let survivors = if conditions.is_empty() {
12521 None
12522 } else {
12523 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
12524 };
12525
12526 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
12533 survivors.clone()
12534 } else if let Some(s) = &survivors {
12535 let mut run_set = s.clone();
12536 run_set.remove_many(overlay_rids.iter().copied());
12537 Some(run_set)
12538 } else {
12539 Some(RowIdSet::from_unsorted(
12540 rids.iter()
12541 .map(|&r| r as u64)
12542 .filter(|r| !overlay_rids.contains(r))
12543 .collect(),
12544 ))
12545 };
12546
12547 let overlay_rows = if overlay_rids.is_empty() {
12548 Vec::new()
12549 } else {
12550 let bound = Self::overlay_materialization_bound(conditions, &survivors);
12551 self.overlay_visible_rows(snapshot, bound)
12552 };
12553
12554 let plans = if positions.is_empty() {
12556 Vec::new()
12557 } else {
12558 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
12559 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
12560 };
12561
12562 let overlay = if overlay_rows.is_empty() {
12564 None
12565 } else {
12566 let filtered =
12567 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
12568 if filtered.is_empty() {
12569 None
12570 } else {
12571 Some(self.materialize_overlay(&filtered, &projection))
12572 }
12573 };
12574
12575 let overlay_row_count = overlay
12576 .as_ref()
12577 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
12578 .unwrap_or(0);
12579 crate::trace::QueryTrace::record(|t| {
12580 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
12581 t.run_count = self.run_refs.len();
12582 t.memtable_rows = self.memtable.len();
12583 t.mutable_run_rows = self.mutable_run.len();
12584 t.overlay_rows = overlay_row_count;
12585 t.conditions_pushed = conditions.len();
12586 t.pages_decoded = plans
12587 .iter()
12588 .map(|p| p.positions.len())
12589 .sum::<usize>()
12590 .min(1);
12591 });
12592
12593 Ok(Some(NativePageCursor::new_with_overlay(
12594 reader, projection, plans, overlay,
12595 )))
12596 }
12597 #[allow(clippy::type_complexity)]
12607 pub fn native_multi_run_cursor(
12608 &self,
12609 snapshot: Snapshot,
12610 projection: Vec<(u16, TypeId)>,
12611 conditions: &[crate::query::Condition],
12612 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
12613 use crate::cursor::{MultiRunCursor, RunStream};
12614 use crate::sorted_run::SYS_ROW_ID;
12615 use std::collections::{BinaryHeap, HashMap, HashSet};
12616 if self.ttl.is_some() {
12617 return Ok(None);
12618 }
12619 if !conditions.is_empty() && !self.indexes_complete {
12622 return Ok(None);
12623 }
12624 if self.run_refs.is_empty() {
12625 return Ok(None);
12626 }
12627
12628 let mut run_meta: Vec<(
12630 RunReader,
12631 Vec<i64>,
12632 Vec<i64>,
12633 Vec<u8>,
12634 Option<Vec<Option<HlcTimestamp>>>,
12635 Vec<usize>,
12636 )> = Vec::with_capacity(self.run_refs.len());
12637 for rr in &self.run_refs {
12638 let mut reader = self.open_reader(rr.run_id)?;
12639 let (rids, eps, del) = reader.system_columns_native()?;
12640 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
12641 let commit_ts: Option<Vec<Option<HlcTimestamp>>> =
12642 if reader.has_column(crate::sorted_run::SYS_COMMIT_TS) {
12643 let col = reader.column_native(crate::sorted_run::SYS_COMMIT_TS)?;
12644 Some(
12645 (0..rids.len())
12646 .map(|i| {
12647 crate::sorted_run::decode_commit_ts_value(col.value_at(i).as_ref())
12648 })
12649 .collect(),
12650 )
12651 } else {
12652 None
12653 };
12654 run_meta.push((reader, rids, eps, del, commit_ts, page_rows));
12655 }
12656
12657 let mut best: HashMap<u64, (Epoch, Option<HlcTimestamp>, usize, usize, bool)> =
12663 HashMap::new();
12664 for (run_idx, (_, rids, eps, del, commit_ts, _)) in run_meta.iter().enumerate() {
12665 for i in 0..rids.len() {
12666 let rid = rids[i] as u64;
12667 let e = Epoch(eps[i] as u64);
12668 let ts = commit_ts.as_ref().and_then(|v| v.get(i).copied().flatten());
12669 if !snapshot.observes_row(e, ts) {
12670 continue;
12671 }
12672 let is_del = del[i] != 0;
12673 best.entry(rid)
12674 .and_modify(|cur| {
12675 if Snapshot::version_is_newer(e, ts, cur.0, cur.1) {
12676 *cur = (e, ts, run_idx, i, is_del);
12677 }
12678 })
12679 .or_insert((e, ts, run_idx, i, is_del));
12680 }
12681 }
12682
12683 let overlay_rids: HashSet<u64> = {
12685 let mut s = HashSet::new();
12686 for row in self.memtable.visible_versions_at(snapshot) {
12687 s.insert(row.row_id.0);
12688 }
12689 for row in self.mutable_run.visible_versions_at(snapshot) {
12690 s.insert(row.row_id.0);
12691 }
12692 s
12693 };
12694
12695 let survivors: Option<RowIdSet> = if conditions.is_empty() {
12697 None
12698 } else {
12699 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12700 for c in conditions {
12701 sets.push(self.resolve_condition(c, snapshot)?);
12702 }
12703 Some(RowIdSet::intersect_many(sets))
12704 };
12705
12706 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
12710 for (rid, (_, _, run_idx, pos, deleted)) in &best {
12711 if *deleted {
12712 continue;
12713 }
12714 if overlay_rids.contains(rid) {
12715 continue;
12716 }
12717 if let Some(s) = &survivors {
12718 if !s.contains(*rid) {
12719 continue;
12720 }
12721 }
12722 per_run[*run_idx].push((*rid, *pos));
12723 }
12724 for v in per_run.iter_mut() {
12725 v.sort_unstable_by_key(|&(rid, _)| rid);
12726 }
12727
12728 let mut streams = Vec::with_capacity(run_meta.len());
12730 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
12731 let mut total = 0usize;
12732 for (run_idx, (reader, _, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
12733 let mut starts = Vec::with_capacity(page_rows.len());
12734 let mut acc = 0usize;
12735 for &r in &page_rows {
12736 starts.push(acc);
12737 acc += r;
12738 }
12739 let mut survivors_vec: Vec<(u64, usize, usize)> =
12740 Vec::with_capacity(per_run[run_idx].len());
12741 for &(rid, pos) in &per_run[run_idx] {
12742 let page_seq = match starts.partition_point(|&s| s <= pos) {
12743 0 => continue,
12744 p => p - 1,
12745 };
12746 let within = pos - starts[page_seq];
12747 survivors_vec.push((rid, page_seq, within));
12748 }
12749 total += survivors_vec.len();
12750 if let Some(&(rid, _, _)) = survivors_vec.first() {
12751 heap.push(std::cmp::Reverse((rid, run_idx)));
12752 }
12753 streams.push(RunStream::new(reader, survivors_vec, page_rows));
12754 }
12755
12756 let overlay_rows = if overlay_rids.is_empty() {
12758 Vec::new()
12759 } else {
12760 let bound = Self::overlay_materialization_bound(conditions, &survivors);
12761 self.overlay_visible_rows(snapshot, bound)
12762 };
12763 let overlay = if overlay_rows.is_empty() {
12764 None
12765 } else {
12766 let filtered =
12767 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
12768 if filtered.is_empty() {
12769 None
12770 } else {
12771 Some(self.materialize_overlay(&filtered, &projection))
12772 }
12773 };
12774
12775 let overlay_row_count = overlay
12776 .as_ref()
12777 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
12778 .unwrap_or(0);
12779 crate::trace::QueryTrace::record(|t| {
12780 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
12781 t.run_count = self.run_refs.len();
12782 t.memtable_rows = self.memtable.len();
12783 t.mutable_run_rows = self.mutable_run.len();
12784 t.overlay_rows = overlay_row_count;
12785 t.conditions_pushed = conditions.len();
12786 t.survivor_count = Some(total);
12787 });
12788
12789 Ok(Some(MultiRunCursor::new(
12790 streams, projection, heap, total, overlay,
12791 )))
12792 }
12793
12794 fn overlay_materialization_bound<'a>(
12806 conditions: &[crate::query::Condition],
12807 survivors: &'a Option<RowIdSet>,
12808 ) -> Option<&'a RowIdSet> {
12809 use crate::query::Condition;
12810 let has_range = conditions
12811 .iter()
12812 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12813 if has_range {
12814 None
12815 } else {
12816 survivors.as_ref()
12817 }
12818 }
12819
12820 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
12832 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
12833 let mut fold = |row: Row| {
12834 if let Some(b) = bound {
12835 if !b.contains(row.row_id.0) {
12836 return;
12837 }
12838 }
12839 best.entry(row.row_id.0)
12840 .and_modify(|(be, br)| {
12841 if row.committed_epoch > *be {
12842 *be = row.committed_epoch;
12843 *br = row.clone();
12844 }
12845 })
12846 .or_insert_with(|| (row.committed_epoch, row));
12847 };
12848 for row in self.memtable.visible_versions_at(snapshot) {
12849 fold(row);
12850 }
12851 for row in self.mutable_run.visible_versions_at(snapshot) {
12852 fold(row);
12853 }
12854 let mut out: Vec<Row> = best
12855 .into_values()
12856 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
12857 .collect();
12858 out.sort_by_key(|r| r.row_id);
12859 out
12860 }
12861
12862 fn filter_overlay_rows(
12870 &self,
12871 rows: Vec<Row>,
12872 conditions: &[crate::query::Condition],
12873 survivors: Option<&RowIdSet>,
12874 snapshot: Snapshot,
12875 ) -> Result<Vec<Row>> {
12876 if conditions.is_empty() {
12877 return Ok(rows);
12878 }
12879 use crate::query::Condition;
12880 let all_index_served = !conditions
12884 .iter()
12885 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12886 if all_index_served {
12887 return Ok(rows
12888 .into_iter()
12889 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
12890 .collect());
12891 }
12892 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12895 for c in conditions {
12896 let s = match c {
12897 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
12898 _ => self.resolve_condition(c, snapshot)?,
12899 };
12900 per_cond_sets.push(s);
12901 }
12902 Ok(rows
12903 .into_iter()
12904 .filter(|row| {
12905 conditions.iter().enumerate().all(|(i, c)| match c {
12906 Condition::Range { column_id, lo, hi } => {
12907 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
12908 }
12909 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
12910 match row.columns.get(column_id) {
12911 Some(Value::Float64(v)) => {
12912 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
12913 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
12914 lo_ok && hi_ok
12915 }
12916 _ => false,
12917 }
12918 }
12919 _ => per_cond_sets[i].contains(row.row_id.0),
12920 })
12921 })
12922 .collect())
12923 }
12924
12925 fn materialize_overlay(
12928 &self,
12929 rows: &[Row],
12930 projection: &[(u16, TypeId)],
12931 ) -> Vec<columnar::NativeColumn> {
12932 if projection.is_empty() {
12933 return vec![columnar::null_native(TypeId::Int64, rows.len())];
12934 }
12935 let mut cols = Vec::with_capacity(projection.len());
12936 for (cid, ty) in projection {
12937 let vals: Vec<Value> = rows
12938 .iter()
12939 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
12940 .collect();
12941 cols.push(columnar::values_to_native(ty.clone(), &vals));
12942 }
12943 cols
12944 }
12945
12946 fn resolve_survivor_rids(
12951 &self,
12952 conditions: &[crate::query::Condition],
12953 reader: &mut RunReader,
12954 snapshot: Snapshot,
12955 ) -> Result<RowIdSet> {
12956 use crate::query::Condition;
12957 let mut sets: Vec<RowIdSet> = Vec::new();
12958 for c in conditions {
12959 self.validate_condition(c)?;
12960 let s: RowIdSet = match c {
12961 Condition::Pk(key) => {
12962 let lookup = self
12963 .schema
12964 .primary_key()
12965 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
12966 .unwrap_or_else(|| key.clone());
12967 self.hot
12968 .get(&lookup)
12969 .map(|r| RowIdSet::one(r.0))
12970 .unwrap_or_else(RowIdSet::empty)
12971 }
12972 Condition::BitmapEq { column_id, value } => {
12973 let lookup = self.index_lookup_key_bytes(*column_id, value);
12974 self.bitmap
12975 .get(column_id)
12976 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
12977 .unwrap_or_else(RowIdSet::empty)
12978 }
12979 Condition::BitmapIn { column_id, values } => {
12980 let bm = self.bitmap.get(column_id);
12981 let mut acc = roaring::RoaringBitmap::new();
12982 if let Some(b) = bm {
12983 for v in values {
12984 let lookup = self.index_lookup_key_bytes(*column_id, v);
12985 acc |= b.get(&lookup);
12986 }
12987 }
12988 RowIdSet::from_roaring(acc)
12989 }
12990 Condition::BytesPrefix { column_id, prefix } => {
12991 if let Some(b) = self.bitmap.get(column_id) {
12992 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
12993 let mut acc = roaring::RoaringBitmap::new();
12994 for key in b.keys() {
12995 if key.starts_with(&lookup_prefix) {
12996 acc |= b.get(&key);
12997 }
12998 }
12999 RowIdSet::from_roaring(acc)
13000 } else {
13001 RowIdSet::empty()
13002 }
13003 }
13004 Condition::FmContains { column_id, pattern } => self
13005 .fm
13006 .get(column_id)
13007 .map(|f| {
13008 RowIdSet::from_unsorted(
13009 f.locate(pattern).into_iter().map(|r| r.0).collect(),
13010 )
13011 })
13012 .unwrap_or_else(RowIdSet::empty),
13013 Condition::FmContainsAll {
13014 column_id,
13015 patterns,
13016 } => {
13017 if let Some(f) = self.fm.get(column_id) {
13018 let sets: Vec<RowIdSet> = patterns
13019 .iter()
13020 .map(|pat| {
13021 RowIdSet::from_unsorted(
13022 f.locate(pat).into_iter().map(|r| r.0).collect(),
13023 )
13024 })
13025 .collect();
13026 RowIdSet::intersect_many(sets)
13027 } else {
13028 RowIdSet::empty()
13029 }
13030 }
13031 Condition::Ann {
13032 column_id,
13033 query,
13034 k,
13035 } => RowIdSet::from_unsorted(
13036 self.retrieve_filtered(
13037 &crate::query::Retriever::Ann {
13038 column_id: *column_id,
13039 query: query.clone(),
13040 k: *k,
13041 },
13042 snapshot,
13043 None,
13044 None,
13045 None,
13046 None,
13047 )?
13048 .into_iter()
13049 .map(|hit| hit.row_id.0)
13050 .collect(),
13051 ),
13052 Condition::SparseMatch {
13053 column_id,
13054 query,
13055 k,
13056 } => RowIdSet::from_unsorted(
13057 self.retrieve_filtered(
13058 &crate::query::Retriever::Sparse {
13059 column_id: *column_id,
13060 query: query.clone(),
13061 k: *k,
13062 },
13063 snapshot,
13064 None,
13065 None,
13066 None,
13067 None,
13068 )?
13069 .into_iter()
13070 .map(|hit| hit.row_id.0)
13071 .collect(),
13072 ),
13073 Condition::MinHashSimilar {
13074 column_id,
13075 query,
13076 k,
13077 } => match self.minhash.get(column_id) {
13078 Some(index) => {
13079 let candidates = index.candidate_row_ids(query);
13080 let eligible =
13081 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
13082 RowIdSet::from_unsorted(
13083 index
13084 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
13085 .into_iter()
13086 .map(|(row_id, _)| row_id.0)
13087 .collect(),
13088 )
13089 }
13090 None => RowIdSet::empty(),
13091 },
13092 Condition::Range { column_id, lo, hi } => {
13093 if let Some(li) = self.learned_range.get(column_id) {
13094 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
13095 } else {
13096 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
13097 }
13098 }
13099 Condition::RangeF64 {
13100 column_id,
13101 lo,
13102 lo_inclusive,
13103 hi,
13104 hi_inclusive,
13105 } => {
13106 if let Some(li) = self.learned_range.get(column_id) {
13107 RowIdSet::from_unsorted(
13108 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
13109 .into_iter()
13110 .collect(),
13111 )
13112 } else {
13113 reader.range_row_id_set_f64(
13114 *column_id,
13115 *lo,
13116 *lo_inclusive,
13117 *hi,
13118 *hi_inclusive,
13119 )?
13120 }
13121 }
13122 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
13123 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
13124 };
13125 sets.push(s);
13126 }
13127 Ok(RowIdSet::intersect_many(sets))
13128 }
13129
13130 pub fn scan_cursor(
13151 &self,
13152 snapshot: Snapshot,
13153 projection: Vec<(u16, TypeId)>,
13154 conditions: &[crate::query::Condition],
13155 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
13156 if self.ttl.is_some() {
13157 return Ok(None);
13158 }
13159 if !conditions.is_empty() && !self.indexes_complete {
13165 return Ok(None);
13166 }
13167 if self.run_refs.len() == 1 {
13168 Ok(self
13169 .native_page_cursor(snapshot, projection, conditions)?
13170 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
13171 } else {
13172 Ok(self
13173 .native_multi_run_cursor(snapshot, projection, conditions)?
13174 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
13175 }
13176 }
13177
13178 pub fn aggregate_native(
13192 &self,
13193 snapshot: Snapshot,
13194 column: Option<u16>,
13195 conditions: &[crate::query::Condition],
13196 agg: NativeAgg,
13197 ) -> Result<Option<NativeAggResult>> {
13198 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
13199 }
13200
13201 pub fn aggregate_native_with_control(
13202 &self,
13203 snapshot: Snapshot,
13204 column: Option<u16>,
13205 conditions: &[crate::query::Condition],
13206 agg: NativeAgg,
13207 control: &crate::ExecutionControl,
13208 ) -> Result<Option<NativeAggResult>> {
13209 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
13210 }
13211
13212 fn aggregate_native_inner(
13213 &self,
13214 snapshot: Snapshot,
13215 column: Option<u16>,
13216 conditions: &[crate::query::Condition],
13217 agg: NativeAgg,
13218 control: Option<&crate::ExecutionControl>,
13219 ) -> Result<Option<NativeAggResult>> {
13220 execution_checkpoint(control, 0)?;
13221 if self.ttl.is_some() {
13222 return Ok(None);
13223 }
13224 if self.run_refs.len() == 1 && conditions.is_empty() {
13226 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
13227 return Ok(Some(res));
13228 }
13229 }
13230 if matches!(agg, NativeAgg::Count) && column.is_none() {
13234 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
13235 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
13236 }
13237 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
13238 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
13239 }
13240 let cid = match column {
13243 Some(c) => c,
13244 None => return Ok(None),
13245 };
13246 let ty = self.column_type(cid);
13247 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
13248 execution_checkpoint(control, 0)?;
13249 return match ty {
13250 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
13251 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
13252 Ok(Some(pack_int(agg, count, sum, mn, mx)))
13253 }
13254 TypeId::Float64 => {
13255 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
13256 Ok(Some(pack_float(agg, count, sum, mn, mx)))
13257 }
13258 _ => Ok(None),
13259 };
13260 }
13261 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
13263 execution_checkpoint(control, 0)?;
13264 match ty {
13265 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
13266 let mut count = 0u64;
13267 let mut sum = 0i128;
13268 let mut mn = i64::MAX;
13269 let mut mx = i64::MIN;
13270 for row in &rows {
13271 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
13272 count += 1;
13273 sum += i128::from(*v);
13274 mn = mn.min(*v);
13275 mx = mx.max(*v);
13276 }
13277 }
13278 Ok(Some(pack_int(agg, count, sum, mn, mx)))
13279 }
13280 TypeId::Float64 => {
13281 let mut count = 0u64;
13282 let mut sum = 0.0f64;
13283 let mut mn = f64::INFINITY;
13284 let mut mx = f64::NEG_INFINITY;
13285 for row in &rows {
13286 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
13287 count += 1;
13288 sum += *v;
13289 mn = mn.min(*v);
13290 mx = mx.max(*v);
13291 }
13292 }
13293 Ok(Some(pack_float(agg, count, sum, mn, mx)))
13294 }
13295 _ => Ok(None),
13296 }
13297 }
13298
13299 fn visible_rows_filtered(
13301 &self,
13302 snapshot: Snapshot,
13303 conditions: &[crate::query::Condition],
13304 control: Option<&crate::ExecutionControl>,
13305 ) -> Result<Vec<Row>> {
13306 let rows = if let Some(control) = control {
13307 self.visible_rows_controlled(snapshot, control)?
13308 } else {
13309 self.visible_rows(snapshot)?
13310 };
13311 if conditions.is_empty() {
13312 return Ok(rows);
13313 }
13314 Ok(rows
13315 .into_iter()
13316 .filter(|row| {
13317 conditions
13318 .iter()
13319 .all(|cond| condition_matches_row(cond, row, &self.schema))
13320 })
13321 .collect())
13322 }
13323
13324 fn aggregate_from_stats(
13332 &self,
13333 snapshot: Snapshot,
13334 column: Option<u16>,
13335 agg: NativeAgg,
13336 ) -> Result<Option<NativeAggResult>> {
13337 let cid = match (agg, column) {
13338 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
13339 _ => return Ok(None), };
13341 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
13342 return Ok(None);
13343 };
13344 let Some(cs) = stats.get(&cid) else {
13345 return Ok(None);
13346 };
13347 match agg {
13348 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
13350 self.live_count.saturating_sub(cs.null_count),
13351 ))),
13352 NativeAgg::Min | NativeAgg::Max => {
13353 let bound = if agg == NativeAgg::Min {
13354 &cs.min
13355 } else {
13356 &cs.max
13357 };
13358 match bound {
13359 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
13360 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
13361 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
13366 None => Ok(None),
13367 }
13368 }
13369 _ => Ok(None),
13370 }
13371 }
13372
13373 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
13382 if self.ttl.is_some() {
13383 return Ok(None);
13384 }
13385 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
13386 return Ok(None);
13387 }
13388 self.ensure_indexes_complete()?;
13391 let reader = self.open_reader(self.run_refs[0].run_id)?;
13392 if self.live_count != reader.row_count() as u64 {
13393 return Ok(None);
13394 }
13395 let Some(bm) = self.bitmap.get(&column_id) else {
13396 return Ok(None); };
13398 let mut distinct = bm.value_count() as u64;
13399 if !bm.get(&Value::Null.encode_key()).is_empty() {
13402 distinct = distinct.saturating_sub(1);
13403 }
13404 Ok(Some(distinct))
13405 }
13406
13407 pub fn aggregate_incremental(
13419 &mut self,
13420 cache_key: u64,
13421 conditions: &[crate::query::Condition],
13422 column: Option<u16>,
13423 agg: NativeAgg,
13424 ) -> Result<IncrementalAggResult> {
13425 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
13426 }
13427
13428 pub fn aggregate_incremental_with_control(
13429 &mut self,
13430 cache_key: u64,
13431 conditions: &[crate::query::Condition],
13432 column: Option<u16>,
13433 agg: NativeAgg,
13434 control: &crate::ExecutionControl,
13435 ) -> Result<IncrementalAggResult> {
13436 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
13437 }
13438
13439 fn aggregate_incremental_inner(
13440 &mut self,
13441 cache_key: u64,
13442 conditions: &[crate::query::Condition],
13443 column: Option<u16>,
13444 agg: NativeAgg,
13445 control: Option<&crate::ExecutionControl>,
13446 ) -> Result<IncrementalAggResult> {
13447 execution_checkpoint(control, 0)?;
13448 let snap = self.snapshot();
13449 let cur_wm = self.allocator.current().0;
13450 let cur_epoch = snap.epoch.0;
13451 let incremental_ok = self.ttl.is_none()
13458 && !self.had_deletes
13459 && self.memtable.is_empty()
13460 && self.mutable_run.is_empty();
13461
13462 if incremental_ok {
13465 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
13466 if cached.epoch == cur_epoch {
13467 return Ok(IncrementalAggResult {
13468 state: cached.state,
13469 incremental: true,
13470 delta_rows: 0,
13471 });
13472 }
13473 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
13474 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
13475 let mut delta_rids = Vec::with_capacity(delta_len);
13476 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
13477 execution_checkpoint(control, index)?;
13478 delta_rids.push(row_id);
13479 }
13480 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
13481 execution_checkpoint(control, 0)?;
13482 let index_sets = self.resolve_index_conditions(conditions, snap)?;
13483 let delta_state = agg_state_from_rows(
13484 &delta_rows,
13485 conditions,
13486 &index_sets,
13487 column,
13488 agg,
13489 &self.schema,
13490 control,
13491 )?;
13492 let merged = cached.state.merge(delta_state);
13493 let delta_n = delta_rids.len() as u64;
13494 Arc::make_mut(&mut self.agg_cache).insert(
13495 cache_key,
13496 CachedAgg {
13497 state: merged.clone(),
13498 watermark: cur_wm,
13499 epoch: cur_epoch,
13500 },
13501 );
13502 return Ok(IncrementalAggResult {
13503 state: merged,
13504 incremental: true,
13505 delta_rows: delta_n,
13506 });
13507 }
13508 }
13509 }
13510
13511 let cursor_ok =
13516 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
13517 let state = if cursor_ok && agg != NativeAgg::Avg {
13518 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
13519 Some(result) => {
13520 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
13521 }
13522 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
13523 }
13524 } else {
13525 self.agg_state_full_scan(conditions, column, agg, snap, control)?
13526 };
13527 if incremental_ok {
13529 Arc::make_mut(&mut self.agg_cache).insert(
13530 cache_key,
13531 CachedAgg {
13532 state: state.clone(),
13533 watermark: cur_wm,
13534 epoch: cur_epoch,
13535 },
13536 );
13537 }
13538 Ok(IncrementalAggResult {
13539 state,
13540 incremental: false,
13541 delta_rows: 0,
13542 })
13543 }
13544
13545 fn agg_state_full_scan(
13548 &self,
13549 conditions: &[crate::query::Condition],
13550 column: Option<u16>,
13551 agg: NativeAgg,
13552 snap: Snapshot,
13553 control: Option<&crate::ExecutionControl>,
13554 ) -> Result<AggState> {
13555 execution_checkpoint(control, 0)?;
13556 let rows = self.visible_rows(snap)?;
13557 execution_checkpoint(control, 0)?;
13558 let index_sets = self.resolve_index_conditions(conditions, snap)?;
13559 agg_state_from_rows(
13560 &rows,
13561 conditions,
13562 &index_sets,
13563 column,
13564 agg,
13565 &self.schema,
13566 control,
13567 )
13568 }
13569
13570 fn resolve_index_conditions(
13573 &self,
13574 conditions: &[crate::query::Condition],
13575 snapshot: Snapshot,
13576 ) -> Result<Vec<RowIdSet>> {
13577 use crate::query::Condition;
13578 let mut sets = Vec::new();
13579 for c in conditions {
13580 if matches!(
13581 c,
13582 Condition::Ann { .. }
13583 | Condition::SparseMatch { .. }
13584 | Condition::MinHashSimilar { .. }
13585 ) {
13586 sets.push(self.resolve_condition(c, snapshot)?);
13587 }
13588 }
13589 Ok(sets)
13590 }
13591
13592 fn column_type(&self, cid: u16) -> TypeId {
13593 self.schema
13594 .columns
13595 .iter()
13596 .find(|c| c.id == cid)
13597 .map(|c| c.ty.clone())
13598 .unwrap_or(TypeId::Bytes)
13599 }
13600
13601 pub fn approx_aggregate(
13610 &mut self,
13611 conditions: &[crate::query::Condition],
13612 column: Option<u16>,
13613 agg: ApproxAgg,
13614 z: f64,
13615 ) -> Result<Option<ApproxResult>> {
13616 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
13617 }
13618
13619 pub fn approx_aggregate_with_candidate_authorization(
13622 &mut self,
13623 conditions: &[crate::query::Condition],
13624 column: Option<u16>,
13625 agg: ApproxAgg,
13626 z: f64,
13627 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
13628 ) -> Result<Option<ApproxResult>> {
13629 use crate::query::Condition;
13630 self.ensure_reservoir_complete()?;
13631 let mut snapshot = self.snapshot();
13636 let n_pop = self.count();
13637 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
13638 if sample_rids.is_empty() {
13639 return Ok(None);
13640 }
13641 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
13643 if live_sample.is_empty() {
13644 snapshot = Snapshot::unbounded();
13645 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
13646 }
13647 let s = live_sample.len();
13648 if s == 0 {
13649 return Ok(None);
13650 }
13651 let authorized = authorization
13652 .map(|authorization| {
13653 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
13654 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
13655 })
13656 .transpose()?;
13657
13658 let mut index_sets: Vec<RowIdSet> = Vec::new();
13661 for c in conditions {
13662 if matches!(
13663 c,
13664 Condition::Ann { .. }
13665 | Condition::SparseMatch { .. }
13666 | Condition::MinHashSimilar { .. }
13667 ) {
13668 index_sets.push(self.resolve_condition(c, snapshot)?);
13669 }
13670 }
13671
13672 let cid = match (agg, column) {
13674 (ApproxAgg::Count, _) => None,
13675 (_, Some(c)) => Some(c),
13676 _ => return Ok(None),
13677 };
13678 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
13679 for r in &live_sample {
13680 if authorized
13681 .as_ref()
13682 .is_some_and(|authorized| !authorized.contains(&r.row_id))
13683 {
13684 continue;
13685 }
13686 if !conditions
13688 .iter()
13689 .all(|c| condition_matches_row(c, r, &self.schema))
13690 {
13691 continue;
13692 }
13693 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
13695 continue;
13696 }
13697 if let Some(cid) = cid {
13698 let mut cells = r
13699 .columns
13700 .get(&cid)
13701 .cloned()
13702 .map(|value| vec![(cid, value)])
13703 .unwrap_or_default();
13704 if let Some(authorization) = authorization {
13705 authorization.security.apply_masks_to_cells(
13706 authorization.table,
13707 &mut cells,
13708 authorization.principal,
13709 );
13710 }
13711 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
13712 passing_vals.push(v);
13713 } } else {
13715 passing_vals.push(0.0); }
13717 }
13718 let m = passing_vals.len();
13719
13720 let (point, half) = match agg {
13721 ApproxAgg::Count => {
13722 let p = m as f64 / s as f64;
13724 let point = n_pop as f64 * p;
13725 let var = if s > 1 {
13726 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
13727 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13728 } else {
13729 0.0
13730 };
13731 (point, z * var.sqrt())
13732 }
13733 ApproxAgg::Sum => {
13734 let y: Vec<f64> = live_sample
13736 .iter()
13737 .map(|r| {
13738 let passes_row = authorized
13739 .as_ref()
13740 .is_none_or(|authorized| authorized.contains(&r.row_id))
13741 && conditions
13742 .iter()
13743 .all(|c| condition_matches_row(c, r, &self.schema))
13744 && index_sets.iter().all(|set| set.contains(r.row_id.0));
13745 if passes_row {
13746 cid.and_then(|cid| {
13747 let mut cells = r
13748 .columns
13749 .get(&cid)
13750 .cloned()
13751 .map(|value| vec![(cid, value)])
13752 .unwrap_or_default();
13753 if let Some(authorization) = authorization {
13754 authorization.security.apply_masks_to_cells(
13755 authorization.table,
13756 &mut cells,
13757 authorization.principal,
13758 );
13759 }
13760 as_f64(cells.first().map(|(_, value)| value))
13761 })
13762 .unwrap_or(0.0)
13763 } else {
13764 0.0
13765 }
13766 })
13767 .collect();
13768 let mean_y = y.iter().sum::<f64>() / s as f64;
13769 let point = n_pop as f64 * mean_y;
13770 let var = if s > 1 {
13771 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
13772 let var_y = ss / (s - 1) as f64;
13773 n_pop as f64 * n_pop as f64 * var_y / s as f64
13774 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13775 } else {
13776 0.0
13777 };
13778 (point, z * var.sqrt())
13779 }
13780 ApproxAgg::Avg => {
13781 if m == 0 {
13782 return Ok(Some(ApproxResult {
13783 point: 0.0,
13784 ci_low: 0.0,
13785 ci_high: 0.0,
13786 n_population: n_pop,
13787 n_sample_live: s,
13788 n_passing: 0,
13789 }));
13790 }
13791 let mean = passing_vals.iter().sum::<f64>() / m as f64;
13792 let half = if m > 1 {
13793 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
13794 let sd = (ss / (m - 1) as f64).sqrt();
13795 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
13796 z * sd / (m as f64).sqrt() * fpc.sqrt()
13797 } else {
13798 0.0
13799 };
13800 (mean, half)
13801 }
13802 };
13803
13804 Ok(Some(ApproxResult {
13805 point,
13806 ci_low: point - half,
13807 ci_high: point + half,
13808 n_population: n_pop,
13809 n_sample_live: s,
13810 n_passing: m,
13811 }))
13812 }
13813
13814 pub fn exact_column_stats(
13822 &self,
13823 _snapshot: Snapshot,
13824 projection: &[u16],
13825 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
13826 if self.ttl.is_some()
13827 || !(self.memtable.is_empty()
13828 && self.mutable_run.is_empty()
13829 && self.run_refs.len() == 1)
13830 {
13831 return Ok(None);
13832 }
13833 let reader = self.open_reader(self.run_refs[0].run_id)?;
13834 if self.live_count != reader.row_count() as u64 {
13835 return Ok(None);
13836 }
13837 let mut out = HashMap::new();
13838 for &cid in projection {
13839 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
13840 Some(c) => c,
13841 None => continue,
13842 };
13843 let Some(stats) = reader.column_page_stats(cid) else {
13845 out.insert(
13846 cid,
13847 ColumnStat {
13848 min: None,
13849 max: None,
13850 null_count: self.live_count,
13851 },
13852 );
13853 continue;
13854 };
13855 let stat = match cdef.ty {
13856 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
13857 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
13858 min: mn.map(Value::Int64),
13859 max: mx.map(Value::Int64),
13860 null_count: n,
13861 })
13862 }
13863 TypeId::Float64 => {
13864 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
13865 min: mn.map(Value::Float64),
13866 max: mx.map(Value::Float64),
13867 null_count: n,
13868 })
13869 }
13870 _ => None,
13871 };
13872 if let Some(s) = stat {
13873 out.insert(cid, s);
13874 }
13875 }
13876 Ok(Some(out))
13877 }
13878
13879 pub fn dir(&self) -> &Path {
13880 &self.dir
13881 }
13882
13883 pub fn schema(&self) -> &Schema {
13884 &self.schema
13885 }
13886
13887 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
13888 self.ann.get(&column_id)
13889 }
13890
13891 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
13892 self.ann.get_mut(&column_id)
13893 }
13894
13895 pub(crate) fn set_catalog_name(&mut self, name: String) {
13896 self.name = name;
13897 }
13898
13899 pub(crate) fn prepare_alter_column(
13900 &mut self,
13901 column_name: &str,
13902 change: &AlterColumn,
13903 ) -> Result<(ColumnDef, Option<Schema>)> {
13904 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
13905 return Err(MongrelError::InvalidArgument(
13906 "ALTER COLUMN requires committing staged writes first".into(),
13907 ));
13908 }
13909 let old = self
13910 .schema
13911 .columns
13912 .iter()
13913 .find(|c| c.name == column_name)
13914 .cloned()
13915 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
13916 let mut next = old.clone();
13917
13918 if let Some(name) = &change.name {
13919 let trimmed = name.trim();
13920 if trimmed.is_empty() {
13921 return Err(MongrelError::InvalidArgument(
13922 "ALTER COLUMN name must not be empty".into(),
13923 ));
13924 }
13925 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
13926 return Err(MongrelError::Schema(format!(
13927 "column {trimmed} already exists"
13928 )));
13929 }
13930 next.name = trimmed.to_string();
13931 }
13932
13933 if let Some(ty) = &change.ty {
13934 next.ty = ty.clone();
13935 }
13936 if let Some(flags) = change.flags {
13937 validate_alter_column_flags(old.flags, flags)?;
13938 next.flags = flags;
13939 }
13940
13941 if let Some(default_change) = &change.default_value {
13942 next.default_value = default_change.clone();
13943 }
13944 if let Some(source_change) = &change.embedding_source {
13945 next.embedding_source = source_change.clone();
13946 }
13947
13948 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
13949 if old.flags.contains(ColumnFlags::NULLABLE)
13950 && !next.flags.contains(ColumnFlags::NULLABLE)
13951 && self.column_has_nulls(old.id)?
13952 {
13953 return Err(MongrelError::InvalidArgument(format!(
13954 "column '{}' contains NULL values",
13955 old.name
13956 )));
13957 }
13958 if next == old {
13959 return Ok((next, None));
13960 }
13961 let mut schema = self.schema.clone();
13962 let index = schema
13963 .columns
13964 .iter()
13965 .position(|column| column.id == next.id)
13966 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
13967 schema.columns[index] = next.clone();
13968 schema.schema_id = schema
13969 .schema_id
13970 .checked_add(1)
13971 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13972 schema.validate_auto_increment()?;
13973 schema.validate_defaults()?;
13974 Ok((next, Some(schema)))
13975 }
13976
13977 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
13978 self.schema = schema;
13979 self.auto_inc = resolve_auto_inc(&self.schema);
13980 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
13981 self.clear_result_cache();
13982 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
13983 }
13984
13985 pub(crate) fn publish_index_schema_change(
13989 &mut self,
13990 schema: Schema,
13991 artifact: SecondaryIndexArtifact,
13992 ) -> Result<()> {
13993 self.apply_altered_schema_prepared(schema);
13994 self.prune_index_maps_to_schema();
13995 match artifact {
13996 SecondaryIndexArtifact::Bitmap(column_id, index) => {
13997 self.bitmap.insert(column_id, index);
13998 }
13999 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
14000 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
14001 }
14002 SecondaryIndexArtifact::Fm(column_id, index) => {
14003 self.fm.insert(column_id, *index);
14004 }
14005 SecondaryIndexArtifact::Ann(column_id, index) => {
14006 self.ann.insert(column_id, *index);
14007 }
14008 SecondaryIndexArtifact::Sparse(column_id, index) => {
14009 self.sparse.insert(column_id, index);
14010 }
14011 SecondaryIndexArtifact::MinHash(column_id, index) => {
14012 self.minhash.insert(column_id, index);
14013 }
14014 }
14015 self.indexes_complete = true;
14016 self.seal_generations();
14017 let view = Arc::new(self.capture_read_generation());
14018 self.published.store(Arc::clone(&view));
14019 checkpoint_current_schema(self)?;
14020 self.invalidate_index_checkpoint();
14023 Ok(())
14024 }
14025
14026 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
14032 self.apply_altered_schema_prepared(schema);
14033 self.prune_index_maps_to_schema();
14034 self.indexes_complete = true;
14035 self.seal_generations();
14036 let view = Arc::new(self.capture_read_generation());
14037 self.published.store(Arc::clone(&view));
14038 checkpoint_current_schema(self)?;
14039 self.invalidate_index_checkpoint();
14041 Ok(())
14042 }
14043
14044 fn prune_index_maps_to_schema(&mut self) {
14045 let keep_bitmap: std::collections::HashSet<u16> = self
14046 .schema
14047 .indexes
14048 .iter()
14049 .filter(|index| index.kind == IndexKind::Bitmap)
14050 .map(|index| index.column_id)
14051 .collect();
14052 let keep_ann: std::collections::HashSet<u16> = self
14053 .schema
14054 .indexes
14055 .iter()
14056 .filter(|index| index.kind == IndexKind::Ann)
14057 .map(|index| index.column_id)
14058 .collect();
14059 let keep_fm: std::collections::HashSet<u16> = self
14060 .schema
14061 .indexes
14062 .iter()
14063 .filter(|index| index.kind == IndexKind::FmIndex)
14064 .map(|index| index.column_id)
14065 .collect();
14066 let keep_sparse: std::collections::HashSet<u16> = self
14067 .schema
14068 .indexes
14069 .iter()
14070 .filter(|index| index.kind == IndexKind::Sparse)
14071 .map(|index| index.column_id)
14072 .collect();
14073 let keep_minhash: std::collections::HashSet<u16> = self
14074 .schema
14075 .indexes
14076 .iter()
14077 .filter(|index| index.kind == IndexKind::MinHash)
14078 .map(|index| index.column_id)
14079 .collect();
14080 let keep_learned: std::collections::HashSet<u16> = self
14081 .schema
14082 .indexes
14083 .iter()
14084 .filter(|index| index.kind == IndexKind::LearnedRange)
14085 .map(|index| index.column_id)
14086 .collect();
14087 self.bitmap
14088 .retain(|column_id, _| keep_bitmap.contains(column_id));
14089 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
14090 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
14091 self.sparse
14092 .retain(|column_id, _| keep_sparse.contains(column_id));
14093 self.minhash
14094 .retain(|column_id, _| keep_minhash.contains(column_id));
14095 {
14096 let learned = Arc::make_mut(&mut self.learned_range);
14097 learned.retain(|column_id, _| keep_learned.contains(column_id));
14098 }
14099 }
14100
14101 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
14102 checkpoint_current_schema(self)
14103 }
14104
14105 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
14106 self.ensure_writable()?;
14107 let previous_schema = self.schema.clone();
14108 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
14109 if let Some(schema) = schema {
14110 self.apply_altered_schema_prepared(schema);
14111 self.checkpoint_standalone_schema_change(previous_schema)?;
14112 }
14113 Ok(column)
14114 }
14115
14116 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
14117 if self.live_count == 0 {
14118 return Ok(false);
14119 }
14120 let snap = self.snapshot();
14121 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
14122 Ok(columns
14123 .first()
14124 .map(|(_, col)| col.null_count(col.len()) != 0)
14125 .unwrap_or(true))
14126 }
14127
14128 fn has_stored_versions(&self) -> bool {
14129 !self.memtable.is_empty()
14130 || !self.mutable_run.is_empty()
14131 || self.run_refs.iter().any(|r| r.row_count > 0)
14132 || !self.retiring.is_empty()
14133 }
14134
14135 pub fn add_column(
14140 &mut self,
14141 name: &str,
14142 ty: TypeId,
14143 flags: ColumnFlags,
14144 default_value: Option<crate::schema::DefaultExpr>,
14145 ) -> Result<u16> {
14146 self.add_column_with_id(name, ty, flags, default_value, None)
14147 }
14148
14149 pub fn add_column_with_id(
14150 &mut self,
14151 name: &str,
14152 ty: TypeId,
14153 flags: ColumnFlags,
14154 default_value: Option<crate::schema::DefaultExpr>,
14155 requested_id: Option<u16>,
14156 ) -> Result<u16> {
14157 self.ensure_writable()?;
14158 let previous_schema = self.schema.clone();
14159 let (column, schema) =
14160 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
14161 self.apply_altered_schema_prepared(schema);
14162 self.checkpoint_standalone_schema_change(previous_schema)?;
14163 Ok(column.id)
14164 }
14165
14166 pub(crate) fn prepare_add_column(
14167 &mut self,
14168 name: &str,
14169 ty: TypeId,
14170 flags: ColumnFlags,
14171 default_value: Option<crate::schema::DefaultExpr>,
14172 requested_id: Option<u16>,
14173 ) -> Result<(ColumnDef, Schema)> {
14174 self.ensure_writable()?;
14175 if self.schema.columns.iter().any(|c| c.name == name) {
14176 return Err(MongrelError::Schema(format!(
14177 "column {name} already exists"
14178 )));
14179 }
14180 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
14181 if self.schema.columns.iter().any(|c| c.id == id) {
14182 return Err(MongrelError::Schema(format!(
14183 "column id {id} already exists"
14184 )));
14185 }
14186 id
14187 } else {
14188 self.schema
14189 .columns
14190 .iter()
14191 .map(|c| c.id)
14192 .max()
14193 .unwrap_or(0)
14194 .checked_add(1)
14195 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
14196 };
14197 let column = ColumnDef {
14198 id,
14199 name: name.to_string(),
14200 ty,
14201 flags,
14202 default_value,
14203 embedding_source: None,
14204 };
14205 let mut next_schema = self.schema.clone();
14206 next_schema.columns.push(column.clone());
14207 next_schema.schema_id = next_schema
14208 .schema_id
14209 .checked_add(1)
14210 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
14211 next_schema.validate_auto_increment()?;
14212 next_schema.validate_defaults()?;
14213 Ok((column, next_schema))
14214 }
14215
14216 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
14225 self.ensure_writable()?;
14226 let cid = self
14227 .schema
14228 .columns
14229 .iter()
14230 .find(|c| c.name == column_name)
14231 .map(|c| c.id)
14232 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
14233 let ty = self
14234 .schema
14235 .columns
14236 .iter()
14237 .find(|c| c.id == cid)
14238 .map(|c| c.ty.clone())
14239 .unwrap_or(TypeId::Int64);
14240 if !matches!(
14241 ty,
14242 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
14243 ) {
14244 return Err(MongrelError::Schema(format!(
14245 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
14246 )));
14247 }
14248 if self
14249 .schema
14250 .indexes
14251 .iter()
14252 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
14253 {
14254 return Ok(()); }
14256 let previous_schema = self.schema.clone();
14257 let previous_learned_range = Arc::clone(&self.learned_range);
14258 let mut next_schema = previous_schema.clone();
14259 next_schema.indexes.push(IndexDef {
14260 name: format!("{}_learned_range", column_name),
14261 column_id: cid,
14262 kind: IndexKind::LearnedRange,
14263 predicate: None,
14264 options: Default::default(),
14265 });
14266 next_schema.schema_id = next_schema
14267 .schema_id
14268 .checked_add(1)
14269 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
14270 self.apply_altered_schema_prepared(next_schema);
14271 if let Err(error) = self.build_learned_ranges() {
14272 self.apply_altered_schema_prepared(previous_schema);
14273 self.learned_range = previous_learned_range;
14274 return Err(error);
14275 }
14276 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
14277 if !matches!(
14278 &error,
14279 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
14280 ) {
14281 self.learned_range = previous_learned_range;
14282 }
14283 return Err(error);
14284 }
14285 Ok(())
14286 }
14287
14288 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
14289 let mut schema_published = false;
14290 let schema_result = match self._root_guard.as_deref() {
14291 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
14292 schema_published = true;
14293 }),
14294 None => write_schema_with_after(&self.dir, &self.schema, || {
14295 schema_published = true;
14296 }),
14297 };
14298 if schema_result.is_err() && !schema_published {
14299 self.apply_altered_schema_prepared(previous_schema);
14300 return schema_result;
14301 }
14302
14303 let manifest_result = self.persist_manifest(self.current_epoch());
14304 match (schema_result, manifest_result) {
14305 (_, Ok(())) => Ok(()),
14306 (Ok(()), Err(error)) => {
14307 self.poison_after_maintenance_publish_failure();
14308 Err(MongrelError::DurableCommit {
14309 epoch: self.current_epoch().0,
14310 message: format!(
14311 "schema is durable but matching manifest publication failed: {error}"
14312 ),
14313 })
14314 }
14315 (Err(schema_error), Err(manifest_error)) => {
14316 self.poison_after_maintenance_publish_failure();
14317 Err(MongrelError::CommitOutcomeUnknown {
14318 epoch: self.current_epoch().0,
14319 message: format!(
14320 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
14321 ),
14322 })
14323 }
14324 }
14325 }
14326
14327 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
14330 self.sync_byte_threshold = threshold;
14331 if let WalSink::Private(w) = &mut self.wal {
14332 w.set_sync_byte_threshold(threshold);
14333 }
14334 }
14335
14336 pub fn page_cache_flush(&self) {
14340 self.page_cache.flush_to_disk();
14341 }
14342
14343 pub fn page_cache_len(&self) -> usize {
14345 self.page_cache.len()
14346 }
14347
14348 pub fn decoded_cache_len(&self) -> usize {
14351 self.decoded_cache.len()
14352 }
14353
14354 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
14357 self.memtable.drain_sorted()
14358 }
14359
14360 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
14361 self.runs_dir().join(format!("r-{run_id}.sr"))
14362 }
14363
14364 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
14365 match self.runs_root.as_deref() {
14366 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
14367 None => Ok(None),
14368 }
14369 }
14370
14371 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
14372 match self.runs_root.as_deref() {
14373 Some(root) => Ok(Some(root.create_regular_new(name)?)),
14374 None => Ok(None),
14375 }
14376 }
14377
14378 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
14379 match self.runs_root.as_deref() {
14380 Some(root) => match root.remove_file(name) {
14381 Ok(()) => Ok(()),
14382 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
14383 Err(error) => Err(error.into()),
14384 },
14385 None => match std::fs::remove_file(self.runs_dir().join(name)) {
14386 Ok(()) => Ok(()),
14387 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
14388 Err(error) => Err(error.into()),
14389 },
14390 }
14391 }
14392
14393 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
14394 match self.runs_root.as_deref() {
14395 Some(root) => root
14396 .rename_file_new(source, destination)
14397 .map_err(Into::into),
14398 None => crate::durable_file::rename(
14399 &self.runs_dir().join(source),
14400 &self.runs_dir().join(destination),
14401 )
14402 .map_err(Into::into),
14403 }
14404 }
14405
14406 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
14407 self.run_refs.iter().map(|run| run.run_id)
14408 }
14409
14410 pub(crate) fn table_dir(&self) -> &Path {
14411 &self.dir
14412 }
14413
14414 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
14415 &self.schema
14416 }
14417
14418 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
14419 let id = self.next_run_id;
14420 self.next_run_id = self
14421 .next_run_id
14422 .checked_add(1)
14423 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
14424 Ok(id)
14425 }
14426
14427 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
14428 self.run_refs.push(run_ref);
14429 }
14430
14431 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
14441 self.retiring.push(crate::manifest::RetiredRun {
14442 run_id,
14443 retire_epoch,
14444 });
14445 }
14446
14447 pub(crate) fn reap_retiring(
14451 &mut self,
14452 min_active: Epoch,
14453 backup_pinned: &std::collections::HashSet<u128>,
14454 ) -> Result<usize> {
14455 if self.retiring.is_empty() {
14456 return Ok(0);
14457 }
14458 let mut reaped = 0;
14459 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
14460 for r in std::mem::take(&mut self.retiring) {
14466 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
14467 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
14468 reaped += 1;
14469 } else {
14470 kept.push(r);
14471 }
14472 }
14473 self.retiring = kept;
14474 if reaped > 0 {
14475 self.persist_manifest(self.current_epoch())?;
14476 }
14477 Ok(reaped)
14478 }
14479
14480 pub(crate) fn has_reapable_retiring(
14481 &self,
14482 min_active: Epoch,
14483 backup_pinned: &std::collections::HashSet<u128>,
14484 ) -> bool {
14485 self.retiring
14486 .iter()
14487 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
14488 }
14489
14490 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
14491 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
14492 return false;
14493 }
14494 self.live_count = self.live_count.saturating_add(run_ref.row_count);
14495 self.run_refs.push(run_ref);
14496 self.indexes_complete = false;
14497 true
14498 }
14499
14500 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
14501 self.kek.as_ref()
14502 }
14503
14504 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
14505 let mut reader = match self.runs_root.as_deref() {
14506 Some(root) => RunReader::open_file_with_cache(
14507 root.open_regular(format!("r-{run_id}.sr"))?,
14508 self.schema.clone(),
14509 self.kek.clone(),
14510 Some(self.page_cache.clone()),
14511 Some(self.decoded_cache.clone()),
14512 self.table_id,
14513 Some(&self.verified_runs),
14514 None,
14515 )?,
14516 None => RunReader::open_with_cache(
14517 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
14518 self.schema.clone(),
14519 self.kek.clone(),
14520 Some(self.page_cache.clone()),
14521 Some(self.decoded_cache.clone()),
14522 self.table_id,
14523 Some(&self.verified_runs),
14524 )?,
14525 };
14526 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
14530 reader.set_uniform_epoch(Epoch(rr.epoch_created));
14531 }
14532 Ok(reader)
14533 }
14534
14535 pub(crate) fn run_refs(&self) -> &[RunRef] {
14536 &self.run_refs
14537 }
14538
14539 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
14540 self.retiring.iter().map(|run| run.run_id)
14541 }
14542
14543 pub(crate) fn runs_dir(&self) -> PathBuf {
14544 self.runs_root
14545 .as_deref()
14546 .and_then(|root| root.io_path().ok())
14547 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
14548 }
14549
14550 pub(crate) fn wal_dir(&self) -> PathBuf {
14551 self.dir.join(WAL_DIR)
14552 }
14553
14554 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
14555 self.run_refs = refs;
14556 self.refresh_run_row_id_ranges();
14557 }
14558
14559 pub(crate) fn run_row_id_range(&self, run_id: u128) -> Option<(u64, u64)> {
14563 self.run_row_id_ranges.get(&run_id).copied()
14564 }
14565
14566 pub(crate) fn refresh_run_row_id_ranges(&mut self) {
14569 let mut ranges = HashMap::with_capacity(self.run_refs.len());
14570 for rr in &self.run_refs {
14571 if let Ok(reader) = self.open_reader(rr.run_id) {
14572 let h = reader.header();
14573 ranges.insert(rr.run_id, (h.min_row_id, h.max_row_id));
14574 }
14575 }
14576 self.run_row_id_ranges = ranges;
14577 }
14578
14579 pub(crate) fn compaction_zstd_level(&self) -> i32 {
14580 self.compaction_zstd_level
14581 }
14582
14583 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
14584 self.kek.clone()
14585 }
14586
14587 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
14591 self.kek.as_ref().map(|k| k.derive_idx_key())
14592 }
14593
14594 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
14598 self.kek.as_ref().map(|k| *k.derive_meta_key())
14599 }
14600
14601 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
14604 self.column_keys
14605 .iter()
14606 .map(|(&id, &(_, scheme))| (id, scheme))
14607 .collect()
14608 }
14609
14610 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
14615 self.tokenize_value_enc(column_id, v)
14616 }
14617
14618 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
14619 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
14620 let (key, scheme) = self.column_keys.get(&column_id)?;
14621 let token: Vec<u8> = match (*scheme, v) {
14622 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
14623 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
14624 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
14625 _ => hmac_token(key, &v.encode_key()).to_vec(),
14626 };
14627 Some(Value::Bytes(token))
14628 }
14629
14630 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
14632 self.index_lookup_key_bytes(column_id, &v.encode_key())
14633 }
14634
14635 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
14638 {
14639 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
14640 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
14641 if *scheme == SCHEME_HMAC_EQ {
14642 return hmac_token(key, encoded).to_vec();
14643 }
14644 }
14645 }
14646 let _ = column_id;
14647 encoded.to_vec()
14648 }
14649}
14650
14651fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
14652 let columnar::NativeColumn::Int64 { data, validity } = col else {
14653 return false;
14654 };
14655 if data.len() < n || !columnar::all_non_null(validity, n) {
14656 return false;
14657 }
14658 data.iter()
14659 .take(n)
14660 .zip(data.iter().skip(1))
14661 .all(|(a, b)| a < b)
14662}
14663
14664#[derive(Debug, Clone)]
14668pub struct ColumnStat {
14669 pub min: Option<Value>,
14670 pub max: Option<Value>,
14671 pub null_count: u64,
14672}
14673
14674#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14676pub enum NativeAgg {
14677 Count,
14678 Sum,
14679 Min,
14680 Max,
14681 Avg,
14682}
14683
14684#[derive(Debug, Clone, PartialEq)]
14686pub enum NativeAggResult {
14687 Count(u64),
14688 Int(i64),
14689 Float(f64),
14690 Null,
14692}
14693
14694#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14696pub enum ApproxAgg {
14697 Count,
14698 Sum,
14699 Avg,
14700}
14701
14702#[derive(Debug, Clone)]
14706pub struct ApproxResult {
14707 pub point: f64,
14709 pub ci_low: f64,
14711 pub ci_high: f64,
14713 pub n_population: u64,
14715 pub n_sample_live: usize,
14717 pub n_passing: usize,
14719}
14720
14721#[derive(Debug, Clone, PartialEq)]
14726pub enum AggState {
14727 Count(u64),
14729 SumI {
14731 sum: i128,
14732 count: u64,
14733 },
14734 SumF {
14736 sum: f64,
14737 count: u64,
14738 },
14739 AvgI {
14741 sum: i128,
14742 count: u64,
14743 },
14744 AvgF {
14746 sum: f64,
14747 count: u64,
14748 },
14749 MinI(i64),
14751 MaxI(i64),
14752 MinF(f64),
14754 MaxF(f64),
14755 Empty,
14757}
14758
14759impl AggState {
14760 pub fn merge(self, other: AggState) -> AggState {
14762 use AggState::*;
14763 match (self, other) {
14764 (Empty, x) | (x, Empty) => x,
14765 (Count(a), Count(b)) => Count(a + b),
14766 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
14767 sum: sa + sb,
14768 count: ca + cb,
14769 },
14770 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
14771 sum: sa + sb,
14772 count: ca + cb,
14773 },
14774 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
14775 sum: sa + sb,
14776 count: ca + cb,
14777 },
14778 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
14779 sum: sa + sb,
14780 count: ca + cb,
14781 },
14782 (MinI(a), MinI(b)) => MinI(a.min(b)),
14783 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
14784 (MinF(a), MinF(b)) => MinF(a.min(b)),
14785 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
14786 _ => Empty, }
14788 }
14789
14790 pub fn point(&self) -> Option<f64> {
14792 match self {
14793 AggState::Count(n) => Some(*n as f64),
14794 AggState::SumI { sum, .. } => Some(*sum as f64),
14795 AggState::SumF { sum, .. } => Some(*sum),
14796 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
14797 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
14798 AggState::MinI(n) => Some(*n as f64),
14799 AggState::MaxI(n) => Some(*n as f64),
14800 AggState::MinF(n) => Some(*n),
14801 AggState::MaxF(n) => Some(*n),
14802 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
14803 }
14804 }
14805
14806 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
14810 let is_float = matches!(ty, Some(TypeId::Float64));
14811 match (agg, result) {
14812 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
14813 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
14814 sum: x as i128,
14815 count: 1, },
14817 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
14818 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
14819 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
14820 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
14821 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
14822 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
14823 (NativeAgg::Count, _) => AggState::Empty,
14824 (_, NativeAggResult::Null) => AggState::Empty,
14825 _ => {
14826 let _ = is_float;
14827 AggState::Empty
14828 }
14829 }
14830 }
14831}
14832
14833#[derive(Debug, Clone)]
14836pub struct CachedAgg {
14837 pub state: AggState,
14838 pub watermark: u64,
14839 pub epoch: u64,
14840}
14841
14842#[derive(Debug, Clone)]
14844pub struct IncrementalAggResult {
14845 pub state: AggState,
14847 pub incremental: bool,
14850 pub delta_rows: u64,
14852}
14853
14854fn agg_state_from_rows(
14858 rows: &[Row],
14859 conditions: &[crate::query::Condition],
14860 index_sets: &[RowIdSet],
14861 column: Option<u16>,
14862 agg: NativeAgg,
14863 schema: &Schema,
14864 control: Option<&crate::ExecutionControl>,
14865) -> Result<AggState> {
14866 let mut count: u64 = 0;
14867 let mut sum_i: i128 = 0;
14868 let mut sum_f: f64 = 0.0;
14869 let mut mn_i: i64 = i64::MAX;
14870 let mut mx_i: i64 = i64::MIN;
14871 let mut mn_f: f64 = f64::INFINITY;
14872 let mut mx_f: f64 = f64::NEG_INFINITY;
14873 let mut saw_int = false;
14874 let mut saw_float = false;
14875 for (index, r) in rows.iter().enumerate() {
14876 execution_checkpoint(control, index)?;
14877 if !conditions
14878 .iter()
14879 .all(|c| condition_matches_row(c, r, schema))
14880 {
14881 continue;
14882 }
14883 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
14884 continue;
14885 }
14886 match agg {
14887 NativeAgg::Count => match column {
14888 None => count += 1,
14890 Some(cid) => match r.columns.get(&cid) {
14893 None | Some(Value::Null) => {}
14894 Some(_) => count += 1,
14895 },
14896 },
14897 _ => match column.and_then(|cid| r.columns.get(&cid)) {
14898 Some(Value::Int64(n)) => {
14899 count += 1;
14900 sum_i += *n as i128;
14901 mn_i = mn_i.min(*n);
14902 mx_i = mx_i.max(*n);
14903 saw_int = true;
14904 }
14905 Some(Value::Float64(f)) => {
14906 count += 1;
14907 sum_f += f;
14908 mn_f = mn_f.min(*f);
14909 mx_f = mx_f.max(*f);
14910 saw_float = true;
14911 }
14912 _ => {}
14913 },
14914 }
14915 }
14916 Ok(match agg {
14917 NativeAgg::Count => {
14918 if count == 0 {
14919 AggState::Empty
14920 } else {
14921 AggState::Count(count)
14922 }
14923 }
14924 NativeAgg::Sum => {
14925 if count == 0 {
14926 AggState::Empty
14927 } else if saw_int {
14928 AggState::SumI { sum: sum_i, count }
14929 } else {
14930 AggState::SumF { sum: sum_f, count }
14931 }
14932 }
14933 NativeAgg::Avg => {
14934 if count == 0 {
14935 AggState::Empty
14936 } else if saw_int {
14937 AggState::AvgI { sum: sum_i, count }
14938 } else {
14939 AggState::AvgF { sum: sum_f, count }
14940 }
14941 }
14942 NativeAgg::Min => {
14943 if !saw_int && !saw_float {
14944 AggState::Empty
14945 } else if saw_int {
14946 AggState::MinI(mn_i)
14947 } else {
14948 AggState::MinF(mn_f)
14949 }
14950 }
14951 NativeAgg::Max => {
14952 if !saw_int && !saw_float {
14953 AggState::Empty
14954 } else if saw_int {
14955 AggState::MaxI(mx_i)
14956 } else {
14957 AggState::MaxF(mx_f)
14958 }
14959 }
14960 })
14961}
14962
14963fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
14967 use crate::query::Condition;
14968 match c {
14969 Condition::Pk(key) => match schema.primary_key() {
14970 Some(pk) => row
14971 .columns
14972 .get(&pk.id)
14973 .map(|v| v.encode_key() == *key)
14974 .unwrap_or(false),
14975 None => false,
14976 },
14977 Condition::BitmapEq { column_id, value } => row
14978 .columns
14979 .get(column_id)
14980 .map(|v| v.encode_key() == *value)
14981 .unwrap_or(false),
14982 Condition::BitmapIn { column_id, values } => {
14983 let key = row.columns.get(column_id).map(|v| v.encode_key());
14984 match key {
14985 Some(k) => values.contains(&k),
14986 None => false,
14987 }
14988 }
14989 Condition::BytesPrefix { column_id, prefix } => row
14990 .columns
14991 .get(column_id)
14992 .map(|v| v.encode_key().starts_with(prefix))
14993 .unwrap_or(false),
14994 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
14995 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
14996 _ => false,
14997 },
14998 Condition::RangeF64 {
14999 column_id,
15000 lo,
15001 lo_inclusive,
15002 hi,
15003 hi_inclusive,
15004 } => match row.columns.get(column_id) {
15005 Some(Value::Float64(n)) => {
15006 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
15007 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
15008 lo_ok && hi_ok
15009 }
15010 _ => false,
15011 },
15012 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
15013 Some(Value::Bytes(b)) => {
15014 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
15015 }
15016 _ => false,
15017 },
15018 Condition::FmContainsAll {
15019 column_id,
15020 patterns,
15021 } => match row.columns.get(column_id) {
15022 Some(Value::Bytes(b)) => patterns
15023 .iter()
15024 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
15025 _ => false,
15026 },
15027 Condition::Ann { .. }
15028 | Condition::SparseMatch { .. }
15029 | Condition::MinHashSimilar { .. } => true,
15030 Condition::IsNull { column_id } => {
15031 matches!(row.columns.get(column_id), Some(Value::Null) | None)
15032 }
15033 Condition::IsNotNull { column_id } => {
15034 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
15035 }
15036 }
15037}
15038
15039fn as_f64(v: Option<&Value>) -> Option<f64> {
15041 match v {
15042 Some(Value::Int64(n)) => Some(*n as f64),
15043 Some(Value::Float64(f)) => Some(*f),
15044 _ => None,
15045 }
15046}
15047
15048fn accumulate_int(
15052 cursor: &mut dyn crate::cursor::Cursor,
15053 control: Option<&crate::ExecutionControl>,
15054) -> Result<(u64, i128, i64, i64)> {
15055 let mut count: u64 = 0;
15056 let mut sum: i128 = 0;
15057 let mut mn: i64 = i64::MAX;
15058 let mut mx: i64 = i64::MIN;
15059 while let Some(cols) = cursor.next_batch()? {
15060 execution_checkpoint(control, 0)?;
15061 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
15062 if crate::columnar::all_non_null(validity, data.len()) {
15063 count += data.len() as u64;
15065 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
15066 execution_checkpoint(control, chunk_index * 1024)?;
15067 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
15068 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
15069 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
15070 }
15071 } else {
15072 for (i, &v) in data.iter().enumerate() {
15073 execution_checkpoint(control, i)?;
15074 if crate::columnar::validity_bit(validity, i) {
15075 count += 1;
15076 sum += v as i128;
15077 mn = mn.min(v);
15078 mx = mx.max(v);
15079 }
15080 }
15081 }
15082 }
15083 }
15084 Ok((count, sum, mn, mx))
15085}
15086
15087fn accumulate_float(
15089 cursor: &mut dyn crate::cursor::Cursor,
15090 control: Option<&crate::ExecutionControl>,
15091) -> Result<(u64, f64, f64, f64)> {
15092 let mut count: u64 = 0;
15093 let mut sum: f64 = 0.0;
15094 let mut mn: f64 = f64::INFINITY;
15095 let mut mx: f64 = f64::NEG_INFINITY;
15096 while let Some(cols) = cursor.next_batch()? {
15097 execution_checkpoint(control, 0)?;
15098 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
15099 if crate::columnar::all_non_null(validity, data.len()) {
15100 count += data.len() as u64;
15101 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
15102 execution_checkpoint(control, chunk_index * 1024)?;
15103 sum += chunk.iter().sum::<f64>();
15104 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
15105 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
15106 }
15107 } else {
15108 for (i, &v) in data.iter().enumerate() {
15109 execution_checkpoint(control, i)?;
15110 if crate::columnar::validity_bit(validity, i) {
15111 count += 1;
15112 sum += v;
15113 mn = mn.min(v);
15114 mx = mx.max(v);
15115 }
15116 }
15117 }
15118 }
15119 }
15120 Ok((count, sum, mn, mx))
15121}
15122
15123#[inline]
15124fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
15125 if index.is_multiple_of(256) {
15126 control
15127 .map(crate::ExecutionControl::checkpoint)
15128 .transpose()?;
15129 }
15130 Ok(())
15131}
15132
15133fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
15134 if count == 0 && !matches!(agg, NativeAgg::Count) {
15135 return NativeAggResult::Null;
15136 }
15137 match agg {
15138 NativeAgg::Count => NativeAggResult::Count(count),
15139 NativeAgg::Sum => match sum.try_into() {
15142 Ok(v) => NativeAggResult::Int(v),
15143 Err(_) => NativeAggResult::Null,
15144 },
15145 NativeAgg::Min => NativeAggResult::Int(mn),
15146 NativeAgg::Max => NativeAggResult::Int(mx),
15147 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
15148 }
15149}
15150
15151fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
15152 if count == 0 && !matches!(agg, NativeAgg::Count) {
15153 return NativeAggResult::Null;
15154 }
15155 match agg {
15156 NativeAgg::Count => NativeAggResult::Count(count),
15157 NativeAgg::Sum => NativeAggResult::Float(sum),
15158 NativeAgg::Min => NativeAggResult::Float(mn),
15159 NativeAgg::Max => NativeAggResult::Float(mx),
15160 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
15161 }
15162}
15163
15164fn agg_int(
15167 stats: &[crate::page::PageStat],
15168 decode: fn(Option<&[u8]>) -> Option<i64>,
15169) -> Option<(Option<i64>, Option<i64>, u64)> {
15170 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
15171 let mut any = false;
15172 for s in stats {
15173 if let Some(v) = decode(s.min.as_deref()) {
15174 mn = mn.min(v);
15175 any = true;
15176 }
15177 if let Some(v) = decode(s.max.as_deref()) {
15178 mx = mx.max(v);
15179 any = true;
15180 }
15181 nulls += s.null_count;
15182 }
15183 any.then_some((Some(mn), Some(mx), nulls))
15184}
15185
15186fn agg_float(
15188 stats: &[crate::page::PageStat],
15189 decode: fn(Option<&[u8]>) -> Option<f64>,
15190) -> Option<(Option<f64>, Option<f64>, u64)> {
15191 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
15192 let mut any = false;
15193 for s in stats {
15194 if let Some(v) = decode(s.min.as_deref()) {
15195 mn = mn.min(v);
15196 any = true;
15197 }
15198 if let Some(v) = decode(s.max.as_deref()) {
15199 mx = mx.max(v);
15200 any = true;
15201 }
15202 nulls += s.null_count;
15203 }
15204 any.then_some((Some(mn), Some(mx), nulls))
15205}
15206
15207type SecondaryIndexes = (
15209 HashMap<u16, BitmapIndex>,
15210 HashMap<u16, AnnIndex>,
15211 HashMap<u16, FmIndex>,
15212 HashMap<u16, SparseIndex>,
15213 HashMap<u16, MinHashIndex>,
15214);
15215
15216fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
15217 let mut bitmap = HashMap::new();
15218 let mut ann = HashMap::new();
15219 let mut fm = HashMap::new();
15220 let mut sparse = HashMap::new();
15221 let mut minhash = HashMap::new();
15222 for idef in &schema.indexes {
15223 match idef.kind {
15224 IndexKind::Bitmap => {
15225 bitmap.insert(idef.column_id, BitmapIndex::new());
15226 }
15227 IndexKind::Ann => {
15228 let dim = schema
15229 .columns
15230 .iter()
15231 .find(|c| c.id == idef.column_id)
15232 .and_then(|c| match c.ty {
15233 TypeId::Embedding { dim } => Some(dim as usize),
15234 _ => None,
15235 })
15236 .unwrap_or(0);
15237 let options = idef.options.ann.clone().unwrap_or_default();
15238 ann.insert(
15239 idef.column_id,
15240 AnnIndex::with_full_options(
15241 dim,
15242 options.m,
15243 options.ef_construction,
15244 options.ef_search,
15245 &options,
15246 ),
15247 );
15248 }
15249 IndexKind::FmIndex => {
15250 fm.insert(idef.column_id, FmIndex::new());
15251 }
15252 IndexKind::Sparse => {
15253 sparse.insert(idef.column_id, SparseIndex::new());
15254 }
15255 IndexKind::MinHash => {
15256 let options = idef.options.minhash.clone().unwrap_or_default();
15257 minhash.insert(
15258 idef.column_id,
15259 MinHashIndex::with_options(options.permutations, options.bands),
15260 );
15261 }
15262 _ => {}
15263 }
15264 }
15265 (bitmap, ann, fm, sparse, minhash)
15266}
15267
15268const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
15269 | ColumnFlags::AUTO_INCREMENT
15270 | ColumnFlags::ENCRYPTED
15271 | ColumnFlags::ENCRYPTED_INDEXABLE
15272 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
15273
15274fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
15275 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
15276 return Err(MongrelError::Schema(
15277 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
15278 ));
15279 }
15280 Ok(())
15281}
15282
15283fn validate_alter_column_type(
15284 schema: &Schema,
15285 old: &ColumnDef,
15286 next: &ColumnDef,
15287 has_stored_versions: bool,
15288) -> Result<()> {
15289 if old.ty == next.ty {
15290 return Ok(());
15291 }
15292 if schema.indexes.iter().any(|i| i.column_id == old.id) {
15293 return Err(MongrelError::Schema(format!(
15294 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
15295 old.name
15296 )));
15297 }
15298 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
15299 return Ok(());
15300 }
15301 Err(MongrelError::Schema(format!(
15302 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
15303 old.ty, next.ty
15304 )))
15305}
15306
15307fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
15308 matches!(
15309 (old, new),
15310 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
15311 )
15312}
15313
15314fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
15320 let mut prev: Option<i64> = None;
15321 for r in rows {
15322 match r.columns.get(&pk_id) {
15323 Some(Value::Int64(v)) => {
15324 if prev.is_some_and(|p| p >= *v) {
15325 return false;
15326 }
15327 prev = Some(*v);
15328 }
15329 _ => return false,
15330 }
15331 }
15332 true
15333}
15334
15335#[allow(clippy::too_many_arguments)]
15336fn index_into(
15337 schema: &Schema,
15338 row: &Row,
15339 hot: &mut HotIndex,
15340 bitmap: &mut HashMap<u16, BitmapIndex>,
15341 ann: &mut HashMap<u16, AnnIndex>,
15342 fm: &mut HashMap<u16, FmIndex>,
15343 sparse: &mut HashMap<u16, SparseIndex>,
15344 minhash: &mut HashMap<u16, MinHashIndex>,
15345) {
15346 for idef in &schema.indexes {
15347 let Some(val) = row.columns.get(&idef.column_id) else {
15348 continue;
15349 };
15350 match idef.kind {
15351 IndexKind::Bitmap => {
15352 if let Some(b) = bitmap.get_mut(&idef.column_id) {
15353 b.insert(val.encode_key(), row.row_id);
15354 }
15355 }
15356 IndexKind::Ann => {
15357 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
15358 if let Some(meta) = val.generated_embedding_metadata() {
15359 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
15361 continue;
15362 }
15363 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
15364 .is_err()
15365 {
15366 continue;
15367 }
15368 }
15369 a.insert_validated(v, row.row_id);
15370 }
15371 }
15372 IndexKind::FmIndex => {
15373 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
15374 f.insert(b.clone(), row.row_id);
15375 }
15376 }
15377 IndexKind::Sparse => {
15378 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
15379 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
15382 s.insert(&terms, row.row_id);
15383 }
15384 }
15385 }
15386 IndexKind::MinHash => {
15387 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
15388 let tokens = crate::index::token_hashes_from_bytes(b);
15391 mh.insert(&tokens, row.row_id);
15392 }
15393 }
15394 _ => {}
15395 }
15396 }
15397 if let Some(pk_col) = schema.primary_key() {
15398 if let Some(pk_val) = row.columns.get(&pk_col.id) {
15399 hot.insert(pk_val.encode_key(), row.row_id);
15400 }
15401 }
15402}
15403
15404#[allow(clippy::too_many_arguments)]
15407fn index_into_single(
15408 idef: &IndexDef,
15409 _schema: &Schema,
15410 row: &Row,
15411 _hot: &mut HotIndex,
15412 bitmap: &mut HashMap<u16, BitmapIndex>,
15413 ann: &mut HashMap<u16, AnnIndex>,
15414 fm: &mut HashMap<u16, FmIndex>,
15415 sparse: &mut HashMap<u16, SparseIndex>,
15416 minhash: &mut HashMap<u16, MinHashIndex>,
15417) {
15418 let Some(val) = row.columns.get(&idef.column_id) else {
15419 return;
15420 };
15421 match idef.kind {
15422 IndexKind::Bitmap => {
15423 if let Some(b) = bitmap.get_mut(&idef.column_id) {
15424 b.insert(val.encode_key(), row.row_id);
15425 }
15426 }
15427 IndexKind::Ann => {
15428 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
15429 if let Some(meta) = val.generated_embedding_metadata() {
15430 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
15432 return;
15433 }
15434 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
15435 .is_err()
15436 {
15437 return;
15438 }
15439 }
15440 a.insert_validated(v, row.row_id);
15441 }
15442 }
15443 IndexKind::FmIndex => {
15444 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
15445 f.insert(b.clone(), row.row_id);
15446 }
15447 }
15448 IndexKind::Sparse => {
15449 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
15450 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
15451 s.insert(&terms, row.row_id);
15452 }
15453 }
15454 }
15455 IndexKind::MinHash => {
15456 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
15457 let tokens = crate::index::token_hashes_from_bytes(b);
15458 mh.insert(&tokens, row.row_id);
15459 }
15460 }
15461 _ => {}
15462 }
15463}
15464
15465fn eval_partial_predicate(
15471 pred: &str,
15472 columns_map: &HashMap<u16, &Value>,
15473 name_to_id: &HashMap<&str, u16>,
15474) -> bool {
15475 let lower = pred.trim().to_ascii_lowercase();
15476 if let Some(rest) = lower.strip_suffix(" is not null") {
15478 let col_name = rest.trim();
15479 if let Some(col_id) = name_to_id.get(col_name) {
15480 return columns_map
15481 .get(col_id)
15482 .is_some_and(|v| !matches!(v, Value::Null));
15483 }
15484 }
15485 if let Some(rest) = lower.strip_suffix(" is null") {
15487 let col_name = rest.trim();
15488 if let Some(col_id) = name_to_id.get(col_name) {
15489 return columns_map
15490 .get(col_id)
15491 .is_none_or(|v| matches!(v, Value::Null));
15492 }
15493 }
15494 true
15497}
15498
15499#[allow(dead_code)]
15505fn bulk_index_key(
15506 column_keys: &HashMap<u16, ([u8; 32], u8)>,
15507 column_id: u16,
15508 ty: TypeId,
15509 col: &columnar::NativeColumn,
15510 i: usize,
15511) -> Option<Vec<u8>> {
15512 let encoded = columnar::encode_key_native(ty, col, i)?;
15513 {
15514 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
15515 if let Some((key, scheme)) = column_keys.get(&column_id) {
15516 return Some(match (*scheme, col) {
15517 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
15518 (_, columnar::NativeColumn::Int64 { data, .. }) => {
15519 ope_token_i64(key, data[i]).to_vec()
15520 }
15521 (_, columnar::NativeColumn::Float64 { data, .. }) => {
15522 ope_token_f64(key, data[i]).to_vec()
15523 }
15524 _ => hmac_token(key, &encoded).to_vec(),
15525 });
15526 }
15527 }
15528 Some(encoded)
15529}
15530
15531pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
15532 write_schema_with_after(dir, schema, || {})
15533}
15534
15535pub(crate) fn write_schema_durable(
15536 root: &crate::durable_file::DurableRoot,
15537 schema: &Schema,
15538) -> Result<()> {
15539 write_schema_durable_with_after(root, schema, || {})
15540}
15541
15542fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
15543where
15544 F: FnOnce(),
15545{
15546 let json = serde_json::to_string_pretty(schema)
15547 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
15548 crate::durable_file::write_atomic_with_after(
15549 &dir.join(SCHEMA_FILENAME),
15550 json.as_bytes(),
15551 after_publish,
15552 )?;
15553 Ok(())
15554}
15555
15556fn write_schema_durable_with_after<F>(
15557 root: &crate::durable_file::DurableRoot,
15558 schema: &Schema,
15559 after_publish: F,
15560) -> Result<()>
15561where
15562 F: FnOnce(),
15563{
15564 let json = serde_json::to_string_pretty(schema)
15565 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
15566 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
15567 Ok(())
15568}
15569
15570fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
15571 let mut schema_published = false;
15572 let schema_result = match table._root_guard.as_deref() {
15573 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
15574 schema_published = true;
15575 }),
15576 None => write_schema_with_after(&table.dir, &table.schema, || {
15577 schema_published = true;
15578 }),
15579 };
15580 if schema_result.is_err() && !schema_published {
15581 return schema_result;
15582 }
15583 match table.persist_manifest(table.current_epoch()) {
15584 Ok(()) => Ok(()),
15585 Err(manifest_error) => Err(match schema_result {
15586 Ok(()) => manifest_error,
15587 Err(schema_error) => MongrelError::Other(format!(
15588 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
15589 )),
15590 }),
15591 }
15592}
15593
15594fn read_schema(dir: &Path) -> Result<Schema> {
15595 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
15596 read_schema_file(file)
15597}
15598
15599fn read_schema_file(file: std::fs::File) -> Result<Schema> {
15600 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
15601 use std::io::Read;
15602
15603 let length = file.metadata()?.len();
15604 if length > MAX_SCHEMA_BYTES {
15605 return Err(MongrelError::ResourceLimitExceeded {
15606 resource: "schema bytes",
15607 requested: usize::try_from(length).unwrap_or(usize::MAX),
15608 limit: MAX_SCHEMA_BYTES as usize,
15609 });
15610 }
15611 let mut bytes = Vec::with_capacity(length as usize);
15612 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
15613 if bytes.len() as u64 != length {
15614 return Err(MongrelError::Schema(
15615 "schema length changed while reading".into(),
15616 ));
15617 }
15618 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
15619}
15620
15621fn preflight_standalone_open(
15622 dir: &Path,
15623 runs_root: Option<&crate::durable_file::DurableRoot>,
15624 idx_root: Option<&crate::durable_file::DurableRoot>,
15625 manifest: &Manifest,
15626 schema: &Schema,
15627 records: &[crate::wal::Record],
15628 kek: Option<Arc<Kek>>,
15629) -> Result<()> {
15630 crate::wal::validate_shared_transaction_framing(records)?;
15631 if manifest.schema_id > schema.schema_id
15632 || manifest.flushed_epoch > manifest.current_epoch
15633 || manifest.global_idx_epoch > manifest.current_epoch
15634 || manifest.next_row_id == u64::MAX
15635 || manifest.auto_inc_next < 0
15636 || manifest.auto_inc_next == i64::MAX
15637 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
15638 {
15639 return Err(MongrelError::InvalidArgument(
15640 "manifest counters or schema identity are invalid".into(),
15641 ));
15642 }
15643 let mut run_ids = HashSet::new();
15644 let mut maximum_row_id = None::<u64>;
15645 for run in &manifest.runs {
15646 if run.run_id >= u64::MAX as u128
15647 || !run_ids.insert(run.run_id)
15648 || run.epoch_created > manifest.current_epoch
15649 {
15650 return Err(MongrelError::InvalidArgument(
15651 "manifest contains an invalid or duplicate active run".into(),
15652 ));
15653 }
15654 let mut reader = match runs_root {
15655 Some(root) => RunReader::open_file(
15656 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
15657 schema.clone(),
15658 kek.clone(),
15659 )?,
15660 None => RunReader::open(
15661 dir.join(RUNS_DIR)
15662 .join(format!("r-{}.sr", run.run_id as u64)),
15663 schema.clone(),
15664 kek.clone(),
15665 )?,
15666 };
15667 let header = reader.header();
15668 if header.run_id != run.run_id
15669 || header.level != run.level
15670 || header.row_count != run.row_count
15671 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
15672 || header.is_uniform_epoch() && header.epoch_created != 0
15673 || header.schema_id > schema.schema_id
15674 {
15675 return Err(MongrelError::InvalidArgument(format!(
15676 "run {} differs from its manifest",
15677 run.run_id
15678 )));
15679 }
15680 if header.row_count != 0 {
15681 maximum_row_id = Some(
15682 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
15683 );
15684 }
15685 reader.validate_all_pages()?;
15686 }
15687 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
15688 return Err(MongrelError::InvalidArgument(
15689 "manifest next_row_id does not advance beyond persisted rows".into(),
15690 ));
15691 }
15692 for run in &manifest.retiring {
15693 if run.run_id >= u64::MAX as u128
15694 || run.retire_epoch > manifest.current_epoch
15695 || !run_ids.insert(run.run_id)
15696 {
15697 return Err(MongrelError::InvalidArgument(
15698 "manifest contains an invalid or duplicate retired run".into(),
15699 ));
15700 }
15701 }
15702 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
15703 match idx_root {
15704 Some(root) => {
15705 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
15706 }
15707 None => {
15708 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
15709 }
15710 }
15711
15712 let committed = records
15713 .iter()
15714 .filter_map(|record| match record.op {
15715 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
15716 _ => None,
15717 })
15718 .collect::<HashMap<_, _>>();
15719 for record in records {
15720 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
15721 continue;
15722 };
15723 match &record.op {
15724 Op::Put { table_id, rows } => {
15725 if *table_id != manifest.table_id {
15726 return Err(MongrelError::CorruptWal {
15727 offset: record.seq.0,
15728 reason: format!(
15729 "private WAL record references table {table_id}, expected {}",
15730 manifest.table_id
15731 ),
15732 });
15733 }
15734 let rows: Vec<Row> =
15735 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
15736 offset: record.seq.0,
15737 reason: format!("committed Put payload could not be decoded: {error}"),
15738 })?;
15739 for row in rows {
15740 if row.deleted || row.row_id.0 == u64::MAX {
15741 return Err(MongrelError::CorruptWal {
15742 offset: record.seq.0,
15743 reason: "committed Put contains an invalid row identity".into(),
15744 });
15745 }
15746 let cells = row.columns.into_iter().collect::<Vec<_>>();
15747 schema
15748 .validate_values(&cells)
15749 .map_err(|error| MongrelError::CorruptWal {
15750 offset: record.seq.0,
15751 reason: format!("committed Put violates table schema: {error}"),
15752 })?;
15753 if schema.auto_increment_column().is_some_and(|column| {
15754 matches!(
15755 cells.iter().find(|(id, _)| *id == column.id),
15756 Some((_, Value::Int64(value))) if *value == i64::MAX
15757 )
15758 }) {
15759 return Err(MongrelError::CorruptWal {
15760 offset: record.seq.0,
15761 reason: "committed Put exhausts AUTO_INCREMENT".into(),
15762 });
15763 }
15764 }
15765 }
15766 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
15767 if *table_id != manifest.table_id =>
15768 {
15769 return Err(MongrelError::CorruptWal {
15770 offset: record.seq.0,
15771 reason: format!(
15772 "private WAL record references table {table_id}, expected {}",
15773 manifest.table_id
15774 ),
15775 });
15776 }
15777 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
15778 return Err(MongrelError::CorruptWal {
15779 offset: record.seq.0,
15780 reason: "private WAL contains shared spilled-run metadata".into(),
15781 });
15782 }
15783 _ => {}
15784 }
15785 }
15786 Ok(())
15787}
15788
15789fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
15790 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
15791}
15792
15793fn wal_segment_number(path: &Path) -> Option<u64> {
15794 path.file_stem()
15795 .and_then(|stem| stem.to_str())
15796 .and_then(|stem| stem.strip_prefix("seg-"))
15797 .and_then(|number| number.parse().ok())
15798}
15799
15800fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
15801 let n = list_wal_numbers(wal_dir)?;
15802 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
15803}
15804
15805fn next_wal_number(wal_dir: &Path) -> Result<u32> {
15806 list_wal_numbers(wal_dir)?
15807 .map(|maximum| {
15808 maximum
15809 .checked_add(1)
15810 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
15811 })
15812 .unwrap_or(Ok(0))
15813}
15814
15815fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
15816 let mut max_n = None;
15817 let entries = match std::fs::read_dir(wal_dir) {
15818 Ok(entries) => entries,
15819 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
15820 Err(error) => return Err(error.into()),
15821 };
15822 for entry in entries {
15823 let entry = entry?;
15824 let fname = entry.file_name();
15825 let Some(s) = fname.to_str() else {
15826 continue;
15827 };
15828 let Some(stripped) = s.strip_prefix("seg-") else {
15829 continue;
15830 };
15831 let Some(number) = stripped.strip_suffix(".wal") else {
15832 return Err(MongrelError::CorruptWal {
15833 offset: 0,
15834 reason: format!("malformed WAL segment name {s:?}"),
15835 });
15836 };
15837 let n = number
15838 .parse::<u32>()
15839 .map_err(|_| MongrelError::CorruptWal {
15840 offset: 0,
15841 reason: format!("malformed WAL segment name {s:?}"),
15842 })?;
15843 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
15844 return Err(MongrelError::CorruptWal {
15845 offset: n as u64,
15846 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
15847 });
15848 }
15849 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
15850 }
15851 Ok(max_n)
15852}
15853
15854#[cfg(test)]
15879mod rem001_multi_run_hlc_authority {
15880 use super::*;
15881 use crate::epoch::{Epoch, Snapshot, VersionStamp};
15882 use crate::memtable::{Row, Value};
15883 use crate::query::{Condition, Query};
15884 use crate::schema::{ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
15885 use mongreldb_types::hlc::HlcTimestamp;
15886 use tempfile::tempdir;
15887
15888 fn hlc(physical_micros: u64) -> HlcTimestamp {
15889 HlcTimestamp {
15890 physical_micros,
15891 logical: 0,
15892 node_tiebreaker: 1,
15893 }
15894 }
15895
15896 fn pk_schema() -> Schema {
15897 Schema {
15898 schema_id: 1,
15899 columns: vec![
15900 ColumnDef {
15901 id: 1,
15902 name: "id".into(),
15903 ty: TypeId::Int64,
15904 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
15905 default_value: None,
15906 embedding_source: None,
15907 },
15908 ColumnDef {
15909 id: 2,
15910 name: "name".into(),
15911 ty: TypeId::Bytes,
15912 flags: ColumnFlags::empty(),
15913 default_value: None,
15914 embedding_source: None,
15915 },
15916 ],
15917 indexes: Vec::new(),
15918 colocation: vec![],
15919 constraints: Default::default(),
15920 clustered: false,
15921 }
15922 }
15923
15924 fn pk_schema_with_bitmap() -> Schema {
15925 Schema {
15926 schema_id: 1,
15927 columns: vec![
15928 ColumnDef {
15929 id: 1,
15930 name: "id".into(),
15931 ty: TypeId::Int64,
15932 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
15933 default_value: None,
15934 embedding_source: None,
15935 },
15936 ColumnDef {
15937 id: 2,
15938 name: "name".into(),
15939 ty: TypeId::Bytes,
15940 flags: ColumnFlags::empty(),
15941 default_value: None,
15942 embedding_source: None,
15943 },
15944 ],
15945 indexes: vec![IndexDef {
15948 name: "name_bitmap".into(),
15949 column_id: 2,
15950 kind: IndexKind::Bitmap,
15951 predicate: None,
15952 options: Default::default(),
15953 }],
15954 colocation: vec![],
15955 constraints: Default::default(),
15956 clustered: false,
15957 }
15958 }
15959
15960 fn encode_put(table_id: u64, row: &Row) -> Vec<u8> {
15961 use crate::database::StagedTxnWrite;
15962 bincode::serialize(&StagedTxnWrite::Put {
15963 table_id,
15964 rows: bincode::serialize(&vec![row.clone()]).expect("encode row"),
15965 })
15966 .expect("encode payload")
15967 }
15968
15969 fn build_inverted_database(dir: &std::path::Path, schema: Schema) -> crate::Database {
15983 let db = crate::Database::create(dir).unwrap();
15984 db.create_table("t", schema).unwrap();
15985 let table_id = db.table_id("t").unwrap();
15986 let rid = RowId(7);
15987
15988 let newer = Row::new_with_hlc(rid, Epoch(1), hlc(500))
15989 .with_column(1, Value::Int64(7))
15990 .with_column(2, Value::Bytes(b"newer-by-HLC".to_vec()));
15991 let older = Row::new_with_hlc(rid, Epoch(2), hlc(400))
15992 .with_column(1, Value::Int64(7))
15993 .with_column(2, Value::Bytes(b"older-by-HLC".to_vec()));
15994
15995 db.apply_staged_txn_writes(1, &[encode_put(table_id, &newer)], hlc(500))
15996 .expect("first apply");
15997 {
15998 let table = db.table("t").unwrap();
15999 table.lock().force_flush().expect("flush run A");
16000 }
16001
16002 db.apply_staged_txn_writes(2, &[encode_put(table_id, &older)], hlc(400))
16003 .expect("second apply");
16004 {
16005 let table = db.table("t").unwrap();
16006 table.lock().force_flush().expect("flush run B");
16007 }
16008
16009 db
16010 }
16011
16012 fn build_inverted_ttl_database(dir: &std::path::Path) -> crate::Database {
16017 let db = crate::Database::create(dir).unwrap();
16018 db.create_table("t", ttl_schema()).unwrap();
16019 let table_id = db.table_id("t").unwrap();
16020 let rid = RowId(7);
16021 let read_time_ns: i64 = 1_000_000_000;
16022 let newer = Row::new_with_hlc(rid, Epoch(1), hlc(500))
16023 .with_column(1, Value::Int64(7))
16024 .with_column(2, Value::Bytes(b"newer-by-HLC".to_vec()))
16025 .with_column(3, Value::Int64(read_time_ns + 60 * 1_000_000_000));
16026 let older = Row::new_with_hlc(rid, Epoch(2), hlc(400))
16027 .with_column(1, Value::Int64(7))
16028 .with_column(2, Value::Bytes(b"older-by-HLC".to_vec()))
16029 .with_column(3, Value::Int64(read_time_ns - 60 * 1_000_000_000));
16030
16031 db.apply_staged_txn_writes(1, &[encode_put(table_id, &newer)], hlc(500))
16032 .expect("first apply");
16033 {
16034 let table = db.table("t").unwrap();
16035 table.lock().force_flush().expect("flush run A");
16036 }
16037 db.apply_staged_txn_writes(2, &[encode_put(table_id, &older)], hlc(400))
16038 .expect("second apply");
16039 {
16040 let table = db.table("t").unwrap();
16041 table.lock().force_flush().expect("flush run B");
16042 }
16043
16044 db
16045 }
16046
16047 fn ttl_schema() -> Schema {
16048 Schema {
16049 schema_id: 1,
16050 columns: vec![
16051 ColumnDef {
16052 id: 1,
16053 name: "id".into(),
16054 ty: TypeId::Int64,
16055 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
16056 default_value: None,
16057 embedding_source: None,
16058 },
16059 ColumnDef {
16060 id: 2,
16061 name: "name".into(),
16062 ty: TypeId::Bytes,
16063 flags: ColumnFlags::empty(),
16064 default_value: None,
16065 embedding_source: None,
16066 },
16067 ColumnDef {
16068 id: 3,
16069 name: "ttl_ts".into(),
16070 ty: TypeId::Int64,
16071 flags: ColumnFlags::empty(),
16072 default_value: None,
16073 embedding_source: None,
16074 },
16075 ],
16076 indexes: Vec::new(),
16077 colocation: vec![],
16078 constraints: Default::default(),
16079 clustered: false,
16080 }
16081 }
16082
16083 #[test]
16084 fn multi_run_rows_for_rids_uses_hlc_winner_when_epoch_is_inverted() {
16085 let dir = tempdir().unwrap();
16086 let db = build_inverted_database(dir.path(), pk_schema());
16087 let table = db.table("t").unwrap();
16088 let table = table.lock();
16089 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16093 let rows = table
16094 .rows_for_rids(&[7], snap)
16095 .expect("rows_for_rids across two runs");
16096 assert_eq!(rows.len(), 1, "one row materialised across two runs");
16097 assert_eq!(
16098 rows[0].columns.get(&2),
16099 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
16100 "multi-run rows_for_rids must honor HLC authority"
16101 );
16102 }
16103
16104 #[test]
16105 fn multi_run_projected_column_uses_hlc_winner_when_epoch_is_inverted() {
16106 let dir = tempdir().unwrap();
16107 let db = build_inverted_database(dir.path(), pk_schema());
16108 let table = db.table("t").unwrap();
16109 let table = table.lock();
16110 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16111 let projected = table
16112 .values_for_rids_at(&[7], 2, snap, i64::MAX)
16113 .expect("values_for_rids_at across two runs");
16114 assert_eq!(projected.len(), 1);
16115 assert_eq!(
16116 projected[0].1,
16117 Value::Bytes(b"newer-by-HLC".to_vec()),
16118 "multi-run projection must pick the HLC-newer column"
16119 );
16120 }
16121
16122 #[test]
16123 fn ann_candidate_eligibility_uses_hlc_winner_across_runs() {
16124 let dir = tempdir().unwrap();
16125 let db = build_inverted_database(dir.path(), pk_schema());
16126 let table = db.table("t").unwrap();
16127 let table = table.lock();
16128 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16129 let eligible = table
16132 .eligible_candidate_ids(&[RowId(7)], 2, snap, None)
16133 .expect("eligibility across runs");
16134 assert!(
16135 eligible.contains(&RowId(7)),
16136 "HLC-newer rid must remain eligible across runs"
16137 );
16138 }
16139
16140 #[test]
16141 fn ttl_inverted_run_is_admitted_via_values_for_rids_at() {
16142 let dir = tempdir().unwrap();
16143 let db = build_inverted_ttl_database(dir.path());
16144 let table = db.table("t").unwrap();
16145 let table = table.lock();
16146 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16147 let now_ns: i64 = 1_000_000_000;
16148 let projected = table
16152 .values_for_rids_at(&[7], 2, snap, now_ns)
16153 .expect("ttl-aware projection");
16154 assert_eq!(projected.len(), 1, "TTL must admit the HLC-newer row");
16155 assert_eq!(
16156 projected[0].1,
16157 Value::Bytes(b"newer-by-HLC".to_vec()),
16158 "TTL must be read from the HLC-newer run; older run's timestamp is expired"
16159 );
16160 }
16161
16162 #[test]
16163 fn hot_candidate_inspection_uses_full_snapshot_for_run_versions() {
16164 let dir = tempdir().unwrap();
16169 let db = build_inverted_database(dir.path(), pk_schema());
16170 let table = db.table("t").unwrap();
16171 let table = table.lock();
16172 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16173 let row = table
16174 .get(RowId(7), snap)
16175 .expect("Table::get across two runs");
16176 assert_eq!(
16177 row.columns.get(&2),
16178 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
16179 "Table::get must surface the HLC-newer winner"
16180 );
16181 }
16182
16183 #[test]
16184 fn learned_range_rebuild_indexes_only_hlc_visible_winners() {
16185 let dir = tempdir().unwrap();
16192 let db = build_inverted_database(dir.path(), pk_schema());
16193 let handle = db.table("t").unwrap();
16194 let mut table = handle.lock();
16195 table.rebuild_indexes_from_runs().expect("rebuild");
16196 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16197 let rows = table
16198 .rows_for_rids(&[7], snap)
16199 .expect("post-rebuild rows_for_rids");
16200 assert_eq!(rows.len(), 1);
16201 assert_eq!(
16202 rows[0].columns.get(&2),
16203 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
16204 "post-rebuild row must be the HLC-newer winner"
16205 );
16206 }
16207
16208 #[test]
16209 fn rebuild_indexes_uses_global_hlc_winners_not_run_iteration_order() {
16210 let dir = tempdir().unwrap();
16223 let db = build_inverted_database(dir.path(), pk_schema_with_bitmap());
16224 let handle = db.table("t").unwrap();
16225 let mut table = handle.lock();
16226 table.rebuild_indexes_from_runs().expect("rebuild");
16227 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
16228 let row = table.get(RowId(7), snap).expect("post-rebuild Table::get");
16230 assert_eq!(
16231 row.columns.get(&2),
16232 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
16233 "PK lookup must resolve to the HLC-newer winner"
16234 );
16235 let query_newer = Query::new().and(Condition::BitmapEq {
16239 column_id: 2,
16240 value: b"newer-by-HLC".to_vec(),
16241 });
16242 let rows_newer = table.query(&query_newer).expect("bitmap eq newer");
16243 let rids_newer: Vec<u64> = rows_newer.iter().map(|r| r.row_id.0).collect();
16244 assert!(
16245 rids_newer.contains(&7),
16246 "bitmap must include the HLC-newer value (got {rids_newer:?})"
16247 );
16248 let query_older = Query::new().and(Condition::BitmapEq {
16249 column_id: 2,
16250 value: b"older-by-HLC".to_vec(),
16251 });
16252 let rows_older = table.query(&query_older).expect("bitmap eq older");
16253 let rids_older: Vec<u64> = rows_older.iter().map(|r| r.row_id.0).collect();
16254 assert!(
16255 !rids_older.contains(&7),
16256 "bitmap must NOT include the HLC-older value (got {rids_older:?})"
16257 );
16258 }
16259
16260 #[test]
16261 fn version_stamp_is_newer_than_matches_snapshot_rule() {
16262 let early = hlc(100);
16265 let late = hlc(200);
16266 let newer = VersionStamp {
16267 epoch: Epoch(9),
16268 commit_ts: Some(late),
16269 };
16270 let older = VersionStamp {
16271 epoch: Epoch(50),
16272 commit_ts: Some(early),
16273 };
16274 assert!(newer.is_newer_than(older));
16275 assert!(!older.is_newer_than(newer));
16276 let a = VersionStamp {
16277 epoch: Epoch(5),
16278 commit_ts: None,
16279 };
16280 let b = VersionStamp {
16281 epoch: Epoch(4),
16282 commit_ts: None,
16283 };
16284 assert!(a.is_newer_than(b));
16285 }
16286}