DataFrame API: allow aggregate functions in select() (#17874)

[cj-zhukov]

Which issue does this PR close?

• Closes #DataFrame API: allow aggregate functions in select() #17874.

Rationale for this change

DataFrame::select currently does not support aggregate expressions directly.
Users must explicitly call DataFrame::aggregate. This PR improves ergonomics by allowing aggregate expressions inside select, while preserving existing behavior and validation rules.

What changes are included in this PR?

• Added support for using aggregate functions directly in DataFrame::select
• Updated snapshot tests to reflect the new logical plan datafusion-cli/tests/snapshots/cli_explain_environment_overrides@explain_plan_environment_overrides.snap

Are these changes tested?

• Existing tests continue to pass
• Added new test for this new use case test_dataframe_api_aggregate_fn_in_select

Are there any user-facing changes?

Yes (behavioral improvement, no API changes).
Users can now use aggregate functions directly in select() without explicitly calling aggregate().

ctx.table("aggregate_test_100")
    .await?
    .select(vec![
        approx_distinct(col("c9")).alias("count_c9"),
        approx_distinct(cast(col("c9"), arrow_schema::DataType::Utf8View))
            .alias("count_c9_str"),
    ])?
    .show()
    .await?;

Previously, this required:

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

Both syntaxes are now valid and can be used interchangeably.

There are no breaking changes to the public API.

[martin-g]

You need to update the Insta snapshot for the datafusion-cli crate

[martin-g]

Hm, you modified the Insta snapshot!
But it is failing now.
So, maybe you have to revert the change ?!

[cj-zhukov]

That helped - thank you!

[martin-g]

find_aggregate_exprs() deduplicates the expressions.
Test like:

let res = df.select(vec![
        count(col("c9")).alias("count_c9"),
        count(col("c9")).alias("count_c9_str"),
    ])?;

fails with:

failures:

---- dataframe::test_dataframe_api_aggregate_fn_in_select2 stdout ----
Error: SchemaError(FieldNotFound { field: Column { relation: None, name: "__agg_1" }, valid_fields: [Column { relation: None, name: "__agg_0" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c1" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c2" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c3" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c4" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c5" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c6" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c7" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c8" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c9" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c10" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c11" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c12" }, Column { relation: Some(Bare { table: "aggregate_test_100" }), name: "c13" }] }, Some(""))

__agg_1 is lost

[martin-g]

What about aggregate expr + non-aggregate one ?
E.g.:

  let res = df.select(vec![
        count(col("c9")).alias("count_c9") + lit(1)
    ])?;

I'd expect 101 but it returns 100

[martin-g]

Please add tests for some more complex queries, e.g. df.select([sum(col("a")) + count(col("b"))]) and something with (qualified and non-qualified) wildcards too.

[martin-g]

__agg_0 could collide with a real column. Is this how it is being done elsewhere ?