colexec: fix an issue with builtin operators and a minor cleanup

pkg/sql/colexec/and_or_projection_tmpl.go

@@ -128,6 +128,13 @@ func (o *_OP_LOWERProjOp) Next(ctx context.Context) coldata.Batch {
 	o.leftFeedOp.batch = batch
 	batch = o.leftProjOpChain.Next(ctx)
 
+	if origLen == 0 {
+		// Run the right-side projection on the remaining tuples.
+		o.rightFeedOp.batch = batch
+		batch = o.rightProjOpChain.Next(ctx)
+		return batch
+	}
+
 	// Now we need to populate a selection vector on the batch in such a way that
 	// those tuples that we already know the result of logical operation for do
 	// not get the projection for the right side.
@@ -184,10 +191,6 @@ func (o *_OP_LOWERProjOp) Next(ctx context.Context) coldata.Batch {
 	}
 	batch.SetLength(origLen)
 
-	if origLen == 0 {
-		return batch
-	}
-
 	rightCol := batch.ColVec(o.rightIdx)
 	outputCol := batch.ColVec(o.outputIdx)
 	rightColVals := rightCol.Bool()

pkg/sql/colexec/builtin_funcs.go

@@ -102,6 +102,9 @@ func (b *defaultBuiltinFuncOperator) Next(ctx context.Context) coldata.Batch {
 			}
 		},
 	)
+	// Although we didn't change the length of the batch, it is necessary to set
+	// the length anyway (this helps maintaining the invariant of flat bytes).
+	batch.SetLength(n)
 	return batch
 }
 
@@ -185,6 +188,9 @@ func (s *substringFunctionOperator) Next(ctx context.Context) coldata.Batch {
 			}
 		},
 	)
+	// Although we didn't change the length of the batch, it is necessary to set
+	// the length anyway (this helps maintaining the invariant of flat bytes).
+	batch.SetLength(n)
 	return batch
 }
 

pkg/sql/colexec/case.go

@@ -114,9 +114,20 @@ func (c *caseOp) Next(ctx context.Context) coldata.Batch {
 		copy(c.origSel, sel)
 	}
 
-	outputCol := c.buffer.batch.Batch.ColVec(c.outputIdx)
-	var batch coldata.Batch
-	c.allocator.PerformOperation([]coldata.Vec{outputCol}, func() {
+	// It is possible that we have a typeless zero-length batch, and we cannot
+	// get a particular columnar vector from it, so we have special casing here
+	// for non-zero length batches.
+	// TODO(yuzefovich): remove it once we can short-circuit.
+	var (
+		outputCol coldata.Vec
+		batch     coldata.Batch
+		destVecs  []coldata.Vec
+	)
+	if origLen > 0 {
+		outputCol = c.buffer.batch.Batch.ColVec(c.outputIdx)
+		destVecs = []coldata.Vec{outputCol}
+	}
+	c.allocator.PerformOperation(destVecs, func() {
 		for i := range c.caseOps {
 			// Run the next case operator chain. It will project its THEN expression
 			// for all tuples that matched its WHEN expression and that were not
@@ -146,73 +157,77 @@ func (c *caseOp) Next(ctx context.Context) coldata.Batch {
 			// current case arm.
 			var subtractIdx int
 			var curIdx uint16
-			inputCol := c.buffer.batch.Batch.ColVec(c.thenIdxs[i])
-			// Copy the results into the output vector, using the toSubtract selection
-			// vector to copy only the elements that we actually wrote according to the
-			// current case arm.
+			if origLen > 0 {
+				inputCol := c.buffer.batch.Batch.ColVec(c.thenIdxs[i])
+				// Copy the results into the output vector, using the toSubtract selection
+				// vector to copy only the elements that we actually wrote according to the
+				// current case arm.
+				outputCol.Copy(
+					coldata.CopySliceArgs{
+						SliceArgs: coldata.SliceArgs{
+							ColType:     c.typ,
+							Src:         inputCol,
+							Sel:         toSubtract,
+							SrcStartIdx: 0,
+							SrcEndIdx:   uint64(len(toSubtract)),
+						},
+						SelOnDest: true,
+					})
+				if oldSel := c.buffer.batch.Batch.Selection(); oldSel != nil {
+					// We have an old selection vector, which represents the tuples that
+					// haven't yet been matched. Remove the ones that just matched from the
+					// old selection vector.
+					for i := range oldSel[:c.buffer.batch.Batch.Length()] {
+						if subtractIdx < len(toSubtract) && toSubtract[subtractIdx] == oldSel[i] {
+							// The ith element of the old selection vector matched the current one
+							// in toSubtract. Skip writing this element, removing it from the old
+							// selection vector.
+							subtractIdx++
+							continue
+						}
+						oldSel[curIdx] = oldSel[i]
+						curIdx++
+					}
+				} else {
+					// 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.buffer.batch.Batch.SetSelection(true)
+					oldSel = c.buffer.batch.Batch.Selection()
+					for i := uint16(0); i < c.buffer.batch.Batch.Length(); i++ {
+						if subtractIdx < len(toSubtract) && toSubtract[subtractIdx] == i {
+							subtractIdx++
+							continue
+						}
+						oldSel[curIdx] = i
+						curIdx++
+					}
+				}
+				c.buffer.batch.Batch.SetLength(curIdx)
+
+				// Now our selection vector is set to exclude all the things that have
+				// matched so far. Reset the buffer and run the next case arm.
+				c.buffer.rewind()
+			}
+		}
+		// Finally, run the else operator, which will project into all tuples that
+		// are remaining in the selection vector (didn't match any case arms). Once
+		// that's done, restore the original selection vector and return the batch.
+		batch = c.elseOp.Next(ctx)
+		if origLen > 0 {
+			inputCol := c.buffer.batch.Batch.ColVec(c.thenIdxs[len(c.thenIdxs)-1])
 			outputCol.Copy(
 				coldata.CopySliceArgs{
 					SliceArgs: coldata.SliceArgs{
 						ColType:     c.typ,
 						Src:         inputCol,
-						Sel:         toSubtract,
+						Sel:         batch.Selection(),
 						SrcStartIdx: 0,
-						SrcEndIdx:   uint64(len(toSubtract)),
+						SrcEndIdx:   uint64(batch.Length()),
 					},
 					SelOnDest: true,
 				})
-			if oldSel := c.buffer.batch.Batch.Selection(); oldSel != nil {
-				// We have an old selection vector, which represents the tuples that
-				// haven't yet been matched. Remove the ones that just matched from the
-				// old selection vector.
-				for i := range oldSel[:c.buffer.batch.Batch.Length()] {
-					if subtractIdx < len(toSubtract) && toSubtract[subtractIdx] == oldSel[i] {
-						// The ith element of the old selection vector matched the current one
-						// in toSubtract. Skip writing this element, removing it from the old
-						// selection vector.
-						subtractIdx++
-						continue
-					}
-					oldSel[curIdx] = oldSel[i]
-					curIdx++
-				}
-			} else {
-				// 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.buffer.batch.Batch.SetSelection(true)
-				oldSel = c.buffer.batch.Batch.Selection()
-				for i := uint16(0); i < c.buffer.batch.Batch.Length(); i++ {
-					if subtractIdx < len(toSubtract) && toSubtract[subtractIdx] == i {
-						subtractIdx++
-						continue
-					}
-					oldSel[curIdx] = i
-					curIdx++
-				}
-			}
-			c.buffer.batch.Batch.SetLength(curIdx)
-
-			// Now our selection vector is set to exclude all the things that have
-			// matched so far. Reset the buffer and run the next case arm.
-			c.buffer.rewind()
 		}
-		// Finally, run the else operator, which will project into all tuples that
-		// are remaining in the selection vector (didn't match any case arms). Once
-		// that's done, restore the original selection vector and return the batch.
-		batch = c.elseOp.Next(ctx)
-		inputCol := c.buffer.batch.Batch.ColVec(c.thenIdxs[len(c.thenIdxs)-1])
-		outputCol.Copy(
-			coldata.CopySliceArgs{
-				SliceArgs: coldata.SliceArgs{
-					ColType:     c.typ,
-					Src:         inputCol,
-					Sel:         batch.Selection(),
-					SrcStartIdx: 0,
-					SrcEndIdx:   uint64(batch.Length()),
-				},
-				SelOnDest: true,
-			})
 	})
 	batch.SetLength(origLen)
 	batch.SetSelection(origHasSel)

pkg/sql/colexec/cfetcher.go

@@ -857,8 +857,7 @@ func (rf *cFetcher) nextBatch(ctx context.Context) (coldata.Batch, error) {
 			return rf.machine.batch, nil
 
 		case stateFinished:
-			rf.machine.batch.SetLength(0)
-			return rf.machine.batch, nil
+			return coldata.ZeroBatch, nil
 		}
 	}
 }

pkg/sql/colexec/count.go

@@ -42,24 +42,22 @@ func NewCountOp(allocator *Allocator, input Operator) Operator {
 
 func (c *countOp) Init() {
 	c.input.Init()
-	// Our output is always just one row.
-	c.internalBatch.SetLength(1)
 	c.count = 0
 	c.done = false
 }
 
 func (c *countOp) Next(ctx context.Context) coldata.Batch {
 	c.internalBatch.ResetInternalBatch()
 	if c.done {
-		c.internalBatch.SetLength(0)
-		return c.internalBatch
+		return coldata.ZeroBatch
 	}
 	for {
 		bat := c.input.Next(ctx)
 		length := bat.Length()
 		if length == 0 {
 			c.done = true
 			c.internalBatch.ColVec(0).Int64()[0] = c.count
+			c.internalBatch.SetLength(1)
 			return c.internalBatch
 		}
 		c.count += int64(length)

pkg/sql/colexec/deselector.go

@@ -53,7 +53,6 @@ func (p *deselectorOp) Next(ctx context.Context) coldata.Batch {
 		return batch
 	}
 
-	p.output.SetLength(batch.Length())
 	p.output.ResetInternalBatch()
 	sel := batch.Selection()
 	p.allocator.PerformOperation(p.output.ColVecs(), func() {
@@ -72,5 +71,6 @@ func (p *deselectorOp) Next(ctx context.Context) coldata.Batch {
 			)
 		}
 	})
+	p.output.SetLength(batch.Length())
 	return p.output
 }

pkg/sql/colexec/operator.go

@@ -195,9 +195,8 @@ func (s *zeroOperator) Init() {
 
 func (s *zeroOperator) Next(ctx context.Context) coldata.Batch {
 	// TODO(solon): Can we avoid calling Next on the input at all?
-	next := s.input.Next(ctx)
-	next.SetLength(0)
-	return next
+	s.input.Next(ctx)
+	return coldata.ZeroBatch
 }
 
 type singleTupleNoInputOperator struct {
@@ -224,8 +223,7 @@ func (s *singleTupleNoInputOperator) Init() {
 func (s *singleTupleNoInputOperator) Next(ctx context.Context) coldata.Batch {
 	s.batch.ResetInternalBatch()
 	if s.nexted {
-		s.batch.SetLength(0)
-		return s.batch
+		return coldata.ZeroBatch
 	}
 	s.nexted = true
 	s.batch.SetLength(1)

pkg/sql/colexec/proj_const_ops_tmpl.go

@@ -106,6 +106,9 @@ func (p _OP_CONST_NAME) Next(ctx context.Context) coldata.Batch {
 	} else {
 		_SET_PROJECTION(false)
 	}
+	// Although we didn't change the length of the batch, it is necessary to set
+	// the length anyway (this helps maintaining the invariant of flat bytes).
+	batch.SetLength(n)
 	return batch
 }
 

pkg/sql/colexec/proj_non_const_ops_tmpl.go

@@ -108,6 +108,9 @@ func (p _OP_NAME) Next(ctx context.Context) coldata.Batch {
 		_SET_PROJECTION(false)
 	}
 
+	// Although we didn't change the length of the batch, it is necessary to set
+	// the length anyway (this helps maintaining the invariant of flat bytes).
+	batch.SetLength(n)
 	return batch
 }
 

pkg/sql/colexec/sort_chunks.go

@@ -365,6 +365,9 @@ func (s *chunker) prepareNextChunks(ctx context.Context) chunkerReadingState {
 // buffer appends all tuples in range [start,end) from s.batch to already
 // buffered tuples.
 func (s *chunker) buffer(start uint16, end uint16) {
+	if start == end {
+		return
+	}
 	s.allocator.PerformOperation(s.bufferedTuples.colVecs, func() {
 		for i := 0; i < len(s.bufferedTuples.colVecs); i++ {
 			s.bufferedTuples.colVecs[i].Append(

pkg/sql/colexec/sorttopk.go

@@ -203,8 +203,7 @@ func (t *topKSorter) emit() coldata.Batch {
 	toEmit := t.topK.Length() - t.emitted
 	if toEmit == 0 {
 		// We're done.
-		t.output.SetLength(0)
-		return t.output
+		return coldata.ZeroBatch
 	}
 	if toEmit > coldata.BatchSize() {
 		toEmit = coldata.BatchSize()

pkg/sql/colexec/testutils.go

@@ -80,10 +80,9 @@ func (s *RepeatableBatchSource) Next(context.Context) coldata.Batch {
 	s.internalBatch.SetSelection(s.sel != nil)
 	s.batchesReturned++
 	if s.batchesToReturn != 0 && s.batchesReturned > s.batchesToReturn {
-		s.internalBatch.SetLength(0)
-	} else {
-		s.internalBatch.SetLength(s.batchLen)
+		return coldata.ZeroBatch
 	}
+	s.internalBatch.SetLength(s.batchLen)
 	if s.sel != nil {
 		// Since selection vectors are mutable, to make sure that we return the
 		// batch with the given selection vector, we need to reset

pkg/sql/colexec/utils_test.go

@@ -224,6 +224,9 @@ func runTestsWithoutAllNullsInjection(
 		// output on its second Next call (we need the first call to Next to get a
 		// reference to a batch to modify), and a second time to modify the batch
 		// and verify that this does not change the operator output.
+		// NOTE: this test makes sense only if the operator returns two non-zero
+		// length batches (if not, we short-circuit the test since the operator
+		// doesn't have to restore anything on a zero-length batch).
 		var (
 			secondBatchHasSelection, secondBatchHasNulls bool
 			inputTypes                                   []coltypes.T
@@ -249,6 +252,9 @@ func runTestsWithoutAllNullsInjection(
 			}
 			ctx := context.Background()
 			b := op.Next(ctx)
+			if b.Length() == 0 {
+				return
+			}
 			if round == 1 {
 				if secondBatchHasSelection {
 					b.SetSelection(false)
@@ -265,6 +271,9 @@ func runTestsWithoutAllNullsInjection(
 				}
 			}
 			b = op.Next(ctx)
+			if b.Length() == 0 {
+				return
+			}
 			if round == 0 {
 				secondBatchHasSelection = b.Selection() != nil
 				secondBatchHasNulls = maybeHasNulls(b)
@@ -488,8 +497,7 @@ func (s *opTestInput) Init() {
 func (s *opTestInput) Next(context.Context) coldata.Batch {
 	s.batch.ResetInternalBatch()
 	if len(s.tuples) == 0 {
-		s.batch.SetLength(0)
-		return s.batch
+		return coldata.ZeroBatch
 	}
 	batchSize := s.batchSize
 	if len(s.tuples) < int(batchSize) {
@@ -700,8 +708,7 @@ func (s *opFixedSelTestInput) Next(context.Context) coldata.Batch {
 		}
 	} else {
 		if s.idx == uint16(len(s.sel)) {
-			s.batch.SetLength(0)
-			return s.batch
+			return coldata.ZeroBatch
 		}
 		batchSize = s.batchSize
 		if uint16(len(s.sel))-s.idx < batchSize {
@@ -1125,7 +1132,7 @@ func (c *chunkingBatchSource) Init() {
 
 func (c *chunkingBatchSource) Next(context.Context) coldata.Batch {
 	if c.curIdx >= c.len {
-		c.batch.SetLength(0)
+		return coldata.ZeroBatch
 	}
 	lastIdx := c.curIdx + uint64(coldata.BatchSize())
 	if lastIdx > c.len {