Collations Integration Ordering

go/vt/vtgate/engine/comparer.go

@@ -25,6 +25,7 @@ import (
 // comparer is the struct that has the logic for comparing two rows in the result set
 type comparer struct {
 	orderBy, weightString, starColFixedIndex int
+	collationID                              collations.ID
 	desc                                     bool
 }
 
@@ -39,8 +40,7 @@ func (c *comparer) compare(r1, r2 []sqltypes.Value) (int, error) {
 	} else {
 		colIndex = c.orderBy
 	}
-	// TODO(king-11) make collation aware
-	cmp, err := evalengine.NullsafeCompare(r1[colIndex], r2[colIndex], collations.Unknown)
+	cmp, err := evalengine.NullsafeCompare(r1[colIndex], r2[colIndex], c.collationID)
 	if err != nil {
 		_, isComparisonErr := err.(evalengine.UnsupportedComparisonError)
 		if !(isComparisonErr && c.weightString != -1) {
@@ -49,8 +49,7 @@ func (c *comparer) compare(r1, r2 []sqltypes.Value) (int, error) {
 		// in case of a comparison error switch to using the weight string column for ordering
 		c.orderBy = c.weightString
 		c.weightString = -1
-		// TODO(king-11) make collation aware
-		cmp, err = evalengine.NullsafeCompare(r1[c.orderBy], r2[c.orderBy], collations.Unknown)
+		cmp, err = evalengine.NullsafeCompare(r1[c.orderBy], r2[c.orderBy], c.collationID)
 		if err != nil {
 			return 0, err
 		}
@@ -71,6 +70,7 @@ func extractSlices(input []OrderByParams) []*comparer {
 			weightString:      order.WeightStringCol,
 			desc:              order.Desc,
 			starColFixedIndex: order.StarColFixedIndex,
+			collationID:       order.CollationID,
 		})
 	}
 	return result

go/vt/vtgate/engine/memory_sort_test.go

@@ -19,6 +19,7 @@ package engine
 import (
 	"testing"
 
+	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/test/utils"
 
 	"github.com/stretchr/testify/require"
@@ -206,6 +207,152 @@ func TestMemorySortExecuteWeightString(t *testing.T) {
 	utils.MustMatch(t, wantResult, result)
 }
 
+func TestMemorySortStreamExecuteCollation(t *testing.T) {
+	fields := sqltypes.MakeTestFields(
+		"normal",
+		"varchar",
+	)
+	fp := &fakePrimitive{
+		results: []*sqltypes.Result{sqltypes.MakeTestResult(
+			fields,
+			"c",
+			"d",
+			"cs",
+			"cs",
+			"c",
+		)},
+	}
+
+	collationID, _ := collations.Default().LookupID("utf8mb4_hu_0900_ai_ci")
+	ms := &MemorySort{
+		OrderBy: []OrderByParams{{
+			Col:         0,
+			CollationID: collationID,
+		}},
+		Input: fp,
+	}
+
+	var results []*sqltypes.Result
+	t.Run("order by collation", func(t *testing.T) {
+		results = nil
+		err := ms.TryStreamExecute(&noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
+			results = append(results, qr)
+			return nil
+		})
+		require.NoError(t, err)
+
+		wantResults := sqltypes.MakeTestStreamingResults(
+			fields,
+			"c",
+			"c",
+			"cs",
+			"cs",
+			"d",
+		)
+		utils.MustMatch(t, wantResults, results)
+	})
+
+	t.Run("Descending order by collation", func(t *testing.T) {
+		ms.OrderBy[0].Desc = true
+		fp.rewind()
+		results = nil
+		err := ms.TryStreamExecute(&noopVCursor{}, nil, true, func(qr *sqltypes.Result) error {
+			results = append(results, qr)
+			return nil
+		})
+		require.NoError(t, err)
+
+		wantResults := sqltypes.MakeTestStreamingResults(
+			fields,
+			"d",
+			"cs",
+			"cs",
+			"c",
+			"c",
+		)
+		utils.MustMatch(t, wantResults, results)
+	})
+
+	t.Run("Limit test", func(t *testing.T) {
+		fp.rewind()
+		upperlimit, err := sqlparser.NewPlanValue(sqlparser.NewArgument("__upper_limit"))
+		require.NoError(t, err)
+		ms.UpperLimit = upperlimit
+		bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
+
+		results = nil
+		err = ms.TryStreamExecute(&noopVCursor{}, bv, true, func(qr *sqltypes.Result) error {
+			results = append(results, qr)
+			return nil
+		})
+		require.NoError(t, err)
+
+		wantResults := sqltypes.MakeTestStreamingResults(
+			fields,
+			"d",
+			"cs",
+			"cs",
+		)
+		utils.MustMatch(t, wantResults, results)
+	})
+}
+
+func TestMemorySortExecuteCollation(t *testing.T) {
+	fields := sqltypes.MakeTestFields(
+		"c1",
+		"varchar",
+	)
+	fp := &fakePrimitive{
+		results: []*sqltypes.Result{sqltypes.MakeTestResult(
+			fields,
+			"c",
+			"d",
+			"cs",
+			"cs",
+			"c",
+		)},
+	}
+
+	collationID, _ := collations.Default().LookupID("utf8mb4_hu_0900_ai_ci")
+	ms := &MemorySort{
+		OrderBy: []OrderByParams{{
+			Col:         0,
+			CollationID: collationID,
+		}},
+		Input: fp,
+	}
+
+	result, err := ms.TryExecute(&noopVCursor{}, nil, false)
+	require.NoError(t, err)
+
+	wantResult := sqltypes.MakeTestResult(
+		fields,
+		"c",
+		"c",
+		"cs",
+		"cs",
+		"d",
+	)
+	utils.MustMatch(t, wantResult, result)
+
+	fp.rewind()
+	upperlimit, err := sqlparser.NewPlanValue(sqlparser.NewArgument("__upper_limit"))
+	require.NoError(t, err)
+	ms.UpperLimit = upperlimit
+	bv := map[string]*querypb.BindVariable{"__upper_limit": sqltypes.Int64BindVariable(3)}
+
+	result, err = ms.TryExecute(&noopVCursor{}, bv, false)
+	require.NoError(t, err)
+
+	wantResult = sqltypes.MakeTestResult(
+		fields,
+		"c",
+		"c",
+		"cs",
+	)
+	utils.MustMatch(t, wantResult, result)
+}
+
 func TestMemorySortStreamExecute(t *testing.T) {
 	fields := sqltypes.MakeTestFields(
 		"c1|c2",

go/vt/vtgate/engine/merge_sort_test.go

@@ -20,6 +20,7 @@ import (
 	"errors"
 	"testing"
 
+	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/test/utils"
 
 	"github.com/stretchr/testify/require"
@@ -147,6 +148,68 @@ func TestMergeSortWeightString(t *testing.T) {
 	utils.MustMatch(t, wantResults, results)
 }
 
+func TestMergeSortCollation(t *testing.T) {
+	idColFields := sqltypes.MakeTestFields("normal", "varchar")
+	shardResults := []*shardResult{{
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"c",
+			"---",
+			"d",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"cs",
+			"---",
+			"d",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"cs",
+			"---",
+			"lu",
+		),
+	}, {
+		results: sqltypes.MakeTestStreamingResults(idColFields,
+			"a",
+			"---",
+			"c",
+		),
+	}}
+
+	collationID, _ := collations.Default().LookupID("utf8mb4_hu_0900_ai_ci")
+	orderBy := []OrderByParams{{
+		Col:         0,
+		CollationID: collationID,
+	}}
+
+	var results []*sqltypes.Result
+	err := testMergeSort(shardResults, orderBy, func(qr *sqltypes.Result) error {
+		results = append(results, qr)
+		return nil
+	})
+	require.NoError(t, err)
+
+	// Results are returned one row at a time.
+	wantResults := sqltypes.MakeTestStreamingResults(idColFields,
+		"a",
+		"---",
+		"c",
+		"---",
+		"c",
+		"---",
+		"cs",
+		"---",
+		"cs",
+		"---",
+		"d",
+		"---",
+		"d",
+		"---",
+		"lu",
+	)
+	utils.MustMatch(t, wantResults, results)
+}
+
 // TestMergeSortDescending tests the normal flow of a merge
 // sort where all shards return descending rows.
 func TestMergeSortDescending(t *testing.T) {

go/vt/vtgate/engine/route.go

@@ -24,6 +24,7 @@ import (
 	"strings"
 	"time"
 
+	"vitess.io/vitess/go/mysql/collations"
 	"vitess.io/vitess/go/vt/log"
 	"vitess.io/vitess/go/vt/sqlparser"
 
@@ -134,6 +135,8 @@ type OrderByParams struct {
 	StarColFixedIndex int
 	// v3 specific boolean. Used to also add weight strings originating from GroupBys to the Group by clause
 	FromGroupBy bool
+	// Collation ID for comparison using collation
+	CollationID collations.ID
 }
 
 // String returns a string. Used for plan descriptions
@@ -150,6 +153,10 @@ func (obp OrderByParams) String() string {
 	} else {
 		val += " ASC"
 	}
+	if obp.CollationID != collations.Unknown {
+		collation := collations.Default().LookupByID(obp.CollationID)
+		val += " COLLATE " + collation.Name()
+	}
 	return val
 }