Repository: src-d/go-mysql-server
Pull request #785 (2019-07-05 10:39:21)

sql: implement count distinct

kuba--: LGTM after the tests pass (TestParseErrors/SELECT_COUNT(DISTINCT_foo)_FROM_b (0.00s))
erizocosmico: It's passing now. (CI is not because openjdk10 stuff

Patch:
diff --git a/.travis.yml b/.travis.yml
index 1b32462..c863a5f 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -68,7 +68,7 @@ jobs:
         - make TEST=ruby integration
 
     - language: java
-      jdk: openjdk10
+      jdk: openjdk8
       before_install:
         - eval "$(gimme 1.12.4)"
       install:
@@ -102,4 +102,4 @@ jobs:
       install:
         - go get ./...
       script:
-        - make TEST=c integration
\ No newline at end of file
+        - make TEST=c integration
diff --git a/SUPPORTED.md b/SUPPORTED.md
index 40a31dd..a316dff 100644
--- a/SUPPORTED.md
+++ b/SUPPORTED.md
@@ -18,7 +18,7 @@
 
 ## Grouping expressions
 - AVG
-- COUNT
+- COUNT and COUNT(DISTINCT)
 - MAX
 - MIN
 - SUM (always returns DOUBLE)
diff --git a/sql/analyzer/resolve_having.go b/sql/analyzer/resolve_having.go
index de520d2..fdac4f6 100644
--- a/sql/analyzer/resolve_having.go
+++ b/sql/analyzer/resolve_having.go
@@ -383,6 +383,11 @@ func aggregationEquals(a, b sql.Expression) bool {
 		// the same.
 		_, ok := b.(*aggregation.Count)
 		return ok
+	case *aggregation.CountDistinct:
+		// it doesn't matter what's inside a Count, the result will be
+		// the same.
+		_, ok := b.(*aggregation.CountDistinct)
+		return ok
 	case *aggregation.Sum:
 		b, ok := b.(*aggregation.Sum)
 		if !ok {
diff --git a/sql/expression/function/aggregation/count.go b/sql/expression/function/aggregation/count.go
index b3f900b..0c6a829 100644
--- a/sql/expression/function/aggregation/count.go
+++ b/sql/expression/function/aggregation/count.go
@@ -3,6 +3,7 @@ package aggregation
 import (
 	"fmt"
 
+	"github.com/mitchellh/hashstructure"
 	"github.com/src-d/go-mysql-server/sql"
 	"github.com/src-d/go-mysql-server/sql/expression"
 )
@@ -87,3 +88,93 @@ func (c *Count) Eval(ctx *sql.Context, buffer sql.Row) (interface{}, error) {
 	count := buffer[0]
 	return count, nil
 }
+
+// CountDistinct node to count how many rows are in the result set.
+type CountDistinct struct {
+	expression.UnaryExpression
+}
+
+// NewCountDistinct creates a new CountDistinct node.
+func NewCountDistinct(e sql.Expression) *CountDistinct {
+	return &CountDistinct{expression.UnaryExpression{Child: e}}
+}
+
+// NewBuffer creates a new buffer for the aggregation.
+func (c *CountDistinct) NewBuffer() sql.Row {
+	return sql.NewRow(make(map[uint64]struct{}))
+}
+
+// Type returns the type of the result.
+func (c *CountDistinct) Type() sql.Type {
+	return sql.Int64
+}
+
+// IsNullable returns whether the return value can be null.
+func (c *CountDistinct) IsNullable() bool {
+	return false
+}
+
+// Resolved implements the Expression interface.
+func (c *CountDistinct) Resolved() bool {
+	if _, ok := c.Child.(*expression.Star); ok {
+		return true
+	}
+
+	return c.Child.Resolved()
+}
+
+func (c *CountDistinct) String() string {
+	return fmt.Sprintf("COUNT(DISTINCT %s)", c.Child)
+}
+
+// WithChildren implements the Expression interface.
+func (c *CountDistinct) WithChildren(children ...sql.Expression) (sql.Expression, error) {
+	if len(children) != 1 {
+		return nil, sql.ErrInvalidChildrenNumber.New(c, len(children), 1)
+	}
+	return NewCountDistinct(children[0]), nil
+}
+
+// Update implements the Aggregation interface.
+func (c *CountDistinct) Update(ctx *sql.Context, buffer, row sql.Row) error {
+	seen := buffer[0].(map[uint64]struct{})
+	var value interface{}
+	if _, ok := c.Child.(*expression.Star); ok {
+		value = row
+	} else {
+		v, err := c.Child.Eval(ctx, row)
+		if v == nil {
+			return nil
+		}
+
+		if err != nil {
+			return err
+		}
+
+		value = v
+	}
+
+	hash, err := hashstructure.Hash(value, nil)
+	if err != nil {
+		return fmt.Errorf("count distinct unable to hash value: %s", err)
+	}
+
+	seen[hash] = struct{}{}
+
+	return nil
+}
+
+// Merge implements the Aggregation interface.
+func (c *CountDistinct) Merge(ctx *sql.Context, buffer, partial sql.Row) error {
+	seen := buffer[0].(map[uint64]struct{})
+	for k := range partial[0].(map[uint64]struct{}) {
+		seen[k] = struct{}{}
+	}
+	return nil
+}
+
+// Eval implements the Aggregation interface.
+func (c *CountDistinct) Eval(ctx *sql.Context, buffer sql.Row) (interface{}, error) {
+	seen := buffer[0].(map[uint64]struct{})
+	return int64(len(seen)), nil
+}
diff --git a/sql/parse/parse.go b/sql/parse/parse.go
index 6dd82e4..407ca05 100644
--- a/sql/parse/parse.go
+++ b/sql/parse/parse.go
@@ -13,6 +13,7 @@ import (
 	"github.com/src-d/go-mysql-server/sql"
 	"github.com/src-d/go-mysql-server/sql/expression"
 	"github.com/src-d/go-mysql-server/sql/expression/function"
+	"github.com/src-d/go-mysql-server/sql/expression/function/aggregation"
 	"github.com/src-d/go-mysql-server/sql/plan"
 	"gopkg.in/src-d/go-errors.v1"
 	"vitess.io/vitess/go/vt/sqlparser"
@@ -659,9 +660,11 @@ func getInt64Value(ctx *sql.Context, expr sqlparser.Expr, errStr string) (int64,
 func isAggregate(e sql.Expression) bool {
 	var isAgg bool
 	expression.Inspect(e, func(e sql.Expression) bool {
-		fn, ok := e.(*expression.UnresolvedFunction)
-		if ok {
-			isAgg = isAgg || fn.IsAggregate
+		switch e := e.(type) {
+		case *expression.UnresolvedFunction:
+			isAgg = isAgg || e.IsAggregate
+		case *aggregation.CountDistinct:
+			isAgg = true
 		}
 
 		return true
@@ -789,7 +792,15 @@ func exprToExpression(e sqlparser.Expr) (sql.Expression, error) {
 		}
 
 		if v.Distinct {
-			return nil, ErrUnsupportedSyntax.New("DISTINCT on aggregations")
+			if v.Name.Lowered() != "count" {
+				return nil, ErrUnsupportedSyntax.New("DISTINCT on non-COUNT aggregations")
+			}
+
+			if len(exprs) != 1 {
+				return nil, ErrUnsupportedSyntax.New("more than one expression in COUNT")
+			}
+
+			return aggregation.NewCountDistinct(exprs[0]), nil
 		}
 
 		return expression.NewUnresolvedFunction(v.Name.Lowered(),


Test patch:
diff --git a/engine_test.go b/engine_test.go
index 49fed93..c86f49b 100644
--- a/engine_test.go
+++ b/engine_test.go
@@ -1291,6 +1291,10 @@ var queries = []struct {
 		`SELECT LAST(i) FROM (SELECT i FROM mytable ORDER BY i) t`,
 		[]sql.Row{{int64(3)}},
 	},
+	{
+		`SELECT COUNT(DISTINCT t.i) FROM tabletest t, mytable t2`,
+		[]sql.Row{{int64(3)}},
+	},
 }
 
 func TestQueries(t *testing.T) {
diff --git a/sql/expression/function/aggregation/count_test.go b/sql/expression/function/aggregation/count_test.go
index 316d931..ea27b05 100644
--- a/sql/expression/function/aggregation/count_test.go
+++ b/sql/expression/function/aggregation/count_test.go
@@ -3,19 +3,12 @@ package aggregation
 import (
 	"testing"
 
-	"github.com/stretchr/testify/require"
 	"github.com/src-d/go-mysql-server/sql"
 	"github.com/src-d/go-mysql-server/sql/expression"
+	"github.com/stretchr/testify/require"
 )
 
-func TestCount_String(t *testing.T) {
-	require := require.New(t)
-
-	c := NewCount(expression.NewLiteral("foo", sql.Text))
-	require.Equal(`COUNT("foo")`, c.String())
-}
-
-func TestCount_Eval_1(t *testing.T) {
+func TestCountEval1(t *testing.T) {
 	require := require.New(t)
 	ctx := sql.NewEmptyContext()
 
@@ -37,7 +30,7 @@ func TestCount_Eval_1(t *testing.T) {
 	require.Equal(int64(7), eval(t, c, b))
 }
 
-func TestCount_Eval_Star(t *testing.T) {
+func TestCountEvalStar(t *testing.T) {
 	require := require.New(t)
 	ctx := sql.NewEmptyContext()
 
@@ -45,21 +38,21 @@ func TestCount_Eval_Star(t *testing.T) {
 	b := c.NewBuffer()
 	require.Equal(int64(0), eval(t, c, b))
 
-	c.Update(ctx, b, nil)
-	c.Update(ctx, b, sql.NewRow("foo"))
-	c.Update(ctx, b, sql.NewRow(1))
-	c.Update(ctx, b, sql.NewRow(nil))
-	c.Update(ctx, b, sql.NewRow(1, 2, 3))
+	require.NoError(c.Update(ctx, b, nil))
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(nil)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1, 2, 3)))
 	require.Equal(int64(5), eval(t, c, b))
 
 	b2 := c.NewBuffer()
-	c.Update(ctx, b2, sql.NewRow())
-	c.Update(ctx, b2, sql.NewRow("foo"))
-	c.Merge(ctx, b, b2)
+	require.NoError(c.Update(ctx, b2, sql.NewRow()))
+	require.NoError(c.Update(ctx, b2, sql.NewRow("foo")))
+	require.NoError(c.Merge(ctx, b, b2))
 	require.Equal(int64(7), eval(t, c, b))
 }
 
-func TestCount_Eval_String(t *testing.T) {
+func TestCountEvalString(t *testing.T) {
 	require := require.New(t)
 	ctx := sql.NewEmptyContext()
 
@@ -67,9 +60,66 @@ func TestCount_Eval_String(t *testing.T) {
 	b := c.NewBuffer()
 	require.Equal(int64(0), eval(t, c, b))
 
-	c.Update(ctx, b, sql.NewRow("foo"))
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
 	require.Equal(int64(1), eval(t, c, b))
 
-	c.Update(ctx, b, sql.NewRow(nil))
+	require.NoError(c.Update(ctx, b, sql.NewRow(nil)))
 	require.Equal(int64(1), eval(t, c, b))
 }
+
+func TestCountDistinctEval1(t *testing.T) {
+	require := require.New(t)
+	ctx := sql.NewEmptyContext()
+
+	c := NewCountDistinct(expression.NewLiteral(1, sql.Int32))
+	b := c.NewBuffer()
+	require.Equal(int64(0), eval(t, c, b))
+
+	require.NoError(c.Update(ctx, b, nil))
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(nil)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1, 2, 3)))
+	require.Equal(int64(1), eval(t, c, b))
+}
+
+func TestCountDistinctEvalStar(t *testing.T) {
+	require := require.New(t)
+	ctx := sql.NewEmptyContext()
+
+	c := NewCountDistinct(expression.NewStar())
+	b := c.NewBuffer()
+	require.Equal(int64(0), eval(t, c, b))
+
+	require.NoError(c.Update(ctx, b, nil))
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(nil)))
+	require.NoError(c.Update(ctx, b, sql.NewRow(1, 2, 3)))
+	require.Equal(int64(5), eval(t, c, b))
+
+	b2 := c.NewBuffer()
+	require.NoError(c.Update(ctx, b2, sql.NewRow(1)))
+	require.NoError(c.Update(ctx, b2, sql.NewRow("foo")))
+	require.NoError(c.Update(ctx, b2, sql.NewRow(5)))
+	require.NoError(c.Merge(ctx, b, b2))
+
+	require.Equal(int64(6), eval(t, c, b))
+}
+
+func TestCountDistinctEvalString(t *testing.T) {
+	require := require.New(t)
+	ctx := sql.NewEmptyContext()
+
+	c := NewCountDistinct(expression.NewGetField(0, sql.Text, "", true))
+	b := c.NewBuffer()
+	require.Equal(int64(0), eval(t, c, b))
+
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
+	require.Equal(int64(1), eval(t, c, b))
+
+	require.NoError(c.Update(ctx, b, sql.NewRow(nil)))
+	require.NoError(c.Update(ctx, b, sql.NewRow("foo")))
+	require.NoError(c.Update(ctx, b, sql.NewRow("bar")))
+	require.Equal(int64(2), eval(t, c, b))
+}
diff --git a/sql/parse/parse_test.go b/sql/parse/parse_test.go
index 069d3a4..849f9cf 100644
--- a/sql/parse/parse_test.go
+++ b/sql/parse/parse_test.go
@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/src-d/go-mysql-server/sql/expression"
+	"github.com/src-d/go-mysql-server/sql/expression/function/aggregation"
 	"github.com/src-d/go-mysql-server/sql/plan"
 	"gopkg.in/src-d/go-errors.v1"
 
@@ -1158,6 +1159,13 @@ var fixtures = map[string]sql.Node{
 		[]sql.Expression{},
 		plan.NewUnresolvedTable("foo", ""),
 	),
+	`SELECT COUNT(DISTINCT i) FROM foo`: plan.NewGroupBy(
+		[]sql.Expression{
+			aggregation.NewCountDistinct(expression.NewUnresolvedColumn("i")),
+		},
+		[]sql.Expression{},
+		plan.NewUnresolvedTable("foo", ""),
+	),
 }
 
 func TestParse(t *testing.T) {
@@ -1191,7 +1199,7 @@ var fixturesErrors = map[string]*errors.Kind{
 	`SELECT '2018-05-01' / INTERVAL 1 DAY`:                    ErrUnsupportedSyntax,
 	`SELECT INTERVAL 1 DAY + INTERVAL 1 DAY`:                  ErrUnsupportedSyntax,
 	`SELECT '2018-05-01' + (INTERVAL 1 DAY + INTERVAL 1 DAY)`: ErrUnsupportedSyntax,
-	`SELECT COUNT(DISTINCT foo) FROM b`:                       ErrUnsupportedSyntax,
+	`SELECT AVG(DISTINCT foo) FROM b`:                         ErrUnsupportedSyntax,
 }
 
 func TestParseErrors(t *testing.T) {