Skip to content

Commit

Permalink
unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@ea-rus
ea-rus committed Nov 2, 2024
1 parent ef30b12 commit 0f5dbb1
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 32 deletions.
24 changes: 12 additions & 12 deletions tests/test_planner/test_join_predictor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@ def test_join_predictor_plan(self):
"""
query = parse_sql(sql)

query_step = parse_sql("select tab1.column1, pred.predicted")
query_step.from_table = Parameter(Result(2))
expected_plan = QueryPlan(
steps=[
FetchDataframeStep(integration='int',
Expand Down Expand Up @@ -75,7 +73,7 @@ def test_join_predictor_plan_aliases(self):
plan = plan_query(query, integrations=['int'], predictor_namespace='mindsdb', predictor_metadata={'pred': {}})

assert plan.steps == expected_plan.steps


def test_join_predictor_plan_limit(self):

Expand Down Expand Up @@ -116,7 +114,7 @@ def test_join_predictor_plan_limit(self):
plan = plan_query(query, integrations=['int'], predictor_namespace='mindsdb', predictor_metadata={'pred': {}})

assert plan.steps == expected_plan.steps


# def test_join_predictor_error_when_filtering_on_predictions(self):
# """
Expand Down Expand Up @@ -673,15 +671,16 @@ def test_complex_subselect(self):

sql = '''
select t2.x, m.id, (select a from int.tab0 where x=0) from int.tab1 t1
join int.tab2 t2 on t1.x = t2.x
join int.tab2 t2 on t1.x = t2.a
join mindsdb.pred m
where m.a=(select a from int.tab3 where x=3)
and t2.x=(select a from int.tab4 where x=4)
and t1.b=1 and t2.b=2 and t1.a = t2.a
'''

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

subquery = parse_sql("""
select t2.x, m.id, x
Expand All @@ -708,22 +707,23 @@ def test_complex_subselect(self):
# tables
FetchDataframeStep(integration='int',
query=parse_sql('select * from tab1 as t1 where b=1')),
SubSelectStep(dataframe=Result(3), query=Select(targets=[Identifier('x')], distinct=True)),
FetchDataframeStep(integration='int', query=q_table2),
JoinStep(left=Result(3), right=Result(4),
JoinStep(left=Result(3), right=Result(5),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
join_type=JoinType.JOIN,
condition=BinaryOperation(op='=', args=[Identifier('t1.x'), Identifier('t2.x')])
condition=BinaryOperation(op='=', args=[Identifier('t1.x'), Identifier('t2.a')])
)
),
# model
ApplyPredictorStep(namespace='mindsdb', dataframe=Result(5),
ApplyPredictorStep(namespace='mindsdb', dataframe=Result(6),
predictor=Identifier('pred', alias=Identifier('m')), row_dict={'a': Result(1)}),
JoinStep(left=Result(5), right=Result(6),
JoinStep(left=Result(6), right=Result(7),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
join_type=JoinType.JOIN)),
QueryStep(subquery, from_table=Result(7)),
QueryStep(subquery, from_table=Result(8)),
],
)
plan = plan_query(query, integrations=['int'], predictor_namespace='mindsdb', predictor_metadata={'pred': {}})
Expand Down
44 changes: 24 additions & 20 deletions tests/test_planner/test_join_tables.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ def test_join_tables_plan(self):
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
from_table=Join(left=Identifier('int.tab1'),
right=Identifier('int2.tab2'),
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
)
)
Expand All @@ -35,7 +35,7 @@ def test_join_tables_plan(self):
JoinStep(left=Result(0), right=Result(1),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'),
Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
Expand All @@ -45,13 +45,13 @@ def test_join_tables_plan(self):
)

assert plan.steps == expected_plan.steps


def test_join_tables_where_plan(self):
query = parse_sql('''
SELECT tab1.column1, tab2.column1, tab2.column2
FROM int.tab1
INNER JOIN int2.tab2 ON tab1.column1 = tab2.column1
INNER JOIN int2.tab2 ON tab1.column1 > tab2.column1
WHERE ((tab1.column1 = 1)
AND (tab2.column1 = 0))
AND (tab1.column3 = tab2.column3)
Expand All @@ -71,7 +71,7 @@ def test_join_tables_where_plan(self):
JoinStep(left=Result(0), right=Result(1),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'),
Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
Expand All @@ -90,7 +90,7 @@ def test_join_tables_plan_groupby(self):
Function('sum', args=[Identifier('tab2.column2')], alias=Identifier('total'))],
from_table=Join(left=Identifier('int.tab1'),
right=Identifier('int2.tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
),
Expand All @@ -117,7 +117,7 @@ def test_join_tables_plan_groupby(self):
JoinStep(left=Result(0), right=Result(1),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'),
Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
Expand All @@ -126,13 +126,13 @@ def test_join_tables_plan_groupby(self):
],
)
assert plan.steps == expected_plan.steps


def test_join_tables_plan_limit_offset(self):
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
from_table=Join(left=Identifier('int.tab1'),
right=Identifier('int2.tab2'),
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
),
limit=Constant(10),
Expand Down Expand Up @@ -161,7 +161,7 @@ def test_join_tables_plan_limit_offset(self):
JoinStep(left=Result(0), right=Result(1),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'),
Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
Expand All @@ -171,13 +171,13 @@ def test_join_tables_plan_limit_offset(self):
)

assert plan.steps == expected_plan.steps


def test_join_tables_plan_order_by(self):
query = Select(targets=[Identifier('tab1.column1'), Identifier('tab2.column1'), Identifier('tab2.column2')],
from_table=Join(left=Identifier('int.tab1'),
right=Identifier('int2.tab2'),
condition=BinaryOperation(op='=', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
condition=BinaryOperation(op='>', args=[Identifier('tab1.column1'), Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
),
limit=Constant(10),
Expand All @@ -203,7 +203,7 @@ def test_join_tables_plan_order_by(self):
JoinStep(left=Result(0), right=Result(1),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(op='=',
condition=BinaryOperation(op='>',
args=[Identifier('tab1.column1'),
Identifier('tab2.column1')]),
join_type=JoinType.INNER_JOIN
Expand All @@ -213,7 +213,7 @@ def test_join_tables_plan_order_by(self):
)

assert plan.steps == expected_plan.steps


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

for i in range(len(plan.steps)):
assert plan.steps[i] == expected_plan.steps[i]


def _disabled_test_join_tables_error_on_unspecified_table_in_condition(self):
# disabled: identifier can be environment of system variable
Expand Down Expand Up @@ -328,7 +328,7 @@ def test_join_tables_plan_default_namespace(self):
def test_complex_join_tables(self):
query = parse_sql('''
select * from int1.tbl1 t1
right join int2.tbl2 t2 on t1.id=t2.id
right join int2.tbl2 t2 on t1.id>t2.id
join pred m
left join tbl3 on tbl3.id=t1.id
where t1.a=1 and t2.b=2 and 1=1
Expand All @@ -337,6 +337,9 @@ def test_complex_join_tables(self):
subquery = copy.deepcopy(query)
subquery.from_table = None

q_table3 = parse_sql('select * from tbl3 where id in 0')
q_table3.where.args[1] = Parameter(Result(5))

plan = plan_query(query, integrations=['int1', 'int2', 'proj'], default_namespace='proj',
predictor_metadata=[{'name': 'pred', 'integration_name': 'proj'}])

Expand All @@ -349,7 +352,7 @@ def test_complex_join_tables(self):
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(
op='=',
op='>',
args=[Identifier('t1.id'),
Identifier('t2.id')]),
join_type=JoinType.RIGHT_JOIN)),
Expand All @@ -359,17 +362,18 @@ def test_complex_join_tables(self):
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
join_type=JoinType.JOIN)),
FetchDataframeStep(integration='proj', query=parse_sql('select * from tbl3')),
SubSelectStep(dataframe=Result(0), query=Select(targets=[Identifier('id')], distinct=True)),
FetchDataframeStep(integration='proj', query=q_table3),
JoinStep(left=Result(4),
right=Result(5),
right=Result(6),
query=Join(left=Identifier('tab1'),
right=Identifier('tab2'),
condition=BinaryOperation(
op='=',
args=[Identifier('tbl3.id'),
Identifier('t1.id')]),
join_type=JoinType.LEFT_JOIN)),
QueryStep(subquery, from_table=Result(6)),
QueryStep(subquery, from_table=Result(7)),
]
)

Expand Down

0 comments on commit 0f5dbb1

Please sign in to comment.