parser: implement Restore for FrameClause, PartitionByClause, WindowS…

…pec and WindowFuncExpr (pingcap#143)
tiancaiamao · Jan 5, 2019 · 1252487 · 1252487
1 parent d7bae25
commit 1252487
Show file tree

Hide file tree

Showing 4 changed files with 155 additions and 4 deletions.
diff --git a/ast/dml.go b/ast/dml.go
@@ -1399,7 +1399,39 @@ type WindowSpec struct {
 
 // Restore implements Node interface.
 func (n *WindowSpec) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	if name := n.Name.String(); name != "" {
+		ctx.WriteName(name)
+		ctx.WriteKeyWord(" AS ")
+	}
+	ctx.WritePlain("(")
+	sep := ""
+	if refName := n.Ref.String(); refName != "" {
+		ctx.WriteName(refName)
+		sep = " "
+	}
+	if n.PartitionBy != nil {
+		ctx.WritePlain(sep)
+		if err := n.PartitionBy.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore WindowSpec.PartitionBy")
+		}
+		sep = " "
+	}
+	if n.OrderBy != nil {
+		ctx.WritePlain(sep)
+		if err := n.OrderBy.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore WindowSpec.OrderBy")
+		}
+		sep = " "
+	}
+	if n.Frame != nil {
+		ctx.WritePlain(sep)
+		if err := n.Frame.Restore(ctx); err != nil {
+			return errors.Annotate(err, "An error occurred while restore WindowSpec.Frame")
+		}
+	}
+	ctx.WritePlain(")")
+
+	return nil
 }
 
 // Accept implements Node Accept interface.
@@ -1442,7 +1474,16 @@ type PartitionByClause struct {
 
 // Restore implements Node interface.
 func (n *PartitionByClause) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord("PARTITION BY ")
+	for i, v := range n.Items {
+		if i != 0 {
+			ctx.WritePlain(", ")
+		}
+		if err := v.Restore(ctx); err != nil {
+			return errors.Annotatef(err, "An error occurred while restore PartitionByClause.Items[%d]", i)
+		}
+	}
+	return nil
 }
 
 // Accept implements Node Accept interface.
@@ -1483,7 +1524,24 @@ type FrameClause struct {
 
 // Restore implements Node interface.
 func (n *FrameClause) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	switch n.Type {
+	case Rows:
+		ctx.WriteKeyWord("ROWS")
+	case Ranges:
+		ctx.WriteKeyWord("RANGE")
+	default:
+		return errors.New("Unsupported window function frame type")
+	}
+	ctx.WriteKeyWord(" BETWEEN ")
+	if err := n.Extent.Start.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore FrameClause.Extent.Start")
+	}
+	ctx.WriteKeyWord(" AND ")
+	if err := n.Extent.End.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore FrameClause.Extent.End")
+	}
+
+	return nil
 }
 
 // Accept implements Node Accept interface.

diff --git a/ast/dml_test.go b/ast/dml_test.go
@@ -329,3 +329,53 @@ func (ts *testDMLSuite) TestFrameBoundRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "select avg(val) over (rows between %s and current row) from t", extractNodeFunc)
 }
+
+func (ts *testDMLSuite) TestFrameClauseRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"ROWS CURRENT ROW", "ROWS BETWEEN CURRENT ROW AND CURRENT ROW"},
+		{"ROWS UNBOUNDED PRECEDING", "ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW"},
+		{"ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING", "ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING"},
+		{"RANGE BETWEEN ? PRECEDING AND ? FOLLOWING", "RANGE BETWEEN ? PRECEDING AND ? FOLLOWING"},
+		{"RANGE BETWEEN INTERVAL 5 DAY PRECEDING AND INTERVAL '2:30' MINUTE_SECOND FOLLOWING", "RANGE BETWEEN INTERVAL 5 DAY PRECEDING AND INTERVAL '2:30' MINUTE_SECOND FOLLOWING"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr.(*WindowFuncExpr).Spec.Frame
+	}
+	RunNodeRestoreTest(c, testCases, "select avg(val) over (%s) from t", extractNodeFunc)
+}
+
+func (ts *testDMLSuite) TestPartitionByClauseRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"PARTITION BY a", "PARTITION BY `a`"},
+		{"PARTITION BY NULL", "PARTITION BY NULL"},
+		{"PARTITION BY a, b", "PARTITION BY `a`, `b`"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr.(*WindowFuncExpr).Spec.PartitionBy
+	}
+	RunNodeRestoreTest(c, testCases, "select avg(val) over (%s rows current row) from t", extractNodeFunc)
+}
+
+func (ts *testDMLSuite) TestWindowSpecRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"w as ()", "`w` AS ()"},
+		{"w as (w1)", "`w` AS (`w1`)"},
+		{"w as (w1 order by country)", "`w` AS (`w1` ORDER BY `country`)"},
+		{"w as (partition by a order by b rows current row)", "`w` AS (PARTITION BY `a` ORDER BY `b` ROWS BETWEEN CURRENT ROW AND CURRENT ROW)"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return &node.(*SelectStmt).WindowSpecs[0]
+	}
+	RunNodeRestoreTest(c, testCases, "select rank() over w from t window %s", extractNodeFunc)
+
+	testCases = []NodeRestoreTestCase{
+		{"w", "(`w`)"},
+		{"()", "()"},
+		{"(w PARTITION BY country)", "(`w` PARTITION BY `country`)"},
+		{"(PARTITION BY a ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING)", "(PARTITION BY `a` ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING)"},
+	}
+	extractNodeFunc = func(node Node) Node {
+		return &node.(*SelectStmt).Fields.Fields[0].Expr.(*WindowFuncExpr).Spec
+	}
+	RunNodeRestoreTest(c, testCases, "select rank() over %s from t window w as (order by a)", extractNodeFunc)
+}
diff --git a/ast/functions.go b/ast/functions.go
@@ -725,7 +725,31 @@ type WindowFuncExpr struct {
 
 // Restore implements Node interface.
 func (n *WindowFuncExpr) Restore(ctx *RestoreCtx) error {
-	return errors.New("Not implemented")
+	ctx.WriteKeyWord(n.F)
+	ctx.WritePlain("(")
+	for i, v := range n.Args {
+		if i != 0 {
+			ctx.WritePlain(", ")
+		} else if n.Distinct {
+			ctx.WriteKeyWord("DISTINCT ")
+		}
+		if err := v.Restore(ctx); err != nil {
+			return errors.Annotatef(err, "An error occurred while restore WindowFuncExpr.Args[%d]", i)
+		}
+	}
+	ctx.WritePlain(")")
+	if n.FromLast {
+		ctx.WriteKeyWord(" FROM LAST")
+	}
+	if n.IgnoreNull {
+		ctx.WriteKeyWord(" IGNORE NULLS")
+	}
+	ctx.WriteKeyWord(" OVER ")
+	if err := n.Spec.Restore(ctx); err != nil {
+		return errors.Annotate(err, "An error occurred while restore WindowFuncExpr.Spec")
+	}
+
+	return nil
 }
 
 // Format formats the window function expression into a Writer.

diff --git a/ast/functions_test.go b/ast/functions_test.go
@@ -120,3 +120,22 @@ func (ts *testFunctionsSuite) TestAggregateFuncExprRestore(c *C) {
 	}
 	RunNodeRestoreTest(c, testCases, "select %s", extractNodeFunc)
 }
+
+func (ts *testDMLSuite) TestWindowFuncExprRestore(c *C) {
+	testCases := []NodeRestoreTestCase{
+		{"RANK() OVER w", "RANK() OVER (`w`)"},
+		{"RANK() OVER (PARTITION BY a)", "RANK() OVER (PARTITION BY `a`)"},
+		{"MAX(DISTINCT a) OVER (PARTITION BY a)", "MAX(DISTINCT `a`) OVER (PARTITION BY `a`)"},
+		{"MAX(DISTINCTROW a) OVER (PARTITION BY a)", "MAX(DISTINCT `a`) OVER (PARTITION BY `a`)"},
+		{"MAX(DISTINCT ALL a) OVER (PARTITION BY a)", "MAX(DISTINCT `a`) OVER (PARTITION BY `a`)"},
+		{"MAX(ALL a) OVER (PARTITION BY a)", "MAX(`a`) OVER (PARTITION BY `a`)"},
+		{"FIRST_VALUE(val) IGNORE NULLS OVER w", "FIRST_VALUE(`val`) IGNORE NULLS OVER (`w`)"},
+		{"FIRST_VALUE(val) RESPECT NULLS OVER w", "FIRST_VALUE(`val`) OVER (`w`)"},
+		{"NTH_VALUE(val, 233) FROM LAST IGNORE NULLS OVER w", "NTH_VALUE(`val`, 233) FROM LAST IGNORE NULLS OVER (`w`)"},
+		{"NTH_VALUE(val, 233) FROM FIRST IGNORE NULLS OVER w", "NTH_VALUE(`val`, 233) IGNORE NULLS OVER (`w`)"},
+	}
+	extractNodeFunc := func(node Node) Node {
+		return node.(*SelectStmt).Fields.Fields[0].Expr
+	}
+	RunNodeRestoreTest(c, testCases, "select %s from t", extractNodeFunc)
+}