exec: add templating for RANK and ROW_NUMBER window functions

Makefile

@@ -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 = \
@@ -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

pkg/sql/exec/.gitignore

@@ -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
@@ -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

pkg/sql/exec/execgen/cmd/execgen/rank_gen.go

@@ -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")
+}

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")
+}

pkg/sql/exec/vecbuiltins/rank.go

@@ -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
@@ -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
 }