More snowflake dialect fixes for unnes

Author: yonatan-sevenai

datafusion/sql/src/unparser/plan.rs

@@ -45,7 +45,7 @@ use crate::utils::UNNEST_PLACEHOLDER;
 use datafusion_common::{
     Column, DataFusionError, Result, ScalarValue, TableReference, assert_or_internal_err,
     internal_err, not_impl_err,
-    tree_node::{TransformedResult, TreeNode},
+    tree_node::{TransformedResult, TreeNode, TreeNodeRecursion},
 };
 use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX};
 use datafusion_expr::{
@@ -420,17 +420,33 @@ impl Unparser<'_> {
                         .select_to_sql_recursively(&new_plan, query, select, relation);
                 }
 
-                // Projection can be top-level plan for unnest relation
+                // Projection can be top-level plan for unnest relation.
                 // The projection generated by the `RecursiveUnnestRewriter` from a UNNEST relation will have
                 // only one expression, which is the placeholder column generated by the rewriter.
-                let unnest_input_type = if p.expr.len() == 1 {
-                    Self::check_unnest_placeholder_with_outer_ref(&p.expr[0])
+                //
+                // Two cases:
+                //  - "bare": the expression IS the placeholder (+ aliases):
+                //      `__unnest_placeholder(...) AS item`
+                //  - "wrapped": the placeholder is inside a function call:
+                //      `json_as_text(__unnest_placeholder(...), 'type') AS target_type`
+                let (unnest_input_type, placeholder_is_bare) = if p.expr.len() == 1 {
+                    if let Some(t) =
+                        Self::check_unnest_placeholder_with_outer_ref(&p.expr[0])
+                    {
+                        (Some(t), true)
+                    } else if let Some(t) =
+                        Self::find_unnest_placeholder_in_expr(&p.expr[0])
+                    {
+                        (Some(t), false)
+                    } else {
+                        (None, false)
+                    }
                 } else {
-                    None
+                    (None, false)
                 };
                 // Extract the outermost user alias (e.g. "c1" from `UNNEST(...) AS c1`).
                 // Internal aliases like "UNNEST(...)" are not user aliases.
-                let user_alias = if unnest_input_type.is_some() {
+                let user_alias = if placeholder_is_bare && unnest_input_type.is_some() {
                     Self::extract_unnest_user_alias(&p.expr[0])
                 } else {
                     None
@@ -472,11 +488,15 @@ impl Unparser<'_> {
                         relation.flatten(flatten_relation);
 
                         if !select.already_projected() {
-                            let value_expr = self.build_flatten_value_select_item(
-                                &alias_name,
-                                user_alias.as_deref(),
-                            );
-                            select.projection(vec![value_expr]);
+                            if placeholder_is_bare {
+                                let value_expr = self.build_flatten_value_select_item(
+                                    &alias_name,
+                                    user_alias.as_deref(),
+                                );
+                                select.projection(vec![value_expr]);
+                            } else {
+                                self.reconstruct_select_statement(plan, p, select)?;
+                            }
                         }
 
                         return self.select_to_sql_recursively(
@@ -509,11 +529,15 @@ impl Unparser<'_> {
                     flatten.outer(unnest.options.preserve_nulls);
 
                     if !select.already_projected() {
-                        let value_expr = self.build_flatten_value_select_item(
-                            &alias_name,
-                            user_alias.as_deref(),
-                        );
-                        select.projection(vec![value_expr]);
+                        if placeholder_is_bare {
+                            let value_expr = self.build_flatten_value_select_item(
+                                &alias_name,
+                                user_alias.as_deref(),
+                            );
+                            select.projection(vec![value_expr]);
+                        } else {
+                            self.reconstruct_select_statement(plan, p, select)?;
+                        }
                     }
 
                     // Recurse into the Unnest → inner source to set the primary
@@ -1356,6 +1380,39 @@ impl Unparser<'_> {
         None
     }
 
+    /// Recursively search the expression tree for an unnest placeholder column.
+    ///
+    /// Unlike [`Self::check_unnest_placeholder_with_outer_ref`] which only
+    /// matches when the placeholder IS the expression (modulo aliases), this
+    /// function finds placeholders buried inside function calls or other
+    /// transformations. For example:
+    ///
+    /// ```text
+    /// Alias("target_type",
+    ///   json_as_text(
+    ///     Column("__unnest_placeholder(...)"),  ← found here
+    ///     Literal("type")))
+    /// ```
+    fn find_unnest_placeholder_in_expr(expr: &Expr) -> Option<UnnestInputType> {
+        let mut result = None;
+        let _ = expr.apply(|e| {
+            if let Expr::Column(Column { name, .. }) = e
+                && let Some(prefix) = name.strip_prefix(UNNEST_PLACEHOLDER)
+            {
+                result = if prefix
+                    .starts_with(&format!("({OUTER_REFERENCE_COLUMN_PREFIX}("))
+                {
+                    Some(UnnestInputType::OuterReference)
+                } else {
+                    Some(UnnestInputType::Scalar)
+                };
+                return Ok(TreeNodeRecursion::Stop);
+            }
+            Ok(TreeNodeRecursion::Continue)
+        });
+        result
+    }
+
     /// Extract the outermost user-provided alias from an unnest expression.
     /// Returns `None` if the outermost alias is DataFusion's internal display
     /// name (e.g. `UNNEST(make_array(...))`), or if there is no alias at all.

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -3283,6 +3283,83 @@ fn snowflake_flatten_limit_between_projection_and_unnest_with_subquery_alias()
     Ok(())
 }
 
+#[test]
+fn snowflake_flatten_composed_expression_wrapping_unnest() -> Result<(), DataFusionError>
+{
+    // Build: Projection(CAST(placeholder AS Int64) AS item_id) → Unnest → Projection → TableScan
+    // The outer Projection wraps the unnest output in a function call.
+    // The FLATTEN code path must detect the placeholder inside the function
+    // and still emit LATERAL FLATTEN.
+    let schema = Schema::new(vec![Field::new(
+        "items",
+        DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))),
+        true,
+    )]);
+
+    let plan = table_scan(Some("source"), &schema, None)?
+        .project(vec![col("items").alias("__unnest_placeholder(items)")])?
+        .unnest_column("__unnest_placeholder(items)")?
+        .project(vec![
+            cast(col("__unnest_placeholder(items)"), DataType::Int64).alias("item_id"),
+        ])?
+        .build()?;
+
+    let snowflake = SnowflakeDialect::new();
+    let unparser = Unparser::new(&snowflake);
+    let result = unparser.plan_to_sql(&plan)?;
+    let actual = result.to_string();
+
+    // Must contain LATERAL FLATTEN despite the placeholder being inside CAST
+    assert!(
+        actual.contains("LATERAL FLATTEN"),
+        "Expected LATERAL FLATTEN in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains("CAST"),
+        "Expected CAST in SQL, got: {actual}"
+    );
+    Ok(())
+}
+
+#[test]
+fn snowflake_flatten_composed_expression_with_limit() -> Result<(), DataFusionError> {
+    // Combines both bugs: composed expression + Limit between Projection and Unnest
+    // Build: Projection(CAST(placeholder AS Int64) AS item_id) → Limit → Unnest → Projection → TableScan
+    let schema = Schema::new(vec![Field::new(
+        "items",
+        DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))),
+        true,
+    )]);
+
+    let plan = table_scan(Some("source"), &schema, None)?
+        .project(vec![col("items").alias("__unnest_placeholder(items)")])?
+        .unnest_column("__unnest_placeholder(items)")?
+        .limit(0, Some(5))?
+        .project(vec![
+            cast(col("__unnest_placeholder(items)"), DataType::Int64).alias("item_id"),
+        ])?
+        .build()?;
+
+    let snowflake = SnowflakeDialect::new();
+    let unparser = Unparser::new(&snowflake);
+    let result = unparser.plan_to_sql(&plan)?;
+    let actual = result.to_string();
+
+    assert!(
+        actual.contains("LATERAL FLATTEN"),
+        "Expected LATERAL FLATTEN in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains("CAST"),
+        "Expected CAST in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains("LIMIT 5"),
+        "Expected LIMIT 5 in SQL, got: {actual}"
+    );
+    Ok(())
+}
+
 #[test]
 fn snowflake_unnest_through_subquery_alias() -> Result<(), DataFusionError> {
     // Build: Projection → Unnest → SubqueryAlias → Projection → TableScan