Skip to content

Commit

Permalink
exec: template out RANK and ROW_NUMBER
Browse files Browse the repository at this point in the history
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
  • Loading branch information
yuzefovich committed Jun 3, 2019
1 parent 4aa32dc commit 02332ed
Show file tree
Hide file tree
Showing 9 changed files with 400 additions and 212 deletions.
4 changes: 4 additions & 0 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -758,6 +758,8 @@ EXECGEN_TARGETS = \
pkg/sql/exec/sum_agg.eg.go \
pkg/sql/exec/tuples_differ.eg.go \
pkg/sql/exec/vec_comparators.eg.go \
pkg/sql/exec/vecbuiltins/rank.eg.go \
pkg/sql/exec/vecbuiltins/row_number.eg.go \
pkg/sql/exec/zerocolumns.eg.go

OPTGEN_TARGETS = \
Expand Down Expand Up @@ -1416,6 +1418,8 @@ 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/vec_comparators.eg.go: pkg/sql/exec/vec_comparators_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
pkg/sql/exec/zerocolumns.eg.go: pkg/sql/exec/zerocolumns_tmpl.go

$(EXECGEN_TARGETS): bin/execgen
Expand Down
3 changes: 2 additions & 1 deletion pkg/sql/exec/.gitignore
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
any_not_null_agg.eg.go
avg_agg.eg.go
coldata/vec.eg.go
colvec.eg.go
const.eg.go
distinct.eg.go
hashjoiner.eg.go
Expand All @@ -16,4 +15,6 @@ sort.eg.go
sum_agg.eg.go
tuples_differ.eg.go
vec_comparators.eg.go
vecbuiltins/rank.eg.go
vecbuiltins/row_number.eg.go
zerocolumns.eg.go
97 changes: 97 additions & 0 deletions pkg/sql/exec/execgen/cmd/execgen/rank_gen.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,97 @@
// 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
}

// UpdateRank is used to encompass the different logic between DENSE_RANK and
// RANK window functions when updating the internal state of rank operators. It
// is used by the template.
func (r rankTmplInfo) UpdateRank() string {
if r.Dense {
return fmt.Sprintf(
`r.rank++`,
)
}
return fmt.Sprintf(
`r.rank += r.rankIncrement
r.rankIncrement = 1`,
)
}

// UpdateRankIncrement is used to encompass the different logic between
// DENSE_RANK and RANK window functions when updating the internal state of
// rank operators. It is used by the template.
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")
}
45 changes: 45 additions & 0 deletions pkg/sql/exec/execgen/cmd/execgen/row_number_gen.go
Original file line number Diff line number Diff line change
@@ -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")
}
170 changes: 14 additions & 156 deletions pkg/sql/exec/vecbuiltins/rank.go
Original file line number Diff line number Diff line change
Expand Up @@ -15,33 +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
Expand All @@ -63,141 +42,20 @@ 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++
}
}
}
}
}
} 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 dense {
if partitionColIdx != -1 {
return &rankDense_true_HasPartition_true_Op{rankInitFields: initFields}, nil
}
return &rankDense_true_HasPartition_false_Op{rankInitFields: initFields}, nil
}
if partitionColIdx != -1 {
return &rankDense_false_HasPartition_true_Op{rankInitFields: initFields}, nil
}
return b
return &rankDense_false_HasPartition_false_Op{rankInitFields: initFields}, nil
}
Loading

0 comments on commit 02332ed

Please sign in to comment.