From 8bf9faddb3b51bd1f1267efbc2037b323b4ba0a4 Mon Sep 17 00:00:00 2001 From: Oleg Nosov Date: Fri, 14 Jan 2022 13:55:55 +0300 Subject: [PATCH] Support `multi_get_*` methods (#572) --- src/db.rs | 44 +++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 2 +- src/snapshot.rs | 49 ++++++++++++++++++++++++++++++++++++++++++++++++ tests/test_db.rs | 49 +++++++++++++++++++++++++++++++++++++++++++----- 4 files changed, 138 insertions(+), 6 deletions(-) diff --git a/src/db.rs b/src/db.rs index 0738b7d..9b673ec 100644 --- a/src/db.rs +++ b/src/db.rs @@ -145,6 +145,25 @@ pub trait DBAccess { key: K, readopts: &ReadOptions, ) -> Result>, Error>; + + fn multi_get_opt( + &self, + keys: I, + readopts: &ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator; + + fn multi_get_cf_opt<'b, K, I, W>( + &self, + keys_cf: I, + readopts: &ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + W: AsColumnFamilyRef + 'b; } impl DBAccess for DBWithThreadMode { @@ -168,6 +187,31 @@ impl DBAccess for DBWithThreadMode { ) -> Result>, Error> { self.get_cf_opt(cf, key, readopts) } + + fn multi_get_opt( + &self, + keys: I, + readopts: &ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + { + self.multi_get_opt(keys, readopts) + } + + fn multi_get_cf_opt<'b, K, I, W>( + &self, + keys_cf: I, + readopts: &ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + W: AsColumnFamilyRef + 'b, + { + self.multi_get_cf_opt(keys_cf, readopts) + } } /// A type alias to DB instance type with the single-threaded column family diff --git a/src/lib.rs b/src/lib.rs index 1b2fba6..c84f834 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -98,7 +98,7 @@ pub use crate::{ ColumnFamilyRef, DEFAULT_COLUMN_FAMILY_NAME, }, compaction_filter::Decision as CompactionDecision, - db::{DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB}, + db::{DBAccess, DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB}, db_iterator::{ DBIterator, DBIteratorWithThreadMode, DBRawIterator, DBRawIteratorWithThreadMode, DBWALIterator, Direction, IteratorMode, diff --git a/src/snapshot.rs b/src/snapshot.rs index a4356ca..0a0ea06 100644 --- a/src/snapshot.rs +++ b/src/snapshot.rs @@ -160,6 +160,55 @@ impl<'a, D: DBAccess> SnapshotWithThreadMode<'a, D> { readopts.set_snapshot(self); self.db.get_cf_opt(cf, key.as_ref(), &readopts) } + + /// Returns the bytes associated with the given key values and default read options. + pub fn multi_get, I>(&self, keys: I) -> Vec>, Error>> + where + I: IntoIterator, + { + let readopts = ReadOptions::default(); + self.multi_get_opt(keys, readopts) + } + + /// Returns the bytes associated with the given key values and default read options. + pub fn multi_get_cf<'b, K, I, W>(&self, keys_cf: I) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + W: AsColumnFamilyRef + 'b, + { + let readopts = ReadOptions::default(); + self.multi_get_cf_opt(keys_cf, readopts) + } + + /// Returns the bytes associated with the given key values and given read options. + pub fn multi_get_opt( + &self, + keys: I, + mut readopts: ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + { + readopts.set_snapshot(self); + self.db.multi_get_opt(keys, &readopts) + } + + /// Returns the bytes associated with the given key values, given column family and read options. + pub fn multi_get_cf_opt<'b, K, I, W>( + &self, + keys_cf: I, + mut readopts: ReadOptions, + ) -> Vec>, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + W: AsColumnFamilyRef + 'b, + { + readopts.set_snapshot(self); + self.db.multi_get_cf_opt(keys_cf, &readopts) + } } impl<'a, D: DBAccess> Drop for SnapshotWithThreadMode<'a, D> { diff --git a/tests/test_db.rs b/tests/test_db.rs index 394d05e..a362d46 100644 --- a/tests/test_db.rs +++ b/tests/test_db.rs @@ -20,7 +20,7 @@ use pretty_assertions::assert_eq; use rocksdb::{ perf::get_memory_usage_stats, BlockBasedOptions, BottommostLevelCompaction, Cache, - CompactOptions, CuckooTableOptions, DBCompactionStyle, DBWithThreadMode, Env, Error, + CompactOptions, CuckooTableOptions, DBAccess, DBCompactionStyle, DBWithThreadMode, Env, Error, FifoCompactOptions, IteratorMode, MultiThreaded, Options, PerfContext, PerfMetric, ReadOptions, SingleThreaded, SliceTransform, Snapshot, UniversalCompactOptions, UniversalCompactionStopStyle, WriteBatch, DB, @@ -926,20 +926,59 @@ fn multi_get() { { let db = DB::open_default(&path).unwrap(); + let initial_snap = db.snapshot(); db.put(b"k1", b"v1").unwrap(); + let k1_snap = db.snapshot(); db.put(b"k2", b"v2").unwrap(); let _ = db.multi_get(&[b"k0"; 40]); + let assert_values = |values: Vec<_>| { + assert_eq!(3, values.len()); + assert_eq!(values[0], None); + assert_eq!(values[1], Some(b"v1".to_vec())); + assert_eq!(values[2], Some(b"v2".to_vec())); + }; + let values = db .multi_get(&[b"k0", b"k1", b"k2"]) .into_iter() .map(Result::unwrap) .collect::>(); - assert_eq!(3, values.len()); - assert_eq!(values[0], None); - assert_eq!(values[1], Some(b"v1".to_vec())); - assert_eq!(values[2], Some(b"v2".to_vec())); + + assert_values(values); + + let values = DBAccess::multi_get_opt(&db, &[b"k0", b"k1", b"k2"], &Default::default()) + .into_iter() + .map(Result::unwrap) + .collect::>(); + + assert_values(values); + + let values = db + .snapshot() + .multi_get(&[b"k0", b"k1", b"k2"]) + .into_iter() + .map(Result::unwrap) + .collect::>(); + + assert_values(values); + + let none_values = initial_snap + .multi_get(&[b"k0", b"k1", b"k2"]) + .into_iter() + .map(Result::unwrap) + .collect::>(); + + assert_eq!(none_values, vec![None; 3]); + + let k1_only = k1_snap + .multi_get(&[b"k0", b"k1", b"k2"]) + .into_iter() + .map(Result::unwrap) + .collect::>(); + + assert_eq!(k1_only, vec![None, Some(b"v1".to_vec()), None]); } }