Bigquery col alias override

datafusion/sql/src/unparser/dialect.rs

@@ -19,6 +19,7 @@ use std::{collections::HashMap, sync::Arc};
 
 use super::{utils::character_length_to_sql, utils::date_part_to_sql, Unparser};
 use arrow::datatypes::TimeUnit;
+use datafusion_common::alias::AliasGenerator;
 use datafusion_common::Result;
 use datafusion_expr::Expr;
 use regex::Regex;
@@ -154,6 +155,14 @@ pub trait Dialect: Send + Sync {
         Ok(None)
     }
 
+    fn col_alias_overrides(
+        &self,
+        _unparser: &Unparser,
+        _alias: &str,
+    ) -> Result<Option<String>> {
+        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
@@ -253,7 +262,44 @@ impl Dialect for DefaultDialect {
         }
     }
 }
+#[derive(Default)]
+pub struct BigQueryDialect {
+    col_alias_generator: AliasGenerator,
+}
+
+impl Dialect for BigQueryDialect {
+    fn identifier_quote_style(&self, _: &str) -> Option<char> {
+        Some('`')
+    }
+
+    fn col_alias_overrides(
+        &self,
+        _unparser: &Unparser,
+        alias: &str,
+    ) -> Result<Option<String>> {
+        // Check if alias contains any special characters not supported by BigQuery col names
+        // https://cloud.google.com/bigquery/docs/schemas#flexible-column-names
+        if alias.contains(
+            &[
+                '!', '"', '$', '(', ')', '*', ',', '.', '/', ';', '?', '@', '[', '\\',
+                ']', '^', '`', '{', '}', '~',
+            ][..],
+        ) {
+            Ok(Some(self.col_alias_generator.next("col")))
+        } else {
+            Ok(Some(alias.to_string()))
+        }
+    }
+}
 
+impl BigQueryDialect {
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            col_alias_generator: AliasGenerator::new(),
+        }
+    }
+}
 pub struct PostgreSqlDialect {}
 
 impl Dialect for PostgreSqlDialect {

datafusion/sql/src/unparser/plan.rs

@@ -1035,9 +1035,18 @@ impl Unparser<'_> {
             Expr::Alias(Alias { expr, name, .. }) => {
                 let inner = self.expr_to_sql(expr)?;
 
+                // Determine the alias name to use
+                let alias_name = if let Some(new_name) =
+                    self.dialect.col_alias_overrides(self, name)?
+                {
+                    new_name.to_string()
+                } else {
+                    name.to_string()
+                };
+
                 Ok(ast::SelectItem::ExprWithAlias {
                     expr: inner,
-                    alias: self.new_ident_quoted_if_needs(name.to_string()),
+                    alias: self.new_ident_quoted_if_needs(alias_name),
                 })
             }
             _ => {

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -31,8 +31,9 @@ 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::{
-    CustomDialectBuilder, DefaultDialect as UnparserDefaultDialect, DefaultDialect,
-    Dialect as UnparserDialect, MySqlDialect as UnparserMySqlDialect, SqliteDialect,
+    BigQueryDialect as UnparserBigqueryDialect, CustomDialectBuilder,
+    DefaultDialect as UnparserDefaultDialect, DefaultDialect, Dialect as UnparserDialect,
+    MySqlDialect as UnparserMySqlDialect, SqliteDialect,
 };
 use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser};
 use sqlparser::ast::Statement;
@@ -54,7 +55,7 @@ use datafusion_sql::unparser::extension_unparser::{
     UnparseToStatementResult, UnparseWithinStatementResult,
     UserDefinedLogicalNodeUnparser,
 };
-use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect};
+use sqlparser::dialect::{Dialect, GenericDialect, MySqlDialect, PostgreSqlDialect};
 use sqlparser::parser::Parser;
 
 #[test]
@@ -536,6 +537,24 @@ fn roundtrip_statement_with_dialect() -> Result<()> {
             parser_dialect: Box::new(GenericDialect {}),
             unparser_dialect: Box::new(SqliteDialect {}),
         },
+        TestStatementWithDialect {
+            sql: "select min(*) as \"min(*)\" from (select 1 as a)",
+            expected: "SELECT min(*) AS `col_1` FROM (SELECT 1 AS `a`)",
+            parser_dialect: Box::new(PostgreSqlDialect {}),
+            unparser_dialect: Box::new(UnparserBigqueryDialect::new()),
+        },
+        TestStatementWithDialect {
+            sql: "select a as \"a*\", b as \"b@\" from (select 1 as a , 2 as b)",
+            expected: "SELECT `a` AS `col_1`, `b` AS `col_2` FROM (SELECT 1 AS `a`, 2 AS `b`)",
+            parser_dialect: Box::new(PostgreSqlDialect {}),
+            unparser_dialect: Box::new(UnparserBigqueryDialect::new()),
+        },
+        TestStatementWithDialect {
+            sql: "select a as \"a*\", b , c as \"c@\" from (select 1 as a , 2 as b, 3 as c)",
+            expected: "SELECT `a` AS `col_1`, `b`, `c` AS `col_2` FROM (SELECT 1 AS `a`, 2 AS `b`, 3 AS `c`)",
+            parser_dialect: Box::new(PostgreSqlDialect {}),
+            unparser_dialect: Box::new(UnparserBigqueryDialect::new()),
+        },
         TestStatementWithDialect {
             sql: "SELECT * FROM UNNEST([1,2,3])",
             expected: r#"SELECT * FROM (SELECT UNNEST([1, 2, 3]) AS "UNNEST(make_array(Int64(1),Int64(2),Int64(3)))")"#,