sqlglot.planner

  1from __future__ import annotations
  2
  3import math
  4import typing as t
  5
  6from sqlglot import alias, exp
  7from sqlglot.helper import name_sequence
  8from sqlglot.optimizer.eliminate_joins import join_condition
  9from sqlglot.optimizer.scope import find_all_in_scope, find_in_scope
 10from collections.abc import Iterator, Sequence, Iterable
 11
 12
 13class Plan:
 14    def __init__(self, expression: exp.Expr) -> None:
 15        self.expression: exp.Expr = expression.copy()
 16        self.root: Step = Step.from_expression(self.expression)
 17        self._dag: dict[Step, set[Step]] = {}
 18
 19    @property
 20    def dag(self) -> dict[Step, set[Step]]:
 21        if not self._dag:
 22            dag: dict[Step, set[Step]] = {}
 23            nodes = {self.root}
 24
 25            while nodes:
 26                node = nodes.pop()
 27                dag[node] = set()
 28
 29                for dep in node.dependencies:
 30                    dag[node].add(dep)
 31                    nodes.add(dep)
 32
 33            self._dag = dag
 34
 35        return self._dag
 36
 37    @property
 38    def leaves(self) -> Iterator[Step]:
 39        return (node for node, deps in self.dag.items() if not deps)
 40
 41    def __repr__(self) -> str:
 42        return f"Plan\n----\n{repr(self.root)}"
 43
 44
 45class Step:
 46    @classmethod
 47    def from_expression(cls, expression: exp.Expr, ctes: dict[str, Step] | None = None) -> Step:
 48        """
 49        Builds a DAG of Steps from a SQL expression so that it's easier to execute in an engine.
 50        Note: the expression's tables and subqueries must be aliased for this method to work. For
 51        example, given the following expression:
 52
 53        SELECT
 54          x.a,
 55          SUM(x.b)
 56        FROM x AS x
 57        JOIN y AS y
 58          ON x.a = y.a
 59        GROUP BY x.a
 60
 61        the following DAG is produced (the expression IDs might differ per execution):
 62
 63        - Aggregate: x (4347984624)
 64            Context:
 65              Aggregations:
 66                - SUM(x.b)
 67              Group:
 68                - x.a
 69            Projections:
 70              - x.a
 71              - "x".""
 72            Dependencies:
 73            - Join: x (4347985296)
 74              Context:
 75                y:
 76                On: x.a = y.a
 77              Projections:
 78              Dependencies:
 79              - Scan: x (4347983136)
 80                Context:
 81                  Source: x AS x
 82                Projections:
 83              - Scan: y (4343416624)
 84                Context:
 85                  Source: y AS y
 86                Projections:
 87
 88        Args:
 89            expression: the expression to build the DAG from.
 90            ctes: a dictionary that maps CTEs to their corresponding Step DAG by name.
 91
 92        Returns:
 93            A Step DAG corresponding to `expression`.
 94        """
 95        ctes = ctes or {}
 96        expression = expression.unnest()
 97        with_: exp.With | None = expression.args.get("with_")
 98
 99        # CTEs break the mold of scope and introduce themselves to all in the context.
100        if with_ is not None:
101            ctes = ctes.copy()
102            for cte in with_.expressions:
103                step = Step.from_expression(cte.this, ctes)
104                step.name = cte.alias
105                ctes[step.name] = step  # type: ignore
106
107        from_ = expression.args.get("from_")
108
109        if isinstance(expression, exp.Select) and from_:
110            step = Scan.from_expression(from_.this, ctes)
111        elif isinstance(expression, exp.SetOperation):
112            step = SetOperation.from_expression(expression, ctes)
113        else:
114            step = Scan()
115
116        joins: list[exp.Join] | None = expression.args.get("joins")
117
118        if joins is not None:
119            join = Join.from_joins(joins, ctes)
120            join.name = step.name
121            join.source_name = step.name
122            join.add_dependency(step)
123            step = join
124        # final selects in this chain of steps representing a select
125        projections: list[exp.Expr] = []
126        # intermediate computations of agg funcs eg x + 1 in SUM(x + 1)
127        operands: dict[exp.Expr, str] = {}
128        aggregations: dict[exp.Expr, None] = {}
129        next_operand_name = name_sequence("_a_")
130
131        def extract_agg_operands(expression: exp.Expr) -> bool:
132            agg_funcs = tuple(find_all_in_scope(expression, exp.AggFunc))
133            if agg_funcs:
134                aggregations[expression] = None
135
136            for agg in agg_funcs:
137                for operand in agg.unnest_operands():
138                    if isinstance(operand, exp.Column):
139                        continue
140                    if operand not in operands:
141                        operands[operand] = next_operand_name()
142
143                    operand.replace(exp.column(operands[operand], quoted=True))
144
145            return bool(agg_funcs)
146
147        def set_ops_and_aggs(step) -> None:
148            step.operands = tuple(alias(operand, alias_) for operand, alias_ in operands.items())
149            step.aggregations = list(aggregations)
150
151        for e in expression.expressions:
152            if find_in_scope(e, exp.AggFunc):
153                projections.append(exp.column(e.alias_or_name, step.name, quoted=True))
154                extract_agg_operands(e)
155            else:
156                projections.append(e)
157
158        where: exp.Where | None = expression.args.get("where")
159
160        if where is not None:
161            step.condition = where.this
162
163        group: exp.Group | None = expression.args.get("group")
164
165        if group is not None or aggregations:
166            aggregate = Aggregate()
167            aggregate.source = step.name
168            aggregate.name = step.name
169
170            having: exp.Having | None = expression.args.get("having")
171
172            if having is not None:
173                if extract_agg_operands(exp.alias_(having.this, "_h", quoted=True)):
174                    aggregate.condition = exp.column("_h", step.name, quoted=True)
175                else:
176                    aggregate.condition = having.this
177
178            set_ops_and_aggs(aggregate)
179
180            # give aggregates names and replace projections with references to them
181            aggregate.group = {
182                f"_g{i}": e for i, e in enumerate(group.expressions if group else [])
183            }
184
185            intermediate: dict[str | exp.Expr, str] = {}
186            for k, v in aggregate.group.items():
187                intermediate[v] = k
188                if isinstance(v, exp.Column):
189                    intermediate[v.name] = k
190
191            for projection in projections:
192                for node in projection.walk():
193                    name = intermediate.get(node)
194                    if name:
195                        node.replace(exp.column(name, step.name))
196
197            if aggregate.condition:
198                for node in aggregate.condition.walk():
199                    name = intermediate.get(node) or intermediate.get(node.name)
200                    if name:
201                        node.replace(exp.column(name, step.name))
202
203            aggregate.add_dependency(step)
204            step = aggregate
205        else:
206            aggregate = None
207
208        order: exp.Order | None = expression.args.get("order")
209
210        if order is not None:
211            if aggregate is not None and isinstance(step, Aggregate):
212                for i, ordered in enumerate(order.expressions):
213                    if extract_agg_operands(exp.alias_(ordered.this, f"_o_{i}", quoted=True)):
214                        ordered.this.replace(exp.column(f"_o_{i}", step.name, quoted=True))
215
216                set_ops_and_aggs(aggregate)
217
218            sort = Sort()
219            sort.name = step.name
220            sort.key = order.expressions
221            sort.add_dependency(step)
222            step = sort
223
224        step.projections = projections
225
226        if isinstance(expression, exp.Select) and expression.args.get("distinct"):
227            distinct = Aggregate()
228            distinct.source = step.name
229            distinct.name = step.name
230            distinct.group = {
231                e.alias_or_name: exp.column(col=e.alias_or_name, table=step.name)
232                for e in projections or expression.expressions
233            }
234            distinct.add_dependency(step)
235            step = distinct
236
237        limit: exp.Limit | None = expression.args.get("limit")
238
239        if limit is not None:
240            step.limit = int(limit.text("expression"))
241
242        return step
243
244    def __init__(self) -> None:
245        self.name: str | None = None
246        self.dependencies: set[Step] = set()
247        self.dependents: set[Step] = set()
248        self.projections: Sequence[exp.Expr] = []
249        self.limit: float = math.inf
250        self.condition: exp.Expr | None = None
251
252    def add_dependency(self, dependency: Step) -> None:
253        self.dependencies.add(dependency)
254        dependency.dependents.add(self)
255
256    def __repr__(self) -> str:
257        return self.to_s()
258
259    def to_s(self, level: int = 0) -> str:
260        indent = "  " * level
261        nested = f"{indent}    "
262
263        context = self._to_s(f"{nested}  ")
264
265        if context:
266            context = [f"{nested}Context:"] + context
267
268        lines = [
269            f"{indent}- {self.id}",
270            *context,
271            f"{nested}Projections:",
272        ]
273
274        for expression in self.projections:
275            lines.append(f"{nested}  - {expression.sql()}")
276
277        if self.condition:
278            lines.append(f"{nested}Condition: {self.condition.sql()}")
279
280        if self.limit is not math.inf:
281            lines.append(f"{nested}Limit: {self.limit}")
282
283        if self.dependencies:
284            lines.append(f"{nested}Dependencies:")
285            for dependency in self.dependencies:
286                lines.append("  " + dependency.to_s(level + 1))
287
288        return "\n".join(lines)
289
290    @property
291    def type_name(self) -> str:
292        return self.__class__.__name__
293
294    @property
295    def id(self) -> str:
296        name = self.name
297        name = f" {name}" if name else ""
298        return f"{self.type_name}:{name} ({id(self)})"
299
300    def _to_s(self, _indent: str) -> list[str]:
301        return []
302
303
304class Scan(Step):
305    @classmethod
306    def from_expression(cls, expression: exp.Expr, ctes: dict[str, Step] | None = None) -> Step:
307        table: exp.Expr = expression
308        alias_ = expression.alias_or_name
309
310        if isinstance(expression, exp.Subquery):
311            table = expression.this
312            step = Step.from_expression(table, ctes)
313            step.name = alias_
314            return step
315
316        step = Scan()
317        step.name = alias_
318        step.source = expression
319        if ctes and table.name in ctes:
320            step.add_dependency(ctes[table.name])
321
322        return step
323
324    def __init__(self) -> None:
325        super().__init__()
326        self.source: exp.Expr | None = None
327
328    def _to_s(self, indent: str) -> list[str]:
329        return [f"{indent}Source: {self.source.sql() if self.source else '-static-'}"]  # type: ignore
330
331
332class Join(Step):
333    @classmethod
334    def from_joins(cls, joins: Iterable[exp.Join], ctes: dict[str, Step] | None = None) -> Join:
335        step = Join()
336
337        for join in joins:
338            source_key, join_key, condition = join_condition(join)
339            step.joins[join.alias_or_name] = {
340                "side": join.side,  # type: ignore
341                "join_key": join_key,
342                "source_key": source_key,
343                "condition": condition,
344            }
345
346            step.add_dependency(Scan.from_expression(join.this, ctes))
347
348        return step
349
350    def __init__(self) -> None:
351        super().__init__()
352        self.source_name: str | None = None
353        self.joins: dict[str, dict[str, list[str] | exp.Expr | list[exp.Expr]]] = {}
354
355    def _to_s(self, indent: str) -> list[str]:
356        lines = [f"{indent}Source: {self.source_name or self.name}"]
357        for name, join in self.joins.items():
358            lines.append(f"{indent}{name}: {join['side'] or 'INNER'}")
359            join_key = ", ".join(str(key) for key in t.cast(list[str], join.get("join_key") or []))
360            if join_key:
361                lines.append(f"{indent}Key: {join_key}")
362            if join.get("condition"):
363                lines.append(f"{indent}On: {join['condition'].sql()}")  # type: ignore
364        return lines
365
366
367class Aggregate(Step):
368    def __init__(self) -> None:
369        super().__init__()
370        self.aggregations: list[exp.Expr] = []
371        self.operands: tuple[exp.Expr, ...] = ()
372        self.group: dict[str, exp.Expr] = {}
373        self.source: str | None = None
374
375    def _to_s(self, indent: str) -> list[str]:
376        lines = [f"{indent}Aggregations:"]
377
378        for expression in self.aggregations:
379            lines.append(f"{indent}  - {expression.sql()}")
380
381        if self.group:
382            lines.append(f"{indent}Group:")
383            for expression in self.group.values():
384                lines.append(f"{indent}  - {expression.sql()}")
385        if self.condition:
386            lines.append(f"{indent}Having:")
387            lines.append(f"{indent}  - {self.condition.sql()}")
388        if self.operands:
389            lines.append(f"{indent}Operands:")
390            for expression in self.operands:
391                lines.append(f"{indent}  - {expression.sql()}")
392
393        return lines
394
395
396class Sort(Step):
397    def __init__(self) -> None:
398        super().__init__()
399        self.key: list[exp.Expr] | None = None
400
401    def _to_s(self, indent: str) -> list[str]:
402        lines = [f"{indent}Key:"]
403
404        for expression in self.key:  # type: ignore
405            lines.append(f"{indent}  - {expression.sql()}")
406
407        return lines
408
409
410class SetOperation(Step):
411    def __init__(self, op: type[exp.Expr], left: str, right: str, distinct: bool = False) -> None:
412        super().__init__()
413        self.op: type[exp.Expr] = op
414        self.left: str = left
415        self.right: str = right
416        self.distinct: bool = distinct
417
418    @classmethod
419    def from_expression(
420        cls, expression: exp.Expr, ctes: dict[str, Step] | None = None
421    ) -> SetOperation:
422        assert isinstance(expression, exp.SetOperation)
423
424        left = Step.from_expression(expression.left, ctes)
425        # SELECT 1 UNION SELECT 2  <-- these subqueries don't have names
426        left.name = left.name or "left"
427        right = Step.from_expression(expression.right, ctes)
428        right.name = right.name or "right"
429        step = cls(
430            op=expression.__class__,
431            left=left.name,
432            right=right.name,
433            distinct=bool(expression.args.get("distinct")),
434        )
435
436        step.add_dependency(left)
437        step.add_dependency(right)
438
439        limit: exp.Limit | None = expression.args.get("limit")
440
441        if limit is not None:
442            step.limit = int(limit.text("expression"))
443
444        return step
445
446    def _to_s(self, indent: str) -> list[str]:
447        lines: list[str] = []
448        if self.distinct:
449            lines.append(f"{indent}Distinct: {self.distinct}")
450        return lines
451
452    @property
453    def type_name(self) -> str:
454        return self.op.__name__