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

     1  // Copyright 2019 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  	"context"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    18  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    24  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    25  	"github.com/stretchr/testify/require"
    26  )
    27  
    28  func TestOrdinality(t *testing.T) {
    29  	defer leaktest.AfterTest(t)()
    30  	ctx := context.Background()
    31  	st := cluster.MakeTestingClusterSettings()
    32  	evalCtx := tree.MakeTestingEvalContext(st)
    33  	defer evalCtx.Stop(ctx)
    34  	flowCtx := &execinfra.FlowCtx{
    35  		EvalCtx: &evalCtx,
    36  		Cfg: &execinfra.ServerConfig{
    37  			Settings: st,
    38  		},
    39  	}
    40  	tcs := []struct {
    41  		tuples     []tuple
    42  		expected   []tuple
    43  		inputTypes []*types.T
    44  	}{
    45  		{
    46  			tuples:     tuples{{1}},
    47  			expected:   tuples{{1, 1}},
    48  			inputTypes: []*types.T{types.Int},
    49  		},
    50  		{
    51  			tuples:     tuples{{}, {}, {}, {}, {}},
    52  			expected:   tuples{{1}, {2}, {3}, {4}, {5}},
    53  			inputTypes: []*types.T{},
    54  		},
    55  		{
    56  			tuples:     tuples{{5}, {6}, {7}, {8}},
    57  			expected:   tuples{{5, 1}, {6, 2}, {7, 3}, {8, 4}},
    58  			inputTypes: []*types.T{types.Int},
    59  		},
    60  		{
    61  			tuples:     tuples{{5, 'a'}, {6, 'b'}, {7, 'c'}, {8, 'd'}},
    62  			expected:   tuples{{5, 'a', 1}, {6, 'b', 2}, {7, 'c', 3}, {8, 'd', 4}},
    63  			inputTypes: []*types.T{types.Int, types.String},
    64  		},
    65  	}
    66  
    67  	for _, tc := range tcs {
    68  		runTests(t, []tuples{tc.tuples}, tc.expected, orderedVerifier,
    69  			func(input []colexecbase.Operator) (colexecbase.Operator, error) {
    70  				return createTestOrdinalityOperator(ctx, flowCtx, input[0], tc.inputTypes)
    71  			})
    72  	}
    73  }
    74  
    75  func BenchmarkOrdinality(b *testing.B) {
    76  	ctx := context.Background()
    77  	st := cluster.MakeTestingClusterSettings()
    78  	evalCtx := tree.MakeTestingEvalContext(st)
    79  	defer evalCtx.Stop(ctx)
    80  	flowCtx := &execinfra.FlowCtx{
    81  		EvalCtx: &evalCtx,
    82  		Cfg: &execinfra.ServerConfig{
    83  			Settings: st,
    84  		},
    85  	}
    86  
    87  	typs := []*types.T{types.Int, types.Int, types.Int}
    88  	batch := testAllocator.NewMemBatch(typs)
    89  	batch.SetLength(coldata.BatchSize())
    90  	source := colexecbase.NewRepeatableBatchSource(testAllocator, batch, typs)
    91  	ordinality, err := createTestOrdinalityOperator(ctx, flowCtx, source, []*types.T{types.Int, types.Int, types.Int})
    92  	require.NoError(b, err)
    93  	ordinality.Init()
    94  
    95  	b.SetBytes(int64(8 * coldata.BatchSize()))
    96  	for i := 0; i < b.N; i++ {
    97  		ordinality.Next(ctx)
    98  	}
    99  }
   100  
   101  func createTestOrdinalityOperator(
   102  	ctx context.Context,
   103  	flowCtx *execinfra.FlowCtx,
   104  	input colexecbase.Operator,
   105  	inputTypes []*types.T,
   106  ) (colexecbase.Operator, error) {
   107  	spec := &execinfrapb.ProcessorSpec{
   108  		Input: []execinfrapb.InputSyncSpec{{ColumnTypes: inputTypes}},
   109  		Core: execinfrapb.ProcessorCoreUnion{
   110  			Ordinality: &execinfrapb.OrdinalitySpec{},
   111  		},
   112  	}
   113  	args := NewColOperatorArgs{
   114  		Spec:                spec,
   115  		Inputs:              []colexecbase.Operator{input},
   116  		StreamingMemAccount: testMemAcc,
   117  	}
   118  	args.TestingKnobs.UseStreamingMemAccountForBuffering = true
   119  	result, err := NewColOperator(ctx, flowCtx, args)
   120  	if err != nil {
   121  		return nil, err
   122  	}
   123  	return result.Op, nil
   124  }