sqlglot.planner

  1from __future__ import annotations
  2
  3import itertools
  4import math
  5import typing as t
  6
  7from sqlglot import alias, exp
  8from sqlglot.errors import UnsupportedError
  9from sqlglot.optimizer.eliminate_joins import join_condition
 10
 11
 12class Plan:
 13    def __init__(self, expression: exp.Expression) -> None:
 14        self.expression = expression.copy()
 15        self.root = Step.from_expression(self.expression)
 16        self._dag: t.Dict[Step, t.Set[Step]] = {}
 17
 18    @property
 19    def dag(self) -> t.Dict[Step, t.Set[Step]]:
 20        if not self._dag:
 21            dag: t.Dict[Step, t.Set[Step]] = {}
 22            nodes = {self.root}
 23
 24            while nodes:
 25                node = nodes.pop()
 26                dag[node] = set()
 27                for dep in node.dependencies:
 28                    dag[node].add(dep)
 29                    nodes.add(dep)
 30            self._dag = dag
 31
 32        return self._dag
 33
 34    @property
 35    def leaves(self) -> t.Iterator[Step]:
 36        return (node for node, deps in self.dag.items() if not deps)
 37
 38    def __repr__(self) -> str:
 39        return f"Plan\n----\n{repr(self.root)}"
 40
 41
 42class Step:
 43    @classmethod
 44    def from_expression(
 45        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
 46    ) -> 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        with_ = expression.args.get("with")
 96
 97        # CTEs break the mold of scope and introduce themselves to all in the context.
 98        if with_:
 99            ctes = ctes.copy()
100            for cte in with_.expressions:
101                step = Step.from_expression(cte.this, ctes)
102                step.name = cte.alias
103                ctes[step.name] = step  # type: ignore
104
105        from_ = expression.args.get("from")
106
107        if isinstance(expression, exp.Select) and from_:
108            from_ = from_.expressions
109            if len(from_) > 1:
110                raise UnsupportedError(
111                    "Multi-from statements are unsupported. Run it through the optimizer"
112                )
113
114            step = Scan.from_expression(from_[0], ctes)
115        elif isinstance(expression, exp.Union):
116            step = SetOperation.from_expression(expression, ctes)
117        else:
118            step = Scan()
119
120        joins = expression.args.get("joins")
121
122        if joins:
123            join = Join.from_joins(joins, ctes)
124            join.name = step.name
125            join.add_dependency(step)
126            step = join
127
128        projections = []  # final selects in this chain of steps representing a select
129        operands = {}  # intermediate computations of agg funcs eg x + 1 in SUM(x + 1)
130        aggregations = []
131        sequence = itertools.count()
132
133        def extract_agg_operands(expression):
134            for agg in expression.find_all(exp.AggFunc):
135                for operand in agg.unnest_operands():
136                    if isinstance(operand, exp.Column):
137                        continue
138                    if operand not in operands:
139                        operands[operand] = f"_a_{next(sequence)}"
140                    operand.replace(exp.column(operands[operand], quoted=True))
141
142        for e in expression.expressions:
143            if e.find(exp.AggFunc):
144                projections.append(exp.column(e.alias_or_name, step.name, quoted=True))
145                aggregations.append(e)
146                extract_agg_operands(e)
147            else:
148                projections.append(e)
149
150        where = expression.args.get("where")
151
152        if where:
153            step.condition = where.this
154
155        group = expression.args.get("group")
156
157        if group or aggregations:
158            aggregate = Aggregate()
159            aggregate.source = step.name
160            aggregate.name = step.name
161
162            having = expression.args.get("having")
163
164            if having:
165                extract_agg_operands(having)
166                aggregate.condition = having.this
167
168            aggregate.operands = tuple(
169                alias(operand, alias_) for operand, alias_ in operands.items()
170            )
171            aggregate.aggregations = aggregations
172            # give aggregates names and replace projections with references to them
173            aggregate.group = {
174                f"_g{i}": e for i, e in enumerate(group.expressions if group else [])
175            }
176            for projection in projections:
177                for i, e in aggregate.group.items():
178                    for child, *_ in projection.walk():
179                        if child == e:
180                            child.replace(exp.column(i, step.name))
181            aggregate.add_dependency(step)
182            step = aggregate
183
184        order = expression.args.get("order")
185
186        if order:
187            sort = Sort()
188            sort.name = step.name
189            sort.key = order.expressions
190            sort.add_dependency(step)
191            step = sort
192
193        step.projections = projections
194
195        if isinstance(expression, exp.Select) and expression.args.get("distinct"):
196            distinct = Aggregate()
197            distinct.source = step.name
198            distinct.name = step.name
199            distinct.group = {
200                e.alias_or_name: exp.column(col=e.alias_or_name, table=step.name)
201                for e in projections or expression.expressions
202            }
203            distinct.add_dependency(step)
204            step = distinct
205
206        limit = expression.args.get("limit")
207
208        if limit:
209            step.limit = int(limit.text("expression"))
210
211        return step
212
213    def __init__(self) -> None:
214        self.name: t.Optional[str] = None
215        self.dependencies: t.Set[Step] = set()
216        self.dependents: t.Set[Step] = set()
217        self.projections: t.Sequence[exp.Expression] = []
218        self.limit: float = math.inf
219        self.condition: t.Optional[exp.Expression] = None
220
221    def add_dependency(self, dependency: Step) -> None:
222        self.dependencies.add(dependency)
223        dependency.dependents.add(self)
224
225    def __repr__(self) -> str:
226        return self.to_s()
227
228    def to_s(self, level: int = 0) -> str:
229        indent = "  " * level
230        nested = f"{indent}    "
231
232        context = self._to_s(f"{nested}  ")
233
234        if context:
235            context = [f"{nested}Context:"] + context
236
237        lines = [
238            f"{indent}- {self.id}",
239            *context,
240            f"{nested}Projections:",
241        ]
242
243        for expression in self.projections:
244            lines.append(f"{nested}  - {expression.sql()}")
245
246        if self.condition:
247            lines.append(f"{nested}Condition: {self.condition.sql()}")
248
249        if self.limit is not math.inf:
250            lines.append(f"{nested}Limit: {self.limit}")
251
252        if self.dependencies:
253            lines.append(f"{nested}Dependencies:")
254            for dependency in self.dependencies:
255                lines.append("  " + dependency.to_s(level + 1))
256
257        return "\n".join(lines)
258
259    @property
260    def type_name(self) -> str:
261        return self.__class__.__name__
262
263    @property
264    def id(self) -> str:
265        name = self.name
266        name = f" {name}" if name else ""
267        return f"{self.type_name}:{name} ({id(self)})"
268
269    def _to_s(self, _indent: str) -> t.List[str]:
270        return []
271
272
273class Scan(Step):
274    @classmethod
275    def from_expression(
276        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
277    ) -> Step:
278        table = expression
279        alias_ = expression.alias_or_name
280
281        if isinstance(expression, exp.Subquery):
282            table = expression.this
283            step = Step.from_expression(table, ctes)
284            step.name = alias_
285            return step
286
287        step = Scan()
288        step.name = alias_
289        step.source = expression
290        if ctes and table.name in ctes:
291            step.add_dependency(ctes[table.name])
292
293        return step
294
295    def __init__(self) -> None:
296        super().__init__()
297        self.source: t.Optional[exp.Expression] = None
298
299    def _to_s(self, indent: str) -> t.List[str]:
300        return [f"{indent}Source: {self.source.sql() if self.source else '-static-'}"]  # type: ignore
301
302
303class Join(Step):
304    @classmethod
305    def from_joins(
306        cls, joins: t.Iterable[exp.Join], ctes: t.Optional[t.Dict[str, Step]] = None
307    ) -> Step:
308        step = Join()
309
310        for join in joins:
311            source_key, join_key, condition = join_condition(join)
312            step.joins[join.this.alias_or_name] = {
313                "side": join.side,
314                "join_key": join_key,
315                "source_key": source_key,
316                "condition": condition,
317            }
318
319            step.add_dependency(Scan.from_expression(join.this, ctes))
320
321        return step
322
323    def __init__(self) -> None:
324        super().__init__()
325        self.joins: t.Dict[str, t.Dict[str, t.List[str] | exp.Expression]] = {}
326
327    def _to_s(self, indent: str) -> t.List[str]:
328        lines = []
329        for name, join in self.joins.items():
330            lines.append(f"{indent}{name}: {join['side']}")
331            if join.get("condition"):
332                lines.append(f"{indent}On: {join['condition'].sql()}")  # type: ignore
333        return lines
334
335
336class Aggregate(Step):
337    def __init__(self) -> None:
338        super().__init__()
339        self.aggregations: t.List[exp.Expression] = []
340        self.operands: t.Tuple[exp.Expression, ...] = ()
341        self.group: t.Dict[str, exp.Expression] = {}
342        self.source: t.Optional[str] = None
343
344    def _to_s(self, indent: str) -> t.List[str]:
345        lines = [f"{indent}Aggregations:"]
346
347        for expression in self.aggregations:
348            lines.append(f"{indent}  - {expression.sql()}")
349
350        if self.group:
351            lines.append(f"{indent}Group:")
352            for expression in self.group.values():
353                lines.append(f"{indent}  - {expression.sql()}")
354        if self.condition:
355            lines.append(f"{indent}Having:")
356            lines.append(f"{indent}  - {self.condition.sql()}")
357        if self.operands:
358            lines.append(f"{indent}Operands:")
359            for expression in self.operands:
360                lines.append(f"{indent}  - {expression.sql()}")
361
362        return lines
363
364
365class Sort(Step):
366    def __init__(self) -> None:
367        super().__init__()
368        self.key = None
369
370    def _to_s(self, indent: str) -> t.List[str]:
371        lines = [f"{indent}Key:"]
372
373        for expression in self.key:  # type: ignore
374            lines.append(f"{indent}  - {expression.sql()}")
375
376        return lines
377
378
379class SetOperation(Step):
380    def __init__(
381        self,
382        op: t.Type[exp.Expression],
383        left: str | None,
384        right: str | None,
385        distinct: bool = False,
386    ) -> None:
387        super().__init__()
388        self.op = op
389        self.left = left
390        self.right = right
391        self.distinct = distinct
392
393    @classmethod
394    def from_expression(
395        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
396    ) -> Step:
397        assert isinstance(expression, exp.Union)
398        left = Step.from_expression(expression.left, ctes)
399        right = Step.from_expression(expression.right, ctes)
400        step = cls(
401            op=expression.__class__,
402            left=left.name,
403            right=right.name,
404            distinct=bool(expression.args.get("distinct")),
405        )
406        step.add_dependency(left)
407        step.add_dependency(right)
408        return step
409
410    def _to_s(self, indent: str) -> t.List[str]:
411        lines = []
412        if self.distinct:
413            lines.append(f"{indent}Distinct: {self.distinct}")
414        return lines
415
416    @property
417    def type_name(self) -> str:
418        return self.op.__name__