Skip to content

Unparser with PostgreSqlDialect fails after Optimization Pass #21066

@cetra3

Description

@cetra3

Describe the bug

The Unparser has a bug in the way that it treats certain PostgreSQL formatted queries, after optimization, which seems to be because requires_derived_table_alias is true:

I.e, optimizing then unparsing this query with postgresql dialect breaks:

WITH base AS (SELECT name, salary FROM t)
SELECT name,
    CASE WHEN SUM(salary) > 0 THEN 1 ELSE 0 END AS x,
    CASE WHEN SUM(salary) > 0 THEN SUM(salary) ELSE 0 END AS y
FROM base GROUP BY name

To Reproduce

Here's a test we can add somewhere that fails:

#[tokio::test]
async fn test_cse_derived_projection_roundtrip() {
    let ctx = SessionContext::new();
    ctx.sql("CREATE TABLE t (name TEXT, salary DOUBLE) AS VALUES ('a', 1.0), ('b', 2.0)")
        .await.unwrap().collect().await.unwrap();

    let sql = "\
        WITH base AS (SELECT name, salary FROM t) \
        SELECT name, \
            CASE WHEN SUM(salary) > 0 THEN 1 ELSE 0 END AS x, \
            CASE WHEN SUM(salary) > 0 THEN SUM(salary) ELSE 0 END AS y \
        FROM base GROUP BY name";

    let df = ctx.sql(sql).await.unwrap();
    let optimized = ctx.state().optimize(df.logical_plan()).unwrap();
    let unparser = Unparser::new(&datafusion_sql::unparser::dialect::PostgreSqlDialect {});
    let unparsed = unparser.plan_to_sql(&optimized).unwrap().to_string();

    ctx.sql(&unparsed).await.unwrap_or_else(|e| {
        panic!("Roundtrip failed.\n\nOriginal: {sql}\nUnparsed: {unparsed}\nError: {e}")
    });
}

Expected behavior

This should parse correctly

Additional context

No response

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions