Support `multi_get_*` methods (#572)

master
Oleg Nosov 3 years ago committed by GitHub
parent 522b496738
commit 8bf9faddb3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 44
      src/db.rs
  2. 2
      src/lib.rs
  3. 49
      src/snapshot.rs
  4. 49
      tests/test_db.rs

@ -145,6 +145,25 @@ pub trait DBAccess {
key: K, key: K,
readopts: &ReadOptions, readopts: &ReadOptions,
) -> Result<Option<Vec<u8>>, Error>; ) -> Result<Option<Vec<u8>>, Error>;
fn multi_get_opt<K, I>(
&self,
keys: I,
readopts: &ReadOptions,
) -> Vec<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>;
fn multi_get_cf_opt<'b, K, I, W>(
&self,
keys_cf: I,
readopts: &ReadOptions,
) -> Vec<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'b W, K)>,
W: AsColumnFamilyRef + 'b;
} }
impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> { impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
@ -168,6 +187,31 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
) -> Result<Option<Vec<u8>>, Error> { ) -> Result<Option<Vec<u8>>, Error> {
self.get_cf_opt(cf, key, readopts) self.get_cf_opt(cf, key, readopts)
} }
fn multi_get_opt<K, I>(
&self,
keys: I,
readopts: &ReadOptions,
) -> Vec<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
self.multi_get_opt(keys, readopts)
}
fn multi_get_cf_opt<'b, K, I, W>(
&self,
keys_cf: I,
readopts: &ReadOptions,
) -> Vec<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'b W, K)>,
W: AsColumnFamilyRef + 'b,
{
self.multi_get_cf_opt(keys_cf, readopts)
}
} }
/// A type alias to DB instance type with the single-threaded column family /// A type alias to DB instance type with the single-threaded column family

@ -98,7 +98,7 @@ pub use crate::{
ColumnFamilyRef, DEFAULT_COLUMN_FAMILY_NAME, ColumnFamilyRef, DEFAULT_COLUMN_FAMILY_NAME,
}, },
compaction_filter::Decision as CompactionDecision, compaction_filter::Decision as CompactionDecision,
db::{DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB}, db::{DBAccess, DBWithThreadMode, LiveFile, MultiThreaded, SingleThreaded, ThreadMode, DB},
db_iterator::{ db_iterator::{
DBIterator, DBIteratorWithThreadMode, DBRawIterator, DBRawIteratorWithThreadMode, DBIterator, DBIteratorWithThreadMode, DBRawIterator, DBRawIteratorWithThreadMode,
DBWALIterator, Direction, IteratorMode, DBWALIterator, Direction, IteratorMode,

@ -160,6 +160,55 @@ impl<'a, D: DBAccess> SnapshotWithThreadMode<'a, D> {
readopts.set_snapshot(self); readopts.set_snapshot(self);
self.db.get_cf_opt(cf, key.as_ref(), &readopts) 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<K: AsRef<[u8]>, I>(&self, keys: I) -> Vec<Result<Option<Vec<u8>>, Error>>
where
I: IntoIterator<Item = K>,
{
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<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'b W, K)>,
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<K, I>(
&self,
keys: I,
mut readopts: ReadOptions,
) -> Vec<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
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<Result<Option<Vec<u8>>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'b W, K)>,
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> { impl<'a, D: DBAccess> Drop for SnapshotWithThreadMode<'a, D> {

@ -20,7 +20,7 @@ use pretty_assertions::assert_eq;
use rocksdb::{ use rocksdb::{
perf::get_memory_usage_stats, BlockBasedOptions, BottommostLevelCompaction, Cache, 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, FifoCompactOptions, IteratorMode, MultiThreaded, Options, PerfContext, PerfMetric, ReadOptions,
SingleThreaded, SliceTransform, Snapshot, UniversalCompactOptions, SingleThreaded, SliceTransform, Snapshot, UniversalCompactOptions,
UniversalCompactionStopStyle, WriteBatch, DB, UniversalCompactionStopStyle, WriteBatch, DB,
@ -926,20 +926,59 @@ fn multi_get() {
{ {
let db = DB::open_default(&path).unwrap(); let db = DB::open_default(&path).unwrap();
let initial_snap = db.snapshot();
db.put(b"k1", b"v1").unwrap(); db.put(b"k1", b"v1").unwrap();
let k1_snap = db.snapshot();
db.put(b"k2", b"v2").unwrap(); db.put(b"k2", b"v2").unwrap();
let _ = db.multi_get(&[b"k0"; 40]); 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 let values = db
.multi_get(&[b"k0", b"k1", b"k2"]) .multi_get(&[b"k0", b"k1", b"k2"])
.into_iter() .into_iter()
.map(Result::unwrap) .map(Result::unwrap)
.collect::<Vec<_>>(); .collect::<Vec<_>>();
assert_eq!(3, values.len());
assert_eq!(values[0], None); assert_values(values);
assert_eq!(values[1], Some(b"v1".to_vec()));
assert_eq!(values[2], Some(b"v2".to_vec())); let values = DBAccess::multi_get_opt(&db, &[b"k0", b"k1", b"k2"], &Default::default())
.into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
assert_values(values);
let values = db
.snapshot()
.multi_get(&[b"k0", b"k1", b"k2"])
.into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
assert_values(values);
let none_values = initial_snap
.multi_get(&[b"k0", b"k1", b"k2"])
.into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
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::<Vec<_>>();
assert_eq!(k1_only, vec![None, Some(b"v1".to_vec()), None]);
} }
} }

Loading…
Cancel
Save