executor: reset groupCount when Close is called in streamagg execut…

…or (#53874) close #53867
pingcap · Jun 18, 2024 · a5a1b67 · a5a1b67
1 parent e47330c
commit a5a1b67
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/pkg/executor/internal/vecgroupchecker/BUILD.bazel b/pkg/executor/internal/vecgroupchecker/BUILD.bazel
@@ -22,7 +22,7 @@ go_test(
     ],
     embed = [":vecgroupchecker"],
     flaky = True,
-    shard_count = 3,
+    shard_count = 4,
     deps = [
         "//pkg/config",
         "//pkg/expression",

diff --git a/pkg/executor/internal/vecgroupchecker/vec_group_checker.go b/pkg/executor/internal/vecgroupchecker/vec_group_checker.go
@@ -496,6 +496,7 @@ func (e *VecGroupChecker) IsExhausted() bool {
 func (e *VecGroupChecker) Reset() {
 	if e.groupOffset != nil {
 		e.groupOffset = e.groupOffset[:0]
+		e.groupCount = 0
 	}
 	if e.sameGroup != nil {
 		e.sameGroup = e.sameGroup[:0]

diff --git a/pkg/executor/internal/vecgroupchecker/vec_group_checker_test.go b/pkg/executor/internal/vecgroupchecker/vec_group_checker_test.go
@@ -268,3 +268,13 @@ func TestVecGroupChecker(t *testing.T) {
 	require.Equal(t, e, 3)
 	require.True(t, groupChecker.IsExhausted())
 }
+
+func TestIssue53867(t *testing.T) {
+	checker := NewVecGroupChecker(nil, true, nil)
+	checker.groupOffset = make([]int, 20)
+	checker.nextGroupID = 10
+	checker.groupCount = 15
+	require.False(t, checker.IsExhausted())
+	checker.Reset()
+	require.True(t, checker.IsExhausted())
+}