Add _with_params variants to query methods

4 files changed, 214 insertions, 12 deletions

diff --git a/log_db/src/common.rs b/log_db/src/common.rs
index 2b40f52..9bafdd4 100644
--- a/log_db/src/common.rs
+++ b/log_db/src/common.rs
@@ -562,3 +562,14 @@ impl<T> RangeBounds<T> for OwnedBounds<T> {
         self.end.as_ref()
     }
 }
+
+#[derive(Debug, Clone)]
+pub struct QueryParams {
+    pub offset: usize,
+    pub limit: usize,
+}
+
+pub static DEFAULT_QUERY_PARAMS: QueryParams = QueryParams {
+    offset: 0,
+    limit: usize::MAX,
+};
diff --git a/log_db/src/engine.rs b/log_db/src/engine.rs
index d593053..c2f9a52 100644
--- a/log_db/src/engine.rs
+++ b/log_db/src/engine.rs
@@ -249,6 +249,7 @@ impl<T, F: Eq + Clone> Engine<T, F> {
         &mut self,
         field: &F,
         values: impl Iterator<Item = &'a Value>,
+        params: &QueryParams,
     ) -> DBResult<Vec<(usize, Record)>> {
         let indexables = values
             .map(|value| {
@@ -305,7 +306,10 @@ impl<T, F: Eq + Clone> Engine<T, F> {
             tagged.extend(mapped);
         }
 
-        let tagged_records = self.read_tagged_log_keys(tagged.into_iter())?;
+        let bound_low = params.offset;
+        let bound_high = (params.offset + params.limit).min(tagged.len());
+        let sliced = &tagged[bound_low..bound_high];
+        let tagged_records = self.read_tagged_log_keys(sliced.into_iter())?;
 
         debug!("Read {} records", tagged_records.len());
 
@@ -316,7 +320,7 @@ impl<T, F: Eq + Clone> Engine<T, F> {
     /// The log keys are accompanied by an integer tag that can be used to identify and group them later.
     fn read_tagged_log_keys<'a>(
         &self,
-        log_keys: impl Iterator<Item = (usize, &'a LogKey)>,
+        log_keys: impl Iterator<Item = &'a (usize, &'a LogKey)>,
     ) -> DBResult<Vec<(usize, Record)>> {
         let mut records = vec![];
         let mut log_keys_map = BTreeMap::new();
@@ -363,7 +367,7 @@ impl<T, F: Eq + Clone> Engine<T, F> {
                 data_file.read_exact(&mut data_buf)?;
 
                 let record = Record::deserialize(&data_buf);
-                records.push((tag, record));
+                records.push((*tag, record));
 
                 current_metadata_offset = new_metadata_offset + row_length;
             }
@@ -376,6 +380,7 @@ impl<T, F: Eq + Clone> Engine<T, F> {
         &mut self,
         field: &F,
         range: B,
+        params: &QueryParams,
     ) -> DBResult<Vec<Record>> {
         fn range_bound_to_indexable(bound: Bound<&Value>) -> DBResult<Bound<IndexableValue>> {
             match bound {
@@ -415,9 +420,13 @@ impl<T, F: Eq + Clone> Engine<T, F> {
             self.secondary_memtables[index].range(indexable_bounds)
         };
 
-        let log_key_batches = log_keys.into_iter().map(|log_key| (0, log_key));
+        let log_key_batches: Vec<(usize, &LogKey)> =
+            log_keys.into_iter().map(|log_key| (0, log_key)).collect();
 
-        let tagged_records = self.read_tagged_log_keys(log_key_batches);
+        let bound_low = params.offset;
+        let bound_high = (params.offset + params.limit).min(log_key_batches.len());
+        let sliced = &log_key_batches[bound_low..bound_high];
+        let tagged_records = self.read_tagged_log_keys(sliced.into_iter());
 
         Ok(tagged_records?.into_iter().map(|(_, rec)| rec).collect())
     }
@@ -451,7 +460,7 @@ impl<T, F: Eq + Clone> Engine<T, F> {
 
     pub fn delete_by_field(&mut self, field: &F, value: &Value) -> DBResult<Vec<Record>> {
         let recs: Vec<Record> = self
-            .batch_find_by_records(field, std::iter::once(value))?
+            .batch_find_by_records(field, std::iter::once(value), &DEFAULT_QUERY_PARAMS)?
             .into_iter()
             .map(|(_, mut rec)| {
                 rec.tombstone = true;
diff --git a/log_db/src/lib.rs b/log_db/src/lib.rs
index 05db3dc..753e41c 100644
--- a/log_db/src/lib.rs
+++ b/log_db/src/lib.rs
@@ -25,7 +25,7 @@ mod memtable_primary;
 mod memtable_secondary;
 mod record;
 
-pub use common::{DBError, DBResult, OwnedBounds, Value};
+pub use common::{DBError, DBResult, OwnedBounds, QueryParams, Value, DEFAULT_QUERY_PARAMS};
 pub use config::{ReadConsistency, Schema, WriteDurability};
 
 use common::*;
@@ -72,6 +72,7 @@ impl<T, F: Eq + Clone> DB<T, F> {
                 // TODO: This clone is only here to appease the borrow checker
                 &engine.config.primary_key.clone(),
                 std::iter::once(value),
+                &DEFAULT_QUERY_PARAMS,
             )
         })?;
 
@@ -86,7 +87,24 @@ impl<T, F: Eq + Clone> DB<T, F> {
     /// Get a collection of records based on an indexed field value.
     pub fn find_by(&mut self, field: &F, value: &Value) -> DBResult<Vec<T>> {
         let recs = self.engine.with_shared_lock(|engine| {
-            engine.batch_find_by_records(field, std::iter::once(value))
+            engine.batch_find_by_records(field, std::iter::once(value), &DEFAULT_QUERY_PARAMS)
+        })?;
+
+        Ok(recs
+            .into_iter()
+            .map(|(_, rec)| (self.engine.config.from_record)(rec.values))
+            .collect())
+    }
+
+    /// Get a collection of records based on an indexed field value, with additional parameters.
+    pub fn find_by_with_params(
+        &mut self,
+        field: &F,
+        value: &Value,
+        params: &QueryParams,
+    ) -> DBResult<Vec<T>> {
+        let recs = self.engine.with_shared_lock(|engine| {
+            engine.batch_find_by_records(field, std::iter::once(value), params)
         })?;
 
         Ok(recs
@@ -99,9 +117,28 @@ impl<T, F: Eq + Clone> DB<T, F> {
     /// Returns a vector of pairs where the first value is an index into the given sequence of values,
     /// and the second value is the record.
     pub fn batch_find_by(&mut self, field: &F, values: &[Value]) -> DBResult<Vec<(usize, T)>> {
-        let recs = self
-            .engine
-            .with_shared_lock(|engine| engine.batch_find_by_records(field, values.iter()))?;
+        let recs = self.engine.with_shared_lock(|engine| {
+            engine.batch_find_by_records(field, values.iter(), &DEFAULT_QUERY_PARAMS)
+        })?;
+
+        Ok(recs
+            .into_iter()
+            .map(|(tag, rec)| (tag, (self.engine.config.from_record)(rec.values)))
+            .collect())
+    }
+
+    /// Get a collection of records based on a sequence of indexed field values, with additional parameters.
+    /// Returns a vector of pairs where the first value is an index into the given sequence of values,
+    /// and the second value is the record.
+    pub fn batch_find_by_with_params(
+        &mut self,
+        field: &F,
+        values: &[Value],
+        params: &QueryParams,
+    ) -> DBResult<Vec<(usize, T)>> {
+        let recs = self.engine.with_shared_lock(|engine| {
+            engine.batch_find_by_records(field, values.iter(), params)
+        })?;
 
         Ok(recs
             .into_iter()
@@ -113,9 +150,28 @@ impl<T, F: Eq + Clone> DB<T, F> {
     /// This method can be used to run comparison-like queries, e.g. `field >= 10`
     /// could be expressed as `db.range_by(Field::Id, 10..)`.
     pub fn range_by<B: RangeBounds<Value>>(&mut self, field: &F, range: B) -> DBResult<Vec<T>> {
+        let recs = self.engine.with_shared_lock(|engine| {
+            engine.range_by_records(field, range, &DEFAULT_QUERY_PARAMS)
+        })?;
+
+        Ok(recs
+            .into_iter()
+            .map(|rec| (self.engine.config.from_record)(rec.values))
+            .collect())
+    }
+
+    /// Get a collection of records based on a range of indexed field values, with additional parameters.
+    /// This method can be used to run comparison-like queries, e.g. `field >= 10`
+    /// could be expressed as `db.range_by(Field::Id, 10..)`.
+    pub fn range_by_with_params<B: RangeBounds<Value>>(
+        &mut self,
+        field: &F,
+        range: B,
+        params: &QueryParams,
+    ) -> DBResult<Vec<T>> {
         let recs = self
             .engine
-            .with_shared_lock(|engine| engine.range_by_records(field, range))?;
+            .with_shared_lock(|engine| engine.range_by_records(field, range, params))?;
 
         Ok(recs
             .into_iter()
diff --git a/log_db/tests/integration.rs b/log_db/tests/integration.rs
index c42f91c..3d71279 100644
--- a/log_db/tests/integration.rs
+++ b/log_db/tests/integration.rs
@@ -865,3 +865,129 @@ fn test_delete_by_multiple_indexes() {
     let result = db.find_by(&Field::Id, &Value::Int(0)).unwrap();
     assert_eq!(result.len(), 0);
 }
+
+#[test]
+fn test_find_by_with_offset_and_limit() {
+    let data_dir = tmp_dir();
+
+    let mut db = DB::configure()
+        .fields(Inst::schema())
+        .primary_key(Inst::primary_key())
+        .secondary_keys(Inst::secondary_keys())
+        .from_record(Inst::from_record)
+        .into_record(Inst::into_record)
+        .data_dir(&data_dir)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    // Insert some records
+    for i in 0..10 {
+        db.upsert(Inst {
+            id: i,
+            name: Some("foo".to_string()),
+            data: vec![],
+        })
+        .unwrap();
+    }
+
+    // Find by name with offset and limit
+    let result = db
+        .find_by_with_params(
+            &Field::Name,
+            &Value::String("foo".to_string()),
+            &QueryParams {
+                offset: 2,
+                limit: 3,
+            },
+        )
+        .unwrap();
+
+    assert_eq!(result.len(), 3);
+    assert_eq!(result[0].id, 2);
+    assert_eq!(result[1].id, 3);
+    assert_eq!(result[2].id, 4);
+}
+
+#[test]
+fn test_batch_find_by_with_offset_and_limit() {
+    let data_dir = tmp_dir();
+
+    let mut db = DB::configure()
+        .fields(Inst::schema())
+        .primary_key(Inst::primary_key())
+        .secondary_keys(Inst::secondary_keys())
+        .from_record(Inst::from_record)
+        .into_record(Inst::into_record)
+        .data_dir(&data_dir)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    // Insert some records
+    for i in 0..10 {
+        db.upsert(Inst {
+            id: i,
+            name: Some("foo".to_string()),
+            data: vec![],
+        })
+        .unwrap();
+    }
+
+    // Batch find by id with offset and limit
+    let batch: Vec<Value> = (2..5).map(Value::Int).collect();
+    let result = db
+        .batch_find_by_with_params(
+            &Field::Id,
+            &batch,
+            &QueryParams {
+                offset: 1,
+                limit: 2,
+            },
+        )
+        .unwrap();
+
+    assert_eq!(result.len(), 2);
+    assert_eq!(result[0].1.id, 3);
+    assert_eq!(result[1].1.id, 4);
+}
+
+#[test]
+fn test_range_by_with_offset_and_limit() {
+    let data_dir = tmp_dir();
+
+    let mut db = DB::configure()
+        .fields(Inst::schema())
+        .primary_key(Inst::primary_key())
+        .secondary_keys(Inst::secondary_keys())
+        .from_record(Inst::from_record)
+        .into_record(Inst::into_record)
+        .data_dir(&data_dir)
+        .initialize()
+        .expect("Failed to initialize DB instance");
+
+    // Insert some records
+    for i in 0..10 {
+        db.upsert(Inst {
+            id: i,
+            name: Some("foo".to_string()),
+            data: vec![],
+        })
+        .unwrap();
+    }
+
+    // Range by id with offset and limit
+    let result = db
+        .range_by_with_params(
+            &Field::Id,
+            &Value::Int(2)..&Value::Int(8),
+            &QueryParams {
+                offset: 1,
+                limit: 3,
+            },
+        )
+        .unwrap();
+
+    assert_eq!(result.len(), 3);
+    assert_eq!(result[0].id, 3);
+    assert_eq!(result[1].id, 4);
+    assert_eq!(result[2].id, 5);
+}