Fix `csv_query_approx_count` once binary type coercion PR is merged

[andygrove]

Describe the bug

csv_query_approx_count fails due to an edge case exposed in #3222

The test fails with:

"Plan(\"Aggregations require unique expression names but the expression \\\"APPROXDISTINCT(#aggregate_test_100.c9)\\\" at position 0 and \\\"APPROXDISTINCT(CAST(#aggregate_test_100.c9 AS Utf8))\\\" at position 1 have the same name. Consider aliasing (\\\"AS\\\") one of them.\")

However, the expressions are already aliased. The alias is dropped before the error check above and this part of the codebase needs some reworking. It is not trivial to fix.

To Reproduce

#[tokio::test]
async fn csv_query_approx_count() -> Result<()> {
    let ctx = SessionContext::new();
    register_aggregate_csv(&ctx).await?;
    let sql = "SELECT approx_distinct(c9) count_c9, approx_distinct(cast(c9 as varchar)) count_c9_str FROM aggregate_test_100";
    let actual = execute_to_batches(&ctx, sql).await;
    let expected = vec![
        "+----------+--------------+",
        "| count_c9 | count_c9_str |",
        "+----------+--------------+",
        "| 100      | 99           |",
        "+----------+--------------+",
    ];
    assert_batches_eq!(expected, &actual);
    Ok(())
}

Expected behavior
Should pass

Additional context
None

[Jefffrey]

Was poking around this a bit, some notes.

Updated context

Test is now located here:

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Lines 1418 to 1423 in 73562e8

	# csv_query_approx_count
	# FIX: https://github.com/apache/datafusion/issues/3353
	# query II
	# SELECT approx_distinct(c9) AS count_c9, approx_distinct(cast(c9 as varchar)) count_c9_str FROM aggregate_test_100
	# ----
	# 100 99

There is also an expected error test here:

datafusion/datafusion/sqllogictest/test_files/aggregate.slt

Lines 142 to 144 in 73562e8

	# https://github.com/apache/datafusion/issues/3353
	statement error DataFusion error: Schema error: Schema contains duplicate unqualified field name "approx_distinct\(aggregate_test_100\.c9\)"
	SELECT approx_distinct(c9) count_c9, approx_distinct(cast(c9 as varchar)) count_c9_str FROM aggregate_test_100

Cause

In the SQL planner we find the aggregate expressions here:

datafusion/datafusion/sql/src/select.rs

Lines 236 to 255 in c89cb70

	// The outer expressions we will search through for aggregates.
	// First, find aggregates in SELECT, HAVING, and QUALIFY
	let select_having_qualify_aggrs = find_aggregate_exprs(
	select_exprs
	.iter()
	.chain(having_expr_opt.iter())
	.chain(qualify_expr_opt.iter()),
	);
	// Find aggregates in ORDER BY
	let order_by_aggrs = find_aggregate_exprs(order_by_rex.iter().map(|s| &s.expr));
	// Combine: all aggregates from SELECT/HAVING/QUALIFY, plus ORDER BY aggregates
	// that aren't already in SELECT/HAVING/QUALIFY
	let mut aggr_exprs = select_having_qualify_aggrs;
	for order_by_aggr in order_by_aggrs {
	if !aggr_exprs.iter().any(|e| e == &order_by_aggr) {
	aggr_exprs.push(order_by_aggr);
	}
	}

See find_aggregate_exprs(), which essentially drops the aliases as mentioned:

datafusion/datafusion/expr/src/utils.rs

Lines 608 to 615 in c89cb70

	/// Collect all deeply nested `Expr::AggregateFunction`.
	/// They are returned in order of occurrence (depth
	/// first), with duplicates omitted.
	pub fn find_aggregate_exprs<'a>(exprs: impl IntoIterator<Item = &'a Expr>) -> Vec<Expr> {
	find_exprs_in_exprs(exprs, &|nested_expr| {
	matches!(nested_expr, Expr::AggregateFunction { .. })
	})
	}

We could alter this to preserve alises like so:

pub fn find_aggregate_exprs<'a>(exprs: impl IntoIterator<Item = &'a Expr>) -> Vec<Expr> {
    find_exprs_in_exprs(exprs, &|nested_expr| {
-        matches!(nested_expr, Expr::AggregateFunction { .. })
+        let is_agg = matches!(nested_expr, Expr::AggregateFunction { .. });
+        let is_alias_agg = if let Expr::Alias(Alias { expr, .. }) = nested_expr {
+            matches!(**expr, Expr::AggregateFunction { .. })
+        } else {
+            false
+        };
+        is_agg || is_alias_agg
    })
}

Which technically fixes this query, however has a knock-on effect, for example breaking a few queries in the tpcds_planning suite.

• There's also some changes to plans since it allows the alias to stay within an aggregate node instead of being a separate projection node

Something interesting is that we can replicate this query via the DataFrame API which allows this query; this makes sense since the bug seems to be in the SQL planning code. See:

    let ctx = SessionContext::new();

    let testdata = datafusion::test_util::arrow_test_data();
    ctx.register_csv(
        "aggregate_test_100",
        format!("{testdata}/csv/aggregate_test_100.csv"),
        CsvReadOptions::default(),
    )
    .await?;

    println!("DF api");
    ctx.table("aggregate_test_100")
        .await?
        .aggregate(
            vec![],
            vec![
                approx_distinct(col("c9")).alias("count_c9"),
                approx_distinct(cast(col("c9"), DataType::Utf8View))
                    .alias("count_c9_str"),
            ],
        )?
        .explain(false, false)?
        .show()
        .await?;

    println!("SQL api");
    let query = "SELECT
      approx_distinct(c9) AS count_c9,
      approx_distinct(cast(c9 as varchar)) count_c9_str
    FROM aggregate_test_100";
    ctx.sql(query).await?.explain(false, false)?.show().await?;

Where the output (on current main) is:

DF api
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                                                                                                      |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan  | Aggregate: groupBy=[[]], aggr=[[approx_distinct(aggregate_test_100.c9) AS count_c9, approx_distinct(CAST(aggregate_test_100.c9 AS Utf8View)) AS count_c9_str]]            |
|               |   TableScan: aggregate_test_100 projection=[c9]                                                                                                                           |
| physical_plan | AggregateExec: mode=Final, gby=[], aggr=[count_c9, count_c9_str]                                                                                                          |
|               |   CoalescePartitionsExec                                                                                                                                                  |
|               |     AggregateExec: mode=Partial, gby=[], aggr=[count_c9, count_c9_str]                                                                                                    |
|               |       RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1                                                                                               |
|               |         DataSourceExec: file_groups={1 group: [[Users/jeffrey/Code/datafusion/testing/data/csv/aggregate_test_100.csv]]}, projection=[c9], file_type=csv, has_header=true |
|               |                                                                                                                                                                           |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
SQL api
Error: SchemaError(DuplicateUnqualifiedField { name: "approx_distinct(aggregate_test_100.c9)" }, Some(""))

So I wonder if we should try fix SQL API to be permissive like DataFrame API (preserving aliases in aggregations) but that seems to require larger scale fixes.

[Jefffrey]

#18975 also raised point that perhaps issue is at cast display level; for example, in the original query:

SELECT
    approx_distinct(c9) AS count_c9,
    approx_distinct(cast(c9 as varchar)) count_c9_str
FROM aggregate_test_100 

If we remove the aliases we'd still run into the same issue, so fixing the aggregations to preserve aliases might not be a complete fix.