Merge cockroachdb#136631

136631: colexec: add cancellation check to top K sort and merge join r=yuzefovich a=yuzefovich

This commit adds the cancellation check to the top K sorter and the merge joiner to be performed on every input batch they read. I think that these two are the only buffering operators that currently don't have any cancellation checks:
- for hash joins and hash aggregation we have the check when performing the hashing
- for ordered aggregator, buffered window functions, cross join, external sort and other disk-backed operators we do the check on each input batch
- for general sort we do the check in the PDQ sort of each column.

Fixes: cockroachdb#136457.

Release note (bug fix): CockroachDB now better respects `statement_timeout` limit on queries involving the top K sort and merge join operations.

Co-authored-by: Yahor Yuzefovich <[email protected]>

pkg/sql/colexec/colexecjoin/mergejoiner.go

@@ -552,6 +552,8 @@ type mergeJoinBase struct {
 	left               mergeJoinInput
 	right              mergeJoinInput
 
+	cancelChecker colexecutils.CancelChecker
+
 	// Output buffer definition.
 	output         coldata.Batch
 	outputCapacity int
@@ -574,6 +576,7 @@ var _ colexecop.Closer = &mergeJoinBase{}
 
 func (o *mergeJoinBase) Reset(ctx context.Context) {
 	o.TwoInputInitHelper.Reset(ctx)
+	o.cancelChecker.Init(ctx)
 	o.state = mjEntry
 	o.bufferedGroup.helper.Reset(ctx)
 	o.proberState.lBatch = nil
@@ -597,6 +600,7 @@ func (o *mergeJoinBase) Init(ctx context.Context) {
 		o.memoryLimit, o.diskQueueCfg, o.fdSemaphore, o.diskAcc, o.diskQueueMemAcc,
 	)
 	o.bufferedGroup.helper.init(o.Ctx)
+	o.cancelChecker.Init(o.Ctx)
 
 	o.builderState.lGroups = make([]group, 1)
 	o.builderState.rGroups = make([]group, 1)
@@ -739,6 +743,7 @@ func (o *mergeJoinBase) sourceFinished() bool {
 // and updates the probing and buffered group states accordingly.
 func (o *mergeJoinBase) continueLeftBufferedGroup() {
 	// Get the next batch from the left.
+	o.cancelChecker.CheckEveryCall()
 	o.proberState.lIdx, o.proberState.lBatch = 0, o.InputOne.Next()
 	o.proberState.lLength = o.proberState.lBatch.Length()
 	o.bufferedGroup.leftGroupStartIdx = 0
@@ -805,6 +810,7 @@ func (o *mergeJoinBase) finishRightBufferedGroup() {
 // (or until we exhaust the right input).
 func (o *mergeJoinBase) completeRightBufferedGroup() {
 	// Get the next batch from the right.
+	o.cancelChecker.CheckEveryCall()
 	o.proberState.rIdx, o.proberState.rBatch = 0, o.InputTwo.Next()
 	o.proberState.rLength = o.proberState.rBatch.Length()
 	// The right input has been fully exhausted.
@@ -870,6 +876,7 @@ func (o *mergeJoinBase) completeRightBufferedGroup() {
 			// The buffered group is still not complete which means that we have
 			// just appended all the tuples from batch to it, so we need to get a
 			// fresh batch from the input.
+			o.cancelChecker.CheckEveryCall()
 			o.proberState.rIdx, o.proberState.rBatch = 0, o.InputTwo.Next()
 			o.proberState.rLength = o.proberState.rBatch.Length()
 			if o.proberState.rLength == 0 {

pkg/sql/colexec/colexecjoin/mergejoiner_tmpl.go

@@ -1366,10 +1366,12 @@ func (o *mergeJoin_JOIN_TYPE_STRINGOp) Next() coldata.Batch {
 			// If this is the first batch or we're done with the current batch,
 			// get the next batch.
 			if o.proberState.lBatch == nil || (o.proberState.lLength != 0 && o.proberState.lIdx == o.proberState.lLength) {
+				o.cancelChecker.CheckEveryCall()
 				o.proberState.lIdx, o.proberState.lBatch = 0, o.InputOne.Next()
 				o.proberState.lLength = o.proberState.lBatch.Length()
 			}
 			if o.proberState.rBatch == nil || (o.proberState.rLength != 0 && o.proberState.rIdx == o.proberState.rLength) {
+				o.cancelChecker.CheckEveryCall()
 				o.proberState.rIdx, o.proberState.rBatch = 0, o.InputTwo.Next()
 				o.proberState.rLength = o.proberState.rBatch.Length()
 			}

pkg/sql/colexec/sorttopk.go

@@ -117,6 +117,8 @@ type topKSorter struct {
 	emitted int
 	output  coldata.Batch
 
+	cancelChecker colexecutils.CancelChecker
+
 	exportedFromTopK  int
 	exportedFromBatch int
 	windowedBatch     coldata.Batch
@@ -150,6 +152,7 @@ func (t *topKSorter) Init(ctx context.Context) {
 		t.orderState.distincter.Init(t.Ctx)
 		t.orderState.group = make([]int, t.k)
 	}
+	t.cancelChecker.Init(t.Ctx)
 }
 
 func (t *topKSorter) Next() coldata.Batch {

pkg/sql/colexec/sorttopk_tmpl.go

@@ -27,6 +27,7 @@ import (
 // execgen:template<partialOrder>
 // execgen:inline
 func nextBatch(t *topKSorter, partialOrder bool) {
+	t.cancelChecker.CheckEveryCall()
 	t.inputBatch = t.Input.Next()
 	if partialOrder {
 		t.orderState.distincterInput.SetBatch(t.inputBatch)