From b4ebcd41fc2aef570c0f0ceaed8e4d0bff4ae4fa Mon Sep 17 00:00:00 2001 From: Yiding Cui Date: Tue, 12 Mar 2019 20:17:31 +0800 Subject: [PATCH 1/5] executor: add window function NTILE --- executor/aggfuncs/builder.go | 11 ++++ executor/aggfuncs/func_ntile.go | 81 +++++++++++++++++++++++++++ executor/window_test.go | 7 +++ expression/aggregation/base_func.go | 9 +++ expression/aggregation/window_func.go | 9 ++- planner/core/logical_plan_test.go | 8 +++ planner/core/rule_fix_empty_schema.go | 20 +++++++ 7 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 executor/aggfuncs/func_ntile.go create mode 100644 planner/core/rule_fix_empty_schema.go diff --git a/executor/aggfuncs/builder.go b/executor/aggfuncs/builder.go index 4ef95f7f4a701..3d811edba423a 100644 --- a/executor/aggfuncs/builder.go +++ b/executor/aggfuncs/builder.go @@ -73,6 +73,8 @@ func BuildWindowFunctions(ctx sessionctx.Context, windowFuncDesc *aggregation.Ag return buildCumeDist(ordinal, orderByCols) case ast.WindowFuncNthValue: return buildNthValue(windowFuncDesc, ordinal) + case ast.WindowFuncNtile: + return buildNtile(windowFuncDesc, ordinal) default: return Build(ctx, windowFuncDesc, ordinal) } @@ -386,3 +388,12 @@ func buildNthValue(aggFuncDesc *aggregation.AggFuncDesc, ordinal int) AggFunc { nth, _, _ := expression.GetUint64FromConstant(aggFuncDesc.Args[1]) return &nthValue{baseAggFunc: base, tp: aggFuncDesc.RetTp, nth: nth} } + +func buildNtile(aggFuncDes *aggregation.AggFuncDesc, ordinal int) AggFunc { + base := baseAggFunc{ + args: aggFuncDes.Args, + ordinal: ordinal, + } + n, _, _ := expression.GetUint64FromConstant(aggFuncDes.Args[0]) + return &ntile{baseAggFunc: base, n: n} +} diff --git a/executor/aggfuncs/func_ntile.go b/executor/aggfuncs/func_ntile.go new file mode 100644 index 0000000000000..8863d92404f96 --- /dev/null +++ b/executor/aggfuncs/func_ntile.go @@ -0,0 +1,81 @@ +// Copyright 2019 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// See the License for the specific language governing permissions and +// limitations under the License. + +package aggfuncs + +import ( + "github.com/pingcap/tidb/sessionctx" + "github.com/pingcap/tidb/util/chunk" +) + +// ntile divides the partition into n ranked groups and returns the group number a row belongs to. +// e.g. We have 11 rows and n = 3. They will be divided into 3 groups.\ +// First 4 rows belongs to group 1. Following 4 rows belongs to group 2. The last 3 rows belongs to group 3. +type ntile struct { + n uint64 + baseAggFunc +} + +type partialResult4Ntile struct { + curIdx uint64 + divisor uint64 + curGroupIdx uint64 + remainder uint64 + quotient uint64 + rows []chunk.Row +} + +func (n *ntile) AllocPartialResult() PartialResult { + return PartialResult(&partialResult4Ntile{divisor: n.n, curGroupIdx: 1}) +} + +func (n *ntile) ResetPartialResult(pr PartialResult) { + p := (*partialResult4Ntile)(pr) + p.curIdx = 0 + p.curGroupIdx = 1 + p.rows = p.rows[:0] +} + +func (n *ntile) UpdatePartialResult(_ sessionctx.Context, rowsInGroup []chunk.Row, pr PartialResult) error { + p := (*partialResult4Ntile)(pr) + p.rows = append(p.rows, rowsInGroup...) + // Update the quotient and remainder. + if p.divisor != 0 { + p.quotient = uint64(len(p.rows)) / p.divisor + p.remainder = uint64(len(p.rows)) % p.divisor + } + return nil +} + +func (n *ntile) AppendFinalResult2Chunk(_ sessionctx.Context, pr PartialResult, chk *chunk.Chunk) error { + p := (*partialResult4Ntile)(pr) + + // If the divisor is 0, the arg of NTILE would be NULL. So we just return NULL. + if p.divisor == 0 { + chk.AppendNull(n.ordinal) + return nil + } + + chk.AppendUint64(n.ordinal, p.curGroupIdx) + + p.curIdx++ + curMaxIdx := p.quotient + if p.curGroupIdx <= p.remainder { + curMaxIdx++ + } + if p.curIdx == curMaxIdx { + p.curIdx = 0 + p.curGroupIdx++ + } + return nil +} diff --git a/executor/window_test.go b/executor/window_test.go index 76cc408a654d9..18a41e8db05bd 100644 --- a/executor/window_test.go +++ b/executor/window_test.go @@ -119,4 +119,11 @@ func (s *testSuite2) TestWindowFunctions(c *C) { result.Check(testkit.Rows("1 2", "1 2", "2 2", "2 2")) result = tk.MustQuery("select a, nth_value(a, 5) over() from t") result.Check(testkit.Rows("1 ", "1 ", "2 ", "2 ")) + + result = tk.MustQuery("select ntile(3) over() from t") + result.Check(testkit.Rows("1", "1", "2", "3")) + result = tk.MustQuery("select ntile(2) over() from t") + result.Check(testkit.Rows("1", "1", "2", "2")) + result = tk.MustQuery("select ntile(null) over() from t") + result.Check(testkit.Rows("", "", "", "")) } diff --git a/expression/aggregation/base_func.go b/expression/aggregation/base_func.go index 9e8db54b5d560..796a6706d9f86 100644 --- a/expression/aggregation/base_func.go +++ b/expression/aggregation/base_func.go @@ -100,6 +100,8 @@ func (a *baseFuncDesc) typeInfer(ctx sessionctx.Context) { a.typeInfer4NumberFuncs() case ast.WindowFuncCumeDist: a.typeInfer4CumeDist() + case ast.WindowFuncNtile: + a.typeInfer4Ntile() default: panic("unsupported agg function: " + a.Name) } @@ -200,6 +202,13 @@ func (a *baseFuncDesc) typeInfer4CumeDist() { a.RetTp.Flen, a.RetTp.Decimal = mysql.MaxRealWidth, mysql.NotFixedDec } +func (a *baseFuncDesc) typeInfer4Ntile() { + a.RetTp = types.NewFieldType(mysql.TypeLonglong) + a.RetTp.Flen = 21 + types.SetBinChsClnFlag(a.RetTp) + a.RetTp.Flag |= mysql.UnsignedFlag +} + // GetDefaultValue gets the default value when the function's input is null. // According to MySQL, default values of the function are listed as follows: // e.g. diff --git a/expression/aggregation/window_func.go b/expression/aggregation/window_func.go index aa436b957aab7..f252fc1cb8b92 100644 --- a/expression/aggregation/window_func.go +++ b/expression/aggregation/window_func.go @@ -28,12 +28,19 @@ type WindowFuncDesc struct { // NewWindowFuncDesc creates a window function signature descriptor. func NewWindowFuncDesc(ctx sessionctx.Context, name string, args []expression.Expression) *WindowFuncDesc { - if strings.ToLower(name) == ast.WindowFuncNthValue { + switch strings.ToLower(name) { + case ast.WindowFuncNthValue: val, isNull, ok := expression.GetUint64FromConstant(args[1]) // nth_value does not allow `0`, but allows `null`. if !ok || (val == 0 && !isNull) { return nil } + case ast.WindowFuncNtile: + val, isNull, ok := expression.GetUint64FromConstant(args[0]) + // ntile does not allow `0`, but allows `null`. + if !ok || (val == 0 && !isNull) { + return nil + } } return &WindowFuncDesc{newBaseFuncDesc(ctx, name, args)} } diff --git a/planner/core/logical_plan_test.go b/planner/core/logical_plan_test.go index 889d4a325bb79..cf475d4f1c901 100644 --- a/planner/core/logical_plan_test.go +++ b/planner/core/logical_plan_test.go @@ -2207,6 +2207,14 @@ func (s *testPlanSuite) TestWindowFunction(c *C) { sql: "select nth_value(a, 0) over() from t", result: "[planner:1210]Incorrect arguments to nth_value", }, + { + sql: "select ntile(0) over() from t", + result: "[planner:1210]Incorrect arguments to ntile", + }, + { + sql: "select ntile(null) over() from t", + result: "TableReader(Table(t))->Window(ntile(null) over())->Projection", + }, } s.Parser.EnableWindowFunc(true) diff --git a/planner/core/rule_fix_empty_schema.go b/planner/core/rule_fix_empty_schema.go new file mode 100644 index 0000000000000..e49803170f499 --- /dev/null +++ b/planner/core/rule_fix_empty_schema.go @@ -0,0 +1,20 @@ +// Copyright 2016 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// See the License for the specific language governing permissions and +// limitations under the License. + +package core + +// A plan's schema may become empty after column pruning. Such as `select count(*) from t`. The schema of the `DataSource` +// is empty. But due to TiDB's execution logic. We need one row to ensure that +func fixEmptySchema(p PhysicalPlan) { + +} From 42f0522e4b73ed65c6cd80d6b4e16139e111949d Mon Sep 17 00:00:00 2001 From: Yiding Cui Date: Wed, 13 Mar 2019 11:15:26 +0800 Subject: [PATCH 2/5] address comments --- planner/core/rule_fix_empty_schema.go | 20 -------------------- 1 file changed, 20 deletions(-) delete mode 100644 planner/core/rule_fix_empty_schema.go diff --git a/planner/core/rule_fix_empty_schema.go b/planner/core/rule_fix_empty_schema.go deleted file mode 100644 index e49803170f499..0000000000000 --- a/planner/core/rule_fix_empty_schema.go +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright 2016 PingCAP, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// See the License for the specific language governing permissions and -// limitations under the License. - -package core - -// A plan's schema may become empty after column pruning. Such as `select count(*) from t`. The schema of the `DataSource` -// is empty. But due to TiDB's execution logic. We need one row to ensure that -func fixEmptySchema(p PhysicalPlan) { - -} From 8e67747b55800109af09dd7478e4c55398c0446e Mon Sep 17 00:00:00 2001 From: Yiding Cui Date: Wed, 13 Mar 2019 11:37:30 +0800 Subject: [PATCH 3/5] fix unit-test --- planner/core/logical_plan_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/planner/core/logical_plan_test.go b/planner/core/logical_plan_test.go index cf475d4f1c901..3c87649e27547 100644 --- a/planner/core/logical_plan_test.go +++ b/planner/core/logical_plan_test.go @@ -2213,7 +2213,7 @@ func (s *testPlanSuite) TestWindowFunction(c *C) { }, { sql: "select ntile(null) over() from t", - result: "TableReader(Table(t))->Window(ntile(null) over())->Projection", + result: "TableReader(Table(t))->Window(ntile() over())->Projection", }, } From 8d262cb4176c81648666a7eb97f7dad6d849d991 Mon Sep 17 00:00:00 2001 From: Yiding Cui Date: Wed, 13 Mar 2019 12:49:53 +0800 Subject: [PATCH 4/5] remove unused char --- executor/aggfuncs/func_ntile.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/executor/aggfuncs/func_ntile.go b/executor/aggfuncs/func_ntile.go index 8863d92404f96..698d1c1a2b694 100644 --- a/executor/aggfuncs/func_ntile.go +++ b/executor/aggfuncs/func_ntile.go @@ -19,7 +19,7 @@ import ( ) // ntile divides the partition into n ranked groups and returns the group number a row belongs to. -// e.g. We have 11 rows and n = 3. They will be divided into 3 groups.\ +// e.g. We have 11 rows and n = 3. They will be divided into 3 groups. // First 4 rows belongs to group 1. Following 4 rows belongs to group 2. The last 3 rows belongs to group 3. type ntile struct { n uint64 From 6cb888c8cae1b9bfc7005d3aaddf5a93962a3971 Mon Sep 17 00:00:00 2001 From: Yiding Cui Date: Fri, 15 Mar 2019 15:00:26 +0800 Subject: [PATCH 5/5] address comments --- executor/aggfuncs/func_ntile.go | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/executor/aggfuncs/func_ntile.go b/executor/aggfuncs/func_ntile.go index 698d1c1a2b694..1adbb326d7609 100644 --- a/executor/aggfuncs/func_ntile.go +++ b/executor/aggfuncs/func_ntile.go @@ -28,7 +28,6 @@ type ntile struct { type partialResult4Ntile struct { curIdx uint64 - divisor uint64 curGroupIdx uint64 remainder uint64 quotient uint64 @@ -36,7 +35,7 @@ type partialResult4Ntile struct { } func (n *ntile) AllocPartialResult() PartialResult { - return PartialResult(&partialResult4Ntile{divisor: n.n, curGroupIdx: 1}) + return PartialResult(&partialResult4Ntile{curGroupIdx: 1}) } func (n *ntile) ResetPartialResult(pr PartialResult) { @@ -50,9 +49,9 @@ func (n *ntile) UpdatePartialResult(_ sessionctx.Context, rowsInGroup []chunk.Ro p := (*partialResult4Ntile)(pr) p.rows = append(p.rows, rowsInGroup...) // Update the quotient and remainder. - if p.divisor != 0 { - p.quotient = uint64(len(p.rows)) / p.divisor - p.remainder = uint64(len(p.rows)) % p.divisor + if n.n != 0 { + p.quotient = uint64(len(p.rows)) / n.n + p.remainder = uint64(len(p.rows)) % n.n } return nil } @@ -61,7 +60,7 @@ func (n *ntile) AppendFinalResult2Chunk(_ sessionctx.Context, pr PartialResult, p := (*partialResult4Ntile)(pr) // If the divisor is 0, the arg of NTILE would be NULL. So we just return NULL. - if p.divisor == 0 { + if n.n == 0 { chk.AppendNull(n.ordinal) return nil }