|
1 | 1 | package plan_node_test |
2 | 2 |
|
3 | 3 | import ( |
4 | | - "context" |
5 | 4 | "encoding/json" |
6 | 5 | "fmt" |
7 | 6 | "github.com/stretchr/testify/assert" |
8 | 7 | "math" |
9 | 8 | "os" |
10 | 9 | "pbench/presto/plan_node" |
11 | | - "strings" |
12 | 10 | "testing" |
13 | 11 | ) |
14 | 12 |
|
15 | | -func traceValue(assignmentMap map[string]plan_node.Value, tableHandle *plan_node.HiveTableHandle, value plan_node.Value) plan_node.Value { |
16 | | - switch typed := value.(type) { |
17 | | - case *plan_node.HiveColumnHandle: |
18 | | - if typed.Table == nil { |
19 | | - typed.Table = tableHandle |
| 13 | +func testParseJoin(t *testing.T, fileName, expected string) { |
| 14 | + t.Helper() |
| 15 | + t.Run(fileName, func(t *testing.T) { |
| 16 | + bytes, err := os.ReadFile(fileName) |
| 17 | + if !assert.Nil(t, err) { |
| 18 | + t.FailNow() |
20 | 19 | } |
21 | | - case *plan_node.TypeCastedValue: |
22 | | - typed.OriginalValue = traceValue(assignmentMap, tableHandle, typed.OriginalValue) |
23 | | - case *plan_node.IdentRef: |
24 | | - return traceValue(assignmentMap, tableHandle, assignmentMap[typed.Ident]) |
25 | | - case *plan_node.FunctionCall: |
26 | | - for i := 0; i < len(typed.Parameters); i++ { |
27 | | - typed.Parameters[i] = traceValue(assignmentMap, tableHandle, typed.Parameters[i]) |
| 20 | + planTree := make(plan_node.PlanTree) |
| 21 | + if !assert.Nil(t, json.Unmarshal(bytes, &planTree)) { |
| 22 | + t.FailNow() |
28 | 23 | } |
29 | | - } |
30 | | - return value |
| 24 | + joins, parseErr := planTree.ParseJoins() |
| 25 | + if assert.Nil(t, parseErr) { |
| 26 | + assert.Equal(t, expected, fmt.Sprint(joins)) |
| 27 | + } |
| 28 | + }) |
31 | 29 | } |
32 | 30 |
|
33 | | -func TestPlanTree(t *testing.T) { |
34 | | - bytes, err := os.ReadFile("sample_plan.json") |
35 | | - if !assert.Nil(t, err) { |
36 | | - t.FailNow() |
37 | | - } |
38 | | - planTree := make(plan_node.PlanTree) |
39 | | - if !assert.Nil(t, json.Unmarshal(bytes, &planTree)) { |
40 | | - t.FailNow() |
41 | | - } |
42 | | - assignmentMap := make(map[string]plan_node.Value) |
43 | | - assert.Nil(t, planTree.Traverse(context.Background(), func(ctx context.Context, node *plan_node.PlanNode) error { |
44 | | - tableHandle := plan_node.ParseHiveTableHandle(node.Identifier) |
45 | | - id := fmt.Sprintf("%s_%s", node.Id, node.Name) |
46 | | - if node.Details != "" { |
47 | | - ast, parseErr := plan_node.PlanNodeDetailParser.ParseString(id, node.Details) |
48 | | - if !assert.Nil(t, parseErr) { |
49 | | - t.FailNow() |
50 | | - } |
51 | | - for i := len(ast.Stmts) - 1; i >= 0; i-- { |
52 | | - if assignment, ok := ast.Stmts[i].(*plan_node.Assignment); ok { |
53 | | - assignmentMap[assignment.Identifier.Ident] = traceValue(assignmentMap, tableHandle, assignment.AssignedValue) |
54 | | - } |
55 | | - } |
56 | | - } |
57 | | - if plan_node.IsJoin[node.Name] { |
58 | | - preds, parseErr := plan_node.PlanNodeJoinPredicatesParser.ParseString(id, node.Identifier) |
59 | | - if !assert.Nil(t, parseErr) { |
60 | | - t.FailNow() |
61 | | - } |
62 | | - b := strings.Builder{} |
63 | | - for _, pred := range preds.Predicates { |
64 | | - b.WriteString(assignmentMap[pred.LeftColumn].String()) |
65 | | - b.WriteString(" x ") |
66 | | - b.WriteString(assignmentMap[pred.RightColumn].String()) |
67 | | - b.WriteString("\n") |
68 | | - } |
69 | | - fmt.Printf("%s\n%s\n", id, b.String()) |
70 | | - } |
71 | | - return nil |
72 | | - }, plan_node.PlanTreeDFSTraverse)) |
| 31 | +func TestParseJoin(t *testing.T) { |
| 32 | + testParseJoin(t, "sample.plan.json", |
| 33 | + `[{RightJoin CAST(glue.ng_public.admin_system_city.id AS bigint) glue.lks.LR_branded_car_enrollment.city_id} {LeftJoin glue.lks.LR_branded_car_enrollment.country glue.lks.LR_admin_system_country.code} {InnerJoin CAST(glue.ng_public.fleet_car.id AS bigint) glue.ng_public.fleet_car_tag_binding.car_id} {LeftJoin glue.ng_public.fleet_car_tag_binding.car_tag_id CAST(glue.ng_public.fleet_car_tag.id AS bigint)} {LeftJoin glue.lks.LR_branded_car_enrollment.car_id CAST(glue.ng_public.fleet_car.id AS bigint)}]`) |
| 34 | + testParseJoin(t, "arithmetics.plan.json", `[{InnerJoin (hive.test_join.t1.a) + (INTEGER '2') ((hive.test_join.t2.b) + (INTEGER '1')) - ((INTEGER '2') * (abs(hive.test_join.t2.c)))}]`) |
73 | 35 | } |
74 | 36 |
|
75 | 37 | func TestJsonFloat64(t *testing.T) { |
|
0 commit comments