sqlglot.optimizer.pushdown_predicates

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