extracted MakeDistinct() from the builder interface

Signed-off-by: Andres Taylor <[email protected]>

go/vt/sqlparser/ast.go

@@ -51,6 +51,7 @@ type (
 		AddOrder(*Order)
 		SetLimit(*Limit)
 		SetLock(lock Lock)
+		MakeDistinct()
 		SQLNode
 	}
 

go/vt/sqlparser/ast_funcs.go

@@ -733,6 +733,11 @@ func (node *Select) SetLock(lock Lock) {
 	node.Lock = lock
 }
 
+// MakeDistinct makes the statement distinct
+func (node *Select) MakeDistinct() {
+	node.Distinct = true
+}
+
 // AddWhere adds the boolean expression to the
 // WHERE clause as an AND condition.
 func (node *Select) AddWhere(expr Expr) {
@@ -780,6 +785,11 @@ func (node *ParenSelect) SetLock(lock Lock) {
 	node.Select.SetLock(lock)
 }
 
+// MakeDistinct implements the SelectStatement interface
+func (node *ParenSelect) MakeDistinct() {
+	node.Select.MakeDistinct()
+}
+
 // AddOrder adds an order by element
 func (node *Union) AddOrder(order *Order) {
 	node.OrderBy = append(node.OrderBy, order)
@@ -795,6 +805,11 @@ func (node *Union) SetLock(lock Lock) {
 	node.Lock = lock
 }
 
+// MakeDistinct implements the SelectStatement interface
+func (node *Union) MakeDistinct() {
+	node.UnionSelects[len(node.UnionSelects)-1].Distinct = true
+}
+
 //Unionize returns a UNION, either creating one or adding SELECT to an existing one
 func Unionize(lhs, rhs SelectStatement, distinct bool, by OrderBy, limit *Limit, lock Lock) *Union {
 	union, isUnion := lhs.(*Union)

go/vt/vtgate/planbuilder/builder.go

@@ -52,9 +52,6 @@ type builder interface {
 	// execute before this one.
 	Reorder(int)
 
-	// MakeDistinct makes the primitive handle the distinct clause.
-	MakeDistinct() (builder, error)
-
 	// PushOrderBy pushes the ORDER BY clause. It returns the
 	// the current primitive or a replacement if a new one was
 	// created.

go/vt/vtgate/planbuilder/concatenate.go

@@ -65,10 +65,6 @@ func (c *concatenate) SupplyWeightString(colNumber int) (weightcolNumber int, er
 	panic("implement me")
 }
 
-func (c *concatenate) MakeDistinct() (builder, error) {
-	return newDistinct(c), nil
-}
-
 func (c *concatenate) PushOrderBy(by sqlparser.OrderBy) (builder, error) {
 	if by == nil {
 		return c, nil

go/vt/vtgate/planbuilder/distinct.go

@@ -40,10 +40,6 @@ func newDistinct(source builder) builder {
 	}
 }
 
-func (d *distinct) MakeDistinct() (builder, error) {
-	return d, nil
-}
-
 func (d *distinct) PushOrderBy(by sqlparser.OrderBy) (builder, error) {
 	orderBy, err := d.input.PushOrderBy(by)
 	if err != nil {

go/vt/vtgate/planbuilder/grouping.go

@@ -93,3 +93,51 @@ func planGroupBy(pb *primitiveBuilder, input builder, groupBy sqlparser.GroupBy)
 	}
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "%T.groupBy: unreachable", input)
 }
+
+// planDistinct makes the output distinct
+func planDistinct(input builder) (builder, error) {
+	switch node := input.(type) {
+	case *mergeSort, *pulloutSubquery:
+		inputs := node.Inputs()
+		input := inputs[0]
+
+		newInput, err := planDistinct(input)
+		if err != nil {
+			return nil, err
+		}
+		inputs[0] = newInput
+		node.Rewrite(inputs...)
+		return node, nil
+	case *route:
+		node.Select.MakeDistinct()
+		return node, nil
+	case *orderedAggregate:
+		for i, rc := range node.resultColumns {
+			// If the column origin is oa (and not the underlying route),
+			// it means that it's an aggregate function supplied by oa.
+			// So, the distinct 'operator' cannot be pushed down into the
+			// route.
+			if rc.column.Origin() == node {
+				return newDistinct(node), nil
+			}
+			node.eaggr.Keys = append(node.eaggr.Keys, i)
+		}
+		newInput, err := planDistinct(node.input)
+		if err != nil {
+			return nil, err
+		}
+		node.input = newInput
+		return node, nil
+
+	case *subquery:
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: distinct on cross-shard subquery")
+	case *concatenate:
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "only union-all is supported for this operator")
+	case *join:
+		return nil, vterrors.New(vtrpcpb.Code_UNIMPLEMENTED, "unsupported: distinct on cross-shard join")
+	case *distinct:
+		return input, nil
+	}
+
+	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "%T.distinct: unreachable", input)
+}

go/vt/vtgate/planbuilder/join.go

@@ -142,11 +142,6 @@ func (jb *join) ResultColumns() []*resultColumn {
 	return jb.resultColumns
 }
 
-// MakeDistinct satisfies the builder interface.
-func (jb *join) MakeDistinct() (builder, error) {
-	return nil, errors.New("unsupported: distinct on cross-shard join")
-}
-
 // PushOrderBy satisfies the builder interface.
 func (jb *join) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	isSpecial := false

go/vt/vtgate/planbuilder/limit.go

@@ -49,11 +49,6 @@ func (l *limit) Primitive() engine.Primitive {
 	return l.elimit
 }
 
-// MakeDistinct satisfies the builder interface.
-func (l *limit) MakeDistinct() (builder, error) {
-	return nil, errors.New("limit.MakeDistinct: unreachable")
-}
-
 // PushGroupBy satisfies the builder interface.
 func (l *limit) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	return nil, errors.New("limit.PushOrderBy: unreachable")

go/vt/vtgate/planbuilder/memory_sort.go

@@ -83,11 +83,6 @@ func (ms *memorySort) Primitive() engine.Primitive {
 	return ms.eMemorySort
 }
 
-// MakeDistinct satisfies the builder interface.
-func (ms *memorySort) MakeDistinct() (builder, error) {
-	return nil, errors.New("memorySort.MakeDistinct: unreachable")
-}
-
 // PushGroupBy satisfies the builder interface.
 func (ms *memorySort) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	return nil, errors.New("memorySort.PushOrderBy: unreachable")

go/vt/vtgate/planbuilder/merge_sort.go

@@ -60,16 +60,6 @@ func (ms *mergeSort) Primitive() engine.Primitive {
 	return ms.input.Primitive()
 }
 
-// MakeDistinct satisfies the builder interface.
-func (ms *mergeSort) MakeDistinct() (builder, error) {
-	distinctSrc, err := ms.input.MakeDistinct()
-	if err != nil {
-		return nil, err
-	}
-	ms.input = distinctSrc
-	return ms, err
-}
-
 // PushOrderBy satisfies the builder interface.
 // A merge sort is created due to the push of an ORDER BY clause.
 // So, this function should never get called.

go/vt/vtgate/planbuilder/ordered_aggregate.go

@@ -323,25 +323,6 @@ func (oa *orderedAggregate) needDistinctHandling(pb *primitiveBuilder, funcExpr
 	return true, innerAliased, nil
 }
 
-func (oa *orderedAggregate) MakeDistinct() (builder, error) {
-	for i, rc := range oa.resultColumns {
-		// If the column origin is oa (and not the underlying route),
-		// it means that it's an aggregate function supplied by oa.
-		// So, the distinct 'operator' cannot be pushed down into the
-		// route.
-		if rc.column.Origin() == oa {
-			return newDistinct(oa), nil
-		}
-		oa.eaggr.Keys = append(oa.eaggr.Keys, i)
-	}
-	distinctSrc, err := oa.input.MakeDistinct()
-	if err != nil {
-		return nil, err
-	}
-	oa.input = distinctSrc
-	return oa, nil
-}
-
 // PushOrderBy pushes the order by expression into the primitive.
 // The requested order must be such that the ordering can be done
 // before the group by, which will allow us to push it down to the

go/vt/vtgate/planbuilder/postprocess.go

@@ -28,7 +28,8 @@ import (
 // and ensures that there are no subqueries.
 func (pb *primitiveBuilder) pushGroupBy(sel *sqlparser.Select) error {
 	if sel.Distinct {
-		newBuilder, err := pb.bldr.MakeDistinct()
+
+		newBuilder, err := planDistinct(pb.bldr)
 		if err != nil {
 			return err
 		}

go/vt/vtgate/planbuilder/pullout_subquery.go

@@ -78,16 +78,6 @@ func (ps *pulloutSubquery) ResultColumns() []*resultColumn {
 	return ps.underlying.ResultColumns()
 }
 
-// MakeDistinct satisfies the builder interface.
-func (ps *pulloutSubquery) MakeDistinct() (builder, error) {
-	distinctUnderlying, err := ps.underlying.MakeDistinct()
-	if err != nil {
-		return nil, err
-	}
-	ps.underlying = distinctUnderlying
-	return ps, err
-}
-
 // PushOrderBy satisfies the builder interface.
 func (ps *pulloutSubquery) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	bldr, err := ps.underlying.PushOrderBy(orderBy)

go/vt/vtgate/planbuilder/route.go

@@ -126,16 +126,6 @@ func (rb *route) PushAnonymous(expr sqlparser.SelectExpr) *resultColumn {
 	return rc
 }
 
-// MakeDistinct satisfies the builder interface.
-func (rb *route) MakeDistinct() (builder, error) {
-	s, ok := rb.Select.(*sqlparser.Select)
-	if !ok {
-		return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unexpected AST struct for query")
-	}
-	s.Distinct = true
-	return rb, nil
-}
-
 // PushOrderBy satisfies the builder interface.
 func (rb *route) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	switch len(orderBy) {

go/vt/vtgate/planbuilder/sql_calc_found_rows.go

@@ -64,11 +64,6 @@ func (s *sqlCalcFoundRows) Reorder(order int) {
 	s.LimitQuery.Reorder(order)
 }
 
-//MakeDistinct implements the builder interface
-func (s *sqlCalcFoundRows) MakeDistinct() (builder, error) {
-	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unreachable: sqlCalcFoundRows.MakeDistinct")
-}
-
 //PushOrderBy implements the builder interface
 func (s *sqlCalcFoundRows) PushOrderBy(sqlparser.OrderBy) (builder, error) {
 	return nil, vterrors.Errorf(vtrpcpb.Code_INTERNAL, "unreachable: sqlCalcFoundRows.PushOrderBy")

go/vt/vtgate/planbuilder/subquery.go

@@ -17,7 +17,6 @@ limitations under the License.
 package planbuilder
 
 import (
-	"errors"
 	"fmt"
 
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -79,11 +78,6 @@ func (sq *subquery) ResultColumns() []*resultColumn {
 	return sq.resultColumns
 }
 
-// MakeDistinct satisfies the builder interface.
-func (sq *subquery) MakeDistinct() (builder, error) {
-	return nil, errors.New("unsupported: distinct on cross-shard subquery")
-}
-
 // PushOrderBy satisfies the builder interface.
 func (sq *subquery) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	if len(orderBy) == 0 {

go/vt/vtgate/planbuilder/vindex_func.go

@@ -89,11 +89,6 @@ func (vf *vindexFunc) ResultColumns() []*resultColumn {
 	return vf.resultColumns
 }
 
-// MakeDistinct satisfies the builder interface.
-func (vf *vindexFunc) MakeDistinct() (builder, error) {
-	return nil, errors.New("unsupported: distinct on vindex function")
-}
-
 // PushOrderBy satisfies the builder interface.
 func (vf *vindexFunc) PushOrderBy(orderBy sqlparser.OrderBy) (builder, error) {
 	if len(orderBy) == 0 {