Edit on GitHub

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

Builds a DAG of Steps from a SQL expression so that it's easier to execute in an engine. Note: the expression's tables and subqueries must be aliased for this method to work. For example, given the following expression:

SELECT x.a, SUM(x.b) FROM x AS x JOIN y AS y ON x.a = y.a GROUP BY x.a

the following DAG is produced (the expression IDs might differ per execution):

  • Aggregate: x (4347984624) Context: Aggregations: - SUM(x.b) Group: - x.a Projections:
    • x.a
    • "x"."" Dependencies:
      • Join: x (4347985296) Context: y: On: x.a = y.a Projections: Dependencies:
    • Scan: x (4347983136) Context: Source: x AS x Projections:
    • Scan: y (4343416624) Context: Source: y AS y Projections:
Arguments:
  • expression: the expression to build the DAG from.
  • ctes: a dictionary that maps CTEs to their corresponding Step DAG by name.
Returns:

A Step DAG corresponding to expression.

name: Optional[str]
dependencies: Set[Step]
dependents: Set[Step]
projections: Sequence[sqlglot.expressions.Expression]
limit: float
condition: Optional[sqlglot.expressions.Expression]
def add_dependency(self, dependency: Step) -> None:
250    def add_dependency(self, dependency: Step) -> None:
251        self.dependencies.add(dependency)
252        dependency.dependents.add(self)
def to_s(self, level: int = 0) -> str:
257    def to_s(self, level: int = 0) -> str:
258        indent = "  " * level
259        nested = f"{indent}    "
260
261        context = self._to_s(f"{nested}  ")
262
263        if context:
264            context = [f"{nested}Context:"] + context
265
266        lines = [
267            f"{indent}- {self.id}",
268            *context,
269            f"{nested}Projections:",
270        ]
271
272        for expression in self.projections:
273            lines.append(f"{nested}  - {expression.sql()}")
274
275        if self.condition:
276            lines.append(f"{nested}Condition: {self.condition.sql()}")
277
278        if self.limit is not math.inf:
279            lines.append(f"{nested}Limit: {self.limit}")
280
281        if self.dependencies:
282            lines.append(f"{nested}Dependencies:")
283            for dependency in self.dependencies:
284                lines.append("  " + dependency.to_s(level + 1))
285
286        return "\n".join(lines)
type_name: str
288    @property
289    def type_name(self) -> str:
290        return self.__class__.__name__
id: str
292    @property
293    def id(self) -> str:
294        name = self.name
295        name = f" {name}" if name else ""
296        return f"{self.type_name}:{name} ({id(self)})"
class Scan(Step):
302class Scan(Step):
303    @classmethod
304    def from_expression(
305        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
306    ) -> Step:
307        table = 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: t.Optional[exp.Expression] = None
327
328    def _to_s(self, indent: str) -> t.List[str]:
329        return [f"{indent}Source: {self.source.sql() if self.source else '-static-'}"]  # type: ignore
@classmethod
def from_expression( cls, expression: sqlglot.expressions.Expression, ctes: Optional[Dict[str, Step]] = None) -> Step:
303    @classmethod
304    def from_expression(
305        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
306    ) -> Step:
307        table = 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

Builds a DAG of Steps from a SQL expression so that it's easier to execute in an engine. Note: the expression's tables and subqueries must be aliased for this method to work. For example, given the following expression:

SELECT x.a, SUM(x.b) FROM x AS x JOIN y AS y ON x.a = y.a GROUP BY x.a

the following DAG is produced (the expression IDs might differ per execution):

  • Aggregate: x (4347984624) Context: Aggregations: - SUM(x.b) Group: - x.a Projections:
    • x.a
    • "x"."" Dependencies:
      • Join: x (4347985296) Context: y: On: x.a = y.a Projections: Dependencies:
    • Scan: x (4347983136) Context: Source: x AS x Projections:
    • Scan: y (4343416624) Context: Source: y AS y Projections:
Arguments:
  • expression: the expression to build the DAG from.
  • ctes: a dictionary that maps CTEs to their corresponding Step DAG by name.
Returns:

A Step DAG corresponding to expression.

source: Optional[sqlglot.expressions.Expression]
class Join(Step):
332class Join(Step):
333    @classmethod
334    def from_joins(
335        cls, joins: t.Iterable[exp.Join], ctes: t.Optional[t.Dict[str, Step]] = None
336    ) -> Join:
337        step = Join()
338
339        for join in joins:
340            source_key, join_key, condition = join_condition(join)
341            step.joins[join.alias_or_name] = {
342                "side": join.side,  # type: ignore
343                "join_key": join_key,
344                "source_key": source_key,
345                "condition": condition,
346            }
347
348            step.add_dependency(Scan.from_expression(join.this, ctes))
349
350        return step
351
352    def __init__(self) -> None:
353        super().__init__()
354        self.source_name: t.Optional[str] = None
355        self.joins: t.Dict[str, t.Dict[str, t.List[str] | exp.Expression]] = {}
356
357    def _to_s(self, indent: str) -> t.List[str]:
358        lines = [f"{indent}Source: {self.source_name or self.name}"]
359        for name, join in self.joins.items():
360            lines.append(f"{indent}{name}: {join['side'] or 'INNER'}")
361            join_key = ", ".join(str(key) for key in t.cast(list, join.get("join_key") or []))
362            if join_key:
363                lines.append(f"{indent}Key: {join_key}")
364            if join.get("condition"):
365                lines.append(f"{indent}On: {join['condition'].sql()}")  # type: ignore
366        return lines
@classmethod
def from_joins( cls, joins: Iterable[sqlglot.expressions.Join], ctes: Optional[Dict[str, Step]] = None) -> Join:
333    @classmethod
334    def from_joins(
335        cls, joins: t.Iterable[exp.Join], ctes: t.Optional[t.Dict[str, Step]] = None
336    ) -> Join:
337        step = Join()
338
339        for join in joins:
340            source_key, join_key, condition = join_condition(join)
341            step.joins[join.alias_or_name] = {
342                "side": join.side,  # type: ignore
343                "join_key": join_key,
344                "source_key": source_key,
345                "condition": condition,
346            }
347
348            step.add_dependency(Scan.from_expression(join.this, ctes))
349
350        return step
source_name: Optional[str]
joins: Dict[str, Dict[str, Union[List[str], sqlglot.expressions.Expression]]]
class Aggregate(Step):
369class Aggregate(Step):
370    def __init__(self) -> None:
371        super().__init__()
372        self.aggregations: t.List[exp.Expression] = []
373        self.operands: t.Tuple[exp.Expression, ...] = ()
374        self.group: t.Dict[str, exp.Expression] = {}
375        self.source: t.Optional[str] = None
376
377    def _to_s(self, indent: str) -> t.List[str]:
378        lines = [f"{indent}Aggregations:"]
379
380        for expression in self.aggregations:
381            lines.append(f"{indent}  - {expression.sql()}")
382
383        if self.group:
384            lines.append(f"{indent}Group:")
385            for expression in self.group.values():
386                lines.append(f"{indent}  - {expression.sql()}")
387        if self.condition:
388            lines.append(f"{indent}Having:")
389            lines.append(f"{indent}  - {self.condition.sql()}")
390        if self.operands:
391            lines.append(f"{indent}Operands:")
392            for expression in self.operands:
393                lines.append(f"{indent}  - {expression.sql()}")
394
395        return lines
aggregations: List[sqlglot.expressions.Expression]
operands: Tuple[sqlglot.expressions.Expression, ...]
group: Dict[str, sqlglot.expressions.Expression]
source: Optional[str]
class Sort(Step):
398class Sort(Step):
399    def __init__(self) -> None:
400        super().__init__()
401        self.key = None
402
403    def _to_s(self, indent: str) -> t.List[str]:
404        lines = [f"{indent}Key:"]
405
406        for expression in self.key:  # type: ignore
407            lines.append(f"{indent}  - {expression.sql()}")
408
409        return lines
key
class SetOperation(Step):
412class SetOperation(Step):
413    def __init__(
414        self,
415        op: t.Type[exp.Expression],
416        left: str | None,
417        right: str | None,
418        distinct: bool = False,
419    ) -> None:
420        super().__init__()
421        self.op = op
422        self.left = left
423        self.right = right
424        self.distinct = distinct
425
426    @classmethod
427    def from_expression(
428        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
429    ) -> SetOperation:
430        assert isinstance(expression, exp.SetOperation)
431
432        left = Step.from_expression(expression.left, ctes)
433        # SELECT 1 UNION SELECT 2  <-- these subqueries don't have names
434        left.name = left.name or "left"
435        right = Step.from_expression(expression.right, ctes)
436        right.name = right.name or "right"
437        step = cls(
438            op=expression.__class__,
439            left=left.name,
440            right=right.name,
441            distinct=bool(expression.args.get("distinct")),
442        )
443
444        step.add_dependency(left)
445        step.add_dependency(right)
446
447        limit = expression.args.get("limit")
448
449        if limit:
450            step.limit = int(limit.text("expression"))
451
452        return step
453
454    def _to_s(self, indent: str) -> t.List[str]:
455        lines = []
456        if self.distinct:
457            lines.append(f"{indent}Distinct: {self.distinct}")
458        return lines
459
460    @property
461    def type_name(self) -> str:
462        return self.op.__name__
SetOperation( op: Type[sqlglot.expressions.Expression], left: str | None, right: str | None, distinct: bool = False)
413    def __init__(
414        self,
415        op: t.Type[exp.Expression],
416        left: str | None,
417        right: str | None,
418        distinct: bool = False,
419    ) -> None:
420        super().__init__()
421        self.op = op
422        self.left = left
423        self.right = right
424        self.distinct = distinct
op
left
right
distinct
@classmethod
def from_expression( cls, expression: sqlglot.expressions.Expression, ctes: Optional[Dict[str, Step]] = None) -> SetOperation:
426    @classmethod
427    def from_expression(
428        cls, expression: exp.Expression, ctes: t.Optional[t.Dict[str, Step]] = None
429    ) -> SetOperation:
430        assert isinstance(expression, exp.SetOperation)
431
432        left = Step.from_expression(expression.left, ctes)
433        # SELECT 1 UNION SELECT 2  <-- these subqueries don't have names
434        left.name = left.name or "left"
435        right = Step.from_expression(expression.right, ctes)
436        right.name = right.name or "right"
437        step = cls(
438            op=expression.__class__,
439            left=left.name,
440            right=right.name,
441            distinct=bool(expression.args.get("distinct")),
442        )
443
444        step.add_dependency(left)
445        step.add_dependency(right)
446
447        limit = expression.args.get("limit")
448
449        if limit:
450            step.limit = int(limit.text("expression"))
451
452        return step

Builds a DAG of Steps from a SQL expression so that it's easier to execute in an engine. Note: the expression's tables and subqueries must be aliased for this method to work. For example, given the following expression:

SELECT x.a, SUM(x.b) FROM x AS x JOIN y AS y ON x.a = y.a GROUP BY x.a

the following DAG is produced (the expression IDs might differ per execution):

  • Aggregate: x (4347984624) Context: Aggregations: - SUM(x.b) Group: - x.a Projections:
    • x.a
    • "x"."" Dependencies:
      • Join: x (4347985296) Context: y: On: x.a = y.a Projections: Dependencies:
    • Scan: x (4347983136) Context: Source: x AS x Projections:
    • Scan: y (4343416624) Context: Source: y AS y Projections:
Arguments:
  • expression: the expression to build the DAG from.
  • ctes: a dictionary that maps CTEs to their corresponding Step DAG by name.
Returns:

A Step DAG corresponding to expression.

type_name: str
460    @property
461    def type_name(self) -> str:
462        return self.op.__name__