This is an automated email from the ASF dual-hosted git repository.
comphead pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new d1308f04c1 fix: order by expr rewrite fix (#14486)
d1308f04c1 is described below
commit d1308f04c1e74bf3b46b075fb239b117bfa64449
Author: Andrey Koshchiy <[email protected]>
AuthorDate: Thu Feb 6 20:52:51 2025 +0300
fix: order by expr rewrite fix (#14486)
* fix: rewrite order by on compound expr
* fmt
* revert transform
* test fix
---
datafusion/expr/src/expr_rewriter/order_by.rs | 18 +++++++++++++++---
datafusion/sql/tests/sql_integration.rs | 15 +++++++++++++++
datafusion/sqllogictest/test_files/order.slt | 13 +++++++++++++
3 files changed, 43 insertions(+), 3 deletions(-)
diff --git a/datafusion/expr/src/expr_rewriter/order_by.rs
b/datafusion/expr/src/expr_rewriter/order_by.rs
index 0044b6cf6f..6db9555550 100644
--- a/datafusion/expr/src/expr_rewriter/order_by.rs
+++ b/datafusion/expr/src/expr_rewriter/order_by.rs
@@ -21,7 +21,9 @@ use crate::expr::Alias;
use crate::expr_rewriter::normalize_col;
use crate::{expr::Sort, Cast, Expr, LogicalPlan, TryCast};
-use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::tree_node::{
+ Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
+};
use datafusion_common::{Column, Result};
/// Rewrite sort on aggregate expressions to sort on the column of aggregate
output
@@ -101,8 +103,18 @@ fn rewrite_in_terms_of_projection(
let search_col = Expr::Column(Column::new_unqualified(name));
// look for the column named the same as this expr
- if let Some(found) = proj_exprs.iter().find(|a|
expr_match(&search_col, a)) {
- let found = found.clone();
+ let mut found = None;
+ for proj_expr in &proj_exprs {
+ proj_expr.apply(|e| {
+ if expr_match(&search_col, e) {
+ found = Some(e.clone());
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ Ok(TreeNodeRecursion::Continue)
+ })?;
+ }
+
+ if let Some(found) = found {
return Ok(Transformed::yes(match normalized_expr {
Expr::Cast(Cast { expr: _, data_type }) => Expr::Cast(Cast {
expr: Box::new(found),
diff --git a/datafusion/sql/tests/sql_integration.rs
b/datafusion/sql/tests/sql_integration.rs
index 21b6bf09fa..0883f9a469 100644
--- a/datafusion/sql/tests/sql_integration.rs
+++ b/datafusion/sql/tests/sql_integration.rs
@@ -2511,6 +2511,21 @@ fn select_groupby_orderby() {
FROM person GROUP BY person.birth_date ORDER BY birth_date;
"#;
quick_test(sql, expected);
+
+ // Use columnized `avg(age)` in the order by
+ let sql = r#"SELECT
+ avg(age) + avg(age),
+ date_trunc('month', person.birth_date) AS "birth_date"
+ FROM person GROUP BY person.birth_date ORDER BY avg(age) + avg(age);
+"#;
+
+ let expected =
+ "Sort: avg(person.age) + avg(person.age) ASC NULLS LAST\
+ \n Projection: avg(person.age) + avg(person.age),
date_trunc(Utf8(\"month\"), person.birth_date) AS birth_date\
+ \n Aggregate: groupBy=[[person.birth_date]],
aggr=[[avg(person.age)]]\
+ \n TableScan: person";
+
+ quick_test(sql, expected);
}
fn logical_plan(sql: &str) -> Result<LogicalPlan> {
diff --git a/datafusion/sqllogictest/test_files/order.slt
b/datafusion/sqllogictest/test_files/order.slt
index d95001b81b..d7da21c58e 100644
--- a/datafusion/sqllogictest/test_files/order.slt
+++ b/datafusion/sqllogictest/test_files/order.slt
@@ -384,6 +384,19 @@ ORDER BY time;
2 2022-01-01T01:00:00
3 2022-01-02T00:00:00
+# Tests for https://github.com/apache/datafusion/issues/14459
+query PI
+select
+ date_trunc('minute',time) AS "time",
+ sum(value) + sum(value)
+FROM t
+GROUP BY time
+ORDER BY sum(value) + sum(value);
+----
+2022-01-01T00:00:00 2
+2022-01-01T01:00:00 4
+2022-01-02T00:00:00 6
+
## SORT BY is not supported
statement error DataFusion error: This feature is not implemented: SORT BY
select * from t SORT BY time;
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
