From ac4ffa3ebe5b92bff3936deac36edabc4c3f24a5 Mon Sep 17 00:00:00 2001
From: Yueh-Hsuan Chiang <93241502+yhchiang-sol@users.noreply.github.com>
Date: Tue, 28 Jun 2022 15:10:07 +0800
Subject: [PATCH] Add batched version MultiGet API (#633)

---
 src/db.rs        | 94 ++++++++++++++++++++++++++++++++++++++++++++++++
 tests/test_db.rs | 27 ++++++++++++++
 2 files changed, 121 insertions(+)

diff --git a/src/db.rs b/src/db.rs
index 3e82a2d..11f1596 100644
--- a/src/db.rs
+++ b/src/db.rs
@@ -177,6 +177,17 @@ pub trait DBAccess {
         K: AsRef<[u8]>,
         I: IntoIterator<Item = (&'b W, K)>,
         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> {
@@ -242,6 +253,20 @@ impl<T: ThreadMode> DBAccess for DBWithThreadMode<T> {
     {
         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
@@ -1022,6 +1047,75 @@ impl<T: ThreadMode> DBWithThreadMode<T> {
         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
     /// `true`. This function uses default `ReadOptions`.
     pub fn key_may_exist<K: AsRef<[u8]>>(&self, key: K) -> bool {
diff --git a/tests/test_db.rs b/tests/test_db.rs
index 2c08993..6ceed76 100644
--- a/tests/test_db.rs
+++ b/tests/test_db.rs
@@ -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]
 fn key_may_exist() {
     let path = DBPath::new("_rust_key_may_exist");