Repository: vitessio/vitess
Pull request #7047 (2020-11-17 10:23:49)

Use distinct primitive to solve more queries
Use distinct primitive to solve more queries.



Patch:
diff --git a/go/vt/vtgate/planbuilder/ordered_aggregate.go b/go/vt/vtgate/planbuilder/ordered_aggregate.go
index e7cf2bc05d..261ba72f44 100644
--- a/go/vt/vtgate/planbuilder/ordered_aggregate.go
+++ b/go/vt/vtgate/planbuilder/ordered_aggregate.go
@@ -70,16 +70,11 @@ func (pb *primitiveBuilder) checkAggregates(sel *sqlparser.Select) error {
 	}
 
 	// Check if we can allow aggregates.
-	hasAggregates := false
-	if sel.Distinct {
-		hasAggregates = true
-	} else {
-		hasAggregates = nodeHasAggregates(sel.SelectExprs)
-	}
+	hasAggregates := nodeHasAggregates(sel.SelectExprs)
 	if len(sel.GroupBy) > 0 {
 		hasAggregates = true
 	}
-	if !hasAggregates {
+	if !hasAggregates && !sel.Distinct {
 		return nil
 	}
 
@@ -88,7 +83,11 @@ func (pb *primitiveBuilder) checkAggregates(sel *sqlparser.Select) error {
 	// we need the ability to push down group by and
 	// order by clauses.
 	if !isRoute {
-		return errors.New("unsupported: cross-shard query with aggregates")
+		if hasAggregates {
+			return errors.New("unsupported: cross-shard query with aggregates")
+		}
+		pb.bldr = newDistinct(pb.bldr)
+		return nil
 	}
 
 	// If there is a distinct clause, we can check the select list
@@ -363,7 +362,7 @@ func (oa *orderedAggregate) MakeDistinct() (builder, error) {
 		// So, the distinct 'operator' cannot be pushed down into the
 		// route.
 		if rc.column.Origin() == oa {
-			return nil, errors.New("unsupported: distinct cannot be combined with aggregate functions")
+			return newDistinct(oa), nil
 		}
 		oa.eaggr.Keys = append(oa.eaggr.Keys, i)
 	}
@@ -372,7 +371,7 @@ func (oa *orderedAggregate) MakeDistinct() (builder, error) {
 		return nil, err
 	}
 	oa.input = distinctSrc
-	return oa, err
+	return oa, nil
 }
 
 // PushGroupBy satisfies the builder interface.


Test patch:
diff --git a/go/test/endtoend/vtgate/aggr_test.go b/go/test/endtoend/vtgate/aggr_test.go
index f7035f7da0..f68a5808c3 100644
--- a/go/test/endtoend/vtgate/aggr_test.go
+++ b/go/test/endtoend/vtgate/aggr_test.go
@@ -18,7 +18,6 @@ package vtgate
 
 import (
 	"context"
-	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -33,27 +32,11 @@ func TestAggregateTypes(t *testing.T) {
 	conn, err := mysql.Connect(ctx, &vtParams)
 	require.Nil(t, err)
 	defer conn.Close()
-
 	exec(t, conn, "insert into aggr_test(id, val1, val2) values(1,'a',1), (2,'A',1), (3,'b',1), (4,'c',3), (5,'c',4)")
 	exec(t, conn, "insert into aggr_test(id, val1, val2) values(6,'d',null), (7,'e',null), (8,'E',1)")
-
-	qr := exec(t, conn, "select val1, count(distinct val2), count(*) from aggr_test group by val1")
-	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("d") INT64(0) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)]]`; got != want {
-		t.Errorf("select:\n%v want\n%v", got, want)
-	}
-
-	qr = exec(t, conn, "select val1, sum(distinct val2), sum(val2) from aggr_test group by val1")
-	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("a") DECIMAL(1) DECIMAL(2)] [VARCHAR("b") DECIMAL(1) DECIMAL(1)] [VARCHAR("c") DECIMAL(7) DECIMAL(7)] [VARCHAR("d") NULL NULL] [VARCHAR("e") DECIMAL(1) DECIMAL(1)]]`; got != want {
-		t.Errorf("select:\n%v want\n%v", got, want)
-	}
-
-	qr = exec(t, conn, "select val1, count(distinct val2) k, count(*) from aggr_test group by val1 order by k desc, val1")
-	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)] [VARCHAR("d") INT64(0) INT64(1)]]`; got != want {
-		t.Errorf("select:\n%v want\n%v", got, want)
-	}
-
-	qr = exec(t, conn, "select val1, count(distinct val2) k, count(*) from aggr_test group by val1 order by k desc, val1 limit 4")
-	if got, want := fmt.Sprintf("%v", qr.Rows), `[[VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)]]`; got != want {
-		t.Errorf("select:\n%v want\n%v", got, want)
-	}
+	assertMatches(t, conn, "select val1, count(distinct val2), count(*) from aggr_test group by val1", `[[VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("d") INT64(0) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)]]`)
+	assertMatches(t, conn, "select val1, sum(distinct val2), sum(val2) from aggr_test group by val1", `[[VARCHAR("a") DECIMAL(1) DECIMAL(2)] [VARCHAR("b") DECIMAL(1) DECIMAL(1)] [VARCHAR("c") DECIMAL(7) DECIMAL(7)] [VARCHAR("d") NULL NULL] [VARCHAR("e") DECIMAL(1) DECIMAL(1)]]`)
+	assertMatches(t, conn, "select val1, count(distinct val2) k, count(*) from aggr_test group by val1 order by k desc, val1", `[[VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)] [VARCHAR("d") INT64(0) INT64(1)]]`)
+	assertMatches(t, conn, "select val1, count(distinct val2) k, count(*) from aggr_test group by val1 order by k desc, val1 limit 4", `[[VARCHAR("c") INT64(2) INT64(2)] [VARCHAR("a") INT64(1) INT64(2)] [VARCHAR("b") INT64(1) INT64(1)] [VARCHAR("e") INT64(1) INT64(2)]]`)
+	exec(t, conn, "delete from aggr_test")
 }
diff --git a/go/test/endtoend/vtgate/misc_test.go b/go/test/endtoend/vtgate/misc_test.go
index 22bb16d434..c0d042e00d 100644
--- a/go/test/endtoend/vtgate/misc_test.go
+++ b/go/test/endtoend/vtgate/misc_test.go
@@ -432,6 +432,24 @@ func TestInsertStmtInOLAP(t *testing.T) {
 	assertMatches(t, conn, `select id1 from t1 order by id1`, `[]`)
 }
 
+func TestDistinct(t *testing.T) {
+	defer cluster.PanicHandler(t)
+
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.Nil(t, err)
+	defer conn.Close()
+	exec(t, conn, "insert into t3(id5,id6,id7) values(1,3,3), (2,3,4), (3,3,6), (4,5,7), (5,5,6)")
+	exec(t, conn, "insert into t7_xxhash(uid,phone) values('1',4), ('2',4), ('3',3), ('4',1), ('5',1)")
+	exec(t, conn, "insert into aggr_test(id, val1, val2) values(1,'a',1), (2,'A',1), (3,'b',1), (4,'c',3), (5,'c',4)")
+	exec(t, conn, "insert into aggr_test(id, val1, val2) values(6,'d',null), (7,'e',null), (8,'E',1)")
+	assertMatches(t, conn, "select distinct val2, count(*) from aggr_test group by val2", `[[NULL INT64(2)] [INT64(1) INT64(4)] [INT64(3) INT64(1)] [INT64(4) INT64(1)]]`)
+	assertMatches(t, conn, "select distinct id6 from t3 join t7_xxhash on t3.id5 = t7_xxhash.phone", `[[INT64(3)] [INT64(5)]]`)
+	exec(t, conn, "delete from t3")
+	exec(t, conn, "delete from t7_xxhash")
+	exec(t, conn, "delete from aggr_test")
+}
+
 func assertMatches(t *testing.T, conn *mysql.Conn, query, expected string) {
 	t.Helper()
 	qr := exec(t, conn, query)
diff --git a/go/vt/vtgate/engine/distinct_test.go b/go/vt/vtgate/engine/distinct_test.go
index daaee343ce..c801cc307e 100644
--- a/go/vt/vtgate/engine/distinct_test.go
+++ b/go/vt/vtgate/engine/distinct_test.go
@@ -47,6 +47,10 @@ func TestDistinct(t *testing.T) {
 		testName:       "int64 numbers, two columns",
 		inputs:         r("a|b", "int64|int64", "0|0", "1|1", "1|1", "null|null", "null|null", "1|2"),
 		expectedResult: r("a|b", "int64|int64", "0|0", "1|1", "null|null", "1|2"),
+	}, {
+		testName:       "int64 numbers, two columns",
+		inputs:         r("a|b", "int64|int64", "3|3", "3|3", "3|4", "5|1", "5|1"),
+		expectedResult: r("a|b", "int64|int64", "3|3", "3|4", "5|1"),
 	}, {
 		testName:       "float64 columns designed to produce the same hashcode but not be equal",
 		inputs:         r("a|b", "float64|float64", "0.1|0.2", "0.1|0.3", "0.1|0.4", "0.1|0.5"),
diff --git a/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt b/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
index 9543131bd5..24f594f49a 100644
--- a/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/aggr_cases.txt
@@ -1202,6 +1202,39 @@
   }
 }
 
+# distinct and aggregate functions
+"select distinct a, count(*) from user"
+{
+  "QueryType": "SELECT",
+  "Original": "select distinct a, count(*) from user",
+  "Instructions": {
+    "OperatorType": "Distinct",
+    "Inputs": [
+      {
+        "OperatorType": "Aggregate",
+        "Variant": "Ordered",
+        "Aggregates": "count(1)",
+        "Distinct": "false",
+        "GroupBy": "0",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select a, count(*) from user where 1 != 1",
+            "OrderBy": "0 ASC",
+            "Query": "select a, count(*) from user order by a asc",
+            "Table": "user"
+          }
+        ]
+      }
+    ]
+  }
+}
+
 # Group by invalid column number (code is duplicated from symab).
 "select id from user group by 1.1"
 "column number is not an int"
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.txt b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
index d93be3a39e..b069f8e85f 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
@@ -1619,3 +1619,45 @@
 # Into outfile s3 in sub-query is incorrect
 "select id from (select id from user into outfile s3 'inner_outfile') as t2"
 "Incorrect usage/placement of 'INTO' (errno 1234) (sqlstate 42000)"
+
+# Distinct with cross shard query
+"select distinct user.a from user join user_extra"
+{
+  "QueryType": "SELECT",
+  "Original": "select distinct user.a from user join user_extra",
+  "Instructions": {
+    "OperatorType": "Distinct",
+    "Inputs": [
+      {
+        "OperatorType": "Join",
+        "Variant": "Join",
+        "JoinColumnIndexes": "-1",
+        "TableName": "user_user_extra",
+        "Inputs": [
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select user.a from user where 1 != 1",
+            "Query": "select user.a from user",
+            "Table": "user"
+          },
+          {
+            "OperatorType": "Route",
+            "Variant": "SelectScatter",
+            "Keyspace": {
+              "Name": "user",
+              "Sharded": true
+            },
+            "FieldQuery": "select 1 from user_extra where 1 != 1",
+            "Query": "select 1 from user_extra",
+            "Table": "user_extra"
+          }
+        ]
+      }
+    ]
+  }
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
index a52531c885..063d3332e6 100644
--- a/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/unsupported_cases.txt
@@ -78,10 +78,6 @@
 "select count(*) a from user having a >10"
 "unsupported: filtering on results of aggregates"
 
-# distinct and aggregate functions
-"select distinct a, count(*) from user"
-"unsupported: distinct cannot be combined with aggregate functions"
-
 # group by must reference select list
 "select a from user group by b"
 "unsupported: in scatter query: group by column must reference column in SELECT list"
@@ -134,10 +130,6 @@
 "select count(*) from user join user_extra"
 "unsupported: cross-shard query with aggregates"
 
-# Aggregate detection (distinct)
-"select distinct user.a from user join user_extra"
-"unsupported: cross-shard query with aggregates"
-
 # Aggregate detection (group_concat)
 "select group_concat(user.a) from user join user_extra"
 "unsupported: cross-shard query with aggregates"