add support_literal_group_by_key for unparser dialect (apache#10)

Author: goldmedal

datafusion/sql/src/unparser/dialect.rs

@@ -245,6 +245,12 @@ pub trait Dialect: Send + Sync {
     fn to_unicode_string_literal(&self, _s: &str) -> Option<ast::Expr> {
         None
     }
+
+    /// Whether the dialect supports using literal values as GROUP BY keys.
+    /// Some dialects like BigQuery do not support using literal values as GROUP BY keys.
+    fn support_literal_group_by_key(&self) -> bool {
+        true
+    }
 }
 
 /// `IntervalStyle` to use for unparsing
@@ -582,6 +588,10 @@ impl Dialect for BigQueryDialect {
     fn unnest_as_table_factor(&self) -> bool {
         true
     }
+
+    fn support_literal_group_by_key(&self) -> bool {
+        false
+    }
 }
 
 impl BigQueryDialect {
@@ -786,6 +796,7 @@ pub struct CustomDialect {
     window_func_support_window_frame: bool,
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
+    support_literal_group_by_key: bool,
 }
 
 impl Default for CustomDialect {
@@ -814,6 +825,7 @@ impl Default for CustomDialect {
             window_func_support_window_frame: true,
             full_qualified_col: false,
             unnest_as_table_factor: false,
+            support_literal_group_by_key: true,
         }
     }
 }
@@ -935,6 +947,10 @@ impl Dialect for CustomDialect {
     fn unnest_as_table_factor(&self) -> bool {
         self.unnest_as_table_factor
     }
+
+    fn support_literal_group_by_key(&self) -> bool {
+        self.support_literal_group_by_key
+    }
 }
 
 /// `CustomDialectBuilder` to build `CustomDialect` using builder pattern
@@ -973,6 +989,7 @@ pub struct CustomDialectBuilder {
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
     unnest_to_flattened_table_factor: bool,
+    support_literal_group_by_key: bool,
 }
 
 impl Default for CustomDialectBuilder {
@@ -1008,6 +1025,7 @@ impl CustomDialectBuilder {
             full_qualified_col: false,
             unnest_as_table_factor: false,
             unnest_to_flattened_table_factor: false,
+            support_literal_group_by_key: true,
         }
     }
 
@@ -1034,6 +1052,7 @@ impl CustomDialectBuilder {
             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,
+            support_literal_group_by_key: self.support_literal_group_by_key,
         }
     }
 
@@ -1182,4 +1201,12 @@ impl CustomDialectBuilder {
         self.unnest_to_flattened_table_factor = unnest_to_flattened_table_factor;
         self
     }
+
+    pub fn with_support_literal_group_by_key(
+        mut self,
+        support_literal_group_by_key: bool,
+    ) -> Self {
+        self.support_literal_group_by_key = support_literal_group_by_key;
+        self
+    }
 }

datafusion/sql/src/unparser/plan.rs

@@ -252,6 +252,10 @@ impl Unparser<'_> {
                 select.group_by(ast::GroupByExpr::Expressions(
                     agg.group_expr
                         .iter()
+                        .filter(|expr| {
+                            self.dialect.support_literal_group_by_key()
+                                || !matches!(expr, Expr::Literal(_, _))
+                        })
                         .map(|expr| self.expr_to_sql(expr))
                         .collect::<Result<Vec<_>>>()?,
                     vec![],
@@ -556,6 +560,10 @@ impl Unparser<'_> {
                     select.group_by(ast::GroupByExpr::Expressions(
                         agg.group_expr
                             .iter()
+                            .filter(|expr| {
+                                self.dialect.support_literal_group_by_key()
+                                    || !matches!(expr, Expr::Literal(_, _))
+                            })
                             .map(|expr| self.expr_to_sql(expr))
                             .collect::<Result<Vec<_>>>()?,
                         vec![],

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -34,9 +34,10 @@ use datafusion_functions_nested::map::map_udf;
 use datafusion_functions_window::rank::rank_udwf;
 use datafusion_sql::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_sql::unparser::dialect::{
-    BigQueryDialect, CustomDialectBuilder, DefaultDialect as UnparserDefaultDialect,
-    DefaultDialect, Dialect as UnparserDialect, MySqlDialect as UnparserMySqlDialect,
-    PostgreSqlDialect as UnparserPostgreSqlDialect, SnowflakeDialect, SqliteDialect,
+    BigQueryDialect, CustomDialect, CustomDialectBuilder,
+    DefaultDialect as UnparserDefaultDialect, DefaultDialect, Dialect as UnparserDialect,
+    MySqlDialect as UnparserMySqlDialect, PostgreSqlDialect as UnparserPostgreSqlDialect,
+    SnowflakeDialect, SqliteDialect,
 };
 use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser};
 use insta::assert_snapshot;
@@ -2708,3 +2709,37 @@ fn test_struct_expr3() {
         @r#"SELECT test.c1."metadata".product."name" FROM (SELECT {"metadata": {product: {"name": 'Product Name'}}} AS c1) AS test"#
     );
 }
+
+#[test]
+fn test_literal_gbk() -> Result<(), DataFusionError> {
+    let sql = r#"
+        select 'ABC' as col1, first_name, min(id) from person group by 1, 2
+
+    "#;
+    let unparser = CustomDialectBuilder::new()
+        .with_support_literal_group_by_key(false)
+        .build();
+    roundtrip_statement_with_dialect_helper!(
+        sql: sql,
+        parser_dialect: GenericDialect {},
+        unparser_dialect: unparser,
+        expected: @r#"SELECT 'ABC' AS col1, person.first_name, min(person.id) FROM person GROUP BY person.first_name"#,
+    );
+
+    let unparser = BigQueryDialect {};
+    roundtrip_statement_with_dialect_helper!(
+        sql: sql,
+        parser_dialect: GenericDialect {},
+        unparser_dialect: unparser,
+        expected: @r#"SELECT 'ABC' AS `col1`, `person`.`first_name`, min(`person`.`id`) FROM `person` GROUP BY `person`.`first_name`"#,
+    );
+
+    let unparser = CustomDialect::default();
+    roundtrip_statement_with_dialect_helper!(
+        sql: sql,
+        parser_dialect: GenericDialect {},
+        unparser_dialect: unparser,
+        expected: @r#"SELECT 'ABC' AS col1, person.first_name, min(person.id) FROM person GROUP BY 'ABC', person.first_name"#,
+    );
+    Ok(())
+}