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";
42const 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;
49
50#[derive(Clone, Copy, Debug)]
65struct AutoIncState {
66 column_id: u16,
67 next: i64,
68 seeded: bool,
69}
70
71type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
72
73fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
76 schema.auto_increment_column().map(|c| AutoIncState {
77 column_id: c.id,
78 next: 0,
79 seeded: false,
80 })
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
96pub enum IndexBuildPolicy {
97 #[default]
99 Deferred,
100 Eager,
102}
103
104pub struct Table {
106 dir: PathBuf,
107 table_id: u64,
108 name: String,
112 auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
117 wal: WalSink,
118 memtable: Memtable,
119 mutable_run: MutableRun,
124 mutable_run_spill_bytes: u64,
126 compaction_zstd_level: i32,
129 allocator: RowIdAllocator,
130 epoch: Arc<EpochAuthority>,
131 persisted_epoch: u64,
134 schema: Schema,
135 hot: HotIndex,
136 kek: Option<Arc<Kek>>,
139 column_keys: HashMap<u16, ([u8; 32], u8)>,
143 run_refs: Vec<RunRef>,
144 retiring: Vec<crate::manifest::RetiredRun>,
147 next_run_id: u64,
148 sync_byte_threshold: u64,
149 current_txn_id: u64,
154 bitmap: HashMap<u16, BitmapIndex>,
155 ann: HashMap<u16, AnnIndex>,
156 fm: HashMap<u16, FmIndex>,
157 sparse: HashMap<u16, SparseIndex>,
158 minhash: HashMap<u16, MinHashIndex>,
159 learned_range: HashMap<u16, ColumnLearnedRange>,
162 pk_by_row: HashMap<RowId, Vec<u8>>,
164 pinned: BTreeMap<Epoch, usize>,
167 pub(crate) live_count: u64,
170 reservoir: crate::reservoir::Reservoir,
173 reservoir_complete: bool,
181 had_deletes: bool,
185 agg_cache: HashMap<u64, CachedAgg>,
189 global_idx_epoch: u64,
193 indexes_complete: bool,
198 index_build_policy: IndexBuildPolicy,
200 pk_by_row_complete: bool,
207 flushed_epoch: u64,
210 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
213 snapshots: Arc<crate::retention::SnapshotRegistry>,
216 commit_lock: Arc<parking_lot::Mutex<()>>,
218 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
221 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
231 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
240 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
242 pending_delete_rids: roaring::RoaringBitmap,
245 pending_put_cols: std::collections::HashSet<u16>,
248 pending_rows: Vec<Row>,
254 pending_rows_auto_inc: Vec<bool>,
255 pending_dels: Vec<RowId>,
258 pending_truncate: Option<Epoch>,
262 auto_inc: Option<AutoIncState>,
265}
266
267const _: () = {
274 const fn assert_sync<T: ?Sized + Sync>() {}
275 assert_sync::<Table>();
276};
277
278enum CachedData {
284 Rows(Arc<Vec<Row>>),
285 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
286}
287
288impl CachedData {
289 fn approx_bytes(&self) -> u64 {
290 match self {
291 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
292 CachedData::Columns(c) => c
293 .iter()
294 .map(|(_, c)| c.approx_bytes())
295 .sum::<u64>()
296 .saturating_add(c.len() as u64 * 16),
297 }
298 }
299}
300
301struct CachedEntry {
305 data: CachedData,
306 footprint: roaring::RoaringBitmap,
307 condition_cols: Vec<u16>,
308}
309
310struct ResultCache {
321 entries: std::collections::HashMap<u64, CachedEntry>,
322 order: std::collections::VecDeque<u64>,
323 bytes: u64,
324 max_bytes: u64,
325 dir: Option<std::path::PathBuf>,
326 #[allow(dead_code)]
327 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
328}
329
330#[derive(serde::Serialize, serde::Deserialize)]
332struct SerializedEntry {
333 condition_cols: Vec<u16>,
334 footprint_bits: Vec<u32>,
335 data: SerializedData,
336}
337
338#[derive(serde::Serialize, serde::Deserialize)]
339enum SerializedData {
340 Rows(Vec<Row>),
341 Columns(Vec<(u16, columnar::NativeColumn)>),
342}
343
344impl SerializedEntry {
345 fn from_entry(entry: &CachedEntry) -> Self {
346 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
347 let data = match &entry.data {
348 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
349 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
350 };
351 Self {
352 condition_cols: entry.condition_cols.clone(),
353 footprint_bits,
354 data,
355 }
356 }
357
358 fn into_entry(self) -> Option<CachedEntry> {
359 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
360 let data = match self.data {
361 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
362 SerializedData::Columns(c) => {
363 if !c.iter().all(|(_, col)| col.validate()) {
366 return None;
367 }
368 CachedData::Columns(Arc::new(c))
369 }
370 };
371 Some(CachedEntry {
372 data,
373 footprint,
374 condition_cols: self.condition_cols,
375 })
376 }
377}
378
379impl ResultCache {
380 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
381
382 fn new() -> Self {
383 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
384 }
385
386 fn with_max_bytes(max_bytes: u64) -> Self {
387 Self {
388 entries: std::collections::HashMap::new(),
389 order: std::collections::VecDeque::new(),
390 bytes: 0,
391 max_bytes,
392 dir: None,
393 cache_dek: None,
394 }
395 }
396
397 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
398 let _ = std::fs::create_dir_all(&dir);
399 self.dir = Some(dir);
400 self
401 }
402
403 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
404 self.cache_dek = dek;
405 self
406 }
407
408 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
409 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
410 }
411
412 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
416 let Some(path) = self.disk_path(key) else {
417 return;
418 };
419 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
420 Ok(s) => s,
421 Err(_) => return,
422 };
423 let on_disk = if let Some(dek) = &self.cache_dek {
425 match self.encrypt_cache(&serialized, dek) {
426 Some(b) => b,
427 None => return,
428 }
429 } else {
430 serialized
431 };
432 let tmp = path.with_extension("tmp");
433 use std::io::Write;
434 let write = || -> std::io::Result<()> {
435 let mut f = std::fs::File::create(&tmp)?;
436 f.write_all(&on_disk)?;
437 f.flush()?;
438 Ok(())
439 };
440 if write().is_err() {
441 let _ = std::fs::remove_file(&tmp);
442 return;
443 }
444 let _ = std::fs::rename(&tmp, &path);
445 }
446
447 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
449 let path = self.disk_path(key)?;
450 let bytes = std::fs::read(&path).ok()?;
451 let plaintext = if let Some(dek) = &self.cache_dek {
452 self.decrypt_cache(&bytes, dek)?
453 } else {
454 bytes
455 };
456 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
457 serialized.into_entry()
458 }
459
460 fn remove_from_disk(&self, key: u64) {
462 if let Some(path) = self.disk_path(key) {
463 let _ = std::fs::remove_file(&path);
464 }
465 }
466
467 #[cfg(feature = "encryption")]
469 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
470 use crate::encryption::Cipher;
471 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
472 let mut nonce = [0u8; 12];
473 crate::encryption::fill_random(&mut nonce);
474 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
475 let mut out = Vec::with_capacity(12 + ct.len());
476 out.extend_from_slice(&nonce);
477 out.extend_from_slice(&ct);
478 Some(out)
479 }
480
481 #[cfg(not(feature = "encryption"))]
482 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
483 None
484 }
485
486 #[cfg(feature = "encryption")]
488 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
489 use crate::encryption::Cipher;
490 if bytes.len() < 28 {
491 return None;
492 }
493 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
494 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
495 let ct = &bytes[12..];
496 cipher.decrypt_page(&nonce, ct).ok()
497 }
498
499 #[cfg(not(feature = "encryption"))]
500 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
501 None
502 }
503
504 fn load_persistent(&mut self) {
507 let Some(dir) = self.dir.as_ref().cloned() else {
508 return;
509 };
510 let entries = match std::fs::read_dir(&dir) {
511 Ok(e) => e,
512 Err(_) => return,
513 };
514 for entry in entries.flatten() {
515 let path = entry.path();
516 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
518 let _ = std::fs::remove_file(&path);
519 continue;
520 }
521 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
522 continue;
523 }
524 let stem = match path.file_stem().and_then(|s| s.to_str()) {
525 Some(s) => s,
526 None => continue,
527 };
528 let key = match u64::from_str_radix(stem, 16) {
529 Ok(k) => k,
530 Err(_) => continue,
531 };
532 let bytes = match std::fs::read(&path) {
533 Ok(b) => b,
534 Err(_) => continue,
535 };
536 let plaintext = if let Some(dek) = &self.cache_dek {
538 match self.decrypt_cache(&bytes, dek) {
539 Some(p) => p,
540 None => {
541 let _ = std::fs::remove_file(&path);
542 continue;
543 }
544 }
545 } else {
546 bytes
547 };
548 match bincode::deserialize::<SerializedEntry>(&plaintext) {
549 Ok(serialized) => {
550 if let Some(entry) = serialized.into_entry() {
551 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
552 self.entries.insert(key, entry);
553 self.order.push_back(key);
554 } else {
555 let _ = std::fs::remove_file(&path);
556 }
557 }
558 Err(_) => {
559 let _ = std::fs::remove_file(&path);
560 }
561 }
562 }
563 self.evict();
564 }
565
566 fn set_max_bytes(&mut self, max_bytes: u64) {
567 self.max_bytes = max_bytes;
568 self.evict();
569 }
570
571 fn touch(&mut self, key: u64) {
573 self.order.retain(|k| *k != key);
574 self.order.push_back(key);
575 }
576
577 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
578 let res = self.entries.get(&key).and_then(|e| match &e.data {
579 CachedData::Rows(r) => Some(r.clone()),
580 CachedData::Columns(_) => None,
581 });
582 if res.is_some() {
583 self.touch(key);
584 return res;
585 }
586 if let Some(entry) = self.load_from_disk(key) {
588 let res = match &entry.data {
589 CachedData::Rows(r) => Some(r.clone()),
590 CachedData::Columns(_) => None,
591 };
592 if res.is_some() {
593 let approx = entry.data.approx_bytes();
594 self.bytes = self.bytes.saturating_add(approx);
595 self.entries.insert(key, entry);
596 self.order.push_back(key);
597 self.evict();
598 return res;
599 }
600 }
601 None
602 }
603
604 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
605 let res = self.entries.get(&key).and_then(|e| match &e.data {
606 CachedData::Columns(c) => Some(c.clone()),
607 CachedData::Rows(_) => None,
608 });
609 if res.is_some() {
610 self.touch(key);
611 return res;
612 }
613 if let Some(entry) = self.load_from_disk(key) {
615 let res = match &entry.data {
616 CachedData::Columns(c) => Some(c.clone()),
617 CachedData::Rows(_) => None,
618 };
619 if res.is_some() {
620 let approx = entry.data.approx_bytes();
621 self.bytes = self.bytes.saturating_add(approx);
622 self.entries.insert(key, entry);
623 self.order.push_back(key);
624 self.evict();
625 return res;
626 }
627 }
628 None
629 }
630
631 fn insert(&mut self, key: u64, entry: CachedEntry) {
632 let approx = entry.data.approx_bytes();
633 if self.entries.remove(&key).is_some() {
634 self.order.retain(|k| *k != key);
635 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
636 }
637 self.store_to_disk(key, &entry);
639 self.bytes = self.bytes.saturating_add(approx);
640 self.entries.insert(key, entry);
641 self.order.push_back(key);
642 self.evict();
643 }
644
645 fn invalidate(
654 &mut self,
655 delete_rids: &roaring::RoaringBitmap,
656 put_cols: &std::collections::HashSet<u16>,
657 ) {
658 if self.entries.is_empty() {
659 return;
660 }
661 let has_deletes = !delete_rids.is_empty();
662 let to_remove: std::collections::HashSet<u64> = self
663 .entries
664 .iter()
665 .filter(|(_, e)| {
666 let delete_hit = if e.footprint.is_empty() {
667 has_deletes
668 } else {
669 e.footprint.intersection_len(delete_rids) > 0
670 };
671 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
672 delete_hit || col_hit
673 })
674 .map(|(&k, _)| k)
675 .collect();
676 for key in &to_remove {
677 if let Some(e) = self.entries.remove(key) {
678 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
679 }
680 self.remove_from_disk(*key);
681 }
682 if !to_remove.is_empty() {
683 self.order.retain(|k| !to_remove.contains(k));
684 }
685 }
686
687 fn clear(&mut self) {
688 if let Some(dir) = &self.dir {
690 if let Ok(entries) = std::fs::read_dir(dir) {
691 for entry in entries.flatten() {
692 let path = entry.path();
693 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
694 let _ = std::fs::remove_file(&path);
695 }
696 }
697 }
698 }
699 self.entries.clear();
700 self.order.clear();
701 self.bytes = 0;
702 }
703
704 fn evict(&mut self) {
705 while self.bytes > self.max_bytes {
706 let Some(k) = self.order.pop_front() else {
707 break;
708 };
709 if let Some(e) = self.entries.remove(&k) {
710 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
711 self.remove_from_disk(k);
715 }
716 }
717 }
718}
719
720type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
727
728fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
729 let _ = kek;
730 #[cfg(feature = "encryption")]
731 {
732 if let Some(k) = kek {
733 return (
734 Some(k.derive_table_wal_key(_table_id)),
735 Some(k.derive_cache_key()),
736 );
737 }
738 }
739 (None, None)
740}
741
742#[cfg(feature = "encryption")]
744fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
745 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
746}
747
748#[cfg(not(feature = "encryption"))]
749fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
750 Box::new(crate::encryption::PlaintextCipher)
751}
752
753fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
754 let Some(kek) = kek else {
755 return HashMap::new();
756 };
757 #[cfg(feature = "encryption")]
758 {
759 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
760 schema
761 .columns
762 .iter()
763 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
764 .map(|c| {
765 let scheme = if schema
766 .indexes
767 .iter()
768 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
769 {
770 SCHEME_OPE_RANGE
771 } else {
772 SCHEME_HMAC_EQ
773 };
774 let key: [u8; 32] = *kek.derive_column_key(c.id);
775 (c.id, (key, scheme))
776 })
777 .collect()
778 }
779 #[cfg(not(feature = "encryption"))]
780 {
781 let _ = (kek, schema);
782 HashMap::new()
783 }
784}
785
786pub(crate) struct SharedCtx {
791 pub epoch: Arc<EpochAuthority>,
792 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
793 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
794 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
795 pub kek: Option<Arc<Kek>>,
796 pub commit_lock: Arc<parking_lot::Mutex<()>>,
802 pub shared: Option<SharedWalCtx>,
806 pub table_name: Option<String>,
809 pub auth: Option<Arc<dyn crate::auth_state::TableAuthChecker>>,
812}
813
814#[derive(Clone)]
820pub(crate) struct SharedWalCtx {
821 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
822 pub group: Arc<GroupCommit>,
823 pub poisoned: Arc<AtomicBool>,
824 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
825}
826
827enum WalSink {
830 Private(Wal),
831 Shared(SharedWalCtx),
832}
833
834impl SharedCtx {
835 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
839 let n_shards = if cache_dir.is_some() {
843 1
844 } else {
845 crate::cache::CACHE_SHARDS
846 };
847 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
848 let page_cache = if let Some(d) = cache_dir {
849 Arc::new(crate::cache::Sharded::new(1, || {
850 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
851 }))
852 } else {
853 Arc::new(crate::cache::Sharded::new(n_shards, || {
854 crate::cache::PageCache::new(per_shard)
855 }))
856 };
857 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
858 let decoded_cache = Arc::new(crate::cache::Sharded::new(
859 crate::cache::CACHE_SHARDS,
860 || crate::cache::DecodedPageCache::new(decoded_per_shard),
861 ));
862 Self {
863 epoch: Arc::new(EpochAuthority::new(0)),
864 page_cache,
865 decoded_cache,
866 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
867 kek,
868 commit_lock: Arc::new(parking_lot::Mutex::new(())),
869 shared: None,
870 table_name: None,
871 auth: None,
872 }
873 }
874}
875
876fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
880 use crate::query::Condition;
881 match c {
882 Condition::Pk(_)
884 | Condition::BitmapEq { .. }
885 | Condition::BitmapIn { .. }
886 | Condition::BytesPrefix { .. }
887 | Condition::IsNull { .. }
888 | Condition::IsNotNull { .. } => 0,
889 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
891 1
892 }
893 Condition::FmContains { .. }
895 | Condition::FmContainsAll { .. }
896 | Condition::Ann { .. }
897 | Condition::SparseMatch { .. } => 2,
898 }
899}
900
901impl Table {
902 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
903 let dir = dir.as_ref().to_path_buf();
904 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
905 Self::create_in(&dir, schema, table_id, ctx)
906 }
907
908 #[cfg(feature = "encryption")]
919 pub fn create_encrypted(
920 dir: impl AsRef<Path>,
921 schema: Schema,
922 table_id: u64,
923 passphrase: &str,
924 ) -> Result<Self> {
925 let dir = dir.as_ref();
926 std::fs::create_dir_all(dir.join(META_DIR))?;
927 let salt = crate::encryption::random_salt();
928 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
929 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
930 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
931 Self::create_in(dir, schema, table_id, ctx)
932 }
933
934 #[cfg(feature = "encryption")]
939 pub fn create_with_key(
940 dir: impl AsRef<Path>,
941 schema: Schema,
942 table_id: u64,
943 key: &[u8],
944 ) -> Result<Self> {
945 let dir = dir.as_ref();
946 std::fs::create_dir_all(dir.join(META_DIR))?;
947 let salt = crate::encryption::random_salt();
948 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
949 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
950 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
951 Self::create_in(dir, schema, table_id, ctx)
952 }
953
954 #[cfg(feature = "encryption")]
956 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
957 let dir = dir.as_ref();
958 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
959 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
960 MongrelError::NotFound(format!(
961 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
962 salt_path
963 ))
964 })?;
965 if salt_bytes.len() != crate::encryption::SALT_LEN {
966 return Err(MongrelError::InvalidArgument(format!(
967 "salt file is {} bytes, expected {}",
968 salt_bytes.len(),
969 crate::encryption::SALT_LEN
970 )));
971 }
972 let mut salt = [0u8; crate::encryption::SALT_LEN];
973 salt.copy_from_slice(&salt_bytes);
974 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
975 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
976 Self::open_in(dir, ctx)
977 }
978
979 pub(crate) fn create_in(
980 dir: impl AsRef<Path>,
981 schema: Schema,
982 table_id: u64,
983 ctx: SharedCtx,
984 ) -> Result<Self> {
985 schema.validate_auto_increment()?;
986 let dir = dir.as_ref().to_path_buf();
987 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
988 write_schema(&dir, &schema)?;
989 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
990 let (wal, current_txn_id) = match ctx.shared.clone() {
993 Some(s) => (WalSink::Shared(s), 0),
994 None => {
995 std::fs::create_dir_all(dir.join(WAL_DIR))?;
996 let mut w = if let Some(ref dk) = wal_dek {
997 Wal::create_with_cipher(
998 dir.join(WAL_DIR).join("seg-000000.wal"),
999 Epoch(0),
1000 Some(make_cipher(dk)),
1001 0,
1002 )?
1003 } else {
1004 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
1005 };
1006 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1007 (WalSink::Private(w), 1)
1008 }
1009 };
1010 let mut manifest = Manifest::new(table_id, schema.schema_id);
1011 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1016 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1017 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1018 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1019 let auto_inc = resolve_auto_inc(&schema);
1020 let rcache_dir = dir.join(RCACHE_DIR);
1021 Ok(Self {
1022 dir,
1023 table_id,
1024 name: ctx.table_name.unwrap_or_default(),
1025 auth: ctx.auth,
1026 wal,
1027 memtable: Memtable::new(),
1028 mutable_run: MutableRun::new(),
1029 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1030 compaction_zstd_level: 3,
1031 allocator: RowIdAllocator::new(0),
1032 epoch: ctx.epoch,
1033 persisted_epoch: 0,
1034 schema,
1035 hot: HotIndex::new(),
1036 kek: ctx.kek,
1037 column_keys,
1038 run_refs: Vec::new(),
1039 retiring: Vec::new(),
1040 next_run_id: 1,
1041 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1042 current_txn_id,
1043 bitmap,
1044 ann,
1045 fm,
1046 sparse,
1047 minhash,
1048 learned_range: HashMap::new(),
1049 pk_by_row: HashMap::new(),
1050 pinned: BTreeMap::new(),
1051 live_count: 0,
1052 reservoir: crate::reservoir::Reservoir::default(),
1053 reservoir_complete: true,
1054 had_deletes: false,
1055 agg_cache: HashMap::new(),
1056 global_idx_epoch: 0,
1057 indexes_complete: true,
1058 index_build_policy: IndexBuildPolicy::default(),
1059 pk_by_row_complete: false,
1060 flushed_epoch: 0,
1061 page_cache: ctx.page_cache,
1062 decoded_cache: ctx.decoded_cache,
1063 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1064 snapshots: ctx.snapshots,
1065 commit_lock: ctx.commit_lock,
1066 result_cache: Arc::new(parking_lot::Mutex::new(
1067 ResultCache::new()
1068 .with_dir(rcache_dir)
1069 .with_cache_dek(cache_dek.clone()),
1070 )),
1071 pending_delete_rids: roaring::RoaringBitmap::new(),
1072 pending_put_cols: std::collections::HashSet::new(),
1073 pending_rows: Vec::new(),
1074 pending_rows_auto_inc: Vec::new(),
1075 pending_dels: Vec::new(),
1076 pending_truncate: None,
1077 wal_dek,
1078 auto_inc,
1079 })
1080 }
1081
1082 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1086 let dir = dir.as_ref();
1087 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1088 Self::open_in(dir, ctx)
1089 }
1090
1091 #[cfg(feature = "encryption")]
1094 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1095 let dir = dir.as_ref();
1096 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1097 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1098 MongrelError::NotFound(format!(
1099 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1100 salt_path
1101 ))
1102 })?;
1103 let salt_len = crate::encryption::SALT_LEN;
1104 if salt_bytes.len() != salt_len {
1105 return Err(MongrelError::InvalidArgument(format!(
1106 "encryption salt is {} bytes, expected {salt_len}",
1107 salt_bytes.len()
1108 )));
1109 }
1110 let mut salt = [0u8; 16];
1111 salt.copy_from_slice(&salt_bytes);
1112 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1113 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1114 let t = Self::open_in(dir, ctx)?;
1115 Ok(t)
1116 }
1117
1118 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1119 let dir = dir.as_ref().to_path_buf();
1120 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1121 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1122 let schema: Schema = read_schema(&dir)?;
1123 let replay_epoch = Epoch(manifest.current_epoch);
1124 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1125 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1129 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1130 None => {
1131 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1132 let replayed = match &active {
1134 Some(path) => {
1135 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1136 crate::wal::replay_with_cipher(path, cipher)?
1137 }
1138 None => Vec::new(),
1139 };
1140 let mut w = match &active {
1141 Some(path) => Wal::create_with_cipher(
1142 path,
1143 replay_epoch,
1144 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1145 0,
1146 )?,
1147 None => Wal::create_with_cipher(
1148 dir.join(WAL_DIR).join("seg-000000.wal"),
1149 replay_epoch,
1150 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1151 0,
1152 )?,
1153 };
1154 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1155 (WalSink::Private(w), replayed, 1)
1156 }
1157 };
1158
1159 let mut memtable = Memtable::new();
1160 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1161 let persisted_epoch = manifest.current_epoch;
1162 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1169 s.next = manifest.auto_inc_next;
1170 s.seeded = manifest.auto_inc_next > 0;
1171 s
1172 });
1173
1174 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1181 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1182 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1183 std::collections::BTreeMap::new();
1184 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1185 let mut saw_delete = false;
1186 for record in replayed {
1187 let txn_id = record.txn_id;
1188 match record.op {
1189 Op::Put { rows, .. } => {
1190 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1191 for row in &rows {
1192 allocator.advance_to(row.row_id);
1193 if let Some(ai) = auto_inc.as_mut() {
1194 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1195 if *n + 1 > ai.next {
1196 ai.next = *n + 1;
1197 }
1198 }
1199 }
1200 }
1201 staged_puts.entry(txn_id).or_default().extend(rows);
1202 }
1203 Op::Delete { row_ids, .. } => {
1204 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1205 }
1206 Op::TxnCommit { epoch, .. } => {
1207 let commit_epoch = Epoch(epoch);
1208 if let Some(puts) = staged_puts.remove(&txn_id) {
1209 for row in &puts {
1210 memtable.upsert(row.clone());
1211 }
1212 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1213 }
1214 if let Some(dels) = staged_deletes.remove(&txn_id) {
1215 saw_delete = true;
1216 for rid in dels {
1217 memtable.tombstone(rid, commit_epoch);
1218 replayed_deletes.push((rid, commit_epoch));
1219 }
1220 }
1221 }
1222 Op::TxnAbort => {
1223 staged_puts.remove(&txn_id);
1224 staged_deletes.remove(&txn_id);
1225 }
1226 Op::TruncateTable { .. }
1227 | Op::ExternalTableState { .. }
1228 | Op::Flush { .. }
1229 | Op::Ddl(_) => {}
1230 }
1231 }
1232
1233 let rcache_dir = dir.join(RCACHE_DIR);
1234 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1235 let mut db = Self {
1236 dir,
1237 table_id: manifest.table_id,
1238 name: ctx.table_name.unwrap_or_default(),
1239 auth: ctx.auth,
1240 wal,
1241 memtable,
1242 mutable_run: MutableRun::new(),
1243 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1244 compaction_zstd_level: 3,
1245 allocator,
1246 epoch: ctx.epoch,
1247 persisted_epoch,
1248 schema,
1249 hot: HotIndex::new(),
1250 kek: ctx.kek,
1251 column_keys,
1252 run_refs: manifest.runs.clone(),
1253 retiring: manifest.retiring.clone(),
1254 next_run_id: manifest
1255 .runs
1256 .iter()
1257 .map(|r| r.run_id as u64 + 1)
1258 .max()
1259 .unwrap_or(1),
1260 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1261 current_txn_id,
1262 bitmap: HashMap::new(),
1263 ann: HashMap::new(),
1264 fm: HashMap::new(),
1265 sparse: HashMap::new(),
1266 minhash: HashMap::new(),
1267 learned_range: HashMap::new(),
1268 pk_by_row: HashMap::new(),
1269 pinned: BTreeMap::new(),
1270 live_count: manifest.live_count,
1271 reservoir: crate::reservoir::Reservoir::default(),
1272 reservoir_complete: false,
1273 had_deletes: saw_delete,
1274 agg_cache: HashMap::new(),
1275 global_idx_epoch: manifest.global_idx_epoch,
1276 indexes_complete: true,
1277 index_build_policy: IndexBuildPolicy::default(),
1278 pk_by_row_complete: false,
1279 flushed_epoch: manifest.flushed_epoch,
1280 page_cache: ctx.page_cache,
1281 decoded_cache: ctx.decoded_cache,
1282 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1283 snapshots: ctx.snapshots,
1284 commit_lock: ctx.commit_lock,
1285 result_cache: Arc::new(parking_lot::Mutex::new(
1286 ResultCache::new()
1287 .with_dir(rcache_dir)
1288 .with_cache_dek(cache_dek.clone()),
1289 )),
1290 pending_delete_rids: roaring::RoaringBitmap::new(),
1291 pending_put_cols: std::collections::HashSet::new(),
1292 pending_rows: Vec::new(),
1293 pending_rows_auto_inc: Vec::new(),
1294 pending_dels: Vec::new(),
1295 pending_truncate: None,
1296 wal_dek,
1297 auto_inc,
1298 };
1299
1300 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1303
1304 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1309 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1310 c.epoch_built == manifest.global_idx_epoch
1311 && manifest.global_idx_epoch > 0
1312 && manifest
1313 .runs
1314 .iter()
1315 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1316 });
1317 if let Some(loaded) = checkpoint {
1318 if checkpoint_valid {
1319 db.hot = loaded.hot;
1320 db.bitmap = loaded.bitmap;
1321 db.ann = loaded.ann;
1322 db.fm = loaded.fm;
1323 db.sparse = loaded.sparse;
1324 db.minhash = loaded.minhash;
1325 db.learned_range = loaded.learned_range;
1326 }
1329 }
1330 if !checkpoint_valid {
1331 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1332 db.bitmap = bitmap;
1333 db.ann = ann;
1334 db.fm = fm;
1335 db.sparse = sparse;
1336 db.minhash = minhash;
1337 db.rebuild_indexes_from_runs()?;
1338 db.build_learned_ranges()?;
1339 }
1340
1341 for (epoch, group) in replayed_puts {
1346 let (losers, winner_pks) = db.partition_pk_winners(&group);
1347 for (key, &row_id) in &winner_pks {
1348 if let Some(old_rid) = db.hot.get(key) {
1349 if old_rid != row_id {
1350 db.tombstone_row(old_rid, epoch, false);
1351 }
1352 }
1353 }
1354 for &loser_rid in &losers {
1355 db.tombstone_row(loser_rid, epoch, false);
1356 }
1357 for (key, row_id) in winner_pks {
1358 db.insert_hot_pk(key, row_id);
1359 }
1360 if db.schema.primary_key().is_none() {
1361 for r in &group {
1362 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1363 }
1364 }
1365 for r in &group {
1366 if !losers.contains(&r.row_id) {
1367 db.index_row(r);
1368 }
1369 }
1370 }
1371 for (rid, epoch) in &replayed_deletes {
1375 db.remove_hot_for_row(*rid, *epoch);
1376 }
1377
1378 db.result_cache.lock().load_persistent();
1385 Ok(db)
1386 }
1387
1388 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1394 if self.reservoir_complete {
1395 return Ok(());
1396 }
1397 self.rebuild_reservoir()?;
1398 self.reservoir_complete = true;
1399 Ok(())
1400 }
1401
1402 fn rebuild_reservoir(&mut self) -> Result<()> {
1405 let snap = self.snapshot();
1406 let rows = self.visible_rows(snap)?;
1407 self.reservoir.reset();
1408 for r in rows {
1409 self.reservoir.offer(r.row_id.0);
1410 }
1411 Ok(())
1412 }
1413
1414 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1415 self.hot = HotIndex::new();
1416 self.pk_by_row.clear();
1417 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1418 self.bitmap = bitmap;
1419 self.ann = ann;
1420 self.fm = fm;
1421 self.sparse = sparse;
1422 self.minhash = minhash;
1423 let snapshot = Epoch(u64::MAX);
1424 for rr in self.run_refs.clone() {
1425 let mut reader = self.open_reader(rr.run_id)?;
1426 for row in reader.visible_rows(snapshot)? {
1427 let tok_row = self.tokenized_for_indexes(&row);
1428 index_into(
1429 &self.schema,
1430 &tok_row,
1431 &mut self.hot,
1432 &mut self.bitmap,
1433 &mut self.ann,
1434 &mut self.fm,
1435 &mut self.sparse,
1436 &mut self.minhash,
1437 );
1438 }
1439 }
1440 for row in self.mutable_run.visible_versions(snapshot) {
1441 if row.deleted {
1442 self.remove_hot_for_row(row.row_id, snapshot);
1443 } else {
1444 self.index_row(&row);
1445 }
1446 }
1447 for row in self.memtable.visible_versions(snapshot) {
1448 if row.deleted {
1449 self.remove_hot_for_row(row.row_id, snapshot);
1450 } else {
1451 self.index_row(&row);
1452 }
1453 }
1454 self.refresh_pk_by_row_from_hot();
1455 Ok(())
1456 }
1457
1458 fn refresh_pk_by_row_from_hot(&mut self) {
1459 self.pk_by_row_complete = true;
1460 if self.schema.primary_key().is_none() {
1461 self.pk_by_row.clear();
1462 return;
1463 }
1464 self.pk_by_row = self
1470 .hot
1471 .entries()
1472 .into_iter()
1473 .map(|(key, row_id)| (row_id, key))
1474 .collect();
1475 }
1476
1477 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1478 if self.schema.primary_key().is_some() {
1479 self.pk_by_row.insert(row_id, key.clone());
1480 }
1481 self.hot.insert(key, row_id);
1482 }
1483
1484 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1488 self.learned_range.clear();
1489 if self.run_refs.len() != 1 {
1490 return Ok(());
1491 }
1492 let cols: Vec<u16> = self
1493 .schema
1494 .indexes
1495 .iter()
1496 .filter(|i| i.kind == IndexKind::LearnedRange)
1497 .map(|i| i.column_id)
1498 .collect();
1499 if cols.is_empty() {
1500 return Ok(());
1501 }
1502 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1503 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1504 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1505 _ => return Ok(()),
1506 };
1507 for cid in cols {
1508 let ty = self
1509 .schema
1510 .columns
1511 .iter()
1512 .find(|c| c.id == cid)
1513 .map(|c| c.ty)
1514 .unwrap_or(TypeId::Int64);
1515 match ty {
1516 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1517 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1518 let pairs: Vec<(i64, u64)> = data
1519 .iter()
1520 .zip(row_ids.iter())
1521 .map(|(v, r)| (*v, *r))
1522 .collect();
1523 self.learned_range
1524 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1525 }
1526 }
1527 TypeId::Float64 => {
1528 if let columnar::NativeColumn::Float64 { data, .. } =
1529 reader.column_native(cid)?
1530 {
1531 let pairs: Vec<(f64, u64)> = data
1532 .iter()
1533 .zip(row_ids.iter())
1534 .map(|(v, r)| (*v, *r))
1535 .collect();
1536 self.learned_range
1537 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1538 }
1539 }
1540 _ => {}
1541 }
1542 }
1543 Ok(())
1544 }
1545
1546 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1553 if self.indexes_complete {
1554 crate::trace::QueryTrace::record(|t| {
1555 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1556 });
1557 return Ok(());
1558 }
1559 crate::trace::QueryTrace::record(|t| {
1560 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1561 });
1562 self.rebuild_indexes_from_runs()?;
1563 self.build_learned_ranges()?;
1564 self.indexes_complete = true;
1565 let epoch = self.current_epoch();
1566 self.checkpoint_indexes(epoch);
1567 Ok(())
1568 }
1569
1570 fn pending_epoch(&self) -> Epoch {
1571 Epoch(self.epoch.visible().0 + 1)
1572 }
1573
1574 fn is_shared(&self) -> bool {
1577 matches!(self.wal, WalSink::Shared(_))
1578 }
1579
1580 fn ensure_txn_id(&mut self) -> u64 {
1584 if self.current_txn_id == 0 {
1585 let id = match &self.wal {
1586 WalSink::Shared(s) => {
1587 let mut g = s.txn_ids.lock();
1588 let v = *g;
1589 *g = g.wrapping_add(1);
1590 v
1591 }
1592 WalSink::Private(_) => 1,
1593 };
1594 self.current_txn_id = id;
1595 }
1596 self.current_txn_id
1597 }
1598
1599 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1602 let txn_id = self.ensure_txn_id();
1603 let table_id = self.table_id;
1604 match &mut self.wal {
1605 WalSink::Private(w) => {
1606 w.append_txn(txn_id, op)?;
1607 }
1608 WalSink::Shared(s) => {
1609 s.wal.lock().append(txn_id, table_id, op)?;
1610 }
1611 }
1612 Ok(())
1613 }
1614
1615 fn require(&self, perm: crate::auth_state::RequiredPermission) -> Result<()> {
1626 match &self.auth {
1627 Some(checker) => checker.check(&self.name, perm),
1628 None => Ok(()),
1629 }
1630 }
1631 pub fn require_select(&self) -> Result<()> {
1636 self.require(crate::auth_state::RequiredPermission::Select)
1637 }
1638 fn require_insert(&self) -> Result<()> {
1639 self.require(crate::auth_state::RequiredPermission::Insert)
1640 }
1641 #[allow(dead_code)]
1645 fn require_update(&self) -> Result<()> {
1646 self.require(crate::auth_state::RequiredPermission::Update)
1647 }
1648 fn require_delete(&self) -> Result<()> {
1649 self.require(crate::auth_state::RequiredPermission::Delete)
1650 }
1651
1652 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1655 self.require_insert()?;
1656 Ok(self.put_returning(columns)?.0)
1657 }
1658
1659 pub fn put_returning(
1664 &mut self,
1665 mut columns: Vec<(u16, Value)>,
1666 ) -> Result<(RowId, Option<i64>)> {
1667 self.require_insert()?;
1668 let assigned = self.fill_auto_inc(&mut columns)?;
1669 self.schema.validate_not_null(&columns)?;
1670 let row_id = if self.schema.clustered {
1675 self.derive_clustered_row_id(&columns)?
1676 } else {
1677 self.allocator.alloc()
1678 };
1679 let epoch = self.pending_epoch();
1680 let mut row = Row::new(row_id, epoch);
1681 for (col_id, val) in columns {
1682 row.columns.insert(col_id, val);
1683 }
1684 self.commit_rows(vec![row], assigned.is_some())?;
1685 Ok((row_id, assigned))
1686 }
1687
1688 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1691 self.require_insert()?;
1692 Ok(self
1693 .put_batch_returning(batch)?
1694 .into_iter()
1695 .map(|(r, _)| r)
1696 .collect())
1697 }
1698
1699 pub fn put_batch_returning(
1702 &mut self,
1703 batch: Vec<Vec<(u16, Value)>>,
1704 ) -> Result<Vec<(RowId, Option<i64>)>> {
1705 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1706 for mut cols in batch {
1707 let assigned = self.fill_auto_inc(&mut cols)?;
1708 filled.push((cols, assigned));
1709 }
1710 for (cols, _) in &filled {
1711 self.schema.validate_not_null(cols)?;
1712 }
1713 let epoch = self.pending_epoch();
1714 let mut rows = Vec::with_capacity(filled.len());
1715 let mut ids = Vec::with_capacity(filled.len());
1716 for (cols, assigned) in filled {
1717 let row_id = if self.schema.clustered {
1718 self.derive_clustered_row_id(&cols)?
1719 } else {
1720 self.allocator.alloc()
1721 };
1722 let mut row = Row::new(row_id, epoch);
1723 for (c, v) in cols {
1724 row.columns.insert(c, v);
1725 }
1726 ids.push((row_id, assigned));
1727 rows.push(row);
1728 }
1729 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1730 self.commit_rows(rows, all_auto_generated)?;
1731 Ok(ids)
1732 }
1733
1734 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1740 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1741 return Ok(None);
1742 };
1743 let pos = columns.iter().position(|(c, _)| *c == cid);
1744 let assigned = match pos {
1745 Some(i) => match &columns[i].1 {
1746 Value::Null => {
1747 let next = self.alloc_auto_inc_value()?;
1748 columns[i].1 = Value::Int64(next);
1749 Some(next)
1750 }
1751 Value::Int64(n) => {
1752 self.advance_auto_inc_past(*n)?;
1753 None
1754 }
1755 other => {
1756 return Err(MongrelError::InvalidArgument(format!(
1757 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1758 other
1759 )))
1760 }
1761 },
1762 None => {
1763 let next = self.alloc_auto_inc_value()?;
1764 columns.push((cid, Value::Int64(next)));
1765 Some(next)
1766 }
1767 };
1768 Ok(assigned)
1769 }
1770
1771 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1773 self.ensure_auto_inc_seeded()?;
1774 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1776 let v = ai.next;
1777 ai.next = ai.next.saturating_add(1);
1778 Ok(v)
1779 }
1780
1781 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1784 self.ensure_auto_inc_seeded()?;
1785 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1786 let floor = used.saturating_add(1).max(1);
1787 if ai.next < floor {
1788 ai.next = floor;
1789 }
1790 Ok(())
1791 }
1792
1793 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1798 let needs_seed = match self.auto_inc {
1799 Some(ai) => !ai.seeded,
1800 None => return Ok(()),
1801 };
1802 if !needs_seed {
1803 return Ok(());
1804 }
1805 if self.seed_empty_auto_inc() {
1806 return Ok(());
1807 }
1808 let cid = self
1809 .auto_inc
1810 .as_ref()
1811 .expect("auto-inc column present")
1812 .column_id;
1813 let max = self.scan_max_int64(cid)?;
1814 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1815 let floor = max.saturating_add(1).max(1);
1816 if ai.next < floor {
1817 ai.next = floor;
1818 }
1819 ai.seeded = true;
1820 Ok(())
1821 }
1822
1823 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1824 if n == 0 || self.auto_inc.is_none() {
1825 return Ok(None);
1826 }
1827 self.ensure_auto_inc_seeded()?;
1828 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1829 let start = ai.next;
1830 ai.next = ai.next.saturating_add(n as i64);
1831 Ok(Some(start))
1832 }
1833
1834 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1838 let mut max: i64 = 0;
1839 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1840 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1841 if *n > max {
1842 max = *n;
1843 }
1844 }
1845 }
1846 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1847 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1848 if *n > max {
1849 max = *n;
1850 }
1851 }
1852 }
1853 for rr in self.run_refs.clone() {
1854 let reader = self.open_reader(rr.run_id)?;
1855 if let Some(stats) = reader.column_page_stats(column_id) {
1856 for s in stats {
1857 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1858 if n > max {
1859 max = n;
1860 }
1861 }
1862 }
1863 } else if reader.has_column(column_id) {
1864 if let columnar::NativeColumn::Int64 { data, validity } =
1865 reader.column_native_shared(column_id)?
1866 {
1867 for (i, n) in data.iter().enumerate() {
1868 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1869 {
1870 max = *n;
1871 }
1872 }
1873 }
1874 }
1875 }
1876 Ok(max)
1877 }
1878
1879 fn seed_empty_auto_inc(&mut self) -> bool {
1880 let Some(ai) = self.auto_inc.as_mut() else {
1881 return false;
1882 };
1883 if ai.seeded || self.live_count != 0 {
1884 return false;
1885 }
1886 if ai.next < 1 {
1887 ai.next = 1;
1888 }
1889 ai.seeded = true;
1890 true
1891 }
1892
1893 fn advance_auto_inc_from_native_columns(
1894 &mut self,
1895 columns: &[(u16, columnar::NativeColumn)],
1896 n: usize,
1897 live_before: u64,
1898 ) -> Result<()> {
1899 let Some(ai) = self.auto_inc.as_mut() else {
1900 return Ok(());
1901 };
1902 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1903 return Ok(());
1904 };
1905 let columnar::NativeColumn::Int64 { data, validity } = col else {
1906 return Err(MongrelError::InvalidArgument(format!(
1907 "AUTO_INCREMENT column {} must be Int64",
1908 ai.column_id
1909 )));
1910 };
1911 let max = if native_int64_strictly_increasing(col, n) {
1912 data.get(n.saturating_sub(1)).copied()
1913 } else {
1914 data.iter()
1915 .take(n)
1916 .enumerate()
1917 .filter_map(|(i, v)| {
1918 if validity.is_empty() || columnar::validity_bit(validity, i) {
1919 Some(*v)
1920 } else {
1921 None
1922 }
1923 })
1924 .max()
1925 };
1926 if let Some(max) = max {
1927 let floor = max.saturating_add(1).max(1);
1928 if ai.next < floor {
1929 ai.next = floor;
1930 }
1931 if ai.seeded || live_before == 0 {
1932 ai.seeded = true;
1933 }
1934 }
1935 Ok(())
1936 }
1937
1938 fn fill_auto_inc_native_columns(
1939 &mut self,
1940 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1941 n: usize,
1942 ) -> Result<()> {
1943 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1944 return Ok(());
1945 };
1946 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1947 if let Some(start) = self.alloc_auto_inc_range(n)? {
1948 columns.push((
1949 cid,
1950 columnar::NativeColumn::Int64 {
1951 data: (start..start.saturating_add(n as i64)).collect(),
1952 validity: vec![0xFF; n.div_ceil(8)],
1953 },
1954 ));
1955 }
1956 return Ok(());
1957 };
1958
1959 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1960 return Err(MongrelError::InvalidArgument(format!(
1961 "AUTO_INCREMENT column {cid} must be Int64"
1962 )));
1963 };
1964 if data.len() < n {
1965 return Err(MongrelError::InvalidArgument(format!(
1966 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1967 data.len()
1968 )));
1969 }
1970 if columnar::all_non_null(validity, n) {
1971 return Ok(());
1972 }
1973 if validity.iter().all(|b| *b == 0) {
1974 if let Some(start) = self.alloc_auto_inc_range(n)? {
1975 for (i, slot) in data.iter_mut().take(n).enumerate() {
1976 *slot = start.saturating_add(i as i64);
1977 }
1978 *validity = vec![0xFF; n.div_ceil(8)];
1979 }
1980 return Ok(());
1981 }
1982
1983 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1984 for (i, slot) in data.iter_mut().enumerate().take(n) {
1985 if columnar::validity_bit(validity, i) {
1986 self.advance_auto_inc_past(*slot)?;
1987 } else {
1988 *slot = self.alloc_auto_inc_value()?;
1989 }
1990 }
1991 *validity = new_validity;
1992 Ok(())
1993 }
1994
1995 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
2009 if self.auto_inc.is_none() {
2010 return Ok(None);
2011 }
2012 Ok(Some(self.alloc_auto_inc_value()?))
2013 }
2014
2015 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
2021 let payload = bincode::serialize(&rows)?;
2022 self.wal_append_data(Op::Put {
2023 table_id: self.table_id,
2024 rows: payload,
2025 })?;
2026 if self.is_shared() {
2027 self.pending_rows_auto_inc
2028 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
2029 self.pending_rows.extend(rows);
2030 } else {
2031 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
2032 }
2033 Ok(())
2034 }
2035
2036 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
2041 self.apply_put_rows_inner(rows, true)
2042 }
2043
2044 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
2045 if check_existing_pk {
2046 self.ensure_indexes_complete()?;
2047 }
2048 if rows.len() == 1 {
2052 let row = rows.into_iter().next().expect("len checked");
2053 return self.apply_put_row_single(row, check_existing_pk);
2054 }
2055 let pk_id = self.schema.primary_key().map(|c| c.id);
2072 let probe = match pk_id {
2073 Some(pid) => {
2074 check_existing_pk
2075 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2076 }
2077 None => false,
2078 };
2079 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2082 for r in rows {
2083 for &cid in r.columns.keys() {
2084 self.pending_put_cols.insert(cid);
2085 }
2086 match pk_id {
2087 Some(pid) if probe || maintain_pk_by_row => {
2088 if let Some(pk_val) = r.columns.get(&pid) {
2089 let key = self.index_lookup_key(pid, pk_val);
2090 if probe {
2091 if let Some(old_rid) = self.hot.get(&key) {
2092 if old_rid != r.row_id {
2093 self.tombstone_row(old_rid, r.committed_epoch, true);
2094 }
2095 }
2096 }
2097 if maintain_pk_by_row {
2098 self.pk_by_row.insert(r.row_id, key);
2099 }
2100 }
2101 }
2102 Some(_) => {}
2103 None => {
2104 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2105 }
2106 }
2107 self.index_row(&r);
2108 self.reservoir.offer(r.row_id.0);
2109 self.memtable.upsert(r);
2110 self.live_count = self.live_count.saturating_add(1);
2113 }
2114 Ok(())
2115 }
2116
2117 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2121 for &cid in row.columns.keys() {
2122 self.pending_put_cols.insert(cid);
2123 }
2124 let epoch = row.committed_epoch;
2125 if let Some(pk_col) = self.schema.primary_key() {
2126 let pk_id = pk_col.id;
2127 if let Some(pk_val) = row.columns.get(&pk_id) {
2128 let maintain_pk_by_row = self.pk_by_row_complete;
2132 if check_existing_pk || maintain_pk_by_row {
2133 let key = self.index_lookup_key(pk_id, pk_val);
2134 if check_existing_pk {
2135 if let Some(old_rid) = self.hot.get(&key) {
2136 if old_rid != row.row_id {
2137 self.tombstone_row(old_rid, epoch, true);
2138 }
2139 }
2140 }
2141 if maintain_pk_by_row {
2142 self.pk_by_row.insert(row.row_id, key);
2143 }
2144 }
2145 }
2146 } else {
2147 self.hot
2148 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2149 }
2150 self.index_row(&row);
2151 self.reservoir.offer(row.row_id.0);
2152 self.memtable.upsert(row);
2153 self.live_count = self.live_count.saturating_add(1);
2154 Ok(())
2155 }
2156
2157 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2160 self.allocator.alloc()
2161 }
2162
2163 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2169 let pk = self.schema.primary_key().ok_or_else(|| {
2170 MongrelError::Schema("clustered table requires a single-column primary key".into())
2171 })?;
2172 let pk_val = columns
2173 .iter()
2174 .find(|(id, _)| *id == pk.id)
2175 .map(|(_, v)| v)
2176 .ok_or_else(|| {
2177 MongrelError::Schema(format!(
2178 "clustered table missing primary key column {} ({})",
2179 pk.id, pk.name
2180 ))
2181 })?;
2182 let key_bytes = pk_val.encode_key();
2183 let mut hash: u64 = 0xcbf29ce484222325;
2185 for &b in &key_bytes {
2186 hash ^= b as u64;
2187 hash = hash.wrapping_mul(0x100000001b3);
2188 }
2189 Ok(RowId(hash.max(1)))
2192 }
2193
2194 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2202 self.ensure_indexes_complete()?;
2203 let n = rows.len();
2204 for r in rows {
2205 for &cid in r.columns.keys() {
2206 self.pending_put_cols.insert(cid);
2207 }
2208 }
2209 let (losers, winner_pks) = self.partition_pk_winners(rows);
2210 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2211 for (key, &row_id) in &winner_pks {
2213 if let Some(old_rid) = self.hot.get(key) {
2214 if old_rid != row_id {
2215 self.tombstone_row(old_rid, epoch, true);
2216 }
2217 }
2218 }
2219 for &loser_rid in &losers {
2222 self.tombstone_row(loser_rid, epoch, false);
2223 }
2224 for (key, row_id) in winner_pks {
2226 self.insert_hot_pk(key, row_id);
2227 }
2228 if self.schema.primary_key().is_none() {
2229 for r in rows {
2230 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2231 }
2232 }
2233 for r in rows {
2234 self.allocator.advance_to(r.row_id);
2235 if !losers.contains(&r.row_id) {
2236 self.index_row(r);
2237 }
2238 }
2239 for r in rows {
2240 if !losers.contains(&r.row_id) {
2241 self.reservoir.offer(r.row_id.0);
2242 }
2243 }
2244 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2245 Ok(())
2246 }
2247
2248 pub(crate) fn recover_apply(
2253 &mut self,
2254 rows: Vec<Row>,
2255 deletes: Vec<(RowId, Epoch)>,
2256 ) -> Result<()> {
2257 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2261 std::collections::BTreeMap::new();
2262 for row in rows {
2263 self.allocator.advance_to(row.row_id);
2264 if let Some(ai) = self.auto_inc.as_mut() {
2269 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2270 if *n + 1 > ai.next {
2271 ai.next = *n + 1;
2272 }
2273 }
2274 }
2275 by_epoch.entry(row.committed_epoch).or_default().push(row);
2276 }
2277 for (epoch, group) in by_epoch {
2278 let (losers, winner_pks) = self.partition_pk_winners(&group);
2279 for (key, &row_id) in &winner_pks {
2281 if let Some(old_rid) = self.hot.get(key) {
2282 if old_rid != row_id {
2283 self.tombstone_row(old_rid, epoch, false);
2284 }
2285 }
2286 }
2287 for (key, row_id) in winner_pks {
2288 self.insert_hot_pk(key, row_id);
2289 }
2290 if self.schema.primary_key().is_none() {
2291 for r in &group {
2292 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2293 }
2294 }
2295 for r in &group {
2296 if !losers.contains(&r.row_id) {
2297 self.memtable.upsert(r.clone());
2298 self.index_row(r);
2299 }
2300 }
2301 }
2302 for (rid, epoch) in deletes {
2303 self.memtable.tombstone(rid, epoch);
2304 self.remove_hot_for_row(rid, epoch);
2305 }
2306 self.reservoir_complete = false;
2309 Ok(())
2310 }
2311
2312 pub(crate) fn flushed_epoch(&self) -> u64 {
2314 self.flushed_epoch
2315 }
2316
2317 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2319 self.schema.validate_not_null(cells)
2320 }
2321
2322 fn validate_columns_not_null(
2326 &self,
2327 columns: &[(u16, columnar::NativeColumn)],
2328 n: usize,
2329 ) -> Result<()> {
2330 let by_id: HashMap<u16, &columnar::NativeColumn> =
2331 columns.iter().map(|(id, c)| (*id, c)).collect();
2332 for col in &self.schema.columns {
2333 if col.flags.contains(ColumnFlags::NULLABLE) {
2334 continue;
2335 }
2336 match by_id.get(&col.id) {
2337 None => {
2338 return Err(MongrelError::InvalidArgument(format!(
2339 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2340 col.name, col.id
2341 )));
2342 }
2343 Some(c) => {
2344 if c.null_count(n) != 0 {
2345 return Err(MongrelError::InvalidArgument(format!(
2346 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2347 col.name, col.id
2348 )));
2349 }
2350 }
2351 }
2352 }
2353 Ok(())
2354 }
2355
2356 fn bulk_pk_winner_indices(
2361 &self,
2362 columns: &[(u16, columnar::NativeColumn)],
2363 n: usize,
2364 ) -> Option<Vec<usize>> {
2365 let pk_col = self.schema.primary_key()?;
2366 let pk_id = pk_col.id;
2367 let pk_ty = pk_col.ty;
2368 let by_id: HashMap<u16, &columnar::NativeColumn> =
2369 columns.iter().map(|(id, c)| (*id, c)).collect();
2370 let pk_native = by_id.get(&pk_id)?;
2371 if native_int64_strictly_increasing(pk_native, n) {
2372 return None;
2373 }
2374 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2376 let mut null_pk_rows: Vec<usize> = Vec::new();
2377 for i in 0..n {
2378 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2379 Some(key) => {
2380 last.insert(key, i);
2381 }
2382 None => null_pk_rows.push(i),
2383 }
2384 }
2385 let mut winners: HashSet<usize> = last.values().copied().collect();
2386 for i in null_pk_rows {
2387 winners.insert(i);
2388 }
2389 Some((0..n).filter(|i| winners.contains(i)).collect())
2390 }
2391
2392 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2394 self.require_delete()?;
2395 let epoch = self.pending_epoch();
2396 self.wal_append_data(Op::Delete {
2397 table_id: self.table_id,
2398 row_ids: vec![row_id],
2399 })?;
2400 if self.is_shared() {
2401 self.pending_dels.push(row_id);
2402 } else {
2403 self.apply_delete(row_id, epoch);
2404 }
2405 Ok(())
2406 }
2407
2408 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2409 let pre = self.get(row_id, self.snapshot());
2410 self.delete(row_id)?;
2411 Ok(pre.map(|row| {
2412 let mut columns: Vec<_> = row.columns.into_iter().collect();
2413 columns.sort_by_key(|(id, _)| *id);
2414 OwnedRow { columns }
2415 }))
2416 }
2417
2418 pub fn truncate(&mut self) -> Result<()> {
2420 self.require_delete()?;
2421 let epoch = self.pending_epoch();
2422 self.wal_append_data(Op::TruncateTable {
2423 table_id: self.table_id,
2424 })?;
2425 self.pending_rows.clear();
2426 self.pending_rows_auto_inc.clear();
2427 self.pending_dels.clear();
2428 self.pending_truncate = Some(epoch);
2429 Ok(())
2430 }
2431
2432 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2434 for rr in std::mem::take(&mut self.run_refs) {
2435 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2436 }
2437 for r in std::mem::take(&mut self.retiring) {
2438 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2439 }
2440 self.memtable = Memtable::new();
2441 self.mutable_run = MutableRun::new();
2442 self.hot = HotIndex::new();
2443 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2444 self.bitmap = bitmap;
2445 self.ann = ann;
2446 self.fm = fm;
2447 self.sparse = sparse;
2448 self.minhash = minhash;
2449 self.learned_range.clear();
2450 self.pk_by_row.clear();
2451 self.pk_by_row_complete = false;
2452 self.live_count = 0;
2453 self.reservoir = crate::reservoir::Reservoir::default();
2454 self.reservoir_complete = true;
2455 self.had_deletes = true;
2456 self.agg_cache.clear();
2457 self.global_idx_epoch = 0;
2458 self.indexes_complete = true;
2459 self.pending_delete_rids.clear();
2460 self.pending_put_cols.clear();
2461 self.pending_rows.clear();
2462 self.pending_rows_auto_inc.clear();
2463 self.pending_dels.clear();
2464 self.clear_result_cache();
2465 self.invalidate_index_checkpoint();
2466 Ok(())
2467 }
2468
2469 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2472 self.remove_hot_for_row(row_id, epoch);
2473 self.tombstone_row(row_id, epoch, true);
2474 }
2475
2476 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2480 let tombstone = Row {
2481 row_id,
2482 committed_epoch: epoch,
2483 columns: std::collections::HashMap::new(),
2484 deleted: true,
2485 };
2486 self.memtable.upsert(tombstone);
2487 self.pk_by_row.remove(&row_id);
2488 if adjust_live_count {
2489 self.live_count = self.live_count.saturating_sub(1);
2490 }
2491 self.pending_delete_rids.insert(row_id.0 as u32);
2493 self.had_deletes = true;
2496 self.agg_cache.clear();
2497 }
2498
2499 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2503 let Some(pk_col) = self.schema.primary_key() else {
2504 return;
2505 };
2506 if self.pk_by_row_complete {
2509 if let Some(key) = self.pk_by_row.remove(&row_id) {
2510 if self.hot.get(&key) == Some(row_id) {
2511 self.hot.remove(&key);
2512 }
2513 }
2514 return;
2515 }
2516 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2535 if self.indexes_complete {
2536 let pk_val = self
2537 .memtable
2538 .get_version(row_id, lookup_epoch)
2539 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2540 .or_else(|| {
2541 self.mutable_run
2542 .get_version(row_id, lookup_epoch)
2543 .filter(|(_, r)| !r.deleted)
2544 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2545 })
2546 .or_else(|| {
2547 self.run_refs.iter().find_map(|rr| {
2548 let mut reader = self.open_reader(rr.run_id).ok()?;
2549 let (_, deleted, val) = reader
2550 .get_version_column(row_id, lookup_epoch, pk_col.id)
2551 .ok()??;
2552 if deleted {
2553 return None;
2554 }
2555 val
2556 })
2557 });
2558 if let Some(pk_val) = pk_val {
2559 let key = self.index_lookup_key(pk_col.id, &pk_val);
2560 if self.hot.get(&key) == Some(row_id) {
2561 self.hot.remove(&key);
2562 }
2563 return;
2564 }
2565 }
2566 self.refresh_pk_by_row_from_hot();
2571 if let Some(key) = self.pk_by_row.remove(&row_id) {
2572 if self.hot.get(&key) == Some(row_id) {
2573 self.hot.remove(&key);
2574 }
2575 }
2576 }
2577
2578 fn partition_pk_winners(
2583 &self,
2584 rows: &[Row],
2585 ) -> (
2586 std::collections::HashSet<RowId>,
2587 std::collections::HashMap<Vec<u8>, RowId>,
2588 ) {
2589 let mut losers = std::collections::HashSet::new();
2590 let Some(pk_col) = self.schema.primary_key() else {
2591 return (losers, std::collections::HashMap::new());
2592 };
2593 let pk_id = pk_col.id;
2594 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2595 std::collections::HashMap::new();
2596 for r in rows {
2597 let Some(pk_val) = r.columns.get(&pk_id) else {
2598 continue;
2599 };
2600 let key = self.index_lookup_key(pk_id, pk_val);
2601 if let Some(&old_rid) = winners.get(&key) {
2602 losers.insert(old_rid);
2603 }
2604 winners.insert(key, r.row_id);
2605 }
2606 (losers, winners)
2607 }
2608
2609 fn index_row(&mut self, row: &Row) {
2610 if row.deleted {
2611 return;
2612 }
2613 let any_predicate = self
2621 .schema
2622 .indexes
2623 .iter()
2624 .any(|idx| idx.predicate.is_some());
2625 if any_predicate {
2626 let columns_map: HashMap<u16, &Value> =
2627 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2628 let name_to_id: HashMap<&str, u16> = self
2629 .schema
2630 .columns
2631 .iter()
2632 .map(|c| (c.name.as_str(), c.id))
2633 .collect();
2634 for idx in &self.schema.indexes {
2635 if let Some(pred) = &idx.predicate {
2636 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2637 continue; }
2639 }
2640 index_into_single(
2642 idx,
2643 &self.schema,
2644 row,
2645 &mut self.hot,
2646 &mut self.bitmap,
2647 &mut self.ann,
2648 &mut self.fm,
2649 &mut self.sparse,
2650 &mut self.minhash,
2651 );
2652 }
2653 return;
2654 }
2655 if self.column_keys.is_empty() {
2659 index_into(
2660 &self.schema,
2661 row,
2662 &mut self.hot,
2663 &mut self.bitmap,
2664 &mut self.ann,
2665 &mut self.fm,
2666 &mut self.sparse,
2667 &mut self.minhash,
2668 );
2669 return;
2670 }
2671 let effective_row = self.tokenized_for_indexes(row);
2672 index_into(
2673 &self.schema,
2674 &effective_row,
2675 &mut self.hot,
2676 &mut self.bitmap,
2677 &mut self.ann,
2678 &mut self.fm,
2679 &mut self.sparse,
2680 &mut self.minhash,
2681 );
2682 }
2683
2684 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2690 if self.column_keys.is_empty() {
2691 return row.clone();
2692 }
2693 #[cfg(feature = "encryption")]
2694 {
2695 use crate::encryption::SCHEME_HMAC_EQ;
2696 let mut tok = row.clone();
2697 for (&cid, &(_, scheme)) in &self.column_keys {
2698 if scheme != SCHEME_HMAC_EQ {
2699 continue;
2700 }
2701 if let Some(v) = tok.columns.get(&cid).cloned() {
2702 if let Some(t) = self.tokenize_value(cid, &v) {
2703 tok.columns.insert(cid, t);
2704 }
2705 }
2706 }
2707 tok
2708 }
2709 #[cfg(not(feature = "encryption"))]
2710 {
2711 row.clone()
2712 }
2713 }
2714
2715 pub fn commit(&mut self) -> Result<Epoch> {
2720 if self.is_shared() {
2721 self.commit_shared()
2722 } else {
2723 self.commit_private()
2724 }
2725 }
2726
2727 fn commit_private(&mut self) -> Result<Epoch> {
2729 let commit_lock = Arc::clone(&self.commit_lock);
2733 let _g = commit_lock.lock();
2734 let new_epoch = self.epoch.bump_assigned();
2735 let txn_id = self.current_txn_id;
2736 match &mut self.wal {
2740 WalSink::Private(w) => {
2741 w.append_txn(
2742 txn_id,
2743 Op::TxnCommit {
2744 epoch: new_epoch.0,
2745 added_runs: Vec::new(),
2746 },
2747 )?;
2748 w.sync()?;
2749 }
2750 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2751 }
2752 if let Some(epoch) = self.pending_truncate.take() {
2754 self.apply_truncate(epoch)?;
2755 }
2756 self.invalidate_pending_cache();
2757 self.persist_manifest(new_epoch)?;
2758 self.epoch.publish_in_order(new_epoch);
2762 self.current_txn_id += 1;
2763 Ok(new_epoch)
2764 }
2765
2766 fn commit_shared(&mut self) -> Result<Epoch> {
2772 use std::sync::atomic::Ordering;
2773 let s = match &self.wal {
2774 WalSink::Shared(s) => s.clone(),
2775 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2776 };
2777 if s.poisoned.load(Ordering::Relaxed) {
2778 return Err(MongrelError::Other(
2779 "database poisoned by fsync error".into(),
2780 ));
2781 }
2782 let commit_lock = Arc::clone(&self.commit_lock);
2789 let _g = commit_lock.lock();
2790 let txn_id = self.ensure_txn_id();
2793 let (new_epoch, commit_seq) = {
2794 let mut wal = s.wal.lock();
2795 let new_epoch = self.epoch.bump_assigned();
2796 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2797 (new_epoch, seq)
2798 };
2799 s.group
2800 .await_durable(&s.wal, commit_seq)
2801 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2802
2803 if self.pending_truncate.take().is_some() {
2806 self.apply_truncate(new_epoch)?;
2807 }
2808 let mut rows = std::mem::take(&mut self.pending_rows);
2809 if !rows.is_empty() {
2810 for r in &mut rows {
2811 r.committed_epoch = new_epoch;
2812 }
2813 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2814 let all_auto_generated =
2815 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2816 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2817 } else {
2818 self.pending_rows_auto_inc.clear();
2819 }
2820 let dels = std::mem::take(&mut self.pending_dels);
2821 for rid in dels {
2822 self.apply_delete(rid, new_epoch);
2823 }
2824
2825 self.invalidate_pending_cache();
2826 self.persist_manifest(new_epoch)?;
2827 self.epoch.publish_in_order(new_epoch);
2828 self.current_txn_id = 0;
2830 Ok(new_epoch)
2831 }
2832
2833 pub fn flush(&mut self) -> Result<Epoch> {
2841 self.ensure_indexes_complete()?;
2842 let epoch = self.commit()?;
2843 let rows = self.memtable.drain_sorted();
2844 if !rows.is_empty() {
2845 self.mutable_run.insert_many(rows);
2846 }
2847 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2848 self.spill_mutable_run(epoch)?;
2849 self.mark_flushed(epoch)?;
2853 self.persist_manifest(epoch)?;
2854 self.build_learned_ranges()?;
2855 self.checkpoint_indexes(epoch);
2858 }
2859 Ok(epoch)
2862 }
2863
2864 pub fn force_flush(&mut self) -> Result<Epoch> {
2873 let saved = self.mutable_run_spill_bytes;
2874 self.mutable_run_spill_bytes = 1;
2875 let result = self.flush();
2876 self.mutable_run_spill_bytes = saved;
2877 result
2878 }
2879
2880 pub fn close(&mut self) -> Result<()> {
2887 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
2888 self.force_flush()?;
2889 }
2890 Ok(())
2891 }
2892
2893 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2900 let op = Op::Flush {
2901 table_id: self.table_id,
2902 flushed_epoch: epoch.0,
2903 };
2904 match &mut self.wal {
2905 WalSink::Private(w) => {
2906 w.append_system(op)?;
2907 w.sync()?;
2908 }
2909 WalSink::Shared(s) => {
2910 s.wal.lock().append_system(op)?;
2915 }
2916 }
2917 self.flushed_epoch = epoch.0;
2918 if matches!(self.wal, WalSink::Private(_)) {
2919 self.rotate_wal(epoch)?;
2920 }
2921 Ok(())
2922 }
2923
2924 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2928 let rows = self.mutable_run.drain_sorted();
2929 if rows.is_empty() {
2930 return Ok(());
2931 }
2932 let run_id = self.next_run_id;
2933 self.next_run_id += 1;
2934 let path = self.run_path(run_id);
2935 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2936 if let Some(kek) = &self.kek {
2937 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2938 }
2939 let header = writer.write(&path, &rows)?;
2940 self.run_refs.push(RunRef {
2941 run_id: run_id as u128,
2942 level: 0,
2943 epoch_created: epoch.0,
2944 row_count: header.row_count,
2945 });
2946 Ok(())
2947 }
2948
2949 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2953 self.mutable_run_spill_bytes = bytes.max(1);
2954 }
2955
2956 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2960 self.compaction_zstd_level = level;
2961 }
2962
2963 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2967 self.result_cache.lock().set_max_bytes(max_bytes);
2968 }
2969
2970 pub(crate) fn clear_result_cache(&mut self) {
2974 self.result_cache.lock().clear();
2975 }
2976
2977 pub fn mutable_run_len(&self) -> usize {
2979 self.mutable_run.len()
2980 }
2981
2982 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2985 self.mutable_run.drain_sorted()
2986 }
2987
2988 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2993 let epoch = self.commit()?;
2994 let n = batch.len();
2995 if n == 0 {
2996 return Ok(epoch);
2997 }
2998 let live_before = self.live_count;
2999 self.spill_mutable_run(epoch)?;
3003 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3004 && self.indexes_complete
3005 && self.run_refs.is_empty()
3006 && self.memtable.is_empty()
3007 && self.mutable_run.is_empty();
3008 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
3014 use rayon::prelude::*;
3015 self.schema
3016 .columns
3017 .par_iter()
3018 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
3019 .collect::<Vec<_>>()
3020 };
3021 drop(batch);
3022 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3027 self.validate_columns_not_null(&user_columns, n)?;
3028 let winner_idx = self
3029 .bulk_pk_winner_indices(&user_columns, n)
3030 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
3031 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3032 match winner_idx.as_deref() {
3033 Some(idx) => {
3034 let compacted = user_columns
3035 .iter()
3036 .map(|(id, c)| (*id, c.gather(idx)))
3037 .collect();
3038 (compacted, idx.len())
3039 }
3040 None => (std::mem::take(&mut user_columns), n),
3041 };
3042 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3043 let first = self.allocator.alloc_range(write_n as u64).0;
3044 for rid in first..first + write_n as u64 {
3045 self.reservoir.offer(rid);
3046 }
3047 let run_id = self.next_run_id;
3048 self.next_run_id += 1;
3049 let path = self.run_path(run_id);
3050 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
3051 .clean(true)
3052 .with_lz4()
3053 .with_native_endian();
3054 if let Some(kek) = &self.kek {
3055 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3056 }
3057 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3058 self.run_refs.push(RunRef {
3059 run_id: run_id as u128,
3060 level: 0,
3061 epoch_created: epoch.0,
3062 row_count: header.row_count,
3063 });
3064 self.live_count = self.live_count.saturating_add(write_n as u64);
3065 if eager_index_build {
3066 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3067 self.index_columns_bulk(&write_columns, &row_ids);
3068 self.indexes_complete = true;
3069 self.build_learned_ranges()?;
3070 } else {
3071 self.indexes_complete = false;
3072 }
3073 self.mark_flushed(epoch)?;
3074 self.persist_manifest(epoch)?;
3075 if eager_index_build {
3076 self.checkpoint_indexes(epoch);
3077 }
3078 self.clear_result_cache();
3079 Ok(epoch)
3080 }
3081
3082 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3085 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3086 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3087 let segment_no = segment
3090 .file_stem()
3091 .and_then(|s| s.to_str())
3092 .and_then(|s| s.strip_prefix("seg-"))
3093 .and_then(|s| s.parse::<u64>().ok())
3094 .unwrap_or(0);
3095 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3096 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3097 wal.sync()?;
3098 self.wal = WalSink::Private(wal);
3099 Ok(())
3100 }
3101
3102 pub(crate) fn invalidate_pending_cache(&mut self) {
3107 self.result_cache
3108 .lock()
3109 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3110 self.pending_delete_rids.clear();
3111 self.pending_put_cols.clear();
3112 }
3113
3114 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3115 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3116 m.current_epoch = epoch.0;
3117 m.next_row_id = self.allocator.current().0;
3118 m.runs = self.run_refs.clone();
3119 m.live_count = self.live_count;
3120 m.global_idx_epoch = self.global_idx_epoch;
3121 m.flushed_epoch = self.flushed_epoch;
3122 m.retiring = self.retiring.clone();
3123 m.auto_inc_next = match self.auto_inc {
3127 Some(ai) if ai.seeded => ai.next,
3128 _ => 0,
3129 };
3130 let meta_dek = self.manifest_meta_dek();
3131 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3132 Ok(())
3133 }
3134
3135 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3141 if !self.indexes_complete {
3144 return;
3145 }
3146 let snap = global_idx::IndexSnapshot {
3147 hot: &self.hot,
3148 bitmap: &self.bitmap,
3149 ann: &self.ann,
3150 fm: &self.fm,
3151 sparse: &self.sparse,
3152 minhash: &self.minhash,
3153 learned_range: &self.learned_range,
3154 };
3155 let idx_dek = self.idx_dek();
3157 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3158 .is_ok()
3159 {
3160 self.global_idx_epoch = epoch.0;
3161 let _ = self.persist_manifest(epoch);
3162 }
3163 }
3164
3165 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3168 self.global_idx_epoch = 0;
3169 global_idx::remove(&self.dir);
3170 let _ = self.persist_manifest(self.epoch.visible());
3171 }
3172
3173 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3176 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3177 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3178 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3179 best = Some((epoch, row));
3180 }
3181 }
3182 for rr in &self.run_refs {
3183 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3184 continue;
3185 };
3186 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3187 continue;
3188 };
3189 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3190 best = Some((epoch, row));
3191 }
3192 }
3193 match best {
3194 Some((_, r)) if r.deleted => None,
3195 Some((_, r)) => Some(r),
3196 None => None,
3197 }
3198 }
3199
3200 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3204 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3205 let mut fold = |row: Row| {
3206 best.entry(row.row_id.0)
3207 .and_modify(|e| {
3208 if row.committed_epoch > e.0 {
3209 *e = (row.committed_epoch, row.clone());
3210 }
3211 })
3212 .or_insert_with(|| (row.committed_epoch, row));
3213 };
3214 for row in self.memtable.visible_versions(snapshot.epoch) {
3215 fold(row);
3216 }
3217 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3218 fold(row);
3219 }
3220 for rr in &self.run_refs {
3221 let mut reader = self.open_reader(rr.run_id)?;
3222 for row in reader.visible_versions(snapshot.epoch)? {
3223 fold(row);
3224 }
3225 }
3226 let mut out: Vec<Row> = best
3227 .into_values()
3228 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3229 .collect();
3230 out.sort_by_key(|r| r.row_id);
3231 Ok(out)
3232 }
3233
3234 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3241 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3242 let rr = self.run_refs[0].clone();
3243 let mut reader = self.open_reader(rr.run_id)?;
3244 let idxs = reader.visible_indices(snapshot.epoch)?;
3245 let mut cols = Vec::with_capacity(self.schema.columns.len());
3246 for cdef in &self.schema.columns {
3247 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3248 }
3249 return Ok(cols);
3250 }
3251 let rows = self.visible_rows(snapshot)?;
3253 let mut cols: Vec<(u16, Vec<Value>)> = self
3254 .schema
3255 .columns
3256 .iter()
3257 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3258 .collect();
3259 for r in &rows {
3260 for (cid, vec) in cols.iter_mut() {
3261 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3262 }
3263 }
3264 Ok(cols)
3265 }
3266
3267 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3269 self.hot.get(key)
3270 }
3271
3272 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3277 self.require_select()?;
3278 self.ensure_indexes_complete()?;
3279 let snapshot = self.snapshot();
3280 crate::trace::QueryTrace::record(|t| {
3281 t.run_count = self.run_refs.len();
3282 t.memtable_rows = self.memtable.len();
3283 t.mutable_run_rows = self.mutable_run.len();
3284 });
3285 if q.conditions.is_empty() {
3289 crate::trace::QueryTrace::record(|t| {
3290 t.scan_mode = crate::trace::ScanMode::Materialized;
3291 t.row_materialized = true;
3292 });
3293 return self.visible_rows(snapshot);
3294 }
3295 crate::trace::QueryTrace::record(|t| {
3296 t.conditions_pushed = q.conditions.len();
3297 t.scan_mode = crate::trace::ScanMode::Materialized;
3298 t.row_materialized = true;
3299 });
3300 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3307 ordered.sort_by_key(|c| condition_cost_rank(c));
3308 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3309 for c in &ordered {
3310 let s = self.resolve_condition(c, snapshot)?;
3311 let empty = s.is_empty();
3312 sets.push(s);
3313 if empty {
3314 break;
3315 }
3316 }
3317 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3318 self.rows_for_rids(&rids, snapshot)
3319 }
3320
3321 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3326 use std::collections::HashMap;
3327 let mut rows = Vec::with_capacity(rids.len());
3328 let tier_size = self.memtable.len() + self.mutable_run.len();
3345 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3346 if rids.len().saturating_mul(24) < tier_size {
3347 for &rid in rids {
3348 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3349 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3350 let newest = match (mem, mrun) {
3351 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3352 (Some((_, mr)), None) => Some(mr),
3353 (None, Some((_, rr))) => Some(rr),
3354 (None, None) => None,
3355 };
3356 if let Some(row) = newest {
3357 overlay.insert(rid, row);
3358 }
3359 }
3360 } else {
3361 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3362 overlay
3363 .entry(row.row_id.0)
3364 .and_modify(|e| {
3365 if row.committed_epoch > e.committed_epoch {
3366 *e = row.clone();
3367 }
3368 })
3369 .or_insert(row);
3370 };
3371 for row in self.memtable.visible_versions(snapshot.epoch) {
3372 fold_newest(row, &mut overlay);
3373 }
3374 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3375 fold_newest(row, &mut overlay);
3376 }
3377 }
3378 if self.run_refs.len() == 1 {
3379 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3380 if rids.len().saturating_mul(24) < reader.row_count() {
3388 for &rid in rids {
3389 if let Some(r) = overlay.get(&rid) {
3390 if !r.deleted {
3391 rows.push(r.clone());
3392 }
3393 continue;
3394 }
3395 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3396 if !row.deleted {
3397 rows.push(row);
3398 }
3399 }
3400 }
3401 return Ok(rows);
3402 }
3403 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3412 enum Src {
3415 Overlay,
3416 Run,
3417 }
3418 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3419 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3420 for rid in rids {
3421 if overlay.contains_key(rid) {
3422 plan.push(Src::Overlay);
3423 continue;
3424 }
3425 match vis_rids.binary_search(&(*rid as i64)) {
3426 Ok(i) => {
3427 plan.push(Src::Run);
3428 fetch.push(positions[i]);
3429 }
3430 Err(_) => { }
3431 }
3432 }
3433 let fetched = reader.materialize_batch(&fetch)?;
3434 let mut fetched_iter = fetched.into_iter();
3435 for (rid, src) in rids.iter().zip(plan) {
3436 match src {
3437 Src::Overlay => {
3438 if let Some(r) = overlay.get(rid) {
3439 if !r.deleted {
3440 rows.push(r.clone());
3441 }
3442 }
3443 }
3444 Src::Run => {
3445 if let Some(row) = fetched_iter.next() {
3446 if !row.deleted {
3447 rows.push(row);
3448 }
3449 }
3450 }
3451 }
3452 }
3453 return Ok(rows);
3454 }
3455 let mut readers: Vec<_> = self
3459 .run_refs
3460 .iter()
3461 .map(|rr| self.open_reader(rr.run_id))
3462 .collect::<Result<Vec<_>>>()?;
3463 for rid in rids {
3464 if let Some(r) = overlay.get(rid) {
3465 if !r.deleted {
3466 rows.push(r.clone());
3467 }
3468 continue;
3469 }
3470 let mut best: Option<(Epoch, Row)> = None;
3471 for reader in readers.iter_mut() {
3472 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3473 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3474 best = Some((epoch, row));
3475 }
3476 }
3477 }
3478 if let Some((_, r)) = best {
3479 if !r.deleted {
3480 rows.push(r);
3481 }
3482 }
3483 }
3484 Ok(rows)
3485 }
3486
3487 pub fn indexes_complete(&self) -> bool {
3497 self.indexes_complete
3498 }
3499
3500 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3502 self.index_build_policy
3503 }
3504
3505 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3509 self.index_build_policy = policy;
3510 }
3511
3512 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3517 if !self.indexes_complete {
3521 return None;
3522 }
3523 let b = self.bitmap.get(&column_id)?;
3524 let result: Vec<Vec<u8>> = b
3525 .keys()
3526 .into_iter()
3527 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3528 .cloned()
3529 .collect();
3530 Some(result)
3531 }
3532
3533 pub fn fk_join_row_ids(
3534 &self,
3535 fk_column_id: u16,
3536 pk_values: &[Vec<u8>],
3537 fk_conditions: &[crate::query::Condition],
3538 snapshot: Snapshot,
3539 ) -> Result<Vec<u64>> {
3540 let Some(b) = self.bitmap.get(&fk_column_id) else {
3541 return Ok(Vec::new());
3542 };
3543 let mut join_set = {
3544 let mut acc = roaring::RoaringBitmap::new();
3545 for v in pk_values {
3546 acc |= b.get(v);
3547 }
3548 RowIdSet::from_roaring(acc)
3549 };
3550 if !fk_conditions.is_empty() {
3551 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3552 sets.push(join_set);
3553 for c in fk_conditions {
3554 sets.push(self.resolve_condition(c, snapshot)?);
3555 }
3556 join_set = RowIdSet::intersect_many(sets);
3557 }
3558 Ok(join_set.into_sorted_vec())
3559 }
3560
3561 pub fn fk_join_count(
3567 &self,
3568 fk_column_id: u16,
3569 pk_values: &[Vec<u8>],
3570 fk_conditions: &[crate::query::Condition],
3571 snapshot: Snapshot,
3572 ) -> Result<u64> {
3573 let Some(b) = self.bitmap.get(&fk_column_id) else {
3574 return Ok(0);
3575 };
3576 let mut acc = roaring::RoaringBitmap::new();
3577 for v in pk_values {
3578 acc |= b.get(v);
3579 }
3580 if fk_conditions.is_empty() {
3581 return Ok(acc.len());
3582 }
3583 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3584 sets.push(RowIdSet::from_roaring(acc));
3585 for c in fk_conditions {
3586 sets.push(self.resolve_condition(c, snapshot)?);
3587 }
3588 Ok(RowIdSet::intersect_many(sets).len() as u64)
3589 }
3590
3591 fn resolve_condition(
3596 &self,
3597 c: &crate::query::Condition,
3598 snapshot: Snapshot,
3599 ) -> Result<RowIdSet> {
3600 use crate::query::Condition;
3601 Ok(match c {
3602 Condition::Pk(key) => {
3603 let lookup = self
3604 .schema
3605 .primary_key()
3606 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3607 .unwrap_or_else(|| key.clone());
3608 self.hot
3609 .get(&lookup)
3610 .map(|r| RowIdSet::one(r.0))
3611 .unwrap_or_else(RowIdSet::empty)
3612 }
3613 Condition::BitmapEq { column_id, value } => {
3614 let lookup = self.index_lookup_key_bytes(*column_id, value);
3615 self.bitmap
3616 .get(column_id)
3617 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3618 .unwrap_or_else(RowIdSet::empty)
3619 }
3620 Condition::BitmapIn { column_id, values } => {
3621 let bm = self.bitmap.get(column_id);
3622 let mut acc = roaring::RoaringBitmap::new();
3623 if let Some(b) = bm {
3624 for v in values {
3625 let lookup = self.index_lookup_key_bytes(*column_id, v);
3626 acc |= b.get(&lookup);
3627 }
3628 }
3629 RowIdSet::from_roaring(acc)
3630 }
3631 Condition::BytesPrefix { column_id, prefix } => {
3632 if let Some(b) = self.bitmap.get(column_id) {
3637 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3638 let mut acc = roaring::RoaringBitmap::new();
3639 for key in b.keys() {
3640 if key.starts_with(&lookup_prefix) {
3641 acc |= b.get(key);
3642 }
3643 }
3644 RowIdSet::from_roaring(acc)
3645 } else {
3646 RowIdSet::empty()
3647 }
3648 }
3649 Condition::FmContains { column_id, pattern } => self
3650 .fm
3651 .get(column_id)
3652 .map(|f| {
3653 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3654 })
3655 .unwrap_or_else(RowIdSet::empty),
3656 Condition::FmContainsAll {
3657 column_id,
3658 patterns,
3659 } => {
3660 if let Some(f) = self.fm.get(column_id) {
3663 let sets: Vec<RowIdSet> = patterns
3664 .iter()
3665 .map(|pat| {
3666 RowIdSet::from_unsorted(
3667 f.locate(pat).into_iter().map(|r| r.0).collect(),
3668 )
3669 })
3670 .collect();
3671 RowIdSet::intersect_many(sets)
3672 } else {
3673 RowIdSet::empty()
3674 }
3675 }
3676 Condition::Ann {
3677 column_id,
3678 query,
3679 k,
3680 } => self
3681 .ann
3682 .get(column_id)
3683 .map(|a| {
3684 RowIdSet::from_unsorted(
3685 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3686 )
3687 })
3688 .unwrap_or_else(RowIdSet::empty),
3689 Condition::SparseMatch {
3690 column_id,
3691 query,
3692 k,
3693 } => self
3694 .sparse
3695 .get(column_id)
3696 .map(|s| {
3697 RowIdSet::from_unsorted(
3698 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3699 )
3700 })
3701 .unwrap_or_else(RowIdSet::empty),
3702 Condition::MinHashSimilar {
3703 column_id,
3704 query,
3705 k,
3706 } => self
3707 .minhash
3708 .get(column_id)
3709 .map(|mh| {
3710 RowIdSet::from_unsorted(
3711 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3712 )
3713 })
3714 .unwrap_or_else(RowIdSet::empty),
3715 Condition::Range { column_id, lo, hi } => {
3716 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3725 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3726 } else if self.run_refs.len() == 1 {
3727 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3728 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3729 } else {
3730 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3731 };
3732 set.remove_many(self.overlay_rid_set(snapshot));
3733 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3734 set
3735 }
3736 Condition::RangeF64 {
3737 column_id,
3738 lo,
3739 lo_inclusive,
3740 hi,
3741 hi_inclusive,
3742 } => {
3743 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3746 RowIdSet::from_unsorted(
3747 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3748 .into_iter()
3749 .collect(),
3750 )
3751 } else if self.run_refs.len() == 1 {
3752 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3753 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3754 } else {
3755 return self.range_scan_f64(
3756 *column_id,
3757 *lo,
3758 *lo_inclusive,
3759 *hi,
3760 *hi_inclusive,
3761 snapshot,
3762 );
3763 };
3764 set.remove_many(self.overlay_rid_set(snapshot));
3765 self.range_scan_overlay_f64(
3766 &mut set,
3767 *column_id,
3768 *lo,
3769 *lo_inclusive,
3770 *hi,
3771 *hi_inclusive,
3772 snapshot,
3773 );
3774 set
3775 }
3776 Condition::IsNull { column_id } => {
3777 let mut set = if self.run_refs.len() == 1 {
3778 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3779 r.null_row_id_set(*column_id, true)?
3780 } else {
3781 return self.null_scan(*column_id, true, snapshot);
3782 };
3783 set.remove_many(self.overlay_rid_set(snapshot));
3784 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3785 set
3786 }
3787 Condition::IsNotNull { column_id } => {
3788 let mut set = if self.run_refs.len() == 1 {
3789 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3790 r.null_row_id_set(*column_id, false)?
3791 } else {
3792 return self.null_scan(*column_id, false, snapshot);
3793 };
3794 set.remove_many(self.overlay_rid_set(snapshot));
3795 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3796 set
3797 }
3798 })
3799 }
3800
3801 fn range_scan_i64(
3809 &self,
3810 column_id: u16,
3811 lo: i64,
3812 hi: i64,
3813 snapshot: Snapshot,
3814 ) -> Result<RowIdSet> {
3815 let mut row_ids = Vec::new();
3816 let overlay_rids = self.overlay_rid_set(snapshot);
3817 for rr in &self.run_refs {
3818 let mut reader = self.open_reader(rr.run_id)?;
3819 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3820 for rid in matched {
3821 if !overlay_rids.contains(&rid) {
3822 row_ids.push(rid);
3823 }
3824 }
3825 }
3826 let mut s = RowIdSet::from_unsorted(row_ids);
3827 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3828 Ok(s)
3829 }
3830
3831 fn range_scan_f64(
3834 &self,
3835 column_id: u16,
3836 lo: f64,
3837 lo_inclusive: bool,
3838 hi: f64,
3839 hi_inclusive: bool,
3840 snapshot: Snapshot,
3841 ) -> Result<RowIdSet> {
3842 let mut row_ids = Vec::new();
3843 let overlay_rids = self.overlay_rid_set(snapshot);
3844 for rr in &self.run_refs {
3845 let mut reader = self.open_reader(rr.run_id)?;
3846 let matched = reader.range_row_ids_visible_f64(
3847 column_id,
3848 lo,
3849 lo_inclusive,
3850 hi,
3851 hi_inclusive,
3852 snapshot.epoch,
3853 )?;
3854 for rid in matched {
3855 if !overlay_rids.contains(&rid) {
3856 row_ids.push(rid);
3857 }
3858 }
3859 }
3860 let mut s = RowIdSet::from_unsorted(row_ids);
3861 self.range_scan_overlay_f64(
3862 &mut s,
3863 column_id,
3864 lo,
3865 lo_inclusive,
3866 hi,
3867 hi_inclusive,
3868 snapshot,
3869 );
3870 Ok(s)
3871 }
3872
3873 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3875 let mut s = HashSet::new();
3876 for row in self.memtable.visible_versions(snapshot.epoch) {
3877 s.insert(row.row_id.0);
3878 }
3879 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3880 s.insert(row.row_id.0);
3881 }
3882 s
3883 }
3884
3885 fn range_scan_overlay_i64(
3886 &self,
3887 s: &mut RowIdSet,
3888 column_id: u16,
3889 lo: i64,
3890 hi: i64,
3891 snapshot: Snapshot,
3892 ) {
3893 let mut newest: HashMap<u64, &Row> = HashMap::new();
3898 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3899 let memtable = self.memtable.visible_versions(snapshot.epoch);
3900 for r in &mutable {
3901 newest.entry(r.row_id.0).or_insert(r);
3902 }
3903 for r in &memtable {
3904 newest.insert(r.row_id.0, r);
3905 }
3906 for row in newest.values() {
3907 if !row.deleted {
3908 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3909 if *v >= lo && *v <= hi {
3910 s.insert(row.row_id.0);
3911 }
3912 }
3913 }
3914 }
3915 }
3916
3917 #[allow(clippy::too_many_arguments)]
3918 fn range_scan_overlay_f64(
3919 &self,
3920 s: &mut RowIdSet,
3921 column_id: u16,
3922 lo: f64,
3923 lo_inclusive: bool,
3924 hi: f64,
3925 hi_inclusive: bool,
3926 snapshot: Snapshot,
3927 ) {
3928 let mut newest: HashMap<u64, &Row> = HashMap::new();
3931 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3932 let memtable = self.memtable.visible_versions(snapshot.epoch);
3933 for r in &mutable {
3934 newest.entry(r.row_id.0).or_insert(r);
3935 }
3936 for r in &memtable {
3937 newest.insert(r.row_id.0, r);
3938 }
3939 for row in newest.values() {
3940 if !row.deleted {
3941 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3942 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3943 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3944 if ok_lo && ok_hi {
3945 s.insert(row.row_id.0);
3946 }
3947 }
3948 }
3949 }
3950 }
3951
3952 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3955 let mut row_ids = Vec::new();
3956 let overlay_rids = self.overlay_rid_set(snapshot);
3957 for rr in &self.run_refs {
3958 let mut reader = self.open_reader(rr.run_id)?;
3959 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3960 for rid in matched {
3961 if !overlay_rids.contains(&rid) {
3962 row_ids.push(rid);
3963 }
3964 }
3965 }
3966 let mut s = RowIdSet::from_unsorted(row_ids);
3967 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3968 Ok(s)
3969 }
3970
3971 fn null_scan_overlay(
3975 &self,
3976 s: &mut RowIdSet,
3977 column_id: u16,
3978 want_nulls: bool,
3979 snapshot: Snapshot,
3980 ) {
3981 let mut newest: HashMap<u64, &Row> = HashMap::new();
3982 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3983 let memtable = self.memtable.visible_versions(snapshot.epoch);
3984 for r in &mutable {
3985 newest.entry(r.row_id.0).or_insert(r);
3986 }
3987 for r in &memtable {
3988 newest.insert(r.row_id.0, r);
3989 }
3990 for row in newest.values() {
3991 if row.deleted {
3992 continue;
3993 }
3994 let is_null = !row.columns.contains_key(&column_id)
3995 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3996 if is_null == want_nulls {
3997 s.insert(row.row_id.0);
3998 }
3999 }
4000 }
4001
4002 pub fn snapshot(&self) -> Snapshot {
4003 Snapshot::at(self.epoch.visible())
4004 }
4005
4006 pub fn pin_snapshot(&mut self) -> Snapshot {
4009 let e = self.epoch.visible();
4010 *self.pinned.entry(e).or_insert(0) += 1;
4011 Snapshot::at(e)
4012 }
4013
4014 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
4016 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
4017 *count -= 1;
4018 if *count == 0 {
4019 self.pinned.remove(&snap.epoch);
4020 }
4021 }
4022 }
4023
4024 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
4034 let local = self.pinned.keys().next().copied();
4035 let global = self.snapshots.min_pinned();
4036 match (local, global) {
4037 (Some(a), Some(b)) => Some(a.min(b)),
4038 (Some(a), None) => Some(a),
4039 (None, b) => b,
4040 }
4041 }
4042
4043 pub fn current_epoch(&self) -> Epoch {
4044 self.epoch.visible()
4045 }
4046
4047 pub fn memtable_len(&self) -> usize {
4048 self.memtable.len()
4049 }
4050
4051 pub fn count(&self) -> u64 {
4054 self.live_count
4055 }
4056
4057 pub fn count_conditions(
4061 &mut self,
4062 conditions: &[crate::query::Condition],
4063 snapshot: Snapshot,
4064 ) -> Result<Option<u64>> {
4065 use crate::query::Condition;
4066 if conditions.is_empty() {
4067 return Ok(Some(self.live_count));
4068 }
4069 let served = |c: &Condition| {
4070 matches!(
4071 c,
4072 Condition::Pk(_)
4073 | Condition::BitmapEq { .. }
4074 | Condition::BitmapIn { .. }
4075 | Condition::BytesPrefix { .. }
4076 | Condition::FmContains { .. }
4077 | Condition::FmContainsAll { .. }
4078 | Condition::Ann { .. }
4079 | Condition::Range { .. }
4080 | Condition::RangeF64 { .. }
4081 | Condition::SparseMatch { .. }
4082 | Condition::MinHashSimilar { .. }
4083 | Condition::IsNull { .. }
4084 | Condition::IsNotNull { .. }
4085 )
4086 };
4087 if !conditions.iter().all(served) {
4088 return Ok(None);
4089 }
4090 self.ensure_indexes_complete()?;
4091 let mut sets = Vec::with_capacity(conditions.len());
4092 for condition in conditions {
4093 sets.push(self.resolve_condition(condition, snapshot)?);
4094 }
4095 let mut rids = RowIdSet::intersect_many(sets);
4096 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4106 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
4107 }
4108 let count = rids.len() as u64;
4109 crate::trace::QueryTrace::record(|t| {
4110 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
4111 t.survivor_count = Some(count as usize);
4112 t.conditions_pushed = conditions.len();
4113 });
4114 Ok(Some(count))
4115 }
4116
4117 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
4122 let mut out = Vec::new();
4123 for row in self.memtable.visible_versions(snapshot.epoch) {
4124 if row.deleted {
4125 out.push(row.row_id.0);
4126 }
4127 }
4128 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4129 if row.deleted {
4130 out.push(row.row_id.0);
4131 }
4132 }
4133 out
4134 }
4135
4136 pub fn bulk_load_columns(
4145 &mut self,
4146 user_columns: Vec<(u16, columnar::NativeColumn)>,
4147 ) -> Result<Epoch> {
4148 self.bulk_load_columns_with(user_columns, 3, false, true)
4149 }
4150
4151 pub fn bulk_load_fast(
4158 &mut self,
4159 user_columns: Vec<(u16, columnar::NativeColumn)>,
4160 ) -> Result<Epoch> {
4161 self.bulk_load_columns_with(user_columns, -1, true, false)
4162 }
4163
4164 fn bulk_load_columns_with(
4165 &mut self,
4166 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4167 zstd_level: i32,
4168 force_plain: bool,
4169 lz4: bool,
4170 ) -> Result<Epoch> {
4171 let epoch = self.commit()?;
4172 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4173 if n == 0 {
4174 return Ok(epoch);
4175 }
4176 let live_before = self.live_count;
4177 self.spill_mutable_run(epoch)?;
4179 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4180 && self.indexes_complete
4181 && self.run_refs.is_empty()
4182 && self.memtable.is_empty()
4183 && self.mutable_run.is_empty();
4184 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4187 self.validate_columns_not_null(&user_columns, n)?;
4188 let winner_idx = self
4189 .bulk_pk_winner_indices(&user_columns, n)
4190 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
4191 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4192 match winner_idx.as_deref() {
4193 Some(idx) => {
4194 let compacted = user_columns
4195 .iter()
4196 .map(|(id, c)| (*id, c.gather(idx)))
4197 .collect();
4198 (compacted, idx.len())
4199 }
4200 None => (user_columns, n),
4201 };
4202 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4203 let first = self.allocator.alloc_range(write_n as u64).0;
4204 for rid in first..first + write_n as u64 {
4205 self.reservoir.offer(rid);
4206 }
4207 let run_id = self.next_run_id;
4208 self.next_run_id += 1;
4209 let path = self.run_path(run_id);
4210 let mut writer =
4211 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4212 if force_plain {
4213 writer = writer.with_plain();
4214 } else if lz4 {
4215 writer = writer.with_lz4();
4218 } else {
4219 writer = writer.with_zstd_level(zstd_level);
4220 }
4221 if let Some(kek) = &self.kek {
4222 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4223 }
4224 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4225 self.run_refs.push(RunRef {
4226 run_id: run_id as u128,
4227 level: 0,
4228 epoch_created: epoch.0,
4229 row_count: header.row_count,
4230 });
4231 self.live_count = self.live_count.saturating_add(write_n as u64);
4232 if eager_index_build {
4233 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4234 self.index_columns_bulk(&write_columns, &row_ids);
4235 self.indexes_complete = true;
4236 self.build_learned_ranges()?;
4237 } else {
4238 self.indexes_complete = false;
4242 }
4243 self.mark_flushed(epoch)?;
4244 self.persist_manifest(epoch)?;
4245 if eager_index_build {
4246 self.checkpoint_indexes(epoch);
4247 }
4248 self.clear_result_cache();
4249 Ok(epoch)
4250 }
4251
4252 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4270 let n = row_ids.len();
4271 if n == 0 {
4272 return;
4273 }
4274 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4275 columns.iter().map(|(id, c)| (*id, c)).collect();
4276 let ty_of: std::collections::HashMap<u16, TypeId> =
4277 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
4278 let pk_id = self.schema.primary_key().map(|c| c.id);
4279
4280 for (i, &rid) in row_ids.iter().enumerate() {
4281 let row_id = RowId(rid);
4282 if let Some(pid) = pk_id {
4283 if let Some(col) = by_id.get(&pid) {
4284 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
4285 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4286 self.insert_hot_pk(key, row_id);
4287 }
4288 }
4289 }
4290 for idef in &self.schema.indexes {
4291 let Some(col) = by_id.get(&idef.column_id) else {
4292 continue;
4293 };
4294 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
4295 match idef.kind {
4296 IndexKind::Bitmap => {
4297 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4298 if let Some(key) =
4299 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4300 {
4301 b.insert(key, row_id);
4302 }
4303 }
4304 }
4305 IndexKind::FmIndex => {
4306 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4307 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4308 f.insert(bytes.to_vec(), row_id);
4309 }
4310 }
4311 }
4312 IndexKind::Sparse => {
4313 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4314 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4315 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4316 s.insert(&terms, row_id);
4317 }
4318 }
4319 }
4320 }
4321 IndexKind::MinHash => {
4322 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4323 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4324 let tokens = crate::index::token_hashes_from_bytes(bytes);
4325 mh.insert(&tokens, row_id);
4326 }
4327 }
4328 }
4329 _ => {}
4330 }
4331 }
4332 }
4333 }
4334
4335 pub fn visible_columns_native(
4340 &self,
4341 snapshot: Snapshot,
4342 projection: Option<&[u16]>,
4343 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4344 let wanted: Vec<u16> = match projection {
4345 Some(p) => p.to_vec(),
4346 None => self.schema.columns.iter().map(|c| c.id).collect(),
4347 };
4348 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4349 let rr = self.run_refs[0].clone();
4350 let mut reader = self.open_reader(rr.run_id)?;
4351 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4352 let all_visible = idxs.len() == reader.row_count();
4353 if reader.has_mmap() {
4359 use rayon::prelude::*;
4360 let valid: Vec<u16> = wanted
4363 .iter()
4364 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4365 .copied()
4366 .collect();
4367 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4369 .par_iter()
4370 .filter_map(|cid| {
4371 reader
4372 .column_native_shared(*cid)
4373 .ok()
4374 .map(|col| (*cid, col))
4375 })
4376 .collect();
4377 let cols = decoded
4378 .into_iter()
4379 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4380 .collect();
4381 return Ok(cols);
4382 }
4383 let mut cols = Vec::with_capacity(wanted.len());
4384 for cid in &wanted {
4385 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4386 Some(c) => c,
4387 None => continue,
4388 };
4389 let col = reader.column_native(cdef.id)?;
4390 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4391 }
4392 return Ok(cols);
4393 }
4394 let vcols = self.visible_columns(snapshot)?;
4395 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4396 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4397 .into_iter()
4398 .filter(|(id, _)| want_set.contains(id))
4399 .map(|(id, vals)| {
4400 let ty = self
4401 .schema
4402 .columns
4403 .iter()
4404 .find(|c| c.id == id)
4405 .map(|c| c.ty)
4406 .unwrap_or(TypeId::Bytes);
4407 (id, columnar::values_to_native(ty, &vals))
4408 })
4409 .collect();
4410 Ok(out)
4411 }
4412
4413 pub fn run_count(&self) -> usize {
4414 self.run_refs.len()
4415 }
4416
4417 pub fn memtable_is_empty(&self) -> bool {
4419 self.memtable.is_empty()
4420 }
4421
4422 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4426 self.page_cache.stats()
4427 }
4428
4429 pub fn reset_page_cache_stats(&self) {
4431 self.page_cache.reset_stats();
4432 }
4433
4434 pub fn run_ids(&self) -> Vec<u128> {
4437 self.run_refs.iter().map(|r| r.run_id).collect()
4438 }
4439
4440 pub fn single_run_is_clean(&self) -> bool {
4444 if self.run_refs.len() != 1 {
4445 return false;
4446 }
4447 self.open_reader(self.run_refs[0].run_id)
4448 .map(|r| r.is_clean())
4449 .unwrap_or(false)
4450 }
4451
4452 fn resolve_footprint(
4459 &self,
4460 conditions: &[crate::query::Condition],
4461 _snapshot: Snapshot,
4462 ) -> roaring::RoaringBitmap {
4463 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4464 return roaring::RoaringBitmap::new();
4465 }
4466 if self.run_refs.is_empty() {
4467 return roaring::RoaringBitmap::new();
4468 }
4469 if self.run_refs.len() == 1 {
4471 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4472 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4473 return rids.to_roaring_lossy();
4474 }
4475 }
4476 }
4477 roaring::RoaringBitmap::new()
4478 }
4479
4480 pub fn query_columns_native_cached(
4491 &mut self,
4492 conditions: &[crate::query::Condition],
4493 projection: Option<&[u16]>,
4494 snapshot: Snapshot,
4495 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4496 if conditions.is_empty() {
4497 return self.query_columns_native(conditions, projection, snapshot);
4498 }
4499 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4503 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4504 crate::trace::QueryTrace::record(|t| {
4505 t.result_cache_hit = true;
4506 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4507 });
4508 return Ok(Some((*hit).clone()));
4509 }
4510 let res = self.query_columns_native(conditions, projection, snapshot)?;
4511 if let Some(cols) = &res {
4512 let footprint = self.resolve_footprint(conditions, snapshot);
4513 let condition_cols = crate::query::condition_columns(conditions);
4514 self.result_cache.lock().insert(
4515 key,
4516 CachedEntry {
4517 data: CachedData::Columns(Arc::new(cols.clone())),
4518 footprint,
4519 condition_cols,
4520 },
4521 );
4522 }
4523 Ok(res)
4524 }
4525
4526 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4531 if q.conditions.is_empty() {
4532 return self.query(q);
4533 }
4534 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4535 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4536 crate::trace::QueryTrace::record(|t| {
4537 t.result_cache_hit = true;
4538 t.scan_mode = crate::trace::ScanMode::Materialized;
4539 });
4540 return Ok((*hit).clone());
4541 }
4542 let rows = self.query(q)?;
4543 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4544 let condition_cols = crate::query::condition_columns(&q.conditions);
4545 self.result_cache.lock().insert(
4546 key,
4547 CachedEntry {
4548 data: CachedData::Rows(Arc::new(rows.clone())),
4549 footprint,
4550 condition_cols,
4551 },
4552 );
4553 Ok(rows)
4554 }
4555
4556 #[allow(clippy::type_complexity)]
4571 pub fn query_columns_native_traced(
4572 &mut self,
4573 conditions: &[crate::query::Condition],
4574 projection: Option<&[u16]>,
4575 snapshot: Snapshot,
4576 ) -> Result<(
4577 Option<Vec<(u16, columnar::NativeColumn)>>,
4578 crate::trace::QueryTrace,
4579 )> {
4580 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4581 self.query_columns_native(conditions, projection, snapshot)
4582 });
4583 Ok((result?, trace))
4584 }
4585
4586 #[allow(clippy::type_complexity)]
4589 pub fn query_columns_native_cached_traced(
4590 &mut self,
4591 conditions: &[crate::query::Condition],
4592 projection: Option<&[u16]>,
4593 snapshot: Snapshot,
4594 ) -> Result<(
4595 Option<Vec<(u16, columnar::NativeColumn)>>,
4596 crate::trace::QueryTrace,
4597 )> {
4598 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4599 self.query_columns_native_cached(conditions, projection, snapshot)
4600 });
4601 Ok((result?, trace))
4602 }
4603
4604 pub fn native_page_cursor_traced(
4606 &self,
4607 snapshot: Snapshot,
4608 projection: Vec<(u16, TypeId)>,
4609 conditions: &[crate::query::Condition],
4610 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4611 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4612 self.native_page_cursor(snapshot, projection, conditions)
4613 });
4614 Ok((result?, trace))
4615 }
4616
4617 pub fn native_multi_run_cursor_traced(
4619 &self,
4620 snapshot: Snapshot,
4621 projection: Vec<(u16, TypeId)>,
4622 conditions: &[crate::query::Condition],
4623 ) -> Result<(
4624 Option<crate::cursor::MultiRunCursor>,
4625 crate::trace::QueryTrace,
4626 )> {
4627 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4628 self.native_multi_run_cursor(snapshot, projection, conditions)
4629 });
4630 Ok((result?, trace))
4631 }
4632
4633 pub fn count_conditions_traced(
4635 &mut self,
4636 conditions: &[crate::query::Condition],
4637 snapshot: Snapshot,
4638 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4639 let (result, trace) =
4640 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4641 Ok((result?, trace))
4642 }
4643
4644 pub fn query_traced(
4646 &mut self,
4647 q: &crate::query::Query,
4648 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4649 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4650 Ok((result?, trace))
4651 }
4652
4653 pub fn query_columns_native(
4658 &mut self,
4659 conditions: &[crate::query::Condition],
4660 projection: Option<&[u16]>,
4661 snapshot: Snapshot,
4662 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4663 use crate::query::Condition;
4664 if conditions.is_empty() {
4665 return Ok(None);
4666 }
4667 self.ensure_indexes_complete()?;
4668
4669 let served = |c: &Condition| {
4674 matches!(
4675 c,
4676 Condition::Pk(_)
4677 | Condition::BitmapEq { .. }
4678 | Condition::BitmapIn { .. }
4679 | Condition::BytesPrefix { .. }
4680 | Condition::FmContains { .. }
4681 | Condition::FmContainsAll { .. }
4682 | Condition::Ann { .. }
4683 | Condition::Range { .. }
4684 | Condition::RangeF64 { .. }
4685 | Condition::SparseMatch { .. }
4686 | Condition::MinHashSimilar { .. }
4687 | Condition::IsNull { .. }
4688 | Condition::IsNotNull { .. }
4689 )
4690 };
4691 if !conditions.iter().all(served) {
4692 return Ok(None);
4693 }
4694 let fast_path =
4695 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4696 crate::trace::QueryTrace::record(|t| {
4697 t.run_count = self.run_refs.len();
4698 t.memtable_rows = self.memtable.len();
4699 t.mutable_run_rows = self.mutable_run.len();
4700 t.conditions_pushed = conditions.len();
4701 t.learned_range_used = conditions.iter().any(|c| match c {
4702 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4703 self.learned_range.contains_key(column_id)
4704 }
4705 _ => false,
4706 });
4707 });
4708 let col_ids: Vec<u16> = projection
4710 .map(|p| p.to_vec())
4711 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4712 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4713 .iter()
4714 .map(|&cid| {
4715 let ty = self
4716 .schema
4717 .columns
4718 .iter()
4719 .find(|c| c.id == cid)
4720 .map(|c| c.ty)
4721 .unwrap_or(TypeId::Bytes);
4722 (cid, ty)
4723 })
4724 .collect();
4725
4726 if fast_path {
4732 let needs_column = conditions.iter().any(|c| match c {
4735 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4736 Condition::RangeF64 { column_id, .. } => {
4737 !self.learned_range.contains_key(column_id)
4738 }
4739 _ => false,
4740 });
4741 let mut reader_opt: Option<RunReader> = if needs_column {
4742 Some(self.open_reader(self.run_refs[0].run_id)?)
4743 } else {
4744 None
4745 };
4746 let mut sets: Vec<RowIdSet> = Vec::new();
4747 for c in conditions {
4748 let s = match c {
4749 Condition::Range { column_id, lo, hi }
4750 if !self.learned_range.contains_key(column_id) =>
4751 {
4752 if reader_opt.is_none() {
4753 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4754 }
4755 reader_opt
4756 .as_mut()
4757 .expect("reader opened for range")
4758 .range_row_id_set_i64(*column_id, *lo, *hi)?
4759 }
4760 Condition::RangeF64 {
4761 column_id,
4762 lo,
4763 lo_inclusive,
4764 hi,
4765 hi_inclusive,
4766 } if !self.learned_range.contains_key(column_id) => {
4767 if reader_opt.is_none() {
4768 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4769 }
4770 reader_opt
4771 .as_mut()
4772 .expect("reader opened for range")
4773 .range_row_id_set_f64(
4774 *column_id,
4775 *lo,
4776 *lo_inclusive,
4777 *hi,
4778 *hi_inclusive,
4779 )?
4780 }
4781 _ => self.resolve_condition(c, snapshot)?,
4782 };
4783 sets.push(s);
4784 }
4785 let candidates = RowIdSet::intersect_many(sets);
4786 crate::trace::QueryTrace::record(|t| {
4787 t.survivor_count = Some(candidates.len());
4788 });
4789 if candidates.is_empty() {
4790 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4791 .iter()
4792 .map(|&id| {
4793 (
4794 id,
4795 columnar::null_native(
4796 proj_pairs
4797 .iter()
4798 .find(|(c, _)| c == &id)
4799 .map(|(_, t)| *t)
4800 .unwrap_or(TypeId::Bytes),
4801 0,
4802 ),
4803 )
4804 })
4805 .collect();
4806 return Ok(Some(cols));
4807 }
4808 let mut reader = match reader_opt.take() {
4809 Some(r) => r,
4810 None => self.open_reader(self.run_refs[0].run_id)?,
4811 };
4812 let candidate_ids = candidates.into_sorted_vec();
4813 let (positions, fast_rid) = if let Some(positions) =
4814 reader.positions_for_row_ids_fast(&candidate_ids)
4815 {
4816 (positions, true)
4817 } else {
4818 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4819 match col {
4820 columnar::NativeColumn::Int64 { data, .. } => {
4821 let mut p: Vec<usize> = candidate_ids
4822 .iter()
4823 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4824 .collect();
4825 p.sort_unstable();
4826 (p, false)
4827 }
4828 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4829 }
4830 };
4831 crate::trace::QueryTrace::record(|t| {
4832 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4833 t.fast_row_id_map = fast_rid;
4834 });
4835 let mut cols = Vec::with_capacity(col_ids.len());
4836 for cid in &col_ids {
4837 let col = reader.column_native(*cid)?;
4838 cols.push((*cid, col.gather(&positions)));
4839 }
4840 return Ok(Some(cols));
4841 }
4842
4843 if !self.run_refs.is_empty() {
4856 use crate::cursor::{drain_cursor_to_columns, Cursor};
4857 let remaining: usize;
4858 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4859 let c = self
4860 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4861 .expect("single-run cursor should build when run_refs.len() == 1");
4862 remaining = c.remaining_rows();
4863 Box::new(c)
4864 } else {
4865 let c = self
4866 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4867 .expect("multi-run cursor should build when run_refs.len() >= 1");
4868 remaining = c.remaining_rows();
4869 Box::new(c)
4870 };
4871 crate::trace::QueryTrace::record(|t| {
4872 if t.survivor_count.is_none() {
4873 t.survivor_count = Some(remaining);
4874 }
4875 });
4876 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4877 return Ok(Some(cols));
4878 }
4879
4880 crate::trace::QueryTrace::record(|t| {
4885 t.scan_mode = crate::trace::ScanMode::Materialized;
4886 t.row_materialized = true;
4887 });
4888 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4889 for c in conditions {
4890 sets.push(self.resolve_condition(c, snapshot)?);
4891 }
4892 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4893 let rows = self.rows_for_rids(&rids, snapshot)?;
4894 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4895 for (cid, ty) in &proj_pairs {
4896 let vals: Vec<Value> = rows
4897 .iter()
4898 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4899 .collect();
4900 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4901 }
4902 Ok(Some(cols))
4903 }
4904
4905 pub fn native_page_cursor(
4920 &self,
4921 snapshot: Snapshot,
4922 projection: Vec<(u16, TypeId)>,
4923 conditions: &[crate::query::Condition],
4924 ) -> Result<Option<NativePageCursor>> {
4925 use crate::cursor::build_page_plans;
4926 if !conditions.is_empty() && !self.indexes_complete {
4929 return Ok(None);
4930 }
4931 if self.run_refs.len() != 1 {
4932 return Ok(None);
4933 }
4934 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4935 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4936
4937 let overlay_rids: HashSet<u64> = {
4940 let mut s = HashSet::new();
4941 for row in self.memtable.visible_versions(snapshot.epoch) {
4942 s.insert(row.row_id.0);
4943 }
4944 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4945 s.insert(row.row_id.0);
4946 }
4947 s
4948 };
4949
4950 let survivors = if conditions.is_empty() {
4954 None
4955 } else {
4956 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4957 };
4958
4959 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4966 survivors.clone()
4967 } else if let Some(s) = &survivors {
4968 let mut run_set = s.clone();
4969 run_set.remove_many(overlay_rids.iter().copied());
4970 Some(run_set)
4971 } else {
4972 Some(RowIdSet::from_unsorted(
4973 rids.iter()
4974 .map(|&r| r as u64)
4975 .filter(|r| !overlay_rids.contains(r))
4976 .collect(),
4977 ))
4978 };
4979
4980 let overlay_rows = if overlay_rids.is_empty() {
4981 Vec::new()
4982 } else {
4983 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4984 self.overlay_visible_rows(snapshot, bound)
4985 };
4986
4987 let plans = if positions.is_empty() {
4989 Vec::new()
4990 } else {
4991 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4992 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4993 };
4994
4995 let overlay = if overlay_rows.is_empty() {
4997 None
4998 } else {
4999 let filtered =
5000 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5001 if filtered.is_empty() {
5002 None
5003 } else {
5004 Some(self.materialize_overlay(&filtered, &projection))
5005 }
5006 };
5007
5008 let overlay_row_count = overlay
5009 .as_ref()
5010 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5011 .unwrap_or(0);
5012 crate::trace::QueryTrace::record(|t| {
5013 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
5014 t.run_count = self.run_refs.len();
5015 t.memtable_rows = self.memtable.len();
5016 t.mutable_run_rows = self.mutable_run.len();
5017 t.overlay_rows = overlay_row_count;
5018 t.conditions_pushed = conditions.len();
5019 t.pages_decoded = plans
5020 .iter()
5021 .map(|p| p.positions.len())
5022 .sum::<usize>()
5023 .min(1);
5024 });
5025
5026 Ok(Some(NativePageCursor::new_with_overlay(
5027 reader, projection, plans, overlay,
5028 )))
5029 }
5030 #[allow(clippy::type_complexity)]
5040 pub fn native_multi_run_cursor(
5041 &self,
5042 snapshot: Snapshot,
5043 projection: Vec<(u16, TypeId)>,
5044 conditions: &[crate::query::Condition],
5045 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
5046 use crate::cursor::{MultiRunCursor, RunStream};
5047 use crate::sorted_run::SYS_ROW_ID;
5048 use std::collections::{BinaryHeap, HashMap, HashSet};
5049 if !conditions.is_empty() && !self.indexes_complete {
5052 return Ok(None);
5053 }
5054 if self.run_refs.is_empty() {
5055 return Ok(None);
5056 }
5057
5058 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
5060 Vec::with_capacity(self.run_refs.len());
5061 for rr in &self.run_refs {
5062 let mut reader = self.open_reader(rr.run_id)?;
5063 let (rids, eps, del) = reader.system_columns_native()?;
5064 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
5065 run_meta.push((reader, rids, eps, del, page_rows));
5066 }
5067
5068 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
5072 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
5073 for i in 0..rids.len() {
5074 let rid = rids[i] as u64;
5075 let e = eps[i] as u64;
5076 if e > snapshot.epoch.0 {
5077 continue;
5078 }
5079 let is_del = del[i] != 0;
5080 best.entry(rid)
5081 .and_modify(|cur| {
5082 if e > cur.0 {
5083 *cur = (e, run_idx, i, is_del);
5084 }
5085 })
5086 .or_insert((e, run_idx, i, is_del));
5087 }
5088 }
5089
5090 let overlay_rids: HashSet<u64> = {
5092 let mut s = HashSet::new();
5093 for row in self.memtable.visible_versions(snapshot.epoch) {
5094 s.insert(row.row_id.0);
5095 }
5096 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5097 s.insert(row.row_id.0);
5098 }
5099 s
5100 };
5101
5102 let survivors: Option<RowIdSet> = if conditions.is_empty() {
5104 None
5105 } else {
5106 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5107 for c in conditions {
5108 sets.push(self.resolve_condition(c, snapshot)?);
5109 }
5110 Some(RowIdSet::intersect_many(sets))
5111 };
5112
5113 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
5117 for (rid, (_, run_idx, pos, deleted)) in &best {
5118 if *deleted {
5119 continue;
5120 }
5121 if overlay_rids.contains(rid) {
5122 continue;
5123 }
5124 if let Some(s) = &survivors {
5125 if !s.contains(*rid) {
5126 continue;
5127 }
5128 }
5129 per_run[*run_idx].push((*rid, *pos));
5130 }
5131 for v in per_run.iter_mut() {
5132 v.sort_unstable_by_key(|&(rid, _)| rid);
5133 }
5134
5135 let mut streams = Vec::with_capacity(run_meta.len());
5137 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
5138 let mut total = 0usize;
5139 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
5140 let mut starts = Vec::with_capacity(page_rows.len());
5141 let mut acc = 0usize;
5142 for &r in &page_rows {
5143 starts.push(acc);
5144 acc += r;
5145 }
5146 let mut survivors_vec: Vec<(u64, usize, usize)> =
5147 Vec::with_capacity(per_run[run_idx].len());
5148 for &(rid, pos) in &per_run[run_idx] {
5149 let page_seq = match starts.partition_point(|&s| s <= pos) {
5150 0 => continue,
5151 p => p - 1,
5152 };
5153 let within = pos - starts[page_seq];
5154 survivors_vec.push((rid, page_seq, within));
5155 }
5156 total += survivors_vec.len();
5157 if let Some(&(rid, _, _)) = survivors_vec.first() {
5158 heap.push(std::cmp::Reverse((rid, run_idx)));
5159 }
5160 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5161 }
5162
5163 let overlay_rows = if overlay_rids.is_empty() {
5165 Vec::new()
5166 } else {
5167 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5168 self.overlay_visible_rows(snapshot, bound)
5169 };
5170 let overlay = if overlay_rows.is_empty() {
5171 None
5172 } else {
5173 let filtered =
5174 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5175 if filtered.is_empty() {
5176 None
5177 } else {
5178 Some(self.materialize_overlay(&filtered, &projection))
5179 }
5180 };
5181
5182 let overlay_row_count = overlay
5183 .as_ref()
5184 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5185 .unwrap_or(0);
5186 crate::trace::QueryTrace::record(|t| {
5187 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5188 t.run_count = self.run_refs.len();
5189 t.memtable_rows = self.memtable.len();
5190 t.mutable_run_rows = self.mutable_run.len();
5191 t.overlay_rows = overlay_row_count;
5192 t.conditions_pushed = conditions.len();
5193 t.survivor_count = Some(total);
5194 });
5195
5196 Ok(Some(MultiRunCursor::new(
5197 streams, projection, heap, total, overlay,
5198 )))
5199 }
5200
5201 fn overlay_materialization_bound<'a>(
5213 conditions: &[crate::query::Condition],
5214 survivors: &'a Option<RowIdSet>,
5215 ) -> Option<&'a RowIdSet> {
5216 use crate::query::Condition;
5217 let has_range = conditions
5218 .iter()
5219 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5220 if has_range {
5221 None
5222 } else {
5223 survivors.as_ref()
5224 }
5225 }
5226
5227 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5239 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5240 let mut fold = |row: Row| {
5241 if let Some(b) = bound {
5242 if !b.contains(row.row_id.0) {
5243 return;
5244 }
5245 }
5246 best.entry(row.row_id.0)
5247 .and_modify(|(be, br)| {
5248 if row.committed_epoch > *be {
5249 *be = row.committed_epoch;
5250 *br = row.clone();
5251 }
5252 })
5253 .or_insert_with(|| (row.committed_epoch, row));
5254 };
5255 for row in self.memtable.visible_versions(snapshot.epoch) {
5256 fold(row);
5257 }
5258 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5259 fold(row);
5260 }
5261 let mut out: Vec<Row> = best
5262 .into_values()
5263 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5264 .collect();
5265 out.sort_by_key(|r| r.row_id);
5266 out
5267 }
5268
5269 fn filter_overlay_rows(
5277 &self,
5278 rows: Vec<Row>,
5279 conditions: &[crate::query::Condition],
5280 survivors: Option<&RowIdSet>,
5281 snapshot: Snapshot,
5282 ) -> Result<Vec<Row>> {
5283 if conditions.is_empty() {
5284 return Ok(rows);
5285 }
5286 use crate::query::Condition;
5287 let all_index_served = !conditions
5291 .iter()
5292 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5293 if all_index_served {
5294 return Ok(rows
5295 .into_iter()
5296 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5297 .collect());
5298 }
5299 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5302 for c in conditions {
5303 let s = match c {
5304 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5305 _ => self.resolve_condition(c, snapshot)?,
5306 };
5307 per_cond_sets.push(s);
5308 }
5309 Ok(rows
5310 .into_iter()
5311 .filter(|row| {
5312 conditions.iter().enumerate().all(|(i, c)| match c {
5313 Condition::Range { column_id, lo, hi } => {
5314 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5315 }
5316 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5317 match row.columns.get(column_id) {
5318 Some(Value::Float64(v)) => {
5319 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5320 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5321 lo_ok && hi_ok
5322 }
5323 _ => false,
5324 }
5325 }
5326 _ => per_cond_sets[i].contains(row.row_id.0),
5327 })
5328 })
5329 .collect())
5330 }
5331
5332 fn materialize_overlay(
5335 &self,
5336 rows: &[Row],
5337 projection: &[(u16, TypeId)],
5338 ) -> Vec<columnar::NativeColumn> {
5339 if projection.is_empty() {
5340 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5341 }
5342 let mut cols = Vec::with_capacity(projection.len());
5343 for (cid, ty) in projection {
5344 let vals: Vec<Value> = rows
5345 .iter()
5346 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5347 .collect();
5348 cols.push(columnar::values_to_native(*ty, &vals));
5349 }
5350 cols
5351 }
5352
5353 fn resolve_survivor_rids(
5358 &self,
5359 conditions: &[crate::query::Condition],
5360 reader: &mut RunReader,
5361 ) -> Result<RowIdSet> {
5362 use crate::query::Condition;
5363 let mut sets: Vec<RowIdSet> = Vec::new();
5364 for c in conditions {
5365 let s: RowIdSet = match c {
5366 Condition::Pk(key) => {
5367 let lookup = self
5368 .schema
5369 .primary_key()
5370 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5371 .unwrap_or_else(|| key.clone());
5372 self.hot
5373 .get(&lookup)
5374 .map(|r| RowIdSet::one(r.0))
5375 .unwrap_or_else(RowIdSet::empty)
5376 }
5377 Condition::BitmapEq { column_id, value } => {
5378 let lookup = self.index_lookup_key_bytes(*column_id, value);
5379 self.bitmap
5380 .get(column_id)
5381 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5382 .unwrap_or_else(RowIdSet::empty)
5383 }
5384 Condition::BitmapIn { column_id, values } => {
5385 let bm = self.bitmap.get(column_id);
5386 let mut acc = roaring::RoaringBitmap::new();
5387 if let Some(b) = bm {
5388 for v in values {
5389 let lookup = self.index_lookup_key_bytes(*column_id, v);
5390 acc |= b.get(&lookup);
5391 }
5392 }
5393 RowIdSet::from_roaring(acc)
5394 }
5395 Condition::BytesPrefix { column_id, prefix } => {
5396 if let Some(b) = self.bitmap.get(column_id) {
5397 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5398 let mut acc = roaring::RoaringBitmap::new();
5399 for key in b.keys() {
5400 if key.starts_with(&lookup_prefix) {
5401 acc |= b.get(key);
5402 }
5403 }
5404 RowIdSet::from_roaring(acc)
5405 } else {
5406 RowIdSet::empty()
5407 }
5408 }
5409 Condition::FmContains { column_id, pattern } => self
5410 .fm
5411 .get(column_id)
5412 .map(|f| {
5413 RowIdSet::from_unsorted(
5414 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5415 )
5416 })
5417 .unwrap_or_else(RowIdSet::empty),
5418 Condition::FmContainsAll {
5419 column_id,
5420 patterns,
5421 } => {
5422 if let Some(f) = self.fm.get(column_id) {
5423 let sets: Vec<RowIdSet> = patterns
5424 .iter()
5425 .map(|pat| {
5426 RowIdSet::from_unsorted(
5427 f.locate(pat).into_iter().map(|r| r.0).collect(),
5428 )
5429 })
5430 .collect();
5431 RowIdSet::intersect_many(sets)
5432 } else {
5433 RowIdSet::empty()
5434 }
5435 }
5436 Condition::Ann {
5437 column_id,
5438 query,
5439 k,
5440 } => self
5441 .ann
5442 .get(column_id)
5443 .map(|a| {
5444 RowIdSet::from_unsorted(
5445 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5446 )
5447 })
5448 .unwrap_or_else(RowIdSet::empty),
5449 Condition::SparseMatch {
5450 column_id,
5451 query,
5452 k,
5453 } => self
5454 .sparse
5455 .get(column_id)
5456 .map(|s| {
5457 RowIdSet::from_unsorted(
5458 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5459 )
5460 })
5461 .unwrap_or_else(RowIdSet::empty),
5462 Condition::MinHashSimilar {
5463 column_id,
5464 query,
5465 k,
5466 } => self
5467 .minhash
5468 .get(column_id)
5469 .map(|mh| {
5470 RowIdSet::from_unsorted(
5471 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5472 )
5473 })
5474 .unwrap_or_else(RowIdSet::empty),
5475 Condition::Range { column_id, lo, hi } => {
5476 if let Some(li) = self.learned_range.get(column_id) {
5477 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5478 } else {
5479 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5480 }
5481 }
5482 Condition::RangeF64 {
5483 column_id,
5484 lo,
5485 lo_inclusive,
5486 hi,
5487 hi_inclusive,
5488 } => {
5489 if let Some(li) = self.learned_range.get(column_id) {
5490 RowIdSet::from_unsorted(
5491 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5492 .into_iter()
5493 .collect(),
5494 )
5495 } else {
5496 reader.range_row_id_set_f64(
5497 *column_id,
5498 *lo,
5499 *lo_inclusive,
5500 *hi,
5501 *hi_inclusive,
5502 )?
5503 }
5504 }
5505 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5506 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5507 };
5508 sets.push(s);
5509 }
5510 Ok(RowIdSet::intersect_many(sets))
5511 }
5512
5513 pub fn scan_cursor(
5534 &self,
5535 snapshot: Snapshot,
5536 projection: Vec<(u16, TypeId)>,
5537 conditions: &[crate::query::Condition],
5538 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5539 if !conditions.is_empty() && !self.indexes_complete {
5545 return Ok(None);
5546 }
5547 if self.run_refs.len() == 1 {
5548 Ok(self
5549 .native_page_cursor(snapshot, projection, conditions)?
5550 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5551 } else {
5552 Ok(self
5553 .native_multi_run_cursor(snapshot, projection, conditions)?
5554 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5555 }
5556 }
5557
5558 pub fn aggregate_native(
5572 &self,
5573 snapshot: Snapshot,
5574 column: Option<u16>,
5575 conditions: &[crate::query::Condition],
5576 agg: NativeAgg,
5577 ) -> Result<Option<NativeAggResult>> {
5578 if self.run_refs.len() == 1 && conditions.is_empty() {
5580 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5581 return Ok(Some(res));
5582 }
5583 }
5584 if matches!(agg, NativeAgg::Count) && column.is_none() {
5586 return Ok(self
5587 .scan_cursor(snapshot, Vec::new(), conditions)?
5588 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5589 }
5590 let cid = match column {
5593 Some(c) => c,
5594 None => return Ok(None),
5595 };
5596 let ty = self.column_type(cid);
5597 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5598 return Ok(None);
5599 };
5600 match ty {
5601 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5602 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5603 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5604 }
5605 TypeId::Float64 => {
5606 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5607 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5608 }
5609 _ => Ok(None),
5610 }
5611 }
5612
5613 fn aggregate_from_stats(
5621 &self,
5622 snapshot: Snapshot,
5623 column: Option<u16>,
5624 agg: NativeAgg,
5625 ) -> Result<Option<NativeAggResult>> {
5626 let cid = match (agg, column) {
5627 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5628 _ => return Ok(None), };
5630 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5631 return Ok(None);
5632 };
5633 let Some(cs) = stats.get(&cid) else {
5634 return Ok(None);
5635 };
5636 match agg {
5637 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5639 self.live_count.saturating_sub(cs.null_count),
5640 ))),
5641 NativeAgg::Min | NativeAgg::Max => {
5642 let bound = if agg == NativeAgg::Min {
5643 &cs.min
5644 } else {
5645 &cs.max
5646 };
5647 match bound {
5648 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5649 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5650 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5655 None => Ok(None),
5656 }
5657 }
5658 _ => Ok(None),
5659 }
5660 }
5661
5662 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5671 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5672 return Ok(None);
5673 }
5674 self.ensure_indexes_complete()?;
5677 let reader = self.open_reader(self.run_refs[0].run_id)?;
5678 if self.live_count != reader.row_count() as u64 {
5679 return Ok(None);
5680 }
5681 let Some(bm) = self.bitmap.get(&column_id) else {
5682 return Ok(None); };
5684 let mut distinct = bm.value_count() as u64;
5685 if !bm.get(&Value::Null.encode_key()).is_empty() {
5688 distinct = distinct.saturating_sub(1);
5689 }
5690 Ok(Some(distinct))
5691 }
5692
5693 pub fn aggregate_incremental(
5705 &mut self,
5706 cache_key: u64,
5707 conditions: &[crate::query::Condition],
5708 column: Option<u16>,
5709 agg: NativeAgg,
5710 ) -> Result<IncrementalAggResult> {
5711 let snap = self.snapshot();
5712 let cur_wm = self.allocator.current().0;
5713 let cur_epoch = snap.epoch.0;
5714 let incremental_ok =
5721 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5722
5723 if incremental_ok {
5726 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5727 if cached.epoch == cur_epoch {
5728 return Ok(IncrementalAggResult {
5729 state: cached.state,
5730 incremental: true,
5731 delta_rows: 0,
5732 });
5733 }
5734 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5735 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5736 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5737 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5738 let delta_state = agg_state_from_rows(
5739 &delta_rows,
5740 conditions,
5741 &index_sets,
5742 column,
5743 agg,
5744 &self.schema,
5745 )?;
5746 let merged = cached.state.merge(delta_state);
5747 let delta_n = delta_rids.len() as u64;
5748 self.agg_cache.insert(
5749 cache_key,
5750 CachedAgg {
5751 state: merged.clone(),
5752 watermark: cur_wm,
5753 epoch: cur_epoch,
5754 },
5755 );
5756 return Ok(IncrementalAggResult {
5757 state: merged,
5758 incremental: true,
5759 delta_rows: delta_n,
5760 });
5761 }
5762 }
5763 }
5764
5765 let cursor_ok =
5770 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5771 let state = if cursor_ok && agg != NativeAgg::Avg {
5772 match self.aggregate_native(snap, column, conditions, agg)? {
5773 Some(result) => {
5774 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5775 }
5776 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5777 }
5778 } else {
5779 self.agg_state_full_scan(conditions, column, agg, snap)?
5780 };
5781 if incremental_ok {
5783 self.agg_cache.insert(
5784 cache_key,
5785 CachedAgg {
5786 state: state.clone(),
5787 watermark: cur_wm,
5788 epoch: cur_epoch,
5789 },
5790 );
5791 }
5792 Ok(IncrementalAggResult {
5793 state,
5794 incremental: false,
5795 delta_rows: 0,
5796 })
5797 }
5798
5799 fn agg_state_full_scan(
5802 &self,
5803 conditions: &[crate::query::Condition],
5804 column: Option<u16>,
5805 agg: NativeAgg,
5806 snap: Snapshot,
5807 ) -> Result<AggState> {
5808 let rows = self.visible_rows(snap)?;
5809 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5810 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5811 }
5812
5813 fn resolve_index_conditions(
5816 &self,
5817 conditions: &[crate::query::Condition],
5818 snapshot: Snapshot,
5819 ) -> Result<Vec<RowIdSet>> {
5820 use crate::query::Condition;
5821 let mut sets = Vec::new();
5822 for c in conditions {
5823 if matches!(
5824 c,
5825 Condition::Ann { .. }
5826 | Condition::SparseMatch { .. }
5827 | Condition::MinHashSimilar { .. }
5828 ) {
5829 sets.push(self.resolve_condition(c, snapshot)?);
5830 }
5831 }
5832 Ok(sets)
5833 }
5834
5835 fn column_type(&self, cid: u16) -> TypeId {
5836 self.schema
5837 .columns
5838 .iter()
5839 .find(|c| c.id == cid)
5840 .map(|c| c.ty)
5841 .unwrap_or(TypeId::Bytes)
5842 }
5843
5844 pub fn approx_aggregate(
5853 &mut self,
5854 conditions: &[crate::query::Condition],
5855 column: Option<u16>,
5856 agg: ApproxAgg,
5857 z: f64,
5858 ) -> Result<Option<ApproxResult>> {
5859 use crate::query::Condition;
5860 self.ensure_reservoir_complete()?;
5861 let snapshot = self.snapshot();
5862 let n_pop = self.live_count;
5863 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5864 if sample_rids.is_empty() {
5865 return Ok(None);
5866 }
5867 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5869 let s = live_sample.len();
5870 if s == 0 {
5871 return Ok(None);
5872 }
5873
5874 let mut index_sets: Vec<RowIdSet> = Vec::new();
5877 for c in conditions {
5878 if matches!(
5879 c,
5880 Condition::Ann { .. }
5881 | Condition::SparseMatch { .. }
5882 | Condition::MinHashSimilar { .. }
5883 ) {
5884 index_sets.push(self.resolve_condition(c, snapshot)?);
5885 }
5886 }
5887
5888 let cid = match (agg, column) {
5890 (ApproxAgg::Count, _) => None,
5891 (_, Some(c)) => Some(c),
5892 _ => return Ok(None),
5893 };
5894 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5895 for r in &live_sample {
5896 if !conditions
5898 .iter()
5899 .all(|c| condition_matches_row(c, r, &self.schema))
5900 {
5901 continue;
5902 }
5903 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5905 continue;
5906 }
5907 if let Some(cid) = cid {
5908 if let Some(v) = as_f64(r.columns.get(&cid)) {
5909 passing_vals.push(v);
5910 } } else {
5912 passing_vals.push(0.0); }
5914 }
5915 let m = passing_vals.len();
5916
5917 let (point, half) = match agg {
5918 ApproxAgg::Count => {
5919 let p = m as f64 / s as f64;
5921 let point = n_pop as f64 * p;
5922 let var = if s > 1 {
5923 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5924 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5925 } else {
5926 0.0
5927 };
5928 (point, z * var.sqrt())
5929 }
5930 ApproxAgg::Sum => {
5931 let y: Vec<f64> = live_sample
5933 .iter()
5934 .map(|r| {
5935 let passes_row = conditions
5936 .iter()
5937 .all(|c| condition_matches_row(c, r, &self.schema))
5938 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5939 if passes_row {
5940 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5941 } else {
5942 0.0
5943 }
5944 })
5945 .collect();
5946 let mean_y = y.iter().sum::<f64>() / s as f64;
5947 let point = n_pop as f64 * mean_y;
5948 let var = if s > 1 {
5949 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5950 let var_y = ss / (s - 1) as f64;
5951 n_pop as f64 * n_pop as f64 * var_y / s as f64
5952 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5953 } else {
5954 0.0
5955 };
5956 (point, z * var.sqrt())
5957 }
5958 ApproxAgg::Avg => {
5959 if m == 0 {
5960 return Ok(Some(ApproxResult {
5961 point: 0.0,
5962 ci_low: 0.0,
5963 ci_high: 0.0,
5964 n_population: n_pop,
5965 n_sample_live: s,
5966 n_passing: 0,
5967 }));
5968 }
5969 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5970 let half = if m > 1 {
5971 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5972 let sd = (ss / (m - 1) as f64).sqrt();
5973 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5974 z * sd / (m as f64).sqrt() * fpc.sqrt()
5975 } else {
5976 0.0
5977 };
5978 (mean, half)
5979 }
5980 };
5981
5982 Ok(Some(ApproxResult {
5983 point,
5984 ci_low: point - half,
5985 ci_high: point + half,
5986 n_population: n_pop,
5987 n_sample_live: s,
5988 n_passing: m,
5989 }))
5990 }
5991
5992 pub fn exact_column_stats(
6000 &self,
6001 _snapshot: Snapshot,
6002 projection: &[u16],
6003 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
6004 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
6005 return Ok(None);
6006 }
6007 let reader = self.open_reader(self.run_refs[0].run_id)?;
6008 if self.live_count != reader.row_count() as u64 {
6009 return Ok(None);
6010 }
6011 let mut out = HashMap::new();
6012 for &cid in projection {
6013 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
6014 Some(c) => c,
6015 None => continue,
6016 };
6017 let Some(stats) = reader.column_page_stats(cid) else {
6019 out.insert(
6020 cid,
6021 ColumnStat {
6022 min: None,
6023 max: None,
6024 null_count: self.live_count,
6025 },
6026 );
6027 continue;
6028 };
6029 let stat = match cdef.ty {
6030 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
6031 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
6032 min: mn.map(Value::Int64),
6033 max: mx.map(Value::Int64),
6034 null_count: n,
6035 })
6036 }
6037 TypeId::Float64 => {
6038 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
6039 min: mn.map(Value::Float64),
6040 max: mx.map(Value::Float64),
6041 null_count: n,
6042 })
6043 }
6044 _ => None,
6045 };
6046 if let Some(s) = stat {
6047 out.insert(cid, s);
6048 }
6049 }
6050 Ok(Some(out))
6051 }
6052
6053 pub fn dir(&self) -> &Path {
6054 &self.dir
6055 }
6056
6057 pub fn schema(&self) -> &Schema {
6058 &self.schema
6059 }
6060
6061 pub(crate) fn prepare_alter_column(
6062 &mut self,
6063 column_name: &str,
6064 change: &AlterColumn,
6065 ) -> Result<ColumnDef> {
6066 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
6067 return Err(MongrelError::InvalidArgument(
6068 "ALTER COLUMN requires committing staged writes first".into(),
6069 ));
6070 }
6071 let old = self
6072 .schema
6073 .columns
6074 .iter()
6075 .find(|c| c.name == column_name)
6076 .cloned()
6077 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6078 let mut next = old.clone();
6079
6080 if let Some(name) = &change.name {
6081 let trimmed = name.trim();
6082 if trimmed.is_empty() {
6083 return Err(MongrelError::InvalidArgument(
6084 "ALTER COLUMN name must not be empty".into(),
6085 ));
6086 }
6087 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
6088 return Err(MongrelError::Schema(format!(
6089 "column {trimmed} already exists"
6090 )));
6091 }
6092 next.name = trimmed.to_string();
6093 }
6094
6095 if let Some(ty) = change.ty {
6096 next.ty = ty;
6097 }
6098 if let Some(flags) = change.flags {
6099 validate_alter_column_flags(old.flags, flags)?;
6100 next.flags = flags;
6101 }
6102
6103 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
6104 if old.flags.contains(ColumnFlags::NULLABLE)
6105 && !next.flags.contains(ColumnFlags::NULLABLE)
6106 && self.column_has_nulls(old.id)?
6107 {
6108 return Err(MongrelError::InvalidArgument(format!(
6109 "column '{}' contains NULL values",
6110 old.name
6111 )));
6112 }
6113 Ok(next)
6114 }
6115
6116 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
6117 let idx = self
6118 .schema
6119 .columns
6120 .iter()
6121 .position(|c| c.id == column.id)
6122 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
6123 if self.schema.columns[idx] == column {
6124 return Ok(());
6125 }
6126 self.schema.columns[idx] = column;
6127 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6128 self.schema.validate_auto_increment()?;
6129 self.auto_inc = resolve_auto_inc(&self.schema);
6130 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
6131 write_schema(&self.dir, &self.schema)?;
6132 self.clear_result_cache();
6133 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6134 self.persist_manifest(self.current_epoch())?;
6135 Ok(())
6136 }
6137
6138 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
6139 let column = self.prepare_alter_column(column_name, &change)?;
6140 self.apply_altered_column(column.clone())?;
6141 Ok(column)
6142 }
6143
6144 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
6145 if self.live_count == 0 {
6146 return Ok(false);
6147 }
6148 let snap = self.snapshot();
6149 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6150 Ok(columns
6151 .first()
6152 .map(|(_, col)| col.null_count(col.len()) != 0)
6153 .unwrap_or(true))
6154 }
6155
6156 fn has_stored_versions(&self) -> bool {
6157 !self.memtable.is_empty()
6158 || !self.mutable_run.is_empty()
6159 || self.run_refs.iter().any(|r| r.row_count > 0)
6160 || !self.retiring.is_empty()
6161 }
6162
6163 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
6168 if self.schema.columns.iter().any(|c| c.name == name) {
6169 return Err(MongrelError::Schema(format!(
6170 "column {name} already exists"
6171 )));
6172 }
6173 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
6174 self.schema.columns.push(ColumnDef {
6175 id,
6176 name: name.to_string(),
6177 ty,
6178 flags,
6179 });
6180 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6181 self.schema.validate_auto_increment()?;
6182 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6183 self.auto_inc = resolve_auto_inc(&self.schema);
6184 }
6185 write_schema(&self.dir, &self.schema)?;
6186 self.clear_result_cache();
6187 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6189 self.persist_manifest(self.current_epoch())?;
6190 Ok(id)
6191 }
6192
6193 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6202 let cid = self
6203 .schema
6204 .columns
6205 .iter()
6206 .find(|c| c.name == column_name)
6207 .map(|c| c.id)
6208 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6209 let ty = self
6210 .schema
6211 .columns
6212 .iter()
6213 .find(|c| c.id == cid)
6214 .map(|c| c.ty)
6215 .unwrap_or(TypeId::Int64);
6216 if !matches!(
6217 ty,
6218 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6219 ) {
6220 return Err(MongrelError::Schema(format!(
6221 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6222 )));
6223 }
6224 if self
6225 .schema
6226 .indexes
6227 .iter()
6228 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6229 {
6230 return Ok(()); }
6232 self.schema.indexes.push(IndexDef {
6233 name: format!("{}_learned_range", column_name),
6234 column_id: cid,
6235 kind: IndexKind::LearnedRange,
6236 predicate: None,
6237 });
6238 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6239 write_schema(&self.dir, &self.schema)?;
6240 self.build_learned_ranges()?;
6241 Ok(())
6242 }
6243
6244 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6247 self.sync_byte_threshold = threshold;
6248 if let WalSink::Private(w) = &mut self.wal {
6249 w.set_sync_byte_threshold(threshold);
6250 }
6251 }
6252
6253 pub fn page_cache_flush(&self) {
6257 self.page_cache.flush_to_disk();
6258 }
6259
6260 pub fn page_cache_len(&self) -> usize {
6262 self.page_cache.len()
6263 }
6264
6265 pub fn decoded_cache_len(&self) -> usize {
6268 self.decoded_cache.len()
6269 }
6270
6271 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6274 self.memtable.drain_sorted()
6275 }
6276
6277 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6278 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6279 }
6280
6281 pub(crate) fn table_dir(&self) -> &Path {
6282 &self.dir
6283 }
6284
6285 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6286 &self.schema
6287 }
6288
6289 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6290 let id = self.next_run_id;
6291 self.next_run_id += 1;
6292 id
6293 }
6294
6295 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6296 self.run_refs.push(run_ref);
6297 }
6298
6299 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6309 self.retiring.push(crate::manifest::RetiredRun {
6310 run_id,
6311 retire_epoch,
6312 });
6313 }
6314
6315 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6319 if self.retiring.is_empty() {
6320 return Ok(0);
6321 }
6322 let mut reaped = 0;
6323 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6324 for r in std::mem::take(&mut self.retiring) {
6330 if min_active.0 >= r.retire_epoch {
6331 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6332 reaped += 1;
6333 } else {
6334 kept.push(r);
6335 }
6336 }
6337 self.retiring = kept;
6338 if reaped > 0 {
6339 self.persist_manifest(self.current_epoch())?;
6340 }
6341 Ok(reaped)
6342 }
6343
6344 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6345 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6346 return false;
6347 }
6348 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6349 self.run_refs.push(run_ref);
6350 self.indexes_complete = false;
6351 true
6352 }
6353
6354 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6355 self.kek.as_ref()
6356 }
6357
6358 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6359 let mut reader = RunReader::open_with_cache(
6360 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6361 self.schema.clone(),
6362 self.kek.clone(),
6363 Some(self.page_cache.clone()),
6364 Some(self.decoded_cache.clone()),
6365 self.table_id,
6366 Some(&self.verified_runs),
6367 )?;
6368 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6372 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6373 }
6374 Ok(reader)
6375 }
6376
6377 pub(crate) fn run_refs(&self) -> &[RunRef] {
6378 &self.run_refs
6379 }
6380
6381 pub(crate) fn runs_dir(&self) -> PathBuf {
6382 self.dir.join(RUNS_DIR)
6383 }
6384
6385 pub(crate) fn wal_dir(&self) -> PathBuf {
6386 self.dir.join(WAL_DIR)
6387 }
6388
6389 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6390 self.run_refs = refs;
6391 }
6392
6393 pub(crate) fn next_run_id(&self) -> u64 {
6394 self.next_run_id
6395 }
6396
6397 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6398 self.compaction_zstd_level
6399 }
6400
6401 pub(crate) fn bump_next_run_id(&mut self) {
6402 self.next_run_id += 1;
6403 }
6404
6405 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6406 self.kek.clone()
6407 }
6408
6409 #[cfg(feature = "encryption")]
6413 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6414 self.kek.as_ref().map(|k| k.derive_idx_key())
6415 }
6416
6417 #[cfg(not(feature = "encryption"))]
6418 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6419 None
6420 }
6421
6422 #[cfg(feature = "encryption")]
6426 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6427 self.kek.as_ref().map(|k| *k.derive_meta_key())
6428 }
6429
6430 #[cfg(not(feature = "encryption"))]
6431 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6432 None
6433 }
6434
6435 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6438 self.column_keys
6439 .iter()
6440 .map(|(&id, &(_, scheme))| (id, scheme))
6441 .collect()
6442 }
6443
6444 #[cfg(feature = "encryption")]
6449 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6450 self.tokenize_value_enc(column_id, v)
6451 }
6452
6453 #[cfg(feature = "encryption")]
6454 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6455 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6456 let (key, scheme) = self.column_keys.get(&column_id)?;
6457 let token: Vec<u8> = match (*scheme, v) {
6458 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6459 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6460 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6461 _ => hmac_token(key, &v.encode_key()).to_vec(),
6462 };
6463 Some(Value::Bytes(token))
6464 }
6465
6466 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6468 self.index_lookup_key_bytes(column_id, &v.encode_key())
6469 }
6470
6471 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6474 #[cfg(feature = "encryption")]
6475 {
6476 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6477 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6478 if *scheme == SCHEME_HMAC_EQ {
6479 return hmac_token(key, encoded).to_vec();
6480 }
6481 }
6482 }
6483 let _ = column_id;
6484 encoded.to_vec()
6485 }
6486}
6487
6488fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6489 let columnar::NativeColumn::Int64 { data, validity } = col else {
6490 return false;
6491 };
6492 if data.len() < n || !columnar::all_non_null(validity, n) {
6493 return false;
6494 }
6495 data.iter()
6496 .take(n)
6497 .zip(data.iter().skip(1))
6498 .all(|(a, b)| a < b)
6499}
6500
6501#[derive(Debug, Clone)]
6505pub struct ColumnStat {
6506 pub min: Option<Value>,
6507 pub max: Option<Value>,
6508 pub null_count: u64,
6509}
6510
6511#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6513pub enum NativeAgg {
6514 Count,
6515 Sum,
6516 Min,
6517 Max,
6518 Avg,
6519}
6520
6521#[derive(Debug, Clone, PartialEq)]
6523pub enum NativeAggResult {
6524 Count(u64),
6525 Int(i64),
6526 Float(f64),
6527 Null,
6529}
6530
6531#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6533pub enum ApproxAgg {
6534 Count,
6535 Sum,
6536 Avg,
6537}
6538
6539#[derive(Debug, Clone)]
6543pub struct ApproxResult {
6544 pub point: f64,
6546 pub ci_low: f64,
6548 pub ci_high: f64,
6550 pub n_population: u64,
6552 pub n_sample_live: usize,
6554 pub n_passing: usize,
6556}
6557
6558#[derive(Debug, Clone, PartialEq)]
6563pub enum AggState {
6564 Count(u64),
6566 SumI {
6568 sum: i128,
6569 count: u64,
6570 },
6571 SumF {
6573 sum: f64,
6574 count: u64,
6575 },
6576 AvgI {
6578 sum: i128,
6579 count: u64,
6580 },
6581 AvgF {
6583 sum: f64,
6584 count: u64,
6585 },
6586 MinI(i64),
6588 MaxI(i64),
6589 MinF(f64),
6591 MaxF(f64),
6592 Empty,
6594}
6595
6596impl AggState {
6597 pub fn merge(self, other: AggState) -> AggState {
6599 use AggState::*;
6600 match (self, other) {
6601 (Empty, x) | (x, Empty) => x,
6602 (Count(a), Count(b)) => Count(a + b),
6603 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6604 sum: sa + sb,
6605 count: ca + cb,
6606 },
6607 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6608 sum: sa + sb,
6609 count: ca + cb,
6610 },
6611 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6612 sum: sa + sb,
6613 count: ca + cb,
6614 },
6615 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6616 sum: sa + sb,
6617 count: ca + cb,
6618 },
6619 (MinI(a), MinI(b)) => MinI(a.min(b)),
6620 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6621 (MinF(a), MinF(b)) => MinF(a.min(b)),
6622 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6623 _ => Empty, }
6625 }
6626
6627 pub fn point(&self) -> Option<f64> {
6629 match self {
6630 AggState::Count(n) => Some(*n as f64),
6631 AggState::SumI { sum, .. } => Some(*sum as f64),
6632 AggState::SumF { sum, .. } => Some(*sum),
6633 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6634 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6635 AggState::MinI(n) => Some(*n as f64),
6636 AggState::MaxI(n) => Some(*n as f64),
6637 AggState::MinF(n) => Some(*n),
6638 AggState::MaxF(n) => Some(*n),
6639 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6640 }
6641 }
6642
6643 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6647 let is_float = matches!(ty, Some(TypeId::Float64));
6648 match (agg, result) {
6649 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6650 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6651 sum: x as i128,
6652 count: 1, },
6654 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6655 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6656 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6657 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6658 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6659 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6660 (NativeAgg::Count, _) => AggState::Empty,
6661 (_, NativeAggResult::Null) => AggState::Empty,
6662 _ => {
6663 let _ = is_float;
6664 AggState::Empty
6665 }
6666 }
6667 }
6668}
6669
6670#[derive(Debug, Clone)]
6673pub struct CachedAgg {
6674 pub state: AggState,
6675 pub watermark: u64,
6676 pub epoch: u64,
6677}
6678
6679#[derive(Debug, Clone)]
6681pub struct IncrementalAggResult {
6682 pub state: AggState,
6684 pub incremental: bool,
6687 pub delta_rows: u64,
6689}
6690
6691fn agg_state_from_rows(
6695 rows: &[Row],
6696 conditions: &[crate::query::Condition],
6697 index_sets: &[RowIdSet],
6698 column: Option<u16>,
6699 agg: NativeAgg,
6700 schema: &Schema,
6701) -> Result<AggState> {
6702 let mut count: u64 = 0;
6703 let mut sum_i: i128 = 0;
6704 let mut sum_f: f64 = 0.0;
6705 let mut mn_i: i64 = i64::MAX;
6706 let mut mx_i: i64 = i64::MIN;
6707 let mut mn_f: f64 = f64::INFINITY;
6708 let mut mx_f: f64 = f64::NEG_INFINITY;
6709 let mut saw_int = false;
6710 let mut saw_float = false;
6711 for r in rows {
6712 if !conditions
6713 .iter()
6714 .all(|c| condition_matches_row(c, r, schema))
6715 {
6716 continue;
6717 }
6718 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6719 continue;
6720 }
6721 match agg {
6722 NativeAgg::Count => match column {
6723 None => count += 1,
6725 Some(cid) => match r.columns.get(&cid) {
6728 None | Some(Value::Null) => {}
6729 Some(_) => count += 1,
6730 },
6731 },
6732 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6733 Some(Value::Int64(n)) => {
6734 count += 1;
6735 sum_i += *n as i128;
6736 mn_i = mn_i.min(*n);
6737 mx_i = mx_i.max(*n);
6738 saw_int = true;
6739 }
6740 Some(Value::Float64(f)) => {
6741 count += 1;
6742 sum_f += f;
6743 mn_f = mn_f.min(*f);
6744 mx_f = mx_f.max(*f);
6745 saw_float = true;
6746 }
6747 _ => {}
6748 },
6749 }
6750 }
6751 Ok(match agg {
6752 NativeAgg::Count => {
6753 if count == 0 {
6754 AggState::Empty
6755 } else {
6756 AggState::Count(count)
6757 }
6758 }
6759 NativeAgg::Sum => {
6760 if count == 0 {
6761 AggState::Empty
6762 } else if saw_int {
6763 AggState::SumI { sum: sum_i, count }
6764 } else {
6765 AggState::SumF { sum: sum_f, count }
6766 }
6767 }
6768 NativeAgg::Avg => {
6769 if count == 0 {
6770 AggState::Empty
6771 } else if saw_int {
6772 AggState::AvgI { sum: sum_i, count }
6773 } else {
6774 AggState::AvgF { sum: sum_f, count }
6775 }
6776 }
6777 NativeAgg::Min => {
6778 if !saw_int && !saw_float {
6779 AggState::Empty
6780 } else if saw_int {
6781 AggState::MinI(mn_i)
6782 } else {
6783 AggState::MinF(mn_f)
6784 }
6785 }
6786 NativeAgg::Max => {
6787 if !saw_int && !saw_float {
6788 AggState::Empty
6789 } else if saw_int {
6790 AggState::MaxI(mx_i)
6791 } else {
6792 AggState::MaxF(mx_f)
6793 }
6794 }
6795 })
6796}
6797
6798fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6802 use crate::query::Condition;
6803 match c {
6804 Condition::Pk(key) => match schema.primary_key() {
6805 Some(pk) => row
6806 .columns
6807 .get(&pk.id)
6808 .map(|v| v.encode_key() == *key)
6809 .unwrap_or(false),
6810 None => false,
6811 },
6812 Condition::BitmapEq { column_id, value } => row
6813 .columns
6814 .get(column_id)
6815 .map(|v| v.encode_key() == *value)
6816 .unwrap_or(false),
6817 Condition::BitmapIn { column_id, values } => {
6818 let key = row.columns.get(column_id).map(|v| v.encode_key());
6819 match key {
6820 Some(k) => values.contains(&k),
6821 None => false,
6822 }
6823 }
6824 Condition::BytesPrefix { column_id, prefix } => row
6825 .columns
6826 .get(column_id)
6827 .map(|v| v.encode_key().starts_with(prefix))
6828 .unwrap_or(false),
6829 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6830 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6831 _ => false,
6832 },
6833 Condition::RangeF64 {
6834 column_id,
6835 lo,
6836 lo_inclusive,
6837 hi,
6838 hi_inclusive,
6839 } => match row.columns.get(column_id) {
6840 Some(Value::Float64(n)) => {
6841 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6842 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6843 lo_ok && hi_ok
6844 }
6845 _ => false,
6846 },
6847 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6848 Some(Value::Bytes(b)) => {
6849 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6850 }
6851 _ => false,
6852 },
6853 Condition::FmContainsAll {
6854 column_id,
6855 patterns,
6856 } => match row.columns.get(column_id) {
6857 Some(Value::Bytes(b)) => patterns
6858 .iter()
6859 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6860 _ => false,
6861 },
6862 Condition::Ann { .. }
6863 | Condition::SparseMatch { .. }
6864 | Condition::MinHashSimilar { .. } => true,
6865 Condition::IsNull { column_id } => {
6866 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6867 }
6868 Condition::IsNotNull { column_id } => {
6869 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6870 }
6871 }
6872}
6873
6874fn as_f64(v: Option<&Value>) -> Option<f64> {
6876 match v {
6877 Some(Value::Int64(n)) => Some(*n as f64),
6878 Some(Value::Float64(f)) => Some(*f),
6879 _ => None,
6880 }
6881}
6882
6883fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6887 let mut count: u64 = 0;
6888 let mut sum: i128 = 0;
6889 let mut mn: i64 = i64::MAX;
6890 let mut mx: i64 = i64::MIN;
6891 while let Some(cols) = cursor.next_batch()? {
6892 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6893 if crate::columnar::all_non_null(validity, data.len()) {
6894 count += data.len() as u64;
6896 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6897 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6898 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6899 } else {
6900 for (i, &v) in data.iter().enumerate() {
6901 if crate::columnar::validity_bit(validity, i) {
6902 count += 1;
6903 sum += v as i128;
6904 mn = mn.min(v);
6905 mx = mx.max(v);
6906 }
6907 }
6908 }
6909 }
6910 }
6911 Ok((count, sum, mn, mx))
6912}
6913
6914fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6916 let mut count: u64 = 0;
6917 let mut sum: f64 = 0.0;
6918 let mut mn: f64 = f64::INFINITY;
6919 let mut mx: f64 = f64::NEG_INFINITY;
6920 while let Some(cols) = cursor.next_batch()? {
6921 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6922 if crate::columnar::all_non_null(validity, data.len()) {
6923 count += data.len() as u64;
6924 sum += data.iter().sum::<f64>();
6925 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6926 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6927 } else {
6928 for (i, &v) in data.iter().enumerate() {
6929 if crate::columnar::validity_bit(validity, i) {
6930 count += 1;
6931 sum += v;
6932 mn = mn.min(v);
6933 mx = mx.max(v);
6934 }
6935 }
6936 }
6937 }
6938 }
6939 Ok((count, sum, mn, mx))
6940}
6941
6942fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6943 if count == 0 && !matches!(agg, NativeAgg::Count) {
6944 return NativeAggResult::Null;
6945 }
6946 match agg {
6947 NativeAgg::Count => NativeAggResult::Count(count),
6948 NativeAgg::Sum => match sum.try_into() {
6951 Ok(v) => NativeAggResult::Int(v),
6952 Err(_) => NativeAggResult::Null,
6953 },
6954 NativeAgg::Min => NativeAggResult::Int(mn),
6955 NativeAgg::Max => NativeAggResult::Int(mx),
6956 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6957 }
6958}
6959
6960fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6961 if count == 0 && !matches!(agg, NativeAgg::Count) {
6962 return NativeAggResult::Null;
6963 }
6964 match agg {
6965 NativeAgg::Count => NativeAggResult::Count(count),
6966 NativeAgg::Sum => NativeAggResult::Float(sum),
6967 NativeAgg::Min => NativeAggResult::Float(mn),
6968 NativeAgg::Max => NativeAggResult::Float(mx),
6969 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6970 }
6971}
6972
6973fn agg_int(
6976 stats: &[crate::page::PageStat],
6977 decode: fn(Option<&[u8]>) -> Option<i64>,
6978) -> Option<(Option<i64>, Option<i64>, u64)> {
6979 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6980 let mut any = false;
6981 for s in stats {
6982 if let Some(v) = decode(s.min.as_deref()) {
6983 mn = mn.min(v);
6984 any = true;
6985 }
6986 if let Some(v) = decode(s.max.as_deref()) {
6987 mx = mx.max(v);
6988 any = true;
6989 }
6990 nulls += s.null_count;
6991 }
6992 any.then_some((Some(mn), Some(mx), nulls))
6993}
6994
6995fn agg_float(
6997 stats: &[crate::page::PageStat],
6998 decode: fn(Option<&[u8]>) -> Option<f64>,
6999) -> Option<(Option<f64>, Option<f64>, u64)> {
7000 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
7001 let mut any = false;
7002 for s in stats {
7003 if let Some(v) = decode(s.min.as_deref()) {
7004 mn = mn.min(v);
7005 any = true;
7006 }
7007 if let Some(v) = decode(s.max.as_deref()) {
7008 mx = mx.max(v);
7009 any = true;
7010 }
7011 nulls += s.null_count;
7012 }
7013 any.then_some((Some(mn), Some(mx), nulls))
7014}
7015
7016type SecondaryIndexes = (
7018 HashMap<u16, BitmapIndex>,
7019 HashMap<u16, AnnIndex>,
7020 HashMap<u16, FmIndex>,
7021 HashMap<u16, SparseIndex>,
7022 HashMap<u16, MinHashIndex>,
7023);
7024
7025fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
7026 let mut bitmap = HashMap::new();
7027 let mut ann = HashMap::new();
7028 let mut fm = HashMap::new();
7029 let mut sparse = HashMap::new();
7030 let mut minhash = HashMap::new();
7031 for idef in &schema.indexes {
7032 match idef.kind {
7033 IndexKind::Bitmap => {
7034 bitmap.insert(idef.column_id, BitmapIndex::new());
7035 }
7036 IndexKind::Ann => {
7037 let dim = schema
7038 .columns
7039 .iter()
7040 .find(|c| c.id == idef.column_id)
7041 .and_then(|c| match c.ty {
7042 TypeId::Embedding { dim } => Some(dim as usize),
7043 _ => None,
7044 })
7045 .unwrap_or(0);
7046 ann.insert(idef.column_id, AnnIndex::new(dim));
7047 }
7048 IndexKind::FmIndex => {
7049 fm.insert(idef.column_id, FmIndex::new());
7050 }
7051 IndexKind::Sparse => {
7052 sparse.insert(idef.column_id, SparseIndex::new());
7053 }
7054 IndexKind::MinHash => {
7055 minhash.insert(idef.column_id, MinHashIndex::new());
7056 }
7057 _ => {}
7058 }
7059 }
7060 (bitmap, ann, fm, sparse, minhash)
7061}
7062
7063const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
7064 | ColumnFlags::AUTO_INCREMENT
7065 | ColumnFlags::ENCRYPTED
7066 | ColumnFlags::ENCRYPTED_INDEXABLE
7067 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
7068
7069fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
7070 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
7071 return Err(MongrelError::Schema(
7072 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
7073 ));
7074 }
7075 Ok(())
7076}
7077
7078fn validate_alter_column_type(
7079 schema: &Schema,
7080 old: &ColumnDef,
7081 next: &ColumnDef,
7082 has_stored_versions: bool,
7083) -> Result<()> {
7084 if old.ty == next.ty {
7085 return Ok(());
7086 }
7087 if schema.indexes.iter().any(|i| i.column_id == old.id) {
7088 return Err(MongrelError::Schema(format!(
7089 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
7090 old.name
7091 )));
7092 }
7093 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
7094 return Ok(());
7095 }
7096 Err(MongrelError::Schema(format!(
7097 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
7098 old.ty, next.ty
7099 )))
7100}
7101
7102fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
7103 matches!(
7104 (old, new),
7105 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
7106 )
7107}
7108
7109fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
7115 let mut prev: Option<i64> = None;
7116 for r in rows {
7117 match r.columns.get(&pk_id) {
7118 Some(Value::Int64(v)) => {
7119 if prev.is_some_and(|p| p >= *v) {
7120 return false;
7121 }
7122 prev = Some(*v);
7123 }
7124 _ => return false,
7125 }
7126 }
7127 true
7128}
7129
7130#[allow(clippy::too_many_arguments)]
7131fn index_into(
7132 schema: &Schema,
7133 row: &Row,
7134 hot: &mut HotIndex,
7135 bitmap: &mut HashMap<u16, BitmapIndex>,
7136 ann: &mut HashMap<u16, AnnIndex>,
7137 fm: &mut HashMap<u16, FmIndex>,
7138 sparse: &mut HashMap<u16, SparseIndex>,
7139 minhash: &mut HashMap<u16, MinHashIndex>,
7140) {
7141 for idef in &schema.indexes {
7142 let Some(val) = row.columns.get(&idef.column_id) else {
7143 continue;
7144 };
7145 match idef.kind {
7146 IndexKind::Bitmap => {
7147 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7148 b.insert(val.encode_key(), row.row_id);
7149 }
7150 }
7151 IndexKind::Ann => {
7152 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7153 a.insert(v, row.row_id);
7154 }
7155 }
7156 IndexKind::FmIndex => {
7157 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7158 f.insert(b.clone(), row.row_id);
7159 }
7160 }
7161 IndexKind::Sparse => {
7162 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7163 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7166 s.insert(&terms, row.row_id);
7167 }
7168 }
7169 }
7170 IndexKind::MinHash => {
7171 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7172 let tokens = crate::index::token_hashes_from_bytes(b);
7175 mh.insert(&tokens, row.row_id);
7176 }
7177 }
7178 _ => {}
7179 }
7180 }
7181 if let Some(pk_col) = schema.primary_key() {
7182 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7183 hot.insert(pk_val.encode_key(), row.row_id);
7184 }
7185 }
7186}
7187
7188#[allow(clippy::too_many_arguments)]
7191fn index_into_single(
7192 idef: &IndexDef,
7193 _schema: &Schema,
7194 row: &Row,
7195 _hot: &mut HotIndex,
7196 bitmap: &mut HashMap<u16, BitmapIndex>,
7197 ann: &mut HashMap<u16, AnnIndex>,
7198 fm: &mut HashMap<u16, FmIndex>,
7199 sparse: &mut HashMap<u16, SparseIndex>,
7200 minhash: &mut HashMap<u16, MinHashIndex>,
7201) {
7202 let Some(val) = row.columns.get(&idef.column_id) else {
7203 return;
7204 };
7205 match idef.kind {
7206 IndexKind::Bitmap => {
7207 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7208 b.insert(val.encode_key(), row.row_id);
7209 }
7210 }
7211 IndexKind::Ann => {
7212 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7213 a.insert(v, row.row_id);
7214 }
7215 }
7216 IndexKind::FmIndex => {
7217 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7218 f.insert(b.clone(), row.row_id);
7219 }
7220 }
7221 IndexKind::Sparse => {
7222 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7223 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7224 s.insert(&terms, row.row_id);
7225 }
7226 }
7227 }
7228 IndexKind::MinHash => {
7229 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7230 let tokens = crate::index::token_hashes_from_bytes(b);
7231 mh.insert(&tokens, row.row_id);
7232 }
7233 }
7234 _ => {}
7235 }
7236}
7237
7238fn eval_partial_predicate(
7244 pred: &str,
7245 columns_map: &HashMap<u16, &Value>,
7246 name_to_id: &HashMap<&str, u16>,
7247) -> bool {
7248 let lower = pred.trim().to_ascii_lowercase();
7249 if let Some(rest) = lower.strip_suffix(" is not null") {
7251 let col_name = rest.trim();
7252 if let Some(col_id) = name_to_id.get(col_name) {
7253 return columns_map
7254 .get(col_id)
7255 .is_some_and(|v| !matches!(v, Value::Null));
7256 }
7257 }
7258 if let Some(rest) = lower.strip_suffix(" is null") {
7260 let col_name = rest.trim();
7261 if let Some(col_id) = name_to_id.get(col_name) {
7262 return columns_map
7263 .get(col_id)
7264 .map_or(true, |v| matches!(v, Value::Null));
7265 }
7266 }
7267 true
7270}
7271
7272#[allow(dead_code)]
7278fn bulk_index_key(
7279 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7280 column_id: u16,
7281 ty: TypeId,
7282 col: &columnar::NativeColumn,
7283 i: usize,
7284) -> Option<Vec<u8>> {
7285 let encoded = columnar::encode_key_native(ty, col, i)?;
7286 #[cfg(feature = "encryption")]
7287 {
7288 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7289 if let Some((key, scheme)) = column_keys.get(&column_id) {
7290 return Some(match (*scheme, col) {
7291 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7292 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7293 ope_token_i64(key, data[i]).to_vec()
7294 }
7295 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7296 ope_token_f64(key, data[i]).to_vec()
7297 }
7298 _ => hmac_token(key, &encoded).to_vec(),
7299 });
7300 }
7301 }
7302 #[cfg(not(feature = "encryption"))]
7303 {
7304 let _ = (column_id, column_keys, col);
7305 }
7306 Some(encoded)
7307}
7308
7309pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7310 let json = serde_json::to_string_pretty(schema)
7311 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7312 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7313 Ok(())
7314}
7315
7316fn read_schema(dir: &Path) -> Result<Schema> {
7317 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7318 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7319}
7320
7321fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7322 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7323}
7324
7325fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7326 let n = list_wal_numbers(wal_dir)?;
7327 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7328}
7329
7330fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7331 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7332}
7333
7334fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7335 let _ = std::fs::create_dir_all(wal_dir);
7336 let mut max_n = None;
7337 for entry in std::fs::read_dir(wal_dir)? {
7338 let entry = entry?;
7339 let fname = entry.file_name();
7340 let Some(s) = fname.to_str() else {
7341 continue;
7342 };
7343 let Some(stripped) = s.strip_prefix("seg-") else {
7344 continue;
7345 };
7346 let Some(stripped) = stripped.strip_suffix(".wal") else {
7347 continue;
7348 };
7349 if let Ok(n) = stripped.parse::<u32>() {
7350 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7351 }
7352 }
7353 Ok(max_n)
7354}