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 May 17, 2019
1 parent bd41ffa commit 48d7011
Show file tree
Hide file tree
Showing 12 changed files with 783 additions and 211 deletions.
6 changes: 5 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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 \
Expand Down Expand Up @@ -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.
Expand Down
91 changes: 91 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,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")
}
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")
}
2 changes: 2 additions & 0 deletions pkg/sql/exec/vecbuiltins/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
rank.eg.go
row_number.eg.go
Loading

0 comments on commit 48d7011

Please sign in to comment.