From e9f81d05ffe90b8058eeaff7bcff40b3cd18ebe1 Mon Sep 17 00:00:00 2001
From: Max Isom <hi@maxisom.me>
Date: Thu, 10 Oct 2024 17:23:27 -0700
Subject: [PATCH] [ENH]: replace `get_*` methods on memory blockfile impl with
 `get_range()`

---
 rust/blockstore/src/memory/reader_writer.rs | 163 +++---
 rust/blockstore/src/memory/storage.rs       | 552 +++-----------------
 rust/blockstore/src/types.rs                |  18 +-
 3 files changed, 159 insertions(+), 574 deletions(-)

diff --git a/rust/blockstore/src/memory/reader_writer.rs b/rust/blockstore/src/memory/reader_writer.rs
index 633a0bb18a1b..63cf52f8d884 100644
--- a/rust/blockstore/src/memory/reader_writer.rs
+++ b/rust/blockstore/src/memory/reader_writer.rs
@@ -1,3 +1,5 @@
+use std::ops::RangeBounds;
+
 use super::{
     super::{BlockfileError, Key, Value},
     storage::{Readable, Storage, StorageBuilder, StorageManager, Writeable},
@@ -95,87 +97,27 @@ impl<
         }
     }
 
-    #[allow(clippy::type_complexity)]
-    pub(crate) fn get_by_prefix(
-        &'storage self,
-        prefix: &str,
-    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
-        let values = V::get_by_prefix_from_storage(prefix, &self.storage);
-        if values.is_empty() {
-            return Err(Box::new(BlockfileError::NotFoundError));
-        }
-        let values = values
-            .iter()
-            .map(|(key, value)| (K::try_from(&key.key).unwrap(), value.clone()))
-            .collect();
-        Ok(values)
-    }
-
-    #[allow(clippy::type_complexity)]
-    pub(crate) fn get_gt(
-        &'storage self,
-        prefix: &str,
-        key: K,
-    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
-        let key = key.into();
-        let values = V::read_gt_from_storage(prefix, key, &self.storage);
-        if values.is_empty() {
-            return Err(Box::new(BlockfileError::NotFoundError));
-        }
-        let values = values
-            .iter()
-            .map(|(key, value)| (K::try_from(&key.key).unwrap(), value.clone()))
-            .collect();
-        Ok(values)
-    }
-
-    #[allow(clippy::type_complexity)]
-    pub(crate) fn get_lt(
-        &'storage self,
-        prefix: &str,
-        key: K,
-    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
-        let key = key.into();
-        let values = V::read_lt_from_storage(prefix, key, &self.storage);
-        if values.is_empty() {
-            return Err(Box::new(BlockfileError::NotFoundError));
-        }
-        let values = values
-            .iter()
-            .map(|(key, value)| (K::try_from(&key.key).unwrap(), value.clone()))
-            .collect();
-        Ok(values)
-    }
-
-    #[allow(clippy::type_complexity)]
-    pub(crate) fn get_gte(
+    pub(crate) fn get_range<'prefix, PrefixRange, KeyRange>(
         &'storage self,
-        prefix: &str,
-        key: K,
-    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
-        let key = key.into();
-        let values = V::read_gte_from_storage(prefix, key, &self.storage);
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
+    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>>
+    where
+        PrefixRange: RangeBounds<&'prefix str>,
+        KeyRange: RangeBounds<K>,
+    {
+        let values = V::read_range_from_storage(
+            prefix_range,
+            (
+                key_range.start_bound().map(|k| k.clone().into()),
+                key_range.end_bound().map(|k| k.clone().into()),
+            ),
+            &self.storage,
+        );
         if values.is_empty() {
             return Err(Box::new(BlockfileError::NotFoundError));
         }
-        let values = values
-            .iter()
-            .map(|(key, value)| (K::try_from(&key.key).unwrap(), value.clone()))
-            .collect();
-        Ok(values)
-    }
 
-    #[allow(clippy::type_complexity)]
-    pub(crate) fn get_lte(
-        &'storage self,
-        prefix: &str,
-        key: K,
-    ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
-        let key = key.into();
-        let values = V::read_lte_from_storage(prefix, key, &self.storage);
-        if values.is_empty() {
-            return Err(Box::new(BlockfileError::NotFoundError));
-        }
         let values = values
             .iter()
             .map(|(key, value)| (K::try_from(&key.key).unwrap(), value.clone()))
@@ -210,6 +152,8 @@ impl<
 
 #[cfg(test)]
 mod tests {
+    use std::ops::Bound;
+
     use super::*;
     use chroma_types::{Chunk, DataRecord, LogRecord, Operation, OperationRecord};
 
@@ -369,7 +313,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<&str, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_by_prefix("prefix").unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -390,7 +334,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 3);
+        let values = reader.get_range("prefix"..="prefix", (Bound::Excluded(3), Bound::Unbounded));
         assert!(values.is_err());
     }
 
@@ -405,7 +349,9 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 0).unwrap();
+        let values = reader
+            .get_range("prefix"..="prefix", (Bound::Excluded(0), Bound::Unbounded))
+            .unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -429,7 +375,9 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 1).unwrap();
+        let values = reader
+            .get_range("prefix"..="prefix", (Bound::Excluded(1), Bound::Unbounded))
+            .unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -450,7 +398,10 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 3.0);
+        let values = reader.get_range(
+            "prefix"..="prefix",
+            (Bound::Excluded(3.0), Bound::Unbounded),
+        );
         assert!(values.is_err());
     }
 
@@ -465,7 +416,12 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 0.0).unwrap();
+        let values = reader
+            .get_range(
+                "prefix"..="prefix",
+                (Bound::Excluded(0.0), Bound::Unbounded),
+            )
+            .unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -489,7 +445,12 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gt("prefix", 1.0).unwrap();
+        let values = reader
+            .get_range(
+                "prefix"..="prefix",
+                (Bound::Excluded(1.0), Bound::Unbounded),
+            )
+            .unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -510,7 +471,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 4);
+        let values = reader.get_range("prefix"..="prefix", 4..);
         assert!(values.is_err());
     }
 
@@ -525,7 +486,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 1).unwrap();
+        let values = reader.get_range("prefix"..="prefix", 1..).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -549,7 +510,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 2).unwrap();
+        let values = reader.get_range("prefix"..="prefix", 2..).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -570,7 +531,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 3.5);
+        let values = reader.get_range("prefix"..="prefix", 3.5..);
         assert!(values.is_err());
     }
 
@@ -585,7 +546,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 0.5).unwrap();
+        let values = reader.get_range("prefix"..="prefix", 0.5..).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -609,7 +570,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_gte("prefix", 1.5).unwrap();
+        let values = reader.get_range("prefix"..="prefix", 1.5..).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -630,7 +591,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 1);
+        let values = reader.get_range("prefix"..="prefix", ..1);
         assert!(values.is_err());
     }
 
@@ -645,7 +606,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 4).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..4).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -669,7 +630,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 3).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..3).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -690,7 +651,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 0.5);
+        let values = reader.get_range("prefix"..="prefix", ..0.5);
         assert!(values.is_err());
     }
 
@@ -705,7 +666,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 3.5).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..3.5).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -729,7 +690,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lt("prefix", 2.5).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..2.5).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -750,7 +711,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 0);
+        let values = reader.get_range("prefix"..="prefix", ..=0);
         assert!(values.is_err());
     }
 
@@ -765,7 +726,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 3).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..=3).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -789,7 +750,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<u32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 2).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..=2).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
@@ -810,7 +771,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 0.5);
+        let values = reader.get_range("prefix"..="prefix", ..=0.5);
         assert!(values.is_err());
     }
 
@@ -825,7 +786,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 3.0).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..=3.0).unwrap();
         assert_eq!(values.len(), 3);
         assert!(values
             .iter()
@@ -849,7 +810,7 @@ mod tests {
 
         let reader: MemoryBlockfileReader<f32, &str> =
             MemoryBlockfileReader::open(writer.id, storage_manager);
-        let values = reader.get_lte("prefix", 2.0).unwrap();
+        let values = reader.get_range("prefix"..="prefix", ..=2.0).unwrap();
         assert_eq!(values.len(), 2);
         assert!(values
             .iter()
diff --git a/rust/blockstore/src/memory/storage.rs b/rust/blockstore/src/memory/storage.rs
index b8b307b41bcf..b2bb2818946f 100644
--- a/rust/blockstore/src/memory/storage.rs
+++ b/rust/blockstore/src/memory/storage.rs
@@ -20,34 +20,14 @@ pub trait Readable<'referred_data>: Sized {
         storage: &'referred_data Storage,
     ) -> Option<Self>;
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)>;
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)>;
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)>;
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)>;
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)>;
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>;
 
     fn get_at_index(
         storage: &'referred_data Storage,
@@ -103,66 +83,21 @@ impl<'referred_data> Readable<'referred_data> for &'referred_data str {
             .map(|s| s.as_str())
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .string_value_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, v.as_str()))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .string_value_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, v.as_str()))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .string_value_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, v.as_str()))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .string_value_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, v.as_str()))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .string_value_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, v.as_str()))
             .collect()
     }
@@ -238,66 +173,21 @@ impl<'referred_data> Readable<'referred_data> for &'referred_data [u32] {
             .map(|a| a.as_slice())
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .uint32_array_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, v.as_slice()))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .uint32_array_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, v.as_slice()))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .uint32_array_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, v.as_slice()))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .uint32_array_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, v.as_slice()))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .uint32_array_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, v.as_slice()))
             .collect()
     }
@@ -368,66 +258,21 @@ impl<'referred_data> Readable<'referred_data> for RoaringBitmap {
             .cloned()
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .roaring_bitmap_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, v.clone()))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .roaring_bitmap_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, v.clone()))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .roaring_bitmap_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, v.clone()))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .roaring_bitmap_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, v.clone()))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .roaring_bitmap_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, v.clone()))
             .collect()
     }
@@ -493,66 +338,21 @@ impl<'referred_data> Readable<'referred_data> for f32 {
             .copied()
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .f32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .f32_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .f32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .f32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .f32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, *v))
             .collect()
     }
@@ -614,66 +414,21 @@ impl<'referred_data> Readable<'referred_data> for u32 {
             .copied()
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .u32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .u32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .u32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .u32_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .u32_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, *v))
             .collect()
     }
@@ -739,66 +494,21 @@ impl<'referred_data> Readable<'referred_data> for bool {
             .copied()
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .bool_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .bool_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .bool_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .bool_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| (k, *v))
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .bool_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
+            })
             .map(|(k, v)| (k, *v))
             .collect()
     }
@@ -897,121 +607,27 @@ impl<'referred_data> Readable<'referred_data> for DataRecord<'referred_data> {
         })
     }
 
-    fn get_by_prefix_from_storage(
-        prefix: &str,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .data_record_id_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix)
-            .map(|(k, v)| {
-                let embedding = storage.data_record_embedding_storage.get(k).unwrap();
-                let id = v;
-                (
-                    k,
-                    DataRecord {
-                        id,
-                        embedding,
-                        metadata: None,
-                        document: None,
-                    },
-                )
-            })
-            .collect()
-    }
-
-    fn read_gt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
+    fn read_range_from_storage<'prefix, PrefixRange, KeyRange>(
+        prefix_range: PrefixRange,
+        key_range: KeyRange,
         storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
+    ) -> Vec<(&'referred_data CompositeKey, Self)>
+    where
+        PrefixRange: std::ops::RangeBounds<&'prefix str>,
+        KeyRange: std::ops::RangeBounds<KeyWrapper>,
+    {
         storage
             .data_record_id_storage
             .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key > key)
-            .map(|(k, v)| {
-                let embedding = storage.data_record_embedding_storage.get(k).unwrap();
-                let id = v;
-                (
-                    k,
-                    DataRecord {
-                        id,
-                        embedding,
-                        metadata: None,
-                        document: None,
-                    },
-                )
+            .filter(|(k, _)| {
+                prefix_range.contains(&k.prefix.as_str()) && key_range.contains(&k.key)
             })
-            .collect()
-    }
-
-    fn read_gte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .data_record_id_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key >= key)
-            .map(|(k, v)| {
-                let embedding = storage.data_record_embedding_storage.get(k).unwrap();
-                let id = v;
-                (
-                    k,
-                    DataRecord {
-                        id,
-                        embedding,
-                        metadata: None,
-                        document: None,
-                    },
-                )
-            })
-            .collect()
-    }
-
-    fn read_lt_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .data_record_id_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key < key)
-            .map(|(k, v)| {
-                let embedding = storage.data_record_embedding_storage.get(k).unwrap();
-                let id = v;
-                (
-                    k,
-                    DataRecord {
-                        id,
-                        embedding,
-                        metadata: None,
-                        document: None,
-                    },
-                )
-            })
-            .collect()
-    }
-
-    fn read_lte_from_storage(
-        prefix: &str,
-        key: KeyWrapper,
-        storage: &'referred_data Storage,
-    ) -> Vec<(&'referred_data CompositeKey, Self)> {
-        storage
-            .data_record_id_storage
-            .iter()
-            .filter(|(k, _)| k.prefix == prefix && k.key <= key)
             .map(|(k, v)| {
                 let embedding = storage.data_record_embedding_storage.get(k).unwrap();
-                let id = v;
                 (
                     k,
                     DataRecord {
-                        id,
+                        id: v,
                         embedding,
                         metadata: None,
                         document: None,
diff --git a/rust/blockstore/src/types.rs b/rust/blockstore/src/types.rs
index f52960ee8bc5..0fdef4f0ac96 100644
--- a/rust/blockstore/src/types.rs
+++ b/rust/blockstore/src/types.rs
@@ -282,7 +282,7 @@ impl<
         prefix: &str,
     ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
         match self {
-            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_by_prefix(prefix),
+            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_range(prefix..=prefix, ..),
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader.get_range(prefix..=prefix, ..).await
             }
@@ -295,7 +295,9 @@ impl<
         key: K,
     ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
         match self {
-            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_gt(prefix, key),
+            BlockfileReader::MemoryBlockfileReader(reader) => {
+                reader.get_range(prefix..=prefix, (Bound::Excluded(key), Bound::Unbounded))
+            }
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader
                     .get_range(prefix..=prefix, (Bound::Excluded(key), Bound::Unbounded))
@@ -310,7 +312,9 @@ impl<
         key: K,
     ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
         match self {
-            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_lt(prefix, key),
+            BlockfileReader::MemoryBlockfileReader(reader) => {
+                reader.get_range(prefix..=prefix, ..key)
+            }
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader.get_range(prefix..=prefix, ..key).await
             }
@@ -323,7 +327,9 @@ impl<
         key: K,
     ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
         match self {
-            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_gte(prefix, key),
+            BlockfileReader::MemoryBlockfileReader(reader) => {
+                reader.get_range(prefix..=prefix, key..)
+            }
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader.get_range(prefix..=prefix, key..).await
             }
@@ -336,7 +342,9 @@ impl<
         key: K,
     ) -> Result<Vec<(K, V)>, Box<dyn ChromaError>> {
         match self {
-            BlockfileReader::MemoryBlockfileReader(reader) => reader.get_lte(prefix, key),
+            BlockfileReader::MemoryBlockfileReader(reader) => {
+                reader.get_range(prefix..=prefix, ..=key)
+            }
             BlockfileReader::ArrowBlockfileReader(reader) => {
                 reader.get_range(prefix..=prefix, ..=key).await
             }