Add multi_get functions family (#487)

master
Stanislav Tkach 4 years ago committed by GitHub
parent 154c23bc99
commit ed69084041
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 1
      CHANGELOG.md
  2. 110
      src/db.rs
  3. 47
      tests/test_db.rs

@ -4,6 +4,7 @@
* Add `DB::cancel_all_background_work` method (stanislav-tkach) * Add `DB::cancel_all_background_work` method (stanislav-tkach)
* Bump `librocksdb-sys` up to 6.13.3 (aleksuss) * Bump `librocksdb-sys` up to 6.13.3 (aleksuss)
* Add `multi_get`, `multi_get_opt`, `multi_get_cf` and `multi_get_cf_opt` `DB` methods (stanislav-tkach)
## 0.15.0 (2020-08-25) ## 0.15.0 (2020-08-25)

@ -564,6 +564,102 @@ impl DB {
self.get_pinned_cf_opt(cf, key, &ReadOptions::default()) self.get_pinned_cf_opt(cf, key, &ReadOptions::default())
} }
/// Return the values associated with the given keys.
pub fn multi_get<K, I>(&self, keys: I) -> Result<Vec<Vec<u8>>, Error>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
self.multi_get_opt(keys, &ReadOptions::default())
}
/// Return the values associated with the given keys using read options.
pub fn multi_get_opt<K, I>(
&self,
keys: I,
readopts: &ReadOptions,
) -> Result<Vec<Vec<u8>>, Error>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
let (keys, keys_sizes): (Vec<Box<[u8]>>, Vec<_>) = keys
.into_iter()
.map(|k| (Box::from(k.as_ref()), k.as_ref().len()))
.unzip();
let ptr_keys: Vec<_> = keys.iter().map(|k| k.as_ptr() as *const c_char).collect();
let mut values = vec![ptr::null_mut(); keys.len()];
let mut values_sizes = vec![0_usize; keys.len()];
unsafe {
ffi_try!(ffi::rocksdb_multi_get(
self.inner,
readopts.inner,
ptr_keys.len(),
ptr_keys.as_ptr(),
keys_sizes.as_ptr(),
values.as_mut_ptr(),
values_sizes.as_mut_ptr(),
));
}
Ok(convert_values(values, values_sizes))
}
/// Return the values associated with the given keys and column families.
pub fn multi_get_cf<'c, K, I>(&self, keys: I) -> Result<Vec<Vec<u8>>, Error>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'c ColumnFamily, K)>,
{
self.multi_get_cf_opt(keys, &ReadOptions::default())
}
/// Return the values associated with the given keys and column families using read options.
pub fn multi_get_cf_opt<'c, K, I>(
&self,
keys: I,
readopts: &ReadOptions,
) -> Result<Vec<Vec<u8>>, Error>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = (&'c ColumnFamily, K)>,
{
let mut boxed_keys: Vec<Box<[u8]>> = Vec::new();
let mut keys_sizes = Vec::new();
let mut column_families = Vec::new();
for (cf, key) in keys {
boxed_keys.push(Box::from(key.as_ref()));
keys_sizes.push(key.as_ref().len());
column_families.push(cf);
}
let ptr_keys: Vec<_> = boxed_keys
.iter()
.map(|k| k.as_ptr() as *const c_char)
.collect();
let ptr_cfs: Vec<_> = column_families
.iter()
.map(|c| c.inner as *const _)
.collect();
let mut values = vec![ptr::null_mut(); boxed_keys.len()];
let mut values_sizes = vec![0_usize; boxed_keys.len()];
unsafe {
ffi_try!(ffi::rocksdb_multi_get_cf(
self.inner,
readopts.inner,
ptr_cfs.as_ptr(),
ptr_keys.len(),
ptr_keys.as_ptr(),
keys_sizes.as_ptr(),
values.as_mut_ptr(),
values_sizes.as_mut_ptr(),
));
}
Ok(convert_values(values, values_sizes))
}
pub fn create_cf<N: AsRef<str>>(&mut self, name: N, opts: &Options) -> Result<(), Error> { pub fn create_cf<N: AsRef<str>>(&mut self, name: N, opts: &Options) -> Result<(), Error> {
let cf_name = if let Ok(c) = CString::new(name.as_ref().as_bytes()) { let cf_name = if let Ok(c) = CString::new(name.as_ref().as_bytes()) {
c c
@ -1428,3 +1524,17 @@ fn convert_options(opts: &[(&str, &str)]) -> Result<Vec<(CString, CString)>, Err
}) })
.collect() .collect()
} }
fn convert_values(values: Vec<*mut c_char>, values_sizes: Vec<usize>) -> Vec<Vec<u8>> {
values
.into_iter()
.zip(values_sizes.into_iter())
.map(|(v, s)| {
let value = unsafe { slice::from_raw_parts(v as *const u8, s) }.into();
unsafe {
ffi::rocksdb_free(v as *mut c_void);
}
value
})
.collect()
}

@ -845,3 +845,50 @@ fn delete_range_test() {
assert!(db.get_cf(cf1, b"k3").unwrap().is_none()); assert!(db.get_cf(cf1, b"k3").unwrap().is_none());
} }
} }
#[test]
fn multi_get() {
let path = DBPath::new("_rust_rocksdb_multi_get");
{
let db = DB::open_default(&path).unwrap();
db.put(b"k1", b"v1").unwrap();
db.put(b"k2", b"v2").unwrap();
let values = db
.multi_get(&[b"k0", b"k1", b"k2"])
.expect("multi_get failed");
assert_eq!(3, values.len());
assert!(values[0].is_empty());
assert_eq!(values[1], b"v1");
assert_eq!(values[2], b"v2");
}
}
#[test]
fn multi_get_cf() {
let path = DBPath::new("_rust_rocksdb_multi_get_cf");
{
let mut opts = Options::default();
opts.create_if_missing(true);
opts.create_missing_column_families(true);
let db = DB::open_cf(&opts, &path, &["cf0", "cf1", "cf2"]).unwrap();
let cf0 = db.cf_handle("cf0").unwrap();
let cf1 = db.cf_handle("cf1").unwrap();
db.put_cf(cf1, b"k1", b"v1").unwrap();
let cf2 = db.cf_handle("cf2").unwrap();
db.put_cf(cf2, b"k2", b"v2").unwrap();
let values = db
.multi_get_cf(vec![(cf0, b"k0"), (cf1, b"k1"), (cf2, b"k2")])
.expect("multi_get failed");
assert_eq!(3, values.len());
assert!(values[0].is_empty());
assert_eq!(values[1], b"v1");
assert_eq!(values[2], b"v2");
}
}

Loading…
Cancel
Save