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::{DrainOutcome, PersistContext, PersistentCacheIdentity, FRAME_MAGIC};
24use crate::row_id_set::RowIdSet;
25use crate::rowid::{RowId, RowIdAllocator};
26use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
27use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
28use crate::txn::{GroupCommit, OwnedRow};
29use crate::wal::{Op, SharedWal, Wal};
30use crate::{MongrelError, Result};
31use arc_swap::ArcSwap;
32use mongreldb_types::hlc::HlcTimestamp;
33use std::borrow::Cow;
34use std::cmp::Reverse;
35use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
36use std::path::{Path, PathBuf};
37use std::sync::atomic::AtomicBool;
38use std::sync::Arc;
39use std::time::Instant;
40use zeroize::Zeroizing;
41
42pub const WAL_DIR: &str = "_wal";
43pub const RUNS_DIR: &str = "_runs";
44pub const CACHE_DIR: &str = "_cache";
45pub const META_DIR: &str = "_meta";
46pub const RCACHE_DIR: &str = "_rcache";
47pub const KEYS_FILENAME: &str = "keys";
48pub const SCHEMA_FILENAME: &str = "schema.json";
49
50fn derive_next_run_id(
51 dir: &Path,
52 runs_root: Option<&crate::durable_file::DurableRoot>,
53 active: &[RunRef],
54 retiring: &[crate::manifest::RetiredRun],
55) -> Result<u64> {
56 let mut maximum = 0_u64;
57 for run_id in active
58 .iter()
59 .map(|run| run.run_id)
60 .chain(retiring.iter().map(|run| run.run_id))
61 {
62 let run_id = u64::try_from(run_id)
63 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
64 maximum = maximum.max(run_id);
65 }
66 let names = match runs_root {
67 Some(root) => root.list_regular_files(".")?,
68 None => std::fs::read_dir(dir.join(RUNS_DIR))?
69 .map(|entry| entry.map(|entry| entry.file_name()))
70 .collect::<std::io::Result<Vec<_>>>()?,
71 };
72 for name in names {
73 let Some(name) = name.to_str() else {
74 continue;
75 };
76 let Some(digits) = name
77 .strip_prefix("r-")
78 .and_then(|name| name.strip_suffix(".sr"))
79 else {
80 continue;
81 };
82 let Ok(run_id) = digits.parse::<u64>() else {
83 continue;
84 };
85 if name == format!("r-{run_id}.sr") {
86 maximum = maximum.max(run_id);
87 }
88 }
89 maximum
90 .checked_add(1)
91 .map(|next| next.max(1))
92 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
93}
94
95enum ControlledVisibleCandidate<'a> {
96 Memory(Row),
97 Memtable(RowId, Epoch, Cow<'a, Row>),
103 MutableRun(RowId, Epoch, &'a Row),
105 Run(RunVisibleVersion),
106}
107
108impl<'a> ControlledVisibleCandidate<'a> {
109 fn row_id(&self) -> RowId {
110 match self {
111 Self::Memory(row) => row.row_id,
112 Self::Memtable(rid, _, _) => *rid,
113 Self::MutableRun(rid, _, _) => *rid,
114 Self::Run(version) => version.row_id,
115 }
116 }
117
118 fn committed_epoch(&self) -> Epoch {
119 match self {
120 Self::Memory(row) => row.committed_epoch,
121 Self::Memtable(_, epoch, _) => *epoch,
122 Self::MutableRun(_, epoch, _) => *epoch,
123 Self::Run(version) => version.committed_epoch,
124 }
125 }
126
127 fn deleted(&self) -> bool {
128 match self {
129 Self::Memory(row) => row.deleted,
130 Self::Memtable(_, _, row) => row.deleted,
131 Self::MutableRun(_, _, row) => row.deleted,
132 Self::Run(version) => version.deleted,
133 }
134 }
135
136 fn commit_ts(&self) -> Option<mongreldb_types::hlc::HlcTimestamp> {
137 match self {
138 Self::Memory(row) => row.commit_ts,
139 Self::Memtable(_, _, row) => row.commit_ts,
140 Self::MutableRun(_, _, row) => row.commit_ts,
141 Self::Run(version) => version.commit_ts,
144 }
145 }
146}
147
148enum ControlledVisibleCursor<'a> {
149 Memory(std::vec::IntoIter<Row>),
151 MemoryStreaming {
155 active: std::vec::IntoIter<Row>,
156 rest: std::collections::btree_map::IntoValues<RowId, Row>,
157 batch_cap: usize,
158 #[allow(dead_code)]
160 peak_active: usize,
161 #[allow(dead_code)]
163 total: usize,
164 },
165 Memtable(MemtableVisibleVersionCursor<'a>),
169 MutableRun(MutableRunVisibleVersionCursor<'a>),
171 Run(Box<RunVisibleVersionCursor>),
172 #[cfg(test)]
173 Synthetic {
174 next: u64,
175 end: u64,
176 },
177}
178
179#[allow(dead_code)] const CONTROLLED_HOT_BATCH: usize = 256;
182
183struct ControlledVisibleSource<'a> {
184 cursor: ControlledVisibleCursor<'a>,
185 current: Option<ControlledVisibleCandidate<'a>>,
186}
187
188impl<'a> ControlledVisibleSource<'a> {
189 #[cfg(test)]
191 fn memory(rows: Vec<Row>) -> Self {
192 Self {
193 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
194 current: None,
195 }
196 }
197
198 #[allow(dead_code)] fn memory_from_map(map: BTreeMap<RowId, Row>) -> Self {
203 let total = map.len();
204 let mut values = map.into_values();
205 if total > CONTROLLED_HOT_BATCH {
206 let active: Vec<Row> = values.by_ref().take(CONTROLLED_HOT_BATCH).collect();
207 let peak = active.len();
208 Self {
209 cursor: ControlledVisibleCursor::MemoryStreaming {
210 active: active.into_iter(),
211 rest: values,
212 batch_cap: CONTROLLED_HOT_BATCH,
213 peak_active: peak,
214 total,
215 },
216 current: None,
217 }
218 } else {
219 let active: Vec<Row> = values.collect();
220 Self {
221 cursor: ControlledVisibleCursor::Memory(active.into_iter()),
222 current: None,
223 }
224 }
225 }
226
227 #[allow(dead_code)] fn memtable_cursor(cursor: MemtableVisibleVersionCursor<'a>) -> Self {
232 Self {
233 cursor: ControlledVisibleCursor::Memtable(cursor),
234 current: None,
235 }
236 }
237
238 #[allow(dead_code)] fn mutable_run_cursor(cursor: MutableRunVisibleVersionCursor<'a>) -> Self {
243 Self {
244 cursor: ControlledVisibleCursor::MutableRun(cursor),
245 current: None,
246 }
247 }
248
249 fn run(cursor: RunVisibleVersionCursor) -> Self {
250 Self {
251 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
252 current: None,
253 }
254 }
255
256 #[cfg(test)]
257 fn synthetic(end: u64) -> Self {
258 Self {
259 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
260 current: None,
261 }
262 }
263
264 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
265 self.current = match &mut self.cursor {
266 ControlledVisibleCursor::Memory(rows) => {
267 rows.next().map(ControlledVisibleCandidate::Memory)
268 }
269 ControlledVisibleCursor::MemoryStreaming {
270 active,
271 rest,
272 batch_cap,
273 peak_active,
274 total: _,
275 } => {
276 if let Some(row) = active.next() {
277 Some(ControlledVisibleCandidate::Memory(row))
278 } else {
279 control.checkpoint()?;
280 let next_batch: Vec<Row> = rest.by_ref().take(*batch_cap).collect();
281 if next_batch.is_empty() {
282 None
283 } else {
284 *peak_active = (*peak_active).max(next_batch.len());
285 *active = next_batch.into_iter();
286 active.next().map(ControlledVisibleCandidate::Memory)
287 }
288 }
289 }
290 ControlledVisibleCursor::Memtable(iter) => {
291 let prev_peak = iter.peak_examined;
292 let next = iter
293 .next()
294 .map(|(rid, epoch, row)| ControlledVisibleCandidate::Memtable(rid, epoch, row));
295 if let Some(peak) = iter.peak_examined.checked_sub(prev_peak) {
296 crate::trace::QueryTrace::record(|t| {
297 t.controlled_scan_peak_same_row_versions = t
298 .controlled_scan_peak_same_row_versions
299 .saturating_add(peak);
300 });
301 }
302 next
303 }
304 ControlledVisibleCursor::MutableRun(iter) => {
305 let prev_peak = iter.peak_examined;
306 let next = iter.next().map(|(rid, epoch, row)| {
307 ControlledVisibleCandidate::MutableRun(rid, epoch, row)
308 });
309 if let Some(peak) = iter.peak_examined.checked_sub(prev_peak) {
310 crate::trace::QueryTrace::record(|t| {
311 t.controlled_scan_peak_same_row_versions = t
312 .controlled_scan_peak_same_row_versions
313 .saturating_add(peak);
314 });
315 }
316 next
317 }
318 ControlledVisibleCursor::Run(cursor) => cursor
319 .next_visible_version(control)?
320 .map(ControlledVisibleCandidate::Run),
321 #[cfg(test)]
322 ControlledVisibleCursor::Synthetic { next, end } => {
323 if *next > *end {
324 None
325 } else {
326 let row = Row::new(RowId(*next), Epoch(1));
327 *next += 1;
328 Some(ControlledVisibleCandidate::Memory(row))
329 }
330 }
331 };
332 Ok(())
333 }
334
335 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate<'a>> {
336 let current = self.current.take().ok_or_else(|| {
337 MongrelError::Other("controlled visible source was not primed".into())
338 })?;
339 self.advance(control)?;
340 Ok(current)
341 }
342
343 fn materialize(
344 &mut self,
345 candidate: ControlledVisibleCandidate<'a>,
346 control: &crate::ExecutionControl,
347 ) -> Result<Row> {
348 match candidate {
349 ControlledVisibleCandidate::Memory(row) => Ok(row),
350 ControlledVisibleCandidate::Memtable(_, _, row) => Ok(row.into_owned()),
351 ControlledVisibleCandidate::MutableRun(_, _, row) => Ok(row.clone()),
352 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
353 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
354 _ => Err(MongrelError::Other(
355 "run candidate escaped its controlled cursor".into(),
356 )),
357 },
358 }
359 }
360
361 #[cfg(test)]
362 fn is_streaming_memory(&self) -> bool {
363 matches!(self.cursor, ControlledVisibleCursor::MemoryStreaming { .. })
364 }
365
366 #[cfg(test)]
367 fn streaming_peak_active(&self) -> Option<usize> {
368 match &self.cursor {
369 ControlledVisibleCursor::MemoryStreaming { peak_active, .. } => Some(*peak_active),
370 _ => None,
371 }
372 }
373
374 #[cfg(test)]
375 fn streaming_total(&self) -> Option<usize> {
376 match &self.cursor {
377 ControlledVisibleCursor::MemoryStreaming { total, .. } => Some(*total),
378 _ => None,
379 }
380 }
381}
382
383fn merge_controlled_visible_sources<'a>(
384 sources: &mut [ControlledVisibleSource<'a>],
385 control: &crate::ExecutionControl,
386 mut expired: impl FnMut(&Row) -> bool,
387 mut visit: impl FnMut(Row) -> Result<()>,
388) -> Result<()> {
389 let mut heap = BinaryHeap::new();
390 for (source_index, source) in sources.iter_mut().enumerate() {
391 source.advance(control)?;
392 if let Some(candidate) = &source.current {
393 heap.push(Reverse((candidate.row_id(), source_index)));
394 }
395 }
396 let mut merged = 0_usize;
397 let mut peak_buffer_rows: usize = heap.len();
398 let mut peak_same_row_versions: usize = 0;
399 let mut cancel_started: Option<std::time::Instant> = None;
400 while let Some(Reverse((row_id, source_index))) = heap.pop() {
401 if merged.is_multiple_of(256) {
402 control.checkpoint()?;
403 }
404 merged += 1;
405 let mut best_source = source_index;
406 let mut best = sources[source_index].pop(control)?;
407 if let Some(next) = &sources[source_index].current {
408 heap.push(Reverse((next.row_id(), source_index)));
409 }
410 let mut same_row_versions: usize = 1;
411 while heap
412 .peek()
413 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
414 {
415 control.checkpoint()?;
416 let Some(Reverse((_, source_index))) = heap.pop() else {
417 break;
418 };
419 let candidate = sources[source_index].pop(control)?;
420 same_row_versions += 1;
421 if Snapshot::version_is_newer(
425 candidate.committed_epoch(),
426 candidate.commit_ts(),
427 best.committed_epoch(),
428 best.commit_ts(),
429 ) {
430 best = candidate;
431 best_source = source_index;
432 }
433 if let Some(next) = &sources[source_index].current {
434 heap.push(Reverse((next.row_id(), source_index)));
435 }
436 }
437 peak_same_row_versions = peak_same_row_versions.max(same_row_versions);
438 peak_buffer_rows = peak_buffer_rows.max(heap.len());
439 if best.deleted() {
440 continue;
441 }
442 let row = sources[best_source].materialize(best, control)?;
443 if !expired(&row) {
444 visit(row)?;
445 }
446 if cancel_started.is_none() && control.is_cancelled() {
447 cancel_started = Some(std::time::Instant::now());
448 }
449 }
450 crate::trace::QueryTrace::record(|t| {
451 t.controlled_scan_source_refills = t.controlled_scan_source_refills.saturating_add(0); t.controlled_scan_peak_source_buffer_rows = t
453 .controlled_scan_peak_source_buffer_rows
454 .saturating_add(peak_buffer_rows);
455 t.controlled_scan_peak_same_row_versions = t
456 .controlled_scan_peak_same_row_versions
457 .saturating_add(peak_same_row_versions);
458 });
459 control.checkpoint()?;
460 Ok(())
461}
462
463#[cfg(test)]
464mod controlled_visible_cursor_tests {
465 use super::*;
466
467 #[test]
468 fn streams_more_than_one_million_rows_without_a_source_cap() {
469 let control = crate::ExecutionControl::new(None);
470 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
471 let mut count = 0_u64;
472 let mut last = 0_u64;
473 merge_controlled_visible_sources(
474 &mut sources,
475 &control,
476 |_| false,
477 |row| {
478 count += 1;
479 assert!(row.row_id.0 > last);
480 last = row.row_id.0;
481 Ok(())
482 },
483 )
484 .unwrap();
485 assert_eq!(count, 1_000_001);
486 assert_eq!(last, 1_000_001);
487 }
488
489 #[test]
490 fn merge_orders_rows_and_honors_newest_tombstones() {
491 let control = crate::ExecutionControl::new(None);
492 let older = vec![
493 Row::new(RowId(1), Epoch(1)),
494 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
495 Row::new(RowId(4), Epoch(1)),
496 ];
497 let mut deleted = Row::new(RowId(1), Epoch(2));
498 deleted.deleted = true;
499 let newer = vec![
500 deleted,
501 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
502 Row::new(RowId(3), Epoch(2)),
503 ];
504 let mut sources = vec![
505 ControlledVisibleSource::memory(older),
506 ControlledVisibleSource::memory(newer),
507 ];
508 let mut rows = Vec::new();
509 merge_controlled_visible_sources(
510 &mut sources,
511 &control,
512 |_| false,
513 |row| {
514 rows.push(row);
515 Ok(())
516 },
517 )
518 .unwrap();
519 assert_eq!(
520 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
521 vec![2, 3, 4]
522 );
523 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
524 }
525
526 #[test]
527 fn controlled_merge_uses_hlc_authority_when_epochs_inverted() {
528 use mongreldb_types::hlc::HlcTimestamp;
529 let hlc_old = HlcTimestamp {
530 physical_micros: 100,
531 logical: 0,
532 node_tiebreaker: 1,
533 };
534 let hlc_new = HlcTimestamp {
535 physical_micros: 200,
536 logical: 0,
537 node_tiebreaker: 1,
538 };
539 let a =
542 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
543 let b =
544 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
545 let control = crate::ExecutionControl::new(None);
546 let mut sources = vec![
547 ControlledVisibleSource::memory(a),
548 ControlledVisibleSource::memory(b),
549 ];
550 let mut rows = Vec::new();
551 merge_controlled_visible_sources(
552 &mut sources,
553 &control,
554 |_| false,
555 |row| {
556 rows.push(row);
557 Ok(())
558 },
559 )
560 .unwrap();
561 assert_eq!(rows.len(), 1);
562 assert_eq!(rows[0].row_id.0, 1);
563 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(11)));
565 assert_eq!(rows[0].commit_ts, Some(hlc_new));
566 }
567
568 #[test]
569 fn controlled_merge_epoch_wins_when_one_side_lacks_hlc() {
570 use mongreldb_types::hlc::HlcTimestamp;
571 let hlc_old = HlcTimestamp {
572 physical_micros: 100,
573 logical: 0,
574 node_tiebreaker: 1,
575 };
576 let a =
579 vec![Row::new_with_hlc(RowId(1), Epoch(10), hlc_old).with_column(1, Value::Int64(10))];
580 let b = vec![Row::new(RowId(1), Epoch(5)).with_column(1, Value::Int64(5))];
581 let control = crate::ExecutionControl::new(None);
582 let mut sources = vec![
583 ControlledVisibleSource::memory(a),
584 ControlledVisibleSource::memory(b),
585 ];
586 let mut rows = Vec::new();
587 merge_controlled_visible_sources(
588 &mut sources,
589 &control,
590 |_| false,
591 |row| {
592 rows.push(row);
593 Ok(())
594 },
595 )
596 .unwrap();
597 assert_eq!(rows.len(), 1);
598 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(10)));
599 }
600
601 #[test]
604 fn controlled_merge_memory_vs_run_uses_hlc_when_run_stamped() {
605 use crate::sorted_run::{RunReader, RunWriter};
606 use mongreldb_types::hlc::HlcTimestamp;
607 use tempfile::tempdir;
608
609 let hlc_old = HlcTimestamp {
610 physical_micros: 100,
611 logical: 0,
612 node_tiebreaker: 1,
613 };
614 let hlc_new = HlcTimestamp {
615 physical_micros: 200,
616 logical: 0,
617 node_tiebreaker: 1,
618 };
619 let schema = Schema {
620 schema_id: 1,
621 columns: vec![ColumnDef {
622 id: 1,
623 name: "v".into(),
624 ty: TypeId::Int64,
625 flags: ColumnFlags::empty(),
626 default_value: None,
627 embedding_source: None,
628 }],
629 indexes: vec![],
630 colocation: vec![],
631 constraints: Default::default(),
632 clustered: false,
633 };
634 let dir = tempdir().unwrap();
635 let path = dir.path().join("r-hlc.sr");
636 let run_rows =
638 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
639 RunWriter::new(&schema, 1, Epoch(50), 0)
640 .write(&path, &run_rows)
641 .unwrap();
642 let reader = RunReader::open(&path, schema, None).unwrap();
643 assert!(reader.has_column(crate::sorted_run::SYS_COMMIT_TS));
644
645 let mem =
647 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
648 let control = crate::ExecutionControl::new(None);
649 let mut sources = vec![
650 ControlledVisibleSource::memory(mem),
651 ControlledVisibleSource::run(
652 reader.into_visible_version_cursor(Epoch(u64::MAX)).unwrap(),
653 ),
654 ];
655 let mut rows = Vec::new();
656 merge_controlled_visible_sources(
657 &mut sources,
658 &control,
659 |_| false,
660 |row| {
661 rows.push(row);
662 Ok(())
663 },
664 )
665 .unwrap();
666 assert_eq!(rows.len(), 1);
667 assert_eq!(
668 rows[0].columns.get(&1),
669 Some(&Value::Int64(11)),
670 "HLC-newer memory version must beat epoch-newer run"
671 );
672 assert_eq!(rows[0].commit_ts, Some(hlc_new));
673 }
674
675 #[test]
676 fn controlled_memory_source_streams_large_overlays_without_full_active_vec() {
677 let mut map = BTreeMap::new();
678 for i in 1..=500u64 {
679 map.insert(
680 RowId(i),
681 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
682 );
683 }
684 let source = ControlledVisibleSource::memory_from_map(map);
685 assert!(
686 source.is_streaming_memory(),
687 "overlays larger than CONTROLLED_HOT_BATCH must use streaming cursor"
688 );
689 assert_eq!(source.streaming_total(), Some(500));
690 assert!(
692 source.streaming_peak_active().unwrap_or(usize::MAX) <= CONTROLLED_HOT_BATCH,
693 "peak active buffer must be <= CONTROLLED_HOT_BATCH"
694 );
695
696 let control = crate::ExecutionControl::new(None);
698 let mut sources = vec![ControlledVisibleSource::memory_from_map({
699 let mut m = BTreeMap::new();
700 for i in 1..=500u64 {
701 m.insert(
702 RowId(i),
703 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
704 );
705 }
706 m
707 })];
708 let mut n = 0usize;
709 merge_controlled_visible_sources(
710 &mut sources,
711 &control,
712 |_| false,
713 |_| {
714 n += 1;
715 Ok(())
716 },
717 )
718 .unwrap();
719 assert_eq!(n, 500);
720 assert!(
721 sources[0].streaming_peak_active().unwrap_or(0) <= CONTROLLED_HOT_BATCH,
722 "after full drain peak active still <= batch cap"
723 );
724 }
725}
726
727fn iso_now_bytes() -> Vec<u8> {
730 let secs = std::time::SystemTime::now()
731 .duration_since(std::time::UNIX_EPOCH)
732 .map(|d| d.as_secs() as i64)
733 .unwrap_or(0);
734 let days = secs.div_euclid(86_400);
735 let rem = secs.rem_euclid(86_400);
736 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
737 let (year, month, day) = civil_from_days(days);
738 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
739}
740
741pub(crate) fn unix_nanos_now() -> i64 {
742 std::time::SystemTime::now()
743 .duration_since(std::time::UNIX_EPOCH)
744 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
745 .unwrap_or(0)
746}
747
748fn ann_candidate_cap(
749 index_len: usize,
750 context: Option<&crate::query::AiExecutionContext>,
751) -> usize {
752 index_len
753 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
754 .min(context.map_or(
755 crate::query::MAX_RAW_INDEX_CANDIDATES,
756 crate::query::AiExecutionContext::max_fused_candidates,
757 ))
758}
759
760#[cfg(test)]
761mod ann_candidate_cap_tests {
762 use super::*;
763
764 #[test]
765 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
766 assert_eq!(
767 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
768 crate::query::MAX_RAW_INDEX_CANDIDATES,
769 );
770 let context = crate::query::AiExecutionContext::with_limits(
771 std::time::Duration::from_secs(1),
772 usize::MAX,
773 17,
774 );
775 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
776 }
777}
778
779fn civil_from_days(z: i64) -> (i64, u32, u32) {
780 let z = z + 719_468;
781 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
782 let doe = z - era * 146_097;
783 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
784 let y = yoe + era * 400;
785 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
786 let mp = (5 * doy + 2) / 153;
787 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
788 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
789 (if m <= 2 { y + 1 } else { y }, m, d)
790}
791
792pub fn clustered_row_id(primary_key: &Value) -> RowId {
799 let mut hash: u64 = 0xcbf29ce484222325;
800 for byte in primary_key.encode_key() {
801 hash ^= u64::from(byte);
802 hash = hash.wrapping_mul(0x100000001b3);
803 }
804 RowId(hash.max(1))
805}
806
807const 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;
814
815#[derive(Clone, Copy, Debug)]
830struct AutoIncState {
831 column_id: u16,
832 next: i64,
833 seeded: bool,
834}
835
836pub(crate) struct RecoveryMetadataPlan {
837 live_count: u64,
838 auto_inc: Option<AutoIncState>,
839 changed: bool,
840}
841
842type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
843
844fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
847 schema.auto_increment_column().map(|c| AutoIncState {
848 column_id: c.id,
849 next: 0,
850 seeded: false,
851 })
852}
853
854#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
867pub enum IndexBuildPolicy {
868 #[default]
870 Deferred,
871 Eager,
873}
874
875#[derive(Clone)]
876struct ReversePkSegment {
877 values: HashMap<RowId, Vec<u8>>,
878 removed: HashSet<RowId>,
879}
880
881#[derive(Clone)]
882struct ReversePkMap {
883 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
884 active: ReversePkSegment,
885}
886
887impl ReversePkMap {
888 fn new() -> Self {
889 Self {
890 frozen: Arc::new(Vec::new()),
891 active: ReversePkSegment {
892 values: HashMap::new(),
893 removed: HashSet::new(),
894 },
895 }
896 }
897
898 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
899 let mut map = Self::new();
900 map.active.values.extend(entries);
901 map
902 }
903
904 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
905 self.active.removed.remove(&row_id);
906 self.active.values.insert(row_id, key);
907 }
908
909 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
910 if let Some(key) = self.active.values.get(row_id) {
911 return Some(key);
912 }
913 if self.active.removed.contains(row_id) {
914 return None;
915 }
916 for segment in self.frozen.iter().rev() {
917 if let Some(key) = segment.values.get(row_id) {
918 return Some(key);
919 }
920 if segment.removed.contains(row_id) {
921 return None;
922 }
923 }
924 None
925 }
926
927 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
928 let previous = self.get(row_id).cloned();
929 self.active.values.remove(row_id);
930 self.active.removed.insert(*row_id);
931 previous
932 }
933
934 fn clear(&mut self) {
935 *self = Self::new();
936 }
937
938 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
939 let mut entries = HashMap::new();
940 for segment in self
941 .frozen
942 .iter()
943 .map(Arc::as_ref)
944 .chain(std::iter::once(&self.active))
945 {
946 for row_id in &segment.removed {
947 entries.remove(row_id);
948 }
949 entries.extend(
950 segment
951 .values
952 .iter()
953 .map(|(row_id, key)| (*row_id, key.clone())),
954 );
955 }
956 entries
957 }
958
959 fn seal(&mut self) {
960 if self.active.values.is_empty() && self.active.removed.is_empty() {
961 return;
962 }
963 let active = std::mem::replace(
964 &mut self.active,
965 ReversePkSegment {
966 values: HashMap::new(),
967 removed: HashSet::new(),
968 },
969 );
970 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
971 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
972 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
973 values: self.entries(),
974 removed: HashSet::new(),
975 })]);
976 }
977 }
978}
979
980#[derive(Clone)]
993pub struct ReadGeneration {
994 schema: Arc<Schema>,
995 base_runs: Arc<Vec<RunRef>>,
996 deltas: TableDeltas,
997 indexes: Arc<IndexGeneration>,
998 visible_through: Epoch,
999}
1000
1001pub(crate) enum SecondaryIndexArtifact {
1005 Bitmap(u16, BitmapIndex),
1006 LearnedRange(u16, ColumnLearnedRange),
1007 Fm(u16, Box<FmIndex>),
1008 Ann(u16, Box<AnnIndex>),
1009 Sparse(u16, SparseIndex),
1010 MinHash(u16, MinHashIndex),
1011}
1012
1013#[derive(Clone)]
1019pub struct TableDeltas {
1020 memtable: Memtable,
1021 mutable_run: MutableRun,
1022 hot: HotIndex,
1023 pk_by_row: ReversePkMap,
1024}
1025
1026impl ReadGeneration {
1027 fn empty(schema: &Schema) -> Self {
1030 Self {
1031 schema: Arc::new(schema.clone()),
1032 base_runs: Arc::new(Vec::new()),
1033 deltas: TableDeltas {
1034 memtable: Memtable::new(),
1035 mutable_run: MutableRun::new(),
1036 hot: HotIndex::new(),
1037 pk_by_row: ReversePkMap::new(),
1038 },
1039 indexes: Arc::new(IndexGeneration::default()),
1040 visible_through: Epoch(0),
1041 }
1042 }
1043
1044 pub fn schema(&self) -> &Arc<Schema> {
1046 &self.schema
1047 }
1048
1049 pub fn base_runs(&self) -> &[RunRef] {
1051 &self.base_runs
1052 }
1053
1054 pub fn indexes(&self) -> &Arc<IndexGeneration> {
1056 &self.indexes
1057 }
1058
1059 pub fn visible_through(&self) -> Epoch {
1061 self.visible_through
1062 }
1063
1064 pub fn deltas(&self) -> &TableDeltas {
1068 &self.deltas
1069 }
1070}
1071
1072impl TableDeltas {
1073 pub fn approx_bytes(&self) -> u64 {
1076 self.memtable
1077 .approx_bytes()
1078 .saturating_add(self.mutable_run.approx_bytes())
1079 }
1080
1081 pub fn memtable_len(&self) -> usize {
1083 self.memtable.len()
1084 }
1085
1086 pub fn mutable_run_len(&self) -> usize {
1088 self.mutable_run.len()
1089 }
1090
1091 pub fn hot_len(&self) -> usize {
1093 self.hot.len()
1094 }
1095
1096 pub fn reverse_pk_len(&self) -> usize {
1098 self.pk_by_row.entries().len()
1099 }
1100}
1101
1102#[derive(Clone)]
1104pub struct Table {
1105 dir: PathBuf,
1106 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1107 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
1108 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
1109 table_id: u64,
1110 name: String,
1114 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1119 read_only: bool,
1122 durable_commit_failed: bool,
1126 wal: WalSink,
1127 memtable: Memtable,
1128 mutable_run: MutableRun,
1133 mutable_run_spill_bytes: u64,
1135 compaction_zstd_level: i32,
1138 allocator: RowIdAllocator,
1139 epoch: Arc<EpochAuthority>,
1140 data_generation: u64,
1143 schema: Schema,
1144 hot: HotIndex,
1145 kek: Option<Arc<Kek>>,
1148 column_keys: HashMap<u16, ([u8; 32], u8)>,
1152 run_refs: Vec<RunRef>,
1153 retiring: Vec<crate::manifest::RetiredRun>,
1156 next_run_id: u64,
1157 sync_byte_threshold: u64,
1158 current_txn_id: u64,
1163 pending_private_mutations: bool,
1167 bitmap: HashMap<u16, BitmapIndex>,
1168 ann: HashMap<u16, AnnIndex>,
1169 fm: HashMap<u16, FmIndex>,
1170 sparse: HashMap<u16, SparseIndex>,
1171 minhash: HashMap<u16, MinHashIndex>,
1172 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
1175 pk_by_row: ReversePkMap,
1177 pinned: BTreeMap<Epoch, usize>,
1180 pub(crate) live_count: u64,
1183 reservoir: crate::reservoir::Reservoir,
1186 reservoir_complete: bool,
1194 had_deletes: bool,
1198 recent_delete_preimages: HashMap<Vec<u8>, Row>,
1203 run_row_id_ranges: HashMap<u128, (u64, u64)>,
1206 agg_cache: Arc<HashMap<u64, CachedAgg>>,
1210 global_idx_epoch: u64,
1214 indexes_complete: bool,
1219 index_build_policy: IndexBuildPolicy,
1221 pk_by_row_complete: bool,
1228 flushed_epoch: u64,
1231 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1234 snapshots: Arc<crate::retention::SnapshotRegistry>,
1237 commit_lock: Arc<parking_lot::Mutex<()>>,
1239 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1242 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
1252 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
1261 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1263 pending_delete_rids: roaring::RoaringBitmap,
1266 pending_put_cols: std::collections::HashSet<u16>,
1269 pending_rows: Vec<Row>,
1275 pending_rows_auto_inc: Vec<bool>,
1276 pending_dels: Vec<RowId>,
1279 pending_truncate: Option<Epoch>,
1283 auto_inc: Option<AutoIncState>,
1286 ttl: Option<TtlPolicy>,
1289 pins: Arc<crate::retention::PinRegistry>,
1296 published: Arc<ArcSwap<ReadGeneration>>,
1301 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
1306 pub(crate) lookup_metrics: LookupMetrics,
1311 run_lookup: RunLookupState,
1315}
1316
1317#[derive(Debug, Default, Clone)]
1325struct RunLookupState {
1326 directory: Option<std::sync::Arc<crate::run_lookup::RunLookupDirectory>>,
1327 fingerprint: u64,
1328 complete: bool,
1329}
1330
1331#[derive(Debug, Default)]
1332pub struct LookupMetrics {
1333 hot_lookup_hit: std::sync::atomic::AtomicU64,
1334 hot_lookup_fallback: std::sync::atomic::AtomicU64,
1335 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64,
1336 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64,
1337 result_cache_memory_hit: std::sync::atomic::AtomicU64,
1338 result_cache_disk_hit: std::sync::atomic::AtomicU64,
1339 result_cache_miss: std::sync::atomic::AtomicU64,
1340 result_cache_persistent_write_us: std::sync::atomic::AtomicU64,
1341 get_run_opened: std::sync::atomic::AtomicU64,
1343 get_run_skipped: std::sync::atomic::AtomicU64,
1344 pub(crate) directory_lookup_hit: std::sync::atomic::AtomicU64,
1346 pub(crate) directory_lookup_fallback: std::sync::atomic::AtomicU64,
1347 pub(crate) directory_incomplete: std::sync::atomic::AtomicU64,
1348 pub(crate) directory_run_readers_opened: std::sync::atomic::AtomicU64,
1349 pub(crate) directory_early_stop_total: std::sync::atomic::AtomicU64,
1350 pub(crate) directory_complete_miss_total: std::sync::atomic::AtomicU64,
1356 pub(crate) result_cache_persist_enqueued_total: std::sync::atomic::AtomicU64,
1358 pub(crate) result_cache_persist_coalesced_total: std::sync::atomic::AtomicU64,
1359 pub(crate) result_cache_persist_dropped_store_total: std::sync::atomic::AtomicU64,
1360 pub(crate) result_cache_persist_remove_total: std::sync::atomic::AtomicU64,
1361 pub(crate) result_cache_persist_stale_store_skipped_total: std::sync::atomic::AtomicU64,
1362 pub(crate) result_cache_persist_errors_total: std::sync::atomic::AtomicU64,
1363 pub(crate) result_cache_persist_shutdown_abandoned_total: std::sync::atomic::AtomicU64,
1364 pub(crate) result_cache_persist_queue_depth: std::sync::atomic::AtomicU64,
1365 pub(crate) hot_fallback_reasons: [std::sync::atomic::AtomicU64; 9],
1367 pub(crate) hot_fallback_overlay_versions_total: std::sync::atomic::AtomicU64,
1368 pub(crate) hot_fallback_runs_considered_total: std::sync::atomic::AtomicU64,
1369 pub(crate) hot_fallback_runs_opened_total: std::sync::atomic::AtomicU64,
1370 pub(crate) hot_fallback_pages_decoded_total: std::sync::atomic::AtomicU64,
1371 pub(crate) hot_fallback_rows_materialized_total: std::sync::atomic::AtomicU64,
1372 pub(crate) hot_lookup_duration_nanos: std::sync::atomic::AtomicU64,
1373 pub(crate) hot_fallback_duration_nanos: std::sync::atomic::AtomicU64,
1374 pub(crate) hot_mapping_rebuild_total: std::sync::atomic::AtomicU64,
1375 pub(crate) hot_checkpoint_rejected_total: std::sync::atomic::AtomicU64,
1376}
1377
1378impl Clone for LookupMetrics {
1379 fn clone(&self) -> Self {
1380 let mut hot_fallback_reasons = [
1381 std::sync::atomic::AtomicU64::new(0),
1382 std::sync::atomic::AtomicU64::new(0),
1383 std::sync::atomic::AtomicU64::new(0),
1384 std::sync::atomic::AtomicU64::new(0),
1385 std::sync::atomic::AtomicU64::new(0),
1386 std::sync::atomic::AtomicU64::new(0),
1387 std::sync::atomic::AtomicU64::new(0),
1388 std::sync::atomic::AtomicU64::new(0),
1389 std::sync::atomic::AtomicU64::new(0),
1390 ];
1391 for (dst, src) in hot_fallback_reasons
1392 .iter_mut()
1393 .zip(self.hot_fallback_reasons.iter())
1394 {
1395 dst.store(
1396 src.load(std::sync::atomic::Ordering::Relaxed),
1397 std::sync::atomic::Ordering::Relaxed,
1398 );
1399 }
1400 let copy_atomic = |src: &std::sync::atomic::AtomicU64| {
1401 std::sync::atomic::AtomicU64::new(src.load(std::sync::atomic::Ordering::Relaxed))
1402 };
1403 Self {
1404 hot_lookup_hit: copy_atomic(&self.hot_lookup_hit),
1405 hot_lookup_fallback: copy_atomic(&self.hot_lookup_fallback),
1406 hot_lookup_fallback_overlay_rows: copy_atomic(&self.hot_lookup_fallback_overlay_rows),
1407 hot_lookup_fallback_runs: copy_atomic(&self.hot_lookup_fallback_runs),
1408 result_cache_memory_hit: copy_atomic(&self.result_cache_memory_hit),
1409 result_cache_disk_hit: copy_atomic(&self.result_cache_disk_hit),
1410 result_cache_miss: copy_atomic(&self.result_cache_miss),
1411 result_cache_persistent_write_us: copy_atomic(&self.result_cache_persistent_write_us),
1412 get_run_opened: copy_atomic(&self.get_run_opened),
1413 get_run_skipped: copy_atomic(&self.get_run_skipped),
1414 directory_lookup_hit: copy_atomic(&self.directory_lookup_hit),
1415 directory_lookup_fallback: copy_atomic(&self.directory_lookup_fallback),
1416 directory_incomplete: copy_atomic(&self.directory_incomplete),
1417 directory_run_readers_opened: copy_atomic(&self.directory_run_readers_opened),
1418 directory_early_stop_total: copy_atomic(&self.directory_early_stop_total),
1419 directory_complete_miss_total: copy_atomic(&self.directory_complete_miss_total),
1420 result_cache_persist_enqueued_total: copy_atomic(
1421 &self.result_cache_persist_enqueued_total,
1422 ),
1423 result_cache_persist_coalesced_total: copy_atomic(
1424 &self.result_cache_persist_coalesced_total,
1425 ),
1426 result_cache_persist_dropped_store_total: copy_atomic(
1427 &self.result_cache_persist_dropped_store_total,
1428 ),
1429 result_cache_persist_remove_total: copy_atomic(&self.result_cache_persist_remove_total),
1430 result_cache_persist_stale_store_skipped_total: copy_atomic(
1431 &self.result_cache_persist_stale_store_skipped_total,
1432 ),
1433 result_cache_persist_errors_total: copy_atomic(&self.result_cache_persist_errors_total),
1434 result_cache_persist_shutdown_abandoned_total: copy_atomic(
1435 &self.result_cache_persist_shutdown_abandoned_total,
1436 ),
1437 result_cache_persist_queue_depth: copy_atomic(&self.result_cache_persist_queue_depth),
1438 hot_fallback_reasons,
1439 hot_fallback_overlay_versions_total: copy_atomic(
1440 &self.hot_fallback_overlay_versions_total,
1441 ),
1442 hot_fallback_runs_considered_total: copy_atomic(
1443 &self.hot_fallback_runs_considered_total,
1444 ),
1445 hot_fallback_runs_opened_total: copy_atomic(&self.hot_fallback_runs_opened_total),
1446 hot_fallback_pages_decoded_total: copy_atomic(&self.hot_fallback_pages_decoded_total),
1447 hot_fallback_rows_materialized_total: copy_atomic(
1448 &self.hot_fallback_rows_materialized_total,
1449 ),
1450 hot_lookup_duration_nanos: copy_atomic(&self.hot_lookup_duration_nanos),
1451 hot_fallback_duration_nanos: copy_atomic(&self.hot_fallback_duration_nanos),
1452 hot_mapping_rebuild_total: copy_atomic(&self.hot_mapping_rebuild_total),
1453 hot_checkpoint_rejected_total: copy_atomic(&self.hot_checkpoint_rejected_total),
1454 }
1455 }
1456}
1457
1458impl LookupMetrics {
1459 fn snapshot(&self) -> LookupMetricsSnapshot {
1460 LookupMetricsSnapshot {
1461 hot_lookup_hit: self
1462 .hot_lookup_hit
1463 .load(std::sync::atomic::Ordering::Relaxed),
1464 hot_lookup_fallback: self
1465 .hot_lookup_fallback
1466 .load(std::sync::atomic::Ordering::Relaxed),
1467 hot_lookup_fallback_overlay_rows: self
1468 .hot_lookup_fallback_overlay_rows
1469 .load(std::sync::atomic::Ordering::Relaxed),
1470 hot_lookup_fallback_runs: self
1471 .hot_lookup_fallback_runs
1472 .load(std::sync::atomic::Ordering::Relaxed),
1473 result_cache_memory_hit: self
1474 .result_cache_memory_hit
1475 .load(std::sync::atomic::Ordering::Relaxed),
1476 result_cache_disk_hit: self
1477 .result_cache_disk_hit
1478 .load(std::sync::atomic::Ordering::Relaxed),
1479 result_cache_miss: self
1480 .result_cache_miss
1481 .load(std::sync::atomic::Ordering::Relaxed),
1482 result_cache_persistent_write_us: self
1483 .result_cache_persistent_write_us
1484 .load(std::sync::atomic::Ordering::Relaxed),
1485 get_run_opened: self
1486 .get_run_opened
1487 .load(std::sync::atomic::Ordering::Relaxed),
1488 get_run_skipped: self
1489 .get_run_skipped
1490 .load(std::sync::atomic::Ordering::Relaxed),
1491 directory_lookup_hit: self
1492 .directory_lookup_hit
1493 .load(std::sync::atomic::Ordering::Relaxed),
1494 directory_lookup_fallback: self
1495 .directory_lookup_fallback
1496 .load(std::sync::atomic::Ordering::Relaxed),
1497 directory_incomplete: self
1498 .directory_incomplete
1499 .load(std::sync::atomic::Ordering::Relaxed),
1500 directory_run_readers_opened: self
1501 .directory_run_readers_opened
1502 .load(std::sync::atomic::Ordering::Relaxed),
1503 directory_early_stop_total: self
1504 .directory_early_stop_total
1505 .load(std::sync::atomic::Ordering::Relaxed),
1506 directory_complete_miss_total: self
1507 .directory_complete_miss_total
1508 .load(std::sync::atomic::Ordering::Relaxed),
1509 result_cache_persist_enqueued_total: self
1510 .result_cache_persist_enqueued_total
1511 .load(std::sync::atomic::Ordering::Relaxed),
1512 result_cache_persist_coalesced_total: self
1513 .result_cache_persist_coalesced_total
1514 .load(std::sync::atomic::Ordering::Relaxed),
1515 result_cache_persist_dropped_store_total: self
1516 .result_cache_persist_dropped_store_total
1517 .load(std::sync::atomic::Ordering::Relaxed),
1518 result_cache_persist_remove_total: self
1519 .result_cache_persist_remove_total
1520 .load(std::sync::atomic::Ordering::Relaxed),
1521 result_cache_persist_stale_store_skipped_total: self
1522 .result_cache_persist_stale_store_skipped_total
1523 .load(std::sync::atomic::Ordering::Relaxed),
1524 result_cache_persist_errors_total: self
1525 .result_cache_persist_errors_total
1526 .load(std::sync::atomic::Ordering::Relaxed),
1527 result_cache_persist_shutdown_abandoned_total: self
1528 .result_cache_persist_shutdown_abandoned_total
1529 .load(std::sync::atomic::Ordering::Relaxed),
1530 result_cache_persist_queue_depth: self
1531 .result_cache_persist_queue_depth
1532 .load(std::sync::atomic::Ordering::Relaxed),
1533 hot_fallback_reasons: [
1534 self.hot_fallback_reasons[0].load(std::sync::atomic::Ordering::Relaxed),
1535 self.hot_fallback_reasons[1].load(std::sync::atomic::Ordering::Relaxed),
1536 self.hot_fallback_reasons[2].load(std::sync::atomic::Ordering::Relaxed),
1537 self.hot_fallback_reasons[3].load(std::sync::atomic::Ordering::Relaxed),
1538 self.hot_fallback_reasons[4].load(std::sync::atomic::Ordering::Relaxed),
1539 self.hot_fallback_reasons[5].load(std::sync::atomic::Ordering::Relaxed),
1540 self.hot_fallback_reasons[6].load(std::sync::atomic::Ordering::Relaxed),
1541 self.hot_fallback_reasons[7].load(std::sync::atomic::Ordering::Relaxed),
1542 self.hot_fallback_reasons[8].load(std::sync::atomic::Ordering::Relaxed),
1543 ],
1544 hot_fallback_overlay_versions_total: self
1545 .hot_fallback_overlay_versions_total
1546 .load(std::sync::atomic::Ordering::Relaxed),
1547 hot_fallback_runs_considered_total: self
1548 .hot_fallback_runs_considered_total
1549 .load(std::sync::atomic::Ordering::Relaxed),
1550 hot_fallback_runs_opened_total: self
1551 .hot_fallback_runs_opened_total
1552 .load(std::sync::atomic::Ordering::Relaxed),
1553 hot_fallback_pages_decoded_total: self
1554 .hot_fallback_pages_decoded_total
1555 .load(std::sync::atomic::Ordering::Relaxed),
1556 hot_fallback_rows_materialized_total: self
1557 .hot_fallback_rows_materialized_total
1558 .load(std::sync::atomic::Ordering::Relaxed),
1559 hot_lookup_duration_nanos: self
1560 .hot_lookup_duration_nanos
1561 .load(std::sync::atomic::Ordering::Relaxed),
1562 hot_fallback_duration_nanos: self
1563 .hot_fallback_duration_nanos
1564 .load(std::sync::atomic::Ordering::Relaxed),
1565 hot_mapping_rebuild_total: self
1566 .hot_mapping_rebuild_total
1567 .load(std::sync::atomic::Ordering::Relaxed),
1568 hot_checkpoint_rejected_total: self
1569 .hot_checkpoint_rejected_total
1570 .load(std::sync::atomic::Ordering::Relaxed),
1571 }
1572 }
1573}
1574
1575#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
1578pub struct LookupMetricsSnapshot {
1579 pub hot_lookup_hit: u64,
1580 pub hot_lookup_fallback: u64,
1581 pub hot_lookup_fallback_overlay_rows: u64,
1582 pub hot_lookup_fallback_runs: u64,
1583 pub result_cache_memory_hit: u64,
1584 pub result_cache_disk_hit: u64,
1585 pub result_cache_miss: u64,
1586 pub result_cache_persistent_write_us: u64,
1587 pub get_run_opened: u64,
1588 pub get_run_skipped: u64,
1589 pub directory_lookup_hit: u64,
1591 pub directory_lookup_fallback: u64,
1592 pub directory_incomplete: u64,
1593 pub directory_run_readers_opened: u64,
1594 pub directory_early_stop_total: u64,
1595 pub directory_complete_miss_total: u64,
1596 pub result_cache_persist_enqueued_total: u64,
1598 pub result_cache_persist_coalesced_total: u64,
1599 pub result_cache_persist_dropped_store_total: u64,
1600 pub result_cache_persist_remove_total: u64,
1601 pub result_cache_persist_stale_store_skipped_total: u64,
1602 pub result_cache_persist_errors_total: u64,
1603 pub result_cache_persist_shutdown_abandoned_total: u64,
1604 pub result_cache_persist_queue_depth: u64,
1605 pub hot_fallback_reasons: [u64; 9],
1607 pub hot_fallback_overlay_versions_total: u64,
1608 pub hot_fallback_runs_considered_total: u64,
1609 pub hot_fallback_runs_opened_total: u64,
1610 pub hot_fallback_pages_decoded_total: u64,
1611 pub hot_fallback_rows_materialized_total: u64,
1612 pub hot_lookup_duration_nanos: u64,
1613 pub hot_fallback_duration_nanos: u64,
1614 pub hot_mapping_rebuild_total: u64,
1615 pub hot_checkpoint_rejected_total: u64,
1616}
1617
1618pub fn hot_fallback_reason_index(r: crate::trace::HotFallbackReason) -> usize {
1621 match r {
1622 crate::trace::HotFallbackReason::MissingMapping => 0,
1623 crate::trace::HotFallbackReason::StaleRowId => 1,
1624 crate::trace::HotFallbackReason::InvisibleAtSnapshot => 2,
1625 crate::trace::HotFallbackReason::HistoricalSnapshot => 3,
1626 crate::trace::HotFallbackReason::Tombstone => 4,
1627 crate::trace::HotFallbackReason::TtlExpired => 5,
1628 crate::trace::HotFallbackReason::PrimaryKeyMismatch => 6,
1629 crate::trace::HotFallbackReason::IndexIncomplete => 7,
1630 crate::trace::HotFallbackReason::CheckpointRejected => 8,
1631 }
1632}
1633
1634const _: () = {
1641 const fn assert_sync<T: ?Sized + Sync>() {}
1642 assert_sync::<Table>();
1643};
1644
1645#[derive(Clone)]
1651enum CachedData {
1652 Rows(Arc<Vec<Row>>),
1653 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
1654}
1655
1656impl CachedData {
1657 fn approx_bytes(&self) -> u64 {
1658 match self {
1659 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
1660 CachedData::Columns(c) => c
1661 .iter()
1662 .map(|(_, c)| c.approx_bytes())
1663 .sum::<u64>()
1664 .saturating_add(c.len() as u64 * 16),
1665 }
1666 }
1667}
1668
1669struct CachedEntry {
1673 data: CachedData,
1674 footprint: roaring::RoaringBitmap,
1675 condition_cols: Vec<u16>,
1676}
1677
1678impl CachedEntry {
1679 fn clone_for_persist(&self) -> Self {
1685 Self {
1686 data: self.data.clone(),
1687 footprint: self.footprint.clone(),
1688 condition_cols: self.condition_cols.clone(),
1689 }
1690 }
1691}
1692
1693struct ResultCache {
1704 entries: std::collections::HashMap<u64, CachedEntry>,
1705 order: BTreeSet<(u64, u64)>,
1706 generations: HashMap<u64, u64>,
1707 next_generation: u64,
1708 condition_index: HashMap<u16, HashSet<u64>>,
1711 empty_footprint_entries: HashSet<u64>,
1714 bytes: u64,
1715 max_bytes: u64,
1716 dir: Option<std::path::PathBuf>,
1717 #[allow(dead_code)]
1718 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1719 memory_hit: std::sync::atomic::AtomicU64,
1723 disk_hit: std::sync::atomic::AtomicU64,
1724 miss: std::sync::atomic::AtomicU64,
1725 persistent_write_us: std::sync::atomic::AtomicU64,
1726 persist_min_bytes: u64,
1731 writer: Option<std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>>,
1737 worker_handle: Option<std::thread::JoinHandle<()>>,
1740 completion_rx: Option<std::sync::mpsc::Receiver<()>>,
1745 cache_payload_cipher: Option<std::sync::Arc<crate::encryption::AesCipher>>,
1749}
1750
1751#[derive(serde::Serialize, serde::Deserialize)]
1753struct SerializedEntry {
1754 condition_cols: Vec<u16>,
1755 footprint_bits: Vec<u32>,
1756 data: SerializedData,
1757}
1758
1759#[derive(serde::Serialize, serde::Deserialize)]
1760enum SerializedData {
1761 Rows(Vec<Row>),
1762 Columns(Vec<(u16, columnar::NativeColumn)>),
1763}
1764
1765#[derive(serde::Serialize)]
1766struct SerializedEntryRef<'a> {
1767 condition_cols: &'a [u16],
1768 footprint_bits: Vec<u32>,
1769 data: SerializedDataRef<'a>,
1770}
1771
1772#[derive(serde::Serialize)]
1773enum SerializedDataRef<'a> {
1774 Rows(&'a [Row]),
1775 Columns(&'a [(u16, columnar::NativeColumn)]),
1776}
1777
1778impl<'a> SerializedEntryRef<'a> {
1779 fn from_entry(entry: &'a CachedEntry) -> Self {
1780 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
1781 let data = match &entry.data {
1782 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
1783 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
1784 };
1785 Self {
1786 condition_cols: &entry.condition_cols,
1787 footprint_bits,
1788 data,
1789 }
1790 }
1791}
1792
1793impl SerializedEntry {
1794 fn into_entry(self) -> Option<CachedEntry> {
1795 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1796 let data = match self.data {
1797 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1798 SerializedData::Columns(c) => {
1799 if !c.iter().all(|(_, col)| col.validate()) {
1802 return None;
1803 }
1804 CachedData::Columns(Arc::new(c))
1805 }
1806 };
1807 Some(CachedEntry {
1808 data,
1809 footprint,
1810 condition_cols: self.condition_cols,
1811 })
1812 }
1813}
1814
1815impl ResultCache {
1816 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1817
1818 fn new() -> Self {
1819 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1820 }
1821
1822 fn with_max_bytes(max_bytes: u64) -> Self {
1823 Self {
1824 entries: std::collections::HashMap::new(),
1825 order: BTreeSet::new(),
1826 generations: HashMap::new(),
1827 next_generation: 0,
1828 condition_index: HashMap::new(),
1829 empty_footprint_entries: HashSet::new(),
1830 bytes: 0,
1831 max_bytes,
1832 dir: None,
1833 cache_dek: None,
1834 memory_hit: std::sync::atomic::AtomicU64::new(0),
1835 disk_hit: std::sync::atomic::AtomicU64::new(0),
1836 miss: std::sync::atomic::AtomicU64::new(0),
1837 persistent_write_us: std::sync::atomic::AtomicU64::new(0),
1838 persist_min_bytes: 4 * 1024,
1841 writer: None,
1842 worker_handle: None,
1843 completion_rx: None,
1844 cache_payload_cipher: None,
1845 }
1846 }
1847
1848 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1849 let _ = std::fs::create_dir_all(&dir);
1850 self.dir = Some(dir);
1851 self
1852 }
1853
1854 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1855 self.cache_dek = dek.clone();
1861 if let Some(dek) = dek {
1862 if let Ok(cipher) = crate::encryption::AesCipher::new(&dek[..]) {
1863 self.cache_payload_cipher = Some(std::sync::Arc::new(cipher));
1864 }
1865 }
1866 self
1867 }
1868
1869 fn install_persistent_writer(
1874 &mut self,
1875 writer: std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
1876 worker_handle: std::thread::JoinHandle<()>,
1877 completion_rx: std::sync::mpsc::Receiver<()>,
1878 ) {
1879 self.writer = Some(writer);
1880 self.worker_handle = Some(worker_handle);
1881 self.completion_rx = Some(completion_rx);
1882 }
1883
1884 fn take_persistent_worker(&mut self) -> Option<std::thread::JoinHandle<()>> {
1887 self.worker_handle.take()
1888 }
1889
1890 fn take_persistent_completion(&mut self) -> Option<std::sync::mpsc::Receiver<()>> {
1893 self.completion_rx.take()
1894 }
1895
1896 fn persistent_writer(
1899 &self,
1900 ) -> Option<&std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>> {
1901 self.writer.as_ref()
1902 }
1903
1904 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1905 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1906 }
1907
1908 fn store_to_disk(&self, context: PersistContext, entry: &CachedEntry) {
1915 let Some(path) = self.disk_path(context.key) else {
1916 return;
1917 };
1918 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1919 Ok(s) => s,
1920 Err(_) => return,
1921 };
1922 let bytes = match crate::result_cache::encode_persisted_entry(
1923 context,
1924 &serialized,
1925 self.cache_payload_cipher.as_deref(),
1926 ) {
1927 Ok(b) => b,
1928 Err(_) => return,
1929 };
1930 let tmp = path.with_extension("tmp");
1931 use std::io::Write;
1932 let write = || -> std::io::Result<()> {
1933 let mut f = std::fs::File::create(&tmp)?;
1934 f.write_all(&bytes)?;
1935 f.flush()?;
1936 f.sync_all()?;
1937 Ok(())
1938 };
1939 if write().is_err() {
1940 let _ = std::fs::remove_file(&tmp);
1941 return;
1942 }
1943 if std::fs::rename(&tmp, &path).is_err() {
1944 let _ = std::fs::remove_file(&tmp);
1945 return;
1946 }
1947 if let Ok(dir) = std::fs::File::open(self.dir.as_ref().expect("dir set").clone()) {
1948 let _ = dir.sync_all();
1949 }
1950 }
1951
1952 fn load_from_disk(&self, key: u64, identity: PersistentCacheIdentity) -> Option<CachedEntry> {
1958 let path = self.disk_path(key)?;
1959 let bytes = std::fs::read(&path).ok()?;
1960 if bytes.len() < FRAME_MAGIC.len() || bytes[..FRAME_MAGIC.len()] != FRAME_MAGIC {
1964 let _ = std::fs::remove_file(&path);
1965 return None;
1966 }
1967 let plaintext = match crate::result_cache::decode_persisted_entry(
1968 identity,
1969 key,
1970 &bytes,
1971 self.cache_payload_cipher.as_deref(),
1972 ) {
1973 Ok(p) => p,
1974 Err(_) => {
1975 let _ = std::fs::remove_file(&path);
1976 return None;
1977 }
1978 };
1979 let serialized: SerializedEntry = match bincode::deserialize(&plaintext) {
1980 Ok(s) => s,
1981 Err(_) => {
1982 let _ = std::fs::remove_file(&path);
1983 return None;
1984 }
1985 };
1986 serialized.into_entry()
1987 }
1988
1989 fn remove_from_disk(&self, key: u64) {
1991 if let Some(path) = self.disk_path(key) {
1992 let _ = std::fs::remove_file(&path);
1993 }
1994 }
1995
1996 fn load_persistent(&mut self, identity: PersistentCacheIdentity) {
2000 let Some(dir) = self.dir.as_ref().cloned() else {
2001 return;
2002 };
2003 let entries = match std::fs::read_dir(&dir) {
2004 Ok(e) => e,
2005 Err(_) => return,
2006 };
2007 for entry in entries.flatten() {
2008 let path = entry.path();
2009 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
2011 let _ = std::fs::remove_file(&path);
2012 continue;
2013 }
2014 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
2015 continue;
2016 }
2017 let stem = match path.file_stem().and_then(|s| s.to_str()) {
2018 Some(s) => s,
2019 None => continue,
2020 };
2021 let key = match u64::from_str_radix(stem, 16) {
2022 Ok(k) => k,
2023 Err(_) => continue,
2024 };
2025 if let Some(entry) = self.load_from_disk(key, identity) {
2029 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
2030 self.index_entry(key, &entry);
2031 self.entries.insert(key, entry);
2032 self.touch(key);
2033 }
2034 }
2035 self.evict();
2036 }
2037
2038 fn set_max_bytes(&mut self, max_bytes: u64) {
2039 self.max_bytes = max_bytes;
2040 self.evict();
2041 }
2042
2043 fn touch(&mut self, key: u64) {
2046 if !self.entries.contains_key(&key) {
2047 return;
2048 }
2049 if let Some(previous) = self.generations.remove(&key) {
2050 self.order.remove(&(previous, key));
2051 }
2052 let generation = self.allocate_generation();
2053 self.generations.insert(key, generation);
2054 self.order.insert((generation, key));
2055 }
2056
2057 fn untrack(&mut self, key: u64) {
2058 if let Some(generation) = self.generations.remove(&key) {
2059 self.order.remove(&(generation, key));
2060 }
2061 }
2062
2063 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
2064 for column in &entry.condition_cols {
2065 self.condition_index.entry(*column).or_default().insert(key);
2066 }
2067 if entry.footprint.is_empty() {
2068 self.empty_footprint_entries.insert(key);
2069 }
2070 }
2071
2072 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
2073 for column in &entry.condition_cols {
2074 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
2075 keys.remove(&key);
2076 keys.is_empty()
2077 });
2078 if remove_column {
2079 self.condition_index.remove(column);
2080 }
2081 }
2082 if entry.footprint.is_empty() {
2083 self.empty_footprint_entries.remove(&key);
2084 }
2085 }
2086
2087 fn allocate_generation(&mut self) -> u64 {
2088 if self.next_generation == u64::MAX {
2089 self.rebase_generations();
2090 }
2091 let generation = self.next_generation;
2092 self.next_generation += 1;
2093 generation
2094 }
2095
2096 fn rebase_generations(&mut self) {
2097 let ordered = std::mem::take(&mut self.order);
2098 self.generations.clear();
2099 for (generation, (_, key)) in ordered.into_iter().enumerate() {
2100 let generation = generation as u64;
2101 self.generations.insert(key, generation);
2102 self.order.insert((generation, key));
2103 }
2104 self.next_generation = self.order.len() as u64;
2105 }
2106
2107 fn get_rows(&mut self, key: u64, identity: PersistentCacheIdentity) -> Option<Arc<Vec<Row>>> {
2108 let res = self.entries.get(&key).and_then(|e| match &e.data {
2109 CachedData::Rows(r) => Some(r.clone()),
2110 CachedData::Columns(_) => None,
2111 });
2112 if res.is_some() {
2113 self.touch(key);
2114 self.memory_hit
2115 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2116 return res;
2117 }
2118 if let Some(entry) = self.load_from_disk(key, identity) {
2120 let res = match &entry.data {
2121 CachedData::Rows(r) => Some(r.clone()),
2122 CachedData::Columns(_) => None,
2123 };
2124 if res.is_some() {
2125 let approx = entry.data.approx_bytes();
2126 self.bytes = self.bytes.saturating_add(approx);
2127 self.index_entry(key, &entry);
2128 self.entries.insert(key, entry);
2129 self.touch(key);
2130 self.evict();
2131 self.disk_hit
2132 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2133 return res;
2134 }
2135 }
2136 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2137 None
2138 }
2139
2140 fn get_columns(
2141 &mut self,
2142 key: u64,
2143 identity: PersistentCacheIdentity,
2144 ) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
2145 let res = self.entries.get(&key).and_then(|e| match &e.data {
2146 CachedData::Columns(c) => Some(c.clone()),
2147 CachedData::Rows(_) => None,
2148 });
2149 if res.is_some() {
2150 self.touch(key);
2151 self.memory_hit
2152 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2153 return res;
2154 }
2155 if let Some(entry) = self.load_from_disk(key, identity) {
2157 let res = match &entry.data {
2158 CachedData::Columns(c) => Some(c.clone()),
2159 CachedData::Rows(_) => None,
2160 };
2161 if res.is_some() {
2162 let approx = entry.data.approx_bytes();
2163 self.bytes = self.bytes.saturating_add(approx);
2164 self.index_entry(key, &entry);
2165 self.entries.insert(key, entry);
2166 self.touch(key);
2167 self.evict();
2168 self.disk_hit
2169 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2170 return res;
2171 }
2172 }
2173 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2174 None
2175 }
2176
2177 fn insert(&mut self, key: u64, entry: CachedEntry) {
2182 let approx = entry.data.approx_bytes();
2183 if let Some(previous) = self.entries.remove(&key) {
2184 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
2185 self.unindex_entry(key, &previous);
2186 self.untrack(key);
2187 }
2188 self.bytes = self.bytes.saturating_add(approx);
2189 self.index_entry(key, &entry);
2190 self.entries.insert(key, entry);
2191 self.touch(key);
2192 self.evict();
2193 }
2194
2195 fn persist_entry(&self, context: PersistContext, entry: &CachedEntry) {
2200 let approx = entry.data.approx_bytes();
2201 if self.dir.is_none() {
2202 return;
2203 }
2204 if self.persist_min_bytes > 0 && approx < self.persist_min_bytes {
2205 return;
2206 }
2207 if let Some(writer) = self.writer.as_ref() {
2208 let entry_clone = entry.clone_for_persist();
2212 let payload_factory: Box<dyn FnOnce() -> Option<Vec<u8>> + Send + 'static> =
2213 Box::new(move || {
2214 bincode::serialize(&SerializedEntryRef::from_entry(&entry_clone)).ok()
2215 });
2216 let persistable = crate::result_cache::PersistableEntry {
2217 key: context.key,
2218 table_id: context.identity.table_id,
2219 schema_id: context.identity.schema_id,
2220 run_generation: context.identity.logical_generation,
2221 entry_generation: context.entry_generation,
2222 bytes: approx as usize,
2223 payload_factory,
2224 };
2225 let write_start = std::time::Instant::now();
2226 writer.enqueue_store(persistable);
2227 let write_us = write_start.elapsed().as_micros() as u64;
2228 self.persistent_write_us
2229 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2230 } else {
2231 let write_start = std::time::Instant::now();
2233 self.store_to_disk(context, entry);
2234 let write_us = write_start.elapsed().as_micros() as u64;
2235 self.persistent_write_us
2236 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2237 }
2238 }
2239
2240 fn allocate_persist_generation(&mut self, key: u64) -> u64 {
2245 if let Some(writer) = self.writer.as_ref() {
2246 let _ = key; writer.bump_persist_generation(key)
2250 } else {
2251 0
2252 }
2253 }
2254
2255 fn enqueue_persist_remove(&self, key: u64) {
2258 if let Some(writer) = self.writer.as_ref() {
2259 writer.enqueue_remove(key);
2260 }
2261 }
2262
2263 fn enqueue_persist_clear(&self) {
2270 if let Some(writer) = self.writer.as_ref() {
2271 writer.enqueue_clear();
2272 }
2273 }
2274
2275 fn flush_persistent_cache(&self, deadline: std::time::Duration) -> u64 {
2281 let Some(writer) = self.persistent_writer() else {
2282 return 0;
2283 };
2284 let start = std::time::Instant::now();
2285 loop {
2286 let depth = writer.queue_depth() as u64;
2287 let in_flight = writer.writes_in_flight();
2293 if depth == 0 && in_flight == 0 {
2294 return 0;
2295 }
2296 if start.elapsed() >= deadline {
2297 return depth;
2298 }
2299 std::thread::sleep(std::time::Duration::from_millis(2));
2300 }
2301 }
2302
2303 fn shutdown_persistent_cache(&mut self, deadline: std::time::Duration) {
2313 let start = Instant::now();
2314 if let Some(writer) = self.persistent_writer() {
2316 loop {
2317 if writer.queue_depth() == 0 {
2318 break;
2319 }
2320 if start.elapsed() >= deadline {
2321 break;
2322 }
2323 std::thread::sleep(std::time::Duration::from_millis(2));
2324 }
2325 }
2326 let writer = self.persistent_writer().cloned();
2330 if let Some(writer) = writer.as_ref() {
2331 writer.shutdown();
2332 let _ = writer.drain_all_as_abandoned();
2333 }
2334 let mut completion = self.take_persistent_completion();
2341 let completed = if let (Some(rx), Some(_writer)) = (completion.as_mut(), writer.as_ref()) {
2342 let elapsed = start.elapsed();
2343 let remaining = deadline.saturating_sub(elapsed);
2344 crate::result_cache::recv_completion_with_timeout(rx, remaining)
2345 } else {
2346 true
2347 };
2348 if completed {
2349 if let Some(handle) = self.take_persistent_worker() {
2350 let _ = handle.join();
2351 }
2352 drop(completion.take());
2354 } else {
2355 if let Some(writer) = writer.as_ref() {
2359 writer.record_outcome(DrainOutcome::Abandoned);
2360 }
2361 self.take_persistent_worker(); }
2363 self.writer = None;
2367 }
2368
2369 pub(crate) fn set_persist_min_bytes(&mut self, min: u64) {
2373 self.persist_min_bytes = min;
2374 }
2375
2376 #[cfg(test)]
2378 fn would_persist(&self, approx: u64) -> bool {
2379 self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
2380 }
2381
2382 fn cache_counters(&self) -> (u64, u64, u64, u64) {
2385 (
2386 self.memory_hit.load(std::sync::atomic::Ordering::Relaxed),
2387 self.disk_hit.load(std::sync::atomic::Ordering::Relaxed),
2388 self.miss.load(std::sync::atomic::Ordering::Relaxed),
2389 self.persistent_write_us
2390 .load(std::sync::atomic::Ordering::Relaxed),
2391 )
2392 }
2393
2394 fn persist_snapshot(&self) -> Option<LookupMetricsSnapshot> {
2397 self.writer.as_ref().map(|w| w.persist_snapshot())
2398 }
2399
2400 fn invalidate(
2409 &mut self,
2410 delete_rids: &roaring::RoaringBitmap,
2411 put_cols: &std::collections::HashSet<u16>,
2412 ) {
2413 if self.entries.is_empty() {
2414 return;
2415 }
2416 let has_deletes = !delete_rids.is_empty();
2417 let mut to_remove = HashSet::new();
2418
2419 for column in put_cols {
2423 if let Some(keys) = self.condition_index.get(column) {
2424 to_remove.extend(keys.iter().copied());
2425 }
2426 }
2427
2428 if has_deletes {
2429 to_remove.extend(self.empty_footprint_entries.iter().copied());
2433 for (&key, entry) in &self.entries {
2434 if !to_remove.contains(&key)
2435 && !entry.footprint.is_empty()
2436 && entry.footprint.intersection_len(delete_rids) > 0
2437 {
2438 to_remove.insert(key);
2439 }
2440 }
2441 }
2442
2443 for key in to_remove {
2444 if let Some(entry) = self.entries.remove(&key) {
2445 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2446 self.unindex_entry(key, &entry);
2447 }
2448 if self.writer.is_some() {
2453 self.enqueue_persist_remove(key);
2454 } else {
2455 self.remove_from_disk(key);
2456 }
2457 self.untrack(key);
2458 }
2459 }
2460
2461 fn clear(&mut self) {
2462 if self.dir.is_some() {
2466 if self.writer.is_some() {
2467 self.enqueue_persist_clear();
2468 } else if let Some(dir) = &self.dir {
2469 if let Ok(entries) = std::fs::read_dir(dir) {
2470 for entry in entries.flatten() {
2471 let path = entry.path();
2472 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
2473 let _ = std::fs::remove_file(&path);
2474 }
2475 }
2476 }
2477 }
2478 }
2479 self.entries.clear();
2480 self.order.clear();
2481 self.generations.clear();
2482 self.next_generation = 0;
2483 self.condition_index.clear();
2484 self.empty_footprint_entries.clear();
2485 self.bytes = 0;
2486 }
2487
2488 fn evict(&mut self) {
2489 while self.bytes > self.max_bytes {
2490 let Some((_, key)) = self.order.pop_first() else {
2491 break;
2492 };
2493 self.generations.remove(&key);
2494 if let Some(entry) = self.entries.remove(&key) {
2495 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2496 self.unindex_entry(key, &entry);
2497 }
2505 }
2506 }
2507}
2508
2509fn spawn_persistent_cache_worker(
2517 dir: std::path::PathBuf,
2518 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
2519 metrics: LookupMetrics,
2520) -> std::io::Result<(
2521 std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
2522 std::thread::JoinHandle<()>,
2523 std::sync::mpsc::Receiver<()>,
2524)> {
2525 use crate::result_cache::{
2526 spawn_persistent_cache_worker as spawn, RealPersistentCacheIo, StalenessGuard,
2527 WorkerConfig, WriterStalenessGuard,
2528 };
2529 let io: std::sync::Arc<dyn crate::result_cache::PersistentCacheIo> =
2530 std::sync::Arc::new(RealPersistentCacheIo::new(dir)?);
2531 let cipher = match cache_dek {
2532 Some(dek) => {
2533 let cipher = crate::encryption::AesCipher::new(&dek[..])
2534 .map_err(|e| std::io::Error::other(format!("aes: {e}")))?;
2535 Some(std::sync::Arc::new(cipher))
2536 }
2537 None => None,
2538 };
2539 let writer = std::sync::Arc::new(crate::result_cache::PersistentResultCacheWriter::new(
2540 metrics,
2541 crate::result_cache::WriterLimits::default(),
2542 ));
2543 let staleness: std::sync::Arc<dyn StalenessGuard> =
2544 std::sync::Arc::new(WriterStalenessGuard::new(writer.clone()));
2545 let (completion_tx, completion_rx) = std::sync::mpsc::channel();
2546 let config = WorkerConfig {
2547 writer: writer.clone(),
2548 io,
2549 cipher,
2550 staleness,
2551 max_staleness_retries: 8,
2552 completion: Some(completion_tx),
2553 };
2554 let handle = spawn(config);
2555 Ok((writer, handle, completion_rx))
2556}
2557
2558#[cfg(test)]
2559mod result_cache_lru_tests {
2560 use super::*;
2561
2562 fn row_entry(row_id: u64) -> CachedEntry {
2563 CachedEntry {
2564 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
2565 footprint: std::iter::once(row_id as u32).collect(),
2566 condition_cols: vec![1],
2567 }
2568 }
2569
2570 #[test]
2571 fn hits_update_recency_without_growing_an_order_queue() {
2572 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2573 let identity = PersistentCacheIdentity {
2574 table_id: 0,
2575 schema_id: 0,
2576 logical_generation: 0,
2577 };
2578 cache.insert(1, row_entry(1));
2579 cache.insert(2, row_entry(2));
2580 assert_eq!(cache.order.len(), cache.entries.len());
2581
2582 for _ in 0..1_000 {
2583 assert!(cache.get_rows(1, identity).is_some());
2584 }
2585
2586 assert_eq!(cache.order.len(), cache.entries.len());
2587 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
2588
2589 let one_entry = cache.entries[&1].data.approx_bytes();
2590 cache.set_max_bytes(one_entry);
2591 assert!(cache.entries.contains_key(&1));
2592 assert!(!cache.entries.contains_key(&2));
2593 assert_eq!(cache.order.len(), cache.entries.len());
2594 }
2595
2596 #[test]
2597 fn tiny_entries_skip_persistent_tier_by_default() {
2598 let dir = tempfile::tempdir().unwrap();
2599 let mut cache = ResultCache::with_max_bytes(u64::MAX).with_dir(dir.path().to_path_buf());
2600 let tiny = row_entry(1).data.approx_bytes();
2601 assert!(
2602 tiny < 4 * 1024,
2603 "row_entry fixture must be below default 4KiB threshold"
2604 );
2605 assert!(
2606 !cache.would_persist(tiny),
2607 "tiny entry must not force synchronous disk publish"
2608 );
2609 assert!(
2610 cache.would_persist(8 * 1024),
2611 "large entry must still persist when dir is set"
2612 );
2613 cache.set_persist_min_bytes(0);
2614 assert!(
2615 cache.would_persist(tiny),
2616 "persist_min_bytes=0 restores always-persist policy"
2617 );
2618 }
2619
2620 #[test]
2621 fn insert_invalidation_uses_the_condition_reverse_index() {
2622 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2623 let mut first = row_entry(1);
2624 first.condition_cols = vec![1];
2625 let mut second = row_entry(2);
2626 second.condition_cols = vec![2];
2627 cache.insert(1, first);
2628 cache.insert(2, second);
2629
2630 let put_cols = std::iter::once(1_u16).collect();
2631 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
2632
2633 assert!(!cache.entries.contains_key(&1));
2634 assert!(cache.entries.contains_key(&2));
2635 assert!(!cache
2636 .condition_index
2637 .get(&1)
2638 .is_some_and(|keys| keys.contains(&1)));
2639 assert!(cache
2640 .condition_index
2641 .get(&2)
2642 .is_some_and(|keys| keys.contains(&2)));
2643 assert_eq!(cache.order.len(), cache.entries.len());
2644 }
2645
2646 #[test]
2647 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
2648 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2649 cache.insert(1, row_entry(1));
2650 cache.insert(2, row_entry(2));
2651 let mut unknown = row_entry(3);
2652 unknown.footprint.clear();
2653 cache.insert(3, unknown);
2654
2655 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
2656 cache.invalidate(&delete_rids, &HashSet::new());
2657
2658 assert!(!cache.entries.contains_key(&1));
2659 assert!(cache.entries.contains_key(&2));
2660 assert!(!cache.entries.contains_key(&3));
2661 assert!(cache.empty_footprint_entries.is_empty());
2662 assert_eq!(cache.order.len(), cache.entries.len());
2663 }
2664
2665 #[test]
2666 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
2667 let entry = row_entry(7);
2668 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
2669 let owned = bincode::serialize(&SerializedEntry {
2670 condition_cols: entry.condition_cols.clone(),
2671 footprint_bits: entry.footprint.iter().collect(),
2672 data: match &entry.data {
2673 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
2674 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
2675 },
2676 })
2677 .unwrap();
2678 assert_eq!(borrowed, owned);
2679 }
2680}
2681
2682type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
2689
2690fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
2691 let _ = kek;
2692 {
2693 if let Some(k) = kek {
2694 return (
2695 Some(k.derive_table_wal_key(_table_id)),
2696 Some(k.derive_cache_key()),
2697 );
2698 }
2699 }
2700 (None, None)
2701}
2702
2703fn read_table_encryption_salt_root(
2704 root: &crate::durable_file::DurableRoot,
2705) -> Result<[u8; crate::encryption::SALT_LEN]> {
2706 use std::io::Read;
2707
2708 let mut file = root
2709 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
2710 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
2711 let length = file.metadata()?.len();
2712 if length != crate::encryption::SALT_LEN as u64 {
2713 return Err(MongrelError::InvalidArgument(format!(
2714 "salt file is {length} bytes, expected {}",
2715 crate::encryption::SALT_LEN
2716 )));
2717 }
2718 let mut salt = [0_u8; crate::encryption::SALT_LEN];
2719 file.read_exact(&mut salt)?;
2720 Ok(salt)
2721}
2722
2723fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
2725 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
2726}
2727
2728fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
2729 let Some(kek) = kek else {
2730 return HashMap::new();
2731 };
2732 {
2733 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
2734 schema
2735 .columns
2736 .iter()
2737 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
2738 .map(|c| {
2739 let scheme = if schema
2740 .indexes
2741 .iter()
2742 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
2743 {
2744 SCHEME_OPE_RANGE
2745 } else {
2746 SCHEME_HMAC_EQ
2747 };
2748 let key: [u8; 32] = *kek.derive_column_key(c.id);
2749 (c.id, (key, scheme))
2750 })
2751 .collect()
2752 }
2753}
2754
2755pub(crate) struct SharedCtx {
2760 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
2761 pub epoch: Arc<EpochAuthority>,
2762 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
2763 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
2764 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
2765 pub kek: Option<Arc<Kek>>,
2766 pub commit_lock: Arc<parking_lot::Mutex<()>>,
2772 pub shared: Option<SharedWalCtx>,
2776 pub table_name: Option<String>,
2779 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
2782 pub read_only: bool,
2784}
2785
2786#[derive(Clone)]
2792pub(crate) struct SharedWalCtx {
2793 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
2794 pub group: Arc<GroupCommit>,
2795 pub poisoned: Arc<AtomicBool>,
2796 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
2797 pub change_wake: tokio::sync::broadcast::Sender<()>,
2798 pub lifecycle: Arc<crate::core::LifecycleController>,
2801 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
2803}
2804
2805enum WalSink {
2808 Private(Wal),
2809 Shared(SharedWalCtx),
2810 ReadOnly,
2811}
2812
2813impl Clone for WalSink {
2814 fn clone(&self) -> Self {
2815 match self {
2816 Self::Shared(shared) => Self::Shared(shared.clone()),
2817 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
2818 }
2819 }
2820}
2821
2822impl SharedCtx {
2823 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
2827 let n_shards = if cache_dir.is_some() {
2831 1
2832 } else {
2833 crate::cache::CACHE_SHARDS
2834 };
2835 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
2836 let page_cache = if let Some(d) = cache_dir {
2837 Arc::new(crate::cache::Sharded::new(1, || {
2838 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
2839 }))
2840 } else {
2841 Arc::new(crate::cache::Sharded::new(n_shards, || {
2842 crate::cache::PageCache::new(per_shard)
2843 }))
2844 };
2845 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
2846 let decoded_cache = Arc::new(crate::cache::Sharded::new(
2847 crate::cache::CACHE_SHARDS,
2848 || crate::cache::DecodedPageCache::new(decoded_per_shard),
2849 ));
2850 Self {
2851 root_guard: None,
2852 epoch: Arc::new(EpochAuthority::new(0)),
2853 page_cache,
2854 decoded_cache,
2855 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
2856 kek,
2857 commit_lock: Arc::new(parking_lot::Mutex::new(())),
2858 shared: None,
2859 table_name: None,
2860 auth: None,
2861 read_only: false,
2862 }
2863 }
2864}
2865
2866fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
2870 use crate::query::Condition;
2871 match c {
2872 Condition::Pk(_)
2874 | Condition::BitmapEq { .. }
2875 | Condition::BitmapIn { .. }
2876 | Condition::BytesPrefix { .. }
2877 | Condition::IsNull { .. }
2878 | Condition::IsNotNull { .. } => 0,
2879 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
2881 1
2882 }
2883 Condition::FmContains { .. }
2885 | Condition::FmContainsAll { .. }
2886 | Condition::Ann { .. }
2887 | Condition::SparseMatch { .. } => 2,
2888 }
2889}
2890
2891impl Table {
2892 pub(crate) fn build_secondary_index_artifact<F>(
2898 &self,
2899 definition: &IndexDef,
2900 rows: &[Row],
2901 mut checkpoint: F,
2902 ) -> Result<SecondaryIndexArtifact>
2903 where
2904 F: FnMut(usize, usize) -> Result<()>,
2905 {
2906 let mut build_schema = self.schema.clone();
2907 build_schema.indexes = vec![definition.clone()];
2908 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
2909 let name_to_id: HashMap<&str, u16> = self
2910 .schema
2911 .columns
2912 .iter()
2913 .map(|column| (column.name.as_str(), column.id))
2914 .collect();
2915 let mut hot = HotIndex::new();
2916 let mut accepted = Vec::with_capacity(rows.len());
2917
2918 for (offset, row) in rows.iter().enumerate() {
2919 checkpoint(offset, rows.len())?;
2920 if row.deleted {
2921 continue;
2922 }
2923 if let Some(predicate) = &definition.predicate {
2924 let columns: HashMap<u16, &Value> =
2925 row.columns.iter().map(|(id, value)| (*id, value)).collect();
2926 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
2927 continue;
2928 }
2929 }
2930 accepted.push(row);
2931 if definition.kind == IndexKind::LearnedRange {
2932 continue;
2933 }
2934 let effective = if self.column_keys.is_empty() {
2935 row.clone()
2936 } else {
2937 self.tokenized_for_indexes(row)
2938 };
2939 if definition.kind == IndexKind::Ann {
2940 if let (Some(index), Some(vector)) = (
2941 ann.get_mut(&definition.column_id),
2942 effective
2943 .columns
2944 .get(&definition.column_id)
2945 .and_then(Value::as_embedding),
2946 ) {
2947 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
2948 checkpoint(offset, rows.len())
2949 })?;
2950 }
2951 } else {
2952 index_into_single(
2953 definition,
2954 &build_schema,
2955 &effective,
2956 &mut hot,
2957 &mut bitmap,
2958 &mut ann,
2959 &mut fm,
2960 &mut sparse,
2961 &mut minhash,
2962 );
2963 }
2964 }
2965 checkpoint(rows.len(), rows.len())?;
2966
2967 if let Some(index) = ann.get_mut(&definition.column_id) {
2968 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
2969 }
2970
2971 let column_id = definition.column_id;
2972 match definition.kind {
2973 IndexKind::Bitmap => bitmap
2974 .remove(&column_id)
2975 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
2976 IndexKind::Ann => ann
2977 .remove(&column_id)
2978 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
2979 IndexKind::FmIndex => fm
2980 .remove(&column_id)
2981 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
2982 IndexKind::Sparse => sparse
2983 .remove(&column_id)
2984 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
2985 IndexKind::MinHash => minhash
2986 .remove(&column_id)
2987 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
2988 IndexKind::LearnedRange => {
2989 let epsilon = definition
2990 .options
2991 .learned_range
2992 .as_ref()
2993 .map(|options| options.epsilon)
2994 .unwrap_or(16);
2995 let column = self
2996 .schema
2997 .columns
2998 .iter()
2999 .find(|column| column.id == column_id)
3000 .ok_or_else(|| {
3001 MongrelError::Schema(format!(
3002 "index {} references unknown column {column_id}",
3003 definition.name
3004 ))
3005 })?;
3006 let index = match column.ty {
3007 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
3008 let pairs: Vec<_> = accepted
3009 .iter()
3010 .filter_map(|row| match row.columns.get(&column_id) {
3011 Some(Value::Int64(value)) if !row.deleted => {
3012 Some((*value, row.row_id.0))
3013 }
3014 _ => None,
3015 })
3016 .collect();
3017 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
3018 }
3019 TypeId::Float32 | TypeId::Float64 => {
3020 let pairs: Vec<_> = accepted
3021 .iter()
3022 .filter_map(|row| match row.columns.get(&column_id) {
3023 Some(Value::Float64(value)) if !row.deleted => {
3024 Some((*value, row.row_id.0))
3025 }
3026 _ => None,
3027 })
3028 .collect();
3029 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
3030 }
3031 ref ty => {
3032 return Err(MongrelError::Schema(format!(
3033 "LearnedRange index {} does not support {ty:?}",
3034 definition.name
3035 )));
3036 }
3037 };
3038 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
3039 }
3040 }
3041 .ok_or_else(|| {
3042 MongrelError::Other(format!(
3043 "failed to construct hidden index {}",
3044 definition.name
3045 ))
3046 })
3047 }
3048
3049 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
3050 let dir = dir.as_ref().to_path_buf();
3051 std::fs::create_dir_all(&dir)?;
3054 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3055 let pinned = root.io_path()?;
3056 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
3057 ctx.root_guard = Some(root.clone());
3058 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3059 root.finalize_deferred_sync_shallow()?;
3065 Ok(table)
3066 }
3067
3068 pub fn create_encrypted(
3079 dir: impl AsRef<Path>,
3080 schema: Schema,
3081 table_id: u64,
3082 passphrase: &str,
3083 ) -> Result<Self> {
3084 let dir = dir.as_ref().to_path_buf();
3085 std::fs::create_dir_all(&dir)?;
3086 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3087 root.create_directory_all(META_DIR)?;
3088 let salt = crate::encryption::random_salt()?;
3089 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
3090 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
3091 let pinned = root.io_path()?;
3092 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3093 ctx.root_guard = Some(root.clone());
3094 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3095 root.finalize_deferred_sync()?;
3096 Ok(table)
3097 }
3098
3099 pub fn create_with_key(
3104 dir: impl AsRef<Path>,
3105 schema: Schema,
3106 table_id: u64,
3107 key: &[u8],
3108 ) -> Result<Self> {
3109 let dir = dir.as_ref().to_path_buf();
3110 std::fs::create_dir_all(&dir)?;
3111 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3112 root.create_directory_all(META_DIR)?;
3113 let salt = crate::encryption::random_salt()?;
3114 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
3115 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
3116 let pinned = root.io_path()?;
3117 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3118 ctx.root_guard = Some(root.clone());
3119 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3120 root.finalize_deferred_sync()?;
3121 Ok(table)
3122 }
3123
3124 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
3126 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3127 let salt = read_table_encryption_salt_root(&root)?;
3128 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
3129 let pinned = root.io_path()?;
3130 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3131 ctx.root_guard = Some(root);
3132 Self::open_in(&pinned, ctx)
3133 }
3134
3135 pub(crate) fn create_in(
3136 dir: impl AsRef<Path>,
3137 schema: Schema,
3138 table_id: u64,
3139 ctx: SharedCtx,
3140 ) -> Result<Self> {
3141 schema.validate_auto_increment()?;
3142 schema.validate_defaults()?;
3143 schema.validate_ai()?;
3144 for index in &schema.indexes {
3145 index.validate_options()?;
3146 }
3147 let dir = dir.as_ref().to_path_buf();
3148 let runs_root = match ctx.root_guard.as_ref() {
3149 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
3150 None => {
3151 crate::durable_file::create_directory_all(&dir)?;
3152 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
3153 None
3154 }
3155 };
3156 match ctx.root_guard.as_deref() {
3157 Some(root) => write_schema_durable(root, &schema)?,
3158 None => write_schema(&dir, &schema)?,
3159 }
3160 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
3161 let (wal, current_txn_id) = match ctx.shared.clone() {
3164 Some(s) => (WalSink::Shared(s), 0),
3165 None => {
3166 let pinned_wal_root = match ctx.root_guard.as_deref() {
3167 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
3168 None => None,
3169 };
3170 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
3171 root.io_path()?
3172 } else {
3173 let wal_dir = dir.join(WAL_DIR);
3174 std::fs::create_dir_all(&wal_dir)?;
3175 wal_dir
3176 };
3177 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
3178 (Some(root), Some(dk)) => {
3179 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
3180 }
3181 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
3182 (None, Some(dk)) => Wal::create_with_cipher(
3183 wal_dir.join("seg-000000.wal"),
3184 Epoch(0),
3185 Some(make_cipher(dk)),
3186 0,
3187 )?,
3188 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
3189 };
3190 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3191 (WalSink::Private(w), 1)
3192 }
3193 };
3194 let mut manifest = Manifest::new(table_id, schema.schema_id);
3195 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3200 match ctx.root_guard.as_deref() {
3201 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
3202 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
3203 }
3204 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
3205 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3206 let auto_inc = resolve_auto_inc(&schema);
3207 let rcache_dir = dir.join(RCACHE_DIR);
3208 let initial_view = ReadGeneration::empty(&schema);
3209 Ok(Self {
3210 dir,
3211 _root_guard: ctx.root_guard,
3212 runs_root,
3213 idx_root: None,
3214 table_id,
3215 name: ctx.table_name.unwrap_or_default(),
3216 auth: ctx.auth,
3217 read_only: ctx.read_only,
3218 durable_commit_failed: false,
3219 wal,
3220 memtable: Memtable::new(),
3221 mutable_run: MutableRun::new(),
3222 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3223 compaction_zstd_level: 3,
3224 allocator: RowIdAllocator::new(0),
3225 epoch: ctx.epoch,
3226 data_generation: 0,
3227 schema,
3228 hot: HotIndex::new(),
3229 kek: ctx.kek,
3230 column_keys,
3231 run_refs: Vec::new(),
3232 retiring: Vec::new(),
3233 next_run_id: 1,
3234 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3235 current_txn_id,
3236 pending_private_mutations: false,
3237 bitmap,
3238 ann,
3239 fm,
3240 sparse,
3241 minhash,
3242 learned_range: Arc::new(HashMap::new()),
3243 pk_by_row: ReversePkMap::new(),
3244 pinned: BTreeMap::new(),
3245 live_count: 0,
3246 reservoir: crate::reservoir::Reservoir::default(),
3247 reservoir_complete: true,
3248 had_deletes: false,
3249 recent_delete_preimages: HashMap::new(),
3250 run_row_id_ranges: HashMap::new(),
3251 agg_cache: Arc::new(HashMap::new()),
3252 global_idx_epoch: 0,
3253 indexes_complete: true,
3254 index_build_policy: IndexBuildPolicy::default(),
3255 pk_by_row_complete: false,
3256 flushed_epoch: 0,
3257 page_cache: ctx.page_cache,
3258 decoded_cache: ctx.decoded_cache,
3259 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3260 snapshots: ctx.snapshots,
3261 commit_lock: ctx.commit_lock,
3262 result_cache: {
3267 let cache = ResultCache::new()
3268 .with_dir(rcache_dir.clone())
3269 .with_cache_dek(cache_dek.clone());
3270 let metrics = Arc::new(LookupMetrics::default());
3274 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3275 if let Ok((writer, handle, completion_rx)) = spawn_persistent_cache_worker(
3276 rcache_dir.clone(),
3277 cache_dek.clone(),
3278 (*metrics).clone(),
3279 ) {
3280 cache_arc
3281 .lock()
3282 .install_persistent_writer(writer, handle, completion_rx);
3283 }
3284 cache_arc
3285 },
3286 pending_delete_rids: roaring::RoaringBitmap::new(),
3287 pending_put_cols: std::collections::HashSet::new(),
3288 pending_rows: Vec::new(),
3289 pending_rows_auto_inc: Vec::new(),
3290 pending_dels: Vec::new(),
3291 pending_truncate: None,
3292 wal_dek,
3293 auto_inc,
3294 ttl: None,
3295 pins: Arc::new(crate::retention::PinRegistry::new()),
3296 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3297 read_generation_pin: None,
3298 lookup_metrics: LookupMetrics::default(),
3299 run_lookup: RunLookupState::default(),
3300 })
3301 }
3302
3303 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
3307 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3308 let pinned = root.io_path()?;
3309 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
3310 ctx.root_guard = Some(root);
3311 Self::open_in(&pinned, ctx)
3312 }
3313
3314 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
3317 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3318 let salt = read_table_encryption_salt_root(&root)?;
3319 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
3320 let pinned = root.io_path()?;
3321 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3322 ctx.root_guard = Some(root);
3323 let t = Self::open_in(&pinned, ctx)?;
3324 Ok(t)
3325 }
3326
3327 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
3328 let dir = dir.as_ref().to_path_buf();
3329 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3330 let mut manifest = match ctx.root_guard.as_ref() {
3331 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
3332 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
3333 };
3334 let schema: Schema = match ctx.root_guard.as_ref() {
3335 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
3336 None => read_schema(&dir)?,
3337 };
3338 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
3344 let runs_root = match ctx.root_guard.as_ref() {
3345 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
3346 None => None,
3347 };
3348 let idx_root = match ctx.root_guard.as_ref() {
3349 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
3350 Ok(root) => Some(Arc::new(root)),
3351 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
3352 Err(error) => return Err(error.into()),
3353 },
3354 None => None,
3355 };
3356 schema.validate_auto_increment()?;
3357 schema.validate_defaults()?;
3358 schema.validate_ai()?;
3359 for index in &schema.indexes {
3360 index.validate_options()?;
3361 }
3362 let replay_epoch = Epoch(manifest.current_epoch);
3363 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
3364 let private_replayed = if ctx.shared.is_none() {
3365 match latest_wal_segment(&dir.join(WAL_DIR))? {
3366 Some(path) => {
3367 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
3368 crate::wal::replay_with_cipher(path, cipher)?
3369 }
3370 None => Vec::new(),
3371 }
3372 } else {
3373 Vec::new()
3374 };
3375 if ctx.shared.is_none() {
3376 preflight_standalone_open(
3377 &dir,
3378 runs_root.as_deref(),
3379 idx_root.as_deref(),
3380 &manifest,
3381 &schema,
3382 &private_replayed,
3383 ctx.kek.clone(),
3384 )?;
3385 }
3386 let next_run_id = derive_next_run_id(
3387 &dir,
3388 runs_root.as_deref(),
3389 &manifest.runs,
3390 &manifest.retiring,
3391 )?;
3392 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
3396 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
3397 None => {
3398 let replayed = private_replayed;
3399 let wal_dir = dir.join(WAL_DIR);
3405 crate::durable_file::create_directory_all(&wal_dir)?;
3406 let segment = next_wal_segment(&wal_dir)?;
3407 let segment_no = wal_segment_number(&segment).unwrap_or(0);
3408 let temporary = wal_dir.join(format!(
3409 ".recovery-{}-{}-{segment_no:06}.tmp",
3410 std::process::id(),
3411 std::time::SystemTime::now()
3412 .duration_since(std::time::UNIX_EPOCH)
3413 .unwrap_or_default()
3414 .as_nanos()
3415 ));
3416 let mut w = Wal::create_with_cipher(
3417 &temporary,
3418 replay_epoch,
3419 wal_dek.as_ref().map(|dk| make_cipher(dk)),
3420 segment_no,
3421 )?;
3422 for record in &replayed {
3423 w.append_txn(record.txn_id, record.op.clone())?;
3424 }
3425 let mut w = w.publish_as(segment)?;
3426 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3427 let next_txn_id = replayed
3428 .iter()
3429 .map(|record| record.txn_id)
3430 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
3431 .max()
3432 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
3433 .unwrap_or(1);
3434 (WalSink::Private(w), replayed, next_txn_id)
3435 }
3436 };
3437
3438 let mut memtable = Memtable::new();
3439 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
3440 let persisted_epoch = manifest.current_epoch;
3441 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
3448 s.next = manifest.auto_inc_next;
3449 s.seeded = manifest.auto_inc_next > 0;
3450 s
3451 });
3452
3453 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
3460 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
3461 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
3462 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
3463 std::collections::BTreeMap::new();
3464 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
3465 let mut recovered_epoch = manifest.current_epoch;
3466 let mut recovered_manifest_dirty = schema_manifest_repair;
3467 let mut saw_delete = false;
3468 for record in replayed {
3469 let txn_id = record.txn_id;
3470 match record.op {
3471 Op::Put { rows, .. } => {
3472 let rows: Vec<Row> = bincode::deserialize(&rows)?;
3473 for row in &rows {
3474 allocator.advance_to(row.row_id)?;
3475 if let Some(ai) = auto_inc.as_mut() {
3476 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3477 let next = n.checked_add(1).ok_or_else(|| {
3478 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3479 })?;
3480 if next > ai.next {
3481 ai.next = next;
3482 }
3483 }
3484 }
3485 }
3486 staged_puts.entry(txn_id).or_default().extend(rows);
3487 }
3488 Op::Delete { row_ids, .. } => {
3489 staged_deletes.entry(txn_id).or_default().extend(row_ids);
3490 }
3491 Op::TxnCommit { epoch, .. } => {
3492 let commit_epoch = Epoch(epoch);
3493 recovered_epoch = recovered_epoch.max(epoch);
3494 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
3495 memtable = Memtable::new();
3496 replayed_puts.clear();
3497 replayed_deletes.clear();
3498 manifest.runs.clear();
3499 manifest.retiring.clear();
3500 manifest.live_count = 0;
3501 manifest.global_idx_epoch = 0;
3502 manifest.current_epoch = manifest.current_epoch.max(epoch);
3503 recovered_manifest_dirty = true;
3504 saw_delete = true;
3505 }
3506 if let Some(puts) = staged_puts.remove(&txn_id) {
3507 if commit_epoch.0 > manifest.flushed_epoch {
3508 for row in &puts {
3509 memtable.upsert(row.clone());
3510 }
3511 replayed_puts.entry(commit_epoch).or_default().extend(puts);
3512 }
3513 }
3514 if let Some(dels) = staged_deletes.remove(&txn_id) {
3515 saw_delete = true;
3516 if commit_epoch.0 > manifest.flushed_epoch {
3517 for rid in dels {
3518 memtable.tombstone(rid, commit_epoch);
3519 replayed_deletes.push((rid, commit_epoch));
3520 }
3521 }
3522 }
3523 }
3524 Op::TxnAbort => {
3525 staged_puts.remove(&txn_id);
3526 staged_deletes.remove(&txn_id);
3527 staged_truncates.remove(&txn_id);
3528 }
3529 Op::TruncateTable { .. } => {
3530 staged_puts.remove(&txn_id);
3531 staged_deletes.remove(&txn_id);
3532 staged_truncates.insert(txn_id);
3533 }
3534 Op::ExternalTableState { .. }
3535 | Op::Flush { .. }
3536 | Op::Ddl(_)
3537 | Op::BeforeImage { .. }
3538 | Op::CommitTimestamp { .. }
3539 | Op::SpilledRows { .. } => {}
3540 }
3541 }
3542
3543 let rcache_dir = dir.join(RCACHE_DIR);
3544 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3545 let initial_view = ReadGeneration::empty(&schema);
3546 let mut db = Self {
3547 dir,
3548 _root_guard: ctx.root_guard,
3549 runs_root,
3550 idx_root,
3551 table_id: manifest.table_id,
3552 name: ctx.table_name.unwrap_or_default(),
3553 auth: ctx.auth,
3554 read_only: ctx.read_only,
3555 durable_commit_failed: false,
3556 wal,
3557 memtable,
3558 mutable_run: MutableRun::new(),
3559 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3560 compaction_zstd_level: 3,
3561 allocator,
3562 epoch: ctx.epoch,
3563 data_generation: persisted_epoch,
3564 schema,
3565 hot: HotIndex::new(),
3566 kek: ctx.kek,
3567 column_keys,
3568 run_refs: manifest.runs.clone(),
3569 retiring: manifest.retiring.clone(),
3570 next_run_id,
3571 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3572 current_txn_id,
3573 pending_private_mutations: false,
3574 bitmap: HashMap::new(),
3575 ann: HashMap::new(),
3576 fm: HashMap::new(),
3577 sparse: HashMap::new(),
3578 minhash: HashMap::new(),
3579 learned_range: Arc::new(HashMap::new()),
3580 pk_by_row: ReversePkMap::new(),
3581 pinned: BTreeMap::new(),
3582 live_count: manifest.live_count,
3583 reservoir: crate::reservoir::Reservoir::default(),
3584 reservoir_complete: false,
3585 had_deletes: saw_delete
3586 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
3587 != manifest.live_count,
3588 recent_delete_preimages: HashMap::new(),
3589 run_row_id_ranges: HashMap::new(),
3590 agg_cache: Arc::new(HashMap::new()),
3591 global_idx_epoch: manifest.global_idx_epoch,
3592 indexes_complete: true,
3593 index_build_policy: IndexBuildPolicy::default(),
3594 pk_by_row_complete: false,
3595 flushed_epoch: manifest.flushed_epoch,
3596 page_cache: ctx.page_cache,
3597 decoded_cache: ctx.decoded_cache,
3598 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3599 snapshots: ctx.snapshots,
3600 commit_lock: ctx.commit_lock,
3601 result_cache: {
3602 let cache = ResultCache::new()
3603 .with_dir(rcache_dir.clone())
3604 .with_cache_dek(cache_dek.clone());
3605 let metrics = LookupMetrics::default();
3606 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3607 if let Ok((writer, handle, completion_rx)) =
3608 spawn_persistent_cache_worker(rcache_dir.clone(), cache_dek.clone(), metrics)
3609 {
3610 cache_arc
3611 .lock()
3612 .install_persistent_writer(writer, handle, completion_rx);
3613 }
3614 cache_arc
3615 },
3616 pending_delete_rids: roaring::RoaringBitmap::new(),
3617 pending_put_cols: std::collections::HashSet::new(),
3618 pending_rows: Vec::new(),
3619 pending_rows_auto_inc: Vec::new(),
3620 pending_dels: Vec::new(),
3621 pending_truncate: None,
3622 wal_dek,
3623 auto_inc,
3624 ttl: manifest.ttl,
3625 pins: Arc::new(crate::retention::PinRegistry::new()),
3626 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3627 read_generation_pin: None,
3628 lookup_metrics: LookupMetrics::default(),
3629 run_lookup: RunLookupState::default(),
3630 };
3631
3632 db.epoch.advance_recovered(Epoch(recovered_epoch));
3635
3636 let checkpoint = match db.idx_root.as_deref() {
3641 Some(root) => {
3642 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
3643 }
3644 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
3645 };
3646 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
3647 c.epoch_built == manifest.global_idx_epoch
3648 && manifest.global_idx_epoch > 0
3649 && manifest
3650 .runs
3651 .iter()
3652 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
3653 });
3654 if let Some(loaded) = checkpoint {
3655 if checkpoint_valid {
3656 db.hot = loaded.hot;
3657 db.bitmap = loaded.bitmap;
3658 db.ann = loaded.ann;
3659 db.fm = loaded.fm;
3660 db.sparse = loaded.sparse;
3661 db.minhash = loaded.minhash;
3662 db.learned_range = Arc::new(loaded.learned_range);
3663 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
3668 for (cid, idx) in bitmap0 {
3669 db.bitmap.entry(cid).or_insert(idx);
3670 }
3671 for (cid, idx) in ann0 {
3672 db.ann.entry(cid).or_insert(idx);
3673 }
3674 for (cid, idx) in fm0 {
3675 db.fm.entry(cid).or_insert(idx);
3676 }
3677 for (cid, idx) in sparse0 {
3678 db.sparse.entry(cid).or_insert(idx);
3679 }
3680 for (cid, idx) in minhash0 {
3681 db.minhash.entry(cid).or_insert(idx);
3682 }
3683 }
3686 }
3687 if !checkpoint_valid {
3688 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
3689 db.bitmap = bitmap;
3690 db.ann = ann;
3691 db.fm = fm;
3692 db.sparse = sparse;
3693 db.minhash = minhash;
3694 db.rebuild_indexes_from_runs()?;
3695 db.build_learned_ranges()?;
3696 }
3697
3698 for (epoch, group) in replayed_puts {
3703 let (losers, winner_pks) = db.partition_pk_winners(&group);
3704 for (key, &row_id) in &winner_pks {
3705 if let Some(old_rid) = db.hot.get(key) {
3706 if old_rid != row_id {
3707 db.tombstone_row(old_rid, epoch, None, false);
3708 }
3709 }
3710 }
3711 for &loser_rid in &losers {
3712 db.tombstone_row(loser_rid, epoch, None, false);
3713 }
3714 for (key, row_id) in winner_pks {
3715 db.insert_hot_pk(key, row_id);
3716 }
3717 if db.schema.primary_key().is_none() {
3718 for r in &group {
3719 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3720 }
3721 }
3722 for r in &group {
3723 if !losers.contains(&r.row_id) {
3724 db.index_row(r);
3725 }
3726 }
3727 }
3728 for (rid, epoch) in &replayed_deletes {
3732 db.remove_hot_for_row(*rid, *epoch);
3733 }
3734
3735 if recovered_manifest_dirty {
3736 let rows = db.visible_rows(Snapshot::unbounded())?;
3737 db.live_count = rows.len() as u64;
3738 db.persist_manifest(Epoch(recovered_epoch))?;
3739 }
3740
3741 let identity = crate::result_cache::PersistentCacheIdentity {
3748 table_id: db.table_id,
3749 schema_id: db.schema.schema_id,
3750 logical_generation: db.epoch.visible().0,
3751 };
3752 db.result_cache.lock().load_persistent(identity);
3753 db.refresh_run_row_id_ranges();
3755 db.load_run_lookup_directory();
3760 Ok(db)
3761 }
3762
3763 fn load_run_lookup_directory(&mut self) {
3768 let active_runs = self.run_refs.clone();
3769 let schema_id = self.schema.schema_id;
3770 let run_generation = active_runs.len() as u64;
3771 let index_generation = self.global_idx_epoch;
3772 let fingerprint = crate::run_lookup::compute_fingerprint(
3773 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
3774 schema_id,
3775 index_generation,
3776 run_generation,
3777 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
3778 );
3779 let runs_dir = match self.runs_root.as_deref() {
3781 Some(root) => match root.io_path() {
3782 Ok(p) => p,
3783 Err(_) => self.dir.join(RUNS_DIR),
3784 },
3785 None => self.dir.join(RUNS_DIR),
3786 };
3787 match crate::run_lookup::RunLookupDirectory::read_checkpoint_sharded(&runs_dir, fingerprint)
3788 {
3789 Ok(dir) => {
3790 self.run_lookup = RunLookupState {
3791 directory: Some(std::sync::Arc::new(dir)),
3792 fingerprint,
3793 complete: true,
3794 };
3795 }
3796 Err(error) => {
3797 self.lookup_metrics
3798 .directory_incomplete
3799 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3800 self.lookup_metrics
3801 .directory_lookup_fallback
3802 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3803 self.run_lookup = RunLookupState {
3804 directory: None,
3805 fingerprint,
3806 complete: false,
3807 };
3808 let _ = error;
3812 }
3813 }
3814 }
3815
3816 fn ensure_reservoir_complete(&mut self) -> Result<()> {
3822 if self.reservoir_complete {
3823 return Ok(());
3824 }
3825 self.rebuild_reservoir()?;
3826 self.reservoir_complete = true;
3827 Ok(())
3828 }
3829
3830 fn rebuild_reservoir(&mut self) -> Result<()> {
3833 let snap = self.snapshot();
3834 let rows = self.visible_rows(snap)?;
3835 self.reservoir.reset();
3836 for r in rows {
3837 self.reservoir.offer(r.row_id.0);
3838 }
3839 Ok(())
3840 }
3841
3842 pub fn rebuild_indexes(&mut self) -> Result<()> {
3846 self.rebuild_indexes_from_runs_inner(None)
3847 }
3848
3849 pub fn __force_hot_map_for_test(&mut self, pk_bytes: &[u8], row_id: RowId) {
3860 self.hot.insert(pk_bytes.to_vec(), row_id);
3861 }
3862
3863 pub fn hot_for_test_remove(&mut self, pk_bytes: &[u8]) {
3866 self.hot.remove(pk_bytes);
3867 }
3868
3869 pub fn set_indexes_incomplete_for_test(&mut self) {
3872 self.indexes_complete = false;
3873 }
3874
3875 pub fn bump_hot_checkpoint_rejected_for_test(&self) {
3879 self.lookup_metrics
3880 .hot_checkpoint_rejected_total
3881 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3882 }
3883
3884 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
3885 self.rebuild_indexes_from_runs_inner(None)
3886 }
3887
3888 fn pk_equality_fallback(
3915 &self,
3916 pk_column_id: u16,
3917 lookup: &[u8],
3918 snapshot: Snapshot,
3919 ) -> Result<(RowIdSet, crate::trace::HotFallbackReason)> {
3920 let mut tombstone_hit = false;
3921 let mut overlay_versions = 0u64;
3922 let now_nanos = unix_nanos_now();
3923 for row in self.memtable.visible_versions_at(snapshot) {
3925 overlay_versions += 1;
3926 self.lookup_metrics
3927 .hot_lookup_fallback_overlay_rows
3928 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3929 if row.deleted {
3930 tombstone_hit = true;
3931 continue;
3932 }
3933 if self.row_expired_at(&row, now_nanos) {
3934 continue;
3935 }
3936 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3937 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3938 self.lookup_metrics
3939 .hot_fallback_overlay_versions_total
3940 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3941 return Ok((
3942 RowIdSet::one(row.row_id.0),
3943 crate::trace::HotFallbackReason::MissingMapping,
3944 ));
3945 }
3946 }
3947 }
3948 for row in self.mutable_run.visible_versions_at(snapshot) {
3949 overlay_versions += 1;
3950 self.lookup_metrics
3951 .hot_lookup_fallback_overlay_rows
3952 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3953 if row.deleted {
3954 tombstone_hit = true;
3955 continue;
3956 }
3957 if self.row_expired_at(&row, now_nanos) {
3958 continue;
3959 }
3960 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3961 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3962 self.lookup_metrics
3963 .hot_fallback_overlay_versions_total
3964 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3965 return Ok((
3966 RowIdSet::one(row.row_id.0),
3967 crate::trace::HotFallbackReason::MissingMapping,
3968 ));
3969 }
3970 }
3971 }
3972 if lookup.len() == 8 {
3975 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
3976 let n = i64::from_be_bytes(arr);
3977 let result = self.range_scan_i64(pk_column_id, n, n, snapshot)?;
3978 self.lookup_metrics
3979 .hot_lookup_fallback_runs
3980 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3981 self.lookup_metrics
3985 .hot_fallback_runs_considered_total
3986 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3987 self.lookup_metrics
3988 .hot_fallback_overlay_versions_total
3989 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3990 let reason = if result.is_empty() {
3991 if tombstone_hit {
3992 crate::trace::HotFallbackReason::Tombstone
3993 } else {
3994 crate::trace::HotFallbackReason::MissingMapping
3995 }
3996 } else {
3997 crate::trace::HotFallbackReason::MissingMapping
3998 };
3999 return Ok((result, reason));
4000 }
4001 }
4002 let mut found: std::collections::BTreeSet<u64> = std::collections::BTreeSet::new();
4005 let overlay = self.overlay_rid_set(snapshot);
4006 let mut runs_considered = 0u64;
4007 for rr in &self.run_refs {
4008 runs_considered += 1;
4009 self.lookup_metrics
4010 .hot_lookup_fallback_runs
4011 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4012 let mut reader = self.open_reader(rr.run_id)?;
4013 for row in reader.visible_versions_at(snapshot)? {
4014 if overlay.contains(&row.row_id.0) || row.deleted {
4015 if row.deleted {
4016 tombstone_hit = true;
4017 }
4018 continue;
4019 }
4020 if self.row_expired_at(&row, now_nanos) {
4021 continue;
4022 }
4023 if let Some(pk_val) = row.columns.get(&pk_column_id) {
4024 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
4025 found.insert(row.row_id.0);
4026 }
4027 }
4028 }
4029 }
4030 let reason = if tombstone_hit {
4031 crate::trace::HotFallbackReason::Tombstone
4032 } else {
4033 crate::trace::HotFallbackReason::MissingMapping
4036 };
4037 self.lookup_metrics
4038 .hot_fallback_overlay_versions_total
4039 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
4040 self.lookup_metrics
4041 .hot_fallback_runs_considered_total
4042 .fetch_add(runs_considered, std::sync::atomic::Ordering::Relaxed);
4043 Ok((RowIdSet::from_unsorted(found.into_iter().collect()), reason))
4044 }
4045
4046 fn record_hot_fallback_reason(&self, reason: crate::trace::HotFallbackReason) {
4050 let idx = hot_fallback_reason_index(reason);
4051 self.lookup_metrics.hot_fallback_reasons[idx]
4052 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4053 self.lookup_metrics
4054 .hot_lookup_fallback
4055 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
4056 crate::trace::QueryTrace::record(|t| {
4057 t.hot_lookup_attempted = true;
4058 t.hot_lookup_hit = false;
4059 t.hot_fallback_reason = Some(reason.as_str());
4060 });
4061 }
4062
4063 fn rebuild_indexes_from_runs_inner(
4064 &mut self,
4065 control: Option<&crate::ExecutionControl>,
4066 ) -> Result<()> {
4067 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
4070 crate::retention::PinSource::OnlineIndexBuild,
4071 self.current_epoch(),
4072 );
4073 self.hot = HotIndex::new();
4074 self.pk_by_row.clear();
4075 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
4076 self.bitmap = bitmap;
4077 self.ann = ann;
4078 self.fm = fm;
4079 self.sparse = sparse;
4080 self.minhash = minhash;
4081 let snapshot = Snapshot::unbounded();
4096 let ttl_now = unix_nanos_now();
4097 let mut scanned = 0_usize;
4098 let mut winners: HashMap<RowId, (VersionStamp, Row)> = HashMap::new();
4099 let fold =
4100 |row: Row, winners: &mut HashMap<RowId, (VersionStamp, Row)>, scanned: &mut usize| {
4101 *scanned += 1;
4102 let stamp = VersionStamp {
4103 epoch: row.committed_epoch,
4104 commit_ts: row.commit_ts,
4105 };
4106 winners
4107 .entry(row.row_id)
4108 .and_modify(|entry| {
4109 if stamp.is_newer_than(entry.0) {
4110 *entry = (stamp, row.clone());
4111 }
4112 })
4113 .or_insert((stamp, row));
4114 };
4115 for rr in self.run_refs.clone() {
4116 if let Some(control) = control {
4117 control.checkpoint()?;
4118 }
4119 let mut reader = self.open_reader(rr.run_id)?;
4120 for row in reader.visible_versions_at(snapshot)? {
4121 if scanned.is_multiple_of(256) {
4122 if let Some(control) = control {
4123 control.checkpoint()?;
4124 }
4125 }
4126 fold(row, &mut winners, &mut scanned);
4130 }
4131 }
4132 for (index, row) in self
4133 .memtable
4134 .visible_versions_at(snapshot)
4135 .into_iter()
4136 .enumerate()
4137 {
4138 if index & 255 == 0 {
4139 if let Some(control) = control {
4140 control.checkpoint()?;
4141 }
4142 }
4143 fold(row, &mut winners, &mut scanned);
4144 }
4145 for (index, row) in self
4146 .mutable_run
4147 .visible_versions_at(snapshot)
4148 .into_iter()
4149 .enumerate()
4150 {
4151 if index & 255 == 0 {
4152 if let Some(control) = control {
4153 control.checkpoint()?;
4154 }
4155 }
4156 fold(row, &mut winners, &mut scanned);
4157 }
4158 for (_rid, (stamp, row)) in winners.drain() {
4159 if row.deleted {
4160 let _ = stamp;
4163 continue;
4164 }
4165 if self.row_expired_at(&row, ttl_now) {
4166 continue;
4167 }
4168 let tok_row = self.tokenized_for_indexes(&row);
4169 index_into(
4170 &self.schema,
4171 &tok_row,
4172 &mut self.hot,
4173 &mut self.bitmap,
4174 &mut self.ann,
4175 &mut self.fm,
4176 &mut self.sparse,
4177 &mut self.minhash,
4178 );
4179 }
4180 let pin_epochs = self.active_pin_epochs_for_rebuild();
4186 if !pin_epochs.is_empty() {
4187 let current_snap = Snapshot::at(Epoch(u64::MAX));
4188 for pin_epoch in pin_epochs {
4189 if let Some(control) = control {
4190 control.checkpoint()?;
4191 }
4192 let pin_snap = Snapshot::at(pin_epoch);
4193 for rr in self.run_refs.clone() {
4194 if let Some(control) = control {
4195 control.checkpoint()?;
4196 }
4197 let mut reader = self.open_reader(rr.run_id)?;
4198 for row in reader.visible_rows(pin_epoch)? {
4199 if row.deleted || self.row_expired_at(&row, ttl_now) {
4200 continue;
4201 }
4202 if self.get(row.row_id, current_snap).is_none() {
4203 self.index_bitmap_membership_only(&row);
4204 }
4205 }
4206 }
4207 for row in self
4208 .mutable_run
4209 .visible_versions_at(pin_snap)
4210 .into_iter()
4211 .chain(self.memtable.visible_versions_at(pin_snap))
4212 {
4213 if row.deleted || self.row_expired_at(&row, ttl_now) {
4214 continue;
4215 }
4216 if self.get(row.row_id, current_snap).is_none() {
4217 self.index_bitmap_membership_only(&row);
4218 }
4219 }
4220 }
4221 }
4222 self.recent_delete_preimages.clear();
4223 self.refresh_pk_by_row_from_hot();
4224 Ok(())
4225 }
4226
4227 fn index_bitmap_membership_only(&mut self, row: &Row) {
4230 if row.deleted {
4231 return;
4232 }
4233 let columns_map: HashMap<u16, &Value> = row.columns.iter().map(|(k, v)| (*k, v)).collect();
4234 let name_to_id: HashMap<&str, u16> = self
4235 .schema
4236 .columns
4237 .iter()
4238 .map(|c| (c.name.as_str(), c.id))
4239 .collect();
4240 for idef in &self.schema.indexes {
4241 if idef.kind != crate::schema::IndexKind::Bitmap {
4242 continue;
4243 }
4244 if let Some(pred) = &idef.predicate {
4245 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4246 continue;
4247 }
4248 }
4249 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
4250 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4251 b.insert(key, row.row_id);
4252 }
4253 }
4254 }
4255 }
4256
4257 fn refresh_pk_by_row_from_hot(&mut self) {
4258 self.pk_by_row_complete = true;
4259 if self.schema.primary_key().is_none() {
4260 self.pk_by_row.clear();
4261 return;
4262 }
4263 self.pk_by_row = ReversePkMap::from_entries(
4269 self.hot
4270 .entries()
4271 .into_iter()
4272 .map(|(key, row_id)| (row_id, key)),
4273 );
4274 }
4275
4276 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
4277 if self.schema.primary_key().is_some() {
4278 self.pk_by_row.insert(row_id, key.clone());
4279 }
4280 self.hot.insert(key, row_id);
4281 }
4282
4283 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
4287 self.build_learned_ranges_inner(None)
4288 }
4289
4290 fn build_learned_ranges_inner(
4291 &mut self,
4292 control: Option<&crate::ExecutionControl>,
4293 ) -> Result<()> {
4294 self.learned_range = Arc::new(HashMap::new());
4295 if self.run_refs.len() != 1 {
4296 return Ok(());
4297 }
4298 let cols: Vec<(u16, usize)> = self
4299 .schema
4300 .indexes
4301 .iter()
4302 .filter(|i| i.kind == IndexKind::LearnedRange)
4303 .map(|i| {
4304 (
4305 i.column_id,
4306 i.options
4307 .learned_range
4308 .as_ref()
4309 .map(|options| options.epsilon)
4310 .unwrap_or(16),
4311 )
4312 })
4313 .collect();
4314 if cols.is_empty() {
4315 return Ok(());
4316 }
4317 let snapshot = Snapshot::unbounded();
4328 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4329 let (visible_positions, visible_rids) = reader.visible_positions_with_rids_at(snapshot)?;
4330 let row_ids: Vec<u64> = visible_rids.iter().map(|r| *r as u64).collect();
4331 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
4332 if column_index % 256 == 0 {
4333 if let Some(control) = control {
4334 control.checkpoint()?;
4335 }
4336 }
4337 let ty = self
4338 .schema
4339 .columns
4340 .iter()
4341 .find(|c| c.id == cid)
4342 .map(|c| c.ty.clone())
4343 .unwrap_or(TypeId::Int64);
4344 match ty {
4345 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
4346 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
4347 let pairs: Vec<(i64, u64)> = visible_positions
4348 .iter()
4349 .zip(row_ids.iter())
4350 .map(|(&p, &r)| (data[p], r))
4351 .collect();
4352 Arc::make_mut(&mut self.learned_range).insert(
4353 cid,
4354 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
4355 );
4356 }
4357 }
4358 TypeId::Float64 => {
4359 if let columnar::NativeColumn::Float64 { data, .. } =
4360 reader.column_native(cid)?
4361 {
4362 let pairs: Vec<(f64, u64)> = visible_positions
4363 .iter()
4364 .zip(row_ids.iter())
4365 .map(|(&p, &r)| (data[p], r))
4366 .collect();
4367 Arc::make_mut(&mut self.learned_range).insert(
4368 cid,
4369 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
4370 );
4371 }
4372 }
4373 _ => {}
4374 }
4375 }
4376 Ok(())
4377 }
4378
4379 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
4386 if self.indexes_complete {
4387 crate::trace::QueryTrace::record(|t| {
4388 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4389 });
4390 return Ok(());
4391 }
4392 crate::trace::QueryTrace::record(|t| {
4393 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4394 });
4395 self.rebuild_indexes_from_runs()?;
4396 self.build_learned_ranges()?;
4397 self.indexes_complete = true;
4398 let epoch = self.current_epoch();
4399 self.checkpoint_indexes(epoch);
4400 Ok(())
4401 }
4402
4403 #[doc(hidden)]
4406 pub fn ensure_indexes_complete_controlled<F>(
4407 &mut self,
4408 control: &crate::ExecutionControl,
4409 before_publish: F,
4410 ) -> Result<bool>
4411 where
4412 F: FnOnce() -> bool,
4413 {
4414 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
4415 .map(|(changed, _)| changed)
4416 }
4417
4418 #[doc(hidden)]
4421 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
4422 &mut self,
4423 control: &crate::ExecutionControl,
4424 before_publish: F,
4425 ) -> Result<(bool, Option<MaintenanceReceipt>)>
4426 where
4427 F: FnOnce() -> bool,
4428 {
4429 if self.indexes_complete {
4430 crate::trace::QueryTrace::record(|trace| {
4431 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4432 });
4433 return Ok((false, None));
4434 }
4435 crate::trace::QueryTrace::record(|trace| {
4436 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4437 });
4438 control.checkpoint()?;
4439 let maintenance_epoch = self.current_epoch();
4440 self.rebuild_indexes_from_runs_inner(Some(control))?;
4441 self.build_learned_ranges_inner(Some(control))?;
4442 control.checkpoint()?;
4443 if !before_publish() {
4444 return Err(MongrelError::Cancelled);
4445 }
4446 self.indexes_complete = true;
4447 self.checkpoint_indexes(maintenance_epoch);
4448 Ok((
4449 true,
4450 Some(MaintenanceReceipt {
4451 epoch: maintenance_epoch,
4452 }),
4453 ))
4454 }
4455
4456 fn pending_epoch(&self) -> Epoch {
4457 Epoch(self.epoch.visible().0 + 1)
4458 }
4459
4460 fn is_shared(&self) -> bool {
4463 matches!(self.wal, WalSink::Shared(_))
4464 }
4465
4466 fn ensure_txn_id(&mut self) -> Result<u64> {
4470 if self.current_txn_id == 0 {
4471 let id = match &self.wal {
4472 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
4473 WalSink::Private(_) => {
4474 return Err(MongrelError::Full(
4475 "standalone transaction id namespace exhausted".into(),
4476 ))
4477 }
4478 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4479 };
4480 self.current_txn_id = id;
4481 }
4482 Ok(self.current_txn_id)
4483 }
4484
4485 fn wal_append_data(&mut self, op: Op) -> Result<()> {
4488 self.ensure_writable()?;
4489 let txn_id = self.ensure_txn_id()?;
4490 let table_id = self.table_id;
4491 match &mut self.wal {
4492 WalSink::Private(w) => {
4493 w.append_txn(txn_id, op)?;
4494 self.pending_private_mutations = true;
4495 }
4496 WalSink::Shared(s) => {
4497 s.wal.lock().append(txn_id, table_id, op)?;
4498 }
4499 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4500 }
4501 Ok(())
4502 }
4503
4504 fn ensure_writable(&self) -> Result<()> {
4505 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
4506 return Err(MongrelError::ReadOnlyReplica);
4507 }
4508 if self.durable_commit_failed {
4509 return Err(MongrelError::Other(
4510 "table poisoned by post-commit failure; reopen required".into(),
4511 ));
4512 }
4513 Ok(())
4514 }
4515
4516 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
4527 match &self.auth {
4528 Some(checker) => checker.check(&self.name, perm),
4529 None => Ok(()),
4530 }
4531 }
4532 pub fn require_select(&self) -> Result<()> {
4537 self.require(crate::auth_state::RequiredPermission::Select)
4538 }
4539 fn require_insert(&self) -> Result<()> {
4540 self.require(crate::auth_state::RequiredPermission::Insert)
4541 }
4542 #[allow(dead_code)]
4546 fn require_update(&self) -> Result<()> {
4547 self.require(crate::auth_state::RequiredPermission::Update)
4548 }
4549 fn require_delete(&self) -> Result<()> {
4550 self.require(crate::auth_state::RequiredPermission::Delete)
4551 }
4552
4553 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
4556 self.require_insert()?;
4557 Ok(self.put_returning(columns)?.0)
4558 }
4559
4560 pub fn put_returning(
4565 &mut self,
4566 mut columns: Vec<(u16, Value)>,
4567 ) -> Result<(RowId, Option<i64>)> {
4568 self.require_insert()?;
4569 let assigned = self.fill_auto_inc(&mut columns)?;
4570 self.apply_defaults(&mut columns)?;
4571 self.schema.validate_values(&columns)?;
4572 let row_id = if self.schema.clustered {
4577 self.derive_clustered_row_id(&columns)?
4578 } else {
4579 self.allocator.alloc()?
4580 };
4581 let epoch = self.pending_epoch();
4582 let mut row = Row::new(row_id, epoch);
4583 for (col_id, val) in columns {
4584 row.columns.insert(col_id, val);
4585 }
4586 self.commit_rows(vec![row], assigned.is_some())?;
4587 Ok((row_id, assigned))
4588 }
4589
4590 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
4593 self.require_insert()?;
4594 Ok(self
4595 .put_batch_returning(batch)?
4596 .into_iter()
4597 .map(|(r, _)| r)
4598 .collect())
4599 }
4600
4601 pub fn put_batch_returning(
4604 &mut self,
4605 batch: Vec<Vec<(u16, Value)>>,
4606 ) -> Result<Vec<(RowId, Option<i64>)>> {
4607 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
4608 for mut cols in batch {
4609 let assigned = self.fill_auto_inc(&mut cols)?;
4610 self.apply_defaults(&mut cols)?;
4611 filled.push((cols, assigned));
4612 }
4613 for (cols, _) in &filled {
4614 self.schema.validate_values(cols)?;
4615 }
4616 let epoch = self.pending_epoch();
4617 let mut rows = Vec::with_capacity(filled.len());
4618 let mut ids = Vec::with_capacity(filled.len());
4619 let first_row_id = if self.schema.clustered {
4620 None
4621 } else {
4622 let count = u64::try_from(filled.len())
4623 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
4624 Some(self.allocator.alloc_range(count)?.0)
4625 };
4626 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
4627 let row_id = match first_row_id {
4628 Some(first) => RowId(first + row_index as u64),
4629 None => self.derive_clustered_row_id(&cols)?,
4630 };
4631 let mut row = Row::new(row_id, epoch);
4632 for (c, v) in cols {
4633 row.columns.insert(c, v);
4634 }
4635 ids.push((row_id, assigned));
4636 rows.push(row);
4637 }
4638 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
4639 self.commit_rows(rows, all_auto_generated)?;
4640 Ok(ids)
4641 }
4642
4643 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
4649 self.ensure_writable()?;
4650 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4651 return Ok(None);
4652 };
4653 let pos = columns.iter().position(|(c, _)| *c == cid);
4654 let assigned = match pos {
4655 Some(i) => match &columns[i].1 {
4656 Value::Null => {
4657 let next = self.alloc_auto_inc_value()?;
4658 columns[i].1 = Value::Int64(next);
4659 Some(next)
4660 }
4661 Value::Int64(n) => {
4662 self.advance_auto_inc_past(*n)?;
4663 None
4664 }
4665 other => {
4666 return Err(MongrelError::InvalidArgument(format!(
4667 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
4668 other
4669 )))
4670 }
4671 },
4672 None => {
4673 let next = self.alloc_auto_inc_value()?;
4674 columns.push((cid, Value::Int64(next)));
4675 Some(next)
4676 }
4677 };
4678 Ok(assigned)
4679 }
4680
4681 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
4687 for col in &self.schema.columns {
4688 let Some(expr) = &col.default_value else {
4689 continue;
4690 };
4691 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
4693 continue;
4694 }
4695 let pos = columns.iter().position(|(c, _)| *c == col.id);
4696 let needs_default = match pos {
4697 None => true,
4698 Some(i) => matches!(columns[i].1, Value::Null),
4699 };
4700 if !needs_default {
4701 continue;
4702 }
4703 let v = match expr {
4704 crate::schema::DefaultExpr::Static(v) => v.clone(),
4705 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
4706 crate::schema::DefaultExpr::Uuid => {
4707 let mut buf = [0u8; 16];
4708 getrandom::getrandom(&mut buf)
4709 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
4710 Value::Uuid(buf)
4711 }
4712 };
4713 match pos {
4714 None => columns.push((col.id, v)),
4715 Some(i) => columns[i].1 = v,
4716 }
4717 }
4718 Ok(())
4719 }
4720
4721 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
4723 self.ensure_auto_inc_seeded()?;
4724 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4726 let v = ai.next;
4727 ai.next = ai
4728 .next
4729 .checked_add(1)
4730 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4731 Ok(v)
4732 }
4733
4734 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
4737 self.ensure_auto_inc_seeded()?;
4738 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4739 let floor = used
4740 .checked_add(1)
4741 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4742 .max(1);
4743 if ai.next < floor {
4744 ai.next = floor;
4745 }
4746 Ok(())
4747 }
4748
4749 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
4754 let needs_seed = match self.auto_inc {
4755 Some(ai) => !ai.seeded,
4756 None => return Ok(()),
4757 };
4758 if !needs_seed {
4759 return Ok(());
4760 }
4761 if self.seed_empty_auto_inc() {
4762 return Ok(());
4763 }
4764 let cid = self
4765 .auto_inc
4766 .as_ref()
4767 .expect("auto-inc column present")
4768 .column_id;
4769 let max = self.scan_max_int64(cid)?;
4770 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4771 let floor = max
4772 .checked_add(1)
4773 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4774 .max(1);
4775 if ai.next < floor {
4776 ai.next = floor;
4777 }
4778 ai.seeded = true;
4779 Ok(())
4780 }
4781
4782 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
4783 if n == 0 || self.auto_inc.is_none() {
4784 return Ok(None);
4785 }
4786 self.ensure_auto_inc_seeded()?;
4787 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4788 let start = ai.next;
4789 let count = i64::try_from(n)
4790 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
4791 ai.next = ai
4792 .next
4793 .checked_add(count)
4794 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4795 Ok(Some(start))
4796 }
4797
4798 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
4802 let mut max: i64 = 0;
4803 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
4804 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4805 if *n > max {
4806 max = *n;
4807 }
4808 }
4809 }
4810 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
4811 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4812 if *n > max {
4813 max = *n;
4814 }
4815 }
4816 }
4817 for rr in self.run_refs.clone() {
4818 let reader = self.open_reader(rr.run_id)?;
4819 if let Some(stats) = reader.column_page_stats(column_id) {
4820 for s in stats {
4821 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
4822 if n > max {
4823 max = n;
4824 }
4825 }
4826 }
4827 } else if reader.has_column(column_id) {
4828 if let columnar::NativeColumn::Int64 { data, validity } =
4829 reader.column_native_shared(column_id)?
4830 {
4831 for (i, n) in data.iter().enumerate() {
4832 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
4833 {
4834 max = *n;
4835 }
4836 }
4837 }
4838 }
4839 }
4840 Ok(max)
4841 }
4842
4843 fn seed_empty_auto_inc(&mut self) -> bool {
4844 let Some(ai) = self.auto_inc.as_mut() else {
4845 return false;
4846 };
4847 if ai.seeded || self.live_count != 0 {
4848 return false;
4849 }
4850 if ai.next < 1 {
4851 ai.next = 1;
4852 }
4853 ai.seeded = true;
4854 true
4855 }
4856
4857 fn advance_auto_inc_from_native_columns(
4858 &mut self,
4859 columns: &[(u16, columnar::NativeColumn)],
4860 n: usize,
4861 live_before: u64,
4862 ) -> Result<()> {
4863 let Some(ai) = self.auto_inc.as_mut() else {
4864 return Ok(());
4865 };
4866 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
4867 return Ok(());
4868 };
4869 let columnar::NativeColumn::Int64 { data, validity } = col else {
4870 return Err(MongrelError::InvalidArgument(format!(
4871 "AUTO_INCREMENT column {} must be Int64",
4872 ai.column_id
4873 )));
4874 };
4875 let max = if native_int64_strictly_increasing(col, n) {
4876 data.get(n.saturating_sub(1)).copied()
4877 } else {
4878 data.iter()
4879 .take(n)
4880 .enumerate()
4881 .filter_map(|(i, v)| {
4882 if validity.is_empty() || columnar::validity_bit(validity, i) {
4883 Some(*v)
4884 } else {
4885 None
4886 }
4887 })
4888 .max()
4889 };
4890 if let Some(max) = max {
4891 let floor = max
4892 .checked_add(1)
4893 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4894 .max(1);
4895 if ai.next < floor {
4896 ai.next = floor;
4897 }
4898 if ai.seeded || live_before == 0 {
4899 ai.seeded = true;
4900 }
4901 }
4902 Ok(())
4903 }
4904
4905 fn fill_auto_inc_native_columns(
4906 &mut self,
4907 columns: &mut Vec<(u16, columnar::NativeColumn)>,
4908 n: usize,
4909 ) -> Result<()> {
4910 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4911 return Ok(());
4912 };
4913 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
4914 if let Some(start) = self.alloc_auto_inc_range(n)? {
4915 columns.push((
4916 cid,
4917 columnar::NativeColumn::Int64 {
4918 data: (start..start.saturating_add(n as i64)).collect(),
4919 validity: vec![0xFF; n.div_ceil(8)],
4920 },
4921 ));
4922 }
4923 return Ok(());
4924 };
4925
4926 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
4927 return Err(MongrelError::InvalidArgument(format!(
4928 "AUTO_INCREMENT column {cid} must be Int64"
4929 )));
4930 };
4931 if data.len() < n {
4932 return Err(MongrelError::InvalidArgument(format!(
4933 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
4934 data.len()
4935 )));
4936 }
4937 if columnar::all_non_null(validity, n) {
4938 return Ok(());
4939 }
4940 if validity.iter().all(|b| *b == 0) {
4941 if let Some(start) = self.alloc_auto_inc_range(n)? {
4942 for (i, slot) in data.iter_mut().take(n).enumerate() {
4943 *slot = start.saturating_add(i as i64);
4944 }
4945 *validity = vec![0xFF; n.div_ceil(8)];
4946 }
4947 return Ok(());
4948 }
4949
4950 let new_validity = vec![0xFF; data.len().div_ceil(8)];
4951 for (i, slot) in data.iter_mut().enumerate().take(n) {
4952 if columnar::validity_bit(validity, i) {
4953 self.advance_auto_inc_past(*slot)?;
4954 } else {
4955 *slot = self.alloc_auto_inc_value()?;
4956 }
4957 }
4958 *validity = new_validity;
4959 Ok(())
4960 }
4961
4962 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
4976 self.ensure_writable()?;
4977 if self.auto_inc.is_none() {
4978 return Ok(None);
4979 }
4980 Ok(Some(self.alloc_auto_inc_value()?))
4981 }
4982
4983 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
4989 let payload = bincode::serialize(&rows)?;
4990 self.wal_append_data(Op::Put {
4991 table_id: self.table_id,
4992 rows: payload,
4993 })?;
4994 if self.is_shared() {
4995 self.pending_rows_auto_inc
4996 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
4997 self.pending_rows.extend(rows);
4998 } else {
4999 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
5000 }
5001 Ok(())
5002 }
5003
5004 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
5007 self.ensure_indexes_complete()
5008 }
5009
5010 pub(crate) fn prepare_durable_publish_controlled(
5011 &mut self,
5012 control: &crate::ExecutionControl,
5013 ) -> Result<()> {
5014 self.ensure_indexes_complete_controlled(control, || true)?;
5015 Ok(())
5016 }
5017
5018 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
5019 self.apply_put_rows_inner_prepared(rows, true);
5020 }
5021
5022 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
5023 if check_existing_pk {
5024 self.ensure_indexes_complete()?;
5025 }
5026 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
5027 Ok(())
5028 }
5029
5030 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
5034 if rows.len() == 1 {
5038 let row = rows.into_iter().next().expect("len checked");
5039 self.apply_put_row_single(row, check_existing_pk);
5040 return;
5041 }
5042 let pk_id = self.schema.primary_key().map(|c| c.id);
5059 let probe = match pk_id {
5060 Some(pid) => {
5061 check_existing_pk
5062 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
5063 }
5064 None => false,
5065 };
5066 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
5069 for r in rows {
5070 for &cid in r.columns.keys() {
5071 self.pending_put_cols.insert(cid);
5072 }
5073 let mut replaced_image: Option<Row> = None;
5074 match pk_id {
5075 Some(pid) if probe || maintain_pk_by_row => {
5076 if let Some(pk_val) = r.columns.get(&pid) {
5077 let key = self.index_lookup_key(pid, pk_val);
5078 if probe {
5079 if let Some(old) = self.recent_delete_preimages.remove(&key) {
5088 replaced_image = Some(old);
5089 if let Some(old_rid) = self.hot.get(&key) {
5090 if old_rid != r.row_id {
5091 self.tombstone_row(
5092 old_rid,
5093 r.committed_epoch,
5094 r.commit_ts,
5095 true,
5096 );
5097 }
5098 }
5099 } else if let Some(old_rid) = self.hot.get(&key) {
5100 if old_rid != r.row_id {
5101 replaced_image = self.get(old_rid, self.snapshot());
5102 self.tombstone_row(
5103 old_rid,
5104 r.committed_epoch,
5105 r.commit_ts,
5106 true,
5107 );
5108 }
5109 }
5110 }
5111 if maintain_pk_by_row {
5112 self.pk_by_row.insert(r.row_id, key);
5113 }
5114 }
5115 }
5116 Some(_) => {}
5117 None => {
5118 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5119 }
5120 }
5121 if let Some(old) = replaced_image {
5122 self.maintain_indexes_on_pk_replace(&old, &r);
5123 } else {
5124 self.index_row(&r);
5125 }
5126 self.reservoir.offer(r.row_id.0);
5127 self.memtable.upsert(r);
5128 self.live_count = self.live_count.saturating_add(1);
5131 }
5132 self.data_generation = self.data_generation.wrapping_add(1);
5133 }
5134
5135 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
5145 for &cid in row.columns.keys() {
5146 self.pending_put_cols.insert(cid);
5147 }
5148 let epoch = row.committed_epoch;
5149 let mut replaced_image: Option<Row> = None;
5150 if let Some(pk_col) = self.schema.primary_key() {
5151 let pk_id = pk_col.id;
5152 if let Some(pk_val) = row.columns.get(&pk_id) {
5153 let maintain_pk_by_row = self.pk_by_row_complete;
5157 if check_existing_pk || maintain_pk_by_row {
5158 let key = self.index_lookup_key(pk_id, pk_val);
5159 if check_existing_pk {
5160 if let Some(old) = self.recent_delete_preimages.remove(&key) {
5169 replaced_image = Some(old);
5170 if let Some(old_rid) = self.hot.get(&key) {
5171 if old_rid != row.row_id {
5172 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5173 }
5174 }
5175 } else if let Some(old_rid) = self.hot.get(&key) {
5176 if old_rid != row.row_id {
5177 replaced_image = self.get(old_rid, self.snapshot());
5181 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5182 }
5183 }
5184 }
5185 if maintain_pk_by_row {
5186 self.pk_by_row.insert(row.row_id, key);
5187 }
5188 }
5189 }
5190 } else {
5191 self.hot
5192 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
5193 }
5194 if let Some(old) = replaced_image {
5195 self.maintain_indexes_on_pk_replace(&old, &row);
5196 } else {
5197 self.index_row(&row);
5198 }
5199 self.reservoir.offer(row.row_id.0);
5200 self.memtable.upsert(row);
5201 self.live_count = self.live_count.saturating_add(1);
5202 self.data_generation = self.data_generation.wrapping_add(1);
5203 }
5204
5205 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
5208 let has_partial = self
5209 .schema
5210 .indexes
5211 .iter()
5212 .any(|idx| idx.predicate.is_some());
5213 if has_partial {
5214 self.unindex_bitmap_membership(old);
5218 self.index_row(new);
5219 return;
5220 }
5221
5222 crate::index::maintain_bitmap_secondary_on_replace(
5223 &self.schema,
5224 &mut self.bitmap,
5225 old,
5226 new,
5227 );
5228 self.index_row_non_bitmap(new);
5229 }
5230
5231 fn index_row_non_bitmap(&mut self, row: &Row) {
5234 if row.deleted {
5235 return;
5236 }
5237 let effective = if self.column_keys.is_empty() {
5238 None
5239 } else {
5240 Some(self.tokenized_for_indexes(row))
5241 };
5242 let source = effective.as_ref().unwrap_or(row);
5243 for idef in &self.schema.indexes {
5244 if idef.kind == crate::schema::IndexKind::Bitmap {
5245 continue;
5246 }
5247 index_into_single(
5248 idef,
5249 &self.schema,
5250 source,
5251 &mut self.hot,
5252 &mut self.bitmap,
5253 &mut self.ann,
5254 &mut self.fm,
5255 &mut self.sparse,
5256 &mut self.minhash,
5257 );
5258 }
5259 if let Some(pk_col) = self.schema.primary_key() {
5260 if let Some(pk_val) = source.columns.get(&pk_col.id) {
5261 let key = if self.column_keys.is_empty() {
5262 pk_val.encode_key()
5263 } else {
5264 self.index_lookup_key(pk_col.id, pk_val)
5265 };
5266 self.hot.insert(key, source.row_id);
5267 }
5268 }
5269 }
5270
5271 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
5274 self.allocator.alloc()
5275 }
5276
5277 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
5283 let pk = self.schema.primary_key().ok_or_else(|| {
5284 MongrelError::Schema("clustered table requires a single-column primary key".into())
5285 })?;
5286 let pk_val = columns
5287 .iter()
5288 .find(|(id, _)| *id == pk.id)
5289 .map(|(_, v)| v)
5290 .ok_or_else(|| {
5291 MongrelError::Schema(format!(
5292 "clustered table missing primary key column {} ({})",
5293 pk.id, pk.name
5294 ))
5295 })?;
5296 Ok(clustered_row_id(pk_val))
5297 }
5298
5299 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
5307 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
5308 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5309 }
5310 let n = rows.len();
5311 for r in rows {
5312 for &cid in r.columns.keys() {
5313 self.pending_put_cols.insert(cid);
5314 }
5315 }
5316 let (losers, winner_pks) = self.partition_pk_winners(rows);
5317 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
5318 let group_ts = rows.first().and_then(|r| r.commit_ts);
5320 for (key, &row_id) in &winner_pks {
5321 if let Some(old_rid) = self.hot.get(key) {
5322 if old_rid != row_id {
5323 self.tombstone_row(old_rid, epoch, group_ts, true);
5324 }
5325 }
5326 }
5327 for &loser_rid in &losers {
5330 self.tombstone_row(loser_rid, epoch, group_ts, false);
5331 }
5332 for (key, row_id) in winner_pks {
5334 self.insert_hot_pk(key, row_id);
5335 }
5336 if self.schema.primary_key().is_none() {
5337 for r in rows {
5338 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5339 }
5340 }
5341 for r in rows {
5342 self.allocator.advance_to(r.row_id)?;
5343 if !losers.contains(&r.row_id) {
5344 self.index_row(r);
5345 }
5346 }
5347 for r in rows {
5348 if !losers.contains(&r.row_id) {
5349 self.reservoir.offer(r.row_id.0);
5350 }
5351 }
5352 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
5353 self.data_generation = self.data_generation.wrapping_add(1);
5354 Ok(())
5355 }
5356
5357 pub(crate) fn recover_apply(
5362 &mut self,
5363 rows: Vec<Row>,
5364 deletes: Vec<(RowId, Epoch)>,
5365 ) -> Result<()> {
5366 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
5370 std::collections::BTreeMap::new();
5371 for row in rows {
5372 if row.row_id.0 >= u64::MAX - 1 {
5373 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5374 }
5375 self.allocator.advance_to(row.row_id)?;
5376 if let Some(ai) = self.auto_inc.as_mut() {
5381 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
5382 let next = n.checked_add(1).ok_or_else(|| {
5383 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
5384 })?;
5385 if next > ai.next {
5386 ai.next = next;
5387 }
5388 }
5389 }
5390 by_epoch.entry(row.committed_epoch).or_default().push(row);
5391 }
5392 for (epoch, group) in by_epoch {
5393 let (losers, winner_pks) = self.partition_pk_winners(&group);
5394 let group_ts = group.first().and_then(|r| r.commit_ts);
5396 for (key, &row_id) in &winner_pks {
5397 if let Some(old_rid) = self.hot.get(key) {
5398 if old_rid != row_id {
5399 self.tombstone_row(old_rid, epoch, group_ts, false);
5400 }
5401 }
5402 }
5403 for (key, row_id) in winner_pks {
5404 self.insert_hot_pk(key, row_id);
5405 }
5406 if self.schema.primary_key().is_none() {
5407 for r in &group {
5408 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5409 }
5410 }
5411 for r in &group {
5412 if !losers.contains(&r.row_id) {
5413 self.memtable.upsert(r.clone());
5414 self.index_row(r);
5415 }
5416 }
5417 }
5418 for (rid, epoch) in deletes {
5419 self.memtable.tombstone(rid, epoch);
5420 self.remove_hot_for_row(rid, epoch);
5421 }
5422 self.reservoir_complete = false;
5425 Ok(())
5426 }
5427
5428 pub(crate) fn flushed_epoch(&self) -> u64 {
5430 self.flushed_epoch
5431 }
5432
5433 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
5434 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
5435 }
5436
5437 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
5439 self.schema.validate_values(cells)
5440 }
5441
5442 fn validate_columns_not_null(
5446 &self,
5447 columns: &[(u16, columnar::NativeColumn)],
5448 n: usize,
5449 ) -> Result<()> {
5450 let by_id: HashMap<u16, &columnar::NativeColumn> =
5451 columns.iter().map(|(id, c)| (*id, c)).collect();
5452 for col in &self.schema.columns {
5453 if !col.flags.contains(ColumnFlags::NULLABLE) {
5454 match by_id.get(&col.id) {
5455 None => {
5456 return Err(MongrelError::InvalidArgument(format!(
5457 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
5458 col.name, col.id
5459 )));
5460 }
5461 Some(c) => {
5462 if c.null_count(n) != 0 {
5463 return Err(MongrelError::InvalidArgument(format!(
5464 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
5465 col.name, col.id
5466 )));
5467 }
5468 }
5469 }
5470 }
5471 if let TypeId::Enum { variants } = &col.ty {
5472 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
5473 if by_id.contains_key(&col.id) {
5474 return Err(MongrelError::InvalidArgument(format!(
5475 "column '{}' ({}) enum requires a bytes column",
5476 col.name, col.id
5477 )));
5478 }
5479 continue;
5480 };
5481 for index in 0..n {
5482 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
5483 continue;
5484 };
5485 if !variants.iter().any(|variant| variant.as_bytes() == value) {
5486 return Err(MongrelError::InvalidArgument(format!(
5487 "column '{}' ({}) enum value {:?} is not one of {:?}",
5488 col.name,
5489 col.id,
5490 String::from_utf8_lossy(value),
5491 variants
5492 )));
5493 }
5494 }
5495 }
5496 }
5497 Ok(())
5498 }
5499
5500 fn bulk_pk_winner_indices(
5505 &self,
5506 columns: &[(u16, columnar::NativeColumn)],
5507 n: usize,
5508 ) -> Option<Vec<usize>> {
5509 let pk_col = self.schema.primary_key()?;
5510 let pk_id = pk_col.id;
5511 let pk_ty = pk_col.ty.clone();
5512 let by_id: HashMap<u16, &columnar::NativeColumn> =
5513 columns.iter().map(|(id, c)| (*id, c)).collect();
5514 let pk_native = by_id.get(&pk_id)?;
5515 if native_int64_strictly_increasing(pk_native, n) {
5516 return None;
5517 }
5518 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
5520 let mut null_pk_rows: Vec<usize> = Vec::new();
5521 for i in 0..n {
5522 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
5523 Some(key) => {
5524 last.insert(key, i);
5525 }
5526 None => null_pk_rows.push(i),
5527 }
5528 }
5529 let mut winners: HashSet<usize> = last.values().copied().collect();
5530 for i in null_pk_rows {
5531 winners.insert(i);
5532 }
5533 Some((0..n).filter(|i| winners.contains(i)).collect())
5534 }
5535
5536 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
5538 self.require_delete()?;
5539 let epoch = self.pending_epoch();
5540 self.wal_append_data(Op::Delete {
5541 table_id: self.table_id,
5542 row_ids: vec![row_id],
5543 })?;
5544 if self.is_shared() {
5545 self.pending_dels.push(row_id);
5546 } else {
5547 self.apply_delete(row_id, epoch);
5548 }
5549 Ok(())
5550 }
5551
5552 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
5553 let pre = self.get(row_id, self.snapshot());
5554 self.delete(row_id)?;
5555 Ok(pre.map(|row| {
5556 let mut columns: Vec<_> = row.columns.into_iter().collect();
5557 columns.sort_by_key(|(id, _)| *id);
5558 OwnedRow { columns }
5559 }))
5560 }
5561
5562 pub fn truncate(&mut self) -> Result<()> {
5564 self.require_delete()?;
5565 let epoch = self.pending_epoch();
5566 self.wal_append_data(Op::TruncateTable {
5567 table_id: self.table_id,
5568 })?;
5569 self.pending_rows.clear();
5570 self.pending_rows_auto_inc.clear();
5571 self.pending_dels.clear();
5572 self.pending_truncate = Some(epoch);
5573 Ok(())
5574 }
5575
5576 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
5578 self.run_refs.clear();
5583 self.retiring.clear();
5584 self.memtable = Memtable::new();
5585 self.mutable_run = MutableRun::new();
5586 self.hot = HotIndex::new();
5587 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
5588 self.bitmap = bitmap;
5589 self.ann = ann;
5590 self.fm = fm;
5591 self.sparse = sparse;
5592 self.minhash = minhash;
5593 self.learned_range = Arc::new(HashMap::new());
5594 self.pk_by_row.clear();
5595 self.pk_by_row_complete = false;
5596 self.live_count = 0;
5597 self.reservoir = crate::reservoir::Reservoir::default();
5598 self.reservoir_complete = true;
5599 self.had_deletes = true;
5600 self.agg_cache = Arc::new(HashMap::new());
5601 self.global_idx_epoch = 0;
5602 self.indexes_complete = true;
5603 self.pending_delete_rids.clear();
5604 self.pending_put_cols.clear();
5605 self.pending_rows.clear();
5606 self.pending_rows_auto_inc.clear();
5607 self.pending_dels.clear();
5608 self.clear_result_cache();
5609 self.invalidate_index_checkpoint();
5610 self.data_generation = self.data_generation.wrapping_add(1);
5611 }
5612
5613 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
5616 self.apply_delete_at(row_id, epoch, None);
5617 }
5618
5619 pub(crate) fn apply_delete_at(
5621 &mut self,
5622 row_id: RowId,
5623 epoch: Epoch,
5624 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5625 ) {
5626 let preimage = self.get(row_id, self.snapshot());
5636 if let Some(row) = preimage {
5637 if let Some(pk_col) = self.schema.primary_key() {
5638 if let Some(pk_val) = row.columns.get(&pk_col.id) {
5639 let key = self.index_lookup_key(pk_col.id, pk_val);
5640 self.recent_delete_preimages.insert(key, row);
5641 }
5642 }
5643 }
5644 self.tombstone_row(row_id, epoch, commit_ts, true);
5645 self.data_generation = self.data_generation.wrapping_add(1);
5646 }
5647
5648 fn unindex_bitmap_membership(&mut self, row: &Row) {
5653 if row.deleted {
5654 return;
5655 }
5656 for idef in &self.schema.indexes {
5657 if idef.kind != crate::schema::IndexKind::Bitmap {
5658 continue;
5659 }
5660 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
5661 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5662 b.remove(&key, row.row_id);
5663 }
5664 }
5665 }
5666 }
5667
5668 fn union_bitmap_point_i64(
5672 &self,
5673 set: &mut RowIdSet,
5674 column_id: u16,
5675 value: i64,
5676 snapshot: Snapshot,
5677 ) {
5678 let Some(b) = self.bitmap.get(&column_id) else {
5679 return;
5680 };
5681 let encoded = Value::Int64(value).encode_key();
5682 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
5683 for rid in b.get(&lookup).iter() {
5684 set.insert(u64::from(rid));
5685 }
5686 set.remove_many(self.overlay_tombstoned_rids(snapshot));
5690 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
5691 }
5692
5693 fn tombstone_row(
5707 &mut self,
5708 row_id: RowId,
5709 epoch: Epoch,
5710 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5711 adjust_live_count: bool,
5712 ) {
5713 let prev_was_live = if adjust_live_count {
5714 match self.memtable.get_version(row_id, epoch) {
5715 Some((_, prev)) => !prev.deleted,
5716 None => true,
5717 }
5718 } else {
5719 false
5720 };
5721 let tombstone = Row {
5722 row_id,
5723 committed_epoch: epoch,
5724 columns: std::collections::HashMap::new(),
5725 deleted: true,
5726 commit_ts,
5727 };
5728 self.memtable.upsert(tombstone);
5729 self.pk_by_row.remove(&row_id);
5730 if prev_was_live {
5731 self.live_count = self.live_count.saturating_sub(1);
5732 }
5733 self.pending_delete_rids.insert(row_id.0 as u32);
5735 self.had_deletes = true;
5738 self.agg_cache = Arc::new(HashMap::new());
5739 }
5740
5741 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
5745 let Some(pk_col) = self.schema.primary_key() else {
5746 return;
5747 };
5748 if self.pk_by_row_complete {
5751 if let Some(key) = self.pk_by_row.remove(&row_id) {
5752 if self.hot.get(&key) == Some(row_id) {
5753 self.hot.remove(&key);
5754 }
5755 }
5756 return;
5757 }
5758 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
5777 if self.indexes_complete {
5778 let pk_val = self
5779 .memtable
5780 .get_version(row_id, lookup_epoch)
5781 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5782 .or_else(|| {
5783 self.mutable_run
5784 .get_version(row_id, lookup_epoch)
5785 .filter(|(_, r)| !r.deleted)
5786 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5787 })
5788 .or_else(|| {
5789 self.run_refs.iter().find_map(|rr| {
5790 let mut reader = self.open_reader(rr.run_id).ok()?;
5791 let (_, deleted, val) = reader
5792 .get_version_column(row_id, lookup_epoch, pk_col.id)
5793 .ok()??;
5794 if deleted {
5795 return None;
5796 }
5797 val
5798 })
5799 });
5800 if let Some(pk_val) = pk_val {
5801 let key = self.index_lookup_key(pk_col.id, &pk_val);
5802 if self.hot.get(&key) == Some(row_id) {
5803 self.hot.remove(&key);
5804 }
5805 return;
5806 }
5807 }
5808 self.refresh_pk_by_row_from_hot();
5813 if let Some(key) = self.pk_by_row.remove(&row_id) {
5814 if self.hot.get(&key) == Some(row_id) {
5815 self.hot.remove(&key);
5816 }
5817 }
5818 }
5819
5820 fn partition_pk_winners(
5825 &self,
5826 rows: &[Row],
5827 ) -> (
5828 std::collections::HashSet<RowId>,
5829 std::collections::HashMap<Vec<u8>, RowId>,
5830 ) {
5831 let mut losers = std::collections::HashSet::new();
5832 let Some(pk_col) = self.schema.primary_key() else {
5833 return (losers, std::collections::HashMap::new());
5834 };
5835 let pk_id = pk_col.id;
5836 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
5837 std::collections::HashMap::new();
5838 for r in rows {
5839 let Some(pk_val) = r.columns.get(&pk_id) else {
5840 continue;
5841 };
5842 let key = self.index_lookup_key(pk_id, pk_val);
5843 if let Some(&old_rid) = winners.get(&key) {
5844 losers.insert(old_rid);
5845 }
5846 winners.insert(key, r.row_id);
5847 }
5848 (losers, winners)
5849 }
5850
5851 fn index_row(&mut self, row: &Row) {
5852 if row.deleted {
5853 return;
5854 }
5855 let any_predicate = self
5863 .schema
5864 .indexes
5865 .iter()
5866 .any(|idx| idx.predicate.is_some());
5867 if any_predicate {
5868 let columns_map: HashMap<u16, &Value> =
5869 row.columns.iter().map(|(k, v)| (*k, v)).collect();
5870 let name_to_id: HashMap<&str, u16> = self
5871 .schema
5872 .columns
5873 .iter()
5874 .map(|c| (c.name.as_str(), c.id))
5875 .collect();
5876 for idx in &self.schema.indexes {
5877 if let Some(pred) = &idx.predicate {
5878 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
5879 continue; }
5881 }
5882 index_into_single(
5884 idx,
5885 &self.schema,
5886 row,
5887 &mut self.hot,
5888 &mut self.bitmap,
5889 &mut self.ann,
5890 &mut self.fm,
5891 &mut self.sparse,
5892 &mut self.minhash,
5893 );
5894 }
5895 return;
5896 }
5897 if self.column_keys.is_empty() {
5901 index_into(
5902 &self.schema,
5903 row,
5904 &mut self.hot,
5905 &mut self.bitmap,
5906 &mut self.ann,
5907 &mut self.fm,
5908 &mut self.sparse,
5909 &mut self.minhash,
5910 );
5911 return;
5912 }
5913 let effective_row = self.tokenized_for_indexes(row);
5914 index_into(
5915 &self.schema,
5916 &effective_row,
5917 &mut self.hot,
5918 &mut self.bitmap,
5919 &mut self.ann,
5920 &mut self.fm,
5921 &mut self.sparse,
5922 &mut self.minhash,
5923 );
5924 }
5925
5926 fn tokenized_for_indexes(&self, row: &Row) -> Row {
5932 if self.column_keys.is_empty() {
5933 return row.clone();
5934 }
5935 {
5936 use crate::encryption::SCHEME_HMAC_EQ;
5937 let mut tok = row.clone();
5938 for (&cid, &(_, scheme)) in &self.column_keys {
5939 if scheme != SCHEME_HMAC_EQ {
5940 continue;
5941 }
5942 if let Some(v) = tok.columns.get(&cid).cloned() {
5943 if let Some(t) = self.tokenize_value(cid, &v) {
5944 tok.columns.insert(cid, t);
5945 }
5946 }
5947 }
5948 tok
5949 }
5950 }
5951
5952 pub fn commit(&mut self) -> Result<Epoch> {
5957 self.commit_inner(None)
5958 }
5959
5960 #[doc(hidden)]
5963 pub fn commit_controlled<F>(
5964 &mut self,
5965 control: &crate::ExecutionControl,
5966 mut before_commit: F,
5967 ) -> Result<Epoch>
5968 where
5969 F: FnMut() -> Result<()>,
5970 {
5971 self.commit_inner(Some((control, &mut before_commit)))
5972 }
5973
5974 fn commit_inner(
5975 &mut self,
5976 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5977 ) -> Result<Epoch> {
5978 self.ensure_writable()?;
5979 if !self.has_pending_mutations() {
5980 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
5981 return Err(MongrelError::Full(
5982 "standalone transaction id namespace exhausted".into(),
5983 ));
5984 }
5985 return Ok(self.epoch.visible());
5986 }
5987 self.commit_new_epoch_inner(controlled)
5988 }
5989
5990 fn commit_new_epoch(&mut self) -> Result<Epoch> {
5994 self.commit_new_epoch_inner(None)
5995 }
5996
5997 fn commit_new_epoch_inner(
5998 &mut self,
5999 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6000 ) -> Result<Epoch> {
6001 self.ensure_writable()?;
6002 if self.is_shared() {
6003 self.commit_shared(controlled)
6004 } else {
6005 self.commit_private(controlled)
6006 }
6007 }
6008
6009 fn commit_private(
6011 &mut self,
6012 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6013 ) -> Result<Epoch> {
6014 let commit_lock = Arc::clone(&self.commit_lock);
6018 let _g = commit_lock.lock();
6019 let txn_id = self.ensure_txn_id()?;
6022 if let Some((control, before_commit)) = controlled {
6023 control.checkpoint()?;
6024 before_commit()?;
6025 }
6026 let new_epoch = self.epoch.bump_assigned();
6027 let epoch_authority = Arc::clone(&self.epoch);
6028 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
6029 let wal_result = match &mut self.wal {
6033 WalSink::Private(w) => w
6034 .append_txn(
6035 txn_id,
6036 Op::TxnCommit {
6037 epoch: new_epoch.0,
6038 added_runs: Vec::new(),
6039 },
6040 )
6041 .and_then(|_| w.sync()),
6042 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
6043 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
6044 };
6045 if let Err(error) = wal_result {
6046 self.durable_commit_failed = true;
6047 return Err(MongrelError::CommitOutcomeUnknown {
6048 epoch: new_epoch.0,
6049 message: error.to_string(),
6050 });
6051 }
6052 if let Some(epoch) = self.pending_truncate.take() {
6055 self.apply_truncate(epoch);
6056 }
6057 self.invalidate_pending_cache();
6058 let publish_result = self.persist_manifest(new_epoch);
6059 self.epoch.publish_in_order(new_epoch);
6063 epoch_guard.disarm();
6064 if let Err(error) = publish_result {
6065 self.durable_commit_failed = true;
6066 return Err(MongrelError::DurableCommit {
6067 epoch: new_epoch.0,
6068 message: error.to_string(),
6069 });
6070 }
6071 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
6072 self.pending_private_mutations = false;
6073 self.data_generation = self.data_generation.wrapping_add(1);
6074 Ok(new_epoch)
6075 }
6076
6077 fn commit_shared(
6083 &mut self,
6084 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6085 ) -> Result<Epoch> {
6086 use std::sync::atomic::Ordering;
6087 let s = match &self.wal {
6088 WalSink::Shared(s) => s.clone(),
6089 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
6090 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
6091 };
6092 if s.poisoned.load(Ordering::Relaxed) {
6093 return Err(MongrelError::Other(
6094 "database poisoned by fsync error".into(),
6095 ));
6096 }
6097 let commit_lock = Arc::clone(&self.commit_lock);
6104 let _g = commit_lock.lock();
6105 if !self.pending_rows.is_empty() {
6106 match controlled.as_ref() {
6107 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
6108 None => self.prepare_durable_publish()?,
6109 }
6110 }
6111 let txn_id = self.ensure_txn_id()?;
6114 let mut wal = s.wal.lock();
6115 if let Some((control, before_commit)) = controlled {
6116 control.checkpoint()?;
6117 before_commit()?;
6118 }
6119 let new_epoch = self.epoch.bump_assigned();
6120 let epoch_authority = Arc::clone(&self.epoch);
6121 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
6122 let commit_ts = s.hlc.now().map_err(|skew| {
6125 MongrelError::Other(format!(
6126 "clock skew rejected commit timestamp allocation: {skew}"
6127 ))
6128 })?;
6129 let commit_seq = match wal.append_commit_at(
6130 txn_id,
6131 new_epoch,
6132 &[],
6133 commit_ts.physical_micros.saturating_mul(1_000),
6134 ) {
6135 Ok(commit_seq) => commit_seq,
6136 Err(error) => {
6137 s.poisoned.store(true, Ordering::Relaxed);
6138 s.lifecycle.poison();
6139 return Err(MongrelError::CommitOutcomeUnknown {
6140 epoch: new_epoch.0,
6141 message: error.to_string(),
6142 });
6143 }
6144 };
6145 drop(wal);
6146 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
6147 s.poisoned.store(true, Ordering::Relaxed);
6148 s.lifecycle.poison();
6149 return Err(MongrelError::CommitOutcomeUnknown {
6150 epoch: new_epoch.0,
6151 message: error.to_string(),
6152 });
6153 }
6154
6155 if self.pending_truncate.take().is_some() {
6158 self.apply_truncate(new_epoch);
6159 }
6160 let mut rows = std::mem::take(&mut self.pending_rows);
6161 if !rows.is_empty() {
6162 for r in &mut rows {
6163 r.committed_epoch = new_epoch;
6164 r.commit_ts = Some(commit_ts);
6165 }
6166 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
6167 let all_auto_generated =
6168 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
6169 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
6170 } else {
6171 self.pending_rows_auto_inc.clear();
6172 }
6173 let dels = std::mem::take(&mut self.pending_dels);
6174 for rid in dels {
6175 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
6176 }
6177
6178 self.invalidate_pending_cache();
6179 let publish_result = self.persist_manifest(new_epoch);
6180 self.epoch.publish_in_order(new_epoch);
6181 epoch_guard.disarm();
6182 let _ = s.change_wake.send(());
6183 if let Err(error) = publish_result {
6184 self.durable_commit_failed = true;
6185 s.poisoned.store(true, Ordering::Relaxed);
6186 s.lifecycle.poison();
6187 return Err(MongrelError::DurableCommit {
6188 epoch: new_epoch.0,
6189 message: error.to_string(),
6190 });
6191 }
6192 self.current_txn_id = 0;
6194 self.data_generation = self.data_generation.wrapping_add(1);
6195 Ok(new_epoch)
6196 }
6197
6198 pub fn flush(&mut self) -> Result<Epoch> {
6206 self.flush_with_outcome().map(|(epoch, _)| epoch)
6207 }
6208
6209 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
6211 self.flush_with_outcome_inner(None)
6212 }
6213
6214 #[doc(hidden)]
6217 pub fn flush_with_outcome_controlled<F>(
6218 &mut self,
6219 control: &crate::ExecutionControl,
6220 mut before_commit: F,
6221 ) -> Result<(Epoch, bool)>
6222 where
6223 F: FnMut() -> Result<()>,
6224 {
6225 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
6226 }
6227
6228 fn flush_with_outcome_inner(
6229 &mut self,
6230 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6231 ) -> Result<(Epoch, bool)> {
6232 match controlled.as_ref() {
6233 Some((control, _)) => {
6234 self.ensure_indexes_complete_controlled(control, || true)?;
6235 }
6236 None => self.ensure_indexes_complete()?,
6237 }
6238 let committed = self.has_pending_mutations();
6239 let epoch = self.commit_inner(controlled)?;
6240 let finish: Result<(Epoch, bool)> = (|| {
6241 let rows = self.memtable.drain_sorted();
6242 if !rows.is_empty() {
6243 self.mutable_run.insert_many(rows);
6244 }
6245 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
6246 self.spill_mutable_run(epoch)?;
6247 self.mark_flushed(epoch)?;
6251 self.persist_manifest(epoch)?;
6252 self.build_learned_ranges()?;
6253 self.checkpoint_indexes(epoch);
6256 let _ = self.publish_run_lookup_directory();
6261 }
6262 Ok((epoch, committed))
6265 })();
6266 let outcome = match finish {
6267 Err(error) if committed => Err(MongrelError::DurableCommit {
6268 epoch: epoch.0,
6269 message: error.to_string(),
6270 }),
6271 result => result,
6272 };
6273 if outcome.is_ok() {
6274 let _ = self.publish_read_generation();
6280 }
6281 outcome
6282 }
6283
6284 fn has_pending_mutations(&self) -> bool {
6285 self.pending_private_mutations
6286 || !self.pending_rows.is_empty()
6287 || !self.pending_dels.is_empty()
6288 || self.pending_truncate.is_some()
6289 }
6290
6291 pub fn has_pending_writes(&self) -> bool {
6292 self.has_pending_mutations()
6293 }
6294
6295 pub fn force_flush(&mut self) -> Result<Epoch> {
6304 let saved = self.mutable_run_spill_bytes;
6305 self.mutable_run_spill_bytes = 1;
6306 let result = self.flush();
6307 self.mutable_run_spill_bytes = saved;
6308 result
6309 }
6310
6311 pub fn close(&mut self) -> Result<()> {
6318 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
6319 self.force_flush()?;
6320 }
6321 self.result_cache
6327 .lock()
6328 .shutdown_persistent_cache(std::time::Duration::from_secs(5));
6329 Ok(())
6330 }
6331
6332 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
6339 let op = Op::Flush {
6340 table_id: self.table_id,
6341 flushed_epoch: epoch.0,
6342 };
6343 match &mut self.wal {
6344 WalSink::Private(w) => {
6345 w.append_system(op)?;
6346 w.sync()?;
6347 }
6348 WalSink::Shared(s) => {
6349 s.wal.lock().append_system(op)?;
6354 }
6355 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
6356 }
6357 self.flushed_epoch = epoch.0;
6358 if matches!(self.wal, WalSink::Private(_)) {
6359 self.rotate_wal(epoch)?;
6360 }
6361 Ok(())
6362 }
6363
6364 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
6368 if self.mutable_run.is_empty() {
6369 return Ok(());
6370 }
6371 let run_id = self.alloc_run_id()?;
6372 let rows = self.mutable_run.drain_sorted();
6373 if rows.is_empty() {
6374 return Ok(());
6375 }
6376 let path = self.run_path(run_id);
6377 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
6378 if let Some(kek) = &self.kek {
6379 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6380 }
6381 let header = match self.create_run_file(run_id)? {
6382 Some(file) => writer.write_file(file, &rows)?,
6383 None => writer.write(&path, &rows)?,
6384 };
6385 self.run_refs.push(RunRef {
6386 run_id: run_id as u128,
6387 level: 0,
6388 epoch_created: epoch.0,
6389 row_count: header.row_count,
6390 });
6391 self.run_row_id_ranges
6392 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6393 Ok(())
6394 }
6395
6396 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
6400 self.mutable_run_spill_bytes = bytes.max(1);
6401 }
6402
6403 pub fn set_compaction_zstd_level(&mut self, level: i32) {
6407 self.compaction_zstd_level = level;
6408 }
6409
6410 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
6414 self.result_cache.lock().set_max_bytes(max_bytes);
6415 }
6416
6417 pub fn flush_persistent_cache(&self, deadline_ms: u64) -> u64 {
6422 self.result_cache
6423 .lock()
6424 .flush_persistent_cache(std::time::Duration::from_millis(deadline_ms))
6425 }
6426
6427 pub fn shutdown_persistent_cache(&mut self, deadline_ms: u64) {
6432 self.result_cache
6433 .lock()
6434 .shutdown_persistent_cache(std::time::Duration::from_millis(deadline_ms));
6435 }
6436
6437 #[doc(hidden)]
6442 pub fn _set_persist_min_bytes_for_test(&mut self, min: u64) {
6443 self.result_cache.lock().set_persist_min_bytes(min);
6444 }
6445
6446 pub(crate) fn clear_result_cache(&mut self) {
6450 self.result_cache.lock().clear();
6451 }
6452
6453 pub fn mutable_run_len(&self) -> usize {
6455 self.mutable_run.len()
6456 }
6457
6458 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
6461 self.mutable_run.drain_sorted()
6462 }
6463
6464 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
6466 let mut snapshot = self.mutable_run.clone();
6467 snapshot.drain_sorted()
6468 }
6469
6470 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
6475 self.ensure_writable()?;
6476 let n = batch.len();
6477 if n == 0 {
6478 return Ok(self.current_epoch());
6479 }
6480 for row in &batch {
6481 self.schema.validate_values(row)?;
6482 }
6483 let epoch = self.commit_new_epoch()?;
6484 let live_before = self.live_count;
6485 self.spill_mutable_run(epoch)?;
6489 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6490 && self.indexes_complete
6491 && self.run_refs.is_empty()
6492 && self.memtable.is_empty()
6493 && self.mutable_run.is_empty();
6494 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
6500 use rayon::prelude::*;
6501 self.schema
6502 .columns
6503 .par_iter()
6504 .map(|cdef| {
6505 (
6506 cdef.id,
6507 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
6508 )
6509 })
6510 .collect::<Vec<_>>()
6511 };
6512 drop(batch);
6513 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6518 self.validate_columns_not_null(&user_columns, n)?;
6519 let winner_idx = self
6520 .bulk_pk_winner_indices(&user_columns, n)
6521 .filter(|idx| idx.len() != n);
6522 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6523 match winner_idx.as_deref() {
6524 Some(idx) => {
6525 let compacted = user_columns
6526 .iter()
6527 .map(|(id, c)| (*id, c.gather(idx)))
6528 .collect();
6529 (compacted, idx.len())
6530 }
6531 None => (std::mem::take(&mut user_columns), n),
6532 };
6533 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6534 let first = self.allocator.alloc_range(write_n as u64)?.0;
6535 for rid in first..first + write_n as u64 {
6536 self.reservoir.offer(rid);
6537 }
6538 let run_id = self.alloc_run_id()?;
6539 let path = self.run_path(run_id);
6540 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
6541 .clean(true)
6542 .with_lz4()
6543 .with_native_endian();
6544 if let Some(kek) = &self.kek {
6545 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6546 }
6547 let header = match self.create_run_file(run_id)? {
6548 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
6549 None => writer.write_native(&path, &write_columns, write_n, first)?,
6550 };
6551 self.run_refs.push(RunRef {
6552 run_id: run_id as u128,
6553 level: 0,
6554 epoch_created: epoch.0,
6555 row_count: header.row_count,
6556 });
6557 self.run_row_id_ranges
6558 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6559 self.live_count = self.live_count.saturating_add(write_n as u64);
6560 if eager_index_build {
6561 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6562 self.index_columns_bulk(&write_columns, &row_ids);
6563 self.indexes_complete = true;
6564 self.build_learned_ranges()?;
6565 } else {
6566 self.indexes_complete = false;
6567 }
6568 self.mark_flushed(epoch)?;
6569 self.persist_manifest(epoch)?;
6570 if eager_index_build {
6571 self.checkpoint_indexes(epoch);
6572 }
6573 self.clear_result_cache();
6574 Ok(epoch)
6575 }
6576
6577 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
6580 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
6581 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
6582 let segment_no = segment
6585 .file_stem()
6586 .and_then(|s| s.to_str())
6587 .and_then(|s| s.strip_prefix("seg-"))
6588 .and_then(|s| s.parse::<u64>().ok())
6589 .unwrap_or(0);
6590 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
6591 wal.set_sync_byte_threshold(self.sync_byte_threshold);
6592 wal.sync()?;
6593 self.wal = WalSink::Private(wal);
6594 Ok(())
6595 }
6596
6597 pub(crate) fn invalidate_pending_cache(&mut self) {
6602 self.result_cache
6603 .lock()
6604 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
6605 self.pending_delete_rids.clear();
6606 self.pending_put_cols.clear();
6607 }
6608
6609 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
6610 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
6611 m.current_epoch = epoch.0;
6612 m.next_row_id = self.allocator.current().0;
6613 m.runs = self.run_refs.clone();
6614 m.live_count = self.live_count;
6615 m.global_idx_epoch = self.global_idx_epoch;
6616 m.flushed_epoch = self.flushed_epoch;
6617 m.retiring = self.retiring.clone();
6618 m.auto_inc_next = match self.auto_inc {
6622 Some(ai) if ai.seeded => ai.next,
6623 _ => 0,
6624 };
6625 m.ttl = self.ttl;
6626 let meta_dek = self.manifest_meta_dek();
6627 match self._root_guard.as_deref() {
6628 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
6629 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
6630 }
6631 Ok(())
6632 }
6633
6634 pub(crate) fn publish_run_lookup_directory(&mut self) -> Result<()> {
6641 mongreldb_fault::inject("manifest.publication")
6645 .map_err(|e| MongrelError::Other(format!("manifest.publication fault: {e}")))?;
6646 let active_runs = self.run_refs.clone();
6647 let schema_id = self.schema.schema_id;
6648 let index_generation = self.global_idx_epoch;
6649 let run_generation = active_runs.len() as u64;
6650 let fingerprint = crate::run_lookup::compute_fingerprint(
6651 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
6652 schema_id,
6653 index_generation,
6654 run_generation,
6655 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
6656 );
6657 let dir = match crate::run_lookup::RunLookupDirectory::rebuild_from_runs(&active_runs, self)
6658 {
6659 Ok(d) => d,
6660 Err(error) => {
6661 self.run_lookup.complete = false;
6664 self.lookup_metrics
6665 .directory_incomplete
6666 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6667 return Err(error);
6668 }
6669 };
6670 let mut dir = dir;
6674 dir.set_fingerprint(fingerprint);
6675 let runs_dir = match self.runs_root.as_deref() {
6676 Some(root) => match root.io_path() {
6677 Ok(p) => p,
6678 Err(_) => self.dir.join(RUNS_DIR),
6679 },
6680 None => self.dir.join(RUNS_DIR),
6681 };
6682 mongreldb_fault::inject("directory.temp_write")
6685 .map_err(|e| MongrelError::Other(format!("directory.temp_write fault: {e}")))?;
6686 let publish =
6687 dir.write_checkpoint_sharded(&runs_dir, crate::run_lookup::DEFAULT_SHARD_BYTES);
6688 mongreldb_fault::inject("directory.sync")
6689 .map_err(|e| MongrelError::Other(format!("directory.sync fault: {e}")))?;
6690 mongreldb_fault::inject("directory.rename")
6691 .map_err(|e| MongrelError::Other(format!("directory.rename fault: {e}")))?;
6692 if let Err(error) = publish {
6693 self.run_lookup = RunLookupState {
6696 directory: None,
6697 fingerprint,
6698 complete: false,
6699 };
6700 self.lookup_metrics
6701 .directory_incomplete
6702 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6703 self.lookup_metrics
6704 .directory_lookup_fallback
6705 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6706 return Err(error.into());
6707 }
6708 self.run_lookup = RunLookupState {
6709 directory: Some(std::sync::Arc::new(dir)),
6710 fingerprint,
6711 complete: true,
6712 };
6713 Ok(())
6714 }
6715
6716 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
6717 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
6721 let live_count = u64::try_from(rows.len())
6722 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
6723 let auto_inc = match self.auto_inc {
6724 Some(mut state) => {
6725 let maximum = self.scan_max_int64(state.column_id)?;
6726 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
6727 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
6728 })?;
6729 state.next = state.next.max(after_maximum).max(1);
6730 state.seeded = true;
6731 Some(state)
6732 }
6733 None => None,
6734 };
6735 Ok(RecoveryMetadataPlan {
6736 live_count,
6737 auto_inc,
6738 changed: live_count != self.live_count
6739 || auto_inc.is_some_and(|planned| {
6740 self.auto_inc.is_none_or(|current| {
6741 current.next != planned.next || current.seeded != planned.seeded
6742 })
6743 }),
6744 })
6745 }
6746
6747 pub(crate) fn apply_recovered_metadata(
6748 &mut self,
6749 plan: RecoveryMetadataPlan,
6750 epoch: Epoch,
6751 ) -> Result<()> {
6752 if !plan.changed {
6753 return Ok(());
6754 }
6755 self.live_count = plan.live_count;
6756 self.auto_inc = plan.auto_inc;
6757 self.persist_manifest(epoch)
6758 }
6759
6760 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
6766 if !self.indexes_complete {
6769 return;
6770 }
6771 if crate::catalog::inject_hook("index.publish.before").is_err() {
6774 return;
6775 }
6776 if self.idx_root.is_none() {
6777 if let Some(root) = self._root_guard.as_ref() {
6778 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
6779 return;
6780 };
6781 self.idx_root = Some(Arc::new(idx_root));
6782 }
6783 }
6784 let snap = global_idx::IndexSnapshot {
6785 hot: &self.hot,
6786 bitmap: &self.bitmap,
6787 ann: &self.ann,
6788 fm: &self.fm,
6789 sparse: &self.sparse,
6790 minhash: &self.minhash,
6791 learned_range: &self.learned_range,
6792 };
6793 let idx_dek = self.idx_dek();
6795 let written = match self.idx_root.as_deref() {
6796 Some(root) => global_idx::write_atomic_root(
6797 root,
6798 self.table_id,
6799 epoch.0,
6800 snap,
6801 idx_dek.as_deref(),
6802 ),
6803 None => global_idx::write_atomic(
6804 &self.dir,
6805 self.table_id,
6806 epoch.0,
6807 snap,
6808 idx_dek.as_deref(),
6809 ),
6810 };
6811 if written.is_ok() {
6812 self.global_idx_epoch = epoch.0;
6813 let _ = self.persist_manifest(epoch);
6814 let _ = crate::catalog::inject_hook("index.publish.after");
6816 }
6817 }
6818
6819 pub(crate) fn invalidate_index_checkpoint(&mut self) {
6822 self.global_idx_epoch = 0;
6823 if let Some(root) = self.idx_root.as_deref() {
6824 let _ = root.remove_file(global_idx::IDX_FILENAME);
6825 } else {
6826 global_idx::remove(&self.dir);
6827 }
6828 let _ = self.persist_manifest(self.epoch.visible());
6829 }
6830
6831 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
6836 self.indexes_complete = false;
6837 self.global_idx_epoch = 0;
6838 if let Some(root) = self.idx_root.as_deref() {
6839 let _ = root.remove_file(global_idx::IDX_FILENAME);
6840 } else {
6841 global_idx::remove(&self.dir);
6842 }
6843 }
6844
6845 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
6846 self.indexes_complete = true;
6847 }
6848
6849 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
6855 self.durable_commit_failed = true;
6856 if let WalSink::Shared(shared) = &self.wal {
6857 shared
6858 .poisoned
6859 .store(true, std::sync::atomic::Ordering::Relaxed);
6860 }
6861 }
6862
6863 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
6867 self.durable_commit_failed = true;
6868 }
6869
6870 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
6895 let mut best: Option<Row> = None;
6896 fn consider(best: &mut Option<Row>, row: Row, snapshot: Snapshot) {
6899 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6900 return;
6901 }
6902 if best.as_ref().is_none_or(|current| {
6903 Snapshot::version_is_newer(
6904 row.committed_epoch,
6905 row.commit_ts,
6906 current.committed_epoch,
6907 current.commit_ts,
6908 )
6909 }) {
6910 *best = Some(row);
6911 }
6912 }
6913 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
6914 consider(&mut best, row, snapshot);
6915 }
6916 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
6917 consider(&mut best, row, snapshot);
6918 }
6919 let decision = if self.run_lookup.complete {
6923 let active_ids: Vec<u128> = self.run_refs.iter().map(|r| r.run_id).collect();
6928 let expected_fp = crate::run_lookup::compute_fingerprint(
6929 &active_ids,
6930 self.schema.schema_id,
6931 self.global_idx_epoch,
6932 active_ids.len() as u64,
6933 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
6934 );
6935 if self.run_lookup.fingerprint == expected_fp {
6936 match self.run_lookup.directory.as_ref() {
6937 Some(dir) => dir.decide(row_id),
6938 None => crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale,
6939 }
6940 } else {
6941 self.lookup_metrics
6942 .directory_lookup_fallback
6943 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6944 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale
6945 }
6946 } else {
6947 self.lookup_metrics
6948 .directory_lookup_fallback
6949 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6950 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale
6951 };
6952 match decision {
6953 crate::run_lookup::DirectoryLookupDecision::CompleteMiss => {
6954 self.lookup_metrics
6959 .directory_lookup_hit
6960 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6961 self.lookup_metrics
6962 .directory_complete_miss_total
6963 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6964 }
6965 crate::run_lookup::DirectoryLookupDecision::Candidates(locators) => {
6966 let mut opened: std::collections::HashSet<u128> = std::collections::HashSet::new();
6972 let mut early_stopped: u64 = 0;
6973 for locator in &locators {
6974 let stamp = best
6977 .as_ref()
6978 .map(|current| crate::run_lookup::VersionStamp {
6979 epoch: current.committed_epoch,
6980 hlc: current.commit_ts,
6981 });
6982 if let Some(s) = stamp {
6983 if !locator.can_contain_version_newer_than(s, snapshot) {
6984 early_stopped += 1;
6985 continue;
6986 }
6987 }
6988 if locator.is_impossible_for(snapshot) {
6989 continue;
6990 }
6991 if !opened.insert(locator.run_id) {
6992 continue;
6993 }
6994 self.lookup_metrics
6995 .directory_run_readers_opened
6996 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6997 self.lookup_metrics
6998 .get_run_opened
6999 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7000 let Ok(mut reader) = self.open_reader(locator.run_id) else {
7001 continue;
7002 };
7003 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
7006 continue;
7007 };
7008 consider(&mut best, row, snapshot);
7009 }
7010 if early_stopped > 0 {
7011 self.lookup_metrics
7012 .directory_early_stop_total
7013 .fetch_add(early_stopped, std::sync::atomic::Ordering::Relaxed);
7014 }
7015 self.lookup_metrics
7016 .directory_lookup_hit
7017 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7018 }
7019 crate::run_lookup::DirectoryLookupDecision::UnavailableOrStale => {
7020 for rr in &self.run_refs {
7022 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
7026 if row_id.0 < min_rid || row_id.0 > max_rid {
7027 self.lookup_metrics
7028 .get_run_skipped
7029 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7030 continue;
7031 }
7032 }
7033 self.lookup_metrics
7034 .get_run_opened
7035 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
7036 let Ok(mut reader) = self.open_reader(rr.run_id) else {
7037 continue;
7038 };
7039 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
7042 continue;
7043 };
7044 consider(&mut best, row, snapshot);
7045 }
7046 }
7047 }
7048 let now_nanos = unix_nanos_now();
7049 match best {
7050 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
7051 Some(r) => Some(r),
7052 None => None,
7053 }
7054 }
7055
7056 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
7060 self.visible_rows_at_time(snapshot, unix_nanos_now())
7061 }
7062
7063 #[doc(hidden)]
7066 pub fn visible_rows_controlled(
7067 &self,
7068 snapshot: Snapshot,
7069 control: &crate::ExecutionControl,
7070 ) -> Result<Vec<Row>> {
7071 let mut out = Vec::new();
7072 self.for_each_visible_row_controlled(snapshot, control, |row| {
7073 out.push(row);
7074 Ok(())
7075 })?;
7076 Ok(out)
7077 }
7078
7079 #[doc(hidden)]
7082 pub fn for_each_visible_row_controlled<F>(
7083 &self,
7084 snapshot: Snapshot,
7085 control: &crate::ExecutionControl,
7086 mut visit: F,
7087 ) -> Result<()>
7088 where
7089 F: FnMut(Row) -> Result<()>,
7090 {
7091 let setup_start = std::time::Instant::now();
7098 let mut versions_examined: u64 = 0;
7099 let mut rows_emitted: u64 = 0;
7100 let mut checkpoints: u64 = 0;
7101 let mut first_row_recorded = false;
7102 let mut sources: Vec<ControlledVisibleSource<'_>> =
7103 Vec::with_capacity(self.run_refs.len() + 2);
7104 control.checkpoint()?;
7105 checkpoints += 1;
7106 if !self.memtable.is_empty() {
7109 sources.push(ControlledVisibleSource::memtable_cursor(
7110 self.memtable.newest_visible_iter(&snapshot),
7111 ));
7112 }
7113 control.checkpoint()?;
7114 checkpoints += 1;
7115 if !self.mutable_run.is_empty() {
7116 sources.push(ControlledVisibleSource::mutable_run_cursor(
7117 self.mutable_run.newest_visible_iter(&snapshot),
7118 ));
7119 }
7120 for run in &self.run_refs {
7121 control.checkpoint()?;
7122 checkpoints += 1;
7123 let reader = self.open_reader(run.run_id)?;
7124 sources.push(ControlledVisibleSource::run(
7125 reader.into_visible_version_cursor_at(snapshot)?,
7126 ));
7127 }
7128 let first_row_start = setup_start;
7131 let _setup_us = setup_start.elapsed().as_micros() as u64;
7132 let now_nanos = unix_nanos_now();
7133 let mut first_row_us: u64 = 0;
7134 let result = merge_controlled_visible_sources(
7135 &mut sources,
7136 control,
7137 |row| self.row_expired_at(row, now_nanos),
7138 |row| {
7139 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
7140 return Ok(());
7141 }
7142 versions_examined += 1;
7143 if !first_row_recorded {
7144 first_row_us = first_row_start.elapsed().as_micros() as u64;
7145 first_row_recorded = true;
7146 }
7147 rows_emitted += 1;
7148 visit(row)
7149 },
7150 );
7151 let time_to_first_row_us = if first_row_recorded { first_row_us } else { 0 };
7152 crate::trace::QueryTrace::record(|t| {
7153 t.controlled_scan_versions_examined = t
7154 .controlled_scan_versions_examined
7155 .saturating_add(versions_examined as usize);
7156 t.controlled_scan_rows_emitted = t
7157 .controlled_scan_rows_emitted
7158 .saturating_add(rows_emitted as usize);
7159 t.controlled_scan_checkpoints = t
7160 .controlled_scan_checkpoints
7161 .saturating_add(checkpoints as usize);
7162 t.controlled_scan_time_to_first_row_us = t
7163 .controlled_scan_time_to_first_row_us
7164 .saturating_add(time_to_first_row_us);
7165 t.controlled_scan_setup_time_us =
7166 t.controlled_scan_setup_time_us.saturating_add(_setup_us);
7167 });
7168 result
7169 }
7170
7171 #[doc(hidden)]
7172 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
7173 let mut best: HashMap<u64, Row> = HashMap::new();
7174 let mut fold = |row: Row| {
7175 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
7176 return;
7177 }
7178 best.entry(row.row_id.0)
7179 .and_modify(|existing| {
7180 if Snapshot::version_is_newer(
7181 row.committed_epoch,
7182 row.commit_ts,
7183 existing.committed_epoch,
7184 existing.commit_ts,
7185 ) {
7186 *existing = row.clone();
7187 }
7188 })
7189 .or_insert(row);
7190 };
7191 for row in self.memtable.visible_versions_at(snapshot) {
7192 fold(row);
7193 }
7194 for row in self.mutable_run.visible_versions_at(snapshot) {
7195 fold(row);
7196 }
7197 for rr in &self.run_refs {
7198 let mut reader = self.open_reader(rr.run_id)?;
7199 for row in reader.visible_versions_at(snapshot)? {
7201 fold(row);
7202 }
7203 }
7204 let mut out: Vec<Row> = best
7205 .into_values()
7206 .filter_map(|r| {
7207 if r.deleted || self.row_expired_at(&r, now_nanos) {
7208 None
7209 } else {
7210 Some(r)
7211 }
7212 })
7213 .collect();
7214 out.sort_by_key(|r| r.row_id);
7215 Ok(out)
7216 }
7217
7218 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
7225 if self.ttl.is_none()
7226 && self.memtable.is_empty()
7227 && self.mutable_run.is_empty()
7228 && self.run_refs.len() == 1
7229 {
7230 let rr = self.run_refs[0].clone();
7231 let mut reader = self.open_reader(rr.run_id)?;
7232 let idxs = reader.visible_indices_at(snapshot)?;
7233 let mut cols = Vec::with_capacity(self.schema.columns.len());
7234 for cdef in &self.schema.columns {
7235 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
7236 }
7237 return Ok(cols);
7238 }
7239 let rows = self.visible_rows(snapshot)?;
7241 let mut cols: Vec<(u16, Vec<Value>)> = self
7242 .schema
7243 .columns
7244 .iter()
7245 .map(|c| (c.id, Vec::with_capacity(rows.len())))
7246 .collect();
7247 for r in &rows {
7248 for (cid, vec) in cols.iter_mut() {
7249 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
7250 }
7251 }
7252 Ok(cols)
7253 }
7254
7255 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
7257 let row_id = self.hot.get(key)?;
7258 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
7259 Some(row_id)
7260 } else {
7261 None
7262 }
7263 }
7264
7265 pub fn lookup_metrics_snapshot(&self) -> LookupMetricsSnapshot {
7269 let mut snap = self.lookup_metrics.snapshot();
7270 let cache = self.result_cache.lock();
7271 let (mem, disk, miss, write_us) = cache.cache_counters();
7272 snap.result_cache_memory_hit = mem;
7273 snap.result_cache_disk_hit = disk;
7274 snap.result_cache_miss = miss;
7275 snap.result_cache_persistent_write_us = write_us;
7276 if let Some(writer_snap) = cache.persist_snapshot() {
7279 snap.result_cache_persist_enqueued_total =
7280 writer_snap.result_cache_persist_enqueued_total;
7281 snap.result_cache_persist_coalesced_total =
7282 writer_snap.result_cache_persist_coalesced_total;
7283 snap.result_cache_persist_dropped_store_total =
7284 writer_snap.result_cache_persist_dropped_store_total;
7285 snap.result_cache_persist_remove_total = writer_snap.result_cache_persist_remove_total;
7286 snap.result_cache_persist_stale_store_skipped_total =
7287 writer_snap.result_cache_persist_stale_store_skipped_total;
7288 snap.result_cache_persist_errors_total = writer_snap.result_cache_persist_errors_total;
7289 snap.result_cache_persist_shutdown_abandoned_total =
7290 writer_snap.result_cache_persist_shutdown_abandoned_total;
7291 snap.result_cache_persist_queue_depth = writer_snap.result_cache_persist_queue_depth;
7292 }
7293 snap
7294 }
7295
7296 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
7301 self.query_at_with_allowed(q, self.snapshot(), None)
7302 }
7303
7304 pub fn query_controlled(
7307 &mut self,
7308 q: &crate::query::Query,
7309 control: &crate::ExecutionControl,
7310 ) -> Result<Vec<Row>> {
7311 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
7312 }
7313
7314 pub fn query_at_with_allowed(
7317 &mut self,
7318 q: &crate::query::Query,
7319 snapshot: Snapshot,
7320 allowed: Option<&std::collections::HashSet<RowId>>,
7321 ) -> Result<Vec<Row>> {
7322 self.query_at_with_allowed_after(q, snapshot, allowed, None)
7323 }
7324
7325 #[doc(hidden)]
7326 pub fn query_at_with_allowed_controlled(
7327 &mut self,
7328 q: &crate::query::Query,
7329 snapshot: Snapshot,
7330 allowed: Option<&std::collections::HashSet<RowId>>,
7331 control: &crate::ExecutionControl,
7332 ) -> Result<Vec<Row>> {
7333 self.require_select()?;
7334 self.ensure_indexes_complete_controlled(control, || true)?;
7335 self.validate_native_query(q)?;
7336 self.query_conditions_at(
7337 &q.conditions,
7338 snapshot,
7339 allowed,
7340 q.limit,
7341 q.offset,
7342 None,
7343 unix_nanos_now(),
7344 Some(control),
7345 )
7346 }
7347
7348 #[doc(hidden)]
7349 pub fn query_at_with_allowed_after(
7350 &mut self,
7351 q: &crate::query::Query,
7352 snapshot: Snapshot,
7353 allowed: Option<&std::collections::HashSet<RowId>>,
7354 after_row_id: Option<RowId>,
7355 ) -> Result<Vec<Row>> {
7356 self.query_at_with_allowed_after_at_time(
7357 q,
7358 snapshot,
7359 allowed,
7360 after_row_id,
7361 unix_nanos_now(),
7362 )
7363 }
7364
7365 #[doc(hidden)]
7366 pub fn query_at_with_allowed_after_at_time(
7367 &mut self,
7368 q: &crate::query::Query,
7369 snapshot: Snapshot,
7370 allowed: Option<&std::collections::HashSet<RowId>>,
7371 after_row_id: Option<RowId>,
7372 query_time_nanos: i64,
7373 ) -> Result<Vec<Row>> {
7374 self.require_select()?;
7375 self.ensure_indexes_complete()?;
7376 self.validate_native_query(q)?;
7377 self.query_conditions_at(
7378 &q.conditions,
7379 snapshot,
7380 allowed,
7381 q.limit,
7382 q.offset,
7383 after_row_id,
7384 query_time_nanos,
7385 None,
7386 )
7387 }
7388
7389 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
7390 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7391 return Err(MongrelError::InvalidArgument(format!(
7392 "query exceeds {} conditions",
7393 crate::query::MAX_HARD_CONDITIONS
7394 )));
7395 }
7396 if let Some(limit) = q.limit {
7397 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
7398 return Err(MongrelError::InvalidArgument(format!(
7399 "query limit must be between 1 and {}",
7400 crate::query::MAX_FINAL_LIMIT
7401 )));
7402 }
7403 }
7404 if q.offset > crate::query::MAX_QUERY_OFFSET {
7405 return Err(MongrelError::InvalidArgument(format!(
7406 "query offset exceeds {}",
7407 crate::query::MAX_QUERY_OFFSET
7408 )));
7409 }
7410 Ok(())
7411 }
7412
7413 #[doc(hidden)]
7416 pub fn query_all_at(
7417 &mut self,
7418 conditions: &[crate::query::Condition],
7419 snapshot: Snapshot,
7420 ) -> Result<Vec<Row>> {
7421 self.require_select()?;
7422 self.ensure_indexes_complete()?;
7423 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7424 return Err(MongrelError::InvalidArgument(format!(
7425 "query exceeds {} conditions",
7426 crate::query::MAX_HARD_CONDITIONS
7427 )));
7428 }
7429 self.query_conditions_at(
7430 conditions,
7431 snapshot,
7432 None,
7433 None,
7434 0,
7435 None,
7436 unix_nanos_now(),
7437 None,
7438 )
7439 }
7440
7441 #[allow(clippy::too_many_arguments)]
7442 fn query_conditions_at(
7443 &self,
7444 conditions: &[crate::query::Condition],
7445 snapshot: Snapshot,
7446 allowed: Option<&std::collections::HashSet<RowId>>,
7447 limit: Option<usize>,
7448 offset: usize,
7449 after_row_id: Option<RowId>,
7450 query_time_nanos: i64,
7451 control: Option<&crate::ExecutionControl>,
7452 ) -> Result<Vec<Row>> {
7453 control
7454 .map(crate::ExecutionControl::checkpoint)
7455 .transpose()?;
7456 crate::trace::QueryTrace::record(|t| {
7457 t.run_count = self.run_refs.len();
7458 t.memtable_rows = self.memtable.len();
7459 t.mutable_run_rows = self.mutable_run.len();
7460 });
7461 if conditions.is_empty() {
7465 crate::trace::QueryTrace::record(|t| {
7466 t.scan_mode = crate::trace::ScanMode::Materialized;
7467 t.row_materialized = true;
7468 });
7469 let mut rows = match control {
7470 Some(control) => self.visible_rows_controlled(snapshot, control)?,
7471 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
7472 };
7473 if let Some(allowed) = allowed {
7474 let mut filtered = Vec::with_capacity(rows.len());
7475 for (index, row) in rows.into_iter().enumerate() {
7476 if index & 255 == 0 {
7477 control
7478 .map(crate::ExecutionControl::checkpoint)
7479 .transpose()?;
7480 }
7481 if allowed.contains(&row.row_id) {
7482 filtered.push(row);
7483 }
7484 }
7485 rows = filtered;
7486 }
7487 if let Some(after_row_id) = after_row_id {
7488 rows.retain(|row| row.row_id > after_row_id);
7489 }
7490 rows.drain(..offset.min(rows.len()));
7491 if let Some(limit) = limit {
7492 rows.truncate(limit);
7493 }
7494 return Ok(rows);
7495 }
7496 crate::trace::QueryTrace::record(|t| {
7497 t.conditions_pushed = conditions.len();
7498 t.scan_mode = crate::trace::ScanMode::Materialized;
7499 t.row_materialized = true;
7500 });
7501 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
7508 ordered.sort_by_key(|c| condition_cost_rank(c));
7509 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
7510 for c in &ordered {
7511 control
7512 .map(crate::ExecutionControl::checkpoint)
7513 .transpose()?;
7514 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
7515 let empty = s.is_empty();
7516 sets.push(s);
7517 if empty {
7518 break;
7519 }
7520 }
7521 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7522 if let Some(allowed) = allowed {
7523 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
7524 }
7525 if let Some(after_row_id) = after_row_id {
7526 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
7527 rids.drain(..first);
7528 }
7529 rids.drain(..offset.min(rids.len()));
7530 if let Some(limit) = limit {
7531 rids.truncate(limit);
7532 }
7533 control
7534 .map(crate::ExecutionControl::checkpoint)
7535 .transpose()?;
7536 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
7537 }
7538
7539 pub fn retrieve(
7541 &mut self,
7542 retriever: &crate::query::Retriever,
7543 ) -> Result<Vec<crate::query::RetrieverHit>> {
7544 self.retrieve_with_allowed(retriever, None)
7545 }
7546
7547 pub fn retrieve_at(
7548 &mut self,
7549 retriever: &crate::query::Retriever,
7550 snapshot: Snapshot,
7551 allowed: Option<&std::collections::HashSet<RowId>>,
7552 ) -> Result<Vec<crate::query::RetrieverHit>> {
7553 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
7554 }
7555
7556 pub fn retrieve_with_allowed(
7559 &mut self,
7560 retriever: &crate::query::Retriever,
7561 allowed: Option<&std::collections::HashSet<RowId>>,
7562 ) -> Result<Vec<crate::query::RetrieverHit>> {
7563 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
7564 }
7565
7566 pub fn retrieve_at_with_allowed(
7567 &mut self,
7568 retriever: &crate::query::Retriever,
7569 snapshot: Snapshot,
7570 allowed: Option<&std::collections::HashSet<RowId>>,
7571 ) -> Result<Vec<crate::query::RetrieverHit>> {
7572 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
7573 }
7574
7575 pub fn retrieve_at_with_allowed_and_context(
7576 &mut self,
7577 retriever: &crate::query::Retriever,
7578 snapshot: Snapshot,
7579 allowed: Option<&std::collections::HashSet<RowId>>,
7580 context: Option<&crate::query::AiExecutionContext>,
7581 ) -> Result<Vec<crate::query::RetrieverHit>> {
7582 self.require_select()?;
7583 self.ensure_indexes_complete()?;
7584 self.validate_retriever(retriever)?;
7585 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
7586 }
7587
7588 pub fn retrieve_at_with_candidate_authorization_and_context(
7589 &mut self,
7590 retriever: &crate::query::Retriever,
7591 snapshot: Snapshot,
7592 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7593 context: Option<&crate::query::AiExecutionContext>,
7594 ) -> Result<Vec<crate::query::RetrieverHit>> {
7595 self.require_select()?;
7596 self.ensure_indexes_complete()?;
7597 self.retrieve_at_with_candidate_authorization_on_generation(
7598 retriever,
7599 snapshot,
7600 authorization,
7601 context,
7602 )
7603 }
7604
7605 #[doc(hidden)]
7606 pub fn retrieve_at_with_candidate_authorization_on_generation(
7607 &self,
7608 retriever: &crate::query::Retriever,
7609 snapshot: Snapshot,
7610 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7611 context: Option<&crate::query::AiExecutionContext>,
7612 ) -> Result<Vec<crate::query::RetrieverHit>> {
7613 self.require_select()?;
7614 self.validate_retriever(retriever)?;
7615 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
7616 }
7617
7618 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
7619 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
7620 let (column_id, k) = match retriever {
7621 Retriever::Ann {
7622 column_id,
7623 query,
7624 k,
7625 } => {
7626 let index = self.ann.get(column_id).ok_or_else(|| {
7627 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
7628 })?;
7629 if query.len() != index.dim() {
7630 return Err(MongrelError::InvalidArgument(format!(
7631 "ANN query dimension must be {}, got {}",
7632 index.dim(),
7633 query.len()
7634 )));
7635 }
7636 if query.iter().any(|value| !value.is_finite()) {
7637 return Err(MongrelError::InvalidArgument(
7638 "ANN query values must be finite".into(),
7639 ));
7640 }
7641 (*column_id, *k)
7642 }
7643 Retriever::Sparse {
7644 column_id,
7645 query,
7646 k,
7647 } => {
7648 if !self.sparse.contains_key(column_id) {
7649 return Err(MongrelError::InvalidArgument(format!(
7650 "column {column_id} has no Sparse index"
7651 )));
7652 }
7653 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
7654 return Err(MongrelError::InvalidArgument(
7655 "Sparse query must be non-empty with finite weights".into(),
7656 ));
7657 }
7658 if query.len() > MAX_SPARSE_TERMS {
7659 return Err(MongrelError::InvalidArgument(format!(
7660 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
7661 )));
7662 }
7663 (*column_id, *k)
7664 }
7665 Retriever::MinHash {
7666 column_id,
7667 members,
7668 k,
7669 } => {
7670 if !self.minhash.contains_key(column_id) {
7671 return Err(MongrelError::InvalidArgument(format!(
7672 "column {column_id} has no MinHash index"
7673 )));
7674 }
7675 if members.is_empty() {
7676 return Err(MongrelError::InvalidArgument(
7677 "MinHash members must not be empty".into(),
7678 ));
7679 }
7680 if members.len() > MAX_SET_MEMBERS {
7681 return Err(MongrelError::InvalidArgument(format!(
7682 "MinHash query exceeds {MAX_SET_MEMBERS} members"
7683 )));
7684 }
7685 let mut total_bytes = 0usize;
7686 for member in members {
7687 let bytes = member.encoded_len();
7688 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
7689 return Err(MongrelError::InvalidArgument(format!(
7690 "MinHash member exceeds {} bytes",
7691 crate::query::MAX_SET_MEMBER_BYTES
7692 )));
7693 }
7694 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
7695 MongrelError::InvalidArgument("MinHash input size overflow".into())
7696 })?;
7697 }
7698 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
7699 return Err(MongrelError::InvalidArgument(format!(
7700 "MinHash input exceeds {} bytes",
7701 crate::query::MAX_SET_INPUT_BYTES
7702 )));
7703 }
7704 (*column_id, *k)
7705 }
7706 };
7707 if k == 0 {
7708 return Err(MongrelError::InvalidArgument(
7709 "retriever k must be > 0".into(),
7710 ));
7711 }
7712 if k > MAX_RETRIEVER_K {
7713 return Err(MongrelError::InvalidArgument(format!(
7714 "retriever k exceeds {MAX_RETRIEVER_K}"
7715 )));
7716 }
7717 debug_assert!(self
7718 .schema
7719 .columns
7720 .iter()
7721 .any(|column| column.id == column_id));
7722 Ok(())
7723 }
7724
7725 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
7726 use crate::query::Condition;
7727 match condition {
7728 Condition::Ann {
7729 column_id,
7730 query,
7731 k,
7732 } => self.validate_retriever(&crate::query::Retriever::Ann {
7733 column_id: *column_id,
7734 query: query.clone(),
7735 k: *k,
7736 }),
7737 Condition::SparseMatch {
7738 column_id,
7739 query,
7740 k,
7741 } => self.validate_retriever(&crate::query::Retriever::Sparse {
7742 column_id: *column_id,
7743 query: query.clone(),
7744 k: *k,
7745 }),
7746 Condition::MinHashSimilar {
7747 column_id,
7748 query,
7749 k,
7750 } => {
7751 if !self.minhash.contains_key(column_id) {
7752 return Err(MongrelError::InvalidArgument(format!(
7753 "column {column_id} has no MinHash index"
7754 )));
7755 }
7756 if query.is_empty() || *k == 0 {
7757 return Err(MongrelError::InvalidArgument(
7758 "MinHash query must be non-empty and k must be > 0".into(),
7759 ));
7760 }
7761 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
7762 {
7763 return Err(MongrelError::InvalidArgument(format!(
7764 "MinHash query must have <= {} members and k <= {}",
7765 crate::query::MAX_SET_MEMBERS,
7766 crate::query::MAX_RETRIEVER_K
7767 )));
7768 }
7769 Ok(())
7770 }
7771 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
7772 Err(MongrelError::InvalidArgument(format!(
7773 "bitmap IN exceeds {} values",
7774 crate::query::MAX_SET_MEMBERS
7775 )))
7776 }
7777 Condition::FmContainsAll { patterns, .. }
7778 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
7779 {
7780 Err(MongrelError::InvalidArgument(format!(
7781 "FM query exceeds {} patterns",
7782 crate::query::MAX_HARD_CONDITIONS
7783 )))
7784 }
7785 _ => Ok(()),
7786 }
7787 }
7788
7789 fn retrieve_filtered(
7790 &self,
7791 retriever: &crate::query::Retriever,
7792 snapshot: Snapshot,
7793 hard_filter: Option<&RowIdSet>,
7794 allowed: Option<&std::collections::HashSet<RowId>>,
7795 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7796 context: Option<&crate::query::AiExecutionContext>,
7797 ) -> Result<Vec<crate::query::RetrieverHit>> {
7798 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
7799 let started = std::time::Instant::now();
7800 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
7801 Retriever::Ann {
7802 column_id,
7803 query,
7804 k,
7805 } => {
7806 let Some(index) = self.ann.get(column_id) else {
7807 return Ok(Vec::new());
7808 };
7809 let cap = ann_candidate_cap(index.len(), context);
7810 crate::trace::QueryTrace::record(|trace| trace.candidate_cap = cap);
7811 if cap == 0 {
7812 return Ok(Vec::new());
7813 }
7814 let mut breadth = (*k).max(1).min(cap);
7815 let mut eligibility = std::collections::HashMap::new();
7816 let mut raw_candidates_total: usize = 0;
7821 let mut visibility_rejected_total: usize = 0;
7822 let mut authorization_rejected_total: usize = 0;
7823 let mut candidate_cap_hit_final = false;
7824 let mut filtered = loop {
7825 let mut seen = std::collections::HashSet::new();
7826 if let Some(context) = context {
7827 context.checkpoint()?;
7828 }
7829 let raw = index.search_with_context(query, breadth, context)?;
7830 eprintln!(
7831 "MASTER DEBUG ANN raw candidates: {:?}",
7832 raw.iter().map(|(r, _)| r.0).collect::<Vec<_>>()
7833 );
7834 crate::trace::QueryTrace::record(|trace| {
7835 trace.raw_candidates = raw.len();
7836 let unique = raw
7837 .iter()
7838 .map(|(row_id, _)| *row_id)
7839 .collect::<std::collections::HashSet<_>>()
7840 .len();
7841 trace.unique_candidates = unique;
7842 trace.duplicate_candidates = raw.len().saturating_sub(unique);
7843 trace.authorization_rejected = raw
7844 .iter()
7845 .filter(|(row_id, _)| {
7846 allowed.is_some_and(|allowed| !allowed.contains(row_id))
7847 })
7848 .count();
7849 trace.hard_filter_rejected = raw
7850 .iter()
7851 .filter(|(row_id, _)| {
7852 hard_filter.is_some_and(|filter| !filter.contains(row_id.0))
7853 })
7854 .count();
7855 });
7856 raw_candidates_total = raw_candidates_total.saturating_add(raw.len());
7857 let unchecked: Vec<_> = raw
7858 .iter()
7859 .map(|(row_id, _)| *row_id)
7860 .filter(|row_id| !eligibility.contains_key(row_id))
7861 .filter(|row_id| {
7862 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7863 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7864 })
7865 .collect();
7866 let eligible = self.eligible_and_authorized_candidate_ids(
7867 &unchecked,
7868 *column_id,
7869 snapshot,
7870 candidate_authorization,
7871 context,
7872 )?;
7873 for row_id in &unchecked {
7874 eligibility.insert(*row_id, eligible.contains(row_id));
7875 }
7876 let new_visibility_or_auth_rejected = unchecked
7879 .iter()
7880 .filter(|row_id| !eligible.contains(row_id))
7881 .count();
7882 let mut new_auth_rejected = 0usize;
7887 if let Some(authorization) = candidate_authorization {
7888 if authorization.security.rls_enabled(authorization.table)
7889 && !authorization.principal.is_admin
7890 {
7891 new_auth_rejected = new_visibility_or_auth_rejected;
7897 }
7898 }
7899 let new_visibility_rejected =
7900 new_visibility_or_auth_rejected.saturating_sub(new_auth_rejected);
7901 visibility_rejected_total =
7902 visibility_rejected_total.saturating_add(new_visibility_rejected);
7903 authorization_rejected_total =
7904 authorization_rejected_total.saturating_add(new_auth_rejected);
7905 let filtered: Vec<_> = raw
7906 .into_iter()
7907 .filter(|(row_id, _)| {
7908 seen.insert(*row_id)
7909 && eligibility.get(row_id).copied().unwrap_or(false)
7910 })
7911 .map(|(row_id, score)| {
7912 let score = match score {
7913 crate::index::AnnDistance::Hamming(d) => {
7914 RetrieverScore::AnnHammingDistance(d)
7915 }
7916 crate::index::AnnDistance::Cosine(d) => {
7917 RetrieverScore::AnnCosineDistance(d)
7918 }
7919 };
7920 (row_id, score)
7921 })
7922 .collect();
7923 if filtered.len() >= *k || breadth >= cap {
7924 if filtered.len() < *k && index.len() > cap && breadth >= cap {
7925 crate::trace::QueryTrace::record(|trace| {
7926 trace.ann_candidate_cap_hit = true;
7927 trace.candidate_cap_hit = true;
7928 });
7929 candidate_cap_hit_final = true;
7930 }
7931 break filtered;
7932 }
7933 breadth = breadth.saturating_mul(2).min(cap);
7934 };
7935 let final_hits = filtered.len();
7938 let index_len = index.len();
7939 crate::trace::QueryTrace::record(|trace| {
7940 trace.candidate_cap = cap;
7941 trace.candidate_cap_hit = candidate_cap_hit_final;
7942 trace.ann_candidate_cap_hit = candidate_cap_hit_final;
7943 trace.raw_candidates =
7944 trace.raw_candidates.saturating_add(raw_candidates_total);
7945 trace.visibility_rejected = trace
7946 .visibility_rejected
7947 .saturating_add(visibility_rejected_total);
7948 trace.authorization_rejected = trace
7949 .authorization_rejected
7950 .saturating_add(authorization_rejected_total);
7951 trace.final_hits = trace.final_hits.saturating_add(final_hits);
7952 let _ = index_len;
7955 });
7956 filtered.truncate(*k);
7957 filtered
7958 }
7959 Retriever::Sparse {
7960 column_id,
7961 query,
7962 k,
7963 } => self
7964 .sparse
7965 .get(column_id)
7966 .map(|index| -> Result<Vec<_>> {
7967 let mut breadth = (*k).max(1);
7968 let mut eligibility = std::collections::HashMap::new();
7969 loop {
7970 if let Some(context) = context {
7971 context.checkpoint()?;
7972 }
7973 let raw = index.search_with_context(query, breadth, context)?;
7974 let unchecked: Vec<_> = raw
7975 .iter()
7976 .map(|(row_id, _)| *row_id)
7977 .filter(|row_id| !eligibility.contains_key(row_id))
7978 .filter(|row_id| {
7979 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7980 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7981 })
7982 .collect();
7983 let eligible = self.eligible_and_authorized_candidate_ids(
7984 &unchecked,
7985 *column_id,
7986 snapshot,
7987 candidate_authorization,
7988 context,
7989 )?;
7990 for row_id in unchecked {
7991 eligibility.insert(row_id, eligible.contains(&row_id));
7992 }
7993 let filtered: Vec<_> = raw
7994 .iter()
7995 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
7996 .take(*k)
7997 .map(|(row_id, score)| {
7998 (*row_id, RetrieverScore::SparseDotProduct(*score))
7999 })
8000 .collect();
8001 if filtered.len() >= *k || raw.len() < breadth {
8002 break Ok(filtered);
8003 }
8004 let next = breadth.saturating_mul(2);
8005 if next == breadth {
8006 break Ok(filtered);
8007 }
8008 breadth = next;
8009 }
8010 })
8011 .transpose()?
8012 .unwrap_or_default(),
8013 Retriever::MinHash {
8014 column_id,
8015 members,
8016 k,
8017 } => self
8018 .minhash
8019 .get(column_id)
8020 .map(|index| -> Result<Vec<_>> {
8021 let mut hashes = Vec::with_capacity(members.len());
8022 for member in members {
8023 if let Some(context) = context {
8024 context.consume(crate::query::work_units(
8025 member.encoded_len(),
8026 crate::query::PARSE_WORK_QUANTUM,
8027 ))?;
8028 }
8029 hashes.push(member.hash_v1());
8030 }
8031 let mut breadth = (*k).max(1);
8032 let mut eligibility = std::collections::HashMap::new();
8033 loop {
8034 if let Some(context) = context {
8035 context.checkpoint()?;
8036 }
8037 let raw = index.search_with_context(&hashes, breadth, context)?;
8038 let unchecked: Vec<_> = raw
8039 .iter()
8040 .map(|(row_id, _)| *row_id)
8041 .filter(|row_id| !eligibility.contains_key(row_id))
8042 .filter(|row_id| {
8043 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
8044 && allowed.is_none_or(|allowed| allowed.contains(row_id))
8045 })
8046 .collect();
8047 let eligible = self.eligible_and_authorized_candidate_ids(
8048 &unchecked,
8049 *column_id,
8050 snapshot,
8051 candidate_authorization,
8052 context,
8053 )?;
8054 for row_id in unchecked {
8055 eligibility.insert(row_id, eligible.contains(&row_id));
8056 }
8057 let filtered: Vec<_> = raw
8058 .iter()
8059 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
8060 .take(*k)
8061 .map(|(row_id, score)| {
8062 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
8063 })
8064 .collect();
8065 if filtered.len() >= *k || raw.len() < breadth {
8066 break Ok(filtered);
8067 }
8068 let next = breadth.saturating_mul(2);
8069 if next == breadth {
8070 break Ok(filtered);
8071 }
8072 breadth = next;
8073 }
8074 })
8075 .transpose()?
8076 .unwrap_or_default(),
8077 };
8078 let requested_k = match retriever {
8079 Retriever::Ann { k, .. }
8080 | Retriever::Sparse { k, .. }
8081 | Retriever::MinHash { k, .. } => *k,
8082 };
8083 crate::trace::QueryTrace::record(|trace| {
8084 trace.final_hits = scored.len();
8085 if scored.len() < requested_k {
8086 trace.underfill_reason = Some(if trace.candidate_cap_hit {
8087 "candidate_cap"
8088 } else {
8089 "eligible_candidates_exhausted"
8090 });
8091 }
8092 });
8093 let elapsed = started.elapsed().as_nanos() as u64;
8094 crate::trace::QueryTrace::record(|trace| {
8095 match retriever {
8096 Retriever::Ann { .. } => {
8097 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
8098 if let Retriever::Ann { column_id, .. } = retriever {
8099 if let Some(index) = self.ann.get(column_id) {
8100 trace.ann_algorithm = Some(index.algorithm());
8101 trace.ann_quantization = Some(index.quantization());
8102 trace.ann_backend = Some(index.backend_name());
8103 }
8104 }
8105 }
8106 Retriever::Sparse { .. } => {
8107 trace.sparse_candidate_nanos =
8108 trace.sparse_candidate_nanos.saturating_add(elapsed)
8109 }
8110 Retriever::MinHash { .. } => {
8111 trace.minhash_candidate_nanos =
8112 trace.minhash_candidate_nanos.saturating_add(elapsed)
8113 }
8114 }
8115 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
8116 });
8117 Ok(scored
8118 .into_iter()
8119 .enumerate()
8120 .map(|(rank, (row_id, score))| RetrieverHit {
8121 row_id,
8122 rank: rank + 1,
8123 score,
8124 })
8125 .collect())
8126 }
8127
8128 fn eligible_candidate_ids(
8129 &self,
8130 candidates: &[RowId],
8131 _column_id: u16,
8132 snapshot: Snapshot,
8133 context: Option<&crate::query::AiExecutionContext>,
8134 ) -> Result<std::collections::HashSet<RowId>> {
8135 let lookup_snapshot = if self.is_shared()
8145 || (self.pending_put_cols.is_empty() && self.pending_delete_rids.is_empty())
8146 || snapshot.epoch.0 >= self.pending_epoch().0
8147 {
8148 snapshot
8149 } else {
8150 Snapshot::at(self.pending_epoch())
8151 };
8152 let mut readers: Vec<_> = self
8153 .run_refs
8154 .iter()
8155 .map(|run| self.open_reader(run.run_id))
8156 .collect::<Result<_>>()?;
8157 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8158 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
8159 for &row_id in candidates {
8160 if let Some(context) = context {
8161 context.consume(1)?;
8162 }
8163 let mem = self.memtable.get_version_at(row_id, lookup_snapshot);
8164 let mutable = self.mutable_run.get_version_at(row_id, lookup_snapshot);
8165 let overlay = match (mem, mutable) {
8166 (Some(left), Some(right)) => Some(
8167 if Snapshot::version_is_newer(
8168 left.1.committed_epoch,
8169 left.1.commit_ts,
8170 right.1.committed_epoch,
8171 right.1.commit_ts,
8172 ) {
8173 left
8174 } else {
8175 right
8176 },
8177 ),
8178 (Some(value), None) | (None, Some(value)) => Some(value),
8179 (None, None) => None,
8180 };
8181 if let Some((_, row)) = overlay {
8182 if !row.deleted && !self.row_expired_at(&row, now) {
8183 eligible.insert(row_id);
8184 }
8185 continue;
8186 }
8187 let mut best: Option<(crate::sorted_run::VersionVisibility, usize)> = None;
8193 for (index, reader) in readers.iter_mut().enumerate() {
8194 if let Some(visibility) =
8195 reader.get_version_visibility_full_at(row_id, lookup_snapshot)?
8196 {
8197 if best
8198 .as_ref()
8199 .map(|(current, _)| visibility.stamp.is_newer_than(current.stamp))
8200 .unwrap_or(true)
8201 {
8202 best = Some((visibility, index));
8203 }
8204 }
8205 }
8206 let Some((visibility, reader_index)) = best else {
8207 continue;
8208 };
8209 if visibility.deleted {
8210 continue;
8211 }
8212 if let Some(ttl) = self.ttl {
8213 if let Some(ttl_value) = readers[reader_index].get_version_column_full_at(
8214 row_id,
8215 lookup_snapshot,
8216 ttl.column_id,
8217 )? {
8218 if let Some(Value::Int64(timestamp)) = ttl_value.value {
8219 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
8220 continue;
8221 }
8222 }
8223 }
8224 }
8225 eligible.insert(row_id);
8226 }
8227 Ok(eligible)
8228 }
8229
8230 fn eligible_and_authorized_candidate_ids(
8231 &self,
8232 candidates: &[RowId],
8233 column_id: u16,
8234 snapshot: Snapshot,
8235 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8236 context: Option<&crate::query::AiExecutionContext>,
8237 ) -> Result<std::collections::HashSet<RowId>> {
8238 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
8239 let Some(authorization) = authorization else {
8240 return Ok(eligible);
8241 };
8242 let candidates: Vec<_> = eligible.into_iter().collect();
8243 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
8244 }
8245
8246 fn policy_allowed_candidate_ids(
8247 &self,
8248 candidates: &[RowId],
8249 snapshot: Snapshot,
8250 authorization: &crate::security::CandidateAuthorization<'_>,
8251 context: Option<&crate::query::AiExecutionContext>,
8252 ) -> Result<std::collections::HashSet<RowId>> {
8253 let started = std::time::Instant::now();
8254 if candidates.is_empty()
8255 || authorization.principal.is_admin
8256 || !authorization.security.rls_enabled(authorization.table)
8257 {
8258 return Ok(candidates.iter().copied().collect());
8259 }
8260 if let Some(context) = context {
8261 context.checkpoint()?;
8262 }
8263 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
8264 let mut rows: std::collections::HashMap<RowId, Row> = candidates
8265 .iter()
8266 .map(|row_id| {
8267 (
8268 *row_id,
8269 Row {
8270 row_id: *row_id,
8271 committed_epoch: snapshot.epoch,
8272 columns: std::collections::HashMap::new(),
8273 deleted: false,
8274 commit_ts: None,
8275 },
8276 )
8277 })
8278 .collect();
8279 let columns = authorization
8280 .security
8281 .select_policy_columns(authorization.table, authorization.principal);
8282 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8283 let mut decoded = 0usize;
8284 for column_id in &columns {
8285 if let Some(context) = context {
8286 context.checkpoint()?;
8287 }
8288 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8289 &row_ids, *column_id, snapshot, query_now, context,
8290 )? {
8291 if let Some(row) = rows.get_mut(&row_id) {
8292 row.columns.insert(*column_id, value);
8293 decoded = decoded.saturating_add(1);
8294 }
8295 }
8296 }
8297 if let Some(context) = context {
8298 context.consume(candidates.len().saturating_add(decoded))?;
8299 }
8300 let allowed = rows
8301 .into_values()
8302 .filter_map(|row| {
8303 authorization
8304 .security
8305 .row_allowed(
8306 authorization.table,
8307 crate::security::PolicyCommand::Select,
8308 &row,
8309 authorization.principal,
8310 false,
8311 )
8312 .then_some(row.row_id)
8313 })
8314 .collect();
8315 crate::trace::QueryTrace::record(|trace| {
8316 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
8317 trace.rls_policy_columns_decoded =
8318 trace.rls_policy_columns_decoded.saturating_add(decoded);
8319 trace.authorization_nanos = trace
8320 .authorization_nanos
8321 .saturating_add(started.elapsed().as_nanos() as u64);
8322 });
8323 Ok(allowed)
8324 }
8325
8326 pub fn search(
8328 &mut self,
8329 request: &crate::query::SearchRequest,
8330 ) -> Result<Vec<crate::query::SearchHit>> {
8331 self.search_with_allowed(request, None)
8332 }
8333
8334 pub fn search_at(
8335 &mut self,
8336 request: &crate::query::SearchRequest,
8337 snapshot: Snapshot,
8338 authorized: Option<&std::collections::HashSet<RowId>>,
8339 ) -> Result<Vec<crate::query::SearchHit>> {
8340 self.search_at_with_allowed(request, snapshot, authorized)
8341 }
8342
8343 pub fn search_with_allowed(
8344 &mut self,
8345 request: &crate::query::SearchRequest,
8346 authorized: Option<&std::collections::HashSet<RowId>>,
8347 ) -> Result<Vec<crate::query::SearchHit>> {
8348 self.search_at_with_allowed(request, self.snapshot(), authorized)
8349 }
8350
8351 pub fn search_at_with_allowed(
8352 &mut self,
8353 request: &crate::query::SearchRequest,
8354 snapshot: Snapshot,
8355 authorized: Option<&std::collections::HashSet<RowId>>,
8356 ) -> Result<Vec<crate::query::SearchHit>> {
8357 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
8358 }
8359
8360 pub fn search_at_with_allowed_and_context(
8361 &mut self,
8362 request: &crate::query::SearchRequest,
8363 snapshot: Snapshot,
8364 authorized: Option<&std::collections::HashSet<RowId>>,
8365 context: Option<&crate::query::AiExecutionContext>,
8366 ) -> Result<Vec<crate::query::SearchHit>> {
8367 self.ensure_indexes_complete()?;
8368 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
8369 }
8370
8371 pub fn search_at_with_candidate_authorization_and_context(
8372 &mut self,
8373 request: &crate::query::SearchRequest,
8374 snapshot: Snapshot,
8375 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8376 context: Option<&crate::query::AiExecutionContext>,
8377 ) -> Result<Vec<crate::query::SearchHit>> {
8378 self.ensure_indexes_complete()?;
8379 self.search_at_with_filters_and_context(
8380 request,
8381 snapshot,
8382 None,
8383 authorization,
8384 context,
8385 None,
8386 )
8387 }
8388
8389 #[doc(hidden)]
8390 pub fn search_at_with_candidate_authorization_on_generation(
8391 &self,
8392 request: &crate::query::SearchRequest,
8393 snapshot: Snapshot,
8394 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8395 context: Option<&crate::query::AiExecutionContext>,
8396 ) -> Result<Vec<crate::query::SearchHit>> {
8397 self.search_at_with_filters_and_context(
8398 request,
8399 snapshot,
8400 None,
8401 authorization,
8402 context,
8403 None,
8404 )
8405 }
8406
8407 #[doc(hidden)]
8408 pub fn search_at_with_candidate_authorization_on_generation_after(
8409 &self,
8410 request: &crate::query::SearchRequest,
8411 snapshot: Snapshot,
8412 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8413 context: Option<&crate::query::AiExecutionContext>,
8414 after: Option<crate::query::SearchAfter>,
8415 ) -> Result<Vec<crate::query::SearchHit>> {
8416 self.search_at_with_filters_and_context(
8417 request,
8418 snapshot,
8419 None,
8420 authorization,
8421 context,
8422 after,
8423 )
8424 }
8425
8426 fn search_at_with_filters_and_context(
8427 &self,
8428 request: &crate::query::SearchRequest,
8429 snapshot: Snapshot,
8430 authorized: Option<&std::collections::HashSet<RowId>>,
8431 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8432 context: Option<&crate::query::AiExecutionContext>,
8433 after: Option<crate::query::SearchAfter>,
8434 ) -> Result<Vec<crate::query::SearchHit>> {
8435 use crate::query::{
8436 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
8437 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
8438 };
8439 let total_started = std::time::Instant::now();
8440 let rank_offset = after.map_or(0, |after| after.returned_count);
8441 self.require_select()?;
8442 if request.limit == 0 {
8443 return Err(MongrelError::InvalidArgument(
8444 "search limit must be > 0".into(),
8445 ));
8446 }
8447 if request.limit > MAX_FINAL_LIMIT {
8448 return Err(MongrelError::InvalidArgument(format!(
8449 "search limit exceeds {MAX_FINAL_LIMIT}"
8450 )));
8451 }
8452 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
8453 return Err(MongrelError::InvalidArgument(
8454 "search-after score must be finite".into(),
8455 ));
8456 }
8457 if request.retrievers.is_empty() {
8458 return Err(MongrelError::InvalidArgument(
8459 "search requires at least one retriever".into(),
8460 ));
8461 }
8462 if request.retrievers.len() > MAX_RETRIEVERS {
8463 return Err(MongrelError::InvalidArgument(format!(
8464 "search exceeds {MAX_RETRIEVERS} retrievers"
8465 )));
8466 }
8467 if request.must.len() > MAX_HARD_CONDITIONS {
8468 return Err(MongrelError::InvalidArgument(format!(
8469 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
8470 )));
8471 }
8472 for condition in &request.must {
8473 self.validate_condition(condition)?;
8474 }
8475 if request.must.iter().any(|condition| {
8476 matches!(
8477 condition,
8478 Condition::Ann { .. }
8479 | Condition::SparseMatch { .. }
8480 | Condition::MinHashSimilar { .. }
8481 )
8482 }) {
8483 return Err(MongrelError::InvalidArgument(
8484 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
8485 .into(),
8486 ));
8487 }
8488 let mut names = std::collections::HashSet::new();
8489 for named in &request.retrievers {
8490 if named.name.is_empty()
8491 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
8492 || !names.insert(named.name.as_str())
8493 {
8494 return Err(MongrelError::InvalidArgument(format!(
8495 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
8496 crate::query::MAX_RETRIEVER_NAME_BYTES
8497 )));
8498 }
8499 if !named.weight.is_finite()
8500 || named.weight < 0.0
8501 || named.weight > MAX_RETRIEVER_WEIGHT
8502 {
8503 return Err(MongrelError::InvalidArgument(format!(
8504 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8505 )));
8506 }
8507 self.validate_retriever(&named.retriever)?;
8508 }
8509 let projection = request
8510 .projection
8511 .clone()
8512 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
8513 if projection.len() > MAX_PROJECTION_COLUMNS {
8514 return Err(MongrelError::InvalidArgument(format!(
8515 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
8516 )));
8517 }
8518 for column_id in &projection {
8519 if !self
8520 .schema
8521 .columns
8522 .iter()
8523 .any(|column| column.id == *column_id)
8524 {
8525 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
8526 }
8527 }
8528 if let Some(crate::query::Rerank::ExactVector {
8529 embedding_column,
8530 query,
8531 candidate_limit,
8532 weight,
8533 ..
8534 }) = &request.rerank
8535 {
8536 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
8537 {
8538 return Err(MongrelError::InvalidArgument(format!(
8539 "rerank candidate_limit must be between search limit and {}",
8540 crate::query::MAX_RETRIEVER_K
8541 )));
8542 }
8543 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
8544 return Err(MongrelError::InvalidArgument(format!(
8545 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8546 )));
8547 }
8548 let column = self
8549 .schema
8550 .columns
8551 .iter()
8552 .find(|column| column.id == *embedding_column)
8553 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
8554 let crate::schema::TypeId::Embedding { dim } = column.ty else {
8555 return Err(MongrelError::InvalidArgument(format!(
8556 "rerank column {embedding_column} is not an embedding"
8557 )));
8558 };
8559 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
8560 return Err(MongrelError::InvalidArgument(format!(
8561 "rerank query must contain {dim} finite values"
8562 )));
8563 }
8564 }
8565
8566 let hard_filter_started = std::time::Instant::now();
8567 let hard_filter = if request.must.is_empty() {
8568 None
8569 } else {
8570 let mut sets = Vec::with_capacity(request.must.len());
8571 for condition in &request.must {
8572 if let Some(context) = context {
8573 context.checkpoint()?;
8574 }
8575 sets.push(self.resolve_condition(condition, snapshot)?);
8576 }
8577 Some(RowIdSet::intersect_many(sets))
8578 };
8579 crate::trace::QueryTrace::record(|trace| {
8580 trace.hard_filter_nanos = trace
8581 .hard_filter_nanos
8582 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
8583 });
8584 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
8585 return Ok(Vec::new());
8586 }
8587
8588 let constant = match request.fusion {
8589 Fusion::ReciprocalRank { constant } => constant,
8590 };
8591 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
8592 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
8593 let mut fusion_nanos = 0u64;
8594 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
8595 std::collections::HashMap::new();
8596 for named in retrievers {
8597 if named.weight == 0.0 {
8598 continue;
8599 }
8600 if let Some(context) = context {
8601 context.checkpoint()?;
8602 }
8603 let hits = self.retrieve_filtered(
8604 &named.retriever,
8605 snapshot,
8606 hard_filter.as_ref(),
8607 authorized,
8608 candidate_authorization,
8609 context,
8610 )?;
8611 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
8612 let fusion_started = std::time::Instant::now();
8613 for hit in hits {
8614 if let Some(context) = context {
8615 context.consume(1)?;
8616 }
8617 let contribution = named.weight / (constant as f64 + hit.rank as f64);
8618 if !contribution.is_finite() {
8619 return Err(MongrelError::InvalidArgument(
8620 "retriever contribution must be finite".into(),
8621 ));
8622 }
8623 let max_fused_candidates = context.map_or(
8624 crate::query::MAX_FUSED_CANDIDATES,
8625 crate::query::AiExecutionContext::max_fused_candidates,
8626 );
8627 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
8628 return Err(MongrelError::WorkBudgetExceeded);
8629 }
8630 let entry = fused.entry(hit.row_id).or_default();
8631 entry.0 += contribution;
8632 if !entry.0.is_finite() {
8633 return Err(MongrelError::InvalidArgument(
8634 "fused score must be finite".into(),
8635 ));
8636 }
8637 entry.1.push(ComponentScore {
8638 retriever_name: retriever_name.clone(),
8639 rank: hit.rank,
8640 raw_score: hit.score,
8641 contribution,
8642 });
8643 }
8644 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
8645 }
8646 let union_size = fused.len();
8647 let mut ranked: Vec<_> = fused
8648 .into_iter()
8649 .map(|(row_id, (fused_score, components))| {
8650 (row_id, fused_score, components, None, fused_score)
8651 })
8652 .collect();
8653 let order = |(a_row, _, _, _, a_score): &(
8654 RowId,
8655 f64,
8656 Vec<ComponentScore>,
8657 Option<f32>,
8658 f64,
8659 ),
8660 (b_row, _, _, _, b_score): &(
8661 RowId,
8662 f64,
8663 Vec<ComponentScore>,
8664 Option<f32>,
8665 f64,
8666 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
8667 if let Some(crate::query::Rerank::ExactVector {
8668 embedding_column,
8669 query,
8670 metric,
8671 candidate_limit,
8672 weight,
8673 }) = &request.rerank
8674 {
8675 let fused_order = |(a_row, a_score, ..): &(
8676 RowId,
8677 f64,
8678 Vec<ComponentScore>,
8679 Option<f32>,
8680 f64,
8681 ),
8682 (b_row, b_score, ..): &(
8683 RowId,
8684 f64,
8685 Vec<ComponentScore>,
8686 Option<f32>,
8687 f64,
8688 )| {
8689 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
8690 };
8691 let selection_started = std::time::Instant::now();
8692 if let Some(context) = context {
8693 context.consume(ranked.len())?;
8694 }
8695 if ranked.len() > *candidate_limit {
8696 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
8697 ranked.truncate(*candidate_limit);
8698 }
8699 ranked.sort_by(fused_order);
8700 fusion_nanos =
8701 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
8702 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
8703 if let Some(context) = context {
8704 context.consume(row_ids.len())?;
8705 }
8706 let query_now =
8707 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8708 let gather_started = std::time::Instant::now();
8709 let vectors = self.values_for_rids_batch_at_with_context(
8710 &row_ids,
8711 *embedding_column,
8712 snapshot,
8713 query_now,
8714 context,
8715 )?;
8716 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
8717 let vector_work =
8718 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
8719 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
8720 if let Some(context) = context {
8721 context.consume(vector_work)?;
8722 }
8723 query
8724 .iter()
8725 .map(|value| f64::from(*value).powi(2))
8726 .sum::<f64>()
8727 .sqrt()
8728 } else {
8729 0.0
8730 };
8731 let score_started = std::time::Instant::now();
8732 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
8733 for (row_id, value) in vectors {
8734 if let Some(meta) = value.generated_embedding_metadata() {
8735 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
8736 continue;
8737 }
8738 }
8739 let Some(vector) = value.as_embedding() else {
8740 continue;
8741 };
8742 let score = match metric {
8743 crate::query::VectorMetric::DotProduct => {
8744 if let Some(context) = context {
8745 context.consume(vector_work)?;
8746 }
8747 query
8748 .iter()
8749 .zip(vector)
8750 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8751 .sum::<f64>()
8752 }
8753 crate::query::VectorMetric::Cosine => {
8754 if let Some(context) = context {
8755 context.consume(vector_work.saturating_mul(2))?;
8756 }
8757 let dot = query
8758 .iter()
8759 .zip(vector)
8760 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8761 .sum::<f64>();
8762 let norm = vector
8763 .iter()
8764 .map(|value| f64::from(*value).powi(2))
8765 .sum::<f64>()
8766 .sqrt();
8767 if query_norm == 0.0 || norm == 0.0 {
8768 0.0
8769 } else {
8770 dot / (query_norm * norm)
8771 }
8772 }
8773 crate::query::VectorMetric::Euclidean => {
8774 if let Some(context) = context {
8775 context.consume(vector_work)?;
8776 }
8777 query
8778 .iter()
8779 .zip(vector)
8780 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
8781 .sum::<f64>()
8782 .sqrt()
8783 }
8784 };
8785 if !score.is_finite() {
8786 return Err(MongrelError::InvalidArgument(
8787 "exact rerank score must be finite".into(),
8788 ));
8789 }
8790 scores.insert(row_id, score as f32);
8791 }
8792 let mut reranked = Vec::with_capacity(ranked.len());
8793 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
8794 let Some(score) = scores.get(&row_id).copied() else {
8795 continue;
8796 };
8797 let ordering_score = match metric {
8798 crate::query::VectorMetric::Euclidean => -f64::from(score),
8799 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
8800 f64::from(score)
8801 }
8802 };
8803 let final_score = fused_score + *weight * ordering_score;
8804 if !final_score.is_finite() {
8805 return Err(MongrelError::InvalidArgument(
8806 "final rerank score must be finite".into(),
8807 ));
8808 }
8809 reranked.push((row_id, fused_score, components, Some(score), final_score));
8810 }
8811 ranked = reranked;
8812 ranked.sort_by(order);
8813 crate::trace::QueryTrace::record(|trace| {
8814 trace.exact_vector_gather_nanos =
8815 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
8816 trace.exact_vector_score_nanos = trace
8817 .exact_vector_score_nanos
8818 .saturating_add(score_started.elapsed().as_nanos() as u64);
8819 });
8820 }
8821 if let Some(after) = after {
8822 ranked.retain(|(row_id, _, _, _, final_score)| {
8823 final_score.total_cmp(&after.final_score).is_lt()
8824 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
8825 });
8826 }
8827 let projection_started = std::time::Instant::now();
8828 let sentinel = projection
8829 .first()
8830 .copied()
8831 .or_else(|| self.schema.columns.first().map(|column| column.id));
8832 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8833 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
8834 let mut projection_rows = 0usize;
8835 let mut projection_cells = 0usize;
8836 while out.len() < request.limit && !ranked.is_empty() {
8837 if let Some(context) = context {
8838 context.checkpoint()?;
8839 context.consume(ranked.len())?;
8840 }
8841 let needed = request.limit - out.len();
8842 let window_size = ranked
8843 .len()
8844 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
8845 let selection_started = std::time::Instant::now();
8846 let mut remainder = if ranked.len() > window_size {
8847 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
8848 ranked.split_off(window_size)
8849 } else {
8850 Vec::new()
8851 };
8852 ranked.sort_by(order);
8853 fusion_nanos =
8854 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
8855 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
8856 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
8857 if let Some(context) = context {
8858 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
8859 }
8860 projection_rows = projection_rows.saturating_add(row_ids.len());
8861 projection_cells =
8862 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
8863 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
8864 std::collections::HashMap::new();
8865 if let Some(column_id) = sentinel {
8866 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8867 &row_ids, column_id, snapshot, query_now, context,
8868 )? {
8869 cells.entry(row_id).or_default().insert(column_id, value);
8870 }
8871 }
8872 for &column_id in &projection {
8873 if Some(column_id) == sentinel {
8874 continue;
8875 }
8876 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8877 &row_ids, column_id, snapshot, query_now, context,
8878 )? {
8879 cells.entry(row_id).or_default().insert(column_id, value);
8880 }
8881 }
8882 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
8883 ranked.drain(..)
8884 {
8885 let Some(row_cells) = cells.remove(&row_id) else {
8886 continue;
8887 };
8888 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
8889 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
8890 out.push(SearchHit {
8891 row_id,
8892 cells: projection
8893 .iter()
8894 .filter_map(|column_id| {
8895 row_cells
8896 .get(column_id)
8897 .cloned()
8898 .map(|value| (*column_id, value))
8899 })
8900 .collect(),
8901 components,
8902 fused_score,
8903 exact_rerank_score,
8904 final_score,
8905 final_rank,
8906 });
8907 if out.len() == request.limit {
8908 break;
8909 }
8910 }
8911 ranked.append(&mut remainder);
8912 }
8913 crate::trace::QueryTrace::record(|trace| {
8914 trace.union_size = union_size;
8915 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
8916 trace.projection_nanos = trace
8917 .projection_nanos
8918 .saturating_add(projection_started.elapsed().as_nanos() as u64);
8919 trace.total_nanos = trace
8920 .total_nanos
8921 .saturating_add(total_started.elapsed().as_nanos() as u64);
8922 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
8923 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
8924 if let Some(context) = context {
8925 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
8926 }
8927 });
8928 Ok(out)
8929 }
8930
8931 pub fn set_similarity(
8934 &mut self,
8935 request: &crate::query::SetSimilarityRequest,
8936 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8937 self.set_similarity_with_allowed(request, None)
8938 }
8939
8940 pub fn set_similarity_at(
8941 &mut self,
8942 request: &crate::query::SetSimilarityRequest,
8943 snapshot: Snapshot,
8944 allowed: Option<&std::collections::HashSet<RowId>>,
8945 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8946 self.set_similarity_explained_at(request, snapshot, allowed)
8947 .map(|(hits, _)| hits)
8948 }
8949
8950 pub fn ann_rerank(
8952 &mut self,
8953 request: &crate::query::AnnRerankRequest,
8954 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8955 self.ann_rerank_with_allowed(request, None)
8956 }
8957
8958 pub fn ann_rerank_with_allowed(
8959 &mut self,
8960 request: &crate::query::AnnRerankRequest,
8961 allowed: Option<&std::collections::HashSet<RowId>>,
8962 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8963 self.ann_rerank_at(request, self.snapshot(), allowed)
8964 }
8965
8966 pub fn ann_rerank_at(
8967 &mut self,
8968 request: &crate::query::AnnRerankRequest,
8969 snapshot: Snapshot,
8970 allowed: Option<&std::collections::HashSet<RowId>>,
8971 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8972 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
8973 }
8974
8975 pub fn ann_rerank_at_with_context(
8976 &mut self,
8977 request: &crate::query::AnnRerankRequest,
8978 snapshot: Snapshot,
8979 allowed: Option<&std::collections::HashSet<RowId>>,
8980 context: Option<&crate::query::AiExecutionContext>,
8981 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8982 self.ensure_indexes_complete()?;
8983 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
8984 }
8985
8986 pub fn ann_rerank_at_with_candidate_authorization_and_context(
8987 &mut self,
8988 request: &crate::query::AnnRerankRequest,
8989 snapshot: Snapshot,
8990 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8991 context: Option<&crate::query::AiExecutionContext>,
8992 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8993 self.ensure_indexes_complete()?;
8994 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
8995 }
8996
8997 #[doc(hidden)]
8998 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
8999 &self,
9000 request: &crate::query::AnnRerankRequest,
9001 snapshot: Snapshot,
9002 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9003 context: Option<&crate::query::AiExecutionContext>,
9004 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9005 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
9006 }
9007
9008 fn ann_rerank_at_with_filters_and_context(
9009 &self,
9010 request: &crate::query::AnnRerankRequest,
9011 snapshot: Snapshot,
9012 allowed: Option<&std::collections::HashSet<RowId>>,
9013 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9014 context: Option<&crate::query::AiExecutionContext>,
9015 ) -> Result<Vec<crate::query::AnnRerankHit>> {
9016 use crate::query::{
9017 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
9018 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
9019 };
9020 if request.candidate_k == 0 || request.limit == 0 {
9021 return Err(MongrelError::InvalidArgument(
9022 "candidate_k and limit must be > 0".into(),
9023 ));
9024 }
9025 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
9026 return Err(MongrelError::InvalidArgument(format!(
9027 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
9028 )));
9029 }
9030 let retriever = Retriever::Ann {
9031 column_id: request.column_id,
9032 query: request.query.clone(),
9033 k: request.candidate_k,
9034 };
9035 self.require_select()?;
9036 self.validate_retriever(&retriever)?;
9037 let hits = self.retrieve_filtered(
9038 &retriever,
9039 snapshot,
9040 None,
9041 allowed,
9042 candidate_authorization,
9043 context,
9044 )?;
9045 let distances: std::collections::HashMap<_, _> = hits
9046 .iter()
9047 .filter_map(|hit| match hit.score {
9048 RetrieverScore::AnnHammingDistance(distance) => {
9049 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
9050 }
9051 RetrieverScore::AnnCosineDistance(distance) => {
9052 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
9053 }
9054 _ => None,
9055 })
9056 .collect();
9057 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
9058 if let Some(context) = context {
9059 context.consume(row_ids.len())?;
9060 }
9061 let gather_started = std::time::Instant::now();
9062 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9063 let values = self.values_for_rids_batch_at_with_context(
9064 &row_ids,
9065 request.column_id,
9066 snapshot,
9067 query_now,
9068 context,
9069 )?;
9070 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
9071 let score_started = std::time::Instant::now();
9072 let vector_work =
9073 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
9074 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
9075 if let Some(context) = context {
9076 context.consume(vector_work)?;
9077 }
9078 request
9079 .query
9080 .iter()
9081 .map(|value| f64::from(*value).powi(2))
9082 .sum::<f64>()
9083 .sqrt()
9084 } else {
9085 0.0
9086 };
9087 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
9088 for (row_id, value) in values {
9089 if let Some(meta) = value.generated_embedding_metadata() {
9090 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
9091 continue;
9092 }
9093 }
9094 let Some(vector) = value.as_embedding() else {
9095 continue;
9096 };
9097 let exact_score = match request.metric {
9098 VectorMetric::DotProduct => {
9099 if let Some(context) = context {
9100 context.consume(vector_work)?;
9101 }
9102 request
9103 .query
9104 .iter()
9105 .zip(vector)
9106 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9107 .sum::<f64>()
9108 }
9109 VectorMetric::Cosine => {
9110 if let Some(context) = context {
9111 context.consume(vector_work.saturating_mul(2))?;
9112 }
9113 let dot = request
9114 .query
9115 .iter()
9116 .zip(vector)
9117 .map(|(left, right)| f64::from(*left) * f64::from(*right))
9118 .sum::<f64>();
9119 let norm = vector
9120 .iter()
9121 .map(|value| f64::from(*value).powi(2))
9122 .sum::<f64>()
9123 .sqrt();
9124 if query_norm == 0.0 || norm == 0.0 {
9125 0.0
9126 } else {
9127 dot / (query_norm * norm)
9128 }
9129 }
9130 VectorMetric::Euclidean => {
9131 if let Some(context) = context {
9132 context.consume(vector_work)?;
9133 }
9134 request
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 let exact_score = exact_score as f32;
9144 if !exact_score.is_finite() {
9145 return Err(MongrelError::InvalidArgument(
9146 "exact ANN score must be finite".into(),
9147 ));
9148 }
9149 let Some(candidate_distance) = distances.get(&row_id).copied() else {
9150 continue;
9151 };
9152 reranked.push(AnnRerankHit {
9153 row_id,
9154 candidate_distance,
9155 exact_score,
9156 });
9157 }
9158 reranked.sort_by(|left, right| {
9159 let score = match request.metric {
9160 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
9161 VectorMetric::Cosine | VectorMetric::DotProduct => {
9162 right.exact_score.total_cmp(&left.exact_score)
9163 }
9164 };
9165 score.then_with(|| left.row_id.cmp(&right.row_id))
9166 });
9167 reranked.truncate(request.limit);
9168 crate::trace::QueryTrace::record(|trace| {
9169 trace.exact_vector_gather_nanos =
9170 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
9171 trace.exact_vector_score_nanos = trace
9172 .exact_vector_score_nanos
9173 .saturating_add(score_started.elapsed().as_nanos() as u64);
9174 });
9175 Ok(reranked)
9176 }
9177
9178 pub fn set_similarity_with_allowed(
9179 &mut self,
9180 request: &crate::query::SetSimilarityRequest,
9181 allowed: Option<&std::collections::HashSet<RowId>>,
9182 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9183 self.set_similarity_explained_at(request, self.snapshot(), allowed)
9184 .map(|(hits, _)| hits)
9185 }
9186
9187 pub fn set_similarity_explained(
9188 &mut self,
9189 request: &crate::query::SetSimilarityRequest,
9190 ) -> Result<(
9191 Vec<crate::query::SetSimilarityHit>,
9192 crate::query::SetSimilarityTrace,
9193 )> {
9194 self.set_similarity_explained_at(request, self.snapshot(), None)
9195 }
9196
9197 fn set_similarity_explained_at(
9198 &mut self,
9199 request: &crate::query::SetSimilarityRequest,
9200 snapshot: Snapshot,
9201 allowed: Option<&std::collections::HashSet<RowId>>,
9202 ) -> Result<(
9203 Vec<crate::query::SetSimilarityHit>,
9204 crate::query::SetSimilarityTrace,
9205 )> {
9206 self.ensure_indexes_complete()?;
9207 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
9208 }
9209
9210 pub fn set_similarity_at_with_context(
9211 &mut self,
9212 request: &crate::query::SetSimilarityRequest,
9213 snapshot: Snapshot,
9214 allowed: Option<&std::collections::HashSet<RowId>>,
9215 context: Option<&crate::query::AiExecutionContext>,
9216 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9217 self.ensure_indexes_complete()?;
9218 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
9219 .map(|(hits, _)| hits)
9220 }
9221
9222 pub fn set_similarity_at_with_candidate_authorization_and_context(
9223 &mut self,
9224 request: &crate::query::SetSimilarityRequest,
9225 snapshot: Snapshot,
9226 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9227 context: Option<&crate::query::AiExecutionContext>,
9228 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9229 self.ensure_indexes_complete()?;
9230 self.set_similarity_explained_at_with_context(
9231 request,
9232 snapshot,
9233 None,
9234 authorization,
9235 context,
9236 )
9237 .map(|(hits, _)| hits)
9238 }
9239
9240 #[doc(hidden)]
9241 pub fn set_similarity_at_with_candidate_authorization_on_generation(
9242 &self,
9243 request: &crate::query::SetSimilarityRequest,
9244 snapshot: Snapshot,
9245 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9246 context: Option<&crate::query::AiExecutionContext>,
9247 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9248 self.set_similarity_explained_at_with_context(
9249 request,
9250 snapshot,
9251 None,
9252 authorization,
9253 context,
9254 )
9255 .map(|(hits, _)| hits)
9256 }
9257
9258 fn set_similarity_explained_at_with_context(
9259 &self,
9260 request: &crate::query::SetSimilarityRequest,
9261 snapshot: Snapshot,
9262 allowed: Option<&std::collections::HashSet<RowId>>,
9263 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9264 context: Option<&crate::query::AiExecutionContext>,
9265 ) -> Result<(
9266 Vec<crate::query::SetSimilarityHit>,
9267 crate::query::SetSimilarityTrace,
9268 )> {
9269 use crate::query::{
9270 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
9271 MAX_SET_MEMBERS,
9272 };
9273 let mut trace = crate::query::SetSimilarityTrace::default();
9274 if request.members.is_empty() {
9275 return Ok((Vec::new(), trace));
9276 }
9277 if request.candidate_k == 0 || request.limit == 0 {
9278 return Err(MongrelError::InvalidArgument(
9279 "candidate_k and limit must be > 0".into(),
9280 ));
9281 }
9282 if request.candidate_k > MAX_RETRIEVER_K
9283 || request.limit > MAX_FINAL_LIMIT
9284 || request.members.len() > MAX_SET_MEMBERS
9285 {
9286 return Err(MongrelError::InvalidArgument(format!(
9287 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
9288 )));
9289 }
9290 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
9291 return Err(MongrelError::InvalidArgument(
9292 "min_jaccard must be finite and between 0 and 1".into(),
9293 ));
9294 }
9295 let started = std::time::Instant::now();
9296 let retriever = Retriever::MinHash {
9297 column_id: request.column_id,
9298 members: request.members.clone(),
9299 k: request.candidate_k,
9300 };
9301 self.require_select()?;
9302 self.validate_retriever(&retriever)?;
9303 let hits = self.retrieve_filtered(
9304 &retriever,
9305 snapshot,
9306 None,
9307 allowed,
9308 candidate_authorization,
9309 context,
9310 )?;
9311 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
9312 trace.candidate_count = hits.len();
9313 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
9314 if let Some(context) = context {
9315 context.consume(row_ids.len())?;
9316 }
9317 let started = std::time::Instant::now();
9318 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9319 let values = self.values_for_rids_batch_at_with_context(
9320 &row_ids,
9321 request.column_id,
9322 snapshot,
9323 query_now,
9324 context,
9325 )?;
9326 trace.gather_us = started.elapsed().as_micros() as u64;
9327 if let Some(context) = context {
9328 context.consume(request.members.len())?;
9329 }
9330 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
9331 let estimates: std::collections::HashMap<_, _> = hits
9332 .into_iter()
9333 .filter_map(|hit| match hit.score {
9334 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
9335 _ => None,
9336 })
9337 .collect();
9338 let started = std::time::Instant::now();
9339 let mut parsed = Vec::with_capacity(values.len());
9340 for (row_id, value) in values {
9341 let Value::Bytes(bytes) = value else {
9342 continue;
9343 };
9344 if let Some(context) = context {
9345 context.consume(crate::query::work_units(
9346 bytes.len(),
9347 crate::query::PARSE_WORK_QUANTUM,
9348 ))?;
9349 }
9350 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
9351 continue;
9352 };
9353 if let Some(context) = context {
9354 context.consume(members.len())?;
9355 }
9356 let stored = members
9357 .into_iter()
9358 .filter_map(|member| match member {
9359 serde_json::Value::String(value) => {
9360 Some(crate::query::SetMember::String(value))
9361 }
9362 serde_json::Value::Number(value) => {
9363 Some(crate::query::SetMember::Number(value))
9364 }
9365 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
9366 _ => None,
9367 })
9368 .collect::<std::collections::HashSet<_>>();
9369 parsed.push((row_id, stored));
9370 }
9371 trace.parse_us = started.elapsed().as_micros() as u64;
9372 trace.verified_count = parsed.len();
9373 let started = std::time::Instant::now();
9374 let mut exact = Vec::new();
9375 for (row_id, stored) in parsed {
9376 if let Some(context) = context {
9377 context.consume(query.len().saturating_add(stored.len()))?;
9378 }
9379 let union = query.union(&stored).count();
9380 let score = if union == 0 {
9381 1.0
9382 } else {
9383 query.intersection(&stored).count() as f32 / union as f32
9384 };
9385 if score >= request.min_jaccard {
9386 exact.push(SetSimilarityHit {
9387 row_id,
9388 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
9389 exact_jaccard: score,
9390 });
9391 }
9392 }
9393 exact.sort_by(|a, b| {
9394 b.exact_jaccard
9395 .total_cmp(&a.exact_jaccard)
9396 .then_with(|| a.row_id.cmp(&b.row_id))
9397 });
9398 exact.truncate(request.limit);
9399 trace.score_us = started.elapsed().as_micros() as u64;
9400 crate::trace::QueryTrace::record(|query_trace| {
9401 query_trace.exact_set_gather_nanos = query_trace
9402 .exact_set_gather_nanos
9403 .saturating_add(trace.gather_us.saturating_mul(1_000));
9404 query_trace.exact_set_parse_nanos = query_trace
9405 .exact_set_parse_nanos
9406 .saturating_add(trace.parse_us.saturating_mul(1_000));
9407 query_trace.exact_set_score_nanos = query_trace
9408 .exact_set_score_nanos
9409 .saturating_add(trace.score_us.saturating_mul(1_000));
9410 });
9411 Ok((exact, trace))
9412 }
9413
9414 fn values_for_rids_batch_at(
9416 &self,
9417 row_ids: &[u64],
9418 column_id: u16,
9419 snapshot: Snapshot,
9420 now: i64,
9421 ) -> Result<Vec<(RowId, Value)>> {
9422 if self.ttl.is_none()
9423 && self.memtable.is_empty()
9424 && self.mutable_run.is_empty()
9425 && self.run_refs.len() == 1
9426 {
9427 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9428 if row_ids.len().saturating_mul(24) < reader.row_count() {
9433 let mut values = Vec::with_capacity(row_ids.len());
9434 for &raw_row_id in row_ids {
9435 let row_id = RowId(raw_row_id);
9436 if let Some((_, false, Some(value))) =
9437 reader.get_version_column_at(row_id, snapshot, column_id)?
9438 {
9439 values.push((row_id, value));
9440 }
9441 }
9442 return Ok(values);
9443 }
9444 let (positions, visible_row_ids) = reader.visible_positions_with_rids_at(snapshot)?;
9445 let requested: Vec<(RowId, usize)> = row_ids
9446 .iter()
9447 .filter_map(|raw| {
9448 visible_row_ids
9449 .binary_search(&(*raw as i64))
9450 .ok()
9451 .map(|index| (RowId(*raw), positions[index]))
9452 })
9453 .collect();
9454 let values = reader.gather_column(
9455 column_id,
9456 &requested
9457 .iter()
9458 .map(|(_, position)| *position)
9459 .collect::<Vec<_>>(),
9460 )?;
9461 return Ok(requested
9462 .into_iter()
9463 .zip(values)
9464 .map(|((row_id, _), value)| (row_id, value))
9465 .collect());
9466 }
9467 self.values_for_rids_at(row_ids, column_id, snapshot, now)
9468 }
9469
9470 fn values_for_rids_batch_at_with_context(
9471 &self,
9472 row_ids: &[u64],
9473 column_id: u16,
9474 snapshot: Snapshot,
9475 now: i64,
9476 context: Option<&crate::query::AiExecutionContext>,
9477 ) -> Result<Vec<(RowId, Value)>> {
9478 let Some(context) = context else {
9479 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
9480 };
9481 let mut values = Vec::with_capacity(row_ids.len());
9482 for chunk in row_ids.chunks(256) {
9483 context.checkpoint()?;
9484 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
9485 }
9486 Ok(values)
9487 }
9488
9489 fn values_for_rids_at(
9491 &self,
9492 row_ids: &[u64],
9493 column_id: u16,
9494 snapshot: Snapshot,
9495 now: i64,
9496 ) -> Result<Vec<(RowId, Value)>> {
9497 let mut readers: Vec<_> = self
9498 .run_refs
9499 .iter()
9500 .map(|run| self.open_reader(run.run_id))
9501 .collect::<Result<_>>()?;
9502 let mut out = Vec::with_capacity(row_ids.len());
9503 for &raw_row_id in row_ids {
9504 let row_id = RowId(raw_row_id);
9505 let mem = self.memtable.get_version_at(row_id, snapshot);
9506 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
9507 let overlay = match (mem, mutable) {
9508 (Some((_, a)), Some((_, b))) => Some(
9509 if Snapshot::version_is_newer(
9510 a.committed_epoch,
9511 a.commit_ts,
9512 b.committed_epoch,
9513 b.commit_ts,
9514 ) {
9515 a
9516 } else {
9517 b
9518 },
9519 ),
9520 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
9521 (None, None) => None,
9522 };
9523 if let Some(row) = overlay {
9524 if !row.deleted && !self.row_expired_at(&row, now) {
9525 if let Some(value) = row.columns.get(&column_id) {
9526 out.push((row_id, value.clone()));
9527 }
9528 }
9529 continue;
9530 }
9531
9532 let mut best: Option<(crate::sorted_run::VersionedColumnValue, usize)> = None;
9537 for (index, reader) in readers.iter_mut().enumerate() {
9538 if let Some(versioned) =
9539 reader.get_version_column_full_at(row_id, snapshot, column_id)?
9540 {
9541 if best
9542 .as_ref()
9543 .map(|(current, _)| versioned.stamp.is_newer_than(current.stamp))
9544 .unwrap_or(true)
9545 {
9546 best = Some((versioned, index));
9547 }
9548 }
9549 }
9550 let Some((candidate, reader_index)) = best else {
9551 continue;
9552 };
9553 if candidate.deleted {
9554 continue;
9555 }
9556 let Some(value) = candidate.value.clone() else {
9557 continue;
9558 };
9559 if let Some(ttl) = self.ttl {
9560 if ttl.column_id != column_id {
9561 if let Some(ttl_value) = readers[reader_index].get_version_column_full_at(
9562 row_id,
9563 snapshot,
9564 ttl.column_id,
9565 )? {
9566 if let Some(Value::Int64(timestamp)) = ttl_value.value {
9567 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9568 continue;
9569 }
9570 }
9571 }
9572 } else if let Value::Int64(timestamp) = &value {
9573 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9574 continue;
9575 }
9576 }
9577 }
9578 out.push((row_id, value));
9579 }
9580 Ok(out)
9581 }
9582
9583 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
9588 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
9589 }
9590
9591 pub fn rows_for_rids_with_context(
9592 &self,
9593 rids: &[u64],
9594 snapshot: Snapshot,
9595 context: &crate::query::AiExecutionContext,
9596 ) -> Result<Vec<Row>> {
9597 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
9598 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
9599 }
9600
9601 fn rows_for_rids_at_time(
9602 &self,
9603 rids: &[u64],
9604 snapshot: Snapshot,
9605 ttl_now: i64,
9606 control: Option<&crate::ExecutionControl>,
9607 ) -> Result<Vec<Row>> {
9608 use std::collections::HashMap;
9609 let mut rows = Vec::with_capacity(rids.len());
9610 let tier_size = self.memtable.len() + self.mutable_run.len();
9625 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
9626 if rids.len().saturating_mul(24) < tier_size {
9627 for &rid in rids {
9628 if overlay.len() & 255 == 0 {
9629 control
9630 .map(crate::ExecutionControl::checkpoint)
9631 .transpose()?;
9632 }
9633 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
9634 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
9635 let newest = match (mem, mrun) {
9636 (Some((_, mr)), Some((_, rr))) => Some(
9637 if Snapshot::version_is_newer(
9638 mr.committed_epoch,
9639 mr.commit_ts,
9640 rr.committed_epoch,
9641 rr.commit_ts,
9642 ) {
9643 mr
9644 } else {
9645 rr
9646 },
9647 ),
9648 (Some((_, mr)), None) => Some(mr),
9649 (None, Some((_, rr))) => Some(rr),
9650 (None, None) => None,
9651 };
9652 if let Some(row) = newest {
9653 overlay.insert(rid, row);
9654 }
9655 }
9656 } else {
9657 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
9658 overlay
9659 .entry(row.row_id.0)
9660 .and_modify(|e| {
9661 if Snapshot::version_is_newer(
9662 row.committed_epoch,
9663 row.commit_ts,
9664 e.committed_epoch,
9665 e.commit_ts,
9666 ) {
9667 *e = row.clone();
9668 }
9669 })
9670 .or_insert(row);
9671 };
9672 for (index, row) in self
9673 .memtable
9674 .visible_versions_at(snapshot)
9675 .into_iter()
9676 .enumerate()
9677 {
9678 if index & 255 == 0 {
9679 control
9680 .map(crate::ExecutionControl::checkpoint)
9681 .transpose()?;
9682 }
9683 fold_newest(row, &mut overlay);
9684 }
9685 for (index, row) in self
9686 .mutable_run
9687 .visible_versions_at(snapshot)
9688 .into_iter()
9689 .enumerate()
9690 {
9691 if index & 255 == 0 {
9692 control
9693 .map(crate::ExecutionControl::checkpoint)
9694 .transpose()?;
9695 }
9696 fold_newest(row, &mut overlay);
9697 }
9698 }
9699 if self.run_refs.len() == 1 {
9700 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9701 if rids.len().saturating_mul(24) < reader.row_count() {
9709 for (index, &rid) in rids.iter().enumerate() {
9710 if index & 255 == 0 {
9711 control
9712 .map(crate::ExecutionControl::checkpoint)
9713 .transpose()?;
9714 }
9715 if let Some(r) = overlay.get(&rid) {
9716 if !r.deleted {
9717 rows.push(r.clone());
9718 }
9719 continue;
9720 }
9721 if let Some((_, row)) = reader.get_version_at(RowId(rid), snapshot)? {
9722 if !row.deleted {
9723 rows.push(row);
9724 }
9725 }
9726 }
9727 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9728 return Ok(rows);
9729 }
9730 let (positions, vis_rids) = reader.visible_positions_with_rids_at(snapshot)?;
9739 enum Src {
9742 Overlay,
9743 Run,
9744 }
9745 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
9746 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
9747 for (index, rid) in rids.iter().enumerate() {
9748 if index & 255 == 0 {
9749 control
9750 .map(crate::ExecutionControl::checkpoint)
9751 .transpose()?;
9752 }
9753 if overlay.contains_key(rid) {
9754 plan.push(Src::Overlay);
9755 continue;
9756 }
9757 match vis_rids.binary_search(&(*rid as i64)) {
9758 Ok(i) => {
9759 plan.push(Src::Run);
9760 fetch.push(positions[i]);
9761 }
9762 Err(_) => { }
9763 }
9764 }
9765 let fetched = reader.materialize_batch(&fetch)?;
9766 let mut fetched_iter = fetched.into_iter();
9767 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
9768 if index & 255 == 0 {
9769 control
9770 .map(crate::ExecutionControl::checkpoint)
9771 .transpose()?;
9772 }
9773 match src {
9774 Src::Overlay => {
9775 if let Some(r) = overlay.get(rid) {
9776 if !r.deleted {
9777 rows.push(r.clone());
9778 }
9779 }
9780 }
9781 Src::Run => {
9782 if let Some(row) = fetched_iter.next() {
9783 if !row.deleted {
9784 rows.push(row);
9785 }
9786 }
9787 }
9788 }
9789 }
9790 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9791 return Ok(rows);
9792 }
9793 let mut readers: Vec<_> = self
9797 .run_refs
9798 .iter()
9799 .map(|rr| self.open_reader(rr.run_id))
9800 .collect::<Result<Vec<_>>>()?;
9801 for (index, rid) in rids.iter().enumerate() {
9802 if index & 255 == 0 {
9803 control
9804 .map(crate::ExecutionControl::checkpoint)
9805 .transpose()?;
9806 }
9807 if let Some(r) = overlay.get(rid) {
9808 if !r.deleted {
9809 rows.push(r.clone());
9810 }
9811 continue;
9812 }
9813 let mut best: Option<(VersionStamp, Row)> = None;
9818 for reader in readers.iter_mut() {
9819 if let Ok(Some((stamp, row))) = reader.get_version_stamp_at(RowId(*rid), snapshot) {
9820 if best
9821 .as_ref()
9822 .map(|(current, _)| stamp.is_newer_than(*current))
9823 .unwrap_or(true)
9824 {
9825 best = Some((stamp, row));
9826 }
9827 }
9828 }
9829 if let Some((_, r)) = best {
9830 if !r.deleted {
9831 rows.push(r);
9832 }
9833 }
9834 }
9835 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9836 Ok(rows)
9837 }
9838
9839 pub fn indexes_complete(&self) -> bool {
9849 self.indexes_complete
9850 }
9851
9852 pub fn index_build_policy(&self) -> IndexBuildPolicy {
9854 self.index_build_policy
9855 }
9856
9857 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
9861 self.index_build_policy = policy;
9862 }
9863
9864 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
9869 if !self.indexes_complete {
9873 return None;
9874 }
9875 let b = self.bitmap.get(&column_id)?;
9876 let result: Vec<Vec<u8>> = b
9877 .keys()
9878 .into_iter()
9879 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
9880 .collect();
9881 Some(result)
9882 }
9883
9884 pub fn fk_join_row_ids(
9885 &self,
9886 fk_column_id: u16,
9887 pk_values: &[Vec<u8>],
9888 fk_conditions: &[crate::query::Condition],
9889 snapshot: Snapshot,
9890 ) -> Result<Vec<u64>> {
9891 let Some(b) = self.bitmap.get(&fk_column_id) else {
9892 return Ok(Vec::new());
9893 };
9894 let mut join_set = {
9895 let mut acc = roaring::RoaringBitmap::new();
9896 for v in pk_values {
9897 acc |= b.get(v);
9898 }
9899 RowIdSet::from_roaring(acc)
9900 };
9901 if !fk_conditions.is_empty() {
9902 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
9903 sets.push(join_set);
9904 for c in fk_conditions {
9905 sets.push(self.resolve_condition(c, snapshot)?);
9906 }
9907 join_set = RowIdSet::intersect_many(sets);
9908 }
9909 Ok(join_set.into_sorted_vec())
9910 }
9911
9912 pub fn fk_join_count(
9918 &self,
9919 fk_column_id: u16,
9920 pk_values: &[Vec<u8>],
9921 fk_conditions: &[crate::query::Condition],
9922 snapshot: Snapshot,
9923 ) -> Result<u64> {
9924 let Some(b) = self.bitmap.get(&fk_column_id) else {
9925 return Ok(0);
9926 };
9927 let mut acc = roaring::RoaringBitmap::new();
9928 for v in pk_values {
9929 acc |= b.get(v);
9930 }
9931 if fk_conditions.is_empty() {
9932 return Ok(acc.len());
9933 }
9934 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
9935 sets.push(RowIdSet::from_roaring(acc));
9936 for c in fk_conditions {
9937 sets.push(self.resolve_condition(c, snapshot)?);
9938 }
9939 Ok(RowIdSet::intersect_many(sets).len() as u64)
9940 }
9941
9942 fn inspect_row_at(&self, rid: RowId, snapshot: Snapshot) -> Option<crate::memtable::Row> {
9947 let mut best: Option<crate::memtable::Row> = None;
9948 let mut consider = |row: crate::memtable::Row| {
9949 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
9950 return;
9951 }
9952 if best.as_ref().is_none_or(|current| {
9953 Snapshot::version_is_newer(
9954 row.committed_epoch,
9955 row.commit_ts,
9956 current.committed_epoch,
9957 current.commit_ts,
9958 )
9959 }) {
9960 best = Some(row);
9961 }
9962 };
9963 if let Some((_, row)) = self.memtable.get_version_at(rid, snapshot) {
9964 consider(row);
9965 }
9966 if let Some((_, row)) = self.mutable_run.get_version_at(rid, snapshot) {
9967 consider(row);
9968 }
9969 for rr in &self.run_refs {
9970 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
9971 if rid.0 < min_rid || rid.0 > max_rid {
9972 continue;
9973 }
9974 }
9975 let Ok(mut reader) = self.open_reader(rr.run_id) else {
9976 continue;
9977 };
9978 let Ok(Some((_stamp, row))) = reader.get_version_stamp_at(rid, snapshot) else {
9982 continue;
9983 };
9984 consider(row);
9985 }
9986 best
9987 }
9988
9989 fn resolve_pk_with_hot_fallback(
9990 &self,
9991 pk_column_id: u16,
9992 lookup: &[u8],
9993 snapshot: Snapshot,
9994 ) -> Result<RowIdSet> {
9995 if !self.indexes_complete {
9999 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10000 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::IndexIncomplete);
10001 return Ok(result);
10002 }
10003 let mapped = self.hot.get(lookup);
10006 let Some(r) = mapped else {
10007 let (result, reason) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10008 self.record_hot_fallback_reason(reason);
10009 return Ok(result);
10010 };
10011 if snapshot.epoch < self.current_epoch() {
10016 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::HistoricalSnapshot);
10017 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10018 return Ok(result);
10019 }
10020 let materialized = self.inspect_row_at(r, snapshot);
10024 let now_nanos = unix_nanos_now();
10025 let inspection = crate::trace::inspect_hot_candidate(
10026 materialized.as_ref(),
10027 snapshot,
10028 self.ttl,
10029 now_nanos,
10030 pk_column_id,
10031 lookup,
10032 |row| {
10033 let pk_value = row.columns.get(&pk_column_id);
10034 pk_value
10035 .map(|v| self.index_lookup_key(pk_column_id, v))
10036 .unwrap_or_default()
10037 },
10038 );
10039 match inspection {
10040 crate::trace::HotCandidateInspection::Hit(row) => {
10041 self.lookup_metrics
10042 .hot_lookup_hit
10043 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
10044 crate::trace::QueryTrace::record(|t| {
10045 t.hot_lookup_attempted = true;
10046 t.hot_lookup_hit = true;
10047 });
10048 Ok(RowIdSet::one(row.row_id.0))
10049 }
10050 failure => {
10051 let reason = match &failure {
10052 crate::trace::HotCandidateInspection::Hit(_) => unreachable!(),
10053 crate::trace::HotCandidateInspection::MissingRow => {
10054 crate::trace::HotFallbackReason::Tombstone
10059 }
10060 crate::trace::HotCandidateInspection::Invisible => {
10061 crate::trace::HotFallbackReason::InvisibleAtSnapshot
10062 }
10063 crate::trace::HotCandidateInspection::Tombstone => {
10064 crate::trace::HotFallbackReason::Tombstone
10065 }
10066 crate::trace::HotCandidateInspection::TtlExpired => {
10067 crate::trace::HotFallbackReason::TtlExpired
10068 }
10069 crate::trace::HotCandidateInspection::PrimaryKeyMismatch { .. } => {
10070 crate::trace::HotFallbackReason::PrimaryKeyMismatch
10071 }
10072 };
10073 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
10079 self.record_hot_fallback_reason(reason);
10080 Ok(result)
10081 }
10082 }
10083 }
10084
10085 fn resolve_condition(
10090 &self,
10091 c: &crate::query::Condition,
10092 snapshot: Snapshot,
10093 ) -> Result<RowIdSet> {
10094 self.resolve_condition_with_allowed(c, snapshot, None)
10095 }
10096
10097 fn resolve_condition_with_allowed(
10098 &self,
10099 c: &crate::query::Condition,
10100 snapshot: Snapshot,
10101 allowed: Option<&std::collections::HashSet<RowId>>,
10102 ) -> Result<RowIdSet> {
10103 use crate::query::Condition;
10104 self.validate_condition(c)?;
10105 Ok(match c {
10106 Condition::Pk(key) => {
10107 let lookup = self
10108 .schema
10109 .primary_key()
10110 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
10111 .unwrap_or_else(|| key.clone());
10112 if let Some(pk_col) = self.schema.primary_key() {
10113 return self.resolve_pk_with_hot_fallback(pk_col.id, &lookup, snapshot);
10114 }
10115 RowIdSet::empty()
10119 }
10120 Condition::BitmapEq { column_id, value } => {
10121 let lookup = self.index_lookup_key_bytes(*column_id, value);
10122 self.bitmap
10123 .get(column_id)
10124 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
10125 .unwrap_or_else(RowIdSet::empty)
10126 }
10127 Condition::BitmapIn { column_id, values } => {
10128 let bm = self.bitmap.get(column_id);
10129 let mut acc = roaring::RoaringBitmap::new();
10130 if let Some(b) = bm {
10131 for v in values {
10132 let lookup = self.index_lookup_key_bytes(*column_id, v);
10133 acc |= b.get(&lookup);
10134 }
10135 }
10136 RowIdSet::from_roaring(acc)
10137 }
10138 Condition::BytesPrefix { column_id, prefix } => {
10139 if let Some(b) = self.bitmap.get(column_id) {
10144 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
10145 let mut acc = roaring::RoaringBitmap::new();
10146 for key in b.keys() {
10147 if key.starts_with(&lookup_prefix) {
10148 acc |= b.get(&key);
10149 }
10150 }
10151 RowIdSet::from_roaring(acc)
10152 } else {
10153 RowIdSet::empty()
10154 }
10155 }
10156 Condition::FmContains { column_id, pattern } => self
10157 .fm
10158 .get(column_id)
10159 .map(|f| {
10160 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
10161 })
10162 .unwrap_or_else(RowIdSet::empty),
10163 Condition::FmContainsAll {
10164 column_id,
10165 patterns,
10166 } => {
10167 if let Some(f) = self.fm.get(column_id) {
10170 let sets: Vec<RowIdSet> = patterns
10171 .iter()
10172 .map(|pat| {
10173 RowIdSet::from_unsorted(
10174 f.locate(pat).into_iter().map(|r| r.0).collect(),
10175 )
10176 })
10177 .collect();
10178 RowIdSet::intersect_many(sets)
10179 } else {
10180 RowIdSet::empty()
10181 }
10182 }
10183 Condition::Ann {
10184 column_id,
10185 query,
10186 k,
10187 } => RowIdSet::from_unsorted(
10188 self.retrieve_filtered(
10189 &crate::query::Retriever::Ann {
10190 column_id: *column_id,
10191 query: query.clone(),
10192 k: *k,
10193 },
10194 snapshot,
10195 None,
10196 allowed,
10197 None,
10198 None,
10199 )?
10200 .into_iter()
10201 .map(|hit| hit.row_id.0)
10202 .collect(),
10203 ),
10204 Condition::SparseMatch {
10205 column_id,
10206 query,
10207 k,
10208 } => RowIdSet::from_unsorted(
10209 self.retrieve_filtered(
10210 &crate::query::Retriever::Sparse {
10211 column_id: *column_id,
10212 query: query.clone(),
10213 k: *k,
10214 },
10215 snapshot,
10216 None,
10217 allowed,
10218 None,
10219 None,
10220 )?
10221 .into_iter()
10222 .map(|hit| hit.row_id.0)
10223 .collect(),
10224 ),
10225 Condition::MinHashSimilar {
10226 column_id,
10227 query,
10228 k,
10229 } => match self.minhash.get(column_id) {
10230 Some(index) => {
10231 let candidates = index.candidate_row_ids(query);
10232 let eligible =
10233 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
10234 RowIdSet::from_unsorted(
10235 index
10236 .search_filtered(query, *k, |row_id| {
10237 eligible.contains(&row_id)
10238 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
10239 })
10240 .into_iter()
10241 .map(|(row_id, _)| row_id.0)
10242 .collect(),
10243 )
10244 }
10245 None => RowIdSet::empty(),
10246 },
10247 Condition::Range { column_id, lo, hi } => {
10248 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10267 if self.run_refs.len() == 1 {
10268 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10271 } else {
10272 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10280 if lo == hi {
10281 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10282 }
10283 return Ok(multi);
10284 }
10285 } else if self.run_refs.len() == 1 {
10286 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10287 r.range_row_id_set_i64(*column_id, *lo, *hi)?
10288 } else {
10289 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10292 if lo == hi {
10293 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10294 }
10295 return Ok(multi);
10296 };
10297 if self.learned_range.get(column_id).is_none() && self.run_refs.len() == 1 {
10298 let candidates: Vec<RowId> =
10299 set.into_sorted_vec().into_iter().map(RowId).collect();
10300 set = RowIdSet::from_unsorted(
10301 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?
10302 .into_iter()
10303 .map(|row_id| row_id.0)
10304 .collect(),
10305 );
10306 }
10307 set.remove_many(self.overlay_rid_set(snapshot));
10308 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
10309 if lo == hi {
10310 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
10311 }
10312 RowIdSet::from_unsorted(
10313 set.into_sorted_vec()
10314 .into_iter()
10315 .filter(|row_id| self.get(RowId(*row_id), snapshot).is_some())
10316 .collect(),
10317 )
10318 }
10319 Condition::RangeF64 {
10320 column_id,
10321 lo,
10322 lo_inclusive,
10323 hi,
10324 hi_inclusive,
10325 } => {
10326 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10329 RowIdSet::from_unsorted(
10330 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10331 .into_iter()
10332 .collect(),
10333 )
10334 } else if self.run_refs.len() == 1 {
10335 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10336 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
10337 } else {
10338 return self.range_scan_f64(
10339 *column_id,
10340 *lo,
10341 *lo_inclusive,
10342 *hi,
10343 *hi_inclusive,
10344 snapshot,
10345 );
10346 };
10347 set.remove_many(self.overlay_rid_set(snapshot));
10348 self.range_scan_overlay_f64(
10349 &mut set,
10350 *column_id,
10351 *lo,
10352 *lo_inclusive,
10353 *hi,
10354 *hi_inclusive,
10355 snapshot,
10356 );
10357 set
10358 }
10359 Condition::IsNull { column_id } => {
10360 let mut set = if self.run_refs.len() == 1 {
10361 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10362 r.null_row_id_set(*column_id, true)?
10363 } else {
10364 return self.null_scan(*column_id, true, snapshot);
10365 };
10366 set.remove_many(self.overlay_rid_set(snapshot));
10367 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
10368 set
10369 }
10370 Condition::IsNotNull { column_id } => {
10371 let mut set = if self.run_refs.len() == 1 {
10372 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10373 r.null_row_id_set(*column_id, false)?
10374 } else {
10375 return self.null_scan(*column_id, false, snapshot);
10376 };
10377 set.remove_many(self.overlay_rid_set(snapshot));
10378 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
10379 set
10380 }
10381 })
10382 }
10383
10384 fn range_scan_i64(
10392 &self,
10393 column_id: u16,
10394 lo: i64,
10395 hi: i64,
10396 snapshot: Snapshot,
10397 ) -> Result<RowIdSet> {
10398 let mut row_ids = Vec::new();
10399 let mut tomb_rids: HashSet<u64> = HashSet::new();
10406 let overlay_rids = self.overlay_rid_set(snapshot);
10407 for rr in &self.run_refs {
10408 let mut reader = self.open_reader(rr.run_id)?;
10409 let matched = reader.range_row_ids_visible_i64_at(column_id, lo, hi, snapshot)?;
10410 for rid in matched {
10411 if !overlay_rids.contains(&rid) {
10412 row_ids.push(rid);
10413 }
10414 }
10415 for rid in reader.tombstoned_row_ids_at(snapshot)? {
10416 tomb_rids.insert(rid);
10417 }
10418 }
10419 for row in self
10420 .memtable
10421 .visible_versions_at(Snapshot::at(self.pending_epoch()))
10422 .into_iter()
10423 .chain(
10424 self.mutable_run
10425 .visible_versions_at(Snapshot::at(self.pending_epoch())),
10426 )
10427 {
10428 if row.deleted {
10429 tomb_rids.insert(row.row_id.0);
10430 }
10431 }
10432 let mut s = RowIdSet::from_unsorted(row_ids);
10433 s.remove_many(tomb_rids);
10434 let candidates: Vec<RowId> = s.into_sorted_vec().into_iter().map(RowId).collect();
10435 let eligible = self.eligible_candidate_ids(&candidates, column_id, snapshot, None)?;
10436 let mut s = RowIdSet::from_unsorted(eligible.into_iter().map(|row_id| row_id.0).collect());
10437 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
10438 Ok(s)
10439 }
10440
10441 fn range_scan_f64(
10444 &self,
10445 column_id: u16,
10446 lo: f64,
10447 lo_inclusive: bool,
10448 hi: f64,
10449 hi_inclusive: bool,
10450 snapshot: Snapshot,
10451 ) -> Result<RowIdSet> {
10452 let mut row_ids = Vec::new();
10453 let overlay_rids = self.overlay_rid_set(snapshot);
10454 for rr in &self.run_refs {
10455 let mut reader = self.open_reader(rr.run_id)?;
10456 let matched = reader.range_row_ids_visible_f64_at(
10457 column_id,
10458 lo,
10459 lo_inclusive,
10460 hi,
10461 hi_inclusive,
10462 snapshot,
10463 )?;
10464 for rid in matched {
10465 if !overlay_rids.contains(&rid) {
10466 row_ids.push(rid);
10467 }
10468 }
10469 }
10470 let mut s = RowIdSet::from_unsorted(row_ids);
10471 self.range_scan_overlay_f64(
10472 &mut s,
10473 column_id,
10474 lo,
10475 lo_inclusive,
10476 hi,
10477 hi_inclusive,
10478 snapshot,
10479 );
10480 Ok(s)
10481 }
10482
10483 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
10485 let mut s = HashSet::new();
10486 for row in self.memtable.visible_versions_at(snapshot) {
10487 s.insert(row.row_id.0);
10488 }
10489 for row in self.mutable_run.visible_versions_at(snapshot) {
10490 s.insert(row.row_id.0);
10491 }
10492 s
10493 }
10494
10495 fn range_scan_overlay_i64(
10496 &self,
10497 s: &mut RowIdSet,
10498 column_id: u16,
10499 lo: i64,
10500 hi: i64,
10501 snapshot: Snapshot,
10502 ) {
10503 let mut newest: HashMap<u64, Row> = HashMap::new();
10510 for r in self.mutable_run.visible_versions_at(snapshot) {
10511 newest.insert(r.row_id.0, r);
10512 }
10513 for r in self.memtable.visible_versions_at(snapshot) {
10514 newest
10515 .entry(r.row_id.0)
10516 .and_modify(|cur| {
10517 if Snapshot::version_is_newer(
10518 r.committed_epoch,
10519 r.commit_ts,
10520 cur.committed_epoch,
10521 cur.commit_ts,
10522 ) {
10523 *cur = r.clone();
10524 }
10525 })
10526 .or_insert(r);
10527 }
10528 for row in newest.values() {
10529 if !row.deleted {
10530 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
10531 if *v >= lo && *v <= hi && !self.row_expired_at(row, unix_nanos_now()) {
10532 s.insert(row.row_id.0);
10533 }
10534 }
10535 }
10536 }
10537 }
10538
10539 #[allow(clippy::too_many_arguments)]
10540 fn range_scan_overlay_f64(
10541 &self,
10542 s: &mut RowIdSet,
10543 column_id: u16,
10544 lo: f64,
10545 lo_inclusive: bool,
10546 hi: f64,
10547 hi_inclusive: bool,
10548 snapshot: Snapshot,
10549 ) {
10550 let mut newest: HashMap<u64, Row> = HashMap::new();
10553 for r in self.mutable_run.visible_versions_at(snapshot) {
10554 newest.insert(r.row_id.0, r);
10555 }
10556 for r in self.memtable.visible_versions_at(snapshot) {
10557 newest
10558 .entry(r.row_id.0)
10559 .and_modify(|cur| {
10560 if Snapshot::version_is_newer(
10561 r.committed_epoch,
10562 r.commit_ts,
10563 cur.committed_epoch,
10564 cur.commit_ts,
10565 ) {
10566 *cur = r.clone();
10567 }
10568 })
10569 .or_insert(r);
10570 }
10571 for row in newest.values() {
10572 if !row.deleted {
10573 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
10574 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
10575 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
10576 if ok_lo && ok_hi && !self.row_expired_at(row, unix_nanos_now()) {
10577 s.insert(row.row_id.0);
10578 }
10579 }
10580 }
10581 }
10582 }
10583
10584 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
10587 let mut row_ids = Vec::new();
10588 let overlay_rids = self.overlay_rid_set(snapshot);
10589 for rr in &self.run_refs {
10590 let mut reader = self.open_reader(rr.run_id)?;
10591 let matched = reader.null_row_ids_visible_at(column_id, want_nulls, snapshot)?;
10592 for rid in matched {
10593 if !overlay_rids.contains(&rid) {
10594 row_ids.push(rid);
10595 }
10596 }
10597 }
10598 let mut s = RowIdSet::from_unsorted(row_ids);
10599 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
10600 Ok(s)
10601 }
10602
10603 fn null_scan_overlay(
10607 &self,
10608 s: &mut RowIdSet,
10609 column_id: u16,
10610 want_nulls: bool,
10611 snapshot: Snapshot,
10612 ) {
10613 let mut newest: HashMap<u64, Row> = HashMap::new();
10614 for r in self.mutable_run.visible_versions_at(snapshot) {
10615 newest.insert(r.row_id.0, r);
10616 }
10617 for r in self.memtable.visible_versions_at(snapshot) {
10618 newest
10619 .entry(r.row_id.0)
10620 .and_modify(|cur| {
10621 if Snapshot::version_is_newer(
10622 r.committed_epoch,
10623 r.commit_ts,
10624 cur.committed_epoch,
10625 cur.commit_ts,
10626 ) {
10627 *cur = r.clone();
10628 }
10629 })
10630 .or_insert(r);
10631 }
10632 for row in newest.values() {
10633 if row.deleted {
10634 continue;
10635 }
10636 let is_null = !row.columns.contains_key(&column_id)
10637 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
10638 if is_null == want_nulls {
10639 s.insert(row.row_id.0);
10640 }
10641 }
10642 }
10643
10644 pub fn snapshot(&self) -> Snapshot {
10645 let epoch = self.epoch.visible();
10646 match &self.wal {
10650 WalSink::Shared(shared) => match shared.hlc.now() {
10651 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
10652 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
10654 },
10655 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
10656 }
10657 }
10658
10659 pub fn data_generation(&self) -> u64 {
10661 self.data_generation
10662 }
10663
10664 pub(crate) fn bump_data_generation(&mut self) {
10665 self.data_generation = self.data_generation.wrapping_add(1);
10666 }
10667
10668 pub fn table_id(&self) -> u64 {
10670 self.table_id
10671 }
10672
10673 fn seal_generations(&mut self) {
10679 self.memtable.seal();
10680 self.mutable_run.seal();
10681 self.hot.seal();
10682 for index in self.bitmap.values_mut() {
10683 index.seal();
10684 }
10685 for index in self.ann.values_mut() {
10686 index.seal();
10687 }
10688 for index in self.fm.values_mut() {
10689 index.seal();
10690 }
10691 for index in self.sparse.values_mut() {
10692 index.seal();
10693 }
10694 for index in self.minhash.values_mut() {
10695 index.seal();
10696 }
10697 self.pk_by_row.seal();
10698 }
10699
10700 fn capture_read_generation(&self) -> ReadGeneration {
10705 let visible_through = self.current_epoch();
10706 ReadGeneration {
10707 schema: Arc::new(self.schema.clone()),
10708 base_runs: Arc::new(self.run_refs.clone()),
10709 deltas: TableDeltas {
10710 memtable: self.memtable.clone(),
10711 mutable_run: self.mutable_run.clone(),
10712 hot: self.hot.clone(),
10713 pk_by_row: self.pk_by_row.clone(),
10714 },
10715 indexes: Arc::new(IndexGeneration::capture(
10716 &self.bitmap,
10717 &self.learned_range,
10718 &self.fm,
10719 &self.ann,
10720 &self.sparse,
10721 &self.minhash,
10722 visible_through,
10723 match &self.wal {
10725 WalSink::Shared(shared) => shared
10726 .hlc
10727 .now()
10728 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
10729 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
10730 },
10731 )),
10732 visible_through,
10733 }
10734 }
10735
10736 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
10741 self.ensure_indexes_complete()?;
10742 self.seal_generations();
10743 let view = Arc::new(self.capture_read_generation());
10744 self.published.store(Arc::clone(&view));
10745 Ok(view)
10746 }
10747
10748 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
10754 self.published.load_full()
10755 }
10756
10757 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
10759 &self.pins
10760 }
10761
10762 pub fn version_gc_floor(&self) -> Epoch {
10767 self.min_active_snapshot()
10768 .unwrap_or_else(|| self.current_epoch())
10769 }
10770
10771 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
10778 let mut report = self.pins.report();
10779 let transaction_floor = [
10780 self.pinned.keys().next().copied(),
10781 self.snapshots.min_pinned(),
10782 ]
10783 .into_iter()
10784 .flatten()
10785 .min();
10786 if let Some(epoch) = transaction_floor {
10787 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
10788 }
10789 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
10790 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
10791 }
10792 report
10793 }
10794
10795 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
10796 self.publish_read_generation()?;
10797 let mut generation = self.clone();
10798 generation.read_only = true;
10799 generation.wal = WalSink::ReadOnly;
10800 generation.pending_delete_rids.clear();
10801 generation.pending_put_cols.clear();
10802 generation.pending_rows.clear();
10803 generation.pending_rows_auto_inc.clear();
10804 generation.pending_dels.clear();
10805 generation.pending_truncate = None;
10806 generation.agg_cache = Arc::new(HashMap::new());
10807 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
10810 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
10813 crate::retention::PinSource::ReadGeneration,
10814 self.current_epoch(),
10815 )));
10816 Ok(generation)
10817 }
10818
10819 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
10820 (std::mem::size_of::<Self>() as u64)
10821 .saturating_add(self.memtable.approx_bytes())
10822 .saturating_add(self.mutable_run.approx_bytes())
10823 .saturating_add(self.live_count.saturating_mul(64))
10824 }
10825
10826 pub fn pin_snapshot(&mut self) -> Snapshot {
10833 let snap = self.snapshot();
10834 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
10835 snap
10836 }
10837
10838 fn active_pin_epochs_for_rebuild(&self) -> Vec<Epoch> {
10850 let mut set = BTreeSet::new();
10851 for epoch in self.pinned.keys().copied() {
10852 set.insert(epoch);
10853 }
10854 for epoch in self.snapshots.live_pinned_epochs() {
10855 set.insert(epoch);
10856 }
10857 for epoch in self.pins.live_pin_epochs() {
10858 set.insert(epoch);
10859 }
10860 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
10861 set.insert(floor);
10862 }
10863 set.into_iter().collect()
10864 }
10865
10866 pub fn hlc_gc_floor(
10876 &self,
10877 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
10878 ) -> crate::epoch::GcFloor {
10879 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
10880 epoch
10881 .and_then(&mut project_epoch)
10882 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
10883 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
10884 };
10885 let transaction = [
10886 self.pinned.keys().next().copied(),
10887 self.snapshots.min_pinned(),
10888 ]
10889 .into_iter()
10890 .flatten()
10891 .min();
10892 let history = self.snapshots.history_floor(self.current_epoch());
10893 crate::epoch::GcFloor {
10894 transaction_snapshot: project(transaction),
10895 history_retention: project(history),
10896 backup_pitr: project(
10897 self.pins
10898 .oldest_for(crate::retention::PinSource::BackupPitr),
10899 ),
10900 replication: project(
10901 self.pins
10902 .oldest_for(crate::retention::PinSource::Replication),
10903 ),
10904 read_generation: project(
10905 self.pins
10906 .oldest_for(crate::retention::PinSource::ReadGeneration),
10907 ),
10908 online_index_build: project(
10909 self.pins
10910 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
10911 ),
10912 }
10913 }
10914
10915 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
10917 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
10918 *count -= 1;
10919 if *count == 0 {
10920 self.pinned.remove(&snap.epoch);
10921 }
10922 }
10923 }
10924
10925 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
10937 let local = self.pinned.keys().next().copied();
10938 let global = self.snapshots.min_pinned();
10939 let history = self.snapshots.history_floor(self.current_epoch());
10940 let pinned = self.pins.oldest_pinned();
10941 [local, global, history, pinned].into_iter().flatten().min()
10942 }
10943
10944 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
10948 self.ensure_writable()?;
10949 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
10950 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
10951 }
10952
10953 pub fn clear_ttl(&mut self) -> Result<()> {
10954 self.ensure_writable()?;
10955 self.apply_ttl_policy_at(None, self.current_epoch())
10956 }
10957
10958 pub fn ttl(&self) -> Option<TtlPolicy> {
10959 self.ttl
10960 }
10961
10962 pub(crate) fn prepare_ttl_policy(
10963 &self,
10964 column_name: &str,
10965 duration_nanos: u64,
10966 ) -> Result<TtlPolicy> {
10967 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
10968 return Err(MongrelError::InvalidArgument(
10969 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
10970 ));
10971 }
10972 let column = self
10973 .schema
10974 .columns
10975 .iter()
10976 .find(|column| column.name == column_name)
10977 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
10978 if column.ty != TypeId::TimestampNanos {
10979 return Err(MongrelError::Schema(format!(
10980 "TTL column {column_name} must be TimestampNanos, is {:?}",
10981 column.ty
10982 )));
10983 }
10984 Ok(TtlPolicy {
10985 column_id: column.id,
10986 duration_nanos,
10987 })
10988 }
10989
10990 pub(crate) fn apply_ttl_policy_at(
10991 &mut self,
10992 policy: Option<TtlPolicy>,
10993 epoch: Epoch,
10994 ) -> Result<()> {
10995 if let Some(policy) = policy {
10996 let column = self
10997 .schema
10998 .columns
10999 .iter()
11000 .find(|column| column.id == policy.column_id)
11001 .ok_or_else(|| {
11002 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
11003 })?;
11004 if column.ty != TypeId::TimestampNanos
11005 || policy.duration_nanos == 0
11006 || policy.duration_nanos > i64::MAX as u64
11007 {
11008 return Err(MongrelError::Schema("invalid TTL policy".into()));
11009 }
11010 }
11011 self.ttl = policy;
11012 self.agg_cache = Arc::new(HashMap::new());
11013 self.clear_result_cache();
11014 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
11015 self.persist_manifest(epoch)
11016 }
11017
11018 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
11019 let Some(policy) = self.ttl else {
11020 return false;
11021 };
11022 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
11023 return false;
11024 };
11025 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
11026 }
11027
11028 pub fn current_epoch(&self) -> Epoch {
11029 self.epoch.visible()
11030 }
11031
11032 pub fn memtable_len(&self) -> usize {
11033 self.memtable.len()
11034 }
11035
11036 pub fn count(&self) -> u64 {
11039 if self.ttl.is_none()
11040 && self.pending_put_cols.is_empty()
11041 && self.pending_delete_rids.is_empty()
11042 && self.pending_rows.is_empty()
11043 && self.pending_dels.is_empty()
11044 && self.pending_truncate.is_none()
11045 {
11046 self.live_count
11047 } else {
11048 self.visible_rows(self.snapshot())
11049 .map(|rows| rows.len() as u64)
11050 .unwrap_or(self.live_count)
11051 }
11052 }
11053
11054 pub fn count_conditions(
11058 &mut self,
11059 conditions: &[crate::query::Condition],
11060 snapshot: Snapshot,
11061 ) -> Result<Option<u64>> {
11062 use crate::query::Condition;
11063 if self.ttl.is_some() {
11064 if conditions.is_empty() {
11065 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
11066 }
11067 let mut sets = Vec::with_capacity(conditions.len());
11068 for condition in conditions {
11069 sets.push(self.resolve_condition(condition, snapshot)?);
11070 }
11071 let survivors = RowIdSet::intersect_many(sets);
11072 let rows = self.visible_rows(snapshot)?;
11073 return Ok(Some(
11074 rows.into_iter()
11075 .filter(|row| survivors.contains(row.row_id.0))
11076 .count() as u64,
11077 ));
11078 }
11079 if conditions.is_empty() {
11080 return Ok(Some(self.count()));
11081 }
11082 let served = |c: &Condition| {
11083 matches!(
11084 c,
11085 Condition::Pk(_)
11086 | Condition::BitmapEq { .. }
11087 | Condition::BitmapIn { .. }
11088 | Condition::BytesPrefix { .. }
11089 | Condition::FmContains { .. }
11090 | Condition::FmContainsAll { .. }
11091 | Condition::Ann { .. }
11092 | Condition::Range { .. }
11093 | Condition::RangeF64 { .. }
11094 | Condition::SparseMatch { .. }
11095 | Condition::MinHashSimilar { .. }
11096 | Condition::IsNull { .. }
11097 | Condition::IsNotNull { .. }
11098 )
11099 };
11100 if !conditions.iter().all(served) {
11101 return Ok(None);
11102 }
11103 self.ensure_indexes_complete()?;
11104 if !self.pending_put_cols.is_empty()
11105 || !self.pending_delete_rids.is_empty()
11106 || !self.pending_rows.is_empty()
11107 || !self.pending_dels.is_empty()
11108 || self.pending_truncate.is_some()
11109 {
11110 let mut sets = Vec::with_capacity(conditions.len());
11111 for condition in conditions {
11112 sets.push(self.resolve_condition(condition, snapshot)?);
11113 }
11114 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
11115 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
11116 }
11117 let mut sets = Vec::with_capacity(conditions.len());
11118 for condition in conditions {
11119 sets.push(self.resolve_condition(condition, snapshot)?);
11120 }
11121 let rids = RowIdSet::intersect_many(sets);
11122 if self.had_deletes || !self.memtable.is_empty() || !self.mutable_run.is_empty() {
11130 let sorted = rids.into_sorted_vec();
11131 let count = self.rows_for_rids(&sorted, snapshot)?.len() as u64;
11132 crate::trace::QueryTrace::record(|t| {
11133 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
11134 t.survivor_count = Some(count as usize);
11135 t.conditions_pushed = conditions.len();
11136 });
11137 return Ok(Some(count));
11138 }
11139 let count = rids.len() as u64;
11140 crate::trace::QueryTrace::record(|t| {
11141 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
11142 t.survivor_count = Some(count as usize);
11143 t.conditions_pushed = conditions.len();
11144 });
11145 Ok(Some(count))
11146 }
11147
11148 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
11154 let mut out = Vec::new();
11155 for row in self.memtable.visible_versions_at(snapshot) {
11156 if row.deleted {
11157 out.push(row.row_id.0);
11158 }
11159 }
11160 for row in self.mutable_run.visible_versions_at(snapshot) {
11161 if row.deleted {
11162 out.push(row.row_id.0);
11163 }
11164 }
11165 out
11166 }
11167
11168 pub fn bulk_load_columns(
11177 &mut self,
11178 user_columns: Vec<(u16, columnar::NativeColumn)>,
11179 ) -> Result<Epoch> {
11180 self.bulk_load_columns_with(user_columns, 3, false, true)
11181 }
11182
11183 pub fn bulk_load_fast(
11190 &mut self,
11191 user_columns: Vec<(u16, columnar::NativeColumn)>,
11192 ) -> Result<Epoch> {
11193 self.bulk_load_columns_with(user_columns, -1, true, false)
11194 }
11195
11196 fn bulk_load_columns_with(
11197 &mut self,
11198 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
11199 zstd_level: i32,
11200 force_plain: bool,
11201 lz4: bool,
11202 ) -> Result<Epoch> {
11203 self.ensure_writable()?;
11204 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
11205 if n == 0 {
11206 return Ok(self.current_epoch());
11207 }
11208 let epoch = self.commit_new_epoch()?;
11209 let live_before = self.live_count;
11210 self.spill_mutable_run(epoch)?;
11212 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
11213 && self.indexes_complete
11214 && self.run_refs.is_empty()
11215 && self.memtable.is_empty()
11216 && self.mutable_run.is_empty();
11217 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
11220 self.validate_columns_not_null(&user_columns, n)?;
11221 let winner_idx = self
11222 .bulk_pk_winner_indices(&user_columns, n)
11223 .filter(|idx| idx.len() != n);
11224 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
11225 match winner_idx.as_deref() {
11226 Some(idx) => {
11227 let compacted = user_columns
11228 .iter()
11229 .map(|(id, c)| (*id, c.gather(idx)))
11230 .collect();
11231 (compacted, idx.len())
11232 }
11233 None => (user_columns, n),
11234 };
11235 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
11236 let first = self.allocator.alloc_range(write_n as u64)?.0;
11237 for rid in first..first + write_n as u64 {
11238 self.reservoir.offer(rid);
11239 }
11240 let run_id = self.alloc_run_id()?;
11241 let path = self.run_path(run_id);
11242 let mut writer =
11243 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
11244 if force_plain {
11245 writer = writer.with_plain();
11246 } else if lz4 {
11247 writer = writer.with_lz4();
11250 } else {
11251 writer = writer.with_zstd_level(zstd_level);
11252 }
11253 if let Some(kek) = &self.kek {
11254 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
11255 }
11256 let header = match self.create_run_file(run_id)? {
11257 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
11258 None => writer.write_native(&path, &write_columns, write_n, first)?,
11259 };
11260 self.run_refs.push(RunRef {
11261 run_id: run_id as u128,
11262 level: 0,
11263 epoch_created: epoch.0,
11264 row_count: header.row_count,
11265 });
11266 self.run_row_id_ranges
11267 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
11268 self.live_count = self.live_count.saturating_add(write_n as u64);
11269 if eager_index_build {
11270 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
11271 self.index_columns_bulk(&write_columns, &row_ids);
11272 self.indexes_complete = true;
11273 self.build_learned_ranges()?;
11274 } else {
11275 self.indexes_complete = false;
11279 }
11280 self.mark_flushed(epoch)?;
11281 self.persist_manifest(epoch)?;
11282 if eager_index_build {
11283 self.checkpoint_indexes(epoch);
11284 }
11285 self.clear_result_cache();
11286 self.data_generation = self.data_generation.wrapping_add(1);
11287 Ok(epoch)
11288 }
11289
11290 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
11308 let n = row_ids.len();
11309 if n == 0 {
11310 return;
11311 }
11312 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
11313 columns.iter().map(|(id, c)| (*id, c)).collect();
11314 let ty_of: std::collections::HashMap<u16, TypeId> = self
11315 .schema
11316 .columns
11317 .iter()
11318 .map(|c| (c.id, c.ty.clone()))
11319 .collect();
11320 let pk_id = self.schema.primary_key().map(|c| c.id);
11321
11322 for (i, &rid) in row_ids.iter().enumerate() {
11323 let row_id = RowId(rid);
11324 if let Some(pid) = pk_id {
11325 if let Some(col) = by_id.get(&pid) {
11326 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
11327 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
11328 self.insert_hot_pk(key, row_id);
11329 }
11330 }
11331 }
11332 for idef in &self.schema.indexes {
11333 let Some(col) = by_id.get(&idef.column_id) else {
11334 continue;
11335 };
11336 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
11337 match idef.kind {
11338 IndexKind::Bitmap => {
11339 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
11340 if let Some(key) =
11341 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
11342 {
11343 b.insert(key, row_id);
11344 }
11345 }
11346 }
11347 IndexKind::FmIndex => {
11348 if let Some(f) = self.fm.get_mut(&idef.column_id) {
11349 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11350 f.insert(bytes.to_vec(), row_id);
11351 }
11352 }
11353 }
11354 IndexKind::Sparse => {
11355 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
11356 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11357 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
11358 s.insert(&terms, row_id);
11359 }
11360 }
11361 }
11362 }
11363 IndexKind::MinHash => {
11364 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
11365 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11366 let tokens = crate::index::token_hashes_from_bytes(bytes);
11367 mh.insert(&tokens, row_id);
11368 }
11369 }
11370 }
11371 _ => {}
11372 }
11373 }
11374 }
11375 }
11376
11377 pub fn visible_columns_native(
11382 &self,
11383 snapshot: Snapshot,
11384 projection: Option<&[u16]>,
11385 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11386 self.visible_columns_native_inner(snapshot, projection, None)
11387 }
11388
11389 pub fn visible_columns_native_with_control(
11390 &self,
11391 snapshot: Snapshot,
11392 projection: Option<&[u16]>,
11393 control: &crate::ExecutionControl,
11394 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11395 self.visible_columns_native_inner(snapshot, projection, Some(control))
11396 }
11397
11398 fn visible_columns_native_inner(
11399 &self,
11400 snapshot: Snapshot,
11401 projection: Option<&[u16]>,
11402 control: Option<&crate::ExecutionControl>,
11403 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11404 execution_checkpoint(control, 0)?;
11405 let wanted: Vec<u16> = match projection {
11406 Some(p) => p.to_vec(),
11407 None => self.schema.columns.iter().map(|c| c.id).collect(),
11408 };
11409 if self.ttl.is_none()
11410 && self.memtable.is_empty()
11411 && self.mutable_run.is_empty()
11412 && self.run_refs.len() == 1
11413 {
11414 let rr = self.run_refs[0].clone();
11415 let mut reader = self.open_reader(rr.run_id)?;
11416 let idxs = reader.visible_indices_native_at(snapshot)?;
11417 execution_checkpoint(control, 0)?;
11418 let all_visible = idxs.len() == reader.row_count();
11419 if reader.has_mmap() && control.is_none() {
11425 use rayon::prelude::*;
11426 let valid: Vec<u16> = wanted
11429 .iter()
11430 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
11431 .copied()
11432 .collect();
11433 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
11435 .par_iter()
11436 .filter_map(|cid| {
11437 reader
11438 .column_native_shared(*cid)
11439 .ok()
11440 .map(|col| (*cid, col))
11441 })
11442 .collect();
11443 let cols = decoded
11444 .into_iter()
11445 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
11446 .collect();
11447 return Ok(cols);
11448 }
11449 let mut cols = Vec::with_capacity(wanted.len());
11450 for (index, cid) in wanted.iter().enumerate() {
11451 execution_checkpoint(control, index)?;
11452 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
11453 Some(c) => c,
11454 None => continue,
11455 };
11456 let col = reader.column_native(cdef.id)?;
11457 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
11458 }
11459 return Ok(cols);
11460 }
11461 let vcols = self.visible_columns(snapshot)?;
11462 execution_checkpoint(control, 0)?;
11463 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
11464 let out: Vec<(u16, columnar::NativeColumn)> = vcols
11465 .into_iter()
11466 .filter(|(id, _)| want_set.contains(id))
11467 .map(|(id, vals)| {
11468 let ty = self
11469 .schema
11470 .columns
11471 .iter()
11472 .find(|c| c.id == id)
11473 .map(|c| c.ty.clone())
11474 .unwrap_or(TypeId::Bytes);
11475 (id, columnar::values_to_native(ty, &vals))
11476 })
11477 .collect();
11478 Ok(out)
11479 }
11480
11481 pub fn run_count(&self) -> usize {
11482 self.run_refs.len()
11483 }
11484
11485 pub fn memtable_is_empty(&self) -> bool {
11487 self.memtable.is_empty()
11488 }
11489
11490 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
11494 self.page_cache.stats()
11495 }
11496
11497 pub fn reset_page_cache_stats(&self) {
11499 self.page_cache.reset_stats();
11500 }
11501
11502 pub fn run_ids(&self) -> Vec<u128> {
11505 self.run_refs.iter().map(|r| r.run_id).collect()
11506 }
11507
11508 pub fn single_run_is_clean(&self) -> bool {
11512 if self.ttl.is_some() || self.run_refs.len() != 1 {
11513 return false;
11514 }
11515 self.open_reader(self.run_refs[0].run_id)
11516 .map(|r| r.is_clean())
11517 .unwrap_or(false)
11518 }
11519
11520 fn resolve_footprint(
11527 &self,
11528 conditions: &[crate::query::Condition],
11529 snapshot: Snapshot,
11530 ) -> roaring::RoaringBitmap {
11531 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
11532 return roaring::RoaringBitmap::new();
11533 }
11534 if self.run_refs.is_empty() {
11535 return roaring::RoaringBitmap::new();
11536 }
11537 if self.run_refs.len() == 1 {
11539 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
11540 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
11541 return rids.to_roaring_lossy();
11542 }
11543 }
11544 }
11545 roaring::RoaringBitmap::new()
11546 }
11547
11548 pub fn query_columns_native_cached(
11559 &mut self,
11560 conditions: &[crate::query::Condition],
11561 projection: Option<&[u16]>,
11562 snapshot: Snapshot,
11563 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11564 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
11565 }
11566
11567 pub fn query_columns_native_cached_with_control(
11568 &mut self,
11569 conditions: &[crate::query::Condition],
11570 projection: Option<&[u16]>,
11571 snapshot: Snapshot,
11572 control: &crate::ExecutionControl,
11573 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11574 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
11575 }
11576
11577 fn query_columns_native_cached_inner(
11578 &mut self,
11579 conditions: &[crate::query::Condition],
11580 projection: Option<&[u16]>,
11581 snapshot: Snapshot,
11582 control: Option<&crate::ExecutionControl>,
11583 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11584 execution_checkpoint(control, 0)?;
11585 if self.ttl.is_some() {
11588 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11589 }
11590 if conditions.is_empty() {
11591 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11592 }
11593 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
11597 let identity = PersistentCacheIdentity {
11598 table_id: self.table_id(),
11599 schema_id: self.schema.schema_id,
11600 logical_generation: self.current_epoch().0,
11601 };
11602 if let Some(hit) = self.result_cache.lock().get_columns(key, identity) {
11603 crate::trace::QueryTrace::record(|t| {
11604 t.result_cache_hit = true;
11605 t.scan_mode = crate::trace::ScanMode::NativePushdown;
11606 });
11607 return Ok(Some((*hit).clone()));
11608 }
11609 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
11610 execution_checkpoint(control, 0)?;
11611 if let Some(cols) = &res {
11612 let footprint = self.resolve_footprint(conditions, snapshot);
11613 let condition_cols = crate::query::condition_columns(conditions);
11614 execution_checkpoint(control, 0)?;
11615 let entry = CachedEntry {
11616 data: CachedData::Columns(Arc::new(cols.clone())),
11617 footprint,
11618 condition_cols,
11619 };
11620 let mut cache = self.result_cache.lock();
11626 let entry_generation = cache.allocate_persist_generation(key);
11627 let context = PersistContext {
11628 identity,
11629 key,
11630 entry_generation,
11631 };
11632 cache.persist_entry(context, &entry);
11633 cache.insert(key, entry);
11634 }
11635 Ok(res)
11636 }
11637
11638 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
11643 if self.ttl.is_some() {
11644 return self.query(q);
11645 }
11646 if q.conditions.is_empty() {
11647 return self.query(q);
11648 }
11649 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
11650 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
11651 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
11652 let identity = PersistentCacheIdentity {
11653 table_id: self.table_id(),
11654 schema_id: self.schema.schema_id,
11655 logical_generation: self.current_epoch().0,
11656 };
11657 if let Some(hit) = self.result_cache.lock().get_rows(key, identity) {
11658 crate::trace::QueryTrace::record(|t| {
11659 t.result_cache_hit = true;
11660 t.scan_mode = crate::trace::ScanMode::Materialized;
11661 });
11662 return Ok((*hit).clone());
11663 }
11664 let rows = self.query(q)?;
11665 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
11666 let condition_cols = crate::query::condition_columns(&q.conditions);
11667 let entry = CachedEntry {
11668 data: CachedData::Rows(Arc::new(rows.clone())),
11669 footprint,
11670 condition_cols,
11671 };
11672 let mut cache = self.result_cache.lock();
11675 let entry_generation = cache.allocate_persist_generation(key);
11676 let context = PersistContext {
11677 identity,
11678 key,
11679 entry_generation,
11680 };
11681 cache.persist_entry(context, &entry);
11682 cache.insert(key, entry);
11683 Ok(rows)
11684 }
11685
11686 #[allow(clippy::type_complexity)]
11701 pub fn query_columns_native_traced(
11702 &mut self,
11703 conditions: &[crate::query::Condition],
11704 projection: Option<&[u16]>,
11705 snapshot: Snapshot,
11706 ) -> Result<(
11707 Option<Vec<(u16, columnar::NativeColumn)>>,
11708 crate::trace::QueryTrace,
11709 )> {
11710 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11711 self.query_columns_native(conditions, projection, snapshot)
11712 });
11713 Ok((result?, trace))
11714 }
11715
11716 #[allow(clippy::type_complexity)]
11719 pub fn query_columns_native_cached_traced(
11720 &mut self,
11721 conditions: &[crate::query::Condition],
11722 projection: Option<&[u16]>,
11723 snapshot: Snapshot,
11724 ) -> Result<(
11725 Option<Vec<(u16, columnar::NativeColumn)>>,
11726 crate::trace::QueryTrace,
11727 )> {
11728 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11729 self.query_columns_native_cached(conditions, projection, snapshot)
11730 });
11731 Ok((result?, trace))
11732 }
11733
11734 pub fn native_page_cursor_traced(
11736 &self,
11737 snapshot: Snapshot,
11738 projection: Vec<(u16, TypeId)>,
11739 conditions: &[crate::query::Condition],
11740 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
11741 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11742 self.native_page_cursor(snapshot, projection, conditions)
11743 });
11744 Ok((result?, trace))
11745 }
11746
11747 pub fn native_multi_run_cursor_traced(
11749 &self,
11750 snapshot: Snapshot,
11751 projection: Vec<(u16, TypeId)>,
11752 conditions: &[crate::query::Condition],
11753 ) -> Result<(
11754 Option<crate::cursor::MultiRunCursor>,
11755 crate::trace::QueryTrace,
11756 )> {
11757 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11758 self.native_multi_run_cursor(snapshot, projection, conditions)
11759 });
11760 Ok((result?, trace))
11761 }
11762
11763 pub fn count_conditions_traced(
11765 &mut self,
11766 conditions: &[crate::query::Condition],
11767 snapshot: Snapshot,
11768 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
11769 let (result, trace) =
11770 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
11771 Ok((result?, trace))
11772 }
11773
11774 pub fn query_traced(
11776 &mut self,
11777 q: &crate::query::Query,
11778 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
11779 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
11780 Ok((result?, trace))
11781 }
11782
11783 pub fn query_columns_native(
11788 &mut self,
11789 conditions: &[crate::query::Condition],
11790 projection: Option<&[u16]>,
11791 snapshot: Snapshot,
11792 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11793 self.query_columns_native_inner(conditions, projection, snapshot, None)
11794 }
11795
11796 pub fn query_columns_native_with_control(
11797 &mut self,
11798 conditions: &[crate::query::Condition],
11799 projection: Option<&[u16]>,
11800 snapshot: Snapshot,
11801 control: &crate::ExecutionControl,
11802 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11803 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
11804 }
11805
11806 fn query_columns_native_inner(
11807 &mut self,
11808 conditions: &[crate::query::Condition],
11809 projection: Option<&[u16]>,
11810 snapshot: Snapshot,
11811 control: Option<&crate::ExecutionControl>,
11812 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11813 use crate::query::Condition;
11814 execution_checkpoint(control, 0)?;
11815 if self.ttl.is_some() {
11818 return Ok(None);
11819 }
11820 if conditions.is_empty() {
11821 return Ok(None);
11822 }
11823 self.ensure_indexes_complete()?;
11824
11825 let served = |c: &Condition| {
11830 matches!(
11831 c,
11832 Condition::Pk(_)
11833 | Condition::BitmapEq { .. }
11834 | Condition::BitmapIn { .. }
11835 | Condition::BytesPrefix { .. }
11836 | Condition::FmContains { .. }
11837 | Condition::FmContainsAll { .. }
11838 | Condition::Ann { .. }
11839 | Condition::Range { .. }
11840 | Condition::RangeF64 { .. }
11841 | Condition::SparseMatch { .. }
11842 | Condition::MinHashSimilar { .. }
11843 | Condition::IsNull { .. }
11844 | Condition::IsNotNull { .. }
11845 )
11846 };
11847 if !conditions.iter().all(served) {
11848 return Ok(None);
11849 }
11850 let fast_path =
11851 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
11852 crate::trace::QueryTrace::record(|t| {
11853 t.run_count = self.run_refs.len();
11854 t.memtable_rows = self.memtable.len();
11855 t.mutable_run_rows = self.mutable_run.len();
11856 t.conditions_pushed = conditions.len();
11857 t.learned_range_used = conditions.iter().any(|c| match c {
11858 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
11859 self.learned_range.contains_key(column_id)
11860 }
11861 _ => false,
11862 });
11863 });
11864 let col_ids: Vec<u16> = projection
11866 .map(|p| p.to_vec())
11867 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
11868 let proj_pairs: Vec<(u16, TypeId)> = col_ids
11869 .iter()
11870 .map(|&cid| {
11871 let ty = self
11872 .schema
11873 .columns
11874 .iter()
11875 .find(|c| c.id == cid)
11876 .map(|c| c.ty.clone())
11877 .unwrap_or(TypeId::Bytes);
11878 (cid, ty)
11879 })
11880 .collect();
11881
11882 if fast_path {
11888 let needs_column = conditions.iter().any(|c| match c {
11891 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
11892 Condition::RangeF64 { column_id, .. } => {
11893 !self.learned_range.contains_key(column_id)
11894 }
11895 _ => false,
11896 });
11897 let mut reader_opt: Option<RunReader> = if needs_column {
11898 Some(self.open_reader(self.run_refs[0].run_id)?)
11899 } else {
11900 None
11901 };
11902 let mut sets: Vec<RowIdSet> = Vec::new();
11903 for (index, c) in conditions.iter().enumerate() {
11904 execution_checkpoint(control, index)?;
11905 let s = match c {
11906 Condition::Range { column_id, lo, hi }
11907 if !self.learned_range.contains_key(column_id) =>
11908 {
11909 if reader_opt.is_none() {
11910 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
11911 }
11912 reader_opt
11913 .as_mut()
11914 .expect("reader opened for range")
11915 .range_row_id_set_i64(*column_id, *lo, *hi)?
11916 }
11917 Condition::RangeF64 {
11918 column_id,
11919 lo,
11920 lo_inclusive,
11921 hi,
11922 hi_inclusive,
11923 } if !self.learned_range.contains_key(column_id) => {
11924 if reader_opt.is_none() {
11925 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
11926 }
11927 reader_opt
11928 .as_mut()
11929 .expect("reader opened for range")
11930 .range_row_id_set_f64(
11931 *column_id,
11932 *lo,
11933 *lo_inclusive,
11934 *hi,
11935 *hi_inclusive,
11936 )?
11937 }
11938 _ => self.resolve_condition(c, snapshot)?,
11939 };
11940 sets.push(s);
11941 }
11942 let candidates = RowIdSet::intersect_many(sets);
11943 crate::trace::QueryTrace::record(|t| {
11944 t.survivor_count = Some(candidates.len());
11945 });
11946 if candidates.is_empty() {
11947 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
11948 .iter()
11949 .map(|&id| {
11950 (
11951 id,
11952 columnar::null_native(
11953 proj_pairs
11954 .iter()
11955 .find(|(c, _)| c == &id)
11956 .map(|(_, t)| t.clone())
11957 .unwrap_or(TypeId::Bytes),
11958 0,
11959 ),
11960 )
11961 })
11962 .collect();
11963 return Ok(Some(cols));
11964 }
11965 let mut reader = match reader_opt.take() {
11966 Some(r) => r,
11967 None => self.open_reader(self.run_refs[0].run_id)?,
11968 };
11969 let candidate_ids = candidates.into_sorted_vec();
11970 let (positions, fast_rid) = if let Some(positions) =
11971 reader.positions_for_row_ids_fast(&candidate_ids)
11972 {
11973 (positions, true)
11974 } else {
11975 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
11976 match col {
11977 columnar::NativeColumn::Int64 { data, .. } => {
11978 let mut p = Vec::with_capacity(candidate_ids.len());
11979 for (index, rid) in candidate_ids.iter().enumerate() {
11980 execution_checkpoint(control, index)?;
11981 if let Ok(position) = data.binary_search(&(*rid as i64)) {
11982 p.push(position);
11983 }
11984 }
11985 p.sort_unstable();
11986 (p, false)
11987 }
11988 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
11989 }
11990 };
11991 crate::trace::QueryTrace::record(|t| {
11992 t.scan_mode = crate::trace::ScanMode::NativePushdown;
11993 t.fast_row_id_map = fast_rid;
11994 });
11995 let mut cols = Vec::with_capacity(col_ids.len());
11996 for (index, cid) in col_ids.iter().enumerate() {
11997 execution_checkpoint(control, index)?;
11998 let col = reader.column_native(*cid)?;
11999 cols.push((*cid, col.gather(&positions)));
12000 }
12001 return Ok(Some(cols));
12002 }
12003
12004 if !self.run_refs.is_empty() {
12017 use crate::cursor::{
12018 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
12019 };
12020 let remaining: usize;
12021 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
12022 let c = self
12023 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
12024 .expect("single-run cursor should build when run_refs.len() == 1");
12025 remaining = c.remaining_rows();
12026 Box::new(c)
12027 } else {
12028 let c = self
12029 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
12030 .expect("multi-run cursor should build when run_refs.len() >= 1");
12031 remaining = c.remaining_rows();
12032 Box::new(c)
12033 };
12034 crate::trace::QueryTrace::record(|t| {
12035 if t.survivor_count.is_none() {
12036 t.survivor_count = Some(remaining);
12037 }
12038 });
12039 let cols = match control {
12040 Some(control) => {
12041 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
12042 }
12043 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
12044 };
12045 return Ok(Some(cols));
12046 }
12047
12048 crate::trace::QueryTrace::record(|t| {
12053 t.scan_mode = crate::trace::ScanMode::Materialized;
12054 t.row_materialized = true;
12055 });
12056 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12057 for (index, c) in conditions.iter().enumerate() {
12058 execution_checkpoint(control, index)?;
12059 sets.push(self.resolve_condition(c, snapshot)?);
12060 }
12061 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
12062 let rows = self.rows_for_rids(&rids, snapshot)?;
12063 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
12064 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
12065 execution_checkpoint(control, index)?;
12066 let vals: Vec<Value> = rows
12067 .iter()
12068 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
12069 .collect();
12070 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
12071 }
12072 Ok(Some(cols))
12073 }
12074
12075 pub fn native_page_cursor(
12090 &self,
12091 snapshot: Snapshot,
12092 projection: Vec<(u16, TypeId)>,
12093 conditions: &[crate::query::Condition],
12094 ) -> Result<Option<NativePageCursor>> {
12095 use crate::cursor::build_page_plans;
12096 if self.ttl.is_some() {
12097 return Ok(None);
12098 }
12099 if !conditions.is_empty() && !self.indexes_complete {
12102 return Ok(None);
12103 }
12104 if self.run_refs.len() != 1 {
12105 return Ok(None);
12106 }
12107 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
12108 let (positions, rids) = reader.visible_positions_with_rids_at(snapshot)?;
12109
12110 let overlay_rids: HashSet<u64> = {
12113 let mut s = HashSet::new();
12114 for row in self.memtable.visible_versions_at(snapshot) {
12115 s.insert(row.row_id.0);
12116 }
12117 for row in self.mutable_run.visible_versions_at(snapshot) {
12118 s.insert(row.row_id.0);
12119 }
12120 s
12121 };
12122
12123 let survivors = if conditions.is_empty() {
12127 None
12128 } else {
12129 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
12130 };
12131
12132 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
12139 survivors.clone()
12140 } else if let Some(s) = &survivors {
12141 let mut run_set = s.clone();
12142 run_set.remove_many(overlay_rids.iter().copied());
12143 Some(run_set)
12144 } else {
12145 Some(RowIdSet::from_unsorted(
12146 rids.iter()
12147 .map(|&r| r as u64)
12148 .filter(|r| !overlay_rids.contains(r))
12149 .collect(),
12150 ))
12151 };
12152
12153 let overlay_rows = if overlay_rids.is_empty() {
12154 Vec::new()
12155 } else {
12156 let bound = Self::overlay_materialization_bound(conditions, &survivors);
12157 self.overlay_visible_rows(snapshot, bound)
12158 };
12159
12160 let plans = if positions.is_empty() {
12162 Vec::new()
12163 } else {
12164 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
12165 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
12166 };
12167
12168 let overlay = if overlay_rows.is_empty() {
12170 None
12171 } else {
12172 let filtered =
12173 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
12174 if filtered.is_empty() {
12175 None
12176 } else {
12177 Some(self.materialize_overlay(&filtered, &projection))
12178 }
12179 };
12180
12181 let overlay_row_count = overlay
12182 .as_ref()
12183 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
12184 .unwrap_or(0);
12185 crate::trace::QueryTrace::record(|t| {
12186 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
12187 t.run_count = self.run_refs.len();
12188 t.memtable_rows = self.memtable.len();
12189 t.mutable_run_rows = self.mutable_run.len();
12190 t.overlay_rows = overlay_row_count;
12191 t.conditions_pushed = conditions.len();
12192 t.pages_decoded = plans
12193 .iter()
12194 .map(|p| p.positions.len())
12195 .sum::<usize>()
12196 .min(1);
12197 });
12198
12199 Ok(Some(NativePageCursor::new_with_overlay(
12200 reader, projection, plans, overlay,
12201 )))
12202 }
12203 #[allow(clippy::type_complexity)]
12213 pub fn native_multi_run_cursor(
12214 &self,
12215 snapshot: Snapshot,
12216 projection: Vec<(u16, TypeId)>,
12217 conditions: &[crate::query::Condition],
12218 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
12219 use crate::cursor::{MultiRunCursor, RunStream};
12220 use crate::sorted_run::SYS_ROW_ID;
12221 use std::collections::{BinaryHeap, HashMap, HashSet};
12222 if self.ttl.is_some() {
12223 return Ok(None);
12224 }
12225 if !conditions.is_empty() && !self.indexes_complete {
12228 return Ok(None);
12229 }
12230 if self.run_refs.is_empty() {
12231 return Ok(None);
12232 }
12233
12234 let mut run_meta: Vec<(
12236 RunReader,
12237 Vec<i64>,
12238 Vec<i64>,
12239 Vec<u8>,
12240 Option<Vec<Option<HlcTimestamp>>>,
12241 Vec<usize>,
12242 )> = Vec::with_capacity(self.run_refs.len());
12243 for rr in &self.run_refs {
12244 let mut reader = self.open_reader(rr.run_id)?;
12245 let (rids, eps, del) = reader.system_columns_native()?;
12246 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
12247 let commit_ts: Option<Vec<Option<HlcTimestamp>>> =
12248 if reader.has_column(crate::sorted_run::SYS_COMMIT_TS) {
12249 let col = reader.column_native(crate::sorted_run::SYS_COMMIT_TS)?;
12250 Some(
12251 (0..rids.len())
12252 .map(|i| {
12253 crate::sorted_run::decode_commit_ts_value(col.value_at(i).as_ref())
12254 })
12255 .collect(),
12256 )
12257 } else {
12258 None
12259 };
12260 run_meta.push((reader, rids, eps, del, commit_ts, page_rows));
12261 }
12262
12263 let mut best: HashMap<u64, (Epoch, Option<HlcTimestamp>, usize, usize, bool)> =
12269 HashMap::new();
12270 for (run_idx, (_, rids, eps, del, commit_ts, _)) in run_meta.iter().enumerate() {
12271 for i in 0..rids.len() {
12272 let rid = rids[i] as u64;
12273 let e = Epoch(eps[i] as u64);
12274 let ts = commit_ts.as_ref().and_then(|v| v.get(i).copied().flatten());
12275 if !snapshot.observes_row(e, ts) {
12276 continue;
12277 }
12278 let is_del = del[i] != 0;
12279 best.entry(rid)
12280 .and_modify(|cur| {
12281 if Snapshot::version_is_newer(e, ts, cur.0, cur.1) {
12282 *cur = (e, ts, run_idx, i, is_del);
12283 }
12284 })
12285 .or_insert((e, ts, run_idx, i, is_del));
12286 }
12287 }
12288
12289 let overlay_rids: HashSet<u64> = {
12291 let mut s = HashSet::new();
12292 for row in self.memtable.visible_versions_at(snapshot) {
12293 s.insert(row.row_id.0);
12294 }
12295 for row in self.mutable_run.visible_versions_at(snapshot) {
12296 s.insert(row.row_id.0);
12297 }
12298 s
12299 };
12300
12301 let survivors: Option<RowIdSet> = if conditions.is_empty() {
12303 None
12304 } else {
12305 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12306 for c in conditions {
12307 sets.push(self.resolve_condition(c, snapshot)?);
12308 }
12309 Some(RowIdSet::intersect_many(sets))
12310 };
12311
12312 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
12316 for (rid, (_, _, run_idx, pos, deleted)) in &best {
12317 if *deleted {
12318 continue;
12319 }
12320 if overlay_rids.contains(rid) {
12321 continue;
12322 }
12323 if let Some(s) = &survivors {
12324 if !s.contains(*rid) {
12325 continue;
12326 }
12327 }
12328 per_run[*run_idx].push((*rid, *pos));
12329 }
12330 for v in per_run.iter_mut() {
12331 v.sort_unstable_by_key(|&(rid, _)| rid);
12332 }
12333
12334 let mut streams = Vec::with_capacity(run_meta.len());
12336 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
12337 let mut total = 0usize;
12338 for (run_idx, (reader, _, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
12339 let mut starts = Vec::with_capacity(page_rows.len());
12340 let mut acc = 0usize;
12341 for &r in &page_rows {
12342 starts.push(acc);
12343 acc += r;
12344 }
12345 let mut survivors_vec: Vec<(u64, usize, usize)> =
12346 Vec::with_capacity(per_run[run_idx].len());
12347 for &(rid, pos) in &per_run[run_idx] {
12348 let page_seq = match starts.partition_point(|&s| s <= pos) {
12349 0 => continue,
12350 p => p - 1,
12351 };
12352 let within = pos - starts[page_seq];
12353 survivors_vec.push((rid, page_seq, within));
12354 }
12355 total += survivors_vec.len();
12356 if let Some(&(rid, _, _)) = survivors_vec.first() {
12357 heap.push(std::cmp::Reverse((rid, run_idx)));
12358 }
12359 streams.push(RunStream::new(reader, survivors_vec, page_rows));
12360 }
12361
12362 let overlay_rows = if overlay_rids.is_empty() {
12364 Vec::new()
12365 } else {
12366 let bound = Self::overlay_materialization_bound(conditions, &survivors);
12367 self.overlay_visible_rows(snapshot, bound)
12368 };
12369 let overlay = if overlay_rows.is_empty() {
12370 None
12371 } else {
12372 let filtered =
12373 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
12374 if filtered.is_empty() {
12375 None
12376 } else {
12377 Some(self.materialize_overlay(&filtered, &projection))
12378 }
12379 };
12380
12381 let overlay_row_count = overlay
12382 .as_ref()
12383 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
12384 .unwrap_or(0);
12385 crate::trace::QueryTrace::record(|t| {
12386 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
12387 t.run_count = self.run_refs.len();
12388 t.memtable_rows = self.memtable.len();
12389 t.mutable_run_rows = self.mutable_run.len();
12390 t.overlay_rows = overlay_row_count;
12391 t.conditions_pushed = conditions.len();
12392 t.survivor_count = Some(total);
12393 });
12394
12395 Ok(Some(MultiRunCursor::new(
12396 streams, projection, heap, total, overlay,
12397 )))
12398 }
12399
12400 fn overlay_materialization_bound<'a>(
12412 conditions: &[crate::query::Condition],
12413 survivors: &'a Option<RowIdSet>,
12414 ) -> Option<&'a RowIdSet> {
12415 use crate::query::Condition;
12416 let has_range = conditions
12417 .iter()
12418 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12419 if has_range {
12420 None
12421 } else {
12422 survivors.as_ref()
12423 }
12424 }
12425
12426 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
12438 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
12439 let mut fold = |row: Row| {
12440 if let Some(b) = bound {
12441 if !b.contains(row.row_id.0) {
12442 return;
12443 }
12444 }
12445 best.entry(row.row_id.0)
12446 .and_modify(|(be, br)| {
12447 if row.committed_epoch > *be {
12448 *be = row.committed_epoch;
12449 *br = row.clone();
12450 }
12451 })
12452 .or_insert_with(|| (row.committed_epoch, row));
12453 };
12454 for row in self.memtable.visible_versions_at(snapshot) {
12455 fold(row);
12456 }
12457 for row in self.mutable_run.visible_versions_at(snapshot) {
12458 fold(row);
12459 }
12460 let mut out: Vec<Row> = best
12461 .into_values()
12462 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
12463 .collect();
12464 out.sort_by_key(|r| r.row_id);
12465 out
12466 }
12467
12468 fn filter_overlay_rows(
12476 &self,
12477 rows: Vec<Row>,
12478 conditions: &[crate::query::Condition],
12479 survivors: Option<&RowIdSet>,
12480 snapshot: Snapshot,
12481 ) -> Result<Vec<Row>> {
12482 if conditions.is_empty() {
12483 return Ok(rows);
12484 }
12485 use crate::query::Condition;
12486 let all_index_served = !conditions
12490 .iter()
12491 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12492 if all_index_served {
12493 return Ok(rows
12494 .into_iter()
12495 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
12496 .collect());
12497 }
12498 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12501 for c in conditions {
12502 let s = match c {
12503 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
12504 _ => self.resolve_condition(c, snapshot)?,
12505 };
12506 per_cond_sets.push(s);
12507 }
12508 Ok(rows
12509 .into_iter()
12510 .filter(|row| {
12511 conditions.iter().enumerate().all(|(i, c)| match c {
12512 Condition::Range { column_id, lo, hi } => {
12513 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
12514 }
12515 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
12516 match row.columns.get(column_id) {
12517 Some(Value::Float64(v)) => {
12518 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
12519 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
12520 lo_ok && hi_ok
12521 }
12522 _ => false,
12523 }
12524 }
12525 _ => per_cond_sets[i].contains(row.row_id.0),
12526 })
12527 })
12528 .collect())
12529 }
12530
12531 fn materialize_overlay(
12534 &self,
12535 rows: &[Row],
12536 projection: &[(u16, TypeId)],
12537 ) -> Vec<columnar::NativeColumn> {
12538 if projection.is_empty() {
12539 return vec![columnar::null_native(TypeId::Int64, rows.len())];
12540 }
12541 let mut cols = Vec::with_capacity(projection.len());
12542 for (cid, ty) in projection {
12543 let vals: Vec<Value> = rows
12544 .iter()
12545 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
12546 .collect();
12547 cols.push(columnar::values_to_native(ty.clone(), &vals));
12548 }
12549 cols
12550 }
12551
12552 fn resolve_survivor_rids(
12557 &self,
12558 conditions: &[crate::query::Condition],
12559 reader: &mut RunReader,
12560 snapshot: Snapshot,
12561 ) -> Result<RowIdSet> {
12562 use crate::query::Condition;
12563 let mut sets: Vec<RowIdSet> = Vec::new();
12564 for c in conditions {
12565 self.validate_condition(c)?;
12566 let s: RowIdSet = match c {
12567 Condition::Pk(key) => {
12568 let lookup = self
12569 .schema
12570 .primary_key()
12571 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
12572 .unwrap_or_else(|| key.clone());
12573 self.hot
12574 .get(&lookup)
12575 .map(|r| RowIdSet::one(r.0))
12576 .unwrap_or_else(RowIdSet::empty)
12577 }
12578 Condition::BitmapEq { column_id, value } => {
12579 let lookup = self.index_lookup_key_bytes(*column_id, value);
12580 self.bitmap
12581 .get(column_id)
12582 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
12583 .unwrap_or_else(RowIdSet::empty)
12584 }
12585 Condition::BitmapIn { column_id, values } => {
12586 let bm = self.bitmap.get(column_id);
12587 let mut acc = roaring::RoaringBitmap::new();
12588 if let Some(b) = bm {
12589 for v in values {
12590 let lookup = self.index_lookup_key_bytes(*column_id, v);
12591 acc |= b.get(&lookup);
12592 }
12593 }
12594 RowIdSet::from_roaring(acc)
12595 }
12596 Condition::BytesPrefix { column_id, prefix } => {
12597 if let Some(b) = self.bitmap.get(column_id) {
12598 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
12599 let mut acc = roaring::RoaringBitmap::new();
12600 for key in b.keys() {
12601 if key.starts_with(&lookup_prefix) {
12602 acc |= b.get(&key);
12603 }
12604 }
12605 RowIdSet::from_roaring(acc)
12606 } else {
12607 RowIdSet::empty()
12608 }
12609 }
12610 Condition::FmContains { column_id, pattern } => self
12611 .fm
12612 .get(column_id)
12613 .map(|f| {
12614 RowIdSet::from_unsorted(
12615 f.locate(pattern).into_iter().map(|r| r.0).collect(),
12616 )
12617 })
12618 .unwrap_or_else(RowIdSet::empty),
12619 Condition::FmContainsAll {
12620 column_id,
12621 patterns,
12622 } => {
12623 if let Some(f) = self.fm.get(column_id) {
12624 let sets: Vec<RowIdSet> = patterns
12625 .iter()
12626 .map(|pat| {
12627 RowIdSet::from_unsorted(
12628 f.locate(pat).into_iter().map(|r| r.0).collect(),
12629 )
12630 })
12631 .collect();
12632 RowIdSet::intersect_many(sets)
12633 } else {
12634 RowIdSet::empty()
12635 }
12636 }
12637 Condition::Ann {
12638 column_id,
12639 query,
12640 k,
12641 } => RowIdSet::from_unsorted(
12642 self.retrieve_filtered(
12643 &crate::query::Retriever::Ann {
12644 column_id: *column_id,
12645 query: query.clone(),
12646 k: *k,
12647 },
12648 snapshot,
12649 None,
12650 None,
12651 None,
12652 None,
12653 )?
12654 .into_iter()
12655 .map(|hit| hit.row_id.0)
12656 .collect(),
12657 ),
12658 Condition::SparseMatch {
12659 column_id,
12660 query,
12661 k,
12662 } => RowIdSet::from_unsorted(
12663 self.retrieve_filtered(
12664 &crate::query::Retriever::Sparse {
12665 column_id: *column_id,
12666 query: query.clone(),
12667 k: *k,
12668 },
12669 snapshot,
12670 None,
12671 None,
12672 None,
12673 None,
12674 )?
12675 .into_iter()
12676 .map(|hit| hit.row_id.0)
12677 .collect(),
12678 ),
12679 Condition::MinHashSimilar {
12680 column_id,
12681 query,
12682 k,
12683 } => match self.minhash.get(column_id) {
12684 Some(index) => {
12685 let candidates = index.candidate_row_ids(query);
12686 let eligible =
12687 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
12688 RowIdSet::from_unsorted(
12689 index
12690 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
12691 .into_iter()
12692 .map(|(row_id, _)| row_id.0)
12693 .collect(),
12694 )
12695 }
12696 None => RowIdSet::empty(),
12697 },
12698 Condition::Range { column_id, lo, hi } => {
12699 if let Some(li) = self.learned_range.get(column_id) {
12700 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
12701 } else {
12702 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
12703 }
12704 }
12705 Condition::RangeF64 {
12706 column_id,
12707 lo,
12708 lo_inclusive,
12709 hi,
12710 hi_inclusive,
12711 } => {
12712 if let Some(li) = self.learned_range.get(column_id) {
12713 RowIdSet::from_unsorted(
12714 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
12715 .into_iter()
12716 .collect(),
12717 )
12718 } else {
12719 reader.range_row_id_set_f64(
12720 *column_id,
12721 *lo,
12722 *lo_inclusive,
12723 *hi,
12724 *hi_inclusive,
12725 )?
12726 }
12727 }
12728 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
12729 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
12730 };
12731 sets.push(s);
12732 }
12733 Ok(RowIdSet::intersect_many(sets))
12734 }
12735
12736 pub fn scan_cursor(
12757 &self,
12758 snapshot: Snapshot,
12759 projection: Vec<(u16, TypeId)>,
12760 conditions: &[crate::query::Condition],
12761 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
12762 if self.ttl.is_some() {
12763 return Ok(None);
12764 }
12765 if !conditions.is_empty() && !self.indexes_complete {
12771 return Ok(None);
12772 }
12773 if self.run_refs.len() == 1 {
12774 Ok(self
12775 .native_page_cursor(snapshot, projection, conditions)?
12776 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
12777 } else {
12778 Ok(self
12779 .native_multi_run_cursor(snapshot, projection, conditions)?
12780 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
12781 }
12782 }
12783
12784 pub fn aggregate_native(
12798 &self,
12799 snapshot: Snapshot,
12800 column: Option<u16>,
12801 conditions: &[crate::query::Condition],
12802 agg: NativeAgg,
12803 ) -> Result<Option<NativeAggResult>> {
12804 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
12805 }
12806
12807 pub fn aggregate_native_with_control(
12808 &self,
12809 snapshot: Snapshot,
12810 column: Option<u16>,
12811 conditions: &[crate::query::Condition],
12812 agg: NativeAgg,
12813 control: &crate::ExecutionControl,
12814 ) -> Result<Option<NativeAggResult>> {
12815 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
12816 }
12817
12818 fn aggregate_native_inner(
12819 &self,
12820 snapshot: Snapshot,
12821 column: Option<u16>,
12822 conditions: &[crate::query::Condition],
12823 agg: NativeAgg,
12824 control: Option<&crate::ExecutionControl>,
12825 ) -> Result<Option<NativeAggResult>> {
12826 execution_checkpoint(control, 0)?;
12827 if self.ttl.is_some() {
12828 return Ok(None);
12829 }
12830 if self.run_refs.len() == 1 && conditions.is_empty() {
12832 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
12833 return Ok(Some(res));
12834 }
12835 }
12836 if matches!(agg, NativeAgg::Count) && column.is_none() {
12840 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
12841 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
12842 }
12843 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
12844 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
12845 }
12846 let cid = match column {
12849 Some(c) => c,
12850 None => return Ok(None),
12851 };
12852 let ty = self.column_type(cid);
12853 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
12854 execution_checkpoint(control, 0)?;
12855 return match ty {
12856 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
12857 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
12858 Ok(Some(pack_int(agg, count, sum, mn, mx)))
12859 }
12860 TypeId::Float64 => {
12861 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
12862 Ok(Some(pack_float(agg, count, sum, mn, mx)))
12863 }
12864 _ => Ok(None),
12865 };
12866 }
12867 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
12869 execution_checkpoint(control, 0)?;
12870 match ty {
12871 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
12872 let mut count = 0u64;
12873 let mut sum = 0i128;
12874 let mut mn = i64::MAX;
12875 let mut mx = i64::MIN;
12876 for row in &rows {
12877 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
12878 count += 1;
12879 sum += i128::from(*v);
12880 mn = mn.min(*v);
12881 mx = mx.max(*v);
12882 }
12883 }
12884 Ok(Some(pack_int(agg, count, sum, mn, mx)))
12885 }
12886 TypeId::Float64 => {
12887 let mut count = 0u64;
12888 let mut sum = 0.0f64;
12889 let mut mn = f64::INFINITY;
12890 let mut mx = f64::NEG_INFINITY;
12891 for row in &rows {
12892 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
12893 count += 1;
12894 sum += *v;
12895 mn = mn.min(*v);
12896 mx = mx.max(*v);
12897 }
12898 }
12899 Ok(Some(pack_float(agg, count, sum, mn, mx)))
12900 }
12901 _ => Ok(None),
12902 }
12903 }
12904
12905 fn visible_rows_filtered(
12907 &self,
12908 snapshot: Snapshot,
12909 conditions: &[crate::query::Condition],
12910 control: Option<&crate::ExecutionControl>,
12911 ) -> Result<Vec<Row>> {
12912 let rows = if let Some(control) = control {
12913 self.visible_rows_controlled(snapshot, control)?
12914 } else {
12915 self.visible_rows(snapshot)?
12916 };
12917 if conditions.is_empty() {
12918 return Ok(rows);
12919 }
12920 Ok(rows
12921 .into_iter()
12922 .filter(|row| {
12923 conditions
12924 .iter()
12925 .all(|cond| condition_matches_row(cond, row, &self.schema))
12926 })
12927 .collect())
12928 }
12929
12930 fn aggregate_from_stats(
12938 &self,
12939 snapshot: Snapshot,
12940 column: Option<u16>,
12941 agg: NativeAgg,
12942 ) -> Result<Option<NativeAggResult>> {
12943 let cid = match (agg, column) {
12944 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
12945 _ => return Ok(None), };
12947 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
12948 return Ok(None);
12949 };
12950 let Some(cs) = stats.get(&cid) else {
12951 return Ok(None);
12952 };
12953 match agg {
12954 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
12956 self.live_count.saturating_sub(cs.null_count),
12957 ))),
12958 NativeAgg::Min | NativeAgg::Max => {
12959 let bound = if agg == NativeAgg::Min {
12960 &cs.min
12961 } else {
12962 &cs.max
12963 };
12964 match bound {
12965 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
12966 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
12967 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
12972 None => Ok(None),
12973 }
12974 }
12975 _ => Ok(None),
12976 }
12977 }
12978
12979 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
12988 if self.ttl.is_some() {
12989 return Ok(None);
12990 }
12991 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
12992 return Ok(None);
12993 }
12994 self.ensure_indexes_complete()?;
12997 let reader = self.open_reader(self.run_refs[0].run_id)?;
12998 if self.live_count != reader.row_count() as u64 {
12999 return Ok(None);
13000 }
13001 let Some(bm) = self.bitmap.get(&column_id) else {
13002 return Ok(None); };
13004 let mut distinct = bm.value_count() as u64;
13005 if !bm.get(&Value::Null.encode_key()).is_empty() {
13008 distinct = distinct.saturating_sub(1);
13009 }
13010 Ok(Some(distinct))
13011 }
13012
13013 pub fn aggregate_incremental(
13025 &mut self,
13026 cache_key: u64,
13027 conditions: &[crate::query::Condition],
13028 column: Option<u16>,
13029 agg: NativeAgg,
13030 ) -> Result<IncrementalAggResult> {
13031 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
13032 }
13033
13034 pub fn aggregate_incremental_with_control(
13035 &mut self,
13036 cache_key: u64,
13037 conditions: &[crate::query::Condition],
13038 column: Option<u16>,
13039 agg: NativeAgg,
13040 control: &crate::ExecutionControl,
13041 ) -> Result<IncrementalAggResult> {
13042 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
13043 }
13044
13045 fn aggregate_incremental_inner(
13046 &mut self,
13047 cache_key: u64,
13048 conditions: &[crate::query::Condition],
13049 column: Option<u16>,
13050 agg: NativeAgg,
13051 control: Option<&crate::ExecutionControl>,
13052 ) -> Result<IncrementalAggResult> {
13053 execution_checkpoint(control, 0)?;
13054 let snap = self.snapshot();
13055 let cur_wm = self.allocator.current().0;
13056 let cur_epoch = snap.epoch.0;
13057 let incremental_ok = self.ttl.is_none()
13064 && !self.had_deletes
13065 && self.memtable.is_empty()
13066 && self.mutable_run.is_empty();
13067
13068 if incremental_ok {
13071 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
13072 if cached.epoch == cur_epoch {
13073 return Ok(IncrementalAggResult {
13074 state: cached.state,
13075 incremental: true,
13076 delta_rows: 0,
13077 });
13078 }
13079 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
13080 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
13081 let mut delta_rids = Vec::with_capacity(delta_len);
13082 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
13083 execution_checkpoint(control, index)?;
13084 delta_rids.push(row_id);
13085 }
13086 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
13087 execution_checkpoint(control, 0)?;
13088 let index_sets = self.resolve_index_conditions(conditions, snap)?;
13089 let delta_state = agg_state_from_rows(
13090 &delta_rows,
13091 conditions,
13092 &index_sets,
13093 column,
13094 agg,
13095 &self.schema,
13096 control,
13097 )?;
13098 let merged = cached.state.merge(delta_state);
13099 let delta_n = delta_rids.len() as u64;
13100 Arc::make_mut(&mut self.agg_cache).insert(
13101 cache_key,
13102 CachedAgg {
13103 state: merged.clone(),
13104 watermark: cur_wm,
13105 epoch: cur_epoch,
13106 },
13107 );
13108 return Ok(IncrementalAggResult {
13109 state: merged,
13110 incremental: true,
13111 delta_rows: delta_n,
13112 });
13113 }
13114 }
13115 }
13116
13117 let cursor_ok =
13122 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
13123 let state = if cursor_ok && agg != NativeAgg::Avg {
13124 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
13125 Some(result) => {
13126 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
13127 }
13128 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
13129 }
13130 } else {
13131 self.agg_state_full_scan(conditions, column, agg, snap, control)?
13132 };
13133 if incremental_ok {
13135 Arc::make_mut(&mut self.agg_cache).insert(
13136 cache_key,
13137 CachedAgg {
13138 state: state.clone(),
13139 watermark: cur_wm,
13140 epoch: cur_epoch,
13141 },
13142 );
13143 }
13144 Ok(IncrementalAggResult {
13145 state,
13146 incremental: false,
13147 delta_rows: 0,
13148 })
13149 }
13150
13151 fn agg_state_full_scan(
13154 &self,
13155 conditions: &[crate::query::Condition],
13156 column: Option<u16>,
13157 agg: NativeAgg,
13158 snap: Snapshot,
13159 control: Option<&crate::ExecutionControl>,
13160 ) -> Result<AggState> {
13161 execution_checkpoint(control, 0)?;
13162 let rows = self.visible_rows(snap)?;
13163 execution_checkpoint(control, 0)?;
13164 let index_sets = self.resolve_index_conditions(conditions, snap)?;
13165 agg_state_from_rows(
13166 &rows,
13167 conditions,
13168 &index_sets,
13169 column,
13170 agg,
13171 &self.schema,
13172 control,
13173 )
13174 }
13175
13176 fn resolve_index_conditions(
13179 &self,
13180 conditions: &[crate::query::Condition],
13181 snapshot: Snapshot,
13182 ) -> Result<Vec<RowIdSet>> {
13183 use crate::query::Condition;
13184 let mut sets = Vec::new();
13185 for c in conditions {
13186 if matches!(
13187 c,
13188 Condition::Ann { .. }
13189 | Condition::SparseMatch { .. }
13190 | Condition::MinHashSimilar { .. }
13191 ) {
13192 sets.push(self.resolve_condition(c, snapshot)?);
13193 }
13194 }
13195 Ok(sets)
13196 }
13197
13198 fn column_type(&self, cid: u16) -> TypeId {
13199 self.schema
13200 .columns
13201 .iter()
13202 .find(|c| c.id == cid)
13203 .map(|c| c.ty.clone())
13204 .unwrap_or(TypeId::Bytes)
13205 }
13206
13207 pub fn approx_aggregate(
13216 &mut self,
13217 conditions: &[crate::query::Condition],
13218 column: Option<u16>,
13219 agg: ApproxAgg,
13220 z: f64,
13221 ) -> Result<Option<ApproxResult>> {
13222 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
13223 }
13224
13225 pub fn approx_aggregate_with_candidate_authorization(
13228 &mut self,
13229 conditions: &[crate::query::Condition],
13230 column: Option<u16>,
13231 agg: ApproxAgg,
13232 z: f64,
13233 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
13234 ) -> Result<Option<ApproxResult>> {
13235 use crate::query::Condition;
13236 self.ensure_reservoir_complete()?;
13237 let mut snapshot = self.snapshot();
13242 let n_pop = self.count();
13243 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
13244 if sample_rids.is_empty() {
13245 return Ok(None);
13246 }
13247 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
13249 if live_sample.is_empty() {
13250 snapshot = Snapshot::unbounded();
13251 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
13252 }
13253 let s = live_sample.len();
13254 if s == 0 {
13255 return Ok(None);
13256 }
13257 let authorized = authorization
13258 .map(|authorization| {
13259 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
13260 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
13261 })
13262 .transpose()?;
13263
13264 let mut index_sets: Vec<RowIdSet> = Vec::new();
13267 for c in conditions {
13268 if matches!(
13269 c,
13270 Condition::Ann { .. }
13271 | Condition::SparseMatch { .. }
13272 | Condition::MinHashSimilar { .. }
13273 ) {
13274 index_sets.push(self.resolve_condition(c, snapshot)?);
13275 }
13276 }
13277
13278 let cid = match (agg, column) {
13280 (ApproxAgg::Count, _) => None,
13281 (_, Some(c)) => Some(c),
13282 _ => return Ok(None),
13283 };
13284 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
13285 for r in &live_sample {
13286 if authorized
13287 .as_ref()
13288 .is_some_and(|authorized| !authorized.contains(&r.row_id))
13289 {
13290 continue;
13291 }
13292 if !conditions
13294 .iter()
13295 .all(|c| condition_matches_row(c, r, &self.schema))
13296 {
13297 continue;
13298 }
13299 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
13301 continue;
13302 }
13303 if let Some(cid) = cid {
13304 let mut cells = r
13305 .columns
13306 .get(&cid)
13307 .cloned()
13308 .map(|value| vec![(cid, value)])
13309 .unwrap_or_default();
13310 if let Some(authorization) = authorization {
13311 authorization.security.apply_masks_to_cells(
13312 authorization.table,
13313 &mut cells,
13314 authorization.principal,
13315 );
13316 }
13317 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
13318 passing_vals.push(v);
13319 } } else {
13321 passing_vals.push(0.0); }
13323 }
13324 let m = passing_vals.len();
13325
13326 let (point, half) = match agg {
13327 ApproxAgg::Count => {
13328 let p = m as f64 / s as f64;
13330 let point = n_pop as f64 * p;
13331 let var = if s > 1 {
13332 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
13333 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13334 } else {
13335 0.0
13336 };
13337 (point, z * var.sqrt())
13338 }
13339 ApproxAgg::Sum => {
13340 let y: Vec<f64> = live_sample
13342 .iter()
13343 .map(|r| {
13344 let passes_row = authorized
13345 .as_ref()
13346 .is_none_or(|authorized| authorized.contains(&r.row_id))
13347 && conditions
13348 .iter()
13349 .all(|c| condition_matches_row(c, r, &self.schema))
13350 && index_sets.iter().all(|set| set.contains(r.row_id.0));
13351 if passes_row {
13352 cid.and_then(|cid| {
13353 let mut cells = r
13354 .columns
13355 .get(&cid)
13356 .cloned()
13357 .map(|value| vec![(cid, value)])
13358 .unwrap_or_default();
13359 if let Some(authorization) = authorization {
13360 authorization.security.apply_masks_to_cells(
13361 authorization.table,
13362 &mut cells,
13363 authorization.principal,
13364 );
13365 }
13366 as_f64(cells.first().map(|(_, value)| value))
13367 })
13368 .unwrap_or(0.0)
13369 } else {
13370 0.0
13371 }
13372 })
13373 .collect();
13374 let mean_y = y.iter().sum::<f64>() / s as f64;
13375 let point = n_pop as f64 * mean_y;
13376 let var = if s > 1 {
13377 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
13378 let var_y = ss / (s - 1) as f64;
13379 n_pop as f64 * n_pop as f64 * var_y / s as f64
13380 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13381 } else {
13382 0.0
13383 };
13384 (point, z * var.sqrt())
13385 }
13386 ApproxAgg::Avg => {
13387 if m == 0 {
13388 return Ok(Some(ApproxResult {
13389 point: 0.0,
13390 ci_low: 0.0,
13391 ci_high: 0.0,
13392 n_population: n_pop,
13393 n_sample_live: s,
13394 n_passing: 0,
13395 }));
13396 }
13397 let mean = passing_vals.iter().sum::<f64>() / m as f64;
13398 let half = if m > 1 {
13399 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
13400 let sd = (ss / (m - 1) as f64).sqrt();
13401 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
13402 z * sd / (m as f64).sqrt() * fpc.sqrt()
13403 } else {
13404 0.0
13405 };
13406 (mean, half)
13407 }
13408 };
13409
13410 Ok(Some(ApproxResult {
13411 point,
13412 ci_low: point - half,
13413 ci_high: point + half,
13414 n_population: n_pop,
13415 n_sample_live: s,
13416 n_passing: m,
13417 }))
13418 }
13419
13420 pub fn exact_column_stats(
13428 &self,
13429 _snapshot: Snapshot,
13430 projection: &[u16],
13431 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
13432 if self.ttl.is_some()
13433 || !(self.memtable.is_empty()
13434 && self.mutable_run.is_empty()
13435 && self.run_refs.len() == 1)
13436 {
13437 return Ok(None);
13438 }
13439 let reader = self.open_reader(self.run_refs[0].run_id)?;
13440 if self.live_count != reader.row_count() as u64 {
13441 return Ok(None);
13442 }
13443 let mut out = HashMap::new();
13444 for &cid in projection {
13445 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
13446 Some(c) => c,
13447 None => continue,
13448 };
13449 let Some(stats) = reader.column_page_stats(cid) else {
13451 out.insert(
13452 cid,
13453 ColumnStat {
13454 min: None,
13455 max: None,
13456 null_count: self.live_count,
13457 },
13458 );
13459 continue;
13460 };
13461 let stat = match cdef.ty {
13462 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
13463 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
13464 min: mn.map(Value::Int64),
13465 max: mx.map(Value::Int64),
13466 null_count: n,
13467 })
13468 }
13469 TypeId::Float64 => {
13470 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
13471 min: mn.map(Value::Float64),
13472 max: mx.map(Value::Float64),
13473 null_count: n,
13474 })
13475 }
13476 _ => None,
13477 };
13478 if let Some(s) = stat {
13479 out.insert(cid, s);
13480 }
13481 }
13482 Ok(Some(out))
13483 }
13484
13485 pub fn dir(&self) -> &Path {
13486 &self.dir
13487 }
13488
13489 pub fn schema(&self) -> &Schema {
13490 &self.schema
13491 }
13492
13493 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
13494 self.ann.get(&column_id)
13495 }
13496
13497 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
13498 self.ann.get_mut(&column_id)
13499 }
13500
13501 pub(crate) fn set_catalog_name(&mut self, name: String) {
13502 self.name = name;
13503 }
13504
13505 pub(crate) fn prepare_alter_column(
13506 &mut self,
13507 column_name: &str,
13508 change: &AlterColumn,
13509 ) -> Result<(ColumnDef, Option<Schema>)> {
13510 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
13511 return Err(MongrelError::InvalidArgument(
13512 "ALTER COLUMN requires committing staged writes first".into(),
13513 ));
13514 }
13515 let old = self
13516 .schema
13517 .columns
13518 .iter()
13519 .find(|c| c.name == column_name)
13520 .cloned()
13521 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
13522 let mut next = old.clone();
13523
13524 if let Some(name) = &change.name {
13525 let trimmed = name.trim();
13526 if trimmed.is_empty() {
13527 return Err(MongrelError::InvalidArgument(
13528 "ALTER COLUMN name must not be empty".into(),
13529 ));
13530 }
13531 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
13532 return Err(MongrelError::Schema(format!(
13533 "column {trimmed} already exists"
13534 )));
13535 }
13536 next.name = trimmed.to_string();
13537 }
13538
13539 if let Some(ty) = &change.ty {
13540 next.ty = ty.clone();
13541 }
13542 if let Some(flags) = change.flags {
13543 validate_alter_column_flags(old.flags, flags)?;
13544 next.flags = flags;
13545 }
13546
13547 if let Some(default_change) = &change.default_value {
13548 next.default_value = default_change.clone();
13549 }
13550 if let Some(source_change) = &change.embedding_source {
13551 next.embedding_source = source_change.clone();
13552 }
13553
13554 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
13555 if old.flags.contains(ColumnFlags::NULLABLE)
13556 && !next.flags.contains(ColumnFlags::NULLABLE)
13557 && self.column_has_nulls(old.id)?
13558 {
13559 return Err(MongrelError::InvalidArgument(format!(
13560 "column '{}' contains NULL values",
13561 old.name
13562 )));
13563 }
13564 if next == old {
13565 return Ok((next, None));
13566 }
13567 let mut schema = self.schema.clone();
13568 let index = schema
13569 .columns
13570 .iter()
13571 .position(|column| column.id == next.id)
13572 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
13573 schema.columns[index] = next.clone();
13574 schema.schema_id = schema
13575 .schema_id
13576 .checked_add(1)
13577 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13578 schema.validate_auto_increment()?;
13579 schema.validate_defaults()?;
13580 Ok((next, Some(schema)))
13581 }
13582
13583 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
13584 self.schema = schema;
13585 self.auto_inc = resolve_auto_inc(&self.schema);
13586 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
13587 self.clear_result_cache();
13588 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
13589 }
13590
13591 pub(crate) fn publish_index_schema_change(
13595 &mut self,
13596 schema: Schema,
13597 artifact: SecondaryIndexArtifact,
13598 ) -> Result<()> {
13599 self.apply_altered_schema_prepared(schema);
13600 self.prune_index_maps_to_schema();
13601 match artifact {
13602 SecondaryIndexArtifact::Bitmap(column_id, index) => {
13603 self.bitmap.insert(column_id, index);
13604 }
13605 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
13606 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
13607 }
13608 SecondaryIndexArtifact::Fm(column_id, index) => {
13609 self.fm.insert(column_id, *index);
13610 }
13611 SecondaryIndexArtifact::Ann(column_id, index) => {
13612 self.ann.insert(column_id, *index);
13613 }
13614 SecondaryIndexArtifact::Sparse(column_id, index) => {
13615 self.sparse.insert(column_id, index);
13616 }
13617 SecondaryIndexArtifact::MinHash(column_id, index) => {
13618 self.minhash.insert(column_id, index);
13619 }
13620 }
13621 self.indexes_complete = true;
13622 self.seal_generations();
13623 let view = Arc::new(self.capture_read_generation());
13624 self.published.store(Arc::clone(&view));
13625 checkpoint_current_schema(self)?;
13626 self.invalidate_index_checkpoint();
13629 Ok(())
13630 }
13631
13632 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
13638 self.apply_altered_schema_prepared(schema);
13639 self.prune_index_maps_to_schema();
13640 self.indexes_complete = true;
13641 self.seal_generations();
13642 let view = Arc::new(self.capture_read_generation());
13643 self.published.store(Arc::clone(&view));
13644 checkpoint_current_schema(self)?;
13645 self.invalidate_index_checkpoint();
13647 Ok(())
13648 }
13649
13650 fn prune_index_maps_to_schema(&mut self) {
13651 let keep_bitmap: std::collections::HashSet<u16> = self
13652 .schema
13653 .indexes
13654 .iter()
13655 .filter(|index| index.kind == IndexKind::Bitmap)
13656 .map(|index| index.column_id)
13657 .collect();
13658 let keep_ann: std::collections::HashSet<u16> = self
13659 .schema
13660 .indexes
13661 .iter()
13662 .filter(|index| index.kind == IndexKind::Ann)
13663 .map(|index| index.column_id)
13664 .collect();
13665 let keep_fm: std::collections::HashSet<u16> = self
13666 .schema
13667 .indexes
13668 .iter()
13669 .filter(|index| index.kind == IndexKind::FmIndex)
13670 .map(|index| index.column_id)
13671 .collect();
13672 let keep_sparse: std::collections::HashSet<u16> = self
13673 .schema
13674 .indexes
13675 .iter()
13676 .filter(|index| index.kind == IndexKind::Sparse)
13677 .map(|index| index.column_id)
13678 .collect();
13679 let keep_minhash: std::collections::HashSet<u16> = self
13680 .schema
13681 .indexes
13682 .iter()
13683 .filter(|index| index.kind == IndexKind::MinHash)
13684 .map(|index| index.column_id)
13685 .collect();
13686 let keep_learned: std::collections::HashSet<u16> = self
13687 .schema
13688 .indexes
13689 .iter()
13690 .filter(|index| index.kind == IndexKind::LearnedRange)
13691 .map(|index| index.column_id)
13692 .collect();
13693 self.bitmap
13694 .retain(|column_id, _| keep_bitmap.contains(column_id));
13695 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
13696 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
13697 self.sparse
13698 .retain(|column_id, _| keep_sparse.contains(column_id));
13699 self.minhash
13700 .retain(|column_id, _| keep_minhash.contains(column_id));
13701 {
13702 let learned = Arc::make_mut(&mut self.learned_range);
13703 learned.retain(|column_id, _| keep_learned.contains(column_id));
13704 }
13705 }
13706
13707 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
13708 checkpoint_current_schema(self)
13709 }
13710
13711 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
13712 self.ensure_writable()?;
13713 let previous_schema = self.schema.clone();
13714 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
13715 if let Some(schema) = schema {
13716 self.apply_altered_schema_prepared(schema);
13717 self.checkpoint_standalone_schema_change(previous_schema)?;
13718 }
13719 Ok(column)
13720 }
13721
13722 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
13723 if self.live_count == 0 {
13724 return Ok(false);
13725 }
13726 let snap = self.snapshot();
13727 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
13728 Ok(columns
13729 .first()
13730 .map(|(_, col)| col.null_count(col.len()) != 0)
13731 .unwrap_or(true))
13732 }
13733
13734 fn has_stored_versions(&self) -> bool {
13735 !self.memtable.is_empty()
13736 || !self.mutable_run.is_empty()
13737 || self.run_refs.iter().any(|r| r.row_count > 0)
13738 || !self.retiring.is_empty()
13739 }
13740
13741 pub fn add_column(
13746 &mut self,
13747 name: &str,
13748 ty: TypeId,
13749 flags: ColumnFlags,
13750 default_value: Option<crate::schema::DefaultExpr>,
13751 ) -> Result<u16> {
13752 self.add_column_with_id(name, ty, flags, default_value, None)
13753 }
13754
13755 pub fn add_column_with_id(
13756 &mut self,
13757 name: &str,
13758 ty: TypeId,
13759 flags: ColumnFlags,
13760 default_value: Option<crate::schema::DefaultExpr>,
13761 requested_id: Option<u16>,
13762 ) -> Result<u16> {
13763 self.ensure_writable()?;
13764 let previous_schema = self.schema.clone();
13765 let (column, schema) =
13766 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
13767 self.apply_altered_schema_prepared(schema);
13768 self.checkpoint_standalone_schema_change(previous_schema)?;
13769 Ok(column.id)
13770 }
13771
13772 pub(crate) fn prepare_add_column(
13773 &mut self,
13774 name: &str,
13775 ty: TypeId,
13776 flags: ColumnFlags,
13777 default_value: Option<crate::schema::DefaultExpr>,
13778 requested_id: Option<u16>,
13779 ) -> Result<(ColumnDef, Schema)> {
13780 self.ensure_writable()?;
13781 if self.schema.columns.iter().any(|c| c.name == name) {
13782 return Err(MongrelError::Schema(format!(
13783 "column {name} already exists"
13784 )));
13785 }
13786 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
13787 if self.schema.columns.iter().any(|c| c.id == id) {
13788 return Err(MongrelError::Schema(format!(
13789 "column id {id} already exists"
13790 )));
13791 }
13792 id
13793 } else {
13794 self.schema
13795 .columns
13796 .iter()
13797 .map(|c| c.id)
13798 .max()
13799 .unwrap_or(0)
13800 .checked_add(1)
13801 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
13802 };
13803 let column = ColumnDef {
13804 id,
13805 name: name.to_string(),
13806 ty,
13807 flags,
13808 default_value,
13809 embedding_source: None,
13810 };
13811 let mut next_schema = self.schema.clone();
13812 next_schema.columns.push(column.clone());
13813 next_schema.schema_id = next_schema
13814 .schema_id
13815 .checked_add(1)
13816 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13817 next_schema.validate_auto_increment()?;
13818 next_schema.validate_defaults()?;
13819 Ok((column, next_schema))
13820 }
13821
13822 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
13831 self.ensure_writable()?;
13832 let cid = self
13833 .schema
13834 .columns
13835 .iter()
13836 .find(|c| c.name == column_name)
13837 .map(|c| c.id)
13838 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
13839 let ty = self
13840 .schema
13841 .columns
13842 .iter()
13843 .find(|c| c.id == cid)
13844 .map(|c| c.ty.clone())
13845 .unwrap_or(TypeId::Int64);
13846 if !matches!(
13847 ty,
13848 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
13849 ) {
13850 return Err(MongrelError::Schema(format!(
13851 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
13852 )));
13853 }
13854 if self
13855 .schema
13856 .indexes
13857 .iter()
13858 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
13859 {
13860 return Ok(()); }
13862 let previous_schema = self.schema.clone();
13863 let previous_learned_range = Arc::clone(&self.learned_range);
13864 let mut next_schema = previous_schema.clone();
13865 next_schema.indexes.push(IndexDef {
13866 name: format!("{}_learned_range", column_name),
13867 column_id: cid,
13868 kind: IndexKind::LearnedRange,
13869 predicate: None,
13870 options: Default::default(),
13871 });
13872 next_schema.schema_id = next_schema
13873 .schema_id
13874 .checked_add(1)
13875 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13876 self.apply_altered_schema_prepared(next_schema);
13877 if let Err(error) = self.build_learned_ranges() {
13878 self.apply_altered_schema_prepared(previous_schema);
13879 self.learned_range = previous_learned_range;
13880 return Err(error);
13881 }
13882 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
13883 if !matches!(
13884 &error,
13885 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
13886 ) {
13887 self.learned_range = previous_learned_range;
13888 }
13889 return Err(error);
13890 }
13891 Ok(())
13892 }
13893
13894 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
13895 let mut schema_published = false;
13896 let schema_result = match self._root_guard.as_deref() {
13897 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
13898 schema_published = true;
13899 }),
13900 None => write_schema_with_after(&self.dir, &self.schema, || {
13901 schema_published = true;
13902 }),
13903 };
13904 if schema_result.is_err() && !schema_published {
13905 self.apply_altered_schema_prepared(previous_schema);
13906 return schema_result;
13907 }
13908
13909 let manifest_result = self.persist_manifest(self.current_epoch());
13910 match (schema_result, manifest_result) {
13911 (_, Ok(())) => Ok(()),
13912 (Ok(()), Err(error)) => {
13913 self.poison_after_maintenance_publish_failure();
13914 Err(MongrelError::DurableCommit {
13915 epoch: self.current_epoch().0,
13916 message: format!(
13917 "schema is durable but matching manifest publication failed: {error}"
13918 ),
13919 })
13920 }
13921 (Err(schema_error), Err(manifest_error)) => {
13922 self.poison_after_maintenance_publish_failure();
13923 Err(MongrelError::CommitOutcomeUnknown {
13924 epoch: self.current_epoch().0,
13925 message: format!(
13926 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
13927 ),
13928 })
13929 }
13930 }
13931 }
13932
13933 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
13936 self.sync_byte_threshold = threshold;
13937 if let WalSink::Private(w) = &mut self.wal {
13938 w.set_sync_byte_threshold(threshold);
13939 }
13940 }
13941
13942 pub fn page_cache_flush(&self) {
13946 self.page_cache.flush_to_disk();
13947 }
13948
13949 pub fn page_cache_len(&self) -> usize {
13951 self.page_cache.len()
13952 }
13953
13954 pub fn decoded_cache_len(&self) -> usize {
13957 self.decoded_cache.len()
13958 }
13959
13960 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
13963 self.memtable.drain_sorted()
13964 }
13965
13966 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
13967 self.runs_dir().join(format!("r-{run_id}.sr"))
13968 }
13969
13970 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
13971 match self.runs_root.as_deref() {
13972 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
13973 None => Ok(None),
13974 }
13975 }
13976
13977 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
13978 match self.runs_root.as_deref() {
13979 Some(root) => Ok(Some(root.create_regular_new(name)?)),
13980 None => Ok(None),
13981 }
13982 }
13983
13984 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
13985 match self.runs_root.as_deref() {
13986 Some(root) => match root.remove_file(name) {
13987 Ok(()) => Ok(()),
13988 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
13989 Err(error) => Err(error.into()),
13990 },
13991 None => match std::fs::remove_file(self.runs_dir().join(name)) {
13992 Ok(()) => Ok(()),
13993 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
13994 Err(error) => Err(error.into()),
13995 },
13996 }
13997 }
13998
13999 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
14000 match self.runs_root.as_deref() {
14001 Some(root) => root
14002 .rename_file_new(source, destination)
14003 .map_err(Into::into),
14004 None => crate::durable_file::rename(
14005 &self.runs_dir().join(source),
14006 &self.runs_dir().join(destination),
14007 )
14008 .map_err(Into::into),
14009 }
14010 }
14011
14012 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
14013 self.run_refs.iter().map(|run| run.run_id)
14014 }
14015
14016 pub(crate) fn table_dir(&self) -> &Path {
14017 &self.dir
14018 }
14019
14020 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
14021 &self.schema
14022 }
14023
14024 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
14025 let id = self.next_run_id;
14026 self.next_run_id = self
14027 .next_run_id
14028 .checked_add(1)
14029 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
14030 Ok(id)
14031 }
14032
14033 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
14034 self.run_refs.push(run_ref);
14035 }
14036
14037 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
14047 self.retiring.push(crate::manifest::RetiredRun {
14048 run_id,
14049 retire_epoch,
14050 });
14051 }
14052
14053 pub(crate) fn reap_retiring(
14057 &mut self,
14058 min_active: Epoch,
14059 backup_pinned: &std::collections::HashSet<u128>,
14060 ) -> Result<usize> {
14061 if self.retiring.is_empty() {
14062 return Ok(0);
14063 }
14064 let mut reaped = 0;
14065 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
14066 for r in std::mem::take(&mut self.retiring) {
14072 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
14073 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
14074 reaped += 1;
14075 } else {
14076 kept.push(r);
14077 }
14078 }
14079 self.retiring = kept;
14080 if reaped > 0 {
14081 self.persist_manifest(self.current_epoch())?;
14082 }
14083 Ok(reaped)
14084 }
14085
14086 pub(crate) fn has_reapable_retiring(
14087 &self,
14088 min_active: Epoch,
14089 backup_pinned: &std::collections::HashSet<u128>,
14090 ) -> bool {
14091 self.retiring
14092 .iter()
14093 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
14094 }
14095
14096 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
14097 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
14098 return false;
14099 }
14100 self.live_count = self.live_count.saturating_add(run_ref.row_count);
14101 self.run_refs.push(run_ref);
14102 self.indexes_complete = false;
14103 true
14104 }
14105
14106 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
14107 self.kek.as_ref()
14108 }
14109
14110 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
14111 let mut reader = match self.runs_root.as_deref() {
14112 Some(root) => RunReader::open_file_with_cache(
14113 root.open_regular(format!("r-{run_id}.sr"))?,
14114 self.schema.clone(),
14115 self.kek.clone(),
14116 Some(self.page_cache.clone()),
14117 Some(self.decoded_cache.clone()),
14118 self.table_id,
14119 Some(&self.verified_runs),
14120 None,
14121 )?,
14122 None => RunReader::open_with_cache(
14123 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
14124 self.schema.clone(),
14125 self.kek.clone(),
14126 Some(self.page_cache.clone()),
14127 Some(self.decoded_cache.clone()),
14128 self.table_id,
14129 Some(&self.verified_runs),
14130 )?,
14131 };
14132 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
14136 reader.set_uniform_epoch(Epoch(rr.epoch_created));
14137 }
14138 Ok(reader)
14139 }
14140
14141 pub(crate) fn run_refs(&self) -> &[RunRef] {
14142 &self.run_refs
14143 }
14144
14145 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
14146 self.retiring.iter().map(|run| run.run_id)
14147 }
14148
14149 pub(crate) fn runs_dir(&self) -> PathBuf {
14150 self.runs_root
14151 .as_deref()
14152 .and_then(|root| root.io_path().ok())
14153 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
14154 }
14155
14156 pub(crate) fn wal_dir(&self) -> PathBuf {
14157 self.dir.join(WAL_DIR)
14158 }
14159
14160 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
14161 self.run_refs = refs;
14162 self.refresh_run_row_id_ranges();
14163 }
14164
14165 pub(crate) fn run_row_id_range(&self, run_id: u128) -> Option<(u64, u64)> {
14169 self.run_row_id_ranges.get(&run_id).copied()
14170 }
14171
14172 pub(crate) fn refresh_run_row_id_ranges(&mut self) {
14175 let mut ranges = HashMap::with_capacity(self.run_refs.len());
14176 for rr in &self.run_refs {
14177 if let Ok(reader) = self.open_reader(rr.run_id) {
14178 let h = reader.header();
14179 ranges.insert(rr.run_id, (h.min_row_id, h.max_row_id));
14180 }
14181 }
14182 self.run_row_id_ranges = ranges;
14183 }
14184
14185 pub(crate) fn compaction_zstd_level(&self) -> i32 {
14186 self.compaction_zstd_level
14187 }
14188
14189 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
14190 self.kek.clone()
14191 }
14192
14193 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
14197 self.kek.as_ref().map(|k| k.derive_idx_key())
14198 }
14199
14200 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
14204 self.kek.as_ref().map(|k| *k.derive_meta_key())
14205 }
14206
14207 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
14210 self.column_keys
14211 .iter()
14212 .map(|(&id, &(_, scheme))| (id, scheme))
14213 .collect()
14214 }
14215
14216 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
14221 self.tokenize_value_enc(column_id, v)
14222 }
14223
14224 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
14225 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
14226 let (key, scheme) = self.column_keys.get(&column_id)?;
14227 let token: Vec<u8> = match (*scheme, v) {
14228 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
14229 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
14230 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
14231 _ => hmac_token(key, &v.encode_key()).to_vec(),
14232 };
14233 Some(Value::Bytes(token))
14234 }
14235
14236 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
14238 self.index_lookup_key_bytes(column_id, &v.encode_key())
14239 }
14240
14241 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
14244 {
14245 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
14246 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
14247 if *scheme == SCHEME_HMAC_EQ {
14248 return hmac_token(key, encoded).to_vec();
14249 }
14250 }
14251 }
14252 let _ = column_id;
14253 encoded.to_vec()
14254 }
14255}
14256
14257fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
14258 let columnar::NativeColumn::Int64 { data, validity } = col else {
14259 return false;
14260 };
14261 if data.len() < n || !columnar::all_non_null(validity, n) {
14262 return false;
14263 }
14264 data.iter()
14265 .take(n)
14266 .zip(data.iter().skip(1))
14267 .all(|(a, b)| a < b)
14268}
14269
14270#[derive(Debug, Clone)]
14274pub struct ColumnStat {
14275 pub min: Option<Value>,
14276 pub max: Option<Value>,
14277 pub null_count: u64,
14278}
14279
14280#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14282pub enum NativeAgg {
14283 Count,
14284 Sum,
14285 Min,
14286 Max,
14287 Avg,
14288}
14289
14290#[derive(Debug, Clone, PartialEq)]
14292pub enum NativeAggResult {
14293 Count(u64),
14294 Int(i64),
14295 Float(f64),
14296 Null,
14298}
14299
14300#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14302pub enum ApproxAgg {
14303 Count,
14304 Sum,
14305 Avg,
14306}
14307
14308#[derive(Debug, Clone)]
14312pub struct ApproxResult {
14313 pub point: f64,
14315 pub ci_low: f64,
14317 pub ci_high: f64,
14319 pub n_population: u64,
14321 pub n_sample_live: usize,
14323 pub n_passing: usize,
14325}
14326
14327#[derive(Debug, Clone, PartialEq)]
14332pub enum AggState {
14333 Count(u64),
14335 SumI {
14337 sum: i128,
14338 count: u64,
14339 },
14340 SumF {
14342 sum: f64,
14343 count: u64,
14344 },
14345 AvgI {
14347 sum: i128,
14348 count: u64,
14349 },
14350 AvgF {
14352 sum: f64,
14353 count: u64,
14354 },
14355 MinI(i64),
14357 MaxI(i64),
14358 MinF(f64),
14360 MaxF(f64),
14361 Empty,
14363}
14364
14365impl AggState {
14366 pub fn merge(self, other: AggState) -> AggState {
14368 use AggState::*;
14369 match (self, other) {
14370 (Empty, x) | (x, Empty) => x,
14371 (Count(a), Count(b)) => Count(a + b),
14372 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
14373 sum: sa + sb,
14374 count: ca + cb,
14375 },
14376 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
14377 sum: sa + sb,
14378 count: ca + cb,
14379 },
14380 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
14381 sum: sa + sb,
14382 count: ca + cb,
14383 },
14384 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
14385 sum: sa + sb,
14386 count: ca + cb,
14387 },
14388 (MinI(a), MinI(b)) => MinI(a.min(b)),
14389 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
14390 (MinF(a), MinF(b)) => MinF(a.min(b)),
14391 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
14392 _ => Empty, }
14394 }
14395
14396 pub fn point(&self) -> Option<f64> {
14398 match self {
14399 AggState::Count(n) => Some(*n as f64),
14400 AggState::SumI { sum, .. } => Some(*sum as f64),
14401 AggState::SumF { sum, .. } => Some(*sum),
14402 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
14403 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
14404 AggState::MinI(n) => Some(*n as f64),
14405 AggState::MaxI(n) => Some(*n as f64),
14406 AggState::MinF(n) => Some(*n),
14407 AggState::MaxF(n) => Some(*n),
14408 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
14409 }
14410 }
14411
14412 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
14416 let is_float = matches!(ty, Some(TypeId::Float64));
14417 match (agg, result) {
14418 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
14419 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
14420 sum: x as i128,
14421 count: 1, },
14423 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
14424 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
14425 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
14426 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
14427 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
14428 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
14429 (NativeAgg::Count, _) => AggState::Empty,
14430 (_, NativeAggResult::Null) => AggState::Empty,
14431 _ => {
14432 let _ = is_float;
14433 AggState::Empty
14434 }
14435 }
14436 }
14437}
14438
14439#[derive(Debug, Clone)]
14442pub struct CachedAgg {
14443 pub state: AggState,
14444 pub watermark: u64,
14445 pub epoch: u64,
14446}
14447
14448#[derive(Debug, Clone)]
14450pub struct IncrementalAggResult {
14451 pub state: AggState,
14453 pub incremental: bool,
14456 pub delta_rows: u64,
14458}
14459
14460fn agg_state_from_rows(
14464 rows: &[Row],
14465 conditions: &[crate::query::Condition],
14466 index_sets: &[RowIdSet],
14467 column: Option<u16>,
14468 agg: NativeAgg,
14469 schema: &Schema,
14470 control: Option<&crate::ExecutionControl>,
14471) -> Result<AggState> {
14472 let mut count: u64 = 0;
14473 let mut sum_i: i128 = 0;
14474 let mut sum_f: f64 = 0.0;
14475 let mut mn_i: i64 = i64::MAX;
14476 let mut mx_i: i64 = i64::MIN;
14477 let mut mn_f: f64 = f64::INFINITY;
14478 let mut mx_f: f64 = f64::NEG_INFINITY;
14479 let mut saw_int = false;
14480 let mut saw_float = false;
14481 for (index, r) in rows.iter().enumerate() {
14482 execution_checkpoint(control, index)?;
14483 if !conditions
14484 .iter()
14485 .all(|c| condition_matches_row(c, r, schema))
14486 {
14487 continue;
14488 }
14489 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
14490 continue;
14491 }
14492 match agg {
14493 NativeAgg::Count => match column {
14494 None => count += 1,
14496 Some(cid) => match r.columns.get(&cid) {
14499 None | Some(Value::Null) => {}
14500 Some(_) => count += 1,
14501 },
14502 },
14503 _ => match column.and_then(|cid| r.columns.get(&cid)) {
14504 Some(Value::Int64(n)) => {
14505 count += 1;
14506 sum_i += *n as i128;
14507 mn_i = mn_i.min(*n);
14508 mx_i = mx_i.max(*n);
14509 saw_int = true;
14510 }
14511 Some(Value::Float64(f)) => {
14512 count += 1;
14513 sum_f += f;
14514 mn_f = mn_f.min(*f);
14515 mx_f = mx_f.max(*f);
14516 saw_float = true;
14517 }
14518 _ => {}
14519 },
14520 }
14521 }
14522 Ok(match agg {
14523 NativeAgg::Count => {
14524 if count == 0 {
14525 AggState::Empty
14526 } else {
14527 AggState::Count(count)
14528 }
14529 }
14530 NativeAgg::Sum => {
14531 if count == 0 {
14532 AggState::Empty
14533 } else if saw_int {
14534 AggState::SumI { sum: sum_i, count }
14535 } else {
14536 AggState::SumF { sum: sum_f, count }
14537 }
14538 }
14539 NativeAgg::Avg => {
14540 if count == 0 {
14541 AggState::Empty
14542 } else if saw_int {
14543 AggState::AvgI { sum: sum_i, count }
14544 } else {
14545 AggState::AvgF { sum: sum_f, count }
14546 }
14547 }
14548 NativeAgg::Min => {
14549 if !saw_int && !saw_float {
14550 AggState::Empty
14551 } else if saw_int {
14552 AggState::MinI(mn_i)
14553 } else {
14554 AggState::MinF(mn_f)
14555 }
14556 }
14557 NativeAgg::Max => {
14558 if !saw_int && !saw_float {
14559 AggState::Empty
14560 } else if saw_int {
14561 AggState::MaxI(mx_i)
14562 } else {
14563 AggState::MaxF(mx_f)
14564 }
14565 }
14566 })
14567}
14568
14569fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
14573 use crate::query::Condition;
14574 match c {
14575 Condition::Pk(key) => match schema.primary_key() {
14576 Some(pk) => row
14577 .columns
14578 .get(&pk.id)
14579 .map(|v| v.encode_key() == *key)
14580 .unwrap_or(false),
14581 None => false,
14582 },
14583 Condition::BitmapEq { column_id, value } => row
14584 .columns
14585 .get(column_id)
14586 .map(|v| v.encode_key() == *value)
14587 .unwrap_or(false),
14588 Condition::BitmapIn { column_id, values } => {
14589 let key = row.columns.get(column_id).map(|v| v.encode_key());
14590 match key {
14591 Some(k) => values.contains(&k),
14592 None => false,
14593 }
14594 }
14595 Condition::BytesPrefix { column_id, prefix } => row
14596 .columns
14597 .get(column_id)
14598 .map(|v| v.encode_key().starts_with(prefix))
14599 .unwrap_or(false),
14600 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
14601 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
14602 _ => false,
14603 },
14604 Condition::RangeF64 {
14605 column_id,
14606 lo,
14607 lo_inclusive,
14608 hi,
14609 hi_inclusive,
14610 } => match row.columns.get(column_id) {
14611 Some(Value::Float64(n)) => {
14612 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
14613 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
14614 lo_ok && hi_ok
14615 }
14616 _ => false,
14617 },
14618 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
14619 Some(Value::Bytes(b)) => {
14620 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
14621 }
14622 _ => false,
14623 },
14624 Condition::FmContainsAll {
14625 column_id,
14626 patterns,
14627 } => match row.columns.get(column_id) {
14628 Some(Value::Bytes(b)) => patterns
14629 .iter()
14630 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
14631 _ => false,
14632 },
14633 Condition::Ann { .. }
14634 | Condition::SparseMatch { .. }
14635 | Condition::MinHashSimilar { .. } => true,
14636 Condition::IsNull { column_id } => {
14637 matches!(row.columns.get(column_id), Some(Value::Null) | None)
14638 }
14639 Condition::IsNotNull { column_id } => {
14640 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
14641 }
14642 }
14643}
14644
14645fn as_f64(v: Option<&Value>) -> Option<f64> {
14647 match v {
14648 Some(Value::Int64(n)) => Some(*n as f64),
14649 Some(Value::Float64(f)) => Some(*f),
14650 _ => None,
14651 }
14652}
14653
14654fn accumulate_int(
14658 cursor: &mut dyn crate::cursor::Cursor,
14659 control: Option<&crate::ExecutionControl>,
14660) -> Result<(u64, i128, i64, i64)> {
14661 let mut count: u64 = 0;
14662 let mut sum: i128 = 0;
14663 let mut mn: i64 = i64::MAX;
14664 let mut mx: i64 = i64::MIN;
14665 while let Some(cols) = cursor.next_batch()? {
14666 execution_checkpoint(control, 0)?;
14667 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
14668 if crate::columnar::all_non_null(validity, data.len()) {
14669 count += data.len() as u64;
14671 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
14672 execution_checkpoint(control, chunk_index * 1024)?;
14673 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
14674 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
14675 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
14676 }
14677 } else {
14678 for (i, &v) in data.iter().enumerate() {
14679 execution_checkpoint(control, i)?;
14680 if crate::columnar::validity_bit(validity, i) {
14681 count += 1;
14682 sum += v as i128;
14683 mn = mn.min(v);
14684 mx = mx.max(v);
14685 }
14686 }
14687 }
14688 }
14689 }
14690 Ok((count, sum, mn, mx))
14691}
14692
14693fn accumulate_float(
14695 cursor: &mut dyn crate::cursor::Cursor,
14696 control: Option<&crate::ExecutionControl>,
14697) -> Result<(u64, f64, f64, f64)> {
14698 let mut count: u64 = 0;
14699 let mut sum: f64 = 0.0;
14700 let mut mn: f64 = f64::INFINITY;
14701 let mut mx: f64 = f64::NEG_INFINITY;
14702 while let Some(cols) = cursor.next_batch()? {
14703 execution_checkpoint(control, 0)?;
14704 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
14705 if crate::columnar::all_non_null(validity, data.len()) {
14706 count += data.len() as u64;
14707 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
14708 execution_checkpoint(control, chunk_index * 1024)?;
14709 sum += chunk.iter().sum::<f64>();
14710 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
14711 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
14712 }
14713 } else {
14714 for (i, &v) in data.iter().enumerate() {
14715 execution_checkpoint(control, i)?;
14716 if crate::columnar::validity_bit(validity, i) {
14717 count += 1;
14718 sum += v;
14719 mn = mn.min(v);
14720 mx = mx.max(v);
14721 }
14722 }
14723 }
14724 }
14725 }
14726 Ok((count, sum, mn, mx))
14727}
14728
14729#[inline]
14730fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
14731 if index.is_multiple_of(256) {
14732 control
14733 .map(crate::ExecutionControl::checkpoint)
14734 .transpose()?;
14735 }
14736 Ok(())
14737}
14738
14739fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
14740 if count == 0 && !matches!(agg, NativeAgg::Count) {
14741 return NativeAggResult::Null;
14742 }
14743 match agg {
14744 NativeAgg::Count => NativeAggResult::Count(count),
14745 NativeAgg::Sum => match sum.try_into() {
14748 Ok(v) => NativeAggResult::Int(v),
14749 Err(_) => NativeAggResult::Null,
14750 },
14751 NativeAgg::Min => NativeAggResult::Int(mn),
14752 NativeAgg::Max => NativeAggResult::Int(mx),
14753 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
14754 }
14755}
14756
14757fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
14758 if count == 0 && !matches!(agg, NativeAgg::Count) {
14759 return NativeAggResult::Null;
14760 }
14761 match agg {
14762 NativeAgg::Count => NativeAggResult::Count(count),
14763 NativeAgg::Sum => NativeAggResult::Float(sum),
14764 NativeAgg::Min => NativeAggResult::Float(mn),
14765 NativeAgg::Max => NativeAggResult::Float(mx),
14766 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
14767 }
14768}
14769
14770fn agg_int(
14773 stats: &[crate::page::PageStat],
14774 decode: fn(Option<&[u8]>) -> Option<i64>,
14775) -> Option<(Option<i64>, Option<i64>, u64)> {
14776 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
14777 let mut any = false;
14778 for s in stats {
14779 if let Some(v) = decode(s.min.as_deref()) {
14780 mn = mn.min(v);
14781 any = true;
14782 }
14783 if let Some(v) = decode(s.max.as_deref()) {
14784 mx = mx.max(v);
14785 any = true;
14786 }
14787 nulls += s.null_count;
14788 }
14789 any.then_some((Some(mn), Some(mx), nulls))
14790}
14791
14792fn agg_float(
14794 stats: &[crate::page::PageStat],
14795 decode: fn(Option<&[u8]>) -> Option<f64>,
14796) -> Option<(Option<f64>, Option<f64>, u64)> {
14797 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
14798 let mut any = false;
14799 for s in stats {
14800 if let Some(v) = decode(s.min.as_deref()) {
14801 mn = mn.min(v);
14802 any = true;
14803 }
14804 if let Some(v) = decode(s.max.as_deref()) {
14805 mx = mx.max(v);
14806 any = true;
14807 }
14808 nulls += s.null_count;
14809 }
14810 any.then_some((Some(mn), Some(mx), nulls))
14811}
14812
14813type SecondaryIndexes = (
14815 HashMap<u16, BitmapIndex>,
14816 HashMap<u16, AnnIndex>,
14817 HashMap<u16, FmIndex>,
14818 HashMap<u16, SparseIndex>,
14819 HashMap<u16, MinHashIndex>,
14820);
14821
14822fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
14823 let mut bitmap = HashMap::new();
14824 let mut ann = HashMap::new();
14825 let mut fm = HashMap::new();
14826 let mut sparse = HashMap::new();
14827 let mut minhash = HashMap::new();
14828 for idef in &schema.indexes {
14829 match idef.kind {
14830 IndexKind::Bitmap => {
14831 bitmap.insert(idef.column_id, BitmapIndex::new());
14832 }
14833 IndexKind::Ann => {
14834 let dim = schema
14835 .columns
14836 .iter()
14837 .find(|c| c.id == idef.column_id)
14838 .and_then(|c| match c.ty {
14839 TypeId::Embedding { dim } => Some(dim as usize),
14840 _ => None,
14841 })
14842 .unwrap_or(0);
14843 let options = idef.options.ann.clone().unwrap_or_default();
14844 ann.insert(
14845 idef.column_id,
14846 AnnIndex::with_full_options(
14847 dim,
14848 options.m,
14849 options.ef_construction,
14850 options.ef_search,
14851 &options,
14852 ),
14853 );
14854 }
14855 IndexKind::FmIndex => {
14856 fm.insert(idef.column_id, FmIndex::new());
14857 }
14858 IndexKind::Sparse => {
14859 sparse.insert(idef.column_id, SparseIndex::new());
14860 }
14861 IndexKind::MinHash => {
14862 let options = idef.options.minhash.clone().unwrap_or_default();
14863 minhash.insert(
14864 idef.column_id,
14865 MinHashIndex::with_options(options.permutations, options.bands),
14866 );
14867 }
14868 _ => {}
14869 }
14870 }
14871 (bitmap, ann, fm, sparse, minhash)
14872}
14873
14874const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
14875 | ColumnFlags::AUTO_INCREMENT
14876 | ColumnFlags::ENCRYPTED
14877 | ColumnFlags::ENCRYPTED_INDEXABLE
14878 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
14879
14880fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
14881 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
14882 return Err(MongrelError::Schema(
14883 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
14884 ));
14885 }
14886 Ok(())
14887}
14888
14889fn validate_alter_column_type(
14890 schema: &Schema,
14891 old: &ColumnDef,
14892 next: &ColumnDef,
14893 has_stored_versions: bool,
14894) -> Result<()> {
14895 if old.ty == next.ty {
14896 return Ok(());
14897 }
14898 if schema.indexes.iter().any(|i| i.column_id == old.id) {
14899 return Err(MongrelError::Schema(format!(
14900 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
14901 old.name
14902 )));
14903 }
14904 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
14905 return Ok(());
14906 }
14907 Err(MongrelError::Schema(format!(
14908 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
14909 old.ty, next.ty
14910 )))
14911}
14912
14913fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
14914 matches!(
14915 (old, new),
14916 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
14917 )
14918}
14919
14920fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
14926 let mut prev: Option<i64> = None;
14927 for r in rows {
14928 match r.columns.get(&pk_id) {
14929 Some(Value::Int64(v)) => {
14930 if prev.is_some_and(|p| p >= *v) {
14931 return false;
14932 }
14933 prev = Some(*v);
14934 }
14935 _ => return false,
14936 }
14937 }
14938 true
14939}
14940
14941#[allow(clippy::too_many_arguments)]
14942fn index_into(
14943 schema: &Schema,
14944 row: &Row,
14945 hot: &mut HotIndex,
14946 bitmap: &mut HashMap<u16, BitmapIndex>,
14947 ann: &mut HashMap<u16, AnnIndex>,
14948 fm: &mut HashMap<u16, FmIndex>,
14949 sparse: &mut HashMap<u16, SparseIndex>,
14950 minhash: &mut HashMap<u16, MinHashIndex>,
14951) {
14952 if row.row_id.0 == 70 || row.row_id.0 == 117 {
14953 eprintln!("DEBUG index_into rid={}", row.row_id.0);
14954 }
14955 for idef in &schema.indexes {
14956 let Some(val) = row.columns.get(&idef.column_id) else {
14957 continue;
14958 };
14959 match idef.kind {
14960 IndexKind::Bitmap => {
14961 if let Some(b) = bitmap.get_mut(&idef.column_id) {
14962 b.insert(val.encode_key(), row.row_id);
14963 }
14964 }
14965 IndexKind::Ann => {
14966 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
14967 if row.row_id.0 == 70 || row.row_id.0 == 117 {
14968 eprintln!("DEBUG index_into ANN index rid={}", row.row_id.0);
14969 }
14970 if let Some(meta) = val.generated_embedding_metadata() {
14971 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
14973 continue;
14974 }
14975 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
14976 .is_err()
14977 {
14978 continue;
14979 }
14980 }
14981 a.insert_validated(v, row.row_id);
14982 } else if row.row_id.0 == 70 || row.row_id.0 == 117 {
14983 eprintln!(
14984 "DEBUG index_into ANN skip rid={} ann_has={} val_is_embed={}",
14985 row.row_id.0,
14986 ann.get_mut(&idef.column_id).is_some(),
14987 val.as_embedding().is_some(),
14988 );
14989 }
14990 }
14991 IndexKind::FmIndex => {
14992 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
14993 f.insert(b.clone(), row.row_id);
14994 }
14995 }
14996 IndexKind::Sparse => {
14997 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
14998 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
15001 s.insert(&terms, row.row_id);
15002 }
15003 }
15004 }
15005 IndexKind::MinHash => {
15006 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
15007 let tokens = crate::index::token_hashes_from_bytes(b);
15010 mh.insert(&tokens, row.row_id);
15011 }
15012 }
15013 _ => {}
15014 }
15015 }
15016 if let Some(pk_col) = schema.primary_key() {
15017 if let Some(pk_val) = row.columns.get(&pk_col.id) {
15018 hot.insert(pk_val.encode_key(), row.row_id);
15019 }
15020 }
15021}
15022
15023#[allow(clippy::too_many_arguments)]
15026fn index_into_single(
15027 idef: &IndexDef,
15028 _schema: &Schema,
15029 row: &Row,
15030 _hot: &mut HotIndex,
15031 bitmap: &mut HashMap<u16, BitmapIndex>,
15032 ann: &mut HashMap<u16, AnnIndex>,
15033 fm: &mut HashMap<u16, FmIndex>,
15034 sparse: &mut HashMap<u16, SparseIndex>,
15035 minhash: &mut HashMap<u16, MinHashIndex>,
15036) {
15037 if row.row_id.0 == 70 || row.row_id.0 == 117 {
15038 eprintln!(
15039 "DEBUG index_into_single rid={} kind={:?} col_id={} has_col={}",
15040 row.row_id.0,
15041 idef.kind,
15042 idef.column_id,
15043 row.columns.contains_key(&idef.column_id),
15044 );
15045 }
15046 let Some(val) = row.columns.get(&idef.column_id) else {
15047 return;
15048 };
15049 match idef.kind {
15050 IndexKind::Bitmap => {
15051 if let Some(b) = bitmap.get_mut(&idef.column_id) {
15052 b.insert(val.encode_key(), row.row_id);
15053 }
15054 }
15055 IndexKind::Ann => {
15056 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
15057 if row.row_id.0 == 70 || row.row_id.0 == 117 {
15058 eprintln!("DEBUG index_into ANN rid={}", row.row_id.0);
15059 }
15060 if let Some(meta) = val.generated_embedding_metadata() {
15061 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
15063 return;
15064 }
15065 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
15066 .is_err()
15067 {
15068 return;
15069 }
15070 }
15071 a.insert_validated(v, row.row_id);
15072 }
15073 }
15074 IndexKind::FmIndex => {
15075 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
15076 f.insert(b.clone(), row.row_id);
15077 }
15078 }
15079 IndexKind::Sparse => {
15080 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
15081 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
15082 s.insert(&terms, row.row_id);
15083 }
15084 }
15085 }
15086 IndexKind::MinHash => {
15087 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
15088 let tokens = crate::index::token_hashes_from_bytes(b);
15089 mh.insert(&tokens, row.row_id);
15090 }
15091 }
15092 _ => {}
15093 }
15094}
15095
15096fn eval_partial_predicate(
15102 pred: &str,
15103 columns_map: &HashMap<u16, &Value>,
15104 name_to_id: &HashMap<&str, u16>,
15105) -> bool {
15106 let lower = pred.trim().to_ascii_lowercase();
15107 if let Some(rest) = lower.strip_suffix(" is not null") {
15109 let col_name = rest.trim();
15110 if let Some(col_id) = name_to_id.get(col_name) {
15111 return columns_map
15112 .get(col_id)
15113 .is_some_and(|v| !matches!(v, Value::Null));
15114 }
15115 }
15116 if let Some(rest) = lower.strip_suffix(" is null") {
15118 let col_name = rest.trim();
15119 if let Some(col_id) = name_to_id.get(col_name) {
15120 return columns_map
15121 .get(col_id)
15122 .is_none_or(|v| matches!(v, Value::Null));
15123 }
15124 }
15125 true
15128}
15129
15130#[allow(dead_code)]
15136fn bulk_index_key(
15137 column_keys: &HashMap<u16, ([u8; 32], u8)>,
15138 column_id: u16,
15139 ty: TypeId,
15140 col: &columnar::NativeColumn,
15141 i: usize,
15142) -> Option<Vec<u8>> {
15143 let encoded = columnar::encode_key_native(ty, col, i)?;
15144 {
15145 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
15146 if let Some((key, scheme)) = column_keys.get(&column_id) {
15147 return Some(match (*scheme, col) {
15148 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
15149 (_, columnar::NativeColumn::Int64 { data, .. }) => {
15150 ope_token_i64(key, data[i]).to_vec()
15151 }
15152 (_, columnar::NativeColumn::Float64 { data, .. }) => {
15153 ope_token_f64(key, data[i]).to_vec()
15154 }
15155 _ => hmac_token(key, &encoded).to_vec(),
15156 });
15157 }
15158 }
15159 Some(encoded)
15160}
15161
15162pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
15163 write_schema_with_after(dir, schema, || {})
15164}
15165
15166pub(crate) fn write_schema_durable(
15167 root: &crate::durable_file::DurableRoot,
15168 schema: &Schema,
15169) -> Result<()> {
15170 write_schema_durable_with_after(root, schema, || {})
15171}
15172
15173fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
15174where
15175 F: FnOnce(),
15176{
15177 let json = serde_json::to_string_pretty(schema)
15178 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
15179 crate::durable_file::write_atomic_with_after(
15180 &dir.join(SCHEMA_FILENAME),
15181 json.as_bytes(),
15182 after_publish,
15183 )?;
15184 Ok(())
15185}
15186
15187fn write_schema_durable_with_after<F>(
15188 root: &crate::durable_file::DurableRoot,
15189 schema: &Schema,
15190 after_publish: F,
15191) -> Result<()>
15192where
15193 F: FnOnce(),
15194{
15195 let json = serde_json::to_string_pretty(schema)
15196 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
15197 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
15198 Ok(())
15199}
15200
15201fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
15202 let mut schema_published = false;
15203 let schema_result = match table._root_guard.as_deref() {
15204 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
15205 schema_published = true;
15206 }),
15207 None => write_schema_with_after(&table.dir, &table.schema, || {
15208 schema_published = true;
15209 }),
15210 };
15211 if schema_result.is_err() && !schema_published {
15212 return schema_result;
15213 }
15214 match table.persist_manifest(table.current_epoch()) {
15215 Ok(()) => Ok(()),
15216 Err(manifest_error) => Err(match schema_result {
15217 Ok(()) => manifest_error,
15218 Err(schema_error) => MongrelError::Other(format!(
15219 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
15220 )),
15221 }),
15222 }
15223}
15224
15225fn read_schema(dir: &Path) -> Result<Schema> {
15226 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
15227 read_schema_file(file)
15228}
15229
15230fn read_schema_file(file: std::fs::File) -> Result<Schema> {
15231 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
15232 use std::io::Read;
15233
15234 let length = file.metadata()?.len();
15235 if length > MAX_SCHEMA_BYTES {
15236 return Err(MongrelError::ResourceLimitExceeded {
15237 resource: "schema bytes",
15238 requested: usize::try_from(length).unwrap_or(usize::MAX),
15239 limit: MAX_SCHEMA_BYTES as usize,
15240 });
15241 }
15242 let mut bytes = Vec::with_capacity(length as usize);
15243 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
15244 if bytes.len() as u64 != length {
15245 return Err(MongrelError::Schema(
15246 "schema length changed while reading".into(),
15247 ));
15248 }
15249 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
15250}
15251
15252fn preflight_standalone_open(
15253 dir: &Path,
15254 runs_root: Option<&crate::durable_file::DurableRoot>,
15255 idx_root: Option<&crate::durable_file::DurableRoot>,
15256 manifest: &Manifest,
15257 schema: &Schema,
15258 records: &[crate::wal::Record],
15259 kek: Option<Arc<Kek>>,
15260) -> Result<()> {
15261 crate::wal::validate_shared_transaction_framing(records)?;
15262 if manifest.schema_id > schema.schema_id
15263 || manifest.flushed_epoch > manifest.current_epoch
15264 || manifest.global_idx_epoch > manifest.current_epoch
15265 || manifest.next_row_id == u64::MAX
15266 || manifest.auto_inc_next < 0
15267 || manifest.auto_inc_next == i64::MAX
15268 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
15269 {
15270 return Err(MongrelError::InvalidArgument(
15271 "manifest counters or schema identity are invalid".into(),
15272 ));
15273 }
15274 let mut run_ids = HashSet::new();
15275 let mut maximum_row_id = None::<u64>;
15276 for run in &manifest.runs {
15277 if run.run_id >= u64::MAX as u128
15278 || !run_ids.insert(run.run_id)
15279 || run.epoch_created > manifest.current_epoch
15280 {
15281 return Err(MongrelError::InvalidArgument(
15282 "manifest contains an invalid or duplicate active run".into(),
15283 ));
15284 }
15285 let mut reader = match runs_root {
15286 Some(root) => RunReader::open_file(
15287 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
15288 schema.clone(),
15289 kek.clone(),
15290 )?,
15291 None => RunReader::open(
15292 dir.join(RUNS_DIR)
15293 .join(format!("r-{}.sr", run.run_id as u64)),
15294 schema.clone(),
15295 kek.clone(),
15296 )?,
15297 };
15298 let header = reader.header();
15299 if header.run_id != run.run_id
15300 || header.level != run.level
15301 || header.row_count != run.row_count
15302 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
15303 || header.is_uniform_epoch() && header.epoch_created != 0
15304 || header.schema_id > schema.schema_id
15305 {
15306 return Err(MongrelError::InvalidArgument(format!(
15307 "run {} differs from its manifest",
15308 run.run_id
15309 )));
15310 }
15311 if header.row_count != 0 {
15312 maximum_row_id = Some(
15313 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
15314 );
15315 }
15316 reader.validate_all_pages()?;
15317 }
15318 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
15319 return Err(MongrelError::InvalidArgument(
15320 "manifest next_row_id does not advance beyond persisted rows".into(),
15321 ));
15322 }
15323 for run in &manifest.retiring {
15324 if run.run_id >= u64::MAX as u128
15325 || run.retire_epoch > manifest.current_epoch
15326 || !run_ids.insert(run.run_id)
15327 {
15328 return Err(MongrelError::InvalidArgument(
15329 "manifest contains an invalid or duplicate retired run".into(),
15330 ));
15331 }
15332 }
15333 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
15334 match idx_root {
15335 Some(root) => {
15336 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
15337 }
15338 None => {
15339 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
15340 }
15341 }
15342
15343 let committed = records
15344 .iter()
15345 .filter_map(|record| match record.op {
15346 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
15347 _ => None,
15348 })
15349 .collect::<HashMap<_, _>>();
15350 for record in records {
15351 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
15352 continue;
15353 };
15354 match &record.op {
15355 Op::Put { table_id, rows } => {
15356 if *table_id != manifest.table_id {
15357 return Err(MongrelError::CorruptWal {
15358 offset: record.seq.0,
15359 reason: format!(
15360 "private WAL record references table {table_id}, expected {}",
15361 manifest.table_id
15362 ),
15363 });
15364 }
15365 let rows: Vec<Row> =
15366 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
15367 offset: record.seq.0,
15368 reason: format!("committed Put payload could not be decoded: {error}"),
15369 })?;
15370 for row in rows {
15371 if row.deleted || row.row_id.0 == u64::MAX {
15372 return Err(MongrelError::CorruptWal {
15373 offset: record.seq.0,
15374 reason: "committed Put contains an invalid row identity".into(),
15375 });
15376 }
15377 let cells = row.columns.into_iter().collect::<Vec<_>>();
15378 schema
15379 .validate_values(&cells)
15380 .map_err(|error| MongrelError::CorruptWal {
15381 offset: record.seq.0,
15382 reason: format!("committed Put violates table schema: {error}"),
15383 })?;
15384 if schema.auto_increment_column().is_some_and(|column| {
15385 matches!(
15386 cells.iter().find(|(id, _)| *id == column.id),
15387 Some((_, Value::Int64(value))) if *value == i64::MAX
15388 )
15389 }) {
15390 return Err(MongrelError::CorruptWal {
15391 offset: record.seq.0,
15392 reason: "committed Put exhausts AUTO_INCREMENT".into(),
15393 });
15394 }
15395 }
15396 }
15397 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
15398 if *table_id != manifest.table_id =>
15399 {
15400 return Err(MongrelError::CorruptWal {
15401 offset: record.seq.0,
15402 reason: format!(
15403 "private WAL record references table {table_id}, expected {}",
15404 manifest.table_id
15405 ),
15406 });
15407 }
15408 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
15409 return Err(MongrelError::CorruptWal {
15410 offset: record.seq.0,
15411 reason: "private WAL contains shared spilled-run metadata".into(),
15412 });
15413 }
15414 _ => {}
15415 }
15416 }
15417 Ok(())
15418}
15419
15420fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
15421 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
15422}
15423
15424fn wal_segment_number(path: &Path) -> Option<u64> {
15425 path.file_stem()
15426 .and_then(|stem| stem.to_str())
15427 .and_then(|stem| stem.strip_prefix("seg-"))
15428 .and_then(|number| number.parse().ok())
15429}
15430
15431fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
15432 let n = list_wal_numbers(wal_dir)?;
15433 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
15434}
15435
15436fn next_wal_number(wal_dir: &Path) -> Result<u32> {
15437 list_wal_numbers(wal_dir)?
15438 .map(|maximum| {
15439 maximum
15440 .checked_add(1)
15441 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
15442 })
15443 .unwrap_or(Ok(0))
15444}
15445
15446fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
15447 let mut max_n = None;
15448 let entries = match std::fs::read_dir(wal_dir) {
15449 Ok(entries) => entries,
15450 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
15451 Err(error) => return Err(error.into()),
15452 };
15453 for entry in entries {
15454 let entry = entry?;
15455 let fname = entry.file_name();
15456 let Some(s) = fname.to_str() else {
15457 continue;
15458 };
15459 let Some(stripped) = s.strip_prefix("seg-") else {
15460 continue;
15461 };
15462 let Some(number) = stripped.strip_suffix(".wal") else {
15463 return Err(MongrelError::CorruptWal {
15464 offset: 0,
15465 reason: format!("malformed WAL segment name {s:?}"),
15466 });
15467 };
15468 let n = number
15469 .parse::<u32>()
15470 .map_err(|_| MongrelError::CorruptWal {
15471 offset: 0,
15472 reason: format!("malformed WAL segment name {s:?}"),
15473 })?;
15474 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
15475 return Err(MongrelError::CorruptWal {
15476 offset: n as u64,
15477 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
15478 });
15479 }
15480 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
15481 }
15482 Ok(max_n)
15483}
15484
15485#[cfg(test)]
15510mod rem001_multi_run_hlc_authority {
15511 use super::*;
15512 use crate::epoch::{Epoch, Snapshot, VersionStamp};
15513 use crate::memtable::{Row, Value};
15514 use crate::query::{Condition, Query};
15515 use crate::schema::{ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
15516 use mongreldb_types::hlc::HlcTimestamp;
15517 use tempfile::tempdir;
15518
15519 fn hlc(physical_micros: u64) -> HlcTimestamp {
15520 HlcTimestamp {
15521 physical_micros,
15522 logical: 0,
15523 node_tiebreaker: 1,
15524 }
15525 }
15526
15527 fn pk_schema() -> Schema {
15528 Schema {
15529 schema_id: 1,
15530 columns: vec![
15531 ColumnDef {
15532 id: 1,
15533 name: "id".into(),
15534 ty: TypeId::Int64,
15535 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
15536 default_value: None,
15537 embedding_source: None,
15538 },
15539 ColumnDef {
15540 id: 2,
15541 name: "name".into(),
15542 ty: TypeId::Bytes,
15543 flags: ColumnFlags::empty(),
15544 default_value: None,
15545 embedding_source: None,
15546 },
15547 ],
15548 indexes: Vec::new(),
15549 colocation: vec![],
15550 constraints: Default::default(),
15551 clustered: false,
15552 }
15553 }
15554
15555 fn pk_schema_with_bitmap() -> Schema {
15556 Schema {
15557 schema_id: 1,
15558 columns: vec![
15559 ColumnDef {
15560 id: 1,
15561 name: "id".into(),
15562 ty: TypeId::Int64,
15563 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
15564 default_value: None,
15565 embedding_source: None,
15566 },
15567 ColumnDef {
15568 id: 2,
15569 name: "name".into(),
15570 ty: TypeId::Bytes,
15571 flags: ColumnFlags::empty(),
15572 default_value: None,
15573 embedding_source: None,
15574 },
15575 ],
15576 indexes: vec![IndexDef {
15579 name: "name_bitmap".into(),
15580 column_id: 2,
15581 kind: IndexKind::Bitmap,
15582 predicate: None,
15583 options: Default::default(),
15584 }],
15585 colocation: vec![],
15586 constraints: Default::default(),
15587 clustered: false,
15588 }
15589 }
15590
15591 fn encode_put(table_id: u64, row: &Row) -> Vec<u8> {
15592 use crate::database::StagedTxnWrite;
15593 bincode::serialize(&StagedTxnWrite::Put {
15594 table_id,
15595 rows: bincode::serialize(&vec![row.clone()]).expect("encode row"),
15596 })
15597 .expect("encode payload")
15598 }
15599
15600 fn build_inverted_database(dir: &std::path::Path, schema: Schema) -> crate::Database {
15614 let db = crate::Database::create(dir).unwrap();
15615 db.create_table("t", schema).unwrap();
15616 let table_id = db.table_id("t").unwrap();
15617 let rid = RowId(7);
15618
15619 let newer = Row::new_with_hlc(rid, Epoch(1), hlc(500))
15620 .with_column(1, Value::Int64(7))
15621 .with_column(2, Value::Bytes(b"newer-by-HLC".to_vec()));
15622 let older = Row::new_with_hlc(rid, Epoch(2), hlc(400))
15623 .with_column(1, Value::Int64(7))
15624 .with_column(2, Value::Bytes(b"older-by-HLC".to_vec()));
15625
15626 db.apply_staged_txn_writes(1, &[encode_put(table_id, &newer)], hlc(500))
15627 .expect("first apply");
15628 {
15629 let table = db.table("t").unwrap();
15630 table.lock().force_flush().expect("flush run A");
15631 }
15632
15633 db.apply_staged_txn_writes(2, &[encode_put(table_id, &older)], hlc(400))
15634 .expect("second apply");
15635 {
15636 let table = db.table("t").unwrap();
15637 table.lock().force_flush().expect("flush run B");
15638 }
15639
15640 db
15641 }
15642
15643 fn build_inverted_ttl_database(dir: &std::path::Path) -> crate::Database {
15648 let db = crate::Database::create(dir).unwrap();
15649 db.create_table("t", ttl_schema()).unwrap();
15650 let table_id = db.table_id("t").unwrap();
15651 let rid = RowId(7);
15652 let read_time_ns: i64 = 1_000_000_000;
15653 let newer = Row::new_with_hlc(rid, Epoch(1), hlc(500))
15654 .with_column(1, Value::Int64(7))
15655 .with_column(2, Value::Bytes(b"newer-by-HLC".to_vec()))
15656 .with_column(3, Value::Int64(read_time_ns + 60 * 1_000_000_000));
15657 let older = Row::new_with_hlc(rid, Epoch(2), hlc(400))
15658 .with_column(1, Value::Int64(7))
15659 .with_column(2, Value::Bytes(b"older-by-HLC".to_vec()))
15660 .with_column(3, Value::Int64(read_time_ns - 60 * 1_000_000_000));
15661
15662 db.apply_staged_txn_writes(1, &[encode_put(table_id, &newer)], hlc(500))
15663 .expect("first apply");
15664 {
15665 let table = db.table("t").unwrap();
15666 table.lock().force_flush().expect("flush run A");
15667 }
15668 db.apply_staged_txn_writes(2, &[encode_put(table_id, &older)], hlc(400))
15669 .expect("second apply");
15670 {
15671 let table = db.table("t").unwrap();
15672 table.lock().force_flush().expect("flush run B");
15673 }
15674
15675 db
15676 }
15677
15678 fn ttl_schema() -> Schema {
15679 Schema {
15680 schema_id: 1,
15681 columns: vec![
15682 ColumnDef {
15683 id: 1,
15684 name: "id".into(),
15685 ty: TypeId::Int64,
15686 flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
15687 default_value: None,
15688 embedding_source: None,
15689 },
15690 ColumnDef {
15691 id: 2,
15692 name: "name".into(),
15693 ty: TypeId::Bytes,
15694 flags: ColumnFlags::empty(),
15695 default_value: None,
15696 embedding_source: None,
15697 },
15698 ColumnDef {
15699 id: 3,
15700 name: "ttl_ts".into(),
15701 ty: TypeId::Int64,
15702 flags: ColumnFlags::empty(),
15703 default_value: None,
15704 embedding_source: None,
15705 },
15706 ],
15707 indexes: Vec::new(),
15708 colocation: vec![],
15709 constraints: Default::default(),
15710 clustered: false,
15711 }
15712 }
15713
15714 #[test]
15715 fn multi_run_rows_for_rids_uses_hlc_winner_when_epoch_is_inverted() {
15716 let dir = tempdir().unwrap();
15717 let db = build_inverted_database(dir.path(), pk_schema());
15718 let table = db.table("t").unwrap();
15719 let table = table.lock();
15720 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15724 let rows = table
15725 .rows_for_rids(&[7], snap)
15726 .expect("rows_for_rids across two runs");
15727 assert_eq!(rows.len(), 1, "one row materialised across two runs");
15728 assert_eq!(
15729 rows[0].columns.get(&2),
15730 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
15731 "multi-run rows_for_rids must honor HLC authority"
15732 );
15733 }
15734
15735 #[test]
15736 fn multi_run_projected_column_uses_hlc_winner_when_epoch_is_inverted() {
15737 let dir = tempdir().unwrap();
15738 let db = build_inverted_database(dir.path(), pk_schema());
15739 let table = db.table("t").unwrap();
15740 let table = table.lock();
15741 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15742 let projected = table
15743 .values_for_rids_at(&[7], 2, snap, i64::MAX)
15744 .expect("values_for_rids_at across two runs");
15745 assert_eq!(projected.len(), 1);
15746 assert_eq!(
15747 projected[0].1,
15748 Value::Bytes(b"newer-by-HLC".to_vec()),
15749 "multi-run projection must pick the HLC-newer column"
15750 );
15751 }
15752
15753 #[test]
15754 fn ann_candidate_eligibility_uses_hlc_winner_across_runs() {
15755 let dir = tempdir().unwrap();
15756 let db = build_inverted_database(dir.path(), pk_schema());
15757 let table = db.table("t").unwrap();
15758 let table = table.lock();
15759 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15760 let eligible = table
15763 .eligible_candidate_ids(&[RowId(7)], 2, snap, None)
15764 .expect("eligibility across runs");
15765 assert!(
15766 eligible.contains(&RowId(7)),
15767 "HLC-newer rid must remain eligible across runs"
15768 );
15769 }
15770
15771 #[test]
15772 fn ttl_inverted_run_is_admitted_via_values_for_rids_at() {
15773 let dir = tempdir().unwrap();
15774 let db = build_inverted_ttl_database(dir.path());
15775 let table = db.table("t").unwrap();
15776 let table = table.lock();
15777 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15778 let now_ns: i64 = 1_000_000_000;
15779 let projected = table
15783 .values_for_rids_at(&[7], 2, snap, now_ns)
15784 .expect("ttl-aware projection");
15785 assert_eq!(projected.len(), 1, "TTL must admit the HLC-newer row");
15786 assert_eq!(
15787 projected[0].1,
15788 Value::Bytes(b"newer-by-HLC".to_vec()),
15789 "TTL must be read from the HLC-newer run; older run's timestamp is expired"
15790 );
15791 }
15792
15793 #[test]
15794 fn hot_candidate_inspection_uses_full_snapshot_for_run_versions() {
15795 let dir = tempdir().unwrap();
15800 let db = build_inverted_database(dir.path(), pk_schema());
15801 let table = db.table("t").unwrap();
15802 let table = table.lock();
15803 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15804 let row = table
15805 .get(RowId(7), snap)
15806 .expect("Table::get across two runs");
15807 assert_eq!(
15808 row.columns.get(&2),
15809 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
15810 "Table::get must surface the HLC-newer winner"
15811 );
15812 }
15813
15814 #[test]
15815 fn learned_range_rebuild_indexes_only_hlc_visible_winners() {
15816 let dir = tempdir().unwrap();
15823 let db = build_inverted_database(dir.path(), pk_schema());
15824 let handle = db.table("t").unwrap();
15825 let mut table = handle.lock();
15826 table.rebuild_indexes_from_runs().expect("rebuild");
15827 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15828 let rows = table
15829 .rows_for_rids(&[7], snap)
15830 .expect("post-rebuild rows_for_rids");
15831 assert_eq!(rows.len(), 1);
15832 assert_eq!(
15833 rows[0].columns.get(&2),
15834 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
15835 "post-rebuild row must be the HLC-newer winner"
15836 );
15837 }
15838
15839 #[test]
15840 fn rebuild_indexes_uses_global_hlc_winners_not_run_iteration_order() {
15841 let dir = tempdir().unwrap();
15854 let db = build_inverted_database(dir.path(), pk_schema_with_bitmap());
15855 let handle = db.table("t").unwrap();
15856 let mut table = handle.lock();
15857 table.rebuild_indexes_from_runs().expect("rebuild");
15858 let snap = Snapshot::at_hlc(Epoch(10), hlc(600));
15859 let row = table.get(RowId(7), snap).expect("post-rebuild Table::get");
15861 assert_eq!(
15862 row.columns.get(&2),
15863 Some(&Value::Bytes(b"newer-by-HLC".to_vec())),
15864 "PK lookup must resolve to the HLC-newer winner"
15865 );
15866 let query_newer = Query::new().and(Condition::BitmapEq {
15870 column_id: 2,
15871 value: b"newer-by-HLC".to_vec(),
15872 });
15873 let rows_newer = table.query(&query_newer).expect("bitmap eq newer");
15874 let rids_newer: Vec<u64> = rows_newer.iter().map(|r| r.row_id.0).collect();
15875 assert!(
15876 rids_newer.contains(&7),
15877 "bitmap must include the HLC-newer value (got {rids_newer:?})"
15878 );
15879 let query_older = Query::new().and(Condition::BitmapEq {
15880 column_id: 2,
15881 value: b"older-by-HLC".to_vec(),
15882 });
15883 let rows_older = table.query(&query_older).expect("bitmap eq older");
15884 let rids_older: Vec<u64> = rows_older.iter().map(|r| r.row_id.0).collect();
15885 assert!(
15886 !rids_older.contains(&7),
15887 "bitmap must NOT include the HLC-older value (got {rids_older:?})"
15888 );
15889 }
15890
15891 #[test]
15892 fn version_stamp_is_newer_than_matches_snapshot_rule() {
15893 let early = hlc(100);
15896 let late = hlc(200);
15897 let newer = VersionStamp {
15898 epoch: Epoch(9),
15899 commit_ts: Some(late),
15900 };
15901 let older = VersionStamp {
15902 epoch: Epoch(50),
15903 commit_ts: Some(early),
15904 };
15905 assert!(newer.is_newer_than(older));
15906 assert!(!older.is_newer_than(newer));
15907 let a = VersionStamp {
15908 epoch: Epoch(5),
15909 commit_ts: None,
15910 };
15911 let b = VersionStamp {
15912 epoch: Epoch(4),
15913 commit_ts: None,
15914 };
15915 assert!(a.is_newer_than(b));
15916 }
15917}