WriteBatchWithIndex's iterator to support SeekToFirst(), SeekToLast() and Prev()

Summary: Support SeekToFirst(), SeekToLast() and Prev() in WBWIIterator, returned by WriteBatchWithIndex::NewIterator().

Test Plan: Write unit test cases to cover the case.

Reviewers: ljin, igor

Reviewed By: igor

Subscribers: rven, yhchiang, leveldb

Differential Revision: https://reviews.facebook.net/D24765

include/rocksdb/utilities/write_batch_with_index.h

@@ -38,10 +38,16 @@ class WBWIIterator {
 
   virtual bool Valid() const = 0;
 
+  virtual void SeekToFirst() = 0;
+
+  virtual void SeekToLast() = 0;
+
   virtual void Seek(const Slice& key) = 0;
 
   virtual void Next() = 0;
 
+  virtual void Prev() = 0;
+
   virtual const WriteEntry& Entry() const = 0;
 
   virtual Status status() const = 0;
@@ -71,29 +77,29 @@ class WriteBatchWithIndex {
 
   WriteBatch* GetWriteBatch();
 
-  virtual void Put(ColumnFamilyHandle* column_family, const Slice& key,
+  void Put(ColumnFamilyHandle* column_family, const Slice& key,
            const Slice& value);
 
-  virtual void Put(const Slice& key, const Slice& value);
+  void Put(const Slice& key, const Slice& value);
 
-  virtual void Merge(ColumnFamilyHandle* column_family, const Slice& key,
+  void Merge(ColumnFamilyHandle* column_family, const Slice& key,
              const Slice& value);
 
-  virtual void Merge(const Slice& key, const Slice& value);
+  void Merge(const Slice& key, const Slice& value);
 
-  virtual void PutLogData(const Slice& blob);
+  void PutLogData(const Slice& blob);
 
-  virtual void Delete(ColumnFamilyHandle* column_family, const Slice& key);
-  virtual void Delete(const Slice& key);
+  void Delete(ColumnFamilyHandle* column_family, const Slice& key);
+  void Delete(const Slice& key);
 
   // Create an iterator of a column family. User can call iterator.Seek() to
   // search to the next entry of or after a key. Keys will be iterated in the
   // order given by index_comparator. For multiple updates on the same key,
   // each update will be returned as a separate entry, in the order of update
   // time.
-  virtual WBWIIterator* NewIterator(ColumnFamilyHandle* column_family);
+  WBWIIterator* NewIterator(ColumnFamilyHandle* column_family);
   // Create an iterator of the default column family.
-  virtual WBWIIterator* NewIterator();
+  WBWIIterator* NewIterator();
 
  private:
   struct Rep;

utilities/write_batch_with_index/write_batch_with_index.cc

@@ -76,6 +76,25 @@ class WBWIIteratorImpl : public WBWIIterator {
 
   virtual bool Valid() const override { return valid_; }
 
+  virtual void SeekToFirst() {
+    valid_ = true;
+    WriteBatchIndexEntry search_entry(nullptr, column_family_id_);
+    skip_list_iter_.Seek(&search_entry);
+    ReadEntry();
+  }
+
+  virtual void SeekToLast() {
+    valid_ = true;
+    WriteBatchIndexEntry search_entry(nullptr, column_family_id_ + 1);
+    skip_list_iter_.Seek(&search_entry);
+    if (!skip_list_iter_.Valid()) {
+      skip_list_iter_.SeekToLast();
+    } else {
+      skip_list_iter_.Prev();
+    }
+    ReadEntry();
+  }
+
   virtual void Seek(const Slice& key) override {
     valid_ = true;
     WriteBatchIndexEntry search_entry(&key, column_family_id_);
@@ -88,6 +107,11 @@ class WBWIIteratorImpl : public WBWIIterator {
     ReadEntry();
   }
 
+  virtual void Prev() override {
+    skip_list_iter_.Prev();
+    ReadEntry();
+  }
+
   virtual const WriteEntry& Entry() const override { return current_; }
 
   virtual Status status() const override { return status_; }

utilities/write_batch_with_index/write_batch_with_index_test.cc

@@ -120,7 +120,12 @@ TEST(WriteBatchWithIndexTest, TestValueAsSecondaryIndex) {
   // Iterator all keys
   {
     std::unique_ptr<WBWIIterator> iter(batch.NewIterator(&data));
+    for (int seek_to_first : {0, 1}) {
+      if (seek_to_first) {
+        iter->SeekToFirst();
+      } else {
         iter->Seek("");
+      }
       for (auto pair : data_map) {
         for (auto v : pair.second) {
           ASSERT_OK(iter->status());
@@ -136,11 +141,32 @@ TEST(WriteBatchWithIndexTest, TestValueAsSecondaryIndex) {
       }
       ASSERT_TRUE(!iter->Valid());
     }
+    iter->SeekToLast();
+    for (auto pair = data_map.rbegin(); pair != data_map.rend(); ++pair) {
+      for (auto v = pair->second.rbegin(); v != pair->second.rend(); v++) {
+        ASSERT_OK(iter->status());
+        ASSERT_TRUE(iter->Valid());
+        auto& write_entry = iter->Entry();
+        ASSERT_EQ(pair->first, write_entry.key.ToString());
+        ASSERT_EQ((*v)->type, write_entry.type);
+        if (write_entry.type != kDeleteRecord) {
+          ASSERT_EQ((*v)->value, write_entry.value.ToString());
+        }
+        iter->Prev();
+      }
+    }
+    ASSERT_TRUE(!iter->Valid());
+  }
 
   // Iterator all indexes
   {
     std::unique_ptr<WBWIIterator> iter(batch.NewIterator(&index));
+    for (int seek_to_first : {0, 1}) {
+      if (seek_to_first) {
+        iter->SeekToFirst();
+      } else {
         iter->Seek("");
+      }
       for (auto pair : index_map) {
         for (auto v : pair.second) {
           ASSERT_OK(iter->status());
@@ -157,6 +183,23 @@ TEST(WriteBatchWithIndexTest, TestValueAsSecondaryIndex) {
       ASSERT_TRUE(!iter->Valid());
     }
 
+    iter->SeekToLast();
+    for (auto pair = index_map.rbegin(); pair != index_map.rend(); ++pair) {
+      for (auto v = pair->second.rbegin(); v != pair->second.rend(); v++) {
+        ASSERT_OK(iter->status());
+        ASSERT_TRUE(iter->Valid());
+        auto& write_entry = iter->Entry();
+        ASSERT_EQ(pair->first, write_entry.key.ToString());
+        if ((*v)->type != kDeleteRecord) {
+          ASSERT_EQ((*v)->key, write_entry.value.ToString());
+          ASSERT_EQ((*v)->value, write_entry.key.ToString());
+        }
+        iter->Prev();
+      }
+    }
+    ASSERT_TRUE(!iter->Valid());
+  }
+
   // Seek to every key
   {
     std::unique_ptr<WBWIIterator> iter(batch.NewIterator(&data));
@@ -357,7 +400,21 @@ TEST(WriteBatchWithIndexTest, TestOverwriteKey) {
 
   {
     std::unique_ptr<WBWIIterator> iter(batch.NewIterator(&cf2));
-    iter->Seek("");
+    iter->SeekToLast();
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ("eee", iter->Entry().key.ToString());
+    ASSERT_EQ("eee", iter->Entry().value.ToString());
+    iter->Prev();
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ("aaa", iter->Entry().key.ToString());
+    ASSERT_EQ("aaa", iter->Entry().value.ToString());
+    iter->Prev();
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(!iter->Valid());
+
+    iter->SeekToFirst();
     ASSERT_OK(iter->status());
     ASSERT_TRUE(iter->Valid());
     ASSERT_EQ("aaa", iter->Entry().key.ToString());
@@ -391,6 +448,19 @@ TEST(WriteBatchWithIndexTest, TestOverwriteKey) {
     iter->Next();
     ASSERT_OK(iter->status());
     ASSERT_TRUE(!iter->Valid());
+
+    iter->SeekToLast();
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ("a11", iter->Entry().key.ToString());
+    ASSERT_EQ("a11", iter->Entry().value.ToString());
+    iter->Prev();
+
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(iter->Valid());
+    ASSERT_EQ("a33", iter->Entry().key.ToString());
+    ASSERT_TRUE(iter->Entry().type == WriteType::kDeleteRecord);
+    iter->Prev();
+    ASSERT_TRUE(!iter->Valid());
   }
 }