This is an automated email from the ASF dual-hosted git repository.
alamb 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 5d3bbaa323 Unparser: Support `ORDER BY` in window function definition
(#10370)
5d3bbaa323 is described below
commit 5d3bbaa3233523819d823e69139cd522db715cae
Author: Junhao Liu <[email protected]>
AuthorDate: Sat May 4 05:15:40 2024 -0600
Unparser: Support `ORDER BY` in window function definition (#10370)
* support OrderBy in Unparser
* resolve comment
---
datafusion/sql/src/unparser/expr.rs | 108 +++++++++++++++++++++++++++++++++---
1 file changed, 99 insertions(+), 9 deletions(-)
diff --git a/datafusion/sql/src/unparser/expr.rs
b/datafusion/sql/src/unparser/expr.rs
index 7194b0a7d8..c619c62668 100644
--- a/datafusion/sql/src/unparser/expr.rs
+++ b/datafusion/sql/src/unparser/expr.rs
@@ -15,10 +15,14 @@
// specific language governing permissions and limitations
// under the License.
+use core::fmt;
+use std::fmt::Display;
+
use arrow_array::{Date32Array, Date64Array};
use arrow_schema::DataType;
use datafusion_common::{
- internal_datafusion_err, not_impl_err, plan_err, Column, Result,
ScalarValue,
+ internal_datafusion_err, internal_err, not_impl_err, plan_err, Column,
Result,
+ ScalarValue,
};
use datafusion_expr::{
expr::{Alias, Exists, InList, ScalarFunction, Sort, WindowFunction},
@@ -30,10 +34,38 @@ use sqlparser::ast::{
use super::Unparser;
+/// DataFusion's Exprs can represent either an `Expr` or an `OrderByExpr`
+pub enum Unparsed {
+ // SQL Expression
+ Expr(ast::Expr),
+ // SQL ORDER BY expression (e.g. `col ASC NULLS FIRST`)
+ OrderByExpr(ast::OrderByExpr),
+}
+
+impl Unparsed {
+ pub fn into_order_by_expr(self) -> Result<ast::OrderByExpr> {
+ if let Unparsed::OrderByExpr(order_by_expr) = self {
+ Ok(order_by_expr)
+ } else {
+ internal_err!("Expected Sort expression to be converted an
OrderByExpr")
+ }
+ }
+}
+
+impl Display for Unparsed {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ match self {
+ Unparsed::Expr(expr) => write!(f, "{}", expr),
+ Unparsed::OrderByExpr(order_by_expr) => write!(f, "{}",
order_by_expr),
+ }
+ }
+}
+
/// Convert a DataFusion [`Expr`] to `sqlparser::ast::Expr`
///
/// This function is the opposite of `SqlToRel::sql_to_expr` and can
-/// be used to, among other things, convert `Expr`s to strings.
+/// be used to, among other things, convert [`Expr`]s to strings.
+/// Throws an error if [`Expr`] can not be represented by an
`sqlparser::ast::Expr`
///
/// # Example
/// ```
@@ -49,6 +81,15 @@ pub fn expr_to_sql(expr: &Expr) -> Result<ast::Expr> {
unparser.expr_to_sql(expr)
}
+/// Convert a DataFusion [`Expr`] to [`Unparsed`]
+///
+/// This function is similar to expr_to_sql, but it supports converting more
[`Expr`] types like
+/// `Sort` expressions to `OrderByExpr` expressions.
+pub fn expr_to_unparsed(expr: &Expr) -> Result<Unparsed> {
+ let unparser = Unparser::default();
+ unparser.expr_to_unparsed(expr)
+}
+
impl Unparser<'_> {
pub fn expr_to_sql(&self, expr: &Expr) -> Result<ast::Expr> {
match expr {
@@ -170,7 +211,7 @@ impl Unparser<'_> {
fun,
args,
partition_by,
- order_by: _,
+ order_by,
window_frame,
null_treatment: _,
}) => {
@@ -189,6 +230,11 @@ impl Unparser<'_> {
ast::WindowFrameUnits::Groups
}
};
+ let order_by: Vec<ast::OrderByExpr> = order_by
+ .iter()
+ .map(|expr| expr_to_unparsed(expr)?.into_order_by_expr())
+ .collect::<Result<Vec<_>>>()?;
+
let start_bound =
self.convert_bound(&window_frame.start_bound);
let end_bound = self.convert_bound(&window_frame.end_bound);
let over = Some(ast::WindowType::WindowSpec(ast::WindowSpec {
@@ -197,13 +243,14 @@ impl Unparser<'_> {
.iter()
.map(|e| self.expr_to_sql(e))
.collect::<Result<Vec<_>>>()?,
- order_by: vec![],
+ order_by,
window_frame: Some(ast::WindowFrame {
units,
start_bound,
end_bound: Option::from(end_bound),
}),
}));
+
Ok(ast::Expr::Function(Function {
name: ast::ObjectName(vec![Ident {
value: func_name.to_string(),
@@ -305,10 +352,10 @@ impl Unparser<'_> {
})
}
Expr::Sort(Sort {
- expr,
+ expr: _,
asc: _,
nulls_first: _,
- }) => self.expr_to_sql(expr),
+ }) => plan_err!("Sort expression should be handled by
expr_to_unparsed"),
Expr::IsNotNull(expr) => {
Ok(ast::Expr::IsNotNull(Box::new(self.expr_to_sql(expr)?)))
}
@@ -366,6 +413,29 @@ impl Unparser<'_> {
}
}
+ /// This function can convert more [`Expr`] types than `expr_to_sql`,
returning an [`Unparsed`]
+ /// like `Sort` expressions to `OrderByExpr` expressions.
+ pub fn expr_to_unparsed(&self, expr: &Expr) -> Result<Unparsed> {
+ match expr {
+ Expr::Sort(Sort {
+ expr,
+ asc,
+ nulls_first,
+ }) => {
+ let sql_parser_expr = self.expr_to_sql(expr)?;
+ Ok(Unparsed::OrderByExpr(ast::OrderByExpr {
+ expr: sql_parser_expr,
+ asc: Some(*asc),
+ nulls_first: Some(*nulls_first),
+ }))
+ }
+ _ => {
+ let sql_parser_expr = self.expr_to_sql(expr)?;
+ Ok(Unparsed::Expr(sql_parser_expr))
+ }
+ }
+ }
+
fn col_to_sql(&self, col: &Column) -> Result<ast::Expr> {
if let Some(table_ref) = &col.relation {
let mut id = table_ref.to_vec();
@@ -992,7 +1062,11 @@ mod tests {
),
args: vec![wildcard()],
partition_by: vec![],
- order_by: vec![],
+ order_by: vec![Expr::Sort(Sort::new(
+ Box::new(col("a")),
+ false,
+ true,
+ ))],
window_frame: WindowFrame::new_bounds(
datafusion_expr::WindowFrameUnits::Range,
datafusion_expr::WindowFrameBound::Preceding(
@@ -1004,7 +1078,7 @@ mod tests {
),
null_treatment: None,
}),
- r#"COUNT(*) OVER (RANGE BETWEEN 6 PRECEDING AND 2 FOLLOWING)"#,
+ r#"COUNT(*) OVER (ORDER BY "a" DESC NULLS FIRST RANGE BETWEEN
6 PRECEDING AND 2 FOLLOWING)"#,
),
(col("a").is_not_null(), r#""a" IS NOT NULL"#),
(
@@ -1041,7 +1115,6 @@ mod tests {
not_exists(Arc::new(dummy_logical_plan.clone())),
r#"NOT EXISTS (SELECT "t"."a" FROM "t" WHERE ("t"."a" = 1))"#,
),
- (col("a").sort(true, true), r#""a""#),
];
for (expr, expected) in tests {
@@ -1055,6 +1128,23 @@ mod tests {
Ok(())
}
+ #[test]
+ fn expr_to_unparsed_ok() -> Result<()> {
+ let tests: Vec<(Expr, &str)> = vec![
+ ((col("a") + col("b")).gt(lit(4)), r#"(("a" + "b") > 4)"#),
+ (col("a").sort(true, true), r#""a" ASC NULLS FIRST"#),
+ ];
+
+ for (expr, expected) in tests {
+ let ast = expr_to_unparsed(&expr)?;
+
+ let actual = format!("{}", ast);
+
+ assert_eq!(actual, expected);
+ }
+
+ Ok(())
+ }
#[test]
fn custom_dialect() -> Result<()> {
let dialect = CustomDialect::new(Some('\''));
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
