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

     1  // Copyright 2020 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 ordering
    12  
    13  import (
    14  	"fmt"
    15  	"testing"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/opt"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/opt/norm"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/opt/props"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/opt/props/physical"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testcat"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/opt/testutils/testexpr"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    25  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    26  )
    27  
    28  func TestDistinctOnProvided(t *testing.T) {
    29  	evalCtx := tree.NewTestingEvalContext(nil /* st */)
    30  	var f norm.Factory
    31  	f.Init(evalCtx, testcat.New())
    32  	md := f.Metadata()
    33  	for i := 1; i <= 5; i++ {
    34  		md.AddColumn(fmt.Sprintf("c%d", i), types.Int)
    35  	}
    36  	c := func(cols ...opt.ColumnID) opt.ColSet {
    37  		return opt.MakeColSet(cols...)
    38  	}
    39  
    40  	fd1eq5 := props.FuncDepSet{}
    41  	fd1eq5.AddEquivalency(1, 5)
    42  
    43  	// DistinctOn might not project all input columns, so we have three sets of
    44  	// columns, corresponding to this SQL:
    45  	//   SELECT <outCols> FROM
    46  	//     SELECT DISTINCT ON(<groupingCols>) <inputCols>
    47  	testCases := []struct {
    48  		inCols       opt.ColSet
    49  		inFDs        props.FuncDepSet
    50  		outCols      opt.ColSet
    51  		groupingCols opt.ColSet
    52  		required     string
    53  		internal     string
    54  		input        string
    55  		expected     string
    56  	}{
    57  		{ // case 1: Internal ordering is stronger; the provided ordering needs
    58  			//         trimming.
    59  			inCols:       c(1, 2, 3, 4, 5),
    60  			inFDs:        props.FuncDepSet{},
    61  			outCols:      c(1, 2, 3, 4, 5),
    62  			groupingCols: c(1, 2),
    63  			required:     "+1",
    64  			internal:     "+1,+5",
    65  			input:        "+1,+5",
    66  			expected:     "+1",
    67  		},
    68  		{ // case 2: Projecting all input columns; ok to pass through provided.
    69  			inCols:       c(1, 2, 3, 4, 5),
    70  			inFDs:        fd1eq5,
    71  			outCols:      c(1, 2, 3, 4, 5),
    72  			groupingCols: c(1, 2),
    73  			required:     "+(1|5)",
    74  			internal:     "",
    75  			input:        "+5",
    76  			expected:     "+5",
    77  		},
    78  		{ // case 3: Not projecting all input columns; the provided ordering
    79  			//         needs remapping.
    80  			inCols:       c(1, 2, 3, 4, 5),
    81  			inFDs:        fd1eq5,
    82  			outCols:      c(1, 2, 3),
    83  			groupingCols: c(1, 2),
    84  			required:     "+(1|5)",
    85  			internal:     "",
    86  			input:        "+5",
    87  			expected:     "+1",
    88  		},
    89  		{ // case 4: The provided ordering needs both trimming and remapping.
    90  			inCols:       c(1, 2, 3, 4, 5),
    91  			inFDs:        fd1eq5,
    92  			outCols:      c(1, 2, 3),
    93  			groupingCols: c(1, 2),
    94  			required:     "+(1|5)",
    95  			internal:     "+5,+4",
    96  			input:        "+5,+4",
    97  			expected:     "+1",
    98  		},
    99  	}
   100  	for tcIdx, tc := range testCases {
   101  		t.Run(fmt.Sprintf("case%d", tcIdx+1), func(t *testing.T) {
   102  			input := &testexpr.Instance{
   103  				Rel: &props.Relational{
   104  					OutputCols: tc.outCols,
   105  					FuncDeps:   fd1eq5,
   106  				},
   107  				Provided: &physical.Provided{
   108  					Ordering: physical.ParseOrdering(tc.input),
   109  				},
   110  			}
   111  			p := memo.GroupingPrivate{
   112  				GroupingCols: tc.groupingCols,
   113  				Ordering:     physical.ParseOrderingChoice(tc.internal),
   114  			}
   115  			var aggs memo.AggregationsExpr
   116  			tc.outCols.Difference(tc.groupingCols).ForEach(func(col opt.ColumnID) {
   117  				aggs = append(aggs, f.ConstructAggregationsItem(
   118  					f.ConstructFirstAgg(f.ConstructVariable(col)),
   119  					col,
   120  				))
   121  			})
   122  			distinctOn := f.Memo().MemoizeDistinctOn(input, aggs, &p)
   123  			req := physical.ParseOrderingChoice(tc.required)
   124  			res := distinctOnBuildProvided(distinctOn, &req).String()
   125  			if res != tc.expected {
   126  				t.Errorf("expected '%s', got '%s'", tc.expected, res)
   127  			}
   128  		})
   129  	}
   130  }