1use crate::columnar;
11use crate::cursor::NativePageCursor;
12use crate::encryption::Kek;
13use crate::encryption::DEK_LEN;
14use crate::epoch::{Epoch, EpochAuthority, Snapshot};
15use crate::global_idx;
16use crate::index::{
17 AnnIndex, BitmapIndex, ColumnLearnedRange, FmIndex, HotIndex, MinHashIndex, SparseIndex,
18};
19use crate::manifest::{self, Manifest, RunRef};
20use crate::memtable::{Memtable, Row, Value};
21use crate::mutable_run::MutableRun;
22use crate::row_id_set::RowIdSet;
23use crate::rowid::{RowId, RowIdAllocator};
24use crate::schema::{AlterColumn, ColumnDef, ColumnFlags, IndexDef, IndexKind, Schema, TypeId};
25use crate::sorted_run::{RunReader, RunWriter};
26use crate::txn::{GroupCommit, OwnedRow};
27use crate::wal::{Op, SharedWal, Wal};
28use crate::{MongrelError, Result};
29use std::collections::{BTreeMap, HashMap, HashSet};
30use std::path::{Path, PathBuf};
31use std::sync::atomic::AtomicBool;
32use std::sync::Arc;
33use zeroize::Zeroizing;
34
35pub const WAL_DIR: &str = "_wal";
36pub const RUNS_DIR: &str = "_runs";
37pub const CACHE_DIR: &str = "_cache";
38pub const META_DIR: &str = "_meta";
39pub const RCACHE_DIR: &str = "_rcache";
40pub const KEYS_FILENAME: &str = "keys";
41pub const SCHEMA_FILENAME: &str = "schema.json";
42const DEFAULT_SYNC_BYTE_THRESHOLD: u64 = 0; pub(crate) const PAGE_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; pub(crate) const DECODED_CACHE_CAPACITY: u64 = 64 * 1024 * 1024; const DEFAULT_MUTABLE_RUN_SPILL_BYTES: u64 = 8 * 1024 * 1024;
49
50#[derive(Clone, Copy, Debug)]
65struct AutoIncState {
66 column_id: u16,
67 next: i64,
68 seeded: bool,
69}
70
71type FilledAutoIncRow = (Vec<(u16, Value)>, Option<i64>);
72
73fn resolve_auto_inc(schema: &Schema) -> Option<AutoIncState> {
76 schema.auto_increment_column().map(|c| AutoIncState {
77 column_id: c.id,
78 next: 0,
79 seeded: false,
80 })
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
96pub enum IndexBuildPolicy {
97 #[default]
99 Deferred,
100 Eager,
102}
103
104pub struct Table {
106 dir: PathBuf,
107 table_id: u64,
108 wal: WalSink,
109 memtable: Memtable,
110 mutable_run: MutableRun,
115 mutable_run_spill_bytes: u64,
117 compaction_zstd_level: i32,
120 allocator: RowIdAllocator,
121 epoch: Arc<EpochAuthority>,
122 persisted_epoch: u64,
125 schema: Schema,
126 hot: HotIndex,
127 kek: Option<Arc<Kek>>,
130 column_keys: HashMap<u16, ([u8; 32], u8)>,
134 run_refs: Vec<RunRef>,
135 retiring: Vec<crate::manifest::RetiredRun>,
138 next_run_id: u64,
139 sync_byte_threshold: u64,
140 current_txn_id: u64,
145 bitmap: HashMap<u16, BitmapIndex>,
146 ann: HashMap<u16, AnnIndex>,
147 fm: HashMap<u16, FmIndex>,
148 sparse: HashMap<u16, SparseIndex>,
149 minhash: HashMap<u16, MinHashIndex>,
150 learned_range: HashMap<u16, ColumnLearnedRange>,
153 pk_by_row: HashMap<RowId, Vec<u8>>,
155 pinned: BTreeMap<Epoch, usize>,
158 pub(crate) live_count: u64,
161 reservoir: crate::reservoir::Reservoir,
164 reservoir_complete: bool,
172 had_deletes: bool,
176 agg_cache: HashMap<u64, CachedAgg>,
180 global_idx_epoch: u64,
184 indexes_complete: bool,
189 index_build_policy: IndexBuildPolicy,
191 pk_by_row_complete: bool,
198 flushed_epoch: u64,
201 page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
204 snapshots: Arc<crate::retention::SnapshotRegistry>,
207 commit_lock: Arc<parking_lot::Mutex<()>>,
209 decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
212 verified_runs: Arc<parking_lot::Mutex<std::collections::HashSet<u128>>>,
222 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
231 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
233 pending_delete_rids: roaring::RoaringBitmap,
236 pending_put_cols: std::collections::HashSet<u16>,
239 pending_rows: Vec<Row>,
245 pending_rows_auto_inc: Vec<bool>,
246 pending_dels: Vec<RowId>,
249 pending_truncate: Option<Epoch>,
253 auto_inc: Option<AutoIncState>,
256}
257
258const _: () = {
265 const fn assert_sync<T: ?Sized + Sync>() {}
266 assert_sync::<Table>();
267};
268
269enum CachedData {
275 Rows(Arc<Vec<Row>>),
276 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
277}
278
279impl CachedData {
280 fn approx_bytes(&self) -> u64 {
281 match self {
282 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
283 CachedData::Columns(c) => c
284 .iter()
285 .map(|(_, c)| c.approx_bytes())
286 .sum::<u64>()
287 .saturating_add(c.len() as u64 * 16),
288 }
289 }
290}
291
292struct CachedEntry {
296 data: CachedData,
297 footprint: roaring::RoaringBitmap,
298 condition_cols: Vec<u16>,
299}
300
301struct ResultCache {
312 entries: std::collections::HashMap<u64, CachedEntry>,
313 order: std::collections::VecDeque<u64>,
314 bytes: u64,
315 max_bytes: u64,
316 dir: Option<std::path::PathBuf>,
317 #[allow(dead_code)]
318 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
319}
320
321#[derive(serde::Serialize, serde::Deserialize)]
323struct SerializedEntry {
324 condition_cols: Vec<u16>,
325 footprint_bits: Vec<u32>,
326 data: SerializedData,
327}
328
329#[derive(serde::Serialize, serde::Deserialize)]
330enum SerializedData {
331 Rows(Vec<Row>),
332 Columns(Vec<(u16, columnar::NativeColumn)>),
333}
334
335impl SerializedEntry {
336 fn from_entry(entry: &CachedEntry) -> Self {
337 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
338 let data = match &entry.data {
339 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
340 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
341 };
342 Self {
343 condition_cols: entry.condition_cols.clone(),
344 footprint_bits,
345 data,
346 }
347 }
348
349 fn into_entry(self) -> Option<CachedEntry> {
350 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
351 let data = match self.data {
352 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
353 SerializedData::Columns(c) => {
354 if !c.iter().all(|(_, col)| col.validate()) {
357 return None;
358 }
359 CachedData::Columns(Arc::new(c))
360 }
361 };
362 Some(CachedEntry {
363 data,
364 footprint,
365 condition_cols: self.condition_cols,
366 })
367 }
368}
369
370impl ResultCache {
371 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
372
373 fn new() -> Self {
374 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
375 }
376
377 fn with_max_bytes(max_bytes: u64) -> Self {
378 Self {
379 entries: std::collections::HashMap::new(),
380 order: std::collections::VecDeque::new(),
381 bytes: 0,
382 max_bytes,
383 dir: None,
384 cache_dek: None,
385 }
386 }
387
388 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
389 let _ = std::fs::create_dir_all(&dir);
390 self.dir = Some(dir);
391 self
392 }
393
394 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
395 self.cache_dek = dek;
396 self
397 }
398
399 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
400 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
401 }
402
403 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
407 let Some(path) = self.disk_path(key) else {
408 return;
409 };
410 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
411 Ok(s) => s,
412 Err(_) => return,
413 };
414 let on_disk = if let Some(dek) = &self.cache_dek {
416 match self.encrypt_cache(&serialized, dek) {
417 Some(b) => b,
418 None => return,
419 }
420 } else {
421 serialized
422 };
423 let tmp = path.with_extension("tmp");
424 use std::io::Write;
425 let write = || -> std::io::Result<()> {
426 let mut f = std::fs::File::create(&tmp)?;
427 f.write_all(&on_disk)?;
428 f.flush()?;
429 Ok(())
430 };
431 if write().is_err() {
432 let _ = std::fs::remove_file(&tmp);
433 return;
434 }
435 let _ = std::fs::rename(&tmp, &path);
436 }
437
438 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
440 let path = self.disk_path(key)?;
441 let bytes = std::fs::read(&path).ok()?;
442 let plaintext = if let Some(dek) = &self.cache_dek {
443 self.decrypt_cache(&bytes, dek)?
444 } else {
445 bytes
446 };
447 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
448 serialized.into_entry()
449 }
450
451 fn remove_from_disk(&self, key: u64) {
453 if let Some(path) = self.disk_path(key) {
454 let _ = std::fs::remove_file(&path);
455 }
456 }
457
458 #[cfg(feature = "encryption")]
460 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
461 use crate::encryption::Cipher;
462 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
463 let mut nonce = [0u8; 12];
464 crate::encryption::fill_random(&mut nonce);
465 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
466 let mut out = Vec::with_capacity(12 + ct.len());
467 out.extend_from_slice(&nonce);
468 out.extend_from_slice(&ct);
469 Some(out)
470 }
471
472 #[cfg(not(feature = "encryption"))]
473 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
474 None
475 }
476
477 #[cfg(feature = "encryption")]
479 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
480 use crate::encryption::Cipher;
481 if bytes.len() < 28 {
482 return None;
483 }
484 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
485 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
486 let ct = &bytes[12..];
487 cipher.decrypt_page(&nonce, ct).ok()
488 }
489
490 #[cfg(not(feature = "encryption"))]
491 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
492 None
493 }
494
495 fn load_persistent(&mut self) {
498 let Some(dir) = self.dir.as_ref().cloned() else {
499 return;
500 };
501 let entries = match std::fs::read_dir(&dir) {
502 Ok(e) => e,
503 Err(_) => return,
504 };
505 for entry in entries.flatten() {
506 let path = entry.path();
507 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
509 let _ = std::fs::remove_file(&path);
510 continue;
511 }
512 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
513 continue;
514 }
515 let stem = match path.file_stem().and_then(|s| s.to_str()) {
516 Some(s) => s,
517 None => continue,
518 };
519 let key = match u64::from_str_radix(stem, 16) {
520 Ok(k) => k,
521 Err(_) => continue,
522 };
523 let bytes = match std::fs::read(&path) {
524 Ok(b) => b,
525 Err(_) => continue,
526 };
527 let plaintext = if let Some(dek) = &self.cache_dek {
529 match self.decrypt_cache(&bytes, dek) {
530 Some(p) => p,
531 None => {
532 let _ = std::fs::remove_file(&path);
533 continue;
534 }
535 }
536 } else {
537 bytes
538 };
539 match bincode::deserialize::<SerializedEntry>(&plaintext) {
540 Ok(serialized) => {
541 if let Some(entry) = serialized.into_entry() {
542 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
543 self.entries.insert(key, entry);
544 self.order.push_back(key);
545 } else {
546 let _ = std::fs::remove_file(&path);
547 }
548 }
549 Err(_) => {
550 let _ = std::fs::remove_file(&path);
551 }
552 }
553 }
554 self.evict();
555 }
556
557 fn set_max_bytes(&mut self, max_bytes: u64) {
558 self.max_bytes = max_bytes;
559 self.evict();
560 }
561
562 fn touch(&mut self, key: u64) {
564 self.order.retain(|k| *k != key);
565 self.order.push_back(key);
566 }
567
568 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
569 let res = self.entries.get(&key).and_then(|e| match &e.data {
570 CachedData::Rows(r) => Some(r.clone()),
571 CachedData::Columns(_) => None,
572 });
573 if res.is_some() {
574 self.touch(key);
575 return res;
576 }
577 if let Some(entry) = self.load_from_disk(key) {
579 let res = match &entry.data {
580 CachedData::Rows(r) => Some(r.clone()),
581 CachedData::Columns(_) => None,
582 };
583 if res.is_some() {
584 let approx = entry.data.approx_bytes();
585 self.bytes = self.bytes.saturating_add(approx);
586 self.entries.insert(key, entry);
587 self.order.push_back(key);
588 self.evict();
589 return res;
590 }
591 }
592 None
593 }
594
595 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
596 let res = self.entries.get(&key).and_then(|e| match &e.data {
597 CachedData::Columns(c) => Some(c.clone()),
598 CachedData::Rows(_) => None,
599 });
600 if res.is_some() {
601 self.touch(key);
602 return res;
603 }
604 if let Some(entry) = self.load_from_disk(key) {
606 let res = match &entry.data {
607 CachedData::Columns(c) => Some(c.clone()),
608 CachedData::Rows(_) => None,
609 };
610 if res.is_some() {
611 let approx = entry.data.approx_bytes();
612 self.bytes = self.bytes.saturating_add(approx);
613 self.entries.insert(key, entry);
614 self.order.push_back(key);
615 self.evict();
616 return res;
617 }
618 }
619 None
620 }
621
622 fn insert(&mut self, key: u64, entry: CachedEntry) {
623 let approx = entry.data.approx_bytes();
624 if self.entries.remove(&key).is_some() {
625 self.order.retain(|k| *k != key);
626 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
627 }
628 self.store_to_disk(key, &entry);
630 self.bytes = self.bytes.saturating_add(approx);
631 self.entries.insert(key, entry);
632 self.order.push_back(key);
633 self.evict();
634 }
635
636 fn invalidate(
645 &mut self,
646 delete_rids: &roaring::RoaringBitmap,
647 put_cols: &std::collections::HashSet<u16>,
648 ) {
649 if self.entries.is_empty() {
650 return;
651 }
652 let has_deletes = !delete_rids.is_empty();
653 let to_remove: std::collections::HashSet<u64> = self
654 .entries
655 .iter()
656 .filter(|(_, e)| {
657 let delete_hit = if e.footprint.is_empty() {
658 has_deletes
659 } else {
660 e.footprint.intersection_len(delete_rids) > 0
661 };
662 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
663 delete_hit || col_hit
664 })
665 .map(|(&k, _)| k)
666 .collect();
667 for key in &to_remove {
668 if let Some(e) = self.entries.remove(key) {
669 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
670 }
671 self.remove_from_disk(*key);
672 }
673 if !to_remove.is_empty() {
674 self.order.retain(|k| !to_remove.contains(k));
675 }
676 }
677
678 fn clear(&mut self) {
679 if let Some(dir) = &self.dir {
681 if let Ok(entries) = std::fs::read_dir(dir) {
682 for entry in entries.flatten() {
683 let path = entry.path();
684 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
685 let _ = std::fs::remove_file(&path);
686 }
687 }
688 }
689 }
690 self.entries.clear();
691 self.order.clear();
692 self.bytes = 0;
693 }
694
695 fn evict(&mut self) {
696 while self.bytes > self.max_bytes {
697 let Some(k) = self.order.pop_front() else {
698 break;
699 };
700 if let Some(e) = self.entries.remove(&k) {
701 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
702 self.remove_from_disk(k);
706 }
707 }
708 }
709}
710
711type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
718
719fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
720 let _ = kek;
721 #[cfg(feature = "encryption")]
722 {
723 if let Some(k) = kek {
724 return (
725 Some(k.derive_table_wal_key(_table_id)),
726 Some(k.derive_cache_key()),
727 );
728 }
729 }
730 (None, None)
731}
732
733#[cfg(feature = "encryption")]
735fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
736 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
737}
738
739#[cfg(not(feature = "encryption"))]
740fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
741 Box::new(crate::encryption::PlaintextCipher)
742}
743
744fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
745 let Some(kek) = kek else {
746 return HashMap::new();
747 };
748 #[cfg(feature = "encryption")]
749 {
750 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
751 schema
752 .columns
753 .iter()
754 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
755 .map(|c| {
756 let scheme = if schema
757 .indexes
758 .iter()
759 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
760 {
761 SCHEME_OPE_RANGE
762 } else {
763 SCHEME_HMAC_EQ
764 };
765 let key: [u8; 32] = *kek.derive_column_key(c.id);
766 (c.id, (key, scheme))
767 })
768 .collect()
769 }
770 #[cfg(not(feature = "encryption"))]
771 {
772 let _ = (kek, schema);
773 HashMap::new()
774 }
775}
776
777pub(crate) struct SharedCtx {
782 pub epoch: Arc<EpochAuthority>,
783 pub page_cache: Arc<crate::cache::Sharded<crate::cache::PageCache>>,
784 pub decoded_cache: Arc<crate::cache::Sharded<crate::cache::DecodedPageCache>>,
785 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
786 pub kek: Option<Arc<Kek>>,
787 pub commit_lock: Arc<parking_lot::Mutex<()>>,
793 pub shared: Option<SharedWalCtx>,
797}
798
799#[derive(Clone)]
805pub(crate) struct SharedWalCtx {
806 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
807 pub group: Arc<GroupCommit>,
808 pub poisoned: Arc<AtomicBool>,
809 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
810}
811
812enum WalSink {
815 Private(Wal),
816 Shared(SharedWalCtx),
817}
818
819impl SharedCtx {
820 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
824 let n_shards = if cache_dir.is_some() {
828 1
829 } else {
830 crate::cache::CACHE_SHARDS
831 };
832 let per_shard = PAGE_CACHE_CAPACITY / n_shards as u64;
833 let page_cache = if let Some(d) = cache_dir {
834 Arc::new(crate::cache::Sharded::new(1, || {
835 crate::cache::PageCache::new(PAGE_CACHE_CAPACITY).with_persistence(d.clone())
836 }))
837 } else {
838 Arc::new(crate::cache::Sharded::new(n_shards, || {
839 crate::cache::PageCache::new(per_shard)
840 }))
841 };
842 let decoded_per_shard = DECODED_CACHE_CAPACITY / crate::cache::CACHE_SHARDS as u64;
843 let decoded_cache = Arc::new(crate::cache::Sharded::new(
844 crate::cache::CACHE_SHARDS,
845 || crate::cache::DecodedPageCache::new(decoded_per_shard),
846 ));
847 Self {
848 epoch: Arc::new(EpochAuthority::new(0)),
849 page_cache,
850 decoded_cache,
851 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
852 kek,
853 commit_lock: Arc::new(parking_lot::Mutex::new(())),
854 shared: None,
855 }
856 }
857}
858
859fn condition_cost_rank(c: &crate::query::Condition) -> u8 {
863 use crate::query::Condition;
864 match c {
865 Condition::Pk(_)
867 | Condition::BitmapEq { .. }
868 | Condition::BitmapIn { .. }
869 | Condition::BytesPrefix { .. }
870 | Condition::IsNull { .. }
871 | Condition::IsNotNull { .. } => 0,
872 Condition::Range { .. } | Condition::RangeF64 { .. } | Condition::MinHashSimilar { .. } => {
874 1
875 }
876 Condition::FmContains { .. }
878 | Condition::FmContainsAll { .. }
879 | Condition::Ann { .. }
880 | Condition::SparseMatch { .. } => 2,
881 }
882}
883
884impl Table {
885 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
886 let dir = dir.as_ref().to_path_buf();
887 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
888 Self::create_in(&dir, schema, table_id, ctx)
889 }
890
891 #[cfg(feature = "encryption")]
902 pub fn create_encrypted(
903 dir: impl AsRef<Path>,
904 schema: Schema,
905 table_id: u64,
906 passphrase: &str,
907 ) -> Result<Self> {
908 let dir = dir.as_ref();
909 std::fs::create_dir_all(dir.join(META_DIR))?;
910 let salt = crate::encryption::random_salt();
911 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
912 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
913 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
914 Self::create_in(dir, schema, table_id, ctx)
915 }
916
917 #[cfg(feature = "encryption")]
922 pub fn create_with_key(
923 dir: impl AsRef<Path>,
924 schema: Schema,
925 table_id: u64,
926 key: &[u8],
927 ) -> Result<Self> {
928 let dir = dir.as_ref();
929 std::fs::create_dir_all(dir.join(META_DIR))?;
930 let salt = crate::encryption::random_salt();
931 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
932 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
933 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
934 Self::create_in(dir, schema, table_id, ctx)
935 }
936
937 #[cfg(feature = "encryption")]
939 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
940 let dir = dir.as_ref();
941 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
942 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
943 MongrelError::NotFound(format!(
944 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
945 salt_path
946 ))
947 })?;
948 if salt_bytes.len() != crate::encryption::SALT_LEN {
949 return Err(MongrelError::InvalidArgument(format!(
950 "salt file is {} bytes, expected {}",
951 salt_bytes.len(),
952 crate::encryption::SALT_LEN
953 )));
954 }
955 let mut salt = [0u8; crate::encryption::SALT_LEN];
956 salt.copy_from_slice(&salt_bytes);
957 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
958 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
959 Self::open_in(dir, ctx)
960 }
961
962 pub(crate) fn create_in(
963 dir: impl AsRef<Path>,
964 schema: Schema,
965 table_id: u64,
966 ctx: SharedCtx,
967 ) -> Result<Self> {
968 schema.validate_auto_increment()?;
969 let dir = dir.as_ref().to_path_buf();
970 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
971 write_schema(&dir, &schema)?;
972 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
973 let (wal, current_txn_id) = match ctx.shared.clone() {
976 Some(s) => (WalSink::Shared(s), 0),
977 None => {
978 std::fs::create_dir_all(dir.join(WAL_DIR))?;
979 let mut w = if let Some(ref dk) = wal_dek {
980 Wal::create_with_cipher(
981 dir.join(WAL_DIR).join("seg-000000.wal"),
982 Epoch(0),
983 Some(make_cipher(dk)),
984 0,
985 )?
986 } else {
987 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
988 };
989 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
990 (WalSink::Private(w), 1)
991 }
992 };
993 let mut manifest = Manifest::new(table_id, schema.schema_id);
994 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
999 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
1000 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
1001 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1002 let auto_inc = resolve_auto_inc(&schema);
1003 let rcache_dir = dir.join(RCACHE_DIR);
1004 Ok(Self {
1005 dir,
1006 table_id,
1007 wal,
1008 memtable: Memtable::new(),
1009 mutable_run: MutableRun::new(),
1010 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1011 compaction_zstd_level: 3,
1012 allocator: RowIdAllocator::new(0),
1013 epoch: ctx.epoch,
1014 persisted_epoch: 0,
1015 schema,
1016 hot: HotIndex::new(),
1017 kek: ctx.kek,
1018 column_keys,
1019 run_refs: Vec::new(),
1020 retiring: Vec::new(),
1021 next_run_id: 1,
1022 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1023 current_txn_id,
1024 bitmap,
1025 ann,
1026 fm,
1027 sparse,
1028 minhash,
1029 learned_range: HashMap::new(),
1030 pk_by_row: HashMap::new(),
1031 pinned: BTreeMap::new(),
1032 live_count: 0,
1033 reservoir: crate::reservoir::Reservoir::default(),
1034 reservoir_complete: true,
1035 had_deletes: false,
1036 agg_cache: HashMap::new(),
1037 global_idx_epoch: 0,
1038 indexes_complete: true,
1039 index_build_policy: IndexBuildPolicy::default(),
1040 pk_by_row_complete: false,
1041 flushed_epoch: 0,
1042 page_cache: ctx.page_cache,
1043 decoded_cache: ctx.decoded_cache,
1044 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1045 snapshots: ctx.snapshots,
1046 commit_lock: ctx.commit_lock,
1047 result_cache: Arc::new(parking_lot::Mutex::new(
1048 ResultCache::new()
1049 .with_dir(rcache_dir)
1050 .with_cache_dek(cache_dek.clone()),
1051 )),
1052 pending_delete_rids: roaring::RoaringBitmap::new(),
1053 pending_put_cols: std::collections::HashSet::new(),
1054 pending_rows: Vec::new(),
1055 pending_rows_auto_inc: Vec::new(),
1056 pending_dels: Vec::new(),
1057 pending_truncate: None,
1058 wal_dek,
1059 auto_inc,
1060 })
1061 }
1062
1063 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
1067 let dir = dir.as_ref();
1068 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
1069 Self::open_in(dir, ctx)
1070 }
1071
1072 #[cfg(feature = "encryption")]
1075 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
1076 let dir = dir.as_ref();
1077 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
1078 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
1079 MongrelError::NotFound(format!(
1080 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
1081 salt_path
1082 ))
1083 })?;
1084 let salt_len = crate::encryption::SALT_LEN;
1085 if salt_bytes.len() != salt_len {
1086 return Err(MongrelError::InvalidArgument(format!(
1087 "encryption salt is {} bytes, expected {salt_len}",
1088 salt_bytes.len()
1089 )));
1090 }
1091 let mut salt = [0u8; 16];
1092 salt.copy_from_slice(&salt_bytes);
1093 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1094 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1095 let t = Self::open_in(dir, ctx)?;
1096 Ok(t)
1097 }
1098
1099 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1100 let dir = dir.as_ref().to_path_buf();
1101 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1102 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1103 let schema: Schema = read_schema(&dir)?;
1104 let replay_epoch = Epoch(manifest.current_epoch);
1105 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1106 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1110 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1111 None => {
1112 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1113 let replayed = match &active {
1115 Some(path) => {
1116 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1117 crate::wal::replay_with_cipher(path, cipher)?
1118 }
1119 None => Vec::new(),
1120 };
1121 let mut w = match &active {
1122 Some(path) => Wal::create_with_cipher(
1123 path,
1124 replay_epoch,
1125 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1126 0,
1127 )?,
1128 None => Wal::create_with_cipher(
1129 dir.join(WAL_DIR).join("seg-000000.wal"),
1130 replay_epoch,
1131 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1132 0,
1133 )?,
1134 };
1135 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1136 (WalSink::Private(w), replayed, 1)
1137 }
1138 };
1139
1140 let mut memtable = Memtable::new();
1141 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1142 let persisted_epoch = manifest.current_epoch;
1143 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1150 s.next = manifest.auto_inc_next;
1151 s.seeded = manifest.auto_inc_next > 0;
1152 s
1153 });
1154
1155 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1162 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1163 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1164 std::collections::BTreeMap::new();
1165 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1166 let mut saw_delete = false;
1167 for record in replayed {
1168 let txn_id = record.txn_id;
1169 match record.op {
1170 Op::Put { rows, .. } => {
1171 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1172 for row in &rows {
1173 allocator.advance_to(row.row_id);
1174 if let Some(ai) = auto_inc.as_mut() {
1175 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1176 if *n + 1 > ai.next {
1177 ai.next = *n + 1;
1178 }
1179 }
1180 }
1181 }
1182 staged_puts.entry(txn_id).or_default().extend(rows);
1183 }
1184 Op::Delete { row_ids, .. } => {
1185 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1186 }
1187 Op::TxnCommit { epoch, .. } => {
1188 let commit_epoch = Epoch(epoch);
1189 if let Some(puts) = staged_puts.remove(&txn_id) {
1190 for row in &puts {
1191 memtable.upsert(row.clone());
1192 }
1193 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1194 }
1195 if let Some(dels) = staged_deletes.remove(&txn_id) {
1196 saw_delete = true;
1197 for rid in dels {
1198 memtable.tombstone(rid, commit_epoch);
1199 replayed_deletes.push((rid, commit_epoch));
1200 }
1201 }
1202 }
1203 Op::TxnAbort => {
1204 staged_puts.remove(&txn_id);
1205 staged_deletes.remove(&txn_id);
1206 }
1207 Op::TruncateTable { .. }
1208 | Op::ExternalTableState { .. }
1209 | Op::Flush { .. }
1210 | Op::Ddl(_) => {}
1211 }
1212 }
1213
1214 let rcache_dir = dir.join(RCACHE_DIR);
1215 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1216 let mut db = Self {
1217 dir,
1218 table_id: manifest.table_id,
1219 wal,
1220 memtable,
1221 mutable_run: MutableRun::new(),
1222 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1223 compaction_zstd_level: 3,
1224 allocator,
1225 epoch: ctx.epoch,
1226 persisted_epoch,
1227 schema,
1228 hot: HotIndex::new(),
1229 kek: ctx.kek,
1230 column_keys,
1231 run_refs: manifest.runs.clone(),
1232 retiring: manifest.retiring.clone(),
1233 next_run_id: manifest
1234 .runs
1235 .iter()
1236 .map(|r| r.run_id as u64 + 1)
1237 .max()
1238 .unwrap_or(1),
1239 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1240 current_txn_id,
1241 bitmap: HashMap::new(),
1242 ann: HashMap::new(),
1243 fm: HashMap::new(),
1244 sparse: HashMap::new(),
1245 minhash: HashMap::new(),
1246 learned_range: HashMap::new(),
1247 pk_by_row: HashMap::new(),
1248 pinned: BTreeMap::new(),
1249 live_count: manifest.live_count,
1250 reservoir: crate::reservoir::Reservoir::default(),
1251 reservoir_complete: false,
1252 had_deletes: saw_delete,
1253 agg_cache: HashMap::new(),
1254 global_idx_epoch: manifest.global_idx_epoch,
1255 indexes_complete: true,
1256 index_build_policy: IndexBuildPolicy::default(),
1257 pk_by_row_complete: false,
1258 flushed_epoch: manifest.flushed_epoch,
1259 page_cache: ctx.page_cache,
1260 decoded_cache: ctx.decoded_cache,
1261 verified_runs: Arc::new(parking_lot::Mutex::new(std::collections::HashSet::new())),
1262 snapshots: ctx.snapshots,
1263 commit_lock: ctx.commit_lock,
1264 result_cache: Arc::new(parking_lot::Mutex::new(
1265 ResultCache::new()
1266 .with_dir(rcache_dir)
1267 .with_cache_dek(cache_dek.clone()),
1268 )),
1269 pending_delete_rids: roaring::RoaringBitmap::new(),
1270 pending_put_cols: std::collections::HashSet::new(),
1271 pending_rows: Vec::new(),
1272 pending_rows_auto_inc: Vec::new(),
1273 pending_dels: Vec::new(),
1274 pending_truncate: None,
1275 wal_dek,
1276 auto_inc,
1277 };
1278
1279 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1282
1283 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1288 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1289 c.epoch_built == manifest.global_idx_epoch
1290 && manifest.global_idx_epoch > 0
1291 && manifest
1292 .runs
1293 .iter()
1294 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1295 });
1296 if let Some(loaded) = checkpoint {
1297 if checkpoint_valid {
1298 db.hot = loaded.hot;
1299 db.bitmap = loaded.bitmap;
1300 db.ann = loaded.ann;
1301 db.fm = loaded.fm;
1302 db.sparse = loaded.sparse;
1303 db.minhash = loaded.minhash;
1304 db.learned_range = loaded.learned_range;
1305 }
1308 }
1309 if !checkpoint_valid {
1310 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1311 db.bitmap = bitmap;
1312 db.ann = ann;
1313 db.fm = fm;
1314 db.sparse = sparse;
1315 db.minhash = minhash;
1316 db.rebuild_indexes_from_runs()?;
1317 db.build_learned_ranges()?;
1318 }
1319
1320 for (epoch, group) in replayed_puts {
1325 let (losers, winner_pks) = db.partition_pk_winners(&group);
1326 for (key, &row_id) in &winner_pks {
1327 if let Some(old_rid) = db.hot.get(key) {
1328 if old_rid != row_id {
1329 db.tombstone_row(old_rid, epoch, false);
1330 }
1331 }
1332 }
1333 for &loser_rid in &losers {
1334 db.tombstone_row(loser_rid, epoch, false);
1335 }
1336 for (key, row_id) in winner_pks {
1337 db.insert_hot_pk(key, row_id);
1338 }
1339 if db.schema.primary_key().is_none() {
1340 for r in &group {
1341 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1342 }
1343 }
1344 for r in &group {
1345 if !losers.contains(&r.row_id) {
1346 db.index_row(r);
1347 }
1348 }
1349 }
1350 for (rid, epoch) in &replayed_deletes {
1354 db.remove_hot_for_row(*rid, *epoch);
1355 }
1356
1357 db.result_cache.lock().load_persistent();
1364 Ok(db)
1365 }
1366
1367 fn ensure_reservoir_complete(&mut self) -> Result<()> {
1373 if self.reservoir_complete {
1374 return Ok(());
1375 }
1376 self.rebuild_reservoir()?;
1377 self.reservoir_complete = true;
1378 Ok(())
1379 }
1380
1381 fn rebuild_reservoir(&mut self) -> Result<()> {
1384 let snap = self.snapshot();
1385 let rows = self.visible_rows(snap)?;
1386 self.reservoir.reset();
1387 for r in rows {
1388 self.reservoir.offer(r.row_id.0);
1389 }
1390 Ok(())
1391 }
1392
1393 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1394 self.hot = HotIndex::new();
1395 self.pk_by_row.clear();
1396 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1397 self.bitmap = bitmap;
1398 self.ann = ann;
1399 self.fm = fm;
1400 self.sparse = sparse;
1401 self.minhash = minhash;
1402 let snapshot = Epoch(u64::MAX);
1403 for rr in self.run_refs.clone() {
1404 let mut reader = self.open_reader(rr.run_id)?;
1405 for row in reader.visible_rows(snapshot)? {
1406 let tok_row = self.tokenized_for_indexes(&row);
1407 index_into(
1408 &self.schema,
1409 &tok_row,
1410 &mut self.hot,
1411 &mut self.bitmap,
1412 &mut self.ann,
1413 &mut self.fm,
1414 &mut self.sparse,
1415 &mut self.minhash,
1416 );
1417 }
1418 }
1419 for row in self.mutable_run.visible_versions(snapshot) {
1420 if row.deleted {
1421 self.remove_hot_for_row(row.row_id, snapshot);
1422 } else {
1423 self.index_row(&row);
1424 }
1425 }
1426 for row in self.memtable.visible_versions(snapshot) {
1427 if row.deleted {
1428 self.remove_hot_for_row(row.row_id, snapshot);
1429 } else {
1430 self.index_row(&row);
1431 }
1432 }
1433 self.refresh_pk_by_row_from_hot();
1434 Ok(())
1435 }
1436
1437 fn refresh_pk_by_row_from_hot(&mut self) {
1438 self.pk_by_row_complete = true;
1439 if self.schema.primary_key().is_none() {
1440 self.pk_by_row.clear();
1441 return;
1442 }
1443 self.pk_by_row = self
1449 .hot
1450 .entries()
1451 .into_iter()
1452 .map(|(key, row_id)| (row_id, key))
1453 .collect();
1454 }
1455
1456 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1457 if self.schema.primary_key().is_some() {
1458 self.pk_by_row.insert(row_id, key.clone());
1459 }
1460 self.hot.insert(key, row_id);
1461 }
1462
1463 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1467 self.learned_range.clear();
1468 if self.run_refs.len() != 1 {
1469 return Ok(());
1470 }
1471 let cols: Vec<u16> = self
1472 .schema
1473 .indexes
1474 .iter()
1475 .filter(|i| i.kind == IndexKind::LearnedRange)
1476 .map(|i| i.column_id)
1477 .collect();
1478 if cols.is_empty() {
1479 return Ok(());
1480 }
1481 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1482 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1483 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1484 _ => return Ok(()),
1485 };
1486 for cid in cols {
1487 let ty = self
1488 .schema
1489 .columns
1490 .iter()
1491 .find(|c| c.id == cid)
1492 .map(|c| c.ty)
1493 .unwrap_or(TypeId::Int64);
1494 match ty {
1495 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1496 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1497 let pairs: Vec<(i64, u64)> = data
1498 .iter()
1499 .zip(row_ids.iter())
1500 .map(|(v, r)| (*v, *r))
1501 .collect();
1502 self.learned_range
1503 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1504 }
1505 }
1506 TypeId::Float64 => {
1507 if let columnar::NativeColumn::Float64 { data, .. } =
1508 reader.column_native(cid)?
1509 {
1510 let pairs: Vec<(f64, u64)> = data
1511 .iter()
1512 .zip(row_ids.iter())
1513 .map(|(v, r)| (*v, *r))
1514 .collect();
1515 self.learned_range
1516 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1517 }
1518 }
1519 _ => {}
1520 }
1521 }
1522 Ok(())
1523 }
1524
1525 pub fn ensure_indexes_complete(&mut self) -> Result<()> {
1532 if self.indexes_complete {
1533 crate::trace::QueryTrace::record(|t| {
1534 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1535 });
1536 return Ok(());
1537 }
1538 crate::trace::QueryTrace::record(|t| {
1539 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1540 });
1541 self.rebuild_indexes_from_runs()?;
1542 self.build_learned_ranges()?;
1543 self.indexes_complete = true;
1544 let epoch = self.current_epoch();
1545 self.checkpoint_indexes(epoch);
1546 Ok(())
1547 }
1548
1549 fn pending_epoch(&self) -> Epoch {
1550 Epoch(self.epoch.visible().0 + 1)
1551 }
1552
1553 fn is_shared(&self) -> bool {
1556 matches!(self.wal, WalSink::Shared(_))
1557 }
1558
1559 fn ensure_txn_id(&mut self) -> u64 {
1563 if self.current_txn_id == 0 {
1564 let id = match &self.wal {
1565 WalSink::Shared(s) => {
1566 let mut g = s.txn_ids.lock();
1567 let v = *g;
1568 *g = g.wrapping_add(1);
1569 v
1570 }
1571 WalSink::Private(_) => 1,
1572 };
1573 self.current_txn_id = id;
1574 }
1575 self.current_txn_id
1576 }
1577
1578 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1581 let txn_id = self.ensure_txn_id();
1582 let table_id = self.table_id;
1583 match &mut self.wal {
1584 WalSink::Private(w) => {
1585 w.append_txn(txn_id, op)?;
1586 }
1587 WalSink::Shared(s) => {
1588 s.wal.lock().append(txn_id, table_id, op)?;
1589 }
1590 }
1591 Ok(())
1592 }
1593
1594 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1602 Ok(self.put_returning(columns)?.0)
1603 }
1604
1605 pub fn put_returning(
1610 &mut self,
1611 mut columns: Vec<(u16, Value)>,
1612 ) -> Result<(RowId, Option<i64>)> {
1613 let assigned = self.fill_auto_inc(&mut columns)?;
1614 self.schema.validate_not_null(&columns)?;
1615 let row_id = if self.schema.clustered {
1620 self.derive_clustered_row_id(&columns)?
1621 } else {
1622 self.allocator.alloc()
1623 };
1624 let epoch = self.pending_epoch();
1625 let mut row = Row::new(row_id, epoch);
1626 for (col_id, val) in columns {
1627 row.columns.insert(col_id, val);
1628 }
1629 self.commit_rows(vec![row], assigned.is_some())?;
1630 Ok((row_id, assigned))
1631 }
1632
1633 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1636 Ok(self
1637 .put_batch_returning(batch)?
1638 .into_iter()
1639 .map(|(r, _)| r)
1640 .collect())
1641 }
1642
1643 pub fn put_batch_returning(
1646 &mut self,
1647 batch: Vec<Vec<(u16, Value)>>,
1648 ) -> Result<Vec<(RowId, Option<i64>)>> {
1649 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1650 for mut cols in batch {
1651 let assigned = self.fill_auto_inc(&mut cols)?;
1652 filled.push((cols, assigned));
1653 }
1654 for (cols, _) in &filled {
1655 self.schema.validate_not_null(cols)?;
1656 }
1657 let epoch = self.pending_epoch();
1658 let mut rows = Vec::with_capacity(filled.len());
1659 let mut ids = Vec::with_capacity(filled.len());
1660 for (cols, assigned) in filled {
1661 let row_id = if self.schema.clustered {
1662 self.derive_clustered_row_id(&cols)?
1663 } else {
1664 self.allocator.alloc()
1665 };
1666 let mut row = Row::new(row_id, epoch);
1667 for (c, v) in cols {
1668 row.columns.insert(c, v);
1669 }
1670 ids.push((row_id, assigned));
1671 rows.push(row);
1672 }
1673 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1674 self.commit_rows(rows, all_auto_generated)?;
1675 Ok(ids)
1676 }
1677
1678 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1684 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1685 return Ok(None);
1686 };
1687 let pos = columns.iter().position(|(c, _)| *c == cid);
1688 let assigned = match pos {
1689 Some(i) => match &columns[i].1 {
1690 Value::Null => {
1691 let next = self.alloc_auto_inc_value()?;
1692 columns[i].1 = Value::Int64(next);
1693 Some(next)
1694 }
1695 Value::Int64(n) => {
1696 self.advance_auto_inc_past(*n)?;
1697 None
1698 }
1699 other => {
1700 return Err(MongrelError::InvalidArgument(format!(
1701 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1702 other
1703 )))
1704 }
1705 },
1706 None => {
1707 let next = self.alloc_auto_inc_value()?;
1708 columns.push((cid, Value::Int64(next)));
1709 Some(next)
1710 }
1711 };
1712 Ok(assigned)
1713 }
1714
1715 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1717 self.ensure_auto_inc_seeded()?;
1718 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1720 let v = ai.next;
1721 ai.next = ai.next.saturating_add(1);
1722 Ok(v)
1723 }
1724
1725 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1728 self.ensure_auto_inc_seeded()?;
1729 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1730 let floor = used.saturating_add(1).max(1);
1731 if ai.next < floor {
1732 ai.next = floor;
1733 }
1734 Ok(())
1735 }
1736
1737 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1742 let needs_seed = match self.auto_inc {
1743 Some(ai) => !ai.seeded,
1744 None => return Ok(()),
1745 };
1746 if !needs_seed {
1747 return Ok(());
1748 }
1749 if self.seed_empty_auto_inc() {
1750 return Ok(());
1751 }
1752 let cid = self
1753 .auto_inc
1754 .as_ref()
1755 .expect("auto-inc column present")
1756 .column_id;
1757 let max = self.scan_max_int64(cid)?;
1758 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1759 let floor = max.saturating_add(1).max(1);
1760 if ai.next < floor {
1761 ai.next = floor;
1762 }
1763 ai.seeded = true;
1764 Ok(())
1765 }
1766
1767 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1768 if n == 0 || self.auto_inc.is_none() {
1769 return Ok(None);
1770 }
1771 self.ensure_auto_inc_seeded()?;
1772 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1773 let start = ai.next;
1774 ai.next = ai.next.saturating_add(n as i64);
1775 Ok(Some(start))
1776 }
1777
1778 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1782 let mut max: i64 = 0;
1783 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1784 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1785 if *n > max {
1786 max = *n;
1787 }
1788 }
1789 }
1790 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1791 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1792 if *n > max {
1793 max = *n;
1794 }
1795 }
1796 }
1797 for rr in self.run_refs.clone() {
1798 let reader = self.open_reader(rr.run_id)?;
1799 if let Some(stats) = reader.column_page_stats(column_id) {
1800 for s in stats {
1801 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1802 if n > max {
1803 max = n;
1804 }
1805 }
1806 }
1807 } else if reader.has_column(column_id) {
1808 if let columnar::NativeColumn::Int64 { data, validity } =
1809 reader.column_native_shared(column_id)?
1810 {
1811 for (i, n) in data.iter().enumerate() {
1812 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1813 {
1814 max = *n;
1815 }
1816 }
1817 }
1818 }
1819 }
1820 Ok(max)
1821 }
1822
1823 fn seed_empty_auto_inc(&mut self) -> bool {
1824 let Some(ai) = self.auto_inc.as_mut() else {
1825 return false;
1826 };
1827 if ai.seeded || self.live_count != 0 {
1828 return false;
1829 }
1830 if ai.next < 1 {
1831 ai.next = 1;
1832 }
1833 ai.seeded = true;
1834 true
1835 }
1836
1837 fn advance_auto_inc_from_native_columns(
1838 &mut self,
1839 columns: &[(u16, columnar::NativeColumn)],
1840 n: usize,
1841 live_before: u64,
1842 ) -> Result<()> {
1843 let Some(ai) = self.auto_inc.as_mut() else {
1844 return Ok(());
1845 };
1846 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1847 return Ok(());
1848 };
1849 let columnar::NativeColumn::Int64 { data, validity } = col else {
1850 return Err(MongrelError::InvalidArgument(format!(
1851 "AUTO_INCREMENT column {} must be Int64",
1852 ai.column_id
1853 )));
1854 };
1855 let max = if native_int64_strictly_increasing(col, n) {
1856 data.get(n.saturating_sub(1)).copied()
1857 } else {
1858 data.iter()
1859 .take(n)
1860 .enumerate()
1861 .filter_map(|(i, v)| {
1862 if validity.is_empty() || columnar::validity_bit(validity, i) {
1863 Some(*v)
1864 } else {
1865 None
1866 }
1867 })
1868 .max()
1869 };
1870 if let Some(max) = max {
1871 let floor = max.saturating_add(1).max(1);
1872 if ai.next < floor {
1873 ai.next = floor;
1874 }
1875 if ai.seeded || live_before == 0 {
1876 ai.seeded = true;
1877 }
1878 }
1879 Ok(())
1880 }
1881
1882 fn fill_auto_inc_native_columns(
1883 &mut self,
1884 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1885 n: usize,
1886 ) -> Result<()> {
1887 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1888 return Ok(());
1889 };
1890 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1891 if let Some(start) = self.alloc_auto_inc_range(n)? {
1892 columns.push((
1893 cid,
1894 columnar::NativeColumn::Int64 {
1895 data: (start..start.saturating_add(n as i64)).collect(),
1896 validity: vec![0xFF; n.div_ceil(8)],
1897 },
1898 ));
1899 }
1900 return Ok(());
1901 };
1902
1903 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1904 return Err(MongrelError::InvalidArgument(format!(
1905 "AUTO_INCREMENT column {cid} must be Int64"
1906 )));
1907 };
1908 if data.len() < n {
1909 return Err(MongrelError::InvalidArgument(format!(
1910 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1911 data.len()
1912 )));
1913 }
1914 if columnar::all_non_null(validity, n) {
1915 return Ok(());
1916 }
1917 if validity.iter().all(|b| *b == 0) {
1918 if let Some(start) = self.alloc_auto_inc_range(n)? {
1919 for (i, slot) in data.iter_mut().take(n).enumerate() {
1920 *slot = start.saturating_add(i as i64);
1921 }
1922 *validity = vec![0xFF; n.div_ceil(8)];
1923 }
1924 return Ok(());
1925 }
1926
1927 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1928 for (i, slot) in data.iter_mut().enumerate().take(n) {
1929 if columnar::validity_bit(validity, i) {
1930 self.advance_auto_inc_past(*slot)?;
1931 } else {
1932 *slot = self.alloc_auto_inc_value()?;
1933 }
1934 }
1935 *validity = new_validity;
1936 Ok(())
1937 }
1938
1939 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1953 if self.auto_inc.is_none() {
1954 return Ok(None);
1955 }
1956 Ok(Some(self.alloc_auto_inc_value()?))
1957 }
1958
1959 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1965 let payload = bincode::serialize(&rows)?;
1966 self.wal_append_data(Op::Put {
1967 table_id: self.table_id,
1968 rows: payload,
1969 })?;
1970 if self.is_shared() {
1971 self.pending_rows_auto_inc
1972 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1973 self.pending_rows.extend(rows);
1974 } else {
1975 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1976 }
1977 Ok(())
1978 }
1979
1980 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1985 self.apply_put_rows_inner(rows, true)
1986 }
1987
1988 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1989 if check_existing_pk {
1990 self.ensure_indexes_complete()?;
1991 }
1992 if rows.len() == 1 {
1996 let row = rows.into_iter().next().expect("len checked");
1997 return self.apply_put_row_single(row, check_existing_pk);
1998 }
1999 let pk_id = self.schema.primary_key().map(|c| c.id);
2016 let probe = match pk_id {
2017 Some(pid) => {
2018 check_existing_pk
2019 && !(self.hot.is_empty() && rows_pk_strictly_increasing(&rows, pid))
2020 }
2021 None => false,
2022 };
2023 let maintain_pk_by_row = pk_id.is_some() && self.pk_by_row_complete;
2026 for r in rows {
2027 for &cid in r.columns.keys() {
2028 self.pending_put_cols.insert(cid);
2029 }
2030 match pk_id {
2031 Some(pid) if probe || maintain_pk_by_row => {
2032 if let Some(pk_val) = r.columns.get(&pid) {
2033 let key = self.index_lookup_key(pid, pk_val);
2034 if probe {
2035 if let Some(old_rid) = self.hot.get(&key) {
2036 if old_rid != r.row_id {
2037 self.tombstone_row(old_rid, r.committed_epoch, true);
2038 }
2039 }
2040 }
2041 if maintain_pk_by_row {
2042 self.pk_by_row.insert(r.row_id, key);
2043 }
2044 }
2045 }
2046 Some(_) => {}
2047 None => {
2048 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2049 }
2050 }
2051 self.index_row(&r);
2052 self.reservoir.offer(r.row_id.0);
2053 self.memtable.upsert(r);
2054 self.live_count = self.live_count.saturating_add(1);
2057 }
2058 Ok(())
2059 }
2060
2061 fn apply_put_row_single(&mut self, row: Row, check_existing_pk: bool) -> Result<()> {
2065 for &cid in row.columns.keys() {
2066 self.pending_put_cols.insert(cid);
2067 }
2068 let epoch = row.committed_epoch;
2069 if let Some(pk_col) = self.schema.primary_key() {
2070 let pk_id = pk_col.id;
2071 if let Some(pk_val) = row.columns.get(&pk_id) {
2072 let maintain_pk_by_row = self.pk_by_row_complete;
2076 if check_existing_pk || maintain_pk_by_row {
2077 let key = self.index_lookup_key(pk_id, pk_val);
2078 if check_existing_pk {
2079 if let Some(old_rid) = self.hot.get(&key) {
2080 if old_rid != row.row_id {
2081 self.tombstone_row(old_rid, epoch, true);
2082 }
2083 }
2084 }
2085 if maintain_pk_by_row {
2086 self.pk_by_row.insert(row.row_id, key);
2087 }
2088 }
2089 }
2090 } else {
2091 self.hot
2092 .insert(row.row_id.0.to_be_bytes().to_vec(), row.row_id);
2093 }
2094 self.index_row(&row);
2095 self.reservoir.offer(row.row_id.0);
2096 self.memtable.upsert(row);
2097 self.live_count = self.live_count.saturating_add(1);
2098 Ok(())
2099 }
2100
2101 pub(crate) fn alloc_row_id(&mut self) -> RowId {
2104 self.allocator.alloc()
2105 }
2106
2107 fn derive_clustered_row_id(&self, columns: &[(u16, Value)]) -> Result<RowId> {
2113 let pk = self.schema.primary_key().ok_or_else(|| {
2114 MongrelError::Schema("clustered table requires a single-column primary key".into())
2115 })?;
2116 let pk_val = columns
2117 .iter()
2118 .find(|(id, _)| *id == pk.id)
2119 .map(|(_, v)| v)
2120 .ok_or_else(|| {
2121 MongrelError::Schema(format!(
2122 "clustered table missing primary key column {} ({})",
2123 pk.id, pk.name
2124 ))
2125 })?;
2126 let key_bytes = pk_val.encode_key();
2127 let mut hash: u64 = 0xcbf29ce484222325;
2129 for &b in &key_bytes {
2130 hash ^= b as u64;
2131 hash = hash.wrapping_mul(0x100000001b3);
2132 }
2133 Ok(RowId(hash.max(1)))
2136 }
2137
2138 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
2146 self.ensure_indexes_complete()?;
2147 let n = rows.len();
2148 for r in rows {
2149 for &cid in r.columns.keys() {
2150 self.pending_put_cols.insert(cid);
2151 }
2152 }
2153 let (losers, winner_pks) = self.partition_pk_winners(rows);
2154 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
2155 for (key, &row_id) in &winner_pks {
2157 if let Some(old_rid) = self.hot.get(key) {
2158 if old_rid != row_id {
2159 self.tombstone_row(old_rid, epoch, true);
2160 }
2161 }
2162 }
2163 for &loser_rid in &losers {
2166 self.tombstone_row(loser_rid, epoch, false);
2167 }
2168 for (key, row_id) in winner_pks {
2170 self.insert_hot_pk(key, row_id);
2171 }
2172 if self.schema.primary_key().is_none() {
2173 for r in rows {
2174 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2175 }
2176 }
2177 for r in rows {
2178 self.allocator.advance_to(r.row_id);
2179 if !losers.contains(&r.row_id) {
2180 self.index_row(r);
2181 }
2182 }
2183 for r in rows {
2184 if !losers.contains(&r.row_id) {
2185 self.reservoir.offer(r.row_id.0);
2186 }
2187 }
2188 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
2189 Ok(())
2190 }
2191
2192 pub(crate) fn recover_apply(
2197 &mut self,
2198 rows: Vec<Row>,
2199 deletes: Vec<(RowId, Epoch)>,
2200 ) -> Result<()> {
2201 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
2205 std::collections::BTreeMap::new();
2206 for row in rows {
2207 self.allocator.advance_to(row.row_id);
2208 if let Some(ai) = self.auto_inc.as_mut() {
2213 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
2214 if *n + 1 > ai.next {
2215 ai.next = *n + 1;
2216 }
2217 }
2218 }
2219 by_epoch.entry(row.committed_epoch).or_default().push(row);
2220 }
2221 for (epoch, group) in by_epoch {
2222 let (losers, winner_pks) = self.partition_pk_winners(&group);
2223 for (key, &row_id) in &winner_pks {
2225 if let Some(old_rid) = self.hot.get(key) {
2226 if old_rid != row_id {
2227 self.tombstone_row(old_rid, epoch, false);
2228 }
2229 }
2230 }
2231 for (key, row_id) in winner_pks {
2232 self.insert_hot_pk(key, row_id);
2233 }
2234 if self.schema.primary_key().is_none() {
2235 for r in &group {
2236 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2237 }
2238 }
2239 for r in &group {
2240 if !losers.contains(&r.row_id) {
2241 self.memtable.upsert(r.clone());
2242 self.index_row(r);
2243 }
2244 }
2245 }
2246 for (rid, epoch) in deletes {
2247 self.memtable.tombstone(rid, epoch);
2248 self.remove_hot_for_row(rid, epoch);
2249 }
2250 self.reservoir_complete = false;
2253 Ok(())
2254 }
2255
2256 pub(crate) fn flushed_epoch(&self) -> u64 {
2258 self.flushed_epoch
2259 }
2260
2261 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2263 self.schema.validate_not_null(cells)
2264 }
2265
2266 fn validate_columns_not_null(
2270 &self,
2271 columns: &[(u16, columnar::NativeColumn)],
2272 n: usize,
2273 ) -> Result<()> {
2274 let by_id: HashMap<u16, &columnar::NativeColumn> =
2275 columns.iter().map(|(id, c)| (*id, c)).collect();
2276 for col in &self.schema.columns {
2277 if col.flags.contains(ColumnFlags::NULLABLE) {
2278 continue;
2279 }
2280 match by_id.get(&col.id) {
2281 None => {
2282 return Err(MongrelError::InvalidArgument(format!(
2283 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2284 col.name, col.id
2285 )));
2286 }
2287 Some(c) => {
2288 if c.null_count(n) != 0 {
2289 return Err(MongrelError::InvalidArgument(format!(
2290 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2291 col.name, col.id
2292 )));
2293 }
2294 }
2295 }
2296 }
2297 Ok(())
2298 }
2299
2300 fn bulk_pk_winner_indices(
2305 &self,
2306 columns: &[(u16, columnar::NativeColumn)],
2307 n: usize,
2308 ) -> Option<Vec<usize>> {
2309 let pk_col = self.schema.primary_key()?;
2310 let pk_id = pk_col.id;
2311 let pk_ty = pk_col.ty;
2312 let by_id: HashMap<u16, &columnar::NativeColumn> =
2313 columns.iter().map(|(id, c)| (*id, c)).collect();
2314 let pk_native = by_id.get(&pk_id)?;
2315 if native_int64_strictly_increasing(pk_native, n) {
2316 return None;
2317 }
2318 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2320 let mut null_pk_rows: Vec<usize> = Vec::new();
2321 for i in 0..n {
2322 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2323 Some(key) => {
2324 last.insert(key, i);
2325 }
2326 None => null_pk_rows.push(i),
2327 }
2328 }
2329 let mut winners: HashSet<usize> = last.values().copied().collect();
2330 for i in null_pk_rows {
2331 winners.insert(i);
2332 }
2333 Some((0..n).filter(|i| winners.contains(i)).collect())
2334 }
2335
2336 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2338 let epoch = self.pending_epoch();
2339 self.wal_append_data(Op::Delete {
2340 table_id: self.table_id,
2341 row_ids: vec![row_id],
2342 })?;
2343 if self.is_shared() {
2344 self.pending_dels.push(row_id);
2345 } else {
2346 self.apply_delete(row_id, epoch);
2347 }
2348 Ok(())
2349 }
2350
2351 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2352 let pre = self.get(row_id, self.snapshot());
2353 self.delete(row_id)?;
2354 Ok(pre.map(|row| {
2355 let mut columns: Vec<_> = row.columns.into_iter().collect();
2356 columns.sort_by_key(|(id, _)| *id);
2357 OwnedRow { columns }
2358 }))
2359 }
2360
2361 pub fn truncate(&mut self) -> Result<()> {
2363 let epoch = self.pending_epoch();
2364 self.wal_append_data(Op::TruncateTable {
2365 table_id: self.table_id,
2366 })?;
2367 self.pending_rows.clear();
2368 self.pending_rows_auto_inc.clear();
2369 self.pending_dels.clear();
2370 self.pending_truncate = Some(epoch);
2371 Ok(())
2372 }
2373
2374 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2376 for rr in std::mem::take(&mut self.run_refs) {
2377 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2378 }
2379 for r in std::mem::take(&mut self.retiring) {
2380 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2381 }
2382 self.memtable = Memtable::new();
2383 self.mutable_run = MutableRun::new();
2384 self.hot = HotIndex::new();
2385 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2386 self.bitmap = bitmap;
2387 self.ann = ann;
2388 self.fm = fm;
2389 self.sparse = sparse;
2390 self.minhash = minhash;
2391 self.learned_range.clear();
2392 self.pk_by_row.clear();
2393 self.pk_by_row_complete = false;
2394 self.live_count = 0;
2395 self.reservoir = crate::reservoir::Reservoir::default();
2396 self.reservoir_complete = true;
2397 self.had_deletes = true;
2398 self.agg_cache.clear();
2399 self.global_idx_epoch = 0;
2400 self.indexes_complete = true;
2401 self.pending_delete_rids.clear();
2402 self.pending_put_cols.clear();
2403 self.pending_rows.clear();
2404 self.pending_rows_auto_inc.clear();
2405 self.pending_dels.clear();
2406 self.clear_result_cache();
2407 self.invalidate_index_checkpoint();
2408 Ok(())
2409 }
2410
2411 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2414 self.remove_hot_for_row(row_id, epoch);
2415 self.tombstone_row(row_id, epoch, true);
2416 }
2417
2418 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2422 let tombstone = Row {
2423 row_id,
2424 committed_epoch: epoch,
2425 columns: std::collections::HashMap::new(),
2426 deleted: true,
2427 };
2428 self.memtable.upsert(tombstone);
2429 self.pk_by_row.remove(&row_id);
2430 if adjust_live_count {
2431 self.live_count = self.live_count.saturating_sub(1);
2432 }
2433 self.pending_delete_rids.insert(row_id.0 as u32);
2435 self.had_deletes = true;
2438 self.agg_cache.clear();
2439 }
2440
2441 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2445 let Some(pk_col) = self.schema.primary_key() else {
2446 return;
2447 };
2448 if self.pk_by_row_complete {
2451 if let Some(key) = self.pk_by_row.remove(&row_id) {
2452 if self.hot.get(&key) == Some(row_id) {
2453 self.hot.remove(&key);
2454 }
2455 }
2456 return;
2457 }
2458 let lookup_epoch = Epoch(epoch.0.saturating_sub(1));
2477 if self.indexes_complete {
2478 let pk_val = self
2479 .memtable
2480 .get_version(row_id, lookup_epoch)
2481 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2482 .or_else(|| {
2483 self.mutable_run
2484 .get_version(row_id, lookup_epoch)
2485 .filter(|(_, r)| !r.deleted)
2486 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2487 })
2488 .or_else(|| {
2489 self.run_refs.iter().find_map(|rr| {
2490 let mut reader = self.open_reader(rr.run_id).ok()?;
2491 let (_, deleted, val) = reader
2492 .get_version_column(row_id, lookup_epoch, pk_col.id)
2493 .ok()??;
2494 if deleted {
2495 return None;
2496 }
2497 val
2498 })
2499 });
2500 if let Some(pk_val) = pk_val {
2501 let key = self.index_lookup_key(pk_col.id, &pk_val);
2502 if self.hot.get(&key) == Some(row_id) {
2503 self.hot.remove(&key);
2504 }
2505 return;
2506 }
2507 }
2508 self.refresh_pk_by_row_from_hot();
2513 if let Some(key) = self.pk_by_row.remove(&row_id) {
2514 if self.hot.get(&key) == Some(row_id) {
2515 self.hot.remove(&key);
2516 }
2517 }
2518 }
2519
2520 fn partition_pk_winners(
2525 &self,
2526 rows: &[Row],
2527 ) -> (
2528 std::collections::HashSet<RowId>,
2529 std::collections::HashMap<Vec<u8>, RowId>,
2530 ) {
2531 let mut losers = std::collections::HashSet::new();
2532 let Some(pk_col) = self.schema.primary_key() else {
2533 return (losers, std::collections::HashMap::new());
2534 };
2535 let pk_id = pk_col.id;
2536 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2537 std::collections::HashMap::new();
2538 for r in rows {
2539 let Some(pk_val) = r.columns.get(&pk_id) else {
2540 continue;
2541 };
2542 let key = self.index_lookup_key(pk_id, pk_val);
2543 if let Some(&old_rid) = winners.get(&key) {
2544 losers.insert(old_rid);
2545 }
2546 winners.insert(key, r.row_id);
2547 }
2548 (losers, winners)
2549 }
2550
2551 fn index_row(&mut self, row: &Row) {
2552 if row.deleted {
2553 return;
2554 }
2555 let any_predicate = self
2563 .schema
2564 .indexes
2565 .iter()
2566 .any(|idx| idx.predicate.is_some());
2567 if any_predicate {
2568 let columns_map: HashMap<u16, &Value> =
2569 row.columns.iter().map(|(k, v)| (*k, v)).collect();
2570 let name_to_id: HashMap<&str, u16> = self
2571 .schema
2572 .columns
2573 .iter()
2574 .map(|c| (c.name.as_str(), c.id))
2575 .collect();
2576 for idx in &self.schema.indexes {
2577 if let Some(pred) = &idx.predicate {
2578 if !eval_partial_predicate(pred, &columns_map, &name_to_id) {
2579 continue; }
2581 }
2582 index_into_single(
2584 idx,
2585 &self.schema,
2586 row,
2587 &mut self.hot,
2588 &mut self.bitmap,
2589 &mut self.ann,
2590 &mut self.fm,
2591 &mut self.sparse,
2592 &mut self.minhash,
2593 );
2594 }
2595 return;
2596 }
2597 if self.column_keys.is_empty() {
2601 index_into(
2602 &self.schema,
2603 row,
2604 &mut self.hot,
2605 &mut self.bitmap,
2606 &mut self.ann,
2607 &mut self.fm,
2608 &mut self.sparse,
2609 &mut self.minhash,
2610 );
2611 return;
2612 }
2613 let effective_row = self.tokenized_for_indexes(row);
2614 index_into(
2615 &self.schema,
2616 &effective_row,
2617 &mut self.hot,
2618 &mut self.bitmap,
2619 &mut self.ann,
2620 &mut self.fm,
2621 &mut self.sparse,
2622 &mut self.minhash,
2623 );
2624 }
2625
2626 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2632 if self.column_keys.is_empty() {
2633 return row.clone();
2634 }
2635 #[cfg(feature = "encryption")]
2636 {
2637 use crate::encryption::SCHEME_HMAC_EQ;
2638 let mut tok = row.clone();
2639 for (&cid, &(_, scheme)) in &self.column_keys {
2640 if scheme != SCHEME_HMAC_EQ {
2641 continue;
2642 }
2643 if let Some(v) = tok.columns.get(&cid).cloned() {
2644 if let Some(t) = self.tokenize_value(cid, &v) {
2645 tok.columns.insert(cid, t);
2646 }
2647 }
2648 }
2649 tok
2650 }
2651 #[cfg(not(feature = "encryption"))]
2652 {
2653 row.clone()
2654 }
2655 }
2656
2657 pub fn commit(&mut self) -> Result<Epoch> {
2662 if self.is_shared() {
2663 self.commit_shared()
2664 } else {
2665 self.commit_private()
2666 }
2667 }
2668
2669 fn commit_private(&mut self) -> Result<Epoch> {
2671 let commit_lock = Arc::clone(&self.commit_lock);
2675 let _g = commit_lock.lock();
2676 let new_epoch = self.epoch.bump_assigned();
2677 let txn_id = self.current_txn_id;
2678 match &mut self.wal {
2682 WalSink::Private(w) => {
2683 w.append_txn(
2684 txn_id,
2685 Op::TxnCommit {
2686 epoch: new_epoch.0,
2687 added_runs: Vec::new(),
2688 },
2689 )?;
2690 w.sync()?;
2691 }
2692 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2693 }
2694 if let Some(epoch) = self.pending_truncate.take() {
2696 self.apply_truncate(epoch)?;
2697 }
2698 self.invalidate_pending_cache();
2699 self.persist_manifest(new_epoch)?;
2700 self.epoch.publish_in_order(new_epoch);
2704 self.current_txn_id += 1;
2705 Ok(new_epoch)
2706 }
2707
2708 fn commit_shared(&mut self) -> Result<Epoch> {
2714 use std::sync::atomic::Ordering;
2715 let s = match &self.wal {
2716 WalSink::Shared(s) => s.clone(),
2717 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2718 };
2719 if s.poisoned.load(Ordering::Relaxed) {
2720 return Err(MongrelError::Other(
2721 "database poisoned by fsync error".into(),
2722 ));
2723 }
2724 let commit_lock = Arc::clone(&self.commit_lock);
2731 let _g = commit_lock.lock();
2732 let txn_id = self.ensure_txn_id();
2735 let (new_epoch, commit_seq) = {
2736 let mut wal = s.wal.lock();
2737 let new_epoch = self.epoch.bump_assigned();
2738 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2739 (new_epoch, seq)
2740 };
2741 s.group
2742 .await_durable(&s.wal, commit_seq)
2743 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2744
2745 if self.pending_truncate.take().is_some() {
2748 self.apply_truncate(new_epoch)?;
2749 }
2750 let mut rows = std::mem::take(&mut self.pending_rows);
2751 if !rows.is_empty() {
2752 for r in &mut rows {
2753 r.committed_epoch = new_epoch;
2754 }
2755 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2756 let all_auto_generated =
2757 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2758 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2759 } else {
2760 self.pending_rows_auto_inc.clear();
2761 }
2762 let dels = std::mem::take(&mut self.pending_dels);
2763 for rid in dels {
2764 self.apply_delete(rid, new_epoch);
2765 }
2766
2767 self.invalidate_pending_cache();
2768 self.persist_manifest(new_epoch)?;
2769 self.epoch.publish_in_order(new_epoch);
2770 self.current_txn_id = 0;
2772 Ok(new_epoch)
2773 }
2774
2775 pub fn flush(&mut self) -> Result<Epoch> {
2783 self.ensure_indexes_complete()?;
2784 let epoch = self.commit()?;
2785 let rows = self.memtable.drain_sorted();
2786 if !rows.is_empty() {
2787 self.mutable_run.insert_many(rows);
2788 }
2789 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2790 self.spill_mutable_run(epoch)?;
2791 self.mark_flushed(epoch)?;
2795 self.persist_manifest(epoch)?;
2796 self.build_learned_ranges()?;
2797 self.checkpoint_indexes(epoch);
2800 }
2801 Ok(epoch)
2804 }
2805
2806 pub fn force_flush(&mut self) -> Result<Epoch> {
2815 let saved = self.mutable_run_spill_bytes;
2816 self.mutable_run_spill_bytes = 1;
2817 let result = self.flush();
2818 self.mutable_run_spill_bytes = saved;
2819 result
2820 }
2821
2822 pub fn close(&mut self) -> Result<()> {
2829 if self.memtable_len() > 0 || self.mutable_run_len() > 0 {
2830 self.force_flush()?;
2831 }
2832 Ok(())
2833 }
2834
2835 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2842 let op = Op::Flush {
2843 table_id: self.table_id,
2844 flushed_epoch: epoch.0,
2845 };
2846 match &mut self.wal {
2847 WalSink::Private(w) => {
2848 w.append_system(op)?;
2849 w.sync()?;
2850 }
2851 WalSink::Shared(s) => {
2852 s.wal.lock().append_system(op)?;
2857 }
2858 }
2859 self.flushed_epoch = epoch.0;
2860 if matches!(self.wal, WalSink::Private(_)) {
2861 self.rotate_wal(epoch)?;
2862 }
2863 Ok(())
2864 }
2865
2866 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2870 let rows = self.mutable_run.drain_sorted();
2871 if rows.is_empty() {
2872 return Ok(());
2873 }
2874 let run_id = self.next_run_id;
2875 self.next_run_id += 1;
2876 let path = self.run_path(run_id);
2877 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2878 if let Some(kek) = &self.kek {
2879 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2880 }
2881 let header = writer.write(&path, &rows)?;
2882 self.run_refs.push(RunRef {
2883 run_id: run_id as u128,
2884 level: 0,
2885 epoch_created: epoch.0,
2886 row_count: header.row_count,
2887 });
2888 Ok(())
2889 }
2890
2891 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2895 self.mutable_run_spill_bytes = bytes.max(1);
2896 }
2897
2898 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2902 self.compaction_zstd_level = level;
2903 }
2904
2905 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2909 self.result_cache.lock().set_max_bytes(max_bytes);
2910 }
2911
2912 pub(crate) fn clear_result_cache(&mut self) {
2916 self.result_cache.lock().clear();
2917 }
2918
2919 pub fn mutable_run_len(&self) -> usize {
2921 self.mutable_run.len()
2922 }
2923
2924 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2927 self.mutable_run.drain_sorted()
2928 }
2929
2930 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2935 let epoch = self.commit()?;
2936 let n = batch.len();
2937 if n == 0 {
2938 return Ok(epoch);
2939 }
2940 let live_before = self.live_count;
2941 self.spill_mutable_run(epoch)?;
2945 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
2946 && self.indexes_complete
2947 && self.run_refs.is_empty()
2948 && self.memtable.is_empty()
2949 && self.mutable_run.is_empty();
2950 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2956 use rayon::prelude::*;
2957 self.schema
2958 .columns
2959 .par_iter()
2960 .map(|cdef| (cdef.id, columnar::rows_to_native(cdef.ty, &batch, cdef.id)))
2961 .collect::<Vec<_>>()
2962 };
2963 drop(batch);
2964 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2969 self.validate_columns_not_null(&user_columns, n)?;
2970 let winner_idx = self
2971 .bulk_pk_winner_indices(&user_columns, n)
2972 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2973 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2974 match winner_idx.as_deref() {
2975 Some(idx) => {
2976 let compacted = user_columns
2977 .iter()
2978 .map(|(id, c)| (*id, c.gather(idx)))
2979 .collect();
2980 (compacted, idx.len())
2981 }
2982 None => (std::mem::take(&mut user_columns), n),
2983 };
2984 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2985 let first = self.allocator.alloc_range(write_n as u64).0;
2986 for rid in first..first + write_n as u64 {
2987 self.reservoir.offer(rid);
2988 }
2989 let run_id = self.next_run_id;
2990 self.next_run_id += 1;
2991 let path = self.run_path(run_id);
2992 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2993 .clean(true)
2994 .with_lz4()
2995 .with_native_endian();
2996 if let Some(kek) = &self.kek {
2997 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2998 }
2999 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3000 self.run_refs.push(RunRef {
3001 run_id: run_id as u128,
3002 level: 0,
3003 epoch_created: epoch.0,
3004 row_count: header.row_count,
3005 });
3006 self.live_count = self.live_count.saturating_add(write_n as u64);
3007 if eager_index_build {
3008 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3009 self.index_columns_bulk(&write_columns, &row_ids);
3010 self.indexes_complete = true;
3011 self.build_learned_ranges()?;
3012 } else {
3013 self.indexes_complete = false;
3014 }
3015 self.mark_flushed(epoch)?;
3016 self.persist_manifest(epoch)?;
3017 if eager_index_build {
3018 self.checkpoint_indexes(epoch);
3019 }
3020 self.clear_result_cache();
3021 Ok(epoch)
3022 }
3023
3024 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
3027 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
3028 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
3029 let segment_no = segment
3032 .file_stem()
3033 .and_then(|s| s.to_str())
3034 .and_then(|s| s.strip_prefix("seg-"))
3035 .and_then(|s| s.parse::<u64>().ok())
3036 .unwrap_or(0);
3037 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
3038 wal.set_sync_byte_threshold(self.sync_byte_threshold);
3039 wal.sync()?;
3040 self.wal = WalSink::Private(wal);
3041 Ok(())
3042 }
3043
3044 pub(crate) fn invalidate_pending_cache(&mut self) {
3049 self.result_cache
3050 .lock()
3051 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
3052 self.pending_delete_rids.clear();
3053 self.pending_put_cols.clear();
3054 }
3055
3056 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
3057 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
3058 m.current_epoch = epoch.0;
3059 m.next_row_id = self.allocator.current().0;
3060 m.runs = self.run_refs.clone();
3061 m.live_count = self.live_count;
3062 m.global_idx_epoch = self.global_idx_epoch;
3063 m.flushed_epoch = self.flushed_epoch;
3064 m.retiring = self.retiring.clone();
3065 m.auto_inc_next = match self.auto_inc {
3069 Some(ai) if ai.seeded => ai.next,
3070 _ => 0,
3071 };
3072 let meta_dek = self.manifest_meta_dek();
3073 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
3074 Ok(())
3075 }
3076
3077 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
3083 if !self.indexes_complete {
3086 return;
3087 }
3088 let snap = global_idx::IndexSnapshot {
3089 hot: &self.hot,
3090 bitmap: &self.bitmap,
3091 ann: &self.ann,
3092 fm: &self.fm,
3093 sparse: &self.sparse,
3094 minhash: &self.minhash,
3095 learned_range: &self.learned_range,
3096 };
3097 let idx_dek = self.idx_dek();
3099 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
3100 .is_ok()
3101 {
3102 self.global_idx_epoch = epoch.0;
3103 let _ = self.persist_manifest(epoch);
3104 }
3105 }
3106
3107 pub(crate) fn invalidate_index_checkpoint(&mut self) {
3110 self.global_idx_epoch = 0;
3111 global_idx::remove(&self.dir);
3112 let _ = self.persist_manifest(self.epoch.visible());
3113 }
3114
3115 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
3118 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
3119 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
3120 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3121 best = Some((epoch, row));
3122 }
3123 }
3124 for rr in &self.run_refs {
3125 let Ok(mut reader) = self.open_reader(rr.run_id) else {
3126 continue;
3127 };
3128 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
3129 continue;
3130 };
3131 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3132 best = Some((epoch, row));
3133 }
3134 }
3135 match best {
3136 Some((_, r)) if r.deleted => None,
3137 Some((_, r)) => Some(r),
3138 None => None,
3139 }
3140 }
3141
3142 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
3146 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
3147 let mut fold = |row: Row| {
3148 best.entry(row.row_id.0)
3149 .and_modify(|e| {
3150 if row.committed_epoch > e.0 {
3151 *e = (row.committed_epoch, row.clone());
3152 }
3153 })
3154 .or_insert_with(|| (row.committed_epoch, row));
3155 };
3156 for row in self.memtable.visible_versions(snapshot.epoch) {
3157 fold(row);
3158 }
3159 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3160 fold(row);
3161 }
3162 for rr in &self.run_refs {
3163 let mut reader = self.open_reader(rr.run_id)?;
3164 for row in reader.visible_versions(snapshot.epoch)? {
3165 fold(row);
3166 }
3167 }
3168 let mut out: Vec<Row> = best
3169 .into_values()
3170 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
3171 .collect();
3172 out.sort_by_key(|r| r.row_id);
3173 Ok(out)
3174 }
3175
3176 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
3183 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3184 let rr = self.run_refs[0].clone();
3185 let mut reader = self.open_reader(rr.run_id)?;
3186 let idxs = reader.visible_indices(snapshot.epoch)?;
3187 let mut cols = Vec::with_capacity(self.schema.columns.len());
3188 for cdef in &self.schema.columns {
3189 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
3190 }
3191 return Ok(cols);
3192 }
3193 let rows = self.visible_rows(snapshot)?;
3195 let mut cols: Vec<(u16, Vec<Value>)> = self
3196 .schema
3197 .columns
3198 .iter()
3199 .map(|c| (c.id, Vec::with_capacity(rows.len())))
3200 .collect();
3201 for r in &rows {
3202 for (cid, vec) in cols.iter_mut() {
3203 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
3204 }
3205 }
3206 Ok(cols)
3207 }
3208
3209 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
3211 self.hot.get(key)
3212 }
3213
3214 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
3219 self.ensure_indexes_complete()?;
3220 let snapshot = self.snapshot();
3221 crate::trace::QueryTrace::record(|t| {
3222 t.run_count = self.run_refs.len();
3223 t.memtable_rows = self.memtable.len();
3224 t.mutable_run_rows = self.mutable_run.len();
3225 });
3226 if q.conditions.is_empty() {
3230 crate::trace::QueryTrace::record(|t| {
3231 t.scan_mode = crate::trace::ScanMode::Materialized;
3232 t.row_materialized = true;
3233 });
3234 return self.visible_rows(snapshot);
3235 }
3236 crate::trace::QueryTrace::record(|t| {
3237 t.conditions_pushed = q.conditions.len();
3238 t.scan_mode = crate::trace::ScanMode::Materialized;
3239 t.row_materialized = true;
3240 });
3241 let mut ordered: Vec<&crate::query::Condition> = q.conditions.iter().collect();
3248 ordered.sort_by_key(|c| condition_cost_rank(c));
3249 let mut sets: Vec<RowIdSet> = Vec::with_capacity(ordered.len());
3250 for c in &ordered {
3251 let s = self.resolve_condition(c, snapshot)?;
3252 let empty = s.is_empty();
3253 sets.push(s);
3254 if empty {
3255 break;
3256 }
3257 }
3258 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
3259 self.rows_for_rids(&rids, snapshot)
3260 }
3261
3262 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
3267 use std::collections::HashMap;
3268 let mut rows = Vec::with_capacity(rids.len());
3269 let tier_size = self.memtable.len() + self.mutable_run.len();
3286 let mut overlay: HashMap<u64, Row> = HashMap::with_capacity(rids.len());
3287 if rids.len().saturating_mul(24) < tier_size {
3288 for &rid in rids {
3289 let mem = self.memtable.get_version(RowId(rid), snapshot.epoch);
3290 let mrun = self.mutable_run.get_version(RowId(rid), snapshot.epoch);
3291 let newest = match (mem, mrun) {
3292 (Some((me, mr)), Some((re, rr))) => Some(if me >= re { mr } else { rr }),
3293 (Some((_, mr)), None) => Some(mr),
3294 (None, Some((_, rr))) => Some(rr),
3295 (None, None) => None,
3296 };
3297 if let Some(row) = newest {
3298 overlay.insert(rid, row);
3299 }
3300 }
3301 } else {
3302 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
3303 overlay
3304 .entry(row.row_id.0)
3305 .and_modify(|e| {
3306 if row.committed_epoch > e.committed_epoch {
3307 *e = row.clone();
3308 }
3309 })
3310 .or_insert(row);
3311 };
3312 for row in self.memtable.visible_versions(snapshot.epoch) {
3313 fold_newest(row, &mut overlay);
3314 }
3315 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3316 fold_newest(row, &mut overlay);
3317 }
3318 }
3319 if self.run_refs.len() == 1 {
3320 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
3321 if rids.len().saturating_mul(24) < reader.row_count() {
3329 for &rid in rids {
3330 if let Some(r) = overlay.get(&rid) {
3331 if !r.deleted {
3332 rows.push(r.clone());
3333 }
3334 continue;
3335 }
3336 if let Some((_, row)) = reader.get_version(RowId(rid), snapshot.epoch)? {
3337 if !row.deleted {
3338 rows.push(row);
3339 }
3340 }
3341 }
3342 return Ok(rows);
3343 }
3344 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
3353 enum Src {
3356 Overlay,
3357 Run,
3358 }
3359 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
3360 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
3361 for rid in rids {
3362 if overlay.contains_key(rid) {
3363 plan.push(Src::Overlay);
3364 continue;
3365 }
3366 match vis_rids.binary_search(&(*rid as i64)) {
3367 Ok(i) => {
3368 plan.push(Src::Run);
3369 fetch.push(positions[i]);
3370 }
3371 Err(_) => { }
3372 }
3373 }
3374 let fetched = reader.materialize_batch(&fetch)?;
3375 let mut fetched_iter = fetched.into_iter();
3376 for (rid, src) in rids.iter().zip(plan) {
3377 match src {
3378 Src::Overlay => {
3379 if let Some(r) = overlay.get(rid) {
3380 if !r.deleted {
3381 rows.push(r.clone());
3382 }
3383 }
3384 }
3385 Src::Run => {
3386 if let Some(row) = fetched_iter.next() {
3387 if !row.deleted {
3388 rows.push(row);
3389 }
3390 }
3391 }
3392 }
3393 }
3394 return Ok(rows);
3395 }
3396 let mut readers: Vec<_> = self
3400 .run_refs
3401 .iter()
3402 .map(|rr| self.open_reader(rr.run_id))
3403 .collect::<Result<Vec<_>>>()?;
3404 for rid in rids {
3405 if let Some(r) = overlay.get(rid) {
3406 if !r.deleted {
3407 rows.push(r.clone());
3408 }
3409 continue;
3410 }
3411 let mut best: Option<(Epoch, Row)> = None;
3412 for reader in readers.iter_mut() {
3413 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3414 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3415 best = Some((epoch, row));
3416 }
3417 }
3418 }
3419 if let Some((_, r)) = best {
3420 if !r.deleted {
3421 rows.push(r);
3422 }
3423 }
3424 }
3425 Ok(rows)
3426 }
3427
3428 pub fn indexes_complete(&self) -> bool {
3438 self.indexes_complete
3439 }
3440
3441 pub fn index_build_policy(&self) -> IndexBuildPolicy {
3443 self.index_build_policy
3444 }
3445
3446 pub fn set_index_build_policy(&mut self, policy: IndexBuildPolicy) {
3450 self.index_build_policy = policy;
3451 }
3452
3453 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3458 if !self.indexes_complete {
3462 return None;
3463 }
3464 let b = self.bitmap.get(&column_id)?;
3465 let result: Vec<Vec<u8>> = b
3466 .keys()
3467 .into_iter()
3468 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3469 .cloned()
3470 .collect();
3471 Some(result)
3472 }
3473
3474 pub fn fk_join_row_ids(
3475 &self,
3476 fk_column_id: u16,
3477 pk_values: &[Vec<u8>],
3478 fk_conditions: &[crate::query::Condition],
3479 snapshot: Snapshot,
3480 ) -> Result<Vec<u64>> {
3481 let Some(b) = self.bitmap.get(&fk_column_id) else {
3482 return Ok(Vec::new());
3483 };
3484 let mut join_set = {
3485 let mut acc = roaring::RoaringBitmap::new();
3486 for v in pk_values {
3487 acc |= b.get(v);
3488 }
3489 RowIdSet::from_roaring(acc)
3490 };
3491 if !fk_conditions.is_empty() {
3492 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3493 sets.push(join_set);
3494 for c in fk_conditions {
3495 sets.push(self.resolve_condition(c, snapshot)?);
3496 }
3497 join_set = RowIdSet::intersect_many(sets);
3498 }
3499 Ok(join_set.into_sorted_vec())
3500 }
3501
3502 pub fn fk_join_count(
3508 &self,
3509 fk_column_id: u16,
3510 pk_values: &[Vec<u8>],
3511 fk_conditions: &[crate::query::Condition],
3512 snapshot: Snapshot,
3513 ) -> Result<u64> {
3514 let Some(b) = self.bitmap.get(&fk_column_id) else {
3515 return Ok(0);
3516 };
3517 let mut acc = roaring::RoaringBitmap::new();
3518 for v in pk_values {
3519 acc |= b.get(v);
3520 }
3521 if fk_conditions.is_empty() {
3522 return Ok(acc.len());
3523 }
3524 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3525 sets.push(RowIdSet::from_roaring(acc));
3526 for c in fk_conditions {
3527 sets.push(self.resolve_condition(c, snapshot)?);
3528 }
3529 Ok(RowIdSet::intersect_many(sets).len() as u64)
3530 }
3531
3532 fn resolve_condition(
3537 &self,
3538 c: &crate::query::Condition,
3539 snapshot: Snapshot,
3540 ) -> Result<RowIdSet> {
3541 use crate::query::Condition;
3542 Ok(match c {
3543 Condition::Pk(key) => {
3544 let lookup = self
3545 .schema
3546 .primary_key()
3547 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3548 .unwrap_or_else(|| key.clone());
3549 self.hot
3550 .get(&lookup)
3551 .map(|r| RowIdSet::one(r.0))
3552 .unwrap_or_else(RowIdSet::empty)
3553 }
3554 Condition::BitmapEq { column_id, value } => {
3555 let lookup = self.index_lookup_key_bytes(*column_id, value);
3556 self.bitmap
3557 .get(column_id)
3558 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3559 .unwrap_or_else(RowIdSet::empty)
3560 }
3561 Condition::BitmapIn { column_id, values } => {
3562 let bm = self.bitmap.get(column_id);
3563 let mut acc = roaring::RoaringBitmap::new();
3564 if let Some(b) = bm {
3565 for v in values {
3566 let lookup = self.index_lookup_key_bytes(*column_id, v);
3567 acc |= b.get(&lookup);
3568 }
3569 }
3570 RowIdSet::from_roaring(acc)
3571 }
3572 Condition::BytesPrefix { column_id, prefix } => {
3573 if let Some(b) = self.bitmap.get(column_id) {
3578 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
3579 let mut acc = roaring::RoaringBitmap::new();
3580 for key in b.keys() {
3581 if key.starts_with(&lookup_prefix) {
3582 acc |= b.get(key);
3583 }
3584 }
3585 RowIdSet::from_roaring(acc)
3586 } else {
3587 RowIdSet::empty()
3588 }
3589 }
3590 Condition::FmContains { column_id, pattern } => self
3591 .fm
3592 .get(column_id)
3593 .map(|f| {
3594 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3595 })
3596 .unwrap_or_else(RowIdSet::empty),
3597 Condition::FmContainsAll {
3598 column_id,
3599 patterns,
3600 } => {
3601 if let Some(f) = self.fm.get(column_id) {
3604 let sets: Vec<RowIdSet> = patterns
3605 .iter()
3606 .map(|pat| {
3607 RowIdSet::from_unsorted(
3608 f.locate(pat).into_iter().map(|r| r.0).collect(),
3609 )
3610 })
3611 .collect();
3612 RowIdSet::intersect_many(sets)
3613 } else {
3614 RowIdSet::empty()
3615 }
3616 }
3617 Condition::Ann {
3618 column_id,
3619 query,
3620 k,
3621 } => self
3622 .ann
3623 .get(column_id)
3624 .map(|a| {
3625 RowIdSet::from_unsorted(
3626 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3627 )
3628 })
3629 .unwrap_or_else(RowIdSet::empty),
3630 Condition::SparseMatch {
3631 column_id,
3632 query,
3633 k,
3634 } => self
3635 .sparse
3636 .get(column_id)
3637 .map(|s| {
3638 RowIdSet::from_unsorted(
3639 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3640 )
3641 })
3642 .unwrap_or_else(RowIdSet::empty),
3643 Condition::MinHashSimilar {
3644 column_id,
3645 query,
3646 k,
3647 } => self
3648 .minhash
3649 .get(column_id)
3650 .map(|mh| {
3651 RowIdSet::from_unsorted(
3652 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3653 )
3654 })
3655 .unwrap_or_else(RowIdSet::empty),
3656 Condition::Range { column_id, lo, hi } => {
3657 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3666 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3667 } else if self.run_refs.len() == 1 {
3668 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3669 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3670 } else {
3671 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3672 };
3673 set.remove_many(self.overlay_rid_set(snapshot));
3674 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3675 set
3676 }
3677 Condition::RangeF64 {
3678 column_id,
3679 lo,
3680 lo_inclusive,
3681 hi,
3682 hi_inclusive,
3683 } => {
3684 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3687 RowIdSet::from_unsorted(
3688 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3689 .into_iter()
3690 .collect(),
3691 )
3692 } else if self.run_refs.len() == 1 {
3693 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3694 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3695 } else {
3696 return self.range_scan_f64(
3697 *column_id,
3698 *lo,
3699 *lo_inclusive,
3700 *hi,
3701 *hi_inclusive,
3702 snapshot,
3703 );
3704 };
3705 set.remove_many(self.overlay_rid_set(snapshot));
3706 self.range_scan_overlay_f64(
3707 &mut set,
3708 *column_id,
3709 *lo,
3710 *lo_inclusive,
3711 *hi,
3712 *hi_inclusive,
3713 snapshot,
3714 );
3715 set
3716 }
3717 Condition::IsNull { column_id } => {
3718 let mut set = if self.run_refs.len() == 1 {
3719 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3720 r.null_row_id_set(*column_id, true)?
3721 } else {
3722 return self.null_scan(*column_id, true, snapshot);
3723 };
3724 set.remove_many(self.overlay_rid_set(snapshot));
3725 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3726 set
3727 }
3728 Condition::IsNotNull { column_id } => {
3729 let mut set = if self.run_refs.len() == 1 {
3730 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3731 r.null_row_id_set(*column_id, false)?
3732 } else {
3733 return self.null_scan(*column_id, false, snapshot);
3734 };
3735 set.remove_many(self.overlay_rid_set(snapshot));
3736 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3737 set
3738 }
3739 })
3740 }
3741
3742 fn range_scan_i64(
3750 &self,
3751 column_id: u16,
3752 lo: i64,
3753 hi: i64,
3754 snapshot: Snapshot,
3755 ) -> Result<RowIdSet> {
3756 let mut row_ids = Vec::new();
3757 let overlay_rids = self.overlay_rid_set(snapshot);
3758 for rr in &self.run_refs {
3759 let mut reader = self.open_reader(rr.run_id)?;
3760 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3761 for rid in matched {
3762 if !overlay_rids.contains(&rid) {
3763 row_ids.push(rid);
3764 }
3765 }
3766 }
3767 let mut s = RowIdSet::from_unsorted(row_ids);
3768 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3769 Ok(s)
3770 }
3771
3772 fn range_scan_f64(
3775 &self,
3776 column_id: u16,
3777 lo: f64,
3778 lo_inclusive: bool,
3779 hi: f64,
3780 hi_inclusive: bool,
3781 snapshot: Snapshot,
3782 ) -> Result<RowIdSet> {
3783 let mut row_ids = Vec::new();
3784 let overlay_rids = self.overlay_rid_set(snapshot);
3785 for rr in &self.run_refs {
3786 let mut reader = self.open_reader(rr.run_id)?;
3787 let matched = reader.range_row_ids_visible_f64(
3788 column_id,
3789 lo,
3790 lo_inclusive,
3791 hi,
3792 hi_inclusive,
3793 snapshot.epoch,
3794 )?;
3795 for rid in matched {
3796 if !overlay_rids.contains(&rid) {
3797 row_ids.push(rid);
3798 }
3799 }
3800 }
3801 let mut s = RowIdSet::from_unsorted(row_ids);
3802 self.range_scan_overlay_f64(
3803 &mut s,
3804 column_id,
3805 lo,
3806 lo_inclusive,
3807 hi,
3808 hi_inclusive,
3809 snapshot,
3810 );
3811 Ok(s)
3812 }
3813
3814 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3816 let mut s = HashSet::new();
3817 for row in self.memtable.visible_versions(snapshot.epoch) {
3818 s.insert(row.row_id.0);
3819 }
3820 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3821 s.insert(row.row_id.0);
3822 }
3823 s
3824 }
3825
3826 fn range_scan_overlay_i64(
3827 &self,
3828 s: &mut RowIdSet,
3829 column_id: u16,
3830 lo: i64,
3831 hi: i64,
3832 snapshot: Snapshot,
3833 ) {
3834 let mut newest: HashMap<u64, &Row> = HashMap::new();
3839 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3840 let memtable = self.memtable.visible_versions(snapshot.epoch);
3841 for r in &mutable {
3842 newest.entry(r.row_id.0).or_insert(r);
3843 }
3844 for r in &memtable {
3845 newest.insert(r.row_id.0, r);
3846 }
3847 for row in newest.values() {
3848 if !row.deleted {
3849 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3850 if *v >= lo && *v <= hi {
3851 s.insert(row.row_id.0);
3852 }
3853 }
3854 }
3855 }
3856 }
3857
3858 #[allow(clippy::too_many_arguments)]
3859 fn range_scan_overlay_f64(
3860 &self,
3861 s: &mut RowIdSet,
3862 column_id: u16,
3863 lo: f64,
3864 lo_inclusive: bool,
3865 hi: f64,
3866 hi_inclusive: bool,
3867 snapshot: Snapshot,
3868 ) {
3869 let mut newest: HashMap<u64, &Row> = HashMap::new();
3872 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3873 let memtable = self.memtable.visible_versions(snapshot.epoch);
3874 for r in &mutable {
3875 newest.entry(r.row_id.0).or_insert(r);
3876 }
3877 for r in &memtable {
3878 newest.insert(r.row_id.0, r);
3879 }
3880 for row in newest.values() {
3881 if !row.deleted {
3882 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3883 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3884 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3885 if ok_lo && ok_hi {
3886 s.insert(row.row_id.0);
3887 }
3888 }
3889 }
3890 }
3891 }
3892
3893 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3896 let mut row_ids = Vec::new();
3897 let overlay_rids = self.overlay_rid_set(snapshot);
3898 for rr in &self.run_refs {
3899 let mut reader = self.open_reader(rr.run_id)?;
3900 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3901 for rid in matched {
3902 if !overlay_rids.contains(&rid) {
3903 row_ids.push(rid);
3904 }
3905 }
3906 }
3907 let mut s = RowIdSet::from_unsorted(row_ids);
3908 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3909 Ok(s)
3910 }
3911
3912 fn null_scan_overlay(
3916 &self,
3917 s: &mut RowIdSet,
3918 column_id: u16,
3919 want_nulls: bool,
3920 snapshot: Snapshot,
3921 ) {
3922 let mut newest: HashMap<u64, &Row> = HashMap::new();
3923 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3924 let memtable = self.memtable.visible_versions(snapshot.epoch);
3925 for r in &mutable {
3926 newest.entry(r.row_id.0).or_insert(r);
3927 }
3928 for r in &memtable {
3929 newest.insert(r.row_id.0, r);
3930 }
3931 for row in newest.values() {
3932 if row.deleted {
3933 continue;
3934 }
3935 let is_null = !row.columns.contains_key(&column_id)
3936 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3937 if is_null == want_nulls {
3938 s.insert(row.row_id.0);
3939 }
3940 }
3941 }
3942
3943 pub fn snapshot(&self) -> Snapshot {
3944 Snapshot::at(self.epoch.visible())
3945 }
3946
3947 pub fn pin_snapshot(&mut self) -> Snapshot {
3950 let e = self.epoch.visible();
3951 *self.pinned.entry(e).or_insert(0) += 1;
3952 Snapshot::at(e)
3953 }
3954
3955 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3957 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3958 *count -= 1;
3959 if *count == 0 {
3960 self.pinned.remove(&snap.epoch);
3961 }
3962 }
3963 }
3964
3965 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3975 let local = self.pinned.keys().next().copied();
3976 let global = self.snapshots.min_pinned();
3977 match (local, global) {
3978 (Some(a), Some(b)) => Some(a.min(b)),
3979 (Some(a), None) => Some(a),
3980 (None, b) => b,
3981 }
3982 }
3983
3984 pub fn current_epoch(&self) -> Epoch {
3985 self.epoch.visible()
3986 }
3987
3988 pub fn memtable_len(&self) -> usize {
3989 self.memtable.len()
3990 }
3991
3992 pub fn count(&self) -> u64 {
3995 self.live_count
3996 }
3997
3998 pub fn count_conditions(
4002 &mut self,
4003 conditions: &[crate::query::Condition],
4004 snapshot: Snapshot,
4005 ) -> Result<Option<u64>> {
4006 use crate::query::Condition;
4007 if conditions.is_empty() {
4008 return Ok(Some(self.live_count));
4009 }
4010 let served = |c: &Condition| {
4011 matches!(
4012 c,
4013 Condition::Pk(_)
4014 | Condition::BitmapEq { .. }
4015 | Condition::BitmapIn { .. }
4016 | Condition::BytesPrefix { .. }
4017 | Condition::FmContains { .. }
4018 | Condition::FmContainsAll { .. }
4019 | Condition::Ann { .. }
4020 | Condition::Range { .. }
4021 | Condition::RangeF64 { .. }
4022 | Condition::SparseMatch { .. }
4023 | Condition::MinHashSimilar { .. }
4024 | Condition::IsNull { .. }
4025 | Condition::IsNotNull { .. }
4026 )
4027 };
4028 if !conditions.iter().all(served) {
4029 return Ok(None);
4030 }
4031 self.ensure_indexes_complete()?;
4032 let mut sets = Vec::with_capacity(conditions.len());
4033 for condition in conditions {
4034 sets.push(self.resolve_condition(condition, snapshot)?);
4035 }
4036 let mut rids = RowIdSet::intersect_many(sets);
4037 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4047 rids.remove_many(self.overlay_tombstoned_rids(snapshot));
4048 }
4049 let count = rids.len() as u64;
4050 crate::trace::QueryTrace::record(|t| {
4051 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
4052 t.survivor_count = Some(count as usize);
4053 t.conditions_pushed = conditions.len();
4054 });
4055 Ok(Some(count))
4056 }
4057
4058 fn overlay_tombstoned_rids(&self, snapshot: Snapshot) -> Vec<u64> {
4063 let mut out = Vec::new();
4064 for row in self.memtable.visible_versions(snapshot.epoch) {
4065 if row.deleted {
4066 out.push(row.row_id.0);
4067 }
4068 }
4069 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4070 if row.deleted {
4071 out.push(row.row_id.0);
4072 }
4073 }
4074 out
4075 }
4076
4077 pub fn bulk_load_columns(
4086 &mut self,
4087 user_columns: Vec<(u16, columnar::NativeColumn)>,
4088 ) -> Result<Epoch> {
4089 self.bulk_load_columns_with(user_columns, 3, false, true)
4090 }
4091
4092 pub fn bulk_load_fast(
4099 &mut self,
4100 user_columns: Vec<(u16, columnar::NativeColumn)>,
4101 ) -> Result<Epoch> {
4102 self.bulk_load_columns_with(user_columns, -1, true, false)
4103 }
4104
4105 fn bulk_load_columns_with(
4106 &mut self,
4107 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
4108 zstd_level: i32,
4109 force_plain: bool,
4110 lz4: bool,
4111 ) -> Result<Epoch> {
4112 let epoch = self.commit()?;
4113 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
4114 if n == 0 {
4115 return Ok(epoch);
4116 }
4117 let live_before = self.live_count;
4118 self.spill_mutable_run(epoch)?;
4120 let eager_index_build = self.index_build_policy == IndexBuildPolicy::Eager
4121 && self.indexes_complete
4122 && self.run_refs.is_empty()
4123 && self.memtable.is_empty()
4124 && self.mutable_run.is_empty();
4125 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
4128 self.validate_columns_not_null(&user_columns, n)?;
4129 let winner_idx = self
4130 .bulk_pk_winner_indices(&user_columns, n)
4131 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
4132 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
4133 match winner_idx.as_deref() {
4134 Some(idx) => {
4135 let compacted = user_columns
4136 .iter()
4137 .map(|(id, c)| (*id, c.gather(idx)))
4138 .collect();
4139 (compacted, idx.len())
4140 }
4141 None => (user_columns, n),
4142 };
4143 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
4144 let first = self.allocator.alloc_range(write_n as u64).0;
4145 for rid in first..first + write_n as u64 {
4146 self.reservoir.offer(rid);
4147 }
4148 let run_id = self.next_run_id;
4149 self.next_run_id += 1;
4150 let path = self.run_path(run_id);
4151 let mut writer =
4152 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
4153 if force_plain {
4154 writer = writer.with_plain();
4155 } else if lz4 {
4156 writer = writer.with_lz4();
4159 } else {
4160 writer = writer.with_zstd_level(zstd_level);
4161 }
4162 if let Some(kek) = &self.kek {
4163 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
4164 }
4165 let header = writer.write_native(&path, &write_columns, write_n, first)?;
4166 self.run_refs.push(RunRef {
4167 run_id: run_id as u128,
4168 level: 0,
4169 epoch_created: epoch.0,
4170 row_count: header.row_count,
4171 });
4172 self.live_count = self.live_count.saturating_add(write_n as u64);
4173 if eager_index_build {
4174 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
4175 self.index_columns_bulk(&write_columns, &row_ids);
4176 self.indexes_complete = true;
4177 self.build_learned_ranges()?;
4178 } else {
4179 self.indexes_complete = false;
4183 }
4184 self.mark_flushed(epoch)?;
4185 self.persist_manifest(epoch)?;
4186 if eager_index_build {
4187 self.checkpoint_indexes(epoch);
4188 }
4189 self.clear_result_cache();
4190 Ok(epoch)
4191 }
4192
4193 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
4211 let n = row_ids.len();
4212 if n == 0 {
4213 return;
4214 }
4215 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
4216 columns.iter().map(|(id, c)| (*id, c)).collect();
4217 let ty_of: std::collections::HashMap<u16, TypeId> =
4218 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
4219 let pk_id = self.schema.primary_key().map(|c| c.id);
4220
4221 for (i, &rid) in row_ids.iter().enumerate() {
4222 let row_id = RowId(rid);
4223 if let Some(pid) = pk_id {
4224 if let Some(col) = by_id.get(&pid) {
4225 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
4226 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
4227 self.insert_hot_pk(key, row_id);
4228 }
4229 }
4230 }
4231 for idef in &self.schema.indexes {
4232 let Some(col) = by_id.get(&idef.column_id) else {
4233 continue;
4234 };
4235 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
4236 match idef.kind {
4237 IndexKind::Bitmap => {
4238 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
4239 if let Some(key) =
4240 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
4241 {
4242 b.insert(key, row_id);
4243 }
4244 }
4245 }
4246 IndexKind::FmIndex => {
4247 if let Some(f) = self.fm.get_mut(&idef.column_id) {
4248 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4249 f.insert(bytes.to_vec(), row_id);
4250 }
4251 }
4252 }
4253 IndexKind::Sparse => {
4254 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
4255 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4256 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
4257 s.insert(&terms, row_id);
4258 }
4259 }
4260 }
4261 }
4262 IndexKind::MinHash => {
4263 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
4264 if let Some(bytes) = columnar::native_bytes_at(col, i) {
4265 let tokens = crate::index::token_hashes_from_bytes(bytes);
4266 mh.insert(&tokens, row_id);
4267 }
4268 }
4269 }
4270 _ => {}
4271 }
4272 }
4273 }
4274 }
4275
4276 pub fn visible_columns_native(
4281 &self,
4282 snapshot: Snapshot,
4283 projection: Option<&[u16]>,
4284 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
4285 let wanted: Vec<u16> = match projection {
4286 Some(p) => p.to_vec(),
4287 None => self.schema.columns.iter().map(|c| c.id).collect(),
4288 };
4289 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
4290 let rr = self.run_refs[0].clone();
4291 let mut reader = self.open_reader(rr.run_id)?;
4292 let idxs = reader.visible_indices_native(snapshot.epoch)?;
4293 let all_visible = idxs.len() == reader.row_count();
4294 if reader.has_mmap() {
4300 use rayon::prelude::*;
4301 let valid: Vec<u16> = wanted
4304 .iter()
4305 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
4306 .copied()
4307 .collect();
4308 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
4310 .par_iter()
4311 .filter_map(|cid| {
4312 reader
4313 .column_native_shared(*cid)
4314 .ok()
4315 .map(|col| (*cid, col))
4316 })
4317 .collect();
4318 let cols = decoded
4319 .into_iter()
4320 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
4321 .collect();
4322 return Ok(cols);
4323 }
4324 let mut cols = Vec::with_capacity(wanted.len());
4325 for cid in &wanted {
4326 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
4327 Some(c) => c,
4328 None => continue,
4329 };
4330 let col = reader.column_native(cdef.id)?;
4331 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
4332 }
4333 return Ok(cols);
4334 }
4335 let vcols = self.visible_columns(snapshot)?;
4336 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
4337 let out: Vec<(u16, columnar::NativeColumn)> = vcols
4338 .into_iter()
4339 .filter(|(id, _)| want_set.contains(id))
4340 .map(|(id, vals)| {
4341 let ty = self
4342 .schema
4343 .columns
4344 .iter()
4345 .find(|c| c.id == id)
4346 .map(|c| c.ty)
4347 .unwrap_or(TypeId::Bytes);
4348 (id, columnar::values_to_native(ty, &vals))
4349 })
4350 .collect();
4351 Ok(out)
4352 }
4353
4354 pub fn run_count(&self) -> usize {
4355 self.run_refs.len()
4356 }
4357
4358 pub fn memtable_is_empty(&self) -> bool {
4360 self.memtable.is_empty()
4361 }
4362
4363 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
4367 self.page_cache.stats()
4368 }
4369
4370 pub fn reset_page_cache_stats(&self) {
4372 self.page_cache.reset_stats();
4373 }
4374
4375 pub fn run_ids(&self) -> Vec<u128> {
4378 self.run_refs.iter().map(|r| r.run_id).collect()
4379 }
4380
4381 pub fn single_run_is_clean(&self) -> bool {
4385 if self.run_refs.len() != 1 {
4386 return false;
4387 }
4388 self.open_reader(self.run_refs[0].run_id)
4389 .map(|r| r.is_clean())
4390 .unwrap_or(false)
4391 }
4392
4393 fn resolve_footprint(
4400 &self,
4401 conditions: &[crate::query::Condition],
4402 _snapshot: Snapshot,
4403 ) -> roaring::RoaringBitmap {
4404 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
4405 return roaring::RoaringBitmap::new();
4406 }
4407 if self.run_refs.is_empty() {
4408 return roaring::RoaringBitmap::new();
4409 }
4410 if self.run_refs.len() == 1 {
4412 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
4413 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
4414 return rids.to_roaring_lossy();
4415 }
4416 }
4417 }
4418 roaring::RoaringBitmap::new()
4419 }
4420
4421 pub fn query_columns_native_cached(
4432 &mut self,
4433 conditions: &[crate::query::Condition],
4434 projection: Option<&[u16]>,
4435 snapshot: Snapshot,
4436 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4437 if conditions.is_empty() {
4438 return self.query_columns_native(conditions, projection, snapshot);
4439 }
4440 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
4444 if let Some(hit) = self.result_cache.lock().get_columns(key) {
4445 crate::trace::QueryTrace::record(|t| {
4446 t.result_cache_hit = true;
4447 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4448 });
4449 return Ok(Some((*hit).clone()));
4450 }
4451 let res = self.query_columns_native(conditions, projection, snapshot)?;
4452 if let Some(cols) = &res {
4453 let footprint = self.resolve_footprint(conditions, snapshot);
4454 let condition_cols = crate::query::condition_columns(conditions);
4455 self.result_cache.lock().insert(
4456 key,
4457 CachedEntry {
4458 data: CachedData::Columns(Arc::new(cols.clone())),
4459 footprint,
4460 condition_cols,
4461 },
4462 );
4463 }
4464 Ok(res)
4465 }
4466
4467 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4472 if q.conditions.is_empty() {
4473 return self.query(q);
4474 }
4475 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4476 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4477 crate::trace::QueryTrace::record(|t| {
4478 t.result_cache_hit = true;
4479 t.scan_mode = crate::trace::ScanMode::Materialized;
4480 });
4481 return Ok((*hit).clone());
4482 }
4483 let rows = self.query(q)?;
4484 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4485 let condition_cols = crate::query::condition_columns(&q.conditions);
4486 self.result_cache.lock().insert(
4487 key,
4488 CachedEntry {
4489 data: CachedData::Rows(Arc::new(rows.clone())),
4490 footprint,
4491 condition_cols,
4492 },
4493 );
4494 Ok(rows)
4495 }
4496
4497 #[allow(clippy::type_complexity)]
4512 pub fn query_columns_native_traced(
4513 &mut self,
4514 conditions: &[crate::query::Condition],
4515 projection: Option<&[u16]>,
4516 snapshot: Snapshot,
4517 ) -> Result<(
4518 Option<Vec<(u16, columnar::NativeColumn)>>,
4519 crate::trace::QueryTrace,
4520 )> {
4521 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4522 self.query_columns_native(conditions, projection, snapshot)
4523 });
4524 Ok((result?, trace))
4525 }
4526
4527 #[allow(clippy::type_complexity)]
4530 pub fn query_columns_native_cached_traced(
4531 &mut self,
4532 conditions: &[crate::query::Condition],
4533 projection: Option<&[u16]>,
4534 snapshot: Snapshot,
4535 ) -> Result<(
4536 Option<Vec<(u16, columnar::NativeColumn)>>,
4537 crate::trace::QueryTrace,
4538 )> {
4539 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4540 self.query_columns_native_cached(conditions, projection, snapshot)
4541 });
4542 Ok((result?, trace))
4543 }
4544
4545 pub fn native_page_cursor_traced(
4547 &self,
4548 snapshot: Snapshot,
4549 projection: Vec<(u16, TypeId)>,
4550 conditions: &[crate::query::Condition],
4551 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4552 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4553 self.native_page_cursor(snapshot, projection, conditions)
4554 });
4555 Ok((result?, trace))
4556 }
4557
4558 pub fn native_multi_run_cursor_traced(
4560 &self,
4561 snapshot: Snapshot,
4562 projection: Vec<(u16, TypeId)>,
4563 conditions: &[crate::query::Condition],
4564 ) -> Result<(
4565 Option<crate::cursor::MultiRunCursor>,
4566 crate::trace::QueryTrace,
4567 )> {
4568 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4569 self.native_multi_run_cursor(snapshot, projection, conditions)
4570 });
4571 Ok((result?, trace))
4572 }
4573
4574 pub fn count_conditions_traced(
4576 &mut self,
4577 conditions: &[crate::query::Condition],
4578 snapshot: Snapshot,
4579 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4580 let (result, trace) =
4581 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4582 Ok((result?, trace))
4583 }
4584
4585 pub fn query_traced(
4587 &mut self,
4588 q: &crate::query::Query,
4589 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4590 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4591 Ok((result?, trace))
4592 }
4593
4594 pub fn query_columns_native(
4599 &mut self,
4600 conditions: &[crate::query::Condition],
4601 projection: Option<&[u16]>,
4602 snapshot: Snapshot,
4603 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4604 use crate::query::Condition;
4605 if conditions.is_empty() {
4606 return Ok(None);
4607 }
4608 self.ensure_indexes_complete()?;
4609
4610 let served = |c: &Condition| {
4615 matches!(
4616 c,
4617 Condition::Pk(_)
4618 | Condition::BitmapEq { .. }
4619 | Condition::BitmapIn { .. }
4620 | Condition::BytesPrefix { .. }
4621 | Condition::FmContains { .. }
4622 | Condition::FmContainsAll { .. }
4623 | Condition::Ann { .. }
4624 | Condition::Range { .. }
4625 | Condition::RangeF64 { .. }
4626 | Condition::SparseMatch { .. }
4627 | Condition::MinHashSimilar { .. }
4628 | Condition::IsNull { .. }
4629 | Condition::IsNotNull { .. }
4630 )
4631 };
4632 if !conditions.iter().all(served) {
4633 return Ok(None);
4634 }
4635 let fast_path =
4636 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4637 crate::trace::QueryTrace::record(|t| {
4638 t.run_count = self.run_refs.len();
4639 t.memtable_rows = self.memtable.len();
4640 t.mutable_run_rows = self.mutable_run.len();
4641 t.conditions_pushed = conditions.len();
4642 t.learned_range_used = conditions.iter().any(|c| match c {
4643 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4644 self.learned_range.contains_key(column_id)
4645 }
4646 _ => false,
4647 });
4648 });
4649 let col_ids: Vec<u16> = projection
4651 .map(|p| p.to_vec())
4652 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4653 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4654 .iter()
4655 .map(|&cid| {
4656 let ty = self
4657 .schema
4658 .columns
4659 .iter()
4660 .find(|c| c.id == cid)
4661 .map(|c| c.ty)
4662 .unwrap_or(TypeId::Bytes);
4663 (cid, ty)
4664 })
4665 .collect();
4666
4667 if fast_path {
4673 let needs_column = conditions.iter().any(|c| match c {
4676 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4677 Condition::RangeF64 { column_id, .. } => {
4678 !self.learned_range.contains_key(column_id)
4679 }
4680 _ => false,
4681 });
4682 let mut reader_opt: Option<RunReader> = if needs_column {
4683 Some(self.open_reader(self.run_refs[0].run_id)?)
4684 } else {
4685 None
4686 };
4687 let mut sets: Vec<RowIdSet> = Vec::new();
4688 for c in conditions {
4689 let s = match c {
4690 Condition::Range { column_id, lo, hi }
4691 if !self.learned_range.contains_key(column_id) =>
4692 {
4693 if reader_opt.is_none() {
4694 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4695 }
4696 reader_opt
4697 .as_mut()
4698 .expect("reader opened for range")
4699 .range_row_id_set_i64(*column_id, *lo, *hi)?
4700 }
4701 Condition::RangeF64 {
4702 column_id,
4703 lo,
4704 lo_inclusive,
4705 hi,
4706 hi_inclusive,
4707 } if !self.learned_range.contains_key(column_id) => {
4708 if reader_opt.is_none() {
4709 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4710 }
4711 reader_opt
4712 .as_mut()
4713 .expect("reader opened for range")
4714 .range_row_id_set_f64(
4715 *column_id,
4716 *lo,
4717 *lo_inclusive,
4718 *hi,
4719 *hi_inclusive,
4720 )?
4721 }
4722 _ => self.resolve_condition(c, snapshot)?,
4723 };
4724 sets.push(s);
4725 }
4726 let candidates = RowIdSet::intersect_many(sets);
4727 crate::trace::QueryTrace::record(|t| {
4728 t.survivor_count = Some(candidates.len());
4729 });
4730 if candidates.is_empty() {
4731 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4732 .iter()
4733 .map(|&id| {
4734 (
4735 id,
4736 columnar::null_native(
4737 proj_pairs
4738 .iter()
4739 .find(|(c, _)| c == &id)
4740 .map(|(_, t)| *t)
4741 .unwrap_or(TypeId::Bytes),
4742 0,
4743 ),
4744 )
4745 })
4746 .collect();
4747 return Ok(Some(cols));
4748 }
4749 let mut reader = match reader_opt.take() {
4750 Some(r) => r,
4751 None => self.open_reader(self.run_refs[0].run_id)?,
4752 };
4753 let candidate_ids = candidates.into_sorted_vec();
4754 let (positions, fast_rid) = if let Some(positions) =
4755 reader.positions_for_row_ids_fast(&candidate_ids)
4756 {
4757 (positions, true)
4758 } else {
4759 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4760 match col {
4761 columnar::NativeColumn::Int64 { data, .. } => {
4762 let mut p: Vec<usize> = candidate_ids
4763 .iter()
4764 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4765 .collect();
4766 p.sort_unstable();
4767 (p, false)
4768 }
4769 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4770 }
4771 };
4772 crate::trace::QueryTrace::record(|t| {
4773 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4774 t.fast_row_id_map = fast_rid;
4775 });
4776 let mut cols = Vec::with_capacity(col_ids.len());
4777 for cid in &col_ids {
4778 let col = reader.column_native(*cid)?;
4779 cols.push((*cid, col.gather(&positions)));
4780 }
4781 return Ok(Some(cols));
4782 }
4783
4784 if !self.run_refs.is_empty() {
4797 use crate::cursor::{drain_cursor_to_columns, Cursor};
4798 let remaining: usize;
4799 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4800 let c = self
4801 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4802 .expect("single-run cursor should build when run_refs.len() == 1");
4803 remaining = c.remaining_rows();
4804 Box::new(c)
4805 } else {
4806 let c = self
4807 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4808 .expect("multi-run cursor should build when run_refs.len() >= 1");
4809 remaining = c.remaining_rows();
4810 Box::new(c)
4811 };
4812 crate::trace::QueryTrace::record(|t| {
4813 if t.survivor_count.is_none() {
4814 t.survivor_count = Some(remaining);
4815 }
4816 });
4817 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4818 return Ok(Some(cols));
4819 }
4820
4821 crate::trace::QueryTrace::record(|t| {
4826 t.scan_mode = crate::trace::ScanMode::Materialized;
4827 t.row_materialized = true;
4828 });
4829 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4830 for c in conditions {
4831 sets.push(self.resolve_condition(c, snapshot)?);
4832 }
4833 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4834 let rows = self.rows_for_rids(&rids, snapshot)?;
4835 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4836 for (cid, ty) in &proj_pairs {
4837 let vals: Vec<Value> = rows
4838 .iter()
4839 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4840 .collect();
4841 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4842 }
4843 Ok(Some(cols))
4844 }
4845
4846 pub fn native_page_cursor(
4861 &self,
4862 snapshot: Snapshot,
4863 projection: Vec<(u16, TypeId)>,
4864 conditions: &[crate::query::Condition],
4865 ) -> Result<Option<NativePageCursor>> {
4866 use crate::cursor::build_page_plans;
4867 if !conditions.is_empty() && !self.indexes_complete {
4870 return Ok(None);
4871 }
4872 if self.run_refs.len() != 1 {
4873 return Ok(None);
4874 }
4875 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4876 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4877
4878 let overlay_rids: HashSet<u64> = {
4881 let mut s = HashSet::new();
4882 for row in self.memtable.visible_versions(snapshot.epoch) {
4883 s.insert(row.row_id.0);
4884 }
4885 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4886 s.insert(row.row_id.0);
4887 }
4888 s
4889 };
4890
4891 let survivors = if conditions.is_empty() {
4895 None
4896 } else {
4897 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4898 };
4899
4900 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4907 survivors.clone()
4908 } else if let Some(s) = &survivors {
4909 let mut run_set = s.clone();
4910 run_set.remove_many(overlay_rids.iter().copied());
4911 Some(run_set)
4912 } else {
4913 Some(RowIdSet::from_unsorted(
4914 rids.iter()
4915 .map(|&r| r as u64)
4916 .filter(|r| !overlay_rids.contains(r))
4917 .collect(),
4918 ))
4919 };
4920
4921 let overlay_rows = if overlay_rids.is_empty() {
4922 Vec::new()
4923 } else {
4924 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4925 self.overlay_visible_rows(snapshot, bound)
4926 };
4927
4928 let plans = if positions.is_empty() {
4930 Vec::new()
4931 } else {
4932 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4933 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4934 };
4935
4936 let overlay = if overlay_rows.is_empty() {
4938 None
4939 } else {
4940 let filtered =
4941 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4942 if filtered.is_empty() {
4943 None
4944 } else {
4945 Some(self.materialize_overlay(&filtered, &projection))
4946 }
4947 };
4948
4949 let overlay_row_count = overlay
4950 .as_ref()
4951 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4952 .unwrap_or(0);
4953 crate::trace::QueryTrace::record(|t| {
4954 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4955 t.run_count = self.run_refs.len();
4956 t.memtable_rows = self.memtable.len();
4957 t.mutable_run_rows = self.mutable_run.len();
4958 t.overlay_rows = overlay_row_count;
4959 t.conditions_pushed = conditions.len();
4960 t.pages_decoded = plans
4961 .iter()
4962 .map(|p| p.positions.len())
4963 .sum::<usize>()
4964 .min(1);
4965 });
4966
4967 Ok(Some(NativePageCursor::new_with_overlay(
4968 reader, projection, plans, overlay,
4969 )))
4970 }
4971 #[allow(clippy::type_complexity)]
4981 pub fn native_multi_run_cursor(
4982 &self,
4983 snapshot: Snapshot,
4984 projection: Vec<(u16, TypeId)>,
4985 conditions: &[crate::query::Condition],
4986 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4987 use crate::cursor::{MultiRunCursor, RunStream};
4988 use crate::sorted_run::SYS_ROW_ID;
4989 use std::collections::{BinaryHeap, HashMap, HashSet};
4990 if !conditions.is_empty() && !self.indexes_complete {
4993 return Ok(None);
4994 }
4995 if self.run_refs.is_empty() {
4996 return Ok(None);
4997 }
4998
4999 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
5001 Vec::with_capacity(self.run_refs.len());
5002 for rr in &self.run_refs {
5003 let mut reader = self.open_reader(rr.run_id)?;
5004 let (rids, eps, del) = reader.system_columns_native()?;
5005 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
5006 run_meta.push((reader, rids, eps, del, page_rows));
5007 }
5008
5009 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
5013 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
5014 for i in 0..rids.len() {
5015 let rid = rids[i] as u64;
5016 let e = eps[i] as u64;
5017 if e > snapshot.epoch.0 {
5018 continue;
5019 }
5020 let is_del = del[i] != 0;
5021 best.entry(rid)
5022 .and_modify(|cur| {
5023 if e > cur.0 {
5024 *cur = (e, run_idx, i, is_del);
5025 }
5026 })
5027 .or_insert((e, run_idx, i, is_del));
5028 }
5029 }
5030
5031 let overlay_rids: HashSet<u64> = {
5033 let mut s = HashSet::new();
5034 for row in self.memtable.visible_versions(snapshot.epoch) {
5035 s.insert(row.row_id.0);
5036 }
5037 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5038 s.insert(row.row_id.0);
5039 }
5040 s
5041 };
5042
5043 let survivors: Option<RowIdSet> = if conditions.is_empty() {
5045 None
5046 } else {
5047 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5048 for c in conditions {
5049 sets.push(self.resolve_condition(c, snapshot)?);
5050 }
5051 Some(RowIdSet::intersect_many(sets))
5052 };
5053
5054 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
5058 for (rid, (_, run_idx, pos, deleted)) in &best {
5059 if *deleted {
5060 continue;
5061 }
5062 if overlay_rids.contains(rid) {
5063 continue;
5064 }
5065 if let Some(s) = &survivors {
5066 if !s.contains(*rid) {
5067 continue;
5068 }
5069 }
5070 per_run[*run_idx].push((*rid, *pos));
5071 }
5072 for v in per_run.iter_mut() {
5073 v.sort_unstable_by_key(|&(rid, _)| rid);
5074 }
5075
5076 let mut streams = Vec::with_capacity(run_meta.len());
5078 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
5079 let mut total = 0usize;
5080 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
5081 let mut starts = Vec::with_capacity(page_rows.len());
5082 let mut acc = 0usize;
5083 for &r in &page_rows {
5084 starts.push(acc);
5085 acc += r;
5086 }
5087 let mut survivors_vec: Vec<(u64, usize, usize)> =
5088 Vec::with_capacity(per_run[run_idx].len());
5089 for &(rid, pos) in &per_run[run_idx] {
5090 let page_seq = match starts.partition_point(|&s| s <= pos) {
5091 0 => continue,
5092 p => p - 1,
5093 };
5094 let within = pos - starts[page_seq];
5095 survivors_vec.push((rid, page_seq, within));
5096 }
5097 total += survivors_vec.len();
5098 if let Some(&(rid, _, _)) = survivors_vec.first() {
5099 heap.push(std::cmp::Reverse((rid, run_idx)));
5100 }
5101 streams.push(RunStream::new(reader, survivors_vec, page_rows));
5102 }
5103
5104 let overlay_rows = if overlay_rids.is_empty() {
5106 Vec::new()
5107 } else {
5108 let bound = Self::overlay_materialization_bound(conditions, &survivors);
5109 self.overlay_visible_rows(snapshot, bound)
5110 };
5111 let overlay = if overlay_rows.is_empty() {
5112 None
5113 } else {
5114 let filtered =
5115 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
5116 if filtered.is_empty() {
5117 None
5118 } else {
5119 Some(self.materialize_overlay(&filtered, &projection))
5120 }
5121 };
5122
5123 let overlay_row_count = overlay
5124 .as_ref()
5125 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
5126 .unwrap_or(0);
5127 crate::trace::QueryTrace::record(|t| {
5128 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
5129 t.run_count = self.run_refs.len();
5130 t.memtable_rows = self.memtable.len();
5131 t.mutable_run_rows = self.mutable_run.len();
5132 t.overlay_rows = overlay_row_count;
5133 t.conditions_pushed = conditions.len();
5134 t.survivor_count = Some(total);
5135 });
5136
5137 Ok(Some(MultiRunCursor::new(
5138 streams, projection, heap, total, overlay,
5139 )))
5140 }
5141
5142 fn overlay_materialization_bound<'a>(
5154 conditions: &[crate::query::Condition],
5155 survivors: &'a Option<RowIdSet>,
5156 ) -> Option<&'a RowIdSet> {
5157 use crate::query::Condition;
5158 let has_range = conditions
5159 .iter()
5160 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5161 if has_range {
5162 None
5163 } else {
5164 survivors.as_ref()
5165 }
5166 }
5167
5168 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
5180 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
5181 let mut fold = |row: Row| {
5182 if let Some(b) = bound {
5183 if !b.contains(row.row_id.0) {
5184 return;
5185 }
5186 }
5187 best.entry(row.row_id.0)
5188 .and_modify(|(be, br)| {
5189 if row.committed_epoch > *be {
5190 *be = row.committed_epoch;
5191 *br = row.clone();
5192 }
5193 })
5194 .or_insert_with(|| (row.committed_epoch, row));
5195 };
5196 for row in self.memtable.visible_versions(snapshot.epoch) {
5197 fold(row);
5198 }
5199 for row in self.mutable_run.visible_versions(snapshot.epoch) {
5200 fold(row);
5201 }
5202 let mut out: Vec<Row> = best
5203 .into_values()
5204 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
5205 .collect();
5206 out.sort_by_key(|r| r.row_id);
5207 out
5208 }
5209
5210 fn filter_overlay_rows(
5218 &self,
5219 rows: Vec<Row>,
5220 conditions: &[crate::query::Condition],
5221 survivors: Option<&RowIdSet>,
5222 snapshot: Snapshot,
5223 ) -> Result<Vec<Row>> {
5224 if conditions.is_empty() {
5225 return Ok(rows);
5226 }
5227 use crate::query::Condition;
5228 let all_index_served = !conditions
5232 .iter()
5233 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
5234 if all_index_served {
5235 return Ok(rows
5236 .into_iter()
5237 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
5238 .collect());
5239 }
5240 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
5243 for c in conditions {
5244 let s = match c {
5245 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
5246 _ => self.resolve_condition(c, snapshot)?,
5247 };
5248 per_cond_sets.push(s);
5249 }
5250 Ok(rows
5251 .into_iter()
5252 .filter(|row| {
5253 conditions.iter().enumerate().all(|(i, c)| match c {
5254 Condition::Range { column_id, lo, hi } => {
5255 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
5256 }
5257 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
5258 match row.columns.get(column_id) {
5259 Some(Value::Float64(v)) => {
5260 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
5261 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
5262 lo_ok && hi_ok
5263 }
5264 _ => false,
5265 }
5266 }
5267 _ => per_cond_sets[i].contains(row.row_id.0),
5268 })
5269 })
5270 .collect())
5271 }
5272
5273 fn materialize_overlay(
5276 &self,
5277 rows: &[Row],
5278 projection: &[(u16, TypeId)],
5279 ) -> Vec<columnar::NativeColumn> {
5280 if projection.is_empty() {
5281 return vec![columnar::null_native(TypeId::Int64, rows.len())];
5282 }
5283 let mut cols = Vec::with_capacity(projection.len());
5284 for (cid, ty) in projection {
5285 let vals: Vec<Value> = rows
5286 .iter()
5287 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
5288 .collect();
5289 cols.push(columnar::values_to_native(*ty, &vals));
5290 }
5291 cols
5292 }
5293
5294 fn resolve_survivor_rids(
5299 &self,
5300 conditions: &[crate::query::Condition],
5301 reader: &mut RunReader,
5302 ) -> Result<RowIdSet> {
5303 use crate::query::Condition;
5304 let mut sets: Vec<RowIdSet> = Vec::new();
5305 for c in conditions {
5306 let s: RowIdSet = match c {
5307 Condition::Pk(key) => {
5308 let lookup = self
5309 .schema
5310 .primary_key()
5311 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
5312 .unwrap_or_else(|| key.clone());
5313 self.hot
5314 .get(&lookup)
5315 .map(|r| RowIdSet::one(r.0))
5316 .unwrap_or_else(RowIdSet::empty)
5317 }
5318 Condition::BitmapEq { column_id, value } => {
5319 let lookup = self.index_lookup_key_bytes(*column_id, value);
5320 self.bitmap
5321 .get(column_id)
5322 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
5323 .unwrap_or_else(RowIdSet::empty)
5324 }
5325 Condition::BitmapIn { column_id, values } => {
5326 let bm = self.bitmap.get(column_id);
5327 let mut acc = roaring::RoaringBitmap::new();
5328 if let Some(b) = bm {
5329 for v in values {
5330 let lookup = self.index_lookup_key_bytes(*column_id, v);
5331 acc |= b.get(&lookup);
5332 }
5333 }
5334 RowIdSet::from_roaring(acc)
5335 }
5336 Condition::BytesPrefix { column_id, prefix } => {
5337 if let Some(b) = self.bitmap.get(column_id) {
5338 let lookup_prefix = self.index_lookup_key_bytes(*column_id, prefix);
5339 let mut acc = roaring::RoaringBitmap::new();
5340 for key in b.keys() {
5341 if key.starts_with(&lookup_prefix) {
5342 acc |= b.get(key);
5343 }
5344 }
5345 RowIdSet::from_roaring(acc)
5346 } else {
5347 RowIdSet::empty()
5348 }
5349 }
5350 Condition::FmContains { column_id, pattern } => self
5351 .fm
5352 .get(column_id)
5353 .map(|f| {
5354 RowIdSet::from_unsorted(
5355 f.locate(pattern).into_iter().map(|r| r.0).collect(),
5356 )
5357 })
5358 .unwrap_or_else(RowIdSet::empty),
5359 Condition::FmContainsAll {
5360 column_id,
5361 patterns,
5362 } => {
5363 if let Some(f) = self.fm.get(column_id) {
5364 let sets: Vec<RowIdSet> = patterns
5365 .iter()
5366 .map(|pat| {
5367 RowIdSet::from_unsorted(
5368 f.locate(pat).into_iter().map(|r| r.0).collect(),
5369 )
5370 })
5371 .collect();
5372 RowIdSet::intersect_many(sets)
5373 } else {
5374 RowIdSet::empty()
5375 }
5376 }
5377 Condition::Ann {
5378 column_id,
5379 query,
5380 k,
5381 } => self
5382 .ann
5383 .get(column_id)
5384 .map(|a| {
5385 RowIdSet::from_unsorted(
5386 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5387 )
5388 })
5389 .unwrap_or_else(RowIdSet::empty),
5390 Condition::SparseMatch {
5391 column_id,
5392 query,
5393 k,
5394 } => self
5395 .sparse
5396 .get(column_id)
5397 .map(|s| {
5398 RowIdSet::from_unsorted(
5399 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5400 )
5401 })
5402 .unwrap_or_else(RowIdSet::empty),
5403 Condition::MinHashSimilar {
5404 column_id,
5405 query,
5406 k,
5407 } => self
5408 .minhash
5409 .get(column_id)
5410 .map(|mh| {
5411 RowIdSet::from_unsorted(
5412 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
5413 )
5414 })
5415 .unwrap_or_else(RowIdSet::empty),
5416 Condition::Range { column_id, lo, hi } => {
5417 if let Some(li) = self.learned_range.get(column_id) {
5418 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
5419 } else {
5420 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
5421 }
5422 }
5423 Condition::RangeF64 {
5424 column_id,
5425 lo,
5426 lo_inclusive,
5427 hi,
5428 hi_inclusive,
5429 } => {
5430 if let Some(li) = self.learned_range.get(column_id) {
5431 RowIdSet::from_unsorted(
5432 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
5433 .into_iter()
5434 .collect(),
5435 )
5436 } else {
5437 reader.range_row_id_set_f64(
5438 *column_id,
5439 *lo,
5440 *lo_inclusive,
5441 *hi,
5442 *hi_inclusive,
5443 )?
5444 }
5445 }
5446 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
5447 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
5448 };
5449 sets.push(s);
5450 }
5451 Ok(RowIdSet::intersect_many(sets))
5452 }
5453
5454 pub fn scan_cursor(
5475 &self,
5476 snapshot: Snapshot,
5477 projection: Vec<(u16, TypeId)>,
5478 conditions: &[crate::query::Condition],
5479 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
5480 if !conditions.is_empty() && !self.indexes_complete {
5486 return Ok(None);
5487 }
5488 if self.run_refs.len() == 1 {
5489 Ok(self
5490 .native_page_cursor(snapshot, projection, conditions)?
5491 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5492 } else {
5493 Ok(self
5494 .native_multi_run_cursor(snapshot, projection, conditions)?
5495 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
5496 }
5497 }
5498
5499 pub fn aggregate_native(
5513 &self,
5514 snapshot: Snapshot,
5515 column: Option<u16>,
5516 conditions: &[crate::query::Condition],
5517 agg: NativeAgg,
5518 ) -> Result<Option<NativeAggResult>> {
5519 if self.run_refs.len() == 1 && conditions.is_empty() {
5521 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5522 return Ok(Some(res));
5523 }
5524 }
5525 if matches!(agg, NativeAgg::Count) && column.is_none() {
5527 return Ok(self
5528 .scan_cursor(snapshot, Vec::new(), conditions)?
5529 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5530 }
5531 let cid = match column {
5534 Some(c) => c,
5535 None => return Ok(None),
5536 };
5537 let ty = self.column_type(cid);
5538 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5539 return Ok(None);
5540 };
5541 match ty {
5542 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5543 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5544 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5545 }
5546 TypeId::Float64 => {
5547 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5548 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5549 }
5550 _ => Ok(None),
5551 }
5552 }
5553
5554 fn aggregate_from_stats(
5562 &self,
5563 snapshot: Snapshot,
5564 column: Option<u16>,
5565 agg: NativeAgg,
5566 ) -> Result<Option<NativeAggResult>> {
5567 let cid = match (agg, column) {
5568 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5569 _ => return Ok(None), };
5571 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5572 return Ok(None);
5573 };
5574 let Some(cs) = stats.get(&cid) else {
5575 return Ok(None);
5576 };
5577 match agg {
5578 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5580 self.live_count.saturating_sub(cs.null_count),
5581 ))),
5582 NativeAgg::Min | NativeAgg::Max => {
5583 let bound = if agg == NativeAgg::Min {
5584 &cs.min
5585 } else {
5586 &cs.max
5587 };
5588 match bound {
5589 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5590 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5591 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5596 None => Ok(None),
5597 }
5598 }
5599 _ => Ok(None),
5600 }
5601 }
5602
5603 pub fn count_distinct_from_bitmap(&mut self, column_id: u16) -> Result<Option<u64>> {
5612 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5613 return Ok(None);
5614 }
5615 self.ensure_indexes_complete()?;
5618 let reader = self.open_reader(self.run_refs[0].run_id)?;
5619 if self.live_count != reader.row_count() as u64 {
5620 return Ok(None);
5621 }
5622 let Some(bm) = self.bitmap.get(&column_id) else {
5623 return Ok(None); };
5625 let mut distinct = bm.value_count() as u64;
5626 if !bm.get(&Value::Null.encode_key()).is_empty() {
5629 distinct = distinct.saturating_sub(1);
5630 }
5631 Ok(Some(distinct))
5632 }
5633
5634 pub fn aggregate_incremental(
5646 &mut self,
5647 cache_key: u64,
5648 conditions: &[crate::query::Condition],
5649 column: Option<u16>,
5650 agg: NativeAgg,
5651 ) -> Result<IncrementalAggResult> {
5652 let snap = self.snapshot();
5653 let cur_wm = self.allocator.current().0;
5654 let cur_epoch = snap.epoch.0;
5655 let incremental_ok =
5662 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5663
5664 if incremental_ok {
5667 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5668 if cached.epoch == cur_epoch {
5669 return Ok(IncrementalAggResult {
5670 state: cached.state,
5671 incremental: true,
5672 delta_rows: 0,
5673 });
5674 }
5675 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5676 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5677 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5678 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5679 let delta_state = agg_state_from_rows(
5680 &delta_rows,
5681 conditions,
5682 &index_sets,
5683 column,
5684 agg,
5685 &self.schema,
5686 )?;
5687 let merged = cached.state.merge(delta_state);
5688 let delta_n = delta_rids.len() as u64;
5689 self.agg_cache.insert(
5690 cache_key,
5691 CachedAgg {
5692 state: merged.clone(),
5693 watermark: cur_wm,
5694 epoch: cur_epoch,
5695 },
5696 );
5697 return Ok(IncrementalAggResult {
5698 state: merged,
5699 incremental: true,
5700 delta_rows: delta_n,
5701 });
5702 }
5703 }
5704 }
5705
5706 let cursor_ok =
5711 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5712 let state = if cursor_ok && agg != NativeAgg::Avg {
5713 match self.aggregate_native(snap, column, conditions, agg)? {
5714 Some(result) => {
5715 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5716 }
5717 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5718 }
5719 } else {
5720 self.agg_state_full_scan(conditions, column, agg, snap)?
5721 };
5722 if incremental_ok {
5724 self.agg_cache.insert(
5725 cache_key,
5726 CachedAgg {
5727 state: state.clone(),
5728 watermark: cur_wm,
5729 epoch: cur_epoch,
5730 },
5731 );
5732 }
5733 Ok(IncrementalAggResult {
5734 state,
5735 incremental: false,
5736 delta_rows: 0,
5737 })
5738 }
5739
5740 fn agg_state_full_scan(
5743 &self,
5744 conditions: &[crate::query::Condition],
5745 column: Option<u16>,
5746 agg: NativeAgg,
5747 snap: Snapshot,
5748 ) -> Result<AggState> {
5749 let rows = self.visible_rows(snap)?;
5750 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5751 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5752 }
5753
5754 fn resolve_index_conditions(
5757 &self,
5758 conditions: &[crate::query::Condition],
5759 snapshot: Snapshot,
5760 ) -> Result<Vec<RowIdSet>> {
5761 use crate::query::Condition;
5762 let mut sets = Vec::new();
5763 for c in conditions {
5764 if matches!(
5765 c,
5766 Condition::Ann { .. }
5767 | Condition::SparseMatch { .. }
5768 | Condition::MinHashSimilar { .. }
5769 ) {
5770 sets.push(self.resolve_condition(c, snapshot)?);
5771 }
5772 }
5773 Ok(sets)
5774 }
5775
5776 fn column_type(&self, cid: u16) -> TypeId {
5777 self.schema
5778 .columns
5779 .iter()
5780 .find(|c| c.id == cid)
5781 .map(|c| c.ty)
5782 .unwrap_or(TypeId::Bytes)
5783 }
5784
5785 pub fn approx_aggregate(
5794 &mut self,
5795 conditions: &[crate::query::Condition],
5796 column: Option<u16>,
5797 agg: ApproxAgg,
5798 z: f64,
5799 ) -> Result<Option<ApproxResult>> {
5800 use crate::query::Condition;
5801 self.ensure_reservoir_complete()?;
5802 let snapshot = self.snapshot();
5803 let n_pop = self.live_count;
5804 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5805 if sample_rids.is_empty() {
5806 return Ok(None);
5807 }
5808 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5810 let s = live_sample.len();
5811 if s == 0 {
5812 return Ok(None);
5813 }
5814
5815 let mut index_sets: Vec<RowIdSet> = Vec::new();
5818 for c in conditions {
5819 if matches!(
5820 c,
5821 Condition::Ann { .. }
5822 | Condition::SparseMatch { .. }
5823 | Condition::MinHashSimilar { .. }
5824 ) {
5825 index_sets.push(self.resolve_condition(c, snapshot)?);
5826 }
5827 }
5828
5829 let cid = match (agg, column) {
5831 (ApproxAgg::Count, _) => None,
5832 (_, Some(c)) => Some(c),
5833 _ => return Ok(None),
5834 };
5835 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5836 for r in &live_sample {
5837 if !conditions
5839 .iter()
5840 .all(|c| condition_matches_row(c, r, &self.schema))
5841 {
5842 continue;
5843 }
5844 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5846 continue;
5847 }
5848 if let Some(cid) = cid {
5849 if let Some(v) = as_f64(r.columns.get(&cid)) {
5850 passing_vals.push(v);
5851 } } else {
5853 passing_vals.push(0.0); }
5855 }
5856 let m = passing_vals.len();
5857
5858 let (point, half) = match agg {
5859 ApproxAgg::Count => {
5860 let p = m as f64 / s as f64;
5862 let point = n_pop as f64 * p;
5863 let var = if s > 1 {
5864 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5865 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5866 } else {
5867 0.0
5868 };
5869 (point, z * var.sqrt())
5870 }
5871 ApproxAgg::Sum => {
5872 let y: Vec<f64> = live_sample
5874 .iter()
5875 .map(|r| {
5876 let passes_row = conditions
5877 .iter()
5878 .all(|c| condition_matches_row(c, r, &self.schema))
5879 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5880 if passes_row {
5881 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5882 } else {
5883 0.0
5884 }
5885 })
5886 .collect();
5887 let mean_y = y.iter().sum::<f64>() / s as f64;
5888 let point = n_pop as f64 * mean_y;
5889 let var = if s > 1 {
5890 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5891 let var_y = ss / (s - 1) as f64;
5892 n_pop as f64 * n_pop as f64 * var_y / s as f64
5893 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5894 } else {
5895 0.0
5896 };
5897 (point, z * var.sqrt())
5898 }
5899 ApproxAgg::Avg => {
5900 if m == 0 {
5901 return Ok(Some(ApproxResult {
5902 point: 0.0,
5903 ci_low: 0.0,
5904 ci_high: 0.0,
5905 n_population: n_pop,
5906 n_sample_live: s,
5907 n_passing: 0,
5908 }));
5909 }
5910 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5911 let half = if m > 1 {
5912 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5913 let sd = (ss / (m - 1) as f64).sqrt();
5914 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5915 z * sd / (m as f64).sqrt() * fpc.sqrt()
5916 } else {
5917 0.0
5918 };
5919 (mean, half)
5920 }
5921 };
5922
5923 Ok(Some(ApproxResult {
5924 point,
5925 ci_low: point - half,
5926 ci_high: point + half,
5927 n_population: n_pop,
5928 n_sample_live: s,
5929 n_passing: m,
5930 }))
5931 }
5932
5933 pub fn exact_column_stats(
5941 &self,
5942 _snapshot: Snapshot,
5943 projection: &[u16],
5944 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5945 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5946 return Ok(None);
5947 }
5948 let reader = self.open_reader(self.run_refs[0].run_id)?;
5949 if self.live_count != reader.row_count() as u64 {
5950 return Ok(None);
5951 }
5952 let mut out = HashMap::new();
5953 for &cid in projection {
5954 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5955 Some(c) => c,
5956 None => continue,
5957 };
5958 let Some(stats) = reader.column_page_stats(cid) else {
5960 out.insert(
5961 cid,
5962 ColumnStat {
5963 min: None,
5964 max: None,
5965 null_count: self.live_count,
5966 },
5967 );
5968 continue;
5969 };
5970 let stat = match cdef.ty {
5971 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5972 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5973 min: mn.map(Value::Int64),
5974 max: mx.map(Value::Int64),
5975 null_count: n,
5976 })
5977 }
5978 TypeId::Float64 => {
5979 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5980 min: mn.map(Value::Float64),
5981 max: mx.map(Value::Float64),
5982 null_count: n,
5983 })
5984 }
5985 _ => None,
5986 };
5987 if let Some(s) = stat {
5988 out.insert(cid, s);
5989 }
5990 }
5991 Ok(Some(out))
5992 }
5993
5994 pub fn dir(&self) -> &Path {
5995 &self.dir
5996 }
5997
5998 pub fn schema(&self) -> &Schema {
5999 &self.schema
6000 }
6001
6002 pub(crate) fn prepare_alter_column(
6003 &mut self,
6004 column_name: &str,
6005 change: &AlterColumn,
6006 ) -> Result<ColumnDef> {
6007 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
6008 return Err(MongrelError::InvalidArgument(
6009 "ALTER COLUMN requires committing staged writes first".into(),
6010 ));
6011 }
6012 let old = self
6013 .schema
6014 .columns
6015 .iter()
6016 .find(|c| c.name == column_name)
6017 .cloned()
6018 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6019 let mut next = old.clone();
6020
6021 if let Some(name) = &change.name {
6022 let trimmed = name.trim();
6023 if trimmed.is_empty() {
6024 return Err(MongrelError::InvalidArgument(
6025 "ALTER COLUMN name must not be empty".into(),
6026 ));
6027 }
6028 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
6029 return Err(MongrelError::Schema(format!(
6030 "column {trimmed} already exists"
6031 )));
6032 }
6033 next.name = trimmed.to_string();
6034 }
6035
6036 if let Some(ty) = change.ty {
6037 next.ty = ty;
6038 }
6039 if let Some(flags) = change.flags {
6040 validate_alter_column_flags(old.flags, flags)?;
6041 next.flags = flags;
6042 }
6043
6044 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
6045 if old.flags.contains(ColumnFlags::NULLABLE)
6046 && !next.flags.contains(ColumnFlags::NULLABLE)
6047 && self.column_has_nulls(old.id)?
6048 {
6049 return Err(MongrelError::InvalidArgument(format!(
6050 "column '{}' contains NULL values",
6051 old.name
6052 )));
6053 }
6054 Ok(next)
6055 }
6056
6057 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
6058 let idx = self
6059 .schema
6060 .columns
6061 .iter()
6062 .position(|c| c.id == column.id)
6063 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
6064 if self.schema.columns[idx] == column {
6065 return Ok(());
6066 }
6067 self.schema.columns[idx] = column;
6068 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6069 self.schema.validate_auto_increment()?;
6070 self.auto_inc = resolve_auto_inc(&self.schema);
6071 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
6072 write_schema(&self.dir, &self.schema)?;
6073 self.clear_result_cache();
6074 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6075 self.persist_manifest(self.current_epoch())?;
6076 Ok(())
6077 }
6078
6079 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
6080 let column = self.prepare_alter_column(column_name, &change)?;
6081 self.apply_altered_column(column.clone())?;
6082 Ok(column)
6083 }
6084
6085 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
6086 if self.live_count == 0 {
6087 return Ok(false);
6088 }
6089 let snap = self.snapshot();
6090 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
6091 Ok(columns
6092 .first()
6093 .map(|(_, col)| col.null_count(col.len()) != 0)
6094 .unwrap_or(true))
6095 }
6096
6097 fn has_stored_versions(&self) -> bool {
6098 !self.memtable.is_empty()
6099 || !self.mutable_run.is_empty()
6100 || self.run_refs.iter().any(|r| r.row_count > 0)
6101 || !self.retiring.is_empty()
6102 }
6103
6104 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
6109 if self.schema.columns.iter().any(|c| c.name == name) {
6110 return Err(MongrelError::Schema(format!(
6111 "column {name} already exists"
6112 )));
6113 }
6114 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
6115 self.schema.columns.push(ColumnDef {
6116 id,
6117 name: name.to_string(),
6118 ty,
6119 flags,
6120 });
6121 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6122 self.schema.validate_auto_increment()?;
6123 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
6124 self.auto_inc = resolve_auto_inc(&self.schema);
6125 }
6126 write_schema(&self.dir, &self.schema)?;
6127 self.clear_result_cache();
6128 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
6130 self.persist_manifest(self.current_epoch())?;
6131 Ok(id)
6132 }
6133
6134 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
6143 let cid = self
6144 .schema
6145 .columns
6146 .iter()
6147 .find(|c| c.name == column_name)
6148 .map(|c| c.id)
6149 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
6150 let ty = self
6151 .schema
6152 .columns
6153 .iter()
6154 .find(|c| c.id == cid)
6155 .map(|c| c.ty)
6156 .unwrap_or(TypeId::Int64);
6157 if !matches!(
6158 ty,
6159 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
6160 ) {
6161 return Err(MongrelError::Schema(format!(
6162 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
6163 )));
6164 }
6165 if self
6166 .schema
6167 .indexes
6168 .iter()
6169 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
6170 {
6171 return Ok(()); }
6173 self.schema.indexes.push(IndexDef {
6174 name: format!("{}_learned_range", column_name),
6175 column_id: cid,
6176 kind: IndexKind::LearnedRange,
6177 predicate: None,
6178 });
6179 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
6180 write_schema(&self.dir, &self.schema)?;
6181 self.build_learned_ranges()?;
6182 Ok(())
6183 }
6184
6185 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
6188 self.sync_byte_threshold = threshold;
6189 if let WalSink::Private(w) = &mut self.wal {
6190 w.set_sync_byte_threshold(threshold);
6191 }
6192 }
6193
6194 pub fn page_cache_flush(&self) {
6198 self.page_cache.flush_to_disk();
6199 }
6200
6201 pub fn page_cache_len(&self) -> usize {
6203 self.page_cache.len()
6204 }
6205
6206 pub fn decoded_cache_len(&self) -> usize {
6209 self.decoded_cache.len()
6210 }
6211
6212 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
6215 self.memtable.drain_sorted()
6216 }
6217
6218 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
6219 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
6220 }
6221
6222 pub(crate) fn table_dir(&self) -> &Path {
6223 &self.dir
6224 }
6225
6226 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
6227 &self.schema
6228 }
6229
6230 pub(crate) fn alloc_run_id(&mut self) -> u64 {
6231 let id = self.next_run_id;
6232 self.next_run_id += 1;
6233 id
6234 }
6235
6236 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
6237 self.run_refs.push(run_ref);
6238 }
6239
6240 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
6250 self.retiring.push(crate::manifest::RetiredRun {
6251 run_id,
6252 retire_epoch,
6253 });
6254 }
6255
6256 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
6260 if self.retiring.is_empty() {
6261 return Ok(0);
6262 }
6263 let mut reaped = 0;
6264 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
6265 for r in std::mem::take(&mut self.retiring) {
6271 if min_active.0 >= r.retire_epoch {
6272 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
6273 reaped += 1;
6274 } else {
6275 kept.push(r);
6276 }
6277 }
6278 self.retiring = kept;
6279 if reaped > 0 {
6280 self.persist_manifest(self.current_epoch())?;
6281 }
6282 Ok(reaped)
6283 }
6284
6285 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
6286 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
6287 return false;
6288 }
6289 self.live_count = self.live_count.saturating_add(run_ref.row_count);
6290 self.run_refs.push(run_ref);
6291 self.indexes_complete = false;
6292 true
6293 }
6294
6295 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
6296 self.kek.as_ref()
6297 }
6298
6299 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
6300 let mut reader = RunReader::open_with_cache(
6301 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
6302 self.schema.clone(),
6303 self.kek.clone(),
6304 Some(self.page_cache.clone()),
6305 Some(self.decoded_cache.clone()),
6306 self.table_id,
6307 Some(&self.verified_runs),
6308 )?;
6309 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
6313 reader.set_uniform_epoch(Epoch(rr.epoch_created));
6314 }
6315 Ok(reader)
6316 }
6317
6318 pub(crate) fn run_refs(&self) -> &[RunRef] {
6319 &self.run_refs
6320 }
6321
6322 pub(crate) fn runs_dir(&self) -> PathBuf {
6323 self.dir.join(RUNS_DIR)
6324 }
6325
6326 pub(crate) fn wal_dir(&self) -> PathBuf {
6327 self.dir.join(WAL_DIR)
6328 }
6329
6330 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
6331 self.run_refs = refs;
6332 }
6333
6334 pub(crate) fn next_run_id(&self) -> u64 {
6335 self.next_run_id
6336 }
6337
6338 pub(crate) fn compaction_zstd_level(&self) -> i32 {
6339 self.compaction_zstd_level
6340 }
6341
6342 pub(crate) fn bump_next_run_id(&mut self) {
6343 self.next_run_id += 1;
6344 }
6345
6346 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
6347 self.kek.clone()
6348 }
6349
6350 #[cfg(feature = "encryption")]
6354 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6355 self.kek.as_ref().map(|k| k.derive_idx_key())
6356 }
6357
6358 #[cfg(not(feature = "encryption"))]
6359 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
6360 None
6361 }
6362
6363 #[cfg(feature = "encryption")]
6367 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6368 self.kek.as_ref().map(|k| *k.derive_meta_key())
6369 }
6370
6371 #[cfg(not(feature = "encryption"))]
6372 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
6373 None
6374 }
6375
6376 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
6379 self.column_keys
6380 .iter()
6381 .map(|(&id, &(_, scheme))| (id, scheme))
6382 .collect()
6383 }
6384
6385 #[cfg(feature = "encryption")]
6390 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
6391 self.tokenize_value_enc(column_id, v)
6392 }
6393
6394 #[cfg(feature = "encryption")]
6395 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
6396 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6397 let (key, scheme) = self.column_keys.get(&column_id)?;
6398 let token: Vec<u8> = match (*scheme, v) {
6399 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
6400 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
6401 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
6402 _ => hmac_token(key, &v.encode_key()).to_vec(),
6403 };
6404 Some(Value::Bytes(token))
6405 }
6406
6407 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
6409 self.index_lookup_key_bytes(column_id, &v.encode_key())
6410 }
6411
6412 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
6415 #[cfg(feature = "encryption")]
6416 {
6417 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
6418 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
6419 if *scheme == SCHEME_HMAC_EQ {
6420 return hmac_token(key, encoded).to_vec();
6421 }
6422 }
6423 }
6424 let _ = column_id;
6425 encoded.to_vec()
6426 }
6427}
6428
6429fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
6430 let columnar::NativeColumn::Int64 { data, validity } = col else {
6431 return false;
6432 };
6433 if data.len() < n || !columnar::all_non_null(validity, n) {
6434 return false;
6435 }
6436 data.iter()
6437 .take(n)
6438 .zip(data.iter().skip(1))
6439 .all(|(a, b)| a < b)
6440}
6441
6442#[derive(Debug, Clone)]
6446pub struct ColumnStat {
6447 pub min: Option<Value>,
6448 pub max: Option<Value>,
6449 pub null_count: u64,
6450}
6451
6452#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6454pub enum NativeAgg {
6455 Count,
6456 Sum,
6457 Min,
6458 Max,
6459 Avg,
6460}
6461
6462#[derive(Debug, Clone, PartialEq)]
6464pub enum NativeAggResult {
6465 Count(u64),
6466 Int(i64),
6467 Float(f64),
6468 Null,
6470}
6471
6472#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6474pub enum ApproxAgg {
6475 Count,
6476 Sum,
6477 Avg,
6478}
6479
6480#[derive(Debug, Clone)]
6484pub struct ApproxResult {
6485 pub point: f64,
6487 pub ci_low: f64,
6489 pub ci_high: f64,
6491 pub n_population: u64,
6493 pub n_sample_live: usize,
6495 pub n_passing: usize,
6497}
6498
6499#[derive(Debug, Clone, PartialEq)]
6504pub enum AggState {
6505 Count(u64),
6507 SumI {
6509 sum: i128,
6510 count: u64,
6511 },
6512 SumF {
6514 sum: f64,
6515 count: u64,
6516 },
6517 AvgI {
6519 sum: i128,
6520 count: u64,
6521 },
6522 AvgF {
6524 sum: f64,
6525 count: u64,
6526 },
6527 MinI(i64),
6529 MaxI(i64),
6530 MinF(f64),
6532 MaxF(f64),
6533 Empty,
6535}
6536
6537impl AggState {
6538 pub fn merge(self, other: AggState) -> AggState {
6540 use AggState::*;
6541 match (self, other) {
6542 (Empty, x) | (x, Empty) => x,
6543 (Count(a), Count(b)) => Count(a + b),
6544 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6545 sum: sa + sb,
6546 count: ca + cb,
6547 },
6548 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6549 sum: sa + sb,
6550 count: ca + cb,
6551 },
6552 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6553 sum: sa + sb,
6554 count: ca + cb,
6555 },
6556 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6557 sum: sa + sb,
6558 count: ca + cb,
6559 },
6560 (MinI(a), MinI(b)) => MinI(a.min(b)),
6561 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6562 (MinF(a), MinF(b)) => MinF(a.min(b)),
6563 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6564 _ => Empty, }
6566 }
6567
6568 pub fn point(&self) -> Option<f64> {
6570 match self {
6571 AggState::Count(n) => Some(*n as f64),
6572 AggState::SumI { sum, .. } => Some(*sum as f64),
6573 AggState::SumF { sum, .. } => Some(*sum),
6574 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6575 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6576 AggState::MinI(n) => Some(*n as f64),
6577 AggState::MaxI(n) => Some(*n as f64),
6578 AggState::MinF(n) => Some(*n),
6579 AggState::MaxF(n) => Some(*n),
6580 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6581 }
6582 }
6583
6584 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6588 let is_float = matches!(ty, Some(TypeId::Float64));
6589 match (agg, result) {
6590 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6591 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6592 sum: x as i128,
6593 count: 1, },
6595 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6596 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6597 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6598 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6599 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6600 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6601 (NativeAgg::Count, _) => AggState::Empty,
6602 (_, NativeAggResult::Null) => AggState::Empty,
6603 _ => {
6604 let _ = is_float;
6605 AggState::Empty
6606 }
6607 }
6608 }
6609}
6610
6611#[derive(Debug, Clone)]
6614pub struct CachedAgg {
6615 pub state: AggState,
6616 pub watermark: u64,
6617 pub epoch: u64,
6618}
6619
6620#[derive(Debug, Clone)]
6622pub struct IncrementalAggResult {
6623 pub state: AggState,
6625 pub incremental: bool,
6628 pub delta_rows: u64,
6630}
6631
6632fn agg_state_from_rows(
6636 rows: &[Row],
6637 conditions: &[crate::query::Condition],
6638 index_sets: &[RowIdSet],
6639 column: Option<u16>,
6640 agg: NativeAgg,
6641 schema: &Schema,
6642) -> Result<AggState> {
6643 let mut count: u64 = 0;
6644 let mut sum_i: i128 = 0;
6645 let mut sum_f: f64 = 0.0;
6646 let mut mn_i: i64 = i64::MAX;
6647 let mut mx_i: i64 = i64::MIN;
6648 let mut mn_f: f64 = f64::INFINITY;
6649 let mut mx_f: f64 = f64::NEG_INFINITY;
6650 let mut saw_int = false;
6651 let mut saw_float = false;
6652 for r in rows {
6653 if !conditions
6654 .iter()
6655 .all(|c| condition_matches_row(c, r, schema))
6656 {
6657 continue;
6658 }
6659 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6660 continue;
6661 }
6662 match agg {
6663 NativeAgg::Count => match column {
6664 None => count += 1,
6666 Some(cid) => match r.columns.get(&cid) {
6669 None | Some(Value::Null) => {}
6670 Some(_) => count += 1,
6671 },
6672 },
6673 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6674 Some(Value::Int64(n)) => {
6675 count += 1;
6676 sum_i += *n as i128;
6677 mn_i = mn_i.min(*n);
6678 mx_i = mx_i.max(*n);
6679 saw_int = true;
6680 }
6681 Some(Value::Float64(f)) => {
6682 count += 1;
6683 sum_f += f;
6684 mn_f = mn_f.min(*f);
6685 mx_f = mx_f.max(*f);
6686 saw_float = true;
6687 }
6688 _ => {}
6689 },
6690 }
6691 }
6692 Ok(match agg {
6693 NativeAgg::Count => {
6694 if count == 0 {
6695 AggState::Empty
6696 } else {
6697 AggState::Count(count)
6698 }
6699 }
6700 NativeAgg::Sum => {
6701 if count == 0 {
6702 AggState::Empty
6703 } else if saw_int {
6704 AggState::SumI { sum: sum_i, count }
6705 } else {
6706 AggState::SumF { sum: sum_f, count }
6707 }
6708 }
6709 NativeAgg::Avg => {
6710 if count == 0 {
6711 AggState::Empty
6712 } else if saw_int {
6713 AggState::AvgI { sum: sum_i, count }
6714 } else {
6715 AggState::AvgF { sum: sum_f, count }
6716 }
6717 }
6718 NativeAgg::Min => {
6719 if !saw_int && !saw_float {
6720 AggState::Empty
6721 } else if saw_int {
6722 AggState::MinI(mn_i)
6723 } else {
6724 AggState::MinF(mn_f)
6725 }
6726 }
6727 NativeAgg::Max => {
6728 if !saw_int && !saw_float {
6729 AggState::Empty
6730 } else if saw_int {
6731 AggState::MaxI(mx_i)
6732 } else {
6733 AggState::MaxF(mx_f)
6734 }
6735 }
6736 })
6737}
6738
6739fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6743 use crate::query::Condition;
6744 match c {
6745 Condition::Pk(key) => match schema.primary_key() {
6746 Some(pk) => row
6747 .columns
6748 .get(&pk.id)
6749 .map(|v| v.encode_key() == *key)
6750 .unwrap_or(false),
6751 None => false,
6752 },
6753 Condition::BitmapEq { column_id, value } => row
6754 .columns
6755 .get(column_id)
6756 .map(|v| v.encode_key() == *value)
6757 .unwrap_or(false),
6758 Condition::BitmapIn { column_id, values } => {
6759 let key = row.columns.get(column_id).map(|v| v.encode_key());
6760 match key {
6761 Some(k) => values.contains(&k),
6762 None => false,
6763 }
6764 }
6765 Condition::BytesPrefix { column_id, prefix } => row
6766 .columns
6767 .get(column_id)
6768 .map(|v| v.encode_key().starts_with(prefix))
6769 .unwrap_or(false),
6770 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6771 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6772 _ => false,
6773 },
6774 Condition::RangeF64 {
6775 column_id,
6776 lo,
6777 lo_inclusive,
6778 hi,
6779 hi_inclusive,
6780 } => match row.columns.get(column_id) {
6781 Some(Value::Float64(n)) => {
6782 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6783 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6784 lo_ok && hi_ok
6785 }
6786 _ => false,
6787 },
6788 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6789 Some(Value::Bytes(b)) => {
6790 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6791 }
6792 _ => false,
6793 },
6794 Condition::FmContainsAll {
6795 column_id,
6796 patterns,
6797 } => match row.columns.get(column_id) {
6798 Some(Value::Bytes(b)) => patterns
6799 .iter()
6800 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6801 _ => false,
6802 },
6803 Condition::Ann { .. }
6804 | Condition::SparseMatch { .. }
6805 | Condition::MinHashSimilar { .. } => true,
6806 Condition::IsNull { column_id } => {
6807 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6808 }
6809 Condition::IsNotNull { column_id } => {
6810 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6811 }
6812 }
6813}
6814
6815fn as_f64(v: Option<&Value>) -> Option<f64> {
6817 match v {
6818 Some(Value::Int64(n)) => Some(*n as f64),
6819 Some(Value::Float64(f)) => Some(*f),
6820 _ => None,
6821 }
6822}
6823
6824fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6828 let mut count: u64 = 0;
6829 let mut sum: i128 = 0;
6830 let mut mn: i64 = i64::MAX;
6831 let mut mx: i64 = i64::MIN;
6832 while let Some(cols) = cursor.next_batch()? {
6833 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6834 if crate::columnar::all_non_null(validity, data.len()) {
6835 count += data.len() as u64;
6837 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6838 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6839 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6840 } else {
6841 for (i, &v) in data.iter().enumerate() {
6842 if crate::columnar::validity_bit(validity, i) {
6843 count += 1;
6844 sum += v as i128;
6845 mn = mn.min(v);
6846 mx = mx.max(v);
6847 }
6848 }
6849 }
6850 }
6851 }
6852 Ok((count, sum, mn, mx))
6853}
6854
6855fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6857 let mut count: u64 = 0;
6858 let mut sum: f64 = 0.0;
6859 let mut mn: f64 = f64::INFINITY;
6860 let mut mx: f64 = f64::NEG_INFINITY;
6861 while let Some(cols) = cursor.next_batch()? {
6862 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6863 if crate::columnar::all_non_null(validity, data.len()) {
6864 count += data.len() as u64;
6865 sum += data.iter().sum::<f64>();
6866 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6867 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6868 } else {
6869 for (i, &v) in data.iter().enumerate() {
6870 if crate::columnar::validity_bit(validity, i) {
6871 count += 1;
6872 sum += v;
6873 mn = mn.min(v);
6874 mx = mx.max(v);
6875 }
6876 }
6877 }
6878 }
6879 }
6880 Ok((count, sum, mn, mx))
6881}
6882
6883fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6884 if count == 0 && !matches!(agg, NativeAgg::Count) {
6885 return NativeAggResult::Null;
6886 }
6887 match agg {
6888 NativeAgg::Count => NativeAggResult::Count(count),
6889 NativeAgg::Sum => match sum.try_into() {
6892 Ok(v) => NativeAggResult::Int(v),
6893 Err(_) => NativeAggResult::Null,
6894 },
6895 NativeAgg::Min => NativeAggResult::Int(mn),
6896 NativeAgg::Max => NativeAggResult::Int(mx),
6897 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6898 }
6899}
6900
6901fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6902 if count == 0 && !matches!(agg, NativeAgg::Count) {
6903 return NativeAggResult::Null;
6904 }
6905 match agg {
6906 NativeAgg::Count => NativeAggResult::Count(count),
6907 NativeAgg::Sum => NativeAggResult::Float(sum),
6908 NativeAgg::Min => NativeAggResult::Float(mn),
6909 NativeAgg::Max => NativeAggResult::Float(mx),
6910 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6911 }
6912}
6913
6914fn agg_int(
6917 stats: &[crate::page::PageStat],
6918 decode: fn(Option<&[u8]>) -> Option<i64>,
6919) -> Option<(Option<i64>, Option<i64>, u64)> {
6920 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6921 let mut any = false;
6922 for s in stats {
6923 if let Some(v) = decode(s.min.as_deref()) {
6924 mn = mn.min(v);
6925 any = true;
6926 }
6927 if let Some(v) = decode(s.max.as_deref()) {
6928 mx = mx.max(v);
6929 any = true;
6930 }
6931 nulls += s.null_count;
6932 }
6933 any.then_some((Some(mn), Some(mx), nulls))
6934}
6935
6936fn agg_float(
6938 stats: &[crate::page::PageStat],
6939 decode: fn(Option<&[u8]>) -> Option<f64>,
6940) -> Option<(Option<f64>, Option<f64>, u64)> {
6941 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6942 let mut any = false;
6943 for s in stats {
6944 if let Some(v) = decode(s.min.as_deref()) {
6945 mn = mn.min(v);
6946 any = true;
6947 }
6948 if let Some(v) = decode(s.max.as_deref()) {
6949 mx = mx.max(v);
6950 any = true;
6951 }
6952 nulls += s.null_count;
6953 }
6954 any.then_some((Some(mn), Some(mx), nulls))
6955}
6956
6957type SecondaryIndexes = (
6959 HashMap<u16, BitmapIndex>,
6960 HashMap<u16, AnnIndex>,
6961 HashMap<u16, FmIndex>,
6962 HashMap<u16, SparseIndex>,
6963 HashMap<u16, MinHashIndex>,
6964);
6965
6966fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6967 let mut bitmap = HashMap::new();
6968 let mut ann = HashMap::new();
6969 let mut fm = HashMap::new();
6970 let mut sparse = HashMap::new();
6971 let mut minhash = HashMap::new();
6972 for idef in &schema.indexes {
6973 match idef.kind {
6974 IndexKind::Bitmap => {
6975 bitmap.insert(idef.column_id, BitmapIndex::new());
6976 }
6977 IndexKind::Ann => {
6978 let dim = schema
6979 .columns
6980 .iter()
6981 .find(|c| c.id == idef.column_id)
6982 .and_then(|c| match c.ty {
6983 TypeId::Embedding { dim } => Some(dim as usize),
6984 _ => None,
6985 })
6986 .unwrap_or(0);
6987 ann.insert(idef.column_id, AnnIndex::new(dim));
6988 }
6989 IndexKind::FmIndex => {
6990 fm.insert(idef.column_id, FmIndex::new());
6991 }
6992 IndexKind::Sparse => {
6993 sparse.insert(idef.column_id, SparseIndex::new());
6994 }
6995 IndexKind::MinHash => {
6996 minhash.insert(idef.column_id, MinHashIndex::new());
6997 }
6998 _ => {}
6999 }
7000 }
7001 (bitmap, ann, fm, sparse, minhash)
7002}
7003
7004const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
7005 | ColumnFlags::AUTO_INCREMENT
7006 | ColumnFlags::ENCRYPTED
7007 | ColumnFlags::ENCRYPTED_INDEXABLE
7008 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
7009
7010fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
7011 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
7012 return Err(MongrelError::Schema(
7013 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
7014 ));
7015 }
7016 Ok(())
7017}
7018
7019fn validate_alter_column_type(
7020 schema: &Schema,
7021 old: &ColumnDef,
7022 next: &ColumnDef,
7023 has_stored_versions: bool,
7024) -> Result<()> {
7025 if old.ty == next.ty {
7026 return Ok(());
7027 }
7028 if schema.indexes.iter().any(|i| i.column_id == old.id) {
7029 return Err(MongrelError::Schema(format!(
7030 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
7031 old.name
7032 )));
7033 }
7034 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
7035 return Ok(());
7036 }
7037 Err(MongrelError::Schema(format!(
7038 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
7039 old.ty, next.ty
7040 )))
7041}
7042
7043fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
7044 matches!(
7045 (old, new),
7046 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
7047 )
7048}
7049
7050fn rows_pk_strictly_increasing(rows: &[Row], pk_id: u16) -> bool {
7056 let mut prev: Option<i64> = None;
7057 for r in rows {
7058 match r.columns.get(&pk_id) {
7059 Some(Value::Int64(v)) => {
7060 if prev.is_some_and(|p| p >= *v) {
7061 return false;
7062 }
7063 prev = Some(*v);
7064 }
7065 _ => return false,
7066 }
7067 }
7068 true
7069}
7070
7071#[allow(clippy::too_many_arguments)]
7072fn index_into(
7073 schema: &Schema,
7074 row: &Row,
7075 hot: &mut HotIndex,
7076 bitmap: &mut HashMap<u16, BitmapIndex>,
7077 ann: &mut HashMap<u16, AnnIndex>,
7078 fm: &mut HashMap<u16, FmIndex>,
7079 sparse: &mut HashMap<u16, SparseIndex>,
7080 minhash: &mut HashMap<u16, MinHashIndex>,
7081) {
7082 for idef in &schema.indexes {
7083 let Some(val) = row.columns.get(&idef.column_id) else {
7084 continue;
7085 };
7086 match idef.kind {
7087 IndexKind::Bitmap => {
7088 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7089 b.insert(val.encode_key(), row.row_id);
7090 }
7091 }
7092 IndexKind::Ann => {
7093 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7094 a.insert(v, row.row_id);
7095 }
7096 }
7097 IndexKind::FmIndex => {
7098 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7099 f.insert(b.clone(), row.row_id);
7100 }
7101 }
7102 IndexKind::Sparse => {
7103 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7104 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7107 s.insert(&terms, row.row_id);
7108 }
7109 }
7110 }
7111 IndexKind::MinHash => {
7112 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7113 let tokens = crate::index::token_hashes_from_bytes(b);
7116 mh.insert(&tokens, row.row_id);
7117 }
7118 }
7119 _ => {}
7120 }
7121 }
7122 if let Some(pk_col) = schema.primary_key() {
7123 if let Some(pk_val) = row.columns.get(&pk_col.id) {
7124 hot.insert(pk_val.encode_key(), row.row_id);
7125 }
7126 }
7127}
7128
7129#[allow(clippy::too_many_arguments)]
7132fn index_into_single(
7133 idef: &IndexDef,
7134 _schema: &Schema,
7135 row: &Row,
7136 _hot: &mut HotIndex,
7137 bitmap: &mut HashMap<u16, BitmapIndex>,
7138 ann: &mut HashMap<u16, AnnIndex>,
7139 fm: &mut HashMap<u16, FmIndex>,
7140 sparse: &mut HashMap<u16, SparseIndex>,
7141 minhash: &mut HashMap<u16, MinHashIndex>,
7142) {
7143 let Some(val) = row.columns.get(&idef.column_id) else {
7144 return;
7145 };
7146 match idef.kind {
7147 IndexKind::Bitmap => {
7148 if let Some(b) = bitmap.get_mut(&idef.column_id) {
7149 b.insert(val.encode_key(), row.row_id);
7150 }
7151 }
7152 IndexKind::Ann => {
7153 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
7154 a.insert(v, row.row_id);
7155 }
7156 }
7157 IndexKind::FmIndex => {
7158 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
7159 f.insert(b.clone(), row.row_id);
7160 }
7161 }
7162 IndexKind::Sparse => {
7163 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
7164 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
7165 s.insert(&terms, row.row_id);
7166 }
7167 }
7168 }
7169 IndexKind::MinHash => {
7170 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
7171 let tokens = crate::index::token_hashes_from_bytes(b);
7172 mh.insert(&tokens, row.row_id);
7173 }
7174 }
7175 _ => {}
7176 }
7177}
7178
7179fn eval_partial_predicate(
7185 pred: &str,
7186 columns_map: &HashMap<u16, &Value>,
7187 name_to_id: &HashMap<&str, u16>,
7188) -> bool {
7189 let lower = pred.trim().to_ascii_lowercase();
7190 if let Some(rest) = lower.strip_suffix(" is not null") {
7192 let col_name = rest.trim();
7193 if let Some(col_id) = name_to_id.get(col_name) {
7194 return columns_map
7195 .get(col_id)
7196 .is_some_and(|v| !matches!(v, Value::Null));
7197 }
7198 }
7199 if let Some(rest) = lower.strip_suffix(" is null") {
7201 let col_name = rest.trim();
7202 if let Some(col_id) = name_to_id.get(col_name) {
7203 return columns_map
7204 .get(col_id)
7205 .map_or(true, |v| matches!(v, Value::Null));
7206 }
7207 }
7208 true
7211}
7212
7213#[allow(dead_code)]
7219fn bulk_index_key(
7220 column_keys: &HashMap<u16, ([u8; 32], u8)>,
7221 column_id: u16,
7222 ty: TypeId,
7223 col: &columnar::NativeColumn,
7224 i: usize,
7225) -> Option<Vec<u8>> {
7226 let encoded = columnar::encode_key_native(ty, col, i)?;
7227 #[cfg(feature = "encryption")]
7228 {
7229 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
7230 if let Some((key, scheme)) = column_keys.get(&column_id) {
7231 return Some(match (*scheme, col) {
7232 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
7233 (_, columnar::NativeColumn::Int64 { data, .. }) => {
7234 ope_token_i64(key, data[i]).to_vec()
7235 }
7236 (_, columnar::NativeColumn::Float64 { data, .. }) => {
7237 ope_token_f64(key, data[i]).to_vec()
7238 }
7239 _ => hmac_token(key, &encoded).to_vec(),
7240 });
7241 }
7242 }
7243 #[cfg(not(feature = "encryption"))]
7244 {
7245 let _ = (column_id, column_keys, col);
7246 }
7247 Some(encoded)
7248}
7249
7250pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
7251 let json = serde_json::to_string_pretty(schema)
7252 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
7253 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
7254 Ok(())
7255}
7256
7257fn read_schema(dir: &Path) -> Result<Schema> {
7258 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
7259 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
7260}
7261
7262fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
7263 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
7264}
7265
7266fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
7267 let n = list_wal_numbers(wal_dir)?;
7268 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
7269}
7270
7271fn next_wal_number(wal_dir: &Path) -> Result<u32> {
7272 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
7273}
7274
7275fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
7276 let _ = std::fs::create_dir_all(wal_dir);
7277 let mut max_n = None;
7278 for entry in std::fs::read_dir(wal_dir)? {
7279 let entry = entry?;
7280 let fname = entry.file_name();
7281 let Some(s) = fname.to_str() else {
7282 continue;
7283 };
7284 let Some(stripped) = s.strip_prefix("seg-") else {
7285 continue;
7286 };
7287 let Some(stripped) = stripped.strip_suffix(".wal") else {
7288 continue;
7289 };
7290 if let Ok(n) = stripped.parse::<u32>() {
7291 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
7292 }
7293 }
7294 Ok(max_n)
7295}