github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/rowexec/project_set_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 rowexec
    12  
    13  import (
    14  	"context"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    23  	"github.com/cockroachdb/cockroach/pkg/testutils/distsqlutils"
    24  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    25  )
    26  
    27  func TestProjectSet(t *testing.T) {
    28  	defer leaktest.AfterTest(t)()
    29  
    30  	v := [10]sqlbase.EncDatum{}
    31  	for i := range v {
    32  		v[i] = sqlbase.IntEncDatum(i)
    33  	}
    34  	null := sqlbase.NullEncDatum()
    35  
    36  	testCases := []struct {
    37  		description string
    38  		spec        execinfrapb.ProjectSetSpec
    39  		input       sqlbase.EncDatumRows
    40  		inputTypes  []*types.T
    41  		expected    sqlbase.EncDatumRows
    42  	}{
    43  		{
    44  			description: "scalar function",
    45  			spec: execinfrapb.ProjectSetSpec{
    46  				Exprs: []execinfrapb.Expression{
    47  					{Expr: "@1 + 1"},
    48  				},
    49  				GeneratedColumns: sqlbase.OneIntCol,
    50  				NumColsPerGen:    []uint32{1},
    51  			},
    52  			input: sqlbase.EncDatumRows{
    53  				{v[2]},
    54  			},
    55  			inputTypes: sqlbase.OneIntCol,
    56  			expected: sqlbase.EncDatumRows{
    57  				{v[2], v[3]},
    58  			},
    59  		},
    60  		{
    61  			description: "set-returning function",
    62  			spec: execinfrapb.ProjectSetSpec{
    63  				Exprs: []execinfrapb.Expression{
    64  					{Expr: "generate_series(@1, 2)"},
    65  				},
    66  				GeneratedColumns: sqlbase.OneIntCol,
    67  				NumColsPerGen:    []uint32{1},
    68  			},
    69  			input: sqlbase.EncDatumRows{
    70  				{v[0]},
    71  				{v[1]},
    72  			},
    73  			inputTypes: sqlbase.OneIntCol,
    74  			expected: sqlbase.EncDatumRows{
    75  				{v[0], v[0]},
    76  				{v[0], v[1]},
    77  				{v[0], v[2]},
    78  				{v[1], v[1]},
    79  				{v[1], v[2]},
    80  			},
    81  		},
    82  		{
    83  			description: "multiple exprs with different lengths",
    84  			spec: execinfrapb.ProjectSetSpec{
    85  				Exprs: []execinfrapb.Expression{
    86  					{Expr: "0"},
    87  					{Expr: "generate_series(0, 0)"},
    88  					{Expr: "generate_series(0, 1)"},
    89  					{Expr: "generate_series(0, 2)"},
    90  				},
    91  				GeneratedColumns: intCols(4),
    92  				NumColsPerGen:    []uint32{1, 1, 1, 1},
    93  			},
    94  			input: sqlbase.EncDatumRows{
    95  				{v[0]},
    96  			},
    97  			inputTypes: sqlbase.OneIntCol,
    98  			expected: sqlbase.EncDatumRows{
    99  				{v[0], v[0], v[0], v[0], v[0]},
   100  				{v[0], null, null, v[1], v[1]},
   101  				{v[0], null, null, null, v[2]},
   102  			},
   103  		},
   104  	}
   105  
   106  	for _, c := range testCases {
   107  		t.Run(c.description, func(t *testing.T) {
   108  			runProcessorTest(
   109  				t,
   110  				execinfrapb.ProcessorCoreUnion{ProjectSet: &c.spec},
   111  				execinfrapb.PostProcessSpec{},
   112  				c.inputTypes,
   113  				c.input,
   114  				append(c.inputTypes, c.spec.GeneratedColumns...), /* outputTypes */
   115  				c.expected,
   116  				nil,
   117  			)
   118  		})
   119  	}
   120  }
   121  
   122  func BenchmarkProjectSet(b *testing.B) {
   123  	defer leaktest.AfterTest(b)()
   124  
   125  	st := cluster.MakeTestingClusterSettings()
   126  	evalCtx := tree.MakeTestingEvalContext(st)
   127  	defer evalCtx.Stop(context.Background())
   128  
   129  	v := [10]sqlbase.EncDatum{}
   130  	for i := range v {
   131  		v[i] = sqlbase.IntEncDatum(i)
   132  	}
   133  
   134  	benchCases := []struct {
   135  		description string
   136  		spec        execinfrapb.ProjectSetSpec
   137  		input       sqlbase.EncDatumRows
   138  		inputTypes  []*types.T
   139  	}{
   140  		{
   141  			description: "generate_series",
   142  			spec: execinfrapb.ProjectSetSpec{
   143  				Exprs: []execinfrapb.Expression{
   144  					{Expr: "generate_series(1, 100000)"},
   145  				},
   146  				GeneratedColumns: sqlbase.OneIntCol,
   147  				NumColsPerGen:    []uint32{1},
   148  			},
   149  			input: sqlbase.EncDatumRows{
   150  				{v[0]},
   151  			},
   152  			inputTypes: sqlbase.OneIntCol,
   153  		},
   154  	}
   155  
   156  	for _, c := range benchCases {
   157  		b.Run(c.description, func(b *testing.B) {
   158  			for i := 0; i < b.N; i++ {
   159  				flowCtx := execinfra.FlowCtx{
   160  					Cfg:     &execinfra.ServerConfig{Settings: st},
   161  					EvalCtx: &evalCtx,
   162  				}
   163  
   164  				in := distsqlutils.NewRowBuffer(c.inputTypes, c.input, distsqlutils.RowBufferArgs{})
   165  				out := &distsqlutils.RowBuffer{}
   166  				p, err := NewProcessor(
   167  					context.Background(), &flowCtx, 0, /* processorID */
   168  					&execinfrapb.ProcessorCoreUnion{ProjectSet: &c.spec}, &execinfrapb.PostProcessSpec{},
   169  					[]execinfra.RowSource{in}, []execinfra.RowReceiver{out}, []execinfra.LocalProcessor{})
   170  				if err != nil {
   171  					b.Fatal(err)
   172  				}
   173  
   174  				p.Run(context.Background())
   175  			}
   176  		})
   177  	}
   178  
   179  }