Merge pull request #119 from fjall-rs/feat/bulk-ingest

Bulk ingestion API

Author: marvin-j97

src/abstract.rs

@@ -18,6 +18,22 @@ pub type RangeItem = crate::Result<KvPair>;
 #[allow(clippy::module_name_repetitions)]
 #[enum_dispatch]
 pub trait AbstractTree {
+    /// 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)>) -> crate::Result<()>;
+
     /// Performs major compaction, blocking the caller until it's done.
     ///
     /// # Errors

src/blob_tree/mod.rs

@@ -230,6 +230,69 @@ impl BlobTree {
 }
 
 impl AbstractTree for BlobTree {
+    fn ingest(&self, iter: impl Iterator<Item = (UserKey, UserValue)>) -> crate::Result<()> {
+        use crate::tree::ingest::Ingestion;
+        use std::time::Instant;
+
+        // NOTE: Lock active memtable so nothing else can be going on while we are bulk loading
+        let lock = self.lock_active_memtable();
+        assert!(
+            lock.is_empty(),
+            "can only perform bulk_ingest on empty trees",
+        );
+
+        let mut segment_writer = Ingestion::new(&self.index)?;
+        let mut blob_writer = self.blobs.get_writer()?;
+
+        let start = Instant::now();
+        let mut count = 0;
+        let mut last_key = None;
+
+        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());
+
+            // NOTE: Values are 32-bit max
+            #[allow(clippy::cast_possible_truncation)]
+            let value_size = value.len() as u32;
+
+            if value_size >= self.index.config.blob_file_separation_threshold {
+                let vhandle = blob_writer.get_next_value_handle();
+
+                let indirection = MaybeInlineValue::Indirect {
+                    vhandle,
+                    size: value_size,
+                };
+                // TODO: use Slice::with_size
+                let mut serialized_indirection = vec![];
+                indirection.encode_into(&mut serialized_indirection)?;
+
+                segment_writer.write(key.clone(), serialized_indirection.into())?;
+
+                blob_writer.write(&key, value)?;
+            } else {
+                // TODO: use Slice::with_size
+                let direct = MaybeInlineValue::Inline(value);
+                let serialized_direct = direct.encode_into_vec();
+                segment_writer.write(key, serialized_direct.into())?;
+            }
+
+            count += 1;
+        }
+
+        self.blobs.register_writer(blob_writer)?;
+        segment_writer.finish()?;
+
+        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)
     }
@@ -367,6 +430,7 @@ impl AbstractTree for BlobTree {
                     vhandle,
                     size: value_size,
                 };
+                // TODO: use Slice::with_size
                 let mut serialized_indirection = vec![];
                 indirection.encode_into(&mut serialized_indirection)?;
 
@@ -375,6 +439,7 @@ impl AbstractTree for BlobTree {
 
                 blob_writer.write(&item.key.user_key, value)?;
             } else {
+                // TODO: use Slice::with_size
                 let direct = MaybeInlineValue::Inline(value);
                 let serialized_direct = direct.encode_into_vec();
                 segment_writer.write(InternalValue::new(item.key, serialized_direct))?;

src/compaction/mod.rs

@@ -8,6 +8,7 @@ pub(crate) mod fifo;
 pub(crate) mod leveled;
 pub(crate) mod maintenance;
 pub(crate) mod major;
+pub(crate) mod movedown;
 pub(crate) mod pulldown;
 pub(crate) mod stream;
 pub(crate) mod tiered;
@@ -22,6 +23,9 @@ use crate::{config::Config, level_manifest::LevelManifest, segment::meta::Segmen
 /// Alias for `Leveled`
 pub type Levelled = Leveled;
 
+#[doc(hidden)]
+pub use movedown::Strategy as MoveDown;
+
 #[doc(hidden)]
 pub use pulldown::Strategy as PullDown;
 

src/compaction/movedown.rs

@@ -0,0 +1,40 @@
+// Copyright (c) 2024-present, fjall-rs
+// This source code is licensed under both the Apache 2.0 and MIT License
+// (found in the LICENSE-* files in the repository)
+
+use super::{Choice, CompactionStrategy, Input};
+use crate::{level_manifest::LevelManifest, Config, HashSet, Segment};
+
+/// Moves down a level into the destination level.
+pub struct Strategy(pub u8, pub u8);
+
+impl CompactionStrategy for Strategy {
+    fn get_name(&self) -> &'static str {
+        "MoveDownCompaction"
+    }
+
+    #[allow(clippy::expect_used)]
+    fn choose(&self, levels: &LevelManifest, _: &Config) -> Choice {
+        let resolved_view = levels.resolved_view();
+
+        let level = resolved_view
+            .get(usize::from(self.0))
+            .expect("level should exist");
+
+        let next_level = resolved_view
+            .get(usize::from(self.1))
+            .expect("next level should exist");
+
+        if next_level.is_empty() {
+            let segment_ids: HashSet<_> = level.segments.iter().map(Segment::id).collect();
+
+            Choice::Move(Input {
+                segment_ids,
+                dest_level: self.1,
+                target_size: 64_000_000,
+            })
+        } else {
+            Choice::DoNothing
+        }
+    }
+}

src/compaction/worker.rs

@@ -43,9 +43,6 @@ pub struct Options {
     /// Levels manifest.
     pub levels: Arc<RwLock<LevelManifest>>,
 
-    /// Sealed memtables (required for temporarily locking).
-    pub sealed_memtables: Arc<RwLock<SealedMemtables>>,
-
     /// Compaction strategy to use.
     pub strategy: Arc<dyn CompactionStrategy>,
 
@@ -63,7 +60,6 @@ impl Options {
             tree_id: tree.id,
             segment_id_generator: tree.segment_id_counter.clone(),
             config: tree.config.clone(),
-            sealed_memtables: tree.sealed_memtables.clone(),
             levels: tree.levels.clone(),
             stop_signal: tree.stop_signal.clone(),
             strategy,

src/key.rs

@@ -11,7 +11,6 @@ use std::{
     cmp::Reverse,
     io::{Read, Write},
 };
-use value_log::Slice;
 use varint_rs::{VarintReader, VarintWriter};
 
 #[derive(Clone, PartialEq, Eq)]

src/tree/ingest.rs

@@ -0,0 +1,136 @@
+// Copyright (c) 2024-present, fjall-rs
+// This source code is licensed under both the Apache 2.0 and MIT License
+// (found in the LICENSE-* files in the repository)
+
+use super::Tree;
+use crate::{
+    file::SEGMENTS_FOLDER,
+    segment::{block_index::BlockIndexImpl, multi_writer::MultiWriter, SegmentInner},
+    AbstractTree, Segment, UserKey, UserValue, ValueType,
+};
+use std::{
+    path::PathBuf,
+    sync::{atomic::AtomicBool, Arc},
+};
+
+pub struct Ingestion<'a> {
+    folder: PathBuf,
+    tree: &'a Tree,
+    writer: MultiWriter,
+}
+
+impl<'a> Ingestion<'a> {
+    pub fn new(tree: &'a Tree) -> crate::Result<Self> {
+        assert_eq!(
+            0,
+            tree.segment_count(),
+            "can only perform bulk_ingest on empty trees",
+        );
+
+        let folder = tree.config.path.join(SEGMENTS_FOLDER);
+        log::debug!("Ingesting into disk segments in {folder:?}");
+
+        let mut writer = MultiWriter::new(
+            tree.segment_id_counter.clone(),
+            128 * 1_024 * 1_024,
+            crate::segment::writer::Options {
+                folder: folder.clone(),
+                data_block_size: tree.config.data_block_size,
+                index_block_size: tree.config.index_block_size,
+                segment_id: 0, /* TODO: unused */
+            },
+        )?
+        .use_compression(tree.config.compression);
+
+        {
+            use crate::segment::writer::BloomConstructionPolicy;
+
+            if tree.config.bloom_bits_per_key >= 0 {
+                writer = writer.use_bloom_policy(BloomConstructionPolicy::BitsPerKey(
+                    tree.config.bloom_bits_per_key.unsigned_abs(),
+                ));
+            } else {
+                writer = writer.use_bloom_policy(BloomConstructionPolicy::BitsPerKey(0));
+            }
+        }
+
+        Ok(Self {
+            folder,
+            tree,
+            writer,
+        })
+    }
+
+    pub fn write(&mut self, key: UserKey, value: UserValue) -> crate::Result<()> {
+        self.writer.write(crate::InternalValue::from_components(
+            key,
+            value,
+            0,
+            ValueType::Value,
+        ))
+    }
+
+    pub fn finish(self) -> crate::Result<()> {
+        use crate::{
+            compaction::MoveDown, segment::block_index::two_level_index::TwoLevelBlockIndex,
+        };
+
+        let results = self.writer.finish()?;
+
+        let created_segments = results
+            .into_iter()
+            .map(|trailer| -> crate::Result<Segment> {
+                let segment_id = trailer.metadata.id;
+                let segment_file_path = self.folder.join(segment_id.to_string());
+
+                let block_index = TwoLevelBlockIndex::from_file(
+                    &segment_file_path,
+                    &trailer.metadata,
+                    trailer.offsets.tli_ptr,
+                    (self.tree.id, segment_id).into(),
+                    self.tree.config.descriptor_table.clone(),
+                    self.tree.config.block_cache.clone(),
+                )?;
+                let block_index = BlockIndexImpl::TwoLevel(block_index);
+                let block_index = Arc::new(block_index);
+
+                Ok(SegmentInner {
+                    tree_id: self.tree.id,
+
+                    descriptor_table: self.tree.config.descriptor_table.clone(),
+                    block_cache: self.tree.config.block_cache.clone(),
+
+                    metadata: trailer.metadata,
+                    offsets: trailer.offsets,
+
+                    #[allow(clippy::needless_borrows_for_generic_args)]
+                    block_index,
+
+                    bloom_filter: Segment::load_bloom(
+                        &segment_file_path,
+                        trailer.offsets.bloom_ptr,
+                    )?,
+
+                    path: segment_file_path,
+                    is_deleted: AtomicBool::default(),
+                }
+                .into())
+            })
+            .collect::<crate::Result<Vec<_>>>()?;
+
+        self.tree.register_segments(&created_segments)?;
+
+        self.tree.compact(Arc::new(MoveDown(0, 6)), 0)?;
+
+        for segment in &created_segments {
+            let segment_file_path = self.folder.join(segment.id().to_string());
+
+            self.tree
+                .config
+                .descriptor_table
+                .insert(&segment_file_path, segment.global_id());
+        }
+
+        Ok(())
+    }
+}