add prefix-based filtering and lazy segment loading to RunReader

Author: zaidoon1

src/run_reader.rs

@@ -29,6 +29,24 @@ impl RunReader {
 
         let (lo, hi) = run.range_indexes(&range)?;
 
+        // Check if any segment in the range might contain data
+        // This provides early filtering for prefix-based queries
+        let segments_in_range = &run[lo..=hi];
+        let mut has_potential_match = false;
+
+        for segment in segments_in_range {
+            // Check if segment might contain data for this range
+            // If segment has no prefix extractor or can't determine, assume it might contain data
+            if segment.might_contain_range(&range) {
+                has_potential_match = true;
+                break;
+            }
+        }
+
+        if !has_potential_match {
+            return None;
+        }
+
         Some(Self::culled(run, range, (Some(lo), Some(hi)), cache_policy))
     }
 
@@ -83,12 +101,22 @@ impl Iterator for RunReader {
                 self.lo += 1;
 
                 if self.lo < self.hi {
-                    self.lo_reader = self
-                        .run
-                        .get(self.lo)
-                        .expect("should exist")
-                        .iter()
-                        .map(|x| Box::new(x) as BoxedIterator); /* .cache_policy(self.cache_policy) */
+                    // Lazily check next segment for potential matches
+                    // This avoids unnecessary I/O for segments that won't contain our prefix
+                    loop {
+                        if self.lo >= self.hi {
+                            break;
+                        }
+
+                        let segment = self.run.get(self.lo).expect("should exist");
+                        if let Some(reader) = segment.iter() {
+                            self.lo_reader = Some(Box::new(reader) as BoxedIterator);
+                            break;
+                        }
+
+                        // Skip this segment as it doesn't contain our range
+                        self.lo += 1;
+                    }
                 }
             } else if let Some(hi_reader) = &mut self.hi_reader {
                 // NOTE: We reached the hi marker, so consume from it instead
@@ -115,12 +143,21 @@ impl DoubleEndedIterator for RunReader {
                 self.hi -= 1;
 
                 if self.lo < self.hi {
-                    self.hi_reader = self
-                        .run
-                        .get(self.hi)
-                        .expect("should exist")
-                        .iter()
-                        .map(|x| Box::new(x) as BoxedIterator); /* .cache_policy(self.cache_policy) */
+                    // Lazily check prev segment for potential matches
+                    loop {
+                        if self.hi <= self.lo {
+                            break;
+                        }
+
+                        let segment = self.run.get(self.hi).expect("should exist");
+                        if let Some(reader) = segment.iter() {
+                            self.hi_reader = Some(Box::new(reader) as BoxedIterator);
+                            break;
+                        }
+
+                        // Skip this segment as it doesn't contain our range
+                        self.hi -= 1;
+                    }
                 }
             } else if let Some(lo_reader) = &mut self.lo_reader {
                 // NOTE: We reached the lo marker, so consume from it instead
@@ -138,7 +175,8 @@ impl DoubleEndedIterator for RunReader {
 #[allow(clippy::expect_used)]
 mod tests {
     use super::*;
-    use crate::{AbstractTree, Slice};
+    use crate::{AbstractTree, Config, Slice};
+    use std::sync::Arc;
     use test_log::test;
 
     #[test]
@@ -299,4 +337,150 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn test_run_reader_prefix_filtering() -> crate::Result<()> {
+        use crate::prefix::FixedPrefixExtractor;
+
+        let tempdir = tempfile::tempdir()?;
+        let tree = Config::new(&tempdir)
+            .prefix_extractor(Arc::new(FixedPrefixExtractor::new(3)))
+            .open()?;
+
+        // Create segments with different prefixes
+        let prefixes = [
+            ["aaa_1", "aaa_2", "aaa_3"],
+            ["bbb_1", "bbb_2", "bbb_3"],
+            ["ccc_1", "ccc_2", "ccc_3"],
+            ["ddd_1", "ddd_2", "ddd_3"],
+        ];
+
+        for batch in prefixes {
+            for id in batch {
+                tree.insert(id, vec![], 0);
+            }
+            tree.flush_active_memtable(0)?;
+        }
+
+        let segments = tree
+            .manifest
+            .read()
+            .expect("lock is poisoned")
+            .iter()
+            .cloned()
+            .collect::<Vec<_>>();
+
+        let run = Arc::new(Run::new(segments));
+
+        // Test 1: Query for non-existent prefix should return None
+        assert!(
+            RunReader::new(
+                run.clone(),
+                UserKey::from("zzz_1")..=UserKey::from("zzz_9"),
+                CachePolicy::Read
+            )
+            .is_none(),
+            "Should return None for non-existent prefix"
+        );
+
+        // Test 2: Query for existing prefix should return reader
+        let reader = RunReader::new(
+            run.clone(),
+            UserKey::from("bbb_1")..=UserKey::from("bbb_3"),
+            CachePolicy::Read,
+        );
+        assert!(reader.is_some(), "Should return reader for existing prefix");
+
+        if let Some(reader) = reader {
+            let items: Vec<_> = reader.flatten().map(|item| item.key.user_key).collect();
+            assert_eq!(items.len(), 3);
+            assert_eq!(items.first(), Some(&Slice::from(*b"bbb_1")));
+            assert_eq!(items.get(1), Some(&Slice::from(*b"bbb_2")));
+            assert_eq!(items.get(2), Some(&Slice::from(*b"bbb_3")));
+        }
+
+        // Test 3: Range query across prefixes with no common prefix
+        let reader = RunReader::new(
+            run,
+            UserKey::from("aaa_3")..=UserKey::from("bbb_1"),
+            CachePolicy::Read,
+        );
+        // Should still work since segments contain the range
+        assert!(reader.is_some());
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_run_reader_lazy_segment_loading() -> crate::Result<()> {
+        use crate::prefix::FixedPrefixExtractor;
+
+        let tempdir = tempfile::tempdir()?;
+        let tree = Config::new(&tempdir)
+            .prefix_extractor(Arc::new(FixedPrefixExtractor::new(4)))
+            .open()?;
+
+        // Create many segments with distinct prefixes
+        let prefixes = [
+            ["pre1_a", "pre1_b", "pre1_c"],
+            ["pre2_a", "pre2_b", "pre2_c"],
+            ["pre3_a", "pre3_b", "pre3_c"],
+            ["pre4_a", "pre4_b", "pre4_c"],
+            ["pre5_a", "pre5_b", "pre5_c"],
+            ["pre6_a", "pre6_b", "pre6_c"],
+        ];
+
+        for batch in prefixes {
+            for id in batch {
+                tree.insert(id, vec![], 0);
+            }
+            tree.flush_active_memtable(0)?;
+        }
+
+        let segments = tree
+            .manifest
+            .read()
+            .expect("lock is poisoned")
+            .iter()
+            .cloned()
+            .collect::<Vec<_>>();
+
+        let run = Arc::new(Run::new(segments));
+
+        // Query for a specific prefix in the middle
+        // Should skip segments without the prefix lazily
+        let reader = RunReader::new(
+            run.clone(),
+            UserKey::from("pre4_a")..=UserKey::from("pre4_c"),
+            CachePolicy::Read,
+        );
+
+        assert!(reader.is_some());
+
+        if let Some(reader) = reader {
+            let items: Vec<_> = reader.flatten().map(|item| item.key.user_key).collect();
+            assert_eq!(items.len(), 3);
+            assert_eq!(items.first(), Some(&Slice::from(*b"pre4_a")));
+            assert_eq!(items.get(1), Some(&Slice::from(*b"pre4_b")));
+            assert_eq!(items.get(2), Some(&Slice::from(*b"pre4_c")));
+        }
+
+        // Query for prefix at the beginning
+        let reader = RunReader::new(
+            run.clone(),
+            UserKey::from("pre1_a")..=UserKey::from("pre1_c"),
+            CachePolicy::Read,
+        );
+        assert!(reader.is_some());
+
+        // Query for prefix at the end
+        let reader = RunReader::new(
+            run,
+            UserKey::from("pre6_a")..=UserKey::from("pre6_c"),
+            CachePolicy::Read,
+        );
+        assert!(reader.is_some());
+
+        Ok(())
+    }
 }

src/segment/mod.rs

@@ -653,10 +653,30 @@ impl Segment {
         self.metadata.key_range.overlaps_with_bounds(bounds)
     }
 
+    /// Returns the seqno range of the `Segment`.
+    #[must_use]
+    pub fn seqno_range(&self) -> (SeqNo, SeqNo) {
+        self.0.metadata.seqnos
+    }
+
     /// Returns the highest sequence number in the segment.
     #[must_use]
     pub fn get_highest_seqno(&self) -> SeqNo {
-        self.metadata.seqnos.1
+        self.0.metadata.seqnos.1
+    }
+
+    /// Checks if this segment might contain data for the given range.
+    /// Returns false only if we can definitively rule out the segment using bloom filters.
+    /// Returns true if the segment might contain data (or if we can't determine).
+    #[must_use]
+    pub fn might_contain_range<R: RangeBounds<UserKey>>(&self, range: &R) -> bool {
+        // If no prefix extractor, we can't use bloom filter optimization
+        if self.prefix_extractor.is_none() {
+            return true;
+        }
+
+        // Check if we can skip this segment based on bloom filter
+        !self.should_skip_range_by_prefix_bloom_filter(range)
     }
 
     /// Returns the amount of tombstone markers in the `Segment`.