colexec: implement lag and lead window functions in vectorized engine

`lag` and `lead` are window functions that return an expression
evaluated at some row `offset` rows from the current row, defaulting
to `default` if such a row does not exist in the partition.

This patch provides vectorized implementations of the `lag` and `lead`
window functions. The logic to buffer each partition into a queue is
now separate from calculation of the window function itself so that
`lag`, `lead`, and `ntile` (and others) can share common logic. As
before, each batch is emitted as soon as its output column has been
entirely filled.

See cockroachdb#37035

Release note (sql change): the vectorized engine now supports the lag
and lead window functions.

Makefile

@@ -909,6 +909,8 @@ EXECGEN_TARGETS = \
   pkg/sql/colexec/colexecsel/default_cmp_sel_ops.eg.go \
   pkg/sql/colexec/colexecsel/selection_ops.eg.go \
   pkg/sql/colexec/colexecsel/sel_like_ops.eg.go \
+  pkg/sql/colexec/colexecwindow/lag.eg.go \
+  pkg/sql/colexec/colexecwindow/lead.eg.go \
   pkg/sql/colexec/colexecwindow/ntile.eg.go \
   pkg/sql/colexec/colexecwindow/rank.eg.go \
   pkg/sql/colexec/colexecwindow/relative_rank.eg.go \

pkg/col/coldata/BUILD.bazel

@@ -44,6 +44,7 @@ go_test(
     deps = [
         "//pkg/col/coldatatestutils",
         "//pkg/sql/colconv",
+        "//pkg/sql/randgen",
         "//pkg/sql/types",
         "//pkg/testutils/buildutil",
         "//pkg/util/leaktest",

pkg/col/coldata/bytes_test.go

@@ -471,7 +471,7 @@ func TestToArrowSerializationFormat(t *testing.T) {
 	rng, _ := randutil.NewPseudoRand()
 	nullChance := 0.2
 	maxStringLength := 10
-	numElements := rng.Intn(BatchSize())
+	numElements := 1 + rng.Intn(BatchSize())
 
 	b := NewBytes(numElements)
 	for i := 0; i < numElements; i++ {
@@ -485,7 +485,10 @@ func TestToArrowSerializationFormat(t *testing.T) {
 	b.UpdateOffsetsToBeNonDecreasing(numElements)
 
 	startIdx := rng.Intn(numElements)
-	endIdx := 1 + startIdx + rng.Intn(numElements-startIdx-1)
+	endIdx := startIdx + rng.Intn(numElements-startIdx)
+	if endIdx == startIdx {
+		endIdx++
+	}
 	wind := b.Window(startIdx, endIdx)
 
 	data, offsets := wind.ToArrowSerializationFormat(wind.Len())

pkg/sql/colexec/colbuilder/execplan.go

@@ -1191,8 +1191,9 @@ func NewColOperator(
 				argIdxs := make([]int, len(wf.ArgsIdxs))
 				for i, idx := range wf.ArgsIdxs {
 					// Retrieve the type of each argument and perform any necessary casting.
-					expectedType := colexecwindow.GetWindowFnArgType(*wf.Func.WindowFunc, i)
-					if !expectedType.Identical(typs[idx]) {
+					needsCast, expectedType := colexecwindow.WindowFnArgNeedsCast(
+						*wf.Func.WindowFunc, typs[idx], i)
+					if needsCast {
 						// We must cast to the expected argument type.
 						castIdx := len(typs)
 						input, err = colexecbase.GetCastOperator(
@@ -1272,10 +1273,8 @@ func NewColOperator(
 					// making sure that we stay within the memory limit, and
 					// they will fall back to disk if necessary.
 					opName := opNamePrefix + "relative-rank"
-					unlimitedAllocator := colmem.NewAllocator(
-						ctx, result.createBufferingUnlimitedMemAccount(ctx, flowCtx, opName, spec.ProcessorID), factory,
-					)
-					diskAcc := result.createDiskAccount(ctx, flowCtx, opName, spec.ProcessorID)
+					unlimitedAllocator, diskAcc := result.getDiskBackedWindowFnFields(
+						ctx, opName, flowCtx, spec.ProcessorID, factory)
 					result.Root, err = colexecwindow.NewRelativeRankOperator(
 						unlimitedAllocator, execinfra.GetWorkMemLimit(flowCtx), args.DiskQueueCfg,
 						args.FDSemaphore, input, typs, windowFn, wf.Ordering.Columns,
@@ -1293,13 +1292,37 @@ func NewColOperator(
 					// making sure that we stay within the memory limit, and
 					// they will fall back to disk if necessary.
 					opName := opNamePrefix + "ntile"
-					unlimitedAllocator := colmem.NewAllocator(
-						ctx, result.createBufferingUnlimitedMemAccount(ctx, flowCtx, opName, spec.ProcessorID), factory,
-					)
-					diskAcc := result.createDiskAccount(ctx, flowCtx, opName, spec.ProcessorID)
+					unlimitedAllocator, diskAcc := result.getDiskBackedWindowFnFields(
+						ctx, opName, flowCtx, spec.ProcessorID, factory)
 					result.Root = colexecwindow.NewNTileOperator(
 						unlimitedAllocator, execinfra.GetWorkMemLimit(flowCtx), args.DiskQueueCfg,
-						args.FDSemaphore, input, typs, outputIdx, partitionColIdx, argIdxs[0], diskAcc,
+						args.FDSemaphore, diskAcc, input, typs, outputIdx, partitionColIdx, argIdxs[0],
+					)
+				case execinfrapb.WindowerSpec_LAG:
+					opName := opNamePrefix + "lag"
+					unlimitedAllocator, diskAcc := result.getDiskBackedWindowFnFields(
+						ctx, opName, flowCtx, spec.ProcessorID, factory)
+					// Lag operators need an extra allocator.
+					bufferAllocator := colmem.NewAllocator(
+						ctx, result.createBufferingUnlimitedMemAccount(ctx, flowCtx, opName, spec.ProcessorID), factory,
+					)
+					result.Root, err = colexecwindow.NewLagOperator(
+						unlimitedAllocator, bufferAllocator, execinfra.GetWorkMemLimit(flowCtx),
+						args.DiskQueueCfg, args.FDSemaphore, diskAcc, input, typs,
+						outputIdx, partitionColIdx, argIdxs[0], argIdxs[1], argIdxs[2],
+					)
+				case execinfrapb.WindowerSpec_LEAD:
+					opName := opNamePrefix + "lead"
+					unlimitedAllocator, diskAcc := result.getDiskBackedWindowFnFields(
+						ctx, opName, flowCtx, spec.ProcessorID, factory)
+					// Lead operators need an extra allocator.
+					bufferAllocator := colmem.NewAllocator(
+						ctx, result.createBufferingUnlimitedMemAccount(ctx, flowCtx, opName, spec.ProcessorID), factory,
+					)
+					result.Root, err = colexecwindow.NewLeadOperator(
+						unlimitedAllocator, bufferAllocator, execinfra.GetWorkMemLimit(flowCtx),
+						args.DiskQueueCfg, args.FDSemaphore, diskAcc, input, typs,
+						outputIdx, partitionColIdx, argIdxs[0], argIdxs[1], argIdxs[2],
 					)
 				default:
 					return r, errors.AssertionFailedf("window function %s is not supported", wf.String())
@@ -1649,6 +1672,22 @@ func (r opResult) updateWithPostProcessResult(ppr postProcessResult) {
 	copy(r.ColumnTypes, ppr.ColumnTypes)
 }
 
+// getDiskBackedWindowFnFields constructs fields common to all disk-backed
+// buffering window operators.
+func (r opResult) getDiskBackedWindowFnFields(
+	ctx context.Context,
+	opName string,
+	flowCtx *execinfra.FlowCtx,
+	processorID int32,
+	factory coldata.ColumnFactory,
+) (*colmem.Allocator, *mon.BoundAccount) {
+	unlimitedAllocator := colmem.NewAllocator(
+		ctx, r.createBufferingUnlimitedMemAccount(ctx, flowCtx, opName, processorID), factory,
+	)
+	diskAcc := r.createDiskAccount(ctx, flowCtx, opName, processorID)
+	return unlimitedAllocator, diskAcc
+}
+
 // planFilterExpr creates all operators to implement filter expression.
 func planFilterExpr(
 	ctx context.Context,

pkg/sql/colexec/colexecutils/BUILD.bazel

@@ -7,6 +7,7 @@ go_library(
         "cancel_checker.go",
         "deselector.go",
         "operator.go",
+        "spilling_buffer.go",
         "spilling_queue.go",
         "utils.go",
     ],
@@ -37,6 +38,7 @@ go_test(
         "dep_test.go",
         "deselector_test.go",
         "main_test.go",
+        "spilling_buffer_test.go",
         "spilling_queue_test.go",
     ],
     embed = [":colexecutils"],
@@ -51,6 +53,7 @@ go_test(
         "//pkg/sql/colexecop",
         "//pkg/sql/colmem",
         "//pkg/sql/execinfra",
+        "//pkg/sql/randgen",
         "//pkg/sql/sem/tree",
         "//pkg/sql/types",
         "//pkg/testutils/buildutil",