Snowflake Dialect

yonatan-sevenai · yonatan-sevenai · commit e70196235e3e · 2026-04-11T13:43:13.000-04:00
diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs
@@ -423,7 +423,10 @@ impl Unparser<'_> {
                 // 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])
+                    Self::check_unnest_placeholder_with_outer_ref(
+                        &p.expr[0],
+                        self.dialect.unnest_as_lateral_flatten(),
+                    )
                 } else {
                     None
                 };
@@ -434,26 +437,98 @@ impl Unparser<'_> {
                         self.try_unnest_to_lateral_flatten_sql(unnest)?
                 {
                     let alias_name = flatten_relation.alias_name().to_string();
-                    relation.flatten(flatten_relation);
 
-                    // Only set the projection if it hasn't already been set by
-                    // an outer Projection (e.g. SELECT * FROM t, UNNEST(...)).
-                    // When the outer Projection already called
-                    // reconstruct_select_statement, the SELECT list is correct
-                    // and we must not replace it with just `_unnest."VALUE"`.
+                    // Check if the Unnest source is a real query (subselect)
+                    // vs an inline array (EmptyRelation).
+                    let inner_projection = match unnest.input.as_ref() {
+                        LogicalPlan::Projection(proj) => proj,
+                        other => {
+                            return internal_err!(
+                                "Unnest input is not a Projection: {other:?}"
+                            );
+                        }
+                    };
+
+                    if matches!(
+                        inner_projection.input.as_ref(),
+                        LogicalPlan::EmptyRelation(_)
+                    ) {
+                        // Inline array case (e.g. UNNEST([1,2,3])):
+                        // FLATTEN is the sole FROM source.
+                        relation.flatten(flatten_relation);
+
+                        if !select.already_projected() {
+                            let value_expr =
+                                self.build_flatten_value_select_item(&alias_name);
+                            select.projection(vec![value_expr]);
+                        }
+
+                        return self.select_to_sql_recursively(
+                            p.input.as_ref(),
+                            query,
+                            select,
+                            relation,
+                        );
+                    }
+
+                    // Non-empty source (table, subquery, etc.):
+                    // The FLATTEN references columns/expressions from the source.
+                    // Convert the inner Projection's expression to SQL for the
+                    // FLATTEN INPUT, then store the FLATTEN to be added as a
+                    // CROSS JOIN after the source relation is processed.
+                    let first_expr = inner_projection.expr.first().ok_or_else(|| {
+                        DataFusionError::Internal(
+                            "Unnest inner projection has no expressions".to_string(),
+                        )
+                    })?;
+                    // Strip the __unnest_placeholder alias to get the raw expression
+                    let raw_expr = match first_expr {
+                        Expr::Alias(Alias { expr, .. }) => expr.as_ref(),
+                        other => other,
+                    };
+                    let input_sql = self.expr_to_sql(raw_expr)?;
+
+                    let mut flatten = FlattenRelationBuilder::default();
+                    flatten.input_expr(input_sql);
+                    flatten.outer(unnest.options.preserve_nulls);
+
                     if !select.already_projected() {
                         let value_expr =
                             self.build_flatten_value_select_item(&alias_name);
                         select.projection(vec![value_expr]);
-                        select.already_projected();
                     }
 
-                    return self.select_to_sql_recursively(
+                    // Recurse into the Unnest → inner source to set the primary
+                    // relation (table scan, subquery, etc.), then add FLATTEN
+                    // as a CROSS JOIN.
+                    self.select_to_sql_recursively(
                         p.input.as_ref(),
                         query,
                         select,
                         relation,
-                    );
+                    )?;
+
+                    let flatten_factor = flatten.build().map_err(|e| {
+                        DataFusionError::Internal(format!("Failed to build FLATTEN: {e}"))
+                    })?;
+                    let cross_join = ast::Join {
+                        relation: flatten_factor,
+                        global: false,
+                        join_operator: ast::JoinOperator::CrossJoin(
+                            ast::JoinConstraint::None,
+                        ),
+                    };
+                    if let Some(mut from) = select.pop_from() {
+                        from.push_join(cross_join);
+                        select.push_from(from);
+                    } else {
+                        // No existing FROM — create one with just the FLATTEN
+                        let mut twj = TableWithJoinsBuilder::default();
+                        twj.push_join(cross_join);
+                        select.push_from(twj);
+                    }
+
+                    return Ok(());
                 }
                 if self.dialect.unnest_as_table_factor()
                     && unnest_input_type.is_some()
@@ -1153,15 +1228,35 @@ impl Unparser<'_> {
     /// - If the column is not a placeholder column, return [None].
     ///
     /// `outer_ref` is the display result of [Expr::OuterReferenceColumn]
-    fn check_unnest_placeholder_with_outer_ref(expr: &Expr) -> Option<UnnestInputType> {
-        if let Expr::Alias(Alias { expr, .. }) = expr
-            && let Expr::Column(Column { name, .. }) = expr.as_ref()
-            && let Some(prefix) = name.strip_prefix(UNNEST_PLACEHOLDER)
-        {
-            if prefix.starts_with(&format!("({OUTER_REFERENCE_COLUMN_PREFIX}(")) {
-                return Some(UnnestInputType::OuterReference);
+    /// When `deep_peel` is true, peels through multiple Alias layers to find
+    /// the inner Column. This is needed for Snowflake FLATTEN where user aliases
+    /// (e.g. `AS items`) add extra Alias wrappers around the placeholder column.
+    fn check_unnest_placeholder_with_outer_ref(
+        expr: &Expr,
+        deep_peel: bool,
+    ) -> Option<UnnestInputType> {
+        let inner = match expr {
+            Expr::Alias(Alias { expr, .. }) => {
+                if deep_peel {
+                    // Peel through all Alias layers
+                    let mut e = expr.as_ref();
+                    while let Expr::Alias(Alias { expr: next, .. }) = e {
+                        e = next.as_ref();
+                    }
+                    e
+                } else {
+                    expr.as_ref()
+                }
+            }
+            _ => return None,
+        };
+        if let Expr::Column(Column { name, .. }) = inner {
+            if let Some(prefix) = name.strip_prefix(UNNEST_PLACEHOLDER) {
+                if prefix.starts_with(&format!("({OUTER_REFERENCE_COLUMN_PREFIX}(")) {
+                    return Some(UnnestInputType::OuterReference);
+                }
+                return Some(UnnestInputType::Scalar);
             }
-            return Some(UnnestInputType::Scalar);
         }
         None
     }
diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs
@@ -26,8 +26,9 @@ use datafusion_expr::test::function_stub::{
     count_udaf, max, max_udaf, min_udaf, sum, sum_udaf,
 };
 use datafusion_expr::{
-    EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder, Union,
-    UserDefinedLogicalNode, UserDefinedLogicalNodeCore, WindowFrame,
+    ColumnarValue, EmptyRelation, Expr, Extension, LogicalPlan, LogicalPlanBuilder,
+    ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Union,
+    UserDefinedLogicalNode, UserDefinedLogicalNodeCore, Volatility, WindowFrame,
     WindowFunctionDefinition, cast, col, lit, table_scan, wildcard,
 };
 use datafusion_functions::unicode;
@@ -3032,7 +3033,7 @@ fn snowflake_unnest_to_lateral_flatten_cross_join_inline() -> Result<(), DataFus
         sql: "SELECT * FROM UNNEST([1,2,3]) u(c1) JOIN j1 ON u.c1 = j1.j1_id",
         parser_dialect: GenericDialect {},
         unparser_dialect: snowflake,
-        expected: @r#"SELECT "u"."c1", "j1"."j1_id", "j1"."j1_string" FROM (SELECT "_unnest"."VALUE" AS "c1" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS _unnest) AS "u" INNER JOIN "j1" ON ("u"."c1" = "j1"."j1_id")"#,
+        expected: @r#"SELECT "u"."c1", "j1"."j1_id", "j1"."j1_string" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "u" INNER JOIN "j1" ON ("u"."c1" = "j1"."j1_id")"#,
     );
     Ok(())
 }
@@ -3072,7 +3073,7 @@ fn snowflake_flatten_select_unnest_with_alias() -> Result<(), DataFusionError> {
         sql: "SELECT UNNEST([1,2,3]) as c1",
         parser_dialect: GenericDialect {},
         unparser_dialect: snowflake,
-        expected: @r#"SELECT "_unnest"."VALUE" AS "c1" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS _unnest"#,
+        expected: @r#"SELECT _unnest."VALUE" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS _unnest"#,
     );
     Ok(())
 }
@@ -3096,7 +3097,89 @@ fn snowflake_flatten_from_unnest_with_table_alias() -> Result<(), DataFusionErro
         sql: "SELECT * FROM UNNEST([1,2,3]) AS t1 (c1)",
         parser_dialect: GenericDialect {},
         unparser_dialect: snowflake,
-        expected: @r#"SELECT "t1"."c1" FROM (SELECT "_unnest"."VALUE" AS "c1" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS _unnest) AS "t1""#,
+        expected: @r#"SELECT "t1"."c1" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "t1""#,
+    );
+    Ok(())
+}
+
+#[test]
+fn snowflake_flatten_unnest_from_subselect() -> Result<(), DataFusionError> {
+    // UNNEST operating on an array column produced by a subselect.
+    // Uses unnest_table which has array_col (List<Int64>).
+    // The filter uses array_col IS NOT NULL — a simple predicate
+    // that doesn't involve struct types (which Snowflake FLATTEN can't handle).
+    let snowflake = SnowflakeDialect::new();
+    roundtrip_statement_with_dialect_helper!(
+        sql: "SELECT UNNEST(array_col) FROM (SELECT array_col FROM unnest_table WHERE array_col IS NOT NULL LIMIT 3)",
+        parser_dialect: GenericDialect {},
+        unparser_dialect: snowflake,
+        expected: @r#"SELECT _unnest."VALUE" FROM (SELECT "unnest_table"."array_col" FROM "unnest_table" WHERE "unnest_table"."array_col" IS NOT NULL LIMIT 3) CROSS JOIN LATERAL FLATTEN(INPUT => "unnest_table"."array_col") AS _unnest"#,
     );
     Ok(())
 }
+
+/// Dummy scalar UDF for testing — takes a string and returns List<Int64>.
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct JsonGetArrayUdf {
+    signature: Signature,
+}
+
+impl JsonGetArrayUdf {
+    fn new() -> Self {
+        Self {
+            signature: Signature::exact(vec![DataType::Utf8], Volatility::Immutable),
+        }
+    }
+}
+
+impl ScalarUDFImpl for JsonGetArrayUdf {
+    fn name(&self) -> &str {
+        "json_get_array"
+    }
+    fn signature(&self) -> &Signature {
+        &self.signature
+    }
+    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
+        Ok(DataType::List(Arc::new(Field::new_list_field(
+            DataType::Int64,
+            true,
+        ))))
+    }
+    fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result<ColumnarValue> {
+        unimplemented!("test stub")
+    }
+}
+
+#[test]
+fn snowflake_flatten_unnest_udf_result() -> Result<(), DataFusionError> {
+    // UNNEST on a UDF result: json_get_array(col) returns List<Int64>,
+    // then UNNEST flattens it. This simulates a common Snowflake pattern
+    // where a UDF parses JSON into an array, then FLATTEN expands it.
+    let sql = "SELECT UNNEST(json_get_array(j1_string)) AS items FROM j1 LIMIT 5";
+
+    let statement = Parser::new(&GenericDialect {})
+        .try_with_sql(sql)?
+        .parse_statement()?;
+
+    let state = MockSessionState::default()
+        .with_aggregate_function(max_udaf())
+        .with_aggregate_function(min_udaf())
+        .with_scalar_function(Arc::new(ScalarUDF::new_from_impl(JsonGetArrayUdf::new())))
+        .with_expr_planner(Arc::new(CoreFunctionPlanner::default()))
+        .with_expr_planner(Arc::new(NestedFunctionPlanner))
+        .with_expr_planner(Arc::new(FieldAccessPlanner));
+
+    let context = MockContextProvider { state };
+    let sql_to_rel = SqlToRel::new(&context);
+    let plan = sql_to_rel
+        .sql_statement_to_plan(statement)
+        .unwrap_or_else(|e| panic!("Failed to parse sql: {sql}\n{e}"));
+
+    let snowflake = SnowflakeDialect::new();
+    let unparser = Unparser::new(&snowflake);
+    let result = unparser.plan_to_sql(&plan)?;
+    let actual = result.to_string();
+
+    insta::assert_snapshot!(actual, @r#"SELECT _unnest."VALUE" FROM "j1" CROSS JOIN LATERAL FLATTEN(INPUT => json_get_array("j1"."j1_string")) AS _unnest LIMIT 5"#);
+    Ok(())
+}
