Working on unnest support for snowflake

Author: yonatan-sevenai

datafusion/sql/src/unparser/ast.rs

@@ -432,11 +432,11 @@ pub struct RelationBuilder {
 }
 
 #[derive(Clone)]
-#[expect(clippy::large_enum_variant)]
 enum TableFactorBuilder {
     Table(TableRelationBuilder),
     Derived(DerivedRelationBuilder),
     Unnest(UnnestRelationBuilder),
+    Flatten(FlattenRelationBuilder),
     Empty,
 }
 
@@ -458,6 +458,11 @@ impl RelationBuilder {
         self
     }
 
+    pub fn flatten(&mut self, value: FlattenRelationBuilder) -> &mut Self {
+        self.relation = Some(TableFactorBuilder::Flatten(value));
+        self
+    }
+
     pub fn empty(&mut self) -> &mut Self {
         self.relation = Some(TableFactorBuilder::Empty);
         self
@@ -474,6 +479,9 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => {
                 rel_builder.alias = value;
             }
+            Some(TableFactorBuilder::Flatten(ref mut rel_builder)) => {
+                rel_builder.alias = value;
+            }
             Some(TableFactorBuilder::Empty) => (),
             None => (),
         }
@@ -484,6 +492,7 @@ impl RelationBuilder {
             Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?),
+            Some(TableFactorBuilder::Flatten(ref value)) => Some(value.build()?),
             Some(TableFactorBuilder::Empty) => None,
             None => return Err(Into::into(UninitializedFieldError::from("relation"))),
         })
@@ -688,6 +697,94 @@ impl Default for UnnestRelationBuilder {
     }
 }
 
+/// Default table alias for FLATTEN table factors.
+/// Snowflake requires an alias to reference output columns (e.g. `_unnest.VALUE`).
+pub const FLATTEN_DEFAULT_ALIAS: &str = "_unnest";
+
+/// Builds a `LATERAL FLATTEN(INPUT => expr, OUTER => bool)` table factor
+/// for Snowflake-style unnesting.
+#[derive(Clone)]
+pub struct FlattenRelationBuilder {
+    pub alias: Option<ast::TableAlias>,
+    /// The input expression to flatten (e.g. a column reference).
+    pub input_expr: Option<ast::Expr>,
+    /// Whether to preserve rows for NULL/empty inputs (Snowflake `OUTER` param).
+    pub outer: bool,
+}
+
+impl FlattenRelationBuilder {
+    pub fn alias(&mut self, value: Option<ast::TableAlias>) -> &mut Self {
+        self.alias = value;
+        self
+    }
+
+    pub fn input_expr(&mut self, value: ast::Expr) -> &mut Self {
+        self.input_expr = Some(value);
+        self
+    }
+
+    pub fn outer(&mut self, value: bool) -> &mut Self {
+        self.outer = value;
+        self
+    }
+
+    pub fn build(&self) -> Result<ast::TableFactor, BuilderError> {
+        let input = self.input_expr.clone().ok_or_else(|| {
+            BuilderError::from(UninitializedFieldError::from("input_expr"))
+        })?;
+
+        let mut args = vec![ast::FunctionArg::Named {
+            name: ast::Ident::new("INPUT"),
+            arg: ast::FunctionArgExpr::Expr(input),
+            operator: ast::FunctionArgOperator::RightArrow,
+        }];
+
+        if self.outer {
+            args.push(ast::FunctionArg::Named {
+                name: ast::Ident::new("OUTER"),
+                arg: ast::FunctionArgExpr::Expr(ast::Expr::Value(
+                    ast::Value::Boolean(true).into(),
+                )),
+                operator: ast::FunctionArgOperator::RightArrow,
+            });
+        }
+
+        Ok(ast::TableFactor::Function {
+            lateral: true,
+            name: ast::ObjectName::from(vec![ast::Ident::new("FLATTEN")]),
+            args,
+            alias: self.alias.clone(),
+        })
+    }
+
+    /// Returns the alias name for this FLATTEN relation.
+    /// Used to build qualified column references like `alias.VALUE`.
+    pub fn alias_name(&self) -> &str {
+        self.alias
+            .as_ref()
+            .map(|a| a.name.value.as_str())
+            .unwrap_or(FLATTEN_DEFAULT_ALIAS)
+    }
+
+    fn create_empty() -> Self {
+        Self {
+            alias: Some(ast::TableAlias {
+                name: ast::Ident::new(FLATTEN_DEFAULT_ALIAS),
+                columns: vec![],
+                explicit: true,
+            }),
+            input_expr: None,
+            outer: false,
+        }
+    }
+}
+
+impl Default for FlattenRelationBuilder {
+    fn default() -> Self {
+        Self::create_empty()
+    }
+}
+
 /// Runtime error when a `build()` method is called and one or more required fields
 /// do not have a value.
 #[derive(Debug, Clone)]

datafusion/sql/src/unparser/dialect.rs

@@ -206,6 +206,15 @@ pub trait Dialect: Send + Sync {
         false
     }
 
+    /// Unparse the unnest plan as `LATERAL FLATTEN(INPUT => expr, ...)`.
+    ///
+    /// Snowflake uses FLATTEN as a table function instead of the SQL-standard UNNEST.
+    /// When this returns `true`, the unparser emits
+    /// `LATERAL FLATTEN(INPUT => <col>, OUTER => <bool>)` in the FROM clause.
+    fn unnest_as_lateral_flatten(&self) -> bool {
+        false
+    }
+
     /// Allows the dialect to override column alias unparsing if the dialect has specific rules.
     /// Returns None if the default unparsing should be used, or Some(String) if there is
     /// a custom implementation for the alias.
@@ -664,6 +673,59 @@ impl BigQueryDialect {
     }
 }
 
+/// Dialect for Snowflake SQL.
+///
+/// Key differences from the default dialect:
+/// - Uses double-quote identifier quoting
+/// - Supports `NULLS FIRST`/`NULLS LAST` in `ORDER BY`
+/// - Does not support empty select lists (`SELECT FROM t`)
+/// - Does not support column aliases in table alias definitions
+///   (Snowflake accepts the syntax but silently ignores the renames in join contexts)
+/// - Unparses `UNNEST` plans as `LATERAL FLATTEN(INPUT => expr, ...)`
+pub struct SnowflakeDialect {}
+
+#[expect(clippy::new_without_default)]
+impl SnowflakeDialect {
+    #[must_use]
+    pub fn new() -> Self {
+        Self {}
+    }
+}
+
+impl Dialect for SnowflakeDialect {
+    fn identifier_quote_style(&self, _: &str) -> Option<char> {
+        Some('"')
+    }
+
+    fn supports_nulls_first_in_sort(&self) -> bool {
+        true
+    }
+
+    fn supports_empty_select_list(&self) -> bool {
+        false
+    }
+
+    fn supports_column_alias_in_table_alias(&self) -> bool {
+        false
+    }
+
+    fn timestamp_cast_dtype(
+        &self,
+        _time_unit: &TimeUnit,
+        tz: &Option<Arc<str>>,
+    ) -> ast::DataType {
+        if tz.is_some() {
+            ast::DataType::Timestamp(None, TimezoneInfo::WithTimeZone)
+        } else {
+            ast::DataType::Timestamp(None, TimezoneInfo::None)
+        }
+    }
+
+    fn unnest_as_lateral_flatten(&self) -> bool {
+        true
+    }
+}
+
 pub struct CustomDialect {
     identifier_quote_style: Option<char>,
     supports_nulls_first_in_sort: bool,
@@ -686,6 +748,7 @@ pub struct CustomDialect {
     window_func_support_window_frame: bool,
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
+    unnest_as_lateral_flatten: bool,
 }
 
 impl Default for CustomDialect {
@@ -715,6 +778,7 @@ impl Default for CustomDialect {
             window_func_support_window_frame: true,
             full_qualified_col: false,
             unnest_as_table_factor: false,
+            unnest_as_lateral_flatten: false,
         }
     }
 }
@@ -829,6 +893,10 @@ impl Dialect for CustomDialect {
     fn unnest_as_table_factor(&self) -> bool {
         self.unnest_as_table_factor
     }
+
+    fn unnest_as_lateral_flatten(&self) -> bool {
+        self.unnest_as_lateral_flatten
+    }
 }
 
 /// `CustomDialectBuilder` to build `CustomDialect` using builder pattern
@@ -867,6 +935,7 @@ pub struct CustomDialectBuilder {
     window_func_support_window_frame: bool,
     full_qualified_col: bool,
     unnest_as_table_factor: bool,
+    unnest_as_lateral_flatten: bool,
 }
 
 impl Default for CustomDialectBuilder {
@@ -902,6 +971,7 @@ impl CustomDialectBuilder {
             window_func_support_window_frame: true,
             full_qualified_col: false,
             unnest_as_table_factor: false,
+            unnest_as_lateral_flatten: false,
         }
     }
 
@@ -929,6 +999,7 @@ impl CustomDialectBuilder {
             window_func_support_window_frame: self.window_func_support_window_frame,
             full_qualified_col: self.full_qualified_col,
             unnest_as_table_factor: self.unnest_as_table_factor,
+            unnest_as_lateral_flatten: self.unnest_as_lateral_flatten,
         }
     }
 
@@ -1075,4 +1146,12 @@ impl CustomDialectBuilder {
         self.unnest_as_table_factor = unnest_as_table_factor;
         self
     }
+
+    pub fn with_unnest_as_lateral_flatten(
+        mut self,
+        unnest_as_lateral_flatten: bool,
+    ) -> Self {
+        self.unnest_as_lateral_flatten = unnest_as_lateral_flatten;
+        self
+    }
 }