Skip to content

Avoid Aliased Window Expr Enter Unreachable Code #14109

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.

Sign up for GitHub

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

Merged
merged 1 commit into from
Jan 14, 2025
Merged

Avoid Aliased Window Expr Enter Unreachable Code #14109

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
38 changes: 25 additions & 13 deletions datafusion/optimizer/src/push_down_filter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ use datafusion_common::tree_node::{
use datafusion_common::{
internal_err, plan_err, qualified_name, Column, DFSchema, Result,
};
use datafusion_expr::expr::WindowFunction;
use datafusion_expr::expr_rewriter::replace_col;
use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, TableScan, Union};
use datafusion_expr::utils::{
Expand Down Expand Up @@ -1001,23 +1002,34 @@ impl OptimizerRule for PushDownFilter {
// multiple window functions, each with potentially different partition keys.
// Therefore, we need to ensure that any potential partition key returned is used in
// ALL window functions. Otherwise, filters cannot be pushed by through that column.
let extract_partition_keys = |func: &WindowFunction| {
func.partition_by
.iter()
.map(|c| Column::from_qualified_name(c.schema_name().to_string()))
.collect::<HashSet<_>>()
};
let potential_partition_keys = window
.window_expr
.iter()
.map(|e| {
if let Expr::WindowFunction(window_expression) = e {
window_expression
.partition_by
.iter()
.map(|c| {
Column::from_qualified_name(
c.schema_name().to_string(),
)
})
.collect::<HashSet<_>>()
} else {
// window functions expressions are only Expr::WindowFunction
unreachable!()
match e {
Expr::WindowFunction(window_func) => {
extract_partition_keys(window_func)
}
Expr::Alias(alias) => {
if let Expr::WindowFunction(window_func) =
alias.expr.as_ref()
{
extract_partition_keys(window_func)
} else {
// window functions expressions are only Expr::WindowFunction
unreachable!()
Copy link
Contributor

Choose a reason for hiding this comment

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

maybe we can turn this into returning an internal error rather than a panic for safety

Copy link
Contributor

Choose a reason for hiding this comment

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

Sounds good. I think there are few of those in the codebase and it would make a good first issue to remove these unreachable panics and replace them with internal errors. Let's create an issue for this

}
}
_ => {
// window functions expressions are only Expr::WindowFunction
unreachable!()
}
}
})
// performs the set intersection of the partition keys of all window functions,
Expand Down
45 changes: 42 additions & 3 deletions datafusion/sqllogictest/test_files/subquery.slt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1306,7 +1306,6 @@ select struct(1, 'b')
----
{c0: 1, c1: b}


query T
select (select struct(1, 'b')['c1']);
----
Expand All @@ -1330,7 +1329,6 @@ WHERE 1+2 = 3 AND column1 IN (SELECT struct(1, 'b')['c0']);
----
1


query I
SELECT * FROM foo
WHERE EXISTS (SELECT * FROM (values (1)) WHERE column1 = foo.x AND struct(1, 'b')['c0'] = 1);
Expand All @@ -1344,6 +1342,47 @@ WHERE 1+2 = 3 AND EXISTS (SELECT * FROM (values (1)) WHERE column1 = foo.x AND s
----
1


statement ok
drop table foo;


# Test for window alias in subquery

# Setup source table
statement ok
CREATE TABLE source_table (
column1 TEXT,
column2 TIMESTAMP,
column3 FLOAT
);

statement ok
INSERT INTO source_table VALUES
('item1', TIMESTAMP '1970-01-01 00:00:01', 50.0),
('item2', TIMESTAMP '1970-01-01 00:00:02', 30.0),
('item1', TIMESTAMP '1970-01-01 00:00:03', 25.0);

# Execute the query
query TPR
WITH SubQuery AS (
SELECT
a.column1,
a.column2 AS ts_column,
a.column3,
SUM(a.column3) OVER (
PARTITION BY a.column1
ORDER BY a.column2 RANGE BETWEEN INTERVAL '10 minutes' PRECEDING AND CURRENT ROW
) AS moving_sum
FROM source_table a
)
SELECT
column1,
ts_column,
moving_sum
FROM SubQuery
WHERE moving_sum > 60;
----
item1 1970-01-01T00:00:03 75

statement ok
drop table source_table;
Loading