github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/rowexec/ordinality_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 rowexec
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"testing"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    24  	"github.com/cockroachdb/cockroach/pkg/testutils/distsqlutils"
    25  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    26  )
    27  
    28  func TestOrdinality(t *testing.T) {
    29  	defer leaktest.AfterTest(t)()
    30  
    31  	v := [15]sqlbase.EncDatum{}
    32  	for i := range v {
    33  		v[i] = sqlbase.DatumToEncDatum(types.Int, tree.NewDInt(tree.DInt(i)))
    34  	}
    35  
    36  	testCases := []struct {
    37  		spec     execinfrapb.OrdinalitySpec
    38  		input    sqlbase.EncDatumRows
    39  		expected sqlbase.EncDatumRows
    40  	}{
    41  		{
    42  			input: sqlbase.EncDatumRows{
    43  				{v[2]},
    44  				{v[5]},
    45  				{v[2]},
    46  				{v[5]},
    47  				{v[2]},
    48  				{v[3]},
    49  				{v[2]},
    50  			},
    51  			expected: sqlbase.EncDatumRows{
    52  				{v[2], v[1]},
    53  				{v[5], v[2]},
    54  				{v[2], v[3]},
    55  				{v[5], v[4]},
    56  				{v[2], v[5]},
    57  				{v[3], v[6]},
    58  				{v[2], v[7]},
    59  			},
    60  		},
    61  		{
    62  			input: sqlbase.EncDatumRows{
    63  				{},
    64  				{},
    65  				{},
    66  				{},
    67  				{},
    68  				{},
    69  				{},
    70  			},
    71  			expected: sqlbase.EncDatumRows{
    72  				{v[1]},
    73  				{v[2]},
    74  				{v[3]},
    75  				{v[4]},
    76  				{v[5]},
    77  				{v[6]},
    78  				{v[7]},
    79  			},
    80  		},
    81  		{
    82  			input: sqlbase.EncDatumRows{
    83  				{v[2], v[1]},
    84  				{v[5], v[2]},
    85  				{v[2], v[3]},
    86  				{v[5], v[4]},
    87  				{v[2], v[5]},
    88  				{v[3], v[6]},
    89  				{v[2], v[7]},
    90  			},
    91  			expected: sqlbase.EncDatumRows{
    92  				{v[2], v[1], v[1]},
    93  				{v[5], v[2], v[2]},
    94  				{v[2], v[3], v[3]},
    95  				{v[5], v[4], v[4]},
    96  				{v[2], v[5], v[5]},
    97  				{v[3], v[6], v[6]},
    98  				{v[2], v[7], v[7]},
    99  			},
   100  		},
   101  	}
   102  
   103  	for _, c := range testCases {
   104  		t.Run("", func(t *testing.T) {
   105  			os := c.spec
   106  
   107  			in := distsqlutils.NewRowBuffer(sqlbase.TwoIntCols, c.input, distsqlutils.RowBufferArgs{})
   108  			out := &distsqlutils.RowBuffer{}
   109  
   110  			st := cluster.MakeTestingClusterSettings()
   111  			evalCtx := tree.MakeTestingEvalContext(st)
   112  			defer evalCtx.Stop(context.Background())
   113  			flowCtx := execinfra.FlowCtx{
   114  				Cfg:     &execinfra.ServerConfig{Settings: st},
   115  				EvalCtx: &evalCtx,
   116  			}
   117  
   118  			d, err := newOrdinalityProcessor(&flowCtx, 0 /* processorID */, &os, in, &execinfrapb.PostProcessSpec{}, out)
   119  			if err != nil {
   120  				t.Fatal(err)
   121  			}
   122  
   123  			d.Run(context.Background())
   124  			if !out.ProducerClosed() {
   125  				t.Fatalf("output RowReceiver not closed")
   126  			}
   127  			var res sqlbase.EncDatumRows
   128  			for {
   129  				row := out.NextNoMeta(t).Copy()
   130  				if row == nil {
   131  					break
   132  				}
   133  				res = append(res, row)
   134  			}
   135  
   136  			var typs []*types.T
   137  			switch len(res[0]) {
   138  			case 1:
   139  				typs = sqlbase.OneIntCol
   140  			case 2:
   141  				typs = sqlbase.TwoIntCols
   142  			case 3:
   143  				typs = sqlbase.ThreeIntCols
   144  			}
   145  			if result := res.String(typs); result != c.expected.String(typs) {
   146  				t.Errorf("invalid results: %s, expected %s'", result, c.expected.String(sqlbase.TwoIntCols))
   147  			}
   148  		})
   149  	}
   150  }
   151  
   152  func BenchmarkOrdinality(b *testing.B) {
   153  	const numCols = 2
   154  
   155  	ctx := context.Background()
   156  	st := cluster.MakeTestingClusterSettings()
   157  	evalCtx := tree.MakeTestingEvalContext(st)
   158  	defer evalCtx.Stop(ctx)
   159  
   160  	flowCtx := &execinfra.FlowCtx{
   161  		Cfg:     &execinfra.ServerConfig{Settings: st},
   162  		EvalCtx: &evalCtx,
   163  	}
   164  	spec := &execinfrapb.OrdinalitySpec{}
   165  
   166  	post := &execinfrapb.PostProcessSpec{}
   167  	for _, numRows := range []int{1 << 4, 1 << 8, 1 << 12, 1 << 16} {
   168  		input := execinfra.NewRepeatableRowSource(sqlbase.TwoIntCols, sqlbase.MakeIntRows(numRows, numCols))
   169  		b.SetBytes(int64(8 * numRows * numCols))
   170  		b.Run(fmt.Sprintf("rows=%d", numRows), func(b *testing.B) {
   171  			for i := 0; i < b.N; i++ {
   172  				o, err := newOrdinalityProcessor(flowCtx, 0 /* processorID */, spec, input, post, &rowDisposer{})
   173  				if err != nil {
   174  					b.Fatal(err)
   175  				}
   176  				o.Run(ctx)
   177  				input.Reset()
   178  			}
   179  		})
   180  	}
   181  }