From 48d7011bb14e19bba1dc6ecb698dfd22f966804d Mon Sep 17 00:00:00 2001 From: Yahor Yuzefovich Date: Thu, 16 May 2019 17:27:28 -0700 Subject: [PATCH] exec: template out RANK and ROW_NUMBER Templates out rankOp into two operators (for dense and non-dense case) and templates out support for PARTITION BY clause for both rankOp and rowNumberOp. The code is undertested and is lacking benchmarks, but that will be addressed soon. Release note: None --- Makefile | 6 +- pkg/sql/exec/execgen/cmd/execgen/rank_gen.go | 91 +++++ .../execgen/cmd/execgen/row_number_gen.go | 45 +++ pkg/sql/exec/vecbuiltins/.gitignore | 2 + pkg/sql/exec/vecbuiltins/rank.eg.go | 313 ++++++++++++++++++ pkg/sql/exec/vecbuiltins/rank.go | 168 +--------- pkg/sql/exec/vecbuiltins/rank_tmpl.go | 147 ++++++++ pkg/sql/exec/vecbuiltins/row_number.eg.go | 75 +++++ pkg/sql/exec/vecbuiltins/row_number.go | 61 +--- pkg/sql/exec/vecbuiltins/row_number_tmpl.go | 84 +++++ .../logictest/testdata/logic_test/exec_window | 1 - pkg/testutils/lint/lint_test.go | 1 + 12 files changed, 783 insertions(+), 211 deletions(-) create mode 100644 pkg/sql/exec/execgen/cmd/execgen/rank_gen.go create mode 100644 pkg/sql/exec/execgen/cmd/execgen/row_number_gen.go create mode 100644 pkg/sql/exec/vecbuiltins/.gitignore create mode 100644 pkg/sql/exec/vecbuiltins/rank.eg.go create mode 100644 pkg/sql/exec/vecbuiltins/rank_tmpl.go create mode 100644 pkg/sql/exec/vecbuiltins/row_number.eg.go create mode 100644 pkg/sql/exec/vecbuiltins/row_number_tmpl.go diff --git a/Makefile b/Makefile index 8913c43b4672..9385f1a2b590 100644 --- a/Makefile +++ b/Makefile @@ -755,7 +755,9 @@ EXECGEN_TARGETS = \ pkg/sql/exec/selection_ops.eg.go \ pkg/sql/exec/sort.eg.go \ pkg/sql/exec/sum_agg.eg.go \ - pkg/sql/exec/tuples_differ.eg.go + pkg/sql/exec/tuples_differ.eg.go \ + pkg/sql/exec/vecbuiltins/rank.eg.go \ + pkg/sql/exec/vecbuiltins/row_number.eg.go OPTGEN_TARGETS = \ pkg/sql/opt/memo/expr.og.go \ @@ -1411,6 +1413,8 @@ pkg/sql/exec/rowstovec.eg.go: pkg/sql/exec/rowstovec_tmpl.go pkg/sql/exec/sort.eg.go: pkg/sql/exec/sort_tmpl.go pkg/sql/exec/sum_agg.eg.go: pkg/sql/exec/sum_agg_tmpl.go pkg/sql/exec/tuples_differ.eg.go: pkg/sql/exec/tuples_differ_tmpl.go +pkg/sql/exec/vecbuiltins/rank.eg.go: pkg/sql/exec/vecbuiltins/rank_tmpl.go +pkg/sql/exec/vecbuiltins/row_number.eg.go: pkg/sql/exec/vecbuiltins/row_number_tmpl.go $(EXECGEN_TARGETS): bin/execgen @# Remove generated files with the old suffix to avoid conflicts. diff --git a/pkg/sql/exec/execgen/cmd/execgen/rank_gen.go b/pkg/sql/exec/execgen/cmd/execgen/rank_gen.go new file mode 100644 index 000000000000..c9d9440e28a3 --- /dev/null +++ b/pkg/sql/exec/execgen/cmd/execgen/rank_gen.go @@ -0,0 +1,91 @@ +// Copyright 2019 The Cockroach Authors. +// +// 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 main + +import ( + "fmt" + "io" + "io/ioutil" + "regexp" + "strings" + "text/template" +) + +type rankTmplInfo struct { + Dense bool + HasPartition bool +} + +func (r rankTmplInfo) UpdateRank() string { + if r.Dense { + return fmt.Sprintf( + `r.rank++`, + ) + } + return fmt.Sprintf( + `r.rank += r.rankIncrement +r.rankIncrement = 1`, + ) +} + +func (r rankTmplInfo) UpdateRankIncrement() string { + if r.Dense { + return `` + } + return fmt.Sprintf( + `r.rankIncrement++`, + ) +} + +// Avoid unused warnings. These methods are used in the template. +var ( + _ = rankTmplInfo{}.UpdateRank() + _ = rankTmplInfo{}.UpdateRankIncrement() +) + +func genRankOps(wr io.Writer) error { + d, err := ioutil.ReadFile("pkg/sql/exec/vecbuiltins/rank_tmpl.go") + if err != nil { + return err + } + + s := string(d) + + s = strings.Replace(s, "_DENSE", "{{.Dense}}", -1) + s = strings.Replace(s, "_PARTITION", "{{.HasPartition}}", -1) + + updateRankRe := regexp.MustCompile(`_UPDATE_RANK\(\)`) + s = updateRankRe.ReplaceAllString(s, "{{.UpdateRank}}") + updateRankIncrementRe := regexp.MustCompile(`_UPDATE_RANK_INCREMENT\(\)`) + s = updateRankIncrementRe.ReplaceAllString(s, "{{.UpdateRankIncrement}}") + + // Now, generate the op, from the template. + tmpl, err := template.New("rank_op").Funcs(template.FuncMap{"buildDict": buildDict}).Parse(s) + if err != nil { + return err + } + + rankTmplInfos := []rankTmplInfo{ + {Dense: false, HasPartition: false}, + {Dense: false, HasPartition: true}, + {Dense: true, HasPartition: false}, + {Dense: true, HasPartition: true}, + } + return tmpl.Execute(wr, rankTmplInfos) +} + +func init() { + registerGenerator(genRankOps, "rank.eg.go") +} diff --git a/pkg/sql/exec/execgen/cmd/execgen/row_number_gen.go b/pkg/sql/exec/execgen/cmd/execgen/row_number_gen.go new file mode 100644 index 000000000000..c7d7642f3ab8 --- /dev/null +++ b/pkg/sql/exec/execgen/cmd/execgen/row_number_gen.go @@ -0,0 +1,45 @@ +// Copyright 2019 The Cockroach Authors. +// +// 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 main + +import ( + "io" + "io/ioutil" + "text/template" +) + +func genRowNumberOp(wr io.Writer) error { + d, err := ioutil.ReadFile("pkg/sql/exec/vecbuiltins/row_number_tmpl.go") + if err != nil { + return err + } + + s := string(d) + + nextRowNumber := makeFunctionRegex("_NEXT_ROW_NUMBER", 1) + s = nextRowNumber.ReplaceAllString(s, `{{template "nextRowNumber" buildDict "Global" $ "HasPartition" $1 }}`) + + // Now, generate the op, from the template. + tmpl, err := template.New("row_number_op").Funcs(template.FuncMap{"buildDict": buildDict}).Parse(s) + if err != nil { + return err + } + + return tmpl.Execute(wr, struct{}{}) +} + +func init() { + registerGenerator(genRowNumberOp, "row_number.eg.go") +} diff --git a/pkg/sql/exec/vecbuiltins/.gitignore b/pkg/sql/exec/vecbuiltins/.gitignore new file mode 100644 index 000000000000..c249aee1c589 --- /dev/null +++ b/pkg/sql/exec/vecbuiltins/.gitignore @@ -0,0 +1,2 @@ +rank.eg.go +row_number.eg.go diff --git a/pkg/sql/exec/vecbuiltins/rank.eg.go b/pkg/sql/exec/vecbuiltins/rank.eg.go new file mode 100644 index 000000000000..9b94c248787a --- /dev/null +++ b/pkg/sql/exec/vecbuiltins/rank.eg.go @@ -0,0 +1,313 @@ +// Code generated by execgen; DO NOT EDIT. +// Copyright 2019 The Cockroach Authors. +// 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 vecbuiltins + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/sql/exec" + "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" + "github.com/cockroachdb/cockroach/pkg/sql/exec/types" +) + +type rankInitFields struct { + input exec.Operator + // distinctCol is the output column of the chain of ordered distinct + // operators in which true will indicate that a new rank needs to be assigned + // to the corresponding tuple. + distinctCol []bool + outputColIdx int + partitionColIdx int +} + +type rankDense_false_HasPartition_false_Op struct { + rankInitFields + + // rank indicates which rank should be assigned to the next tuple. + rank int64 + // rankIncrement indicates by how much rank should be incremented when a + // tuple distinct from the previous one on the ordering columns is seen. It + // is used only in case of a regular rank function (i.e. not dense). + rankIncrement int64 +} + +var _ exec.Operator = &rankDense_false_HasPartition_false_Op{} + +func (r *rankDense_false_HasPartition_false_Op) Init() { + r.input.Init() + // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a + // tuple in the batch, we first increment r.rank, so setting this + // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as + // desired). + r.rankIncrement = 1 +} + +func (r *rankDense_false_HasPartition_false_Op) Next(ctx context.Context) coldata.Batch { + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch + } + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rankCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + // TODO(yuzefovich): template out sel vs non-sel cases. + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + if r.distinctCol[sel[i]] { + r.rank += r.rankIncrement + r.rankIncrement = 1 + rankCol[sel[i]] = r.rank + } else { + rankCol[sel[i]] = r.rank + r.rankIncrement++ + } + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + if r.distinctCol[i] { + r.rank += r.rankIncrement + r.rankIncrement = 1 + rankCol[i] = r.rank + } else { + rankCol[i] = r.rank + r.rankIncrement++ + } + } + } + return batch +} + +type rankDense_false_HasPartition_true_Op struct { + rankInitFields + + // rank indicates which rank should be assigned to the next tuple. + rank int64 + // rankIncrement indicates by how much rank should be incremented when a + // tuple distinct from the previous one on the ordering columns is seen. It + // is used only in case of a regular rank function (i.e. not dense). + rankIncrement int64 +} + +var _ exec.Operator = &rankDense_false_HasPartition_true_Op{} + +func (r *rankDense_false_HasPartition_true_Op) Init() { + r.input.Init() + // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a + // tuple in the batch, we first increment r.rank, so setting this + // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as + // desired). + r.rankIncrement = 1 +} + +func (r *rankDense_false_HasPartition_true_Op) Next(ctx context.Context) coldata.Batch { + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch + } + + if r.partitionColIdx == batch.Width() { + batch.AppendCol(types.Bool) + } else if r.partitionColIdx > batch.Width() { + panic("unexpected: column partitionColIdx is neither present nor the next to be appended") + } + partitionCol := batch.ColVec(r.partitionColIdx).Bool() + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rankCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + // TODO(yuzefovich): template out sel vs non-sel cases. + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[sel[i]] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[sel[i]] = 1 + continue + } + if r.distinctCol[sel[i]] { + r.rank += r.rankIncrement + r.rankIncrement = 1 + rankCol[sel[i]] = r.rank + } else { + rankCol[sel[i]] = r.rank + r.rankIncrement++ + } + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[i] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[i] = 1 + continue + } + if r.distinctCol[i] { + r.rank += r.rankIncrement + r.rankIncrement = 1 + rankCol[i] = r.rank + } else { + rankCol[i] = r.rank + r.rankIncrement++ + } + } + } + return batch +} + +type rankDense_true_HasPartition_false_Op struct { + rankInitFields + + // rank indicates which rank should be assigned to the next tuple. + rank int64 + // rankIncrement indicates by how much rank should be incremented when a + // tuple distinct from the previous one on the ordering columns is seen. It + // is used only in case of a regular rank function (i.e. not dense). + rankIncrement int64 +} + +var _ exec.Operator = &rankDense_true_HasPartition_false_Op{} + +func (r *rankDense_true_HasPartition_false_Op) Init() { + r.input.Init() + // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a + // tuple in the batch, we first increment r.rank, so setting this + // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as + // desired). + r.rankIncrement = 1 +} + +func (r *rankDense_true_HasPartition_false_Op) Next(ctx context.Context) coldata.Batch { + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch + } + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rankCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + // TODO(yuzefovich): template out sel vs non-sel cases. + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + if r.distinctCol[sel[i]] { + r.rank++ + rankCol[sel[i]] = r.rank + } else { + rankCol[sel[i]] = r.rank + + } + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + if r.distinctCol[i] { + r.rank++ + rankCol[i] = r.rank + } else { + rankCol[i] = r.rank + + } + } + } + return batch +} + +type rankDense_true_HasPartition_true_Op struct { + rankInitFields + + // rank indicates which rank should be assigned to the next tuple. + rank int64 + // rankIncrement indicates by how much rank should be incremented when a + // tuple distinct from the previous one on the ordering columns is seen. It + // is used only in case of a regular rank function (i.e. not dense). + rankIncrement int64 +} + +var _ exec.Operator = &rankDense_true_HasPartition_true_Op{} + +func (r *rankDense_true_HasPartition_true_Op) Init() { + r.input.Init() + // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a + // tuple in the batch, we first increment r.rank, so setting this + // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as + // desired). + r.rankIncrement = 1 +} + +func (r *rankDense_true_HasPartition_true_Op) Next(ctx context.Context) coldata.Batch { + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch + } + + if r.partitionColIdx == batch.Width() { + batch.AppendCol(types.Bool) + } else if r.partitionColIdx > batch.Width() { + panic("unexpected: column partitionColIdx is neither present nor the next to be appended") + } + partitionCol := batch.ColVec(r.partitionColIdx).Bool() + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rankCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + // TODO(yuzefovich): template out sel vs non-sel cases. + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[sel[i]] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[sel[i]] = 1 + continue + } + if r.distinctCol[sel[i]] { + r.rank++ + rankCol[sel[i]] = r.rank + } else { + rankCol[sel[i]] = r.rank + + } + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[i] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[i] = 1 + continue + } + if r.distinctCol[i] { + r.rank++ + rankCol[i] = r.rank + } else { + rankCol[i] = r.rank + + } + } + } + return batch +} diff --git a/pkg/sql/exec/vecbuiltins/rank.go b/pkg/sql/exec/vecbuiltins/rank.go index 3c67f73a5ff5..893731237a7b 100644 --- a/pkg/sql/exec/vecbuiltins/rank.go +++ b/pkg/sql/exec/vecbuiltins/rank.go @@ -15,34 +15,12 @@ package vecbuiltins import ( - "context" - "github.com/cockroachdb/cockroach/pkg/sql/exec" - "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" "github.com/cockroachdb/cockroach/pkg/sql/exec/types" ) // TODO(yuzefovich): add randomized tests. -type rankOp struct { - input exec.Operator - dense bool - // distinctCol is the output column of the chain of ordered distinct - // operators in which true will indicate that a new rank needs to be assigned - // to the corresponding tuple. - distinctCol []bool - outputColIdx int - partitionColIdx int - - // rank indicates which rank should be assigned to the next tuple. - rank int64 - // rankIncrement indicates by how much rank should be incremented when a - // tuple distinct from the previous one on the ordering columns is seen. It - // is used only in case of a regular rank function (i.e. not dense). - rankIncrement int64 -} - -var _ exec.Operator = &rankOp{} - +// TODO(yuzefovich): add benchmarks. // NewRankOperator creates a new exec.Operator that computes window function // RANK or DENSE_RANK. dense distinguishes between the two functions. input // *must* already be ordered on orderingCols (which should not be empty). @@ -63,141 +41,21 @@ func NewRankOperator( if err != nil { return nil, err } - return &rankOp{input: op, dense: dense, distinctCol: outputCol, outputColIdx: outputColIdx, partitionColIdx: partitionColIdx}, nil -} - -func (r *rankOp) Init() { - r.input.Init() - // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a - // tuple in the batch, we first increment r.rank, so setting this - // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as - // desired). - r.rankIncrement = 1 -} - -func (r *rankOp) Next(ctx context.Context) coldata.Batch { - b := r.input.Next(ctx) - if b.Length() == 0 { - return b + initFields := rankInitFields{ + input: op, + distinctCol: outputCol, + outputColIdx: outputColIdx, + partitionColIdx: partitionColIdx, } - if r.partitionColIdx != -1 { - if r.partitionColIdx == b.Width() { - b.AppendCol(types.Bool) - } else if r.partitionColIdx > b.Width() { - panic("unexpected: column partitionColIdx is neither present nor the next to be appended") - } - if r.outputColIdx == b.Width() { - b.AppendCol(types.Int64) - } else if r.outputColIdx > b.Width() { - panic("unexpected: column outputColIdx is neither present nor the next to be appended") - } - partitionCol := b.ColVec(r.partitionColIdx).Bool() - rankCol := b.ColVec(r.outputColIdx).Int64() - if r.distinctCol == nil { - panic("unexpected: distinctCol is nil in rankOp") - } - sel := b.Selection() - if sel != nil { - for i := uint16(0); i < b.Length(); i++ { - if partitionCol[sel[i]] { - r.rank = 1 - r.rankIncrement = 1 - rankCol[sel[i]] = 1 - } else { - if r.distinctCol[sel[i]] { - // TODO(yuzefovich): template this part out to generate two different - // rank operators. - if r.dense { - r.rank++ - } else { - r.rank += r.rankIncrement - r.rankIncrement = 1 - } - rankCol[sel[i]] = r.rank - } else { - rankCol[sel[i]] = r.rank - if !r.dense { - r.rankIncrement++ - } - } - } - } - } else { - for i := uint16(0); i < b.Length(); i++ { - if partitionCol[i] { - r.rank = 1 - r.rankIncrement = 1 - rankCol[i] = 1 - } else { - if r.distinctCol[i] { - // TODO(yuzefovich): template this part out to generate two different - // rank operators. - if r.dense { - r.rank++ - } else { - r.rank += r.rankIncrement - r.rankIncrement = 1 - } - rankCol[i] = r.rank - } else { - rankCol[i] = r.rank - if !r.dense { - r.rankIncrement++ - } - } - } - } + if dense { + if partitionColIdx != -1 { + return &rankDense_true_HasPartition_true_Op{rankInitFields: initFields}, nil } + return &rankDense_true_HasPartition_false_Op{rankInitFields: initFields}, nil } else { - if r.outputColIdx == b.Width() { - b.AppendCol(types.Int64) - } else if r.outputColIdx > b.Width() { - panic("unexpected: column outputColIdx is neither present nor the next to be appended") - } - rankCol := b.ColVec(r.outputColIdx).Int64() - if r.distinctCol == nil { - panic("unexpected: distinctCol is nil in rankOp") - } - sel := b.Selection() - if sel != nil { - for i := uint16(0); i < b.Length(); i++ { - if r.distinctCol[sel[i]] { - // TODO(yuzefovich): template this part out to generate two different - // rank operators. - if r.dense { - r.rank++ - } else { - r.rank += r.rankIncrement - r.rankIncrement = 1 - } - rankCol[sel[i]] = r.rank - } else { - rankCol[sel[i]] = r.rank - if !r.dense { - r.rankIncrement++ - } - } - } - } else { - for i := uint16(0); i < b.Length(); i++ { - if r.distinctCol[i] { - // TODO(yuzefovich): template this part out to generate two different - // rank operators. - if r.dense { - r.rank++ - } else { - r.rank += r.rankIncrement - r.rankIncrement = 1 - } - rankCol[i] = r.rank - } else { - rankCol[i] = r.rank - if !r.dense { - r.rankIncrement++ - } - } - } + if partitionColIdx != -1 { + return &rankDense_false_HasPartition_true_Op{rankInitFields: initFields}, nil } + return &rankDense_false_HasPartition_false_Op{rankInitFields: initFields}, nil } - return b } diff --git a/pkg/sql/exec/vecbuiltins/rank_tmpl.go b/pkg/sql/exec/vecbuiltins/rank_tmpl.go new file mode 100644 index 000000000000..09e4fb5c148e --- /dev/null +++ b/pkg/sql/exec/vecbuiltins/rank_tmpl.go @@ -0,0 +1,147 @@ +// Copyright 2019 The Cockroach Authors. +// +// 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. + +// {{/* +// +build execgen_template +// +// This file is the execgen template for rank.eg.go. It's formatted in a +// special way, so it's both valid Go and a valid text/template input. This +// permits editing this file with editor support. +// +// */}} + +package vecbuiltins + +import ( + "context" + + "github.com/cockroachdb/cockroach/pkg/sql/exec" + "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" + "github.com/cockroachdb/cockroach/pkg/sql/exec/types" +) + +// {{/* + +// _UPDATE_RANK is the template function for updating the state of rank +// operators. +func _UPDATE_RANK() { + panic("") +} + +// _UPDATE_RANK_INCREMENT is the template function for updating the state of +// rank operators. +func _UPDATE_RANK_INCREMENT() { + panic("") +} + +// */}} + +type rankInitFields struct { + input exec.Operator + // distinctCol is the output column of the chain of ordered distinct + // operators in which true will indicate that a new rank needs to be assigned + // to the corresponding tuple. + distinctCol []bool + outputColIdx int + partitionColIdx int +} + +// {{range .}} + +type rankDense__DENSE_HasPartition__PARTITION_Op struct { + rankInitFields + + // rank indicates which rank should be assigned to the next tuple. + rank int64 + // rankIncrement indicates by how much rank should be incremented when a + // tuple distinct from the previous one on the ordering columns is seen. It + // is used only in case of a regular rank function (i.e. not dense). + rankIncrement int64 +} + +var _ exec.Operator = &rankDense__DENSE_HasPartition__PARTITION_Op{} + +func (r *rankDense__DENSE_HasPartition__PARTITION_Op) Init() { + r.input.Init() + // RANK and DENSE_RANK start counting from 1. Before we assign the rank to a + // tuple in the batch, we first increment r.rank, so setting this + // rankIncrement to 1 will update r.rank to 1 on the very first tuple (as + // desired). + r.rankIncrement = 1 +} + +func (r *rankDense__DENSE_HasPartition__PARTITION_Op) Next(ctx context.Context) coldata.Batch { + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch + } + + // {{ if .HasPartition }} + if r.partitionColIdx == batch.Width() { + batch.AppendCol(types.Bool) + } else if r.partitionColIdx > batch.Width() { + panic("unexpected: column partitionColIdx is neither present nor the next to be appended") + } + partitionCol := batch.ColVec(r.partitionColIdx).Bool() + // {{ end }} + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rankCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + // TODO(yuzefovich): template out sel vs non-sel cases. + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + // {{ if .HasPartition }} + if partitionCol[sel[i]] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[sel[i]] = 1 + continue + } + // {{end}} + if r.distinctCol[sel[i]] { + _UPDATE_RANK() + rankCol[sel[i]] = r.rank + } else { + rankCol[sel[i]] = r.rank + _UPDATE_RANK_INCREMENT() + } + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + // {{ if .HasPartition }} + if partitionCol[i] { + r.rank = 1 + r.rankIncrement = 1 + rankCol[i] = 1 + continue + } + // {{end}} + if r.distinctCol[i] { + _UPDATE_RANK() + rankCol[i] = r.rank + } else { + rankCol[i] = r.rank + _UPDATE_RANK_INCREMENT() + } + } + } + return batch +} + +// {{end}} diff --git a/pkg/sql/exec/vecbuiltins/row_number.eg.go b/pkg/sql/exec/vecbuiltins/row_number.eg.go new file mode 100644 index 000000000000..bd826f448790 --- /dev/null +++ b/pkg/sql/exec/vecbuiltins/row_number.eg.go @@ -0,0 +1,75 @@ +// Code generated by execgen; DO NOT EDIT. +// Copyright 2019 The Cockroach Authors. +// 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 vecbuiltins + +import ( + "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" + "github.com/cockroachdb/cockroach/pkg/sql/exec/types" +) + +func (r *rowNumberOp) nextBodyWithPartition(batch coldata.Batch) { + + if r.partitionColIdx == batch.Width() { + batch.AppendCol(types.Bool) + } else if r.partitionColIdx > batch.Width() { + panic("unexpected: column partitionColIdx is neither present nor the next to be appended") + } + partitionCol := batch.ColVec(r.partitionColIdx).Bool() + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rowNumberCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[sel[i]] { + r.rowNumber = 1 + } + rowNumberCol[sel[i]] = r.rowNumber + r.rowNumber++ + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + if partitionCol[i] { + r.rowNumber = 1 + } + rowNumberCol[i] = r.rowNumber + r.rowNumber++ + } + } +} + +func (r *rowNumberOp) nextBodyNoPartition(batch coldata.Batch) { + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rowNumberCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + rowNumberCol[sel[i]] = r.rowNumber + r.rowNumber++ + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + rowNumberCol[i] = r.rowNumber + r.rowNumber++ + } + } +} diff --git a/pkg/sql/exec/vecbuiltins/row_number.go b/pkg/sql/exec/vecbuiltins/row_number.go index 4c67b5b0c5eb..e2f3b218be8e 100644 --- a/pkg/sql/exec/vecbuiltins/row_number.go +++ b/pkg/sql/exec/vecbuiltins/row_number.go @@ -19,10 +19,10 @@ import ( "github.com/cockroachdb/cockroach/pkg/sql/exec" "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" - "github.com/cockroachdb/cockroach/pkg/sql/exec/types" ) // TODO(yuzefovich): add randomized tests. +// TODO(yuzefovich): add benchmarks. type rowNumberOp struct { input exec.Operator outputColIdx int @@ -50,61 +50,14 @@ func (r *rowNumberOp) Init() { } func (r *rowNumberOp) Next(ctx context.Context) coldata.Batch { - b := r.input.Next(ctx) - if b.Length() == 0 { - return b + batch := r.input.Next(ctx) + if batch.Length() == 0 { + return batch } - // TODO(yuzefovich): template partition out. if r.partitionColIdx != -1 { - if r.partitionColIdx == b.Width() { - b.AppendCol(types.Bool) - } else if r.partitionColIdx > b.Width() { - panic("unexpected: column partitionColIdx is neither present nor the next to be appended") - } - if r.outputColIdx == b.Width() { - b.AppendCol(types.Int64) - } else if r.outputColIdx > b.Width() { - panic("unexpected: column outputColIdx is neither present nor the next to be appended") - } - partitionCol := b.ColVec(r.partitionColIdx).Bool() - rowNumberCol := b.ColVec(r.outputColIdx).Int64() - sel := b.Selection() - if sel != nil { - for i := uint16(0); i < b.Length(); i++ { - if partitionCol[sel[i]] { - r.rowNumber = 1 - } - rowNumberCol[sel[i]] = r.rowNumber - r.rowNumber++ - } - } else { - for i := uint16(0); i < b.Length(); i++ { - if partitionCol[i] { - r.rowNumber = 1 - } - rowNumberCol[i] = r.rowNumber - r.rowNumber++ - } - } + r.nextBodyWithPartition(batch) } else { - if r.outputColIdx == b.Width() { - b.AppendCol(types.Int64) - } else if r.outputColIdx > b.Width() { - panic("unexpected: column outputColIdx is neither present nor the next to be appended") - } - rowNumberCol := b.ColVec(r.outputColIdx).Int64() - sel := b.Selection() - if sel != nil { - for i := uint16(0); i < b.Length(); i++ { - rowNumberCol[sel[i]] = r.rowNumber - r.rowNumber++ - } - } else { - for i := uint16(0); i < b.Length(); i++ { - rowNumberCol[i] = r.rowNumber - r.rowNumber++ - } - } + r.nextBodyNoPartition(batch) } - return b + return batch } diff --git a/pkg/sql/exec/vecbuiltins/row_number_tmpl.go b/pkg/sql/exec/vecbuiltins/row_number_tmpl.go new file mode 100644 index 000000000000..60c123565912 --- /dev/null +++ b/pkg/sql/exec/vecbuiltins/row_number_tmpl.go @@ -0,0 +1,84 @@ +// Copyright 2019 The Cockroach Authors. +// +// 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. + +// {{/* +// +build execgen_template +// +// This file is the execgen template for row_number.eg.go. It's formatted in a +// special way, so it's both valid Go and a valid text/template input. This +// permits editing this file with editor support. +// +// */}} + +package vecbuiltins + +import ( + "github.com/cockroachdb/cockroach/pkg/sql/exec/coldata" + "github.com/cockroachdb/cockroach/pkg/sql/exec/types" +) + +// {{/* +func _NEXT_ROW_NUMBER(hasPartition bool) { // */}} + // {{define "nextRowNumber"}} + + // {{ if $.HasPartition }} + if r.partitionColIdx == batch.Width() { + batch.AppendCol(types.Bool) + } else if r.partitionColIdx > batch.Width() { + panic("unexpected: column partitionColIdx is neither present nor the next to be appended") + } + partitionCol := batch.ColVec(r.partitionColIdx).Bool() + // {{ end }} + + if r.outputColIdx == batch.Width() { + batch.AppendCol(types.Int64) + } else if r.outputColIdx > batch.Width() { + panic("unexpected: column outputColIdx is neither present nor the next to be appended") + } + rowNumberCol := batch.ColVec(r.outputColIdx).Int64() + sel := batch.Selection() + if sel != nil { + for i := uint16(0); i < batch.Length(); i++ { + // {{ if $.HasPartition }} + if partitionCol[sel[i]] { + r.rowNumber = 1 + } + // {{ end }} + rowNumberCol[sel[i]] = r.rowNumber + r.rowNumber++ + } + } else { + for i := uint16(0); i < batch.Length(); i++ { + // {{ if $.HasPartition }} + if partitionCol[i] { + r.rowNumber = 1 + } + // {{ end }} + rowNumberCol[i] = r.rowNumber + r.rowNumber++ + } + } + // {{end}} + // {{/* +} + +// */}} + +func (r *rowNumberOp) nextBodyWithPartition(batch coldata.Batch) { + _NEXT_ROW_NUMBER(true) +} + +func (r *rowNumberOp) nextBodyNoPartition(batch coldata.Batch) { + _NEXT_ROW_NUMBER(false) +} diff --git a/pkg/sql/logictest/testdata/logic_test/exec_window b/pkg/sql/logictest/testdata/logic_test/exec_window index 17b332b714d2..e697a52c3882 100644 --- a/pkg/sql/logictest/testdata/logic_test/exec_window +++ b/pkg/sql/logictest/testdata/logic_test/exec_window @@ -75,7 +75,6 @@ SELECT a, b, dense_rank() OVER () FROM t ORDER BY b, a 0 b 1 1 b 1 - query ITI SELECT a, b, dense_rank() OVER (ORDER BY a) FROM t ORDER BY b, a ---- diff --git a/pkg/testutils/lint/lint_test.go b/pkg/testutils/lint/lint_test.go index c90e00988e71..50058392f5d6 100644 --- a/pkg/testutils/lint/lint_test.go +++ b/pkg/testutils/lint/lint_test.go @@ -1421,6 +1421,7 @@ func TestLint(t *testing.T) { stream.GrepNot("sql/.*exported func .* returns unexported type sql.planNode"), stream.GrepNot("struct field (XXX_NoUnkeyedLiteral|XXX_sizecache) should be"), stream.GrepNot("pkg/sql/types/types.go.* var Uuid should be UUID"), + stream.GrepNot("pkg/sql/exec/vecbuiltins/.*_tmpl.go.* use lowercase for SQL built-in functions"), ), func(s string) { t.Errorf("\n%s", s) }); err != nil {