fix: Remove dangling table references in unparser

datafusion/sql/src/unparser/ast.rs

@@ -25,6 +25,8 @@ use core::fmt;
 
 use sqlparser::ast;
 
+use super::rewrite::remove_dangling_expr;
+
 #[derive(Clone)]
 pub(super) struct QueryBuilder {
     with: Option<ast::With>,
@@ -238,7 +240,101 @@ impl SelectBuilder {
         self.value_table_mode = value;
         self
     }
-    pub fn build(&self) -> Result<ast::Select, BuilderError> {
+    fn collect_valid_idents(&self, relation_builder: &RelationBuilder) -> Vec<String> {
+        let mut all_idents = Vec::new();
+        if let Some(source_alias) = relation_builder.get_alias() {
+            all_idents.push(source_alias);
+        } else if let Some(source_name) = relation_builder.get_name() {
+            let full_ident = source_name.to_string();
+            if let Some(name) = source_name.0.last() {
+                if full_ident != name.to_string() {
+                    // supports identifiers that contain the entire path, as well as just the end table leaf
+                    // like catalog.schema.table and table
+                    all_idents.push(name.to_string());
+                }
+            }
+            all_idents.push(full_ident);
+        }
+
+        if let Some(twg) = self.from.last() {
+            twg.joins.iter().for_each(|join| match &join.relation {
+                ast::TableFactor::Table { alias, name, .. } => {
+                    if let Some(alias) = alias {
+                        all_idents.push(alias.name.to_string());
+                    } else {
+                        let full_ident = name.to_string();
+                        if let Some(name) = name.0.last() {
+                            if full_ident != name.to_string() {
+                                // supports identifiers that contain the entire path, as well as just the end table leaf
+                                // like catalog.schema.table and table
+                                all_idents.push(name.to_string());
+                            }
+                        }
+                        all_idents.push(full_ident);
+                    }
+                }
+                ast::TableFactor::Derived {
+                    alias: Some(alias), ..
+                } => {
+                    all_idents.push(alias.name.to_string());
+                }
+                _ => {}
+            });
+        }
+
+        all_idents
+    }
+
+    /// Remove any dangling table identifiers from the projection, selection, group by, order by and function arguments
+    /// This removes any references to tables that are not part of any from/source or join, as they would be invalid
+    fn remove_dangling_identifiers(
+        &mut self,
+        query: &mut Option<QueryBuilder>,
+        relation_builder: &RelationBuilder,
+    ) {
+        let all_idents = self.collect_valid_idents(relation_builder);
+
+        // Ensure that the projection contains references to sources that actually exist
+        self.projection.iter_mut().for_each(|select_item| {
+            if let ast::SelectItem::UnnamedExpr(expr) = select_item {
+                *expr = remove_dangling_expr(expr.clone(), &all_idents);
+            }
+        });
+
+        // replace dangling references in the selection
+        if let Some(expr) = self.selection.as_ref() {
+            self.selection = Some(remove_dangling_expr(expr.clone(), &all_idents));
+        }
+
+        // Check the order by as well
+        if let Some(query) = query.as_mut() {
+            query.order_by.iter_mut().for_each(|sort_item| {
+                sort_item.expr =
+                    remove_dangling_expr(sort_item.expr.clone(), &all_idents);
+            });
+        }
+
+        // Order by could be a sort in the select builder
+        self.sort_by.iter_mut().for_each(|sort_item| {
+            *sort_item = remove_dangling_expr(sort_item.clone(), &all_idents);
+        });
+
+        // check the group by as well
+        if let Some(ast::GroupByExpr::Expressions(ref mut group_by, _)) =
+            self.group_by.as_mut()
+        {
+            group_by.iter_mut().for_each(|expr| {
+                *expr = remove_dangling_expr(expr.clone(), &all_idents);
+            });
+        }
+    }
+    pub fn build(
+        &mut self,
+        query: &mut Option<QueryBuilder>,
+        relation_builder: &RelationBuilder,
+    ) -> Result<ast::Select, BuilderError> {
+        self.remove_dangling_identifiers(query, relation_builder);
+
         Ok(ast::Select {
             distinct: self.distinct.clone(),
             top: self.top.clone(),
@@ -307,7 +403,6 @@ impl TableWithJoinsBuilder {
         self.relation = Some(value);
         self
     }
-
     pub fn joins(&mut self, value: Vec<ast::Join>) -> &mut Self {
         self.joins = value;
         self
@@ -360,6 +455,23 @@ impl RelationBuilder {
     pub fn has_relation(&self) -> bool {
         self.relation.is_some()
     }
+    pub fn get_name(&self) -> Option<&ast::ObjectName> {
+        match self.relation {
+            Some(TableFactorBuilder::Table(ref value)) => value.name.as_ref(),
+            _ => None,
+        }
+    }
+    pub fn get_alias(&self) -> Option<String> {
+        match self.relation {
+            Some(TableFactorBuilder::Table(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            Some(TableFactorBuilder::Derived(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            _ => None,
+        }
+    }
     pub fn table(&mut self, value: TableRelationBuilder) -> &mut Self {
         self.relation = Some(TableFactorBuilder::Table(value));
         self

datafusion/sql/src/unparser/plan.rs

@@ -158,10 +158,12 @@ impl Unparser<'_> {
         }
 
         let mut twj = select_builder.pop_from().unwrap();
-        twj.relation(relation_builder);
+        twj.relation(relation_builder.clone());
         select_builder.push_from(twj);
 
-        Ok(SetExpr::Select(Box::new(select_builder.build()?)))
+        Ok(SetExpr::Select(Box::new(
+            select_builder.build(query, &relation_builder)?,
+        )))
     }
 
     /// Reconstructs a SELECT SQL statement from a logical plan by unprojecting column expressions

datafusion/sql/src/unparser/rewrite.rs

@@ -24,7 +24,7 @@ use datafusion_common::{
 };
 use datafusion_expr::{expr::Alias, tree_node::transform_sort_vec};
 use datafusion_expr::{Expr, LogicalPlan, Projection, Sort, SortExpr};
-use sqlparser::ast::Ident;
+use sqlparser::ast::{self, display_separated, Ident};
 
 /// Normalize the schema of a union plan to remove qualifiers from the schema fields and sort expressions.
 ///
@@ -363,3 +363,138 @@ impl TreeNodeRewriter for TableAliasRewriter<'_> {
         }
     }
 }
+
+/// Takes an input list of identifiers and a list of identifiers that are available from relations or joins.
+/// Removes any table identifiers that are not present in the list of available identifiers, retains original column names.
+pub fn remove_dangling_identifiers(idents: &mut Vec<Ident>, available_idents: &[String]) {
+    if idents.len() > 1 {
+        let ident_source = display_separated(
+            &idents
+                .clone()
+                .into_iter()
+                .take(idents.len() - 1)
+                .collect::<Vec<Ident>>(),
+            ".",
+        )
+        .to_string();
+        // If the identifier is not present in the list of all identifiers, it refers to a table that does not exist
+        if !available_idents.contains(&ident_source) {
+            let Some(last) = idents.last() else {
+                unreachable!("CompoundIdentifier must have a last element");
+            };
+            // Reset the identifiers to only the last element, which is the column name
+            *idents = vec![last.clone()];
+        }
+    }
+}
+
+/// Handle removing dangling identifiers from an expression
+/// This function can call itself recursively to handle nested expressions
+/// Like binary ops or functions which contain nested expressions/arguments
+pub fn remove_dangling_expr(
+    expr: ast::Expr,
+    available_idents: &Vec<String>,
+) -> ast::Expr {
+    match expr {
+        ast::Expr::BinaryOp { left, op, right } => {
+            let left = remove_dangling_expr(*left, available_idents);
+            let right = remove_dangling_expr(*right, available_idents);
+            ast::Expr::BinaryOp {
+                left: Box::new(left),
+                op,
+                right: Box::new(right),
+            }
+        }
+        ast::Expr::Nested(expr) => {
+            let expr = remove_dangling_expr(*expr, available_idents);
+            ast::Expr::Nested(Box::new(expr))
+        }
+        ast::Expr::CompoundIdentifier(idents) => {
+            let mut idents = idents.clone();
+            remove_dangling_identifiers(&mut idents, available_idents);
+
+            if idents.is_empty() {
+                unreachable!("Identifier must have at least one element");
+            } else if idents.len() == 1 {
+                ast::Expr::Identifier(idents[0].clone())
+            } else {
+                ast::Expr::CompoundIdentifier(idents)
+            }
+        }
+        ast::Expr::Function(ast::Function {
+            args,
+            name,
+            parameters,
+            filter,
+            null_treatment,
+            over,
+            within_group,
+        }) => {
+            let args = if let ast::FunctionArguments::List(mut args) = args {
+                args.args.iter_mut().for_each(|arg| match arg {
+                    ast::FunctionArg::Named {
+                        arg: ast::FunctionArgExpr::Expr(expr),
+                        ..
+                    }
+                    | ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(expr)) => {
+                        *expr = remove_dangling_expr(expr.clone(), available_idents);
+                    }
+                    _ => {}
+                });
+
+                ast::FunctionArguments::List(args)
+            } else {
+                args
+            };
+
+            ast::Expr::Function(ast::Function {
+                args,
+                name,
+                parameters,
+                filter,
+                null_treatment,
+                over,
+                within_group,
+            })
+        }
+        _ => expr,
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_remove_dangling_identifiers() {
+        let tests = vec![
+            (vec![], vec![Ident::new("column1".to_string())]),
+            (
+                vec!["table1.table2".to_string()],
+                vec![
+                    Ident::new("table1".to_string()),
+                    Ident::new("table2".to_string()),
+                    Ident::new("column1".to_string()),
+                ],
+            ),
+            (
+                vec!["table1".to_string()],
+                vec![Ident::new("column1".to_string())],
+            ),
+        ];
+
+        for test in tests {
+            let test_in = test.0;
+            let test_out = test.1;
+
+            let mut idents = vec![
+                Ident::new("table1".to_string()),
+                Ident::new("table2".to_string()),
+                Ident::new("column1".to_string()),
+            ];
+
+            remove_dangling_identifiers(&mut idents, &test_in);
+            assert_eq!(idents, test_out);
+        }
+    }
+}