exec: template partition out for rank in vectorized engine

Release note: None

Makefile

@@ -753,6 +753,7 @@ EXECGEN_TARGETS = \
   pkg/sql/exec/sort.eg.go \
   pkg/sql/exec/sum_agg.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 = \
@@ -1409,6 +1410,7 @@ 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

pkg/sql/exec/execgen/cmd/execgen/rank_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 genRankOps(wr io.Writer) error {
+	d, err := ioutil.ReadFile("pkg/sql/exec/vecbuiltins/rank_tmpl.go")
+	if err != nil {
+		return err
+	}
+
+	s := string(d)
+
+	nextRank := makeFunctionRegex("_NEXT_RANK", 1)
+	s = nextRank.ReplaceAllString(s, `{{template "nextRank" buildDict "Global" $ "HasPartition" $1 }}`)
+
+	// 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
+	}
+
+	return tmpl.Execute(wr, struct{}{})
+}
+
+func init() {
+	registerGenerator(genRankOps, "rank.eg.go")
+}

pkg/sql/exec/vecbuiltins/rank.go

@@ -24,6 +24,7 @@ import (
 
 type rankOp struct {
 	input exec.Operator
+	batch coldata.Batch
 	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
@@ -75,128 +76,14 @@ func (r *rankOp) Init() {
 }
 
 func (r *rankOp) Next(ctx context.Context) coldata.Batch {
-	b := r.input.Next(ctx)
-	if b.Length() == 0 {
-		return b
+	r.batch = r.input.Next(ctx)
+	if r.batch.Length() == 0 {
+		return r.batch
 	}
 	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++
-						}
-					}
-				}
-			}
-		}
+		r.nextBodyWithPartition()
 	} 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++
-					}
-				}
-			}
-		}
+		r.nextBodyNoPartition()
 	}
-	return b
+	return r.batch
 }