fix: fix streaming dataframe output

wudidapaopao · wudidapaopao · commit 58388ceb3c46 · 2025-12-06T07:36:53.000+08:00
diff --git a/programs/local/ChdbClient.cpp b/programs/local/ChdbClient.cpp
@@ -14,6 +14,9 @@
 
 #if USE_PYTHON
 #include <PythonTableCache.h>
+#include <PandasDataFrameBuilder.h>
+#include <pybind11/pybind11.h>
+namespace py = pybind11;
 #endif
 
 namespace DB
@@ -318,14 +321,21 @@ CHDB::QueryResultPtr ChdbClient::executeStreamingIterate(void * streaming_result
 #if USE_PYTHON
             if (Poco::toLower(default_output_format) == "dataframe")
             {
-                res = std::make_unique<CHDB::ChunkQueryResult>(
+                auto rows_read = processed_rows - old_processed_rows;
+                auto chunk_result = std::make_unique<CHDB::ChunkQueryResult>(
                     std::move(collected_chunks),
                     std::move(collected_chunks_header),
                     elapsed_time - old_elapsed_time,
-                    processed_rows - old_processed_rows,
+                    rows_read,
                     processed_bytes - old_processed_bytes,
                     storage_rows_read - old_storage_rows_read,
                     storage_bytes_read - old_storage_bytes_read);
+
+                py::gil_scoped_acquire acquire;
+                CHDB::PandasDataFrameBuilder builder(*chunk_result);
+                py::handle df = builder.getDataFrame().release();
+
+                res = std::make_unique<CHDB::DataFrameQueryResult>(df, rows_read);
             }
             else
 #endif
@@ -350,7 +360,7 @@ CHDB::QueryResultPtr ChdbClient::executeStreamingIterate(void * streaming_result
             }
         }
 
-        // Check if query should end based on result type
+        /// Check if query should end based on result type
         bool is_end = !res->getError().empty() || is_canceled || res->isEmpty();
         if (is_end)
         {
diff --git a/programs/local/EmbeddedServer.cpp b/programs/local/EmbeddedServer.cpp
@@ -74,8 +74,6 @@
 #    include <azure/storage/common/internal/xml_wrapper.hpp>
 #endif
 
-bool chdb_embedded_server_initialized = false;
-
 namespace fs = std::filesystem;
 
 namespace CurrentMetrics
@@ -513,8 +511,6 @@ try
         global_register_once_flag,
         []()
         {
-            chdb_embedded_server_initialized = true;
-
             registerInterpreters();
             /// Don't initialize DateLUT
             registerFunctions();
diff --git a/programs/local/LocalChdb.cpp b/programs/local/LocalChdb.cpp
@@ -3,7 +3,6 @@
 #include "PandasDataFrameBuilder.h"
 #include "ChunkCollectorOutputFormat.h"
 #include "PythonImporter.h"
-#include "PythonTableCache.h"
 #include "StoragePython.h"
 
 #include <pybind11/detail/non_limited_api.h>
@@ -274,13 +273,14 @@ query_result * connection_wrapper::query(const std::string & query_str, const st
 
     auto * result = chdb_query_n(*conn, query_str.data(), query_str.size(), format.data(), format.size());
 
-    auto error_msg = CHDB::chdb_result_error_string(result);
+    const auto & error_msg = CHDB::chdb_result_error_string(result);
     if (!error_msg.empty())
     {
         std::string msg_copy(error_msg);
         chdb_destroy_query_result(result);
         throw std::runtime_error(msg_copy);
     }
+
     return new query_result(result, false);
 }
 
@@ -298,7 +298,7 @@ py::object connection_wrapper::query_df(const std::string & query_str)
 
         result = chdb_query_n(*conn, query_str.data(), query_str.size(), format.data(), format.size());
 
-        auto error_msg = CHDB::chdb_result_error_string(result);
+        const auto & error_msg = CHDB::chdb_result_error_string(result);
         if (!error_msg.empty())
         {
             std::string msg_copy(error_msg);
@@ -322,7 +322,7 @@ streaming_query_result * connection_wrapper::send_query(const std::string & quer
     CHDB::cachePythonTablesFromQuery(reinterpret_cast<chdb_conn *>(*conn), query_str);
     py::gil_scoped_release release;
     auto * result = chdb_stream_query_n(*conn, query_str.data(), query_str.size(), format.data(), format.size());
-    auto error_msg = CHDB::chdb_result_error_string(result);
+    const auto & error_msg = CHDB::chdb_result_error_string(result);
     if (!error_msg.empty())
     {
         std::string msg_copy(error_msg);
@@ -342,7 +342,7 @@ query_result * connection_wrapper::streaming_fetch_result(streaming_query_result
 
     auto * result  = chdb_stream_fetch_result(*conn, streaming_result->get_result());
 
-    const auto error_msg = CHDB::chdb_result_error_string(result);
+    const auto & error_msg = CHDB::chdb_result_error_string(result);
     if (!error_msg.empty())
     {
         std::string msg_copy(error_msg);
@@ -359,30 +359,29 @@ py::object connection_wrapper::streaming_fetch_df(streaming_query_result * strea
         return py::none();
 
     chdb_result * result = nullptr;
-    CHDB::ChunkQueryResult * chunk_result = nullptr;
+    CHDB::DataFrameQueryResult * chunk_result = nullptr;
 
     {
         py::gil_scoped_release release;
 
-        result  = chdb_stream_fetch_result(*conn, streaming_result->get_result());
+        result = chdb_stream_fetch_result(*conn, streaming_result->get_result());
 
-        auto error_msg = CHDB::chdb_result_error_string(result);
+        const auto & error_msg = CHDB::chdb_result_error_string(result);
         if (!error_msg.empty())
         {
             std::string msg_copy(error_msg);
             chdb_destroy_query_result(result);
             throw std::runtime_error(msg_copy);
         }
 
-        if (!(chunk_result = dynamic_cast<CHDB::ChunkQueryResult *>(reinterpret_cast<CHDB::QueryResult*>(result))))
-            throw std::runtime_error("Expected ChunkQueryResult for dataframe format");
+        if (!(chunk_result = dynamic_cast<CHDB::DataFrameQueryResult *>(reinterpret_cast<CHDB::QueryResult*>(result))))
+            throw std::runtime_error("Expected DataFrameQueryResult for dataframe format");
     }
 
-    CHDB::PandasDataFrameBuilder builder(*chunk_result);
-    auto df = builder.getDataFrame();
+    py::handle df_handle = chunk_result->dataframe;
     chdb_destroy_query_result(result);
 
-    return df;
+    return py::reinterpret_steal<py::object>(df_handle);
 }
 
 void connection_wrapper::streaming_cancel_query(streaming_query_result * streaming_result)
diff --git a/programs/local/LocalServer.cpp b/programs/local/LocalServer.cpp
@@ -73,8 +73,6 @@
 #   include <azure/storage/common/internal/xml_wrapper.hpp>
 #endif
 
-extern bool chdb_embedded_server_initialized;
-
 namespace fs = std::filesystem;
 
 namespace CurrentMetrics
@@ -650,8 +648,6 @@ try
         global_register_once_flag,
         []()
         {
-            chdb_embedded_server_initialized = true;
-
             registerInterpreters();
             /// Don't initialize DateLUT
             registerFunctions();
diff --git a/programs/local/QueryResult.h b/programs/local/QueryResult.h
@@ -9,6 +9,8 @@
 
 #if USE_PYTHON
 #include <Processors/Chunk.h>
+#include <pybind11/pybind11.h>
+namespace py = pybind11;
 namespace DB
 {
     class Block;
@@ -23,7 +25,8 @@ enum class QueryResultType : uint8_t
     RESULT_TYPE_MATERIALIZED = 0,
     RESULT_TYPE_STREAMING = 1,
     RESULT_TYPE_CHUNK = 2,
-    RESULT_TYPE_NONE = 3
+    RESULT_TYPE_DATAFRAME = 3,
+    RESULT_TYPE_NONE = 4
 };
 
 class QueryResult
@@ -144,13 +147,34 @@ class ChunkQueryResult : public QueryResult
     uint64_t storage_rows_read;
     uint64_t storage_bytes_read;
 };
+
+class DataFrameQueryResult : public QueryResult
+{
+public:
+    explicit DataFrameQueryResult(
+        py::handle dataframe_,
+        uint64_t rows_read)
+        : QueryResult(QueryResultType::RESULT_TYPE_DATAFRAME),
+        dataframe(dataframe_),
+        is_empty(rows_read == 0)
+    {}
+
+    bool isEmpty() const override
+    {
+        return is_empty;
+    }
+
+    py::handle dataframe;
+    bool is_empty;
+};
 #endif
 
 using QueryResultPtr = std::unique_ptr<QueryResult>;
 using MaterializedQueryResultPtr = std::unique_ptr<MaterializedQueryResult>;
 using StreamQueryResultPtr = std::unique_ptr<StreamQueryResult>;
 #if USE_PYTHON
 using ChunkQueryResultPtr = std::unique_ptr<ChunkQueryResult>;
+using DataFrameQueryResultPtr = std::unique_ptr<DataFrameQueryResult>;
 #endif
 
 } // namespace CHDB
diff --git a/programs/local/chdb-arrow.cpp b/programs/local/chdb-arrow.cpp
@@ -134,13 +134,17 @@ chdb_state chdb_arrow_scan(
     chdb_connection conn, const char * table_name,
     chdb_arrow_stream arrow_stream)
 {
+    CHDB::ChdbMemoryTrackingGuard guard;
+
     return chdb_inner_arrow_scan(conn, table_name, arrow_stream, false);
 }
 
 chdb_state chdb_arrow_array_scan(
     chdb_connection conn, const char * table_name,
     chdb_arrow_schema arrow_schema, chdb_arrow_array arrow_array)
 {
+    CHDB::ChdbMemoryTrackingGuard guard;
+
     auto * private_data = new CHDB::PrivateData();
 	private_data->schema = reinterpret_cast<ArrowSchema *>(arrow_schema);
 	private_data->array = reinterpret_cast<ArrowArray *>(arrow_array);
@@ -158,6 +162,8 @@ chdb_state chdb_arrow_array_scan(
 
 chdb_state chdb_arrow_unregister_table(chdb_connection conn, const char * table_name)
 {
+    CHDB::ChdbMemoryTrackingGuard guard;
+
     if (!table_name)
         return CHDBError;
 
diff --git a/programs/local/chdb-internal.h b/programs/local/chdb-internal.h
@@ -18,9 +18,29 @@ inline bool checkConnectionValidity(chdb_conn * connection)
     return connection && connection->connected;
 }
 
+extern thread_local bool chdb_memory_tracking;
+
 namespace CHDB
 {
 
+class ChdbMemoryTrackingGuard
+{
+public:
+    ChdbMemoryTrackingGuard()
+        : previous_value(chdb_memory_tracking)
+    {
+        chdb_memory_tracking = true;
+    }
+
+    ~ChdbMemoryTrackingGuard()
+    {
+        chdb_memory_tracking = previous_value;
+    }
+
+private:
+    bool previous_value;
+};
+
 std::unique_ptr<MaterializedQueryResult> pyEntryClickHouseLocal(int argc, char ** argv);
 
 const std::string & chdb_result_error_string(chdb_result * result);
diff --git a/programs/local/chdb.cpp b/programs/local/chdb.cpp
diff --git a/src/Common/CurrentMemoryTracker.cpp b/src/Common/CurrentMemoryTracker.cpp
diff --git a/tests/test_streaming_query.py b/tests/test_streaming_query.py

Original file line number	Diff line number	Diff line change
`@@ -134,13 +134,17 @@ chdb_state chdb_arrow_scan(`
`134`	`134`	`chdb_connection conn, const char * table_name,`
`135`	`135`	`chdb_arrow_stream arrow_stream)`
`136`	`136`	`{`
	`137`	`+ CHDB::ChdbMemoryTrackingGuard guard;`
	`138`	`+`
`137`	`139`	`return chdb_inner_arrow_scan(conn, table_name, arrow_stream, false);`
`138`	`140`	`}`
`139`	`141`
`140`	`142`	`chdb_state chdb_arrow_array_scan(`
`141`	`143`	`chdb_connection conn, const char * table_name,`
`142`	`144`	`chdb_arrow_schema arrow_schema, chdb_arrow_array arrow_array)`
`143`	`145`	`{`
	`146`	`+ CHDB::ChdbMemoryTrackingGuard guard;`
	`147`	`+`
`144`	`148`	`auto * private_data = new CHDB::PrivateData();`
`145`	`149`	`private_data->schema = reinterpret_cast<ArrowSchema *>(arrow_schema);`
`146`	`150`	`private_data->array = reinterpret_cast<ArrowArray *>(arrow_array);`
`@@ -158,6 +162,8 @@ chdb_state chdb_arrow_array_scan(`
`158`	`162`
`159`	`163`	`chdb_state chdb_arrow_unregister_table(chdb_connection conn, const char * table_name)`
`160`	`164`	`{`
	`165`	`+ CHDB::ChdbMemoryTrackingGuard guard;`
	`166`	`+`
`161`	`167`	`if (!table_name)`
`162`	`168`	`return CHDBError;`
`163`	`169`