Adding a user comparator for comparing Uint64 slices.

Summary:
- New Uint64 comparator
- Modify Reader and Builder to take custom user comparators instead of bytewise comparator
- Modify logic for choosing unused user key in builder
- Modify iterator logic in reader
- test changes

Test Plan:
cuckoo_table_{builder,reader,db}_test
make check all

Reviewers: ljin, sdong

Reviewed By: ljin

Subscribers: dhruba, leveldb

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

db/cuckoo_table_db_test.cc

@@ -131,8 +131,6 @@ TEST(CuckooTableDBTest, Flush) {
   ASSERT_EQ("v2", Get("key2"));
   ASSERT_EQ("v3", Get("key3"));
   ASSERT_EQ("NOT_FOUND", Get("key4"));
-  ASSERT_EQ("Invalid argument: Length of key is invalid.", Get("somelongkey"));
-  ASSERT_EQ("Invalid argument: Length of key is invalid.", Get("s"));
 
   // Now add more keys and flush.
   ASSERT_OK(Put("key4", "v4"));
@@ -195,6 +193,38 @@ static std::string Key(int i) {
   snprintf(buf, sizeof(buf), "key_______%06d", i);
   return std::string(buf);
 }
+static std::string Uint64Key(uint64_t i) {
+  std::string str;
+  str.resize(8);
+  memcpy(&str[0], static_cast<void*>(&i), 8);
+  return str;
+}
+}  // namespace.
+
+TEST(CuckooTableDBTest, Uint64Comparator) {
+  Options options = CurrentOptions();
+  options.comparator = test::Uint64Comparator();
+  Reopen(&options);
+
+  ASSERT_OK(Put(Uint64Key(1), "v1"));
+  ASSERT_OK(Put(Uint64Key(2), "v2"));
+  ASSERT_OK(Put(Uint64Key(3), "v3"));
+  dbfull()->TEST_FlushMemTable();
+
+  ASSERT_EQ("v1", Get(Uint64Key(1)));
+  ASSERT_EQ("v2", Get(Uint64Key(2)));
+  ASSERT_EQ("v3", Get(Uint64Key(3)));
+  ASSERT_EQ("NOT_FOUND", Get(Uint64Key(4)));
+
+  // Add more keys.
+  ASSERT_OK(Delete(Uint64Key(2)));  // Delete.
+  ASSERT_OK(Put(Uint64Key(3), "v0"));  // Update.
+  ASSERT_OK(Put(Uint64Key(4), "v4"));
+  dbfull()->TEST_FlushMemTable();
+  ASSERT_EQ("v1", Get(Uint64Key(1)));
+  ASSERT_EQ("NOT_FOUND", Get(Uint64Key(2)));
+  ASSERT_EQ("v0", Get(Uint64Key(3)));
+  ASSERT_EQ("v4", Get(Uint64Key(4)));
 }
 
 TEST(CuckooTableDBTest, CompactionTrigger) {

db/db_bench.cc

@@ -163,6 +163,7 @@ DEFINE_int32(duration, 0, "Time in seconds for the random-ops tests to run."
 
 DEFINE_int32(value_size, 100, "Size of each value");
 
+DEFINE_bool(use_uint64_comparator, false, "use Uint64 user comparator");
 
 static bool ValidateKeySize(const char* flagname, int32_t value) {
   return true;
@@ -1650,6 +1651,13 @@ class Benchmark {
       options.prefix_extractor.reset(
           NewFixedPrefixTransform(FLAGS_prefix_size));
     }
+    if (FLAGS_use_uint64_comparator) {
+      options.comparator = test::Uint64Comparator();
+      if (FLAGS_key_size != 8) {
+        fprintf(stderr, "Using Uint64 comparator but key size is not 8.\n");
+        exit(1);
+      }
+    }
     options.memtable_prefix_bloom_bits = FLAGS_memtable_bloom_bits;
     options.bloom_locality = FLAGS_bloom_locality;
     options.max_open_files = FLAGS_open_files;

table/cuckoo_table_builder.cc

@@ -39,6 +39,7 @@ extern const uint64_t kCuckooTableMagicNumber = 0x926789d0c5f17873ull;
 CuckooTableBuilder::CuckooTableBuilder(
     WritableFile* file, double hash_table_ratio,
     uint32_t max_num_hash_table, uint32_t max_search_depth,
+    const Comparator* user_comparator,
     uint64_t (*get_slice_hash)(const Slice&, uint32_t, uint64_t))
     : num_hash_table_(2),
       file_(file),
@@ -47,6 +48,7 @@ CuckooTableBuilder::CuckooTableBuilder(
       max_search_depth_(max_search_depth),
       is_last_level_file_(false),
       has_seen_first_key_(false),
+      ucomp_(user_comparator),
       get_slice_hash_(get_slice_hash),
       closed_(false) {
   properties_.num_entries = 0;
@@ -73,6 +75,8 @@ void CuckooTableBuilder::Add(const Slice& key, const Slice& value) {
   if (!has_seen_first_key_) {
     is_last_level_file_ = ikey.sequence == 0;
     has_seen_first_key_ = true;
+    smallest_user_key_.assign(ikey.user_key.data(), ikey.user_key.size());
+    largest_user_key_.assign(ikey.user_key.data(), ikey.user_key.size());
   }
   // Even if one sequence number is non-zero, then it is not last level.
   assert(!is_last_level_file_ || ikey.sequence == 0);
@@ -82,23 +86,17 @@ void CuckooTableBuilder::Add(const Slice& key, const Slice& value) {
   } else {
     kvs_.emplace_back(std::make_pair(key.ToString(), value.ToString()));
   }
-
   properties_.num_entries++;
 
-  // We assume that the keys are inserted in sorted order as determined by
-  // Byte-wise comparator. To identify an unused key, which will be used in
-  // filling empty buckets in the table, we try to find gaps between successive
-  // keys inserted (ie, latest key and previous in kvs_).
-  if (unused_user_key_.empty() && kvs_.size() > 1) {
-    std::string prev_key = is_last_level_file_ ? kvs_[kvs_.size()-1].first
-      : ExtractUserKey(kvs_[kvs_.size()-1].first).ToString();
-    std::string new_user_key = prev_key;
-    new_user_key.back()++;
-    // We ignore carry-overs and check that it is larger than previous key.
-    if (Slice(new_user_key).compare(Slice(prev_key)) > 0 &&
-        Slice(new_user_key).compare(ikey.user_key) < 0) {
-      unused_user_key_ = new_user_key;
-    }
+  // In order to fill the empty buckets in the hash table, we identify a
+  // key which is not used so far (unused_user_key). We determine this by
+  // maintaining smallest and largest keys inserted so far in bytewise order
+  // and use them to find a key outside this range in Finish() operation.
+  // Note that this strategy is independent of user comparator used here.
+  if (ikey.user_key.compare(smallest_user_key_) < 0) {
+    smallest_user_key_.assign(ikey.user_key.data(), ikey.user_key.size());
+  } else if (ikey.user_key.compare(largest_user_key_) > 0) {
+    largest_user_key_.assign(ikey.user_key.data(), ikey.user_key.size());
   }
 }
 
@@ -119,7 +117,7 @@ Status CuckooTableBuilder::MakeHashTable(std::vector<CuckooBucket>* buckets) {
         bucket_found = true;
         break;
       } else {
-        if (user_key.compare(is_last_level_file_
+        if (ucomp_->Compare(user_key, is_last_level_file_
               ? Slice(kvs_[(*buckets)[hash_val].vector_idx].first)
               : ExtractUserKey(
                 kvs_[(*buckets)[hash_val].vector_idx].first)) == 0) {
@@ -160,31 +158,37 @@ Status CuckooTableBuilder::Finish() {
   if (!s.ok()) {
     return s;
   }
-  if (unused_user_key_.empty() && !kvs_.empty()) {
-    // Try to find the key next to last key by handling carryovers.
-    std::string last_key =
-      is_last_level_file_ ? kvs_[kvs_.size()-1].first
-      : ExtractUserKey(kvs_[kvs_.size()-1].first).ToString();
-    std::string new_user_key = last_key;
-    int curr_pos = new_user_key.size() - 1;
+  // Determine unused_user_key to fill empty buckets.
+  std::string unused_bucket;
+  if (!kvs_.empty()) {
+    std::string unused_user_key = smallest_user_key_;
+    int curr_pos = unused_user_key.size() - 1;
     while (curr_pos >= 0) {
-      ++new_user_key[curr_pos];
-      if (new_user_key > last_key) {
-        unused_user_key_ = new_user_key;
+      --unused_user_key[curr_pos];
+      if (Slice(unused_user_key).compare(smallest_user_key_) < 0) {
         break;
       }
       --curr_pos;
     }
+    if (curr_pos < 0) {
+      // Try using the largest key to identify an unused key.
+      unused_user_key = largest_user_key_;
+      curr_pos = unused_user_key.size() - 1;
+      while (curr_pos >= 0) {
+        ++unused_user_key[curr_pos];
+        if (Slice(unused_user_key).compare(largest_user_key_) > 0) {
+          break;
+        }
+        --curr_pos;
+      }
+    }
     if (curr_pos < 0) {
       return Status::Corruption("Unable to find unused key");
     }
-  }
-  std::string unused_bucket;
-  if (!kvs_.empty()) {
     if (is_last_level_file_) {
-      unused_bucket = unused_user_key_;
+      unused_bucket = unused_user_key;
     } else {
-      ParsedInternalKey ikey(unused_user_key_, 0, kTypeValue);
+      ParsedInternalKey ikey(unused_user_key, 0, kTypeValue);
       AppendInternalKey(&unused_bucket, ikey);
     }
   }

table/cuckoo_table_builder.h

@@ -22,7 +22,7 @@ class CuckooTableBuilder: public TableBuilder {
  public:
   CuckooTableBuilder(
       WritableFile* file, double hash_table_ratio, uint32_t max_num_hash_table,
-      uint32_t max_search_depth,
+      uint32_t max_search_depth, const Comparator* user_comparator,
       uint64_t (*get_slice_hash)(const Slice&, uint32_t, uint64_t));
 
   // REQUIRES: Either Finish() or Abandon() has been called.
@@ -83,9 +83,11 @@ class CuckooTableBuilder: public TableBuilder {
   std::vector<std::pair<std::string, std::string>> kvs_;
   TableProperties properties_;
   bool has_seen_first_key_;
+  const Comparator* ucomp_;
   uint64_t (*get_slice_hash_)(const Slice& s, uint32_t index,
     uint64_t max_num_buckets);
-  std::string unused_user_key_ = "";
+  std::string largest_user_key_ = "";
+  std::string smallest_user_key_ = "";
 
   bool closed_;  // Either Finish() or Abandon() has been called.
 

table/cuckoo_table_builder_test.cc

@@ -119,7 +119,7 @@ TEST(CuckooBuilderTest, SuccessWithEmptyFile) {
   fname = test::TmpDir() + "/NoCollisionFullKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      4, 100, GetSliceHash);
+      4, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   ASSERT_OK(builder.Finish());
   ASSERT_OK(writable_file->Close());
@@ -146,7 +146,7 @@ TEST(CuckooBuilderTest, WriteSuccessNoCollisionFullKey) {
   fname = test::TmpDir() + "/NoCollisionFullKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(keys[i]), Slice(values[i]));
@@ -157,7 +157,7 @@ TEST(CuckooBuilderTest, WriteSuccessNoCollisionFullKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = GetInternalKey("key05", true);
+  std::string expected_unused_bucket = GetInternalKey("key00", true);
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 2, false);
@@ -183,7 +183,7 @@ TEST(CuckooBuilderTest, WriteSuccessWithCollisionFullKey) {
   fname = test::TmpDir() + "/WithCollisionFullKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(keys[i]), Slice(values[i]));
@@ -194,7 +194,7 @@ TEST(CuckooBuilderTest, WriteSuccessWithCollisionFullKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = GetInternalKey("key05", true);
+  std::string expected_unused_bucket = GetInternalKey("key00", true);
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 4, false);
@@ -225,7 +225,7 @@ TEST(CuckooBuilderTest, WithCollisionPathFullKey) {
   fname = test::TmpDir() + "/WithCollisionPathFullKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(keys[i]), Slice(values[i]));
@@ -236,7 +236,7 @@ TEST(CuckooBuilderTest, WithCollisionPathFullKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = GetInternalKey("key06", true);
+  std::string expected_unused_bucket = GetInternalKey("key00", true);
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 2, false);
@@ -258,7 +258,7 @@ TEST(CuckooBuilderTest, WriteSuccessNoCollisionUserKey) {
   fname = test::TmpDir() + "/NoCollisionUserKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(GetInternalKey(user_keys[i], true)), Slice(values[i]));
@@ -269,7 +269,7 @@ TEST(CuckooBuilderTest, WriteSuccessNoCollisionUserKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = user_keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = "key05";
+  std::string expected_unused_bucket = "key00";
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(user_keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 2, true);
@@ -291,7 +291,7 @@ TEST(CuckooBuilderTest, WriteSuccessWithCollisionUserKey) {
   fname = test::TmpDir() + "/WithCollisionUserKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(GetInternalKey(user_keys[i], true)), Slice(values[i]));
@@ -302,7 +302,7 @@ TEST(CuckooBuilderTest, WriteSuccessWithCollisionUserKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = user_keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = "key05";
+  std::string expected_unused_bucket = "key00";
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(user_keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 4, true);
@@ -326,7 +326,7 @@ TEST(CuckooBuilderTest, WithCollisionPathUserKey) {
   fname = test::TmpDir() + "/WithCollisionPathUserKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 2, GetSliceHash);
+      num_hash_fun, 2, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(GetInternalKey(user_keys[i], true)), Slice(values[i]));
@@ -337,7 +337,7 @@ TEST(CuckooBuilderTest, WithCollisionPathUserKey) {
   ASSERT_OK(writable_file->Close());
 
   uint32_t expected_max_buckets = user_keys.size() / kHashTableRatio;
-  std::string expected_unused_bucket = "key06";
+  std::string expected_unused_bucket = "key00";
   expected_unused_bucket += std::string(values[0].size(), 'a');
   CheckFileContents(user_keys, values, expected_locations,
       expected_unused_bucket, expected_max_buckets, 2, true);
@@ -362,7 +362,7 @@ TEST(CuckooBuilderTest, FailWhenCollisionPathTooLong) {
   fname = test::TmpDir() + "/WithCollisionPathUserKey";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 2, GetSliceHash);
+      num_hash_fun, 2, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
   for (uint32_t i = 0; i < user_keys.size(); i++) {
     builder.Add(Slice(GetInternalKey(user_keys[i], false)), Slice("value"));
@@ -382,7 +382,7 @@ TEST(CuckooBuilderTest, FailWhenSameKeyInserted) {
   fname = test::TmpDir() + "/FailWhenSameKeyInserted";
   ASSERT_OK(env_->NewWritableFile(fname, &writable_file, env_options_));
   CuckooTableBuilder builder(writable_file.get(), kHashTableRatio,
-      num_hash_fun, 100, GetSliceHash);
+      num_hash_fun, 100, BytewiseComparator(), GetSliceHash);
   ASSERT_OK(builder.status());
 
   builder.Add(Slice(GetInternalKey(user_key, false)), Slice("value1"));

table/cuckoo_table_factory.cc

@@ -36,7 +36,7 @@ Status CuckooTableFactory::NewTableReader(const Options& options,
     std::unique_ptr<RandomAccessFile>&& file, uint64_t file_size,
     std::unique_ptr<TableReader>* table) const {
   std::unique_ptr<CuckooTableReader> new_reader(new CuckooTableReader(options,
-      std::move(file), file_size, GetSliceMurmurHash));
+      std::move(file), file_size, icomp.user_comparator(), GetSliceMurmurHash));
   Status s = new_reader->status();
   if (s.ok()) {
     *table = std::move(new_reader);
@@ -48,7 +48,8 @@ TableBuilder* CuckooTableFactory::NewTableBuilder(
     const Options& options, const InternalKeyComparator& internal_comparator,
     WritableFile* file, CompressionType compression_type) const {
   return new CuckooTableBuilder(file, hash_table_ratio_, kMaxNumHashTable,
-      max_search_depth_, GetSliceMurmurHash);
+      max_search_depth_, internal_comparator.user_comparator(),
+      GetSliceMurmurHash);
 }
 
 std::string CuckooTableFactory::GetPrintableTableOptions() const {

table/cuckoo_table_factory.h

@@ -21,7 +21,6 @@ extern uint64_t GetSliceMurmurHash(const Slice& s, uint32_t index,
 // - Key length and Value length are fixed.
 // - Does not support Snapshot.
 // - Does not support Merge operations.
-// - Only supports Bytewise comparators.
 class CuckooTableFactory : public TableFactory {
  public:
   CuckooTableFactory(double hash_table_ratio, uint32_t max_search_depth)

table/cuckoo_table_reader.cc

@@ -28,8 +28,10 @@ CuckooTableReader::CuckooTableReader(
     const Options& options,
     std::unique_ptr<RandomAccessFile>&& file,
     uint64_t file_size,
+    const Comparator* comparator,
     uint64_t (*get_slice_hash)(const Slice&, uint32_t, uint64_t))
     : file_(std::move(file)),
+      ucomp_(comparator),
       get_slice_hash_(get_slice_hash) {
   if (!options.allow_mmap_reads) {
     status_ = Status::InvalidArgument("File is not mmaped");
@@ -85,24 +87,19 @@ Status CuckooTableReader::Get(
     bool (*result_handler)(void* arg, const ParsedInternalKey& k,
                            const Slice& v),
     void (*mark_key_may_exist_handler)(void* handle_context)) {
-  ParsedInternalKey ikey;
-  if (!ParseInternalKey(key, &ikey)) {
-    return Status::Corruption("Unable to parse key into inernal key.");
-  }
-  if ((is_last_level_ && key.size() != key_length_ + 8) ||
-      (!is_last_level_ && key.size() != key_length_)) {
-    return Status::InvalidArgument("Length of key is invalid.");
-  }
+  assert(key.size() == key_length_ + (is_last_level_ ? 8 : 0));
+  Slice user_key = ExtractUserKey(key);
   for (uint32_t hash_cnt = 0; hash_cnt < num_hash_fun_; ++hash_cnt) {
-    uint64_t hash_val = get_slice_hash_(ikey.user_key, hash_cnt, num_buckets_);
+    uint64_t hash_val = get_slice_hash_(user_key, hash_cnt, num_buckets_);
     assert(hash_val < num_buckets_);
     const char* bucket = &file_data_.data()[hash_val * bucket_length_];
-    if (unused_key_.compare(0, key_length_, bucket, key_length_) == 0) {
+    if (ucomp_->Compare(Slice(unused_key_.data(), user_key.size()),
+          Slice(bucket, user_key.size())) == 0) {
       return Status::OK();
     }
     // Here, we compare only the user key part as we support only one entry
     // per user key and we don't support sanpshot.
-    if (ikey.user_key.compare(Slice(bucket, ikey.user_key.size())) == 0) {
+    if (ucomp_->Compare(user_key, Slice(bucket, user_key.size())) == 0) {
       Slice value = Slice(&bucket[key_length_], value_length_);
       if (is_last_level_) {
         ParsedInternalKey found_ikey(Slice(bucket, key_length_), 0, kTypeValue);
@@ -121,10 +118,10 @@ Status CuckooTableReader::Get(
 }
 
 void CuckooTableReader::Prepare(const Slice& key) {
+  Slice user_key = ExtractUserKey(key);
   // Prefetching first location also helps improve Get performance.
   for (uint32_t hash_cnt = 0; hash_cnt < num_hash_fun_; ++hash_cnt) {
-    uint64_t hash_val = get_slice_hash_(ExtractUserKey(key),
-        hash_cnt, num_buckets_);
+    uint64_t hash_val = get_slice_hash_(user_key, hash_cnt, num_buckets_);
     PREFETCH(&file_data_.data()[hash_val * bucket_length_], 0, 3);
   }
 }
@@ -145,17 +142,29 @@ class CuckooTableIterator : public Iterator {
   void LoadKeysFromReader();
 
  private:
-  struct {
+  struct CompareKeys {
+    CompareKeys(const Comparator* ucomp, const bool last_level)
+      : ucomp_(ucomp),
+        is_last_level_(last_level) {}
     bool operator()(const std::pair<Slice, uint32_t>& first,
         const std::pair<Slice, uint32_t>& second) const {
-      return first.first.compare(second.first) < 0;
+      if (is_last_level_) {
+        return ucomp_->Compare(first.first, second.first) < 0;
+      } else {
+        return ucomp_->Compare(ExtractUserKey(first.first),
+            ExtractUserKey(second.first)) < 0;
+      }
     }
-  } CompareKeys;
+
+   private:
+    const Comparator* ucomp_;
+    const bool is_last_level_;
+  };
+  const CompareKeys comparator_;
   void PrepareKVAtCurrIdx();
   CuckooTableReader* reader_;
   Status status_;
   // Contains a map of keys to bucket_id sorted in key order.
-  // We assume byte-wise comparison for key ordering.
   std::vector<std::pair<Slice, uint32_t>> key_to_bucket_id_;
   // We assume that the number of items can be stored in uint32 (4 Billion).
   uint32_t curr_key_idx_;
@@ -167,7 +176,8 @@ class CuckooTableIterator : public Iterator {
 };
 
 CuckooTableIterator::CuckooTableIterator(CuckooTableReader* reader)
-  : reader_(reader),
+  : comparator_(reader->ucomp_, reader->is_last_level_),
+    reader_(reader),
     curr_key_idx_(std::numeric_limits<int32_t>::max()) {
   key_to_bucket_id_.clear();
   curr_value_.clear();
@@ -186,7 +196,7 @@ void CuckooTableIterator::LoadKeysFromReader() {
   }
   assert(key_to_bucket_id_.size() ==
       reader_->GetTableProperties()->num_entries);
-  std::sort(key_to_bucket_id_.begin(), key_to_bucket_id_.end(), CompareKeys);
+  std::sort(key_to_bucket_id_.begin(), key_to_bucket_id_.end(), comparator_);
   curr_key_idx_ = key_to_bucket_id_.size();
 }
 
@@ -208,7 +218,7 @@ void CuckooTableIterator::Seek(const Slice& target) {
   auto seek_it = std::lower_bound(key_to_bucket_id_.begin(),
       key_to_bucket_id_.end(),
       std::make_pair(target_to_search, 0),
-      CompareKeys);
+      comparator_);
   curr_key_idx_ = std::distance(key_to_bucket_id_.begin(), seek_it);
   PrepareKVAtCurrIdx();
 }

table/cuckoo_table_reader.h

@@ -29,6 +29,7 @@ class CuckooTableReader: public TableReader {
       const Options& options,
       std::unique_ptr<RandomAccessFile>&& file,
       uint64_t file_size,
+      const Comparator* user_comparator,
       uint64_t (*get_slice_hash)(const Slice&, uint32_t, uint64_t));
   ~CuckooTableReader() {}
 
@@ -70,6 +71,7 @@ class CuckooTableReader: public TableReader {
   uint32_t value_length_;
   uint32_t bucket_length_;
   uint64_t num_buckets_;
+  const Comparator* ucomp_;
   uint64_t (*get_slice_hash_)(const Slice& s, uint32_t index,
       uint64_t max_num_buckets);
 };

table/cuckoo_table_reader_test.cc

@@ -104,11 +104,12 @@ class CuckooReaderTest {
     return std::string(reinterpret_cast<char*>(&i), sizeof(i));
   }
 
-  void CreateCuckooFileAndCheckReader() {
+  void CreateCuckooFileAndCheckReader(
+      const Comparator* ucomp = BytewiseComparator()) {
     std::unique_ptr<WritableFile> writable_file;
     ASSERT_OK(env->NewWritableFile(fname, &writable_file, env_options));
     CuckooTableBuilder builder(
-        writable_file.get(), 0.9, kNumHashFunc, 100, GetSliceHash);
+        writable_file.get(), 0.9, kNumHashFunc, 100, ucomp, GetSliceHash);
     ASSERT_OK(builder.status());
     for (uint32_t key_idx = 0; key_idx < num_items; ++key_idx) {
       builder.Add(Slice(keys[key_idx]), Slice(values[key_idx]));
@@ -127,6 +128,7 @@ class CuckooReaderTest {
         options,
         std::move(read_file),
         file_size,
+        ucomp,
         GetSliceHash);
     ASSERT_OK(reader.status());
     for (uint32_t i = 0; i < num_items; ++i) {
@@ -145,13 +147,14 @@ class CuckooReaderTest {
     }
   }
 
-  void CheckIterator() {
+  void CheckIterator(const Comparator* ucomp = BytewiseComparator()) {
     std::unique_ptr<RandomAccessFile> read_file;
     ASSERT_OK(env->NewRandomAccessFile(fname, &read_file, env_options));
     CuckooTableReader reader(
         options,
         std::move(read_file),
         file_size,
+        ucomp,
         GetSliceHash);
     ASSERT_OK(reader.status());
     Iterator* it = reader.NewIterator(ReadOptions(), nullptr);
@@ -243,12 +246,40 @@ TEST(CuckooReaderTest, WhenKeyExists) {
   CreateCuckooFileAndCheckReader();
 }
 
+TEST(CuckooReaderTest, WhenKeyExistsWithUint64Comparator) {
+  SetUp(kNumHashFunc);
+  fname = test::TmpDir() + "/CuckooReaderUint64_WhenKeyExists";
+  for (uint64_t i = 0; i < num_items; i++) {
+    user_keys[i].resize(8);
+    memcpy(&user_keys[i][0], static_cast<void*>(&i), 8);
+    ParsedInternalKey ikey(user_keys[i], i + 1000, kTypeValue);
+    AppendInternalKey(&keys[i], ikey);
+    values[i] = "value" + NumToStr(i);
+    // Give disjoint hash values.
+    AddHashLookups(user_keys[i], i, kNumHashFunc);
+  }
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+  // Last level file.
+  UpdateKeys(true);
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+  // Test with collision. Make all hash values collide.
+  hash_map.clear();
+  for (uint32_t i = 0; i < num_items; i++) {
+    AddHashLookups(user_keys[i], 0, kNumHashFunc);
+  }
+  UpdateKeys(false);
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+  // Last level file.
+  UpdateKeys(true);
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+}
+
 TEST(CuckooReaderTest, CheckIterator) {
   SetUp(2*kNumHashFunc);
   fname = test::TmpDir() + "/CuckooReader_CheckIterator";
   for (uint64_t i = 0; i < num_items; i++) {
     user_keys[i] = "key" + NumToStr(i);
-    ParsedInternalKey ikey(user_keys[i], 0, kTypeValue);
+    ParsedInternalKey ikey(user_keys[i], 1000, kTypeValue);
     AppendInternalKey(&keys[i], ikey);
     values[i] = "value" + NumToStr(i);
     // Give disjoint hash values, in reverse order.
@@ -262,6 +293,26 @@ TEST(CuckooReaderTest, CheckIterator) {
   CheckIterator();
 }
 
+TEST(CuckooReaderTest, CheckIteratorUint64) {
+  SetUp(2*kNumHashFunc);
+  fname = test::TmpDir() + "/CuckooReader_CheckIterator";
+  for (uint64_t i = 0; i < num_items; i++) {
+    user_keys[i].resize(8);
+    memcpy(&user_keys[i][0], static_cast<void*>(&i), 8);
+    ParsedInternalKey ikey(user_keys[i], 1000, kTypeValue);
+    AppendInternalKey(&keys[i], ikey);
+    values[i] = "value" + NumToStr(i);
+    // Give disjoint hash values, in reverse order.
+    AddHashLookups(user_keys[i], num_items-i-1, kNumHashFunc);
+  }
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+  CheckIterator(test::Uint64Comparator());
+  // Last level file.
+  UpdateKeys(true);
+  CreateCuckooFileAndCheckReader(test::Uint64Comparator());
+  CheckIterator(test::Uint64Comparator());
+}
+
 TEST(CuckooReaderTest, WhenKeyNotFound) {
   // Add keys with colliding hash values.
   SetUp(kNumHashFunc);
@@ -281,6 +332,7 @@ TEST(CuckooReaderTest, WhenKeyNotFound) {
       options,
       std::move(read_file),
       file_size,
+      BytewiseComparator(),
       GetSliceHash);
   ASSERT_OK(reader.status());
   // Search for a key with colliding hash values.
@@ -305,31 +357,6 @@ TEST(CuckooReaderTest, WhenKeyNotFound) {
   ASSERT_EQ(0, v.call_count);
   ASSERT_OK(reader.status());
 
-  // Test read with corrupted key.
-  Slice corrupt_key("corrupt_ikey");
-  ASSERT_TRUE(!ParseInternalKey(corrupt_key, &ikey));
-  ASSERT_TRUE(reader.Get(
-        ReadOptions(), corrupt_key, &v,
-        AssertValues, nullptr).IsCorruption());
-  ASSERT_EQ(0, v.call_count);
-  ASSERT_OK(reader.status());
-
-  // Test read with key of invalid length.
-  IterKey k;
-  k.SetInternalKey("very_long_key", 0, kTypeValue);
-  ASSERT_TRUE(reader.Get(
-        ReadOptions(), k.GetKey(), &v,
-        AssertValues, nullptr).IsInvalidArgument());
-  ASSERT_EQ(0, v.call_count);
-  ASSERT_OK(reader.status());
-  k.Clear();
-  k.SetInternalKey("s", 0, kTypeValue);
-  ASSERT_TRUE(reader.Get(
-        ReadOptions(), k.GetKey(), &v,
-        AssertValues, nullptr).IsInvalidArgument());
-  ASSERT_EQ(0, v.call_count);
-  ASSERT_OK(reader.status());
-
   // Test read when key is unused key.
   std::string unused_key =
     reader.GetTableProperties()->user_collected_properties.at(
@@ -393,7 +420,7 @@ void WriteFile(const std::vector<std::string>& keys,
   ASSERT_OK(env->NewWritableFile(fname, &writable_file, env_options));
   CuckooTableBuilder builder(
       writable_file.get(), hash_ratio,
-      kMaxNumHashTable, 1000, GetSliceMurmurHash);
+      kMaxNumHashTable, 1000, test::Uint64Comparator(), GetSliceMurmurHash);
   ASSERT_OK(builder.status());
   for (uint64_t key_idx = 0; key_idx < num; ++key_idx) {
     // Value is just a part of key.
@@ -411,7 +438,8 @@ void WriteFile(const std::vector<std::string>& keys,
   ASSERT_OK(env->NewRandomAccessFile(fname, &read_file, env_options));
 
   CuckooTableReader reader(
-      options, std::move(read_file), file_size, GetSliceMurmurHash);
+      options, std::move(read_file), file_size,
+      test::Uint64Comparator(), GetSliceMurmurHash);
   ASSERT_OK(reader.status());
   ReadOptions r_options;
   for (const auto& key : keys) {
@@ -439,7 +467,8 @@ void ReadKeys(const std::vector<std::string>& keys, uint64_t num,
   ASSERT_OK(env->NewRandomAccessFile(fname, &read_file, env_options));
 
   CuckooTableReader reader(
-      options, std::move(read_file), file_size, GetSliceMurmurHash);
+      options, std::move(read_file), file_size, test::Uint64Comparator(),
+      GetSliceMurmurHash);
   ASSERT_OK(reader.status());
   const UserCollectedProperties user_props =
     reader.GetTableProperties()->user_collected_properties;

util/testutil.cc

@@ -9,6 +9,7 @@
 
 #include "util/testutil.h"
 
+#include "port/port.h"
 #include "util/random.h"
 
 namespace rocksdb {
@@ -52,5 +53,50 @@ extern Slice CompressibleString(Random* rnd, double compressed_fraction,
   return Slice(*dst);
 }
 
+namespace {
+class Uint64ComparatorImpl : public Comparator {
+ public:
+  Uint64ComparatorImpl() { }
+
+  virtual const char* Name() const override {
+    return "rocksdb.Uint64Comparator";
+  }
+
+  virtual int Compare(const Slice& a, const Slice& b) const override {
+    assert(a.size() == sizeof(uint64_t) && b.size() == sizeof(uint64_t));
+    const uint64_t* left = reinterpret_cast<const uint64_t*>(a.data());
+    const uint64_t* right = reinterpret_cast<const uint64_t*>(b.data());
+    if (*left == *right) {
+      return 0;
+    } else if (*left < *right) {
+      return -1;
+    } else {
+      return 1;
+    }
+  }
+
+  virtual void FindShortestSeparator(std::string* start,
+      const Slice& limit) const override {
+    return;
+  }
+
+  virtual void FindShortSuccessor(std::string* key) const override {
+    return;
+  }
+};
+}  // namespace
+
+static port::OnceType once = LEVELDB_ONCE_INIT;
+static const Comparator* uint64comp;
+
+static void InitModule() {
+  uint64comp = new Uint64ComparatorImpl;
+}
+
+const Comparator* Uint64Comparator() {
+  port::InitOnce(&once, InitModule);
+  return uint64comp;
+}
+
 }  // namespace test
 }  // namespace rocksdb

util/testutil.h

@@ -76,5 +76,12 @@ class PlainInternalKeyComparator : public InternalKeyComparator {
   }
 };
 
+// Returns a user key comparator that can be used for comparing two uint64_t
+// slices. Instead of comparing slices byte-wise, it compares all the 8 bytes
+// at once. Assumes same endian-ness is used though the database's lifetime.
+// Symantics of comparison would differ from Bytewise comparator in little
+// endian machines.
+extern const Comparator* Uint64Comparator();
+
 }  // namespace test
 }  // namespace rocksdb