Fix and Make aggregation planner handle aggregation functions better

go/test/endtoend/utils/utils.go

@@ -220,15 +220,15 @@ func AssertMatchesWithTimeout(t *testing.T, conn *mysql.Conn, query, expected st
 }
 
 // WaitForAuthoritative waits for a table to become authoritative
-func WaitForAuthoritative(t *testing.T, vtgateProcess cluster.VtgateProcess, ks, tbl string) error {
+func WaitForAuthoritative(t *testing.T, ks, tbl string, readVSchema func() (*interface{}, error)) error {
 	timeout := time.After(10 * time.Second)
 	for {
 		select {
 		case <-timeout:
 			return fmt.Errorf("schema tracking didn't mark table t2 as authoritative until timeout")
 		default:
 			time.Sleep(1 * time.Second)
-			res, err := vtgateProcess.ReadVSchema()
+			res, err := readVSchema()
 			require.NoError(t, err, res)
 			t2Map := getTableT2Map(res, ks, tbl)
 			authoritative, fieldPresent := t2Map["column_list_authoritative"]

go/test/endtoend/vtgate/gen4/gen4_test.go

@@ -430,9 +430,9 @@ func TestOuterJoin(t *testing.T) {
 }
 
 func TestUsingJoin(t *testing.T) {
-	require.NoError(t, utils.WaitForAuthoritative(t, clusterInstance.VtgateProcess, shardedKs, "t1"))
-	require.NoError(t, utils.WaitForAuthoritative(t, clusterInstance.VtgateProcess, shardedKs, "t2"))
-	require.NoError(t, utils.WaitForAuthoritative(t, clusterInstance.VtgateProcess, shardedKs, "t3"))
+	require.NoError(t, utils.WaitForAuthoritative(t, shardedKs, "t1", clusterInstance.VtgateProcess.ReadVSchema))
+	require.NoError(t, utils.WaitForAuthoritative(t, shardedKs, "t2", clusterInstance.VtgateProcess.ReadVSchema))
+	require.NoError(t, utils.WaitForAuthoritative(t, shardedKs, "t3", clusterInstance.VtgateProcess.ReadVSchema))
 
 	mcmp, closer := start(t)
 	defer closer()

go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go

@@ -33,7 +33,7 @@ func start(t *testing.T) (utils.MySQLCompare, func()) {
 	deleteAll := func() {
 		_, _ = utils.ExecAllowError(t, mcmp.VtConn, "set workload = oltp")
 
-		tables := []string{"t9", "aggr_test", "t3", "t7_xxhash", "aggr_test_dates", "t7_xxhash_idx", "t1", "t2"}
+		tables := []string{"t9", "aggr_test", "t3", "t7_xxhash", "aggr_test_dates", "t7_xxhash_idx", "t1", "t2", "t10"}
 		for _, table := range tables {
 			_, _ = mcmp.ExecAndIgnore("delete from " + table)
 		}
@@ -424,3 +424,20 @@ func TestAggregationRandomOnAnAggregatedValue(t *testing.T) {
 	mcmp.AssertMatchesNoOrder("select /*vt+ PLANNER=gen4 */ A.a, A.b, (A.a / A.b) as d from (select sum(a) as a, sum(b) as b from t10 where a = 100) A;",
 		`[[DECIMAL(100) DECIMAL(10) DECIMAL(10.0000)]]`)
 }
+
+func TestBuggyQueries(t *testing.T) {
+	// These queries have been found to be producing the wrong results by the query fuzzer
+	// Adding them as end2end tests to make sure we never get them wrong again
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into t10(k, a, b) values (0, 100, 10), (10, 200, 20), (20, null, null)")
+
+	mcmp.AssertMatches("select /*vt+ PLANNER=Gen4 */ sum(t1.a) from t10 as t1, t10 as t2",
+		`[[DECIMAL(900)]]`)
+
+	mcmp.AssertMatches("select /*vt+ PLANNER=gen4 */t1.a, sum(t1.a), count(*), t1.a, sum(t1.a), count(*) from t10 as t1, t10 as t2 group by t1.a",
+		"[[NULL NULL INT64(3) NULL NULL INT64(3)] "+
+			"[INT32(100) DECIMAL(300) INT64(3) INT32(100) DECIMAL(300) INT64(3)] "+
+			"[INT32(200) DECIMAL(600) INT64(3) INT32(200) DECIMAL(600) INT64(3)]]")
+}

go/test/endtoend/vtgate/queries/aggregation/distinct_test.go

@@ -49,5 +49,6 @@ func TestDistinctIt(t *testing.T) {
 		mcmp.AssertMatches("select /*vt+ PLANNER=Gen4 */ distinct val1 from aggr_test order by val1 desc", `[[VARCHAR("e")] [VARCHAR("d")] [VARCHAR("c")] [VARCHAR("b")] [VARCHAR("a")]]`)
 		mcmp.AssertMatchesNoOrder("select /*vt+ PLANNER=Gen4 */ distinct val1, count(*) from aggr_test group by val1", `[[VARCHAR("a") INT64(2)] [VARCHAR("b") INT64(1)] [VARCHAR("c") INT64(2)] [VARCHAR("d") INT64(1)] [VARCHAR("e") INT64(2)]]`)
 		mcmp.AssertMatchesNoOrder("select /*vt+ PLANNER=Gen4 */ distinct val1+val2 from aggr_test", `[[NULL] [FLOAT64(1)] [FLOAT64(3)] [FLOAT64(4)]]`)
+		mcmp.AssertMatchesNoOrder("select /*vt+ PLANNER=Gen4 */ distinct count(*) from aggr_test group by val1", `[[INT64(2)] [INT64(1)]]`)
 	}
 }

go/test/endtoend/vtgate/queries/foundrows/found_rows_test.go

@@ -44,7 +44,7 @@ func TestFoundRows(t *testing.T) {
 	// Wait for schema tracking to run and mark t2 as authoritative before we try out the queries.
 	// Some of the queries depend on schema tracking to run successfully to be able to replace the StarExpr
 	// in the select clause with the definitive column list.
-	err = utils.WaitForAuthoritative(t, clusterInstance.VtgateProcess, keyspaceName, "t2")
+	err = utils.WaitForAuthoritative(t, keyspaceName, "t2", clusterInstance.VtgateProcess.ReadVSchema)
 	require.NoError(t, err)
 	runTests := func(workload string) {
 		mcmp.AssertFoundRowsValue("select * from t2", workload, 5)

go/vt/vtgate/endtoend/last_insert_id_test.go

@@ -21,15 +21,18 @@ import (
 	"fmt"
 	"testing"
 
-	"vitess.io/vitess/go/vt/vtgate/evalengine"
-
 	"github.com/google/go-cmp/cmp"
 	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/utils"
+	"vitess.io/vitess/go/vt/vtgate/evalengine"
 )
 
 func TestLastInsertId(t *testing.T) {
+	require.NoError(t,
+		utils.WaitForAuthoritative(t, "ks", "t1_last_insert_id", cluster.VTProcess().ReadVSchema))
+
 	ctx := context.Background()
 	conn, err := mysql.Connect(ctx, &vtParams)
 	require.NoError(t, err)
@@ -53,6 +56,9 @@ func TestLastInsertId(t *testing.T) {
 }
 
 func TestLastInsertIdWithRollback(t *testing.T) {
+	require.NoError(t,
+		utils.WaitForAuthoritative(t, "ks", "t1_last_insert_id", cluster.VTProcess().ReadVSchema))
+
 	ctx := context.Background()
 	conn, err := mysql.Connect(ctx, &vtParams)
 	require.NoError(t, err)

go/vt/vtgate/endtoend/main_test.go

@@ -18,17 +18,17 @@ package endtoend
 
 import (
 	"context"
+	_ "embed"
 	"fmt"
 	"os"
 	"testing"
 
 	_flag "vitess.io/vitess/go/internal/flag"
 	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/sqltypes"
-	"vitess.io/vitess/go/vt/vttest"
-
 	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
 	vttestpb "vitess.io/vitess/go/vt/proto/vttest"
+	"vitess.io/vitess/go/vt/vttest"
 )
 
 var (
@@ -37,82 +37,8 @@ var (
 	mysqlParams mysql.ConnParams
 	grpcAddress string
 
-	schema = `
-create table t1(
-	id1 bigint,
-	id2 bigint,
-	primary key(id1)
-) Engine=InnoDB;
-
-create table t1_copy_basic(
-	id1 bigint,
-	id2 bigint,
-	primary key(id1)
-) Engine=InnoDB;
-
-create table t1_copy_all(
-	id1 bigint,
-	id2 bigint,
-	primary key(id1)
-) Engine=InnoDB;
-
-create table t1_copy_resume(
-	id1 bigint,
-	id2 bigint,
-	primary key(id1)
-) Engine=InnoDB;
-
-create table t1_id2_idx(
-	id2 bigint,
-	keyspace_id varbinary(10),
-	primary key(id2)
-) Engine=InnoDB;
-
-create table vstream_test(
-	id bigint,
-	val bigint,
-	primary key(id)
-) Engine=InnoDB;
-
-create table aggr_test(
-	id bigint,
-	val1 varchar(16),
-	val2 bigint,
-	primary key(id)
-) Engine=InnoDB;
-
-create table t2(
-	id3 bigint,
-	id4 bigint,
-	primary key(id3)
-) Engine=InnoDB;
-
-create table t2_id4_idx(
-	id bigint not null auto_increment,
-	id4 bigint,
-	id3 bigint,
-	primary key(id),
-	key idx_id4(id4)
-) Engine=InnoDB;
-
-create table t1_last_insert_id(
-	id bigint not null auto_increment,
-	id1 bigint,
-	primary key(id)
-) Engine=InnoDB;
-
-create table t1_row_count(
-	id bigint not null,
-	id1 bigint,
-	primary key(id)
-) Engine=InnoDB;
-
-create table t1_sharded(
-	id1 bigint,
-	id2 bigint,
-	primary key(id1)
-) Engine=InnoDB;
-`
+	//go:embed schema.sql
+	Schema string
 
 	vschema = &vschemapb.Keyspace{
 		Sharded: true,
@@ -281,7 +207,7 @@ func TestMain(m *testing.M) {
 				},
 			},
 		}
-		if err := cfg.InitSchemas("ks", schema, vschema); err != nil {
+		if err := cfg.InitSchemas("ks", Schema, vschema); err != nil {
 			fmt.Fprintf(os.Stderr, "%v\n", err)
 			os.RemoveAll(cfg.SchemaDir)
 			return 1

go/vt/vtgate/endtoend/schema.sql

@@ -0,0 +1,74 @@
+create table t1(
+                   id1 bigint,
+                   id2 bigint,
+                   primary key(id1)
+) Engine=InnoDB;
+
+create table t1_copy_basic(
+                              id1 bigint,
+                              id2 bigint,
+                              primary key(id1)
+) Engine=InnoDB;
+
+create table t1_copy_all(
+                            id1 bigint,
+                            id2 bigint,
+                            primary key(id1)
+) Engine=InnoDB;
+
+create table t1_copy_resume(
+                               id1 bigint,
+                               id2 bigint,
+                               primary key(id1)
+) Engine=InnoDB;
+
+create table t1_id2_idx(
+                           id2 bigint,
+                           keyspace_id varbinary(10),
+                           primary key(id2)
+) Engine=InnoDB;
+
+create table vstream_test(
+                             id bigint,
+                             val bigint,
+                             primary key(id)
+) Engine=InnoDB;
+
+create table aggr_test(
+                          id bigint,
+                          val1 varchar(16),
+                          val2 bigint,
+                          primary key(id)
+) Engine=InnoDB;
+
+create table t2(
+                   id3 bigint,
+                   id4 bigint,
+                   primary key(id3)
+) Engine=InnoDB;
+
+create table t2_id4_idx(
+                           id bigint not null auto_increment,
+                           id4 bigint,
+                           id3 bigint,
+                           primary key(id),
+                           key idx_id4(id4)
+) Engine=InnoDB;
+
+create table t1_last_insert_id(
+                                  id bigint not null auto_increment,
+                                  id1 bigint,
+                                  primary key(id)
+) Engine=InnoDB;
+
+create table t1_row_count(
+                             id bigint not null,
+                             id1 bigint,
+                             primary key(id)
+) Engine=InnoDB;
+
+create table t1_sharded(
+                           id1 bigint,
+                           id2 bigint,
+                           primary key(id1)
+) Engine=InnoDB;

go/vt/vtgate/engine/distinct.go

@@ -35,7 +35,7 @@ type (
 	Distinct struct {
 		Source    Primitive
 		CheckCols []CheckCol
-		Truncate  bool
+		Truncate  int
 	}
 	CheckCol struct {
 		Col       int
@@ -189,8 +189,8 @@ func (d *Distinct) TryExecute(ctx context.Context, vcursor VCursor, bindVars map
 			result.Rows = append(result.Rows, row)
 		}
 	}
-	if d.Truncate {
-		return result.Truncate(len(d.CheckCols)), nil
+	if d.Truncate > 0 {
+		return result.Truncate(d.Truncate), nil
 	}
 	return result, err
 }
@@ -260,8 +260,8 @@ func (d *Distinct) description() PrimitiveDescription {
 		other["Collations"] = colls
 	}
 
-	if d.Truncate {
-		other["ResultColumns"] = len(d.CheckCols)
+	if d.Truncate > 0 {
+		other["ResultColumns"] = d.Truncate
 	}
 	return PrimitiveDescription{
 		Other:        other,

go/vt/vtgate/engine/distinct_test.go

@@ -96,7 +96,6 @@ func TestDistinct(t *testing.T) {
 			distinct := &Distinct{
 				Source:    &fakePrimitive{results: []*sqltypes.Result{tc.inputs}},
 				CheckCols: checkCols,
-				Truncate:  false,
 			}
 
 			qr, err := distinct.TryExecute(context.Background(), &noopVCursor{}, nil, true)
@@ -145,7 +144,7 @@ func TestWeightStringFallBack(t *testing.T) {
 	distinct := &Distinct{
 		Source:    &fakePrimitive{results: []*sqltypes.Result{input}},
 		CheckCols: checkCols,
-		Truncate:  true,
+		Truncate:  1,
 	}
 
 	qr, err := distinct.TryExecute(context.Background(), &noopVCursor{}, nil, true)