Skip to content

Commit 29ab045

Browse files
ea-rusea-rus
committed
Merge branch 'optimize-join' into cte-support
2 parents e324f45 + 0f5dbb1 commit 29ab045

2 files changed

Lines changed: 34 additions & 30 deletions

File tree

‎tests/test_planner/test_join_predictor.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ def test_join_predictor_plan(self):
2222
"""
2323
query = parse_sql(sql)
2424

25-
query_step = parse_sql("select tab1.column1, pred.predicted")
26-
query_step.from_table = Parameter(Result(2))
2725
expected_plan = QueryPlan(
2826
steps=[
2927
FetchDataframeStep(integration='int',
@@ -665,15 +663,16 @@ def test_complex_subselect(self):
665663

666664
sql = '''
667665
select t2.x, m.id, (select a from int.tab0 where x=0) from int.tab1 t1
668-
join int.tab2 t2 on t1.x = t2.x
666+
join int.tab2 t2 on t1.x = t2.a
669667
join mindsdb.pred m
670668
where m.a=(select a from int.tab3 where x=3)
671669
and t2.x=(select a from int.tab4 where x=4)
672670
and t1.b=1 and t2.b=2 and t1.a = t2.a
673671
'''
674672

675-
q_table2 = parse_sql('select * from tab2 as t2 where x=0 and b=2 ')
676-
q_table2.where.args[0].args[1] = Parameter(Result(2))
673+
q_table2 = parse_sql('select * from tab2 as t2 where x=0 and b=2 AND a IN 1')
674+
q_table2.where.args[0].args[0].args[1] = Parameter(Result(2))
675+
q_table2.where.args[1].args[1] = Parameter(Result(4))
677676

678677
subquery = parse_sql("""
679678
select t2.x, m.id, x
@@ -700,22 +699,23 @@ def test_complex_subselect(self):
700699
# tables
701700
FetchDataframeStep(integration='int',
702701
query=parse_sql('select * from tab1 as t1 where b=1')),
702+
SubSelectStep(dataframe=Result(3), query=Select(targets=[Identifier('x')], distinct=True)),
703703
FetchDataframeStep(integration='int', query=q_table2),
704-
JoinStep(left=Result(3), right=Result(4),
704+
JoinStep(left=Result(3), right=Result(5),
705705
query=Join(left=Identifier('tab1'),
706706
right=Identifier('tab2'),
707707
join_type=JoinType.JOIN,
708-
condition=BinaryOperation(op='=', args=[Identifier('t1.x'), Identifier('t2.x')])
708+
condition=BinaryOperation(op='=', args=[Identifier('t1.x'), Identifier('t2.a')])
709709
)
710710
),
711711
# model
712-
ApplyPredictorStep(namespace='mindsdb', dataframe=Result(5),
712+
ApplyPredictorStep(namespace='mindsdb', dataframe=Result(6),
713713
predictor=Identifier('pred', alias=Identifier('m')), row_dict={'a': Result(1)}),
714-
JoinStep(left=Result(5), right=Result(6),
714+
JoinStep(left=Result(6), right=Result(7),
715715
query=Join(left=Identifier('tab1'),
716716
right=Identifier('tab2'),
717717
join_type=JoinType.JOIN)),
718-
QueryStep(subquery, from_table=Result(7)),
718+
QueryStep(subquery, from_table=Result(8)),
719719
],
720720
)
721721
plan = plan_query(query, integrations=['int'], predictor_namespace='mindsdb', predictor_metadata={'pred': {}})

‎tests/test_planner/test_join_tables.py‎

Lines changed: 24 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ def test_join_tables_plan(self):
1616
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
1717
from_table=Join(left=Identifier('int.tab1'),
1818
right=Identifier('int2.tab2'),
19-
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
19+
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
2020
join_type=JoinType.INNER_JOIN
2121
)
2222
)
@@ -35,7 +35,7 @@ def test_join_tables_plan(self):
3535
JoinStep(left=Result(0), right=Result(1),
3636
query=Join(left=Identifier('tab1'),
3737
right=Identifier('tab2'),
38-
condition=BinaryOperation(op='=',
38+
condition=BinaryOperation(op='>',
3939
args=[Identifier('tab1.column1'),
4040
Identifier('tab2.column1')]),
4141
join_type=JoinType.INNER_JOIN
@@ -45,13 +45,13 @@ def test_join_tables_plan(self):
4545
)
4646

4747
assert plan.steps == expected_plan.steps
48-
48+
4949

5050
def test_join_tables_where_plan(self):
5151
query = parse_sql('''
5252
SELECT tab1.column1, tab2.column1, tab2.column2
5353
FROM int.tab1
54-
INNER JOIN int2.tab2 ON tab1.column1 = tab2.column1
54+
INNER JOIN int2.tab2 ON tab1.column1 > tab2.column1
5555
WHERE ((tab1.column1 = 1)
5656
AND (tab2.column1 = 0))
5757
AND (tab1.column3 = tab2.column3)
@@ -71,7 +71,7 @@ def test_join_tables_where_plan(self):
7171
JoinStep(left=Result(0), right=Result(1),
7272
query=Join(left=Identifier('tab1'),
7373
right=Identifier('tab2'),
74-
condition=BinaryOperation(op='=',
74+
condition=BinaryOperation(op='>',
7575
args=[Identifier('tab1.column1'),
7676
Identifier('tab2.column1')]),
7777
join_type=JoinType.INNER_JOIN
@@ -90,7 +90,7 @@ def test_join_tables_plan_groupby(self):
9090
Function('sum', args=[Identifier('tab2.column2')], alias=Identifier('total'))],
9191
from_table=Join(left=Identifier('int.tab1'),
9292
right=Identifier('int2.tab2'),
93-
condition=BinaryOperation(op='=',
93+
condition=BinaryOperation(op='>',
9494
args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
9595
join_type=JoinType.INNER_JOIN
9696
),
@@ -117,7 +117,7 @@ def test_join_tables_plan_groupby(self):
117117
JoinStep(left=Result(0), right=Result(1),
118118
query=Join(left=Identifier('tab1'),
119119
right=Identifier('tab2'),
120-
condition=BinaryOperation(op='=',
120+
condition=BinaryOperation(op='>',
121121
args=[Identifier('tab1.column1'),
122122
Identifier('tab2.column1')]),
123123
join_type=JoinType.INNER_JOIN
@@ -126,13 +126,13 @@ def test_join_tables_plan_groupby(self):
126126
],
127127
)
128128
assert plan.steps == expected_plan.steps
129-
129+
130130

131131
def test_join_tables_plan_limit_offset(self):
132132
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
133133
from_table=Join(left=Identifier('int.tab1'),
134134
right=Identifier('int2.tab2'),
135-
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
135+
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
136136
join_type=JoinType.INNER_JOIN
137137
),
138138
limit=Constant(10),
@@ -161,7 +161,7 @@ def test_join_tables_plan_limit_offset(self):
161161
JoinStep(left=Result(0), right=Result(1),
162162
query=Join(left=Identifier('tab1'),
163163
right=Identifier('tab2'),
164-
condition=BinaryOperation(op='=',
164+
condition=BinaryOperation(op='>',
165165
args=[Identifier('tab1.column1'),
166166
Identifier('tab2.column1')]),
167167
join_type=JoinType.INNER_JOIN
@@ -171,13 +171,13 @@ def test_join_tables_plan_limit_offset(self):
171171
)
172172

173173
assert plan.steps == expected_plan.steps
174-
174+
175175

176176
def test_join_tables_plan_order_by(self):
177177
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
178178
from_table=Join(left=Identifier('int.tab1'),
179179
right=Identifier('int2.tab2'),
180-
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
180+
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
181181
join_type=JoinType.INNER_JOIN
182182
),
183183
limit=Constant(10),
@@ -203,7 +203,7 @@ def test_join_tables_plan_order_by(self):
203203
JoinStep(left=Result(0), right=Result(1),
204204
query=Join(left=Identifier('tab1'),
205205
right=Identifier('tab2'),
206-
condition=BinaryOperation(op='=',
206+
condition=BinaryOperation(op='>',
207207
args=[Identifier('tab1.column1'),
208208
Identifier('tab2.column1')]),
209209
join_type=JoinType.INNER_JOIN
@@ -213,7 +213,7 @@ def test_join_tables_plan_order_by(self):
213213
)
214214

215215
assert plan.steps == expected_plan.steps
216-
216+
217217

218218
# This quiery should be sent to integration without raising exception
219219
# def test_join_tables_where_ambigous_column_error(self):
@@ -278,7 +278,7 @@ def test_join_tables_disambiguate_identifiers_in_condition(self):
278278

279279
for i in range(len(plan.steps)):
280280
assert plan.steps[i] == expected_plan.steps[i]
281-
281+
282282

283283
def _disabled_test_join_tables_error_on_unspecified_table_in_condition(self):
284284
# disabled: identifier can be environment of system variable
@@ -328,7 +328,7 @@ def test_join_tables_plan_default_namespace(self):
328328
def test_complex_join_tables(self):
329329
query = parse_sql('''
330330
select * from int1.tbl1 t1
331-
right join int2.tbl2 t2 on t1.id=t2.id
331+
right join int2.tbl2 t2 on t1.id>t2.id
332332
join pred m
333333
left join tbl3 on tbl3.id=t1.id
334334
where t1.a=1 and t2.b=2 and 1=1
@@ -337,6 +337,9 @@ def test_complex_join_tables(self):
337337
subquery = copy.deepcopy(query)
338338
subquery.from_table = None
339339

340+
q_table3 = parse_sql('select * from tbl3 where id in 0')
341+
q_table3.where.args[1] = Parameter(Result(5))
342+
340343
plan = plan_query(query, integrations=['int1', 'int2', 'proj'], default_namespace='proj',
341344
predictor_metadata=[{'name': 'pred', 'integration_name': 'proj'}])
342345

@@ -349,7 +352,7 @@ def test_complex_join_tables(self):
349352
query=Join(left=Identifier('tab1'),
350353
right=Identifier('tab2'),
351354
condition=BinaryOperation(
352-
op='=',
355+
op='>',
353356
args=[Identifier('t1.id'),
354357
Identifier('t2.id')]),
355358
join_type=JoinType.RIGHT_JOIN)),
@@ -359,17 +362,18 @@ def test_complex_join_tables(self):
359362
query=Join(left=Identifier('tab1'),
360363
right=Identifier('tab2'),
361364
join_type=JoinType.JOIN)),
362-
FetchDataframeStep(integration='proj', query=parse_sql('select * from tbl3')),
365+
SubSelectStep(dataframe=Result(0), query=Select(targets=[Identifier('id')], distinct=True)),
366+
FetchDataframeStep(integration='proj', query=q_table3),
363367
JoinStep(left=Result(4),
364-
right=Result(5),
368+
right=Result(6),
365369
query=Join(left=Identifier('tab1'),
366370
right=Identifier('tab2'),
367371
condition=BinaryOperation(
368372
op='=',
369373
args=[Identifier('tbl3.id'),
370374
Identifier('t1.id')]),
371375
join_type=JoinType.LEFT_JOIN)),
372-
QueryStep(subquery, from_table=Result(6)),
376+
QueryStep(subquery, from_table=Result(7)),
373377
]
374378
)
375379

0 commit comments

Comments
 (0)