github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/simple_project_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  	"sync"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    19  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    20  	"github.com/stretchr/testify/require"
    21  )
    22  
    23  func TestSimpleProjectOp(t *testing.T) {
    24  	defer leaktest.AfterTest(t)()
    25  	tcs := []struct {
    26  		tuples     tuples
    27  		expected   tuples
    28  		colsToKeep []uint32
    29  	}{
    30  		{
    31  			colsToKeep: []uint32{0, 2},
    32  			tuples: tuples{
    33  				{1, 2, 3},
    34  				{1, 2, 3},
    35  			},
    36  			expected: tuples{
    37  				{1, 3},
    38  				{1, 3},
    39  			},
    40  		},
    41  		{
    42  			colsToKeep: []uint32{0, 1},
    43  			tuples: tuples{
    44  				{1, 2, 3},
    45  				{1, 2, 3},
    46  			},
    47  			expected: tuples{
    48  				{1, 2},
    49  				{1, 2},
    50  			},
    51  		},
    52  		{
    53  			colsToKeep: []uint32{2, 1},
    54  			tuples: tuples{
    55  				{1, 2, 3},
    56  				{1, 2, 3},
    57  			},
    58  			expected: tuples{
    59  				{3, 2},
    60  				{3, 2},
    61  			},
    62  		},
    63  	}
    64  	for _, tc := range tcs {
    65  		runTests(t, []tuples{tc.tuples}, tc.expected, orderedVerifier, func(input []colexecbase.Operator) (colexecbase.Operator, error) {
    66  			return NewSimpleProjectOp(input[0], len(tc.tuples[0]), tc.colsToKeep), nil
    67  		})
    68  	}
    69  
    70  	// Empty projection. The all nulls injection test case will also return
    71  	// nothing.
    72  	runTestsWithoutAllNullsInjection(t, []tuples{{{1, 2, 3}, {1, 2, 3}}}, nil /* typs */, tuples{{}, {}}, orderedVerifier,
    73  		func(input []colexecbase.Operator) (colexecbase.Operator, error) {
    74  			return NewSimpleProjectOp(input[0], 3 /* numInputCols */, nil), nil
    75  		})
    76  
    77  	t.Run("RedundantProjectionIsNotPlanned", func(t *testing.T) {
    78  		typs := []*types.T{types.Int, types.Int}
    79  		input := newFiniteBatchSource(testAllocator.NewMemBatch(typs), typs, 1 /* usableCount */)
    80  		projectOp := NewSimpleProjectOp(input, len(typs), []uint32{0, 1})
    81  		require.IsType(t, input, projectOp)
    82  	})
    83  }
    84  
    85  // TestSimpleProjectOpWithUnorderedSynchronizer sets up the following
    86  // structure:
    87  //
    88  //  input 1 --
    89  //            | --> unordered synchronizer --> simpleProjectOp --> constInt64Op
    90  //  input 2 --
    91  //
    92  // and makes sure that the output is as expected. The idea is to test
    93  // simpleProjectOp in case when it receives multiple "different internally"
    94  // batches. See #45686 for detailed discussion.
    95  func TestSimpleProjectOpWithUnorderedSynchronizer(t *testing.T) {
    96  	defer leaktest.AfterTest(t)()
    97  	inputTypes := []*types.T{types.Bytes, types.Float}
    98  	constVal := int64(42)
    99  	var wg sync.WaitGroup
   100  	inputTuples := []tuples{
   101  		{{"a", 1.0}, {"aa", 10.0}},
   102  		{{"b", 2.0}, {"bb", 20.0}},
   103  	}
   104  	expected := tuples{
   105  		{"a", constVal},
   106  		{"aa", constVal},
   107  		{"b", constVal},
   108  		{"bb", constVal},
   109  	}
   110  	runTestsWithoutAllNullsInjection(t, inputTuples, [][]*types.T{inputTypes, inputTypes}, expected,
   111  		unorderedVerifier, func(inputs []colexecbase.Operator) (colexecbase.Operator, error) {
   112  			var input colexecbase.Operator
   113  			input = NewParallelUnorderedSynchronizer(inputs, inputTypes, &wg)
   114  			input = NewSimpleProjectOp(input, len(inputTypes), []uint32{0})
   115  			return NewConstOp(testAllocator, input, types.Int, constVal, 1)
   116  		})
   117  	wg.Wait()
   118  }