Add batched version MultiGet API (#633)

master
Yueh-Hsuan Chiang 2 years ago committed by GitHub
parent da4b7143b9
commit ac4ffa3ebe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 94
      src/db.rs
  2. 27
      tests/test_db.rs

@ -177,6 +177,17 @@ pub trait DBAccess {
K: AsRef<[u8]>, K: AsRef<[u8]>,
I: IntoIterator<Item = (&'b W, K)>, I: IntoIterator<Item = (&'b W, K)>,
W: AsColumnFamilyRef + 'b; W: AsColumnFamilyRef + 'b;
fn batched_multi_get_cf_opt<K, I>(
&self,
cf: &impl AsColumnFamilyRef,
keys: I,
sorted_input: bool,
readopts: &ReadOptions,
) -> Vec<Result<Option<DBPinnableSlice>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>;
} }
impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> { impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
@ -242,6 +253,20 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
{ {
self.multi_get_cf_opt(keys_cf, readopts) self.multi_get_cf_opt(keys_cf, readopts)
} }
fn batched_multi_get_cf_opt<K, I>(
&self,
cf: &impl AsColumnFamilyRef,
keys: I,
sorted_input: bool,
readopts: &ReadOptions,
) -> Vec<Result<Option<DBPinnableSlice>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
self.batched_multi_get_cf_opt(cf, keys, sorted_input, 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
@ -1022,6 +1047,75 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
convert_values(values, values_sizes, errors) convert_values(values, values_sizes, errors)
} }
/// Return the values associated with the given keys and the specified column family
/// where internally the read requests are processed in batch if block-based table
/// SST format is used. It is a more optimized version of multi_get_cf.
pub fn batched_multi_get_cf<K, I>(
&self,
cf: &impl AsColumnFamilyRef,
keys: I,
sorted_input: bool,
) -> Vec<Result<Option<DBPinnableSlice>, Error>>
where
K: AsRef<[u8]>,
I: IntoIterator<Item = K>,
{
self.batched_multi_get_cf_opt(cf, keys, sorted_input, &ReadOptions::default())
}
/// Return the values associated with the given keys and the specified column family
/// where internally the read requests are processed in batch if block-based table
/// SST format is used. It is a more optimized version of multi_get_cf_opt.
pub fn batched_multi_get_cf_opt<K, I>(
&self,
cf: &impl AsColumnFamilyRef,
keys: I,
sorted_input: bool,
readopts: &ReadOptions,
) -> Vec<Result<Option<DBPinnableSlice>, 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 pinned_values = vec![ptr::null_mut(); ptr_keys.len()];
let mut errors = vec![ptr::null_mut(); ptr_keys.len()];
unsafe {
ffi::rocksdb_batched_multi_get_cf(
self.inner,
readopts.inner,
cf.inner(),
ptr_keys.len(),
ptr_keys.as_ptr(),
keys_sizes.as_ptr(),
pinned_values.as_mut_ptr(),
errors.as_mut_ptr(),
sorted_input,
);
pinned_values
.into_iter()
.zip(errors.into_iter())
.map(|(v, e)| {
if e.is_null() {
if v.is_null() {
Ok(None)
} else {
Ok(Some(DBPinnableSlice::from_c(v)))
}
} else {
Err(Error::new(crate::ffi_util::error_message(e)))
}
})
.collect()
}
}
/// Returns `false` if the given key definitely doesn't exist in the database, otherwise returns /// Returns `false` if the given key definitely doesn't exist in the database, otherwise returns
/// `true`. This function uses default `ReadOptions`. /// `true`. This function uses default `ReadOptions`.
pub fn key_may_exist<K: AsRef<[u8]>>(&self, key: K) -> bool { pub fn key_may_exist<K: AsRef<[u8]>>(&self, key: K) -> bool {

@ -1244,6 +1244,33 @@ fn multi_get_cf() {
} }
} }
#[test]
fn batched_multi_get_cf() {
let path = DBPath::new("_rust_rocksdb_batched_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"]).unwrap();
let cf = db.cf_handle("cf0").unwrap();
db.put_cf(&cf, b"k1", b"v1").unwrap();
db.put_cf(&cf, b"k2", b"v2").unwrap();
let values = db
.batched_multi_get_cf(&cf, vec![b"k0", b"k1", b"k2"], true) // sorted_input
.into_iter()
.map(Result::unwrap)
.collect::<Vec<_>>();
assert_eq!(3, values.len());
assert!(values[0].is_none());
assert!(values[1].is_some());
assert_eq!(&(values[1].as_ref().unwrap())[0..2], b"v1");
assert_eq!(&(values[2].as_ref().unwrap())[0..2], b"v2");
}
}
#[test] #[test]
fn key_may_exist() { fn key_may_exist() {
let path = DBPath::new("_rust_key_may_exist"); let path = DBPath::new("_rust_key_may_exist");

Loading…
Cancel
Save