github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/sorttopk_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package colexec
    12  
    13  import (
    14  	"testing"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    17  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    19  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    20  )
    21  
    22  var topKSortTestCases []sortTestCase
    23  
    24  func init() {
    25  	topKSortTestCases = []sortTestCase{
    26  		{
    27  			description: "k < input length",
    28  			tuples:      tuples{{1}, {2}, {3}, {4}, {5}, {6}, {7}},
    29  			expected:    tuples{{1}, {2}, {3}},
    30  			typs:        []*types.T{types.Int},
    31  			ordCols:     []execinfrapb.Ordering_Column{{ColIdx: 0}},
    32  			k:           3,
    33  		},
    34  		{
    35  			description: "k > input length",
    36  			tuples:      tuples{{1}, {2}, {3}, {4}, {5}, {6}, {7}},
    37  			expected:    tuples{{1}, {2}, {3}, {4}, {5}, {6}, {7}},
    38  			typs:        []*types.T{types.Int},
    39  			ordCols:     []execinfrapb.Ordering_Column{{ColIdx: 0}},
    40  			k:           10,
    41  		},
    42  		{
    43  			description: "nulls",
    44  			tuples:      tuples{{1}, {2}, {nil}, {3}, {4}, {5}, {6}, {7}, {nil}},
    45  			expected:    tuples{{nil}, {nil}, {1}},
    46  			typs:        []*types.T{types.Int},
    47  			ordCols:     []execinfrapb.Ordering_Column{{ColIdx: 0}},
    48  			k:           3,
    49  		},
    50  		{
    51  			description: "descending",
    52  			tuples:      tuples{{0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5}, {1, 5}},
    53  			expected:    tuples{{0, 5}, {1, 5}, {0, 4}},
    54  			typs:        []*types.T{types.Int, types.Int},
    55  			ordCols: []execinfrapb.Ordering_Column{
    56  				{ColIdx: 1, Direction: execinfrapb.Ordering_Column_DESC},
    57  				{ColIdx: 0, Direction: execinfrapb.Ordering_Column_ASC},
    58  			},
    59  			k: 3,
    60  		},
    61  	}
    62  }
    63  
    64  func TestTopKSorter(t *testing.T) {
    65  	defer leaktest.AfterTest(t)()
    66  
    67  	for _, tc := range topKSortTestCases {
    68  		t.Run(tc.description, func(t *testing.T) {
    69  			runTests(t, []tuples{tc.tuples}, tc.expected, orderedVerifier, func(input []colexecbase.Operator) (colexecbase.Operator, error) {
    70  				return NewTopKSorter(testAllocator, input[0], tc.typs, tc.ordCols, tc.k), nil
    71  			})
    72  		})
    73  	}
    74  }