github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/col/coldata/batch_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 coldata_test
    12  
    13  import (
    14  	"reflect"
    15  	"testing"
    16  	"unsafe"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    20  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    21  	"github.com/stretchr/testify/assert"
    22  )
    23  
    24  func TestBatchReset(t *testing.T) {
    25  	defer leaktest.AfterTest(t)()
    26  
    27  	resetAndCheck := func(b coldata.Batch, typs []*types.T, n int, shouldReuse bool) {
    28  		t.Helper()
    29  		// Use the data backing the ColVecs slice as a proxy for when things get
    30  		// reallocated.
    31  		vecsBefore := b.ColVecs()
    32  		ptrBefore := (*reflect.SliceHeader)(unsafe.Pointer(&vecsBefore))
    33  		b.Reset(typs, n, coldata.StandardColumnFactory)
    34  		vecsAfter := b.ColVecs()
    35  		ptrAfter := (*reflect.SliceHeader)(unsafe.Pointer(&vecsAfter))
    36  		assert.Equal(t, shouldReuse, ptrBefore.Data == ptrAfter.Data)
    37  		assert.Equal(t, n, b.Length())
    38  		assert.Equal(t, len(typs), b.Width())
    39  
    40  		assert.Nil(t, b.Selection())
    41  		b.SetSelection(true)
    42  		// Invariant: selection vector length matches batch length
    43  		// Invariant: all cap(column) is >= cap(selection vector)
    44  		assert.Equal(t, n, len(b.Selection()))
    45  		selCap := cap(b.Selection())
    46  
    47  		for i, vec := range b.ColVecs() {
    48  			assert.False(t, vec.MaybeHasNulls())
    49  			assert.False(t, vec.Nulls().NullAt(0))
    50  			assert.True(t, typs[i].Identical(vec.Type()))
    51  			// Sanity check that we can actually use the column. This is mostly for
    52  			// making sure a flat bytes column gets reset.
    53  			vec.Nulls().SetNull(0)
    54  			assert.True(t, vec.Nulls().NullAt(0))
    55  			switch vec.CanonicalTypeFamily() {
    56  			case types.IntFamily:
    57  				x := vec.Int64()
    58  				assert.True(t, len(x) >= n)
    59  				assert.True(t, cap(x) >= selCap)
    60  				x[0] = 1
    61  			case types.BytesFamily:
    62  				x := vec.Bytes()
    63  				assert.True(t, x.Len() >= n)
    64  				x.Set(0, []byte{1})
    65  			default:
    66  				panic(vec.Type())
    67  			}
    68  		}
    69  	}
    70  
    71  	typsInt := []*types.T{types.Int}
    72  	typsBytes := []*types.T{types.Bytes}
    73  	typsIntBytes := []*types.T{types.Int, types.Bytes}
    74  	var b coldata.Batch
    75  
    76  	// Simple case, reuse
    77  	b = coldata.NewMemBatch(typsInt, coldata.StandardColumnFactory)
    78  	resetAndCheck(b, typsInt, 1, true)
    79  
    80  	// Types don't match, don't reuse
    81  	b = coldata.NewMemBatch(typsInt, coldata.StandardColumnFactory)
    82  	resetAndCheck(b, typsBytes, 1, false)
    83  
    84  	// Columns are a prefix, reuse
    85  	b = coldata.NewMemBatch(typsIntBytes, coldata.StandardColumnFactory)
    86  	resetAndCheck(b, typsInt, 1, true)
    87  
    88  	// Exact length, reuse
    89  	b = coldata.NewMemBatchWithSize(typsInt, 1, coldata.StandardColumnFactory)
    90  	resetAndCheck(b, typsInt, 1, true)
    91  
    92  	// Insufficient capacity, don't reuse
    93  	b = coldata.NewMemBatchWithSize(typsInt, 1, coldata.StandardColumnFactory)
    94  	resetAndCheck(b, typsInt, 2, false)
    95  
    96  	// Selection vector gets reset
    97  	b = coldata.NewMemBatchWithSize(typsInt, 1, coldata.StandardColumnFactory)
    98  	b.SetSelection(true)
    99  	b.Selection()[0] = 7
   100  	resetAndCheck(b, typsInt, 1, true)
   101  
   102  	// Nulls gets reset
   103  	b = coldata.NewMemBatchWithSize(typsInt, 1, coldata.StandardColumnFactory)
   104  	b.ColVec(0).Nulls().SetNull(0)
   105  	resetAndCheck(b, typsInt, 1, true)
   106  
   107  	// Bytes columns use a different impl than everything else
   108  	b = coldata.NewMemBatch(typsBytes, coldata.StandardColumnFactory)
   109  	resetAndCheck(b, typsBytes, 1, true)
   110  }