address more pr feedback

Author: zaidoon1

src/compaction/pulldown.rs

@@ -2,8 +2,8 @@
 // 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, segment::Segment, Config, HashSet};
+use super::{Choice, CompactionStrategy};
+use crate::{config::Config, level_manifest::LevelManifest};
 
 /// Pulls down and merges a level into the destination level.
 ///

src/compaction/tiered.rs

@@ -2,8 +2,8 @@
 // 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 as CompactionInput};
-use crate::{level_manifest::LevelManifest, segment::Segment, Config, HashSet};
+use super::{Choice, CompactionStrategy};
+use crate::{config::Config, level_manifest::LevelManifest};
 
 fn desired_level_size_in_bytes(level_idx: u8, ratio: u8, base_size: u32) -> usize {
     (ratio as usize).pow(u32::from(level_idx + 1)) * (base_size as usize)

src/compaction/worker.rs

@@ -12,7 +12,7 @@ use crate::{
     level_manifest::LevelManifest,
     merge::Merger,
     run_scanner::RunScanner,
-    segment::{multi_writer::MultiWriter, Segment},
+    segment::{multi_writer::MultiWriter, RecoverParams, Segment},
     stop_signal::StopSignal,
     tree::inner::TreeId,
     Config, InternalValue, SegmentId, SeqNo,
@@ -359,17 +359,17 @@ fn merge_segments(
     let created_segments = writer_results
         .into_iter()
         .map(|segment_id| -> crate::Result<Segment> {
-            Segment::recover(
-                segments_base_folder.join(segment_id.to_string()),
-                opts.tree_id,
-                opts.config.cache.clone(),
-                opts.config.descriptor_table.clone(),
-                opts.config.prefix_extractor.clone(),
-                payload.dest_level <= 1, // TODO: look at configuration
-                payload.dest_level <= 2, // TODO: look at configuration
+            Segment::recover(RecoverParams {
+                file_path: segments_base_folder.join(segment_id.to_string()),
+                tree_id: opts.tree_id,
+                cache: opts.config.cache.clone(),
+                descriptor_table: opts.config.descriptor_table.clone(),
+                prefix_extractor: opts.config.prefix_extractor.clone(),
+                pin_filter: payload.dest_level <= 1, // TODO: look at configuration
+                pin_index: payload.dest_level <= 2,  // TODO: look at configuration
                 #[cfg(feature = "metrics")]
-                opts.metrics.clone(),
-            )
+                metrics: opts.metrics.clone(),
+            })
 
             /* let segment_id = trailer.metadata.id;
             let segment_file_path = segments_base_folder.join(segment_id.to_string());

src/prefix.rs

@@ -29,30 +29,27 @@ use std::sync::Arc;
 /// }
 /// ```
 ///
-/// ## Segmented prefixes (e.g., `account_id#user_id)`:
+/// ## Domain-based prefixes (e.g., `example.com@user`):
 /// ```
 /// use lsm_tree::prefix::PrefixExtractor;
 ///
-/// struct SegmentedPrefixExtractor;
+/// struct DomainPrefixExtractor;
 ///
-/// impl PrefixExtractor for SegmentedPrefixExtractor {
+/// impl PrefixExtractor for DomainPrefixExtractor {
 ///     fn extract<'a>(&self, key: &'a [u8]) -> Box<dyn Iterator<Item = &'a [u8]> + 'a> {
-///         let mut prefixes = vec![];
-///         let mut end = 0;
-///         for (i, &byte) in key.iter().enumerate() {
-///             if byte == b'#' {
-///                 prefixes.push(&key[0..i]);
-///                 end = i;
-///             }
-///         }
-///         if end < key.len() {
-///             prefixes.push(key);
+///         // Extract domain prefix before '@' separator
+///         // This allows efficient scans over all users in a domain
+///         // e.g., "example.com@alice" -> "example.com"
+///         if let Some(pos) = key.iter().position(|&b| b == b'@') {
+///             Box::new(std::iter::once(&key[0..pos]))
+///         } else {
+///             // If no separator, use the full key
+///             Box::new(std::iter::once(key))
 ///         }
-///         Box::new(prefixes.into_iter())
 ///     }
 ///     
 ///     fn name(&self) -> &str {
-///         "segmented_prefix"
+///         "domain_prefix"
 ///     }
 /// }
 /// ```
@@ -66,6 +63,12 @@ pub trait PrefixExtractor: Send + Sync {
 
     /// Returns a unique name for this prefix extractor.
     fn name(&self) -> &str;
+
+    /// Returns the first prefix for a key, if any.
+    /// This is an optimization to avoid boxing for the common case.
+    fn extract_first<'a>(&self, key: &'a [u8]) -> Option<&'a [u8]> {
+        self.extract(key).next()
+    }
 }
 
 /// A prefix extractor that returns the full key.
@@ -81,6 +84,10 @@ impl PrefixExtractor for FullKeyExtractor {
     fn name(&self) -> &'static str {
         "full_key"
     }
+
+    fn extract_first<'a>(&self, key: &'a [u8]) -> Option<&'a [u8]> {
+        Some(key)
+    }
 }
 
 /// A prefix extractor that returns a fixed-length prefix.
@@ -112,6 +119,14 @@ impl PrefixExtractor for FixedPrefixExtractor {
     fn name(&self) -> &'static str {
         "fixed_prefix"
     }
+
+    fn extract_first<'a>(&self, key: &'a [u8]) -> Option<&'a [u8]> {
+        if key.len() <= self.length {
+            Some(key)
+        } else {
+            key.get(0..self.length)
+        }
+    }
 }
 
 /// A prefix extractor that requires keys to be at least a certain length.