From 37470d341cad23b9b56c02be0ad677a69f50ca47 Mon Sep 17 00:00:00 2001 From: Martin Ek Date: Fri, 6 Apr 2018 14:56:36 +0200 Subject: [PATCH] Add db.full_iterator() --- src/db.rs | 19 ++++++++++++++++ tests/test_iterator.rs | 50 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/src/db.rs b/src/db.rs index a26fae9..78fd8c7 100644 --- a/src/db.rs +++ b/src/db.rs @@ -888,6 +888,15 @@ impl DB { DBIterator::new(self, &opts, mode) } + /// Opens an interator with `set_total_order_seek` enabled. + /// This must be used to iterate across prefixes when `set_memtable_factory` has been called + /// with a Hash-based implementation. + pub fn full_iterator(&self, mode: IteratorMode) -> DBIterator { + let mut opts = ReadOptions::default(); + opts.set_total_order_seek(true); + DBIterator::new(self, &opts, mode) + } + pub fn prefix_iterator<'a>(&self, prefix: &'a [u8]) -> DBIterator { let mut opts = ReadOptions::default(); opts.set_prefix_same_as_start(true); @@ -903,6 +912,16 @@ impl DB { DBIterator::new_cf(self, cf_handle, &opts, mode) } + pub fn full_iterator_cf( + &self, + cf_handle: ColumnFamily, + mode: IteratorMode, + ) -> Result { + let mut opts = ReadOptions::default(); + opts.set_total_order_seek(true); + DBIterator::new_cf(self, cf_handle, &opts, mode) + } + pub fn prefix_iterator_cf<'a>( &self, cf_handle: ColumnFamily, diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs index 3cb368f..5f8b1bb 100644 --- a/tests/test_iterator.rs +++ b/tests/test_iterator.rs @@ -14,7 +14,7 @@ // extern crate rocksdb; -use rocksdb::{DB, Direction, IteratorMode, Options}; +use rocksdb::{DB, Direction, IteratorMode, MemtableFactory, Options}; fn cba(input: &Box<[u8]>) -> Box<[u8]> { input.iter().cloned().collect::>().into_boxed_slice() @@ -194,3 +194,51 @@ pub fn test_prefix_iterator() { } } } + +#[test] +pub fn test_full_iterator() { + let path = "_rust_rocksdb_fulliteratortest"; + { + let a1: Box<[u8]> = key(b"aaa1"); + let a2: Box<[u8]> = key(b"aaa2"); + let b1: Box<[u8]> = key(b"bbb1"); + let b2: Box<[u8]> = key(b"bbb2"); + + let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3); + let factory = MemtableFactory::HashSkipList { + bucket_count: 1_000_000, + height: 4, + branching_factor: 4, + }; + + let mut opts = Options::default(); + opts.create_if_missing(true); + opts.set_prefix_extractor(prefix_extractor); + opts.set_allow_concurrent_memtable_write(false); + opts.set_memtable_factory(factory); + + let db = DB::open(&opts, path).unwrap(); + + assert!(db.put(&*a1, &*a1).is_ok()); + assert!(db.put(&*a2, &*a2).is_ok()); + assert!(db.put(&*b1, &*b1).is_ok()); + assert!(db.put(&*b2, &*b2).is_ok()); + + // A normal iterator won't work here since we're using a HashSkipList for our memtable + // implementation (which buckets keys based on their prefix): + let bad_iterator = db.iterator(IteratorMode::Start); + assert_eq!(bad_iterator.collect::>(), vec![]); + + let expected = vec![ + (cba(&a1), cba(&a1)), + (cba(&a2), cba(&a2)), + (cba(&b1), cba(&b1)), + (cba(&b2), cba(&b2)), + ]; + + let a_iterator = db.full_iterator(IteratorMode::Start); + assert_eq!(a_iterator.collect::>(), expected) + } + let opts = Options::default(); + assert!(DB::destroy(&opts, path).is_ok()); +}