@nxtedition/rocksdb 13.1.4 → 13.2.0

package/deps/rocksdb/rocksdb/memtable/inlineskiplist.h

@@ -141,8 +141,9 @@ class InlineSkipList {
   // Returns true iff an entry that compares equal to key is in the list.
   bool Contains(const char* key) const;
 
-  // Return estimated number of entries smaller than `key`.
-  uint64_t EstimateCount(const char* key) const;
+  // Return estimated number of entries from `start_ikey` to `end_ikey`.
+  uint64_t ApproximateNumEntries(const Slice& start_ikey,
+                                 const Slice& end_ikey) const;
 
   // Validate correctness of the skip-list.
   void TEST_Validate() const;
@@ -673,31 +674,88 @@ InlineSkipList<Comparator>::FindRandomEntry() const {
 }
 
 template <class Comparator>
-uint64_t InlineSkipList<Comparator>::EstimateCount(const char* key) const {
+uint64_t InlineSkipList<Comparator>::ApproximateNumEntries(
+    const Slice& start_ikey, const Slice& end_ikey) const {
+  // The number of entries at a given level for the given range, in terms of
+  // the actual number of entries in that range (level 0), follows a binomial
+  // distribution, which is very well approximated by the Poisson distribution.
+  // That has stddev sqrt(x) where x is the expected number of entries (mean)
+  // at this level, and the best predictor of x is the number of observed
+  // entries (at this level). To predict the number of entries on level 0 we use
+  // x * kBranchinng ^ level. From the standard deviation, the P99+ relative
+  // error is roughly 3 * sqrt(x) / x. Thus, a reasonable approach would be to
+  // find the smallest level with at least some moderate constant number entries
+  // in range. E.g. with at least ~40 entries, we expect P99+ relative error
+  // (approximation accuracy) of ~ 50% = 3 * sqrt(40) / 40; P95 error of
+  // ~30%; P75 error of < 20%.
+  //
+  // However, there are two issues with this approach, and an observation:
+  // * Pointer chasing on the larger (bottom) levels is much slower because of
+  // cache hierarchy effects, so when the result is smaller, getting the result
+  // will be substantially slower, despite traversing a similar number of
+  // entries. (We could be clever about pipelining our pointer chasing but
+  // that's complicated.)
+  // * The larger (bottom) levels also have lower variance because there's a
+  // chance (or certainty) that we reach level 0 and return the exact answer.
+  // * For applications in query planning, we can also tolerate more variance on
+  // small results because the impact of misestimating is likely smaller.
+  //
+  // These factors point us to an approach in which we have a higher minimum
+  // threshold number of samples for higher levels and lower for lower levels
+  // (see sufficient_samples below). This seems to yield roughly consistent
+  // relative error (stddev around 20%, less for large results) and roughly
+  // consistent query time around the time of two memtable point queries.
+  //
+  // Engineering observation: it is tempting to think that taking into account
+  // what we already found in how many entries occur on higher levels, not just
+  // the first iterated level with a sufficient number of samples, would yield
+  // a more accurate estimate. But that doesn't work because of the particular
+  // correlations and independences of the data: each level higher is just an
+  // independently probabilistic filtering of the level below it. That
+  // filtering from level l to l+1 has no more information about levels
+  // 0 .. l-1 than we can get from level l. The structure of RandomHeight() is
+  // a clue to these correlations and independences.
+
+  Node* lb = head_;
+  Node* ub = nullptr;
   uint64_t count = 0;
-
-  Node* x = head_;
-  int level = GetMaxHeight() - 1;
-  const DecodedKey key_decoded = compare_.decode_key(key);
-  while (true) {
-    assert(x == head_ || compare_(x->Key(), key_decoded) < 0);
-    Node* next = x->Next(level);
-    if (next != nullptr) {
-      PREFETCH(next->Next(level), 0, 1);
+  for (int level = GetMaxHeight() - 1; level >= 0; level--) {
+    auto sufficient_samples = static_cast<uint64_t>(level) * kBranching_ + 10U;
+    if (count >= sufficient_samples) {
+      // No more counting; apply powers of kBranching and avoid floating point
+      count *= kBranching_;
+      continue;
     }
-    if (next == nullptr || compare_(next->Key(), key_decoded) >= 0) {
-      if (level == 0) {
-        return count;
-      } else {
-        // Switch to next list
-        count *= kBranching_;
-        level--;
+    count = 0;
+    Node* next;
+    // Get a more precise lower bound (for start key)
+    for (;;) {
+      next = lb->Next(level);
+      if (next == ub) {
+        break;
+      }
+      assert(next != nullptr);
+      if (compare_(next->Key(), start_ikey) >= 0) {
+        break;
+      }
+      lb = next;
+    }
+    // Count entries on this level until upper bound (for end key)
+    for (;;) {
+      if (next == ub) {
+        break;
+      }
+      assert(next != nullptr);
+      if (compare_(next->Key(), end_ikey) >= 0) {
+        // Save refined upper bound to potentially save key comparison
+        ub = next;
+        break;
       }
-    } else {
-      x = next;
       count++;
+      next = next->Next(level);
     }
   }
+  return count;
 }
 
 template <class Comparator>

package/deps/rocksdb/rocksdb/memtable/skiplist.h

@@ -64,8 +64,9 @@ class SkipList {
   // Returns true iff an entry that compares equal to key is in the list.
   bool Contains(const Key& key) const;
 
-  // Return estimated number of entries smaller than `key`.
-  uint64_t EstimateCount(const Key& key) const;
+  // Return estimated number of entries from `start_ikey` to `end_ikey`.
+  uint64_t ApproximateNumEntries(const Slice& start_ikey,
+                                 const Slice& end_ikey) const;
 
   // Iteration over the contents of a skip list
   class Iterator {
@@ -383,27 +384,49 @@ typename SkipList<Key, Comparator>::Node* SkipList<Key, Comparator>::FindLast()
 }
 
 template <typename Key, class Comparator>
-uint64_t SkipList<Key, Comparator>::EstimateCount(const Key& key) const {
+uint64_t SkipList<Key, Comparator>::ApproximateNumEntries(
+    const Slice& start_ikey, const Slice& end_ikey) const {
+  // See InlineSkipList<Comparator>::ApproximateNumEntries() (copy-paste)
+  Node* lb = head_;
+  Node* ub = nullptr;
   uint64_t count = 0;
-
-  Node* x = head_;
-  int level = GetMaxHeight() - 1;
-  while (true) {
-    assert(x == head_ || compare_(x->key, key) < 0);
-    Node* next = x->Next(level);
-    if (next == nullptr || compare_(next->key, key) >= 0) {
-      if (level == 0) {
-        return count;
-      } else {
-        // Switch to next list
-        count *= kBranching_;
-        level--;
+  for (int level = GetMaxHeight() - 1; level >= 0; level--) {
+    auto sufficient_samples = static_cast<uint64_t>(level) * kBranching_ + 10U;
+    if (count >= sufficient_samples) {
+      // No more counting; apply powers of kBranching and avoid floating point
+      count *= kBranching_;
+      continue;
+    }
+    count = 0;
+    Node* next;
+    // Get a more precise lower bound (for start key)
+    for (;;) {
+      next = lb->Next(level);
+      if (next == ub) {
+        break;
+      }
+      assert(next != nullptr);
+      if (compare_(next->Key(), start_ikey) >= 0) {
+        break;
+      }
+      lb = next;
+    }
+    // Count entries on this level until upper bound (for end key)
+    for (;;) {
+      if (next == ub) {
+        break;
+      }
+      assert(next != nullptr);
+      if (compare_(next->Key(), end_ikey) >= 0) {
+        // Save refined upper bound to potentially save key comparison
+        ub = next;
+        break;
       }
-    } else {
-      x = next;
       count++;
+      next = next->Next(level);
     }
   }
+  return count;
 }
 
 template <typename Key, class Comparator>

package/deps/rocksdb/rocksdb/memtable/skiplistrep.cc

@@ -108,11 +108,7 @@ class SkipListRep : public MemTableRep {
 
   uint64_t ApproximateNumEntries(const Slice& start_ikey,
                                  const Slice& end_ikey) override {
-    std::string tmp;
-    uint64_t start_count =
-        skip_list_.EstimateCount(EncodeKey(&tmp, start_ikey));
-    uint64_t end_count = skip_list_.EstimateCount(EncodeKey(&tmp, end_ikey));
-    return (end_count >= start_count) ? (end_count - start_count) : 0;
+    return skip_list_.ApproximateNumEntries(start_ikey, end_ikey);
   }
 
   void UniqueRandomSample(const uint64_t num_entries,

package/deps/rocksdb/rocksdb/memtable/wbwi_memtable.cc

@@ -0,0 +1,169 @@
+//  Copyright (c) Meta Platforms, Inc. and affiliates.
+//  This source code is licensed under both the GPLv2 (found in the
+//  COPYING file in the root directory) and Apache 2.0 License
+//  (found in the LICENSE.Apache file in the root directory).
+
+#include "memtable/wbwi_memtable.h"
+
+#include "db/memtable.h"
+
+namespace ROCKSDB_NAMESPACE {
+
+const std::unordered_map<WriteType, ValueType>
+    WBWIMemTableIterator::WriteTypeToValueTypeMap = {
+        {kPutRecord, kTypeValue},
+        {kMergeRecord, kTypeMerge},
+        {kDeleteRecord, kTypeDeletion},
+        {kSingleDeleteRecord, kTypeSingleDeletion},
+        {kDeleteRangeRecord, kTypeRangeDeletion},
+        {kPutEntityRecord, kTypeWideColumnEntity},
+        // Only the above record types are added to WBWI.
+        // kLogDataRecord, kXIDRecord, kUnknownRecord
+};
+
+InternalIterator* WBWIMemTable::NewIterator(
+    const ReadOptions&, UnownedPtr<const SeqnoToTimeMapping>, Arena* arena,
+    const SliceTransform* /* prefix_extractor */, bool for_flush) {
+  // Ingested WBWIMemTable should have an assigned seqno
+  assert(assigned_seqno_.upper_bound != kMaxSequenceNumber);
+  assert(assigned_seqno_.lower_bound != kMaxSequenceNumber);
+  assert(arena);
+  auto mem = arena->AllocateAligned(sizeof(WBWIMemTableIterator));
+  return new (mem) WBWIMemTableIterator(
+      std::unique_ptr<WBWIIterator>(wbwi_->NewIterator(cf_id_)),
+      assigned_seqno_, comparator_, for_flush);
+}
+
+inline InternalIterator* WBWIMemTable::NewIterator() const {
+  assert(assigned_seqno_.upper_bound != kMaxSequenceNumber);
+  assert(assigned_seqno_.lower_bound != kMaxSequenceNumber);
+  return new WBWIMemTableIterator(
+      std::unique_ptr<WBWIIterator>(wbwi_->NewIterator(cf_id_)),
+      assigned_seqno_, comparator_, /*for_flush=*/false);
+}
+
+bool WBWIMemTable::Get(const LookupKey& key, std::string* value,
+                       PinnableWideColumns* columns, std::string* timestamp,
+                       Status* s, MergeContext* merge_context,
+                       SequenceNumber* max_covering_tombstone_seq,
+                       SequenceNumber* out_seq, const ReadOptions&,
+                       bool immutable_memtable, ReadCallback* callback,
+                       bool* is_blob_index, bool do_merge) {
+  (void)immutable_memtable;
+  (void)timestamp;
+  (void)columns;
+  assert(immutable_memtable);
+  assert(!timestamp);  // TODO: support UDT
+  assert(!columns);    // TODO: support WideColumn
+  assert(assigned_seqno_.upper_bound != kMaxSequenceNumber);
+  assert(assigned_seqno_.lower_bound != kMaxSequenceNumber);
+  // WBWI does not support DeleteRange yet.
+  assert(!wbwi_->GetWriteBatch()->HasDeleteRange());
+
+  [[maybe_unused]] SequenceNumber read_seq =
+      GetInternalKeySeqno(key.internal_key());
+  std::unique_ptr<InternalIterator> iter{NewIterator()};
+  iter->Seek(key.internal_key());
+  const Slice lookup_user_key = key.user_key();
+
+  while (iter->Valid() && comparator_->EqualWithoutTimestamp(
+                              ExtractUserKey(iter->key()), lookup_user_key)) {
+    uint64_t tag = ExtractInternalKeyFooter(iter->key());
+    ValueType type;
+    SequenceNumber seq;
+    UnPackSequenceAndType(tag, &seq, &type);
+    // Unsupported operations.
+    assert(type != kTypeBlobIndex);
+    assert(type != kTypeWideColumnEntity);
+    assert(type != kTypeValuePreferredSeqno);
+    assert(type != kTypeDeletionWithTimestamp);
+    assert(type != kTypeMerge);
+    if (!callback || callback->IsVisible(seq)) {
+      if (*out_seq == kMaxSequenceNumber) {
+        *out_seq = std::max(seq, *max_covering_tombstone_seq);
+      }
+      if (*max_covering_tombstone_seq > seq) {
+        type = kTypeRangeDeletion;
+      }
+      switch (type) {
+        case kTypeValue: {
+          HandleTypeValue(lookup_user_key, iter->value(), iter->IsValuePinned(),
+                          do_merge, s->IsMergeInProgress(), merge_context,
+                          moptions_.merge_operator, clock_,
+                          moptions_.statistics, moptions_.info_log, s, value,
+                          columns, is_blob_index);
+          assert(seq <= read_seq);
+          return /*found_final_value=*/true;
+        }
+        case kTypeDeletion:
+        case kTypeSingleDeletion:
+        case kTypeRangeDeletion: {
+          HandleTypeDeletion(lookup_user_key, s->IsMergeInProgress(),
+                             merge_context, moptions_.merge_operator, clock_,
+                             moptions_.statistics, moptions_.info_log, s, value,
+                             columns);
+          assert(seq <= read_seq);
+          return /*found_final_value=*/true;
+        }
+        default: {
+          std::string msg("Unrecognized or unsupported value type: " +
+                          std::to_string(static_cast<int>(type)) + ". ");
+          msg.append("User key: " +
+                     ExtractUserKey(iter->key()).ToString(/*hex=*/true) + ". ");
+          msg.append("seq: " + std::to_string(seq) + ".");
+          *s = Status::Corruption(msg.c_str());
+          return /*found_final_value=*/true;
+        }
+      }
+    }
+    // Current key not visible or we read a merge key
+    assert(s->IsMergeInProgress() || (callback && !callback->IsVisible(seq)));
+    iter->Next();
+  }
+  if (!iter->status().ok() &&
+      (s->ok() || s->IsMergeInProgress() || s->IsNotFound())) {
+    *s = iter->status();
+    // stop further look up
+    return true;
+  }
+  return /*found_final_value=*/false;
+}
+
+void WBWIMemTable::MultiGet(const ReadOptions& read_options,
+                            MultiGetRange* range, ReadCallback* callback,
+                            bool immutable_memtable) {
+  (void)immutable_memtable;
+  // Should only be used as immutable memtable.
+  assert(immutable_memtable);
+  // TODO: reuse the InternalIterator created in Get().
+  for (auto iter = range->begin(); iter != range->end(); ++iter) {
+    SequenceNumber dummy_seq = 0;
+    bool found_final_value =
+        Get(*iter->lkey, iter->value ? iter->value->GetSelf() : nullptr,
+            iter->columns, iter->timestamp, iter->s, &(iter->merge_context),
+            &(iter->max_covering_tombstone_seq), &dummy_seq, read_options, true,
+            callback, nullptr, true);
+    if (found_final_value) {
+      if (iter->s->ok() || iter->s->IsNotFound()) {
+        if (iter->value) {
+          iter->value->PinSelf();
+          range->AddValueSize(iter->value->size());
+        } else {
+          assert(iter->columns);
+          range->AddValueSize(iter->columns->serialized_size());
+        }
+      }
+      range->MarkKeyDone(iter);
+      if (range->GetValueSize() > read_options.value_size_soft_limit) {
+        // Set all remaining keys in range to Abort
+        for (auto range_iter = range->begin(); range_iter != range->end();
+             ++range_iter) {
+          range->MarkKeyDone(range_iter);
+          *(range_iter->s) = Status::Aborted();
+        }
+        break;
+      }
+    }
+  }
+}
+}  // namespace ROCKSDB_NAMESPACE