kosiew commented on code in PR #25482:
URL: https://github.com/apache/datafusion/pull/25482#discussion_r4205069179
##########
datafusion/core/tests/dataframe/mod.rs:
##########
@@ -7800,3 +7803,463 @@ async fn test_unresolved_lambda_variable() ->
Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_dataframe_api_select_semantics() -> Result<()> {
Review Comment:
Could you also add a small empty-input case using
`filter(lit(false)).select(...)` with distinct aliases? It would be useful to
verify that global aggregation still returns exactly one row, with `COUNT = 0`,
`SUM = NULL`, and the literal value preserved.
##########
datafusion/core/src/dataframe/mod.rs:
##########
@@ -380,53 +510,96 @@ impl DataFrame {
self.select(expr_list)
}
- /// Project arbitrary expressions (like SQL SELECT expressions) into a new
- /// `DataFrame`.
+ /// Project expressions into a new [`DataFrame`], like SQL SELECT.
+ ///
+ /// The output has one column for each element in expr_list.
///
- /// The output `DataFrame` has one column for each element in `expr_list`.
+ /// A list of only aggregates is treated as a global aggregation
+ /// ([`Self::aggregate`] with no group columns). To group, use
+ /// [`Self::aggregate`].
///
/// # Example
/// ```
/// # use datafusion::prelude::*;
/// # use datafusion::error::Result;
+ /// # use datafusion::functions_aggregate::expr_fn::{count, sum};
/// # use datafusion_common::assert_batches_sorted_eq;
/// # #[tokio::main]
/// # async fn main() -> Result<()> {
/// let ctx = SessionContext::new();
/// let df = ctx
/// .read_csv("tests/data/example.csv", CsvReadOptions::new())
/// .await?;
- /// let df = df.select(vec![col("a"), col("b") * col("c")])?;
+ ///
+ /// // Per-row projection
+ /// let res = df.clone().select(vec![col("a"), col("b") * col("c")])?;
/// let expected = vec![
/// "+---+-----------------------+",
/// "| a | ?table?.b * ?table?.c |",
/// "+---+-----------------------+",
/// "| 1 | 6 |",
/// "+---+-----------------------+",
/// ];
- /// # assert_batches_sorted_eq!(expected, &df.collect().await?);
+ /// # assert_batches_sorted_eq!(expected, &res.collect().await?);
+ ///
+ /// // Global aggregation (no GROUP BY) — same as aggregate(vec![], ...)
+ /// let res = df.select(vec![
+ /// count(col("a")).alias("cnt"),
+ /// sum(col("b")).alias("sum_b"),
+ /// ])?;
+ /// let expected = vec![
+ /// "+-----+-------+",
+ /// "| cnt | sum_b |",
+ /// "+-----+-------+",
+ /// "| 1 | 2 |",
+ /// "+-----+-------+",
+ /// ];
+ /// # assert_batches_sorted_eq!(expected, &res.collect().await?);
/// # Ok(())
/// # }
/// ```
pub fn select(
self,
expr_list: impl IntoIterator<Item = impl Into<SelectExpr>>,
) -> Result<DataFrame> {
- let expr_list: Vec<SelectExpr> =
+ let mut expr_list: Vec<SelectExpr> =
expr_list.into_iter().map(|e| e.into()).collect::<Vec<_>>();
- let expressions = expr_list.iter().filter_map(|e| match e {
- SelectExpr::Expression(expr) => Some(expr),
- _ => None,
- });
+ let expressions: Vec<&Expr> = expr_list
+ .iter()
+ .filter_map(|e| match e {
+ SelectExpr::Expression(expr) => Some(expr),
+ _ => None,
+ })
+ .collect();
- let window_func_exprs = find_window_exprs(expressions);
- let plan = if window_func_exprs.is_empty() {
- self.plan
- } else {
+ // 1. Windows (row-preserving)
+ let window_func_exprs = find_window_exprs(expressions.iter().copied());
+ let has_windows = !window_func_exprs.is_empty();
+ let mut plan = if has_windows {
LogicalPlanBuilder::window_plan(self.plan, window_func_exprs)?
+ } else {
+ self.plan
};
+ // 2. Global aggregate only (no GROUP BY). Mixed lists and
+ // aggregate().select() reshapes stay on the projection path.
+ let aggr_exprs = find_aggregate_exprs(expressions.iter().copied());
+ if !aggr_exprs.is_empty()
+ && !has_windows
+ && should_apply_global_aggregate(&expr_list, &plan)?
+ {
+ let input = plan;
+ let (aggr_with_alias, rewrite_map) =
+ alias_aggregate_exprs(aggr_exprs, &input)?;
+ plan = LogicalPlanBuilder::from(input.clone())
+ .aggregate(Vec::<Expr>::new(), aggr_with_alias)?
+ .build()?;
+ expr_list = rewrite_select_aggs(expr_list, &input, &rewrite_map)?;
+ expr_list = uniquify_select_expr_names(expr_list)?;
Review Comment:
This now silently renames duplicate explicit aliases before the existing
projection validator can reject them, so `select` behaves differently from
ordinary projections and explicit `aggregate`. Please remove the public
output-name uniquification and let the existing projection validation reject
duplicate names, while keeping unique internal aggregate aliases where needed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
