sqlglot.optimizer.normalize

  1from __future__ import annotations
  2
  3import logging
  4
  5from sqlglot import exp
  6from sqlglot.errors import OptimizeError
  7from sqlglot.helper import while_changing
  8from sqlglot.optimizer.scope import find_all_in_scope
  9from sqlglot.optimizer.simplify import flatten, rewrite_between, uniq_sort
 10
 11logger = logging.getLogger("sqlglot")
 12
 13
 14def normalize(expression: exp.Expression, dnf: bool = False, max_distance: int = 128):
 15    """
 16    Rewrite sqlglot AST into conjunctive normal form or disjunctive normal form.
 17
 18    Example:
 19        >>> import sqlglot
 20        >>> expression = sqlglot.parse_one("(x AND y) OR z")
 21        >>> normalize(expression, dnf=False).sql()
 22        '(x OR z) AND (y OR z)'
 23
 24    Args:
 25        expression: expression to normalize
 26        dnf: rewrite in disjunctive normal form instead.
 27        max_distance (int): the maximal estimated distance from cnf/dnf to attempt conversion
 28    Returns:
 29        sqlglot.Expression: normalized expression
 30    """
 31    for node in tuple(expression.walk(prune=lambda e: isinstance(e, exp.Connector))):
 32        if isinstance(node, exp.Connector):
 33            if normalized(node, dnf=dnf):
 34                continue
 35            root = node is expression
 36            original = node.copy()
 37
 38            node.transform(rewrite_between, copy=False)
 39            distance = normalization_distance(node, dnf=dnf, max_=max_distance)
 40
 41            if distance > max_distance:
 42                logger.info(
 43                    f"Skipping normalization because distance {distance} exceeds max {max_distance}"
 44                )
 45                return expression
 46
 47            try:
 48                node = node.replace(
 49                    while_changing(node, lambda e: distributive_law(e, dnf, max_distance))
 50                )
 51            except OptimizeError as e:
 52                logger.info(e)
 53                node.replace(original)
 54                if root:
 55                    return original
 56                return expression
 57
 58            if root:
 59                expression = node
 60
 61    return expression
 62
 63
 64def normalized(expression: exp.Expression, dnf: bool = False) -> bool:
 65    """
 66    Checks whether a given expression is in a normal form of interest.
 67
 68    Example:
 69        >>> from sqlglot import parse_one
 70        >>> normalized(parse_one("(a AND b) OR c OR (d AND e)"), dnf=True)
 71        True
 72        >>> normalized(parse_one("(a OR b) AND c"))  # Checks CNF by default
 73        True
 74        >>> normalized(parse_one("a AND (b OR c)"), dnf=True)
 75        False
 76
 77    Args:
 78        expression: The expression to check if it's normalized.
 79        dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF).
 80            Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF).
 81    """
 82    ancestor, root = (exp.And, exp.Or) if dnf else (exp.Or, exp.And)
 83    return not any(
 84        connector.find_ancestor(ancestor) for connector in find_all_in_scope(expression, root)
 85    )
 86
 87
 88def normalization_distance(
 89    expression: exp.Expression, dnf: bool = False, max_: float = float("inf")
 90) -> int:
 91    """
 92    The difference in the number of predicates between a given expression and its normalized form.
 93
 94    This is used as an estimate of the cost of the conversion which is exponential in complexity.
 95
 96    Example:
 97        >>> import sqlglot
 98        >>> expression = sqlglot.parse_one("(a AND b) OR (c AND d)")
 99        >>> normalization_distance(expression)
100        4
101
102    Args:
103        expression: The expression to compute the normalization distance for.
104        dnf: Whether to check if the expression is in Disjunctive Normal Form (DNF).
105            Default: False, i.e. we check if it's in Conjunctive Normal Form (CNF).
106        max_: stop early if count exceeds this.
107
108    Returns:
109        The normalization distance.
110    """
111    total = -(sum(1 for _ in expression.find_all(exp.Connector)) + 1)
112
113    for length in _predicate_lengths(expression, dnf, max_):
114        total += length
115        if total > max_:
116            return total
117
118    return total
119
120
121def _predicate_lengths(expression, dnf, max_=float("inf"), depth=0):
122    """
123    Returns a list of predicate lengths when expanded to normalized form.
124
125    (A AND B) OR C -> [2, 2] because len(A OR C), len(B OR C).
126    """
127    if depth > max_:
128        yield depth
129        return
130
131    expression = expression.unnest()
132
133    if not isinstance(expression, exp.Connector):
134        yield 1
135        return
136
137    depth += 1
138    left, right = expression.args.values()
139
140    if isinstance(expression, exp.And if dnf else exp.Or):
141        for a in _predicate_lengths(left, dnf, max_, depth):
142            for b in _predicate_lengths(right, dnf, max_, depth):
143                yield a + b
144    else:
145        yield from _predicate_lengths(left, dnf, max_, depth)
146        yield from _predicate_lengths(right, dnf, max_, depth)
147
148
149def distributive_law(expression, dnf, max_distance):
150    """
151    x OR (y AND z) -> (x OR y) AND (x OR z)
152    (x AND y) OR (y AND z) -> (x OR y) AND (x OR z) AND (y OR y) AND (y OR z)
153    """
154    if normalized(expression, dnf=dnf):
155        return expression
156
157    distance = normalization_distance(expression, dnf=dnf, max_=max_distance)
158
159    if distance > max_distance:
160        raise OptimizeError(f"Normalization distance {distance} exceeds max {max_distance}")
161
162    exp.replace_children(expression, lambda e: distributive_law(e, dnf, max_distance))
163    to_exp, from_exp = (exp.Or, exp.And) if dnf else (exp.And, exp.Or)
164
165    if isinstance(expression, from_exp):
166        a, b = expression.unnest_operands()
167
168        from_func = exp.and_ if from_exp == exp.And else exp.or_
169        to_func = exp.and_ if to_exp == exp.And else exp.or_
170
171        if isinstance(a, to_exp) and isinstance(b, to_exp):
172            if len(tuple(a.find_all(exp.Connector))) > len(tuple(b.find_all(exp.Connector))):
173                return _distribute(a, b, from_func, to_func)
174            return _distribute(b, a, from_func, to_func)
175        if isinstance(a, to_exp):
176            return _distribute(b, a, from_func, to_func)
177        if isinstance(b, to_exp):
178            return _distribute(a, b, from_func, to_func)
179
180    return expression
181
182
183def _distribute(a, b, from_func, to_func):
184    if isinstance(a, exp.Connector):
185        exp.replace_children(
186            a,
187            lambda c: to_func(
188                uniq_sort(flatten(from_func(c, b.left))),
189                uniq_sort(flatten(from_func(c, b.right))),
190                copy=False,
191            ),
192        )
193    else:
194        a = to_func(
195            uniq_sort(flatten(from_func(a, b.left))),
196            uniq_sort(flatten(from_func(a, b.right))),
197            copy=False,
198        )
199
200    return a