From b30a916e5fc3d0b8076899a2b683762f1865bb87 Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Tue, 29 Mar 2022 04:31:33 +0000
Subject: [PATCH] sql: propagate limit for top K sort correctly in tests

In 22.1 time frame we started propagating the value of K for top K sort
in the spec of the processor, and not in the post-processing spec, but
we forgot to update some of the tests accordingly.

Release note: None
---
 pkg/sql/colexec/external_sort_test.go      | 4 +---
 pkg/sql/distsql/columnar_operators_test.go | 6 +++---
 pkg/sql/rowexec/sorter_test.go             | 7 ++++---
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/colexec/external_sort_test.go b/pkg/sql/colexec/external_sort_test.go
index 38ad233f6208..8a41b32ecac3 100644
--- a/pkg/sql/colexec/external_sort_test.go
+++ b/pkg/sql/colexec/external_sort_test.go
@@ -454,15 +454,13 @@ func createDiskBackedSorter(
 	sorterSpec := &execinfrapb.SorterSpec{
 		OutputOrdering:   execinfrapb.Ordering{Columns: ordCols},
 		OrderingMatchLen: uint32(matchLen),
+		Limit:            int64(k),
 	}
 	spec := &execinfrapb.ProcessorSpec{
 		Input: []execinfrapb.InputSyncSpec{{ColumnTypes: typs}},
 		Core: execinfrapb.ProcessorCoreUnion{
 			Sorter: sorterSpec,
 		},
-		Post: execinfrapb.PostProcessSpec{
-			Limit: k,
-		},
 		ResultTypes: typs,
 	}
 	args := &colexecargs.NewColOperatorArgs{
diff --git a/pkg/sql/distsql/columnar_operators_test.go b/pkg/sql/distsql/columnar_operators_test.go
index 6204e7d35582..3204b4a49384 100644
--- a/pkg/sql/distsql/columnar_operators_test.go
+++ b/pkg/sql/distsql/columnar_operators_test.go
@@ -548,15 +548,15 @@ func TestSorterAgainstProcessor(t *testing.T) {
 					sorterSpec := &execinfrapb.SorterSpec{
 						OutputOrdering: execinfrapb.Ordering{Columns: orderingCols},
 					}
-					var limit, offset uint64
+					var offset uint64
 					if topK > 0 {
 						offset = uint64(rng.Intn(int(topK)))
-						limit = topK - offset
+						sorterSpec.Limit = int64(topK - offset)
 					}
 					pspec := &execinfrapb.ProcessorSpec{
 						Input:       []execinfrapb.InputSyncSpec{{ColumnTypes: inputTypes}},
 						Core:        execinfrapb.ProcessorCoreUnion{Sorter: sorterSpec},
-						Post:        execinfrapb.PostProcessSpec{Limit: limit, Offset: offset},
+						Post:        execinfrapb.PostProcessSpec{Offset: offset},
 						ResultTypes: inputTypes,
 					}
 					args := verifyColOperatorArgs{
diff --git a/pkg/sql/rowexec/sorter_test.go b/pkg/sql/rowexec/sorter_test.go
index d6fd0d336cf1..dcb8f7b3daf9 100644
--- a/pkg/sql/rowexec/sorter_test.go
+++ b/pkg/sql/rowexec/sorter_test.go
@@ -455,14 +455,15 @@ func BenchmarkSortLimit(b *testing.B) {
 	const numRows = 1 << 16
 	b.Run(fmt.Sprintf("rows=%d", numRows), func(b *testing.B) {
 		input := execinfra.NewRepeatableRowSource(types.TwoIntCols, randgen.MakeRandIntRows(rng, numRows, numCols))
-		for _, limit := range []uint64{1 << 4, 1 << 8, 1 << 12, 1 << 16} {
-			post := execinfrapb.PostProcessSpec{Limit: limit}
+		for _, limit := range []int64{1 << 4, 1 << 8, 1 << 12, 1 << 16} {
+			spec.Limit = limit
 			b.Run(fmt.Sprintf("Limit=%d", limit), func(b *testing.B) {
 				b.SetBytes(int64(numRows * numCols * 8))
 				b.ResetTimer()
 				for i := 0; i < b.N; i++ {
 					s, err := newSorter(
-						context.Background(), &flowCtx, 0 /* processorID */, &spec, input, &post, &rowDisposer{},
+						context.Background(), &flowCtx, 0, /* processorID */
+						&spec, input, &execinfrapb.PostProcessSpec{Limit: 0}, &rowDisposer{},
 					)
 					if err != nil {
 						b.Fatal(err)