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 { .. }
1208 | Op::ExternalTableState { .. }
1209 | Op::Flush { .. }
1210 | Op::Ddl(_) => {}
1211 }
1212 }
1213
1214 let rcache_dir = dir.join(RCACHE_DIR);
1215 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1216 let mut db = Self {
1217 dir,
1218 table_id: manifest.table_id,
1219 wal,
1220 memtable,
1221 mutable_run: MutableRun::new(),
1222 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1223 compaction_zstd_level: 3,
1224 allocator,
1225 epoch: ctx.epoch,
1226 persisted_epoch,
1227 schema,
1228 hot: HotIndex::new(),
1229 kek: ctx.kek,
1230 column_keys,
1231 run_refs: manifest.runs.clone(),
1232 retiring: manifest.retiring.clone(),
1233 next_run_id: manifest
1234 .runs
1235 .iter()
1236 .map(|r| r.run_id as u64 + 1)
1237 .max()
1238 .unwrap_or(1),
1239 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1240 current_txn_id,
1241 bitmap: HashMap::new(),
1242 ann: HashMap::new(),
1243 fm: HashMap::new(),
1244 sparse: HashMap::new(),
1245 minhash: HashMap::new(),
1246 learned_range: HashMap::new(),
1247 pk_by_row: HashMap::new(),
1248 pinned: BTreeMap::new(),
1249 live_count: manifest.live_count,
1250 reservoir: crate::reservoir::Reservoir::default(),
1251 reservoir_complete: false,
1252 had_deletes: saw_delete,
1253 agg_cache: HashMap::new(),
1254 global_idx_epoch: manifest.global_idx_epoch,
1255 indexes_complete: true,
1256 index_build_policy: IndexBuildPolicy::default(),
1257 pk_by_row_complete: false,
1258 flushed_epoch: manifest.flushed_epoch,
1259 page_cache: ctx.page_cache,
1260 decoded_cache: ctx.decoded_cache,
1261 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1262 snapshots: ctx.snapshots,
1263 commit_lock: ctx.commit_lock,
1264 result_cache: Arc::new(parking_lot::Mutex::new(
1265 ResultCache::new()
1266 .with_dir(rcache_dir)
1267 .with_cache_dek(cache_dek.clone()),
1268 )),
1269 pending_delete_rids: roaring::RoaringBitmap::new(),
1270 pending_put_cols: std::collections::HashSet::new(),
1271 pending_rows: Vec::new(),
1272 pending_rows_auto_inc: Vec::new(),
1273 pending_dels: Vec::new(),
1274 pending_truncate: None,
1275 wal_dek,
1276 auto_inc,
1277 };
1278
1279 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1282
1283 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1288 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1289 c.epoch_built == manifest.global_idx_epoch
1290 && manifest.global_idx_epoch > 0
1291 && manifest
1292 .runs
1293 .iter()
1294 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1295 });
1296 if let Some(loaded) = checkpoint {
1297 if checkpoint_valid {
1298 db.hot = loaded.hot;
1299 db.bitmap = loaded.bitmap;
1300 db.ann = loaded.ann;
1301 db.fm = loaded.fm;
1302 db.sparse = loaded.sparse;
1303 db.minhash = loaded.minhash;
1304 db.learned_range = loaded.learned_range;
1305 }
1308 }
1309 if !checkpoint_valid {
1310 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1311 db.bitmap = bitmap;
1312 db.ann = ann;
1313 db.fm = fm;
1314 db.sparse = sparse;
1315 db.minhash = minhash;
1316 db.rebuild_indexes_from_runs()?;
1317 db.build_learned_ranges()?;
1318 }
1319
1320 for (epoch, group) in replayed_puts {
1325 let (losers, winner_pks) = db.partition_pk_winners(&group);
1326 for (key, &row_id) in &winner_pks {
1327 if let Some(old_rid) = db.hot.get(key) {
1328 if old_rid != row_id {
1329 db.tombstone_row(old_rid, epoch, false);
1330 }
1331 }
1332 }
1333 for &loser_rid in &losers {
1334 db.tombstone_row(loser_rid, epoch, false);
1335 }
1336 for (key, row_id) in winner_pks {
1337 db.insert_hot_pk(key, row_id);
1338 }
1339 if db.schema.primary_key().is_none() {
1340 for r in &group {
1341 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1342 }
1343 }
1344 for r in &group {
1345 if !losers.contains(&r.row_id) {
1346 db.index_row(r);
1347 }
1348 }
1349 }
1350 for (rid, epoch) in &replayed_deletes {
1354 db.remove_hot_for_row(*rid, *epoch);
1355 }
1356
1357 db.result_cache.lock().load_persistent();
1364 Ok(db)
1365 }
1366
1367 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1373 if self.reservoir_complete {
1374 return Ok(());
1375 }
1376 self.rebuild_reservoir()?;
1377 self.reservoir_complete = true;
1378 Ok(())
1379 }
1380
1381 fn rebuild_reservoir(&mut self) -> Result<()> {
1384 let snap = self.snapshot();
1385 let rows = self.visible_rows(snap)?;
1386 self.reservoir.reset();
1387 for r in rows {
1388 self.reservoir.offer(r.row_id.0);
1389 }
1390 Ok(())
1391 }
1392
1393 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1394 self.hot = HotIndex::new();
1395 self.pk_by_row.clear();
1396 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1397 self.bitmap = bitmap;
1398 self.ann = ann;
1399 self.fm = fm;
1400 self.sparse = sparse;
1401 self.minhash = minhash;
1402 let snapshot = Epoch(u64::MAX);
1403 for rr in self.run_refs.clone() {
1404 let mut reader = self.open_reader(rr.run_id)?;
1405 for row in reader.visible_rows(snapshot)? {
1406 let tok_row = self.tokenized_for_indexes(&row);
1407 index_into(
1408 &self.schema,
1409 &tok_row,
1410 &mut self.hot,
1411 &mut self.bitmap,
1412 &mut self.ann,
1413 &mut self.fm,
1414 &mut self.sparse,
1415 &mut self.minhash,
1416 );
1417 }
1418 }
1419 for row in self.mutable_run.visible_versions(snapshot) {
1420 if row.deleted {
1421 self.remove_hot_for_row(row.row_id, snapshot);
1422 } else {
1423 self.index_row(&row);
1424 }
1425 }
1426 for row in self.memtable.visible_versions(snapshot) {
1427 if row.deleted {
1428 self.remove_hot_for_row(row.row_id, snapshot);
1429 } else {
1430 self.index_row(&row);
1431 }
1432 }
1433 self.refresh_pk_by_row_from_hot();
1434 Ok(())
1435 }
1436
1437 fn refresh_pk_by_row_from_hot(&mut self) {
1438 self.pk_by_row_complete = true;
1439 if self.schema.primary_key().is_none() {
1440 self.pk_by_row.clear();
1441 return;
1442 }
1443 self.pk_by_row = self
1449 .hot
1450 .entries()
1451 .into_iter()
1452 .map(|(key, row_id)| (row_id, key))
1453 .collect();
1454 }
1455
1456 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1457 if self.schema.primary_key().is_some() {
1458 self.pk_by_row.insert(row_id, key.clone());
1459 }
1460 self.hot.insert(key, row_id);
1461 }
1462
1463 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1467 self.learned_range.clear();
1468 if self.run_refs.len() != 1 {
1469 return Ok(());
1470 }
1471 let cols: Vec<u16> = self
1472 .schema
1473 .indexes
1474 .iter()
1475 .filter(|i| i.kind == IndexKind::LearnedRange)
1476 .map(|i| i.column_id)
1477 .collect();
1478 if cols.is_empty() {
1479 return Ok(());
1480 }
1481 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1482 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1483 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1484 _ => return Ok(()),
1485 };
1486 for cid in cols {
1487 let ty = self
1488 .schema
1489 .columns
1490 .iter()
1491 .find(|c| c.id == cid)
1492 .map(|c| c.ty)
1493 .unwrap_or(TypeId::Int64);
1494 match ty {
1495 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1496 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1497 let pairs: Vec<(i64, u64)> = data
1498 .iter()
1499 .zip(row_ids.iter())
1500 .map(|(v, r)| (*v, *r))
1501 .collect();
1502 self.learned_range
1503 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1504 }
1505 }
1506 TypeId::Float64 => {
1507 if let columnar::NativeColumn::Float64 { data, .. } =
1508 reader.column_native(cid)?
1509 {
1510 let pairs: Vec<(f64, u64)> = data
1511 .iter()
1512 .zip(row_ids.iter())
1513 .map(|(v, r)| (*v, *r))
1514 .collect();
1515 self.learned_range
1516 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1517 }
1518 }
1519 _ => {}
1520 }
1521 }
1522 Ok(())
1523 }
1524
1525 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1532 if self.indexes_complete {
1533 crate::trace::QueryTrace::record(|t| {
1534 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1535 });
1536 return Ok(());
1537 }
1538 crate::trace::QueryTrace::record(|t| {
1539 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1540 });
1541 self.rebuild_indexes_from_runs()?;
1542 self.build_learned_ranges()?;
1543 self.indexes_complete = true;
1544 let epoch = self.current_epoch();
1545 self.checkpoint_indexes(epoch);
1546 Ok(())
1547 }
1548
1549 fn pending_epoch(&self) -> Epoch {
1550 Epoch(self.epoch.visible().0 + 1)
1551 }
1552
1553 fn is_shared(&self) -> bool {
1556 matches!(self.wal, WalSink::Shared(_))
1557 }
1558
1559 fn ensure_txn_id(&mut self) -> u64 {
1563 if self.current_txn_id == 0 {
1564 let id = match &self.wal {
1565 WalSink::Shared(s) => {
1566 let mut g = s.txn_ids.lock();
1567 let v = *g;
1568 *g = g.wrapping_add(1);
1569 v
1570 }
1571 WalSink::Private(_) => 1,
1572 };
1573 self.current_txn_id = id;
1574 }
1575 self.current_txn_id
1576 }
1577
1578 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1581 let txn_id = self.ensure_txn_id();
1582 let table_id = self.table_id;
1583 match &mut self.wal {
1584 WalSink::Private(w) => {
1585 w.append_txn(txn_id, op)?;
1586 }
1587 WalSink::Shared(s) => {
1588 s.wal.lock().append(txn_id, table_id, op)?;
1589 }
1590 }
1591 Ok(())
1592 }
1593
1594 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1602 Ok(self.put_returning(columns)?.0)
1603 }
1604
1605 pub fn put_returning(
1610 &mut self,
1611 mut columns: Vec<(u16, Value)>,
1612 ) -> Result<(RowId, Option<i64>)> {
1613 let assigned = self.fill_auto_inc(&mut columns)?;
1614 self.schema.validate_not_null(&columns)?;
1615 let row_id = self.allocator.alloc();
1616 let epoch = self.pending_epoch();
1617 let mut row = Row::new(row_id, epoch);
1618 for (col_id, val) in columns {
1619 row.columns.insert(col_id, val);
1620 }
1621 self.commit_rows(vec![row], assigned.is_some())?;
1622 Ok((row_id, assigned))
1623 }
1624
1625 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1628 Ok(self
1629 .put_batch_returning(batch)?
1630 .into_iter()
1631 .map(|(r, _)| r)
1632 .collect())
1633 }
1634
1635 pub fn put_batch_returning(
1638 &mut self,
1639 batch: Vec<Vec<(u16, Value)>>,
1640 ) -> Result<Vec<(RowId, Option<i64>)>> {
1641 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1642 for mut cols in batch {
1643 let assigned = self.fill_auto_inc(&mut cols)?;
1644 filled.push((cols, assigned));
1645 }
1646 for (cols, _) in &filled {
1647 self.schema.validate_not_null(cols)?;
1648 }
1649 let epoch = self.pending_epoch();
1650 let mut rows = Vec::with_capacity(filled.len());
1651 let mut ids = Vec::with_capacity(filled.len());
1652 for (cols, assigned) in filled {
1653 let row_id = self.allocator.alloc();
1654 let mut row = Row::new(row_id, epoch);
1655 for (c, v) in cols {
1656 row.columns.insert(c, v);
1657 }
1658 ids.push((row_id, assigned));
1659 rows.push(row);
1660 }
1661 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1662 self.commit_rows(rows, all_auto_generated)?;
1663 Ok(ids)
1664 }
1665
1666 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1672 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1673 return Ok(None);
1674 };
1675 let pos = columns.iter().position(|(c, _)| *c == cid);
1676 let assigned = match pos {
1677 Some(i) => match &columns[i].1 {
1678 Value::Null => {
1679 let next = self.alloc_auto_inc_value()?;
1680 columns[i].1 = Value::Int64(next);
1681 Some(next)
1682 }
1683 Value::Int64(n) => {
1684 self.advance_auto_inc_past(*n)?;
1685 None
1686 }
1687 other => {
1688 return Err(MongrelError::InvalidArgument(format!(
1689 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1690 other
1691 )))
1692 }
1693 },
1694 None => {
1695 let next = self.alloc_auto_inc_value()?;
1696 columns.push((cid, Value::Int64(next)));
1697 Some(next)
1698 }
1699 };
1700 Ok(assigned)
1701 }
1702
1703 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1705 self.ensure_auto_inc_seeded()?;
1706 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1708 let v = ai.next;
1709 ai.next = ai.next.saturating_add(1);
1710 Ok(v)
1711 }
1712
1713 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1716 self.ensure_auto_inc_seeded()?;
1717 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1718 let floor = used.saturating_add(1).max(1);
1719 if ai.next < floor {
1720 ai.next = floor;
1721 }
1722 Ok(())
1723 }
1724
1725 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1730 let needs_seed = match self.auto_inc {
1731 Some(ai) => !ai.seeded,
1732 None => return Ok(()),
1733 };
1734 if !needs_seed {
1735 return Ok(());
1736 }
1737 if self.seed_empty_auto_inc() {
1738 return Ok(());
1739 }
1740 let cid = self
1741 .auto_inc
1742 .as_ref()
1743 .expect("auto-inc column present")
1744 .column_id;
1745 let max = self.scan_max_int64(cid)?;
1746 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1747 let floor = max.saturating_add(1).max(1);
1748 if ai.next < floor {
1749 ai.next = floor;
1750 }
1751 ai.seeded = true;
1752 Ok(())
1753 }
1754
1755 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1756 if n == 0 || self.auto_inc.is_none() {
1757 return Ok(None);
1758 }
1759 self.ensure_auto_inc_seeded()?;
1760 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1761 let start = ai.next;
1762 ai.next = ai.next.saturating_add(n as i64);
1763 Ok(Some(start))
1764 }
1765
1766 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1770 let mut max: i64 = 0;
1771 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1772 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1773 if *n > max {
1774 max = *n;
1775 }
1776 }
1777 }
1778 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1779 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1780 if *n > max {
1781 max = *n;
1782 }
1783 }
1784 }
1785 for rr in self.run_refs.clone() {
1786 let reader = self.open_reader(rr.run_id)?;
1787 if let Some(stats) = reader.column_page_stats(column_id) {
1788 for s in stats {
1789 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1790 if n > max {
1791 max = n;
1792 }
1793 }
1794 }
1795 } else if reader.has_column(column_id) {
1796 if let columnar::NativeColumn::Int64 { data, validity } =
1797 reader.column_native_shared(column_id)?
1798 {
1799 for (i, n) in data.iter().enumerate() {
1800 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1801 {
1802 max = *n;
1803 }
1804 }
1805 }
1806 }
1807 }
1808 Ok(max)
1809 }
1810
1811 fn seed_empty_auto_inc(&mut self) -> bool {
1812 let Some(ai) = self.auto_inc.as_mut() else {
1813 return false;
1814 };
1815 if ai.seeded || self.live_count != 0 {
1816 return false;
1817 }
1818 if ai.next < 1 {
1819 ai.next = 1;
1820 }
1821 ai.seeded = true;
1822 true
1823 }
1824
1825 fn advance_auto_inc_from_native_columns(
1826 &mut self,
1827 columns: &[(u16, columnar::NativeColumn)],
1828 n: usize,
1829 live_before: u64,
1830 ) -> Result<()> {
1831 let Some(ai) = self.auto_inc.as_mut() else {
1832 return Ok(());
1833 };
1834 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1835 return Ok(());
1836 };
1837 let columnar::NativeColumn::Int64 { data, validity } = col else {
1838 return Err(MongrelError::InvalidArgument(format!(
1839 "AUTO_INCREMENT column {} must be Int64",
1840 ai.column_id
1841 )));
1842 };
1843 let max = if native_int64_strictly_increasing(col, n) {
1844 data.get(n.saturating_sub(1)).copied()
1845 } else {
1846 data.iter()
1847 .take(n)
1848 .enumerate()
1849 .filter_map(|(i, v)| {
1850 if validity.is_empty() || columnar::validity_bit(validity, i) {
1851 Some(*v)
1852 } else {
1853 None
1854 }
1855 })
1856 .max()
1857 };
1858 if let Some(max) = max {
1859 let floor = max.saturating_add(1).max(1);
1860 if ai.next < floor {
1861 ai.next = floor;
1862 }
1863 if ai.seeded || live_before == 0 {
1864 ai.seeded = true;
1865 }
1866 }
1867 Ok(())
1868 }
1869
1870 fn fill_auto_inc_native_columns(
1871 &mut self,
1872 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1873 n: usize,
1874 ) -> Result<()> {
1875 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1876 return Ok(());
1877 };
1878 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1879 if let Some(start) = self.alloc_auto_inc_range(n)? {
1880 columns.push((
1881 cid,
1882 columnar::NativeColumn::Int64 {
1883 data: (start..start.saturating_add(n as i64)).collect(),
1884 validity: vec![0xFF; n.div_ceil(8)],
1885 },
1886 ));
1887 }
1888 return Ok(());
1889 };
1890
1891 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1892 return Err(MongrelError::InvalidArgument(format!(
1893 "AUTO_INCREMENT column {cid} must be Int64"
1894 )));
1895 };
1896 if data.len() < n {
1897 return Err(MongrelError::InvalidArgument(format!(
1898 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1899 data.len()
1900 )));
1901 }
1902 if columnar::all_non_null(validity, n) {
1903 return Ok(());
1904 }
1905 if validity.iter().all(|b| *b == 0) {
1906 if let Some(start) = self.alloc_auto_inc_range(n)? {
1907 for (i, slot) in data.iter_mut().take(n).enumerate() {
1908 *slot = start.saturating_add(i as i64);
1909 }
1910 *validity = vec![0xFF; n.div_ceil(8)];
1911 }
1912 return Ok(());
1913 }
1914
1915 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1916 for (i, slot) in data.iter_mut().enumerate().take(n) {
1917 if columnar::validity_bit(validity, i) {
1918 self.advance_auto_inc_past(*slot)?;
1919 } else {
1920 *slot = self.alloc_auto_inc_value()?;
1921 }
1922 }
1923 *validity = new_validity;
1924 Ok(())
1925 }
1926
1927 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1941 if self.auto_inc.is_none() {
1942 return Ok(None);
1943 }
1944 Ok(Some(self.alloc_auto_inc_value()?))
1945 }
1946
1947 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1953 let payload = bincode::serialize(&rows)?;
1954 self.wal_append_data(Op::Put {
1955 table_id: self.table_id,
1956 rows: payload,
1957 })?;
1958 if self.is_shared() {
1959 self.pending_rows_auto_inc
1960 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1961 self.pending_rows.extend(rows);
1962 } else {
1963 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1964 }
1965 Ok(())
1966 }
1967
1968 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1973 self.apply_put_rows_inner(rows, true)
1974 }
1975
1976 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1977 if check_existing_pk {
1978 self.ensure_indexes_complete()?;
1979 }
1980 if rows.len() == 1 {
1984 let row = rows.into_iter().next().expect("len checked");
1985 return self.apply_put_row_single(row, check_existing_pk);
1986 }
1987 let pk_id = self.schema.primary_key().map(|c| c.id);
2004 let probe = match pk_id {
2005 Some(pid) => {
2006 check_existing_pk
2007 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2008 }
2009 None => false,
2010 };
2011 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2014 for r in rows {
2015 for &cid in r.columns.keys() {
2016 self.pending_put_cols.insert(cid);
2017 }
2018 match pk_id {
2019 Some(pid) if probe || maintain_pk_by_row => {
2020 if let Some(pk_val) = r.columns.get(&pid) {
2021 let key = self.index_lookup_key(pid, pk_val);
2022 if probe {
2023 if let Some(old_rid) = self.hot.get(&key) {
2024 if old_rid != r.row_id {
2025 self.tombstone_row(old_rid, r.committed_epoch, true);
2026 }
2027 }
2028 }
2029 if maintain_pk_by_row {
2030 self.pk_by_row.insert(r.row_id, key);
2031 }
2032 }
2033 }
2034 Some(_) => {}
2035 None => {
2036 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2037 }
2038 }
2039 self.index_row(&r);
2040 self.reservoir.offer(r.row_id.0);
2041 self.memtable.upsert(r);
2042 self.live_count = self.live_count.saturating_add(1);
2045 }
2046 Ok(())
2047 }
2048
2049 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2053 for &cid in row.columns.keys() {
2054 self.pending_put_cols.insert(cid);
2055 }
2056 let epoch = row.committed_epoch;
2057 if let Some(pk_col) = self.schema.primary_key() {
2058 let pk_id = pk_col.id;
2059 if let Some(pk_val) = row.columns.get(&pk_id) {
2060 let maintain_pk_by_row = self.pk_by_row_complete;
2064 if check_existing_pk || maintain_pk_by_row {
2065 let key = self.index_lookup_key(pk_id, pk_val);
2066 if check_existing_pk {
2067 if let Some(old_rid) = self.hot.get(&key) {
2068 if old_rid != row.row_id {
2069 self.tombstone_row(old_rid, epoch, true);
2070 }
2071 }
2072 }
2073 if maintain_pk_by_row {
2074 self.pk_by_row.insert(row.row_id, key);
2075 }
2076 }
2077 }
2078 } else {
2079 self.hot
2080 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2081 }
2082 self.index_row(&row);
2083 self.reservoir.offer(row.row_id.0);
2084 self.memtable.upsert(row);
2085 self.live_count = self.live_count.saturating_add(1);
2086 Ok(())
2087 }
2088
2089 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2092 self.allocator.alloc()
2093 }
2094
2095 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2103 self.ensure_indexes_complete()?;
2104 let n = rows.len();
2105 for r in rows {
2106 for &cid in r.columns.keys() {
2107 self.pending_put_cols.insert(cid);
2108 }
2109 }
2110 let (losers, winner_pks) = self.partition_pk_winners(rows);
2111 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2112 for (key, &row_id) in &winner_pks {
2114 if let Some(old_rid) = self.hot.get(key) {
2115 if old_rid != row_id {
2116 self.tombstone_row(old_rid, epoch, true);
2117 }
2118 }
2119 }
2120 for &loser_rid in &losers {
2123 self.tombstone_row(loser_rid, epoch, false);
2124 }
2125 for (key, row_id) in winner_pks {
2127 self.insert_hot_pk(key, row_id);
2128 }
2129 if self.schema.primary_key().is_none() {
2130 for r in rows {
2131 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2132 }
2133 }
2134 for r in rows {
2135 self.allocator.advance_to(r.row_id);
2136 if !losers.contains(&r.row_id) {
2137 self.index_row(r);
2138 }
2139 }
2140 for r in rows {
2141 if !losers.contains(&r.row_id) {
2142 self.reservoir.offer(r.row_id.0);
2143 }
2144 }
2145 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2146 Ok(())
2147 }
2148
2149 pub(crate) fn recover_apply(
2154 &mut self,
2155 rows: Vec<Row>,
2156 deletes: Vec<(RowId, Epoch)>,
2157 ) -> Result<()> {
2158 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2162 std::collections::BTreeMap::new();
2163 for row in rows {
2164 self.allocator.advance_to(row.row_id);
2165 if let Some(ai) = self.auto_inc.as_mut() {
2170 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2171 if *n + 1 > ai.next {
2172 ai.next = *n + 1;
2173 }
2174 }
2175 }
2176 by_epoch.entry(row.committed_epoch).or_default().push(row);
2177 }
2178 for (epoch, group) in by_epoch {
2179 let (losers, winner_pks) = self.partition_pk_winners(&group);
2180 for (key, &row_id) in &winner_pks {
2182 if let Some(old_rid) = self.hot.get(key) {
2183 if old_rid != row_id {
2184 self.tombstone_row(old_rid, epoch, false);
2185 }
2186 }
2187 }
2188 for (key, row_id) in winner_pks {
2189 self.insert_hot_pk(key, row_id);
2190 }
2191 if self.schema.primary_key().is_none() {
2192 for r in &group {
2193 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2194 }
2195 }
2196 for r in &group {
2197 if !losers.contains(&r.row_id) {
2198 self.memtable.upsert(r.clone());
2199 self.index_row(r);
2200 }
2201 }
2202 }
2203 for (rid, epoch) in deletes {
2204 self.memtable.tombstone(rid, epoch);
2205 self.remove_hot_for_row(rid, epoch);
2206 }
2207 self.reservoir_complete = false;
2210 Ok(())
2211 }
2212
2213 pub(crate) fn flushed_epoch(&self) -> u64 {
2215 self.flushed_epoch
2216 }
2217
2218 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2220 self.schema.validate_not_null(cells)
2221 }
2222
2223 fn validate_columns_not_null(
2227 &self,
2228 columns: &[(u16, columnar::NativeColumn)],
2229 n: usize,
2230 ) -> Result<()> {
2231 let by_id: HashMap<u16, &columnar::NativeColumn> =
2232 columns.iter().map(|(id, c)| (*id, c)).collect();
2233 for col in &self.schema.columns {
2234 if col.flags.contains(ColumnFlags::NULLABLE) {
2235 continue;
2236 }
2237 match by_id.get(&col.id) {
2238 None => {
2239 return Err(MongrelError::InvalidArgument(format!(
2240 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2241 col.name, col.id
2242 )));
2243 }
2244 Some(c) => {
2245 if c.null_count(n) != 0 {
2246 return Err(MongrelError::InvalidArgument(format!(
2247 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2248 col.name, col.id
2249 )));
2250 }
2251 }
2252 }
2253 }
2254 Ok(())
2255 }
2256
2257 fn bulk_pk_winner_indices(
2262 &self,
2263 columns: &[(u16, columnar::NativeColumn)],
2264 n: usize,
2265 ) -> Option<Vec<usize>> {
2266 let pk_col = self.schema.primary_key()?;
2267 let pk_id = pk_col.id;
2268 let pk_ty = pk_col.ty;
2269 let by_id: HashMap<u16, &columnar::NativeColumn> =
2270 columns.iter().map(|(id, c)| (*id, c)).collect();
2271 let pk_native = by_id.get(&pk_id)?;
2272 if native_int64_strictly_increasing(pk_native, n) {
2273 return None;
2274 }
2275 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2277 let mut null_pk_rows: Vec<usize> = Vec::new();
2278 for i in 0..n {
2279 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2280 Some(key) => {
2281 last.insert(key, i);
2282 }
2283 None => null_pk_rows.push(i),
2284 }
2285 }
2286 let mut winners: HashSet<usize> = last.values().copied().collect();
2287 for i in null_pk_rows {
2288 winners.insert(i);
2289 }
2290 Some((0..n).filter(|i| winners.contains(i)).collect())
2291 }
2292
2293 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2295 let epoch = self.pending_epoch();
2296 self.wal_append_data(Op::Delete {
2297 table_id: self.table_id,
2298 row_ids: vec![row_id],
2299 })?;
2300 if self.is_shared() {
2301 self.pending_dels.push(row_id);
2302 } else {
2303 self.apply_delete(row_id, epoch);
2304 }
2305 Ok(())
2306 }
2307
2308 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2309 let pre = self.get(row_id, self.snapshot());
2310 self.delete(row_id)?;
2311 Ok(pre.map(|row| {
2312 let mut columns: Vec<_> = row.columns.into_iter().collect();
2313 columns.sort_by_key(|(id, _)| *id);
2314 OwnedRow { columns }
2315 }))
2316 }
2317
2318 pub fn truncate(&mut self) -> Result<()> {
2320 let epoch = self.pending_epoch();
2321 self.wal_append_data(Op::TruncateTable {
2322 table_id: self.table_id,
2323 })?;
2324 self.pending_rows.clear();
2325 self.pending_rows_auto_inc.clear();
2326 self.pending_dels.clear();
2327 self.pending_truncate = Some(epoch);
2328 Ok(())
2329 }
2330
2331 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2333 for rr in std::mem::take(&mut self.run_refs) {
2334 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2335 }
2336 for r in std::mem::take(&mut self.retiring) {
2337 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2338 }
2339 self.memtable = Memtable::new();
2340 self.mutable_run = MutableRun::new();
2341 self.hot = HotIndex::new();
2342 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2343 self.bitmap = bitmap;
2344 self.ann = ann;
2345 self.fm = fm;
2346 self.sparse = sparse;
2347 self.minhash = minhash;
2348 self.learned_range.clear();
2349 self.pk_by_row.clear();
2350 self.pk_by_row_complete = false;
2351 self.live_count = 0;
2352 self.reservoir = crate::reservoir::Reservoir::default();
2353 self.reservoir_complete = true;
2354 self.had_deletes = true;
2355 self.agg_cache.clear();
2356 self.global_idx_epoch = 0;
2357 self.indexes_complete = true;
2358 self.pending_delete_rids.clear();
2359 self.pending_put_cols.clear();
2360 self.pending_rows.clear();
2361 self.pending_rows_auto_inc.clear();
2362 self.pending_dels.clear();
2363 self.clear_result_cache();
2364 self.invalidate_index_checkpoint();
2365 Ok(())
2366 }
2367
2368 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2371 self.remove_hot_for_row(row_id, epoch);
2372 self.tombstone_row(row_id, epoch, true);
2373 }
2374
2375 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2379 let tombstone = Row {
2380 row_id,
2381 committed_epoch: epoch,
2382 columns: std::collections::HashMap::new(),
2383 deleted: true,
2384 };
2385 self.memtable.upsert(tombstone);
2386 self.pk_by_row.remove(&row_id);
2387 if adjust_live_count {
2388 self.live_count = self.live_count.saturating_sub(1);
2389 }
2390 self.pending_delete_rids.insert(row_id.0 as u32);
2392 self.had_deletes = true;
2395 self.agg_cache.clear();
2396 }
2397
2398 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2402 let Some(pk_col) = self.schema.primary_key() else {
2403 return;
2404 };
2405 if self.pk_by_row_complete {
2408 if let Some(key) = self.pk_by_row.remove(&row_id) {
2409 if self.hot.get(&key) == Some(row_id) {
2410 self.hot.remove(&key);
2411 }
2412 }
2413 return;
2414 }
2415 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2434 if self.indexes_complete {
2435 let pk_val = self
2436 .memtable
2437 .get_version(row_id, lookup_epoch)
2438 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2439 .or_else(|| {
2440 self.mutable_run
2441 .get_version(row_id, lookup_epoch)
2442 .filter(|(_, r)| !r.deleted)
2443 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2444 })
2445 .or_else(|| {
2446 self.run_refs.iter().find_map(|rr| {
2447 let mut reader = self.open_reader(rr.run_id).ok()?;
2448 let (_, deleted, val) = reader
2449 .get_version_column(row_id, lookup_epoch, pk_col.id)
2450 .ok()??;
2451 if deleted {
2452 return None;
2453 }
2454 val
2455 })
2456 });
2457 if let Some(pk_val) = pk_val {
2458 let key = self.index_lookup_key(pk_col.id, &pk_val);
2459 if self.hot.get(&key) == Some(row_id) {
2460 self.hot.remove(&key);
2461 }
2462 return;
2463 }
2464 }
2465 self.refresh_pk_by_row_from_hot();
2470 if let Some(key) = self.pk_by_row.remove(&row_id) {
2471 if self.hot.get(&key) == Some(row_id) {
2472 self.hot.remove(&key);
2473 }
2474 }
2475 }
2476
2477 fn partition_pk_winners(
2482 &self,
2483 rows: &[Row],
2484 ) -> (
2485 std::collections::HashSet<RowId>,
2486 std::collections::HashMap<Vec<u8>, RowId>,
2487 ) {
2488 let mut losers = std::collections::HashSet::new();
2489 let Some(pk_col) = self.schema.primary_key() else {
2490 return (losers, std::collections::HashMap::new());
2491 };
2492 let pk_id = pk_col.id;
2493 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2494 std::collections::HashMap::new();
2495 for r in rows {
2496 let Some(pk_val) = r.columns.get(&pk_id) else {
2497 continue;
2498 };
2499 let key = self.index_lookup_key(pk_id, pk_val);
2500 if let Some(&old_rid) = winners.get(&key) {
2501 losers.insert(old_rid);
2502 }
2503 winners.insert(key, r.row_id);
2504 }
2505 (losers, winners)
2506 }
2507
2508 fn index_row(&mut self, row: &Row) {
2509 if row.deleted {
2510 return;
2511 }
2512 if self.column_keys.is_empty() {
2516 index_into(
2517 &self.schema,
2518 row,
2519 &mut self.hot,
2520 &mut self.bitmap,
2521 &mut self.ann,
2522 &mut self.fm,
2523 &mut self.sparse,
2524 &mut self.minhash,
2525 );
2526 return;
2527 }
2528 let effective_row = self.tokenized_for_indexes(row);
2529 index_into(
2530 &self.schema,
2531 &effective_row,
2532 &mut self.hot,
2533 &mut self.bitmap,
2534 &mut self.ann,
2535 &mut self.fm,
2536 &mut self.sparse,
2537 &mut self.minhash,
2538 );
2539 }
2540
2541 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2547 if self.column_keys.is_empty() {
2548 return row.clone();
2549 }
2550 #[cfg(feature = "encryption")]
2551 {
2552 use crate::encryption::SCHEME_HMAC_EQ;
2553 let mut tok = row.clone();
2554 for (&cid, &(_, scheme)) in &self.column_keys {
2555 if scheme != SCHEME_HMAC_EQ {
2556 continue;
2557 }
2558 if let Some(v) = tok.columns.get(&cid).cloned() {
2559 if let Some(t) = self.tokenize_value(cid, &v) {
2560 tok.columns.insert(cid, t);
2561 }
2562 }
2563 }
2564 tok
2565 }
2566 #[cfg(not(feature = "encryption"))]
2567 {
2568 row.clone()
2569 }
2570 }
2571
2572 pub fn commit(&mut self) -> Result<Epoch> {
2577 if self.is_shared() {
2578 self.commit_shared()
2579 } else {
2580 self.commit_private()
2581 }
2582 }
2583
2584 fn commit_private(&mut self) -> Result<Epoch> {
2586 let commit_lock = Arc::clone(&self.commit_lock);
2590 let _g = commit_lock.lock();
2591 let new_epoch = self.epoch.bump_assigned();
2592 let txn_id = self.current_txn_id;
2593 match &mut self.wal {
2597 WalSink::Private(w) => {
2598 w.append_txn(
2599 txn_id,
2600 Op::TxnCommit {
2601 epoch: new_epoch.0,
2602 added_runs: Vec::new(),
2603 },
2604 )?;
2605 w.sync()?;
2606 }
2607 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2608 }
2609 if let Some(epoch) = self.pending_truncate.take() {
2611 self.apply_truncate(epoch)?;
2612 }
2613 self.invalidate_pending_cache();
2614 self.persist_manifest(new_epoch)?;
2615 self.epoch.publish_in_order(new_epoch);
2619 self.current_txn_id += 1;
2620 Ok(new_epoch)
2621 }
2622
2623 fn commit_shared(&mut self) -> Result<Epoch> {
2629 use std::sync::atomic::Ordering;
2630 let s = match &self.wal {
2631 WalSink::Shared(s) => s.clone(),
2632 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2633 };
2634 if s.poisoned.load(Ordering::Relaxed) {
2635 return Err(MongrelError::Other(
2636 "database poisoned by fsync error".into(),
2637 ));
2638 }
2639 let commit_lock = Arc::clone(&self.commit_lock);
2646 let _g = commit_lock.lock();
2647 let txn_id = self.ensure_txn_id();
2650 let (new_epoch, commit_seq) = {
2651 let mut wal = s.wal.lock();
2652 let new_epoch = self.epoch.bump_assigned();
2653 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2654 (new_epoch, seq)
2655 };
2656 s.group
2657 .await_durable(&s.wal, commit_seq)
2658 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2659
2660 if self.pending_truncate.take().is_some() {
2663 self.apply_truncate(new_epoch)?;
2664 }
2665 let mut rows = std::mem::take(&mut self.pending_rows);
2666 if !rows.is_empty() {
2667 for r in &mut rows {
2668 r.committed_epoch = new_epoch;
2669 }
2670 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2671 let all_auto_generated =
2672 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2673 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2674 } else {
2675 self.pending_rows_auto_inc.clear();
2676 }
2677 let dels = std::mem::take(&mut self.pending_dels);
2678 for rid in dels {
2679 self.apply_delete(rid, new_epoch);
2680 }
2681
2682 self.invalidate_pending_cache();
2683 self.persist_manifest(new_epoch)?;
2684 self.epoch.publish_in_order(new_epoch);
2685 self.current_txn_id = 0;
2687 Ok(new_epoch)
2688 }
2689
2690 pub fn flush(&mut self) -> Result<Epoch> {
2698 self.ensure_indexes_complete()?;
2699 let epoch = self.commit()?;
2700 let rows = self.memtable.drain_sorted();
2701 if !rows.is_empty() {
2702 self.mutable_run.insert_many(rows);
2703 }
2704 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2705 self.spill_mutable_run(epoch)?;
2706 self.mark_flushed(epoch)?;
2710 self.persist_manifest(epoch)?;
2711 self.build_learned_ranges()?;
2712 self.checkpoint_indexes(epoch);
2715 }
2716 Ok(epoch)
2719 }
2720
2721 pub fn force_flush(&mut self) -> Result<Epoch> {
2730 let saved = self.mutable_run_spill_bytes;
2731 self.mutable_run_spill_bytes = 1;
2732 let result = self.flush();
2733 self.mutable_run_spill_bytes = saved;
2734 result
2735 }
2736
2737 pub fn close(&mut self) -> Result<()> {
2744 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
2745 self.force_flush()?;
2746 }
2747 Ok(())
2748 }
2749
2750 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2757 let op = Op::Flush {
2758 table_id: self.table_id,
2759 flushed_epoch: epoch.0,
2760 };
2761 match &mut self.wal {
2762 WalSink::Private(w) => {
2763 w.append_system(op)?;
2764 w.sync()?;
2765 }
2766 WalSink::Shared(s) => {
2767 s.wal.lock().append_system(op)?;
2772 }
2773 }
2774 self.flushed_epoch = epoch.0;
2775 if matches!(self.wal, WalSink::Private(_)) {
2776 self.rotate_wal(epoch)?;
2777 }
2778 Ok(())
2779 }
2780
2781 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2785 let rows = self.mutable_run.drain_sorted();
2786 if rows.is_empty() {
2787 return Ok(());
2788 }
2789 let run_id = self.next_run_id;
2790 self.next_run_id += 1;
2791 let path = self.run_path(run_id);
2792 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2793 if let Some(kek) = &self.kek {
2794 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2795 }
2796 let header = writer.write(&path, &rows)?;
2797 self.run_refs.push(RunRef {
2798 run_id: run_id as u128,
2799 level: 0,
2800 epoch_created: epoch.0,
2801 row_count: header.row_count,
2802 });
2803 Ok(())
2804 }
2805
2806 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2810 self.mutable_run_spill_bytes = bytes.max(1);
2811 }
2812
2813 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2817 self.compaction_zstd_level = level;
2818 }
2819
2820 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2824 self.result_cache.lock().set_max_bytes(max_bytes);
2825 }
2826
2827 pub(crate) fn clear_result_cache(&mut self) {
2831 self.result_cache.lock().clear();
2832 }
2833
2834 pub fn mutable_run_len(&self) -> usize {
2836 self.mutable_run.len()
2837 }
2838
2839 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2842 self.mutable_run.drain_sorted()
2843 }
2844
2845 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2850 let epoch = self.commit()?;
2851 let n = batch.len();
2852 if n == 0 {
2853 return Ok(epoch);
2854 }
2855 let live_before = self.live_count;
2856 self.spill_mutable_run(epoch)?;
2860 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
2861 && self.indexes_complete
2862 && self.run_refs.is_empty()
2863 && self.memtable.is_empty()
2864 && self.mutable_run.is_empty();
2865 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2871 use rayon::prelude::*;
2872 self.schema
2873 .columns
2874 .par_iter()
2875 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
2876 .collect::<Vec<_>>()
2877 };
2878 drop(batch);
2879 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2884 self.validate_columns_not_null(&user_columns, n)?;
2885 let winner_idx = self
2886 .bulk_pk_winner_indices(&user_columns, n)
2887 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2888 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2889 match winner_idx.as_deref() {
2890 Some(idx) => {
2891 let compacted = user_columns
2892 .iter()
2893 .map(|(id, c)| (*id, c.gather(idx)))
2894 .collect();
2895 (compacted, idx.len())
2896 }
2897 None => (std::mem::take(&mut user_columns), n),
2898 };
2899 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2900 let first = self.allocator.alloc_range(write_n as u64).0;
2901 for rid in first..first + write_n as u64 {
2902 self.reservoir.offer(rid);
2903 }
2904 let run_id = self.next_run_id;
2905 self.next_run_id += 1;
2906 let path = self.run_path(run_id);
2907 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2908 .clean(true)
2909 .with_lz4()
2910 .with_native_endian();
2911 if let Some(kek) = &self.kek {
2912 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2913 }
2914 let header = writer.write_native(&path, &write_columns, write_n, first)?;
2915 self.run_refs.push(RunRef {
2916 run_id: run_id as u128,
2917 level: 0,
2918 epoch_created: epoch.0,
2919 row_count: header.row_count,
2920 });
2921 self.live_count = self.live_count.saturating_add(write_n as u64);
2922 if eager_index_build {
2923 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
2924 self.index_columns_bulk(&write_columns, &row_ids);
2925 self.indexes_complete = true;
2926 self.build_learned_ranges()?;
2927 } else {
2928 self.indexes_complete = false;
2929 }
2930 self.mark_flushed(epoch)?;
2931 self.persist_manifest(epoch)?;
2932 if eager_index_build {
2933 self.checkpoint_indexes(epoch);
2934 }
2935 self.clear_result_cache();
2936 Ok(epoch)
2937 }
2938
2939 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
2942 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
2943 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
2944 let segment_no = segment
2947 .file_stem()
2948 .and_then(|s| s.to_str())
2949 .and_then(|s| s.strip_prefix("seg-"))
2950 .and_then(|s| s.parse::<u64>().ok())
2951 .unwrap_or(0);
2952 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
2953 wal.set_sync_byte_threshold(self.sync_byte_threshold);
2954 wal.sync()?;
2955 self.wal = WalSink::Private(wal);
2956 Ok(())
2957 }
2958
2959 pub(crate) fn invalidate_pending_cache(&mut self) {
2964 self.result_cache
2965 .lock()
2966 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
2967 self.pending_delete_rids.clear();
2968 self.pending_put_cols.clear();
2969 }
2970
2971 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
2972 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
2973 m.current_epoch = epoch.0;
2974 m.next_row_id = self.allocator.current().0;
2975 m.runs = self.run_refs.clone();
2976 m.live_count = self.live_count;
2977 m.global_idx_epoch = self.global_idx_epoch;
2978 m.flushed_epoch = self.flushed_epoch;
2979 m.retiring = self.retiring.clone();
2980 m.auto_inc_next = match self.auto_inc {
2984 Some(ai) if ai.seeded => ai.next,
2985 _ => 0,
2986 };
2987 let meta_dek = self.manifest_meta_dek();
2988 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
2989 Ok(())
2990 }
2991
2992 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
2998 if !self.indexes_complete {
3001 return;
3002 }
3003 let snap = global_idx::IndexSnapshot {
3004 hot: &self.hot,
3005 bitmap: &self.bitmap,
3006 ann: &self.ann,
3007 fm: &self.fm,
3008 sparse: &self.sparse,
3009 minhash: &self.minhash,
3010 learned_range: &self.learned_range,
3011 };
3012 let idx_dek = self.idx_dek();
3014 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3015 .is_ok()
3016 {
3017 self.global_idx_epoch = epoch.0;
3018 let _ = self.persist_manifest(epoch);
3019 }
3020 }
3021
3022 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3025 self.global_idx_epoch = 0;
3026 global_idx::remove(&self.dir);
3027 let _ = self.persist_manifest(self.epoch.visible());
3028 }
3029
3030 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3033 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3034 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3035 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3036 best = Some((epoch, row));
3037 }
3038 }
3039 for rr in &self.run_refs {
3040 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3041 continue;
3042 };
3043 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3044 continue;
3045 };
3046 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3047 best = Some((epoch, row));
3048 }
3049 }
3050 match best {
3051 Some((_, r)) if r.deleted => None,
3052 Some((_, r)) => Some(r),
3053 None => None,
3054 }
3055 }
3056
3057 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3061 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3062 let mut fold = |row: Row| {
3063 best.entry(row.row_id.0)
3064 .and_modify(|e| {
3065 if row.committed_epoch > e.0 {
3066 *e = (row.committed_epoch, row.clone());
3067 }
3068 })
3069 .or_insert_with(|| (row.committed_epoch, row));
3070 };
3071 for row in self.memtable.visible_versions(snapshot.epoch) {
3072 fold(row);
3073 }
3074 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3075 fold(row);
3076 }
3077 for rr in &self.run_refs {
3078 let mut reader = self.open_reader(rr.run_id)?;
3079 for row in reader.visible_versions(snapshot.epoch)? {
3080 fold(row);
3081 }
3082 }
3083 let mut out: Vec<Row> = best
3084 .into_values()
3085 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3086 .collect();
3087 out.sort_by_key(|r| r.row_id);
3088 Ok(out)
3089 }
3090
3091 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3098 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3099 let rr = self.run_refs[0].clone();
3100 let mut reader = self.open_reader(rr.run_id)?;
3101 let idxs = reader.visible_indices(snapshot.epoch)?;
3102 let mut cols = Vec::with_capacity(self.schema.columns.len());
3103 for cdef in &self.schema.columns {
3104 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3105 }
3106 return Ok(cols);
3107 }
3108 let rows = self.visible_rows(snapshot)?;
3110 let mut cols: Vec<(u16, Vec<Value>)> = self
3111 .schema
3112 .columns
3113 .iter()
3114 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3115 .collect();
3116 for r in &rows {
3117 for (cid, vec) in cols.iter_mut() {
3118 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3119 }
3120 }
3121 Ok(cols)
3122 }
3123
3124 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3126 self.hot.get(key)
3127 }
3128
3129 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3134 self.ensure_indexes_complete()?;
3135 let snapshot = self.snapshot();
3136 crate::trace::QueryTrace::record(|t| {
3137 t.run_count = self.run_refs.len();
3138 t.memtable_rows = self.memtable.len();
3139 t.mutable_run_rows = self.mutable_run.len();
3140 });
3141 if q.conditions.is_empty() {
3145 crate::trace::QueryTrace::record(|t| {
3146 t.scan_mode = crate::trace::ScanMode::Materialized;
3147 t.row_materialized = true;
3148 });
3149 return self.visible_rows(snapshot);
3150 }
3151 crate::trace::QueryTrace::record(|t| {
3152 t.conditions_pushed = q.conditions.len();
3153 t.scan_mode = crate::trace::ScanMode::Materialized;
3154 t.row_materialized = true;
3155 });
3156 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3163 ordered.sort_by_key(|c| condition_cost_rank(c));
3164 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3165 for c in &ordered {
3166 let s = self.resolve_condition(c, snapshot)?;
3167 let empty = s.is_empty();
3168 sets.push(s);
3169 if empty {
3170 break;
3171 }
3172 }
3173 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3174 self.rows_for_rids(&rids, snapshot)
3175 }
3176
3177 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3182 use std::collections::HashMap;
3183 let mut rows = Vec::with_capacity(rids.len());
3184 let tier_size = self.memtable.len() + self.mutable_run.len();
3201 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3202 if rids.len().saturating_mul(24) < tier_size {
3203 for &rid in rids {
3204 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3205 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3206 let newest = match (mem, mrun) {
3207 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3208 (Some((_, mr)), None) => Some(mr),
3209 (None, Some((_, rr))) => Some(rr),
3210 (None, None) => None,
3211 };
3212 if let Some(row) = newest {
3213 overlay.insert(rid, row);
3214 }
3215 }
3216 } else {
3217 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3218 overlay
3219 .entry(row.row_id.0)
3220 .and_modify(|e| {
3221 if row.committed_epoch > e.committed_epoch {
3222 *e = row.clone();
3223 }
3224 })
3225 .or_insert(row);
3226 };
3227 for row in self.memtable.visible_versions(snapshot.epoch) {
3228 fold_newest(row, &mut overlay);
3229 }
3230 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3231 fold_newest(row, &mut overlay);
3232 }
3233 }
3234 if self.run_refs.len() == 1 {
3235 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3236 if rids.len().saturating_mul(24) < reader.row_count() {
3244 for &rid in rids {
3245 if let Some(r) = overlay.get(&rid) {
3246 if !r.deleted {
3247 rows.push(r.clone());
3248 }
3249 continue;
3250 }
3251 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3252 if !row.deleted {
3253 rows.push(row);
3254 }
3255 }
3256 }
3257 return Ok(rows);
3258 }
3259 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3268 enum Src {
3271 Overlay,
3272 Run,
3273 }
3274 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3275 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3276 for rid in rids {
3277 if overlay.contains_key(rid) {
3278 plan.push(Src::Overlay);
3279 continue;
3280 }
3281 match vis_rids.binary_search(&(*rid as i64)) {
3282 Ok(i) => {
3283 plan.push(Src::Run);
3284 fetch.push(positions[i]);
3285 }
3286 Err(_) => { }
3287 }
3288 }
3289 let fetched = reader.materialize_batch(&fetch)?;
3290 let mut fetched_iter = fetched.into_iter();
3291 for (rid, src) in rids.iter().zip(plan) {
3292 match src {
3293 Src::Overlay => {
3294 if let Some(r) = overlay.get(rid) {
3295 if !r.deleted {
3296 rows.push(r.clone());
3297 }
3298 }
3299 }
3300 Src::Run => {
3301 if let Some(row) = fetched_iter.next() {
3302 if !row.deleted {
3303 rows.push(row);
3304 }
3305 }
3306 }
3307 }
3308 }
3309 return Ok(rows);
3310 }
3311 let mut readers: Vec<_> = self
3315 .run_refs
3316 .iter()
3317 .map(|rr| self.open_reader(rr.run_id))
3318 .collect::<Result<Vec<_>>>()?;
3319 for rid in rids {
3320 if let Some(r) = overlay.get(rid) {
3321 if !r.deleted {
3322 rows.push(r.clone());
3323 }
3324 continue;
3325 }
3326 let mut best: Option<(Epoch, Row)> = None;
3327 for reader in readers.iter_mut() {
3328 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3329 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3330 best = Some((epoch, row));
3331 }
3332 }
3333 }
3334 if let Some((_, r)) = best {
3335 if !r.deleted {
3336 rows.push(r);
3337 }
3338 }
3339 }
3340 Ok(rows)
3341 }
3342
3343 pub fn indexes_complete(&self) -> bool {
3353 self.indexes_complete
3354 }
3355
3356 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3358 self.index_build_policy
3359 }
3360
3361 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3365 self.index_build_policy = policy;
3366 }
3367
3368 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3373 if !self.indexes_complete {
3377 return None;
3378 }
3379 let b = self.bitmap.get(&column_id)?;
3380 let result: Vec<Vec<u8>> = b
3381 .keys()
3382 .into_iter()
3383 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3384 .cloned()
3385 .collect();
3386 Some(result)
3387 }
3388
3389 pub fn fk_join_row_ids(
3390 &self,
3391 fk_column_id: u16,
3392 pk_values: &[Vec<u8>],
3393 fk_conditions: &[crate::query::Condition],
3394 snapshot: Snapshot,
3395 ) -> Result<Vec<u64>> {
3396 let Some(b) = self.bitmap.get(&fk_column_id) else {
3397 return Ok(Vec::new());
3398 };
3399 let mut join_set = {
3400 let mut acc = roaring::RoaringBitmap::new();
3401 for v in pk_values {
3402 acc |= b.get(v);
3403 }
3404 RowIdSet::from_roaring(acc)
3405 };
3406 if !fk_conditions.is_empty() {
3407 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3408 sets.push(join_set);
3409 for c in fk_conditions {
3410 sets.push(self.resolve_condition(c, snapshot)?);
3411 }
3412 join_set = RowIdSet::intersect_many(sets);
3413 }
3414 Ok(join_set.into_sorted_vec())
3415 }
3416
3417 pub fn fk_join_count(
3423 &self,
3424 fk_column_id: u16,
3425 pk_values: &[Vec<u8>],
3426 fk_conditions: &[crate::query::Condition],
3427 snapshot: Snapshot,
3428 ) -> Result<u64> {
3429 let Some(b) = self.bitmap.get(&fk_column_id) else {
3430 return Ok(0);
3431 };
3432 let mut acc = roaring::RoaringBitmap::new();
3433 for v in pk_values {
3434 acc |= b.get(v);
3435 }
3436 if fk_conditions.is_empty() {
3437 return Ok(acc.len());
3438 }
3439 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3440 sets.push(RowIdSet::from_roaring(acc));
3441 for c in fk_conditions {
3442 sets.push(self.resolve_condition(c, snapshot)?);
3443 }
3444 Ok(RowIdSet::intersect_many(sets).len() as u64)
3445 }
3446
3447 fn resolve_condition(
3452 &self,
3453 c: &crate::query::Condition,
3454 snapshot: Snapshot,
3455 ) -> Result<RowIdSet> {
3456 use crate::query::Condition;
3457 Ok(match c {
3458 Condition::Pk(key) => {
3459 let lookup = self
3460 .schema
3461 .primary_key()
3462 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3463 .unwrap_or_else(|| key.clone());
3464 self.hot
3465 .get(&lookup)
3466 .map(|r| RowIdSet::one(r.0))
3467 .unwrap_or_else(RowIdSet::empty)
3468 }
3469 Condition::BitmapEq { column_id, value } => {
3470 let lookup = self.index_lookup_key_bytes(*column_id, value);
3471 self.bitmap
3472 .get(column_id)
3473 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3474 .unwrap_or_else(RowIdSet::empty)
3475 }
3476 Condition::BitmapIn { column_id, values } => {
3477 let bm = self.bitmap.get(column_id);
3478 let mut acc = roaring::RoaringBitmap::new();
3479 if let Some(b) = bm {
3480 for v in values {
3481 let lookup = self.index_lookup_key_bytes(*column_id, v);
3482 acc |= b.get(&lookup);
3483 }
3484 }
3485 RowIdSet::from_roaring(acc)
3486 }
3487 Condition::BytesPrefix { column_id, prefix } => {
3488 if let Some(b) = self.bitmap.get(column_id) {
3493 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3494 let mut acc = roaring::RoaringBitmap::new();
3495 for key in b.keys() {
3496 if key.starts_with(&lookup_prefix) {
3497 acc |= b.get(key);
3498 }
3499 }
3500 RowIdSet::from_roaring(acc)
3501 } else {
3502 RowIdSet::empty()
3503 }
3504 }
3505 Condition::FmContains { column_id, pattern } => self
3506 .fm
3507 .get(column_id)
3508 .map(|f| {
3509 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3510 })
3511 .unwrap_or_else(RowIdSet::empty),
3512 Condition::FmContainsAll {
3513 column_id,
3514 patterns,
3515 } => {
3516 if let Some(f) = self.fm.get(column_id) {
3519 let sets: Vec<RowIdSet> = patterns
3520 .iter()
3521 .map(|pat| {
3522 RowIdSet::from_unsorted(
3523 f.locate(pat).into_iter().map(|r| r.0).collect(),
3524 )
3525 })
3526 .collect();
3527 RowIdSet::intersect_many(sets)
3528 } else {
3529 RowIdSet::empty()
3530 }
3531 }
3532 Condition::Ann {
3533 column_id,
3534 query,
3535 k,
3536 } => self
3537 .ann
3538 .get(column_id)
3539 .map(|a| {
3540 RowIdSet::from_unsorted(
3541 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3542 )
3543 })
3544 .unwrap_or_else(RowIdSet::empty),
3545 Condition::SparseMatch {
3546 column_id,
3547 query,
3548 k,
3549 } => self
3550 .sparse
3551 .get(column_id)
3552 .map(|s| {
3553 RowIdSet::from_unsorted(
3554 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3555 )
3556 })
3557 .unwrap_or_else(RowIdSet::empty),
3558 Condition::MinHashSimilar {
3559 column_id,
3560 query,
3561 k,
3562 } => self
3563 .minhash
3564 .get(column_id)
3565 .map(|mh| {
3566 RowIdSet::from_unsorted(
3567 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3568 )
3569 })
3570 .unwrap_or_else(RowIdSet::empty),
3571 Condition::Range { column_id, lo, hi } => {
3572 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3581 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3582 } else if self.run_refs.len() == 1 {
3583 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3584 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3585 } else {
3586 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3587 };
3588 set.remove_many(self.overlay_rid_set(snapshot));
3589 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3590 set
3591 }
3592 Condition::RangeF64 {
3593 column_id,
3594 lo,
3595 lo_inclusive,
3596 hi,
3597 hi_inclusive,
3598 } => {
3599 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3602 RowIdSet::from_unsorted(
3603 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3604 .into_iter()
3605 .collect(),
3606 )
3607 } else if self.run_refs.len() == 1 {
3608 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3609 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3610 } else {
3611 return self.range_scan_f64(
3612 *column_id,
3613 *lo,
3614 *lo_inclusive,
3615 *hi,
3616 *hi_inclusive,
3617 snapshot,
3618 );
3619 };
3620 set.remove_many(self.overlay_rid_set(snapshot));
3621 self.range_scan_overlay_f64(
3622 &mut set,
3623 *column_id,
3624 *lo,
3625 *lo_inclusive,
3626 *hi,
3627 *hi_inclusive,
3628 snapshot,
3629 );
3630 set
3631 }
3632 Condition::IsNull { column_id } => {
3633 let mut set = if self.run_refs.len() == 1 {
3634 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3635 r.null_row_id_set(*column_id, true)?
3636 } else {
3637 return self.null_scan(*column_id, true, snapshot);
3638 };
3639 set.remove_many(self.overlay_rid_set(snapshot));
3640 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3641 set
3642 }
3643 Condition::IsNotNull { column_id } => {
3644 let mut set = if self.run_refs.len() == 1 {
3645 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3646 r.null_row_id_set(*column_id, false)?
3647 } else {
3648 return self.null_scan(*column_id, false, snapshot);
3649 };
3650 set.remove_many(self.overlay_rid_set(snapshot));
3651 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3652 set
3653 }
3654 })
3655 }
3656
3657 fn range_scan_i64(
3665 &self,
3666 column_id: u16,
3667 lo: i64,
3668 hi: i64,
3669 snapshot: Snapshot,
3670 ) -> Result<RowIdSet> {
3671 let mut row_ids = Vec::new();
3672 let overlay_rids = self.overlay_rid_set(snapshot);
3673 for rr in &self.run_refs {
3674 let mut reader = self.open_reader(rr.run_id)?;
3675 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3676 for rid in matched {
3677 if !overlay_rids.contains(&rid) {
3678 row_ids.push(rid);
3679 }
3680 }
3681 }
3682 let mut s = RowIdSet::from_unsorted(row_ids);
3683 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3684 Ok(s)
3685 }
3686
3687 fn range_scan_f64(
3690 &self,
3691 column_id: u16,
3692 lo: f64,
3693 lo_inclusive: bool,
3694 hi: f64,
3695 hi_inclusive: bool,
3696 snapshot: Snapshot,
3697 ) -> Result<RowIdSet> {
3698 let mut row_ids = Vec::new();
3699 let overlay_rids = self.overlay_rid_set(snapshot);
3700 for rr in &self.run_refs {
3701 let mut reader = self.open_reader(rr.run_id)?;
3702 let matched = reader.range_row_ids_visible_f64(
3703 column_id,
3704 lo,
3705 lo_inclusive,
3706 hi,
3707 hi_inclusive,
3708 snapshot.epoch,
3709 )?;
3710 for rid in matched {
3711 if !overlay_rids.contains(&rid) {
3712 row_ids.push(rid);
3713 }
3714 }
3715 }
3716 let mut s = RowIdSet::from_unsorted(row_ids);
3717 self.range_scan_overlay_f64(
3718 &mut s,
3719 column_id,
3720 lo,
3721 lo_inclusive,
3722 hi,
3723 hi_inclusive,
3724 snapshot,
3725 );
3726 Ok(s)
3727 }
3728
3729 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3731 let mut s = HashSet::new();
3732 for row in self.memtable.visible_versions(snapshot.epoch) {
3733 s.insert(row.row_id.0);
3734 }
3735 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3736 s.insert(row.row_id.0);
3737 }
3738 s
3739 }
3740
3741 fn range_scan_overlay_i64(
3742 &self,
3743 s: &mut RowIdSet,
3744 column_id: u16,
3745 lo: i64,
3746 hi: i64,
3747 snapshot: Snapshot,
3748 ) {
3749 let mut newest: HashMap<u64, &Row> = HashMap::new();
3754 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3755 let memtable = self.memtable.visible_versions(snapshot.epoch);
3756 for r in &mutable {
3757 newest.entry(r.row_id.0).or_insert(r);
3758 }
3759 for r in &memtable {
3760 newest.insert(r.row_id.0, r);
3761 }
3762 for row in newest.values() {
3763 if !row.deleted {
3764 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3765 if *v >= lo && *v <= hi {
3766 s.insert(row.row_id.0);
3767 }
3768 }
3769 }
3770 }
3771 }
3772
3773 #[allow(clippy::too_many_arguments)]
3774 fn range_scan_overlay_f64(
3775 &self,
3776 s: &mut RowIdSet,
3777 column_id: u16,
3778 lo: f64,
3779 lo_inclusive: bool,
3780 hi: f64,
3781 hi_inclusive: bool,
3782 snapshot: Snapshot,
3783 ) {
3784 let mut newest: HashMap<u64, &Row> = HashMap::new();
3787 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3788 let memtable = self.memtable.visible_versions(snapshot.epoch);
3789 for r in &mutable {
3790 newest.entry(r.row_id.0).or_insert(r);
3791 }
3792 for r in &memtable {
3793 newest.insert(r.row_id.0, r);
3794 }
3795 for row in newest.values() {
3796 if !row.deleted {
3797 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3798 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3799 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3800 if ok_lo && ok_hi {
3801 s.insert(row.row_id.0);
3802 }
3803 }
3804 }
3805 }
3806 }
3807
3808 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3811 let mut row_ids = Vec::new();
3812 let overlay_rids = self.overlay_rid_set(snapshot);
3813 for rr in &self.run_refs {
3814 let mut reader = self.open_reader(rr.run_id)?;
3815 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3816 for rid in matched {
3817 if !overlay_rids.contains(&rid) {
3818 row_ids.push(rid);
3819 }
3820 }
3821 }
3822 let mut s = RowIdSet::from_unsorted(row_ids);
3823 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3824 Ok(s)
3825 }
3826
3827 fn null_scan_overlay(
3831 &self,
3832 s: &mut RowIdSet,
3833 column_id: u16,
3834 want_nulls: bool,
3835 snapshot: Snapshot,
3836 ) {
3837 let mut newest: HashMap<u64, &Row> = HashMap::new();
3838 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3839 let memtable = self.memtable.visible_versions(snapshot.epoch);
3840 for r in &mutable {
3841 newest.entry(r.row_id.0).or_insert(r);
3842 }
3843 for r in &memtable {
3844 newest.insert(r.row_id.0, r);
3845 }
3846 for row in newest.values() {
3847 if row.deleted {
3848 continue;
3849 }
3850 let is_null = !row.columns.contains_key(&column_id)
3851 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3852 if is_null == want_nulls {
3853 s.insert(row.row_id.0);
3854 }
3855 }
3856 }
3857
3858 pub fn snapshot(&self) -> Snapshot {
3859 Snapshot::at(self.epoch.visible())
3860 }
3861
3862 pub fn pin_snapshot(&mut self) -> Snapshot {
3865 let e = self.epoch.visible();
3866 *self.pinned.entry(e).or_insert(0) += 1;
3867 Snapshot::at(e)
3868 }
3869
3870 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3872 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3873 *count -= 1;
3874 if *count == 0 {
3875 self.pinned.remove(&snap.epoch);
3876 }
3877 }
3878 }
3879
3880 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3890 let local = self.pinned.keys().next().copied();
3891 let global = self.snapshots.min_pinned();
3892 match (local, global) {
3893 (Some(a), Some(b)) => Some(a.min(b)),
3894 (Some(a), None) => Some(a),
3895 (None, b) => b,
3896 }
3897 }
3898
3899 pub fn current_epoch(&self) -> Epoch {
3900 self.epoch.visible()
3901 }
3902
3903 pub fn memtable_len(&self) -> usize {
3904 self.memtable.len()
3905 }
3906
3907 pub fn count(&self) -> u64 {
3910 self.live_count
3911 }
3912
3913 pub fn count_conditions(
3917 &mut self,
3918 conditions: &[crate::query::Condition],
3919 snapshot: Snapshot,
3920 ) -> Result<Option<u64>> {
3921 use crate::query::Condition;
3922 if conditions.is_empty() {
3923 return Ok(Some(self.live_count));
3924 }
3925 let served = |c: &Condition| {
3926 matches!(
3927 c,
3928 Condition::Pk(_)
3929 | Condition::BitmapEq { .. }
3930 | Condition::BitmapIn { .. }
3931 | Condition::BytesPrefix { .. }
3932 | Condition::FmContains { .. }
3933 | Condition::FmContainsAll { .. }
3934 | Condition::Ann { .. }
3935 | Condition::Range { .. }
3936 | Condition::RangeF64 { .. }
3937 | Condition::SparseMatch { .. }
3938 | Condition::MinHashSimilar { .. }
3939 | Condition::IsNull { .. }
3940 | Condition::IsNotNull { .. }
3941 )
3942 };
3943 if !conditions.iter().all(served) {
3944 return Ok(None);
3945 }
3946 self.ensure_indexes_complete()?;
3947 let mut sets = Vec::with_capacity(conditions.len());
3948 for condition in conditions {
3949 sets.push(self.resolve_condition(condition, snapshot)?);
3950 }
3951 let mut rids = RowIdSet::intersect_many(sets);
3952 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
3962 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
3963 }
3964 let count = rids.len() as u64;
3965 crate::trace::QueryTrace::record(|t| {
3966 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
3967 t.survivor_count = Some(count as usize);
3968 t.conditions_pushed = conditions.len();
3969 });
3970 Ok(Some(count))
3971 }
3972
3973 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
3978 let mut out = Vec::new();
3979 for row in self.memtable.visible_versions(snapshot.epoch) {
3980 if row.deleted {
3981 out.push(row.row_id.0);
3982 }
3983 }
3984 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3985 if row.deleted {
3986 out.push(row.row_id.0);
3987 }
3988 }
3989 out
3990 }
3991
3992 pub fn bulk_load_columns(
4001 &mut self,
4002 user_columns: Vec<(u16, columnar::NativeColumn)>,
4003 ) -> Result<Epoch> {
4004 self.bulk_load_columns_with(user_columns, 3, false, true)
4005 }
4006
4007 pub fn bulk_load_fast(
4014 &mut self,
4015 user_columns: Vec<(u16, columnar::NativeColumn)>,
4016 ) -> Result<Epoch> {
4017 self.bulk_load_columns_with(user_columns, -1, true, false)
4018 }
4019
4020 fn bulk_load_columns_with(
4021 &mut self,
4022 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4023 zstd_level: i32,
4024 force_plain: bool,
4025 lz4: bool,
4026 ) -> Result<Epoch> {
4027 let epoch = self.commit()?;
4028 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4029 if n == 0 {
4030 return Ok(epoch);
4031 }
4032 let live_before = self.live_count;
4033 self.spill_mutable_run(epoch)?;
4035 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4036 && self.indexes_complete
4037 && self.run_refs.is_empty()
4038 && self.memtable.is_empty()
4039 && self.mutable_run.is_empty();
4040 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4043 self.validate_columns_not_null(&user_columns, n)?;
4044 let winner_idx = self
4045 .bulk_pk_winner_indices(&user_columns, n)
4046 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
4047 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4048 match winner_idx.as_deref() {
4049 Some(idx) => {
4050 let compacted = user_columns
4051 .iter()
4052 .map(|(id, c)| (*id, c.gather(idx)))
4053 .collect();
4054 (compacted, idx.len())
4055 }
4056 None => (user_columns, n),
4057 };
4058 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4059 let first = self.allocator.alloc_range(write_n as u64).0;
4060 for rid in first..first + write_n as u64 {
4061 self.reservoir.offer(rid);
4062 }
4063 let run_id = self.next_run_id;
4064 self.next_run_id += 1;
4065 let path = self.run_path(run_id);
4066 let mut writer =
4067 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4068 if force_plain {
4069 writer = writer.with_plain();
4070 } else if lz4 {
4071 writer = writer.with_lz4();
4074 } else {
4075 writer = writer.with_zstd_level(zstd_level);
4076 }
4077 if let Some(kek) = &self.kek {
4078 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4079 }
4080 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4081 self.run_refs.push(RunRef {
4082 run_id: run_id as u128,
4083 level: 0,
4084 epoch_created: epoch.0,
4085 row_count: header.row_count,
4086 });
4087 self.live_count = self.live_count.saturating_add(write_n as u64);
4088 if eager_index_build {
4089 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4090 self.index_columns_bulk(&write_columns, &row_ids);
4091 self.indexes_complete = true;
4092 self.build_learned_ranges()?;
4093 } else {
4094 self.indexes_complete = false;
4098 }
4099 self.mark_flushed(epoch)?;
4100 self.persist_manifest(epoch)?;
4101 if eager_index_build {
4102 self.checkpoint_indexes(epoch);
4103 }
4104 self.clear_result_cache();
4105 Ok(epoch)
4106 }
4107
4108 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4126 let n = row_ids.len();
4127 if n == 0 {
4128 return;
4129 }
4130 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4131 columns.iter().map(|(id, c)| (*id, c)).collect();
4132 let ty_of: std::collections::HashMap<u16, TypeId> =
4133 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
4134 let pk_id = self.schema.primary_key().map(|c| c.id);
4135
4136 for (i, &rid) in row_ids.iter().enumerate() {
4137 let row_id = RowId(rid);
4138 if let Some(pid) = pk_id {
4139 if let Some(col) = by_id.get(&pid) {
4140 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
4141 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4142 self.insert_hot_pk(key, row_id);
4143 }
4144 }
4145 }
4146 for idef in &self.schema.indexes {
4147 let Some(col) = by_id.get(&idef.column_id) else {
4148 continue;
4149 };
4150 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
4151 match idef.kind {
4152 IndexKind::Bitmap => {
4153 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4154 if let Some(key) =
4155 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4156 {
4157 b.insert(key, row_id);
4158 }
4159 }
4160 }
4161 IndexKind::FmIndex => {
4162 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4163 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4164 f.insert(bytes.to_vec(), row_id);
4165 }
4166 }
4167 }
4168 IndexKind::Sparse => {
4169 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4170 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4171 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4172 s.insert(&terms, row_id);
4173 }
4174 }
4175 }
4176 }
4177 IndexKind::MinHash => {
4178 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4179 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4180 let tokens = crate::index::token_hashes_from_bytes(bytes);
4181 mh.insert(&tokens, row_id);
4182 }
4183 }
4184 }
4185 _ => {}
4186 }
4187 }
4188 }
4189 }
4190
4191 pub fn visible_columns_native(
4196 &self,
4197 snapshot: Snapshot,
4198 projection: Option<&[u16]>,
4199 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4200 let wanted: Vec<u16> = match projection {
4201 Some(p) => p.to_vec(),
4202 None => self.schema.columns.iter().map(|c| c.id).collect(),
4203 };
4204 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4205 let rr = self.run_refs[0].clone();
4206 let mut reader = self.open_reader(rr.run_id)?;
4207 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4208 let all_visible = idxs.len() == reader.row_count();
4209 if reader.has_mmap() {
4215 use rayon::prelude::*;
4216 let valid: Vec<u16> = wanted
4219 .iter()
4220 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4221 .copied()
4222 .collect();
4223 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4225 .par_iter()
4226 .filter_map(|cid| {
4227 reader
4228 .column_native_shared(*cid)
4229 .ok()
4230 .map(|col| (*cid, col))
4231 })
4232 .collect();
4233 let cols = decoded
4234 .into_iter()
4235 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4236 .collect();
4237 return Ok(cols);
4238 }
4239 let mut cols = Vec::with_capacity(wanted.len());
4240 for cid in &wanted {
4241 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4242 Some(c) => c,
4243 None => continue,
4244 };
4245 let col = reader.column_native(cdef.id)?;
4246 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4247 }
4248 return Ok(cols);
4249 }
4250 let vcols = self.visible_columns(snapshot)?;
4251 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4252 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4253 .into_iter()
4254 .filter(|(id, _)| want_set.contains(id))
4255 .map(|(id, vals)| {
4256 let ty = self
4257 .schema
4258 .columns
4259 .iter()
4260 .find(|c| c.id == id)
4261 .map(|c| c.ty)
4262 .unwrap_or(TypeId::Bytes);
4263 (id, columnar::values_to_native(ty, &vals))
4264 })
4265 .collect();
4266 Ok(out)
4267 }
4268
4269 pub fn run_count(&self) -> usize {
4270 self.run_refs.len()
4271 }
4272
4273 pub fn memtable_is_empty(&self) -> bool {
4275 self.memtable.is_empty()
4276 }
4277
4278 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4282 self.page_cache.stats()
4283 }
4284
4285 pub fn reset_page_cache_stats(&self) {
4287 self.page_cache.reset_stats();
4288 }
4289
4290 pub fn run_ids(&self) -> Vec<u128> {
4293 self.run_refs.iter().map(|r| r.run_id).collect()
4294 }
4295
4296 pub fn single_run_is_clean(&self) -> bool {
4300 if self.run_refs.len() != 1 {
4301 return false;
4302 }
4303 self.open_reader(self.run_refs[0].run_id)
4304 .map(|r| r.is_clean())
4305 .unwrap_or(false)
4306 }
4307
4308 fn resolve_footprint(
4315 &self,
4316 conditions: &[crate::query::Condition],
4317 _snapshot: Snapshot,
4318 ) -> roaring::RoaringBitmap {
4319 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4320 return roaring::RoaringBitmap::new();
4321 }
4322 if self.run_refs.is_empty() {
4323 return roaring::RoaringBitmap::new();
4324 }
4325 if self.run_refs.len() == 1 {
4327 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4328 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4329 return rids.to_roaring_lossy();
4330 }
4331 }
4332 }
4333 roaring::RoaringBitmap::new()
4334 }
4335
4336 pub fn query_columns_native_cached(
4347 &mut self,
4348 conditions: &[crate::query::Condition],
4349 projection: Option<&[u16]>,
4350 snapshot: Snapshot,
4351 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4352 if conditions.is_empty() {
4353 return self.query_columns_native(conditions, projection, snapshot);
4354 }
4355 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4359 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4360 crate::trace::QueryTrace::record(|t| {
4361 t.result_cache_hit = true;
4362 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4363 });
4364 return Ok(Some((*hit).clone()));
4365 }
4366 let res = self.query_columns_native(conditions, projection, snapshot)?;
4367 if let Some(cols) = &res {
4368 let footprint = self.resolve_footprint(conditions, snapshot);
4369 let condition_cols = crate::query::condition_columns(conditions);
4370 self.result_cache.lock().insert(
4371 key,
4372 CachedEntry {
4373 data: CachedData::Columns(Arc::new(cols.clone())),
4374 footprint,
4375 condition_cols,
4376 },
4377 );
4378 }
4379 Ok(res)
4380 }
4381
4382 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4387 if q.conditions.is_empty() {
4388 return self.query(q);
4389 }
4390 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4391 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4392 crate::trace::QueryTrace::record(|t| {
4393 t.result_cache_hit = true;
4394 t.scan_mode = crate::trace::ScanMode::Materialized;
4395 });
4396 return Ok((*hit).clone());
4397 }
4398 let rows = self.query(q)?;
4399 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4400 let condition_cols = crate::query::condition_columns(&q.conditions);
4401 self.result_cache.lock().insert(
4402 key,
4403 CachedEntry {
4404 data: CachedData::Rows(Arc::new(rows.clone())),
4405 footprint,
4406 condition_cols,
4407 },
4408 );
4409 Ok(rows)
4410 }
4411
4412 #[allow(clippy::type_complexity)]
4427 pub fn query_columns_native_traced(
4428 &mut self,
4429 conditions: &[crate::query::Condition],
4430 projection: Option<&[u16]>,
4431 snapshot: Snapshot,
4432 ) -> Result<(
4433 Option<Vec<(u16, columnar::NativeColumn)>>,
4434 crate::trace::QueryTrace,
4435 )> {
4436 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4437 self.query_columns_native(conditions, projection, snapshot)
4438 });
4439 Ok((result?, trace))
4440 }
4441
4442 #[allow(clippy::type_complexity)]
4445 pub fn query_columns_native_cached_traced(
4446 &mut self,
4447 conditions: &[crate::query::Condition],
4448 projection: Option<&[u16]>,
4449 snapshot: Snapshot,
4450 ) -> Result<(
4451 Option<Vec<(u16, columnar::NativeColumn)>>,
4452 crate::trace::QueryTrace,
4453 )> {
4454 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4455 self.query_columns_native_cached(conditions, projection, snapshot)
4456 });
4457 Ok((result?, trace))
4458 }
4459
4460 pub fn native_page_cursor_traced(
4462 &self,
4463 snapshot: Snapshot,
4464 projection: Vec<(u16, TypeId)>,
4465 conditions: &[crate::query::Condition],
4466 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4467 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4468 self.native_page_cursor(snapshot, projection, conditions)
4469 });
4470 Ok((result?, trace))
4471 }
4472
4473 pub fn native_multi_run_cursor_traced(
4475 &self,
4476 snapshot: Snapshot,
4477 projection: Vec<(u16, TypeId)>,
4478 conditions: &[crate::query::Condition],
4479 ) -> Result<(
4480 Option<crate::cursor::MultiRunCursor>,
4481 crate::trace::QueryTrace,
4482 )> {
4483 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4484 self.native_multi_run_cursor(snapshot, projection, conditions)
4485 });
4486 Ok((result?, trace))
4487 }
4488
4489 pub fn count_conditions_traced(
4491 &mut self,
4492 conditions: &[crate::query::Condition],
4493 snapshot: Snapshot,
4494 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4495 let (result, trace) =
4496 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4497 Ok((result?, trace))
4498 }
4499
4500 pub fn query_traced(
4502 &mut self,
4503 q: &crate::query::Query,
4504 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4505 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4506 Ok((result?, trace))
4507 }
4508
4509 pub fn query_columns_native(
4514 &mut self,
4515 conditions: &[crate::query::Condition],
4516 projection: Option<&[u16]>,
4517 snapshot: Snapshot,
4518 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4519 use crate::query::Condition;
4520 if conditions.is_empty() {
4521 return Ok(None);
4522 }
4523 self.ensure_indexes_complete()?;
4524
4525 let served = |c: &Condition| {
4530 matches!(
4531 c,
4532 Condition::Pk(_)
4533 | Condition::BitmapEq { .. }
4534 | Condition::BitmapIn { .. }
4535 | Condition::BytesPrefix { .. }
4536 | Condition::FmContains { .. }
4537 | Condition::FmContainsAll { .. }
4538 | Condition::Ann { .. }
4539 | Condition::Range { .. }
4540 | Condition::RangeF64 { .. }
4541 | Condition::SparseMatch { .. }
4542 | Condition::MinHashSimilar { .. }
4543 | Condition::IsNull { .. }
4544 | Condition::IsNotNull { .. }
4545 )
4546 };
4547 if !conditions.iter().all(served) {
4548 return Ok(None);
4549 }
4550 let fast_path =
4551 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4552 crate::trace::QueryTrace::record(|t| {
4553 t.run_count = self.run_refs.len();
4554 t.memtable_rows = self.memtable.len();
4555 t.mutable_run_rows = self.mutable_run.len();
4556 t.conditions_pushed = conditions.len();
4557 t.learned_range_used = conditions.iter().any(|c| match c {
4558 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4559 self.learned_range.contains_key(column_id)
4560 }
4561 _ => false,
4562 });
4563 });
4564 let col_ids: Vec<u16> = projection
4566 .map(|p| p.to_vec())
4567 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4568 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4569 .iter()
4570 .map(|&cid| {
4571 let ty = self
4572 .schema
4573 .columns
4574 .iter()
4575 .find(|c| c.id == cid)
4576 .map(|c| c.ty)
4577 .unwrap_or(TypeId::Bytes);
4578 (cid, ty)
4579 })
4580 .collect();
4581
4582 if fast_path {
4588 let needs_column = conditions.iter().any(|c| match c {
4591 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4592 Condition::RangeF64 { column_id, .. } => {
4593 !self.learned_range.contains_key(column_id)
4594 }
4595 _ => false,
4596 });
4597 let mut reader_opt: Option<RunReader> = if needs_column {
4598 Some(self.open_reader(self.run_refs[0].run_id)?)
4599 } else {
4600 None
4601 };
4602 let mut sets: Vec<RowIdSet> = Vec::new();
4603 for c in conditions {
4604 let s = match c {
4605 Condition::Range { column_id, lo, hi }
4606 if !self.learned_range.contains_key(column_id) =>
4607 {
4608 if reader_opt.is_none() {
4609 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4610 }
4611 reader_opt
4612 .as_mut()
4613 .expect("reader opened for range")
4614 .range_row_id_set_i64(*column_id, *lo, *hi)?
4615 }
4616 Condition::RangeF64 {
4617 column_id,
4618 lo,
4619 lo_inclusive,
4620 hi,
4621 hi_inclusive,
4622 } if !self.learned_range.contains_key(column_id) => {
4623 if reader_opt.is_none() {
4624 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4625 }
4626 reader_opt
4627 .as_mut()
4628 .expect("reader opened for range")
4629 .range_row_id_set_f64(
4630 *column_id,
4631 *lo,
4632 *lo_inclusive,
4633 *hi,
4634 *hi_inclusive,
4635 )?
4636 }
4637 _ => self.resolve_condition(c, snapshot)?,
4638 };
4639 sets.push(s);
4640 }
4641 let candidates = RowIdSet::intersect_many(sets);
4642 crate::trace::QueryTrace::record(|t| {
4643 t.survivor_count = Some(candidates.len());
4644 });
4645 if candidates.is_empty() {
4646 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4647 .iter()
4648 .map(|&id| {
4649 (
4650 id,
4651 columnar::null_native(
4652 proj_pairs
4653 .iter()
4654 .find(|(c, _)| c == &id)
4655 .map(|(_, t)| *t)
4656 .unwrap_or(TypeId::Bytes),
4657 0,
4658 ),
4659 )
4660 })
4661 .collect();
4662 return Ok(Some(cols));
4663 }
4664 let mut reader = match reader_opt.take() {
4665 Some(r) => r,
4666 None => self.open_reader(self.run_refs[0].run_id)?,
4667 };
4668 let candidate_ids = candidates.into_sorted_vec();
4669 let (positions, fast_rid) = if let Some(positions) =
4670 reader.positions_for_row_ids_fast(&candidate_ids)
4671 {
4672 (positions, true)
4673 } else {
4674 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4675 match col {
4676 columnar::NativeColumn::Int64 { data, .. } => {
4677 let mut p: Vec<usize> = candidate_ids
4678 .iter()
4679 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4680 .collect();
4681 p.sort_unstable();
4682 (p, false)
4683 }
4684 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4685 }
4686 };
4687 crate::trace::QueryTrace::record(|t| {
4688 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4689 t.fast_row_id_map = fast_rid;
4690 });
4691 let mut cols = Vec::with_capacity(col_ids.len());
4692 for cid in &col_ids {
4693 let col = reader.column_native(*cid)?;
4694 cols.push((*cid, col.gather(&positions)));
4695 }
4696 return Ok(Some(cols));
4697 }
4698
4699 if !self.run_refs.is_empty() {
4712 use crate::cursor::{drain_cursor_to_columns, Cursor};
4713 let remaining: usize;
4714 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4715 let c = self
4716 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4717 .expect("single-run cursor should build when run_refs.len() == 1");
4718 remaining = c.remaining_rows();
4719 Box::new(c)
4720 } else {
4721 let c = self
4722 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4723 .expect("multi-run cursor should build when run_refs.len() >= 1");
4724 remaining = c.remaining_rows();
4725 Box::new(c)
4726 };
4727 crate::trace::QueryTrace::record(|t| {
4728 if t.survivor_count.is_none() {
4729 t.survivor_count = Some(remaining);
4730 }
4731 });
4732 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4733 return Ok(Some(cols));
4734 }
4735
4736 crate::trace::QueryTrace::record(|t| {
4741 t.scan_mode = crate::trace::ScanMode::Materialized;
4742 t.row_materialized = true;
4743 });
4744 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4745 for c in conditions {
4746 sets.push(self.resolve_condition(c, snapshot)?);
4747 }
4748 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4749 let rows = self.rows_for_rids(&rids, snapshot)?;
4750 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4751 for (cid, ty) in &proj_pairs {
4752 let vals: Vec<Value> = rows
4753 .iter()
4754 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4755 .collect();
4756 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4757 }
4758 Ok(Some(cols))
4759 }
4760
4761 pub fn native_page_cursor(
4776 &self,
4777 snapshot: Snapshot,
4778 projection: Vec<(u16, TypeId)>,
4779 conditions: &[crate::query::Condition],
4780 ) -> Result<Option<NativePageCursor>> {
4781 use crate::cursor::build_page_plans;
4782 if !conditions.is_empty() && !self.indexes_complete {
4785 return Ok(None);
4786 }
4787 if self.run_refs.len() != 1 {
4788 return Ok(None);
4789 }
4790 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4791 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4792
4793 let overlay_rids: HashSet<u64> = {
4796 let mut s = HashSet::new();
4797 for row in self.memtable.visible_versions(snapshot.epoch) {
4798 s.insert(row.row_id.0);
4799 }
4800 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4801 s.insert(row.row_id.0);
4802 }
4803 s
4804 };
4805
4806 let survivors = if conditions.is_empty() {
4810 None
4811 } else {
4812 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4813 };
4814
4815 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4822 survivors.clone()
4823 } else if let Some(s) = &survivors {
4824 let mut run_set = s.clone();
4825 run_set.remove_many(overlay_rids.iter().copied());
4826 Some(run_set)
4827 } else {
4828 Some(RowIdSet::from_unsorted(
4829 rids.iter()
4830 .map(|&r| r as u64)
4831 .filter(|r| !overlay_rids.contains(r))
4832 .collect(),
4833 ))
4834 };
4835
4836 let overlay_rows = if overlay_rids.is_empty() {
4837 Vec::new()
4838 } else {
4839 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4840 self.overlay_visible_rows(snapshot, bound)
4841 };
4842
4843 let plans = if positions.is_empty() {
4845 Vec::new()
4846 } else {
4847 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4848 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4849 };
4850
4851 let overlay = if overlay_rows.is_empty() {
4853 None
4854 } else {
4855 let filtered =
4856 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4857 if filtered.is_empty() {
4858 None
4859 } else {
4860 Some(self.materialize_overlay(&filtered, &projection))
4861 }
4862 };
4863
4864 let overlay_row_count = overlay
4865 .as_ref()
4866 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4867 .unwrap_or(0);
4868 crate::trace::QueryTrace::record(|t| {
4869 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4870 t.run_count = self.run_refs.len();
4871 t.memtable_rows = self.memtable.len();
4872 t.mutable_run_rows = self.mutable_run.len();
4873 t.overlay_rows = overlay_row_count;
4874 t.conditions_pushed = conditions.len();
4875 t.pages_decoded = plans
4876 .iter()
4877 .map(|p| p.positions.len())
4878 .sum::<usize>()
4879 .min(1);
4880 });
4881
4882 Ok(Some(NativePageCursor::new_with_overlay(
4883 reader, projection, plans, overlay,
4884 )))
4885 }
4886 #[allow(clippy::type_complexity)]
4896 pub fn native_multi_run_cursor(
4897 &self,
4898 snapshot: Snapshot,
4899 projection: Vec<(u16, TypeId)>,
4900 conditions: &[crate::query::Condition],
4901 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4902 use crate::cursor::{MultiRunCursor, RunStream};
4903 use crate::sorted_run::SYS_ROW_ID;
4904 use std::collections::{BinaryHeap, HashMap, HashSet};
4905 if !conditions.is_empty() && !self.indexes_complete {
4908 return Ok(None);
4909 }
4910 if self.run_refs.is_empty() {
4911 return Ok(None);
4912 }
4913
4914 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
4916 Vec::with_capacity(self.run_refs.len());
4917 for rr in &self.run_refs {
4918 let mut reader = self.open_reader(rr.run_id)?;
4919 let (rids, eps, del) = reader.system_columns_native()?;
4920 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
4921 run_meta.push((reader, rids, eps, del, page_rows));
4922 }
4923
4924 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
4928 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
4929 for i in 0..rids.len() {
4930 let rid = rids[i] as u64;
4931 let e = eps[i] as u64;
4932 if e > snapshot.epoch.0 {
4933 continue;
4934 }
4935 let is_del = del[i] != 0;
4936 best.entry(rid)
4937 .and_modify(|cur| {
4938 if e > cur.0 {
4939 *cur = (e, run_idx, i, is_del);
4940 }
4941 })
4942 .or_insert((e, run_idx, i, is_del));
4943 }
4944 }
4945
4946 let overlay_rids: HashSet<u64> = {
4948 let mut s = HashSet::new();
4949 for row in self.memtable.visible_versions(snapshot.epoch) {
4950 s.insert(row.row_id.0);
4951 }
4952 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4953 s.insert(row.row_id.0);
4954 }
4955 s
4956 };
4957
4958 let survivors: Option<RowIdSet> = if conditions.is_empty() {
4960 None
4961 } else {
4962 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4963 for c in conditions {
4964 sets.push(self.resolve_condition(c, snapshot)?);
4965 }
4966 Some(RowIdSet::intersect_many(sets))
4967 };
4968
4969 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
4973 for (rid, (_, run_idx, pos, deleted)) in &best {
4974 if *deleted {
4975 continue;
4976 }
4977 if overlay_rids.contains(rid) {
4978 continue;
4979 }
4980 if let Some(s) = &survivors {
4981 if !s.contains(*rid) {
4982 continue;
4983 }
4984 }
4985 per_run[*run_idx].push((*rid, *pos));
4986 }
4987 for v in per_run.iter_mut() {
4988 v.sort_unstable_by_key(|&(rid, _)| rid);
4989 }
4990
4991 let mut streams = Vec::with_capacity(run_meta.len());
4993 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
4994 let mut total = 0usize;
4995 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
4996 let mut starts = Vec::with_capacity(page_rows.len());
4997 let mut acc = 0usize;
4998 for &r in &page_rows {
4999 starts.push(acc);
5000 acc += r;
5001 }
5002 let mut survivors_vec: Vec<(u64, usize, usize)> =
5003 Vec::with_capacity(per_run[run_idx].len());
5004 for &(rid, pos) in &per_run[run_idx] {
5005 let page_seq = match starts.partition_point(|&s| s <= pos) {
5006 0 => continue,
5007 p => p - 1,
5008 };
5009 let within = pos - starts[page_seq];
5010 survivors_vec.push((rid, page_seq, within));
5011 }
5012 total += survivors_vec.len();
5013 if let Some(&(rid, _, _)) = survivors_vec.first() {
5014 heap.push(std::cmp::Reverse((rid, run_idx)));
5015 }
5016 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5017 }
5018
5019 let overlay_rows = if overlay_rids.is_empty() {
5021 Vec::new()
5022 } else {
5023 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5024 self.overlay_visible_rows(snapshot, bound)
5025 };
5026 let overlay = if overlay_rows.is_empty() {
5027 None
5028 } else {
5029 let filtered =
5030 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5031 if filtered.is_empty() {
5032 None
5033 } else {
5034 Some(self.materialize_overlay(&filtered, &projection))
5035 }
5036 };
5037
5038 let overlay_row_count = overlay
5039 .as_ref()
5040 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5041 .unwrap_or(0);
5042 crate::trace::QueryTrace::record(|t| {
5043 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5044 t.run_count = self.run_refs.len();
5045 t.memtable_rows = self.memtable.len();
5046 t.mutable_run_rows = self.mutable_run.len();
5047 t.overlay_rows = overlay_row_count;
5048 t.conditions_pushed = conditions.len();
5049 t.survivor_count = Some(total);
5050 });
5051
5052 Ok(Some(MultiRunCursor::new(
5053 streams, projection, heap, total, overlay,
5054 )))
5055 }
5056
5057 fn overlay_materialization_bound<'a>(
5069 conditions: &[crate::query::Condition],
5070 survivors: &'a Option<RowIdSet>,
5071 ) -> Option<&'a RowIdSet> {
5072 use crate::query::Condition;
5073 let has_range = conditions
5074 .iter()
5075 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5076 if has_range {
5077 None
5078 } else {
5079 survivors.as_ref()
5080 }
5081 }
5082
5083 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5095 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5096 let mut fold = |row: Row| {
5097 if let Some(b) = bound {
5098 if !b.contains(row.row_id.0) {
5099 return;
5100 }
5101 }
5102 best.entry(row.row_id.0)
5103 .and_modify(|(be, br)| {
5104 if row.committed_epoch > *be {
5105 *be = row.committed_epoch;
5106 *br = row.clone();
5107 }
5108 })
5109 .or_insert_with(|| (row.committed_epoch, row));
5110 };
5111 for row in self.memtable.visible_versions(snapshot.epoch) {
5112 fold(row);
5113 }
5114 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5115 fold(row);
5116 }
5117 let mut out: Vec<Row> = best
5118 .into_values()
5119 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5120 .collect();
5121 out.sort_by_key(|r| r.row_id);
5122 out
5123 }
5124
5125 fn filter_overlay_rows(
5133 &self,
5134 rows: Vec<Row>,
5135 conditions: &[crate::query::Condition],
5136 survivors: Option<&RowIdSet>,
5137 snapshot: Snapshot,
5138 ) -> Result<Vec<Row>> {
5139 if conditions.is_empty() {
5140 return Ok(rows);
5141 }
5142 use crate::query::Condition;
5143 let all_index_served = !conditions
5147 .iter()
5148 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5149 if all_index_served {
5150 return Ok(rows
5151 .into_iter()
5152 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5153 .collect());
5154 }
5155 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5158 for c in conditions {
5159 let s = match c {
5160 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5161 _ => self.resolve_condition(c, snapshot)?,
5162 };
5163 per_cond_sets.push(s);
5164 }
5165 Ok(rows
5166 .into_iter()
5167 .filter(|row| {
5168 conditions.iter().enumerate().all(|(i, c)| match c {
5169 Condition::Range { column_id, lo, hi } => {
5170 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5171 }
5172 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5173 match row.columns.get(column_id) {
5174 Some(Value::Float64(v)) => {
5175 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5176 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5177 lo_ok && hi_ok
5178 }
5179 _ => false,
5180 }
5181 }
5182 _ => per_cond_sets[i].contains(row.row_id.0),
5183 })
5184 })
5185 .collect())
5186 }
5187
5188 fn materialize_overlay(
5191 &self,
5192 rows: &[Row],
5193 projection: &[(u16, TypeId)],
5194 ) -> Vec<columnar::NativeColumn> {
5195 if projection.is_empty() {
5196 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5197 }
5198 let mut cols = Vec::with_capacity(projection.len());
5199 for (cid, ty) in projection {
5200 let vals: Vec<Value> = rows
5201 .iter()
5202 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5203 .collect();
5204 cols.push(columnar::values_to_native(*ty, &vals));
5205 }
5206 cols
5207 }
5208
5209 fn resolve_survivor_rids(
5214 &self,
5215 conditions: &[crate::query::Condition],
5216 reader: &mut RunReader,
5217 ) -> Result<RowIdSet> {
5218 use crate::query::Condition;
5219 let mut sets: Vec<RowIdSet> = Vec::new();
5220 for c in conditions {
5221 let s: RowIdSet = match c {
5222 Condition::Pk(key) => {
5223 let lookup = self
5224 .schema
5225 .primary_key()
5226 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5227 .unwrap_or_else(|| key.clone());
5228 self.hot
5229 .get(&lookup)
5230 .map(|r| RowIdSet::one(r.0))
5231 .unwrap_or_else(RowIdSet::empty)
5232 }
5233 Condition::BitmapEq { column_id, value } => {
5234 let lookup = self.index_lookup_key_bytes(*column_id, value);
5235 self.bitmap
5236 .get(column_id)
5237 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5238 .unwrap_or_else(RowIdSet::empty)
5239 }
5240 Condition::BitmapIn { column_id, values } => {
5241 let bm = self.bitmap.get(column_id);
5242 let mut acc = roaring::RoaringBitmap::new();
5243 if let Some(b) = bm {
5244 for v in values {
5245 let lookup = self.index_lookup_key_bytes(*column_id, v);
5246 acc |= b.get(&lookup);
5247 }
5248 }
5249 RowIdSet::from_roaring(acc)
5250 }
5251 Condition::BytesPrefix { column_id, prefix } => {
5252 if let Some(b) = self.bitmap.get(column_id) {
5253 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5254 let mut acc = roaring::RoaringBitmap::new();
5255 for key in b.keys() {
5256 if key.starts_with(&lookup_prefix) {
5257 acc |= b.get(key);
5258 }
5259 }
5260 RowIdSet::from_roaring(acc)
5261 } else {
5262 RowIdSet::empty()
5263 }
5264 }
5265 Condition::FmContains { column_id, pattern } => self
5266 .fm
5267 .get(column_id)
5268 .map(|f| {
5269 RowIdSet::from_unsorted(
5270 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5271 )
5272 })
5273 .unwrap_or_else(RowIdSet::empty),
5274 Condition::FmContainsAll {
5275 column_id,
5276 patterns,
5277 } => {
5278 if let Some(f) = self.fm.get(column_id) {
5279 let sets: Vec<RowIdSet> = patterns
5280 .iter()
5281 .map(|pat| {
5282 RowIdSet::from_unsorted(
5283 f.locate(pat).into_iter().map(|r| r.0).collect(),
5284 )
5285 })
5286 .collect();
5287 RowIdSet::intersect_many(sets)
5288 } else {
5289 RowIdSet::empty()
5290 }
5291 }
5292 Condition::Ann {
5293 column_id,
5294 query,
5295 k,
5296 } => self
5297 .ann
5298 .get(column_id)
5299 .map(|a| {
5300 RowIdSet::from_unsorted(
5301 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5302 )
5303 })
5304 .unwrap_or_else(RowIdSet::empty),
5305 Condition::SparseMatch {
5306 column_id,
5307 query,
5308 k,
5309 } => self
5310 .sparse
5311 .get(column_id)
5312 .map(|s| {
5313 RowIdSet::from_unsorted(
5314 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5315 )
5316 })
5317 .unwrap_or_else(RowIdSet::empty),
5318 Condition::MinHashSimilar {
5319 column_id,
5320 query,
5321 k,
5322 } => self
5323 .minhash
5324 .get(column_id)
5325 .map(|mh| {
5326 RowIdSet::from_unsorted(
5327 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5328 )
5329 })
5330 .unwrap_or_else(RowIdSet::empty),
5331 Condition::Range { column_id, lo, hi } => {
5332 if let Some(li) = self.learned_range.get(column_id) {
5333 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5334 } else {
5335 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5336 }
5337 }
5338 Condition::RangeF64 {
5339 column_id,
5340 lo,
5341 lo_inclusive,
5342 hi,
5343 hi_inclusive,
5344 } => {
5345 if let Some(li) = self.learned_range.get(column_id) {
5346 RowIdSet::from_unsorted(
5347 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5348 .into_iter()
5349 .collect(),
5350 )
5351 } else {
5352 reader.range_row_id_set_f64(
5353 *column_id,
5354 *lo,
5355 *lo_inclusive,
5356 *hi,
5357 *hi_inclusive,
5358 )?
5359 }
5360 }
5361 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5362 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5363 };
5364 sets.push(s);
5365 }
5366 Ok(RowIdSet::intersect_many(sets))
5367 }
5368
5369 pub fn scan_cursor(
5390 &self,
5391 snapshot: Snapshot,
5392 projection: Vec<(u16, TypeId)>,
5393 conditions: &[crate::query::Condition],
5394 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5395 if !conditions.is_empty() && !self.indexes_complete {
5401 return Ok(None);
5402 }
5403 if self.run_refs.len() == 1 {
5404 Ok(self
5405 .native_page_cursor(snapshot, projection, conditions)?
5406 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5407 } else {
5408 Ok(self
5409 .native_multi_run_cursor(snapshot, projection, conditions)?
5410 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5411 }
5412 }
5413
5414 pub fn aggregate_native(
5428 &self,
5429 snapshot: Snapshot,
5430 column: Option<u16>,
5431 conditions: &[crate::query::Condition],
5432 agg: NativeAgg,
5433 ) -> Result<Option<NativeAggResult>> {
5434 if self.run_refs.len() == 1 && conditions.is_empty() {
5436 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5437 return Ok(Some(res));
5438 }
5439 }
5440 if matches!(agg, NativeAgg::Count) && column.is_none() {
5442 return Ok(self
5443 .scan_cursor(snapshot, Vec::new(), conditions)?
5444 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5445 }
5446 let cid = match column {
5449 Some(c) => c,
5450 None => return Ok(None),
5451 };
5452 let ty = self.column_type(cid);
5453 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5454 return Ok(None);
5455 };
5456 match ty {
5457 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5458 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5459 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5460 }
5461 TypeId::Float64 => {
5462 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5463 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5464 }
5465 _ => Ok(None),
5466 }
5467 }
5468
5469 fn aggregate_from_stats(
5477 &self,
5478 snapshot: Snapshot,
5479 column: Option<u16>,
5480 agg: NativeAgg,
5481 ) -> Result<Option<NativeAggResult>> {
5482 let cid = match (agg, column) {
5483 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5484 _ => return Ok(None), };
5486 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5487 return Ok(None);
5488 };
5489 let Some(cs) = stats.get(&cid) else {
5490 return Ok(None);
5491 };
5492 match agg {
5493 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5495 self.live_count.saturating_sub(cs.null_count),
5496 ))),
5497 NativeAgg::Min | NativeAgg::Max => {
5498 let bound = if agg == NativeAgg::Min {
5499 &cs.min
5500 } else {
5501 &cs.max
5502 };
5503 match bound {
5504 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5505 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5506 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5511 None => Ok(None),
5512 }
5513 }
5514 _ => Ok(None),
5515 }
5516 }
5517
5518 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5527 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5528 return Ok(None);
5529 }
5530 self.ensure_indexes_complete()?;
5533 let reader = self.open_reader(self.run_refs[0].run_id)?;
5534 if self.live_count != reader.row_count() as u64 {
5535 return Ok(None);
5536 }
5537 let Some(bm) = self.bitmap.get(&column_id) else {
5538 return Ok(None); };
5540 let mut distinct = bm.value_count() as u64;
5541 if !bm.get(&Value::Null.encode_key()).is_empty() {
5544 distinct = distinct.saturating_sub(1);
5545 }
5546 Ok(Some(distinct))
5547 }
5548
5549 pub fn aggregate_incremental(
5561 &mut self,
5562 cache_key: u64,
5563 conditions: &[crate::query::Condition],
5564 column: Option<u16>,
5565 agg: NativeAgg,
5566 ) -> Result<IncrementalAggResult> {
5567 let snap = self.snapshot();
5568 let cur_wm = self.allocator.current().0;
5569 let cur_epoch = snap.epoch.0;
5570 let incremental_ok =
5577 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5578
5579 if incremental_ok {
5582 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5583 if cached.epoch == cur_epoch {
5584 return Ok(IncrementalAggResult {
5585 state: cached.state,
5586 incremental: true,
5587 delta_rows: 0,
5588 });
5589 }
5590 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5591 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5592 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5593 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5594 let delta_state = agg_state_from_rows(
5595 &delta_rows,
5596 conditions,
5597 &index_sets,
5598 column,
5599 agg,
5600 &self.schema,
5601 )?;
5602 let merged = cached.state.merge(delta_state);
5603 let delta_n = delta_rids.len() as u64;
5604 self.agg_cache.insert(
5605 cache_key,
5606 CachedAgg {
5607 state: merged.clone(),
5608 watermark: cur_wm,
5609 epoch: cur_epoch,
5610 },
5611 );
5612 return Ok(IncrementalAggResult {
5613 state: merged,
5614 incremental: true,
5615 delta_rows: delta_n,
5616 });
5617 }
5618 }
5619 }
5620
5621 let cursor_ok =
5626 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5627 let state = if cursor_ok && agg != NativeAgg::Avg {
5628 match self.aggregate_native(snap, column, conditions, agg)? {
5629 Some(result) => {
5630 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5631 }
5632 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5633 }
5634 } else {
5635 self.agg_state_full_scan(conditions, column, agg, snap)?
5636 };
5637 if incremental_ok {
5639 self.agg_cache.insert(
5640 cache_key,
5641 CachedAgg {
5642 state: state.clone(),
5643 watermark: cur_wm,
5644 epoch: cur_epoch,
5645 },
5646 );
5647 }
5648 Ok(IncrementalAggResult {
5649 state,
5650 incremental: false,
5651 delta_rows: 0,
5652 })
5653 }
5654
5655 fn agg_state_full_scan(
5658 &self,
5659 conditions: &[crate::query::Condition],
5660 column: Option<u16>,
5661 agg: NativeAgg,
5662 snap: Snapshot,
5663 ) -> Result<AggState> {
5664 let rows = self.visible_rows(snap)?;
5665 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5666 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5667 }
5668
5669 fn resolve_index_conditions(
5672 &self,
5673 conditions: &[crate::query::Condition],
5674 snapshot: Snapshot,
5675 ) -> Result<Vec<RowIdSet>> {
5676 use crate::query::Condition;
5677 let mut sets = Vec::new();
5678 for c in conditions {
5679 if matches!(
5680 c,
5681 Condition::Ann { .. }
5682 | Condition::SparseMatch { .. }
5683 | Condition::MinHashSimilar { .. }
5684 ) {
5685 sets.push(self.resolve_condition(c, snapshot)?);
5686 }
5687 }
5688 Ok(sets)
5689 }
5690
5691 fn column_type(&self, cid: u16) -> TypeId {
5692 self.schema
5693 .columns
5694 .iter()
5695 .find(|c| c.id == cid)
5696 .map(|c| c.ty)
5697 .unwrap_or(TypeId::Bytes)
5698 }
5699
5700 pub fn approx_aggregate(
5709 &mut self,
5710 conditions: &[crate::query::Condition],
5711 column: Option<u16>,
5712 agg: ApproxAgg,
5713 z: f64,
5714 ) -> Result<Option<ApproxResult>> {
5715 use crate::query::Condition;
5716 self.ensure_reservoir_complete()?;
5717 let snapshot = self.snapshot();
5718 let n_pop = self.live_count;
5719 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5720 if sample_rids.is_empty() {
5721 return Ok(None);
5722 }
5723 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5725 let s = live_sample.len();
5726 if s == 0 {
5727 return Ok(None);
5728 }
5729
5730 let mut index_sets: Vec<RowIdSet> = Vec::new();
5733 for c in conditions {
5734 if matches!(
5735 c,
5736 Condition::Ann { .. }
5737 | Condition::SparseMatch { .. }
5738 | Condition::MinHashSimilar { .. }
5739 ) {
5740 index_sets.push(self.resolve_condition(c, snapshot)?);
5741 }
5742 }
5743
5744 let cid = match (agg, column) {
5746 (ApproxAgg::Count, _) => None,
5747 (_, Some(c)) => Some(c),
5748 _ => return Ok(None),
5749 };
5750 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5751 for r in &live_sample {
5752 if !conditions
5754 .iter()
5755 .all(|c| condition_matches_row(c, r, &self.schema))
5756 {
5757 continue;
5758 }
5759 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5761 continue;
5762 }
5763 if let Some(cid) = cid {
5764 if let Some(v) = as_f64(r.columns.get(&cid)) {
5765 passing_vals.push(v);
5766 } } else {
5768 passing_vals.push(0.0); }
5770 }
5771 let m = passing_vals.len();
5772
5773 let (point, half) = match agg {
5774 ApproxAgg::Count => {
5775 let p = m as f64 / s as f64;
5777 let point = n_pop as f64 * p;
5778 let var = if s > 1 {
5779 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5780 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5781 } else {
5782 0.0
5783 };
5784 (point, z * var.sqrt())
5785 }
5786 ApproxAgg::Sum => {
5787 let y: Vec<f64> = live_sample
5789 .iter()
5790 .map(|r| {
5791 let passes_row = conditions
5792 .iter()
5793 .all(|c| condition_matches_row(c, r, &self.schema))
5794 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5795 if passes_row {
5796 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5797 } else {
5798 0.0
5799 }
5800 })
5801 .collect();
5802 let mean_y = y.iter().sum::<f64>() / s as f64;
5803 let point = n_pop as f64 * mean_y;
5804 let var = if s > 1 {
5805 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5806 let var_y = ss / (s - 1) as f64;
5807 n_pop as f64 * n_pop as f64 * var_y / s as f64
5808 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5809 } else {
5810 0.0
5811 };
5812 (point, z * var.sqrt())
5813 }
5814 ApproxAgg::Avg => {
5815 if m == 0 {
5816 return Ok(Some(ApproxResult {
5817 point: 0.0,
5818 ci_low: 0.0,
5819 ci_high: 0.0,
5820 n_population: n_pop,
5821 n_sample_live: s,
5822 n_passing: 0,
5823 }));
5824 }
5825 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5826 let half = if m > 1 {
5827 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5828 let sd = (ss / (m - 1) as f64).sqrt();
5829 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5830 z * sd / (m as f64).sqrt() * fpc.sqrt()
5831 } else {
5832 0.0
5833 };
5834 (mean, half)
5835 }
5836 };
5837
5838 Ok(Some(ApproxResult {
5839 point,
5840 ci_low: point - half,
5841 ci_high: point + half,
5842 n_population: n_pop,
5843 n_sample_live: s,
5844 n_passing: m,
5845 }))
5846 }
5847
5848 pub fn exact_column_stats(
5856 &self,
5857 _snapshot: Snapshot,
5858 projection: &[u16],
5859 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5860 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5861 return Ok(None);
5862 }
5863 let reader = self.open_reader(self.run_refs[0].run_id)?;
5864 if self.live_count != reader.row_count() as u64 {
5865 return Ok(None);
5866 }
5867 let mut out = HashMap::new();
5868 for &cid in projection {
5869 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5870 Some(c) => c,
5871 None => continue,
5872 };
5873 let Some(stats) = reader.column_page_stats(cid) else {
5875 out.insert(
5876 cid,
5877 ColumnStat {
5878 min: None,
5879 max: None,
5880 null_count: self.live_count,
5881 },
5882 );
5883 continue;
5884 };
5885 let stat = match cdef.ty {
5886 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5887 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5888 min: mn.map(Value::Int64),
5889 max: mx.map(Value::Int64),
5890 null_count: n,
5891 })
5892 }
5893 TypeId::Float64 => {
5894 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5895 min: mn.map(Value::Float64),
5896 max: mx.map(Value::Float64),
5897 null_count: n,
5898 })
5899 }
5900 _ => None,
5901 };
5902 if let Some(s) = stat {
5903 out.insert(cid, s);
5904 }
5905 }
5906 Ok(Some(out))
5907 }
5908
5909 pub fn dir(&self) -> &Path {
5910 &self.dir
5911 }
5912
5913 pub fn schema(&self) -> &Schema {
5914 &self.schema
5915 }
5916
5917 pub(crate) fn prepare_alter_column(
5918 &mut self,
5919 column_name: &str,
5920 change: &AlterColumn,
5921 ) -> Result<ColumnDef> {
5922 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
5923 return Err(MongrelError::InvalidArgument(
5924 "ALTER COLUMN requires committing staged writes first".into(),
5925 ));
5926 }
5927 let old = self
5928 .schema
5929 .columns
5930 .iter()
5931 .find(|c| c.name == column_name)
5932 .cloned()
5933 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5934 let mut next = old.clone();
5935
5936 if let Some(name) = &change.name {
5937 let trimmed = name.trim();
5938 if trimmed.is_empty() {
5939 return Err(MongrelError::InvalidArgument(
5940 "ALTER COLUMN name must not be empty".into(),
5941 ));
5942 }
5943 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
5944 return Err(MongrelError::Schema(format!(
5945 "column {trimmed} already exists"
5946 )));
5947 }
5948 next.name = trimmed.to_string();
5949 }
5950
5951 if let Some(ty) = change.ty {
5952 next.ty = ty;
5953 }
5954 if let Some(flags) = change.flags {
5955 validate_alter_column_flags(old.flags, flags)?;
5956 next.flags = flags;
5957 }
5958
5959 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
5960 if old.flags.contains(ColumnFlags::NULLABLE)
5961 && !next.flags.contains(ColumnFlags::NULLABLE)
5962 && self.column_has_nulls(old.id)?
5963 {
5964 return Err(MongrelError::InvalidArgument(format!(
5965 "column '{}' contains NULL values",
5966 old.name
5967 )));
5968 }
5969 Ok(next)
5970 }
5971
5972 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
5973 let idx = self
5974 .schema
5975 .columns
5976 .iter()
5977 .position(|c| c.id == column.id)
5978 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
5979 if self.schema.columns[idx] == column {
5980 return Ok(());
5981 }
5982 self.schema.columns[idx] = column;
5983 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5984 self.schema.validate_auto_increment()?;
5985 self.auto_inc = resolve_auto_inc(&self.schema);
5986 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
5987 write_schema(&self.dir, &self.schema)?;
5988 self.clear_result_cache();
5989 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5990 self.persist_manifest(self.current_epoch())?;
5991 Ok(())
5992 }
5993
5994 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
5995 let column = self.prepare_alter_column(column_name, &change)?;
5996 self.apply_altered_column(column.clone())?;
5997 Ok(column)
5998 }
5999
6000 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
6001 if self.live_count == 0 {
6002 return Ok(false);
6003 }
6004 let snap = self.snapshot();
6005 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6006 Ok(columns
6007 .first()
6008 .map(|(_, col)| col.null_count(col.len()) != 0)
6009 .unwrap_or(true))
6010 }
6011
6012 fn has_stored_versions(&self) -> bool {
6013 !self.memtable.is_empty()
6014 || !self.mutable_run.is_empty()
6015 || self.run_refs.iter().any(|r| r.row_count > 0)
6016 || !self.retiring.is_empty()
6017 }
6018
6019 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
6024 if self.schema.columns.iter().any(|c| c.name == name) {
6025 return Err(MongrelError::Schema(format!(
6026 "column {name} already exists"
6027 )));
6028 }
6029 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
6030 self.schema.columns.push(ColumnDef {
6031 id,
6032 name: name.to_string(),
6033 ty,
6034 flags,
6035 });
6036 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6037 self.schema.validate_auto_increment()?;
6038 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6039 self.auto_inc = resolve_auto_inc(&self.schema);
6040 }
6041 write_schema(&self.dir, &self.schema)?;
6042 self.clear_result_cache();
6043 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6045 self.persist_manifest(self.current_epoch())?;
6046 Ok(id)
6047 }
6048
6049 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6058 let cid = self
6059 .schema
6060 .columns
6061 .iter()
6062 .find(|c| c.name == column_name)
6063 .map(|c| c.id)
6064 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6065 let ty = self
6066 .schema
6067 .columns
6068 .iter()
6069 .find(|c| c.id == cid)
6070 .map(|c| c.ty)
6071 .unwrap_or(TypeId::Int64);
6072 if !matches!(
6073 ty,
6074 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6075 ) {
6076 return Err(MongrelError::Schema(format!(
6077 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6078 )));
6079 }
6080 if self
6081 .schema
6082 .indexes
6083 .iter()
6084 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6085 {
6086 return Ok(()); }
6088 self.schema.indexes.push(IndexDef {
6089 name: format!("{}_learned_range", column_name),
6090 column_id: cid,
6091 kind: IndexKind::LearnedRange,
6092 });
6093 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6094 write_schema(&self.dir, &self.schema)?;
6095 self.build_learned_ranges()?;
6096 Ok(())
6097 }
6098
6099 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6102 self.sync_byte_threshold = threshold;
6103 if let WalSink::Private(w) = &mut self.wal {
6104 w.set_sync_byte_threshold(threshold);
6105 }
6106 }
6107
6108 pub fn page_cache_flush(&self) {
6112 self.page_cache.flush_to_disk();
6113 }
6114
6115 pub fn page_cache_len(&self) -> usize {
6117 self.page_cache.len()
6118 }
6119
6120 pub fn decoded_cache_len(&self) -> usize {
6123 self.decoded_cache.len()
6124 }
6125
6126 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6129 self.memtable.drain_sorted()
6130 }
6131
6132 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6133 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6134 }
6135
6136 pub(crate) fn table_dir(&self) -> &Path {
6137 &self.dir
6138 }
6139
6140 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6141 &self.schema
6142 }
6143
6144 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6145 let id = self.next_run_id;
6146 self.next_run_id += 1;
6147 id
6148 }
6149
6150 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6151 self.run_refs.push(run_ref);
6152 }
6153
6154 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6164 self.retiring.push(crate::manifest::RetiredRun {
6165 run_id,
6166 retire_epoch,
6167 });
6168 }
6169
6170 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6174 if self.retiring.is_empty() {
6175 return Ok(0);
6176 }
6177 let mut reaped = 0;
6178 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6179 for r in std::mem::take(&mut self.retiring) {
6185 if min_active.0 >= r.retire_epoch {
6186 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6187 reaped += 1;
6188 } else {
6189 kept.push(r);
6190 }
6191 }
6192 self.retiring = kept;
6193 if reaped > 0 {
6194 self.persist_manifest(self.current_epoch())?;
6195 }
6196 Ok(reaped)
6197 }
6198
6199 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6200 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6201 return false;
6202 }
6203 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6204 self.run_refs.push(run_ref);
6205 self.indexes_complete = false;
6206 true
6207 }
6208
6209 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6210 self.kek.as_ref()
6211 }
6212
6213 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6214 let mut reader = RunReader::open_with_cache(
6215 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6216 self.schema.clone(),
6217 self.kek.clone(),
6218 Some(self.page_cache.clone()),
6219 Some(self.decoded_cache.clone()),
6220 self.table_id,
6221 Some(&self.verified_runs),
6222 )?;
6223 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6227 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6228 }
6229 Ok(reader)
6230 }
6231
6232 pub(crate) fn run_refs(&self) -> &[RunRef] {
6233 &self.run_refs
6234 }
6235
6236 pub(crate) fn runs_dir(&self) -> PathBuf {
6237 self.dir.join(RUNS_DIR)
6238 }
6239
6240 pub(crate) fn wal_dir(&self) -> PathBuf {
6241 self.dir.join(WAL_DIR)
6242 }
6243
6244 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6245 self.run_refs = refs;
6246 }
6247
6248 pub(crate) fn next_run_id(&self) -> u64 {
6249 self.next_run_id
6250 }
6251
6252 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6253 self.compaction_zstd_level
6254 }
6255
6256 pub(crate) fn bump_next_run_id(&mut self) {
6257 self.next_run_id += 1;
6258 }
6259
6260 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6261 self.kek.clone()
6262 }
6263
6264 #[cfg(feature = "encryption")]
6268 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6269 self.kek.as_ref().map(|k| k.derive_idx_key())
6270 }
6271
6272 #[cfg(not(feature = "encryption"))]
6273 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6274 None
6275 }
6276
6277 #[cfg(feature = "encryption")]
6281 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6282 self.kek.as_ref().map(|k| *k.derive_meta_key())
6283 }
6284
6285 #[cfg(not(feature = "encryption"))]
6286 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6287 None
6288 }
6289
6290 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6293 self.column_keys
6294 .iter()
6295 .map(|(&id, &(_, scheme))| (id, scheme))
6296 .collect()
6297 }
6298
6299 #[cfg(feature = "encryption")]
6304 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6305 self.tokenize_value_enc(column_id, v)
6306 }
6307
6308 #[cfg(feature = "encryption")]
6309 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6310 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6311 let (key, scheme) = self.column_keys.get(&column_id)?;
6312 let token: Vec<u8> = match (*scheme, v) {
6313 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6314 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6315 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6316 _ => hmac_token(key, &v.encode_key()).to_vec(),
6317 };
6318 Some(Value::Bytes(token))
6319 }
6320
6321 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6323 self.index_lookup_key_bytes(column_id, &v.encode_key())
6324 }
6325
6326 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6329 #[cfg(feature = "encryption")]
6330 {
6331 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6332 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6333 if *scheme == SCHEME_HMAC_EQ {
6334 return hmac_token(key, encoded).to_vec();
6335 }
6336 }
6337 }
6338 let _ = column_id;
6339 encoded.to_vec()
6340 }
6341}
6342
6343fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6344 let columnar::NativeColumn::Int64 { data, validity } = col else {
6345 return false;
6346 };
6347 if data.len() < n || !columnar::all_non_null(validity, n) {
6348 return false;
6349 }
6350 data.iter()
6351 .take(n)
6352 .zip(data.iter().skip(1))
6353 .all(|(a, b)| a < b)
6354}
6355
6356#[derive(Debug, Clone)]
6360pub struct ColumnStat {
6361 pub min: Option<Value>,
6362 pub max: Option<Value>,
6363 pub null_count: u64,
6364}
6365
6366#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6368pub enum NativeAgg {
6369 Count,
6370 Sum,
6371 Min,
6372 Max,
6373 Avg,
6374}
6375
6376#[derive(Debug, Clone, PartialEq)]
6378pub enum NativeAggResult {
6379 Count(u64),
6380 Int(i64),
6381 Float(f64),
6382 Null,
6384}
6385
6386#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6388pub enum ApproxAgg {
6389 Count,
6390 Sum,
6391 Avg,
6392}
6393
6394#[derive(Debug, Clone)]
6398pub struct ApproxResult {
6399 pub point: f64,
6401 pub ci_low: f64,
6403 pub ci_high: f64,
6405 pub n_population: u64,
6407 pub n_sample_live: usize,
6409 pub n_passing: usize,
6411}
6412
6413#[derive(Debug, Clone, PartialEq)]
6418pub enum AggState {
6419 Count(u64),
6421 SumI {
6423 sum: i128,
6424 count: u64,
6425 },
6426 SumF {
6428 sum: f64,
6429 count: u64,
6430 },
6431 AvgI {
6433 sum: i128,
6434 count: u64,
6435 },
6436 AvgF {
6438 sum: f64,
6439 count: u64,
6440 },
6441 MinI(i64),
6443 MaxI(i64),
6444 MinF(f64),
6446 MaxF(f64),
6447 Empty,
6449}
6450
6451impl AggState {
6452 pub fn merge(self, other: AggState) -> AggState {
6454 use AggState::*;
6455 match (self, other) {
6456 (Empty, x) | (x, Empty) => x,
6457 (Count(a), Count(b)) => Count(a + b),
6458 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6459 sum: sa + sb,
6460 count: ca + cb,
6461 },
6462 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6463 sum: sa + sb,
6464 count: ca + cb,
6465 },
6466 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6467 sum: sa + sb,
6468 count: ca + cb,
6469 },
6470 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6471 sum: sa + sb,
6472 count: ca + cb,
6473 },
6474 (MinI(a), MinI(b)) => MinI(a.min(b)),
6475 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6476 (MinF(a), MinF(b)) => MinF(a.min(b)),
6477 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6478 _ => Empty, }
6480 }
6481
6482 pub fn point(&self) -> Option<f64> {
6484 match self {
6485 AggState::Count(n) => Some(*n as f64),
6486 AggState::SumI { sum, .. } => Some(*sum as f64),
6487 AggState::SumF { sum, .. } => Some(*sum),
6488 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6489 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6490 AggState::MinI(n) => Some(*n as f64),
6491 AggState::MaxI(n) => Some(*n as f64),
6492 AggState::MinF(n) => Some(*n),
6493 AggState::MaxF(n) => Some(*n),
6494 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6495 }
6496 }
6497
6498 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6502 let is_float = matches!(ty, Some(TypeId::Float64));
6503 match (agg, result) {
6504 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6505 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6506 sum: x as i128,
6507 count: 1, },
6509 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6510 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6511 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6512 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6513 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6514 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6515 (NativeAgg::Count, _) => AggState::Empty,
6516 (_, NativeAggResult::Null) => AggState::Empty,
6517 _ => {
6518 let _ = is_float;
6519 AggState::Empty
6520 }
6521 }
6522 }
6523}
6524
6525#[derive(Debug, Clone)]
6528pub struct CachedAgg {
6529 pub state: AggState,
6530 pub watermark: u64,
6531 pub epoch: u64,
6532}
6533
6534#[derive(Debug, Clone)]
6536pub struct IncrementalAggResult {
6537 pub state: AggState,
6539 pub incremental: bool,
6542 pub delta_rows: u64,
6544}
6545
6546fn agg_state_from_rows(
6550 rows: &[Row],
6551 conditions: &[crate::query::Condition],
6552 index_sets: &[RowIdSet],
6553 column: Option<u16>,
6554 agg: NativeAgg,
6555 schema: &Schema,
6556) -> Result<AggState> {
6557 let mut count: u64 = 0;
6558 let mut sum_i: i128 = 0;
6559 let mut sum_f: f64 = 0.0;
6560 let mut mn_i: i64 = i64::MAX;
6561 let mut mx_i: i64 = i64::MIN;
6562 let mut mn_f: f64 = f64::INFINITY;
6563 let mut mx_f: f64 = f64::NEG_INFINITY;
6564 let mut saw_int = false;
6565 let mut saw_float = false;
6566 for r in rows {
6567 if !conditions
6568 .iter()
6569 .all(|c| condition_matches_row(c, r, schema))
6570 {
6571 continue;
6572 }
6573 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6574 continue;
6575 }
6576 match agg {
6577 NativeAgg::Count => match column {
6578 None => count += 1,
6580 Some(cid) => match r.columns.get(&cid) {
6583 None | Some(Value::Null) => {}
6584 Some(_) => count += 1,
6585 },
6586 },
6587 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6588 Some(Value::Int64(n)) => {
6589 count += 1;
6590 sum_i += *n as i128;
6591 mn_i = mn_i.min(*n);
6592 mx_i = mx_i.max(*n);
6593 saw_int = true;
6594 }
6595 Some(Value::Float64(f)) => {
6596 count += 1;
6597 sum_f += f;
6598 mn_f = mn_f.min(*f);
6599 mx_f = mx_f.max(*f);
6600 saw_float = true;
6601 }
6602 _ => {}
6603 },
6604 }
6605 }
6606 Ok(match agg {
6607 NativeAgg::Count => {
6608 if count == 0 {
6609 AggState::Empty
6610 } else {
6611 AggState::Count(count)
6612 }
6613 }
6614 NativeAgg::Sum => {
6615 if count == 0 {
6616 AggState::Empty
6617 } else if saw_int {
6618 AggState::SumI { sum: sum_i, count }
6619 } else {
6620 AggState::SumF { sum: sum_f, count }
6621 }
6622 }
6623 NativeAgg::Avg => {
6624 if count == 0 {
6625 AggState::Empty
6626 } else if saw_int {
6627 AggState::AvgI { sum: sum_i, count }
6628 } else {
6629 AggState::AvgF { sum: sum_f, count }
6630 }
6631 }
6632 NativeAgg::Min => {
6633 if !saw_int && !saw_float {
6634 AggState::Empty
6635 } else if saw_int {
6636 AggState::MinI(mn_i)
6637 } else {
6638 AggState::MinF(mn_f)
6639 }
6640 }
6641 NativeAgg::Max => {
6642 if !saw_int && !saw_float {
6643 AggState::Empty
6644 } else if saw_int {
6645 AggState::MaxI(mx_i)
6646 } else {
6647 AggState::MaxF(mx_f)
6648 }
6649 }
6650 })
6651}
6652
6653fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6657 use crate::query::Condition;
6658 match c {
6659 Condition::Pk(key) => match schema.primary_key() {
6660 Some(pk) => row
6661 .columns
6662 .get(&pk.id)
6663 .map(|v| v.encode_key() == *key)
6664 .unwrap_or(false),
6665 None => false,
6666 },
6667 Condition::BitmapEq { column_id, value } => row
6668 .columns
6669 .get(column_id)
6670 .map(|v| v.encode_key() == *value)
6671 .unwrap_or(false),
6672 Condition::BitmapIn { column_id, values } => {
6673 let key = row.columns.get(column_id).map(|v| v.encode_key());
6674 match key {
6675 Some(k) => values.contains(&k),
6676 None => false,
6677 }
6678 }
6679 Condition::BytesPrefix { column_id, prefix } => row
6680 .columns
6681 .get(column_id)
6682 .map(|v| v.encode_key().starts_with(prefix))
6683 .unwrap_or(false),
6684 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6685 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6686 _ => false,
6687 },
6688 Condition::RangeF64 {
6689 column_id,
6690 lo,
6691 lo_inclusive,
6692 hi,
6693 hi_inclusive,
6694 } => match row.columns.get(column_id) {
6695 Some(Value::Float64(n)) => {
6696 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6697 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6698 lo_ok && hi_ok
6699 }
6700 _ => false,
6701 },
6702 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6703 Some(Value::Bytes(b)) => {
6704 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6705 }
6706 _ => false,
6707 },
6708 Condition::FmContainsAll {
6709 column_id,
6710 patterns,
6711 } => match row.columns.get(column_id) {
6712 Some(Value::Bytes(b)) => patterns
6713 .iter()
6714 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6715 _ => false,
6716 },
6717 Condition::Ann { .. }
6718 | Condition::SparseMatch { .. }
6719 | Condition::MinHashSimilar { .. } => true,
6720 Condition::IsNull { column_id } => {
6721 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6722 }
6723 Condition::IsNotNull { column_id } => {
6724 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6725 }
6726 }
6727}
6728
6729fn as_f64(v: Option<&Value>) -> Option<f64> {
6731 match v {
6732 Some(Value::Int64(n)) => Some(*n as f64),
6733 Some(Value::Float64(f)) => Some(*f),
6734 _ => None,
6735 }
6736}
6737
6738fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6742 let mut count: u64 = 0;
6743 let mut sum: i128 = 0;
6744 let mut mn: i64 = i64::MAX;
6745 let mut mx: i64 = i64::MIN;
6746 while let Some(cols) = cursor.next_batch()? {
6747 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6748 if crate::columnar::all_non_null(validity, data.len()) {
6749 count += data.len() as u64;
6751 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6752 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6753 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6754 } else {
6755 for (i, &v) in data.iter().enumerate() {
6756 if crate::columnar::validity_bit(validity, i) {
6757 count += 1;
6758 sum += v as i128;
6759 mn = mn.min(v);
6760 mx = mx.max(v);
6761 }
6762 }
6763 }
6764 }
6765 }
6766 Ok((count, sum, mn, mx))
6767}
6768
6769fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6771 let mut count: u64 = 0;
6772 let mut sum: f64 = 0.0;
6773 let mut mn: f64 = f64::INFINITY;
6774 let mut mx: f64 = f64::NEG_INFINITY;
6775 while let Some(cols) = cursor.next_batch()? {
6776 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6777 if crate::columnar::all_non_null(validity, data.len()) {
6778 count += data.len() as u64;
6779 sum += data.iter().sum::<f64>();
6780 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6781 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6782 } else {
6783 for (i, &v) in data.iter().enumerate() {
6784 if crate::columnar::validity_bit(validity, i) {
6785 count += 1;
6786 sum += v;
6787 mn = mn.min(v);
6788 mx = mx.max(v);
6789 }
6790 }
6791 }
6792 }
6793 }
6794 Ok((count, sum, mn, mx))
6795}
6796
6797fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6798 if count == 0 && !matches!(agg, NativeAgg::Count) {
6799 return NativeAggResult::Null;
6800 }
6801 match agg {
6802 NativeAgg::Count => NativeAggResult::Count(count),
6803 NativeAgg::Sum => match sum.try_into() {
6806 Ok(v) => NativeAggResult::Int(v),
6807 Err(_) => NativeAggResult::Null,
6808 },
6809 NativeAgg::Min => NativeAggResult::Int(mn),
6810 NativeAgg::Max => NativeAggResult::Int(mx),
6811 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6812 }
6813}
6814
6815fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6816 if count == 0 && !matches!(agg, NativeAgg::Count) {
6817 return NativeAggResult::Null;
6818 }
6819 match agg {
6820 NativeAgg::Count => NativeAggResult::Count(count),
6821 NativeAgg::Sum => NativeAggResult::Float(sum),
6822 NativeAgg::Min => NativeAggResult::Float(mn),
6823 NativeAgg::Max => NativeAggResult::Float(mx),
6824 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6825 }
6826}
6827
6828fn agg_int(
6831 stats: &[crate::page::PageStat],
6832 decode: fn(Option<&[u8]>) -> Option<i64>,
6833) -> Option<(Option<i64>, Option<i64>, u64)> {
6834 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6835 let mut any = false;
6836 for s in stats {
6837 if let Some(v) = decode(s.min.as_deref()) {
6838 mn = mn.min(v);
6839 any = true;
6840 }
6841 if let Some(v) = decode(s.max.as_deref()) {
6842 mx = mx.max(v);
6843 any = true;
6844 }
6845 nulls += s.null_count;
6846 }
6847 any.then_some((Some(mn), Some(mx), nulls))
6848}
6849
6850fn agg_float(
6852 stats: &[crate::page::PageStat],
6853 decode: fn(Option<&[u8]>) -> Option<f64>,
6854) -> Option<(Option<f64>, Option<f64>, u64)> {
6855 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6856 let mut any = false;
6857 for s in stats {
6858 if let Some(v) = decode(s.min.as_deref()) {
6859 mn = mn.min(v);
6860 any = true;
6861 }
6862 if let Some(v) = decode(s.max.as_deref()) {
6863 mx = mx.max(v);
6864 any = true;
6865 }
6866 nulls += s.null_count;
6867 }
6868 any.then_some((Some(mn), Some(mx), nulls))
6869}
6870
6871type SecondaryIndexes = (
6873 HashMap<u16, BitmapIndex>,
6874 HashMap<u16, AnnIndex>,
6875 HashMap<u16, FmIndex>,
6876 HashMap<u16, SparseIndex>,
6877 HashMap<u16, MinHashIndex>,
6878);
6879
6880fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6881 let mut bitmap = HashMap::new();
6882 let mut ann = HashMap::new();
6883 let mut fm = HashMap::new();
6884 let mut sparse = HashMap::new();
6885 let mut minhash = HashMap::new();
6886 for idef in &schema.indexes {
6887 match idef.kind {
6888 IndexKind::Bitmap => {
6889 bitmap.insert(idef.column_id, BitmapIndex::new());
6890 }
6891 IndexKind::Ann => {
6892 let dim = schema
6893 .columns
6894 .iter()
6895 .find(|c| c.id == idef.column_id)
6896 .and_then(|c| match c.ty {
6897 TypeId::Embedding { dim } => Some(dim as usize),
6898 _ => None,
6899 })
6900 .unwrap_or(0);
6901 ann.insert(idef.column_id, AnnIndex::new(dim));
6902 }
6903 IndexKind::FmIndex => {
6904 fm.insert(idef.column_id, FmIndex::new());
6905 }
6906 IndexKind::Sparse => {
6907 sparse.insert(idef.column_id, SparseIndex::new());
6908 }
6909 IndexKind::MinHash => {
6910 minhash.insert(idef.column_id, MinHashIndex::new());
6911 }
6912 _ => {}
6913 }
6914 }
6915 (bitmap, ann, fm, sparse, minhash)
6916}
6917
6918const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
6919 | ColumnFlags::AUTO_INCREMENT
6920 | ColumnFlags::ENCRYPTED
6921 | ColumnFlags::ENCRYPTED_INDEXABLE
6922 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
6923
6924fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
6925 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
6926 return Err(MongrelError::Schema(
6927 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
6928 ));
6929 }
6930 Ok(())
6931}
6932
6933fn validate_alter_column_type(
6934 schema: &Schema,
6935 old: &ColumnDef,
6936 next: &ColumnDef,
6937 has_stored_versions: bool,
6938) -> Result<()> {
6939 if old.ty == next.ty {
6940 return Ok(());
6941 }
6942 if schema.indexes.iter().any(|i| i.column_id == old.id) {
6943 return Err(MongrelError::Schema(format!(
6944 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
6945 old.name
6946 )));
6947 }
6948 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
6949 return Ok(());
6950 }
6951 Err(MongrelError::Schema(format!(
6952 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
6953 old.ty, next.ty
6954 )))
6955}
6956
6957fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
6958 matches!(
6959 (old, new),
6960 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
6961 )
6962}
6963
6964fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
6970 let mut prev: Option<i64> = None;
6971 for r in rows {
6972 match r.columns.get(&pk_id) {
6973 Some(Value::Int64(v)) => {
6974 if prev.is_some_and(|p| p >= *v) {
6975 return false;
6976 }
6977 prev = Some(*v);
6978 }
6979 _ => return false,
6980 }
6981 }
6982 true
6983}
6984
6985#[allow(clippy::too_many_arguments)]
6986fn index_into(
6987 schema: &Schema,
6988 row: &Row,
6989 hot: &mut HotIndex,
6990 bitmap: &mut HashMap<u16, BitmapIndex>,
6991 ann: &mut HashMap<u16, AnnIndex>,
6992 fm: &mut HashMap<u16, FmIndex>,
6993 sparse: &mut HashMap<u16, SparseIndex>,
6994 minhash: &mut HashMap<u16, MinHashIndex>,
6995) {
6996 for idef in &schema.indexes {
6997 let Some(val) = row.columns.get(&idef.column_id) else {
6998 continue;
6999 };
7000 match idef.kind {
7001 IndexKind::Bitmap => {
7002 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7003 b.insert(val.encode_key(), row.row_id);
7004 }
7005 }
7006 IndexKind::Ann => {
7007 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7008 a.insert(v, row.row_id);
7009 }
7010 }
7011 IndexKind::FmIndex => {
7012 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7013 f.insert(b.clone(), row.row_id);
7014 }
7015 }
7016 IndexKind::Sparse => {
7017 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7018 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7021 s.insert(&terms, row.row_id);
7022 }
7023 }
7024 }
7025 IndexKind::MinHash => {
7026 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7027 let tokens = crate::index::token_hashes_from_bytes(b);
7030 mh.insert(&tokens, row.row_id);
7031 }
7032 }
7033 _ => {}
7034 }
7035 }
7036 if let Some(pk_col) = schema.primary_key() {
7037 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7038 hot.insert(pk_val.encode_key(), row.row_id);
7039 }
7040 }
7041}
7042
7043#[allow(dead_code)]
7049fn bulk_index_key(
7050 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7051 column_id: u16,
7052 ty: TypeId,
7053 col: &columnar::NativeColumn,
7054 i: usize,
7055) -> Option<Vec<u8>> {
7056 let encoded = columnar::encode_key_native(ty, col, i)?;
7057 #[cfg(feature = "encryption")]
7058 {
7059 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7060 if let Some((key, scheme)) = column_keys.get(&column_id) {
7061 return Some(match (*scheme, col) {
7062 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7063 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7064 ope_token_i64(key, data[i]).to_vec()
7065 }
7066 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7067 ope_token_f64(key, data[i]).to_vec()
7068 }
7069 _ => hmac_token(key, &encoded).to_vec(),
7070 });
7071 }
7072 }
7073 #[cfg(not(feature = "encryption"))]
7074 {
7075 let _ = (column_id, column_keys, col);
7076 }
7077 Some(encoded)
7078}
7079
7080pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7081 let json = serde_json::to_string_pretty(schema)
7082 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7083 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7084 Ok(())
7085}
7086
7087fn read_schema(dir: &Path) -> Result<Schema> {
7088 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7089 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7090}
7091
7092fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7093 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7094}
7095
7096fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7097 let n = list_wal_numbers(wal_dir)?;
7098 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7099}
7100
7101fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7102 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7103}
7104
7105fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7106 let _ = std::fs::create_dir_all(wal_dir);
7107 let mut max_n = None;
7108 for entry in std::fs::read_dir(wal_dir)? {
7109 let entry = entry?;
7110 let fname = entry.file_name();
7111 let Some(s) = fname.to_str() else {
7112 continue;
7113 };
7114 let Some(stripped) = s.strip_prefix("seg-") else {
7115 continue;
7116 };
7117 let Some(stripped) = stripped.strip_suffix(".wal") else {
7118 continue;
7119 };
7120 if let Ok(n) = stripped.parse::<u32>() {
7121 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7122 }
7123 }
7124 Ok(max_n)
7125}