Skip to content

Commit

Permalink
use new AST
Browse files Browse the repository at this point in the history
  • Loading branch information
Yaiba committed Jun 24, 2024
1 parent a723dda commit 3bed38b
Show file tree
Hide file tree
Showing 14 changed files with 498 additions and 229 deletions.
75 changes: 46 additions & 29 deletions internal/engine/cost/costmodel/rel_expr_test.go
Original file line number Diff line number Diff line change
@@ -1,11 +1,14 @@
package costmodel

import (
"testing"

"github.com/stretchr/testify/assert"

"github.com/kwilteam/kwil-db/core/types"
"github.com/kwilteam/kwil-db/internal/engine/cost/internal/testkit"
"github.com/kwilteam/kwil-db/internal/engine/cost/query_planner"
sqlparser "github.com/kwilteam/kwil-db/parse/sql"
"github.com/stretchr/testify/assert"
"testing"
"github.com/kwilteam/kwil-db/parse"
)

func Test_RelExpr_String(t *testing.T) {
Expand All @@ -17,7 +20,7 @@ func Test_RelExpr_String(t *testing.T) {
{
name: "test",
r: &RelExpr{},
want: "test\n\n stat: &{0 []}\n cost: 0",
want: "Unknown LogicalPlan type <nil>, Stat: (<nil>), Cost: 0",
},
}
for _, tt := range tests {
Expand Down Expand Up @@ -67,8 +70,9 @@ func Test_NewRelExpr(t *testing.T) {
{
name: "select wildcard",
sql: "SELECT * FROM users",
wt: "Projection: users.id, users.username, users.age, users.state, users.wallet, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.id, users.username, users.age, users.state, users.wallet, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
//{ // TODO?
// name: "select wildcard, deduplication",
Expand All @@ -79,34 +83,39 @@ func Test_NewRelExpr(t *testing.T) {
{
name: "select columns",
sql: "select username, age from users",
wt: "Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select column with alias",
sql: "select username as name from users",
wt: "Projection: users.username AS name, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username AS name, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select column expression",
sql: "select username, age+10 from users",
wt: "Projection: users.username, users.age + 10, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age + 10, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select with where",
sql: "select username, age from users where age > 20",
wt: "Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Filter: users.age > 20, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Filter: users.age > 20, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select with multiple where",
sql: "select username, age from users where age > 20 and state = 'CA'",
wt: "Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Filter: users.age > 20 AND users.state = 'CA', Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Filter: users.age > 20 AND users.state = 'CA', Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
//{
// name: "select with group by",
Expand All @@ -117,48 +126,56 @@ func Test_NewRelExpr(t *testing.T) {
name: "select with limit, without offset",
sql: "select username, age from users limit 10",
wt: "Limit: skip=0, fetch=10, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
" Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select with limit and offset",
sql: "select username, age from users limit 10 offset 5",
wt: "Limit: skip=5, fetch=10, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
" Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select with order by default",
sql: "select username, age from users order by age",
wt: "Sort: age ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
wt: "Sort: age ASC NULLS LAST, id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
{
name: "select with order by desc",
sql: "select username, age from users order by age desc",
wt: "Sort: age DESC NULLS FIRST, Stat: (RowCount: 0), Cost: 0\n" +
wt: "Sort: age DESC NULLS LAST, id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
/////////////////////// subquery
{
name: "select with subquery",
sql: "select username, age from (select * from users) as u",
wt: "Projection: users.username, users.age, Stat: (RowCount: 5), Cost: 0\n" +
" Projection: users.id, users.username, users.age, users.state, users.wallet, Stat: (RowCount: 5), Cost: 0\n" +
" Scan: users, Stat: (RowCount: 5), Cost: 0\n",
wt: "Sort: id ASC NULLS LAST, username ASC NULLS LAST, age ASC NULLS LAST, state ASC NULLS LAST, wallet ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.username, users.age, Stat: (RowCount: 0), Cost: 0\n" +
" Sort: id ASC NULLS LAST, Stat: (RowCount: 0), Cost: 0\n" +
" Projection: users.id, users.username, users.age, users.state, users.wallet, Stat: (RowCount: 5), Cost: 0\n Scan: users, Stat: (RowCount: 5), Cost: 0\n",
},
/////////////////////// two relations

}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
stmt, err := sqlparser.Parse(tt.sql)
pr, err := parse.ParseSQL(tt.sql, &types.Schema{
Name: "",
Tables: []*types.Table{testkit.MockUsersSchemaTable},
})

assert.NoError(t, err)
assert.NoError(t, pr.ParseErrs.Err())

q := query_planner.NewPlanner(cat)
plan := q.ToPlan(stmt)
plan := q.ToPlan(pr.AST)
rel := BuildRelExpr(plan)
assert.Equal(t, tt.wt, Format(rel, 0))
})
Expand Down
28 changes: 24 additions & 4 deletions internal/engine/cost/demo/demo_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,25 +3,45 @@ package demo
import (
"context"
"fmt"
"github.com/kwilteam/kwil-db/core/types"
"github.com/kwilteam/kwil-db/internal/engine/cost/internal/testkit"

"github.com/kwilteam/kwil-db/internal/engine/cost/query_planner"
sqlparser "github.com/kwilteam/kwil-db/parse/sql"
"github.com/kwilteam/kwil-db/parse"
)

// ExampleDemo demonstrates how a SQL is parsed, planned and executed, served as
// a simple mental model to understand the flow of the engine.
//
// For cost estimation, execution will be carried out by internal/engine, not by
// the virtual plan in this pkg.
func ExampleDemo() {
// NOTE!!!: this will fail, as 'parser' add sorting to every AST,
// but virutal plan does not support sorting yet.
//
// See Example_ExecutionContext_execute instead, although it only shows how
// a logical plan is executed.

// enter engine
rawSql := "SELECT state, username FROM users WHERE age = 20"
stmt, err := sqlparser.Parse(rawSql)

pr, err := parse.ParseSQL(rawSql, &types.Schema{
Name: "",
Tables: []*types.Table{testkit.MockUsersSchemaTable},
})

if err != nil {
panic(err)
}

if pr.ParseErrs.Err() != nil {
panic(pr.ParseErrs.Err())
}

// load into engine
catalog := testkit.InitMockCatalog()

planner := query_planner.NewPlanner(catalog)
plan := planner.ToPlan(stmt)
plan := planner.ToPlan(pr.AST)

//opt := optimizer.NewOptimizer()
//plan := opt.Optimize(plan)
Expand Down
1 change: 0 additions & 1 deletion internal/engine/cost/demo/execution.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package demo

import (
"context"

"github.com/kwilteam/kwil-db/internal/engine/cost/datasource"
dt "github.com/kwilteam/kwil-db/internal/engine/cost/datatypes"
"github.com/kwilteam/kwil-db/internal/engine/cost/logical_plan"
Expand Down
2 changes: 2 additions & 0 deletions internal/engine/cost/demo/execution_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@ import (

var stubTable = &dt.TableRef{Table: "users"}

// Example_ExecutionContext_execute demonstrates how to use the ExecutionContext
// to execute a logical plan.
func Example_ExecutionContext_execute() {
ctx := NewExecutionContext()
df := ctx.Csv("users", "../testdata/users.csv").
Expand Down
68 changes: 68 additions & 0 deletions internal/engine/cost/internal/testkit/mockCat.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@ package testkit

import (
"fmt"
"github.com/kwilteam/kwil-db/core/types"

ds "github.com/kwilteam/kwil-db/internal/engine/cost/datasource"
dt "github.com/kwilteam/kwil-db/internal/engine/cost/datatypes"
Expand Down Expand Up @@ -44,3 +45,70 @@ func InitMockCatalog() *mockCatalog {
},
}
}

var (
MockUsersSchemaTable = &types.Table{
Name: "users",
Columns: []*types.Column{
{
Name: "id",
Type: types.IntType,
Attributes: []*types.Attribute{
{
Type: types.PRIMARY_KEY,
},
{
Type: types.NOT_NULL,
},
},
},
{
Name: "username",
Type: types.TextType,
Attributes: []*types.Attribute{
{
Type: types.NOT_NULL,
},
{
Type: types.UNIQUE,
},
{
Type: types.MIN_LENGTH,
Value: "5",
},
{
Type: types.MAX_LENGTH,
Value: "32",
},
},
},
{
Name: "age",
Type: types.IntType,
Attributes: []*types.Attribute{
{
Type: types.NOT_NULL,
},
},
},
{
Name: "state",
Type: types.TextType,
Attributes: []*types.Attribute{
{
Type: types.NOT_NULL,
},
},
},
{
Name: "wallet",
Type: types.TextType,
Attributes: []*types.Attribute{
{
Type: types.NOT_NULL,
},
},
},
},
}
)
Loading

0 comments on commit 3bed38b

Please sign in to comment.