sqlglot.optimizer.simplify

  1import datetime
  2import functools
  3import itertools
  4from collections import deque
  5from decimal import Decimal
  6
  7from sqlglot import exp
  8from sqlglot.expressions import FALSE, NULL, TRUE
  9from sqlglot.generator import Generator
 10from sqlglot.helper import first, while_changing
 11
 12GENERATOR = Generator(normalize=True, identify=True)
 13
 14
 15def simplify(expression):
 16    """
 17    Rewrite sqlglot AST to simplify expressions.
 18
 19    Example:
 20        >>> import sqlglot
 21        >>> expression = sqlglot.parse_one("TRUE AND TRUE")
 22        >>> simplify(expression).sql()
 23        'TRUE'
 24
 25    Args:
 26        expression (sqlglot.Expression): expression to simplify
 27    Returns:
 28        sqlglot.Expression: simplified expression
 29    """
 30
 31    def _simplify(expression, root=True):
 32        node = expression
 33        node = rewrite_between(node)
 34        node = uniq_sort(node)
 35        node = absorb_and_eliminate(node)
 36        exp.replace_children(node, lambda e: _simplify(e, False))
 37        node = simplify_not(node)
 38        node = flatten(node)
 39        node = simplify_connectors(node)
 40        node = remove_compliments(node)
 41        node.parent = expression.parent
 42        node = simplify_literals(node)
 43        node = simplify_parens(node)
 44        if root:
 45            expression.replace(node)
 46        return node
 47
 48    expression = while_changing(expression, _simplify)
 49    remove_where_true(expression)
 50    return expression
 51
 52
 53def rewrite_between(expression: exp.Expression) -> exp.Expression:
 54    """Rewrite x between y and z to x >= y AND x <= z.
 55
 56    This is done because comparison simplification is only done on lt/lte/gt/gte.
 57    """
 58    if isinstance(expression, exp.Between):
 59        return exp.and_(
 60            exp.GTE(this=expression.this.copy(), expression=expression.args["low"]),
 61            exp.LTE(this=expression.this.copy(), expression=expression.args["high"]),
 62        )
 63    return expression
 64
 65
 66def simplify_not(expression):
 67    """
 68    Demorgan's Law
 69    NOT (x OR y) -> NOT x AND NOT y
 70    NOT (x AND y) -> NOT x OR NOT y
 71    """
 72    if isinstance(expression, exp.Not):
 73        if isinstance(expression.this, exp.Null):
 74            return exp.null()
 75        if isinstance(expression.this, exp.Paren):
 76            condition = expression.this.unnest()
 77            if isinstance(condition, exp.And):
 78                return exp.or_(exp.not_(condition.left), exp.not_(condition.right))
 79            if isinstance(condition, exp.Or):
 80                return exp.and_(exp.not_(condition.left), exp.not_(condition.right))
 81            if isinstance(condition, exp.Null):
 82                return exp.null()
 83        if always_true(expression.this):
 84            return exp.false()
 85        if expression.this == FALSE:
 86            return exp.true()
 87        if isinstance(expression.this, exp.Not):
 88            # double negation
 89            # NOT NOT x -> x
 90            return expression.this.this
 91    return expression
 92
 93
 94def flatten(expression):
 95    """
 96    A AND (B AND C) -> A AND B AND C
 97    A OR (B OR C) -> A OR B OR C
 98    """
 99    if isinstance(expression, exp.Connector):
100        for node in expression.args.values():
101            child = node.unnest()
102            if isinstance(child, expression.__class__):
103                node.replace(child)
104    return expression
105
106
107def simplify_connectors(expression):
108    def _simplify_connectors(expression, left, right):
109        if isinstance(expression, exp.Connector):
110            if left == right:
111                return left
112            if isinstance(expression, exp.And):
113                if FALSE in (left, right):
114                    return exp.false()
115                if NULL in (left, right):
116                    return exp.null()
117                if always_true(left) and always_true(right):
118                    return exp.true()
119                if always_true(left):
120                    return right
121                if always_true(right):
122                    return left
123                return _simplify_comparison(expression, left, right)
124            elif isinstance(expression, exp.Or):
125                if always_true(left) or always_true(right):
126                    return exp.true()
127                if left == FALSE and right == FALSE:
128                    return exp.false()
129                if (
130                    (left == NULL and right == NULL)
131                    or (left == NULL and right == FALSE)
132                    or (left == FALSE and right == NULL)
133                ):
134                    return exp.null()
135                if left == FALSE:
136                    return right
137                if right == FALSE:
138                    return left
139                return _simplify_comparison(expression, left, right, or_=True)
140        return None
141
142    return _flat_simplify(expression, _simplify_connectors)
143
144
145LT_LTE = (exp.LT, exp.LTE)
146GT_GTE = (exp.GT, exp.GTE)
147
148COMPARISONS = (
149    *LT_LTE,
150    *GT_GTE,
151    exp.EQ,
152    exp.NEQ,
153)
154
155INVERSE_COMPARISONS = {
156    exp.LT: exp.GT,
157    exp.GT: exp.LT,
158    exp.LTE: exp.GTE,
159    exp.GTE: exp.LTE,
160}
161
162
163def _simplify_comparison(expression, left, right, or_=False):
164    if isinstance(left, COMPARISONS) and isinstance(right, COMPARISONS):
165        ll, lr = left.args.values()
166        rl, rr = right.args.values()
167
168        largs = {ll, lr}
169        rargs = {rl, rr}
170
171        matching = largs & rargs
172        columns = {m for m in matching if isinstance(m, exp.Column)}
173
174        if matching and columns:
175            try:
176                l = first(largs - columns)
177                r = first(rargs - columns)
178            except StopIteration:
179                return expression
180
181            # make sure the comparison is always of the form x > 1 instead of 1 < x
182            if left.__class__ in INVERSE_COMPARISONS and l == ll:
183                left = INVERSE_COMPARISONS[left.__class__](this=lr, expression=ll)
184            if right.__class__ in INVERSE_COMPARISONS and r == rl:
185                right = INVERSE_COMPARISONS[right.__class__](this=rr, expression=rl)
186
187            if l.is_number and r.is_number:
188                l = float(l.name)
189                r = float(r.name)
190            elif l.is_string and r.is_string:
191                l = l.name
192                r = r.name
193            else:
194                return None
195
196            for (a, av), (b, bv) in itertools.permutations(((left, l), (right, r))):
197                if isinstance(a, LT_LTE) and isinstance(b, LT_LTE):
198                    return left if (av > bv if or_ else av <= bv) else right
199                if isinstance(a, GT_GTE) and isinstance(b, GT_GTE):
200                    return left if (av < bv if or_ else av >= bv) else right
201
202                # we can't ever shortcut to true because the column could be null
203                if isinstance(a, exp.LT) and isinstance(b, GT_GTE):
204                    if not or_ and av <= bv:
205                        return exp.false()
206                elif isinstance(a, exp.GT) and isinstance(b, LT_LTE):
207                    if not or_ and av >= bv:
208                        return exp.false()
209                elif isinstance(a, exp.EQ):
210                    if isinstance(b, exp.LT):
211                        return exp.false() if av >= bv else a
212                    if isinstance(b, exp.LTE):
213                        return exp.false() if av > bv else a
214                    if isinstance(b, exp.GT):
215                        return exp.false() if av <= bv else a
216                    if isinstance(b, exp.GTE):
217                        return exp.false() if av < bv else a
218                    if isinstance(b, exp.NEQ):
219                        return exp.false() if av == bv else a
220    return None
221
222
223def remove_compliments(expression):
224    """
225    Removing compliments.
226
227    A AND NOT A -> FALSE
228    A OR NOT A -> TRUE
229    """
230    if isinstance(expression, exp.Connector):
231        compliment = exp.false() if isinstance(expression, exp.And) else exp.true()
232
233        for a, b in itertools.permutations(expression.flatten(), 2):
234            if is_complement(a, b):
235                return compliment
236    return expression
237
238
239def uniq_sort(expression):
240    """
241    Uniq and sort a connector.
242
243    C AND A AND B AND B -> A AND B AND C
244    """
245    if isinstance(expression, exp.Connector):
246        result_func = exp.and_ if isinstance(expression, exp.And) else exp.or_
247        flattened = tuple(expression.flatten())
248        deduped = {GENERATOR.generate(e): e for e in flattened}
249        arr = tuple(deduped.items())
250
251        # check if the operands are already sorted, if not sort them
252        # A AND C AND B -> A AND B AND C
253        for i, (sql, e) in enumerate(arr[1:]):
254            if sql < arr[i][0]:
255                expression = result_func(*(deduped[sql] for sql in sorted(deduped)))
256                break
257        else:
258            # we didn't have to sort but maybe we need to dedup
259            if len(deduped) < len(flattened):
260                expression = result_func(*deduped.values())
261
262    return expression
263
264
265def absorb_and_eliminate(expression):
266    """
267    absorption:
268        A AND (A OR B) -> A
269        A OR (A AND B) -> A
270        A AND (NOT A OR B) -> A AND B
271        A OR (NOT A AND B) -> A OR B
272    elimination:
273        (A AND B) OR (A AND NOT B) -> A
274        (A OR B) AND (A OR NOT B) -> A
275    """
276    if isinstance(expression, exp.Connector):
277        kind = exp.Or if isinstance(expression, exp.And) else exp.And
278
279        for a, b in itertools.permutations(expression.flatten(), 2):
280            if isinstance(a, kind):
281                aa, ab = a.unnest_operands()
282
283                # absorb
284                if is_complement(b, aa):
285                    aa.replace(exp.true() if kind == exp.And else exp.false())
286                elif is_complement(b, ab):
287                    ab.replace(exp.true() if kind == exp.And else exp.false())
288                elif (set(b.flatten()) if isinstance(b, kind) else {b}) < set(a.flatten()):
289                    a.replace(exp.false() if kind == exp.And else exp.true())
290                elif isinstance(b, kind):
291                    # eliminate
292                    rhs = b.unnest_operands()
293                    ba, bb = rhs
294
295                    if aa in rhs and (is_complement(ab, ba) or is_complement(ab, bb)):
296                        a.replace(aa)
297                        b.replace(aa)
298                    elif ab in rhs and (is_complement(aa, ba) or is_complement(aa, bb)):
299                        a.replace(ab)
300                        b.replace(ab)
301
302    return expression
303
304
305def simplify_literals(expression):
306    if isinstance(expression, exp.Binary):
307        return _flat_simplify(expression, _simplify_binary)
308    elif isinstance(expression, exp.Neg):
309        this = expression.this
310        if this.is_number:
311            value = this.name
312            if value[0] == "-":
313                return exp.Literal.number(value[1:])
314            return exp.Literal.number(f"-{value}")
315
316    return expression
317
318
319def _simplify_binary(expression, a, b):
320    if isinstance(expression, exp.Is):
321        if isinstance(b, exp.Not):
322            c = b.this
323            not_ = True
324        else:
325            c = b
326            not_ = False
327
328        if c == NULL:
329            if isinstance(a, exp.Literal):
330                return exp.true() if not_ else exp.false()
331            if a == NULL:
332                return exp.false() if not_ else exp.true()
333    elif isinstance(expression, (exp.NullSafeEQ, exp.NullSafeNEQ)):
334        return None
335    elif NULL in (a, b):
336        return exp.null()
337
338    if a.is_number and b.is_number:
339        a = int(a.name) if a.is_int else Decimal(a.name)
340        b = int(b.name) if b.is_int else Decimal(b.name)
341
342        if isinstance(expression, exp.Add):
343            return exp.Literal.number(a + b)
344        if isinstance(expression, exp.Sub):
345            return exp.Literal.number(a - b)
346        if isinstance(expression, exp.Mul):
347            return exp.Literal.number(a * b)
348        if isinstance(expression, exp.Div):
349            if isinstance(a, int) and isinstance(b, int):
350                return exp.Literal.number(a // b)
351            return exp.Literal.number(a / b)
352
353        boolean = eval_boolean(expression, a, b)
354
355        if boolean:
356            return boolean
357    elif a.is_string and b.is_string:
358        boolean = eval_boolean(expression, a, b)
359
360        if boolean:
361            return boolean
362    elif isinstance(a, exp.Cast) and isinstance(b, exp.Interval):
363        a, b = extract_date(a), extract_interval(b)
364        if a and b:
365            if isinstance(expression, exp.Add):
366                return date_literal(a + b)
367            if isinstance(expression, exp.Sub):
368                return date_literal(a - b)
369    elif isinstance(a, exp.Interval) and isinstance(b, exp.Cast):
370        a, b = extract_interval(a), extract_date(b)
371        # you cannot subtract a date from an interval
372        if a and b and isinstance(expression, exp.Add):
373            return date_literal(a + b)
374
375    return None
376
377
378def simplify_parens(expression):
379    if (
380        isinstance(expression, exp.Paren)
381        and not isinstance(expression.this, exp.Select)
382        and (
383            not isinstance(expression.parent, (exp.Condition, exp.Binary))
384            or isinstance(expression.this, (exp.Is, exp.Like))
385            or not isinstance(expression.this, exp.Binary)
386        )
387    ):
388        return expression.this
389    return expression
390
391
392def remove_where_true(expression):
393    for where in expression.find_all(exp.Where):
394        if always_true(where.this):
395            where.parent.set("where", None)
396    for join in expression.find_all(exp.Join):
397        if always_true(join.args.get("on")):
398            join.set("kind", "CROSS")
399            join.set("on", None)
400
401
402def always_true(expression):
403    return expression == TRUE or isinstance(expression, exp.Literal)
404
405
406def is_complement(a, b):
407    return isinstance(b, exp.Not) and b.this == a
408
409
410def eval_boolean(expression, a, b):
411    if isinstance(expression, (exp.EQ, exp.Is)):
412        return boolean_literal(a == b)
413    if isinstance(expression, exp.NEQ):
414        return boolean_literal(a != b)
415    if isinstance(expression, exp.GT):
416        return boolean_literal(a > b)
417    if isinstance(expression, exp.GTE):
418        return boolean_literal(a >= b)
419    if isinstance(expression, exp.LT):
420        return boolean_literal(a < b)
421    if isinstance(expression, exp.LTE):
422        return boolean_literal(a <= b)
423    return None
424
425
426def extract_date(cast):
427    # The "fromisoformat" conversion could fail if the cast is used on an identifier,
428    # so in that case we can't extract the date.
429    try:
430        if cast.args["to"].this == exp.DataType.Type.DATE:
431            return datetime.date.fromisoformat(cast.name)
432        if cast.args["to"].this == exp.DataType.Type.DATETIME:
433            return datetime.datetime.fromisoformat(cast.name)
434    except ValueError:
435        return None
436
437
438def extract_interval(interval):
439    try:
440        from dateutil.relativedelta import relativedelta  # type: ignore
441    except ModuleNotFoundError:
442        return None
443
444    n = int(interval.name)
445    unit = interval.text("unit").lower()
446
447    if unit == "year":
448        return relativedelta(years=n)
449    if unit == "month":
450        return relativedelta(months=n)
451    if unit == "week":
452        return relativedelta(weeks=n)
453    if unit == "day":
454        return relativedelta(days=n)
455    return None
456
457
458def date_literal(date):
459    return exp.cast(
460        exp.Literal.string(date),
461        "DATETIME" if isinstance(date, datetime.datetime) else "DATE",
462    )
463
464
465def boolean_literal(condition):
466    return exp.true() if condition else exp.false()
467
468
469def _flat_simplify(expression, simplifier):
470    operands = []
471    queue = deque(expression.flatten(unnest=False))
472    size = len(queue)
473
474    while queue:
475        a = queue.popleft()
476
477        for b in queue:
478            result = simplifier(expression, a, b)
479
480            if result:
481                queue.remove(b)
482                queue.append(result)
483                break
484        else:
485            operands.append(a)
486
487    if len(operands) < size:
488        return functools.reduce(lambda a, b: expression.__class__(this=a, expression=b), operands)
489    return expression