sqlglot.optimizer.simplify

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