BloomBitsReader v1

src/storage/v2/disk/rocksdb_storage.cpp

@@ -10,7 +10,11 @@
 // licenses/APL.txt.
 
 #include "rocksdb_storage.hpp"
+#include <rocksdb/slice.h>
 #include <string_view>
+#include "utils/algorithm.hpp"
+#include "utils/disk_utils.hpp"
+#include "utils/exceptions.hpp"
 #include "utils/rocksdb_serialization.hpp"
 
 namespace memgraph::storage {
@@ -18,6 +22,7 @@ namespace memgraph::storage {
 namespace {
 
 inline rocksdb::Slice StripTimestampFromUserKey(const rocksdb::Slice &user_key, size_t ts_sz) {
+  // spdlog::debug("StripTimestampFromUserKey: {}", user_key.ToString());
   rocksdb::Slice ret = user_key;
   ret.remove_suffix(ts_sz);
   return ret;
@@ -30,10 +35,14 @@ inline rocksdb::Slice ExtractTimestampFromUserKey(const rocksdb::Slice &user_key
 }
 
 // Extracts global id from user key. User key must be without timestamp.
-std::string_view ExtractGidFromUserKey(const rocksdb::Slice &key) {
-  assert(key.size() >= 2);
-  auto keyStrView = key.ToStringView();
-  return keyStrView.substr(keyStrView.find_last_of('|') + 1);
+inline std::string_view ExtractGidFromUserKey(const rocksdb::Slice &key) {
+  std::string_view keyStrView = key.ToStringView();
+  // spdlog::debug("ExtractGidFromKey: {}", keyStrView);
+  if (utils::Contains(keyStrView, '|')) {
+    assert(key.size() >= 2);
+    return keyStrView.substr(keyStrView.find_last_of('|') + 1);
+  }
+  return keyStrView;
 }
 
 }  // namespace
@@ -43,26 +52,66 @@ ComparatorWithU64TsImpl::ComparatorWithU64TsImpl()
   assert(cmp_without_ts_->timestamp_size() == 0);
 }
 
+/// TODO: try to write somehow unit test for this
 int ComparatorWithU64TsImpl::Compare(const rocksdb::Slice &a, const rocksdb::Slice &b) const {
-  int ret = CompareWithoutTimestamp(a, b);
-  if (ret != 0) {
-    return ret;
+  std::string a_str = a.ToString();
+  std::string b_str = b.ToString();
+  // spdlog::debug("Received a: {} b: {} in compare method", a_str, b_str);
+  uint32_t num_separators_a = std::count(a_str.begin(), a_str.end(), '|');
+  uint32_t num_separators_b = std::count(b_str.begin(), b_str.end(), '|');
+  utils::RocksDBType a_key_type = utils::GetRocksDBKeyType(num_separators_a);
+  utils::RocksDBType b_key_type = utils::GetRocksDBKeyType(num_separators_b);
+  if (utils::ComparingVertexWithVertex(a_key_type, b_key_type) ||
+      utils::ComparingEdgeWithEdge(a_key_type, b_key_type)) {
+    int ret = CompareWithoutTimestamp(a, b);
+    if (ret != 0) {
+      return ret;
+    }
+    // Compare timestamp.
+    // For the same user key with different timestamps, larger (newer) timestamp
+    // comes first.
+    return CompareTimestamp(ExtractTimestampFromUserKey(b), ExtractTimestampFromUserKey(a));
   }
-  // Compare timestamp.
-  // For the same user key with different timestamps, larger (newer) timestamp
-  // comes first.
-  return CompareTimestamp(ExtractTimestampFromUserKey(b), ExtractTimestampFromUserKey(a));
+  if (utils::ComparingEdgeWithGID(a_key_type, b_key_type)) {
+    return CompareEdgeWithGidForPrefixSearch(a_str, b_str);
+  }
+  throw utils::BasicException(
+      "Cannot handle specific compare use-case when one of the keys are neither vertex nor edge.");
 }
 
+int ComparatorWithU64TsImpl::CompareEdgeWithGidForPrefixSearch(const std::string_view edge,
+                                                               const std::string_view source_vertex_gid) const {
+  // spdlog::debug("Compare edge with gid for prefix search START");
+  rocksdb::Slice stripped_src_vertex_gid = StripTimestampFromUserKey(source_vertex_gid, timestamp_size());
+  // rocksdb::Slice stripped_src_vertex_gid = source_vertex_gid;
+  std::string_view edge_source_vertex_gid = edge.substr(0, edge.find('|'));
+  // spdlog::debug("Edge: {} Edge source vertex gid: {} Size: {} Source vertex gid: {} Size: {}", edge,
+  //               edge_source_vertex_gid, edge_source_vertex_gid.size(), stripped_src_vertex_gid.ToString(),
+  //               stripped_src_vertex_gid.ToString().size());
+  int cmp_res = cmp_without_ts_->Compare(edge_source_vertex_gid, stripped_src_vertex_gid);
+  // spdlog::debug("Compare result: {}", cmp_res);
+  return cmp_res;
+}
+
+/// TODO: entered CompareWithoutTimestamp with GID. Handle something like in a compare function.
 int ComparatorWithU64TsImpl::CompareWithoutTimestamp(const rocksdb::Slice &a, bool a_has_ts, const rocksdb::Slice &b,
                                                      bool b_has_ts) const {
   const size_t ts_sz = timestamp_size();
+  // spdlog::debug("Timestamp size: {}", ts_sz);
+  // spdlog::debug("a_has_ts: {} b_has_ts: {}", a_has_ts, b_has_ts);
+  // spdlog::debug("a size: {} b size: {}", a.size(), b.size());
   assert(!a_has_ts || a.size() >= ts_sz);
   assert(!b_has_ts || b.size() >= ts_sz);
+  // spdlog::debug("a key: {}", a.ToString());
+  // spdlog::debug("b key: {}", b.ToString());
   rocksdb::Slice lhsUserKey = a_has_ts ? StripTimestampFromUserKey(a, ts_sz) : a;
   rocksdb::Slice rhsUserKey = b_has_ts ? StripTimestampFromUserKey(b, ts_sz) : b;
+  // spdlog::debug("lhsUserKey: {}", lhsUserKey.ToString());
+  // spdlog::debug("rhsUserKey: {}", rhsUserKey.ToString());
   rocksdb::Slice lhsGid = ExtractGidFromUserKey(lhsUserKey);
   rocksdb::Slice rhsGid = ExtractGidFromUserKey(rhsUserKey);
+  // spdlog::debug("lhsGid: {}", lhsGid.ToString());
+  // spdlog::debug("rhsGid: {}", rhsGid.ToString());
   return cmp_without_ts_->Compare(lhsGid, rhsGid);
 }
 

src/storage/v2/disk/rocksdb_storage.hpp

@@ -13,14 +13,23 @@
 
 #include <rocksdb/comparator.h>
 #include <rocksdb/db.h>
+#include <rocksdb/filter_policy.h>
 #include <rocksdb/iterator.h>
 #include <rocksdb/options.h>
+#include <rocksdb/slice.h>
+#include <rocksdb/slice_transform.h>
 #include <rocksdb/status.h>
+#include <rocksdb/table/block_based/filter_policy_internal.h>
+#include <rocksdb/util/bloom_impl.h>
+#include <rocksdb/util/hash.h>
 #include <rocksdb/utilities/transaction_db.h>
+#include <string_view>
 
 #include "storage/v2/id_types.hpp"
 #include "storage/v2/property_store.hpp"
+#include "utils/algorithm.hpp"
 #include "utils/logging.hpp"
+#include "utils/string.hpp"
 
 namespace memgraph::storage {
 
@@ -84,6 +93,97 @@ class ComparatorWithU64TsImpl : public rocksdb::Comparator {
 
  private:
   const Comparator *cmp_without_ts_{nullptr};
+
+  int CompareEdgeWithGidForPrefixSearch(std::string_view edge, std::string_view source_vertex_gid) const;
+};
+
+class VerticalLinePrefixTransform : public rocksdb::SliceTransform {
+ private:
+  std::string id_{"vertical_line_separator_prefix_transform"};
+
+ public:
+  explicit VerticalLinePrefixTransform() {}
+
+  static const char *kClassName() { return "memgraph.VerticalLine"; }
+  static const char *kNickName() { return "vertical_line"; }
+  const char *Name() const override { return kClassName(); }
+  const char *NickName() const override { return kNickName(); }
+
+  bool IsInstanceOf(const std::string &name) const override {
+    if (name == id_) {
+      return true;
+    }
+    return rocksdb::SliceTransform::IsInstanceOf(name);
+  }
+
+  std::string GetId() const override { return id_; }
+
+  rocksdb::Slice Transform(const rocksdb::Slice &src) const override {
+    // spdlog::debug("Received request for transform: {} {}", src.ToString(), src.ToString().size());
+    const std::string src_str = src.ToString();
+    if (utils::Contains(src.ToString(), '|')) {
+      assert(InDomain(src));
+      auto res = rocksdb::Slice(src.data(), src_str.find('|'));
+      // spdlog::debug("Transform result: {} from: {}", res.ToString(), src_str);
+      return res;
+    }
+    return src;
+  }
+
+  bool InDomain(const rocksdb::Slice &src) const override { return (utils::Contains(src.ToString(), '|')); }
+
+  // deprecated and implemented here just for backwards compatibility
+  bool InRange(const rocksdb::Slice & /*dst*/) const override { return true; }
+
+  bool FullLengthEnabled(size_t * /*len*/) const override { return false; }
+
+  bool SameResultWhenAppended(const rocksdb::Slice &prefix) const override { return InDomain(prefix); }
+};
+
+class MemgraphBloomBitsReader : public rocksdb::BuiltinFilterBitsReader {
+ public:
+  MemgraphBloomBitsReader(const char *data, int num_probes, uint32_t len_bytes)
+      : data_(data), num_probes_(num_probes), len_bytes_(len_bytes) {}
+
+  // No Copy allowed
+  MemgraphBloomBitsReader(const MemgraphBloomBitsReader &) = delete;
+  void operator=(const MemgraphBloomBitsReader &) = delete;
+  MemgraphBloomBitsReader(const MemgraphBloomBitsReader &&) = delete;
+  void operator=(const MemgraphBloomBitsReader &&) = delete;
+
+  ~MemgraphBloomBitsReader() override {}
+
+  bool MayMatch(const rocksdb::Slice &key) override {
+    uint64_t h = rocksdb::GetSliceHash64(key);
+    uint32_t byte_offset;
+    rocksdb::FastLocalBloomImpl::PrepareHash(rocksdb::Lower32of64(h), len_bytes_, data_,
+                                             /*out*/ &byte_offset);
+    return rocksdb::FastLocalBloomImpl::HashMayMatchPrepared(rocksdb::Upper32of64(h), num_probes_, data_ + byte_offset);
+  }
+
+  void MayMatch(int num_keys, rocksdb::Slice **keys, bool *may_match) override {
+    std::array<uint32_t, 32> hashes;
+    std::array<uint32_t, 32> byte_offsets;
+    for (int i = 0; i < num_keys; ++i) {
+      uint64_t h = rocksdb::GetSliceHash64(*keys[i]);
+      rocksdb::FastLocalBloomImpl::PrepareHash(rocksdb::Lower32of64(h), len_bytes_, data_,
+                                               /*out*/ &byte_offsets[i]);
+      hashes[i] = rocksdb::Upper32of64(h);
+    }
+    for (int i = 0; i < num_keys; ++i) {
+      may_match[i] = rocksdb::FastLocalBloomImpl::HashMayMatchPrepared(hashes[i], num_probes_, data_ + byte_offsets[i]);
+    }
+  }
+
+  bool HashMayMatch(const uint64_t h) override {
+    return rocksdb::FastLocalBloomImpl::HashMayMatch(rocksdb::Lower32of64(h), rocksdb::Upper32of64(h), len_bytes_,
+                                                     num_probes_, data_);
+  }
+
+ private:
+  const char *data_;
+  const int num_probes_;
+  const uint32_t len_bytes_;
 };
 
 }  // namespace memgraph::storage

src/storage/v2/disk/storage.cpp

@@ -16,7 +16,11 @@
 
 #include <rocksdb/comparator.h>
 #include <rocksdb/db.h>
+#include <rocksdb/filter_policy.h>
+#include <rocksdb/memtablerep.h>
 #include <rocksdb/slice.h>
+#include <rocksdb/slice_transform.h>
+#include <rocksdb/table.h>
 
 #include <rocksdb/options.h>
 #include <rocksdb/utilities/transaction.h>
@@ -241,6 +245,12 @@ DiskStorage::DiskStorage(Config config)
   LoadIndexInfoIfExists();
   LoadConstraintsInfoIfExists();
   kvstore_->options_.create_if_missing = true;
+  rocksdb::BlockBasedTableOptions table_options;
+  // TODO: have to implement our own FilterPolicy
+  // table_options.filter_policy.reset(rocksdb::NewBloomFilterPolicy(10, false));
+  table_options.filter_policy.reset(rocksdb::NewRibbonFilterPolicy(10, false));
+  kvstore_->options_.table_factory.reset(rocksdb::NewBlockBasedTableFactory(table_options));
+  kvstore_->options_.prefix_extractor.reset(new VerticalLinePrefixTransform());
   kvstore_->options_.comparator = new ComparatorWithU64TsImpl();
   kvstore_->options_.compression = rocksdb::kNoCompression;
   kvstore_->options_.wal_recovery_mode = rocksdb::WALRecoveryMode::kPointInTimeRecovery;
@@ -382,6 +392,7 @@ std::optional<EdgeAccessor> DiskStorage::DiskAccessor::DeserializeEdge(const roc
 }
 
 VerticesIterable DiskStorage::DiskAccessor::Vertices(View view) {
+  // spdlog::debug("Scanning all vertices");
   auto *disk_storage = static_cast<DiskStorage *>(storage_);
   rocksdb::ReadOptions ro;
   std::string strTs = utils::StringTimestamp(transaction_.start_timestamp);
@@ -886,6 +897,7 @@ void DiskStorage::DiskAccessor::PrefetchEdges(const auto &prefetch_edge_filter)
 }
 
 void DiskStorage::DiskAccessor::PrefetchInEdges(const VertexAccessor &vertex_acc) {
+  spdlog::debug("Prefetch input edges");
   PrefetchEdges([&vertex_acc](const std::vector<std::string> &disk_edge_parts) -> bool {
     auto disk_vertex_in_edge_gid = disk_edge_parts[1];
     auto edge_gid = disk_edge_parts[4];
@@ -902,21 +914,46 @@ void DiskStorage::DiskAccessor::PrefetchInEdges(const VertexAccessor &vertex_acc
   });
 }
 
+// void DiskStorage::DiskAccessor::PrefetchOutEdgesTmp(const VertexAccessor &vertex_acc) {
+//   spdlog::debug("Prefetch output edges");
+//   PrefetchEdges([&vertex_acc](const std::vector<std::string> &disk_edge_parts) -> bool {
+//     auto disk_vertex_out_edge_gid = disk_edge_parts[0];
+//     auto edge_gid = disk_edge_parts[4];
+//     auto out_edges_res = vertex_acc.OutEdges(storage::View::NEW);
+//     if (out_edges_res.HasValue()) {
+//       for (const auto &edge_acc : out_edges_res.GetValue()) {
+//         if (utils::SerializeIdType(edge_acc.Gid()) == edge_gid) {
+//           // We already inserted this edge into the vertex's out_edges list.
+//           return false;
+//         }
+//       }
+//     }
+//     return disk_vertex_out_edge_gid == utils::SerializeIdType(vertex_acc.Gid());
+//   });
+// }
+
 void DiskStorage::DiskAccessor::PrefetchOutEdges(const VertexAccessor &vertex_acc) {
-  PrefetchEdges([&vertex_acc](const std::vector<std::string> &disk_edge_parts) -> bool {
-    auto disk_vertex_out_edge_gid = disk_edge_parts[0];
-    auto edge_gid = disk_edge_parts[4];
-    auto out_edges_res = vertex_acc.OutEdges(storage::View::NEW);
-    if (out_edges_res.HasValue()) {
-      for (const auto &edge_acc : out_edges_res.GetValue()) {
-        if (utils::SerializeIdType(edge_acc.Gid()) == edge_gid) {
-          // We already inserted this edge into the vertex's out_edges list.
-          return false;
-        }
-      }
-    }
-    return disk_vertex_out_edge_gid == utils::SerializeIdType(vertex_acc.Gid());
-  });
+  spdlog::debug("Prefetch out edges");
+  rocksdb::ReadOptions read_opts;
+  read_opts.total_order_seek = true;
+  read_opts.prefix_same_as_start = true;
+  auto strTs = utils::StringTimestamp(transaction_.start_timestamp);
+  rocksdb::Slice ts(strTs);
+  read_opts.timestamp = &ts;
+  auto *disk_storage = static_cast<DiskStorage *>(storage_);
+  auto it = std::unique_ptr<rocksdb::Iterator>(
+      disk_transaction_->GetIterator(read_opts, disk_storage->kvstore_->edge_chandle));
+  auto gid_str = std::to_string(vertex_acc.Gid().AsUint());
+  spdlog::debug("Gid string: {}", gid_str);
+  it->Seek(gid_str);
+  while (it->Valid()) {
+    spdlog::debug("In loop execution");
+    auto key_str = it->key().ToStringView();
+    if (!key_str.starts_with(gid_str)) break;
+    spdlog::debug("Key str: {} gid str: {}", key_str, gid_str);
+    DeserializeEdge(key_str, it->value());
+    it->Next();
+  }
 }
 
 Result<EdgeAccessor> DiskStorage::DiskAccessor::CreateEdge(const VertexAccessor *from, const VertexAccessor *to,

src/utils/disk_utils.hpp

@@ -15,6 +15,8 @@
 
 namespace memgraph::utils {
 
+enum RocksDBType { VERTEX, EDGE, GID };
+
 inline std::optional<std::string> GetOldDiskKeyOrNull(storage::Delta *head) {
   while (head->next != nullptr) {
     head = head->next;
@@ -25,4 +27,26 @@ inline std::optional<std::string> GetOldDiskKeyOrNull(storage::Delta *head) {
   return std::nullopt;
 }
 
+inline RocksDBType GetRocksDBKeyType(uint32_t num_separators) {
+  if (num_separators == 0) {
+    return RocksDBType::GID;
+  }
+  if (num_separators == 4) {
+    return RocksDBType::EDGE;
+  }
+  return RocksDBType::VERTEX;
+}
+
+inline bool ComparingVertexWithVertex(RocksDBType type_a, RocksDBType type_b) {
+  return type_a == RocksDBType::VERTEX && type_b == RocksDBType::VERTEX;
+}
+
+inline bool ComparingEdgeWithEdge(RocksDBType type_a, RocksDBType type_b) {
+  return type_a == RocksDBType::EDGE && type_b == RocksDBType::EDGE;
+}
+
+inline bool ComparingEdgeWithGID(RocksDBType type_a, RocksDBType type_b) {
+  return type_a == RocksDBType::EDGE && type_b == RocksDBType::GID;
+}
+
 }  // namespace memgraph::utils

tests/unit/storage_v2_edge_ondisk.cpp

@@ -4247,9 +4247,6 @@ TEST_P(StorageEdgeTest, VertexDetachDeleteSingleAbort) {
     auto vertex_to = acc->FindVertex(gid_to, memgraph::storage::View::NEW);
     ASSERT_FALSE(vertex_from);
     ASSERT_TRUE(vertex_to);
-    // We prefetch edges implicitly when go thorough query Accessor
-    acc->PrefetchOutEdges(*vertex_from);
-    acc->PrefetchInEdges(*vertex_from);
     acc->PrefetchOutEdges(*vertex_to);
     acc->PrefetchInEdges(*vertex_to);