diff options
| -rw-r--r-- | CHANGELOG.md | 1 | ||||
| -rw-r--r-- | src/db.rs | 110 | ||||
| -rw-r--r-- | tests/test_db.rs | 47 |
3 files changed, 158 insertions, 0 deletions
diff --git a/CHANGELOG.md b/CHANGELOG.md index d8cc210..5f36a23 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ * Add `DB::cancel_all_background_work` method (stanislav-tkach) * 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) @@ -564,6 +564,102 @@ impl DB { 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> { let cf_name = if let Ok(c) = CString::new(name.as_ref().as_bytes()) { c @@ -1428,3 +1524,17 @@ fn convert_options(opts: &[(&str, &str)]) -> Result<Vec<(CString, CString)>, Err }) .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() +} diff --git a/tests/test_db.rs b/tests/test_db.rs index 8006341..eed57dc 100644 --- a/tests/test_db.rs +++ b/tests/test_db.rs @@ -845,3 +845,50 @@ fn delete_range_test() { 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"); + } +} |