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 {