Handle DISTINCT with the new operators (vitessio#13201)

* handle distinct with the new operators Signed-off-by: Andres Taylor <andres@planetscale.com> * feat: do not fail limit when offsets are present Signed-off-by: Manan Gupta <manan@planetscale.com> * test: add more DISTINCT end to end tests Signed-off-by: Andres Taylor <andres@planetscale.com> * remove limitation around aggregation and distinct Signed-off-by: Andres Taylor <andres@planetscale.com> --------- Signed-off-by: Andres Taylor <andres@planetscale.com> Signed-off-by: Manan Gupta <manan@planetscale.com> Co-authored-by: Manan Gupta <manan@planetscale.com>
planetscale · Jun 8, 2023 · 4a71f3d · 4a71f3d
1 parent 435b4cd
commit 4a71f3d
Show file tree

Hide file tree

Showing 14 changed files with 410 additions and 124 deletions.
diff --git a/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go b/go/test/endtoend/vtgate/queries/aggregation/aggregation_test.go
@@ -81,20 +81,6 @@ func TestGroupBy(t *testing.T) {
 	mcmp.AssertMatches("select /*vt+ PLANNER=Gen4 */ id6+id7, count(*) k from t3 group by id6+id7 order by k", `[[INT64(9) INT64(1)] [INT64(6) INT64(2)] [INT64(3) INT64(3)]]`)
 }
 
-func TestDistinct(t *testing.T) {
-	mcmp, closer := start(t)
-	defer closer()
-	mcmp.Exec("insert into t3(id5,id6,id7) values(1,3,3), (2,3,4), (3,3,6), (4,5,7), (5,5,6)")
-	mcmp.Exec("insert into t7_xxhash(uid,phone) values('1',4), ('2',4), ('3',3), ('4',1), ('5',1)")
-	mcmp.Exec("insert into aggr_test(id, val1, val2) values(1,'a',1), (2,'A',1), (3,'b',1), (4,'c',3), (5,'c',4)")
-	mcmp.Exec("insert into aggr_test(id, val1, val2) values(6,'d',null), (7,'e',null), (8,'E',1)")
-	mcmp.AssertMatches("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)]]`)
-	mcmp.AssertMatches("select distinct id6 from t3 join t7_xxhash on t3.id5 = t7_xxhash.phone", `[[INT64(3)] [INT64(5)]]`)
-	mcmp.Exec("delete from t3")
-	mcmp.Exec("delete from t7_xxhash")
-	mcmp.Exec("delete from aggr_test")
-}
-
 func TestEqualFilterOnScatter(t *testing.T) {
 	mcmp, closer := start(t)
 	defer closer()

diff --git a/go/test/endtoend/vtgate/queries/aggregation/distinct_test.go b/go/test/endtoend/vtgate/queries/aggregation/distinct_test.go
@@ -0,0 +1,53 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package aggregation
+
+import (
+	"testing"
+
+	"vitess.io/vitess/go/test/endtoend/utils"
+)
+
+func TestDistinct(t *testing.T) {
+	mcmp, closer := start(t)
+	defer closer()
+	mcmp.Exec("insert into t3(id5,id6,id7) values(1,3,3), (2,3,4), (3,3,6), (4,5,7), (5,5,6)")
+	mcmp.Exec("insert into t7_xxhash(uid,phone) values('1',4), ('2',4), ('3',3), ('4',1), ('5',1)")
+	mcmp.Exec("insert into aggr_test(id, val1, val2) values(1,'a',1), (2,'A',1), (3,'b',1), (4,'c',3), (5,'c',4)")
+	mcmp.Exec("insert into aggr_test(id, val1, val2) values(6,'d',null), (7,'e',null), (8,'E',1)")
+	mcmp.AssertMatches("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)]]`)
+	mcmp.AssertMatches("select distinct id6 from t3 join t7_xxhash on t3.id5 = t7_xxhash.phone", `[[INT64(3)] [INT64(5)]]`)
+}
+
+func TestDistinctIt(t *testing.T) {
+	// tests more variations of DISTINCT
+	mcmp, closer := start(t)
+	defer closer()
+
+	mcmp.Exec("insert into aggr_test(id, val1, val2) values(1,'a',1), (2,'A',1), (3,'b',1), (4,'c',3), (5,'c',4)")
+	mcmp.Exec("insert into aggr_test(id, val1, val2) values(6,'d',null), (7,'e',null), (8,'E',1)")
+
+	mcmp.AssertMatchesNoOrder("select distinct val1 from aggr_test", `[[VARCHAR("c")] [VARCHAR("d")] [VARCHAR("e")] [VARCHAR("a")] [VARCHAR("b")]]`)
+	mcmp.AssertMatchesNoOrder("select distinct val2 from aggr_test", `[[INT64(1)] [INT64(4)] [INT64(3)] [NULL]]`)
+	mcmp.AssertMatchesNoOrder("select distinct id from aggr_test", `[[INT64(1)] [INT64(2)] [INT64(3)] [INT64(5)] [INT64(4)] [INT64(6)] [INT64(7)] [INT64(8)]]`)
+
+	if utils.BinaryIsAtVersion(17, "vtgate") {
+		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)]]`)
+	}
+}
diff --git a/go/vt/vtgate/planbuilder/operator_transformers.go b/go/vt/vtgate/planbuilder/operator_transformers.go
@@ -65,6 +65,8 @@ func transformToLogicalPlan(ctx *plancontext.PlanningContext, op ops.Operator, i
 		return transformOrdering(ctx, op)
 	case *operators.Aggregator:
 		return transformAggregator(ctx, op)
+	case *operators.Distinct:
+		return transformDistinct(ctx, op)
 	}
 
 	return nil, vterrors.VT13001(fmt.Sprintf("unknown type encountered: %T (transformToLogicalPlan)", op))
@@ -112,6 +114,14 @@ func transformAggregator(ctx *plancontext.PlanningContext, op *operators.Aggrega
 	return oa, nil
 }
 
+func transformDistinct(ctx *plancontext.PlanningContext, op *operators.Distinct) (logicalPlan, error) {
+	src, err := transformToLogicalPlan(ctx, op.Source, false)
+	if err != nil {
+		return nil, err
+	}
+	return newDistinct(src, op.Columns /*needToTruncate*/, false), nil
+}
+
 func transformOrdering(ctx *plancontext.PlanningContext, op *operators.Ordering) (logicalPlan, error) {
 	plan, err := transformToLogicalPlan(ctx, op.Source, false)
 	if err != nil {

diff --git a/go/vt/vtgate/planbuilder/operators/aggregation_pushing.go b/go/vt/vtgate/planbuilder/operators/aggregation_pushing.go
@@ -121,7 +121,7 @@ withNextColumn:
 				continue withNextColumn
 			}
 		}
-		pushedAggr.addNoPushCol(aeWrap(col), true)
+		pushedAggr.addColumnWithoutPushing(aeWrap(col), true)
 	}
 
 	// Set the source of the filter to the new aggregator placed below the route.

diff --git a/go/vt/vtgate/planbuilder/operators/aggregator.go b/go/vt/vtgate/planbuilder/operators/aggregator.go
@@ -89,7 +89,7 @@ func (a *Aggregator) AddPredicate(ctx *plancontext.PlanningContext, expr sqlpars
 	return a, nil
 }
 
-func (a *Aggregator) addNoPushCol(expr *sqlparser.AliasedExpr, addToGroupBy bool) int {
+func (a *Aggregator) addColumnWithoutPushing(expr *sqlparser.AliasedExpr, addToGroupBy bool) int {
 	offset := len(a.Columns)
 	a.Columns = append(a.Columns, expr)
 

diff --git a/go/vt/vtgate/planbuilder/operators/distinct.go b/go/vt/vtgate/planbuilder/operators/distinct.go
@@ -0,0 +1,124 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package operators
+
+import (
+	"golang.org/x/exp/slices"
+
+	"vitess.io/vitess/go/vt/sqlparser"
+	"vitess.io/vitess/go/vt/vtgate/engine"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/operators/ops"
+	"vitess.io/vitess/go/vt/vtgate/planbuilder/plancontext"
+)
+
+type (
+	Distinct struct {
+		Source ops.Operator
+		QP     *QueryProjection
+		Pushed bool
+
+		// When offset planning, we'll fill in this field
+		Columns []engine.CheckCol
+	}
+)
+
+func (d *Distinct) planOffsets(ctx *plancontext.PlanningContext) error {
+	columns, err := d.GetColumns()
+	if err != nil {
+		return err
+	}
+	d.Columns = nil
+	var exprs []sqlparser.Expr
+	for _, col := range columns {
+		newSrc, offset, err := d.Source.AddColumn(ctx, col, true, false)
+		if err != nil {
+			return err
+		}
+		d.Source = newSrc
+		e := d.QP.GetSimplifiedExpr(col.Expr)
+		exprs = append(exprs, e)
+		d.Columns = append(d.Columns, engine.CheckCol{
+			Col:       offset,
+			Collation: ctx.SemTable.CollationForExpr(e),
+		})
+	}
+	for i, e := range exprs {
+		if !ctx.SemTable.NeedsWeightString(e) {
+			continue
+		}
+		newSrc, offset, err := d.Source.AddColumn(ctx, aeWrap(weightStringFor(e)), true, false)
+		if err != nil {
+			return err
+		}
+		d.Source = newSrc
+		d.Columns[i].WsCol = &offset
+	}
+	return nil
+}
+
+func (d *Distinct) Clone(inputs []ops.Operator) ops.Operator {
+	return &Distinct{
+		Source:  inputs[0],
+		Columns: slices.Clone(d.Columns),
+		QP:      d.QP,
+		Pushed:  d.Pushed,
+	}
+}
+
+func (d *Distinct) Inputs() []ops.Operator {
+	return []ops.Operator{d.Source}
+}
+
+func (d *Distinct) SetInputs(operators []ops.Operator) {
+	d.Source = operators[0]
+}
+
+func (d *Distinct) AddPredicate(ctx *plancontext.PlanningContext, expr sqlparser.Expr) (ops.Operator, error) {
+	newSrc, err := d.Source.AddPredicate(ctx, expr)
+	if err != nil {
+		return nil, err
+	}
+	d.Source = newSrc
+	return d, nil
+}
+
+func (d *Distinct) AddColumn(ctx *plancontext.PlanningContext, expr *sqlparser.AliasedExpr, reuseExisting, addToGroupBy bool) (ops.Operator, int, error) {
+	newSrc, offset, err := d.Source.AddColumn(ctx, expr, reuseExisting, addToGroupBy)
+	if err != nil {
+		return nil, 0, err
+	}
+	d.Source = newSrc
+	return d, offset, nil
+}
+
+func (d *Distinct) GetColumns() ([]*sqlparser.AliasedExpr, error) {
+	return d.Source.GetColumns()
+}
+
+func (d *Distinct) Description() ops.OpDescription {
+	return ops.OpDescription{
+		OperatorType: "Distinct",
+	}
+}
+
+func (d *Distinct) ShortDescription() string {
+	return ""
+}
+
+func (d *Distinct) GetOrdering() ([]ops.OrderBy, error) {
+	return d.Source.GetOrdering()
+}