fix(unparser): keep a qualifier that an UNNEST or FLATTEN alias, a schema-qualified table or a struct field path binds (refs spiceai/spiceai#13442)

datafusion/sql/src/expr/mod.rs

@@ -1165,12 +1165,18 @@ impl<S: ContextProvider> SqlToRel<'_, S> {
         schema: &DFSchema,
         planner_context: &mut PlannerContext,
     ) -> Result<(Expr, Vec<AccessExpr>)> {
-        let SQLExpr::Identifier(root_ident) = root else {
-            let root = self.sql_expr_to_logical_expr(root, schema, planner_context)?;
-            return Ok((root, access_chain));
+        // An identifier chain arrives either as a plain or a compound identifier
+        // root followed by `Dot` identifiers; either way the whole chain is one
+        // qualified reference, resolved as the parser's `a.b.c` would be.
+        let mut compound_idents = match root {
+            SQLExpr::Identifier(root_ident) => vec![root_ident],
+            SQLExpr::CompoundIdentifier(idents) => idents,
+            root => {
+                let root =
+                    self.sql_expr_to_logical_expr(root, schema, planner_context)?;
+                return Ok((root, access_chain));
+            }
         };
-
-        let mut compound_idents = vec![root_ident];
         let first_non_ident = access_chain
             .iter()
             .position(|access| !matches!(access, AccessExpr::Dot(SQLExpr::Identifier(_))))

datafusion/sql/src/unparser/ast.rs

@@ -952,6 +952,19 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Derived(ref value)) => {
                 value.alias.as_ref().map(|a| a.name.to_string())
             }
+            Some(TableFactorBuilder::Unnest(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            Some(TableFactorBuilder::Flatten(ref value)) => {
+                value.alias.as_ref().map(|a| a.name.to_string())
+            }
+            _ => None,
+        }
+    }
+    /// The name of the bare table this relation reads, if it is one.
+    pub(super) fn table_name(&self) -> Option<&ast::ObjectName> {
+        match self.relation {
+            Some(TableFactorBuilder::Table(ref value)) => value.name.as_ref(),
             _ => None,
         }
     }

datafusion/sql/src/unparser/expr.rs

@@ -908,40 +908,30 @@ impl Unparser<'_> {
             fields.push(field);
         }
 
-        match &args[0] {
-            Expr::Column(col) => {
-                let mut id = match self.col_to_sql(col)? {
-                    ast::Expr::Identifier(ident) => vec![ident],
-                    ast::Expr::CompoundIdentifier(idents) => idents,
-                    other => {
-                        return internal_err!(
-                            "expected col_to_sql to return an Identifier or CompoundIdentifier, but received: {:?}",
-                            other
-                        );
-                    }
-                };
-                id.extend(fields);
-                Ok(ast::Expr::CompoundIdentifier(id))
-            }
+        let root = match &args[0] {
+            Expr::Column(col) => self.col_to_sql(col)?,
             Expr::ScalarFunction(struct_expr) => {
-                let root = self
-                    .scalar_function_to_sql(struct_expr.func.name(), &struct_expr.args)?;
-                let access_chain = fields
-                    .into_iter()
-                    .map(|field| ast::AccessExpr::Dot(ast::Expr::Identifier(field)))
-                    .collect();
-                Ok(ast::Expr::CompoundFieldAccess {
-                    root: Box::new(root),
-                    access_chain,
-                })
+                self.scalar_function_to_sql(struct_expr.func.name(), &struct_expr.args)?
             }
             _ => {
-                internal_err!(
+                return internal_err!(
                     "get_field expects first argument to be column or scalar function, but received: {:?}",
                     &args[0]
-                )
+                );
             }
-        }
+        };
+        // The column keeps its own node rather than being flattened into one compound
+        // identifier with the field path, so a pass over the finished statement can
+        // tell the column's relation qualifier from the fields after it. Both print
+        // the same.
+        let access_chain = fields
+            .into_iter()
+            .map(|field| ast::AccessExpr::Dot(ast::Expr::Identifier(field)))
+            .collect();
+        Ok(ast::Expr::CompoundFieldAccess {
+            root: Box::new(root),
+            access_chain,
+        })
     }
 
     fn map_to_sql(&self, args: &[Expr]) -> Result<ast::Expr> {

datafusion/sql/src/unparser/plan.rs

@@ -27,7 +27,7 @@ use super::{
     },
     rewrite::{
         TableAliasRewriter, inject_column_aliases_into_subquery, normalize_union_schema,
-        remove_dangling_identifiers, requalify_column_onto_derived_table,
+        remove_dangling_qualifier, requalify_column_onto_derived_table,
         rewrite_plan_for_sort_on_non_projected_fields,
         subquery_alias_inner_query_and_columns,
     },
@@ -290,22 +290,29 @@ fn relations_capturable_by(plan: &LogicalPlan, hoisted: &str) -> Result<HashSet<
 /// joined relations.
 ///
 /// An unaliased nested join (`a JOIN (b JOIN c ON …) ON …`) does not hide the
-/// relations inside it, so their names stay in scope.
+/// relations inside it, so their names stay in scope. An `UNNEST(…) AS u` or a
+/// `LATERAL FLATTEN(…) AS u` is in scope by its alias like a derived table.
 #[cfg_attr(feature = "recursive_protection", recursive::recursive)]
 fn collect_relation_idents(relation: &ast::TableFactor, idents: &mut Vec<String>) {
     match relation {
         ast::TableFactor::Table { alias, name, .. } => {
             if let Some(alias) = alias {
                 idents.push(alias.name.to_string());
             } else {
-                idents.push(name.to_string());
+                collect_table_name_idents(name, idents);
             }
         }
         ast::TableFactor::Derived {
             alias: Some(alias), ..
         }
         | ast::TableFactor::NestedJoin {
             alias: Some(alias), ..
+        }
+        | ast::TableFactor::UNNEST {
+            alias: Some(alias), ..
+        }
+        | ast::TableFactor::Function {
+            alias: Some(alias), ..
         } => {
             idents.push(alias.name.to_string());
         }
@@ -322,6 +329,18 @@ fn collect_relation_idents(relation: &ast::TableFactor, idents: &mut Vec<String>
     }
 }
 
+/// Names a `SELECT` can qualify a column with through the bare table `name`: its
+/// full name, and, for a multi-part name, its last part, which is how a column
+/// over it is qualified unless the dialect writes every part.
+fn collect_table_name_idents(name: &ast::ObjectName, idents: &mut Vec<String>) {
+    idents.push(name.to_string());
+    if name.0.len() > 1
+        && let Some(last) = name.0.last()
+    {
+        idents.push(last.to_string());
+    }
+}
+
 /// The refusal for a predicate below a FULL JOIN that no scan's own derived
 /// table applies: `WHERE`, `HAVING` and `QUALIFY` all run after the join and
 /// would discard the rows it preserves.
@@ -929,8 +948,8 @@ impl Unparser<'_> {
         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() {
-            all_idents.push(source_name);
+        } else if let Some(source_name) = relation_builder.table_name() {
+            collect_table_name_idents(source_name, &mut all_idents);
         }
 
         let mut twj = select_builder.pop_from().unwrap();
@@ -944,10 +963,8 @@ impl Unparser<'_> {
         // Ensure that the projection contains references to sources that actually exist
         let mut projection = select_builder.get_projection();
         projection.iter_mut().for_each(|select_item| {
-            if let ast::SelectItem::UnnamedExpr(ast::Expr::CompoundIdentifier(idents)) =
-                select_item
-            {
-                remove_dangling_identifiers(idents, &all_idents);
+            if let ast::SelectItem::UnnamedExpr(expr) = select_item {
+                remove_dangling_qualifier(expr, &all_idents);
             }
         });
 
@@ -956,9 +973,7 @@ impl Unparser<'_> {
             && let Some(OrderByKind::Expressions(mut order_by)) = query.get_order_by()
         {
             order_by.iter_mut().for_each(|sort_item| {
-                if let ast::Expr::CompoundIdentifier(idents) = &mut sort_item.expr {
-                    remove_dangling_identifiers(idents, &all_idents);
-                }
+                remove_dangling_qualifier(&mut sort_item.expr, &all_idents);
             });
 
             query.order_by(OrderByKind::Expressions(order_by));
@@ -967,9 +982,7 @@ impl Unparser<'_> {
         // Order by could be a sort in the select builder
         let mut sort = select_builder.get_sort_by();
         sort.iter_mut().for_each(|sort_item| {
-            if let ast::Expr::CompoundIdentifier(idents) = &mut sort_item.expr {
-                remove_dangling_identifiers(idents, &all_idents);
-            }
+            remove_dangling_qualifier(&mut sort_item.expr, &all_idents);
         });
 
         select_builder.projection(projection);

datafusion/sql/src/unparser/rewrite.rs

@@ -24,7 +24,7 @@ use datafusion_common::{
 };
 use datafusion_expr::expr::{Alias, UNNEST_COLUMN_PREFIX};
 use datafusion_expr::{Distinct, Expr, LogicalPlan, Projection, Sort, SortExpr};
-use sqlparser::ast::Ident;
+use sqlparser::ast::{self, Ident};
 
 /// Normalize the schema of a union plan to remove qualifiers from the schema fields and sort expressions.
 ///
@@ -683,6 +683,30 @@ pub fn remove_dangling_identifiers(idents: &mut Vec<Ident>, available_idents: &[
     }
 }
 
+/// [`remove_dangling_identifiers`] over a column reference, or over the column a
+/// struct field path (`t.s.f`) is rooted at, keeping the path.
+///
+/// Only a path of `.field` steps is judged. That is the shape `get_field` is emitted
+/// in, and its root is always a column of this query's own relations: a correlated
+/// column is refused there. A subscript (`t.a[1]`) is left alone, as its root can be
+/// a correlated reference, which binds in the enclosing query whatever this `FROM`
+/// reads.
+pub fn remove_dangling_qualifier(expr: &mut ast::Expr, available_idents: &[String]) {
+    match expr {
+        ast::Expr::CompoundIdentifier(idents) => {
+            remove_dangling_identifiers(idents, available_idents);
+        }
+        ast::Expr::CompoundFieldAccess { root, access_chain }
+            if access_chain
+                .iter()
+                .all(|access| matches!(access, ast::AccessExpr::Dot(_))) =>
+        {
+            remove_dangling_qualifier(root, available_idents);
+        }
+        _ => {}
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -722,6 +746,41 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_remove_dangling_qualifier() {
+        let column = |parts: &[&str]| {
+            ast::Expr::CompoundIdentifier(
+                parts.iter().map(|part| Ident::new(*part)).collect(),
+            )
+        };
+        let field_path = |root: ast::Expr| ast::Expr::CompoundFieldAccess {
+            root: Box::new(root),
+            access_chain: vec![ast::AccessExpr::Dot(ast::Expr::Identifier(Ident::new(
+                "product",
+            )))],
+        };
+        let subscript = |root: ast::Expr| ast::Expr::CompoundFieldAccess {
+            root: Box::new(root),
+            access_chain: vec![ast::AccessExpr::Subscript(ast::Subscript::Index {
+                index: ast::Expr::value(ast::Value::Number("1".to_string(), false)),
+            })],
+        };
+
+        // A field path is judged by the column it is rooted at, and keeps its path.
+        let mut expr = field_path(column(&["t", "metadata"]));
+        remove_dangling_qualifier(&mut expr, &[]);
+        assert_eq!(expr, field_path(column(&["metadata"])));
+
+        let mut expr = field_path(column(&["t", "metadata"]));
+        remove_dangling_qualifier(&mut expr, &["t".to_string()]);
+        assert_eq!(expr, field_path(column(&["t", "metadata"])));
+
+        // A subscript's root may be a correlated reference, so it is left alone.
+        let mut expr = subscript(column(&["o", "array_col"]));
+        remove_dangling_qualifier(&mut expr, &[]);
+        assert_eq!(expr, subscript(column(&["o", "array_col"])));
+    }
+
     // this is a regression test: when the outer projection has fewer expressions than
     // the inner projection, `subquery_alias_inner_query_and_columns` must not panic
     // with an index oob error