sqlglot.optimizer.pushdown_predicates

  1from sqlglot import exp
  2from sqlglot.optimizer.normalize import normalized
  3from sqlglot.optimizer.scope import build_scope
  4from sqlglot.optimizer.simplify import simplify
  5
  6
  7def pushdown_predicates(expression):
  8    """
  9    Rewrite sqlglot AST to pushdown predicates in FROMS and JOINS
 10
 11    Example:
 12        >>> import sqlglot
 13        >>> sql = "SELECT y.a AS a FROM (SELECT x.a AS a FROM x AS x) AS y WHERE y.a = 1"
 14        >>> expression = sqlglot.parse_one(sql)
 15        >>> pushdown_predicates(expression).sql()
 16        'SELECT y.a AS a FROM (SELECT x.a AS a FROM x AS x WHERE x.a = 1) AS y WHERE TRUE'
 17
 18    Args:
 19        expression (sqlglot.Expression): expression to optimize
 20    Returns:
 21        sqlglot.Expression: optimized expression
 22    """
 23    root = build_scope(expression)
 24
 25    if root:
 26        scope_ref_count = root.ref_count()
 27
 28        for scope in reversed(list(root.traverse())):
 29            select = scope.expression
 30            where = select.args.get("where")
 31            if where:
 32                selected_sources = scope.selected_sources
 33                # a right join can only push down to itself and not the source FROM table
 34                for k, (node, source) in selected_sources.items():
 35                    parent = node.find_ancestor(exp.Join, exp.From)
 36                    if isinstance(parent, exp.Join) and parent.side == "RIGHT":
 37                        selected_sources = {k: (node, source)}
 38                        break
 39                pushdown(where.this, selected_sources, scope_ref_count)
 40
 41            # joins should only pushdown into itself, not to other joins
 42            # so we limit the selected sources to only itself
 43            for join in select.args.get("joins") or []:
 44                name = join.alias_or_name
 45                if name in scope.selected_sources:
 46                    pushdown(
 47                        join.args.get("on"), {name: scope.selected_sources[name]}, scope_ref_count
 48                    )
 49
 50    return expression
 51
 52
 53def pushdown(condition, sources, scope_ref_count):
 54    if not condition:
 55        return
 56
 57    condition = condition.replace(simplify(condition))
 58    cnf_like = normalized(condition) or not normalized(condition, dnf=True)
 59
 60    predicates = list(
 61        condition.flatten()
 62        if isinstance(condition, exp.And if cnf_like else exp.Or)
 63        else [condition]
 64    )
 65
 66    if cnf_like:
 67        pushdown_cnf(predicates, sources, scope_ref_count)
 68    else:
 69        pushdown_dnf(predicates, sources, scope_ref_count)
 70
 71
 72def pushdown_cnf(predicates, scope, scope_ref_count):
 73    """
 74    If the predicates are in CNF like form, we can simply replace each block in the parent.
 75    """
 76    for predicate in predicates:
 77        for node in nodes_for_predicate(predicate, scope, scope_ref_count).values():
 78            if isinstance(node, exp.Join):
 79                predicate.replace(exp.true())
 80                node.on(predicate, copy=False)
 81                break
 82            if isinstance(node, exp.Select):
 83                predicate.replace(exp.true())
 84                node.where(replace_aliases(node, predicate), copy=False)
 85
 86
 87def pushdown_dnf(predicates, scope, scope_ref_count):
 88    """
 89    If the predicates are in DNF form, we can only push down conditions that are in all blocks.
 90    Additionally, we can't remove predicates from their original form.
 91    """
 92    # find all the tables that can be pushdown too
 93    # these are tables that are referenced in all blocks of a DNF
 94    # (a.x AND b.x) OR (a.y AND c.y)
 95    # only table a can be push down
 96    pushdown_tables = set()
 97
 98    for a in predicates:
 99        a_tables = exp.column_table_names(a)
100
101        for b in predicates:
102            a_tables &= exp.column_table_names(b)
103
104        pushdown_tables.update(a_tables)
105
106    conditions = {}
107
108    # for every pushdown table, find all related conditions in all predicates
109    # combine them with ORS
110    # (a.x AND and a.y AND b.x) OR (a.z AND c.y) -> (a.x AND a.y) OR (a.z)
111    for table in sorted(pushdown_tables):
112        for predicate in predicates:
113            nodes = nodes_for_predicate(predicate, scope, scope_ref_count)
114
115            if table not in nodes:
116                continue
117
118            predicate_condition = None
119
120            for column in predicate.find_all(exp.Column):
121                if column.table == table:
122                    condition = column.find_ancestor(exp.Condition)
123                    predicate_condition = (
124                        exp.and_(predicate_condition, condition)
125                        if predicate_condition
126                        else condition
127                    )
128
129            if predicate_condition:
130                conditions[table] = (
131                    exp.or_(conditions[table], predicate_condition)
132                    if table in conditions
133                    else predicate_condition
134                )
135
136        for name, node in nodes.items():
137            if name not in conditions:
138                continue
139
140            predicate = conditions[name]
141
142            if isinstance(node, exp.Join):
143                node.on(predicate, copy=False)
144            elif isinstance(node, exp.Select):
145                node.where(replace_aliases(node, predicate), copy=False)
146
147
148def nodes_for_predicate(predicate, sources, scope_ref_count):
149    nodes = {}
150    tables = exp.column_table_names(predicate)
151    where_condition = isinstance(predicate.find_ancestor(exp.Join, exp.Where), exp.Where)
152
153    for table in sorted(tables):
154        node, source = sources.get(table) or (None, None)
155
156        # if the predicate is in a where statement we can try to push it down
157        # we want to find the root join or from statement
158        if node and where_condition:
159            node = node.find_ancestor(exp.Join, exp.From)
160
161        # a node can reference a CTE which should be pushed down
162        if isinstance(node, exp.From) and not isinstance(source, exp.Table):
163            with_ = source.parent.expression.args.get("with")
164            if with_ and with_.recursive:
165                return {}
166            node = source.expression
167
168        if isinstance(node, exp.Join):
169            if node.side and node.side != "RIGHT":
170                return {}
171            nodes[table] = node
172        elif isinstance(node, exp.Select) and len(tables) == 1:
173            # We can't push down window expressions
174            has_window_expression = any(
175                select for select in node.selects if select.find(exp.Window)
176            )
177            # we can't push down predicates to select statements if they are referenced in
178            # multiple places.
179            if (
180                not node.args.get("group")
181                and scope_ref_count[id(source)] < 2
182                and not has_window_expression
183            ):
184                nodes[table] = node
185    return nodes
186
187
188def replace_aliases(source, predicate):
189    aliases = {}
190
191    for select in source.selects:
192        if isinstance(select, exp.Alias):
193            aliases[select.alias] = select.this
194        else:
195            aliases[select.name] = select
196
197    def _replace_alias(column):
198        if isinstance(column, exp.Column) and column.name in aliases:
199            return aliases[column.name].copy()
200        return column
201
202    return predicate.transform(_replace_alias)