vitess.io/vitess@v0.16.2/go/vt/vtgate/engine/memory_sort_test.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package engine
    18  
    19  import (
    20  	"context"
    21  	"testing"
    22  
    23  	"vitess.io/vitess/go/vt/servenv"
    24  	"vitess.io/vitess/go/vt/vtgate/evalengine"
    25  
    26  	"vitess.io/vitess/go/mysql/collations"
    27  	"vitess.io/vitess/go/test/utils"
    28  
    29  	"github.com/stretchr/testify/require"
    30  
    31  	"vitess.io/vitess/go/sqltypes"
    32  	querypb "vitess.io/vitess/go/vt/proto/query"
    33  )
    34  
    35  func init() {
    36  	// We require MySQL 8.0 collations for the comparisons in the tests
    37  	mySQLVersion := "8.0.0"
    38  	servenv.SetMySQLServerVersionForTest(mySQLVersion)
    39  	collationEnv = collations.NewEnvironment(mySQLVersion)
    40  }
    41  
    42  func TestMemorySortExecute(t *testing.T) {
    43  	fields := sqltypes.MakeTestFields(
    44  		"c1|c2",
    45  		"varbinary|decimal",
    46  	)
    47  	fp := &fakePrimitive{
    48  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
    49  			fields,
    50  			"a|1",
    51  			"g|2",
    52  			"a|1",
    53  			"c|4",
    54  			"c|3",
    55  		)},
    56  	}
    57  
    58  	ms := &MemorySort{
    59  		OrderBy: []OrderByParams{{
    60  			WeightStringCol: -1,
    61  			Col:             1,
    62  		}},
    63  		Input: fp,
    64  	}
    65  
    66  	result, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
    67  	require.NoError(t, err)
    68  
    69  	wantResult := sqltypes.MakeTestResult(
    70  		fields,
    71  		"a|1",
    72  		"a|1",
    73  		"g|2",
    74  		"c|3",
    75  		"c|4",
    76  	)
    77  	utils.MustMatch(t, wantResult, result)
    78  
    79  	fp.rewind()
    80  	ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
    81  	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
    82  
    83  	result, err = ms.TryExecute(context.Background(), &noopVCursor{}, bv, false)
    84  	require.NoError(t, err)
    85  
    86  	wantResult = sqltypes.MakeTestResult(
    87  		fields,
    88  		"a|1",
    89  		"a|1",
    90  		"g|2",
    91  	)
    92  	utils.MustMatch(t, wantResult, result)
    93  }
    94  
    95  func TestMemorySortStreamExecuteWeightString(t *testing.T) {
    96  	fields := sqltypes.MakeTestFields(
    97  		"weightString|normal",
    98  		"varbinary|varchar",
    99  	)
   100  	fp := &fakePrimitive{
   101  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   102  			fields,
   103  			"null|x",
   104  			"g|d",
   105  			"a|a",
   106  			"c|t",
   107  			"f|p",
   108  		)},
   109  	}
   110  
   111  	ms := &MemorySort{
   112  		OrderBy: []OrderByParams{{
   113  			WeightStringCol: 0,
   114  			Col:             1,
   115  		}},
   116  		Input: fp,
   117  	}
   118  	var results []*sqltypes.Result
   119  
   120  	t.Run("order by weight string", func(t *testing.T) {
   121  
   122  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
   123  			results = append(results, qr)
   124  			return nil
   125  		})
   126  		require.NoError(t, err)
   127  
   128  		wantResults := sqltypes.MakeTestStreamingResults(
   129  			fields,
   130  			"null|x",
   131  			"a|a",
   132  			"c|t",
   133  			"f|p",
   134  			"g|d",
   135  		)
   136  		utils.MustMatch(t, wantResults, results)
   137  	})
   138  
   139  	t.Run("Limit test", func(t *testing.T) {
   140  		fp.rewind()
   141  		ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   142  		bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   143  
   144  		results = nil
   145  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, bv, true, func(qr *sqltypes.Result) error {
   146  			results = append(results, qr)
   147  			return nil
   148  		})
   149  		require.NoError(t, err)
   150  
   151  		wantResults := sqltypes.MakeTestStreamingResults(
   152  			fields,
   153  			"null|x",
   154  			"a|a",
   155  			"c|t",
   156  		)
   157  		utils.MustMatch(t, wantResults, results)
   158  	})
   159  }
   160  
   161  func TestMemorySortExecuteWeightString(t *testing.T) {
   162  	fields := sqltypes.MakeTestFields(
   163  		"c1|c2",
   164  		"varchar|varbinary",
   165  	)
   166  	fp := &fakePrimitive{
   167  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   168  			fields,
   169  			"a|1",
   170  			"g|2",
   171  			"a|1",
   172  			"c|4",
   173  			"c|3",
   174  		)},
   175  	}
   176  
   177  	ms := &MemorySort{
   178  		OrderBy: []OrderByParams{{
   179  			WeightStringCol: 1,
   180  			Col:             0,
   181  		}},
   182  		Input: fp,
   183  	}
   184  
   185  	result, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
   186  	require.NoError(t, err)
   187  
   188  	wantResult := sqltypes.MakeTestResult(
   189  		fields,
   190  		"a|1",
   191  		"a|1",
   192  		"g|2",
   193  		"c|3",
   194  		"c|4",
   195  	)
   196  	utils.MustMatch(t, wantResult, result)
   197  
   198  	fp.rewind()
   199  	ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   200  	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   201  
   202  	result, err = ms.TryExecute(context.Background(), &noopVCursor{}, bv, false)
   203  	require.NoError(t, err)
   204  
   205  	wantResult = sqltypes.MakeTestResult(
   206  		fields,
   207  		"a|1",
   208  		"a|1",
   209  		"g|2",
   210  	)
   211  	utils.MustMatch(t, wantResult, result)
   212  }
   213  
   214  func TestMemorySortStreamExecuteCollation(t *testing.T) {
   215  	fields := sqltypes.MakeTestFields(
   216  		"normal",
   217  		"varchar",
   218  	)
   219  	fp := &fakePrimitive{
   220  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   221  			fields,
   222  			"c",
   223  			"d",
   224  			"cs",
   225  			"cs",
   226  			"c",
   227  		)},
   228  	}
   229  
   230  	collationID, _ := collations.Local().LookupID("utf8mb4_hu_0900_ai_ci")
   231  	ms := &MemorySort{
   232  		OrderBy: []OrderByParams{{
   233  			Col:         0,
   234  			CollationID: collationID,
   235  		}},
   236  		Input: fp,
   237  	}
   238  
   239  	var results []*sqltypes.Result
   240  	t.Run("order by collation", func(t *testing.T) {
   241  		results = nil
   242  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
   243  			results = append(results, qr)
   244  			return nil
   245  		})
   246  		require.NoError(t, err)
   247  
   248  		wantResults := sqltypes.MakeTestStreamingResults(
   249  			fields,
   250  			"c",
   251  			"c",
   252  			"cs",
   253  			"cs",
   254  			"d",
   255  		)
   256  		utils.MustMatch(t, wantResults, results)
   257  	})
   258  
   259  	t.Run("Descending order by collation", func(t *testing.T) {
   260  		ms.OrderBy[0].Desc = true
   261  		fp.rewind()
   262  		results = nil
   263  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
   264  			results = append(results, qr)
   265  			return nil
   266  		})
   267  		require.NoError(t, err)
   268  
   269  		wantResults := sqltypes.MakeTestStreamingResults(
   270  			fields,
   271  			"d",
   272  			"cs",
   273  			"cs",
   274  			"c",
   275  			"c",
   276  		)
   277  		utils.MustMatch(t, wantResults, results)
   278  	})
   279  
   280  	t.Run("Limit test", func(t *testing.T) {
   281  		fp.rewind()
   282  		ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   283  		bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   284  
   285  		results = nil
   286  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, bv, true, func(qr *sqltypes.Result) error {
   287  			results = append(results, qr)
   288  			return nil
   289  		})
   290  		require.NoError(t, err)
   291  
   292  		wantResults := sqltypes.MakeTestStreamingResults(
   293  			fields,
   294  			"d",
   295  			"cs",
   296  			"cs",
   297  		)
   298  		utils.MustMatch(t, wantResults, results)
   299  	})
   300  }
   301  
   302  func TestMemorySortExecuteCollation(t *testing.T) {
   303  	fields := sqltypes.MakeTestFields(
   304  		"c1",
   305  		"varchar",
   306  	)
   307  	fp := &fakePrimitive{
   308  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   309  			fields,
   310  			"c",
   311  			"d",
   312  			"cs",
   313  			"cs",
   314  			"c",
   315  		)},
   316  	}
   317  
   318  	collationID, _ := collations.Local().LookupID("utf8mb4_hu_0900_ai_ci")
   319  	ms := &MemorySort{
   320  		OrderBy: []OrderByParams{{
   321  			Col:         0,
   322  			CollationID: collationID,
   323  		}},
   324  		Input: fp,
   325  	}
   326  
   327  	result, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
   328  	require.NoError(t, err)
   329  
   330  	wantResult := sqltypes.MakeTestResult(
   331  		fields,
   332  		"c",
   333  		"c",
   334  		"cs",
   335  		"cs",
   336  		"d",
   337  	)
   338  	utils.MustMatch(t, wantResult, result)
   339  
   340  	fp.rewind()
   341  	ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   342  	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   343  
   344  	result, err = ms.TryExecute(context.Background(), &noopVCursor{}, bv, false)
   345  	require.NoError(t, err)
   346  
   347  	wantResult = sqltypes.MakeTestResult(
   348  		fields,
   349  		"c",
   350  		"c",
   351  		"cs",
   352  	)
   353  	utils.MustMatch(t, wantResult, result)
   354  }
   355  
   356  func TestMemorySortStreamExecute(t *testing.T) {
   357  	fields := sqltypes.MakeTestFields(
   358  		"c1|c2",
   359  		"varbinary|decimal",
   360  	)
   361  	fp := &fakePrimitive{
   362  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   363  			fields,
   364  			"a|1",
   365  			"g|2",
   366  			"a|1",
   367  			"c|4",
   368  			"c|3",
   369  		)},
   370  	}
   371  
   372  	ms := &MemorySort{
   373  		OrderBy: []OrderByParams{{
   374  			WeightStringCol: -1,
   375  			Col:             1,
   376  		}},
   377  		Input: fp,
   378  	}
   379  
   380  	var results []*sqltypes.Result
   381  	err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
   382  		results = append(results, qr)
   383  		return nil
   384  	})
   385  	require.NoError(t, err)
   386  
   387  	wantResults := sqltypes.MakeTestStreamingResults(
   388  		fields,
   389  		"a|1",
   390  		"a|1",
   391  		"g|2",
   392  		"c|3",
   393  		"c|4",
   394  	)
   395  	utils.MustMatch(t, wantResults, results)
   396  
   397  	fp.rewind()
   398  	ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   399  	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   400  
   401  	results = nil
   402  	err = ms.TryStreamExecute(context.Background(), &noopVCursor{}, bv, true, func(qr *sqltypes.Result) error {
   403  		results = append(results, qr)
   404  		return nil
   405  	})
   406  	require.NoError(t, err)
   407  
   408  	wantResults = sqltypes.MakeTestStreamingResults(
   409  		fields,
   410  		"a|1",
   411  		"a|1",
   412  		"g|2",
   413  	)
   414  	utils.MustMatch(t, wantResults, results)
   415  }
   416  
   417  func TestMemorySortGetFields(t *testing.T) {
   418  	result := sqltypes.MakeTestResult(
   419  		sqltypes.MakeTestFields(
   420  			"col1|col2",
   421  			"int64|varchar",
   422  		),
   423  	)
   424  	fp := &fakePrimitive{results: []*sqltypes.Result{result}}
   425  
   426  	ms := &MemorySort{Input: fp}
   427  
   428  	got, err := ms.GetFields(context.Background(), nil, nil)
   429  	require.NoError(t, err)
   430  	utils.MustMatch(t, result, got)
   431  }
   432  
   433  func TestMemorySortExecuteTruncate(t *testing.T) {
   434  	fields := sqltypes.MakeTestFields(
   435  		"c1|c2|c3",
   436  		"varbinary|decimal|int64",
   437  	)
   438  	fp := &fakePrimitive{
   439  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   440  			fields,
   441  			"a|1|1",
   442  			"g|2|1",
   443  			"a|1|1",
   444  			"c|4|1",
   445  			"c|3|1",
   446  		)},
   447  	}
   448  
   449  	ms := &MemorySort{
   450  		OrderBy: []OrderByParams{{
   451  			WeightStringCol: -1,
   452  			Col:             1,
   453  		}},
   454  		Input:               fp,
   455  		TruncateColumnCount: 2,
   456  	}
   457  
   458  	result, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
   459  	require.NoError(t, err)
   460  
   461  	wantResult := sqltypes.MakeTestResult(
   462  		fields[:2],
   463  		"a|1",
   464  		"a|1",
   465  		"g|2",
   466  		"c|3",
   467  		"c|4",
   468  	)
   469  	utils.MustMatch(t, wantResult, result)
   470  }
   471  
   472  func TestMemorySortStreamExecuteTruncate(t *testing.T) {
   473  	fields := sqltypes.MakeTestFields(
   474  		"c1|c2|c3",
   475  		"varbinary|decimal|int64",
   476  	)
   477  	fp := &fakePrimitive{
   478  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   479  			fields,
   480  			"a|1|1",
   481  			"g|2|1",
   482  			"a|1|1",
   483  			"c|4|1",
   484  			"c|3|1",
   485  		)},
   486  	}
   487  
   488  	ms := &MemorySort{
   489  		OrderBy: []OrderByParams{{
   490  			WeightStringCol: -1,
   491  			Col:             1,
   492  		}},
   493  		Input:               fp,
   494  		TruncateColumnCount: 2,
   495  	}
   496  
   497  	var results []*sqltypes.Result
   498  	err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
   499  		results = append(results, qr)
   500  		return nil
   501  	})
   502  	require.NoError(t, err)
   503  
   504  	wantResults := sqltypes.MakeTestStreamingResults(
   505  		fields[:2],
   506  		"a|1",
   507  		"a|1",
   508  		"g|2",
   509  		"c|3",
   510  		"c|4",
   511  	)
   512  	utils.MustMatch(t, wantResults, results)
   513  }
   514  
   515  func TestMemorySortMultiColumn(t *testing.T) {
   516  	fields := sqltypes.MakeTestFields(
   517  		"c1|c2",
   518  		"varbinary|decimal",
   519  	)
   520  	fp := &fakePrimitive{
   521  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   522  			fields,
   523  			"a|1",
   524  			"b|2",
   525  			"b|1",
   526  			"c|4",
   527  			"c|3",
   528  		)},
   529  	}
   530  
   531  	ms := &MemorySort{
   532  		OrderBy: []OrderByParams{{
   533  			Col:             1,
   534  			WeightStringCol: -1,
   535  		}, {
   536  			Col:             0,
   537  			WeightStringCol: -1,
   538  			Desc:            true,
   539  		}},
   540  		Input: fp,
   541  	}
   542  
   543  	result, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
   544  	require.NoError(t, err)
   545  
   546  	wantResult := sqltypes.MakeTestResult(
   547  		fields,
   548  		"b|1",
   549  		"a|1",
   550  		"b|2",
   551  		"c|3",
   552  		"c|4",
   553  	)
   554  	utils.MustMatch(t, wantResult, result)
   555  
   556  	fp.rewind()
   557  	ms.UpperLimit = evalengine.NewBindVar("__upper_limit", collations.TypedCollation{})
   558  	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
   559  
   560  	result, err = ms.TryExecute(context.Background(), &noopVCursor{}, bv, false)
   561  	require.NoError(t, err)
   562  
   563  	wantResult = sqltypes.MakeTestResult(
   564  		fields,
   565  		"b|1",
   566  		"a|1",
   567  		"b|2",
   568  	)
   569  	utils.MustMatch(t, wantResult, result)
   570  }
   571  
   572  func TestMemorySortMaxMemoryRows(t *testing.T) {
   573  	saveMax := testMaxMemoryRows
   574  	saveIgnore := testIgnoreMaxMemoryRows
   575  	testMaxMemoryRows = 3
   576  	defer func() {
   577  		testMaxMemoryRows = saveMax
   578  		testIgnoreMaxMemoryRows = saveIgnore
   579  	}()
   580  
   581  	testCases := []struct {
   582  		ignoreMaxMemoryRows bool
   583  		err                 string
   584  	}{
   585  		{true, ""},
   586  		{false, "in-memory row count exceeded allowed limit of 3"},
   587  	}
   588  	fields := sqltypes.MakeTestFields(
   589  		"c1|c2",
   590  		"varbinary|decimal",
   591  	)
   592  	for _, test := range testCases {
   593  		fp := &fakePrimitive{
   594  			results: []*sqltypes.Result{sqltypes.MakeTestResult(
   595  				fields,
   596  				"a|1",
   597  				"b|2",
   598  				"a|1",
   599  				"c|4",
   600  				"c|3",
   601  			)},
   602  		}
   603  
   604  		ms := &MemorySort{
   605  			OrderBy: []OrderByParams{{
   606  				WeightStringCol: -1,
   607  				Col:             1,
   608  			}},
   609  			Input: fp,
   610  		}
   611  
   612  		testIgnoreMaxMemoryRows = test.ignoreMaxMemoryRows
   613  		err := ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, false, func(qr *sqltypes.Result) error {
   614  			return nil
   615  		})
   616  		if testIgnoreMaxMemoryRows {
   617  			require.NoError(t, err)
   618  		} else {
   619  			require.EqualError(t, err, test.err)
   620  		}
   621  	}
   622  }
   623  
   624  func TestMemorySortExecuteNoVarChar(t *testing.T) {
   625  	fields := sqltypes.MakeTestFields(
   626  		"c1|c2",
   627  		"varchar|decimal",
   628  	)
   629  	fp := &fakePrimitive{
   630  		results: []*sqltypes.Result{sqltypes.MakeTestResult(
   631  			fields,
   632  			"a|1",
   633  			"b|2",
   634  			"a|1",
   635  			"c|4",
   636  			"c|3",
   637  		)},
   638  	}
   639  
   640  	ms := &MemorySort{
   641  		OrderBy: []OrderByParams{{
   642  			WeightStringCol: -1,
   643  			Col:             0,
   644  		}},
   645  		Input: fp,
   646  	}
   647  
   648  	_, err := ms.TryExecute(context.Background(), &noopVCursor{}, nil, false)
   649  	want := "cannot compare strings, collation is unknown or unsupported (collation ID: 0)"
   650  	if err == nil || err.Error() != want {
   651  		t.Errorf("Execute err: %v, want %v", err, want)
   652  	}
   653  
   654  	fp.rewind()
   655  	err = ms.TryStreamExecute(context.Background(), &noopVCursor{}, nil, false, func(qr *sqltypes.Result) error {
   656  		return nil
   657  	})
   658  	if err == nil || err.Error() != want {
   659  		t.Errorf("StreamExecute err: %v, want %v", err, want)
   660  	}
   661  }