colexec: optimize CASE, AND, and OR projection operators

Previously, in CASE, AND, and OR projection operators we would always
allocate internal selection vectors of maximum capacity in the
constructor. This is inefficient when the query processes only small
data sets. This commit makes the allocations lazy and of tight sizes
based on the batch lengths.

Release note: None

pkg/sql/colexec/and_or_projection_tmpl.go

@@ -23,6 +23,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
+	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
 	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
@@ -118,7 +119,6 @@ func new_OP_TITLEProjOp(
 		leftIdx:          leftIdx,
 		rightIdx:         rightIdx,
 		outputIdx:        outputIdx,
-		origSel:          make([]int, coldata.BatchSize()),
 	}
 }
 
@@ -163,7 +163,8 @@ func (o *_OP_LOWERProjOp) Next(ctx context.Context) coldata.Batch {
 	}
 	usesSel := false
 	if sel := batch.Selection(); sel != nil {
-		copy(o.origSel[:origLen], sel[:origLen])
+		o.origSel = colexecutils.EnsureSelectionVectorLength(o.origSel, origLen)
+		copy(o.origSel, sel)
 		usesSel = true
 	}
 

pkg/sql/colexec/case.go

@@ -14,6 +14,7 @@ import (
 	"context"
 
 	"github.com/cockroachdb/cockroach/pkg/col/coldata"
+	"github.com/cockroachdb/cockroach/pkg/sql/colexec/colexecutils"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/colexecop"
 	"github.com/cockroachdb/cockroach/pkg/sql/colmem"
@@ -94,8 +95,6 @@ func NewCaseOp(
 		thenIdxs:  thenIdxs,
 		outputIdx: outputIdx,
 		typ:       typ,
-		origSel:   make([]int, coldata.BatchSize()),
-		prevSel:   make([]int, coldata.BatchSize()),
 	}
 }
 
@@ -115,15 +114,16 @@ func (c *caseOp) Next(ctx context.Context) coldata.Batch {
 	var origHasSel bool
 	if sel := c.buffer.batch.Selection(); sel != nil {
 		origHasSel = true
+		c.origSel = colexecutils.EnsureSelectionVectorLength(c.origSel, origLen)
 		copy(c.origSel, sel)
 	}
 
 	prevLen := origLen
 	prevHasSel := false
 	if sel := c.buffer.batch.Selection(); sel != nil {
 		prevHasSel = true
-		c.prevSel = c.prevSel[:origLen]
-		copy(c.prevSel[:origLen], sel[:origLen])
+		c.prevSel = colexecutils.EnsureSelectionVectorLength(c.prevSel, origLen)
+		copy(c.prevSel, sel)
 	}
 	outputCol := c.buffer.batch.ColVec(c.outputIdx)
 	if outputCol.MaybeHasNulls() {
@@ -199,8 +199,12 @@ func (c *caseOp) Next(ctx context.Context) coldata.Batch {
 					// No selection vector means there have been no matches yet, and we were
 					// considering the entire batch of tuples for this case arm. Make a new
 					// selection vector with all of the tuples but the ones that just matched.
-					c.prevSel = c.prevSel[:cap(c.prevSel)]
+					c.prevSel = colexecutils.EnsureSelectionVectorLength(c.prevSel, origLen)
 					for i := 0; i < origLen; i++ {
+						// Note that here we rely on the assumption that
+						// toSubtract is an increasing sequence (because our
+						// selection vectors are such) to optimize the
+						// subtraction.
 						if subtractIdx < len(toSubtract) && toSubtract[subtractIdx] == i {
 							subtractIdx++
 							continue

pkg/sql/colexec/colexecutils/utils.go

@@ -281,3 +281,12 @@ func HandleErrorFromDiskQueue(err error) {
 		colexecerror.InternalError(err)
 	}
 }
+
+// EnsureSelectionVectorLength returns an int slice that is guaranteed to have
+// the specified length. old is reused if possible but is *not* zeroed out.
+func EnsureSelectionVectorLength(old []int, length int) []int {
+	if cap(old) >= length {
+		return old[:length]
+	}
+	return make([]int, length)
+}