From b91b16de963b3c7b2cd2202dc2b5602aa0d4e915 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=C3=A1nos=20Benjamin=20Antal?= <benjamin.antal@memgraph.io>
Date: Tue, 17 Jan 2023 07:06:25 +0100
Subject: [PATCH] Fix `Interpreter::PullMultiple` for queries that return some
 values

---
 src/query/v2/interpreter.cpp | 20 +++++++++++---------
 src/query/v2/multiframe.cpp  |  4 ----
 src/query/v2/multiframe.hpp  |  1 -
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/src/query/v2/interpreter.cpp b/src/query/v2/interpreter.cpp
index 594942aec..23386bd1c 100644
--- a/src/query/v2/interpreter.cpp
+++ b/src/query/v2/interpreter.cpp
@@ -704,7 +704,6 @@ PullPlan::PullPlan(const std::shared_ptr<CachedPlan> plan, const Parameters &par
   ctx_.request_router = request_router;
   ctx_.edge_ids_alloc = &interpreter_context->edge_ids_alloc;
 }
-
 std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::PullMultiple(AnyStream *stream, std::optional<int> n,
                                                                         const std::vector<Symbol> &output_symbols,
                                                                         std::map<std::string, TypedValue> *summary) {
@@ -734,7 +733,7 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::PullMultiple(AnyStrea
   // Returns true if a result was pulled.
   const auto pull_result = [&]() -> bool {
     cursor_->PullMultiple(multi_frame_, ctx_);
-    return !multi_frame_.HasInvalidFrame();
+    return multi_frame_.HasValidFrame();
   };
 
   const auto stream_values = [&output_symbols, &stream](const Frame &frame) {
@@ -755,13 +754,14 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::PullMultiple(AnyStrea
   int i = 0;
   if (has_unsent_results_ && !output_symbols.empty()) {
     // stream unsent results from previous pull
-
-    auto iterator_for_valid_frame_only = multi_frame_.GetValidFramesReader();
-    for (const auto &frame : iterator_for_valid_frame_only) {
+    for (auto &frame : multi_frame_.GetValidFramesConsumer()) {
       stream_values(frame);
+      frame.MakeInvalid();
       ++i;
+      if (i == n) {
+        break;
+      }
     }
-    multi_frame_.MakeAllFramesInvalid();
   }
 
   for (; !n || i < n;) {
@@ -770,13 +770,15 @@ std::optional<plan::ProfilingStatsWithTotalTime> PullPlan::PullMultiple(AnyStrea
     }
 
     if (!output_symbols.empty()) {
-      auto iterator_for_valid_frame_only = multi_frame_.GetValidFramesReader();
-      for (const auto &frame : iterator_for_valid_frame_only) {
+      for (auto &frame : multi_frame_.GetValidFramesConsumer()) {
         stream_values(frame);
+        frame.MakeInvalid();
         ++i;
+        if (i == n) {
+          break;
+        }
       }
     }
-    multi_frame_.MakeAllFramesInvalid();
   }
 
   // If we finished because we streamed the requested n results,
diff --git a/src/query/v2/multiframe.cpp b/src/query/v2/multiframe.cpp
index 0ddfd3aa7..835cdbc0f 100644
--- a/src/query/v2/multiframe.cpp
+++ b/src/query/v2/multiframe.cpp
@@ -48,10 +48,6 @@ bool MultiFrame::HasValidFrame() const noexcept {
   return std::any_of(frames_.begin(), frames_.end(), [](auto &frame) { return frame.IsValid(); });
 }
 
-bool MultiFrame::HasInvalidFrame() const noexcept {
-  return std::any_of(frames_.rbegin(), frames_.rend(), [](auto &frame) { return !frame.IsValid(); });
-}
-
 // NOLINTNEXTLINE (bugprone-exception-escape)
 void MultiFrame::DefragmentValidFrames() noexcept {
   /*
diff --git a/src/query/v2/multiframe.hpp b/src/query/v2/multiframe.hpp
index f84fe6421..5f821bb6b 100644
--- a/src/query/v2/multiframe.hpp
+++ b/src/query/v2/multiframe.hpp
@@ -81,7 +81,6 @@ class MultiFrame {
   void MakeAllFramesInvalid() noexcept;
 
   bool HasValidFrame() const noexcept;
-  bool HasInvalidFrame() const noexcept;
 
   inline utils::MemoryResource *GetMemoryResource() { return frames_[0].GetMemoryResource(); }