diff --git a/executor.go b/executor.go index 8cd24f33e..69d96928f 100644 --- a/executor.go +++ b/executor.go @@ -5939,7 +5939,7 @@ func (e *executor) mapperLocal(ctx context.Context, shards []uint64, mapFn mapFu ch := make(chan mapResponse, len(shards)) expected := 0 - shardLoop: +shardLoop: for _, shard := range shards { j := job{ shard: shard, diff --git a/planner.go b/planner.go index b8482a6af..e1eb6471d 100644 --- a/planner.go +++ b/planner.go @@ -19,6 +19,7 @@ import ( "database/sql" "fmt" "io" + "strconv" "strings" "github.com/molecula/featurebase/v2/pql" @@ -73,6 +74,12 @@ func (p *Planner) planAggregateSelectStatement(ctx context.Context, stmt *sql2.S return nil, fmt.Errorf("unexpected source type in aggregate query: %T", source) } + // Convert WHERE clause. + cond, err := p.planExprPQL(ctx, stmt, stmt.WhereExpr) + if err != nil { + return nil, err + } + // TODO: Support multiple aggregate calls. if len(stmt.Columns) > 1 { return nil, fmt.Errorf("only one call allowed in aggregate query") @@ -91,7 +98,7 @@ func (p *Planner) planAggregateSelectStatement(ctx context.Context, stmt *sql2.S callName := strings.ToUpper(sql2.IdentName(call.Name)) switch callName { case "COUNT": - return NewCountNode(p.executor, sql2.IdentName(source.Name)), nil + return NewCountNode(p.executor, sql2.IdentName(source.Name), cond), nil default: return nil, fmt.Errorf("unsupported call in aggregate query: %T", callName) } @@ -103,6 +110,183 @@ func (p *Planner) planNonAggregateSelectStatement(ctx context.Context, stmt *sql return nil, fmt.Errorf("cannot plan non-aggregate SELECT query") } +// planExprPQL returns a PQL call tree for a given expression. +func (p *Planner) planExprPQL(ctx context.Context, stmt *sql2.SelectStatement, expr sql2.Expr) (_ *pql.Call, err error) { + if expr == nil { + return nil, nil + } + + switch expr := expr.(type) { + case *sql2.BinaryExpr: + return p.planBinaryExprPQL(ctx, stmt, expr) + case *sql2.BindExpr: + return nil, fmt.Errorf("bind expressions are not supported") + case *sql2.BlobLit: + return nil, fmt.Errorf("blob literals are not supported") + case *sql2.BoolLit: + return nil, fmt.Errorf("boolean literals are not supported") + case *sql2.Call: + return nil, fmt.Errorf("call expressions are not supported") + case *sql2.CaseExpr: + return nil, fmt.Errorf("case expressions are not supported") + case *sql2.CastExpr: + return nil, fmt.Errorf("cast expressions are not supported") + case *sql2.Exists: + return nil, fmt.Errorf("exists expressions are not supported") + case *sql2.ExprList: + return nil, fmt.Errorf("expression lists are not supported") + case *sql2.Ident: + return nil, fmt.Errorf("identifiers are not supported") + case *sql2.NullLit: + return nil, fmt.Errorf("NULL expressions are not supported") + case *sql2.NumberLit: + return nil, fmt.Errorf("number expressions are not supported") + case *sql2.ParenExpr: + return p.planExprPQL(ctx, stmt, expr.X) + case *sql2.QualifiedRef: + return nil, fmt.Errorf("qualified references are not supported") + case *sql2.Raise: + return nil, fmt.Errorf("raise expressions are not supported") + case *sql2.Range: + return nil, fmt.Errorf("range expressions are not supported") + case *sql2.StringLit: + return nil, fmt.Errorf("string literals are not supported") + case *sql2.UnaryExpr: + return nil, fmt.Errorf("unary expressions are not supported") + default: + return nil, fmt.Errorf("unexpected SQL expression type: %T", expr) + } +} + +func (p *Planner) planBinaryExprPQL(ctx context.Context, stmt *sql2.SelectStatement, expr *sql2.BinaryExpr) (_ *pql.Call, err error) { + switch op := expr.Op; op { + case sql2.AND, sql2.OR: + name := "Intersect" + if op == sql2.OR { + name = "Union" + } + + x, err := p.planExprPQL(ctx, stmt, expr.X) + if err != nil { + return nil, err + } + y, err := p.planExprPQL(ctx, stmt, expr.Y) + if err != nil { + return nil, err + } + + return &pql.Call{ + Name: name, + Children: []*pql.Call{x, y}, + }, nil + + case sql2.EQ, sql2.NE, sql2.LT, sql2.LE, sql2.GT, sql2.GE: + // Ensure field reference exists in binary expression. + x, y := expr.X, expr.Y + xIdent, xOk := x.(*sql2.Ident) + yIdent, yOk := y.(*sql2.Ident) + if xOk && yOk { + return nil, fmt.Errorf("cannot compare fields in a WHERE clause") + } else if !xOk && !yOk { + return nil, fmt.Errorf("expression must reference one field") + } + + // Rewrite expression so field ref is LHS. + if !xOk && yOk { + xIdent, y = yIdent, x + switch op { + case sql2.LT: + op = sql2.GT + case sql2.LE: + op = sql2.GE + case sql2.GT: + op = sql2.LT + case sql2.GE: + op = sql2.LE + } + } + + pqlValue, err := sqlToPQLValue(y) + if err != nil { + return nil, err + } + + isBSI := true // TODO: Check field if it is a BSI field. + if !isBSI { + return &pql.Call{ + Name: "Row", + Args: map[string]interface{}{ + sql2.IdentName(xIdent): pqlValue, + }, + }, nil + } + + pqlOp, err := sqlToPQLOp(op) + if err != nil { + return nil, err + } + return &pql.Call{ + Name: "Row", + Args: map[string]interface{}{ + sql2.IdentName(xIdent): &pql.Condition{ + Op: pqlOp, + Value: pqlValue, + }, + }, + }, nil + + case sql2.BITAND, sql2.BITOR, sql2.BITNOT, sql2.LSHIFT, sql2.RSHIFT: + return nil, fmt.Errorf("bitwise operators are not supported in WHERE clause") + case sql2.PLUS, sql2.MINUS, sql2.STAR, sql2.SLASH, sql2.REM: // + + return nil, fmt.Errorf("arithmetic operators are not supported in WHERE clause") + case sql2.CONCAT: + return nil, fmt.Errorf("concatenation operator is not supported in WHERE clause") + case sql2.IN, sql2.NOTIN: + return nil, fmt.Errorf("IN operator is not supported") + case sql2.BETWEEN, sql2.NOTBETWEEN: + return nil, fmt.Errorf("BETWEEN operator is not supported") + default: + return nil, fmt.Errorf("unexpected binary expression operator: %s", expr.Op) + } +} + +// sqlToPQLOp converts a SQL2 operation token to PQL. +func sqlToPQLOp(op sql2.Token) (pql.Token, error) { + switch op { + case sql2.EQ: + return pql.EQ, nil + case sql2.NE: + return pql.NEQ, nil + case sql2.LT: + return pql.LT, nil + case sql2.LE: + return pql.LTE, nil + case sql2.GT: + return pql.GT, nil + case sql2.GE: + return pql.GTE, nil + default: + return pql.ILLEGAL, fmt.Errorf("cannot convert SQL op %q to PQL", op) + } +} + +// sqlToPQLValue converts a literal SQL2 expression node to a PQL Go value. +func sqlToPQLValue(expr sql2.Expr) (interface{}, error) { + switch expr := expr.(type) { + case *sql2.StringLit: + return expr.Value, nil + case *sql2.NumberLit: + if expr.IsFloat() { + return strconv.ParseFloat(expr.Value, 64) + } + return strconv.ParseInt(expr.Value, 10, 64) + case *sql2.BoolLit: + return expr.Value, nil + default: + return nil, fmt.Errorf("cannot convert SQL expression %T to a literal value", expr) + } +} + type Stmt struct { node StmtNode } @@ -266,14 +450,19 @@ var _ StmtNode = (*CountNode)(nil) type CountNode struct { executor *executor indexName string + cond *pql.Call // conditional row []interface{} } -func NewCountNode(executor *executor, indexName string) *CountNode { +func NewCountNode(executor *executor, indexName string, cond *pql.Call) *CountNode { + if cond == nil { + cond = &pql.Call{Name: "All"} + } return &CountNode{ executor: executor, indexName: indexName, + cond: cond, } } @@ -286,11 +475,14 @@ func (n *CountNode) Next(ctx context.Context) error { if n.row != nil { return io.EOF } - result, err := n.executor.Execute(ctx, n.indexName, &pql.Query{ + + q := &pql.Query{ Calls: []*pql.Call{ - {Name: "Count", Children: []*pql.Call{{Name: "All"}}}, + {Name: "Count", Children: []*pql.Call{n.cond}}, }, - }, nil, nil) + } + + result, err := n.executor.Execute(ctx, n.indexName, q, nil, nil) if err != nil { return err } diff --git a/planner_test.go b/planner_test.go index e3ce18aa7..3d451a343 100644 --- a/planner_test.go +++ b/planner_test.go @@ -32,7 +32,9 @@ func TestPlanner_Count(t *testing.T) { } defer index.Close() - if _, err := index.CreateField("f"); err != nil { + if _, err := index.CreateField("f", pilosa.OptFieldTypeInt(0, 1000)); err != nil { + t.Fatal(err) + } else if _, err := index.CreateField("x", pilosa.OptFieldTypeInt(0, 1000)); err != nil { t.Fatal(err) } @@ -43,27 +45,124 @@ func TestPlanner_Count(t *testing.T) { Set(1, f=10) Set(2, f=10) Set(3, f=11) + Set(4, f=12) + Set(5, f=12) + Set(6, f=13) + + Set(1, x=100) + Set(2, x=200) `}); err != nil { t.Fatal(err) } - // Parse SQL into AST. - q := `SELECT COUNT(*) AS "count" FROM i` - stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) - if err != nil { - t.Fatal(err) - } - defer stmt.Close() + t.Run("ALL", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() - // Scan first row from result set. - var n int - if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { - t.Fatal(err) - } else if got, want := n, 3; got != want { - t.Fatalf("Scan()=%d, want %d", got, want) - } + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 6; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) - if err := stmt.Close(); err != nil { - t.Fatal(err) - } + t.Run("WHERE", func(t *testing.T) { + t.Run("EQ", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f = 10` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 2; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + t.Run("NE", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f != 10` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 4; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + t.Run("LT", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f < 12` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 3; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + t.Run("GT", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f > 12` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 1; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + + t.Run("AND", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f = 10 AND x = 100` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 1; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + + t.Run("OR", func(t *testing.T) { + q := `SELECT COUNT(*) AS "count" FROM i WHERE f = 10 OR x = 200 OR f = 12` + stmt, err := c.GetNode(0).Server.PlanSQL(context.Background(), q) + if err != nil { + t.Fatal(err) + } + defer stmt.Close() + + var n int + if err := stmt.QueryRowContext(context.Background()).Scan(&n); err != nil { + t.Fatal(err) + } else if got, want := n, 4; got != want { + t.Fatalf("Scan()=%d, want %d", got, want) + } + }) + }) } diff --git a/sql2/ast.go b/sql2/ast.go index f94c4a1c6..83ecf00a2 100644 --- a/sql2/ast.go +++ b/sql2/ast.go @@ -1330,6 +1330,13 @@ type NumberLit struct { // IsAggregate returns false. func (expr *NumberLit) IsAggregate() bool { return false } +// IsFloat returns true if literal contains a dot or 'e'. +func (expr *NumberLit) IsFloat() bool { + return strings.Contains(expr.Value, ".") || + strings.Contains(expr.Value, "e") || + strings.Contains(expr.Value, "E") +} + // Clone returns a deep copy of lit. func (lit *NumberLit) Clone() *NumberLit { if lit == nil {