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<crate::cache::Sharded<crate::cache::PageCache>>,
204 snapshots: Arc<crate::retention::SnapshotRegistry>,
207 commit_lock: Arc<parking_lot::Mutex<()>>,
209 decoded_cache: Arc<crate::cache::Sharded<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<crate::cache::Sharded<crate::cache::PageCache>>,
784 pub decoded_cache: Arc<crate::cache::Sharded<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 n_shards = if cache_dir.is_some() {
828 1
829 } else {
830 crate::cache::CACHE_SHARDS
831 };
832 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
833 let page_cache = if let Some(d) = cache_dir {
834 Arc::new(crate::cache::Sharded::new(1, || {
835 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
836 }))
837 } else {
838 Arc::new(crate::cache::Sharded::new(n_shards, || {
839 crate::cache::PageCache::new(per_shard)
840 }))
841 };
842 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
843 let decoded_cache = Arc::new(crate::cache::Sharded::new(
844 crate::cache::CACHE_SHARDS,
845 || crate::cache::DecodedPageCache::new(decoded_per_shard),
846 ));
847 Self {
848 epoch: Arc::new(EpochAuthority::new(0)),
849 page_cache,
850 decoded_cache,
851 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
852 kek,
853 commit_lock: Arc::new(parking_lot::Mutex::new(())),
854 shared: None,
855 }
856 }
857}
858
859fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
863 use crate::query::Condition;
864 match c {
865 Condition::Pk(_)
867 | Condition::BitmapEq { .. }
868 | Condition::BitmapIn { .. }
869 | Condition::BytesPrefix { .. }
870 | Condition::IsNull { .. }
871 | Condition::IsNotNull { .. } => 0,
872 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
874 1
875 }
876 Condition::FmContains { .. }
878 | Condition::FmContainsAll { .. }
879 | Condition::Ann { .. }
880 | Condition::SparseMatch { .. } => 2,
881 }
882}
883
884impl Table {
885 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
886 let dir = dir.as_ref().to_path_buf();
887 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
888 Self::create_in(&dir, schema, table_id, ctx)
889 }
890
891 #[cfg(feature = "encryption")]
902 pub fn create_encrypted(
903 dir: impl AsRef<Path>,
904 schema: Schema,
905 table_id: u64,
906 passphrase: &str,
907 ) -> Result<Self> {
908 let dir = dir.as_ref();
909 std::fs::create_dir_all(dir.join(META_DIR))?;
910 let salt = crate::encryption::random_salt();
911 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
912 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
913 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
914 Self::create_in(dir, schema, table_id, ctx)
915 }
916
917 #[cfg(feature = "encryption")]
922 pub fn create_with_key(
923 dir: impl AsRef<Path>,
924 schema: Schema,
925 table_id: u64,
926 key: &[u8],
927 ) -> Result<Self> {
928 let dir = dir.as_ref();
929 std::fs::create_dir_all(dir.join(META_DIR))?;
930 let salt = crate::encryption::random_salt();
931 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
932 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
933 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
934 Self::create_in(dir, schema, table_id, ctx)
935 }
936
937 #[cfg(feature = "encryption")]
939 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
940 let dir = dir.as_ref();
941 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
942 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
943 MongrelError::NotFound(format!(
944 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
945 salt_path
946 ))
947 })?;
948 if salt_bytes.len() != crate::encryption::SALT_LEN {
949 return Err(MongrelError::InvalidArgument(format!(
950 "salt file is {} bytes, expected {}",
951 salt_bytes.len(),
952 crate::encryption::SALT_LEN
953 )));
954 }
955 let mut salt = [0u8; crate::encryption::SALT_LEN];
956 salt.copy_from_slice(&salt_bytes);
957 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
958 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
959 Self::open_in(dir, ctx)
960 }
961
962 pub(crate) fn create_in(
963 dir: impl AsRef<Path>,
964 schema: Schema,
965 table_id: u64,
966 ctx: SharedCtx,
967 ) -> Result<Self> {
968 schema.validate_auto_increment()?;
969 let dir = dir.as_ref().to_path_buf();
970 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
971 write_schema(&dir, &schema)?;
972 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
973 let (wal, current_txn_id) = match ctx.shared.clone() {
976 Some(s) => (WalSink::Shared(s), 0),
977 None => {
978 std::fs::create_dir_all(dir.join(WAL_DIR))?;
979 let mut w = if let Some(ref dk) = wal_dek {
980 Wal::create_with_cipher(
981 dir.join(WAL_DIR).join("seg-000000.wal"),
982 Epoch(0),
983 Some(make_cipher(dk)),
984 0,
985 )?
986 } else {
987 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
988 };
989 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
990 (WalSink::Private(w), 1)
991 }
992 };
993 let mut manifest = Manifest::new(table_id, schema.schema_id);
994 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
999 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1000 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1001 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1002 let auto_inc = resolve_auto_inc(&schema);
1003 let rcache_dir = dir.join(RCACHE_DIR);
1004 Ok(Self {
1005 dir,
1006 table_id,
1007 wal,
1008 memtable: Memtable::new(),
1009 mutable_run: MutableRun::new(),
1010 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1011 compaction_zstd_level: 3,
1012 allocator: RowIdAllocator::new(0),
1013 epoch: ctx.epoch,
1014 persisted_epoch: 0,
1015 schema,
1016 hot: HotIndex::new(),
1017 kek: ctx.kek,
1018 column_keys,
1019 run_refs: Vec::new(),
1020 retiring: Vec::new(),
1021 next_run_id: 1,
1022 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1023 current_txn_id,
1024 bitmap,
1025 ann,
1026 fm,
1027 sparse,
1028 minhash,
1029 learned_range: HashMap::new(),
1030 pk_by_row: HashMap::new(),
1031 pinned: BTreeMap::new(),
1032 live_count: 0,
1033 reservoir: crate::reservoir::Reservoir::default(),
1034 reservoir_complete: true,
1035 had_deletes: false,
1036 agg_cache: HashMap::new(),
1037 global_idx_epoch: 0,
1038 indexes_complete: true,
1039 index_build_policy: IndexBuildPolicy::default(),
1040 pk_by_row_complete: false,
1041 flushed_epoch: 0,
1042 page_cache: ctx.page_cache,
1043 decoded_cache: ctx.decoded_cache,
1044 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1045 snapshots: ctx.snapshots,
1046 commit_lock: ctx.commit_lock,
1047 result_cache: Arc::new(parking_lot::Mutex::new(
1048 ResultCache::new()
1049 .with_dir(rcache_dir)
1050 .with_cache_dek(cache_dek.clone()),
1051 )),
1052 pending_delete_rids: roaring::RoaringBitmap::new(),
1053 pending_put_cols: std::collections::HashSet::new(),
1054 pending_rows: Vec::new(),
1055 pending_rows_auto_inc: Vec::new(),
1056 pending_dels: Vec::new(),
1057 pending_truncate: None,
1058 wal_dek,
1059 auto_inc,
1060 })
1061 }
1062
1063 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1067 let dir = dir.as_ref();
1068 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1069 Self::open_in(dir, ctx)
1070 }
1071
1072 #[cfg(feature = "encryption")]
1075 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1076 let dir = dir.as_ref();
1077 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1078 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1079 MongrelError::NotFound(format!(
1080 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1081 salt_path
1082 ))
1083 })?;
1084 let salt_len = crate::encryption::SALT_LEN;
1085 if salt_bytes.len() != salt_len {
1086 return Err(MongrelError::InvalidArgument(format!(
1087 "encryption salt is {} bytes, expected {salt_len}",
1088 salt_bytes.len()
1089 )));
1090 }
1091 let mut salt = [0u8; 16];
1092 salt.copy_from_slice(&salt_bytes);
1093 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1094 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1095 let t = Self::open_in(dir, ctx)?;
1096 Ok(t)
1097 }
1098
1099 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1100 let dir = dir.as_ref().to_path_buf();
1101 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1102 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1103 let schema: Schema = read_schema(&dir)?;
1104 let replay_epoch = Epoch(manifest.current_epoch);
1105 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1106 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1110 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1111 None => {
1112 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1113 let replayed = match &active {
1115 Some(path) => {
1116 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1117 crate::wal::replay_with_cipher(path, cipher)?
1118 }
1119 None => Vec::new(),
1120 };
1121 let mut w = match &active {
1122 Some(path) => Wal::create_with_cipher(
1123 path,
1124 replay_epoch,
1125 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1126 0,
1127 )?,
1128 None => Wal::create_with_cipher(
1129 dir.join(WAL_DIR).join("seg-000000.wal"),
1130 replay_epoch,
1131 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1132 0,
1133 )?,
1134 };
1135 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1136 (WalSink::Private(w), replayed, 1)
1137 }
1138 };
1139
1140 let mut memtable = Memtable::new();
1141 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1142 let persisted_epoch = manifest.current_epoch;
1143 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1150 s.next = manifest.auto_inc_next;
1151 s.seeded = manifest.auto_inc_next > 0;
1152 s
1153 });
1154
1155 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1162 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1163 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1164 std::collections::BTreeMap::new();
1165 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1166 let mut saw_delete = false;
1167 for record in replayed {
1168 let txn_id = record.txn_id;
1169 match record.op {
1170 Op::Put { rows, .. } => {
1171 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1172 for row in &rows {
1173 allocator.advance_to(row.row_id);
1174 if let Some(ai) = auto_inc.as_mut() {
1175 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1176 if *n + 1 > ai.next {
1177 ai.next = *n + 1;
1178 }
1179 }
1180 }
1181 }
1182 staged_puts.entry(txn_id).or_default().extend(rows);
1183 }
1184 Op::Delete { row_ids, .. } => {
1185 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1186 }
1187 Op::TxnCommit { epoch, .. } => {
1188 let commit_epoch = Epoch(epoch);
1189 if let Some(puts) = staged_puts.remove(&txn_id) {
1190 for row in &puts {
1191 memtable.upsert(row.clone());
1192 }
1193 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1194 }
1195 if let Some(dels) = staged_deletes.remove(&txn_id) {
1196 saw_delete = true;
1197 for rid in dels {
1198 memtable.tombstone(rid, commit_epoch);
1199 replayed_deletes.push((rid, commit_epoch));
1200 }
1201 }
1202 }
1203 Op::TxnAbort => {
1204 staged_puts.remove(&txn_id);
1205 staged_deletes.remove(&txn_id);
1206 }
1207 Op::TruncateTable { .. } | Op::Flush { .. } | Op::Ddl(_) => {}
1208 }
1209 }
1210
1211 let rcache_dir = dir.join(RCACHE_DIR);
1212 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1213 let mut db = Self {
1214 dir,
1215 table_id: manifest.table_id,
1216 wal,
1217 memtable,
1218 mutable_run: MutableRun::new(),
1219 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1220 compaction_zstd_level: 3,
1221 allocator,
1222 epoch: ctx.epoch,
1223 persisted_epoch,
1224 schema,
1225 hot: HotIndex::new(),
1226 kek: ctx.kek,
1227 column_keys,
1228 run_refs: manifest.runs.clone(),
1229 retiring: manifest.retiring.clone(),
1230 next_run_id: manifest
1231 .runs
1232 .iter()
1233 .map(|r| r.run_id as u64 + 1)
1234 .max()
1235 .unwrap_or(1),
1236 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1237 current_txn_id,
1238 bitmap: HashMap::new(),
1239 ann: HashMap::new(),
1240 fm: HashMap::new(),
1241 sparse: HashMap::new(),
1242 minhash: HashMap::new(),
1243 learned_range: HashMap::new(),
1244 pk_by_row: HashMap::new(),
1245 pinned: BTreeMap::new(),
1246 live_count: manifest.live_count,
1247 reservoir: crate::reservoir::Reservoir::default(),
1248 reservoir_complete: false,
1249 had_deletes: saw_delete,
1250 agg_cache: HashMap::new(),
1251 global_idx_epoch: manifest.global_idx_epoch,
1252 indexes_complete: true,
1253 index_build_policy: IndexBuildPolicy::default(),
1254 pk_by_row_complete: false,
1255 flushed_epoch: manifest.flushed_epoch,
1256 page_cache: ctx.page_cache,
1257 decoded_cache: ctx.decoded_cache,
1258 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1259 snapshots: ctx.snapshots,
1260 commit_lock: ctx.commit_lock,
1261 result_cache: Arc::new(parking_lot::Mutex::new(
1262 ResultCache::new()
1263 .with_dir(rcache_dir)
1264 .with_cache_dek(cache_dek.clone()),
1265 )),
1266 pending_delete_rids: roaring::RoaringBitmap::new(),
1267 pending_put_cols: std::collections::HashSet::new(),
1268 pending_rows: Vec::new(),
1269 pending_rows_auto_inc: Vec::new(),
1270 pending_dels: Vec::new(),
1271 pending_truncate: None,
1272 wal_dek,
1273 auto_inc,
1274 };
1275
1276 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1279
1280 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1285 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1286 c.epoch_built == manifest.global_idx_epoch
1287 && manifest.global_idx_epoch > 0
1288 && manifest
1289 .runs
1290 .iter()
1291 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1292 });
1293 if let Some(loaded) = checkpoint {
1294 if checkpoint_valid {
1295 db.hot = loaded.hot;
1296 db.bitmap = loaded.bitmap;
1297 db.ann = loaded.ann;
1298 db.fm = loaded.fm;
1299 db.sparse = loaded.sparse;
1300 db.minhash = loaded.minhash;
1301 db.learned_range = loaded.learned_range;
1302 }
1305 }
1306 if !checkpoint_valid {
1307 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1308 db.bitmap = bitmap;
1309 db.ann = ann;
1310 db.fm = fm;
1311 db.sparse = sparse;
1312 db.minhash = minhash;
1313 db.rebuild_indexes_from_runs()?;
1314 db.build_learned_ranges()?;
1315 }
1316
1317 for (epoch, group) in replayed_puts {
1322 let (losers, winner_pks) = db.partition_pk_winners(&group);
1323 for (key, &row_id) in &winner_pks {
1324 if let Some(old_rid) = db.hot.get(key) {
1325 if old_rid != row_id {
1326 db.tombstone_row(old_rid, epoch, false);
1327 }
1328 }
1329 }
1330 for &loser_rid in &losers {
1331 db.tombstone_row(loser_rid, epoch, false);
1332 }
1333 for (key, row_id) in winner_pks {
1334 db.insert_hot_pk(key, row_id);
1335 }
1336 if db.schema.primary_key().is_none() {
1337 for r in &group {
1338 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1339 }
1340 }
1341 for r in &group {
1342 if !losers.contains(&r.row_id) {
1343 db.index_row(r);
1344 }
1345 }
1346 }
1347 for (rid, epoch) in &replayed_deletes {
1351 db.remove_hot_for_row(*rid, *epoch);
1352 }
1353
1354 db.result_cache.lock().load_persistent();
1361 Ok(db)
1362 }
1363
1364 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1370 if self.reservoir_complete {
1371 return Ok(());
1372 }
1373 self.rebuild_reservoir()?;
1374 self.reservoir_complete = true;
1375 Ok(())
1376 }
1377
1378 fn rebuild_reservoir(&mut self) -> Result<()> {
1381 let snap = self.snapshot();
1382 let rows = self.visible_rows(snap)?;
1383 self.reservoir.reset();
1384 for r in rows {
1385 self.reservoir.offer(r.row_id.0);
1386 }
1387 Ok(())
1388 }
1389
1390 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1391 self.hot = HotIndex::new();
1392 self.pk_by_row.clear();
1393 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1394 self.bitmap = bitmap;
1395 self.ann = ann;
1396 self.fm = fm;
1397 self.sparse = sparse;
1398 self.minhash = minhash;
1399 let snapshot = Epoch(u64::MAX);
1400 for rr in self.run_refs.clone() {
1401 let mut reader = self.open_reader(rr.run_id)?;
1402 for row in reader.visible_rows(snapshot)? {
1403 let tok_row = self.tokenized_for_indexes(&row);
1404 index_into(
1405 &self.schema,
1406 &tok_row,
1407 &mut self.hot,
1408 &mut self.bitmap,
1409 &mut self.ann,
1410 &mut self.fm,
1411 &mut self.sparse,
1412 &mut self.minhash,
1413 );
1414 }
1415 }
1416 for row in self.mutable_run.visible_versions(snapshot) {
1417 if row.deleted {
1418 self.remove_hot_for_row(row.row_id, snapshot);
1419 } else {
1420 self.index_row(&row);
1421 }
1422 }
1423 for row in self.memtable.visible_versions(snapshot) {
1424 if row.deleted {
1425 self.remove_hot_for_row(row.row_id, snapshot);
1426 } else {
1427 self.index_row(&row);
1428 }
1429 }
1430 self.refresh_pk_by_row_from_hot();
1431 Ok(())
1432 }
1433
1434 fn refresh_pk_by_row_from_hot(&mut self) {
1435 self.pk_by_row_complete = true;
1436 if self.schema.primary_key().is_none() {
1437 self.pk_by_row.clear();
1438 return;
1439 }
1440 self.pk_by_row = self
1446 .hot
1447 .entries()
1448 .into_iter()
1449 .map(|(key, row_id)| (row_id, key))
1450 .collect();
1451 }
1452
1453 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1454 if self.schema.primary_key().is_some() {
1455 self.pk_by_row.insert(row_id, key.clone());
1456 }
1457 self.hot.insert(key, row_id);
1458 }
1459
1460 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1464 self.learned_range.clear();
1465 if self.run_refs.len() != 1 {
1466 return Ok(());
1467 }
1468 let cols: Vec<u16> = self
1469 .schema
1470 .indexes
1471 .iter()
1472 .filter(|i| i.kind == IndexKind::LearnedRange)
1473 .map(|i| i.column_id)
1474 .collect();
1475 if cols.is_empty() {
1476 return Ok(());
1477 }
1478 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1479 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1480 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1481 _ => return Ok(()),
1482 };
1483 for cid in cols {
1484 let ty = self
1485 .schema
1486 .columns
1487 .iter()
1488 .find(|c| c.id == cid)
1489 .map(|c| c.ty)
1490 .unwrap_or(TypeId::Int64);
1491 match ty {
1492 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1493 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1494 let pairs: Vec<(i64, u64)> = data
1495 .iter()
1496 .zip(row_ids.iter())
1497 .map(|(v, r)| (*v, *r))
1498 .collect();
1499 self.learned_range
1500 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1501 }
1502 }
1503 TypeId::Float64 => {
1504 if let columnar::NativeColumn::Float64 { data, .. } =
1505 reader.column_native(cid)?
1506 {
1507 let pairs: Vec<(f64, u64)> = data
1508 .iter()
1509 .zip(row_ids.iter())
1510 .map(|(v, r)| (*v, *r))
1511 .collect();
1512 self.learned_range
1513 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1514 }
1515 }
1516 _ => {}
1517 }
1518 }
1519 Ok(())
1520 }
1521
1522 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1529 if self.indexes_complete {
1530 crate::trace::QueryTrace::record(|t| {
1531 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1532 });
1533 return Ok(());
1534 }
1535 crate::trace::QueryTrace::record(|t| {
1536 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1537 });
1538 self.rebuild_indexes_from_runs()?;
1539 self.build_learned_ranges()?;
1540 self.indexes_complete = true;
1541 let epoch = self.current_epoch();
1542 self.checkpoint_indexes(epoch);
1543 Ok(())
1544 }
1545
1546 fn pending_epoch(&self) -> Epoch {
1547 Epoch(self.epoch.visible().0 + 1)
1548 }
1549
1550 fn is_shared(&self) -> bool {
1553 matches!(self.wal, WalSink::Shared(_))
1554 }
1555
1556 fn ensure_txn_id(&mut self) -> u64 {
1560 if self.current_txn_id == 0 {
1561 let id = match &self.wal {
1562 WalSink::Shared(s) => {
1563 let mut g = s.txn_ids.lock();
1564 let v = *g;
1565 *g = g.wrapping_add(1);
1566 v
1567 }
1568 WalSink::Private(_) => 1,
1569 };
1570 self.current_txn_id = id;
1571 }
1572 self.current_txn_id
1573 }
1574
1575 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1578 let txn_id = self.ensure_txn_id();
1579 let table_id = self.table_id;
1580 match &mut self.wal {
1581 WalSink::Private(w) => {
1582 w.append_txn(txn_id, op)?;
1583 }
1584 WalSink::Shared(s) => {
1585 s.wal.lock().append(txn_id, table_id, op)?;
1586 }
1587 }
1588 Ok(())
1589 }
1590
1591 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1599 Ok(self.put_returning(columns)?.0)
1600 }
1601
1602 pub fn put_returning(
1607 &mut self,
1608 mut columns: Vec<(u16, Value)>,
1609 ) -> Result<(RowId, Option<i64>)> {
1610 let assigned = self.fill_auto_inc(&mut columns)?;
1611 self.schema.validate_not_null(&columns)?;
1612 let row_id = self.allocator.alloc();
1613 let epoch = self.pending_epoch();
1614 let mut row = Row::new(row_id, epoch);
1615 for (col_id, val) in columns {
1616 row.columns.insert(col_id, val);
1617 }
1618 self.commit_rows(vec![row], assigned.is_some())?;
1619 Ok((row_id, assigned))
1620 }
1621
1622 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1625 Ok(self
1626 .put_batch_returning(batch)?
1627 .into_iter()
1628 .map(|(r, _)| r)
1629 .collect())
1630 }
1631
1632 pub fn put_batch_returning(
1635 &mut self,
1636 batch: Vec<Vec<(u16, Value)>>,
1637 ) -> Result<Vec<(RowId, Option<i64>)>> {
1638 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1639 for mut cols in batch {
1640 let assigned = self.fill_auto_inc(&mut cols)?;
1641 filled.push((cols, assigned));
1642 }
1643 for (cols, _) in &filled {
1644 self.schema.validate_not_null(cols)?;
1645 }
1646 let epoch = self.pending_epoch();
1647 let mut rows = Vec::with_capacity(filled.len());
1648 let mut ids = Vec::with_capacity(filled.len());
1649 for (cols, assigned) in filled {
1650 let row_id = self.allocator.alloc();
1651 let mut row = Row::new(row_id, epoch);
1652 for (c, v) in cols {
1653 row.columns.insert(c, v);
1654 }
1655 ids.push((row_id, assigned));
1656 rows.push(row);
1657 }
1658 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1659 self.commit_rows(rows, all_auto_generated)?;
1660 Ok(ids)
1661 }
1662
1663 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1669 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1670 return Ok(None);
1671 };
1672 let pos = columns.iter().position(|(c, _)| *c == cid);
1673 let assigned = match pos {
1674 Some(i) => match &columns[i].1 {
1675 Value::Null => {
1676 let next = self.alloc_auto_inc_value()?;
1677 columns[i].1 = Value::Int64(next);
1678 Some(next)
1679 }
1680 Value::Int64(n) => {
1681 self.advance_auto_inc_past(*n)?;
1682 None
1683 }
1684 other => {
1685 return Err(MongrelError::InvalidArgument(format!(
1686 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1687 other
1688 )))
1689 }
1690 },
1691 None => {
1692 let next = self.alloc_auto_inc_value()?;
1693 columns.push((cid, Value::Int64(next)));
1694 Some(next)
1695 }
1696 };
1697 Ok(assigned)
1698 }
1699
1700 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1702 self.ensure_auto_inc_seeded()?;
1703 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1705 let v = ai.next;
1706 ai.next = ai.next.saturating_add(1);
1707 Ok(v)
1708 }
1709
1710 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1713 self.ensure_auto_inc_seeded()?;
1714 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1715 let floor = used.saturating_add(1).max(1);
1716 if ai.next < floor {
1717 ai.next = floor;
1718 }
1719 Ok(())
1720 }
1721
1722 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1727 let needs_seed = match self.auto_inc {
1728 Some(ai) => !ai.seeded,
1729 None => return Ok(()),
1730 };
1731 if !needs_seed {
1732 return Ok(());
1733 }
1734 if self.seed_empty_auto_inc() {
1735 return Ok(());
1736 }
1737 let cid = self
1738 .auto_inc
1739 .as_ref()
1740 .expect("auto-inc column present")
1741 .column_id;
1742 let max = self.scan_max_int64(cid)?;
1743 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1744 let floor = max.saturating_add(1).max(1);
1745 if ai.next < floor {
1746 ai.next = floor;
1747 }
1748 ai.seeded = true;
1749 Ok(())
1750 }
1751
1752 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1753 if n == 0 || self.auto_inc.is_none() {
1754 return Ok(None);
1755 }
1756 self.ensure_auto_inc_seeded()?;
1757 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1758 let start = ai.next;
1759 ai.next = ai.next.saturating_add(n as i64);
1760 Ok(Some(start))
1761 }
1762
1763 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1767 let mut max: i64 = 0;
1768 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1769 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1770 if *n > max {
1771 max = *n;
1772 }
1773 }
1774 }
1775 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1776 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1777 if *n > max {
1778 max = *n;
1779 }
1780 }
1781 }
1782 for rr in self.run_refs.clone() {
1783 let reader = self.open_reader(rr.run_id)?;
1784 if let Some(stats) = reader.column_page_stats(column_id) {
1785 for s in stats {
1786 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1787 if n > max {
1788 max = n;
1789 }
1790 }
1791 }
1792 } else if reader.has_column(column_id) {
1793 if let columnar::NativeColumn::Int64 { data, validity } =
1794 reader.column_native_shared(column_id)?
1795 {
1796 for (i, n) in data.iter().enumerate() {
1797 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1798 {
1799 max = *n;
1800 }
1801 }
1802 }
1803 }
1804 }
1805 Ok(max)
1806 }
1807
1808 fn seed_empty_auto_inc(&mut self) -> bool {
1809 let Some(ai) = self.auto_inc.as_mut() else {
1810 return false;
1811 };
1812 if ai.seeded || self.live_count != 0 {
1813 return false;
1814 }
1815 if ai.next < 1 {
1816 ai.next = 1;
1817 }
1818 ai.seeded = true;
1819 true
1820 }
1821
1822 fn advance_auto_inc_from_native_columns(
1823 &mut self,
1824 columns: &[(u16, columnar::NativeColumn)],
1825 n: usize,
1826 live_before: u64,
1827 ) -> Result<()> {
1828 let Some(ai) = self.auto_inc.as_mut() else {
1829 return Ok(());
1830 };
1831 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1832 return Ok(());
1833 };
1834 let columnar::NativeColumn::Int64 { data, validity } = col else {
1835 return Err(MongrelError::InvalidArgument(format!(
1836 "AUTO_INCREMENT column {} must be Int64",
1837 ai.column_id
1838 )));
1839 };
1840 let max = if native_int64_strictly_increasing(col, n) {
1841 data.get(n.saturating_sub(1)).copied()
1842 } else {
1843 data.iter()
1844 .take(n)
1845 .enumerate()
1846 .filter_map(|(i, v)| {
1847 if validity.is_empty() || columnar::validity_bit(validity, i) {
1848 Some(*v)
1849 } else {
1850 None
1851 }
1852 })
1853 .max()
1854 };
1855 if let Some(max) = max {
1856 let floor = max.saturating_add(1).max(1);
1857 if ai.next < floor {
1858 ai.next = floor;
1859 }
1860 if ai.seeded || live_before == 0 {
1861 ai.seeded = true;
1862 }
1863 }
1864 Ok(())
1865 }
1866
1867 fn fill_auto_inc_native_columns(
1868 &mut self,
1869 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1870 n: usize,
1871 ) -> Result<()> {
1872 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1873 return Ok(());
1874 };
1875 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1876 if let Some(start) = self.alloc_auto_inc_range(n)? {
1877 columns.push((
1878 cid,
1879 columnar::NativeColumn::Int64 {
1880 data: (start..start.saturating_add(n as i64)).collect(),
1881 validity: vec![0xFF; n.div_ceil(8)],
1882 },
1883 ));
1884 }
1885 return Ok(());
1886 };
1887
1888 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1889 return Err(MongrelError::InvalidArgument(format!(
1890 "AUTO_INCREMENT column {cid} must be Int64"
1891 )));
1892 };
1893 if data.len() < n {
1894 return Err(MongrelError::InvalidArgument(format!(
1895 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1896 data.len()
1897 )));
1898 }
1899 if columnar::all_non_null(validity, n) {
1900 return Ok(());
1901 }
1902 if validity.iter().all(|b| *b == 0) {
1903 if let Some(start) = self.alloc_auto_inc_range(n)? {
1904 for (i, slot) in data.iter_mut().take(n).enumerate() {
1905 *slot = start.saturating_add(i as i64);
1906 }
1907 *validity = vec![0xFF; n.div_ceil(8)];
1908 }
1909 return Ok(());
1910 }
1911
1912 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1913 for (i, slot) in data.iter_mut().enumerate().take(n) {
1914 if columnar::validity_bit(validity, i) {
1915 self.advance_auto_inc_past(*slot)?;
1916 } else {
1917 *slot = self.alloc_auto_inc_value()?;
1918 }
1919 }
1920 *validity = new_validity;
1921 Ok(())
1922 }
1923
1924 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1938 if self.auto_inc.is_none() {
1939 return Ok(None);
1940 }
1941 Ok(Some(self.alloc_auto_inc_value()?))
1942 }
1943
1944 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1950 let payload = bincode::serialize(&rows)?;
1951 self.wal_append_data(Op::Put {
1952 table_id: self.table_id,
1953 rows: payload,
1954 })?;
1955 if self.is_shared() {
1956 self.pending_rows_auto_inc
1957 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1958 self.pending_rows.extend(rows);
1959 } else {
1960 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1961 }
1962 Ok(())
1963 }
1964
1965 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1970 self.apply_put_rows_inner(rows, true)
1971 }
1972
1973 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1974 if check_existing_pk {
1975 self.ensure_indexes_complete()?;
1976 }
1977 if rows.len() == 1 {
1981 let row = rows.into_iter().next().expect("len checked");
1982 return self.apply_put_row_single(row, check_existing_pk);
1983 }
1984 let pk_id = self.schema.primary_key().map(|c| c.id);
2001 let probe = match pk_id {
2002 Some(pid) => {
2003 check_existing_pk
2004 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2005 }
2006 None => false,
2007 };
2008 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2011 for r in rows {
2012 for &cid in r.columns.keys() {
2013 self.pending_put_cols.insert(cid);
2014 }
2015 match pk_id {
2016 Some(pid) if probe || maintain_pk_by_row => {
2017 if let Some(pk_val) = r.columns.get(&pid) {
2018 let key = self.index_lookup_key(pid, pk_val);
2019 if probe {
2020 if let Some(old_rid) = self.hot.get(&key) {
2021 if old_rid != r.row_id {
2022 self.tombstone_row(old_rid, r.committed_epoch, true);
2023 }
2024 }
2025 }
2026 if maintain_pk_by_row {
2027 self.pk_by_row.insert(r.row_id, key);
2028 }
2029 }
2030 }
2031 Some(_) => {}
2032 None => {
2033 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2034 }
2035 }
2036 self.index_row(&r);
2037 self.reservoir.offer(r.row_id.0);
2038 self.memtable.upsert(r);
2039 self.live_count = self.live_count.saturating_add(1);
2042 }
2043 Ok(())
2044 }
2045
2046 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2050 for &cid in row.columns.keys() {
2051 self.pending_put_cols.insert(cid);
2052 }
2053 let epoch = row.committed_epoch;
2054 if let Some(pk_col) = self.schema.primary_key() {
2055 let pk_id = pk_col.id;
2056 if let Some(pk_val) = row.columns.get(&pk_id) {
2057 let maintain_pk_by_row = self.pk_by_row_complete;
2061 if check_existing_pk || maintain_pk_by_row {
2062 let key = self.index_lookup_key(pk_id, pk_val);
2063 if check_existing_pk {
2064 if let Some(old_rid) = self.hot.get(&key) {
2065 if old_rid != row.row_id {
2066 self.tombstone_row(old_rid, epoch, true);
2067 }
2068 }
2069 }
2070 if maintain_pk_by_row {
2071 self.pk_by_row.insert(row.row_id, key);
2072 }
2073 }
2074 }
2075 } else {
2076 self.hot
2077 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2078 }
2079 self.index_row(&row);
2080 self.reservoir.offer(row.row_id.0);
2081 self.memtable.upsert(row);
2082 self.live_count = self.live_count.saturating_add(1);
2083 Ok(())
2084 }
2085
2086 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2089 self.allocator.alloc()
2090 }
2091
2092 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2100 self.ensure_indexes_complete()?;
2101 let n = rows.len();
2102 for r in rows {
2103 for &cid in r.columns.keys() {
2104 self.pending_put_cols.insert(cid);
2105 }
2106 }
2107 let (losers, winner_pks) = self.partition_pk_winners(rows);
2108 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2109 for (key, &row_id) in &winner_pks {
2111 if let Some(old_rid) = self.hot.get(key) {
2112 if old_rid != row_id {
2113 self.tombstone_row(old_rid, epoch, true);
2114 }
2115 }
2116 }
2117 for &loser_rid in &losers {
2120 self.tombstone_row(loser_rid, epoch, false);
2121 }
2122 for (key, row_id) in winner_pks {
2124 self.insert_hot_pk(key, row_id);
2125 }
2126 if self.schema.primary_key().is_none() {
2127 for r in rows {
2128 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2129 }
2130 }
2131 for r in rows {
2132 self.allocator.advance_to(r.row_id);
2133 if !losers.contains(&r.row_id) {
2134 self.index_row(r);
2135 }
2136 }
2137 for r in rows {
2138 if !losers.contains(&r.row_id) {
2139 self.reservoir.offer(r.row_id.0);
2140 }
2141 }
2142 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2143 Ok(())
2144 }
2145
2146 pub(crate) fn recover_apply(
2151 &mut self,
2152 rows: Vec<Row>,
2153 deletes: Vec<(RowId, Epoch)>,
2154 ) -> Result<()> {
2155 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2159 std::collections::BTreeMap::new();
2160 for row in rows {
2161 self.allocator.advance_to(row.row_id);
2162 if let Some(ai) = self.auto_inc.as_mut() {
2167 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2168 if *n + 1 > ai.next {
2169 ai.next = *n + 1;
2170 }
2171 }
2172 }
2173 by_epoch.entry(row.committed_epoch).or_default().push(row);
2174 }
2175 for (epoch, group) in by_epoch {
2176 let (losers, winner_pks) = self.partition_pk_winners(&group);
2177 for (key, &row_id) in &winner_pks {
2179 if let Some(old_rid) = self.hot.get(key) {
2180 if old_rid != row_id {
2181 self.tombstone_row(old_rid, epoch, false);
2182 }
2183 }
2184 }
2185 for (key, row_id) in winner_pks {
2186 self.insert_hot_pk(key, row_id);
2187 }
2188 if self.schema.primary_key().is_none() {
2189 for r in &group {
2190 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2191 }
2192 }
2193 for r in &group {
2194 if !losers.contains(&r.row_id) {
2195 self.memtable.upsert(r.clone());
2196 self.index_row(r);
2197 }
2198 }
2199 }
2200 for (rid, epoch) in deletes {
2201 self.memtable.tombstone(rid, epoch);
2202 self.remove_hot_for_row(rid, epoch);
2203 }
2204 self.reservoir_complete = false;
2207 Ok(())
2208 }
2209
2210 pub(crate) fn flushed_epoch(&self) -> u64 {
2212 self.flushed_epoch
2213 }
2214
2215 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2217 self.schema.validate_not_null(cells)
2218 }
2219
2220 fn validate_columns_not_null(
2224 &self,
2225 columns: &[(u16, columnar::NativeColumn)],
2226 n: usize,
2227 ) -> Result<()> {
2228 let by_id: HashMap<u16, &columnar::NativeColumn> =
2229 columns.iter().map(|(id, c)| (*id, c)).collect();
2230 for col in &self.schema.columns {
2231 if col.flags.contains(ColumnFlags::NULLABLE) {
2232 continue;
2233 }
2234 match by_id.get(&col.id) {
2235 None => {
2236 return Err(MongrelError::InvalidArgument(format!(
2237 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2238 col.name, col.id
2239 )));
2240 }
2241 Some(c) => {
2242 if c.null_count(n) != 0 {
2243 return Err(MongrelError::InvalidArgument(format!(
2244 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2245 col.name, col.id
2246 )));
2247 }
2248 }
2249 }
2250 }
2251 Ok(())
2252 }
2253
2254 fn bulk_pk_winner_indices(
2259 &self,
2260 columns: &[(u16, columnar::NativeColumn)],
2261 n: usize,
2262 ) -> Option<Vec<usize>> {
2263 let pk_col = self.schema.primary_key()?;
2264 let pk_id = pk_col.id;
2265 let pk_ty = pk_col.ty;
2266 let by_id: HashMap<u16, &columnar::NativeColumn> =
2267 columns.iter().map(|(id, c)| (*id, c)).collect();
2268 let pk_native = by_id.get(&pk_id)?;
2269 if native_int64_strictly_increasing(pk_native, n) {
2270 return None;
2271 }
2272 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2274 let mut null_pk_rows: Vec<usize> = Vec::new();
2275 for i in 0..n {
2276 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2277 Some(key) => {
2278 last.insert(key, i);
2279 }
2280 None => null_pk_rows.push(i),
2281 }
2282 }
2283 let mut winners: HashSet<usize> = last.values().copied().collect();
2284 for i in null_pk_rows {
2285 winners.insert(i);
2286 }
2287 Some((0..n).filter(|i| winners.contains(i)).collect())
2288 }
2289
2290 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2292 let epoch = self.pending_epoch();
2293 self.wal_append_data(Op::Delete {
2294 table_id: self.table_id,
2295 row_ids: vec![row_id],
2296 })?;
2297 if self.is_shared() {
2298 self.pending_dels.push(row_id);
2299 } else {
2300 self.apply_delete(row_id, epoch);
2301 }
2302 Ok(())
2303 }
2304
2305 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2306 let pre = self.get(row_id, self.snapshot());
2307 self.delete(row_id)?;
2308 Ok(pre.map(|row| {
2309 let mut columns: Vec<_> = row.columns.into_iter().collect();
2310 columns.sort_by_key(|(id, _)| *id);
2311 OwnedRow { columns }
2312 }))
2313 }
2314
2315 pub fn truncate(&mut self) -> Result<()> {
2317 let epoch = self.pending_epoch();
2318 self.wal_append_data(Op::TruncateTable {
2319 table_id: self.table_id,
2320 })?;
2321 self.pending_rows.clear();
2322 self.pending_rows_auto_inc.clear();
2323 self.pending_dels.clear();
2324 self.pending_truncate = Some(epoch);
2325 Ok(())
2326 }
2327
2328 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2330 for rr in std::mem::take(&mut self.run_refs) {
2331 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2332 }
2333 for r in std::mem::take(&mut self.retiring) {
2334 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2335 }
2336 self.memtable = Memtable::new();
2337 self.mutable_run = MutableRun::new();
2338 self.hot = HotIndex::new();
2339 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2340 self.bitmap = bitmap;
2341 self.ann = ann;
2342 self.fm = fm;
2343 self.sparse = sparse;
2344 self.minhash = minhash;
2345 self.learned_range.clear();
2346 self.pk_by_row.clear();
2347 self.pk_by_row_complete = false;
2348 self.live_count = 0;
2349 self.reservoir = crate::reservoir::Reservoir::default();
2350 self.reservoir_complete = true;
2351 self.had_deletes = true;
2352 self.agg_cache.clear();
2353 self.global_idx_epoch = 0;
2354 self.indexes_complete = true;
2355 self.pending_delete_rids.clear();
2356 self.pending_put_cols.clear();
2357 self.pending_rows.clear();
2358 self.pending_rows_auto_inc.clear();
2359 self.pending_dels.clear();
2360 self.clear_result_cache();
2361 self.invalidate_index_checkpoint();
2362 Ok(())
2363 }
2364
2365 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2368 self.remove_hot_for_row(row_id, epoch);
2369 self.tombstone_row(row_id, epoch, true);
2370 }
2371
2372 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2376 let tombstone = Row {
2377 row_id,
2378 committed_epoch: epoch,
2379 columns: std::collections::HashMap::new(),
2380 deleted: true,
2381 };
2382 self.memtable.upsert(tombstone);
2383 self.pk_by_row.remove(&row_id);
2384 if adjust_live_count {
2385 self.live_count = self.live_count.saturating_sub(1);
2386 }
2387 self.pending_delete_rids.insert(row_id.0 as u32);
2389 self.had_deletes = true;
2392 self.agg_cache.clear();
2393 }
2394
2395 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2399 let Some(pk_col) = self.schema.primary_key() else {
2400 return;
2401 };
2402 if self.pk_by_row_complete {
2405 if let Some(key) = self.pk_by_row.remove(&row_id) {
2406 if self.hot.get(&key) == Some(row_id) {
2407 self.hot.remove(&key);
2408 }
2409 }
2410 return;
2411 }
2412 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2431 if self.indexes_complete {
2432 let pk_val = self
2433 .memtable
2434 .get_version(row_id, lookup_epoch)
2435 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2436 .or_else(|| {
2437 self.mutable_run
2438 .get_version(row_id, lookup_epoch)
2439 .filter(|(_, r)| !r.deleted)
2440 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2441 })
2442 .or_else(|| {
2443 self.run_refs.iter().find_map(|rr| {
2444 let mut reader = self.open_reader(rr.run_id).ok()?;
2445 let (_, deleted, val) = reader
2446 .get_version_column(row_id, lookup_epoch, pk_col.id)
2447 .ok()??;
2448 if deleted {
2449 return None;
2450 }
2451 val
2452 })
2453 });
2454 if let Some(pk_val) = pk_val {
2455 let key = self.index_lookup_key(pk_col.id, &pk_val);
2456 if self.hot.get(&key) == Some(row_id) {
2457 self.hot.remove(&key);
2458 }
2459 return;
2460 }
2461 }
2462 self.refresh_pk_by_row_from_hot();
2467 if let Some(key) = self.pk_by_row.remove(&row_id) {
2468 if self.hot.get(&key) == Some(row_id) {
2469 self.hot.remove(&key);
2470 }
2471 }
2472 }
2473
2474 fn partition_pk_winners(
2479 &self,
2480 rows: &[Row],
2481 ) -> (
2482 std::collections::HashSet<RowId>,
2483 std::collections::HashMap<Vec<u8>, RowId>,
2484 ) {
2485 let mut losers = std::collections::HashSet::new();
2486 let Some(pk_col) = self.schema.primary_key() else {
2487 return (losers, std::collections::HashMap::new());
2488 };
2489 let pk_id = pk_col.id;
2490 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2491 std::collections::HashMap::new();
2492 for r in rows {
2493 let Some(pk_val) = r.columns.get(&pk_id) else {
2494 continue;
2495 };
2496 let key = self.index_lookup_key(pk_id, pk_val);
2497 if let Some(&old_rid) = winners.get(&key) {
2498 losers.insert(old_rid);
2499 }
2500 winners.insert(key, r.row_id);
2501 }
2502 (losers, winners)
2503 }
2504
2505 fn index_row(&mut self, row: &Row) {
2506 if row.deleted {
2507 return;
2508 }
2509 if self.column_keys.is_empty() {
2513 index_into(
2514 &self.schema,
2515 row,
2516 &mut self.hot,
2517 &mut self.bitmap,
2518 &mut self.ann,
2519 &mut self.fm,
2520 &mut self.sparse,
2521 &mut self.minhash,
2522 );
2523 return;
2524 }
2525 let effective_row = self.tokenized_for_indexes(row);
2526 index_into(
2527 &self.schema,
2528 &effective_row,
2529 &mut self.hot,
2530 &mut self.bitmap,
2531 &mut self.ann,
2532 &mut self.fm,
2533 &mut self.sparse,
2534 &mut self.minhash,
2535 );
2536 }
2537
2538 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2544 if self.column_keys.is_empty() {
2545 return row.clone();
2546 }
2547 #[cfg(feature = "encryption")]
2548 {
2549 use crate::encryption::SCHEME_HMAC_EQ;
2550 let mut tok = row.clone();
2551 for (&cid, &(_, scheme)) in &self.column_keys {
2552 if scheme != SCHEME_HMAC_EQ {
2553 continue;
2554 }
2555 if let Some(v) = tok.columns.get(&cid).cloned() {
2556 if let Some(t) = self.tokenize_value(cid, &v) {
2557 tok.columns.insert(cid, t);
2558 }
2559 }
2560 }
2561 tok
2562 }
2563 #[cfg(not(feature = "encryption"))]
2564 {
2565 row.clone()
2566 }
2567 }
2568
2569 pub fn commit(&mut self) -> Result<Epoch> {
2574 if self.is_shared() {
2575 self.commit_shared()
2576 } else {
2577 self.commit_private()
2578 }
2579 }
2580
2581 fn commit_private(&mut self) -> Result<Epoch> {
2583 let commit_lock = Arc::clone(&self.commit_lock);
2587 let _g = commit_lock.lock();
2588 let new_epoch = self.epoch.bump_assigned();
2589 let txn_id = self.current_txn_id;
2590 match &mut self.wal {
2594 WalSink::Private(w) => {
2595 w.append_txn(
2596 txn_id,
2597 Op::TxnCommit {
2598 epoch: new_epoch.0,
2599 added_runs: Vec::new(),
2600 },
2601 )?;
2602 w.sync()?;
2603 }
2604 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2605 }
2606 if let Some(epoch) = self.pending_truncate.take() {
2608 self.apply_truncate(epoch)?;
2609 }
2610 self.invalidate_pending_cache();
2611 self.persist_manifest(new_epoch)?;
2612 self.epoch.publish_in_order(new_epoch);
2616 self.current_txn_id += 1;
2617 Ok(new_epoch)
2618 }
2619
2620 fn commit_shared(&mut self) -> Result<Epoch> {
2626 use std::sync::atomic::Ordering;
2627 let s = match &self.wal {
2628 WalSink::Shared(s) => s.clone(),
2629 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2630 };
2631 if s.poisoned.load(Ordering::Relaxed) {
2632 return Err(MongrelError::Other(
2633 "database poisoned by fsync error".into(),
2634 ));
2635 }
2636 let commit_lock = Arc::clone(&self.commit_lock);
2643 let _g = commit_lock.lock();
2644 let txn_id = self.ensure_txn_id();
2647 let (new_epoch, commit_seq) = {
2648 let mut wal = s.wal.lock();
2649 let new_epoch = self.epoch.bump_assigned();
2650 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2651 (new_epoch, seq)
2652 };
2653 s.group
2654 .await_durable(&s.wal, commit_seq)
2655 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2656
2657 if self.pending_truncate.take().is_some() {
2660 self.apply_truncate(new_epoch)?;
2661 }
2662 let mut rows = std::mem::take(&mut self.pending_rows);
2663 if !rows.is_empty() {
2664 for r in &mut rows {
2665 r.committed_epoch = new_epoch;
2666 }
2667 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2668 let all_auto_generated =
2669 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2670 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2671 } else {
2672 self.pending_rows_auto_inc.clear();
2673 }
2674 let dels = std::mem::take(&mut self.pending_dels);
2675 for rid in dels {
2676 self.apply_delete(rid, new_epoch);
2677 }
2678
2679 self.invalidate_pending_cache();
2680 self.persist_manifest(new_epoch)?;
2681 self.epoch.publish_in_order(new_epoch);
2682 self.current_txn_id = 0;
2684 Ok(new_epoch)
2685 }
2686
2687 pub fn flush(&mut self) -> Result<Epoch> {
2695 self.ensure_indexes_complete()?;
2696 let epoch = self.commit()?;
2697 let rows = self.memtable.drain_sorted();
2698 if !rows.is_empty() {
2699 self.mutable_run.insert_many(rows);
2700 }
2701 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2702 self.spill_mutable_run(epoch)?;
2703 self.mark_flushed(epoch)?;
2707 self.persist_manifest(epoch)?;
2708 self.build_learned_ranges()?;
2709 self.checkpoint_indexes(epoch);
2712 }
2713 Ok(epoch)
2716 }
2717
2718 pub fn force_flush(&mut self) -> Result<Epoch> {
2727 let saved = self.mutable_run_spill_bytes;
2728 self.mutable_run_spill_bytes = 1;
2729 let result = self.flush();
2730 self.mutable_run_spill_bytes = saved;
2731 result
2732 }
2733
2734 pub fn close(&mut self) -> Result<()> {
2741 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
2742 self.force_flush()?;
2743 }
2744 Ok(())
2745 }
2746
2747 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2754 let op = Op::Flush {
2755 table_id: self.table_id,
2756 flushed_epoch: epoch.0,
2757 };
2758 match &mut self.wal {
2759 WalSink::Private(w) => {
2760 w.append_system(op)?;
2761 w.sync()?;
2762 }
2763 WalSink::Shared(s) => {
2764 s.wal.lock().append_system(op)?;
2769 }
2770 }
2771 self.flushed_epoch = epoch.0;
2772 if matches!(self.wal, WalSink::Private(_)) {
2773 self.rotate_wal(epoch)?;
2774 }
2775 Ok(())
2776 }
2777
2778 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2782 let rows = self.mutable_run.drain_sorted();
2783 if rows.is_empty() {
2784 return Ok(());
2785 }
2786 let run_id = self.next_run_id;
2787 self.next_run_id += 1;
2788 let path = self.run_path(run_id);
2789 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2790 if let Some(kek) = &self.kek {
2791 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2792 }
2793 let header = writer.write(&path, &rows)?;
2794 self.run_refs.push(RunRef {
2795 run_id: run_id as u128,
2796 level: 0,
2797 epoch_created: epoch.0,
2798 row_count: header.row_count,
2799 });
2800 Ok(())
2801 }
2802
2803 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2807 self.mutable_run_spill_bytes = bytes.max(1);
2808 }
2809
2810 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2814 self.compaction_zstd_level = level;
2815 }
2816
2817 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2821 self.result_cache.lock().set_max_bytes(max_bytes);
2822 }
2823
2824 pub(crate) fn clear_result_cache(&mut self) {
2828 self.result_cache.lock().clear();
2829 }
2830
2831 pub fn mutable_run_len(&self) -> usize {
2833 self.mutable_run.len()
2834 }
2835
2836 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2839 self.mutable_run.drain_sorted()
2840 }
2841
2842 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2847 let epoch = self.commit()?;
2848 let n = batch.len();
2849 if n == 0 {
2850 return Ok(epoch);
2851 }
2852 let live_before = self.live_count;
2853 self.spill_mutable_run(epoch)?;
2857 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
2858 && self.indexes_complete
2859 && self.run_refs.is_empty()
2860 && self.memtable.is_empty()
2861 && self.mutable_run.is_empty();
2862 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2868 use rayon::prelude::*;
2869 self.schema
2870 .columns
2871 .par_iter()
2872 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
2873 .collect::<Vec<_>>()
2874 };
2875 drop(batch);
2876 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2881 self.validate_columns_not_null(&user_columns, n)?;
2882 let winner_idx = self
2883 .bulk_pk_winner_indices(&user_columns, n)
2884 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2885 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2886 match winner_idx.as_deref() {
2887 Some(idx) => {
2888 let compacted = user_columns
2889 .iter()
2890 .map(|(id, c)| (*id, c.gather(idx)))
2891 .collect();
2892 (compacted, idx.len())
2893 }
2894 None => (std::mem::take(&mut user_columns), n),
2895 };
2896 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2897 let first = self.allocator.alloc_range(write_n as u64).0;
2898 for rid in first..first + write_n as u64 {
2899 self.reservoir.offer(rid);
2900 }
2901 let run_id = self.next_run_id;
2902 self.next_run_id += 1;
2903 let path = self.run_path(run_id);
2904 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2905 .clean(true)
2906 .with_lz4()
2907 .with_native_endian();
2908 if let Some(kek) = &self.kek {
2909 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2910 }
2911 let header = writer.write_native(&path, &write_columns, write_n, first)?;
2912 self.run_refs.push(RunRef {
2913 run_id: run_id as u128,
2914 level: 0,
2915 epoch_created: epoch.0,
2916 row_count: header.row_count,
2917 });
2918 self.live_count = self.live_count.saturating_add(write_n as u64);
2919 if eager_index_build {
2920 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
2921 self.index_columns_bulk(&write_columns, &row_ids);
2922 self.indexes_complete = true;
2923 self.build_learned_ranges()?;
2924 } else {
2925 self.indexes_complete = false;
2926 }
2927 self.mark_flushed(epoch)?;
2928 self.persist_manifest(epoch)?;
2929 if eager_index_build {
2930 self.checkpoint_indexes(epoch);
2931 }
2932 self.clear_result_cache();
2933 Ok(epoch)
2934 }
2935
2936 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
2939 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
2940 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
2941 let segment_no = segment
2944 .file_stem()
2945 .and_then(|s| s.to_str())
2946 .and_then(|s| s.strip_prefix("seg-"))
2947 .and_then(|s| s.parse::<u64>().ok())
2948 .unwrap_or(0);
2949 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
2950 wal.set_sync_byte_threshold(self.sync_byte_threshold);
2951 wal.sync()?;
2952 self.wal = WalSink::Private(wal);
2953 Ok(())
2954 }
2955
2956 pub(crate) fn invalidate_pending_cache(&mut self) {
2961 self.result_cache
2962 .lock()
2963 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
2964 self.pending_delete_rids.clear();
2965 self.pending_put_cols.clear();
2966 }
2967
2968 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
2969 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
2970 m.current_epoch = epoch.0;
2971 m.next_row_id = self.allocator.current().0;
2972 m.runs = self.run_refs.clone();
2973 m.live_count = self.live_count;
2974 m.global_idx_epoch = self.global_idx_epoch;
2975 m.flushed_epoch = self.flushed_epoch;
2976 m.retiring = self.retiring.clone();
2977 m.auto_inc_next = match self.auto_inc {
2981 Some(ai) if ai.seeded => ai.next,
2982 _ => 0,
2983 };
2984 let meta_dek = self.manifest_meta_dek();
2985 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
2986 Ok(())
2987 }
2988
2989 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
2995 if !self.indexes_complete {
2998 return;
2999 }
3000 let snap = global_idx::IndexSnapshot {
3001 hot: &self.hot,
3002 bitmap: &self.bitmap,
3003 ann: &self.ann,
3004 fm: &self.fm,
3005 sparse: &self.sparse,
3006 minhash: &self.minhash,
3007 learned_range: &self.learned_range,
3008 };
3009 let idx_dek = self.idx_dek();
3011 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3012 .is_ok()
3013 {
3014 self.global_idx_epoch = epoch.0;
3015 let _ = self.persist_manifest(epoch);
3016 }
3017 }
3018
3019 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3022 self.global_idx_epoch = 0;
3023 global_idx::remove(&self.dir);
3024 let _ = self.persist_manifest(self.epoch.visible());
3025 }
3026
3027 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3030 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3031 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3032 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3033 best = Some((epoch, row));
3034 }
3035 }
3036 for rr in &self.run_refs {
3037 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3038 continue;
3039 };
3040 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3041 continue;
3042 };
3043 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3044 best = Some((epoch, row));
3045 }
3046 }
3047 match best {
3048 Some((_, r)) if r.deleted => None,
3049 Some((_, r)) => Some(r),
3050 None => None,
3051 }
3052 }
3053
3054 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3058 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3059 let mut fold = |row: Row| {
3060 best.entry(row.row_id.0)
3061 .and_modify(|e| {
3062 if row.committed_epoch > e.0 {
3063 *e = (row.committed_epoch, row.clone());
3064 }
3065 })
3066 .or_insert_with(|| (row.committed_epoch, row));
3067 };
3068 for row in self.memtable.visible_versions(snapshot.epoch) {
3069 fold(row);
3070 }
3071 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3072 fold(row);
3073 }
3074 for rr in &self.run_refs {
3075 let mut reader = self.open_reader(rr.run_id)?;
3076 for row in reader.visible_versions(snapshot.epoch)? {
3077 fold(row);
3078 }
3079 }
3080 let mut out: Vec<Row> = best
3081 .into_values()
3082 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3083 .collect();
3084 out.sort_by_key(|r| r.row_id);
3085 Ok(out)
3086 }
3087
3088 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3095 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3096 let rr = self.run_refs[0].clone();
3097 let mut reader = self.open_reader(rr.run_id)?;
3098 let idxs = reader.visible_indices(snapshot.epoch)?;
3099 let mut cols = Vec::with_capacity(self.schema.columns.len());
3100 for cdef in &self.schema.columns {
3101 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3102 }
3103 return Ok(cols);
3104 }
3105 let rows = self.visible_rows(snapshot)?;
3107 let mut cols: Vec<(u16, Vec<Value>)> = self
3108 .schema
3109 .columns
3110 .iter()
3111 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3112 .collect();
3113 for r in &rows {
3114 for (cid, vec) in cols.iter_mut() {
3115 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3116 }
3117 }
3118 Ok(cols)
3119 }
3120
3121 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3123 self.hot.get(key)
3124 }
3125
3126 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3131 self.ensure_indexes_complete()?;
3132 let snapshot = self.snapshot();
3133 crate::trace::QueryTrace::record(|t| {
3134 t.run_count = self.run_refs.len();
3135 t.memtable_rows = self.memtable.len();
3136 t.mutable_run_rows = self.mutable_run.len();
3137 });
3138 if q.conditions.is_empty() {
3142 crate::trace::QueryTrace::record(|t| {
3143 t.scan_mode = crate::trace::ScanMode::Materialized;
3144 t.row_materialized = true;
3145 });
3146 return self.visible_rows(snapshot);
3147 }
3148 crate::trace::QueryTrace::record(|t| {
3149 t.conditions_pushed = q.conditions.len();
3150 t.scan_mode = crate::trace::ScanMode::Materialized;
3151 t.row_materialized = true;
3152 });
3153 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3160 ordered.sort_by_key(|c| condition_cost_rank(c));
3161 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3162 for c in &ordered {
3163 let s = self.resolve_condition(c, snapshot)?;
3164 let empty = s.is_empty();
3165 sets.push(s);
3166 if empty {
3167 break;
3168 }
3169 }
3170 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3171 self.rows_for_rids(&rids, snapshot)
3172 }
3173
3174 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3179 use std::collections::HashMap;
3180 let mut rows = Vec::with_capacity(rids.len());
3181 let tier_size = self.memtable.len() + self.mutable_run.len();
3198 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3199 if rids.len().saturating_mul(24) < tier_size {
3200 for &rid in rids {
3201 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3202 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3203 let newest = match (mem, mrun) {
3204 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3205 (Some((_, mr)), None) => Some(mr),
3206 (None, Some((_, rr))) => Some(rr),
3207 (None, None) => None,
3208 };
3209 if let Some(row) = newest {
3210 overlay.insert(rid, row);
3211 }
3212 }
3213 } else {
3214 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3215 overlay
3216 .entry(row.row_id.0)
3217 .and_modify(|e| {
3218 if row.committed_epoch > e.committed_epoch {
3219 *e = row.clone();
3220 }
3221 })
3222 .or_insert(row);
3223 };
3224 for row in self.memtable.visible_versions(snapshot.epoch) {
3225 fold_newest(row, &mut overlay);
3226 }
3227 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3228 fold_newest(row, &mut overlay);
3229 }
3230 }
3231 if self.run_refs.len() == 1 {
3232 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3233 if rids.len().saturating_mul(24) < reader.row_count() {
3241 for &rid in rids {
3242 if let Some(r) = overlay.get(&rid) {
3243 if !r.deleted {
3244 rows.push(r.clone());
3245 }
3246 continue;
3247 }
3248 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3249 if !row.deleted {
3250 rows.push(row);
3251 }
3252 }
3253 }
3254 return Ok(rows);
3255 }
3256 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3265 enum Src {
3268 Overlay,
3269 Run,
3270 }
3271 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3272 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3273 for rid in rids {
3274 if overlay.contains_key(rid) {
3275 plan.push(Src::Overlay);
3276 continue;
3277 }
3278 match vis_rids.binary_search(&(*rid as i64)) {
3279 Ok(i) => {
3280 plan.push(Src::Run);
3281 fetch.push(positions[i]);
3282 }
3283 Err(_) => { }
3284 }
3285 }
3286 let fetched = reader.materialize_batch(&fetch)?;
3287 let mut fetched_iter = fetched.into_iter();
3288 for (rid, src) in rids.iter().zip(plan) {
3289 match src {
3290 Src::Overlay => {
3291 if let Some(r) = overlay.get(rid) {
3292 if !r.deleted {
3293 rows.push(r.clone());
3294 }
3295 }
3296 }
3297 Src::Run => {
3298 if let Some(row) = fetched_iter.next() {
3299 if !row.deleted {
3300 rows.push(row);
3301 }
3302 }
3303 }
3304 }
3305 }
3306 return Ok(rows);
3307 }
3308 let mut readers: Vec<_> = self
3312 .run_refs
3313 .iter()
3314 .map(|rr| self.open_reader(rr.run_id))
3315 .collect::<Result<Vec<_>>>()?;
3316 for rid in rids {
3317 if let Some(r) = overlay.get(rid) {
3318 if !r.deleted {
3319 rows.push(r.clone());
3320 }
3321 continue;
3322 }
3323 let mut best: Option<(Epoch, Row)> = None;
3324 for reader in readers.iter_mut() {
3325 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3326 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3327 best = Some((epoch, row));
3328 }
3329 }
3330 }
3331 if let Some((_, r)) = best {
3332 if !r.deleted {
3333 rows.push(r);
3334 }
3335 }
3336 }
3337 Ok(rows)
3338 }
3339
3340 pub fn indexes_complete(&self) -> bool {
3350 self.indexes_complete
3351 }
3352
3353 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3355 self.index_build_policy
3356 }
3357
3358 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3362 self.index_build_policy = policy;
3363 }
3364
3365 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3370 if !self.indexes_complete {
3374 return None;
3375 }
3376 let b = self.bitmap.get(&column_id)?;
3377 let result: Vec<Vec<u8>> = b
3378 .keys()
3379 .into_iter()
3380 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3381 .cloned()
3382 .collect();
3383 Some(result)
3384 }
3385
3386 pub fn fk_join_row_ids(
3387 &self,
3388 fk_column_id: u16,
3389 pk_values: &[Vec<u8>],
3390 fk_conditions: &[crate::query::Condition],
3391 snapshot: Snapshot,
3392 ) -> Result<Vec<u64>> {
3393 let Some(b) = self.bitmap.get(&fk_column_id) else {
3394 return Ok(Vec::new());
3395 };
3396 let mut join_set = {
3397 let mut acc = roaring::RoaringBitmap::new();
3398 for v in pk_values {
3399 acc |= b.get(v);
3400 }
3401 RowIdSet::from_roaring(acc)
3402 };
3403 if !fk_conditions.is_empty() {
3404 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3405 sets.push(join_set);
3406 for c in fk_conditions {
3407 sets.push(self.resolve_condition(c, snapshot)?);
3408 }
3409 join_set = RowIdSet::intersect_many(sets);
3410 }
3411 Ok(join_set.into_sorted_vec())
3412 }
3413
3414 pub fn fk_join_count(
3420 &self,
3421 fk_column_id: u16,
3422 pk_values: &[Vec<u8>],
3423 fk_conditions: &[crate::query::Condition],
3424 snapshot: Snapshot,
3425 ) -> Result<u64> {
3426 let Some(b) = self.bitmap.get(&fk_column_id) else {
3427 return Ok(0);
3428 };
3429 let mut acc = roaring::RoaringBitmap::new();
3430 for v in pk_values {
3431 acc |= b.get(v);
3432 }
3433 if fk_conditions.is_empty() {
3434 return Ok(acc.len());
3435 }
3436 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3437 sets.push(RowIdSet::from_roaring(acc));
3438 for c in fk_conditions {
3439 sets.push(self.resolve_condition(c, snapshot)?);
3440 }
3441 Ok(RowIdSet::intersect_many(sets).len() as u64)
3442 }
3443
3444 fn resolve_condition(
3449 &self,
3450 c: &crate::query::Condition,
3451 snapshot: Snapshot,
3452 ) -> Result<RowIdSet> {
3453 use crate::query::Condition;
3454 Ok(match c {
3455 Condition::Pk(key) => {
3456 let lookup = self
3457 .schema
3458 .primary_key()
3459 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3460 .unwrap_or_else(|| key.clone());
3461 self.hot
3462 .get(&lookup)
3463 .map(|r| RowIdSet::one(r.0))
3464 .unwrap_or_else(RowIdSet::empty)
3465 }
3466 Condition::BitmapEq { column_id, value } => {
3467 let lookup = self.index_lookup_key_bytes(*column_id, value);
3468 self.bitmap
3469 .get(column_id)
3470 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3471 .unwrap_or_else(RowIdSet::empty)
3472 }
3473 Condition::BitmapIn { column_id, values } => {
3474 let bm = self.bitmap.get(column_id);
3475 let mut acc = roaring::RoaringBitmap::new();
3476 if let Some(b) = bm {
3477 for v in values {
3478 let lookup = self.index_lookup_key_bytes(*column_id, v);
3479 acc |= b.get(&lookup);
3480 }
3481 }
3482 RowIdSet::from_roaring(acc)
3483 }
3484 Condition::BytesPrefix { column_id, prefix } => {
3485 if let Some(b) = self.bitmap.get(column_id) {
3490 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3491 let mut acc = roaring::RoaringBitmap::new();
3492 for key in b.keys() {
3493 if key.starts_with(&lookup_prefix) {
3494 acc |= b.get(key);
3495 }
3496 }
3497 RowIdSet::from_roaring(acc)
3498 } else {
3499 RowIdSet::empty()
3500 }
3501 }
3502 Condition::FmContains { column_id, pattern } => self
3503 .fm
3504 .get(column_id)
3505 .map(|f| {
3506 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3507 })
3508 .unwrap_or_else(RowIdSet::empty),
3509 Condition::FmContainsAll {
3510 column_id,
3511 patterns,
3512 } => {
3513 if let Some(f) = self.fm.get(column_id) {
3516 let sets: Vec<RowIdSet> = patterns
3517 .iter()
3518 .map(|pat| {
3519 RowIdSet::from_unsorted(
3520 f.locate(pat).into_iter().map(|r| r.0).collect(),
3521 )
3522 })
3523 .collect();
3524 RowIdSet::intersect_many(sets)
3525 } else {
3526 RowIdSet::empty()
3527 }
3528 }
3529 Condition::Ann {
3530 column_id,
3531 query,
3532 k,
3533 } => self
3534 .ann
3535 .get(column_id)
3536 .map(|a| {
3537 RowIdSet::from_unsorted(
3538 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3539 )
3540 })
3541 .unwrap_or_else(RowIdSet::empty),
3542 Condition::SparseMatch {
3543 column_id,
3544 query,
3545 k,
3546 } => self
3547 .sparse
3548 .get(column_id)
3549 .map(|s| {
3550 RowIdSet::from_unsorted(
3551 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3552 )
3553 })
3554 .unwrap_or_else(RowIdSet::empty),
3555 Condition::MinHashSimilar {
3556 column_id,
3557 query,
3558 k,
3559 } => self
3560 .minhash
3561 .get(column_id)
3562 .map(|mh| {
3563 RowIdSet::from_unsorted(
3564 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3565 )
3566 })
3567 .unwrap_or_else(RowIdSet::empty),
3568 Condition::Range { column_id, lo, hi } => {
3569 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3578 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3579 } else if self.run_refs.len() == 1 {
3580 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3581 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3582 } else {
3583 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3584 };
3585 set.remove_many(self.overlay_rid_set(snapshot));
3586 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3587 set
3588 }
3589 Condition::RangeF64 {
3590 column_id,
3591 lo,
3592 lo_inclusive,
3593 hi,
3594 hi_inclusive,
3595 } => {
3596 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3599 RowIdSet::from_unsorted(
3600 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3601 .into_iter()
3602 .collect(),
3603 )
3604 } else if self.run_refs.len() == 1 {
3605 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3606 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3607 } else {
3608 return self.range_scan_f64(
3609 *column_id,
3610 *lo,
3611 *lo_inclusive,
3612 *hi,
3613 *hi_inclusive,
3614 snapshot,
3615 );
3616 };
3617 set.remove_many(self.overlay_rid_set(snapshot));
3618 self.range_scan_overlay_f64(
3619 &mut set,
3620 *column_id,
3621 *lo,
3622 *lo_inclusive,
3623 *hi,
3624 *hi_inclusive,
3625 snapshot,
3626 );
3627 set
3628 }
3629 Condition::IsNull { column_id } => {
3630 let mut set = if self.run_refs.len() == 1 {
3631 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3632 r.null_row_id_set(*column_id, true)?
3633 } else {
3634 return self.null_scan(*column_id, true, snapshot);
3635 };
3636 set.remove_many(self.overlay_rid_set(snapshot));
3637 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3638 set
3639 }
3640 Condition::IsNotNull { column_id } => {
3641 let mut set = if self.run_refs.len() == 1 {
3642 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3643 r.null_row_id_set(*column_id, false)?
3644 } else {
3645 return self.null_scan(*column_id, false, snapshot);
3646 };
3647 set.remove_many(self.overlay_rid_set(snapshot));
3648 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3649 set
3650 }
3651 })
3652 }
3653
3654 fn range_scan_i64(
3662 &self,
3663 column_id: u16,
3664 lo: i64,
3665 hi: i64,
3666 snapshot: Snapshot,
3667 ) -> Result<RowIdSet> {
3668 let mut row_ids = Vec::new();
3669 let overlay_rids = self.overlay_rid_set(snapshot);
3670 for rr in &self.run_refs {
3671 let mut reader = self.open_reader(rr.run_id)?;
3672 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3673 for rid in matched {
3674 if !overlay_rids.contains(&rid) {
3675 row_ids.push(rid);
3676 }
3677 }
3678 }
3679 let mut s = RowIdSet::from_unsorted(row_ids);
3680 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3681 Ok(s)
3682 }
3683
3684 fn range_scan_f64(
3687 &self,
3688 column_id: u16,
3689 lo: f64,
3690 lo_inclusive: bool,
3691 hi: f64,
3692 hi_inclusive: bool,
3693 snapshot: Snapshot,
3694 ) -> Result<RowIdSet> {
3695 let mut row_ids = Vec::new();
3696 let overlay_rids = self.overlay_rid_set(snapshot);
3697 for rr in &self.run_refs {
3698 let mut reader = self.open_reader(rr.run_id)?;
3699 let matched = reader.range_row_ids_visible_f64(
3700 column_id,
3701 lo,
3702 lo_inclusive,
3703 hi,
3704 hi_inclusive,
3705 snapshot.epoch,
3706 )?;
3707 for rid in matched {
3708 if !overlay_rids.contains(&rid) {
3709 row_ids.push(rid);
3710 }
3711 }
3712 }
3713 let mut s = RowIdSet::from_unsorted(row_ids);
3714 self.range_scan_overlay_f64(
3715 &mut s,
3716 column_id,
3717 lo,
3718 lo_inclusive,
3719 hi,
3720 hi_inclusive,
3721 snapshot,
3722 );
3723 Ok(s)
3724 }
3725
3726 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3728 let mut s = HashSet::new();
3729 for row in self.memtable.visible_versions(snapshot.epoch) {
3730 s.insert(row.row_id.0);
3731 }
3732 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3733 s.insert(row.row_id.0);
3734 }
3735 s
3736 }
3737
3738 fn range_scan_overlay_i64(
3739 &self,
3740 s: &mut RowIdSet,
3741 column_id: u16,
3742 lo: i64,
3743 hi: i64,
3744 snapshot: Snapshot,
3745 ) {
3746 let mut newest: HashMap<u64, &Row> = HashMap::new();
3751 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3752 let memtable = self.memtable.visible_versions(snapshot.epoch);
3753 for r in &mutable {
3754 newest.entry(r.row_id.0).or_insert(r);
3755 }
3756 for r in &memtable {
3757 newest.insert(r.row_id.0, r);
3758 }
3759 for row in newest.values() {
3760 if !row.deleted {
3761 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3762 if *v >= lo && *v <= hi {
3763 s.insert(row.row_id.0);
3764 }
3765 }
3766 }
3767 }
3768 }
3769
3770 #[allow(clippy::too_many_arguments)]
3771 fn range_scan_overlay_f64(
3772 &self,
3773 s: &mut RowIdSet,
3774 column_id: u16,
3775 lo: f64,
3776 lo_inclusive: bool,
3777 hi: f64,
3778 hi_inclusive: bool,
3779 snapshot: Snapshot,
3780 ) {
3781 let mut newest: HashMap<u64, &Row> = HashMap::new();
3784 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3785 let memtable = self.memtable.visible_versions(snapshot.epoch);
3786 for r in &mutable {
3787 newest.entry(r.row_id.0).or_insert(r);
3788 }
3789 for r in &memtable {
3790 newest.insert(r.row_id.0, r);
3791 }
3792 for row in newest.values() {
3793 if !row.deleted {
3794 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3795 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3796 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3797 if ok_lo && ok_hi {
3798 s.insert(row.row_id.0);
3799 }
3800 }
3801 }
3802 }
3803 }
3804
3805 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3808 let mut row_ids = Vec::new();
3809 let overlay_rids = self.overlay_rid_set(snapshot);
3810 for rr in &self.run_refs {
3811 let mut reader = self.open_reader(rr.run_id)?;
3812 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3813 for rid in matched {
3814 if !overlay_rids.contains(&rid) {
3815 row_ids.push(rid);
3816 }
3817 }
3818 }
3819 let mut s = RowIdSet::from_unsorted(row_ids);
3820 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3821 Ok(s)
3822 }
3823
3824 fn null_scan_overlay(
3828 &self,
3829 s: &mut RowIdSet,
3830 column_id: u16,
3831 want_nulls: bool,
3832 snapshot: Snapshot,
3833 ) {
3834 let mut newest: HashMap<u64, &Row> = HashMap::new();
3835 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3836 let memtable = self.memtable.visible_versions(snapshot.epoch);
3837 for r in &mutable {
3838 newest.entry(r.row_id.0).or_insert(r);
3839 }
3840 for r in &memtable {
3841 newest.insert(r.row_id.0, r);
3842 }
3843 for row in newest.values() {
3844 if row.deleted {
3845 continue;
3846 }
3847 let is_null = !row.columns.contains_key(&column_id)
3848 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3849 if is_null == want_nulls {
3850 s.insert(row.row_id.0);
3851 }
3852 }
3853 }
3854
3855 pub fn snapshot(&self) -> Snapshot {
3856 Snapshot::at(self.epoch.visible())
3857 }
3858
3859 pub fn pin_snapshot(&mut self) -> Snapshot {
3862 let e = self.epoch.visible();
3863 *self.pinned.entry(e).or_insert(0) += 1;
3864 Snapshot::at(e)
3865 }
3866
3867 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3869 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3870 *count -= 1;
3871 if *count == 0 {
3872 self.pinned.remove(&snap.epoch);
3873 }
3874 }
3875 }
3876
3877 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3887 let local = self.pinned.keys().next().copied();
3888 let global = self.snapshots.min_pinned();
3889 match (local, global) {
3890 (Some(a), Some(b)) => Some(a.min(b)),
3891 (Some(a), None) => Some(a),
3892 (None, b) => b,
3893 }
3894 }
3895
3896 pub fn current_epoch(&self) -> Epoch {
3897 self.epoch.visible()
3898 }
3899
3900 pub fn memtable_len(&self) -> usize {
3901 self.memtable.len()
3902 }
3903
3904 pub fn count(&self) -> u64 {
3907 self.live_count
3908 }
3909
3910 pub fn count_conditions(
3914 &mut self,
3915 conditions: &[crate::query::Condition],
3916 snapshot: Snapshot,
3917 ) -> Result<Option<u64>> {
3918 use crate::query::Condition;
3919 if conditions.is_empty() {
3920 return Ok(Some(self.live_count));
3921 }
3922 let served = |c: &Condition| {
3923 matches!(
3924 c,
3925 Condition::Pk(_)
3926 | Condition::BitmapEq { .. }
3927 | Condition::BitmapIn { .. }
3928 | Condition::BytesPrefix { .. }
3929 | Condition::FmContains { .. }
3930 | Condition::FmContainsAll { .. }
3931 | Condition::Ann { .. }
3932 | Condition::Range { .. }
3933 | Condition::RangeF64 { .. }
3934 | Condition::SparseMatch { .. }
3935 | Condition::MinHashSimilar { .. }
3936 | Condition::IsNull { .. }
3937 | Condition::IsNotNull { .. }
3938 )
3939 };
3940 if !conditions.iter().all(served) {
3941 return Ok(None);
3942 }
3943 self.ensure_indexes_complete()?;
3944 let mut sets = Vec::with_capacity(conditions.len());
3945 for condition in conditions {
3946 sets.push(self.resolve_condition(condition, snapshot)?);
3947 }
3948 let mut rids = RowIdSet::intersect_many(sets);
3949 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
3959 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
3960 }
3961 let count = rids.len() as u64;
3962 crate::trace::QueryTrace::record(|t| {
3963 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
3964 t.survivor_count = Some(count as usize);
3965 t.conditions_pushed = conditions.len();
3966 });
3967 Ok(Some(count))
3968 }
3969
3970 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
3975 let mut out = Vec::new();
3976 for row in self.memtable.visible_versions(snapshot.epoch) {
3977 if row.deleted {
3978 out.push(row.row_id.0);
3979 }
3980 }
3981 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3982 if row.deleted {
3983 out.push(row.row_id.0);
3984 }
3985 }
3986 out
3987 }
3988
3989 pub fn bulk_load_columns(
3998 &mut self,
3999 user_columns: Vec<(u16, columnar::NativeColumn)>,
4000 ) -> Result<Epoch> {
4001 self.bulk_load_columns_with(user_columns, 3, false, true)
4002 }
4003
4004 pub fn bulk_load_fast(
4011 &mut self,
4012 user_columns: Vec<(u16, columnar::NativeColumn)>,
4013 ) -> Result<Epoch> {
4014 self.bulk_load_columns_with(user_columns, -1, true, false)
4015 }
4016
4017 fn bulk_load_columns_with(
4018 &mut self,
4019 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4020 zstd_level: i32,
4021 force_plain: bool,
4022 lz4: bool,
4023 ) -> Result<Epoch> {
4024 let epoch = self.commit()?;
4025 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4026 if n == 0 {
4027 return Ok(epoch);
4028 }
4029 let live_before = self.live_count;
4030 self.spill_mutable_run(epoch)?;
4032 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4033 && self.indexes_complete
4034 && self.run_refs.is_empty()
4035 && self.memtable.is_empty()
4036 && self.mutable_run.is_empty();
4037 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4040 self.validate_columns_not_null(&user_columns, n)?;
4041 let winner_idx = self
4042 .bulk_pk_winner_indices(&user_columns, n)
4043 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
4044 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4045 match winner_idx.as_deref() {
4046 Some(idx) => {
4047 let compacted = user_columns
4048 .iter()
4049 .map(|(id, c)| (*id, c.gather(idx)))
4050 .collect();
4051 (compacted, idx.len())
4052 }
4053 None => (user_columns, n),
4054 };
4055 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4056 let first = self.allocator.alloc_range(write_n as u64).0;
4057 for rid in first..first + write_n as u64 {
4058 self.reservoir.offer(rid);
4059 }
4060 let run_id = self.next_run_id;
4061 self.next_run_id += 1;
4062 let path = self.run_path(run_id);
4063 let mut writer =
4064 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4065 if force_plain {
4066 writer = writer.with_plain();
4067 } else if lz4 {
4068 writer = writer.with_lz4();
4071 } else {
4072 writer = writer.with_zstd_level(zstd_level);
4073 }
4074 if let Some(kek) = &self.kek {
4075 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4076 }
4077 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4078 self.run_refs.push(RunRef {
4079 run_id: run_id as u128,
4080 level: 0,
4081 epoch_created: epoch.0,
4082 row_count: header.row_count,
4083 });
4084 self.live_count = self.live_count.saturating_add(write_n as u64);
4085 if eager_index_build {
4086 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4087 self.index_columns_bulk(&write_columns, &row_ids);
4088 self.indexes_complete = true;
4089 self.build_learned_ranges()?;
4090 } else {
4091 self.indexes_complete = false;
4095 }
4096 self.mark_flushed(epoch)?;
4097 self.persist_manifest(epoch)?;
4098 if eager_index_build {
4099 self.checkpoint_indexes(epoch);
4100 }
4101 self.clear_result_cache();
4102 Ok(epoch)
4103 }
4104
4105 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4123 let n = row_ids.len();
4124 if n == 0 {
4125 return;
4126 }
4127 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4128 columns.iter().map(|(id, c)| (*id, c)).collect();
4129 let ty_of: std::collections::HashMap<u16, TypeId> =
4130 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
4131 let pk_id = self.schema.primary_key().map(|c| c.id);
4132
4133 for (i, &rid) in row_ids.iter().enumerate() {
4134 let row_id = RowId(rid);
4135 if let Some(pid) = pk_id {
4136 if let Some(col) = by_id.get(&pid) {
4137 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
4138 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4139 self.insert_hot_pk(key, row_id);
4140 }
4141 }
4142 }
4143 for idef in &self.schema.indexes {
4144 let Some(col) = by_id.get(&idef.column_id) else {
4145 continue;
4146 };
4147 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
4148 match idef.kind {
4149 IndexKind::Bitmap => {
4150 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4151 if let Some(key) =
4152 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4153 {
4154 b.insert(key, row_id);
4155 }
4156 }
4157 }
4158 IndexKind::FmIndex => {
4159 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4160 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4161 f.insert(bytes.to_vec(), row_id);
4162 }
4163 }
4164 }
4165 IndexKind::Sparse => {
4166 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4167 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4168 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4169 s.insert(&terms, row_id);
4170 }
4171 }
4172 }
4173 }
4174 IndexKind::MinHash => {
4175 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4176 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4177 let tokens = crate::index::token_hashes_from_bytes(bytes);
4178 mh.insert(&tokens, row_id);
4179 }
4180 }
4181 }
4182 _ => {}
4183 }
4184 }
4185 }
4186 }
4187
4188 pub fn visible_columns_native(
4193 &self,
4194 snapshot: Snapshot,
4195 projection: Option<&[u16]>,
4196 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4197 let wanted: Vec<u16> = match projection {
4198 Some(p) => p.to_vec(),
4199 None => self.schema.columns.iter().map(|c| c.id).collect(),
4200 };
4201 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4202 let rr = self.run_refs[0].clone();
4203 let mut reader = self.open_reader(rr.run_id)?;
4204 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4205 let all_visible = idxs.len() == reader.row_count();
4206 if reader.has_mmap() {
4212 use rayon::prelude::*;
4213 let valid: Vec<u16> = wanted
4216 .iter()
4217 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4218 .copied()
4219 .collect();
4220 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4222 .par_iter()
4223 .filter_map(|cid| {
4224 reader
4225 .column_native_shared(*cid)
4226 .ok()
4227 .map(|col| (*cid, col))
4228 })
4229 .collect();
4230 let cols = decoded
4231 .into_iter()
4232 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4233 .collect();
4234 return Ok(cols);
4235 }
4236 let mut cols = Vec::with_capacity(wanted.len());
4237 for cid in &wanted {
4238 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4239 Some(c) => c,
4240 None => continue,
4241 };
4242 let col = reader.column_native(cdef.id)?;
4243 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4244 }
4245 return Ok(cols);
4246 }
4247 let vcols = self.visible_columns(snapshot)?;
4248 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4249 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4250 .into_iter()
4251 .filter(|(id, _)| want_set.contains(id))
4252 .map(|(id, vals)| {
4253 let ty = self
4254 .schema
4255 .columns
4256 .iter()
4257 .find(|c| c.id == id)
4258 .map(|c| c.ty)
4259 .unwrap_or(TypeId::Bytes);
4260 (id, columnar::values_to_native(ty, &vals))
4261 })
4262 .collect();
4263 Ok(out)
4264 }
4265
4266 pub fn run_count(&self) -> usize {
4267 self.run_refs.len()
4268 }
4269
4270 pub fn memtable_is_empty(&self) -> bool {
4272 self.memtable.is_empty()
4273 }
4274
4275 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4279 self.page_cache.stats()
4280 }
4281
4282 pub fn reset_page_cache_stats(&self) {
4284 self.page_cache.reset_stats();
4285 }
4286
4287 pub fn run_ids(&self) -> Vec<u128> {
4290 self.run_refs.iter().map(|r| r.run_id).collect()
4291 }
4292
4293 pub fn single_run_is_clean(&self) -> bool {
4297 if self.run_refs.len() != 1 {
4298 return false;
4299 }
4300 self.open_reader(self.run_refs[0].run_id)
4301 .map(|r| r.is_clean())
4302 .unwrap_or(false)
4303 }
4304
4305 fn resolve_footprint(
4312 &self,
4313 conditions: &[crate::query::Condition],
4314 _snapshot: Snapshot,
4315 ) -> roaring::RoaringBitmap {
4316 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4317 return roaring::RoaringBitmap::new();
4318 }
4319 if self.run_refs.is_empty() {
4320 return roaring::RoaringBitmap::new();
4321 }
4322 if self.run_refs.len() == 1 {
4324 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4325 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4326 return rids.to_roaring_lossy();
4327 }
4328 }
4329 }
4330 roaring::RoaringBitmap::new()
4331 }
4332
4333 pub fn query_columns_native_cached(
4344 &mut self,
4345 conditions: &[crate::query::Condition],
4346 projection: Option<&[u16]>,
4347 snapshot: Snapshot,
4348 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4349 if conditions.is_empty() {
4350 return self.query_columns_native(conditions, projection, snapshot);
4351 }
4352 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4356 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4357 crate::trace::QueryTrace::record(|t| {
4358 t.result_cache_hit = true;
4359 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4360 });
4361 return Ok(Some((*hit).clone()));
4362 }
4363 let res = self.query_columns_native(conditions, projection, snapshot)?;
4364 if let Some(cols) = &res {
4365 let footprint = self.resolve_footprint(conditions, snapshot);
4366 let condition_cols = crate::query::condition_columns(conditions);
4367 self.result_cache.lock().insert(
4368 key,
4369 CachedEntry {
4370 data: CachedData::Columns(Arc::new(cols.clone())),
4371 footprint,
4372 condition_cols,
4373 },
4374 );
4375 }
4376 Ok(res)
4377 }
4378
4379 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4384 if q.conditions.is_empty() {
4385 return self.query(q);
4386 }
4387 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4388 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4389 crate::trace::QueryTrace::record(|t| {
4390 t.result_cache_hit = true;
4391 t.scan_mode = crate::trace::ScanMode::Materialized;
4392 });
4393 return Ok((*hit).clone());
4394 }
4395 let rows = self.query(q)?;
4396 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4397 let condition_cols = crate::query::condition_columns(&q.conditions);
4398 self.result_cache.lock().insert(
4399 key,
4400 CachedEntry {
4401 data: CachedData::Rows(Arc::new(rows.clone())),
4402 footprint,
4403 condition_cols,
4404 },
4405 );
4406 Ok(rows)
4407 }
4408
4409 #[allow(clippy::type_complexity)]
4424 pub fn query_columns_native_traced(
4425 &mut self,
4426 conditions: &[crate::query::Condition],
4427 projection: Option<&[u16]>,
4428 snapshot: Snapshot,
4429 ) -> Result<(
4430 Option<Vec<(u16, columnar::NativeColumn)>>,
4431 crate::trace::QueryTrace,
4432 )> {
4433 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4434 self.query_columns_native(conditions, projection, snapshot)
4435 });
4436 Ok((result?, trace))
4437 }
4438
4439 #[allow(clippy::type_complexity)]
4442 pub fn query_columns_native_cached_traced(
4443 &mut self,
4444 conditions: &[crate::query::Condition],
4445 projection: Option<&[u16]>,
4446 snapshot: Snapshot,
4447 ) -> Result<(
4448 Option<Vec<(u16, columnar::NativeColumn)>>,
4449 crate::trace::QueryTrace,
4450 )> {
4451 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4452 self.query_columns_native_cached(conditions, projection, snapshot)
4453 });
4454 Ok((result?, trace))
4455 }
4456
4457 pub fn native_page_cursor_traced(
4459 &self,
4460 snapshot: Snapshot,
4461 projection: Vec<(u16, TypeId)>,
4462 conditions: &[crate::query::Condition],
4463 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4464 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4465 self.native_page_cursor(snapshot, projection, conditions)
4466 });
4467 Ok((result?, trace))
4468 }
4469
4470 pub fn native_multi_run_cursor_traced(
4472 &self,
4473 snapshot: Snapshot,
4474 projection: Vec<(u16, TypeId)>,
4475 conditions: &[crate::query::Condition],
4476 ) -> Result<(
4477 Option<crate::cursor::MultiRunCursor>,
4478 crate::trace::QueryTrace,
4479 )> {
4480 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4481 self.native_multi_run_cursor(snapshot, projection, conditions)
4482 });
4483 Ok((result?, trace))
4484 }
4485
4486 pub fn count_conditions_traced(
4488 &mut self,
4489 conditions: &[crate::query::Condition],
4490 snapshot: Snapshot,
4491 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4492 let (result, trace) =
4493 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4494 Ok((result?, trace))
4495 }
4496
4497 pub fn query_traced(
4499 &mut self,
4500 q: &crate::query::Query,
4501 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4502 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4503 Ok((result?, trace))
4504 }
4505
4506 pub fn query_columns_native(
4511 &mut self,
4512 conditions: &[crate::query::Condition],
4513 projection: Option<&[u16]>,
4514 snapshot: Snapshot,
4515 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4516 use crate::query::Condition;
4517 if conditions.is_empty() {
4518 return Ok(None);
4519 }
4520 self.ensure_indexes_complete()?;
4521
4522 let served = |c: &Condition| {
4527 matches!(
4528 c,
4529 Condition::Pk(_)
4530 | Condition::BitmapEq { .. }
4531 | Condition::BitmapIn { .. }
4532 | Condition::BytesPrefix { .. }
4533 | Condition::FmContains { .. }
4534 | Condition::FmContainsAll { .. }
4535 | Condition::Ann { .. }
4536 | Condition::Range { .. }
4537 | Condition::RangeF64 { .. }
4538 | Condition::SparseMatch { .. }
4539 | Condition::MinHashSimilar { .. }
4540 | Condition::IsNull { .. }
4541 | Condition::IsNotNull { .. }
4542 )
4543 };
4544 if !conditions.iter().all(served) {
4545 return Ok(None);
4546 }
4547 let fast_path =
4548 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4549 crate::trace::QueryTrace::record(|t| {
4550 t.run_count = self.run_refs.len();
4551 t.memtable_rows = self.memtable.len();
4552 t.mutable_run_rows = self.mutable_run.len();
4553 t.conditions_pushed = conditions.len();
4554 t.learned_range_used = conditions.iter().any(|c| match c {
4555 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4556 self.learned_range.contains_key(column_id)
4557 }
4558 _ => false,
4559 });
4560 });
4561 let col_ids: Vec<u16> = projection
4563 .map(|p| p.to_vec())
4564 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4565 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4566 .iter()
4567 .map(|&cid| {
4568 let ty = self
4569 .schema
4570 .columns
4571 .iter()
4572 .find(|c| c.id == cid)
4573 .map(|c| c.ty)
4574 .unwrap_or(TypeId::Bytes);
4575 (cid, ty)
4576 })
4577 .collect();
4578
4579 if fast_path {
4585 let needs_column = conditions.iter().any(|c| match c {
4588 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4589 Condition::RangeF64 { column_id, .. } => {
4590 !self.learned_range.contains_key(column_id)
4591 }
4592 _ => false,
4593 });
4594 let mut reader_opt: Option<RunReader> = if needs_column {
4595 Some(self.open_reader(self.run_refs[0].run_id)?)
4596 } else {
4597 None
4598 };
4599 let mut sets: Vec<RowIdSet> = Vec::new();
4600 for c in conditions {
4601 let s = match c {
4602 Condition::Range { column_id, lo, hi }
4603 if !self.learned_range.contains_key(column_id) =>
4604 {
4605 if reader_opt.is_none() {
4606 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4607 }
4608 reader_opt
4609 .as_mut()
4610 .expect("reader opened for range")
4611 .range_row_id_set_i64(*column_id, *lo, *hi)?
4612 }
4613 Condition::RangeF64 {
4614 column_id,
4615 lo,
4616 lo_inclusive,
4617 hi,
4618 hi_inclusive,
4619 } if !self.learned_range.contains_key(column_id) => {
4620 if reader_opt.is_none() {
4621 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4622 }
4623 reader_opt
4624 .as_mut()
4625 .expect("reader opened for range")
4626 .range_row_id_set_f64(
4627 *column_id,
4628 *lo,
4629 *lo_inclusive,
4630 *hi,
4631 *hi_inclusive,
4632 )?
4633 }
4634 _ => self.resolve_condition(c, snapshot)?,
4635 };
4636 sets.push(s);
4637 }
4638 let candidates = RowIdSet::intersect_many(sets);
4639 crate::trace::QueryTrace::record(|t| {
4640 t.survivor_count = Some(candidates.len());
4641 });
4642 if candidates.is_empty() {
4643 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4644 .iter()
4645 .map(|&id| {
4646 (
4647 id,
4648 columnar::null_native(
4649 proj_pairs
4650 .iter()
4651 .find(|(c, _)| c == &id)
4652 .map(|(_, t)| *t)
4653 .unwrap_or(TypeId::Bytes),
4654 0,
4655 ),
4656 )
4657 })
4658 .collect();
4659 return Ok(Some(cols));
4660 }
4661 let mut reader = match reader_opt.take() {
4662 Some(r) => r,
4663 None => self.open_reader(self.run_refs[0].run_id)?,
4664 };
4665 let candidate_ids = candidates.into_sorted_vec();
4666 let (positions, fast_rid) = if let Some(positions) =
4667 reader.positions_for_row_ids_fast(&candidate_ids)
4668 {
4669 (positions, true)
4670 } else {
4671 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4672 match col {
4673 columnar::NativeColumn::Int64 { data, .. } => {
4674 let mut p: Vec<usize> = candidate_ids
4675 .iter()
4676 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4677 .collect();
4678 p.sort_unstable();
4679 (p, false)
4680 }
4681 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4682 }
4683 };
4684 crate::trace::QueryTrace::record(|t| {
4685 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4686 t.fast_row_id_map = fast_rid;
4687 });
4688 let mut cols = Vec::with_capacity(col_ids.len());
4689 for cid in &col_ids {
4690 let col = reader.column_native(*cid)?;
4691 cols.push((*cid, col.gather(&positions)));
4692 }
4693 return Ok(Some(cols));
4694 }
4695
4696 if !self.run_refs.is_empty() {
4709 use crate::cursor::{drain_cursor_to_columns, Cursor};
4710 let remaining: usize;
4711 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4712 let c = self
4713 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4714 .expect("single-run cursor should build when run_refs.len() == 1");
4715 remaining = c.remaining_rows();
4716 Box::new(c)
4717 } else {
4718 let c = self
4719 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4720 .expect("multi-run cursor should build when run_refs.len() >= 1");
4721 remaining = c.remaining_rows();
4722 Box::new(c)
4723 };
4724 crate::trace::QueryTrace::record(|t| {
4725 if t.survivor_count.is_none() {
4726 t.survivor_count = Some(remaining);
4727 }
4728 });
4729 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4730 return Ok(Some(cols));
4731 }
4732
4733 crate::trace::QueryTrace::record(|t| {
4738 t.scan_mode = crate::trace::ScanMode::Materialized;
4739 t.row_materialized = true;
4740 });
4741 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4742 for c in conditions {
4743 sets.push(self.resolve_condition(c, snapshot)?);
4744 }
4745 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4746 let rows = self.rows_for_rids(&rids, snapshot)?;
4747 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4748 for (cid, ty) in &proj_pairs {
4749 let vals: Vec<Value> = rows
4750 .iter()
4751 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4752 .collect();
4753 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4754 }
4755 Ok(Some(cols))
4756 }
4757
4758 pub fn native_page_cursor(
4773 &self,
4774 snapshot: Snapshot,
4775 projection: Vec<(u16, TypeId)>,
4776 conditions: &[crate::query::Condition],
4777 ) -> Result<Option<NativePageCursor>> {
4778 use crate::cursor::build_page_plans;
4779 if !conditions.is_empty() && !self.indexes_complete {
4782 return Ok(None);
4783 }
4784 if self.run_refs.len() != 1 {
4785 return Ok(None);
4786 }
4787 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4788 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4789
4790 let overlay_rids: HashSet<u64> = {
4793 let mut s = HashSet::new();
4794 for row in self.memtable.visible_versions(snapshot.epoch) {
4795 s.insert(row.row_id.0);
4796 }
4797 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4798 s.insert(row.row_id.0);
4799 }
4800 s
4801 };
4802
4803 let survivors = if conditions.is_empty() {
4807 None
4808 } else {
4809 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4810 };
4811
4812 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4819 survivors.clone()
4820 } else if let Some(s) = &survivors {
4821 let mut run_set = s.clone();
4822 run_set.remove_many(overlay_rids.iter().copied());
4823 Some(run_set)
4824 } else {
4825 Some(RowIdSet::from_unsorted(
4826 rids.iter()
4827 .map(|&r| r as u64)
4828 .filter(|r| !overlay_rids.contains(r))
4829 .collect(),
4830 ))
4831 };
4832
4833 let overlay_rows = if overlay_rids.is_empty() {
4834 Vec::new()
4835 } else {
4836 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4837 self.overlay_visible_rows(snapshot, bound)
4838 };
4839
4840 let plans = if positions.is_empty() {
4842 Vec::new()
4843 } else {
4844 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4845 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4846 };
4847
4848 let overlay = if overlay_rows.is_empty() {
4850 None
4851 } else {
4852 let filtered =
4853 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4854 if filtered.is_empty() {
4855 None
4856 } else {
4857 Some(self.materialize_overlay(&filtered, &projection))
4858 }
4859 };
4860
4861 let overlay_row_count = overlay
4862 .as_ref()
4863 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4864 .unwrap_or(0);
4865 crate::trace::QueryTrace::record(|t| {
4866 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4867 t.run_count = self.run_refs.len();
4868 t.memtable_rows = self.memtable.len();
4869 t.mutable_run_rows = self.mutable_run.len();
4870 t.overlay_rows = overlay_row_count;
4871 t.conditions_pushed = conditions.len();
4872 t.pages_decoded = plans
4873 .iter()
4874 .map(|p| p.positions.len())
4875 .sum::<usize>()
4876 .min(1);
4877 });
4878
4879 Ok(Some(NativePageCursor::new_with_overlay(
4880 reader, projection, plans, overlay,
4881 )))
4882 }
4883 #[allow(clippy::type_complexity)]
4893 pub fn native_multi_run_cursor(
4894 &self,
4895 snapshot: Snapshot,
4896 projection: Vec<(u16, TypeId)>,
4897 conditions: &[crate::query::Condition],
4898 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4899 use crate::cursor::{MultiRunCursor, RunStream};
4900 use crate::sorted_run::SYS_ROW_ID;
4901 use std::collections::{BinaryHeap, HashMap, HashSet};
4902 if !conditions.is_empty() && !self.indexes_complete {
4905 return Ok(None);
4906 }
4907 if self.run_refs.is_empty() {
4908 return Ok(None);
4909 }
4910
4911 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
4913 Vec::with_capacity(self.run_refs.len());
4914 for rr in &self.run_refs {
4915 let mut reader = self.open_reader(rr.run_id)?;
4916 let (rids, eps, del) = reader.system_columns_native()?;
4917 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
4918 run_meta.push((reader, rids, eps, del, page_rows));
4919 }
4920
4921 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
4925 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
4926 for i in 0..rids.len() {
4927 let rid = rids[i] as u64;
4928 let e = eps[i] as u64;
4929 if e > snapshot.epoch.0 {
4930 continue;
4931 }
4932 let is_del = del[i] != 0;
4933 best.entry(rid)
4934 .and_modify(|cur| {
4935 if e > cur.0 {
4936 *cur = (e, run_idx, i, is_del);
4937 }
4938 })
4939 .or_insert((e, run_idx, i, is_del));
4940 }
4941 }
4942
4943 let overlay_rids: HashSet<u64> = {
4945 let mut s = HashSet::new();
4946 for row in self.memtable.visible_versions(snapshot.epoch) {
4947 s.insert(row.row_id.0);
4948 }
4949 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4950 s.insert(row.row_id.0);
4951 }
4952 s
4953 };
4954
4955 let survivors: Option<RowIdSet> = if conditions.is_empty() {
4957 None
4958 } else {
4959 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4960 for c in conditions {
4961 sets.push(self.resolve_condition(c, snapshot)?);
4962 }
4963 Some(RowIdSet::intersect_many(sets))
4964 };
4965
4966 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
4970 for (rid, (_, run_idx, pos, deleted)) in &best {
4971 if *deleted {
4972 continue;
4973 }
4974 if overlay_rids.contains(rid) {
4975 continue;
4976 }
4977 if let Some(s) = &survivors {
4978 if !s.contains(*rid) {
4979 continue;
4980 }
4981 }
4982 per_run[*run_idx].push((*rid, *pos));
4983 }
4984 for v in per_run.iter_mut() {
4985 v.sort_unstable_by_key(|&(rid, _)| rid);
4986 }
4987
4988 let mut streams = Vec::with_capacity(run_meta.len());
4990 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
4991 let mut total = 0usize;
4992 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
4993 let mut starts = Vec::with_capacity(page_rows.len());
4994 let mut acc = 0usize;
4995 for &r in &page_rows {
4996 starts.push(acc);
4997 acc += r;
4998 }
4999 let mut survivors_vec: Vec<(u64, usize, usize)> =
5000 Vec::with_capacity(per_run[run_idx].len());
5001 for &(rid, pos) in &per_run[run_idx] {
5002 let page_seq = match starts.partition_point(|&s| s <= pos) {
5003 0 => continue,
5004 p => p - 1,
5005 };
5006 let within = pos - starts[page_seq];
5007 survivors_vec.push((rid, page_seq, within));
5008 }
5009 total += survivors_vec.len();
5010 if let Some(&(rid, _, _)) = survivors_vec.first() {
5011 heap.push(std::cmp::Reverse((rid, run_idx)));
5012 }
5013 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5014 }
5015
5016 let overlay_rows = if overlay_rids.is_empty() {
5018 Vec::new()
5019 } else {
5020 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5021 self.overlay_visible_rows(snapshot, bound)
5022 };
5023 let overlay = if overlay_rows.is_empty() {
5024 None
5025 } else {
5026 let filtered =
5027 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5028 if filtered.is_empty() {
5029 None
5030 } else {
5031 Some(self.materialize_overlay(&filtered, &projection))
5032 }
5033 };
5034
5035 let overlay_row_count = overlay
5036 .as_ref()
5037 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5038 .unwrap_or(0);
5039 crate::trace::QueryTrace::record(|t| {
5040 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5041 t.run_count = self.run_refs.len();
5042 t.memtable_rows = self.memtable.len();
5043 t.mutable_run_rows = self.mutable_run.len();
5044 t.overlay_rows = overlay_row_count;
5045 t.conditions_pushed = conditions.len();
5046 t.survivor_count = Some(total);
5047 });
5048
5049 Ok(Some(MultiRunCursor::new(
5050 streams, projection, heap, total, overlay,
5051 )))
5052 }
5053
5054 fn overlay_materialization_bound<'a>(
5066 conditions: &[crate::query::Condition],
5067 survivors: &'a Option<RowIdSet>,
5068 ) -> Option<&'a RowIdSet> {
5069 use crate::query::Condition;
5070 let has_range = conditions
5071 .iter()
5072 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5073 if has_range {
5074 None
5075 } else {
5076 survivors.as_ref()
5077 }
5078 }
5079
5080 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5092 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5093 let mut fold = |row: Row| {
5094 if let Some(b) = bound {
5095 if !b.contains(row.row_id.0) {
5096 return;
5097 }
5098 }
5099 best.entry(row.row_id.0)
5100 .and_modify(|(be, br)| {
5101 if row.committed_epoch > *be {
5102 *be = row.committed_epoch;
5103 *br = row.clone();
5104 }
5105 })
5106 .or_insert_with(|| (row.committed_epoch, row));
5107 };
5108 for row in self.memtable.visible_versions(snapshot.epoch) {
5109 fold(row);
5110 }
5111 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5112 fold(row);
5113 }
5114 let mut out: Vec<Row> = best
5115 .into_values()
5116 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5117 .collect();
5118 out.sort_by_key(|r| r.row_id);
5119 out
5120 }
5121
5122 fn filter_overlay_rows(
5130 &self,
5131 rows: Vec<Row>,
5132 conditions: &[crate::query::Condition],
5133 survivors: Option<&RowIdSet>,
5134 snapshot: Snapshot,
5135 ) -> Result<Vec<Row>> {
5136 if conditions.is_empty() {
5137 return Ok(rows);
5138 }
5139 use crate::query::Condition;
5140 let all_index_served = !conditions
5144 .iter()
5145 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5146 if all_index_served {
5147 return Ok(rows
5148 .into_iter()
5149 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5150 .collect());
5151 }
5152 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5155 for c in conditions {
5156 let s = match c {
5157 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5158 _ => self.resolve_condition(c, snapshot)?,
5159 };
5160 per_cond_sets.push(s);
5161 }
5162 Ok(rows
5163 .into_iter()
5164 .filter(|row| {
5165 conditions.iter().enumerate().all(|(i, c)| match c {
5166 Condition::Range { column_id, lo, hi } => {
5167 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5168 }
5169 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5170 match row.columns.get(column_id) {
5171 Some(Value::Float64(v)) => {
5172 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5173 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5174 lo_ok && hi_ok
5175 }
5176 _ => false,
5177 }
5178 }
5179 _ => per_cond_sets[i].contains(row.row_id.0),
5180 })
5181 })
5182 .collect())
5183 }
5184
5185 fn materialize_overlay(
5188 &self,
5189 rows: &[Row],
5190 projection: &[(u16, TypeId)],
5191 ) -> Vec<columnar::NativeColumn> {
5192 if projection.is_empty() {
5193 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5194 }
5195 let mut cols = Vec::with_capacity(projection.len());
5196 for (cid, ty) in projection {
5197 let vals: Vec<Value> = rows
5198 .iter()
5199 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5200 .collect();
5201 cols.push(columnar::values_to_native(*ty, &vals));
5202 }
5203 cols
5204 }
5205
5206 fn resolve_survivor_rids(
5211 &self,
5212 conditions: &[crate::query::Condition],
5213 reader: &mut RunReader,
5214 ) -> Result<RowIdSet> {
5215 use crate::query::Condition;
5216 let mut sets: Vec<RowIdSet> = Vec::new();
5217 for c in conditions {
5218 let s: RowIdSet = match c {
5219 Condition::Pk(key) => {
5220 let lookup = self
5221 .schema
5222 .primary_key()
5223 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5224 .unwrap_or_else(|| key.clone());
5225 self.hot
5226 .get(&lookup)
5227 .map(|r| RowIdSet::one(r.0))
5228 .unwrap_or_else(RowIdSet::empty)
5229 }
5230 Condition::BitmapEq { column_id, value } => {
5231 let lookup = self.index_lookup_key_bytes(*column_id, value);
5232 self.bitmap
5233 .get(column_id)
5234 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5235 .unwrap_or_else(RowIdSet::empty)
5236 }
5237 Condition::BitmapIn { column_id, values } => {
5238 let bm = self.bitmap.get(column_id);
5239 let mut acc = roaring::RoaringBitmap::new();
5240 if let Some(b) = bm {
5241 for v in values {
5242 let lookup = self.index_lookup_key_bytes(*column_id, v);
5243 acc |= b.get(&lookup);
5244 }
5245 }
5246 RowIdSet::from_roaring(acc)
5247 }
5248 Condition::BytesPrefix { column_id, prefix } => {
5249 if let Some(b) = self.bitmap.get(column_id) {
5250 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5251 let mut acc = roaring::RoaringBitmap::new();
5252 for key in b.keys() {
5253 if key.starts_with(&lookup_prefix) {
5254 acc |= b.get(key);
5255 }
5256 }
5257 RowIdSet::from_roaring(acc)
5258 } else {
5259 RowIdSet::empty()
5260 }
5261 }
5262 Condition::FmContains { column_id, pattern } => self
5263 .fm
5264 .get(column_id)
5265 .map(|f| {
5266 RowIdSet::from_unsorted(
5267 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5268 )
5269 })
5270 .unwrap_or_else(RowIdSet::empty),
5271 Condition::FmContainsAll {
5272 column_id,
5273 patterns,
5274 } => {
5275 if let Some(f) = self.fm.get(column_id) {
5276 let sets: Vec<RowIdSet> = patterns
5277 .iter()
5278 .map(|pat| {
5279 RowIdSet::from_unsorted(
5280 f.locate(pat).into_iter().map(|r| r.0).collect(),
5281 )
5282 })
5283 .collect();
5284 RowIdSet::intersect_many(sets)
5285 } else {
5286 RowIdSet::empty()
5287 }
5288 }
5289 Condition::Ann {
5290 column_id,
5291 query,
5292 k,
5293 } => self
5294 .ann
5295 .get(column_id)
5296 .map(|a| {
5297 RowIdSet::from_unsorted(
5298 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5299 )
5300 })
5301 .unwrap_or_else(RowIdSet::empty),
5302 Condition::SparseMatch {
5303 column_id,
5304 query,
5305 k,
5306 } => self
5307 .sparse
5308 .get(column_id)
5309 .map(|s| {
5310 RowIdSet::from_unsorted(
5311 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5312 )
5313 })
5314 .unwrap_or_else(RowIdSet::empty),
5315 Condition::MinHashSimilar {
5316 column_id,
5317 query,
5318 k,
5319 } => self
5320 .minhash
5321 .get(column_id)
5322 .map(|mh| {
5323 RowIdSet::from_unsorted(
5324 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5325 )
5326 })
5327 .unwrap_or_else(RowIdSet::empty),
5328 Condition::Range { column_id, lo, hi } => {
5329 if let Some(li) = self.learned_range.get(column_id) {
5330 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5331 } else {
5332 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5333 }
5334 }
5335 Condition::RangeF64 {
5336 column_id,
5337 lo,
5338 lo_inclusive,
5339 hi,
5340 hi_inclusive,
5341 } => {
5342 if let Some(li) = self.learned_range.get(column_id) {
5343 RowIdSet::from_unsorted(
5344 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5345 .into_iter()
5346 .collect(),
5347 )
5348 } else {
5349 reader.range_row_id_set_f64(
5350 *column_id,
5351 *lo,
5352 *lo_inclusive,
5353 *hi,
5354 *hi_inclusive,
5355 )?
5356 }
5357 }
5358 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5359 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5360 };
5361 sets.push(s);
5362 }
5363 Ok(RowIdSet::intersect_many(sets))
5364 }
5365
5366 pub fn scan_cursor(
5387 &self,
5388 snapshot: Snapshot,
5389 projection: Vec<(u16, TypeId)>,
5390 conditions: &[crate::query::Condition],
5391 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5392 if !conditions.is_empty() && !self.indexes_complete {
5398 return Ok(None);
5399 }
5400 if self.run_refs.len() == 1 {
5401 Ok(self
5402 .native_page_cursor(snapshot, projection, conditions)?
5403 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5404 } else {
5405 Ok(self
5406 .native_multi_run_cursor(snapshot, projection, conditions)?
5407 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5408 }
5409 }
5410
5411 pub fn aggregate_native(
5425 &self,
5426 snapshot: Snapshot,
5427 column: Option<u16>,
5428 conditions: &[crate::query::Condition],
5429 agg: NativeAgg,
5430 ) -> Result<Option<NativeAggResult>> {
5431 if self.run_refs.len() == 1 && conditions.is_empty() {
5433 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5434 return Ok(Some(res));
5435 }
5436 }
5437 if matches!(agg, NativeAgg::Count) && column.is_none() {
5439 return Ok(self
5440 .scan_cursor(snapshot, Vec::new(), conditions)?
5441 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5442 }
5443 let cid = match column {
5446 Some(c) => c,
5447 None => return Ok(None),
5448 };
5449 let ty = self.column_type(cid);
5450 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5451 return Ok(None);
5452 };
5453 match ty {
5454 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5455 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5456 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5457 }
5458 TypeId::Float64 => {
5459 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5460 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5461 }
5462 _ => Ok(None),
5463 }
5464 }
5465
5466 fn aggregate_from_stats(
5474 &self,
5475 snapshot: Snapshot,
5476 column: Option<u16>,
5477 agg: NativeAgg,
5478 ) -> Result<Option<NativeAggResult>> {
5479 let cid = match (agg, column) {
5480 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5481 _ => return Ok(None), };
5483 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5484 return Ok(None);
5485 };
5486 let Some(cs) = stats.get(&cid) else {
5487 return Ok(None);
5488 };
5489 match agg {
5490 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5492 self.live_count.saturating_sub(cs.null_count),
5493 ))),
5494 NativeAgg::Min | NativeAgg::Max => {
5495 let bound = if agg == NativeAgg::Min {
5496 &cs.min
5497 } else {
5498 &cs.max
5499 };
5500 match bound {
5501 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5502 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5503 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5508 None => Ok(None),
5509 }
5510 }
5511 _ => Ok(None),
5512 }
5513 }
5514
5515 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5524 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5525 return Ok(None);
5526 }
5527 self.ensure_indexes_complete()?;
5530 let reader = self.open_reader(self.run_refs[0].run_id)?;
5531 if self.live_count != reader.row_count() as u64 {
5532 return Ok(None);
5533 }
5534 let Some(bm) = self.bitmap.get(&column_id) else {
5535 return Ok(None); };
5537 let mut distinct = bm.value_count() as u64;
5538 if !bm.get(&Value::Null.encode_key()).is_empty() {
5541 distinct = distinct.saturating_sub(1);
5542 }
5543 Ok(Some(distinct))
5544 }
5545
5546 pub fn aggregate_incremental(
5558 &mut self,
5559 cache_key: u64,
5560 conditions: &[crate::query::Condition],
5561 column: Option<u16>,
5562 agg: NativeAgg,
5563 ) -> Result<IncrementalAggResult> {
5564 let snap = self.snapshot();
5565 let cur_wm = self.allocator.current().0;
5566 let cur_epoch = snap.epoch.0;
5567 let incremental_ok =
5574 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5575
5576 if incremental_ok {
5579 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5580 if cached.epoch == cur_epoch {
5581 return Ok(IncrementalAggResult {
5582 state: cached.state,
5583 incremental: true,
5584 delta_rows: 0,
5585 });
5586 }
5587 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5588 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5589 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5590 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5591 let delta_state = agg_state_from_rows(
5592 &delta_rows,
5593 conditions,
5594 &index_sets,
5595 column,
5596 agg,
5597 &self.schema,
5598 )?;
5599 let merged = cached.state.merge(delta_state);
5600 let delta_n = delta_rids.len() as u64;
5601 self.agg_cache.insert(
5602 cache_key,
5603 CachedAgg {
5604 state: merged.clone(),
5605 watermark: cur_wm,
5606 epoch: cur_epoch,
5607 },
5608 );
5609 return Ok(IncrementalAggResult {
5610 state: merged,
5611 incremental: true,
5612 delta_rows: delta_n,
5613 });
5614 }
5615 }
5616 }
5617
5618 let cursor_ok =
5623 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5624 let state = if cursor_ok && agg != NativeAgg::Avg {
5625 match self.aggregate_native(snap, column, conditions, agg)? {
5626 Some(result) => {
5627 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5628 }
5629 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5630 }
5631 } else {
5632 self.agg_state_full_scan(conditions, column, agg, snap)?
5633 };
5634 if incremental_ok {
5636 self.agg_cache.insert(
5637 cache_key,
5638 CachedAgg {
5639 state: state.clone(),
5640 watermark: cur_wm,
5641 epoch: cur_epoch,
5642 },
5643 );
5644 }
5645 Ok(IncrementalAggResult {
5646 state,
5647 incremental: false,
5648 delta_rows: 0,
5649 })
5650 }
5651
5652 fn agg_state_full_scan(
5655 &self,
5656 conditions: &[crate::query::Condition],
5657 column: Option<u16>,
5658 agg: NativeAgg,
5659 snap: Snapshot,
5660 ) -> Result<AggState> {
5661 let rows = self.visible_rows(snap)?;
5662 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5663 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5664 }
5665
5666 fn resolve_index_conditions(
5669 &self,
5670 conditions: &[crate::query::Condition],
5671 snapshot: Snapshot,
5672 ) -> Result<Vec<RowIdSet>> {
5673 use crate::query::Condition;
5674 let mut sets = Vec::new();
5675 for c in conditions {
5676 if matches!(
5677 c,
5678 Condition::Ann { .. }
5679 | Condition::SparseMatch { .. }
5680 | Condition::MinHashSimilar { .. }
5681 ) {
5682 sets.push(self.resolve_condition(c, snapshot)?);
5683 }
5684 }
5685 Ok(sets)
5686 }
5687
5688 fn column_type(&self, cid: u16) -> TypeId {
5689 self.schema
5690 .columns
5691 .iter()
5692 .find(|c| c.id == cid)
5693 .map(|c| c.ty)
5694 .unwrap_or(TypeId::Bytes)
5695 }
5696
5697 pub fn approx_aggregate(
5706 &mut self,
5707 conditions: &[crate::query::Condition],
5708 column: Option<u16>,
5709 agg: ApproxAgg,
5710 z: f64,
5711 ) -> Result<Option<ApproxResult>> {
5712 use crate::query::Condition;
5713 self.ensure_reservoir_complete()?;
5714 let snapshot = self.snapshot();
5715 let n_pop = self.live_count;
5716 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5717 if sample_rids.is_empty() {
5718 return Ok(None);
5719 }
5720 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5722 let s = live_sample.len();
5723 if s == 0 {
5724 return Ok(None);
5725 }
5726
5727 let mut index_sets: Vec<RowIdSet> = Vec::new();
5730 for c in conditions {
5731 if matches!(
5732 c,
5733 Condition::Ann { .. }
5734 | Condition::SparseMatch { .. }
5735 | Condition::MinHashSimilar { .. }
5736 ) {
5737 index_sets.push(self.resolve_condition(c, snapshot)?);
5738 }
5739 }
5740
5741 let cid = match (agg, column) {
5743 (ApproxAgg::Count, _) => None,
5744 (_, Some(c)) => Some(c),
5745 _ => return Ok(None),
5746 };
5747 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5748 for r in &live_sample {
5749 if !conditions
5751 .iter()
5752 .all(|c| condition_matches_row(c, r, &self.schema))
5753 {
5754 continue;
5755 }
5756 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5758 continue;
5759 }
5760 if let Some(cid) = cid {
5761 if let Some(v) = as_f64(r.columns.get(&cid)) {
5762 passing_vals.push(v);
5763 } } else {
5765 passing_vals.push(0.0); }
5767 }
5768 let m = passing_vals.len();
5769
5770 let (point, half) = match agg {
5771 ApproxAgg::Count => {
5772 let p = m as f64 / s as f64;
5774 let point = n_pop as f64 * p;
5775 let var = if s > 1 {
5776 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5777 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5778 } else {
5779 0.0
5780 };
5781 (point, z * var.sqrt())
5782 }
5783 ApproxAgg::Sum => {
5784 let y: Vec<f64> = live_sample
5786 .iter()
5787 .map(|r| {
5788 let passes_row = conditions
5789 .iter()
5790 .all(|c| condition_matches_row(c, r, &self.schema))
5791 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5792 if passes_row {
5793 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5794 } else {
5795 0.0
5796 }
5797 })
5798 .collect();
5799 let mean_y = y.iter().sum::<f64>() / s as f64;
5800 let point = n_pop as f64 * mean_y;
5801 let var = if s > 1 {
5802 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5803 let var_y = ss / (s - 1) as f64;
5804 n_pop as f64 * n_pop as f64 * var_y / s as f64
5805 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5806 } else {
5807 0.0
5808 };
5809 (point, z * var.sqrt())
5810 }
5811 ApproxAgg::Avg => {
5812 if m == 0 {
5813 return Ok(Some(ApproxResult {
5814 point: 0.0,
5815 ci_low: 0.0,
5816 ci_high: 0.0,
5817 n_population: n_pop,
5818 n_sample_live: s,
5819 n_passing: 0,
5820 }));
5821 }
5822 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5823 let half = if m > 1 {
5824 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5825 let sd = (ss / (m - 1) as f64).sqrt();
5826 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5827 z * sd / (m as f64).sqrt() * fpc.sqrt()
5828 } else {
5829 0.0
5830 };
5831 (mean, half)
5832 }
5833 };
5834
5835 Ok(Some(ApproxResult {
5836 point,
5837 ci_low: point - half,
5838 ci_high: point + half,
5839 n_population: n_pop,
5840 n_sample_live: s,
5841 n_passing: m,
5842 }))
5843 }
5844
5845 pub fn exact_column_stats(
5853 &self,
5854 _snapshot: Snapshot,
5855 projection: &[u16],
5856 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5857 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5858 return Ok(None);
5859 }
5860 let reader = self.open_reader(self.run_refs[0].run_id)?;
5861 if self.live_count != reader.row_count() as u64 {
5862 return Ok(None);
5863 }
5864 let mut out = HashMap::new();
5865 for &cid in projection {
5866 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5867 Some(c) => c,
5868 None => continue,
5869 };
5870 let Some(stats) = reader.column_page_stats(cid) else {
5872 out.insert(
5873 cid,
5874 ColumnStat {
5875 min: None,
5876 max: None,
5877 null_count: self.live_count,
5878 },
5879 );
5880 continue;
5881 };
5882 let stat = match cdef.ty {
5883 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5884 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5885 min: mn.map(Value::Int64),
5886 max: mx.map(Value::Int64),
5887 null_count: n,
5888 })
5889 }
5890 TypeId::Float64 => {
5891 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5892 min: mn.map(Value::Float64),
5893 max: mx.map(Value::Float64),
5894 null_count: n,
5895 })
5896 }
5897 _ => None,
5898 };
5899 if let Some(s) = stat {
5900 out.insert(cid, s);
5901 }
5902 }
5903 Ok(Some(out))
5904 }
5905
5906 pub fn dir(&self) -> &Path {
5907 &self.dir
5908 }
5909
5910 pub fn schema(&self) -> &Schema {
5911 &self.schema
5912 }
5913
5914 pub(crate) fn prepare_alter_column(
5915 &mut self,
5916 column_name: &str,
5917 change: &AlterColumn,
5918 ) -> Result<ColumnDef> {
5919 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
5920 return Err(MongrelError::InvalidArgument(
5921 "ALTER COLUMN requires committing staged writes first".into(),
5922 ));
5923 }
5924 let old = self
5925 .schema
5926 .columns
5927 .iter()
5928 .find(|c| c.name == column_name)
5929 .cloned()
5930 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5931 let mut next = old.clone();
5932
5933 if let Some(name) = &change.name {
5934 let trimmed = name.trim();
5935 if trimmed.is_empty() {
5936 return Err(MongrelError::InvalidArgument(
5937 "ALTER COLUMN name must not be empty".into(),
5938 ));
5939 }
5940 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
5941 return Err(MongrelError::Schema(format!(
5942 "column {trimmed} already exists"
5943 )));
5944 }
5945 next.name = trimmed.to_string();
5946 }
5947
5948 if let Some(ty) = change.ty {
5949 next.ty = ty;
5950 }
5951 if let Some(flags) = change.flags {
5952 validate_alter_column_flags(old.flags, flags)?;
5953 next.flags = flags;
5954 }
5955
5956 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
5957 if old.flags.contains(ColumnFlags::NULLABLE)
5958 && !next.flags.contains(ColumnFlags::NULLABLE)
5959 && self.column_has_nulls(old.id)?
5960 {
5961 return Err(MongrelError::InvalidArgument(format!(
5962 "column '{}' contains NULL values",
5963 old.name
5964 )));
5965 }
5966 Ok(next)
5967 }
5968
5969 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
5970 let idx = self
5971 .schema
5972 .columns
5973 .iter()
5974 .position(|c| c.id == column.id)
5975 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
5976 if self.schema.columns[idx] == column {
5977 return Ok(());
5978 }
5979 self.schema.columns[idx] = column;
5980 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5981 self.schema.validate_auto_increment()?;
5982 self.auto_inc = resolve_auto_inc(&self.schema);
5983 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
5984 write_schema(&self.dir, &self.schema)?;
5985 self.clear_result_cache();
5986 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5987 self.persist_manifest(self.current_epoch())?;
5988 Ok(())
5989 }
5990
5991 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
5992 let column = self.prepare_alter_column(column_name, &change)?;
5993 self.apply_altered_column(column.clone())?;
5994 Ok(column)
5995 }
5996
5997 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
5998 if self.live_count == 0 {
5999 return Ok(false);
6000 }
6001 let snap = self.snapshot();
6002 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6003 Ok(columns
6004 .first()
6005 .map(|(_, col)| col.null_count(col.len()) != 0)
6006 .unwrap_or(true))
6007 }
6008
6009 fn has_stored_versions(&self) -> bool {
6010 !self.memtable.is_empty()
6011 || !self.mutable_run.is_empty()
6012 || self.run_refs.iter().any(|r| r.row_count > 0)
6013 || !self.retiring.is_empty()
6014 }
6015
6016 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
6021 if self.schema.columns.iter().any(|c| c.name == name) {
6022 return Err(MongrelError::Schema(format!(
6023 "column {name} already exists"
6024 )));
6025 }
6026 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
6027 self.schema.columns.push(ColumnDef {
6028 id,
6029 name: name.to_string(),
6030 ty,
6031 flags,
6032 });
6033 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6034 self.schema.validate_auto_increment()?;
6035 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6036 self.auto_inc = resolve_auto_inc(&self.schema);
6037 }
6038 write_schema(&self.dir, &self.schema)?;
6039 self.clear_result_cache();
6040 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6042 self.persist_manifest(self.current_epoch())?;
6043 Ok(id)
6044 }
6045
6046 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6055 let cid = self
6056 .schema
6057 .columns
6058 .iter()
6059 .find(|c| c.name == column_name)
6060 .map(|c| c.id)
6061 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6062 let ty = self
6063 .schema
6064 .columns
6065 .iter()
6066 .find(|c| c.id == cid)
6067 .map(|c| c.ty)
6068 .unwrap_or(TypeId::Int64);
6069 if !matches!(
6070 ty,
6071 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6072 ) {
6073 return Err(MongrelError::Schema(format!(
6074 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6075 )));
6076 }
6077 if self
6078 .schema
6079 .indexes
6080 .iter()
6081 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6082 {
6083 return Ok(()); }
6085 self.schema.indexes.push(IndexDef {
6086 name: format!("{}_learned_range", column_name),
6087 column_id: cid,
6088 kind: IndexKind::LearnedRange,
6089 });
6090 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6091 write_schema(&self.dir, &self.schema)?;
6092 self.build_learned_ranges()?;
6093 Ok(())
6094 }
6095
6096 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6099 self.sync_byte_threshold = threshold;
6100 if let WalSink::Private(w) = &mut self.wal {
6101 w.set_sync_byte_threshold(threshold);
6102 }
6103 }
6104
6105 pub fn page_cache_flush(&self) {
6109 self.page_cache.flush_to_disk();
6110 }
6111
6112 pub fn page_cache_len(&self) -> usize {
6114 self.page_cache.len()
6115 }
6116
6117 pub fn decoded_cache_len(&self) -> usize {
6120 self.decoded_cache.len()
6121 }
6122
6123 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6126 self.memtable.drain_sorted()
6127 }
6128
6129 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6130 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6131 }
6132
6133 pub(crate) fn table_dir(&self) -> &Path {
6134 &self.dir
6135 }
6136
6137 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6138 &self.schema
6139 }
6140
6141 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6142 let id = self.next_run_id;
6143 self.next_run_id += 1;
6144 id
6145 }
6146
6147 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6148 self.run_refs.push(run_ref);
6149 }
6150
6151 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6161 self.retiring.push(crate::manifest::RetiredRun {
6162 run_id,
6163 retire_epoch,
6164 });
6165 }
6166
6167 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6171 if self.retiring.is_empty() {
6172 return Ok(0);
6173 }
6174 let mut reaped = 0;
6175 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6176 for r in std::mem::take(&mut self.retiring) {
6182 if min_active.0 >= r.retire_epoch {
6183 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6184 reaped += 1;
6185 } else {
6186 kept.push(r);
6187 }
6188 }
6189 self.retiring = kept;
6190 if reaped > 0 {
6191 self.persist_manifest(self.current_epoch())?;
6192 }
6193 Ok(reaped)
6194 }
6195
6196 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6197 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6198 return false;
6199 }
6200 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6201 self.run_refs.push(run_ref);
6202 self.indexes_complete = false;
6203 true
6204 }
6205
6206 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6207 self.kek.as_ref()
6208 }
6209
6210 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6211 let mut reader = RunReader::open_with_cache(
6212 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6213 self.schema.clone(),
6214 self.kek.clone(),
6215 Some(self.page_cache.clone()),
6216 Some(self.decoded_cache.clone()),
6217 self.table_id,
6218 Some(&self.verified_runs),
6219 )?;
6220 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6224 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6225 }
6226 Ok(reader)
6227 }
6228
6229 pub(crate) fn run_refs(&self) -> &[RunRef] {
6230 &self.run_refs
6231 }
6232
6233 pub(crate) fn runs_dir(&self) -> PathBuf {
6234 self.dir.join(RUNS_DIR)
6235 }
6236
6237 pub(crate) fn wal_dir(&self) -> PathBuf {
6238 self.dir.join(WAL_DIR)
6239 }
6240
6241 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6242 self.run_refs = refs;
6243 }
6244
6245 pub(crate) fn next_run_id(&self) -> u64 {
6246 self.next_run_id
6247 }
6248
6249 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6250 self.compaction_zstd_level
6251 }
6252
6253 pub(crate) fn bump_next_run_id(&mut self) {
6254 self.next_run_id += 1;
6255 }
6256
6257 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6258 self.kek.clone()
6259 }
6260
6261 #[cfg(feature = "encryption")]
6265 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6266 self.kek.as_ref().map(|k| k.derive_idx_key())
6267 }
6268
6269 #[cfg(not(feature = "encryption"))]
6270 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6271 None
6272 }
6273
6274 #[cfg(feature = "encryption")]
6278 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6279 self.kek.as_ref().map(|k| *k.derive_meta_key())
6280 }
6281
6282 #[cfg(not(feature = "encryption"))]
6283 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6284 None
6285 }
6286
6287 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6290 self.column_keys
6291 .iter()
6292 .map(|(&id, &(_, scheme))| (id, scheme))
6293 .collect()
6294 }
6295
6296 #[cfg(feature = "encryption")]
6301 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6302 self.tokenize_value_enc(column_id, v)
6303 }
6304
6305 #[cfg(feature = "encryption")]
6306 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6307 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6308 let (key, scheme) = self.column_keys.get(&column_id)?;
6309 let token: Vec<u8> = match (*scheme, v) {
6310 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6311 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6312 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6313 _ => hmac_token(key, &v.encode_key()).to_vec(),
6314 };
6315 Some(Value::Bytes(token))
6316 }
6317
6318 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6320 self.index_lookup_key_bytes(column_id, &v.encode_key())
6321 }
6322
6323 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6326 #[cfg(feature = "encryption")]
6327 {
6328 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6329 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6330 if *scheme == SCHEME_HMAC_EQ {
6331 return hmac_token(key, encoded).to_vec();
6332 }
6333 }
6334 }
6335 let _ = column_id;
6336 encoded.to_vec()
6337 }
6338}
6339
6340fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6341 let columnar::NativeColumn::Int64 { data, validity } = col else {
6342 return false;
6343 };
6344 if data.len() < n || !columnar::all_non_null(validity, n) {
6345 return false;
6346 }
6347 data.iter()
6348 .take(n)
6349 .zip(data.iter().skip(1))
6350 .all(|(a, b)| a < b)
6351}
6352
6353#[derive(Debug, Clone)]
6357pub struct ColumnStat {
6358 pub min: Option<Value>,
6359 pub max: Option<Value>,
6360 pub null_count: u64,
6361}
6362
6363#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6365pub enum NativeAgg {
6366 Count,
6367 Sum,
6368 Min,
6369 Max,
6370 Avg,
6371}
6372
6373#[derive(Debug, Clone, PartialEq)]
6375pub enum NativeAggResult {
6376 Count(u64),
6377 Int(i64),
6378 Float(f64),
6379 Null,
6381}
6382
6383#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6385pub enum ApproxAgg {
6386 Count,
6387 Sum,
6388 Avg,
6389}
6390
6391#[derive(Debug, Clone)]
6395pub struct ApproxResult {
6396 pub point: f64,
6398 pub ci_low: f64,
6400 pub ci_high: f64,
6402 pub n_population: u64,
6404 pub n_sample_live: usize,
6406 pub n_passing: usize,
6408}
6409
6410#[derive(Debug, Clone, PartialEq)]
6415pub enum AggState {
6416 Count(u64),
6418 SumI {
6420 sum: i128,
6421 count: u64,
6422 },
6423 SumF {
6425 sum: f64,
6426 count: u64,
6427 },
6428 AvgI {
6430 sum: i128,
6431 count: u64,
6432 },
6433 AvgF {
6435 sum: f64,
6436 count: u64,
6437 },
6438 MinI(i64),
6440 MaxI(i64),
6441 MinF(f64),
6443 MaxF(f64),
6444 Empty,
6446}
6447
6448impl AggState {
6449 pub fn merge(self, other: AggState) -> AggState {
6451 use AggState::*;
6452 match (self, other) {
6453 (Empty, x) | (x, Empty) => x,
6454 (Count(a), Count(b)) => Count(a + b),
6455 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6456 sum: sa + sb,
6457 count: ca + cb,
6458 },
6459 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6460 sum: sa + sb,
6461 count: ca + cb,
6462 },
6463 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6464 sum: sa + sb,
6465 count: ca + cb,
6466 },
6467 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6468 sum: sa + sb,
6469 count: ca + cb,
6470 },
6471 (MinI(a), MinI(b)) => MinI(a.min(b)),
6472 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6473 (MinF(a), MinF(b)) => MinF(a.min(b)),
6474 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6475 _ => Empty, }
6477 }
6478
6479 pub fn point(&self) -> Option<f64> {
6481 match self {
6482 AggState::Count(n) => Some(*n as f64),
6483 AggState::SumI { sum, .. } => Some(*sum as f64),
6484 AggState::SumF { sum, .. } => Some(*sum),
6485 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6486 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6487 AggState::MinI(n) => Some(*n as f64),
6488 AggState::MaxI(n) => Some(*n as f64),
6489 AggState::MinF(n) => Some(*n),
6490 AggState::MaxF(n) => Some(*n),
6491 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6492 }
6493 }
6494
6495 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6499 let is_float = matches!(ty, Some(TypeId::Float64));
6500 match (agg, result) {
6501 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6502 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6503 sum: x as i128,
6504 count: 1, },
6506 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6507 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6508 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6509 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6510 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6511 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6512 (NativeAgg::Count, _) => AggState::Empty,
6513 (_, NativeAggResult::Null) => AggState::Empty,
6514 _ => {
6515 let _ = is_float;
6516 AggState::Empty
6517 }
6518 }
6519 }
6520}
6521
6522#[derive(Debug, Clone)]
6525pub struct CachedAgg {
6526 pub state: AggState,
6527 pub watermark: u64,
6528 pub epoch: u64,
6529}
6530
6531#[derive(Debug, Clone)]
6533pub struct IncrementalAggResult {
6534 pub state: AggState,
6536 pub incremental: bool,
6539 pub delta_rows: u64,
6541}
6542
6543fn agg_state_from_rows(
6547 rows: &[Row],
6548 conditions: &[crate::query::Condition],
6549 index_sets: &[RowIdSet],
6550 column: Option<u16>,
6551 agg: NativeAgg,
6552 schema: &Schema,
6553) -> Result<AggState> {
6554 let mut count: u64 = 0;
6555 let mut sum_i: i128 = 0;
6556 let mut sum_f: f64 = 0.0;
6557 let mut mn_i: i64 = i64::MAX;
6558 let mut mx_i: i64 = i64::MIN;
6559 let mut mn_f: f64 = f64::INFINITY;
6560 let mut mx_f: f64 = f64::NEG_INFINITY;
6561 let mut saw_int = false;
6562 let mut saw_float = false;
6563 for r in rows {
6564 if !conditions
6565 .iter()
6566 .all(|c| condition_matches_row(c, r, schema))
6567 {
6568 continue;
6569 }
6570 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6571 continue;
6572 }
6573 match agg {
6574 NativeAgg::Count => match column {
6575 None => count += 1,
6577 Some(cid) => match r.columns.get(&cid) {
6580 None | Some(Value::Null) => {}
6581 Some(_) => count += 1,
6582 },
6583 },
6584 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6585 Some(Value::Int64(n)) => {
6586 count += 1;
6587 sum_i += *n as i128;
6588 mn_i = mn_i.min(*n);
6589 mx_i = mx_i.max(*n);
6590 saw_int = true;
6591 }
6592 Some(Value::Float64(f)) => {
6593 count += 1;
6594 sum_f += f;
6595 mn_f = mn_f.min(*f);
6596 mx_f = mx_f.max(*f);
6597 saw_float = true;
6598 }
6599 _ => {}
6600 },
6601 }
6602 }
6603 Ok(match agg {
6604 NativeAgg::Count => {
6605 if count == 0 {
6606 AggState::Empty
6607 } else {
6608 AggState::Count(count)
6609 }
6610 }
6611 NativeAgg::Sum => {
6612 if count == 0 {
6613 AggState::Empty
6614 } else if saw_int {
6615 AggState::SumI { sum: sum_i, count }
6616 } else {
6617 AggState::SumF { sum: sum_f, count }
6618 }
6619 }
6620 NativeAgg::Avg => {
6621 if count == 0 {
6622 AggState::Empty
6623 } else if saw_int {
6624 AggState::AvgI { sum: sum_i, count }
6625 } else {
6626 AggState::AvgF { sum: sum_f, count }
6627 }
6628 }
6629 NativeAgg::Min => {
6630 if !saw_int && !saw_float {
6631 AggState::Empty
6632 } else if saw_int {
6633 AggState::MinI(mn_i)
6634 } else {
6635 AggState::MinF(mn_f)
6636 }
6637 }
6638 NativeAgg::Max => {
6639 if !saw_int && !saw_float {
6640 AggState::Empty
6641 } else if saw_int {
6642 AggState::MaxI(mx_i)
6643 } else {
6644 AggState::MaxF(mx_f)
6645 }
6646 }
6647 })
6648}
6649
6650fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6654 use crate::query::Condition;
6655 match c {
6656 Condition::Pk(key) => match schema.primary_key() {
6657 Some(pk) => row
6658 .columns
6659 .get(&pk.id)
6660 .map(|v| v.encode_key() == *key)
6661 .unwrap_or(false),
6662 None => false,
6663 },
6664 Condition::BitmapEq { column_id, value } => row
6665 .columns
6666 .get(column_id)
6667 .map(|v| v.encode_key() == *value)
6668 .unwrap_or(false),
6669 Condition::BitmapIn { column_id, values } => {
6670 let key = row.columns.get(column_id).map(|v| v.encode_key());
6671 match key {
6672 Some(k) => values.contains(&k),
6673 None => false,
6674 }
6675 }
6676 Condition::BytesPrefix { column_id, prefix } => row
6677 .columns
6678 .get(column_id)
6679 .map(|v| v.encode_key().starts_with(prefix))
6680 .unwrap_or(false),
6681 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6682 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6683 _ => false,
6684 },
6685 Condition::RangeF64 {
6686 column_id,
6687 lo,
6688 lo_inclusive,
6689 hi,
6690 hi_inclusive,
6691 } => match row.columns.get(column_id) {
6692 Some(Value::Float64(n)) => {
6693 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6694 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6695 lo_ok && hi_ok
6696 }
6697 _ => false,
6698 },
6699 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6700 Some(Value::Bytes(b)) => {
6701 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6702 }
6703 _ => false,
6704 },
6705 Condition::FmContainsAll {
6706 column_id,
6707 patterns,
6708 } => match row.columns.get(column_id) {
6709 Some(Value::Bytes(b)) => patterns
6710 .iter()
6711 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6712 _ => false,
6713 },
6714 Condition::Ann { .. }
6715 | Condition::SparseMatch { .. }
6716 | Condition::MinHashSimilar { .. } => true,
6717 Condition::IsNull { column_id } => {
6718 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6719 }
6720 Condition::IsNotNull { column_id } => {
6721 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6722 }
6723 }
6724}
6725
6726fn as_f64(v: Option<&Value>) -> Option<f64> {
6728 match v {
6729 Some(Value::Int64(n)) => Some(*n as f64),
6730 Some(Value::Float64(f)) => Some(*f),
6731 _ => None,
6732 }
6733}
6734
6735fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6739 let mut count: u64 = 0;
6740 let mut sum: i128 = 0;
6741 let mut mn: i64 = i64::MAX;
6742 let mut mx: i64 = i64::MIN;
6743 while let Some(cols) = cursor.next_batch()? {
6744 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6745 if crate::columnar::all_non_null(validity, data.len()) {
6746 count += data.len() as u64;
6748 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6749 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6750 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6751 } else {
6752 for (i, &v) in data.iter().enumerate() {
6753 if crate::columnar::validity_bit(validity, i) {
6754 count += 1;
6755 sum += v as i128;
6756 mn = mn.min(v);
6757 mx = mx.max(v);
6758 }
6759 }
6760 }
6761 }
6762 }
6763 Ok((count, sum, mn, mx))
6764}
6765
6766fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6768 let mut count: u64 = 0;
6769 let mut sum: f64 = 0.0;
6770 let mut mn: f64 = f64::INFINITY;
6771 let mut mx: f64 = f64::NEG_INFINITY;
6772 while let Some(cols) = cursor.next_batch()? {
6773 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6774 if crate::columnar::all_non_null(validity, data.len()) {
6775 count += data.len() as u64;
6776 sum += data.iter().sum::<f64>();
6777 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6778 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6779 } else {
6780 for (i, &v) in data.iter().enumerate() {
6781 if crate::columnar::validity_bit(validity, i) {
6782 count += 1;
6783 sum += v;
6784 mn = mn.min(v);
6785 mx = mx.max(v);
6786 }
6787 }
6788 }
6789 }
6790 }
6791 Ok((count, sum, mn, mx))
6792}
6793
6794fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6795 if count == 0 && !matches!(agg, NativeAgg::Count) {
6796 return NativeAggResult::Null;
6797 }
6798 match agg {
6799 NativeAgg::Count => NativeAggResult::Count(count),
6800 NativeAgg::Sum => match sum.try_into() {
6803 Ok(v) => NativeAggResult::Int(v),
6804 Err(_) => NativeAggResult::Null,
6805 },
6806 NativeAgg::Min => NativeAggResult::Int(mn),
6807 NativeAgg::Max => NativeAggResult::Int(mx),
6808 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6809 }
6810}
6811
6812fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6813 if count == 0 && !matches!(agg, NativeAgg::Count) {
6814 return NativeAggResult::Null;
6815 }
6816 match agg {
6817 NativeAgg::Count => NativeAggResult::Count(count),
6818 NativeAgg::Sum => NativeAggResult::Float(sum),
6819 NativeAgg::Min => NativeAggResult::Float(mn),
6820 NativeAgg::Max => NativeAggResult::Float(mx),
6821 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6822 }
6823}
6824
6825fn agg_int(
6828 stats: &[crate::page::PageStat],
6829 decode: fn(Option<&[u8]>) -> Option<i64>,
6830) -> Option<(Option<i64>, Option<i64>, u64)> {
6831 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6832 let mut any = false;
6833 for s in stats {
6834 if let Some(v) = decode(s.min.as_deref()) {
6835 mn = mn.min(v);
6836 any = true;
6837 }
6838 if let Some(v) = decode(s.max.as_deref()) {
6839 mx = mx.max(v);
6840 any = true;
6841 }
6842 nulls += s.null_count;
6843 }
6844 any.then_some((Some(mn), Some(mx), nulls))
6845}
6846
6847fn agg_float(
6849 stats: &[crate::page::PageStat],
6850 decode: fn(Option<&[u8]>) -> Option<f64>,
6851) -> Option<(Option<f64>, Option<f64>, u64)> {
6852 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6853 let mut any = false;
6854 for s in stats {
6855 if let Some(v) = decode(s.min.as_deref()) {
6856 mn = mn.min(v);
6857 any = true;
6858 }
6859 if let Some(v) = decode(s.max.as_deref()) {
6860 mx = mx.max(v);
6861 any = true;
6862 }
6863 nulls += s.null_count;
6864 }
6865 any.then_some((Some(mn), Some(mx), nulls))
6866}
6867
6868type SecondaryIndexes = (
6870 HashMap<u16, BitmapIndex>,
6871 HashMap<u16, AnnIndex>,
6872 HashMap<u16, FmIndex>,
6873 HashMap<u16, SparseIndex>,
6874 HashMap<u16, MinHashIndex>,
6875);
6876
6877fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6878 let mut bitmap = HashMap::new();
6879 let mut ann = HashMap::new();
6880 let mut fm = HashMap::new();
6881 let mut sparse = HashMap::new();
6882 let mut minhash = HashMap::new();
6883 for idef in &schema.indexes {
6884 match idef.kind {
6885 IndexKind::Bitmap => {
6886 bitmap.insert(idef.column_id, BitmapIndex::new());
6887 }
6888 IndexKind::Ann => {
6889 let dim = schema
6890 .columns
6891 .iter()
6892 .find(|c| c.id == idef.column_id)
6893 .and_then(|c| match c.ty {
6894 TypeId::Embedding { dim } => Some(dim as usize),
6895 _ => None,
6896 })
6897 .unwrap_or(0);
6898 ann.insert(idef.column_id, AnnIndex::new(dim));
6899 }
6900 IndexKind::FmIndex => {
6901 fm.insert(idef.column_id, FmIndex::new());
6902 }
6903 IndexKind::Sparse => {
6904 sparse.insert(idef.column_id, SparseIndex::new());
6905 }
6906 IndexKind::MinHash => {
6907 minhash.insert(idef.column_id, MinHashIndex::new());
6908 }
6909 _ => {}
6910 }
6911 }
6912 (bitmap, ann, fm, sparse, minhash)
6913}
6914
6915const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
6916 | ColumnFlags::AUTO_INCREMENT
6917 | ColumnFlags::ENCRYPTED
6918 | ColumnFlags::ENCRYPTED_INDEXABLE
6919 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
6920
6921fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
6922 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
6923 return Err(MongrelError::Schema(
6924 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
6925 ));
6926 }
6927 Ok(())
6928}
6929
6930fn validate_alter_column_type(
6931 schema: &Schema,
6932 old: &ColumnDef,
6933 next: &ColumnDef,
6934 has_stored_versions: bool,
6935) -> Result<()> {
6936 if old.ty == next.ty {
6937 return Ok(());
6938 }
6939 if schema.indexes.iter().any(|i| i.column_id == old.id) {
6940 return Err(MongrelError::Schema(format!(
6941 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
6942 old.name
6943 )));
6944 }
6945 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
6946 return Ok(());
6947 }
6948 Err(MongrelError::Schema(format!(
6949 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
6950 old.ty, next.ty
6951 )))
6952}
6953
6954fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
6955 matches!(
6956 (old, new),
6957 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
6958 )
6959}
6960
6961fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
6967 let mut prev: Option<i64> = None;
6968 for r in rows {
6969 match r.columns.get(&pk_id) {
6970 Some(Value::Int64(v)) => {
6971 if prev.is_some_and(|p| p >= *v) {
6972 return false;
6973 }
6974 prev = Some(*v);
6975 }
6976 _ => return false,
6977 }
6978 }
6979 true
6980}
6981
6982#[allow(clippy::too_many_arguments)]
6983fn index_into(
6984 schema: &Schema,
6985 row: &Row,
6986 hot: &mut HotIndex,
6987 bitmap: &mut HashMap<u16, BitmapIndex>,
6988 ann: &mut HashMap<u16, AnnIndex>,
6989 fm: &mut HashMap<u16, FmIndex>,
6990 sparse: &mut HashMap<u16, SparseIndex>,
6991 minhash: &mut HashMap<u16, MinHashIndex>,
6992) {
6993 for idef in &schema.indexes {
6994 let Some(val) = row.columns.get(&idef.column_id) else {
6995 continue;
6996 };
6997 match idef.kind {
6998 IndexKind::Bitmap => {
6999 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7000 b.insert(val.encode_key(), row.row_id);
7001 }
7002 }
7003 IndexKind::Ann => {
7004 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7005 a.insert(v, row.row_id);
7006 }
7007 }
7008 IndexKind::FmIndex => {
7009 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7010 f.insert(b.clone(), row.row_id);
7011 }
7012 }
7013 IndexKind::Sparse => {
7014 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7015 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7018 s.insert(&terms, row.row_id);
7019 }
7020 }
7021 }
7022 IndexKind::MinHash => {
7023 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7024 let tokens = crate::index::token_hashes_from_bytes(b);
7027 mh.insert(&tokens, row.row_id);
7028 }
7029 }
7030 _ => {}
7031 }
7032 }
7033 if let Some(pk_col) = schema.primary_key() {
7034 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7035 hot.insert(pk_val.encode_key(), row.row_id);
7036 }
7037 }
7038}
7039
7040#[allow(dead_code)]
7046fn bulk_index_key(
7047 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7048 column_id: u16,
7049 ty: TypeId,
7050 col: &columnar::NativeColumn,
7051 i: usize,
7052) -> Option<Vec<u8>> {
7053 let encoded = columnar::encode_key_native(ty, col, i)?;
7054 #[cfg(feature = "encryption")]
7055 {
7056 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7057 if let Some((key, scheme)) = column_keys.get(&column_id) {
7058 return Some(match (*scheme, col) {
7059 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7060 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7061 ope_token_i64(key, data[i]).to_vec()
7062 }
7063 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7064 ope_token_f64(key, data[i]).to_vec()
7065 }
7066 _ => hmac_token(key, &encoded).to_vec(),
7067 });
7068 }
7069 }
7070 #[cfg(not(feature = "encryption"))]
7071 {
7072 let _ = (column_id, column_keys, col);
7073 }
7074 Some(encoded)
7075}
7076
7077pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7078 let json = serde_json::to_string_pretty(schema)
7079 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7080 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7081 Ok(())
7082}
7083
7084fn read_schema(dir: &Path) -> Result<Schema> {
7085 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7086 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7087}
7088
7089fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7090 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7091}
7092
7093fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7094 let n = list_wal_numbers(wal_dir)?;
7095 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7096}
7097
7098fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7099 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7100}
7101
7102fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7103 let _ = std::fs::create_dir_all(wal_dir);
7104 let mut max_n = None;
7105 for entry in std::fs::read_dir(wal_dir)? {
7106 let entry = entry?;
7107 let fname = entry.file_name();
7108 let Some(s) = fname.to_str() else {
7109 continue;
7110 };
7111 let Some(stripped) = s.strip_prefix("seg-") else {
7112 continue;
7113 };
7114 let Some(stripped) = stripped.strip_suffix(".wal") else {
7115 continue;
7116 };
7117 if let Ok(n) = stripped.parse::<u32>() {
7118 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7119 }
7120 }
7121 Ok(max_n)
7122}