Merge pull request #8453 from planetscale/gen4-groupby

Gen4: groupby support
vitessio · Jul 15, 2021 · 1312ec8 · 1312ec8
2 parents 4c4e895 + 041bb58
commit 1312ec8
Show file tree

Hide file tree

Showing 37 changed files with 1,359 additions and 442 deletions.
diff --git a/go/mysql/sql_error.go b/go/mysql/sql_error.go
@@ -176,6 +176,7 @@ var stateToMysqlCode = map[vterrors.State]struct {
 	vterrors.NoSuchTable:                  {num: ERNoSuchTable, state: SSUnknownTable},
 	vterrors.NotSupportedYet:              {num: ERNotSupportedYet, state: SSClientError},
 	vterrors.ForbidSchemaChange:           {num: ERForbidSchemaChange, state: SSUnknownSQLState},
+	vterrors.MixOfGroupFuncAndFields:      {num: ERMixOfGroupFuncAndFields, state: SSClientError},
 	vterrors.NetPacketTooLarge:            {num: ERNetPacketTooLarge, state: SSNetError},
 	vterrors.NonUniqError:                 {num: ERNonUniq, state: SSConstraintViolation},
 	vterrors.NonUniqTable:                 {num: ERNonUniqTable, state: SSClientError},

diff --git a/go/test/endtoend/vtgate/gen4/gen4_test.go b/go/test/endtoend/vtgate/gen4/gen4_test.go
@@ -34,6 +34,10 @@ func TestOrderBy(t *testing.T) {
 	require.NoError(t, err)
 	defer conn.Close()
 
+	defer func() {
+		_, _ = exec(t, conn, `delete from t1`)
+	}()
+
 	// insert some data.
 	checkedExec(t, conn, `insert into t1(id, col) values (100, 123),(10, 12),(1, 13),(1000, 1234)`)
 
@@ -48,6 +52,40 @@ func TestOrderBy(t *testing.T) {
 	require.Error(t, err)
 }
 
+func TestGroupBy(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	defer func() {
+		_, _ = exec(t, conn, `delete from t1`)
+		_, _ = exec(t, conn, `delete from t2`)
+	}()
+
+	// insert some data.
+	checkedExec(t, conn, `insert into t1(id, col) values (1, 123),(2, 12),(3, 13),(4, 1234)`)
+	checkedExec(t, conn, `insert into t2(id, tcol1, tcol2) values (1, 'A', 'A'),(2, 'B', 'C'),(3, 'A', 'C'),(4, 'C', 'A'),(5, 'A', 'A'),(6, 'B', 'C'),(7, 'B', 'A'),(8, 'C', 'B')`)
+
+	// Gen4 only supported query.
+	assertMatches(t, conn, `select tcol2, tcol1, count(id) from t2 group by tcol2, tcol1`,
+		`[[VARCHAR("A") VARCHAR("A") INT64(2)] [VARCHAR("A") VARCHAR("B") INT64(1)] [VARCHAR("A") VARCHAR("C") INT64(1)] [VARCHAR("B") VARCHAR("C") INT64(1)] [VARCHAR("C") VARCHAR("A") INT64(1)] [VARCHAR("C") VARCHAR("B") INT64(2)]]`)
+
+	assertMatches(t, conn, `select tcol1, tcol1 from t2 order by tcol1`,
+		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")] [VARCHAR("C") VARCHAR("C")]]`)
+
+	assertMatches(t, conn, `select tcol1, tcol1 from t1 join t2 on t1.id = t2.id order by tcol1`,
+		`[[VARCHAR("A") VARCHAR("A")] [VARCHAR("A") VARCHAR("A")] [VARCHAR("B") VARCHAR("B")] [VARCHAR("C") VARCHAR("C")]]`)
+
+	assertMatches(t, conn, `select count(*) k, tcol1, tcol2, "abc" b from t2 group by tcol1, tcol2, b order by k, tcol2, tcol1`,
+		`[[INT64(1) VARCHAR("B") VARCHAR("A") VARCHAR("abc")] `+
+			`[INT64(1) VARCHAR("C") VARCHAR("A") VARCHAR("abc")] `+
+			`[INT64(1) VARCHAR("C") VARCHAR("B") VARCHAR("abc")] `+
+			`[INT64(1) VARCHAR("A") VARCHAR("C") VARCHAR("abc")] `+
+			`[INT64(2) VARCHAR("A") VARCHAR("A") VARCHAR("abc")] `+
+			`[INT64(2) VARCHAR("B") VARCHAR("C") VARCHAR("abc")]]`)
+}
+
 func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
 	t.Helper()
 	qr := checkedExec(t, conn, query)

diff --git a/go/test/endtoend/vtgate/gen4/main_test.go b/go/test/endtoend/vtgate/gen4/main_test.go
@@ -35,6 +35,13 @@ var (
 	col bigint,
 	primary key(id)
 ) Engine=InnoDB;
+
+create table t2(
+	id bigint,
+	tcol1 varchar(50),
+	tcol2 varchar(50),
+	primary key(id)
+) Engine=InnoDB;
 `
 
 	VSchema = `
@@ -53,6 +60,20 @@ var (
           "name": "xxhash"
         }
       ]
+    },
+    "t2": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ],
+      "columns": [
+        {
+          "name": "tcol1",
+          "type": "VARCHAR"
+        }
+      ]
     }
   }
 }`
@@ -84,7 +105,7 @@ func TestMain(m *testing.M) {
 		}
 
 		// Start vtgate
-		clusterInstance.VtGateExtraArgs = []string{"-planner_version", "Gen4Fallback"} // enable Gen4 planner.
+		clusterInstance.VtGateExtraArgs = []string{"-planner_version", "Gen4"} // enable Gen4 planner.
 		err = clusterInstance.StartVtgate()
 		if err != nil {
 			return 1

diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go
@@ -1366,3 +1366,22 @@ func ToString(exprs []TableExpr) string {
 	}
 	return buf.String()
 }
+
+// ContainsAggregation returns true if the expression contains aggregation
+func ContainsAggregation(e Expr) bool {
+	hasAggregates := false
+	_ = Walk(func(node SQLNode) (kontinue bool, err error) {
+		switch node := node.(type) {
+		case *FuncExpr:
+			if node.IsAggregate() {
+				hasAggregates = true
+				return false, nil
+			}
+		case *GroupConcatExpr:
+			hasAggregates = true
+			return false, nil
+		}
+		return true, nil
+	}, e)
+	return hasAggregates
+}
diff --git a/go/vt/vterrors/state.go b/go/vt/vterrors/state.go
@@ -39,6 +39,7 @@ const (
 	WrongTypeForVar
 	WrongValueForVar
 	LockOrActiveTransaction
+	MixOfGroupFuncAndFields
 
 	// failed precondition
 	NoDB

diff --git a/go/vt/vtgate/engine/cached_size.go b/go/vt/vtgate/engine/cached_size.go
diff --git a/go/vt/vtgate/engine/comparer.go b/go/vt/vtgate/engine/comparer.go
@@ -59,8 +59,8 @@ func (c *comparer) compare(r1, r2 []sqltypes.Value) (int, error) {
 	return cmp, nil
 }
 
-// extractSlices extracts the three fields of OrderbyParams into a slice of comparers
-func extractSlices(input []OrderbyParams) []*comparer {
+// extractSlices extracts the three fields of OrderByParams into a slice of comparers
+func extractSlices(input []OrderByParams) []*comparer {
 	var result []*comparer
 	for _, order := range input {
 		result = append(result, &comparer{

diff --git a/go/vt/vtgate/engine/memory_sort.go b/go/vt/vtgate/engine/memory_sort.go
@@ -35,7 +35,7 @@ var _ Primitive = (*MemorySort)(nil)
 // MemorySort is a primitive that performs in-memory sorting.
 type MemorySort struct {
 	UpperLimit sqltypes.PlanValue
-	OrderBy    []OrderbyParams
+	OrderBy    []OrderByParams
 	Input      Primitive
 
 	// TruncateColumnCount specifies the number of columns to return
@@ -183,6 +183,9 @@ func (ms *MemorySort) description() PrimitiveDescription {
 	if !value.IsNull() {
 		other["UpperLimit"] = value.String()
 	}
+	if ms.TruncateColumnCount > 0 {
+		other["ResultColumns"] = ms.TruncateColumnCount
+	}
 	return PrimitiveDescription{
 		OperatorType: "Sort",
 		Variant:      "Memory",
@@ -191,7 +194,7 @@ func (ms *MemorySort) description() PrimitiveDescription {
 }
 
 func orderByParamsToString(i interface{}) string {
-	return i.(OrderbyParams).String()
+	return i.(OrderByParams).String()
 }
 
 //GenericJoin will iterate over arrays, slices or maps, and executes the f function to get a

diff --git a/go/vt/vtgate/engine/memory_sort_test.go b/go/vt/vtgate/engine/memory_sort_test.go
@@ -45,7 +45,7 @@ func TestMemorySortExecute(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: -1,
 			Col:             1,
 		}},
@@ -100,7 +100,7 @@ func TestMemorySortStreamExecuteWeightString(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: 0,
 			Col:             1,
 		}},
@@ -168,7 +168,7 @@ func TestMemorySortExecuteWeightString(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: 1,
 			Col:             0,
 		}},
@@ -223,7 +223,7 @@ func TestMemorySortStreamExecute(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: -1,
 			Col:             1,
 		}},
@@ -302,7 +302,7 @@ func TestMemorySortExecuteTruncate(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: -1,
 			Col:             1,
 		}},
@@ -341,7 +341,7 @@ func TestMemorySortStreamExecuteTruncate(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: -1,
 			Col:             1,
 		}},
@@ -384,7 +384,7 @@ func TestMemorySortMultiColumn(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			Col:             1,
 			WeightStringCol: -1,
 		}, {
@@ -459,7 +459,7 @@ func TestMemorySortMaxMemoryRows(t *testing.T) {
 		}
 
 		ms := &MemorySort{
-			OrderBy: []OrderbyParams{{
+			OrderBy: []OrderByParams{{
 				WeightStringCol: -1,
 				Col:             1,
 			}},
@@ -495,7 +495,7 @@ func TestMemorySortExecuteNoVarChar(t *testing.T) {
 	}
 
 	ms := &MemorySort{
-		OrderBy: []OrderbyParams{{
+		OrderBy: []OrderByParams{{
 			WeightStringCol: -1,
 			Col:             0,
 		}},

diff --git a/go/vt/vtgate/engine/merge_sort.go b/go/vt/vtgate/engine/merge_sort.go
@@ -50,7 +50,7 @@ var _ Primitive = (*MergeSort)(nil)
 // so that vdiff can use it. In that situation, only StreamExecute is used.
 type MergeSort struct {
 	Primitives              []StreamExecutor
-	OrderBy                 []OrderbyParams
+	OrderBy                 []OrderByParams
 	ScatterErrorsAsWarnings bool
 	noInputs
 	noTxNeeded