Set EdgeFilterPolicy

src/storage/v2/disk/edge_filter_policy.hpp

@@ -0,0 +1,109 @@
+// Copyright 2023 Memgraph Ltd.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+// License, and you may not use this file except in compliance with the Business Source License.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+#pragma once
+
+#include "rocksdb/filter_policy.h"
+#include "rocksdb/table/block_based/filter_policy_internal.h"
+
+#include "storage/v2/edge.hpp"
+#include "utils/exceptions.hpp"
+
+namespace memgraph::storage {
+
+/// FilterBitsBuilder and FilterBitsReader are just forward declarations if I don't include filter_policy_internal.h,
+/// I cannot inherit from them.
+
+class EdgeFilterPolicy : public rocksdb::FilterPolicy {
+ private:
+  class EdgeFilterBitsReader : public rocksdb::FilterBitsReader {
+   public:
+    explicit EdgeFilterBitsReader(rocksdb::FilterBitsReader *orig_reader) : orig_reader_(orig_reader) {}
+
+    EdgeFilterBitsReader(const EdgeFilterBitsReader &) = delete;
+    EdgeFilterBitsReader &operator=(const EdgeFilterBitsReader &) = delete;
+    EdgeFilterBitsReader(EdgeFilterBitsReader &&) = delete;
+    EdgeFilterBitsReader &operator=(EdgeFilterBitsReader &&) = delete;
+    ~EdgeFilterBitsReader() override {}
+
+    bool MayMatch(const rocksdb::Slice &entry) override {
+      spdlog::trace("Entry received in MayMatch: {}", entry.ToString());
+      return orig_reader_->MayMatch(entry);
+    }
+
+   private:
+    rocksdb::FilterBitsReader *orig_reader_;
+  };
+
+  class EdgeFilterBitsBuilder : public rocksdb::FilterBitsBuilder {
+   public:
+    explicit EdgeFilterBitsBuilder(rocksdb::FilterBitsBuilder *orig_builder) : orig_builder_(orig_builder) {}
+
+    EdgeFilterBitsBuilder(const EdgeFilterBitsBuilder &) = delete;
+    EdgeFilterBitsBuilder &operator=(const EdgeFilterBitsBuilder &) = delete;
+    EdgeFilterBitsBuilder(EdgeFilterBitsBuilder &&) = delete;
+    EdgeFilterBitsBuilder &operator=(EdgeFilterBitsBuilder &&) = delete;
+    ~EdgeFilterBitsBuilder() override {}
+
+    // Add a key (or prefix) to the filter. Typically, a builder will keep
+    // a set of 64-bit key hashes and only build the filter in Finish
+    // when the final number of keys is known. Keys are added in sorted order
+    // and duplicated keys are possible, so typically, the builder will
+    // only add this key if its hash is different from the most recently
+    // added.
+    void AddKey(const rocksdb::Slice &key) override {
+      spdlog::trace("Request for AddKey: {}", key.ToString());
+      orig_builder_->AddKey(key);
+    }
+
+    // Called by RocksDB before Finish to populate
+    // TableProperties::num_filter_entries, so should represent the
+    // number of unique keys (and/or prefixes) added, but does not have
+    // to be exact. `return 0;` may be used to conspicuously indicate "unknown".
+    size_t EstimateEntriesAdded() override { return orig_builder_->EstimateEntriesAdded(); }
+
+    // Generate the filter using the keys that are added
+    // The return value of this function would be the filter bits,
+    // The ownership of actual data is set to buf
+    rocksdb::Slice Finish(std::unique_ptr<const char[]> *buf) override { return orig_builder_->Finish(buf); }
+
+    size_t ApproximateNumEntries(size_t bytes) override { return orig_builder_->ApproximateNumEntries(bytes); }
+
+   private:
+    rocksdb::FilterBitsBuilder *orig_builder_;
+  };
+
+ public:
+  explicit EdgeFilterPolicy() {}
+
+  EdgeFilterPolicy(const EdgeFilterPolicy &) = delete;
+  EdgeFilterPolicy &operator=(const EdgeFilterPolicy &) = delete;
+  EdgeFilterPolicy(EdgeFilterPolicy &&) = delete;
+  EdgeFilterPolicy &operator=(EdgeFilterPolicy &&) = delete;
+  ~EdgeFilterPolicy() override { delete orig_policy_; }
+
+  const char *Name() const override { return "EdgeFilterPolicy"; }
+
+  const char *CompatibilityName() const override { return "EdgeFilterPolicy"; }
+
+  rocksdb::FilterBitsBuilder *GetBuilderWithContext(const rocksdb::FilterBuildingContext &context) const override {
+    return new EdgeFilterBitsBuilder(orig_policy_->GetBuilderWithContext(context));
+  }
+
+  rocksdb::FilterBitsReader *GetFilterBitsReader(const rocksdb::Slice &contents) const override {
+    return new EdgeFilterBitsReader(orig_policy_->GetFilterBitsReader(contents));
+  }
+
+ private:
+  const rocksdb::FilterPolicy *orig_policy_ = rocksdb::NewBloomFilterPolicy(10);
+};
+
+}  // namespace memgraph::storage

src/storage/v2/disk/storage.cpp

@@ -25,6 +25,7 @@
 #include "kvstore/kvstore.hpp"
 #include "spdlog/spdlog.h"
 #include "storage/v2/constraints/unique_constraints.hpp"
+#include "storage/v2/disk/edge_filter_policy.hpp"
 #include "storage/v2/disk/rocksdb_storage.hpp"
 #include "storage/v2/disk/storage.hpp"
 #include "storage/v2/disk/unique_constraints.hpp"
@@ -248,6 +249,12 @@ DiskStorage::DiskStorage(Config config)
   kvstore_->options_.wal_compression = rocksdb::kNoCompression;
   std::vector<rocksdb::ColumnFamilyHandle *> column_handles;
   std::vector<rocksdb::ColumnFamilyDescriptor> column_families;
+  // Set up bloom filter
+  rocksdb::BlockBasedTableOptions table_options;
+  table_options.filter_policy.reset(new EdgeFilterPolicy());
+  kvstore_->options_.table_factory.reset(
+      rocksdb::NewBlockBasedTableFactory(table_options));  // For multiple column family setting, set up spec
+
   if (utils::DirExists(config.disk.main_storage_directory)) {
     column_families.emplace_back(vertexHandle, kvstore_->options_);
     column_families.emplace_back(edgeHandle, kvstore_->options_);
@@ -742,9 +749,25 @@ VertexAccessor DiskStorage::DiskAccessor::CreateVertex(utils::SkipList<Vertex>::
 }
 
 std::optional<VertexAccessor> DiskStorage::DiskAccessor::FindVertex(storage::Gid gid, View view) {
+  rocksdb::ReadOptions read_opts_1;
+  auto strTs_1 = utils::StringTimestamp(transaction_.start_timestamp);
+  rocksdb::Slice ts_1(strTs_1);
+  read_opts_1.timestamp = &ts_1;
+  auto *disk_storage = static_cast<DiskStorage *>(storage_);
+
+  std::string value;
+  auto maybe_res =
+      disk_transaction_->Get(read_opts_1, disk_storage->kvstore_->vertex_chandle, utils::SerializeIdType(gid), &value);
+  if (maybe_res.ok()) {
+    spdlog::trace("Found vertex with gid using Get: {}", value);
+  } else {
+    spdlog::trace("Nothing found with single get, status: {}", maybe_res.getState());
+  }
+
   auto acc = vertices_.access();
   auto vertex_it = acc.find(gid);
   if (vertex_it != acc.end()) {
+    spdlog::trace("Found in main cache");
     return VertexAccessor::Create(&*vertex_it, &transaction_, &storage_->indices_, &storage_->constraints_, config_,
                                   view);
   }
@@ -752,6 +775,7 @@ std::optional<VertexAccessor> DiskStorage::DiskAccessor::FindVertex(storage::Gid
     acc = vec->access();
     auto index_it = acc.find(gid);
     if (index_it != acc.end()) {
+      spdlog::trace("Found in index cache");
       return VertexAccessor::Create(&*index_it, &transaction_, &storage_->indices_, &storage_->constraints_, config_,
                                     view);
     }
@@ -761,9 +785,9 @@ std::optional<VertexAccessor> DiskStorage::DiskAccessor::FindVertex(storage::Gid
   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_->vertex_chandle));
+
   for (it->SeekToFirst(); it->Valid(); it->Next()) {
     const auto &key = it->key();
     if (Gid::FromUint(std::stoull(utils::ExtractGidFromKey(key.ToString()))) == gid) {