github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/colexec/columnarizer_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  	"context"
    15  	"testing"
    16  	"unsafe"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    19  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    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/sqlbase"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    25  	"github.com/cockroachdb/cockroach/pkg/testutils"
    26  	"github.com/cockroachdb/cockroach/pkg/testutils/distsqlutils"
    27  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    28  	"github.com/cockroachdb/errors"
    29  	"github.com/stretchr/testify/require"
    30  )
    31  
    32  func TestColumnarizerResetsInternalBatch(t *testing.T) {
    33  	defer leaktest.AfterTest(t)()
    34  	typs := []*types.T{types.Int}
    35  	// There will be at least two batches of rows so that we can see whether the
    36  	// internal batch is reset.
    37  	nRows := coldata.BatchSize() * 2
    38  	nCols := len(typs)
    39  	rows := sqlbase.MakeIntRows(nRows, nCols)
    40  	input := execinfra.NewRepeatableRowSource(typs, rows)
    41  
    42  	ctx := context.Background()
    43  	st := cluster.MakeTestingClusterSettings()
    44  	evalCtx := tree.MakeTestingEvalContext(st)
    45  	defer evalCtx.Stop(ctx)
    46  	flowCtx := &execinfra.FlowCtx{
    47  		Cfg:     &execinfra.ServerConfig{Settings: st},
    48  		EvalCtx: &evalCtx,
    49  	}
    50  
    51  	c, err := NewColumnarizer(ctx, testAllocator, flowCtx, 0, input)
    52  	if err != nil {
    53  		t.Fatal(err)
    54  	}
    55  	c.Init()
    56  	foundRows := 0
    57  	for {
    58  		batch := c.Next(ctx)
    59  		require.Nil(t, batch.Selection(), "Columnarizer didn't reset the internal batch")
    60  		if batch.Length() == 0 {
    61  			break
    62  		}
    63  		foundRows += batch.Length()
    64  		// The "meat" of the test - we're updating the batch that the Columnarizer
    65  		// owns.
    66  		batch.SetSelection(true)
    67  	}
    68  	require.Equal(t, nRows, foundRows)
    69  }
    70  
    71  func TestColumnarizerDrainsAndClosesInput(t *testing.T) {
    72  	defer leaktest.AfterTest(t)()
    73  
    74  	ctx := context.Background()
    75  	st := cluster.MakeTestingClusterSettings()
    76  	evalCtx := tree.MakeTestingEvalContext(st)
    77  	defer evalCtx.Stop(ctx)
    78  
    79  	rb := distsqlutils.NewRowBuffer([]*types.T{types.Int}, nil /* rows */, distsqlutils.RowBufferArgs{})
    80  	flowCtx := &execinfra.FlowCtx{EvalCtx: &evalCtx}
    81  
    82  	const errMsg = "artificial error"
    83  	rb.Push(nil, &execinfrapb.ProducerMetadata{Err: errors.New(errMsg)})
    84  	c, err := NewColumnarizer(ctx, testAllocator, flowCtx, 0 /* processorID */, rb)
    85  	require.NoError(t, err)
    86  
    87  	c.Init()
    88  
    89  	// Calling DrainMeta from the vectorized execution engine should propagate to
    90  	// non-vectorized components as calling ConsumerDone and then draining their
    91  	// metadata.
    92  	metaCh := make(chan []execinfrapb.ProducerMetadata)
    93  	go func() {
    94  		metaCh <- c.DrainMeta(ctx)
    95  	}()
    96  
    97  	// Make Next race with DrainMeta, this should be supported by the
    98  	// Columnarizer. If the metadata is obtained through this Next call, the
    99  	// Columnarizer still returns it in DrainMeta.
   100  	_ = c.Next(ctx)
   101  
   102  	meta := <-metaCh
   103  	require.True(t, len(meta) == 1)
   104  	require.True(t, testutils.IsError(meta[0].Err, errMsg))
   105  	require.True(t, rb.Done)
   106  }
   107  
   108  func BenchmarkColumnarize(b *testing.B) {
   109  	types := []*types.T{types.Int, types.Int}
   110  	nRows := 10000
   111  	nCols := 2
   112  	rows := sqlbase.MakeIntRows(nRows, nCols)
   113  	input := execinfra.NewRepeatableRowSource(types, rows)
   114  
   115  	ctx := context.Background()
   116  	st := cluster.MakeTestingClusterSettings()
   117  	evalCtx := tree.MakeTestingEvalContext(st)
   118  	defer evalCtx.Stop(ctx)
   119  	flowCtx := &execinfra.FlowCtx{
   120  		Cfg:     &execinfra.ServerConfig{Settings: st},
   121  		EvalCtx: &evalCtx,
   122  	}
   123  
   124  	b.SetBytes(int64(nRows * nCols * int(unsafe.Sizeof(int64(0)))))
   125  
   126  	c, err := NewColumnarizer(ctx, testAllocator, flowCtx, 0, input)
   127  	if err != nil {
   128  		b.Fatal(err)
   129  	}
   130  	c.Init()
   131  	for i := 0; i < b.N; i++ {
   132  		foundRows := 0
   133  		for {
   134  			batch := c.Next(ctx)
   135  			if batch.Length() == 0 {
   136  				break
   137  			}
   138  			foundRows += batch.Length()
   139  		}
   140  		if foundRows != nRows {
   141  			b.Fatalf("found %d rows, expected %d", foundRows, nRows)
   142  		}
   143  		input.Reset()
   144  	}
   145  }