From f29a9a65cb50e6d3a7d6416de576603eccd0eb30 Mon Sep 17 00:00:00 2001 From: Jan Tuomi Date: Mon, 10 Feb 2025 22:05:20 +0200 Subject: Add _with_params variants to query methods --- log_db/src/common.rs | 11 ++++ log_db/src/engine.rs | 21 +++++--- log_db/src/lib.rs | 68 +++++++++++++++++++++--- log_db/tests/integration.rs | 126 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 214 insertions(+), 12 deletions(-) (limited to 'log_db') 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 RangeBounds for OwnedBounds { 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 Engine { &mut self, field: &F, values: impl Iterator, + params: &QueryParams, ) -> DBResult> { let indexables = values .map(|value| { @@ -305,7 +306,10 @@ impl Engine { 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 Engine { /// 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, + log_keys: impl Iterator, ) -> DBResult> { let mut records = vec![]; let mut log_keys_map = BTreeMap::new(); @@ -363,7 +367,7 @@ impl Engine { 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 Engine { &mut self, field: &F, range: B, + params: &QueryParams, ) -> DBResult> { fn range_bound_to_indexable(bound: Bound<&Value>) -> DBResult> { match bound { @@ -415,9 +420,13 @@ impl Engine { 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 Engine { pub fn delete_by_field(&mut self, field: &F, value: &Value) -> DBResult> { let recs: Vec = 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 DB { // 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 DB { /// Get a collection of records based on an indexed field value. pub fn find_by(&mut self, field: &F, value: &Value) -> DBResult> { 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> { + 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 DB { /// 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> { - 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> { + 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 DB { /// 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>(&mut self, field: &F, range: B) -> DBResult> { + 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>( + &mut self, + field: &F, + range: B, + params: &QueryParams, + ) -> DBResult> { 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 = (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); +} -- cgit v1.3