Edit on GitHub

sqlglot.optimizer.optimize_joins

View Source

 1from __future__ import annotations
 2
 3import typing as t
 4
 5from sqlglot import exp
 6from sqlglot.helper import tsort
 7
 8JOIN_ATTRS = ("on", "side", "kind", "using", "method")
 9
10
11def optimize_joins(expression):
12    """
13    Removes cross joins if possible and reorder joins based on predicate dependencies.
14
15    Example:
16        >>> from sqlglot import parse_one
17        >>> optimize_joins(parse_one("SELECT * FROM x CROSS JOIN y JOIN z ON x.a = z.a AND y.a = z.a")).sql()
18        'SELECT * FROM x JOIN z ON x.a = z.a AND TRUE JOIN y ON y.a = z.a'
19    """
20
21    for select in expression.find_all(exp.Select):
22        references = {}
23        cross_joins = []
24
25        for join in select.args.get("joins", []):
26            tables = other_table_names(join)
27
28            if tables:
29                for table in tables:
30                    references[table] = references.get(table, []) + [join]
31            else:
32                cross_joins.append((join.alias_or_name, join))
33
34        for name, join in cross_joins:
35            for dep in references.get(name, []):
36                on = dep.args["on"]
37
38                if isinstance(on, exp.Connector):
39                    if len(other_table_names(dep)) < 2:
40                        continue
41
42                    for predicate in on.flatten():
43                        if name in exp.column_table_names(predicate):
44                            predicate.replace(exp.true())
45                            join.on(predicate, copy=False)
46
47    expression = reorder_joins(expression)
48    expression = normalize(expression)
49    return expression
50
51
52def reorder_joins(expression):
53    """
54    Reorder joins by topological sort order based on predicate references.
55    """
56    for from_ in expression.find_all(exp.From):
57        parent = from_.parent
58        joins = {join.alias_or_name: join for join in parent.args.get("joins", [])}
59        dag = {name: other_table_names(join) for name, join in joins.items()}
60        parent.set(
61            "joins",
62            [joins[name] for name in tsort(dag) if name != from_.alias_or_name and name in joins],
63        )
64    return expression
65
66
67def normalize(expression):
68    """
69    Remove INNER and OUTER from joins as they are optional.
70    """
71    for join in expression.find_all(exp.Join):
72        if not any(join.args.get(k) for k in JOIN_ATTRS):
73            join.set("kind", "CROSS")
74
75        if join.kind == "CROSS":
76            join.set("on", None)
77        else:
78            join.set("kind", None)
79
80            if not join.args.get("on") and not join.args.get("using"):
81                join.set("on", exp.true())
82    return expression
83
84
85def other_table_names(join: exp.Join) -> t.Set[str]:
86    on = join.args.get("on")
87    return exp.column_table_names(on, join.alias_or_name) if on else set()

JOIN_ATTRS = ('on', 'side', 'kind', 'using', 'method')

def optimize_joins(expression): View Source

12def optimize_joins(expression):
13    """
14    Removes cross joins if possible and reorder joins based on predicate dependencies.
15
16    Example:
17        >>> from sqlglot import parse_one
18        >>> optimize_joins(parse_one("SELECT * FROM x CROSS JOIN y JOIN z ON x.a = z.a AND y.a = z.a")).sql()
19        'SELECT * FROM x JOIN z ON x.a = z.a AND TRUE JOIN y ON y.a = z.a'
20    """
21
22    for select in expression.find_all(exp.Select):
23        references = {}
24        cross_joins = []
25
26        for join in select.args.get("joins", []):
27            tables = other_table_names(join)
28
29            if tables:
30                for table in tables:
31                    references[table] = references.get(table, []) + [join]
32            else:
33                cross_joins.append((join.alias_or_name, join))
34
35        for name, join in cross_joins:
36            for dep in references.get(name, []):
37                on = dep.args["on"]
38
39                if isinstance(on, exp.Connector):
40                    if len(other_table_names(dep)) < 2:
41                        continue
42
43                    for predicate in on.flatten():
44                        if name in exp.column_table_names(predicate):
45                            predicate.replace(exp.true())
46                            join.on(predicate, copy=False)
47
48    expression = reorder_joins(expression)
49    expression = normalize(expression)
50    return expression

Removes cross joins if possible and reorder joins based on predicate dependencies.

Example:

>>> from sqlglot import parse_one
>>> optimize_joins(parse_one("SELECT * FROM x CROSS JOIN y JOIN z ON x.a = z.a AND y.a = z.a")).sql()
'SELECT * FROM x JOIN z ON x.a = z.a AND TRUE JOIN y ON y.a = z.a'

def reorder_joins(expression): View Source

53def reorder_joins(expression):
54    """
55    Reorder joins by topological sort order based on predicate references.
56    """
57    for from_ in expression.find_all(exp.From):
58        parent = from_.parent
59        joins = {join.alias_or_name: join for join in parent.args.get("joins", [])}
60        dag = {name: other_table_names(join) for name, join in joins.items()}
61        parent.set(
62            "joins",
63            [joins[name] for name in tsort(dag) if name != from_.alias_or_name and name in joins],
64        )
65    return expression

Reorder joins by topological sort order based on predicate references.

def normalize(expression): View Source

68def normalize(expression):
69    """
70    Remove INNER and OUTER from joins as they are optional.
71    """
72    for join in expression.find_all(exp.Join):
73        if not any(join.args.get(k) for k in JOIN_ATTRS):
74            join.set("kind", "CROSS")
75
76        if join.kind == "CROSS":
77            join.set("on", None)
78        else:
79            join.set("kind", None)
80
81            if not join.args.get("on") and not join.args.get("using"):
82                join.set("on", exp.true())
83    return expression

Remove INNER and OUTER from joins as they are optional.

def other_table_names(join: sqlglot.expressions.Join) -> Set[str]: View Source

86def other_table_names(join: exp.Join) -> t.Set[str]:
87    on = join.args.get("on")
88    return exp.column_table_names(on, join.alias_or_name) if on else set()