Merge pull request #326 from wudidapaopao/yuxiaozhe-dev

Add support for streaming query in chDB

Author: auxten

README-zh.md

@@ -147,6 +147,55 @@ print(query("select sum_udf(12,22)"))
 参见: [test_udf.py](tests/test_udf.py).
 </details>
 
+<details>
+    <summary><h4>🗂️ 流式查询</h4></summary>
+
+通过分块流式处理大数据集，保持内存使用恒定。
+
+```python
+from chdb import session as chs
+
+sess = chs.Session()
+
+# 示例1：流式查询基础用法
+rows_cnt = 0
+with sess.send_query("SELECT * FROM numbers(200000)", "CSV") as stream_result:
+    for chunk in stream_result:
+        rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 200000
+
+# 示例2：使用fetch()手动迭代
+rows_cnt = 0
+stream_result = sess.send_query("SELECT * FROM numbers(200000)", "CSV")
+while True:
+    chunk = stream_result.fetch()
+    if chunk is None:
+        break
+    rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 200000
+
+# 示例3：提前取消查询
+rows_cnt = 0
+stream_result = sess.send_query("SELECT * FROM numbers(200000)", "CSV")
+while True:
+    chunk = stream_result.fetch()
+    if chunk is None:
+        break
+    if rows_cnt > 0:
+        stream_result.cancel()
+        break
+    rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 65409
+
+sess.close()
+```
+
+参见: [test_streaming_query.py](tests/test_streaming_query.py)。
+</details>
+
 更多示例，请参见 [examples](examples) 和 [tests](tests)。
 
 ## 演示和示例

README.md

@@ -223,6 +223,56 @@ see also: [test_udf.py](tests/test_udf.py).
 </details>
 
 
+<details>
+    <summary><h4>🗂️ Streaming Query</h4></summary>
+
+Process large datasets with constant memory usage through chunked streaming.
+
+```python
+from chdb import session as chs
+
+sess = chs.Session()
+
+# Example 1: Basic example of using streaming query
+rows_cnt = 0
+with sess.send_query("SELECT * FROM numbers(200000)", "CSV") as stream_result:
+    for chunk in stream_result:
+        rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 200000
+
+# Example 2: Manual iteration with fetch()
+rows_cnt = 0
+stream_result = sess.send_query("SELECT * FROM numbers(200000)", "CSV")
+while True:
+    chunk = stream_result.fetch()
+    if chunk is None:
+        break
+    rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 200000
+
+# Example 3: Early cancellation demo
+rows_cnt = 0
+stream_result = sess.send_query("SELECT * FROM numbers(200000)", "CSV")
+while True:
+    chunk = stream_result.fetch()
+    if chunk is None:
+        break
+    if rows_cnt > 0:
+        stream_result.cancel()
+        break
+    rows_cnt += chunk.rows_read()
+
+print(rows_cnt) # 65409
+
+sess.close()
+```
+
+For more details, see [test_streaming_query.py](tests/test_streaming_query.py).
+</details>
+
+
 <details>
     <summary><h4>🗂️ Python Table Engine</h4></summary>
 

chdb/session/state.py

@@ -4,7 +4,7 @@
 
 import chdb
 from ..state import sqlitelike as chdb_stateful
-
+from ..state.sqlitelike import StreamingResult
 
 g_session = None
 g_session_path = None
@@ -120,3 +120,16 @@ def query(self, sql, fmt="CSV", udf_path=""):
 
     # alias sql = query
     sql = query
+
+    def send_query(self, sql, fmt="CSV") -> StreamingResult:
+        """
+        Execute a streaming query.
+        """
+        if fmt == "Debug":
+            warnings.warn(
+                """Debug format is not supported in Session.query
+Please try use parameters in connection string instead:
+Eg: conn = connect(f"db_path?verbose&log-level=test")"""
+            )
+            fmt = "CSV"
+        return self._conn.send_query(sql, fmt)

chdb/state/sqlitelike.py

@@ -40,6 +40,57 @@ def to_df(r):
     return t.to_pandas(use_threads=True)
 
 
+class StreamingResult:
+    def __init__(self, c_result, conn, result_func):
+        self._result = c_result
+        self._result_func = result_func
+        self._conn = conn
+        self._exhausted = False
+
+    def fetch(self):
+        """Fetch next chunk of streaming results"""
+        if self._exhausted:
+            return None
+
+        try:
+            result = self._conn.streaming_fetch_result(self._result)
+            if result is None or result.rows_read() == 0:
+                self._exhausted = True
+                return None
+            return self._result_func(result)
+        except Exception as e:
+            self._exhausted = True
+            raise RuntimeError(f"Streaming query failed: {str(e)}") from e
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._exhausted:
+            raise StopIteration
+
+        chunk = self.fetch()
+        if chunk is None:
+            self._exhausted = True
+            raise StopIteration
+
+        return chunk
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def cancel(self):
+        self._exhausted = True
+
+        try:
+            self._conn.streaming_cancel_query(self._result)
+        except Exception as e:
+            raise RuntimeError(f"Failed to cancel streaming query: {str(e)}") from e
+
+
 class Connection:
     def __init__(self, connection_string: str):
         # print("Connection", connection_string)
@@ -59,6 +110,15 @@ def query(self, query: str, format: str = "CSV") -> Any:
         result = self._conn.query(query, format)
         return result_func(result)
 
+    def send_query(self, query: str, format: str = "CSV") -> StreamingResult:
+        lower_output_format = format.lower()
+        result_func = _process_result_format_funs.get(lower_output_format, lambda x: x)
+        if lower_output_format in _arrow_format:
+            format = "Arrow"
+
+        c_stream_result = self._conn.send_query(query, format)
+        return StreamingResult(c_stream_result, self._conn, result_func)
+
     def close(self) -> None:
         # print("close")
         if self._cursor:

programs/local/LocalChdb.cpp

@@ -315,11 +315,62 @@ query_result * connection_wrapper::query(const std::string & query_str, const st
     }
     if (result->error_message)
     {
-        throw std::runtime_error(result->error_message);
+        std::string msg_copy(result->error_message);
+        free_result_v2(result);
+        throw std::runtime_error(msg_copy);
     }
     return new query_result(result, false);
 }
 
+streaming_query_result * connection_wrapper::send_query(const std::string & query_str, const std::string & format)
+{
+    global_query_obj = findQueryableObjFromQuery(query_str);
+
+    py::gil_scoped_release release;
+    auto * result = query_conn_streaming(*conn, query_str.c_str(), format.c_str());
+    const auto * error_msg = chdb_streaming_result_error(result);
+    if (error_msg)
+    {
+        std::string msg_copy(error_msg);
+        chdb_destroy_result(result);
+        throw std::runtime_error(msg_copy);
+    }
+
+    return new streaming_query_result(result);
+}
+
+query_result * connection_wrapper::streaming_fetch_result(streaming_query_result * streaming_result)
+{
+    py::gil_scoped_release release;
+
+    if (!streaming_result || !streaming_result->get_result())
+        return nullptr;
+
+    auto * result  = chdb_streaming_fetch_result(*conn, streaming_result->get_result());
+
+    if (result->len == 0)
+        LOG_DEBUG(getLogger("CHDB"), "Empty result returned for streaming query");
+
+    if (result->error_message)
+    {
+        std::string msg_copy(result->error_message);
+        free_result_v2(result);
+        throw std::runtime_error(msg_copy);
+    }
+
+    return new query_result(result, false);
+}
+
+void connection_wrapper::streaming_cancel_query(streaming_query_result * streaming_result)
+{
+    py::gil_scoped_release release;
+
+    if (!streaming_result || !streaming_result->get_result())
+        return;
+
+    chdb_streaming_cancel_query(*conn, streaming_result->get_result());
+}
+
 void cursor_wrapper::execute(const std::string & query_str)
 {
     release_result();
@@ -407,6 +458,11 @@ PYBIND11_MODULE(_chdb, m)
         .def("has_error", &query_result::has_error)
         .def("error_message", &query_result::error_message);
 
+    py::class_<streaming_query_result>(m, "streaming_query_result")
+        .def(py::init<chdb_streaming_result *>(), py::return_value_policy::take_ownership)
+        .def("has_error", &streaming_query_result::has_error)
+        .def("error_message", &streaming_query_result::error_message);
+
     py::class_<DB::PyReader, std::shared_ptr<DB::PyReader>>(m, "PyReader")
         .def(
             py::init<const py::object &>(),
@@ -460,7 +516,23 @@ PYBIND11_MODULE(_chdb, m)
             &connection_wrapper::query,
             py::arg("query_str"),
             py::arg("format") = "CSV",
-            "Execute a query and return a query_result object");
+            "Execute a query and return a query_result object")
+        .def(
+            "send_query",
+            &connection_wrapper::send_query,
+            py::arg("query_str"),
+            py::arg("format") = "CSV",
+            "Send a streaming query and return a streaming query result object")
+        .def(
+            "streaming_fetch_result",
+            &connection_wrapper::streaming_fetch_result,
+                py::arg("streaming_result"),
+                "Fetches a data chunk from the streaming result. This function should be called repeatedly until the result is exhausted")
+        .def(
+            "streaming_cancel_query",
+            &connection_wrapper::streaming_cancel_query,
+            py::arg("streaming_result"),
+            "Cancel a streaming query");
 
     m.def(
         "query",

programs/local/LocalChdb.h

@@ -25,6 +25,7 @@ class __attribute__((visibility("default"))) connection_wrapper;
 class __attribute__((visibility("default"))) cursor_wrapper;
 class __attribute__((visibility("default"))) memoryview_wrapper;
 class __attribute__((visibility("default"))) query_result;
+class __attribute__((visibility("default"))) streaming_query_result;
 
 class connection_wrapper
 {
@@ -42,6 +43,9 @@ class connection_wrapper
     void commit();
     void close();
     query_result * query(const std::string & query_str, const std::string & format = "CSV");
+    streaming_query_result * send_query(const std::string & query_str, const std::string & format = "CSV");
+    query_result * streaming_fetch_result(streaming_query_result * streaming_result);
+    void streaming_cancel_query(streaming_query_result * streaming_result);
 
     // Move the private methods declarations here
     std::pair<std::string, std::map<std::string, std::string>> parse_connection_string(const std::string & conn_str);
@@ -169,6 +173,22 @@ class query_result
     memoryview_wrapper * get_memview();
 };
 
+class streaming_query_result
+{
+private:
+    chdb_streaming_result * result;
+
+public:
+    streaming_query_result(chdb_streaming_result * result_) : result(result_) {}
+    ~streaming_query_result()
+    {
+        chdb_destroy_result(result);
+    }
+    bool has_error() { return chdb_streaming_result_error(result) != nullptr; }
+    py::str error_message() { return chdb_streaming_result_error(result); }
+    chdb_streaming_result * get_result() { return result; }
+};
+
 class memoryview_wrapper
 {
 private: