From 8d3819369f9704b705053a1f2112dbf1d7036c20 Mon Sep 17 00:00:00 2001
From: Igor Canadi <icanadi@fb.com>
Date: Mon, 2 Feb 2015 22:29:43 -0800
Subject: [PATCH] NewIteratorWithBase() for default column family

Summary: I'm moving mongo to a single column family, so I need DeltaBase iterator with default column family.

Test Plan: Added unit test

Reviewers: sdong

Reviewed By: sdong

Subscribers: dhruba, leveldb

Differential Revision: https://reviews.facebook.net/D32589
---
 .../utilities/write_batch_with_index.h        |  2 ++
 .../write_batch_with_index.cc                 | 16 +++++++--
 .../write_batch_with_index_test.cc            | 33 +++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/include/rocksdb/utilities/write_batch_with_index.h b/include/rocksdb/utilities/write_batch_with_index.h
index ee5ec198e..566934b70 100644
--- a/include/rocksdb/utilities/write_batch_with_index.h
+++ b/include/rocksdb/utilities/write_batch_with_index.h
@@ -106,6 +106,8 @@ class WriteBatchWithIndex {
   // base_iterator as base
   Iterator* NewIteratorWithBase(ColumnFamilyHandle* column_family,
                                 Iterator* base_iterator);
+  // default column family
+  Iterator* NewIteratorWithBase(Iterator* base_iterator);
 
  private:
   struct Rep;
diff --git a/utilities/write_batch_with_index/write_batch_with_index.cc b/utilities/write_batch_with_index/write_batch_with_index.cc
index 761b955a1..160e7dac7 100644
--- a/utilities/write_batch_with_index/write_batch_with_index.cc
+++ b/utilities/write_batch_with_index/write_batch_with_index.cc
@@ -317,9 +317,9 @@ struct WriteBatchIndexEntry {
 
 class WriteBatchEntryComparator {
  public:
-  WriteBatchEntryComparator(const Comparator* default_comparator,
+  WriteBatchEntryComparator(const Comparator* _default_comparator,
                             const ReadableWriteBatch* write_batch)
-      : default_comparator_(default_comparator), write_batch_(write_batch) {}
+      : default_comparator_(_default_comparator), write_batch_(write_batch) {}
   // Compare a and b. Return a negative value if a is less than b, 0 if they
   // are equal, and a positive value if a is greater than b
   int operator()(const WriteBatchIndexEntry* entry1,
@@ -333,6 +333,8 @@ class WriteBatchEntryComparator {
     cf_comparator_map_[column_family_id] = comparator;
   }
 
+  const Comparator* default_comparator() { return default_comparator_; }
+
  private:
   const Comparator* default_comparator_;
   std::unordered_map<uint32_t, const Comparator*> cf_comparator_map_;
@@ -590,6 +592,16 @@ Iterator* WriteBatchWithIndex::NewIteratorWithBase(
                                GetColumnFamilyUserComparator(column_family));
 }
 
+Iterator* WriteBatchWithIndex::NewIteratorWithBase(Iterator* base_iterator) {
+  if (rep->overwrite_key == false) {
+    assert(false);
+    return nullptr;
+  }
+  // default column family's comparator
+  return new BaseDeltaIterator(base_iterator, NewIterator(),
+                               rep->comparator.default_comparator());
+}
+
 void WriteBatchWithIndex::Put(ColumnFamilyHandle* column_family,
                               const Slice& key, const Slice& value) {
   rep->SetLastEntryOffset();
diff --git a/utilities/write_batch_with_index/write_batch_with_index_test.cc b/utilities/write_batch_with_index/write_batch_with_index_test.cc
index b573699db..f5d6a55a3 100644
--- a/utilities/write_batch_with_index/write_batch_with_index_test.cc
+++ b/utilities/write_batch_with_index/write_batch_with_index_test.cc
@@ -842,6 +842,39 @@ TEST(WriteBatchWithIndexTest, TestIteraratorWithBaseReverseCmp) {
     iter->Seek("a");
     AssertIter(iter.get(), "a", "aa");
   }
+
+  // default column family
+  batch.Put("a", "b");
+  {
+    KVMap map;
+    map["b"] = "";
+    std::unique_ptr<Iterator> iter(batch.NewIteratorWithBase(new KVIter(&map)));
+
+    iter->SeekToFirst();
+    AssertIter(iter.get(), "a", "b");
+    iter->Next();
+    AssertIter(iter.get(), "b", "");
+    iter->Next();
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(!iter->Valid());
+
+    iter->SeekToLast();
+    AssertIter(iter.get(), "b", "");
+    iter->Prev();
+    AssertIter(iter.get(), "a", "b");
+    iter->Prev();
+    ASSERT_OK(iter->status());
+    ASSERT_TRUE(!iter->Valid());
+
+    iter->Seek("b");
+    AssertIter(iter.get(), "b", "");
+
+    iter->Prev();
+    AssertIter(iter.get(), "a", "b");
+
+    iter->Seek("0");
+    AssertIter(iter.get(), "a", "b");
+  }
 }
 
 }  // namespace