Check-in local shard manager sender code

2023-01-30 11:35:55 +00:00
22 changed files with 180 additions and 700 deletions
--- a/src/coordinator/shard_map.hpp
+++ b/src/coordinator/shard_map.hpp
@@ -1,4 +1,4 @@
-// Copyright 2023 Memgraph Ltd.
+// Copyright 2022 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
@@ -128,9 +128,6 @@ struct LabelSpace {
  // Maps between the smallest primary key stored in the shard and the shard
  std::map<PrimaryKey, ShardMetadata> shards;
  size_t replication_factor;
-  // TODO
-  // Stub value. Should be replaced once the shard-split logic is in place.
-  int64_t split_threshold{10000};

  friend std::ostream &operator<<(std::ostream &in, const LabelSpace &label_space) {
    using utils::print_helpers::operator<<;
--- a/src/io/transport.hpp
+++ b/src/io/transport.hpp
@@ -49,10 +49,10 @@ template <Message M>
 using ResponseResult = BasicResult<TimedOut, ResponseEnvelope<M>>;

 template <Message M>
-using ResponseFuture = memgraph::io::Future<ResponseResult<M>>;
+using ResponseFuture = Future<ResponseResult<M>>;

 template <Message M>
-using ResponsePromise = memgraph::io::Promise<ResponseResult<M>>;
+using ResponsePromise = Promise<ResponseResult<M>>;

 template <Message... Ms>
 struct RequestEnvelope {
@@ -65,6 +65,19 @@ struct RequestEnvelope {
 template <Message... Ms>
 using RequestResult = BasicResult<TimedOut, RequestEnvelope<Ms...>>;

+/// This is a concrete type that allows one message type to be
+/// sent to a single address. Initially intended to be used by the
+/// Shard to send messages to the local ShardManager.
+template <Message M>
+class Sender {
+  std::function<void(M)> sender_;
+
+ public:
+  Sender(std::function<void(M)> sender) : sender_(sender) {}
+
+  void Send(M message) { sender_(message); }
+};
+
 template <typename I>
 class Io {
  I implementation_;
@@ -173,5 +186,16 @@ class Io {
  }

  LatencyHistogramSummaries ResponseLatencies() { return implementation_.ResponseLatencies(); }
+
+  template <Message M>
+  Sender<M> GetSender(Address address) {
+    Io<I> io_copy = Io(implementation_, address_);
+
+    std::function<void(M)> sender = [address, io_copy](M message) mutable {
+      io_copy.template Send<M>(address, 0, message);
+    };
+
+    return Sender{sender};
+  }
 };
 };  // namespace memgraph::io
--- a/src/query/v2/interpreter.cpp
+++ b/src/query/v2/interpreter.cpp
@@ -64,11 +64,7 @@
 #include "utils/tsc.hpp"
 #include "utils/variant_helpers.hpp"

-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-DEFINE_bool(use_multi_frame, false, "Whether to use MultiFrame or not");
-
 namespace EventCounter {
-
 extern Event ReadQuery;
 extern Event WriteQuery;
 extern Event ReadWriteQuery;
@@ -78,7 +74,6 @@ extern const Event LabelPropertyIndexCreated;

 extern const Event StreamsCreated;
 extern const Event TriggersCreated;
-
 }  // namespace EventCounter

 namespace memgraph::query::v2 {
@@ -693,7 +688,7 @@ PullPlan::PullPlan(const std::shared_ptr<CachedPlan> plan, const Parameters &par
    : plan_(plan),
      cursor_(plan->plan().MakeCursor(execution_memory)),
      frame_(plan->symbol_table().max_position(), execution_memory),
-      multi_frame_(plan->symbol_table().max_position(), FLAGS_default_multi_frame_size, execution_memory),
+      multi_frame_(plan->symbol_table().max_position(), kNumberOfFramesInMultiframe, execution_memory),
      memory_limit_(memory_limit) {
  ctx_.db_accessor = dba;
  ctx_.symbol_table = plan->symbol_table();
@@ -817,7 +812,8 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::PullMultiple(AnyStrea
 std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::Pull(AnyStream *stream, std::optional<int> n,
                                                                const std::vector<Symbol> &output_symbols,
                                                                std::map<std::string, TypedValue> *summary) {
-  if (FLAGS_use_multi_frame) {
+  auto should_pull_multiple = false;  // TODO on the long term, we will only use PullMultiple
+  if (should_pull_multiple) {
    return PullMultiple(stream, n, output_symbols, summary);
  }
  // Set up temporary memory for a single Pull. Initial memory comes from the
--- a/src/query/v2/multiframe.cpp
+++ b/src/query/v2/multiframe.cpp
@@ -17,9 +17,6 @@
 #include "query/v2/bindings/frame.hpp"
 #include "utils/pmr/vector.hpp"

-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-DEFINE_uint64(default_multi_frame_size, 100, "Default size of MultiFrame");
-
 namespace memgraph::query::v2 {

 static_assert(std::forward_iterator<ValidFramesReader::Iterator>);
--- a/src/query/v2/multiframe.hpp
+++ b/src/query/v2/multiframe.hpp
@@ -13,14 +13,10 @@

 #include <iterator>

-#include <gflags/gflags.h>
-
 #include "query/v2/bindings/frame.hpp"

-// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
-DECLARE_uint64(default_multi_frame_size);
-
 namespace memgraph::query::v2 {
+constexpr uint64_t kNumberOfFramesInMultiframe = 1000;  // TODO have it configurable

 class ValidFramesConsumer;
 class ValidFramesModifier;
--- a/src/query/v2/plan/operator.cpp
+++ b/src/query/v2/plan/operator.cpp
@@ -202,7 +202,7 @@ class DistributedCreateNodeCursor : public Cursor {
      request_router->CreateVertices(NodeCreationInfoToRequests(context, multi_frame));
    }
    PlaceNodesOnTheMultiFrame(multi_frame, context);
-    return true;
+    return false;
  }

  void Shutdown() override { input_cursor_->Shutdown(); }
@@ -218,7 +218,6 @@ class DistributedCreateNodeCursor : public Cursor {
  }

  std::vector<msgs::NewVertex> NodeCreationInfoToRequest(ExecutionContext &context, Frame &frame) {
-    primary_keys_.clear();
    std::vector<msgs::NewVertex> requests;
    msgs::PrimaryKey pk;
    msgs::NewVertex rqst;
@@ -228,27 +227,22 @@ class DistributedCreateNodeCursor : public Cursor {
    ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, nullptr,
                                  storage::v3::View::NEW);
    if (const auto *node_info_properties = std::get_if<PropertiesMapList>(&node_info_.properties)) {
-      for (const auto &[property, value_expression] : *node_info_properties) {
+      for (const auto &[key, value_expression] : *node_info_properties) {
        TypedValue val = value_expression->Accept(evaluator);
-        auto msgs_value = TypedValueToValue(val);
-        if (context.request_router->IsPrimaryProperty(primary_label, property)) {
-          rqst.primary_key.push_back(msgs_value);
-          pk.push_back(std::move(msgs_value));
-        } else {
-          rqst.properties.emplace_back(property, std::move(msgs_value));
+        if (context.request_router->IsPrimaryKey(primary_label, key)) {
+          rqst.primary_key.push_back(TypedValueToValue(val));
+          pk.push_back(TypedValueToValue(val));
        }
      }
    } else {
      auto property_map = evaluator.Visit(*std::get<ParameterLookup *>(node_info_.properties)).ValueMap();
-      for (const auto &[property, typed_value] : property_map) {
-        auto property_str = std::string(property);
-        auto property_id = context.request_router->NameToProperty(property_str);
-        auto msgs_value = TypedValueToValue(typed_value);
-        if (context.request_router->IsPrimaryProperty(primary_label, property_id)) {
-          rqst.primary_key.push_back(msgs_value);
-          pk.push_back(std::move(msgs_value));
-        } else
-          rqst.properties.emplace_back(property_id, std::move(msgs_value));
+      for (const auto &[key, value] : property_map) {
+        auto key_str = std::string(key);
+        auto property_id = context.request_router->NameToProperty(key_str);
+        if (context.request_router->IsPrimaryKey(primary_label, property_id)) {
+          rqst.primary_key.push_back(TypedValueToValue(value));
+          pk.push_back(TypedValueToValue(value));
+        }
      }
    }

@@ -274,7 +268,6 @@ class DistributedCreateNodeCursor : public Cursor {
  }

  std::vector<msgs::NewVertex> NodeCreationInfoToRequests(ExecutionContext &context, MultiFrame &multi_frame) {
-    primary_keys_.clear();
    std::vector<msgs::NewVertex> requests;
    auto multi_frame_modifier = multi_frame.GetValidFramesModifier();
    for (auto &frame : multi_frame_modifier) {
@@ -287,27 +280,22 @@ class DistributedCreateNodeCursor : public Cursor {
      ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, nullptr,
                                    storage::v3::View::NEW);
      if (const auto *node_info_properties = std::get_if<PropertiesMapList>(&node_info_.properties)) {
-        for (const auto &[property, value_expression] : *node_info_properties) {
+        for (const auto &[key, value_expression] : *node_info_properties) {
          TypedValue val = value_expression->Accept(evaluator);
-          auto msgs_value = TypedValueToValue(val);
-          if (context.request_router->IsPrimaryProperty(primary_label, property)) {
-            rqst.primary_key.push_back(msgs_value);
-            pk.push_back(std::move(msgs_value));
-          } else {
-            rqst.properties.emplace_back(property, std::move(msgs_value));
+          if (context.request_router->IsPrimaryKey(primary_label, key)) {
+            rqst.primary_key.push_back(TypedValueToValue(val));
+            pk.push_back(TypedValueToValue(val));
          }
        }
      } else {
        auto property_map = evaluator.Visit(*std::get<ParameterLookup *>(node_info_.properties)).ValueMap();
-        for (const auto &[property, typed_value] : property_map) {
-          auto property_str = std::string(property);
-          auto property_id = context.request_router->NameToProperty(property_str);
-          auto msgs_value = TypedValueToValue(typed_value);
-          if (context.request_router->IsPrimaryProperty(primary_label, property_id)) {
-            rqst.primary_key.push_back(msgs_value);
-            pk.push_back(std::move(msgs_value));
-          } else
-            rqst.properties.emplace_back(property_id, std::move(msgs_value));
+        for (const auto &[key, value] : property_map) {
+          auto key_str = std::string(key);
+          auto property_id = context.request_router->NameToProperty(key_str);
+          if (context.request_router->IsPrimaryKey(primary_label, property_id)) {
+            rqst.primary_key.push_back(TypedValueToValue(value));
+            pk.push_back(TypedValueToValue(value));
+          }
        }
      }

@@ -509,7 +497,7 @@ class DistributedScanAllAndFilterCursor : public Cursor {

    if (!own_multi_frame_.has_value()) {
      own_multi_frame_.emplace(MultiFrame(output_multi_frame.GetFirstFrame().elems().size(),
-                                          FLAGS_default_multi_frame_size, output_multi_frame.GetMemoryResource()));
+                                          kNumberOfFramesInMultiframe, output_multi_frame.GetMemoryResource()));
      own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
      own_frames_it_ = own_frames_consumer_->begin();
    }
@@ -615,6 +603,8 @@ class DistributedScanByPrimaryKeyCursor : public Cursor {
        filter_expressions_(filter_expressions),
        primary_key_(primary_key) {}

+  enum class State : int8_t { INITIALIZING, COMPLETED };
+
  using VertexAccessor = accessors::VertexAccessor;

  std::optional<VertexAccessor> MakeRequestSingleFrame(Frame &frame, RequestRouterInterface &request_router,
@@ -647,43 +637,6 @@ class DistributedScanByPrimaryKeyCursor : public Cursor {
    return VertexAccessor(vertex, properties, &request_router);
  }

-  void MakeRequestMultiFrame(MultiFrame &multi_frame, RequestRouterInterface &request_router,
-                             ExecutionContext &context) {
-    msgs::GetPropertiesRequest req;
-    const msgs::Label label = {.id = msgs::LabelId::FromUint(label_.AsUint())};
-
-    std::unordered_set<msgs::VertexId> used_vertex_ids;
-
-    for (auto &frame : multi_frame.GetValidFramesModifier()) {
-      ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, context.request_router,
-                                    storage::v3::View::NEW);
-
-      std::vector<msgs::Value> pk;
-      for (auto *primary_property : primary_key_) {
-        pk.push_back(TypedValueToValue(primary_property->Accept(evaluator)));
-      }
-
-      auto vertex_id = std::make_pair(label, std::move(pk));
-      auto [it, inserted] = used_vertex_ids.emplace(std::move(vertex_id));
-      if (inserted) {
-        req.vertex_ids.emplace_back(*it);
-      }
-    }
-
-    auto get_prop_result = std::invoke([&context, &request_router, &req]() mutable {
-      SCOPED_REQUEST_WAIT_PROFILE;
-      return request_router.GetProperties(req);
-    });
-
-    for (auto &result : get_prop_result) {
-      // TODO (gvolfing) figure out labels when relevant.
-      msgs::Vertex vertex = {.id = result.vertex, .labels = {}};
-
-      id_to_accessor_mapping_.emplace(result.vertex,
-                                      VertexAccessor(std::move(vertex), std::move(result.props), &request_router));
-    }
-  }
-
  bool Pull(Frame &frame, ExecutionContext &context) override {
    SCOPED_PROFILE_OP(op_name_);

@@ -702,108 +655,17 @@ class DistributedScanByPrimaryKeyCursor : public Cursor {
    return false;
  }

-  void EnsureOwnMultiFrameIsGood(MultiFrame &output_multi_frame) {
-    if (!own_multi_frame_.has_value()) {
-      own_multi_frame_.emplace(MultiFrame(output_multi_frame.GetFirstFrame().elems().size(),
-                                          FLAGS_default_multi_frame_size, output_multi_frame.GetMemoryResource()));
-      own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
-      own_frames_it_ = own_frames_consumer_->begin();
-    }
-    MG_ASSERT(output_multi_frame.GetFirstFrame().elems().size() == own_multi_frame_->GetFirstFrame().elems().size());
-  }
-
-  bool PullMultiple(MultiFrame &output_multi_frame, ExecutionContext &context) override {
-    SCOPED_PROFILE_OP(op_name_);
-    EnsureOwnMultiFrameIsGood(output_multi_frame);
-
-    auto output_frames_populator = output_multi_frame.GetInvalidFramesPopulator();
-    auto populated_any = false;
-
-    while (true) {
-      switch (state_) {
-        case State::PullInput: {
-          id_to_accessor_mapping_.clear();
-          if (!input_cursor_->PullMultiple(*own_multi_frame_, context)) {
-            state_ = State::Exhausted;
-            return populated_any;
-          }
-          own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
-          own_frames_it_ = own_frames_consumer_->begin();
-
-          if (own_frames_it_ == own_frames_consumer_->end()) {
-            continue;
-          }
-
-          MakeRequestMultiFrame(*own_multi_frame_, *context.request_router, context);
-
-          state_ = State::PopulateOutput;
-          break;
-        }
-        case State::PopulateOutput: {
-          if (!output_multi_frame.HasInvalidFrame()) {
-            if (own_frames_it_ == own_frames_consumer_->end()) {
-              id_to_accessor_mapping_.clear();
-            }
-            return populated_any;
-          }
-
-          if (own_frames_it_ == own_frames_consumer_->end()) {
-            state_ = State::PullInput;
-            continue;
-          }
-
-          for (auto output_frame_it = output_frames_populator.begin();
-               output_frame_it != output_frames_populator.end() && own_frames_it_ != own_frames_consumer_->end();
-               ++own_frames_it_) {
-            auto &output_frame = *output_frame_it;
-
-            ExpressionEvaluator evaluator(&*own_frames_it_, context.symbol_table, context.evaluation_context,
-                                          context.request_router, storage::v3::View::NEW);
-
-            std::vector<msgs::Value> pk;
-            for (auto *primary_property : primary_key_) {
-              pk.push_back(TypedValueToValue(primary_property->Accept(evaluator)));
-            }
-
-            const msgs::Label label = {.id = msgs::LabelId::FromUint(label_.AsUint())};
-            auto vertex_id = std::make_pair(label, std::move(pk));
-
-            if (const auto it = id_to_accessor_mapping_.find(vertex_id); it != id_to_accessor_mapping_.end()) {
-              output_frame = *own_frames_it_;
-              output_frame[output_symbol_] = TypedValue(it->second);
-              populated_any = true;
-              ++output_frame_it;
-            }
-            own_frames_it_->MakeInvalid();
-          }
-          break;
-        }
-        case State::Exhausted: {
-          return populated_any;
-        }
-      }
-    }
-    return populated_any;
-  };
-
  void Reset() override { input_cursor_->Reset(); }

  void Shutdown() override { input_cursor_->Shutdown(); }

 private:
-  enum class State { PullInput, PopulateOutput, Exhausted };
-
-  State state_{State::PullInput};
  const Symbol output_symbol_;
  const UniqueCursorPtr input_cursor_;
  const char *op_name_;
  storage::v3::LabelId label_;
  std::optional<std::vector<Expression *>> filter_expressions_;
  std::vector<Expression *> primary_key_;
-  std::optional<MultiFrame> own_multi_frame_;
-  std::optional<ValidFramesConsumer> own_frames_consumer_;
-  ValidFramesConsumer::Iterator own_frames_it_;
-  std::unordered_map<msgs::VertexId, VertexAccessor> id_to_accessor_mapping_;
 };

 ScanAll::ScanAll(const std::shared_ptr<LogicalOperator> &input, Symbol output_symbol, storage::v3::View view)
@@ -1345,47 +1207,28 @@ bool ContainsSameEdge(const TypedValue &a, const TypedValue &b) {

  return a.ValueEdge() == b.ValueEdge();
 }
-
-bool IsExpansionOk(Frame &frame, const Symbol &expand_symbol, const std::vector<Symbol> &previous_symbols) {
-  // This shouldn't raise a TypedValueException, because the planner
-  // makes sure these are all of the expected type. In case they are not
-  // an error should be raised long before this code is executed.
-  return std::ranges::all_of(previous_symbols,
-                             [&frame, &expand_value = frame[expand_symbol]](const auto &previous_symbol) {
-                               const auto &previous_value = frame[previous_symbol];
-                               return !ContainsSameEdge(previous_value, expand_value);
-                             });
-}
-
 }  // namespace

 bool EdgeUniquenessFilter::EdgeUniquenessFilterCursor::Pull(Frame &frame, ExecutionContext &context) {
  SCOPED_PROFILE_OP("EdgeUniquenessFilter");
+
+  auto expansion_ok = [&]() {
+    const auto &expand_value = frame[self_.expand_symbol_];
+    for (const auto &previous_symbol : self_.previous_symbols_) {
+      const auto &previous_value = frame[previous_symbol];
+      // This shouldn't raise a TypedValueException, because the planner
+      // makes sure these are all of the expected type. In case they are not
+      // an error should be raised long before this code is executed.
+      if (ContainsSameEdge(previous_value, expand_value)) return false;
+    }
+    return true;
+  };
+
  while (input_cursor_->Pull(frame, context))
-    if (IsExpansionOk(frame, self_.expand_symbol_, self_.previous_symbols_)) return true;
+    if (expansion_ok()) return true;
  return false;
 }

-bool EdgeUniquenessFilter::EdgeUniquenessFilterCursor::PullMultiple(MultiFrame &output_multi_frame,
-                                                                    ExecutionContext &context) {
-  SCOPED_PROFILE_OP("EdgeUniquenessFilterMF");
-  auto populated_any = false;
-
-  while (output_multi_frame.HasInvalidFrame()) {
-    if (!input_cursor_->PullMultiple(output_multi_frame, context)) {
-      return populated_any;
-    }
-    for (auto &frame : output_multi_frame.GetValidFramesConsumer()) {
-      if (IsExpansionOk(frame, self_.expand_symbol_, self_.previous_symbols_)) {
-        populated_any = true;
-      } else {
-        frame.MakeInvalid();
-      }
-    }
-  }
-  return populated_any;
-}
-
 void EdgeUniquenessFilter::EdgeUniquenessFilterCursor::Shutdown() { input_cursor_->Shutdown(); }

 void EdgeUniquenessFilter::EdgeUniquenessFilterCursor::Reset() { input_cursor_->Reset(); }
@@ -1477,55 +1320,6 @@ class AggregateCursor : public Cursor {
    auto remember_values_it = aggregation_it_->second.remember_.begin();
    for (const Symbol &remember_sym : self_.remember_) frame[remember_sym] = *remember_values_it++;

-    ++aggregation_it_;
-    return true;
-  }
-
-  bool PullMultiple(MultiFrame &multi_frame, ExecutionContext &context) override {
-    SCOPED_PROFILE_OP("AggregateMF");
-
-    if (!pulled_all_input_) {
-      ProcessAll(multi_frame, &context);
-      pulled_all_input_ = true;
-      MG_ASSERT(!multi_frame.HasValidFrame(), "ProcessAll didn't consumed all input frames!");
-      aggregation_it_ = aggregation_.begin();
-
-      // in case there is no input and no group_bys we need to return true
-      // just this once
-      if (aggregation_.empty() && self_.group_by_.empty()) {
-        auto frame = multi_frame.GetFirstFrame();
-        frame.MakeValid();
-        auto *pull_memory = context.evaluation_context.memory;
-        // place default aggregation values on the frame
-        for (const auto &elem : self_.aggregations_) {
-          frame[elem.output_sym] = DefaultAggregationOpValue(elem, pull_memory);
-        }
-        // place null as remember values on the frame
-        for (const Symbol &remember_sym : self_.remember_) {
-          frame[remember_sym] = TypedValue(pull_memory);
-        }
-        return true;
-      }
-    }
-
-    if (aggregation_it_ == aggregation_.end()) {
-      return false;
-    }
-
-    // place aggregation values on the frame
-    auto &frame = multi_frame.GetFirstFrame();
-    frame.MakeValid();
-    auto aggregation_values_it = aggregation_it_->second.values_.begin();
-    for (const auto &aggregation_elem : self_.aggregations_) {
-      frame[aggregation_elem.output_sym] = *aggregation_values_it++;
-    }
-
-    // place remember values on the frame
-    auto remember_values_it = aggregation_it_->second.remember_.begin();
-    for (const Symbol &remember_sym : self_.remember_) {
-      frame[remember_sym] = *remember_values_it++;
-    }
-
    aggregation_it_++;
    return true;
  }
@@ -1594,23 +1388,18 @@ class AggregateCursor : public Cursor {
      ProcessOne(*frame, &evaluator);
    }

-    CalculateAverages(*context);
-  }
-
-  void ProcessAll(MultiFrame &multi_frame, ExecutionContext *context) {
-    while (input_cursor_->PullMultiple(multi_frame, *context)) {
-      auto valid_frames_modifier =
-          multi_frame.GetValidFramesConsumer();  // consumer is needed i.o. reader because of the evaluator
-
-      for (auto &frame : valid_frames_modifier) {
-        ExpressionEvaluator evaluator(&frame, context->symbol_table, context->evaluation_context,
-                                      context->request_router, storage::v3::View::NEW);
-        ProcessOne(frame, &evaluator);
-        frame.MakeInvalid();
+    // calculate AVG aggregations (so far they have only been summed)
+    for (size_t pos = 0; pos < self_.aggregations_.size(); ++pos) {
+      if (self_.aggregations_[pos].op != Aggregation::Op::AVG) continue;
+      for (auto &kv : aggregation_) {
+        AggregationValue &agg_value = kv.second;
+        auto count = agg_value.counts_[pos];
+        auto *pull_memory = context->evaluation_context.memory;
+        if (count > 0) {
+          agg_value.values_[pos] = agg_value.values_[pos] / TypedValue(static_cast<double>(count), pull_memory);
+        }
      }
    }
-
-    CalculateAverages(*context);
  }

  /**
@@ -1628,20 +1417,6 @@ class AggregateCursor : public Cursor {
    Update(evaluator, &agg_value);
  }

-  void CalculateAverages(ExecutionContext &context) {
-    for (size_t pos = 0; pos < self_.aggregations_.size(); ++pos) {
-      if (self_.aggregations_[pos].op != Aggregation::Op::AVG) continue;
-      for (auto &kv : aggregation_) {
-        AggregationValue &agg_value = kv.second;
-        auto count = agg_value.counts_[pos];
-        auto *pull_memory = context.evaluation_context.memory;
-        if (count > 0) {
-          agg_value.values_[pos] = agg_value.values_[pos] / TypedValue(static_cast<double>(count), pull_memory);
-        }
-      }
-    }
-  }
-
  /** Ensures the new AggregationValue has been initialized. This means
   * that the value vectors are filled with an appropriate number of Nulls,
   * counts are set to 0 and remember values are remembered.
@@ -1675,7 +1450,7 @@ class AggregateCursor : public Cursor {
    for (; count_it < agg_value->counts_.end(); count_it++, value_it++, agg_elem_it++) {
      // COUNT(*) is the only case where input expression is optional
      // handle it here
-      auto *input_expr_ptr = agg_elem_it->value;
+      auto input_expr_ptr = agg_elem_it->value;
      if (!input_expr_ptr) {
        *count_it += 1;
        *value_it = *count_it;
@@ -1766,7 +1541,7 @@ class AggregateCursor : public Cursor {

  /** Checks if the given TypedValue is legal in MIN and MAX. If not
   * an appropriate exception is thrown. */
-  static void EnsureOkForMinMax(const TypedValue &value) {
+  void EnsureOkForMinMax(const TypedValue &value) const {
    switch (value.type()) {
      case TypedValue::Type::Bool:
      case TypedValue::Type::Int:
@@ -1782,7 +1557,7 @@ class AggregateCursor : public Cursor {

  /** Checks if the given TypedValue is legal in AVG and SUM. If not
   * an appropriate exception is thrown. */
-  static void EnsureOkForAvgSum(const TypedValue &value) {
+  void EnsureOkForAvgSum(const TypedValue &value) const {
    switch (value.type()) {
      case TypedValue::Type::Int:
      case TypedValue::Type::Double:
@@ -2197,7 +1972,14 @@ class UnwindCursor : public Cursor {
        if (!input_cursor_->Pull(frame, context)) return false;

        // successful pull from input, initialize value and iterator
-        SetInputValue(frame, context);
+        ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, context.request_router,
+                                      storage::v3::View::OLD);
+        TypedValue input_value = self_.input_expression_->Accept(evaluator);
+        if (input_value.type() != TypedValue::Type::List)
+          throw QueryRuntimeException("Argument of UNWIND must be a list, but '{}' was provided.", input_value.type());
+        // Copy the evaluted input_value_list to our vector.
+        input_value_ = input_value.ValueList();
+        input_value_it_ = input_value_.begin();
      }

      // if we reached the end of our list of values goto back to top
@@ -2208,70 +1990,6 @@ class UnwindCursor : public Cursor {
    }
  }

-  bool PullMultiple(MultiFrame &output_multi_frame, ExecutionContext &context) override {
-    SCOPED_PROFILE_OP("UnwindMF");
-
-    if (!own_multi_frame_.has_value()) {
-      own_multi_frame_.emplace(MultiFrame(output_multi_frame.GetFirstFrame().elems().size(),
-                                          FLAGS_default_multi_frame_size, output_multi_frame.GetMemoryResource()));
-      own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
-      own_frames_it_ = own_frames_consumer_->begin();
-    }
-
-    auto output_frames_populator = output_multi_frame.GetInvalidFramesPopulator();
-    auto populated_any = false;
-
-    while (true) {
-      switch (state_) {
-        case State::PullInput: {
-          if (!input_cursor_->PullMultiple(*own_multi_frame_, context)) {
-            state_ = State::Exhausted;
-            return populated_any;
-          }
-          own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
-          own_frames_it_ = own_frames_consumer_->begin();
-          state_ = State::InitializeInputValue;
-          break;
-        }
-        case State::InitializeInputValue: {
-          if (own_frames_it_ == own_frames_consumer_->end()) {
-            state_ = State::PullInput;
-            continue;
-          }
-          SetInputValue(*own_frames_it_, context);
-          state_ = State::PopulateOutput;
-          break;
-        }
-        case State::PopulateOutput: {
-          if (!output_multi_frame.HasInvalidFrame()) {
-            return populated_any;
-          }
-          if (input_value_it_ == input_value_.end()) {
-            own_frames_it_->MakeInvalid();
-            ++own_frames_it_;
-            state_ = State::InitializeInputValue;
-            continue;
-          }
-
-          for (auto output_frame_it = output_frames_populator.begin();
-               output_frame_it != output_frames_populator.end() && input_value_it_ != input_value_.end();
-               ++output_frame_it) {
-            auto &output_frame = *output_frame_it;
-            output_frame = *own_frames_it_;
-            output_frame[self_.output_symbol_] = std::move(*input_value_it_);
-            input_value_it_++;
-            populated_any = true;
-          }
-          break;
-        }
-        case State::Exhausted: {
-          return populated_any;
-        }
-      }
-    }
-    return populated_any;
-  }
-
  void Shutdown() override { input_cursor_->Shutdown(); }

  void Reset() override {
@@ -2280,36 +1998,13 @@ class UnwindCursor : public Cursor {
    input_value_it_ = input_value_.end();
  }

-  void SetInputValue(Frame &frame, ExecutionContext &context) {
-    ExpressionEvaluator evaluator(&frame, context.symbol_table, context.evaluation_context, context.request_router,
-                                  storage::v3::View::OLD);
-    TypedValue input_value = self_.input_expression_->Accept(evaluator);
-    if (input_value.type() != TypedValue::Type::List) {
-      throw QueryRuntimeException("Argument of UNWIND must be a list, but '{}' was provided.", input_value.type());
-    }
-    // It would be nice if we could move it, however it can be tricky to make it work because of allocators and
-    // different memory resources, be careful.
-    input_value_ = std::move(input_value.ValueList());
-    input_value_it_ = input_value_.begin();
-  }
-
 private:
-  using InputVector = utils::pmr::vector<TypedValue>;
-  using InputIterator = InputVector::iterator;
-
  const Unwind &self_;
  const UniqueCursorPtr input_cursor_;
  // typed values we are unwinding and yielding
-  InputVector input_value_;
+  utils::pmr::vector<TypedValue> input_value_;
  // current position in input_value_
-  InputIterator input_value_it_ = input_value_.end();
-
-  enum class State { PullInput, InitializeInputValue, PopulateOutput, Exhausted };
-
-  State state_{State::PullInput};
-  std::optional<MultiFrame> own_multi_frame_;
-  std::optional<ValidFramesConsumer> own_frames_consumer_;
-  ValidFramesConsumer::Iterator own_frames_it_;
+  decltype(input_value_)::iterator input_value_it_ = input_value_.end();
 };

 UniqueCursorPtr Unwind::MakeCursor(utils::MemoryResource *mem) const {
@@ -2972,16 +2667,27 @@ class DistributedCreateExpandCursor : public Cursor {
      const auto &v1 = v1_value.ValueVertex();
      const auto &v2 = OtherVertex(frame);

+      // Set src and dest vertices
+      // TODO(jbajic) Currently we are only handling scenario where vertices
+      // are matched
+      const auto set_vertex = [&context](const auto &vertex, auto &vertex_id) {
+        vertex_id.first = vertex.PrimaryLabel();
+        for (const auto &[key, val] : vertex.Properties()) {
+          if (context.request_router->IsPrimaryKey(vertex_id.first.id, key)) {
+            vertex_id.second.push_back(val);
+          }
+        }
+      };
      std::invoke([&]() {
        switch (edge_info.direction) {
          case EdgeAtom::Direction::IN: {
-            request.src_vertex = v2.Id();
-            request.dest_vertex = v1.Id();
+            set_vertex(v2, request.src_vertex);
+            set_vertex(v1, request.dest_vertex);
            break;
          }
          case EdgeAtom::Direction::OUT: {
-            request.src_vertex = v1.Id();
-            request.dest_vertex = v2.Id();
+            set_vertex(v1, request.src_vertex);
+            set_vertex(v2, request.dest_vertex);
            break;
          }
          case EdgeAtom::Direction::BOTH:
@@ -3123,9 +2829,6 @@ class DistributedExpandCursor : public Cursor {
      auto &vertex = vertex_value.ValueVertex();
      msgs::ExpandOneRequest request;
      request.direction = DirectionToMsgsDirection(self_.common_.direction);
-      std::transform(self_.common_.edge_types.begin(), self_.common_.edge_types.end(),
-                     std::back_inserter(request.edge_types),
-                     [](const storage::v3::EdgeTypeId edge_type_id) { return msgs::EdgeType{edge_type_id}; });
      // to not fetch any properties of the edges
      request.edge_properties.emplace();
      request.src_vertices.push_back(vertex.Id());
@@ -3266,9 +2969,6 @@ class DistributedExpandCursor : public Cursor {

    msgs::ExpandOneRequest request;
    request.direction = DirectionToMsgsDirection(self_.common_.direction);
-    std::transform(self_.common_.edge_types.begin(), self_.common_.edge_types.end(),
-                   std::back_inserter(request.edge_types),
-                   [](const storage::v3::EdgeTypeId edge_type_id) { return msgs::EdgeType{edge_type_id}; });
    // to not fetch any properties of the edges
    request.edge_properties.emplace();
    for (const auto &frame : own_multi_frame_->GetValidFramesReader()) {
@@ -3382,7 +3082,7 @@ class DistributedExpandCursor : public Cursor {
  void EnsureOwnMultiFrameIsGood(MultiFrame &output_multi_frame) {
    if (!own_multi_frame_.has_value()) {
      own_multi_frame_.emplace(MultiFrame(output_multi_frame.GetFirstFrame().elems().size(),
-                                          FLAGS_default_multi_frame_size, output_multi_frame.GetMemoryResource()));
+                                          kNumberOfFramesInMultiframe, output_multi_frame.GetMemoryResource()));
      own_frames_consumer_.emplace(own_multi_frame_->GetValidFramesConsumer());
      own_frames_it_ = own_frames_consumer_->begin();
    }
--- a/src/query/v2/plan/operator.lcp
+++ b/src/query/v2/plan/operator.lcp
@@ -1570,7 +1570,6 @@ edge lists).")
     EdgeUniquenessFilterCursor(const EdgeUniquenessFilter &,
                                utils::MemoryResource *);
     bool Pull(Frame &, ExecutionContext &) override;
-     bool PullMultiple(MultiFrame &, ExecutionContext &) override;
     void Shutdown() override;
     void Reset() override;

--- a/src/query/v2/plan/rewrite/index_lookup.hpp
+++ b/src/query/v2/plan/rewrite/index_lookup.hpp
@@ -17,7 +17,6 @@
 #pragma once

 #include <algorithm>
-#include <limits>
 #include <memory>
 #include <optional>
 #include <unordered_map>
@@ -562,12 +561,8 @@ class IndexLookupRewriter final : public HierarchicalLogicalOperatorVisitor {
  // `max_vertex_count` controls, whether no operator should be created if the
  // vertex count in the best index exceeds this number. In such a case,
  // `nullptr` is returned and `input` is not chained.
-  // std::unique_ptr<ScanAll> GenScanByIndex(const ScanAll &scan, const std::optional<int64_t> &max_vertex_count =
-  // std::nullopt) {
-  std::unique_ptr<ScanAll> GenScanByIndex(const ScanAll &scan, std::optional<int64_t> max_vertex_count = std::nullopt) {
-    // debug (gvolfing)
-    max_vertex_count = std::numeric_limits<int64_t>::max();
-
+  std::unique_ptr<ScanAll> GenScanByIndex(const ScanAll &scan,
+                                          const std::optional<int64_t> &max_vertex_count = std::nullopt) {
    const auto &input = scan.input();
    const auto &node_symbol = scan.output_symbol_;
    const auto &view = scan.view_;
@@ -602,9 +597,6 @@ class IndexLookupRewriter final : public HierarchicalLogicalOperatorVisitor {
                     [](const auto &schema_elem) { return schema_elem.property_id; });

      for (const auto &property_filter : property_filters) {
-        if (property_filter.property_filter->type_ != PropertyFilter::Type::EQUAL) {
-          continue;
-        }
        const auto &property_id = db_->NameToProperty(property_filter.property_filter->property_.name);
        if (std::find(schema_properties.begin(), schema_properties.end(), property_id) != schema_properties.end()) {
          pk_temp.emplace_back(std::make_pair(property_filter.expression, property_filter));
--- a/src/query/v2/plan/vertex_count_cache.hpp
+++ b/src/query/v2/plan/vertex_count_cache.hpp
@@ -38,15 +38,12 @@ class VertexCountCache {
  auto NameToProperty(const std::string &name) { return request_router_->NameToProperty(name); }
  auto NameToEdgeType(const std::string &name) { return request_router_->NameToEdgeType(name); }

-  int64_t VerticesCount() { return request_router_->GetApproximateVertexCount(); }
+  int64_t VerticesCount() { return 1; }

-  int64_t VerticesCount(storage::v3::LabelId label) { return request_router_->GetApproximateVertexCount(label); }
+  int64_t VerticesCount(storage::v3::LabelId /*label*/) { return 1; }

-  int64_t VerticesCount(storage::v3::LabelId label, storage::v3::PropertyId property) {
-    return request_router_->GetApproximateVertexCount(label, property);
-  }
+  int64_t VerticesCount(storage::v3::LabelId /*label*/, storage::v3::PropertyId /*property*/) { return 1; }

-  // TODO(gvolfing) check if we actually use these overloads...
  int64_t VerticesCount(storage::v3::LabelId /*label*/, storage::v3::PropertyId /*property*/,
                        const storage::v3::PropertyValue & /*value*/) {
    return 1;
--- a/src/query/v2/request_router.hpp
+++ b/src/query/v2/request_router.hpp
@@ -117,15 +117,11 @@ class RequestRouterInterface {
  virtual std::optional<storage::v3::EdgeTypeId> MaybeNameToEdgeType(const std::string &name) const = 0;
  virtual std::optional<storage::v3::LabelId> MaybeNameToLabel(const std::string &name) const = 0;
  virtual bool IsPrimaryLabel(storage::v3::LabelId label) const = 0;
-  virtual bool IsPrimaryProperty(storage::v3::LabelId primary_label, storage::v3::PropertyId property) const = 0;
+  virtual bool IsPrimaryKey(storage::v3::LabelId primary_label, storage::v3::PropertyId property) const = 0;

  virtual std::optional<std::pair<uint64_t, uint64_t>> AllocateInitialEdgeIds(io::Address coordinator_address) = 0;
  virtual void InstallSimulatorTicker(std::function<bool()> tick_simulator) = 0;
  virtual const std::vector<coordinator::SchemaProperty> &GetSchemaForLabel(storage::v3::LabelId label) const = 0;
-
-  virtual int64_t GetApproximateVertexCount() const = 0;
-  virtual int64_t GetApproximateVertexCount(storage::v3::LabelId label) const = 0;
-  virtual int64_t GetApproximateVertexCount(storage::v3::LabelId label, storage::v3::PropertyId property) const = 0;
 };

 // TODO(kostasrim)rename this class template
@@ -235,7 +231,7 @@ class RequestRouter : public RequestRouterInterface {
    return edge_types_.IdToName(id.AsUint());
  }

-  bool IsPrimaryProperty(storage::v3::LabelId primary_label, storage::v3::PropertyId property) const override {
+  bool IsPrimaryKey(storage::v3::LabelId primary_label, storage::v3::PropertyId property) const override {
    const auto schema_it = shards_map_.schemas.find(primary_label);
    MG_ASSERT(schema_it != shards_map_.schemas.end(), "Invalid primary label id: {}", primary_label.AsUint());

@@ -419,30 +415,6 @@ class RequestRouter : public RequestRouterInterface {
    return shards_map_.GetLabelId(name);
  }

-  int64_t GetApproximateVertexCount() const override {
-    int64_t vertex_count = 0;
-
-    for (const auto &label_space : shards_map_.label_spaces) {
-      const auto split_threshold = label_space.second.split_threshold;
-      const auto shard_count = static_cast<int64_t>(label_space.second.shards.size());
-      vertex_count += split_threshold * shard_count;
-    }
-
-    return vertex_count;
-  }
-
-  int64_t GetApproximateVertexCount(storage::v3::LabelId label) const override {
-    const auto &label_space = shards_map_.label_spaces.at(label);
-    return label_space.split_threshold * label_space.shards.size();
-  }
-
-  int64_t GetApproximateVertexCount(storage::v3::LabelId label, storage::v3::PropertyId /*property*/) const override {
-    // TODO(gvolfing)
-    // Once we have reliable metadata to approximate the
-    // vertex count -based on properties- rework this function.
-    return GetApproximateVertexCount(label);
-  }
-
 private:
  std::vector<ShardRequestState<msgs::CreateVerticesRequest>> RequestsForCreateVertices(
      const std::vector<msgs::NewVertex> &new_vertices) {
--- a/src/storage/v3/shard_worker.hpp
+++ b/src/storage/v3/shard_worker.hpp
@@ -196,6 +196,13 @@ class ShardWorker {
    // TODO(tyler) get peers from Coordinator in HeartbeatResponse
    std::vector<Address> rsm_peers = {};

+    Address local_shard_manager_address = io_.GetAddress().ForkLocalShardManager();
+    io::Sender<io::messages::ShardManagerMessages> local_shard_manager_sender =
+        io_.template GetSender<io::messages::ShardManagerMessages>(local_shard_manager_address);
+
+    // TODO(tyler) pass this local_shard_manager_sender to the Shard so that it can communicate back to the local
+    // manager from split code
+
    std::unique_ptr<Shard> shard = std::make_unique<Shard>(to_init.label_id, to_init.min_key, to_init.max_key,
                                                           to_init.schema, to_init.config, to_init.id_to_names);

--- a/tests/mgbench/compare_results.py
+++ b/tests/mgbench/compare_results.py
@@ -14,6 +14,7 @@
 import argparse
 import json

+
 FIELDS = [
    {
        "name": "throughput",
@@ -84,32 +85,39 @@ def compare_results(results_from, results_to, fields):
                if group == "__import__":
                    continue
                for scenario, summary_to in scenarios.items():
-                    summary_from = recursive_get(results_from, dataset, variant, group, scenario, value={})
-                    if (
-                        len(summary_from) > 0
-                        and summary_to["count"] != summary_from["count"]
-                        or summary_to["num_workers"] != summary_from["num_workers"]
-                    ):
+                    summary_from = recursive_get(
+                        results_from, dataset, variant, group, scenario,
+                        value={})
+                    if len(summary_from) > 0 and \
+                            summary_to["count"] != summary_from["count"] or \
+                            summary_to["num_workers"] != \
+                            summary_from["num_workers"]:
                        raise Exception("Incompatible results!")
-                    testcode = "/".join([dataset, variant, group, scenario, "{:02d}".format(summary_to["num_workers"])])
+                    testcode = "/".join([dataset, variant, group, scenario,
+                                         "{:02d}".format(
+                                             summary_to["num_workers"])])
                    row = {}
                    performance_changed = False
                    for field in fields:
                        key = field["name"]
                        if key in summary_to:
-                            row[key] = compute_diff(summary_from.get(key, None), summary_to[key])
+                            row[key] = compute_diff(
+                                summary_from.get(key, None),
+                                summary_to[key])
                        elif key in summary_to["database"]:
                            row[key] = compute_diff(
-                                recursive_get(summary_from, "database", key, value=None), summary_to["database"][key]
-                            )
+                                recursive_get(summary_from, "database", key,
+                                              value=None),
+                                summary_to["database"][key])
                        else:
                            row[key] = compute_diff(
-                                recursive_get(summary_from, "metadata", key, "average", value=None),
-                                summary_to["metadata"][key]["average"],
-                            )
-                        if "diff" not in row[key] or (
-                            "diff_treshold" in field and abs(row[key]["diff"]) >= field["diff_treshold"]
-                        ):
+                                recursive_get(summary_from, "metadata", key,
+                                              "average", value=None),
+                                summary_to["metadata"][key]["average"])
+                        if "diff" not in row[key] or \
+                                ("diff_treshold" in field and
+                                 abs(row[key]["diff"]) >=
+                                 field["diff_treshold"]):
                            performance_changed = True
                    if performance_changed:
                        ret[testcode] = row
@@ -122,36 +130,29 @@ def generate_remarkup(fields, data):
        ret += "<table>\n"
        ret += "  <tr>\n"
        ret += "    <th>Testcode</th>\n"
-        ret += (
-            "\n".join(
-                map(
-                    lambda x: "    <th>{}</th>".format(x["name"].replace("_", " ").capitalize()),
-                    fields,
-                )
-            )
-            + "\n"
-        )
+        ret += "\n".join(map(lambda x: "    <th>{}</th>".format(
+            x["name"].replace("_", " ").capitalize()), fields)) + "\n"
        ret += "  </tr>\n"
        for testcode in sorted(data.keys()):
            ret += "  <tr>\n"
            ret += "    <td>{}</td>\n".format(testcode)
            for field in fields:
-                result = data[testcode].get(field["name"])
-                if result != None:
-                    value = result["value"] * field["scaling"]
-                    if "diff" in result:
-                        diff = result["diff"]
-                        arrow = "arrow-up" if diff >= 0 else "arrow-down"
-                        if not (field["positive_diff_better"] ^ (diff >= 0)):
-                            color = "green"
-                        else:
-                            color = "red"
-                        sign = "{{icon {} color={}}}".format(arrow, color)
-                        ret += '    <td bgcolor="{}">{:.3f}{} ({:+.2%})</td>\n'.format(
-                            color, value, field["unit"], diff
-                        )
+                result = data[testcode][field["name"]]
+                value = result["value"] * field["scaling"]
+                if "diff" in result:
+                    diff = result["diff"]
+                    arrow = "arrow-up" if diff >= 0 else "arrow-down"
+                    if not (field["positive_diff_better"] ^ (diff >= 0)):
+                        color = "green"
                    else:
-                        ret += '<td bgcolor="blue">{:.3f}{} //(new)// </td>\n'.format(value, field["unit"])
+                        color = "red"
+                    sign = "{{icon {} color={}}}".format(arrow, color)
+                    ret += "    <td>{:.3f}{} //({:+.2%})// {}</td>\n".format(
+                        value, field["unit"], diff, sign)
+                else:
+                    ret += "    <td>{:.3f}{} //(new)// " \
+                           "{{icon plus color=blue}}</td>\n".format(
+                               value, field["unit"])
            ret += "  </tr>\n"
        ret += "</table>\n"
    else:
@@ -160,14 +161,11 @@ def generate_remarkup(fields, data):


 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Compare results of multiple benchmark runs.")
-    parser.add_argument(
-        "--compare",
-        action="append",
-        nargs=2,
-        metavar=("from", "to"),
-        help="compare results between `from` and `to` files",
-    )
+    parser = argparse.ArgumentParser(
+        description="Compare results of multiple benchmark runs.")
+    parser.add_argument("--compare", action="append", nargs=2,
+                        metavar=("from", "to"),
+                        help="compare results between `from` and `to` files")
    parser.add_argument("--output", default="", help="output file name")
    args = parser.parse_args()

--- a/tests/mgbench/dataset_creator_unwind.py
+++ b/tests/mgbench/dataset_creator_unwind.py
@@ -1,139 +0,0 @@
-# Copyright 2022 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.
-
-import argparse
-import random
-
-import helpers
-
-# Explaination of datasets:
-#   - empty_only_index: contains index; contains no data
-#   - small: contains index; contains data (small dataset)
-#
-# Datamodel is as follow:
-#
-#                               ┌──────────────┐
-#                               │ Permission   │
-#         ┌────────────────┐    │  Schema:uuid │   ┌────────────┐
-#         │:IS_FOR_IDENTITY├────┤  Index:name  ├───┤:IS_FOR_FILE│
-#         └┬───────────────┘    └──────────────┘   └────────────┤
-#          │                                                    │
-#   ┌──────▼──────────────┐                                  ┌──▼────────────────┐
-#   │  Identity           │                                  │ File              │
-#   │   Schema:uuid       │                                  │  Schema:uuid      │
-#   │   Index:email       │                                  │  Index:name       │
-#   └─────────────────────┘                                  │  Index:platformId │
-#                                                            └───────────────────┘
-#
-#   - File: attributes: ["uuid", "name", "platformId"]
-#   - Permission: attributes: ["uuid", "name"]
-#   - Identity: attributes: ["uuid", "email"]
-#
-# Indexes:
-#   - File: [File(uuid), File(platformId), File(name)]
-#   - Permission: [Permission(uuid), Permission(name)]
-#   - Identity: [Identity(uuid), Identity(email)]
-#
-# Edges:
-#   - (:Permission)-[:IS_FOR_FILE]->(:File)
-#   - (:Permission)-[:IS_FOR_IDENTITYR]->(:Identity)
-#
-# AccessControl specific: uuid is the schema
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--number_of_identities", type=int, default=10)
-    parser.add_argument("--number_of_files", type=int, default=10)
-    parser.add_argument("--percentage_of_permissions", type=float, default=1.0)
-    parser.add_argument("--filename", default="dataset.cypher")
-
-    args = parser.parse_args()
-
-    number_of_identities = args.number_of_identities
-    number_of_files = args.number_of_files
-    percentage_of_permissions = args.percentage_of_permissions
-    filename = args.filename
-
-    assert number_of_identities >= 0
-    assert number_of_files >= 0
-    assert percentage_of_permissions > 0.0 and percentage_of_permissions <= 1.0
-    assert filename != ""
-
-    with open(filename, "w") as f:
-        f.write("MATCH (n) DETACH DELETE n;\n")
-
-        # Create the indexes
-        f.write("CREATE INDEX ON :File;\n")
-        f.write("CREATE INDEX ON :Permission;\n")
-        f.write("CREATE INDEX ON :Identity;\n")
-        f.write("CREATE INDEX ON :File(platformId);\n")
-        f.write("CREATE INDEX ON :File(name);\n")
-        f.write("CREATE INDEX ON :Permission(name);\n")
-        f.write("CREATE INDEX ON :Identity(email);\n")
-
-        # Create extra index: in distributed, this will be the schema
-        f.write("CREATE INDEX ON :File(uuid);\n")
-        f.write("CREATE INDEX ON :Permission(uuid);\n")
-        f.write("CREATE INDEX ON :Identity(uuid);\n")
-
-        uuid = 1
-
-        # Create the nodes File
-        f.write("UNWIND [")
-        for index in range(0, number_of_files):
-            if index != 0:
-                f.write(",")
-            f.write(f' {{uuid: {uuid}, platformId: "platform_id", name: "name_file_{uuid}"}}')
-            uuid += 1
-        f.write("] AS props CREATE (:File {uuid: props.uuid, platformId: props.platformId, name: props.name});\n")
-
-        identities = []
-        f.write("UNWIND [")
-        # Create the nodes Identity
-        for index in range(0, number_of_identities):
-            if index != 0:
-                f.write(",")
-            f.write(f' {{uuid: {uuid}, name: "mail_{uuid}@something.com"}}')
-            uuid += 1
-        f.write("] AS props CREATE (:Identity {uuid: props.uuid, name: props.name});\n")
-
-        f.write("UNWIND [")
-        created = 0
-        for outer_index in range(0, number_of_files):
-            for inner_index in range(0, number_of_identities):
-
-                file_uuid = outer_index + 1
-                identity_uuid = number_of_files + inner_index + 1
-
-                if random.random() <= percentage_of_permissions:
-
-                    if created > 0:
-                        f.write(",")
-
-                    f.write(
-                        f' {{permUuid: {uuid}, permName: "name_permission_{uuid}", fileUuid: {file_uuid}, identityUuid: {identity_uuid}}}'
-                    )
-                    created += 1
-                    uuid += 1
-
-                    if created == 5000:
-                        f.write(
-                            "] AS props MATCH (file:File {uuid:props.fileUuid}), (identity:Identity {uuid: props.identityUuid}) CREATE (permission:Permission {uuid: props.permUuid, name: props.permName}) CREATE (permission)-[: IS_FOR_FILE]->(file) CREATE (permission)-[: IS_FOR_IDENTITY]->(identity);\nUNWIND ["
-                        )
-                        created = 0
-        f.write(
-            "] AS props MATCH (file:File {uuid:props.fileUuid}), (identity:Identity {uuid: props.identityUuid}) CREATE (permission:Permission {uuid: props.permUuid, name: props.permName}) CREATE (permission)-[: IS_FOR_FILE]->(file) CREATE (permission)-[: IS_FOR_IDENTITY]->(identity);\n"
-        )
-
-
-if __name__ == "__main__":
-    main()
--- a/tests/mgbench/datasets.py
+++ b/tests/mgbench/datasets.py
@@ -353,7 +353,7 @@ class AccessControl(Dataset):

    def benchmark__create__vertex(self):
        self.next_value_idx += 1
-        query = ("CREATE (:File {uuid: $uuid})", {"uuid": self.next_value_idx})
+        query = (f"CREATE (:File {{uuid: {self.next_value_idx}}});", {})
        return query

    def benchmark__create__edges(self):
@@ -379,24 +379,6 @@ class AccessControl(Dataset):
        return query

    def benchmark__match__match_all_vertices_with_edges(self):
+        self.next_value_idx += 1
        query = ("MATCH (permission:Permission)-[e:IS_FOR_FILE]->(file:File) RETURN *", {})
        return query
-
-    def benchmark__match__match_users_with_permission_for_files(self):
-        file_uuid_1 = self._get_random_uuid("File")
-        file_uuid_2 = self._get_random_uuid("File")
-        min_file_uuid = min(file_uuid_1, file_uuid_2)
-        max_file_uuid = max(file_uuid_1, file_uuid_2)
-        query = (
-            "MATCH (f:File)<-[ff:IS_FOR_FILE]-(p:Permission)-[fi:IS_FOR_IDENTITY]->(i:Identity) WHERE f.uuid >= $min_file_uuid  AND f.uuid <= $max_file_uuid RETURN *",
-            {"min_file_uuid": min_file_uuid, "max_file_uuid": max_file_uuid},
-        )
-        return query
-
-    def benchmark__match__match_users_with_permission_for_specific_file(self):
-        file_uuid = self._get_random_uuid("File")
-        query = (
-            "MATCH (f:File {uuid: $file_uuid})<-[ff:IS_FOR_FILE]-(p:Permission)-[fi:IS_FOR_IDENTITY]->(i:Identity) RETURN *",
-            {"file_uuid": file_uuid},
-        )
-        return query
--- a/tests/mgbench/runners.py
+++ b/tests/mgbench/runners.py
@@ -68,15 +68,6 @@ class Memgraph:
        self._cleanup()
        atexit.unregister(self._cleanup)

-    # Returns None if string_value is not true or false, casing doesn't matter
-    def _get_bool_value(self, string_value):
-        lower_string_value = string_value.lower()
-        if lower_string_value == "true":
-            return True
-        if lower_string_value == "false":
-            return False
-        return None
-
    def _get_args(self, **kwargs):
        data_directory = os.path.join(self._directory.name, "memgraph")
        if self._memgraph_version >= (0, 50, 0):
@@ -92,13 +83,7 @@ class Memgraph:
            args_list = self._extra_args.split(" ")
            assert len(args_list) % 2 == 0
            for i in range(0, len(args_list), 2):
-                key = args_list[i]
-                value = args_list[i + 1]
-                maybe_bool_value = self._get_bool_value(value)
-                if maybe_bool_value is not None:
-                    kwargs[key] = maybe_bool_value
-                else:
-                    kwargs[key] = value
+                kwargs[args_list[i]] = args_list[i + 1]

        return _convert_args_to_flags(self._memgraph_binary, **kwargs)

--- a/tests/mgbench/splitfiles/accesscontrol_large.shard_configuration
+++ b/tests/mgbench/splitfiles/accesscontrol_large.shard_configuration
@@ -1,12 +1,8 @@
-8
+4
 uuid
 email
 name
 platformId
-permUuid
-permName
-fileUuid
-identityUuid
 2
 IS_FOR_IDENTITY
 IS_FOR_FILE
--- a/tests/mgbench/splitfiles/accesscontrol_medium.shard_configuration
+++ b/tests/mgbench/splitfiles/accesscontrol_medium.shard_configuration
@@ -1,12 +1,8 @@
-8
+4
 uuid
 email
 name
 platformId
-permUuid
-permName
-fileUuid
-identityUuid
 2
 IS_FOR_IDENTITY
 IS_FOR_FILE
--- a/tests/mgbench/splitfiles/accesscontrol_small.shard_configuration
+++ b/tests/mgbench/splitfiles/accesscontrol_small.shard_configuration
@@ -1,12 +1,8 @@
-8
+4
 uuid
 email
 name
 platformId
-permUuid
-permName
-fileUuid
-identityUuid
 2
 IS_FOR_IDENTITY
 IS_FOR_FILE
--- a/tests/unit/mock_helpers.hpp
+++ b/tests/unit/mock_helpers.hpp
@@ -41,14 +41,10 @@ class MockedRequestRouter : public RequestRouterInterface {
  MOCK_METHOD(std::optional<storage::v3::EdgeTypeId>, MaybeNameToEdgeType, (const std::string &), (const));
  MOCK_METHOD(std::optional<storage::v3::LabelId>, MaybeNameToLabel, (const std::string &), (const));
  MOCK_METHOD(bool, IsPrimaryLabel, (storage::v3::LabelId), (const));
-  MOCK_METHOD(bool, IsPrimaryProperty, (storage::v3::LabelId, storage::v3::PropertyId), (const));
+  MOCK_METHOD(bool, IsPrimaryKey, (storage::v3::LabelId, storage::v3::PropertyId), (const));
  MOCK_METHOD((std::optional<std::pair<uint64_t, uint64_t>>), AllocateInitialEdgeIds, (io::Address));
  MOCK_METHOD(void, InstallSimulatorTicker, (std::function<bool()>));
  MOCK_METHOD(const std::vector<coordinator::SchemaProperty> &, GetSchemaForLabel, (storage::v3::LabelId), (const));
-  MOCK_METHOD(int64_t, GetApproximateVertexCount, (), (const));
-  MOCK_METHOD(int64_t, GetApproximateVertexCount, (storage::v3::LabelId label), (const));
-  MOCK_METHOD(int64_t, GetApproximateVertexCount, (storage::v3::LabelId label, storage::v3::PropertyId property),
-              (const));
 };

 class MockedLogicalOperator : public plan::LogicalOperator {
--- a/tests/unit/query_v2_create_node_multiframe.cpp
+++ b/tests/unit/query_v2_create_node_multiframe.cpp
@@ -58,7 +58,7 @@ TEST(CreateNodeTest, CreateNodeCursor) {
  MockedRequestRouter router;
  EXPECT_CALL(router, CreateVertices(_)).Times(1).WillOnce(Return(std::vector<msgs::CreateVerticesResponse>{}));
  EXPECT_CALL(router, IsPrimaryLabel(_)).WillRepeatedly(Return(true));
-  EXPECT_CALL(router, IsPrimaryProperty(_, _)).WillRepeatedly(Return(true));
+  EXPECT_CALL(router, IsPrimaryKey(_, _)).WillRepeatedly(Return(true));
  auto context = MakeContext(ast, symbol_table, &router, &id_alloc);
  auto multi_frame = CreateMultiFrame(context.symbol_table.max_position());
  cursor->PullMultiple(multi_frame, context);
--- a/tests/unit/query_v2_expression_evaluator.cpp
+++ b/tests/unit/query_v2_expression_evaluator.cpp
@@ -123,7 +123,7 @@ class MockedRequestRouter : public RequestRouterInterface {

  bool IsPrimaryLabel(LabelId label) const override { return true; }

-  bool IsPrimaryProperty(LabelId primary_label, PropertyId property) const override { return true; }
+  bool IsPrimaryKey(LabelId primary_label, PropertyId property) const override { return true; }

  std::optional<std::pair<uint64_t, uint64_t>> AllocateInitialEdgeIds(io::Address coordinator_address) override {
    return {};
@@ -135,13 +135,6 @@ class MockedRequestRouter : public RequestRouterInterface {
    return schema;
  };

-  // TODO(gvolfing) once the real implementation is done make sure these are solved as well.
-  int64_t GetApproximateVertexCount() const override { return 1; }
-  int64_t GetApproximateVertexCount(storage::v3::LabelId label) const override { return 1; }
-  int64_t GetApproximateVertexCount(storage::v3::LabelId label, storage::v3::PropertyId property) const override {
-    return 1;
-  }
-
 private:
  void SetUpNameIdMappers() {
    std::unordered_map<uint64_t, std::string> id_to_name;
--- a/tests/unit/query_v2_plan.cpp
+++ b/tests/unit/query_v2_plan.cpp
@@ -86,7 +86,7 @@ class TestPlanner : public ::testing::Test {};

 using PlannerTypes = ::testing::Types<Planner>;

-TYPED_TEST_SUITE(TestPlanner, PlannerTypes);
+TYPED_TEST_CASE(TestPlanner, PlannerTypes);

 TYPED_TEST(TestPlanner, MatchFilterPropIsNotNull) {
  const char *prim_label_name = "prim_label_one";