Skip to content

Commit 65b3d05

Browse files
ea-rusea-rus
committed
Merge branch 'window-fix' into parser_fixes
2 parents 4a67a14 + b65132d commit 65b3d05

File tree

11 files changed

+166
-83
lines changed

11 files changed

+166
-83
lines changed

‎mindsdb_sql/parser/ast/select/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from .select import Select
22
from .common_table_expression import CommonTableExpression
3-
from .union import Union
3+
from .union import Union, Except, Intersect
44
from .constant import Constant, NullConstant, Last
55
from .star import Star
66
from .identifier import Identifier

‎mindsdb_sql/parser/ast/select/case.py

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,14 @@
44

55

66
class Case(ASTNode):
7-
def __init__(self, rules, default=None, *args, **kwargs):
7+
def __init__(self, rules, default=None, arg=None, *args, **kwargs):
88
super().__init__(*args, **kwargs)
99

1010
# structure:
1111
# [
1212
# [ condition, result ]
1313
# ]
14+
self.arg = arg
1415
self.rules = rules
1516
self.default = default
1617

@@ -36,7 +37,12 @@ def to_tree(self, *args, level=0, **kwargs):
3637
if self.default is not None:
3738
default_str = f'{ind1}default => {self.default.to_string()}\n'
3839

40+
arg_str = ''
41+
if self.arg is not None:
42+
arg_str = f'{ind1}arg => {self.arg.to_string()}\n'
43+
3944
return f'{ind}Case(\n' \
45+
f'{arg_str}'\
4046
f'{rules_str}\n' \
4147
f'{default_str}' \
4248
f'{ind})'
@@ -53,4 +59,8 @@ def get_string(self, *args, alias=True, **kwargs):
5359
default_str = ''
5460
if self.default is not None:
5561
default_str = f' ELSE {self.default.to_string()}'
56-
return f"CASE {rules_str}{default_str} END"
62+
63+
arg_str = ''
64+
if self.arg is not None:
65+
arg_str = f'{self.arg.to_string()} '
66+
return f"CASE {arg_str}{rules_str}{default_str} END"

‎mindsdb_sql/parser/ast/select/operation.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -98,12 +98,13 @@ def get_string(self, *args, **kwargs):
9898

9999

100100
class WindowFunction(ASTNode):
101-
def __init__(self, function, partition=None, order_by=None, alias=None):
101+
def __init__(self, function, partition=None, order_by=None, alias=None, modifier=None):
102102
super().__init__()
103103
self.function = function
104104
self.partition = partition
105105
self.order_by = order_by
106106
self.alias = alias
107+
self.modifier = modifier
107108

108109
def to_tree(self, *args, level=0, **kwargs):
109110
fnc_str = self.function.to_tree(level=level+2)
@@ -143,7 +144,8 @@ def to_string(self, *args, **kwargs):
143144
alias_str = self.alias.to_string()
144145
else:
145146
alias_str = ''
146-
return f'{fnc_str} over({partition_str} {order_str}) {alias_str}'
147+
modifier_str = self.modifier if self.modifier else ''
148+
return f'{fnc_str} over({partition_str} {order_str}{modifier_str}) {alias_str}'
147149

148150

149151
class Object(ASTNode):

‎mindsdb_sql/parser/ast/select/union.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
from mindsdb_sql.parser.utils import indent
33

44

5-
class Union(ASTNode):
5+
class CombiningQuery(ASTNode):
6+
operation = None
67

78
def __init__(self,
89
left,
@@ -24,7 +25,8 @@ def to_tree(self, *args, level=0, **kwargs):
2425
left_str = f'\n{ind1}left=\n{self.left.to_tree(level=level + 2)},'
2526
right_str = f'\n{ind1}right=\n{self.right.to_tree(level=level + 2)},'
2627

27-
out_str = f'{ind}Union(unique={repr(self.unique)},' \
28+
cls_name = self.__class__.__name__
29+
out_str = f'{ind}{cls_name}(unique={repr(self.unique)},' \
2830
f'{left_str}' \
2931
f'{right_str}' \
3032
f'\n{ind})'
@@ -33,7 +35,21 @@ def to_tree(self, *args, level=0, **kwargs):
3335
def get_string(self, *args, **kwargs):
3436
left_str = str(self.left)
3537
right_str = str(self.right)
36-
keyword = 'UNION' if self.unique else 'UNION ALL'
38+
keyword = self.operation
39+
if not self.unique:
40+
keyword += ' ALL'
3741
out_str = f"""{left_str}\n{keyword}\n{right_str}"""
3842

3943
return out_str
44+
45+
46+
class Union(CombiningQuery):
47+
operation = 'UNION'
48+
49+
50+
class Intersect(CombiningQuery):
51+
operation = 'INTERSECT'
52+
53+
54+
class Except(CombiningQuery):
55+
operation = 'EXCEPT'

‎mindsdb_sql/parser/dialects/mindsdb/lexer.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ class MindsDBLexer(Lexer):
5555

5656
JOIN, INNER, OUTER, CROSS, LEFT, RIGHT, ON,
5757

58-
UNION, ALL,
58+
UNION, ALL, INTERSECT, EXCEPT,
5959

6060
# CASE
6161
CASE, ELSE, END, THEN, WHEN,
@@ -238,6 +238,8 @@ class MindsDBLexer(Lexer):
238238
# UNION
239239

240240
UNION = r'\bUNION\b'
241+
INTERSECT = r'\bINTERSECT\b'
242+
EXCEPT = r'\bEXCEPT\b'
241243
ALL = r'\bALL\b'
242244

243245
# CASE

‎mindsdb_sql/parser/dialects/mindsdb/parser.py

Lines changed: 50 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ class MindsDBParser(Parser):
7070
'drop_dataset',
7171
'select',
7272
'insert',
73+
'union',
7374
'update',
7475
'delete',
7576
'evaluate',
@@ -614,10 +615,13 @@ def update(self, p):
614615

615616
# INSERT
616617
@_('INSERT INTO identifier LPAREN column_list RPAREN select',
617-
'INSERT INTO identifier select')
618+
'INSERT INTO identifier LPAREN column_list RPAREN union',
619+
'INSERT INTO identifier select',
620+
'INSERT INTO identifier union')
618621
def insert(self, p):
619622
columns = getattr(p, 'column_list', None)
620-
return Insert(table=p.identifier, columns=columns, from_select=p.select)
623+
query = p.select if hasattr(p, 'select') else p.union
624+
return Insert(table=p.identifier, columns=columns, from_select=query)
621625

622626
@_('INSERT INTO identifier LPAREN column_list RPAREN VALUES expr_list_set',
623627
'INSERT INTO identifier VALUES expr_list_set')
@@ -998,19 +1002,35 @@ def database_engine(self, p):
9981002
engine = p.string
9991003
return {'identifier':p.identifier, 'engine':engine, 'if_not_exists':p.if_not_exists_or_empty}
10001004

1001-
# UNION / UNION ALL
1002-
@_('select UNION select')
1003-
def select(self, p):
1004-
return Union(left=p[0], right=p[2], unique=True)
1005-
1006-
@_('select UNION ALL select')
1007-
def select(self, p):
1008-
return Union(left=p[0], right=p[3], unique=False)
1005+
# Combining
1006+
@_('select UNION select',
1007+
'union UNION select',
1008+
'select UNION ALL select',
1009+
'union UNION ALL select')
1010+
def union(self, p):
1011+
unique = not hasattr(p, 'ALL')
1012+
return Union(left=p[0], right=p[2] if unique else p[3], unique=unique)
1013+
1014+
@_('select INTERSECT select',
1015+
'union INTERSECT select',
1016+
'select INTERSECT ALL select',
1017+
'union INTERSECT ALL select')
1018+
def union(self, p):
1019+
unique = not hasattr(p, 'ALL')
1020+
return Intersect(left=p[0], right=p[2] if unique else p[3], unique=unique)
1021+
@_('select EXCEPT select',
1022+
'union EXCEPT select',
1023+
'select EXCEPT ALL select',
1024+
'union EXCEPT ALL select')
1025+
def union(self, p):
1026+
unique = not hasattr(p, 'ALL')
1027+
return Except(left=p[0], right=p[2] if unique else p[3], unique=unique)
10091028

10101029
# tableau
10111030
@_('LPAREN select RPAREN')
1031+
@_('LPAREN union RPAREN')
10121032
def select(self, p):
1013-
return p.select
1033+
return p[1]
10141034

10151035
# WITH
10161036
@_('ctes select')
@@ -1030,13 +1050,14 @@ def ctes(self, p):
10301050
]
10311051
return ctes
10321052

1033-
@_('WITH identifier cte_columns_or_nothing AS LPAREN select RPAREN')
1053+
@_('WITH identifier cte_columns_or_nothing AS LPAREN select RPAREN',
1054+
'WITH identifier cte_columns_or_nothing AS LPAREN union RPAREN')
10341055
def ctes(self, p):
10351056
return [
10361057
CommonTableExpression(
10371058
name=p.identifier,
10381059
columns=p.cte_columns_or_nothing,
1039-
query=p.select)
1060+
query=p[5])
10401061
]
10411062

10421063
@_('empty')
@@ -1331,6 +1352,15 @@ def column_list(self, p):
13311352
def case(self, p):
13321353
return Case(rules=p.case_conditions, default=getattr(p, 'expr', None))
13331354

1355+
@_('CASE expr case_conditions ELSE expr END',
1356+
'CASE expr case_conditions END')
1357+
def case(self, p):
1358+
if hasattr(p, 'expr'):
1359+
arg, default = p.expr, None
1360+
else:
1361+
arg, default = p.expr0, p.expr1
1362+
return Case(rules=p.case_conditions, default=default, arg=arg)
1363+
13341364
@_('case_condition',
13351365
'case_conditions case_condition')
13361366
def case_conditions(self, p):
@@ -1343,13 +1373,18 @@ def case_condition(self, p):
13431373
return [p.expr0, p.expr1]
13441374

13451375
# Window function
1346-
@_('function OVER LPAREN window RPAREN')
1376+
@_('expr OVER LPAREN window RPAREN',
1377+
'expr OVER LPAREN window id BETWEEN id id AND id id RPAREN')
13471378
def window_function(self, p):
13481379

1380+
modifier = None
1381+
if hasattr(p, 'BETWEEN'):
1382+
modifier = f'{p.id0} BETWEEN {p.id1} {p.id2} AND {p.id3} {p.id4}'
13491383
return WindowFunction(
1350-
function=p.function,
1384+
function=p.expr,
13511385
order_by=p.window.get('order_by'),
13521386
partition=p.window.get('partition'),
1387+
modifier=modifier,
13531388
)
13541389

13551390
@_('window PARTITION_BY expr_list')

‎mindsdb_sql/planner/query_planner.py

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from mindsdb_sql.exceptions import PlanningException
44
from mindsdb_sql.parser import ast
55
from mindsdb_sql.parser.ast import (Select, Identifier, Join, Star, BinaryOperation, Constant, Union, CreateTable,
6-
Function, Insert,
6+
Function, Insert, Except, Intersect,
77
Update, NativeQuery, Parameter, Delete)
88
from mindsdb_sql.planner import utils
99
from mindsdb_sql.planner.query_plan import QueryPlan
@@ -229,6 +229,19 @@ def find_objects(node, is_table, **kwargs):
229229
mdb_entities.append(node)
230230

231231
query_traversal(query, find_objects)
232+
233+
# cte names are not mdb objects
234+
if query.cte:
235+
cte_names = [
236+
cte.name.parts[-1]
237+
for cte in query.cte
238+
]
239+
mdb_entities = [
240+
item
241+
for item in mdb_entities
242+
if '.'.join(item.parts) not in cte_names
243+
]
244+
232245
return {
233246
'mdb_entities': mdb_entities,
234247
'integrations': integrations,
@@ -672,13 +685,23 @@ def plan_delete(self, query: Delete):
672685
))
673686

674687
def plan_cte(self, query):
688+
query_info = self.get_query_info(query)
689+
690+
if (
691+
len(query_info['integrations']) == 1
692+
and len(query_info['mdb_entities']) == 0
693+
and len(query_info['user_functions']) == 0
694+
):
695+
# single integration, will be planned later
696+
return
697+
675698
for cte in query.cte:
676699
step = self.plan_select(cte.query)
677700
name = cte.name.parts[-1]
678701
self.cte_results[name] = step.result
679702

680703
def plan_select(self, query, integration=None):
681-
if isinstance(query, Union):
704+
if isinstance(query, (Union, Except, Intersect)):
682705
return self.plan_union(query, integration=integration)
683706

684707
if query.cte is not None:
@@ -734,14 +757,15 @@ def plan_sub_select(self, query, prev_step, add_absent_cols=False):
734757
return prev_step
735758

736759
def plan_union(self, query, integration=None):
737-
if isinstance(query.left, Union):
738-
step1 = self.plan_union(query.left, integration=integration)
739-
else:
740-
# it is select
741-
step1 = self.plan_select(query.left, integration=integration)
760+
step1 = self.plan_select(query.left, integration=integration)
742761
step2 = self.plan_select(query.right, integration=integration)
762+
operation = 'union'
763+
if isinstance(query, Except):
764+
operation = 'except'
765+
elif isinstance(query, Intersect):
766+
operation = 'intersect'
743767

744-
return self.plan.add_step(UnionStep(left=step1.result, right=step2.result, unique=query.unique))
768+
return self.plan.add_step(UnionStep(left=step1.result, right=step2.result, unique=query.unique, operation=operation))
745769

746770
# method for compatibility
747771
def from_query(self, query=None):
@@ -750,7 +774,7 @@ def from_query(self, query=None):
750774
if query is None:
751775
query = self.query
752776

753-
if isinstance(query, (Select, Union)):
777+
if isinstance(query, (Select, Union, Except, Intersect)):
754778
self.plan_select(query)
755779
elif isinstance(query, CreateTable):
756780
self.plan_create_table(query)

‎mindsdb_sql/planner/steps.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,11 +75,12 @@ def __init__(self, left, right, query, *args, **kwargs):
7575

7676
class UnionStep(PlanStep):
7777
"""Union of two dataframes, producing a new dataframe"""
78-
def __init__(self, left, right, unique, *args, **kwargs):
78+
def __init__(self, left, right, unique, operation='union', *args, **kwargs):
7979
super().__init__(*args, **kwargs)
8080
self.left = left
8181
self.right = right
8282
self.unique = unique
83+
self.operation = operation
8384

8485

8586
# TODO remove

‎mindsdb_sql/planner/utils.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ def query_traversal(node, callback, is_table=False, is_target=False, parent_quer
145145
array.append(node_out)
146146
node.order_by = array
147147

148-
elif isinstance(node, ast.Union):
148+
elif isinstance(node, (ast.Union, ast.Intersect, ast.Except)):
149149
node_out = query_traversal(node.left, callback, parent_query=node)
150150
if node_out is not None:
151151
node.left = node_out

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy