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 wal: WalSink,
109 memtable: Memtable,
110 mutable_run: MutableRun,
115 mutable_run_spill_bytes: u64,
117 compaction_zstd_level: i32,
120 allocator: RowIdAllocator,
121 epoch: Arc<EpochAuthority>,
122 persisted_epoch: u64,
125 schema: Schema,
126 hot: HotIndex,
127 kek: Option<Arc<Kek>>,
130 column_keys: HashMap<u16, ([u8; 32], u8)>,
134 run_refs: Vec<RunRef>,
135 retiring: Vec<crate::manifest::RetiredRun>,
138 next_run_id: u64,
139 sync_byte_threshold: u64,
140 current_txn_id: u64,
145 bitmap: HashMap<u16, BitmapIndex>,
146 ann: HashMap<u16, AnnIndex>,
147 fm: HashMap<u16, FmIndex>,
148 sparse: HashMap<u16, SparseIndex>,
149 minhash: HashMap<u16, MinHashIndex>,
150 learned_range: HashMap<u16, ColumnLearnedRange>,
153 pk_by_row: HashMap<RowId, Vec<u8>>,
155 pinned: BTreeMap<Epoch, usize>,
158 pub(crate) live_count: u64,
161 reservoir: crate::reservoir::Reservoir,
164 reservoir_complete: bool,
172 had_deletes: bool,
176 agg_cache: HashMap<u64, CachedAgg>,
180 global_idx_epoch: u64,
184 indexes_complete: bool,
189 index_build_policy: IndexBuildPolicy,
191 pk_by_row_complete: bool,
198 flushed_epoch: u64,
201 page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
204 snapshots: Arc<crate::retention::SnapshotRegistry>,
207 commit_lock: Arc<parking_lot::Mutex<()>>,
209 decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
212 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
222 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
231 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
233 pending_delete_rids: roaring::RoaringBitmap,
236 pending_put_cols: std::collections::HashSet<u16>,
239 pending_rows: Vec<Row>,
245 pending_rows_auto_inc: Vec<bool>,
246 pending_dels: Vec<RowId>,
249 pending_truncate: Option<Epoch>,
253 auto_inc: Option<AutoIncState>,
256}
257
258const _: () = {
265 const fn assert_sync<T: ?Sized + Sync>() {}
266 assert_sync::<Table>();
267};
268
269enum CachedData {
275 Rows(Arc<Vec<Row>>),
276 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
277}
278
279impl CachedData {
280 fn approx_bytes(&self) -> u64 {
281 match self {
282 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
283 CachedData::Columns(c) => c
284 .iter()
285 .map(|(_, c)| c.approx_bytes())
286 .sum::<u64>()
287 .saturating_add(c.len() as u64 * 16),
288 }
289 }
290}
291
292struct CachedEntry {
296 data: CachedData,
297 footprint: roaring::RoaringBitmap,
298 condition_cols: Vec<u16>,
299}
300
301struct ResultCache {
312 entries: std::collections::HashMap<u64, CachedEntry>,
313 order: std::collections::VecDeque<u64>,
314 bytes: u64,
315 max_bytes: u64,
316 dir: Option<std::path::PathBuf>,
317 #[allow(dead_code)]
318 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
319}
320
321#[derive(serde::Serialize, serde::Deserialize)]
323struct SerializedEntry {
324 condition_cols: Vec<u16>,
325 footprint_bits: Vec<u32>,
326 data: SerializedData,
327}
328
329#[derive(serde::Serialize, serde::Deserialize)]
330enum SerializedData {
331 Rows(Vec<Row>),
332 Columns(Vec<(u16, columnar::NativeColumn)>),
333}
334
335impl SerializedEntry {
336 fn from_entry(entry: &CachedEntry) -> Self {
337 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
338 let data = match &entry.data {
339 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
340 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
341 };
342 Self {
343 condition_cols: entry.condition_cols.clone(),
344 footprint_bits,
345 data,
346 }
347 }
348
349 fn into_entry(self) -> Option<CachedEntry> {
350 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
351 let data = match self.data {
352 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
353 SerializedData::Columns(c) => {
354 if !c.iter().all(|(_, col)| col.validate()) {
357 return None;
358 }
359 CachedData::Columns(Arc::new(c))
360 }
361 };
362 Some(CachedEntry {
363 data,
364 footprint,
365 condition_cols: self.condition_cols,
366 })
367 }
368}
369
370impl ResultCache {
371 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
372
373 fn new() -> Self {
374 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
375 }
376
377 fn with_max_bytes(max_bytes: u64) -> Self {
378 Self {
379 entries: std::collections::HashMap::new(),
380 order: std::collections::VecDeque::new(),
381 bytes: 0,
382 max_bytes,
383 dir: None,
384 cache_dek: None,
385 }
386 }
387
388 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
389 let _ = std::fs::create_dir_all(&dir);
390 self.dir = Some(dir);
391 self
392 }
393
394 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
395 self.cache_dek = dek;
396 self
397 }
398
399 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
400 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
401 }
402
403 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
407 let Some(path) = self.disk_path(key) else {
408 return;
409 };
410 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
411 Ok(s) => s,
412 Err(_) => return,
413 };
414 let on_disk = if let Some(dek) = &self.cache_dek {
416 match self.encrypt_cache(&serialized, dek) {
417 Some(b) => b,
418 None => return,
419 }
420 } else {
421 serialized
422 };
423 let tmp = path.with_extension("tmp");
424 use std::io::Write;
425 let write = || -> std::io::Result<()> {
426 let mut f = std::fs::File::create(&tmp)?;
427 f.write_all(&on_disk)?;
428 f.flush()?;
429 Ok(())
430 };
431 if write().is_err() {
432 let _ = std::fs::remove_file(&tmp);
433 return;
434 }
435 let _ = std::fs::rename(&tmp, &path);
436 }
437
438 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
440 let path = self.disk_path(key)?;
441 let bytes = std::fs::read(&path).ok()?;
442 let plaintext = if let Some(dek) = &self.cache_dek {
443 self.decrypt_cache(&bytes, dek)?
444 } else {
445 bytes
446 };
447 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
448 serialized.into_entry()
449 }
450
451 fn remove_from_disk(&self, key: u64) {
453 if let Some(path) = self.disk_path(key) {
454 let _ = std::fs::remove_file(&path);
455 }
456 }
457
458 #[cfg(feature = "encryption")]
460 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
461 use crate::encryption::Cipher;
462 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
463 let mut nonce = [0u8; 12];
464 crate::encryption::fill_random(&mut nonce);
465 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
466 let mut out = Vec::with_capacity(12 + ct.len());
467 out.extend_from_slice(&nonce);
468 out.extend_from_slice(&ct);
469 Some(out)
470 }
471
472 #[cfg(not(feature = "encryption"))]
473 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
474 None
475 }
476
477 #[cfg(feature = "encryption")]
479 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
480 use crate::encryption::Cipher;
481 if bytes.len() < 28 {
482 return None;
483 }
484 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
485 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
486 let ct = &bytes[12..];
487 cipher.decrypt_page(&nonce, ct).ok()
488 }
489
490 #[cfg(not(feature = "encryption"))]
491 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
492 None
493 }
494
495 fn load_persistent(&mut self) {
498 let Some(dir) = self.dir.as_ref().cloned() else {
499 return;
500 };
501 let entries = match std::fs::read_dir(&dir) {
502 Ok(e) => e,
503 Err(_) => return,
504 };
505 for entry in entries.flatten() {
506 let path = entry.path();
507 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
509 let _ = std::fs::remove_file(&path);
510 continue;
511 }
512 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
513 continue;
514 }
515 let stem = match path.file_stem().and_then(|s| s.to_str()) {
516 Some(s) => s,
517 None => continue,
518 };
519 let key = match u64::from_str_radix(stem, 16) {
520 Ok(k) => k,
521 Err(_) => continue,
522 };
523 let bytes = match std::fs::read(&path) {
524 Ok(b) => b,
525 Err(_) => continue,
526 };
527 let plaintext = if let Some(dek) = &self.cache_dek {
529 match self.decrypt_cache(&bytes, dek) {
530 Some(p) => p,
531 None => {
532 let _ = std::fs::remove_file(&path);
533 continue;
534 }
535 }
536 } else {
537 bytes
538 };
539 match bincode::deserialize::<SerializedEntry>(&plaintext) {
540 Ok(serialized) => {
541 if let Some(entry) = serialized.into_entry() {
542 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
543 self.entries.insert(key, entry);
544 self.order.push_back(key);
545 } else {
546 let _ = std::fs::remove_file(&path);
547 }
548 }
549 Err(_) => {
550 let _ = std::fs::remove_file(&path);
551 }
552 }
553 }
554 self.evict();
555 }
556
557 fn set_max_bytes(&mut self, max_bytes: u64) {
558 self.max_bytes = max_bytes;
559 self.evict();
560 }
561
562 fn touch(&mut self, key: u64) {
564 self.order.retain(|k| *k != key);
565 self.order.push_back(key);
566 }
567
568 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
569 let res = self.entries.get(&key).and_then(|e| match &e.data {
570 CachedData::Rows(r) => Some(r.clone()),
571 CachedData::Columns(_) => None,
572 });
573 if res.is_some() {
574 self.touch(key);
575 return res;
576 }
577 if let Some(entry) = self.load_from_disk(key) {
579 let res = match &entry.data {
580 CachedData::Rows(r) => Some(r.clone()),
581 CachedData::Columns(_) => None,
582 };
583 if res.is_some() {
584 let approx = entry.data.approx_bytes();
585 self.bytes = self.bytes.saturating_add(approx);
586 self.entries.insert(key, entry);
587 self.order.push_back(key);
588 self.evict();
589 return res;
590 }
591 }
592 None
593 }
594
595 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
596 let res = self.entries.get(&key).and_then(|e| match &e.data {
597 CachedData::Columns(c) => Some(c.clone()),
598 CachedData::Rows(_) => None,
599 });
600 if res.is_some() {
601 self.touch(key);
602 return res;
603 }
604 if let Some(entry) = self.load_from_disk(key) {
606 let res = match &entry.data {
607 CachedData::Columns(c) => Some(c.clone()),
608 CachedData::Rows(_) => None,
609 };
610 if res.is_some() {
611 let approx = entry.data.approx_bytes();
612 self.bytes = self.bytes.saturating_add(approx);
613 self.entries.insert(key, entry);
614 self.order.push_back(key);
615 self.evict();
616 return res;
617 }
618 }
619 None
620 }
621
622 fn insert(&mut self, key: u64, entry: CachedEntry) {
623 let approx = entry.data.approx_bytes();
624 if self.entries.remove(&key).is_some() {
625 self.order.retain(|k| *k != key);
626 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
627 }
628 self.store_to_disk(key, &entry);
630 self.bytes = self.bytes.saturating_add(approx);
631 self.entries.insert(key, entry);
632 self.order.push_back(key);
633 self.evict();
634 }
635
636 fn invalidate(
645 &mut self,
646 delete_rids: &roaring::RoaringBitmap,
647 put_cols: &std::collections::HashSet<u16>,
648 ) {
649 if self.entries.is_empty() {
650 return;
651 }
652 let has_deletes = !delete_rids.is_empty();
653 let to_remove: std::collections::HashSet<u64> = self
654 .entries
655 .iter()
656 .filter(|(_, e)| {
657 let delete_hit = if e.footprint.is_empty() {
658 has_deletes
659 } else {
660 e.footprint.intersection_len(delete_rids) > 0
661 };
662 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
663 delete_hit || col_hit
664 })
665 .map(|(&k, _)| k)
666 .collect();
667 for key in &to_remove {
668 if let Some(e) = self.entries.remove(key) {
669 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
670 }
671 self.remove_from_disk(*key);
672 }
673 if !to_remove.is_empty() {
674 self.order.retain(|k| !to_remove.contains(k));
675 }
676 }
677
678 fn clear(&mut self) {
679 if let Some(dir) = &self.dir {
681 if let Ok(entries) = std::fs::read_dir(dir) {
682 for entry in entries.flatten() {
683 let path = entry.path();
684 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
685 let _ = std::fs::remove_file(&path);
686 }
687 }
688 }
689 }
690 self.entries.clear();
691 self.order.clear();
692 self.bytes = 0;
693 }
694
695 fn evict(&mut self) {
696 while self.bytes > self.max_bytes {
697 let Some(k) = self.order.pop_front() else {
698 break;
699 };
700 if let Some(e) = self.entries.remove(&k) {
701 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
702 self.remove_from_disk(k);
706 }
707 }
708 }
709}
710
711type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
718
719fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
720 let _ = kek;
721 #[cfg(feature = "encryption")]
722 {
723 if let Some(k) = kek {
724 return (
725 Some(k.derive_table_wal_key(_table_id)),
726 Some(k.derive_cache_key()),
727 );
728 }
729 }
730 (None, None)
731}
732
733#[cfg(feature = "encryption")]
735fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
736 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
737}
738
739#[cfg(not(feature = "encryption"))]
740fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
741 Box::new(crate::encryption::PlaintextCipher)
742}
743
744fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
745 let Some(kek) = kek else {
746 return HashMap::new();
747 };
748 #[cfg(feature = "encryption")]
749 {
750 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
751 schema
752 .columns
753 .iter()
754 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
755 .map(|c| {
756 let scheme = if schema
757 .indexes
758 .iter()
759 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
760 {
761 SCHEME_OPE_RANGE
762 } else {
763 SCHEME_HMAC_EQ
764 };
765 let key: [u8; 32] = *kek.derive_column_key(c.id);
766 (c.id, (key, scheme))
767 })
768 .collect()
769 }
770 #[cfg(not(feature = "encryption"))]
771 {
772 let _ = (kek, schema);
773 HashMap::new()
774 }
775}
776
777pub(crate) struct SharedCtx {
782 pub epoch: Arc<EpochAuthority>,
783 pub page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
784 pub decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
785 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
786 pub kek: Option<Arc<Kek>>,
787 pub commit_lock: Arc<parking_lot::Mutex<()>>,
793 pub shared: Option<SharedWalCtx>,
797}
798
799#[derive(Clone)]
805pub(crate) struct SharedWalCtx {
806 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
807 pub group: Arc<GroupCommit>,
808 pub poisoned: Arc<AtomicBool>,
809 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
810}
811
812enum WalSink {
815 Private(Wal),
816 Shared(SharedWalCtx),
817}
818
819impl SharedCtx {
820 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
824 let mut cache = crate::cache::PageCache::new(PAGE_CACHE_CAPACITY);
825 if let Some(d) = cache_dir {
826 cache = cache.with_persistence(d);
827 }
828 Self {
829 epoch: Arc::new(EpochAuthority::new(0)),
830 page_cache: Arc::new(parking_lot::Mutex::new(cache)),
831 decoded_cache: Arc::new(parking_lot::Mutex::new(
832 crate::cache::DecodedPageCache::new(DECODED_CACHE_CAPACITY),
833 )),
834 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
835 kek,
836 commit_lock: Arc::new(parking_lot::Mutex::new(())),
837 shared: None,
838 }
839 }
840}
841
842impl Table {
843 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
844 let dir = dir.as_ref().to_path_buf();
845 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
846 Self::create_in(&dir, schema, table_id, ctx)
847 }
848
849 #[cfg(feature = "encryption")]
860 pub fn create_encrypted(
861 dir: impl AsRef<Path>,
862 schema: Schema,
863 table_id: u64,
864 passphrase: &str,
865 ) -> Result<Self> {
866 let dir = dir.as_ref();
867 std::fs::create_dir_all(dir.join(META_DIR))?;
868 let salt = crate::encryption::random_salt();
869 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
870 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
871 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
872 Self::create_in(dir, schema, table_id, ctx)
873 }
874
875 #[cfg(feature = "encryption")]
880 pub fn create_with_key(
881 dir: impl AsRef<Path>,
882 schema: Schema,
883 table_id: u64,
884 key: &[u8],
885 ) -> Result<Self> {
886 let dir = dir.as_ref();
887 std::fs::create_dir_all(dir.join(META_DIR))?;
888 let salt = crate::encryption::random_salt();
889 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
890 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
891 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
892 Self::create_in(dir, schema, table_id, ctx)
893 }
894
895 #[cfg(feature = "encryption")]
897 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
898 let dir = dir.as_ref();
899 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
900 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
901 MongrelError::NotFound(format!(
902 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
903 salt_path
904 ))
905 })?;
906 if salt_bytes.len() != crate::encryption::SALT_LEN {
907 return Err(MongrelError::InvalidArgument(format!(
908 "salt file is {} bytes, expected {}",
909 salt_bytes.len(),
910 crate::encryption::SALT_LEN
911 )));
912 }
913 let mut salt = [0u8; crate::encryption::SALT_LEN];
914 salt.copy_from_slice(&salt_bytes);
915 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
916 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
917 Self::open_in(dir, ctx)
918 }
919
920 pub(crate) fn create_in(
921 dir: impl AsRef<Path>,
922 schema: Schema,
923 table_id: u64,
924 ctx: SharedCtx,
925 ) -> Result<Self> {
926 schema.validate_auto_increment()?;
927 let dir = dir.as_ref().to_path_buf();
928 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
929 write_schema(&dir, &schema)?;
930 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
931 let (wal, current_txn_id) = match ctx.shared.clone() {
934 Some(s) => (WalSink::Shared(s), 0),
935 None => {
936 std::fs::create_dir_all(dir.join(WAL_DIR))?;
937 let mut w = if let Some(ref dk) = wal_dek {
938 Wal::create_with_cipher(
939 dir.join(WAL_DIR).join("seg-000000.wal"),
940 Epoch(0),
941 Some(make_cipher(dk)),
942 0,
943 )?
944 } else {
945 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
946 };
947 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
948 (WalSink::Private(w), 1)
949 }
950 };
951 let mut manifest = Manifest::new(table_id, schema.schema_id);
952 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
957 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
958 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
959 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
960 let auto_inc = resolve_auto_inc(&schema);
961 let rcache_dir = dir.join(RCACHE_DIR);
962 Ok(Self {
963 dir,
964 table_id,
965 wal,
966 memtable: Memtable::new(),
967 mutable_run: MutableRun::new(),
968 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
969 compaction_zstd_level: 3,
970 allocator: RowIdAllocator::new(0),
971 epoch: ctx.epoch,
972 persisted_epoch: 0,
973 schema,
974 hot: HotIndex::new(),
975 kek: ctx.kek,
976 column_keys,
977 run_refs: Vec::new(),
978 retiring: Vec::new(),
979 next_run_id: 1,
980 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
981 current_txn_id,
982 bitmap,
983 ann,
984 fm,
985 sparse,
986 minhash,
987 learned_range: HashMap::new(),
988 pk_by_row: HashMap::new(),
989 pinned: BTreeMap::new(),
990 live_count: 0,
991 reservoir: crate::reservoir::Reservoir::default(),
992 reservoir_complete: true,
993 had_deletes: false,
994 agg_cache: HashMap::new(),
995 global_idx_epoch: 0,
996 indexes_complete: true,
997 index_build_policy: IndexBuildPolicy::default(),
998 pk_by_row_complete: false,
999 flushed_epoch: 0,
1000 page_cache: ctx.page_cache,
1001 decoded_cache: ctx.decoded_cache,
1002 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1003 snapshots: ctx.snapshots,
1004 commit_lock: ctx.commit_lock,
1005 result_cache: Arc::new(parking_lot::Mutex::new(
1006 ResultCache::new()
1007 .with_dir(rcache_dir)
1008 .with_cache_dek(cache_dek.clone()),
1009 )),
1010 pending_delete_rids: roaring::RoaringBitmap::new(),
1011 pending_put_cols: std::collections::HashSet::new(),
1012 pending_rows: Vec::new(),
1013 pending_rows_auto_inc: Vec::new(),
1014 pending_dels: Vec::new(),
1015 pending_truncate: None,
1016 wal_dek,
1017 auto_inc,
1018 })
1019 }
1020
1021 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1025 let dir = dir.as_ref();
1026 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1027 Self::open_in(dir, ctx)
1028 }
1029
1030 #[cfg(feature = "encryption")]
1033 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1034 let dir = dir.as_ref();
1035 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1036 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1037 MongrelError::NotFound(format!(
1038 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1039 salt_path
1040 ))
1041 })?;
1042 let salt_len = crate::encryption::SALT_LEN;
1043 if salt_bytes.len() != salt_len {
1044 return Err(MongrelError::InvalidArgument(format!(
1045 "encryption salt is {} bytes, expected {salt_len}",
1046 salt_bytes.len()
1047 )));
1048 }
1049 let mut salt = [0u8; 16];
1050 salt.copy_from_slice(&salt_bytes);
1051 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1052 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1053 let t = Self::open_in(dir, ctx)?;
1054 Ok(t)
1055 }
1056
1057 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1058 let dir = dir.as_ref().to_path_buf();
1059 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1060 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1061 let schema: Schema = read_schema(&dir)?;
1062 let replay_epoch = Epoch(manifest.current_epoch);
1063 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1064 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1068 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1069 None => {
1070 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1071 let replayed = match &active {
1073 Some(path) => {
1074 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1075 crate::wal::replay_with_cipher(path, cipher)?
1076 }
1077 None => Vec::new(),
1078 };
1079 let mut w = match &active {
1080 Some(path) => Wal::create_with_cipher(
1081 path,
1082 replay_epoch,
1083 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1084 0,
1085 )?,
1086 None => Wal::create_with_cipher(
1087 dir.join(WAL_DIR).join("seg-000000.wal"),
1088 replay_epoch,
1089 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1090 0,
1091 )?,
1092 };
1093 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1094 (WalSink::Private(w), replayed, 1)
1095 }
1096 };
1097
1098 let mut memtable = Memtable::new();
1099 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1100 let persisted_epoch = manifest.current_epoch;
1101 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1108 s.next = manifest.auto_inc_next;
1109 s.seeded = manifest.auto_inc_next > 0;
1110 s
1111 });
1112
1113 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1120 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1121 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1122 std::collections::BTreeMap::new();
1123 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1124 let mut saw_delete = false;
1125 for record in replayed {
1126 let txn_id = record.txn_id;
1127 match record.op {
1128 Op::Put { rows, .. } => {
1129 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1130 for row in &rows {
1131 allocator.advance_to(row.row_id);
1132 if let Some(ai) = auto_inc.as_mut() {
1133 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1134 if *n + 1 > ai.next {
1135 ai.next = *n + 1;
1136 }
1137 }
1138 }
1139 }
1140 staged_puts.entry(txn_id).or_default().extend(rows);
1141 }
1142 Op::Delete { row_ids, .. } => {
1143 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1144 }
1145 Op::TxnCommit { epoch, .. } => {
1146 let commit_epoch = Epoch(epoch);
1147 if let Some(puts) = staged_puts.remove(&txn_id) {
1148 for row in &puts {
1149 memtable.upsert(row.clone());
1150 }
1151 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1152 }
1153 if let Some(dels) = staged_deletes.remove(&txn_id) {
1154 saw_delete = true;
1155 for rid in dels {
1156 memtable.tombstone(rid, commit_epoch);
1157 replayed_deletes.push((rid, commit_epoch));
1158 }
1159 }
1160 }
1161 Op::TxnAbort => {
1162 staged_puts.remove(&txn_id);
1163 staged_deletes.remove(&txn_id);
1164 }
1165 Op::TruncateTable { .. } | Op::Flush { .. } | Op::Ddl(_) => {}
1166 }
1167 }
1168
1169 let rcache_dir = dir.join(RCACHE_DIR);
1170 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1171 let mut db = Self {
1172 dir,
1173 table_id: manifest.table_id,
1174 wal,
1175 memtable,
1176 mutable_run: MutableRun::new(),
1177 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1178 compaction_zstd_level: 3,
1179 allocator,
1180 epoch: ctx.epoch,
1181 persisted_epoch,
1182 schema,
1183 hot: HotIndex::new(),
1184 kek: ctx.kek,
1185 column_keys,
1186 run_refs: manifest.runs.clone(),
1187 retiring: manifest.retiring.clone(),
1188 next_run_id: manifest
1189 .runs
1190 .iter()
1191 .map(|r| r.run_id as u64 + 1)
1192 .max()
1193 .unwrap_or(1),
1194 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1195 current_txn_id,
1196 bitmap: HashMap::new(),
1197 ann: HashMap::new(),
1198 fm: HashMap::new(),
1199 sparse: HashMap::new(),
1200 minhash: HashMap::new(),
1201 learned_range: HashMap::new(),
1202 pk_by_row: HashMap::new(),
1203 pinned: BTreeMap::new(),
1204 live_count: manifest.live_count,
1205 reservoir: crate::reservoir::Reservoir::default(),
1206 reservoir_complete: false,
1207 had_deletes: saw_delete,
1208 agg_cache: HashMap::new(),
1209 global_idx_epoch: manifest.global_idx_epoch,
1210 indexes_complete: true,
1211 index_build_policy: IndexBuildPolicy::default(),
1212 pk_by_row_complete: false,
1213 flushed_epoch: manifest.flushed_epoch,
1214 page_cache: ctx.page_cache,
1215 decoded_cache: ctx.decoded_cache,
1216 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1217 snapshots: ctx.snapshots,
1218 commit_lock: ctx.commit_lock,
1219 result_cache: Arc::new(parking_lot::Mutex::new(
1220 ResultCache::new()
1221 .with_dir(rcache_dir)
1222 .with_cache_dek(cache_dek.clone()),
1223 )),
1224 pending_delete_rids: roaring::RoaringBitmap::new(),
1225 pending_put_cols: std::collections::HashSet::new(),
1226 pending_rows: Vec::new(),
1227 pending_rows_auto_inc: Vec::new(),
1228 pending_dels: Vec::new(),
1229 pending_truncate: None,
1230 wal_dek,
1231 auto_inc,
1232 };
1233
1234 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1237
1238 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1243 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1244 c.epoch_built == manifest.global_idx_epoch
1245 && manifest.global_idx_epoch > 0
1246 && manifest
1247 .runs
1248 .iter()
1249 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1250 });
1251 if let Some(loaded) = checkpoint {
1252 if checkpoint_valid {
1253 db.hot = loaded.hot;
1254 db.bitmap = loaded.bitmap;
1255 db.ann = loaded.ann;
1256 db.fm = loaded.fm;
1257 db.sparse = loaded.sparse;
1258 db.minhash = loaded.minhash;
1259 db.learned_range = loaded.learned_range;
1260 }
1263 }
1264 if !checkpoint_valid {
1265 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1266 db.bitmap = bitmap;
1267 db.ann = ann;
1268 db.fm = fm;
1269 db.sparse = sparse;
1270 db.minhash = minhash;
1271 db.rebuild_indexes_from_runs()?;
1272 db.build_learned_ranges()?;
1273 }
1274
1275 for (epoch, group) in replayed_puts {
1280 let (losers, winner_pks) = db.partition_pk_winners(&group);
1281 for (key, &row_id) in &winner_pks {
1282 if let Some(old_rid) = db.hot.get(key) {
1283 if old_rid != row_id {
1284 db.tombstone_row(old_rid, epoch, false);
1285 }
1286 }
1287 }
1288 for &loser_rid in &losers {
1289 db.tombstone_row(loser_rid, epoch, false);
1290 }
1291 for (key, row_id) in winner_pks {
1292 db.insert_hot_pk(key, row_id);
1293 }
1294 if db.schema.primary_key().is_none() {
1295 for r in &group {
1296 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1297 }
1298 }
1299 for r in &group {
1300 if !losers.contains(&r.row_id) {
1301 db.index_row(r);
1302 }
1303 }
1304 }
1305 for (rid, epoch) in &replayed_deletes {
1309 db.remove_hot_for_row(*rid, *epoch);
1310 }
1311
1312 db.result_cache.lock().load_persistent();
1319 Ok(db)
1320 }
1321
1322 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1328 if self.reservoir_complete {
1329 return Ok(());
1330 }
1331 self.rebuild_reservoir()?;
1332 self.reservoir_complete = true;
1333 Ok(())
1334 }
1335
1336 fn rebuild_reservoir(&mut self) -> Result<()> {
1339 let snap = self.snapshot();
1340 let rows = self.visible_rows(snap)?;
1341 self.reservoir.reset();
1342 for r in rows {
1343 self.reservoir.offer(r.row_id.0);
1344 }
1345 Ok(())
1346 }
1347
1348 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1349 self.hot = HotIndex::new();
1350 self.pk_by_row.clear();
1351 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1352 self.bitmap = bitmap;
1353 self.ann = ann;
1354 self.fm = fm;
1355 self.sparse = sparse;
1356 self.minhash = minhash;
1357 let snapshot = Epoch(u64::MAX);
1358 for rr in self.run_refs.clone() {
1359 let mut reader = self.open_reader(rr.run_id)?;
1360 for row in reader.visible_rows(snapshot)? {
1361 let tok_row = self.tokenized_for_indexes(&row);
1362 index_into(
1363 &self.schema,
1364 &tok_row,
1365 &mut self.hot,
1366 &mut self.bitmap,
1367 &mut self.ann,
1368 &mut self.fm,
1369 &mut self.sparse,
1370 &mut self.minhash,
1371 );
1372 }
1373 }
1374 for row in self.mutable_run.visible_versions(snapshot) {
1375 if row.deleted {
1376 self.remove_hot_for_row(row.row_id, snapshot);
1377 } else {
1378 self.index_row(&row);
1379 }
1380 }
1381 for row in self.memtable.visible_versions(snapshot) {
1382 if row.deleted {
1383 self.remove_hot_for_row(row.row_id, snapshot);
1384 } else {
1385 self.index_row(&row);
1386 }
1387 }
1388 self.refresh_pk_by_row_from_hot();
1389 Ok(())
1390 }
1391
1392 fn refresh_pk_by_row_from_hot(&mut self) {
1393 self.pk_by_row_complete = true;
1394 if self.schema.primary_key().is_none() {
1395 self.pk_by_row.clear();
1396 return;
1397 }
1398 self.pk_by_row = self
1404 .hot
1405 .entries()
1406 .into_iter()
1407 .map(|(key, row_id)| (row_id, key))
1408 .collect();
1409 }
1410
1411 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1412 if self.schema.primary_key().is_some() {
1413 self.pk_by_row.insert(row_id, key.clone());
1414 }
1415 self.hot.insert(key, row_id);
1416 }
1417
1418 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1422 self.learned_range.clear();
1423 if self.run_refs.len() != 1 {
1424 return Ok(());
1425 }
1426 let cols: Vec<u16> = self
1427 .schema
1428 .indexes
1429 .iter()
1430 .filter(|i| i.kind == IndexKind::LearnedRange)
1431 .map(|i| i.column_id)
1432 .collect();
1433 if cols.is_empty() {
1434 return Ok(());
1435 }
1436 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1437 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1438 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1439 _ => return Ok(()),
1440 };
1441 for cid in cols {
1442 let ty = self
1443 .schema
1444 .columns
1445 .iter()
1446 .find(|c| c.id == cid)
1447 .map(|c| c.ty)
1448 .unwrap_or(TypeId::Int64);
1449 match ty {
1450 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1451 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1452 let pairs: Vec<(i64, u64)> = data
1453 .iter()
1454 .zip(row_ids.iter())
1455 .map(|(v, r)| (*v, *r))
1456 .collect();
1457 self.learned_range
1458 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1459 }
1460 }
1461 TypeId::Float64 => {
1462 if let columnar::NativeColumn::Float64 { data, .. } =
1463 reader.column_native(cid)?
1464 {
1465 let pairs: Vec<(f64, u64)> = data
1466 .iter()
1467 .zip(row_ids.iter())
1468 .map(|(v, r)| (*v, *r))
1469 .collect();
1470 self.learned_range
1471 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1472 }
1473 }
1474 _ => {}
1475 }
1476 }
1477 Ok(())
1478 }
1479
1480 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1487 if self.indexes_complete {
1488 crate::trace::QueryTrace::record(|t| {
1489 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1490 });
1491 return Ok(());
1492 }
1493 crate::trace::QueryTrace::record(|t| {
1494 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1495 });
1496 self.rebuild_indexes_from_runs()?;
1497 self.build_learned_ranges()?;
1498 self.indexes_complete = true;
1499 let epoch = self.current_epoch();
1500 self.checkpoint_indexes(epoch);
1501 Ok(())
1502 }
1503
1504 fn pending_epoch(&self) -> Epoch {
1505 Epoch(self.epoch.visible().0 + 1)
1506 }
1507
1508 fn is_shared(&self) -> bool {
1511 matches!(self.wal, WalSink::Shared(_))
1512 }
1513
1514 fn ensure_txn_id(&mut self) -> u64 {
1518 if self.current_txn_id == 0 {
1519 let id = match &self.wal {
1520 WalSink::Shared(s) => {
1521 let mut g = s.txn_ids.lock();
1522 let v = *g;
1523 *g = g.wrapping_add(1);
1524 v
1525 }
1526 WalSink::Private(_) => 1,
1527 };
1528 self.current_txn_id = id;
1529 }
1530 self.current_txn_id
1531 }
1532
1533 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1536 let txn_id = self.ensure_txn_id();
1537 let table_id = self.table_id;
1538 match &mut self.wal {
1539 WalSink::Private(w) => {
1540 w.append_txn(txn_id, op)?;
1541 }
1542 WalSink::Shared(s) => {
1543 s.wal.lock().append(txn_id, table_id, op)?;
1544 }
1545 }
1546 Ok(())
1547 }
1548
1549 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1557 Ok(self.put_returning(columns)?.0)
1558 }
1559
1560 pub fn put_returning(
1565 &mut self,
1566 mut columns: Vec<(u16, Value)>,
1567 ) -> Result<(RowId, Option<i64>)> {
1568 let assigned = self.fill_auto_inc(&mut columns)?;
1569 self.schema.validate_not_null(&columns)?;
1570 let row_id = self.allocator.alloc();
1571 let epoch = self.pending_epoch();
1572 let mut row = Row::new(row_id, epoch);
1573 for (col_id, val) in columns {
1574 row.columns.insert(col_id, val);
1575 }
1576 self.commit_rows(vec![row], assigned.is_some())?;
1577 Ok((row_id, assigned))
1578 }
1579
1580 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1583 Ok(self
1584 .put_batch_returning(batch)?
1585 .into_iter()
1586 .map(|(r, _)| r)
1587 .collect())
1588 }
1589
1590 pub fn put_batch_returning(
1593 &mut self,
1594 batch: Vec<Vec<(u16, Value)>>,
1595 ) -> Result<Vec<(RowId, Option<i64>)>> {
1596 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1597 for mut cols in batch {
1598 let assigned = self.fill_auto_inc(&mut cols)?;
1599 filled.push((cols, assigned));
1600 }
1601 for (cols, _) in &filled {
1602 self.schema.validate_not_null(cols)?;
1603 }
1604 let epoch = self.pending_epoch();
1605 let mut rows = Vec::with_capacity(filled.len());
1606 let mut ids = Vec::with_capacity(filled.len());
1607 for (cols, assigned) in filled {
1608 let row_id = self.allocator.alloc();
1609 let mut row = Row::new(row_id, epoch);
1610 for (c, v) in cols {
1611 row.columns.insert(c, v);
1612 }
1613 ids.push((row_id, assigned));
1614 rows.push(row);
1615 }
1616 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1617 self.commit_rows(rows, all_auto_generated)?;
1618 Ok(ids)
1619 }
1620
1621 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1627 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1628 return Ok(None);
1629 };
1630 let pos = columns.iter().position(|(c, _)| *c == cid);
1631 let assigned = match pos {
1632 Some(i) => match &columns[i].1 {
1633 Value::Null => {
1634 let next = self.alloc_auto_inc_value()?;
1635 columns[i].1 = Value::Int64(next);
1636 Some(next)
1637 }
1638 Value::Int64(n) => {
1639 self.advance_auto_inc_past(*n)?;
1640 None
1641 }
1642 other => {
1643 return Err(MongrelError::InvalidArgument(format!(
1644 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1645 other
1646 )))
1647 }
1648 },
1649 None => {
1650 let next = self.alloc_auto_inc_value()?;
1651 columns.push((cid, Value::Int64(next)));
1652 Some(next)
1653 }
1654 };
1655 Ok(assigned)
1656 }
1657
1658 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1660 self.ensure_auto_inc_seeded()?;
1661 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1663 let v = ai.next;
1664 ai.next = ai.next.saturating_add(1);
1665 Ok(v)
1666 }
1667
1668 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1671 self.ensure_auto_inc_seeded()?;
1672 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1673 let floor = used.saturating_add(1).max(1);
1674 if ai.next < floor {
1675 ai.next = floor;
1676 }
1677 Ok(())
1678 }
1679
1680 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1685 let needs_seed = match self.auto_inc {
1686 Some(ai) => !ai.seeded,
1687 None => return Ok(()),
1688 };
1689 if !needs_seed {
1690 return Ok(());
1691 }
1692 if self.seed_empty_auto_inc() {
1693 return Ok(());
1694 }
1695 let cid = self
1696 .auto_inc
1697 .as_ref()
1698 .expect("auto-inc column present")
1699 .column_id;
1700 let max = self.scan_max_int64(cid)?;
1701 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1702 let floor = max.saturating_add(1).max(1);
1703 if ai.next < floor {
1704 ai.next = floor;
1705 }
1706 ai.seeded = true;
1707 Ok(())
1708 }
1709
1710 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1711 if n == 0 || self.auto_inc.is_none() {
1712 return Ok(None);
1713 }
1714 self.ensure_auto_inc_seeded()?;
1715 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1716 let start = ai.next;
1717 ai.next = ai.next.saturating_add(n as i64);
1718 Ok(Some(start))
1719 }
1720
1721 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1725 let mut max: i64 = 0;
1726 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1727 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1728 if *n > max {
1729 max = *n;
1730 }
1731 }
1732 }
1733 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1734 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1735 if *n > max {
1736 max = *n;
1737 }
1738 }
1739 }
1740 for rr in self.run_refs.clone() {
1741 let reader = self.open_reader(rr.run_id)?;
1742 if let Some(stats) = reader.column_page_stats(column_id) {
1743 for s in stats {
1744 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1745 if n > max {
1746 max = n;
1747 }
1748 }
1749 }
1750 } else if reader.has_column(column_id) {
1751 if let columnar::NativeColumn::Int64 { data, validity } =
1752 reader.column_native_shared(column_id)?
1753 {
1754 for (i, n) in data.iter().enumerate() {
1755 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1756 {
1757 max = *n;
1758 }
1759 }
1760 }
1761 }
1762 }
1763 Ok(max)
1764 }
1765
1766 fn seed_empty_auto_inc(&mut self) -> bool {
1767 let Some(ai) = self.auto_inc.as_mut() else {
1768 return false;
1769 };
1770 if ai.seeded || self.live_count != 0 {
1771 return false;
1772 }
1773 if ai.next < 1 {
1774 ai.next = 1;
1775 }
1776 ai.seeded = true;
1777 true
1778 }
1779
1780 fn advance_auto_inc_from_native_columns(
1781 &mut self,
1782 columns: &[(u16, columnar::NativeColumn)],
1783 n: usize,
1784 live_before: u64,
1785 ) -> Result<()> {
1786 let Some(ai) = self.auto_inc.as_mut() else {
1787 return Ok(());
1788 };
1789 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1790 return Ok(());
1791 };
1792 let columnar::NativeColumn::Int64 { data, validity } = col else {
1793 return Err(MongrelError::InvalidArgument(format!(
1794 "AUTO_INCREMENT column {} must be Int64",
1795 ai.column_id
1796 )));
1797 };
1798 let max = if native_int64_strictly_increasing(col, n) {
1799 data.get(n.saturating_sub(1)).copied()
1800 } else {
1801 data.iter()
1802 .take(n)
1803 .enumerate()
1804 .filter_map(|(i, v)| {
1805 if validity.is_empty() || columnar::validity_bit(validity, i) {
1806 Some(*v)
1807 } else {
1808 None
1809 }
1810 })
1811 .max()
1812 };
1813 if let Some(max) = max {
1814 let floor = max.saturating_add(1).max(1);
1815 if ai.next < floor {
1816 ai.next = floor;
1817 }
1818 if ai.seeded || live_before == 0 {
1819 ai.seeded = true;
1820 }
1821 }
1822 Ok(())
1823 }
1824
1825 fn fill_auto_inc_native_columns(
1826 &mut self,
1827 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1828 n: usize,
1829 ) -> Result<()> {
1830 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1831 return Ok(());
1832 };
1833 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1834 if let Some(start) = self.alloc_auto_inc_range(n)? {
1835 columns.push((
1836 cid,
1837 columnar::NativeColumn::Int64 {
1838 data: (start..start.saturating_add(n as i64)).collect(),
1839 validity: vec![0xFF; n.div_ceil(8)],
1840 },
1841 ));
1842 }
1843 return Ok(());
1844 };
1845
1846 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1847 return Err(MongrelError::InvalidArgument(format!(
1848 "AUTO_INCREMENT column {cid} must be Int64"
1849 )));
1850 };
1851 if data.len() < n {
1852 return Err(MongrelError::InvalidArgument(format!(
1853 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1854 data.len()
1855 )));
1856 }
1857 if columnar::all_non_null(validity, n) {
1858 return Ok(());
1859 }
1860 if validity.iter().all(|b| *b == 0) {
1861 if let Some(start) = self.alloc_auto_inc_range(n)? {
1862 for (i, slot) in data.iter_mut().take(n).enumerate() {
1863 *slot = start.saturating_add(i as i64);
1864 }
1865 *validity = vec![0xFF; n.div_ceil(8)];
1866 }
1867 return Ok(());
1868 }
1869
1870 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1871 for (i, slot) in data.iter_mut().enumerate().take(n) {
1872 if columnar::validity_bit(validity, i) {
1873 self.advance_auto_inc_past(*slot)?;
1874 } else {
1875 *slot = self.alloc_auto_inc_value()?;
1876 }
1877 }
1878 *validity = new_validity;
1879 Ok(())
1880 }
1881
1882 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1896 if self.auto_inc.is_none() {
1897 return Ok(None);
1898 }
1899 Ok(Some(self.alloc_auto_inc_value()?))
1900 }
1901
1902 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1908 let payload = bincode::serialize(&rows)?;
1909 self.wal_append_data(Op::Put {
1910 table_id: self.table_id,
1911 rows: payload,
1912 })?;
1913 if self.is_shared() {
1914 self.pending_rows_auto_inc
1915 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1916 self.pending_rows.extend(rows);
1917 } else {
1918 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1919 }
1920 Ok(())
1921 }
1922
1923 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1928 self.apply_put_rows_inner(rows, true)
1929 }
1930
1931 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1932 if check_existing_pk {
1933 self.ensure_indexes_complete()?;
1934 }
1935 if rows.len() == 1 {
1939 let row = rows.into_iter().next().expect("len checked");
1940 return self.apply_put_row_single(row, check_existing_pk);
1941 }
1942 let pk_id = self.schema.primary_key().map(|c| c.id);
1959 let probe = match pk_id {
1960 Some(pid) => {
1961 check_existing_pk
1962 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
1963 }
1964 None => false,
1965 };
1966 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
1969 for r in rows {
1970 for &cid in r.columns.keys() {
1971 self.pending_put_cols.insert(cid);
1972 }
1973 match pk_id {
1974 Some(pid) if probe || maintain_pk_by_row => {
1975 if let Some(pk_val) = r.columns.get(&pid) {
1976 let key = self.index_lookup_key(pid, pk_val);
1977 if probe {
1978 if let Some(old_rid) = self.hot.get(&key) {
1979 if old_rid != r.row_id {
1980 self.tombstone_row(old_rid, r.committed_epoch, true);
1981 }
1982 }
1983 }
1984 if maintain_pk_by_row {
1985 self.pk_by_row.insert(r.row_id, key);
1986 }
1987 }
1988 }
1989 Some(_) => {}
1990 None => {
1991 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1992 }
1993 }
1994 self.index_row(&r);
1995 self.reservoir.offer(r.row_id.0);
1996 self.memtable.upsert(r);
1997 self.live_count = self.live_count.saturating_add(1);
2000 }
2001 Ok(())
2002 }
2003
2004 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2008 for &cid in row.columns.keys() {
2009 self.pending_put_cols.insert(cid);
2010 }
2011 let epoch = row.committed_epoch;
2012 if let Some(pk_col) = self.schema.primary_key() {
2013 let pk_id = pk_col.id;
2014 if let Some(pk_val) = row.columns.get(&pk_id) {
2015 let maintain_pk_by_row = self.pk_by_row_complete;
2019 if check_existing_pk || maintain_pk_by_row {
2020 let key = self.index_lookup_key(pk_id, pk_val);
2021 if check_existing_pk {
2022 if let Some(old_rid) = self.hot.get(&key) {
2023 if old_rid != row.row_id {
2024 self.tombstone_row(old_rid, epoch, true);
2025 }
2026 }
2027 }
2028 if maintain_pk_by_row {
2029 self.pk_by_row.insert(row.row_id, key);
2030 }
2031 }
2032 }
2033 } else {
2034 self.hot
2035 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2036 }
2037 self.index_row(&row);
2038 self.reservoir.offer(row.row_id.0);
2039 self.memtable.upsert(row);
2040 self.live_count = self.live_count.saturating_add(1);
2041 Ok(())
2042 }
2043
2044 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2047 self.allocator.alloc()
2048 }
2049
2050 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2058 self.ensure_indexes_complete()?;
2059 let n = rows.len();
2060 for r in rows {
2061 for &cid in r.columns.keys() {
2062 self.pending_put_cols.insert(cid);
2063 }
2064 }
2065 let (losers, winner_pks) = self.partition_pk_winners(rows);
2066 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2067 for (key, &row_id) in &winner_pks {
2069 if let Some(old_rid) = self.hot.get(key) {
2070 if old_rid != row_id {
2071 self.tombstone_row(old_rid, epoch, true);
2072 }
2073 }
2074 }
2075 for &loser_rid in &losers {
2078 self.tombstone_row(loser_rid, epoch, false);
2079 }
2080 for (key, row_id) in winner_pks {
2082 self.insert_hot_pk(key, row_id);
2083 }
2084 if self.schema.primary_key().is_none() {
2085 for r in rows {
2086 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2087 }
2088 }
2089 for r in rows {
2090 self.allocator.advance_to(r.row_id);
2091 if !losers.contains(&r.row_id) {
2092 self.index_row(r);
2093 }
2094 }
2095 for r in rows {
2096 if !losers.contains(&r.row_id) {
2097 self.reservoir.offer(r.row_id.0);
2098 }
2099 }
2100 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2101 Ok(())
2102 }
2103
2104 pub(crate) fn recover_apply(
2109 &mut self,
2110 rows: Vec<Row>,
2111 deletes: Vec<(RowId, Epoch)>,
2112 ) -> Result<()> {
2113 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2117 std::collections::BTreeMap::new();
2118 for row in rows {
2119 self.allocator.advance_to(row.row_id);
2120 if let Some(ai) = self.auto_inc.as_mut() {
2125 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2126 if *n + 1 > ai.next {
2127 ai.next = *n + 1;
2128 }
2129 }
2130 }
2131 by_epoch.entry(row.committed_epoch).or_default().push(row);
2132 }
2133 for (epoch, group) in by_epoch {
2134 let (losers, winner_pks) = self.partition_pk_winners(&group);
2135 for (key, &row_id) in &winner_pks {
2137 if let Some(old_rid) = self.hot.get(key) {
2138 if old_rid != row_id {
2139 self.tombstone_row(old_rid, epoch, false);
2140 }
2141 }
2142 }
2143 for (key, row_id) in winner_pks {
2144 self.insert_hot_pk(key, row_id);
2145 }
2146 if self.schema.primary_key().is_none() {
2147 for r in &group {
2148 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2149 }
2150 }
2151 for r in &group {
2152 if !losers.contains(&r.row_id) {
2153 self.memtable.upsert(r.clone());
2154 self.index_row(r);
2155 }
2156 }
2157 }
2158 for (rid, epoch) in deletes {
2159 self.memtable.tombstone(rid, epoch);
2160 self.remove_hot_for_row(rid, epoch);
2161 }
2162 self.reservoir_complete = false;
2165 Ok(())
2166 }
2167
2168 pub(crate) fn flushed_epoch(&self) -> u64 {
2170 self.flushed_epoch
2171 }
2172
2173 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2175 self.schema.validate_not_null(cells)
2176 }
2177
2178 fn validate_columns_not_null(
2182 &self,
2183 columns: &[(u16, columnar::NativeColumn)],
2184 n: usize,
2185 ) -> Result<()> {
2186 let by_id: HashMap<u16, &columnar::NativeColumn> =
2187 columns.iter().map(|(id, c)| (*id, c)).collect();
2188 for col in &self.schema.columns {
2189 if col.flags.contains(ColumnFlags::NULLABLE) {
2190 continue;
2191 }
2192 match by_id.get(&col.id) {
2193 None => {
2194 return Err(MongrelError::InvalidArgument(format!(
2195 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2196 col.name, col.id
2197 )));
2198 }
2199 Some(c) => {
2200 if c.null_count(n) != 0 {
2201 return Err(MongrelError::InvalidArgument(format!(
2202 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2203 col.name, col.id
2204 )));
2205 }
2206 }
2207 }
2208 }
2209 Ok(())
2210 }
2211
2212 fn bulk_pk_winner_indices(
2217 &self,
2218 columns: &[(u16, columnar::NativeColumn)],
2219 n: usize,
2220 ) -> Option<Vec<usize>> {
2221 let pk_col = self.schema.primary_key()?;
2222 let pk_id = pk_col.id;
2223 let pk_ty = pk_col.ty;
2224 let by_id: HashMap<u16, &columnar::NativeColumn> =
2225 columns.iter().map(|(id, c)| (*id, c)).collect();
2226 let pk_native = by_id.get(&pk_id)?;
2227 if native_int64_strictly_increasing(pk_native, n) {
2228 return None;
2229 }
2230 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2232 let mut null_pk_rows: Vec<usize> = Vec::new();
2233 for i in 0..n {
2234 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2235 Some(key) => {
2236 last.insert(key, i);
2237 }
2238 None => null_pk_rows.push(i),
2239 }
2240 }
2241 let mut winners: HashSet<usize> = last.values().copied().collect();
2242 for i in null_pk_rows {
2243 winners.insert(i);
2244 }
2245 Some((0..n).filter(|i| winners.contains(i)).collect())
2246 }
2247
2248 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2250 let epoch = self.pending_epoch();
2251 self.wal_append_data(Op::Delete {
2252 table_id: self.table_id,
2253 row_ids: vec![row_id],
2254 })?;
2255 if self.is_shared() {
2256 self.pending_dels.push(row_id);
2257 } else {
2258 self.apply_delete(row_id, epoch);
2259 }
2260 Ok(())
2261 }
2262
2263 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2264 let pre = self.get(row_id, self.snapshot());
2265 self.delete(row_id)?;
2266 Ok(pre.map(|row| {
2267 let mut columns: Vec<_> = row.columns.into_iter().collect();
2268 columns.sort_by_key(|(id, _)| *id);
2269 OwnedRow { columns }
2270 }))
2271 }
2272
2273 pub fn truncate(&mut self) -> Result<()> {
2275 let epoch = self.pending_epoch();
2276 self.wal_append_data(Op::TruncateTable {
2277 table_id: self.table_id,
2278 })?;
2279 self.pending_rows.clear();
2280 self.pending_rows_auto_inc.clear();
2281 self.pending_dels.clear();
2282 self.pending_truncate = Some(epoch);
2283 Ok(())
2284 }
2285
2286 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2288 for rr in std::mem::take(&mut self.run_refs) {
2289 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2290 }
2291 for r in std::mem::take(&mut self.retiring) {
2292 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2293 }
2294 self.memtable = Memtable::new();
2295 self.mutable_run = MutableRun::new();
2296 self.hot = HotIndex::new();
2297 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2298 self.bitmap = bitmap;
2299 self.ann = ann;
2300 self.fm = fm;
2301 self.sparse = sparse;
2302 self.minhash = minhash;
2303 self.learned_range.clear();
2304 self.pk_by_row.clear();
2305 self.pk_by_row_complete = false;
2306 self.live_count = 0;
2307 self.reservoir = crate::reservoir::Reservoir::default();
2308 self.reservoir_complete = true;
2309 self.had_deletes = true;
2310 self.agg_cache.clear();
2311 self.global_idx_epoch = 0;
2312 self.indexes_complete = true;
2313 self.pending_delete_rids.clear();
2314 self.pending_put_cols.clear();
2315 self.pending_rows.clear();
2316 self.pending_rows_auto_inc.clear();
2317 self.pending_dels.clear();
2318 self.clear_result_cache();
2319 self.invalidate_index_checkpoint();
2320 Ok(())
2321 }
2322
2323 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2326 self.remove_hot_for_row(row_id, epoch);
2327 self.tombstone_row(row_id, epoch, true);
2328 }
2329
2330 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2334 let tombstone = Row {
2335 row_id,
2336 committed_epoch: epoch,
2337 columns: std::collections::HashMap::new(),
2338 deleted: true,
2339 };
2340 self.memtable.upsert(tombstone);
2341 self.pk_by_row.remove(&row_id);
2342 if adjust_live_count {
2343 self.live_count = self.live_count.saturating_sub(1);
2344 }
2345 self.pending_delete_rids.insert(row_id.0 as u32);
2347 self.had_deletes = true;
2350 self.agg_cache.clear();
2351 }
2352
2353 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2357 let Some(pk_col) = self.schema.primary_key() else {
2358 return;
2359 };
2360 if self.pk_by_row_complete {
2363 if let Some(key) = self.pk_by_row.remove(&row_id) {
2364 if self.hot.get(&key) == Some(row_id) {
2365 self.hot.remove(&key);
2366 }
2367 }
2368 return;
2369 }
2370 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2389 if self.indexes_complete {
2390 let pk_val = self
2391 .memtable
2392 .get_version(row_id, lookup_epoch)
2393 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2394 .or_else(|| {
2395 self.mutable_run
2396 .get_version(row_id, lookup_epoch)
2397 .filter(|(_, r)| !r.deleted)
2398 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2399 })
2400 .or_else(|| {
2401 self.run_refs.iter().find_map(|rr| {
2402 let mut reader = self.open_reader(rr.run_id).ok()?;
2403 let (_, deleted, val) = reader
2404 .get_version_column(row_id, lookup_epoch, pk_col.id)
2405 .ok()??;
2406 if deleted {
2407 return None;
2408 }
2409 val
2410 })
2411 });
2412 if let Some(pk_val) = pk_val {
2413 let key = self.index_lookup_key(pk_col.id, &pk_val);
2414 if self.hot.get(&key) == Some(row_id) {
2415 self.hot.remove(&key);
2416 }
2417 return;
2418 }
2419 }
2420 self.refresh_pk_by_row_from_hot();
2425 if let Some(key) = self.pk_by_row.remove(&row_id) {
2426 if self.hot.get(&key) == Some(row_id) {
2427 self.hot.remove(&key);
2428 }
2429 }
2430 }
2431
2432 fn partition_pk_winners(
2437 &self,
2438 rows: &[Row],
2439 ) -> (
2440 std::collections::HashSet<RowId>,
2441 std::collections::HashMap<Vec<u8>, RowId>,
2442 ) {
2443 let mut losers = std::collections::HashSet::new();
2444 let Some(pk_col) = self.schema.primary_key() else {
2445 return (losers, std::collections::HashMap::new());
2446 };
2447 let pk_id = pk_col.id;
2448 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2449 std::collections::HashMap::new();
2450 for r in rows {
2451 let Some(pk_val) = r.columns.get(&pk_id) else {
2452 continue;
2453 };
2454 let key = self.index_lookup_key(pk_id, pk_val);
2455 if let Some(&old_rid) = winners.get(&key) {
2456 losers.insert(old_rid);
2457 }
2458 winners.insert(key, r.row_id);
2459 }
2460 (losers, winners)
2461 }
2462
2463 fn index_row(&mut self, row: &Row) {
2464 if row.deleted {
2465 return;
2466 }
2467 if self.column_keys.is_empty() {
2471 index_into(
2472 &self.schema,
2473 row,
2474 &mut self.hot,
2475 &mut self.bitmap,
2476 &mut self.ann,
2477 &mut self.fm,
2478 &mut self.sparse,
2479 &mut self.minhash,
2480 );
2481 return;
2482 }
2483 let effective_row = self.tokenized_for_indexes(row);
2484 index_into(
2485 &self.schema,
2486 &effective_row,
2487 &mut self.hot,
2488 &mut self.bitmap,
2489 &mut self.ann,
2490 &mut self.fm,
2491 &mut self.sparse,
2492 &mut self.minhash,
2493 );
2494 }
2495
2496 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2502 if self.column_keys.is_empty() {
2503 return row.clone();
2504 }
2505 #[cfg(feature = "encryption")]
2506 {
2507 use crate::encryption::SCHEME_HMAC_EQ;
2508 let mut tok = row.clone();
2509 for (&cid, &(_, scheme)) in &self.column_keys {
2510 if scheme != SCHEME_HMAC_EQ {
2511 continue;
2512 }
2513 if let Some(v) = tok.columns.get(&cid).cloned() {
2514 if let Some(t) = self.tokenize_value(cid, &v) {
2515 tok.columns.insert(cid, t);
2516 }
2517 }
2518 }
2519 tok
2520 }
2521 #[cfg(not(feature = "encryption"))]
2522 {
2523 row.clone()
2524 }
2525 }
2526
2527 pub fn commit(&mut self) -> Result<Epoch> {
2532 if self.is_shared() {
2533 self.commit_shared()
2534 } else {
2535 self.commit_private()
2536 }
2537 }
2538
2539 fn commit_private(&mut self) -> Result<Epoch> {
2541 let commit_lock = Arc::clone(&self.commit_lock);
2545 let _g = commit_lock.lock();
2546 let new_epoch = self.epoch.bump_assigned();
2547 let txn_id = self.current_txn_id;
2548 match &mut self.wal {
2552 WalSink::Private(w) => {
2553 w.append_txn(
2554 txn_id,
2555 Op::TxnCommit {
2556 epoch: new_epoch.0,
2557 added_runs: Vec::new(),
2558 },
2559 )?;
2560 w.sync()?;
2561 }
2562 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2563 }
2564 if let Some(epoch) = self.pending_truncate.take() {
2566 self.apply_truncate(epoch)?;
2567 }
2568 self.invalidate_pending_cache();
2569 self.persist_manifest(new_epoch)?;
2570 self.epoch.publish_in_order(new_epoch);
2574 self.current_txn_id += 1;
2575 Ok(new_epoch)
2576 }
2577
2578 fn commit_shared(&mut self) -> Result<Epoch> {
2584 use std::sync::atomic::Ordering;
2585 let s = match &self.wal {
2586 WalSink::Shared(s) => s.clone(),
2587 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2588 };
2589 if s.poisoned.load(Ordering::Relaxed) {
2590 return Err(MongrelError::Other(
2591 "database poisoned by fsync error".into(),
2592 ));
2593 }
2594 let commit_lock = Arc::clone(&self.commit_lock);
2601 let _g = commit_lock.lock();
2602 let txn_id = self.ensure_txn_id();
2605 let (new_epoch, commit_seq) = {
2606 let mut wal = s.wal.lock();
2607 let new_epoch = self.epoch.bump_assigned();
2608 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2609 (new_epoch, seq)
2610 };
2611 s.group
2612 .await_durable(&s.wal, commit_seq)
2613 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2614
2615 if self.pending_truncate.take().is_some() {
2618 self.apply_truncate(new_epoch)?;
2619 }
2620 let mut rows = std::mem::take(&mut self.pending_rows);
2621 if !rows.is_empty() {
2622 for r in &mut rows {
2623 r.committed_epoch = new_epoch;
2624 }
2625 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2626 let all_auto_generated =
2627 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2628 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2629 } else {
2630 self.pending_rows_auto_inc.clear();
2631 }
2632 let dels = std::mem::take(&mut self.pending_dels);
2633 for rid in dels {
2634 self.apply_delete(rid, new_epoch);
2635 }
2636
2637 self.invalidate_pending_cache();
2638 self.persist_manifest(new_epoch)?;
2639 self.epoch.publish_in_order(new_epoch);
2640 self.current_txn_id = 0;
2642 Ok(new_epoch)
2643 }
2644
2645 pub fn flush(&mut self) -> Result<Epoch> {
2653 self.ensure_indexes_complete()?;
2654 let epoch = self.commit()?;
2655 let rows = self.memtable.drain_sorted();
2656 if !rows.is_empty() {
2657 self.mutable_run.insert_many(rows);
2658 }
2659 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2660 self.spill_mutable_run(epoch)?;
2661 self.mark_flushed(epoch)?;
2665 self.persist_manifest(epoch)?;
2666 self.build_learned_ranges()?;
2667 self.checkpoint_indexes(epoch);
2670 }
2671 Ok(epoch)
2674 }
2675
2676 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2683 let op = Op::Flush {
2684 table_id: self.table_id,
2685 flushed_epoch: epoch.0,
2686 };
2687 match &mut self.wal {
2688 WalSink::Private(w) => {
2689 w.append_system(op)?;
2690 w.sync()?;
2691 }
2692 WalSink::Shared(s) => {
2693 s.wal.lock().append_system(op)?;
2698 }
2699 }
2700 self.flushed_epoch = epoch.0;
2701 if matches!(self.wal, WalSink::Private(_)) {
2702 self.rotate_wal(epoch)?;
2703 }
2704 Ok(())
2705 }
2706
2707 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2711 let rows = self.mutable_run.drain_sorted();
2712 if rows.is_empty() {
2713 return Ok(());
2714 }
2715 let run_id = self.next_run_id;
2716 self.next_run_id += 1;
2717 let path = self.run_path(run_id);
2718 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2719 if let Some(kek) = &self.kek {
2720 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2721 }
2722 let header = writer.write(&path, &rows)?;
2723 self.run_refs.push(RunRef {
2724 run_id: run_id as u128,
2725 level: 0,
2726 epoch_created: epoch.0,
2727 row_count: header.row_count,
2728 });
2729 Ok(())
2730 }
2731
2732 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2736 self.mutable_run_spill_bytes = bytes.max(1);
2737 }
2738
2739 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2743 self.compaction_zstd_level = level;
2744 }
2745
2746 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2750 self.result_cache.lock().set_max_bytes(max_bytes);
2751 }
2752
2753 pub(crate) fn clear_result_cache(&mut self) {
2757 self.result_cache.lock().clear();
2758 }
2759
2760 pub fn mutable_run_len(&self) -> usize {
2762 self.mutable_run.len()
2763 }
2764
2765 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2768 self.mutable_run.drain_sorted()
2769 }
2770
2771 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2776 let epoch = self.commit()?;
2777 let n = batch.len();
2778 if n == 0 {
2779 return Ok(epoch);
2780 }
2781 let live_before = self.live_count;
2782 self.spill_mutable_run(epoch)?;
2786 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
2787 && self.indexes_complete
2788 && self.run_refs.is_empty()
2789 && self.memtable.is_empty()
2790 && self.mutable_run.is_empty();
2791 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2797 use rayon::prelude::*;
2798 self.schema
2799 .columns
2800 .par_iter()
2801 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
2802 .collect::<Vec<_>>()
2803 };
2804 drop(batch);
2805 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2810 self.validate_columns_not_null(&user_columns, n)?;
2811 let winner_idx = self
2812 .bulk_pk_winner_indices(&user_columns, n)
2813 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2814 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2815 match winner_idx.as_deref() {
2816 Some(idx) => {
2817 let compacted = user_columns
2818 .iter()
2819 .map(|(id, c)| (*id, c.gather(idx)))
2820 .collect();
2821 (compacted, idx.len())
2822 }
2823 None => (std::mem::take(&mut user_columns), n),
2824 };
2825 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2826 let first = self.allocator.alloc_range(write_n as u64).0;
2827 for rid in first..first + write_n as u64 {
2828 self.reservoir.offer(rid);
2829 }
2830 let run_id = self.next_run_id;
2831 self.next_run_id += 1;
2832 let path = self.run_path(run_id);
2833 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2834 .clean(true)
2835 .with_lz4()
2836 .with_native_endian();
2837 if let Some(kek) = &self.kek {
2838 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2839 }
2840 let header = writer.write_native(&path, &write_columns, write_n, first)?;
2841 self.run_refs.push(RunRef {
2842 run_id: run_id as u128,
2843 level: 0,
2844 epoch_created: epoch.0,
2845 row_count: header.row_count,
2846 });
2847 self.live_count = self.live_count.saturating_add(write_n as u64);
2848 if eager_index_build {
2849 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
2850 self.index_columns_bulk(&write_columns, &row_ids);
2851 self.indexes_complete = true;
2852 self.build_learned_ranges()?;
2853 } else {
2854 self.indexes_complete = false;
2855 }
2856 self.mark_flushed(epoch)?;
2857 self.persist_manifest(epoch)?;
2858 if eager_index_build {
2859 self.checkpoint_indexes(epoch);
2860 }
2861 self.clear_result_cache();
2862 Ok(epoch)
2863 }
2864
2865 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
2868 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
2869 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
2870 let segment_no = segment
2873 .file_stem()
2874 .and_then(|s| s.to_str())
2875 .and_then(|s| s.strip_prefix("seg-"))
2876 .and_then(|s| s.parse::<u64>().ok())
2877 .unwrap_or(0);
2878 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
2879 wal.set_sync_byte_threshold(self.sync_byte_threshold);
2880 wal.sync()?;
2881 self.wal = WalSink::Private(wal);
2882 Ok(())
2883 }
2884
2885 pub(crate) fn invalidate_pending_cache(&mut self) {
2890 self.result_cache
2891 .lock()
2892 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
2893 self.pending_delete_rids.clear();
2894 self.pending_put_cols.clear();
2895 }
2896
2897 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
2898 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
2899 m.current_epoch = epoch.0;
2900 m.next_row_id = self.allocator.current().0;
2901 m.runs = self.run_refs.clone();
2902 m.live_count = self.live_count;
2903 m.global_idx_epoch = self.global_idx_epoch;
2904 m.flushed_epoch = self.flushed_epoch;
2905 m.retiring = self.retiring.clone();
2906 m.auto_inc_next = match self.auto_inc {
2910 Some(ai) if ai.seeded => ai.next,
2911 _ => 0,
2912 };
2913 let meta_dek = self.manifest_meta_dek();
2914 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
2915 Ok(())
2916 }
2917
2918 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
2924 if !self.indexes_complete {
2927 return;
2928 }
2929 let snap = global_idx::IndexSnapshot {
2930 hot: &self.hot,
2931 bitmap: &self.bitmap,
2932 ann: &self.ann,
2933 fm: &self.fm,
2934 sparse: &self.sparse,
2935 minhash: &self.minhash,
2936 learned_range: &self.learned_range,
2937 };
2938 let idx_dek = self.idx_dek();
2940 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
2941 .is_ok()
2942 {
2943 self.global_idx_epoch = epoch.0;
2944 let _ = self.persist_manifest(epoch);
2945 }
2946 }
2947
2948 pub(crate) fn invalidate_index_checkpoint(&mut self) {
2951 self.global_idx_epoch = 0;
2952 global_idx::remove(&self.dir);
2953 let _ = self.persist_manifest(self.epoch.visible());
2954 }
2955
2956 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
2959 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
2960 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
2961 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2962 best = Some((epoch, row));
2963 }
2964 }
2965 for rr in &self.run_refs {
2966 let Ok(mut reader) = self.open_reader(rr.run_id) else {
2967 continue;
2968 };
2969 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
2970 continue;
2971 };
2972 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2973 best = Some((epoch, row));
2974 }
2975 }
2976 match best {
2977 Some((_, r)) if r.deleted => None,
2978 Some((_, r)) => Some(r),
2979 None => None,
2980 }
2981 }
2982
2983 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
2987 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
2988 let mut fold = |row: Row| {
2989 best.entry(row.row_id.0)
2990 .and_modify(|e| {
2991 if row.committed_epoch > e.0 {
2992 *e = (row.committed_epoch, row.clone());
2993 }
2994 })
2995 .or_insert_with(|| (row.committed_epoch, row));
2996 };
2997 for row in self.memtable.visible_versions(snapshot.epoch) {
2998 fold(row);
2999 }
3000 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3001 fold(row);
3002 }
3003 for rr in &self.run_refs {
3004 let mut reader = self.open_reader(rr.run_id)?;
3005 for row in reader.visible_versions(snapshot.epoch)? {
3006 fold(row);
3007 }
3008 }
3009 let mut out: Vec<Row> = best
3010 .into_values()
3011 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3012 .collect();
3013 out.sort_by_key(|r| r.row_id);
3014 Ok(out)
3015 }
3016
3017 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3024 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3025 let rr = self.run_refs[0].clone();
3026 let mut reader = self.open_reader(rr.run_id)?;
3027 let idxs = reader.visible_indices(snapshot.epoch)?;
3028 let mut cols = Vec::with_capacity(self.schema.columns.len());
3029 for cdef in &self.schema.columns {
3030 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3031 }
3032 return Ok(cols);
3033 }
3034 let rows = self.visible_rows(snapshot)?;
3036 let mut cols: Vec<(u16, Vec<Value>)> = self
3037 .schema
3038 .columns
3039 .iter()
3040 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3041 .collect();
3042 for r in &rows {
3043 for (cid, vec) in cols.iter_mut() {
3044 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3045 }
3046 }
3047 Ok(cols)
3048 }
3049
3050 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3052 self.hot.get(key)
3053 }
3054
3055 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3060 self.ensure_indexes_complete()?;
3061 let snapshot = self.snapshot();
3062 crate::trace::QueryTrace::record(|t| {
3063 t.run_count = self.run_refs.len();
3064 t.memtable_rows = self.memtable.len();
3065 t.mutable_run_rows = self.mutable_run.len();
3066 });
3067 if q.conditions.is_empty() {
3071 crate::trace::QueryTrace::record(|t| {
3072 t.scan_mode = crate::trace::ScanMode::Materialized;
3073 t.row_materialized = true;
3074 });
3075 return self.visible_rows(snapshot);
3076 }
3077 crate::trace::QueryTrace::record(|t| {
3078 t.conditions_pushed = q.conditions.len();
3079 t.scan_mode = crate::trace::ScanMode::Materialized;
3080 t.row_materialized = true;
3081 });
3082 let mut sets: Vec<RowIdSet> = Vec::with_capacity(q.conditions.len());
3083 for c in &q.conditions {
3084 sets.push(self.resolve_condition(c, snapshot)?);
3085 }
3086 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3087 self.rows_for_rids(&rids, snapshot)
3088 }
3089
3090 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3095 use std::collections::HashMap;
3096 let mut rows = Vec::with_capacity(rids.len());
3097 let tier_size = self.memtable.len() + self.mutable_run.len();
3114 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3115 if rids.len().saturating_mul(24) < tier_size {
3116 for &rid in rids {
3117 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3118 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3119 let newest = match (mem, mrun) {
3120 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3121 (Some((_, mr)), None) => Some(mr),
3122 (None, Some((_, rr))) => Some(rr),
3123 (None, None) => None,
3124 };
3125 if let Some(row) = newest {
3126 overlay.insert(rid, row);
3127 }
3128 }
3129 } else {
3130 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3131 overlay
3132 .entry(row.row_id.0)
3133 .and_modify(|e| {
3134 if row.committed_epoch > e.committed_epoch {
3135 *e = row.clone();
3136 }
3137 })
3138 .or_insert(row);
3139 };
3140 for row in self.memtable.visible_versions(snapshot.epoch) {
3141 fold_newest(row, &mut overlay);
3142 }
3143 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3144 fold_newest(row, &mut overlay);
3145 }
3146 }
3147 if self.run_refs.len() == 1 {
3148 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3149 if rids.len().saturating_mul(24) < reader.row_count() {
3157 for &rid in rids {
3158 if let Some(r) = overlay.get(&rid) {
3159 if !r.deleted {
3160 rows.push(r.clone());
3161 }
3162 continue;
3163 }
3164 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3165 if !row.deleted {
3166 rows.push(row);
3167 }
3168 }
3169 }
3170 return Ok(rows);
3171 }
3172 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3181 enum Src {
3184 Overlay,
3185 Run,
3186 }
3187 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3188 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3189 for rid in rids {
3190 if overlay.contains_key(rid) {
3191 plan.push(Src::Overlay);
3192 continue;
3193 }
3194 match vis_rids.binary_search(&(*rid as i64)) {
3195 Ok(i) => {
3196 plan.push(Src::Run);
3197 fetch.push(positions[i]);
3198 }
3199 Err(_) => { }
3200 }
3201 }
3202 let fetched = reader.materialize_batch(&fetch)?;
3203 let mut fetched_iter = fetched.into_iter();
3204 for (rid, src) in rids.iter().zip(plan) {
3205 match src {
3206 Src::Overlay => {
3207 if let Some(r) = overlay.get(rid) {
3208 if !r.deleted {
3209 rows.push(r.clone());
3210 }
3211 }
3212 }
3213 Src::Run => {
3214 if let Some(row) = fetched_iter.next() {
3215 if !row.deleted {
3216 rows.push(row);
3217 }
3218 }
3219 }
3220 }
3221 }
3222 return Ok(rows);
3223 }
3224 let mut readers: Vec<_> = self
3228 .run_refs
3229 .iter()
3230 .map(|rr| self.open_reader(rr.run_id))
3231 .collect::<Result<Vec<_>>>()?;
3232 for rid in rids {
3233 if let Some(r) = overlay.get(rid) {
3234 if !r.deleted {
3235 rows.push(r.clone());
3236 }
3237 continue;
3238 }
3239 let mut best: Option<(Epoch, Row)> = None;
3240 for reader in readers.iter_mut() {
3241 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3242 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3243 best = Some((epoch, row));
3244 }
3245 }
3246 }
3247 if let Some((_, r)) = best {
3248 if !r.deleted {
3249 rows.push(r);
3250 }
3251 }
3252 }
3253 Ok(rows)
3254 }
3255
3256 pub fn indexes_complete(&self) -> bool {
3266 self.indexes_complete
3267 }
3268
3269 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3271 self.index_build_policy
3272 }
3273
3274 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3278 self.index_build_policy = policy;
3279 }
3280
3281 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3286 if !self.indexes_complete {
3290 return None;
3291 }
3292 let b = self.bitmap.get(&column_id)?;
3293 let result: Vec<Vec<u8>> = b
3294 .keys()
3295 .into_iter()
3296 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3297 .cloned()
3298 .collect();
3299 Some(result)
3300 }
3301
3302 pub fn fk_join_row_ids(
3303 &self,
3304 fk_column_id: u16,
3305 pk_values: &[Vec<u8>],
3306 fk_conditions: &[crate::query::Condition],
3307 snapshot: Snapshot,
3308 ) -> Result<Vec<u64>> {
3309 let Some(b) = self.bitmap.get(&fk_column_id) else {
3310 return Ok(Vec::new());
3311 };
3312 let mut join_set = {
3313 let mut acc = roaring::RoaringBitmap::new();
3314 for v in pk_values {
3315 acc |= b.get(v);
3316 }
3317 RowIdSet::from_roaring(acc)
3318 };
3319 if !fk_conditions.is_empty() {
3320 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3321 sets.push(join_set);
3322 for c in fk_conditions {
3323 sets.push(self.resolve_condition(c, snapshot)?);
3324 }
3325 join_set = RowIdSet::intersect_many(sets);
3326 }
3327 Ok(join_set.into_sorted_vec())
3328 }
3329
3330 pub fn fk_join_count(
3336 &self,
3337 fk_column_id: u16,
3338 pk_values: &[Vec<u8>],
3339 fk_conditions: &[crate::query::Condition],
3340 snapshot: Snapshot,
3341 ) -> Result<u64> {
3342 let Some(b) = self.bitmap.get(&fk_column_id) else {
3343 return Ok(0);
3344 };
3345 let mut acc = roaring::RoaringBitmap::new();
3346 for v in pk_values {
3347 acc |= b.get(v);
3348 }
3349 if fk_conditions.is_empty() {
3350 return Ok(acc.len());
3351 }
3352 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3353 sets.push(RowIdSet::from_roaring(acc));
3354 for c in fk_conditions {
3355 sets.push(self.resolve_condition(c, snapshot)?);
3356 }
3357 Ok(RowIdSet::intersect_many(sets).len() as u64)
3358 }
3359
3360 fn resolve_condition(
3365 &self,
3366 c: &crate::query::Condition,
3367 snapshot: Snapshot,
3368 ) -> Result<RowIdSet> {
3369 use crate::query::Condition;
3370 Ok(match c {
3371 Condition::Pk(key) => {
3372 let lookup = self
3373 .schema
3374 .primary_key()
3375 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3376 .unwrap_or_else(|| key.clone());
3377 self.hot
3378 .get(&lookup)
3379 .map(|r| RowIdSet::one(r.0))
3380 .unwrap_or_else(RowIdSet::empty)
3381 }
3382 Condition::BitmapEq { column_id, value } => {
3383 let lookup = self.index_lookup_key_bytes(*column_id, value);
3384 self.bitmap
3385 .get(column_id)
3386 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3387 .unwrap_or_else(RowIdSet::empty)
3388 }
3389 Condition::BitmapIn { column_id, values } => {
3390 let bm = self.bitmap.get(column_id);
3391 let mut acc = roaring::RoaringBitmap::new();
3392 if let Some(b) = bm {
3393 for v in values {
3394 let lookup = self.index_lookup_key_bytes(*column_id, v);
3395 acc |= b.get(&lookup);
3396 }
3397 }
3398 RowIdSet::from_roaring(acc)
3399 }
3400 Condition::FmContains { column_id, pattern } => self
3401 .fm
3402 .get(column_id)
3403 .map(|f| {
3404 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3405 })
3406 .unwrap_or_else(RowIdSet::empty),
3407 Condition::FmContainsAll {
3408 column_id,
3409 patterns,
3410 } => {
3411 if let Some(f) = self.fm.get(column_id) {
3414 let sets: Vec<RowIdSet> = patterns
3415 .iter()
3416 .map(|pat| {
3417 RowIdSet::from_unsorted(
3418 f.locate(pat).into_iter().map(|r| r.0).collect(),
3419 )
3420 })
3421 .collect();
3422 RowIdSet::intersect_many(sets)
3423 } else {
3424 RowIdSet::empty()
3425 }
3426 }
3427 Condition::Ann {
3428 column_id,
3429 query,
3430 k,
3431 } => self
3432 .ann
3433 .get(column_id)
3434 .map(|a| {
3435 RowIdSet::from_unsorted(
3436 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3437 )
3438 })
3439 .unwrap_or_else(RowIdSet::empty),
3440 Condition::SparseMatch {
3441 column_id,
3442 query,
3443 k,
3444 } => self
3445 .sparse
3446 .get(column_id)
3447 .map(|s| {
3448 RowIdSet::from_unsorted(
3449 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3450 )
3451 })
3452 .unwrap_or_else(RowIdSet::empty),
3453 Condition::MinHashSimilar {
3454 column_id,
3455 query,
3456 k,
3457 } => self
3458 .minhash
3459 .get(column_id)
3460 .map(|mh| {
3461 RowIdSet::from_unsorted(
3462 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3463 )
3464 })
3465 .unwrap_or_else(RowIdSet::empty),
3466 Condition::Range { column_id, lo, hi } => {
3467 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3476 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3477 } else if self.run_refs.len() == 1 {
3478 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3479 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3480 } else {
3481 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3482 };
3483 set.remove_many(self.overlay_rid_set(snapshot));
3484 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3485 set
3486 }
3487 Condition::RangeF64 {
3488 column_id,
3489 lo,
3490 lo_inclusive,
3491 hi,
3492 hi_inclusive,
3493 } => {
3494 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3497 RowIdSet::from_unsorted(
3498 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3499 .into_iter()
3500 .collect(),
3501 )
3502 } else if self.run_refs.len() == 1 {
3503 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3504 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3505 } else {
3506 return self.range_scan_f64(
3507 *column_id,
3508 *lo,
3509 *lo_inclusive,
3510 *hi,
3511 *hi_inclusive,
3512 snapshot,
3513 );
3514 };
3515 set.remove_many(self.overlay_rid_set(snapshot));
3516 self.range_scan_overlay_f64(
3517 &mut set,
3518 *column_id,
3519 *lo,
3520 *lo_inclusive,
3521 *hi,
3522 *hi_inclusive,
3523 snapshot,
3524 );
3525 set
3526 }
3527 Condition::IsNull { column_id } => {
3528 let mut set = if self.run_refs.len() == 1 {
3529 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3530 r.null_row_id_set(*column_id, true)?
3531 } else {
3532 return self.null_scan(*column_id, true, snapshot);
3533 };
3534 set.remove_many(self.overlay_rid_set(snapshot));
3535 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3536 set
3537 }
3538 Condition::IsNotNull { column_id } => {
3539 let mut set = if self.run_refs.len() == 1 {
3540 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3541 r.null_row_id_set(*column_id, false)?
3542 } else {
3543 return self.null_scan(*column_id, false, snapshot);
3544 };
3545 set.remove_many(self.overlay_rid_set(snapshot));
3546 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3547 set
3548 }
3549 })
3550 }
3551
3552 fn range_scan_i64(
3560 &self,
3561 column_id: u16,
3562 lo: i64,
3563 hi: i64,
3564 snapshot: Snapshot,
3565 ) -> Result<RowIdSet> {
3566 let mut row_ids = Vec::new();
3567 let overlay_rids = self.overlay_rid_set(snapshot);
3568 for rr in &self.run_refs {
3569 let mut reader = self.open_reader(rr.run_id)?;
3570 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3571 for rid in matched {
3572 if !overlay_rids.contains(&rid) {
3573 row_ids.push(rid);
3574 }
3575 }
3576 }
3577 let mut s = RowIdSet::from_unsorted(row_ids);
3578 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3579 Ok(s)
3580 }
3581
3582 fn range_scan_f64(
3585 &self,
3586 column_id: u16,
3587 lo: f64,
3588 lo_inclusive: bool,
3589 hi: f64,
3590 hi_inclusive: bool,
3591 snapshot: Snapshot,
3592 ) -> Result<RowIdSet> {
3593 let mut row_ids = Vec::new();
3594 let overlay_rids = self.overlay_rid_set(snapshot);
3595 for rr in &self.run_refs {
3596 let mut reader = self.open_reader(rr.run_id)?;
3597 let matched = reader.range_row_ids_visible_f64(
3598 column_id,
3599 lo,
3600 lo_inclusive,
3601 hi,
3602 hi_inclusive,
3603 snapshot.epoch,
3604 )?;
3605 for rid in matched {
3606 if !overlay_rids.contains(&rid) {
3607 row_ids.push(rid);
3608 }
3609 }
3610 }
3611 let mut s = RowIdSet::from_unsorted(row_ids);
3612 self.range_scan_overlay_f64(
3613 &mut s,
3614 column_id,
3615 lo,
3616 lo_inclusive,
3617 hi,
3618 hi_inclusive,
3619 snapshot,
3620 );
3621 Ok(s)
3622 }
3623
3624 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3626 let mut s = HashSet::new();
3627 for row in self.memtable.visible_versions(snapshot.epoch) {
3628 s.insert(row.row_id.0);
3629 }
3630 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3631 s.insert(row.row_id.0);
3632 }
3633 s
3634 }
3635
3636 fn range_scan_overlay_i64(
3637 &self,
3638 s: &mut RowIdSet,
3639 column_id: u16,
3640 lo: i64,
3641 hi: i64,
3642 snapshot: Snapshot,
3643 ) {
3644 let mut newest: HashMap<u64, &Row> = HashMap::new();
3649 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3650 let memtable = self.memtable.visible_versions(snapshot.epoch);
3651 for r in &mutable {
3652 newest.entry(r.row_id.0).or_insert(r);
3653 }
3654 for r in &memtable {
3655 newest.insert(r.row_id.0, r);
3656 }
3657 for row in newest.values() {
3658 if !row.deleted {
3659 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3660 if *v >= lo && *v <= hi {
3661 s.insert(row.row_id.0);
3662 }
3663 }
3664 }
3665 }
3666 }
3667
3668 #[allow(clippy::too_many_arguments)]
3669 fn range_scan_overlay_f64(
3670 &self,
3671 s: &mut RowIdSet,
3672 column_id: u16,
3673 lo: f64,
3674 lo_inclusive: bool,
3675 hi: f64,
3676 hi_inclusive: bool,
3677 snapshot: Snapshot,
3678 ) {
3679 let mut newest: HashMap<u64, &Row> = HashMap::new();
3682 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3683 let memtable = self.memtable.visible_versions(snapshot.epoch);
3684 for r in &mutable {
3685 newest.entry(r.row_id.0).or_insert(r);
3686 }
3687 for r in &memtable {
3688 newest.insert(r.row_id.0, r);
3689 }
3690 for row in newest.values() {
3691 if !row.deleted {
3692 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3693 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3694 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3695 if ok_lo && ok_hi {
3696 s.insert(row.row_id.0);
3697 }
3698 }
3699 }
3700 }
3701 }
3702
3703 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3706 let mut row_ids = Vec::new();
3707 let overlay_rids = self.overlay_rid_set(snapshot);
3708 for rr in &self.run_refs {
3709 let mut reader = self.open_reader(rr.run_id)?;
3710 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3711 for rid in matched {
3712 if !overlay_rids.contains(&rid) {
3713 row_ids.push(rid);
3714 }
3715 }
3716 }
3717 let mut s = RowIdSet::from_unsorted(row_ids);
3718 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3719 Ok(s)
3720 }
3721
3722 fn null_scan_overlay(
3726 &self,
3727 s: &mut RowIdSet,
3728 column_id: u16,
3729 want_nulls: bool,
3730 snapshot: Snapshot,
3731 ) {
3732 let mut newest: HashMap<u64, &Row> = HashMap::new();
3733 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3734 let memtable = self.memtable.visible_versions(snapshot.epoch);
3735 for r in &mutable {
3736 newest.entry(r.row_id.0).or_insert(r);
3737 }
3738 for r in &memtable {
3739 newest.insert(r.row_id.0, r);
3740 }
3741 for row in newest.values() {
3742 if row.deleted {
3743 continue;
3744 }
3745 let is_null = !row.columns.contains_key(&column_id)
3746 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3747 if is_null == want_nulls {
3748 s.insert(row.row_id.0);
3749 }
3750 }
3751 }
3752
3753 pub fn snapshot(&self) -> Snapshot {
3754 Snapshot::at(self.epoch.visible())
3755 }
3756
3757 pub fn pin_snapshot(&mut self) -> Snapshot {
3760 let e = self.epoch.visible();
3761 *self.pinned.entry(e).or_insert(0) += 1;
3762 Snapshot::at(e)
3763 }
3764
3765 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3767 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3768 *count -= 1;
3769 if *count == 0 {
3770 self.pinned.remove(&snap.epoch);
3771 }
3772 }
3773 }
3774
3775 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3785 let local = self.pinned.keys().next().copied();
3786 let global = self.snapshots.min_pinned();
3787 match (local, global) {
3788 (Some(a), Some(b)) => Some(a.min(b)),
3789 (Some(a), None) => Some(a),
3790 (None, b) => b,
3791 }
3792 }
3793
3794 pub fn current_epoch(&self) -> Epoch {
3795 self.epoch.visible()
3796 }
3797
3798 pub fn memtable_len(&self) -> usize {
3799 self.memtable.len()
3800 }
3801
3802 pub fn count(&self) -> u64 {
3805 self.live_count
3806 }
3807
3808 pub fn count_conditions(
3812 &mut self,
3813 conditions: &[crate::query::Condition],
3814 snapshot: Snapshot,
3815 ) -> Result<Option<u64>> {
3816 use crate::query::Condition;
3817 if conditions.is_empty() {
3818 return Ok(Some(self.live_count));
3819 }
3820 let served = |c: &Condition| {
3821 matches!(
3822 c,
3823 Condition::Pk(_)
3824 | Condition::BitmapEq { .. }
3825 | Condition::BitmapIn { .. }
3826 | Condition::FmContains { .. }
3827 | Condition::FmContainsAll { .. }
3828 | Condition::Ann { .. }
3829 | Condition::Range { .. }
3830 | Condition::RangeF64 { .. }
3831 | Condition::SparseMatch { .. }
3832 | Condition::MinHashSimilar { .. }
3833 | Condition::IsNull { .. }
3834 | Condition::IsNotNull { .. }
3835 )
3836 };
3837 if !conditions.iter().all(served) {
3838 return Ok(None);
3839 }
3840 self.ensure_indexes_complete()?;
3841 let mut sets = Vec::with_capacity(conditions.len());
3842 for condition in conditions {
3843 sets.push(self.resolve_condition(condition, snapshot)?);
3844 }
3845 let count = RowIdSet::intersect_many(sets).len() as u64;
3846 crate::trace::QueryTrace::record(|t| {
3847 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
3848 t.survivor_count = Some(count as usize);
3849 t.conditions_pushed = conditions.len();
3850 });
3851 Ok(Some(count))
3852 }
3853
3854 pub fn bulk_load_columns(
3863 &mut self,
3864 user_columns: Vec<(u16, columnar::NativeColumn)>,
3865 ) -> Result<Epoch> {
3866 self.bulk_load_columns_with(user_columns, 3, false, true)
3867 }
3868
3869 pub fn bulk_load_fast(
3876 &mut self,
3877 user_columns: Vec<(u16, columnar::NativeColumn)>,
3878 ) -> Result<Epoch> {
3879 self.bulk_load_columns_with(user_columns, -1, true, false)
3880 }
3881
3882 fn bulk_load_columns_with(
3883 &mut self,
3884 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
3885 zstd_level: i32,
3886 force_plain: bool,
3887 lz4: bool,
3888 ) -> Result<Epoch> {
3889 let epoch = self.commit()?;
3890 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
3891 if n == 0 {
3892 return Ok(epoch);
3893 }
3894 let live_before = self.live_count;
3895 self.spill_mutable_run(epoch)?;
3897 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3898 && self.indexes_complete
3899 && self.run_refs.is_empty()
3900 && self.memtable.is_empty()
3901 && self.mutable_run.is_empty();
3902 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3905 self.validate_columns_not_null(&user_columns, n)?;
3906 let winner_idx = self
3907 .bulk_pk_winner_indices(&user_columns, n)
3908 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
3909 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3910 match winner_idx.as_deref() {
3911 Some(idx) => {
3912 let compacted = user_columns
3913 .iter()
3914 .map(|(id, c)| (*id, c.gather(idx)))
3915 .collect();
3916 (compacted, idx.len())
3917 }
3918 None => (user_columns, n),
3919 };
3920 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3921 let first = self.allocator.alloc_range(write_n as u64).0;
3922 for rid in first..first + write_n as u64 {
3923 self.reservoir.offer(rid);
3924 }
3925 let run_id = self.next_run_id;
3926 self.next_run_id += 1;
3927 let path = self.run_path(run_id);
3928 let mut writer =
3929 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
3930 if force_plain {
3931 writer = writer.with_plain();
3932 } else if lz4 {
3933 writer = writer.with_lz4();
3936 } else {
3937 writer = writer.with_zstd_level(zstd_level);
3938 }
3939 if let Some(kek) = &self.kek {
3940 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3941 }
3942 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3943 self.run_refs.push(RunRef {
3944 run_id: run_id as u128,
3945 level: 0,
3946 epoch_created: epoch.0,
3947 row_count: header.row_count,
3948 });
3949 self.live_count = self.live_count.saturating_add(write_n as u64);
3950 if eager_index_build {
3951 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3952 self.index_columns_bulk(&write_columns, &row_ids);
3953 self.indexes_complete = true;
3954 self.build_learned_ranges()?;
3955 } else {
3956 self.indexes_complete = false;
3960 }
3961 self.mark_flushed(epoch)?;
3962 self.persist_manifest(epoch)?;
3963 if eager_index_build {
3964 self.checkpoint_indexes(epoch);
3965 }
3966 self.clear_result_cache();
3967 Ok(epoch)
3968 }
3969
3970 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
3988 let n = row_ids.len();
3989 if n == 0 {
3990 return;
3991 }
3992 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
3993 columns.iter().map(|(id, c)| (*id, c)).collect();
3994 let ty_of: std::collections::HashMap<u16, TypeId> =
3995 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
3996 let pk_id = self.schema.primary_key().map(|c| c.id);
3997
3998 for (i, &rid) in row_ids.iter().enumerate() {
3999 let row_id = RowId(rid);
4000 if let Some(pid) = pk_id {
4001 if let Some(col) = by_id.get(&pid) {
4002 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
4003 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4004 self.insert_hot_pk(key, row_id);
4005 }
4006 }
4007 }
4008 for idef in &self.schema.indexes {
4009 let Some(col) = by_id.get(&idef.column_id) else {
4010 continue;
4011 };
4012 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
4013 match idef.kind {
4014 IndexKind::Bitmap => {
4015 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4016 if let Some(key) =
4017 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4018 {
4019 b.insert(key, row_id);
4020 }
4021 }
4022 }
4023 IndexKind::FmIndex => {
4024 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4025 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4026 f.insert(bytes.to_vec(), row_id);
4027 }
4028 }
4029 }
4030 IndexKind::Sparse => {
4031 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4032 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4033 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4034 s.insert(&terms, row_id);
4035 }
4036 }
4037 }
4038 }
4039 IndexKind::MinHash => {
4040 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4041 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4042 let tokens = crate::index::token_hashes_from_bytes(bytes);
4043 mh.insert(&tokens, row_id);
4044 }
4045 }
4046 }
4047 _ => {}
4048 }
4049 }
4050 }
4051 }
4052
4053 pub fn visible_columns_native(
4058 &self,
4059 snapshot: Snapshot,
4060 projection: Option<&[u16]>,
4061 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4062 let wanted: Vec<u16> = match projection {
4063 Some(p) => p.to_vec(),
4064 None => self.schema.columns.iter().map(|c| c.id).collect(),
4065 };
4066 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4067 let rr = self.run_refs[0].clone();
4068 let mut reader = self.open_reader(rr.run_id)?;
4069 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4070 let all_visible = idxs.len() == reader.row_count();
4071 if reader.has_mmap() {
4077 use rayon::prelude::*;
4078 let valid: Vec<u16> = wanted
4081 .iter()
4082 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4083 .copied()
4084 .collect();
4085 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4087 .par_iter()
4088 .filter_map(|cid| {
4089 reader
4090 .column_native_shared(*cid)
4091 .ok()
4092 .map(|col| (*cid, col))
4093 })
4094 .collect();
4095 let cols = decoded
4096 .into_iter()
4097 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4098 .collect();
4099 return Ok(cols);
4100 }
4101 let mut cols = Vec::with_capacity(wanted.len());
4102 for cid in &wanted {
4103 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4104 Some(c) => c,
4105 None => continue,
4106 };
4107 let col = reader.column_native(cdef.id)?;
4108 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4109 }
4110 return Ok(cols);
4111 }
4112 let vcols = self.visible_columns(snapshot)?;
4113 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4114 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4115 .into_iter()
4116 .filter(|(id, _)| want_set.contains(id))
4117 .map(|(id, vals)| {
4118 let ty = self
4119 .schema
4120 .columns
4121 .iter()
4122 .find(|c| c.id == id)
4123 .map(|c| c.ty)
4124 .unwrap_or(TypeId::Bytes);
4125 (id, columnar::values_to_native(ty, &vals))
4126 })
4127 .collect();
4128 Ok(out)
4129 }
4130
4131 pub fn run_count(&self) -> usize {
4132 self.run_refs.len()
4133 }
4134
4135 pub fn memtable_is_empty(&self) -> bool {
4137 self.memtable.is_empty()
4138 }
4139
4140 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4144 self.page_cache.lock().stats()
4145 }
4146
4147 pub fn reset_page_cache_stats(&self) {
4149 self.page_cache.lock().reset_stats();
4150 }
4151
4152 pub fn run_ids(&self) -> Vec<u128> {
4155 self.run_refs.iter().map(|r| r.run_id).collect()
4156 }
4157
4158 pub fn single_run_is_clean(&self) -> bool {
4162 if self.run_refs.len() != 1 {
4163 return false;
4164 }
4165 self.open_reader(self.run_refs[0].run_id)
4166 .map(|r| r.is_clean())
4167 .unwrap_or(false)
4168 }
4169
4170 fn resolve_footprint(
4177 &self,
4178 conditions: &[crate::query::Condition],
4179 _snapshot: Snapshot,
4180 ) -> roaring::RoaringBitmap {
4181 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4182 return roaring::RoaringBitmap::new();
4183 }
4184 if self.run_refs.is_empty() {
4185 return roaring::RoaringBitmap::new();
4186 }
4187 if self.run_refs.len() == 1 {
4189 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4190 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4191 return rids.to_roaring_lossy();
4192 }
4193 }
4194 }
4195 roaring::RoaringBitmap::new()
4196 }
4197
4198 pub fn query_columns_native_cached(
4209 &mut self,
4210 conditions: &[crate::query::Condition],
4211 projection: Option<&[u16]>,
4212 snapshot: Snapshot,
4213 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4214 if conditions.is_empty() {
4215 return self.query_columns_native(conditions, projection, snapshot);
4216 }
4217 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4221 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4222 crate::trace::QueryTrace::record(|t| {
4223 t.result_cache_hit = true;
4224 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4225 });
4226 return Ok(Some((*hit).clone()));
4227 }
4228 let res = self.query_columns_native(conditions, projection, snapshot)?;
4229 if let Some(cols) = &res {
4230 let footprint = self.resolve_footprint(conditions, snapshot);
4231 let condition_cols = crate::query::condition_columns(conditions);
4232 self.result_cache.lock().insert(
4233 key,
4234 CachedEntry {
4235 data: CachedData::Columns(Arc::new(cols.clone())),
4236 footprint,
4237 condition_cols,
4238 },
4239 );
4240 }
4241 Ok(res)
4242 }
4243
4244 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4249 if q.conditions.is_empty() {
4250 return self.query(q);
4251 }
4252 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4253 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4254 crate::trace::QueryTrace::record(|t| {
4255 t.result_cache_hit = true;
4256 t.scan_mode = crate::trace::ScanMode::Materialized;
4257 });
4258 return Ok((*hit).clone());
4259 }
4260 let rows = self.query(q)?;
4261 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4262 let condition_cols = crate::query::condition_columns(&q.conditions);
4263 self.result_cache.lock().insert(
4264 key,
4265 CachedEntry {
4266 data: CachedData::Rows(Arc::new(rows.clone())),
4267 footprint,
4268 condition_cols,
4269 },
4270 );
4271 Ok(rows)
4272 }
4273
4274 #[allow(clippy::type_complexity)]
4289 pub fn query_columns_native_traced(
4290 &mut self,
4291 conditions: &[crate::query::Condition],
4292 projection: Option<&[u16]>,
4293 snapshot: Snapshot,
4294 ) -> Result<(
4295 Option<Vec<(u16, columnar::NativeColumn)>>,
4296 crate::trace::QueryTrace,
4297 )> {
4298 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4299 self.query_columns_native(conditions, projection, snapshot)
4300 });
4301 Ok((result?, trace))
4302 }
4303
4304 #[allow(clippy::type_complexity)]
4307 pub fn query_columns_native_cached_traced(
4308 &mut self,
4309 conditions: &[crate::query::Condition],
4310 projection: Option<&[u16]>,
4311 snapshot: Snapshot,
4312 ) -> Result<(
4313 Option<Vec<(u16, columnar::NativeColumn)>>,
4314 crate::trace::QueryTrace,
4315 )> {
4316 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4317 self.query_columns_native_cached(conditions, projection, snapshot)
4318 });
4319 Ok((result?, trace))
4320 }
4321
4322 pub fn native_page_cursor_traced(
4324 &self,
4325 snapshot: Snapshot,
4326 projection: Vec<(u16, TypeId)>,
4327 conditions: &[crate::query::Condition],
4328 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4329 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4330 self.native_page_cursor(snapshot, projection, conditions)
4331 });
4332 Ok((result?, trace))
4333 }
4334
4335 pub fn native_multi_run_cursor_traced(
4337 &self,
4338 snapshot: Snapshot,
4339 projection: Vec<(u16, TypeId)>,
4340 conditions: &[crate::query::Condition],
4341 ) -> Result<(
4342 Option<crate::cursor::MultiRunCursor>,
4343 crate::trace::QueryTrace,
4344 )> {
4345 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4346 self.native_multi_run_cursor(snapshot, projection, conditions)
4347 });
4348 Ok((result?, trace))
4349 }
4350
4351 pub fn count_conditions_traced(
4353 &mut self,
4354 conditions: &[crate::query::Condition],
4355 snapshot: Snapshot,
4356 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4357 let (result, trace) =
4358 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4359 Ok((result?, trace))
4360 }
4361
4362 pub fn query_traced(
4364 &mut self,
4365 q: &crate::query::Query,
4366 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4367 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4368 Ok((result?, trace))
4369 }
4370
4371 pub fn query_columns_native(
4376 &mut self,
4377 conditions: &[crate::query::Condition],
4378 projection: Option<&[u16]>,
4379 snapshot: Snapshot,
4380 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4381 use crate::query::Condition;
4382 if conditions.is_empty() {
4383 return Ok(None);
4384 }
4385 self.ensure_indexes_complete()?;
4386
4387 let served = |c: &Condition| {
4392 matches!(
4393 c,
4394 Condition::Pk(_)
4395 | Condition::BitmapEq { .. }
4396 | Condition::BitmapIn { .. }
4397 | Condition::FmContains { .. }
4398 | Condition::FmContainsAll { .. }
4399 | Condition::Ann { .. }
4400 | Condition::Range { .. }
4401 | Condition::RangeF64 { .. }
4402 | Condition::SparseMatch { .. }
4403 | Condition::MinHashSimilar { .. }
4404 | Condition::IsNull { .. }
4405 | Condition::IsNotNull { .. }
4406 )
4407 };
4408 if !conditions.iter().all(served) {
4409 return Ok(None);
4410 }
4411 let fast_path =
4412 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4413 crate::trace::QueryTrace::record(|t| {
4414 t.run_count = self.run_refs.len();
4415 t.memtable_rows = self.memtable.len();
4416 t.mutable_run_rows = self.mutable_run.len();
4417 t.conditions_pushed = conditions.len();
4418 t.learned_range_used = conditions.iter().any(|c| match c {
4419 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4420 self.learned_range.contains_key(column_id)
4421 }
4422 _ => false,
4423 });
4424 });
4425 let col_ids: Vec<u16> = projection
4427 .map(|p| p.to_vec())
4428 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4429 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4430 .iter()
4431 .map(|&cid| {
4432 let ty = self
4433 .schema
4434 .columns
4435 .iter()
4436 .find(|c| c.id == cid)
4437 .map(|c| c.ty)
4438 .unwrap_or(TypeId::Bytes);
4439 (cid, ty)
4440 })
4441 .collect();
4442
4443 if fast_path {
4449 let needs_column = conditions.iter().any(|c| match c {
4452 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4453 Condition::RangeF64 { column_id, .. } => {
4454 !self.learned_range.contains_key(column_id)
4455 }
4456 _ => false,
4457 });
4458 let mut reader_opt: Option<RunReader> = if needs_column {
4459 Some(self.open_reader(self.run_refs[0].run_id)?)
4460 } else {
4461 None
4462 };
4463 let mut sets: Vec<RowIdSet> = Vec::new();
4464 for c in conditions {
4465 let s = match c {
4466 Condition::Range { column_id, lo, hi }
4467 if !self.learned_range.contains_key(column_id) =>
4468 {
4469 if reader_opt.is_none() {
4470 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4471 }
4472 reader_opt
4473 .as_mut()
4474 .expect("reader opened for range")
4475 .range_row_id_set_i64(*column_id, *lo, *hi)?
4476 }
4477 Condition::RangeF64 {
4478 column_id,
4479 lo,
4480 lo_inclusive,
4481 hi,
4482 hi_inclusive,
4483 } if !self.learned_range.contains_key(column_id) => {
4484 if reader_opt.is_none() {
4485 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4486 }
4487 reader_opt
4488 .as_mut()
4489 .expect("reader opened for range")
4490 .range_row_id_set_f64(
4491 *column_id,
4492 *lo,
4493 *lo_inclusive,
4494 *hi,
4495 *hi_inclusive,
4496 )?
4497 }
4498 _ => self.resolve_condition(c, snapshot)?,
4499 };
4500 sets.push(s);
4501 }
4502 let candidates = RowIdSet::intersect_many(sets);
4503 crate::trace::QueryTrace::record(|t| {
4504 t.survivor_count = Some(candidates.len());
4505 });
4506 if candidates.is_empty() {
4507 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4508 .iter()
4509 .map(|&id| {
4510 (
4511 id,
4512 columnar::null_native(
4513 proj_pairs
4514 .iter()
4515 .find(|(c, _)| c == &id)
4516 .map(|(_, t)| *t)
4517 .unwrap_or(TypeId::Bytes),
4518 0,
4519 ),
4520 )
4521 })
4522 .collect();
4523 return Ok(Some(cols));
4524 }
4525 let mut reader = match reader_opt.take() {
4526 Some(r) => r,
4527 None => self.open_reader(self.run_refs[0].run_id)?,
4528 };
4529 let candidate_ids = candidates.into_sorted_vec();
4530 let (positions, fast_rid) = if let Some(positions) =
4531 reader.positions_for_row_ids_fast(&candidate_ids)
4532 {
4533 (positions, true)
4534 } else {
4535 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4536 match col {
4537 columnar::NativeColumn::Int64 { data, .. } => {
4538 let mut p: Vec<usize> = candidate_ids
4539 .iter()
4540 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4541 .collect();
4542 p.sort_unstable();
4543 (p, false)
4544 }
4545 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4546 }
4547 };
4548 crate::trace::QueryTrace::record(|t| {
4549 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4550 t.fast_row_id_map = fast_rid;
4551 });
4552 let mut cols = Vec::with_capacity(col_ids.len());
4553 for cid in &col_ids {
4554 let col = reader.column_native(*cid)?;
4555 cols.push((*cid, col.gather(&positions)));
4556 }
4557 return Ok(Some(cols));
4558 }
4559
4560 if !self.run_refs.is_empty() {
4573 use crate::cursor::{drain_cursor_to_columns, Cursor};
4574 let remaining: usize;
4575 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4576 let c = self
4577 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4578 .expect("single-run cursor should build when run_refs.len() == 1");
4579 remaining = c.remaining_rows();
4580 Box::new(c)
4581 } else {
4582 let c = self
4583 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4584 .expect("multi-run cursor should build when run_refs.len() >= 1");
4585 remaining = c.remaining_rows();
4586 Box::new(c)
4587 };
4588 crate::trace::QueryTrace::record(|t| {
4589 if t.survivor_count.is_none() {
4590 t.survivor_count = Some(remaining);
4591 }
4592 });
4593 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4594 return Ok(Some(cols));
4595 }
4596
4597 crate::trace::QueryTrace::record(|t| {
4602 t.scan_mode = crate::trace::ScanMode::Materialized;
4603 t.row_materialized = true;
4604 });
4605 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4606 for c in conditions {
4607 sets.push(self.resolve_condition(c, snapshot)?);
4608 }
4609 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4610 let rows = self.rows_for_rids(&rids, snapshot)?;
4611 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4612 for (cid, ty) in &proj_pairs {
4613 let vals: Vec<Value> = rows
4614 .iter()
4615 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4616 .collect();
4617 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4618 }
4619 Ok(Some(cols))
4620 }
4621
4622 pub fn native_page_cursor(
4637 &self,
4638 snapshot: Snapshot,
4639 projection: Vec<(u16, TypeId)>,
4640 conditions: &[crate::query::Condition],
4641 ) -> Result<Option<NativePageCursor>> {
4642 use crate::cursor::build_page_plans;
4643 if !conditions.is_empty() && !self.indexes_complete {
4646 return Ok(None);
4647 }
4648 if self.run_refs.len() != 1 {
4649 return Ok(None);
4650 }
4651 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4652 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4653
4654 let overlay_rids: HashSet<u64> = {
4657 let mut s = HashSet::new();
4658 for row in self.memtable.visible_versions(snapshot.epoch) {
4659 s.insert(row.row_id.0);
4660 }
4661 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4662 s.insert(row.row_id.0);
4663 }
4664 s
4665 };
4666
4667 let survivors = if conditions.is_empty() {
4671 None
4672 } else {
4673 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4674 };
4675
4676 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4683 survivors.clone()
4684 } else if let Some(s) = &survivors {
4685 let mut run_set = s.clone();
4686 run_set.remove_many(overlay_rids.iter().copied());
4687 Some(run_set)
4688 } else {
4689 Some(RowIdSet::from_unsorted(
4690 rids.iter()
4691 .map(|&r| r as u64)
4692 .filter(|r| !overlay_rids.contains(r))
4693 .collect(),
4694 ))
4695 };
4696
4697 let overlay_rows = if overlay_rids.is_empty() {
4698 Vec::new()
4699 } else {
4700 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4701 self.overlay_visible_rows(snapshot, bound)
4702 };
4703
4704 let plans = if positions.is_empty() {
4706 Vec::new()
4707 } else {
4708 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4709 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4710 };
4711
4712 let overlay = if overlay_rows.is_empty() {
4714 None
4715 } else {
4716 let filtered =
4717 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4718 if filtered.is_empty() {
4719 None
4720 } else {
4721 Some(self.materialize_overlay(&filtered, &projection))
4722 }
4723 };
4724
4725 let overlay_row_count = overlay
4726 .as_ref()
4727 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4728 .unwrap_or(0);
4729 crate::trace::QueryTrace::record(|t| {
4730 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4731 t.run_count = self.run_refs.len();
4732 t.memtable_rows = self.memtable.len();
4733 t.mutable_run_rows = self.mutable_run.len();
4734 t.overlay_rows = overlay_row_count;
4735 t.conditions_pushed = conditions.len();
4736 t.pages_decoded = plans
4737 .iter()
4738 .map(|p| p.positions.len())
4739 .sum::<usize>()
4740 .min(1);
4741 });
4742
4743 Ok(Some(NativePageCursor::new_with_overlay(
4744 reader, projection, plans, overlay,
4745 )))
4746 }
4747 #[allow(clippy::type_complexity)]
4757 pub fn native_multi_run_cursor(
4758 &self,
4759 snapshot: Snapshot,
4760 projection: Vec<(u16, TypeId)>,
4761 conditions: &[crate::query::Condition],
4762 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4763 use crate::cursor::{MultiRunCursor, RunStream};
4764 use crate::sorted_run::SYS_ROW_ID;
4765 use std::collections::{BinaryHeap, HashMap, HashSet};
4766 if !conditions.is_empty() && !self.indexes_complete {
4769 return Ok(None);
4770 }
4771 if self.run_refs.is_empty() {
4772 return Ok(None);
4773 }
4774
4775 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
4777 Vec::with_capacity(self.run_refs.len());
4778 for rr in &self.run_refs {
4779 let mut reader = self.open_reader(rr.run_id)?;
4780 let (rids, eps, del) = reader.system_columns_native()?;
4781 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
4782 run_meta.push((reader, rids, eps, del, page_rows));
4783 }
4784
4785 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
4789 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
4790 for i in 0..rids.len() {
4791 let rid = rids[i] as u64;
4792 let e = eps[i] as u64;
4793 if e > snapshot.epoch.0 {
4794 continue;
4795 }
4796 let is_del = del[i] != 0;
4797 best.entry(rid)
4798 .and_modify(|cur| {
4799 if e > cur.0 {
4800 *cur = (e, run_idx, i, is_del);
4801 }
4802 })
4803 .or_insert((e, run_idx, i, is_del));
4804 }
4805 }
4806
4807 let overlay_rids: HashSet<u64> = {
4809 let mut s = HashSet::new();
4810 for row in self.memtable.visible_versions(snapshot.epoch) {
4811 s.insert(row.row_id.0);
4812 }
4813 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4814 s.insert(row.row_id.0);
4815 }
4816 s
4817 };
4818
4819 let survivors: Option<RowIdSet> = if conditions.is_empty() {
4821 None
4822 } else {
4823 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4824 for c in conditions {
4825 sets.push(self.resolve_condition(c, snapshot)?);
4826 }
4827 Some(RowIdSet::intersect_many(sets))
4828 };
4829
4830 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
4834 for (rid, (_, run_idx, pos, deleted)) in &best {
4835 if *deleted {
4836 continue;
4837 }
4838 if overlay_rids.contains(rid) {
4839 continue;
4840 }
4841 if let Some(s) = &survivors {
4842 if !s.contains(*rid) {
4843 continue;
4844 }
4845 }
4846 per_run[*run_idx].push((*rid, *pos));
4847 }
4848 for v in per_run.iter_mut() {
4849 v.sort_unstable_by_key(|&(rid, _)| rid);
4850 }
4851
4852 let mut streams = Vec::with_capacity(run_meta.len());
4854 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
4855 let mut total = 0usize;
4856 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
4857 let mut starts = Vec::with_capacity(page_rows.len());
4858 let mut acc = 0usize;
4859 for &r in &page_rows {
4860 starts.push(acc);
4861 acc += r;
4862 }
4863 let mut survivors_vec: Vec<(u64, usize, usize)> =
4864 Vec::with_capacity(per_run[run_idx].len());
4865 for &(rid, pos) in &per_run[run_idx] {
4866 let page_seq = match starts.partition_point(|&s| s <= pos) {
4867 0 => continue,
4868 p => p - 1,
4869 };
4870 let within = pos - starts[page_seq];
4871 survivors_vec.push((rid, page_seq, within));
4872 }
4873 total += survivors_vec.len();
4874 if let Some(&(rid, _, _)) = survivors_vec.first() {
4875 heap.push(std::cmp::Reverse((rid, run_idx)));
4876 }
4877 streams.push(RunStream::new(reader, survivors_vec, page_rows));
4878 }
4879
4880 let overlay_rows = if overlay_rids.is_empty() {
4882 Vec::new()
4883 } else {
4884 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4885 self.overlay_visible_rows(snapshot, bound)
4886 };
4887 let overlay = if overlay_rows.is_empty() {
4888 None
4889 } else {
4890 let filtered =
4891 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4892 if filtered.is_empty() {
4893 None
4894 } else {
4895 Some(self.materialize_overlay(&filtered, &projection))
4896 }
4897 };
4898
4899 let overlay_row_count = overlay
4900 .as_ref()
4901 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4902 .unwrap_or(0);
4903 crate::trace::QueryTrace::record(|t| {
4904 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
4905 t.run_count = self.run_refs.len();
4906 t.memtable_rows = self.memtable.len();
4907 t.mutable_run_rows = self.mutable_run.len();
4908 t.overlay_rows = overlay_row_count;
4909 t.conditions_pushed = conditions.len();
4910 t.survivor_count = Some(total);
4911 });
4912
4913 Ok(Some(MultiRunCursor::new(
4914 streams, projection, heap, total, overlay,
4915 )))
4916 }
4917
4918 fn overlay_materialization_bound<'a>(
4930 conditions: &[crate::query::Condition],
4931 survivors: &'a Option<RowIdSet>,
4932 ) -> Option<&'a RowIdSet> {
4933 use crate::query::Condition;
4934 let has_range = conditions
4935 .iter()
4936 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
4937 if has_range {
4938 None
4939 } else {
4940 survivors.as_ref()
4941 }
4942 }
4943
4944 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
4956 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4957 let mut fold = |row: Row| {
4958 if let Some(b) = bound {
4959 if !b.contains(row.row_id.0) {
4960 return;
4961 }
4962 }
4963 best.entry(row.row_id.0)
4964 .and_modify(|(be, br)| {
4965 if row.committed_epoch > *be {
4966 *be = row.committed_epoch;
4967 *br = row.clone();
4968 }
4969 })
4970 .or_insert_with(|| (row.committed_epoch, row));
4971 };
4972 for row in self.memtable.visible_versions(snapshot.epoch) {
4973 fold(row);
4974 }
4975 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4976 fold(row);
4977 }
4978 let mut out: Vec<Row> = best
4979 .into_values()
4980 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
4981 .collect();
4982 out.sort_by_key(|r| r.row_id);
4983 out
4984 }
4985
4986 fn filter_overlay_rows(
4994 &self,
4995 rows: Vec<Row>,
4996 conditions: &[crate::query::Condition],
4997 survivors: Option<&RowIdSet>,
4998 snapshot: Snapshot,
4999 ) -> Result<Vec<Row>> {
5000 if conditions.is_empty() {
5001 return Ok(rows);
5002 }
5003 use crate::query::Condition;
5004 let all_index_served = !conditions
5008 .iter()
5009 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5010 if all_index_served {
5011 return Ok(rows
5012 .into_iter()
5013 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5014 .collect());
5015 }
5016 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5019 for c in conditions {
5020 let s = match c {
5021 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5022 _ => self.resolve_condition(c, snapshot)?,
5023 };
5024 per_cond_sets.push(s);
5025 }
5026 Ok(rows
5027 .into_iter()
5028 .filter(|row| {
5029 conditions.iter().enumerate().all(|(i, c)| match c {
5030 Condition::Range { column_id, lo, hi } => {
5031 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5032 }
5033 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5034 match row.columns.get(column_id) {
5035 Some(Value::Float64(v)) => {
5036 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5037 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5038 lo_ok && hi_ok
5039 }
5040 _ => false,
5041 }
5042 }
5043 _ => per_cond_sets[i].contains(row.row_id.0),
5044 })
5045 })
5046 .collect())
5047 }
5048
5049 fn materialize_overlay(
5052 &self,
5053 rows: &[Row],
5054 projection: &[(u16, TypeId)],
5055 ) -> Vec<columnar::NativeColumn> {
5056 if projection.is_empty() {
5057 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5058 }
5059 let mut cols = Vec::with_capacity(projection.len());
5060 for (cid, ty) in projection {
5061 let vals: Vec<Value> = rows
5062 .iter()
5063 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5064 .collect();
5065 cols.push(columnar::values_to_native(*ty, &vals));
5066 }
5067 cols
5068 }
5069
5070 fn resolve_survivor_rids(
5075 &self,
5076 conditions: &[crate::query::Condition],
5077 reader: &mut RunReader,
5078 ) -> Result<RowIdSet> {
5079 use crate::query::Condition;
5080 let mut sets: Vec<RowIdSet> = Vec::new();
5081 for c in conditions {
5082 let s: RowIdSet = match c {
5083 Condition::Pk(key) => {
5084 let lookup = self
5085 .schema
5086 .primary_key()
5087 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5088 .unwrap_or_else(|| key.clone());
5089 self.hot
5090 .get(&lookup)
5091 .map(|r| RowIdSet::one(r.0))
5092 .unwrap_or_else(RowIdSet::empty)
5093 }
5094 Condition::BitmapEq { column_id, value } => {
5095 let lookup = self.index_lookup_key_bytes(*column_id, value);
5096 self.bitmap
5097 .get(column_id)
5098 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5099 .unwrap_or_else(RowIdSet::empty)
5100 }
5101 Condition::BitmapIn { column_id, values } => {
5102 let bm = self.bitmap.get(column_id);
5103 let mut acc = roaring::RoaringBitmap::new();
5104 if let Some(b) = bm {
5105 for v in values {
5106 let lookup = self.index_lookup_key_bytes(*column_id, v);
5107 acc |= b.get(&lookup);
5108 }
5109 }
5110 RowIdSet::from_roaring(acc)
5111 }
5112 Condition::FmContains { column_id, pattern } => self
5113 .fm
5114 .get(column_id)
5115 .map(|f| {
5116 RowIdSet::from_unsorted(
5117 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5118 )
5119 })
5120 .unwrap_or_else(RowIdSet::empty),
5121 Condition::FmContainsAll {
5122 column_id,
5123 patterns,
5124 } => {
5125 if let Some(f) = self.fm.get(column_id) {
5126 let sets: Vec<RowIdSet> = patterns
5127 .iter()
5128 .map(|pat| {
5129 RowIdSet::from_unsorted(
5130 f.locate(pat).into_iter().map(|r| r.0).collect(),
5131 )
5132 })
5133 .collect();
5134 RowIdSet::intersect_many(sets)
5135 } else {
5136 RowIdSet::empty()
5137 }
5138 }
5139 Condition::Ann {
5140 column_id,
5141 query,
5142 k,
5143 } => self
5144 .ann
5145 .get(column_id)
5146 .map(|a| {
5147 RowIdSet::from_unsorted(
5148 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5149 )
5150 })
5151 .unwrap_or_else(RowIdSet::empty),
5152 Condition::SparseMatch {
5153 column_id,
5154 query,
5155 k,
5156 } => self
5157 .sparse
5158 .get(column_id)
5159 .map(|s| {
5160 RowIdSet::from_unsorted(
5161 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5162 )
5163 })
5164 .unwrap_or_else(RowIdSet::empty),
5165 Condition::MinHashSimilar {
5166 column_id,
5167 query,
5168 k,
5169 } => self
5170 .minhash
5171 .get(column_id)
5172 .map(|mh| {
5173 RowIdSet::from_unsorted(
5174 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5175 )
5176 })
5177 .unwrap_or_else(RowIdSet::empty),
5178 Condition::Range { column_id, lo, hi } => {
5179 if let Some(li) = self.learned_range.get(column_id) {
5180 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5181 } else {
5182 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5183 }
5184 }
5185 Condition::RangeF64 {
5186 column_id,
5187 lo,
5188 lo_inclusive,
5189 hi,
5190 hi_inclusive,
5191 } => {
5192 if let Some(li) = self.learned_range.get(column_id) {
5193 RowIdSet::from_unsorted(
5194 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5195 .into_iter()
5196 .collect(),
5197 )
5198 } else {
5199 reader.range_row_id_set_f64(
5200 *column_id,
5201 *lo,
5202 *lo_inclusive,
5203 *hi,
5204 *hi_inclusive,
5205 )?
5206 }
5207 }
5208 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5209 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5210 };
5211 sets.push(s);
5212 }
5213 Ok(RowIdSet::intersect_many(sets))
5214 }
5215
5216 pub fn scan_cursor(
5237 &self,
5238 snapshot: Snapshot,
5239 projection: Vec<(u16, TypeId)>,
5240 conditions: &[crate::query::Condition],
5241 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5242 if !conditions.is_empty() && !self.indexes_complete {
5248 return Ok(None);
5249 }
5250 if self.run_refs.len() == 1 {
5251 Ok(self
5252 .native_page_cursor(snapshot, projection, conditions)?
5253 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5254 } else {
5255 Ok(self
5256 .native_multi_run_cursor(snapshot, projection, conditions)?
5257 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5258 }
5259 }
5260
5261 pub fn aggregate_native(
5275 &self,
5276 snapshot: Snapshot,
5277 column: Option<u16>,
5278 conditions: &[crate::query::Condition],
5279 agg: NativeAgg,
5280 ) -> Result<Option<NativeAggResult>> {
5281 if self.run_refs.len() == 1 && conditions.is_empty() {
5283 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5284 return Ok(Some(res));
5285 }
5286 }
5287 if matches!(agg, NativeAgg::Count) && column.is_none() {
5289 return Ok(self
5290 .scan_cursor(snapshot, Vec::new(), conditions)?
5291 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5292 }
5293 let cid = match column {
5296 Some(c) => c,
5297 None => return Ok(None),
5298 };
5299 let ty = self.column_type(cid);
5300 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5301 return Ok(None);
5302 };
5303 match ty {
5304 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5305 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5306 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5307 }
5308 TypeId::Float64 => {
5309 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5310 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5311 }
5312 _ => Ok(None),
5313 }
5314 }
5315
5316 fn aggregate_from_stats(
5324 &self,
5325 snapshot: Snapshot,
5326 column: Option<u16>,
5327 agg: NativeAgg,
5328 ) -> Result<Option<NativeAggResult>> {
5329 let cid = match (agg, column) {
5330 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5331 _ => return Ok(None), };
5333 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5334 return Ok(None);
5335 };
5336 let Some(cs) = stats.get(&cid) else {
5337 return Ok(None);
5338 };
5339 match agg {
5340 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5342 self.live_count.saturating_sub(cs.null_count),
5343 ))),
5344 NativeAgg::Min | NativeAgg::Max => {
5345 let bound = if agg == NativeAgg::Min {
5346 &cs.min
5347 } else {
5348 &cs.max
5349 };
5350 match bound {
5351 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5352 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5353 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5358 None => Ok(None),
5359 }
5360 }
5361 _ => Ok(None),
5362 }
5363 }
5364
5365 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5374 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5375 return Ok(None);
5376 }
5377 self.ensure_indexes_complete()?;
5380 let reader = self.open_reader(self.run_refs[0].run_id)?;
5381 if self.live_count != reader.row_count() as u64 {
5382 return Ok(None);
5383 }
5384 let Some(bm) = self.bitmap.get(&column_id) else {
5385 return Ok(None); };
5387 let mut distinct = bm.value_count() as u64;
5388 if !bm.get(&Value::Null.encode_key()).is_empty() {
5391 distinct = distinct.saturating_sub(1);
5392 }
5393 Ok(Some(distinct))
5394 }
5395
5396 pub fn aggregate_incremental(
5408 &mut self,
5409 cache_key: u64,
5410 conditions: &[crate::query::Condition],
5411 column: Option<u16>,
5412 agg: NativeAgg,
5413 ) -> Result<IncrementalAggResult> {
5414 let snap = self.snapshot();
5415 let cur_wm = self.allocator.current().0;
5416 let cur_epoch = snap.epoch.0;
5417 let incremental_ok =
5424 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5425
5426 if incremental_ok {
5429 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5430 if cached.epoch == cur_epoch {
5431 return Ok(IncrementalAggResult {
5432 state: cached.state,
5433 incremental: true,
5434 delta_rows: 0,
5435 });
5436 }
5437 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5438 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5439 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5440 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5441 let delta_state = agg_state_from_rows(
5442 &delta_rows,
5443 conditions,
5444 &index_sets,
5445 column,
5446 agg,
5447 &self.schema,
5448 )?;
5449 let merged = cached.state.merge(delta_state);
5450 let delta_n = delta_rids.len() as u64;
5451 self.agg_cache.insert(
5452 cache_key,
5453 CachedAgg {
5454 state: merged.clone(),
5455 watermark: cur_wm,
5456 epoch: cur_epoch,
5457 },
5458 );
5459 return Ok(IncrementalAggResult {
5460 state: merged,
5461 incremental: true,
5462 delta_rows: delta_n,
5463 });
5464 }
5465 }
5466 }
5467
5468 let cursor_ok =
5473 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5474 let state = if cursor_ok && agg != NativeAgg::Avg {
5475 match self.aggregate_native(snap, column, conditions, agg)? {
5476 Some(result) => {
5477 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5478 }
5479 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5480 }
5481 } else {
5482 self.agg_state_full_scan(conditions, column, agg, snap)?
5483 };
5484 if incremental_ok {
5486 self.agg_cache.insert(
5487 cache_key,
5488 CachedAgg {
5489 state: state.clone(),
5490 watermark: cur_wm,
5491 epoch: cur_epoch,
5492 },
5493 );
5494 }
5495 Ok(IncrementalAggResult {
5496 state,
5497 incremental: false,
5498 delta_rows: 0,
5499 })
5500 }
5501
5502 fn agg_state_full_scan(
5505 &self,
5506 conditions: &[crate::query::Condition],
5507 column: Option<u16>,
5508 agg: NativeAgg,
5509 snap: Snapshot,
5510 ) -> Result<AggState> {
5511 let rows = self.visible_rows(snap)?;
5512 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5513 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5514 }
5515
5516 fn resolve_index_conditions(
5519 &self,
5520 conditions: &[crate::query::Condition],
5521 snapshot: Snapshot,
5522 ) -> Result<Vec<RowIdSet>> {
5523 use crate::query::Condition;
5524 let mut sets = Vec::new();
5525 for c in conditions {
5526 if matches!(
5527 c,
5528 Condition::Ann { .. }
5529 | Condition::SparseMatch { .. }
5530 | Condition::MinHashSimilar { .. }
5531 ) {
5532 sets.push(self.resolve_condition(c, snapshot)?);
5533 }
5534 }
5535 Ok(sets)
5536 }
5537
5538 fn column_type(&self, cid: u16) -> TypeId {
5539 self.schema
5540 .columns
5541 .iter()
5542 .find(|c| c.id == cid)
5543 .map(|c| c.ty)
5544 .unwrap_or(TypeId::Bytes)
5545 }
5546
5547 pub fn approx_aggregate(
5556 &mut self,
5557 conditions: &[crate::query::Condition],
5558 column: Option<u16>,
5559 agg: ApproxAgg,
5560 z: f64,
5561 ) -> Result<Option<ApproxResult>> {
5562 use crate::query::Condition;
5563 self.ensure_reservoir_complete()?;
5564 let snapshot = self.snapshot();
5565 let n_pop = self.live_count;
5566 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5567 if sample_rids.is_empty() {
5568 return Ok(None);
5569 }
5570 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5572 let s = live_sample.len();
5573 if s == 0 {
5574 return Ok(None);
5575 }
5576
5577 let mut index_sets: Vec<RowIdSet> = Vec::new();
5580 for c in conditions {
5581 if matches!(
5582 c,
5583 Condition::Ann { .. }
5584 | Condition::SparseMatch { .. }
5585 | Condition::MinHashSimilar { .. }
5586 ) {
5587 index_sets.push(self.resolve_condition(c, snapshot)?);
5588 }
5589 }
5590
5591 let cid = match (agg, column) {
5593 (ApproxAgg::Count, _) => None,
5594 (_, Some(c)) => Some(c),
5595 _ => return Ok(None),
5596 };
5597 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5598 for r in &live_sample {
5599 if !conditions
5601 .iter()
5602 .all(|c| condition_matches_row(c, r, &self.schema))
5603 {
5604 continue;
5605 }
5606 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5608 continue;
5609 }
5610 if let Some(cid) = cid {
5611 if let Some(v) = as_f64(r.columns.get(&cid)) {
5612 passing_vals.push(v);
5613 } } else {
5615 passing_vals.push(0.0); }
5617 }
5618 let m = passing_vals.len();
5619
5620 let (point, half) = match agg {
5621 ApproxAgg::Count => {
5622 let p = m as f64 / s as f64;
5624 let point = n_pop as f64 * p;
5625 let var = if s > 1 {
5626 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5627 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5628 } else {
5629 0.0
5630 };
5631 (point, z * var.sqrt())
5632 }
5633 ApproxAgg::Sum => {
5634 let y: Vec<f64> = live_sample
5636 .iter()
5637 .map(|r| {
5638 let passes_row = conditions
5639 .iter()
5640 .all(|c| condition_matches_row(c, r, &self.schema))
5641 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5642 if passes_row {
5643 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5644 } else {
5645 0.0
5646 }
5647 })
5648 .collect();
5649 let mean_y = y.iter().sum::<f64>() / s as f64;
5650 let point = n_pop as f64 * mean_y;
5651 let var = if s > 1 {
5652 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5653 let var_y = ss / (s - 1) as f64;
5654 n_pop as f64 * n_pop as f64 * var_y / s as f64
5655 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5656 } else {
5657 0.0
5658 };
5659 (point, z * var.sqrt())
5660 }
5661 ApproxAgg::Avg => {
5662 if m == 0 {
5663 return Ok(Some(ApproxResult {
5664 point: 0.0,
5665 ci_low: 0.0,
5666 ci_high: 0.0,
5667 n_population: n_pop,
5668 n_sample_live: s,
5669 n_passing: 0,
5670 }));
5671 }
5672 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5673 let half = if m > 1 {
5674 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5675 let sd = (ss / (m - 1) as f64).sqrt();
5676 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5677 z * sd / (m as f64).sqrt() * fpc.sqrt()
5678 } else {
5679 0.0
5680 };
5681 (mean, half)
5682 }
5683 };
5684
5685 Ok(Some(ApproxResult {
5686 point,
5687 ci_low: point - half,
5688 ci_high: point + half,
5689 n_population: n_pop,
5690 n_sample_live: s,
5691 n_passing: m,
5692 }))
5693 }
5694
5695 pub fn exact_column_stats(
5703 &self,
5704 _snapshot: Snapshot,
5705 projection: &[u16],
5706 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5707 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5708 return Ok(None);
5709 }
5710 let reader = self.open_reader(self.run_refs[0].run_id)?;
5711 if self.live_count != reader.row_count() as u64 {
5712 return Ok(None);
5713 }
5714 let mut out = HashMap::new();
5715 for &cid in projection {
5716 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5717 Some(c) => c,
5718 None => continue,
5719 };
5720 let Some(stats) = reader.column_page_stats(cid) else {
5722 out.insert(
5723 cid,
5724 ColumnStat {
5725 min: None,
5726 max: None,
5727 null_count: self.live_count,
5728 },
5729 );
5730 continue;
5731 };
5732 let stat = match cdef.ty {
5733 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5734 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5735 min: mn.map(Value::Int64),
5736 max: mx.map(Value::Int64),
5737 null_count: n,
5738 })
5739 }
5740 TypeId::Float64 => {
5741 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5742 min: mn.map(Value::Float64),
5743 max: mx.map(Value::Float64),
5744 null_count: n,
5745 })
5746 }
5747 _ => None,
5748 };
5749 if let Some(s) = stat {
5750 out.insert(cid, s);
5751 }
5752 }
5753 Ok(Some(out))
5754 }
5755
5756 pub fn dir(&self) -> &Path {
5757 &self.dir
5758 }
5759
5760 pub fn schema(&self) -> &Schema {
5761 &self.schema
5762 }
5763
5764 pub(crate) fn prepare_alter_column(
5765 &mut self,
5766 column_name: &str,
5767 change: &AlterColumn,
5768 ) -> Result<ColumnDef> {
5769 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
5770 return Err(MongrelError::InvalidArgument(
5771 "ALTER COLUMN requires committing staged writes first".into(),
5772 ));
5773 }
5774 let old = self
5775 .schema
5776 .columns
5777 .iter()
5778 .find(|c| c.name == column_name)
5779 .cloned()
5780 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5781 let mut next = old.clone();
5782
5783 if let Some(name) = &change.name {
5784 let trimmed = name.trim();
5785 if trimmed.is_empty() {
5786 return Err(MongrelError::InvalidArgument(
5787 "ALTER COLUMN name must not be empty".into(),
5788 ));
5789 }
5790 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
5791 return Err(MongrelError::Schema(format!(
5792 "column {trimmed} already exists"
5793 )));
5794 }
5795 next.name = trimmed.to_string();
5796 }
5797
5798 if let Some(ty) = change.ty {
5799 next.ty = ty;
5800 }
5801 if let Some(flags) = change.flags {
5802 validate_alter_column_flags(old.flags, flags)?;
5803 next.flags = flags;
5804 }
5805
5806 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
5807 if old.flags.contains(ColumnFlags::NULLABLE)
5808 && !next.flags.contains(ColumnFlags::NULLABLE)
5809 && self.column_has_nulls(old.id)?
5810 {
5811 return Err(MongrelError::InvalidArgument(format!(
5812 "column '{}' contains NULL values",
5813 old.name
5814 )));
5815 }
5816 Ok(next)
5817 }
5818
5819 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
5820 let idx = self
5821 .schema
5822 .columns
5823 .iter()
5824 .position(|c| c.id == column.id)
5825 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
5826 if self.schema.columns[idx] == column {
5827 return Ok(());
5828 }
5829 self.schema.columns[idx] = column;
5830 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5831 self.schema.validate_auto_increment()?;
5832 self.auto_inc = resolve_auto_inc(&self.schema);
5833 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
5834 write_schema(&self.dir, &self.schema)?;
5835 self.clear_result_cache();
5836 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5837 self.persist_manifest(self.current_epoch())?;
5838 Ok(())
5839 }
5840
5841 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
5842 let column = self.prepare_alter_column(column_name, &change)?;
5843 self.apply_altered_column(column.clone())?;
5844 Ok(column)
5845 }
5846
5847 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
5848 if self.live_count == 0 {
5849 return Ok(false);
5850 }
5851 let snap = self.snapshot();
5852 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
5853 Ok(columns
5854 .first()
5855 .map(|(_, col)| col.null_count(col.len()) != 0)
5856 .unwrap_or(true))
5857 }
5858
5859 fn has_stored_versions(&self) -> bool {
5860 !self.memtable.is_empty()
5861 || !self.mutable_run.is_empty()
5862 || self.run_refs.iter().any(|r| r.row_count > 0)
5863 || !self.retiring.is_empty()
5864 }
5865
5866 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
5871 if self.schema.columns.iter().any(|c| c.name == name) {
5872 return Err(MongrelError::Schema(format!(
5873 "column {name} already exists"
5874 )));
5875 }
5876 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
5877 self.schema.columns.push(ColumnDef {
5878 id,
5879 name: name.to_string(),
5880 ty,
5881 flags,
5882 });
5883 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5884 self.schema.validate_auto_increment()?;
5885 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
5886 self.auto_inc = resolve_auto_inc(&self.schema);
5887 }
5888 write_schema(&self.dir, &self.schema)?;
5889 self.clear_result_cache();
5890 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5892 self.persist_manifest(self.current_epoch())?;
5893 Ok(id)
5894 }
5895
5896 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
5905 let cid = self
5906 .schema
5907 .columns
5908 .iter()
5909 .find(|c| c.name == column_name)
5910 .map(|c| c.id)
5911 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5912 let ty = self
5913 .schema
5914 .columns
5915 .iter()
5916 .find(|c| c.id == cid)
5917 .map(|c| c.ty)
5918 .unwrap_or(TypeId::Int64);
5919 if !matches!(
5920 ty,
5921 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
5922 ) {
5923 return Err(MongrelError::Schema(format!(
5924 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
5925 )));
5926 }
5927 if self
5928 .schema
5929 .indexes
5930 .iter()
5931 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
5932 {
5933 return Ok(()); }
5935 self.schema.indexes.push(IndexDef {
5936 name: format!("{}_learned_range", column_name),
5937 column_id: cid,
5938 kind: IndexKind::LearnedRange,
5939 });
5940 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5941 write_schema(&self.dir, &self.schema)?;
5942 self.build_learned_ranges()?;
5943 Ok(())
5944 }
5945
5946 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
5949 self.sync_byte_threshold = threshold;
5950 if let WalSink::Private(w) = &mut self.wal {
5951 w.set_sync_byte_threshold(threshold);
5952 }
5953 }
5954
5955 pub fn page_cache_flush(&self) {
5959 self.page_cache.lock().flush_to_disk();
5960 }
5961
5962 pub fn page_cache_len(&self) -> usize {
5964 self.page_cache.lock().len()
5965 }
5966
5967 pub fn decoded_cache_len(&self) -> usize {
5970 self.decoded_cache.lock().len()
5971 }
5972
5973 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
5976 self.memtable.drain_sorted()
5977 }
5978
5979 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
5980 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
5981 }
5982
5983 pub(crate) fn table_dir(&self) -> &Path {
5984 &self.dir
5985 }
5986
5987 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
5988 &self.schema
5989 }
5990
5991 pub(crate) fn alloc_run_id(&mut self) -> u64 {
5992 let id = self.next_run_id;
5993 self.next_run_id += 1;
5994 id
5995 }
5996
5997 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
5998 self.run_refs.push(run_ref);
5999 }
6000
6001 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6011 self.retiring.push(crate::manifest::RetiredRun {
6012 run_id,
6013 retire_epoch,
6014 });
6015 }
6016
6017 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6021 if self.retiring.is_empty() {
6022 return Ok(0);
6023 }
6024 let mut reaped = 0;
6025 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6026 for r in std::mem::take(&mut self.retiring) {
6032 if min_active.0 >= r.retire_epoch {
6033 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6034 reaped += 1;
6035 } else {
6036 kept.push(r);
6037 }
6038 }
6039 self.retiring = kept;
6040 if reaped > 0 {
6041 self.persist_manifest(self.current_epoch())?;
6042 }
6043 Ok(reaped)
6044 }
6045
6046 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6047 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6048 return false;
6049 }
6050 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6051 self.run_refs.push(run_ref);
6052 self.indexes_complete = false;
6053 true
6054 }
6055
6056 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6057 self.kek.as_ref()
6058 }
6059
6060 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6061 let mut reader = RunReader::open_with_cache(
6062 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6063 self.schema.clone(),
6064 self.kek.clone(),
6065 Some(self.page_cache.clone()),
6066 Some(self.decoded_cache.clone()),
6067 self.table_id,
6068 Some(&self.verified_runs),
6069 )?;
6070 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6074 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6075 }
6076 Ok(reader)
6077 }
6078
6079 pub(crate) fn run_refs(&self) -> &[RunRef] {
6080 &self.run_refs
6081 }
6082
6083 pub(crate) fn runs_dir(&self) -> PathBuf {
6084 self.dir.join(RUNS_DIR)
6085 }
6086
6087 pub(crate) fn wal_dir(&self) -> PathBuf {
6088 self.dir.join(WAL_DIR)
6089 }
6090
6091 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6092 self.run_refs = refs;
6093 }
6094
6095 pub(crate) fn next_run_id(&self) -> u64 {
6096 self.next_run_id
6097 }
6098
6099 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6100 self.compaction_zstd_level
6101 }
6102
6103 pub(crate) fn bump_next_run_id(&mut self) {
6104 self.next_run_id += 1;
6105 }
6106
6107 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6108 self.kek.clone()
6109 }
6110
6111 #[cfg(feature = "encryption")]
6115 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6116 self.kek.as_ref().map(|k| k.derive_idx_key())
6117 }
6118
6119 #[cfg(not(feature = "encryption"))]
6120 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6121 None
6122 }
6123
6124 #[cfg(feature = "encryption")]
6128 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6129 self.kek.as_ref().map(|k| *k.derive_meta_key())
6130 }
6131
6132 #[cfg(not(feature = "encryption"))]
6133 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6134 None
6135 }
6136
6137 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6140 self.column_keys
6141 .iter()
6142 .map(|(&id, &(_, scheme))| (id, scheme))
6143 .collect()
6144 }
6145
6146 #[cfg(feature = "encryption")]
6151 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6152 self.tokenize_value_enc(column_id, v)
6153 }
6154
6155 #[cfg(feature = "encryption")]
6156 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6157 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6158 let (key, scheme) = self.column_keys.get(&column_id)?;
6159 let token: Vec<u8> = match (*scheme, v) {
6160 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6161 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6162 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6163 _ => hmac_token(key, &v.encode_key()).to_vec(),
6164 };
6165 Some(Value::Bytes(token))
6166 }
6167
6168 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6170 self.index_lookup_key_bytes(column_id, &v.encode_key())
6171 }
6172
6173 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6176 #[cfg(feature = "encryption")]
6177 {
6178 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6179 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6180 if *scheme == SCHEME_HMAC_EQ {
6181 return hmac_token(key, encoded).to_vec();
6182 }
6183 }
6184 }
6185 let _ = column_id;
6186 encoded.to_vec()
6187 }
6188}
6189
6190fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6191 let columnar::NativeColumn::Int64 { data, validity } = col else {
6192 return false;
6193 };
6194 if data.len() < n || !columnar::all_non_null(validity, n) {
6195 return false;
6196 }
6197 data.iter()
6198 .take(n)
6199 .zip(data.iter().skip(1))
6200 .all(|(a, b)| a < b)
6201}
6202
6203#[derive(Debug, Clone)]
6207pub struct ColumnStat {
6208 pub min: Option<Value>,
6209 pub max: Option<Value>,
6210 pub null_count: u64,
6211}
6212
6213#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6215pub enum NativeAgg {
6216 Count,
6217 Sum,
6218 Min,
6219 Max,
6220 Avg,
6221}
6222
6223#[derive(Debug, Clone, PartialEq)]
6225pub enum NativeAggResult {
6226 Count(u64),
6227 Int(i64),
6228 Float(f64),
6229 Null,
6231}
6232
6233#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6235pub enum ApproxAgg {
6236 Count,
6237 Sum,
6238 Avg,
6239}
6240
6241#[derive(Debug, Clone)]
6245pub struct ApproxResult {
6246 pub point: f64,
6248 pub ci_low: f64,
6250 pub ci_high: f64,
6252 pub n_population: u64,
6254 pub n_sample_live: usize,
6256 pub n_passing: usize,
6258}
6259
6260#[derive(Debug, Clone, PartialEq)]
6265pub enum AggState {
6266 Count(u64),
6268 SumI {
6270 sum: i128,
6271 count: u64,
6272 },
6273 SumF {
6275 sum: f64,
6276 count: u64,
6277 },
6278 AvgI {
6280 sum: i128,
6281 count: u64,
6282 },
6283 AvgF {
6285 sum: f64,
6286 count: u64,
6287 },
6288 MinI(i64),
6290 MaxI(i64),
6291 MinF(f64),
6293 MaxF(f64),
6294 Empty,
6296}
6297
6298impl AggState {
6299 pub fn merge(self, other: AggState) -> AggState {
6301 use AggState::*;
6302 match (self, other) {
6303 (Empty, x) | (x, Empty) => x,
6304 (Count(a), Count(b)) => Count(a + b),
6305 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6306 sum: sa + sb,
6307 count: ca + cb,
6308 },
6309 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6310 sum: sa + sb,
6311 count: ca + cb,
6312 },
6313 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6314 sum: sa + sb,
6315 count: ca + cb,
6316 },
6317 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6318 sum: sa + sb,
6319 count: ca + cb,
6320 },
6321 (MinI(a), MinI(b)) => MinI(a.min(b)),
6322 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6323 (MinF(a), MinF(b)) => MinF(a.min(b)),
6324 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6325 _ => Empty, }
6327 }
6328
6329 pub fn point(&self) -> Option<f64> {
6331 match self {
6332 AggState::Count(n) => Some(*n as f64),
6333 AggState::SumI { sum, .. } => Some(*sum as f64),
6334 AggState::SumF { sum, .. } => Some(*sum),
6335 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6336 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6337 AggState::MinI(n) => Some(*n as f64),
6338 AggState::MaxI(n) => Some(*n as f64),
6339 AggState::MinF(n) => Some(*n),
6340 AggState::MaxF(n) => Some(*n),
6341 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6342 }
6343 }
6344
6345 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6349 let is_float = matches!(ty, Some(TypeId::Float64));
6350 match (agg, result) {
6351 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6352 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6353 sum: x as i128,
6354 count: 1, },
6356 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6357 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6358 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6359 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6360 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6361 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6362 (NativeAgg::Count, _) => AggState::Empty,
6363 (_, NativeAggResult::Null) => AggState::Empty,
6364 _ => {
6365 let _ = is_float;
6366 AggState::Empty
6367 }
6368 }
6369 }
6370}
6371
6372#[derive(Debug, Clone)]
6375pub struct CachedAgg {
6376 pub state: AggState,
6377 pub watermark: u64,
6378 pub epoch: u64,
6379}
6380
6381#[derive(Debug, Clone)]
6383pub struct IncrementalAggResult {
6384 pub state: AggState,
6386 pub incremental: bool,
6389 pub delta_rows: u64,
6391}
6392
6393fn agg_state_from_rows(
6397 rows: &[Row],
6398 conditions: &[crate::query::Condition],
6399 index_sets: &[RowIdSet],
6400 column: Option<u16>,
6401 agg: NativeAgg,
6402 schema: &Schema,
6403) -> Result<AggState> {
6404 let mut count: u64 = 0;
6405 let mut sum_i: i128 = 0;
6406 let mut sum_f: f64 = 0.0;
6407 let mut mn_i: i64 = i64::MAX;
6408 let mut mx_i: i64 = i64::MIN;
6409 let mut mn_f: f64 = f64::INFINITY;
6410 let mut mx_f: f64 = f64::NEG_INFINITY;
6411 let mut saw_int = false;
6412 let mut saw_float = false;
6413 for r in rows {
6414 if !conditions
6415 .iter()
6416 .all(|c| condition_matches_row(c, r, schema))
6417 {
6418 continue;
6419 }
6420 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6421 continue;
6422 }
6423 match agg {
6424 NativeAgg::Count => match column {
6425 None => count += 1,
6427 Some(cid) => match r.columns.get(&cid) {
6430 None | Some(Value::Null) => {}
6431 Some(_) => count += 1,
6432 },
6433 },
6434 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6435 Some(Value::Int64(n)) => {
6436 count += 1;
6437 sum_i += *n as i128;
6438 mn_i = mn_i.min(*n);
6439 mx_i = mx_i.max(*n);
6440 saw_int = true;
6441 }
6442 Some(Value::Float64(f)) => {
6443 count += 1;
6444 sum_f += f;
6445 mn_f = mn_f.min(*f);
6446 mx_f = mx_f.max(*f);
6447 saw_float = true;
6448 }
6449 _ => {}
6450 },
6451 }
6452 }
6453 Ok(match agg {
6454 NativeAgg::Count => {
6455 if count == 0 {
6456 AggState::Empty
6457 } else {
6458 AggState::Count(count)
6459 }
6460 }
6461 NativeAgg::Sum => {
6462 if count == 0 {
6463 AggState::Empty
6464 } else if saw_int {
6465 AggState::SumI { sum: sum_i, count }
6466 } else {
6467 AggState::SumF { sum: sum_f, count }
6468 }
6469 }
6470 NativeAgg::Avg => {
6471 if count == 0 {
6472 AggState::Empty
6473 } else if saw_int {
6474 AggState::AvgI { sum: sum_i, count }
6475 } else {
6476 AggState::AvgF { sum: sum_f, count }
6477 }
6478 }
6479 NativeAgg::Min => {
6480 if !saw_int && !saw_float {
6481 AggState::Empty
6482 } else if saw_int {
6483 AggState::MinI(mn_i)
6484 } else {
6485 AggState::MinF(mn_f)
6486 }
6487 }
6488 NativeAgg::Max => {
6489 if !saw_int && !saw_float {
6490 AggState::Empty
6491 } else if saw_int {
6492 AggState::MaxI(mx_i)
6493 } else {
6494 AggState::MaxF(mx_f)
6495 }
6496 }
6497 })
6498}
6499
6500fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6504 use crate::query::Condition;
6505 match c {
6506 Condition::Pk(key) => match schema.primary_key() {
6507 Some(pk) => row
6508 .columns
6509 .get(&pk.id)
6510 .map(|v| v.encode_key() == *key)
6511 .unwrap_or(false),
6512 None => false,
6513 },
6514 Condition::BitmapEq { column_id, value } => row
6515 .columns
6516 .get(column_id)
6517 .map(|v| v.encode_key() == *value)
6518 .unwrap_or(false),
6519 Condition::BitmapIn { column_id, values } => {
6520 let key = row.columns.get(column_id).map(|v| v.encode_key());
6521 match key {
6522 Some(k) => values.contains(&k),
6523 None => false,
6524 }
6525 }
6526 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6527 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6528 _ => false,
6529 },
6530 Condition::RangeF64 {
6531 column_id,
6532 lo,
6533 lo_inclusive,
6534 hi,
6535 hi_inclusive,
6536 } => match row.columns.get(column_id) {
6537 Some(Value::Float64(n)) => {
6538 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6539 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6540 lo_ok && hi_ok
6541 }
6542 _ => false,
6543 },
6544 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6545 Some(Value::Bytes(b)) => {
6546 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6547 }
6548 _ => false,
6549 },
6550 Condition::FmContainsAll {
6551 column_id,
6552 patterns,
6553 } => match row.columns.get(column_id) {
6554 Some(Value::Bytes(b)) => patterns
6555 .iter()
6556 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6557 _ => false,
6558 },
6559 Condition::Ann { .. }
6560 | Condition::SparseMatch { .. }
6561 | Condition::MinHashSimilar { .. } => true,
6562 Condition::IsNull { column_id } => {
6563 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6564 }
6565 Condition::IsNotNull { column_id } => {
6566 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6567 }
6568 }
6569}
6570
6571fn as_f64(v: Option<&Value>) -> Option<f64> {
6573 match v {
6574 Some(Value::Int64(n)) => Some(*n as f64),
6575 Some(Value::Float64(f)) => Some(*f),
6576 _ => None,
6577 }
6578}
6579
6580fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6584 let mut count: u64 = 0;
6585 let mut sum: i128 = 0;
6586 let mut mn: i64 = i64::MAX;
6587 let mut mx: i64 = i64::MIN;
6588 while let Some(cols) = cursor.next_batch()? {
6589 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6590 if crate::columnar::all_non_null(validity, data.len()) {
6591 count += data.len() as u64;
6593 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6594 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6595 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6596 } else {
6597 for (i, &v) in data.iter().enumerate() {
6598 if crate::columnar::validity_bit(validity, i) {
6599 count += 1;
6600 sum += v as i128;
6601 mn = mn.min(v);
6602 mx = mx.max(v);
6603 }
6604 }
6605 }
6606 }
6607 }
6608 Ok((count, sum, mn, mx))
6609}
6610
6611fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6613 let mut count: u64 = 0;
6614 let mut sum: f64 = 0.0;
6615 let mut mn: f64 = f64::INFINITY;
6616 let mut mx: f64 = f64::NEG_INFINITY;
6617 while let Some(cols) = cursor.next_batch()? {
6618 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6619 if crate::columnar::all_non_null(validity, data.len()) {
6620 count += data.len() as u64;
6621 sum += data.iter().sum::<f64>();
6622 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6623 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6624 } else {
6625 for (i, &v) in data.iter().enumerate() {
6626 if crate::columnar::validity_bit(validity, i) {
6627 count += 1;
6628 sum += v;
6629 mn = mn.min(v);
6630 mx = mx.max(v);
6631 }
6632 }
6633 }
6634 }
6635 }
6636 Ok((count, sum, mn, mx))
6637}
6638
6639fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6640 if count == 0 && !matches!(agg, NativeAgg::Count) {
6641 return NativeAggResult::Null;
6642 }
6643 match agg {
6644 NativeAgg::Count => NativeAggResult::Count(count),
6645 NativeAgg::Sum => match sum.try_into() {
6648 Ok(v) => NativeAggResult::Int(v),
6649 Err(_) => NativeAggResult::Null,
6650 },
6651 NativeAgg::Min => NativeAggResult::Int(mn),
6652 NativeAgg::Max => NativeAggResult::Int(mx),
6653 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6654 }
6655}
6656
6657fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6658 if count == 0 && !matches!(agg, NativeAgg::Count) {
6659 return NativeAggResult::Null;
6660 }
6661 match agg {
6662 NativeAgg::Count => NativeAggResult::Count(count),
6663 NativeAgg::Sum => NativeAggResult::Float(sum),
6664 NativeAgg::Min => NativeAggResult::Float(mn),
6665 NativeAgg::Max => NativeAggResult::Float(mx),
6666 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6667 }
6668}
6669
6670fn agg_int(
6673 stats: &[crate::page::PageStat],
6674 decode: fn(Option<&[u8]>) -> Option<i64>,
6675) -> Option<(Option<i64>, Option<i64>, u64)> {
6676 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6677 let mut any = false;
6678 for s in stats {
6679 if let Some(v) = decode(s.min.as_deref()) {
6680 mn = mn.min(v);
6681 any = true;
6682 }
6683 if let Some(v) = decode(s.max.as_deref()) {
6684 mx = mx.max(v);
6685 any = true;
6686 }
6687 nulls += s.null_count;
6688 }
6689 any.then_some((Some(mn), Some(mx), nulls))
6690}
6691
6692fn agg_float(
6694 stats: &[crate::page::PageStat],
6695 decode: fn(Option<&[u8]>) -> Option<f64>,
6696) -> Option<(Option<f64>, Option<f64>, u64)> {
6697 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6698 let mut any = false;
6699 for s in stats {
6700 if let Some(v) = decode(s.min.as_deref()) {
6701 mn = mn.min(v);
6702 any = true;
6703 }
6704 if let Some(v) = decode(s.max.as_deref()) {
6705 mx = mx.max(v);
6706 any = true;
6707 }
6708 nulls += s.null_count;
6709 }
6710 any.then_some((Some(mn), Some(mx), nulls))
6711}
6712
6713type SecondaryIndexes = (
6715 HashMap<u16, BitmapIndex>,
6716 HashMap<u16, AnnIndex>,
6717 HashMap<u16, FmIndex>,
6718 HashMap<u16, SparseIndex>,
6719 HashMap<u16, MinHashIndex>,
6720);
6721
6722fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6723 let mut bitmap = HashMap::new();
6724 let mut ann = HashMap::new();
6725 let mut fm = HashMap::new();
6726 let mut sparse = HashMap::new();
6727 let mut minhash = HashMap::new();
6728 for idef in &schema.indexes {
6729 match idef.kind {
6730 IndexKind::Bitmap => {
6731 bitmap.insert(idef.column_id, BitmapIndex::new());
6732 }
6733 IndexKind::Ann => {
6734 let dim = schema
6735 .columns
6736 .iter()
6737 .find(|c| c.id == idef.column_id)
6738 .and_then(|c| match c.ty {
6739 TypeId::Embedding { dim } => Some(dim as usize),
6740 _ => None,
6741 })
6742 .unwrap_or(0);
6743 ann.insert(idef.column_id, AnnIndex::new(dim));
6744 }
6745 IndexKind::FmIndex => {
6746 fm.insert(idef.column_id, FmIndex::new());
6747 }
6748 IndexKind::Sparse => {
6749 sparse.insert(idef.column_id, SparseIndex::new());
6750 }
6751 IndexKind::MinHash => {
6752 minhash.insert(idef.column_id, MinHashIndex::new());
6753 }
6754 _ => {}
6755 }
6756 }
6757 (bitmap, ann, fm, sparse, minhash)
6758}
6759
6760const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
6761 | ColumnFlags::AUTO_INCREMENT
6762 | ColumnFlags::ENCRYPTED
6763 | ColumnFlags::ENCRYPTED_INDEXABLE
6764 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
6765
6766fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
6767 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
6768 return Err(MongrelError::Schema(
6769 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
6770 ));
6771 }
6772 Ok(())
6773}
6774
6775fn validate_alter_column_type(
6776 schema: &Schema,
6777 old: &ColumnDef,
6778 next: &ColumnDef,
6779 has_stored_versions: bool,
6780) -> Result<()> {
6781 if old.ty == next.ty {
6782 return Ok(());
6783 }
6784 if schema.indexes.iter().any(|i| i.column_id == old.id) {
6785 return Err(MongrelError::Schema(format!(
6786 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
6787 old.name
6788 )));
6789 }
6790 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
6791 return Ok(());
6792 }
6793 Err(MongrelError::Schema(format!(
6794 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
6795 old.ty, next.ty
6796 )))
6797}
6798
6799fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
6800 matches!(
6801 (old, new),
6802 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
6803 )
6804}
6805
6806fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
6812 let mut prev: Option<i64> = None;
6813 for r in rows {
6814 match r.columns.get(&pk_id) {
6815 Some(Value::Int64(v)) => {
6816 if prev.is_some_and(|p| p >= *v) {
6817 return false;
6818 }
6819 prev = Some(*v);
6820 }
6821 _ => return false,
6822 }
6823 }
6824 true
6825}
6826
6827#[allow(clippy::too_many_arguments)]
6828fn index_into(
6829 schema: &Schema,
6830 row: &Row,
6831 hot: &mut HotIndex,
6832 bitmap: &mut HashMap<u16, BitmapIndex>,
6833 ann: &mut HashMap<u16, AnnIndex>,
6834 fm: &mut HashMap<u16, FmIndex>,
6835 sparse: &mut HashMap<u16, SparseIndex>,
6836 minhash: &mut HashMap<u16, MinHashIndex>,
6837) {
6838 for idef in &schema.indexes {
6839 let Some(val) = row.columns.get(&idef.column_id) else {
6840 continue;
6841 };
6842 match idef.kind {
6843 IndexKind::Bitmap => {
6844 if let Some(b) = bitmap.get_mut(&idef.column_id) {
6845 b.insert(val.encode_key(), row.row_id);
6846 }
6847 }
6848 IndexKind::Ann => {
6849 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
6850 a.insert(v, row.row_id);
6851 }
6852 }
6853 IndexKind::FmIndex => {
6854 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
6855 f.insert(b.clone(), row.row_id);
6856 }
6857 }
6858 IndexKind::Sparse => {
6859 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
6860 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
6863 s.insert(&terms, row.row_id);
6864 }
6865 }
6866 }
6867 IndexKind::MinHash => {
6868 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
6869 let tokens = crate::index::token_hashes_from_bytes(b);
6872 mh.insert(&tokens, row.row_id);
6873 }
6874 }
6875 _ => {}
6876 }
6877 }
6878 if let Some(pk_col) = schema.primary_key() {
6879 if let Some(pk_val) = row.columns.get(&pk_col.id) {
6880 hot.insert(pk_val.encode_key(), row.row_id);
6881 }
6882 }
6883}
6884
6885#[allow(dead_code)]
6891fn bulk_index_key(
6892 column_keys: &HashMap<u16, ([u8; 32], u8)>,
6893 column_id: u16,
6894 ty: TypeId,
6895 col: &columnar::NativeColumn,
6896 i: usize,
6897) -> Option<Vec<u8>> {
6898 let encoded = columnar::encode_key_native(ty, col, i)?;
6899 #[cfg(feature = "encryption")]
6900 {
6901 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6902 if let Some((key, scheme)) = column_keys.get(&column_id) {
6903 return Some(match (*scheme, col) {
6904 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
6905 (_, columnar::NativeColumn::Int64 { data, .. }) => {
6906 ope_token_i64(key, data[i]).to_vec()
6907 }
6908 (_, columnar::NativeColumn::Float64 { data, .. }) => {
6909 ope_token_f64(key, data[i]).to_vec()
6910 }
6911 _ => hmac_token(key, &encoded).to_vec(),
6912 });
6913 }
6914 }
6915 #[cfg(not(feature = "encryption"))]
6916 {
6917 let _ = (column_id, column_keys, col);
6918 }
6919 Some(encoded)
6920}
6921
6922pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
6923 let json = serde_json::to_string_pretty(schema)
6924 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
6925 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
6926 Ok(())
6927}
6928
6929fn read_schema(dir: &Path) -> Result<Schema> {
6930 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
6931 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
6932}
6933
6934fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
6935 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
6936}
6937
6938fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
6939 let n = list_wal_numbers(wal_dir)?;
6940 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
6941}
6942
6943fn next_wal_number(wal_dir: &Path) -> Result<u32> {
6944 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
6945}
6946
6947fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
6948 let _ = std::fs::create_dir_all(wal_dir);
6949 let mut max_n = None;
6950 for entry in std::fs::read_dir(wal_dir)? {
6951 let entry = entry?;
6952 let fname = entry.file_name();
6953 let Some(s) = fname.to_str() else {
6954 continue;
6955 };
6956 let Some(stripped) = s.strip_prefix("seg-") else {
6957 continue;
6958 };
6959 let Some(stripped) = stripped.strip_suffix(".wal") else {
6960 continue;
6961 };
6962 if let Ok(n) = stripped.parse::<u32>() {
6963 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
6964 }
6965 }
6966 Ok(max_n)
6967}