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