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, W: AsColumnFamilyRef + 'b; + + fn batched_multi_get_cf_opt( + &self, + cf: &impl AsColumnFamilyRef, + keys: I, + sorted_input: bool, + readopts: &ReadOptions, + ) -> Vec, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator; } impl DBAccess for DBWithThreadMode { @@ -242,6 +253,20 @@ impl DBAccess for DBWithThreadMode { { self.multi_get_cf_opt(keys_cf, readopts) } + + fn batched_multi_get_cf_opt( + &self, + cf: &impl AsColumnFamilyRef, + keys: I, + sorted_input: bool, + readopts: &ReadOptions, + ) -> Vec, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + { + 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 DBWithThreadMode { 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( + &self, + cf: &impl AsColumnFamilyRef, + keys: I, + sorted_input: bool, + ) -> Vec, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + { + 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( + &self, + cf: &impl AsColumnFamilyRef, + keys: I, + sorted_input: bool, + readopts: &ReadOptions, + ) -> Vec, Error>> + where + K: AsRef<[u8]>, + I: IntoIterator, + { + let (keys, keys_sizes): (Vec>, 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>(&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::>(); + 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");