planner: make var_samp can be used as a window function (#53130)

close #52933
pingcap · May 15, 2024 · 2d552c0 · 2d552c0
1 parent 6aef624
commit 2d552c0
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 2 deletions.
diff --git a/pkg/executor/window_test.go b/pkg/executor/window_test.go
@@ -470,3 +470,12 @@ func TestIssue45964And46050(t *testing.T) {
 	testReturnColumnNullableAttribute(tk, "cume_dist()", false)
 	testReturnColumnNullableAttribute(tk, "percent_rank()", false)
 }
+
+func TestVarSampAsAWindowFunction(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1 (c1 int)")
+	tk.MustExec("select var_samp(c1) from t1")
+	tk.MustExec("select c1, var_samp(c1) over (partition by c1) from t1")
+}
diff --git a/pkg/parser/parser.go b/pkg/parser/parser.go
@@ -18309,7 +18309,11 @@ yynewstate:
 		}
 	case 1539:
 		{
-			parser.yyVAL.expr = &ast.AggregateFuncExpr{F: yyS[yypt-5].ident, Args: []ast.ExprNode{yyS[yypt-2].expr}, Distinct: yyS[yypt-3].item.(bool)}
+			if yyS[yypt-0].item != nil {
+				parser.yyVAL.expr = &ast.WindowFuncExpr{Name: yyS[yypt-5].ident, Args: []ast.ExprNode{yyS[yypt-2].expr}, Distinct: yyS[yypt-3].item.(bool), Spec: *(yyS[yypt-0].item.(*ast.WindowSpec))}
+			} else {
+				parser.yyVAL.expr = &ast.AggregateFuncExpr{F: yyS[yypt-5].ident, Args: []ast.ExprNode{yyS[yypt-2].expr}, Distinct: yyS[yypt-3].item.(bool)}
+			}
 		}
 	case 1540:
 		{

diff --git a/pkg/parser/parser.y b/pkg/parser/parser.y
@@ -8364,7 +8364,11 @@ SumExpr:
 	}
 |	builtinVarSamp '(' BuggyDefaultFalseDistinctOpt Expression ')' OptWindowingClause
 	{
-		$$ = &ast.AggregateFuncExpr{F: $1, Args: []ast.ExprNode{$4}, Distinct: $3.(bool)}
+		if $6 != nil {
+			$$ = &ast.WindowFuncExpr{Name: $1, Args: []ast.ExprNode{$4}, Distinct: $3.(bool), Spec: *($6.(*ast.WindowSpec))}
+		} else {
+			$$ = &ast.AggregateFuncExpr{F: $1, Args: []ast.ExprNode{$4}, Distinct: $3.(bool)}
+		}
 	}
 |	"JSON_ARRAYAGG" '(' Expression ')' OptWindowingClause
 	{