fix(colexec): honor cancellation in window pipelines (#26080)

fix(colexec): honor cancellation in window pipelines

Approved by: @XuPeng-SH

Author: jiangxinmeng1

pkg/sql/colexec/partition/partition.go

@@ -30,6 +30,18 @@ import (
 
 const opName = "partition"
 
+const cancellationCheckInterval = 1024
+
+func checkCanceled(proc *process.Process, iteration int) error {
+	if iteration&(cancellationCheckInterval-1) != 0 {
+		return nil
+	}
+	if err, canceled := vm.CancelCheck(proc); canceled {
+		return err
+	}
+	return nil
+}
+
 func (partition *Partition) String(buf *bytes.Buffer) {
 	buf.WriteString(opName)
 	buf.WriteString(": partition([")
@@ -191,6 +203,9 @@ func (ctr *container) pickAndSend(proc *process.Process, result *vm.CallResult)
 	var cols []*vector.Vector
 	fromRemoveBatch := false
 	for {
+		if err = checkCanceled(proc, wholeLength); err != nil {
+			return false, err
+		}
 		if wholeLength == 0 || fromRemoveBatch {
 			choice = ctr.pickFirstRow()
 		} else {

pkg/sql/colexec/partition/partition_test.go

@@ -16,6 +16,7 @@ package partition
 
 import (
 	"bytes"
+	"context"
 	"testing"
 
 	"github.com/matrixorigin/matrixone/pkg/common/mpool"
@@ -82,6 +83,31 @@ func TestPartition(t *testing.T) {
 	}
 }
 
+func TestPartitionOutputHonorsCancellation(t *testing.T) {
+	proc := testutil.NewProcessWithMPool(t, "", mpool.MustNewZero())
+	arg := &Partition{
+		OrderBySpecs: []*plan.OrderBySpec{{Expr: newExpression(0, types.T_int32)}},
+	}
+	require.NoError(t, arg.Prepare(proc))
+
+	bat := colexec.MakeMockPartitionBatchs(1, proc.Mp())
+	arg.ctr.batchList = append(arg.ctr.batchList, bat)
+	require.NoError(t, arg.ctr.evaluateOrderColumn(proc, 0))
+	arg.ctr.indexList = []int64{0}
+	arg.ctr.status = eval
+
+	ctx, cancel := context.WithCancel(proc.Ctx)
+	proc.Ctx = ctx
+	cancel()
+
+	_, err := arg.Call(proc)
+	require.ErrorIs(t, err, context.Canceled)
+
+	arg.Free(proc, true, err)
+	proc.Free()
+	require.Equal(t, int64(0), proc.Mp().CurrNB())
+}
+
 func newExpression(pos int32, typeID types.T) *plan.Expr {
 	return &plan.Expr{
 		Expr: &plan.Expr_Col{

pkg/sql/colexec/window/value_window_test.go

@@ -15,6 +15,7 @@
 package window
 
 import (
+	"context"
 	"math"
 	"testing"
 
@@ -213,6 +214,43 @@ func TestProcessValueFunc_NthValue(t *testing.T) {
 	}
 }
 
+func TestProcessValueFuncHonorsCancellation(t *testing.T) {
+	testCases := []struct {
+		name string
+		spec func() *plan.Expr
+	}{
+		{name: "lag", spec: makeLagWindowSpec},
+		{name: "lead", spec: makeLeadWindowSpec},
+		{name: "first_value", spec: makeFirstValueWindowSpec},
+		{name: "last_value", spec: makeLastValueWindowSpec},
+		{name: "nth_value", spec: makeNthValueWindowSpec},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			mp := mpool.MustNewZero()
+			proc := testutil.NewProcessWithMPool(t, "", mp)
+			bat := makeInt32Batch(mp, []int32{10, 20})
+			ctr := &container{bat: bat}
+			ctr.aggVecs = make([]colexec.ExprEvalVector, 1)
+			ctr.aggVecs[0].Vec = []*vector.Vector{bat.Vecs[0]}
+			arg := &Window{WinSpecList: []*plan.Expr{tc.spec()}}
+
+			ctx, cancel := context.WithCancel(proc.Ctx)
+			proc.Ctx = ctx
+			cancel()
+
+			result, err := ctr.processValueFunc(0, arg, proc)
+			require.ErrorIs(t, err, context.Canceled)
+			require.Nil(t, result)
+
+			bat.Clean(mp)
+			proc.Free()
+			require.Equal(t, int64(0), mp.CurrNB())
+		})
+	}
+}
+
 func TestProcessValueFunc_ErrorPathFreesLocalResult(t *testing.T) {
 	mp := mpool.MustNewZero()
 	resultMP, err := mpool.NewMPool("value-window-error", 1<<20, mpool.NoFixed)

pkg/sql/colexec/window/window.go

@@ -37,6 +37,22 @@ import (
 
 const opName = "window"
 
+// A Window call can spend a long time inside one frame evaluation, especially
+// for running frames whose aggregate work is quadratic in the partition size.
+// Keep the cancellation polling overhead bounded while still allowing KILL
+// QUERY / KILL CONNECTION to interrupt that work promptly.
+const cancellationCheckInterval = 1024
+
+func checkCanceled(proc *process.Process, iteration int) error {
+	if iteration&(cancellationCheckInterval-1) != 0 {
+		return nil
+	}
+	if err, canceled := vm.CancelCheck(proc); canceled {
+		return err
+	}
+	return nil
+}
+
 func (window *Window) String(buf *bytes.Buffer) {
 	buf.WriteString(opName)
 	buf.WriteString(": window")
@@ -80,6 +96,9 @@ func (window *Window) Call(proc *process.Process) (vm.CallResult, error) {
 	ctr := &window.ctr
 
 	for {
+		if err, canceled := vm.CancelCheck(proc); canceled {
+			return vm.CancelResult, err
+		}
 		switch ctr.status {
 		case receiveAll:
 			for {
@@ -154,6 +173,9 @@ func (window *Window) Call(proc *process.Process) (vm.CallResult, error) {
 			}
 			// calculate
 			for i, w := range window.WinSpecList {
+				if err = checkCanceled(proc, 0); err != nil {
+					return result, err
+				}
 				// sort and partitions
 				if window.Fs = makeOrderBy(w); window.Fs != nil {
 					if len(ctr.orderVecs) == 0 {
@@ -277,7 +299,13 @@ func (ctr *container) processFunc(idx int, ap *Window, proc *process.Process, an
 
 		o := 0
 		for p := 1; p < len(ctr.ps); p++ {
+			if err = checkCanceled(proc, p); err != nil {
+				return err
+			}
 			for ; o < len(ctr.os); o++ {
+				if err = checkCanceled(proc, o); err != nil {
+					return err
+				}
 
 				if ctr.os[o] <= ctr.ps[p] {
 					// For NTILE, pass both os vector and bucket count vector
@@ -302,6 +330,9 @@ func (ctr *container) processFunc(idx int, ap *Window, proc *process.Process, an
 	} else {
 		// plan.Function_AGG
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return err
+			}
 
 			start, end := 0, n
 
@@ -326,6 +357,9 @@ func (ctr *container) processFunc(idx int, ap *Window, proc *process.Process, an
 			}
 
 			for k := left; k < right; k++ {
+				if err = checkCanceled(proc, k-left); err != nil {
+					return err
+				}
 				if err = ctr.batAggs[idx].Fill(j, k, ctr.aggVecs[idx].Vec); err != nil {
 					return err
 				}
@@ -390,6 +424,9 @@ func (ctr *container) processValueFunc(idx int, ap *Window, proc *process.Proces
 			defaultVec = ctr.aggVecs[idx].Vec[2]
 		}
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return nil, err
+			}
 			offset, ok := constOffset, constOK
 			if offsetVec != nil && !offsetVec.IsConst() {
 				offset, ok = getInt64FromVec(offsetVec, j)
@@ -430,6 +467,9 @@ func (ctr *container) processValueFunc(idx int, ap *Window, proc *process.Proces
 			defaultVec = ctr.aggVecs[idx].Vec[2]
 		}
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return nil, err
+			}
 			offset, ok := constOffset, constOK
 			if offsetVec != nil && !offsetVec.IsConst() {
 				offset, ok = getInt64FromVec(offsetVec, j)
@@ -458,6 +498,9 @@ func (ctr *container) processValueFunc(idx int, ap *Window, proc *process.Proces
 
 	case "first_value":
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return nil, err
+			}
 			start, end := 0, n
 			if ctr.ps != nil {
 				start, end = buildPartitionInterval(ctr.ps, j, n)
@@ -485,6 +528,9 @@ func (ctr *container) processValueFunc(idx int, ap *Window, proc *process.Proces
 
 	case "last_value":
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return nil, err
+			}
 			start, end := 0, n
 			if ctr.ps != nil {
 				start, end = buildPartitionInterval(ctr.ps, j, n)
@@ -521,6 +567,9 @@ func (ctr *container) processValueFunc(idx int, ap *Window, proc *process.Proces
 			}
 		}
 		for j := 0; j < n; j++ {
+			if err = checkCanceled(proc, j); err != nil {
+				return nil, err
+			}
 			nthVal, ok := constNth, constOK
 			if nthVec != nil && !nthVec.IsConst() {
 				nthVal, ok = getInt64FromVec(nthVec, j)
@@ -834,15 +883,24 @@ func (ctr *container) processOrder(idx int, ap *Window, bat *batch.Batch, proc *
 	// }
 	ctr.sels = make([]int64, rowCount)
 	for i := 0; i < rowCount; i++ {
+		if err := checkCanceled(proc, i); err != nil {
+			return false, err
+		}
 		ctr.sels[i] = int64(i)
 	}
 
 	// skip sort for const vector
 	if !ovec.IsConst() {
+		if err := checkCanceled(proc, 0); err != nil {
+			return false, err
+		}
 		nullCnt := ovec.GetNulls().Count()
 		if nullCnt < ovec.Length() {
 			sort.Sort(ctr.desc[0], ctr.nullsLast[0], nullCnt > 0, ctr.sels, ovec)
 		}
+		if err := checkCanceled(proc, 0); err != nil {
+			return false, err
+		}
 	}
 
 	ps := make([]int64, 0, 16)
@@ -853,6 +911,9 @@ func (ctr *container) processOrder(idx int, ap *Window, bat *batch.Batch, proc *
 
 	i, j := 1, len(ctr.orderVecs)
 	for ; i < j; i++ {
+		if err := checkCanceled(proc, 0); err != nil {
+			return false, err
+		}
 		desc := ctr.desc[i]
 		nullsLast := ctr.nullsLast[i]
 		ps = partition.Partition(ctr.sels, ds, ps, ovec)
@@ -862,6 +923,9 @@ func (ctr *container) processOrder(idx int, ap *Window, bat *batch.Batch, proc *
 			nullCnt := vec.GetNulls().Count()
 			if nullCnt < vec.Length() {
 				for i, j := 0, len(ps); i < j; i++ {
+					if err := checkCanceled(proc, i); err != nil {
+						return false, err
+					}
 					if i == j-1 {
 						sort.Sort(desc, nullsLast, nullCnt > 0, ctr.sels[ps[i]:], vec)
 					} else {
@@ -870,6 +934,9 @@ func (ctr *container) processOrder(idx int, ap *Window, bat *batch.Batch, proc *
 				}
 			}
 		}
+		if err := checkCanceled(proc, 0); err != nil {
+			return false, err
+		}
 		ovec = vec
 		if n == i {
 			ctr.ps = make([]int64, len(ps))