Add support for data block hash index (#358)

src/db_options.rs

@@ -292,6 +292,41 @@ impl BlockBasedOptions {
             ffi::rocksdb_block_based_options_set_index_block_restart_interval(self.inner, interval);
         }
     }
+
+    /// Set the data block index type for point lookups:
+    ///  `DataBlockIndexType::BinarySearch` to use binary search within the data block.
+    ///  `DataBlockIndexType::BinaryAndHash` to use the data block hash index in combination with
+    ///  the normal binary search.
+    ///
+    /// The hash table utilization ratio is adjustable using [`set_data_block_hash_ratio`](#method.set_data_block_hash_ratio), which is
+    /// valid only when using `DataBlockIndexType::BinaryAndHash`.
+    ///
+    /// Default: `BinarySearch`
+    /// # Example
+    ///
+    /// ```
+    /// use rocksdb::{BlockBasedOptions, DataBlockIndexType, Options};
+    ///
+    /// let mut opts = Options::default();
+    /// let mut block_opts = BlockBasedOptions::default();
+    /// block_opts.set_data_block_index_type(DataBlockIndexType::BinaryAndHash);
+    /// block_opts.set_data_block_hash_ratio(0.85);
+    /// ```
+    pub fn set_data_block_index_type(&mut self, index_type: DataBlockIndexType) {
+        let index_t = index_type as i32;
+        unsafe { ffi::rocksdb_block_based_options_set_data_block_index_type(self.inner, index_t) }
+    }
+
+    /// Set the data block hash index utilization ratio.
+    ///
+    /// The smaller the utilization ratio, the less hash collisions happen, and so reduce the risk for a
+    /// point lookup to fall back to binary search due to the collisions. A small ratio means faster
+    /// lookup at the price of more space overhead.
+    ///
+    /// Default: 0.75
+    pub fn set_data_block_hash_ratio(&mut self, ratio: f64) {
+        unsafe { ffi::rocksdb_block_based_options_set_data_block_hash_ratio(self.inner, ratio) }
+    }
 }
 
 impl Default for BlockBasedOptions {
@@ -1724,6 +1759,19 @@ pub enum BlockBasedIndexType {
     TwoLevelIndexSearch,
 }
 
+/// Used by BlockBasedOptions::set_data_block_index_type.
+#[repr(C)]
+pub enum DataBlockIndexType {
+    /// Use binary search when performing point lookup for keys in data blocks.
+    /// This is the default.
+    BinarySearch = 0,
+
+    /// Appends a compact hash table to the end of the data block for efficient indexing. Backwards
+    /// compatible with databases created without this feature. Once turned on, existing data will
+    /// be gradually converted to the hash index format.
+    BinaryAndHash = 1,
+}
+
 /// Defines the underlying memtable implementation.
 /// See https://github.com/facebook/rocksdb/wiki/MemTable for more information.
 pub enum MemtableFactory {

src/lib.rs

@@ -78,8 +78,8 @@ pub use crate::{
     db_iterator::{DBIterator, DBRawIterator, DBWALIterator, Direction, IteratorMode},
     db_options::{
         BlockBasedIndexType, BlockBasedOptions, DBCompactionStyle, DBCompressionType,
-        DBRecoveryMode, FlushOptions, MemtableFactory, Options, PlainTableFactoryOptions,
+        DBRecoveryMode, DataBlockIndexType, FlushOptions, MemtableFactory, Options,
-        ReadOptions, WriteOptions,
+        PlainTableFactoryOptions, ReadOptions, WriteOptions,
     },
     db_pinnable_slice::DBPinnableSlice,
     merge_operator::MergeOperands,

tests/test_rocksdb_options.rs

@@ -15,7 +15,7 @@
 mod util;
 
 use crate::util::DBPath;
-use rocksdb::{BlockBasedOptions, Options, ReadOptions, DB};
+use rocksdb::{BlockBasedOptions, DataBlockIndexType, Options, ReadOptions, DB};
 use std::{fs, io::Read as _};
 
 #[test]
@@ -87,3 +87,49 @@ fn test_read_options() {
     let mut read_opts = ReadOptions::default();
     read_opts.set_verify_checksums(false);
 }
+
+#[test]
+fn test_set_data_block_index_type() {
+    let path = "_rust_rocksdb_test_set_data_block_index_type";
+    let n = DBPath::new(path);
+
+    // Default is `BinarySearch`
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+
+        let block_opts = BlockBasedOptions::default();
+        opts.set_block_based_table_factory(&block_opts);
+        let _db = DB::open(&opts, &n).expect("open a db works");
+
+        let mut rocksdb_log = fs::File::open(format!("{}/LOG", (&n).as_ref().to_str().unwrap()))
+            .expect("rocksdb creates a LOG file");
+        let mut settings = String::new();
+        rocksdb_log
+            .read_to_string(&mut settings)
+            .expect("can read the LOG file");
+        assert!(settings.contains("data_block_index_type: 0"));
+        assert!(settings.contains("data_block_hash_table_util_ratio: 0.750000"));
+    }
+
+    // Setting the index type and hash table utilization ratio works
+    {
+        let mut opts = Options::default();
+        opts.create_if_missing(false);
+
+        let mut block_opts = BlockBasedOptions::default();
+        block_opts.set_data_block_index_type(DataBlockIndexType::BinaryAndHash);
+        block_opts.set_data_block_hash_ratio(0.35);
+        opts.set_block_based_table_factory(&block_opts);
+        let _db = DB::open(&opts, &n).expect("open a db works");
+
+        let mut rocksdb_log = fs::File::open(format!("{}/LOG", (&n).as_ref().to_str().unwrap()))
+            .expect("rocksdb creates a LOG file");
+        let mut settings = String::new();
+        rocksdb_log
+            .read_to_string(&mut settings)
+            .expect("can read the LOG file");
+        assert!(settings.contains("data_block_index_type: 1"));
+        assert!(settings.contains("data_block_hash_table_util_ratio: 0.350000"));
+    }
+}