Implement top operator

Fixes #26

parser/ast.go

@@ -136,6 +136,22 @@ func (op *TakeOperator) Span() Span {
 	return newSpan(op.Pipe.Start, op.RowCount.Span().End)
 }
 
+// TopOperator represents a `| top` operator in a [TabularExpr].
+// It implements [TabularOperator].
+type TopOperator struct {
+	Pipe     Span
+	Keyword  Span
+	RowCount Expr
+	By       Span
+	Col      *SortTerm
+}
+
+func (op *TopOperator) tabularOperator() {}
+
+func (op *TopOperator) Span() Span {
+	return newSpan(op.Pipe.Start, op.Col.Span().End)
+}
+
 // ProjectOperator represents a `| project` operator in a [TabularExpr].
 // It implements [TabularOperator].
 type ProjectOperator struct {

parser/parser.go

@@ -127,6 +127,12 @@ func (p *parser) tabularExpr() (*TabularExpr, error) {
 				expr.Operators = append(expr.Operators, op)
 			}
 			returnedError = joinErrors(returnedError, err)
+		case "top":
+			op, err := p.topOperator(pipeToken, operatorName)
+			if op != nil {
+				expr.Operators = append(expr.Operators, op)
+			}
+			returnedError = joinErrors(returnedError, err)
 		case "project":
 			op, err := p.projectOperator(pipeToken, operatorName)
 			if op != nil {
@@ -188,86 +194,89 @@ func (p *parser) sortOperator(pipe, keyword Token) (*SortOperator, error) {
 		Keyword: newSpan(keyword.Span.Start, by.Span.End),
 	}
 	for {
-		x, err := p.expr()
+		term, err := p.sortTerm()
+		if term != nil {
+			op.Terms = append(op.Terms, term)
+		}
 		if err != nil {
 			return op, makeErrorOpaque(err)
 		}
-		term := &SortTerm{
-			X:           x,
-			AscDescSpan: nullSpan(),
-			NullsSpan:   nullSpan(),
-		}
-		op.Terms = append(op.Terms, term)
 
-		// asc/desc
-		tok, ok := p.next()
-		if !ok {
-			return op, nil
-		}
-		switch tok.Kind {
-		case TokenComma:
-			continue
-		case TokenIdentifier:
-			switch tok.Value {
-			case "asc":
-				term.Asc = true
-				term.AscDescSpan = tok.Span
-				term.NullsFirst = true
-			case "desc":
-				term.Asc = false
-				term.AscDescSpan = tok.Span
-				term.NullsFirst = false
-			case "nulls":
-				// Good, but wait until next switch statement.
-				p.prev()
-			default:
-				p.prev()
-				return op, nil
-			}
-		default:
+		// Check for a comma to see if we should proceed.
+		if tok, _ := p.next(); tok.Kind != TokenComma {
 			p.prev()
 			return op, nil
 		}
+	}
+}
 
-		// nulls first/last
-		tok, ok = p.next()
-		if !ok {
-			return op, nil
-		}
-		switch {
-		case tok.Kind == TokenComma:
-			continue
-		case tok.Kind == TokenIdentifier && tok.Value == "nulls":
-			switch tok2, _ := p.next(); {
-			case tok2.Kind == TokenIdentifier && tok2.Value == "first":
-				term.NullsFirst = true
-				term.NullsSpan = newSpan(tok.Span.Start, tok2.Span.End)
-			case tok2.Kind == TokenIdentifier && tok2.Value == "last":
-				term.NullsFirst = false
-				term.NullsSpan = newSpan(tok.Span.Start, tok2.Span.End)
-			default:
-				p.prev()
-				return op, &parseError{
-					source: p.source,
-					span:   tok2.Span,
-					err:    fmt.Errorf("expected 'first' or 'last', got %s", formatToken(p.source, tok2)),
-				}
-			}
+func (p *parser) sortTerm() (*SortTerm, error) {
+	x, err := p.expr()
+	if err != nil {
+		return nil, err
+	}
+	term := &SortTerm{
+		X:           x,
+		AscDescSpan: nullSpan(),
+		NullsSpan:   nullSpan(),
+	}
+
+	// asc/desc
+	tok, ok := p.next()
+	if !ok {
+		return term, nil
+	}
+	switch tok.Kind {
+	case TokenIdentifier:
+		switch tok.Value {
+		case "asc":
+			term.Asc = true
+			term.AscDescSpan = tok.Span
+			term.NullsFirst = true
+		case "desc":
+			term.Asc = false
+			term.AscDescSpan = tok.Span
+			term.NullsFirst = false
+		case "nulls":
+			// Good, but wait until next switch statement.
+			p.prev()
 		default:
 			p.prev()
-			return op, nil
+			return term, nil
 		}
+	default:
+		p.prev()
+		return term, nil
+	}
 
-		// Check for a comma to see if we should proceed.
-		tok, ok = p.next()
-		if !ok {
-			return op, nil
-		}
-		if tok.Kind != TokenComma {
+	// nulls first/last
+	tok, ok = p.next()
+	if !ok {
+		return term, nil
+	}
+	switch {
+	case tok.Kind == TokenIdentifier && tok.Value == "nulls":
+		switch tok2, _ := p.next(); {
+		case tok2.Kind == TokenIdentifier && tok2.Value == "first":
+			term.NullsFirst = true
+			term.NullsSpan = newSpan(tok.Span.Start, tok2.Span.End)
+		case tok2.Kind == TokenIdentifier && tok2.Value == "last":
+			term.NullsFirst = false
+			term.NullsSpan = newSpan(tok.Span.Start, tok2.Span.End)
+		default:
 			p.prev()
-			return op, nil
+			return term, &parseError{
+				source: p.source,
+				span:   tok2.Span,
+				err:    fmt.Errorf("expected 'first' or 'last', got %s", formatToken(p.source, tok2)),
+			}
 		}
+	default:
+		p.prev()
+		return term, nil
 	}
+
+	return term, nil
 }
 
 func (p *parser) takeOperator(pipe, keyword Token) (*TakeOperator, error) {
@@ -300,6 +309,52 @@ func (p *parser) takeOperator(pipe, keyword Token) (*TakeOperator, error) {
 	return op, nil
 }
 
+func (p *parser) topOperator(pipe, keyword Token) (*TopOperator, error) {
+	op := &TopOperator{
+		Pipe:    pipe.Span,
+		Keyword: keyword.Span,
+		By:      nullSpan(),
+	}
+
+	tok, _ := p.next()
+	if tok.Kind != TokenNumber {
+		p.prev()
+		return op, &parseError{
+			source: p.source,
+			span:   tok.Span,
+			err:    fmt.Errorf("expected integer, got %s", formatToken(p.source, tok)),
+		}
+	}
+	rowCount := &BasicLit{
+		Kind:      tok.Kind,
+		Value:     tok.Value,
+		ValueSpan: tok.Span,
+	}
+	op.RowCount = rowCount
+	if !rowCount.IsInteger() {
+		return op, &parseError{
+			source: p.source,
+			span:   tok.Span,
+			err:    fmt.Errorf("expected integer, got %s", formatToken(p.source, tok)),
+		}
+	}
+
+	tok, _ = p.next()
+	if tok.Kind != TokenBy {
+		p.prev()
+		return op, &parseError{
+			source: p.source,
+			span:   tok.Span,
+			err:    fmt.Errorf("expected 'by', got %s", formatToken(p.source, tok)),
+		}
+	}
+	op.By = tok.Span
+
+	var err error
+	op.Col, err = p.sortTerm()
+	return op, makeErrorOpaque(err)
+}
+
 func (p *parser) projectOperator(pipe, keyword Token) (*ProjectOperator, error) {
 	op := &ProjectOperator{
 		Pipe:    pipe.Span,

parser/parser_test.go

@@ -1079,6 +1079,38 @@ func TestParse(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:  "Top",
+			query: "StormEvents | top 3 by InjuriesDirect",
+			want: &TabularExpr{
+				Source: &TableRef{
+					Table: &Ident{
+						Name:     "StormEvents",
+						NameSpan: newSpan(0, 11),
+					},
+				},
+				Operators: []TabularOperator{
+					&TopOperator{
+						Pipe:    newSpan(12, 13),
+						Keyword: newSpan(14, 17),
+						RowCount: &BasicLit{
+							Kind:      TokenNumber,
+							Value:     "3",
+							ValueSpan: newSpan(18, 19),
+						},
+						By: newSpan(20, 22),
+						Col: &SortTerm{
+							X: &Ident{
+								Name:     "InjuriesDirect",
+								NameSpan: newSpan(23, 37),
+							},
+							AscDescSpan: nullSpan(),
+							NullsSpan:   nullSpan(),
+						},
+					},
+				},
+			},
+		},
 	}
 
 	equateInvalidSpans := cmp.FilterValues(func(span1, span2 Span) bool {

pql.go

@@ -79,6 +79,21 @@ func splitQueries(expr *parser.TabularExpr) ([]*subquery, error) {
 				}
 				subqueries = append(subqueries, lastSubquery)
 			}
+		case *parser.TopOperator:
+			if lastSubquery == nil || !canAttachSort(lastSubquery.op) || lastSubquery.sort != nil || lastSubquery.take != nil {
+				lastSubquery = new(subquery)
+				subqueries = append(subqueries, lastSubquery)
+			}
+			lastSubquery.sort = &parser.SortOperator{
+				Pipe:    op.Pipe,
+				Keyword: op.Keyword,
+				Terms:   []*parser.SortTerm{op.Col},
+			}
+			lastSubquery.take = &parser.TakeOperator{
+				Pipe:     op.Pipe,
+				Keyword:  op.Keyword,
+				RowCount: op.RowCount,
+			}
 		default:
 			lastSubquery = &subquery{
 				op: op,

testdata/Goldens/Top/input.pql

@@ -0,0 +1,2 @@
+SourceFiles
+| top 3 by LineCount

testdata/Goldens/Top/output.csv

@@ -0,0 +1,3 @@
+parser,parser_test.go,1108
+parser,parser.go,878
+parser,lex.go,681

testdata/Goldens/Top/output.sql

@@ -0,0 +1 @@
+SELECT * FROM "SourceFiles" ORDER BY "LineCount" DESC NULLS LAST LIMIT 3;