Query stats (#91)

* Add query result statistics rows_read(), bytes_read(), elapsed()

* Add example and statistics test

Author: auxten

README-zh.md

@@ -61,6 +61,7 @@ res = chdb.query('select version()', 'Pretty'); print(res)
 # 查看更多数据类型格式，请参见 tests/format_output.py
 res = chdb.query('select * from file("data.parquet", Parquet)', 'JSON'); print(res)
 res = chdb.query('select * from file("data.csv", CSV)', 'CSV');  print(res)
+print(f"SQL read {res.rows_read()} rows, {res.bytes_read()} bytes, elapsed {res.elapsed()} seconds")
 ```
 
 ### Pandas DataFrame 输出

README.md

@@ -66,6 +66,7 @@ res = chdb.query('select version()', 'Pretty'); print(res)
 # See more data type format in tests/format_output.py
 res = chdb.query('select * from file("data.parquet", Parquet)', 'JSON'); print(res)
 res = chdb.query('select * from file("data.csv", CSV)', 'CSV');  print(res)
+print(f"SQL read {res.rows_read()} rows, {res.bytes_read()} bytes, elapsed {res.elapsed()} seconds")
 ```
 
 ### Pandas dataframe output

programs/local/LocalChdb.cpp

@@ -126,6 +126,9 @@ PYBIND11_MODULE(_chdb, m)
         .def("__len__", &query_result::size)
         .def("__repr__", &query_result::str)
         .def("size", &query_result::size)
+        .def("rows_read", &query_result::rows_read)
+        .def("bytes_read", &query_result::bytes_read)
+        .def("elapsed", &query_result::elapsed)
         .def("get_memview", &query_result::get_memview);
 
 

programs/local/LocalChdb.h

@@ -56,6 +56,31 @@ class local_result_wrapper
         }
         return py::str(result->buf, result->len);
     }
+    // Query statistics
+    size_t rows_read()
+    {
+        if (result == nullptr)
+        {
+            return 0;
+        }
+        return result->rows_read;
+    }
+    size_t bytes_read()
+    {
+        if (result == nullptr)
+        {
+            return 0;
+        }
+        return result->bytes_read;
+    }
+    double elapsed()
+    {
+        if (result == nullptr)
+        {
+            return 0;
+        }
+        return result->elapsed;
+    }
 };
 
 class query_result
@@ -70,6 +95,9 @@ class query_result
     py::bytes bytes() { return result_wrapper->bytes(); }
     py::str str() { return result_wrapper->str(); }
     size_t size() { return result_wrapper->size(); }
+    size_t rows_read() { return result_wrapper->rows_read(); }
+    size_t bytes_read() { return result_wrapper->bytes_read(); }
+    double elapsed() { return result_wrapper->elapsed(); }
     memoryview_wrapper * get_memview();
 };
 

programs/local/LocalServer.cpp

@@ -933,7 +933,16 @@ void LocalServer::readArguments(int argc, char ** argv, Arguments & common_argum
 //     }
 // }
 
-std::vector<char> * pyEntryClickHouseLocal(int argc, char ** argv)
+class query_result_
+{
+public:
+    uint64_t rows;
+    uint64_t bytes;
+    double elapsed;
+    std::vector<char> * buf;
+};
+
+std::unique_ptr<query_result_> pyEntryClickHouseLocal(int argc, char ** argv)
 {
     try
     {
@@ -942,10 +951,14 @@ std::vector<char> * pyEntryClickHouseLocal(int argc, char ** argv)
         int ret = app.run();
         if (ret == 0)
         {
-            auto buf = app.getQueryOutputVector();
+            auto result = std::make_unique<query_result_>();
+            result->buf = app.getQueryOutputVector();
+            result->rows = app.getProcessedRows();
+            result->bytes = app.getProcessedBytes();
+            result->elapsed = app.getElapsedTime();
 
             // std::cerr << std::string(out->begin(), out->end()) << std::endl;
-            return buf;
+            return result;
         }
         else
         {
@@ -970,15 +983,18 @@ std::vector<char> * pyEntryClickHouseLocal(int argc, char ** argv)
 // todo fix the memory leak and unnecessary copy
 local_result * query_stable(int argc, char ** argv)
 {
-    std::vector<char> * result = pyEntryClickHouseLocal(argc, argv);
+    auto result = pyEntryClickHouseLocal(argc, argv);
     if (!result)
     {
         return nullptr;
     }
     local_result * res = new local_result;
-    res->len = result->size();
-    res->buf = result->data();
-    res->_vec = result;
+    res->len = result->buf->size();
+    res->buf = result->buf->data();
+    res->_vec = result->buf;
+    res->rows_read = result->rows;
+    res->bytes_read = result->bytes;
+    res->elapsed = result->elapsed;
     return res;
 }
 
@@ -996,10 +1012,10 @@ void free_result(local_result * result)
 
 int mainEntryClickHouseLocal(int argc, char ** argv)
 {
-    auto buf = pyEntryClickHouseLocal(argc, argv);
-    if (buf)
+    auto result = pyEntryClickHouseLocal(argc, argv);
+    if (result)
     {
-        std::cout << std::string(buf->begin(), buf->end()) << std::endl;
+        std::cout << std::string(result->buf->begin(), result->buf->end()) << std::endl;
         return 0;
     }
     else

programs/local/chdb.h

@@ -1,4 +1,5 @@
 #pragma once
+#include <cstdint>
 #include <stddef.h>
 
 extern "C" {
@@ -7,6 +8,9 @@ struct local_result
     char * buf;
     size_t len;
     void * _vec; // std::vector<char> *, for freeing
+    double elapsed;
+    uint64_t rows_read;
+    uint64_t bytes_read;
 };
 
 local_result * query_stable(int argc, char ** argv);

src/Client/ClientBase.cpp

@@ -437,6 +437,7 @@ void ClientBase::onData(Block & block, ASTPtr parsed_query)
         return;
 
     processed_rows += block.rows();
+    processed_bytes += block.bytes();
     /// Even if all blocks are empty, we still need to initialize the output stream to write empty resultset.
     initOutputFormat(block, parsed_query);
 
@@ -1575,6 +1576,7 @@ try
         {
             connection->sendData(block, /* name */"", /* scalar */false);
             processed_rows += block.rows();
+            processed_bytes += block.bytes();
         }
     }
 
@@ -1722,6 +1724,7 @@ void ClientBase::processParsedSingleQuery(const String & full_query, const Strin
     }
 
     processed_rows = 0;
+    processed_bytes = 0;
     written_first_block = false;
     progress_indication.resetProgress();
     profile_events.watch.restart();

src/Client/ClientBase.h

@@ -80,6 +80,9 @@ class ClientBase : public Poco::Util::Application, public IHints<2, ClientBase>
         // std::vector<char> vec(buf.begin(), buf.end());
         return query_result_memory;
     }
+    size_t getProcessedRows() const { return processed_rows; }
+    size_t getProcessedBytes() const { return processed_bytes; }
+    double getElapsedTime() const { return progress_indication.elapsedSeconds(); }
 
     std::vector<String> getAllRegisteredNames() const override { return cmd_options; }
 
@@ -281,6 +284,7 @@ class ClientBase : public Poco::Util::Application, public IHints<2, ClientBase>
     bool need_render_profile_events = true;
     bool written_first_block = false;
     size_t processed_rows = 0; /// How many rows have been read or written.
+    size_t processed_bytes = 0; /// How many bytes have been read or written.
     bool print_num_processed_rows = false; /// Whether to print the number of processed rows at
 
     bool print_stack_trace = false;

tests/test_gc.py

@@ -37,5 +37,6 @@ def test_gc(self):
         self.assertEqual(mv3.tobytes(), b'123,"adbcdefg"\n')
         self.assertEqual(len(mv3), 15)
 
+
 if __name__ == '__main__':
     unittest.main()

tests/test_statistics.py

@@ -0,0 +1,39 @@
+#!python3
+
+import os
+import unittest
+import chdb
+
+N = 1000
+
+
+class TestQueryStatistics(unittest.TestCase):
+    def setUp(self) -> None:
+        # create tmp csv file
+        with open(".test.csv", "w") as f:
+            f.write("a,b,c\n")
+            for i in range(N):
+                f.write(f"{i},{i*2},{i*3}\n")
+        return super().setUp()
+
+    def tearDown(self) -> None:
+        # remove tmp csv file
+        os.remove(".test.csv")
+        return super().tearDown()
+
+    def test_csv_stats(self):
+        ret = chdb.query("SELECT * FROM file('.test.csv', CSV)", "CSV")
+        self.assertEqual(ret.rows_read(), N)
+        self.assertGreater(ret.elapsed(), 0.000001)
+        self.assertEqual(ret.bytes_read(), 27000)
+        print(f"SQL read {ret.rows_read()} rows, {ret.bytes_read()} bytes, elapsed {ret.elapsed()} seconds")
+
+    def test_non_exist_stats(self):
+        ret = chdb.query("SELECT * FROM file('notexist.parquet', Parquet)", "Parquet")
+        self.assertEqual(ret.rows_read(), 0)
+        self.assertEqual(ret.bytes_read(), 0)
+        print(f"SQL read {ret.rows_read()} rows, {ret.bytes_read()} bytes, elapsed {ret.elapsed()} seconds")
+
+
+if __name__ == "__main__":
+    unittest.main()