Add in operator

Fixes #21

parser/ast.go

@@ -260,6 +260,30 @@ func (expr *UnaryExpr) Span() Span {
 
 func (expr *UnaryExpr) expression() {}
 
+// An InExpr represents an "in" operator expression.
+type InExpr struct {
+	X      Expr
+	In     Span
+	Lparen Span
+	Vals   []Expr
+	Rparen Span
+}
+
+func (expr *InExpr) Span() Span {
+	switch {
+	case expr.Rparen.IsValid():
+		return newSpan(expr.X.Span().Start, expr.Rparen.End)
+	case len(expr.Vals) > 0:
+		return newSpan(expr.X.Span().Start, expr.Vals[len(expr.Vals)-1].Span().End)
+	case expr.Lparen.IsValid():
+		return newSpan(expr.X.Span().Start, expr.Lparen.End)
+	default:
+		return newSpan(expr.X.Span().Start, expr.In.End)
+	}
+}
+
+func (expr *InExpr) expression() {}
+
 // A ParenExpr represents a parenthized expression.
 type ParenExpr struct {
 	Lparen Span

parser/lex.go

@@ -97,6 +97,9 @@ const (
 	// The Value will be the empty string.
 	TokenRParen
 
+	// TokenBy is the keyword "in".
+	// The Value will be the empty string.
+	TokenIn
 	// TokenBy is the keyword "by".
 	// The Value will be the empty string.
 	TokenBy
@@ -299,8 +302,9 @@ func Scan(query string) []Token {
 
 var keywords = map[string]TokenKind{
 	"and": TokenAnd,
-	"or":  TokenOr,
 	"by":  TokenBy,
+	"in":  TokenIn,
+	"or":  TokenOr,
 }
 
 func (s *scanner) ident() Token {

parser/parser.go

@@ -555,6 +555,55 @@ func (p *parser) exprBinaryTrail(x Expr, minPrecedence int) (Expr, error) {
 			p.prev()
 			return x, finalError
 		}
+
+		if op1.Kind == TokenIn {
+			lparen, _ := p.next()
+			if lparen.Kind != TokenLParen {
+				x = &InExpr{
+					X:      x,
+					In:     op1.Span,
+					Lparen: nullSpan(),
+					Rparen: nullSpan(),
+				}
+				finalError = joinErrors(finalError, &parseError{
+					source: p.source,
+					span:   lparen.Span,
+					err:    fmt.Errorf("expected '(', got %s", formatToken(p.source, lparen)),
+				})
+				return x, finalError
+			}
+			vals, err := p.exprList()
+			if err != nil {
+				finalError = joinErrors(finalError, makeErrorOpaque(err))
+				p.skipTo(TokenRParen)
+			}
+			rparen, _ := p.next()
+			if rparen.Kind != TokenRParen {
+				x = &InExpr{
+					X:      x,
+					In:     op1.Span,
+					Lparen: lparen.Span,
+					Vals:   vals,
+					Rparen: nullSpan(),
+				}
+				finalError = joinErrors(finalError, &parseError{
+					source: p.source,
+					span:   lparen.Span,
+					err:    fmt.Errorf("expected ')', got %s", formatToken(p.source, rparen)),
+				})
+				return x, finalError
+			}
+
+			x = &InExpr{
+				X:      x,
+				In:     op1.Span,
+				Lparen: lparen.Span,
+				Vals:   vals,
+				Rparen: rparen.Span,
+			}
+			continue
+		}
+
 		y, err := p.unaryExpr()
 		if err != nil {
 			finalError = joinErrors(finalError, makeErrorOpaque(err))
@@ -595,7 +644,7 @@ func operatorPrecedence(op TokenKind) int {
 	case TokenPlus, TokenMinus:
 		return 3
 	case TokenEq, TokenNE, TokenLT, TokenLE, TokenGT, TokenGE,
-		TokenCaseInsensitiveEq, TokenCaseInsensitiveNE:
+		TokenCaseInsensitiveEq, TokenCaseInsensitiveNE, TokenIn:
 		return 2
 	case TokenAnd:
 		return 1

parser/parser_test.go

@@ -522,6 +522,157 @@ func TestParse(t *testing.T) {
 				},
 			},
 		},
+		{
+			name:  "In",
+			query: `StormEvents | where State in ("GEORGIA", "MISSISSIPPI")`,
+			want: &TabularExpr{
+				Source: &TableRef{
+					Table: &Ident{
+						Name:     "StormEvents",
+						NameSpan: newSpan(0, 11),
+					},
+				},
+				Operators: []TabularOperator{
+					&WhereOperator{
+						Pipe:    newSpan(12, 13),
+						Keyword: newSpan(14, 19),
+						Predicate: &InExpr{
+							X: &Ident{
+								Name:     "State",
+								NameSpan: newSpan(20, 25),
+							},
+							In:     newSpan(26, 28),
+							Lparen: newSpan(29, 30),
+							Vals: []Expr{
+								&BasicLit{
+									Kind:      TokenString,
+									ValueSpan: newSpan(30, 39),
+									Value:     "GEORGIA",
+								},
+								&BasicLit{
+									Kind:      TokenString,
+									ValueSpan: newSpan(41, 54),
+									Value:     "MISSISSIPPI",
+								},
+							},
+							Rparen: newSpan(54, 55),
+						},
+					},
+				},
+			},
+		},
+		{
+			name:  "InAnd",
+			query: `StormEvents | where State in ("GEORGIA", "MISSISSIPPI") and DamageProperty > 10000`,
+			want: &TabularExpr{
+				Source: &TableRef{
+					Table: &Ident{
+						Name:     "StormEvents",
+						NameSpan: newSpan(0, 11),
+					},
+				},
+				Operators: []TabularOperator{
+					&WhereOperator{
+						Pipe:    newSpan(12, 13),
+						Keyword: newSpan(14, 19),
+						Predicate: &BinaryExpr{
+							X: &InExpr{
+								X: &Ident{
+									Name:     "State",
+									NameSpan: newSpan(20, 25),
+								},
+								In:     newSpan(26, 28),
+								Lparen: newSpan(29, 30),
+								Vals: []Expr{
+									&BasicLit{
+										Kind:      TokenString,
+										ValueSpan: newSpan(30, 39),
+										Value:     "GEORGIA",
+									},
+									&BasicLit{
+										Kind:      TokenString,
+										ValueSpan: newSpan(41, 54),
+										Value:     "MISSISSIPPI",
+									},
+								},
+								Rparen: newSpan(54, 55),
+							},
+							Op:     TokenAnd,
+							OpSpan: newSpan(56, 59),
+							Y: &BinaryExpr{
+								X: &Ident{
+									Name:     "DamageProperty",
+									NameSpan: newSpan(60, 74),
+								},
+								Op:     TokenGT,
+								OpSpan: newSpan(75, 76),
+								Y: &BasicLit{
+									Kind:      TokenNumber,
+									Value:     "10000",
+									ValueSpan: newSpan(77, 82),
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name:  "InAndFlipped",
+			query: `StormEvents | where DamageProperty > 10000 and State in ("GEORGIA", "MISSISSIPPI")`,
+			want: &TabularExpr{
+				Source: &TableRef{
+					Table: &Ident{
+						Name:     "StormEvents",
+						NameSpan: newSpan(0, 11),
+					},
+				},
+				Operators: []TabularOperator{
+					&WhereOperator{
+						Pipe:    newSpan(12, 13),
+						Keyword: newSpan(14, 19),
+						Predicate: &BinaryExpr{
+							X: &BinaryExpr{
+								X: &Ident{
+									Name:     "DamageProperty",
+									NameSpan: newSpan(20, 34),
+								},
+								Op:     TokenGT,
+								OpSpan: newSpan(35, 36),
+								Y: &BasicLit{
+									Kind:      TokenNumber,
+									Value:     "10000",
+									ValueSpan: newSpan(37, 42),
+								},
+							},
+							Op:     TokenAnd,
+							OpSpan: newSpan(43, 46),
+							Y: &InExpr{
+								X: &Ident{
+									Name:     "State",
+									NameSpan: newSpan(47, 52),
+								},
+								In:     newSpan(53, 55),
+								Lparen: newSpan(56, 57),
+								Vals: []Expr{
+									&BasicLit{
+										Kind:      TokenString,
+										ValueSpan: newSpan(57, 66),
+										Value:     "GEORGIA",
+									},
+									&BasicLit{
+										Kind:      TokenString,
+										ValueSpan: newSpan(68, 81),
+										Value:     "MISSISSIPPI",
+									},
+								},
+								Rparen: newSpan(81, 82),
+							},
+						},
+					},
+				},
+			},
+		},
 		{
 			name:  "BadArgument",
 			query: "foo | where strcat('a', .bork, 'x', 'y')",

pql.go

@@ -401,6 +401,20 @@ func writeExpression(sb *strings.Builder, source string, x parser.Expr) error {
 				fmt.Fprintf(sb, "NULL /* unhandled %s binary op */ ", x.Op)
 			}
 		}
+	case *parser.InExpr:
+		if err := writeExpressionMaybeParen(sb, source, x.X); err != nil {
+			return err
+		}
+		sb.WriteString(" IN (")
+		for i, y := range x.Vals {
+			if i > 0 {
+				sb.WriteString(", ")
+			}
+			if err := writeExpressionMaybeParen(sb, source, y); err != nil {
+				return err
+			}
+		}
+		sb.WriteString(")")
 	case *parser.CallExpr:
 		if f := initKnownFunctions()[x.Func.Name]; f != nil {
 			if err := f.write(sb, source, x); err != nil {

testdata/Goldens/In/input.pql

@@ -0,0 +1,2 @@
+StormEvents
+| where State in ("GEORGIA", "MISSISSIPPI")

testdata/Goldens/In/output.csv

@@ -0,0 +1,2 @@
+11503,GEORGIA,Thunderstorm Wind,2000
+13913,MISSISSIPPI,Thunderstorm Wind,20000

testdata/Goldens/In/output.sql

@@ -0,0 +1 @@
+SELECT * FROM "StormEvents" WHERE "State" IN ('GEORGIA', 'MISSISSIPPI');

testdata/Goldens/InAnd/input.pql

@@ -0,0 +1,2 @@
+StormEvents
+| where State in ("GEORGIA", "MISSISSIPPI") and DamageProperty > 10000

testdata/Goldens/InAnd/output.csv

@@ -0,0 +1 @@
+13913,MISSISSIPPI,Thunderstorm Wind,20000

testdata/Goldens/InAnd/output.sql

@@ -0,0 +1 @@
+SELECT * FROM "StormEvents" WHERE ("State" IN ('GEORGIA', 'MISSISSIPPI')) AND ("DamageProperty" > 10000);