Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 67 additions & 32 deletions sqlglot/optimizer/unnest_subqueries.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,8 @@ def decorrelate(select, parent_select, external_columns, next_alias_name):

table_alias = next_alias_name()
keys = []
eq_count = 0
external_ids = set()

# for all external columns in the where statement, find the relevant predicate
# keys to convert it into a join
Expand Down Expand Up @@ -196,8 +198,14 @@ def decorrelate(select, parent_select, external_columns, next_alias_name):
return

keys.append((key, column, predicate))

if not any(isinstance(predicate, exp.EQ) for *_, predicate in keys):
external_ids.add(id(column))
eq_count += isinstance(predicate, exp.EQ)

# Non-EQ predicates are replaced with TRUE in the subquery, so they no longer filter the rows
# feeding its projections. Their keys are instead collected with ARRAY_AGG per EQ group and
# re-checked in the outer query with ARRAY_ANY, which is only correct for EXISTS: a projected
# value like SUM would otherwise be computed over the unfiltered rows
if not eq_count or (len(keys) > eq_count and not isinstance(parent_predicate, exp.Exists)):
Comment thread
geooo109 marked this conversation as resolved.
return
Comment on lines +208 to 209

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This makes sure non-EXISTS filters over decorrelated subqueries are left as-is to avoid the issues I mentioned in the changed tests. For EXISTS filters we only care about row existence, which is easier to answer, so we continue.


is_subquery_projection = any(
Expand All @@ -211,8 +219,16 @@ def decorrelate(select, parent_select, external_columns, next_alias_name):
group_by = []

for key, _, predicate in keys:
# The key is projected by the subquery and the other side is moved out of it, so
# neither can reference columns from the opposite scope
other = predicate.right if key is predicate.left else predicate.left
if any(id(c) in external_ids for c in key.find_all(exp.Column)) or any(
id(c) not in external_ids for c in other.find_all(exp.Column)
):
return

# if we filter on the value of the subquery, it needs to be unique
if key == value.this:
if key == value.this and isinstance(predicate, exp.EQ):
Comment on lines -215 to +231

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A non-EQ predicate whose key happens to match the projection would be routed to the GROUP BY path and then silently dropped, e.g.:

SELECT x.a FROM x WHERE EXISTS (SELECT y.b AS b FROM y WHERE y.a = x.a AND y.b > x.b)

Here, the projection y.b matches the key y.b from y.b > x.b, so without this we'd end up with:

SELECT x.a FROM x
LEFT JOIN (SELECT y.a AS _u_1, y.b AS b FROM y WHERE TRUE AND TRUE GROUP BY y.a, y.b) AS _u_0
  ON _u_0._u_1 = x.a
WHERE NOT _u_0._u_1 IS NULL

This doesn't even have the y.b > x.b filter, and also groups by both a and b, which is wrong because it duplicates outer rows.

key_aliases[key] = value.alias
group_by.append(key)
else:
Expand Down Expand Up @@ -250,21 +266,38 @@ def decorrelate(select, parent_select, external_columns, next_alias_name):
if isinstance(parent_predicate, exp.Exists):
select.set("expressions", [])

for key, alias in key_aliases.items():
if key in group_by:
# add all keys to the projections of the subquery
# so that we can use it as a join key
if isinstance(parent_predicate, exp.Exists) or key != value.this:
select.select(f"{key} AS {alias}", copy=False)
else:
select.select(exp.alias_(agg_func(this=key.copy()), alias, quoted=False), copy=False)
for key in group_by:
# add all keys to the projections of the subquery so that we can use it as a join key
if isinstance(parent_predicate, exp.Exists) or key != value.this:
select.select(exp.alias_(key, key_aliases[key]), copy=False)

array_keys = [key for key in key_aliases if key not in group_by]
use_struct = len(array_keys) > 1

if array_keys:
# Multiple keys are collected as one struct per row, so that all of their predicates are
# checked against the same row below
Comment on lines +278 to +279

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Without aggregating structs instead of plain keys, predicates over non-EQ keys could end up being satisfied across different rows, resulting in incorrect results.

For example, given:

WITH x AS (
  SELECT * FROM (VALUES (1, 5, 15)) AS t(a, b, c)
), y AS (
  SELECT * FROM (VALUES (1, 1, 10), (1, 9, 20)) AS t(a, b, c)
)

This query produces no results, because every row in y satisfies one inequality and fails the other:

SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b AND y.c < x.c);

If we simply aggregated the keys in separate arrays, we'd get this, which produces [(1,)]:

SELECT x.a FROM x
LEFT JOIN (
  SELECT y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2, ARRAY_AGG(y.c) AS _u_3
  FROM y WHERE TRUE AND TRUE AND TRUE GROUP BY y.a
) AS _u_0 ON _u_0._u_1 = x.a
WHERE NOT _u_0._u_1 IS NULL
  -- _u_2: [1, 9], _u_3: [10, 20]
  AND ARRAY_ANY(_u_0._u_2, _x -> _x > x.b)
  AND ARRAY_ANY(_u_0._u_3, _x -> _x < x.c);

As the comment points out, the struct is constructed to make sure the comparison happens over values appearing in the same original row in y.

array_alias = next_alias_name() if use_struct else key_aliases[array_keys[0]]
array_item = (
exp.Struct(
expressions=[
exp.PropertyEQ(this=exp.to_identifier(key_aliases[key]), expression=key.copy())
for key in array_keys
]
)
if use_struct
else array_keys[0].copy()
)
select.select(
exp.alias_(exp.ArrayAgg(this=array_item), array_alias, quoted=False), copy=False
)

alias = exp.column(value.alias, table_alias)
other = _other_operand(parent_predicate)
op_type = type(parent_predicate.parent) if parent_predicate else None

if isinstance(parent_predicate, exp.Exists):
alias = exp.column(list(key_aliases.values())[0], table_alias)
alias = exp.column(next(key_aliases[key] for key in group_by), table_alias)
Comment on lines -267 to +300

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This needed to happen because the first key may not be in an EQ, so we'd get an invalid join key.

parent_predicate = _replace(parent_predicate, f"NOT {alias} IS NULL")
elif isinstance(parent_predicate, exp.All):
assert issubclass(op_type, exp.Binary)
Expand Down Expand Up @@ -307,34 +340,36 @@ def remove_aggs(node):

select.parent.replace(alias)

for key, column, predicate in keys:
predicate.replace(exp.true())
nested = exp.column(key_aliases[key], table_alias)
array_predicates = []

if is_subquery_projection:
key.replace(nested)
if not isinstance(predicate, exp.EQ):
parent_select.where(predicate, copy=False)
continue
for key, _, predicate in keys:
predicate.replace(exp.true())

if key in group_by:
key.replace(nested)
key.replace(exp.column(key_aliases[key], table_alias))
else:
# Built as AST rather than a SQL string, because dialect-specific operators such as
# Postgres' `@>` can't be round-tripped through the default dialect's parser.
key.replace(exp.to_identifier("_x"))
right = exp.ArrayAny(
this=nested,
expression=exp.Lambda(this=predicate.copy(), expressions=[exp.to_identifier("_x")]),
)
parent_predicate = _replace(
parent_predicate,
exp.paren(exp.and_(parent_predicate.copy(), right, copy=False)),
key.replace(
exp.column(key_aliases[key], "_x") if use_struct else exp.to_identifier("_x")
)
array_predicates.append(predicate)
Comment thread
georgesittas marked this conversation as resolved.

if array_predicates:
# Built as AST rather than a SQL string, because dialect-specific operators such as
# Postgres' `@>` can't be round-tripped through the default dialect's parser.
right = exp.ArrayAny(
this=exp.column(array_alias, table_alias),
expression=exp.Lambda(
this=exp.and_(*array_predicates, copy=False), expressions=[exp.to_identifier("_x")]
),
)
parent_predicate = _replace(
parent_predicate, exp.paren(exp.and_(parent_predicate.copy(), right, copy=False))
)

parent_select.join(
select.group_by(*group_by, copy=False),
on=[predicate for *_, predicate in keys if isinstance(predicate, exp.EQ)],
# A grouped key is constant per group, so any predicate on it can be checked in the join
on=[predicate for key, _, predicate in keys if key in group_by],
join_type="LEFT",
join_alias=table_alias,
copy=False,
Expand Down
36 changes: 34 additions & 2 deletions tests/fixtures/optimizer/unnest_subqueries.sql
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ SELECT * FROM x WHERE x.a IN (SELECT y.a AS a FROM y WHERE y.b = x.a);
SELECT * FROM x LEFT JOIN (SELECT ARRAY_AGG(y.a) AS a, y.b AS _u_1 FROM y WHERE TRUE GROUP BY y.b) AS _u_0 ON _u_0._u_1 = x.a WHERE ARRAY_ANY(_u_0.a, _x -> _x = x.a);

SELECT * FROM x WHERE x.a < (SELECT SUM(y.a) AS a FROM y WHERE y.a = x.a and y.a = x.b and y.b <> x.d);
SELECT * FROM x LEFT JOIN (SELECT SUM(y.a) AS a, y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2 FROM y WHERE TRUE AND TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a AND _u_0._u_1 = x.b WHERE (x.a < _u_0.a AND ARRAY_ANY(_u_0._u_2, _x -> _x <> x.d));
SELECT * FROM x WHERE x.a < (SELECT SUM(y.a) AS a FROM y WHERE y.a = x.a AND y.a = x.b AND y.b <> x.d);
Comment on lines 32 to +33

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WITH x AS (
  SELECT * FROM (VALUES (2, 2, 1, 3)) AS t(a, b, c, d)
), y AS (
  SELECT * FROM (VALUES (2, 3, 0, 0), (2, 1, 0, 0)) AS t(a, b, c, d)
)
-- original, returns [] because the sum is evaluated over (2, 1, 0, 0) and produces 2 which is equal to x.a
SELECT x.a, x.b, x.c, x.d FROM x
WHERE x.a < (SELECT SUM(y.a) AS a FROM y WHERE y.a = x.a AND y.a = x.b AND y.b <> x.d);

-- optimized (main), returns [(2, 2, 1, 3)] because the sum is evaluated over both y rows, producing 4
-- which passes the x.a < _u_0.a check and array_agg produces [3, 1] and 1 passes _x <> x.d
SELECT x.a, x.b, x.c, x.d FROM x
LEFT JOIN (
  SELECT SUM(y.a) AS a, y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2
  FROM y WHERE TRUE AND TRUE AND TRUE GROUP BY y.a
) AS _u_0 ON _u_0._u_1 = x.a AND _u_0._u_1 = x.b
WHERE (x.a < _u_0.a AND ARRAY_ANY(_u_0._u_2, _x -> _x <> x.d));

The optimized form in the PR matches the original query's output: y.b <> x.d should exclude the row with b = 3, leaving a sum of 2, so that 2 < 2 is false.


SELECT * FROM x WHERE EXISTS (SELECT y.a AS a, y.b AS b FROM y WHERE x.a = y.a);
SELECT * FROM x LEFT JOIN (SELECT y.a AS a FROM y WHERE TRUE GROUP BY y.a) AS _u_0 ON x.a = _u_0.a WHERE NOT _u_0.a IS NULL;
Expand Down Expand Up @@ -137,7 +137,7 @@ SELECT x.a > (SELECT SUM(y.a) AS b FROM y) FROM x;
SELECT x.a > _u_0.b FROM x CROSS JOIN (SELECT SUM(y.a) AS b FROM y) AS _u_0;

SELECT (SELECT MAX(t2.c1) AS c1 FROM t2 WHERE t2.c2 = t1.c2 AND t2.c3 <= TRUNC(t1.c3)) AS c FROM t1;
SELECT _u_0.c1 AS c FROM t1 LEFT JOIN (SELECT MAX(t2.c1) AS c1, t2.c2 AS _u_1, MAX(t2.c3) AS _u_2 FROM t2 WHERE TRUE AND TRUE GROUP BY t2.c2) AS _u_0 ON _u_0._u_1 = t1.c2 WHERE _u_0._u_2 <= TRUNC(t1.c3);
SELECT (SELECT MAX(t2.c1) AS c1 FROM t2 WHERE t2.c2 = t1.c2 AND t2.c3 <= TRUNC(t1.c3)) AS c FROM t1;
Comment on lines 139 to +140

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

WITH t1 AS (
  SELECT * FROM (VALUES (1, 1, 5), (2, 2, 5)) AS t(c1, c2, c3)
), t2 AS (
  SELECT * FROM (VALUES (10, 1, 3), (20, 1, 9)) AS t(c1, c2, c3)
)
-- original, returns [(10,), (NULL,)] because y: (10, 1, 3), x: (1, 1, 5) satisfy the predicates, producing a max
-- of 10, and y: (20, 1, 9) doesn't satisfy t2.c3 <= trunc(t1.c3), so max is evaluated over {}, producing null
SELECT (SELECT MAX(t2.c1) AS c1 FROM t2 WHERE t2.c2 = t1.c2 AND t2.c3 <= TRUNC(t1.c3)) AS c
FROM t1;

-- optimized (main), returns []  because the rhs of the join produces (20, 1, 9) and the outer filter is
-- not satisfied, since 9 > 5
SELECT _u_0.c1 AS c FROM t1
LEFT JOIN (
  SELECT MAX(t2.c1) AS c1, t2.c2 AS _u_1, MAX(t2.c3) AS _u_2
  FROM t2 WHERE TRUE AND TRUE GROUP BY t2.c2
) AS _u_0 ON _u_0._u_1 = t1.c2
WHERE _u_0._u_2 <= TRUNC(t1.c3);

The optimized form in the PR again preserves the original query's output.


SELECT s.t AS t FROM s WHERE 1 IN (SELECT t.a AS a FROM t WHERE t.b > 1);
SELECT s.t AS t FROM s LEFT JOIN (SELECT t.a AS a FROM t WHERE t.b > 1 GROUP BY t.a) AS _u_0 ON 1 = _u_0.a WHERE NOT _u_0.a IS NULL;
Expand Down Expand Up @@ -211,3 +211,35 @@ SELECT x.id FROM x WHERE NOT EXISTS(SELECT 1 FROM y WHERE NOT (y.id = x.id));
# title: positive equality with NOT operand is unnested
SELECT x.flag FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.flag = (NOT x.flag));
SELECT x.flag FROM x LEFT JOIN (SELECT y.flag AS _u_1 FROM y WHERE TRUE GROUP BY y.flag) AS _u_0 ON _u_0._u_1 = (NOT x.flag) WHERE NOT _u_0._u_1 IS NULL;

# title: exists with a single non-equality key is unnested
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b);
SELECT x.a FROM x LEFT JOIN (SELECT y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2 FROM y WHERE TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE (NOT _u_0._u_1 IS NULL AND ARRAY_ANY(_u_0._u_2, _x -> _x > x.b));

# title: exists with multiple non-equality keys pairs them in a struct
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b AND y.c < x.c);
SELECT x.a FROM x LEFT JOIN (SELECT y.a AS _u_1, ARRAY_AGG(STRUCT(y.b AS _u_2, y.c AS _u_3)) AS _u_4 FROM y WHERE TRUE AND TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE (NOT _u_0._u_1 IS NULL AND ARRAY_ANY(_u_0._u_4, _x -> _x._u_2 > x.b AND _x._u_3 < x.c));

# title: in with a non-equality key is not unnested
SELECT x.a FROM x WHERE x.c IN (SELECT y.c FROM y WHERE y.a = x.a AND y.b > x.b);
SELECT x.a FROM x WHERE x.c IN (SELECT y.c FROM y WHERE y.a = x.a AND y.b > x.b);

# title: exists with multiple non-equality predicates on the same key
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b AND y.b < x.c);
SELECT x.a FROM x LEFT JOIN (SELECT y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2 FROM y WHERE TRUE AND TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE (NOT _u_0._u_1 IS NULL AND ARRAY_ANY(_u_0._u_2, _x -> _x > x.b AND _x < x.c));

# title: exists with a non-equality key that is also the projected value
SELECT x.a FROM x WHERE EXISTS (SELECT y.b AS b FROM y WHERE y.a = x.a AND y.b > x.b);
SELECT x.a FROM x LEFT JOIN (SELECT y.a AS _u_1, ARRAY_AGG(y.b) AS _u_2 FROM y WHERE TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a WHERE (NOT _u_0._u_1 IS NULL AND ARRAY_ANY(_u_0._u_2, _x -> _x > x.b));

# title: exists with a non-equality predicate on an equality key checks it on the join
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.a > x.c);
SELECT x.a FROM x LEFT JOIN (SELECT y.a AS _u_1 FROM y WHERE TRUE AND TRUE GROUP BY y.a) AS _u_0 ON _u_0._u_1 = x.a AND _u_0._u_1 > x.c WHERE NOT _u_0._u_1 IS NULL;

# title: predicate with an inner column on the outer side is not unnested
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b + y.c);
SELECT x.a FROM x WHERE EXISTS(SELECT 1 FROM y WHERE y.a = x.a AND y.b > x.b + y.c);

# title: predicate with an outer column on the key side is not unnested
SELECT x.a FROM x WHERE EXISTS (SELECT 1 FROM y WHERE y.a + x.b = x.a);
SELECT x.a FROM x WHERE EXISTS(SELECT 1 FROM y WHERE y.a + x.b = x.a);
Loading