Cleanup multiple implementations of VectorIterator (#8901)

Summary:
There were three implementations of VectorIterator (util/vector_iterator, test_util/testutil.h and LoggingForwardVectorIterator).  Merged them into one class to increase code coverage/testing and reduce duplication.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/8901

Reviewed By: pdillinger

Differential Revision: D31022673

Pulled By: mrambacher

fbshipit-source-id: 8e3acbd2dfd60b4df609d02cc72846de2389d531

db/blob/blob_counting_iterator_test.cc

@@ -14,6 +14,7 @@
 #include "db/dbformat.h"
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
+#include "util/vector_iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -67,7 +68,7 @@ TEST(BlobCountingIteratorTest, CountBlobs) {
 
   assert(keys.size() == values.size());
 
-  test::VectorIterator input(keys, values);
+  VectorIterator input(keys, values);
   BlobGarbageMeter blob_garbage_meter;
 
   BlobCountingIterator blob_counter(&input, &blob_garbage_meter);
@@ -307,7 +308,7 @@ TEST(BlobCountingIteratorTest, CorruptBlobIndex) {
 
   assert(keys.size() == values.size());
 
-  test::VectorIterator input(keys, values);
+  VectorIterator input(keys, values);
   BlobGarbageMeter blob_garbage_meter;
 
   BlobCountingIterator blob_counter(&input, &blob_garbage_meter);

db/compaction/clipping_iterator_test.cc

@@ -10,22 +10,24 @@
 #include <string>
 #include <vector>
 
+#include "db/dbformat.h"
 #include "rocksdb/comparator.h"
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
+#include "util/vector_iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
 // A vector iterator which does its own bounds checking. This is for testing the
 // optimizations in the clipping iterator where we bypass the bounds checking if
 // the input iterator has already performed it.
-class BoundsCheckingVectorIterator : public test::VectorIterator {
+class BoundsCheckingVectorIterator : public VectorIterator {
  public:
   BoundsCheckingVectorIterator(const std::vector<std::string>& keys,
                                const std::vector<std::string>& values,
                                const Slice* start, const Slice* end,
                                const Comparator* cmp)
-      : VectorIterator(keys, values), start_(start), end_(end), cmp_(cmp) {
+      : VectorIterator(keys, values, cmp), start_(start), end_(end), cmp_(cmp) {
     assert(cmp_);
   }
 
@@ -105,7 +107,7 @@ TEST_P(ClippingIteratorTest, Clip) {
       use_bounds_checking_vec_it
           ? new BoundsCheckingVectorIterator(input_keys, input_values, &start,
                                              &end, BytewiseComparator())
-          : new test::VectorIterator(input_keys, input_values));
+          : new VectorIterator(input_keys, input_values, BytewiseComparator()));
 
   ClippingIterator clip(input.get(), &start, &end, BytewiseComparator());
 

db/compaction/compaction_iterator_test.cc

@@ -3,15 +3,17 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
+#include "db/compaction/compaction_iterator.h"
 
 #include <string>
 #include <vector>
 
-#include "db/compaction/compaction_iterator.h"
+#include "db/dbformat.h"
 #include "port/port.h"
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
 #include "util/string_util.h"
+#include "util/vector_iterator.h"
 #include "utilities/merge_operators.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -86,7 +88,7 @@ class FilterAllKeysCompactionFilter : public CompactionFilter {
   const char* Name() const override { return "AllKeysCompactionFilter"; }
 };
 
-class LoggingForwardVectorIterator : public InternalIterator {
+class LoggingForwardVectorIterator : public VectorIterator {
  public:
   struct Action {
     enum class Type {
@@ -108,22 +110,19 @@ class LoggingForwardVectorIterator : public InternalIterator {
 
   LoggingForwardVectorIterator(const std::vector<std::string>& keys,
                                const std::vector<std::string>& values)
-      : keys_(keys), values_(values), current_(keys.size()) {
-    assert(keys_.size() == values_.size());
+      : VectorIterator(keys, values) {
+    current_ = keys_.size();
   }
 
-  bool Valid() const override { return current_ < keys_.size(); }
-
   void SeekToFirst() override {
     log.emplace_back(Action::Type::SEEK_TO_FIRST);
-    current_ = 0;
+    VectorIterator::SeekToFirst();
   }
   void SeekToLast() override { assert(false); }
 
   void Seek(const Slice& target) override {
     log.emplace_back(Action::Type::SEEK, target.ToString());
-    current_ = std::lower_bound(keys_.begin(), keys_.end(), target.ToString()) -
-               keys_.begin();
+    VectorIterator::Seek(target);
   }
 
   void SeekForPrev(const Slice& /*target*/) override { assert(false); }
@@ -131,27 +130,20 @@ class LoggingForwardVectorIterator : public InternalIterator {
   void Next() override {
     assert(Valid());
     log.emplace_back(Action::Type::NEXT);
-    current_++;
+    VectorIterator::Next();
   }
   void Prev() override { assert(false); }
 
   Slice key() const override {
     assert(Valid());
-    return Slice(keys_[current_]);
+    return VectorIterator::key();
   }
   Slice value() const override {
     assert(Valid());
-    return Slice(values_[current_]);
+    return VectorIterator::value();
   }
 
-  Status status() const override { return Status::OK(); }
-
   std::vector<Action> log;
-
- private:
-  std::vector<std::string> keys_;
-  std::vector<std::string> values_;
-  size_t current_;
 };
 
 class FakeCompaction : public CompactionIterator::CompactionProxy {
@@ -244,7 +236,7 @@ class CompactionIteratorTest : public testing::TestWithParam<bool> {
       bool key_not_exists_beyond_output_level = false,
       const std::string* full_history_ts_low = nullptr) {
     std::unique_ptr<InternalIterator> unfragmented_range_del_iter(
-        new test::VectorIterator(range_del_ks, range_del_vs));
+        new VectorIterator(range_del_ks, range_del_vs, &icmp_));
     auto tombstone_list = std::make_shared<FragmentedRangeTombstoneList>(
         std::move(unfragmented_range_del_iter), icmp_);
     std::unique_ptr<FragmentedRangeTombstoneIterator> range_del_iter(

db/merge_helper_test.cc

@@ -3,30 +3,33 @@
 //  COPYING file in the root directory) and Apache 2.0 License
 //  (found in the LICENSE.Apache file in the root directory).
 
+#include "db/merge_helper.h"
+
 #include <algorithm>
 #include <string>
 #include <vector>
 
-#include "db/merge_helper.h"
+#include "db/dbformat.h"
 #include "rocksdb/comparator.h"
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
 #include "util/coding.h"
+#include "util/vector_iterator.h"
 #include "utilities/merge_operators.h"
 
 namespace ROCKSDB_NAMESPACE {
 
 class MergeHelperTest : public testing::Test {
  public:
-  MergeHelperTest() { env_ = Env::Default(); }
+  MergeHelperTest() : icmp_(BytewiseComparator()) { env_ = Env::Default(); }
 
   ~MergeHelperTest() override = default;
 
   Status Run(SequenceNumber stop_before, bool at_bottom,
              SequenceNumber latest_snapshot = 0) {
-    iter_.reset(new test::VectorIterator(ks_, vs_));
+    iter_.reset(new VectorIterator(ks_, vs_, &icmp_));
     iter_->SeekToFirst();
-    merge_helper_.reset(new MergeHelper(env_, BytewiseComparator(),
+    merge_helper_.reset(new MergeHelper(env_, icmp_.user_comparator(),
                                         merge_op_.get(), filter_.get(), nullptr,
                                         false, latest_snapshot));
     return merge_helper_->MergeUntil(iter_.get(), nullptr /* range_del_agg */,
@@ -45,7 +48,8 @@ class MergeHelperTest : public testing::Test {
   }
 
   Env* env_;
-  std::unique_ptr<test::VectorIterator> iter_;
+  InternalKeyComparator icmp_;
+  std::unique_ptr<VectorIterator> iter_;
   std::shared_ptr<MergeOperator> merge_op_;
   std::unique_ptr<MergeHelper> merge_helper_;
   std::vector<std::string> ks_;

db/range_del_aggregator_bench.cc

@@ -19,15 +19,16 @@ int main() {
 #include <string>
 #include <vector>
 
+#include "db/dbformat.h"
 #include "db/range_del_aggregator.h"
 #include "db/range_tombstone_fragmenter.h"
 #include "rocksdb/comparator.h"
 #include "rocksdb/system_clock.h"
-#include "test_util/testutil.h"
 #include "util/coding.h"
 #include "util/gflags_compat.h"
 #include "util/random.h"
 #include "util/stop_watch.h"
+#include "util/vector_iterator.h"
 
 using GFLAGS_NAMESPACE::ParseCommandLineFlags;
 
@@ -146,8 +147,8 @@ std::unique_ptr<InternalIterator> MakeRangeDelIterator(
     keys.push_back(key_and_value.first.Encode().ToString());
     values.push_back(key_and_value.second.ToString());
   }
-  return std::unique_ptr<test::VectorIterator>(
-      new test::VectorIterator(keys, values));
+  return std::unique_ptr<VectorIterator>(
+      new VectorIterator(keys, values, &icmp));
 }
 
 // convert long to a big-endian slice key
@@ -207,8 +208,6 @@ int main(int argc, char** argv) {
                 ROCKSDB_NAMESPACE::Key(start), ROCKSDB_NAMESPACE::Key(end), j);
       }
 
-      auto range_del_iter =
-          ROCKSDB_NAMESPACE::MakeRangeDelIterator(persistent_range_tombstones);
       fragmented_range_tombstone_lists.emplace_back(
           new ROCKSDB_NAMESPACE::FragmentedRangeTombstoneList(
               ROCKSDB_NAMESPACE::MakeRangeDelIterator(

db/range_del_aggregator_test.cc

@@ -13,6 +13,7 @@
 #include "db/dbformat.h"
 #include "db/range_tombstone_fragmenter.h"
 #include "test_util/testutil.h"
+#include "util/vector_iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -30,8 +31,8 @@ std::unique_ptr<InternalIterator> MakeRangeDelIter(
     keys.push_back(key_and_value.first.Encode().ToString());
     values.push_back(key_and_value.second.ToString());
   }
-  return std::unique_ptr<test::VectorIterator>(
-      new test::VectorIterator(keys, values));
+  return std::unique_ptr<VectorIterator>(
+      new VectorIterator(keys, values, &bytewise_icmp));
 }
 
 std::vector<std::unique_ptr<FragmentedRangeTombstoneList>>

db/range_tombstone_fragmenter_test.cc

@@ -6,8 +6,10 @@
 #include "db/range_tombstone_fragmenter.h"
 
 #include "db/db_test_util.h"
+#include "db/dbformat.h"
 #include "rocksdb/comparator.h"
 #include "test_util/testutil.h"
+#include "util/vector_iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -25,8 +27,8 @@ std::unique_ptr<InternalIterator> MakeRangeDelIter(
     keys.push_back(key_and_value.first.Encode().ToString());
     values.push_back(key_and_value.second.ToString());
   }
-  return std::unique_ptr<test::VectorIterator>(
-      new test::VectorIterator(keys, values));
+  return std::unique_ptr<VectorIterator>(
+      new VectorIterator(keys, values, &bytewise_icmp));
 }
 
 void CheckIterPosition(const RangeTombstone& tombstone,

table/merger_test.cc

@@ -10,6 +10,7 @@
 #include "test_util/testharness.h"
 #include "test_util/testutil.h"
 #include "util/random.h"
+#include "util/vector_iterator.h"
 
 namespace ROCKSDB_NAMESPACE {
 
@@ -101,14 +102,14 @@ class MergerTest : public testing::Test {
     std::vector<InternalIterator*> small_iterators;
     for (size_t i = 0; i < num_iterators; ++i) {
       auto strings = GenerateStrings(strings_per_iterator, letters_per_string);
-      small_iterators.push_back(new test::VectorIterator(strings));
+      small_iterators.push_back(new VectorIterator(strings, strings, &icomp_));
       all_keys_.insert(all_keys_.end(), strings.begin(), strings.end());
     }
 
     merging_iterator_.reset(
         NewMergingIterator(&icomp_, &small_iterators[0],
                            static_cast<int>(small_iterators.size())));
-    single_iterator_.reset(new test::VectorIterator(all_keys_));
+    single_iterator_.reset(new VectorIterator(all_keys_, all_keys_, &icomp_));
   }
 
   InternalKeyComparator icomp_;

test_util/testutil.h

@@ -128,58 +128,6 @@ class SimpleSuffixReverseComparator : public Comparator {
 // endian machines.
 extern const Comparator* Uint64Comparator();
 
-// Iterator over a vector of keys/values
-class VectorIterator : public InternalIterator {
- public:
-  explicit VectorIterator(const std::vector<std::string>& keys)
-      : keys_(keys), current_(keys.size()) {
-    std::sort(keys_.begin(), keys_.end());
-    values_.resize(keys.size());
-  }
-
-  VectorIterator(const std::vector<std::string>& keys,
-      const std::vector<std::string>& values)
-    : keys_(keys), values_(values), current_(keys.size()) {
-    assert(keys_.size() == values_.size());
-  }
-
-  virtual bool Valid() const override { return current_ < keys_.size(); }
-
-  virtual void SeekToFirst() override { current_ = 0; }
-  virtual void SeekToLast() override { current_ = keys_.size() - 1; }
-
-  virtual void Seek(const Slice& target) override {
-    current_ = std::lower_bound(keys_.begin(), keys_.end(), target.ToString()) -
-               keys_.begin();
-  }
-
-  virtual void SeekForPrev(const Slice& target) override {
-    current_ = std::upper_bound(keys_.begin(), keys_.end(), target.ToString()) -
-               keys_.begin();
-    if (!Valid()) {
-      SeekToLast();
-    } else {
-      Prev();
-    }
-  }
-
-  virtual void Next() override { current_++; }
-  virtual void Prev() override { current_--; }
-
-  virtual Slice key() const override { return Slice(keys_[current_]); }
-  virtual Slice value() const override { return Slice(values_[current_]); }
-
-  virtual Status status() const override { return Status::OK(); }
-
-  virtual bool IsKeyPinned() const override { return true; }
-  virtual bool IsValuePinned() const override { return true; }
-
- private:
-  std::vector<std::string> keys_;
-  std::vector<std::string> values_;
-  size_t current_;
-};
-
 class StringSink : public FSWritableFile {
  public:
   std::string contents_;

util/vector_iterator.h

@@ -16,18 +16,20 @@ namespace ROCKSDB_NAMESPACE {
 class VectorIterator : public InternalIterator {
  public:
   VectorIterator(std::vector<std::string> keys, std::vector<std::string> values,
-                 const InternalKeyComparator* icmp)
+                 const Comparator* icmp = nullptr)
       : keys_(std::move(keys)),
         values_(std::move(values)),
-        indexed_cmp_(icmp, &keys_),
-        current_(0) {
+        current_(keys_.size()),
+        indexed_cmp_(icmp, &keys_) {
     assert(keys_.size() == values_.size());
 
     indices_.reserve(keys_.size());
     for (size_t i = 0; i < keys_.size(); i++) {
       indices_.push_back(i);
     }
-    std::sort(indices_.begin(), indices_.end(), indexed_cmp_);
+    if (icmp != nullptr) {
+      std::sort(indices_.begin(), indices_.end(), indexed_cmp_);
+    }
   }
 
   virtual bool Valid() const override {
@@ -38,15 +40,27 @@ class VectorIterator : public InternalIterator {
   virtual void SeekToLast() override { current_ = indices_.size() - 1; }
 
   virtual void Seek(const Slice& target) override {
-    current_ = std::lower_bound(indices_.begin(), indices_.end(), target,
-                                indexed_cmp_) -
-               indices_.begin();
+    if (indexed_cmp_.cmp != nullptr) {
+      current_ = std::lower_bound(indices_.begin(), indices_.end(), target,
+                                  indexed_cmp_) -
+                 indices_.begin();
+    } else {
+      current_ =
+          std::lower_bound(keys_.begin(), keys_.end(), target.ToString()) -
+          keys_.begin();
+    }
   }
 
   virtual void SeekForPrev(const Slice& target) override {
-    current_ = std::lower_bound(indices_.begin(), indices_.end(), target,
-                                indexed_cmp_) -
-               indices_.begin();
+    if (indexed_cmp_.cmp != nullptr) {
+      current_ = std::upper_bound(indices_.begin(), indices_.end(), target,
+                                  indexed_cmp_) -
+                 indices_.begin();
+    } else {
+      current_ =
+          std::upper_bound(keys_.begin(), keys_.end(), target.ToString()) -
+          keys_.begin();
+    }
     if (!Valid()) {
       SeekToLast();
     } else {
@@ -69,9 +83,14 @@ class VectorIterator : public InternalIterator {
   virtual bool IsKeyPinned() const override { return true; }
   virtual bool IsValuePinned() const override { return true; }
 
+ protected:
+  std::vector<std::string> keys_;
+  std::vector<std::string> values_;
+  size_t current_;
+
  private:
   struct IndexedKeyComparator {
-    IndexedKeyComparator(const InternalKeyComparator* c,
+    IndexedKeyComparator(const Comparator* c,
                          const std::vector<std::string>* ks)
         : cmp(c), keys(ks) {}
 
@@ -87,15 +106,12 @@ class VectorIterator : public InternalIterator {
       return cmp->Compare(a, (*keys)[b]) < 0;
     }
 
-    const InternalKeyComparator* cmp;
+    const Comparator* cmp;
     const std::vector<std::string>* keys;
   };
 
-  std::vector<std::string> keys_;
-  std::vector<std::string> values_;
   IndexedKeyComparator indexed_cmp_;
   std::vector<size_t> indices_;
-  size_t current_;
 };
 
 }  // namespace ROCKSDB_NAMESPACE