-
Notifications
You must be signed in to change notification settings - Fork 20
Expand file tree
/
Copy pathRewriter.py
More file actions
293 lines (248 loc) · 11.7 KB
/
Copy pathRewriter.py
File metadata and controls
293 lines (248 loc) · 11.7 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
from pythonql.Ast import *
from pythonql.sources.source import RDBMSTable
# Helper function
def ensure_list(x):
if isinstance(x,list):
return x
else:
return [x]
# Provide a list of variable that were defined in the given
# list of clauses
def clause_vars(clauses):
vs = set()
for c in clauses:
if c['name'] == 'join':
vs = vs.union(clause_vars(ensure_list(c['left'])))
vs = vs.union(clause_vars(ensure_list(c['right'])))
elif c['name'] in ['for','let','window','match']:
vs = vs.union(set(c['vars']))
elif c['name'] =='groupby':
vs = vs.union({v[1] for v in c['groupby_list']})
return vs
# Compute the list of variables that are referenced in these
# clauses
def clause_live_vars(clauses):
vs = set()
for c in clauses:
if c['name'] == 'select':
vs = vs.union(get_all_vars(get_ast(c['expr'])))
elif c['name'] == 'join' and 'cond' in c:
vs = vs.union(get_all_vars(get_ast(c['cond'])))
elif c['name'] in ['for','let','where']:
e = get_ast(c['expr'])
vs = vs.union(get_all_vars(e))
elif c['name'] == 'orderby':
for ex in c['orderby_list']:
e = get_ast(ex)
vs = vs.union(get_all_vars(e))
elif c['name'] == 'groupby':
for (ex,_) in c['groupby_list']:
e = get_ast(ex)
vs = vs.union(get_all_vars(e))
return vs
# Check if an expression looks like a join condition
# This just checks the syntax of an expression
def is_join_cond(e):
if isinstance(e,compare_e):
if len(e.ops)==1 and e.ops[0] == '==':
if all([type(x) in [name_e,attribute_e] for x in visit(e.left) ]):
if all([type(x) in [name_e,attribute_e] for x in visit(e.comparators[0]) ]):
return True
return False
# Visit all joins that are children of a given join
def visit_joins(j):
yield j
if not isinstance(j['left'],list):
for jj in visit_joins(j['left']):
yield jj
if not isinstance(j['right'],list):
for jj in visit_joins(j['right']):
yield jj
# Main rewriter routine, pushed as much work as possible into the database sources,
# creates joins and join condition and handles query hints.
def rewrite(clauses,visible_vars):
source_id = 0
databases = {}
source_meta = {}
source_clauses = {}
rest_clauses = []
# Have we seen any group-bys in the plan?
# In this case we can't push fors or lets any longer
groupbys_seen = False
# All visible variables at this point in the plan
current_vars = set()
# Live variables - all variables that are needed above in the plan
live_vars = set()
# Variables that were turned into a list by group-by
list_vars = set()
hints = []
join_conds = []
for c in clauses:
# Compute all defined vars
current_vars = current_vars.union(clause_vars([c]))
# If we see a for clause, we try to find its source, if there's none,
# we'll create a new source. We'll then push the for clause into the
# source clause. This can be done for the clauses that don't depend upon
# any variables.
if c['name'] == 'for' and not groupbys_seen:
source = get_ast(c['expr'])
if (len(c['vars']) == 1
and isinstance(source,name_e)
and isinstance(visible_vars.get(source.id), RDBMSTable) ):
database = visible_vars[source.id]
meta = {"type":"database",
"database":database.engine,
"source":database
}
c['database'] = meta
if not database.engine.url in databases:
source_meta[source_id] = meta
databases[database.engine.url] = source_id
source_clauses[source_id] = []
source_id += 1
source_clauses[databases[database.engine.url]].append(c)
# elif (isinstance(source,call_e) and
# isinstance(source.func,name_e) and
# source.func.id == 'Postgres'):
# database = source.args[0]
# if not database in databases:
# databases[database] = source_id
# source_meta[source_id] = {"type":"database", "database":database}
# source_clauses[source_id] = []
# source_id += 1
# source_clauses[databases[database]].append(c)
elif len(c['vars'])==1 and set(get_all_vars(source)).intersection(current_vars) == set():
source_meta[source_id] = {"type":"expr", "expr":c['expr']}
source_clauses[source_id] = [c]
source_id += 1
else:
rest_clauses.append(c)
# We can push let clause into one of the sources, if that can't happen, we'll create
# a new source for it. It can be pushed into an existing source only if it only depends
# upon the variables in only the source or has no dependencies.
elif c['name'] == 'let' and not groupbys_seen:
expr = get_ast(c['expr'])
let_vars = get_all_vars(expr)
srcs = [s for s in range(source_id) if let_vars - clause_vars(source_clauses[s]) == set()]
if len(c['vars'])==1 and srcs:
src = srcs[0]
if source_meta[src]['type'] == 'database':
if source_meta[src]['source'].supports(source_clauses[src],get_ast(c['expr'])):
source_clauses[src].append(c)
else:
rest_clauses.append(c)
else:
source_clauses[src].append(c)
else:
rest_clauses.append(c)
# When we see a group-by, we mark all the variables not in the group-by key as list
# variables. This knowledge will help to figure out whether we can send further clauses
# that depend on these variables to the source.
elif c['name'] == 'groupby':
list_vars = current_vars - {x[1] for x in c['groupby_list']}
rest_clauses.append(c)
groupbys_seen = True
# The where clause is especially important for us, since it includes conditions that
# we can push to the sources, including join conditions, and also various hints.
elif c['name'] == 'where':
expr = get_ast(c['expr'])
exprs = [expr]
if isinstance(expr,boolOp_e) and expr.op == 'and':
exprs = expr.args
remaining_exprs = []
for e in exprs:
# If this is a hint, record the hint
if isinstance(e,call_e) and isinstance(e.func,name_e) and e.func.id=='hint':
hints.append(e)
else:
# If the entire expression can be pushed to a specific source, do so
srcs = [s for s in range(source_id) if get_all_vars(e) - clause_vars(source_clauses[s]) == set()]
if srcs:
src = srcs[0]
if source_meta[src]['type'] == 'database':
if source_meta[src]['source'].supports(source_clauses[src],e):
source_clauses[src].append({'name':'where', 'expr':print_ast(e)})
else:
remaining_exprs.append(e)
else:
source_clauses[src].append({'name':'where', 'expr':print_ast(e)})
# If this looks like a join condition, we'll record it separately. However, we need to
# to check that its a real join condition, i.e. doesn't include a refence from the
# local variables
elif is_join_cond(e) and not get_all_vars(e).intersection(visible_vars):
join_conds.append(e)
else:
remaining_exprs.append(e)
if remaining_exprs:
e = remaining_exprs[0] if len(remaining_exprs)==1 else boolOp_e('and',remaining_exprs)
rest_clauses.append( {'name':'where', 'expr':print_ast(e)})
else:
rest_clauses.append(c)
join = None
# Create a tree of joins if there is more than one source
if len(source_meta)>1:
last_join = {'name':'join',
'left':source_clauses[source_id-2],
'right':source_clauses[source_id-1],
'left_conds':[],
'right_conds':[]}
for s in range(source_id-3,-1,-1):
last_join = {'name':'join',
'left':source_clauses[s],
'right':last_join,
'left_conds':[],
'right_conds':[]}
join = last_join
# Push join condition to the deepest level
for cond in join_conds:
all_cond_vars = get_all_vars(cond)
deepest_join = join
while True:
left_vars = clause_vars(ensure_list(deepest_join['left']))
right_vars = clause_vars(ensure_list(deepest_join['right']))
if all_cond_vars.intersection(left_vars) == all_cond_vars:
deepest_join == deepest_join['left']
elif all_cond_vars.intersection(right_vars) == all_cond_vars:
deepest_join == deepest_join['right']
else:
break
left_cond_vars = clause_vars(ensure_list(deepest_join['left']))
c1 = cond.left
c2 = cond.comparators[0]
if get_all_vars(c1).intersection(left_cond_vars):
deepest_join['left_conds'].append(print_ast(c1))
deepest_join['right_conds'].append(print_ast(c2))
else:
deepest_join['left_conds'].append(print_ast(c2))
deepest_join['right_conds'].append(print_ast(c1))
# Push in join hints to the level with join conditions
for hint in hints:
join_type = hint.args[0].value
left_var = hint.args[1].value
right_var = hint.args[2].value
for j in visit_joins(join):
l_vars = clause_vars(ensure_list(j['left']))
r_vars = clause_vars(ensure_list(j['right']))
if left_var in l_vars and right_var in r_vars:
j['hint'] = {'join_type':join_type, 'dir':'right'}
elif left_var in r_vars and right_var in l_vars:
j['hint'] = {'join_type':join_type, 'dir':'left'}
# Iterate over the database sources and translate the queries into
# database-specific dialects
for db in databases:
src_id = databases[db]
src_clauses = source_clauses[src_id]
src_meta = source_meta[src_id]
# Compute the project list
live_var_list = clause_live_vars(rest_clauses)
if join:
live_var_list = live_var_list.union(clause_live_vars([join]))
vars = clause_vars(src_clauses)
project_list = live_var_list.intersection(vars)
wrapped = src_meta['source'].wrap(src_clauses,project_list)
src_clauses.clear()
src_clauses.append( wrapped )
if join:
return [join] + rest_clauses
else:
return ( source_clauses[0] if source_clauses else [] ) + rest_clauses