From 59b2983a86b6b89c793e52a1f6c76b6e713c66ca Mon Sep 17 00:00:00 2001 From: James Cor Date: Fri, 11 Apr 2025 00:54:09 -0700 Subject: [PATCH 1/2] implement ntile --- enginetest/queries/script_queries.go | 202 ++++++++++++++++++ go.mod | 2 +- go.sum | 4 +- optgen/cmd/source/unary_aggs.yaml | 16 +- .../function/aggregation/window/ntile.go | 142 ++++++++++++ .../function/aggregation/window_functions.go | 89 ++++++++ sql/expression/function/registry.go | 1 + sql/planbuilder/aggregates.go | 1 + 8 files changed, 446 insertions(+), 11 deletions(-) create mode 100644 sql/expression/function/aggregation/window/ntile.go diff --git a/enginetest/queries/script_queries.go b/enginetest/queries/script_queries.go index e169d0901c..c998a7c0cb 100644 --- a/enginetest/queries/script_queries.go +++ b/enginetest/queries/script_queries.go @@ -8107,6 +8107,208 @@ where }, }, }, + { + Name: "ntile tests", + Dialect: "mysql", + SetUpScript: []string{ + "create table t (i int primary key, j int);", + "insert into t values (1, 1), (2, 1), (3, 1), (4, 1), (5, 1), (6, 2), (7, 2), (8, 2), (9, 2), (10, 2);", + }, + Assertions: []ScriptTestAssertion{ + { + Query: "select i, ntile(null) over() from t;", + ExpectedErr: sql.ErrInvalidArgument, + }, + { + Query: "select i, ntile(0) over() from t;", + ExpectedErr: sql.ErrInvalidArgument, + }, + { + Query: "select i, ntile(-1) over() from t;", + ExpectedErr: sql.ErrInvalidArgument, + }, + { + Query: "select i, ntile(100) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(2)}, + {3, uint64(3)}, + {4, uint64(4)}, + {5, uint64(5)}, + {6, uint64(6)}, + {7, uint64(7)}, + {8, uint64(8)}, + {9, uint64(9)}, + {10, uint64(10)}, + }, + }, + { + Query: "select i, ntile(10) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(2)}, + {3, uint64(3)}, + {4, uint64(4)}, + {5, uint64(5)}, + {6, uint64(6)}, + {7, uint64(7)}, + {8, uint64(8)}, + {9, uint64(9)}, + {10, uint64(10)}, + }, + }, + { + Query: "select i, ntile(9) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(3)}, + {5, uint64(4)}, + {6, uint64(5)}, + {7, uint64(6)}, + {8, uint64(7)}, + {9, uint64(8)}, + {10, uint64(9)}, + }, + }, + { + Query: "select i, ntile(8) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(4)}, + {7, uint64(5)}, + {8, uint64(6)}, + {9, uint64(7)}, + {10, uint64(8)}, + }, + }, + { + Query: "select i, ntile(7) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(5)}, + {9, uint64(6)}, + {10, uint64(7)}, + }, + }, + { + Query: "select i, ntile(6) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(4)}, + {9, uint64(5)}, + {10, uint64(6)}, + }, + }, + { + Query: "select i, ntile(5) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(4)}, + {9, uint64(5)}, + {10, uint64(5)}, + }, + }, + { + Query: "select i, ntile(4) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(2)}, + {5, uint64(2)}, + {6, uint64(2)}, + {7, uint64(3)}, + {8, uint64(3)}, + {9, uint64(4)}, + {10, uint64(4)}, + }, + }, + { + Query: "select i, ntile(3) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(2)}, + {6, uint64(2)}, + {7, uint64(2)}, + {8, uint64(3)}, + {9, uint64(3)}, + {10, uint64(3)}, + }, + }, + { + Query: "select i, ntile(2) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(1)}, + {6, uint64(2)}, + {7, uint64(2)}, + {8, uint64(2)}, + {9, uint64(2)}, + {10, uint64(2)}, + }, + }, + { + Query: "select i, ntile(1) over() from t;", + Expected: []sql.Row{ + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(1)}, + {6, uint64(1)}, + {7, uint64(1)}, + {8, uint64(1)}, + {9, uint64(1)}, + {10, uint64(1)}, + }, + }, + { + Query: "select i, j, ntile(2) over(partition by j) from t;", + Expected: []sql.Row{ + {1, 1, uint64(1)}, + {2, 1, uint64(1)}, + {3, 1, uint64(1)}, + {4, 1, uint64(2)}, + {5, 1, uint64(2)}, + {6, 2, uint64(1)}, + {7, 2, uint64(1)}, + {8, 2, uint64(1)}, + {9, 2, uint64(2)}, + {10, 2, uint64(2)}, + }, + }, + }, + }, } var SpatialScriptTests = []ScriptTest{ diff --git a/go.mod b/go.mod index aed4846e2d..d1ad5f897a 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( github.com/dolthub/go-icu-regex v0.0.0-20250327004329-6799764f2dad github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 - github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4 + github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b github.com/go-kit/kit v0.10.0 github.com/go-sql-driver/mysql v1.7.2-0.20231213112541-0004702b931d github.com/gocraft/dbr/v2 v2.7.2 diff --git a/go.sum b/go.sum index 391e1a6454..b525e231fd 100644 --- a/go.sum +++ b/go.sum @@ -58,8 +58,8 @@ github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71 h1:bMGS25NWAGTE github.com/dolthub/jsonpath v0.0.2-0.20240227200619-19675ab05c71/go.mod h1:2/2zjLQ/JOOSbbSboojeg+cAwcRV0fDLzIiWch/lhqI= github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81 h1:7/v8q9XGFa6q5Ap4Z/OhNkAMBaK5YeuEzwJt+NZdhiE= github.com/dolthub/sqllogictest/go v0.0.0-20201107003712-816f3ae12d81/go.mod h1:siLfyv2c92W1eN/R4QqG/+RjjX5W2+gCTRjZxBjI3TY= -github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4 h1:LGTt2LtYX8vaai32d+c9L0sMcP+Dg9w1kO6+lbsxxYg= -github.com/dolthub/vitess v0.0.0-20250410090211-143e6b272ad4/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70= +github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b h1:2wE+qJwJ5SRIzz+dJQT8XbkpK+g8/pFt34AU/iJ5K+Y= +github.com/dolthub/vitess v0.0.0-20250410233614-8d8c7a5b3d6b/go.mod h1:1gQZs/byeHLMSul3Lvl3MzioMtOW1je79QYGyi2fd70= github.com/dustin/go-humanize v0.0.0-20171111073723-bb3d318650d4/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk= github.com/eapache/go-resiliency v1.1.0/go.mod h1:kFI+JgMyC7bLPUVY133qvEBtVayf5mFgVsvEsIPBvNs= github.com/eapache/go-xerial-snappy v0.0.0-20180814174437-776d5712da21/go.mod h1:+020luEh2TKB4/GOp8oxxtq0Daoen/Cii55CzbTV6DU= diff --git a/optgen/cmd/source/unary_aggs.yaml b/optgen/cmd/source/unary_aggs.yaml index 7e40fe762f..996d416197 100644 --- a/optgen/cmd/source/unary_aggs.yaml +++ b/optgen/cmd/source/unary_aggs.yaml @@ -32,11 +32,11 @@ unaryAggs: - name: "Sum" desc: "returns the sum of expr in all rows" nullable: false -- name: "StdDevPop" - desc: "returns the population standard deviation of expr" -- name: "StdDevSamp" - desc: "returns the sample standard deviation of expr" -- name: "VarPop" - desc: "returns the population variance of expr" -- name: "VarSamp" - desc: "returns the sample variance of expr" \ No newline at end of file +- name: "StdDevPop" + desc: "returns the population standard deviation of expr" +- name: "StdDevSamp" + desc: "returns the sample standard deviation of expr" +- name: "VarPop" + desc: "returns the population variance of expr" +- name: "VarSamp" + desc: "returns the sample variance of expr" \ No newline at end of file diff --git a/sql/expression/function/aggregation/window/ntile.go b/sql/expression/function/aggregation/window/ntile.go new file mode 100644 index 0000000000..c49583ce30 --- /dev/null +++ b/sql/expression/function/aggregation/window/ntile.go @@ -0,0 +1,142 @@ +// Copyright 2025 Dolthub, 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, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package window + +import ( + "strings" + + "github.com/dolthub/go-mysql-server/sql" + "github.com/dolthub/go-mysql-server/sql/expression/function/aggregation" + "github.com/dolthub/go-mysql-server/sql/types" +) + +type NTile struct { + pos uint64 + count uint64 + + bucketExpr sql.Expression + bucketSize uint64 + + id sql.ColumnId + window *sql.WindowDefinition +} + +var _ sql.FunctionExpression = (*NTile)(nil) +var _ sql.WindowAggregation = (*NTile)(nil) +var _ sql.WindowAdaptableExpression = (*NTile)(nil) +var _ sql.CollationCoercible = (*NTile)(nil) + +func NewNTile(expr sql.Expression) sql.Expression { + return &NTile{ + bucketExpr: expr, + } +} + +// Id implements sql.IdExpression +func (n *NTile) Id() sql.ColumnId { + return n.id +} + +// WithId implements sql.IdExpression +func (n *NTile) WithId(id sql.ColumnId) sql.IdExpression { + ret := *n + ret.id = id + return &ret +} + +// Description implements sql.FunctionExpression +func (n *NTile) Description() string { + return "returns percentage rank value." +} + +// Window implements sql.WindowExpression +func (n *NTile) Window() *sql.WindowDefinition { + return n.window +} + +func (n *NTile) Resolved() bool { + return windowResolved(n.window) +} + +func (n *NTile) String() string { + sb := strings.Builder{} + sb.WriteString("ntile()") + if n.window != nil { + sb.WriteString(" ") + sb.WriteString(n.window.String()) + } + return sb.String() +} + +func (n *NTile) DebugString() string { + sb := strings.Builder{} + sb.WriteString("ntile()") + if n.window != nil { + sb.WriteString(" ") + sb.WriteString(sql.DebugString(n.window)) + } + return sb.String() +} + +// FunctionName implements sql.FunctionExpression +func (n *NTile) FunctionName() string { + return "NTILE" +} + +// Type implements sql.Expression +func (n *NTile) Type() sql.Type { + return types.Float64 +} + +// CollationCoercibility implements the interface sql.CollationCoercible. +func (*NTile) CollationCoercibility(ctx *sql.Context) (collation sql.CollationID, coercibility byte) { + return sql.Collation_binary, 5 +} + +// IsNullable implements sql.Expression +func (n *NTile) IsNullable() bool { + return false +} + +// Eval implements sql.Expression +func (n *NTile) Eval(ctx *sql.Context, row sql.Row) (interface{}, error) { + return nil, sql.ErrWindowUnsupported.New(n.FunctionName()) +} + +// Children implements sql.Expression +func (n *NTile) Children() []sql.Expression { + return n.window.ToExpressions() +} + +// WithChildren implements sql.Expression +func (n *NTile) WithChildren(children ...sql.Expression) (sql.Expression, error) { + window, err := n.window.FromExpressions(children) + if err != nil { + return nil, err + } + + return n.WithWindow(window), nil +} + +// WithWindow implements sql.WindowAggregation +func (n *NTile) WithWindow(window *sql.WindowDefinition) sql.WindowAdaptableExpression { + nr := *n + nr.window = window + return &nr +} + +func (n *NTile) NewWindowFunction() (sql.WindowFunction, error) { + return aggregation.NewNTile(n.bucketExpr), nil +} diff --git a/sql/expression/function/aggregation/window_functions.go b/sql/expression/function/aggregation/window_functions.go index 3265781d75..8d7cb0f578 100644 --- a/sql/expression/function/aggregation/window_functions.go +++ b/sql/expression/function/aggregation/window_functions.go @@ -1335,6 +1335,95 @@ func (a *DenseRank) Compute(ctx *sql.Context, interval sql.WindowInterval, buf s return a.denseRank } +type NTile struct { + numBucketsExpr sql.Expression + + pos uint64 + bucketSize uint64 + bigBuckets uint64 + bucket uint64 + + // orderBy tracks peer group increments + orderBy []sql.Expression +} + +func NewNTile(expr sql.Expression) *NTile { + return &NTile{ + numBucketsExpr: expr, + } +} + +func (n *NTile) WithWindow(w *sql.WindowDefinition) (sql.WindowFunction, error) { + na := *n + na.orderBy = w.OrderBy.ToExpressions() + return &na, nil +} + +func (n *NTile) Dispose() { + return +} + +func (n *NTile) DefaultFramer() sql.WindowFramer { + return NewPeerGroupFramer(n.orderBy) +} + +func (n *NTile) StartPartition(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) error { + n.Dispose() + if interval.End < interval.Start { + return nil + } + + numBucketsVal, err := n.numBucketsExpr.Eval(ctx, nil) + if err != nil { + return err + } + numBucketsVal, _, err = types.Int64.Convert(ctx, numBucketsVal) + if err != nil { + numBucketsVal = uint64(0) + } + if numBucketsVal == nil { + return sql.ErrInvalidArgument.New("NTILE") + } + if numBucketsVal.(int64) <= 0 { + return sql.ErrInvalidArgument.New("NTILE") + } + + count := uint64(interval.End - interval.Start) + numBuckets := uint64(numBucketsVal.(int64)) + if numBuckets > count { + n.bucketSize = 1 + } else { + n.bucketSize = count / numBuckets + n.bigBuckets = count % numBuckets + } + n.pos = 0 + n.bucket = 1 + return nil +} + +// Compute returns the appropriate bucket for the current row. +func (n *NTile) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.WindowBuffer) interface{} { + defer func() { n.pos++ }() + if n.pos == 0 { + return n.bucket + } + + // the first n.bigBuckets buckets are of size n.bucketSize + 1 + // the remaining buckets are of size n.bucketSize + if n.bigBuckets > 0 && n.pos % (n.bucketSize + 1) == 0 { + n.bucket++ + n.bigBuckets-- + if n.bigBuckets == 0 { + n.pos = 0 + } + } else if n.bigBuckets == 0 && n.pos % n.bucketSize == 0 { + n.bucket++ + } + + return n.bucket +} + + type Lag struct { leadLagBase } diff --git a/sql/expression/function/registry.go b/sql/expression/function/registry.go index 478d47e5ff..90188102b2 100644 --- a/sql/expression/function/registry.go +++ b/sql/expression/function/registry.go @@ -181,6 +181,7 @@ var BuiltIns = []sql.Function{ sql.FunctionN{Name: "mod", Fn: NewMod}, sql.Function1{Name: "month", Fn: NewMonth}, sql.Function1{Name: "monthname", Fn: NewMonthName}, + sql.Function1{Name: "ntile", Fn: window.NewNTile}, sql.FunctionN{Name: "now", Fn: NewNow}, sql.Function2{Name: "nullif", Fn: NewNullIf}, sql.Function1{Name: "octet_length", Fn: NewLength}, diff --git a/sql/planbuilder/aggregates.go b/sql/planbuilder/aggregates.go index c5f0b13d9f..662b1f6027 100644 --- a/sql/planbuilder/aggregates.go +++ b/sql/planbuilder/aggregates.go @@ -505,6 +505,7 @@ func isWindowFunc(name string) bool { "row_number", "percent_rank", "lead", "lag", "first_value", "last_value", "rank", "dense_rank", + "ntile", "std", "stddev", "stddev_pop", "stddev_samp", "variance", "var_pop", "var_samp": return true From be970f316030b24434f9f54a4e1544799271936e Mon Sep 17 00:00:00 2001 From: jycor Date: Fri, 11 Apr 2025 07:56:43 +0000 Subject: [PATCH 2/2] [ga-format-pr] Run ./format_repo.sh to fix formatting --- enginetest/queries/script_queries.go | 246 +++++++++--------- .../function/aggregation/window_functions.go | 5 +- 2 files changed, 125 insertions(+), 126 deletions(-) diff --git a/enginetest/queries/script_queries.go b/enginetest/queries/script_queries.go index c998a7c0cb..3f874eb3ba 100644 --- a/enginetest/queries/script_queries.go +++ b/enginetest/queries/script_queries.go @@ -8108,7 +8108,7 @@ where }, }, { - Name: "ntile tests", + Name: "ntile tests", Dialect: "mysql", SetUpScript: []string{ "create table t (i int primary key, j int);", @@ -8120,190 +8120,190 @@ where ExpectedErr: sql.ErrInvalidArgument, }, { - Query: "select i, ntile(0) over() from t;", + Query: "select i, ntile(0) over() from t;", ExpectedErr: sql.ErrInvalidArgument, }, { - Query: "select i, ntile(-1) over() from t;", + Query: "select i, ntile(-1) over() from t;", ExpectedErr: sql.ErrInvalidArgument, }, { - Query: "select i, ntile(100) over() from t;", + Query: "select i, ntile(100) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(2)}, - {3, uint64(3)}, - {4, uint64(4)}, - {5, uint64(5)}, - {6, uint64(6)}, - {7, uint64(7)}, - {8, uint64(8)}, - {9, uint64(9)}, + {1, uint64(1)}, + {2, uint64(2)}, + {3, uint64(3)}, + {4, uint64(4)}, + {5, uint64(5)}, + {6, uint64(6)}, + {7, uint64(7)}, + {8, uint64(8)}, + {9, uint64(9)}, {10, uint64(10)}, }, }, { - Query: "select i, ntile(10) over() from t;", + Query: "select i, ntile(10) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(2)}, - {3, uint64(3)}, - {4, uint64(4)}, - {5, uint64(5)}, - {6, uint64(6)}, - {7, uint64(7)}, - {8, uint64(8)}, - {9, uint64(9)}, + {1, uint64(1)}, + {2, uint64(2)}, + {3, uint64(3)}, + {4, uint64(4)}, + {5, uint64(5)}, + {6, uint64(6)}, + {7, uint64(7)}, + {8, uint64(8)}, + {9, uint64(9)}, {10, uint64(10)}, }, }, { - Query: "select i, ntile(9) over() from t;", + Query: "select i, ntile(9) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(2)}, - {4, uint64(3)}, - {5, uint64(4)}, - {6, uint64(5)}, - {7, uint64(6)}, - {8, uint64(7)}, - {9, uint64(8)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(3)}, + {5, uint64(4)}, + {6, uint64(5)}, + {7, uint64(6)}, + {8, uint64(7)}, + {9, uint64(8)}, {10, uint64(9)}, }, }, { - Query: "select i, ntile(8) over() from t;", + Query: "select i, ntile(8) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(2)}, - {4, uint64(2)}, - {5, uint64(3)}, - {6, uint64(4)}, - {7, uint64(5)}, - {8, uint64(6)}, - {9, uint64(7)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(4)}, + {7, uint64(5)}, + {8, uint64(6)}, + {9, uint64(7)}, {10, uint64(8)}, }, }, { - Query: "select i, ntile(7) over() from t;", + Query: "select i, ntile(7) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(2)}, - {4, uint64(2)}, - {5, uint64(3)}, - {6, uint64(3)}, - {7, uint64(4)}, - {8, uint64(5)}, - {9, uint64(6)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(5)}, + {9, uint64(6)}, {10, uint64(7)}, }, }, { - Query: "select i, ntile(6) over() from t;", + Query: "select i, ntile(6) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(2)}, - {4, uint64(2)}, - {5, uint64(3)}, - {6, uint64(3)}, - {7, uint64(4)}, - {8, uint64(4)}, - {9, uint64(5)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(4)}, + {9, uint64(5)}, {10, uint64(6)}, }, }, { - Query: "select i, ntile(5) over() from t;", + Query: "select i, ntile(5) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(2)}, - {4, uint64(2)}, - {5, uint64(3)}, - {6, uint64(3)}, - {7, uint64(4)}, - {8, uint64(4)}, - {9, uint64(5)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(2)}, + {4, uint64(2)}, + {5, uint64(3)}, + {6, uint64(3)}, + {7, uint64(4)}, + {8, uint64(4)}, + {9, uint64(5)}, {10, uint64(5)}, }, }, { - Query: "select i, ntile(4) over() from t;", + Query: "select i, ntile(4) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(1)}, - {4, uint64(2)}, - {5, uint64(2)}, - {6, uint64(2)}, - {7, uint64(3)}, - {8, uint64(3)}, - {9, uint64(4)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(2)}, + {5, uint64(2)}, + {6, uint64(2)}, + {7, uint64(3)}, + {8, uint64(3)}, + {9, uint64(4)}, {10, uint64(4)}, }, }, { - Query: "select i, ntile(3) over() from t;", + Query: "select i, ntile(3) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(1)}, - {4, uint64(1)}, - {5, uint64(2)}, - {6, uint64(2)}, - {7, uint64(2)}, - {8, uint64(3)}, - {9, uint64(3)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(2)}, + {6, uint64(2)}, + {7, uint64(2)}, + {8, uint64(3)}, + {9, uint64(3)}, {10, uint64(3)}, }, }, { - Query: "select i, ntile(2) over() from t;", + Query: "select i, ntile(2) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(1)}, - {4, uint64(1)}, - {5, uint64(1)}, - {6, uint64(2)}, - {7, uint64(2)}, - {8, uint64(2)}, - {9, uint64(2)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(1)}, + {6, uint64(2)}, + {7, uint64(2)}, + {8, uint64(2)}, + {9, uint64(2)}, {10, uint64(2)}, }, }, { - Query: "select i, ntile(1) over() from t;", + Query: "select i, ntile(1) over() from t;", Expected: []sql.Row{ - {1, uint64(1)}, - {2, uint64(1)}, - {3, uint64(1)}, - {4, uint64(1)}, - {5, uint64(1)}, - {6, uint64(1)}, - {7, uint64(1)}, - {8, uint64(1)}, - {9, uint64(1)}, + {1, uint64(1)}, + {2, uint64(1)}, + {3, uint64(1)}, + {4, uint64(1)}, + {5, uint64(1)}, + {6, uint64(1)}, + {7, uint64(1)}, + {8, uint64(1)}, + {9, uint64(1)}, {10, uint64(1)}, }, }, { - Query: "select i, j, ntile(2) over(partition by j) from t;", + Query: "select i, j, ntile(2) over(partition by j) from t;", Expected: []sql.Row{ - {1, 1, uint64(1)}, - {2, 1, uint64(1)}, - {3, 1, uint64(1)}, - {4, 1, uint64(2)}, - {5, 1, uint64(2)}, - {6, 2, uint64(1)}, - {7, 2, uint64(1)}, - {8, 2, uint64(1)}, - {9, 2, uint64(2)}, + {1, 1, uint64(1)}, + {2, 1, uint64(1)}, + {3, 1, uint64(1)}, + {4, 1, uint64(2)}, + {5, 1, uint64(2)}, + {6, 2, uint64(1)}, + {7, 2, uint64(1)}, + {8, 2, uint64(1)}, + {9, 2, uint64(2)}, {10, 2, uint64(2)}, }, }, diff --git a/sql/expression/function/aggregation/window_functions.go b/sql/expression/function/aggregation/window_functions.go index 8d7cb0f578..4316d9b899 100644 --- a/sql/expression/function/aggregation/window_functions.go +++ b/sql/expression/function/aggregation/window_functions.go @@ -1410,20 +1410,19 @@ func (n *NTile) Compute(ctx *sql.Context, interval sql.WindowInterval, buf sql.W // the first n.bigBuckets buckets are of size n.bucketSize + 1 // the remaining buckets are of size n.bucketSize - if n.bigBuckets > 0 && n.pos % (n.bucketSize + 1) == 0 { + if n.bigBuckets > 0 && n.pos%(n.bucketSize+1) == 0 { n.bucket++ n.bigBuckets-- if n.bigBuckets == 0 { n.pos = 0 } - } else if n.bigBuckets == 0 && n.pos % n.bucketSize == 0 { + } else if n.bigBuckets == 0 && n.pos%n.bucketSize == 0 { n.bucket++ } return n.bucket } - type Lag struct { leadLagBase }