From 7af80b114c0ef66a0ee118ec8dfcc27a39400747 Mon Sep 17 00:00:00 2001
From: sundyli <543950155@qq.com>
Date: Mon, 13 Jan 2025 13:00:40 +0800
Subject: [PATCH] fix(query): increase state_rows in copy agg state rows
 (#17252)

* fix(query): increase state_rows in copy agg state rows

* fix(query): increase state_rows in copy agg state rows

* fix(query): increase state_rows in copy agg state rows
---
 src/query/expression/src/aggregate/payload.rs |  1 +
 .../20+_others/20_0022_agg_memory.result      |  6 +++
 .../20+_others/20_0022_agg_memory.sh          | 46 +++++++++++++++++++
 3 files changed, 53 insertions(+)
 create mode 100755 tests/suites/0_stateless/20+_others/20_0022_agg_memory.result
 create mode 100755 tests/suites/0_stateless/20+_others/20_0022_agg_memory.sh

diff --git a/src/query/expression/src/aggregate/payload.rs b/src/query/expression/src/aggregate/payload.rs
index 504699715e466..c3cde64a989f4 100644
--- a/src/query/expression/src/aggregate/payload.rs
+++ b/src/query/expression/src/aggregate/payload.rs
@@ -339,6 +339,7 @@ impl Payload {
                 )
             }
             page.rows += 1;
+            page.state_rows += 1;
 
             if page.rows == page.capacity {
                 page = self.writable_page();
diff --git a/tests/suites/0_stateless/20+_others/20_0022_agg_memory.result b/tests/suites/0_stateless/20+_others/20_0022_agg_memory.result
new file mode 100755
index 0000000000000..f32baf89a114b
--- /dev/null
+++ b/tests/suites/0_stateless/20+_others/20_0022_agg_memory.result
@@ -0,0 +1,6 @@
+executing 1
+executing 2
+executing 3
+executing 4
+executing 5
+Memory usage difference is less than 5%
diff --git a/tests/suites/0_stateless/20+_others/20_0022_agg_memory.sh b/tests/suites/0_stateless/20+_others/20_0022_agg_memory.sh
new file mode 100755
index 0000000000000..e8602a4ae66a0
--- /dev/null
+++ b/tests/suites/0_stateless/20+_others/20_0022_agg_memory.sh
@@ -0,0 +1,46 @@
+#!/usr/bin/env bash
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+
+## warmup
+for i in `seq 1 2`;do
+	$BENDSQL_CLIENT_CONNECT --query="""
+	select number::string , max(number::string), min(number::string), count(distinct number) from numbers(10000000) group by 1 ignore_result;
+	"""
+done
+
+
+PIDS=($(pgrep databend-query))
+BEFORE_MEM=0
+for PID in "${PIDS[@]}"; do
+    MEM=$(ps -o rss= -p $PID | tail -n 1)
+    BEFORE_MEM=$((BEFORE_MEM + MEM))
+done
+
+
+for i in `seq 1 5`;do
+	echo "executing $i"
+	$BENDSQL_CLIENT_CONNECT --query="""
+	select number::string , max(number::string), min(number::string), count(distinct number) from numbers(10000000) group by 1 ignore_result;
+	"""
+done
+
+sleep 15
+
+
+AFTER_MEM=0
+for PID in "${PIDS[@]}"; do
+    MEM=$(ps -o rss= -p $PID | tail -n 1)
+    AFTER_MEM=$((AFTER_MEM + MEM))
+done
+
+# Calculate the difference in percentage
+DIFF=$(awk -v before=$BEFORE_MEM -v after=$AFTER_MEM 'BEGIN {print (after-before)/before * 100}')
+
+# Check if the difference is less than 5%
+if (( $(awk -v diff=$DIFF 'BEGIN {print (diff < 5)}') )); then
+    echo "Memory usage difference is less than 5%"
+else
+    echo "Memory usage difference is greater than 5%, before ${BEFORE_MEM} ${AFTER_MEM}"
+fi
\ No newline at end of file