total order seek (#69)

librocksdb_sys/src/lib.rs

@@ -231,6 +231,7 @@ extern "C" {
                                                         kLen: size_t);
     pub fn crocksdb_readoptions_set_read_tier(readopts: *mut DBReadOptions, tier: c_int);
     pub fn crocksdb_readoptions_set_tailing(readopts: *mut DBReadOptions, v: bool);
+    pub fn crocksdb_readoptions_set_total_order_seek(readopts: *mut DBReadOptions, v: bool);
 
     pub fn crocksdb_get(db: *const DBInstance,
                         readopts: *const DBReadOptions,

src/rocksdb.rs

@@ -544,6 +544,10 @@ impl DB {
         DBIterator::new_cf(self, cf_handle, opts)
     }
 
+    pub fn iter_cf_opt(&self, cf_handle: &CFHandle, opts: ReadOptions) -> DBIterator {
+        DBIterator::new_cf(self, cf_handle, opts)
+    }
+
     pub fn snapshot(&self) -> Snapshot {
         Snapshot::new(self)
     }

src/rocksdb_options.rs

@@ -166,6 +166,12 @@ impl ReadOptions {
         }
     }
 
+    pub fn set_total_order_seek(&mut self, v: bool) {
+        unsafe {
+            crocksdb_ffi::crocksdb_readoptions_set_total_order_seek(self.inner, v);
+        }
+    }
+
     pub unsafe fn get_inner(&self) -> *const DBReadOptions {
         self.inner
     }

src/slice_transform.rs

@@ -33,7 +33,7 @@ pub trait SliceTransform {
     fn in_domain(&mut self, key: &[u8]) -> bool;
 
     // This is currently not used and remains here for backward compatibility.
-    fn in_range(&mut self, key: &[u8]) -> bool {
+    fn in_range(&mut self, _: &[u8]) -> bool {
         true
     }
 }

test/test_iterator.rs

 use rocksdb::*;
 use tempdir::TempDir;
 
+struct FixedSuffixTransform {
+    pub suffix_len: usize,
+}
+
+impl SliceTransform for FixedSuffixTransform {
+    fn transform<'a>(&mut self, key: &'a [u8]) -> &'a [u8] {
+        &key[..self.suffix_len]
+    }
+
+    fn in_domain(&mut self, key: &[u8]) -> bool {
+        key.len() >= self.suffix_len
+    }
+}
+
 fn prev_collect<'a>(iter: &mut DBIterator<'a>) -> Vec<Kv> {
     let mut buf = vec![];
     while iter.valid() {
@@ -186,3 +200,65 @@ fn read_with_upper_bound() {
         assert_eq!(count, 2);
     }
 }
+
+#[test]
+fn test_total_order_seek() {
+    let path = TempDir::new("_rust_rocksdb_total_order_seek").expect("");
+    let mut bbto = BlockBasedOptions::new();
+    bbto.set_bloom_filter(10, false);
+    bbto.set_whole_key_filtering(false);
+    let mut opts = Options::new();
+    opts.create_if_missing(true);
+    opts.set_block_based_table_factory(&bbto);
+    opts.set_prefix_extractor("FixedSuffixTransform",
+                              Box::new(FixedSuffixTransform { suffix_len: 2 }))
+        .unwrap();
+    // also create prefix bloom for memtable
+    opts.set_memtable_prefix_bloom_size_ratio(0.1 as f64);
+
+    let keys = vec![b"k1-0", b"k1-1", b"k1-2", b"k2-0", b"k2-1", b"k2-2", b"k3-0", b"k3-1",
+                    b"k3-2"];
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    let wopts = WriteOptions::new();
+
+    // sst1
+    db.put_opt(b"k1-0", b"a", &wopts).unwrap();
+    db.put_opt(b"k1-1", b"b", &wopts).unwrap();
+    db.put_opt(b"k1-2", b"c", &wopts).unwrap();
+    db.flush(true /* sync */).unwrap(); // flush memtable to sst file.
+
+    // sst2
+    db.put_opt(b"k2-0", b"a", &wopts).unwrap();
+    db.put_opt(b"k2-1", b"b", &wopts).unwrap();
+    db.put_opt(b"k2-2", b"c", &wopts).unwrap();
+    db.flush(true /* sync */).unwrap(); // flush memtable to sst file.
+
+    // memtable
+    db.put_opt(b"k3-0", b"a", &wopts).unwrap();
+    db.put_opt(b"k3-1", b"b", &wopts).unwrap();
+    db.put_opt(b"k3-2", b"c", &wopts).unwrap();
+
+    let mut iter = db.iter();
+    iter.seek(SeekKey::Key(b"k1-0"));
+    let mut key_count = 0;
+    while iter.valid() {
+        // only iterator sst files and memtable that contain keys has the same prefix with b"k1-0".
+        assert_eq!(keys[key_count], iter.key());
+        key_count = key_count + 1;
+        iter.next();
+    }
+    assert!(key_count == 3);
+
+    let mut ropts = ReadOptions::new();
+    ropts.set_total_order_seek(true);
+    let mut iter = db.iter_opt(ropts);
+    iter.seek(SeekKey::Key(b"k1-0"));
+    let mut key_count = 0;
+    while iter.valid() {
+        // iterator all sst files and memtables
+        assert_eq!(keys[key_count], iter.key());
+        key_count = key_count + 1;
+        iter.next();
+    }
+    assert!(key_count == 9);
+}