More snowflake dialect fixes for unnest

Author: yonatan-sevenai

datafusion/sql/src/unparser/plan.rs

@@ -437,30 +437,31 @@ impl Unparser<'_> {
                 };
                 if self.dialect.unnest_as_lateral_flatten()
                     && unnest_input_type.is_some()
-                    && let LogicalPlan::Unnest(unnest) = p.input.as_ref()
+                    && let Some(unnest) = Self::peel_to_unnest(p.input.as_ref())
                     && let Some(flatten_relation) =
                         self.try_unnest_to_lateral_flatten_sql(unnest)?
                 {
                     let alias_name = flatten_relation.alias_name().to_string();
 
                     // 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,
-                        LogicalPlan::SubqueryAlias(alias) => match alias.input.as_ref() {
-                            LogicalPlan::Projection(proj) => proj,
-                            other => {
-                                return internal_err!(
-                                    "Unnest input (through SubqueryAlias) is not a Projection: {other:?}"
-                                );
-                            }
-                        },
-                        other => {
-                            return internal_err!(
-                                "Unnest input is not a Projection: {other:?}"
-                            );
-                        }
-                    };
+                    let inner_projection =
+                        Self::peel_to_inner_projection(unnest.input.as_ref())
+                            .ok_or_else(|| {
+                                DataFusionError::Internal(format!(
+                                    "Unnest input is not a Projection: {:?}",
+                                    unnest.input
+                                ))
+                            })?;
+
+                    // Apply any intermediate Limit/Sort modifiers and get
+                    // a reference to the Unnest LogicalPlan node for recursion.
+                    // This bypasses the Limit/Sort handlers (which would create
+                    // unwanted derived subqueries when already_projected is set).
+                    let unnest_plan = self.apply_transparent_and_find_unnest_plan(
+                        p.input.as_ref(),
+                        query,
+                    )?;
 
                     if matches!(
                         inner_projection.input.as_ref(),
@@ -479,7 +480,7 @@ impl Unparser<'_> {
                         }
 
                         return self.select_to_sql_recursively(
-                            p.input.as_ref(),
+                            unnest_plan,
                             query,
                             select,
                             relation,
@@ -518,12 +519,7 @@ impl Unparser<'_> {
                     // 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,
-                    )?;
+                    self.select_to_sql_recursively(unnest_plan, query, select, relation)?;
 
                     let flatten_factor = flatten.build().map_err(|e| {
                         DataFusionError::Internal(format!("Failed to build FLATTEN: {e}"))
@@ -550,14 +546,18 @@ impl Unparser<'_> {
                 if self.dialect.unnest_as_table_factor()
                     && unnest_input_type.is_some()
                     && user_alias.is_none() // Skip if user alias present — fall through to reconstruct_select_statement which preserves aliases
-                    && let LogicalPlan::Unnest(unnest) = &p.input.as_ref()
+                    && let Some(unnest) = Self::peel_to_unnest(p.input.as_ref())
                     && let Some(unnest_relation) =
                         self.try_unnest_to_table_factor_sql(unnest)?
                 {
                     relation.unnest(unnest_relation);
-                    return self.select_to_sql_recursively(
+                    let unnest_plan = self.apply_transparent_and_find_unnest_plan(
                         p.input.as_ref(),
                         query,
+                    )?;
+                    return self.select_to_sql_recursively(
+                        unnest_plan,
+                        query,
                         select,
                         relation,
                     );
@@ -1201,13 +1201,9 @@ impl Unparser<'_> {
                 // |     Projection: table.col1, table.col2 AS UNNEST(table.col2)
                 // |       Filter: table.col3 = Int64(3)
                 // |         TableScan: table projection=None
-                if let LogicalPlan::Projection(p) = unnest.input.as_ref() {
-                    // continue with projection input
-                    self.select_to_sql_recursively(&p.input, query, select, relation)
-                } else if let LogicalPlan::SubqueryAlias(alias) = unnest.input.as_ref()
-                    && let LogicalPlan::Projection(p) = alias.input.as_ref()
-                {
-                    // SubqueryAlias wraps the Projection (e.g. passthrough tables)
+                if let Some(p) = Self::peel_to_inner_projection(unnest.input.as_ref()) {
+                    // Skip the inner Projection (synthetic rewriter node)
+                    // and continue with its input.
                     self.select_to_sql_recursively(&p.input, query, select, relation)
                 } else {
                     internal_err!("Unnest input is not a Projection: {unnest:?}")
@@ -1242,6 +1238,96 @@ impl Unparser<'_> {
         }
     }
 
+    /// Walk through "transparent" plan nodes (Limit, Sort) to find an Unnest.
+    ///
+    /// The DataFusion optimizer may insert Limit or Sort between the outer
+    /// Projection and the Unnest node. These nodes modify result quantity or
+    /// ordering but do not change the plan shape for unnest detection. The
+    /// normal recursion in [`Self::select_to_sql_recursively`] handles their
+    /// SQL rendering (LIMIT/OFFSET/ORDER BY); this helper only needs to
+    /// locate the Unnest so the FLATTEN / table-factor code path can fire.
+    fn peel_to_unnest(plan: &LogicalPlan) -> Option<&Unnest> {
+        match plan {
+            LogicalPlan::Unnest(unnest) => Some(unnest),
+            LogicalPlan::Limit(limit) => Self::peel_to_unnest(limit.input.as_ref()),
+            LogicalPlan::Sort(sort) => Self::peel_to_unnest(sort.input.as_ref()),
+            _ => None,
+        }
+    }
+
+    /// Walk through "transparent" plan nodes (SubqueryAlias, Limit, Sort) to
+    /// find the inner Projection that feeds an Unnest node.
+    ///
+    /// The inner Projection is created by the `RecursiveUnnestRewriter` and
+    /// contains the array expression that the Unnest operates on. A
+    /// `SubqueryAlias` (e.g. from a virtual/passthrough table) may wrap the
+    /// Projection; Limit/Sort are also handled for robustness.
+    fn peel_to_inner_projection(plan: &LogicalPlan) -> Option<&Projection> {
+        match plan {
+            LogicalPlan::Projection(p) => Some(p),
+            LogicalPlan::SubqueryAlias(alias) => {
+                Self::peel_to_inner_projection(alias.input.as_ref())
+            }
+            LogicalPlan::Limit(limit) => {
+                Self::peel_to_inner_projection(limit.input.as_ref())
+            }
+            LogicalPlan::Sort(sort) => {
+                Self::peel_to_inner_projection(sort.input.as_ref())
+            }
+            _ => None,
+        }
+    }
+
+    /// Walk through transparent nodes (Limit, Sort) between the outer
+    /// Projection and the Unnest, applying their SQL modifiers (LIMIT,
+    /// OFFSET, ORDER BY) to the query builder. Returns a reference to the
+    /// Unnest `LogicalPlan` node so the caller can recurse into it directly,
+    /// bypassing the intermediate handlers that would otherwise create
+    /// unwanted derived subqueries.
+    fn apply_transparent_and_find_unnest_plan<'a>(
+        &self,
+        plan: &'a LogicalPlan,
+        query: &mut Option<QueryBuilder>,
+    ) -> Result<&'a LogicalPlan> {
+        match plan {
+            LogicalPlan::Unnest(_) => Ok(plan),
+            LogicalPlan::Limit(limit) => {
+                if let Some(fetch) = &limit.fetch
+                    && let Some(q) = query.as_mut()
+                {
+                    q.limit(Some(self.expr_to_sql(fetch)?));
+                }
+                if let Some(skip) = &limit.skip
+                    && let Some(q) = query.as_mut()
+                {
+                    q.offset(Some(ast::Offset {
+                        rows: ast::OffsetRows::None,
+                        value: self.expr_to_sql(skip)?,
+                    }));
+                }
+                self.apply_transparent_and_find_unnest_plan(limit.input.as_ref(), query)
+            }
+            LogicalPlan::Sort(sort) => {
+                let Some(query_ref) = query.as_mut() else {
+                    return internal_err!(
+                        "Sort between Projection and Unnest requires a statement context."
+                    );
+                };
+                if let Some(fetch) = sort.fetch {
+                    query_ref.limit(Some(ast::Expr::value(ast::Value::Number(
+                        fetch.to_string(),
+                        false,
+                    ))));
+                }
+                query_ref.order_by(self.sorts_to_sql(&sort.expr)?);
+                self.apply_transparent_and_find_unnest_plan(sort.input.as_ref(), query)
+            }
+            other => {
+                internal_err!("Unexpected node between Projection and Unnest: {other:?}")
+            }
+        }
+    }
+
     /// Try to find the placeholder column name generated by `RecursiveUnnestRewriter`.
     ///
     /// - If the column is a placeholder column match the pattern `Expr::Alias(Expr::Column("__unnest_placeholder(...)"))`,
@@ -1338,16 +1424,9 @@ impl Unparser<'_> {
         &self,
         unnest: &Unnest,
     ) -> Result<Option<FlattenRelationBuilder>> {
-        let projection = match unnest.input.as_ref() {
-            LogicalPlan::Projection(p) => p,
-            LogicalPlan::SubqueryAlias(alias) => {
-                if let LogicalPlan::Projection(p) = alias.input.as_ref() {
-                    p
-                } else {
-                    return Ok(None);
-                }
-            }
-            _ => return Ok(None),
+        let Some(projection) = Self::peel_to_inner_projection(unnest.input.as_ref())
+        else {
+            return Ok(None);
         };
 
         // For now, handle the simple case of a single expression to flatten.

datafusion/sql/tests/cases/plan_to_sql.rs

@@ -3184,6 +3184,105 @@ fn snowflake_flatten_unnest_udf_result() -> Result<(), DataFusionError> {
     Ok(())
 }
 
+#[test]
+fn snowflake_flatten_limit_between_projection_and_unnest() -> Result<(), DataFusionError>
+{
+    // Build: Projection → Limit → Unnest → Projection → TableScan
+    // The optimizer can insert a Limit between the outer Projection and the
+    // Unnest. The FLATTEN code path must look through transparent nodes
+    // (Limit, Sort) to find the 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(5))? // Limit BETWEEN outer Projection and Unnest
+        .project(vec![col("__unnest_placeholder(items)").alias("item")])?
+        .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 — the Limit must not prevent FLATTEN detection
+    insta::assert_snapshot!(actual, @r#"SELECT _unnest."VALUE" AS "item" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS _unnest LIMIT 5"#);
+    Ok(())
+}
+
+#[test]
+fn snowflake_flatten_sort_between_projection_and_unnest() -> Result<(), DataFusionError> {
+    // Build: Projection → Sort → Unnest → Projection → TableScan
+    // Same as Limit test but with Sort instead.
+    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)")?
+        .sort(vec![col("__unnest_placeholder(items)").sort(true, true)])?
+        .project(vec![col("__unnest_placeholder(items)").alias("item")])?
+        .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 — the Sort must not prevent FLATTEN detection
+    assert!(
+        actual.contains("LATERAL FLATTEN"),
+        "Expected LATERAL FLATTEN in SQL, got: {actual}"
+    );
+    assert!(
+        actual.contains("ORDER BY"),
+        "Expected ORDER BY in SQL, got: {actual}"
+    );
+    Ok(())
+}
+
+#[test]
+fn snowflake_flatten_limit_between_projection_and_unnest_with_subquery_alias()
+-> Result<(), DataFusionError> {
+    // Build: Projection → Limit → Unnest → SubqueryAlias → Projection → TableScan
+    // Combines the Limit and SubqueryAlias transparent node patterns.
+    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)")])?
+        .alias("t")?
+        .unnest_column("__unnest_placeholder(items)")?
+        .limit(0, Some(10))?
+        .project(vec![col("__unnest_placeholder(items)").alias("item")])?
+        .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