Extension of PyQueryResult (#1700)

lehners · mewim · web-flow · commit 15febbaada30 · 2023-06-22T14:49:16.000-04:00
* add cmake-build-release to gitignore

* add getExecutionTime, getCompilingTime, and getNumTuples to python api

* Add tests

---------

Co-authored-by: Chang Liu &lt;liuc223@gmail.com&gt;
diff --git a/src_cpp/include/py_query_result.h b/src_cpp/include/py_query_result.h
@@ -40,6 +40,12 @@ class PyQueryResult {
 
     bool isSuccess() const;
 
+    double getExecutionTime();
+
+    double getCompilingTime();
+
+    size_t getNumTuples();
+
 private:
     static py::dict getPyDictFromProperties(
         const std::vector<std::pair<std::string, std::unique_ptr<kuzu::common::Value>>>&
diff --git a/src_cpp/py_query_result.cpp b/src_cpp/py_query_result.cpp
@@ -24,7 +24,10 @@ void PyQueryResult::initialize(py::handle& m) {
         .def("getColumnNames", &PyQueryResult::getColumnNames)
         .def("getColumnDataTypes", &PyQueryResult::getColumnDataTypes)
         .def("resetIterator", &PyQueryResult::resetIterator)
-        .def("isSuccess", &PyQueryResult::isSuccess);
+        .def("isSuccess", &PyQueryResult::isSuccess)
+        .def("getCompilingTime", &PyQueryResult::getCompilingTime)
+        .def("getExecutionTime", &PyQueryResult::getExecutionTime)
+        .def("getNumTuples", &PyQueryResult::getNumTuples);
     // PyDateTime_IMPORT is a macro that must be invoked before calling any other cpython datetime
     // macros. One could also invoke this in a separate function like constructor. See
     // https://docs.python.org/3/c-api/datetime.html for details.
@@ -236,3 +239,15 @@ py::dict PyQueryResult::convertNodeIdToPyDict(const nodeID_t& nodeId) {
     idDict["table"] = py::cast(nodeId.tableID);
     return idDict;
 }
+
+double PyQueryResult::getExecutionTime() {
+    return queryResult->getQuerySummary()->getExecutionTime();
+}
+
+double PyQueryResult::getCompilingTime() {
+    return queryResult->getQuerySummary()->getCompilingTime();
+}
+
+size_t PyQueryResult::getNumTuples() {
+    return queryResult->getNumTuples();
+}
diff --git a/src_py/query_result.py b/src_py/query_result.py
@@ -44,6 +44,15 @@ class QueryResult:
     get_as_torch_geometric()
         Converts the nodes and rels in query result into a PyTorch Geometric graph representation
         torch_geometric.data.Data or torch_geometric.data.HeteroData.
+
+    get_execution_time()
+        Get the time in ms which was required for executing the query.
+
+    get_compiling_time()
+        Get the time in ms which was required for compiling the query.
+
+    def get_num_tuples(self):
+        Get the number of tuples which the query returned.
     """
 
     def __init__(self, connection, query_result):
@@ -359,3 +368,39 @@ def get_as_torch_geometric(self):
 
         converter = TorchGeometricResultConverter(self)
         return converter.get_as_torch_geometric()
+
+    def get_execution_time(self):
+        """
+        Get the time in ms which was required for executing the query.
+
+        Returns
+        -------
+        double
+            Query execution time as double in ms.
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getExecutionTime()
+
+    def get_compiling_time(self):
+        """
+        Get the time in ms which was required for compiling the query.
+
+        Returns
+        -------
+        double
+            Query compile time as double in ms.
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getCompilingTime()
+
+    def get_num_tuples(self):
+        """
+        Get the number of tuples which the query returned.
+
+        Returns
+        -------
+        int
+            Number of tuples.
+        """
+        self.check_for_query_result_close()
+        return self._query_result.getNumTuples()
diff --git a/test/test_main.py b/test/test_main.py
@@ -8,6 +8,7 @@
 from test_networkx import *
 from test_parameter import *
 from test_prepared_statement import *
+from test_query_result import *
 from test_query_result_close import *
 from test_timeout import *
 from test_torch_geometric import *
diff --git a/test/test_query_result.py b/test/test_query_result.py
@@ -0,0 +1,19 @@
+def test_get_execution_time(establish_connection):
+    conn, db = establish_connection
+    result = conn.execute("MATCH (a:person) WHERE a.ID = 0 RETURN a")
+    assert result.get_execution_time() > 0
+    result.close()
+
+
+def test_get_compiling_time(establish_connection):
+    conn, db = establish_connection
+    result = conn.execute("MATCH (a:person) WHERE a.ID = 0 RETURN a")
+    assert result.get_compiling_time() > 0
+    result.close()
+
+
+def test_get_num_tuples(establish_connection):
+    conn, db = establish_connection
+    result = conn.execute("MATCH (a:person) WHERE a.ID = 0 RETURN a")
+    assert result.get_num_tuples() == 1
+    result.close()
