github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/offset_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/sql/colexecbase"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  )
    22  
    23  func TestOffset(t *testing.T) {
    24  	defer leaktest.AfterTest(t)()
    25  	tcs := []struct {
    26  		offset   int
    27  		tuples   []tuple
    28  		expected []tuple
    29  	}{
    30  		{
    31  			offset:   0,
    32  			tuples:   tuples{{1}, {2}, {3}, {4}},
    33  			expected: tuples{{1}, {2}, {3}, {4}},
    34  		},
    35  		{
    36  			offset:   1,
    37  			tuples:   tuples{{1}, {2}, {3}, {4}},
    38  			expected: tuples{{2}, {3}, {4}},
    39  		},
    40  		{
    41  			offset:   2,
    42  			tuples:   tuples{{1}, {2}, {3}, {4}},
    43  			expected: tuples{{3}, {4}},
    44  		},
    45  		{
    46  			offset:   4,
    47  			tuples:   tuples{{1}, {2}, {3}, {4}},
    48  			expected: tuples{},
    49  		},
    50  		{
    51  			offset:   100000,
    52  			tuples:   tuples{{1}, {2}, {3}, {4}},
    53  			expected: tuples{},
    54  		},
    55  	}
    56  
    57  	for _, tc := range tcs {
    58  		// The tuples consisting of all nulls still count as separate rows, so if
    59  		// we replace all values with nulls, we should get the same output.
    60  		runTestsWithoutAllNullsInjection(t, []tuples{tc.tuples}, nil /* typs */, tc.expected, unorderedVerifier, func(input []colexecbase.Operator) (colexecbase.Operator, error) {
    61  			return NewOffsetOp(input[0], tc.offset), nil
    62  		})
    63  	}
    64  }
    65  
    66  func BenchmarkOffset(b *testing.B) {
    67  	ctx := context.Background()
    68  	typs := []*types.T{types.Int, types.Int, types.Int}
    69  	batch := testAllocator.NewMemBatch(typs)
    70  	batch.SetLength(coldata.BatchSize())
    71  	source := colexecbase.NewRepeatableBatchSource(testAllocator, batch, typs)
    72  	source.Init()
    73  
    74  	o := NewOffsetOp(source, 1)
    75  	// Set throughput proportional to size of the selection vector.
    76  	b.SetBytes(int64(2 * coldata.BatchSize()))
    77  	for i := 0; i < b.N; i++ {
    78  		o.(*offsetOp).Reset()
    79  		o.Next(ctx)
    80  	}
    81  }