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
83pub struct Table {
85 dir: PathBuf,
86 table_id: u64,
87 wal: WalSink,
88 memtable: Memtable,
89 mutable_run: MutableRun,
94 mutable_run_spill_bytes: u64,
96 compaction_zstd_level: i32,
99 allocator: RowIdAllocator,
100 epoch: Arc<EpochAuthority>,
101 persisted_epoch: u64,
104 schema: Schema,
105 hot: HotIndex,
106 kek: Option<Arc<Kek>>,
109 column_keys: HashMap<u16, ([u8; 32], u8)>,
113 run_refs: Vec<RunRef>,
114 retiring: Vec<crate::manifest::RetiredRun>,
117 next_run_id: u64,
118 sync_byte_threshold: u64,
119 current_txn_id: u64,
124 bitmap: HashMap<u16, BitmapIndex>,
125 ann: HashMap<u16, AnnIndex>,
126 fm: HashMap<u16, FmIndex>,
127 sparse: HashMap<u16, SparseIndex>,
128 minhash: HashMap<u16, MinHashIndex>,
129 learned_range: HashMap<u16, ColumnLearnedRange>,
132 pk_by_row: HashMap<RowId, Vec<u8>>,
134 pinned: BTreeMap<Epoch, usize>,
137 pub(crate) live_count: u64,
140 reservoir: crate::reservoir::Reservoir,
143 had_deletes: bool,
147 agg_cache: HashMap<u64, CachedAgg>,
151 global_idx_epoch: u64,
155 indexes_complete: bool,
160 flushed_epoch: u64,
163 page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
166 snapshots: Arc<crate::retention::SnapshotRegistry>,
169 commit_lock: Arc<parking_lot::Mutex<()>>,
171 decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
174 result_cache: Arc<parking_lot::Mutex<ResultCache>>,
183 wal_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
185 pending_delete_rids: roaring::RoaringBitmap,
188 pending_put_cols: std::collections::HashSet<u16>,
191 pending_rows: Vec<Row>,
197 pending_rows_auto_inc: Vec<bool>,
198 pending_dels: Vec<RowId>,
201 pending_truncate: Option<Epoch>,
205 auto_inc: Option<AutoIncState>,
208}
209
210const _: () = {
217 const fn assert_sync<T: ?Sized + Sync>() {}
218 assert_sync::<Table>();
219};
220
221enum CachedData {
227 Rows(Arc<Vec<Row>>),
228 Columns(Arc<Vec<(u16, columnar::NativeColumn)>>),
229}
230
231impl CachedData {
232 fn approx_bytes(&self) -> u64 {
233 match self {
234 CachedData::Rows(r) => r.iter().map(|r| r.estimated_bytes()).sum::<u64>(),
235 CachedData::Columns(c) => c
236 .iter()
237 .map(|(_, c)| c.approx_bytes())
238 .sum::<u64>()
239 .saturating_add(c.len() as u64 * 16),
240 }
241 }
242}
243
244struct CachedEntry {
248 data: CachedData,
249 footprint: roaring::RoaringBitmap,
250 condition_cols: Vec<u16>,
251}
252
253struct ResultCache {
264 entries: std::collections::HashMap<u64, CachedEntry>,
265 order: std::collections::VecDeque<u64>,
266 bytes: u64,
267 max_bytes: u64,
268 dir: Option<std::path::PathBuf>,
269 #[allow(dead_code)]
270 cache_dek: Option<Zeroizing<[u8; DEK_LEN]>>,
271}
272
273#[derive(serde::Serialize, serde::Deserialize)]
275struct SerializedEntry {
276 condition_cols: Vec<u16>,
277 footprint_bits: Vec<u32>,
278 data: SerializedData,
279}
280
281#[derive(serde::Serialize, serde::Deserialize)]
282enum SerializedData {
283 Rows(Vec<Row>),
284 Columns(Vec<(u16, columnar::NativeColumn)>),
285}
286
287impl SerializedEntry {
288 fn from_entry(entry: &CachedEntry) -> Self {
289 let footprint_bits: Vec<u32> = entry.footprint.iter().collect();
290 let data = match &entry.data {
291 CachedData::Rows(r) => SerializedData::Rows((**r).clone()),
292 CachedData::Columns(c) => SerializedData::Columns((**c).clone()),
293 };
294 Self {
295 condition_cols: entry.condition_cols.clone(),
296 footprint_bits,
297 data,
298 }
299 }
300
301 fn into_entry(self) -> Option<CachedEntry> {
302 let footprint: roaring::RoaringBitmap = self.footprint_bits.into_iter().collect();
303 let data = match self.data {
304 SerializedData::Rows(r) => CachedData::Rows(Arc::new(r)),
305 SerializedData::Columns(c) => {
306 if !c.iter().all(|(_, col)| col.validate()) {
309 return None;
310 }
311 CachedData::Columns(Arc::new(c))
312 }
313 };
314 Some(CachedEntry {
315 data,
316 footprint,
317 condition_cols: self.condition_cols,
318 })
319 }
320}
321
322impl ResultCache {
323 const DEFAULT_MAX_BYTES: u64 = 256 * 1024 * 1024;
324
325 fn new() -> Self {
326 Self::with_max_bytes(Self::DEFAULT_MAX_BYTES)
327 }
328
329 fn with_max_bytes(max_bytes: u64) -> Self {
330 Self {
331 entries: std::collections::HashMap::new(),
332 order: std::collections::VecDeque::new(),
333 bytes: 0,
334 max_bytes,
335 dir: None,
336 cache_dek: None,
337 }
338 }
339
340 fn with_dir(mut self, dir: std::path::PathBuf) -> Self {
341 let _ = std::fs::create_dir_all(&dir);
342 self.dir = Some(dir);
343 self
344 }
345
346 fn with_cache_dek(mut self, dek: Option<Zeroizing<[u8; DEK_LEN]>>) -> Self {
347 self.cache_dek = dek;
348 self
349 }
350
351 fn disk_path(&self, key: u64) -> Option<std::path::PathBuf> {
352 self.dir.as_ref().map(|d| d.join(format!("{key:016x}.bin")))
353 }
354
355 fn store_to_disk(&self, key: u64, entry: &CachedEntry) {
359 let Some(path) = self.disk_path(key) else {
360 return;
361 };
362 let serialized = match bincode::serialize(&SerializedEntry::from_entry(entry)) {
363 Ok(s) => s,
364 Err(_) => return,
365 };
366 let on_disk = if let Some(dek) = &self.cache_dek {
368 match self.encrypt_cache(&serialized, dek) {
369 Some(b) => b,
370 None => return,
371 }
372 } else {
373 serialized
374 };
375 let tmp = path.with_extension("tmp");
376 use std::io::Write;
377 let write = || -> std::io::Result<()> {
378 let mut f = std::fs::File::create(&tmp)?;
379 f.write_all(&on_disk)?;
380 f.flush()?;
381 Ok(())
382 };
383 if write().is_err() {
384 let _ = std::fs::remove_file(&tmp);
385 return;
386 }
387 let _ = std::fs::rename(&tmp, &path);
388 }
389
390 fn load_from_disk(&self, key: u64) -> Option<CachedEntry> {
392 let path = self.disk_path(key)?;
393 let bytes = std::fs::read(&path).ok()?;
394 let plaintext = if let Some(dek) = &self.cache_dek {
395 self.decrypt_cache(&bytes, dek)?
396 } else {
397 bytes
398 };
399 let serialized: SerializedEntry = bincode::deserialize(&plaintext).ok()?;
400 serialized.into_entry()
401 }
402
403 fn remove_from_disk(&self, key: u64) {
405 if let Some(path) = self.disk_path(key) {
406 let _ = std::fs::remove_file(&path);
407 }
408 }
409
410 #[cfg(feature = "encryption")]
412 fn encrypt_cache(&self, plaintext: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
413 use crate::encryption::Cipher;
414 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
415 let mut nonce = [0u8; 12];
416 crate::encryption::fill_random(&mut nonce);
417 let ct = cipher.encrypt_page(&nonce, plaintext).ok()?;
418 let mut out = Vec::with_capacity(12 + ct.len());
419 out.extend_from_slice(&nonce);
420 out.extend_from_slice(&ct);
421 Some(out)
422 }
423
424 #[cfg(not(feature = "encryption"))]
425 fn encrypt_cache(&self, _plaintext: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
426 None
427 }
428
429 #[cfg(feature = "encryption")]
431 fn decrypt_cache(&self, bytes: &[u8], dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
432 use crate::encryption::Cipher;
433 if bytes.len() < 28 {
434 return None;
435 }
436 let cipher = crate::encryption::AesCipher::new(&dek[..]).ok()?;
437 let nonce: [u8; 12] = bytes[..12].try_into().ok()?;
438 let ct = &bytes[12..];
439 cipher.decrypt_page(&nonce, ct).ok()
440 }
441
442 #[cfg(not(feature = "encryption"))]
443 fn decrypt_cache(&self, _bytes: &[u8], _dek: &Zeroizing<[u8; DEK_LEN]>) -> Option<Vec<u8>> {
444 None
445 }
446
447 fn load_persistent(&mut self) {
450 let Some(dir) = self.dir.as_ref().cloned() else {
451 return;
452 };
453 let entries = match std::fs::read_dir(&dir) {
454 Ok(e) => e,
455 Err(_) => return,
456 };
457 for entry in entries.flatten() {
458 let path = entry.path();
459 if path.extension().and_then(|e| e.to_str()) == Some("tmp") {
461 let _ = std::fs::remove_file(&path);
462 continue;
463 }
464 if path.extension().and_then(|e| e.to_str()) != Some("bin") {
465 continue;
466 }
467 let stem = match path.file_stem().and_then(|s| s.to_str()) {
468 Some(s) => s,
469 None => continue,
470 };
471 let key = match u64::from_str_radix(stem, 16) {
472 Ok(k) => k,
473 Err(_) => continue,
474 };
475 let bytes = match std::fs::read(&path) {
476 Ok(b) => b,
477 Err(_) => continue,
478 };
479 let plaintext = if let Some(dek) = &self.cache_dek {
481 match self.decrypt_cache(&bytes, dek) {
482 Some(p) => p,
483 None => {
484 let _ = std::fs::remove_file(&path);
485 continue;
486 }
487 }
488 } else {
489 bytes
490 };
491 match bincode::deserialize::<SerializedEntry>(&plaintext) {
492 Ok(serialized) => {
493 if let Some(entry) = serialized.into_entry() {
494 self.bytes = self.bytes.saturating_add(entry.data.approx_bytes());
495 self.entries.insert(key, entry);
496 self.order.push_back(key);
497 } else {
498 let _ = std::fs::remove_file(&path);
499 }
500 }
501 Err(_) => {
502 let _ = std::fs::remove_file(&path);
503 }
504 }
505 }
506 self.evict();
507 }
508
509 fn set_max_bytes(&mut self, max_bytes: u64) {
510 self.max_bytes = max_bytes;
511 self.evict();
512 }
513
514 fn touch(&mut self, key: u64) {
516 self.order.retain(|k| *k != key);
517 self.order.push_back(key);
518 }
519
520 fn get_rows(&mut self, key: u64) -> Option<Arc<Vec<Row>>> {
521 let res = self.entries.get(&key).and_then(|e| match &e.data {
522 CachedData::Rows(r) => Some(r.clone()),
523 CachedData::Columns(_) => None,
524 });
525 if res.is_some() {
526 self.touch(key);
527 return res;
528 }
529 if let Some(entry) = self.load_from_disk(key) {
531 let res = match &entry.data {
532 CachedData::Rows(r) => Some(r.clone()),
533 CachedData::Columns(_) => None,
534 };
535 if res.is_some() {
536 let approx = entry.data.approx_bytes();
537 self.bytes = self.bytes.saturating_add(approx);
538 self.entries.insert(key, entry);
539 self.order.push_back(key);
540 self.evict();
541 return res;
542 }
543 }
544 None
545 }
546
547 fn get_columns(&mut self, key: u64) -> Option<Arc<Vec<(u16, columnar::NativeColumn)>>> {
548 let res = self.entries.get(&key).and_then(|e| match &e.data {
549 CachedData::Columns(c) => Some(c.clone()),
550 CachedData::Rows(_) => None,
551 });
552 if res.is_some() {
553 self.touch(key);
554 return res;
555 }
556 if let Some(entry) = self.load_from_disk(key) {
558 let res = match &entry.data {
559 CachedData::Columns(c) => Some(c.clone()),
560 CachedData::Rows(_) => None,
561 };
562 if res.is_some() {
563 let approx = entry.data.approx_bytes();
564 self.bytes = self.bytes.saturating_add(approx);
565 self.entries.insert(key, entry);
566 self.order.push_back(key);
567 self.evict();
568 return res;
569 }
570 }
571 None
572 }
573
574 fn insert(&mut self, key: u64, entry: CachedEntry) {
575 let approx = entry.data.approx_bytes();
576 if self.entries.remove(&key).is_some() {
577 self.order.retain(|k| *k != key);
578 self.bytes = self.entries.values().map(|e| e.data.approx_bytes()).sum();
579 }
580 self.store_to_disk(key, &entry);
582 self.bytes = self.bytes.saturating_add(approx);
583 self.entries.insert(key, entry);
584 self.order.push_back(key);
585 self.evict();
586 }
587
588 fn invalidate(
597 &mut self,
598 delete_rids: &roaring::RoaringBitmap,
599 put_cols: &std::collections::HashSet<u16>,
600 ) {
601 if self.entries.is_empty() {
602 return;
603 }
604 let has_deletes = !delete_rids.is_empty();
605 let to_remove: std::collections::HashSet<u64> = self
606 .entries
607 .iter()
608 .filter(|(_, e)| {
609 let delete_hit = if e.footprint.is_empty() {
610 has_deletes
611 } else {
612 e.footprint.intersection_len(delete_rids) > 0
613 };
614 let col_hit = e.condition_cols.iter().any(|c| put_cols.contains(c));
615 delete_hit || col_hit
616 })
617 .map(|(&k, _)| k)
618 .collect();
619 for key in &to_remove {
620 if let Some(e) = self.entries.remove(key) {
621 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
622 }
623 self.remove_from_disk(*key);
624 }
625 if !to_remove.is_empty() {
626 self.order.retain(|k| !to_remove.contains(k));
627 }
628 }
629
630 fn clear(&mut self) {
631 if let Some(dir) = &self.dir {
633 if let Ok(entries) = std::fs::read_dir(dir) {
634 for entry in entries.flatten() {
635 let path = entry.path();
636 if path.extension().and_then(|e| e.to_str()) == Some("bin") {
637 let _ = std::fs::remove_file(&path);
638 }
639 }
640 }
641 }
642 self.entries.clear();
643 self.order.clear();
644 self.bytes = 0;
645 }
646
647 fn evict(&mut self) {
648 while self.bytes > self.max_bytes {
649 let Some(k) = self.order.pop_front() else {
650 break;
651 };
652 if let Some(e) = self.entries.remove(&k) {
653 self.bytes = self.bytes.saturating_sub(e.data.approx_bytes());
654 self.remove_from_disk(k);
658 }
659 }
660 }
661}
662
663type DekaOpt = Option<Zeroizing<[u8; DEK_LEN]>>;
670
671fn derive_subkeys(kek: Option<&Kek>, _table_id: u64) -> (DekaOpt, DekaOpt) {
672 let _ = kek;
673 #[cfg(feature = "encryption")]
674 {
675 if let Some(k) = kek {
676 return (
677 Some(k.derive_table_wal_key(_table_id)),
678 Some(k.derive_cache_key()),
679 );
680 }
681 }
682 (None, None)
683}
684
685#[cfg(feature = "encryption")]
687fn make_cipher(dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
688 Box::new(crate::encryption::AesCipher::new(&dek[..]).expect("DEK is 32 bytes"))
689}
690
691#[cfg(not(feature = "encryption"))]
692fn make_cipher(_dek: &Zeroizing<[u8; DEK_LEN]>) -> Box<dyn crate::encryption::Cipher> {
693 Box::new(crate::encryption::PlaintextCipher)
694}
695
696fn build_column_keys(kek: Option<&Kek>, schema: &Schema) -> HashMap<u16, ([u8; 32], u8)> {
697 let Some(kek) = kek else {
698 return HashMap::new();
699 };
700 #[cfg(feature = "encryption")]
701 {
702 use crate::encryption::{SCHEME_HMAC_EQ, SCHEME_OPE_RANGE};
703 schema
704 .columns
705 .iter()
706 .filter(|c| c.flags.contains(ColumnFlags::ENCRYPTED_INDEXABLE))
707 .map(|c| {
708 let scheme = if schema
709 .indexes
710 .iter()
711 .any(|i| i.column_id == c.id && i.kind == IndexKind::LearnedRange)
712 {
713 SCHEME_OPE_RANGE
714 } else {
715 SCHEME_HMAC_EQ
716 };
717 let key: [u8; 32] = *kek.derive_column_key(c.id);
718 (c.id, (key, scheme))
719 })
720 .collect()
721 }
722 #[cfg(not(feature = "encryption"))]
723 {
724 let _ = (kek, schema);
725 HashMap::new()
726 }
727}
728
729pub(crate) struct SharedCtx {
734 pub epoch: Arc<EpochAuthority>,
735 pub page_cache: Arc<parking_lot::Mutex<crate::cache::PageCache>>,
736 pub decoded_cache: Arc<parking_lot::Mutex<crate::cache::DecodedPageCache>>,
737 pub snapshots: Arc<crate::retention::SnapshotRegistry>,
738 pub kek: Option<Arc<Kek>>,
739 pub commit_lock: Arc<parking_lot::Mutex<()>>,
745 pub shared: Option<SharedWalCtx>,
749}
750
751#[derive(Clone)]
757pub(crate) struct SharedWalCtx {
758 pub wal: Arc<parking_lot::Mutex<SharedWal>>,
759 pub group: Arc<GroupCommit>,
760 pub poisoned: Arc<AtomicBool>,
761 pub txn_ids: Arc<parking_lot::Mutex<u64>>,
762}
763
764enum WalSink {
767 Private(Wal),
768 Shared(SharedWalCtx),
769}
770
771impl SharedCtx {
772 pub(crate) fn new(kek: Option<Arc<Kek>>, cache_dir: Option<PathBuf>) -> Self {
776 let mut cache = crate::cache::PageCache::new(PAGE_CACHE_CAPACITY);
777 if let Some(d) = cache_dir {
778 cache = cache.with_persistence(d);
779 }
780 Self {
781 epoch: Arc::new(EpochAuthority::new(0)),
782 page_cache: Arc::new(parking_lot::Mutex::new(cache)),
783 decoded_cache: Arc::new(parking_lot::Mutex::new(
784 crate::cache::DecodedPageCache::new(DECODED_CACHE_CAPACITY),
785 )),
786 snapshots: Arc::new(crate::retention::SnapshotRegistry::new()),
787 kek,
788 commit_lock: Arc::new(parking_lot::Mutex::new(())),
789 shared: None,
790 }
791 }
792}
793
794impl Table {
795 pub fn create(dir: impl AsRef<Path>, schema: Schema, table_id: u64) -> Result<Self> {
796 let dir = dir.as_ref().to_path_buf();
797 let ctx = SharedCtx::new(None, Some(dir.join(CACHE_DIR)));
798 Self::create_in(&dir, schema, table_id, ctx)
799 }
800
801 #[cfg(feature = "encryption")]
812 pub fn create_encrypted(
813 dir: impl AsRef<Path>,
814 schema: Schema,
815 table_id: u64,
816 passphrase: &str,
817 ) -> Result<Self> {
818 let dir = dir.as_ref();
819 std::fs::create_dir_all(dir.join(META_DIR))?;
820 let salt = crate::encryption::random_salt();
821 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
822 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
823 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
824 Self::create_in(dir, schema, table_id, ctx)
825 }
826
827 #[cfg(feature = "encryption")]
832 pub fn create_with_key(
833 dir: impl AsRef<Path>,
834 schema: Schema,
835 table_id: u64,
836 key: &[u8],
837 ) -> Result<Self> {
838 let dir = dir.as_ref();
839 std::fs::create_dir_all(dir.join(META_DIR))?;
840 let salt = crate::encryption::random_salt();
841 std::fs::write(dir.join(META_DIR).join(KEYS_FILENAME), salt)?;
842 let kek: Arc<Kek> = Arc::new(Kek::from_raw_key(key, &salt)?);
843 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
844 Self::create_in(dir, schema, table_id, ctx)
845 }
846
847 #[cfg(feature = "encryption")]
849 pub fn open_with_key(dir: impl AsRef<Path>, key: &[u8]) -> Result<Self> {
850 let dir = dir.as_ref();
851 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
852 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
853 MongrelError::NotFound(format!(
854 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
855 salt_path
856 ))
857 })?;
858 if salt_bytes.len() != crate::encryption::SALT_LEN {
859 return Err(MongrelError::InvalidArgument(format!(
860 "salt file is {} bytes, expected {}",
861 salt_bytes.len(),
862 crate::encryption::SALT_LEN
863 )));
864 }
865 let mut salt = [0u8; crate::encryption::SALT_LEN];
866 salt.copy_from_slice(&salt_bytes);
867 let kek = Arc::new(Kek::from_raw_key(key, &salt)?);
868 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
869 Self::open_in(dir, ctx)
870 }
871
872 pub(crate) fn create_in(
873 dir: impl AsRef<Path>,
874 schema: Schema,
875 table_id: u64,
876 ctx: SharedCtx,
877 ) -> Result<Self> {
878 schema.validate_auto_increment()?;
879 let dir = dir.as_ref().to_path_buf();
880 std::fs::create_dir_all(dir.join(RUNS_DIR))?;
881 write_schema(&dir, &schema)?;
882 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), table_id);
883 let (wal, current_txn_id) = match ctx.shared.clone() {
886 Some(s) => (WalSink::Shared(s), 0),
887 None => {
888 std::fs::create_dir_all(dir.join(WAL_DIR))?;
889 let mut w = if let Some(ref dk) = wal_dek {
890 Wal::create_with_cipher(
891 dir.join(WAL_DIR).join("seg-000000.wal"),
892 Epoch(0),
893 Some(make_cipher(dk)),
894 0,
895 )?
896 } else {
897 Wal::create(dir.join(WAL_DIR).join("seg-000000.wal"), Epoch(0))?
898 };
899 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
900 (WalSink::Private(w), 1)
901 }
902 };
903 let mut manifest = Manifest::new(table_id, schema.schema_id);
904 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
909 manifest::write_atomic(&dir, &mut manifest, manifest_meta_dek.as_ref())?;
910 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&schema);
911 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
912 let auto_inc = resolve_auto_inc(&schema);
913 let rcache_dir = dir.join(RCACHE_DIR);
914 Ok(Self {
915 dir,
916 table_id,
917 wal,
918 memtable: Memtable::new(),
919 mutable_run: MutableRun::new(),
920 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
921 compaction_zstd_level: 3,
922 allocator: RowIdAllocator::new(0),
923 epoch: ctx.epoch,
924 persisted_epoch: 0,
925 schema,
926 hot: HotIndex::new(),
927 kek: ctx.kek,
928 column_keys,
929 run_refs: Vec::new(),
930 retiring: Vec::new(),
931 next_run_id: 1,
932 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
933 current_txn_id,
934 bitmap,
935 ann,
936 fm,
937 sparse,
938 minhash,
939 learned_range: HashMap::new(),
940 pk_by_row: HashMap::new(),
941 pinned: BTreeMap::new(),
942 live_count: 0,
943 reservoir: crate::reservoir::Reservoir::default(),
944 had_deletes: false,
945 agg_cache: HashMap::new(),
946 global_idx_epoch: 0,
947 indexes_complete: true,
948 flushed_epoch: 0,
949 page_cache: ctx.page_cache,
950 decoded_cache: ctx.decoded_cache,
951 snapshots: ctx.snapshots,
952 commit_lock: ctx.commit_lock,
953 result_cache: Arc::new(parking_lot::Mutex::new(
954 ResultCache::new()
955 .with_dir(rcache_dir)
956 .with_cache_dek(cache_dek.clone()),
957 )),
958 pending_delete_rids: roaring::RoaringBitmap::new(),
959 pending_put_cols: std::collections::HashSet::new(),
960 pending_rows: Vec::new(),
961 pending_rows_auto_inc: Vec::new(),
962 pending_dels: Vec::new(),
963 pending_truncate: None,
964 wal_dek,
965 auto_inc,
966 })
967 }
968
969 pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
973 let dir = dir.as_ref();
974 let ctx = SharedCtx::new(None, Some(dir.to_path_buf().join(CACHE_DIR)));
975 Self::open_in(dir, ctx)
976 }
977
978 #[cfg(feature = "encryption")]
981 pub fn open_encrypted(dir: impl AsRef<Path>, passphrase: &str) -> Result<Self> {
982 let dir = dir.as_ref();
983 let salt_path = dir.join(META_DIR).join(KEYS_FILENAME);
984 let salt_bytes = std::fs::read(&salt_path).map_err(|e| {
985 MongrelError::NotFound(format!(
986 "encryption salt file {:?}: {e} (table not encrypted, or corrupted)",
987 salt_path
988 ))
989 })?;
990 let salt_len = crate::encryption::SALT_LEN;
991 if salt_bytes.len() != salt_len {
992 return Err(MongrelError::InvalidArgument(format!(
993 "encryption salt is {} bytes, expected {salt_len}",
994 salt_bytes.len()
995 )));
996 }
997 let mut salt = [0u8; 16];
998 salt.copy_from_slice(&salt_bytes);
999 let kek: Arc<Kek> = Arc::new(Kek::derive(passphrase, &salt)?);
1000 let ctx = SharedCtx::new(Some(kek), Some(dir.to_path_buf().join(CACHE_DIR)));
1001 let t = Self::open_in(dir, ctx)?;
1002 Ok(t)
1003 }
1004
1005 pub(crate) fn open_in(dir: impl AsRef<Path>, ctx: SharedCtx) -> Result<Self> {
1006 let dir = dir.as_ref().to_path_buf();
1007 let manifest_meta_dek = crate::encryption::meta_dek_for(ctx.kek.as_deref());
1008 let manifest = manifest::read(&dir, manifest_meta_dek.as_ref())?;
1009 let schema: Schema = read_schema(&dir)?;
1010 let replay_epoch = Epoch(manifest.current_epoch);
1011 let (wal_dek, cache_dek) = derive_subkeys(ctx.kek.as_deref(), manifest.table_id);
1012 let (wal, replayed, current_txn_id) = match ctx.shared.clone() {
1016 Some(s) => (WalSink::Shared(s), Vec::new(), 0),
1017 None => {
1018 let active = latest_wal_segment(&dir.join(WAL_DIR))?;
1019 let replayed = match &active {
1021 Some(path) => {
1022 let cipher = wal_dek.as_ref().map(|dk| make_cipher(dk));
1023 crate::wal::replay_with_cipher(path, cipher)?
1024 }
1025 None => Vec::new(),
1026 };
1027 let mut w = match &active {
1028 Some(path) => Wal::create_with_cipher(
1029 path,
1030 replay_epoch,
1031 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1032 0,
1033 )?,
1034 None => Wal::create_with_cipher(
1035 dir.join(WAL_DIR).join("seg-000000.wal"),
1036 replay_epoch,
1037 wal_dek.as_ref().map(|dk| make_cipher(dk)),
1038 0,
1039 )?,
1040 };
1041 w.set_sync_byte_threshold(DEFAULT_SYNC_BYTE_THRESHOLD);
1042 (WalSink::Private(w), replayed, 1)
1043 }
1044 };
1045
1046 let mut memtable = Memtable::new();
1047 let mut allocator = RowIdAllocator::new(manifest.next_row_id);
1048 let persisted_epoch = manifest.current_epoch;
1049 let mut auto_inc = resolve_auto_inc(&schema).map(|mut s| {
1056 s.next = manifest.auto_inc_next;
1057 s.seeded = manifest.auto_inc_next > 0;
1058 s
1059 });
1060
1061 let mut staged_puts: HashMap<u64, Vec<Row>> = HashMap::new();
1068 let mut staged_deletes: HashMap<u64, Vec<RowId>> = HashMap::new();
1069 let mut replayed_puts: std::collections::BTreeMap<Epoch, Vec<Row>> =
1070 std::collections::BTreeMap::new();
1071 let mut replayed_deletes: Vec<(RowId, Epoch)> = Vec::new();
1072 let mut saw_delete = false;
1073 for record in replayed {
1074 let txn_id = record.txn_id;
1075 match record.op {
1076 Op::Put { rows, .. } => {
1077 let rows: Vec<Row> = bincode::deserialize(&rows)?;
1078 for row in &rows {
1079 allocator.advance_to(row.row_id);
1080 if let Some(ai) = auto_inc.as_mut() {
1081 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1082 if *n + 1 > ai.next {
1083 ai.next = *n + 1;
1084 }
1085 }
1086 }
1087 }
1088 staged_puts.entry(txn_id).or_default().extend(rows);
1089 }
1090 Op::Delete { row_ids, .. } => {
1091 staged_deletes.entry(txn_id).or_default().extend(row_ids);
1092 }
1093 Op::TxnCommit { epoch, .. } => {
1094 let commit_epoch = Epoch(epoch);
1095 if let Some(puts) = staged_puts.remove(&txn_id) {
1096 for row in &puts {
1097 memtable.upsert(row.clone());
1098 }
1099 replayed_puts.entry(commit_epoch).or_default().extend(puts);
1100 }
1101 if let Some(dels) = staged_deletes.remove(&txn_id) {
1102 saw_delete = true;
1103 for rid in dels {
1104 memtable.tombstone(rid, commit_epoch);
1105 replayed_deletes.push((rid, commit_epoch));
1106 }
1107 }
1108 }
1109 Op::TxnAbort => {
1110 staged_puts.remove(&txn_id);
1111 staged_deletes.remove(&txn_id);
1112 }
1113 Op::TruncateTable { .. } | Op::Flush { .. } | Op::Ddl(_) => {}
1114 }
1115 }
1116
1117 let rcache_dir = dir.join(RCACHE_DIR);
1118 let column_keys = build_column_keys(ctx.kek.as_deref(), &schema);
1119 let mut db = Self {
1120 dir,
1121 table_id: manifest.table_id,
1122 wal,
1123 memtable,
1124 mutable_run: MutableRun::new(),
1125 mutable_run_spill_bytes: DEFAULT_MUTABLE_RUN_SPILL_BYTES,
1126 compaction_zstd_level: 3,
1127 allocator,
1128 epoch: ctx.epoch,
1129 persisted_epoch,
1130 schema,
1131 hot: HotIndex::new(),
1132 kek: ctx.kek,
1133 column_keys,
1134 run_refs: manifest.runs.clone(),
1135 retiring: manifest.retiring.clone(),
1136 next_run_id: manifest
1137 .runs
1138 .iter()
1139 .map(|r| r.run_id as u64 + 1)
1140 .max()
1141 .unwrap_or(1),
1142 sync_byte_threshold: DEFAULT_SYNC_BYTE_THRESHOLD,
1143 current_txn_id,
1144 bitmap: HashMap::new(),
1145 ann: HashMap::new(),
1146 fm: HashMap::new(),
1147 sparse: HashMap::new(),
1148 minhash: HashMap::new(),
1149 learned_range: HashMap::new(),
1150 pk_by_row: HashMap::new(),
1151 pinned: BTreeMap::new(),
1152 live_count: manifest.live_count,
1153 reservoir: crate::reservoir::Reservoir::default(),
1154 had_deletes: saw_delete,
1155 agg_cache: HashMap::new(),
1156 global_idx_epoch: manifest.global_idx_epoch,
1157 indexes_complete: true,
1158 flushed_epoch: manifest.flushed_epoch,
1159 page_cache: ctx.page_cache,
1160 decoded_cache: ctx.decoded_cache,
1161 snapshots: ctx.snapshots,
1162 commit_lock: ctx.commit_lock,
1163 result_cache: Arc::new(parking_lot::Mutex::new(
1164 ResultCache::new()
1165 .with_dir(rcache_dir)
1166 .with_cache_dek(cache_dek.clone()),
1167 )),
1168 pending_delete_rids: roaring::RoaringBitmap::new(),
1169 pending_put_cols: std::collections::HashSet::new(),
1170 pending_rows: Vec::new(),
1171 pending_rows_auto_inc: Vec::new(),
1172 pending_dels: Vec::new(),
1173 pending_truncate: None,
1174 wal_dek,
1175 auto_inc,
1176 };
1177
1178 db.epoch.advance_recovered(Epoch(db.persisted_epoch));
1181
1182 let checkpoint = global_idx::read(&db.dir, db.idx_dek().as_deref())?;
1187 let checkpoint_valid = checkpoint.as_ref().is_some_and(|c| {
1188 c.epoch_built == manifest.global_idx_epoch
1189 && manifest.global_idx_epoch > 0
1190 && manifest
1191 .runs
1192 .iter()
1193 .all(|r| r.epoch_created <= manifest.global_idx_epoch)
1194 });
1195 if let Some(loaded) = checkpoint {
1196 if checkpoint_valid {
1197 db.hot = loaded.hot;
1198 db.bitmap = loaded.bitmap;
1199 db.ann = loaded.ann;
1200 db.fm = loaded.fm;
1201 db.sparse = loaded.sparse;
1202 db.minhash = loaded.minhash;
1203 db.learned_range = loaded.learned_range;
1204 db.refresh_pk_by_row_from_hot();
1205 }
1206 }
1207 if !checkpoint_valid {
1208 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&db.schema);
1209 db.bitmap = bitmap;
1210 db.ann = ann;
1211 db.fm = fm;
1212 db.sparse = sparse;
1213 db.minhash = minhash;
1214 db.rebuild_indexes_from_runs()?;
1215 db.build_learned_ranges()?;
1216 }
1217
1218 for (epoch, group) in replayed_puts {
1223 let (losers, winner_pks) = db.partition_pk_winners(&group);
1224 for (key, &row_id) in &winner_pks {
1225 if let Some(old_rid) = db.hot.get(key) {
1226 if old_rid != row_id {
1227 db.tombstone_row(old_rid, epoch, false);
1228 }
1229 }
1230 }
1231 for &loser_rid in &losers {
1232 db.tombstone_row(loser_rid, epoch, false);
1233 }
1234 for (key, row_id) in winner_pks {
1235 db.insert_hot_pk(key, row_id);
1236 }
1237 if db.schema.primary_key().is_none() {
1238 for r in &group {
1239 db.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1240 }
1241 }
1242 for r in &group {
1243 if !losers.contains(&r.row_id) {
1244 db.index_row(r);
1245 }
1246 }
1247 }
1248 for (rid, epoch) in &replayed_deletes {
1252 db.remove_hot_for_row(*rid, *epoch);
1253 }
1254
1255 let _ = db.rebuild_reservoir();
1256 db.result_cache.lock().load_persistent();
1259 Ok(db)
1260 }
1261
1262 fn rebuild_reservoir(&mut self) -> Result<()> {
1265 let snap = self.snapshot();
1266 let rows = self.visible_rows(snap)?;
1267 self.reservoir.reset();
1268 for r in rows {
1269 self.reservoir.offer(r.row_id.0);
1270 }
1271 Ok(())
1272 }
1273
1274 fn rebuild_indexes_from_runs(&mut self) -> Result<()> {
1275 self.hot = HotIndex::new();
1276 self.pk_by_row.clear();
1277 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
1278 self.bitmap = bitmap;
1279 self.ann = ann;
1280 self.fm = fm;
1281 self.sparse = sparse;
1282 self.minhash = minhash;
1283 let snapshot = Epoch(u64::MAX);
1284 for rr in self.run_refs.clone() {
1285 let mut reader = self.open_reader(rr.run_id)?;
1286 for row in reader.visible_rows(snapshot)? {
1287 let tok_row = self.tokenized_for_indexes(&row);
1288 index_into(
1289 &self.schema,
1290 &tok_row,
1291 &mut self.hot,
1292 &mut self.bitmap,
1293 &mut self.ann,
1294 &mut self.fm,
1295 &mut self.sparse,
1296 &mut self.minhash,
1297 );
1298 }
1299 }
1300 for row in self.mutable_run.visible_versions(snapshot) {
1301 if row.deleted {
1302 self.remove_hot_for_row(row.row_id, snapshot);
1303 } else {
1304 self.index_row(&row);
1305 }
1306 }
1307 for row in self.memtable.visible_versions(snapshot) {
1308 if row.deleted {
1309 self.remove_hot_for_row(row.row_id, snapshot);
1310 } else {
1311 self.index_row(&row);
1312 }
1313 }
1314 self.refresh_pk_by_row_from_hot();
1315 Ok(())
1316 }
1317
1318 fn refresh_pk_by_row_from_hot(&mut self) {
1319 self.pk_by_row.clear();
1320 if self.schema.primary_key().is_none() {
1321 return;
1322 }
1323 for (key, row_id) in self.hot.entries() {
1324 self.pk_by_row.insert(row_id, key);
1325 }
1326 }
1327
1328 fn insert_hot_pk(&mut self, key: Vec<u8>, row_id: RowId) {
1329 if self.schema.primary_key().is_some() {
1330 self.pk_by_row.insert(row_id, key.clone());
1331 }
1332 self.hot.insert(key, row_id);
1333 }
1334
1335 pub(crate) fn build_learned_ranges(&mut self) -> Result<()> {
1339 self.learned_range.clear();
1340 if self.run_refs.len() != 1 {
1341 return Ok(());
1342 }
1343 let cols: Vec<u16> = self
1344 .schema
1345 .indexes
1346 .iter()
1347 .filter(|i| i.kind == IndexKind::LearnedRange)
1348 .map(|i| i.column_id)
1349 .collect();
1350 if cols.is_empty() {
1351 return Ok(());
1352 }
1353 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
1354 let row_ids: Vec<u64> = match reader.column_native(crate::sorted_run::SYS_ROW_ID)? {
1355 columnar::NativeColumn::Int64 { data, .. } => data.iter().map(|x| *x as u64).collect(),
1356 _ => return Ok(()),
1357 };
1358 for cid in cols {
1359 let ty = self
1360 .schema
1361 .columns
1362 .iter()
1363 .find(|c| c.id == cid)
1364 .map(|c| c.ty)
1365 .unwrap_or(TypeId::Int64);
1366 match ty {
1367 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
1368 if let columnar::NativeColumn::Int64 { data, .. } = reader.column_native(cid)? {
1369 let pairs: Vec<(i64, u64)> = data
1370 .iter()
1371 .zip(row_ids.iter())
1372 .map(|(v, r)| (*v, *r))
1373 .collect();
1374 self.learned_range
1375 .insert(cid, ColumnLearnedRange::build_i64(&pairs));
1376 }
1377 }
1378 TypeId::Float64 => {
1379 if let columnar::NativeColumn::Float64 { data, .. } =
1380 reader.column_native(cid)?
1381 {
1382 let pairs: Vec<(f64, u64)> = data
1383 .iter()
1384 .zip(row_ids.iter())
1385 .map(|(v, r)| (*v, *r))
1386 .collect();
1387 self.learned_range
1388 .insert(cid, ColumnLearnedRange::build_f64(&pairs));
1389 }
1390 }
1391 _ => {}
1392 }
1393 }
1394 Ok(())
1395 }
1396
1397 fn ensure_indexes_complete(&mut self) -> Result<()> {
1401 if self.indexes_complete {
1402 crate::trace::QueryTrace::record(|t| {
1403 t.index_rebuild = crate::trace::IndexRebuild::AlreadyComplete;
1404 });
1405 return Ok(());
1406 }
1407 crate::trace::QueryTrace::record(|t| {
1408 t.index_rebuild = crate::trace::IndexRebuild::Rebuilt;
1409 });
1410 self.rebuild_indexes_from_runs()?;
1411 self.build_learned_ranges()?;
1412 self.indexes_complete = true;
1413 let epoch = self.current_epoch();
1414 self.checkpoint_indexes(epoch);
1415 Ok(())
1416 }
1417
1418 fn pending_epoch(&self) -> Epoch {
1419 Epoch(self.epoch.visible().0 + 1)
1420 }
1421
1422 fn is_shared(&self) -> bool {
1425 matches!(self.wal, WalSink::Shared(_))
1426 }
1427
1428 fn ensure_txn_id(&mut self) -> u64 {
1432 if self.current_txn_id == 0 {
1433 let id = match &self.wal {
1434 WalSink::Shared(s) => {
1435 let mut g = s.txn_ids.lock();
1436 let v = *g;
1437 *g = g.wrapping_add(1);
1438 v
1439 }
1440 WalSink::Private(_) => 1,
1441 };
1442 self.current_txn_id = id;
1443 }
1444 self.current_txn_id
1445 }
1446
1447 fn wal_append_data(&mut self, op: Op) -> Result<()> {
1450 let txn_id = self.ensure_txn_id();
1451 let table_id = self.table_id;
1452 match &mut self.wal {
1453 WalSink::Private(w) => {
1454 w.append_txn(txn_id, op)?;
1455 }
1456 WalSink::Shared(s) => {
1457 s.wal.lock().append(txn_id, table_id, op)?;
1458 }
1459 }
1460 Ok(())
1461 }
1462
1463 pub fn put(&mut self, columns: Vec<(u16, Value)>) -> Result<RowId> {
1471 Ok(self.put_returning(columns)?.0)
1472 }
1473
1474 pub fn put_returning(
1479 &mut self,
1480 mut columns: Vec<(u16, Value)>,
1481 ) -> Result<(RowId, Option<i64>)> {
1482 let assigned = self.fill_auto_inc(&mut columns)?;
1483 let mut col_map = std::collections::HashMap::with_capacity(columns.len());
1484 for (c, v) in &columns {
1485 col_map.insert(*c, v.clone());
1486 }
1487 self.schema.validate_not_null(&col_map)?;
1488 let row_id = self.allocator.alloc();
1489 let epoch = self.pending_epoch();
1490 let mut row = Row::new(row_id, epoch);
1491 for (col_id, val) in columns {
1492 row.columns.insert(col_id, val);
1493 }
1494 self.commit_rows(vec![row], assigned.is_some())?;
1495 Ok((row_id, assigned))
1496 }
1497
1498 pub fn put_batch(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Vec<RowId>> {
1501 Ok(self
1502 .put_batch_returning(batch)?
1503 .into_iter()
1504 .map(|(r, _)| r)
1505 .collect())
1506 }
1507
1508 pub fn put_batch_returning(
1511 &mut self,
1512 batch: Vec<Vec<(u16, Value)>>,
1513 ) -> Result<Vec<(RowId, Option<i64>)>> {
1514 let mut filled: Vec<FilledAutoIncRow> = Vec::with_capacity(batch.len());
1515 for mut cols in batch {
1516 let assigned = self.fill_auto_inc(&mut cols)?;
1517 filled.push((cols, assigned));
1518 }
1519 for (cols, _) in &filled {
1520 let mut col_map = std::collections::HashMap::with_capacity(cols.len());
1521 for (c, v) in cols {
1522 col_map.insert(*c, v.clone());
1523 }
1524 self.schema.validate_not_null(&col_map)?;
1525 }
1526 let epoch = self.pending_epoch();
1527 let mut rows = Vec::with_capacity(filled.len());
1528 let mut ids = Vec::with_capacity(filled.len());
1529 for (cols, assigned) in filled {
1530 let row_id = self.allocator.alloc();
1531 let mut row = Row::new(row_id, epoch);
1532 for (c, v) in cols {
1533 row.columns.insert(c, v);
1534 }
1535 ids.push((row_id, assigned));
1536 rows.push(row);
1537 }
1538 let all_auto_generated = ids.iter().all(|(_, assigned)| assigned.is_some());
1539 self.commit_rows(rows, all_auto_generated)?;
1540 Ok(ids)
1541 }
1542
1543 pub fn fill_auto_inc(&mut self, columns: &mut Vec<(u16, Value)>) -> Result<Option<i64>> {
1549 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1550 return Ok(None);
1551 };
1552 let pos = columns.iter().position(|(c, _)| *c == cid);
1553 let assigned = match pos {
1554 Some(i) => match &columns[i].1 {
1555 Value::Null => {
1556 let next = self.alloc_auto_inc_value()?;
1557 columns[i].1 = Value::Int64(next);
1558 Some(next)
1559 }
1560 Value::Int64(n) => {
1561 self.advance_auto_inc_past(*n)?;
1562 None
1563 }
1564 other => {
1565 return Err(MongrelError::InvalidArgument(format!(
1566 "AUTO_INCREMENT column {cid} must be Int64 or NULL, got {:?}",
1567 other
1568 )))
1569 }
1570 },
1571 None => {
1572 let next = self.alloc_auto_inc_value()?;
1573 columns.push((cid, Value::Int64(next)));
1574 Some(next)
1575 }
1576 };
1577 Ok(assigned)
1578 }
1579
1580 fn alloc_auto_inc_value(&mut self) -> Result<i64> {
1582 self.ensure_auto_inc_seeded()?;
1583 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1585 let v = ai.next;
1586 ai.next = ai.next.saturating_add(1);
1587 Ok(v)
1588 }
1589
1590 fn advance_auto_inc_past(&mut self, used: i64) -> Result<()> {
1593 self.ensure_auto_inc_seeded()?;
1594 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1595 let floor = used.saturating_add(1).max(1);
1596 if ai.next < floor {
1597 ai.next = floor;
1598 }
1599 Ok(())
1600 }
1601
1602 fn ensure_auto_inc_seeded(&mut self) -> Result<()> {
1607 let needs_seed = match self.auto_inc {
1608 Some(ai) => !ai.seeded,
1609 None => return Ok(()),
1610 };
1611 if !needs_seed {
1612 return Ok(());
1613 }
1614 if self.seed_empty_auto_inc() {
1615 return Ok(());
1616 }
1617 let cid = self
1618 .auto_inc
1619 .as_ref()
1620 .expect("auto-inc column present")
1621 .column_id;
1622 let max = self.scan_max_int64(cid)?;
1623 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1624 let floor = max.saturating_add(1).max(1);
1625 if ai.next < floor {
1626 ai.next = floor;
1627 }
1628 ai.seeded = true;
1629 Ok(())
1630 }
1631
1632 fn alloc_auto_inc_range(&mut self, n: usize) -> Result<Option<i64>> {
1633 if n == 0 || self.auto_inc.is_none() {
1634 return Ok(None);
1635 }
1636 self.ensure_auto_inc_seeded()?;
1637 let ai = self.auto_inc.as_mut().expect("auto-inc column present");
1638 let start = ai.next;
1639 ai.next = ai.next.saturating_add(n as i64);
1640 Ok(Some(start))
1641 }
1642
1643 fn scan_max_int64(&mut self, column_id: u16) -> Result<i64> {
1647 let mut max: i64 = 0;
1648 for r in self.memtable.visible_versions(Epoch(u64::MAX)) {
1649 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1650 if *n > max {
1651 max = *n;
1652 }
1653 }
1654 }
1655 for r in self.mutable_run.visible_versions(Epoch(u64::MAX)) {
1656 if let Some(Value::Int64(n)) = r.columns.get(&column_id) {
1657 if *n > max {
1658 max = *n;
1659 }
1660 }
1661 }
1662 for rr in self.run_refs.clone() {
1663 let reader = self.open_reader(rr.run_id)?;
1664 if let Some(stats) = reader.column_page_stats(column_id) {
1665 for s in stats {
1666 if let Some(n) = crate::sorted_run::be_i64(s.max.as_deref()) {
1667 if n > max {
1668 max = n;
1669 }
1670 }
1671 }
1672 } else if reader.has_column(column_id) {
1673 if let columnar::NativeColumn::Int64 { data, validity } =
1674 reader.column_native_shared(column_id)?
1675 {
1676 for (i, n) in data.iter().enumerate() {
1677 if (validity.is_empty() || columnar::validity_bit(&validity, i)) && *n > max
1678 {
1679 max = *n;
1680 }
1681 }
1682 }
1683 }
1684 }
1685 Ok(max)
1686 }
1687
1688 fn seed_empty_auto_inc(&mut self) -> bool {
1689 let Some(ai) = self.auto_inc.as_mut() else {
1690 return false;
1691 };
1692 if ai.seeded || self.live_count != 0 {
1693 return false;
1694 }
1695 if ai.next < 1 {
1696 ai.next = 1;
1697 }
1698 ai.seeded = true;
1699 true
1700 }
1701
1702 fn advance_auto_inc_from_native_columns(
1703 &mut self,
1704 columns: &[(u16, columnar::NativeColumn)],
1705 n: usize,
1706 live_before: u64,
1707 ) -> Result<()> {
1708 let Some(ai) = self.auto_inc.as_mut() else {
1709 return Ok(());
1710 };
1711 let Some((_, col)) = columns.iter().find(|(cid, _)| *cid == ai.column_id) else {
1712 return Ok(());
1713 };
1714 let columnar::NativeColumn::Int64 { data, validity } = col else {
1715 return Err(MongrelError::InvalidArgument(format!(
1716 "AUTO_INCREMENT column {} must be Int64",
1717 ai.column_id
1718 )));
1719 };
1720 let max = if native_int64_strictly_increasing(col, n) {
1721 data.get(n.saturating_sub(1)).copied()
1722 } else {
1723 data.iter()
1724 .take(n)
1725 .enumerate()
1726 .filter_map(|(i, v)| {
1727 if validity.is_empty() || columnar::validity_bit(validity, i) {
1728 Some(*v)
1729 } else {
1730 None
1731 }
1732 })
1733 .max()
1734 };
1735 if let Some(max) = max {
1736 let floor = max.saturating_add(1).max(1);
1737 if ai.next < floor {
1738 ai.next = floor;
1739 }
1740 if ai.seeded || live_before == 0 {
1741 ai.seeded = true;
1742 }
1743 }
1744 Ok(())
1745 }
1746
1747 fn fill_auto_inc_native_columns(
1748 &mut self,
1749 columns: &mut Vec<(u16, columnar::NativeColumn)>,
1750 n: usize,
1751 ) -> Result<()> {
1752 let Some(cid) = self.auto_inc.as_ref().map(|a| a.column_id) else {
1753 return Ok(());
1754 };
1755 let Some(pos) = columns.iter().position(|(id, _)| *id == cid) else {
1756 if let Some(start) = self.alloc_auto_inc_range(n)? {
1757 columns.push((
1758 cid,
1759 columnar::NativeColumn::Int64 {
1760 data: (start..start.saturating_add(n as i64)).collect(),
1761 validity: vec![0xFF; n.div_ceil(8)],
1762 },
1763 ));
1764 }
1765 return Ok(());
1766 };
1767
1768 let columnar::NativeColumn::Int64 { data, validity } = &mut columns[pos].1 else {
1769 return Err(MongrelError::InvalidArgument(format!(
1770 "AUTO_INCREMENT column {cid} must be Int64"
1771 )));
1772 };
1773 if data.len() < n {
1774 return Err(MongrelError::InvalidArgument(format!(
1775 "AUTO_INCREMENT column {cid} has {} rows, expected {n}",
1776 data.len()
1777 )));
1778 }
1779 if columnar::all_non_null(validity, n) {
1780 return Ok(());
1781 }
1782 if validity.iter().all(|b| *b == 0) {
1783 if let Some(start) = self.alloc_auto_inc_range(n)? {
1784 for (i, slot) in data.iter_mut().take(n).enumerate() {
1785 *slot = start.saturating_add(i as i64);
1786 }
1787 *validity = vec![0xFF; n.div_ceil(8)];
1788 }
1789 return Ok(());
1790 }
1791
1792 let new_validity = vec![0xFF; data.len().div_ceil(8)];
1793 for (i, slot) in data.iter_mut().enumerate().take(n) {
1794 if columnar::validity_bit(validity, i) {
1795 self.advance_auto_inc_past(*slot)?;
1796 } else {
1797 *slot = self.alloc_auto_inc_value()?;
1798 }
1799 }
1800 *validity = new_validity;
1801 Ok(())
1802 }
1803
1804 pub fn reserve_auto_inc(&mut self) -> Result<Option<i64>> {
1818 if self.auto_inc.is_none() {
1819 return Ok(None);
1820 }
1821 Ok(Some(self.alloc_auto_inc_value()?))
1822 }
1823
1824 fn commit_rows(&mut self, rows: Vec<Row>, auto_inc_generated: bool) -> Result<()> {
1830 let payload = bincode::serialize(&rows)?;
1831 self.wal_append_data(Op::Put {
1832 table_id: self.table_id,
1833 rows: payload,
1834 })?;
1835 if self.is_shared() {
1836 self.pending_rows_auto_inc
1837 .extend(std::iter::repeat(auto_inc_generated).take(rows.len()));
1838 self.pending_rows.extend(rows);
1839 } else {
1840 self.apply_put_rows_inner(rows, !auto_inc_generated)?;
1841 }
1842 Ok(())
1843 }
1844
1845 pub(crate) fn apply_put_rows(&mut self, rows: Vec<Row>) -> Result<()> {
1850 self.apply_put_rows_inner(rows, true)
1851 }
1852
1853 fn apply_put_rows_inner(&mut self, rows: Vec<Row>, check_existing_pk: bool) -> Result<()> {
1854 if check_existing_pk {
1855 self.ensure_indexes_complete()?;
1856 }
1857 let n = rows.len();
1858 for r in &rows {
1860 for &cid in r.columns.keys() {
1861 self.pending_put_cols.insert(cid);
1862 }
1863 }
1864 let (losers, winner_pks) = self.partition_pk_winners(&rows);
1865 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
1866 if check_existing_pk {
1868 for (key, &row_id) in &winner_pks {
1869 if let Some(old_rid) = self.hot.get(key) {
1870 if old_rid != row_id {
1871 self.tombstone_row(old_rid, epoch, true);
1872 }
1873 }
1874 }
1875 }
1876 for (key, row_id) in winner_pks {
1878 self.insert_hot_pk(key, row_id);
1879 }
1880 if self.schema.primary_key().is_none() {
1881 for r in &rows {
1882 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1883 }
1884 }
1885 for r in &rows {
1887 if !losers.contains(&r.row_id) {
1888 self.index_row(r);
1889 }
1890 }
1891 for r in &rows {
1892 if !losers.contains(&r.row_id) {
1893 self.reservoir.offer(r.row_id.0);
1894 }
1895 }
1896 for r in rows {
1897 if !losers.contains(&r.row_id) {
1898 self.memtable.upsert(r);
1899 }
1900 }
1901 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
1902 Ok(())
1903 }
1904
1905 pub(crate) fn alloc_row_id(&mut self) -> RowId {
1908 self.allocator.alloc()
1909 }
1910
1911 pub(crate) fn apply_run_metadata(&mut self, rows: &[Row]) -> Result<()> {
1919 self.ensure_indexes_complete()?;
1920 let n = rows.len();
1921 for r in rows {
1922 for &cid in r.columns.keys() {
1923 self.pending_put_cols.insert(cid);
1924 }
1925 }
1926 let (losers, winner_pks) = self.partition_pk_winners(rows);
1927 let epoch = rows.first().map(|r| r.committed_epoch).unwrap_or(Epoch(0));
1928 for (key, &row_id) in &winner_pks {
1930 if let Some(old_rid) = self.hot.get(key) {
1931 if old_rid != row_id {
1932 self.tombstone_row(old_rid, epoch, true);
1933 }
1934 }
1935 }
1936 for &loser_rid in &losers {
1939 self.tombstone_row(loser_rid, epoch, false);
1940 }
1941 for (key, row_id) in winner_pks {
1943 self.insert_hot_pk(key, row_id);
1944 }
1945 if self.schema.primary_key().is_none() {
1946 for r in rows {
1947 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
1948 }
1949 }
1950 for r in rows {
1951 self.allocator.advance_to(r.row_id);
1952 if !losers.contains(&r.row_id) {
1953 self.index_row(r);
1954 }
1955 }
1956 for r in rows {
1957 if !losers.contains(&r.row_id) {
1958 self.reservoir.offer(r.row_id.0);
1959 }
1960 }
1961 self.live_count = self.live_count.saturating_add((n - losers.len()) as u64);
1962 Ok(())
1963 }
1964
1965 pub(crate) fn recover_apply(
1970 &mut self,
1971 rows: Vec<Row>,
1972 deletes: Vec<(RowId, Epoch)>,
1973 ) -> Result<()> {
1974 let mut by_epoch: std::collections::BTreeMap<Epoch, Vec<Row>> =
1978 std::collections::BTreeMap::new();
1979 for row in rows {
1980 self.allocator.advance_to(row.row_id);
1981 if let Some(ai) = self.auto_inc.as_mut() {
1986 if let Some(Value::Int64(n)) = row.columns.get(&ai.column_id) {
1987 if *n + 1 > ai.next {
1988 ai.next = *n + 1;
1989 }
1990 }
1991 }
1992 by_epoch.entry(row.committed_epoch).or_default().push(row);
1993 }
1994 for (epoch, group) in by_epoch {
1995 let (losers, winner_pks) = self.partition_pk_winners(&group);
1996 for (key, &row_id) in &winner_pks {
1998 if let Some(old_rid) = self.hot.get(key) {
1999 if old_rid != row_id {
2000 self.tombstone_row(old_rid, epoch, false);
2001 }
2002 }
2003 }
2004 for (key, row_id) in winner_pks {
2005 self.insert_hot_pk(key, row_id);
2006 }
2007 if self.schema.primary_key().is_none() {
2008 for r in &group {
2009 self.hot.insert(r.row_id.0.to_be_bytes().to_vec(), r.row_id);
2010 }
2011 }
2012 for r in &group {
2013 if !losers.contains(&r.row_id) {
2014 self.memtable.upsert(r.clone());
2015 self.index_row(r);
2016 }
2017 }
2018 }
2019 for (rid, epoch) in deletes {
2020 self.memtable.tombstone(rid, epoch);
2021 self.remove_hot_for_row(rid, epoch);
2022 }
2023 let _ = self.rebuild_reservoir();
2024 Ok(())
2025 }
2026
2027 pub(crate) fn flushed_epoch(&self) -> u64 {
2029 self.flushed_epoch
2030 }
2031
2032 pub(crate) fn validate_cells_not_null(&self, cells: &[(u16, Value)]) -> Result<()> {
2034 let mut col_map = std::collections::HashMap::with_capacity(cells.len());
2035 for (c, v) in cells {
2036 col_map.insert(*c, v.clone());
2037 }
2038 self.schema.validate_not_null(&col_map)
2039 }
2040
2041 fn validate_columns_not_null(
2045 &self,
2046 columns: &[(u16, columnar::NativeColumn)],
2047 n: usize,
2048 ) -> Result<()> {
2049 let by_id: HashMap<u16, &columnar::NativeColumn> =
2050 columns.iter().map(|(id, c)| (*id, c)).collect();
2051 for col in &self.schema.columns {
2052 if col.flags.contains(ColumnFlags::NULLABLE) {
2053 continue;
2054 }
2055 match by_id.get(&col.id) {
2056 None => {
2057 return Err(MongrelError::InvalidArgument(format!(
2058 "column '{}' ({}) is NOT NULL but was omitted from the bulk load",
2059 col.name, col.id
2060 )));
2061 }
2062 Some(c) => {
2063 if c.null_count(n) != 0 {
2064 return Err(MongrelError::InvalidArgument(format!(
2065 "column '{}' ({}) is NOT NULL but the bulk load contains nulls",
2066 col.name, col.id
2067 )));
2068 }
2069 }
2070 }
2071 }
2072 Ok(())
2073 }
2074
2075 fn bulk_pk_winner_indices(
2080 &self,
2081 columns: &[(u16, columnar::NativeColumn)],
2082 n: usize,
2083 ) -> Option<Vec<usize>> {
2084 let pk_col = self.schema.primary_key()?;
2085 let pk_id = pk_col.id;
2086 let pk_ty = pk_col.ty;
2087 let by_id: HashMap<u16, &columnar::NativeColumn> =
2088 columns.iter().map(|(id, c)| (*id, c)).collect();
2089 let pk_native = by_id.get(&pk_id)?;
2090 if native_int64_strictly_increasing(pk_native, n) {
2091 return None;
2092 }
2093 let mut last: HashMap<Vec<u8>, usize> = HashMap::new();
2095 let mut null_pk_rows: Vec<usize> = Vec::new();
2096 for i in 0..n {
2097 match bulk_index_key(&self.column_keys, pk_id, pk_ty, pk_native, i) {
2098 Some(key) => {
2099 last.insert(key, i);
2100 }
2101 None => null_pk_rows.push(i),
2102 }
2103 }
2104 let mut winners: HashSet<usize> = last.values().copied().collect();
2105 for i in null_pk_rows {
2106 winners.insert(i);
2107 }
2108 Some((0..n).filter(|i| winners.contains(i)).collect())
2109 }
2110
2111 pub fn delete(&mut self, row_id: RowId) -> Result<()> {
2113 let epoch = self.pending_epoch();
2114 self.wal_append_data(Op::Delete {
2115 table_id: self.table_id,
2116 row_ids: vec![row_id],
2117 })?;
2118 if self.is_shared() {
2119 self.pending_dels.push(row_id);
2120 } else {
2121 self.apply_delete(row_id, epoch);
2122 }
2123 Ok(())
2124 }
2125
2126 pub fn delete_returning(&mut self, row_id: RowId) -> Result<Option<OwnedRow>> {
2127 let pre = self.get(row_id, self.snapshot());
2128 self.delete(row_id)?;
2129 Ok(pre.map(|row| {
2130 let mut columns: Vec<_> = row.columns.into_iter().collect();
2131 columns.sort_by_key(|(id, _)| *id);
2132 OwnedRow { columns }
2133 }))
2134 }
2135
2136 pub fn truncate(&mut self) -> Result<()> {
2138 let epoch = self.pending_epoch();
2139 self.wal_append_data(Op::TruncateTable {
2140 table_id: self.table_id,
2141 })?;
2142 self.pending_rows.clear();
2143 self.pending_rows_auto_inc.clear();
2144 self.pending_dels.clear();
2145 self.pending_truncate = Some(epoch);
2146 Ok(())
2147 }
2148
2149 pub(crate) fn apply_truncate(&mut self, _epoch: Epoch) -> Result<()> {
2151 for rr in std::mem::take(&mut self.run_refs) {
2152 let _ = std::fs::remove_file(self.run_path(rr.run_id as u64));
2153 }
2154 for r in std::mem::take(&mut self.retiring) {
2155 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
2156 }
2157 self.memtable = Memtable::new();
2158 self.mutable_run = MutableRun::new();
2159 self.hot = HotIndex::new();
2160 let (bitmap, ann, fm, sparse, minhash) = empty_indexes(&self.schema);
2161 self.bitmap = bitmap;
2162 self.ann = ann;
2163 self.fm = fm;
2164 self.sparse = sparse;
2165 self.minhash = minhash;
2166 self.learned_range.clear();
2167 self.pk_by_row.clear();
2168 self.live_count = 0;
2169 self.reservoir = crate::reservoir::Reservoir::default();
2170 self.had_deletes = true;
2171 self.agg_cache.clear();
2172 self.global_idx_epoch = 0;
2173 self.indexes_complete = true;
2174 self.pending_delete_rids.clear();
2175 self.pending_put_cols.clear();
2176 self.pending_rows.clear();
2177 self.pending_rows_auto_inc.clear();
2178 self.pending_dels.clear();
2179 self.clear_result_cache();
2180 self.invalidate_index_checkpoint();
2181 Ok(())
2182 }
2183
2184 pub(crate) fn apply_delete(&mut self, row_id: RowId, epoch: Epoch) {
2187 self.remove_hot_for_row(row_id, epoch);
2188 self.tombstone_row(row_id, epoch, true);
2189 }
2190
2191 fn tombstone_row(&mut self, row_id: RowId, epoch: Epoch, adjust_live_count: bool) {
2195 let tombstone = Row {
2196 row_id,
2197 committed_epoch: epoch,
2198 columns: std::collections::HashMap::new(),
2199 deleted: true,
2200 };
2201 self.memtable.upsert(tombstone);
2202 self.pk_by_row.remove(&row_id);
2203 if adjust_live_count {
2204 self.live_count = self.live_count.saturating_sub(1);
2205 }
2206 self.pending_delete_rids.insert(row_id.0 as u32);
2208 self.had_deletes = true;
2211 self.agg_cache.clear();
2212 }
2213
2214 fn remove_hot_for_row(&mut self, row_id: RowId, epoch: Epoch) {
2218 let Some(pk_col) = self.schema.primary_key() else {
2219 return;
2220 };
2221 if let Some(key) = self.pk_by_row.remove(&row_id) {
2222 if self.hot.get(&key) == Some(row_id) {
2223 self.hot.remove(&key);
2224 }
2225 return;
2226 }
2227 if !self.indexes_complete {
2228 return;
2229 }
2230 let pk_val = self
2233 .memtable
2234 .get_version(row_id, epoch)
2235 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2236 .or_else(|| {
2237 self.mutable_run
2238 .get_version(row_id, epoch)
2239 .filter(|(_, r)| !r.deleted)
2240 .and_then(|(_, r)| r.columns.get(&pk_col.id).cloned())
2241 })
2242 .or_else(|| {
2243 self.run_refs.iter().find_map(|rr| {
2244 let mut reader = self.open_reader(rr.run_id).ok()?;
2245 let (_, r) = reader.get_version(row_id, epoch).ok()??;
2246 if r.deleted {
2247 return None;
2248 }
2249 r.columns.get(&pk_col.id).cloned()
2250 })
2251 });
2252 if let Some(pk_val) = pk_val {
2253 let key = self.index_lookup_key(pk_col.id, &pk_val);
2254 if self.hot.get(&key) == Some(row_id) {
2255 self.hot.remove(&key);
2256 }
2257 }
2258 }
2259
2260 fn partition_pk_winners(
2265 &self,
2266 rows: &[Row],
2267 ) -> (
2268 std::collections::HashSet<RowId>,
2269 std::collections::HashMap<Vec<u8>, RowId>,
2270 ) {
2271 let mut losers = std::collections::HashSet::new();
2272 let Some(pk_col) = self.schema.primary_key() else {
2273 return (losers, std::collections::HashMap::new());
2274 };
2275 let pk_id = pk_col.id;
2276 let mut winners: std::collections::HashMap<Vec<u8>, RowId> =
2277 std::collections::HashMap::new();
2278 for r in rows {
2279 let Some(pk_val) = r.columns.get(&pk_id) else {
2280 continue;
2281 };
2282 let key = self.index_lookup_key(pk_id, pk_val);
2283 if let Some(&old_rid) = winners.get(&key) {
2284 losers.insert(old_rid);
2285 }
2286 winners.insert(key, r.row_id);
2287 }
2288 (losers, winners)
2289 }
2290
2291 fn index_row(&mut self, row: &Row) {
2292 if row.deleted {
2293 return;
2294 }
2295 let effective_row = self.tokenized_for_indexes(row);
2296 index_into(
2297 &self.schema,
2298 &effective_row,
2299 &mut self.hot,
2300 &mut self.bitmap,
2301 &mut self.ann,
2302 &mut self.fm,
2303 &mut self.sparse,
2304 &mut self.minhash,
2305 );
2306 }
2307
2308 fn tokenized_for_indexes(&self, row: &Row) -> Row {
2314 if self.column_keys.is_empty() {
2315 return row.clone();
2316 }
2317 #[cfg(feature = "encryption")]
2318 {
2319 use crate::encryption::SCHEME_HMAC_EQ;
2320 let mut tok = row.clone();
2321 for (&cid, &(_, scheme)) in &self.column_keys {
2322 if scheme != SCHEME_HMAC_EQ {
2323 continue;
2324 }
2325 if let Some(v) = tok.columns.get(&cid).cloned() {
2326 if let Some(t) = self.tokenize_value(cid, &v) {
2327 tok.columns.insert(cid, t);
2328 }
2329 }
2330 }
2331 tok
2332 }
2333 #[cfg(not(feature = "encryption"))]
2334 {
2335 row.clone()
2336 }
2337 }
2338
2339 pub fn commit(&mut self) -> Result<Epoch> {
2344 if self.is_shared() {
2345 self.commit_shared()
2346 } else {
2347 self.commit_private()
2348 }
2349 }
2350
2351 fn commit_private(&mut self) -> Result<Epoch> {
2353 let commit_lock = Arc::clone(&self.commit_lock);
2357 let _g = commit_lock.lock();
2358 let new_epoch = self.epoch.bump_assigned();
2359 let txn_id = self.current_txn_id;
2360 match &mut self.wal {
2364 WalSink::Private(w) => {
2365 w.append_txn(
2366 txn_id,
2367 Op::TxnCommit {
2368 epoch: new_epoch.0,
2369 added_runs: Vec::new(),
2370 },
2371 )?;
2372 w.sync()?;
2373 }
2374 WalSink::Shared(_) => unreachable!("commit_private on a shared sink"),
2375 }
2376 if let Some(epoch) = self.pending_truncate.take() {
2378 self.apply_truncate(epoch)?;
2379 }
2380 self.invalidate_pending_cache();
2381 self.persist_manifest(new_epoch)?;
2382 self.epoch.publish_in_order(new_epoch);
2386 self.current_txn_id += 1;
2387 Ok(new_epoch)
2388 }
2389
2390 fn commit_shared(&mut self) -> Result<Epoch> {
2396 use std::sync::atomic::Ordering;
2397 let s = match &self.wal {
2398 WalSink::Shared(s) => s.clone(),
2399 WalSink::Private(_) => unreachable!("commit_shared on a private sink"),
2400 };
2401 if s.poisoned.load(Ordering::Relaxed) {
2402 return Err(MongrelError::Other(
2403 "database poisoned by fsync error".into(),
2404 ));
2405 }
2406 let commit_lock = Arc::clone(&self.commit_lock);
2413 let _g = commit_lock.lock();
2414 let txn_id = self.ensure_txn_id();
2417 let (new_epoch, commit_seq) = {
2418 let mut wal = s.wal.lock();
2419 let new_epoch = self.epoch.bump_assigned();
2420 let seq = wal.append_commit(txn_id, new_epoch, &[])?;
2421 (new_epoch, seq)
2422 };
2423 s.group
2424 .await_durable(&s.wal, commit_seq)
2425 .inspect_err(|_| s.poisoned.store(true, Ordering::Relaxed))?;
2426
2427 if self.pending_truncate.take().is_some() {
2430 self.apply_truncate(new_epoch)?;
2431 }
2432 let mut rows = std::mem::take(&mut self.pending_rows);
2433 if !rows.is_empty() {
2434 for r in &mut rows {
2435 r.committed_epoch = new_epoch;
2436 }
2437 let auto_inc_flags = std::mem::take(&mut self.pending_rows_auto_inc);
2438 let all_auto_generated =
2439 auto_inc_flags.len() == rows.len() && auto_inc_flags.iter().all(|b| *b);
2440 self.apply_put_rows_inner(rows, !all_auto_generated)?;
2441 } else {
2442 self.pending_rows_auto_inc.clear();
2443 }
2444 let dels = std::mem::take(&mut self.pending_dels);
2445 for rid in dels {
2446 self.apply_delete(rid, new_epoch);
2447 }
2448
2449 self.invalidate_pending_cache();
2450 self.persist_manifest(new_epoch)?;
2451 self.epoch.publish_in_order(new_epoch);
2452 self.current_txn_id = 0;
2454 Ok(new_epoch)
2455 }
2456
2457 pub fn flush(&mut self) -> Result<Epoch> {
2465 self.ensure_indexes_complete()?;
2466 let epoch = self.commit()?;
2467 let rows = self.memtable.drain_sorted();
2468 if !rows.is_empty() {
2469 self.mutable_run.insert_many(rows);
2470 }
2471 if self.mutable_run.approx_bytes() >= self.mutable_run_spill_bytes {
2472 self.spill_mutable_run(epoch)?;
2473 self.mark_flushed(epoch)?;
2477 self.persist_manifest(epoch)?;
2478 self.build_learned_ranges()?;
2479 self.checkpoint_indexes(epoch);
2482 }
2483 Ok(epoch)
2486 }
2487
2488 fn mark_flushed(&mut self, epoch: Epoch) -> Result<()> {
2495 let op = Op::Flush {
2496 table_id: self.table_id,
2497 flushed_epoch: epoch.0,
2498 };
2499 match &mut self.wal {
2500 WalSink::Private(w) => {
2501 w.append_system(op)?;
2502 w.sync()?;
2503 }
2504 WalSink::Shared(s) => {
2505 s.wal.lock().append_system(op)?;
2510 }
2511 }
2512 self.flushed_epoch = epoch.0;
2513 if matches!(self.wal, WalSink::Private(_)) {
2514 self.rotate_wal(epoch)?;
2515 }
2516 Ok(())
2517 }
2518
2519 fn spill_mutable_run(&mut self, epoch: Epoch) -> Result<()> {
2523 let rows = self.mutable_run.drain_sorted();
2524 if rows.is_empty() {
2525 return Ok(());
2526 }
2527 let run_id = self.next_run_id;
2528 self.next_run_id += 1;
2529 let path = self.run_path(run_id);
2530 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0);
2531 if let Some(kek) = &self.kek {
2532 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2533 }
2534 let header = writer.write(&path, &rows)?;
2535 self.run_refs.push(RunRef {
2536 run_id: run_id as u128,
2537 level: 0,
2538 epoch_created: epoch.0,
2539 row_count: header.row_count,
2540 });
2541 Ok(())
2542 }
2543
2544 pub fn set_mutable_run_spill_bytes(&mut self, bytes: u64) {
2548 self.mutable_run_spill_bytes = bytes.max(1);
2549 }
2550
2551 pub fn set_compaction_zstd_level(&mut self, level: i32) {
2555 self.compaction_zstd_level = level;
2556 }
2557
2558 pub fn set_result_cache_max_bytes(&mut self, max_bytes: u64) {
2562 self.result_cache.lock().set_max_bytes(max_bytes);
2563 }
2564
2565 pub(crate) fn clear_result_cache(&mut self) {
2569 self.result_cache.lock().clear();
2570 }
2571
2572 pub fn mutable_run_len(&self) -> usize {
2574 self.mutable_run.len()
2575 }
2576
2577 pub(crate) fn drain_mutable_run(&mut self) -> Vec<Row> {
2580 self.mutable_run.drain_sorted()
2581 }
2582
2583 pub fn bulk_load(&mut self, batch: Vec<Vec<(u16, Value)>>) -> Result<Epoch> {
2588 let epoch = self.commit()?;
2589 let n = batch.len();
2590 if n == 0 {
2591 return Ok(epoch);
2592 }
2593 let live_before = self.live_count;
2594 self.spill_mutable_run(epoch)?;
2598 let eager_index_build = self.indexes_complete
2599 && self.run_refs.is_empty()
2600 && self.memtable.is_empty()
2601 && self.mutable_run.is_empty();
2602 let mut user_columns: Vec<(u16, columnar::NativeColumn)> = {
2608 use rayon::prelude::*;
2609 use std::collections::HashMap;
2610 let mut by_col: HashMap<u16, Vec<Value>> = HashMap::new();
2611 for cdef in &self.schema.columns {
2612 by_col.insert(cdef.id, vec![Value::Null; n]);
2613 }
2614 for (i, row) in batch.iter().enumerate() {
2615 for (id, v) in row {
2616 if let Some(col) = by_col.get_mut(id) {
2617 col[i] = v.clone();
2618 }
2619 }
2620 }
2621 self.schema
2622 .columns
2623 .par_iter()
2624 .map(|cdef| {
2625 let vals = by_col.get(&cdef.id).map(|v| v.as_slice()).unwrap_or(&[]);
2626 (cdef.id, columnar::values_to_native(cdef.ty, vals))
2627 })
2628 .collect::<Vec<_>>()
2629 };
2630 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
2635 self.validate_columns_not_null(&user_columns, n)?;
2636 let winner_idx = self
2637 .bulk_pk_winner_indices(&user_columns, n)
2638 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
2639 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
2640 match winner_idx.as_deref() {
2641 Some(idx) => {
2642 let compacted = user_columns
2643 .iter()
2644 .map(|(id, c)| (*id, c.gather(idx)))
2645 .collect();
2646 (compacted, idx.len())
2647 }
2648 None => (std::mem::take(&mut user_columns), n),
2649 };
2650 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
2651 let first = self.allocator.alloc_range(write_n as u64).0;
2652 for rid in first..first + write_n as u64 {
2653 self.reservoir.offer(rid);
2654 }
2655 let run_id = self.next_run_id;
2656 self.next_run_id += 1;
2657 let path = self.run_path(run_id);
2658 let mut writer = RunWriter::new(&self.schema, run_id as u128, epoch, 0)
2659 .clean(true)
2660 .with_lz4()
2661 .with_native_endian();
2662 if let Some(kek) = &self.kek {
2663 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
2664 }
2665 let header = writer.write_native(&path, &write_columns, write_n, first)?;
2666 self.run_refs.push(RunRef {
2667 run_id: run_id as u128,
2668 level: 0,
2669 epoch_created: epoch.0,
2670 row_count: header.row_count,
2671 });
2672 self.live_count = self.live_count.saturating_add(write_n as u64);
2673 if eager_index_build {
2674 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
2675 self.index_columns_bulk(&write_columns, &row_ids);
2676 self.indexes_complete = true;
2677 self.build_learned_ranges()?;
2678 } else {
2679 self.indexes_complete = false;
2680 }
2681 self.mark_flushed(epoch)?;
2682 self.persist_manifest(epoch)?;
2683 if eager_index_build {
2684 self.checkpoint_indexes(epoch);
2685 }
2686 self.clear_result_cache();
2687 Ok(epoch)
2688 }
2689
2690 fn rotate_wal(&mut self, epoch: Epoch) -> Result<()> {
2693 let segment = next_wal_segment(&self.dir.join(WAL_DIR))?;
2694 let cipher = self.wal_dek.as_ref().map(|dk| make_cipher(dk));
2695 let segment_no = segment
2698 .file_stem()
2699 .and_then(|s| s.to_str())
2700 .and_then(|s| s.strip_prefix("seg-"))
2701 .and_then(|s| s.parse::<u64>().ok())
2702 .unwrap_or(0);
2703 let mut wal = Wal::create_with_cipher(segment, epoch, cipher, segment_no)?;
2704 wal.set_sync_byte_threshold(self.sync_byte_threshold);
2705 wal.sync()?;
2706 self.wal = WalSink::Private(wal);
2707 Ok(())
2708 }
2709
2710 pub(crate) fn invalidate_pending_cache(&mut self) {
2715 self.result_cache
2716 .lock()
2717 .invalidate(&self.pending_delete_rids, &self.pending_put_cols);
2718 self.pending_delete_rids.clear();
2719 self.pending_put_cols.clear();
2720 }
2721
2722 pub(crate) fn persist_manifest(&self, epoch: Epoch) -> Result<()> {
2723 let mut m = Manifest::new(self.table_id, self.schema.schema_id);
2724 m.current_epoch = epoch.0;
2725 m.next_row_id = self.allocator.current().0;
2726 m.runs = self.run_refs.clone();
2727 m.live_count = self.live_count;
2728 m.global_idx_epoch = self.global_idx_epoch;
2729 m.flushed_epoch = self.flushed_epoch;
2730 m.retiring = self.retiring.clone();
2731 m.auto_inc_next = match self.auto_inc {
2735 Some(ai) if ai.seeded => ai.next,
2736 _ => 0,
2737 };
2738 let meta_dek = self.manifest_meta_dek();
2739 manifest::write_atomic(&self.dir, &mut m, meta_dek.as_ref())?;
2740 Ok(())
2741 }
2742
2743 pub(crate) fn checkpoint_indexes(&mut self, epoch: Epoch) {
2749 if !self.indexes_complete {
2752 return;
2753 }
2754 let snap = global_idx::IndexSnapshot {
2755 hot: &self.hot,
2756 bitmap: &self.bitmap,
2757 ann: &self.ann,
2758 fm: &self.fm,
2759 sparse: &self.sparse,
2760 minhash: &self.minhash,
2761 learned_range: &self.learned_range,
2762 };
2763 let idx_dek = self.idx_dek();
2765 if global_idx::write_atomic(&self.dir, self.table_id, epoch.0, snap, idx_dek.as_deref())
2766 .is_ok()
2767 {
2768 self.global_idx_epoch = epoch.0;
2769 let _ = self.persist_manifest(epoch);
2770 }
2771 }
2772
2773 pub(crate) fn invalidate_index_checkpoint(&mut self) {
2776 self.global_idx_epoch = 0;
2777 global_idx::remove(&self.dir);
2778 let _ = self.persist_manifest(self.epoch.visible());
2779 }
2780
2781 pub fn get(&self, row_id: RowId, snapshot: Snapshot) -> Option<Row> {
2784 let mut best: Option<(Epoch, Row)> = self.memtable.get_version(row_id, snapshot.epoch);
2785 if let Some((epoch, row)) = self.mutable_run.get_version(row_id, snapshot.epoch) {
2786 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2787 best = Some((epoch, row));
2788 }
2789 }
2790 for rr in &self.run_refs {
2791 let Ok(mut reader) = self.open_reader(rr.run_id) else {
2792 continue;
2793 };
2794 let Ok(Some((epoch, row))) = reader.get_version(row_id, snapshot.epoch) else {
2795 continue;
2796 };
2797 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
2798 best = Some((epoch, row));
2799 }
2800 }
2801 match best {
2802 Some((_, r)) if r.deleted => None,
2803 Some((_, r)) => Some(r),
2804 None => None,
2805 }
2806 }
2807
2808 pub fn visible_rows(&self, snapshot: Snapshot) -> Result<Vec<Row>> {
2812 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
2813 let mut fold = |row: Row| {
2814 best.entry(row.row_id.0)
2815 .and_modify(|e| {
2816 if row.committed_epoch > e.0 {
2817 *e = (row.committed_epoch, row.clone());
2818 }
2819 })
2820 .or_insert_with(|| (row.committed_epoch, row));
2821 };
2822 for row in self.memtable.visible_versions(snapshot.epoch) {
2823 fold(row);
2824 }
2825 for row in self.mutable_run.visible_versions(snapshot.epoch) {
2826 fold(row);
2827 }
2828 for rr in &self.run_refs {
2829 let mut reader = self.open_reader(rr.run_id)?;
2830 for row in reader.visible_versions(snapshot.epoch)? {
2831 fold(row);
2832 }
2833 }
2834 let mut out: Vec<Row> = best
2835 .into_values()
2836 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
2837 .collect();
2838 out.sort_by_key(|r| r.row_id);
2839 Ok(out)
2840 }
2841
2842 pub fn visible_columns(&self, snapshot: Snapshot) -> Result<Vec<(u16, Vec<Value>)>> {
2849 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
2850 let rr = self.run_refs[0].clone();
2851 let mut reader = self.open_reader(rr.run_id)?;
2852 let idxs = reader.visible_indices(snapshot.epoch)?;
2853 let mut cols = Vec::with_capacity(self.schema.columns.len());
2854 for cdef in &self.schema.columns {
2855 cols.push((cdef.id, reader.gather_column(cdef.id, &idxs)?));
2856 }
2857 return Ok(cols);
2858 }
2859 let rows = self.visible_rows(snapshot)?;
2861 let mut cols: Vec<(u16, Vec<Value>)> = self
2862 .schema
2863 .columns
2864 .iter()
2865 .map(|c| (c.id, Vec::with_capacity(rows.len())))
2866 .collect();
2867 for r in &rows {
2868 for (cid, vec) in cols.iter_mut() {
2869 vec.push(r.columns.get(cid).cloned().unwrap_or(Value::Null));
2870 }
2871 }
2872 Ok(cols)
2873 }
2874
2875 pub fn lookup_pk(&self, key: &[u8]) -> Option<RowId> {
2877 self.hot.get(key)
2878 }
2879
2880 pub fn query(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
2885 self.ensure_indexes_complete()?;
2886 let snapshot = self.snapshot();
2887 crate::trace::QueryTrace::record(|t| {
2888 t.run_count = self.run_refs.len();
2889 t.memtable_rows = self.memtable.len();
2890 t.mutable_run_rows = self.mutable_run.len();
2891 });
2892 if q.conditions.is_empty() {
2896 crate::trace::QueryTrace::record(|t| {
2897 t.scan_mode = crate::trace::ScanMode::Materialized;
2898 t.row_materialized = true;
2899 });
2900 return self.visible_rows(snapshot);
2901 }
2902 crate::trace::QueryTrace::record(|t| {
2903 t.conditions_pushed = q.conditions.len();
2904 t.scan_mode = crate::trace::ScanMode::Materialized;
2905 t.row_materialized = true;
2906 });
2907 let mut sets: Vec<RowIdSet> = Vec::with_capacity(q.conditions.len());
2908 for c in &q.conditions {
2909 sets.push(self.resolve_condition(c, snapshot)?);
2910 }
2911 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
2912 self.rows_for_rids(&rids, snapshot)
2913 }
2914
2915 pub fn rows_for_rids(&self, rids: &[u64], snapshot: Snapshot) -> Result<Vec<Row>> {
2920 use std::collections::HashMap;
2921 let mut rows = Vec::with_capacity(rids.len());
2922 let mut overlay: HashMap<u64, Row> = HashMap::new();
2928 let fold_newest = |row: Row, overlay: &mut HashMap<u64, Row>| {
2929 overlay
2930 .entry(row.row_id.0)
2931 .and_modify(|e| {
2932 if row.committed_epoch > e.committed_epoch {
2933 *e = row.clone();
2934 }
2935 })
2936 .or_insert(row);
2937 };
2938 for row in self.memtable.visible_versions(snapshot.epoch) {
2939 fold_newest(row, &mut overlay);
2940 }
2941 for row in self.mutable_run.visible_versions(snapshot.epoch) {
2942 fold_newest(row, &mut overlay);
2943 }
2944 if self.run_refs.len() == 1 {
2945 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
2954 let (positions, vis_rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
2955 enum Src {
2958 Overlay,
2959 Run,
2960 }
2961 let mut plan: Vec<Src> = Vec::with_capacity(rids.len());
2962 let mut fetch: Vec<usize> = Vec::with_capacity(rids.len());
2963 for rid in rids {
2964 if overlay.contains_key(rid) {
2965 plan.push(Src::Overlay);
2966 continue;
2967 }
2968 match vis_rids.binary_search(&(*rid as i64)) {
2969 Ok(i) => {
2970 plan.push(Src::Run);
2971 fetch.push(positions[i]);
2972 }
2973 Err(_) => { }
2974 }
2975 }
2976 let fetched = reader.materialize_batch(&fetch)?;
2977 let mut fetched_iter = fetched.into_iter();
2978 for (rid, src) in rids.iter().zip(plan) {
2979 match src {
2980 Src::Overlay => {
2981 if let Some(r) = overlay.get(rid) {
2982 if !r.deleted {
2983 rows.push(r.clone());
2984 }
2985 }
2986 }
2987 Src::Run => {
2988 if let Some(row) = fetched_iter.next() {
2989 if !row.deleted {
2990 rows.push(row);
2991 }
2992 }
2993 }
2994 }
2995 }
2996 return Ok(rows);
2997 }
2998 let mut readers: Vec<_> = self
3002 .run_refs
3003 .iter()
3004 .map(|rr| self.open_reader(rr.run_id))
3005 .collect::<Result<Vec<_>>>()?;
3006 for rid in rids {
3007 if let Some(r) = overlay.get(rid) {
3008 if !r.deleted {
3009 rows.push(r.clone());
3010 }
3011 continue;
3012 }
3013 let mut best: Option<(Epoch, Row)> = None;
3014 for reader in readers.iter_mut() {
3015 if let Ok(Some((epoch, row))) = reader.get_version(RowId(*rid), snapshot.epoch) {
3016 if best.as_ref().map(|(be, _)| epoch > *be).unwrap_or(true) {
3017 best = Some((epoch, row));
3018 }
3019 }
3020 }
3021 if let Some((_, r)) = best {
3022 if !r.deleted {
3023 rows.push(r);
3024 }
3025 }
3026 }
3027 Ok(rows)
3028 }
3029
3030 pub fn indexes_complete(&self) -> bool {
3040 self.indexes_complete
3041 }
3042
3043 pub fn broadcast_join_values(&self, column_id: u16, pk_db: &Table) -> Option<Vec<Vec<u8>>> {
3048 let b = self.bitmap.get(&column_id)?;
3049 let result: Vec<Vec<u8>> = b
3050 .keys()
3051 .into_iter()
3052 .filter(|k| pk_db.hot.get(k.as_slice()).is_some())
3053 .cloned()
3054 .collect();
3055 Some(result)
3056 }
3057
3058 pub fn fk_join_row_ids(
3059 &self,
3060 fk_column_id: u16,
3061 pk_values: &[Vec<u8>],
3062 fk_conditions: &[crate::query::Condition],
3063 snapshot: Snapshot,
3064 ) -> Result<Vec<u64>> {
3065 let Some(b) = self.bitmap.get(&fk_column_id) else {
3066 return Ok(Vec::new());
3067 };
3068 let mut join_set = {
3069 let mut acc = roaring::RoaringBitmap::new();
3070 for v in pk_values {
3071 acc |= b.get(v);
3072 }
3073 RowIdSet::from_roaring(acc)
3074 };
3075 if !fk_conditions.is_empty() {
3076 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3077 sets.push(join_set);
3078 for c in fk_conditions {
3079 sets.push(self.resolve_condition(c, snapshot)?);
3080 }
3081 join_set = RowIdSet::intersect_many(sets);
3082 }
3083 Ok(join_set.into_sorted_vec())
3084 }
3085
3086 pub fn fk_join_count(
3092 &self,
3093 fk_column_id: u16,
3094 pk_values: &[Vec<u8>],
3095 fk_conditions: &[crate::query::Condition],
3096 snapshot: Snapshot,
3097 ) -> Result<u64> {
3098 let Some(b) = self.bitmap.get(&fk_column_id) else {
3099 return Ok(0);
3100 };
3101 let mut acc = roaring::RoaringBitmap::new();
3102 for v in pk_values {
3103 acc |= b.get(v);
3104 }
3105 if fk_conditions.is_empty() {
3106 return Ok(acc.len());
3107 }
3108 let mut sets: Vec<RowIdSet> = Vec::with_capacity(fk_conditions.len() + 1);
3109 sets.push(RowIdSet::from_roaring(acc));
3110 for c in fk_conditions {
3111 sets.push(self.resolve_condition(c, snapshot)?);
3112 }
3113 Ok(RowIdSet::intersect_many(sets).len() as u64)
3114 }
3115
3116 fn resolve_condition(
3121 &self,
3122 c: &crate::query::Condition,
3123 snapshot: Snapshot,
3124 ) -> Result<RowIdSet> {
3125 use crate::query::Condition;
3126 Ok(match c {
3127 Condition::Pk(key) => {
3128 let lookup = self
3129 .schema
3130 .primary_key()
3131 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
3132 .unwrap_or_else(|| key.clone());
3133 self.hot
3134 .get(&lookup)
3135 .map(|r| RowIdSet::one(r.0))
3136 .unwrap_or_else(RowIdSet::empty)
3137 }
3138 Condition::BitmapEq { column_id, value } => {
3139 let lookup = self.index_lookup_key_bytes(*column_id, value);
3140 self.bitmap
3141 .get(column_id)
3142 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
3143 .unwrap_or_else(RowIdSet::empty)
3144 }
3145 Condition::BitmapIn { column_id, values } => {
3146 let bm = self.bitmap.get(column_id);
3147 let mut acc = roaring::RoaringBitmap::new();
3148 if let Some(b) = bm {
3149 for v in values {
3150 let lookup = self.index_lookup_key_bytes(*column_id, v);
3151 acc |= b.get(&lookup);
3152 }
3153 }
3154 RowIdSet::from_roaring(acc)
3155 }
3156 Condition::FmContains { column_id, pattern } => self
3157 .fm
3158 .get(column_id)
3159 .map(|f| {
3160 RowIdSet::from_unsorted(f.locate(pattern).into_iter().map(|r| r.0).collect())
3161 })
3162 .unwrap_or_else(RowIdSet::empty),
3163 Condition::FmContainsAll {
3164 column_id,
3165 patterns,
3166 } => {
3167 if let Some(f) = self.fm.get(column_id) {
3170 let sets: Vec<RowIdSet> = patterns
3171 .iter()
3172 .map(|pat| {
3173 RowIdSet::from_unsorted(
3174 f.locate(pat).into_iter().map(|r| r.0).collect(),
3175 )
3176 })
3177 .collect();
3178 RowIdSet::intersect_many(sets)
3179 } else {
3180 RowIdSet::empty()
3181 }
3182 }
3183 Condition::Ann {
3184 column_id,
3185 query,
3186 k,
3187 } => self
3188 .ann
3189 .get(column_id)
3190 .map(|a| {
3191 RowIdSet::from_unsorted(
3192 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3193 )
3194 })
3195 .unwrap_or_else(RowIdSet::empty),
3196 Condition::SparseMatch {
3197 column_id,
3198 query,
3199 k,
3200 } => self
3201 .sparse
3202 .get(column_id)
3203 .map(|s| {
3204 RowIdSet::from_unsorted(
3205 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3206 )
3207 })
3208 .unwrap_or_else(RowIdSet::empty),
3209 Condition::MinHashSimilar {
3210 column_id,
3211 query,
3212 k,
3213 } => self
3214 .minhash
3215 .get(column_id)
3216 .map(|mh| {
3217 RowIdSet::from_unsorted(
3218 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
3219 )
3220 })
3221 .unwrap_or_else(RowIdSet::empty),
3222 Condition::Range { column_id, lo, hi } => {
3223 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3232 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
3233 } else if self.run_refs.len() == 1 {
3234 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3235 r.range_row_id_set_i64(*column_id, *lo, *hi)?
3236 } else {
3237 return self.range_scan_i64(*column_id, *lo, *hi, snapshot);
3238 };
3239 set.remove_many(self.overlay_rid_set(snapshot));
3240 self.range_scan_overlay_i64(&mut set, *column_id, *lo, *hi, snapshot);
3241 set
3242 }
3243 Condition::RangeF64 {
3244 column_id,
3245 lo,
3246 lo_inclusive,
3247 hi,
3248 hi_inclusive,
3249 } => {
3250 let mut set = if let Some(li) = self.learned_range.get(column_id) {
3253 RowIdSet::from_unsorted(
3254 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
3255 .into_iter()
3256 .collect(),
3257 )
3258 } else if self.run_refs.len() == 1 {
3259 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3260 r.range_row_id_set_f64(*column_id, *lo, *lo_inclusive, *hi, *hi_inclusive)?
3261 } else {
3262 return self.range_scan_f64(
3263 *column_id,
3264 *lo,
3265 *lo_inclusive,
3266 *hi,
3267 *hi_inclusive,
3268 snapshot,
3269 );
3270 };
3271 set.remove_many(self.overlay_rid_set(snapshot));
3272 self.range_scan_overlay_f64(
3273 &mut set,
3274 *column_id,
3275 *lo,
3276 *lo_inclusive,
3277 *hi,
3278 *hi_inclusive,
3279 snapshot,
3280 );
3281 set
3282 }
3283 Condition::IsNull { column_id } => {
3284 let mut set = if self.run_refs.len() == 1 {
3285 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3286 r.null_row_id_set(*column_id, true)?
3287 } else {
3288 return self.null_scan(*column_id, true, snapshot);
3289 };
3290 set.remove_many(self.overlay_rid_set(snapshot));
3291 self.null_scan_overlay(&mut set, *column_id, true, snapshot);
3292 set
3293 }
3294 Condition::IsNotNull { column_id } => {
3295 let mut set = if self.run_refs.len() == 1 {
3296 let mut r = self.open_reader(self.run_refs[0].run_id)?;
3297 r.null_row_id_set(*column_id, false)?
3298 } else {
3299 return self.null_scan(*column_id, false, snapshot);
3300 };
3301 set.remove_many(self.overlay_rid_set(snapshot));
3302 self.null_scan_overlay(&mut set, *column_id, false, snapshot);
3303 set
3304 }
3305 })
3306 }
3307
3308 fn range_scan_i64(
3316 &self,
3317 column_id: u16,
3318 lo: i64,
3319 hi: i64,
3320 snapshot: Snapshot,
3321 ) -> Result<RowIdSet> {
3322 let mut row_ids = Vec::new();
3323 let overlay_rids = self.overlay_rid_set(snapshot);
3324 for rr in &self.run_refs {
3325 let mut reader = self.open_reader(rr.run_id)?;
3326 let matched = reader.range_row_ids_visible_i64(column_id, lo, hi, snapshot.epoch)?;
3327 for rid in matched {
3328 if !overlay_rids.contains(&rid) {
3329 row_ids.push(rid);
3330 }
3331 }
3332 }
3333 let mut s = RowIdSet::from_unsorted(row_ids);
3334 self.range_scan_overlay_i64(&mut s, column_id, lo, hi, snapshot);
3335 Ok(s)
3336 }
3337
3338 fn range_scan_f64(
3341 &self,
3342 column_id: u16,
3343 lo: f64,
3344 lo_inclusive: bool,
3345 hi: f64,
3346 hi_inclusive: bool,
3347 snapshot: Snapshot,
3348 ) -> Result<RowIdSet> {
3349 let mut row_ids = Vec::new();
3350 let overlay_rids = self.overlay_rid_set(snapshot);
3351 for rr in &self.run_refs {
3352 let mut reader = self.open_reader(rr.run_id)?;
3353 let matched = reader.range_row_ids_visible_f64(
3354 column_id,
3355 lo,
3356 lo_inclusive,
3357 hi,
3358 hi_inclusive,
3359 snapshot.epoch,
3360 )?;
3361 for rid in matched {
3362 if !overlay_rids.contains(&rid) {
3363 row_ids.push(rid);
3364 }
3365 }
3366 }
3367 let mut s = RowIdSet::from_unsorted(row_ids);
3368 self.range_scan_overlay_f64(
3369 &mut s,
3370 column_id,
3371 lo,
3372 lo_inclusive,
3373 hi,
3374 hi_inclusive,
3375 snapshot,
3376 );
3377 Ok(s)
3378 }
3379
3380 fn overlay_rid_set(&self, snapshot: Snapshot) -> HashSet<u64> {
3382 let mut s = HashSet::new();
3383 for row in self.memtable.visible_versions(snapshot.epoch) {
3384 s.insert(row.row_id.0);
3385 }
3386 for row in self.mutable_run.visible_versions(snapshot.epoch) {
3387 s.insert(row.row_id.0);
3388 }
3389 s
3390 }
3391
3392 fn range_scan_overlay_i64(
3393 &self,
3394 s: &mut RowIdSet,
3395 column_id: u16,
3396 lo: i64,
3397 hi: i64,
3398 snapshot: Snapshot,
3399 ) {
3400 let mut newest: HashMap<u64, &Row> = HashMap::new();
3405 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3406 let memtable = self.memtable.visible_versions(snapshot.epoch);
3407 for r in &mutable {
3408 newest.entry(r.row_id.0).or_insert(r);
3409 }
3410 for r in &memtable {
3411 newest.insert(r.row_id.0, r);
3412 }
3413 for row in newest.values() {
3414 if !row.deleted {
3415 if let Some(Value::Int64(v)) = row.columns.get(&column_id) {
3416 if *v >= lo && *v <= hi {
3417 s.insert(row.row_id.0);
3418 }
3419 }
3420 }
3421 }
3422 }
3423
3424 #[allow(clippy::too_many_arguments)]
3425 fn range_scan_overlay_f64(
3426 &self,
3427 s: &mut RowIdSet,
3428 column_id: u16,
3429 lo: f64,
3430 lo_inclusive: bool,
3431 hi: f64,
3432 hi_inclusive: bool,
3433 snapshot: Snapshot,
3434 ) {
3435 let mut newest: HashMap<u64, &Row> = HashMap::new();
3438 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3439 let memtable = self.memtable.visible_versions(snapshot.epoch);
3440 for r in &mutable {
3441 newest.entry(r.row_id.0).or_insert(r);
3442 }
3443 for r in &memtable {
3444 newest.insert(r.row_id.0, r);
3445 }
3446 for row in newest.values() {
3447 if !row.deleted {
3448 if let Some(Value::Float64(v)) = row.columns.get(&column_id) {
3449 let ok_lo = if lo_inclusive { *v >= lo } else { *v > lo };
3450 let ok_hi = if hi_inclusive { *v <= hi } else { *v < hi };
3451 if ok_lo && ok_hi {
3452 s.insert(row.row_id.0);
3453 }
3454 }
3455 }
3456 }
3457 }
3458
3459 fn null_scan(&self, column_id: u16, want_nulls: bool, snapshot: Snapshot) -> Result<RowIdSet> {
3462 let mut row_ids = Vec::new();
3463 let overlay_rids = self.overlay_rid_set(snapshot);
3464 for rr in &self.run_refs {
3465 let mut reader = self.open_reader(rr.run_id)?;
3466 let matched = reader.null_row_ids_visible(column_id, want_nulls, snapshot.epoch)?;
3467 for rid in matched {
3468 if !overlay_rids.contains(&rid) {
3469 row_ids.push(rid);
3470 }
3471 }
3472 }
3473 let mut s = RowIdSet::from_unsorted(row_ids);
3474 self.null_scan_overlay(&mut s, column_id, want_nulls, snapshot);
3475 Ok(s)
3476 }
3477
3478 fn null_scan_overlay(
3482 &self,
3483 s: &mut RowIdSet,
3484 column_id: u16,
3485 want_nulls: bool,
3486 snapshot: Snapshot,
3487 ) {
3488 let mut newest: HashMap<u64, &Row> = HashMap::new();
3489 let mutable = self.mutable_run.visible_versions(snapshot.epoch);
3490 let memtable = self.memtable.visible_versions(snapshot.epoch);
3491 for r in &mutable {
3492 newest.entry(r.row_id.0).or_insert(r);
3493 }
3494 for r in &memtable {
3495 newest.insert(r.row_id.0, r);
3496 }
3497 for row in newest.values() {
3498 if row.deleted {
3499 continue;
3500 }
3501 let is_null = !row.columns.contains_key(&column_id)
3502 || matches!(row.columns.get(&column_id), Some(Value::Null) | None);
3503 if is_null == want_nulls {
3504 s.insert(row.row_id.0);
3505 }
3506 }
3507 }
3508
3509 pub fn snapshot(&self) -> Snapshot {
3510 Snapshot::at(self.epoch.visible())
3511 }
3512
3513 pub fn pin_snapshot(&mut self) -> Snapshot {
3516 let e = self.epoch.visible();
3517 *self.pinned.entry(e).or_insert(0) += 1;
3518 Snapshot::at(e)
3519 }
3520
3521 pub fn unpin_snapshot(&mut self, snap: Snapshot) {
3523 if let Some(count) = self.pinned.get_mut(&snap.epoch) {
3524 *count -= 1;
3525 if *count == 0 {
3526 self.pinned.remove(&snap.epoch);
3527 }
3528 }
3529 }
3530
3531 pub(crate) fn min_active_snapshot(&self) -> Option<Epoch> {
3541 let local = self.pinned.keys().next().copied();
3542 let global = self.snapshots.min_pinned();
3543 match (local, global) {
3544 (Some(a), Some(b)) => Some(a.min(b)),
3545 (Some(a), None) => Some(a),
3546 (None, b) => b,
3547 }
3548 }
3549
3550 pub fn current_epoch(&self) -> Epoch {
3551 self.epoch.visible()
3552 }
3553
3554 pub fn memtable_len(&self) -> usize {
3555 self.memtable.len()
3556 }
3557
3558 pub fn count(&self) -> u64 {
3561 self.live_count
3562 }
3563
3564 pub fn count_conditions(
3568 &mut self,
3569 conditions: &[crate::query::Condition],
3570 snapshot: Snapshot,
3571 ) -> Result<Option<u64>> {
3572 use crate::query::Condition;
3573 if conditions.is_empty() {
3574 return Ok(Some(self.live_count));
3575 }
3576 let served = |c: &Condition| {
3577 matches!(
3578 c,
3579 Condition::Pk(_)
3580 | Condition::BitmapEq { .. }
3581 | Condition::BitmapIn { .. }
3582 | Condition::FmContains { .. }
3583 | Condition::FmContainsAll { .. }
3584 | Condition::Ann { .. }
3585 | Condition::Range { .. }
3586 | Condition::RangeF64 { .. }
3587 | Condition::SparseMatch { .. }
3588 | Condition::MinHashSimilar { .. }
3589 | Condition::IsNull { .. }
3590 | Condition::IsNotNull { .. }
3591 )
3592 };
3593 if !conditions.iter().all(served) {
3594 return Ok(None);
3595 }
3596 self.ensure_indexes_complete()?;
3597 let mut sets = Vec::with_capacity(conditions.len());
3598 for condition in conditions {
3599 sets.push(self.resolve_condition(condition, snapshot)?);
3600 }
3601 let count = RowIdSet::intersect_many(sets).len() as u64;
3602 crate::trace::QueryTrace::record(|t| {
3603 t.scan_mode = crate::trace::ScanMode::CountSurvivors;
3604 t.survivor_count = Some(count as usize);
3605 t.conditions_pushed = conditions.len();
3606 });
3607 Ok(Some(count))
3608 }
3609
3610 pub fn bulk_load_columns(
3617 &mut self,
3618 user_columns: Vec<(u16, columnar::NativeColumn)>,
3619 ) -> Result<Epoch> {
3620 self.bulk_load_columns_with(user_columns, 3, false, true)
3621 }
3622
3623 pub fn bulk_load_fast(
3630 &mut self,
3631 user_columns: Vec<(u16, columnar::NativeColumn)>,
3632 ) -> Result<Epoch> {
3633 self.bulk_load_columns_with(user_columns, -1, true, false)
3634 }
3635
3636 fn bulk_load_columns_with(
3637 &mut self,
3638 mut user_columns: Vec<(u16, columnar::NativeColumn)>,
3639 zstd_level: i32,
3640 force_plain: bool,
3641 lz4: bool,
3642 ) -> Result<Epoch> {
3643 let epoch = self.commit()?;
3644 let n = user_columns.first().map(|(_, c)| c.len()).unwrap_or(0);
3645 if n == 0 {
3646 return Ok(epoch);
3647 }
3648 let live_before = self.live_count;
3649 self.spill_mutable_run(epoch)?;
3651 let eager_index_build = self.indexes_complete
3652 && self.run_refs.is_empty()
3653 && self.memtable.is_empty()
3654 && self.mutable_run.is_empty();
3655 self.fill_auto_inc_native_columns(&mut user_columns, n)?;
3658 self.validate_columns_not_null(&user_columns, n)?;
3659 let winner_idx = self
3660 .bulk_pk_winner_indices(&user_columns, n)
3661 .and_then(|idx| if idx.len() == n { None } else { Some(idx) });
3662 let (write_columns, write_n): (Vec<(u16, columnar::NativeColumn)>, usize) =
3663 match winner_idx.as_deref() {
3664 Some(idx) => {
3665 let compacted = user_columns
3666 .iter()
3667 .map(|(id, c)| (*id, c.gather(idx)))
3668 .collect();
3669 (compacted, idx.len())
3670 }
3671 None => (user_columns, n),
3672 };
3673 self.advance_auto_inc_from_native_columns(&write_columns, write_n, live_before)?;
3674 let first = self.allocator.alloc_range(write_n as u64).0;
3675 for rid in first..first + write_n as u64 {
3676 self.reservoir.offer(rid);
3677 }
3678 let run_id = self.next_run_id;
3679 self.next_run_id += 1;
3680 let path = self.run_path(run_id);
3681 let mut writer =
3682 RunWriter::new(&self.schema, run_id as u128, epoch, 0).with_native_endian();
3683 if force_plain {
3684 writer = writer.with_plain();
3685 } else if lz4 {
3686 writer = writer.with_lz4();
3689 } else {
3690 writer = writer.with_zstd_level(zstd_level);
3691 }
3692 if let Some(kek) = &self.kek {
3693 writer = writer.with_encryption(kek.as_ref(), self.indexable_column_specs());
3694 }
3695 let header = writer.write_native(&path, &write_columns, write_n, first)?;
3696 self.run_refs.push(RunRef {
3697 run_id: run_id as u128,
3698 level: 0,
3699 epoch_created: epoch.0,
3700 row_count: header.row_count,
3701 });
3702 self.live_count = self.live_count.saturating_add(write_n as u64);
3703 if eager_index_build {
3704 let row_ids: Vec<u64> = (first..first + write_n as u64).collect();
3705 self.index_columns_bulk(&write_columns, &row_ids);
3706 self.indexes_complete = true;
3707 self.build_learned_ranges()?;
3708 } else {
3709 self.indexes_complete = false;
3713 }
3714 self.mark_flushed(epoch)?;
3715 self.persist_manifest(epoch)?;
3716 if eager_index_build {
3717 self.checkpoint_indexes(epoch);
3718 }
3719 self.clear_result_cache();
3720 Ok(epoch)
3721 }
3722
3723 fn index_columns_bulk(&mut self, columns: &[(u16, columnar::NativeColumn)], row_ids: &[u64]) {
3741 let n = row_ids.len();
3742 if n == 0 {
3743 return;
3744 }
3745 let by_id: std::collections::HashMap<u16, &columnar::NativeColumn> =
3746 columns.iter().map(|(id, c)| (*id, c)).collect();
3747 let ty_of: std::collections::HashMap<u16, TypeId> =
3748 self.schema.columns.iter().map(|c| (c.id, c.ty)).collect();
3749 let pk_id = self.schema.primary_key().map(|c| c.id);
3750
3751 for (i, &rid) in row_ids.iter().enumerate() {
3752 let row_id = RowId(rid);
3753 if let Some(pid) = pk_id {
3754 if let Some(col) = by_id.get(&pid) {
3755 let ty = ty_of.get(&pid).copied().unwrap_or(TypeId::Int64);
3756 if let Some(key) = bulk_index_key(&self.column_keys, pid, ty, col, i) {
3757 self.insert_hot_pk(key, row_id);
3758 }
3759 }
3760 }
3761 for idef in &self.schema.indexes {
3762 let Some(col) = by_id.get(&idef.column_id) else {
3763 continue;
3764 };
3765 let ty = ty_of.get(&idef.column_id).copied().unwrap_or(TypeId::Int64);
3766 match idef.kind {
3767 IndexKind::Bitmap => {
3768 if let Some(b) = self.bitmap.get_mut(&idef.column_id) {
3769 if let Some(key) =
3770 bulk_index_key(&self.column_keys, idef.column_id, ty, col, i)
3771 {
3772 b.insert(key, row_id);
3773 }
3774 }
3775 }
3776 IndexKind::FmIndex => {
3777 if let Some(f) = self.fm.get_mut(&idef.column_id) {
3778 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3779 f.insert(bytes.to_vec(), row_id);
3780 }
3781 }
3782 }
3783 IndexKind::Sparse => {
3784 if let Some(s) = self.sparse.get_mut(&idef.column_id) {
3785 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3786 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(bytes) {
3787 s.insert(&terms, row_id);
3788 }
3789 }
3790 }
3791 }
3792 IndexKind::MinHash => {
3793 if let Some(mh) = self.minhash.get_mut(&idef.column_id) {
3794 if let Some(bytes) = columnar::native_bytes_at(col, i) {
3795 let tokens = crate::index::token_hashes_from_bytes(bytes);
3796 mh.insert(&tokens, row_id);
3797 }
3798 }
3799 }
3800 _ => {}
3801 }
3802 }
3803 }
3804 }
3805
3806 pub fn visible_columns_native(
3811 &self,
3812 snapshot: Snapshot,
3813 projection: Option<&[u16]>,
3814 ) -> Result<Vec<(u16, columnar::NativeColumn)>> {
3815 let wanted: Vec<u16> = match projection {
3816 Some(p) => p.to_vec(),
3817 None => self.schema.columns.iter().map(|c| c.id).collect(),
3818 };
3819 if self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1 {
3820 let rr = self.run_refs[0].clone();
3821 let mut reader = self.open_reader(rr.run_id)?;
3822 let idxs = reader.visible_indices_native(snapshot.epoch)?;
3823 let all_visible = idxs.len() == reader.row_count();
3824 if reader.has_mmap() {
3830 use rayon::prelude::*;
3831 let valid: Vec<u16> = wanted
3834 .iter()
3835 .filter(|cid| self.schema.columns.iter().any(|c| c.id == **cid))
3836 .copied()
3837 .collect();
3838 let decoded: Vec<(u16, columnar::NativeColumn)> = valid
3840 .par_iter()
3841 .filter_map(|cid| {
3842 reader
3843 .column_native_shared(*cid)
3844 .ok()
3845 .map(|col| (*cid, col))
3846 })
3847 .collect();
3848 let cols = decoded
3849 .into_iter()
3850 .map(|(id, col)| (id, if all_visible { col } else { col.gather(&idxs) }))
3851 .collect();
3852 return Ok(cols);
3853 }
3854 let mut cols = Vec::with_capacity(wanted.len());
3855 for cid in &wanted {
3856 let cdef = match self.schema.columns.iter().find(|c| c.id == *cid) {
3857 Some(c) => c,
3858 None => continue,
3859 };
3860 let col = reader.column_native(cdef.id)?;
3861 cols.push((cdef.id, if all_visible { col } else { col.gather(&idxs) }));
3862 }
3863 return Ok(cols);
3864 }
3865 let vcols = self.visible_columns(snapshot)?;
3866 let want_set: std::collections::HashSet<u16> = wanted.iter().copied().collect();
3867 let out: Vec<(u16, columnar::NativeColumn)> = vcols
3868 .into_iter()
3869 .filter(|(id, _)| want_set.contains(id))
3870 .map(|(id, vals)| {
3871 let ty = self
3872 .schema
3873 .columns
3874 .iter()
3875 .find(|c| c.id == id)
3876 .map(|c| c.ty)
3877 .unwrap_or(TypeId::Bytes);
3878 (id, columnar::values_to_native(ty, &vals))
3879 })
3880 .collect();
3881 Ok(out)
3882 }
3883
3884 pub fn run_count(&self) -> usize {
3885 self.run_refs.len()
3886 }
3887
3888 pub fn memtable_is_empty(&self) -> bool {
3890 self.memtable.is_empty()
3891 }
3892
3893 pub fn page_cache_stats(&self) -> crate::cache::CacheStats {
3897 self.page_cache.lock().stats()
3898 }
3899
3900 pub fn reset_page_cache_stats(&self) {
3902 self.page_cache.lock().reset_stats();
3903 }
3904
3905 pub fn run_ids(&self) -> Vec<u128> {
3908 self.run_refs.iter().map(|r| r.run_id).collect()
3909 }
3910
3911 pub fn single_run_is_clean(&self) -> bool {
3915 if self.run_refs.len() != 1 {
3916 return false;
3917 }
3918 self.open_reader(self.run_refs[0].run_id)
3919 .map(|r| r.is_clean())
3920 .unwrap_or(false)
3921 }
3922
3923 fn resolve_footprint(
3930 &self,
3931 conditions: &[crate::query::Condition],
3932 _snapshot: Snapshot,
3933 ) -> roaring::RoaringBitmap {
3934 if !self.memtable.is_empty() || !self.mutable_run.is_empty() {
3935 return roaring::RoaringBitmap::new();
3936 }
3937 if self.run_refs.is_empty() {
3938 return roaring::RoaringBitmap::new();
3939 }
3940 if self.run_refs.len() == 1 {
3942 if let Ok(mut reader) = self.open_reader(self.run_refs[0].run_id) {
3943 if let Ok(rids) = self.resolve_survivor_rids(conditions, &mut reader) {
3944 return rids.to_roaring_lossy();
3945 }
3946 }
3947 }
3948 roaring::RoaringBitmap::new()
3949 }
3950
3951 pub fn query_columns_native_cached(
3962 &mut self,
3963 conditions: &[crate::query::Condition],
3964 projection: Option<&[u16]>,
3965 snapshot: Snapshot,
3966 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
3967 if conditions.is_empty() {
3968 return self.query_columns_native(conditions, projection, snapshot);
3969 }
3970 let key = crate::query::canonical_query_key(conditions, projection, snapshot.epoch.0);
3974 if let Some(hit) = self.result_cache.lock().get_columns(key) {
3975 crate::trace::QueryTrace::record(|t| {
3976 t.result_cache_hit = true;
3977 t.scan_mode = crate::trace::ScanMode::NativePushdown;
3978 });
3979 return Ok(Some((*hit).clone()));
3980 }
3981 let res = self.query_columns_native(conditions, projection, snapshot)?;
3982 if let Some(cols) = &res {
3983 let footprint = self.resolve_footprint(conditions, snapshot);
3984 let condition_cols = crate::query::condition_columns(conditions);
3985 self.result_cache.lock().insert(
3986 key,
3987 CachedEntry {
3988 data: CachedData::Columns(Arc::new(cols.clone())),
3989 footprint,
3990 condition_cols,
3991 },
3992 );
3993 }
3994 Ok(res)
3995 }
3996
3997 pub fn query_cached(&mut self, q: &crate::query::Query) -> Result<Vec<Row>> {
4002 if q.conditions.is_empty() {
4003 return self.query(q);
4004 }
4005 let key = crate::query::canonical_query_key(&q.conditions, None, 0);
4006 if let Some(hit) = self.result_cache.lock().get_rows(key) {
4007 crate::trace::QueryTrace::record(|t| {
4008 t.result_cache_hit = true;
4009 t.scan_mode = crate::trace::ScanMode::Materialized;
4010 });
4011 return Ok((*hit).clone());
4012 }
4013 let rows = self.query(q)?;
4014 let footprint = rows.iter().map(|r| r.row_id.0 as u32).collect();
4015 let condition_cols = crate::query::condition_columns(&q.conditions);
4016 self.result_cache.lock().insert(
4017 key,
4018 CachedEntry {
4019 data: CachedData::Rows(Arc::new(rows.clone())),
4020 footprint,
4021 condition_cols,
4022 },
4023 );
4024 Ok(rows)
4025 }
4026
4027 #[allow(clippy::type_complexity)]
4042 pub fn query_columns_native_traced(
4043 &mut self,
4044 conditions: &[crate::query::Condition],
4045 projection: Option<&[u16]>,
4046 snapshot: Snapshot,
4047 ) -> Result<(
4048 Option<Vec<(u16, columnar::NativeColumn)>>,
4049 crate::trace::QueryTrace,
4050 )> {
4051 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4052 self.query_columns_native(conditions, projection, snapshot)
4053 });
4054 Ok((result?, trace))
4055 }
4056
4057 #[allow(clippy::type_complexity)]
4060 pub fn query_columns_native_cached_traced(
4061 &mut self,
4062 conditions: &[crate::query::Condition],
4063 projection: Option<&[u16]>,
4064 snapshot: Snapshot,
4065 ) -> Result<(
4066 Option<Vec<(u16, columnar::NativeColumn)>>,
4067 crate::trace::QueryTrace,
4068 )> {
4069 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4070 self.query_columns_native_cached(conditions, projection, snapshot)
4071 });
4072 Ok((result?, trace))
4073 }
4074
4075 pub fn native_page_cursor_traced(
4077 &self,
4078 snapshot: Snapshot,
4079 projection: Vec<(u16, TypeId)>,
4080 conditions: &[crate::query::Condition],
4081 ) -> Result<(Option<NativePageCursor>, crate::trace::QueryTrace)> {
4082 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4083 self.native_page_cursor(snapshot, projection, conditions)
4084 });
4085 Ok((result?, trace))
4086 }
4087
4088 pub fn native_multi_run_cursor_traced(
4090 &self,
4091 snapshot: Snapshot,
4092 projection: Vec<(u16, TypeId)>,
4093 conditions: &[crate::query::Condition],
4094 ) -> Result<(
4095 Option<crate::cursor::MultiRunCursor>,
4096 crate::trace::QueryTrace,
4097 )> {
4098 let (result, trace) = crate::trace::QueryTrace::capture(|| {
4099 self.native_multi_run_cursor(snapshot, projection, conditions)
4100 });
4101 Ok((result?, trace))
4102 }
4103
4104 pub fn count_conditions_traced(
4106 &mut self,
4107 conditions: &[crate::query::Condition],
4108 snapshot: Snapshot,
4109 ) -> Result<(Option<u64>, crate::trace::QueryTrace)> {
4110 let (result, trace) =
4111 crate::trace::QueryTrace::capture(|| self.count_conditions(conditions, snapshot));
4112 Ok((result?, trace))
4113 }
4114
4115 pub fn query_traced(
4117 &mut self,
4118 q: &crate::query::Query,
4119 ) -> Result<(Vec<Row>, crate::trace::QueryTrace)> {
4120 let (result, trace) = crate::trace::QueryTrace::capture(|| self.query(q));
4121 Ok((result?, trace))
4122 }
4123
4124 pub fn query_columns_native(
4129 &mut self,
4130 conditions: &[crate::query::Condition],
4131 projection: Option<&[u16]>,
4132 snapshot: Snapshot,
4133 ) -> Result<Option<Vec<(u16, columnar::NativeColumn)>>> {
4134 use crate::query::Condition;
4135 if conditions.is_empty() {
4136 return Ok(None);
4137 }
4138 self.ensure_indexes_complete()?;
4139
4140 let served = |c: &Condition| {
4145 matches!(
4146 c,
4147 Condition::Pk(_)
4148 | Condition::BitmapEq { .. }
4149 | Condition::BitmapIn { .. }
4150 | Condition::FmContains { .. }
4151 | Condition::FmContainsAll { .. }
4152 | Condition::Ann { .. }
4153 | Condition::Range { .. }
4154 | Condition::RangeF64 { .. }
4155 | Condition::SparseMatch { .. }
4156 | Condition::MinHashSimilar { .. }
4157 | Condition::IsNull { .. }
4158 | Condition::IsNotNull { .. }
4159 )
4160 };
4161 if !conditions.iter().all(served) {
4162 return Ok(None);
4163 }
4164 let fast_path =
4165 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
4166 crate::trace::QueryTrace::record(|t| {
4167 t.run_count = self.run_refs.len();
4168 t.memtable_rows = self.memtable.len();
4169 t.mutable_run_rows = self.mutable_run.len();
4170 t.conditions_pushed = conditions.len();
4171 t.learned_range_used = conditions.iter().any(|c| match c {
4172 Condition::Range { column_id, .. } | Condition::RangeF64 { column_id, .. } => {
4173 self.learned_range.contains_key(column_id)
4174 }
4175 _ => false,
4176 });
4177 });
4178 let col_ids: Vec<u16> = projection
4180 .map(|p| p.to_vec())
4181 .unwrap_or_else(|| self.schema.columns.iter().map(|c| c.id).collect());
4182 let proj_pairs: Vec<(u16, TypeId)> = col_ids
4183 .iter()
4184 .map(|&cid| {
4185 let ty = self
4186 .schema
4187 .columns
4188 .iter()
4189 .find(|c| c.id == cid)
4190 .map(|c| c.ty)
4191 .unwrap_or(TypeId::Bytes);
4192 (cid, ty)
4193 })
4194 .collect();
4195
4196 if fast_path {
4202 let needs_column = conditions.iter().any(|c| match c {
4205 Condition::Range { column_id, .. } => !self.learned_range.contains_key(column_id),
4206 Condition::RangeF64 { column_id, .. } => {
4207 !self.learned_range.contains_key(column_id)
4208 }
4209 _ => false,
4210 });
4211 let mut reader_opt: Option<RunReader> = if needs_column {
4212 Some(self.open_reader(self.run_refs[0].run_id)?)
4213 } else {
4214 None
4215 };
4216 let mut sets: Vec<RowIdSet> = Vec::new();
4217 for c in conditions {
4218 let s = match c {
4219 Condition::Range { column_id, lo, hi }
4220 if !self.learned_range.contains_key(column_id) =>
4221 {
4222 if reader_opt.is_none() {
4223 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4224 }
4225 reader_opt
4226 .as_mut()
4227 .expect("reader opened for range")
4228 .range_row_id_set_i64(*column_id, *lo, *hi)?
4229 }
4230 Condition::RangeF64 {
4231 column_id,
4232 lo,
4233 lo_inclusive,
4234 hi,
4235 hi_inclusive,
4236 } if !self.learned_range.contains_key(column_id) => {
4237 if reader_opt.is_none() {
4238 reader_opt = Some(self.open_reader(self.run_refs[0].run_id)?);
4239 }
4240 reader_opt
4241 .as_mut()
4242 .expect("reader opened for range")
4243 .range_row_id_set_f64(
4244 *column_id,
4245 *lo,
4246 *lo_inclusive,
4247 *hi,
4248 *hi_inclusive,
4249 )?
4250 }
4251 _ => self.resolve_condition(c, snapshot)?,
4252 };
4253 sets.push(s);
4254 }
4255 let candidates = RowIdSet::intersect_many(sets);
4256 crate::trace::QueryTrace::record(|t| {
4257 t.survivor_count = Some(candidates.len());
4258 });
4259 if candidates.is_empty() {
4260 let cols: Vec<(u16, columnar::NativeColumn)> = col_ids
4261 .iter()
4262 .map(|&id| {
4263 (
4264 id,
4265 columnar::null_native(
4266 proj_pairs
4267 .iter()
4268 .find(|(c, _)| c == &id)
4269 .map(|(_, t)| *t)
4270 .unwrap_or(TypeId::Bytes),
4271 0,
4272 ),
4273 )
4274 })
4275 .collect();
4276 return Ok(Some(cols));
4277 }
4278 let mut reader = match reader_opt.take() {
4279 Some(r) => r,
4280 None => self.open_reader(self.run_refs[0].run_id)?,
4281 };
4282 let candidate_ids = candidates.into_sorted_vec();
4283 let (positions, fast_rid) = if let Some(positions) =
4284 reader.positions_for_row_ids_fast(&candidate_ids)
4285 {
4286 (positions, true)
4287 } else {
4288 let col = reader.column_native(crate::sorted_run::SYS_ROW_ID)?;
4289 match col {
4290 columnar::NativeColumn::Int64 { data, .. } => {
4291 let mut p: Vec<usize> = candidate_ids
4292 .iter()
4293 .filter_map(|rid| data.binary_search(&(*rid as i64)).ok())
4294 .collect();
4295 p.sort_unstable();
4296 (p, false)
4297 }
4298 _ => return Err(MongrelError::InvalidArgument("sys row_id not int64".into())),
4299 }
4300 };
4301 crate::trace::QueryTrace::record(|t| {
4302 t.scan_mode = crate::trace::ScanMode::NativePushdown;
4303 t.fast_row_id_map = fast_rid;
4304 });
4305 let mut cols = Vec::with_capacity(col_ids.len());
4306 for cid in &col_ids {
4307 let col = reader.column_native(*cid)?;
4308 cols.push((*cid, col.gather(&positions)));
4309 }
4310 return Ok(Some(cols));
4311 }
4312
4313 if !self.run_refs.is_empty() {
4326 use crate::cursor::{drain_cursor_to_columns, Cursor};
4327 let remaining: usize;
4328 let mut cursor: Box<dyn crate::cursor::Cursor> = if self.run_refs.len() == 1 {
4329 let c = self
4330 .native_page_cursor(snapshot, proj_pairs.clone(), conditions)?
4331 .expect("single-run cursor should build when run_refs.len() == 1");
4332 remaining = c.remaining_rows();
4333 Box::new(c)
4334 } else {
4335 let c = self
4336 .native_multi_run_cursor(snapshot, proj_pairs.clone(), conditions)?
4337 .expect("multi-run cursor should build when run_refs.len() >= 1");
4338 remaining = c.remaining_rows();
4339 Box::new(c)
4340 };
4341 crate::trace::QueryTrace::record(|t| {
4342 if t.survivor_count.is_none() {
4343 t.survivor_count = Some(remaining);
4344 }
4345 });
4346 let cols = drain_cursor_to_columns(cursor.as_mut(), &proj_pairs)?;
4347 return Ok(Some(cols));
4348 }
4349
4350 crate::trace::QueryTrace::record(|t| {
4355 t.scan_mode = crate::trace::ScanMode::Materialized;
4356 t.row_materialized = true;
4357 });
4358 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4359 for c in conditions {
4360 sets.push(self.resolve_condition(c, snapshot)?);
4361 }
4362 let rids = RowIdSet::intersect_many(sets).into_sorted_vec();
4363 let rows = self.rows_for_rids(&rids, snapshot)?;
4364 let mut cols: Vec<(u16, columnar::NativeColumn)> = Vec::with_capacity(col_ids.len());
4365 for (cid, ty) in &proj_pairs {
4366 let vals: Vec<Value> = rows
4367 .iter()
4368 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4369 .collect();
4370 cols.push((*cid, columnar::values_to_native(*ty, &vals)));
4371 }
4372 Ok(Some(cols))
4373 }
4374
4375 pub fn native_page_cursor(
4390 &self,
4391 snapshot: Snapshot,
4392 projection: Vec<(u16, TypeId)>,
4393 conditions: &[crate::query::Condition],
4394 ) -> Result<Option<NativePageCursor>> {
4395 use crate::cursor::build_page_plans;
4396 if self.run_refs.len() != 1 {
4397 return Ok(None);
4398 }
4399 let mut reader = self.open_reader(self.run_refs[0].run_id)?;
4400 let (positions, rids) = reader.visible_positions_with_rids(snapshot.epoch)?;
4401
4402 let overlay_rids: HashSet<u64> = {
4405 let mut s = HashSet::new();
4406 for row in self.memtable.visible_versions(snapshot.epoch) {
4407 s.insert(row.row_id.0);
4408 }
4409 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4410 s.insert(row.row_id.0);
4411 }
4412 s
4413 };
4414
4415 let survivors = if conditions.is_empty() {
4419 None
4420 } else {
4421 Some(self.resolve_survivor_rids(conditions, &mut reader)?)
4422 };
4423
4424 let run_survivors: Option<RowIdSet> = if overlay_rids.is_empty() {
4431 survivors.clone()
4432 } else if let Some(s) = &survivors {
4433 let mut run_set = s.clone();
4434 run_set.remove_many(overlay_rids.iter().copied());
4435 Some(run_set)
4436 } else {
4437 Some(RowIdSet::from_unsorted(
4438 rids.iter()
4439 .map(|&r| r as u64)
4440 .filter(|r| !overlay_rids.contains(r))
4441 .collect(),
4442 ))
4443 };
4444
4445 let overlay_rows = if overlay_rids.is_empty() {
4446 Vec::new()
4447 } else {
4448 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4449 self.overlay_visible_rows(snapshot, bound)
4450 };
4451
4452 let plans = if positions.is_empty() {
4454 Vec::new()
4455 } else {
4456 let page_rows = reader.page_row_counts(crate::sorted_run::SYS_ROW_ID)?;
4457 build_page_plans(&positions, &rids, &page_rows, run_survivors.as_ref())
4458 };
4459
4460 let overlay = if overlay_rows.is_empty() {
4462 None
4463 } else {
4464 let filtered =
4465 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4466 if filtered.is_empty() {
4467 None
4468 } else {
4469 Some(self.materialize_overlay(&filtered, &projection))
4470 }
4471 };
4472
4473 let overlay_row_count = overlay
4474 .as_ref()
4475 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4476 .unwrap_or(0);
4477 crate::trace::QueryTrace::record(|t| {
4478 t.scan_mode = crate::trace::ScanMode::NativePageCursor;
4479 t.run_count = self.run_refs.len();
4480 t.memtable_rows = self.memtable.len();
4481 t.mutable_run_rows = self.mutable_run.len();
4482 t.overlay_rows = overlay_row_count;
4483 t.conditions_pushed = conditions.len();
4484 t.pages_decoded = plans
4485 .iter()
4486 .map(|p| p.positions.len())
4487 .sum::<usize>()
4488 .min(1);
4489 });
4490
4491 Ok(Some(NativePageCursor::new_with_overlay(
4492 reader, projection, plans, overlay,
4493 )))
4494 }
4495 #[allow(clippy::type_complexity)]
4505 pub fn native_multi_run_cursor(
4506 &self,
4507 snapshot: Snapshot,
4508 projection: Vec<(u16, TypeId)>,
4509 conditions: &[crate::query::Condition],
4510 ) -> Result<Option<crate::cursor::MultiRunCursor>> {
4511 use crate::cursor::{MultiRunCursor, RunStream};
4512 use crate::sorted_run::SYS_ROW_ID;
4513 use std::collections::{BinaryHeap, HashMap, HashSet};
4514 if self.run_refs.is_empty() {
4515 return Ok(None);
4516 }
4517
4518 let mut run_meta: Vec<(RunReader, Vec<i64>, Vec<i64>, Vec<u8>, Vec<usize>)> =
4520 Vec::with_capacity(self.run_refs.len());
4521 for rr in &self.run_refs {
4522 let mut reader = self.open_reader(rr.run_id)?;
4523 let (rids, eps, del) = reader.system_columns_native()?;
4524 let page_rows = reader.page_row_counts(SYS_ROW_ID)?;
4525 run_meta.push((reader, rids, eps, del, page_rows));
4526 }
4527
4528 let mut best: HashMap<u64, (u64, usize, usize, bool)> = HashMap::new();
4532 for (run_idx, (_, rids, eps, del, _)) in run_meta.iter().enumerate() {
4533 for i in 0..rids.len() {
4534 let rid = rids[i] as u64;
4535 let e = eps[i] as u64;
4536 if e > snapshot.epoch.0 {
4537 continue;
4538 }
4539 let is_del = del[i] != 0;
4540 best.entry(rid)
4541 .and_modify(|cur| {
4542 if e > cur.0 {
4543 *cur = (e, run_idx, i, is_del);
4544 }
4545 })
4546 .or_insert((e, run_idx, i, is_del));
4547 }
4548 }
4549
4550 let overlay_rids: HashSet<u64> = {
4552 let mut s = HashSet::new();
4553 for row in self.memtable.visible_versions(snapshot.epoch) {
4554 s.insert(row.row_id.0);
4555 }
4556 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4557 s.insert(row.row_id.0);
4558 }
4559 s
4560 };
4561
4562 let survivors: Option<RowIdSet> = if conditions.is_empty() {
4564 None
4565 } else {
4566 let mut sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4567 for c in conditions {
4568 sets.push(self.resolve_condition(c, snapshot)?);
4569 }
4570 Some(RowIdSet::intersect_many(sets))
4571 };
4572
4573 let mut per_run: Vec<Vec<(u64, usize)>> = vec![Vec::new(); run_meta.len()];
4577 for (rid, (_, run_idx, pos, deleted)) in &best {
4578 if *deleted {
4579 continue;
4580 }
4581 if overlay_rids.contains(rid) {
4582 continue;
4583 }
4584 if let Some(s) = &survivors {
4585 if !s.contains(*rid) {
4586 continue;
4587 }
4588 }
4589 per_run[*run_idx].push((*rid, *pos));
4590 }
4591 for v in per_run.iter_mut() {
4592 v.sort_unstable_by_key(|&(rid, _)| rid);
4593 }
4594
4595 let mut streams = Vec::with_capacity(run_meta.len());
4597 let mut heap: BinaryHeap<std::cmp::Reverse<(u64, usize)>> = BinaryHeap::new();
4598 let mut total = 0usize;
4599 for (run_idx, (reader, _, _, _, page_rows)) in run_meta.into_iter().enumerate() {
4600 let mut starts = Vec::with_capacity(page_rows.len());
4601 let mut acc = 0usize;
4602 for &r in &page_rows {
4603 starts.push(acc);
4604 acc += r;
4605 }
4606 let mut survivors_vec: Vec<(u64, usize, usize)> =
4607 Vec::with_capacity(per_run[run_idx].len());
4608 for &(rid, pos) in &per_run[run_idx] {
4609 let page_seq = match starts.partition_point(|&s| s <= pos) {
4610 0 => continue,
4611 p => p - 1,
4612 };
4613 let within = pos - starts[page_seq];
4614 survivors_vec.push((rid, page_seq, within));
4615 }
4616 total += survivors_vec.len();
4617 if let Some(&(rid, _, _)) = survivors_vec.first() {
4618 heap.push(std::cmp::Reverse((rid, run_idx)));
4619 }
4620 streams.push(RunStream::new(reader, survivors_vec, page_rows));
4621 }
4622
4623 let overlay_rows = if overlay_rids.is_empty() {
4625 Vec::new()
4626 } else {
4627 let bound = Self::overlay_materialization_bound(conditions, &survivors);
4628 self.overlay_visible_rows(snapshot, bound)
4629 };
4630 let overlay = if overlay_rows.is_empty() {
4631 None
4632 } else {
4633 let filtered =
4634 self.filter_overlay_rows(overlay_rows, conditions, survivors.as_ref(), snapshot)?;
4635 if filtered.is_empty() {
4636 None
4637 } else {
4638 Some(self.materialize_overlay(&filtered, &projection))
4639 }
4640 };
4641
4642 let overlay_row_count = overlay
4643 .as_ref()
4644 .map(|c| c.first().map(|c| c.len()).unwrap_or(0))
4645 .unwrap_or(0);
4646 crate::trace::QueryTrace::record(|t| {
4647 t.scan_mode = crate::trace::ScanMode::MultiRunCursor;
4648 t.run_count = self.run_refs.len();
4649 t.memtable_rows = self.memtable.len();
4650 t.mutable_run_rows = self.mutable_run.len();
4651 t.overlay_rows = overlay_row_count;
4652 t.conditions_pushed = conditions.len();
4653 t.survivor_count = Some(total);
4654 });
4655
4656 Ok(Some(MultiRunCursor::new(
4657 streams, projection, heap, total, overlay,
4658 )))
4659 }
4660
4661 fn overlay_materialization_bound<'a>(
4673 conditions: &[crate::query::Condition],
4674 survivors: &'a Option<RowIdSet>,
4675 ) -> Option<&'a RowIdSet> {
4676 use crate::query::Condition;
4677 let has_range = conditions
4678 .iter()
4679 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
4680 if has_range {
4681 None
4682 } else {
4683 survivors.as_ref()
4684 }
4685 }
4686
4687 fn overlay_visible_rows(&self, snapshot: Snapshot, bound: Option<&RowIdSet>) -> Vec<Row> {
4699 let mut best: HashMap<u64, (Epoch, Row)> = HashMap::new();
4700 let mut fold = |row: Row| {
4701 if let Some(b) = bound {
4702 if !b.contains(row.row_id.0) {
4703 return;
4704 }
4705 }
4706 best.entry(row.row_id.0)
4707 .and_modify(|(be, br)| {
4708 if row.committed_epoch > *be {
4709 *be = row.committed_epoch;
4710 *br = row.clone();
4711 }
4712 })
4713 .or_insert_with(|| (row.committed_epoch, row));
4714 };
4715 for row in self.memtable.visible_versions(snapshot.epoch) {
4716 fold(row);
4717 }
4718 for row in self.mutable_run.visible_versions(snapshot.epoch) {
4719 fold(row);
4720 }
4721 let mut out: Vec<Row> = best
4722 .into_values()
4723 .filter_map(|(_, r)| if r.deleted { None } else { Some(r) })
4724 .collect();
4725 out.sort_by_key(|r| r.row_id);
4726 out
4727 }
4728
4729 fn filter_overlay_rows(
4737 &self,
4738 rows: Vec<Row>,
4739 conditions: &[crate::query::Condition],
4740 survivors: Option<&RowIdSet>,
4741 snapshot: Snapshot,
4742 ) -> Result<Vec<Row>> {
4743 if conditions.is_empty() {
4744 return Ok(rows);
4745 }
4746 use crate::query::Condition;
4747 let all_index_served = !conditions
4751 .iter()
4752 .any(|c| matches!(c, Condition::Range { .. } | Condition::RangeF64 { .. }));
4753 if all_index_served {
4754 return Ok(rows
4755 .into_iter()
4756 .filter(|r| survivors.map_or(true, |s| s.contains(r.row_id.0)))
4757 .collect());
4758 }
4759 let mut per_cond_sets: Vec<RowIdSet> = Vec::with_capacity(conditions.len());
4762 for c in conditions {
4763 let s = match c {
4764 Condition::Range { .. } | Condition::RangeF64 { .. } => RowIdSet::empty(),
4765 _ => self.resolve_condition(c, snapshot)?,
4766 };
4767 per_cond_sets.push(s);
4768 }
4769 Ok(rows
4770 .into_iter()
4771 .filter(|row| {
4772 conditions.iter().enumerate().all(|(i, c)| match c {
4773 Condition::Range { column_id, lo, hi } => {
4774 matches!(row.columns.get(column_id), Some(Value::Int64(v)) if *v >= *lo && *v <= *hi)
4775 }
4776 Condition::RangeF64 { column_id, lo, lo_inclusive, hi, hi_inclusive } => {
4777 match row.columns.get(column_id) {
4778 Some(Value::Float64(v)) => {
4779 let lo_ok = if *lo_inclusive { *v >= *lo } else { *v > *lo };
4780 let hi_ok = if *hi_inclusive { *v <= *hi } else { *v < *hi };
4781 lo_ok && hi_ok
4782 }
4783 _ => false,
4784 }
4785 }
4786 _ => per_cond_sets[i].contains(row.row_id.0),
4787 })
4788 })
4789 .collect())
4790 }
4791
4792 fn materialize_overlay(
4795 &self,
4796 rows: &[Row],
4797 projection: &[(u16, TypeId)],
4798 ) -> Vec<columnar::NativeColumn> {
4799 if projection.is_empty() {
4800 return vec![columnar::null_native(TypeId::Int64, rows.len())];
4801 }
4802 let mut cols = Vec::with_capacity(projection.len());
4803 for (cid, ty) in projection {
4804 let vals: Vec<Value> = rows
4805 .iter()
4806 .map(|r| r.columns.get(cid).cloned().unwrap_or(Value::Null))
4807 .collect();
4808 cols.push(columnar::values_to_native(*ty, &vals));
4809 }
4810 cols
4811 }
4812
4813 fn resolve_survivor_rids(
4818 &self,
4819 conditions: &[crate::query::Condition],
4820 reader: &mut RunReader,
4821 ) -> Result<RowIdSet> {
4822 use crate::query::Condition;
4823 let mut sets: Vec<RowIdSet> = Vec::new();
4824 for c in conditions {
4825 let s: RowIdSet = match c {
4826 Condition::Pk(key) => {
4827 let lookup = self
4828 .schema
4829 .primary_key()
4830 .map(|pk| self.index_lookup_key_bytes(pk.id, key))
4831 .unwrap_or_else(|| key.clone());
4832 self.hot
4833 .get(&lookup)
4834 .map(|r| RowIdSet::one(r.0))
4835 .unwrap_or_else(RowIdSet::empty)
4836 }
4837 Condition::BitmapEq { column_id, value } => {
4838 let lookup = self.index_lookup_key_bytes(*column_id, value);
4839 self.bitmap
4840 .get(column_id)
4841 .map(|b| RowIdSet::from_roaring(b.get(&lookup)))
4842 .unwrap_or_else(RowIdSet::empty)
4843 }
4844 Condition::BitmapIn { column_id, values } => {
4845 let bm = self.bitmap.get(column_id);
4846 let mut acc = roaring::RoaringBitmap::new();
4847 if let Some(b) = bm {
4848 for v in values {
4849 let lookup = self.index_lookup_key_bytes(*column_id, v);
4850 acc |= b.get(&lookup);
4851 }
4852 }
4853 RowIdSet::from_roaring(acc)
4854 }
4855 Condition::FmContains { column_id, pattern } => self
4856 .fm
4857 .get(column_id)
4858 .map(|f| {
4859 RowIdSet::from_unsorted(
4860 f.locate(pattern).into_iter().map(|r| r.0).collect(),
4861 )
4862 })
4863 .unwrap_or_else(RowIdSet::empty),
4864 Condition::FmContainsAll {
4865 column_id,
4866 patterns,
4867 } => {
4868 if let Some(f) = self.fm.get(column_id) {
4869 let sets: Vec<RowIdSet> = patterns
4870 .iter()
4871 .map(|pat| {
4872 RowIdSet::from_unsorted(
4873 f.locate(pat).into_iter().map(|r| r.0).collect(),
4874 )
4875 })
4876 .collect();
4877 RowIdSet::intersect_many(sets)
4878 } else {
4879 RowIdSet::empty()
4880 }
4881 }
4882 Condition::Ann {
4883 column_id,
4884 query,
4885 k,
4886 } => self
4887 .ann
4888 .get(column_id)
4889 .map(|a| {
4890 RowIdSet::from_unsorted(
4891 a.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
4892 )
4893 })
4894 .unwrap_or_else(RowIdSet::empty),
4895 Condition::SparseMatch {
4896 column_id,
4897 query,
4898 k,
4899 } => self
4900 .sparse
4901 .get(column_id)
4902 .map(|s| {
4903 RowIdSet::from_unsorted(
4904 s.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
4905 )
4906 })
4907 .unwrap_or_else(RowIdSet::empty),
4908 Condition::MinHashSimilar {
4909 column_id,
4910 query,
4911 k,
4912 } => self
4913 .minhash
4914 .get(column_id)
4915 .map(|mh| {
4916 RowIdSet::from_unsorted(
4917 mh.search(query, *k).into_iter().map(|(r, _)| r.0).collect(),
4918 )
4919 })
4920 .unwrap_or_else(RowIdSet::empty),
4921 Condition::Range { column_id, lo, hi } => {
4922 if let Some(li) = self.learned_range.get(column_id) {
4923 RowIdSet::from_unsorted(li.range(*lo, *hi).into_iter().collect())
4924 } else {
4925 reader.range_row_id_set_i64(*column_id, *lo, *hi)?
4926 }
4927 }
4928 Condition::RangeF64 {
4929 column_id,
4930 lo,
4931 lo_inclusive,
4932 hi,
4933 hi_inclusive,
4934 } => {
4935 if let Some(li) = self.learned_range.get(column_id) {
4936 RowIdSet::from_unsorted(
4937 li.range_f64(*lo, *lo_inclusive, *hi, *hi_inclusive)
4938 .into_iter()
4939 .collect(),
4940 )
4941 } else {
4942 reader.range_row_id_set_f64(
4943 *column_id,
4944 *lo,
4945 *lo_inclusive,
4946 *hi,
4947 *hi_inclusive,
4948 )?
4949 }
4950 }
4951 Condition::IsNull { column_id } => reader.null_row_id_set(*column_id, true)?,
4952 Condition::IsNotNull { column_id } => reader.null_row_id_set(*column_id, false)?,
4953 };
4954 sets.push(s);
4955 }
4956 Ok(RowIdSet::intersect_many(sets))
4957 }
4958
4959 pub fn scan_cursor(
4980 &self,
4981 snapshot: Snapshot,
4982 projection: Vec<(u16, TypeId)>,
4983 conditions: &[crate::query::Condition],
4984 ) -> Result<Option<Box<dyn crate::cursor::Cursor>>> {
4985 if self.run_refs.len() == 1 {
4986 Ok(self
4987 .native_page_cursor(snapshot, projection, conditions)?
4988 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
4989 } else {
4990 Ok(self
4991 .native_multi_run_cursor(snapshot, projection, conditions)?
4992 .map(|c| Box::new(c) as Box<dyn crate::cursor::Cursor>))
4993 }
4994 }
4995
4996 pub fn aggregate_native(
5010 &self,
5011 snapshot: Snapshot,
5012 column: Option<u16>,
5013 conditions: &[crate::query::Condition],
5014 agg: NativeAgg,
5015 ) -> Result<Option<NativeAggResult>> {
5016 if self.run_refs.len() == 1 && conditions.is_empty() {
5018 if let Some(res) = self.aggregate_from_stats(snapshot, column, agg)? {
5019 return Ok(Some(res));
5020 }
5021 }
5022 if matches!(agg, NativeAgg::Count) && column.is_none() {
5024 return Ok(self
5025 .scan_cursor(snapshot, Vec::new(), conditions)?
5026 .map(|c| NativeAggResult::Count(c.remaining_rows() as u64)));
5027 }
5028 let cid = match column {
5031 Some(c) => c,
5032 None => return Ok(None),
5033 };
5034 let ty = self.column_type(cid);
5035 let Some(mut cursor) = self.scan_cursor(snapshot, vec![(cid, ty)], conditions)? else {
5036 return Ok(None);
5037 };
5038 match ty {
5039 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5040 let (count, sum, mn, mx) = accumulate_int(cursor.as_mut())?;
5041 Ok(Some(pack_int(agg, count, sum, mn, mx)))
5042 }
5043 TypeId::Float64 => {
5044 let (count, sum, mn, mx) = accumulate_float(cursor.as_mut())?;
5045 Ok(Some(pack_float(agg, count, sum, mn, mx)))
5046 }
5047 _ => Ok(None),
5048 }
5049 }
5050
5051 fn aggregate_from_stats(
5059 &self,
5060 snapshot: Snapshot,
5061 column: Option<u16>,
5062 agg: NativeAgg,
5063 ) -> Result<Option<NativeAggResult>> {
5064 let cid = match (agg, column) {
5065 (NativeAgg::Count | NativeAgg::Min | NativeAgg::Max, Some(c)) => c,
5066 _ => return Ok(None), };
5068 let Some(stats) = self.exact_column_stats(snapshot, &[cid])? else {
5069 return Ok(None);
5070 };
5071 let Some(cs) = stats.get(&cid) else {
5072 return Ok(None);
5073 };
5074 match agg {
5075 NativeAgg::Count => Ok(Some(NativeAggResult::Count(
5077 self.live_count.saturating_sub(cs.null_count),
5078 ))),
5079 NativeAgg::Min | NativeAgg::Max => {
5080 let bound = if agg == NativeAgg::Min {
5081 &cs.min
5082 } else {
5083 &cs.max
5084 };
5085 match bound {
5086 Some(Value::Int64(x)) => Ok(Some(NativeAggResult::Int(*x))),
5087 Some(Value::Float64(x)) => Ok(Some(NativeAggResult::Float(*x))),
5088 Some(_) => Ok(None), None if cs.null_count >= self.live_count => Ok(Some(NativeAggResult::Null)),
5093 None => Ok(None),
5094 }
5095 }
5096 _ => Ok(None),
5097 }
5098 }
5099
5100 pub fn count_distinct_from_bitmap(&self, column_id: u16) -> Result<Option<u64>> {
5109 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5110 return Ok(None);
5111 }
5112 let reader = self.open_reader(self.run_refs[0].run_id)?;
5113 if self.live_count != reader.row_count() as u64 {
5114 return Ok(None);
5115 }
5116 let Some(bm) = self.bitmap.get(&column_id) else {
5117 return Ok(None); };
5119 let mut distinct = bm.value_count() as u64;
5120 if !bm.get(&Value::Null.encode_key()).is_empty() {
5123 distinct = distinct.saturating_sub(1);
5124 }
5125 Ok(Some(distinct))
5126 }
5127
5128 pub fn aggregate_incremental(
5140 &mut self,
5141 cache_key: u64,
5142 conditions: &[crate::query::Condition],
5143 column: Option<u16>,
5144 agg: NativeAgg,
5145 ) -> Result<IncrementalAggResult> {
5146 let snap = self.snapshot();
5147 let cur_wm = self.allocator.current().0;
5148 let cur_epoch = snap.epoch.0;
5149 let incremental_ok =
5156 !self.had_deletes && self.memtable.is_empty() && self.mutable_run.is_empty();
5157
5158 if incremental_ok {
5161 if let Some(cached) = self.agg_cache.get(&cache_key).cloned() {
5162 if cached.epoch == cur_epoch {
5163 return Ok(IncrementalAggResult {
5164 state: cached.state,
5165 incremental: true,
5166 delta_rows: 0,
5167 });
5168 }
5169 if cached.epoch < cur_epoch && cached.watermark <= cur_wm {
5170 let delta_rids: Vec<u64> = (cached.watermark..cur_wm).collect();
5171 let delta_rows = self.rows_for_rids(&delta_rids, snap)?;
5172 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5173 let delta_state = agg_state_from_rows(
5174 &delta_rows,
5175 conditions,
5176 &index_sets,
5177 column,
5178 agg,
5179 &self.schema,
5180 )?;
5181 let merged = cached.state.merge(delta_state);
5182 let delta_n = delta_rids.len() as u64;
5183 self.agg_cache.insert(
5184 cache_key,
5185 CachedAgg {
5186 state: merged.clone(),
5187 watermark: cur_wm,
5188 epoch: cur_epoch,
5189 },
5190 );
5191 return Ok(IncrementalAggResult {
5192 state: merged,
5193 incremental: true,
5194 delta_rows: delta_n,
5195 });
5196 }
5197 }
5198 }
5199
5200 let cursor_ok =
5205 self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1;
5206 let state = if cursor_ok && agg != NativeAgg::Avg {
5207 match self.aggregate_native(snap, column, conditions, agg)? {
5208 Some(result) => {
5209 AggState::from_native(result, agg, column.map(|c| self.column_type(c)))
5210 }
5211 None => self.agg_state_full_scan(conditions, column, agg, snap)?,
5212 }
5213 } else {
5214 self.agg_state_full_scan(conditions, column, agg, snap)?
5215 };
5216 if incremental_ok {
5218 self.agg_cache.insert(
5219 cache_key,
5220 CachedAgg {
5221 state: state.clone(),
5222 watermark: cur_wm,
5223 epoch: cur_epoch,
5224 },
5225 );
5226 }
5227 Ok(IncrementalAggResult {
5228 state,
5229 incremental: false,
5230 delta_rows: 0,
5231 })
5232 }
5233
5234 fn agg_state_full_scan(
5237 &self,
5238 conditions: &[crate::query::Condition],
5239 column: Option<u16>,
5240 agg: NativeAgg,
5241 snap: Snapshot,
5242 ) -> Result<AggState> {
5243 let rows = self.visible_rows(snap)?;
5244 let index_sets = self.resolve_index_conditions(conditions, snap)?;
5245 agg_state_from_rows(&rows, conditions, &index_sets, column, agg, &self.schema)
5246 }
5247
5248 fn resolve_index_conditions(
5251 &self,
5252 conditions: &[crate::query::Condition],
5253 snapshot: Snapshot,
5254 ) -> Result<Vec<RowIdSet>> {
5255 use crate::query::Condition;
5256 let mut sets = Vec::new();
5257 for c in conditions {
5258 if matches!(
5259 c,
5260 Condition::Ann { .. }
5261 | Condition::SparseMatch { .. }
5262 | Condition::MinHashSimilar { .. }
5263 ) {
5264 sets.push(self.resolve_condition(c, snapshot)?);
5265 }
5266 }
5267 Ok(sets)
5268 }
5269
5270 fn column_type(&self, cid: u16) -> TypeId {
5271 self.schema
5272 .columns
5273 .iter()
5274 .find(|c| c.id == cid)
5275 .map(|c| c.ty)
5276 .unwrap_or(TypeId::Bytes)
5277 }
5278
5279 pub fn approx_aggregate(
5288 &self,
5289 conditions: &[crate::query::Condition],
5290 column: Option<u16>,
5291 agg: ApproxAgg,
5292 z: f64,
5293 ) -> Result<Option<ApproxResult>> {
5294 use crate::query::Condition;
5295 let snapshot = self.snapshot();
5296 let n_pop = self.live_count;
5297 let sample_rids: Vec<u64> = self.reservoir.row_ids().to_vec();
5298 if sample_rids.is_empty() {
5299 return Ok(None);
5300 }
5301 let live_sample = self.rows_for_rids(&sample_rids, snapshot)?;
5303 let s = live_sample.len();
5304 if s == 0 {
5305 return Ok(None);
5306 }
5307
5308 let mut index_sets: Vec<RowIdSet> = Vec::new();
5311 for c in conditions {
5312 if matches!(
5313 c,
5314 Condition::Ann { .. }
5315 | Condition::SparseMatch { .. }
5316 | Condition::MinHashSimilar { .. }
5317 ) {
5318 index_sets.push(self.resolve_condition(c, snapshot)?);
5319 }
5320 }
5321
5322 let cid = match (agg, column) {
5324 (ApproxAgg::Count, _) => None,
5325 (_, Some(c)) => Some(c),
5326 _ => return Ok(None),
5327 };
5328 let mut passing_vals: Vec<f64> = Vec::with_capacity(s);
5329 for r in &live_sample {
5330 if !conditions
5332 .iter()
5333 .all(|c| condition_matches_row(c, r, &self.schema))
5334 {
5335 continue;
5336 }
5337 if !index_sets.iter().all(|set| set.contains(r.row_id.0)) {
5339 continue;
5340 }
5341 if let Some(cid) = cid {
5342 if let Some(v) = as_f64(r.columns.get(&cid)) {
5343 passing_vals.push(v);
5344 } } else {
5346 passing_vals.push(0.0); }
5348 }
5349 let m = passing_vals.len();
5350
5351 let (point, half) = match agg {
5352 ApproxAgg::Count => {
5353 let p = m as f64 / s as f64;
5355 let point = n_pop as f64 * p;
5356 let var = if s > 1 {
5357 n_pop as f64 * n_pop as f64 * p * (1.0 - p) / s as f64
5358 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5359 } else {
5360 0.0
5361 };
5362 (point, z * var.sqrt())
5363 }
5364 ApproxAgg::Sum => {
5365 let y: Vec<f64> = live_sample
5367 .iter()
5368 .map(|r| {
5369 let passes_row = conditions
5370 .iter()
5371 .all(|c| condition_matches_row(c, r, &self.schema))
5372 && index_sets.iter().all(|set| set.contains(r.row_id.0));
5373 if passes_row {
5374 cid.and_then(|c| as_f64(r.columns.get(&c))).unwrap_or(0.0)
5375 } else {
5376 0.0
5377 }
5378 })
5379 .collect();
5380 let mean_y = y.iter().sum::<f64>() / s as f64;
5381 let point = n_pop as f64 * mean_y;
5382 let var = if s > 1 {
5383 let ss: f64 = y.iter().map(|v| (v - mean_y).powi(2)).sum();
5384 let var_y = ss / (s - 1) as f64;
5385 n_pop as f64 * n_pop as f64 * var_y / s as f64
5386 * (1.0 - s as f64 / n_pop as f64).max(0.0)
5387 } else {
5388 0.0
5389 };
5390 (point, z * var.sqrt())
5391 }
5392 ApproxAgg::Avg => {
5393 if m == 0 {
5394 return Ok(Some(ApproxResult {
5395 point: 0.0,
5396 ci_low: 0.0,
5397 ci_high: 0.0,
5398 n_population: n_pop,
5399 n_sample_live: s,
5400 n_passing: 0,
5401 }));
5402 }
5403 let mean = passing_vals.iter().sum::<f64>() / m as f64;
5404 let half = if m > 1 {
5405 let ss: f64 = passing_vals.iter().map(|v| (v - mean).powi(2)).sum();
5406 let sd = (ss / (m - 1) as f64).sqrt();
5407 let fpc = (1.0 - s as f64 / n_pop as f64).max(0.0);
5408 z * sd / (m as f64).sqrt() * fpc.sqrt()
5409 } else {
5410 0.0
5411 };
5412 (mean, half)
5413 }
5414 };
5415
5416 Ok(Some(ApproxResult {
5417 point,
5418 ci_low: point - half,
5419 ci_high: point + half,
5420 n_population: n_pop,
5421 n_sample_live: s,
5422 n_passing: m,
5423 }))
5424 }
5425
5426 pub fn exact_column_stats(
5434 &self,
5435 _snapshot: Snapshot,
5436 projection: &[u16],
5437 ) -> Result<Option<HashMap<u16, ColumnStat>>> {
5438 if !(self.memtable.is_empty() && self.mutable_run.is_empty() && self.run_refs.len() == 1) {
5439 return Ok(None);
5440 }
5441 let reader = self.open_reader(self.run_refs[0].run_id)?;
5442 if self.live_count != reader.row_count() as u64 {
5443 return Ok(None);
5444 }
5445 let mut out = HashMap::new();
5446 for &cid in projection {
5447 let cdef = match self.schema.columns.iter().find(|c| c.id == cid) {
5448 Some(c) => c,
5449 None => continue,
5450 };
5451 let Some(stats) = reader.column_page_stats(cid) else {
5453 out.insert(
5454 cid,
5455 ColumnStat {
5456 min: None,
5457 max: None,
5458 null_count: self.live_count,
5459 },
5460 );
5461 continue;
5462 };
5463 let stat = match cdef.ty {
5464 TypeId::Int64 | TypeId::TimestampNanos | TypeId::Date32 => {
5465 agg_int(stats, crate::sorted_run::be_i64).map(|(mn, mx, n)| ColumnStat {
5466 min: mn.map(Value::Int64),
5467 max: mx.map(Value::Int64),
5468 null_count: n,
5469 })
5470 }
5471 TypeId::Float64 => {
5472 agg_float(stats, crate::sorted_run::be_f64).map(|(mn, mx, n)| ColumnStat {
5473 min: mn.map(Value::Float64),
5474 max: mx.map(Value::Float64),
5475 null_count: n,
5476 })
5477 }
5478 _ => None,
5479 };
5480 if let Some(s) = stat {
5481 out.insert(cid, s);
5482 }
5483 }
5484 Ok(Some(out))
5485 }
5486
5487 pub fn dir(&self) -> &Path {
5488 &self.dir
5489 }
5490
5491 pub fn schema(&self) -> &Schema {
5492 &self.schema
5493 }
5494
5495 pub(crate) fn prepare_alter_column(
5496 &mut self,
5497 column_name: &str,
5498 change: &AlterColumn,
5499 ) -> Result<ColumnDef> {
5500 if !self.pending_rows.is_empty() || !self.pending_dels.is_empty() {
5501 return Err(MongrelError::InvalidArgument(
5502 "ALTER COLUMN requires committing staged writes first".into(),
5503 ));
5504 }
5505 let old = self
5506 .schema
5507 .columns
5508 .iter()
5509 .find(|c| c.name == column_name)
5510 .cloned()
5511 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5512 let mut next = old.clone();
5513
5514 if let Some(name) = &change.name {
5515 let trimmed = name.trim();
5516 if trimmed.is_empty() {
5517 return Err(MongrelError::InvalidArgument(
5518 "ALTER COLUMN name must not be empty".into(),
5519 ));
5520 }
5521 if trimmed != old.name && self.schema.columns.iter().any(|c| c.name == trimmed) {
5522 return Err(MongrelError::Schema(format!(
5523 "column {trimmed} already exists"
5524 )));
5525 }
5526 next.name = trimmed.to_string();
5527 }
5528
5529 if let Some(ty) = change.ty {
5530 next.ty = ty;
5531 }
5532 if let Some(flags) = change.flags {
5533 validate_alter_column_flags(old.flags, flags)?;
5534 next.flags = flags;
5535 }
5536
5537 validate_alter_column_type(&self.schema, &old, &next, self.has_stored_versions())?;
5538 if old.flags.contains(ColumnFlags::NULLABLE)
5539 && !next.flags.contains(ColumnFlags::NULLABLE)
5540 && self.column_has_nulls(old.id)?
5541 {
5542 return Err(MongrelError::InvalidArgument(format!(
5543 "column '{}' contains NULL values",
5544 old.name
5545 )));
5546 }
5547 Ok(next)
5548 }
5549
5550 pub(crate) fn apply_altered_column(&mut self, column: ColumnDef) -> Result<()> {
5551 let idx = self
5552 .schema
5553 .columns
5554 .iter()
5555 .position(|c| c.id == column.id)
5556 .ok_or_else(|| MongrelError::Schema(format!("unknown column {}", column.id)))?;
5557 if self.schema.columns[idx] == column {
5558 return Ok(());
5559 }
5560 self.schema.columns[idx] = column;
5561 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5562 self.schema.validate_auto_increment()?;
5563 self.auto_inc = resolve_auto_inc(&self.schema);
5564 self.column_keys = build_column_keys(self.kek.as_deref(), &self.schema);
5565 write_schema(&self.dir, &self.schema)?;
5566 self.clear_result_cache();
5567 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5568 self.persist_manifest(self.current_epoch())?;
5569 Ok(())
5570 }
5571
5572 pub fn alter_column(&mut self, column_name: &str, change: AlterColumn) -> Result<ColumnDef> {
5573 let column = self.prepare_alter_column(column_name, &change)?;
5574 self.apply_altered_column(column.clone())?;
5575 Ok(column)
5576 }
5577
5578 fn column_has_nulls(&mut self, column_id: u16) -> Result<bool> {
5579 if self.live_count == 0 {
5580 return Ok(false);
5581 }
5582 let snap = self.snapshot();
5583 let columns = self.visible_columns_native(snap, Some(&[column_id]))?;
5584 Ok(columns
5585 .first()
5586 .map(|(_, col)| col.null_count(col.len()) != 0)
5587 .unwrap_or(true))
5588 }
5589
5590 fn has_stored_versions(&self) -> bool {
5591 !self.memtable.is_empty()
5592 || !self.mutable_run.is_empty()
5593 || self.run_refs.iter().any(|r| r.row_count > 0)
5594 || !self.retiring.is_empty()
5595 }
5596
5597 pub fn add_column(&mut self, name: &str, ty: TypeId, flags: ColumnFlags) -> Result<u16> {
5602 if self.schema.columns.iter().any(|c| c.name == name) {
5603 return Err(MongrelError::Schema(format!(
5604 "column {name} already exists"
5605 )));
5606 }
5607 let id = self.schema.columns.iter().map(|c| c.id).max().unwrap_or(0) + 1;
5608 self.schema.columns.push(ColumnDef {
5609 id,
5610 name: name.to_string(),
5611 ty,
5612 flags,
5613 });
5614 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5615 self.schema.validate_auto_increment()?;
5616 if flags.contains(ColumnFlags::AUTO_INCREMENT) {
5617 self.auto_inc = resolve_auto_inc(&self.schema);
5618 }
5619 write_schema(&self.dir, &self.schema)?;
5620 self.clear_result_cache();
5621 let _ = std::fs::remove_dir_all(self.dir.join("_shadow"));
5623 self.persist_manifest(self.current_epoch())?;
5624 Ok(id)
5625 }
5626
5627 pub fn add_learned_range_index(&mut self, column_name: &str) -> Result<()> {
5636 let cid = self
5637 .schema
5638 .columns
5639 .iter()
5640 .find(|c| c.name == column_name)
5641 .map(|c| c.id)
5642 .ok_or_else(|| MongrelError::Schema(format!("unknown column {column_name}")))?;
5643 let ty = self
5644 .schema
5645 .columns
5646 .iter()
5647 .find(|c| c.id == cid)
5648 .map(|c| c.ty)
5649 .unwrap_or(TypeId::Int64);
5650 if !matches!(
5651 ty,
5652 TypeId::Int64 | TypeId::Float64 | TypeId::TimestampNanos | TypeId::Date32
5653 ) {
5654 return Err(MongrelError::Schema(format!(
5655 "LearnedRange requires a numeric column; {column_name} is {ty:?}"
5656 )));
5657 }
5658 if self
5659 .schema
5660 .indexes
5661 .iter()
5662 .any(|i| i.column_id == cid && i.kind == IndexKind::LearnedRange)
5663 {
5664 return Ok(()); }
5666 self.schema.indexes.push(IndexDef {
5667 name: format!("{}_learned_range", column_name),
5668 column_id: cid,
5669 kind: IndexKind::LearnedRange,
5670 });
5671 self.schema.schema_id = self.schema.schema_id.saturating_add(1);
5672 write_schema(&self.dir, &self.schema)?;
5673 self.build_learned_ranges()?;
5674 Ok(())
5675 }
5676
5677 pub fn set_sync_byte_threshold(&mut self, threshold: u64) {
5680 self.sync_byte_threshold = threshold;
5681 if let WalSink::Private(w) = &mut self.wal {
5682 w.set_sync_byte_threshold(threshold);
5683 }
5684 }
5685
5686 pub fn page_cache_flush(&self) {
5690 self.page_cache.lock().flush_to_disk();
5691 }
5692
5693 pub fn page_cache_len(&self) -> usize {
5695 self.page_cache.lock().len()
5696 }
5697
5698 pub fn decoded_cache_len(&self) -> usize {
5701 self.decoded_cache.lock().len()
5702 }
5703
5704 pub fn drain_memtable_sorted(&mut self) -> Vec<Row> {
5707 self.memtable.drain_sorted()
5708 }
5709
5710 pub(crate) fn run_path(&self, run_id: u64) -> PathBuf {
5711 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr"))
5712 }
5713
5714 pub(crate) fn table_dir(&self) -> &Path {
5715 &self.dir
5716 }
5717
5718 pub(crate) fn schema_ref(&self) -> &crate::schema::Schema {
5719 &self.schema
5720 }
5721
5722 pub(crate) fn alloc_run_id(&mut self) -> u64 {
5723 let id = self.next_run_id;
5724 self.next_run_id += 1;
5725 id
5726 }
5727
5728 pub(crate) fn link_run(&mut self, run_ref: crate::manifest::RunRef) {
5729 self.run_refs.push(run_ref);
5730 }
5731
5732 pub(crate) fn retire_run(&mut self, run_id: u128, retire_epoch: u64) {
5742 self.retiring.push(crate::manifest::RetiredRun {
5743 run_id,
5744 retire_epoch,
5745 });
5746 }
5747
5748 pub(crate) fn reap_retiring(&mut self, min_active: Epoch) -> Result<usize> {
5752 if self.retiring.is_empty() {
5753 return Ok(0);
5754 }
5755 let mut reaped = 0;
5756 let mut kept: Vec<crate::manifest::RetiredRun> = Vec::new();
5757 for r in std::mem::take(&mut self.retiring) {
5763 if min_active.0 >= r.retire_epoch {
5764 let _ = std::fs::remove_file(self.run_path(r.run_id as u64));
5765 reaped += 1;
5766 } else {
5767 kept.push(r);
5768 }
5769 }
5770 self.retiring = kept;
5771 if reaped > 0 {
5772 self.persist_manifest(self.current_epoch())?;
5773 }
5774 Ok(reaped)
5775 }
5776
5777 pub(crate) fn recover_spilled_run(&mut self, run_ref: crate::manifest::RunRef) -> bool {
5778 if self.run_refs.iter().any(|r| r.run_id == run_ref.run_id) {
5779 return false;
5780 }
5781 self.live_count = self.live_count.saturating_add(run_ref.row_count);
5782 self.run_refs.push(run_ref);
5783 self.indexes_complete = false;
5784 true
5785 }
5786
5787 pub(crate) fn kek_ref(&self) -> Option<&Arc<Kek>> {
5788 self.kek.as_ref()
5789 }
5790
5791 pub(crate) fn open_reader(&self, run_id: u128) -> Result<RunReader> {
5792 let mut reader = RunReader::open_with_cache(
5793 self.dir.join(RUNS_DIR).join(format!("r-{run_id}.sr")),
5794 self.schema.clone(),
5795 self.kek.clone(),
5796 Some(self.page_cache.clone()),
5797 Some(self.decoded_cache.clone()),
5798 self.table_id,
5799 )?;
5800 if let Some(rr) = self.run_refs.iter().find(|r| r.run_id == run_id) {
5804 reader.set_uniform_epoch(Epoch(rr.epoch_created));
5805 }
5806 Ok(reader)
5807 }
5808
5809 pub(crate) fn run_refs(&self) -> &[RunRef] {
5810 &self.run_refs
5811 }
5812
5813 pub(crate) fn runs_dir(&self) -> PathBuf {
5814 self.dir.join(RUNS_DIR)
5815 }
5816
5817 pub(crate) fn wal_dir(&self) -> PathBuf {
5818 self.dir.join(WAL_DIR)
5819 }
5820
5821 pub(crate) fn set_run_refs(&mut self, refs: Vec<RunRef>) {
5822 self.run_refs = refs;
5823 }
5824
5825 pub(crate) fn next_run_id(&self) -> u64 {
5826 self.next_run_id
5827 }
5828
5829 pub(crate) fn compaction_zstd_level(&self) -> i32 {
5830 self.compaction_zstd_level
5831 }
5832
5833 pub(crate) fn bump_next_run_id(&mut self) {
5834 self.next_run_id += 1;
5835 }
5836
5837 pub(crate) fn kek(&self) -> Option<Arc<Kek>> {
5838 self.kek.clone()
5839 }
5840
5841 #[cfg(feature = "encryption")]
5845 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
5846 self.kek.as_ref().map(|k| k.derive_idx_key())
5847 }
5848
5849 #[cfg(not(feature = "encryption"))]
5850 fn idx_dek(&self) -> Option<Zeroizing<[u8; DEK_LEN]>> {
5851 None
5852 }
5853
5854 #[cfg(feature = "encryption")]
5858 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
5859 self.kek.as_ref().map(|k| *k.derive_meta_key())
5860 }
5861
5862 #[cfg(not(feature = "encryption"))]
5863 fn manifest_meta_dek(&self) -> Option<[u8; DEK_LEN]> {
5864 None
5865 }
5866
5867 pub(crate) fn indexable_column_specs(&self) -> Vec<(u16, u8)> {
5870 self.column_keys
5871 .iter()
5872 .map(|(&id, &(_, scheme))| (id, scheme))
5873 .collect()
5874 }
5875
5876 #[cfg(feature = "encryption")]
5881 fn tokenize_value(&self, column_id: u16, v: &Value) -> Option<Value> {
5882 self.tokenize_value_enc(column_id, v)
5883 }
5884
5885 #[cfg(feature = "encryption")]
5886 fn tokenize_value_enc(&self, column_id: u16, v: &Value) -> Option<Value> {
5887 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
5888 let (key, scheme) = self.column_keys.get(&column_id)?;
5889 let token: Vec<u8> = match (*scheme, v) {
5890 (SCHEME_HMAC_EQ, _) => hmac_token(key, &v.encode_key()).to_vec(),
5891 (_, Value::Int64(x)) => ope_token_i64(key, *x).to_vec(),
5892 (_, Value::Float64(x)) => ope_token_f64(key, *x).to_vec(),
5893 _ => hmac_token(key, &v.encode_key()).to_vec(),
5894 };
5895 Some(Value::Bytes(token))
5896 }
5897
5898 fn index_lookup_key(&self, column_id: u16, v: &Value) -> Vec<u8> {
5900 self.index_lookup_key_bytes(column_id, &v.encode_key())
5901 }
5902
5903 fn index_lookup_key_bytes(&self, column_id: u16, encoded: &[u8]) -> Vec<u8> {
5906 #[cfg(feature = "encryption")]
5907 {
5908 use crate::encryption::{hmac_token, SCHEME_HMAC_EQ};
5909 if let Some((key, scheme)) = self.column_keys.get(&column_id) {
5910 if *scheme == SCHEME_HMAC_EQ {
5911 return hmac_token(key, encoded).to_vec();
5912 }
5913 }
5914 }
5915 let _ = column_id;
5916 encoded.to_vec()
5917 }
5918}
5919
5920fn native_int64_strictly_increasing(col: &columnar::NativeColumn, n: usize) -> bool {
5921 let columnar::NativeColumn::Int64 { data, validity } = col else {
5922 return false;
5923 };
5924 if data.len() < n || !columnar::all_non_null(validity, n) {
5925 return false;
5926 }
5927 data.iter()
5928 .take(n)
5929 .zip(data.iter().skip(1))
5930 .all(|(a, b)| a < b)
5931}
5932
5933#[derive(Debug, Clone)]
5937pub struct ColumnStat {
5938 pub min: Option<Value>,
5939 pub max: Option<Value>,
5940 pub null_count: u64,
5941}
5942
5943#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5945pub enum NativeAgg {
5946 Count,
5947 Sum,
5948 Min,
5949 Max,
5950 Avg,
5951}
5952
5953#[derive(Debug, Clone, PartialEq)]
5955pub enum NativeAggResult {
5956 Count(u64),
5957 Int(i64),
5958 Float(f64),
5959 Null,
5961}
5962
5963#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5965pub enum ApproxAgg {
5966 Count,
5967 Sum,
5968 Avg,
5969}
5970
5971#[derive(Debug, Clone)]
5975pub struct ApproxResult {
5976 pub point: f64,
5978 pub ci_low: f64,
5980 pub ci_high: f64,
5982 pub n_population: u64,
5984 pub n_sample_live: usize,
5986 pub n_passing: usize,
5988}
5989
5990#[derive(Debug, Clone, PartialEq)]
5995pub enum AggState {
5996 Count(u64),
5998 SumI {
6000 sum: i128,
6001 count: u64,
6002 },
6003 SumF {
6005 sum: f64,
6006 count: u64,
6007 },
6008 AvgI {
6010 sum: i128,
6011 count: u64,
6012 },
6013 AvgF {
6015 sum: f64,
6016 count: u64,
6017 },
6018 MinI(i64),
6020 MaxI(i64),
6021 MinF(f64),
6023 MaxF(f64),
6024 Empty,
6026}
6027
6028impl AggState {
6029 pub fn merge(self, other: AggState) -> AggState {
6031 use AggState::*;
6032 match (self, other) {
6033 (Empty, x) | (x, Empty) => x,
6034 (Count(a), Count(b)) => Count(a + b),
6035 (SumI { sum: sa, count: ca }, SumI { sum: sb, count: cb }) => SumI {
6036 sum: sa + sb,
6037 count: ca + cb,
6038 },
6039 (SumF { sum: sa, count: ca }, SumF { sum: sb, count: cb }) => SumF {
6040 sum: sa + sb,
6041 count: ca + cb,
6042 },
6043 (AvgI { sum: sa, count: ca }, AvgI { sum: sb, count: cb }) => AvgI {
6044 sum: sa + sb,
6045 count: ca + cb,
6046 },
6047 (AvgF { sum: sa, count: ca }, AvgF { sum: sb, count: cb }) => AvgF {
6048 sum: sa + sb,
6049 count: ca + cb,
6050 },
6051 (MinI(a), MinI(b)) => MinI(a.min(b)),
6052 (MaxI(a), MaxI(b)) => MaxI(a.max(b)),
6053 (MinF(a), MinF(b)) => MinF(a.min(b)),
6054 (MaxF(a), MaxF(b)) => MaxF(a.max(b)),
6055 _ => Empty, }
6057 }
6058
6059 pub fn point(&self) -> Option<f64> {
6061 match self {
6062 AggState::Count(n) => Some(*n as f64),
6063 AggState::SumI { sum, .. } => Some(*sum as f64),
6064 AggState::SumF { sum, .. } => Some(*sum),
6065 AggState::AvgI { sum, count } if *count > 0 => Some(*sum as f64 / *count as f64),
6066 AggState::AvgF { sum, count } if *count > 0 => Some(*sum / *count as f64),
6067 AggState::MinI(n) => Some(*n as f64),
6068 AggState::MaxI(n) => Some(*n as f64),
6069 AggState::MinF(n) => Some(*n),
6070 AggState::MaxF(n) => Some(*n),
6071 AggState::AvgI { .. } | AggState::AvgF { .. } | AggState::Empty => None,
6072 }
6073 }
6074
6075 pub fn from_native(result: NativeAggResult, agg: NativeAgg, ty: Option<TypeId>) -> Self {
6079 let is_float = matches!(ty, Some(TypeId::Float64));
6080 match (agg, result) {
6081 (NativeAgg::Count, NativeAggResult::Count(n)) => AggState::Count(n),
6082 (NativeAgg::Sum, NativeAggResult::Int(x)) => AggState::SumI {
6083 sum: x as i128,
6084 count: 1, },
6086 (NativeAgg::Sum, NativeAggResult::Float(x)) => AggState::SumF { sum: x, count: 1 },
6087 (NativeAgg::Avg, NativeAggResult::Float(x)) => AggState::AvgF { sum: x, count: 1 },
6088 (NativeAgg::Min, NativeAggResult::Int(x)) => AggState::MinI(x),
6089 (NativeAgg::Max, NativeAggResult::Int(x)) => AggState::MaxI(x),
6090 (NativeAgg::Min, NativeAggResult::Float(x)) => AggState::MinF(x),
6091 (NativeAgg::Max, NativeAggResult::Float(x)) => AggState::MaxF(x),
6092 (NativeAgg::Count, _) => AggState::Empty,
6093 (_, NativeAggResult::Null) => AggState::Empty,
6094 _ => {
6095 let _ = is_float;
6096 AggState::Empty
6097 }
6098 }
6099 }
6100}
6101
6102#[derive(Debug, Clone)]
6105pub struct CachedAgg {
6106 pub state: AggState,
6107 pub watermark: u64,
6108 pub epoch: u64,
6109}
6110
6111#[derive(Debug, Clone)]
6113pub struct IncrementalAggResult {
6114 pub state: AggState,
6116 pub incremental: bool,
6119 pub delta_rows: u64,
6121}
6122
6123fn agg_state_from_rows(
6127 rows: &[Row],
6128 conditions: &[crate::query::Condition],
6129 index_sets: &[RowIdSet],
6130 column: Option<u16>,
6131 agg: NativeAgg,
6132 schema: &Schema,
6133) -> Result<AggState> {
6134 let mut count: u64 = 0;
6135 let mut sum_i: i128 = 0;
6136 let mut sum_f: f64 = 0.0;
6137 let mut mn_i: i64 = i64::MAX;
6138 let mut mx_i: i64 = i64::MIN;
6139 let mut mn_f: f64 = f64::INFINITY;
6140 let mut mx_f: f64 = f64::NEG_INFINITY;
6141 let mut saw_int = false;
6142 let mut saw_float = false;
6143 for r in rows {
6144 if !conditions
6145 .iter()
6146 .all(|c| condition_matches_row(c, r, schema))
6147 {
6148 continue;
6149 }
6150 if !index_sets.iter().all(|s| s.contains(r.row_id.0)) {
6151 continue;
6152 }
6153 match agg {
6154 NativeAgg::Count => match column {
6155 None => count += 1,
6157 Some(cid) => match r.columns.get(&cid) {
6160 None | Some(Value::Null) => {}
6161 Some(_) => count += 1,
6162 },
6163 },
6164 _ => match column.and_then(|cid| r.columns.get(&cid)) {
6165 Some(Value::Int64(n)) => {
6166 count += 1;
6167 sum_i += *n as i128;
6168 mn_i = mn_i.min(*n);
6169 mx_i = mx_i.max(*n);
6170 saw_int = true;
6171 }
6172 Some(Value::Float64(f)) => {
6173 count += 1;
6174 sum_f += f;
6175 mn_f = mn_f.min(*f);
6176 mx_f = mx_f.max(*f);
6177 saw_float = true;
6178 }
6179 _ => {}
6180 },
6181 }
6182 }
6183 Ok(match agg {
6184 NativeAgg::Count => {
6185 if count == 0 {
6186 AggState::Empty
6187 } else {
6188 AggState::Count(count)
6189 }
6190 }
6191 NativeAgg::Sum => {
6192 if count == 0 {
6193 AggState::Empty
6194 } else if saw_int {
6195 AggState::SumI { sum: sum_i, count }
6196 } else {
6197 AggState::SumF { sum: sum_f, count }
6198 }
6199 }
6200 NativeAgg::Avg => {
6201 if count == 0 {
6202 AggState::Empty
6203 } else if saw_int {
6204 AggState::AvgI { sum: sum_i, count }
6205 } else {
6206 AggState::AvgF { sum: sum_f, count }
6207 }
6208 }
6209 NativeAgg::Min => {
6210 if !saw_int && !saw_float {
6211 AggState::Empty
6212 } else if saw_int {
6213 AggState::MinI(mn_i)
6214 } else {
6215 AggState::MinF(mn_f)
6216 }
6217 }
6218 NativeAgg::Max => {
6219 if !saw_int && !saw_float {
6220 AggState::Empty
6221 } else if saw_int {
6222 AggState::MaxI(mx_i)
6223 } else {
6224 AggState::MaxF(mx_f)
6225 }
6226 }
6227 })
6228}
6229
6230fn condition_matches_row(c: &crate::query::Condition, row: &Row, schema: &Schema) -> bool {
6234 use crate::query::Condition;
6235 match c {
6236 Condition::Pk(key) => match schema.primary_key() {
6237 Some(pk) => row
6238 .columns
6239 .get(&pk.id)
6240 .map(|v| v.encode_key() == *key)
6241 .unwrap_or(false),
6242 None => false,
6243 },
6244 Condition::BitmapEq { column_id, value } => row
6245 .columns
6246 .get(column_id)
6247 .map(|v| v.encode_key() == *value)
6248 .unwrap_or(false),
6249 Condition::BitmapIn { column_id, values } => {
6250 let key = row.columns.get(column_id).map(|v| v.encode_key());
6251 match key {
6252 Some(k) => values.contains(&k),
6253 None => false,
6254 }
6255 }
6256 Condition::Range { column_id, lo, hi } => match row.columns.get(column_id) {
6257 Some(Value::Int64(n)) => *n >= *lo && *n <= *hi,
6258 _ => false,
6259 },
6260 Condition::RangeF64 {
6261 column_id,
6262 lo,
6263 lo_inclusive,
6264 hi,
6265 hi_inclusive,
6266 } => match row.columns.get(column_id) {
6267 Some(Value::Float64(n)) => {
6268 let lo_ok = if *lo_inclusive { *n >= *lo } else { *n > *lo };
6269 let hi_ok = if *hi_inclusive { *n <= *hi } else { *n < *hi };
6270 lo_ok && hi_ok
6271 }
6272 _ => false,
6273 },
6274 Condition::FmContains { column_id, pattern } => match row.columns.get(column_id) {
6275 Some(Value::Bytes(b)) => {
6276 !pattern.is_empty() && b.windows(pattern.len()).any(|w| w == &pattern[..])
6277 }
6278 _ => false,
6279 },
6280 Condition::FmContainsAll {
6281 column_id,
6282 patterns,
6283 } => match row.columns.get(column_id) {
6284 Some(Value::Bytes(b)) => patterns
6285 .iter()
6286 .all(|pat| !pat.is_empty() && b.windows(pat.len()).any(|w| w == &pat[..])),
6287 _ => false,
6288 },
6289 Condition::Ann { .. }
6290 | Condition::SparseMatch { .. }
6291 | Condition::MinHashSimilar { .. } => true,
6292 Condition::IsNull { column_id } => {
6293 matches!(row.columns.get(column_id), Some(Value::Null) | None)
6294 }
6295 Condition::IsNotNull { column_id } => {
6296 !matches!(row.columns.get(column_id), Some(Value::Null) | None)
6297 }
6298 }
6299}
6300
6301fn as_f64(v: Option<&Value>) -> Option<f64> {
6303 match v {
6304 Some(Value::Int64(n)) => Some(*n as f64),
6305 Some(Value::Float64(f)) => Some(*f),
6306 _ => None,
6307 }
6308}
6309
6310fn accumulate_int(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, i128, i64, i64)> {
6314 let mut count: u64 = 0;
6315 let mut sum: i128 = 0;
6316 let mut mn: i64 = i64::MAX;
6317 let mut mx: i64 = i64::MIN;
6318 while let Some(cols) = cursor.next_batch()? {
6319 if let Some(crate::columnar::NativeColumn::Int64 { data, validity }) = cols.first() {
6320 if crate::columnar::all_non_null(validity, data.len()) {
6321 count += data.len() as u64;
6323 sum += data.iter().map(|&v| v as i128).sum::<i128>();
6324 mn = mn.min(*data.iter().min().unwrap_or(&mn));
6325 mx = mx.max(*data.iter().max().unwrap_or(&mx));
6326 } else {
6327 for (i, &v) in data.iter().enumerate() {
6328 if crate::columnar::validity_bit(validity, i) {
6329 count += 1;
6330 sum += v as i128;
6331 mn = mn.min(v);
6332 mx = mx.max(v);
6333 }
6334 }
6335 }
6336 }
6337 }
6338 Ok((count, sum, mn, mx))
6339}
6340
6341fn accumulate_float(cursor: &mut dyn crate::cursor::Cursor) -> Result<(u64, f64, f64, f64)> {
6343 let mut count: u64 = 0;
6344 let mut sum: f64 = 0.0;
6345 let mut mn: f64 = f64::INFINITY;
6346 let mut mx: f64 = f64::NEG_INFINITY;
6347 while let Some(cols) = cursor.next_batch()? {
6348 if let Some(crate::columnar::NativeColumn::Float64 { data, validity }) = cols.first() {
6349 if crate::columnar::all_non_null(validity, data.len()) {
6350 count += data.len() as u64;
6351 sum += data.iter().sum::<f64>();
6352 mn = mn.min(data.iter().copied().fold(f64::INFINITY, f64::min));
6353 mx = mx.max(data.iter().copied().fold(f64::NEG_INFINITY, f64::max));
6354 } else {
6355 for (i, &v) in data.iter().enumerate() {
6356 if crate::columnar::validity_bit(validity, i) {
6357 count += 1;
6358 sum += v;
6359 mn = mn.min(v);
6360 mx = mx.max(v);
6361 }
6362 }
6363 }
6364 }
6365 }
6366 Ok((count, sum, mn, mx))
6367}
6368
6369fn pack_int(agg: NativeAgg, count: u64, sum: i128, mn: i64, mx: i64) -> NativeAggResult {
6370 if count == 0 && !matches!(agg, NativeAgg::Count) {
6371 return NativeAggResult::Null;
6372 }
6373 match agg {
6374 NativeAgg::Count => NativeAggResult::Count(count),
6375 NativeAgg::Sum => match sum.try_into() {
6378 Ok(v) => NativeAggResult::Int(v),
6379 Err(_) => NativeAggResult::Null,
6380 },
6381 NativeAgg::Min => NativeAggResult::Int(mn),
6382 NativeAgg::Max => NativeAggResult::Int(mx),
6383 NativeAgg::Avg => NativeAggResult::Float((sum as f64) / (count as f64)),
6384 }
6385}
6386
6387fn pack_float(agg: NativeAgg, count: u64, sum: f64, mn: f64, mx: f64) -> NativeAggResult {
6388 if count == 0 && !matches!(agg, NativeAgg::Count) {
6389 return NativeAggResult::Null;
6390 }
6391 match agg {
6392 NativeAgg::Count => NativeAggResult::Count(count),
6393 NativeAgg::Sum => NativeAggResult::Float(sum),
6394 NativeAgg::Min => NativeAggResult::Float(mn),
6395 NativeAgg::Max => NativeAggResult::Float(mx),
6396 NativeAgg::Avg => NativeAggResult::Float(sum / (count as f64)),
6397 }
6398}
6399
6400fn agg_int(
6403 stats: &[crate::page::PageStat],
6404 decode: fn(Option<&[u8]>) -> Option<i64>,
6405) -> Option<(Option<i64>, Option<i64>, u64)> {
6406 let (mut mn, mut mx, mut nulls) = (i64::MAX, i64::MIN, 0u64);
6407 let mut any = false;
6408 for s in stats {
6409 if let Some(v) = decode(s.min.as_deref()) {
6410 mn = mn.min(v);
6411 any = true;
6412 }
6413 if let Some(v) = decode(s.max.as_deref()) {
6414 mx = mx.max(v);
6415 any = true;
6416 }
6417 nulls += s.null_count;
6418 }
6419 any.then_some((Some(mn), Some(mx), nulls))
6420}
6421
6422fn agg_float(
6424 stats: &[crate::page::PageStat],
6425 decode: fn(Option<&[u8]>) -> Option<f64>,
6426) -> Option<(Option<f64>, Option<f64>, u64)> {
6427 let (mut mn, mut mx, mut nulls) = (f64::INFINITY, f64::NEG_INFINITY, 0u64);
6428 let mut any = false;
6429 for s in stats {
6430 if let Some(v) = decode(s.min.as_deref()) {
6431 mn = mn.min(v);
6432 any = true;
6433 }
6434 if let Some(v) = decode(s.max.as_deref()) {
6435 mx = mx.max(v);
6436 any = true;
6437 }
6438 nulls += s.null_count;
6439 }
6440 any.then_some((Some(mn), Some(mx), nulls))
6441}
6442
6443type SecondaryIndexes = (
6445 HashMap<u16, BitmapIndex>,
6446 HashMap<u16, AnnIndex>,
6447 HashMap<u16, FmIndex>,
6448 HashMap<u16, SparseIndex>,
6449 HashMap<u16, MinHashIndex>,
6450);
6451
6452fn empty_indexes(schema: &Schema) -> SecondaryIndexes {
6453 let mut bitmap = HashMap::new();
6454 let mut ann = HashMap::new();
6455 let mut fm = HashMap::new();
6456 let mut sparse = HashMap::new();
6457 let mut minhash = HashMap::new();
6458 for idef in &schema.indexes {
6459 match idef.kind {
6460 IndexKind::Bitmap => {
6461 bitmap.insert(idef.column_id, BitmapIndex::new());
6462 }
6463 IndexKind::Ann => {
6464 let dim = schema
6465 .columns
6466 .iter()
6467 .find(|c| c.id == idef.column_id)
6468 .and_then(|c| match c.ty {
6469 TypeId::Embedding { dim } => Some(dim as usize),
6470 _ => None,
6471 })
6472 .unwrap_or(0);
6473 ann.insert(idef.column_id, AnnIndex::new(dim));
6474 }
6475 IndexKind::FmIndex => {
6476 fm.insert(idef.column_id, FmIndex::new());
6477 }
6478 IndexKind::Sparse => {
6479 sparse.insert(idef.column_id, SparseIndex::new());
6480 }
6481 IndexKind::MinHash => {
6482 minhash.insert(idef.column_id, MinHashIndex::new());
6483 }
6484 _ => {}
6485 }
6486 }
6487 (bitmap, ann, fm, sparse, minhash)
6488}
6489
6490const ALTER_COLUMN_PROTECTED_FLAGS: u32 = ColumnFlags::PRIMARY_KEY
6491 | ColumnFlags::AUTO_INCREMENT
6492 | ColumnFlags::ENCRYPTED
6493 | ColumnFlags::ENCRYPTED_INDEXABLE
6494 | ColumnFlags::EMBEDDING_BINARY_QUANTIZED;
6495
6496fn validate_alter_column_flags(old: ColumnFlags, new: ColumnFlags) -> Result<()> {
6497 if (old.bits() ^ new.bits()) & ALTER_COLUMN_PROTECTED_FLAGS != 0 {
6498 return Err(MongrelError::Schema(
6499 "ALTER COLUMN may only change NULLABLE; primary key, auto-increment, encryption, and embedding flags are immutable".into(),
6500 ));
6501 }
6502 Ok(())
6503}
6504
6505fn validate_alter_column_type(
6506 schema: &Schema,
6507 old: &ColumnDef,
6508 next: &ColumnDef,
6509 has_stored_versions: bool,
6510) -> Result<()> {
6511 if old.ty == next.ty {
6512 return Ok(());
6513 }
6514 if schema.indexes.iter().any(|i| i.column_id == old.id) {
6515 return Err(MongrelError::Schema(format!(
6516 "ALTER COLUMN TYPE is not supported for indexed column '{}'",
6517 old.name
6518 )));
6519 }
6520 if !has_stored_versions || storage_compatible_type_change(old.ty, next.ty) {
6521 return Ok(());
6522 }
6523 Err(MongrelError::Schema(format!(
6524 "ALTER COLUMN TYPE from {:?} to {:?} requires an empty column or a representation-compatible type",
6525 old.ty, next.ty
6526 )))
6527}
6528
6529fn storage_compatible_type_change(old: TypeId, new: TypeId) -> bool {
6530 matches!(
6531 (old, new),
6532 (TypeId::Int64, TypeId::TimestampNanos) | (TypeId::TimestampNanos, TypeId::Int64)
6533 )
6534}
6535
6536#[allow(clippy::too_many_arguments)]
6537fn index_into(
6538 schema: &Schema,
6539 row: &Row,
6540 hot: &mut HotIndex,
6541 bitmap: &mut HashMap<u16, BitmapIndex>,
6542 ann: &mut HashMap<u16, AnnIndex>,
6543 fm: &mut HashMap<u16, FmIndex>,
6544 sparse: &mut HashMap<u16, SparseIndex>,
6545 minhash: &mut HashMap<u16, MinHashIndex>,
6546) {
6547 for idef in &schema.indexes {
6548 let Some(val) = row.columns.get(&idef.column_id) else {
6549 continue;
6550 };
6551 match idef.kind {
6552 IndexKind::Bitmap => {
6553 if let Some(b) = bitmap.get_mut(&idef.column_id) {
6554 b.insert(val.encode_key(), row.row_id);
6555 }
6556 }
6557 IndexKind::Ann => {
6558 if let (Some(a), Value::Embedding(v)) = (ann.get_mut(&idef.column_id), val) {
6559 a.insert(v, row.row_id);
6560 }
6561 }
6562 IndexKind::FmIndex => {
6563 if let (Some(f), Value::Bytes(b)) = (fm.get_mut(&idef.column_id), val) {
6564 f.insert(b.clone(), row.row_id);
6565 }
6566 }
6567 IndexKind::Sparse => {
6568 if let (Some(s), Value::Bytes(b)) = (sparse.get_mut(&idef.column_id), val) {
6569 if let Ok(terms) = bincode::deserialize::<Vec<(u32, f32)>>(b) {
6572 s.insert(&terms, row.row_id);
6573 }
6574 }
6575 }
6576 IndexKind::MinHash => {
6577 if let (Some(mh), Value::Bytes(b)) = (minhash.get_mut(&idef.column_id), val) {
6578 let tokens = crate::index::token_hashes_from_bytes(b);
6581 mh.insert(&tokens, row.row_id);
6582 }
6583 }
6584 _ => {}
6585 }
6586 }
6587 if let Some(pk_col) = schema.primary_key() {
6588 if let Some(pk_val) = row.columns.get(&pk_col.id) {
6589 hot.insert(pk_val.encode_key(), row.row_id);
6590 }
6591 }
6592}
6593
6594#[allow(dead_code)]
6600fn bulk_index_key(
6601 column_keys: &HashMap<u16, ([u8; 32], u8)>,
6602 column_id: u16,
6603 ty: TypeId,
6604 col: &columnar::NativeColumn,
6605 i: usize,
6606) -> Option<Vec<u8>> {
6607 let encoded = columnar::encode_key_native(ty, col, i)?;
6608 #[cfg(feature = "encryption")]
6609 {
6610 use crate::encryption::{hmac_token, ope_token_f64, ope_token_i64, SCHEME_HMAC_EQ};
6611 if let Some((key, scheme)) = column_keys.get(&column_id) {
6612 return Some(match (*scheme, col) {
6613 (SCHEME_HMAC_EQ, _) => hmac_token(key, &encoded).to_vec(),
6614 (_, columnar::NativeColumn::Int64 { data, .. }) => {
6615 ope_token_i64(key, data[i]).to_vec()
6616 }
6617 (_, columnar::NativeColumn::Float64 { data, .. }) => {
6618 ope_token_f64(key, data[i]).to_vec()
6619 }
6620 _ => hmac_token(key, &encoded).to_vec(),
6621 });
6622 }
6623 }
6624 #[cfg(not(feature = "encryption"))]
6625 {
6626 let _ = (column_id, column_keys, col);
6627 }
6628 Some(encoded)
6629}
6630
6631pub(crate) fn write_schema(dir: &Path, schema: &Schema) -> Result<()> {
6632 let json = serde_json::to_string_pretty(schema)
6633 .map_err(|e| MongrelError::Schema(format!("encode schema: {e}")))?;
6634 std::fs::write(dir.join(SCHEMA_FILENAME), json)?;
6635 Ok(())
6636}
6637
6638fn read_schema(dir: &Path) -> Result<Schema> {
6639 serde_json::from_str(&std::fs::read_to_string(dir.join(SCHEMA_FILENAME))?)
6640 .map_err(|e| MongrelError::Schema(format!("decode schema: {e}")))
6641}
6642
6643fn next_wal_segment(wal_dir: &Path) -> Result<PathBuf> {
6644 Ok(wal_dir.join(format!("seg-{:06}.wal", next_wal_number(wal_dir)?)))
6645}
6646
6647fn latest_wal_segment(wal_dir: &Path) -> Result<Option<PathBuf>> {
6648 let n = list_wal_numbers(wal_dir)?;
6649 Ok(n.map(|max| wal_dir.join(format!("seg-{max:06}.wal"))))
6650}
6651
6652fn next_wal_number(wal_dir: &Path) -> Result<u32> {
6653 Ok(list_wal_numbers(wal_dir)?.map(|m| m + 1).unwrap_or(0))
6654}
6655
6656fn list_wal_numbers(wal_dir: &Path) -> Result<Option<u32>> {
6657 let _ = std::fs::create_dir_all(wal_dir);
6658 let mut max_n = None;
6659 for entry in std::fs::read_dir(wal_dir)? {
6660 let entry = entry?;
6661 let fname = entry.file_name();
6662 let Some(s) = fname.to_str() else {
6663 continue;
6664 };
6665 let Some(stripped) = s.strip_prefix("seg-") else {
6666 continue;
6667 };
6668 let Some(stripped) = stripped.strip_suffix(".wal") else {
6669 continue;
6670 };
6671 if let Ok(n) = stripped.parse::<u32>() {
6672 max_n = Some(max_n.map(|m: u32| m.max(n)).unwrap_or(n));
6673 }
6674 }
6675 Ok(max_n)
6676}