Extend expression planners for Snowflake compatibility

Cargo.toml

@@ -196,7 +196,10 @@ regex = "1.12"
 rstest = "0.26.1"
 serde_json = "1"
 sha2 = "^0.11.0"
-sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor"] }
+sqlparser = { git = "https://github.com/Embucket/datafusion-sqlparser-rs.git", rev = "1a3f48f60802d0f135a968a27c89a3351e78b159", default-features = false, features = [
+    "std",
+    "visitor",
+] }
 stacker = "0.1.24"
 strum = "0.28.0"
 strum_macros = "0.28.0"

datafusion/common/src/config.rs

@@ -271,6 +271,10 @@ config_namespace! {
         /// When set to true, SQL parser will parse float as decimal type
         pub parse_float_as_decimal: bool, default = false
 
+        /// When set to true, insignificant trailing zeros are removed from decimal literals.
+        /// For example, `10.00` is planned as `DECIMAL(2, 0)` instead of `DECIMAL(4, 2)`.
+        pub trim_decimal_literal_trailing_zeros: bool, default = false
+
         /// When set to true, SQL parser will normalize ident (convert ident to lowercase when not quoted)
         pub enable_ident_normalization: bool, default = true
 

datafusion/core/src/datasource/view_test.rs

@@ -289,10 +289,10 @@ mod tests {
 
         insta::assert_snapshot!(batches_to_string(&results),@r"
         +---------+---------+---------+
-        | column2 | column1 | column3 |
+        | column1 | column2 | column3 |
         +---------+---------+---------+
-        | 2       | 1       | 3       |
-        | 5       | 4       | 6       |
+        | 1       | 2       | 3       |
+        | 4       | 5       | 6       |
         +---------+---------+---------+
         ");
 

datafusion/core/src/execution/session_state.rs

@@ -581,6 +581,8 @@ impl SessionState {
 
         ParserOptions {
             parse_float_as_decimal: sql_parser_options.parse_float_as_decimal,
+            trim_decimal_literal_trailing_zeros: sql_parser_options
+                .trim_decimal_literal_trailing_zeros,
             enable_ident_normalization: sql_parser_options.enable_ident_normalization,
             enable_options_value_normalization: sql_parser_options
                 .enable_options_value_normalization,

datafusion/core/tests/sql/joins.rs

@@ -25,6 +25,33 @@ use datafusion_sql::unparser::plan_to_sql;
 
 use super::*;
 
+#[tokio::test]
+async fn natural_full_join_wildcard_coalesces_join_key() -> Result<()> {
+    let ctx = SessionContext::new();
+    let dataframe = ctx
+        .sql(
+            "WITH d1(id, name) AS (VALUES (1, 'a'), (2, 'b'), (4, 'c')),
+                  d2(id, value) AS (VALUES (1, 'xx'), (2, 'yy'), (5, 'zz'))
+             SELECT * FROM d1 NATURAL FULL OUTER JOIN d2 ORDER BY id",
+        )
+        .await?;
+
+    assert_batches_eq!(
+        [
+            "+----+------+-------+",
+            "| id | name | value |",
+            "+----+------+-------+",
+            "| 1  | a    | xx    |",
+            "| 2  | b    | yy    |",
+            "| 4  | c    |       |",
+            "| 5  |      | zz    |",
+            "+----+------+-------+",
+        ],
+        &dataframe.collect().await?
+    );
+    Ok(())
+}
+
 #[tokio::test]
 async fn join_change_in_planner() -> Result<()> {
     let config = SessionConfig::new().with_target_partitions(8);

datafusion/core/tests/sql/pivot_unpivot.rs

@@ -33,12 +33,48 @@ async fn pivot_list_is_lowered_to_filtered_aggregates() -> Result<()> {
         .await?;
 
     insta::assert_snapshot!(batches_to_string(&batches), @r"
-    +----+----+----+
-    | id | a  | b  |
-    +----+----+----+
-    | 1  | 10 | 20 |
-    | 2  | 7  |    |
-    +----+----+----+
+    +----+-----+-----+
+    | id | 'a' | 'b' |
+    +----+-----+-----+
+    | 1  | 10  | 20  |
+    | 2  | 7   |     |
+    +----+-----+-----+
+    ");
+    Ok(())
+}
+
+#[tokio::test]
+async fn chained_pivots_apply_explicit_column_aliases() -> Result<()> {
+    let batches = SessionContext::new()
+        .sql(
+            "SELECT SUM(q1_sales) AS q1_sales,
+                    SUM(q2_sales) AS q2_sales,
+                    MAX(q1_discount) AS q1_discount,
+                    MAX(q2_discount) AS q2_discount
+             FROM (
+                 SELECT amount,
+                        quarter AS sales_quarter,
+                        quarter AS discount_quarter,
+                        discount
+                 FROM (VALUES
+                     (100, 'Q1', 10),
+                     (200, 'Q2', 20)
+                 ) sales(amount, quarter, discount)
+             )
+             PIVOT(SUM(amount) FOR sales_quarter IN ('Q1', 'Q2'))
+             PIVOT(MAX(discount) FOR discount_quarter IN ('Q1', 'Q2'))
+             AS p(q1_sales, q2_sales, q1_discount, q2_discount)",
+        )
+        .await?
+        .collect()
+        .await?;
+
+    insta::assert_snapshot!(batches_to_string(&batches), @r"
+    +----------+----------+-------------+-------------+
+    | q1_sales | q2_sales | q1_discount | q2_discount |
+    +----------+----------+-------------+-------------+
+    | 100      | 200      | 10          | 20          |
+    +----------+----------+-------------+-------------+
     ");
     Ok(())
 }

datafusion/core/tests/sql/unparser.rs

@@ -378,7 +378,7 @@ async fn optimized_duckdb_unparse_qualifies_nested_passthrough_column() -> Resul
     // `o` (which is only the base-table alias one level deeper). The bug emitted
     // `"o"."order_id"` inside that derived table; the fix emits a bare column.
     let expected = concat!(
-        r#"SELECT "o"."order_id", "o"."discount_pct_2" "#,
+        r#"SELECT "oi"."order_id", "o"."discount_pct_2" "#,
         r#"FROM "warehouse"."main"."order_items" AS "oi" "#,
         r#"INNER JOIN (SELECT "order_id", "#,
         r#"CASE WHEN "__common_expr_1" IS NOT NULL "#,

datafusion/core/tests/user_defined/user_defined_scalar_functions.rs

@@ -609,6 +609,24 @@ async fn volatile_scalar_udf_with_params() -> Result<()> {
         ];
         assert_batches_eq!(expected, &result);
 
+        let result = plan_and_collect(
+            &ctx,
+            "SELECT * FROM (SELECT add_index_to_string('test') AS str FROM t) \
+             ORDER BY add_index_to_string('test') LIMIT 2 OFFSET 5",
+        )
+        .await?;
+        assert_batches_eq!(
+            [
+                "+---------+",
+                "| str     |",
+                "+---------+",
+                "| 6) test |",
+                "| 7) test |",
+                "+---------+",
+            ],
+            &result
+        );
+
         let result =
             plan_and_collect(&ctx, "select add_index_to_string('test_value') as str") // with fixed function parameters
                 .await?;

datafusion/expr/src/planner.rs

@@ -154,6 +154,17 @@ pub trait ContextProvider {
 ///
 /// [Extending SQL in DataFusion: from ->> to TABLESAMPLE blog]: https://datafusion.apache.org/blog/2026/01/12/extending-sql
 pub trait ExprPlanner: Debug + Send + Sync {
+    /// Plans scalar functions, such as `CONCAT(<expr>, ...)`, with access to the input schema.
+    ///
+    /// Returns the original scalar function if planning is not possible.
+    fn plan_scalar_with_schema(
+        &self,
+        expr: RawScalarExpr,
+        _schema: &DFSchema,
+    ) -> Result<PlannerResult<RawScalarExpr>> {
+        Ok(PlannerResult::Original(expr))
+    }
+
     /// Plan the binary operation between two expressions, returns original
     /// BinaryExpr if not possible
     fn plan_binary_op(
@@ -273,6 +284,19 @@ pub trait ExprPlanner: Debug + Send + Sync {
         Ok(PlannerResult::Original(expr))
     }
 
+    /// Plans aggregate functions with access to the input schema.
+    ///
+    /// The default implementation delegates to [`Self::plan_aggregate`] so existing planners do
+    /// not need to change. Planners that need to resolve schema-dependent arguments, such as a
+    /// qualified wildcard, can override this method instead.
+    fn plan_aggregate_with_schema(
+        &self,
+        expr: RawAggregateExpr,
+        _schema: &DFSchema,
+    ) -> Result<PlannerResult<RawAggregateExpr>> {
+        self.plan_aggregate(expr)
+    }
+
     /// Plans window functions, such as `COUNT(<expr>)`
     ///
     /// Returns original expression arguments if not possible
@@ -330,6 +354,13 @@ pub struct RawAggregateExpr {
     pub null_treatment: Option<NullTreatment>,
 }
 
+/// This structure is used by scalar function expression planners.
+#[derive(Debug, Clone)]
+pub struct RawScalarExpr {
+    pub func: Arc<ScalarUDF>,
+    pub args: Vec<Expr>,
+}
+
 /// This structure is used by `WindowFunctionPlanner` to plan operators with
 /// custom expressions.
 #[derive(Debug, Clone)]