Add multi_get functions family (#487)

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)
 

src/db.rs

@@ -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()
+}

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");
+    }
+}