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

…53285) close #52933
pingcap · Jun 8, 2024 · 5722930 · 5722930
1 parent eb27609
commit 5722930
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 23 deletions.
diff --git a/executor/window_test.go b/executor/window_test.go
@@ -498,3 +498,12 @@ func TestIssue29947(t *testing.T) {
 	result.Check(testkit.Rows("2", "3"))
 	tk.MustExec("commit")
 }
+
+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/parser/parser.go b/parser/parser.go
@@ -16719,7 +16719,11 @@ yynewstate:
 		}
 	case 1437:
 		{
-			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{F: 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 1438:
 		{

diff --git a/parser/parser.y b/parser/parser.y
@@ -50,10 +50,11 @@ import (
 %token	<ident>
 
 	/*yy:token "%c"     */
-	identifier           "identifier"
-	asof                 "AS OF"
-	toTimestamp          "TO TIMESTAMP"
-	toTSO                "TO TSO"
+	identifier  "identifier"
+	asof        "AS OF"
+	toTimestamp "TO TIMESTAMP"
+	toTSO       "TO TSO"
+
 	/*yy:token "_%c"    */
 	underscoreCS "UNDERSCORE_CHARSET"
 
@@ -2621,44 +2622,44 @@ FlashbackToTimestampStmt:
 	"FLASHBACK" "CLUSTER" toTimestamp stringLit
 	{
 		$$ = &ast.FlashBackToTimestampStmt{
-			FlashbackTS: ast.NewValueExpr($4, "", ""),
+			FlashbackTS:  ast.NewValueExpr($4, "", ""),
 			FlashbackTSO: 0,
 		}
 	}
 |	"FLASHBACK" "TABLE" TableNameList toTimestamp stringLit
 	{
 		$$ = &ast.FlashBackToTimestampStmt{
-			Tables:      $3.([]*ast.TableName),
-			FlashbackTS: ast.NewValueExpr($5, "", ""),
+			Tables:       $3.([]*ast.TableName),
+			FlashbackTS:  ast.NewValueExpr($5, "", ""),
 			FlashbackTSO: 0,
 		}
 	}
 |	"FLASHBACK" DatabaseSym DBName toTimestamp stringLit
 	{
 		$$ = &ast.FlashBackToTimestampStmt{
-			DBName:      model.NewCIStr($3),
-			FlashbackTS: ast.NewValueExpr($5, "", ""),
+			DBName:       model.NewCIStr($3),
+			FlashbackTS:  ast.NewValueExpr($5, "", ""),
 			FlashbackTSO: 0,
 		}
 	}
 |	"FLASHBACK" "CLUSTER" toTSO LengthNum
 	{
 		if tsoValue, ok := $4.(uint64); ok && tsoValue > 0 {
 			$$ = &ast.FlashBackToTimestampStmt{
-        		FlashbackTSO: tsoValue,
-        	}
+				FlashbackTSO: tsoValue,
+			}
 		} else {
-    		yylex.AppendError(yylex.Errorf("Invalid TSO value provided: %d", $4))
-    		return 1
+			yylex.AppendError(yylex.Errorf("Invalid TSO value provided: %d", $4))
+			return 1
 		}
 	}
 |	"FLASHBACK" "TABLE" TableNameList toTSO LengthNum
 	{
 		if tsoValue, ok := $5.(uint64); ok && tsoValue > 0 {
 			$$ = &ast.FlashBackToTimestampStmt{
-            	Tables:      $3.([]*ast.TableName),
-            	FlashbackTSO: tsoValue,
-            }
+				Tables:       $3.([]*ast.TableName),
+				FlashbackTSO: tsoValue,
+			}
 		} else {
 			yylex.AppendError(yylex.Errorf("Invalid TSO value provided: %d", $5))
 			return 1
@@ -2668,16 +2669,15 @@ FlashbackToTimestampStmt:
 	{
 		if tsoValue, ok := $5.(uint64); ok && tsoValue > 0 {
 			$$ = &ast.FlashBackToTimestampStmt{
-            	DBName:      model.NewCIStr($3),
-            	FlashbackTSO: tsoValue,
+				DBName:       model.NewCIStr($3),
+				FlashbackTSO: tsoValue,
 			}
 		} else {
 			yylex.AppendError(yylex.Errorf("Invalid TSO value provided: %d", $5))
 			return 1
 		}
 	}
 
-
 /*******************************************************************
  *
  *  Flush Back Table Statement
@@ -7868,7 +7868,11 @@ SumExpr:
 	}
 |	builtinVarSamp '(' BuggyDefaultFalseDistinctOpt Expression ')' OptWindowingClause
 	{
-		$$ = &ast.AggregateFuncExpr{F: $1, Args: []ast.ExprNode{$4}, Distinct: $3.(bool)}
+		if $6 != nil {
+			$$ = &ast.WindowFuncExpr{F: $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
 	{
@@ -9772,8 +9776,8 @@ SetOprStmtWoutLimitOrderBy:
 			setOprList2 = []ast.Node{x}
 			with2 = x.With
 		case *ast.SetOprStmt:
-		    // child setOprStmt's limit and order should also make sense
-		    // we should separate it out from other normal SetOprSelectList.
+			// child setOprStmt's limit and order should also make sense
+			// we should separate it out from other normal SetOprSelectList.
 			setOprList2 = x.SelectList.Selects
 			with2 = x.With
 			limit2 = x.Limit