Repository navigation
Expand file tree
/
Copy pathnormalize.py
More file actions
228 lines (183 loc) · 7.79 KB
/
Copy pathnormalize.py
File metadata and controls
228 lines (183 loc) · 7.79 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
from __future__ import annotations
from collections.abc import Callable, Iterator
import logging
from sqlglot import exp
from sqlglot.errors import OptimizeError
from sqlglot.helper import while_changing
from sqlglot.optimizer.scope import find_all_in_scope
from sqlglot.optimizer.simplify import Simplifier, flatten
logger = logging.getLogger("sqlglot")
def normalize(expression: exp.Expr, dnf: bool = False, max_distance: int = 128) -> exp.Expr:
"""
Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.
Example:
>>> import sqlglot
>>> expression = sqlglot.parse_one("(x AND y) OR z")
>>> normalize(expression, dnf=False).sql()
'(x OR z) AND (y OR z)'
Args:
expression: expression to normalize
dnf: rewrite in disjunctive normal form instead.
max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion
Returns:
sqlglot.Expr: normalized expression
"""
simplifier = Simplifier(annotate_new_expressions=False)
for node in tuple(expression.walk(prune=lambda e: isinstance(e, exp.Connector))):
if isinstance(node, exp.Connector):
if normalized(node, dnf=dnf):
continue
root = node is expression
original = node.copy()
node.transform(simplifier.rewrite_between, copy=False)
distance = normalization_distance(node, dnf=dnf, max_=max_distance)
if distance > max_distance:
logger.info(
f"Skipping normalization because distance {distance} exceeds max {max_distance}"
)
return expression
try:
node = node.replace(
while_changing(
node,
lambda e: distributive_law(e, dnf, max_distance, simplifier=simplifier),
)
)
except OptimizeError as e:
logger.info(e)
node.replace(original)
if root:
return original
return expression
if root:
expression = node
return expression
def normalized(expression: exp.Expr, dnf: bool = False) -> bool:
"""
Checks whether a given expression is in a normal form of interest.
Example:
>>> from sqlglot import parse_one
>>> normalized(parse_one("(a AND b) OR c OR (d AND e)"), dnf=True)
True
>>> normalized(parse_one("(a OR b) AND c")) # Checks CNF by default
True
>>> normalized(parse_one("a AND (b OR c)"), dnf=True)
False
Args:
expression: The expression to check if it's normalized.
dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF).
Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF).
"""
ancestor, root = (exp.And, exp.Or) if dnf else (exp.Or, exp.And)
return not any(
connector.find_ancestor(ancestor) for connector in find_all_in_scope(expression, root)
)
def normalization_distance(
expression: exp.Expr, dnf: bool = False, max_: float = float("inf")
) -> int:
"""
The difference in the number of predicates between a given expression and its normalized form.
This is used as an estimate of the cost of the conversion which is exponential in complexity.
Example:
>>> import sqlglot
>>> expression = sqlglot.parse_one("(a AND b) OR (c AND d)")
>>> normalization_distance(expression)
4
Args:
expression: The expression to compute the normalization distance for.
dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF).
Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF).
max_: stop early if count exceeds this.
Returns:
The normalization distance.
"""
total = -(sum(1 for _ in expression.find_all(exp.Connector)) + 1)
for length in _predicate_lengths(expression, dnf, max_):
total += length
if total > max_:
return total
return total
def _predicate_lengths(
expression: exp.Expr, dnf: bool, max_: float = float("inf"), depth: int = 0
) -> Iterator[int]:
"""
Returns a list of predicate lengths when expanded to normalized form.
(A AND B) OR C -> [2, 2] because len(A OR C), len(B OR C).
"""
if depth > max_:
yield depth
return
expression = expression.unnest()
if not isinstance(expression, exp.Connector):
yield 1
return
depth += 1
left, right = expression.left, expression.right
if isinstance(expression, exp.And if dnf else exp.Or):
for a in _predicate_lengths(left, dnf, max_, depth):
for b in _predicate_lengths(right, dnf, max_, depth):
yield a + b
else:
yield from _predicate_lengths(left, dnf, max_, depth)
yield from _predicate_lengths(right, dnf, max_, depth)
def distributive_law(
expression: exp.Expr, dnf: bool, max_distance: float, simplifier: Simplifier | None = None
):
"""
x OR (y AND z) -> (x OR y) AND (x OR z)
(x AND y) OR (y AND z) -> (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)
"""
if normalized(expression, dnf=dnf):
return expression
distance = normalization_distance(expression, dnf=dnf, max_=max_distance)
if distance > max_distance:
raise OptimizeError(f"Normalization distance {distance} exceeds max {max_distance}")
exp.replace_children(expression, lambda e: distributive_law(e, dnf, max_distance))
to_exp, from_exp = (exp.Or, exp.And) if dnf else (exp.And, exp.Or)
if isinstance(expression, from_exp):
a, b = expression.unnest_operands()
from_func = exp.and_ if from_exp == exp.And else exp.or_
to_func = exp.and_ if to_exp == exp.And else exp.or_
simplifier = simplifier or Simplifier(annotate_new_expressions=False)
if isinstance(a, to_exp) and isinstance(b, to_exp):
if len(tuple(a.find_all(exp.Connector))) > len(tuple(b.find_all(exp.Connector))):
return _distribute(a, b, from_func, to_func, simplifier)
return _distribute(b, a, from_func, to_func, simplifier)
if isinstance(a, to_exp):
return _distribute(b, a, from_func, to_func, simplifier)
if isinstance(b, to_exp):
return _distribute(a, b, from_func, to_func, simplifier)
return expression
def _distribute(
a,
b,
from_func: Callable[..., exp.Condition],
to_func: Callable[..., exp.Condition],
simplifier: Simplifier,
):
# When `a` and `b` are connectors of the SAME polarity (e.g. both AND in
# CNF mode), `b` is distributed across `a`'s children so the existing
# cross-product handling can simplify them.
if isinstance(a, exp.Connector) and isinstance(a, type(b)):
exp.replace_children(
a,
lambda c: to_func(
simplifier.uniq_sort(flatten(from_func(c, b.left))),
simplifier.uniq_sort(flatten(from_func(c, b.right))),
copy=False,
),
)
return a
# Otherwise apply the textbook rule
# ``a OR (b.left AND b.right) -> (a OR b.left) AND (a OR b.right)`` (or
# the AND/OR dual). When `a` is a connector of the OPPOSITE polarity to
# `b` (e.g. `Or(...) ∨ And(...)` in CNF mode) the previous behaviour was
# to push `b` into each child of `a`, leaving an intermediate
# ``Or(And, And)`` that `while_changing` would later cross-product into
# redundant superset and duplicate clauses. Returning the direct rewrite
# avoids that detour.
return to_func(
simplifier.uniq_sort(flatten(from_func(a, b.left))),
simplifier.uniq_sort(flatten(from_func(a, b.right))),
copy=False,
)