Skip to content

Commit 0ea55bd

Browse files
authored
implement ntile (#2935)
1 parent ecbe045 commit 0ea55bd

File tree

8 files changed

+445
-11
lines changed

8 files changed

+445
-11
lines changed

enginetest/queries/script_queries.go

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8107,6 +8107,208 @@ where
81078107
},
81088108
},
81098109
},
8110+
{
8111+
Name: "ntile tests",
8112+
Dialect: "mysql",
8113+
SetUpScript: []string{
8114+
"create table t (i int primary key, j int);",
8115+
"insert into t values (1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (6, 2), (7, 2), (8, 2), (9, 2), (10, 2);",
8116+
},
8117+
Assertions: []ScriptTestAssertion{
8118+
{
8119+
Query: "select i, ntile(null) over() from t;",
8120+
ExpectedErr: sql.ErrInvalidArgument,
8121+
},
8122+
{
8123+
Query: "select i, ntile(0) over() from t;",
8124+
ExpectedErr: sql.ErrInvalidArgument,
8125+
},
8126+
{
8127+
Query: "select i, ntile(-1) over() from t;",
8128+
ExpectedErr: sql.ErrInvalidArgument,
8129+
},
8130+
{
8131+
Query: "select i, ntile(100) over() from t;",
8132+
Expected: []sql.Row{
8133+
{1, uint64(1)},
8134+
{2, uint64(2)},
8135+
{3, uint64(3)},
8136+
{4, uint64(4)},
8137+
{5, uint64(5)},
8138+
{6, uint64(6)},
8139+
{7, uint64(7)},
8140+
{8, uint64(8)},
8141+
{9, uint64(9)},
8142+
{10, uint64(10)},
8143+
},
8144+
},
8145+
{
8146+
Query: "select i, ntile(10) over() from t;",
8147+
Expected: []sql.Row{
8148+
{1, uint64(1)},
8149+
{2, uint64(2)},
8150+
{3, uint64(3)},
8151+
{4, uint64(4)},
8152+
{5, uint64(5)},
8153+
{6, uint64(6)},
8154+
{7, uint64(7)},
8155+
{8, uint64(8)},
8156+
{9, uint64(9)},
8157+
{10, uint64(10)},
8158+
},
8159+
},
8160+
{
8161+
Query: "select i, ntile(9) over() from t;",
8162+
Expected: []sql.Row{
8163+
{1, uint64(1)},
8164+
{2, uint64(1)},
8165+
{3, uint64(2)},
8166+
{4, uint64(3)},
8167+
{5, uint64(4)},
8168+
{6, uint64(5)},
8169+
{7, uint64(6)},
8170+
{8, uint64(7)},
8171+
{9, uint64(8)},
8172+
{10, uint64(9)},
8173+
},
8174+
},
8175+
{
8176+
Query: "select i, ntile(8) over() from t;",
8177+
Expected: []sql.Row{
8178+
{1, uint64(1)},
8179+
{2, uint64(1)},
8180+
{3, uint64(2)},
8181+
{4, uint64(2)},
8182+
{5, uint64(3)},
8183+
{6, uint64(4)},
8184+
{7, uint64(5)},
8185+
{8, uint64(6)},
8186+
{9, uint64(7)},
8187+
{10, uint64(8)},
8188+
},
8189+
},
8190+
{
8191+
Query: "select i, ntile(7) over() from t;",
8192+
Expected: []sql.Row{
8193+
{1, uint64(1)},
8194+
{2, uint64(1)},
8195+
{3, uint64(2)},
8196+
{4, uint64(2)},
8197+
{5, uint64(3)},
8198+
{6, uint64(3)},
8199+
{7, uint64(4)},
8200+
{8, uint64(5)},
8201+
{9, uint64(6)},
8202+
{10, uint64(7)},
8203+
},
8204+
},
8205+
{
8206+
Query: "select i, ntile(6) over() from t;",
8207+
Expected: []sql.Row{
8208+
{1, uint64(1)},
8209+
{2, uint64(1)},
8210+
{3, uint64(2)},
8211+
{4, uint64(2)},
8212+
{5, uint64(3)},
8213+
{6, uint64(3)},
8214+
{7, uint64(4)},
8215+
{8, uint64(4)},
8216+
{9, uint64(5)},
8217+
{10, uint64(6)},
8218+
},
8219+
},
8220+
{
8221+
Query: "select i, ntile(5) over() from t;",
8222+
Expected: []sql.Row{
8223+
{1, uint64(1)},
8224+
{2, uint64(1)},
8225+
{3, uint64(2)},
8226+
{4, uint64(2)},
8227+
{5, uint64(3)},
8228+
{6, uint64(3)},
8229+
{7, uint64(4)},
8230+
{8, uint64(4)},
8231+
{9, uint64(5)},
8232+
{10, uint64(5)},
8233+
},
8234+
},
8235+
{
8236+
Query: "select i, ntile(4) over() from t;",
8237+
Expected: []sql.Row{
8238+
{1, uint64(1)},
8239+
{2, uint64(1)},
8240+
{3, uint64(1)},
8241+
{4, uint64(2)},
8242+
{5, uint64(2)},
8243+
{6, uint64(2)},
8244+
{7, uint64(3)},
8245+
{8, uint64(3)},
8246+
{9, uint64(4)},
8247+
{10, uint64(4)},
8248+
},
8249+
},
8250+
{
8251+
Query: "select i, ntile(3) over() from t;",
8252+
Expected: []sql.Row{
8253+
{1, uint64(1)},
8254+
{2, uint64(1)},
8255+
{3, uint64(1)},
8256+
{4, uint64(1)},
8257+
{5, uint64(2)},
8258+
{6, uint64(2)},
8259+
{7, uint64(2)},
8260+
{8, uint64(3)},
8261+
{9, uint64(3)},
8262+
{10, uint64(3)},
8263+
},
8264+
},
8265+
{
8266+
Query: "select i, ntile(2) over() from t;",
8267+
Expected: []sql.Row{
8268+
{1, uint64(1)},
8269+
{2, uint64(1)},
8270+
{3, uint64(1)},
8271+
{4, uint64(1)},
8272+
{5, uint64(1)},
8273+
{6, uint64(2)},
8274+
{7, uint64(2)},
8275+
{8, uint64(2)},
8276+
{9, uint64(2)},
8277+
{10, uint64(2)},
8278+
},
8279+
},
8280+
{
8281+
Query: "select i, ntile(1) over() from t;",
8282+
Expected: []sql.Row{
8283+
{1, uint64(1)},
8284+
{2, uint64(1)},
8285+
{3, uint64(1)},
8286+
{4, uint64(1)},
8287+
{5, uint64(1)},
8288+
{6, uint64(1)},
8289+
{7, uint64(1)},
8290+
{8, uint64(1)},
8291+
{9, uint64(1)},
8292+
{10, uint64(1)},
8293+
},
8294+
},
8295+
{
8296+
Query: "select i, j, ntile(2) over(partition by j) from t;",
8297+
Expected: []sql.Row{
8298+
{1, 1, uint64(1)},
8299+
{2, 1, uint64(1)},
8300+
{3, 1, uint64(1)},
8301+
{4, 1, uint64(2)},
8302+
{5, 1, uint64(2)},
8303+
{6, 2, uint64(1)},
8304+
{7, 2, uint64(1)},
8305+
{8, 2, uint64(1)},
8306+
{9, 2, uint64(2)},
8307+
{10, 2, uint64(2)},
8308+
},
8309+
},
8310+
},
8311+
},
81108312
}
81118313

81128314
var SpatialScriptTests = []ScriptTest{

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ require (
66
github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad
77
github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71
88
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81
9-
github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4
9+
github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b
1010
github.com/go-kit/kit v0.10.0
1111
github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d
1212
github.com/gocraft/dbr/v2 v2.7.2

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,8 @@ github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 h1:bMGS25NWAGTE
5858
github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71/go.mod h1:2/2zjLQ/JOOSbbSboojeg+cAwcRV0fDLzIiWch/lhqI=
5959
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 h1:7/v8q9XGFa6q5Ap4Z/OhNkAMBaK5YeuEzwJt+NZdhiE=
6060
github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY=
61-
github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4 h1:LGTt2LtYX8vaai32d+c9L0sMcP+Dg9w1kO6+lbsxxYg=
62-
github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
61+
github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b h1:2wE+qJwJ5SRIzz+dJQT8XbkpK+g8/pFt34AU/iJ5K+Y=
62+
github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70=
6363
github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
6464
github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs=
6565
github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU=

optgen/cmd/source/unary_aggs.yaml

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,11 @@ unaryAggs:
3232
- name: "Sum"
3333
desc: "returns the sum of expr in all rows"
3434
nullable: false
35-
- name: "StdDevPop"
36-
desc: "returns the population standard deviation of expr"
37-
- name: "StdDevSamp"
38-
desc: "returns the sample standard deviation of expr"
39-
- name: "VarPop"
40-
desc: "returns the population variance of expr"
41-
- name: "VarSamp"
42-
desc: "returns the sample variance of expr"
35+
- name: "StdDevPop"
36+
desc: "returns the population standard deviation of expr"
37+
- name: "StdDevSamp"
38+
desc: "returns the sample standard deviation of expr"
39+
- name: "VarPop"
40+
desc: "returns the population variance of expr"
41+
- name: "VarSamp"
42+
desc: "returns the sample variance of expr"
Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
// Copyright 2025 Dolthub, Inc.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package window
16+
17+
import (
18+
"strings"
19+
20+
"github.com/dolthub/go-mysql-server/sql"
21+
"github.com/dolthub/go-mysql-server/sql/expression/function/aggregation"
22+
"github.com/dolthub/go-mysql-server/sql/types"
23+
)
24+
25+
type NTile struct {
26+
pos uint64
27+
count uint64
28+
29+
bucketExpr sql.Expression
30+
bucketSize uint64
31+
32+
id sql.ColumnId
33+
window *sql.WindowDefinition
34+
}
35+
36+
var _ sql.FunctionExpression = (*NTile)(nil)
37+
var _ sql.WindowAggregation = (*NTile)(nil)
38+
var _ sql.WindowAdaptableExpression = (*NTile)(nil)
39+
var _ sql.CollationCoercible = (*NTile)(nil)
40+
41+
func NewNTile(expr sql.Expression) sql.Expression {
42+
return &NTile{
43+
bucketExpr: expr,
44+
}
45+
}
46+
47+
// Id implements sql.IdExpression
48+
func (n *NTile) Id() sql.ColumnId {
49+
return n.id
50+
}
51+
52+
// WithId implements sql.IdExpression
53+
func (n *NTile) WithId(id sql.ColumnId) sql.IdExpression {
54+
ret := *n
55+
ret.id = id
56+
return &ret
57+
}
58+
59+
// Description implements sql.FunctionExpression
60+
func (n *NTile) Description() string {
61+
return "returns percentage rank value."
62+
}
63+
64+
// Window implements sql.WindowExpression
65+
func (n *NTile) Window() *sql.WindowDefinition {
66+
return n.window
67+
}
68+
69+
func (n *NTile) Resolved() bool {
70+
return windowResolved(n.window)
71+
}
72+
73+
func (n *NTile) String() string {
74+
sb := strings.Builder{}
75+
sb.WriteString("ntile()")
76+
if n.window != nil {
77+
sb.WriteString(" ")
78+
sb.WriteString(n.window.String())
79+
}
80+
return sb.String()
81+
}
82+
83+
func (n *NTile) DebugString() string {
84+
sb := strings.Builder{}
85+
sb.WriteString("ntile()")
86+
if n.window != nil {
87+
sb.WriteString(" ")
88+
sb.WriteString(sql.DebugString(n.window))
89+
}
90+
return sb.String()
91+
}
92+
93+
// FunctionName implements sql.FunctionExpression
94+
func (n *NTile) FunctionName() string {
95+
return "NTILE"
96+
}
97+
98+
// Type implements sql.Expression
99+
func (n *NTile) Type() sql.Type {
100+
return types.Float64
101+
}
102+
103+
// CollationCoercibility implements the interface sql.CollationCoercible.
104+
func (*NTile) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) {
105+
return sql.Collation_binary, 5
106+
}
107+
108+
// IsNullable implements sql.Expression
109+
func (n *NTile) IsNullable() bool {
110+
return false
111+
}
112+
113+
// Eval implements sql.Expression
114+
func (n *NTile) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) {
115+
return nil, sql.ErrWindowUnsupported.New(n.FunctionName())
116+
}
117+
118+
// Children implements sql.Expression
119+
func (n *NTile) Children() []sql.Expression {
120+
return n.window.ToExpressions()
121+
}
122+
123+
// WithChildren implements sql.Expression
124+
func (n *NTile) WithChildren(children ...sql.Expression) (sql.Expression, error) {
125+
window, err := n.window.FromExpressions(children)
126+
if err != nil {
127+
return nil, err
128+
}
129+
130+
return n.WithWindow(window), nil
131+
}
132+
133+
// WithWindow implements sql.WindowAggregation
134+
func (n *NTile) WithWindow(window *sql.WindowDefinition) sql.WindowAdaptableExpression {
135+
nr := *n
136+
nr.window = window
137+
return &nr
138+
}
139+
140+
func (n *NTile) NewWindowFunction() (sql.WindowFunction, error) {
141+
return aggregation.NewNTile(n.bucketExpr), nil
142+
}

0 commit comments

Comments
 (0)