cj-zhukov commented on code in PR #25482:
URL: https://github.com/apache/datafusion/pull/25482#discussion_r4209156396
##########
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:
I agree with you
--
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]
