Expose Ingestion API with Inversion of Control

Author: zaidoon1

src/abstract.rs

@@ -6,7 +6,7 @@ use crate::{
     blob_tree::FragmentationMap, compaction::CompactionStrategy, config::TreeType,
     iter_guard::IterGuardImpl, table::Table, tree::inner::MemtableId, version::Version,
     vlog::BlobFile, AnyTree, BlobTree, Config, Guard, InternalValue, KvPair, Memtable, SeqNo,
-    SequenceNumberCounter, TableId, Tree, TreeId, UserKey, UserValue,
+    TableId, Tree, TreeId, UserKey, UserValue,
 };
 use enum_dispatch::enum_dispatch;
 use std::{ops::RangeBounds, sync::Arc};
@@ -72,27 +72,6 @@ pub trait AbstractTree {
         index: Option<Arc<Memtable>>,
     ) -> Box<dyn DoubleEndedIterator<Item = IterGuardImpl> + Send + 'static>;
 
-    /// Ingests a sorted stream of key-value pairs into the tree.
-    ///
-    /// Can only be called on a new fresh, empty tree.
-    ///
-    /// # Errors
-    ///
-    /// Will return `Err` if an IO error occurs.
-    ///
-    /// # Panics
-    ///
-    /// Panics if the tree is **not** initially empty.
-    ///
-    /// Will panic if the input iterator is not sorted in ascending order.
-    #[doc(hidden)]
-    fn ingest(
-        &self,
-        iter: impl Iterator<Item = (UserKey, UserValue)>,
-        seqno_generator: &SequenceNumberCounter,
-        visible_seqno: &SequenceNumberCounter,
-    ) -> crate::Result<()>;
-
     /// Returns the approximate number of tombstones in the tree.
     fn tombstone_count(&self) -> u64;
 

src/blob_tree/ingest.rs

@@ -0,0 +1,194 @@
+use crate::{
+    blob_tree::handle::BlobIndirection,
+    file::BLOBS_FOLDER,
+    table::Table,
+    tree::ingest::Ingestion as TableIngestion,
+    vlog::{BlobFileWriter, ValueHandle},
+    SeqNo, UserKey, UserValue,
+};
+
+/// Bulk ingestion for BlobTree
+///
+/// Items NEED to be added in ascending key order.
+pub struct BlobIngestion<'a> {
+    tree: &'a crate::BlobTree,
+    pub(crate) table: TableIngestion<'a>,
+    pub(crate) blob: BlobFileWriter,
+    seqno: SeqNo,
+    separation_threshold: u32,
+    last_key: Option<UserKey>,
+}
+
+impl<'a> BlobIngestion<'a> {
+    /// Creates a new ingestion.
+    ///
+    /// # Errors
+    ///
+    /// Will return `Err` if an IO error occurs.
+    pub fn new(tree: &'a crate::BlobTree) -> crate::Result<Self> {
+        let blob_file_size = tree
+            .index
+            .config
+            .kv_separation_opts
+            .as_ref()
+            .expect("kv separation options should exist")
+            .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;
+
+        Ok(Self {
+            tree,
+            table,
+            blob,
+            seqno: 0,
+            separation_threshold,
+            last_key: None,
+        })
+    }
+
+    /// Sets the ingestion seqno.
+    #[must_use]
+    pub fn with_seqno(mut self, seqno: SeqNo) -> Self {
+        self.seqno = seqno;
+        self.table = self.table.with_seqno(seqno);
+        self
+    }
+
+    /// Writes a key-value pair.
+    ///
+    /// # Errors
+    ///
+    /// Will return `Err` if an IO error occurs.
+    pub fn write(&mut self, key: UserKey, value: UserValue) -> crate::Result<()> {
+        // Check order before any blob I/O to avoid partial writes on failure
+        if let Some(prev) = &self.last_key {
+            assert!(
+                key > *prev,
+                "next key in ingestion must be greater than last key"
+            );
+        }
+
+        #[allow(clippy::cast_possible_truncation)]
+        let value_size = value.len() as u32;
+
+        if value_size >= self.separation_threshold {
+            let offset = self.blob.offset();
+            let blob_file_id = self.blob.blob_file_id();
+            let on_disk_size = self.blob.write(&key, self.seqno, &value)?;
+
+            let indirection = BlobIndirection {
+                vhandle: ValueHandle {
+                    blob_file_id,
+                    offset,
+                    on_disk_size,
+                },
+                size: value_size,
+            };
+
+            let cloned_key = key.clone();
+            let res = self.table.write_indirection(key, indirection);
+            if res.is_ok() {
+                self.last_key = Some(cloned_key);
+            }
+            res
+        } else {
+            let cloned_key = key.clone();
+            let res = self.table.write(key, value);
+            if res.is_ok() {
+                self.last_key = Some(cloned_key);
+            }
+            res
+        }
+    }
+
+    /// Writes a tombstone for a key.
+    ///
+    /// # Errors
+    ///
+    /// Will return `Err` if an IO error occurs.
+    pub fn write_tombstone(&mut self, key: UserKey) -> crate::Result<()> {
+        if let Some(prev) = &self.last_key {
+            assert!(
+                key > *prev,
+                "next key in ingestion must be greater than last key"
+            );
+        }
+
+        let cloned_key = key.clone();
+        let res = self.table.write_tombstone(key);
+        if res.is_ok() {
+            self.last_key = Some(cloned_key);
+        }
+        res
+    }
+
+    /// Finishes the ingestion.
+    ///
+    /// # Errors
+    ///
+    /// Will return `Err` if an IO error occurs.
+    pub fn finish(self) -> crate::Result<()> {
+        use crate::AbstractTree;
+
+        // Capture required handles before consuming fields during finalization
+        let index = self.index().clone();
+        let tree = self.tree.clone();
+
+        let blob_files = self.blob.finish()?;
+        let results = self.table.writer.finish()?;
+
+        let pin_filter = index.config.filter_block_pinning_policy.get(0);
+        let pin_index = index.config.index_block_pinning_policy.get(0);
+
+        let created_tables = results
+            .into_iter()
+            .map(|(table_id, checksum)| -> crate::Result<Table> {
+                Table::recover(
+                    index
+                        .config
+                        .path
+                        .join(crate::file::TABLES_FOLDER)
+                        .join(table_id.to_string()),
+                    checksum,
+                    index.id,
+                    index.config.cache.clone(),
+                    index.config.descriptor_table.clone(),
+                    pin_filter,
+                    pin_index,
+                    #[cfg(feature = "metrics")]
+                    index.metrics.clone(),
+                )
+            })
+            .collect::<crate::Result<Vec<_>>>()?;
+
+        tree.register_tables(&created_tables, Some(&blob_files), None)?;
+
+        Ok(())
+    }
+
+    #[inline]
+    fn index(&self) -> &crate::Tree {
+        &self.tree.index
+    }
+}

src/blob_tree/mod.rs

@@ -4,6 +4,7 @@
 
 mod gc;
 pub mod handle;
+pub mod ingest;
 
 #[doc(hidden)]
 pub use gc::{FragmentationEntry, FragmentationMap};
@@ -251,122 +252,6 @@ impl AbstractTree for BlobTree {
         self.index.drop_range(range)
     }
 
-    fn ingest(
-        &self,
-        iter: impl Iterator<Item = (UserKey, UserValue)>,
-        seqno_generator: &SequenceNumberCounter,
-        visible_seqno: &SequenceNumberCounter,
-    ) -> crate::Result<()> {
-        use crate::{compaction::MoveDown, tree::ingest::Ingestion};
-        use std::time::Instant;
-
-        let seqno = seqno_generator.next();
-
-        let blob_file_size = self
-            .index
-            .config
-            .kv_separation_opts
-            .as_ref()
-            .expect("kv separation options should exist")
-            .file_target_size;
-
-        let mut table_writer = Ingestion::new(&self.index)?.with_seqno(seqno);
-        let mut blob_writer = BlobFileWriter::new(
-            self.index.0.blob_file_id_generator.clone(),
-            blob_file_size,
-            self.index.config.path.join(BLOBS_FOLDER),
-        )?
-        .use_compression(
-            self.index
-                .config
-                .kv_separation_opts
-                .as_ref()
-                .expect("blob options should exist")
-                .compression,
-        );
-
-        let start = Instant::now();
-        let mut count = 0;
-        let mut last_key = None;
-
-        let separation_threshold = self
-            .index
-            .config
-            .kv_separation_opts
-            .as_ref()
-            .expect("kv separation options should exist")
-            .separation_threshold;
-
-        for (key, value) in iter {
-            if let Some(last_key) = &last_key {
-                assert!(
-                    key > last_key,
-                    "next key in bulk ingest was not greater than last key",
-                );
-            }
-            last_key = Some(key.clone());
-
-            #[expect(clippy::cast_possible_truncation, reason = "values are 32-bit max")]
-            let value_size = value.len() as u32;
-
-            if value_size >= separation_threshold {
-                let offset = blob_writer.offset();
-                let blob_file_id = blob_writer.blob_file_id();
-                let on_disk_size = blob_writer.write(&key, seqno, &value)?;
-
-                let indirection = BlobIndirection {
-                    vhandle: ValueHandle {
-                        blob_file_id,
-                        offset,
-                        on_disk_size,
-                    },
-                    size: value_size,
-                };
-
-                table_writer.write_indirection(key, indirection)?;
-            } else {
-                table_writer.write(key, value)?;
-            }
-
-            count += 1;
-        }
-
-        let blob_files = blob_writer.finish()?;
-        let results = table_writer.writer.finish()?;
-
-        let pin_filter = self.index.config.filter_block_pinning_policy.get(0);
-        let pin_index = self.index.config.filter_block_pinning_policy.get(0);
-
-        let created_tables = results
-            .into_iter()
-            .map(|(table_id, checksum)| -> crate::Result<Table> {
-                Table::recover(
-                    self.index
-                        .config
-                        .path
-                        .join(crate::file::TABLES_FOLDER)
-                        .join(table_id.to_string()),
-                    checksum,
-                    self.index.id,
-                    self.index.config.cache.clone(),
-                    self.index.config.descriptor_table.clone(),
-                    pin_filter,
-                    pin_index,
-                    #[cfg(feature = "metrics")]
-                    self.index.metrics.clone(),
-                )
-            })
-            .collect::<crate::Result<Vec<_>>>()?;
-
-        self.register_tables(&created_tables, Some(&blob_files), None)?;
-
-        visible_seqno.fetch_max(seqno + 1);
-
-        log::info!("Ingested {count} items in {:?}", start.elapsed());
-
-        Ok(())
-    }
-
     fn major_compact(&self, target_size: u64, seqno_threshold: SeqNo) -> crate::Result<()> {
         self.index.major_compact(target_size, seqno_threshold)
     }

src/lib.rs

@@ -221,13 +221,13 @@ pub use {
     merge::BoxedIterator,
     slice::Builder,
     table::{GlobalTableId, Table, TableId},
-    tree::ingest::Ingestion,
     tree::inner::TreeId,
     value::InternalValue,
 };
 
 pub use {
     any_tree::AnyTree,
+    blob_tree::ingest::BlobIngestion,
     blob_tree::BlobTree,
     cache::Cache,
     compression::CompressionType,
@@ -240,6 +240,7 @@ pub use {
     r#abstract::AbstractTree,
     seqno::SequenceNumberCounter,
     slice::Slice,
+    tree::ingest::Ingestion,
     tree::Tree,
     value::SeqNo,
     value_type::ValueType,