Introduce full_qualified_col option for the unparser dialect (#13241)

* introduce `full_qualified_col` for unparser dialect * fix test and clippy
apache · Nov 5, 2024 · 7c6f891 · 7c6f891
1 parent e8520ab
commit 7c6f891
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 9 deletions.
diff --git a/datafusion/sql/src/unparser/dialect.rs b/datafusion/sql/src/unparser/dialect.rs
@@ -137,6 +137,14 @@ pub trait Dialect: Send + Sync {
     ) -> Result<Option<ast::Expr>> {
         Ok(None)
     }
+
+    /// Allow to unparse a qualified column with a full qualified name
+    /// (e.g. catalog_name.schema_name.table_name.column_name)
+    /// Otherwise, the column will be unparsed with only the table name and colum name
+    /// (e.g. table_name.column_name)
+    fn full_qualified_col(&self) -> bool {
+        false
+    }
 }
 
 /// `IntervalStyle` to use for unparsing
@@ -373,6 +381,7 @@ pub struct CustomDialect {
     date32_cast_dtype: ast::DataType,
     supports_column_alias_in_table_alias: bool,
     requires_derived_table_alias: bool,
+    full_qualified_col: bool,
 }
 
 impl Default for CustomDialect {
@@ -396,6 +405,7 @@ impl Default for CustomDialect {
             date32_cast_dtype: ast::DataType::Date,
             supports_column_alias_in_table_alias: true,
             requires_derived_table_alias: false,
+            full_qualified_col: false,
         }
     }
 }
@@ -488,6 +498,10 @@ impl Dialect for CustomDialect {
     fn requires_derived_table_alias(&self) -> bool {
         self.requires_derived_table_alias
     }
+
+    fn full_qualified_col(&self) -> bool {
+        self.full_qualified_col
+    }
 }
 
 /// `CustomDialectBuilder` to build `CustomDialect` using builder pattern
@@ -520,6 +534,7 @@ pub struct CustomDialectBuilder {
     date32_cast_dtype: ast::DataType,
     supports_column_alias_in_table_alias: bool,
     requires_derived_table_alias: bool,
+    full_qualified_col: bool,
 }
 
 impl Default for CustomDialectBuilder {
@@ -549,6 +564,7 @@ impl CustomDialectBuilder {
             date32_cast_dtype: ast::DataType::Date,
             supports_column_alias_in_table_alias: true,
             requires_derived_table_alias: false,
+            full_qualified_col: false,
         }
     }
 
@@ -570,6 +586,7 @@ impl CustomDialectBuilder {
             supports_column_alias_in_table_alias: self
                 .supports_column_alias_in_table_alias,
             requires_derived_table_alias: self.requires_derived_table_alias,
+            full_qualified_col: self.full_qualified_col,
         }
     }
 
@@ -677,4 +694,10 @@ impl CustomDialectBuilder {
         self.requires_derived_table_alias = requires_derived_table_alias;
         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;
+        self
+    }
 }
diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs
@@ -527,7 +527,11 @@ impl Unparser<'_> {
 
     fn col_to_sql(&self, col: &Column) -> Result<ast::Expr> {
         if let Some(table_ref) = &col.relation {
-            let mut id = table_ref.to_vec();
+            let mut id = if self.dialect.full_qualified_col() {
+                table_ref.to_vec()
+            } else {
+                vec![table_ref.table().to_string()]
+            };
             id.push(col.name.to_string());
             return Ok(ast::Expr::CompoundIdentifier(
                 id.iter()
@@ -1545,7 +1549,7 @@ mod tests {
                     name: "c".to_string(),
                 })
                 .gt(lit(4)),
-                r#"(a.b.c > 4)"#,
+                r#"(b.c > 4)"#,
             ),
             (
                 case(col("a"))
@@ -1882,7 +1886,7 @@ mod tests {
                         name: "array_col".to_string(),
                     })),
                 }),
-                r#"UNNEST("schema"."table".array_col)"#,
+                r#"UNNEST("table".array_col)"#,
             ),
         ];
 

diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -28,8 +28,8 @@ use datafusion_functions_nested::make_array::make_array_udf;
 use datafusion_functions_window::rank::rank_udwf;
 use datafusion_sql::planner::{ContextProvider, PlannerContext, SqlToRel};
 use datafusion_sql::unparser::dialect::{
-    DefaultDialect as UnparserDefaultDialect, Dialect as UnparserDialect,
-    MySqlDialect as UnparserMySqlDialect, SqliteDialect,
+    CustomDialectBuilder, DefaultDialect as UnparserDefaultDialect, DefaultDialect,
+    Dialect as UnparserDialect, MySqlDialect as UnparserMySqlDialect, SqliteDialect,
 };
 use datafusion_sql::unparser::{expr_to_sql, plan_to_sql, Unparser};
 
@@ -565,7 +565,7 @@ Projection: unnest_placeholder(unnest_table.struct_col).field1, unnest_placehold
 
 #[test]
 fn test_table_references_in_plan_to_sql() {
-    fn test(table_name: &str, expected_sql: &str) {
+    fn test(table_name: &str, expected_sql: &str, dialect: &impl UnparserDialect) {
         let schema = Schema::new(vec![
             Field::new("id", DataType::Utf8, false),
             Field::new("value", DataType::Utf8, false),
@@ -576,22 +576,48 @@ fn test_table_references_in_plan_to_sql() {
             .unwrap()
             .build()
             .unwrap();
-        let sql = plan_to_sql(&plan).unwrap();
+
+        let unparser = Unparser::new(dialect);
+        let sql = unparser.plan_to_sql(&plan).unwrap();
 
         assert_eq!(sql.to_string(), expected_sql)
     }
 
     test(
         "catalog.schema.table",
-        r#"SELECT "catalog"."schema"."table".id, "catalog"."schema"."table"."value" FROM "catalog"."schema"."table""#,
+        r#"SELECT "table".id, "table"."value" FROM "catalog"."schema"."table""#,
+        &DefaultDialect {},
     );
     test(
         "schema.table",
-        r#"SELECT "schema"."table".id, "schema"."table"."value" FROM "schema"."table""#,
+        r#"SELECT "table".id, "table"."value" FROM "schema"."table""#,
+        &DefaultDialect {},
     );
     test(
         "table",
         r#"SELECT "table".id, "table"."value" FROM "table""#,
+        &DefaultDialect {},
+    );
+
+    let custom_dialect = CustomDialectBuilder::default()
+        .with_full_qualified_col(true)
+        .with_identifier_quote_style('"')
+        .build();
+
+    test(
+        "catalog.schema.table",
+        r#"SELECT "catalog"."schema"."table"."id", "catalog"."schema"."table"."value" FROM "catalog"."schema"."table""#,
+        &custom_dialect,
+    );
+    test(
+        "schema.table",
+        r#"SELECT "schema"."table"."id", "schema"."table"."value" FROM "schema"."table""#,
+        &custom_dialect,
+    );
+    test(
+        "table",
+        r#"SELECT "table"."id", "table"."value" FROM "table""#,
+        &custom_dialect,
     );
 }