1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::collections::{BTreeMap, HashMap, HashSet};
30use std::path::{Path, PathBuf};
31use std::sync::atomic::AtomicBool;
32use std::sync::Arc;
33use zeroize::Zeroizing;
34
35pub const WAL_DIR: &str = "_wal";
36pub const RUNS_DIR: &str = "_runs";
37pub const CACHE_DIR: &str = "_cache";
38pub const META_DIR: &str = "_meta";
39pub const RCACHE_DIR: &str = "_rcache";
40pub const KEYS_FILENAME: &str = "keys";
41pub const SCHEMA_FILENAME: &str = "schema.json";
42
43fn iso_now_bytes() -> Vec<u8> {
46 let secs = std::time::SystemTime::now()
47 .duration_since(std::time::UNIX_EPOCH)
48 .map(|d| d.as_secs() as i64)
49 .unwrap_or(0);
50 let days = secs.div_euclid(86_400);
51 let rem = secs.rem_euclid(86_400);
52 let (hour, minute, second) = (rem / 3600, (rem % 3600) / 60, rem % 60);
53 let (year, month, day) = civil_from_days(days);
54 format!("{year:04}-{month:02}-{day:02}T{hour:02}:{minute:02}:{second:02}Z").into_bytes()
55}
56
57fn civil_from_days(z: i64) -> (i64, u32, u32) {
58 let z = z + 719_468;
59 let era = if z >= 0 { z } else { z - 146_096 } / 146_097;
60 let doe = z - era * 146_097;
61 let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
62 let y = yoe + era * 400;
63 let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
64 let mp = (5 * doy + 2) / 153;
65 let d = (doy - (153 * mp + 2) / 5 + 1) as u32;
66 let m = if mp < 10 { mp + 3 } else { mp - 9 } as u32;
67 (if m <= 2 { y + 1 } else { y }, m, d)
68}
69
70const 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;
77
78#[derive(Clone, Copy, Debug)]
93struct AutoIncState {
94 column_id: u16,
95 next: i64,
96 seeded: bool,
97}
98
99type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
100
101fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
104 schema.auto_increment_column().map(|c| AutoIncState {
105 column_id: c.id,
106 next: 0,
107 seeded: false,
108 })
109}
110
111#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
124pub enum IndexBuildPolicy {
125 #[default]
127 Deferred,
128 Eager,
130}
131
132pub struct Table {
134 dir: PathBuf,
135 table_id: u64,
136 name: String,
140 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
145 wal: WalSink,
146 memtable: Memtable,
147 mutable_run: MutableRun,
152 mutable_run_spill_bytes: u64,
154 compaction_zstd_level: i32,
157 allocator: RowIdAllocator,
158 epoch: Arc<EpochAuthority>,
159 persisted_epoch: u64,
162 schema: Schema,
163 hot: HotIndex,
164 kek: Option<Arc<Kek>>,
167 column_keys: HashMap<u16, ([u8; 32], u8)>,
171 run_refs: Vec<RunRef>,
172 retiring: Vec<crate::manifest::RetiredRun>,
175 next_run_id: u64,
176 sync_byte_threshold: u64,
177 current_txn_id: u64,
182 bitmap: HashMap<u16, BitmapIndex>,
183 ann: HashMap<u16, AnnIndex>,
184 fm: HashMap<u16, FmIndex>,
185 sparse: HashMap<u16, SparseIndex>,
186 minhash: HashMap<u16, MinHashIndex>,
187 learned_range: HashMap<u16, ColumnLearnedRange>,
190 pk_by_row: HashMap<RowId, Vec<u8>>,
192 pinned: BTreeMap<Epoch, usize>,
195 pub(crate) live_count: u64,
198 reservoir: crate::reservoir::Reservoir,
201 reservoir_complete: bool,
209 had_deletes: bool,
213 agg_cache: HashMap<u64, CachedAgg>,
217 global_idx_epoch: u64,
221 indexes_complete: bool,
226 index_build_policy: IndexBuildPolicy,
228 pk_by_row_complete: bool,
235 flushed_epoch: u64,
238 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
241 snapshots: Arc<crate::retention::SnapshotRegistry>,
244 commit_lock: Arc<parking_lot::Mutex<()>>,
246 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
249 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
259 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
268 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
270 pending_delete_rids: roaring::RoaringBitmap,
273 pending_put_cols: std::collections::HashSet<u16>,
276 pending_rows: Vec<Row>,
282 pending_rows_auto_inc: Vec<bool>,
283 pending_dels: Vec<RowId>,
286 pending_truncate: Option<Epoch>,
290 auto_inc: Option<AutoIncState>,
293}
294
295const _: () = {
302 const fn assert_sync<T: ?Sized + Sync>() {}
303 assert_sync::<Table>();
304};
305
306enum CachedData {
312 Rows(Arc<Vec<Row>>),
313 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
314}
315
316impl CachedData {
317 fn approx_bytes(&self) -> u64 {
318 match self {
319 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
320 CachedData::Columns(c) => c
321 .iter()
322 .map(|(_, c)| c.approx_bytes())
323 .sum::<u64>()
324 .saturating_add(c.len() as u64 * 16),
325 }
326 }
327}
328
329struct CachedEntry {
333 data: CachedData,
334 footprint: roaring::RoaringBitmap,
335 condition_cols: Vec<u16>,
336}
337
338struct ResultCache {
349 entries: std::collections::HashMap<u64, CachedEntry>,
350 order: std::collections::VecDeque<u64>,
351 bytes: u64,
352 max_bytes: u64,
353 dir: Option<std::path::PathBuf>,
354 #[allow(dead_code)]
355 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
356}
357
358#[derive(serde::Serialize, serde::Deserialize)]
360struct SerializedEntry {
361 condition_cols: Vec<u16>,
362 footprint_bits: Vec<u32>,
363 data: SerializedData,
364}
365
366#[derive(serde::Serialize, serde::Deserialize)]
367enum SerializedData {
368 Rows(Vec<Row>),
369 Columns(Vec<(u16, columnar::NativeColumn)>),
370}
371
372impl SerializedEntry {
373 fn from_entry(entry: &CachedEntry) -> Self {
374 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
375 let data = match &entry.data {
376 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
377 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
378 };
379 Self {
380 condition_cols: entry.condition_cols.clone(),
381 footprint_bits,
382 data,
383 }
384 }
385
386 fn into_entry(self) -> Option<CachedEntry> {
387 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
388 let data = match self.data {
389 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
390 SerializedData::Columns(c) => {
391 if !c.iter().all(|(_, col)| col.validate()) {
394 return None;
395 }
396 CachedData::Columns(Arc::new(c))
397 }
398 };
399 Some(CachedEntry {
400 data,
401 footprint,
402 condition_cols: self.condition_cols,
403 })
404 }
405}
406
407impl ResultCache {
408 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
409
410 fn new() -> Self {
411 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
412 }
413
414 fn with_max_bytes(max_bytes: u64) -> Self {
415 Self {
416 entries: std::collections::HashMap::new(),
417 order: std::collections::VecDeque::new(),
418 bytes: 0,
419 max_bytes,
420 dir: None,
421 cache_dek: None,
422 }
423 }
424
425 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
426 let _ = std::fs::create_dir_all(&dir);
427 self.dir = Some(dir);
428 self
429 }
430
431 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
432 self.cache_dek = dek;
433 self
434 }
435
436 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
437 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
438 }
439
440 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
444 let Some(path) = self.disk_path(key) else {
445 return;
446 };
447 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
448 Ok(s) => s,
449 Err(_) => return,
450 };
451 let on_disk = if let Some(dek) = &self.cache_dek {
453 match self.encrypt_cache(&serialized, dek) {
454 Some(b) => b,
455 None => return,
456 }
457 } else {
458 serialized
459 };
460 let tmp = path.with_extension("tmp");
461 use std::io::Write;
462 let write = || -> std::io::Result<()> {
463 let mut f = std::fs::File::create(&tmp)?;
464 f.write_all(&on_disk)?;
465 f.flush()?;
466 Ok(())
467 };
468 if write().is_err() {
469 let _ = std::fs::remove_file(&tmp);
470 return;
471 }
472 let _ = std::fs::rename(&tmp, &path);
473 }
474
475 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
477 let path = self.disk_path(key)?;
478 let bytes = std::fs::read(&path).ok()?;
479 let plaintext = if let Some(dek) = &self.cache_dek {
480 self.decrypt_cache(&bytes, dek)?
481 } else {
482 bytes
483 };
484 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
485 serialized.into_entry()
486 }
487
488 fn remove_from_disk(&self, key: u64) {
490 if let Some(path) = self.disk_path(key) {
491 let _ = std::fs::remove_file(&path);
492 }
493 }
494
495 #[cfg(feature = "encryption")]
497 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
498 use crate::encryption::Cipher;
499 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
500 let mut nonce = [0u8; 12];
501 crate::encryption::fill_random(&mut nonce);
502 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
503 let mut out = Vec::with_capacity(12 + ct.len());
504 out.extend_from_slice(&nonce);
505 out.extend_from_slice(&ct);
506 Some(out)
507 }
508
509 #[cfg(not(feature = "encryption"))]
510 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
511 None
512 }
513
514 #[cfg(feature = "encryption")]
516 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
517 use crate::encryption::Cipher;
518 if bytes.len() < 28 {
519 return None;
520 }
521 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
522 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
523 let ct = &bytes[12..];
524 cipher.decrypt_page(&nonce, ct).ok()
525 }
526
527 #[cfg(not(feature = "encryption"))]
528 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
529 None
530 }
531
532 fn load_persistent(&mut self) {
535 let Some(dir) = self.dir.as_ref().cloned() else {
536 return;
537 };
538 let entries = match std::fs::read_dir(&dir) {
539 Ok(e) => e,
540 Err(_) => return,
541 };
542 for entry in entries.flatten() {
543 let path = entry.path();
544 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
546 let _ = std::fs::remove_file(&path);
547 continue;
548 }
549 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
550 continue;
551 }
552 let stem = match path.file_stem().and_then(|s| s.to_str()) {
553 Some(s) => s,
554 None => continue,
555 };
556 let key = match u64::from_str_radix(stem, 16) {
557 Ok(k) => k,
558 Err(_) => continue,
559 };
560 let bytes = match std::fs::read(&path) {
561 Ok(b) => b,
562 Err(_) => continue,
563 };
564 let plaintext = if let Some(dek) = &self.cache_dek {
566 match self.decrypt_cache(&bytes, dek) {
567 Some(p) => p,
568 None => {
569 let _ = std::fs::remove_file(&path);
570 continue;
571 }
572 }
573 } else {
574 bytes
575 };
576 match bincode::deserialize::<SerializedEntry>(&plaintext) {
577 Ok(serialized) => {
578 if let Some(entry) = serialized.into_entry() {
579 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
580 self.entries.insert(key, entry);
581 self.order.push_back(key);
582 } else {
583 let _ = std::fs::remove_file(&path);
584 }
585 }
586 Err(_) => {
587 let _ = std::fs::remove_file(&path);
588 }
589 }
590 }
591 self.evict();
592 }
593
594 fn set_max_bytes(&mut self, max_bytes: u64) {
595 self.max_bytes = max_bytes;
596 self.evict();
597 }
598
599 fn touch(&mut self, key: u64) {
601 self.order.retain(|k| *k != key);
602 self.order.push_back(key);
603 }
604
605 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
606 let res = self.entries.get(&key).and_then(|e| match &e.data {
607 CachedData::Rows(r) => Some(r.clone()),
608 CachedData::Columns(_) => None,
609 });
610 if res.is_some() {
611 self.touch(key);
612 return res;
613 }
614 if let Some(entry) = self.load_from_disk(key) {
616 let res = match &entry.data {
617 CachedData::Rows(r) => Some(r.clone()),
618 CachedData::Columns(_) => None,
619 };
620 if res.is_some() {
621 let approx = entry.data.approx_bytes();
622 self.bytes = self.bytes.saturating_add(approx);
623 self.entries.insert(key, entry);
624 self.order.push_back(key);
625 self.evict();
626 return res;
627 }
628 }
629 None
630 }
631
632 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
633 let res = self.entries.get(&key).and_then(|e| match &e.data {
634 CachedData::Columns(c) => Some(c.clone()),
635 CachedData::Rows(_) => None,
636 });
637 if res.is_some() {
638 self.touch(key);
639 return res;
640 }
641 if let Some(entry) = self.load_from_disk(key) {
643 let res = match &entry.data {
644 CachedData::Columns(c) => Some(c.clone()),
645 CachedData::Rows(_) => None,
646 };
647 if res.is_some() {
648 let approx = entry.data.approx_bytes();
649 self.bytes = self.bytes.saturating_add(approx);
650 self.entries.insert(key, entry);
651 self.order.push_back(key);
652 self.evict();
653 return res;
654 }
655 }
656 None
657 }
658
659 fn insert(&mut self, key: u64, entry: CachedEntry) {
660 let approx = entry.data.approx_bytes();
661 if self.entries.remove(&key).is_some() {
662 self.order.retain(|k| *k != key);
663 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
664 }
665 self.store_to_disk(key, &entry);
667 self.bytes = self.bytes.saturating_add(approx);
668 self.entries.insert(key, entry);
669 self.order.push_back(key);
670 self.evict();
671 }
672
673 fn invalidate(
682 &mut self,
683 delete_rids: &roaring::RoaringBitmap,
684 put_cols: &std::collections::HashSet<u16>,
685 ) {
686 if self.entries.is_empty() {
687 return;
688 }
689 let has_deletes = !delete_rids.is_empty();
690 let to_remove: std::collections::HashSet<u64> = self
691 .entries
692 .iter()
693 .filter(|(_, e)| {
694 let delete_hit = if e.footprint.is_empty() {
695 has_deletes
696 } else {
697 e.footprint.intersection_len(delete_rids) > 0
698 };
699 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
700 delete_hit || col_hit
701 })
702 .map(|(&k, _)| k)
703 .collect();
704 for key in &to_remove {
705 if let Some(e) = self.entries.remove(key) {
706 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
707 }
708 self.remove_from_disk(*key);
709 }
710 if !to_remove.is_empty() {
711 self.order.retain(|k| !to_remove.contains(k));
712 }
713 }
714
715 fn clear(&mut self) {
716 if let Some(dir) = &self.dir {
718 if let Ok(entries) = std::fs::read_dir(dir) {
719 for entry in entries.flatten() {
720 let path = entry.path();
721 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
722 let _ = std::fs::remove_file(&path);
723 }
724 }
725 }
726 }
727 self.entries.clear();
728 self.order.clear();
729 self.bytes = 0;
730 }
731
732 fn evict(&mut self) {
733 while self.bytes > self.max_bytes {
734 let Some(k) = self.order.pop_front() else {
735 break;
736 };
737 if let Some(e) = self.entries.remove(&k) {
738 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
739 self.remove_from_disk(k);
743 }
744 }
745 }
746}
747
748type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
755
756fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
757 let _ = kek;
758 #[cfg(feature = "encryption")]
759 {
760 if let Some(k) = kek {
761 return (
762 Some(k.derive_table_wal_key(_table_id)),
763 Some(k.derive_cache_key()),
764 );
765 }
766 }
767 (None, None)
768}
769
770#[cfg(feature = "encryption")]
772fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
773 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
774}
775
776#[cfg(not(feature = "encryption"))]
777fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
778 Box::new(crate::encryption::PlaintextCipher)
779}
780
781fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
782 let Some(kek) = kek else {
783 return HashMap::new();
784 };
785 #[cfg(feature = "encryption")]
786 {
787 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
788 schema
789 .columns
790 .iter()
791 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
792 .map(|c| {
793 let scheme = if schema
794 .indexes
795 .iter()
796 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
797 {
798 SCHEME_OPE_RANGE
799 } else {
800 SCHEME_HMAC_EQ
801 };
802 let key: [u8; 32] = *kek.derive_column_key(c.id);
803 (c.id, (key, scheme))
804 })
805 .collect()
806 }
807 #[cfg(not(feature = "encryption"))]
808 {
809 let _ = (kek, schema);
810 HashMap::new()
811 }
812}
813
814pub(crate) struct SharedCtx {
819 pub epoch: Arc<EpochAuthority>,
820 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
821 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
822 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
823 pub kek: Option<Arc<Kek>>,
824 pub commit_lock: Arc<parking_lot::Mutex<()>>,
830 pub shared: Option<SharedWalCtx>,
834 pub table_name: Option<String>,
837 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
840}
841
842#[derive(Clone)]
848pub(crate) struct SharedWalCtx {
849 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
850 pub group: Arc<GroupCommit>,
851 pub poisoned: Arc<AtomicBool>,
852 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
853}
854
855enum WalSink {
858 Private(Wal),
859 Shared(SharedWalCtx),
860}
861
862impl SharedCtx {
863 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
867 let n_shards = if cache_dir.is_some() {
871 1
872 } else {
873 crate::cache::CACHE_SHARDS
874 };
875 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
876 let page_cache = if let Some(d) = cache_dir {
877 Arc::new(crate::cache::Sharded::new(1, || {
878 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
879 }))
880 } else {
881 Arc::new(crate::cache::Sharded::new(n_shards, || {
882 crate::cache::PageCache::new(per_shard)
883 }))
884 };
885 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
886 let decoded_cache = Arc::new(crate::cache::Sharded::new(
887 crate::cache::CACHE_SHARDS,
888 || crate::cache::DecodedPageCache::new(decoded_per_shard),
889 ));
890 Self {
891 epoch: Arc::new(EpochAuthority::new(0)),
892 page_cache,
893 decoded_cache,
894 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
895 kek,
896 commit_lock: Arc::new(parking_lot::Mutex::new(())),
897 shared: None,
898 table_name: None,
899 auth: None,
900 }
901 }
902}
903
904fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
908 use crate::query::Condition;
909 match c {
910 Condition::Pk(_)
912 | Condition::BitmapEq { .. }
913 | Condition::BitmapIn { .. }
914 | Condition::BytesPrefix { .. }
915 | Condition::IsNull { .. }
916 | Condition::IsNotNull { .. } => 0,
917 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
919 1
920 }
921 Condition::FmContains { .. }
923 | Condition::FmContainsAll { .. }
924 | Condition::Ann { .. }
925 | Condition::SparseMatch { .. } => 2,
926 }
927}
928
929impl Table {
930 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
931 let dir = dir.as_ref().to_path_buf();
932 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
933 Self::create_in(&dir, schema, table_id, ctx)
934 }
935
936 #[cfg(feature = "encryption")]
947 pub fn create_encrypted(
948 dir: impl AsRef<Path>,
949 schema: Schema,
950 table_id: u64,
951 passphrase: &str,
952 ) -> Result<Self> {
953 let dir = dir.as_ref();
954 std::fs::create_dir_all(dir.join(META_DIR))?;
955 let salt = crate::encryption::random_salt();
956 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
957 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
958 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
959 Self::create_in(dir, schema, table_id, ctx)
960 }
961
962 #[cfg(feature = "encryption")]
967 pub fn create_with_key(
968 dir: impl AsRef<Path>,
969 schema: Schema,
970 table_id: u64,
971 key: &[u8],
972 ) -> Result<Self> {
973 let dir = dir.as_ref();
974 std::fs::create_dir_all(dir.join(META_DIR))?;
975 let salt = crate::encryption::random_salt();
976 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
977 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
978 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
979 Self::create_in(dir, schema, table_id, ctx)
980 }
981
982 #[cfg(feature = "encryption")]
984 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
985 let dir = dir.as_ref();
986 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
987 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
988 MongrelError::NotFound(format!(
989 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
990 salt_path
991 ))
992 })?;
993 if salt_bytes.len() != crate::encryption::SALT_LEN {
994 return Err(MongrelError::InvalidArgument(format!(
995 "salt file is {} bytes, expected {}",
996 salt_bytes.len(),
997 crate::encryption::SALT_LEN
998 )));
999 }
1000 let mut salt = [0u8; crate::encryption::SALT_LEN];
1001 salt.copy_from_slice(&salt_bytes);
1002 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
1003 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1004 Self::open_in(dir, ctx)
1005 }
1006
1007 pub(crate) fn create_in(
1008 dir: impl AsRef<Path>,
1009 schema: Schema,
1010 table_id: u64,
1011 ctx: SharedCtx,
1012 ) -> Result<Self> {
1013 schema.validate_auto_increment()?;
1014 schema.validate_defaults()?;
1015 let dir = dir.as_ref().to_path_buf();
1016 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
1017 write_schema(&dir, &schema)?;
1018 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
1019 let (wal, current_txn_id) = match ctx.shared.clone() {
1022 Some(s) => (WalSink::Shared(s), 0),
1023 None => {
1024 std::fs::create_dir_all(dir.join(WAL_DIR))?;
1025 let mut w = if let Some(ref dk) = wal_dek {
1026 Wal::create_with_cipher(
1027 dir.join(WAL_DIR).join("seg-000000.wal"),
1028 Epoch(0),
1029 Some(make_cipher(dk)),
1030 0,
1031 )?
1032 } else {
1033 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1034 };
1035 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1036 (WalSink::Private(w), 1)
1037 }
1038 };
1039 let mut manifest = Manifest::new(table_id, schema.schema_id);
1040 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1045 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1046 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1047 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1048 let auto_inc = resolve_auto_inc(&schema);
1049 let rcache_dir = dir.join(RCACHE_DIR);
1050 Ok(Self {
1051 dir,
1052 table_id,
1053 name: ctx.table_name.unwrap_or_default(),
1054 auth: ctx.auth,
1055 wal,
1056 memtable: Memtable::new(),
1057 mutable_run: MutableRun::new(),
1058 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1059 compaction_zstd_level: 3,
1060 allocator: RowIdAllocator::new(0),
1061 epoch: ctx.epoch,
1062 persisted_epoch: 0,
1063 schema,
1064 hot: HotIndex::new(),
1065 kek: ctx.kek,
1066 column_keys,
1067 run_refs: Vec::new(),
1068 retiring: Vec::new(),
1069 next_run_id: 1,
1070 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1071 current_txn_id,
1072 bitmap,
1073 ann,
1074 fm,
1075 sparse,
1076 minhash,
1077 learned_range: HashMap::new(),
1078 pk_by_row: HashMap::new(),
1079 pinned: BTreeMap::new(),
1080 live_count: 0,
1081 reservoir: crate::reservoir::Reservoir::default(),
1082 reservoir_complete: true,
1083 had_deletes: false,
1084 agg_cache: HashMap::new(),
1085 global_idx_epoch: 0,
1086 indexes_complete: true,
1087 index_build_policy: IndexBuildPolicy::default(),
1088 pk_by_row_complete: false,
1089 flushed_epoch: 0,
1090 page_cache: ctx.page_cache,
1091 decoded_cache: ctx.decoded_cache,
1092 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1093 snapshots: ctx.snapshots,
1094 commit_lock: ctx.commit_lock,
1095 result_cache: Arc::new(parking_lot::Mutex::new(
1096 ResultCache::new()
1097 .with_dir(rcache_dir)
1098 .with_cache_dek(cache_dek.clone()),
1099 )),
1100 pending_delete_rids: roaring::RoaringBitmap::new(),
1101 pending_put_cols: std::collections::HashSet::new(),
1102 pending_rows: Vec::new(),
1103 pending_rows_auto_inc: Vec::new(),
1104 pending_dels: Vec::new(),
1105 pending_truncate: None,
1106 wal_dek,
1107 auto_inc,
1108 })
1109 }
1110
1111 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1115 let dir = dir.as_ref();
1116 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1117 Self::open_in(dir, ctx)
1118 }
1119
1120 #[cfg(feature = "encryption")]
1123 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1124 let dir = dir.as_ref();
1125 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1126 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1127 MongrelError::NotFound(format!(
1128 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1129 salt_path
1130 ))
1131 })?;
1132 let salt_len = crate::encryption::SALT_LEN;
1133 if salt_bytes.len() != salt_len {
1134 return Err(MongrelError::InvalidArgument(format!(
1135 "encryption salt is {} bytes, expected {salt_len}",
1136 salt_bytes.len()
1137 )));
1138 }
1139 let mut salt = [0u8; 16];
1140 salt.copy_from_slice(&salt_bytes);
1141 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1142 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1143 let t = Self::open_in(dir, ctx)?;
1144 Ok(t)
1145 }
1146
1147 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1148 let dir = dir.as_ref().to_path_buf();
1149 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1150 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1151 let schema: Schema = read_schema(&dir)?;
1152 let replay_epoch = Epoch(manifest.current_epoch);
1153 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1154 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1158 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1159 None => {
1160 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1161 let replayed = match &active {
1163 Some(path) => {
1164 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1165 crate::wal::replay_with_cipher(path, cipher)?
1166 }
1167 None => Vec::new(),
1168 };
1169 let mut w = match &active {
1170 Some(path) => Wal::create_with_cipher(
1171 path,
1172 replay_epoch,
1173 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1174 0,
1175 )?,
1176 None => Wal::create_with_cipher(
1177 dir.join(WAL_DIR).join("seg-000000.wal"),
1178 replay_epoch,
1179 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1180 0,
1181 )?,
1182 };
1183 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1184 (WalSink::Private(w), replayed, 1)
1185 }
1186 };
1187
1188 let mut memtable = Memtable::new();
1189 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1190 let persisted_epoch = manifest.current_epoch;
1191 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1198 s.next = manifest.auto_inc_next;
1199 s.seeded = manifest.auto_inc_next > 0;
1200 s
1201 });
1202
1203 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1210 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1211 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1212 std::collections::BTreeMap::new();
1213 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1214 let mut saw_delete = false;
1215 for record in replayed {
1216 let txn_id = record.txn_id;
1217 match record.op {
1218 Op::Put { rows, .. } => {
1219 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1220 for row in &rows {
1221 allocator.advance_to(row.row_id);
1222 if let Some(ai) = auto_inc.as_mut() {
1223 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1224 if *n + 1 > ai.next {
1225 ai.next = *n + 1;
1226 }
1227 }
1228 }
1229 }
1230 staged_puts.entry(txn_id).or_default().extend(rows);
1231 }
1232 Op::Delete { row_ids, .. } => {
1233 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1234 }
1235 Op::TxnCommit { epoch, .. } => {
1236 let commit_epoch = Epoch(epoch);
1237 if let Some(puts) = staged_puts.remove(&txn_id) {
1238 for row in &puts {
1239 memtable.upsert(row.clone());
1240 }
1241 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1242 }
1243 if let Some(dels) = staged_deletes.remove(&txn_id) {
1244 saw_delete = true;
1245 for rid in dels {
1246 memtable.tombstone(rid, commit_epoch);
1247 replayed_deletes.push((rid, commit_epoch));
1248 }
1249 }
1250 }
1251 Op::TxnAbort => {
1252 staged_puts.remove(&txn_id);
1253 staged_deletes.remove(&txn_id);
1254 }
1255 Op::TruncateTable { .. }
1256 | Op::ExternalTableState { .. }
1257 | Op::Flush { .. }
1258 | Op::Ddl(_) => {}
1259 }
1260 }
1261
1262 let rcache_dir = dir.join(RCACHE_DIR);
1263 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1264 let mut db = Self {
1265 dir,
1266 table_id: manifest.table_id,
1267 name: ctx.table_name.unwrap_or_default(),
1268 auth: ctx.auth,
1269 wal,
1270 memtable,
1271 mutable_run: MutableRun::new(),
1272 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1273 compaction_zstd_level: 3,
1274 allocator,
1275 epoch: ctx.epoch,
1276 persisted_epoch,
1277 schema,
1278 hot: HotIndex::new(),
1279 kek: ctx.kek,
1280 column_keys,
1281 run_refs: manifest.runs.clone(),
1282 retiring: manifest.retiring.clone(),
1283 next_run_id: manifest
1284 .runs
1285 .iter()
1286 .map(|r| r.run_id as u64 + 1)
1287 .max()
1288 .unwrap_or(1),
1289 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1290 current_txn_id,
1291 bitmap: HashMap::new(),
1292 ann: HashMap::new(),
1293 fm: HashMap::new(),
1294 sparse: HashMap::new(),
1295 minhash: HashMap::new(),
1296 learned_range: HashMap::new(),
1297 pk_by_row: HashMap::new(),
1298 pinned: BTreeMap::new(),
1299 live_count: manifest.live_count,
1300 reservoir: crate::reservoir::Reservoir::default(),
1301 reservoir_complete: false,
1302 had_deletes: saw_delete,
1303 agg_cache: HashMap::new(),
1304 global_idx_epoch: manifest.global_idx_epoch,
1305 indexes_complete: true,
1306 index_build_policy: IndexBuildPolicy::default(),
1307 pk_by_row_complete: false,
1308 flushed_epoch: manifest.flushed_epoch,
1309 page_cache: ctx.page_cache,
1310 decoded_cache: ctx.decoded_cache,
1311 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1312 snapshots: ctx.snapshots,
1313 commit_lock: ctx.commit_lock,
1314 result_cache: Arc::new(parking_lot::Mutex::new(
1315 ResultCache::new()
1316 .with_dir(rcache_dir)
1317 .with_cache_dek(cache_dek.clone()),
1318 )),
1319 pending_delete_rids: roaring::RoaringBitmap::new(),
1320 pending_put_cols: std::collections::HashSet::new(),
1321 pending_rows: Vec::new(),
1322 pending_rows_auto_inc: Vec::new(),
1323 pending_dels: Vec::new(),
1324 pending_truncate: None,
1325 wal_dek,
1326 auto_inc,
1327 };
1328
1329 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1332
1333 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1338 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1339 c.epoch_built == manifest.global_idx_epoch
1340 && manifest.global_idx_epoch > 0
1341 && manifest
1342 .runs
1343 .iter()
1344 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1345 });
1346 if let Some(loaded) = checkpoint {
1347 if checkpoint_valid {
1348 db.hot = loaded.hot;
1349 db.bitmap = loaded.bitmap;
1350 db.ann = loaded.ann;
1351 db.fm = loaded.fm;
1352 db.sparse = loaded.sparse;
1353 db.minhash = loaded.minhash;
1354 db.learned_range = loaded.learned_range;
1355 }
1358 }
1359 if !checkpoint_valid {
1360 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1361 db.bitmap = bitmap;
1362 db.ann = ann;
1363 db.fm = fm;
1364 db.sparse = sparse;
1365 db.minhash = minhash;
1366 db.rebuild_indexes_from_runs()?;
1367 db.build_learned_ranges()?;
1368 }
1369
1370 for (epoch, group) in replayed_puts {
1375 let (losers, winner_pks) = db.partition_pk_winners(&group);
1376 for (key, &row_id) in &winner_pks {
1377 if let Some(old_rid) = db.hot.get(key) {
1378 if old_rid != row_id {
1379 db.tombstone_row(old_rid, epoch, false);
1380 }
1381 }
1382 }
1383 for &loser_rid in &losers {
1384 db.tombstone_row(loser_rid, epoch, false);
1385 }
1386 for (key, row_id) in winner_pks {
1387 db.insert_hot_pk(key, row_id);
1388 }
1389 if db.schema.primary_key().is_none() {
1390 for r in &group {
1391 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1392 }
1393 }
1394 for r in &group {
1395 if !losers.contains(&r.row_id) {
1396 db.index_row(r);
1397 }
1398 }
1399 }
1400 for (rid, epoch) in &replayed_deletes {
1404 db.remove_hot_for_row(*rid, *epoch);
1405 }
1406
1407 db.result_cache.lock().load_persistent();
1414 Ok(db)
1415 }
1416
1417 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1423 if self.reservoir_complete {
1424 return Ok(());
1425 }
1426 self.rebuild_reservoir()?;
1427 self.reservoir_complete = true;
1428 Ok(())
1429 }
1430
1431 fn rebuild_reservoir(&mut self) -> Result<()> {
1434 let snap = self.snapshot();
1435 let rows = self.visible_rows(snap)?;
1436 self.reservoir.reset();
1437 for r in rows {
1438 self.reservoir.offer(r.row_id.0);
1439 }
1440 Ok(())
1441 }
1442
1443 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1444 self.hot = HotIndex::new();
1445 self.pk_by_row.clear();
1446 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1447 self.bitmap = bitmap;
1448 self.ann = ann;
1449 self.fm = fm;
1450 self.sparse = sparse;
1451 self.minhash = minhash;
1452 let snapshot = Epoch(u64::MAX);
1453 for rr in self.run_refs.clone() {
1454 let mut reader = self.open_reader(rr.run_id)?;
1455 for row in reader.visible_rows(snapshot)? {
1456 let tok_row = self.tokenized_for_indexes(&row);
1457 index_into(
1458 &self.schema,
1459 &tok_row,
1460 &mut self.hot,
1461 &mut self.bitmap,
1462 &mut self.ann,
1463 &mut self.fm,
1464 &mut self.sparse,
1465 &mut self.minhash,
1466 );
1467 }
1468 }
1469 for row in self.mutable_run.visible_versions(snapshot) {
1470 if row.deleted {
1471 self.remove_hot_for_row(row.row_id, snapshot);
1472 } else {
1473 self.index_row(&row);
1474 }
1475 }
1476 for row in self.memtable.visible_versions(snapshot) {
1477 if row.deleted {
1478 self.remove_hot_for_row(row.row_id, snapshot);
1479 } else {
1480 self.index_row(&row);
1481 }
1482 }
1483 self.refresh_pk_by_row_from_hot();
1484 Ok(())
1485 }
1486
1487 fn refresh_pk_by_row_from_hot(&mut self) {
1488 self.pk_by_row_complete = true;
1489 if self.schema.primary_key().is_none() {
1490 self.pk_by_row.clear();
1491 return;
1492 }
1493 self.pk_by_row = self
1499 .hot
1500 .entries()
1501 .into_iter()
1502 .map(|(key, row_id)| (row_id, key))
1503 .collect();
1504 }
1505
1506 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1507 if self.schema.primary_key().is_some() {
1508 self.pk_by_row.insert(row_id, key.clone());
1509 }
1510 self.hot.insert(key, row_id);
1511 }
1512
1513 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1517 self.learned_range.clear();
1518 if self.run_refs.len() != 1 {
1519 return Ok(());
1520 }
1521 let cols: Vec<u16> = self
1522 .schema
1523 .indexes
1524 .iter()
1525 .filter(|i| i.kind == IndexKind::LearnedRange)
1526 .map(|i| i.column_id)
1527 .collect();
1528 if cols.is_empty() {
1529 return Ok(());
1530 }
1531 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1532 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1533 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1534 _ => return Ok(()),
1535 };
1536 for cid in cols {
1537 let ty = self
1538 .schema
1539 .columns
1540 .iter()
1541 .find(|c| c.id == cid)
1542 .map(|c| c.ty.clone())
1543 .unwrap_or(TypeId::Int64);
1544 match ty {
1545 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1546 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1547 let pairs: Vec<(i64, u64)> = data
1548 .iter()
1549 .zip(row_ids.iter())
1550 .map(|(v, r)| (*v, *r))
1551 .collect();
1552 self.learned_range
1553 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1554 }
1555 }
1556 TypeId::Float64 => {
1557 if let columnar::NativeColumn::Float64 { data, .. } =
1558 reader.column_native(cid)?
1559 {
1560 let pairs: Vec<(f64, u64)> = data
1561 .iter()
1562 .zip(row_ids.iter())
1563 .map(|(v, r)| (*v, *r))
1564 .collect();
1565 self.learned_range
1566 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1567 }
1568 }
1569 _ => {}
1570 }
1571 }
1572 Ok(())
1573 }
1574
1575 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1582 if self.indexes_complete {
1583 crate::trace::QueryTrace::record(|t| {
1584 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1585 });
1586 return Ok(());
1587 }
1588 crate::trace::QueryTrace::record(|t| {
1589 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1590 });
1591 self.rebuild_indexes_from_runs()?;
1592 self.build_learned_ranges()?;
1593 self.indexes_complete = true;
1594 let epoch = self.current_epoch();
1595 self.checkpoint_indexes(epoch);
1596 Ok(())
1597 }
1598
1599 fn pending_epoch(&self) -> Epoch {
1600 Epoch(self.epoch.visible().0 + 1)
1601 }
1602
1603 fn is_shared(&self) -> bool {
1606 matches!(self.wal, WalSink::Shared(_))
1607 }
1608
1609 fn ensure_txn_id(&mut self) -> u64 {
1613 if self.current_txn_id == 0 {
1614 let id = match &self.wal {
1615 WalSink::Shared(s) => {
1616 let mut g = s.txn_ids.lock();
1617 let v = *g;
1618 *g = g.wrapping_add(1);
1619 v
1620 }
1621 WalSink::Private(_) => 1,
1622 };
1623 self.current_txn_id = id;
1624 }
1625 self.current_txn_id
1626 }
1627
1628 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1631 let txn_id = self.ensure_txn_id();
1632 let table_id = self.table_id;
1633 match &mut self.wal {
1634 WalSink::Private(w) => {
1635 w.append_txn(txn_id, op)?;
1636 }
1637 WalSink::Shared(s) => {
1638 s.wal.lock().append(txn_id, table_id, op)?;
1639 }
1640 }
1641 Ok(())
1642 }
1643
1644 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1655 match &self.auth {
1656 Some(checker) => checker.check(&self.name, perm),
1657 None => Ok(()),
1658 }
1659 }
1660 pub fn require_select(&self) -> Result<()> {
1665 self.require(crate::auth_state::RequiredPermission::Select)
1666 }
1667 fn require_insert(&self) -> Result<()> {
1668 self.require(crate::auth_state::RequiredPermission::Insert)
1669 }
1670 #[allow(dead_code)]
1674 fn require_update(&self) -> Result<()> {
1675 self.require(crate::auth_state::RequiredPermission::Update)
1676 }
1677 fn require_delete(&self) -> Result<()> {
1678 self.require(crate::auth_state::RequiredPermission::Delete)
1679 }
1680
1681 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1684 self.require_insert()?;
1685 Ok(self.put_returning(columns)?.0)
1686 }
1687
1688 pub fn put_returning(
1693 &mut self,
1694 mut columns: Vec<(u16, Value)>,
1695 ) -> Result<(RowId, Option<i64>)> {
1696 self.require_insert()?;
1697 let assigned = self.fill_auto_inc(&mut columns)?;
1698 self.apply_defaults(&mut columns)?;
1699 self.schema.validate_not_null(&columns)?;
1700 let row_id = if self.schema.clustered {
1705 self.derive_clustered_row_id(&columns)?
1706 } else {
1707 self.allocator.alloc()
1708 };
1709 let epoch = self.pending_epoch();
1710 let mut row = Row::new(row_id, epoch);
1711 for (col_id, val) in columns {
1712 row.columns.insert(col_id, val);
1713 }
1714 self.commit_rows(vec![row], assigned.is_some())?;
1715 Ok((row_id, assigned))
1716 }
1717
1718 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1721 self.require_insert()?;
1722 Ok(self
1723 .put_batch_returning(batch)?
1724 .into_iter()
1725 .map(|(r, _)| r)
1726 .collect())
1727 }
1728
1729 pub fn put_batch_returning(
1732 &mut self,
1733 batch: Vec<Vec<(u16, Value)>>,
1734 ) -> Result<Vec<(RowId, Option<i64>)>> {
1735 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1736 for mut cols in batch {
1737 let assigned = self.fill_auto_inc(&mut cols)?;
1738 self.apply_defaults(&mut cols)?;
1739 filled.push((cols, assigned));
1740 }
1741 for (cols, _) in &filled {
1742 self.schema.validate_not_null(cols)?;
1743 }
1744 let epoch = self.pending_epoch();
1745 let mut rows = Vec::with_capacity(filled.len());
1746 let mut ids = Vec::with_capacity(filled.len());
1747 for (cols, assigned) in filled {
1748 let row_id = if self.schema.clustered {
1749 self.derive_clustered_row_id(&cols)?
1750 } else {
1751 self.allocator.alloc()
1752 };
1753 let mut row = Row::new(row_id, epoch);
1754 for (c, v) in cols {
1755 row.columns.insert(c, v);
1756 }
1757 ids.push((row_id, assigned));
1758 rows.push(row);
1759 }
1760 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1761 self.commit_rows(rows, all_auto_generated)?;
1762 Ok(ids)
1763 }
1764
1765 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1771 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1772 return Ok(None);
1773 };
1774 let pos = columns.iter().position(|(c, _)| *c == cid);
1775 let assigned = match pos {
1776 Some(i) => match &columns[i].1 {
1777 Value::Null => {
1778 let next = self.alloc_auto_inc_value()?;
1779 columns[i].1 = Value::Int64(next);
1780 Some(next)
1781 }
1782 Value::Int64(n) => {
1783 self.advance_auto_inc_past(*n)?;
1784 None
1785 }
1786 other => {
1787 return Err(MongrelError::InvalidArgument(format!(
1788 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1789 other
1790 )))
1791 }
1792 },
1793 None => {
1794 let next = self.alloc_auto_inc_value()?;
1795 columns.push((cid, Value::Int64(next)));
1796 Some(next)
1797 }
1798 };
1799 Ok(assigned)
1800 }
1801
1802 pub fn apply_defaults(&self, columns: &mut Vec<(u16, Value)>) -> Result<()> {
1808 for col in &self.schema.columns {
1809 let Some(expr) = &col.default_value else {
1810 continue;
1811 };
1812 if col.flags.contains(ColumnFlags::AUTO_INCREMENT) {
1814 continue;
1815 }
1816 let pos = columns.iter().position(|(c, _)| *c == col.id);
1817 let needs_default = match pos {
1818 None => true,
1819 Some(i) => matches!(columns[i].1, Value::Null),
1820 };
1821 if !needs_default {
1822 continue;
1823 }
1824 let v = match expr {
1825 crate::schema::DefaultExpr::Static(v) => v.clone(),
1826 crate::schema::DefaultExpr::Now => Value::Bytes(iso_now_bytes()),
1827 crate::schema::DefaultExpr::Uuid => {
1828 let mut buf = [0u8; 16];
1829 getrandom::getrandom(&mut buf)
1830 .map_err(|e| MongrelError::Other(format!("UUID generation failed: {e}")))?;
1831 Value::Uuid(buf)
1832 }
1833 };
1834 match pos {
1835 None => columns.push((col.id, v)),
1836 Some(i) => columns[i].1 = v,
1837 }
1838 }
1839 Ok(())
1840 }
1841
1842 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1844 self.ensure_auto_inc_seeded()?;
1845 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1847 let v = ai.next;
1848 ai.next = ai.next.saturating_add(1);
1849 Ok(v)
1850 }
1851
1852 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1855 self.ensure_auto_inc_seeded()?;
1856 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1857 let floor = used.saturating_add(1).max(1);
1858 if ai.next < floor {
1859 ai.next = floor;
1860 }
1861 Ok(())
1862 }
1863
1864 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1869 let needs_seed = match self.auto_inc {
1870 Some(ai) => !ai.seeded,
1871 None => return Ok(()),
1872 };
1873 if !needs_seed {
1874 return Ok(());
1875 }
1876 if self.seed_empty_auto_inc() {
1877 return Ok(());
1878 }
1879 let cid = self
1880 .auto_inc
1881 .as_ref()
1882 .expect("auto-inc column present")
1883 .column_id;
1884 let max = self.scan_max_int64(cid)?;
1885 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1886 let floor = max.saturating_add(1).max(1);
1887 if ai.next < floor {
1888 ai.next = floor;
1889 }
1890 ai.seeded = true;
1891 Ok(())
1892 }
1893
1894 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1895 if n == 0 || self.auto_inc.is_none() {
1896 return Ok(None);
1897 }
1898 self.ensure_auto_inc_seeded()?;
1899 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1900 let start = ai.next;
1901 ai.next = ai.next.saturating_add(n as i64);
1902 Ok(Some(start))
1903 }
1904
1905 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1909 let mut max: i64 = 0;
1910 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1911 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1912 if *n > max {
1913 max = *n;
1914 }
1915 }
1916 }
1917 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1918 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1919 if *n > max {
1920 max = *n;
1921 }
1922 }
1923 }
1924 for rr in self.run_refs.clone() {
1925 let reader = self.open_reader(rr.run_id)?;
1926 if let Some(stats) = reader.column_page_stats(column_id) {
1927 for s in stats {
1928 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1929 if n > max {
1930 max = n;
1931 }
1932 }
1933 }
1934 } else if reader.has_column(column_id) {
1935 if let columnar::NativeColumn::Int64 { data, validity } =
1936 reader.column_native_shared(column_id)?
1937 {
1938 for (i, n) in data.iter().enumerate() {
1939 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1940 {
1941 max = *n;
1942 }
1943 }
1944 }
1945 }
1946 }
1947 Ok(max)
1948 }
1949
1950 fn seed_empty_auto_inc(&mut self) -> bool {
1951 let Some(ai) = self.auto_inc.as_mut() else {
1952 return false;
1953 };
1954 if ai.seeded || self.live_count != 0 {
1955 return false;
1956 }
1957 if ai.next < 1 {
1958 ai.next = 1;
1959 }
1960 ai.seeded = true;
1961 true
1962 }
1963
1964 fn advance_auto_inc_from_native_columns(
1965 &mut self,
1966 columns: &[(u16, columnar::NativeColumn)],
1967 n: usize,
1968 live_before: u64,
1969 ) -> Result<()> {
1970 let Some(ai) = self.auto_inc.as_mut() else {
1971 return Ok(());
1972 };
1973 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1974 return Ok(());
1975 };
1976 let columnar::NativeColumn::Int64 { data, validity } = col else {
1977 return Err(MongrelError::InvalidArgument(format!(
1978 "AUTO_INCREMENT column {} must be Int64",
1979 ai.column_id
1980 )));
1981 };
1982 let max = if native_int64_strictly_increasing(col, n) {
1983 data.get(n.saturating_sub(1)).copied()
1984 } else {
1985 data.iter()
1986 .take(n)
1987 .enumerate()
1988 .filter_map(|(i, v)| {
1989 if validity.is_empty() || columnar::validity_bit(validity, i) {
1990 Some(*v)
1991 } else {
1992 None
1993 }
1994 })
1995 .max()
1996 };
1997 if let Some(max) = max {
1998 let floor = max.saturating_add(1).max(1);
1999 if ai.next < floor {
2000 ai.next = floor;
2001 }
2002 if ai.seeded || live_before == 0 {
2003 ai.seeded = true;
2004 }
2005 }
2006 Ok(())
2007 }
2008
2009 fn fill_auto_inc_native_columns(
2010 &mut self,
2011 columns: &mut Vec<(u16, columnar::NativeColumn)>,
2012 n: usize,
2013 ) -> Result<()> {
2014 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
2015 return Ok(());
2016 };
2017 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
2018 if let Some(start) = self.alloc_auto_inc_range(n)? {
2019 columns.push((
2020 cid,
2021 columnar::NativeColumn::Int64 {
2022 data: (start..start.saturating_add(n as i64)).collect(),
2023 validity: vec![0xFF; n.div_ceil(8)],
2024 },
2025 ));
2026 }
2027 return Ok(());
2028 };
2029
2030 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
2031 return Err(MongrelError::InvalidArgument(format!(
2032 "AUTO_INCREMENT column {cid} must be Int64"
2033 )));
2034 };
2035 if data.len() < n {
2036 return Err(MongrelError::InvalidArgument(format!(
2037 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
2038 data.len()
2039 )));
2040 }
2041 if columnar::all_non_null(validity, n) {
2042 return Ok(());
2043 }
2044 if validity.iter().all(|b| *b == 0) {
2045 if let Some(start) = self.alloc_auto_inc_range(n)? {
2046 for (i, slot) in data.iter_mut().take(n).enumerate() {
2047 *slot = start.saturating_add(i as i64);
2048 }
2049 *validity = vec![0xFF; n.div_ceil(8)];
2050 }
2051 return Ok(());
2052 }
2053
2054 let new_validity = vec![0xFF; data.len().div_ceil(8)];
2055 for (i, slot) in data.iter_mut().enumerate().take(n) {
2056 if columnar::validity_bit(validity, i) {
2057 self.advance_auto_inc_past(*slot)?;
2058 } else {
2059 *slot = self.alloc_auto_inc_value()?;
2060 }
2061 }
2062 *validity = new_validity;
2063 Ok(())
2064 }
2065
2066 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2080 if self.auto_inc.is_none() {
2081 return Ok(None);
2082 }
2083 Ok(Some(self.alloc_auto_inc_value()?))
2084 }
2085
2086 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2092 let payload = bincode::serialize(&rows)?;
2093 self.wal_append_data(Op::Put {
2094 table_id: self.table_id,
2095 rows: payload,
2096 })?;
2097 if self.is_shared() {
2098 self.pending_rows_auto_inc
2099 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2100 self.pending_rows.extend(rows);
2101 } else {
2102 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2103 }
2104 Ok(())
2105 }
2106
2107 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2112 self.apply_put_rows_inner(rows, true)
2113 }
2114
2115 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2116 if check_existing_pk {
2117 self.ensure_indexes_complete()?;
2118 }
2119 if rows.len() == 1 {
2123 let row = rows.into_iter().next().expect("len checked");
2124 return self.apply_put_row_single(row, check_existing_pk);
2125 }
2126 let pk_id = self.schema.primary_key().map(|c| c.id);
2143 let probe = match pk_id {
2144 Some(pid) => {
2145 check_existing_pk
2146 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2147 }
2148 None => false,
2149 };
2150 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2153 for r in rows {
2154 for &cid in r.columns.keys() {
2155 self.pending_put_cols.insert(cid);
2156 }
2157 match pk_id {
2158 Some(pid) if probe || maintain_pk_by_row => {
2159 if let Some(pk_val) = r.columns.get(&pid) {
2160 let key = self.index_lookup_key(pid, pk_val);
2161 if probe {
2162 if let Some(old_rid) = self.hot.get(&key) {
2163 if old_rid != r.row_id {
2164 self.tombstone_row(old_rid, r.committed_epoch, true);
2165 }
2166 }
2167 }
2168 if maintain_pk_by_row {
2169 self.pk_by_row.insert(r.row_id, key);
2170 }
2171 }
2172 }
2173 Some(_) => {}
2174 None => {
2175 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2176 }
2177 }
2178 self.index_row(&r);
2179 self.reservoir.offer(r.row_id.0);
2180 self.memtable.upsert(r);
2181 self.live_count = self.live_count.saturating_add(1);
2184 }
2185 Ok(())
2186 }
2187
2188 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2192 for &cid in row.columns.keys() {
2193 self.pending_put_cols.insert(cid);
2194 }
2195 let epoch = row.committed_epoch;
2196 if let Some(pk_col) = self.schema.primary_key() {
2197 let pk_id = pk_col.id;
2198 if let Some(pk_val) = row.columns.get(&pk_id) {
2199 let maintain_pk_by_row = self.pk_by_row_complete;
2203 if check_existing_pk || maintain_pk_by_row {
2204 let key = self.index_lookup_key(pk_id, pk_val);
2205 if check_existing_pk {
2206 if let Some(old_rid) = self.hot.get(&key) {
2207 if old_rid != row.row_id {
2208 self.tombstone_row(old_rid, epoch, true);
2209 }
2210 }
2211 }
2212 if maintain_pk_by_row {
2213 self.pk_by_row.insert(row.row_id, key);
2214 }
2215 }
2216 }
2217 } else {
2218 self.hot
2219 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2220 }
2221 self.index_row(&row);
2222 self.reservoir.offer(row.row_id.0);
2223 self.memtable.upsert(row);
2224 self.live_count = self.live_count.saturating_add(1);
2225 Ok(())
2226 }
2227
2228 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2231 self.allocator.alloc()
2232 }
2233
2234 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2240 let pk = self.schema.primary_key().ok_or_else(|| {
2241 MongrelError::Schema("clustered table requires a single-column primary key".into())
2242 })?;
2243 let pk_val = columns
2244 .iter()
2245 .find(|(id, _)| *id == pk.id)
2246 .map(|(_, v)| v)
2247 .ok_or_else(|| {
2248 MongrelError::Schema(format!(
2249 "clustered table missing primary key column {} ({})",
2250 pk.id, pk.name
2251 ))
2252 })?;
2253 let key_bytes = pk_val.encode_key();
2254 let mut hash: u64 = 0xcbf29ce484222325;
2256 for &b in &key_bytes {
2257 hash ^= b as u64;
2258 hash = hash.wrapping_mul(0x100000001b3);
2259 }
2260 Ok(RowId(hash.max(1)))
2263 }
2264
2265 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2273 self.ensure_indexes_complete()?;
2274 let n = rows.len();
2275 for r in rows {
2276 for &cid in r.columns.keys() {
2277 self.pending_put_cols.insert(cid);
2278 }
2279 }
2280 let (losers, winner_pks) = self.partition_pk_winners(rows);
2281 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2282 for (key, &row_id) in &winner_pks {
2284 if let Some(old_rid) = self.hot.get(key) {
2285 if old_rid != row_id {
2286 self.tombstone_row(old_rid, epoch, true);
2287 }
2288 }
2289 }
2290 for &loser_rid in &losers {
2293 self.tombstone_row(loser_rid, epoch, false);
2294 }
2295 for (key, row_id) in winner_pks {
2297 self.insert_hot_pk(key, row_id);
2298 }
2299 if self.schema.primary_key().is_none() {
2300 for r in rows {
2301 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2302 }
2303 }
2304 for r in rows {
2305 self.allocator.advance_to(r.row_id);
2306 if !losers.contains(&r.row_id) {
2307 self.index_row(r);
2308 }
2309 }
2310 for r in rows {
2311 if !losers.contains(&r.row_id) {
2312 self.reservoir.offer(r.row_id.0);
2313 }
2314 }
2315 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2316 Ok(())
2317 }
2318
2319 pub(crate) fn recover_apply(
2324 &mut self,
2325 rows: Vec<Row>,
2326 deletes: Vec<(RowId, Epoch)>,
2327 ) -> Result<()> {
2328 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2332 std::collections::BTreeMap::new();
2333 for row in rows {
2334 self.allocator.advance_to(row.row_id);
2335 if let Some(ai) = self.auto_inc.as_mut() {
2340 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2341 if *n + 1 > ai.next {
2342 ai.next = *n + 1;
2343 }
2344 }
2345 }
2346 by_epoch.entry(row.committed_epoch).or_default().push(row);
2347 }
2348 for (epoch, group) in by_epoch {
2349 let (losers, winner_pks) = self.partition_pk_winners(&group);
2350 for (key, &row_id) in &winner_pks {
2352 if let Some(old_rid) = self.hot.get(key) {
2353 if old_rid != row_id {
2354 self.tombstone_row(old_rid, epoch, false);
2355 }
2356 }
2357 }
2358 for (key, row_id) in winner_pks {
2359 self.insert_hot_pk(key, row_id);
2360 }
2361 if self.schema.primary_key().is_none() {
2362 for r in &group {
2363 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2364 }
2365 }
2366 for r in &group {
2367 if !losers.contains(&r.row_id) {
2368 self.memtable.upsert(r.clone());
2369 self.index_row(r);
2370 }
2371 }
2372 }
2373 for (rid, epoch) in deletes {
2374 self.memtable.tombstone(rid, epoch);
2375 self.remove_hot_for_row(rid, epoch);
2376 }
2377 self.reservoir_complete = false;
2380 Ok(())
2381 }
2382
2383 pub(crate) fn flushed_epoch(&self) -> u64 {
2385 self.flushed_epoch
2386 }
2387
2388 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2390 self.schema.validate_not_null(cells)
2391 }
2392
2393 fn validate_columns_not_null(
2397 &self,
2398 columns: &[(u16, columnar::NativeColumn)],
2399 n: usize,
2400 ) -> Result<()> {
2401 let by_id: HashMap<u16, &columnar::NativeColumn> =
2402 columns.iter().map(|(id, c)| (*id, c)).collect();
2403 for col in &self.schema.columns {
2404 if col.flags.contains(ColumnFlags::NULLABLE) {
2405 continue;
2406 }
2407 match by_id.get(&col.id) {
2408 None => {
2409 return Err(MongrelError::InvalidArgument(format!(
2410 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2411 col.name, col.id
2412 )));
2413 }
2414 Some(c) => {
2415 if c.null_count(n) != 0 {
2416 return Err(MongrelError::InvalidArgument(format!(
2417 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2418 col.name, col.id
2419 )));
2420 }
2421 }
2422 }
2423 }
2424 Ok(())
2425 }
2426
2427 fn bulk_pk_winner_indices(
2432 &self,
2433 columns: &[(u16, columnar::NativeColumn)],
2434 n: usize,
2435 ) -> Option<Vec<usize>> {
2436 let pk_col = self.schema.primary_key()?;
2437 let pk_id = pk_col.id;
2438 let pk_ty = pk_col.ty.clone();
2439 let by_id: HashMap<u16, &columnar::NativeColumn> =
2440 columns.iter().map(|(id, c)| (*id, c)).collect();
2441 let pk_native = by_id.get(&pk_id)?;
2442 if native_int64_strictly_increasing(pk_native, n) {
2443 return None;
2444 }
2445 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2447 let mut null_pk_rows: Vec<usize> = Vec::new();
2448 for i in 0..n {
2449 match bulk_index_key(&self.column_keys, pk_id, pk_ty.clone(), pk_native, i) {
2450 Some(key) => {
2451 last.insert(key, i);
2452 }
2453 None => null_pk_rows.push(i),
2454 }
2455 }
2456 let mut winners: HashSet<usize> = last.values().copied().collect();
2457 for i in null_pk_rows {
2458 winners.insert(i);
2459 }
2460 Some((0..n).filter(|i| winners.contains(i)).collect())
2461 }
2462
2463 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2465 self.require_delete()?;
2466 let epoch = self.pending_epoch();
2467 self.wal_append_data(Op::Delete {
2468 table_id: self.table_id,
2469 row_ids: vec![row_id],
2470 })?;
2471 if self.is_shared() {
2472 self.pending_dels.push(row_id);
2473 } else {
2474 self.apply_delete(row_id, epoch);
2475 }
2476 Ok(())
2477 }
2478
2479 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2480 let pre = self.get(row_id, self.snapshot());
2481 self.delete(row_id)?;
2482 Ok(pre.map(|row| {
2483 let mut columns: Vec<_> = row.columns.into_iter().collect();
2484 columns.sort_by_key(|(id, _)| *id);
2485 OwnedRow { columns }
2486 }))
2487 }
2488
2489 pub fn truncate(&mut self) -> Result<()> {
2491 self.require_delete()?;
2492 let epoch = self.pending_epoch();
2493 self.wal_append_data(Op::TruncateTable {
2494 table_id: self.table_id,
2495 })?;
2496 self.pending_rows.clear();
2497 self.pending_rows_auto_inc.clear();
2498 self.pending_dels.clear();
2499 self.pending_truncate = Some(epoch);
2500 Ok(())
2501 }
2502
2503 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2505 for rr in std::mem::take(&mut self.run_refs) {
2506 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2507 }
2508 for r in std::mem::take(&mut self.retiring) {
2509 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2510 }
2511 self.memtable = Memtable::new();
2512 self.mutable_run = MutableRun::new();
2513 self.hot = HotIndex::new();
2514 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2515 self.bitmap = bitmap;
2516 self.ann = ann;
2517 self.fm = fm;
2518 self.sparse = sparse;
2519 self.minhash = minhash;
2520 self.learned_range.clear();
2521 self.pk_by_row.clear();
2522 self.pk_by_row_complete = false;
2523 self.live_count = 0;
2524 self.reservoir = crate::reservoir::Reservoir::default();
2525 self.reservoir_complete = true;
2526 self.had_deletes = true;
2527 self.agg_cache.clear();
2528 self.global_idx_epoch = 0;
2529 self.indexes_complete = true;
2530 self.pending_delete_rids.clear();
2531 self.pending_put_cols.clear();
2532 self.pending_rows.clear();
2533 self.pending_rows_auto_inc.clear();
2534 self.pending_dels.clear();
2535 self.clear_result_cache();
2536 self.invalidate_index_checkpoint();
2537 Ok(())
2538 }
2539
2540 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2543 self.remove_hot_for_row(row_id, epoch);
2544 self.tombstone_row(row_id, epoch, true);
2545 }
2546
2547 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2551 let tombstone = Row {
2552 row_id,
2553 committed_epoch: epoch,
2554 columns: std::collections::HashMap::new(),
2555 deleted: true,
2556 };
2557 self.memtable.upsert(tombstone);
2558 self.pk_by_row.remove(&row_id);
2559 if adjust_live_count {
2560 self.live_count = self.live_count.saturating_sub(1);
2561 }
2562 self.pending_delete_rids.insert(row_id.0 as u32);
2564 self.had_deletes = true;
2567 self.agg_cache.clear();
2568 }
2569
2570 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2574 let Some(pk_col) = self.schema.primary_key() else {
2575 return;
2576 };
2577 if self.pk_by_row_complete {
2580 if let Some(key) = self.pk_by_row.remove(&row_id) {
2581 if self.hot.get(&key) == Some(row_id) {
2582 self.hot.remove(&key);
2583 }
2584 }
2585 return;
2586 }
2587 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2606 if self.indexes_complete {
2607 let pk_val = self
2608 .memtable
2609 .get_version(row_id, lookup_epoch)
2610 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2611 .or_else(|| {
2612 self.mutable_run
2613 .get_version(row_id, lookup_epoch)
2614 .filter(|(_, r)| !r.deleted)
2615 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2616 })
2617 .or_else(|| {
2618 self.run_refs.iter().find_map(|rr| {
2619 let mut reader = self.open_reader(rr.run_id).ok()?;
2620 let (_, deleted, val) = reader
2621 .get_version_column(row_id, lookup_epoch, pk_col.id)
2622 .ok()??;
2623 if deleted {
2624 return None;
2625 }
2626 val
2627 })
2628 });
2629 if let Some(pk_val) = pk_val {
2630 let key = self.index_lookup_key(pk_col.id, &pk_val);
2631 if self.hot.get(&key) == Some(row_id) {
2632 self.hot.remove(&key);
2633 }
2634 return;
2635 }
2636 }
2637 self.refresh_pk_by_row_from_hot();
2642 if let Some(key) = self.pk_by_row.remove(&row_id) {
2643 if self.hot.get(&key) == Some(row_id) {
2644 self.hot.remove(&key);
2645 }
2646 }
2647 }
2648
2649 fn partition_pk_winners(
2654 &self,
2655 rows: &[Row],
2656 ) -> (
2657 std::collections::HashSet<RowId>,
2658 std::collections::HashMap<Vec<u8>, RowId>,
2659 ) {
2660 let mut losers = std::collections::HashSet::new();
2661 let Some(pk_col) = self.schema.primary_key() else {
2662 return (losers, std::collections::HashMap::new());
2663 };
2664 let pk_id = pk_col.id;
2665 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2666 std::collections::HashMap::new();
2667 for r in rows {
2668 let Some(pk_val) = r.columns.get(&pk_id) else {
2669 continue;
2670 };
2671 let key = self.index_lookup_key(pk_id, pk_val);
2672 if let Some(&old_rid) = winners.get(&key) {
2673 losers.insert(old_rid);
2674 }
2675 winners.insert(key, r.row_id);
2676 }
2677 (losers, winners)
2678 }
2679
2680 fn index_row(&mut self, row: &Row) {
2681 if row.deleted {
2682 return;
2683 }
2684 let any_predicate = self
2692 .schema
2693 .indexes
2694 .iter()
2695 .any(|idx| idx.predicate.is_some());
2696 if any_predicate {
2697 let columns_map: HashMap<u16, &Value> =
2698 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2699 let name_to_id: HashMap<&str, u16> = self
2700 .schema
2701 .columns
2702 .iter()
2703 .map(|c| (c.name.as_str(), c.id))
2704 .collect();
2705 for idx in &self.schema.indexes {
2706 if let Some(pred) = &idx.predicate {
2707 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2708 continue; }
2710 }
2711 index_into_single(
2713 idx,
2714 &self.schema,
2715 row,
2716 &mut self.hot,
2717 &mut self.bitmap,
2718 &mut self.ann,
2719 &mut self.fm,
2720 &mut self.sparse,
2721 &mut self.minhash,
2722 );
2723 }
2724 return;
2725 }
2726 if self.column_keys.is_empty() {
2730 index_into(
2731 &self.schema,
2732 row,
2733 &mut self.hot,
2734 &mut self.bitmap,
2735 &mut self.ann,
2736 &mut self.fm,
2737 &mut self.sparse,
2738 &mut self.minhash,
2739 );
2740 return;
2741 }
2742 let effective_row = self.tokenized_for_indexes(row);
2743 index_into(
2744 &self.schema,
2745 &effective_row,
2746 &mut self.hot,
2747 &mut self.bitmap,
2748 &mut self.ann,
2749 &mut self.fm,
2750 &mut self.sparse,
2751 &mut self.minhash,
2752 );
2753 }
2754
2755 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2761 if self.column_keys.is_empty() {
2762 return row.clone();
2763 }
2764 #[cfg(feature = "encryption")]
2765 {
2766 use crate::encryption::SCHEME_HMAC_EQ;
2767 let mut tok = row.clone();
2768 for (&cid, &(_, scheme)) in &self.column_keys {
2769 if scheme != SCHEME_HMAC_EQ {
2770 continue;
2771 }
2772 if let Some(v) = tok.columns.get(&cid).cloned() {
2773 if let Some(t) = self.tokenize_value(cid, &v) {
2774 tok.columns.insert(cid, t);
2775 }
2776 }
2777 }
2778 tok
2779 }
2780 #[cfg(not(feature = "encryption"))]
2781 {
2782 row.clone()
2783 }
2784 }
2785
2786 pub fn commit(&mut self) -> Result<Epoch> {
2791 if self.is_shared() {
2792 self.commit_shared()
2793 } else {
2794 self.commit_private()
2795 }
2796 }
2797
2798 fn commit_private(&mut self) -> Result<Epoch> {
2800 let commit_lock = Arc::clone(&self.commit_lock);
2804 let _g = commit_lock.lock();
2805 let new_epoch = self.epoch.bump_assigned();
2806 let txn_id = self.current_txn_id;
2807 match &mut self.wal {
2811 WalSink::Private(w) => {
2812 w.append_txn(
2813 txn_id,
2814 Op::TxnCommit {
2815 epoch: new_epoch.0,
2816 added_runs: Vec::new(),
2817 },
2818 )?;
2819 w.sync()?;
2820 }
2821 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2822 }
2823 if let Some(epoch) = self.pending_truncate.take() {
2825 self.apply_truncate(epoch)?;
2826 }
2827 self.invalidate_pending_cache();
2828 self.persist_manifest(new_epoch)?;
2829 self.epoch.publish_in_order(new_epoch);
2833 self.current_txn_id += 1;
2834 Ok(new_epoch)
2835 }
2836
2837 fn commit_shared(&mut self) -> Result<Epoch> {
2843 use std::sync::atomic::Ordering;
2844 let s = match &self.wal {
2845 WalSink::Shared(s) => s.clone(),
2846 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2847 };
2848 if s.poisoned.load(Ordering::Relaxed) {
2849 return Err(MongrelError::Other(
2850 "database poisoned by fsync error".into(),
2851 ));
2852 }
2853 let commit_lock = Arc::clone(&self.commit_lock);
2860 let _g = commit_lock.lock();
2861 let txn_id = self.ensure_txn_id();
2864 let (new_epoch, commit_seq) = {
2865 let mut wal = s.wal.lock();
2866 let new_epoch = self.epoch.bump_assigned();
2867 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2868 (new_epoch, seq)
2869 };
2870 s.group
2871 .await_durable(&s.wal, commit_seq)
2872 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2873
2874 if self.pending_truncate.take().is_some() {
2877 self.apply_truncate(new_epoch)?;
2878 }
2879 let mut rows = std::mem::take(&mut self.pending_rows);
2880 if !rows.is_empty() {
2881 for r in &mut rows {
2882 r.committed_epoch = new_epoch;
2883 }
2884 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2885 let all_auto_generated =
2886 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2887 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2888 } else {
2889 self.pending_rows_auto_inc.clear();
2890 }
2891 let dels = std::mem::take(&mut self.pending_dels);
2892 for rid in dels {
2893 self.apply_delete(rid, new_epoch);
2894 }
2895
2896 self.invalidate_pending_cache();
2897 self.persist_manifest(new_epoch)?;
2898 self.epoch.publish_in_order(new_epoch);
2899 self.current_txn_id = 0;
2901 Ok(new_epoch)
2902 }
2903
2904 pub fn flush(&mut self) -> Result<Epoch> {
2912 self.ensure_indexes_complete()?;
2913 let epoch = self.commit()?;
2914 let rows = self.memtable.drain_sorted();
2915 if !rows.is_empty() {
2916 self.mutable_run.insert_many(rows);
2917 }
2918 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2919 self.spill_mutable_run(epoch)?;
2920 self.mark_flushed(epoch)?;
2924 self.persist_manifest(epoch)?;
2925 self.build_learned_ranges()?;
2926 self.checkpoint_indexes(epoch);
2929 }
2930 Ok(epoch)
2933 }
2934
2935 pub fn force_flush(&mut self) -> Result<Epoch> {
2944 let saved = self.mutable_run_spill_bytes;
2945 self.mutable_run_spill_bytes = 1;
2946 let result = self.flush();
2947 self.mutable_run_spill_bytes = saved;
2948 result
2949 }
2950
2951 pub fn close(&mut self) -> Result<()> {
2958 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
2959 self.force_flush()?;
2960 }
2961 Ok(())
2962 }
2963
2964 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2971 let op = Op::Flush {
2972 table_id: self.table_id,
2973 flushed_epoch: epoch.0,
2974 };
2975 match &mut self.wal {
2976 WalSink::Private(w) => {
2977 w.append_system(op)?;
2978 w.sync()?;
2979 }
2980 WalSink::Shared(s) => {
2981 s.wal.lock().append_system(op)?;
2986 }
2987 }
2988 self.flushed_epoch = epoch.0;
2989 if matches!(self.wal, WalSink::Private(_)) {
2990 self.rotate_wal(epoch)?;
2991 }
2992 Ok(())
2993 }
2994
2995 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2999 let rows = self.mutable_run.drain_sorted();
3000 if rows.is_empty() {
3001 return Ok(());
3002 }
3003 let run_id = self.next_run_id;
3004 self.next_run_id += 1;
3005 let path = self.run_path(run_id);
3006 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
3007 if let Some(kek) = &self.kek {
3008 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3009 }
3010 let header = writer.write(&path, &rows)?;
3011 self.run_refs.push(RunRef {
3012 run_id: run_id as u128,
3013 level: 0,
3014 epoch_created: epoch.0,
3015 row_count: header.row_count,
3016 });
3017 Ok(())
3018 }
3019
3020 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
3024 self.mutable_run_spill_bytes = bytes.max(1);
3025 }
3026
3027 pub fn set_compaction_zstd_level(&mut self, level: i32) {
3031 self.compaction_zstd_level = level;
3032 }
3033
3034 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
3038 self.result_cache.lock().set_max_bytes(max_bytes);
3039 }
3040
3041 pub(crate) fn clear_result_cache(&mut self) {
3045 self.result_cache.lock().clear();
3046 }
3047
3048 pub fn mutable_run_len(&self) -> usize {
3050 self.mutable_run.len()
3051 }
3052
3053 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
3056 self.mutable_run.drain_sorted()
3057 }
3058
3059 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
3064 let epoch = self.commit()?;
3065 let n = batch.len();
3066 if n == 0 {
3067 return Ok(epoch);
3068 }
3069 let live_before = self.live_count;
3070 self.spill_mutable_run(epoch)?;
3074 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3075 && self.indexes_complete
3076 && self.run_refs.is_empty()
3077 && self.memtable.is_empty()
3078 && self.mutable_run.is_empty();
3079 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3085 use rayon::prelude::*;
3086 self.schema
3087 .columns
3088 .par_iter()
3089 .map(|cdef| {
3090 (
3091 cdef.id,
3092 columnar::rows_to_native(cdef.ty.clone(), &batch, cdef.id),
3093 )
3094 })
3095 .collect::<Vec<_>>()
3096 };
3097 drop(batch);
3098 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3103 self.validate_columns_not_null(&user_columns, n)?;
3104 let winner_idx = self
3105 .bulk_pk_winner_indices(&user_columns, n)
3106 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
3107 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3108 match winner_idx.as_deref() {
3109 Some(idx) => {
3110 let compacted = user_columns
3111 .iter()
3112 .map(|(id, c)| (*id, c.gather(idx)))
3113 .collect();
3114 (compacted, idx.len())
3115 }
3116 None => (std::mem::take(&mut user_columns), n),
3117 };
3118 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3119 let first = self.allocator.alloc_range(write_n as u64).0;
3120 for rid in first..first + write_n as u64 {
3121 self.reservoir.offer(rid);
3122 }
3123 let run_id = self.next_run_id;
3124 self.next_run_id += 1;
3125 let path = self.run_path(run_id);
3126 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3127 .clean(true)
3128 .with_lz4()
3129 .with_native_endian();
3130 if let Some(kek) = &self.kek {
3131 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3132 }
3133 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3134 self.run_refs.push(RunRef {
3135 run_id: run_id as u128,
3136 level: 0,
3137 epoch_created: epoch.0,
3138 row_count: header.row_count,
3139 });
3140 self.live_count = self.live_count.saturating_add(write_n as u64);
3141 if eager_index_build {
3142 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3143 self.index_columns_bulk(&write_columns, &row_ids);
3144 self.indexes_complete = true;
3145 self.build_learned_ranges()?;
3146 } else {
3147 self.indexes_complete = false;
3148 }
3149 self.mark_flushed(epoch)?;
3150 self.persist_manifest(epoch)?;
3151 if eager_index_build {
3152 self.checkpoint_indexes(epoch);
3153 }
3154 self.clear_result_cache();
3155 Ok(epoch)
3156 }
3157
3158 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3161 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3162 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3163 let segment_no = segment
3166 .file_stem()
3167 .and_then(|s| s.to_str())
3168 .and_then(|s| s.strip_prefix("seg-"))
3169 .and_then(|s| s.parse::<u64>().ok())
3170 .unwrap_or(0);
3171 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3172 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3173 wal.sync()?;
3174 self.wal = WalSink::Private(wal);
3175 Ok(())
3176 }
3177
3178 pub(crate) fn invalidate_pending_cache(&mut self) {
3183 self.result_cache
3184 .lock()
3185 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3186 self.pending_delete_rids.clear();
3187 self.pending_put_cols.clear();
3188 }
3189
3190 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3191 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3192 m.current_epoch = epoch.0;
3193 m.next_row_id = self.allocator.current().0;
3194 m.runs = self.run_refs.clone();
3195 m.live_count = self.live_count;
3196 m.global_idx_epoch = self.global_idx_epoch;
3197 m.flushed_epoch = self.flushed_epoch;
3198 m.retiring = self.retiring.clone();
3199 m.auto_inc_next = match self.auto_inc {
3203 Some(ai) if ai.seeded => ai.next,
3204 _ => 0,
3205 };
3206 let meta_dek = self.manifest_meta_dek();
3207 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3208 Ok(())
3209 }
3210
3211 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3217 if !self.indexes_complete {
3220 return;
3221 }
3222 let snap = global_idx::IndexSnapshot {
3223 hot: &self.hot,
3224 bitmap: &self.bitmap,
3225 ann: &self.ann,
3226 fm: &self.fm,
3227 sparse: &self.sparse,
3228 minhash: &self.minhash,
3229 learned_range: &self.learned_range,
3230 };
3231 let idx_dek = self.idx_dek();
3233 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3234 .is_ok()
3235 {
3236 self.global_idx_epoch = epoch.0;
3237 let _ = self.persist_manifest(epoch);
3238 }
3239 }
3240
3241 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3244 self.global_idx_epoch = 0;
3245 global_idx::remove(&self.dir);
3246 let _ = self.persist_manifest(self.epoch.visible());
3247 }
3248
3249 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3252 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3253 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3254 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3255 best = Some((epoch, row));
3256 }
3257 }
3258 for rr in &self.run_refs {
3259 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3260 continue;
3261 };
3262 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3263 continue;
3264 };
3265 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3266 best = Some((epoch, row));
3267 }
3268 }
3269 match best {
3270 Some((_, r)) if r.deleted => None,
3271 Some((_, r)) => Some(r),
3272 None => None,
3273 }
3274 }
3275
3276 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3280 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3281 let mut fold = |row: Row| {
3282 best.entry(row.row_id.0)
3283 .and_modify(|e| {
3284 if row.committed_epoch > e.0 {
3285 *e = (row.committed_epoch, row.clone());
3286 }
3287 })
3288 .or_insert_with(|| (row.committed_epoch, row));
3289 };
3290 for row in self.memtable.visible_versions(snapshot.epoch) {
3291 fold(row);
3292 }
3293 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3294 fold(row);
3295 }
3296 for rr in &self.run_refs {
3297 let mut reader = self.open_reader(rr.run_id)?;
3298 for row in reader.visible_versions(snapshot.epoch)? {
3299 fold(row);
3300 }
3301 }
3302 let mut out: Vec<Row> = best
3303 .into_values()
3304 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3305 .collect();
3306 out.sort_by_key(|r| r.row_id);
3307 Ok(out)
3308 }
3309
3310 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3317 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3318 let rr = self.run_refs[0].clone();
3319 let mut reader = self.open_reader(rr.run_id)?;
3320 let idxs = reader.visible_indices(snapshot.epoch)?;
3321 let mut cols = Vec::with_capacity(self.schema.columns.len());
3322 for cdef in &self.schema.columns {
3323 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3324 }
3325 return Ok(cols);
3326 }
3327 let rows = self.visible_rows(snapshot)?;
3329 let mut cols: Vec<(u16, Vec<Value>)> = self
3330 .schema
3331 .columns
3332 .iter()
3333 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3334 .collect();
3335 for r in &rows {
3336 for (cid, vec) in cols.iter_mut() {
3337 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3338 }
3339 }
3340 Ok(cols)
3341 }
3342
3343 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3345 self.hot.get(key)
3346 }
3347
3348 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3353 self.require_select()?;
3354 self.ensure_indexes_complete()?;
3355 let snapshot = self.snapshot();
3356 crate::trace::QueryTrace::record(|t| {
3357 t.run_count = self.run_refs.len();
3358 t.memtable_rows = self.memtable.len();
3359 t.mutable_run_rows = self.mutable_run.len();
3360 });
3361 if q.conditions.is_empty() {
3365 crate::trace::QueryTrace::record(|t| {
3366 t.scan_mode = crate::trace::ScanMode::Materialized;
3367 t.row_materialized = true;
3368 });
3369 return self.visible_rows(snapshot);
3370 }
3371 crate::trace::QueryTrace::record(|t| {
3372 t.conditions_pushed = q.conditions.len();
3373 t.scan_mode = crate::trace::ScanMode::Materialized;
3374 t.row_materialized = true;
3375 });
3376 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3383 ordered.sort_by_key(|c| condition_cost_rank(c));
3384 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3385 for c in &ordered {
3386 let s = self.resolve_condition(c, snapshot)?;
3387 let empty = s.is_empty();
3388 sets.push(s);
3389 if empty {
3390 break;
3391 }
3392 }
3393 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3394 self.rows_for_rids(&rids, snapshot)
3395 }
3396
3397 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3402 use std::collections::HashMap;
3403 let mut rows = Vec::with_capacity(rids.len());
3404 let tier_size = self.memtable.len() + self.mutable_run.len();
3421 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3422 if rids.len().saturating_mul(24) < tier_size {
3423 for &rid in rids {
3424 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3425 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3426 let newest = match (mem, mrun) {
3427 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3428 (Some((_, mr)), None) => Some(mr),
3429 (None, Some((_, rr))) => Some(rr),
3430 (None, None) => None,
3431 };
3432 if let Some(row) = newest {
3433 overlay.insert(rid, row);
3434 }
3435 }
3436 } else {
3437 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3438 overlay
3439 .entry(row.row_id.0)
3440 .and_modify(|e| {
3441 if row.committed_epoch > e.committed_epoch {
3442 *e = row.clone();
3443 }
3444 })
3445 .or_insert(row);
3446 };
3447 for row in self.memtable.visible_versions(snapshot.epoch) {
3448 fold_newest(row, &mut overlay);
3449 }
3450 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3451 fold_newest(row, &mut overlay);
3452 }
3453 }
3454 if self.run_refs.len() == 1 {
3455 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3456 if rids.len().saturating_mul(24) < reader.row_count() {
3464 for &rid in rids {
3465 if let Some(r) = overlay.get(&rid) {
3466 if !r.deleted {
3467 rows.push(r.clone());
3468 }
3469 continue;
3470 }
3471 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3472 if !row.deleted {
3473 rows.push(row);
3474 }
3475 }
3476 }
3477 return Ok(rows);
3478 }
3479 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3488 enum Src {
3491 Overlay,
3492 Run,
3493 }
3494 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3495 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3496 for rid in rids {
3497 if overlay.contains_key(rid) {
3498 plan.push(Src::Overlay);
3499 continue;
3500 }
3501 match vis_rids.binary_search(&(*rid as i64)) {
3502 Ok(i) => {
3503 plan.push(Src::Run);
3504 fetch.push(positions[i]);
3505 }
3506 Err(_) => { }
3507 }
3508 }
3509 let fetched = reader.materialize_batch(&fetch)?;
3510 let mut fetched_iter = fetched.into_iter();
3511 for (rid, src) in rids.iter().zip(plan) {
3512 match src {
3513 Src::Overlay => {
3514 if let Some(r) = overlay.get(rid) {
3515 if !r.deleted {
3516 rows.push(r.clone());
3517 }
3518 }
3519 }
3520 Src::Run => {
3521 if let Some(row) = fetched_iter.next() {
3522 if !row.deleted {
3523 rows.push(row);
3524 }
3525 }
3526 }
3527 }
3528 }
3529 return Ok(rows);
3530 }
3531 let mut readers: Vec<_> = self
3535 .run_refs
3536 .iter()
3537 .map(|rr| self.open_reader(rr.run_id))
3538 .collect::<Result<Vec<_>>>()?;
3539 for rid in rids {
3540 if let Some(r) = overlay.get(rid) {
3541 if !r.deleted {
3542 rows.push(r.clone());
3543 }
3544 continue;
3545 }
3546 let mut best: Option<(Epoch, Row)> = None;
3547 for reader in readers.iter_mut() {
3548 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3549 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3550 best = Some((epoch, row));
3551 }
3552 }
3553 }
3554 if let Some((_, r)) = best {
3555 if !r.deleted {
3556 rows.push(r);
3557 }
3558 }
3559 }
3560 Ok(rows)
3561 }
3562
3563 pub fn indexes_complete(&self) -> bool {
3573 self.indexes_complete
3574 }
3575
3576 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3578 self.index_build_policy
3579 }
3580
3581 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3585 self.index_build_policy = policy;
3586 }
3587
3588 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3593 if !self.indexes_complete {
3597 return None;
3598 }
3599 let b = self.bitmap.get(&column_id)?;
3600 let result: Vec<Vec<u8>> = b
3601 .keys()
3602 .into_iter()
3603 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3604 .cloned()
3605 .collect();
3606 Some(result)
3607 }
3608
3609 pub fn fk_join_row_ids(
3610 &self,
3611 fk_column_id: u16,
3612 pk_values: &[Vec<u8>],
3613 fk_conditions: &[crate::query::Condition],
3614 snapshot: Snapshot,
3615 ) -> Result<Vec<u64>> {
3616 let Some(b) = self.bitmap.get(&fk_column_id) else {
3617 return Ok(Vec::new());
3618 };
3619 let mut join_set = {
3620 let mut acc = roaring::RoaringBitmap::new();
3621 for v in pk_values {
3622 acc |= b.get(v);
3623 }
3624 RowIdSet::from_roaring(acc)
3625 };
3626 if !fk_conditions.is_empty() {
3627 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3628 sets.push(join_set);
3629 for c in fk_conditions {
3630 sets.push(self.resolve_condition(c, snapshot)?);
3631 }
3632 join_set = RowIdSet::intersect_many(sets);
3633 }
3634 Ok(join_set.into_sorted_vec())
3635 }
3636
3637 pub fn fk_join_count(
3643 &self,
3644 fk_column_id: u16,
3645 pk_values: &[Vec<u8>],
3646 fk_conditions: &[crate::query::Condition],
3647 snapshot: Snapshot,
3648 ) -> Result<u64> {
3649 let Some(b) = self.bitmap.get(&fk_column_id) else {
3650 return Ok(0);
3651 };
3652 let mut acc = roaring::RoaringBitmap::new();
3653 for v in pk_values {
3654 acc |= b.get(v);
3655 }
3656 if fk_conditions.is_empty() {
3657 return Ok(acc.len());
3658 }
3659 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3660 sets.push(RowIdSet::from_roaring(acc));
3661 for c in fk_conditions {
3662 sets.push(self.resolve_condition(c, snapshot)?);
3663 }
3664 Ok(RowIdSet::intersect_many(sets).len() as u64)
3665 }
3666
3667 fn resolve_condition(
3672 &self,
3673 c: &crate::query::Condition,
3674 snapshot: Snapshot,
3675 ) -> Result<RowIdSet> {
3676 use crate::query::Condition;
3677 Ok(match c {
3678 Condition::Pk(key) => {
3679 let lookup = self
3680 .schema
3681 .primary_key()
3682 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3683 .unwrap_or_else(|| key.clone());
3684 self.hot
3685 .get(&lookup)
3686 .map(|r| RowIdSet::one(r.0))
3687 .unwrap_or_else(RowIdSet::empty)
3688 }
3689 Condition::BitmapEq { column_id, value } => {
3690 let lookup = self.index_lookup_key_bytes(*column_id, value);
3691 self.bitmap
3692 .get(column_id)
3693 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3694 .unwrap_or_else(RowIdSet::empty)
3695 }
3696 Condition::BitmapIn { column_id, values } => {
3697 let bm = self.bitmap.get(column_id);
3698 let mut acc = roaring::RoaringBitmap::new();
3699 if let Some(b) = bm {
3700 for v in values {
3701 let lookup = self.index_lookup_key_bytes(*column_id, v);
3702 acc |= b.get(&lookup);
3703 }
3704 }
3705 RowIdSet::from_roaring(acc)
3706 }
3707 Condition::BytesPrefix { column_id, prefix } => {
3708 if let Some(b) = self.bitmap.get(column_id) {
3713 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3714 let mut acc = roaring::RoaringBitmap::new();
3715 for key in b.keys() {
3716 if key.starts_with(&lookup_prefix) {
3717 acc |= b.get(key);
3718 }
3719 }
3720 RowIdSet::from_roaring(acc)
3721 } else {
3722 RowIdSet::empty()
3723 }
3724 }
3725 Condition::FmContains { column_id, pattern } => self
3726 .fm
3727 .get(column_id)
3728 .map(|f| {
3729 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3730 })
3731 .unwrap_or_else(RowIdSet::empty),
3732 Condition::FmContainsAll {
3733 column_id,
3734 patterns,
3735 } => {
3736 if let Some(f) = self.fm.get(column_id) {
3739 let sets: Vec<RowIdSet> = patterns
3740 .iter()
3741 .map(|pat| {
3742 RowIdSet::from_unsorted(
3743 f.locate(pat).into_iter().map(|r| r.0).collect(),
3744 )
3745 })
3746 .collect();
3747 RowIdSet::intersect_many(sets)
3748 } else {
3749 RowIdSet::empty()
3750 }
3751 }
3752 Condition::Ann {
3753 column_id,
3754 query,
3755 k,
3756 } => self
3757 .ann
3758 .get(column_id)
3759 .map(|a| {
3760 RowIdSet::from_unsorted(
3761 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3762 )
3763 })
3764 .unwrap_or_else(RowIdSet::empty),
3765 Condition::SparseMatch {
3766 column_id,
3767 query,
3768 k,
3769 } => self
3770 .sparse
3771 .get(column_id)
3772 .map(|s| {
3773 RowIdSet::from_unsorted(
3774 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3775 )
3776 })
3777 .unwrap_or_else(RowIdSet::empty),
3778 Condition::MinHashSimilar {
3779 column_id,
3780 query,
3781 k,
3782 } => self
3783 .minhash
3784 .get(column_id)
3785 .map(|mh| {
3786 RowIdSet::from_unsorted(
3787 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3788 )
3789 })
3790 .unwrap_or_else(RowIdSet::empty),
3791 Condition::Range { column_id, lo, hi } => {
3792 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3801 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3802 } else if self.run_refs.len() == 1 {
3803 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3804 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3805 } else {
3806 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3807 };
3808 set.remove_many(self.overlay_rid_set(snapshot));
3809 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3810 set
3811 }
3812 Condition::RangeF64 {
3813 column_id,
3814 lo,
3815 lo_inclusive,
3816 hi,
3817 hi_inclusive,
3818 } => {
3819 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3822 RowIdSet::from_unsorted(
3823 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3824 .into_iter()
3825 .collect(),
3826 )
3827 } else if self.run_refs.len() == 1 {
3828 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3829 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3830 } else {
3831 return self.range_scan_f64(
3832 *column_id,
3833 *lo,
3834 *lo_inclusive,
3835 *hi,
3836 *hi_inclusive,
3837 snapshot,
3838 );
3839 };
3840 set.remove_many(self.overlay_rid_set(snapshot));
3841 self.range_scan_overlay_f64(
3842 &mut set,
3843 *column_id,
3844 *lo,
3845 *lo_inclusive,
3846 *hi,
3847 *hi_inclusive,
3848 snapshot,
3849 );
3850 set
3851 }
3852 Condition::IsNull { column_id } => {
3853 let mut set = if self.run_refs.len() == 1 {
3854 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3855 r.null_row_id_set(*column_id, true)?
3856 } else {
3857 return self.null_scan(*column_id, true, snapshot);
3858 };
3859 set.remove_many(self.overlay_rid_set(snapshot));
3860 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3861 set
3862 }
3863 Condition::IsNotNull { column_id } => {
3864 let mut set = if self.run_refs.len() == 1 {
3865 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3866 r.null_row_id_set(*column_id, false)?
3867 } else {
3868 return self.null_scan(*column_id, false, snapshot);
3869 };
3870 set.remove_many(self.overlay_rid_set(snapshot));
3871 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3872 set
3873 }
3874 })
3875 }
3876
3877 fn range_scan_i64(
3885 &self,
3886 column_id: u16,
3887 lo: i64,
3888 hi: i64,
3889 snapshot: Snapshot,
3890 ) -> Result<RowIdSet> {
3891 let mut row_ids = Vec::new();
3892 let overlay_rids = self.overlay_rid_set(snapshot);
3893 for rr in &self.run_refs {
3894 let mut reader = self.open_reader(rr.run_id)?;
3895 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3896 for rid in matched {
3897 if !overlay_rids.contains(&rid) {
3898 row_ids.push(rid);
3899 }
3900 }
3901 }
3902 let mut s = RowIdSet::from_unsorted(row_ids);
3903 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3904 Ok(s)
3905 }
3906
3907 fn range_scan_f64(
3910 &self,
3911 column_id: u16,
3912 lo: f64,
3913 lo_inclusive: bool,
3914 hi: f64,
3915 hi_inclusive: bool,
3916 snapshot: Snapshot,
3917 ) -> Result<RowIdSet> {
3918 let mut row_ids = Vec::new();
3919 let overlay_rids = self.overlay_rid_set(snapshot);
3920 for rr in &self.run_refs {
3921 let mut reader = self.open_reader(rr.run_id)?;
3922 let matched = reader.range_row_ids_visible_f64(
3923 column_id,
3924 lo,
3925 lo_inclusive,
3926 hi,
3927 hi_inclusive,
3928 snapshot.epoch,
3929 )?;
3930 for rid in matched {
3931 if !overlay_rids.contains(&rid) {
3932 row_ids.push(rid);
3933 }
3934 }
3935 }
3936 let mut s = RowIdSet::from_unsorted(row_ids);
3937 self.range_scan_overlay_f64(
3938 &mut s,
3939 column_id,
3940 lo,
3941 lo_inclusive,
3942 hi,
3943 hi_inclusive,
3944 snapshot,
3945 );
3946 Ok(s)
3947 }
3948
3949 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3951 let mut s = HashSet::new();
3952 for row in self.memtable.visible_versions(snapshot.epoch) {
3953 s.insert(row.row_id.0);
3954 }
3955 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3956 s.insert(row.row_id.0);
3957 }
3958 s
3959 }
3960
3961 fn range_scan_overlay_i64(
3962 &self,
3963 s: &mut RowIdSet,
3964 column_id: u16,
3965 lo: i64,
3966 hi: i64,
3967 snapshot: Snapshot,
3968 ) {
3969 let mut newest: HashMap<u64, &Row> = HashMap::new();
3974 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3975 let memtable = self.memtable.visible_versions(snapshot.epoch);
3976 for r in &mutable {
3977 newest.entry(r.row_id.0).or_insert(r);
3978 }
3979 for r in &memtable {
3980 newest.insert(r.row_id.0, r);
3981 }
3982 for row in newest.values() {
3983 if !row.deleted {
3984 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3985 if *v >= lo && *v <= hi {
3986 s.insert(row.row_id.0);
3987 }
3988 }
3989 }
3990 }
3991 }
3992
3993 #[allow(clippy::too_many_arguments)]
3994 fn range_scan_overlay_f64(
3995 &self,
3996 s: &mut RowIdSet,
3997 column_id: u16,
3998 lo: f64,
3999 lo_inclusive: bool,
4000 hi: f64,
4001 hi_inclusive: bool,
4002 snapshot: Snapshot,
4003 ) {
4004 let mut newest: HashMap<u64, &Row> = HashMap::new();
4007 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4008 let memtable = self.memtable.visible_versions(snapshot.epoch);
4009 for r in &mutable {
4010 newest.entry(r.row_id.0).or_insert(r);
4011 }
4012 for r in &memtable {
4013 newest.insert(r.row_id.0, r);
4014 }
4015 for row in newest.values() {
4016 if !row.deleted {
4017 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
4018 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
4019 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
4020 if ok_lo && ok_hi {
4021 s.insert(row.row_id.0);
4022 }
4023 }
4024 }
4025 }
4026 }
4027
4028 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
4031 let mut row_ids = Vec::new();
4032 let overlay_rids = self.overlay_rid_set(snapshot);
4033 for rr in &self.run_refs {
4034 let mut reader = self.open_reader(rr.run_id)?;
4035 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
4036 for rid in matched {
4037 if !overlay_rids.contains(&rid) {
4038 row_ids.push(rid);
4039 }
4040 }
4041 }
4042 let mut s = RowIdSet::from_unsorted(row_ids);
4043 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
4044 Ok(s)
4045 }
4046
4047 fn null_scan_overlay(
4051 &self,
4052 s: &mut RowIdSet,
4053 column_id: u16,
4054 want_nulls: bool,
4055 snapshot: Snapshot,
4056 ) {
4057 let mut newest: HashMap<u64, &Row> = HashMap::new();
4058 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
4059 let memtable = self.memtable.visible_versions(snapshot.epoch);
4060 for r in &mutable {
4061 newest.entry(r.row_id.0).or_insert(r);
4062 }
4063 for r in &memtable {
4064 newest.insert(r.row_id.0, r);
4065 }
4066 for row in newest.values() {
4067 if row.deleted {
4068 continue;
4069 }
4070 let is_null = !row.columns.contains_key(&column_id)
4071 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
4072 if is_null == want_nulls {
4073 s.insert(row.row_id.0);
4074 }
4075 }
4076 }
4077
4078 pub fn snapshot(&self) -> Snapshot {
4079 Snapshot::at(self.epoch.visible())
4080 }
4081
4082 pub fn pin_snapshot(&mut self) -> Snapshot {
4085 let e = self.epoch.visible();
4086 *self.pinned.entry(e).or_insert(0) += 1;
4087 Snapshot::at(e)
4088 }
4089
4090 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
4092 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
4093 *count -= 1;
4094 if *count == 0 {
4095 self.pinned.remove(&snap.epoch);
4096 }
4097 }
4098 }
4099
4100 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
4110 let local = self.pinned.keys().next().copied();
4111 let global = self.snapshots.min_pinned();
4112 match (local, global) {
4113 (Some(a), Some(b)) => Some(a.min(b)),
4114 (Some(a), None) => Some(a),
4115 (None, b) => b,
4116 }
4117 }
4118
4119 pub fn current_epoch(&self) -> Epoch {
4120 self.epoch.visible()
4121 }
4122
4123 pub fn memtable_len(&self) -> usize {
4124 self.memtable.len()
4125 }
4126
4127 pub fn count(&self) -> u64 {
4130 self.live_count
4131 }
4132
4133 pub fn count_conditions(
4137 &mut self,
4138 conditions: &[crate::query::Condition],
4139 snapshot: Snapshot,
4140 ) -> Result<Option<u64>> {
4141 use crate::query::Condition;
4142 if conditions.is_empty() {
4143 return Ok(Some(self.live_count));
4144 }
4145 let served = |c: &Condition| {
4146 matches!(
4147 c,
4148 Condition::Pk(_)
4149 | Condition::BitmapEq { .. }
4150 | Condition::BitmapIn { .. }
4151 | Condition::BytesPrefix { .. }
4152 | Condition::FmContains { .. }
4153 | Condition::FmContainsAll { .. }
4154 | Condition::Ann { .. }
4155 | Condition::Range { .. }
4156 | Condition::RangeF64 { .. }
4157 | Condition::SparseMatch { .. }
4158 | Condition::MinHashSimilar { .. }
4159 | Condition::IsNull { .. }
4160 | Condition::IsNotNull { .. }
4161 )
4162 };
4163 if !conditions.iter().all(served) {
4164 return Ok(None);
4165 }
4166 self.ensure_indexes_complete()?;
4167 let mut sets = Vec::with_capacity(conditions.len());
4168 for condition in conditions {
4169 sets.push(self.resolve_condition(condition, snapshot)?);
4170 }
4171 let mut rids = RowIdSet::intersect_many(sets);
4172 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4182 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
4183 }
4184 let count = rids.len() as u64;
4185 crate::trace::QueryTrace::record(|t| {
4186 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
4187 t.survivor_count = Some(count as usize);
4188 t.conditions_pushed = conditions.len();
4189 });
4190 Ok(Some(count))
4191 }
4192
4193 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
4198 let mut out = Vec::new();
4199 for row in self.memtable.visible_versions(snapshot.epoch) {
4200 if row.deleted {
4201 out.push(row.row_id.0);
4202 }
4203 }
4204 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4205 if row.deleted {
4206 out.push(row.row_id.0);
4207 }
4208 }
4209 out
4210 }
4211
4212 pub fn bulk_load_columns(
4221 &mut self,
4222 user_columns: Vec<(u16, columnar::NativeColumn)>,
4223 ) -> Result<Epoch> {
4224 self.bulk_load_columns_with(user_columns, 3, false, true)
4225 }
4226
4227 pub fn bulk_load_fast(
4234 &mut self,
4235 user_columns: Vec<(u16, columnar::NativeColumn)>,
4236 ) -> Result<Epoch> {
4237 self.bulk_load_columns_with(user_columns, -1, true, false)
4238 }
4239
4240 fn bulk_load_columns_with(
4241 &mut self,
4242 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4243 zstd_level: i32,
4244 force_plain: bool,
4245 lz4: bool,
4246 ) -> Result<Epoch> {
4247 let epoch = self.commit()?;
4248 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4249 if n == 0 {
4250 return Ok(epoch);
4251 }
4252 let live_before = self.live_count;
4253 self.spill_mutable_run(epoch)?;
4255 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4256 && self.indexes_complete
4257 && self.run_refs.is_empty()
4258 && self.memtable.is_empty()
4259 && self.mutable_run.is_empty();
4260 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4263 self.validate_columns_not_null(&user_columns, n)?;
4264 let winner_idx = self
4265 .bulk_pk_winner_indices(&user_columns, n)
4266 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
4267 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4268 match winner_idx.as_deref() {
4269 Some(idx) => {
4270 let compacted = user_columns
4271 .iter()
4272 .map(|(id, c)| (*id, c.gather(idx)))
4273 .collect();
4274 (compacted, idx.len())
4275 }
4276 None => (user_columns, n),
4277 };
4278 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4279 let first = self.allocator.alloc_range(write_n as u64).0;
4280 for rid in first..first + write_n as u64 {
4281 self.reservoir.offer(rid);
4282 }
4283 let run_id = self.next_run_id;
4284 self.next_run_id += 1;
4285 let path = self.run_path(run_id);
4286 let mut writer =
4287 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4288 if force_plain {
4289 writer = writer.with_plain();
4290 } else if lz4 {
4291 writer = writer.with_lz4();
4294 } else {
4295 writer = writer.with_zstd_level(zstd_level);
4296 }
4297 if let Some(kek) = &self.kek {
4298 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4299 }
4300 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4301 self.run_refs.push(RunRef {
4302 run_id: run_id as u128,
4303 level: 0,
4304 epoch_created: epoch.0,
4305 row_count: header.row_count,
4306 });
4307 self.live_count = self.live_count.saturating_add(write_n as u64);
4308 if eager_index_build {
4309 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4310 self.index_columns_bulk(&write_columns, &row_ids);
4311 self.indexes_complete = true;
4312 self.build_learned_ranges()?;
4313 } else {
4314 self.indexes_complete = false;
4318 }
4319 self.mark_flushed(epoch)?;
4320 self.persist_manifest(epoch)?;
4321 if eager_index_build {
4322 self.checkpoint_indexes(epoch);
4323 }
4324 self.clear_result_cache();
4325 Ok(epoch)
4326 }
4327
4328 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4346 let n = row_ids.len();
4347 if n == 0 {
4348 return;
4349 }
4350 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4351 columns.iter().map(|(id, c)| (*id, c)).collect();
4352 let ty_of: std::collections::HashMap<u16, TypeId> = self
4353 .schema
4354 .columns
4355 .iter()
4356 .map(|c| (c.id, c.ty.clone()))
4357 .collect();
4358 let pk_id = self.schema.primary_key().map(|c| c.id);
4359
4360 for (i, &rid) in row_ids.iter().enumerate() {
4361 let row_id = RowId(rid);
4362 if let Some(pid) = pk_id {
4363 if let Some(col) = by_id.get(&pid) {
4364 let ty = ty_of.get(&pid).cloned().unwrap_or(TypeId::Int64);
4365 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4366 self.insert_hot_pk(key, row_id);
4367 }
4368 }
4369 }
4370 for idef in &self.schema.indexes {
4371 let Some(col) = by_id.get(&idef.column_id) else {
4372 continue;
4373 };
4374 let ty = ty_of.get(&idef.column_id).cloned().unwrap_or(TypeId::Int64);
4375 match idef.kind {
4376 IndexKind::Bitmap => {
4377 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4378 if let Some(key) =
4379 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4380 {
4381 b.insert(key, row_id);
4382 }
4383 }
4384 }
4385 IndexKind::FmIndex => {
4386 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4387 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4388 f.insert(bytes.to_vec(), row_id);
4389 }
4390 }
4391 }
4392 IndexKind::Sparse => {
4393 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4394 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4395 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4396 s.insert(&terms, row_id);
4397 }
4398 }
4399 }
4400 }
4401 IndexKind::MinHash => {
4402 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4403 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4404 let tokens = crate::index::token_hashes_from_bytes(bytes);
4405 mh.insert(&tokens, row_id);
4406 }
4407 }
4408 }
4409 _ => {}
4410 }
4411 }
4412 }
4413 }
4414
4415 pub fn visible_columns_native(
4420 &self,
4421 snapshot: Snapshot,
4422 projection: Option<&[u16]>,
4423 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4424 let wanted: Vec<u16> = match projection {
4425 Some(p) => p.to_vec(),
4426 None => self.schema.columns.iter().map(|c| c.id).collect(),
4427 };
4428 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4429 let rr = self.run_refs[0].clone();
4430 let mut reader = self.open_reader(rr.run_id)?;
4431 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4432 let all_visible = idxs.len() == reader.row_count();
4433 if reader.has_mmap() {
4439 use rayon::prelude::*;
4440 let valid: Vec<u16> = wanted
4443 .iter()
4444 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4445 .copied()
4446 .collect();
4447 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4449 .par_iter()
4450 .filter_map(|cid| {
4451 reader
4452 .column_native_shared(*cid)
4453 .ok()
4454 .map(|col| (*cid, col))
4455 })
4456 .collect();
4457 let cols = decoded
4458 .into_iter()
4459 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4460 .collect();
4461 return Ok(cols);
4462 }
4463 let mut cols = Vec::with_capacity(wanted.len());
4464 for cid in &wanted {
4465 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4466 Some(c) => c,
4467 None => continue,
4468 };
4469 let col = reader.column_native(cdef.id)?;
4470 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4471 }
4472 return Ok(cols);
4473 }
4474 let vcols = self.visible_columns(snapshot)?;
4475 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4476 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4477 .into_iter()
4478 .filter(|(id, _)| want_set.contains(id))
4479 .map(|(id, vals)| {
4480 let ty = self
4481 .schema
4482 .columns
4483 .iter()
4484 .find(|c| c.id == id)
4485 .map(|c| c.ty.clone())
4486 .unwrap_or(TypeId::Bytes);
4487 (id, columnar::values_to_native(ty, &vals))
4488 })
4489 .collect();
4490 Ok(out)
4491 }
4492
4493 pub fn run_count(&self) -> usize {
4494 self.run_refs.len()
4495 }
4496
4497 pub fn memtable_is_empty(&self) -> bool {
4499 self.memtable.is_empty()
4500 }
4501
4502 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4506 self.page_cache.stats()
4507 }
4508
4509 pub fn reset_page_cache_stats(&self) {
4511 self.page_cache.reset_stats();
4512 }
4513
4514 pub fn run_ids(&self) -> Vec<u128> {
4517 self.run_refs.iter().map(|r| r.run_id).collect()
4518 }
4519
4520 pub fn single_run_is_clean(&self) -> bool {
4524 if self.run_refs.len() != 1 {
4525 return false;
4526 }
4527 self.open_reader(self.run_refs[0].run_id)
4528 .map(|r| r.is_clean())
4529 .unwrap_or(false)
4530 }
4531
4532 fn resolve_footprint(
4539 &self,
4540 conditions: &[crate::query::Condition],
4541 _snapshot: Snapshot,
4542 ) -> roaring::RoaringBitmap {
4543 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4544 return roaring::RoaringBitmap::new();
4545 }
4546 if self.run_refs.is_empty() {
4547 return roaring::RoaringBitmap::new();
4548 }
4549 if self.run_refs.len() == 1 {
4551 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4552 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4553 return rids.to_roaring_lossy();
4554 }
4555 }
4556 }
4557 roaring::RoaringBitmap::new()
4558 }
4559
4560 pub fn query_columns_native_cached(
4571 &mut self,
4572 conditions: &[crate::query::Condition],
4573 projection: Option<&[u16]>,
4574 snapshot: Snapshot,
4575 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4576 if conditions.is_empty() {
4577 return self.query_columns_native(conditions, projection, snapshot);
4578 }
4579 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4583 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4584 crate::trace::QueryTrace::record(|t| {
4585 t.result_cache_hit = true;
4586 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4587 });
4588 return Ok(Some((*hit).clone()));
4589 }
4590 let res = self.query_columns_native(conditions, projection, snapshot)?;
4591 if let Some(cols) = &res {
4592 let footprint = self.resolve_footprint(conditions, snapshot);
4593 let condition_cols = crate::query::condition_columns(conditions);
4594 self.result_cache.lock().insert(
4595 key,
4596 CachedEntry {
4597 data: CachedData::Columns(Arc::new(cols.clone())),
4598 footprint,
4599 condition_cols,
4600 },
4601 );
4602 }
4603 Ok(res)
4604 }
4605
4606 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4611 if q.conditions.is_empty() {
4612 return self.query(q);
4613 }
4614 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4615 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4616 crate::trace::QueryTrace::record(|t| {
4617 t.result_cache_hit = true;
4618 t.scan_mode = crate::trace::ScanMode::Materialized;
4619 });
4620 return Ok((*hit).clone());
4621 }
4622 let rows = self.query(q)?;
4623 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4624 let condition_cols = crate::query::condition_columns(&q.conditions);
4625 self.result_cache.lock().insert(
4626 key,
4627 CachedEntry {
4628 data: CachedData::Rows(Arc::new(rows.clone())),
4629 footprint,
4630 condition_cols,
4631 },
4632 );
4633 Ok(rows)
4634 }
4635
4636 #[allow(clippy::type_complexity)]
4651 pub fn query_columns_native_traced(
4652 &mut self,
4653 conditions: &[crate::query::Condition],
4654 projection: Option<&[u16]>,
4655 snapshot: Snapshot,
4656 ) -> Result<(
4657 Option<Vec<(u16, columnar::NativeColumn)>>,
4658 crate::trace::QueryTrace,
4659 )> {
4660 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4661 self.query_columns_native(conditions, projection, snapshot)
4662 });
4663 Ok((result?, trace))
4664 }
4665
4666 #[allow(clippy::type_complexity)]
4669 pub fn query_columns_native_cached_traced(
4670 &mut self,
4671 conditions: &[crate::query::Condition],
4672 projection: Option<&[u16]>,
4673 snapshot: Snapshot,
4674 ) -> Result<(
4675 Option<Vec<(u16, columnar::NativeColumn)>>,
4676 crate::trace::QueryTrace,
4677 )> {
4678 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4679 self.query_columns_native_cached(conditions, projection, snapshot)
4680 });
4681 Ok((result?, trace))
4682 }
4683
4684 pub fn native_page_cursor_traced(
4686 &self,
4687 snapshot: Snapshot,
4688 projection: Vec<(u16, TypeId)>,
4689 conditions: &[crate::query::Condition],
4690 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4691 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4692 self.native_page_cursor(snapshot, projection, conditions)
4693 });
4694 Ok((result?, trace))
4695 }
4696
4697 pub fn native_multi_run_cursor_traced(
4699 &self,
4700 snapshot: Snapshot,
4701 projection: Vec<(u16, TypeId)>,
4702 conditions: &[crate::query::Condition],
4703 ) -> Result<(
4704 Option<crate::cursor::MultiRunCursor>,
4705 crate::trace::QueryTrace,
4706 )> {
4707 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4708 self.native_multi_run_cursor(snapshot, projection, conditions)
4709 });
4710 Ok((result?, trace))
4711 }
4712
4713 pub fn count_conditions_traced(
4715 &mut self,
4716 conditions: &[crate::query::Condition],
4717 snapshot: Snapshot,
4718 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4719 let (result, trace) =
4720 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4721 Ok((result?, trace))
4722 }
4723
4724 pub fn query_traced(
4726 &mut self,
4727 q: &crate::query::Query,
4728 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4729 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4730 Ok((result?, trace))
4731 }
4732
4733 pub fn query_columns_native(
4738 &mut self,
4739 conditions: &[crate::query::Condition],
4740 projection: Option<&[u16]>,
4741 snapshot: Snapshot,
4742 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4743 use crate::query::Condition;
4744 if conditions.is_empty() {
4745 return Ok(None);
4746 }
4747 self.ensure_indexes_complete()?;
4748
4749 let served = |c: &Condition| {
4754 matches!(
4755 c,
4756 Condition::Pk(_)
4757 | Condition::BitmapEq { .. }
4758 | Condition::BitmapIn { .. }
4759 | Condition::BytesPrefix { .. }
4760 | Condition::FmContains { .. }
4761 | Condition::FmContainsAll { .. }
4762 | Condition::Ann { .. }
4763 | Condition::Range { .. }
4764 | Condition::RangeF64 { .. }
4765 | Condition::SparseMatch { .. }
4766 | Condition::MinHashSimilar { .. }
4767 | Condition::IsNull { .. }
4768 | Condition::IsNotNull { .. }
4769 )
4770 };
4771 if !conditions.iter().all(served) {
4772 return Ok(None);
4773 }
4774 let fast_path =
4775 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4776 crate::trace::QueryTrace::record(|t| {
4777 t.run_count = self.run_refs.len();
4778 t.memtable_rows = self.memtable.len();
4779 t.mutable_run_rows = self.mutable_run.len();
4780 t.conditions_pushed = conditions.len();
4781 t.learned_range_used = conditions.iter().any(|c| match c {
4782 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4783 self.learned_range.contains_key(column_id)
4784 }
4785 _ => false,
4786 });
4787 });
4788 let col_ids: Vec<u16> = projection
4790 .map(|p| p.to_vec())
4791 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4792 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4793 .iter()
4794 .map(|&cid| {
4795 let ty = self
4796 .schema
4797 .columns
4798 .iter()
4799 .find(|c| c.id == cid)
4800 .map(|c| c.ty.clone())
4801 .unwrap_or(TypeId::Bytes);
4802 (cid, ty)
4803 })
4804 .collect();
4805
4806 if fast_path {
4812 let needs_column = conditions.iter().any(|c| match c {
4815 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4816 Condition::RangeF64 { column_id, .. } => {
4817 !self.learned_range.contains_key(column_id)
4818 }
4819 _ => false,
4820 });
4821 let mut reader_opt: Option<RunReader> = if needs_column {
4822 Some(self.open_reader(self.run_refs[0].run_id)?)
4823 } else {
4824 None
4825 };
4826 let mut sets: Vec<RowIdSet> = Vec::new();
4827 for c in conditions {
4828 let s = match c {
4829 Condition::Range { column_id, lo, hi }
4830 if !self.learned_range.contains_key(column_id) =>
4831 {
4832 if reader_opt.is_none() {
4833 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4834 }
4835 reader_opt
4836 .as_mut()
4837 .expect("reader opened for range")
4838 .range_row_id_set_i64(*column_id, *lo, *hi)?
4839 }
4840 Condition::RangeF64 {
4841 column_id,
4842 lo,
4843 lo_inclusive,
4844 hi,
4845 hi_inclusive,
4846 } if !self.learned_range.contains_key(column_id) => {
4847 if reader_opt.is_none() {
4848 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4849 }
4850 reader_opt
4851 .as_mut()
4852 .expect("reader opened for range")
4853 .range_row_id_set_f64(
4854 *column_id,
4855 *lo,
4856 *lo_inclusive,
4857 *hi,
4858 *hi_inclusive,
4859 )?
4860 }
4861 _ => self.resolve_condition(c, snapshot)?,
4862 };
4863 sets.push(s);
4864 }
4865 let candidates = RowIdSet::intersect_many(sets);
4866 crate::trace::QueryTrace::record(|t| {
4867 t.survivor_count = Some(candidates.len());
4868 });
4869 if candidates.is_empty() {
4870 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4871 .iter()
4872 .map(|&id| {
4873 (
4874 id,
4875 columnar::null_native(
4876 proj_pairs
4877 .iter()
4878 .find(|(c, _)| c == &id)
4879 .map(|(_, t)| t.clone())
4880 .unwrap_or(TypeId::Bytes),
4881 0,
4882 ),
4883 )
4884 })
4885 .collect();
4886 return Ok(Some(cols));
4887 }
4888 let mut reader = match reader_opt.take() {
4889 Some(r) => r,
4890 None => self.open_reader(self.run_refs[0].run_id)?,
4891 };
4892 let candidate_ids = candidates.into_sorted_vec();
4893 let (positions, fast_rid) = if let Some(positions) =
4894 reader.positions_for_row_ids_fast(&candidate_ids)
4895 {
4896 (positions, true)
4897 } else {
4898 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4899 match col {
4900 columnar::NativeColumn::Int64 { data, .. } => {
4901 let mut p: Vec<usize> = candidate_ids
4902 .iter()
4903 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4904 .collect();
4905 p.sort_unstable();
4906 (p, false)
4907 }
4908 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4909 }
4910 };
4911 crate::trace::QueryTrace::record(|t| {
4912 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4913 t.fast_row_id_map = fast_rid;
4914 });
4915 let mut cols = Vec::with_capacity(col_ids.len());
4916 for cid in &col_ids {
4917 let col = reader.column_native(*cid)?;
4918 cols.push((*cid, col.gather(&positions)));
4919 }
4920 return Ok(Some(cols));
4921 }
4922
4923 if !self.run_refs.is_empty() {
4936 use crate::cursor::{drain_cursor_to_columns, Cursor};
4937 let remaining: usize;
4938 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4939 let c = self
4940 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4941 .expect("single-run cursor should build when run_refs.len() == 1");
4942 remaining = c.remaining_rows();
4943 Box::new(c)
4944 } else {
4945 let c = self
4946 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4947 .expect("multi-run cursor should build when run_refs.len() >= 1");
4948 remaining = c.remaining_rows();
4949 Box::new(c)
4950 };
4951 crate::trace::QueryTrace::record(|t| {
4952 if t.survivor_count.is_none() {
4953 t.survivor_count = Some(remaining);
4954 }
4955 });
4956 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4957 return Ok(Some(cols));
4958 }
4959
4960 crate::trace::QueryTrace::record(|t| {
4965 t.scan_mode = crate::trace::ScanMode::Materialized;
4966 t.row_materialized = true;
4967 });
4968 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4969 for c in conditions {
4970 sets.push(self.resolve_condition(c, snapshot)?);
4971 }
4972 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4973 let rows = self.rows_for_rids(&rids, snapshot)?;
4974 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4975 for (cid, ty) in &proj_pairs {
4976 let vals: Vec<Value> = rows
4977 .iter()
4978 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4979 .collect();
4980 cols.push((*cid, columnar::values_to_native(ty.clone(), &vals)));
4981 }
4982 Ok(Some(cols))
4983 }
4984
4985 pub fn native_page_cursor(
5000 &self,
5001 snapshot: Snapshot,
5002 projection: Vec<(u16, TypeId)>,
5003 conditions: &[crate::query::Condition],
5004 ) -> Result<Option<NativePageCursor>> {
5005 use crate::cursor::build_page_plans;
5006 if !conditions.is_empty() && !self.indexes_complete {
5009 return Ok(None);
5010 }
5011 if self.run_refs.len() != 1 {
5012 return Ok(None);
5013 }
5014 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
5015 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
5016
5017 let overlay_rids: HashSet<u64> = {
5020 let mut s = HashSet::new();
5021 for row in self.memtable.visible_versions(snapshot.epoch) {
5022 s.insert(row.row_id.0);
5023 }
5024 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5025 s.insert(row.row_id.0);
5026 }
5027 s
5028 };
5029
5030 let survivors = if conditions.is_empty() {
5034 None
5035 } else {
5036 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
5037 };
5038
5039 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
5046 survivors.clone()
5047 } else if let Some(s) = &survivors {
5048 let mut run_set = s.clone();
5049 run_set.remove_many(overlay_rids.iter().copied());
5050 Some(run_set)
5051 } else {
5052 Some(RowIdSet::from_unsorted(
5053 rids.iter()
5054 .map(|&r| r as u64)
5055 .filter(|r| !overlay_rids.contains(r))
5056 .collect(),
5057 ))
5058 };
5059
5060 let overlay_rows = if overlay_rids.is_empty() {
5061 Vec::new()
5062 } else {
5063 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5064 self.overlay_visible_rows(snapshot, bound)
5065 };
5066
5067 let plans = if positions.is_empty() {
5069 Vec::new()
5070 } else {
5071 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
5072 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
5073 };
5074
5075 let overlay = if overlay_rows.is_empty() {
5077 None
5078 } else {
5079 let filtered =
5080 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5081 if filtered.is_empty() {
5082 None
5083 } else {
5084 Some(self.materialize_overlay(&filtered, &projection))
5085 }
5086 };
5087
5088 let overlay_row_count = overlay
5089 .as_ref()
5090 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5091 .unwrap_or(0);
5092 crate::trace::QueryTrace::record(|t| {
5093 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
5094 t.run_count = self.run_refs.len();
5095 t.memtable_rows = self.memtable.len();
5096 t.mutable_run_rows = self.mutable_run.len();
5097 t.overlay_rows = overlay_row_count;
5098 t.conditions_pushed = conditions.len();
5099 t.pages_decoded = plans
5100 .iter()
5101 .map(|p| p.positions.len())
5102 .sum::<usize>()
5103 .min(1);
5104 });
5105
5106 Ok(Some(NativePageCursor::new_with_overlay(
5107 reader, projection, plans, overlay,
5108 )))
5109 }
5110 #[allow(clippy::type_complexity)]
5120 pub fn native_multi_run_cursor(
5121 &self,
5122 snapshot: Snapshot,
5123 projection: Vec<(u16, TypeId)>,
5124 conditions: &[crate::query::Condition],
5125 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
5126 use crate::cursor::{MultiRunCursor, RunStream};
5127 use crate::sorted_run::SYS_ROW_ID;
5128 use std::collections::{BinaryHeap, HashMap, HashSet};
5129 if !conditions.is_empty() && !self.indexes_complete {
5132 return Ok(None);
5133 }
5134 if self.run_refs.is_empty() {
5135 return Ok(None);
5136 }
5137
5138 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
5140 Vec::with_capacity(self.run_refs.len());
5141 for rr in &self.run_refs {
5142 let mut reader = self.open_reader(rr.run_id)?;
5143 let (rids, eps, del) = reader.system_columns_native()?;
5144 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
5145 run_meta.push((reader, rids, eps, del, page_rows));
5146 }
5147
5148 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
5152 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
5153 for i in 0..rids.len() {
5154 let rid = rids[i] as u64;
5155 let e = eps[i] as u64;
5156 if e > snapshot.epoch.0 {
5157 continue;
5158 }
5159 let is_del = del[i] != 0;
5160 best.entry(rid)
5161 .and_modify(|cur| {
5162 if e > cur.0 {
5163 *cur = (e, run_idx, i, is_del);
5164 }
5165 })
5166 .or_insert((e, run_idx, i, is_del));
5167 }
5168 }
5169
5170 let overlay_rids: HashSet<u64> = {
5172 let mut s = HashSet::new();
5173 for row in self.memtable.visible_versions(snapshot.epoch) {
5174 s.insert(row.row_id.0);
5175 }
5176 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5177 s.insert(row.row_id.0);
5178 }
5179 s
5180 };
5181
5182 let survivors: Option<RowIdSet> = if conditions.is_empty() {
5184 None
5185 } else {
5186 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5187 for c in conditions {
5188 sets.push(self.resolve_condition(c, snapshot)?);
5189 }
5190 Some(RowIdSet::intersect_many(sets))
5191 };
5192
5193 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
5197 for (rid, (_, run_idx, pos, deleted)) in &best {
5198 if *deleted {
5199 continue;
5200 }
5201 if overlay_rids.contains(rid) {
5202 continue;
5203 }
5204 if let Some(s) = &survivors {
5205 if !s.contains(*rid) {
5206 continue;
5207 }
5208 }
5209 per_run[*run_idx].push((*rid, *pos));
5210 }
5211 for v in per_run.iter_mut() {
5212 v.sort_unstable_by_key(|&(rid, _)| rid);
5213 }
5214
5215 let mut streams = Vec::with_capacity(run_meta.len());
5217 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
5218 let mut total = 0usize;
5219 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
5220 let mut starts = Vec::with_capacity(page_rows.len());
5221 let mut acc = 0usize;
5222 for &r in &page_rows {
5223 starts.push(acc);
5224 acc += r;
5225 }
5226 let mut survivors_vec: Vec<(u64, usize, usize)> =
5227 Vec::with_capacity(per_run[run_idx].len());
5228 for &(rid, pos) in &per_run[run_idx] {
5229 let page_seq = match starts.partition_point(|&s| s <= pos) {
5230 0 => continue,
5231 p => p - 1,
5232 };
5233 let within = pos - starts[page_seq];
5234 survivors_vec.push((rid, page_seq, within));
5235 }
5236 total += survivors_vec.len();
5237 if let Some(&(rid, _, _)) = survivors_vec.first() {
5238 heap.push(std::cmp::Reverse((rid, run_idx)));
5239 }
5240 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5241 }
5242
5243 let overlay_rows = if overlay_rids.is_empty() {
5245 Vec::new()
5246 } else {
5247 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5248 self.overlay_visible_rows(snapshot, bound)
5249 };
5250 let overlay = if overlay_rows.is_empty() {
5251 None
5252 } else {
5253 let filtered =
5254 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5255 if filtered.is_empty() {
5256 None
5257 } else {
5258 Some(self.materialize_overlay(&filtered, &projection))
5259 }
5260 };
5261
5262 let overlay_row_count = overlay
5263 .as_ref()
5264 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5265 .unwrap_or(0);
5266 crate::trace::QueryTrace::record(|t| {
5267 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5268 t.run_count = self.run_refs.len();
5269 t.memtable_rows = self.memtable.len();
5270 t.mutable_run_rows = self.mutable_run.len();
5271 t.overlay_rows = overlay_row_count;
5272 t.conditions_pushed = conditions.len();
5273 t.survivor_count = Some(total);
5274 });
5275
5276 Ok(Some(MultiRunCursor::new(
5277 streams, projection, heap, total, overlay,
5278 )))
5279 }
5280
5281 fn overlay_materialization_bound<'a>(
5293 conditions: &[crate::query::Condition],
5294 survivors: &'a Option<RowIdSet>,
5295 ) -> Option<&'a RowIdSet> {
5296 use crate::query::Condition;
5297 let has_range = conditions
5298 .iter()
5299 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5300 if has_range {
5301 None
5302 } else {
5303 survivors.as_ref()
5304 }
5305 }
5306
5307 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5319 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5320 let mut fold = |row: Row| {
5321 if let Some(b) = bound {
5322 if !b.contains(row.row_id.0) {
5323 return;
5324 }
5325 }
5326 best.entry(row.row_id.0)
5327 .and_modify(|(be, br)| {
5328 if row.committed_epoch > *be {
5329 *be = row.committed_epoch;
5330 *br = row.clone();
5331 }
5332 })
5333 .or_insert_with(|| (row.committed_epoch, row));
5334 };
5335 for row in self.memtable.visible_versions(snapshot.epoch) {
5336 fold(row);
5337 }
5338 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5339 fold(row);
5340 }
5341 let mut out: Vec<Row> = best
5342 .into_values()
5343 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5344 .collect();
5345 out.sort_by_key(|r| r.row_id);
5346 out
5347 }
5348
5349 fn filter_overlay_rows(
5357 &self,
5358 rows: Vec<Row>,
5359 conditions: &[crate::query::Condition],
5360 survivors: Option<&RowIdSet>,
5361 snapshot: Snapshot,
5362 ) -> Result<Vec<Row>> {
5363 if conditions.is_empty() {
5364 return Ok(rows);
5365 }
5366 use crate::query::Condition;
5367 let all_index_served = !conditions
5371 .iter()
5372 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5373 if all_index_served {
5374 return Ok(rows
5375 .into_iter()
5376 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5377 .collect());
5378 }
5379 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5382 for c in conditions {
5383 let s = match c {
5384 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5385 _ => self.resolve_condition(c, snapshot)?,
5386 };
5387 per_cond_sets.push(s);
5388 }
5389 Ok(rows
5390 .into_iter()
5391 .filter(|row| {
5392 conditions.iter().enumerate().all(|(i, c)| match c {
5393 Condition::Range { column_id, lo, hi } => {
5394 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5395 }
5396 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5397 match row.columns.get(column_id) {
5398 Some(Value::Float64(v)) => {
5399 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5400 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5401 lo_ok && hi_ok
5402 }
5403 _ => false,
5404 }
5405 }
5406 _ => per_cond_sets[i].contains(row.row_id.0),
5407 })
5408 })
5409 .collect())
5410 }
5411
5412 fn materialize_overlay(
5415 &self,
5416 rows: &[Row],
5417 projection: &[(u16, TypeId)],
5418 ) -> Vec<columnar::NativeColumn> {
5419 if projection.is_empty() {
5420 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5421 }
5422 let mut cols = Vec::with_capacity(projection.len());
5423 for (cid, ty) in projection {
5424 let vals: Vec<Value> = rows
5425 .iter()
5426 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5427 .collect();
5428 cols.push(columnar::values_to_native(ty.clone(), &vals));
5429 }
5430 cols
5431 }
5432
5433 fn resolve_survivor_rids(
5438 &self,
5439 conditions: &[crate::query::Condition],
5440 reader: &mut RunReader,
5441 ) -> Result<RowIdSet> {
5442 use crate::query::Condition;
5443 let mut sets: Vec<RowIdSet> = Vec::new();
5444 for c in conditions {
5445 let s: RowIdSet = match c {
5446 Condition::Pk(key) => {
5447 let lookup = self
5448 .schema
5449 .primary_key()
5450 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5451 .unwrap_or_else(|| key.clone());
5452 self.hot
5453 .get(&lookup)
5454 .map(|r| RowIdSet::one(r.0))
5455 .unwrap_or_else(RowIdSet::empty)
5456 }
5457 Condition::BitmapEq { column_id, value } => {
5458 let lookup = self.index_lookup_key_bytes(*column_id, value);
5459 self.bitmap
5460 .get(column_id)
5461 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5462 .unwrap_or_else(RowIdSet::empty)
5463 }
5464 Condition::BitmapIn { column_id, values } => {
5465 let bm = self.bitmap.get(column_id);
5466 let mut acc = roaring::RoaringBitmap::new();
5467 if let Some(b) = bm {
5468 for v in values {
5469 let lookup = self.index_lookup_key_bytes(*column_id, v);
5470 acc |= b.get(&lookup);
5471 }
5472 }
5473 RowIdSet::from_roaring(acc)
5474 }
5475 Condition::BytesPrefix { column_id, prefix } => {
5476 if let Some(b) = self.bitmap.get(column_id) {
5477 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5478 let mut acc = roaring::RoaringBitmap::new();
5479 for key in b.keys() {
5480 if key.starts_with(&lookup_prefix) {
5481 acc |= b.get(key);
5482 }
5483 }
5484 RowIdSet::from_roaring(acc)
5485 } else {
5486 RowIdSet::empty()
5487 }
5488 }
5489 Condition::FmContains { column_id, pattern } => self
5490 .fm
5491 .get(column_id)
5492 .map(|f| {
5493 RowIdSet::from_unsorted(
5494 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5495 )
5496 })
5497 .unwrap_or_else(RowIdSet::empty),
5498 Condition::FmContainsAll {
5499 column_id,
5500 patterns,
5501 } => {
5502 if let Some(f) = self.fm.get(column_id) {
5503 let sets: Vec<RowIdSet> = patterns
5504 .iter()
5505 .map(|pat| {
5506 RowIdSet::from_unsorted(
5507 f.locate(pat).into_iter().map(|r| r.0).collect(),
5508 )
5509 })
5510 .collect();
5511 RowIdSet::intersect_many(sets)
5512 } else {
5513 RowIdSet::empty()
5514 }
5515 }
5516 Condition::Ann {
5517 column_id,
5518 query,
5519 k,
5520 } => self
5521 .ann
5522 .get(column_id)
5523 .map(|a| {
5524 RowIdSet::from_unsorted(
5525 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5526 )
5527 })
5528 .unwrap_or_else(RowIdSet::empty),
5529 Condition::SparseMatch {
5530 column_id,
5531 query,
5532 k,
5533 } => self
5534 .sparse
5535 .get(column_id)
5536 .map(|s| {
5537 RowIdSet::from_unsorted(
5538 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5539 )
5540 })
5541 .unwrap_or_else(RowIdSet::empty),
5542 Condition::MinHashSimilar {
5543 column_id,
5544 query,
5545 k,
5546 } => self
5547 .minhash
5548 .get(column_id)
5549 .map(|mh| {
5550 RowIdSet::from_unsorted(
5551 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5552 )
5553 })
5554 .unwrap_or_else(RowIdSet::empty),
5555 Condition::Range { column_id, lo, hi } => {
5556 if let Some(li) = self.learned_range.get(column_id) {
5557 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5558 } else {
5559 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5560 }
5561 }
5562 Condition::RangeF64 {
5563 column_id,
5564 lo,
5565 lo_inclusive,
5566 hi,
5567 hi_inclusive,
5568 } => {
5569 if let Some(li) = self.learned_range.get(column_id) {
5570 RowIdSet::from_unsorted(
5571 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5572 .into_iter()
5573 .collect(),
5574 )
5575 } else {
5576 reader.range_row_id_set_f64(
5577 *column_id,
5578 *lo,
5579 *lo_inclusive,
5580 *hi,
5581 *hi_inclusive,
5582 )?
5583 }
5584 }
5585 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5586 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5587 };
5588 sets.push(s);
5589 }
5590 Ok(RowIdSet::intersect_many(sets))
5591 }
5592
5593 pub fn scan_cursor(
5614 &self,
5615 snapshot: Snapshot,
5616 projection: Vec<(u16, TypeId)>,
5617 conditions: &[crate::query::Condition],
5618 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5619 if !conditions.is_empty() && !self.indexes_complete {
5625 return Ok(None);
5626 }
5627 if self.run_refs.len() == 1 {
5628 Ok(self
5629 .native_page_cursor(snapshot, projection, conditions)?
5630 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5631 } else {
5632 Ok(self
5633 .native_multi_run_cursor(snapshot, projection, conditions)?
5634 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5635 }
5636 }
5637
5638 pub fn aggregate_native(
5652 &self,
5653 snapshot: Snapshot,
5654 column: Option<u16>,
5655 conditions: &[crate::query::Condition],
5656 agg: NativeAgg,
5657 ) -> Result<Option<NativeAggResult>> {
5658 if self.run_refs.len() == 1 && conditions.is_empty() {
5660 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5661 return Ok(Some(res));
5662 }
5663 }
5664 if matches!(agg, NativeAgg::Count) && column.is_none() {
5666 return Ok(self
5667 .scan_cursor(snapshot, Vec::new(), conditions)?
5668 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5669 }
5670 let cid = match column {
5673 Some(c) => c,
5674 None => return Ok(None),
5675 };
5676 let ty = self.column_type(cid);
5677 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty.clone())], conditions)?
5678 else {
5679 return Ok(None);
5680 };
5681 match ty {
5682 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5683 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5684 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5685 }
5686 TypeId::Float64 => {
5687 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5688 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5689 }
5690 _ => Ok(None),
5691 }
5692 }
5693
5694 fn aggregate_from_stats(
5702 &self,
5703 snapshot: Snapshot,
5704 column: Option<u16>,
5705 agg: NativeAgg,
5706 ) -> Result<Option<NativeAggResult>> {
5707 let cid = match (agg, column) {
5708 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5709 _ => return Ok(None), };
5711 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5712 return Ok(None);
5713 };
5714 let Some(cs) = stats.get(&cid) else {
5715 return Ok(None);
5716 };
5717 match agg {
5718 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5720 self.live_count.saturating_sub(cs.null_count),
5721 ))),
5722 NativeAgg::Min | NativeAgg::Max => {
5723 let bound = if agg == NativeAgg::Min {
5724 &cs.min
5725 } else {
5726 &cs.max
5727 };
5728 match bound {
5729 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5730 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5731 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5736 None => Ok(None),
5737 }
5738 }
5739 _ => Ok(None),
5740 }
5741 }
5742
5743 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5752 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5753 return Ok(None);
5754 }
5755 self.ensure_indexes_complete()?;
5758 let reader = self.open_reader(self.run_refs[0].run_id)?;
5759 if self.live_count != reader.row_count() as u64 {
5760 return Ok(None);
5761 }
5762 let Some(bm) = self.bitmap.get(&column_id) else {
5763 return Ok(None); };
5765 let mut distinct = bm.value_count() as u64;
5766 if !bm.get(&Value::Null.encode_key()).is_empty() {
5769 distinct = distinct.saturating_sub(1);
5770 }
5771 Ok(Some(distinct))
5772 }
5773
5774 pub fn aggregate_incremental(
5786 &mut self,
5787 cache_key: u64,
5788 conditions: &[crate::query::Condition],
5789 column: Option<u16>,
5790 agg: NativeAgg,
5791 ) -> Result<IncrementalAggResult> {
5792 let snap = self.snapshot();
5793 let cur_wm = self.allocator.current().0;
5794 let cur_epoch = snap.epoch.0;
5795 let incremental_ok =
5802 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5803
5804 if incremental_ok {
5807 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5808 if cached.epoch == cur_epoch {
5809 return Ok(IncrementalAggResult {
5810 state: cached.state,
5811 incremental: true,
5812 delta_rows: 0,
5813 });
5814 }
5815 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5816 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5817 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5818 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5819 let delta_state = agg_state_from_rows(
5820 &delta_rows,
5821 conditions,
5822 &index_sets,
5823 column,
5824 agg,
5825 &self.schema,
5826 )?;
5827 let merged = cached.state.merge(delta_state);
5828 let delta_n = delta_rids.len() as u64;
5829 self.agg_cache.insert(
5830 cache_key,
5831 CachedAgg {
5832 state: merged.clone(),
5833 watermark: cur_wm,
5834 epoch: cur_epoch,
5835 },
5836 );
5837 return Ok(IncrementalAggResult {
5838 state: merged,
5839 incremental: true,
5840 delta_rows: delta_n,
5841 });
5842 }
5843 }
5844 }
5845
5846 let cursor_ok =
5851 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5852 let state = if cursor_ok && agg != NativeAgg::Avg {
5853 match self.aggregate_native(snap, column, conditions, agg)? {
5854 Some(result) => {
5855 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5856 }
5857 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5858 }
5859 } else {
5860 self.agg_state_full_scan(conditions, column, agg, snap)?
5861 };
5862 if incremental_ok {
5864 self.agg_cache.insert(
5865 cache_key,
5866 CachedAgg {
5867 state: state.clone(),
5868 watermark: cur_wm,
5869 epoch: cur_epoch,
5870 },
5871 );
5872 }
5873 Ok(IncrementalAggResult {
5874 state,
5875 incremental: false,
5876 delta_rows: 0,
5877 })
5878 }
5879
5880 fn agg_state_full_scan(
5883 &self,
5884 conditions: &[crate::query::Condition],
5885 column: Option<u16>,
5886 agg: NativeAgg,
5887 snap: Snapshot,
5888 ) -> Result<AggState> {
5889 let rows = self.visible_rows(snap)?;
5890 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5891 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5892 }
5893
5894 fn resolve_index_conditions(
5897 &self,
5898 conditions: &[crate::query::Condition],
5899 snapshot: Snapshot,
5900 ) -> Result<Vec<RowIdSet>> {
5901 use crate::query::Condition;
5902 let mut sets = Vec::new();
5903 for c in conditions {
5904 if matches!(
5905 c,
5906 Condition::Ann { .. }
5907 | Condition::SparseMatch { .. }
5908 | Condition::MinHashSimilar { .. }
5909 ) {
5910 sets.push(self.resolve_condition(c, snapshot)?);
5911 }
5912 }
5913 Ok(sets)
5914 }
5915
5916 fn column_type(&self, cid: u16) -> TypeId {
5917 self.schema
5918 .columns
5919 .iter()
5920 .find(|c| c.id == cid)
5921 .map(|c| c.ty.clone())
5922 .unwrap_or(TypeId::Bytes)
5923 }
5924
5925 pub fn approx_aggregate(
5934 &mut self,
5935 conditions: &[crate::query::Condition],
5936 column: Option<u16>,
5937 agg: ApproxAgg,
5938 z: f64,
5939 ) -> Result<Option<ApproxResult>> {
5940 use crate::query::Condition;
5941 self.ensure_reservoir_complete()?;
5942 let snapshot = self.snapshot();
5943 let n_pop = self.live_count;
5944 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5945 if sample_rids.is_empty() {
5946 return Ok(None);
5947 }
5948 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5950 let s = live_sample.len();
5951 if s == 0 {
5952 return Ok(None);
5953 }
5954
5955 let mut index_sets: Vec<RowIdSet> = Vec::new();
5958 for c in conditions {
5959 if matches!(
5960 c,
5961 Condition::Ann { .. }
5962 | Condition::SparseMatch { .. }
5963 | Condition::MinHashSimilar { .. }
5964 ) {
5965 index_sets.push(self.resolve_condition(c, snapshot)?);
5966 }
5967 }
5968
5969 let cid = match (agg, column) {
5971 (ApproxAgg::Count, _) => None,
5972 (_, Some(c)) => Some(c),
5973 _ => return Ok(None),
5974 };
5975 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5976 for r in &live_sample {
5977 if !conditions
5979 .iter()
5980 .all(|c| condition_matches_row(c, r, &self.schema))
5981 {
5982 continue;
5983 }
5984 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5986 continue;
5987 }
5988 if let Some(cid) = cid {
5989 if let Some(v) = as_f64(r.columns.get(&cid)) {
5990 passing_vals.push(v);
5991 } } else {
5993 passing_vals.push(0.0); }
5995 }
5996 let m = passing_vals.len();
5997
5998 let (point, half) = match agg {
5999 ApproxAgg::Count => {
6000 let p = m as f64 / s as f64;
6002 let point = n_pop as f64 * p;
6003 let var = if s > 1 {
6004 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
6005 * (1.0 - s as f64 / n_pop as f64).max(0.0)
6006 } else {
6007 0.0
6008 };
6009 (point, z * var.sqrt())
6010 }
6011 ApproxAgg::Sum => {
6012 let y: Vec<f64> = live_sample
6014 .iter()
6015 .map(|r| {
6016 let passes_row = conditions
6017 .iter()
6018 .all(|c| condition_matches_row(c, r, &self.schema))
6019 && index_sets.iter().all(|set| set.contains(r.row_id.0));
6020 if passes_row {
6021 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
6022 } else {
6023 0.0
6024 }
6025 })
6026 .collect();
6027 let mean_y = y.iter().sum::<f64>() / s as f64;
6028 let point = n_pop as f64 * mean_y;
6029 let var = if s > 1 {
6030 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
6031 let var_y = ss / (s - 1) as f64;
6032 n_pop as f64 * n_pop as f64 * var_y / s as f64
6033 * (1.0 - s as f64 / n_pop as f64).max(0.0)
6034 } else {
6035 0.0
6036 };
6037 (point, z * var.sqrt())
6038 }
6039 ApproxAgg::Avg => {
6040 if m == 0 {
6041 return Ok(Some(ApproxResult {
6042 point: 0.0,
6043 ci_low: 0.0,
6044 ci_high: 0.0,
6045 n_population: n_pop,
6046 n_sample_live: s,
6047 n_passing: 0,
6048 }));
6049 }
6050 let mean = passing_vals.iter().sum::<f64>() / m as f64;
6051 let half = if m > 1 {
6052 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
6053 let sd = (ss / (m - 1) as f64).sqrt();
6054 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
6055 z * sd / (m as f64).sqrt() * fpc.sqrt()
6056 } else {
6057 0.0
6058 };
6059 (mean, half)
6060 }
6061 };
6062
6063 Ok(Some(ApproxResult {
6064 point,
6065 ci_low: point - half,
6066 ci_high: point + half,
6067 n_population: n_pop,
6068 n_sample_live: s,
6069 n_passing: m,
6070 }))
6071 }
6072
6073 pub fn exact_column_stats(
6081 &self,
6082 _snapshot: Snapshot,
6083 projection: &[u16],
6084 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
6085 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
6086 return Ok(None);
6087 }
6088 let reader = self.open_reader(self.run_refs[0].run_id)?;
6089 if self.live_count != reader.row_count() as u64 {
6090 return Ok(None);
6091 }
6092 let mut out = HashMap::new();
6093 for &cid in projection {
6094 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
6095 Some(c) => c,
6096 None => continue,
6097 };
6098 let Some(stats) = reader.column_page_stats(cid) else {
6100 out.insert(
6101 cid,
6102 ColumnStat {
6103 min: None,
6104 max: None,
6105 null_count: self.live_count,
6106 },
6107 );
6108 continue;
6109 };
6110 let stat = match cdef.ty {
6111 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
6112 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
6113 min: mn.map(Value::Int64),
6114 max: mx.map(Value::Int64),
6115 null_count: n,
6116 })
6117 }
6118 TypeId::Float64 => {
6119 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
6120 min: mn.map(Value::Float64),
6121 max: mx.map(Value::Float64),
6122 null_count: n,
6123 })
6124 }
6125 _ => None,
6126 };
6127 if let Some(s) = stat {
6128 out.insert(cid, s);
6129 }
6130 }
6131 Ok(Some(out))
6132 }
6133
6134 pub fn dir(&self) -> &Path {
6135 &self.dir
6136 }
6137
6138 pub fn schema(&self) -> &Schema {
6139 &self.schema
6140 }
6141
6142 pub(crate) fn prepare_alter_column(
6143 &mut self,
6144 column_name: &str,
6145 change: &AlterColumn,
6146 ) -> Result<ColumnDef> {
6147 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
6148 return Err(MongrelError::InvalidArgument(
6149 "ALTER COLUMN requires committing staged writes first".into(),
6150 ));
6151 }
6152 let old = self
6153 .schema
6154 .columns
6155 .iter()
6156 .find(|c| c.name == column_name)
6157 .cloned()
6158 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6159 let mut next = old.clone();
6160
6161 if let Some(name) = &change.name {
6162 let trimmed = name.trim();
6163 if trimmed.is_empty() {
6164 return Err(MongrelError::InvalidArgument(
6165 "ALTER COLUMN name must not be empty".into(),
6166 ));
6167 }
6168 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
6169 return Err(MongrelError::Schema(format!(
6170 "column {trimmed} already exists"
6171 )));
6172 }
6173 next.name = trimmed.to_string();
6174 }
6175
6176 if let Some(ty) = &change.ty {
6177 next.ty = ty.clone();
6178 }
6179 if let Some(flags) = change.flags {
6180 validate_alter_column_flags(old.flags, flags)?;
6181 next.flags = flags;
6182 }
6183
6184 if let Some(default_change) = &change.default_value {
6185 next.default_value = default_change.clone();
6186 }
6187
6188 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
6189 if old.flags.contains(ColumnFlags::NULLABLE)
6190 && !next.flags.contains(ColumnFlags::NULLABLE)
6191 && self.column_has_nulls(old.id)?
6192 {
6193 return Err(MongrelError::InvalidArgument(format!(
6194 "column '{}' contains NULL values",
6195 old.name
6196 )));
6197 }
6198 Ok(next)
6199 }
6200
6201 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
6202 let idx = self
6203 .schema
6204 .columns
6205 .iter()
6206 .position(|c| c.id == column.id)
6207 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
6208 if self.schema.columns[idx] == column {
6209 return Ok(());
6210 }
6211 self.schema.columns[idx] = column;
6212 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6213 self.schema.validate_auto_increment()?;
6214 self.schema.validate_defaults()?;
6215 self.auto_inc = resolve_auto_inc(&self.schema);
6216 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
6217 write_schema(&self.dir, &self.schema)?;
6218 self.clear_result_cache();
6219 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6220 self.persist_manifest(self.current_epoch())?;
6221 Ok(())
6222 }
6223
6224 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
6225 let column = self.prepare_alter_column(column_name, &change)?;
6226 self.apply_altered_column(column.clone())?;
6227 Ok(column)
6228 }
6229
6230 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
6231 if self.live_count == 0 {
6232 return Ok(false);
6233 }
6234 let snap = self.snapshot();
6235 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6236 Ok(columns
6237 .first()
6238 .map(|(_, col)| col.null_count(col.len()) != 0)
6239 .unwrap_or(true))
6240 }
6241
6242 fn has_stored_versions(&self) -> bool {
6243 !self.memtable.is_empty()
6244 || !self.mutable_run.is_empty()
6245 || self.run_refs.iter().any(|r| r.row_count > 0)
6246 || !self.retiring.is_empty()
6247 }
6248
6249 pub fn add_column(
6254 &mut self,
6255 name: &str,
6256 ty: TypeId,
6257 flags: ColumnFlags,
6258 default_value: Option<crate::schema::DefaultExpr>,
6259 ) -> Result<u16> {
6260 if self.schema.columns.iter().any(|c| c.name == name) {
6261 return Err(MongrelError::Schema(format!(
6262 "column {name} already exists"
6263 )));
6264 }
6265 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
6266 self.schema.columns.push(ColumnDef {
6267 id,
6268 name: name.to_string(),
6269 ty,
6270 flags,
6271 default_value,
6272 });
6273 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6274 self.schema.validate_auto_increment()?;
6275 self.schema.validate_defaults()?;
6276 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6277 self.auto_inc = resolve_auto_inc(&self.schema);
6278 }
6279 write_schema(&self.dir, &self.schema)?;
6280 self.clear_result_cache();
6281 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6283 self.persist_manifest(self.current_epoch())?;
6284 Ok(id)
6285 }
6286
6287 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6296 let cid = self
6297 .schema
6298 .columns
6299 .iter()
6300 .find(|c| c.name == column_name)
6301 .map(|c| c.id)
6302 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6303 let ty = self
6304 .schema
6305 .columns
6306 .iter()
6307 .find(|c| c.id == cid)
6308 .map(|c| c.ty.clone())
6309 .unwrap_or(TypeId::Int64);
6310 if !matches!(
6311 ty,
6312 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6313 ) {
6314 return Err(MongrelError::Schema(format!(
6315 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6316 )));
6317 }
6318 if self
6319 .schema
6320 .indexes
6321 .iter()
6322 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6323 {
6324 return Ok(()); }
6326 self.schema.indexes.push(IndexDef {
6327 name: format!("{}_learned_range", column_name),
6328 column_id: cid,
6329 kind: IndexKind::LearnedRange,
6330 predicate: None,
6331 });
6332 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6333 write_schema(&self.dir, &self.schema)?;
6334 self.build_learned_ranges()?;
6335 Ok(())
6336 }
6337
6338 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6341 self.sync_byte_threshold = threshold;
6342 if let WalSink::Private(w) = &mut self.wal {
6343 w.set_sync_byte_threshold(threshold);
6344 }
6345 }
6346
6347 pub fn page_cache_flush(&self) {
6351 self.page_cache.flush_to_disk();
6352 }
6353
6354 pub fn page_cache_len(&self) -> usize {
6356 self.page_cache.len()
6357 }
6358
6359 pub fn decoded_cache_len(&self) -> usize {
6362 self.decoded_cache.len()
6363 }
6364
6365 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6368 self.memtable.drain_sorted()
6369 }
6370
6371 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6372 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6373 }
6374
6375 pub(crate) fn table_dir(&self) -> &Path {
6376 &self.dir
6377 }
6378
6379 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6380 &self.schema
6381 }
6382
6383 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6384 let id = self.next_run_id;
6385 self.next_run_id += 1;
6386 id
6387 }
6388
6389 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6390 self.run_refs.push(run_ref);
6391 }
6392
6393 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6403 self.retiring.push(crate::manifest::RetiredRun {
6404 run_id,
6405 retire_epoch,
6406 });
6407 }
6408
6409 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6413 if self.retiring.is_empty() {
6414 return Ok(0);
6415 }
6416 let mut reaped = 0;
6417 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6418 for r in std::mem::take(&mut self.retiring) {
6424 if min_active.0 >= r.retire_epoch {
6425 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6426 reaped += 1;
6427 } else {
6428 kept.push(r);
6429 }
6430 }
6431 self.retiring = kept;
6432 if reaped > 0 {
6433 self.persist_manifest(self.current_epoch())?;
6434 }
6435 Ok(reaped)
6436 }
6437
6438 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6439 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6440 return false;
6441 }
6442 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6443 self.run_refs.push(run_ref);
6444 self.indexes_complete = false;
6445 true
6446 }
6447
6448 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6449 self.kek.as_ref()
6450 }
6451
6452 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6453 let mut reader = RunReader::open_with_cache(
6454 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6455 self.schema.clone(),
6456 self.kek.clone(),
6457 Some(self.page_cache.clone()),
6458 Some(self.decoded_cache.clone()),
6459 self.table_id,
6460 Some(&self.verified_runs),
6461 )?;
6462 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6466 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6467 }
6468 Ok(reader)
6469 }
6470
6471 pub(crate) fn run_refs(&self) -> &[RunRef] {
6472 &self.run_refs
6473 }
6474
6475 pub(crate) fn runs_dir(&self) -> PathBuf {
6476 self.dir.join(RUNS_DIR)
6477 }
6478
6479 pub(crate) fn wal_dir(&self) -> PathBuf {
6480 self.dir.join(WAL_DIR)
6481 }
6482
6483 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6484 self.run_refs = refs;
6485 }
6486
6487 pub(crate) fn next_run_id(&self) -> u64 {
6488 self.next_run_id
6489 }
6490
6491 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6492 self.compaction_zstd_level
6493 }
6494
6495 pub(crate) fn bump_next_run_id(&mut self) {
6496 self.next_run_id += 1;
6497 }
6498
6499 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6500 self.kek.clone()
6501 }
6502
6503 #[cfg(feature = "encryption")]
6507 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6508 self.kek.as_ref().map(|k| k.derive_idx_key())
6509 }
6510
6511 #[cfg(not(feature = "encryption"))]
6512 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6513 None
6514 }
6515
6516 #[cfg(feature = "encryption")]
6520 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6521 self.kek.as_ref().map(|k| *k.derive_meta_key())
6522 }
6523
6524 #[cfg(not(feature = "encryption"))]
6525 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6526 None
6527 }
6528
6529 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6532 self.column_keys
6533 .iter()
6534 .map(|(&id, &(_, scheme))| (id, scheme))
6535 .collect()
6536 }
6537
6538 #[cfg(feature = "encryption")]
6543 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6544 self.tokenize_value_enc(column_id, v)
6545 }
6546
6547 #[cfg(feature = "encryption")]
6548 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6549 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6550 let (key, scheme) = self.column_keys.get(&column_id)?;
6551 let token: Vec<u8> = match (*scheme, v) {
6552 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6553 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6554 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6555 _ => hmac_token(key, &v.encode_key()).to_vec(),
6556 };
6557 Some(Value::Bytes(token))
6558 }
6559
6560 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6562 self.index_lookup_key_bytes(column_id, &v.encode_key())
6563 }
6564
6565 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6568 #[cfg(feature = "encryption")]
6569 {
6570 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6571 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6572 if *scheme == SCHEME_HMAC_EQ {
6573 return hmac_token(key, encoded).to_vec();
6574 }
6575 }
6576 }
6577 let _ = column_id;
6578 encoded.to_vec()
6579 }
6580}
6581
6582fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6583 let columnar::NativeColumn::Int64 { data, validity } = col else {
6584 return false;
6585 };
6586 if data.len() < n || !columnar::all_non_null(validity, n) {
6587 return false;
6588 }
6589 data.iter()
6590 .take(n)
6591 .zip(data.iter().skip(1))
6592 .all(|(a, b)| a < b)
6593}
6594
6595#[derive(Debug, Clone)]
6599pub struct ColumnStat {
6600 pub min: Option<Value>,
6601 pub max: Option<Value>,
6602 pub null_count: u64,
6603}
6604
6605#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6607pub enum NativeAgg {
6608 Count,
6609 Sum,
6610 Min,
6611 Max,
6612 Avg,
6613}
6614
6615#[derive(Debug, Clone, PartialEq)]
6617pub enum NativeAggResult {
6618 Count(u64),
6619 Int(i64),
6620 Float(f64),
6621 Null,
6623}
6624
6625#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6627pub enum ApproxAgg {
6628 Count,
6629 Sum,
6630 Avg,
6631}
6632
6633#[derive(Debug, Clone)]
6637pub struct ApproxResult {
6638 pub point: f64,
6640 pub ci_low: f64,
6642 pub ci_high: f64,
6644 pub n_population: u64,
6646 pub n_sample_live: usize,
6648 pub n_passing: usize,
6650}
6651
6652#[derive(Debug, Clone, PartialEq)]
6657pub enum AggState {
6658 Count(u64),
6660 SumI {
6662 sum: i128,
6663 count: u64,
6664 },
6665 SumF {
6667 sum: f64,
6668 count: u64,
6669 },
6670 AvgI {
6672 sum: i128,
6673 count: u64,
6674 },
6675 AvgF {
6677 sum: f64,
6678 count: u64,
6679 },
6680 MinI(i64),
6682 MaxI(i64),
6683 MinF(f64),
6685 MaxF(f64),
6686 Empty,
6688}
6689
6690impl AggState {
6691 pub fn merge(self, other: AggState) -> AggState {
6693 use AggState::*;
6694 match (self, other) {
6695 (Empty, x) | (x, Empty) => x,
6696 (Count(a), Count(b)) => Count(a + b),
6697 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6698 sum: sa + sb,
6699 count: ca + cb,
6700 },
6701 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6702 sum: sa + sb,
6703 count: ca + cb,
6704 },
6705 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6706 sum: sa + sb,
6707 count: ca + cb,
6708 },
6709 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6710 sum: sa + sb,
6711 count: ca + cb,
6712 },
6713 (MinI(a), MinI(b)) => MinI(a.min(b)),
6714 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6715 (MinF(a), MinF(b)) => MinF(a.min(b)),
6716 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6717 _ => Empty, }
6719 }
6720
6721 pub fn point(&self) -> Option<f64> {
6723 match self {
6724 AggState::Count(n) => Some(*n as f64),
6725 AggState::SumI { sum, .. } => Some(*sum as f64),
6726 AggState::SumF { sum, .. } => Some(*sum),
6727 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6728 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6729 AggState::MinI(n) => Some(*n as f64),
6730 AggState::MaxI(n) => Some(*n as f64),
6731 AggState::MinF(n) => Some(*n),
6732 AggState::MaxF(n) => Some(*n),
6733 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6734 }
6735 }
6736
6737 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6741 let is_float = matches!(ty, Some(TypeId::Float64));
6742 match (agg, result) {
6743 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6744 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6745 sum: x as i128,
6746 count: 1, },
6748 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6749 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6750 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6751 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6752 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6753 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6754 (NativeAgg::Count, _) => AggState::Empty,
6755 (_, NativeAggResult::Null) => AggState::Empty,
6756 _ => {
6757 let _ = is_float;
6758 AggState::Empty
6759 }
6760 }
6761 }
6762}
6763
6764#[derive(Debug, Clone)]
6767pub struct CachedAgg {
6768 pub state: AggState,
6769 pub watermark: u64,
6770 pub epoch: u64,
6771}
6772
6773#[derive(Debug, Clone)]
6775pub struct IncrementalAggResult {
6776 pub state: AggState,
6778 pub incremental: bool,
6781 pub delta_rows: u64,
6783}
6784
6785fn agg_state_from_rows(
6789 rows: &[Row],
6790 conditions: &[crate::query::Condition],
6791 index_sets: &[RowIdSet],
6792 column: Option<u16>,
6793 agg: NativeAgg,
6794 schema: &Schema,
6795) -> Result<AggState> {
6796 let mut count: u64 = 0;
6797 let mut sum_i: i128 = 0;
6798 let mut sum_f: f64 = 0.0;
6799 let mut mn_i: i64 = i64::MAX;
6800 let mut mx_i: i64 = i64::MIN;
6801 let mut mn_f: f64 = f64::INFINITY;
6802 let mut mx_f: f64 = f64::NEG_INFINITY;
6803 let mut saw_int = false;
6804 let mut saw_float = false;
6805 for r in rows {
6806 if !conditions
6807 .iter()
6808 .all(|c| condition_matches_row(c, r, schema))
6809 {
6810 continue;
6811 }
6812 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6813 continue;
6814 }
6815 match agg {
6816 NativeAgg::Count => match column {
6817 None => count += 1,
6819 Some(cid) => match r.columns.get(&cid) {
6822 None | Some(Value::Null) => {}
6823 Some(_) => count += 1,
6824 },
6825 },
6826 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6827 Some(Value::Int64(n)) => {
6828 count += 1;
6829 sum_i += *n as i128;
6830 mn_i = mn_i.min(*n);
6831 mx_i = mx_i.max(*n);
6832 saw_int = true;
6833 }
6834 Some(Value::Float64(f)) => {
6835 count += 1;
6836 sum_f += f;
6837 mn_f = mn_f.min(*f);
6838 mx_f = mx_f.max(*f);
6839 saw_float = true;
6840 }
6841 _ => {}
6842 },
6843 }
6844 }
6845 Ok(match agg {
6846 NativeAgg::Count => {
6847 if count == 0 {
6848 AggState::Empty
6849 } else {
6850 AggState::Count(count)
6851 }
6852 }
6853 NativeAgg::Sum => {
6854 if count == 0 {
6855 AggState::Empty
6856 } else if saw_int {
6857 AggState::SumI { sum: sum_i, count }
6858 } else {
6859 AggState::SumF { sum: sum_f, count }
6860 }
6861 }
6862 NativeAgg::Avg => {
6863 if count == 0 {
6864 AggState::Empty
6865 } else if saw_int {
6866 AggState::AvgI { sum: sum_i, count }
6867 } else {
6868 AggState::AvgF { sum: sum_f, count }
6869 }
6870 }
6871 NativeAgg::Min => {
6872 if !saw_int && !saw_float {
6873 AggState::Empty
6874 } else if saw_int {
6875 AggState::MinI(mn_i)
6876 } else {
6877 AggState::MinF(mn_f)
6878 }
6879 }
6880 NativeAgg::Max => {
6881 if !saw_int && !saw_float {
6882 AggState::Empty
6883 } else if saw_int {
6884 AggState::MaxI(mx_i)
6885 } else {
6886 AggState::MaxF(mx_f)
6887 }
6888 }
6889 })
6890}
6891
6892fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6896 use crate::query::Condition;
6897 match c {
6898 Condition::Pk(key) => match schema.primary_key() {
6899 Some(pk) => row
6900 .columns
6901 .get(&pk.id)
6902 .map(|v| v.encode_key() == *key)
6903 .unwrap_or(false),
6904 None => false,
6905 },
6906 Condition::BitmapEq { column_id, value } => row
6907 .columns
6908 .get(column_id)
6909 .map(|v| v.encode_key() == *value)
6910 .unwrap_or(false),
6911 Condition::BitmapIn { column_id, values } => {
6912 let key = row.columns.get(column_id).map(|v| v.encode_key());
6913 match key {
6914 Some(k) => values.contains(&k),
6915 None => false,
6916 }
6917 }
6918 Condition::BytesPrefix { column_id, prefix } => row
6919 .columns
6920 .get(column_id)
6921 .map(|v| v.encode_key().starts_with(prefix))
6922 .unwrap_or(false),
6923 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6924 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6925 _ => false,
6926 },
6927 Condition::RangeF64 {
6928 column_id,
6929 lo,
6930 lo_inclusive,
6931 hi,
6932 hi_inclusive,
6933 } => match row.columns.get(column_id) {
6934 Some(Value::Float64(n)) => {
6935 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6936 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6937 lo_ok && hi_ok
6938 }
6939 _ => false,
6940 },
6941 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6942 Some(Value::Bytes(b)) => {
6943 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6944 }
6945 _ => false,
6946 },
6947 Condition::FmContainsAll {
6948 column_id,
6949 patterns,
6950 } => match row.columns.get(column_id) {
6951 Some(Value::Bytes(b)) => patterns
6952 .iter()
6953 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6954 _ => false,
6955 },
6956 Condition::Ann { .. }
6957 | Condition::SparseMatch { .. }
6958 | Condition::MinHashSimilar { .. } => true,
6959 Condition::IsNull { column_id } => {
6960 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6961 }
6962 Condition::IsNotNull { column_id } => {
6963 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6964 }
6965 }
6966}
6967
6968fn as_f64(v: Option<&Value>) -> Option<f64> {
6970 match v {
6971 Some(Value::Int64(n)) => Some(*n as f64),
6972 Some(Value::Float64(f)) => Some(*f),
6973 _ => None,
6974 }
6975}
6976
6977fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6981 let mut count: u64 = 0;
6982 let mut sum: i128 = 0;
6983 let mut mn: i64 = i64::MAX;
6984 let mut mx: i64 = i64::MIN;
6985 while let Some(cols) = cursor.next_batch()? {
6986 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6987 if crate::columnar::all_non_null(validity, data.len()) {
6988 count += data.len() as u64;
6990 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6991 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6992 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6993 } else {
6994 for (i, &v) in data.iter().enumerate() {
6995 if crate::columnar::validity_bit(validity, i) {
6996 count += 1;
6997 sum += v as i128;
6998 mn = mn.min(v);
6999 mx = mx.max(v);
7000 }
7001 }
7002 }
7003 }
7004 }
7005 Ok((count, sum, mn, mx))
7006}
7007
7008fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
7010 let mut count: u64 = 0;
7011 let mut sum: f64 = 0.0;
7012 let mut mn: f64 = f64::INFINITY;
7013 let mut mx: f64 = f64::NEG_INFINITY;
7014 while let Some(cols) = cursor.next_batch()? {
7015 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
7016 if crate::columnar::all_non_null(validity, data.len()) {
7017 count += data.len() as u64;
7018 sum += data.iter().sum::<f64>();
7019 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
7020 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
7021 } else {
7022 for (i, &v) in data.iter().enumerate() {
7023 if crate::columnar::validity_bit(validity, i) {
7024 count += 1;
7025 sum += v;
7026 mn = mn.min(v);
7027 mx = mx.max(v);
7028 }
7029 }
7030 }
7031 }
7032 }
7033 Ok((count, sum, mn, mx))
7034}
7035
7036fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
7037 if count == 0 && !matches!(agg, NativeAgg::Count) {
7038 return NativeAggResult::Null;
7039 }
7040 match agg {
7041 NativeAgg::Count => NativeAggResult::Count(count),
7042 NativeAgg::Sum => match sum.try_into() {
7045 Ok(v) => NativeAggResult::Int(v),
7046 Err(_) => NativeAggResult::Null,
7047 },
7048 NativeAgg::Min => NativeAggResult::Int(mn),
7049 NativeAgg::Max => NativeAggResult::Int(mx),
7050 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
7051 }
7052}
7053
7054fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
7055 if count == 0 && !matches!(agg, NativeAgg::Count) {
7056 return NativeAggResult::Null;
7057 }
7058 match agg {
7059 NativeAgg::Count => NativeAggResult::Count(count),
7060 NativeAgg::Sum => NativeAggResult::Float(sum),
7061 NativeAgg::Min => NativeAggResult::Float(mn),
7062 NativeAgg::Max => NativeAggResult::Float(mx),
7063 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
7064 }
7065}
7066
7067fn agg_int(
7070 stats: &[crate::page::PageStat],
7071 decode: fn(Option<&[u8]>) -> Option<i64>,
7072) -> Option<(Option<i64>, Option<i64>, u64)> {
7073 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
7074 let mut any = false;
7075 for s in stats {
7076 if let Some(v) = decode(s.min.as_deref()) {
7077 mn = mn.min(v);
7078 any = true;
7079 }
7080 if let Some(v) = decode(s.max.as_deref()) {
7081 mx = mx.max(v);
7082 any = true;
7083 }
7084 nulls += s.null_count;
7085 }
7086 any.then_some((Some(mn), Some(mx), nulls))
7087}
7088
7089fn agg_float(
7091 stats: &[crate::page::PageStat],
7092 decode: fn(Option<&[u8]>) -> Option<f64>,
7093) -> Option<(Option<f64>, Option<f64>, u64)> {
7094 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
7095 let mut any = false;
7096 for s in stats {
7097 if let Some(v) = decode(s.min.as_deref()) {
7098 mn = mn.min(v);
7099 any = true;
7100 }
7101 if let Some(v) = decode(s.max.as_deref()) {
7102 mx = mx.max(v);
7103 any = true;
7104 }
7105 nulls += s.null_count;
7106 }
7107 any.then_some((Some(mn), Some(mx), nulls))
7108}
7109
7110type SecondaryIndexes = (
7112 HashMap<u16, BitmapIndex>,
7113 HashMap<u16, AnnIndex>,
7114 HashMap<u16, FmIndex>,
7115 HashMap<u16, SparseIndex>,
7116 HashMap<u16, MinHashIndex>,
7117);
7118
7119fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
7120 let mut bitmap = HashMap::new();
7121 let mut ann = HashMap::new();
7122 let mut fm = HashMap::new();
7123 let mut sparse = HashMap::new();
7124 let mut minhash = HashMap::new();
7125 for idef in &schema.indexes {
7126 match idef.kind {
7127 IndexKind::Bitmap => {
7128 bitmap.insert(idef.column_id, BitmapIndex::new());
7129 }
7130 IndexKind::Ann => {
7131 let dim = schema
7132 .columns
7133 .iter()
7134 .find(|c| c.id == idef.column_id)
7135 .and_then(|c| match c.ty {
7136 TypeId::Embedding { dim } => Some(dim as usize),
7137 _ => None,
7138 })
7139 .unwrap_or(0);
7140 ann.insert(idef.column_id, AnnIndex::new(dim));
7141 }
7142 IndexKind::FmIndex => {
7143 fm.insert(idef.column_id, FmIndex::new());
7144 }
7145 IndexKind::Sparse => {
7146 sparse.insert(idef.column_id, SparseIndex::new());
7147 }
7148 IndexKind::MinHash => {
7149 minhash.insert(idef.column_id, MinHashIndex::new());
7150 }
7151 _ => {}
7152 }
7153 }
7154 (bitmap, ann, fm, sparse, minhash)
7155}
7156
7157const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
7158 | ColumnFlags::AUTO_INCREMENT
7159 | ColumnFlags::ENCRYPTED
7160 | ColumnFlags::ENCRYPTED_INDEXABLE
7161 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
7162
7163fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
7164 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
7165 return Err(MongrelError::Schema(
7166 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
7167 ));
7168 }
7169 Ok(())
7170}
7171
7172fn validate_alter_column_type(
7173 schema: &Schema,
7174 old: &ColumnDef,
7175 next: &ColumnDef,
7176 has_stored_versions: bool,
7177) -> Result<()> {
7178 if old.ty == next.ty {
7179 return Ok(());
7180 }
7181 if schema.indexes.iter().any(|i| i.column_id == old.id) {
7182 return Err(MongrelError::Schema(format!(
7183 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
7184 old.name
7185 )));
7186 }
7187 if !has_stored_versions || storage_compatible_type_change(old.ty.clone(), next.ty.clone()) {
7188 return Ok(());
7189 }
7190 Err(MongrelError::Schema(format!(
7191 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
7192 old.ty, next.ty
7193 )))
7194}
7195
7196fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
7197 matches!(
7198 (old, new),
7199 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
7200 )
7201}
7202
7203fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
7209 let mut prev: Option<i64> = None;
7210 for r in rows {
7211 match r.columns.get(&pk_id) {
7212 Some(Value::Int64(v)) => {
7213 if prev.is_some_and(|p| p >= *v) {
7214 return false;
7215 }
7216 prev = Some(*v);
7217 }
7218 _ => return false,
7219 }
7220 }
7221 true
7222}
7223
7224#[allow(clippy::too_many_arguments)]
7225fn index_into(
7226 schema: &Schema,
7227 row: &Row,
7228 hot: &mut HotIndex,
7229 bitmap: &mut HashMap<u16, BitmapIndex>,
7230 ann: &mut HashMap<u16, AnnIndex>,
7231 fm: &mut HashMap<u16, FmIndex>,
7232 sparse: &mut HashMap<u16, SparseIndex>,
7233 minhash: &mut HashMap<u16, MinHashIndex>,
7234) {
7235 for idef in &schema.indexes {
7236 let Some(val) = row.columns.get(&idef.column_id) else {
7237 continue;
7238 };
7239 match idef.kind {
7240 IndexKind::Bitmap => {
7241 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7242 b.insert(val.encode_key(), row.row_id);
7243 }
7244 }
7245 IndexKind::Ann => {
7246 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7247 a.insert(v, row.row_id);
7248 }
7249 }
7250 IndexKind::FmIndex => {
7251 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7252 f.insert(b.clone(), row.row_id);
7253 }
7254 }
7255 IndexKind::Sparse => {
7256 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7257 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7260 s.insert(&terms, row.row_id);
7261 }
7262 }
7263 }
7264 IndexKind::MinHash => {
7265 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7266 let tokens = crate::index::token_hashes_from_bytes(b);
7269 mh.insert(&tokens, row.row_id);
7270 }
7271 }
7272 _ => {}
7273 }
7274 }
7275 if let Some(pk_col) = schema.primary_key() {
7276 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7277 hot.insert(pk_val.encode_key(), row.row_id);
7278 }
7279 }
7280}
7281
7282#[allow(clippy::too_many_arguments)]
7285fn index_into_single(
7286 idef: &IndexDef,
7287 _schema: &Schema,
7288 row: &Row,
7289 _hot: &mut HotIndex,
7290 bitmap: &mut HashMap<u16, BitmapIndex>,
7291 ann: &mut HashMap<u16, AnnIndex>,
7292 fm: &mut HashMap<u16, FmIndex>,
7293 sparse: &mut HashMap<u16, SparseIndex>,
7294 minhash: &mut HashMap<u16, MinHashIndex>,
7295) {
7296 let Some(val) = row.columns.get(&idef.column_id) else {
7297 return;
7298 };
7299 match idef.kind {
7300 IndexKind::Bitmap => {
7301 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7302 b.insert(val.encode_key(), row.row_id);
7303 }
7304 }
7305 IndexKind::Ann => {
7306 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7307 a.insert(v, row.row_id);
7308 }
7309 }
7310 IndexKind::FmIndex => {
7311 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7312 f.insert(b.clone(), row.row_id);
7313 }
7314 }
7315 IndexKind::Sparse => {
7316 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7317 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7318 s.insert(&terms, row.row_id);
7319 }
7320 }
7321 }
7322 IndexKind::MinHash => {
7323 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7324 let tokens = crate::index::token_hashes_from_bytes(b);
7325 mh.insert(&tokens, row.row_id);
7326 }
7327 }
7328 _ => {}
7329 }
7330}
7331
7332fn eval_partial_predicate(
7338 pred: &str,
7339 columns_map: &HashMap<u16, &Value>,
7340 name_to_id: &HashMap<&str, u16>,
7341) -> bool {
7342 let lower = pred.trim().to_ascii_lowercase();
7343 if let Some(rest) = lower.strip_suffix(" is not null") {
7345 let col_name = rest.trim();
7346 if let Some(col_id) = name_to_id.get(col_name) {
7347 return columns_map
7348 .get(col_id)
7349 .is_some_and(|v| !matches!(v, Value::Null));
7350 }
7351 }
7352 if let Some(rest) = lower.strip_suffix(" is null") {
7354 let col_name = rest.trim();
7355 if let Some(col_id) = name_to_id.get(col_name) {
7356 return columns_map
7357 .get(col_id)
7358 .map_or(true, |v| matches!(v, Value::Null));
7359 }
7360 }
7361 true
7364}
7365
7366#[allow(dead_code)]
7372fn bulk_index_key(
7373 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7374 column_id: u16,
7375 ty: TypeId,
7376 col: &columnar::NativeColumn,
7377 i: usize,
7378) -> Option<Vec<u8>> {
7379 let encoded = columnar::encode_key_native(ty, col, i)?;
7380 #[cfg(feature = "encryption")]
7381 {
7382 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7383 if let Some((key, scheme)) = column_keys.get(&column_id) {
7384 return Some(match (*scheme, col) {
7385 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7386 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7387 ope_token_i64(key, data[i]).to_vec()
7388 }
7389 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7390 ope_token_f64(key, data[i]).to_vec()
7391 }
7392 _ => hmac_token(key, &encoded).to_vec(),
7393 });
7394 }
7395 }
7396 #[cfg(not(feature = "encryption"))]
7397 {
7398 let _ = (column_id, column_keys, col);
7399 }
7400 Some(encoded)
7401}
7402
7403pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7404 let json = serde_json::to_string_pretty(schema)
7405 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7406 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7407 Ok(())
7408}
7409
7410fn read_schema(dir: &Path) -> Result<Schema> {
7411 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7412 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7413}
7414
7415fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7416 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7417}
7418
7419fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7420 let n = list_wal_numbers(wal_dir)?;
7421 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7422}
7423
7424fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7425 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7426}
7427
7428fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7429 let _ = std::fs::create_dir_all(wal_dir);
7430 let mut max_n = None;
7431 for entry in std::fs::read_dir(wal_dir)? {
7432 let entry = entry?;
7433 let fname = entry.file_name();
7434 let Some(s) = fname.to_str() else {
7435 continue;
7436 };
7437 let Some(stripped) = s.strip_prefix("seg-") else {
7438 continue;
7439 };
7440 let Some(stripped) = stripped.strip_suffix(".wal") else {
7441 continue;
7442 };
7443 if let Ok(n) = stripped.parse::<u32>() {
7444 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7445 }
7446 }
7447 Ok(max_n)
7448}