[rts] Seek iterators don't need to return facts in lexicographic order

I want to relax the restriction that seek() must create an iterator that
returns facts in lexicographic order, because this gives more
flexibility in the storage layer. In particular we don't need to sort
all the keys, which can be problematic since keys can be large.

It turns out that we don't rely on the lexicographic ordering except for
saving and resuming query continuations, and this can be done by using
the current fact as the checkpoint instead.

We can remove the merge() operation on FactIterator and use only
append(), and this allows some optimisations when resuming iterators.

Author: simonmar

glean/rocksdb/database-impl.cpp

@@ -306,8 +306,7 @@ struct SeekIterator final : rts::FactIterator {
 } // namespace
 
 std::unique_ptr<rts::FactIterator>
-DatabaseImpl::seek(Pid type, folly::ByteRange start, size_t prefix_size) {
-  assert(prefix_size <= start.size());
+DatabaseImpl::seek(Pid type, folly::ByteRange prefix, std::optional<rts::Fact::Ref> restart) {
   if (count(type).high() == 0) {
     return std::make_unique<EmptyIterator>();
   }
@@ -316,22 +315,30 @@ DatabaseImpl::seek(Pid type, folly::ByteRange start, size_t prefix_size) {
   binary::Output out;
   out.fixed(type);
   const auto type_size = out.size();
-  out.put(start);
-  return std::make_unique<SeekIterator>(
-      out.bytes(), type_size + prefix_size, type, this);
+
+  if (restart.has_value()) {
+      auto fact = restart.value();
+      out.put(fact.clause.key());
+      return std::make_unique<SeekIterator>(
+          out.bytes(), type_size + prefix.size(), type, this);
+  } else {
+      out.put(prefix);
+      return std::make_unique<SeekIterator>(
+          out.bytes(), type_size + prefix.size(), type, this);
+  }
 }
 
 std::unique_ptr<rts::FactIterator> DatabaseImpl::seekWithinSection(
     Pid type,
-    folly::ByteRange start,
-    size_t prefix_size,
+    folly::ByteRange prefix,
     Id from,
-    Id upto) {
+    Id upto,
+    std::optional<rts::Fact::Ref> restart) {
   if (upto <= startingId() || firstFreeId() <= from) {
     return std::make_unique<EmptyIterator>();
   }
 
-  return Section(this, from, upto).seek(type, start, prefix_size);
+  return Section(this, from, upto).seek(type, prefix, restart);
 }
 
 namespace {

glean/rocksdb/database-impl.h

@@ -84,13 +84,13 @@ struct DatabaseImpl final : Database {
   std::unique_ptr<rts::FactIterator> enumerate(Id from, Id upto) override;
   std::unique_ptr<rts::FactIterator> enumerateBack(Id from, Id downto) override;
   std::unique_ptr<rts::FactIterator>
-  seek(Pid type, folly::ByteRange start, size_t prefix_size) override;
+  seek(Pid type, folly::ByteRange prefix, std::optional<rts::Fact::Ref> restart) override;
   std::unique_ptr<rts::FactIterator> seekWithinSection(
       Pid type,
-      folly::ByteRange start,
-      size_t prefix_size,
+      folly::ByteRange prefix,
       Id from,
-      Id upto) override;
+      Id upto,
+      std::optional<rts::Fact::Ref> restart) override;
 
   rts::UsetId getOwner(Id id) override;
 

glean/rts/benchmarking/factblock.cpp

@@ -91,7 +91,7 @@ bool FactBlock::lookupEachByKey(Lookup& lookup) const {
 
 bool FactBlock::seekToEach(Lookup& lookup) const {
   for (const auto fact : *this) {
-    auto iter = lookup.seek(fact.type, fact.key(), 0);
+    auto iter = lookup.seek(fact.type, fact.key());
     auto ref = iter->get(FactIterator::Demand::KeyOnly);
     if (ref.id != fact.id) {
       return false;

glean/rts/benchmarking/ffi.cpp

@@ -83,7 +83,7 @@ const char* glean_benchmarking_seek_count(
   return ffi::wrap([=] {
     size_t n = 0;
     for (size_t i = 0; i < pids_count; ++i) {
-      auto iter = lookup->seek(Pid::fromThrift(pids[i]), {}, 0);
+      auto iter = lookup->seek(Pid::fromThrift(pids[i]), {});
       while (auto ref = iter->get(FactIterator::Demand::KeyOnly)) {
         ++n;
         iter->next();

glean/rts/cache.cpp

@@ -227,18 +227,18 @@ std::unique_ptr<FactIterator> LookupCache::Anchor::enumerateBack(
 
 std::unique_ptr<FactIterator> LookupCache::Anchor::seek(
     Pid type,
-    folly::ByteRange start,
-    size_t prefix_size) {
-  return base->seek(type, start, prefix_size);
+    folly::ByteRange prefix,
+    std::optional<Fact::Ref> restart) {
+  return base->seek(type, prefix, restart);
 }
 
 std::unique_ptr<FactIterator> LookupCache::Anchor::seekWithinSection(
     Pid type,
-    folly::ByteRange start,
-    size_t prefix_size,
+    folly::ByteRange prefix,
     Id from,
-    Id to) {
-  return base->seekWithinSection(type, start, prefix_size, from, to);
+    Id to,
+    std::optional<Fact::Ref> restart) {
+  return base->seekWithinSection(type, prefix, from, to, restart);
 }
 
 void LookupCache::insert(Fact::unique_ptr owned) {

glean/rts/cache.h

@@ -151,14 +151,14 @@ struct LookupCache {
     }
 
     std::unique_ptr<FactIterator>
-    seek(Pid type, folly::ByteRange start, size_t prefix_size) override;
+    seek(Pid type, folly::ByteRange prefix, std::optional<Fact::Ref>) override;
 
     std::unique_ptr<FactIterator> seekWithinSection(
         Pid type,
-        folly::ByteRange start,
-        size_t prefix_size,
+        folly::ByteRange prefix,
         Id from,
-        Id to) override;
+        Id to,
+        std::optional<Fact::Ref>) override;
 
     UsetId getOwner(Id id) override {
       return base->getOwner(id);

glean/rts/factset.cpp

@@ -7,6 +7,7 @@
  */
 
 #include "glean/rts/factset.h"
+#include "glean/rts/binary.h"
 
 namespace facebook {
 namespace glean {
@@ -242,14 +243,28 @@ FactSet::seek(Pid type, folly::ByteRange start, size_t prefix_size) {
   }
 }
 
+std::unique_ptr<FactIterator>
+FactSet::seek(Pid type, folly::ByteRange prefix, std::optional<Fact::Ref> restart) {
+  if (restart.has_value()) {
+      auto id = restart.value().id;
+      if (id < facts.startingId() || id >= facts.startingId() + facts.size()) {
+          error("seek: restart Id not within set");
+      }
+      auto key = facts[id - facts.startingId()].clause.key();
+      return seek(type, key, prefix.size());
+  } else {
+      return seek(type, prefix, prefix.size());
+  }
+}
+
 std::unique_ptr<FactIterator> FactSet::seekWithinSection(
     Pid type,
-    folly::ByteRange start,
-    size_t prefix_size,
+    folly::ByteRange prefix,
     Id from,
-    Id to) {
+    Id to,
+    std::optional<Fact::Ref> restart) {
   if (from <= startingId() && firstFreeId() <= to) {
-    return seek(type, start, prefix_size);
+    return seek(type, prefix, restart);
   }
 
   // We have no use case for actually performing a bounded

glean/rts/factset.h

@@ -284,14 +284,14 @@ struct FactSet final : public Define {
   /// slow. The first call for each predicate will be especially slow as it will
   /// need to create an index.
   std::unique_ptr<FactIterator>
-  seek(Pid type, folly::ByteRange start, size_t prefix_size) override;
+  seek(Pid type, folly::ByteRange prefix, std::optional<Fact::Ref>) override;
 
   std::unique_ptr<FactIterator> seekWithinSection(
       Pid type,
-      folly::ByteRange start,
-      size_t prefix_size,
+      folly::ByteRange prefix,
       Id from,
-      Id to) override;
+      Id to,
+      std::optional<Fact::Ref>) override;
 
   UsetId getOwner(Id) override {
     return INVALID_USET;
@@ -350,6 +350,9 @@ struct FactSet final : public Define {
   /// don't do seeks on FactSets.
   struct Index;
   OnDemand<Index> index;
+
+  std::unique_ptr<FactIterator>
+  seek(Pid type, folly::ByteRange start, size_t prefix_size);
 };
 
 } // namespace rts

glean/rts/lookup.cpp

@@ -100,20 +100,6 @@ struct MergeIterator final : public FactIterator {
 
 } // namespace
 
-std::unique_ptr<FactIterator> FactIterator::merge(
-    std::unique_ptr<FactIterator> left,
-    std::unique_ptr<FactIterator> right,
-    size_t prefix_size) {
-  if (left->get(Demand::KeyOnly)) {
-    return right->get(Demand::KeyOnly)
-        ? std::make_unique<MergeIterator>(
-              std::move(left), std::move(right), prefix_size)
-        : std::move(left);
-  } else {
-    return right;
-  }
-}
-
 std::unique_ptr<FactIterator> Section::enumerate(Id from, Id upto) {
   if (upto <= lowBoundary() || highBoundary() <= from) {
     return std::make_unique<EmptyIterator>();
@@ -135,7 +121,7 @@ std::unique_ptr<FactIterator> Section::enumerateBack(Id from, Id downto) {
 }
 
 std::unique_ptr<FactIterator>
-Section::seek(Pid type, folly::ByteRange start, size_t prefix_size) {
+Section::seek(Pid type, folly::ByteRange prefix, std::optional<Fact::Ref> restart) {
   struct Iterator final : FactIterator {
     Iterator(std::unique_ptr<FactIterator> base, Id upto, Id from)
         : base_(std::move(base)), high_boundary_(upto), low_boundary_(from) {}
@@ -170,19 +156,19 @@ Section::seek(Pid type, folly::ByteRange start, size_t prefix_size) {
     Id low_boundary_;
   };
   return std::make_unique<Iterator>(
-      base()->seek(type, start, prefix_size), highBoundary(), lowBoundary());
+      base()->seek(type, prefix, restart), highBoundary(), lowBoundary());
 }
 
 std::unique_ptr<FactIterator> Section::seekWithinSection(
     Pid type,
-    folly::ByteRange start,
-    size_t prefix_size,
+    folly::ByteRange prefix,
     Id from,
-    Id upto) {
+    Id upto,
+    std::optional<Fact::Ref> restart) {
   if (from <= lowBoundary() && highBoundary() <= upto) {
-    return seek(type, start, prefix_size);
+    return seek(type, prefix, restart);
   } else {
-    return Section(base_, from, upto).seek(type, start, prefix_size);
+    return Section(base_, from, upto).seek(type, prefix, restart);
   }
 }
 

glean/rts/lookup.h

@@ -40,11 +40,6 @@ struct FactIterator {
   virtual std::optional<Id> lower_bound() = 0;
   virtual std::optional<Id> upper_bound() = 0;
 
-  static std::unique_ptr<FactIterator> merge(
-      std::unique_ptr<FactIterator> left,
-      std::unique_ptr<FactIterator> right,
-      size_t prefix_size);
-
   static std::unique_ptr<FactIterator> append(
       std::unique_ptr<FactIterator> left,
       std::unique_ptr<FactIterator> right);
@@ -127,23 +122,27 @@ struct Lookup {
       Id from = Id::invalid(),
       Id downto = Id::invalid()) = 0;
 
-  // Obtain a  prefix iterator for the given predicate. The iterator covers keys
-  // which begin with the first 'prefix_size' bytes of 'start', starting with
-  // the first key not lexicographically less than 'start', and produces facts
-  // in lexicographic order. In particular, setting 'prefix_size' to
-  // 'start.size()' iterates over all keys with the prefix 'start'.
+  // Obtain a  prefix iterator for the given predicate. The iterator will return
+  // each key with the given prefix exactly once, in no specified order.
+  //
+  // If `restart` is non-empty, the fact must be within the iterator's
+  // range; the iterator will be positioned at the given fact.  This
+  // can be used to resume a previous iterator.
   virtual std::unique_ptr<FactIterator>
-  seek(Pid type, folly::ByteRange start, size_t prefix_size) = 0;
+  seek(Pid type,
+       folly::ByteRange prefix,
+       std::optional<Fact::Ref> restart = {}) = 0;
 
   // Perform a seek on a section of the Lookup.
   // Results will have Ids within the range [from, utpto).
   // Facts can reference Ids outside of the specified range.
+  // See `seek()` for the meaning of `restart`.
   virtual std::unique_ptr<FactIterator> seekWithinSection(
       Pid type,
-      folly::ByteRange start,
-      size_t prefix_size,
+      folly::ByteRange prefix,
       Id from,
-      Id to) = 0;
+      Id to,
+      std::optional<Fact::Ref> restart = {}) = 0;
 
   virtual UsetId getOwner(Id id) = 0;
 };
@@ -183,12 +182,20 @@ struct EmptyLookup final : Lookup {
     return std::make_unique<EmptyIterator>();
   }
 
-  std::unique_ptr<FactIterator> seek(Pid, folly::ByteRange, size_t) override {
+  std::unique_ptr<FactIterator> seek(
+      Pid,
+      folly::ByteRange,
+      std::optional<Fact::Ref>) override {
     return std::make_unique<EmptyIterator>();
   }
 
   std::unique_ptr<FactIterator>
-  seekWithinSection(Pid, folly::ByteRange, size_t, Id, Id) override {
+  seekWithinSection(
+      Pid,
+      folly::ByteRange,
+      Id,
+      Id,
+      std::optional<Fact::Ref>) override {
     return std::make_unique<EmptyIterator>();
   }
 
@@ -239,14 +246,14 @@ struct Section : Lookup {
   std::unique_ptr<FactIterator> enumerateBack(Id from, Id downto) override;
 
   std::unique_ptr<FactIterator>
-  seek(Pid type, folly::ByteRange start, size_t prefix_size) override;
+  seek(Pid type, folly::ByteRange prefix, std::optional<Fact::Ref>) override;
 
   std::unique_ptr<FactIterator> seekWithinSection(
       Pid type,
-      folly::ByteRange start,
-      size_t prefix_size,
+      folly::ByteRange prefix,
       Id from,
-      Id to) override;
+      Id to,
+      std::optional<Fact::Ref>) override;
 
   UsetId getOwner(Id id) override {
     return isWithinBounds(id) ? base()->getOwner(id) : INVALID_USET;