Repository: proullon/ramsql
Pull request #51 (2021-11-04 12:02:16)

Support DISTINCT and DISTINCT ON 
This PR adds support for `DISTINCT` and PostgreSQL specific `DISTINCT ON` variant. 


Patch:
diff --git a/engine/distinct.go b/engine/distinct.go
new file mode 100644
index 0000000..09b4c02
--- /dev/null
+++ b/engine/distinct.go
@@ -0,0 +1,110 @@
+package engine
+
+import (
+	"github.com/proullon/ramsql/engine/log"
+	"github.com/proullon/ramsql/engine/protocol"
+)
+
+type distinct struct {
+	realConn protocol.EngineConn
+	seen     seen
+	len      int
+}
+
+func distinctedConn(conn protocol.EngineConn, len int) protocol.EngineConn {
+	return &distinct{
+		realConn: conn,
+		len:      len,
+		seen:     make(seen),
+	}
+}
+
+// Not needed
+func (l *distinct) ReadStatement() (string, error) {
+	log.Debug("limit.ReadStatement: should not be used\n")
+	return "", nil
+}
+
+// Not needed
+func (l *distinct) WriteResult(last int64, ra int64) error {
+	log.Debug("limit.WriteResult: should not be used\n")
+	return nil
+}
+
+func (l *distinct) WriteError(err error) error {
+	return l.realConn.WriteError(err)
+}
+
+func (l *distinct) WriteRowHeader(header []string) error {
+	if l.len > 0 {
+		// Postgres returns only columns outside of DISTINCT ON
+		return l.realConn.WriteRowHeader(header[l.len:])
+	}
+	return l.realConn.WriteRowHeader(header)
+}
+
+func (l *distinct) WriteRow(row []string) error {
+	if l.len > 0 {
+		if l.seen.exists(row[:l.len]) {
+			return nil
+		}
+		// Postgres returns only columns outside of DISTINCT ON
+		return l.realConn.WriteRow(row[l.len:])
+	} else {
+		if l.seen.exists(row) {
+			return nil
+		}
+		return l.realConn.WriteRow(row)
+	}
+}
+
+func (l *distinct) WriteRowEnd() error {
+	return l.realConn.WriteRowEnd()
+}
+
+func (l *distinct) equalRows(a, b []string) bool {
+	if l.len > 0 {
+		if len(a) < l.len || len(b) < l.len {
+			return false
+		}
+
+		for idx := 0; idx < l.len; idx++ {
+			if a[idx] != b[idx] {
+				return false
+			}
+		}
+
+		return true
+	}
+
+	if len(a) != len(b) {
+		return false
+	}
+	for idx := range a {
+		if a[idx] != b[idx] {
+			return false
+		}
+	}
+
+	return true
+}
+
+type seen map[string]seen
+
+func (s seen) exists(r []string) bool {
+	if c, ok := s[r[0]]; ok {
+		if len(r) == 1 {
+			return true
+		}
+
+		return c.exists(r[1:])
+	}
+
+	s[r[0]] = make(seen)
+	if len(r) == 1 {
+		return false
+	}
+
+	// does not exists, but we want to populate the tree fully
+	return s[r[0]].exists(r[1:])
+}
diff --git a/engine/parser/lexer.go b/engine/parser/lexer.go
index fcf906d..0468fac 100644
--- a/engine/parser/lexer.go
+++ b/engine/parser/lexer.go
@@ -42,6 +42,7 @@ const (
 	TruncateToken
 	DropToken
 	GrantToken
+	DistinctToken
 
 	// Second order Token
 
@@ -141,6 +142,7 @@ func (l *lexer) lex(instruction []byte) ([]Token, error) {
 	matchers = append(matchers, l.MatchTruncateToken)
 	matchers = append(matchers, l.MatchDropToken)
 	matchers = append(matchers, l.MatchGrantToken)
+	matchers = append(matchers, l.MatchDistinctToken)
 	// Second order Matcher
 	matchers = append(matchers, l.MatchTableToken)
 	matchers = append(matchers, l.MatchFromToken)
@@ -329,6 +331,10 @@ func (l *lexer) MatchSelectToken() bool {
 	return l.Match([]byte("select"), SelectToken)
 }
 
+func (l *lexer) MatchDistinctToken() bool {
+	return l.Match([]byte("distinct"), DistinctToken)
+}
+
 func (l *lexer) MatchInsertToken() bool {
 	return l.Match([]byte("insert"), InsertToken)
 }
diff --git a/engine/parser/parser.go b/engine/parser/parser.go
index dd07d84..0bdcc84 100644
--- a/engine/parser/parser.go
+++ b/engine/parser/parser.go
@@ -464,7 +464,7 @@ func (p *parser) parseAttribute() (*Decl, error) {
 		}
 	}
 
-	// shoud be a StringToken here
+	// should be a StringToken here
 	// If there is a point after, it's a table name,
 	// if not, it's the attribute
 	if !p.is(StringToken, StarToken) {
diff --git a/engine/parser/select.go b/engine/parser/select.go
index fb485e9..1e59617 100644
--- a/engine/parser/select.go
+++ b/engine/parser/select.go
@@ -20,28 +20,65 @@ func (p *parser) parseSelect(tokens []Token) (*Instruction, error) {
 		return nil, fmt.Errorf("SELECT token must be followed by attributes to select")
 	}
 
+	var (
+		distinctDecl *Decl
+		distinctOpen bool
+	)
+	if p.is(DistinctToken) {
+		distinctDecl, err = p.consumeToken(DistinctToken)
+		if err != nil {
+			return nil, err
+		}
+
+		if p.is(OnToken) {
+			if err := p.next(); err != nil {
+				return nil, err
+			}
+			if !p.is(BracketOpeningToken) {
+				return nil, fmt.Errorf("Syntax error near %v, opening bracket expected\n", tokens[p.index])
+			}
+			if err := p.next(); err != nil {
+				return nil, err
+			}
+			distinctOpen = true
+		}
+
+		selectDecl.Add(distinctDecl)
+	}
+
 	for {
-		if p.is(CountToken) {
+		switch {
+		case p.is(CountToken):
 			attrDecl, err := p.parseBuiltinFunc()
 			if err != nil {
 				return nil, err
 			}
 			selectDecl.Add(attrDecl)
-		} else {
+		default:
 			attrDecl, err := p.parseAttribute()
 			if err != nil {
 				return nil, err
 			}
+			if distinctOpen {
+				distinctDecl.Add(attrDecl)
+			}
 			selectDecl.Add(attrDecl)
 		}
 
-		// If comma, loop again.
-		if p.is(CommaToken) {
+		switch {
+		case distinctOpen && p.is(BracketClosingToken):
+			if err := p.next(); err != nil {
+				return nil, err
+			}
+			distinctOpen = false
+			continue
+		case p.is(CommaToken):
 			if err := p.next(); err != nil {
 				return nil, err
 			}
 			continue
 		}
+
 		break
 	}
 
diff --git a/engine/select.go b/engine/select.go
index 63ac478..be1e1bd 100644
--- a/engine/select.go
+++ b/engine/select.go
@@ -123,6 +123,8 @@ func selectExecutor(e *Engine, selectDecl *parser.Decl, conn protocol.EngineConn
 				return fmt.Errorf("wrong offset value: %s", err)
 			}
 			conn = offsetedConn(conn, offset)
+		case parser.DistinctToken:
+			conn = distinctedConn(conn, len(selectDecl.Decl[i].Decl))
 		}
 	}
 
@@ -143,7 +145,7 @@ func selectExecutor(e *Engine, selectDecl *parser.Decl, conn protocol.EngineConn
 	}
 
 	if len(functors) == 0 {
-		// Instanciate a new select functor
+		// Instantiate a new select functor
 		functors, err = getSelectFunctors(selectDecl)
 		if err != nil {
 			return err
@@ -164,7 +166,7 @@ type selectFunctor interface {
 	Done() error
 }
 
-// getSelectFunctors instanciate new functors for COUNT, MAX, MIN, AVG, ... and default select functor that return rows to client
+// getSelectFunctors instantiate new functors for COUNT, MAX, MIN, AVG, ... and default select functor that return rows to client
 // If a functor is specified, no attribute can be selected ?
 func getSelectFunctors(attr *parser.Decl) ([]selectFunctor, error) {
 	var functors []selectFunctor


Test patch:
diff --git a/driver/driver_test.go b/driver/driver_test.go
index 0b0eecb..5fad4f4 100644
--- a/driver/driver_test.go
+++ b/driver/driver_test.go
@@ -305,7 +305,7 @@ func TestCompareDateGT(t *testing.T) {
 
 	query := "SELECT dat FROM comp WHERE dat > '2018-03-03'"
 
-	rows, err := db.Query(query, )
+	rows, err := db.Query(query)
 	if err != nil {
 		t.Fatalf("sql.Query: %s", err)
 	}
@@ -362,7 +362,7 @@ func TestCompareDateLT(t *testing.T) {
 
 	query := "SELECT dat FROM comp WHERE dat < '2019-03-03'"
 
-	rows, err := db.Query(query, )
+	rows, err := db.Query(query)
 	if err != nil {
 		t.Fatalf("sql.Query: %s", err)
 	}
@@ -552,7 +552,6 @@ func TestEqualAndDistinct(t *testing.T) {
 	if nb != 2 {
 		t.Fatalf("Expected 2 rows, got %d", nb)
 	}
-
 }
 
 func TestGreaterThanOrEqualAndLessThanOrEqual(t *testing.T) {
@@ -952,3 +951,62 @@ func TestJSON(t *testing.T) {
 		t.Fatalf("Unexpected values (second unmarshal): %+v\n", s)
 	}
 }
+
+func TestDistinct(t *testing.T) {
+	log.UseTestLogger(t)
+
+	batch := []string{
+		`CREATE TABLE user (name TEXT, surname TEXT, age INT);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Foo', 'Bar', 20);`,
+		`INSERT INTO user (name, surname, age) VALUES ('John', 'Doe', 20);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Jane', 'Doe', 40);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Joe', 'Doe', 10);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Homer', 'Simpson', 40);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Marge', 'Simpson', 40);`,
+		`INSERT INTO user (name, surname, age) VALUES ('Bruce', 'Wayne', 3333);`,
+	}
+
+	db, err := sql.Open("ramsql", "TestDistinct")
+	if err != nil {
+		t.Fatalf("sql.Open : Error : %s\n", err)
+	}
+	defer db.Close()
+
+	for _, b := range batch {
+		_, err = db.Exec(b)
+		if err != nil {
+			t.Fatalf("sql.Exec: Error: %s\n", err)
+		}
+	}
+
+	testDistinct := func(t *testing.T, query string, exp int, dest ...interface{}) {
+		rows, err := db.Query(query)
+		if err != nil {
+			t.Fatalf("sql.Query: %s", err)
+		}
+
+		var got int
+		for rows.Next() {
+			got++
+			if err := rows.Scan(dest...); err != nil {
+				t.Fatal(err)
+			}
+		}
+		if err := rows.Err(); err != nil {
+			t.Fatal(err)
+		}
+
+		if got != exp {
+			t.Fatalf("Expected %d rows, got %d", exp, got)
+		}
+	}
+
+	t.Run("distinct", func(t *testing.T) {
+		var surname string
+		testDistinct(t, `SELECT DISTINCT surname FROM user`, 4, &surname)
+	})
+	t.Run("distinct-on", func(t *testing.T) {
+		var name string
+		testDistinct(t, `SELECT DISTINCT ON (surname, age) name FROM user`, 6, &name)
+	})
+}