unify ingestion API across tree types

zaidoon1 · zaidoon1 · commit fda555310cf1 · 2025-11-08T15:03:29.000-05:00
diff --git a/src/any_tree.rs b/src/any_tree.rs
@@ -2,7 +2,10 @@
 // This source code is licensed under both the Apache 2.0 and MIT License
 // (found in the LICENSE-* files in the repository)
 
-use crate::{BlobTree, Tree};
+use crate::{
+    blob_tree::ingest::BlobIngestion, tree::ingest::Ingestion, BlobTree, SeqNo, Tree, UserKey,
+    UserValue,
+};
 use enum_dispatch::enum_dispatch;
 
 /// May be a standard [`Tree`] or a [`BlobTree`]
@@ -15,3 +18,56 @@ pub enum AnyTree {
     /// Key-value separated LSM-tree, see [`BlobTree`]
     Blob(BlobTree),
 }
+
+/// Unified ingestion builder over AnyTree
+pub enum AnyIngestion<'a> {
+    /// Ingestion for a standard LSM-tree
+    Standard(Ingestion<'a>),
+    /// Ingestion for a BlobTree with KV separation
+    Blob(BlobIngestion<'a>),
+}
+
+impl<'a> AnyIngestion<'a> {
+    #[must_use]
+    /// Sets the sequence number used for subsequent writes
+    pub fn with_seqno(self, seqno: SeqNo) -> Self {
+        match self {
+            AnyIngestion::Standard(i) => AnyIngestion::Standard(i.with_seqno(seqno)),
+            AnyIngestion::Blob(b) => AnyIngestion::Blob(b.with_seqno(seqno)),
+        }
+    }
+
+    /// Writes a key-value pair
+    pub fn write(&mut self, key: UserKey, value: UserValue) -> crate::Result<()> {
+        match self {
+            AnyIngestion::Standard(i) => i.write(key, value),
+            AnyIngestion::Blob(b) => b.write(key, value),
+        }
+    }
+
+    /// Writes a tombstone for a key
+    pub fn write_tombstone(&mut self, key: UserKey) -> crate::Result<()> {
+        match self {
+            AnyIngestion::Standard(i) => i.write_tombstone(key),
+            AnyIngestion::Blob(b) => b.write_tombstone(key),
+        }
+    }
+
+    /// Finalizes ingestion and registers created tables (and blob files if present)
+    pub fn finish(self) -> crate::Result<()> {
+        match self {
+            AnyIngestion::Standard(i) => i.finish(),
+            AnyIngestion::Blob(b) => b.finish(),
+        }
+    }
+}
+
+impl AnyTree {
+    /// Starts an ingestion for any tree type (standard or blob)
+    pub fn ingestion(&self) -> crate::Result<AnyIngestion<'_>> {
+        match self {
+            AnyTree::Standard(t) => Ok(AnyIngestion::Standard(Ingestion::new(t)?)),
+            AnyTree::Blob(b) => Ok(AnyIngestion::Blob(BlobIngestion::new(b)?)),
+        }
+    }
+}
diff --git a/src/blob_tree/ingest.rs b/src/blob_tree/ingest.rs
@@ -26,36 +26,24 @@ impl<'a> BlobIngestion<'a> {
     ///
     /// Will return `Err` if an IO error occurs.
     pub fn new(tree: &'a crate::BlobTree) -> crate::Result<Self> {
-        let blob_file_size = tree
+        let kv = tree
             .index
             .config
             .kv_separation_opts
             .as_ref()
-            .expect("kv separation options should exist")
-            .file_target_size;
+            .expect("kv separation options should exist");
+
+        let blob_file_size = kv.file_target_size;
 
         let table = TableIngestion::new(&tree.index)?;
         let blob = BlobFileWriter::new(
             tree.index.0.blob_file_id_generator.clone(),
             blob_file_size,
             tree.index.config.path.join(BLOBS_FOLDER),
         )?
-        .use_compression(
-            tree.index
-                .config
-                .kv_separation_opts
-                .as_ref()
-                .expect("blob options should exist")
-                .compression,
-        );
-
-        let separation_threshold = tree
-            .index
-            .config
-            .kv_separation_opts
-            .as_ref()
-            .expect("kv separation options should exist")
-            .separation_threshold;
+        .use_compression(kv.compression);
+
+        let separation_threshold = kv.separation_threshold;
 
         Ok(Self {
             tree,
diff --git a/src/lib.rs b/src/lib.rs
@@ -226,7 +226,7 @@ pub use {
 };
 
 pub use {
-    any_tree::AnyTree,
+    any_tree::{AnyIngestion, AnyTree},
     blob_tree::ingest::BlobIngestion,
     blob_tree::BlobTree,
     cache::Cache,
diff --git a/src/tree/mod.rs b/src/tree/mod.rs
@@ -701,7 +701,7 @@ impl Tree {
         None
     }
 
-    // Look through runs/tables and keep the newest visible version for the key
+    // Scan levels top-down and runs newest-first; return the first table hit
     fn get_internal_entry_from_tables(
         &self,
         version: &Version,
@@ -712,18 +712,13 @@ impl Tree {
         // https://fjall-rs.github.io/post/bloom-filter-hash-sharing/
         let key_hash = crate::table::filter::standard_bloom::Builder::get_hash(key);
 
-        let mut best: Option<InternalValue> = None;
-
         for level in version.iter_levels() {
             for run in level.iter() {
                 // NOTE: Based on benchmarking, binary search is only worth it with ~4 tables
                 if run.len() >= 4 {
                     if let Some(table) = run.get_for_key(key) {
                         if let Some(item) = table.get(key, seqno, key_hash)? {
-                            match &best {
-                                Some(b) if b.key.seqno >= item.key.seqno => {}
-                                _ => best = Some(item),
-                            }
+                            return Ok(Some(item));
                         }
                     }
                 } else {
@@ -734,17 +729,14 @@ impl Tree {
                         }
 
                         if let Some(item) = table.get(key, seqno, key_hash)? {
-                            match &best {
-                                Some(b) if b.key.seqno >= item.key.seqno => {}
-                                _ => best = Some(item),
-                            }
+                            return Ok(Some(item));
                         }
                     }
                 }
             }
         }
 
-        Ok(best)
+        Ok(None)
     }
 
     fn inner_compact(
diff --git a/tests/ingestion_api.rs b/tests/ingestion_api.rs
@@ -1,22 +1,16 @@
-use lsm_tree::{
-    AbstractTree, AnyTree, BlobIngestion, Config, Ingestion, KvSeparationOptions, SeqNo,
-};
+use lsm_tree::{AbstractTree, Config, KvSeparationOptions, SeqNo};
 
 #[test]
 fn tree_ingestion_tombstones_delete_existing_keys() -> lsm_tree::Result<()> {
     let folder = tempfile::tempdir()?;
     let tree = Config::new(folder, Default::default()).open()?;
-    let tree = match tree {
-        AnyTree::Standard(t) => t,
-        _ => unreachable!(),
-    };
 
     for i in 0..10u32 {
         let key = format!("k{:03}", i);
         tree.insert(key.as_bytes(), b"v", 0);
     }
 
-    let mut ingest = Ingestion::new(&tree)?.with_seqno(10);
+    let mut ingest = tree.ingestion()?.with_seqno(10);
     for i in 0..10u32 {
         let key = format!("k{:03}", i);
         ingest.write_tombstone(key.as_bytes().into())?;
@@ -38,17 +32,13 @@ fn blob_tree_ingestion_tombstones_delete_existing_keys() -> lsm_tree::Result<()>
     let tree = Config::new(folder, Default::default())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let tree = match tree {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     for i in 0..8u32 {
         let key = format!("b{:03}", i);
         tree.insert(key.as_bytes(), b"x", 0);
     }
 
-    let mut ingest = BlobIngestion::new(&tree)?.with_seqno(10);
+    let mut ingest = tree.ingestion()?.with_seqno(10);
     for i in 0..8u32 {
         let key = format!("b{:03}", i);
         ingest.write_tombstone(key.as_bytes().into())?;
@@ -68,13 +58,9 @@ fn blob_tree_ingestion_tombstones_delete_existing_keys() -> lsm_tree::Result<()>
 fn tree_ingestion_finish_no_writes_noop() -> lsm_tree::Result<()> {
     let folder = tempfile::tempdir()?;
     let tree = Config::new(folder, Default::default()).open()?;
-    let tree = match tree {
-        AnyTree::Standard(t) => t,
-        _ => unreachable!(),
-    };
 
     let before_tables = tree.table_count();
-    Ingestion::new(&tree)?.finish()?;
+    tree.ingestion()?.finish()?;
     let after_tables = tree.table_count();
 
     assert_eq!(before_tables, after_tables);
@@ -89,10 +75,6 @@ fn blob_ingestion_only_tombstones_does_not_create_blob_files() -> lsm_tree::Resu
     let tree = Config::new(folder, Default::default())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let tree = match tree {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     for i in 0..5u32 {
         let key = format!("d{:03}", i);
@@ -101,7 +83,7 @@ fn blob_ingestion_only_tombstones_does_not_create_blob_files() -> lsm_tree::Resu
 
     let before_blobs = tree.blob_file_count();
 
-    let mut ingest = BlobIngestion::new(&tree)?.with_seqno(10);
+    let mut ingest = tree.ingestion()?.with_seqno(10);
     for i in 0..5u32 {
         let key = format!("d{:03}", i);
         ingest.write_tombstone(key.as_bytes().into())?;
@@ -125,15 +107,11 @@ fn blob_ingestion_finish_no_writes_noop() -> lsm_tree::Result<()> {
     let tree = Config::new(folder, Default::default())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let tree = match tree {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     let before_tables = tree.table_count();
     let before_blobs = tree.blob_file_count();
 
-    BlobIngestion::new(&tree)?.finish()?;
+    tree.ingestion()?.finish()?;
 
     let after_tables = tree.table_count();
     let after_blobs = tree.blob_file_count();
@@ -151,12 +129,8 @@ fn blob_ingestion_separates_large_values_and_reads_ok() -> lsm_tree::Result<()>
     let tree = Config::new(folder, Default::default())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(8)))
         .open()?;
-    let tree = match tree {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
-    let mut ingest = BlobIngestion::new(&tree)?.with_seqno(1);
+    let mut ingest = tree.ingestion()?.with_seqno(1);
     ingest.write("k_big1".as_bytes().into(), vec![1u8; 16].into())?;
     ingest.write("k_big2".as_bytes().into(), vec![2u8; 32].into())?;
     ingest.write("k_small".as_bytes().into(), b"abc".as_slice().into())?;
diff --git a/tests/tree_bulk_ingest.rs b/tests/tree_bulk_ingest.rs
@@ -1,7 +1,4 @@
-use lsm_tree::{
-    AbstractTree, AnyTree, BlobIngestion, Config, Guard, Ingestion, KvSeparationOptions, SeqNo,
-    SequenceNumberCounter,
-};
+use lsm_tree::{AbstractTree, Config, Guard, KvSeparationOptions, SeqNo, SequenceNumberCounter};
 use test_log::test;
 
 const ITEM_COUNT: usize = 100_000;
@@ -14,13 +11,9 @@ fn tree_bulk_ingest() -> lsm_tree::Result<()> {
     let visible_seqno = SequenceNumberCounter::default();
 
     let tree = Config::new(folder, seqno.clone()).open()?;
-    let tree = match tree {
-        AnyTree::Standard(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = Ingestion::new(&tree)?.with_seqno(seq);
+    let mut ingestion = tree.ingestion()?.with_seqno(seq);
     for x in 0..ITEM_COUNT as u64 {
         let k = x.to_be_bytes();
         let v = nanoid::nanoid!();
@@ -53,13 +46,9 @@ fn tree_copy() -> lsm_tree::Result<()> {
     let visible_seqno = SequenceNumberCounter::default();
 
     let src = Config::new(folder, seqno.clone()).open()?;
-    let src = match src {
-        AnyTree::Standard(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = Ingestion::new(&src)?.with_seqno(seq);
+    let mut ingestion = src.ingestion()?.with_seqno(seq);
     for x in 0..ITEM_COUNT as u64 {
         let k = x.to_be_bytes();
         let v = nanoid::nanoid!();
@@ -83,13 +72,9 @@ fn tree_copy() -> lsm_tree::Result<()> {
 
     let folder = tempfile::tempdir()?;
     let dest = Config::new(folder, seqno.clone()).open()?;
-    let dest = match dest {
-        AnyTree::Standard(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = Ingestion::new(&dest)?.with_seqno(seq);
+    let mut ingestion = dest.ingestion()?.with_seqno(seq);
     for item in src.iter(SeqNo::MAX, None) {
         let (k, v) = item.into_inner().unwrap();
         ingestion.write(k, v)?;
@@ -123,13 +108,9 @@ fn blob_tree_bulk_ingest() -> lsm_tree::Result<()> {
     let tree = Config::new(folder, seqno.clone())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let tree = match tree {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = BlobIngestion::new(&tree)?.with_seqno(seq);
+    let mut ingestion = tree.ingestion()?.with_seqno(seq);
     for x in 0..ITEM_COUNT as u64 {
         let k = x.to_be_bytes();
         let v = nanoid::nanoid!();
@@ -165,13 +146,9 @@ fn blob_tree_copy() -> lsm_tree::Result<()> {
     let src = Config::new(folder, seqno.clone())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let src = match src {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = BlobIngestion::new(&src)?.with_seqno(seq);
+    let mut ingestion = src.ingestion()?.with_seqno(seq);
     for x in 0..ITEM_COUNT as u64 {
         let k = x.to_be_bytes();
         let v = nanoid::nanoid!();
@@ -198,13 +175,9 @@ fn blob_tree_copy() -> lsm_tree::Result<()> {
     let dest = Config::new(folder, seqno.clone())
         .with_kv_separation(Some(KvSeparationOptions::default().separation_threshold(1)))
         .open()?;
-    let dest = match dest {
-        AnyTree::Blob(t) => t,
-        _ => unreachable!(),
-    };
 
     let seq = seqno.next();
-    let mut ingestion = BlobIngestion::new(&dest)?.with_seqno(seq);
+    let mut ingestion = dest.ingestion()?.with_seqno(seq);
     for item in src.iter(SeqNo::MAX, None) {
         let (k, v) = item.into_inner().unwrap();
         ingestion.write(k, v)?;

Original file line number	Diff line number	Diff line change
`@@ -701,7 +701,7 @@ impl Tree {`
`701`	`701`	`None`
`702`	`702`	`}`
`703`	`703`
`704`		`- // Look through runs/tables and keep the newest visible version for the key`
	`704`	`+ // Scan levels top-down and runs newest-first; return the first table hit`
`705`	`705`	`fn get_internal_entry_from_tables(`
`706`	`706`	`&self,`
`707`	`707`	`version: &Version,`
`@@ -712,18 +712,13 @@ impl Tree {`
`712`	`712`	`// https://fjall-rs.github.io/post/bloom-filter-hash-sharing/`
`713`	`713`	`let key_hash = crate::table::filter::standard_bloom::Builder::get_hash(key);`
`714`	`714`
`715`		`- let mut best: Option<InternalValue> = None;`
`716`		`-`
`717`	`715`	`for level in version.iter_levels() {`
`718`	`716`	`for run in level.iter() {`
`719`	`717`	`// NOTE: Based on benchmarking, binary search is only worth it with ~4 tables`
`720`	`718`	`if run.len() >= 4 {`
`721`	`719`	`if let Some(table) = run.get_for_key(key) {`
`722`	`720`	`if let Some(item) = table.get(key, seqno, key_hash)? {`
`723`		`- match &best {`
`724`		`- Some(b) if b.key.seqno >= item.key.seqno => {}`
`725`		`- _ => best = Some(item),`
`726`		`- }`
	`721`	`+ return Ok(Some(item));`
`727`	`722`	`}`
`728`	`723`	`}`
`729`	`724`	`} else {`
`@@ -734,17 +729,14 @@ impl Tree {`
`734`	`729`	`}`
`735`	`730`
`736`	`731`	`if let Some(item) = table.get(key, seqno, key_hash)? {`
`737`		`- match &best {`
`738`		`- Some(b) if b.key.seqno >= item.key.seqno => {}`
`739`		`- _ => best = Some(item),`
`740`		`- }`
	`732`	`+ return Ok(Some(item));`
`741`	`733`	`}`
`742`	`734`	`}`
`743`	`735`	`}`
`744`	`736`	`}`
`745`	`737`	`}`
`746`	`738`
`747`		`- Ok(best)`
	`739`	`+ Ok(None)`
`748`	`740`	`}`
`749`	`741`
`750`	`742`	`fn inner_compact(`