Customize window frame support for dialect (apache#14288)

* Customize window frame support for dialect (#70)

* Customize window frame support for dialect

* fix: ignore frame only when frame implies no frame

* Add comments. move the window frame determine logic to dialect method

* Update datafusion/sql/src/unparser/dialect.rs

* Update test case

* fix

---------

Co-authored-by: Phillip LeBlanc <[email protected]>

* fix clippy

---------

Co-authored-by: Phillip LeBlanc <[email protected]>

Author: Sevenannn

datafusion/sql/src/unparser/dialect.rs

@@ -23,7 +23,9 @@ use datafusion_expr::Expr;
 use regex::Regex;
 use sqlparser::tokenizer::Span;
 use sqlparser::{
-    ast::{self, BinaryOperator, Function, Ident, ObjectName, TimezoneInfo},
+    ast::{
+        self, BinaryOperator, Function, Ident, ObjectName, TimezoneInfo, WindowFrameBound,
+    },
     keywords::ALL_KEYWORDS,
 };
 
@@ -153,6 +155,18 @@ pub trait Dialect: Send + Sync {
         Ok(None)
     }
 
+    /// Allows the dialect to choose to omit window frame in unparsing
+    /// based on function name and window frame bound
+    /// Returns false if specific function name / window frame bound indicates no window frame is needed in unparsing
+    fn window_func_support_window_frame(
+        &self,
+        _func_name: &str,
+        _start_bound: &WindowFrameBound,
+        _end_bound: &WindowFrameBound,
+    ) -> bool {
+        true
+    }
+
     /// Extends the dialect's default rules for unparsing scalar functions.
     /// This is useful for supporting application-specific UDFs or custom engine extensions.
     fn with_custom_scalar_overrides(
@@ -500,6 +514,7 @@ pub struct CustomDialect {
     supports_column_alias_in_table_alias: bool,
     requires_derived_table_alias: bool,
     division_operator: BinaryOperator,
+    window_func_support_window_frame: bool,
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
 }
@@ -527,6 +542,7 @@ impl Default for CustomDialect {
             supports_column_alias_in_table_alias: true,
             requires_derived_table_alias: false,
             division_operator: BinaryOperator::Divide,
+            window_func_support_window_frame: true,
             full_qualified_col: false,
             unnest_as_table_factor: false,
         }
@@ -634,6 +650,15 @@ impl Dialect for CustomDialect {
         self.division_operator.clone()
     }
 
+    fn window_func_support_window_frame(
+        &self,
+        _func_name: &str,
+        _start_bound: &WindowFrameBound,
+        _end_bound: &WindowFrameBound,
+    ) -> bool {
+        self.window_func_support_window_frame
+    }
+
     fn full_qualified_col(&self) -> bool {
         self.full_qualified_col
     }
@@ -675,6 +700,7 @@ pub struct CustomDialectBuilder {
     supports_column_alias_in_table_alias: bool,
     requires_derived_table_alias: bool,
     division_operator: BinaryOperator,
+    window_func_support_window_frame: bool,
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
 }
@@ -708,6 +734,7 @@ impl CustomDialectBuilder {
             supports_column_alias_in_table_alias: true,
             requires_derived_table_alias: false,
             division_operator: BinaryOperator::Divide,
+            window_func_support_window_frame: true,
             full_qualified_col: false,
             unnest_as_table_factor: false,
         }
@@ -733,6 +760,7 @@ impl CustomDialectBuilder {
                 .supports_column_alias_in_table_alias,
             requires_derived_table_alias: self.requires_derived_table_alias,
             division_operator: self.division_operator,
+            window_func_support_window_frame: self.window_func_support_window_frame,
             full_qualified_col: self.full_qualified_col,
             unnest_as_table_factor: self.unnest_as_table_factor,
         }
@@ -857,6 +885,14 @@ impl CustomDialectBuilder {
         self
     }
 
+    pub fn with_window_func_support_window_frame(
+        mut self,
+        window_func_support_window_frame: bool,
+    ) -> Self {
+        self.window_func_support_window_frame = window_func_support_window_frame;
+        self
+    }
+
     /// Customize the dialect to allow full qualified column names
     pub fn with_full_qualified_col(mut self, full_qualified_col: bool) -> Self {
         self.full_qualified_col = full_qualified_col;

datafusion/sql/src/unparser/expr.rs

@@ -217,18 +217,29 @@ impl Unparser<'_> {
 
                 let start_bound = self.convert_bound(&window_frame.start_bound)?;
                 let end_bound = self.convert_bound(&window_frame.end_bound)?;
+
+                let window_frame = if self.dialect.window_func_support_window_frame(
+                    func_name,
+                    &start_bound,
+                    &end_bound,
+                ) {
+                    Some(ast::WindowFrame {
+                        units,
+                        start_bound,
+                        end_bound: Some(end_bound),
+                    })
+                } else {
+                    None
+                };
+
                 let over = Some(ast::WindowType::WindowSpec(ast::WindowSpec {
                     window_name: None,
                     partition_by: partition_by
                         .iter()
                         .map(|e| self.expr_to_sql_inner(e))
                         .collect::<Result<Vec<_>>>()?,
                     order_by,
-                    window_frame: Some(ast::WindowFrame {
-                        units,
-                        start_bound,
-                        end_bound: Option::from(end_bound),
-                    }),
+                    window_frame,
                 }));
 
                 Ok(ast::Expr::Function(Function {
@@ -1632,6 +1643,7 @@ mod tests {
     use datafusion_functions_aggregate::expr_fn::sum;
     use datafusion_functions_nested::expr_fn::{array_element, make_array};
     use datafusion_functions_nested::map::map;
+    use datafusion_functions_window::rank::rank_udwf;
     use datafusion_functions_window::row_number::row_number_udwf;
 
     use crate::unparser::dialect::{
@@ -2677,6 +2689,39 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_window_func_support_window_frame() -> Result<()> {
+        let default_dialect: Arc<dyn Dialect> =
+            Arc::new(CustomDialectBuilder::new().build());
+
+        let test_dialect: Arc<dyn Dialect> = Arc::new(
+            CustomDialectBuilder::new()
+                .with_window_func_support_window_frame(false)
+                .build(),
+        );
+
+        for (dialect, expected) in [
+            (
+                default_dialect,
+                "rank() OVER (ORDER BY a ASC NULLS FIRST ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)",
+            ),
+            (test_dialect, "rank() OVER (ORDER BY a ASC NULLS FIRST)"),
+        ] {
+            let unparser = Unparser::new(dialect.as_ref());
+            let func = WindowFunctionDefinition::WindowUDF(rank_udwf());
+            let mut window_func = WindowFunction::new(func, vec![]);
+            window_func.order_by = vec![Sort::new(col("a"), true, true)];
+            let expr = Expr::WindowFunction(window_func);
+            let ast = unparser.expr_to_sql(&expr)?;
+
+            let actual = ast.to_string();
+            let expected = expected.to_string();
+
+            assert_eq!(actual, expected);
+        }
+        Ok(())
+    }
+
     #[test]
     fn test_utf8_view_to_sql() -> Result<()> {
         let dialect = CustomDialectBuilder::new()