From 7af6bf833cfcb95ee51a0371d1bd7f411685d1c9 Mon Sep 17 00:00:00 2001
From: Yang Xiufeng <yangxiufeng.c@gmail.com>
Date: Tue, 21 Jan 2025 15:26:00 +0800
Subject: [PATCH] fix: support returning rows larger than 10MB. (#17344)

---
 .../src/servers/http/v1/query/page_manager.rs        |  9 ++++++++-
 .../09_http_handler/09_0004_large_row.result         |  3 +++
 .../1_stateful/09_http_handler/09_0004_large_row.sh  | 12 ++++++++++++
 3 files changed, 23 insertions(+), 1 deletion(-)
 create mode 100755 tests/suites/1_stateful/09_http_handler/09_0004_large_row.result
 create mode 100755 tests/suites/1_stateful/09_http_handler/09_0004_large_row.sh

diff --git a/src/query/service/src/servers/http/v1/query/page_manager.rs b/src/query/service/src/servers/http/v1/query/page_manager.rs
index bd9909a84d5f5..4da62af8e4d35 100644
--- a/src/query/service/src/servers/http/v1/query/page_manager.rs
+++ b/src/query/service/src/servers/http/v1/query/page_manager.rs
@@ -143,6 +143,9 @@ impl PageManager {
                 i += 1;
             } else {
                 *remain_size = 0;
+                if res.is_empty() && i == 0 {
+                    i += 1
+                }
             }
         }
         res.extend_from_slice(&rows[..i]);
@@ -163,7 +166,11 @@ impl PageManager {
                     remain_size -= size;
                     remain_rows -= 1;
                 } else {
-                    self.row_buffer.push_front(row);
+                    if res.is_empty() {
+                        res.push(row);
+                    } else {
+                        self.row_buffer.push_front(row);
+                    }
                     remain_size = 0;
                 }
             } else {
diff --git a/tests/suites/1_stateful/09_http_handler/09_0004_large_row.result b/tests/suites/1_stateful/09_http_handler/09_0004_large_row.result
new file mode 100755
index 0000000000000..052f39dfb8524
--- /dev/null
+++ b/tests/suites/1_stateful/09_http_handler/09_0004_large_row.result
@@ -0,0 +1,3 @@
+1
+1
+"/v1/query/09_004/final"
diff --git a/tests/suites/1_stateful/09_http_handler/09_0004_large_row.sh b/tests/suites/1_stateful/09_http_handler/09_0004_large_row.sh
new file mode 100755
index 0000000000000..f1e5489fabf45
--- /dev/null
+++ b/tests/suites/1_stateful/09_http_handler/09_0004_large_row.sh
@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+qid="09_004"
+curl -s --header 'Content-Type: application/json' --header "X-DATABEND-QUERY-ID: $qid" --request POST '127.0.0.1:8000/v1/query/' \
+ --data-raw $'{"sql": "select json_array_agg(json_object(\'num\',number)), (number % 2) as s from numbers(2000000) group by s;", "pagination": { "wait_time_secs": 5}}' \
+ -u root: | jq '.data | length'
+
+curl -s --header 'Content-Type: application/json'  "127.0.0.1:8000/v1/query/$qid/page/1" \
+ -u root: | jq '(.data | length), .next_uri'