-
Notifications
You must be signed in to change notification settings - Fork 1.3k
fix(optimizer)!: decorrelation when non-EQ keys are present #8317
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
de5d6e5
7014947
6c7fb5f
ff27b28
b82e529
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 | ||
|
|
@@ -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)): | ||
| return | ||
|
Comment on lines
+208
to
209
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This makes sure non- |
||
|
|
||
| is_subquery_projection = any( | ||
|
|
@@ -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
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. A non- 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 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 NULLThis doesn't even have the |
||
| key_aliases[key] = value.alias | ||
| group_by.append(key) | ||
| else: | ||
|
|
@@ -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
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 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
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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 |
||
| parent_predicate = _replace(parent_predicate, f"NOT {alias} IS NULL") | ||
| elif isinstance(parent_predicate, exp.All): | ||
| assert issubclass(op_type, exp.Binary) | ||
|
|
@@ -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) | ||
|
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, | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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: |
||
|
|
||
| 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; | ||
|
|
@@ -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
Collaborator
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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; | ||
|
|
@@ -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); | ||
Uh oh!
There was an error while loading. Please reload this page.