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 had_deletes: bool,
168 agg_cache: HashMap<u64, CachedAgg>,
172 global_idx_epoch: u64,
176 indexes_complete: bool,
181 index_build_policy: IndexBuildPolicy,
183 pk_by_row_complete: bool,
190 flushed_epoch: u64,
193 page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
196 snapshots: Arc<crate::retention::SnapshotRegistry>,
199 commit_lock: Arc<parking_lot::Mutex<()>>,
201 decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
204 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
213 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
215 pending_delete_rids: roaring::RoaringBitmap,
218 pending_put_cols: std::collections::HashSet<u16>,
221 pending_rows: Vec<Row>,
227 pending_rows_auto_inc: Vec<bool>,
228 pending_dels: Vec<RowId>,
231 pending_truncate: Option<Epoch>,
235 auto_inc: Option<AutoIncState>,
238}
239
240const _: () = {
247 const fn assert_sync<T: ?Sized + Sync>() {}
248 assert_sync::<Table>();
249};
250
251enum CachedData {
257 Rows(Arc<Vec<Row>>),
258 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
259}
260
261impl CachedData {
262 fn approx_bytes(&self) -> u64 {
263 match self {
264 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
265 CachedData::Columns(c) => c
266 .iter()
267 .map(|(_, c)| c.approx_bytes())
268 .sum::<u64>()
269 .saturating_add(c.len() as u64 * 16),
270 }
271 }
272}
273
274struct CachedEntry {
278 data: CachedData,
279 footprint: roaring::RoaringBitmap,
280 condition_cols: Vec<u16>,
281}
282
283struct ResultCache {
294 entries: std::collections::HashMap<u64, CachedEntry>,
295 order: std::collections::VecDeque<u64>,
296 bytes: u64,
297 max_bytes: u64,
298 dir: Option<std::path::PathBuf>,
299 #[allow(dead_code)]
300 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
301}
302
303#[derive(serde::Serialize, serde::Deserialize)]
305struct SerializedEntry {
306 condition_cols: Vec<u16>,
307 footprint_bits: Vec<u32>,
308 data: SerializedData,
309}
310
311#[derive(serde::Serialize, serde::Deserialize)]
312enum SerializedData {
313 Rows(Vec<Row>),
314 Columns(Vec<(u16, columnar::NativeColumn)>),
315}
316
317impl SerializedEntry {
318 fn from_entry(entry: &CachedEntry) -> Self {
319 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
320 let data = match &entry.data {
321 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
322 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
323 };
324 Self {
325 condition_cols: entry.condition_cols.clone(),
326 footprint_bits,
327 data,
328 }
329 }
330
331 fn into_entry(self) -> Option<CachedEntry> {
332 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
333 let data = match self.data {
334 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
335 SerializedData::Columns(c) => {
336 if !c.iter().all(|(_, col)| col.validate()) {
339 return None;
340 }
341 CachedData::Columns(Arc::new(c))
342 }
343 };
344 Some(CachedEntry {
345 data,
346 footprint,
347 condition_cols: self.condition_cols,
348 })
349 }
350}
351
352impl ResultCache {
353 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
354
355 fn new() -> Self {
356 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
357 }
358
359 fn with_max_bytes(max_bytes: u64) -> Self {
360 Self {
361 entries: std::collections::HashMap::new(),
362 order: std::collections::VecDeque::new(),
363 bytes: 0,
364 max_bytes,
365 dir: None,
366 cache_dek: None,
367 }
368 }
369
370 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
371 let _ = std::fs::create_dir_all(&dir);
372 self.dir = Some(dir);
373 self
374 }
375
376 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
377 self.cache_dek = dek;
378 self
379 }
380
381 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
382 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
383 }
384
385 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
389 let Some(path) = self.disk_path(key) else {
390 return;
391 };
392 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
393 Ok(s) => s,
394 Err(_) => return,
395 };
396 let on_disk = if let Some(dek) = &self.cache_dek {
398 match self.encrypt_cache(&serialized, dek) {
399 Some(b) => b,
400 None => return,
401 }
402 } else {
403 serialized
404 };
405 let tmp = path.with_extension("tmp");
406 use std::io::Write;
407 let write = || -> std::io::Result<()> {
408 let mut f = std::fs::File::create(&tmp)?;
409 f.write_all(&on_disk)?;
410 f.flush()?;
411 Ok(())
412 };
413 if write().is_err() {
414 let _ = std::fs::remove_file(&tmp);
415 return;
416 }
417 let _ = std::fs::rename(&tmp, &path);
418 }
419
420 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
422 let path = self.disk_path(key)?;
423 let bytes = std::fs::read(&path).ok()?;
424 let plaintext = if let Some(dek) = &self.cache_dek {
425 self.decrypt_cache(&bytes, dek)?
426 } else {
427 bytes
428 };
429 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
430 serialized.into_entry()
431 }
432
433 fn remove_from_disk(&self, key: u64) {
435 if let Some(path) = self.disk_path(key) {
436 let _ = std::fs::remove_file(&path);
437 }
438 }
439
440 #[cfg(feature = "encryption")]
442 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
443 use crate::encryption::Cipher;
444 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
445 let mut nonce = [0u8; 12];
446 crate::encryption::fill_random(&mut nonce);
447 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
448 let mut out = Vec::with_capacity(12 + ct.len());
449 out.extend_from_slice(&nonce);
450 out.extend_from_slice(&ct);
451 Some(out)
452 }
453
454 #[cfg(not(feature = "encryption"))]
455 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
456 None
457 }
458
459 #[cfg(feature = "encryption")]
461 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
462 use crate::encryption::Cipher;
463 if bytes.len() < 28 {
464 return None;
465 }
466 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
467 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
468 let ct = &bytes[12..];
469 cipher.decrypt_page(&nonce, ct).ok()
470 }
471
472 #[cfg(not(feature = "encryption"))]
473 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
474 None
475 }
476
477 fn load_persistent(&mut self) {
480 let Some(dir) = self.dir.as_ref().cloned() else {
481 return;
482 };
483 let entries = match std::fs::read_dir(&dir) {
484 Ok(e) => e,
485 Err(_) => return,
486 };
487 for entry in entries.flatten() {
488 let path = entry.path();
489 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
491 let _ = std::fs::remove_file(&path);
492 continue;
493 }
494 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
495 continue;
496 }
497 let stem = match path.file_stem().and_then(|s| s.to_str()) {
498 Some(s) => s,
499 None => continue,
500 };
501 let key = match u64::from_str_radix(stem, 16) {
502 Ok(k) => k,
503 Err(_) => continue,
504 };
505 let bytes = match std::fs::read(&path) {
506 Ok(b) => b,
507 Err(_) => continue,
508 };
509 let plaintext = if let Some(dek) = &self.cache_dek {
511 match self.decrypt_cache(&bytes, dek) {
512 Some(p) => p,
513 None => {
514 let _ = std::fs::remove_file(&path);
515 continue;
516 }
517 }
518 } else {
519 bytes
520 };
521 match bincode::deserialize::<SerializedEntry>(&plaintext) {
522 Ok(serialized) => {
523 if let Some(entry) = serialized.into_entry() {
524 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
525 self.entries.insert(key, entry);
526 self.order.push_back(key);
527 } else {
528 let _ = std::fs::remove_file(&path);
529 }
530 }
531 Err(_) => {
532 let _ = std::fs::remove_file(&path);
533 }
534 }
535 }
536 self.evict();
537 }
538
539 fn set_max_bytes(&mut self, max_bytes: u64) {
540 self.max_bytes = max_bytes;
541 self.evict();
542 }
543
544 fn touch(&mut self, key: u64) {
546 self.order.retain(|k| *k != key);
547 self.order.push_back(key);
548 }
549
550 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
551 let res = self.entries.get(&key).and_then(|e| match &e.data {
552 CachedData::Rows(r) => Some(r.clone()),
553 CachedData::Columns(_) => None,
554 });
555 if res.is_some() {
556 self.touch(key);
557 return res;
558 }
559 if let Some(entry) = self.load_from_disk(key) {
561 let res = match &entry.data {
562 CachedData::Rows(r) => Some(r.clone()),
563 CachedData::Columns(_) => None,
564 };
565 if res.is_some() {
566 let approx = entry.data.approx_bytes();
567 self.bytes = self.bytes.saturating_add(approx);
568 self.entries.insert(key, entry);
569 self.order.push_back(key);
570 self.evict();
571 return res;
572 }
573 }
574 None
575 }
576
577 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
578 let res = self.entries.get(&key).and_then(|e| match &e.data {
579 CachedData::Columns(c) => Some(c.clone()),
580 CachedData::Rows(_) => None,
581 });
582 if res.is_some() {
583 self.touch(key);
584 return res;
585 }
586 if let Some(entry) = self.load_from_disk(key) {
588 let res = match &entry.data {
589 CachedData::Columns(c) => Some(c.clone()),
590 CachedData::Rows(_) => None,
591 };
592 if res.is_some() {
593 let approx = entry.data.approx_bytes();
594 self.bytes = self.bytes.saturating_add(approx);
595 self.entries.insert(key, entry);
596 self.order.push_back(key);
597 self.evict();
598 return res;
599 }
600 }
601 None
602 }
603
604 fn insert(&mut self, key: u64, entry: CachedEntry) {
605 let approx = entry.data.approx_bytes();
606 if self.entries.remove(&key).is_some() {
607 self.order.retain(|k| *k != key);
608 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
609 }
610 self.store_to_disk(key, &entry);
612 self.bytes = self.bytes.saturating_add(approx);
613 self.entries.insert(key, entry);
614 self.order.push_back(key);
615 self.evict();
616 }
617
618 fn invalidate(
627 &mut self,
628 delete_rids: &roaring::RoaringBitmap,
629 put_cols: &std::collections::HashSet<u16>,
630 ) {
631 if self.entries.is_empty() {
632 return;
633 }
634 let has_deletes = !delete_rids.is_empty();
635 let to_remove: std::collections::HashSet<u64> = self
636 .entries
637 .iter()
638 .filter(|(_, e)| {
639 let delete_hit = if e.footprint.is_empty() {
640 has_deletes
641 } else {
642 e.footprint.intersection_len(delete_rids) > 0
643 };
644 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
645 delete_hit || col_hit
646 })
647 .map(|(&k, _)| k)
648 .collect();
649 for key in &to_remove {
650 if let Some(e) = self.entries.remove(key) {
651 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
652 }
653 self.remove_from_disk(*key);
654 }
655 if !to_remove.is_empty() {
656 self.order.retain(|k| !to_remove.contains(k));
657 }
658 }
659
660 fn clear(&mut self) {
661 if let Some(dir) = &self.dir {
663 if let Ok(entries) = std::fs::read_dir(dir) {
664 for entry in entries.flatten() {
665 let path = entry.path();
666 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
667 let _ = std::fs::remove_file(&path);
668 }
669 }
670 }
671 }
672 self.entries.clear();
673 self.order.clear();
674 self.bytes = 0;
675 }
676
677 fn evict(&mut self) {
678 while self.bytes > self.max_bytes {
679 let Some(k) = self.order.pop_front() else {
680 break;
681 };
682 if let Some(e) = self.entries.remove(&k) {
683 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
684 self.remove_from_disk(k);
688 }
689 }
690 }
691}
692
693type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
700
701fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
702 let _ = kek;
703 #[cfg(feature = "encryption")]
704 {
705 if let Some(k) = kek {
706 return (
707 Some(k.derive_table_wal_key(_table_id)),
708 Some(k.derive_cache_key()),
709 );
710 }
711 }
712 (None, None)
713}
714
715#[cfg(feature = "encryption")]
717fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
718 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
719}
720
721#[cfg(not(feature = "encryption"))]
722fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
723 Box::new(crate::encryption::PlaintextCipher)
724}
725
726fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
727 let Some(kek) = kek else {
728 return HashMap::new();
729 };
730 #[cfg(feature = "encryption")]
731 {
732 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
733 schema
734 .columns
735 .iter()
736 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
737 .map(|c| {
738 let scheme = if schema
739 .indexes
740 .iter()
741 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
742 {
743 SCHEME_OPE_RANGE
744 } else {
745 SCHEME_HMAC_EQ
746 };
747 let key: [u8; 32] = *kek.derive_column_key(c.id);
748 (c.id, (key, scheme))
749 })
750 .collect()
751 }
752 #[cfg(not(feature = "encryption"))]
753 {
754 let _ = (kek, schema);
755 HashMap::new()
756 }
757}
758
759pub(crate) struct SharedCtx {
764 pub epoch: Arc<EpochAuthority>,
765 pub page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
766 pub decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
767 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
768 pub kek: Option<Arc<Kek>>,
769 pub commit_lock: Arc<parking_lot::Mutex<()>>,
775 pub shared: Option<SharedWalCtx>,
779}
780
781#[derive(Clone)]
787pub(crate) struct SharedWalCtx {
788 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
789 pub group: Arc<GroupCommit>,
790 pub poisoned: Arc<AtomicBool>,
791 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
792}
793
794enum WalSink {
797 Private(Wal),
798 Shared(SharedWalCtx),
799}
800
801impl SharedCtx {
802 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
806 let mut cache = crate::cache::PageCache::new(PAGE_CACHE_CAPACITY);
807 if let Some(d) = cache_dir {
808 cache = cache.with_persistence(d);
809 }
810 Self {
811 epoch: Arc::new(EpochAuthority::new(0)),
812 page_cache: Arc::new(parking_lot::Mutex::new(cache)),
813 decoded_cache: Arc::new(parking_lot::Mutex::new(
814 crate::cache::DecodedPageCache::new(DECODED_CACHE_CAPACITY),
815 )),
816 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
817 kek,
818 commit_lock: Arc::new(parking_lot::Mutex::new(())),
819 shared: None,
820 }
821 }
822}
823
824impl Table {
825 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
826 let dir = dir.as_ref().to_path_buf();
827 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
828 Self::create_in(&dir, schema, table_id, ctx)
829 }
830
831 #[cfg(feature = "encryption")]
842 pub fn create_encrypted(
843 dir: impl AsRef<Path>,
844 schema: Schema,
845 table_id: u64,
846 passphrase: &str,
847 ) -> Result<Self> {
848 let dir = dir.as_ref();
849 std::fs::create_dir_all(dir.join(META_DIR))?;
850 let salt = crate::encryption::random_salt();
851 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
852 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
853 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
854 Self::create_in(dir, schema, table_id, ctx)
855 }
856
857 #[cfg(feature = "encryption")]
862 pub fn create_with_key(
863 dir: impl AsRef<Path>,
864 schema: Schema,
865 table_id: u64,
866 key: &[u8],
867 ) -> Result<Self> {
868 let dir = dir.as_ref();
869 std::fs::create_dir_all(dir.join(META_DIR))?;
870 let salt = crate::encryption::random_salt();
871 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
872 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
873 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
874 Self::create_in(dir, schema, table_id, ctx)
875 }
876
877 #[cfg(feature = "encryption")]
879 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
880 let dir = dir.as_ref();
881 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
882 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
883 MongrelError::NotFound(format!(
884 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
885 salt_path
886 ))
887 })?;
888 if salt_bytes.len() != crate::encryption::SALT_LEN {
889 return Err(MongrelError::InvalidArgument(format!(
890 "salt file is {} bytes, expected {}",
891 salt_bytes.len(),
892 crate::encryption::SALT_LEN
893 )));
894 }
895 let mut salt = [0u8; crate::encryption::SALT_LEN];
896 salt.copy_from_slice(&salt_bytes);
897 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
898 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
899 Self::open_in(dir, ctx)
900 }
901
902 pub(crate) fn create_in(
903 dir: impl AsRef<Path>,
904 schema: Schema,
905 table_id: u64,
906 ctx: SharedCtx,
907 ) -> Result<Self> {
908 schema.validate_auto_increment()?;
909 let dir = dir.as_ref().to_path_buf();
910 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
911 write_schema(&dir, &schema)?;
912 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
913 let (wal, current_txn_id) = match ctx.shared.clone() {
916 Some(s) => (WalSink::Shared(s), 0),
917 None => {
918 std::fs::create_dir_all(dir.join(WAL_DIR))?;
919 let mut w = if let Some(ref dk) = wal_dek {
920 Wal::create_with_cipher(
921 dir.join(WAL_DIR).join("seg-000000.wal"),
922 Epoch(0),
923 Some(make_cipher(dk)),
924 0,
925 )?
926 } else {
927 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
928 };
929 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
930 (WalSink::Private(w), 1)
931 }
932 };
933 let mut manifest = Manifest::new(table_id, schema.schema_id);
934 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
939 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
940 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
941 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
942 let auto_inc = resolve_auto_inc(&schema);
943 let rcache_dir = dir.join(RCACHE_DIR);
944 Ok(Self {
945 dir,
946 table_id,
947 wal,
948 memtable: Memtable::new(),
949 mutable_run: MutableRun::new(),
950 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
951 compaction_zstd_level: 3,
952 allocator: RowIdAllocator::new(0),
953 epoch: ctx.epoch,
954 persisted_epoch: 0,
955 schema,
956 hot: HotIndex::new(),
957 kek: ctx.kek,
958 column_keys,
959 run_refs: Vec::new(),
960 retiring: Vec::new(),
961 next_run_id: 1,
962 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
963 current_txn_id,
964 bitmap,
965 ann,
966 fm,
967 sparse,
968 minhash,
969 learned_range: HashMap::new(),
970 pk_by_row: HashMap::new(),
971 pinned: BTreeMap::new(),
972 live_count: 0,
973 reservoir: crate::reservoir::Reservoir::default(),
974 had_deletes: false,
975 agg_cache: HashMap::new(),
976 global_idx_epoch: 0,
977 indexes_complete: true,
978 index_build_policy: IndexBuildPolicy::default(),
979 pk_by_row_complete: false,
980 flushed_epoch: 0,
981 page_cache: ctx.page_cache,
982 decoded_cache: ctx.decoded_cache,
983 snapshots: ctx.snapshots,
984 commit_lock: ctx.commit_lock,
985 result_cache: Arc::new(parking_lot::Mutex::new(
986 ResultCache::new()
987 .with_dir(rcache_dir)
988 .with_cache_dek(cache_dek.clone()),
989 )),
990 pending_delete_rids: roaring::RoaringBitmap::new(),
991 pending_put_cols: std::collections::HashSet::new(),
992 pending_rows: Vec::new(),
993 pending_rows_auto_inc: Vec::new(),
994 pending_dels: Vec::new(),
995 pending_truncate: None,
996 wal_dek,
997 auto_inc,
998 })
999 }
1000
1001 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1005 let dir = dir.as_ref();
1006 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1007 Self::open_in(dir, ctx)
1008 }
1009
1010 #[cfg(feature = "encryption")]
1013 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1014 let dir = dir.as_ref();
1015 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1016 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1017 MongrelError::NotFound(format!(
1018 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1019 salt_path
1020 ))
1021 })?;
1022 let salt_len = crate::encryption::SALT_LEN;
1023 if salt_bytes.len() != salt_len {
1024 return Err(MongrelError::InvalidArgument(format!(
1025 "encryption salt is {} bytes, expected {salt_len}",
1026 salt_bytes.len()
1027 )));
1028 }
1029 let mut salt = [0u8; 16];
1030 salt.copy_from_slice(&salt_bytes);
1031 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1032 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1033 let t = Self::open_in(dir, ctx)?;
1034 Ok(t)
1035 }
1036
1037 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1038 let dir = dir.as_ref().to_path_buf();
1039 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1040 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1041 let schema: Schema = read_schema(&dir)?;
1042 let replay_epoch = Epoch(manifest.current_epoch);
1043 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1044 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1048 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1049 None => {
1050 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1051 let replayed = match &active {
1053 Some(path) => {
1054 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1055 crate::wal::replay_with_cipher(path, cipher)?
1056 }
1057 None => Vec::new(),
1058 };
1059 let mut w = match &active {
1060 Some(path) => Wal::create_with_cipher(
1061 path,
1062 replay_epoch,
1063 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1064 0,
1065 )?,
1066 None => Wal::create_with_cipher(
1067 dir.join(WAL_DIR).join("seg-000000.wal"),
1068 replay_epoch,
1069 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1070 0,
1071 )?,
1072 };
1073 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1074 (WalSink::Private(w), replayed, 1)
1075 }
1076 };
1077
1078 let mut memtable = Memtable::new();
1079 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1080 let persisted_epoch = manifest.current_epoch;
1081 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1088 s.next = manifest.auto_inc_next;
1089 s.seeded = manifest.auto_inc_next > 0;
1090 s
1091 });
1092
1093 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1100 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1101 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1102 std::collections::BTreeMap::new();
1103 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1104 let mut saw_delete = false;
1105 for record in replayed {
1106 let txn_id = record.txn_id;
1107 match record.op {
1108 Op::Put { rows, .. } => {
1109 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1110 for row in &rows {
1111 allocator.advance_to(row.row_id);
1112 if let Some(ai) = auto_inc.as_mut() {
1113 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1114 if *n + 1 > ai.next {
1115 ai.next = *n + 1;
1116 }
1117 }
1118 }
1119 }
1120 staged_puts.entry(txn_id).or_default().extend(rows);
1121 }
1122 Op::Delete { row_ids, .. } => {
1123 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1124 }
1125 Op::TxnCommit { epoch, .. } => {
1126 let commit_epoch = Epoch(epoch);
1127 if let Some(puts) = staged_puts.remove(&txn_id) {
1128 for row in &puts {
1129 memtable.upsert(row.clone());
1130 }
1131 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1132 }
1133 if let Some(dels) = staged_deletes.remove(&txn_id) {
1134 saw_delete = true;
1135 for rid in dels {
1136 memtable.tombstone(rid, commit_epoch);
1137 replayed_deletes.push((rid, commit_epoch));
1138 }
1139 }
1140 }
1141 Op::TxnAbort => {
1142 staged_puts.remove(&txn_id);
1143 staged_deletes.remove(&txn_id);
1144 }
1145 Op::TruncateTable { .. } | Op::Flush { .. } | Op::Ddl(_) => {}
1146 }
1147 }
1148
1149 let rcache_dir = dir.join(RCACHE_DIR);
1150 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1151 let mut db = Self {
1152 dir,
1153 table_id: manifest.table_id,
1154 wal,
1155 memtable,
1156 mutable_run: MutableRun::new(),
1157 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1158 compaction_zstd_level: 3,
1159 allocator,
1160 epoch: ctx.epoch,
1161 persisted_epoch,
1162 schema,
1163 hot: HotIndex::new(),
1164 kek: ctx.kek,
1165 column_keys,
1166 run_refs: manifest.runs.clone(),
1167 retiring: manifest.retiring.clone(),
1168 next_run_id: manifest
1169 .runs
1170 .iter()
1171 .map(|r| r.run_id as u64 + 1)
1172 .max()
1173 .unwrap_or(1),
1174 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1175 current_txn_id,
1176 bitmap: HashMap::new(),
1177 ann: HashMap::new(),
1178 fm: HashMap::new(),
1179 sparse: HashMap::new(),
1180 minhash: HashMap::new(),
1181 learned_range: HashMap::new(),
1182 pk_by_row: HashMap::new(),
1183 pinned: BTreeMap::new(),
1184 live_count: manifest.live_count,
1185 reservoir: crate::reservoir::Reservoir::default(),
1186 had_deletes: saw_delete,
1187 agg_cache: HashMap::new(),
1188 global_idx_epoch: manifest.global_idx_epoch,
1189 indexes_complete: true,
1190 index_build_policy: IndexBuildPolicy::default(),
1191 pk_by_row_complete: false,
1192 flushed_epoch: manifest.flushed_epoch,
1193 page_cache: ctx.page_cache,
1194 decoded_cache: ctx.decoded_cache,
1195 snapshots: ctx.snapshots,
1196 commit_lock: ctx.commit_lock,
1197 result_cache: Arc::new(parking_lot::Mutex::new(
1198 ResultCache::new()
1199 .with_dir(rcache_dir)
1200 .with_cache_dek(cache_dek.clone()),
1201 )),
1202 pending_delete_rids: roaring::RoaringBitmap::new(),
1203 pending_put_cols: std::collections::HashSet::new(),
1204 pending_rows: Vec::new(),
1205 pending_rows_auto_inc: Vec::new(),
1206 pending_dels: Vec::new(),
1207 pending_truncate: None,
1208 wal_dek,
1209 auto_inc,
1210 };
1211
1212 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1215
1216 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1221 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1222 c.epoch_built == manifest.global_idx_epoch
1223 && manifest.global_idx_epoch > 0
1224 && manifest
1225 .runs
1226 .iter()
1227 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1228 });
1229 if let Some(loaded) = checkpoint {
1230 if checkpoint_valid {
1231 db.hot = loaded.hot;
1232 db.bitmap = loaded.bitmap;
1233 db.ann = loaded.ann;
1234 db.fm = loaded.fm;
1235 db.sparse = loaded.sparse;
1236 db.minhash = loaded.minhash;
1237 db.learned_range = loaded.learned_range;
1238 }
1241 }
1242 if !checkpoint_valid {
1243 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1244 db.bitmap = bitmap;
1245 db.ann = ann;
1246 db.fm = fm;
1247 db.sparse = sparse;
1248 db.minhash = minhash;
1249 db.rebuild_indexes_from_runs()?;
1250 db.build_learned_ranges()?;
1251 }
1252
1253 for (epoch, group) in replayed_puts {
1258 let (losers, winner_pks) = db.partition_pk_winners(&group);
1259 for (key, &row_id) in &winner_pks {
1260 if let Some(old_rid) = db.hot.get(key) {
1261 if old_rid != row_id {
1262 db.tombstone_row(old_rid, epoch, false);
1263 }
1264 }
1265 }
1266 for &loser_rid in &losers {
1267 db.tombstone_row(loser_rid, epoch, false);
1268 }
1269 for (key, row_id) in winner_pks {
1270 db.insert_hot_pk(key, row_id);
1271 }
1272 if db.schema.primary_key().is_none() {
1273 for r in &group {
1274 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1275 }
1276 }
1277 for r in &group {
1278 if !losers.contains(&r.row_id) {
1279 db.index_row(r);
1280 }
1281 }
1282 }
1283 for (rid, epoch) in &replayed_deletes {
1287 db.remove_hot_for_row(*rid, *epoch);
1288 }
1289
1290 let _ = db.rebuild_reservoir();
1291 db.result_cache.lock().load_persistent();
1294 Ok(db)
1295 }
1296
1297 fn rebuild_reservoir(&mut self) -> Result<()> {
1300 let snap = self.snapshot();
1301 let rows = self.visible_rows(snap)?;
1302 self.reservoir.reset();
1303 for r in rows {
1304 self.reservoir.offer(r.row_id.0);
1305 }
1306 Ok(())
1307 }
1308
1309 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1310 self.hot = HotIndex::new();
1311 self.pk_by_row.clear();
1312 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1313 self.bitmap = bitmap;
1314 self.ann = ann;
1315 self.fm = fm;
1316 self.sparse = sparse;
1317 self.minhash = minhash;
1318 let snapshot = Epoch(u64::MAX);
1319 for rr in self.run_refs.clone() {
1320 let mut reader = self.open_reader(rr.run_id)?;
1321 for row in reader.visible_rows(snapshot)? {
1322 let tok_row = self.tokenized_for_indexes(&row);
1323 index_into(
1324 &self.schema,
1325 &tok_row,
1326 &mut self.hot,
1327 &mut self.bitmap,
1328 &mut self.ann,
1329 &mut self.fm,
1330 &mut self.sparse,
1331 &mut self.minhash,
1332 );
1333 }
1334 }
1335 for row in self.mutable_run.visible_versions(snapshot) {
1336 if row.deleted {
1337 self.remove_hot_for_row(row.row_id, snapshot);
1338 } else {
1339 self.index_row(&row);
1340 }
1341 }
1342 for row in self.memtable.visible_versions(snapshot) {
1343 if row.deleted {
1344 self.remove_hot_for_row(row.row_id, snapshot);
1345 } else {
1346 self.index_row(&row);
1347 }
1348 }
1349 self.refresh_pk_by_row_from_hot();
1350 Ok(())
1351 }
1352
1353 fn refresh_pk_by_row_from_hot(&mut self) {
1354 self.pk_by_row.clear();
1355 self.pk_by_row_complete = true;
1356 if self.schema.primary_key().is_none() {
1357 return;
1358 }
1359 for (key, row_id) in self.hot.entries() {
1360 self.pk_by_row.insert(row_id, key);
1361 }
1362 }
1363
1364 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1365 if self.schema.primary_key().is_some() {
1366 self.pk_by_row.insert(row_id, key.clone());
1367 }
1368 self.hot.insert(key, row_id);
1369 }
1370
1371 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1375 self.learned_range.clear();
1376 if self.run_refs.len() != 1 {
1377 return Ok(());
1378 }
1379 let cols: Vec<u16> = self
1380 .schema
1381 .indexes
1382 .iter()
1383 .filter(|i| i.kind == IndexKind::LearnedRange)
1384 .map(|i| i.column_id)
1385 .collect();
1386 if cols.is_empty() {
1387 return Ok(());
1388 }
1389 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1390 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1391 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1392 _ => return Ok(()),
1393 };
1394 for cid in cols {
1395 let ty = self
1396 .schema
1397 .columns
1398 .iter()
1399 .find(|c| c.id == cid)
1400 .map(|c| c.ty)
1401 .unwrap_or(TypeId::Int64);
1402 match ty {
1403 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1404 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1405 let pairs: Vec<(i64, u64)> = data
1406 .iter()
1407 .zip(row_ids.iter())
1408 .map(|(v, r)| (*v, *r))
1409 .collect();
1410 self.learned_range
1411 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1412 }
1413 }
1414 TypeId::Float64 => {
1415 if let columnar::NativeColumn::Float64 { data, .. } =
1416 reader.column_native(cid)?
1417 {
1418 let pairs: Vec<(f64, u64)> = data
1419 .iter()
1420 .zip(row_ids.iter())
1421 .map(|(v, r)| (*v, *r))
1422 .collect();
1423 self.learned_range
1424 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1425 }
1426 }
1427 _ => {}
1428 }
1429 }
1430 Ok(())
1431 }
1432
1433 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1440 if self.indexes_complete {
1441 crate::trace::QueryTrace::record(|t| {
1442 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1443 });
1444 return Ok(());
1445 }
1446 crate::trace::QueryTrace::record(|t| {
1447 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1448 });
1449 self.rebuild_indexes_from_runs()?;
1450 self.build_learned_ranges()?;
1451 self.indexes_complete = true;
1452 let epoch = self.current_epoch();
1453 self.checkpoint_indexes(epoch);
1454 Ok(())
1455 }
1456
1457 fn pending_epoch(&self) -> Epoch {
1458 Epoch(self.epoch.visible().0 + 1)
1459 }
1460
1461 fn is_shared(&self) -> bool {
1464 matches!(self.wal, WalSink::Shared(_))
1465 }
1466
1467 fn ensure_txn_id(&mut self) -> u64 {
1471 if self.current_txn_id == 0 {
1472 let id = match &self.wal {
1473 WalSink::Shared(s) => {
1474 let mut g = s.txn_ids.lock();
1475 let v = *g;
1476 *g = g.wrapping_add(1);
1477 v
1478 }
1479 WalSink::Private(_) => 1,
1480 };
1481 self.current_txn_id = id;
1482 }
1483 self.current_txn_id
1484 }
1485
1486 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1489 let txn_id = self.ensure_txn_id();
1490 let table_id = self.table_id;
1491 match &mut self.wal {
1492 WalSink::Private(w) => {
1493 w.append_txn(txn_id, op)?;
1494 }
1495 WalSink::Shared(s) => {
1496 s.wal.lock().append(txn_id, table_id, op)?;
1497 }
1498 }
1499 Ok(())
1500 }
1501
1502 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1510 Ok(self.put_returning(columns)?.0)
1511 }
1512
1513 pub fn put_returning(
1518 &mut self,
1519 mut columns: Vec<(u16, Value)>,
1520 ) -> Result<(RowId, Option<i64>)> {
1521 let assigned = self.fill_auto_inc(&mut columns)?;
1522 self.schema.validate_not_null(&columns)?;
1523 let row_id = self.allocator.alloc();
1524 let epoch = self.pending_epoch();
1525 let mut row = Row::new(row_id, epoch);
1526 for (col_id, val) in columns {
1527 row.columns.insert(col_id, val);
1528 }
1529 self.commit_rows(vec![row], assigned.is_some())?;
1530 Ok((row_id, assigned))
1531 }
1532
1533 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1536 Ok(self
1537 .put_batch_returning(batch)?
1538 .into_iter()
1539 .map(|(r, _)| r)
1540 .collect())
1541 }
1542
1543 pub fn put_batch_returning(
1546 &mut self,
1547 batch: Vec<Vec<(u16, Value)>>,
1548 ) -> Result<Vec<(RowId, Option<i64>)>> {
1549 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1550 for mut cols in batch {
1551 let assigned = self.fill_auto_inc(&mut cols)?;
1552 filled.push((cols, assigned));
1553 }
1554 for (cols, _) in &filled {
1555 self.schema.validate_not_null(cols)?;
1556 }
1557 let epoch = self.pending_epoch();
1558 let mut rows = Vec::with_capacity(filled.len());
1559 let mut ids = Vec::with_capacity(filled.len());
1560 for (cols, assigned) in filled {
1561 let row_id = self.allocator.alloc();
1562 let mut row = Row::new(row_id, epoch);
1563 for (c, v) in cols {
1564 row.columns.insert(c, v);
1565 }
1566 ids.push((row_id, assigned));
1567 rows.push(row);
1568 }
1569 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1570 self.commit_rows(rows, all_auto_generated)?;
1571 Ok(ids)
1572 }
1573
1574 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1580 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1581 return Ok(None);
1582 };
1583 let pos = columns.iter().position(|(c, _)| *c == cid);
1584 let assigned = match pos {
1585 Some(i) => match &columns[i].1 {
1586 Value::Null => {
1587 let next = self.alloc_auto_inc_value()?;
1588 columns[i].1 = Value::Int64(next);
1589 Some(next)
1590 }
1591 Value::Int64(n) => {
1592 self.advance_auto_inc_past(*n)?;
1593 None
1594 }
1595 other => {
1596 return Err(MongrelError::InvalidArgument(format!(
1597 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1598 other
1599 )))
1600 }
1601 },
1602 None => {
1603 let next = self.alloc_auto_inc_value()?;
1604 columns.push((cid, Value::Int64(next)));
1605 Some(next)
1606 }
1607 };
1608 Ok(assigned)
1609 }
1610
1611 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1613 self.ensure_auto_inc_seeded()?;
1614 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1616 let v = ai.next;
1617 ai.next = ai.next.saturating_add(1);
1618 Ok(v)
1619 }
1620
1621 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1624 self.ensure_auto_inc_seeded()?;
1625 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1626 let floor = used.saturating_add(1).max(1);
1627 if ai.next < floor {
1628 ai.next = floor;
1629 }
1630 Ok(())
1631 }
1632
1633 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1638 let needs_seed = match self.auto_inc {
1639 Some(ai) => !ai.seeded,
1640 None => return Ok(()),
1641 };
1642 if !needs_seed {
1643 return Ok(());
1644 }
1645 if self.seed_empty_auto_inc() {
1646 return Ok(());
1647 }
1648 let cid = self
1649 .auto_inc
1650 .as_ref()
1651 .expect("auto-inc column present")
1652 .column_id;
1653 let max = self.scan_max_int64(cid)?;
1654 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1655 let floor = max.saturating_add(1).max(1);
1656 if ai.next < floor {
1657 ai.next = floor;
1658 }
1659 ai.seeded = true;
1660 Ok(())
1661 }
1662
1663 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1664 if n == 0 || self.auto_inc.is_none() {
1665 return Ok(None);
1666 }
1667 self.ensure_auto_inc_seeded()?;
1668 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1669 let start = ai.next;
1670 ai.next = ai.next.saturating_add(n as i64);
1671 Ok(Some(start))
1672 }
1673
1674 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1678 let mut max: i64 = 0;
1679 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1680 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1681 if *n > max {
1682 max = *n;
1683 }
1684 }
1685 }
1686 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1687 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1688 if *n > max {
1689 max = *n;
1690 }
1691 }
1692 }
1693 for rr in self.run_refs.clone() {
1694 let reader = self.open_reader(rr.run_id)?;
1695 if let Some(stats) = reader.column_page_stats(column_id) {
1696 for s in stats {
1697 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1698 if n > max {
1699 max = n;
1700 }
1701 }
1702 }
1703 } else if reader.has_column(column_id) {
1704 if let columnar::NativeColumn::Int64 { data, validity } =
1705 reader.column_native_shared(column_id)?
1706 {
1707 for (i, n) in data.iter().enumerate() {
1708 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1709 {
1710 max = *n;
1711 }
1712 }
1713 }
1714 }
1715 }
1716 Ok(max)
1717 }
1718
1719 fn seed_empty_auto_inc(&mut self) -> bool {
1720 let Some(ai) = self.auto_inc.as_mut() else {
1721 return false;
1722 };
1723 if ai.seeded || self.live_count != 0 {
1724 return false;
1725 }
1726 if ai.next < 1 {
1727 ai.next = 1;
1728 }
1729 ai.seeded = true;
1730 true
1731 }
1732
1733 fn advance_auto_inc_from_native_columns(
1734 &mut self,
1735 columns: &[(u16, columnar::NativeColumn)],
1736 n: usize,
1737 live_before: u64,
1738 ) -> Result<()> {
1739 let Some(ai) = self.auto_inc.as_mut() else {
1740 return Ok(());
1741 };
1742 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1743 return Ok(());
1744 };
1745 let columnar::NativeColumn::Int64 { data, validity } = col else {
1746 return Err(MongrelError::InvalidArgument(format!(
1747 "AUTO_INCREMENT column {} must be Int64",
1748 ai.column_id
1749 )));
1750 };
1751 let max = if native_int64_strictly_increasing(col, n) {
1752 data.get(n.saturating_sub(1)).copied()
1753 } else {
1754 data.iter()
1755 .take(n)
1756 .enumerate()
1757 .filter_map(|(i, v)| {
1758 if validity.is_empty() || columnar::validity_bit(validity, i) {
1759 Some(*v)
1760 } else {
1761 None
1762 }
1763 })
1764 .max()
1765 };
1766 if let Some(max) = max {
1767 let floor = max.saturating_add(1).max(1);
1768 if ai.next < floor {
1769 ai.next = floor;
1770 }
1771 if ai.seeded || live_before == 0 {
1772 ai.seeded = true;
1773 }
1774 }
1775 Ok(())
1776 }
1777
1778 fn fill_auto_inc_native_columns(
1779 &mut self,
1780 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1781 n: usize,
1782 ) -> Result<()> {
1783 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1784 return Ok(());
1785 };
1786 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1787 if let Some(start) = self.alloc_auto_inc_range(n)? {
1788 columns.push((
1789 cid,
1790 columnar::NativeColumn::Int64 {
1791 data: (start..start.saturating_add(n as i64)).collect(),
1792 validity: vec![0xFF; n.div_ceil(8)],
1793 },
1794 ));
1795 }
1796 return Ok(());
1797 };
1798
1799 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1800 return Err(MongrelError::InvalidArgument(format!(
1801 "AUTO_INCREMENT column {cid} must be Int64"
1802 )));
1803 };
1804 if data.len() < n {
1805 return Err(MongrelError::InvalidArgument(format!(
1806 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1807 data.len()
1808 )));
1809 }
1810 if columnar::all_non_null(validity, n) {
1811 return Ok(());
1812 }
1813 if validity.iter().all(|b| *b == 0) {
1814 if let Some(start) = self.alloc_auto_inc_range(n)? {
1815 for (i, slot) in data.iter_mut().take(n).enumerate() {
1816 *slot = start.saturating_add(i as i64);
1817 }
1818 *validity = vec![0xFF; n.div_ceil(8)];
1819 }
1820 return Ok(());
1821 }
1822
1823 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1824 for (i, slot) in data.iter_mut().enumerate().take(n) {
1825 if columnar::validity_bit(validity, i) {
1826 self.advance_auto_inc_past(*slot)?;
1827 } else {
1828 *slot = self.alloc_auto_inc_value()?;
1829 }
1830 }
1831 *validity = new_validity;
1832 Ok(())
1833 }
1834
1835 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1849 if self.auto_inc.is_none() {
1850 return Ok(None);
1851 }
1852 Ok(Some(self.alloc_auto_inc_value()?))
1853 }
1854
1855 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1861 let payload = bincode::serialize(&rows)?;
1862 self.wal_append_data(Op::Put {
1863 table_id: self.table_id,
1864 rows: payload,
1865 })?;
1866 if self.is_shared() {
1867 self.pending_rows_auto_inc
1868 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1869 self.pending_rows.extend(rows);
1870 } else {
1871 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1872 }
1873 Ok(())
1874 }
1875
1876 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1881 self.apply_put_rows_inner(rows, true)
1882 }
1883
1884 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1885 if check_existing_pk {
1886 self.ensure_indexes_complete()?;
1887 }
1888 if rows.len() == 1 {
1892 let row = rows.into_iter().next().expect("len checked");
1893 return self.apply_put_row_single(row, check_existing_pk);
1894 }
1895 let pk_id = self.schema.primary_key().map(|c| c.id);
1912 let probe = match pk_id {
1913 Some(pid) => {
1914 check_existing_pk
1915 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
1916 }
1917 None => false,
1918 };
1919 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
1922 for r in rows {
1923 for &cid in r.columns.keys() {
1924 self.pending_put_cols.insert(cid);
1925 }
1926 match pk_id {
1927 Some(pid) if probe || maintain_pk_by_row => {
1928 if let Some(pk_val) = r.columns.get(&pid) {
1929 let key = self.index_lookup_key(pid, pk_val);
1930 if probe {
1931 if let Some(old_rid) = self.hot.get(&key) {
1932 if old_rid != r.row_id {
1933 self.tombstone_row(old_rid, r.committed_epoch, true);
1934 }
1935 }
1936 }
1937 if maintain_pk_by_row {
1938 self.pk_by_row.insert(r.row_id, key);
1939 }
1940 }
1941 }
1942 Some(_) => {}
1943 None => {
1944 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1945 }
1946 }
1947 self.index_row(&r);
1948 self.reservoir.offer(r.row_id.0);
1949 self.memtable.upsert(r);
1950 self.live_count = self.live_count.saturating_add(1);
1953 }
1954 Ok(())
1955 }
1956
1957 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
1961 for &cid in row.columns.keys() {
1962 self.pending_put_cols.insert(cid);
1963 }
1964 let epoch = row.committed_epoch;
1965 if let Some(pk_col) = self.schema.primary_key() {
1966 let pk_id = pk_col.id;
1967 if let Some(pk_val) = row.columns.get(&pk_id) {
1968 let maintain_pk_by_row = self.pk_by_row_complete;
1972 if check_existing_pk || maintain_pk_by_row {
1973 let key = self.index_lookup_key(pk_id, pk_val);
1974 if check_existing_pk {
1975 if let Some(old_rid) = self.hot.get(&key) {
1976 if old_rid != row.row_id {
1977 self.tombstone_row(old_rid, epoch, true);
1978 }
1979 }
1980 }
1981 if maintain_pk_by_row {
1982 self.pk_by_row.insert(row.row_id, key);
1983 }
1984 }
1985 }
1986 } else {
1987 self.hot
1988 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
1989 }
1990 self.index_row(&row);
1991 self.reservoir.offer(row.row_id.0);
1992 self.memtable.upsert(row);
1993 self.live_count = self.live_count.saturating_add(1);
1994 Ok(())
1995 }
1996
1997 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2000 self.allocator.alloc()
2001 }
2002
2003 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2011 self.ensure_indexes_complete()?;
2012 let n = rows.len();
2013 for r in rows {
2014 for &cid in r.columns.keys() {
2015 self.pending_put_cols.insert(cid);
2016 }
2017 }
2018 let (losers, winner_pks) = self.partition_pk_winners(rows);
2019 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2020 for (key, &row_id) in &winner_pks {
2022 if let Some(old_rid) = self.hot.get(key) {
2023 if old_rid != row_id {
2024 self.tombstone_row(old_rid, epoch, true);
2025 }
2026 }
2027 }
2028 for &loser_rid in &losers {
2031 self.tombstone_row(loser_rid, epoch, false);
2032 }
2033 for (key, row_id) in winner_pks {
2035 self.insert_hot_pk(key, row_id);
2036 }
2037 if self.schema.primary_key().is_none() {
2038 for r in rows {
2039 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2040 }
2041 }
2042 for r in rows {
2043 self.allocator.advance_to(r.row_id);
2044 if !losers.contains(&r.row_id) {
2045 self.index_row(r);
2046 }
2047 }
2048 for r in rows {
2049 if !losers.contains(&r.row_id) {
2050 self.reservoir.offer(r.row_id.0);
2051 }
2052 }
2053 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2054 Ok(())
2055 }
2056
2057 pub(crate) fn recover_apply(
2062 &mut self,
2063 rows: Vec<Row>,
2064 deletes: Vec<(RowId, Epoch)>,
2065 ) -> Result<()> {
2066 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2070 std::collections::BTreeMap::new();
2071 for row in rows {
2072 self.allocator.advance_to(row.row_id);
2073 if let Some(ai) = self.auto_inc.as_mut() {
2078 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2079 if *n + 1 > ai.next {
2080 ai.next = *n + 1;
2081 }
2082 }
2083 }
2084 by_epoch.entry(row.committed_epoch).or_default().push(row);
2085 }
2086 for (epoch, group) in by_epoch {
2087 let (losers, winner_pks) = self.partition_pk_winners(&group);
2088 for (key, &row_id) in &winner_pks {
2090 if let Some(old_rid) = self.hot.get(key) {
2091 if old_rid != row_id {
2092 self.tombstone_row(old_rid, epoch, false);
2093 }
2094 }
2095 }
2096 for (key, row_id) in winner_pks {
2097 self.insert_hot_pk(key, row_id);
2098 }
2099 if self.schema.primary_key().is_none() {
2100 for r in &group {
2101 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2102 }
2103 }
2104 for r in &group {
2105 if !losers.contains(&r.row_id) {
2106 self.memtable.upsert(r.clone());
2107 self.index_row(r);
2108 }
2109 }
2110 }
2111 for (rid, epoch) in deletes {
2112 self.memtable.tombstone(rid, epoch);
2113 self.remove_hot_for_row(rid, epoch);
2114 }
2115 let _ = self.rebuild_reservoir();
2116 Ok(())
2117 }
2118
2119 pub(crate) fn flushed_epoch(&self) -> u64 {
2121 self.flushed_epoch
2122 }
2123
2124 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2126 self.schema.validate_not_null(cells)
2127 }
2128
2129 fn validate_columns_not_null(
2133 &self,
2134 columns: &[(u16, columnar::NativeColumn)],
2135 n: usize,
2136 ) -> Result<()> {
2137 let by_id: HashMap<u16, &columnar::NativeColumn> =
2138 columns.iter().map(|(id, c)| (*id, c)).collect();
2139 for col in &self.schema.columns {
2140 if col.flags.contains(ColumnFlags::NULLABLE) {
2141 continue;
2142 }
2143 match by_id.get(&col.id) {
2144 None => {
2145 return Err(MongrelError::InvalidArgument(format!(
2146 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2147 col.name, col.id
2148 )));
2149 }
2150 Some(c) => {
2151 if c.null_count(n) != 0 {
2152 return Err(MongrelError::InvalidArgument(format!(
2153 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2154 col.name, col.id
2155 )));
2156 }
2157 }
2158 }
2159 }
2160 Ok(())
2161 }
2162
2163 fn bulk_pk_winner_indices(
2168 &self,
2169 columns: &[(u16, columnar::NativeColumn)],
2170 n: usize,
2171 ) -> Option<Vec<usize>> {
2172 let pk_col = self.schema.primary_key()?;
2173 let pk_id = pk_col.id;
2174 let pk_ty = pk_col.ty;
2175 let by_id: HashMap<u16, &columnar::NativeColumn> =
2176 columns.iter().map(|(id, c)| (*id, c)).collect();
2177 let pk_native = by_id.get(&pk_id)?;
2178 if native_int64_strictly_increasing(pk_native, n) {
2179 return None;
2180 }
2181 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2183 let mut null_pk_rows: Vec<usize> = Vec::new();
2184 for i in 0..n {
2185 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2186 Some(key) => {
2187 last.insert(key, i);
2188 }
2189 None => null_pk_rows.push(i),
2190 }
2191 }
2192 let mut winners: HashSet<usize> = last.values().copied().collect();
2193 for i in null_pk_rows {
2194 winners.insert(i);
2195 }
2196 Some((0..n).filter(|i| winners.contains(i)).collect())
2197 }
2198
2199 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2201 let epoch = self.pending_epoch();
2202 self.wal_append_data(Op::Delete {
2203 table_id: self.table_id,
2204 row_ids: vec![row_id],
2205 })?;
2206 if self.is_shared() {
2207 self.pending_dels.push(row_id);
2208 } else {
2209 self.apply_delete(row_id, epoch);
2210 }
2211 Ok(())
2212 }
2213
2214 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2215 let pre = self.get(row_id, self.snapshot());
2216 self.delete(row_id)?;
2217 Ok(pre.map(|row| {
2218 let mut columns: Vec<_> = row.columns.into_iter().collect();
2219 columns.sort_by_key(|(id, _)| *id);
2220 OwnedRow { columns }
2221 }))
2222 }
2223
2224 pub fn truncate(&mut self) -> Result<()> {
2226 let epoch = self.pending_epoch();
2227 self.wal_append_data(Op::TruncateTable {
2228 table_id: self.table_id,
2229 })?;
2230 self.pending_rows.clear();
2231 self.pending_rows_auto_inc.clear();
2232 self.pending_dels.clear();
2233 self.pending_truncate = Some(epoch);
2234 Ok(())
2235 }
2236
2237 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2239 for rr in std::mem::take(&mut self.run_refs) {
2240 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2241 }
2242 for r in std::mem::take(&mut self.retiring) {
2243 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2244 }
2245 self.memtable = Memtable::new();
2246 self.mutable_run = MutableRun::new();
2247 self.hot = HotIndex::new();
2248 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2249 self.bitmap = bitmap;
2250 self.ann = ann;
2251 self.fm = fm;
2252 self.sparse = sparse;
2253 self.minhash = minhash;
2254 self.learned_range.clear();
2255 self.pk_by_row.clear();
2256 self.pk_by_row_complete = false;
2257 self.live_count = 0;
2258 self.reservoir = crate::reservoir::Reservoir::default();
2259 self.had_deletes = true;
2260 self.agg_cache.clear();
2261 self.global_idx_epoch = 0;
2262 self.indexes_complete = true;
2263 self.pending_delete_rids.clear();
2264 self.pending_put_cols.clear();
2265 self.pending_rows.clear();
2266 self.pending_rows_auto_inc.clear();
2267 self.pending_dels.clear();
2268 self.clear_result_cache();
2269 self.invalidate_index_checkpoint();
2270 Ok(())
2271 }
2272
2273 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2276 self.remove_hot_for_row(row_id, epoch);
2277 self.tombstone_row(row_id, epoch, true);
2278 }
2279
2280 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2284 let tombstone = Row {
2285 row_id,
2286 committed_epoch: epoch,
2287 columns: std::collections::HashMap::new(),
2288 deleted: true,
2289 };
2290 self.memtable.upsert(tombstone);
2291 self.pk_by_row.remove(&row_id);
2292 if adjust_live_count {
2293 self.live_count = self.live_count.saturating_sub(1);
2294 }
2295 self.pending_delete_rids.insert(row_id.0 as u32);
2297 self.had_deletes = true;
2300 self.agg_cache.clear();
2301 }
2302
2303 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2307 if !self.pk_by_row_complete {
2310 self.refresh_pk_by_row_from_hot();
2311 }
2312 let Some(pk_col) = self.schema.primary_key() else {
2313 return;
2314 };
2315 if let Some(key) = self.pk_by_row.remove(&row_id) {
2316 if self.hot.get(&key) == Some(row_id) {
2317 self.hot.remove(&key);
2318 }
2319 return;
2320 }
2321 if !self.indexes_complete {
2322 return;
2323 }
2324 let pk_val = self
2327 .memtable
2328 .get_version(row_id, epoch)
2329 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2330 .or_else(|| {
2331 self.mutable_run
2332 .get_version(row_id, epoch)
2333 .filter(|(_, r)| !r.deleted)
2334 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2335 })
2336 .or_else(|| {
2337 self.run_refs.iter().find_map(|rr| {
2338 let mut reader = self.open_reader(rr.run_id).ok()?;
2339 let (_, r) = reader.get_version(row_id, epoch).ok()??;
2340 if r.deleted {
2341 return None;
2342 }
2343 r.columns.get(&pk_col.id).cloned()
2344 })
2345 });
2346 if let Some(pk_val) = pk_val {
2347 let key = self.index_lookup_key(pk_col.id, &pk_val);
2348 if self.hot.get(&key) == Some(row_id) {
2349 self.hot.remove(&key);
2350 }
2351 }
2352 }
2353
2354 fn partition_pk_winners(
2359 &self,
2360 rows: &[Row],
2361 ) -> (
2362 std::collections::HashSet<RowId>,
2363 std::collections::HashMap<Vec<u8>, RowId>,
2364 ) {
2365 let mut losers = std::collections::HashSet::new();
2366 let Some(pk_col) = self.schema.primary_key() else {
2367 return (losers, std::collections::HashMap::new());
2368 };
2369 let pk_id = pk_col.id;
2370 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2371 std::collections::HashMap::new();
2372 for r in rows {
2373 let Some(pk_val) = r.columns.get(&pk_id) else {
2374 continue;
2375 };
2376 let key = self.index_lookup_key(pk_id, pk_val);
2377 if let Some(&old_rid) = winners.get(&key) {
2378 losers.insert(old_rid);
2379 }
2380 winners.insert(key, r.row_id);
2381 }
2382 (losers, winners)
2383 }
2384
2385 fn index_row(&mut self, row: &Row) {
2386 if row.deleted {
2387 return;
2388 }
2389 if self.column_keys.is_empty() {
2393 index_into(
2394 &self.schema,
2395 row,
2396 &mut self.hot,
2397 &mut self.bitmap,
2398 &mut self.ann,
2399 &mut self.fm,
2400 &mut self.sparse,
2401 &mut self.minhash,
2402 );
2403 return;
2404 }
2405 let effective_row = self.tokenized_for_indexes(row);
2406 index_into(
2407 &self.schema,
2408 &effective_row,
2409 &mut self.hot,
2410 &mut self.bitmap,
2411 &mut self.ann,
2412 &mut self.fm,
2413 &mut self.sparse,
2414 &mut self.minhash,
2415 );
2416 }
2417
2418 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2424 if self.column_keys.is_empty() {
2425 return row.clone();
2426 }
2427 #[cfg(feature = "encryption")]
2428 {
2429 use crate::encryption::SCHEME_HMAC_EQ;
2430 let mut tok = row.clone();
2431 for (&cid, &(_, scheme)) in &self.column_keys {
2432 if scheme != SCHEME_HMAC_EQ {
2433 continue;
2434 }
2435 if let Some(v) = tok.columns.get(&cid).cloned() {
2436 if let Some(t) = self.tokenize_value(cid, &v) {
2437 tok.columns.insert(cid, t);
2438 }
2439 }
2440 }
2441 tok
2442 }
2443 #[cfg(not(feature = "encryption"))]
2444 {
2445 row.clone()
2446 }
2447 }
2448
2449 pub fn commit(&mut self) -> Result<Epoch> {
2454 if self.is_shared() {
2455 self.commit_shared()
2456 } else {
2457 self.commit_private()
2458 }
2459 }
2460
2461 fn commit_private(&mut self) -> Result<Epoch> {
2463 let commit_lock = Arc::clone(&self.commit_lock);
2467 let _g = commit_lock.lock();
2468 let new_epoch = self.epoch.bump_assigned();
2469 let txn_id = self.current_txn_id;
2470 match &mut self.wal {
2474 WalSink::Private(w) => {
2475 w.append_txn(
2476 txn_id,
2477 Op::TxnCommit {
2478 epoch: new_epoch.0,
2479 added_runs: Vec::new(),
2480 },
2481 )?;
2482 w.sync()?;
2483 }
2484 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2485 }
2486 if let Some(epoch) = self.pending_truncate.take() {
2488 self.apply_truncate(epoch)?;
2489 }
2490 self.invalidate_pending_cache();
2491 self.persist_manifest(new_epoch)?;
2492 self.epoch.publish_in_order(new_epoch);
2496 self.current_txn_id += 1;
2497 Ok(new_epoch)
2498 }
2499
2500 fn commit_shared(&mut self) -> Result<Epoch> {
2506 use std::sync::atomic::Ordering;
2507 let s = match &self.wal {
2508 WalSink::Shared(s) => s.clone(),
2509 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2510 };
2511 if s.poisoned.load(Ordering::Relaxed) {
2512 return Err(MongrelError::Other(
2513 "database poisoned by fsync error".into(),
2514 ));
2515 }
2516 let commit_lock = Arc::clone(&self.commit_lock);
2523 let _g = commit_lock.lock();
2524 let txn_id = self.ensure_txn_id();
2527 let (new_epoch, commit_seq) = {
2528 let mut wal = s.wal.lock();
2529 let new_epoch = self.epoch.bump_assigned();
2530 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2531 (new_epoch, seq)
2532 };
2533 s.group
2534 .await_durable(&s.wal, commit_seq)
2535 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2536
2537 if self.pending_truncate.take().is_some() {
2540 self.apply_truncate(new_epoch)?;
2541 }
2542 let mut rows = std::mem::take(&mut self.pending_rows);
2543 if !rows.is_empty() {
2544 for r in &mut rows {
2545 r.committed_epoch = new_epoch;
2546 }
2547 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2548 let all_auto_generated =
2549 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2550 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2551 } else {
2552 self.pending_rows_auto_inc.clear();
2553 }
2554 let dels = std::mem::take(&mut self.pending_dels);
2555 for rid in dels {
2556 self.apply_delete(rid, new_epoch);
2557 }
2558
2559 self.invalidate_pending_cache();
2560 self.persist_manifest(new_epoch)?;
2561 self.epoch.publish_in_order(new_epoch);
2562 self.current_txn_id = 0;
2564 Ok(new_epoch)
2565 }
2566
2567 pub fn flush(&mut self) -> Result<Epoch> {
2575 self.ensure_indexes_complete()?;
2576 let epoch = self.commit()?;
2577 let rows = self.memtable.drain_sorted();
2578 if !rows.is_empty() {
2579 self.mutable_run.insert_many(rows);
2580 }
2581 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2582 self.spill_mutable_run(epoch)?;
2583 self.mark_flushed(epoch)?;
2587 self.persist_manifest(epoch)?;
2588 self.build_learned_ranges()?;
2589 self.checkpoint_indexes(epoch);
2592 }
2593 Ok(epoch)
2596 }
2597
2598 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2605 let op = Op::Flush {
2606 table_id: self.table_id,
2607 flushed_epoch: epoch.0,
2608 };
2609 match &mut self.wal {
2610 WalSink::Private(w) => {
2611 w.append_system(op)?;
2612 w.sync()?;
2613 }
2614 WalSink::Shared(s) => {
2615 s.wal.lock().append_system(op)?;
2620 }
2621 }
2622 self.flushed_epoch = epoch.0;
2623 if matches!(self.wal, WalSink::Private(_)) {
2624 self.rotate_wal(epoch)?;
2625 }
2626 Ok(())
2627 }
2628
2629 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2633 let rows = self.mutable_run.drain_sorted();
2634 if rows.is_empty() {
2635 return Ok(());
2636 }
2637 let run_id = self.next_run_id;
2638 self.next_run_id += 1;
2639 let path = self.run_path(run_id);
2640 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2641 if let Some(kek) = &self.kek {
2642 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2643 }
2644 let header = writer.write(&path, &rows)?;
2645 self.run_refs.push(RunRef {
2646 run_id: run_id as u128,
2647 level: 0,
2648 epoch_created: epoch.0,
2649 row_count: header.row_count,
2650 });
2651 Ok(())
2652 }
2653
2654 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2658 self.mutable_run_spill_bytes = bytes.max(1);
2659 }
2660
2661 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2665 self.compaction_zstd_level = level;
2666 }
2667
2668 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2672 self.result_cache.lock().set_max_bytes(max_bytes);
2673 }
2674
2675 pub(crate) fn clear_result_cache(&mut self) {
2679 self.result_cache.lock().clear();
2680 }
2681
2682 pub fn mutable_run_len(&self) -> usize {
2684 self.mutable_run.len()
2685 }
2686
2687 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2690 self.mutable_run.drain_sorted()
2691 }
2692
2693 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2698 let epoch = self.commit()?;
2699 let n = batch.len();
2700 if n == 0 {
2701 return Ok(epoch);
2702 }
2703 let live_before = self.live_count;
2704 self.spill_mutable_run(epoch)?;
2708 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
2709 && self.indexes_complete
2710 && self.run_refs.is_empty()
2711 && self.memtable.is_empty()
2712 && self.mutable_run.is_empty();
2713 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2719 use rayon::prelude::*;
2720 self.schema
2721 .columns
2722 .par_iter()
2723 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
2724 .collect::<Vec<_>>()
2725 };
2726 drop(batch);
2727 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2732 self.validate_columns_not_null(&user_columns, n)?;
2733 let winner_idx = self
2734 .bulk_pk_winner_indices(&user_columns, n)
2735 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2736 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2737 match winner_idx.as_deref() {
2738 Some(idx) => {
2739 let compacted = user_columns
2740 .iter()
2741 .map(|(id, c)| (*id, c.gather(idx)))
2742 .collect();
2743 (compacted, idx.len())
2744 }
2745 None => (std::mem::take(&mut user_columns), n),
2746 };
2747 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2748 let first = self.allocator.alloc_range(write_n as u64).0;
2749 for rid in first..first + write_n as u64 {
2750 self.reservoir.offer(rid);
2751 }
2752 let run_id = self.next_run_id;
2753 self.next_run_id += 1;
2754 let path = self.run_path(run_id);
2755 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2756 .clean(true)
2757 .with_lz4()
2758 .with_native_endian();
2759 if let Some(kek) = &self.kek {
2760 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2761 }
2762 let header = writer.write_native(&path, &write_columns, write_n, first)?;
2763 self.run_refs.push(RunRef {
2764 run_id: run_id as u128,
2765 level: 0,
2766 epoch_created: epoch.0,
2767 row_count: header.row_count,
2768 });
2769 self.live_count = self.live_count.saturating_add(write_n as u64);
2770 if eager_index_build {
2771 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
2772 self.index_columns_bulk(&write_columns, &row_ids);
2773 self.indexes_complete = true;
2774 self.build_learned_ranges()?;
2775 } else {
2776 self.indexes_complete = false;
2777 }
2778 self.mark_flushed(epoch)?;
2779 self.persist_manifest(epoch)?;
2780 if eager_index_build {
2781 self.checkpoint_indexes(epoch);
2782 }
2783 self.clear_result_cache();
2784 Ok(epoch)
2785 }
2786
2787 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
2790 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
2791 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
2792 let segment_no = segment
2795 .file_stem()
2796 .and_then(|s| s.to_str())
2797 .and_then(|s| s.strip_prefix("seg-"))
2798 .and_then(|s| s.parse::<u64>().ok())
2799 .unwrap_or(0);
2800 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
2801 wal.set_sync_byte_threshold(self.sync_byte_threshold);
2802 wal.sync()?;
2803 self.wal = WalSink::Private(wal);
2804 Ok(())
2805 }
2806
2807 pub(crate) fn invalidate_pending_cache(&mut self) {
2812 self.result_cache
2813 .lock()
2814 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
2815 self.pending_delete_rids.clear();
2816 self.pending_put_cols.clear();
2817 }
2818
2819 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
2820 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
2821 m.current_epoch = epoch.0;
2822 m.next_row_id = self.allocator.current().0;
2823 m.runs = self.run_refs.clone();
2824 m.live_count = self.live_count;
2825 m.global_idx_epoch = self.global_idx_epoch;
2826 m.flushed_epoch = self.flushed_epoch;
2827 m.retiring = self.retiring.clone();
2828 m.auto_inc_next = match self.auto_inc {
2832 Some(ai) if ai.seeded => ai.next,
2833 _ => 0,
2834 };
2835 let meta_dek = self.manifest_meta_dek();
2836 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
2837 Ok(())
2838 }
2839
2840 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
2846 if !self.indexes_complete {
2849 return;
2850 }
2851 let snap = global_idx::IndexSnapshot {
2852 hot: &self.hot,
2853 bitmap: &self.bitmap,
2854 ann: &self.ann,
2855 fm: &self.fm,
2856 sparse: &self.sparse,
2857 minhash: &self.minhash,
2858 learned_range: &self.learned_range,
2859 };
2860 let idx_dek = self.idx_dek();
2862 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
2863 .is_ok()
2864 {
2865 self.global_idx_epoch = epoch.0;
2866 let _ = self.persist_manifest(epoch);
2867 }
2868 }
2869
2870 pub(crate) fn invalidate_index_checkpoint(&mut self) {
2873 self.global_idx_epoch = 0;
2874 global_idx::remove(&self.dir);
2875 let _ = self.persist_manifest(self.epoch.visible());
2876 }
2877
2878 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
2881 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
2882 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
2883 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2884 best = Some((epoch, row));
2885 }
2886 }
2887 for rr in &self.run_refs {
2888 let Ok(mut reader) = self.open_reader(rr.run_id) else {
2889 continue;
2890 };
2891 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
2892 continue;
2893 };
2894 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2895 best = Some((epoch, row));
2896 }
2897 }
2898 match best {
2899 Some((_, r)) if r.deleted => None,
2900 Some((_, r)) => Some(r),
2901 None => None,
2902 }
2903 }
2904
2905 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
2909 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
2910 let mut fold = |row: Row| {
2911 best.entry(row.row_id.0)
2912 .and_modify(|e| {
2913 if row.committed_epoch > e.0 {
2914 *e = (row.committed_epoch, row.clone());
2915 }
2916 })
2917 .or_insert_with(|| (row.committed_epoch, row));
2918 };
2919 for row in self.memtable.visible_versions(snapshot.epoch) {
2920 fold(row);
2921 }
2922 for row in self.mutable_run.visible_versions(snapshot.epoch) {
2923 fold(row);
2924 }
2925 for rr in &self.run_refs {
2926 let mut reader = self.open_reader(rr.run_id)?;
2927 for row in reader.visible_versions(snapshot.epoch)? {
2928 fold(row);
2929 }
2930 }
2931 let mut out: Vec<Row> = best
2932 .into_values()
2933 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
2934 .collect();
2935 out.sort_by_key(|r| r.row_id);
2936 Ok(out)
2937 }
2938
2939 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
2946 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
2947 let rr = self.run_refs[0].clone();
2948 let mut reader = self.open_reader(rr.run_id)?;
2949 let idxs = reader.visible_indices(snapshot.epoch)?;
2950 let mut cols = Vec::with_capacity(self.schema.columns.len());
2951 for cdef in &self.schema.columns {
2952 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
2953 }
2954 return Ok(cols);
2955 }
2956 let rows = self.visible_rows(snapshot)?;
2958 let mut cols: Vec<(u16, Vec<Value>)> = self
2959 .schema
2960 .columns
2961 .iter()
2962 .map(|c| (c.id, Vec::with_capacity(rows.len())))
2963 .collect();
2964 for r in &rows {
2965 for (cid, vec) in cols.iter_mut() {
2966 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
2967 }
2968 }
2969 Ok(cols)
2970 }
2971
2972 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
2974 self.hot.get(key)
2975 }
2976
2977 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
2982 self.ensure_indexes_complete()?;
2983 let snapshot = self.snapshot();
2984 crate::trace::QueryTrace::record(|t| {
2985 t.run_count = self.run_refs.len();
2986 t.memtable_rows = self.memtable.len();
2987 t.mutable_run_rows = self.mutable_run.len();
2988 });
2989 if q.conditions.is_empty() {
2993 crate::trace::QueryTrace::record(|t| {
2994 t.scan_mode = crate::trace::ScanMode::Materialized;
2995 t.row_materialized = true;
2996 });
2997 return self.visible_rows(snapshot);
2998 }
2999 crate::trace::QueryTrace::record(|t| {
3000 t.conditions_pushed = q.conditions.len();
3001 t.scan_mode = crate::trace::ScanMode::Materialized;
3002 t.row_materialized = true;
3003 });
3004 let mut sets: Vec<RowIdSet> = Vec::with_capacity(q.conditions.len());
3005 for c in &q.conditions {
3006 sets.push(self.resolve_condition(c, snapshot)?);
3007 }
3008 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3009 self.rows_for_rids(&rids, snapshot)
3010 }
3011
3012 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3017 use std::collections::HashMap;
3018 let mut rows = Vec::with_capacity(rids.len());
3019 let mut overlay: HashMap<u64, Row> = HashMap::new();
3025 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3026 overlay
3027 .entry(row.row_id.0)
3028 .and_modify(|e| {
3029 if row.committed_epoch > e.committed_epoch {
3030 *e = row.clone();
3031 }
3032 })
3033 .or_insert(row);
3034 };
3035 for row in self.memtable.visible_versions(snapshot.epoch) {
3036 fold_newest(row, &mut overlay);
3037 }
3038 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3039 fold_newest(row, &mut overlay);
3040 }
3041 if self.run_refs.len() == 1 {
3042 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3051 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3052 enum Src {
3055 Overlay,
3056 Run,
3057 }
3058 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3059 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3060 for rid in rids {
3061 if overlay.contains_key(rid) {
3062 plan.push(Src::Overlay);
3063 continue;
3064 }
3065 match vis_rids.binary_search(&(*rid as i64)) {
3066 Ok(i) => {
3067 plan.push(Src::Run);
3068 fetch.push(positions[i]);
3069 }
3070 Err(_) => { }
3071 }
3072 }
3073 let fetched = reader.materialize_batch(&fetch)?;
3074 let mut fetched_iter = fetched.into_iter();
3075 for (rid, src) in rids.iter().zip(plan) {
3076 match src {
3077 Src::Overlay => {
3078 if let Some(r) = overlay.get(rid) {
3079 if !r.deleted {
3080 rows.push(r.clone());
3081 }
3082 }
3083 }
3084 Src::Run => {
3085 if let Some(row) = fetched_iter.next() {
3086 if !row.deleted {
3087 rows.push(row);
3088 }
3089 }
3090 }
3091 }
3092 }
3093 return Ok(rows);
3094 }
3095 let mut readers: Vec<_> = self
3099 .run_refs
3100 .iter()
3101 .map(|rr| self.open_reader(rr.run_id))
3102 .collect::<Result<Vec<_>>>()?;
3103 for rid in rids {
3104 if let Some(r) = overlay.get(rid) {
3105 if !r.deleted {
3106 rows.push(r.clone());
3107 }
3108 continue;
3109 }
3110 let mut best: Option<(Epoch, Row)> = None;
3111 for reader in readers.iter_mut() {
3112 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3113 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3114 best = Some((epoch, row));
3115 }
3116 }
3117 }
3118 if let Some((_, r)) = best {
3119 if !r.deleted {
3120 rows.push(r);
3121 }
3122 }
3123 }
3124 Ok(rows)
3125 }
3126
3127 pub fn indexes_complete(&self) -> bool {
3137 self.indexes_complete
3138 }
3139
3140 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3142 self.index_build_policy
3143 }
3144
3145 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3149 self.index_build_policy = policy;
3150 }
3151
3152 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3157 if !self.indexes_complete {
3161 return None;
3162 }
3163 let b = self.bitmap.get(&column_id)?;
3164 let result: Vec<Vec<u8>> = b
3165 .keys()
3166 .into_iter()
3167 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3168 .cloned()
3169 .collect();
3170 Some(result)
3171 }
3172
3173 pub fn fk_join_row_ids(
3174 &self,
3175 fk_column_id: u16,
3176 pk_values: &[Vec<u8>],
3177 fk_conditions: &[crate::query::Condition],
3178 snapshot: Snapshot,
3179 ) -> Result<Vec<u64>> {
3180 let Some(b) = self.bitmap.get(&fk_column_id) else {
3181 return Ok(Vec::new());
3182 };
3183 let mut join_set = {
3184 let mut acc = roaring::RoaringBitmap::new();
3185 for v in pk_values {
3186 acc |= b.get(v);
3187 }
3188 RowIdSet::from_roaring(acc)
3189 };
3190 if !fk_conditions.is_empty() {
3191 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3192 sets.push(join_set);
3193 for c in fk_conditions {
3194 sets.push(self.resolve_condition(c, snapshot)?);
3195 }
3196 join_set = RowIdSet::intersect_many(sets);
3197 }
3198 Ok(join_set.into_sorted_vec())
3199 }
3200
3201 pub fn fk_join_count(
3207 &self,
3208 fk_column_id: u16,
3209 pk_values: &[Vec<u8>],
3210 fk_conditions: &[crate::query::Condition],
3211 snapshot: Snapshot,
3212 ) -> Result<u64> {
3213 let Some(b) = self.bitmap.get(&fk_column_id) else {
3214 return Ok(0);
3215 };
3216 let mut acc = roaring::RoaringBitmap::new();
3217 for v in pk_values {
3218 acc |= b.get(v);
3219 }
3220 if fk_conditions.is_empty() {
3221 return Ok(acc.len());
3222 }
3223 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3224 sets.push(RowIdSet::from_roaring(acc));
3225 for c in fk_conditions {
3226 sets.push(self.resolve_condition(c, snapshot)?);
3227 }
3228 Ok(RowIdSet::intersect_many(sets).len() as u64)
3229 }
3230
3231 fn resolve_condition(
3236 &self,
3237 c: &crate::query::Condition,
3238 snapshot: Snapshot,
3239 ) -> Result<RowIdSet> {
3240 use crate::query::Condition;
3241 Ok(match c {
3242 Condition::Pk(key) => {
3243 let lookup = self
3244 .schema
3245 .primary_key()
3246 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3247 .unwrap_or_else(|| key.clone());
3248 self.hot
3249 .get(&lookup)
3250 .map(|r| RowIdSet::one(r.0))
3251 .unwrap_or_else(RowIdSet::empty)
3252 }
3253 Condition::BitmapEq { column_id, value } => {
3254 let lookup = self.index_lookup_key_bytes(*column_id, value);
3255 self.bitmap
3256 .get(column_id)
3257 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3258 .unwrap_or_else(RowIdSet::empty)
3259 }
3260 Condition::BitmapIn { column_id, values } => {
3261 let bm = self.bitmap.get(column_id);
3262 let mut acc = roaring::RoaringBitmap::new();
3263 if let Some(b) = bm {
3264 for v in values {
3265 let lookup = self.index_lookup_key_bytes(*column_id, v);
3266 acc |= b.get(&lookup);
3267 }
3268 }
3269 RowIdSet::from_roaring(acc)
3270 }
3271 Condition::FmContains { column_id, pattern } => self
3272 .fm
3273 .get(column_id)
3274 .map(|f| {
3275 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3276 })
3277 .unwrap_or_else(RowIdSet::empty),
3278 Condition::FmContainsAll {
3279 column_id,
3280 patterns,
3281 } => {
3282 if let Some(f) = self.fm.get(column_id) {
3285 let sets: Vec<RowIdSet> = patterns
3286 .iter()
3287 .map(|pat| {
3288 RowIdSet::from_unsorted(
3289 f.locate(pat).into_iter().map(|r| r.0).collect(),
3290 )
3291 })
3292 .collect();
3293 RowIdSet::intersect_many(sets)
3294 } else {
3295 RowIdSet::empty()
3296 }
3297 }
3298 Condition::Ann {
3299 column_id,
3300 query,
3301 k,
3302 } => self
3303 .ann
3304 .get(column_id)
3305 .map(|a| {
3306 RowIdSet::from_unsorted(
3307 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3308 )
3309 })
3310 .unwrap_or_else(RowIdSet::empty),
3311 Condition::SparseMatch {
3312 column_id,
3313 query,
3314 k,
3315 } => self
3316 .sparse
3317 .get(column_id)
3318 .map(|s| {
3319 RowIdSet::from_unsorted(
3320 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3321 )
3322 })
3323 .unwrap_or_else(RowIdSet::empty),
3324 Condition::MinHashSimilar {
3325 column_id,
3326 query,
3327 k,
3328 } => self
3329 .minhash
3330 .get(column_id)
3331 .map(|mh| {
3332 RowIdSet::from_unsorted(
3333 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3334 )
3335 })
3336 .unwrap_or_else(RowIdSet::empty),
3337 Condition::Range { column_id, lo, hi } => {
3338 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3347 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3348 } else if self.run_refs.len() == 1 {
3349 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3350 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3351 } else {
3352 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3353 };
3354 set.remove_many(self.overlay_rid_set(snapshot));
3355 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3356 set
3357 }
3358 Condition::RangeF64 {
3359 column_id,
3360 lo,
3361 lo_inclusive,
3362 hi,
3363 hi_inclusive,
3364 } => {
3365 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3368 RowIdSet::from_unsorted(
3369 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3370 .into_iter()
3371 .collect(),
3372 )
3373 } else if self.run_refs.len() == 1 {
3374 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3375 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3376 } else {
3377 return self.range_scan_f64(
3378 *column_id,
3379 *lo,
3380 *lo_inclusive,
3381 *hi,
3382 *hi_inclusive,
3383 snapshot,
3384 );
3385 };
3386 set.remove_many(self.overlay_rid_set(snapshot));
3387 self.range_scan_overlay_f64(
3388 &mut set,
3389 *column_id,
3390 *lo,
3391 *lo_inclusive,
3392 *hi,
3393 *hi_inclusive,
3394 snapshot,
3395 );
3396 set
3397 }
3398 Condition::IsNull { column_id } => {
3399 let mut set = if self.run_refs.len() == 1 {
3400 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3401 r.null_row_id_set(*column_id, true)?
3402 } else {
3403 return self.null_scan(*column_id, true, snapshot);
3404 };
3405 set.remove_many(self.overlay_rid_set(snapshot));
3406 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3407 set
3408 }
3409 Condition::IsNotNull { column_id } => {
3410 let mut set = if self.run_refs.len() == 1 {
3411 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3412 r.null_row_id_set(*column_id, false)?
3413 } else {
3414 return self.null_scan(*column_id, false, snapshot);
3415 };
3416 set.remove_many(self.overlay_rid_set(snapshot));
3417 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3418 set
3419 }
3420 })
3421 }
3422
3423 fn range_scan_i64(
3431 &self,
3432 column_id: u16,
3433 lo: i64,
3434 hi: i64,
3435 snapshot: Snapshot,
3436 ) -> Result<RowIdSet> {
3437 let mut row_ids = Vec::new();
3438 let overlay_rids = self.overlay_rid_set(snapshot);
3439 for rr in &self.run_refs {
3440 let mut reader = self.open_reader(rr.run_id)?;
3441 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3442 for rid in matched {
3443 if !overlay_rids.contains(&rid) {
3444 row_ids.push(rid);
3445 }
3446 }
3447 }
3448 let mut s = RowIdSet::from_unsorted(row_ids);
3449 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3450 Ok(s)
3451 }
3452
3453 fn range_scan_f64(
3456 &self,
3457 column_id: u16,
3458 lo: f64,
3459 lo_inclusive: bool,
3460 hi: f64,
3461 hi_inclusive: bool,
3462 snapshot: Snapshot,
3463 ) -> Result<RowIdSet> {
3464 let mut row_ids = Vec::new();
3465 let overlay_rids = self.overlay_rid_set(snapshot);
3466 for rr in &self.run_refs {
3467 let mut reader = self.open_reader(rr.run_id)?;
3468 let matched = reader.range_row_ids_visible_f64(
3469 column_id,
3470 lo,
3471 lo_inclusive,
3472 hi,
3473 hi_inclusive,
3474 snapshot.epoch,
3475 )?;
3476 for rid in matched {
3477 if !overlay_rids.contains(&rid) {
3478 row_ids.push(rid);
3479 }
3480 }
3481 }
3482 let mut s = RowIdSet::from_unsorted(row_ids);
3483 self.range_scan_overlay_f64(
3484 &mut s,
3485 column_id,
3486 lo,
3487 lo_inclusive,
3488 hi,
3489 hi_inclusive,
3490 snapshot,
3491 );
3492 Ok(s)
3493 }
3494
3495 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3497 let mut s = HashSet::new();
3498 for row in self.memtable.visible_versions(snapshot.epoch) {
3499 s.insert(row.row_id.0);
3500 }
3501 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3502 s.insert(row.row_id.0);
3503 }
3504 s
3505 }
3506
3507 fn range_scan_overlay_i64(
3508 &self,
3509 s: &mut RowIdSet,
3510 column_id: u16,
3511 lo: i64,
3512 hi: i64,
3513 snapshot: Snapshot,
3514 ) {
3515 let mut newest: HashMap<u64, &Row> = HashMap::new();
3520 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3521 let memtable = self.memtable.visible_versions(snapshot.epoch);
3522 for r in &mutable {
3523 newest.entry(r.row_id.0).or_insert(r);
3524 }
3525 for r in &memtable {
3526 newest.insert(r.row_id.0, r);
3527 }
3528 for row in newest.values() {
3529 if !row.deleted {
3530 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3531 if *v >= lo && *v <= hi {
3532 s.insert(row.row_id.0);
3533 }
3534 }
3535 }
3536 }
3537 }
3538
3539 #[allow(clippy::too_many_arguments)]
3540 fn range_scan_overlay_f64(
3541 &self,
3542 s: &mut RowIdSet,
3543 column_id: u16,
3544 lo: f64,
3545 lo_inclusive: bool,
3546 hi: f64,
3547 hi_inclusive: bool,
3548 snapshot: Snapshot,
3549 ) {
3550 let mut newest: HashMap<u64, &Row> = HashMap::new();
3553 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3554 let memtable = self.memtable.visible_versions(snapshot.epoch);
3555 for r in &mutable {
3556 newest.entry(r.row_id.0).or_insert(r);
3557 }
3558 for r in &memtable {
3559 newest.insert(r.row_id.0, r);
3560 }
3561 for row in newest.values() {
3562 if !row.deleted {
3563 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3564 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3565 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3566 if ok_lo && ok_hi {
3567 s.insert(row.row_id.0);
3568 }
3569 }
3570 }
3571 }
3572 }
3573
3574 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3577 let mut row_ids = Vec::new();
3578 let overlay_rids = self.overlay_rid_set(snapshot);
3579 for rr in &self.run_refs {
3580 let mut reader = self.open_reader(rr.run_id)?;
3581 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3582 for rid in matched {
3583 if !overlay_rids.contains(&rid) {
3584 row_ids.push(rid);
3585 }
3586 }
3587 }
3588 let mut s = RowIdSet::from_unsorted(row_ids);
3589 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3590 Ok(s)
3591 }
3592
3593 fn null_scan_overlay(
3597 &self,
3598 s: &mut RowIdSet,
3599 column_id: u16,
3600 want_nulls: bool,
3601 snapshot: Snapshot,
3602 ) {
3603 let mut newest: HashMap<u64, &Row> = HashMap::new();
3604 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3605 let memtable = self.memtable.visible_versions(snapshot.epoch);
3606 for r in &mutable {
3607 newest.entry(r.row_id.0).or_insert(r);
3608 }
3609 for r in &memtable {
3610 newest.insert(r.row_id.0, r);
3611 }
3612 for row in newest.values() {
3613 if row.deleted {
3614 continue;
3615 }
3616 let is_null = !row.columns.contains_key(&column_id)
3617 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3618 if is_null == want_nulls {
3619 s.insert(row.row_id.0);
3620 }
3621 }
3622 }
3623
3624 pub fn snapshot(&self) -> Snapshot {
3625 Snapshot::at(self.epoch.visible())
3626 }
3627
3628 pub fn pin_snapshot(&mut self) -> Snapshot {
3631 let e = self.epoch.visible();
3632 *self.pinned.entry(e).or_insert(0) += 1;
3633 Snapshot::at(e)
3634 }
3635
3636 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3638 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3639 *count -= 1;
3640 if *count == 0 {
3641 self.pinned.remove(&snap.epoch);
3642 }
3643 }
3644 }
3645
3646 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3656 let local = self.pinned.keys().next().copied();
3657 let global = self.snapshots.min_pinned();
3658 match (local, global) {
3659 (Some(a), Some(b)) => Some(a.min(b)),
3660 (Some(a), None) => Some(a),
3661 (None, b) => b,
3662 }
3663 }
3664
3665 pub fn current_epoch(&self) -> Epoch {
3666 self.epoch.visible()
3667 }
3668
3669 pub fn memtable_len(&self) -> usize {
3670 self.memtable.len()
3671 }
3672
3673 pub fn count(&self) -> u64 {
3676 self.live_count
3677 }
3678
3679 pub fn count_conditions(
3683 &mut self,
3684 conditions: &[crate::query::Condition],
3685 snapshot: Snapshot,
3686 ) -> Result<Option<u64>> {
3687 use crate::query::Condition;
3688 if conditions.is_empty() {
3689 return Ok(Some(self.live_count));
3690 }
3691 let served = |c: &Condition| {
3692 matches!(
3693 c,
3694 Condition::Pk(_)
3695 | Condition::BitmapEq { .. }
3696 | Condition::BitmapIn { .. }
3697 | Condition::FmContains { .. }
3698 | Condition::FmContainsAll { .. }
3699 | Condition::Ann { .. }
3700 | Condition::Range { .. }
3701 | Condition::RangeF64 { .. }
3702 | Condition::SparseMatch { .. }
3703 | Condition::MinHashSimilar { .. }
3704 | Condition::IsNull { .. }
3705 | Condition::IsNotNull { .. }
3706 )
3707 };
3708 if !conditions.iter().all(served) {
3709 return Ok(None);
3710 }
3711 self.ensure_indexes_complete()?;
3712 let mut sets = Vec::with_capacity(conditions.len());
3713 for condition in conditions {
3714 sets.push(self.resolve_condition(condition, snapshot)?);
3715 }
3716 let count = RowIdSet::intersect_many(sets).len() as u64;
3717 crate::trace::QueryTrace::record(|t| {
3718 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
3719 t.survivor_count = Some(count as usize);
3720 t.conditions_pushed = conditions.len();
3721 });
3722 Ok(Some(count))
3723 }
3724
3725 pub fn bulk_load_columns(
3734 &mut self,
3735 user_columns: Vec<(u16, columnar::NativeColumn)>,
3736 ) -> Result<Epoch> {
3737 self.bulk_load_columns_with(user_columns, 3, false, true)
3738 }
3739
3740 pub fn bulk_load_fast(
3747 &mut self,
3748 user_columns: Vec<(u16, columnar::NativeColumn)>,
3749 ) -> Result<Epoch> {
3750 self.bulk_load_columns_with(user_columns, -1, true, false)
3751 }
3752
3753 fn bulk_load_columns_with(
3754 &mut self,
3755 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
3756 zstd_level: i32,
3757 force_plain: bool,
3758 lz4: bool,
3759 ) -> Result<Epoch> {
3760 let epoch = self.commit()?;
3761 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
3762 if n == 0 {
3763 return Ok(epoch);
3764 }
3765 let live_before = self.live_count;
3766 self.spill_mutable_run(epoch)?;
3768 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
3769 && self.indexes_complete
3770 && self.run_refs.is_empty()
3771 && self.memtable.is_empty()
3772 && self.mutable_run.is_empty();
3773 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3776 self.validate_columns_not_null(&user_columns, n)?;
3777 let winner_idx = self
3778 .bulk_pk_winner_indices(&user_columns, n)
3779 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
3780 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3781 match winner_idx.as_deref() {
3782 Some(idx) => {
3783 let compacted = user_columns
3784 .iter()
3785 .map(|(id, c)| (*id, c.gather(idx)))
3786 .collect();
3787 (compacted, idx.len())
3788 }
3789 None => (user_columns, n),
3790 };
3791 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3792 let first = self.allocator.alloc_range(write_n as u64).0;
3793 for rid in first..first + write_n as u64 {
3794 self.reservoir.offer(rid);
3795 }
3796 let run_id = self.next_run_id;
3797 self.next_run_id += 1;
3798 let path = self.run_path(run_id);
3799 let mut writer =
3800 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
3801 if force_plain {
3802 writer = writer.with_plain();
3803 } else if lz4 {
3804 writer = writer.with_lz4();
3807 } else {
3808 writer = writer.with_zstd_level(zstd_level);
3809 }
3810 if let Some(kek) = &self.kek {
3811 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3812 }
3813 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3814 self.run_refs.push(RunRef {
3815 run_id: run_id as u128,
3816 level: 0,
3817 epoch_created: epoch.0,
3818 row_count: header.row_count,
3819 });
3820 self.live_count = self.live_count.saturating_add(write_n as u64);
3821 if eager_index_build {
3822 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3823 self.index_columns_bulk(&write_columns, &row_ids);
3824 self.indexes_complete = true;
3825 self.build_learned_ranges()?;
3826 } else {
3827 self.indexes_complete = false;
3831 }
3832 self.mark_flushed(epoch)?;
3833 self.persist_manifest(epoch)?;
3834 if eager_index_build {
3835 self.checkpoint_indexes(epoch);
3836 }
3837 self.clear_result_cache();
3838 Ok(epoch)
3839 }
3840
3841 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
3859 let n = row_ids.len();
3860 if n == 0 {
3861 return;
3862 }
3863 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
3864 columns.iter().map(|(id, c)| (*id, c)).collect();
3865 let ty_of: std::collections::HashMap<u16, TypeId> =
3866 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
3867 let pk_id = self.schema.primary_key().map(|c| c.id);
3868
3869 for (i, &rid) in row_ids.iter().enumerate() {
3870 let row_id = RowId(rid);
3871 if let Some(pid) = pk_id {
3872 if let Some(col) = by_id.get(&pid) {
3873 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
3874 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
3875 self.insert_hot_pk(key, row_id);
3876 }
3877 }
3878 }
3879 for idef in &self.schema.indexes {
3880 let Some(col) = by_id.get(&idef.column_id) else {
3881 continue;
3882 };
3883 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
3884 match idef.kind {
3885 IndexKind::Bitmap => {
3886 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3887 if let Some(key) =
3888 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
3889 {
3890 b.insert(key, row_id);
3891 }
3892 }
3893 }
3894 IndexKind::FmIndex => {
3895 if let Some(f) = self.fm.get_mut(&idef.column_id) {
3896 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3897 f.insert(bytes.to_vec(), row_id);
3898 }
3899 }
3900 }
3901 IndexKind::Sparse => {
3902 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
3903 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3904 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
3905 s.insert(&terms, row_id);
3906 }
3907 }
3908 }
3909 }
3910 IndexKind::MinHash => {
3911 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
3912 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3913 let tokens = crate::index::token_hashes_from_bytes(bytes);
3914 mh.insert(&tokens, row_id);
3915 }
3916 }
3917 }
3918 _ => {}
3919 }
3920 }
3921 }
3922 }
3923
3924 pub fn visible_columns_native(
3929 &self,
3930 snapshot: Snapshot,
3931 projection: Option<&[u16]>,
3932 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
3933 let wanted: Vec<u16> = match projection {
3934 Some(p) => p.to_vec(),
3935 None => self.schema.columns.iter().map(|c| c.id).collect(),
3936 };
3937 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3938 let rr = self.run_refs[0].clone();
3939 let mut reader = self.open_reader(rr.run_id)?;
3940 let idxs = reader.visible_indices_native(snapshot.epoch)?;
3941 let all_visible = idxs.len() == reader.row_count();
3942 if reader.has_mmap() {
3948 use rayon::prelude::*;
3949 let valid: Vec<u16> = wanted
3952 .iter()
3953 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
3954 .copied()
3955 .collect();
3956 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
3958 .par_iter()
3959 .filter_map(|cid| {
3960 reader
3961 .column_native_shared(*cid)
3962 .ok()
3963 .map(|col| (*cid, col))
3964 })
3965 .collect();
3966 let cols = decoded
3967 .into_iter()
3968 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
3969 .collect();
3970 return Ok(cols);
3971 }
3972 let mut cols = Vec::with_capacity(wanted.len());
3973 for cid in &wanted {
3974 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
3975 Some(c) => c,
3976 None => continue,
3977 };
3978 let col = reader.column_native(cdef.id)?;
3979 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
3980 }
3981 return Ok(cols);
3982 }
3983 let vcols = self.visible_columns(snapshot)?;
3984 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
3985 let out: Vec<(u16, columnar::NativeColumn)> = vcols
3986 .into_iter()
3987 .filter(|(id, _)| want_set.contains(id))
3988 .map(|(id, vals)| {
3989 let ty = self
3990 .schema
3991 .columns
3992 .iter()
3993 .find(|c| c.id == id)
3994 .map(|c| c.ty)
3995 .unwrap_or(TypeId::Bytes);
3996 (id, columnar::values_to_native(ty, &vals))
3997 })
3998 .collect();
3999 Ok(out)
4000 }
4001
4002 pub fn run_count(&self) -> usize {
4003 self.run_refs.len()
4004 }
4005
4006 pub fn memtable_is_empty(&self) -> bool {
4008 self.memtable.is_empty()
4009 }
4010
4011 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4015 self.page_cache.lock().stats()
4016 }
4017
4018 pub fn reset_page_cache_stats(&self) {
4020 self.page_cache.lock().reset_stats();
4021 }
4022
4023 pub fn run_ids(&self) -> Vec<u128> {
4026 self.run_refs.iter().map(|r| r.run_id).collect()
4027 }
4028
4029 pub fn single_run_is_clean(&self) -> bool {
4033 if self.run_refs.len() != 1 {
4034 return false;
4035 }
4036 self.open_reader(self.run_refs[0].run_id)
4037 .map(|r| r.is_clean())
4038 .unwrap_or(false)
4039 }
4040
4041 fn resolve_footprint(
4048 &self,
4049 conditions: &[crate::query::Condition],
4050 _snapshot: Snapshot,
4051 ) -> roaring::RoaringBitmap {
4052 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4053 return roaring::RoaringBitmap::new();
4054 }
4055 if self.run_refs.is_empty() {
4056 return roaring::RoaringBitmap::new();
4057 }
4058 if self.run_refs.len() == 1 {
4060 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4061 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4062 return rids.to_roaring_lossy();
4063 }
4064 }
4065 }
4066 roaring::RoaringBitmap::new()
4067 }
4068
4069 pub fn query_columns_native_cached(
4080 &mut self,
4081 conditions: &[crate::query::Condition],
4082 projection: Option<&[u16]>,
4083 snapshot: Snapshot,
4084 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4085 if conditions.is_empty() {
4086 return self.query_columns_native(conditions, projection, snapshot);
4087 }
4088 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4092 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4093 crate::trace::QueryTrace::record(|t| {
4094 t.result_cache_hit = true;
4095 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4096 });
4097 return Ok(Some((*hit).clone()));
4098 }
4099 let res = self.query_columns_native(conditions, projection, snapshot)?;
4100 if let Some(cols) = &res {
4101 let footprint = self.resolve_footprint(conditions, snapshot);
4102 let condition_cols = crate::query::condition_columns(conditions);
4103 self.result_cache.lock().insert(
4104 key,
4105 CachedEntry {
4106 data: CachedData::Columns(Arc::new(cols.clone())),
4107 footprint,
4108 condition_cols,
4109 },
4110 );
4111 }
4112 Ok(res)
4113 }
4114
4115 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4120 if q.conditions.is_empty() {
4121 return self.query(q);
4122 }
4123 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4124 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4125 crate::trace::QueryTrace::record(|t| {
4126 t.result_cache_hit = true;
4127 t.scan_mode = crate::trace::ScanMode::Materialized;
4128 });
4129 return Ok((*hit).clone());
4130 }
4131 let rows = self.query(q)?;
4132 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4133 let condition_cols = crate::query::condition_columns(&q.conditions);
4134 self.result_cache.lock().insert(
4135 key,
4136 CachedEntry {
4137 data: CachedData::Rows(Arc::new(rows.clone())),
4138 footprint,
4139 condition_cols,
4140 },
4141 );
4142 Ok(rows)
4143 }
4144
4145 #[allow(clippy::type_complexity)]
4160 pub fn query_columns_native_traced(
4161 &mut self,
4162 conditions: &[crate::query::Condition],
4163 projection: Option<&[u16]>,
4164 snapshot: Snapshot,
4165 ) -> Result<(
4166 Option<Vec<(u16, columnar::NativeColumn)>>,
4167 crate::trace::QueryTrace,
4168 )> {
4169 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4170 self.query_columns_native(conditions, projection, snapshot)
4171 });
4172 Ok((result?, trace))
4173 }
4174
4175 #[allow(clippy::type_complexity)]
4178 pub fn query_columns_native_cached_traced(
4179 &mut self,
4180 conditions: &[crate::query::Condition],
4181 projection: Option<&[u16]>,
4182 snapshot: Snapshot,
4183 ) -> Result<(
4184 Option<Vec<(u16, columnar::NativeColumn)>>,
4185 crate::trace::QueryTrace,
4186 )> {
4187 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4188 self.query_columns_native_cached(conditions, projection, snapshot)
4189 });
4190 Ok((result?, trace))
4191 }
4192
4193 pub fn native_page_cursor_traced(
4195 &self,
4196 snapshot: Snapshot,
4197 projection: Vec<(u16, TypeId)>,
4198 conditions: &[crate::query::Condition],
4199 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4200 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4201 self.native_page_cursor(snapshot, projection, conditions)
4202 });
4203 Ok((result?, trace))
4204 }
4205
4206 pub fn native_multi_run_cursor_traced(
4208 &self,
4209 snapshot: Snapshot,
4210 projection: Vec<(u16, TypeId)>,
4211 conditions: &[crate::query::Condition],
4212 ) -> Result<(
4213 Option<crate::cursor::MultiRunCursor>,
4214 crate::trace::QueryTrace,
4215 )> {
4216 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4217 self.native_multi_run_cursor(snapshot, projection, conditions)
4218 });
4219 Ok((result?, trace))
4220 }
4221
4222 pub fn count_conditions_traced(
4224 &mut self,
4225 conditions: &[crate::query::Condition],
4226 snapshot: Snapshot,
4227 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4228 let (result, trace) =
4229 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4230 Ok((result?, trace))
4231 }
4232
4233 pub fn query_traced(
4235 &mut self,
4236 q: &crate::query::Query,
4237 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4238 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4239 Ok((result?, trace))
4240 }
4241
4242 pub fn query_columns_native(
4247 &mut self,
4248 conditions: &[crate::query::Condition],
4249 projection: Option<&[u16]>,
4250 snapshot: Snapshot,
4251 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4252 use crate::query::Condition;
4253 if conditions.is_empty() {
4254 return Ok(None);
4255 }
4256 self.ensure_indexes_complete()?;
4257
4258 let served = |c: &Condition| {
4263 matches!(
4264 c,
4265 Condition::Pk(_)
4266 | Condition::BitmapEq { .. }
4267 | Condition::BitmapIn { .. }
4268 | Condition::FmContains { .. }
4269 | Condition::FmContainsAll { .. }
4270 | Condition::Ann { .. }
4271 | Condition::Range { .. }
4272 | Condition::RangeF64 { .. }
4273 | Condition::SparseMatch { .. }
4274 | Condition::MinHashSimilar { .. }
4275 | Condition::IsNull { .. }
4276 | Condition::IsNotNull { .. }
4277 )
4278 };
4279 if !conditions.iter().all(served) {
4280 return Ok(None);
4281 }
4282 let fast_path =
4283 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4284 crate::trace::QueryTrace::record(|t| {
4285 t.run_count = self.run_refs.len();
4286 t.memtable_rows = self.memtable.len();
4287 t.mutable_run_rows = self.mutable_run.len();
4288 t.conditions_pushed = conditions.len();
4289 t.learned_range_used = conditions.iter().any(|c| match c {
4290 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4291 self.learned_range.contains_key(column_id)
4292 }
4293 _ => false,
4294 });
4295 });
4296 let col_ids: Vec<u16> = projection
4298 .map(|p| p.to_vec())
4299 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4300 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4301 .iter()
4302 .map(|&cid| {
4303 let ty = self
4304 .schema
4305 .columns
4306 .iter()
4307 .find(|c| c.id == cid)
4308 .map(|c| c.ty)
4309 .unwrap_or(TypeId::Bytes);
4310 (cid, ty)
4311 })
4312 .collect();
4313
4314 if fast_path {
4320 let needs_column = conditions.iter().any(|c| match c {
4323 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4324 Condition::RangeF64 { column_id, .. } => {
4325 !self.learned_range.contains_key(column_id)
4326 }
4327 _ => false,
4328 });
4329 let mut reader_opt: Option<RunReader> = if needs_column {
4330 Some(self.open_reader(self.run_refs[0].run_id)?)
4331 } else {
4332 None
4333 };
4334 let mut sets: Vec<RowIdSet> = Vec::new();
4335 for c in conditions {
4336 let s = match c {
4337 Condition::Range { column_id, lo, hi }
4338 if !self.learned_range.contains_key(column_id) =>
4339 {
4340 if reader_opt.is_none() {
4341 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4342 }
4343 reader_opt
4344 .as_mut()
4345 .expect("reader opened for range")
4346 .range_row_id_set_i64(*column_id, *lo, *hi)?
4347 }
4348 Condition::RangeF64 {
4349 column_id,
4350 lo,
4351 lo_inclusive,
4352 hi,
4353 hi_inclusive,
4354 } if !self.learned_range.contains_key(column_id) => {
4355 if reader_opt.is_none() {
4356 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4357 }
4358 reader_opt
4359 .as_mut()
4360 .expect("reader opened for range")
4361 .range_row_id_set_f64(
4362 *column_id,
4363 *lo,
4364 *lo_inclusive,
4365 *hi,
4366 *hi_inclusive,
4367 )?
4368 }
4369 _ => self.resolve_condition(c, snapshot)?,
4370 };
4371 sets.push(s);
4372 }
4373 let candidates = RowIdSet::intersect_many(sets);
4374 crate::trace::QueryTrace::record(|t| {
4375 t.survivor_count = Some(candidates.len());
4376 });
4377 if candidates.is_empty() {
4378 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4379 .iter()
4380 .map(|&id| {
4381 (
4382 id,
4383 columnar::null_native(
4384 proj_pairs
4385 .iter()
4386 .find(|(c, _)| c == &id)
4387 .map(|(_, t)| *t)
4388 .unwrap_or(TypeId::Bytes),
4389 0,
4390 ),
4391 )
4392 })
4393 .collect();
4394 return Ok(Some(cols));
4395 }
4396 let mut reader = match reader_opt.take() {
4397 Some(r) => r,
4398 None => self.open_reader(self.run_refs[0].run_id)?,
4399 };
4400 let candidate_ids = candidates.into_sorted_vec();
4401 let (positions, fast_rid) = if let Some(positions) =
4402 reader.positions_for_row_ids_fast(&candidate_ids)
4403 {
4404 (positions, true)
4405 } else {
4406 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4407 match col {
4408 columnar::NativeColumn::Int64 { data, .. } => {
4409 let mut p: Vec<usize> = candidate_ids
4410 .iter()
4411 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4412 .collect();
4413 p.sort_unstable();
4414 (p, false)
4415 }
4416 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4417 }
4418 };
4419 crate::trace::QueryTrace::record(|t| {
4420 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4421 t.fast_row_id_map = fast_rid;
4422 });
4423 let mut cols = Vec::with_capacity(col_ids.len());
4424 for cid in &col_ids {
4425 let col = reader.column_native(*cid)?;
4426 cols.push((*cid, col.gather(&positions)));
4427 }
4428 return Ok(Some(cols));
4429 }
4430
4431 if !self.run_refs.is_empty() {
4444 use crate::cursor::{drain_cursor_to_columns, Cursor};
4445 let remaining: usize;
4446 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4447 let c = self
4448 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4449 .expect("single-run cursor should build when run_refs.len() == 1");
4450 remaining = c.remaining_rows();
4451 Box::new(c)
4452 } else {
4453 let c = self
4454 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4455 .expect("multi-run cursor should build when run_refs.len() >= 1");
4456 remaining = c.remaining_rows();
4457 Box::new(c)
4458 };
4459 crate::trace::QueryTrace::record(|t| {
4460 if t.survivor_count.is_none() {
4461 t.survivor_count = Some(remaining);
4462 }
4463 });
4464 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4465 return Ok(Some(cols));
4466 }
4467
4468 crate::trace::QueryTrace::record(|t| {
4473 t.scan_mode = crate::trace::ScanMode::Materialized;
4474 t.row_materialized = true;
4475 });
4476 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4477 for c in conditions {
4478 sets.push(self.resolve_condition(c, snapshot)?);
4479 }
4480 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4481 let rows = self.rows_for_rids(&rids, snapshot)?;
4482 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4483 for (cid, ty) in &proj_pairs {
4484 let vals: Vec<Value> = rows
4485 .iter()
4486 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4487 .collect();
4488 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4489 }
4490 Ok(Some(cols))
4491 }
4492
4493 pub fn native_page_cursor(
4508 &self,
4509 snapshot: Snapshot,
4510 projection: Vec<(u16, TypeId)>,
4511 conditions: &[crate::query::Condition],
4512 ) -> Result<Option<NativePageCursor>> {
4513 use crate::cursor::build_page_plans;
4514 if !conditions.is_empty() && !self.indexes_complete {
4517 return Ok(None);
4518 }
4519 if self.run_refs.len() != 1 {
4520 return Ok(None);
4521 }
4522 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4523 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4524
4525 let overlay_rids: HashSet<u64> = {
4528 let mut s = HashSet::new();
4529 for row in self.memtable.visible_versions(snapshot.epoch) {
4530 s.insert(row.row_id.0);
4531 }
4532 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4533 s.insert(row.row_id.0);
4534 }
4535 s
4536 };
4537
4538 let survivors = if conditions.is_empty() {
4542 None
4543 } else {
4544 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4545 };
4546
4547 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4554 survivors.clone()
4555 } else if let Some(s) = &survivors {
4556 let mut run_set = s.clone();
4557 run_set.remove_many(overlay_rids.iter().copied());
4558 Some(run_set)
4559 } else {
4560 Some(RowIdSet::from_unsorted(
4561 rids.iter()
4562 .map(|&r| r as u64)
4563 .filter(|r| !overlay_rids.contains(r))
4564 .collect(),
4565 ))
4566 };
4567
4568 let overlay_rows = if overlay_rids.is_empty() {
4569 Vec::new()
4570 } else {
4571 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4572 self.overlay_visible_rows(snapshot, bound)
4573 };
4574
4575 let plans = if positions.is_empty() {
4577 Vec::new()
4578 } else {
4579 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4580 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4581 };
4582
4583 let overlay = if overlay_rows.is_empty() {
4585 None
4586 } else {
4587 let filtered =
4588 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4589 if filtered.is_empty() {
4590 None
4591 } else {
4592 Some(self.materialize_overlay(&filtered, &projection))
4593 }
4594 };
4595
4596 let overlay_row_count = overlay
4597 .as_ref()
4598 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4599 .unwrap_or(0);
4600 crate::trace::QueryTrace::record(|t| {
4601 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4602 t.run_count = self.run_refs.len();
4603 t.memtable_rows = self.memtable.len();
4604 t.mutable_run_rows = self.mutable_run.len();
4605 t.overlay_rows = overlay_row_count;
4606 t.conditions_pushed = conditions.len();
4607 t.pages_decoded = plans
4608 .iter()
4609 .map(|p| p.positions.len())
4610 .sum::<usize>()
4611 .min(1);
4612 });
4613
4614 Ok(Some(NativePageCursor::new_with_overlay(
4615 reader, projection, plans, overlay,
4616 )))
4617 }
4618 #[allow(clippy::type_complexity)]
4628 pub fn native_multi_run_cursor(
4629 &self,
4630 snapshot: Snapshot,
4631 projection: Vec<(u16, TypeId)>,
4632 conditions: &[crate::query::Condition],
4633 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4634 use crate::cursor::{MultiRunCursor, RunStream};
4635 use crate::sorted_run::SYS_ROW_ID;
4636 use std::collections::{BinaryHeap, HashMap, HashSet};
4637 if !conditions.is_empty() && !self.indexes_complete {
4640 return Ok(None);
4641 }
4642 if self.run_refs.is_empty() {
4643 return Ok(None);
4644 }
4645
4646 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
4648 Vec::with_capacity(self.run_refs.len());
4649 for rr in &self.run_refs {
4650 let mut reader = self.open_reader(rr.run_id)?;
4651 let (rids, eps, del) = reader.system_columns_native()?;
4652 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
4653 run_meta.push((reader, rids, eps, del, page_rows));
4654 }
4655
4656 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
4660 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
4661 for i in 0..rids.len() {
4662 let rid = rids[i] as u64;
4663 let e = eps[i] as u64;
4664 if e > snapshot.epoch.0 {
4665 continue;
4666 }
4667 let is_del = del[i] != 0;
4668 best.entry(rid)
4669 .and_modify(|cur| {
4670 if e > cur.0 {
4671 *cur = (e, run_idx, i, is_del);
4672 }
4673 })
4674 .or_insert((e, run_idx, i, is_del));
4675 }
4676 }
4677
4678 let overlay_rids: HashSet<u64> = {
4680 let mut s = HashSet::new();
4681 for row in self.memtable.visible_versions(snapshot.epoch) {
4682 s.insert(row.row_id.0);
4683 }
4684 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4685 s.insert(row.row_id.0);
4686 }
4687 s
4688 };
4689
4690 let survivors: Option<RowIdSet> = if conditions.is_empty() {
4692 None
4693 } else {
4694 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4695 for c in conditions {
4696 sets.push(self.resolve_condition(c, snapshot)?);
4697 }
4698 Some(RowIdSet::intersect_many(sets))
4699 };
4700
4701 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
4705 for (rid, (_, run_idx, pos, deleted)) in &best {
4706 if *deleted {
4707 continue;
4708 }
4709 if overlay_rids.contains(rid) {
4710 continue;
4711 }
4712 if let Some(s) = &survivors {
4713 if !s.contains(*rid) {
4714 continue;
4715 }
4716 }
4717 per_run[*run_idx].push((*rid, *pos));
4718 }
4719 for v in per_run.iter_mut() {
4720 v.sort_unstable_by_key(|&(rid, _)| rid);
4721 }
4722
4723 let mut streams = Vec::with_capacity(run_meta.len());
4725 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
4726 let mut total = 0usize;
4727 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
4728 let mut starts = Vec::with_capacity(page_rows.len());
4729 let mut acc = 0usize;
4730 for &r in &page_rows {
4731 starts.push(acc);
4732 acc += r;
4733 }
4734 let mut survivors_vec: Vec<(u64, usize, usize)> =
4735 Vec::with_capacity(per_run[run_idx].len());
4736 for &(rid, pos) in &per_run[run_idx] {
4737 let page_seq = match starts.partition_point(|&s| s <= pos) {
4738 0 => continue,
4739 p => p - 1,
4740 };
4741 let within = pos - starts[page_seq];
4742 survivors_vec.push((rid, page_seq, within));
4743 }
4744 total += survivors_vec.len();
4745 if let Some(&(rid, _, _)) = survivors_vec.first() {
4746 heap.push(std::cmp::Reverse((rid, run_idx)));
4747 }
4748 streams.push(RunStream::new(reader, survivors_vec, page_rows));
4749 }
4750
4751 let overlay_rows = if overlay_rids.is_empty() {
4753 Vec::new()
4754 } else {
4755 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4756 self.overlay_visible_rows(snapshot, bound)
4757 };
4758 let overlay = if overlay_rows.is_empty() {
4759 None
4760 } else {
4761 let filtered =
4762 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4763 if filtered.is_empty() {
4764 None
4765 } else {
4766 Some(self.materialize_overlay(&filtered, &projection))
4767 }
4768 };
4769
4770 let overlay_row_count = overlay
4771 .as_ref()
4772 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4773 .unwrap_or(0);
4774 crate::trace::QueryTrace::record(|t| {
4775 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
4776 t.run_count = self.run_refs.len();
4777 t.memtable_rows = self.memtable.len();
4778 t.mutable_run_rows = self.mutable_run.len();
4779 t.overlay_rows = overlay_row_count;
4780 t.conditions_pushed = conditions.len();
4781 t.survivor_count = Some(total);
4782 });
4783
4784 Ok(Some(MultiRunCursor::new(
4785 streams, projection, heap, total, overlay,
4786 )))
4787 }
4788
4789 fn overlay_materialization_bound<'a>(
4801 conditions: &[crate::query::Condition],
4802 survivors: &'a Option<RowIdSet>,
4803 ) -> Option<&'a RowIdSet> {
4804 use crate::query::Condition;
4805 let has_range = conditions
4806 .iter()
4807 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
4808 if has_range {
4809 None
4810 } else {
4811 survivors.as_ref()
4812 }
4813 }
4814
4815 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
4827 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4828 let mut fold = |row: Row| {
4829 if let Some(b) = bound {
4830 if !b.contains(row.row_id.0) {
4831 return;
4832 }
4833 }
4834 best.entry(row.row_id.0)
4835 .and_modify(|(be, br)| {
4836 if row.committed_epoch > *be {
4837 *be = row.committed_epoch;
4838 *br = row.clone();
4839 }
4840 })
4841 .or_insert_with(|| (row.committed_epoch, row));
4842 };
4843 for row in self.memtable.visible_versions(snapshot.epoch) {
4844 fold(row);
4845 }
4846 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4847 fold(row);
4848 }
4849 let mut out: Vec<Row> = best
4850 .into_values()
4851 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
4852 .collect();
4853 out.sort_by_key(|r| r.row_id);
4854 out
4855 }
4856
4857 fn filter_overlay_rows(
4865 &self,
4866 rows: Vec<Row>,
4867 conditions: &[crate::query::Condition],
4868 survivors: Option<&RowIdSet>,
4869 snapshot: Snapshot,
4870 ) -> Result<Vec<Row>> {
4871 if conditions.is_empty() {
4872 return Ok(rows);
4873 }
4874 use crate::query::Condition;
4875 let all_index_served = !conditions
4879 .iter()
4880 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
4881 if all_index_served {
4882 return Ok(rows
4883 .into_iter()
4884 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
4885 .collect());
4886 }
4887 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4890 for c in conditions {
4891 let s = match c {
4892 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
4893 _ => self.resolve_condition(c, snapshot)?,
4894 };
4895 per_cond_sets.push(s);
4896 }
4897 Ok(rows
4898 .into_iter()
4899 .filter(|row| {
4900 conditions.iter().enumerate().all(|(i, c)| match c {
4901 Condition::Range { column_id, lo, hi } => {
4902 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
4903 }
4904 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
4905 match row.columns.get(column_id) {
4906 Some(Value::Float64(v)) => {
4907 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
4908 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
4909 lo_ok && hi_ok
4910 }
4911 _ => false,
4912 }
4913 }
4914 _ => per_cond_sets[i].contains(row.row_id.0),
4915 })
4916 })
4917 .collect())
4918 }
4919
4920 fn materialize_overlay(
4923 &self,
4924 rows: &[Row],
4925 projection: &[(u16, TypeId)],
4926 ) -> Vec<columnar::NativeColumn> {
4927 if projection.is_empty() {
4928 return vec![columnar::null_native(TypeId::Int64, rows.len())];
4929 }
4930 let mut cols = Vec::with_capacity(projection.len());
4931 for (cid, ty) in projection {
4932 let vals: Vec<Value> = rows
4933 .iter()
4934 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4935 .collect();
4936 cols.push(columnar::values_to_native(*ty, &vals));
4937 }
4938 cols
4939 }
4940
4941 fn resolve_survivor_rids(
4946 &self,
4947 conditions: &[crate::query::Condition],
4948 reader: &mut RunReader,
4949 ) -> Result<RowIdSet> {
4950 use crate::query::Condition;
4951 let mut sets: Vec<RowIdSet> = Vec::new();
4952 for c in conditions {
4953 let s: RowIdSet = match c {
4954 Condition::Pk(key) => {
4955 let lookup = self
4956 .schema
4957 .primary_key()
4958 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
4959 .unwrap_or_else(|| key.clone());
4960 self.hot
4961 .get(&lookup)
4962 .map(|r| RowIdSet::one(r.0))
4963 .unwrap_or_else(RowIdSet::empty)
4964 }
4965 Condition::BitmapEq { column_id, value } => {
4966 let lookup = self.index_lookup_key_bytes(*column_id, value);
4967 self.bitmap
4968 .get(column_id)
4969 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
4970 .unwrap_or_else(RowIdSet::empty)
4971 }
4972 Condition::BitmapIn { column_id, values } => {
4973 let bm = self.bitmap.get(column_id);
4974 let mut acc = roaring::RoaringBitmap::new();
4975 if let Some(b) = bm {
4976 for v in values {
4977 let lookup = self.index_lookup_key_bytes(*column_id, v);
4978 acc |= b.get(&lookup);
4979 }
4980 }
4981 RowIdSet::from_roaring(acc)
4982 }
4983 Condition::FmContains { column_id, pattern } => self
4984 .fm
4985 .get(column_id)
4986 .map(|f| {
4987 RowIdSet::from_unsorted(
4988 f.locate(pattern).into_iter().map(|r| r.0).collect(),
4989 )
4990 })
4991 .unwrap_or_else(RowIdSet::empty),
4992 Condition::FmContainsAll {
4993 column_id,
4994 patterns,
4995 } => {
4996 if let Some(f) = self.fm.get(column_id) {
4997 let sets: Vec<RowIdSet> = patterns
4998 .iter()
4999 .map(|pat| {
5000 RowIdSet::from_unsorted(
5001 f.locate(pat).into_iter().map(|r| r.0).collect(),
5002 )
5003 })
5004 .collect();
5005 RowIdSet::intersect_many(sets)
5006 } else {
5007 RowIdSet::empty()
5008 }
5009 }
5010 Condition::Ann {
5011 column_id,
5012 query,
5013 k,
5014 } => self
5015 .ann
5016 .get(column_id)
5017 .map(|a| {
5018 RowIdSet::from_unsorted(
5019 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5020 )
5021 })
5022 .unwrap_or_else(RowIdSet::empty),
5023 Condition::SparseMatch {
5024 column_id,
5025 query,
5026 k,
5027 } => self
5028 .sparse
5029 .get(column_id)
5030 .map(|s| {
5031 RowIdSet::from_unsorted(
5032 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5033 )
5034 })
5035 .unwrap_or_else(RowIdSet::empty),
5036 Condition::MinHashSimilar {
5037 column_id,
5038 query,
5039 k,
5040 } => self
5041 .minhash
5042 .get(column_id)
5043 .map(|mh| {
5044 RowIdSet::from_unsorted(
5045 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5046 )
5047 })
5048 .unwrap_or_else(RowIdSet::empty),
5049 Condition::Range { column_id, lo, hi } => {
5050 if let Some(li) = self.learned_range.get(column_id) {
5051 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5052 } else {
5053 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5054 }
5055 }
5056 Condition::RangeF64 {
5057 column_id,
5058 lo,
5059 lo_inclusive,
5060 hi,
5061 hi_inclusive,
5062 } => {
5063 if let Some(li) = self.learned_range.get(column_id) {
5064 RowIdSet::from_unsorted(
5065 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5066 .into_iter()
5067 .collect(),
5068 )
5069 } else {
5070 reader.range_row_id_set_f64(
5071 *column_id,
5072 *lo,
5073 *lo_inclusive,
5074 *hi,
5075 *hi_inclusive,
5076 )?
5077 }
5078 }
5079 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5080 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5081 };
5082 sets.push(s);
5083 }
5084 Ok(RowIdSet::intersect_many(sets))
5085 }
5086
5087 pub fn scan_cursor(
5108 &self,
5109 snapshot: Snapshot,
5110 projection: Vec<(u16, TypeId)>,
5111 conditions: &[crate::query::Condition],
5112 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5113 if !conditions.is_empty() && !self.indexes_complete {
5119 return Ok(None);
5120 }
5121 if self.run_refs.len() == 1 {
5122 Ok(self
5123 .native_page_cursor(snapshot, projection, conditions)?
5124 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5125 } else {
5126 Ok(self
5127 .native_multi_run_cursor(snapshot, projection, conditions)?
5128 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5129 }
5130 }
5131
5132 pub fn aggregate_native(
5146 &self,
5147 snapshot: Snapshot,
5148 column: Option<u16>,
5149 conditions: &[crate::query::Condition],
5150 agg: NativeAgg,
5151 ) -> Result<Option<NativeAggResult>> {
5152 if self.run_refs.len() == 1 && conditions.is_empty() {
5154 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5155 return Ok(Some(res));
5156 }
5157 }
5158 if matches!(agg, NativeAgg::Count) && column.is_none() {
5160 return Ok(self
5161 .scan_cursor(snapshot, Vec::new(), conditions)?
5162 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5163 }
5164 let cid = match column {
5167 Some(c) => c,
5168 None => return Ok(None),
5169 };
5170 let ty = self.column_type(cid);
5171 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5172 return Ok(None);
5173 };
5174 match ty {
5175 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5176 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5177 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5178 }
5179 TypeId::Float64 => {
5180 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5181 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5182 }
5183 _ => Ok(None),
5184 }
5185 }
5186
5187 fn aggregate_from_stats(
5195 &self,
5196 snapshot: Snapshot,
5197 column: Option<u16>,
5198 agg: NativeAgg,
5199 ) -> Result<Option<NativeAggResult>> {
5200 let cid = match (agg, column) {
5201 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5202 _ => return Ok(None), };
5204 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5205 return Ok(None);
5206 };
5207 let Some(cs) = stats.get(&cid) else {
5208 return Ok(None);
5209 };
5210 match agg {
5211 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5213 self.live_count.saturating_sub(cs.null_count),
5214 ))),
5215 NativeAgg::Min | NativeAgg::Max => {
5216 let bound = if agg == NativeAgg::Min {
5217 &cs.min
5218 } else {
5219 &cs.max
5220 };
5221 match bound {
5222 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5223 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5224 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5229 None => Ok(None),
5230 }
5231 }
5232 _ => Ok(None),
5233 }
5234 }
5235
5236 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5245 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5246 return Ok(None);
5247 }
5248 self.ensure_indexes_complete()?;
5251 let reader = self.open_reader(self.run_refs[0].run_id)?;
5252 if self.live_count != reader.row_count() as u64 {
5253 return Ok(None);
5254 }
5255 let Some(bm) = self.bitmap.get(&column_id) else {
5256 return Ok(None); };
5258 let mut distinct = bm.value_count() as u64;
5259 if !bm.get(&Value::Null.encode_key()).is_empty() {
5262 distinct = distinct.saturating_sub(1);
5263 }
5264 Ok(Some(distinct))
5265 }
5266
5267 pub fn aggregate_incremental(
5279 &mut self,
5280 cache_key: u64,
5281 conditions: &[crate::query::Condition],
5282 column: Option<u16>,
5283 agg: NativeAgg,
5284 ) -> Result<IncrementalAggResult> {
5285 let snap = self.snapshot();
5286 let cur_wm = self.allocator.current().0;
5287 let cur_epoch = snap.epoch.0;
5288 let incremental_ok =
5295 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5296
5297 if incremental_ok {
5300 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5301 if cached.epoch == cur_epoch {
5302 return Ok(IncrementalAggResult {
5303 state: cached.state,
5304 incremental: true,
5305 delta_rows: 0,
5306 });
5307 }
5308 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5309 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5310 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5311 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5312 let delta_state = agg_state_from_rows(
5313 &delta_rows,
5314 conditions,
5315 &index_sets,
5316 column,
5317 agg,
5318 &self.schema,
5319 )?;
5320 let merged = cached.state.merge(delta_state);
5321 let delta_n = delta_rids.len() as u64;
5322 self.agg_cache.insert(
5323 cache_key,
5324 CachedAgg {
5325 state: merged.clone(),
5326 watermark: cur_wm,
5327 epoch: cur_epoch,
5328 },
5329 );
5330 return Ok(IncrementalAggResult {
5331 state: merged,
5332 incremental: true,
5333 delta_rows: delta_n,
5334 });
5335 }
5336 }
5337 }
5338
5339 let cursor_ok =
5344 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5345 let state = if cursor_ok && agg != NativeAgg::Avg {
5346 match self.aggregate_native(snap, column, conditions, agg)? {
5347 Some(result) => {
5348 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5349 }
5350 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5351 }
5352 } else {
5353 self.agg_state_full_scan(conditions, column, agg, snap)?
5354 };
5355 if incremental_ok {
5357 self.agg_cache.insert(
5358 cache_key,
5359 CachedAgg {
5360 state: state.clone(),
5361 watermark: cur_wm,
5362 epoch: cur_epoch,
5363 },
5364 );
5365 }
5366 Ok(IncrementalAggResult {
5367 state,
5368 incremental: false,
5369 delta_rows: 0,
5370 })
5371 }
5372
5373 fn agg_state_full_scan(
5376 &self,
5377 conditions: &[crate::query::Condition],
5378 column: Option<u16>,
5379 agg: NativeAgg,
5380 snap: Snapshot,
5381 ) -> Result<AggState> {
5382 let rows = self.visible_rows(snap)?;
5383 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5384 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5385 }
5386
5387 fn resolve_index_conditions(
5390 &self,
5391 conditions: &[crate::query::Condition],
5392 snapshot: Snapshot,
5393 ) -> Result<Vec<RowIdSet>> {
5394 use crate::query::Condition;
5395 let mut sets = Vec::new();
5396 for c in conditions {
5397 if matches!(
5398 c,
5399 Condition::Ann { .. }
5400 | Condition::SparseMatch { .. }
5401 | Condition::MinHashSimilar { .. }
5402 ) {
5403 sets.push(self.resolve_condition(c, snapshot)?);
5404 }
5405 }
5406 Ok(sets)
5407 }
5408
5409 fn column_type(&self, cid: u16) -> TypeId {
5410 self.schema
5411 .columns
5412 .iter()
5413 .find(|c| c.id == cid)
5414 .map(|c| c.ty)
5415 .unwrap_or(TypeId::Bytes)
5416 }
5417
5418 pub fn approx_aggregate(
5427 &self,
5428 conditions: &[crate::query::Condition],
5429 column: Option<u16>,
5430 agg: ApproxAgg,
5431 z: f64,
5432 ) -> Result<Option<ApproxResult>> {
5433 use crate::query::Condition;
5434 let snapshot = self.snapshot();
5435 let n_pop = self.live_count;
5436 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5437 if sample_rids.is_empty() {
5438 return Ok(None);
5439 }
5440 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5442 let s = live_sample.len();
5443 if s == 0 {
5444 return Ok(None);
5445 }
5446
5447 let mut index_sets: Vec<RowIdSet> = Vec::new();
5450 for c in conditions {
5451 if matches!(
5452 c,
5453 Condition::Ann { .. }
5454 | Condition::SparseMatch { .. }
5455 | Condition::MinHashSimilar { .. }
5456 ) {
5457 index_sets.push(self.resolve_condition(c, snapshot)?);
5458 }
5459 }
5460
5461 let cid = match (agg, column) {
5463 (ApproxAgg::Count, _) => None,
5464 (_, Some(c)) => Some(c),
5465 _ => return Ok(None),
5466 };
5467 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5468 for r in &live_sample {
5469 if !conditions
5471 .iter()
5472 .all(|c| condition_matches_row(c, r, &self.schema))
5473 {
5474 continue;
5475 }
5476 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5478 continue;
5479 }
5480 if let Some(cid) = cid {
5481 if let Some(v) = as_f64(r.columns.get(&cid)) {
5482 passing_vals.push(v);
5483 } } else {
5485 passing_vals.push(0.0); }
5487 }
5488 let m = passing_vals.len();
5489
5490 let (point, half) = match agg {
5491 ApproxAgg::Count => {
5492 let p = m as f64 / s as f64;
5494 let point = n_pop as f64 * p;
5495 let var = if s > 1 {
5496 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5497 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5498 } else {
5499 0.0
5500 };
5501 (point, z * var.sqrt())
5502 }
5503 ApproxAgg::Sum => {
5504 let y: Vec<f64> = live_sample
5506 .iter()
5507 .map(|r| {
5508 let passes_row = conditions
5509 .iter()
5510 .all(|c| condition_matches_row(c, r, &self.schema))
5511 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5512 if passes_row {
5513 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5514 } else {
5515 0.0
5516 }
5517 })
5518 .collect();
5519 let mean_y = y.iter().sum::<f64>() / s as f64;
5520 let point = n_pop as f64 * mean_y;
5521 let var = if s > 1 {
5522 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5523 let var_y = ss / (s - 1) as f64;
5524 n_pop as f64 * n_pop as f64 * var_y / s as f64
5525 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5526 } else {
5527 0.0
5528 };
5529 (point, z * var.sqrt())
5530 }
5531 ApproxAgg::Avg => {
5532 if m == 0 {
5533 return Ok(Some(ApproxResult {
5534 point: 0.0,
5535 ci_low: 0.0,
5536 ci_high: 0.0,
5537 n_population: n_pop,
5538 n_sample_live: s,
5539 n_passing: 0,
5540 }));
5541 }
5542 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5543 let half = if m > 1 {
5544 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5545 let sd = (ss / (m - 1) as f64).sqrt();
5546 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5547 z * sd / (m as f64).sqrt() * fpc.sqrt()
5548 } else {
5549 0.0
5550 };
5551 (mean, half)
5552 }
5553 };
5554
5555 Ok(Some(ApproxResult {
5556 point,
5557 ci_low: point - half,
5558 ci_high: point + half,
5559 n_population: n_pop,
5560 n_sample_live: s,
5561 n_passing: m,
5562 }))
5563 }
5564
5565 pub fn exact_column_stats(
5573 &self,
5574 _snapshot: Snapshot,
5575 projection: &[u16],
5576 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5577 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5578 return Ok(None);
5579 }
5580 let reader = self.open_reader(self.run_refs[0].run_id)?;
5581 if self.live_count != reader.row_count() as u64 {
5582 return Ok(None);
5583 }
5584 let mut out = HashMap::new();
5585 for &cid in projection {
5586 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5587 Some(c) => c,
5588 None => continue,
5589 };
5590 let Some(stats) = reader.column_page_stats(cid) else {
5592 out.insert(
5593 cid,
5594 ColumnStat {
5595 min: None,
5596 max: None,
5597 null_count: self.live_count,
5598 },
5599 );
5600 continue;
5601 };
5602 let stat = match cdef.ty {
5603 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5604 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5605 min: mn.map(Value::Int64),
5606 max: mx.map(Value::Int64),
5607 null_count: n,
5608 })
5609 }
5610 TypeId::Float64 => {
5611 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5612 min: mn.map(Value::Float64),
5613 max: mx.map(Value::Float64),
5614 null_count: n,
5615 })
5616 }
5617 _ => None,
5618 };
5619 if let Some(s) = stat {
5620 out.insert(cid, s);
5621 }
5622 }
5623 Ok(Some(out))
5624 }
5625
5626 pub fn dir(&self) -> &Path {
5627 &self.dir
5628 }
5629
5630 pub fn schema(&self) -> &Schema {
5631 &self.schema
5632 }
5633
5634 pub(crate) fn prepare_alter_column(
5635 &mut self,
5636 column_name: &str,
5637 change: &AlterColumn,
5638 ) -> Result<ColumnDef> {
5639 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
5640 return Err(MongrelError::InvalidArgument(
5641 "ALTER COLUMN requires committing staged writes first".into(),
5642 ));
5643 }
5644 let old = self
5645 .schema
5646 .columns
5647 .iter()
5648 .find(|c| c.name == column_name)
5649 .cloned()
5650 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5651 let mut next = old.clone();
5652
5653 if let Some(name) = &change.name {
5654 let trimmed = name.trim();
5655 if trimmed.is_empty() {
5656 return Err(MongrelError::InvalidArgument(
5657 "ALTER COLUMN name must not be empty".into(),
5658 ));
5659 }
5660 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
5661 return Err(MongrelError::Schema(format!(
5662 "column {trimmed} already exists"
5663 )));
5664 }
5665 next.name = trimmed.to_string();
5666 }
5667
5668 if let Some(ty) = change.ty {
5669 next.ty = ty;
5670 }
5671 if let Some(flags) = change.flags {
5672 validate_alter_column_flags(old.flags, flags)?;
5673 next.flags = flags;
5674 }
5675
5676 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
5677 if old.flags.contains(ColumnFlags::NULLABLE)
5678 && !next.flags.contains(ColumnFlags::NULLABLE)
5679 && self.column_has_nulls(old.id)?
5680 {
5681 return Err(MongrelError::InvalidArgument(format!(
5682 "column '{}' contains NULL values",
5683 old.name
5684 )));
5685 }
5686 Ok(next)
5687 }
5688
5689 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
5690 let idx = self
5691 .schema
5692 .columns
5693 .iter()
5694 .position(|c| c.id == column.id)
5695 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
5696 if self.schema.columns[idx] == column {
5697 return Ok(());
5698 }
5699 self.schema.columns[idx] = column;
5700 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5701 self.schema.validate_auto_increment()?;
5702 self.auto_inc = resolve_auto_inc(&self.schema);
5703 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
5704 write_schema(&self.dir, &self.schema)?;
5705 self.clear_result_cache();
5706 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5707 self.persist_manifest(self.current_epoch())?;
5708 Ok(())
5709 }
5710
5711 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
5712 let column = self.prepare_alter_column(column_name, &change)?;
5713 self.apply_altered_column(column.clone())?;
5714 Ok(column)
5715 }
5716
5717 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
5718 if self.live_count == 0 {
5719 return Ok(false);
5720 }
5721 let snap = self.snapshot();
5722 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
5723 Ok(columns
5724 .first()
5725 .map(|(_, col)| col.null_count(col.len()) != 0)
5726 .unwrap_or(true))
5727 }
5728
5729 fn has_stored_versions(&self) -> bool {
5730 !self.memtable.is_empty()
5731 || !self.mutable_run.is_empty()
5732 || self.run_refs.iter().any(|r| r.row_count > 0)
5733 || !self.retiring.is_empty()
5734 }
5735
5736 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
5741 if self.schema.columns.iter().any(|c| c.name == name) {
5742 return Err(MongrelError::Schema(format!(
5743 "column {name} already exists"
5744 )));
5745 }
5746 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
5747 self.schema.columns.push(ColumnDef {
5748 id,
5749 name: name.to_string(),
5750 ty,
5751 flags,
5752 });
5753 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5754 self.schema.validate_auto_increment()?;
5755 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
5756 self.auto_inc = resolve_auto_inc(&self.schema);
5757 }
5758 write_schema(&self.dir, &self.schema)?;
5759 self.clear_result_cache();
5760 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5762 self.persist_manifest(self.current_epoch())?;
5763 Ok(id)
5764 }
5765
5766 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
5775 let cid = self
5776 .schema
5777 .columns
5778 .iter()
5779 .find(|c| c.name == column_name)
5780 .map(|c| c.id)
5781 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5782 let ty = self
5783 .schema
5784 .columns
5785 .iter()
5786 .find(|c| c.id == cid)
5787 .map(|c| c.ty)
5788 .unwrap_or(TypeId::Int64);
5789 if !matches!(
5790 ty,
5791 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
5792 ) {
5793 return Err(MongrelError::Schema(format!(
5794 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
5795 )));
5796 }
5797 if self
5798 .schema
5799 .indexes
5800 .iter()
5801 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
5802 {
5803 return Ok(()); }
5805 self.schema.indexes.push(IndexDef {
5806 name: format!("{}_learned_range", column_name),
5807 column_id: cid,
5808 kind: IndexKind::LearnedRange,
5809 });
5810 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5811 write_schema(&self.dir, &self.schema)?;
5812 self.build_learned_ranges()?;
5813 Ok(())
5814 }
5815
5816 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
5819 self.sync_byte_threshold = threshold;
5820 if let WalSink::Private(w) = &mut self.wal {
5821 w.set_sync_byte_threshold(threshold);
5822 }
5823 }
5824
5825 pub fn page_cache_flush(&self) {
5829 self.page_cache.lock().flush_to_disk();
5830 }
5831
5832 pub fn page_cache_len(&self) -> usize {
5834 self.page_cache.lock().len()
5835 }
5836
5837 pub fn decoded_cache_len(&self) -> usize {
5840 self.decoded_cache.lock().len()
5841 }
5842
5843 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
5846 self.memtable.drain_sorted()
5847 }
5848
5849 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
5850 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
5851 }
5852
5853 pub(crate) fn table_dir(&self) -> &Path {
5854 &self.dir
5855 }
5856
5857 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
5858 &self.schema
5859 }
5860
5861 pub(crate) fn alloc_run_id(&mut self) -> u64 {
5862 let id = self.next_run_id;
5863 self.next_run_id += 1;
5864 id
5865 }
5866
5867 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
5868 self.run_refs.push(run_ref);
5869 }
5870
5871 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
5881 self.retiring.push(crate::manifest::RetiredRun {
5882 run_id,
5883 retire_epoch,
5884 });
5885 }
5886
5887 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
5891 if self.retiring.is_empty() {
5892 return Ok(0);
5893 }
5894 let mut reaped = 0;
5895 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
5896 for r in std::mem::take(&mut self.retiring) {
5902 if min_active.0 >= r.retire_epoch {
5903 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
5904 reaped += 1;
5905 } else {
5906 kept.push(r);
5907 }
5908 }
5909 self.retiring = kept;
5910 if reaped > 0 {
5911 self.persist_manifest(self.current_epoch())?;
5912 }
5913 Ok(reaped)
5914 }
5915
5916 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
5917 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
5918 return false;
5919 }
5920 self.live_count = self.live_count.saturating_add(run_ref.row_count);
5921 self.run_refs.push(run_ref);
5922 self.indexes_complete = false;
5923 true
5924 }
5925
5926 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
5927 self.kek.as_ref()
5928 }
5929
5930 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
5931 let mut reader = RunReader::open_with_cache(
5932 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
5933 self.schema.clone(),
5934 self.kek.clone(),
5935 Some(self.page_cache.clone()),
5936 Some(self.decoded_cache.clone()),
5937 self.table_id,
5938 )?;
5939 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
5943 reader.set_uniform_epoch(Epoch(rr.epoch_created));
5944 }
5945 Ok(reader)
5946 }
5947
5948 pub(crate) fn run_refs(&self) -> &[RunRef] {
5949 &self.run_refs
5950 }
5951
5952 pub(crate) fn runs_dir(&self) -> PathBuf {
5953 self.dir.join(RUNS_DIR)
5954 }
5955
5956 pub(crate) fn wal_dir(&self) -> PathBuf {
5957 self.dir.join(WAL_DIR)
5958 }
5959
5960 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
5961 self.run_refs = refs;
5962 }
5963
5964 pub(crate) fn next_run_id(&self) -> u64 {
5965 self.next_run_id
5966 }
5967
5968 pub(crate) fn compaction_zstd_level(&self) -> i32 {
5969 self.compaction_zstd_level
5970 }
5971
5972 pub(crate) fn bump_next_run_id(&mut self) {
5973 self.next_run_id += 1;
5974 }
5975
5976 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
5977 self.kek.clone()
5978 }
5979
5980 #[cfg(feature = "encryption")]
5984 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
5985 self.kek.as_ref().map(|k| k.derive_idx_key())
5986 }
5987
5988 #[cfg(not(feature = "encryption"))]
5989 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
5990 None
5991 }
5992
5993 #[cfg(feature = "encryption")]
5997 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
5998 self.kek.as_ref().map(|k| *k.derive_meta_key())
5999 }
6000
6001 #[cfg(not(feature = "encryption"))]
6002 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6003 None
6004 }
6005
6006 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6009 self.column_keys
6010 .iter()
6011 .map(|(&id, &(_, scheme))| (id, scheme))
6012 .collect()
6013 }
6014
6015 #[cfg(feature = "encryption")]
6020 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6021 self.tokenize_value_enc(column_id, v)
6022 }
6023
6024 #[cfg(feature = "encryption")]
6025 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6026 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6027 let (key, scheme) = self.column_keys.get(&column_id)?;
6028 let token: Vec<u8> = match (*scheme, v) {
6029 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6030 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6031 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6032 _ => hmac_token(key, &v.encode_key()).to_vec(),
6033 };
6034 Some(Value::Bytes(token))
6035 }
6036
6037 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6039 self.index_lookup_key_bytes(column_id, &v.encode_key())
6040 }
6041
6042 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6045 #[cfg(feature = "encryption")]
6046 {
6047 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6048 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6049 if *scheme == SCHEME_HMAC_EQ {
6050 return hmac_token(key, encoded).to_vec();
6051 }
6052 }
6053 }
6054 let _ = column_id;
6055 encoded.to_vec()
6056 }
6057}
6058
6059fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6060 let columnar::NativeColumn::Int64 { data, validity } = col else {
6061 return false;
6062 };
6063 if data.len() < n || !columnar::all_non_null(validity, n) {
6064 return false;
6065 }
6066 data.iter()
6067 .take(n)
6068 .zip(data.iter().skip(1))
6069 .all(|(a, b)| a < b)
6070}
6071
6072#[derive(Debug, Clone)]
6076pub struct ColumnStat {
6077 pub min: Option<Value>,
6078 pub max: Option<Value>,
6079 pub null_count: u64,
6080}
6081
6082#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6084pub enum NativeAgg {
6085 Count,
6086 Sum,
6087 Min,
6088 Max,
6089 Avg,
6090}
6091
6092#[derive(Debug, Clone, PartialEq)]
6094pub enum NativeAggResult {
6095 Count(u64),
6096 Int(i64),
6097 Float(f64),
6098 Null,
6100}
6101
6102#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6104pub enum ApproxAgg {
6105 Count,
6106 Sum,
6107 Avg,
6108}
6109
6110#[derive(Debug, Clone)]
6114pub struct ApproxResult {
6115 pub point: f64,
6117 pub ci_low: f64,
6119 pub ci_high: f64,
6121 pub n_population: u64,
6123 pub n_sample_live: usize,
6125 pub n_passing: usize,
6127}
6128
6129#[derive(Debug, Clone, PartialEq)]
6134pub enum AggState {
6135 Count(u64),
6137 SumI {
6139 sum: i128,
6140 count: u64,
6141 },
6142 SumF {
6144 sum: f64,
6145 count: u64,
6146 },
6147 AvgI {
6149 sum: i128,
6150 count: u64,
6151 },
6152 AvgF {
6154 sum: f64,
6155 count: u64,
6156 },
6157 MinI(i64),
6159 MaxI(i64),
6160 MinF(f64),
6162 MaxF(f64),
6163 Empty,
6165}
6166
6167impl AggState {
6168 pub fn merge(self, other: AggState) -> AggState {
6170 use AggState::*;
6171 match (self, other) {
6172 (Empty, x) | (x, Empty) => x,
6173 (Count(a), Count(b)) => Count(a + b),
6174 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6175 sum: sa + sb,
6176 count: ca + cb,
6177 },
6178 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6179 sum: sa + sb,
6180 count: ca + cb,
6181 },
6182 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6183 sum: sa + sb,
6184 count: ca + cb,
6185 },
6186 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6187 sum: sa + sb,
6188 count: ca + cb,
6189 },
6190 (MinI(a), MinI(b)) => MinI(a.min(b)),
6191 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6192 (MinF(a), MinF(b)) => MinF(a.min(b)),
6193 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6194 _ => Empty, }
6196 }
6197
6198 pub fn point(&self) -> Option<f64> {
6200 match self {
6201 AggState::Count(n) => Some(*n as f64),
6202 AggState::SumI { sum, .. } => Some(*sum as f64),
6203 AggState::SumF { sum, .. } => Some(*sum),
6204 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6205 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6206 AggState::MinI(n) => Some(*n as f64),
6207 AggState::MaxI(n) => Some(*n as f64),
6208 AggState::MinF(n) => Some(*n),
6209 AggState::MaxF(n) => Some(*n),
6210 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6211 }
6212 }
6213
6214 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6218 let is_float = matches!(ty, Some(TypeId::Float64));
6219 match (agg, result) {
6220 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6221 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6222 sum: x as i128,
6223 count: 1, },
6225 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6226 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6227 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6228 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6229 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6230 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6231 (NativeAgg::Count, _) => AggState::Empty,
6232 (_, NativeAggResult::Null) => AggState::Empty,
6233 _ => {
6234 let _ = is_float;
6235 AggState::Empty
6236 }
6237 }
6238 }
6239}
6240
6241#[derive(Debug, Clone)]
6244pub struct CachedAgg {
6245 pub state: AggState,
6246 pub watermark: u64,
6247 pub epoch: u64,
6248}
6249
6250#[derive(Debug, Clone)]
6252pub struct IncrementalAggResult {
6253 pub state: AggState,
6255 pub incremental: bool,
6258 pub delta_rows: u64,
6260}
6261
6262fn agg_state_from_rows(
6266 rows: &[Row],
6267 conditions: &[crate::query::Condition],
6268 index_sets: &[RowIdSet],
6269 column: Option<u16>,
6270 agg: NativeAgg,
6271 schema: &Schema,
6272) -> Result<AggState> {
6273 let mut count: u64 = 0;
6274 let mut sum_i: i128 = 0;
6275 let mut sum_f: f64 = 0.0;
6276 let mut mn_i: i64 = i64::MAX;
6277 let mut mx_i: i64 = i64::MIN;
6278 let mut mn_f: f64 = f64::INFINITY;
6279 let mut mx_f: f64 = f64::NEG_INFINITY;
6280 let mut saw_int = false;
6281 let mut saw_float = false;
6282 for r in rows {
6283 if !conditions
6284 .iter()
6285 .all(|c| condition_matches_row(c, r, schema))
6286 {
6287 continue;
6288 }
6289 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6290 continue;
6291 }
6292 match agg {
6293 NativeAgg::Count => match column {
6294 None => count += 1,
6296 Some(cid) => match r.columns.get(&cid) {
6299 None | Some(Value::Null) => {}
6300 Some(_) => count += 1,
6301 },
6302 },
6303 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6304 Some(Value::Int64(n)) => {
6305 count += 1;
6306 sum_i += *n as i128;
6307 mn_i = mn_i.min(*n);
6308 mx_i = mx_i.max(*n);
6309 saw_int = true;
6310 }
6311 Some(Value::Float64(f)) => {
6312 count += 1;
6313 sum_f += f;
6314 mn_f = mn_f.min(*f);
6315 mx_f = mx_f.max(*f);
6316 saw_float = true;
6317 }
6318 _ => {}
6319 },
6320 }
6321 }
6322 Ok(match agg {
6323 NativeAgg::Count => {
6324 if count == 0 {
6325 AggState::Empty
6326 } else {
6327 AggState::Count(count)
6328 }
6329 }
6330 NativeAgg::Sum => {
6331 if count == 0 {
6332 AggState::Empty
6333 } else if saw_int {
6334 AggState::SumI { sum: sum_i, count }
6335 } else {
6336 AggState::SumF { sum: sum_f, count }
6337 }
6338 }
6339 NativeAgg::Avg => {
6340 if count == 0 {
6341 AggState::Empty
6342 } else if saw_int {
6343 AggState::AvgI { sum: sum_i, count }
6344 } else {
6345 AggState::AvgF { sum: sum_f, count }
6346 }
6347 }
6348 NativeAgg::Min => {
6349 if !saw_int && !saw_float {
6350 AggState::Empty
6351 } else if saw_int {
6352 AggState::MinI(mn_i)
6353 } else {
6354 AggState::MinF(mn_f)
6355 }
6356 }
6357 NativeAgg::Max => {
6358 if !saw_int && !saw_float {
6359 AggState::Empty
6360 } else if saw_int {
6361 AggState::MaxI(mx_i)
6362 } else {
6363 AggState::MaxF(mx_f)
6364 }
6365 }
6366 })
6367}
6368
6369fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6373 use crate::query::Condition;
6374 match c {
6375 Condition::Pk(key) => match schema.primary_key() {
6376 Some(pk) => row
6377 .columns
6378 .get(&pk.id)
6379 .map(|v| v.encode_key() == *key)
6380 .unwrap_or(false),
6381 None => false,
6382 },
6383 Condition::BitmapEq { column_id, value } => row
6384 .columns
6385 .get(column_id)
6386 .map(|v| v.encode_key() == *value)
6387 .unwrap_or(false),
6388 Condition::BitmapIn { column_id, values } => {
6389 let key = row.columns.get(column_id).map(|v| v.encode_key());
6390 match key {
6391 Some(k) => values.contains(&k),
6392 None => false,
6393 }
6394 }
6395 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6396 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6397 _ => false,
6398 },
6399 Condition::RangeF64 {
6400 column_id,
6401 lo,
6402 lo_inclusive,
6403 hi,
6404 hi_inclusive,
6405 } => match row.columns.get(column_id) {
6406 Some(Value::Float64(n)) => {
6407 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6408 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6409 lo_ok && hi_ok
6410 }
6411 _ => false,
6412 },
6413 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6414 Some(Value::Bytes(b)) => {
6415 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6416 }
6417 _ => false,
6418 },
6419 Condition::FmContainsAll {
6420 column_id,
6421 patterns,
6422 } => match row.columns.get(column_id) {
6423 Some(Value::Bytes(b)) => patterns
6424 .iter()
6425 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6426 _ => false,
6427 },
6428 Condition::Ann { .. }
6429 | Condition::SparseMatch { .. }
6430 | Condition::MinHashSimilar { .. } => true,
6431 Condition::IsNull { column_id } => {
6432 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6433 }
6434 Condition::IsNotNull { column_id } => {
6435 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6436 }
6437 }
6438}
6439
6440fn as_f64(v: Option<&Value>) -> Option<f64> {
6442 match v {
6443 Some(Value::Int64(n)) => Some(*n as f64),
6444 Some(Value::Float64(f)) => Some(*f),
6445 _ => None,
6446 }
6447}
6448
6449fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6453 let mut count: u64 = 0;
6454 let mut sum: i128 = 0;
6455 let mut mn: i64 = i64::MAX;
6456 let mut mx: i64 = i64::MIN;
6457 while let Some(cols) = cursor.next_batch()? {
6458 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6459 if crate::columnar::all_non_null(validity, data.len()) {
6460 count += data.len() as u64;
6462 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6463 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6464 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6465 } else {
6466 for (i, &v) in data.iter().enumerate() {
6467 if crate::columnar::validity_bit(validity, i) {
6468 count += 1;
6469 sum += v as i128;
6470 mn = mn.min(v);
6471 mx = mx.max(v);
6472 }
6473 }
6474 }
6475 }
6476 }
6477 Ok((count, sum, mn, mx))
6478}
6479
6480fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6482 let mut count: u64 = 0;
6483 let mut sum: f64 = 0.0;
6484 let mut mn: f64 = f64::INFINITY;
6485 let mut mx: f64 = f64::NEG_INFINITY;
6486 while let Some(cols) = cursor.next_batch()? {
6487 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6488 if crate::columnar::all_non_null(validity, data.len()) {
6489 count += data.len() as u64;
6490 sum += data.iter().sum::<f64>();
6491 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6492 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6493 } else {
6494 for (i, &v) in data.iter().enumerate() {
6495 if crate::columnar::validity_bit(validity, i) {
6496 count += 1;
6497 sum += v;
6498 mn = mn.min(v);
6499 mx = mx.max(v);
6500 }
6501 }
6502 }
6503 }
6504 }
6505 Ok((count, sum, mn, mx))
6506}
6507
6508fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6509 if count == 0 && !matches!(agg, NativeAgg::Count) {
6510 return NativeAggResult::Null;
6511 }
6512 match agg {
6513 NativeAgg::Count => NativeAggResult::Count(count),
6514 NativeAgg::Sum => match sum.try_into() {
6517 Ok(v) => NativeAggResult::Int(v),
6518 Err(_) => NativeAggResult::Null,
6519 },
6520 NativeAgg::Min => NativeAggResult::Int(mn),
6521 NativeAgg::Max => NativeAggResult::Int(mx),
6522 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6523 }
6524}
6525
6526fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6527 if count == 0 && !matches!(agg, NativeAgg::Count) {
6528 return NativeAggResult::Null;
6529 }
6530 match agg {
6531 NativeAgg::Count => NativeAggResult::Count(count),
6532 NativeAgg::Sum => NativeAggResult::Float(sum),
6533 NativeAgg::Min => NativeAggResult::Float(mn),
6534 NativeAgg::Max => NativeAggResult::Float(mx),
6535 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6536 }
6537}
6538
6539fn agg_int(
6542 stats: &[crate::page::PageStat],
6543 decode: fn(Option<&[u8]>) -> Option<i64>,
6544) -> Option<(Option<i64>, Option<i64>, u64)> {
6545 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6546 let mut any = false;
6547 for s in stats {
6548 if let Some(v) = decode(s.min.as_deref()) {
6549 mn = mn.min(v);
6550 any = true;
6551 }
6552 if let Some(v) = decode(s.max.as_deref()) {
6553 mx = mx.max(v);
6554 any = true;
6555 }
6556 nulls += s.null_count;
6557 }
6558 any.then_some((Some(mn), Some(mx), nulls))
6559}
6560
6561fn agg_float(
6563 stats: &[crate::page::PageStat],
6564 decode: fn(Option<&[u8]>) -> Option<f64>,
6565) -> Option<(Option<f64>, Option<f64>, u64)> {
6566 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6567 let mut any = false;
6568 for s in stats {
6569 if let Some(v) = decode(s.min.as_deref()) {
6570 mn = mn.min(v);
6571 any = true;
6572 }
6573 if let Some(v) = decode(s.max.as_deref()) {
6574 mx = mx.max(v);
6575 any = true;
6576 }
6577 nulls += s.null_count;
6578 }
6579 any.then_some((Some(mn), Some(mx), nulls))
6580}
6581
6582type SecondaryIndexes = (
6584 HashMap<u16, BitmapIndex>,
6585 HashMap<u16, AnnIndex>,
6586 HashMap<u16, FmIndex>,
6587 HashMap<u16, SparseIndex>,
6588 HashMap<u16, MinHashIndex>,
6589);
6590
6591fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6592 let mut bitmap = HashMap::new();
6593 let mut ann = HashMap::new();
6594 let mut fm = HashMap::new();
6595 let mut sparse = HashMap::new();
6596 let mut minhash = HashMap::new();
6597 for idef in &schema.indexes {
6598 match idef.kind {
6599 IndexKind::Bitmap => {
6600 bitmap.insert(idef.column_id, BitmapIndex::new());
6601 }
6602 IndexKind::Ann => {
6603 let dim = schema
6604 .columns
6605 .iter()
6606 .find(|c| c.id == idef.column_id)
6607 .and_then(|c| match c.ty {
6608 TypeId::Embedding { dim } => Some(dim as usize),
6609 _ => None,
6610 })
6611 .unwrap_or(0);
6612 ann.insert(idef.column_id, AnnIndex::new(dim));
6613 }
6614 IndexKind::FmIndex => {
6615 fm.insert(idef.column_id, FmIndex::new());
6616 }
6617 IndexKind::Sparse => {
6618 sparse.insert(idef.column_id, SparseIndex::new());
6619 }
6620 IndexKind::MinHash => {
6621 minhash.insert(idef.column_id, MinHashIndex::new());
6622 }
6623 _ => {}
6624 }
6625 }
6626 (bitmap, ann, fm, sparse, minhash)
6627}
6628
6629const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
6630 | ColumnFlags::AUTO_INCREMENT
6631 | ColumnFlags::ENCRYPTED
6632 | ColumnFlags::ENCRYPTED_INDEXABLE
6633 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
6634
6635fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
6636 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
6637 return Err(MongrelError::Schema(
6638 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
6639 ));
6640 }
6641 Ok(())
6642}
6643
6644fn validate_alter_column_type(
6645 schema: &Schema,
6646 old: &ColumnDef,
6647 next: &ColumnDef,
6648 has_stored_versions: bool,
6649) -> Result<()> {
6650 if old.ty == next.ty {
6651 return Ok(());
6652 }
6653 if schema.indexes.iter().any(|i| i.column_id == old.id) {
6654 return Err(MongrelError::Schema(format!(
6655 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
6656 old.name
6657 )));
6658 }
6659 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
6660 return Ok(());
6661 }
6662 Err(MongrelError::Schema(format!(
6663 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
6664 old.ty, next.ty
6665 )))
6666}
6667
6668fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
6669 matches!(
6670 (old, new),
6671 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
6672 )
6673}
6674
6675fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
6681 let mut prev: Option<i64> = None;
6682 for r in rows {
6683 match r.columns.get(&pk_id) {
6684 Some(Value::Int64(v)) => {
6685 if prev.is_some_and(|p| p >= *v) {
6686 return false;
6687 }
6688 prev = Some(*v);
6689 }
6690 _ => return false,
6691 }
6692 }
6693 true
6694}
6695
6696#[allow(clippy::too_many_arguments)]
6697fn index_into(
6698 schema: &Schema,
6699 row: &Row,
6700 hot: &mut HotIndex,
6701 bitmap: &mut HashMap<u16, BitmapIndex>,
6702 ann: &mut HashMap<u16, AnnIndex>,
6703 fm: &mut HashMap<u16, FmIndex>,
6704 sparse: &mut HashMap<u16, SparseIndex>,
6705 minhash: &mut HashMap<u16, MinHashIndex>,
6706) {
6707 for idef in &schema.indexes {
6708 let Some(val) = row.columns.get(&idef.column_id) else {
6709 continue;
6710 };
6711 match idef.kind {
6712 IndexKind::Bitmap => {
6713 if let Some(b) = bitmap.get_mut(&idef.column_id) {
6714 b.insert(val.encode_key(), row.row_id);
6715 }
6716 }
6717 IndexKind::Ann => {
6718 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
6719 a.insert(v, row.row_id);
6720 }
6721 }
6722 IndexKind::FmIndex => {
6723 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
6724 f.insert(b.clone(), row.row_id);
6725 }
6726 }
6727 IndexKind::Sparse => {
6728 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
6729 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
6732 s.insert(&terms, row.row_id);
6733 }
6734 }
6735 }
6736 IndexKind::MinHash => {
6737 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
6738 let tokens = crate::index::token_hashes_from_bytes(b);
6741 mh.insert(&tokens, row.row_id);
6742 }
6743 }
6744 _ => {}
6745 }
6746 }
6747 if let Some(pk_col) = schema.primary_key() {
6748 if let Some(pk_val) = row.columns.get(&pk_col.id) {
6749 hot.insert(pk_val.encode_key(), row.row_id);
6750 }
6751 }
6752}
6753
6754#[allow(dead_code)]
6760fn bulk_index_key(
6761 column_keys: &HashMap<u16, ([u8; 32], u8)>,
6762 column_id: u16,
6763 ty: TypeId,
6764 col: &columnar::NativeColumn,
6765 i: usize,
6766) -> Option<Vec<u8>> {
6767 let encoded = columnar::encode_key_native(ty, col, i)?;
6768 #[cfg(feature = "encryption")]
6769 {
6770 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6771 if let Some((key, scheme)) = column_keys.get(&column_id) {
6772 return Some(match (*scheme, col) {
6773 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
6774 (_, columnar::NativeColumn::Int64 { data, .. }) => {
6775 ope_token_i64(key, data[i]).to_vec()
6776 }
6777 (_, columnar::NativeColumn::Float64 { data, .. }) => {
6778 ope_token_f64(key, data[i]).to_vec()
6779 }
6780 _ => hmac_token(key, &encoded).to_vec(),
6781 });
6782 }
6783 }
6784 #[cfg(not(feature = "encryption"))]
6785 {
6786 let _ = (column_id, column_keys, col);
6787 }
6788 Some(encoded)
6789}
6790
6791pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
6792 let json = serde_json::to_string_pretty(schema)
6793 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
6794 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
6795 Ok(())
6796}
6797
6798fn read_schema(dir: &Path) -> Result<Schema> {
6799 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
6800 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
6801}
6802
6803fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
6804 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
6805}
6806
6807fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
6808 let n = list_wal_numbers(wal_dir)?;
6809 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
6810}
6811
6812fn next_wal_number(wal_dir: &Path) -> Result<u32> {
6813 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
6814}
6815
6816fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
6817 let _ = std::fs::create_dir_all(wal_dir);
6818 let mut max_n = None;
6819 for entry in std::fs::read_dir(wal_dir)? {
6820 let entry = entry?;
6821 let fname = entry.file_name();
6822 let Some(s) = fname.to_str() else {
6823 continue;
6824 };
6825 let Some(stripped) = s.strip_prefix("seg-") else {
6826 continue;
6827 };
6828 let Some(stripped) = stripped.strip_suffix(".wal") else {
6829 continue;
6830 };
6831 if let Ok(n) = stripped.parse::<u32>() {
6832 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
6833 }
6834 }
6835 Ok(max_n)
6836}