diff --git a/parser/ast/dml.go b/parser/ast/dml.go index 18d274dc1fbcc..05dc733a92d7b 100644 --- a/parser/ast/dml.go +++ b/parser/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/parser/ast/dml_test.go b/parser/ast/dml_test.go index 9dfb802559aa2..1bb2864480e78 100644 --- a/parser/ast/dml_test.go +++ b/parser/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/parser/ast/functions.go b/parser/ast/functions.go index cb04061b4d378..d1008eb854240 100644 --- a/parser/ast/functions.go +++ b/parser/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/parser/ast/functions_test.go b/parser/ast/functions_test.go index bca4d830e950a..15b5a8134e9c1 100644 --- a/parser/ast/functions_test.go +++ b/parser/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) +}