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