github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/deselector_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  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    19  	"github.com/cockroachdb/cockroach/pkg/col/coldatatestutils"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/colexecbase"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    22  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    23  	"github.com/cockroachdb/cockroach/pkg/util/randutil"
    24  )
    25  
    26  func TestDeselector(t *testing.T) {
    27  	defer leaktest.AfterTest(t)()
    28  	tcs := []struct {
    29  		typs     []*types.T
    30  		tuples   []tuple
    31  		sel      []int
    32  		expected []tuple
    33  	}{
    34  		{
    35  			typs:     []*types.T{types.Int},
    36  			tuples:   tuples{{0}, {1}, {2}},
    37  			sel:      nil,
    38  			expected: tuples{{0}, {1}, {2}},
    39  		},
    40  		{
    41  			typs:     []*types.T{types.Int},
    42  			tuples:   tuples{{0}, {1}, {2}},
    43  			sel:      []int{},
    44  			expected: tuples{},
    45  		},
    46  		{
    47  			typs:     []*types.T{types.Int},
    48  			tuples:   tuples{{0}, {1}, {2}},
    49  			sel:      []int{1},
    50  			expected: tuples{{1}},
    51  		},
    52  		{
    53  			typs:     []*types.T{types.Int},
    54  			tuples:   tuples{{0}, {1}, {2}},
    55  			sel:      []int{0, 2},
    56  			expected: tuples{{0}, {2}},
    57  		},
    58  		{
    59  			typs:     []*types.T{types.Int},
    60  			tuples:   tuples{{0}, {1}, {2}},
    61  			sel:      []int{0, 1, 2},
    62  			expected: tuples{{0}, {1}, {2}},
    63  		},
    64  	}
    65  
    66  	for _, tc := range tcs {
    67  		runTestsWithFixedSel(t, []tuples{tc.tuples}, tc.typs, tc.sel, func(t *testing.T, input []colexecbase.Operator) {
    68  			op := NewDeselectorOp(testAllocator, input[0], tc.typs)
    69  			out := newOpTestOutput(op, tc.expected)
    70  
    71  			if err := out.Verify(); err != nil {
    72  				t.Fatal(err)
    73  			}
    74  		})
    75  	}
    76  }
    77  
    78  func BenchmarkDeselector(b *testing.B) {
    79  	rng, _ := randutil.NewPseudoRand()
    80  	ctx := context.Background()
    81  
    82  	nCols := 1
    83  	inputTypes := make([]*types.T, nCols)
    84  
    85  	for colIdx := 0; colIdx < nCols; colIdx++ {
    86  		inputTypes[colIdx] = types.Int
    87  	}
    88  
    89  	batch := testAllocator.NewMemBatch(inputTypes)
    90  
    91  	for colIdx := 0; colIdx < nCols; colIdx++ {
    92  		col := batch.ColVec(colIdx).Int64()
    93  		for i := 0; i < coldata.BatchSize(); i++ {
    94  			col[i] = int64(i)
    95  		}
    96  	}
    97  	for _, probOfOmitting := range []float64{0.1, 0.9} {
    98  		sel := coldatatestutils.RandomSel(rng, coldata.BatchSize(), probOfOmitting)
    99  		batchLen := len(sel)
   100  
   101  		for _, nBatches := range []int{1 << 1, 1 << 2, 1 << 4, 1 << 8} {
   102  			b.Run(fmt.Sprintf("rows=%d/after selection=%d", nBatches*coldata.BatchSize(), nBatches*batchLen), func(b *testing.B) {
   103  				// We're measuring the amount of data that is not selected out.
   104  				b.SetBytes(int64(8 * nBatches * batchLen * nCols))
   105  				batch.SetSelection(true)
   106  				copy(batch.Selection(), sel)
   107  				batch.SetLength(batchLen)
   108  				input := colexecbase.NewRepeatableBatchSource(testAllocator, batch, inputTypes)
   109  				op := NewDeselectorOp(testAllocator, input, inputTypes)
   110  				op.Init()
   111  				b.ResetTimer()
   112  				for i := 0; i < b.N; i++ {
   113  					input.ResetBatchesToReturn(nBatches)
   114  					for b := op.Next(ctx); b.Length() != 0; b = op.Next(ctx) {
   115  					}
   116  					// We don't need to reset the deselector because it doesn't keep any
   117  					// state. We do, however, want to keep its already allocated memory
   118  					// so that this memory allocation doesn't impact the benchmark.
   119  				}
   120  				b.StopTimer()
   121  			})
   122  		}
   123  	}
   124  }