1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, EpochGuard, MaintenanceReceipt, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, IndexGeneration, MinHashIndex,
18 SparseIndex,
19};
20use crate::manifest::{self, Manifest, RunRef, TtlPolicy};
21use crate::memtable::{Memtable, MemtableVisibleVersionCursor, Row, Value};
22use crate::mutable_run::{MutableRun, MutableRunVisibleVersionCursor};
23use crate::row_id_set::RowIdSet;
24use crate::rowid::{RowId, RowIdAllocator};
25use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
26use crate::sorted_run::{RunReader, RunVisibleVersion, RunVisibleVersionCursor, RunWriter};
27use crate::txn::{GroupCommit, OwnedRow};
28use crate::wal::{Op, SharedWal, Wal};
29use crate::{MongrelError, Result};
30use arc_swap::ArcSwap;
31use mongreldb_types::hlc::HlcTimestamp;
32use std::borrow::Cow;
33use std::cmp::Reverse;
34use std::collections::{BTreeMap, BTreeSet, BinaryHeap, HashMap, HashSet};
35use std::path::{Path, PathBuf};
36use std::sync::atomic::AtomicBool;
37use std::sync::Arc;
38use zeroize::Zeroizing;
39
40pub const WAL_DIR: &str = "_wal";
41pub const RUNS_DIR: &str = "_runs";
42pub const CACHE_DIR: &str = "_cache";
43pub const META_DIR: &str = "_meta";
44pub const RCACHE_DIR: &str = "_rcache";
45pub const KEYS_FILENAME: &str = "keys";
46pub const SCHEMA_FILENAME: &str = "schema.json";
47
48fn derive_next_run_id(
49 dir: &Path,
50 runs_root: Option<&crate::durable_file::DurableRoot>,
51 active: &[RunRef],
52 retiring: &[crate::manifest::RetiredRun],
53) -> Result<u64> {
54 let mut maximum = 0_u64;
55 for run_id in active
56 .iter()
57 .map(|run| run.run_id)
58 .chain(retiring.iter().map(|run| run.run_id))
59 {
60 let run_id = u64::try_from(run_id)
61 .map_err(|_| MongrelError::Full("run-id namespace exhausted".into()))?;
62 maximum = maximum.max(run_id);
63 }
64 let names = match runs_root {
65 Some(root) => root.list_regular_files(".")?,
66 None => std::fs::read_dir(dir.join(RUNS_DIR))?
67 .map(|entry| entry.map(|entry| entry.file_name()))
68 .collect::<std::io::Result<Vec<_>>>()?,
69 };
70 for name in names {
71 let Some(name) = name.to_str() else {
72 continue;
73 };
74 let Some(digits) = name
75 .strip_prefix("r-")
76 .and_then(|name| name.strip_suffix(".sr"))
77 else {
78 continue;
79 };
80 let Ok(run_id) = digits.parse::<u64>() else {
81 continue;
82 };
83 if name == format!("r-{run_id}.sr") {
84 maximum = maximum.max(run_id);
85 }
86 }
87 maximum
88 .checked_add(1)
89 .map(|next| next.max(1))
90 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))
91}
92
93enum ControlledVisibleCandidate<'a> {
94 Memory(Row),
95 Memtable(RowId, Epoch, Cow<'a, Row>),
101 MutableRun(RowId, Epoch, &'a Row),
103 Run(RunVisibleVersion),
104}
105
106impl<'a> ControlledVisibleCandidate<'a> {
107 fn row_id(&self) -> RowId {
108 match self {
109 Self::Memory(row) => row.row_id,
110 Self::Memtable(rid, _, _) => *rid,
111 Self::MutableRun(rid, _, _) => *rid,
112 Self::Run(version) => version.row_id,
113 }
114 }
115
116 fn committed_epoch(&self) -> Epoch {
117 match self {
118 Self::Memory(row) => row.committed_epoch,
119 Self::Memtable(_, epoch, _) => *epoch,
120 Self::MutableRun(_, epoch, _) => *epoch,
121 Self::Run(version) => version.committed_epoch,
122 }
123 }
124
125 fn deleted(&self) -> bool {
126 match self {
127 Self::Memory(row) => row.deleted,
128 Self::Memtable(_, _, row) => row.deleted,
129 Self::MutableRun(_, _, row) => row.deleted,
130 Self::Run(version) => version.deleted,
131 }
132 }
133
134 fn commit_ts(&self) -> Option<mongreldb_types::hlc::HlcTimestamp> {
135 match self {
136 Self::Memory(row) => row.commit_ts,
137 Self::Memtable(_, _, row) => row.commit_ts,
138 Self::MutableRun(_, _, row) => row.commit_ts,
139 Self::Run(version) => version.commit_ts,
142 }
143 }
144}
145
146enum ControlledVisibleCursor<'a> {
147 Memory(std::vec::IntoIter<Row>),
149 MemoryStreaming {
153 active: std::vec::IntoIter<Row>,
154 rest: std::collections::btree_map::IntoValues<RowId, Row>,
155 batch_cap: usize,
156 #[allow(dead_code)]
158 peak_active: usize,
159 #[allow(dead_code)]
161 total: usize,
162 },
163 Memtable(MemtableVisibleVersionCursor<'a>),
167 MutableRun(MutableRunVisibleVersionCursor<'a>),
169 Run(Box<RunVisibleVersionCursor>),
170 #[cfg(test)]
171 Synthetic {
172 next: u64,
173 end: u64,
174 },
175}
176
177#[allow(dead_code)] const CONTROLLED_HOT_BATCH: usize = 256;
180
181struct ControlledVisibleSource<'a> {
182 cursor: ControlledVisibleCursor<'a>,
183 current: Option<ControlledVisibleCandidate<'a>>,
184}
185
186impl<'a> ControlledVisibleSource<'a> {
187 #[cfg(test)]
189 fn memory(rows: Vec<Row>) -> Self {
190 Self {
191 cursor: ControlledVisibleCursor::Memory(rows.into_iter()),
192 current: None,
193 }
194 }
195
196 #[allow(dead_code)] fn memory_from_map(map: BTreeMap<RowId, Row>) -> Self {
201 let total = map.len();
202 let mut values = map.into_values();
203 if total > CONTROLLED_HOT_BATCH {
204 let active: Vec<Row> = values.by_ref().take(CONTROLLED_HOT_BATCH).collect();
205 let peak = active.len();
206 Self {
207 cursor: ControlledVisibleCursor::MemoryStreaming {
208 active: active.into_iter(),
209 rest: values,
210 batch_cap: CONTROLLED_HOT_BATCH,
211 peak_active: peak,
212 total,
213 },
214 current: None,
215 }
216 } else {
217 let active: Vec<Row> = values.collect();
218 Self {
219 cursor: ControlledVisibleCursor::Memory(active.into_iter()),
220 current: None,
221 }
222 }
223 }
224
225 #[allow(dead_code)] fn memtable_cursor(cursor: MemtableVisibleVersionCursor<'a>) -> Self {
230 Self {
231 cursor: ControlledVisibleCursor::Memtable(cursor),
232 current: None,
233 }
234 }
235
236 #[allow(dead_code)] fn mutable_run_cursor(cursor: MutableRunVisibleVersionCursor<'a>) -> Self {
241 Self {
242 cursor: ControlledVisibleCursor::MutableRun(cursor),
243 current: None,
244 }
245 }
246
247 fn run(cursor: RunVisibleVersionCursor) -> Self {
248 Self {
249 cursor: ControlledVisibleCursor::Run(Box::new(cursor)),
250 current: None,
251 }
252 }
253
254 #[cfg(test)]
255 fn synthetic(end: u64) -> Self {
256 Self {
257 cursor: ControlledVisibleCursor::Synthetic { next: 1, end },
258 current: None,
259 }
260 }
261
262 fn advance(&mut self, control: &crate::ExecutionControl) -> Result<()> {
263 self.current = match &mut self.cursor {
264 ControlledVisibleCursor::Memory(rows) => {
265 rows.next().map(ControlledVisibleCandidate::Memory)
266 }
267 ControlledVisibleCursor::MemoryStreaming {
268 active,
269 rest,
270 batch_cap,
271 peak_active,
272 total: _,
273 } => {
274 if let Some(row) = active.next() {
275 Some(ControlledVisibleCandidate::Memory(row))
276 } else {
277 control.checkpoint()?;
278 let next_batch: Vec<Row> = rest.by_ref().take(*batch_cap).collect();
279 if next_batch.is_empty() {
280 None
281 } else {
282 *peak_active = (*peak_active).max(next_batch.len());
283 *active = next_batch.into_iter();
284 active.next().map(ControlledVisibleCandidate::Memory)
285 }
286 }
287 }
288 ControlledVisibleCursor::Memtable(iter) => iter
289 .next()
290 .map(|(rid, epoch, row)| ControlledVisibleCandidate::Memtable(rid, epoch, row)),
291 ControlledVisibleCursor::MutableRun(iter) => iter
292 .next()
293 .map(|(rid, epoch, row)| ControlledVisibleCandidate::MutableRun(rid, epoch, row)),
294 ControlledVisibleCursor::Run(cursor) => cursor
295 .next_visible_version(control)?
296 .map(ControlledVisibleCandidate::Run),
297 #[cfg(test)]
298 ControlledVisibleCursor::Synthetic { next, end } => {
299 if *next > *end {
300 None
301 } else {
302 let row = Row::new(RowId(*next), Epoch(1));
303 *next += 1;
304 Some(ControlledVisibleCandidate::Memory(row))
305 }
306 }
307 };
308 Ok(())
309 }
310
311 fn pop(&mut self, control: &crate::ExecutionControl) -> Result<ControlledVisibleCandidate<'a>> {
312 let current = self.current.take().ok_or_else(|| {
313 MongrelError::Other("controlled visible source was not primed".into())
314 })?;
315 self.advance(control)?;
316 Ok(current)
317 }
318
319 fn materialize(
320 &mut self,
321 candidate: ControlledVisibleCandidate<'a>,
322 control: &crate::ExecutionControl,
323 ) -> Result<Row> {
324 match candidate {
325 ControlledVisibleCandidate::Memory(row) => Ok(row),
326 ControlledVisibleCandidate::Memtable(_, _, row) => Ok(row.into_owned()),
327 ControlledVisibleCandidate::MutableRun(_, _, row) => Ok(row.clone()),
328 ControlledVisibleCandidate::Run(version) => match &mut self.cursor {
329 ControlledVisibleCursor::Run(cursor) => cursor.materialize(version, control),
330 _ => Err(MongrelError::Other(
331 "run candidate escaped its controlled cursor".into(),
332 )),
333 },
334 }
335 }
336
337 #[cfg(test)]
338 fn is_streaming_memory(&self) -> bool {
339 matches!(self.cursor, ControlledVisibleCursor::MemoryStreaming { .. })
340 }
341
342 #[cfg(test)]
343 fn streaming_peak_active(&self) -> Option<usize> {
344 match &self.cursor {
345 ControlledVisibleCursor::MemoryStreaming { peak_active, .. } => Some(*peak_active),
346 _ => None,
347 }
348 }
349
350 #[cfg(test)]
351 fn streaming_total(&self) -> Option<usize> {
352 match &self.cursor {
353 ControlledVisibleCursor::MemoryStreaming { total, .. } => Some(*total),
354 _ => None,
355 }
356 }
357}
358
359fn merge_controlled_visible_sources<'a>(
360 sources: &mut [ControlledVisibleSource<'a>],
361 control: &crate::ExecutionControl,
362 mut expired: impl FnMut(&Row) -> bool,
363 mut visit: impl FnMut(Row) -> Result<()>,
364) -> Result<()> {
365 let mut heap = BinaryHeap::new();
366 for (source_index, source) in sources.iter_mut().enumerate() {
367 source.advance(control)?;
368 if let Some(candidate) = &source.current {
369 heap.push(Reverse((candidate.row_id(), source_index)));
370 }
371 }
372 let mut merged = 0_usize;
373 while let Some(Reverse((row_id, source_index))) = heap.pop() {
374 if merged.is_multiple_of(256) {
375 control.checkpoint()?;
376 }
377 merged += 1;
378 let mut best_source = source_index;
379 let mut best = sources[source_index].pop(control)?;
380 if let Some(next) = &sources[source_index].current {
381 heap.push(Reverse((next.row_id(), source_index)));
382 }
383 while heap
384 .peek()
385 .is_some_and(|Reverse((candidate, _))| *candidate == row_id)
386 {
387 control.checkpoint()?;
388 let Some(Reverse((_, source_index))) = heap.pop() else {
389 break;
390 };
391 let candidate = sources[source_index].pop(control)?;
392 if Snapshot::version_is_newer(
396 candidate.committed_epoch(),
397 candidate.commit_ts(),
398 best.committed_epoch(),
399 best.commit_ts(),
400 ) {
401 best = candidate;
402 best_source = source_index;
403 }
404 if let Some(next) = &sources[source_index].current {
405 heap.push(Reverse((next.row_id(), source_index)));
406 }
407 }
408 if best.deleted() {
409 continue;
410 }
411 let row = sources[best_source].materialize(best, control)?;
412 if !expired(&row) {
413 visit(row)?;
414 }
415 }
416 control.checkpoint()
417}
418
419#[cfg(test)]
420mod controlled_visible_cursor_tests {
421 use super::*;
422
423 #[test]
424 fn streams_more_than_one_million_rows_without_a_source_cap() {
425 let control = crate::ExecutionControl::new(None);
426 let mut sources = vec![ControlledVisibleSource::synthetic(1_000_001)];
427 let mut count = 0_u64;
428 let mut last = 0_u64;
429 merge_controlled_visible_sources(
430 &mut sources,
431 &control,
432 |_| false,
433 |row| {
434 count += 1;
435 assert!(row.row_id.0 > last);
436 last = row.row_id.0;
437 Ok(())
438 },
439 )
440 .unwrap();
441 assert_eq!(count, 1_000_001);
442 assert_eq!(last, 1_000_001);
443 }
444
445 #[test]
446 fn merge_orders_rows_and_honors_newest_tombstones() {
447 let control = crate::ExecutionControl::new(None);
448 let older = vec![
449 Row::new(RowId(1), Epoch(1)),
450 Row::new(RowId(2), Epoch(1)).with_column(1, Value::Int64(20)),
451 Row::new(RowId(4), Epoch(1)),
452 ];
453 let mut deleted = Row::new(RowId(1), Epoch(2));
454 deleted.deleted = true;
455 let newer = vec![
456 deleted,
457 Row::new(RowId(2), Epoch(2)).with_column(1, Value::Int64(22)),
458 Row::new(RowId(3), Epoch(2)),
459 ];
460 let mut sources = vec![
461 ControlledVisibleSource::memory(older),
462 ControlledVisibleSource::memory(newer),
463 ];
464 let mut rows = Vec::new();
465 merge_controlled_visible_sources(
466 &mut sources,
467 &control,
468 |_| false,
469 |row| {
470 rows.push(row);
471 Ok(())
472 },
473 )
474 .unwrap();
475 assert_eq!(
476 rows.iter().map(|row| row.row_id.0).collect::<Vec<_>>(),
477 vec![2, 3, 4]
478 );
479 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(22)));
480 }
481
482 #[test]
483 fn controlled_merge_uses_hlc_authority_when_epochs_inverted() {
484 use mongreldb_types::hlc::HlcTimestamp;
485 let hlc_old = HlcTimestamp {
486 physical_micros: 100,
487 logical: 0,
488 node_tiebreaker: 1,
489 };
490 let hlc_new = HlcTimestamp {
491 physical_micros: 200,
492 logical: 0,
493 node_tiebreaker: 1,
494 };
495 let a =
498 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
499 let b =
500 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
501 let control = crate::ExecutionControl::new(None);
502 let mut sources = vec![
503 ControlledVisibleSource::memory(a),
504 ControlledVisibleSource::memory(b),
505 ];
506 let mut rows = Vec::new();
507 merge_controlled_visible_sources(
508 &mut sources,
509 &control,
510 |_| false,
511 |row| {
512 rows.push(row);
513 Ok(())
514 },
515 )
516 .unwrap();
517 assert_eq!(rows.len(), 1);
518 assert_eq!(rows[0].row_id.0, 1);
519 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(11)));
521 assert_eq!(rows[0].commit_ts, Some(hlc_new));
522 }
523
524 #[test]
525 fn controlled_merge_epoch_wins_when_one_side_lacks_hlc() {
526 use mongreldb_types::hlc::HlcTimestamp;
527 let hlc_old = HlcTimestamp {
528 physical_micros: 100,
529 logical: 0,
530 node_tiebreaker: 1,
531 };
532 let a =
535 vec![Row::new_with_hlc(RowId(1), Epoch(10), hlc_old).with_column(1, Value::Int64(10))];
536 let b = vec![Row::new(RowId(1), Epoch(5)).with_column(1, Value::Int64(5))];
537 let control = crate::ExecutionControl::new(None);
538 let mut sources = vec![
539 ControlledVisibleSource::memory(a),
540 ControlledVisibleSource::memory(b),
541 ];
542 let mut rows = Vec::new();
543 merge_controlled_visible_sources(
544 &mut sources,
545 &control,
546 |_| false,
547 |row| {
548 rows.push(row);
549 Ok(())
550 },
551 )
552 .unwrap();
553 assert_eq!(rows.len(), 1);
554 assert_eq!(rows[0].columns.get(&1), Some(&Value::Int64(10)));
555 }
556
557 #[test]
560 fn controlled_merge_memory_vs_run_uses_hlc_when_run_stamped() {
561 use crate::sorted_run::{RunReader, RunWriter};
562 use mongreldb_types::hlc::HlcTimestamp;
563 use tempfile::tempdir;
564
565 let hlc_old = HlcTimestamp {
566 physical_micros: 100,
567 logical: 0,
568 node_tiebreaker: 1,
569 };
570 let hlc_new = HlcTimestamp {
571 physical_micros: 200,
572 logical: 0,
573 node_tiebreaker: 1,
574 };
575 let schema = Schema {
576 schema_id: 1,
577 columns: vec![ColumnDef {
578 id: 1,
579 name: "v".into(),
580 ty: TypeId::Int64,
581 flags: ColumnFlags::empty(),
582 default_value: None,
583 embedding_source: None,
584 }],
585 indexes: vec![],
586 colocation: vec![],
587 constraints: Default::default(),
588 clustered: false,
589 };
590 let dir = tempdir().unwrap();
591 let path = dir.path().join("r-hlc.sr");
592 let run_rows =
594 vec![Row::new_with_hlc(RowId(1), Epoch(50), hlc_old).with_column(1, Value::Int64(999))];
595 RunWriter::new(&schema, 1, Epoch(50), 0)
596 .write(&path, &run_rows)
597 .unwrap();
598 let reader = RunReader::open(&path, schema, None).unwrap();
599 assert!(reader.has_column(crate::sorted_run::SYS_COMMIT_TS));
600
601 let mem =
603 vec![Row::new_with_hlc(RowId(1), Epoch(1), hlc_new).with_column(1, Value::Int64(11))];
604 let control = crate::ExecutionControl::new(None);
605 let mut sources = vec![
606 ControlledVisibleSource::memory(mem),
607 ControlledVisibleSource::run(
608 reader.into_visible_version_cursor(Epoch(u64::MAX)).unwrap(),
609 ),
610 ];
611 let mut rows = Vec::new();
612 merge_controlled_visible_sources(
613 &mut sources,
614 &control,
615 |_| false,
616 |row| {
617 rows.push(row);
618 Ok(())
619 },
620 )
621 .unwrap();
622 assert_eq!(rows.len(), 1);
623 assert_eq!(
624 rows[0].columns.get(&1),
625 Some(&Value::Int64(11)),
626 "HLC-newer memory version must beat epoch-newer run"
627 );
628 assert_eq!(rows[0].commit_ts, Some(hlc_new));
629 }
630
631 #[test]
632 fn controlled_memory_source_streams_large_overlays_without_full_active_vec() {
633 let mut map = BTreeMap::new();
634 for i in 1..=500u64 {
635 map.insert(
636 RowId(i),
637 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
638 );
639 }
640 let source = ControlledVisibleSource::memory_from_map(map);
641 assert!(
642 source.is_streaming_memory(),
643 "overlays larger than CONTROLLED_HOT_BATCH must use streaming cursor"
644 );
645 assert_eq!(source.streaming_total(), Some(500));
646 assert!(
648 source.streaming_peak_active().unwrap_or(usize::MAX) <= CONTROLLED_HOT_BATCH,
649 "peak active buffer must be <= CONTROLLED_HOT_BATCH"
650 );
651
652 let control = crate::ExecutionControl::new(None);
654 let mut sources = vec![ControlledVisibleSource::memory_from_map({
655 let mut m = BTreeMap::new();
656 for i in 1..=500u64 {
657 m.insert(
658 RowId(i),
659 Row::new(RowId(i), Epoch(1)).with_column(1, Value::Int64(i as i64)),
660 );
661 }
662 m
663 })];
664 let mut n = 0usize;
665 merge_controlled_visible_sources(
666 &mut sources,
667 &control,
668 |_| false,
669 |_| {
670 n += 1;
671 Ok(())
672 },
673 )
674 .unwrap();
675 assert_eq!(n, 500);
676 assert!(
677 sources[0].streaming_peak_active().unwrap_or(0) <= CONTROLLED_HOT_BATCH,
678 "after full drain peak active still <= batch cap"
679 );
680 }
681}
682
683fn iso_now_bytes() -> Vec<u8> {
686 let secs = std::time::SystemTime::now()
687 .duration_since(std::time::UNIX_EPOCH)
688 .map(|d| d.as_secs() as i64)
689 .unwrap_or(0);
690 let days = secs.div_euclid(86_400);
691 let rem = secs.rem_euclid(86_400);
692 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
693 let (year, month, day) = civil_from_days(days);
694 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
695}
696
697pub(crate) fn unix_nanos_now() -> i64 {
698 std::time::SystemTime::now()
699 .duration_since(std::time::UNIX_EPOCH)
700 .map(|d| d.as_nanos().min(i64::MAX as u128) as i64)
701 .unwrap_or(0)
702}
703
704fn ann_candidate_cap(
705 index_len: usize,
706 context: Option<&crate::query::AiExecutionContext>,
707) -> usize {
708 index_len
709 .min(crate::query::MAX_RAW_INDEX_CANDIDATES)
710 .min(context.map_or(
711 crate::query::MAX_RAW_INDEX_CANDIDATES,
712 crate::query::AiExecutionContext::max_fused_candidates,
713 ))
714}
715
716#[cfg(test)]
717mod ann_candidate_cap_tests {
718 use super::*;
719
720 #[test]
721 fn raw_and_request_candidate_ceilings_are_both_hard_bounds() {
722 assert_eq!(
723 ann_candidate_cap(crate::query::MAX_RAW_INDEX_CANDIDATES + 1, None),
724 crate::query::MAX_RAW_INDEX_CANDIDATES,
725 );
726 let context = crate::query::AiExecutionContext::with_limits(
727 std::time::Duration::from_secs(1),
728 usize::MAX,
729 17,
730 );
731 assert_eq!(ann_candidate_cap(1_000_000, Some(&context)), 17);
732 }
733}
734
735fn civil_from_days(z: i64) -> (i64, u32, u32) {
736 let z = z + 719_468;
737 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
738 let doe = z - era * 146_097;
739 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
740 let y = yoe + era * 400;
741 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
742 let mp = (5 * doy + 2) / 153;
743 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
744 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
745 (if m <= 2 { y + 1 } else { y }, m, d)
746}
747
748pub fn clustered_row_id(primary_key: &Value) -> RowId {
755 let mut hash: u64 = 0xcbf29ce484222325;
756 for byte in primary_key.encode_key() {
757 hash ^= u64::from(byte);
758 hash = hash.wrapping_mul(0x100000001b3);
759 }
760 RowId(hash.max(1))
761}
762
763const 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;
770
771#[derive(Clone, Copy, Debug)]
786struct AutoIncState {
787 column_id: u16,
788 next: i64,
789 seeded: bool,
790}
791
792pub(crate) struct RecoveryMetadataPlan {
793 live_count: u64,
794 auto_inc: Option<AutoIncState>,
795 changed: bool,
796}
797
798type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
799
800fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
803 schema.auto_increment_column().map(|c| AutoIncState {
804 column_id: c.id,
805 next: 0,
806 seeded: false,
807 })
808}
809
810#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
823pub enum IndexBuildPolicy {
824 #[default]
826 Deferred,
827 Eager,
829}
830
831#[derive(Clone)]
832struct ReversePkSegment {
833 values: HashMap<RowId, Vec<u8>>,
834 removed: HashSet<RowId>,
835}
836
837#[derive(Clone)]
838struct ReversePkMap {
839 frozen: Arc<Vec<Arc<ReversePkSegment>>>,
840 active: ReversePkSegment,
841}
842
843impl ReversePkMap {
844 fn new() -> Self {
845 Self {
846 frozen: Arc::new(Vec::new()),
847 active: ReversePkSegment {
848 values: HashMap::new(),
849 removed: HashSet::new(),
850 },
851 }
852 }
853
854 fn from_entries(entries: impl IntoIterator<Item = (RowId, Vec<u8>)>) -> Self {
855 let mut map = Self::new();
856 map.active.values.extend(entries);
857 map
858 }
859
860 fn insert(&mut self, row_id: RowId, key: Vec<u8>) {
861 self.active.removed.remove(&row_id);
862 self.active.values.insert(row_id, key);
863 }
864
865 fn get(&self, row_id: &RowId) -> Option<&Vec<u8>> {
866 if let Some(key) = self.active.values.get(row_id) {
867 return Some(key);
868 }
869 if self.active.removed.contains(row_id) {
870 return None;
871 }
872 for segment in self.frozen.iter().rev() {
873 if let Some(key) = segment.values.get(row_id) {
874 return Some(key);
875 }
876 if segment.removed.contains(row_id) {
877 return None;
878 }
879 }
880 None
881 }
882
883 fn remove(&mut self, row_id: &RowId) -> Option<Vec<u8>> {
884 let previous = self.get(row_id).cloned();
885 self.active.values.remove(row_id);
886 self.active.removed.insert(*row_id);
887 previous
888 }
889
890 fn clear(&mut self) {
891 *self = Self::new();
892 }
893
894 fn entries(&self) -> HashMap<RowId, Vec<u8>> {
895 let mut entries = HashMap::new();
896 for segment in self
897 .frozen
898 .iter()
899 .map(Arc::as_ref)
900 .chain(std::iter::once(&self.active))
901 {
902 for row_id in &segment.removed {
903 entries.remove(row_id);
904 }
905 entries.extend(
906 segment
907 .values
908 .iter()
909 .map(|(row_id, key)| (*row_id, key.clone())),
910 );
911 }
912 entries
913 }
914
915 fn seal(&mut self) {
916 if self.active.values.is_empty() && self.active.removed.is_empty() {
917 return;
918 }
919 let active = std::mem::replace(
920 &mut self.active,
921 ReversePkSegment {
922 values: HashMap::new(),
923 removed: HashSet::new(),
924 },
925 );
926 Arc::make_mut(&mut self.frozen).push(Arc::new(active));
927 if self.frozen.len() >= crate::MAX_READ_GENERATION_LAYERS {
928 self.frozen = Arc::new(vec![Arc::new(ReversePkSegment {
929 values: self.entries(),
930 removed: HashSet::new(),
931 })]);
932 }
933 }
934}
935
936#[derive(Clone)]
949pub struct ReadGeneration {
950 schema: Arc<Schema>,
951 base_runs: Arc<Vec<RunRef>>,
952 deltas: TableDeltas,
953 indexes: Arc<IndexGeneration>,
954 visible_through: Epoch,
955}
956
957pub(crate) enum SecondaryIndexArtifact {
961 Bitmap(u16, BitmapIndex),
962 LearnedRange(u16, ColumnLearnedRange),
963 Fm(u16, Box<FmIndex>),
964 Ann(u16, Box<AnnIndex>),
965 Sparse(u16, SparseIndex),
966 MinHash(u16, MinHashIndex),
967}
968
969#[derive(Clone)]
975pub struct TableDeltas {
976 memtable: Memtable,
977 mutable_run: MutableRun,
978 hot: HotIndex,
979 pk_by_row: ReversePkMap,
980}
981
982impl ReadGeneration {
983 fn empty(schema: &Schema) -> Self {
986 Self {
987 schema: Arc::new(schema.clone()),
988 base_runs: Arc::new(Vec::new()),
989 deltas: TableDeltas {
990 memtable: Memtable::new(),
991 mutable_run: MutableRun::new(),
992 hot: HotIndex::new(),
993 pk_by_row: ReversePkMap::new(),
994 },
995 indexes: Arc::new(IndexGeneration::default()),
996 visible_through: Epoch(0),
997 }
998 }
999
1000 pub fn schema(&self) -> &Arc<Schema> {
1002 &self.schema
1003 }
1004
1005 pub fn base_runs(&self) -> &[RunRef] {
1007 &self.base_runs
1008 }
1009
1010 pub fn indexes(&self) -> &Arc<IndexGeneration> {
1012 &self.indexes
1013 }
1014
1015 pub fn visible_through(&self) -> Epoch {
1017 self.visible_through
1018 }
1019
1020 pub fn deltas(&self) -> &TableDeltas {
1024 &self.deltas
1025 }
1026}
1027
1028impl TableDeltas {
1029 pub fn approx_bytes(&self) -> u64 {
1032 self.memtable
1033 .approx_bytes()
1034 .saturating_add(self.mutable_run.approx_bytes())
1035 }
1036
1037 pub fn memtable_len(&self) -> usize {
1039 self.memtable.len()
1040 }
1041
1042 pub fn mutable_run_len(&self) -> usize {
1044 self.mutable_run.len()
1045 }
1046
1047 pub fn hot_len(&self) -> usize {
1049 self.hot.len()
1050 }
1051
1052 pub fn reverse_pk_len(&self) -> usize {
1054 self.pk_by_row.entries().len()
1055 }
1056}
1057
1058#[derive(Clone)]
1060pub struct Table {
1061 dir: PathBuf,
1062 _root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
1063 runs_root: Option<Arc<crate::durable_file::DurableRoot>>,
1064 idx_root: Option<Arc<crate::durable_file::DurableRoot>>,
1065 table_id: u64,
1066 name: String,
1070 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
1075 read_only: bool,
1078 durable_commit_failed: bool,
1082 wal: WalSink,
1083 memtable: Memtable,
1084 mutable_run: MutableRun,
1089 mutable_run_spill_bytes: u64,
1091 compaction_zstd_level: i32,
1094 allocator: RowIdAllocator,
1095 epoch: Arc<EpochAuthority>,
1096 data_generation: u64,
1099 schema: Schema,
1100 hot: HotIndex,
1101 kek: Option<Arc<Kek>>,
1104 column_keys: HashMap<u16, ([u8; 32], u8)>,
1108 run_refs: Vec<RunRef>,
1109 retiring: Vec<crate::manifest::RetiredRun>,
1112 next_run_id: u64,
1113 sync_byte_threshold: u64,
1114 current_txn_id: u64,
1119 pending_private_mutations: bool,
1123 bitmap: HashMap<u16, BitmapIndex>,
1124 ann: HashMap<u16, AnnIndex>,
1125 fm: HashMap<u16, FmIndex>,
1126 sparse: HashMap<u16, SparseIndex>,
1127 minhash: HashMap<u16, MinHashIndex>,
1128 learned_range: Arc<HashMap<u16, ColumnLearnedRange>>,
1131 pk_by_row: ReversePkMap,
1133 pinned: BTreeMap<Epoch, usize>,
1136 pub(crate) live_count: u64,
1139 reservoir: crate::reservoir::Reservoir,
1142 reservoir_complete: bool,
1150 had_deletes: bool,
1154 recent_delete_preimages: HashMap<Vec<u8>, Row>,
1159 run_row_id_ranges: HashMap<u128, (u64, u64)>,
1162 agg_cache: Arc<HashMap<u64, CachedAgg>>,
1166 global_idx_epoch: u64,
1170 indexes_complete: bool,
1175 index_build_policy: IndexBuildPolicy,
1177 pk_by_row_complete: bool,
1184 flushed_epoch: u64,
1187 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
1190 snapshots: Arc<crate::retention::SnapshotRegistry>,
1193 commit_lock: Arc<parking_lot::Mutex<()>>,
1195 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
1198 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
1208 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
1217 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1219 pending_delete_rids: roaring::RoaringBitmap,
1222 pending_put_cols: std::collections::HashSet<u16>,
1225 pending_rows: Vec<Row>,
1231 pending_rows_auto_inc: Vec<bool>,
1232 pending_dels: Vec<RowId>,
1235 pending_truncate: Option<Epoch>,
1239 auto_inc: Option<AutoIncState>,
1242 ttl: Option<TtlPolicy>,
1245 pins: Arc<crate::retention::PinRegistry>,
1252 published: Arc<ArcSwap<ReadGeneration>>,
1257 read_generation_pin: Option<Arc<crate::retention::PinGuard>>,
1262 pub(crate) lookup_metrics: LookupMetrics,
1267 run_lookup: RunLookupState,
1271}
1272
1273#[derive(Debug, Default, Clone)]
1281struct RunLookupState {
1282 directory: Option<std::sync::Arc<crate::run_lookup::RunLookupDirectory>>,
1283 fingerprint: u64,
1284 complete: bool,
1285}
1286
1287#[derive(Debug, Default)]
1288pub struct LookupMetrics {
1289 hot_lookup_hit: std::sync::atomic::AtomicU64,
1290 hot_lookup_fallback: std::sync::atomic::AtomicU64,
1291 hot_lookup_fallback_overlay_rows: std::sync::atomic::AtomicU64,
1292 hot_lookup_fallback_runs: std::sync::atomic::AtomicU64,
1293 result_cache_memory_hit: std::sync::atomic::AtomicU64,
1294 result_cache_disk_hit: std::sync::atomic::AtomicU64,
1295 result_cache_miss: std::sync::atomic::AtomicU64,
1296 result_cache_persistent_write_us: std::sync::atomic::AtomicU64,
1297 get_run_opened: std::sync::atomic::AtomicU64,
1299 get_run_skipped: std::sync::atomic::AtomicU64,
1300 pub(crate) directory_lookup_hit: std::sync::atomic::AtomicU64,
1302 pub(crate) directory_lookup_fallback: std::sync::atomic::AtomicU64,
1303 pub(crate) directory_incomplete: std::sync::atomic::AtomicU64,
1304 pub(crate) directory_run_readers_opened: std::sync::atomic::AtomicU64,
1305 pub(crate) directory_early_stop_total: std::sync::atomic::AtomicU64,
1306 pub(crate) result_cache_persist_enqueued_total: std::sync::atomic::AtomicU64,
1308 pub(crate) result_cache_persist_coalesced_total: std::sync::atomic::AtomicU64,
1309 pub(crate) result_cache_persist_dropped_store_total: std::sync::atomic::AtomicU64,
1310 pub(crate) result_cache_persist_remove_total: std::sync::atomic::AtomicU64,
1311 pub(crate) result_cache_persist_stale_store_skipped_total: std::sync::atomic::AtomicU64,
1312 pub(crate) result_cache_persist_errors_total: std::sync::atomic::AtomicU64,
1313 pub(crate) result_cache_persist_shutdown_abandoned_total: std::sync::atomic::AtomicU64,
1314 pub(crate) result_cache_persist_queue_depth: std::sync::atomic::AtomicU64,
1315 pub(crate) hot_fallback_reasons: [std::sync::atomic::AtomicU64; 9],
1317 pub(crate) hot_fallback_overlay_versions_total: std::sync::atomic::AtomicU64,
1318 pub(crate) hot_fallback_runs_considered_total: std::sync::atomic::AtomicU64,
1319 pub(crate) hot_fallback_runs_opened_total: std::sync::atomic::AtomicU64,
1320 pub(crate) hot_fallback_pages_decoded_total: std::sync::atomic::AtomicU64,
1321 pub(crate) hot_fallback_rows_materialized_total: std::sync::atomic::AtomicU64,
1322 pub(crate) hot_lookup_duration_nanos: std::sync::atomic::AtomicU64,
1323 pub(crate) hot_fallback_duration_nanos: std::sync::atomic::AtomicU64,
1324 pub(crate) hot_mapping_rebuild_total: std::sync::atomic::AtomicU64,
1325 pub(crate) hot_checkpoint_rejected_total: std::sync::atomic::AtomicU64,
1326}
1327
1328impl Clone for LookupMetrics {
1329 fn clone(&self) -> Self {
1330 let mut hot_fallback_reasons = [
1331 std::sync::atomic::AtomicU64::new(0),
1332 std::sync::atomic::AtomicU64::new(0),
1333 std::sync::atomic::AtomicU64::new(0),
1334 std::sync::atomic::AtomicU64::new(0),
1335 std::sync::atomic::AtomicU64::new(0),
1336 std::sync::atomic::AtomicU64::new(0),
1337 std::sync::atomic::AtomicU64::new(0),
1338 std::sync::atomic::AtomicU64::new(0),
1339 std::sync::atomic::AtomicU64::new(0),
1340 ];
1341 for (dst, src) in hot_fallback_reasons
1342 .iter_mut()
1343 .zip(self.hot_fallback_reasons.iter())
1344 {
1345 dst.store(
1346 src.load(std::sync::atomic::Ordering::Relaxed),
1347 std::sync::atomic::Ordering::Relaxed,
1348 );
1349 }
1350 let copy_atomic = |src: &std::sync::atomic::AtomicU64| {
1351 std::sync::atomic::AtomicU64::new(src.load(std::sync::atomic::Ordering::Relaxed))
1352 };
1353 Self {
1354 hot_lookup_hit: copy_atomic(&self.hot_lookup_hit),
1355 hot_lookup_fallback: copy_atomic(&self.hot_lookup_fallback),
1356 hot_lookup_fallback_overlay_rows: copy_atomic(&self.hot_lookup_fallback_overlay_rows),
1357 hot_lookup_fallback_runs: copy_atomic(&self.hot_lookup_fallback_runs),
1358 result_cache_memory_hit: copy_atomic(&self.result_cache_memory_hit),
1359 result_cache_disk_hit: copy_atomic(&self.result_cache_disk_hit),
1360 result_cache_miss: copy_atomic(&self.result_cache_miss),
1361 result_cache_persistent_write_us: copy_atomic(&self.result_cache_persistent_write_us),
1362 get_run_opened: copy_atomic(&self.get_run_opened),
1363 get_run_skipped: copy_atomic(&self.get_run_skipped),
1364 directory_lookup_hit: copy_atomic(&self.directory_lookup_hit),
1365 directory_lookup_fallback: copy_atomic(&self.directory_lookup_fallback),
1366 directory_incomplete: copy_atomic(&self.directory_incomplete),
1367 directory_run_readers_opened: copy_atomic(&self.directory_run_readers_opened),
1368 directory_early_stop_total: copy_atomic(&self.directory_early_stop_total),
1369 result_cache_persist_enqueued_total: copy_atomic(
1370 &self.result_cache_persist_enqueued_total,
1371 ),
1372 result_cache_persist_coalesced_total: copy_atomic(
1373 &self.result_cache_persist_coalesced_total,
1374 ),
1375 result_cache_persist_dropped_store_total: copy_atomic(
1376 &self.result_cache_persist_dropped_store_total,
1377 ),
1378 result_cache_persist_remove_total: copy_atomic(&self.result_cache_persist_remove_total),
1379 result_cache_persist_stale_store_skipped_total: copy_atomic(
1380 &self.result_cache_persist_stale_store_skipped_total,
1381 ),
1382 result_cache_persist_errors_total: copy_atomic(&self.result_cache_persist_errors_total),
1383 result_cache_persist_shutdown_abandoned_total: copy_atomic(
1384 &self.result_cache_persist_shutdown_abandoned_total,
1385 ),
1386 result_cache_persist_queue_depth: copy_atomic(&self.result_cache_persist_queue_depth),
1387 hot_fallback_reasons,
1388 hot_fallback_overlay_versions_total: copy_atomic(
1389 &self.hot_fallback_overlay_versions_total,
1390 ),
1391 hot_fallback_runs_considered_total: copy_atomic(
1392 &self.hot_fallback_runs_considered_total,
1393 ),
1394 hot_fallback_runs_opened_total: copy_atomic(&self.hot_fallback_runs_opened_total),
1395 hot_fallback_pages_decoded_total: copy_atomic(&self.hot_fallback_pages_decoded_total),
1396 hot_fallback_rows_materialized_total: copy_atomic(
1397 &self.hot_fallback_rows_materialized_total,
1398 ),
1399 hot_lookup_duration_nanos: copy_atomic(&self.hot_lookup_duration_nanos),
1400 hot_fallback_duration_nanos: copy_atomic(&self.hot_fallback_duration_nanos),
1401 hot_mapping_rebuild_total: copy_atomic(&self.hot_mapping_rebuild_total),
1402 hot_checkpoint_rejected_total: copy_atomic(&self.hot_checkpoint_rejected_total),
1403 }
1404 }
1405}
1406
1407impl LookupMetrics {
1408 fn snapshot(&self) -> LookupMetricsSnapshot {
1409 LookupMetricsSnapshot {
1410 hot_lookup_hit: self
1411 .hot_lookup_hit
1412 .load(std::sync::atomic::Ordering::Relaxed),
1413 hot_lookup_fallback: self
1414 .hot_lookup_fallback
1415 .load(std::sync::atomic::Ordering::Relaxed),
1416 hot_lookup_fallback_overlay_rows: self
1417 .hot_lookup_fallback_overlay_rows
1418 .load(std::sync::atomic::Ordering::Relaxed),
1419 hot_lookup_fallback_runs: self
1420 .hot_lookup_fallback_runs
1421 .load(std::sync::atomic::Ordering::Relaxed),
1422 result_cache_memory_hit: self
1423 .result_cache_memory_hit
1424 .load(std::sync::atomic::Ordering::Relaxed),
1425 result_cache_disk_hit: self
1426 .result_cache_disk_hit
1427 .load(std::sync::atomic::Ordering::Relaxed),
1428 result_cache_miss: self
1429 .result_cache_miss
1430 .load(std::sync::atomic::Ordering::Relaxed),
1431 result_cache_persistent_write_us: self
1432 .result_cache_persistent_write_us
1433 .load(std::sync::atomic::Ordering::Relaxed),
1434 get_run_opened: self
1435 .get_run_opened
1436 .load(std::sync::atomic::Ordering::Relaxed),
1437 get_run_skipped: self
1438 .get_run_skipped
1439 .load(std::sync::atomic::Ordering::Relaxed),
1440 directory_lookup_hit: self
1441 .directory_lookup_hit
1442 .load(std::sync::atomic::Ordering::Relaxed),
1443 directory_lookup_fallback: self
1444 .directory_lookup_fallback
1445 .load(std::sync::atomic::Ordering::Relaxed),
1446 directory_incomplete: self
1447 .directory_incomplete
1448 .load(std::sync::atomic::Ordering::Relaxed),
1449 directory_run_readers_opened: self
1450 .directory_run_readers_opened
1451 .load(std::sync::atomic::Ordering::Relaxed),
1452 directory_early_stop_total: self
1453 .directory_early_stop_total
1454 .load(std::sync::atomic::Ordering::Relaxed),
1455 result_cache_persist_enqueued_total: self
1456 .result_cache_persist_enqueued_total
1457 .load(std::sync::atomic::Ordering::Relaxed),
1458 result_cache_persist_coalesced_total: self
1459 .result_cache_persist_coalesced_total
1460 .load(std::sync::atomic::Ordering::Relaxed),
1461 result_cache_persist_dropped_store_total: self
1462 .result_cache_persist_dropped_store_total
1463 .load(std::sync::atomic::Ordering::Relaxed),
1464 result_cache_persist_remove_total: self
1465 .result_cache_persist_remove_total
1466 .load(std::sync::atomic::Ordering::Relaxed),
1467 result_cache_persist_stale_store_skipped_total: self
1468 .result_cache_persist_stale_store_skipped_total
1469 .load(std::sync::atomic::Ordering::Relaxed),
1470 result_cache_persist_errors_total: self
1471 .result_cache_persist_errors_total
1472 .load(std::sync::atomic::Ordering::Relaxed),
1473 result_cache_persist_shutdown_abandoned_total: self
1474 .result_cache_persist_shutdown_abandoned_total
1475 .load(std::sync::atomic::Ordering::Relaxed),
1476 result_cache_persist_queue_depth: self
1477 .result_cache_persist_queue_depth
1478 .load(std::sync::atomic::Ordering::Relaxed),
1479 hot_fallback_reasons: [
1480 self.hot_fallback_reasons[0].load(std::sync::atomic::Ordering::Relaxed),
1481 self.hot_fallback_reasons[1].load(std::sync::atomic::Ordering::Relaxed),
1482 self.hot_fallback_reasons[2].load(std::sync::atomic::Ordering::Relaxed),
1483 self.hot_fallback_reasons[3].load(std::sync::atomic::Ordering::Relaxed),
1484 self.hot_fallback_reasons[4].load(std::sync::atomic::Ordering::Relaxed),
1485 self.hot_fallback_reasons[5].load(std::sync::atomic::Ordering::Relaxed),
1486 self.hot_fallback_reasons[6].load(std::sync::atomic::Ordering::Relaxed),
1487 self.hot_fallback_reasons[7].load(std::sync::atomic::Ordering::Relaxed),
1488 self.hot_fallback_reasons[8].load(std::sync::atomic::Ordering::Relaxed),
1489 ],
1490 hot_fallback_overlay_versions_total: self
1491 .hot_fallback_overlay_versions_total
1492 .load(std::sync::atomic::Ordering::Relaxed),
1493 hot_fallback_runs_considered_total: self
1494 .hot_fallback_runs_considered_total
1495 .load(std::sync::atomic::Ordering::Relaxed),
1496 hot_fallback_runs_opened_total: self
1497 .hot_fallback_runs_opened_total
1498 .load(std::sync::atomic::Ordering::Relaxed),
1499 hot_fallback_pages_decoded_total: self
1500 .hot_fallback_pages_decoded_total
1501 .load(std::sync::atomic::Ordering::Relaxed),
1502 hot_fallback_rows_materialized_total: self
1503 .hot_fallback_rows_materialized_total
1504 .load(std::sync::atomic::Ordering::Relaxed),
1505 hot_lookup_duration_nanos: self
1506 .hot_lookup_duration_nanos
1507 .load(std::sync::atomic::Ordering::Relaxed),
1508 hot_fallback_duration_nanos: self
1509 .hot_fallback_duration_nanos
1510 .load(std::sync::atomic::Ordering::Relaxed),
1511 hot_mapping_rebuild_total: self
1512 .hot_mapping_rebuild_total
1513 .load(std::sync::atomic::Ordering::Relaxed),
1514 hot_checkpoint_rejected_total: self
1515 .hot_checkpoint_rejected_total
1516 .load(std::sync::atomic::Ordering::Relaxed),
1517 }
1518 }
1519}
1520
1521#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
1524pub struct LookupMetricsSnapshot {
1525 pub hot_lookup_hit: u64,
1526 pub hot_lookup_fallback: u64,
1527 pub hot_lookup_fallback_overlay_rows: u64,
1528 pub hot_lookup_fallback_runs: u64,
1529 pub result_cache_memory_hit: u64,
1530 pub result_cache_disk_hit: u64,
1531 pub result_cache_miss: u64,
1532 pub result_cache_persistent_write_us: u64,
1533 pub get_run_opened: u64,
1534 pub get_run_skipped: u64,
1535 pub directory_lookup_hit: u64,
1537 pub directory_lookup_fallback: u64,
1538 pub directory_incomplete: u64,
1539 pub directory_run_readers_opened: u64,
1540 pub directory_early_stop_total: u64,
1541 pub result_cache_persist_enqueued_total: u64,
1543 pub result_cache_persist_coalesced_total: u64,
1544 pub result_cache_persist_dropped_store_total: u64,
1545 pub result_cache_persist_remove_total: u64,
1546 pub result_cache_persist_stale_store_skipped_total: u64,
1547 pub result_cache_persist_errors_total: u64,
1548 pub result_cache_persist_shutdown_abandoned_total: u64,
1549 pub result_cache_persist_queue_depth: u64,
1550 pub hot_fallback_reasons: [u64; 9],
1552 pub hot_fallback_overlay_versions_total: u64,
1553 pub hot_fallback_runs_considered_total: u64,
1554 pub hot_fallback_runs_opened_total: u64,
1555 pub hot_fallback_pages_decoded_total: u64,
1556 pub hot_fallback_rows_materialized_total: u64,
1557 pub hot_lookup_duration_nanos: u64,
1558 pub hot_fallback_duration_nanos: u64,
1559 pub hot_mapping_rebuild_total: u64,
1560 pub hot_checkpoint_rejected_total: u64,
1561}
1562
1563pub fn hot_fallback_reason_index(r: crate::trace::HotFallbackReason) -> usize {
1566 match r {
1567 crate::trace::HotFallbackReason::MissingMapping => 0,
1568 crate::trace::HotFallbackReason::StaleRowId => 1,
1569 crate::trace::HotFallbackReason::InvisibleAtSnapshot => 2,
1570 crate::trace::HotFallbackReason::HistoricalSnapshot => 3,
1571 crate::trace::HotFallbackReason::Tombstone => 4,
1572 crate::trace::HotFallbackReason::TtlExpired => 5,
1573 crate::trace::HotFallbackReason::PrimaryKeyMismatch => 6,
1574 crate::trace::HotFallbackReason::IndexIncomplete => 7,
1575 crate::trace::HotFallbackReason::CheckpointRejected => 8,
1576 }
1577}
1578
1579const _: () = {
1586 const fn assert_sync<T: ?Sized + Sync>() {}
1587 assert_sync::<Table>();
1588};
1589
1590#[derive(Clone)]
1596enum CachedData {
1597 Rows(Arc<Vec<Row>>),
1598 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
1599}
1600
1601impl CachedData {
1602 fn approx_bytes(&self) -> u64 {
1603 match self {
1604 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
1605 CachedData::Columns(c) => c
1606 .iter()
1607 .map(|(_, c)| c.approx_bytes())
1608 .sum::<u64>()
1609 .saturating_add(c.len() as u64 * 16),
1610 }
1611 }
1612}
1613
1614struct CachedEntry {
1618 data: CachedData,
1619 footprint: roaring::RoaringBitmap,
1620 condition_cols: Vec<u16>,
1621}
1622
1623impl CachedEntry {
1624 fn clone_for_persist(&self) -> Self {
1630 Self {
1631 data: self.data.clone(),
1632 footprint: self.footprint.clone(),
1633 condition_cols: self.condition_cols.clone(),
1634 }
1635 }
1636}
1637
1638struct ResultCache {
1649 entries: std::collections::HashMap<u64, CachedEntry>,
1650 order: BTreeSet<(u64, u64)>,
1651 generations: HashMap<u64, u64>,
1652 next_generation: u64,
1653 condition_index: HashMap<u16, HashSet<u64>>,
1656 empty_footprint_entries: HashSet<u64>,
1659 bytes: u64,
1660 max_bytes: u64,
1661 dir: Option<std::path::PathBuf>,
1662 #[allow(dead_code)]
1663 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
1664 memory_hit: std::sync::atomic::AtomicU64,
1668 disk_hit: std::sync::atomic::AtomicU64,
1669 miss: std::sync::atomic::AtomicU64,
1670 persistent_write_us: std::sync::atomic::AtomicU64,
1671 persist_min_bytes: u64,
1676 writer: Option<std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>>,
1682 worker_handle: Option<std::thread::JoinHandle<()>>,
1685}
1686
1687#[derive(serde::Serialize, serde::Deserialize)]
1689struct SerializedEntry {
1690 condition_cols: Vec<u16>,
1691 footprint_bits: Vec<u32>,
1692 data: SerializedData,
1693}
1694
1695#[derive(serde::Serialize, serde::Deserialize)]
1696enum SerializedData {
1697 Rows(Vec<Row>),
1698 Columns(Vec<(u16, columnar::NativeColumn)>),
1699}
1700
1701#[derive(serde::Serialize)]
1702struct SerializedEntryRef<'a> {
1703 condition_cols: &'a [u16],
1704 footprint_bits: Vec<u32>,
1705 data: SerializedDataRef<'a>,
1706}
1707
1708#[derive(serde::Serialize)]
1709enum SerializedDataRef<'a> {
1710 Rows(&'a [Row]),
1711 Columns(&'a [(u16, columnar::NativeColumn)]),
1712}
1713
1714impl<'a> SerializedEntryRef<'a> {
1715 fn from_entry(entry: &'a CachedEntry) -> Self {
1716 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
1717 let data = match &entry.data {
1718 CachedData::Rows(rows) => SerializedDataRef::Rows(rows.as_slice()),
1719 CachedData::Columns(columns) => SerializedDataRef::Columns(columns.as_slice()),
1720 };
1721 Self {
1722 condition_cols: &entry.condition_cols,
1723 footprint_bits,
1724 data,
1725 }
1726 }
1727}
1728
1729impl SerializedEntry {
1730 fn into_entry(self) -> Option<CachedEntry> {
1731 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
1732 let data = match self.data {
1733 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
1734 SerializedData::Columns(c) => {
1735 if !c.iter().all(|(_, col)| col.validate()) {
1738 return None;
1739 }
1740 CachedData::Columns(Arc::new(c))
1741 }
1742 };
1743 Some(CachedEntry {
1744 data,
1745 footprint,
1746 condition_cols: self.condition_cols,
1747 })
1748 }
1749}
1750
1751impl ResultCache {
1752 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
1753
1754 fn new() -> Self {
1755 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
1756 }
1757
1758 fn with_max_bytes(max_bytes: u64) -> Self {
1759 Self {
1760 entries: std::collections::HashMap::new(),
1761 order: BTreeSet::new(),
1762 generations: HashMap::new(),
1763 next_generation: 0,
1764 condition_index: HashMap::new(),
1765 empty_footprint_entries: HashSet::new(),
1766 bytes: 0,
1767 max_bytes,
1768 dir: None,
1769 cache_dek: None,
1770 memory_hit: std::sync::atomic::AtomicU64::new(0),
1771 disk_hit: std::sync::atomic::AtomicU64::new(0),
1772 miss: std::sync::atomic::AtomicU64::new(0),
1773 persistent_write_us: std::sync::atomic::AtomicU64::new(0),
1774 persist_min_bytes: 4 * 1024,
1777 writer: None,
1778 worker_handle: None,
1779 }
1780 }
1781
1782 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
1783 let _ = std::fs::create_dir_all(&dir);
1784 self.dir = Some(dir);
1785 self
1786 }
1787
1788 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
1789 self.cache_dek = dek;
1790 self
1791 }
1792
1793 fn install_persistent_writer(
1798 &mut self,
1799 writer: std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
1800 worker_handle: std::thread::JoinHandle<()>,
1801 ) {
1802 self.writer = Some(writer);
1803 self.worker_handle = Some(worker_handle);
1804 }
1805
1806 fn take_persistent_worker(&mut self) -> Option<std::thread::JoinHandle<()>> {
1809 self.worker_handle.take()
1810 }
1811
1812 fn persistent_writer(&self) -> Option<&std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>> {
1815 self.writer.as_ref()
1816 }
1817
1818 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
1819 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
1820 }
1821
1822 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
1826 let Some(path) = self.disk_path(key) else {
1827 return;
1828 };
1829 let serialized = match bincode::serialize(&SerializedEntryRef::from_entry(entry)) {
1830 Ok(s) => s,
1831 Err(_) => return,
1832 };
1833 let on_disk = if let Some(dek) = &self.cache_dek {
1835 match self.encrypt_cache(&serialized, dek) {
1836 Some(b) => b,
1837 None => return,
1838 }
1839 } else {
1840 serialized
1841 };
1842 let tmp = path.with_extension("tmp");
1843 use std::io::Write;
1844 let write = || -> std::io::Result<()> {
1845 let mut f = std::fs::File::create(&tmp)?;
1846 f.write_all(&on_disk)?;
1847 f.flush()?;
1848 Ok(())
1849 };
1850 if write().is_err() {
1851 let _ = std::fs::remove_file(&tmp);
1852 return;
1853 }
1854 let _ = std::fs::rename(&tmp, &path);
1855 }
1856
1857 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
1859 let path = self.disk_path(key)?;
1860 let bytes = std::fs::read(&path).ok()?;
1861 let plaintext = if let Some(dek) = &self.cache_dek {
1862 self.decrypt_cache(&bytes, dek)?
1863 } else {
1864 bytes
1865 };
1866 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
1867 serialized.into_entry()
1868 }
1869
1870 fn remove_from_disk(&self, key: u64) {
1872 if let Some(path) = self.disk_path(key) {
1873 let _ = std::fs::remove_file(&path);
1874 }
1875 }
1876
1877 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1879 use crate::encryption::Cipher;
1880 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1881 let mut nonce = [0u8; 12];
1882 crate::encryption::fill_random(&mut nonce).ok()?;
1883 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
1884 let mut out = Vec::with_capacity(12 + ct.len());
1885 out.extend_from_slice(&nonce);
1886 out.extend_from_slice(&ct);
1887 Some(out)
1888 }
1889
1890 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
1892 use crate::encryption::Cipher;
1893 if bytes.len() < 28 {
1894 return None;
1895 }
1896 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
1897 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
1898 let ct = &bytes[12..];
1899 cipher.decrypt_page(&nonce, ct).ok()
1900 }
1901
1902 fn load_persistent(&mut self) {
1905 let Some(dir) = self.dir.as_ref().cloned() else {
1906 return;
1907 };
1908 let entries = match std::fs::read_dir(&dir) {
1909 Ok(e) => e,
1910 Err(_) => return,
1911 };
1912 for entry in entries.flatten() {
1913 let path = entry.path();
1914 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
1916 let _ = std::fs::remove_file(&path);
1917 continue;
1918 }
1919 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
1920 continue;
1921 }
1922 let stem = match path.file_stem().and_then(|s| s.to_str()) {
1923 Some(s) => s,
1924 None => continue,
1925 };
1926 let key = match u64::from_str_radix(stem, 16) {
1927 Ok(k) => k,
1928 Err(_) => continue,
1929 };
1930 let bytes = match std::fs::read(&path) {
1931 Ok(b) => b,
1932 Err(_) => continue,
1933 };
1934 let plaintext = if let Some(dek) = &self.cache_dek {
1936 match self.decrypt_cache(&bytes, dek) {
1937 Some(p) => p,
1938 None => {
1939 let _ = std::fs::remove_file(&path);
1940 continue;
1941 }
1942 }
1943 } else {
1944 bytes
1945 };
1946 match bincode::deserialize::<SerializedEntry>(&plaintext) {
1947 Ok(serialized) => {
1948 if let Some(entry) = serialized.into_entry() {
1949 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
1950 self.index_entry(key, &entry);
1951 self.entries.insert(key, entry);
1952 self.touch(key);
1953 } else {
1954 let _ = std::fs::remove_file(&path);
1955 }
1956 }
1957 Err(_) => {
1958 let _ = std::fs::remove_file(&path);
1959 }
1960 }
1961 }
1962 self.evict();
1963 }
1964
1965 fn set_max_bytes(&mut self, max_bytes: u64) {
1966 self.max_bytes = max_bytes;
1967 self.evict();
1968 }
1969
1970 fn touch(&mut self, key: u64) {
1973 if !self.entries.contains_key(&key) {
1974 return;
1975 }
1976 if let Some(previous) = self.generations.remove(&key) {
1977 self.order.remove(&(previous, key));
1978 }
1979 let generation = self.allocate_generation();
1980 self.generations.insert(key, generation);
1981 self.order.insert((generation, key));
1982 }
1983
1984 fn untrack(&mut self, key: u64) {
1985 if let Some(generation) = self.generations.remove(&key) {
1986 self.order.remove(&(generation, key));
1987 }
1988 }
1989
1990 fn index_entry(&mut self, key: u64, entry: &CachedEntry) {
1991 for column in &entry.condition_cols {
1992 self.condition_index.entry(*column).or_default().insert(key);
1993 }
1994 if entry.footprint.is_empty() {
1995 self.empty_footprint_entries.insert(key);
1996 }
1997 }
1998
1999 fn unindex_entry(&mut self, key: u64, entry: &CachedEntry) {
2000 for column in &entry.condition_cols {
2001 let remove_column = self.condition_index.get_mut(column).is_some_and(|keys| {
2002 keys.remove(&key);
2003 keys.is_empty()
2004 });
2005 if remove_column {
2006 self.condition_index.remove(column);
2007 }
2008 }
2009 if entry.footprint.is_empty() {
2010 self.empty_footprint_entries.remove(&key);
2011 }
2012 }
2013
2014 fn allocate_generation(&mut self) -> u64 {
2015 if self.next_generation == u64::MAX {
2016 self.rebase_generations();
2017 }
2018 let generation = self.next_generation;
2019 self.next_generation += 1;
2020 generation
2021 }
2022
2023 fn rebase_generations(&mut self) {
2024 let ordered = std::mem::take(&mut self.order);
2025 self.generations.clear();
2026 for (generation, (_, key)) in ordered.into_iter().enumerate() {
2027 let generation = generation as u64;
2028 self.generations.insert(key, generation);
2029 self.order.insert((generation, key));
2030 }
2031 self.next_generation = self.order.len() as u64;
2032 }
2033
2034 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
2035 let res = self.entries.get(&key).and_then(|e| match &e.data {
2036 CachedData::Rows(r) => Some(r.clone()),
2037 CachedData::Columns(_) => None,
2038 });
2039 if res.is_some() {
2040 self.touch(key);
2041 self.memory_hit
2042 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2043 return res;
2044 }
2045 if let Some(entry) = self.load_from_disk(key) {
2047 let res = match &entry.data {
2048 CachedData::Rows(r) => Some(r.clone()),
2049 CachedData::Columns(_) => None,
2050 };
2051 if res.is_some() {
2052 let approx = entry.data.approx_bytes();
2053 self.bytes = self.bytes.saturating_add(approx);
2054 self.index_entry(key, &entry);
2055 self.entries.insert(key, entry);
2056 self.touch(key);
2057 self.evict();
2058 self.disk_hit
2059 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2060 return res;
2061 }
2062 }
2063 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2064 None
2065 }
2066
2067 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
2068 let res = self.entries.get(&key).and_then(|e| match &e.data {
2069 CachedData::Columns(c) => Some(c.clone()),
2070 CachedData::Rows(_) => None,
2071 });
2072 if res.is_some() {
2073 self.touch(key);
2074 self.memory_hit
2075 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2076 return res;
2077 }
2078 if let Some(entry) = self.load_from_disk(key) {
2080 let res = match &entry.data {
2081 CachedData::Columns(c) => Some(c.clone()),
2082 CachedData::Rows(_) => None,
2083 };
2084 if res.is_some() {
2085 let approx = entry.data.approx_bytes();
2086 self.bytes = self.bytes.saturating_add(approx);
2087 self.index_entry(key, &entry);
2088 self.entries.insert(key, entry);
2089 self.touch(key);
2090 self.evict();
2091 self.disk_hit
2092 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2093 return res;
2094 }
2095 }
2096 self.miss.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
2097 None
2098 }
2099
2100 fn insert(&mut self, key: u64, entry: CachedEntry) {
2101 let approx = entry.data.approx_bytes();
2102 if let Some(previous) = self.entries.remove(&key) {
2103 self.bytes = self.bytes.saturating_sub(previous.data.approx_bytes());
2104 self.unindex_entry(key, &previous);
2105 self.untrack(key);
2106 }
2107 if self.dir.is_some()
2119 && self.writer.is_none()
2120 && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
2121 {
2122 let write_start = std::time::Instant::now();
2123 self.store_to_disk(key, &entry);
2124 let write_us = write_start.elapsed().as_micros() as u64;
2125 self.persistent_write_us
2126 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2127 }
2128 self.bytes = self.bytes.saturating_add(approx);
2129 self.index_entry(key, &entry);
2130 self.entries.insert(key, entry);
2131 self.touch(key);
2132 self.evict();
2133 }
2134
2135 fn allocate_persist_generation(&mut self, key: u64) -> u64 {
2140 if let Some(writer) = self.writer.as_ref() {
2141 let _ = key; writer.bump_persist_generation(key)
2145 } else {
2146 0
2147 }
2148 }
2149
2150 fn enqueue_persist(
2156 &self,
2157 key: u64,
2158 entry: &CachedEntry,
2159 table_id: u64,
2160 schema_id: u64,
2161 run_generation: u64,
2162 entry_generation: u64,
2163 ) {
2164 let Some(writer) = self.writer.as_ref() else {
2165 return;
2166 };
2167 let approx = entry.data.approx_bytes();
2168 if self.persist_min_bytes > 0 && approx < self.persist_min_bytes {
2169 return;
2170 }
2171 let entry_clone = entry.clone_for_persist();
2172 let payload_factory: Box<dyn FnOnce() -> Option<Vec<u8>> + Send + 'static> =
2173 Box::new(move || {
2174 bincode::serialize(&SerializedEntryRef::from_entry(&entry_clone)).ok()
2175 });
2176 let persistable = crate::result_cache::PersistableEntry {
2177 key,
2178 table_id,
2179 schema_id,
2180 run_generation,
2181 entry_generation,
2182 bytes: approx as usize,
2183 payload_factory,
2184 };
2185 let write_start = std::time::Instant::now();
2186 writer.enqueue_store(persistable);
2187 let write_us = write_start.elapsed().as_micros() as u64;
2188 self.persistent_write_us
2189 .fetch_add(write_us, std::sync::atomic::Ordering::Relaxed);
2190 }
2191
2192 fn enqueue_persist_remove(&self, key: u64) {
2195 if let Some(writer) = self.writer.as_ref() {
2196 writer.enqueue_remove(key);
2197 }
2198 }
2199
2200 fn enqueue_persist_clear(&self) {
2203 if let Some(writer) = self.writer.as_ref() {
2204 writer.enqueue_clear();
2205 }
2206 }
2207
2208 fn has_persistent_writer(&self) -> bool {
2210 self.writer.is_some()
2211 }
2212
2213 fn flush_persistent_cache(&self, deadline: std::time::Duration) -> u64 {
2219 let Some(writer) = self.persistent_writer() else {
2220 return 0;
2221 };
2222 let start = std::time::Instant::now();
2223 loop {
2224 let depth = writer.queue_depth() as u64;
2225 let in_flight = writer.writes_in_flight();
2231 if depth == 0 && in_flight == 0 {
2232 return 0;
2233 }
2234 if start.elapsed() >= deadline {
2235 return depth;
2236 }
2237 std::thread::sleep(std::time::Duration::from_millis(2));
2238 }
2239 }
2240
2241 fn shutdown_persistent_cache(&mut self, deadline: std::time::Duration) {
2246 let Some(writer) = self.persistent_writer().cloned() else {
2247 return;
2248 };
2249 let start = std::time::Instant::now();
2250 loop {
2251 if writer.queue_depth() == 0 {
2252 break;
2253 }
2254 if start.elapsed() >= deadline {
2255 break;
2256 }
2257 std::thread::sleep(std::time::Duration::from_millis(2));
2258 }
2259 writer.shutdown();
2260 let _ = writer.drain_all_as_abandoned();
2263 if let Some(handle) = self.take_persistent_worker() {
2265 let _ = handle.join();
2266 }
2267 }
2268
2269 #[cfg(test)]
2270 fn set_persist_min_bytes(&mut self, min: u64) {
2271 self.persist_min_bytes = min;
2272 }
2273
2274 #[cfg(test)]
2276 fn would_persist(&self, approx: u64) -> bool {
2277 self.dir.is_some() && (self.persist_min_bytes == 0 || approx >= self.persist_min_bytes)
2278 }
2279
2280 fn cache_counters(&self) -> (u64, u64, u64, u64) {
2283 (
2284 self.memory_hit.load(std::sync::atomic::Ordering::Relaxed),
2285 self.disk_hit.load(std::sync::atomic::Ordering::Relaxed),
2286 self.miss.load(std::sync::atomic::Ordering::Relaxed),
2287 self.persistent_write_us
2288 .load(std::sync::atomic::Ordering::Relaxed),
2289 )
2290 }
2291
2292 fn persist_snapshot(&self) -> Option<LookupMetricsSnapshot> {
2295 self.writer
2296 .as_ref()
2297 .map(|w| w.persist_snapshot())
2298 }
2299
2300 fn invalidate(
2309 &mut self,
2310 delete_rids: &roaring::RoaringBitmap,
2311 put_cols: &std::collections::HashSet<u16>,
2312 ) {
2313 if self.entries.is_empty() {
2314 return;
2315 }
2316 let has_deletes = !delete_rids.is_empty();
2317 let mut to_remove = HashSet::new();
2318
2319 for column in put_cols {
2323 if let Some(keys) = self.condition_index.get(column) {
2324 to_remove.extend(keys.iter().copied());
2325 }
2326 }
2327
2328 if has_deletes {
2329 to_remove.extend(self.empty_footprint_entries.iter().copied());
2333 for (&key, entry) in &self.entries {
2334 if !to_remove.contains(&key)
2335 && !entry.footprint.is_empty()
2336 && entry.footprint.intersection_len(delete_rids) > 0
2337 {
2338 to_remove.insert(key);
2339 }
2340 }
2341 }
2342
2343 for key in to_remove {
2344 if let Some(entry) = self.entries.remove(&key) {
2345 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2346 self.unindex_entry(key, &entry);
2347 }
2348 if self.writer.is_some() {
2353 self.enqueue_persist_remove(key);
2354 } else {
2355 self.remove_from_disk(key);
2356 }
2357 self.untrack(key);
2358 }
2359 }
2360
2361 fn clear(&mut self) {
2362 if self.dir.is_some() {
2366 if self.writer.is_some() {
2367 self.enqueue_persist_clear();
2368 } else if let Some(dir) = &self.dir {
2369 if let Ok(entries) = std::fs::read_dir(dir) {
2370 for entry in entries.flatten() {
2371 let path = entry.path();
2372 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
2373 let _ = std::fs::remove_file(&path);
2374 }
2375 }
2376 }
2377 }
2378 }
2379 self.entries.clear();
2380 self.order.clear();
2381 self.generations.clear();
2382 self.next_generation = 0;
2383 self.condition_index.clear();
2384 self.empty_footprint_entries.clear();
2385 self.bytes = 0;
2386 }
2387
2388 fn evict(&mut self) {
2389 while self.bytes > self.max_bytes {
2390 let Some((_, key)) = self.order.pop_first() else {
2391 break;
2392 };
2393 self.generations.remove(&key);
2394 if let Some(entry) = self.entries.remove(&key) {
2395 self.bytes = self.bytes.saturating_sub(entry.data.approx_bytes());
2396 self.unindex_entry(key, &entry);
2397 if self.writer.is_some() {
2403 self.enqueue_persist_remove(key);
2404 } else {
2405 self.remove_from_disk(key);
2406 }
2407 }
2408 }
2409 }
2410}
2411
2412fn spawn_persistent_cache_worker(
2417 dir: std::path::PathBuf,
2418 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
2419 metrics: LookupMetrics,
2420) -> std::io::Result<(
2421 std::sync::Arc<crate::result_cache::PersistentResultCacheWriter>,
2422 std::thread::JoinHandle<()>,
2423)> {
2424 use crate::result_cache::{
2425 RealPersistentCacheIo, StalenessGuard, WorkerConfig,
2426 WriterStalenessGuard, spawn_persistent_cache_worker as spawn,
2427 };
2428 let io: std::sync::Arc<dyn crate::result_cache::PersistentCacheIo> =
2429 std::sync::Arc::new(RealPersistentCacheIo::new(dir)?);
2430 let cipher = match cache_dek {
2431 Some(dek) => {
2432 let cipher = crate::encryption::AesCipher::new(&dek[..])
2433 .map_err(|e| std::io::Error::other(format!("aes: {e}")))?;
2434 Some(std::sync::Arc::new(cipher))
2435 }
2436 None => None,
2437 };
2438 let writer =
2439 std::sync::Arc::new(crate::result_cache::PersistentResultCacheWriter::new(
2440 metrics,
2441 crate::result_cache::WriterLimits::default(),
2442 ));
2443 let staleness: std::sync::Arc<dyn StalenessGuard> = std::sync::Arc::new(
2444 WriterStalenessGuard::new(writer.clone()),
2445 );
2446 let config = WorkerConfig {
2447 writer: writer.clone(),
2448 io,
2449 cipher,
2450 staleness,
2451 max_staleness_retries: 8,
2452 };
2453 let handle = spawn(config);
2454 Ok((writer, handle))
2455}
2456
2457#[cfg(test)]
2458mod result_cache_lru_tests {
2459 use super::*;
2460
2461 fn row_entry(row_id: u64) -> CachedEntry {
2462 CachedEntry {
2463 data: CachedData::Rows(Arc::new(vec![Row::new(RowId(row_id), Epoch(1))])),
2464 footprint: std::iter::once(row_id as u32).collect(),
2465 condition_cols: vec![1],
2466 }
2467 }
2468
2469 #[test]
2470 fn hits_update_recency_without_growing_an_order_queue() {
2471 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2472 cache.insert(1, row_entry(1));
2473 cache.insert(2, row_entry(2));
2474 assert_eq!(cache.order.len(), cache.entries.len());
2475
2476 for _ in 0..1_000 {
2477 assert!(cache.get_rows(1).is_some());
2478 }
2479
2480 assert_eq!(cache.order.len(), cache.entries.len());
2481 assert_eq!(cache.order.first().map(|(_, key)| *key), Some(2));
2482
2483 let one_entry = cache.entries[&1].data.approx_bytes();
2484 cache.set_max_bytes(one_entry);
2485 assert!(cache.entries.contains_key(&1));
2486 assert!(!cache.entries.contains_key(&2));
2487 assert_eq!(cache.order.len(), cache.entries.len());
2488 }
2489
2490 #[test]
2491 fn tiny_entries_skip_persistent_tier_by_default() {
2492 let dir = tempfile::tempdir().unwrap();
2493 let mut cache = ResultCache::with_max_bytes(u64::MAX).with_dir(dir.path().to_path_buf());
2494 let tiny = row_entry(1).data.approx_bytes();
2495 assert!(
2496 tiny < 4 * 1024,
2497 "row_entry fixture must be below default 4KiB threshold"
2498 );
2499 assert!(
2500 !cache.would_persist(tiny),
2501 "tiny entry must not force synchronous disk publish"
2502 );
2503 assert!(
2504 cache.would_persist(8 * 1024),
2505 "large entry must still persist when dir is set"
2506 );
2507 cache.set_persist_min_bytes(0);
2508 assert!(
2509 cache.would_persist(tiny),
2510 "persist_min_bytes=0 restores always-persist policy"
2511 );
2512 }
2513
2514 #[test]
2515 fn insert_invalidation_uses_the_condition_reverse_index() {
2516 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2517 let mut first = row_entry(1);
2518 first.condition_cols = vec![1];
2519 let mut second = row_entry(2);
2520 second.condition_cols = vec![2];
2521 cache.insert(1, first);
2522 cache.insert(2, second);
2523
2524 let put_cols = std::iter::once(1_u16).collect();
2525 cache.invalidate(&roaring::RoaringBitmap::new(), &put_cols);
2526
2527 assert!(!cache.entries.contains_key(&1));
2528 assert!(cache.entries.contains_key(&2));
2529 assert!(!cache
2530 .condition_index
2531 .get(&1)
2532 .is_some_and(|keys| keys.contains(&1)));
2533 assert!(cache
2534 .condition_index
2535 .get(&2)
2536 .is_some_and(|keys| keys.contains(&2)));
2537 assert_eq!(cache.order.len(), cache.entries.len());
2538 }
2539
2540 #[test]
2541 fn delete_invalidation_preserves_known_and_unknown_footprint_semantics() {
2542 let mut cache = ResultCache::with_max_bytes(u64::MAX);
2543 cache.insert(1, row_entry(1));
2544 cache.insert(2, row_entry(2));
2545 let mut unknown = row_entry(3);
2546 unknown.footprint.clear();
2547 cache.insert(3, unknown);
2548
2549 let delete_rids: roaring::RoaringBitmap = std::iter::once(1_u32).collect();
2550 cache.invalidate(&delete_rids, &HashSet::new());
2551
2552 assert!(!cache.entries.contains_key(&1));
2553 assert!(cache.entries.contains_key(&2));
2554 assert!(!cache.entries.contains_key(&3));
2555 assert!(cache.empty_footprint_entries.is_empty());
2556 assert_eq!(cache.order.len(), cache.entries.len());
2557 }
2558
2559 #[test]
2560 fn borrowed_persistence_encoding_matches_the_existing_owned_format() {
2561 let entry = row_entry(7);
2562 let borrowed = bincode::serialize(&SerializedEntryRef::from_entry(&entry)).unwrap();
2563 let owned = bincode::serialize(&SerializedEntry {
2564 condition_cols: entry.condition_cols.clone(),
2565 footprint_bits: entry.footprint.iter().collect(),
2566 data: match &entry.data {
2567 CachedData::Rows(rows) => SerializedData::Rows((**rows).clone()),
2568 CachedData::Columns(columns) => SerializedData::Columns((**columns).clone()),
2569 },
2570 })
2571 .unwrap();
2572 assert_eq!(borrowed, owned);
2573 }
2574}
2575
2576type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
2583
2584fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
2585 let _ = kek;
2586 {
2587 if let Some(k) = kek {
2588 return (
2589 Some(k.derive_table_wal_key(_table_id)),
2590 Some(k.derive_cache_key()),
2591 );
2592 }
2593 }
2594 (None, None)
2595}
2596
2597fn read_table_encryption_salt_root(
2598 root: &crate::durable_file::DurableRoot,
2599) -> Result<[u8; crate::encryption::SALT_LEN]> {
2600 use std::io::Read;
2601
2602 let mut file = root
2603 .open_regular(Path::new(META_DIR).join(KEYS_FILENAME))
2604 .map_err(|error| MongrelError::NotFound(format!("encryption salt file: {error}")))?;
2605 let length = file.metadata()?.len();
2606 if length != crate::encryption::SALT_LEN as u64 {
2607 return Err(MongrelError::InvalidArgument(format!(
2608 "salt file is {length} bytes, expected {}",
2609 crate::encryption::SALT_LEN
2610 )));
2611 }
2612 let mut salt = [0_u8; crate::encryption::SALT_LEN];
2613 file.read_exact(&mut salt)?;
2614 Ok(salt)
2615}
2616
2617fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
2619 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
2620}
2621
2622fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
2623 let Some(kek) = kek else {
2624 return HashMap::new();
2625 };
2626 {
2627 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
2628 schema
2629 .columns
2630 .iter()
2631 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
2632 .map(|c| {
2633 let scheme = if schema
2634 .indexes
2635 .iter()
2636 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
2637 {
2638 SCHEME_OPE_RANGE
2639 } else {
2640 SCHEME_HMAC_EQ
2641 };
2642 let key: [u8; 32] = *kek.derive_column_key(c.id);
2643 (c.id, (key, scheme))
2644 })
2645 .collect()
2646 }
2647}
2648
2649pub(crate) struct SharedCtx {
2654 pub root_guard: Option<Arc<crate::durable_file::DurableRoot>>,
2655 pub epoch: Arc<EpochAuthority>,
2656 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
2657 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
2658 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
2659 pub kek: Option<Arc<Kek>>,
2660 pub commit_lock: Arc<parking_lot::Mutex<()>>,
2666 pub shared: Option<SharedWalCtx>,
2670 pub table_name: Option<String>,
2673 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
2676 pub read_only: bool,
2678}
2679
2680#[derive(Clone)]
2686pub(crate) struct SharedWalCtx {
2687 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
2688 pub group: Arc<GroupCommit>,
2689 pub poisoned: Arc<AtomicBool>,
2690 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
2691 pub change_wake: tokio::sync::broadcast::Sender<()>,
2692 pub lifecycle: Arc<crate::core::LifecycleController>,
2695 pub hlc: Arc<mongreldb_types::hlc::HlcClock>,
2697}
2698
2699enum WalSink {
2702 Private(Wal),
2703 Shared(SharedWalCtx),
2704 ReadOnly,
2705}
2706
2707impl Clone for WalSink {
2708 fn clone(&self) -> Self {
2709 match self {
2710 Self::Shared(shared) => Self::Shared(shared.clone()),
2711 Self::Private(_) | Self::ReadOnly => Self::ReadOnly,
2712 }
2713 }
2714}
2715
2716impl SharedCtx {
2717 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
2721 let n_shards = if cache_dir.is_some() {
2725 1
2726 } else {
2727 crate::cache::CACHE_SHARDS
2728 };
2729 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
2730 let page_cache = if let Some(d) = cache_dir {
2731 Arc::new(crate::cache::Sharded::new(1, || {
2732 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
2733 }))
2734 } else {
2735 Arc::new(crate::cache::Sharded::new(n_shards, || {
2736 crate::cache::PageCache::new(per_shard)
2737 }))
2738 };
2739 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
2740 let decoded_cache = Arc::new(crate::cache::Sharded::new(
2741 crate::cache::CACHE_SHARDS,
2742 || crate::cache::DecodedPageCache::new(decoded_per_shard),
2743 ));
2744 Self {
2745 root_guard: None,
2746 epoch: Arc::new(EpochAuthority::new(0)),
2747 page_cache,
2748 decoded_cache,
2749 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
2750 kek,
2751 commit_lock: Arc::new(parking_lot::Mutex::new(())),
2752 shared: None,
2753 table_name: None,
2754 auth: None,
2755 read_only: false,
2756 }
2757 }
2758}
2759
2760fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
2764 use crate::query::Condition;
2765 match c {
2766 Condition::Pk(_)
2768 | Condition::BitmapEq { .. }
2769 | Condition::BitmapIn { .. }
2770 | Condition::BytesPrefix { .. }
2771 | Condition::IsNull { .. }
2772 | Condition::IsNotNull { .. } => 0,
2773 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
2775 1
2776 }
2777 Condition::FmContains { .. }
2779 | Condition::FmContainsAll { .. }
2780 | Condition::Ann { .. }
2781 | Condition::SparseMatch { .. } => 2,
2782 }
2783}
2784
2785impl Table {
2786 pub(crate) fn build_secondary_index_artifact<F>(
2792 &self,
2793 definition: &IndexDef,
2794 rows: &[Row],
2795 mut checkpoint: F,
2796 ) -> Result<SecondaryIndexArtifact>
2797 where
2798 F: FnMut(usize, usize) -> Result<()>,
2799 {
2800 let mut build_schema = self.schema.clone();
2801 build_schema.indexes = vec![definition.clone()];
2802 let (mut bitmap, mut ann, mut fm, mut sparse, mut minhash) = empty_indexes(&build_schema);
2803 let name_to_id: HashMap<&str, u16> = self
2804 .schema
2805 .columns
2806 .iter()
2807 .map(|column| (column.name.as_str(), column.id))
2808 .collect();
2809 let mut hot = HotIndex::new();
2810 let mut accepted = Vec::with_capacity(rows.len());
2811
2812 for (offset, row) in rows.iter().enumerate() {
2813 checkpoint(offset, rows.len())?;
2814 if row.deleted {
2815 continue;
2816 }
2817 if let Some(predicate) = &definition.predicate {
2818 let columns: HashMap<u16, &Value> =
2819 row.columns.iter().map(|(id, value)| (*id, value)).collect();
2820 if !eval_partial_predicate(predicate, &columns, &name_to_id) {
2821 continue;
2822 }
2823 }
2824 accepted.push(row);
2825 if definition.kind == IndexKind::LearnedRange {
2826 continue;
2827 }
2828 let effective = if self.column_keys.is_empty() {
2829 row.clone()
2830 } else {
2831 self.tokenized_for_indexes(row)
2832 };
2833 if definition.kind == IndexKind::Ann {
2834 if let (Some(index), Some(vector)) = (
2835 ann.get_mut(&definition.column_id),
2836 effective
2837 .columns
2838 .get(&definition.column_id)
2839 .and_then(Value::as_embedding),
2840 ) {
2841 index.insert_validated_with_checkpoint(vector, effective.row_id, || {
2842 checkpoint(offset, rows.len())
2843 })?;
2844 }
2845 } else {
2846 index_into_single(
2847 definition,
2848 &build_schema,
2849 &effective,
2850 &mut hot,
2851 &mut bitmap,
2852 &mut ann,
2853 &mut fm,
2854 &mut sparse,
2855 &mut minhash,
2856 );
2857 }
2858 }
2859 checkpoint(rows.len(), rows.len())?;
2860
2861 if let Some(index) = ann.get_mut(&definition.column_id) {
2862 index.seal_with_checkpoint(|| checkpoint(rows.len(), rows.len()))?;
2863 }
2864
2865 let column_id = definition.column_id;
2866 match definition.kind {
2867 IndexKind::Bitmap => bitmap
2868 .remove(&column_id)
2869 .map(|index| SecondaryIndexArtifact::Bitmap(column_id, index)),
2870 IndexKind::Ann => ann
2871 .remove(&column_id)
2872 .map(|index| SecondaryIndexArtifact::Ann(column_id, Box::new(index))),
2873 IndexKind::FmIndex => fm
2874 .remove(&column_id)
2875 .map(|index| SecondaryIndexArtifact::Fm(column_id, Box::new(index))),
2876 IndexKind::Sparse => sparse
2877 .remove(&column_id)
2878 .map(|index| SecondaryIndexArtifact::Sparse(column_id, index)),
2879 IndexKind::MinHash => minhash
2880 .remove(&column_id)
2881 .map(|index| SecondaryIndexArtifact::MinHash(column_id, index)),
2882 IndexKind::LearnedRange => {
2883 let epsilon = definition
2884 .options
2885 .learned_range
2886 .as_ref()
2887 .map(|options| options.epsilon)
2888 .unwrap_or(16);
2889 let column = self
2890 .schema
2891 .columns
2892 .iter()
2893 .find(|column| column.id == column_id)
2894 .ok_or_else(|| {
2895 MongrelError::Schema(format!(
2896 "index {} references unknown column {column_id}",
2897 definition.name
2898 ))
2899 })?;
2900 let index = match column.ty {
2901 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
2902 let pairs: Vec<_> = accepted
2903 .iter()
2904 .filter_map(|row| match row.columns.get(&column_id) {
2905 Some(Value::Int64(value)) if !row.deleted => {
2906 Some((*value, row.row_id.0))
2907 }
2908 _ => None,
2909 })
2910 .collect();
2911 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon)
2912 }
2913 TypeId::Float32 | TypeId::Float64 => {
2914 let pairs: Vec<_> = accepted
2915 .iter()
2916 .filter_map(|row| match row.columns.get(&column_id) {
2917 Some(Value::Float64(value)) if !row.deleted => {
2918 Some((*value, row.row_id.0))
2919 }
2920 _ => None,
2921 })
2922 .collect();
2923 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon)
2924 }
2925 ref ty => {
2926 return Err(MongrelError::Schema(format!(
2927 "LearnedRange index {} does not support {ty:?}",
2928 definition.name
2929 )));
2930 }
2931 };
2932 Some(SecondaryIndexArtifact::LearnedRange(column_id, index))
2933 }
2934 }
2935 .ok_or_else(|| {
2936 MongrelError::Other(format!(
2937 "failed to construct hidden index {}",
2938 definition.name
2939 ))
2940 })
2941 }
2942
2943 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
2944 let dir = dir.as_ref().to_path_buf();
2945 std::fs::create_dir_all(&dir)?;
2948 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2949 let pinned = root.io_path()?;
2950 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
2951 ctx.root_guard = Some(root.clone());
2952 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2953 root.finalize_deferred_sync_shallow()?;
2959 Ok(table)
2960 }
2961
2962 pub fn create_encrypted(
2973 dir: impl AsRef<Path>,
2974 schema: Schema,
2975 table_id: u64,
2976 passphrase: &str,
2977 ) -> Result<Self> {
2978 let dir = dir.as_ref().to_path_buf();
2979 std::fs::create_dir_all(&dir)?;
2980 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
2981 root.create_directory_all(META_DIR)?;
2982 let salt = crate::encryption::random_salt()?;
2983 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
2984 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
2985 let pinned = root.io_path()?;
2986 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
2987 ctx.root_guard = Some(root.clone());
2988 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
2989 root.finalize_deferred_sync()?;
2990 Ok(table)
2991 }
2992
2993 pub fn create_with_key(
2998 dir: impl AsRef<Path>,
2999 schema: Schema,
3000 table_id: u64,
3001 key: &[u8],
3002 ) -> Result<Self> {
3003 let dir = dir.as_ref().to_path_buf();
3004 std::fs::create_dir_all(&dir)?;
3005 let root = Arc::new(crate::durable_file::DurableRoot::open_deferred(&dir)?);
3006 root.create_directory_all(META_DIR)?;
3007 let salt = crate::encryption::random_salt()?;
3008 root.write_atomic(Path::new(META_DIR).join(KEYS_FILENAME), &salt)?;
3009 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
3010 let pinned = root.io_path()?;
3011 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3012 ctx.root_guard = Some(root.clone());
3013 let table = Self::create_in(&pinned, schema, table_id, ctx)?;
3014 root.finalize_deferred_sync()?;
3015 Ok(table)
3016 }
3017
3018 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
3020 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3021 let salt = read_table_encryption_salt_root(&root)?;
3022 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
3023 let pinned = root.io_path()?;
3024 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3025 ctx.root_guard = Some(root);
3026 Self::open_in(&pinned, ctx)
3027 }
3028
3029 pub(crate) fn create_in(
3030 dir: impl AsRef<Path>,
3031 schema: Schema,
3032 table_id: u64,
3033 ctx: SharedCtx,
3034 ) -> Result<Self> {
3035 schema.validate_auto_increment()?;
3036 schema.validate_defaults()?;
3037 schema.validate_ai()?;
3038 for index in &schema.indexes {
3039 index.validate_options()?;
3040 }
3041 let dir = dir.as_ref().to_path_buf();
3042 let runs_root = match ctx.root_guard.as_ref() {
3043 Some(root) => Some(Arc::new(root.create_directory_all_pinned(RUNS_DIR)?)),
3044 None => {
3045 crate::durable_file::create_directory_all(&dir)?;
3046 crate::durable_file::create_directory_all(&dir.join(RUNS_DIR))?;
3047 None
3048 }
3049 };
3050 match ctx.root_guard.as_deref() {
3051 Some(root) => write_schema_durable(root, &schema)?,
3052 None => write_schema(&dir, &schema)?,
3053 }
3054 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
3055 let (wal, current_txn_id) = match ctx.shared.clone() {
3058 Some(s) => (WalSink::Shared(s), 0),
3059 None => {
3060 let pinned_wal_root = match ctx.root_guard.as_deref() {
3061 Some(root) => Some(root.create_directory_all_pinned(WAL_DIR)?),
3062 None => None,
3063 };
3064 let wal_dir = if let Some(root) = pinned_wal_root.as_ref() {
3065 root.io_path()?
3066 } else {
3067 let wal_dir = dir.join(WAL_DIR);
3068 std::fs::create_dir_all(&wal_dir)?;
3069 wal_dir
3070 };
3071 let mut w = match (pinned_wal_root.as_ref(), wal_dek.as_ref()) {
3072 (Some(root), Some(dk)) => {
3073 Wal::create_in_root(root, 0, Epoch(0), Some(make_cipher(dk)))?
3074 }
3075 (Some(root), None) => Wal::create_in_root(root, 0, Epoch(0), None)?,
3076 (None, Some(dk)) => Wal::create_with_cipher(
3077 wal_dir.join("seg-000000.wal"),
3078 Epoch(0),
3079 Some(make_cipher(dk)),
3080 0,
3081 )?,
3082 (None, None) => Wal::create(wal_dir.join("seg-000000.wal"), Epoch(0))?,
3083 };
3084 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3085 (WalSink::Private(w), 1)
3086 }
3087 };
3088 let mut manifest = Manifest::new(table_id, schema.schema_id);
3089 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3094 match ctx.root_guard.as_deref() {
3095 Some(root) => manifest::write_durable(root, &mut manifest, manifest_meta_dek.as_ref())?,
3096 None => manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?,
3097 }
3098 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
3099 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3100 let auto_inc = resolve_auto_inc(&schema);
3101 let rcache_dir = dir.join(RCACHE_DIR);
3102 let initial_view = ReadGeneration::empty(&schema);
3103 Ok(Self {
3104 dir,
3105 _root_guard: ctx.root_guard,
3106 runs_root,
3107 idx_root: None,
3108 table_id,
3109 name: ctx.table_name.unwrap_or_default(),
3110 auth: ctx.auth,
3111 read_only: ctx.read_only,
3112 durable_commit_failed: false,
3113 wal,
3114 memtable: Memtable::new(),
3115 mutable_run: MutableRun::new(),
3116 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3117 compaction_zstd_level: 3,
3118 allocator: RowIdAllocator::new(0),
3119 epoch: ctx.epoch,
3120 data_generation: 0,
3121 schema,
3122 hot: HotIndex::new(),
3123 kek: ctx.kek,
3124 column_keys,
3125 run_refs: Vec::new(),
3126 retiring: Vec::new(),
3127 next_run_id: 1,
3128 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3129 current_txn_id,
3130 pending_private_mutations: false,
3131 bitmap,
3132 ann,
3133 fm,
3134 sparse,
3135 minhash,
3136 learned_range: Arc::new(HashMap::new()),
3137 pk_by_row: ReversePkMap::new(),
3138 pinned: BTreeMap::new(),
3139 live_count: 0,
3140 reservoir: crate::reservoir::Reservoir::default(),
3141 reservoir_complete: true,
3142 had_deletes: false,
3143 recent_delete_preimages: HashMap::new(),
3144 run_row_id_ranges: HashMap::new(),
3145 agg_cache: Arc::new(HashMap::new()),
3146 global_idx_epoch: 0,
3147 indexes_complete: true,
3148 index_build_policy: IndexBuildPolicy::default(),
3149 pk_by_row_complete: false,
3150 flushed_epoch: 0,
3151 page_cache: ctx.page_cache,
3152 decoded_cache: ctx.decoded_cache,
3153 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3154 snapshots: ctx.snapshots,
3155 commit_lock: ctx.commit_lock,
3156 result_cache: {
3161 let cache = ResultCache::new()
3162 .with_dir(rcache_dir.clone())
3163 .with_cache_dek(cache_dek.clone());
3164 let metrics = Arc::new(LookupMetrics::default());
3168 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3169 if let Ok((writer, handle)) =
3170 spawn_persistent_cache_worker(rcache_dir.clone(), cache_dek.clone(), (*metrics).clone())
3171 {
3172 cache_arc.lock().install_persistent_writer(writer, handle);
3173 }
3174 cache_arc
3175 },
3176 pending_delete_rids: roaring::RoaringBitmap::new(),
3177 pending_put_cols: std::collections::HashSet::new(),
3178 pending_rows: Vec::new(),
3179 pending_rows_auto_inc: Vec::new(),
3180 pending_dels: Vec::new(),
3181 pending_truncate: None,
3182 wal_dek,
3183 auto_inc,
3184 ttl: None,
3185 pins: Arc::new(crate::retention::PinRegistry::new()),
3186 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3187 read_generation_pin: None,
3188 lookup_metrics: LookupMetrics::default(),
3189 run_lookup: RunLookupState::default(),
3190 })
3191 }
3192
3193 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
3197 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3198 let pinned = root.io_path()?;
3199 let mut ctx = SharedCtx::new(None, Some(pinned.join(CACHE_DIR)));
3200 ctx.root_guard = Some(root);
3201 Self::open_in(&pinned, ctx)
3202 }
3203
3204 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
3207 let root = Arc::new(crate::durable_file::DurableRoot::open(dir.as_ref())?);
3208 let salt = read_table_encryption_salt_root(&root)?;
3209 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
3210 let pinned = root.io_path()?;
3211 let mut ctx = SharedCtx::new(Some(kek), Some(pinned.join(CACHE_DIR)));
3212 ctx.root_guard = Some(root);
3213 let t = Self::open_in(&pinned, ctx)?;
3214 Ok(t)
3215 }
3216
3217 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
3218 let dir = dir.as_ref().to_path_buf();
3219 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
3220 let mut manifest = match ctx.root_guard.as_ref() {
3221 Some(root) => manifest::read_durable(root, "", manifest_meta_dek.as_ref())?,
3222 None => manifest::read(&dir, manifest_meta_dek.as_ref())?,
3223 };
3224 let schema: Schema = match ctx.root_guard.as_ref() {
3225 Some(root) => read_schema_file(root.open_regular(SCHEMA_FILENAME)?)?,
3226 None => read_schema(&dir)?,
3227 };
3228 let schema_manifest_repair = manifest.schema_id < schema.schema_id;
3234 let runs_root = match ctx.root_guard.as_ref() {
3235 Some(root) => Some(Arc::new(root.open_directory(RUNS_DIR)?)),
3236 None => None,
3237 };
3238 let idx_root = match ctx.root_guard.as_ref() {
3239 Some(root) => match root.open_directory(global_idx::IDX_DIR) {
3240 Ok(root) => Some(Arc::new(root)),
3241 Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
3242 Err(error) => return Err(error.into()),
3243 },
3244 None => None,
3245 };
3246 schema.validate_auto_increment()?;
3247 schema.validate_defaults()?;
3248 schema.validate_ai()?;
3249 for index in &schema.indexes {
3250 index.validate_options()?;
3251 }
3252 let replay_epoch = Epoch(manifest.current_epoch);
3253 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
3254 let private_replayed = if ctx.shared.is_none() {
3255 match latest_wal_segment(&dir.join(WAL_DIR))? {
3256 Some(path) => {
3257 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
3258 crate::wal::replay_with_cipher(path, cipher)?
3259 }
3260 None => Vec::new(),
3261 }
3262 } else {
3263 Vec::new()
3264 };
3265 if ctx.shared.is_none() {
3266 preflight_standalone_open(
3267 &dir,
3268 runs_root.as_deref(),
3269 idx_root.as_deref(),
3270 &manifest,
3271 &schema,
3272 &private_replayed,
3273 ctx.kek.clone(),
3274 )?;
3275 }
3276 let next_run_id = derive_next_run_id(
3277 &dir,
3278 runs_root.as_deref(),
3279 &manifest.runs,
3280 &manifest.retiring,
3281 )?;
3282 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
3286 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
3287 None => {
3288 let replayed = private_replayed;
3289 let wal_dir = dir.join(WAL_DIR);
3295 crate::durable_file::create_directory_all(&wal_dir)?;
3296 let segment = next_wal_segment(&wal_dir)?;
3297 let segment_no = wal_segment_number(&segment).unwrap_or(0);
3298 let temporary = wal_dir.join(format!(
3299 ".recovery-{}-{}-{segment_no:06}.tmp",
3300 std::process::id(),
3301 std::time::SystemTime::now()
3302 .duration_since(std::time::UNIX_EPOCH)
3303 .unwrap_or_default()
3304 .as_nanos()
3305 ));
3306 let mut w = Wal::create_with_cipher(
3307 &temporary,
3308 replay_epoch,
3309 wal_dek.as_ref().map(|dk| make_cipher(dk)),
3310 segment_no,
3311 )?;
3312 for record in &replayed {
3313 w.append_txn(record.txn_id, record.op.clone())?;
3314 }
3315 let mut w = w.publish_as(segment)?;
3316 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
3317 let next_txn_id = replayed
3318 .iter()
3319 .map(|record| record.txn_id)
3320 .filter(|txn_id| *txn_id != crate::wal::SYSTEM_TXN_ID)
3321 .max()
3322 .map(|txn_id| txn_id.checked_add(1).unwrap_or(0))
3323 .unwrap_or(1);
3324 (WalSink::Private(w), replayed, next_txn_id)
3325 }
3326 };
3327
3328 let mut memtable = Memtable::new();
3329 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
3330 let persisted_epoch = manifest.current_epoch;
3331 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
3338 s.next = manifest.auto_inc_next;
3339 s.seeded = manifest.auto_inc_next > 0;
3340 s
3341 });
3342
3343 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
3350 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
3351 let mut staged_truncates: std::collections::HashSet<u64> = std::collections::HashSet::new();
3352 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
3353 std::collections::BTreeMap::new();
3354 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
3355 let mut recovered_epoch = manifest.current_epoch;
3356 let mut recovered_manifest_dirty = schema_manifest_repair;
3357 let mut saw_delete = false;
3358 for record in replayed {
3359 let txn_id = record.txn_id;
3360 match record.op {
3361 Op::Put { rows, .. } => {
3362 let rows: Vec<Row> = bincode::deserialize(&rows)?;
3363 for row in &rows {
3364 allocator.advance_to(row.row_id)?;
3365 if let Some(ai) = auto_inc.as_mut() {
3366 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
3367 let next = n.checked_add(1).ok_or_else(|| {
3368 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
3369 })?;
3370 if next > ai.next {
3371 ai.next = next;
3372 }
3373 }
3374 }
3375 }
3376 staged_puts.entry(txn_id).or_default().extend(rows);
3377 }
3378 Op::Delete { row_ids, .. } => {
3379 staged_deletes.entry(txn_id).or_default().extend(row_ids);
3380 }
3381 Op::TxnCommit { epoch, .. } => {
3382 let commit_epoch = Epoch(epoch);
3383 recovered_epoch = recovered_epoch.max(epoch);
3384 if staged_truncates.remove(&txn_id) && commit_epoch.0 > manifest.flushed_epoch {
3385 memtable = Memtable::new();
3386 replayed_puts.clear();
3387 replayed_deletes.clear();
3388 manifest.runs.clear();
3389 manifest.retiring.clear();
3390 manifest.live_count = 0;
3391 manifest.global_idx_epoch = 0;
3392 manifest.current_epoch = manifest.current_epoch.max(epoch);
3393 recovered_manifest_dirty = true;
3394 saw_delete = true;
3395 }
3396 if let Some(puts) = staged_puts.remove(&txn_id) {
3397 if commit_epoch.0 > manifest.flushed_epoch {
3398 for row in &puts {
3399 memtable.upsert(row.clone());
3400 }
3401 replayed_puts.entry(commit_epoch).or_default().extend(puts);
3402 }
3403 }
3404 if let Some(dels) = staged_deletes.remove(&txn_id) {
3405 saw_delete = true;
3406 if commit_epoch.0 > manifest.flushed_epoch {
3407 for rid in dels {
3408 memtable.tombstone(rid, commit_epoch);
3409 replayed_deletes.push((rid, commit_epoch));
3410 }
3411 }
3412 }
3413 }
3414 Op::TxnAbort => {
3415 staged_puts.remove(&txn_id);
3416 staged_deletes.remove(&txn_id);
3417 staged_truncates.remove(&txn_id);
3418 }
3419 Op::TruncateTable { .. } => {
3420 staged_puts.remove(&txn_id);
3421 staged_deletes.remove(&txn_id);
3422 staged_truncates.insert(txn_id);
3423 }
3424 Op::ExternalTableState { .. }
3425 | Op::Flush { .. }
3426 | Op::Ddl(_)
3427 | Op::BeforeImage { .. }
3428 | Op::CommitTimestamp { .. }
3429 | Op::SpilledRows { .. } => {}
3430 }
3431 }
3432
3433 let rcache_dir = dir.join(RCACHE_DIR);
3434 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
3435 let initial_view = ReadGeneration::empty(&schema);
3436 let mut db = Self {
3437 dir,
3438 _root_guard: ctx.root_guard,
3439 runs_root,
3440 idx_root,
3441 table_id: manifest.table_id,
3442 name: ctx.table_name.unwrap_or_default(),
3443 auth: ctx.auth,
3444 read_only: ctx.read_only,
3445 durable_commit_failed: false,
3446 wal,
3447 memtable,
3448 mutable_run: MutableRun::new(),
3449 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
3450 compaction_zstd_level: 3,
3451 allocator,
3452 epoch: ctx.epoch,
3453 data_generation: persisted_epoch,
3454 schema,
3455 hot: HotIndex::new(),
3456 kek: ctx.kek,
3457 column_keys,
3458 run_refs: manifest.runs.clone(),
3459 retiring: manifest.retiring.clone(),
3460 next_run_id,
3461 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
3462 current_txn_id,
3463 pending_private_mutations: false,
3464 bitmap: HashMap::new(),
3465 ann: HashMap::new(),
3466 fm: HashMap::new(),
3467 sparse: HashMap::new(),
3468 minhash: HashMap::new(),
3469 learned_range: Arc::new(HashMap::new()),
3470 pk_by_row: ReversePkMap::new(),
3471 pinned: BTreeMap::new(),
3472 live_count: manifest.live_count,
3473 reservoir: crate::reservoir::Reservoir::default(),
3474 reservoir_complete: false,
3475 had_deletes: saw_delete
3476 || manifest.runs.iter().map(|run| run.row_count).sum::<u64>()
3477 != manifest.live_count,
3478 recent_delete_preimages: HashMap::new(),
3479 run_row_id_ranges: HashMap::new(),
3480 agg_cache: Arc::new(HashMap::new()),
3481 global_idx_epoch: manifest.global_idx_epoch,
3482 indexes_complete: true,
3483 index_build_policy: IndexBuildPolicy::default(),
3484 pk_by_row_complete: false,
3485 flushed_epoch: manifest.flushed_epoch,
3486 page_cache: ctx.page_cache,
3487 decoded_cache: ctx.decoded_cache,
3488 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
3489 snapshots: ctx.snapshots,
3490 commit_lock: ctx.commit_lock,
3491 result_cache: {
3492 let cache = ResultCache::new()
3493 .with_dir(rcache_dir.clone())
3494 .with_cache_dek(cache_dek.clone());
3495 let metrics = LookupMetrics::default();
3496 let cache_arc = Arc::new(parking_lot::Mutex::new(cache));
3497 if let Ok((writer, handle)) =
3498 spawn_persistent_cache_worker(rcache_dir.clone(), cache_dek.clone(), metrics)
3499 {
3500 cache_arc.lock().install_persistent_writer(writer, handle);
3501 }
3502 cache_arc
3503 },
3504 pending_delete_rids: roaring::RoaringBitmap::new(),
3505 pending_put_cols: std::collections::HashSet::new(),
3506 pending_rows: Vec::new(),
3507 pending_rows_auto_inc: Vec::new(),
3508 pending_dels: Vec::new(),
3509 pending_truncate: None,
3510 wal_dek,
3511 auto_inc,
3512 ttl: manifest.ttl,
3513 pins: Arc::new(crate::retention::PinRegistry::new()),
3514 published: Arc::new(ArcSwap::from_pointee(initial_view)),
3515 read_generation_pin: None,
3516 lookup_metrics: LookupMetrics::default(),
3517 run_lookup: RunLookupState::default(),
3518 };
3519
3520 db.epoch.advance_recovered(Epoch(recovered_epoch));
3523
3524 let checkpoint = match db.idx_root.as_deref() {
3529 Some(root) => {
3530 global_idx::read_root(root, db.table_id, &db.schema, db.idx_dek().as_deref())?
3531 }
3532 None => global_idx::read(&db.dir, db.table_id, &db.schema, db.idx_dek().as_deref())?,
3533 };
3534 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
3535 c.epoch_built == manifest.global_idx_epoch
3536 && manifest.global_idx_epoch > 0
3537 && manifest
3538 .runs
3539 .iter()
3540 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
3541 });
3542 if let Some(loaded) = checkpoint {
3543 if checkpoint_valid {
3544 db.hot = loaded.hot;
3545 db.bitmap = loaded.bitmap;
3546 db.ann = loaded.ann;
3547 db.fm = loaded.fm;
3548 db.sparse = loaded.sparse;
3549 db.minhash = loaded.minhash;
3550 db.learned_range = Arc::new(loaded.learned_range);
3551 let (bitmap0, ann0, fm0, sparse0, minhash0) = empty_indexes(&db.schema);
3556 for (cid, idx) in bitmap0 {
3557 db.bitmap.entry(cid).or_insert(idx);
3558 }
3559 for (cid, idx) in ann0 {
3560 db.ann.entry(cid).or_insert(idx);
3561 }
3562 for (cid, idx) in fm0 {
3563 db.fm.entry(cid).or_insert(idx);
3564 }
3565 for (cid, idx) in sparse0 {
3566 db.sparse.entry(cid).or_insert(idx);
3567 }
3568 for (cid, idx) in minhash0 {
3569 db.minhash.entry(cid).or_insert(idx);
3570 }
3571 }
3574 }
3575 if !checkpoint_valid {
3576 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
3577 db.bitmap = bitmap;
3578 db.ann = ann;
3579 db.fm = fm;
3580 db.sparse = sparse;
3581 db.minhash = minhash;
3582 db.rebuild_indexes_from_runs()?;
3583 db.build_learned_ranges()?;
3584 }
3585
3586 for (epoch, group) in replayed_puts {
3591 let (losers, winner_pks) = db.partition_pk_winners(&group);
3592 for (key, &row_id) in &winner_pks {
3593 if let Some(old_rid) = db.hot.get(key) {
3594 if old_rid != row_id {
3595 db.tombstone_row(old_rid, epoch, None, false);
3596 }
3597 }
3598 }
3599 for &loser_rid in &losers {
3600 db.tombstone_row(loser_rid, epoch, None, false);
3601 }
3602 for (key, row_id) in winner_pks {
3603 db.insert_hot_pk(key, row_id);
3604 }
3605 if db.schema.primary_key().is_none() {
3606 for r in &group {
3607 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
3608 }
3609 }
3610 for r in &group {
3611 if !losers.contains(&r.row_id) {
3612 db.index_row(r);
3613 }
3614 }
3615 }
3616 for (rid, epoch) in &replayed_deletes {
3620 db.remove_hot_for_row(*rid, *epoch);
3621 }
3622
3623 if recovered_manifest_dirty {
3624 let rows = db.visible_rows(Snapshot::unbounded())?;
3625 db.live_count = rows.len() as u64;
3626 db.persist_manifest(Epoch(recovered_epoch))?;
3627 }
3628
3629 db.result_cache.lock().load_persistent();
3636 db.refresh_run_row_id_ranges();
3638 db.load_run_lookup_directory();
3643 Ok(db)
3644 }
3645
3646 fn load_run_lookup_directory(&mut self) {
3651 let active_runs = self.run_refs.clone();
3652 let schema_id = self.schema.schema_id;
3653 let run_generation = active_runs.len() as u64;
3654 let index_generation = self.global_idx_epoch;
3655 let fingerprint = crate::run_lookup::compute_fingerprint(
3656 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
3657 schema_id,
3658 index_generation,
3659 run_generation,
3660 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
3661 );
3662 let runs_dir = match self.runs_root.as_deref() {
3664 Some(root) => match root.io_path() {
3665 Ok(p) => p,
3666 Err(_) => self.dir.join(RUNS_DIR),
3667 },
3668 None => self.dir.join(RUNS_DIR),
3669 };
3670 match crate::run_lookup::RunLookupDirectory::read_checkpoint_sharded(&runs_dir, fingerprint)
3671 {
3672 Ok(dir) => {
3673 self.run_lookup = RunLookupState {
3674 directory: Some(std::sync::Arc::new(dir)),
3675 fingerprint,
3676 complete: true,
3677 };
3678 }
3679 Err(error) => {
3680 self.lookup_metrics
3681 .directory_incomplete
3682 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3683 self.lookup_metrics
3684 .directory_lookup_fallback
3685 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3686 self.run_lookup = RunLookupState {
3687 directory: None,
3688 fingerprint,
3689 complete: false,
3690 };
3691 let _ = error;
3695 }
3696 }
3697 }
3698
3699 fn ensure_reservoir_complete(&mut self) -> Result<()> {
3705 if self.reservoir_complete {
3706 return Ok(());
3707 }
3708 self.rebuild_reservoir()?;
3709 self.reservoir_complete = true;
3710 Ok(())
3711 }
3712
3713 fn rebuild_reservoir(&mut self) -> Result<()> {
3716 let snap = self.snapshot();
3717 let rows = self.visible_rows(snap)?;
3718 self.reservoir.reset();
3719 for r in rows {
3720 self.reservoir.offer(r.row_id.0);
3721 }
3722 Ok(())
3723 }
3724
3725 pub fn rebuild_indexes(&mut self) -> Result<()> {
3729 self.rebuild_indexes_from_runs_inner(None)
3730 }
3731
3732 pub fn __force_hot_map_for_test(&mut self, pk_bytes: &[u8], row_id: RowId) {
3743 self.hot.insert(pk_bytes.to_vec(), row_id);
3744 }
3745
3746 pub fn hot_for_test_remove(&mut self, pk_bytes: &[u8]) {
3749 self.hot.remove(pk_bytes);
3750 }
3751
3752 pub fn set_indexes_incomplete_for_test(&mut self) {
3755 self.indexes_complete = false;
3756 }
3757
3758 pub fn bump_hot_checkpoint_rejected_for_test(&self) {
3762 self.lookup_metrics
3763 .hot_checkpoint_rejected_total
3764 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3765 }
3766
3767 pub(crate) fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
3768 self.rebuild_indexes_from_runs_inner(None)
3769 }
3770
3771 fn pk_equality_fallback(
3798 &self,
3799 pk_column_id: u16,
3800 lookup: &[u8],
3801 snapshot: Snapshot,
3802 ) -> Result<(RowIdSet, crate::trace::HotFallbackReason)> {
3803 let mut tombstone_hit = false;
3804 let mut overlay_versions = 0u64;
3805 let now_nanos = unix_nanos_now();
3806 for row in self.memtable.visible_versions_at(snapshot) {
3808 overlay_versions += 1;
3809 self.lookup_metrics
3810 .hot_lookup_fallback_overlay_rows
3811 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3812 if row.deleted {
3813 tombstone_hit = true;
3814 continue;
3815 }
3816 if self.row_expired_at(&row, now_nanos) {
3817 continue;
3818 }
3819 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3820 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3821 self.lookup_metrics
3822 .hot_fallback_overlay_versions_total
3823 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3824 return Ok((
3825 RowIdSet::one(row.row_id.0),
3826 crate::trace::HotFallbackReason::MissingMapping,
3827 ));
3828 }
3829 }
3830 }
3831 for row in self.mutable_run.visible_versions_at(snapshot) {
3832 overlay_versions += 1;
3833 self.lookup_metrics
3834 .hot_lookup_fallback_overlay_rows
3835 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3836 if row.deleted {
3837 tombstone_hit = true;
3838 continue;
3839 }
3840 if self.row_expired_at(&row, now_nanos) {
3841 continue;
3842 }
3843 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3844 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3845 self.lookup_metrics
3846 .hot_fallback_overlay_versions_total
3847 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3848 return Ok((
3849 RowIdSet::one(row.row_id.0),
3850 crate::trace::HotFallbackReason::MissingMapping,
3851 ));
3852 }
3853 }
3854 }
3855 if lookup.len() == 8 {
3858 if let Ok(arr) = <[u8; 8]>::try_from(lookup) {
3859 let n = i64::from_be_bytes(arr);
3860 let result = self.range_scan_i64(pk_column_id, n, n, snapshot)?;
3861 self.lookup_metrics
3862 .hot_lookup_fallback_runs
3863 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3864 self.lookup_metrics
3868 .hot_fallback_runs_considered_total
3869 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3870 self.lookup_metrics
3871 .hot_fallback_overlay_versions_total
3872 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3873 let reason = if result.is_empty() {
3874 if tombstone_hit {
3875 crate::trace::HotFallbackReason::Tombstone
3876 } else {
3877 crate::trace::HotFallbackReason::MissingMapping
3878 }
3879 } else {
3880 crate::trace::HotFallbackReason::MissingMapping
3881 };
3882 return Ok((result, reason));
3883 }
3884 }
3885 let mut found: std::collections::BTreeSet<u64> = std::collections::BTreeSet::new();
3888 let overlay = self.overlay_rid_set(snapshot);
3889 let mut runs_considered = 0u64;
3890 for rr in &self.run_refs {
3891 runs_considered += 1;
3892 self.lookup_metrics
3893 .hot_lookup_fallback_runs
3894 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3895 let mut reader = self.open_reader(rr.run_id)?;
3896 for row in reader.visible_versions_at(snapshot)? {
3897 if overlay.contains(&row.row_id.0) || row.deleted {
3898 if row.deleted {
3899 tombstone_hit = true;
3900 }
3901 continue;
3902 }
3903 if self.row_expired_at(&row, now_nanos) {
3904 continue;
3905 }
3906 if let Some(pk_val) = row.columns.get(&pk_column_id) {
3907 if self.index_lookup_key(pk_column_id, pk_val) == lookup {
3908 found.insert(row.row_id.0);
3909 }
3910 }
3911 }
3912 }
3913 let reason = if tombstone_hit {
3914 crate::trace::HotFallbackReason::Tombstone
3915 } else {
3916 crate::trace::HotFallbackReason::MissingMapping
3919 };
3920 self.lookup_metrics
3921 .hot_fallback_overlay_versions_total
3922 .fetch_add(overlay_versions, std::sync::atomic::Ordering::Relaxed);
3923 self.lookup_metrics
3924 .hot_fallback_runs_considered_total
3925 .fetch_add(runs_considered, std::sync::atomic::Ordering::Relaxed);
3926 Ok((RowIdSet::from_unsorted(found.into_iter().collect()), reason))
3927 }
3928
3929 fn record_hot_fallback_reason(&self, reason: crate::trace::HotFallbackReason) {
3933 let idx = hot_fallback_reason_index(reason);
3934 self.lookup_metrics.hot_fallback_reasons[idx]
3935 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3936 self.lookup_metrics
3937 .hot_lookup_fallback
3938 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
3939 crate::trace::QueryTrace::record(|t| {
3940 t.hot_lookup_attempted = true;
3941 t.hot_lookup_hit = false;
3942 t.hot_fallback_reason = Some(reason.as_str());
3943 });
3944 }
3945
3946 fn rebuild_indexes_from_runs_inner(
3947 &mut self,
3948 control: Option<&crate::ExecutionControl>,
3949 ) -> Result<()> {
3950 let _index_build_pin = Arc::clone(self.pin_registry()).pin(
3953 crate::retention::PinSource::OnlineIndexBuild,
3954 self.current_epoch(),
3955 );
3956 self.hot = HotIndex::new();
3957 self.pk_by_row.clear();
3958 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
3959 self.bitmap = bitmap;
3960 self.ann = ann;
3961 self.fm = fm;
3962 self.sparse = sparse;
3963 self.minhash = minhash;
3964 let snapshot = Epoch(u64::MAX);
3965 let ttl_now = unix_nanos_now();
3966 let mut scanned = 0_usize;
3967 for rr in self.run_refs.clone() {
3968 if let Some(control) = control {
3969 control.checkpoint()?;
3970 }
3971 let mut reader = self.open_reader(rr.run_id)?;
3972 for row in reader.visible_versions_at(Snapshot::at(snapshot))? {
3973 if row.deleted {
3974 continue;
3975 }
3976 if scanned.is_multiple_of(256) {
3977 if let Some(control) = control {
3978 control.checkpoint()?;
3979 }
3980 }
3981 scanned += 1;
3982 if self.row_expired_at(&row, ttl_now) {
3983 continue;
3984 }
3985 let tok_row = self.tokenized_for_indexes(&row);
3986 index_into(
3987 &self.schema,
3988 &tok_row,
3989 &mut self.hot,
3990 &mut self.bitmap,
3991 &mut self.ann,
3992 &mut self.fm,
3993 &mut self.sparse,
3994 &mut self.minhash,
3995 );
3996 }
3997 }
3998 for row in self.mutable_run.visible_versions(snapshot) {
3999 if scanned.is_multiple_of(256) {
4000 if let Some(control) = control {
4001 control.checkpoint()?;
4002 }
4003 }
4004 scanned += 1;
4005 if row.deleted {
4006 self.remove_hot_for_row(row.row_id, snapshot);
4007 } else if !self.row_expired_at(&row, ttl_now) {
4008 self.index_row(&row);
4009 }
4010 }
4011 for row in self.memtable.visible_versions(snapshot) {
4012 if scanned.is_multiple_of(256) {
4013 if let Some(control) = control {
4014 control.checkpoint()?;
4015 }
4016 }
4017 scanned += 1;
4018 if row.deleted {
4019 self.remove_hot_for_row(row.row_id, snapshot);
4020 } else if !self.row_expired_at(&row, ttl_now) {
4021 self.index_row(&row);
4022 }
4023 }
4024 let pin_epochs = self.active_pin_epochs_for_rebuild();
4030 if !pin_epochs.is_empty() {
4031 let current_snap = Snapshot::at(Epoch(u64::MAX));
4032 for pin_epoch in pin_epochs {
4033 if let Some(control) = control {
4034 control.checkpoint()?;
4035 }
4036 let pin_snap = Snapshot::at(pin_epoch);
4037 for rr in self.run_refs.clone() {
4038 if let Some(control) = control {
4039 control.checkpoint()?;
4040 }
4041 let mut reader = self.open_reader(rr.run_id)?;
4042 for row in reader.visible_rows(pin_epoch)? {
4043 if row.deleted || self.row_expired_at(&row, ttl_now) {
4044 continue;
4045 }
4046 if self.get(row.row_id, current_snap).is_none() {
4047 self.index_bitmap_membership_only(&row);
4048 }
4049 }
4050 }
4051 for row in self
4052 .mutable_run
4053 .visible_versions_at(pin_snap)
4054 .into_iter()
4055 .chain(self.memtable.visible_versions_at(pin_snap))
4056 {
4057 if row.deleted || self.row_expired_at(&row, ttl_now) {
4058 continue;
4059 }
4060 if self.get(row.row_id, current_snap).is_none() {
4061 self.index_bitmap_membership_only(&row);
4062 }
4063 }
4064 }
4065 }
4066 self.recent_delete_preimages.clear();
4067 self.refresh_pk_by_row_from_hot();
4068 Ok(())
4069 }
4070
4071 fn index_bitmap_membership_only(&mut self, row: &Row) {
4074 if row.deleted {
4075 return;
4076 }
4077 let columns_map: HashMap<u16, &Value> = row.columns.iter().map(|(k, v)| (*k, v)).collect();
4078 let name_to_id: HashMap<&str, u16> = self
4079 .schema
4080 .columns
4081 .iter()
4082 .map(|c| (c.name.as_str(), c.id))
4083 .collect();
4084 for idef in &self.schema.indexes {
4085 if idef.kind != crate::schema::IndexKind::Bitmap {
4086 continue;
4087 }
4088 if let Some(pred) = &idef.predicate {
4089 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
4090 continue;
4091 }
4092 }
4093 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
4094 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4095 b.insert(key, row.row_id);
4096 }
4097 }
4098 }
4099 }
4100
4101 fn refresh_pk_by_row_from_hot(&mut self) {
4102 self.pk_by_row_complete = true;
4103 if self.schema.primary_key().is_none() {
4104 self.pk_by_row.clear();
4105 return;
4106 }
4107 self.pk_by_row = ReversePkMap::from_entries(
4113 self.hot
4114 .entries()
4115 .into_iter()
4116 .map(|(key, row_id)| (row_id, key)),
4117 );
4118 }
4119
4120 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
4121 if self.schema.primary_key().is_some() {
4122 self.pk_by_row.insert(row_id, key.clone());
4123 }
4124 self.hot.insert(key, row_id);
4125 }
4126
4127 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
4131 self.build_learned_ranges_inner(None)
4132 }
4133
4134 fn build_learned_ranges_inner(
4135 &mut self,
4136 control: Option<&crate::ExecutionControl>,
4137 ) -> Result<()> {
4138 self.learned_range = Arc::new(HashMap::new());
4139 if self.run_refs.len() != 1 {
4140 return Ok(());
4141 }
4142 let cols: Vec<(u16, usize)> = self
4143 .schema
4144 .indexes
4145 .iter()
4146 .filter(|i| i.kind == IndexKind::LearnedRange)
4147 .map(|i| {
4148 (
4149 i.column_id,
4150 i.options
4151 .learned_range
4152 .as_ref()
4153 .map(|options| options.epsilon)
4154 .unwrap_or(16),
4155 )
4156 })
4157 .collect();
4158 if cols.is_empty() {
4159 return Ok(());
4160 }
4161 let snapshot_epoch = self.current_epoch();
4168 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4169 let (visible_positions, visible_rids) =
4170 reader.visible_positions_with_rids(snapshot_epoch)?;
4171 let row_ids: Vec<u64> = visible_rids.iter().map(|r| *r as u64).collect();
4172 for (column_index, (cid, epsilon)) in cols.into_iter().enumerate() {
4173 if column_index % 256 == 0 {
4174 if let Some(control) = control {
4175 control.checkpoint()?;
4176 }
4177 }
4178 let ty = self
4179 .schema
4180 .columns
4181 .iter()
4182 .find(|c| c.id == cid)
4183 .map(|c| c.ty.clone())
4184 .unwrap_or(TypeId::Int64);
4185 match ty {
4186 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
4187 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
4188 let pairs: Vec<(i64, u64)> = visible_positions
4189 .iter()
4190 .zip(row_ids.iter())
4191 .map(|(&p, &r)| (data[p], r))
4192 .collect();
4193 Arc::make_mut(&mut self.learned_range).insert(
4194 cid,
4195 ColumnLearnedRange::build_i64_with_epsilon(&pairs, epsilon),
4196 );
4197 }
4198 }
4199 TypeId::Float64 => {
4200 if let columnar::NativeColumn::Float64 { data, .. } =
4201 reader.column_native(cid)?
4202 {
4203 let pairs: Vec<(f64, u64)> = visible_positions
4204 .iter()
4205 .zip(row_ids.iter())
4206 .map(|(&p, &r)| (data[p], r))
4207 .collect();
4208 Arc::make_mut(&mut self.learned_range).insert(
4209 cid,
4210 ColumnLearnedRange::build_f64_with_epsilon(&pairs, epsilon),
4211 );
4212 }
4213 }
4214 _ => {}
4215 }
4216 }
4217 Ok(())
4218 }
4219
4220 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
4227 if self.indexes_complete {
4228 crate::trace::QueryTrace::record(|t| {
4229 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4230 });
4231 return Ok(());
4232 }
4233 crate::trace::QueryTrace::record(|t| {
4234 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4235 });
4236 self.rebuild_indexes_from_runs()?;
4237 self.build_learned_ranges()?;
4238 self.indexes_complete = true;
4239 let epoch = self.current_epoch();
4240 self.checkpoint_indexes(epoch);
4241 Ok(())
4242 }
4243
4244 #[doc(hidden)]
4247 pub fn ensure_indexes_complete_controlled<F>(
4248 &mut self,
4249 control: &crate::ExecutionControl,
4250 before_publish: F,
4251 ) -> Result<bool>
4252 where
4253 F: FnOnce() -> bool,
4254 {
4255 self.ensure_indexes_complete_controlled_with_receipt(control, before_publish)
4256 .map(|(changed, _)| changed)
4257 }
4258
4259 #[doc(hidden)]
4262 pub fn ensure_indexes_complete_controlled_with_receipt<F>(
4263 &mut self,
4264 control: &crate::ExecutionControl,
4265 before_publish: F,
4266 ) -> Result<(bool, Option<MaintenanceReceipt>)>
4267 where
4268 F: FnOnce() -> bool,
4269 {
4270 if self.indexes_complete {
4271 crate::trace::QueryTrace::record(|trace| {
4272 trace.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
4273 });
4274 return Ok((false, None));
4275 }
4276 crate::trace::QueryTrace::record(|trace| {
4277 trace.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
4278 });
4279 control.checkpoint()?;
4280 let maintenance_epoch = self.current_epoch();
4281 self.rebuild_indexes_from_runs_inner(Some(control))?;
4282 self.build_learned_ranges_inner(Some(control))?;
4283 control.checkpoint()?;
4284 if !before_publish() {
4285 return Err(MongrelError::Cancelled);
4286 }
4287 self.indexes_complete = true;
4288 self.checkpoint_indexes(maintenance_epoch);
4289 Ok((
4290 true,
4291 Some(MaintenanceReceipt {
4292 epoch: maintenance_epoch,
4293 }),
4294 ))
4295 }
4296
4297 fn pending_epoch(&self) -> Epoch {
4298 Epoch(self.epoch.visible().0 + 1)
4299 }
4300
4301 fn is_shared(&self) -> bool {
4304 matches!(self.wal, WalSink::Shared(_))
4305 }
4306
4307 fn ensure_txn_id(&mut self) -> Result<u64> {
4311 if self.current_txn_id == 0 {
4312 let id = match &self.wal {
4313 WalSink::Shared(s) => crate::txn::allocate_txn_id(&s.txn_ids)?,
4314 WalSink::Private(_) => {
4315 return Err(MongrelError::Full(
4316 "standalone transaction id namespace exhausted".into(),
4317 ))
4318 }
4319 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4320 };
4321 self.current_txn_id = id;
4322 }
4323 Ok(self.current_txn_id)
4324 }
4325
4326 fn wal_append_data(&mut self, op: Op) -> Result<()> {
4329 self.ensure_writable()?;
4330 let txn_id = self.ensure_txn_id()?;
4331 let table_id = self.table_id;
4332 match &mut self.wal {
4333 WalSink::Private(w) => {
4334 w.append_txn(txn_id, op)?;
4335 self.pending_private_mutations = true;
4336 }
4337 WalSink::Shared(s) => {
4338 s.wal.lock().append(txn_id, table_id, op)?;
4339 }
4340 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
4341 }
4342 Ok(())
4343 }
4344
4345 fn ensure_writable(&self) -> Result<()> {
4346 if self.read_only || matches!(self.wal, WalSink::ReadOnly) {
4347 return Err(MongrelError::ReadOnlyReplica);
4348 }
4349 if self.durable_commit_failed {
4350 return Err(MongrelError::Other(
4351 "table poisoned by post-commit failure; reopen required".into(),
4352 ));
4353 }
4354 Ok(())
4355 }
4356
4357 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
4368 match &self.auth {
4369 Some(checker) => checker.check(&self.name, perm),
4370 None => Ok(()),
4371 }
4372 }
4373 pub fn require_select(&self) -> Result<()> {
4378 self.require(crate::auth_state::RequiredPermission::Select)
4379 }
4380 fn require_insert(&self) -> Result<()> {
4381 self.require(crate::auth_state::RequiredPermission::Insert)
4382 }
4383 #[allow(dead_code)]
4387 fn require_update(&self) -> Result<()> {
4388 self.require(crate::auth_state::RequiredPermission::Update)
4389 }
4390 fn require_delete(&self) -> Result<()> {
4391 self.require(crate::auth_state::RequiredPermission::Delete)
4392 }
4393
4394 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
4397 self.require_insert()?;
4398 Ok(self.put_returning(columns)?.0)
4399 }
4400
4401 pub fn put_returning(
4406 &mut self,
4407 mut columns: Vec<(u16, Value)>,
4408 ) -> Result<(RowId, Option<i64>)> {
4409 self.require_insert()?;
4410 let assigned = self.fill_auto_inc(&mut columns)?;
4411 self.apply_defaults(&mut columns)?;
4412 self.schema.validate_values(&columns)?;
4413 let row_id = if self.schema.clustered {
4418 self.derive_clustered_row_id(&columns)?
4419 } else {
4420 self.allocator.alloc()?
4421 };
4422 let epoch = self.pending_epoch();
4423 let mut row = Row::new(row_id, epoch);
4424 for (col_id, val) in columns {
4425 row.columns.insert(col_id, val);
4426 }
4427 self.commit_rows(vec![row], assigned.is_some())?;
4428 Ok((row_id, assigned))
4429 }
4430
4431 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
4434 self.require_insert()?;
4435 Ok(self
4436 .put_batch_returning(batch)?
4437 .into_iter()
4438 .map(|(r, _)| r)
4439 .collect())
4440 }
4441
4442 pub fn put_batch_returning(
4445 &mut self,
4446 batch: Vec<Vec<(u16, Value)>>,
4447 ) -> Result<Vec<(RowId, Option<i64>)>> {
4448 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
4449 for mut cols in batch {
4450 let assigned = self.fill_auto_inc(&mut cols)?;
4451 self.apply_defaults(&mut cols)?;
4452 filled.push((cols, assigned));
4453 }
4454 for (cols, _) in &filled {
4455 self.schema.validate_values(cols)?;
4456 }
4457 let epoch = self.pending_epoch();
4458 let mut rows = Vec::with_capacity(filled.len());
4459 let mut ids = Vec::with_capacity(filled.len());
4460 let first_row_id = if self.schema.clustered {
4461 None
4462 } else {
4463 let count = u64::try_from(filled.len())
4464 .map_err(|_| MongrelError::Full("row-id allocation request is too large".into()))?;
4465 Some(self.allocator.alloc_range(count)?.0)
4466 };
4467 for (row_index, (cols, assigned)) in filled.into_iter().enumerate() {
4468 let row_id = match first_row_id {
4469 Some(first) => RowId(first + row_index as u64),
4470 None => self.derive_clustered_row_id(&cols)?,
4471 };
4472 let mut row = Row::new(row_id, epoch);
4473 for (c, v) in cols {
4474 row.columns.insert(c, v);
4475 }
4476 ids.push((row_id, assigned));
4477 rows.push(row);
4478 }
4479 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
4480 self.commit_rows(rows, all_auto_generated)?;
4481 Ok(ids)
4482 }
4483
4484 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
4490 self.ensure_writable()?;
4491 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4492 return Ok(None);
4493 };
4494 let pos = columns.iter().position(|(c, _)| *c == cid);
4495 let assigned = match pos {
4496 Some(i) => match &columns[i].1 {
4497 Value::Null => {
4498 let next = self.alloc_auto_inc_value()?;
4499 columns[i].1 = Value::Int64(next);
4500 Some(next)
4501 }
4502 Value::Int64(n) => {
4503 self.advance_auto_inc_past(*n)?;
4504 None
4505 }
4506 other => {
4507 return Err(MongrelError::InvalidArgument(format!(
4508 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
4509 other
4510 )))
4511 }
4512 },
4513 None => {
4514 let next = self.alloc_auto_inc_value()?;
4515 columns.push((cid, Value::Int64(next)));
4516 Some(next)
4517 }
4518 };
4519 Ok(assigned)
4520 }
4521
4522 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
4528 for col in &self.schema.columns {
4529 let Some(expr) = &col.default_value else {
4530 continue;
4531 };
4532 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
4534 continue;
4535 }
4536 let pos = columns.iter().position(|(c, _)| *c == col.id);
4537 let needs_default = match pos {
4538 None => true,
4539 Some(i) => matches!(columns[i].1, Value::Null),
4540 };
4541 if !needs_default {
4542 continue;
4543 }
4544 let v = match expr {
4545 crate::schema::DefaultExpr::Static(v) => v.clone(),
4546 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
4547 crate::schema::DefaultExpr::Uuid => {
4548 let mut buf = [0u8; 16];
4549 getrandom::getrandom(&mut buf)
4550 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
4551 Value::Uuid(buf)
4552 }
4553 };
4554 match pos {
4555 None => columns.push((col.id, v)),
4556 Some(i) => columns[i].1 = v,
4557 }
4558 }
4559 Ok(())
4560 }
4561
4562 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
4564 self.ensure_auto_inc_seeded()?;
4565 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4567 let v = ai.next;
4568 ai.next = ai
4569 .next
4570 .checked_add(1)
4571 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4572 Ok(v)
4573 }
4574
4575 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
4578 self.ensure_auto_inc_seeded()?;
4579 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4580 let floor = used
4581 .checked_add(1)
4582 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4583 .max(1);
4584 if ai.next < floor {
4585 ai.next = floor;
4586 }
4587 Ok(())
4588 }
4589
4590 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
4595 let needs_seed = match self.auto_inc {
4596 Some(ai) => !ai.seeded,
4597 None => return Ok(()),
4598 };
4599 if !needs_seed {
4600 return Ok(());
4601 }
4602 if self.seed_empty_auto_inc() {
4603 return Ok(());
4604 }
4605 let cid = self
4606 .auto_inc
4607 .as_ref()
4608 .expect("auto-inc column present")
4609 .column_id;
4610 let max = self.scan_max_int64(cid)?;
4611 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4612 let floor = max
4613 .checked_add(1)
4614 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4615 .max(1);
4616 if ai.next < floor {
4617 ai.next = floor;
4618 }
4619 ai.seeded = true;
4620 Ok(())
4621 }
4622
4623 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
4624 if n == 0 || self.auto_inc.is_none() {
4625 return Ok(None);
4626 }
4627 self.ensure_auto_inc_seeded()?;
4628 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
4629 let start = ai.next;
4630 let count = i64::try_from(n)
4631 .map_err(|_| MongrelError::Full("AUTO_INCREMENT range is too large".into()))?;
4632 ai.next = ai
4633 .next
4634 .checked_add(count)
4635 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?;
4636 Ok(Some(start))
4637 }
4638
4639 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
4643 let mut max: i64 = 0;
4644 for r in self.memtable.visible_versions_at(Snapshot::unbounded()) {
4645 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4646 if *n > max {
4647 max = *n;
4648 }
4649 }
4650 }
4651 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
4652 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
4653 if *n > max {
4654 max = *n;
4655 }
4656 }
4657 }
4658 for rr in self.run_refs.clone() {
4659 let reader = self.open_reader(rr.run_id)?;
4660 if let Some(stats) = reader.column_page_stats(column_id) {
4661 for s in stats {
4662 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
4663 if n > max {
4664 max = n;
4665 }
4666 }
4667 }
4668 } else if reader.has_column(column_id) {
4669 if let columnar::NativeColumn::Int64 { data, validity } =
4670 reader.column_native_shared(column_id)?
4671 {
4672 for (i, n) in data.iter().enumerate() {
4673 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
4674 {
4675 max = *n;
4676 }
4677 }
4678 }
4679 }
4680 }
4681 Ok(max)
4682 }
4683
4684 fn seed_empty_auto_inc(&mut self) -> bool {
4685 let Some(ai) = self.auto_inc.as_mut() else {
4686 return false;
4687 };
4688 if ai.seeded || self.live_count != 0 {
4689 return false;
4690 }
4691 if ai.next < 1 {
4692 ai.next = 1;
4693 }
4694 ai.seeded = true;
4695 true
4696 }
4697
4698 fn advance_auto_inc_from_native_columns(
4699 &mut self,
4700 columns: &[(u16, columnar::NativeColumn)],
4701 n: usize,
4702 live_before: u64,
4703 ) -> Result<()> {
4704 let Some(ai) = self.auto_inc.as_mut() else {
4705 return Ok(());
4706 };
4707 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
4708 return Ok(());
4709 };
4710 let columnar::NativeColumn::Int64 { data, validity } = col else {
4711 return Err(MongrelError::InvalidArgument(format!(
4712 "AUTO_INCREMENT column {} must be Int64",
4713 ai.column_id
4714 )));
4715 };
4716 let max = if native_int64_strictly_increasing(col, n) {
4717 data.get(n.saturating_sub(1)).copied()
4718 } else {
4719 data.iter()
4720 .take(n)
4721 .enumerate()
4722 .filter_map(|(i, v)| {
4723 if validity.is_empty() || columnar::validity_bit(validity, i) {
4724 Some(*v)
4725 } else {
4726 None
4727 }
4728 })
4729 .max()
4730 };
4731 if let Some(max) = max {
4732 let floor = max
4733 .checked_add(1)
4734 .ok_or_else(|| MongrelError::Full("AUTO_INCREMENT namespace exhausted".into()))?
4735 .max(1);
4736 if ai.next < floor {
4737 ai.next = floor;
4738 }
4739 if ai.seeded || live_before == 0 {
4740 ai.seeded = true;
4741 }
4742 }
4743 Ok(())
4744 }
4745
4746 fn fill_auto_inc_native_columns(
4747 &mut self,
4748 columns: &mut Vec<(u16, columnar::NativeColumn)>,
4749 n: usize,
4750 ) -> Result<()> {
4751 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
4752 return Ok(());
4753 };
4754 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
4755 if let Some(start) = self.alloc_auto_inc_range(n)? {
4756 columns.push((
4757 cid,
4758 columnar::NativeColumn::Int64 {
4759 data: (start..start.saturating_add(n as i64)).collect(),
4760 validity: vec![0xFF; n.div_ceil(8)],
4761 },
4762 ));
4763 }
4764 return Ok(());
4765 };
4766
4767 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
4768 return Err(MongrelError::InvalidArgument(format!(
4769 "AUTO_INCREMENT column {cid} must be Int64"
4770 )));
4771 };
4772 if data.len() < n {
4773 return Err(MongrelError::InvalidArgument(format!(
4774 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
4775 data.len()
4776 )));
4777 }
4778 if columnar::all_non_null(validity, n) {
4779 return Ok(());
4780 }
4781 if validity.iter().all(|b| *b == 0) {
4782 if let Some(start) = self.alloc_auto_inc_range(n)? {
4783 for (i, slot) in data.iter_mut().take(n).enumerate() {
4784 *slot = start.saturating_add(i as i64);
4785 }
4786 *validity = vec![0xFF; n.div_ceil(8)];
4787 }
4788 return Ok(());
4789 }
4790
4791 let new_validity = vec![0xFF; data.len().div_ceil(8)];
4792 for (i, slot) in data.iter_mut().enumerate().take(n) {
4793 if columnar::validity_bit(validity, i) {
4794 self.advance_auto_inc_past(*slot)?;
4795 } else {
4796 *slot = self.alloc_auto_inc_value()?;
4797 }
4798 }
4799 *validity = new_validity;
4800 Ok(())
4801 }
4802
4803 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
4817 self.ensure_writable()?;
4818 if self.auto_inc.is_none() {
4819 return Ok(None);
4820 }
4821 Ok(Some(self.alloc_auto_inc_value()?))
4822 }
4823
4824 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
4830 let payload = bincode::serialize(&rows)?;
4831 self.wal_append_data(Op::Put {
4832 table_id: self.table_id,
4833 rows: payload,
4834 })?;
4835 if self.is_shared() {
4836 self.pending_rows_auto_inc
4837 .extend(std::iter::repeat_n(auto_inc_generated, rows.len()));
4838 self.pending_rows.extend(rows);
4839 } else {
4840 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
4841 }
4842 Ok(())
4843 }
4844
4845 pub(crate) fn prepare_durable_publish(&mut self) -> Result<()> {
4848 self.ensure_indexes_complete()
4849 }
4850
4851 pub(crate) fn prepare_durable_publish_controlled(
4852 &mut self,
4853 control: &crate::ExecutionControl,
4854 ) -> Result<()> {
4855 self.ensure_indexes_complete_controlled(control, || true)?;
4856 Ok(())
4857 }
4858
4859 pub(crate) fn apply_put_rows_prepared(&mut self, rows: Vec<Row>) {
4860 self.apply_put_rows_inner_prepared(rows, true);
4861 }
4862
4863 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
4864 if check_existing_pk {
4865 self.ensure_indexes_complete()?;
4866 }
4867 self.apply_put_rows_inner_prepared(rows, check_existing_pk);
4868 Ok(())
4869 }
4870
4871 fn apply_put_rows_inner_prepared(&mut self, rows: Vec<Row>, check_existing_pk: bool) {
4875 if rows.len() == 1 {
4879 let row = rows.into_iter().next().expect("len checked");
4880 self.apply_put_row_single(row, check_existing_pk);
4881 return;
4882 }
4883 let pk_id = self.schema.primary_key().map(|c| c.id);
4900 let probe = match pk_id {
4901 Some(pid) => {
4902 check_existing_pk
4903 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
4904 }
4905 None => false,
4906 };
4907 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
4910 for r in rows {
4911 for &cid in r.columns.keys() {
4912 self.pending_put_cols.insert(cid);
4913 }
4914 let mut replaced_image: Option<Row> = None;
4915 match pk_id {
4916 Some(pid) if probe || maintain_pk_by_row => {
4917 if let Some(pk_val) = r.columns.get(&pid) {
4918 let key = self.index_lookup_key(pid, pk_val);
4919 if probe {
4920 if let Some(old) = self.recent_delete_preimages.remove(&key) {
4929 replaced_image = Some(old);
4930 if let Some(old_rid) = self.hot.get(&key) {
4931 if old_rid != r.row_id {
4932 self.tombstone_row(
4933 old_rid,
4934 r.committed_epoch,
4935 r.commit_ts,
4936 true,
4937 );
4938 }
4939 }
4940 } else if let Some(old_rid) = self.hot.get(&key) {
4941 if old_rid != r.row_id {
4942 replaced_image = self.get(old_rid, self.snapshot());
4943 self.tombstone_row(
4944 old_rid,
4945 r.committed_epoch,
4946 r.commit_ts,
4947 true,
4948 );
4949 }
4950 }
4951 }
4952 if maintain_pk_by_row {
4953 self.pk_by_row.insert(r.row_id, key);
4954 }
4955 }
4956 }
4957 Some(_) => {}
4958 None => {
4959 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
4960 }
4961 }
4962 if let Some(old) = replaced_image {
4963 self.maintain_indexes_on_pk_replace(&old, &r);
4964 } else {
4965 self.index_row(&r);
4966 }
4967 self.reservoir.offer(r.row_id.0);
4968 self.memtable.upsert(r);
4969 self.live_count = self.live_count.saturating_add(1);
4972 }
4973 self.data_generation = self.data_generation.wrapping_add(1);
4974 }
4975
4976 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) {
4986 for &cid in row.columns.keys() {
4987 self.pending_put_cols.insert(cid);
4988 }
4989 let epoch = row.committed_epoch;
4990 let mut replaced_image: Option<Row> = None;
4991 if let Some(pk_col) = self.schema.primary_key() {
4992 let pk_id = pk_col.id;
4993 if let Some(pk_val) = row.columns.get(&pk_id) {
4994 let maintain_pk_by_row = self.pk_by_row_complete;
4998 if check_existing_pk || maintain_pk_by_row {
4999 let key = self.index_lookup_key(pk_id, pk_val);
5000 if check_existing_pk {
5001 if let Some(old) = self.recent_delete_preimages.remove(&key) {
5010 replaced_image = Some(old);
5011 if let Some(old_rid) = self.hot.get(&key) {
5012 if old_rid != row.row_id {
5013 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5014 }
5015 }
5016 } else if let Some(old_rid) = self.hot.get(&key) {
5017 if old_rid != row.row_id {
5018 replaced_image = self.get(old_rid, self.snapshot());
5022 self.tombstone_row(old_rid, epoch, row.commit_ts, true);
5023 }
5024 }
5025 }
5026 if maintain_pk_by_row {
5027 self.pk_by_row.insert(row.row_id, key);
5028 }
5029 }
5030 }
5031 } else {
5032 self.hot
5033 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
5034 }
5035 if let Some(old) = replaced_image {
5036 self.maintain_indexes_on_pk_replace(&old, &row);
5037 } else {
5038 self.index_row(&row);
5039 }
5040 self.reservoir.offer(row.row_id.0);
5041 self.memtable.upsert(row);
5042 self.live_count = self.live_count.saturating_add(1);
5043 self.data_generation = self.data_generation.wrapping_add(1);
5044 }
5045
5046 fn maintain_indexes_on_pk_replace(&mut self, old: &Row, new: &Row) {
5049 let has_partial = self
5050 .schema
5051 .indexes
5052 .iter()
5053 .any(|idx| idx.predicate.is_some());
5054 if has_partial {
5055 self.unindex_bitmap_membership(old);
5059 self.index_row(new);
5060 return;
5061 }
5062
5063 crate::index::maintain_bitmap_secondary_on_replace(
5064 &self.schema,
5065 &mut self.bitmap,
5066 old,
5067 new,
5068 );
5069 self.index_row_non_bitmap(new);
5070 }
5071
5072 fn index_row_non_bitmap(&mut self, row: &Row) {
5075 if row.deleted {
5076 return;
5077 }
5078 let effective = if self.column_keys.is_empty() {
5079 None
5080 } else {
5081 Some(self.tokenized_for_indexes(row))
5082 };
5083 let source = effective.as_ref().unwrap_or(row);
5084 for idef in &self.schema.indexes {
5085 if idef.kind == crate::schema::IndexKind::Bitmap {
5086 continue;
5087 }
5088 index_into_single(
5089 idef,
5090 &self.schema,
5091 source,
5092 &mut self.hot,
5093 &mut self.bitmap,
5094 &mut self.ann,
5095 &mut self.fm,
5096 &mut self.sparse,
5097 &mut self.minhash,
5098 );
5099 }
5100 if let Some(pk_col) = self.schema.primary_key() {
5101 if let Some(pk_val) = source.columns.get(&pk_col.id) {
5102 let key = if self.column_keys.is_empty() {
5103 pk_val.encode_key()
5104 } else {
5105 self.index_lookup_key(pk_col.id, pk_val)
5106 };
5107 self.hot.insert(key, source.row_id);
5108 }
5109 }
5110 }
5111
5112 pub(crate) fn alloc_row_id(&mut self) -> Result<RowId> {
5115 self.allocator.alloc()
5116 }
5117
5118 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
5124 let pk = self.schema.primary_key().ok_or_else(|| {
5125 MongrelError::Schema("clustered table requires a single-column primary key".into())
5126 })?;
5127 let pk_val = columns
5128 .iter()
5129 .find(|(id, _)| *id == pk.id)
5130 .map(|(_, v)| v)
5131 .ok_or_else(|| {
5132 MongrelError::Schema(format!(
5133 "clustered table missing primary key column {} ({})",
5134 pk.id, pk.name
5135 ))
5136 })?;
5137 Ok(clustered_row_id(pk_val))
5138 }
5139
5140 pub(crate) fn apply_run_metadata_prepared(&mut self, rows: &[Row]) -> Result<()> {
5148 if rows.iter().any(|row| row.row_id.0 >= u64::MAX - 1) {
5149 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5150 }
5151 let n = rows.len();
5152 for r in rows {
5153 for &cid in r.columns.keys() {
5154 self.pending_put_cols.insert(cid);
5155 }
5156 }
5157 let (losers, winner_pks) = self.partition_pk_winners(rows);
5158 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
5159 let group_ts = rows.first().and_then(|r| r.commit_ts);
5161 for (key, &row_id) in &winner_pks {
5162 if let Some(old_rid) = self.hot.get(key) {
5163 if old_rid != row_id {
5164 self.tombstone_row(old_rid, epoch, group_ts, true);
5165 }
5166 }
5167 }
5168 for &loser_rid in &losers {
5171 self.tombstone_row(loser_rid, epoch, group_ts, false);
5172 }
5173 for (key, row_id) in winner_pks {
5175 self.insert_hot_pk(key, row_id);
5176 }
5177 if self.schema.primary_key().is_none() {
5178 for r in rows {
5179 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5180 }
5181 }
5182 for r in rows {
5183 self.allocator.advance_to(r.row_id)?;
5184 if !losers.contains(&r.row_id) {
5185 self.index_row(r);
5186 }
5187 }
5188 for r in rows {
5189 if !losers.contains(&r.row_id) {
5190 self.reservoir.offer(r.row_id.0);
5191 }
5192 }
5193 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
5194 self.data_generation = self.data_generation.wrapping_add(1);
5195 Ok(())
5196 }
5197
5198 pub(crate) fn recover_apply(
5203 &mut self,
5204 rows: Vec<Row>,
5205 deletes: Vec<(RowId, Epoch)>,
5206 ) -> Result<()> {
5207 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
5211 std::collections::BTreeMap::new();
5212 for row in rows {
5213 if row.row_id.0 >= u64::MAX - 1 {
5214 return Err(MongrelError::Full("row-id namespace exhausted".into()));
5215 }
5216 self.allocator.advance_to(row.row_id)?;
5217 if let Some(ai) = self.auto_inc.as_mut() {
5222 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
5223 let next = n.checked_add(1).ok_or_else(|| {
5224 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
5225 })?;
5226 if next > ai.next {
5227 ai.next = next;
5228 }
5229 }
5230 }
5231 by_epoch.entry(row.committed_epoch).or_default().push(row);
5232 }
5233 for (epoch, group) in by_epoch {
5234 let (losers, winner_pks) = self.partition_pk_winners(&group);
5235 let group_ts = group.first().and_then(|r| r.commit_ts);
5237 for (key, &row_id) in &winner_pks {
5238 if let Some(old_rid) = self.hot.get(key) {
5239 if old_rid != row_id {
5240 self.tombstone_row(old_rid, epoch, group_ts, false);
5241 }
5242 }
5243 }
5244 for (key, row_id) in winner_pks {
5245 self.insert_hot_pk(key, row_id);
5246 }
5247 if self.schema.primary_key().is_none() {
5248 for r in &group {
5249 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
5250 }
5251 }
5252 for r in &group {
5253 if !losers.contains(&r.row_id) {
5254 self.memtable.upsert(r.clone());
5255 self.index_row(r);
5256 }
5257 }
5258 }
5259 for (rid, epoch) in deletes {
5260 self.memtable.tombstone(rid, epoch);
5261 self.remove_hot_for_row(rid, epoch);
5262 }
5263 self.reservoir_complete = false;
5266 Ok(())
5267 }
5268
5269 pub(crate) fn flushed_epoch(&self) -> u64 {
5271 self.flushed_epoch
5272 }
5273
5274 pub(crate) fn set_flushed_epoch(&mut self, epoch: Epoch) {
5275 self.flushed_epoch = self.flushed_epoch.max(epoch.0);
5276 }
5277
5278 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
5280 self.schema.validate_values(cells)
5281 }
5282
5283 fn validate_columns_not_null(
5287 &self,
5288 columns: &[(u16, columnar::NativeColumn)],
5289 n: usize,
5290 ) -> Result<()> {
5291 let by_id: HashMap<u16, &columnar::NativeColumn> =
5292 columns.iter().map(|(id, c)| (*id, c)).collect();
5293 for col in &self.schema.columns {
5294 if !col.flags.contains(ColumnFlags::NULLABLE) {
5295 match by_id.get(&col.id) {
5296 None => {
5297 return Err(MongrelError::InvalidArgument(format!(
5298 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
5299 col.name, col.id
5300 )));
5301 }
5302 Some(c) => {
5303 if c.null_count(n) != 0 {
5304 return Err(MongrelError::InvalidArgument(format!(
5305 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
5306 col.name, col.id
5307 )));
5308 }
5309 }
5310 }
5311 }
5312 if let TypeId::Enum { variants } = &col.ty {
5313 let Some(columnar::NativeColumn::Bytes { .. }) = by_id.get(&col.id).copied() else {
5314 if by_id.contains_key(&col.id) {
5315 return Err(MongrelError::InvalidArgument(format!(
5316 "column '{}' ({}) enum requires a bytes column",
5317 col.name, col.id
5318 )));
5319 }
5320 continue;
5321 };
5322 for index in 0..n {
5323 let Some(value) = columnar::native_bytes_at(by_id[&col.id], index) else {
5324 continue;
5325 };
5326 if !variants.iter().any(|variant| variant.as_bytes() == value) {
5327 return Err(MongrelError::InvalidArgument(format!(
5328 "column '{}' ({}) enum value {:?} is not one of {:?}",
5329 col.name,
5330 col.id,
5331 String::from_utf8_lossy(value),
5332 variants
5333 )));
5334 }
5335 }
5336 }
5337 }
5338 Ok(())
5339 }
5340
5341 fn bulk_pk_winner_indices(
5346 &self,
5347 columns: &[(u16, columnar::NativeColumn)],
5348 n: usize,
5349 ) -> Option<Vec<usize>> {
5350 let pk_col = self.schema.primary_key()?;
5351 let pk_id = pk_col.id;
5352 let pk_ty = pk_col.ty.clone();
5353 let by_id: HashMap<u16, &columnar::NativeColumn> =
5354 columns.iter().map(|(id, c)| (*id, c)).collect();
5355 let pk_native = by_id.get(&pk_id)?;
5356 if native_int64_strictly_increasing(pk_native, n) {
5357 return None;
5358 }
5359 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
5361 let mut null_pk_rows: Vec<usize> = Vec::new();
5362 for i in 0..n {
5363 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
5364 Some(key) => {
5365 last.insert(key, i);
5366 }
5367 None => null_pk_rows.push(i),
5368 }
5369 }
5370 let mut winners: HashSet<usize> = last.values().copied().collect();
5371 for i in null_pk_rows {
5372 winners.insert(i);
5373 }
5374 Some((0..n).filter(|i| winners.contains(i)).collect())
5375 }
5376
5377 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
5379 self.require_delete()?;
5380 let epoch = self.pending_epoch();
5381 self.wal_append_data(Op::Delete {
5382 table_id: self.table_id,
5383 row_ids: vec![row_id],
5384 })?;
5385 if self.is_shared() {
5386 self.pending_dels.push(row_id);
5387 } else {
5388 self.apply_delete(row_id, epoch);
5389 }
5390 Ok(())
5391 }
5392
5393 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
5394 let pre = self.get(row_id, self.snapshot());
5395 self.delete(row_id)?;
5396 Ok(pre.map(|row| {
5397 let mut columns: Vec<_> = row.columns.into_iter().collect();
5398 columns.sort_by_key(|(id, _)| *id);
5399 OwnedRow { columns }
5400 }))
5401 }
5402
5403 pub fn truncate(&mut self) -> Result<()> {
5405 self.require_delete()?;
5406 let epoch = self.pending_epoch();
5407 self.wal_append_data(Op::TruncateTable {
5408 table_id: self.table_id,
5409 })?;
5410 self.pending_rows.clear();
5411 self.pending_rows_auto_inc.clear();
5412 self.pending_dels.clear();
5413 self.pending_truncate = Some(epoch);
5414 Ok(())
5415 }
5416
5417 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) {
5419 self.run_refs.clear();
5424 self.retiring.clear();
5425 self.memtable = Memtable::new();
5426 self.mutable_run = MutableRun::new();
5427 self.hot = HotIndex::new();
5428 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
5429 self.bitmap = bitmap;
5430 self.ann = ann;
5431 self.fm = fm;
5432 self.sparse = sparse;
5433 self.minhash = minhash;
5434 self.learned_range = Arc::new(HashMap::new());
5435 self.pk_by_row.clear();
5436 self.pk_by_row_complete = false;
5437 self.live_count = 0;
5438 self.reservoir = crate::reservoir::Reservoir::default();
5439 self.reservoir_complete = true;
5440 self.had_deletes = true;
5441 self.agg_cache = Arc::new(HashMap::new());
5442 self.global_idx_epoch = 0;
5443 self.indexes_complete = true;
5444 self.pending_delete_rids.clear();
5445 self.pending_put_cols.clear();
5446 self.pending_rows.clear();
5447 self.pending_rows_auto_inc.clear();
5448 self.pending_dels.clear();
5449 self.clear_result_cache();
5450 self.invalidate_index_checkpoint();
5451 self.data_generation = self.data_generation.wrapping_add(1);
5452 }
5453
5454 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
5457 self.apply_delete_at(row_id, epoch, None);
5458 }
5459
5460 pub(crate) fn apply_delete_at(
5462 &mut self,
5463 row_id: RowId,
5464 epoch: Epoch,
5465 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5466 ) {
5467 let preimage = self.get(row_id, self.snapshot());
5477 if let Some(row) = preimage {
5478 if let Some(pk_col) = self.schema.primary_key() {
5479 if let Some(pk_val) = row.columns.get(&pk_col.id) {
5480 let key = self.index_lookup_key(pk_col.id, pk_val);
5481 self.recent_delete_preimages.insert(key, row);
5482 }
5483 }
5484 }
5485 self.tombstone_row(row_id, epoch, commit_ts, true);
5486 self.data_generation = self.data_generation.wrapping_add(1);
5487 }
5488
5489 fn unindex_bitmap_membership(&mut self, row: &Row) {
5494 if row.deleted {
5495 return;
5496 }
5497 for idef in &self.schema.indexes {
5498 if idef.kind != crate::schema::IndexKind::Bitmap {
5499 continue;
5500 }
5501 if let Some(key) = crate::index::maintain::bitmap_key_for_column(row, idef.column_id) {
5502 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
5503 b.remove(&key, row.row_id);
5504 }
5505 }
5506 }
5507 }
5508
5509 fn union_bitmap_point_i64(
5513 &self,
5514 set: &mut RowIdSet,
5515 column_id: u16,
5516 value: i64,
5517 snapshot: Snapshot,
5518 ) {
5519 let Some(b) = self.bitmap.get(&column_id) else {
5520 return;
5521 };
5522 let encoded = Value::Int64(value).encode_key();
5523 let lookup = self.index_lookup_key_bytes(column_id, &encoded);
5524 for rid in b.get(&lookup).iter() {
5525 set.insert(u64::from(rid));
5526 }
5527 set.remove_many(self.overlay_tombstoned_rids(snapshot));
5531 self.range_scan_overlay_i64(set, column_id, value, value, snapshot);
5532 }
5533
5534 fn tombstone_row(
5548 &mut self,
5549 row_id: RowId,
5550 epoch: Epoch,
5551 commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
5552 adjust_live_count: bool,
5553 ) {
5554 let prev_was_live = if adjust_live_count {
5555 match self.memtable.get_version(row_id, epoch) {
5556 Some((_, prev)) => !prev.deleted,
5557 None => true,
5558 }
5559 } else {
5560 false
5561 };
5562 let tombstone = Row {
5563 row_id,
5564 committed_epoch: epoch,
5565 columns: std::collections::HashMap::new(),
5566 deleted: true,
5567 commit_ts,
5568 };
5569 self.memtable.upsert(tombstone);
5570 self.pk_by_row.remove(&row_id);
5571 if prev_was_live {
5572 self.live_count = self.live_count.saturating_sub(1);
5573 }
5574 self.pending_delete_rids.insert(row_id.0 as u32);
5576 self.had_deletes = true;
5579 self.agg_cache = Arc::new(HashMap::new());
5580 }
5581
5582 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
5586 let Some(pk_col) = self.schema.primary_key() else {
5587 return;
5588 };
5589 if self.pk_by_row_complete {
5592 if let Some(key) = self.pk_by_row.remove(&row_id) {
5593 if self.hot.get(&key) == Some(row_id) {
5594 self.hot.remove(&key);
5595 }
5596 }
5597 return;
5598 }
5599 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
5618 if self.indexes_complete {
5619 let pk_val = self
5620 .memtable
5621 .get_version(row_id, lookup_epoch)
5622 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5623 .or_else(|| {
5624 self.mutable_run
5625 .get_version(row_id, lookup_epoch)
5626 .filter(|(_, r)| !r.deleted)
5627 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
5628 })
5629 .or_else(|| {
5630 self.run_refs.iter().find_map(|rr| {
5631 let mut reader = self.open_reader(rr.run_id).ok()?;
5632 let (_, deleted, val) = reader
5633 .get_version_column(row_id, lookup_epoch, pk_col.id)
5634 .ok()??;
5635 if deleted {
5636 return None;
5637 }
5638 val
5639 })
5640 });
5641 if let Some(pk_val) = pk_val {
5642 let key = self.index_lookup_key(pk_col.id, &pk_val);
5643 if self.hot.get(&key) == Some(row_id) {
5644 self.hot.remove(&key);
5645 }
5646 return;
5647 }
5648 }
5649 self.refresh_pk_by_row_from_hot();
5654 if let Some(key) = self.pk_by_row.remove(&row_id) {
5655 if self.hot.get(&key) == Some(row_id) {
5656 self.hot.remove(&key);
5657 }
5658 }
5659 }
5660
5661 fn partition_pk_winners(
5666 &self,
5667 rows: &[Row],
5668 ) -> (
5669 std::collections::HashSet<RowId>,
5670 std::collections::HashMap<Vec<u8>, RowId>,
5671 ) {
5672 let mut losers = std::collections::HashSet::new();
5673 let Some(pk_col) = self.schema.primary_key() else {
5674 return (losers, std::collections::HashMap::new());
5675 };
5676 let pk_id = pk_col.id;
5677 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
5678 std::collections::HashMap::new();
5679 for r in rows {
5680 let Some(pk_val) = r.columns.get(&pk_id) else {
5681 continue;
5682 };
5683 let key = self.index_lookup_key(pk_id, pk_val);
5684 if let Some(&old_rid) = winners.get(&key) {
5685 losers.insert(old_rid);
5686 }
5687 winners.insert(key, r.row_id);
5688 }
5689 (losers, winners)
5690 }
5691
5692 fn index_row(&mut self, row: &Row) {
5693 if row.deleted {
5694 return;
5695 }
5696 let any_predicate = self
5704 .schema
5705 .indexes
5706 .iter()
5707 .any(|idx| idx.predicate.is_some());
5708 if any_predicate {
5709 let columns_map: HashMap<u16, &Value> =
5710 row.columns.iter().map(|(k, v)| (*k, v)).collect();
5711 let name_to_id: HashMap<&str, u16> = self
5712 .schema
5713 .columns
5714 .iter()
5715 .map(|c| (c.name.as_str(), c.id))
5716 .collect();
5717 for idx in &self.schema.indexes {
5718 if let Some(pred) = &idx.predicate {
5719 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
5720 continue; }
5722 }
5723 index_into_single(
5725 idx,
5726 &self.schema,
5727 row,
5728 &mut self.hot,
5729 &mut self.bitmap,
5730 &mut self.ann,
5731 &mut self.fm,
5732 &mut self.sparse,
5733 &mut self.minhash,
5734 );
5735 }
5736 return;
5737 }
5738 if self.column_keys.is_empty() {
5742 index_into(
5743 &self.schema,
5744 row,
5745 &mut self.hot,
5746 &mut self.bitmap,
5747 &mut self.ann,
5748 &mut self.fm,
5749 &mut self.sparse,
5750 &mut self.minhash,
5751 );
5752 return;
5753 }
5754 let effective_row = self.tokenized_for_indexes(row);
5755 index_into(
5756 &self.schema,
5757 &effective_row,
5758 &mut self.hot,
5759 &mut self.bitmap,
5760 &mut self.ann,
5761 &mut self.fm,
5762 &mut self.sparse,
5763 &mut self.minhash,
5764 );
5765 }
5766
5767 fn tokenized_for_indexes(&self, row: &Row) -> Row {
5773 if self.column_keys.is_empty() {
5774 return row.clone();
5775 }
5776 {
5777 use crate::encryption::SCHEME_HMAC_EQ;
5778 let mut tok = row.clone();
5779 for (&cid, &(_, scheme)) in &self.column_keys {
5780 if scheme != SCHEME_HMAC_EQ {
5781 continue;
5782 }
5783 if let Some(v) = tok.columns.get(&cid).cloned() {
5784 if let Some(t) = self.tokenize_value(cid, &v) {
5785 tok.columns.insert(cid, t);
5786 }
5787 }
5788 }
5789 tok
5790 }
5791 }
5792
5793 pub fn commit(&mut self) -> Result<Epoch> {
5798 self.commit_inner(None)
5799 }
5800
5801 #[doc(hidden)]
5804 pub fn commit_controlled<F>(
5805 &mut self,
5806 control: &crate::ExecutionControl,
5807 mut before_commit: F,
5808 ) -> Result<Epoch>
5809 where
5810 F: FnMut() -> Result<()>,
5811 {
5812 self.commit_inner(Some((control, &mut before_commit)))
5813 }
5814
5815 fn commit_inner(
5816 &mut self,
5817 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5818 ) -> Result<Epoch> {
5819 self.ensure_writable()?;
5820 if !self.has_pending_mutations() {
5821 if self.current_txn_id == 0 && matches!(&self.wal, WalSink::Private(_)) {
5822 return Err(MongrelError::Full(
5823 "standalone transaction id namespace exhausted".into(),
5824 ));
5825 }
5826 return Ok(self.epoch.visible());
5827 }
5828 self.commit_new_epoch_inner(controlled)
5829 }
5830
5831 fn commit_new_epoch(&mut self) -> Result<Epoch> {
5835 self.commit_new_epoch_inner(None)
5836 }
5837
5838 fn commit_new_epoch_inner(
5839 &mut self,
5840 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5841 ) -> Result<Epoch> {
5842 self.ensure_writable()?;
5843 if self.is_shared() {
5844 self.commit_shared(controlled)
5845 } else {
5846 self.commit_private(controlled)
5847 }
5848 }
5849
5850 fn commit_private(
5852 &mut self,
5853 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5854 ) -> Result<Epoch> {
5855 let commit_lock = Arc::clone(&self.commit_lock);
5859 let _g = commit_lock.lock();
5860 let txn_id = self.ensure_txn_id()?;
5863 if let Some((control, before_commit)) = controlled {
5864 control.checkpoint()?;
5865 before_commit()?;
5866 }
5867 let new_epoch = self.epoch.bump_assigned();
5868 let epoch_authority = Arc::clone(&self.epoch);
5869 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5870 let wal_result = match &mut self.wal {
5874 WalSink::Private(w) => w
5875 .append_txn(
5876 txn_id,
5877 Op::TxnCommit {
5878 epoch: new_epoch.0,
5879 added_runs: Vec::new(),
5880 },
5881 )
5882 .and_then(|_| w.sync()),
5883 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
5884 WalSink::ReadOnly => Err(MongrelError::ReadOnlyReplica),
5885 };
5886 if let Err(error) = wal_result {
5887 self.durable_commit_failed = true;
5888 return Err(MongrelError::CommitOutcomeUnknown {
5889 epoch: new_epoch.0,
5890 message: error.to_string(),
5891 });
5892 }
5893 if let Some(epoch) = self.pending_truncate.take() {
5896 self.apply_truncate(epoch);
5897 }
5898 self.invalidate_pending_cache();
5899 let publish_result = self.persist_manifest(new_epoch);
5900 self.epoch.publish_in_order(new_epoch);
5904 epoch_guard.disarm();
5905 if let Err(error) = publish_result {
5906 self.durable_commit_failed = true;
5907 return Err(MongrelError::DurableCommit {
5908 epoch: new_epoch.0,
5909 message: error.to_string(),
5910 });
5911 }
5912 self.current_txn_id = txn_id.checked_add(1).unwrap_or(0);
5913 self.pending_private_mutations = false;
5914 self.data_generation = self.data_generation.wrapping_add(1);
5915 Ok(new_epoch)
5916 }
5917
5918 fn commit_shared(
5924 &mut self,
5925 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
5926 ) -> Result<Epoch> {
5927 use std::sync::atomic::Ordering;
5928 let s = match &self.wal {
5929 WalSink::Shared(s) => s.clone(),
5930 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
5931 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
5932 };
5933 if s.poisoned.load(Ordering::Relaxed) {
5934 return Err(MongrelError::Other(
5935 "database poisoned by fsync error".into(),
5936 ));
5937 }
5938 let commit_lock = Arc::clone(&self.commit_lock);
5945 let _g = commit_lock.lock();
5946 if !self.pending_rows.is_empty() {
5947 match controlled.as_ref() {
5948 Some((control, _)) => self.prepare_durable_publish_controlled(control)?,
5949 None => self.prepare_durable_publish()?,
5950 }
5951 }
5952 let txn_id = self.ensure_txn_id()?;
5955 let mut wal = s.wal.lock();
5956 if let Some((control, before_commit)) = controlled {
5957 control.checkpoint()?;
5958 before_commit()?;
5959 }
5960 let new_epoch = self.epoch.bump_assigned();
5961 let epoch_authority = Arc::clone(&self.epoch);
5962 let mut epoch_guard = EpochGuard::new(epoch_authority.as_ref(), new_epoch);
5963 let commit_ts = s.hlc.now().map_err(|skew| {
5966 MongrelError::Other(format!(
5967 "clock skew rejected commit timestamp allocation: {skew}"
5968 ))
5969 })?;
5970 let commit_seq = match wal.append_commit_at(
5971 txn_id,
5972 new_epoch,
5973 &[],
5974 commit_ts.physical_micros.saturating_mul(1_000),
5975 ) {
5976 Ok(commit_seq) => commit_seq,
5977 Err(error) => {
5978 s.poisoned.store(true, Ordering::Relaxed);
5979 s.lifecycle.poison();
5980 return Err(MongrelError::CommitOutcomeUnknown {
5981 epoch: new_epoch.0,
5982 message: error.to_string(),
5983 });
5984 }
5985 };
5986 drop(wal);
5987 if let Err(error) = s.group.await_durable(&s.wal, commit_seq) {
5988 s.poisoned.store(true, Ordering::Relaxed);
5989 s.lifecycle.poison();
5990 return Err(MongrelError::CommitOutcomeUnknown {
5991 epoch: new_epoch.0,
5992 message: error.to_string(),
5993 });
5994 }
5995
5996 if self.pending_truncate.take().is_some() {
5999 self.apply_truncate(new_epoch);
6000 }
6001 let mut rows = std::mem::take(&mut self.pending_rows);
6002 if !rows.is_empty() {
6003 for r in &mut rows {
6004 r.committed_epoch = new_epoch;
6005 r.commit_ts = Some(commit_ts);
6006 }
6007 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
6008 let all_auto_generated =
6009 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
6010 self.apply_put_rows_inner_prepared(rows, !all_auto_generated);
6011 } else {
6012 self.pending_rows_auto_inc.clear();
6013 }
6014 let dels = std::mem::take(&mut self.pending_dels);
6015 for rid in dels {
6016 self.apply_delete_at(rid, new_epoch, Some(commit_ts));
6017 }
6018
6019 self.invalidate_pending_cache();
6020 let publish_result = self.persist_manifest(new_epoch);
6021 self.epoch.publish_in_order(new_epoch);
6022 epoch_guard.disarm();
6023 let _ = s.change_wake.send(());
6024 if let Err(error) = publish_result {
6025 self.durable_commit_failed = true;
6026 s.poisoned.store(true, Ordering::Relaxed);
6027 s.lifecycle.poison();
6028 return Err(MongrelError::DurableCommit {
6029 epoch: new_epoch.0,
6030 message: error.to_string(),
6031 });
6032 }
6033 self.current_txn_id = 0;
6035 self.data_generation = self.data_generation.wrapping_add(1);
6036 Ok(new_epoch)
6037 }
6038
6039 pub fn flush(&mut self) -> Result<Epoch> {
6047 self.flush_with_outcome().map(|(epoch, _)| epoch)
6048 }
6049
6050 pub fn flush_with_outcome(&mut self) -> Result<(Epoch, bool)> {
6052 self.flush_with_outcome_inner(None)
6053 }
6054
6055 #[doc(hidden)]
6058 pub fn flush_with_outcome_controlled<F>(
6059 &mut self,
6060 control: &crate::ExecutionControl,
6061 mut before_commit: F,
6062 ) -> Result<(Epoch, bool)>
6063 where
6064 F: FnMut() -> Result<()>,
6065 {
6066 self.flush_with_outcome_inner(Some((control, &mut before_commit)))
6067 }
6068
6069 fn flush_with_outcome_inner(
6070 &mut self,
6071 controlled: Option<(&crate::ExecutionControl, &mut dyn FnMut() -> Result<()>)>,
6072 ) -> Result<(Epoch, bool)> {
6073 match controlled.as_ref() {
6074 Some((control, _)) => {
6075 self.ensure_indexes_complete_controlled(control, || true)?;
6076 }
6077 None => self.ensure_indexes_complete()?,
6078 }
6079 let committed = self.has_pending_mutations();
6080 let epoch = self.commit_inner(controlled)?;
6081 let finish: Result<(Epoch, bool)> = (|| {
6082 let rows = self.memtable.drain_sorted();
6083 if !rows.is_empty() {
6084 self.mutable_run.insert_many(rows);
6085 }
6086 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
6087 self.spill_mutable_run(epoch)?;
6088 self.mark_flushed(epoch)?;
6092 self.persist_manifest(epoch)?;
6093 self.build_learned_ranges()?;
6094 self.checkpoint_indexes(epoch);
6097 let _ = self.publish_run_lookup_directory();
6102 }
6103 Ok((epoch, committed))
6106 })();
6107 let outcome = match finish {
6108 Err(error) if committed => Err(MongrelError::DurableCommit {
6109 epoch: epoch.0,
6110 message: error.to_string(),
6111 }),
6112 result => result,
6113 };
6114 if outcome.is_ok() {
6115 let _ = self.publish_read_generation();
6121 }
6122 outcome
6123 }
6124
6125 fn has_pending_mutations(&self) -> bool {
6126 self.pending_private_mutations
6127 || !self.pending_rows.is_empty()
6128 || !self.pending_dels.is_empty()
6129 || self.pending_truncate.is_some()
6130 }
6131
6132 pub fn has_pending_writes(&self) -> bool {
6133 self.has_pending_mutations()
6134 }
6135
6136 pub fn force_flush(&mut self) -> Result<Epoch> {
6145 let saved = self.mutable_run_spill_bytes;
6146 self.mutable_run_spill_bytes = 1;
6147 let result = self.flush();
6148 self.mutable_run_spill_bytes = saved;
6149 result
6150 }
6151
6152 pub fn close(&mut self) -> Result<()> {
6159 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
6160 self.force_flush()?;
6161 }
6162 self.result_cache
6168 .lock()
6169 .shutdown_persistent_cache(std::time::Duration::from_secs(5));
6170 Ok(())
6171 }
6172
6173 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
6180 let op = Op::Flush {
6181 table_id: self.table_id,
6182 flushed_epoch: epoch.0,
6183 };
6184 match &mut self.wal {
6185 WalSink::Private(w) => {
6186 w.append_system(op)?;
6187 w.sync()?;
6188 }
6189 WalSink::Shared(s) => {
6190 s.wal.lock().append_system(op)?;
6195 }
6196 WalSink::ReadOnly => return Err(MongrelError::ReadOnlyReplica),
6197 }
6198 self.flushed_epoch = epoch.0;
6199 if matches!(self.wal, WalSink::Private(_)) {
6200 self.rotate_wal(epoch)?;
6201 }
6202 Ok(())
6203 }
6204
6205 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
6209 if self.mutable_run.is_empty() {
6210 return Ok(());
6211 }
6212 let run_id = self.alloc_run_id()?;
6213 let rows = self.mutable_run.drain_sorted();
6214 if rows.is_empty() {
6215 return Ok(());
6216 }
6217 let path = self.run_path(run_id);
6218 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
6219 if let Some(kek) = &self.kek {
6220 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6221 }
6222 let header = match self.create_run_file(run_id)? {
6223 Some(file) => writer.write_file(file, &rows)?,
6224 None => writer.write(&path, &rows)?,
6225 };
6226 self.run_refs.push(RunRef {
6227 run_id: run_id as u128,
6228 level: 0,
6229 epoch_created: epoch.0,
6230 row_count: header.row_count,
6231 });
6232 self.run_row_id_ranges
6233 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6234 Ok(())
6235 }
6236
6237 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
6241 self.mutable_run_spill_bytes = bytes.max(1);
6242 }
6243
6244 pub fn set_compaction_zstd_level(&mut self, level: i32) {
6248 self.compaction_zstd_level = level;
6249 }
6250
6251 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
6255 self.result_cache.lock().set_max_bytes(max_bytes);
6256 }
6257
6258 pub fn flush_persistent_cache(&self, deadline_ms: u64) -> u64 {
6263 self.result_cache
6264 .lock()
6265 .flush_persistent_cache(std::time::Duration::from_millis(deadline_ms))
6266 }
6267
6268 pub fn shutdown_persistent_cache(&mut self, deadline_ms: u64) {
6273 self.result_cache
6274 .lock()
6275 .shutdown_persistent_cache(std::time::Duration::from_millis(deadline_ms));
6276 }
6277
6278 pub(crate) fn clear_result_cache(&mut self) {
6282 self.result_cache.lock().clear();
6283 }
6284
6285 pub fn mutable_run_len(&self) -> usize {
6287 self.mutable_run.len()
6288 }
6289
6290 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
6293 self.mutable_run.drain_sorted()
6294 }
6295
6296 pub(crate) fn snapshot_mutable_run(&self) -> Vec<Row> {
6298 let mut snapshot = self.mutable_run.clone();
6299 snapshot.drain_sorted()
6300 }
6301
6302 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
6307 self.ensure_writable()?;
6308 let n = batch.len();
6309 if n == 0 {
6310 return Ok(self.current_epoch());
6311 }
6312 for row in &batch {
6313 self.schema.validate_values(row)?;
6314 }
6315 let epoch = self.commit_new_epoch()?;
6316 let live_before = self.live_count;
6317 self.spill_mutable_run(epoch)?;
6321 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
6322 && self.indexes_complete
6323 && self.run_refs.is_empty()
6324 && self.memtable.is_empty()
6325 && self.mutable_run.is_empty();
6326 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
6332 use rayon::prelude::*;
6333 self.schema
6334 .columns
6335 .par_iter()
6336 .map(|cdef| {
6337 (
6338 cdef.id,
6339 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
6340 )
6341 })
6342 .collect::<Vec<_>>()
6343 };
6344 drop(batch);
6345 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
6350 self.validate_columns_not_null(&user_columns, n)?;
6351 let winner_idx = self
6352 .bulk_pk_winner_indices(&user_columns, n)
6353 .filter(|idx| idx.len() != n);
6354 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
6355 match winner_idx.as_deref() {
6356 Some(idx) => {
6357 let compacted = user_columns
6358 .iter()
6359 .map(|(id, c)| (*id, c.gather(idx)))
6360 .collect();
6361 (compacted, idx.len())
6362 }
6363 None => (std::mem::take(&mut user_columns), n),
6364 };
6365 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
6366 let first = self.allocator.alloc_range(write_n as u64)?.0;
6367 for rid in first..first + write_n as u64 {
6368 self.reservoir.offer(rid);
6369 }
6370 let run_id = self.alloc_run_id()?;
6371 let path = self.run_path(run_id);
6372 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
6373 .clean(true)
6374 .with_lz4()
6375 .with_native_endian();
6376 if let Some(kek) = &self.kek {
6377 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
6378 }
6379 let header = match self.create_run_file(run_id)? {
6380 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
6381 None => writer.write_native(&path, &write_columns, write_n, first)?,
6382 };
6383 self.run_refs.push(RunRef {
6384 run_id: run_id as u128,
6385 level: 0,
6386 epoch_created: epoch.0,
6387 row_count: header.row_count,
6388 });
6389 self.run_row_id_ranges
6390 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
6391 self.live_count = self.live_count.saturating_add(write_n as u64);
6392 if eager_index_build {
6393 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
6394 self.index_columns_bulk(&write_columns, &row_ids);
6395 self.indexes_complete = true;
6396 self.build_learned_ranges()?;
6397 } else {
6398 self.indexes_complete = false;
6399 }
6400 self.mark_flushed(epoch)?;
6401 self.persist_manifest(epoch)?;
6402 if eager_index_build {
6403 self.checkpoint_indexes(epoch);
6404 }
6405 self.clear_result_cache();
6406 Ok(epoch)
6407 }
6408
6409 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
6412 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
6413 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
6414 let segment_no = segment
6417 .file_stem()
6418 .and_then(|s| s.to_str())
6419 .and_then(|s| s.strip_prefix("seg-"))
6420 .and_then(|s| s.parse::<u64>().ok())
6421 .unwrap_or(0);
6422 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
6423 wal.set_sync_byte_threshold(self.sync_byte_threshold);
6424 wal.sync()?;
6425 self.wal = WalSink::Private(wal);
6426 Ok(())
6427 }
6428
6429 pub(crate) fn invalidate_pending_cache(&mut self) {
6434 self.result_cache
6435 .lock()
6436 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
6437 self.pending_delete_rids.clear();
6438 self.pending_put_cols.clear();
6439 }
6440
6441 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
6442 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
6443 m.current_epoch = epoch.0;
6444 m.next_row_id = self.allocator.current().0;
6445 m.runs = self.run_refs.clone();
6446 m.live_count = self.live_count;
6447 m.global_idx_epoch = self.global_idx_epoch;
6448 m.flushed_epoch = self.flushed_epoch;
6449 m.retiring = self.retiring.clone();
6450 m.auto_inc_next = match self.auto_inc {
6454 Some(ai) if ai.seeded => ai.next,
6455 _ => 0,
6456 };
6457 m.ttl = self.ttl;
6458 let meta_dek = self.manifest_meta_dek();
6459 match self._root_guard.as_deref() {
6460 Some(root) => manifest::write_durable(root, &mut m, meta_dek.as_ref())?,
6461 None => manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?,
6462 }
6463 Ok(())
6464 }
6465
6466 pub(crate) fn publish_run_lookup_directory(&mut self) -> Result<()> {
6473 mongreldb_fault::inject("manifest.publication")
6477 .map_err(|e| MongrelError::Other(format!("manifest.publication fault: {e}")))?;
6478 let active_runs = self.run_refs.clone();
6479 let schema_id = self.schema.schema_id;
6480 let index_generation = self.global_idx_epoch;
6481 let run_generation = active_runs.len() as u64;
6482 let fingerprint = crate::run_lookup::compute_fingerprint(
6483 &active_runs.iter().map(|r| r.run_id).collect::<Vec<_>>(),
6484 schema_id,
6485 index_generation,
6486 run_generation,
6487 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
6488 );
6489 let dir = match crate::run_lookup::RunLookupDirectory::rebuild_from_runs(&active_runs, self)
6490 {
6491 Ok(d) => d,
6492 Err(error) => {
6493 self.run_lookup.complete = false;
6496 self.lookup_metrics
6497 .directory_incomplete
6498 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6499 return Err(error);
6500 }
6501 };
6502 let mut dir = dir;
6506 dir.set_fingerprint(fingerprint);
6507 let runs_dir = match self.runs_root.as_deref() {
6508 Some(root) => match root.io_path() {
6509 Ok(p) => p,
6510 Err(_) => self.dir.join(RUNS_DIR),
6511 },
6512 None => self.dir.join(RUNS_DIR),
6513 };
6514 mongreldb_fault::inject("directory.temp_write")
6517 .map_err(|e| MongrelError::Other(format!("directory.temp_write fault: {e}")))?;
6518 let publish =
6519 dir.write_checkpoint_sharded(&runs_dir, crate::run_lookup::DEFAULT_SHARD_BYTES);
6520 mongreldb_fault::inject("directory.sync")
6521 .map_err(|e| MongrelError::Other(format!("directory.sync fault: {e}")))?;
6522 mongreldb_fault::inject("directory.rename")
6523 .map_err(|e| MongrelError::Other(format!("directory.rename fault: {e}")))?;
6524 if let Err(error) = publish {
6525 self.run_lookup = RunLookupState {
6528 directory: None,
6529 fingerprint,
6530 complete: false,
6531 };
6532 self.lookup_metrics
6533 .directory_incomplete
6534 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6535 self.lookup_metrics
6536 .directory_lookup_fallback
6537 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6538 return Err(error.into());
6539 }
6540 self.run_lookup = RunLookupState {
6541 directory: Some(std::sync::Arc::new(dir)),
6542 fingerprint,
6543 complete: true,
6544 };
6545 Ok(())
6546 }
6547
6548 pub(crate) fn plan_recovered_metadata(&mut self) -> Result<RecoveryMetadataPlan> {
6549 let rows = self.visible_rows_at_time(Snapshot::unbounded(), i64::MIN)?;
6553 let live_count = u64::try_from(rows.len())
6554 .map_err(|_| MongrelError::Full("table live-row count exceeds u64".into()))?;
6555 let auto_inc = match self.auto_inc {
6556 Some(mut state) => {
6557 let maximum = self.scan_max_int64(state.column_id)?;
6558 let after_maximum = maximum.checked_add(1).ok_or_else(|| {
6559 MongrelError::Full("AUTO_INCREMENT namespace exhausted".into())
6560 })?;
6561 state.next = state.next.max(after_maximum).max(1);
6562 state.seeded = true;
6563 Some(state)
6564 }
6565 None => None,
6566 };
6567 Ok(RecoveryMetadataPlan {
6568 live_count,
6569 auto_inc,
6570 changed: live_count != self.live_count
6571 || auto_inc.is_some_and(|planned| {
6572 self.auto_inc.is_none_or(|current| {
6573 current.next != planned.next || current.seeded != planned.seeded
6574 })
6575 }),
6576 })
6577 }
6578
6579 pub(crate) fn apply_recovered_metadata(
6580 &mut self,
6581 plan: RecoveryMetadataPlan,
6582 epoch: Epoch,
6583 ) -> Result<()> {
6584 if !plan.changed {
6585 return Ok(());
6586 }
6587 self.live_count = plan.live_count;
6588 self.auto_inc = plan.auto_inc;
6589 self.persist_manifest(epoch)
6590 }
6591
6592 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
6598 if !self.indexes_complete {
6601 return;
6602 }
6603 if crate::catalog::inject_hook("index.publish.before").is_err() {
6606 return;
6607 }
6608 if self.idx_root.is_none() {
6609 if let Some(root) = self._root_guard.as_ref() {
6610 let Ok(idx_root) = root.create_directory_all_pinned(global_idx::IDX_DIR) else {
6611 return;
6612 };
6613 self.idx_root = Some(Arc::new(idx_root));
6614 }
6615 }
6616 let snap = global_idx::IndexSnapshot {
6617 hot: &self.hot,
6618 bitmap: &self.bitmap,
6619 ann: &self.ann,
6620 fm: &self.fm,
6621 sparse: &self.sparse,
6622 minhash: &self.minhash,
6623 learned_range: &self.learned_range,
6624 };
6625 let idx_dek = self.idx_dek();
6627 let written = match self.idx_root.as_deref() {
6628 Some(root) => global_idx::write_atomic_root(
6629 root,
6630 self.table_id,
6631 epoch.0,
6632 snap,
6633 idx_dek.as_deref(),
6634 ),
6635 None => global_idx::write_atomic(
6636 &self.dir,
6637 self.table_id,
6638 epoch.0,
6639 snap,
6640 idx_dek.as_deref(),
6641 ),
6642 };
6643 if written.is_ok() {
6644 self.global_idx_epoch = epoch.0;
6645 let _ = self.persist_manifest(epoch);
6646 let _ = crate::catalog::inject_hook("index.publish.after");
6648 }
6649 }
6650
6651 pub(crate) fn invalidate_index_checkpoint(&mut self) {
6654 self.global_idx_epoch = 0;
6655 if let Some(root) = self.idx_root.as_deref() {
6656 let _ = root.remove_file(global_idx::IDX_FILENAME);
6657 } else {
6658 global_idx::remove(&self.dir);
6659 }
6660 let _ = self.persist_manifest(self.epoch.visible());
6661 }
6662
6663 pub(crate) fn prepare_indexes_for_run_replacement(&mut self) {
6668 self.indexes_complete = false;
6669 self.global_idx_epoch = 0;
6670 if let Some(root) = self.idx_root.as_deref() {
6671 let _ = root.remove_file(global_idx::IDX_FILENAME);
6672 } else {
6673 global_idx::remove(&self.dir);
6674 }
6675 }
6676
6677 pub(crate) fn finish_indexes_for_run_replacement(&mut self) {
6678 self.indexes_complete = true;
6679 }
6680
6681 pub(crate) fn poison_after_maintenance_publish_failure(&mut self) {
6687 self.durable_commit_failed = true;
6688 if let WalSink::Shared(shared) = &self.wal {
6689 shared
6690 .poisoned
6691 .store(true, std::sync::atomic::Ordering::Relaxed);
6692 }
6693 }
6694
6695 pub(crate) fn mark_unavailable_after_quarantine(&mut self) {
6699 self.durable_commit_failed = true;
6700 }
6701
6702 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
6716 let mut best: Option<Row> = None;
6717 let mut consider = |row: Row| {
6718 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6719 return;
6720 }
6721 if best.as_ref().is_none_or(|current| {
6722 Snapshot::version_is_newer(
6723 row.committed_epoch,
6724 row.commit_ts,
6725 current.committed_epoch,
6726 current.commit_ts,
6727 )
6728 }) {
6729 best = Some(row);
6730 }
6731 };
6732 if let Some((_, row)) = self.memtable.get_version_at(row_id, snapshot) {
6733 consider(row);
6734 }
6735 if let Some((_, row)) = self.mutable_run.get_version_at(row_id, snapshot) {
6736 consider(row);
6737 }
6738 let dir_locators = if self.run_lookup.complete {
6742 let active_ids: Vec<u128> = self.run_refs.iter().map(|r| r.run_id).collect();
6747 let expected_fp = crate::run_lookup::compute_fingerprint(
6748 &active_ids,
6749 self.schema.schema_id,
6750 self.global_idx_epoch,
6751 active_ids.len() as u64,
6752 crate::run_lookup::DIRECTORY_FORMAT_VERSION,
6753 );
6754 if self.run_lookup.fingerprint == expected_fp {
6755 self.run_lookup
6756 .directory
6757 .as_ref()
6758 .map(|d| d.locate(row_id).locators.clone())
6759 } else {
6760 self.lookup_metrics
6761 .directory_lookup_fallback
6762 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6763 None
6764 }
6765 } else {
6766 self.lookup_metrics
6767 .directory_lookup_fallback
6768 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6769 None
6770 };
6771 match dir_locators {
6772 Some(locators) if !locators.is_empty() => {
6773 let mut opened: std::collections::HashSet<u128> = std::collections::HashSet::new();
6779 for locator in &locators {
6780 if locator.is_impossible_for(snapshot) {
6781 continue;
6782 }
6783 if !opened.insert(locator.run_id) {
6784 continue;
6785 }
6786 self.lookup_metrics
6787 .directory_run_readers_opened
6788 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6789 self.lookup_metrics
6790 .get_run_opened
6791 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6792 let Ok(mut reader) = self.open_reader(locator.run_id) else {
6793 continue;
6794 };
6795 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
6798 continue;
6799 };
6800 consider(row);
6801 }
6802 if opened.is_empty() {
6803 self.lookup_metrics
6804 .directory_lookup_fallback
6805 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6806 } else {
6807 self.lookup_metrics
6808 .directory_lookup_hit
6809 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6810 }
6811 }
6812 Some(_) | None => {
6813 for rr in &self.run_refs {
6815 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
6819 if row_id.0 < min_rid || row_id.0 > max_rid {
6820 self.lookup_metrics
6821 .get_run_skipped
6822 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6823 continue;
6824 }
6825 }
6826 self.lookup_metrics
6827 .get_run_opened
6828 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
6829 let Ok(mut reader) = self.open_reader(rr.run_id) else {
6830 continue;
6831 };
6832 let Ok(Some((_, row))) = reader.get_version_at(row_id, snapshot) else {
6835 continue;
6836 };
6837 consider(row);
6838 }
6839 }
6840 }
6841 let now_nanos = unix_nanos_now();
6842 match best {
6843 Some(r) if r.deleted || self.row_expired_at(&r, now_nanos) => None,
6844 Some(r) => Some(r),
6845 None => None,
6846 }
6847 }
6848
6849 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
6853 self.visible_rows_at_time(snapshot, unix_nanos_now())
6854 }
6855
6856 #[doc(hidden)]
6859 pub fn visible_rows_controlled(
6860 &self,
6861 snapshot: Snapshot,
6862 control: &crate::ExecutionControl,
6863 ) -> Result<Vec<Row>> {
6864 let mut out = Vec::new();
6865 self.for_each_visible_row_controlled(snapshot, control, |row| {
6866 out.push(row);
6867 Ok(())
6868 })?;
6869 Ok(out)
6870 }
6871
6872 #[doc(hidden)]
6875 pub fn for_each_visible_row_controlled<F>(
6876 &self,
6877 snapshot: Snapshot,
6878 control: &crate::ExecutionControl,
6879 mut visit: F,
6880 ) -> Result<()>
6881 where
6882 F: FnMut(Row) -> Result<()>,
6883 {
6884 let setup_start = std::time::Instant::now();
6891 let mut versions_examined: u64 = 0;
6892 let mut rows_emitted: u64 = 0;
6893 let mut checkpoints: u64 = 0;
6894 let mut first_row_recorded = false;
6895 let mut sources: Vec<ControlledVisibleSource<'_>> =
6896 Vec::with_capacity(self.run_refs.len() + 2);
6897 control.checkpoint()?;
6898 checkpoints += 1;
6899 if !self.memtable.is_empty() {
6902 sources.push(ControlledVisibleSource::memtable_cursor(
6903 self.memtable.newest_visible_iter(&snapshot),
6904 ));
6905 }
6906 control.checkpoint()?;
6907 checkpoints += 1;
6908 if !self.mutable_run.is_empty() {
6909 sources.push(ControlledVisibleSource::mutable_run_cursor(
6910 self.mutable_run.newest_visible_iter(&snapshot),
6911 ));
6912 }
6913 for run in &self.run_refs {
6914 control.checkpoint()?;
6915 checkpoints += 1;
6916 let reader = self.open_reader(run.run_id)?;
6917 sources.push(ControlledVisibleSource::run(
6918 reader.into_visible_version_cursor_at(snapshot)?,
6919 ));
6920 }
6921 let first_row_start = setup_start;
6924 let _setup_us = setup_start.elapsed().as_micros() as u64;
6925 let now_nanos = unix_nanos_now();
6926 let mut first_row_us: u64 = 0;
6927 let result = merge_controlled_visible_sources(
6928 &mut sources,
6929 control,
6930 |row| self.row_expired_at(row, now_nanos),
6931 |row| {
6932 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6933 return Ok(());
6934 }
6935 versions_examined += 1;
6936 if !first_row_recorded {
6937 first_row_us = first_row_start.elapsed().as_micros() as u64;
6938 first_row_recorded = true;
6939 }
6940 rows_emitted += 1;
6941 visit(row)
6942 },
6943 );
6944 let time_to_first_row_us = if first_row_recorded { first_row_us } else { 0 };
6945 crate::trace::QueryTrace::record(|t| {
6946 t.controlled_scan_versions_examined = t
6947 .controlled_scan_versions_examined
6948 .saturating_add(versions_examined as usize);
6949 t.controlled_scan_rows_emitted = t
6950 .controlled_scan_rows_emitted
6951 .saturating_add(rows_emitted as usize);
6952 t.controlled_scan_checkpoints = t
6953 .controlled_scan_checkpoints
6954 .saturating_add(checkpoints as usize);
6955 t.controlled_scan_time_to_first_row_us = t
6956 .controlled_scan_time_to_first_row_us
6957 .saturating_add(time_to_first_row_us);
6958 t.controlled_scan_setup_time_us =
6959 t.controlled_scan_setup_time_us.saturating_add(_setup_us);
6960 });
6961 result
6962 }
6963
6964 #[doc(hidden)]
6965 pub fn visible_rows_at_time(&self, snapshot: Snapshot, now_nanos: i64) -> Result<Vec<Row>> {
6966 let mut best: HashMap<u64, Row> = HashMap::new();
6967 let mut fold = |row: Row| {
6968 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
6969 return;
6970 }
6971 best.entry(row.row_id.0)
6972 .and_modify(|existing| {
6973 if Snapshot::version_is_newer(
6974 row.committed_epoch,
6975 row.commit_ts,
6976 existing.committed_epoch,
6977 existing.commit_ts,
6978 ) {
6979 *existing = row.clone();
6980 }
6981 })
6982 .or_insert(row);
6983 };
6984 for row in self.memtable.visible_versions_at(snapshot) {
6985 fold(row);
6986 }
6987 for row in self.mutable_run.visible_versions_at(snapshot) {
6988 fold(row);
6989 }
6990 for rr in &self.run_refs {
6991 let mut reader = self.open_reader(rr.run_id)?;
6992 for row in reader.visible_versions_at(snapshot)? {
6994 fold(row);
6995 }
6996 }
6997 let mut out: Vec<Row> = best
6998 .into_values()
6999 .filter_map(|r| {
7000 if r.deleted || self.row_expired_at(&r, now_nanos) {
7001 None
7002 } else {
7003 Some(r)
7004 }
7005 })
7006 .collect();
7007 out.sort_by_key(|r| r.row_id);
7008 Ok(out)
7009 }
7010
7011 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
7018 if self.ttl.is_none()
7019 && self.memtable.is_empty()
7020 && self.mutable_run.is_empty()
7021 && self.run_refs.len() == 1
7022 {
7023 let rr = self.run_refs[0].clone();
7024 let mut reader = self.open_reader(rr.run_id)?;
7025 let idxs = reader.visible_indices_at(snapshot)?;
7026 let mut cols = Vec::with_capacity(self.schema.columns.len());
7027 for cdef in &self.schema.columns {
7028 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
7029 }
7030 return Ok(cols);
7031 }
7032 let rows = self.visible_rows(snapshot)?;
7034 let mut cols: Vec<(u16, Vec<Value>)> = self
7035 .schema
7036 .columns
7037 .iter()
7038 .map(|c| (c.id, Vec::with_capacity(rows.len())))
7039 .collect();
7040 for r in &rows {
7041 for (cid, vec) in cols.iter_mut() {
7042 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
7043 }
7044 }
7045 Ok(cols)
7046 }
7047
7048 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
7050 let row_id = self.hot.get(key)?;
7051 if self.ttl.is_none() || self.get(row_id, Snapshot::unbounded()).is_some() {
7052 Some(row_id)
7053 } else {
7054 None
7055 }
7056 }
7057
7058 pub fn lookup_metrics_snapshot(&self) -> LookupMetricsSnapshot {
7062 let mut snap = self.lookup_metrics.snapshot();
7063 let cache = self.result_cache.lock();
7064 let (mem, disk, miss, write_us) = cache.cache_counters();
7065 snap.result_cache_memory_hit = mem;
7066 snap.result_cache_disk_hit = disk;
7067 snap.result_cache_miss = miss;
7068 snap.result_cache_persistent_write_us = write_us;
7069 if let Some(writer_snap) = cache.persist_snapshot() {
7072 snap.result_cache_persist_enqueued_total = writer_snap.result_cache_persist_enqueued_total;
7073 snap.result_cache_persist_coalesced_total = writer_snap.result_cache_persist_coalesced_total;
7074 snap.result_cache_persist_dropped_store_total = writer_snap.result_cache_persist_dropped_store_total;
7075 snap.result_cache_persist_remove_total = writer_snap.result_cache_persist_remove_total;
7076 snap.result_cache_persist_stale_store_skipped_total = writer_snap.result_cache_persist_stale_store_skipped_total;
7077 snap.result_cache_persist_errors_total = writer_snap.result_cache_persist_errors_total;
7078 snap.result_cache_persist_shutdown_abandoned_total =
7079 writer_snap.result_cache_persist_shutdown_abandoned_total;
7080 snap.result_cache_persist_queue_depth = writer_snap.result_cache_persist_queue_depth;
7081 }
7082 snap
7083 }
7084
7085 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
7090 self.query_at_with_allowed(q, self.snapshot(), None)
7091 }
7092
7093 pub fn query_controlled(
7096 &mut self,
7097 q: &crate::query::Query,
7098 control: &crate::ExecutionControl,
7099 ) -> Result<Vec<Row>> {
7100 self.query_at_with_allowed_controlled(q, self.snapshot(), None, control)
7101 }
7102
7103 pub fn query_at_with_allowed(
7106 &mut self,
7107 q: &crate::query::Query,
7108 snapshot: Snapshot,
7109 allowed: Option<&std::collections::HashSet<RowId>>,
7110 ) -> Result<Vec<Row>> {
7111 self.query_at_with_allowed_after(q, snapshot, allowed, None)
7112 }
7113
7114 #[doc(hidden)]
7115 pub fn query_at_with_allowed_controlled(
7116 &mut self,
7117 q: &crate::query::Query,
7118 snapshot: Snapshot,
7119 allowed: Option<&std::collections::HashSet<RowId>>,
7120 control: &crate::ExecutionControl,
7121 ) -> Result<Vec<Row>> {
7122 self.require_select()?;
7123 self.ensure_indexes_complete_controlled(control, || true)?;
7124 self.validate_native_query(q)?;
7125 self.query_conditions_at(
7126 &q.conditions,
7127 snapshot,
7128 allowed,
7129 q.limit,
7130 q.offset,
7131 None,
7132 unix_nanos_now(),
7133 Some(control),
7134 )
7135 }
7136
7137 #[doc(hidden)]
7138 pub fn query_at_with_allowed_after(
7139 &mut self,
7140 q: &crate::query::Query,
7141 snapshot: Snapshot,
7142 allowed: Option<&std::collections::HashSet<RowId>>,
7143 after_row_id: Option<RowId>,
7144 ) -> Result<Vec<Row>> {
7145 self.query_at_with_allowed_after_at_time(
7146 q,
7147 snapshot,
7148 allowed,
7149 after_row_id,
7150 unix_nanos_now(),
7151 )
7152 }
7153
7154 #[doc(hidden)]
7155 pub fn query_at_with_allowed_after_at_time(
7156 &mut self,
7157 q: &crate::query::Query,
7158 snapshot: Snapshot,
7159 allowed: Option<&std::collections::HashSet<RowId>>,
7160 after_row_id: Option<RowId>,
7161 query_time_nanos: i64,
7162 ) -> Result<Vec<Row>> {
7163 self.require_select()?;
7164 self.ensure_indexes_complete()?;
7165 self.validate_native_query(q)?;
7166 self.query_conditions_at(
7167 &q.conditions,
7168 snapshot,
7169 allowed,
7170 q.limit,
7171 q.offset,
7172 after_row_id,
7173 query_time_nanos,
7174 None,
7175 )
7176 }
7177
7178 fn validate_native_query(&self, q: &crate::query::Query) -> Result<()> {
7179 if q.conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7180 return Err(MongrelError::InvalidArgument(format!(
7181 "query exceeds {} conditions",
7182 crate::query::MAX_HARD_CONDITIONS
7183 )));
7184 }
7185 if let Some(limit) = q.limit {
7186 if limit == 0 || limit > crate::query::MAX_FINAL_LIMIT {
7187 return Err(MongrelError::InvalidArgument(format!(
7188 "query limit must be between 1 and {}",
7189 crate::query::MAX_FINAL_LIMIT
7190 )));
7191 }
7192 }
7193 if q.offset > crate::query::MAX_QUERY_OFFSET {
7194 return Err(MongrelError::InvalidArgument(format!(
7195 "query offset exceeds {}",
7196 crate::query::MAX_QUERY_OFFSET
7197 )));
7198 }
7199 Ok(())
7200 }
7201
7202 #[doc(hidden)]
7205 pub fn query_all_at(
7206 &mut self,
7207 conditions: &[crate::query::Condition],
7208 snapshot: Snapshot,
7209 ) -> Result<Vec<Row>> {
7210 self.require_select()?;
7211 self.ensure_indexes_complete()?;
7212 if conditions.len() > crate::query::MAX_HARD_CONDITIONS {
7213 return Err(MongrelError::InvalidArgument(format!(
7214 "query exceeds {} conditions",
7215 crate::query::MAX_HARD_CONDITIONS
7216 )));
7217 }
7218 self.query_conditions_at(
7219 conditions,
7220 snapshot,
7221 None,
7222 None,
7223 0,
7224 None,
7225 unix_nanos_now(),
7226 None,
7227 )
7228 }
7229
7230 #[allow(clippy::too_many_arguments)]
7231 fn query_conditions_at(
7232 &self,
7233 conditions: &[crate::query::Condition],
7234 snapshot: Snapshot,
7235 allowed: Option<&std::collections::HashSet<RowId>>,
7236 limit: Option<usize>,
7237 offset: usize,
7238 after_row_id: Option<RowId>,
7239 query_time_nanos: i64,
7240 control: Option<&crate::ExecutionControl>,
7241 ) -> Result<Vec<Row>> {
7242 control
7243 .map(crate::ExecutionControl::checkpoint)
7244 .transpose()?;
7245 crate::trace::QueryTrace::record(|t| {
7246 t.run_count = self.run_refs.len();
7247 t.memtable_rows = self.memtable.len();
7248 t.mutable_run_rows = self.mutable_run.len();
7249 });
7250 if conditions.is_empty() {
7254 crate::trace::QueryTrace::record(|t| {
7255 t.scan_mode = crate::trace::ScanMode::Materialized;
7256 t.row_materialized = true;
7257 });
7258 let mut rows = match control {
7259 Some(control) => self.visible_rows_controlled(snapshot, control)?,
7260 None => self.visible_rows_at_time(snapshot, query_time_nanos)?,
7261 };
7262 if let Some(allowed) = allowed {
7263 let mut filtered = Vec::with_capacity(rows.len());
7264 for (index, row) in rows.into_iter().enumerate() {
7265 if index & 255 == 0 {
7266 control
7267 .map(crate::ExecutionControl::checkpoint)
7268 .transpose()?;
7269 }
7270 if allowed.contains(&row.row_id) {
7271 filtered.push(row);
7272 }
7273 }
7274 rows = filtered;
7275 }
7276 if let Some(after_row_id) = after_row_id {
7277 rows.retain(|row| row.row_id > after_row_id);
7278 }
7279 rows.drain(..offset.min(rows.len()));
7280 if let Some(limit) = limit {
7281 rows.truncate(limit);
7282 }
7283 return Ok(rows);
7284 }
7285 crate::trace::QueryTrace::record(|t| {
7286 t.conditions_pushed = conditions.len();
7287 t.scan_mode = crate::trace::ScanMode::Materialized;
7288 t.row_materialized = true;
7289 });
7290 let mut ordered: Vec<&crate::query::Condition> = conditions.iter().collect();
7297 ordered.sort_by_key(|c| condition_cost_rank(c));
7298 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
7299 for c in &ordered {
7300 control
7301 .map(crate::ExecutionControl::checkpoint)
7302 .transpose()?;
7303 let s = self.resolve_condition_with_allowed(c, snapshot, allowed)?;
7304 let empty = s.is_empty();
7305 sets.push(s);
7306 if empty {
7307 break;
7308 }
7309 }
7310 let mut rids = RowIdSet::intersect_many(sets).into_sorted_vec();
7311 if let Some(allowed) = allowed {
7312 rids.retain(|row_id| allowed.contains(&RowId(*row_id)));
7313 }
7314 if let Some(after_row_id) = after_row_id {
7315 let first = rids.partition_point(|row_id| *row_id <= after_row_id.0);
7316 rids.drain(..first);
7317 }
7318 rids.drain(..offset.min(rids.len()));
7319 if let Some(limit) = limit {
7320 rids.truncate(limit);
7321 }
7322 control
7323 .map(crate::ExecutionControl::checkpoint)
7324 .transpose()?;
7325 self.rows_for_rids_at_time(&rids, snapshot, query_time_nanos, control)
7326 }
7327
7328 pub fn retrieve(
7330 &mut self,
7331 retriever: &crate::query::Retriever,
7332 ) -> Result<Vec<crate::query::RetrieverHit>> {
7333 self.retrieve_with_allowed(retriever, None)
7334 }
7335
7336 pub fn retrieve_at(
7337 &mut self,
7338 retriever: &crate::query::Retriever,
7339 snapshot: Snapshot,
7340 allowed: Option<&std::collections::HashSet<RowId>>,
7341 ) -> Result<Vec<crate::query::RetrieverHit>> {
7342 self.retrieve_at_with_allowed(retriever, snapshot, allowed)
7343 }
7344
7345 pub fn retrieve_with_allowed(
7348 &mut self,
7349 retriever: &crate::query::Retriever,
7350 allowed: Option<&std::collections::HashSet<RowId>>,
7351 ) -> Result<Vec<crate::query::RetrieverHit>> {
7352 self.retrieve_at_with_allowed(retriever, self.snapshot(), allowed)
7353 }
7354
7355 pub fn retrieve_at_with_allowed(
7356 &mut self,
7357 retriever: &crate::query::Retriever,
7358 snapshot: Snapshot,
7359 allowed: Option<&std::collections::HashSet<RowId>>,
7360 ) -> Result<Vec<crate::query::RetrieverHit>> {
7361 self.retrieve_at_with_allowed_and_context(retriever, snapshot, allowed, None)
7362 }
7363
7364 pub fn retrieve_at_with_allowed_and_context(
7365 &mut self,
7366 retriever: &crate::query::Retriever,
7367 snapshot: Snapshot,
7368 allowed: Option<&std::collections::HashSet<RowId>>,
7369 context: Option<&crate::query::AiExecutionContext>,
7370 ) -> Result<Vec<crate::query::RetrieverHit>> {
7371 self.require_select()?;
7372 self.ensure_indexes_complete()?;
7373 self.validate_retriever(retriever)?;
7374 self.retrieve_filtered(retriever, snapshot, None, allowed, None, context)
7375 }
7376
7377 pub fn retrieve_at_with_candidate_authorization_and_context(
7378 &mut self,
7379 retriever: &crate::query::Retriever,
7380 snapshot: Snapshot,
7381 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7382 context: Option<&crate::query::AiExecutionContext>,
7383 ) -> Result<Vec<crate::query::RetrieverHit>> {
7384 self.require_select()?;
7385 self.ensure_indexes_complete()?;
7386 self.retrieve_at_with_candidate_authorization_on_generation(
7387 retriever,
7388 snapshot,
7389 authorization,
7390 context,
7391 )
7392 }
7393
7394 #[doc(hidden)]
7395 pub fn retrieve_at_with_candidate_authorization_on_generation(
7396 &self,
7397 retriever: &crate::query::Retriever,
7398 snapshot: Snapshot,
7399 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7400 context: Option<&crate::query::AiExecutionContext>,
7401 ) -> Result<Vec<crate::query::RetrieverHit>> {
7402 self.require_select()?;
7403 self.validate_retriever(retriever)?;
7404 self.retrieve_filtered(retriever, snapshot, None, None, authorization, context)
7405 }
7406
7407 fn validate_retriever(&self, retriever: &crate::query::Retriever) -> Result<()> {
7408 use crate::query::{Retriever, MAX_RETRIEVER_K, MAX_SET_MEMBERS, MAX_SPARSE_TERMS};
7409 let (column_id, k) = match retriever {
7410 Retriever::Ann {
7411 column_id,
7412 query,
7413 k,
7414 } => {
7415 let index = self.ann.get(column_id).ok_or_else(|| {
7416 MongrelError::InvalidArgument(format!("column {column_id} has no ANN index"))
7417 })?;
7418 if query.len() != index.dim() {
7419 return Err(MongrelError::InvalidArgument(format!(
7420 "ANN query dimension must be {}, got {}",
7421 index.dim(),
7422 query.len()
7423 )));
7424 }
7425 if query.iter().any(|value| !value.is_finite()) {
7426 return Err(MongrelError::InvalidArgument(
7427 "ANN query values must be finite".into(),
7428 ));
7429 }
7430 (*column_id, *k)
7431 }
7432 Retriever::Sparse {
7433 column_id,
7434 query,
7435 k,
7436 } => {
7437 if !self.sparse.contains_key(column_id) {
7438 return Err(MongrelError::InvalidArgument(format!(
7439 "column {column_id} has no Sparse index"
7440 )));
7441 }
7442 if query.is_empty() || query.iter().any(|(_, weight)| !weight.is_finite()) {
7443 return Err(MongrelError::InvalidArgument(
7444 "Sparse query must be non-empty with finite weights".into(),
7445 ));
7446 }
7447 if query.len() > MAX_SPARSE_TERMS {
7448 return Err(MongrelError::InvalidArgument(format!(
7449 "Sparse query exceeds {MAX_SPARSE_TERMS} terms"
7450 )));
7451 }
7452 (*column_id, *k)
7453 }
7454 Retriever::MinHash {
7455 column_id,
7456 members,
7457 k,
7458 } => {
7459 if !self.minhash.contains_key(column_id) {
7460 return Err(MongrelError::InvalidArgument(format!(
7461 "column {column_id} has no MinHash index"
7462 )));
7463 }
7464 if members.is_empty() {
7465 return Err(MongrelError::InvalidArgument(
7466 "MinHash members must not be empty".into(),
7467 ));
7468 }
7469 if members.len() > MAX_SET_MEMBERS {
7470 return Err(MongrelError::InvalidArgument(format!(
7471 "MinHash query exceeds {MAX_SET_MEMBERS} members"
7472 )));
7473 }
7474 let mut total_bytes = 0usize;
7475 for member in members {
7476 let bytes = member.encoded_len();
7477 if bytes > crate::query::MAX_SET_MEMBER_BYTES {
7478 return Err(MongrelError::InvalidArgument(format!(
7479 "MinHash member exceeds {} bytes",
7480 crate::query::MAX_SET_MEMBER_BYTES
7481 )));
7482 }
7483 total_bytes = total_bytes.checked_add(bytes).ok_or_else(|| {
7484 MongrelError::InvalidArgument("MinHash input size overflow".into())
7485 })?;
7486 }
7487 if total_bytes > crate::query::MAX_SET_INPUT_BYTES {
7488 return Err(MongrelError::InvalidArgument(format!(
7489 "MinHash input exceeds {} bytes",
7490 crate::query::MAX_SET_INPUT_BYTES
7491 )));
7492 }
7493 (*column_id, *k)
7494 }
7495 };
7496 if k == 0 {
7497 return Err(MongrelError::InvalidArgument(
7498 "retriever k must be > 0".into(),
7499 ));
7500 }
7501 if k > MAX_RETRIEVER_K {
7502 return Err(MongrelError::InvalidArgument(format!(
7503 "retriever k exceeds {MAX_RETRIEVER_K}"
7504 )));
7505 }
7506 debug_assert!(self
7507 .schema
7508 .columns
7509 .iter()
7510 .any(|column| column.id == column_id));
7511 Ok(())
7512 }
7513
7514 fn validate_condition(&self, condition: &crate::query::Condition) -> Result<()> {
7515 use crate::query::Condition;
7516 match condition {
7517 Condition::Ann {
7518 column_id,
7519 query,
7520 k,
7521 } => self.validate_retriever(&crate::query::Retriever::Ann {
7522 column_id: *column_id,
7523 query: query.clone(),
7524 k: *k,
7525 }),
7526 Condition::SparseMatch {
7527 column_id,
7528 query,
7529 k,
7530 } => self.validate_retriever(&crate::query::Retriever::Sparse {
7531 column_id: *column_id,
7532 query: query.clone(),
7533 k: *k,
7534 }),
7535 Condition::MinHashSimilar {
7536 column_id,
7537 query,
7538 k,
7539 } => {
7540 if !self.minhash.contains_key(column_id) {
7541 return Err(MongrelError::InvalidArgument(format!(
7542 "column {column_id} has no MinHash index"
7543 )));
7544 }
7545 if query.is_empty() || *k == 0 {
7546 return Err(MongrelError::InvalidArgument(
7547 "MinHash query must be non-empty and k must be > 0".into(),
7548 ));
7549 }
7550 if query.len() > crate::query::MAX_SET_MEMBERS || *k > crate::query::MAX_RETRIEVER_K
7551 {
7552 return Err(MongrelError::InvalidArgument(format!(
7553 "MinHash query must have <= {} members and k <= {}",
7554 crate::query::MAX_SET_MEMBERS,
7555 crate::query::MAX_RETRIEVER_K
7556 )));
7557 }
7558 Ok(())
7559 }
7560 Condition::BitmapIn { values, .. } if values.len() > crate::query::MAX_SET_MEMBERS => {
7561 Err(MongrelError::InvalidArgument(format!(
7562 "bitmap IN exceeds {} values",
7563 crate::query::MAX_SET_MEMBERS
7564 )))
7565 }
7566 Condition::FmContainsAll { patterns, .. }
7567 if patterns.len() > crate::query::MAX_HARD_CONDITIONS =>
7568 {
7569 Err(MongrelError::InvalidArgument(format!(
7570 "FM query exceeds {} patterns",
7571 crate::query::MAX_HARD_CONDITIONS
7572 )))
7573 }
7574 _ => Ok(()),
7575 }
7576 }
7577
7578 fn retrieve_filtered(
7579 &self,
7580 retriever: &crate::query::Retriever,
7581 snapshot: Snapshot,
7582 hard_filter: Option<&RowIdSet>,
7583 allowed: Option<&std::collections::HashSet<RowId>>,
7584 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
7585 context: Option<&crate::query::AiExecutionContext>,
7586 ) -> Result<Vec<crate::query::RetrieverHit>> {
7587 use crate::query::{Retriever, RetrieverHit, RetrieverScore};
7588 let started = std::time::Instant::now();
7589 let scored: Vec<(RowId, RetrieverScore)> = match retriever {
7590 Retriever::Ann {
7591 column_id,
7592 query,
7593 k,
7594 } => {
7595 let Some(index) = self.ann.get(column_id) else {
7596 return Ok(Vec::new());
7597 };
7598 let cap = ann_candidate_cap(index.len(), context);
7599 crate::trace::QueryTrace::record(|trace| trace.candidate_cap = cap);
7600 if cap == 0 {
7601 return Ok(Vec::new());
7602 }
7603 let mut breadth = (*k).max(1).min(cap);
7604 let mut eligibility = std::collections::HashMap::new();
7605 let mut raw_candidates_total: usize = 0;
7610 let mut visibility_rejected_total: usize = 0;
7611 let mut authorization_rejected_total: usize = 0;
7612 let mut candidate_cap_hit_final = false;
7613 let mut filtered = loop {
7614 let mut seen = std::collections::HashSet::new();
7615 if let Some(context) = context {
7616 context.checkpoint()?;
7617 }
7618 let raw = index.search_with_context(query, breadth, context)?;
7619 crate::trace::QueryTrace::record(|trace| {
7620 trace.raw_candidates = raw.len();
7621 let unique = raw
7622 .iter()
7623 .map(|(row_id, _)| *row_id)
7624 .collect::<std::collections::HashSet<_>>()
7625 .len();
7626 trace.unique_candidates = unique;
7627 trace.duplicate_candidates = raw.len().saturating_sub(unique);
7628 trace.authorization_rejected = raw
7629 .iter()
7630 .filter(|(row_id, _)| {
7631 allowed.is_some_and(|allowed| !allowed.contains(row_id))
7632 })
7633 .count();
7634 trace.hard_filter_rejected = raw
7635 .iter()
7636 .filter(|(row_id, _)| {
7637 hard_filter.is_some_and(|filter| !filter.contains(row_id.0))
7638 })
7639 .count();
7640 });
7641 raw_candidates_total = raw_candidates_total.saturating_add(raw.len());
7642 let unchecked: Vec<_> = raw
7643 .iter()
7644 .map(|(row_id, _)| *row_id)
7645 .filter(|row_id| !eligibility.contains_key(row_id))
7646 .filter(|row_id| {
7647 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7648 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7649 })
7650 .collect();
7651 let eligible = self.eligible_and_authorized_candidate_ids(
7652 &unchecked,
7653 *column_id,
7654 snapshot,
7655 candidate_authorization,
7656 context,
7657 )?;
7658 for row_id in &unchecked {
7659 eligibility.insert(*row_id, eligible.contains(row_id));
7660 }
7661 let new_visibility_or_auth_rejected = unchecked
7664 .iter()
7665 .filter(|row_id| !eligible.contains(row_id))
7666 .count();
7667 let mut new_auth_rejected = 0usize;
7672 if let Some(authorization) = candidate_authorization {
7673 if authorization.security.rls_enabled(authorization.table)
7674 && !authorization.principal.is_admin
7675 {
7676 new_auth_rejected = new_visibility_or_auth_rejected;
7682 }
7683 }
7684 let new_visibility_rejected =
7685 new_visibility_or_auth_rejected.saturating_sub(new_auth_rejected);
7686 visibility_rejected_total =
7687 visibility_rejected_total.saturating_add(new_visibility_rejected);
7688 authorization_rejected_total =
7689 authorization_rejected_total.saturating_add(new_auth_rejected);
7690 let filtered: Vec<_> = raw
7691 .into_iter()
7692 .filter(|(row_id, _)| {
7693 seen.insert(*row_id)
7694 && eligibility.get(row_id).copied().unwrap_or(false)
7695 })
7696 .map(|(row_id, score)| {
7697 let score = match score {
7698 crate::index::AnnDistance::Hamming(d) => {
7699 RetrieverScore::AnnHammingDistance(d)
7700 }
7701 crate::index::AnnDistance::Cosine(d) => {
7702 RetrieverScore::AnnCosineDistance(d)
7703 }
7704 };
7705 (row_id, score)
7706 })
7707 .collect();
7708 if filtered.len() >= *k || breadth >= cap {
7709 if filtered.len() < *k && index.len() > cap && breadth >= cap {
7710 crate::trace::QueryTrace::record(|trace| {
7711 trace.ann_candidate_cap_hit = true;
7712 trace.candidate_cap_hit = true;
7713 });
7714 candidate_cap_hit_final = true;
7715 }
7716 break filtered;
7717 }
7718 breadth = breadth.saturating_mul(2).min(cap);
7719 };
7720 let final_hits = filtered.len();
7723 let index_len = index.len();
7724 crate::trace::QueryTrace::record(|trace| {
7725 trace.candidate_cap = cap;
7726 trace.candidate_cap_hit = candidate_cap_hit_final;
7727 trace.ann_candidate_cap_hit = candidate_cap_hit_final;
7728 trace.raw_candidates =
7729 trace.raw_candidates.saturating_add(raw_candidates_total);
7730 trace.visibility_rejected = trace
7731 .visibility_rejected
7732 .saturating_add(visibility_rejected_total);
7733 trace.authorization_rejected = trace
7734 .authorization_rejected
7735 .saturating_add(authorization_rejected_total);
7736 trace.final_hits = trace.final_hits.saturating_add(final_hits);
7737 let _ = index_len;
7740 });
7741 filtered.truncate(*k);
7742 filtered
7743 }
7744 Retriever::Sparse {
7745 column_id,
7746 query,
7747 k,
7748 } => self
7749 .sparse
7750 .get(column_id)
7751 .map(|index| -> Result<Vec<_>> {
7752 let mut breadth = (*k).max(1);
7753 let mut eligibility = std::collections::HashMap::new();
7754 loop {
7755 if let Some(context) = context {
7756 context.checkpoint()?;
7757 }
7758 let raw = index.search_with_context(query, breadth, context)?;
7759 let unchecked: Vec<_> = raw
7760 .iter()
7761 .map(|(row_id, _)| *row_id)
7762 .filter(|row_id| !eligibility.contains_key(row_id))
7763 .filter(|row_id| {
7764 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7765 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7766 })
7767 .collect();
7768 let eligible = self.eligible_and_authorized_candidate_ids(
7769 &unchecked,
7770 *column_id,
7771 snapshot,
7772 candidate_authorization,
7773 context,
7774 )?;
7775 for row_id in unchecked {
7776 eligibility.insert(row_id, eligible.contains(&row_id));
7777 }
7778 let filtered: Vec<_> = raw
7779 .iter()
7780 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
7781 .take(*k)
7782 .map(|(row_id, score)| {
7783 (*row_id, RetrieverScore::SparseDotProduct(*score))
7784 })
7785 .collect();
7786 if filtered.len() >= *k || raw.len() < breadth {
7787 break Ok(filtered);
7788 }
7789 let next = breadth.saturating_mul(2);
7790 if next == breadth {
7791 break Ok(filtered);
7792 }
7793 breadth = next;
7794 }
7795 })
7796 .transpose()?
7797 .unwrap_or_default(),
7798 Retriever::MinHash {
7799 column_id,
7800 members,
7801 k,
7802 } => self
7803 .minhash
7804 .get(column_id)
7805 .map(|index| -> Result<Vec<_>> {
7806 let mut hashes = Vec::with_capacity(members.len());
7807 for member in members {
7808 if let Some(context) = context {
7809 context.consume(crate::query::work_units(
7810 member.encoded_len(),
7811 crate::query::PARSE_WORK_QUANTUM,
7812 ))?;
7813 }
7814 hashes.push(member.hash_v1());
7815 }
7816 let mut breadth = (*k).max(1);
7817 let mut eligibility = std::collections::HashMap::new();
7818 loop {
7819 if let Some(context) = context {
7820 context.checkpoint()?;
7821 }
7822 let raw = index.search_with_context(&hashes, breadth, context)?;
7823 let unchecked: Vec<_> = raw
7824 .iter()
7825 .map(|(row_id, _)| *row_id)
7826 .filter(|row_id| !eligibility.contains_key(row_id))
7827 .filter(|row_id| {
7828 hard_filter.is_none_or(|filter| filter.contains(row_id.0))
7829 && allowed.is_none_or(|allowed| allowed.contains(row_id))
7830 })
7831 .collect();
7832 let eligible = self.eligible_and_authorized_candidate_ids(
7833 &unchecked,
7834 *column_id,
7835 snapshot,
7836 candidate_authorization,
7837 context,
7838 )?;
7839 for row_id in unchecked {
7840 eligibility.insert(row_id, eligible.contains(&row_id));
7841 }
7842 let filtered: Vec<_> = raw
7843 .iter()
7844 .filter(|(row_id, _)| eligibility.get(row_id).copied().unwrap_or(false))
7845 .take(*k)
7846 .map(|(row_id, score)| {
7847 (*row_id, RetrieverScore::MinHashEstimatedJaccard(*score))
7848 })
7849 .collect();
7850 if filtered.len() >= *k || raw.len() < breadth {
7851 break Ok(filtered);
7852 }
7853 let next = breadth.saturating_mul(2);
7854 if next == breadth {
7855 break Ok(filtered);
7856 }
7857 breadth = next;
7858 }
7859 })
7860 .transpose()?
7861 .unwrap_or_default(),
7862 };
7863 let requested_k = match retriever {
7864 Retriever::Ann { k, .. }
7865 | Retriever::Sparse { k, .. }
7866 | Retriever::MinHash { k, .. } => *k,
7867 };
7868 crate::trace::QueryTrace::record(|trace| {
7869 trace.final_hits = scored.len();
7870 if scored.len() < requested_k {
7871 trace.underfill_reason = Some(if trace.candidate_cap_hit {
7872 "candidate_cap"
7873 } else {
7874 "eligible_candidates_exhausted"
7875 });
7876 }
7877 });
7878 let elapsed = started.elapsed().as_nanos() as u64;
7879 crate::trace::QueryTrace::record(|trace| {
7880 match retriever {
7881 Retriever::Ann { .. } => {
7882 trace.ann_candidate_nanos = trace.ann_candidate_nanos.saturating_add(elapsed);
7883 if let Retriever::Ann { column_id, .. } = retriever {
7884 if let Some(index) = self.ann.get(column_id) {
7885 trace.ann_algorithm = Some(index.algorithm());
7886 trace.ann_quantization = Some(index.quantization());
7887 trace.ann_backend = Some(index.backend_name());
7888 }
7889 }
7890 }
7891 Retriever::Sparse { .. } => {
7892 trace.sparse_candidate_nanos =
7893 trace.sparse_candidate_nanos.saturating_add(elapsed)
7894 }
7895 Retriever::MinHash { .. } => {
7896 trace.minhash_candidate_nanos =
7897 trace.minhash_candidate_nanos.saturating_add(elapsed)
7898 }
7899 }
7900 trace.candidate_count = trace.candidate_count.saturating_add(scored.len());
7901 });
7902 Ok(scored
7903 .into_iter()
7904 .enumerate()
7905 .map(|(rank, (row_id, score))| RetrieverHit {
7906 row_id,
7907 rank: rank + 1,
7908 score,
7909 })
7910 .collect())
7911 }
7912
7913 fn eligible_candidate_ids(
7914 &self,
7915 candidates: &[RowId],
7916 _column_id: u16,
7917 snapshot: Snapshot,
7918 context: Option<&crate::query::AiExecutionContext>,
7919 ) -> Result<std::collections::HashSet<RowId>> {
7920 let lookup_snapshot = if self.is_shared()
7930 || (self.pending_put_cols.is_empty() && self.pending_delete_rids.is_empty())
7931 || snapshot.epoch.0 >= self.pending_epoch().0
7932 {
7933 snapshot
7934 } else {
7935 Snapshot::at(self.pending_epoch())
7936 };
7937 let mut readers: Vec<_> = self
7938 .run_refs
7939 .iter()
7940 .map(|run| self.open_reader(run.run_id))
7941 .collect::<Result<_>>()?;
7942 let now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
7943 let mut eligible = std::collections::HashSet::with_capacity(candidates.len());
7944 for &row_id in candidates {
7945 if let Some(context) = context {
7946 context.consume(1)?;
7947 }
7948 let mem = self.memtable.get_version_at(row_id, lookup_snapshot);
7949 let mutable = self.mutable_run.get_version_at(row_id, lookup_snapshot);
7950 let overlay = match (mem, mutable) {
7951 (Some(left), Some(right)) => Some(
7952 if Snapshot::version_is_newer(
7953 left.1.committed_epoch,
7954 left.1.commit_ts,
7955 right.1.committed_epoch,
7956 right.1.commit_ts,
7957 ) {
7958 left
7959 } else {
7960 right
7961 },
7962 ),
7963 (Some(value), None) | (None, Some(value)) => Some(value),
7964 (None, None) => None,
7965 };
7966 if let Some((_, row)) = overlay {
7967 if !row.deleted && !self.row_expired_at(&row, now) {
7968 eligible.insert(row_id);
7969 }
7970 continue;
7971 }
7972 let mut best: Option<(Epoch, bool, usize)> = None;
7973 for (index, reader) in readers.iter_mut().enumerate() {
7974 if let Some((epoch, deleted)) =
7975 reader.get_version_visibility_at(row_id, lookup_snapshot)?
7976 {
7977 if best
7978 .as_ref()
7979 .map(|(best_epoch, ..)| epoch > *best_epoch)
7980 .unwrap_or(true)
7981 {
7982 best = Some((epoch, deleted, index));
7983 }
7984 }
7985 }
7986 let Some((_, false, reader_index)) = best else {
7987 continue;
7988 };
7989 if let Some(ttl) = self.ttl {
7990 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
7991 .get_version_column_at(row_id, lookup_snapshot, ttl.column_id)?
7992 {
7993 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
7994 continue;
7995 }
7996 }
7997 }
7998 eligible.insert(row_id);
7999 }
8000 Ok(eligible)
8001 }
8002
8003 fn eligible_and_authorized_candidate_ids(
8004 &self,
8005 candidates: &[RowId],
8006 column_id: u16,
8007 snapshot: Snapshot,
8008 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8009 context: Option<&crate::query::AiExecutionContext>,
8010 ) -> Result<std::collections::HashSet<RowId>> {
8011 let eligible = self.eligible_candidate_ids(candidates, column_id, snapshot, context)?;
8012 let Some(authorization) = authorization else {
8013 return Ok(eligible);
8014 };
8015 let candidates: Vec<_> = eligible.into_iter().collect();
8016 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, context)
8017 }
8018
8019 fn policy_allowed_candidate_ids(
8020 &self,
8021 candidates: &[RowId],
8022 snapshot: Snapshot,
8023 authorization: &crate::security::CandidateAuthorization<'_>,
8024 context: Option<&crate::query::AiExecutionContext>,
8025 ) -> Result<std::collections::HashSet<RowId>> {
8026 let started = std::time::Instant::now();
8027 if candidates.is_empty()
8028 || authorization.principal.is_admin
8029 || !authorization.security.rls_enabled(authorization.table)
8030 {
8031 return Ok(candidates.iter().copied().collect());
8032 }
8033 if let Some(context) = context {
8034 context.checkpoint()?;
8035 }
8036 let row_ids: Vec<_> = candidates.iter().map(|row_id| row_id.0).collect();
8037 let mut rows: std::collections::HashMap<RowId, Row> = candidates
8038 .iter()
8039 .map(|row_id| {
8040 (
8041 *row_id,
8042 Row {
8043 row_id: *row_id,
8044 committed_epoch: snapshot.epoch,
8045 columns: std::collections::HashMap::new(),
8046 deleted: false,
8047 commit_ts: None,
8048 },
8049 )
8050 })
8051 .collect();
8052 let columns = authorization
8053 .security
8054 .select_policy_columns(authorization.table, authorization.principal);
8055 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8056 let mut decoded = 0usize;
8057 for column_id in &columns {
8058 if let Some(context) = context {
8059 context.checkpoint()?;
8060 }
8061 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8062 &row_ids, *column_id, snapshot, query_now, context,
8063 )? {
8064 if let Some(row) = rows.get_mut(&row_id) {
8065 row.columns.insert(*column_id, value);
8066 decoded = decoded.saturating_add(1);
8067 }
8068 }
8069 }
8070 if let Some(context) = context {
8071 context.consume(candidates.len().saturating_add(decoded))?;
8072 }
8073 let allowed = rows
8074 .into_values()
8075 .filter_map(|row| {
8076 authorization
8077 .security
8078 .row_allowed(
8079 authorization.table,
8080 crate::security::PolicyCommand::Select,
8081 &row,
8082 authorization.principal,
8083 false,
8084 )
8085 .then_some(row.row_id)
8086 })
8087 .collect();
8088 crate::trace::QueryTrace::record(|trace| {
8089 trace.rls_rows_evaluated = trace.rls_rows_evaluated.saturating_add(candidates.len());
8090 trace.rls_policy_columns_decoded =
8091 trace.rls_policy_columns_decoded.saturating_add(decoded);
8092 trace.authorization_nanos = trace
8093 .authorization_nanos
8094 .saturating_add(started.elapsed().as_nanos() as u64);
8095 });
8096 Ok(allowed)
8097 }
8098
8099 pub fn search(
8101 &mut self,
8102 request: &crate::query::SearchRequest,
8103 ) -> Result<Vec<crate::query::SearchHit>> {
8104 self.search_with_allowed(request, None)
8105 }
8106
8107 pub fn search_at(
8108 &mut self,
8109 request: &crate::query::SearchRequest,
8110 snapshot: Snapshot,
8111 authorized: Option<&std::collections::HashSet<RowId>>,
8112 ) -> Result<Vec<crate::query::SearchHit>> {
8113 self.search_at_with_allowed(request, snapshot, authorized)
8114 }
8115
8116 pub fn search_with_allowed(
8117 &mut self,
8118 request: &crate::query::SearchRequest,
8119 authorized: Option<&std::collections::HashSet<RowId>>,
8120 ) -> Result<Vec<crate::query::SearchHit>> {
8121 self.search_at_with_allowed(request, self.snapshot(), authorized)
8122 }
8123
8124 pub fn search_at_with_allowed(
8125 &mut self,
8126 request: &crate::query::SearchRequest,
8127 snapshot: Snapshot,
8128 authorized: Option<&std::collections::HashSet<RowId>>,
8129 ) -> Result<Vec<crate::query::SearchHit>> {
8130 self.search_at_with_allowed_and_context(request, snapshot, authorized, None)
8131 }
8132
8133 pub fn search_at_with_allowed_and_context(
8134 &mut self,
8135 request: &crate::query::SearchRequest,
8136 snapshot: Snapshot,
8137 authorized: Option<&std::collections::HashSet<RowId>>,
8138 context: Option<&crate::query::AiExecutionContext>,
8139 ) -> Result<Vec<crate::query::SearchHit>> {
8140 self.ensure_indexes_complete()?;
8141 self.search_at_with_filters_and_context(request, snapshot, authorized, None, context, None)
8142 }
8143
8144 pub fn search_at_with_candidate_authorization_and_context(
8145 &mut self,
8146 request: &crate::query::SearchRequest,
8147 snapshot: Snapshot,
8148 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8149 context: Option<&crate::query::AiExecutionContext>,
8150 ) -> Result<Vec<crate::query::SearchHit>> {
8151 self.ensure_indexes_complete()?;
8152 self.search_at_with_filters_and_context(
8153 request,
8154 snapshot,
8155 None,
8156 authorization,
8157 context,
8158 None,
8159 )
8160 }
8161
8162 #[doc(hidden)]
8163 pub fn search_at_with_candidate_authorization_on_generation(
8164 &self,
8165 request: &crate::query::SearchRequest,
8166 snapshot: Snapshot,
8167 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8168 context: Option<&crate::query::AiExecutionContext>,
8169 ) -> Result<Vec<crate::query::SearchHit>> {
8170 self.search_at_with_filters_and_context(
8171 request,
8172 snapshot,
8173 None,
8174 authorization,
8175 context,
8176 None,
8177 )
8178 }
8179
8180 #[doc(hidden)]
8181 pub fn search_at_with_candidate_authorization_on_generation_after(
8182 &self,
8183 request: &crate::query::SearchRequest,
8184 snapshot: Snapshot,
8185 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8186 context: Option<&crate::query::AiExecutionContext>,
8187 after: Option<crate::query::SearchAfter>,
8188 ) -> Result<Vec<crate::query::SearchHit>> {
8189 self.search_at_with_filters_and_context(
8190 request,
8191 snapshot,
8192 None,
8193 authorization,
8194 context,
8195 after,
8196 )
8197 }
8198
8199 fn search_at_with_filters_and_context(
8200 &self,
8201 request: &crate::query::SearchRequest,
8202 snapshot: Snapshot,
8203 authorized: Option<&std::collections::HashSet<RowId>>,
8204 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8205 context: Option<&crate::query::AiExecutionContext>,
8206 after: Option<crate::query::SearchAfter>,
8207 ) -> Result<Vec<crate::query::SearchHit>> {
8208 use crate::query::{
8209 ComponentScore, Condition, Fusion, SearchHit, MAX_FINAL_LIMIT, MAX_HARD_CONDITIONS,
8210 MAX_PROJECTION_COLUMNS, MAX_RETRIEVERS, MAX_RETRIEVER_WEIGHT,
8211 };
8212 let total_started = std::time::Instant::now();
8213 let rank_offset = after.map_or(0, |after| after.returned_count);
8214 self.require_select()?;
8215 if request.limit == 0 {
8216 return Err(MongrelError::InvalidArgument(
8217 "search limit must be > 0".into(),
8218 ));
8219 }
8220 if request.limit > MAX_FINAL_LIMIT {
8221 return Err(MongrelError::InvalidArgument(format!(
8222 "search limit exceeds {MAX_FINAL_LIMIT}"
8223 )));
8224 }
8225 if after.is_some_and(|cursor| !cursor.final_score.is_finite()) {
8226 return Err(MongrelError::InvalidArgument(
8227 "search-after score must be finite".into(),
8228 ));
8229 }
8230 if request.retrievers.is_empty() {
8231 return Err(MongrelError::InvalidArgument(
8232 "search requires at least one retriever".into(),
8233 ));
8234 }
8235 if request.retrievers.len() > MAX_RETRIEVERS {
8236 return Err(MongrelError::InvalidArgument(format!(
8237 "search exceeds {MAX_RETRIEVERS} retrievers"
8238 )));
8239 }
8240 if request.must.len() > MAX_HARD_CONDITIONS {
8241 return Err(MongrelError::InvalidArgument(format!(
8242 "search exceeds {MAX_HARD_CONDITIONS} hard conditions"
8243 )));
8244 }
8245 for condition in &request.must {
8246 self.validate_condition(condition)?;
8247 }
8248 if request.must.iter().any(|condition| {
8249 matches!(
8250 condition,
8251 Condition::Ann { .. }
8252 | Condition::SparseMatch { .. }
8253 | Condition::MinHashSimilar { .. }
8254 )
8255 }) {
8256 return Err(MongrelError::InvalidArgument(
8257 "ranked ANN, Sparse, and MinHash conditions must be retrievers, not must filters"
8258 .into(),
8259 ));
8260 }
8261 let mut names = std::collections::HashSet::new();
8262 for named in &request.retrievers {
8263 if named.name.is_empty()
8264 || named.name.len() > crate::query::MAX_RETRIEVER_NAME_BYTES
8265 || !names.insert(named.name.as_str())
8266 {
8267 return Err(MongrelError::InvalidArgument(format!(
8268 "retriever names must be non-empty, unique, and at most {} UTF-8 bytes",
8269 crate::query::MAX_RETRIEVER_NAME_BYTES
8270 )));
8271 }
8272 if !named.weight.is_finite()
8273 || named.weight < 0.0
8274 || named.weight > MAX_RETRIEVER_WEIGHT
8275 {
8276 return Err(MongrelError::InvalidArgument(format!(
8277 "retriever weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8278 )));
8279 }
8280 self.validate_retriever(&named.retriever)?;
8281 }
8282 let projection = request
8283 .projection
8284 .clone()
8285 .unwrap_or_else(|| self.schema.columns.iter().map(|column| column.id).collect());
8286 if projection.len() > MAX_PROJECTION_COLUMNS {
8287 return Err(MongrelError::InvalidArgument(format!(
8288 "projection exceeds {MAX_PROJECTION_COLUMNS} columns"
8289 )));
8290 }
8291 for column_id in &projection {
8292 if !self
8293 .schema
8294 .columns
8295 .iter()
8296 .any(|column| column.id == *column_id)
8297 {
8298 return Err(MongrelError::ColumnNotFound(column_id.to_string()));
8299 }
8300 }
8301 if let Some(crate::query::Rerank::ExactVector {
8302 embedding_column,
8303 query,
8304 candidate_limit,
8305 weight,
8306 ..
8307 }) = &request.rerank
8308 {
8309 if *candidate_limit < request.limit || *candidate_limit > crate::query::MAX_RETRIEVER_K
8310 {
8311 return Err(MongrelError::InvalidArgument(format!(
8312 "rerank candidate_limit must be between search limit and {}",
8313 crate::query::MAX_RETRIEVER_K
8314 )));
8315 }
8316 if !weight.is_finite() || *weight < 0.0 || *weight > MAX_RETRIEVER_WEIGHT {
8317 return Err(MongrelError::InvalidArgument(format!(
8318 "rerank weight must be finite, non-negative, and <= {MAX_RETRIEVER_WEIGHT}"
8319 )));
8320 }
8321 let column = self
8322 .schema
8323 .columns
8324 .iter()
8325 .find(|column| column.id == *embedding_column)
8326 .ok_or_else(|| MongrelError::ColumnNotFound(embedding_column.to_string()))?;
8327 let crate::schema::TypeId::Embedding { dim } = column.ty else {
8328 return Err(MongrelError::InvalidArgument(format!(
8329 "rerank column {embedding_column} is not an embedding"
8330 )));
8331 };
8332 if query.len() != dim as usize || query.iter().any(|value| !value.is_finite()) {
8333 return Err(MongrelError::InvalidArgument(format!(
8334 "rerank query must contain {dim} finite values"
8335 )));
8336 }
8337 }
8338
8339 let hard_filter_started = std::time::Instant::now();
8340 let hard_filter = if request.must.is_empty() {
8341 None
8342 } else {
8343 let mut sets = Vec::with_capacity(request.must.len());
8344 for condition in &request.must {
8345 if let Some(context) = context {
8346 context.checkpoint()?;
8347 }
8348 sets.push(self.resolve_condition(condition, snapshot)?);
8349 }
8350 Some(RowIdSet::intersect_many(sets))
8351 };
8352 crate::trace::QueryTrace::record(|trace| {
8353 trace.hard_filter_nanos = trace
8354 .hard_filter_nanos
8355 .saturating_add(hard_filter_started.elapsed().as_nanos() as u64);
8356 });
8357 if hard_filter.as_ref().is_some_and(RowIdSet::is_empty) {
8358 return Ok(Vec::new());
8359 }
8360
8361 let constant = match request.fusion {
8362 Fusion::ReciprocalRank { constant } => constant,
8363 };
8364 let mut retrievers: Vec<_> = request.retrievers.iter().collect();
8365 retrievers.sort_by(|a, b| a.name.cmp(&b.name));
8366 let mut fusion_nanos = 0u64;
8367 let mut fused: std::collections::HashMap<RowId, (f64, Vec<ComponentScore>)> =
8368 std::collections::HashMap::new();
8369 for named in retrievers {
8370 if named.weight == 0.0 {
8371 continue;
8372 }
8373 if let Some(context) = context {
8374 context.checkpoint()?;
8375 }
8376 let hits = self.retrieve_filtered(
8377 &named.retriever,
8378 snapshot,
8379 hard_filter.as_ref(),
8380 authorized,
8381 candidate_authorization,
8382 context,
8383 )?;
8384 let retriever_name: std::sync::Arc<str> = named.name.as_str().into();
8385 let fusion_started = std::time::Instant::now();
8386 for hit in hits {
8387 if let Some(context) = context {
8388 context.consume(1)?;
8389 }
8390 let contribution = named.weight / (constant as f64 + hit.rank as f64);
8391 if !contribution.is_finite() {
8392 return Err(MongrelError::InvalidArgument(
8393 "retriever contribution must be finite".into(),
8394 ));
8395 }
8396 let max_fused_candidates = context.map_or(
8397 crate::query::MAX_FUSED_CANDIDATES,
8398 crate::query::AiExecutionContext::max_fused_candidates,
8399 );
8400 if !fused.contains_key(&hit.row_id) && fused.len() >= max_fused_candidates {
8401 return Err(MongrelError::WorkBudgetExceeded);
8402 }
8403 let entry = fused.entry(hit.row_id).or_default();
8404 entry.0 += contribution;
8405 if !entry.0.is_finite() {
8406 return Err(MongrelError::InvalidArgument(
8407 "fused score must be finite".into(),
8408 ));
8409 }
8410 entry.1.push(ComponentScore {
8411 retriever_name: retriever_name.clone(),
8412 rank: hit.rank,
8413 raw_score: hit.score,
8414 contribution,
8415 });
8416 }
8417 fusion_nanos = fusion_nanos.saturating_add(fusion_started.elapsed().as_nanos() as u64);
8418 }
8419 let union_size = fused.len();
8420 let mut ranked: Vec<_> = fused
8421 .into_iter()
8422 .map(|(row_id, (fused_score, components))| {
8423 (row_id, fused_score, components, None, fused_score)
8424 })
8425 .collect();
8426 let order = |(a_row, _, _, _, a_score): &(
8427 RowId,
8428 f64,
8429 Vec<ComponentScore>,
8430 Option<f32>,
8431 f64,
8432 ),
8433 (b_row, _, _, _, b_score): &(
8434 RowId,
8435 f64,
8436 Vec<ComponentScore>,
8437 Option<f32>,
8438 f64,
8439 )| { b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row)) };
8440 if let Some(crate::query::Rerank::ExactVector {
8441 embedding_column,
8442 query,
8443 metric,
8444 candidate_limit,
8445 weight,
8446 }) = &request.rerank
8447 {
8448 let fused_order = |(a_row, a_score, ..): &(
8449 RowId,
8450 f64,
8451 Vec<ComponentScore>,
8452 Option<f32>,
8453 f64,
8454 ),
8455 (b_row, b_score, ..): &(
8456 RowId,
8457 f64,
8458 Vec<ComponentScore>,
8459 Option<f32>,
8460 f64,
8461 )| {
8462 b_score.total_cmp(a_score).then_with(|| a_row.cmp(b_row))
8463 };
8464 let selection_started = std::time::Instant::now();
8465 if let Some(context) = context {
8466 context.consume(ranked.len())?;
8467 }
8468 if ranked.len() > *candidate_limit {
8469 let (_, _, _) = ranked.select_nth_unstable_by(*candidate_limit, fused_order);
8470 ranked.truncate(*candidate_limit);
8471 }
8472 ranked.sort_by(fused_order);
8473 fusion_nanos =
8474 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
8475 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
8476 if let Some(context) = context {
8477 context.consume(row_ids.len())?;
8478 }
8479 let query_now =
8480 context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8481 let gather_started = std::time::Instant::now();
8482 let vectors = self.values_for_rids_batch_at_with_context(
8483 &row_ids,
8484 *embedding_column,
8485 snapshot,
8486 query_now,
8487 context,
8488 )?;
8489 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
8490 let vector_work =
8491 crate::query::work_units(query.len(), crate::query::VECTOR_WORK_QUANTUM);
8492 let query_norm = if matches!(metric, crate::query::VectorMetric::Cosine) {
8493 if let Some(context) = context {
8494 context.consume(vector_work)?;
8495 }
8496 query
8497 .iter()
8498 .map(|value| f64::from(*value).powi(2))
8499 .sum::<f64>()
8500 .sqrt()
8501 } else {
8502 0.0
8503 };
8504 let score_started = std::time::Instant::now();
8505 let mut scores = std::collections::HashMap::with_capacity(vectors.len());
8506 for (row_id, value) in vectors {
8507 if let Some(meta) = value.generated_embedding_metadata() {
8508 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
8509 continue;
8510 }
8511 }
8512 let Some(vector) = value.as_embedding() else {
8513 continue;
8514 };
8515 let score = match metric {
8516 crate::query::VectorMetric::DotProduct => {
8517 if let Some(context) = context {
8518 context.consume(vector_work)?;
8519 }
8520 query
8521 .iter()
8522 .zip(vector)
8523 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8524 .sum::<f64>()
8525 }
8526 crate::query::VectorMetric::Cosine => {
8527 if let Some(context) = context {
8528 context.consume(vector_work.saturating_mul(2))?;
8529 }
8530 let dot = query
8531 .iter()
8532 .zip(vector)
8533 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8534 .sum::<f64>();
8535 let norm = vector
8536 .iter()
8537 .map(|value| f64::from(*value).powi(2))
8538 .sum::<f64>()
8539 .sqrt();
8540 if query_norm == 0.0 || norm == 0.0 {
8541 0.0
8542 } else {
8543 dot / (query_norm * norm)
8544 }
8545 }
8546 crate::query::VectorMetric::Euclidean => {
8547 if let Some(context) = context {
8548 context.consume(vector_work)?;
8549 }
8550 query
8551 .iter()
8552 .zip(vector)
8553 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
8554 .sum::<f64>()
8555 .sqrt()
8556 }
8557 };
8558 if !score.is_finite() {
8559 return Err(MongrelError::InvalidArgument(
8560 "exact rerank score must be finite".into(),
8561 ));
8562 }
8563 scores.insert(row_id, score as f32);
8564 }
8565 let mut reranked = Vec::with_capacity(ranked.len());
8566 for (row_id, fused_score, components, _, _) in ranked.drain(..) {
8567 let Some(score) = scores.get(&row_id).copied() else {
8568 continue;
8569 };
8570 let ordering_score = match metric {
8571 crate::query::VectorMetric::Euclidean => -f64::from(score),
8572 crate::query::VectorMetric::Cosine | crate::query::VectorMetric::DotProduct => {
8573 f64::from(score)
8574 }
8575 };
8576 let final_score = fused_score + *weight * ordering_score;
8577 if !final_score.is_finite() {
8578 return Err(MongrelError::InvalidArgument(
8579 "final rerank score must be finite".into(),
8580 ));
8581 }
8582 reranked.push((row_id, fused_score, components, Some(score), final_score));
8583 }
8584 ranked = reranked;
8585 ranked.sort_by(order);
8586 crate::trace::QueryTrace::record(|trace| {
8587 trace.exact_vector_gather_nanos =
8588 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
8589 trace.exact_vector_score_nanos = trace
8590 .exact_vector_score_nanos
8591 .saturating_add(score_started.elapsed().as_nanos() as u64);
8592 });
8593 }
8594 if let Some(after) = after {
8595 ranked.retain(|(row_id, _, _, _, final_score)| {
8596 final_score.total_cmp(&after.final_score).is_lt()
8597 || (final_score.total_cmp(&after.final_score).is_eq() && *row_id > after.row_id)
8598 });
8599 }
8600 let projection_started = std::time::Instant::now();
8601 let sentinel = projection
8602 .first()
8603 .copied()
8604 .or_else(|| self.schema.columns.first().map(|column| column.id));
8605 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8606 let mut out = Vec::with_capacity(request.limit.min(ranked.len()));
8607 let mut projection_rows = 0usize;
8608 let mut projection_cells = 0usize;
8609 while out.len() < request.limit && !ranked.is_empty() {
8610 if let Some(context) = context {
8611 context.checkpoint()?;
8612 context.consume(ranked.len())?;
8613 }
8614 let needed = request.limit - out.len();
8615 let window_size = ranked
8616 .len()
8617 .min(needed.saturating_mul(2).max(needed.saturating_add(8)));
8618 let selection_started = std::time::Instant::now();
8619 let mut remainder = if ranked.len() > window_size {
8620 let (_, _, _) = ranked.select_nth_unstable_by(window_size, order);
8621 ranked.split_off(window_size)
8622 } else {
8623 Vec::new()
8624 };
8625 ranked.sort_by(order);
8626 fusion_nanos =
8627 fusion_nanos.saturating_add(selection_started.elapsed().as_nanos() as u64);
8628 let row_ids: Vec<_> = ranked.iter().map(|(row_id, ..)| row_id.0).collect();
8629 let gathered_columns = projection.len().max(usize::from(sentinel.is_some()));
8630 if let Some(context) = context {
8631 context.consume(row_ids.len().saturating_mul(gathered_columns))?;
8632 }
8633 projection_rows = projection_rows.saturating_add(row_ids.len());
8634 projection_cells =
8635 projection_cells.saturating_add(row_ids.len().saturating_mul(gathered_columns));
8636 let mut cells: std::collections::HashMap<RowId, std::collections::HashMap<u16, Value>> =
8637 std::collections::HashMap::new();
8638 if let Some(column_id) = sentinel {
8639 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8640 &row_ids, column_id, snapshot, query_now, context,
8641 )? {
8642 cells.entry(row_id).or_default().insert(column_id, value);
8643 }
8644 }
8645 for &column_id in &projection {
8646 if Some(column_id) == sentinel {
8647 continue;
8648 }
8649 for (row_id, value) in self.values_for_rids_batch_at_with_context(
8650 &row_ids, column_id, snapshot, query_now, context,
8651 )? {
8652 cells.entry(row_id).or_default().insert(column_id, value);
8653 }
8654 }
8655 for (row_id, fused_score, mut components, exact_rerank_score, final_score) in
8656 ranked.drain(..)
8657 {
8658 let Some(row_cells) = cells.remove(&row_id) else {
8659 continue;
8660 };
8661 components.sort_by(|a, b| a.retriever_name.cmp(&b.retriever_name));
8662 let final_rank = rank_offset.saturating_add(out.len()).saturating_add(1);
8663 out.push(SearchHit {
8664 row_id,
8665 cells: projection
8666 .iter()
8667 .filter_map(|column_id| {
8668 row_cells
8669 .get(column_id)
8670 .cloned()
8671 .map(|value| (*column_id, value))
8672 })
8673 .collect(),
8674 components,
8675 fused_score,
8676 exact_rerank_score,
8677 final_score,
8678 final_rank,
8679 });
8680 if out.len() == request.limit {
8681 break;
8682 }
8683 }
8684 ranked.append(&mut remainder);
8685 }
8686 crate::trace::QueryTrace::record(|trace| {
8687 trace.union_size = union_size;
8688 trace.fusion_nanos = trace.fusion_nanos.saturating_add(fusion_nanos);
8689 trace.projection_nanos = trace
8690 .projection_nanos
8691 .saturating_add(projection_started.elapsed().as_nanos() as u64);
8692 trace.total_nanos = trace
8693 .total_nanos
8694 .saturating_add(total_started.elapsed().as_nanos() as u64);
8695 trace.projection_rows = trace.projection_rows.saturating_add(projection_rows);
8696 trace.projection_cells = trace.projection_cells.saturating_add(projection_cells);
8697 if let Some(context) = context {
8698 trace.work_consumed = trace.work_consumed.saturating_add(context.consumed_work());
8699 }
8700 });
8701 Ok(out)
8702 }
8703
8704 pub fn set_similarity(
8707 &mut self,
8708 request: &crate::query::SetSimilarityRequest,
8709 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8710 self.set_similarity_with_allowed(request, None)
8711 }
8712
8713 pub fn set_similarity_at(
8714 &mut self,
8715 request: &crate::query::SetSimilarityRequest,
8716 snapshot: Snapshot,
8717 allowed: Option<&std::collections::HashSet<RowId>>,
8718 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8719 self.set_similarity_explained_at(request, snapshot, allowed)
8720 .map(|(hits, _)| hits)
8721 }
8722
8723 pub fn ann_rerank(
8725 &mut self,
8726 request: &crate::query::AnnRerankRequest,
8727 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8728 self.ann_rerank_with_allowed(request, None)
8729 }
8730
8731 pub fn ann_rerank_with_allowed(
8732 &mut self,
8733 request: &crate::query::AnnRerankRequest,
8734 allowed: Option<&std::collections::HashSet<RowId>>,
8735 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8736 self.ann_rerank_at(request, self.snapshot(), allowed)
8737 }
8738
8739 pub fn ann_rerank_at(
8740 &mut self,
8741 request: &crate::query::AnnRerankRequest,
8742 snapshot: Snapshot,
8743 allowed: Option<&std::collections::HashSet<RowId>>,
8744 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8745 self.ann_rerank_at_with_context(request, snapshot, allowed, None)
8746 }
8747
8748 pub fn ann_rerank_at_with_context(
8749 &mut self,
8750 request: &crate::query::AnnRerankRequest,
8751 snapshot: Snapshot,
8752 allowed: Option<&std::collections::HashSet<RowId>>,
8753 context: Option<&crate::query::AiExecutionContext>,
8754 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8755 self.ensure_indexes_complete()?;
8756 self.ann_rerank_at_with_filters_and_context(request, snapshot, allowed, None, context)
8757 }
8758
8759 pub fn ann_rerank_at_with_candidate_authorization_and_context(
8760 &mut self,
8761 request: &crate::query::AnnRerankRequest,
8762 snapshot: Snapshot,
8763 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8764 context: Option<&crate::query::AiExecutionContext>,
8765 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8766 self.ensure_indexes_complete()?;
8767 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
8768 }
8769
8770 #[doc(hidden)]
8771 pub fn ann_rerank_at_with_candidate_authorization_on_generation(
8772 &self,
8773 request: &crate::query::AnnRerankRequest,
8774 snapshot: Snapshot,
8775 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8776 context: Option<&crate::query::AiExecutionContext>,
8777 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8778 self.ann_rerank_at_with_filters_and_context(request, snapshot, None, authorization, context)
8779 }
8780
8781 fn ann_rerank_at_with_filters_and_context(
8782 &self,
8783 request: &crate::query::AnnRerankRequest,
8784 snapshot: Snapshot,
8785 allowed: Option<&std::collections::HashSet<RowId>>,
8786 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
8787 context: Option<&crate::query::AiExecutionContext>,
8788 ) -> Result<Vec<crate::query::AnnRerankHit>> {
8789 use crate::query::{
8790 AnnCandidateDistance, AnnRerankHit, Retriever, RetrieverScore, VectorMetric,
8791 MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
8792 };
8793 if request.candidate_k == 0 || request.limit == 0 {
8794 return Err(MongrelError::InvalidArgument(
8795 "candidate_k and limit must be > 0".into(),
8796 ));
8797 }
8798 if request.candidate_k > MAX_RETRIEVER_K || request.limit > MAX_FINAL_LIMIT {
8799 return Err(MongrelError::InvalidArgument(format!(
8800 "candidate_k must be <= {MAX_RETRIEVER_K} and limit <= {MAX_FINAL_LIMIT}"
8801 )));
8802 }
8803 let retriever = Retriever::Ann {
8804 column_id: request.column_id,
8805 query: request.query.clone(),
8806 k: request.candidate_k,
8807 };
8808 self.require_select()?;
8809 self.validate_retriever(&retriever)?;
8810 let hits = self.retrieve_filtered(
8811 &retriever,
8812 snapshot,
8813 None,
8814 allowed,
8815 candidate_authorization,
8816 context,
8817 )?;
8818 let distances: std::collections::HashMap<_, _> = hits
8819 .iter()
8820 .filter_map(|hit| match hit.score {
8821 RetrieverScore::AnnHammingDistance(distance) => {
8822 Some((hit.row_id, AnnCandidateDistance::Hamming(distance)))
8823 }
8824 RetrieverScore::AnnCosineDistance(distance) => {
8825 Some((hit.row_id, AnnCandidateDistance::Cosine(distance)))
8826 }
8827 _ => None,
8828 })
8829 .collect();
8830 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
8831 if let Some(context) = context {
8832 context.consume(row_ids.len())?;
8833 }
8834 let gather_started = std::time::Instant::now();
8835 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
8836 let values = self.values_for_rids_batch_at_with_context(
8837 &row_ids,
8838 request.column_id,
8839 snapshot,
8840 query_now,
8841 context,
8842 )?;
8843 let gather_nanos = gather_started.elapsed().as_nanos() as u64;
8844 let score_started = std::time::Instant::now();
8845 let vector_work =
8846 crate::query::work_units(request.query.len(), crate::query::VECTOR_WORK_QUANTUM);
8847 let query_norm = if matches!(request.metric, VectorMetric::Cosine) {
8848 if let Some(context) = context {
8849 context.consume(vector_work)?;
8850 }
8851 request
8852 .query
8853 .iter()
8854 .map(|value| f64::from(*value).powi(2))
8855 .sum::<f64>()
8856 .sqrt()
8857 } else {
8858 0.0
8859 };
8860 let mut reranked = Vec::with_capacity(values.len().min(request.limit));
8861 for (row_id, value) in values {
8862 if let Some(meta) = value.generated_embedding_metadata() {
8863 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
8864 continue;
8865 }
8866 }
8867 let Some(vector) = value.as_embedding() else {
8868 continue;
8869 };
8870 let exact_score = match request.metric {
8871 VectorMetric::DotProduct => {
8872 if let Some(context) = context {
8873 context.consume(vector_work)?;
8874 }
8875 request
8876 .query
8877 .iter()
8878 .zip(vector)
8879 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8880 .sum::<f64>()
8881 }
8882 VectorMetric::Cosine => {
8883 if let Some(context) = context {
8884 context.consume(vector_work.saturating_mul(2))?;
8885 }
8886 let dot = request
8887 .query
8888 .iter()
8889 .zip(vector)
8890 .map(|(left, right)| f64::from(*left) * f64::from(*right))
8891 .sum::<f64>();
8892 let norm = vector
8893 .iter()
8894 .map(|value| f64::from(*value).powi(2))
8895 .sum::<f64>()
8896 .sqrt();
8897 if query_norm == 0.0 || norm == 0.0 {
8898 0.0
8899 } else {
8900 dot / (query_norm * norm)
8901 }
8902 }
8903 VectorMetric::Euclidean => {
8904 if let Some(context) = context {
8905 context.consume(vector_work)?;
8906 }
8907 request
8908 .query
8909 .iter()
8910 .zip(vector)
8911 .map(|(left, right)| (f64::from(*left) - f64::from(*right)).powi(2))
8912 .sum::<f64>()
8913 .sqrt()
8914 }
8915 };
8916 let exact_score = exact_score as f32;
8917 if !exact_score.is_finite() {
8918 return Err(MongrelError::InvalidArgument(
8919 "exact ANN score must be finite".into(),
8920 ));
8921 }
8922 let Some(candidate_distance) = distances.get(&row_id).copied() else {
8923 continue;
8924 };
8925 reranked.push(AnnRerankHit {
8926 row_id,
8927 candidate_distance,
8928 exact_score,
8929 });
8930 }
8931 reranked.sort_by(|left, right| {
8932 let score = match request.metric {
8933 VectorMetric::Euclidean => left.exact_score.total_cmp(&right.exact_score),
8934 VectorMetric::Cosine | VectorMetric::DotProduct => {
8935 right.exact_score.total_cmp(&left.exact_score)
8936 }
8937 };
8938 score.then_with(|| left.row_id.cmp(&right.row_id))
8939 });
8940 reranked.truncate(request.limit);
8941 crate::trace::QueryTrace::record(|trace| {
8942 trace.exact_vector_gather_nanos =
8943 trace.exact_vector_gather_nanos.saturating_add(gather_nanos);
8944 trace.exact_vector_score_nanos = trace
8945 .exact_vector_score_nanos
8946 .saturating_add(score_started.elapsed().as_nanos() as u64);
8947 });
8948 Ok(reranked)
8949 }
8950
8951 pub fn set_similarity_with_allowed(
8952 &mut self,
8953 request: &crate::query::SetSimilarityRequest,
8954 allowed: Option<&std::collections::HashSet<RowId>>,
8955 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8956 self.set_similarity_explained_at(request, self.snapshot(), allowed)
8957 .map(|(hits, _)| hits)
8958 }
8959
8960 pub fn set_similarity_explained(
8961 &mut self,
8962 request: &crate::query::SetSimilarityRequest,
8963 ) -> Result<(
8964 Vec<crate::query::SetSimilarityHit>,
8965 crate::query::SetSimilarityTrace,
8966 )> {
8967 self.set_similarity_explained_at(request, self.snapshot(), None)
8968 }
8969
8970 fn set_similarity_explained_at(
8971 &mut self,
8972 request: &crate::query::SetSimilarityRequest,
8973 snapshot: Snapshot,
8974 allowed: Option<&std::collections::HashSet<RowId>>,
8975 ) -> Result<(
8976 Vec<crate::query::SetSimilarityHit>,
8977 crate::query::SetSimilarityTrace,
8978 )> {
8979 self.ensure_indexes_complete()?;
8980 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, None)
8981 }
8982
8983 pub fn set_similarity_at_with_context(
8984 &mut self,
8985 request: &crate::query::SetSimilarityRequest,
8986 snapshot: Snapshot,
8987 allowed: Option<&std::collections::HashSet<RowId>>,
8988 context: Option<&crate::query::AiExecutionContext>,
8989 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
8990 self.ensure_indexes_complete()?;
8991 self.set_similarity_explained_at_with_context(request, snapshot, allowed, None, context)
8992 .map(|(hits, _)| hits)
8993 }
8994
8995 pub fn set_similarity_at_with_candidate_authorization_and_context(
8996 &mut self,
8997 request: &crate::query::SetSimilarityRequest,
8998 snapshot: Snapshot,
8999 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9000 context: Option<&crate::query::AiExecutionContext>,
9001 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9002 self.ensure_indexes_complete()?;
9003 self.set_similarity_explained_at_with_context(
9004 request,
9005 snapshot,
9006 None,
9007 authorization,
9008 context,
9009 )
9010 .map(|(hits, _)| hits)
9011 }
9012
9013 #[doc(hidden)]
9014 pub fn set_similarity_at_with_candidate_authorization_on_generation(
9015 &self,
9016 request: &crate::query::SetSimilarityRequest,
9017 snapshot: Snapshot,
9018 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9019 context: Option<&crate::query::AiExecutionContext>,
9020 ) -> Result<Vec<crate::query::SetSimilarityHit>> {
9021 self.set_similarity_explained_at_with_context(
9022 request,
9023 snapshot,
9024 None,
9025 authorization,
9026 context,
9027 )
9028 .map(|(hits, _)| hits)
9029 }
9030
9031 fn set_similarity_explained_at_with_context(
9032 &self,
9033 request: &crate::query::SetSimilarityRequest,
9034 snapshot: Snapshot,
9035 allowed: Option<&std::collections::HashSet<RowId>>,
9036 candidate_authorization: Option<&crate::security::CandidateAuthorization<'_>>,
9037 context: Option<&crate::query::AiExecutionContext>,
9038 ) -> Result<(
9039 Vec<crate::query::SetSimilarityHit>,
9040 crate::query::SetSimilarityTrace,
9041 )> {
9042 use crate::query::{
9043 Retriever, RetrieverScore, SetSimilarityHit, MAX_FINAL_LIMIT, MAX_RETRIEVER_K,
9044 MAX_SET_MEMBERS,
9045 };
9046 let mut trace = crate::query::SetSimilarityTrace::default();
9047 if request.members.is_empty() {
9048 return Ok((Vec::new(), trace));
9049 }
9050 if request.candidate_k == 0 || request.limit == 0 {
9051 return Err(MongrelError::InvalidArgument(
9052 "candidate_k and limit must be > 0".into(),
9053 ));
9054 }
9055 if request.candidate_k > MAX_RETRIEVER_K
9056 || request.limit > MAX_FINAL_LIMIT
9057 || request.members.len() > MAX_SET_MEMBERS
9058 {
9059 return Err(MongrelError::InvalidArgument(format!(
9060 "candidate_k must be <= {MAX_RETRIEVER_K}, limit <= {MAX_FINAL_LIMIT}, and members <= {MAX_SET_MEMBERS}"
9061 )));
9062 }
9063 if !request.min_jaccard.is_finite() || !(0.0..=1.0).contains(&request.min_jaccard) {
9064 return Err(MongrelError::InvalidArgument(
9065 "min_jaccard must be finite and between 0 and 1".into(),
9066 ));
9067 }
9068 let started = std::time::Instant::now();
9069 let retriever = Retriever::MinHash {
9070 column_id: request.column_id,
9071 members: request.members.clone(),
9072 k: request.candidate_k,
9073 };
9074 self.require_select()?;
9075 self.validate_retriever(&retriever)?;
9076 let hits = self.retrieve_filtered(
9077 &retriever,
9078 snapshot,
9079 None,
9080 allowed,
9081 candidate_authorization,
9082 context,
9083 )?;
9084 trace.candidate_generation_us = started.elapsed().as_micros() as u64;
9085 trace.candidate_count = hits.len();
9086 let row_ids: Vec<_> = hits.iter().map(|hit| hit.row_id.0).collect();
9087 if let Some(context) = context {
9088 context.consume(row_ids.len())?;
9089 }
9090 let started = std::time::Instant::now();
9091 let query_now = context.map_or_else(unix_nanos_now, |context| context.query_time_nanos());
9092 let values = self.values_for_rids_batch_at_with_context(
9093 &row_ids,
9094 request.column_id,
9095 snapshot,
9096 query_now,
9097 context,
9098 )?;
9099 trace.gather_us = started.elapsed().as_micros() as u64;
9100 if let Some(context) = context {
9101 context.consume(request.members.len())?;
9102 }
9103 let query: std::collections::HashSet<_> = request.members.iter().cloned().collect();
9104 let estimates: std::collections::HashMap<_, _> = hits
9105 .into_iter()
9106 .filter_map(|hit| match hit.score {
9107 RetrieverScore::MinHashEstimatedJaccard(score) => Some((hit.row_id, score)),
9108 _ => None,
9109 })
9110 .collect();
9111 let started = std::time::Instant::now();
9112 let mut parsed = Vec::with_capacity(values.len());
9113 for (row_id, value) in values {
9114 let Value::Bytes(bytes) = value else {
9115 continue;
9116 };
9117 if let Some(context) = context {
9118 context.consume(crate::query::work_units(
9119 bytes.len(),
9120 crate::query::PARSE_WORK_QUANTUM,
9121 ))?;
9122 }
9123 let Ok(serde_json::Value::Array(members)) = serde_json::from_slice(&bytes) else {
9124 continue;
9125 };
9126 if let Some(context) = context {
9127 context.consume(members.len())?;
9128 }
9129 let stored = members
9130 .into_iter()
9131 .filter_map(|member| match member {
9132 serde_json::Value::String(value) => {
9133 Some(crate::query::SetMember::String(value))
9134 }
9135 serde_json::Value::Number(value) => {
9136 Some(crate::query::SetMember::Number(value))
9137 }
9138 serde_json::Value::Bool(value) => Some(crate::query::SetMember::Boolean(value)),
9139 _ => None,
9140 })
9141 .collect::<std::collections::HashSet<_>>();
9142 parsed.push((row_id, stored));
9143 }
9144 trace.parse_us = started.elapsed().as_micros() as u64;
9145 trace.verified_count = parsed.len();
9146 let started = std::time::Instant::now();
9147 let mut exact = Vec::new();
9148 for (row_id, stored) in parsed {
9149 if let Some(context) = context {
9150 context.consume(query.len().saturating_add(stored.len()))?;
9151 }
9152 let union = query.union(&stored).count();
9153 let score = if union == 0 {
9154 1.0
9155 } else {
9156 query.intersection(&stored).count() as f32 / union as f32
9157 };
9158 if score >= request.min_jaccard {
9159 exact.push(SetSimilarityHit {
9160 row_id,
9161 estimated_jaccard: estimates.get(&row_id).copied().unwrap_or_default(),
9162 exact_jaccard: score,
9163 });
9164 }
9165 }
9166 exact.sort_by(|a, b| {
9167 b.exact_jaccard
9168 .total_cmp(&a.exact_jaccard)
9169 .then_with(|| a.row_id.cmp(&b.row_id))
9170 });
9171 exact.truncate(request.limit);
9172 trace.score_us = started.elapsed().as_micros() as u64;
9173 crate::trace::QueryTrace::record(|query_trace| {
9174 query_trace.exact_set_gather_nanos = query_trace
9175 .exact_set_gather_nanos
9176 .saturating_add(trace.gather_us.saturating_mul(1_000));
9177 query_trace.exact_set_parse_nanos = query_trace
9178 .exact_set_parse_nanos
9179 .saturating_add(trace.parse_us.saturating_mul(1_000));
9180 query_trace.exact_set_score_nanos = query_trace
9181 .exact_set_score_nanos
9182 .saturating_add(trace.score_us.saturating_mul(1_000));
9183 });
9184 Ok((exact, trace))
9185 }
9186
9187 fn values_for_rids_batch_at(
9189 &self,
9190 row_ids: &[u64],
9191 column_id: u16,
9192 snapshot: Snapshot,
9193 now: i64,
9194 ) -> Result<Vec<(RowId, Value)>> {
9195 if self.ttl.is_none()
9196 && self.memtable.is_empty()
9197 && self.mutable_run.is_empty()
9198 && self.run_refs.len() == 1
9199 {
9200 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9201 if row_ids.len().saturating_mul(24) < reader.row_count() {
9206 let mut values = Vec::with_capacity(row_ids.len());
9207 for &raw_row_id in row_ids {
9208 let row_id = RowId(raw_row_id);
9209 if let Some((_, false, Some(value))) =
9210 reader.get_version_column_at(row_id, snapshot, column_id)?
9211 {
9212 values.push((row_id, value));
9213 }
9214 }
9215 return Ok(values);
9216 }
9217 let (positions, visible_row_ids) = reader.visible_positions_with_rids_at(snapshot)?;
9218 let requested: Vec<(RowId, usize)> = row_ids
9219 .iter()
9220 .filter_map(|raw| {
9221 visible_row_ids
9222 .binary_search(&(*raw as i64))
9223 .ok()
9224 .map(|index| (RowId(*raw), positions[index]))
9225 })
9226 .collect();
9227 let values = reader.gather_column(
9228 column_id,
9229 &requested
9230 .iter()
9231 .map(|(_, position)| *position)
9232 .collect::<Vec<_>>(),
9233 )?;
9234 return Ok(requested
9235 .into_iter()
9236 .zip(values)
9237 .map(|((row_id, _), value)| (row_id, value))
9238 .collect());
9239 }
9240 self.values_for_rids_at(row_ids, column_id, snapshot, now)
9241 }
9242
9243 fn values_for_rids_batch_at_with_context(
9244 &self,
9245 row_ids: &[u64],
9246 column_id: u16,
9247 snapshot: Snapshot,
9248 now: i64,
9249 context: Option<&crate::query::AiExecutionContext>,
9250 ) -> Result<Vec<(RowId, Value)>> {
9251 let Some(context) = context else {
9252 return self.values_for_rids_batch_at(row_ids, column_id, snapshot, now);
9253 };
9254 let mut values = Vec::with_capacity(row_ids.len());
9255 for chunk in row_ids.chunks(256) {
9256 context.checkpoint()?;
9257 values.extend(self.values_for_rids_batch_at(chunk, column_id, snapshot, now)?);
9258 }
9259 Ok(values)
9260 }
9261
9262 fn values_for_rids_at(
9264 &self,
9265 row_ids: &[u64],
9266 column_id: u16,
9267 snapshot: Snapshot,
9268 now: i64,
9269 ) -> Result<Vec<(RowId, Value)>> {
9270 let mut readers: Vec<_> = self
9271 .run_refs
9272 .iter()
9273 .map(|run| self.open_reader(run.run_id))
9274 .collect::<Result<_>>()?;
9275 let mut out = Vec::with_capacity(row_ids.len());
9276 for &raw_row_id in row_ids {
9277 let row_id = RowId(raw_row_id);
9278 let mem = self.memtable.get_version_at(row_id, snapshot);
9279 let mutable = self.mutable_run.get_version_at(row_id, snapshot);
9280 let overlay = match (mem, mutable) {
9281 (Some((_, a)), Some((_, b))) => Some(
9282 if Snapshot::version_is_newer(
9283 a.committed_epoch,
9284 a.commit_ts,
9285 b.committed_epoch,
9286 b.commit_ts,
9287 ) {
9288 a
9289 } else {
9290 b
9291 },
9292 ),
9293 (Some((_, value)), None) | (None, Some((_, value))) => Some(value),
9294 (None, None) => None,
9295 };
9296 if let Some(row) = overlay {
9297 if !row.deleted && !self.row_expired_at(&row, now) {
9298 if let Some(value) = row.columns.get(&column_id) {
9299 out.push((row_id, value.clone()));
9300 }
9301 }
9302 continue;
9303 }
9304
9305 let mut best: Option<(Epoch, bool, Option<Value>, usize)> = None;
9306 for (index, reader) in readers.iter_mut().enumerate() {
9307 if let Some((epoch, deleted, value)) =
9308 reader.get_version_column_at(row_id, snapshot, column_id)?
9309 {
9310 if best
9311 .as_ref()
9312 .map(|(best_epoch, ..)| epoch > *best_epoch)
9313 .unwrap_or(true)
9314 {
9315 best = Some((epoch, deleted, value, index));
9316 }
9317 }
9318 }
9319 let Some((_, false, Some(value), reader_index)) = best else {
9320 continue;
9321 };
9322 if let Some(ttl) = self.ttl {
9323 if ttl.column_id != column_id {
9324 if let Some((_, _, Some(Value::Int64(timestamp)))) = readers[reader_index]
9325 .get_version_column_at(row_id, snapshot, ttl.column_id)?
9326 {
9327 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9328 continue;
9329 }
9330 }
9331 } else if let Value::Int64(timestamp) = value {
9332 if timestamp.saturating_add(ttl.duration_nanos as i64) <= now {
9333 continue;
9334 }
9335 }
9336 }
9337 out.push((row_id, value));
9338 }
9339 Ok(out)
9340 }
9341
9342 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
9347 self.rows_for_rids_at_time(rids, snapshot, unix_nanos_now(), None)
9348 }
9349
9350 pub fn rows_for_rids_with_context(
9351 &self,
9352 rids: &[u64],
9353 snapshot: Snapshot,
9354 context: &crate::query::AiExecutionContext,
9355 ) -> Result<Vec<Row>> {
9356 context.consume(rids.len().saturating_mul(self.schema.columns.len()))?;
9357 self.rows_for_rids_at_time(rids, snapshot, context.query_time_nanos(), None)
9358 }
9359
9360 fn rows_for_rids_at_time(
9361 &self,
9362 rids: &[u64],
9363 snapshot: Snapshot,
9364 ttl_now: i64,
9365 control: Option<&crate::ExecutionControl>,
9366 ) -> Result<Vec<Row>> {
9367 use std::collections::HashMap;
9368 let mut rows = Vec::with_capacity(rids.len());
9369 let tier_size = self.memtable.len() + self.mutable_run.len();
9384 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
9385 if rids.len().saturating_mul(24) < tier_size {
9386 for &rid in rids {
9387 if overlay.len() & 255 == 0 {
9388 control
9389 .map(crate::ExecutionControl::checkpoint)
9390 .transpose()?;
9391 }
9392 let mem = self.memtable.get_version_at(RowId(rid), snapshot);
9393 let mrun = self.mutable_run.get_version_at(RowId(rid), snapshot);
9394 let newest = match (mem, mrun) {
9395 (Some((_, mr)), Some((_, rr))) => Some(
9396 if Snapshot::version_is_newer(
9397 mr.committed_epoch,
9398 mr.commit_ts,
9399 rr.committed_epoch,
9400 rr.commit_ts,
9401 ) {
9402 mr
9403 } else {
9404 rr
9405 },
9406 ),
9407 (Some((_, mr)), None) => Some(mr),
9408 (None, Some((_, rr))) => Some(rr),
9409 (None, None) => None,
9410 };
9411 if let Some(row) = newest {
9412 overlay.insert(rid, row);
9413 }
9414 }
9415 } else {
9416 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
9417 overlay
9418 .entry(row.row_id.0)
9419 .and_modify(|e| {
9420 if Snapshot::version_is_newer(
9421 row.committed_epoch,
9422 row.commit_ts,
9423 e.committed_epoch,
9424 e.commit_ts,
9425 ) {
9426 *e = row.clone();
9427 }
9428 })
9429 .or_insert(row);
9430 };
9431 for (index, row) in self
9432 .memtable
9433 .visible_versions_at(snapshot)
9434 .into_iter()
9435 .enumerate()
9436 {
9437 if index & 255 == 0 {
9438 control
9439 .map(crate::ExecutionControl::checkpoint)
9440 .transpose()?;
9441 }
9442 fold_newest(row, &mut overlay);
9443 }
9444 for (index, row) in self
9445 .mutable_run
9446 .visible_versions_at(snapshot)
9447 .into_iter()
9448 .enumerate()
9449 {
9450 if index & 255 == 0 {
9451 control
9452 .map(crate::ExecutionControl::checkpoint)
9453 .transpose()?;
9454 }
9455 fold_newest(row, &mut overlay);
9456 }
9457 }
9458 if self.run_refs.len() == 1 {
9459 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
9460 if rids.len().saturating_mul(24) < reader.row_count() {
9468 for (index, &rid) in rids.iter().enumerate() {
9469 if index & 255 == 0 {
9470 control
9471 .map(crate::ExecutionControl::checkpoint)
9472 .transpose()?;
9473 }
9474 if let Some(r) = overlay.get(&rid) {
9475 if !r.deleted {
9476 rows.push(r.clone());
9477 }
9478 continue;
9479 }
9480 if let Some((_, row)) = reader.get_version_at(RowId(rid), snapshot)? {
9481 if !row.deleted {
9482 rows.push(row);
9483 }
9484 }
9485 }
9486 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9487 return Ok(rows);
9488 }
9489 let (positions, vis_rids) = reader.visible_positions_with_rids_at(snapshot)?;
9498 enum Src {
9501 Overlay,
9502 Run,
9503 }
9504 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
9505 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
9506 for (index, rid) in rids.iter().enumerate() {
9507 if index & 255 == 0 {
9508 control
9509 .map(crate::ExecutionControl::checkpoint)
9510 .transpose()?;
9511 }
9512 if overlay.contains_key(rid) {
9513 plan.push(Src::Overlay);
9514 continue;
9515 }
9516 match vis_rids.binary_search(&(*rid as i64)) {
9517 Ok(i) => {
9518 plan.push(Src::Run);
9519 fetch.push(positions[i]);
9520 }
9521 Err(_) => { }
9522 }
9523 }
9524 let fetched = reader.materialize_batch(&fetch)?;
9525 let mut fetched_iter = fetched.into_iter();
9526 for (index, (rid, src)) in rids.iter().zip(plan).enumerate() {
9527 if index & 255 == 0 {
9528 control
9529 .map(crate::ExecutionControl::checkpoint)
9530 .transpose()?;
9531 }
9532 match src {
9533 Src::Overlay => {
9534 if let Some(r) = overlay.get(rid) {
9535 if !r.deleted {
9536 rows.push(r.clone());
9537 }
9538 }
9539 }
9540 Src::Run => {
9541 if let Some(row) = fetched_iter.next() {
9542 if !row.deleted {
9543 rows.push(row);
9544 }
9545 }
9546 }
9547 }
9548 }
9549 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9550 return Ok(rows);
9551 }
9552 let mut readers: Vec<_> = self
9556 .run_refs
9557 .iter()
9558 .map(|rr| self.open_reader(rr.run_id))
9559 .collect::<Result<Vec<_>>>()?;
9560 for (index, rid) in rids.iter().enumerate() {
9561 if index & 255 == 0 {
9562 control
9563 .map(crate::ExecutionControl::checkpoint)
9564 .transpose()?;
9565 }
9566 if let Some(r) = overlay.get(rid) {
9567 if !r.deleted {
9568 rows.push(r.clone());
9569 }
9570 continue;
9571 }
9572 let mut best: Option<(Epoch, Row)> = None;
9573 for reader in readers.iter_mut() {
9574 if let Ok(Some((epoch, row))) = reader.get_version_at(RowId(*rid), snapshot) {
9575 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
9576 best = Some((epoch, row));
9577 }
9578 }
9579 }
9580 if let Some((_, r)) = best {
9581 if !r.deleted {
9582 rows.push(r);
9583 }
9584 }
9585 }
9586 rows.retain(|row| !self.row_expired_at(row, ttl_now));
9587 Ok(rows)
9588 }
9589
9590 pub fn indexes_complete(&self) -> bool {
9600 self.indexes_complete
9601 }
9602
9603 pub fn index_build_policy(&self) -> IndexBuildPolicy {
9605 self.index_build_policy
9606 }
9607
9608 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
9612 self.index_build_policy = policy;
9613 }
9614
9615 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
9620 if !self.indexes_complete {
9624 return None;
9625 }
9626 let b = self.bitmap.get(&column_id)?;
9627 let result: Vec<Vec<u8>> = b
9628 .keys()
9629 .into_iter()
9630 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
9631 .collect();
9632 Some(result)
9633 }
9634
9635 pub fn fk_join_row_ids(
9636 &self,
9637 fk_column_id: u16,
9638 pk_values: &[Vec<u8>],
9639 fk_conditions: &[crate::query::Condition],
9640 snapshot: Snapshot,
9641 ) -> Result<Vec<u64>> {
9642 let Some(b) = self.bitmap.get(&fk_column_id) else {
9643 return Ok(Vec::new());
9644 };
9645 let mut join_set = {
9646 let mut acc = roaring::RoaringBitmap::new();
9647 for v in pk_values {
9648 acc |= b.get(v);
9649 }
9650 RowIdSet::from_roaring(acc)
9651 };
9652 if !fk_conditions.is_empty() {
9653 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
9654 sets.push(join_set);
9655 for c in fk_conditions {
9656 sets.push(self.resolve_condition(c, snapshot)?);
9657 }
9658 join_set = RowIdSet::intersect_many(sets);
9659 }
9660 Ok(join_set.into_sorted_vec())
9661 }
9662
9663 pub fn fk_join_count(
9669 &self,
9670 fk_column_id: u16,
9671 pk_values: &[Vec<u8>],
9672 fk_conditions: &[crate::query::Condition],
9673 snapshot: Snapshot,
9674 ) -> Result<u64> {
9675 let Some(b) = self.bitmap.get(&fk_column_id) else {
9676 return Ok(0);
9677 };
9678 let mut acc = roaring::RoaringBitmap::new();
9679 for v in pk_values {
9680 acc |= b.get(v);
9681 }
9682 if fk_conditions.is_empty() {
9683 return Ok(acc.len());
9684 }
9685 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
9686 sets.push(RowIdSet::from_roaring(acc));
9687 for c in fk_conditions {
9688 sets.push(self.resolve_condition(c, snapshot)?);
9689 }
9690 Ok(RowIdSet::intersect_many(sets).len() as u64)
9691 }
9692
9693 fn inspect_row_at(&self, rid: RowId, snapshot: Snapshot) -> Option<crate::memtable::Row> {
9698 let mut best: Option<crate::memtable::Row> = None;
9699 let mut consider = |row: crate::memtable::Row| {
9700 if !snapshot.observes_row(row.committed_epoch, row.commit_ts) {
9701 return;
9702 }
9703 if best.as_ref().is_none_or(|current| {
9704 Snapshot::version_is_newer(
9705 row.committed_epoch,
9706 row.commit_ts,
9707 current.committed_epoch,
9708 current.commit_ts,
9709 )
9710 }) {
9711 best = Some(row);
9712 }
9713 };
9714 if let Some((_, row)) = self.memtable.get_version_at(rid, snapshot) {
9715 consider(row);
9716 }
9717 if let Some((_, row)) = self.mutable_run.get_version_at(rid, snapshot) {
9718 consider(row);
9719 }
9720 for rr in &self.run_refs {
9721 if let Some(&(min_rid, max_rid)) = self.run_row_id_ranges.get(&rr.run_id) {
9722 if rid.0 < min_rid || rid.0 > max_rid {
9723 continue;
9724 }
9725 }
9726 let Ok(mut reader) = self.open_reader(rr.run_id) else {
9727 continue;
9728 };
9729 let Ok(Some((_, row))) = reader.get_version(rid, snapshot.epoch) else {
9730 continue;
9731 };
9732 consider(row);
9733 }
9734 best
9735 }
9736
9737 fn resolve_pk_with_hot_fallback(
9738 &self,
9739 pk_column_id: u16,
9740 lookup: &[u8],
9741 snapshot: Snapshot,
9742 ) -> Result<RowIdSet> {
9743 if !self.indexes_complete {
9747 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
9748 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::IndexIncomplete);
9749 return Ok(result);
9750 }
9751 let mapped = self.hot.get(lookup);
9754 let Some(r) = mapped else {
9755 let (result, reason) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
9756 self.record_hot_fallback_reason(reason);
9757 return Ok(result);
9758 };
9759 if snapshot.epoch < self.current_epoch() {
9764 self.record_hot_fallback_reason(crate::trace::HotFallbackReason::HistoricalSnapshot);
9765 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
9766 return Ok(result);
9767 }
9768 let materialized = self.inspect_row_at(r, snapshot);
9772 let now_nanos = unix_nanos_now();
9773 let inspection = crate::trace::inspect_hot_candidate(
9774 materialized.as_ref(),
9775 snapshot,
9776 self.ttl,
9777 now_nanos,
9778 pk_column_id,
9779 lookup,
9780 |row| {
9781 let pk_value = row.columns.get(&pk_column_id);
9782 pk_value
9783 .map(|v| self.index_lookup_key(pk_column_id, v))
9784 .unwrap_or_default()
9785 },
9786 );
9787 match inspection {
9788 crate::trace::HotCandidateInspection::Hit(row) => {
9789 self.lookup_metrics
9790 .hot_lookup_hit
9791 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
9792 crate::trace::QueryTrace::record(|t| {
9793 t.hot_lookup_attempted = true;
9794 t.hot_lookup_hit = true;
9795 });
9796 Ok(RowIdSet::one(row.row_id.0))
9797 }
9798 failure => {
9799 let reason = match &failure {
9800 crate::trace::HotCandidateInspection::Hit(_) => unreachable!(),
9801 crate::trace::HotCandidateInspection::MissingRow => {
9802 crate::trace::HotFallbackReason::Tombstone
9807 }
9808 crate::trace::HotCandidateInspection::Invisible => {
9809 crate::trace::HotFallbackReason::InvisibleAtSnapshot
9810 }
9811 crate::trace::HotCandidateInspection::Tombstone => {
9812 crate::trace::HotFallbackReason::Tombstone
9813 }
9814 crate::trace::HotCandidateInspection::TtlExpired => {
9815 crate::trace::HotFallbackReason::TtlExpired
9816 }
9817 crate::trace::HotCandidateInspection::PrimaryKeyMismatch { .. } => {
9818 crate::trace::HotFallbackReason::PrimaryKeyMismatch
9819 }
9820 };
9821 let (result, _inner) = self.pk_equality_fallback(pk_column_id, lookup, snapshot)?;
9827 self.record_hot_fallback_reason(reason);
9828 Ok(result)
9829 }
9830 }
9831 }
9832
9833 fn resolve_condition(
9838 &self,
9839 c: &crate::query::Condition,
9840 snapshot: Snapshot,
9841 ) -> Result<RowIdSet> {
9842 self.resolve_condition_with_allowed(c, snapshot, None)
9843 }
9844
9845 fn resolve_condition_with_allowed(
9846 &self,
9847 c: &crate::query::Condition,
9848 snapshot: Snapshot,
9849 allowed: Option<&std::collections::HashSet<RowId>>,
9850 ) -> Result<RowIdSet> {
9851 use crate::query::Condition;
9852 self.validate_condition(c)?;
9853 Ok(match c {
9854 Condition::Pk(key) => {
9855 let lookup = self
9856 .schema
9857 .primary_key()
9858 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
9859 .unwrap_or_else(|| key.clone());
9860 if let Some(pk_col) = self.schema.primary_key() {
9861 return self.resolve_pk_with_hot_fallback(pk_col.id, &lookup, snapshot);
9862 }
9863 RowIdSet::empty()
9867 }
9868 Condition::BitmapEq { column_id, value } => {
9869 let lookup = self.index_lookup_key_bytes(*column_id, value);
9870 self.bitmap
9871 .get(column_id)
9872 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
9873 .unwrap_or_else(RowIdSet::empty)
9874 }
9875 Condition::BitmapIn { column_id, values } => {
9876 let bm = self.bitmap.get(column_id);
9877 let mut acc = roaring::RoaringBitmap::new();
9878 if let Some(b) = bm {
9879 for v in values {
9880 let lookup = self.index_lookup_key_bytes(*column_id, v);
9881 acc |= b.get(&lookup);
9882 }
9883 }
9884 RowIdSet::from_roaring(acc)
9885 }
9886 Condition::BytesPrefix { column_id, prefix } => {
9887 if let Some(b) = self.bitmap.get(column_id) {
9892 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
9893 let mut acc = roaring::RoaringBitmap::new();
9894 for key in b.keys() {
9895 if key.starts_with(&lookup_prefix) {
9896 acc |= b.get(&key);
9897 }
9898 }
9899 RowIdSet::from_roaring(acc)
9900 } else {
9901 RowIdSet::empty()
9902 }
9903 }
9904 Condition::FmContains { column_id, pattern } => self
9905 .fm
9906 .get(column_id)
9907 .map(|f| {
9908 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
9909 })
9910 .unwrap_or_else(RowIdSet::empty),
9911 Condition::FmContainsAll {
9912 column_id,
9913 patterns,
9914 } => {
9915 if let Some(f) = self.fm.get(column_id) {
9918 let sets: Vec<RowIdSet> = patterns
9919 .iter()
9920 .map(|pat| {
9921 RowIdSet::from_unsorted(
9922 f.locate(pat).into_iter().map(|r| r.0).collect(),
9923 )
9924 })
9925 .collect();
9926 RowIdSet::intersect_many(sets)
9927 } else {
9928 RowIdSet::empty()
9929 }
9930 }
9931 Condition::Ann {
9932 column_id,
9933 query,
9934 k,
9935 } => RowIdSet::from_unsorted(
9936 self.retrieve_filtered(
9937 &crate::query::Retriever::Ann {
9938 column_id: *column_id,
9939 query: query.clone(),
9940 k: *k,
9941 },
9942 snapshot,
9943 None,
9944 allowed,
9945 None,
9946 None,
9947 )?
9948 .into_iter()
9949 .map(|hit| hit.row_id.0)
9950 .collect(),
9951 ),
9952 Condition::SparseMatch {
9953 column_id,
9954 query,
9955 k,
9956 } => RowIdSet::from_unsorted(
9957 self.retrieve_filtered(
9958 &crate::query::Retriever::Sparse {
9959 column_id: *column_id,
9960 query: query.clone(),
9961 k: *k,
9962 },
9963 snapshot,
9964 None,
9965 allowed,
9966 None,
9967 None,
9968 )?
9969 .into_iter()
9970 .map(|hit| hit.row_id.0)
9971 .collect(),
9972 ),
9973 Condition::MinHashSimilar {
9974 column_id,
9975 query,
9976 k,
9977 } => match self.minhash.get(column_id) {
9978 Some(index) => {
9979 let candidates = index.candidate_row_ids(query);
9980 let eligible =
9981 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
9982 RowIdSet::from_unsorted(
9983 index
9984 .search_filtered(query, *k, |row_id| {
9985 eligible.contains(&row_id)
9986 && allowed.is_none_or(|allowed| allowed.contains(&row_id))
9987 })
9988 .into_iter()
9989 .map(|(row_id, _)| row_id.0)
9990 .collect(),
9991 )
9992 }
9993 None => RowIdSet::empty(),
9994 },
9995 Condition::Range { column_id, lo, hi } => {
9996 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10015 if self.run_refs.len() == 1 {
10016 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
10019 } else {
10020 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10028 if lo == hi {
10029 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10030 }
10031 return Ok(multi);
10032 }
10033 } else if self.run_refs.len() == 1 {
10034 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10035 r.range_row_id_set_i64(*column_id, *lo, *hi)?
10036 } else {
10037 let mut multi = self.range_scan_i64(*column_id, *lo, *hi, snapshot)?;
10040 if lo == hi {
10041 self.union_bitmap_point_i64(&mut multi, *column_id, *lo, snapshot);
10042 }
10043 return Ok(multi);
10044 };
10045 if self.learned_range.get(column_id).is_none() && self.run_refs.len() == 1 {
10046 let candidates: Vec<RowId> =
10047 set.into_sorted_vec().into_iter().map(RowId).collect();
10048 set = RowIdSet::from_unsorted(
10049 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?
10050 .into_iter()
10051 .map(|row_id| row_id.0)
10052 .collect(),
10053 );
10054 }
10055 set.remove_many(self.overlay_rid_set(snapshot));
10056 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
10057 if lo == hi {
10058 self.union_bitmap_point_i64(&mut set, *column_id, *lo, snapshot);
10059 }
10060 RowIdSet::from_unsorted(
10061 set.into_sorted_vec()
10062 .into_iter()
10063 .filter(|row_id| self.get(RowId(*row_id), snapshot).is_some())
10064 .collect(),
10065 )
10066 }
10067 Condition::RangeF64 {
10068 column_id,
10069 lo,
10070 lo_inclusive,
10071 hi,
10072 hi_inclusive,
10073 } => {
10074 let mut set = if let Some(li) = self.learned_range.get(column_id) {
10077 RowIdSet::from_unsorted(
10078 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
10079 .into_iter()
10080 .collect(),
10081 )
10082 } else if self.run_refs.len() == 1 {
10083 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10084 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
10085 } else {
10086 return self.range_scan_f64(
10087 *column_id,
10088 *lo,
10089 *lo_inclusive,
10090 *hi,
10091 *hi_inclusive,
10092 snapshot,
10093 );
10094 };
10095 set.remove_many(self.overlay_rid_set(snapshot));
10096 self.range_scan_overlay_f64(
10097 &mut set,
10098 *column_id,
10099 *lo,
10100 *lo_inclusive,
10101 *hi,
10102 *hi_inclusive,
10103 snapshot,
10104 );
10105 set
10106 }
10107 Condition::IsNull { column_id } => {
10108 let mut set = if self.run_refs.len() == 1 {
10109 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10110 r.null_row_id_set(*column_id, true)?
10111 } else {
10112 return self.null_scan(*column_id, true, snapshot);
10113 };
10114 set.remove_many(self.overlay_rid_set(snapshot));
10115 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
10116 set
10117 }
10118 Condition::IsNotNull { column_id } => {
10119 let mut set = if self.run_refs.len() == 1 {
10120 let mut r = self.open_reader(self.run_refs[0].run_id)?;
10121 r.null_row_id_set(*column_id, false)?
10122 } else {
10123 return self.null_scan(*column_id, false, snapshot);
10124 };
10125 set.remove_many(self.overlay_rid_set(snapshot));
10126 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
10127 set
10128 }
10129 })
10130 }
10131
10132 fn range_scan_i64(
10140 &self,
10141 column_id: u16,
10142 lo: i64,
10143 hi: i64,
10144 snapshot: Snapshot,
10145 ) -> Result<RowIdSet> {
10146 let mut row_ids = Vec::new();
10147 let mut tomb_rids: HashSet<u64> = HashSet::new();
10154 let overlay_rids = self.overlay_rid_set(snapshot);
10155 for rr in &self.run_refs {
10156 let mut reader = self.open_reader(rr.run_id)?;
10157 let matched = reader.range_row_ids_visible_i64_at(column_id, lo, hi, snapshot)?;
10158 for rid in matched {
10159 if !overlay_rids.contains(&rid) {
10160 row_ids.push(rid);
10161 }
10162 }
10163 for rid in reader.tombstoned_row_ids_at(snapshot)? {
10164 tomb_rids.insert(rid);
10165 }
10166 }
10167 for row in self
10168 .memtable
10169 .visible_versions_at(Snapshot::at(self.pending_epoch()))
10170 .into_iter()
10171 .chain(
10172 self.mutable_run
10173 .visible_versions_at(Snapshot::at(self.pending_epoch())),
10174 )
10175 {
10176 if row.deleted {
10177 tomb_rids.insert(row.row_id.0);
10178 }
10179 }
10180 let mut s = RowIdSet::from_unsorted(row_ids);
10181 s.remove_many(tomb_rids);
10182 let candidates: Vec<RowId> = s.into_sorted_vec().into_iter().map(RowId).collect();
10183 let eligible = self.eligible_candidate_ids(&candidates, column_id, snapshot, None)?;
10184 let mut s = RowIdSet::from_unsorted(eligible.into_iter().map(|row_id| row_id.0).collect());
10185 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
10186 Ok(s)
10187 }
10188
10189 fn range_scan_f64(
10192 &self,
10193 column_id: u16,
10194 lo: f64,
10195 lo_inclusive: bool,
10196 hi: f64,
10197 hi_inclusive: bool,
10198 snapshot: Snapshot,
10199 ) -> Result<RowIdSet> {
10200 let mut row_ids = Vec::new();
10201 let overlay_rids = self.overlay_rid_set(snapshot);
10202 for rr in &self.run_refs {
10203 let mut reader = self.open_reader(rr.run_id)?;
10204 let matched = reader.range_row_ids_visible_f64_at(
10205 column_id,
10206 lo,
10207 lo_inclusive,
10208 hi,
10209 hi_inclusive,
10210 snapshot,
10211 )?;
10212 for rid in matched {
10213 if !overlay_rids.contains(&rid) {
10214 row_ids.push(rid);
10215 }
10216 }
10217 }
10218 let mut s = RowIdSet::from_unsorted(row_ids);
10219 self.range_scan_overlay_f64(
10220 &mut s,
10221 column_id,
10222 lo,
10223 lo_inclusive,
10224 hi,
10225 hi_inclusive,
10226 snapshot,
10227 );
10228 Ok(s)
10229 }
10230
10231 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
10233 let mut s = HashSet::new();
10234 for row in self.memtable.visible_versions_at(snapshot) {
10235 s.insert(row.row_id.0);
10236 }
10237 for row in self.mutable_run.visible_versions_at(snapshot) {
10238 s.insert(row.row_id.0);
10239 }
10240 s
10241 }
10242
10243 fn range_scan_overlay_i64(
10244 &self,
10245 s: &mut RowIdSet,
10246 column_id: u16,
10247 lo: i64,
10248 hi: i64,
10249 snapshot: Snapshot,
10250 ) {
10251 let mut newest: HashMap<u64, Row> = HashMap::new();
10258 for r in self.mutable_run.visible_versions_at(snapshot) {
10259 newest.insert(r.row_id.0, r);
10260 }
10261 for r in self.memtable.visible_versions_at(snapshot) {
10262 newest
10263 .entry(r.row_id.0)
10264 .and_modify(|cur| {
10265 if Snapshot::version_is_newer(
10266 r.committed_epoch,
10267 r.commit_ts,
10268 cur.committed_epoch,
10269 cur.commit_ts,
10270 ) {
10271 *cur = r.clone();
10272 }
10273 })
10274 .or_insert(r);
10275 }
10276 for row in newest.values() {
10277 if !row.deleted {
10278 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
10279 if *v >= lo && *v <= hi && !self.row_expired_at(row, unix_nanos_now()) {
10280 s.insert(row.row_id.0);
10281 }
10282 }
10283 }
10284 }
10285 }
10286
10287 #[allow(clippy::too_many_arguments)]
10288 fn range_scan_overlay_f64(
10289 &self,
10290 s: &mut RowIdSet,
10291 column_id: u16,
10292 lo: f64,
10293 lo_inclusive: bool,
10294 hi: f64,
10295 hi_inclusive: bool,
10296 snapshot: Snapshot,
10297 ) {
10298 let mut newest: HashMap<u64, Row> = HashMap::new();
10301 for r in self.mutable_run.visible_versions_at(snapshot) {
10302 newest.insert(r.row_id.0, r);
10303 }
10304 for r in self.memtable.visible_versions_at(snapshot) {
10305 newest
10306 .entry(r.row_id.0)
10307 .and_modify(|cur| {
10308 if Snapshot::version_is_newer(
10309 r.committed_epoch,
10310 r.commit_ts,
10311 cur.committed_epoch,
10312 cur.commit_ts,
10313 ) {
10314 *cur = r.clone();
10315 }
10316 })
10317 .or_insert(r);
10318 }
10319 for row in newest.values() {
10320 if !row.deleted {
10321 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
10322 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
10323 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
10324 if ok_lo && ok_hi && !self.row_expired_at(row, unix_nanos_now()) {
10325 s.insert(row.row_id.0);
10326 }
10327 }
10328 }
10329 }
10330 }
10331
10332 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
10335 let mut row_ids = Vec::new();
10336 let overlay_rids = self.overlay_rid_set(snapshot);
10337 for rr in &self.run_refs {
10338 let mut reader = self.open_reader(rr.run_id)?;
10339 let matched = reader.null_row_ids_visible_at(column_id, want_nulls, snapshot)?;
10340 for rid in matched {
10341 if !overlay_rids.contains(&rid) {
10342 row_ids.push(rid);
10343 }
10344 }
10345 }
10346 let mut s = RowIdSet::from_unsorted(row_ids);
10347 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
10348 Ok(s)
10349 }
10350
10351 fn null_scan_overlay(
10355 &self,
10356 s: &mut RowIdSet,
10357 column_id: u16,
10358 want_nulls: bool,
10359 snapshot: Snapshot,
10360 ) {
10361 let mut newest: HashMap<u64, Row> = HashMap::new();
10362 for r in self.mutable_run.visible_versions_at(snapshot) {
10363 newest.insert(r.row_id.0, r);
10364 }
10365 for r in self.memtable.visible_versions_at(snapshot) {
10366 newest
10367 .entry(r.row_id.0)
10368 .and_modify(|cur| {
10369 if Snapshot::version_is_newer(
10370 r.committed_epoch,
10371 r.commit_ts,
10372 cur.committed_epoch,
10373 cur.commit_ts,
10374 ) {
10375 *cur = r.clone();
10376 }
10377 })
10378 .or_insert(r);
10379 }
10380 for row in newest.values() {
10381 if row.deleted {
10382 continue;
10383 }
10384 let is_null = !row.columns.contains_key(&column_id)
10385 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
10386 if is_null == want_nulls {
10387 s.insert(row.row_id.0);
10388 }
10389 }
10390 }
10391
10392 pub fn snapshot(&self) -> Snapshot {
10393 let epoch = self.epoch.visible();
10394 match &self.wal {
10398 WalSink::Shared(shared) => match shared.hlc.now() {
10399 Ok(commit_ts) => Snapshot::at_hlc(epoch, commit_ts),
10400 Err(_) => Snapshot::at_hlc(epoch, mongreldb_types::hlc::HlcTimestamp::MAX),
10402 },
10403 WalSink::Private(_) | WalSink::ReadOnly => Snapshot::at(epoch),
10404 }
10405 }
10406
10407 pub fn data_generation(&self) -> u64 {
10409 self.data_generation
10410 }
10411
10412 pub(crate) fn bump_data_generation(&mut self) {
10413 self.data_generation = self.data_generation.wrapping_add(1);
10414 }
10415
10416 pub fn table_id(&self) -> u64 {
10418 self.table_id
10419 }
10420
10421 fn seal_generations(&mut self) {
10427 self.memtable.seal();
10428 self.mutable_run.seal();
10429 self.hot.seal();
10430 for index in self.bitmap.values_mut() {
10431 index.seal();
10432 }
10433 for index in self.ann.values_mut() {
10434 index.seal();
10435 }
10436 for index in self.fm.values_mut() {
10437 index.seal();
10438 }
10439 for index in self.sparse.values_mut() {
10440 index.seal();
10441 }
10442 for index in self.minhash.values_mut() {
10443 index.seal();
10444 }
10445 self.pk_by_row.seal();
10446 }
10447
10448 fn capture_read_generation(&self) -> ReadGeneration {
10453 let visible_through = self.current_epoch();
10454 ReadGeneration {
10455 schema: Arc::new(self.schema.clone()),
10456 base_runs: Arc::new(self.run_refs.clone()),
10457 deltas: TableDeltas {
10458 memtable: self.memtable.clone(),
10459 mutable_run: self.mutable_run.clone(),
10460 hot: self.hot.clone(),
10461 pk_by_row: self.pk_by_row.clone(),
10462 },
10463 indexes: Arc::new(IndexGeneration::capture(
10464 &self.bitmap,
10465 &self.learned_range,
10466 &self.fm,
10467 &self.ann,
10468 &self.sparse,
10469 &self.minhash,
10470 visible_through,
10471 match &self.wal {
10473 WalSink::Shared(shared) => shared
10474 .hlc
10475 .now()
10476 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::MAX),
10477 _ => mongreldb_types::hlc::HlcTimestamp::MAX,
10478 },
10479 )),
10480 visible_through,
10481 }
10482 }
10483
10484 pub fn publish_read_generation(&mut self) -> Result<Arc<ReadGeneration>> {
10489 self.ensure_indexes_complete()?;
10490 self.seal_generations();
10491 let view = Arc::new(self.capture_read_generation());
10492 self.published.store(Arc::clone(&view));
10493 Ok(view)
10494 }
10495
10496 pub fn published_read_generation(&self) -> Arc<ReadGeneration> {
10502 self.published.load_full()
10503 }
10504
10505 pub fn pin_registry(&self) -> &Arc<crate::retention::PinRegistry> {
10507 &self.pins
10508 }
10509
10510 pub fn version_gc_floor(&self) -> Epoch {
10515 self.min_active_snapshot()
10516 .unwrap_or_else(|| self.current_epoch())
10517 }
10518
10519 pub fn version_pins_report(&self) -> crate::retention::PinsReport {
10526 let mut report = self.pins.report();
10527 let transaction_floor = [
10528 self.pinned.keys().next().copied(),
10529 self.snapshots.min_pinned(),
10530 ]
10531 .into_iter()
10532 .flatten()
10533 .min();
10534 if let Some(epoch) = transaction_floor {
10535 report.record_projection(crate::retention::PinSource::TransactionSnapshot, epoch);
10536 }
10537 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
10538 report.record_projection(crate::retention::PinSource::HistoryRetention, floor);
10539 }
10540 report
10541 }
10542
10543 pub(crate) fn clone_read_generation(&mut self) -> Result<Self> {
10544 self.publish_read_generation()?;
10545 let mut generation = self.clone();
10546 generation.read_only = true;
10547 generation.wal = WalSink::ReadOnly;
10548 generation.pending_delete_rids.clear();
10549 generation.pending_put_cols.clear();
10550 generation.pending_rows.clear();
10551 generation.pending_rows_auto_inc.clear();
10552 generation.pending_dels.clear();
10553 generation.pending_truncate = None;
10554 generation.agg_cache = Arc::new(HashMap::new());
10555 generation.published = Arc::new(ArcSwap::new(self.published.load_full()));
10558 generation.read_generation_pin = Some(Arc::new(self.pins.pin(
10561 crate::retention::PinSource::ReadGeneration,
10562 self.current_epoch(),
10563 )));
10564 Ok(generation)
10565 }
10566
10567 pub(crate) fn estimated_clone_bytes(&self) -> u64 {
10568 (std::mem::size_of::<Self>() as u64)
10569 .saturating_add(self.memtable.approx_bytes())
10570 .saturating_add(self.mutable_run.approx_bytes())
10571 .saturating_add(self.live_count.saturating_mul(64))
10572 }
10573
10574 pub fn pin_snapshot(&mut self) -> Snapshot {
10581 let snap = self.snapshot();
10582 *self.pinned.entry(snap.epoch).or_insert(0) += 1;
10583 snap
10584 }
10585
10586 fn active_pin_epochs_for_rebuild(&self) -> Vec<Epoch> {
10598 let mut set = BTreeSet::new();
10599 for epoch in self.pinned.keys().copied() {
10600 set.insert(epoch);
10601 }
10602 for epoch in self.snapshots.live_pinned_epochs() {
10603 set.insert(epoch);
10604 }
10605 for epoch in self.pins.live_pin_epochs() {
10606 set.insert(epoch);
10607 }
10608 if let Some(floor) = self.snapshots.history_floor(self.current_epoch()) {
10609 set.insert(floor);
10610 }
10611 set.into_iter().collect()
10612 }
10613
10614 pub fn hlc_gc_floor(
10624 &self,
10625 mut project_epoch: impl FnMut(Epoch) -> Option<mongreldb_types::hlc::HlcTimestamp>,
10626 ) -> crate::epoch::GcFloor {
10627 let mut project = |epoch: Option<Epoch>| -> mongreldb_types::hlc::HlcTimestamp {
10628 epoch
10629 .and_then(&mut project_epoch)
10630 .filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
10631 .unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
10632 };
10633 let transaction = [
10634 self.pinned.keys().next().copied(),
10635 self.snapshots.min_pinned(),
10636 ]
10637 .into_iter()
10638 .flatten()
10639 .min();
10640 let history = self.snapshots.history_floor(self.current_epoch());
10641 crate::epoch::GcFloor {
10642 transaction_snapshot: project(transaction),
10643 history_retention: project(history),
10644 backup_pitr: project(
10645 self.pins
10646 .oldest_for(crate::retention::PinSource::BackupPitr),
10647 ),
10648 replication: project(
10649 self.pins
10650 .oldest_for(crate::retention::PinSource::Replication),
10651 ),
10652 read_generation: project(
10653 self.pins
10654 .oldest_for(crate::retention::PinSource::ReadGeneration),
10655 ),
10656 online_index_build: project(
10657 self.pins
10658 .oldest_for(crate::retention::PinSource::OnlineIndexBuild),
10659 ),
10660 }
10661 }
10662
10663 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
10665 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
10666 *count -= 1;
10667 if *count == 0 {
10668 self.pinned.remove(&snap.epoch);
10669 }
10670 }
10671 }
10672
10673 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
10685 let local = self.pinned.keys().next().copied();
10686 let global = self.snapshots.min_pinned();
10687 let history = self.snapshots.history_floor(self.current_epoch());
10688 let pinned = self.pins.oldest_pinned();
10689 [local, global, history, pinned].into_iter().flatten().min()
10690 }
10691
10692 pub fn set_ttl(&mut self, column_name: &str, duration_nanos: u64) -> Result<()> {
10696 self.ensure_writable()?;
10697 let policy = self.prepare_ttl_policy(column_name, duration_nanos)?;
10698 self.apply_ttl_policy_at(Some(policy), self.current_epoch())
10699 }
10700
10701 pub fn clear_ttl(&mut self) -> Result<()> {
10702 self.ensure_writable()?;
10703 self.apply_ttl_policy_at(None, self.current_epoch())
10704 }
10705
10706 pub fn ttl(&self) -> Option<TtlPolicy> {
10707 self.ttl
10708 }
10709
10710 pub(crate) fn prepare_ttl_policy(
10711 &self,
10712 column_name: &str,
10713 duration_nanos: u64,
10714 ) -> Result<TtlPolicy> {
10715 if duration_nanos == 0 || duration_nanos > i64::MAX as u64 {
10716 return Err(MongrelError::InvalidArgument(
10717 "TTL duration must be between 1 and i64::MAX nanoseconds".into(),
10718 ));
10719 }
10720 let column = self
10721 .schema
10722 .columns
10723 .iter()
10724 .find(|column| column.name == column_name)
10725 .ok_or_else(|| MongrelError::Schema(format!("unknown TTL column {column_name}")))?;
10726 if column.ty != TypeId::TimestampNanos {
10727 return Err(MongrelError::Schema(format!(
10728 "TTL column {column_name} must be TimestampNanos, is {:?}",
10729 column.ty
10730 )));
10731 }
10732 Ok(TtlPolicy {
10733 column_id: column.id,
10734 duration_nanos,
10735 })
10736 }
10737
10738 pub(crate) fn apply_ttl_policy_at(
10739 &mut self,
10740 policy: Option<TtlPolicy>,
10741 epoch: Epoch,
10742 ) -> Result<()> {
10743 if let Some(policy) = policy {
10744 let column = self
10745 .schema
10746 .columns
10747 .iter()
10748 .find(|column| column.id == policy.column_id)
10749 .ok_or_else(|| {
10750 MongrelError::Schema(format!("unknown TTL column id {}", policy.column_id))
10751 })?;
10752 if column.ty != TypeId::TimestampNanos
10753 || policy.duration_nanos == 0
10754 || policy.duration_nanos > i64::MAX as u64
10755 {
10756 return Err(MongrelError::Schema("invalid TTL policy".into()));
10757 }
10758 }
10759 self.ttl = policy;
10760 self.agg_cache = Arc::new(HashMap::new());
10761 self.clear_result_cache();
10762 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
10763 self.persist_manifest(epoch)
10764 }
10765
10766 pub(crate) fn row_expired_at(&self, row: &Row, now_nanos: i64) -> bool {
10767 let Some(policy) = self.ttl else {
10768 return false;
10769 };
10770 let Some(Value::Int64(timestamp)) = row.columns.get(&policy.column_id) else {
10771 return false;
10772 };
10773 timestamp.saturating_add(policy.duration_nanos as i64) <= now_nanos
10774 }
10775
10776 pub fn current_epoch(&self) -> Epoch {
10777 self.epoch.visible()
10778 }
10779
10780 pub fn memtable_len(&self) -> usize {
10781 self.memtable.len()
10782 }
10783
10784 pub fn count(&self) -> u64 {
10787 if self.ttl.is_none()
10788 && self.pending_put_cols.is_empty()
10789 && self.pending_delete_rids.is_empty()
10790 && self.pending_rows.is_empty()
10791 && self.pending_dels.is_empty()
10792 && self.pending_truncate.is_none()
10793 {
10794 self.live_count
10795 } else {
10796 self.visible_rows(self.snapshot())
10797 .map(|rows| rows.len() as u64)
10798 .unwrap_or(self.live_count)
10799 }
10800 }
10801
10802 pub fn count_conditions(
10806 &mut self,
10807 conditions: &[crate::query::Condition],
10808 snapshot: Snapshot,
10809 ) -> Result<Option<u64>> {
10810 use crate::query::Condition;
10811 if self.ttl.is_some() {
10812 if conditions.is_empty() {
10813 return Ok(Some(self.visible_rows(snapshot)?.len() as u64));
10814 }
10815 let mut sets = Vec::with_capacity(conditions.len());
10816 for condition in conditions {
10817 sets.push(self.resolve_condition(condition, snapshot)?);
10818 }
10819 let survivors = RowIdSet::intersect_many(sets);
10820 let rows = self.visible_rows(snapshot)?;
10821 return Ok(Some(
10822 rows.into_iter()
10823 .filter(|row| survivors.contains(row.row_id.0))
10824 .count() as u64,
10825 ));
10826 }
10827 if conditions.is_empty() {
10828 return Ok(Some(self.count()));
10829 }
10830 let served = |c: &Condition| {
10831 matches!(
10832 c,
10833 Condition::Pk(_)
10834 | Condition::BitmapEq { .. }
10835 | Condition::BitmapIn { .. }
10836 | Condition::BytesPrefix { .. }
10837 | Condition::FmContains { .. }
10838 | Condition::FmContainsAll { .. }
10839 | Condition::Ann { .. }
10840 | Condition::Range { .. }
10841 | Condition::RangeF64 { .. }
10842 | Condition::SparseMatch { .. }
10843 | Condition::MinHashSimilar { .. }
10844 | Condition::IsNull { .. }
10845 | Condition::IsNotNull { .. }
10846 )
10847 };
10848 if !conditions.iter().all(served) {
10849 return Ok(None);
10850 }
10851 self.ensure_indexes_complete()?;
10852 if !self.pending_put_cols.is_empty()
10853 || !self.pending_delete_rids.is_empty()
10854 || !self.pending_rows.is_empty()
10855 || !self.pending_dels.is_empty()
10856 || self.pending_truncate.is_some()
10857 {
10858 let mut sets = Vec::with_capacity(conditions.len());
10859 for condition in conditions {
10860 sets.push(self.resolve_condition(condition, snapshot)?);
10861 }
10862 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
10863 return Ok(Some(self.rows_for_rids(&rids, snapshot)?.len() as u64));
10864 }
10865 let mut sets = Vec::with_capacity(conditions.len());
10866 for condition in conditions {
10867 sets.push(self.resolve_condition(condition, snapshot)?);
10868 }
10869 let rids = RowIdSet::intersect_many(sets);
10870 if self.had_deletes || !self.memtable.is_empty() || !self.mutable_run.is_empty() {
10878 let sorted = rids.into_sorted_vec();
10879 let count = self.rows_for_rids(&sorted, snapshot)?.len() as u64;
10880 crate::trace::QueryTrace::record(|t| {
10881 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
10882 t.survivor_count = Some(count as usize);
10883 t.conditions_pushed = conditions.len();
10884 });
10885 return Ok(Some(count));
10886 }
10887 let count = rids.len() as u64;
10888 crate::trace::QueryTrace::record(|t| {
10889 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
10890 t.survivor_count = Some(count as usize);
10891 t.conditions_pushed = conditions.len();
10892 });
10893 Ok(Some(count))
10894 }
10895
10896 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
10902 let mut out = Vec::new();
10903 for row in self.memtable.visible_versions_at(snapshot) {
10904 if row.deleted {
10905 out.push(row.row_id.0);
10906 }
10907 }
10908 for row in self.mutable_run.visible_versions_at(snapshot) {
10909 if row.deleted {
10910 out.push(row.row_id.0);
10911 }
10912 }
10913 out
10914 }
10915
10916 pub fn bulk_load_columns(
10925 &mut self,
10926 user_columns: Vec<(u16, columnar::NativeColumn)>,
10927 ) -> Result<Epoch> {
10928 self.bulk_load_columns_with(user_columns, 3, false, true)
10929 }
10930
10931 pub fn bulk_load_fast(
10938 &mut self,
10939 user_columns: Vec<(u16, columnar::NativeColumn)>,
10940 ) -> Result<Epoch> {
10941 self.bulk_load_columns_with(user_columns, -1, true, false)
10942 }
10943
10944 fn bulk_load_columns_with(
10945 &mut self,
10946 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
10947 zstd_level: i32,
10948 force_plain: bool,
10949 lz4: bool,
10950 ) -> Result<Epoch> {
10951 self.ensure_writable()?;
10952 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
10953 if n == 0 {
10954 return Ok(self.current_epoch());
10955 }
10956 let epoch = self.commit_new_epoch()?;
10957 let live_before = self.live_count;
10958 self.spill_mutable_run(epoch)?;
10960 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
10961 && self.indexes_complete
10962 && self.run_refs.is_empty()
10963 && self.memtable.is_empty()
10964 && self.mutable_run.is_empty();
10965 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
10968 self.validate_columns_not_null(&user_columns, n)?;
10969 let winner_idx = self
10970 .bulk_pk_winner_indices(&user_columns, n)
10971 .filter(|idx| idx.len() != n);
10972 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
10973 match winner_idx.as_deref() {
10974 Some(idx) => {
10975 let compacted = user_columns
10976 .iter()
10977 .map(|(id, c)| (*id, c.gather(idx)))
10978 .collect();
10979 (compacted, idx.len())
10980 }
10981 None => (user_columns, n),
10982 };
10983 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
10984 let first = self.allocator.alloc_range(write_n as u64)?.0;
10985 for rid in first..first + write_n as u64 {
10986 self.reservoir.offer(rid);
10987 }
10988 let run_id = self.alloc_run_id()?;
10989 let path = self.run_path(run_id);
10990 let mut writer =
10991 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
10992 if force_plain {
10993 writer = writer.with_plain();
10994 } else if lz4 {
10995 writer = writer.with_lz4();
10998 } else {
10999 writer = writer.with_zstd_level(zstd_level);
11000 }
11001 if let Some(kek) = &self.kek {
11002 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
11003 }
11004 let header = match self.create_run_file(run_id)? {
11005 Some(file) => writer.write_native_file(file, &write_columns, write_n, first)?,
11006 None => writer.write_native(&path, &write_columns, write_n, first)?,
11007 };
11008 self.run_refs.push(RunRef {
11009 run_id: run_id as u128,
11010 level: 0,
11011 epoch_created: epoch.0,
11012 row_count: header.row_count,
11013 });
11014 self.run_row_id_ranges
11015 .insert(run_id as u128, (header.min_row_id, header.max_row_id));
11016 self.live_count = self.live_count.saturating_add(write_n as u64);
11017 if eager_index_build {
11018 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
11019 self.index_columns_bulk(&write_columns, &row_ids);
11020 self.indexes_complete = true;
11021 self.build_learned_ranges()?;
11022 } else {
11023 self.indexes_complete = false;
11027 }
11028 self.mark_flushed(epoch)?;
11029 self.persist_manifest(epoch)?;
11030 if eager_index_build {
11031 self.checkpoint_indexes(epoch);
11032 }
11033 self.clear_result_cache();
11034 self.data_generation = self.data_generation.wrapping_add(1);
11035 Ok(epoch)
11036 }
11037
11038 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
11056 let n = row_ids.len();
11057 if n == 0 {
11058 return;
11059 }
11060 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
11061 columns.iter().map(|(id, c)| (*id, c)).collect();
11062 let ty_of: std::collections::HashMap<u16, TypeId> = self
11063 .schema
11064 .columns
11065 .iter()
11066 .map(|c| (c.id, c.ty.clone()))
11067 .collect();
11068 let pk_id = self.schema.primary_key().map(|c| c.id);
11069
11070 for (i, &rid) in row_ids.iter().enumerate() {
11071 let row_id = RowId(rid);
11072 if let Some(pid) = pk_id {
11073 if let Some(col) = by_id.get(&pid) {
11074 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
11075 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
11076 self.insert_hot_pk(key, row_id);
11077 }
11078 }
11079 }
11080 for idef in &self.schema.indexes {
11081 let Some(col) = by_id.get(&idef.column_id) else {
11082 continue;
11083 };
11084 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
11085 match idef.kind {
11086 IndexKind::Bitmap => {
11087 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
11088 if let Some(key) =
11089 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
11090 {
11091 b.insert(key, row_id);
11092 }
11093 }
11094 }
11095 IndexKind::FmIndex => {
11096 if let Some(f) = self.fm.get_mut(&idef.column_id) {
11097 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11098 f.insert(bytes.to_vec(), row_id);
11099 }
11100 }
11101 }
11102 IndexKind::Sparse => {
11103 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
11104 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11105 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
11106 s.insert(&terms, row_id);
11107 }
11108 }
11109 }
11110 }
11111 IndexKind::MinHash => {
11112 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
11113 if let Some(bytes) = columnar::native_bytes_at(col, i) {
11114 let tokens = crate::index::token_hashes_from_bytes(bytes);
11115 mh.insert(&tokens, row_id);
11116 }
11117 }
11118 }
11119 _ => {}
11120 }
11121 }
11122 }
11123 }
11124
11125 pub fn visible_columns_native(
11130 &self,
11131 snapshot: Snapshot,
11132 projection: Option<&[u16]>,
11133 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11134 self.visible_columns_native_inner(snapshot, projection, None)
11135 }
11136
11137 pub fn visible_columns_native_with_control(
11138 &self,
11139 snapshot: Snapshot,
11140 projection: Option<&[u16]>,
11141 control: &crate::ExecutionControl,
11142 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11143 self.visible_columns_native_inner(snapshot, projection, Some(control))
11144 }
11145
11146 fn visible_columns_native_inner(
11147 &self,
11148 snapshot: Snapshot,
11149 projection: Option<&[u16]>,
11150 control: Option<&crate::ExecutionControl>,
11151 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
11152 execution_checkpoint(control, 0)?;
11153 let wanted: Vec<u16> = match projection {
11154 Some(p) => p.to_vec(),
11155 None => self.schema.columns.iter().map(|c| c.id).collect(),
11156 };
11157 if self.ttl.is_none()
11158 && self.memtable.is_empty()
11159 && self.mutable_run.is_empty()
11160 && self.run_refs.len() == 1
11161 {
11162 let rr = self.run_refs[0].clone();
11163 let mut reader = self.open_reader(rr.run_id)?;
11164 let idxs = reader.visible_indices_native_at(snapshot)?;
11165 execution_checkpoint(control, 0)?;
11166 let all_visible = idxs.len() == reader.row_count();
11167 if reader.has_mmap() && control.is_none() {
11173 use rayon::prelude::*;
11174 let valid: Vec<u16> = wanted
11177 .iter()
11178 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
11179 .copied()
11180 .collect();
11181 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
11183 .par_iter()
11184 .filter_map(|cid| {
11185 reader
11186 .column_native_shared(*cid)
11187 .ok()
11188 .map(|col| (*cid, col))
11189 })
11190 .collect();
11191 let cols = decoded
11192 .into_iter()
11193 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
11194 .collect();
11195 return Ok(cols);
11196 }
11197 let mut cols = Vec::with_capacity(wanted.len());
11198 for (index, cid) in wanted.iter().enumerate() {
11199 execution_checkpoint(control, index)?;
11200 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
11201 Some(c) => c,
11202 None => continue,
11203 };
11204 let col = reader.column_native(cdef.id)?;
11205 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
11206 }
11207 return Ok(cols);
11208 }
11209 let vcols = self.visible_columns(snapshot)?;
11210 execution_checkpoint(control, 0)?;
11211 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
11212 let out: Vec<(u16, columnar::NativeColumn)> = vcols
11213 .into_iter()
11214 .filter(|(id, _)| want_set.contains(id))
11215 .map(|(id, vals)| {
11216 let ty = self
11217 .schema
11218 .columns
11219 .iter()
11220 .find(|c| c.id == id)
11221 .map(|c| c.ty.clone())
11222 .unwrap_or(TypeId::Bytes);
11223 (id, columnar::values_to_native(ty, &vals))
11224 })
11225 .collect();
11226 Ok(out)
11227 }
11228
11229 pub fn run_count(&self) -> usize {
11230 self.run_refs.len()
11231 }
11232
11233 pub fn memtable_is_empty(&self) -> bool {
11235 self.memtable.is_empty()
11236 }
11237
11238 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
11242 self.page_cache.stats()
11243 }
11244
11245 pub fn reset_page_cache_stats(&self) {
11247 self.page_cache.reset_stats();
11248 }
11249
11250 pub fn run_ids(&self) -> Vec<u128> {
11253 self.run_refs.iter().map(|r| r.run_id).collect()
11254 }
11255
11256 pub fn single_run_is_clean(&self) -> bool {
11260 if self.ttl.is_some() || self.run_refs.len() != 1 {
11261 return false;
11262 }
11263 self.open_reader(self.run_refs[0].run_id)
11264 .map(|r| r.is_clean())
11265 .unwrap_or(false)
11266 }
11267
11268 fn resolve_footprint(
11275 &self,
11276 conditions: &[crate::query::Condition],
11277 snapshot: Snapshot,
11278 ) -> roaring::RoaringBitmap {
11279 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
11280 return roaring::RoaringBitmap::new();
11281 }
11282 if self.run_refs.is_empty() {
11283 return roaring::RoaringBitmap::new();
11284 }
11285 if self.run_refs.len() == 1 {
11287 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
11288 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader, snapshot) {
11289 return rids.to_roaring_lossy();
11290 }
11291 }
11292 }
11293 roaring::RoaringBitmap::new()
11294 }
11295
11296 pub fn query_columns_native_cached(
11307 &mut self,
11308 conditions: &[crate::query::Condition],
11309 projection: Option<&[u16]>,
11310 snapshot: Snapshot,
11311 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11312 self.query_columns_native_cached_inner(conditions, projection, snapshot, None)
11313 }
11314
11315 pub fn query_columns_native_cached_with_control(
11316 &mut self,
11317 conditions: &[crate::query::Condition],
11318 projection: Option<&[u16]>,
11319 snapshot: Snapshot,
11320 control: &crate::ExecutionControl,
11321 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11322 self.query_columns_native_cached_inner(conditions, projection, snapshot, Some(control))
11323 }
11324
11325 fn query_columns_native_cached_inner(
11326 &mut self,
11327 conditions: &[crate::query::Condition],
11328 projection: Option<&[u16]>,
11329 snapshot: Snapshot,
11330 control: Option<&crate::ExecutionControl>,
11331 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11332 execution_checkpoint(control, 0)?;
11333 if self.ttl.is_some() {
11336 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11337 }
11338 if conditions.is_empty() {
11339 return self.query_columns_native_inner(conditions, projection, snapshot, control);
11340 }
11341 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
11345 if let Some(hit) = self.result_cache.lock().get_columns(key) {
11346 crate::trace::QueryTrace::record(|t| {
11347 t.result_cache_hit = true;
11348 t.scan_mode = crate::trace::ScanMode::NativePushdown;
11349 });
11350 return Ok(Some((*hit).clone()));
11351 }
11352 let res = self.query_columns_native_inner(conditions, projection, snapshot, control)?;
11353 execution_checkpoint(control, 0)?;
11354 if let Some(cols) = &res {
11355 let footprint = self.resolve_footprint(conditions, snapshot);
11356 let condition_cols = crate::query::condition_columns(conditions);
11357 execution_checkpoint(control, 0)?;
11358 let entry = CachedEntry {
11359 data: CachedData::Columns(Arc::new(cols.clone())),
11360 footprint,
11361 condition_cols,
11362 };
11363 let table_id = self.table_id();
11364 let schema_id = self.schema.schema_id;
11365 let data_generation = self.data_generation;
11366 let mut cache = self.result_cache.lock();
11371 let entry_generation = cache.allocate_persist_generation(key);
11372 cache.enqueue_persist(key, &entry, table_id, schema_id, data_generation, entry_generation);
11373 cache.insert(key, entry);
11374 }
11375 Ok(res)
11376 }
11377
11378 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
11383 if self.ttl.is_some() {
11384 return self.query(q);
11385 }
11386 if q.conditions.is_empty() {
11387 return self.query(q);
11388 }
11389 let key = crate::query::canonical_query_key(&q.conditions, None, 0)
11390 ^ (q.limit.unwrap_or(usize::MAX) as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15)
11391 ^ (q.offset as u64).wrapping_mul(0xC2B2_AE3D_27D4_EB4F);
11392 if let Some(hit) = self.result_cache.lock().get_rows(key) {
11393 crate::trace::QueryTrace::record(|t| {
11394 t.result_cache_hit = true;
11395 t.scan_mode = crate::trace::ScanMode::Materialized;
11396 });
11397 return Ok((*hit).clone());
11398 }
11399 let rows = self.query(q)?;
11400 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
11401 let condition_cols = crate::query::condition_columns(&q.conditions);
11402 let entry = CachedEntry {
11403 data: CachedData::Rows(Arc::new(rows.clone())),
11404 footprint,
11405 condition_cols,
11406 };
11407 let table_id = self.table_id();
11408 let schema_id = self.schema.schema_id;
11409 let data_generation = self.data_generation;
11410 let mut cache = self.result_cache.lock();
11413 let entry_generation = cache.allocate_persist_generation(key);
11414 cache.enqueue_persist(key, &entry, table_id, schema_id, data_generation, entry_generation);
11415 cache.insert(key, entry);
11416 Ok(rows)
11417 }
11418
11419 #[allow(clippy::type_complexity)]
11434 pub fn query_columns_native_traced(
11435 &mut self,
11436 conditions: &[crate::query::Condition],
11437 projection: Option<&[u16]>,
11438 snapshot: Snapshot,
11439 ) -> Result<(
11440 Option<Vec<(u16, columnar::NativeColumn)>>,
11441 crate::trace::QueryTrace,
11442 )> {
11443 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11444 self.query_columns_native(conditions, projection, snapshot)
11445 });
11446 Ok((result?, trace))
11447 }
11448
11449 #[allow(clippy::type_complexity)]
11452 pub fn query_columns_native_cached_traced(
11453 &mut self,
11454 conditions: &[crate::query::Condition],
11455 projection: Option<&[u16]>,
11456 snapshot: Snapshot,
11457 ) -> Result<(
11458 Option<Vec<(u16, columnar::NativeColumn)>>,
11459 crate::trace::QueryTrace,
11460 )> {
11461 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11462 self.query_columns_native_cached(conditions, projection, snapshot)
11463 });
11464 Ok((result?, trace))
11465 }
11466
11467 pub fn native_page_cursor_traced(
11469 &self,
11470 snapshot: Snapshot,
11471 projection: Vec<(u16, TypeId)>,
11472 conditions: &[crate::query::Condition],
11473 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
11474 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11475 self.native_page_cursor(snapshot, projection, conditions)
11476 });
11477 Ok((result?, trace))
11478 }
11479
11480 pub fn native_multi_run_cursor_traced(
11482 &self,
11483 snapshot: Snapshot,
11484 projection: Vec<(u16, TypeId)>,
11485 conditions: &[crate::query::Condition],
11486 ) -> Result<(
11487 Option<crate::cursor::MultiRunCursor>,
11488 crate::trace::QueryTrace,
11489 )> {
11490 let (result, trace) = crate::trace::QueryTrace::capture(|| {
11491 self.native_multi_run_cursor(snapshot, projection, conditions)
11492 });
11493 Ok((result?, trace))
11494 }
11495
11496 pub fn count_conditions_traced(
11498 &mut self,
11499 conditions: &[crate::query::Condition],
11500 snapshot: Snapshot,
11501 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
11502 let (result, trace) =
11503 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
11504 Ok((result?, trace))
11505 }
11506
11507 pub fn query_traced(
11509 &mut self,
11510 q: &crate::query::Query,
11511 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
11512 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
11513 Ok((result?, trace))
11514 }
11515
11516 pub fn query_columns_native(
11521 &mut self,
11522 conditions: &[crate::query::Condition],
11523 projection: Option<&[u16]>,
11524 snapshot: Snapshot,
11525 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11526 self.query_columns_native_inner(conditions, projection, snapshot, None)
11527 }
11528
11529 pub fn query_columns_native_with_control(
11530 &mut self,
11531 conditions: &[crate::query::Condition],
11532 projection: Option<&[u16]>,
11533 snapshot: Snapshot,
11534 control: &crate::ExecutionControl,
11535 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11536 self.query_columns_native_inner(conditions, projection, snapshot, Some(control))
11537 }
11538
11539 fn query_columns_native_inner(
11540 &mut self,
11541 conditions: &[crate::query::Condition],
11542 projection: Option<&[u16]>,
11543 snapshot: Snapshot,
11544 control: Option<&crate::ExecutionControl>,
11545 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
11546 use crate::query::Condition;
11547 execution_checkpoint(control, 0)?;
11548 if self.ttl.is_some() {
11551 return Ok(None);
11552 }
11553 if conditions.is_empty() {
11554 return Ok(None);
11555 }
11556 self.ensure_indexes_complete()?;
11557
11558 let served = |c: &Condition| {
11563 matches!(
11564 c,
11565 Condition::Pk(_)
11566 | Condition::BitmapEq { .. }
11567 | Condition::BitmapIn { .. }
11568 | Condition::BytesPrefix { .. }
11569 | Condition::FmContains { .. }
11570 | Condition::FmContainsAll { .. }
11571 | Condition::Ann { .. }
11572 | Condition::Range { .. }
11573 | Condition::RangeF64 { .. }
11574 | Condition::SparseMatch { .. }
11575 | Condition::MinHashSimilar { .. }
11576 | Condition::IsNull { .. }
11577 | Condition::IsNotNull { .. }
11578 )
11579 };
11580 if !conditions.iter().all(served) {
11581 return Ok(None);
11582 }
11583 let fast_path =
11584 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
11585 crate::trace::QueryTrace::record(|t| {
11586 t.run_count = self.run_refs.len();
11587 t.memtable_rows = self.memtable.len();
11588 t.mutable_run_rows = self.mutable_run.len();
11589 t.conditions_pushed = conditions.len();
11590 t.learned_range_used = conditions.iter().any(|c| match c {
11591 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
11592 self.learned_range.contains_key(column_id)
11593 }
11594 _ => false,
11595 });
11596 });
11597 let col_ids: Vec<u16> = projection
11599 .map(|p| p.to_vec())
11600 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
11601 let proj_pairs: Vec<(u16, TypeId)> = col_ids
11602 .iter()
11603 .map(|&cid| {
11604 let ty = self
11605 .schema
11606 .columns
11607 .iter()
11608 .find(|c| c.id == cid)
11609 .map(|c| c.ty.clone())
11610 .unwrap_or(TypeId::Bytes);
11611 (cid, ty)
11612 })
11613 .collect();
11614
11615 if fast_path {
11621 let needs_column = conditions.iter().any(|c| match c {
11624 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
11625 Condition::RangeF64 { column_id, .. } => {
11626 !self.learned_range.contains_key(column_id)
11627 }
11628 _ => false,
11629 });
11630 let mut reader_opt: Option<RunReader> = if needs_column {
11631 Some(self.open_reader(self.run_refs[0].run_id)?)
11632 } else {
11633 None
11634 };
11635 let mut sets: Vec<RowIdSet> = Vec::new();
11636 for (index, c) in conditions.iter().enumerate() {
11637 execution_checkpoint(control, index)?;
11638 let s = match c {
11639 Condition::Range { column_id, lo, hi }
11640 if !self.learned_range.contains_key(column_id) =>
11641 {
11642 if reader_opt.is_none() {
11643 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
11644 }
11645 reader_opt
11646 .as_mut()
11647 .expect("reader opened for range")
11648 .range_row_id_set_i64(*column_id, *lo, *hi)?
11649 }
11650 Condition::RangeF64 {
11651 column_id,
11652 lo,
11653 lo_inclusive,
11654 hi,
11655 hi_inclusive,
11656 } if !self.learned_range.contains_key(column_id) => {
11657 if reader_opt.is_none() {
11658 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
11659 }
11660 reader_opt
11661 .as_mut()
11662 .expect("reader opened for range")
11663 .range_row_id_set_f64(
11664 *column_id,
11665 *lo,
11666 *lo_inclusive,
11667 *hi,
11668 *hi_inclusive,
11669 )?
11670 }
11671 _ => self.resolve_condition(c, snapshot)?,
11672 };
11673 sets.push(s);
11674 }
11675 let candidates = RowIdSet::intersect_many(sets);
11676 crate::trace::QueryTrace::record(|t| {
11677 t.survivor_count = Some(candidates.len());
11678 });
11679 if candidates.is_empty() {
11680 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
11681 .iter()
11682 .map(|&id| {
11683 (
11684 id,
11685 columnar::null_native(
11686 proj_pairs
11687 .iter()
11688 .find(|(c, _)| c == &id)
11689 .map(|(_, t)| t.clone())
11690 .unwrap_or(TypeId::Bytes),
11691 0,
11692 ),
11693 )
11694 })
11695 .collect();
11696 return Ok(Some(cols));
11697 }
11698 let mut reader = match reader_opt.take() {
11699 Some(r) => r,
11700 None => self.open_reader(self.run_refs[0].run_id)?,
11701 };
11702 let candidate_ids = candidates.into_sorted_vec();
11703 let (positions, fast_rid) = if let Some(positions) =
11704 reader.positions_for_row_ids_fast(&candidate_ids)
11705 {
11706 (positions, true)
11707 } else {
11708 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
11709 match col {
11710 columnar::NativeColumn::Int64 { data, .. } => {
11711 let mut p = Vec::with_capacity(candidate_ids.len());
11712 for (index, rid) in candidate_ids.iter().enumerate() {
11713 execution_checkpoint(control, index)?;
11714 if let Ok(position) = data.binary_search(&(*rid as i64)) {
11715 p.push(position);
11716 }
11717 }
11718 p.sort_unstable();
11719 (p, false)
11720 }
11721 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
11722 }
11723 };
11724 crate::trace::QueryTrace::record(|t| {
11725 t.scan_mode = crate::trace::ScanMode::NativePushdown;
11726 t.fast_row_id_map = fast_rid;
11727 });
11728 let mut cols = Vec::with_capacity(col_ids.len());
11729 for (index, cid) in col_ids.iter().enumerate() {
11730 execution_checkpoint(control, index)?;
11731 let col = reader.column_native(*cid)?;
11732 cols.push((*cid, col.gather(&positions)));
11733 }
11734 return Ok(Some(cols));
11735 }
11736
11737 if !self.run_refs.is_empty() {
11750 use crate::cursor::{
11751 drain_cursor_to_columns, drain_cursor_to_columns_with_control, Cursor,
11752 };
11753 let remaining: usize;
11754 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
11755 let c = self
11756 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
11757 .expect("single-run cursor should build when run_refs.len() == 1");
11758 remaining = c.remaining_rows();
11759 Box::new(c)
11760 } else {
11761 let c = self
11762 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
11763 .expect("multi-run cursor should build when run_refs.len() >= 1");
11764 remaining = c.remaining_rows();
11765 Box::new(c)
11766 };
11767 crate::trace::QueryTrace::record(|t| {
11768 if t.survivor_count.is_none() {
11769 t.survivor_count = Some(remaining);
11770 }
11771 });
11772 let cols = match control {
11773 Some(control) => {
11774 drain_cursor_to_columns_with_control(cursor.as_mut(), &proj_pairs, control)?
11775 }
11776 None => drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?,
11777 };
11778 return Ok(Some(cols));
11779 }
11780
11781 crate::trace::QueryTrace::record(|t| {
11786 t.scan_mode = crate::trace::ScanMode::Materialized;
11787 t.row_materialized = true;
11788 });
11789 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
11790 for (index, c) in conditions.iter().enumerate() {
11791 execution_checkpoint(control, index)?;
11792 sets.push(self.resolve_condition(c, snapshot)?);
11793 }
11794 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
11795 let rows = self.rows_for_rids(&rids, snapshot)?;
11796 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
11797 for (index, (cid, ty)) in proj_pairs.iter().enumerate() {
11798 execution_checkpoint(control, index)?;
11799 let vals: Vec<Value> = rows
11800 .iter()
11801 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
11802 .collect();
11803 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
11804 }
11805 Ok(Some(cols))
11806 }
11807
11808 pub fn native_page_cursor(
11823 &self,
11824 snapshot: Snapshot,
11825 projection: Vec<(u16, TypeId)>,
11826 conditions: &[crate::query::Condition],
11827 ) -> Result<Option<NativePageCursor>> {
11828 use crate::cursor::build_page_plans;
11829 if self.ttl.is_some() {
11830 return Ok(None);
11831 }
11832 if !conditions.is_empty() && !self.indexes_complete {
11835 return Ok(None);
11836 }
11837 if self.run_refs.len() != 1 {
11838 return Ok(None);
11839 }
11840 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
11841 let (positions, rids) = reader.visible_positions_with_rids_at(snapshot)?;
11842
11843 let overlay_rids: HashSet<u64> = {
11846 let mut s = HashSet::new();
11847 for row in self.memtable.visible_versions_at(snapshot) {
11848 s.insert(row.row_id.0);
11849 }
11850 for row in self.mutable_run.visible_versions_at(snapshot) {
11851 s.insert(row.row_id.0);
11852 }
11853 s
11854 };
11855
11856 let survivors = if conditions.is_empty() {
11860 None
11861 } else {
11862 Some(self.resolve_survivor_rids(conditions, &mut reader, snapshot)?)
11863 };
11864
11865 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
11872 survivors.clone()
11873 } else if let Some(s) = &survivors {
11874 let mut run_set = s.clone();
11875 run_set.remove_many(overlay_rids.iter().copied());
11876 Some(run_set)
11877 } else {
11878 Some(RowIdSet::from_unsorted(
11879 rids.iter()
11880 .map(|&r| r as u64)
11881 .filter(|r| !overlay_rids.contains(r))
11882 .collect(),
11883 ))
11884 };
11885
11886 let overlay_rows = if overlay_rids.is_empty() {
11887 Vec::new()
11888 } else {
11889 let bound = Self::overlay_materialization_bound(conditions, &survivors);
11890 self.overlay_visible_rows(snapshot, bound)
11891 };
11892
11893 let plans = if positions.is_empty() {
11895 Vec::new()
11896 } else {
11897 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
11898 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
11899 };
11900
11901 let overlay = if overlay_rows.is_empty() {
11903 None
11904 } else {
11905 let filtered =
11906 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
11907 if filtered.is_empty() {
11908 None
11909 } else {
11910 Some(self.materialize_overlay(&filtered, &projection))
11911 }
11912 };
11913
11914 let overlay_row_count = overlay
11915 .as_ref()
11916 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
11917 .unwrap_or(0);
11918 crate::trace::QueryTrace::record(|t| {
11919 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
11920 t.run_count = self.run_refs.len();
11921 t.memtable_rows = self.memtable.len();
11922 t.mutable_run_rows = self.mutable_run.len();
11923 t.overlay_rows = overlay_row_count;
11924 t.conditions_pushed = conditions.len();
11925 t.pages_decoded = plans
11926 .iter()
11927 .map(|p| p.positions.len())
11928 .sum::<usize>()
11929 .min(1);
11930 });
11931
11932 Ok(Some(NativePageCursor::new_with_overlay(
11933 reader, projection, plans, overlay,
11934 )))
11935 }
11936 #[allow(clippy::type_complexity)]
11946 pub fn native_multi_run_cursor(
11947 &self,
11948 snapshot: Snapshot,
11949 projection: Vec<(u16, TypeId)>,
11950 conditions: &[crate::query::Condition],
11951 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
11952 use crate::cursor::{MultiRunCursor, RunStream};
11953 use crate::sorted_run::SYS_ROW_ID;
11954 use std::collections::{BinaryHeap, HashMap, HashSet};
11955 if self.ttl.is_some() {
11956 return Ok(None);
11957 }
11958 if !conditions.is_empty() && !self.indexes_complete {
11961 return Ok(None);
11962 }
11963 if self.run_refs.is_empty() {
11964 return Ok(None);
11965 }
11966
11967 let mut run_meta: Vec<(
11969 RunReader,
11970 Vec<i64>,
11971 Vec<i64>,
11972 Vec<u8>,
11973 Option<Vec<Option<HlcTimestamp>>>,
11974 Vec<usize>,
11975 )> = Vec::with_capacity(self.run_refs.len());
11976 for rr in &self.run_refs {
11977 let mut reader = self.open_reader(rr.run_id)?;
11978 let (rids, eps, del) = reader.system_columns_native()?;
11979 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
11980 let commit_ts: Option<Vec<Option<HlcTimestamp>>> =
11981 if reader.has_column(crate::sorted_run::SYS_COMMIT_TS) {
11982 let col = reader.column_native(crate::sorted_run::SYS_COMMIT_TS)?;
11983 Some(
11984 (0..rids.len())
11985 .map(|i| {
11986 crate::sorted_run::decode_commit_ts_value(col.value_at(i).as_ref())
11987 })
11988 .collect(),
11989 )
11990 } else {
11991 None
11992 };
11993 run_meta.push((reader, rids, eps, del, commit_ts, page_rows));
11994 }
11995
11996 let mut best: HashMap<u64, (Epoch, Option<HlcTimestamp>, usize, usize, bool)> =
12002 HashMap::new();
12003 for (run_idx, (_, rids, eps, del, commit_ts, _)) in run_meta.iter().enumerate() {
12004 for i in 0..rids.len() {
12005 let rid = rids[i] as u64;
12006 let e = Epoch(eps[i] as u64);
12007 let ts = commit_ts.as_ref().and_then(|v| v.get(i).copied().flatten());
12008 if !snapshot.observes_row(e, ts) {
12009 continue;
12010 }
12011 let is_del = del[i] != 0;
12012 best.entry(rid)
12013 .and_modify(|cur| {
12014 if Snapshot::version_is_newer(e, ts, cur.0, cur.1) {
12015 *cur = (e, ts, run_idx, i, is_del);
12016 }
12017 })
12018 .or_insert((e, ts, run_idx, i, is_del));
12019 }
12020 }
12021
12022 let overlay_rids: HashSet<u64> = {
12024 let mut s = HashSet::new();
12025 for row in self.memtable.visible_versions_at(snapshot) {
12026 s.insert(row.row_id.0);
12027 }
12028 for row in self.mutable_run.visible_versions_at(snapshot) {
12029 s.insert(row.row_id.0);
12030 }
12031 s
12032 };
12033
12034 let survivors: Option<RowIdSet> = if conditions.is_empty() {
12036 None
12037 } else {
12038 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12039 for c in conditions {
12040 sets.push(self.resolve_condition(c, snapshot)?);
12041 }
12042 Some(RowIdSet::intersect_many(sets))
12043 };
12044
12045 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
12049 for (rid, (_, _, run_idx, pos, deleted)) in &best {
12050 if *deleted {
12051 continue;
12052 }
12053 if overlay_rids.contains(rid) {
12054 continue;
12055 }
12056 if let Some(s) = &survivors {
12057 if !s.contains(*rid) {
12058 continue;
12059 }
12060 }
12061 per_run[*run_idx].push((*rid, *pos));
12062 }
12063 for v in per_run.iter_mut() {
12064 v.sort_unstable_by_key(|&(rid, _)| rid);
12065 }
12066
12067 let mut streams = Vec::with_capacity(run_meta.len());
12069 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
12070 let mut total = 0usize;
12071 for (run_idx, (reader, _, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
12072 let mut starts = Vec::with_capacity(page_rows.len());
12073 let mut acc = 0usize;
12074 for &r in &page_rows {
12075 starts.push(acc);
12076 acc += r;
12077 }
12078 let mut survivors_vec: Vec<(u64, usize, usize)> =
12079 Vec::with_capacity(per_run[run_idx].len());
12080 for &(rid, pos) in &per_run[run_idx] {
12081 let page_seq = match starts.partition_point(|&s| s <= pos) {
12082 0 => continue,
12083 p => p - 1,
12084 };
12085 let within = pos - starts[page_seq];
12086 survivors_vec.push((rid, page_seq, within));
12087 }
12088 total += survivors_vec.len();
12089 if let Some(&(rid, _, _)) = survivors_vec.first() {
12090 heap.push(std::cmp::Reverse((rid, run_idx)));
12091 }
12092 streams.push(RunStream::new(reader, survivors_vec, page_rows));
12093 }
12094
12095 let overlay_rows = if overlay_rids.is_empty() {
12097 Vec::new()
12098 } else {
12099 let bound = Self::overlay_materialization_bound(conditions, &survivors);
12100 self.overlay_visible_rows(snapshot, bound)
12101 };
12102 let overlay = if overlay_rows.is_empty() {
12103 None
12104 } else {
12105 let filtered =
12106 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
12107 if filtered.is_empty() {
12108 None
12109 } else {
12110 Some(self.materialize_overlay(&filtered, &projection))
12111 }
12112 };
12113
12114 let overlay_row_count = overlay
12115 .as_ref()
12116 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
12117 .unwrap_or(0);
12118 crate::trace::QueryTrace::record(|t| {
12119 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
12120 t.run_count = self.run_refs.len();
12121 t.memtable_rows = self.memtable.len();
12122 t.mutable_run_rows = self.mutable_run.len();
12123 t.overlay_rows = overlay_row_count;
12124 t.conditions_pushed = conditions.len();
12125 t.survivor_count = Some(total);
12126 });
12127
12128 Ok(Some(MultiRunCursor::new(
12129 streams, projection, heap, total, overlay,
12130 )))
12131 }
12132
12133 fn overlay_materialization_bound<'a>(
12145 conditions: &[crate::query::Condition],
12146 survivors: &'a Option<RowIdSet>,
12147 ) -> Option<&'a RowIdSet> {
12148 use crate::query::Condition;
12149 let has_range = conditions
12150 .iter()
12151 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12152 if has_range {
12153 None
12154 } else {
12155 survivors.as_ref()
12156 }
12157 }
12158
12159 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
12171 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
12172 let mut fold = |row: Row| {
12173 if let Some(b) = bound {
12174 if !b.contains(row.row_id.0) {
12175 return;
12176 }
12177 }
12178 best.entry(row.row_id.0)
12179 .and_modify(|(be, br)| {
12180 if row.committed_epoch > *be {
12181 *be = row.committed_epoch;
12182 *br = row.clone();
12183 }
12184 })
12185 .or_insert_with(|| (row.committed_epoch, row));
12186 };
12187 for row in self.memtable.visible_versions_at(snapshot) {
12188 fold(row);
12189 }
12190 for row in self.mutable_run.visible_versions_at(snapshot) {
12191 fold(row);
12192 }
12193 let mut out: Vec<Row> = best
12194 .into_values()
12195 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
12196 .collect();
12197 out.sort_by_key(|r| r.row_id);
12198 out
12199 }
12200
12201 fn filter_overlay_rows(
12209 &self,
12210 rows: Vec<Row>,
12211 conditions: &[crate::query::Condition],
12212 survivors: Option<&RowIdSet>,
12213 snapshot: Snapshot,
12214 ) -> Result<Vec<Row>> {
12215 if conditions.is_empty() {
12216 return Ok(rows);
12217 }
12218 use crate::query::Condition;
12219 let all_index_served = !conditions
12223 .iter()
12224 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
12225 if all_index_served {
12226 return Ok(rows
12227 .into_iter()
12228 .filter(|r| survivors.is_none_or(|s| s.contains(r.row_id.0)))
12229 .collect());
12230 }
12231 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
12234 for c in conditions {
12235 let s = match c {
12236 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
12237 _ => self.resolve_condition(c, snapshot)?,
12238 };
12239 per_cond_sets.push(s);
12240 }
12241 Ok(rows
12242 .into_iter()
12243 .filter(|row| {
12244 conditions.iter().enumerate().all(|(i, c)| match c {
12245 Condition::Range { column_id, lo, hi } => {
12246 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
12247 }
12248 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
12249 match row.columns.get(column_id) {
12250 Some(Value::Float64(v)) => {
12251 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
12252 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
12253 lo_ok && hi_ok
12254 }
12255 _ => false,
12256 }
12257 }
12258 _ => per_cond_sets[i].contains(row.row_id.0),
12259 })
12260 })
12261 .collect())
12262 }
12263
12264 fn materialize_overlay(
12267 &self,
12268 rows: &[Row],
12269 projection: &[(u16, TypeId)],
12270 ) -> Vec<columnar::NativeColumn> {
12271 if projection.is_empty() {
12272 return vec![columnar::null_native(TypeId::Int64, rows.len())];
12273 }
12274 let mut cols = Vec::with_capacity(projection.len());
12275 for (cid, ty) in projection {
12276 let vals: Vec<Value> = rows
12277 .iter()
12278 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
12279 .collect();
12280 cols.push(columnar::values_to_native(ty.clone(), &vals));
12281 }
12282 cols
12283 }
12284
12285 fn resolve_survivor_rids(
12290 &self,
12291 conditions: &[crate::query::Condition],
12292 reader: &mut RunReader,
12293 snapshot: Snapshot,
12294 ) -> Result<RowIdSet> {
12295 use crate::query::Condition;
12296 let mut sets: Vec<RowIdSet> = Vec::new();
12297 for c in conditions {
12298 self.validate_condition(c)?;
12299 let s: RowIdSet = match c {
12300 Condition::Pk(key) => {
12301 let lookup = self
12302 .schema
12303 .primary_key()
12304 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
12305 .unwrap_or_else(|| key.clone());
12306 self.hot
12307 .get(&lookup)
12308 .map(|r| RowIdSet::one(r.0))
12309 .unwrap_or_else(RowIdSet::empty)
12310 }
12311 Condition::BitmapEq { column_id, value } => {
12312 let lookup = self.index_lookup_key_bytes(*column_id, value);
12313 self.bitmap
12314 .get(column_id)
12315 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
12316 .unwrap_or_else(RowIdSet::empty)
12317 }
12318 Condition::BitmapIn { column_id, values } => {
12319 let bm = self.bitmap.get(column_id);
12320 let mut acc = roaring::RoaringBitmap::new();
12321 if let Some(b) = bm {
12322 for v in values {
12323 let lookup = self.index_lookup_key_bytes(*column_id, v);
12324 acc |= b.get(&lookup);
12325 }
12326 }
12327 RowIdSet::from_roaring(acc)
12328 }
12329 Condition::BytesPrefix { column_id, prefix } => {
12330 if let Some(b) = self.bitmap.get(column_id) {
12331 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
12332 let mut acc = roaring::RoaringBitmap::new();
12333 for key in b.keys() {
12334 if key.starts_with(&lookup_prefix) {
12335 acc |= b.get(&key);
12336 }
12337 }
12338 RowIdSet::from_roaring(acc)
12339 } else {
12340 RowIdSet::empty()
12341 }
12342 }
12343 Condition::FmContains { column_id, pattern } => self
12344 .fm
12345 .get(column_id)
12346 .map(|f| {
12347 RowIdSet::from_unsorted(
12348 f.locate(pattern).into_iter().map(|r| r.0).collect(),
12349 )
12350 })
12351 .unwrap_or_else(RowIdSet::empty),
12352 Condition::FmContainsAll {
12353 column_id,
12354 patterns,
12355 } => {
12356 if let Some(f) = self.fm.get(column_id) {
12357 let sets: Vec<RowIdSet> = patterns
12358 .iter()
12359 .map(|pat| {
12360 RowIdSet::from_unsorted(
12361 f.locate(pat).into_iter().map(|r| r.0).collect(),
12362 )
12363 })
12364 .collect();
12365 RowIdSet::intersect_many(sets)
12366 } else {
12367 RowIdSet::empty()
12368 }
12369 }
12370 Condition::Ann {
12371 column_id,
12372 query,
12373 k,
12374 } => RowIdSet::from_unsorted(
12375 self.retrieve_filtered(
12376 &crate::query::Retriever::Ann {
12377 column_id: *column_id,
12378 query: query.clone(),
12379 k: *k,
12380 },
12381 snapshot,
12382 None,
12383 None,
12384 None,
12385 None,
12386 )?
12387 .into_iter()
12388 .map(|hit| hit.row_id.0)
12389 .collect(),
12390 ),
12391 Condition::SparseMatch {
12392 column_id,
12393 query,
12394 k,
12395 } => RowIdSet::from_unsorted(
12396 self.retrieve_filtered(
12397 &crate::query::Retriever::Sparse {
12398 column_id: *column_id,
12399 query: query.clone(),
12400 k: *k,
12401 },
12402 snapshot,
12403 None,
12404 None,
12405 None,
12406 None,
12407 )?
12408 .into_iter()
12409 .map(|hit| hit.row_id.0)
12410 .collect(),
12411 ),
12412 Condition::MinHashSimilar {
12413 column_id,
12414 query,
12415 k,
12416 } => match self.minhash.get(column_id) {
12417 Some(index) => {
12418 let candidates = index.candidate_row_ids(query);
12419 let eligible =
12420 self.eligible_candidate_ids(&candidates, *column_id, snapshot, None)?;
12421 RowIdSet::from_unsorted(
12422 index
12423 .search_filtered(query, *k, |row_id| eligible.contains(&row_id))
12424 .into_iter()
12425 .map(|(row_id, _)| row_id.0)
12426 .collect(),
12427 )
12428 }
12429 None => RowIdSet::empty(),
12430 },
12431 Condition::Range { column_id, lo, hi } => {
12432 if let Some(li) = self.learned_range.get(column_id) {
12433 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
12434 } else {
12435 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
12436 }
12437 }
12438 Condition::RangeF64 {
12439 column_id,
12440 lo,
12441 lo_inclusive,
12442 hi,
12443 hi_inclusive,
12444 } => {
12445 if let Some(li) = self.learned_range.get(column_id) {
12446 RowIdSet::from_unsorted(
12447 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
12448 .into_iter()
12449 .collect(),
12450 )
12451 } else {
12452 reader.range_row_id_set_f64(
12453 *column_id,
12454 *lo,
12455 *lo_inclusive,
12456 *hi,
12457 *hi_inclusive,
12458 )?
12459 }
12460 }
12461 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
12462 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
12463 };
12464 sets.push(s);
12465 }
12466 Ok(RowIdSet::intersect_many(sets))
12467 }
12468
12469 pub fn scan_cursor(
12490 &self,
12491 snapshot: Snapshot,
12492 projection: Vec<(u16, TypeId)>,
12493 conditions: &[crate::query::Condition],
12494 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
12495 if self.ttl.is_some() {
12496 return Ok(None);
12497 }
12498 if !conditions.is_empty() && !self.indexes_complete {
12504 return Ok(None);
12505 }
12506 if self.run_refs.len() == 1 {
12507 Ok(self
12508 .native_page_cursor(snapshot, projection, conditions)?
12509 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
12510 } else {
12511 Ok(self
12512 .native_multi_run_cursor(snapshot, projection, conditions)?
12513 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
12514 }
12515 }
12516
12517 pub fn aggregate_native(
12531 &self,
12532 snapshot: Snapshot,
12533 column: Option<u16>,
12534 conditions: &[crate::query::Condition],
12535 agg: NativeAgg,
12536 ) -> Result<Option<NativeAggResult>> {
12537 self.aggregate_native_inner(snapshot, column, conditions, agg, None)
12538 }
12539
12540 pub fn aggregate_native_with_control(
12541 &self,
12542 snapshot: Snapshot,
12543 column: Option<u16>,
12544 conditions: &[crate::query::Condition],
12545 agg: NativeAgg,
12546 control: &crate::ExecutionControl,
12547 ) -> Result<Option<NativeAggResult>> {
12548 self.aggregate_native_inner(snapshot, column, conditions, agg, Some(control))
12549 }
12550
12551 fn aggregate_native_inner(
12552 &self,
12553 snapshot: Snapshot,
12554 column: Option<u16>,
12555 conditions: &[crate::query::Condition],
12556 agg: NativeAgg,
12557 control: Option<&crate::ExecutionControl>,
12558 ) -> Result<Option<NativeAggResult>> {
12559 execution_checkpoint(control, 0)?;
12560 if self.ttl.is_some() {
12561 return Ok(None);
12562 }
12563 if self.run_refs.len() == 1 && conditions.is_empty() {
12565 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
12566 return Ok(Some(res));
12567 }
12568 }
12569 if matches!(agg, NativeAgg::Count) && column.is_none() {
12573 if let Some(c) = self.scan_cursor(snapshot, Vec::new(), conditions)? {
12574 return Ok(Some(NativeAggResult::Count(c.remaining_rows() as u64)));
12575 }
12576 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
12577 return Ok(Some(NativeAggResult::Count(rows.len() as u64)));
12578 }
12579 let cid = match column {
12582 Some(c) => c,
12583 None => return Ok(None),
12584 };
12585 let ty = self.column_type(cid);
12586 if let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)? {
12587 execution_checkpoint(control, 0)?;
12588 return match ty {
12589 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
12590 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut(), control)?;
12591 Ok(Some(pack_int(agg, count, sum, mn, mx)))
12592 }
12593 TypeId::Float64 => {
12594 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut(), control)?;
12595 Ok(Some(pack_float(agg, count, sum, mn, mx)))
12596 }
12597 _ => Ok(None),
12598 };
12599 }
12600 let rows = self.visible_rows_filtered(snapshot, conditions, control)?;
12602 execution_checkpoint(control, 0)?;
12603 match ty {
12604 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
12605 let mut count = 0u64;
12606 let mut sum = 0i128;
12607 let mut mn = i64::MAX;
12608 let mut mx = i64::MIN;
12609 for row in &rows {
12610 if let Some(Value::Int64(v)) = row.columns.get(&cid) {
12611 count += 1;
12612 sum += i128::from(*v);
12613 mn = mn.min(*v);
12614 mx = mx.max(*v);
12615 }
12616 }
12617 Ok(Some(pack_int(agg, count, sum, mn, mx)))
12618 }
12619 TypeId::Float64 => {
12620 let mut count = 0u64;
12621 let mut sum = 0.0f64;
12622 let mut mn = f64::INFINITY;
12623 let mut mx = f64::NEG_INFINITY;
12624 for row in &rows {
12625 if let Some(Value::Float64(v)) = row.columns.get(&cid) {
12626 count += 1;
12627 sum += *v;
12628 mn = mn.min(*v);
12629 mx = mx.max(*v);
12630 }
12631 }
12632 Ok(Some(pack_float(agg, count, sum, mn, mx)))
12633 }
12634 _ => Ok(None),
12635 }
12636 }
12637
12638 fn visible_rows_filtered(
12640 &self,
12641 snapshot: Snapshot,
12642 conditions: &[crate::query::Condition],
12643 control: Option<&crate::ExecutionControl>,
12644 ) -> Result<Vec<Row>> {
12645 let rows = if let Some(control) = control {
12646 self.visible_rows_controlled(snapshot, control)?
12647 } else {
12648 self.visible_rows(snapshot)?
12649 };
12650 if conditions.is_empty() {
12651 return Ok(rows);
12652 }
12653 Ok(rows
12654 .into_iter()
12655 .filter(|row| {
12656 conditions
12657 .iter()
12658 .all(|cond| condition_matches_row(cond, row, &self.schema))
12659 })
12660 .collect())
12661 }
12662
12663 fn aggregate_from_stats(
12671 &self,
12672 snapshot: Snapshot,
12673 column: Option<u16>,
12674 agg: NativeAgg,
12675 ) -> Result<Option<NativeAggResult>> {
12676 let cid = match (agg, column) {
12677 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
12678 _ => return Ok(None), };
12680 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
12681 return Ok(None);
12682 };
12683 let Some(cs) = stats.get(&cid) else {
12684 return Ok(None);
12685 };
12686 match agg {
12687 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
12689 self.live_count.saturating_sub(cs.null_count),
12690 ))),
12691 NativeAgg::Min | NativeAgg::Max => {
12692 let bound = if agg == NativeAgg::Min {
12693 &cs.min
12694 } else {
12695 &cs.max
12696 };
12697 match bound {
12698 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
12699 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
12700 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
12705 None => Ok(None),
12706 }
12707 }
12708 _ => Ok(None),
12709 }
12710 }
12711
12712 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
12721 if self.ttl.is_some() {
12722 return Ok(None);
12723 }
12724 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
12725 return Ok(None);
12726 }
12727 self.ensure_indexes_complete()?;
12730 let reader = self.open_reader(self.run_refs[0].run_id)?;
12731 if self.live_count != reader.row_count() as u64 {
12732 return Ok(None);
12733 }
12734 let Some(bm) = self.bitmap.get(&column_id) else {
12735 return Ok(None); };
12737 let mut distinct = bm.value_count() as u64;
12738 if !bm.get(&Value::Null.encode_key()).is_empty() {
12741 distinct = distinct.saturating_sub(1);
12742 }
12743 Ok(Some(distinct))
12744 }
12745
12746 pub fn aggregate_incremental(
12758 &mut self,
12759 cache_key: u64,
12760 conditions: &[crate::query::Condition],
12761 column: Option<u16>,
12762 agg: NativeAgg,
12763 ) -> Result<IncrementalAggResult> {
12764 self.aggregate_incremental_inner(cache_key, conditions, column, agg, None)
12765 }
12766
12767 pub fn aggregate_incremental_with_control(
12768 &mut self,
12769 cache_key: u64,
12770 conditions: &[crate::query::Condition],
12771 column: Option<u16>,
12772 agg: NativeAgg,
12773 control: &crate::ExecutionControl,
12774 ) -> Result<IncrementalAggResult> {
12775 self.aggregate_incremental_inner(cache_key, conditions, column, agg, Some(control))
12776 }
12777
12778 fn aggregate_incremental_inner(
12779 &mut self,
12780 cache_key: u64,
12781 conditions: &[crate::query::Condition],
12782 column: Option<u16>,
12783 agg: NativeAgg,
12784 control: Option<&crate::ExecutionControl>,
12785 ) -> Result<IncrementalAggResult> {
12786 execution_checkpoint(control, 0)?;
12787 let snap = self.snapshot();
12788 let cur_wm = self.allocator.current().0;
12789 let cur_epoch = snap.epoch.0;
12790 let incremental_ok = self.ttl.is_none()
12797 && !self.had_deletes
12798 && self.memtable.is_empty()
12799 && self.mutable_run.is_empty();
12800
12801 if incremental_ok {
12804 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
12805 if cached.epoch == cur_epoch {
12806 return Ok(IncrementalAggResult {
12807 state: cached.state,
12808 incremental: true,
12809 delta_rows: 0,
12810 });
12811 }
12812 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
12813 let delta_len = cur_wm.saturating_sub(cached.watermark) as usize;
12814 let mut delta_rids = Vec::with_capacity(delta_len);
12815 for (index, row_id) in (cached.watermark..cur_wm).enumerate() {
12816 execution_checkpoint(control, index)?;
12817 delta_rids.push(row_id);
12818 }
12819 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
12820 execution_checkpoint(control, 0)?;
12821 let index_sets = self.resolve_index_conditions(conditions, snap)?;
12822 let delta_state = agg_state_from_rows(
12823 &delta_rows,
12824 conditions,
12825 &index_sets,
12826 column,
12827 agg,
12828 &self.schema,
12829 control,
12830 )?;
12831 let merged = cached.state.merge(delta_state);
12832 let delta_n = delta_rids.len() as u64;
12833 Arc::make_mut(&mut self.agg_cache).insert(
12834 cache_key,
12835 CachedAgg {
12836 state: merged.clone(),
12837 watermark: cur_wm,
12838 epoch: cur_epoch,
12839 },
12840 );
12841 return Ok(IncrementalAggResult {
12842 state: merged,
12843 incremental: true,
12844 delta_rows: delta_n,
12845 });
12846 }
12847 }
12848 }
12849
12850 let cursor_ok =
12855 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
12856 let state = if cursor_ok && agg != NativeAgg::Avg {
12857 match self.aggregate_native_inner(snap, column, conditions, agg, control)? {
12858 Some(result) => {
12859 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
12860 }
12861 None => self.agg_state_full_scan(conditions, column, agg, snap, control)?,
12862 }
12863 } else {
12864 self.agg_state_full_scan(conditions, column, agg, snap, control)?
12865 };
12866 if incremental_ok {
12868 Arc::make_mut(&mut self.agg_cache).insert(
12869 cache_key,
12870 CachedAgg {
12871 state: state.clone(),
12872 watermark: cur_wm,
12873 epoch: cur_epoch,
12874 },
12875 );
12876 }
12877 Ok(IncrementalAggResult {
12878 state,
12879 incremental: false,
12880 delta_rows: 0,
12881 })
12882 }
12883
12884 fn agg_state_full_scan(
12887 &self,
12888 conditions: &[crate::query::Condition],
12889 column: Option<u16>,
12890 agg: NativeAgg,
12891 snap: Snapshot,
12892 control: Option<&crate::ExecutionControl>,
12893 ) -> Result<AggState> {
12894 execution_checkpoint(control, 0)?;
12895 let rows = self.visible_rows(snap)?;
12896 execution_checkpoint(control, 0)?;
12897 let index_sets = self.resolve_index_conditions(conditions, snap)?;
12898 agg_state_from_rows(
12899 &rows,
12900 conditions,
12901 &index_sets,
12902 column,
12903 agg,
12904 &self.schema,
12905 control,
12906 )
12907 }
12908
12909 fn resolve_index_conditions(
12912 &self,
12913 conditions: &[crate::query::Condition],
12914 snapshot: Snapshot,
12915 ) -> Result<Vec<RowIdSet>> {
12916 use crate::query::Condition;
12917 let mut sets = Vec::new();
12918 for c in conditions {
12919 if matches!(
12920 c,
12921 Condition::Ann { .. }
12922 | Condition::SparseMatch { .. }
12923 | Condition::MinHashSimilar { .. }
12924 ) {
12925 sets.push(self.resolve_condition(c, snapshot)?);
12926 }
12927 }
12928 Ok(sets)
12929 }
12930
12931 fn column_type(&self, cid: u16) -> TypeId {
12932 self.schema
12933 .columns
12934 .iter()
12935 .find(|c| c.id == cid)
12936 .map(|c| c.ty.clone())
12937 .unwrap_or(TypeId::Bytes)
12938 }
12939
12940 pub fn approx_aggregate(
12949 &mut self,
12950 conditions: &[crate::query::Condition],
12951 column: Option<u16>,
12952 agg: ApproxAgg,
12953 z: f64,
12954 ) -> Result<Option<ApproxResult>> {
12955 self.approx_aggregate_with_candidate_authorization(conditions, column, agg, z, None)
12956 }
12957
12958 pub fn approx_aggregate_with_candidate_authorization(
12961 &mut self,
12962 conditions: &[crate::query::Condition],
12963 column: Option<u16>,
12964 agg: ApproxAgg,
12965 z: f64,
12966 authorization: Option<&crate::security::CandidateAuthorization<'_>>,
12967 ) -> Result<Option<ApproxResult>> {
12968 use crate::query::Condition;
12969 self.ensure_reservoir_complete()?;
12970 let mut snapshot = self.snapshot();
12975 let n_pop = self.count();
12976 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
12977 if sample_rids.is_empty() {
12978 return Ok(None);
12979 }
12980 let mut live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
12982 if live_sample.is_empty() {
12983 snapshot = Snapshot::unbounded();
12984 live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
12985 }
12986 let s = live_sample.len();
12987 if s == 0 {
12988 return Ok(None);
12989 }
12990 let authorized = authorization
12991 .map(|authorization| {
12992 let candidates = live_sample.iter().map(|row| row.row_id).collect::<Vec<_>>();
12993 self.policy_allowed_candidate_ids(&candidates, snapshot, authorization, None)
12994 })
12995 .transpose()?;
12996
12997 let mut index_sets: Vec<RowIdSet> = Vec::new();
13000 for c in conditions {
13001 if matches!(
13002 c,
13003 Condition::Ann { .. }
13004 | Condition::SparseMatch { .. }
13005 | Condition::MinHashSimilar { .. }
13006 ) {
13007 index_sets.push(self.resolve_condition(c, snapshot)?);
13008 }
13009 }
13010
13011 let cid = match (agg, column) {
13013 (ApproxAgg::Count, _) => None,
13014 (_, Some(c)) => Some(c),
13015 _ => return Ok(None),
13016 };
13017 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
13018 for r in &live_sample {
13019 if authorized
13020 .as_ref()
13021 .is_some_and(|authorized| !authorized.contains(&r.row_id))
13022 {
13023 continue;
13024 }
13025 if !conditions
13027 .iter()
13028 .all(|c| condition_matches_row(c, r, &self.schema))
13029 {
13030 continue;
13031 }
13032 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
13034 continue;
13035 }
13036 if let Some(cid) = cid {
13037 let mut cells = r
13038 .columns
13039 .get(&cid)
13040 .cloned()
13041 .map(|value| vec![(cid, value)])
13042 .unwrap_or_default();
13043 if let Some(authorization) = authorization {
13044 authorization.security.apply_masks_to_cells(
13045 authorization.table,
13046 &mut cells,
13047 authorization.principal,
13048 );
13049 }
13050 if let Some(v) = as_f64(cells.first().map(|(_, value)| value)) {
13051 passing_vals.push(v);
13052 } } else {
13054 passing_vals.push(0.0); }
13056 }
13057 let m = passing_vals.len();
13058
13059 let (point, half) = match agg {
13060 ApproxAgg::Count => {
13061 let p = m as f64 / s as f64;
13063 let point = n_pop as f64 * p;
13064 let var = if s > 1 {
13065 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
13066 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13067 } else {
13068 0.0
13069 };
13070 (point, z * var.sqrt())
13071 }
13072 ApproxAgg::Sum => {
13073 let y: Vec<f64> = live_sample
13075 .iter()
13076 .map(|r| {
13077 let passes_row = authorized
13078 .as_ref()
13079 .is_none_or(|authorized| authorized.contains(&r.row_id))
13080 && conditions
13081 .iter()
13082 .all(|c| condition_matches_row(c, r, &self.schema))
13083 && index_sets.iter().all(|set| set.contains(r.row_id.0));
13084 if passes_row {
13085 cid.and_then(|cid| {
13086 let mut cells = r
13087 .columns
13088 .get(&cid)
13089 .cloned()
13090 .map(|value| vec![(cid, value)])
13091 .unwrap_or_default();
13092 if let Some(authorization) = authorization {
13093 authorization.security.apply_masks_to_cells(
13094 authorization.table,
13095 &mut cells,
13096 authorization.principal,
13097 );
13098 }
13099 as_f64(cells.first().map(|(_, value)| value))
13100 })
13101 .unwrap_or(0.0)
13102 } else {
13103 0.0
13104 }
13105 })
13106 .collect();
13107 let mean_y = y.iter().sum::<f64>() / s as f64;
13108 let point = n_pop as f64 * mean_y;
13109 let var = if s > 1 {
13110 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
13111 let var_y = ss / (s - 1) as f64;
13112 n_pop as f64 * n_pop as f64 * var_y / s as f64
13113 * (1.0 - s as f64 / n_pop as f64).max(0.0)
13114 } else {
13115 0.0
13116 };
13117 (point, z * var.sqrt())
13118 }
13119 ApproxAgg::Avg => {
13120 if m == 0 {
13121 return Ok(Some(ApproxResult {
13122 point: 0.0,
13123 ci_low: 0.0,
13124 ci_high: 0.0,
13125 n_population: n_pop,
13126 n_sample_live: s,
13127 n_passing: 0,
13128 }));
13129 }
13130 let mean = passing_vals.iter().sum::<f64>() / m as f64;
13131 let half = if m > 1 {
13132 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
13133 let sd = (ss / (m - 1) as f64).sqrt();
13134 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
13135 z * sd / (m as f64).sqrt() * fpc.sqrt()
13136 } else {
13137 0.0
13138 };
13139 (mean, half)
13140 }
13141 };
13142
13143 Ok(Some(ApproxResult {
13144 point,
13145 ci_low: point - half,
13146 ci_high: point + half,
13147 n_population: n_pop,
13148 n_sample_live: s,
13149 n_passing: m,
13150 }))
13151 }
13152
13153 pub fn exact_column_stats(
13161 &self,
13162 _snapshot: Snapshot,
13163 projection: &[u16],
13164 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
13165 if self.ttl.is_some()
13166 || !(self.memtable.is_empty()
13167 && self.mutable_run.is_empty()
13168 && self.run_refs.len() == 1)
13169 {
13170 return Ok(None);
13171 }
13172 let reader = self.open_reader(self.run_refs[0].run_id)?;
13173 if self.live_count != reader.row_count() as u64 {
13174 return Ok(None);
13175 }
13176 let mut out = HashMap::new();
13177 for &cid in projection {
13178 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
13179 Some(c) => c,
13180 None => continue,
13181 };
13182 let Some(stats) = reader.column_page_stats(cid) else {
13184 out.insert(
13185 cid,
13186 ColumnStat {
13187 min: None,
13188 max: None,
13189 null_count: self.live_count,
13190 },
13191 );
13192 continue;
13193 };
13194 let stat = match cdef.ty {
13195 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
13196 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
13197 min: mn.map(Value::Int64),
13198 max: mx.map(Value::Int64),
13199 null_count: n,
13200 })
13201 }
13202 TypeId::Float64 => {
13203 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
13204 min: mn.map(Value::Float64),
13205 max: mx.map(Value::Float64),
13206 null_count: n,
13207 })
13208 }
13209 _ => None,
13210 };
13211 if let Some(s) = stat {
13212 out.insert(cid, s);
13213 }
13214 }
13215 Ok(Some(out))
13216 }
13217
13218 pub fn dir(&self) -> &Path {
13219 &self.dir
13220 }
13221
13222 pub fn schema(&self) -> &Schema {
13223 &self.schema
13224 }
13225
13226 pub fn ann_index(&self, column_id: u16) -> Option<&crate::index::AnnIndex> {
13227 self.ann.get(&column_id)
13228 }
13229
13230 pub fn ann_index_mut(&mut self, column_id: u16) -> Option<&mut crate::index::AnnIndex> {
13231 self.ann.get_mut(&column_id)
13232 }
13233
13234 pub(crate) fn set_catalog_name(&mut self, name: String) {
13235 self.name = name;
13236 }
13237
13238 pub(crate) fn prepare_alter_column(
13239 &mut self,
13240 column_name: &str,
13241 change: &AlterColumn,
13242 ) -> Result<(ColumnDef, Option<Schema>)> {
13243 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
13244 return Err(MongrelError::InvalidArgument(
13245 "ALTER COLUMN requires committing staged writes first".into(),
13246 ));
13247 }
13248 let old = self
13249 .schema
13250 .columns
13251 .iter()
13252 .find(|c| c.name == column_name)
13253 .cloned()
13254 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
13255 let mut next = old.clone();
13256
13257 if let Some(name) = &change.name {
13258 let trimmed = name.trim();
13259 if trimmed.is_empty() {
13260 return Err(MongrelError::InvalidArgument(
13261 "ALTER COLUMN name must not be empty".into(),
13262 ));
13263 }
13264 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
13265 return Err(MongrelError::Schema(format!(
13266 "column {trimmed} already exists"
13267 )));
13268 }
13269 next.name = trimmed.to_string();
13270 }
13271
13272 if let Some(ty) = &change.ty {
13273 next.ty = ty.clone();
13274 }
13275 if let Some(flags) = change.flags {
13276 validate_alter_column_flags(old.flags, flags)?;
13277 next.flags = flags;
13278 }
13279
13280 if let Some(default_change) = &change.default_value {
13281 next.default_value = default_change.clone();
13282 }
13283 if let Some(source_change) = &change.embedding_source {
13284 next.embedding_source = source_change.clone();
13285 }
13286
13287 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
13288 if old.flags.contains(ColumnFlags::NULLABLE)
13289 && !next.flags.contains(ColumnFlags::NULLABLE)
13290 && self.column_has_nulls(old.id)?
13291 {
13292 return Err(MongrelError::InvalidArgument(format!(
13293 "column '{}' contains NULL values",
13294 old.name
13295 )));
13296 }
13297 if next == old {
13298 return Ok((next, None));
13299 }
13300 let mut schema = self.schema.clone();
13301 let index = schema
13302 .columns
13303 .iter()
13304 .position(|column| column.id == next.id)
13305 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", next.id)))?;
13306 schema.columns[index] = next.clone();
13307 schema.schema_id = schema
13308 .schema_id
13309 .checked_add(1)
13310 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13311 schema.validate_auto_increment()?;
13312 schema.validate_defaults()?;
13313 Ok((next, Some(schema)))
13314 }
13315
13316 pub(crate) fn apply_altered_schema_prepared(&mut self, schema: Schema) {
13317 self.schema = schema;
13318 self.auto_inc = resolve_auto_inc(&self.schema);
13319 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
13320 self.clear_result_cache();
13321 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
13322 }
13323
13324 pub(crate) fn publish_index_schema_change(
13328 &mut self,
13329 schema: Schema,
13330 artifact: SecondaryIndexArtifact,
13331 ) -> Result<()> {
13332 self.apply_altered_schema_prepared(schema);
13333 self.prune_index_maps_to_schema();
13334 match artifact {
13335 SecondaryIndexArtifact::Bitmap(column_id, index) => {
13336 self.bitmap.insert(column_id, index);
13337 }
13338 SecondaryIndexArtifact::LearnedRange(column_id, index) => {
13339 Arc::make_mut(&mut self.learned_range).insert(column_id, index);
13340 }
13341 SecondaryIndexArtifact::Fm(column_id, index) => {
13342 self.fm.insert(column_id, *index);
13343 }
13344 SecondaryIndexArtifact::Ann(column_id, index) => {
13345 self.ann.insert(column_id, *index);
13346 }
13347 SecondaryIndexArtifact::Sparse(column_id, index) => {
13348 self.sparse.insert(column_id, index);
13349 }
13350 SecondaryIndexArtifact::MinHash(column_id, index) => {
13351 self.minhash.insert(column_id, index);
13352 }
13353 }
13354 self.indexes_complete = true;
13355 self.seal_generations();
13356 let view = Arc::new(self.capture_read_generation());
13357 self.published.store(Arc::clone(&view));
13358 checkpoint_current_schema(self)?;
13359 self.invalidate_index_checkpoint();
13362 Ok(())
13363 }
13364
13365 pub(crate) fn publish_index_drop(&mut self, schema: Schema) -> Result<()> {
13371 self.apply_altered_schema_prepared(schema);
13372 self.prune_index_maps_to_schema();
13373 self.indexes_complete = true;
13374 self.seal_generations();
13375 let view = Arc::new(self.capture_read_generation());
13376 self.published.store(Arc::clone(&view));
13377 checkpoint_current_schema(self)?;
13378 self.invalidate_index_checkpoint();
13380 Ok(())
13381 }
13382
13383 fn prune_index_maps_to_schema(&mut self) {
13384 let keep_bitmap: std::collections::HashSet<u16> = self
13385 .schema
13386 .indexes
13387 .iter()
13388 .filter(|index| index.kind == IndexKind::Bitmap)
13389 .map(|index| index.column_id)
13390 .collect();
13391 let keep_ann: std::collections::HashSet<u16> = self
13392 .schema
13393 .indexes
13394 .iter()
13395 .filter(|index| index.kind == IndexKind::Ann)
13396 .map(|index| index.column_id)
13397 .collect();
13398 let keep_fm: std::collections::HashSet<u16> = self
13399 .schema
13400 .indexes
13401 .iter()
13402 .filter(|index| index.kind == IndexKind::FmIndex)
13403 .map(|index| index.column_id)
13404 .collect();
13405 let keep_sparse: std::collections::HashSet<u16> = self
13406 .schema
13407 .indexes
13408 .iter()
13409 .filter(|index| index.kind == IndexKind::Sparse)
13410 .map(|index| index.column_id)
13411 .collect();
13412 let keep_minhash: std::collections::HashSet<u16> = self
13413 .schema
13414 .indexes
13415 .iter()
13416 .filter(|index| index.kind == IndexKind::MinHash)
13417 .map(|index| index.column_id)
13418 .collect();
13419 let keep_learned: std::collections::HashSet<u16> = self
13420 .schema
13421 .indexes
13422 .iter()
13423 .filter(|index| index.kind == IndexKind::LearnedRange)
13424 .map(|index| index.column_id)
13425 .collect();
13426 self.bitmap
13427 .retain(|column_id, _| keep_bitmap.contains(column_id));
13428 self.ann.retain(|column_id, _| keep_ann.contains(column_id));
13429 self.fm.retain(|column_id, _| keep_fm.contains(column_id));
13430 self.sparse
13431 .retain(|column_id, _| keep_sparse.contains(column_id));
13432 self.minhash
13433 .retain(|column_id, _| keep_minhash.contains(column_id));
13434 {
13435 let learned = Arc::make_mut(&mut self.learned_range);
13436 learned.retain(|column_id, _| keep_learned.contains(column_id));
13437 }
13438 }
13439
13440 pub(crate) fn checkpoint_altered_schema(&mut self) -> Result<()> {
13441 checkpoint_current_schema(self)
13442 }
13443
13444 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
13445 self.ensure_writable()?;
13446 let previous_schema = self.schema.clone();
13447 let (column, schema) = self.prepare_alter_column(column_name, &change)?;
13448 if let Some(schema) = schema {
13449 self.apply_altered_schema_prepared(schema);
13450 self.checkpoint_standalone_schema_change(previous_schema)?;
13451 }
13452 Ok(column)
13453 }
13454
13455 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
13456 if self.live_count == 0 {
13457 return Ok(false);
13458 }
13459 let snap = self.snapshot();
13460 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
13461 Ok(columns
13462 .first()
13463 .map(|(_, col)| col.null_count(col.len()) != 0)
13464 .unwrap_or(true))
13465 }
13466
13467 fn has_stored_versions(&self) -> bool {
13468 !self.memtable.is_empty()
13469 || !self.mutable_run.is_empty()
13470 || self.run_refs.iter().any(|r| r.row_count > 0)
13471 || !self.retiring.is_empty()
13472 }
13473
13474 pub fn add_column(
13479 &mut self,
13480 name: &str,
13481 ty: TypeId,
13482 flags: ColumnFlags,
13483 default_value: Option<crate::schema::DefaultExpr>,
13484 ) -> Result<u16> {
13485 self.add_column_with_id(name, ty, flags, default_value, None)
13486 }
13487
13488 pub fn add_column_with_id(
13489 &mut self,
13490 name: &str,
13491 ty: TypeId,
13492 flags: ColumnFlags,
13493 default_value: Option<crate::schema::DefaultExpr>,
13494 requested_id: Option<u16>,
13495 ) -> Result<u16> {
13496 self.ensure_writable()?;
13497 let previous_schema = self.schema.clone();
13498 let (column, schema) =
13499 self.prepare_add_column(name, ty, flags, default_value, requested_id)?;
13500 self.apply_altered_schema_prepared(schema);
13501 self.checkpoint_standalone_schema_change(previous_schema)?;
13502 Ok(column.id)
13503 }
13504
13505 pub(crate) fn prepare_add_column(
13506 &mut self,
13507 name: &str,
13508 ty: TypeId,
13509 flags: ColumnFlags,
13510 default_value: Option<crate::schema::DefaultExpr>,
13511 requested_id: Option<u16>,
13512 ) -> Result<(ColumnDef, Schema)> {
13513 self.ensure_writable()?;
13514 if self.schema.columns.iter().any(|c| c.name == name) {
13515 return Err(MongrelError::Schema(format!(
13516 "column {name} already exists"
13517 )));
13518 }
13519 let id = if let Some(id) = requested_id.filter(|id| *id != 0) {
13520 if self.schema.columns.iter().any(|c| c.id == id) {
13521 return Err(MongrelError::Schema(format!(
13522 "column id {id} already exists"
13523 )));
13524 }
13525 id
13526 } else {
13527 self.schema
13528 .columns
13529 .iter()
13530 .map(|c| c.id)
13531 .max()
13532 .unwrap_or(0)
13533 .checked_add(1)
13534 .ok_or_else(|| MongrelError::Schema("column id space exhausted".into()))?
13535 };
13536 let column = ColumnDef {
13537 id,
13538 name: name.to_string(),
13539 ty,
13540 flags,
13541 default_value,
13542 embedding_source: None,
13543 };
13544 let mut next_schema = self.schema.clone();
13545 next_schema.columns.push(column.clone());
13546 next_schema.schema_id = next_schema
13547 .schema_id
13548 .checked_add(1)
13549 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13550 next_schema.validate_auto_increment()?;
13551 next_schema.validate_defaults()?;
13552 Ok((column, next_schema))
13553 }
13554
13555 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
13564 self.ensure_writable()?;
13565 let cid = self
13566 .schema
13567 .columns
13568 .iter()
13569 .find(|c| c.name == column_name)
13570 .map(|c| c.id)
13571 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
13572 let ty = self
13573 .schema
13574 .columns
13575 .iter()
13576 .find(|c| c.id == cid)
13577 .map(|c| c.ty.clone())
13578 .unwrap_or(TypeId::Int64);
13579 if !matches!(
13580 ty,
13581 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
13582 ) {
13583 return Err(MongrelError::Schema(format!(
13584 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
13585 )));
13586 }
13587 if self
13588 .schema
13589 .indexes
13590 .iter()
13591 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
13592 {
13593 return Ok(()); }
13595 let previous_schema = self.schema.clone();
13596 let previous_learned_range = Arc::clone(&self.learned_range);
13597 let mut next_schema = previous_schema.clone();
13598 next_schema.indexes.push(IndexDef {
13599 name: format!("{}_learned_range", column_name),
13600 column_id: cid,
13601 kind: IndexKind::LearnedRange,
13602 predicate: None,
13603 options: Default::default(),
13604 });
13605 next_schema.schema_id = next_schema
13606 .schema_id
13607 .checked_add(1)
13608 .ok_or_else(|| MongrelError::Schema("schema id space exhausted".into()))?;
13609 self.apply_altered_schema_prepared(next_schema);
13610 if let Err(error) = self.build_learned_ranges() {
13611 self.apply_altered_schema_prepared(previous_schema);
13612 self.learned_range = previous_learned_range;
13613 return Err(error);
13614 }
13615 if let Err(error) = self.checkpoint_standalone_schema_change(previous_schema) {
13616 if !matches!(
13617 &error,
13618 MongrelError::DurableCommit { .. } | MongrelError::CommitOutcomeUnknown { .. }
13619 ) {
13620 self.learned_range = previous_learned_range;
13621 }
13622 return Err(error);
13623 }
13624 Ok(())
13625 }
13626
13627 fn checkpoint_standalone_schema_change(&mut self, previous_schema: Schema) -> Result<()> {
13628 let mut schema_published = false;
13629 let schema_result = match self._root_guard.as_deref() {
13630 Some(root) => write_schema_durable_with_after(root, &self.schema, || {
13631 schema_published = true;
13632 }),
13633 None => write_schema_with_after(&self.dir, &self.schema, || {
13634 schema_published = true;
13635 }),
13636 };
13637 if schema_result.is_err() && !schema_published {
13638 self.apply_altered_schema_prepared(previous_schema);
13639 return schema_result;
13640 }
13641
13642 let manifest_result = self.persist_manifest(self.current_epoch());
13643 match (schema_result, manifest_result) {
13644 (_, Ok(())) => Ok(()),
13645 (Ok(()), Err(error)) => {
13646 self.poison_after_maintenance_publish_failure();
13647 Err(MongrelError::DurableCommit {
13648 epoch: self.current_epoch().0,
13649 message: format!(
13650 "schema is durable but matching manifest publication failed: {error}"
13651 ),
13652 })
13653 }
13654 (Err(schema_error), Err(manifest_error)) => {
13655 self.poison_after_maintenance_publish_failure();
13656 Err(MongrelError::CommitOutcomeUnknown {
13657 epoch: self.current_epoch().0,
13658 message: format!(
13659 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
13660 ),
13661 })
13662 }
13663 }
13664 }
13665
13666 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
13669 self.sync_byte_threshold = threshold;
13670 if let WalSink::Private(w) = &mut self.wal {
13671 w.set_sync_byte_threshold(threshold);
13672 }
13673 }
13674
13675 pub fn page_cache_flush(&self) {
13679 self.page_cache.flush_to_disk();
13680 }
13681
13682 pub fn page_cache_len(&self) -> usize {
13684 self.page_cache.len()
13685 }
13686
13687 pub fn decoded_cache_len(&self) -> usize {
13690 self.decoded_cache.len()
13691 }
13692
13693 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
13696 self.memtable.drain_sorted()
13697 }
13698
13699 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
13700 self.runs_dir().join(format!("r-{run_id}.sr"))
13701 }
13702
13703 pub(crate) fn create_run_file(&self, run_id: u64) -> Result<Option<std::fs::File>> {
13704 match self.runs_root.as_deref() {
13705 Some(root) => Ok(Some(root.create_regular_new(format!("r-{run_id}.sr"))?)),
13706 None => Ok(None),
13707 }
13708 }
13709
13710 pub(crate) fn create_run_entry(&self, name: &Path) -> Result<Option<std::fs::File>> {
13711 match self.runs_root.as_deref() {
13712 Some(root) => Ok(Some(root.create_regular_new(name)?)),
13713 None => Ok(None),
13714 }
13715 }
13716
13717 pub(crate) fn remove_run_entry(&self, name: &Path) -> Result<()> {
13718 match self.runs_root.as_deref() {
13719 Some(root) => match root.remove_file(name) {
13720 Ok(()) => Ok(()),
13721 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
13722 Err(error) => Err(error.into()),
13723 },
13724 None => match std::fs::remove_file(self.runs_dir().join(name)) {
13725 Ok(()) => Ok(()),
13726 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
13727 Err(error) => Err(error.into()),
13728 },
13729 }
13730 }
13731
13732 pub(crate) fn publish_run_entry(&self, source: &Path, destination: &Path) -> Result<()> {
13733 match self.runs_root.as_deref() {
13734 Some(root) => root
13735 .rename_file_new(source, destination)
13736 .map_err(Into::into),
13737 None => crate::durable_file::rename(
13738 &self.runs_dir().join(source),
13739 &self.runs_dir().join(destination),
13740 )
13741 .map_err(Into::into),
13742 }
13743 }
13744
13745 pub(crate) fn active_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
13746 self.run_refs.iter().map(|run| run.run_id)
13747 }
13748
13749 pub(crate) fn table_dir(&self) -> &Path {
13750 &self.dir
13751 }
13752
13753 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
13754 &self.schema
13755 }
13756
13757 pub(crate) fn alloc_run_id(&mut self) -> Result<u64> {
13758 let id = self.next_run_id;
13759 self.next_run_id = self
13760 .next_run_id
13761 .checked_add(1)
13762 .ok_or_else(|| MongrelError::Full("run-id namespace exhausted".into()))?;
13763 Ok(id)
13764 }
13765
13766 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
13767 self.run_refs.push(run_ref);
13768 }
13769
13770 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
13780 self.retiring.push(crate::manifest::RetiredRun {
13781 run_id,
13782 retire_epoch,
13783 });
13784 }
13785
13786 pub(crate) fn reap_retiring(
13790 &mut self,
13791 min_active: Epoch,
13792 backup_pinned: &std::collections::HashSet<u128>,
13793 ) -> Result<usize> {
13794 if self.retiring.is_empty() {
13795 return Ok(0);
13796 }
13797 let mut reaped = 0;
13798 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
13799 for r in std::mem::take(&mut self.retiring) {
13805 if min_active.0 >= r.retire_epoch && !backup_pinned.contains(&r.run_id) {
13806 let _ = self.remove_run_entry(Path::new(&format!("r-{}.sr", r.run_id)));
13807 reaped += 1;
13808 } else {
13809 kept.push(r);
13810 }
13811 }
13812 self.retiring = kept;
13813 if reaped > 0 {
13814 self.persist_manifest(self.current_epoch())?;
13815 }
13816 Ok(reaped)
13817 }
13818
13819 pub(crate) fn has_reapable_retiring(
13820 &self,
13821 min_active: Epoch,
13822 backup_pinned: &std::collections::HashSet<u128>,
13823 ) -> bool {
13824 self.retiring
13825 .iter()
13826 .any(|run| min_active.0 >= run.retire_epoch && !backup_pinned.contains(&run.run_id))
13827 }
13828
13829 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
13830 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
13831 return false;
13832 }
13833 self.live_count = self.live_count.saturating_add(run_ref.row_count);
13834 self.run_refs.push(run_ref);
13835 self.indexes_complete = false;
13836 true
13837 }
13838
13839 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
13840 self.kek.as_ref()
13841 }
13842
13843 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
13844 let mut reader = match self.runs_root.as_deref() {
13845 Some(root) => RunReader::open_file_with_cache(
13846 root.open_regular(format!("r-{run_id}.sr"))?,
13847 self.schema.clone(),
13848 self.kek.clone(),
13849 Some(self.page_cache.clone()),
13850 Some(self.decoded_cache.clone()),
13851 self.table_id,
13852 Some(&self.verified_runs),
13853 None,
13854 )?,
13855 None => RunReader::open_with_cache(
13856 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
13857 self.schema.clone(),
13858 self.kek.clone(),
13859 Some(self.page_cache.clone()),
13860 Some(self.decoded_cache.clone()),
13861 self.table_id,
13862 Some(&self.verified_runs),
13863 )?,
13864 };
13865 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
13869 reader.set_uniform_epoch(Epoch(rr.epoch_created));
13870 }
13871 Ok(reader)
13872 }
13873
13874 pub(crate) fn run_refs(&self) -> &[RunRef] {
13875 &self.run_refs
13876 }
13877
13878 pub(crate) fn retiring_run_ids(&self) -> impl Iterator<Item = u128> + '_ {
13879 self.retiring.iter().map(|run| run.run_id)
13880 }
13881
13882 pub(crate) fn runs_dir(&self) -> PathBuf {
13883 self.runs_root
13884 .as_deref()
13885 .and_then(|root| root.io_path().ok())
13886 .unwrap_or_else(|| self.dir.join(RUNS_DIR))
13887 }
13888
13889 pub(crate) fn wal_dir(&self) -> PathBuf {
13890 self.dir.join(WAL_DIR)
13891 }
13892
13893 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
13894 self.run_refs = refs;
13895 self.refresh_run_row_id_ranges();
13896 }
13897
13898 pub(crate) fn run_row_id_range(&self, run_id: u128) -> Option<(u64, u64)> {
13902 self.run_row_id_ranges.get(&run_id).copied()
13903 }
13904
13905 pub(crate) fn refresh_run_row_id_ranges(&mut self) {
13908 let mut ranges = HashMap::with_capacity(self.run_refs.len());
13909 for rr in &self.run_refs {
13910 if let Ok(reader) = self.open_reader(rr.run_id) {
13911 let h = reader.header();
13912 ranges.insert(rr.run_id, (h.min_row_id, h.max_row_id));
13913 }
13914 }
13915 self.run_row_id_ranges = ranges;
13916 }
13917
13918 pub(crate) fn compaction_zstd_level(&self) -> i32 {
13919 self.compaction_zstd_level
13920 }
13921
13922 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
13923 self.kek.clone()
13924 }
13925
13926 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
13930 self.kek.as_ref().map(|k| k.derive_idx_key())
13931 }
13932
13933 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
13937 self.kek.as_ref().map(|k| *k.derive_meta_key())
13938 }
13939
13940 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
13943 self.column_keys
13944 .iter()
13945 .map(|(&id, &(_, scheme))| (id, scheme))
13946 .collect()
13947 }
13948
13949 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
13954 self.tokenize_value_enc(column_id, v)
13955 }
13956
13957 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
13958 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
13959 let (key, scheme) = self.column_keys.get(&column_id)?;
13960 let token: Vec<u8> = match (*scheme, v) {
13961 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
13962 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
13963 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
13964 _ => hmac_token(key, &v.encode_key()).to_vec(),
13965 };
13966 Some(Value::Bytes(token))
13967 }
13968
13969 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
13971 self.index_lookup_key_bytes(column_id, &v.encode_key())
13972 }
13973
13974 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
13977 {
13978 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
13979 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
13980 if *scheme == SCHEME_HMAC_EQ {
13981 return hmac_token(key, encoded).to_vec();
13982 }
13983 }
13984 }
13985 let _ = column_id;
13986 encoded.to_vec()
13987 }
13988}
13989
13990fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
13991 let columnar::NativeColumn::Int64 { data, validity } = col else {
13992 return false;
13993 };
13994 if data.len() < n || !columnar::all_non_null(validity, n) {
13995 return false;
13996 }
13997 data.iter()
13998 .take(n)
13999 .zip(data.iter().skip(1))
14000 .all(|(a, b)| a < b)
14001}
14002
14003#[derive(Debug, Clone)]
14007pub struct ColumnStat {
14008 pub min: Option<Value>,
14009 pub max: Option<Value>,
14010 pub null_count: u64,
14011}
14012
14013#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14015pub enum NativeAgg {
14016 Count,
14017 Sum,
14018 Min,
14019 Max,
14020 Avg,
14021}
14022
14023#[derive(Debug, Clone, PartialEq)]
14025pub enum NativeAggResult {
14026 Count(u64),
14027 Int(i64),
14028 Float(f64),
14029 Null,
14031}
14032
14033#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14035pub enum ApproxAgg {
14036 Count,
14037 Sum,
14038 Avg,
14039}
14040
14041#[derive(Debug, Clone)]
14045pub struct ApproxResult {
14046 pub point: f64,
14048 pub ci_low: f64,
14050 pub ci_high: f64,
14052 pub n_population: u64,
14054 pub n_sample_live: usize,
14056 pub n_passing: usize,
14058}
14059
14060#[derive(Debug, Clone, PartialEq)]
14065pub enum AggState {
14066 Count(u64),
14068 SumI {
14070 sum: i128,
14071 count: u64,
14072 },
14073 SumF {
14075 sum: f64,
14076 count: u64,
14077 },
14078 AvgI {
14080 sum: i128,
14081 count: u64,
14082 },
14083 AvgF {
14085 sum: f64,
14086 count: u64,
14087 },
14088 MinI(i64),
14090 MaxI(i64),
14091 MinF(f64),
14093 MaxF(f64),
14094 Empty,
14096}
14097
14098impl AggState {
14099 pub fn merge(self, other: AggState) -> AggState {
14101 use AggState::*;
14102 match (self, other) {
14103 (Empty, x) | (x, Empty) => x,
14104 (Count(a), Count(b)) => Count(a + b),
14105 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
14106 sum: sa + sb,
14107 count: ca + cb,
14108 },
14109 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
14110 sum: sa + sb,
14111 count: ca + cb,
14112 },
14113 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
14114 sum: sa + sb,
14115 count: ca + cb,
14116 },
14117 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
14118 sum: sa + sb,
14119 count: ca + cb,
14120 },
14121 (MinI(a), MinI(b)) => MinI(a.min(b)),
14122 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
14123 (MinF(a), MinF(b)) => MinF(a.min(b)),
14124 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
14125 _ => Empty, }
14127 }
14128
14129 pub fn point(&self) -> Option<f64> {
14131 match self {
14132 AggState::Count(n) => Some(*n as f64),
14133 AggState::SumI { sum, .. } => Some(*sum as f64),
14134 AggState::SumF { sum, .. } => Some(*sum),
14135 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
14136 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
14137 AggState::MinI(n) => Some(*n as f64),
14138 AggState::MaxI(n) => Some(*n as f64),
14139 AggState::MinF(n) => Some(*n),
14140 AggState::MaxF(n) => Some(*n),
14141 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
14142 }
14143 }
14144
14145 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
14149 let is_float = matches!(ty, Some(TypeId::Float64));
14150 match (agg, result) {
14151 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
14152 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
14153 sum: x as i128,
14154 count: 1, },
14156 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
14157 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
14158 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
14159 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
14160 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
14161 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
14162 (NativeAgg::Count, _) => AggState::Empty,
14163 (_, NativeAggResult::Null) => AggState::Empty,
14164 _ => {
14165 let _ = is_float;
14166 AggState::Empty
14167 }
14168 }
14169 }
14170}
14171
14172#[derive(Debug, Clone)]
14175pub struct CachedAgg {
14176 pub state: AggState,
14177 pub watermark: u64,
14178 pub epoch: u64,
14179}
14180
14181#[derive(Debug, Clone)]
14183pub struct IncrementalAggResult {
14184 pub state: AggState,
14186 pub incremental: bool,
14189 pub delta_rows: u64,
14191}
14192
14193fn agg_state_from_rows(
14197 rows: &[Row],
14198 conditions: &[crate::query::Condition],
14199 index_sets: &[RowIdSet],
14200 column: Option<u16>,
14201 agg: NativeAgg,
14202 schema: &Schema,
14203 control: Option<&crate::ExecutionControl>,
14204) -> Result<AggState> {
14205 let mut count: u64 = 0;
14206 let mut sum_i: i128 = 0;
14207 let mut sum_f: f64 = 0.0;
14208 let mut mn_i: i64 = i64::MAX;
14209 let mut mx_i: i64 = i64::MIN;
14210 let mut mn_f: f64 = f64::INFINITY;
14211 let mut mx_f: f64 = f64::NEG_INFINITY;
14212 let mut saw_int = false;
14213 let mut saw_float = false;
14214 for (index, r) in rows.iter().enumerate() {
14215 execution_checkpoint(control, index)?;
14216 if !conditions
14217 .iter()
14218 .all(|c| condition_matches_row(c, r, schema))
14219 {
14220 continue;
14221 }
14222 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
14223 continue;
14224 }
14225 match agg {
14226 NativeAgg::Count => match column {
14227 None => count += 1,
14229 Some(cid) => match r.columns.get(&cid) {
14232 None | Some(Value::Null) => {}
14233 Some(_) => count += 1,
14234 },
14235 },
14236 _ => match column.and_then(|cid| r.columns.get(&cid)) {
14237 Some(Value::Int64(n)) => {
14238 count += 1;
14239 sum_i += *n as i128;
14240 mn_i = mn_i.min(*n);
14241 mx_i = mx_i.max(*n);
14242 saw_int = true;
14243 }
14244 Some(Value::Float64(f)) => {
14245 count += 1;
14246 sum_f += f;
14247 mn_f = mn_f.min(*f);
14248 mx_f = mx_f.max(*f);
14249 saw_float = true;
14250 }
14251 _ => {}
14252 },
14253 }
14254 }
14255 Ok(match agg {
14256 NativeAgg::Count => {
14257 if count == 0 {
14258 AggState::Empty
14259 } else {
14260 AggState::Count(count)
14261 }
14262 }
14263 NativeAgg::Sum => {
14264 if count == 0 {
14265 AggState::Empty
14266 } else if saw_int {
14267 AggState::SumI { sum: sum_i, count }
14268 } else {
14269 AggState::SumF { sum: sum_f, count }
14270 }
14271 }
14272 NativeAgg::Avg => {
14273 if count == 0 {
14274 AggState::Empty
14275 } else if saw_int {
14276 AggState::AvgI { sum: sum_i, count }
14277 } else {
14278 AggState::AvgF { sum: sum_f, count }
14279 }
14280 }
14281 NativeAgg::Min => {
14282 if !saw_int && !saw_float {
14283 AggState::Empty
14284 } else if saw_int {
14285 AggState::MinI(mn_i)
14286 } else {
14287 AggState::MinF(mn_f)
14288 }
14289 }
14290 NativeAgg::Max => {
14291 if !saw_int && !saw_float {
14292 AggState::Empty
14293 } else if saw_int {
14294 AggState::MaxI(mx_i)
14295 } else {
14296 AggState::MaxF(mx_f)
14297 }
14298 }
14299 })
14300}
14301
14302fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
14306 use crate::query::Condition;
14307 match c {
14308 Condition::Pk(key) => match schema.primary_key() {
14309 Some(pk) => row
14310 .columns
14311 .get(&pk.id)
14312 .map(|v| v.encode_key() == *key)
14313 .unwrap_or(false),
14314 None => false,
14315 },
14316 Condition::BitmapEq { column_id, value } => row
14317 .columns
14318 .get(column_id)
14319 .map(|v| v.encode_key() == *value)
14320 .unwrap_or(false),
14321 Condition::BitmapIn { column_id, values } => {
14322 let key = row.columns.get(column_id).map(|v| v.encode_key());
14323 match key {
14324 Some(k) => values.contains(&k),
14325 None => false,
14326 }
14327 }
14328 Condition::BytesPrefix { column_id, prefix } => row
14329 .columns
14330 .get(column_id)
14331 .map(|v| v.encode_key().starts_with(prefix))
14332 .unwrap_or(false),
14333 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
14334 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
14335 _ => false,
14336 },
14337 Condition::RangeF64 {
14338 column_id,
14339 lo,
14340 lo_inclusive,
14341 hi,
14342 hi_inclusive,
14343 } => match row.columns.get(column_id) {
14344 Some(Value::Float64(n)) => {
14345 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
14346 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
14347 lo_ok && hi_ok
14348 }
14349 _ => false,
14350 },
14351 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
14352 Some(Value::Bytes(b)) => {
14353 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
14354 }
14355 _ => false,
14356 },
14357 Condition::FmContainsAll {
14358 column_id,
14359 patterns,
14360 } => match row.columns.get(column_id) {
14361 Some(Value::Bytes(b)) => patterns
14362 .iter()
14363 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
14364 _ => false,
14365 },
14366 Condition::Ann { .. }
14367 | Condition::SparseMatch { .. }
14368 | Condition::MinHashSimilar { .. } => true,
14369 Condition::IsNull { column_id } => {
14370 matches!(row.columns.get(column_id), Some(Value::Null) | None)
14371 }
14372 Condition::IsNotNull { column_id } => {
14373 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
14374 }
14375 }
14376}
14377
14378fn as_f64(v: Option<&Value>) -> Option<f64> {
14380 match v {
14381 Some(Value::Int64(n)) => Some(*n as f64),
14382 Some(Value::Float64(f)) => Some(*f),
14383 _ => None,
14384 }
14385}
14386
14387fn accumulate_int(
14391 cursor: &mut dyn crate::cursor::Cursor,
14392 control: Option<&crate::ExecutionControl>,
14393) -> Result<(u64, i128, i64, i64)> {
14394 let mut count: u64 = 0;
14395 let mut sum: i128 = 0;
14396 let mut mn: i64 = i64::MAX;
14397 let mut mx: i64 = i64::MIN;
14398 while let Some(cols) = cursor.next_batch()? {
14399 execution_checkpoint(control, 0)?;
14400 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
14401 if crate::columnar::all_non_null(validity, data.len()) {
14402 count += data.len() as u64;
14404 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
14405 execution_checkpoint(control, chunk_index * 1024)?;
14406 sum += chunk.iter().map(|&v| v as i128).sum::<i128>();
14407 mn = mn.min(*chunk.iter().min().unwrap_or(&mn));
14408 mx = mx.max(*chunk.iter().max().unwrap_or(&mx));
14409 }
14410 } else {
14411 for (i, &v) in data.iter().enumerate() {
14412 execution_checkpoint(control, i)?;
14413 if crate::columnar::validity_bit(validity, i) {
14414 count += 1;
14415 sum += v as i128;
14416 mn = mn.min(v);
14417 mx = mx.max(v);
14418 }
14419 }
14420 }
14421 }
14422 }
14423 Ok((count, sum, mn, mx))
14424}
14425
14426fn accumulate_float(
14428 cursor: &mut dyn crate::cursor::Cursor,
14429 control: Option<&crate::ExecutionControl>,
14430) -> Result<(u64, f64, f64, f64)> {
14431 let mut count: u64 = 0;
14432 let mut sum: f64 = 0.0;
14433 let mut mn: f64 = f64::INFINITY;
14434 let mut mx: f64 = f64::NEG_INFINITY;
14435 while let Some(cols) = cursor.next_batch()? {
14436 execution_checkpoint(control, 0)?;
14437 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
14438 if crate::columnar::all_non_null(validity, data.len()) {
14439 count += data.len() as u64;
14440 for (chunk_index, chunk) in data.chunks(1024).enumerate() {
14441 execution_checkpoint(control, chunk_index * 1024)?;
14442 sum += chunk.iter().sum::<f64>();
14443 mn = mn.min(chunk.iter().copied().fold(f64::INFINITY, f64::min));
14444 mx = mx.max(chunk.iter().copied().fold(f64::NEG_INFINITY, f64::max));
14445 }
14446 } else {
14447 for (i, &v) in data.iter().enumerate() {
14448 execution_checkpoint(control, i)?;
14449 if crate::columnar::validity_bit(validity, i) {
14450 count += 1;
14451 sum += v;
14452 mn = mn.min(v);
14453 mx = mx.max(v);
14454 }
14455 }
14456 }
14457 }
14458 }
14459 Ok((count, sum, mn, mx))
14460}
14461
14462#[inline]
14463fn execution_checkpoint(control: Option<&crate::ExecutionControl>, index: usize) -> Result<()> {
14464 if index.is_multiple_of(256) {
14465 control
14466 .map(crate::ExecutionControl::checkpoint)
14467 .transpose()?;
14468 }
14469 Ok(())
14470}
14471
14472fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
14473 if count == 0 && !matches!(agg, NativeAgg::Count) {
14474 return NativeAggResult::Null;
14475 }
14476 match agg {
14477 NativeAgg::Count => NativeAggResult::Count(count),
14478 NativeAgg::Sum => match sum.try_into() {
14481 Ok(v) => NativeAggResult::Int(v),
14482 Err(_) => NativeAggResult::Null,
14483 },
14484 NativeAgg::Min => NativeAggResult::Int(mn),
14485 NativeAgg::Max => NativeAggResult::Int(mx),
14486 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
14487 }
14488}
14489
14490fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
14491 if count == 0 && !matches!(agg, NativeAgg::Count) {
14492 return NativeAggResult::Null;
14493 }
14494 match agg {
14495 NativeAgg::Count => NativeAggResult::Count(count),
14496 NativeAgg::Sum => NativeAggResult::Float(sum),
14497 NativeAgg::Min => NativeAggResult::Float(mn),
14498 NativeAgg::Max => NativeAggResult::Float(mx),
14499 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
14500 }
14501}
14502
14503fn agg_int(
14506 stats: &[crate::page::PageStat],
14507 decode: fn(Option<&[u8]>) -> Option<i64>,
14508) -> Option<(Option<i64>, Option<i64>, u64)> {
14509 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
14510 let mut any = false;
14511 for s in stats {
14512 if let Some(v) = decode(s.min.as_deref()) {
14513 mn = mn.min(v);
14514 any = true;
14515 }
14516 if let Some(v) = decode(s.max.as_deref()) {
14517 mx = mx.max(v);
14518 any = true;
14519 }
14520 nulls += s.null_count;
14521 }
14522 any.then_some((Some(mn), Some(mx), nulls))
14523}
14524
14525fn agg_float(
14527 stats: &[crate::page::PageStat],
14528 decode: fn(Option<&[u8]>) -> Option<f64>,
14529) -> Option<(Option<f64>, Option<f64>, u64)> {
14530 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
14531 let mut any = false;
14532 for s in stats {
14533 if let Some(v) = decode(s.min.as_deref()) {
14534 mn = mn.min(v);
14535 any = true;
14536 }
14537 if let Some(v) = decode(s.max.as_deref()) {
14538 mx = mx.max(v);
14539 any = true;
14540 }
14541 nulls += s.null_count;
14542 }
14543 any.then_some((Some(mn), Some(mx), nulls))
14544}
14545
14546type SecondaryIndexes = (
14548 HashMap<u16, BitmapIndex>,
14549 HashMap<u16, AnnIndex>,
14550 HashMap<u16, FmIndex>,
14551 HashMap<u16, SparseIndex>,
14552 HashMap<u16, MinHashIndex>,
14553);
14554
14555fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
14556 let mut bitmap = HashMap::new();
14557 let mut ann = HashMap::new();
14558 let mut fm = HashMap::new();
14559 let mut sparse = HashMap::new();
14560 let mut minhash = HashMap::new();
14561 for idef in &schema.indexes {
14562 match idef.kind {
14563 IndexKind::Bitmap => {
14564 bitmap.insert(idef.column_id, BitmapIndex::new());
14565 }
14566 IndexKind::Ann => {
14567 let dim = schema
14568 .columns
14569 .iter()
14570 .find(|c| c.id == idef.column_id)
14571 .and_then(|c| match c.ty {
14572 TypeId::Embedding { dim } => Some(dim as usize),
14573 _ => None,
14574 })
14575 .unwrap_or(0);
14576 let options = idef.options.ann.clone().unwrap_or_default();
14577 ann.insert(
14578 idef.column_id,
14579 AnnIndex::with_full_options(
14580 dim,
14581 options.m,
14582 options.ef_construction,
14583 options.ef_search,
14584 &options,
14585 ),
14586 );
14587 }
14588 IndexKind::FmIndex => {
14589 fm.insert(idef.column_id, FmIndex::new());
14590 }
14591 IndexKind::Sparse => {
14592 sparse.insert(idef.column_id, SparseIndex::new());
14593 }
14594 IndexKind::MinHash => {
14595 let options = idef.options.minhash.clone().unwrap_or_default();
14596 minhash.insert(
14597 idef.column_id,
14598 MinHashIndex::with_options(options.permutations, options.bands),
14599 );
14600 }
14601 _ => {}
14602 }
14603 }
14604 (bitmap, ann, fm, sparse, minhash)
14605}
14606
14607const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
14608 | ColumnFlags::AUTO_INCREMENT
14609 | ColumnFlags::ENCRYPTED
14610 | ColumnFlags::ENCRYPTED_INDEXABLE
14611 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
14612
14613fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
14614 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
14615 return Err(MongrelError::Schema(
14616 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
14617 ));
14618 }
14619 Ok(())
14620}
14621
14622fn validate_alter_column_type(
14623 schema: &Schema,
14624 old: &ColumnDef,
14625 next: &ColumnDef,
14626 has_stored_versions: bool,
14627) -> Result<()> {
14628 if old.ty == next.ty {
14629 return Ok(());
14630 }
14631 if schema.indexes.iter().any(|i| i.column_id == old.id) {
14632 return Err(MongrelError::Schema(format!(
14633 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
14634 old.name
14635 )));
14636 }
14637 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
14638 return Ok(());
14639 }
14640 Err(MongrelError::Schema(format!(
14641 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
14642 old.ty, next.ty
14643 )))
14644}
14645
14646fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
14647 matches!(
14648 (old, new),
14649 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
14650 )
14651}
14652
14653fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
14659 let mut prev: Option<i64> = None;
14660 for r in rows {
14661 match r.columns.get(&pk_id) {
14662 Some(Value::Int64(v)) => {
14663 if prev.is_some_and(|p| p >= *v) {
14664 return false;
14665 }
14666 prev = Some(*v);
14667 }
14668 _ => return false,
14669 }
14670 }
14671 true
14672}
14673
14674#[allow(clippy::too_many_arguments)]
14675fn index_into(
14676 schema: &Schema,
14677 row: &Row,
14678 hot: &mut HotIndex,
14679 bitmap: &mut HashMap<u16, BitmapIndex>,
14680 ann: &mut HashMap<u16, AnnIndex>,
14681 fm: &mut HashMap<u16, FmIndex>,
14682 sparse: &mut HashMap<u16, SparseIndex>,
14683 minhash: &mut HashMap<u16, MinHashIndex>,
14684) {
14685 for idef in &schema.indexes {
14686 let Some(val) = row.columns.get(&idef.column_id) else {
14687 continue;
14688 };
14689 match idef.kind {
14690 IndexKind::Bitmap => {
14691 if let Some(b) = bitmap.get_mut(&idef.column_id) {
14692 b.insert(val.encode_key(), row.row_id);
14693 }
14694 }
14695 IndexKind::Ann => {
14696 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
14697 if let Some(meta) = val.generated_embedding_metadata() {
14698 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
14700 continue;
14701 }
14702 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
14703 .is_err()
14704 {
14705 continue;
14706 }
14707 }
14708 a.insert_validated(v, row.row_id);
14709 }
14710 }
14711 IndexKind::FmIndex => {
14712 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
14713 f.insert(b.clone(), row.row_id);
14714 }
14715 }
14716 IndexKind::Sparse => {
14717 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
14718 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
14721 s.insert(&terms, row.row_id);
14722 }
14723 }
14724 }
14725 IndexKind::MinHash => {
14726 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
14727 let tokens = crate::index::token_hashes_from_bytes(b);
14730 mh.insert(&tokens, row.row_id);
14731 }
14732 }
14733 _ => {}
14734 }
14735 }
14736 if let Some(pk_col) = schema.primary_key() {
14737 if let Some(pk_val) = row.columns.get(&pk_col.id) {
14738 hot.insert(pk_val.encode_key(), row.row_id);
14739 }
14740 }
14741}
14742
14743#[allow(clippy::too_many_arguments)]
14746fn index_into_single(
14747 idef: &IndexDef,
14748 _schema: &Schema,
14749 row: &Row,
14750 _hot: &mut HotIndex,
14751 bitmap: &mut HashMap<u16, BitmapIndex>,
14752 ann: &mut HashMap<u16, AnnIndex>,
14753 fm: &mut HashMap<u16, FmIndex>,
14754 sparse: &mut HashMap<u16, SparseIndex>,
14755 minhash: &mut HashMap<u16, MinHashIndex>,
14756) {
14757 let Some(val) = row.columns.get(&idef.column_id) else {
14758 return;
14759 };
14760 match idef.kind {
14761 IndexKind::Bitmap => {
14762 if let Some(b) = bitmap.get_mut(&idef.column_id) {
14763 b.insert(val.encode_key(), row.row_id);
14764 }
14765 }
14766 IndexKind::Ann => {
14767 if let (Some(a), Some(v)) = (ann.get_mut(&idef.column_id), val.as_embedding()) {
14768 if let Some(meta) = val.generated_embedding_metadata() {
14769 if !crate::embedding_jobs::embedding_status_is_ann_eligible(meta.status) {
14771 return;
14772 }
14773 if a.bind_or_check_semantic_identity(&meta.semantic_identity)
14774 .is_err()
14775 {
14776 return;
14777 }
14778 }
14779 a.insert_validated(v, row.row_id);
14780 }
14781 }
14782 IndexKind::FmIndex => {
14783 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
14784 f.insert(b.clone(), row.row_id);
14785 }
14786 }
14787 IndexKind::Sparse => {
14788 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
14789 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
14790 s.insert(&terms, row.row_id);
14791 }
14792 }
14793 }
14794 IndexKind::MinHash => {
14795 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
14796 let tokens = crate::index::token_hashes_from_bytes(b);
14797 mh.insert(&tokens, row.row_id);
14798 }
14799 }
14800 _ => {}
14801 }
14802}
14803
14804fn eval_partial_predicate(
14810 pred: &str,
14811 columns_map: &HashMap<u16, &Value>,
14812 name_to_id: &HashMap<&str, u16>,
14813) -> bool {
14814 let lower = pred.trim().to_ascii_lowercase();
14815 if let Some(rest) = lower.strip_suffix(" is not null") {
14817 let col_name = rest.trim();
14818 if let Some(col_id) = name_to_id.get(col_name) {
14819 return columns_map
14820 .get(col_id)
14821 .is_some_and(|v| !matches!(v, Value::Null));
14822 }
14823 }
14824 if let Some(rest) = lower.strip_suffix(" is null") {
14826 let col_name = rest.trim();
14827 if let Some(col_id) = name_to_id.get(col_name) {
14828 return columns_map
14829 .get(col_id)
14830 .is_none_or(|v| matches!(v, Value::Null));
14831 }
14832 }
14833 true
14836}
14837
14838#[allow(dead_code)]
14844fn bulk_index_key(
14845 column_keys: &HashMap<u16, ([u8; 32], u8)>,
14846 column_id: u16,
14847 ty: TypeId,
14848 col: &columnar::NativeColumn,
14849 i: usize,
14850) -> Option<Vec<u8>> {
14851 let encoded = columnar::encode_key_native(ty, col, i)?;
14852 {
14853 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
14854 if let Some((key, scheme)) = column_keys.get(&column_id) {
14855 return Some(match (*scheme, col) {
14856 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
14857 (_, columnar::NativeColumn::Int64 { data, .. }) => {
14858 ope_token_i64(key, data[i]).to_vec()
14859 }
14860 (_, columnar::NativeColumn::Float64 { data, .. }) => {
14861 ope_token_f64(key, data[i]).to_vec()
14862 }
14863 _ => hmac_token(key, &encoded).to_vec(),
14864 });
14865 }
14866 }
14867 Some(encoded)
14868}
14869
14870pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
14871 write_schema_with_after(dir, schema, || {})
14872}
14873
14874pub(crate) fn write_schema_durable(
14875 root: &crate::durable_file::DurableRoot,
14876 schema: &Schema,
14877) -> Result<()> {
14878 write_schema_durable_with_after(root, schema, || {})
14879}
14880
14881fn write_schema_with_after<F>(dir: &Path, schema: &Schema, after_publish: F) -> Result<()>
14882where
14883 F: FnOnce(),
14884{
14885 let json = serde_json::to_string_pretty(schema)
14886 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
14887 crate::durable_file::write_atomic_with_after(
14888 &dir.join(SCHEMA_FILENAME),
14889 json.as_bytes(),
14890 after_publish,
14891 )?;
14892 Ok(())
14893}
14894
14895fn write_schema_durable_with_after<F>(
14896 root: &crate::durable_file::DurableRoot,
14897 schema: &Schema,
14898 after_publish: F,
14899) -> Result<()>
14900where
14901 F: FnOnce(),
14902{
14903 let json = serde_json::to_string_pretty(schema)
14904 .map_err(|error| MongrelError::Schema(format!("encode schema: {error}")))?;
14905 root.write_atomic_with_after(SCHEMA_FILENAME, json.as_bytes(), after_publish)?;
14906 Ok(())
14907}
14908
14909fn checkpoint_current_schema(table: &mut Table) -> Result<()> {
14910 let mut schema_published = false;
14911 let schema_result = match table._root_guard.as_deref() {
14912 Some(root) => write_schema_durable_with_after(root, &table.schema, || {
14913 schema_published = true;
14914 }),
14915 None => write_schema_with_after(&table.dir, &table.schema, || {
14916 schema_published = true;
14917 }),
14918 };
14919 if schema_result.is_err() && !schema_published {
14920 return schema_result;
14921 }
14922 match table.persist_manifest(table.current_epoch()) {
14923 Ok(()) => Ok(()),
14924 Err(manifest_error) => Err(match schema_result {
14925 Ok(()) => manifest_error,
14926 Err(schema_error) => MongrelError::Other(format!(
14927 "schema publication sync failed ({schema_error}); matching manifest publication also failed ({manifest_error})"
14928 )),
14929 }),
14930 }
14931}
14932
14933fn read_schema(dir: &Path) -> Result<Schema> {
14934 let file = crate::durable_file::open_regular_nofollow(&dir.join(SCHEMA_FILENAME))?;
14935 read_schema_file(file)
14936}
14937
14938fn read_schema_file(file: std::fs::File) -> Result<Schema> {
14939 const MAX_SCHEMA_BYTES: u64 = 16 * 1024 * 1024;
14940 use std::io::Read;
14941
14942 let length = file.metadata()?.len();
14943 if length > MAX_SCHEMA_BYTES {
14944 return Err(MongrelError::ResourceLimitExceeded {
14945 resource: "schema bytes",
14946 requested: usize::try_from(length).unwrap_or(usize::MAX),
14947 limit: MAX_SCHEMA_BYTES as usize,
14948 });
14949 }
14950 let mut bytes = Vec::with_capacity(length as usize);
14951 file.take(MAX_SCHEMA_BYTES + 1).read_to_end(&mut bytes)?;
14952 if bytes.len() as u64 != length {
14953 return Err(MongrelError::Schema(
14954 "schema length changed while reading".into(),
14955 ));
14956 }
14957 serde_json::from_slice(&bytes).map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
14958}
14959
14960fn preflight_standalone_open(
14961 dir: &Path,
14962 runs_root: Option<&crate::durable_file::DurableRoot>,
14963 idx_root: Option<&crate::durable_file::DurableRoot>,
14964 manifest: &Manifest,
14965 schema: &Schema,
14966 records: &[crate::wal::Record],
14967 kek: Option<Arc<Kek>>,
14968) -> Result<()> {
14969 crate::wal::validate_shared_transaction_framing(records)?;
14970 if manifest.schema_id > schema.schema_id
14971 || manifest.flushed_epoch > manifest.current_epoch
14972 || manifest.global_idx_epoch > manifest.current_epoch
14973 || manifest.next_row_id == u64::MAX
14974 || manifest.auto_inc_next < 0
14975 || manifest.auto_inc_next == i64::MAX
14976 || (schema.auto_increment_column().is_none() && manifest.auto_inc_next != 0)
14977 {
14978 return Err(MongrelError::InvalidArgument(
14979 "manifest counters or schema identity are invalid".into(),
14980 ));
14981 }
14982 let mut run_ids = HashSet::new();
14983 let mut maximum_row_id = None::<u64>;
14984 for run in &manifest.runs {
14985 if run.run_id >= u64::MAX as u128
14986 || !run_ids.insert(run.run_id)
14987 || run.epoch_created > manifest.current_epoch
14988 {
14989 return Err(MongrelError::InvalidArgument(
14990 "manifest contains an invalid or duplicate active run".into(),
14991 ));
14992 }
14993 let mut reader = match runs_root {
14994 Some(root) => RunReader::open_file(
14995 root.open_regular(format!("r-{}.sr", run.run_id as u64))?,
14996 schema.clone(),
14997 kek.clone(),
14998 )?,
14999 None => RunReader::open(
15000 dir.join(RUNS_DIR)
15001 .join(format!("r-{}.sr", run.run_id as u64)),
15002 schema.clone(),
15003 kek.clone(),
15004 )?,
15005 };
15006 let header = reader.header();
15007 if header.run_id != run.run_id
15008 || header.level != run.level
15009 || header.row_count != run.row_count
15010 || !header.is_uniform_epoch() && header.epoch_created != run.epoch_created
15011 || header.is_uniform_epoch() && header.epoch_created != 0
15012 || header.schema_id > schema.schema_id
15013 {
15014 return Err(MongrelError::InvalidArgument(format!(
15015 "run {} differs from its manifest",
15016 run.run_id
15017 )));
15018 }
15019 if header.row_count != 0 {
15020 maximum_row_id = Some(
15021 maximum_row_id.map_or(header.max_row_id, |value| value.max(header.max_row_id)),
15022 );
15023 }
15024 reader.validate_all_pages()?;
15025 }
15026 if maximum_row_id.is_some_and(|maximum| manifest.next_row_id <= maximum) {
15027 return Err(MongrelError::InvalidArgument(
15028 "manifest next_row_id does not advance beyond persisted rows".into(),
15029 ));
15030 }
15031 for run in &manifest.retiring {
15032 if run.run_id >= u64::MAX as u128
15033 || run.retire_epoch > manifest.current_epoch
15034 || !run_ids.insert(run.run_id)
15035 {
15036 return Err(MongrelError::InvalidArgument(
15037 "manifest contains an invalid or duplicate retired run".into(),
15038 ));
15039 }
15040 }
15041 let idx_dek = kek.as_ref().map(|key| key.derive_idx_key());
15042 match idx_root {
15043 Some(root) => {
15044 global_idx::read_root(root, manifest.table_id, schema, idx_dek.as_deref())?;
15045 }
15046 None => {
15047 global_idx::read(dir, manifest.table_id, schema, idx_dek.as_deref())?;
15048 }
15049 }
15050
15051 let committed = records
15052 .iter()
15053 .filter_map(|record| match record.op {
15054 Op::TxnCommit { epoch, .. } => Some((record.txn_id, epoch)),
15055 _ => None,
15056 })
15057 .collect::<HashMap<_, _>>();
15058 for record in records {
15059 let Some(&_commit_epoch) = committed.get(&record.txn_id) else {
15060 continue;
15061 };
15062 match &record.op {
15063 Op::Put { table_id, rows } => {
15064 if *table_id != manifest.table_id {
15065 return Err(MongrelError::CorruptWal {
15066 offset: record.seq.0,
15067 reason: format!(
15068 "private WAL record references table {table_id}, expected {}",
15069 manifest.table_id
15070 ),
15071 });
15072 }
15073 let rows: Vec<Row> =
15074 bincode::deserialize(rows).map_err(|error| MongrelError::CorruptWal {
15075 offset: record.seq.0,
15076 reason: format!("committed Put payload could not be decoded: {error}"),
15077 })?;
15078 for row in rows {
15079 if row.deleted || row.row_id.0 == u64::MAX {
15080 return Err(MongrelError::CorruptWal {
15081 offset: record.seq.0,
15082 reason: "committed Put contains an invalid row identity".into(),
15083 });
15084 }
15085 let cells = row.columns.into_iter().collect::<Vec<_>>();
15086 schema
15087 .validate_values(&cells)
15088 .map_err(|error| MongrelError::CorruptWal {
15089 offset: record.seq.0,
15090 reason: format!("committed Put violates table schema: {error}"),
15091 })?;
15092 if schema.auto_increment_column().is_some_and(|column| {
15093 matches!(
15094 cells.iter().find(|(id, _)| *id == column.id),
15095 Some((_, Value::Int64(value))) if *value == i64::MAX
15096 )
15097 }) {
15098 return Err(MongrelError::CorruptWal {
15099 offset: record.seq.0,
15100 reason: "committed Put exhausts AUTO_INCREMENT".into(),
15101 });
15102 }
15103 }
15104 }
15105 Op::Delete { table_id, .. } | Op::TruncateTable { table_id }
15106 if *table_id != manifest.table_id =>
15107 {
15108 return Err(MongrelError::CorruptWal {
15109 offset: record.seq.0,
15110 reason: format!(
15111 "private WAL record references table {table_id}, expected {}",
15112 manifest.table_id
15113 ),
15114 });
15115 }
15116 Op::TxnCommit { added_runs, .. } if !added_runs.is_empty() => {
15117 return Err(MongrelError::CorruptWal {
15118 offset: record.seq.0,
15119 reason: "private WAL contains shared spilled-run metadata".into(),
15120 });
15121 }
15122 _ => {}
15123 }
15124 }
15125 Ok(())
15126}
15127
15128fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
15129 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
15130}
15131
15132fn wal_segment_number(path: &Path) -> Option<u64> {
15133 path.file_stem()
15134 .and_then(|stem| stem.to_str())
15135 .and_then(|stem| stem.strip_prefix("seg-"))
15136 .and_then(|number| number.parse().ok())
15137}
15138
15139fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
15140 let n = list_wal_numbers(wal_dir)?;
15141 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
15142}
15143
15144fn next_wal_number(wal_dir: &Path) -> Result<u32> {
15145 list_wal_numbers(wal_dir)?
15146 .map(|maximum| {
15147 maximum
15148 .checked_add(1)
15149 .ok_or_else(|| MongrelError::Full("WAL segment namespace exhausted".into()))
15150 })
15151 .unwrap_or(Ok(0))
15152}
15153
15154fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
15155 let mut max_n = None;
15156 let entries = match std::fs::read_dir(wal_dir) {
15157 Ok(entries) => entries,
15158 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
15159 Err(error) => return Err(error.into()),
15160 };
15161 for entry in entries {
15162 let entry = entry?;
15163 let fname = entry.file_name();
15164 let Some(s) = fname.to_str() else {
15165 continue;
15166 };
15167 let Some(stripped) = s.strip_prefix("seg-") else {
15168 continue;
15169 };
15170 let Some(number) = stripped.strip_suffix(".wal") else {
15171 return Err(MongrelError::CorruptWal {
15172 offset: 0,
15173 reason: format!("malformed WAL segment name {s:?}"),
15174 });
15175 };
15176 let n = number
15177 .parse::<u32>()
15178 .map_err(|_| MongrelError::CorruptWal {
15179 offset: 0,
15180 reason: format!("malformed WAL segment name {s:?}"),
15181 })?;
15182 if s != format!("seg-{n:06}.wal") || !entry.file_type()?.is_file() {
15183 return Err(MongrelError::CorruptWal {
15184 offset: n as u64,
15185 reason: format!("noncanonical or nonregular WAL segment {s:?}"),
15186 });
15187 }
15188 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
15189 }
15190 Ok(max_n)
15191}