Unnest fixes for multi-arguments

Author: yonatan-sevenai

datafusion/sql/src/unparser/plan.rs

@@ -421,20 +421,25 @@ impl Unparser<'_> {
                 }
 
                 // Projection can be top-level plan for unnest relation.
-                // The projection generated by the `RecursiveUnnestRewriter`
-                // will have exactly one expression referencing the unnest
-                // placeholder column. The placeholder may be at the top level
-                // ("bare": `__unnest_placeholder(...) AS item`) or wrapped
-                // inside a function call ("wrapped":
-                // `json_as_text(__unnest_placeholder(...), 'type') AS t`).
-                let (unnest_input_type, placeholder_is_bare) = if p.expr.len() == 1 {
-                    match Self::find_unnest_placeholder(&p.expr[0]) {
-                        Some((t, is_bare)) => (Some(t), is_bare),
-                        None => (None, false),
-                    }
-                } else {
-                    (None, false)
-                };
+                // At least one expression will reference the unnest
+                // placeholder column. The placeholder may be:
+                //  - "bare": the sole expression IS the placeholder (+ aliases)
+                //  - "wrapped": inside a function call, or one of several exprs
+                // When bare, we emit `_unnest."VALUE" [AS alias]`.
+                // Otherwise `reconstruct_select_statement` renders all
+                // expressions and rewrites the placeholder via
+                // `unproject_unnest_expr_as_flatten_value`.
+                let (unnest_input_type, placeholder_is_bare) = p
+                    .expr
+                    .iter()
+                    .find_map(Self::find_unnest_placeholder)
+                    .map(|(t, is_bare)| {
+                        // Only bare when there is a single expression that
+                        // IS the placeholder — multi-expression projections
+                        // always need reconstruct_select_statement.
+                        (Some(t), is_bare && p.expr.len() == 1)
+                    })
+                    .unwrap_or((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 placeholder_is_bare && unnest_input_type.is_some() {
@@ -524,8 +529,11 @@ impl Unparser<'_> {
                     return Ok(());
                 }
                 // Standard UNNEST table factor path (BigQuery, etc.).
+                // Only fires for single-expression projections — multi-expr
+                // falls through to reconstruct_select_statement.
                 if self.dialect.unnest_as_table_factor()
                     && unnest_input_type.is_some()
+                    && p.expr.len() == 1
                     && user_alias.is_none() // Skip if user alias present — fall through to reconstruct_select_statement which preserves aliases
                     && let Some((unnest, unnest_plan)) =
                         self.peel_to_unnest_with_modifiers(p.input.as_ref(), query)?

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -3360,6 +3360,87 @@ fn snowflake_flatten_composed_expression_with_limit() -> Result<(), DataFusionEr
     Ok(())
 }
 
+#[test]
+fn snowflake_flatten_multi_expression_projection() -> Result<(), DataFusionError> {
+    // Build: Projection([CAST(placeholder AS Int64) AS a, CAST(placeholder AS Utf8) AS b])
+    //          → Unnest → Projection → TableScan
+    // The outer Projection has TWO expressions — both reference the placeholder.
+    // The FLATTEN code path must fire even when p.expr.len() > 1.
+    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("a"),
+            cast(col("__unnest_placeholder(items)"), DataType::Utf8).alias("b"),
+        ])?
+        .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}"
+    );
+    // Both expressions should be present
+    assert!(
+        actual.contains("CAST"),
+        "Expected CAST in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains(r#"AS "a""#),
+        "Expected alias 'a' in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains(r#"AS "b""#),
+        "Expected alias 'b' in SQL, got: {actual}"
+    );
+    Ok(())
+}
+
+#[test]
+fn snowflake_flatten_multi_expression_with_limit() -> Result<(), DataFusionError> {
+    // Multi-expression + Limit between Projection and Unnest
+    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(10))?
+        .project(vec![
+            cast(col("__unnest_placeholder(items)"), DataType::Int64).alias("a"),
+            cast(col("__unnest_placeholder(items)"), DataType::Utf8).alias("b"),
+        ])?
+        .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("LIMIT 10"),
+        "Expected LIMIT 10 in SQL, got: {actual}"
+    );
+    Ok(())
+}
+
 #[test]
 fn snowflake_unnest_through_subquery_alias() -> Result<(), DataFusionError> {
     // Build: Projection → Unnest → SubqueryAlias → Projection → TableScan