feat: add support for SELECT DISTINCT ON(…) (#194)

Author: bripkens

parser/ast.go

@@ -6739,6 +6739,7 @@ type SelectQuery struct {
 	With          *WithClause
 	Top           *TopClause
 	HasDistinct   bool
+	DistinctOn    *DistinctOn
 	SelectItems   []*SelectItem
 	From          *FromClause
 	ArrayJoin     *ArrayJoinClause
@@ -6782,6 +6783,11 @@ func (s *SelectQuery) String() string { // nolint: funlen
 	builder.WriteString("SELECT ")
 	if s.HasDistinct {
 		builder.WriteString("DISTINCT ")
+
+		if s.DistinctOn != nil {
+			builder.WriteString(s.DistinctOn.String())
+			builder.WriteString(" ")
+		}
 	}
 	if s.Top != nil {
 		builder.WriteString(s.Top.String())
@@ -6952,6 +6958,44 @@ func (s *SelectQuery) Accept(visitor ASTVisitor) error {
 	return visitor.VisitSelectQuery(s)
 }
 
+type DistinctOn struct {
+	Idents        []*Ident
+	DistinctOnPos Pos
+	DistinctOnEnd Pos
+}
+
+func (s *DistinctOn) Pos() Pos {
+	return s.DistinctOnPos
+}
+
+func (s *DistinctOn) End() Pos {
+	return s.DistinctOnEnd
+}
+
+func (s *DistinctOn) String() string {
+	var builder strings.Builder
+	builder.WriteString("ON (")
+	for i, ident := range s.Idents {
+		if i > 0 {
+			builder.WriteString(", ")
+		}
+		builder.WriteString(ident.String())
+	}
+	builder.WriteByte(')')
+	return builder.String()
+}
+
+func (s *DistinctOn) Accept(visitor ASTVisitor) error {
+	visitor.Enter(s)
+	defer visitor.Leave(s)
+	for _, ident := range s.Idents {
+		if err := ident.Accept(visitor); err != nil {
+			return err
+		}
+	}
+	return visitor.VisitDistinctOn(s)
+}
+
 type SubQuery struct {
 	HasParen bool
 	Select   *SelectQuery

parser/ast_visitor.go

@@ -184,6 +184,7 @@ type ASTVisitor interface {
 	VisitDescribeExpr(expr *DescribeStmt) error
 	VisitSelectItem(expr *SelectItem) error
 	VisitTargetPairExpr(expr *TargetPair) error
+	VisitDistinctOn(expr *DistinctOn) error
 
 	Enter(expr Expr)
 	Leave(expr Expr)
@@ -1476,6 +1477,13 @@ func (v *DefaultASTVisitor) VisitTargetPairExpr(expr *TargetPair) error {
 	return nil
 }
 
+func (v *DefaultASTVisitor) VisitDistinctOn(expr *DistinctOn) error {
+	if v.Visit != nil {
+		return v.Visit(expr)
+	}
+	return nil
+}
+
 func (v *DefaultASTVisitor) Enter(expr Expr) {}
 
 func (v *DefaultASTVisitor) Leave(expr Expr) {}

parser/parser_query.go

@@ -71,6 +71,49 @@ func (p *Parser) parseTopClause(pos Pos) (*TopClause, error) {
 	}, nil
 }
 
+func (p *Parser) tryParseDistinctOn(pos Pos) (*DistinctOn, error) {
+	if !p.matchKeyword(KeywordOn) {
+		return nil, nil
+	}
+	return p.parseDistinctOn(pos)
+}
+
+func (p *Parser) parseDistinctOn(pos Pos) (*DistinctOn, error) {
+	if err := p.expectKeyword(KeywordOn); err != nil {
+		return nil, err
+	}
+
+	if err := p.expectTokenKind(TokenKindLParen); err != nil {
+		return nil, err
+	}
+
+	ident, err := p.parseIdent()
+	if err != nil {
+		return nil, err
+	}
+	idents := []*Ident{ident}
+
+	for p.matchTokenKind(TokenKindComma) {
+		_ = p.lexer.consumeToken()
+
+		ident, err = p.parseIdent()
+		if err != nil {
+			return nil, err
+		}
+		idents = append(idents, ident)
+	}
+
+	if err := p.expectTokenKind(TokenKindRParen); err != nil {
+		return nil, err
+	}
+
+	return &DistinctOn{
+		Idents:        idents,
+		DistinctOnPos: pos,
+		DistinctOnEnd: p.Pos(),
+	}, nil
+}
+
 func (p *Parser) tryParseFromClause(pos Pos) (*FromClause, error) {
 	if !p.matchKeyword(KeywordFrom) {
 		return nil, nil
@@ -835,6 +878,10 @@ func (p *Parser) parseSelectStmt(pos Pos) (*SelectQuery, error) { // nolint: fun
 	}
 	// DISTINCT?
 	hasDistinct := p.tryConsumeKeywords(KeywordDistinct)
+	distinctOn, err := p.tryParseDistinctOn(p.Pos())
+	if err != nil {
+		return nil, err
+	}
 
 	top, err := p.tryParseTopClause(p.Pos())
 	if err != nil {
@@ -961,6 +1008,7 @@ func (p *Parser) parseSelectStmt(pos Pos) (*SelectQuery, error) { // nolint: fun
 		StatementEnd: statementEnd,
 		Top:          top,
 		HasDistinct:  hasDistinct,
+		DistinctOn:   distinctOn,
 		SelectItems:  selectItems,
 		From:         from,
 		ArrayJoin:    arrayJoin,

parser/testdata/basic/output/quantile_functions.sql.golden.json

@@ -5,6 +5,7 @@
     "With": null,
     "Top": null,
     "HasDistinct": false,
+    "DistinctOn": null,
     "SelectItems": [
       {
         "Expr": {

parser/testdata/ddl/output/bug_001.sql.golden.json

@@ -58,6 +58,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {

parser/testdata/ddl/output/create_live_view_basic.sql.golden.json

@@ -83,6 +83,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {

parser/testdata/ddl/output/create_materialized_view_basic.sql.golden.json

@@ -255,6 +255,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {

parser/testdata/ddl/output/create_materialized_view_with_definer.sql.golden.json

@@ -181,6 +181,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {

parser/testdata/ddl/output/create_materialized_view_with_empty_table_schema.sql.golden.json

@@ -152,6 +152,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {
@@ -247,6 +248,7 @@
                     "With": null,
                     "Top": null,
                     "HasDistinct": false,
+                    "DistinctOn": null,
                     "SelectItems": [
                       {
                         "Expr": {

parser/testdata/ddl/output/create_materialized_view_with_refresh.sql.golden.json

@@ -148,6 +148,7 @@
         "With": null,
         "Top": null,
         "HasDistinct": false,
+        "DistinctOn": null,
         "SelectItems": [
           {
             "Expr": {