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

     1  // Copyright 2016 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  // Note that this file is not in pkg/sql/colexec because it instantiates a
    12  // server, and if it were moved into sql/colexec, that would create a cycle
    13  // with pkg/server.
    14  
    15  package colflow_test
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"testing"
    21  
    22  	"github.com/cockroachdb/cockroach/pkg/base"
    23  	"github.com/cockroachdb/cockroach/pkg/keys"
    24  	"github.com/cockroachdb/cockroach/pkg/kv"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/colexec"
    26  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    27  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    28  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    29  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    30  	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
    31  	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
    32  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    33  	"github.com/cockroachdb/cockroach/pkg/util/log"
    34  )
    35  
    36  func BenchmarkColBatchScan(b *testing.B) {
    37  	defer leaktest.AfterTest(b)()
    38  	logScope := log.Scope(b)
    39  	defer logScope.Close(b)
    40  	ctx := context.Background()
    41  
    42  	s, sqlDB, kvDB := serverutils.StartServer(b, base.TestServerArgs{})
    43  	defer s.Stopper().Stop(ctx)
    44  
    45  	const numCols = 2
    46  	for _, numRows := range []int{1 << 4, 1 << 8, 1 << 12, 1 << 16} {
    47  		tableName := fmt.Sprintf("t%d", numRows)
    48  		sqlutils.CreateTable(
    49  			b, sqlDB, tableName,
    50  			"k INT PRIMARY KEY, v INT",
    51  			numRows,
    52  			sqlutils.ToRowFn(sqlutils.RowIdxFn, sqlutils.RowModuloFn(42)),
    53  		)
    54  		tableDesc := sqlbase.GetTableDescriptor(kvDB, keys.SystemSQLCodec, "test", tableName)
    55  		b.Run(fmt.Sprintf("rows=%d", numRows), func(b *testing.B) {
    56  			spec := execinfrapb.ProcessorSpec{
    57  				Core: execinfrapb.ProcessorCoreUnion{
    58  					TableReader: &execinfrapb.TableReaderSpec{
    59  						Table: *tableDesc,
    60  						Spans: []execinfrapb.TableReaderSpan{
    61  							{Span: tableDesc.PrimaryIndexSpan(keys.SystemSQLCodec)},
    62  						},
    63  					}},
    64  				Post: execinfrapb.PostProcessSpec{
    65  					Projection:    true,
    66  					OutputColumns: []uint32{0, 1},
    67  				},
    68  			}
    69  
    70  			evalCtx := tree.MakeTestingEvalContext(s.ClusterSettings())
    71  			defer evalCtx.Stop(ctx)
    72  
    73  			flowCtx := execinfra.FlowCtx{
    74  				EvalCtx: &evalCtx,
    75  				Cfg:     &execinfra.ServerConfig{Settings: s.ClusterSettings()},
    76  				Txn:     kv.NewTxn(ctx, s.DB(), s.NodeID()),
    77  				NodeID:  evalCtx.NodeID,
    78  			}
    79  
    80  			b.SetBytes(int64(numRows * numCols * 8))
    81  			b.ResetTimer()
    82  			for i := 0; i < b.N; i++ {
    83  				b.StopTimer()
    84  				args := colexec.NewColOperatorArgs{
    85  					Spec:                &spec,
    86  					StreamingMemAccount: testMemAcc,
    87  				}
    88  				args.TestingKnobs.UseStreamingMemAccountForBuffering = true
    89  				res, err := colexec.NewColOperator(ctx, &flowCtx, args)
    90  				if err != nil {
    91  					b.Fatal(err)
    92  				}
    93  				tr := res.Op
    94  				tr.Init()
    95  				b.StartTimer()
    96  				for {
    97  					bat := tr.Next(ctx)
    98  					if bat.Length() == 0 {
    99  						break
   100  					}
   101  				}
   102  			}
   103  		})
   104  	}
   105  }