From 4d930fb73b7d7d3112e2a1144966efc7abdbf8d6 Mon Sep 17 00:00:00 2001
From: Josip Mrden <josip.mrden@memgraph.io>
Date: Thu, 29 Feb 2024 12:49:59 +0100
Subject: [PATCH] Model execution rows

---
 src/query/procedure/mg_procedure_impl.cpp | 45 ++++++++++++-----------
 src/query/procedure/mg_procedure_impl.hpp | 14 +++++--
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/src/query/procedure/mg_procedure_impl.cpp b/src/query/procedure/mg_procedure_impl.cpp
index a5cbf0e3e..451c4cb45 100644
--- a/src/query/procedure/mg_procedure_impl.cpp
+++ b/src/query/procedure/mg_procedure_impl.cpp
@@ -4025,20 +4025,6 @@ mgp_error mgp_untrack_current_thread_allocations(mgp_graph *graph) {
   });
 }
 
-std::vector<std::vector<mgp_value>> GetMgpRowsFromTypedValues(
-    std::vector<std::vector<memgraph::query::TypedValue>> tv_rows, mgp_graph *graph, mgp_memory *memory) {
-  std::vector<std::vector<mgp_value>> mgp_rows;
-  for (auto &row : tv_rows) {
-    std::vector<mgp_value> vals;
-    for (auto &val : row) {
-      vals.emplace_back(val, graph, memory->impl);
-    }
-    mgp_rows.emplace_back(std::move(vals));
-  }
-
-  return mgp_rows;
-}
-
 mgp_execution_headers::mgp_execution_headers(memgraph::utils::pmr::vector<memgraph::utils::pmr::string> &&storage)
     : headers(std::move(storage)){};
 
@@ -4059,16 +4045,31 @@ mgp_error mgp_execution_headers_at(mgp_execution_headers *headers, size_t index,
       result);
 }
 
-mgp_execution_result::mgp_execution_result(mgp_execution_headers &&headers,
-                                           std::vector<std::vector<memgraph::query::TypedValue>> tv_rows,
-                                           mgp_graph *graph, mgp_memory *memory)
-    : headers(std::move(headers)), rows(GetMgpRowsFromTypedValues(std::move(tv_rows), graph, memory)) {}
+mgp_execution_rows::mgp_execution_rows(
+    memgraph::utils::pmr::vector<memgraph::utils::pmr::vector<memgraph::query::TypedValue>> &&tv_rows)
+    : rows(std::move(tv_rows)) {}
+
+mgp_execution_result::mgp_execution_result(mgp_execution_headers &&headers, mgp_execution_rows &&rows)
+    : headers(std::move(headers)), rows(std::move(rows)) {}
 
 struct MgProcedureResultStream final {
+  explicit MgProcedureResultStream(mgp_memory *memory) : rows(memory->impl), memory(memory) {}
   using Row = std::vector<memgraph::query::TypedValue>;
   using Rows = std::vector<Row>;
-  Rows rows;
-  void Result(const Row &row) { rows.push_back(row); }
+  using PmrRow = memgraph::utils::pmr::vector<memgraph::query::TypedValue>;
+  using PmrRows = memgraph::utils::pmr::vector<PmrRow>;
+
+  PmrRows rows;
+  mgp_memory *memory;
+
+  void Result(const Row &row) {
+    PmrRow pmr_row(memory->impl);
+    for (auto &val : row) {
+      pmr_row.emplace_back(std::move(val));
+    }
+
+    rows.emplace_back(std::move(pmr_row));
+  }
 };
 
 mgp_error mgp_execute_query(mgp_graph *graph, mgp_memory *memory, const char *query, mgp_execution_result **result) {
@@ -4087,7 +4088,7 @@ mgp_error mgp_execute_query(mgp_graph *graph, mgp_memory *memory, const char *qu
         });
 
         auto prepare_query_result = interpreter.Prepare(query_string, {}, {});
-        MgProcedureResultStream stream;
+        MgProcedureResultStream stream(memory);
         interpreter.Pull(&stream, {}, prepare_query_result.qid);
 
         memgraph::utils::pmr::vector<memgraph::utils::pmr::string> headers(memory->impl);
@@ -4096,7 +4097,7 @@ mgp_error mgp_execute_query(mgp_graph *graph, mgp_memory *memory, const char *qu
         }
 
         return NewRawMgpObject<mgp_execution_result>(memory, mgp_execution_headers{std::move(headers)},
-                                                     std::move(stream.rows), graph, memory);
+                                                     mgp_execution_rows{std::move(stream.rows)});
       },
       result);
 }
diff --git a/src/query/procedure/mg_procedure_impl.hpp b/src/query/procedure/mg_procedure_impl.hpp
index 1a8bf425b..f0efc1605 100644
--- a/src/query/procedure/mg_procedure_impl.hpp
+++ b/src/query/procedure/mg_procedure_impl.hpp
@@ -1005,13 +1005,19 @@ struct mgp_execution_headers {
   storage_type headers;
 };
 
+struct mgp_execution_rows {
+  explicit mgp_execution_rows(
+      memgraph::utils::pmr::vector<memgraph::utils::pmr::vector<memgraph::query::TypedValue>> &&tv_rows);
+  ~mgp_execution_rows() = default;
+
+  memgraph::utils::pmr::vector<memgraph::utils::pmr::vector<memgraph::query::TypedValue>> rows;
+};
+
 struct mgp_execution_result {
-  explicit mgp_execution_result(mgp_execution_headers &&headers,
-                                std::vector<std::vector<memgraph::query::TypedValue>> tv_rows, mgp_graph *graph,
-                                mgp_memory *memory);
+  explicit mgp_execution_result(mgp_execution_headers &&headers, mgp_execution_rows &&rows);
 
   ~mgp_execution_result() = default;
 
   mgp_execution_headers headers;
-  std::vector<std::vector<mgp_value>> rows;
+  mgp_execution_rows rows;
 };