You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/22 02:53:48 UTC
(datafusion) branch gh-readonly-queue/main/pr-25482-83d2e9d12f66c61a1259c80f6df0c242299ad4b3 created (now 0cba642cef)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-25482-83d2e9d12f66c61a1259c80f6df0c242299ad4b3
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 0cba642cef DataFrame API: allow aggregate functions in select() 2 (#17874) (#25482)
This branch includes the following new commits:
new 0cba642cef DataFrame API: allow aggregate functions in select() 2 (#17874) (#25482)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: DataFrame API: allow aggregate functions in select() 2 (#17874) (#25482)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-25482-83d2e9d12f66c61a1259c80f6df0c242299ad4b3
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 0cba642ceff7c75e701f88101a7d032f00426174
Author: Sergey Zhukov <62...@users.noreply.github.com>
AuthorDate: Thu Oct 22 02:53:28 2026 +0000
DataFrame API: allow aggregate functions in select() 2 (#17874) (#25482)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Closes #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.
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
Please explain the problem you are trying to solve in terms of the
user-visible
behavior, rather than the implementation.
For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of
the
implementation. "COUNT(DISTINCT) returns wrong results when the column
contains
nulls" is the user-visible problem.
-->
## What changes are included in this PR?
- `DataFrame::select` lifts a global Aggregate when every expression is
an aggregate or a scalar over aggregates (and literals).
- Existing tests were not rewritten.
- Examples are left for a follow-up.
<!--
There is no need to duplicate the description in the issue here, but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
## What is the testing strategy for this PR?
- Existing `DataFrame` select/aggregate tests continue to pass without
edits.
- Added `test_dataframe_api_select_semantics` (from #21021, plus two
cases for this design).
- `roundtrip_expr_api` was not changed and still passes (mixed aggregate
+ scalar lists stay a `Projection`).
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
Briefly describe how this PR is tested, and point to the specific tests
you added. For example: 'This new feature is covered by the
`sqllogictest` cases added in `foo.slt`'.
If this PR does not add tests, explain why. For example, if the change
is already covered by existing tests, please mention it.
You should also check the `codecov` bot reply on this PR to confirm the
changed code is exercised.
-->
## Are there any user-facing changes?
Yes — additive behavior only. The public `select()` signature is
unchanged. There are no breaking API changes.
Users can write a global aggregation with `select()`:
```rust
ctx.table("aggregate_test_100")
.await?
.select(vec![
approx_distinct(col("c9")).alias("count_c9"),
approx_distinct(cast(col("c9"), DataType::Utf8View))
.alias("count_c9_str"),
])?
.show()
.await?;
```
Previously this required:
```rust
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"),
],
)?
.show()
.await?;
```
Both remain valid. Grouped aggregation is still aggregate(group, aggs).
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---------
Co-authored-by: kosiew <ko...@gmail.com>
---
datafusion/core/src/dataframe/mod.rs | 188 ++++++++++++-
datafusion/core/tests/dataframe/mod.rs | 494 ++++++++++++++++++++++++++++++++-
2 files changed, 667 insertions(+), 15 deletions(-)
diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs
index 0cf86aa969..d0738b9897 100644
--- a/datafusion/core/src/dataframe/mod.rs
+++ b/datafusion/core/src/dataframe/mod.rs
@@ -51,12 +51,15 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::util::display::{ArrayFormatter, FormatOptions};
use arrow_schema::FieldRef;
use datafusion_common::config::{CsvOptions, JsonOptions};
+use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
use datafusion_common::{
Column, DFSchema, DataFusionError, ParamValues, ScalarValue, SchemaError,
TableReference, UnnestOptions, exec_err, internal_datafusion_err, not_impl_err,
plan_datafusion_err, plan_err, unqualified_field_not_found,
};
+use datafusion_expr::expr_rewriter::{normalize_col, normalize_cols};
use datafusion_expr::select_expr::SelectExpr;
+use datafusion_expr::utils::{columnize_expr, find_aggregate_exprs};
use datafusion_expr::{
ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case,
dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION,
@@ -249,6 +252,113 @@ pub struct DataFrame {
projection_requires_validation: bool,
}
+/// True when select() should insert a global Aggregate before projecting.
+///
+/// Requires a pure aggregate list (no wildcards or input columns outside
+/// aggregate calls). Aggregates that already columnize onto plan (for
+/// example after [`DataFrame::aggregate`]) stay on the projection path.
+fn should_apply_global_aggregate(
+ expr_list: &[SelectExpr],
+ plan: &LogicalPlan,
+) -> Result<bool> {
+ let mut remaining_aggs = false;
+ for expr in expr_list {
+ match expr {
+ SelectExpr::Wildcard(_) | SelectExpr::QualifiedWildcard(_, _) => {
+ return Ok(false);
+ }
+ SelectExpr::Expression(expr) => {
+ if has_column_outside_aggregate(expr) {
+ return Ok(false);
+ }
+ let columnized =
+ columnize_expr(normalize_col(expr.clone(), plan)?, plan)?;
+ if !find_aggregate_exprs([&columnized]).is_empty() {
+ remaining_aggs = true;
+ }
+ }
+ }
+ }
+ // False when every aggregate already columnizes onto plan.
+ Ok(remaining_aggs)
+}
+
+/// Alias each aggregate uniquely so the Aggregate schema has distinct names.
+///
+/// The aliases are internal: the projection built on top of the Aggregate
+/// re-aliases its outputs, so user-facing names are unaffected. Identical
+/// aggregates are already deduplicated by find_aggregate_exprs; the suffix
+/// loop only guards against distinct aggregates whose names collide.
+///
+/// Returns the aliased aggregates and a map from the normalized original
+/// expression to the output column.
+fn alias_aggregate_exprs(
+ exprs: Vec<Expr>,
+ plan: &LogicalPlan,
+) -> Result<(Vec<Expr>, HashMap<Expr, Expr>)> {
+ let exprs = normalize_cols(exprs, plan)?;
+ let mut used = HashSet::new();
+ let mut rewrite_map = HashMap::new();
+ let mut aliased = Vec::with_capacity(exprs.len());
+ for expr in exprs {
+ if rewrite_map.contains_key(&expr) {
+ continue;
+ }
+ let base = expr.name_for_alias()?;
+ let mut name = base.clone();
+ let mut i = 1;
+ while !used.insert(name.clone()) {
+ name = format!("{base}_{i}");
+ i += 1;
+ }
+ rewrite_map.insert(expr.clone(), Expr::Column(Column::from_name(&name)));
+ aliased.push(expr.alias(name));
+ }
+ Ok((aliased, rewrite_map))
+}
+
+/// Replace aggregate subexpressions with columns produced by the Aggregate node.
+fn rewrite_select_aggs(
+ expr_list: Vec<SelectExpr>,
+ input: &LogicalPlan,
+ rewrite_map: &HashMap<Expr, Expr>,
+) -> Result<Vec<SelectExpr>> {
+ expr_list
+ .into_iter()
+ .map(|select_expr| match select_expr {
+ SelectExpr::Expression(expr) => {
+ let expr = normalize_col(expr, input)?;
+ let expr = expr
+ .transform(|node| {
+ Ok(match rewrite_map.get(&node) {
+ Some(col) => Transformed::yes(col.clone()),
+ None => Transformed::no(node),
+ })
+ })?
+ .data;
+ Ok(SelectExpr::Expression(expr))
+ }
+ other => Ok(other),
+ })
+ .collect()
+}
+
+/// True if expr references an input column that is not inside an aggregate.
+fn has_column_outside_aggregate(expr: &Expr) -> bool {
+ let mut found = false;
+ let _ = expr.apply(|e| {
+ if matches!(e, Expr::AggregateFunction(_)) {
+ return Ok(TreeNodeRecursion::Jump);
+ }
+ if matches!(e, Expr::Column(_)) {
+ found = true;
+ return Ok(TreeNodeRecursion::Stop);
+ }
+ Ok(TreeNodeRecursion::Continue)
+ });
+ found
+}
+
impl DataFrame {
/// Create a new `DataFrame ` based on an existing `LogicalPlan`
///
@@ -380,15 +490,23 @@ 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 `DataFrame` has one column for each element in `expr_list`.
+ /// The output 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`]. Mixing aggregates with non-aggregated columns is
+ /// not treated as an implicit GROUP BY.
+ ///
+ /// Output names must be unique, as with any projection; duplicate names
+ /// are rejected rather than renamed.
///
/// # 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<()> {
@@ -396,7 +514,9 @@ impl DataFrame {
/// 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 |",
@@ -404,7 +524,21 @@ impl DataFrame {
/// "| 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(())
/// # }
/// ```
@@ -412,21 +546,47 @@ impl DataFrame {
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.
+ // Aggregates get unique internal aliases here; the user-facing
+ // output names are validated by project() below, like any
+ // other projection.
+ 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)?;
+ }
+
+ // 3. Project — columnize_expr maps window/aggregate exprs to columns
+ // and duplicate output names are rejected
let project_plan = LogicalPlanBuilder::from(plan).project(expr_list)?.build()?;
Ok(DataFrame {
diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs
index 1a4f6ffcb4..4bbf6e2bcc 100644
--- a/datafusion/core/tests/dataframe/mod.rs
+++ b/datafusion/core/tests/dataframe/mod.rs
@@ -35,6 +35,7 @@ use arrow::util::pretty::pretty_format_batches;
use arrow_schema::{SortOptions, TimeUnit};
use datafusion::{assert_batches_eq, dataframe};
use datafusion_common::metadata::FieldMetadata;
+use datafusion_expr::select_expr::SelectExpr;
use datafusion_functions_aggregate::count::{count_all, count_all_window};
use datafusion_functions_aggregate::expr_fn::{
array_agg, avg, avg_distinct, count, count_distinct, max, median, min, sum,
@@ -74,7 +75,9 @@ use datafusion_common_runtime::SpawnedTask;
use datafusion_datasource::file_format::format_as_file_type;
use datafusion_execution::config::SessionConfig;
use datafusion_execution::runtime_env::RuntimeEnv;
-use datafusion_expr::expr::{GroupingSet, NullTreatment, Sort, WindowFunction};
+use datafusion_expr::expr::{
+ GroupingSet, NullTreatment, Sort, WildcardOptions, WindowFunction,
+};
use datafusion_expr::var_provider::{VarProvider, VarType};
use datafusion_expr::{
CreateMemoryTable, CreateView, DdlStatement, Expr, ExprFunctionExt, ExprSchemable,
@@ -7998,3 +8001,492 @@ async fn test_unresolved_lambda_variable() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_dataframe_api_select_semantics() -> Result<()> {
+ let df = test_table().await?;
+
+ // ----------------------------------------------------------------------
+ // Aggregate functions in SELECT (no GROUP BY)
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT
+ // COUNT(c9) AS count_c9,
+ // COUNT(CAST(c9 AS Utf8View)) AS count_c9_str,
+ // SUM(c9) AS sum_c9,
+ // COUNT(c8) AS count_c8,
+ // SUM(c9) + COUNT(c8) AS total1,
+ // (COUNT(c9) + 1) * 2 AS total2,
+ // COUNT(c9) + 1 AS count_c9_add_1
+ // FROM t
+ let res = df.clone().select(vec![
+ count(col("c9")).alias("count_c9"),
+ count(cast(col("c9"), DataType::Utf8View)).alias("count_c9_str"),
+ sum(col("c9")).alias("sum_c9"),
+ count(col("c8")).alias("count_c8"),
+ (sum(col("c9")) + count(col("c8"))).alias("total1"),
+ ((count(col("c9")) + lit(1)) * lit(2)).alias("total2"),
+ (count(col("c9")) + lit(1)).alias("count_c9_add_1"),
+ ])?;
+
+ assert_batches_eq!(
+ &[
+ "+----------+--------------+--------------+----------+--------------+--------+----------------+",
+ "| count_c9 | count_c9_str | sum_c9 | count_c8 | total1 | total2 | count_c9_add_1 |",
+ "+----------+--------------+--------------+----------+--------------+--------+----------------+",
+ "| 100 | 100 | 222089770060 | 100 | 222089770160 | 202 | 101 |",
+ "+----------+--------------+--------------+----------+--------------+--------+----------------+",
+ ],
+ &res.collect().await?
+ );
+
+ // ----------------------------------------------------------------------
+ // Aggregate + literal (still one row)
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT SUM(c9) AS sum_c9, 1 AS one FROM t
+ let res = df
+ .clone()
+ .select(vec![sum(col("c9")).alias("sum_c9"), lit(1).alias("one")])?;
+
+ assert_batches_eq!(
+ &[
+ "+--------------+-----+",
+ "| sum_c9 | one |",
+ "+--------------+-----+",
+ "| 222089770060 | 1 |",
+ "+--------------+-----+",
+ ],
+ &res.collect().await?
+ );
+
+ // ----------------------------------------------------------------------
+ // Empty input (global aggregate still returns exactly one row)
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT COUNT(c9) AS cnt, SUM(c9) AS sum_c9, 1 AS one FROM t WHERE false
+ let res = df.clone().filter(lit(false))?.select(vec![
+ count(col("c9")).alias("cnt"),
+ sum(col("c9")).alias("sum_c9"),
+ lit(1).alias("one"),
+ ])?;
+
+ assert_batches_eq!(
+ &[
+ "+-----+--------+-----+",
+ "| cnt | sum_c9 | one |",
+ "+-----+--------+-----+",
+ "| 0 | | 1 |",
+ "+-----+--------+-----+",
+ ],
+ &res.collect().await?
+ );
+
+ // ----------------------------------------------------------------------
+ // Duplicate output names are rejected, like any other projection
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT COUNT(c9) AS count_c9, COUNT(c9) AS count_c9 FROM t
+ let err = df
+ .clone()
+ .select(vec![
+ count(col("c9")).alias("count_c9"),
+ count(col("c9")).alias("count_c9"),
+ ])
+ .unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("require unique expression names"),
+ "duplicate aliases must be rejected at planning time, got: {msg}"
+ );
+
+ // SQL:
+ // SELECT COUNT(c9), COUNT(c9) FROM t
+ let err = df
+ .clone()
+ .select(vec![count(col("c9")), count(col("c9"))])
+ .unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("require unique expression names"),
+ "duplicate unaliased aggregates must be rejected at planning time, got: {msg}"
+ );
+
+ // ----------------------------------------------------------------------
+ // Mixed column + aggregate (no implicit GROUP BY)
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT c1, SUM(c9) AS sum_c9 FROM t
+ // Must not become GROUP BY c1. Planning succeeds; execution still fails.
+ let res = df
+ .clone()
+ .select(vec![col("c1"), sum(col("c9")).alias("sum_c9")])?;
+ let err = res.collect().await.unwrap_err();
+ let msg = err.to_string();
+ assert!(
+ msg.contains("not supported in this position"),
+ "mixed list must not run as grouped aggregation, got: {msg}"
+ );
+
+ // ----------------------------------------------------------------------
+ // Wildcard handling
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT *, 42 FROM t
+ let res = df
+ .clone()
+ .select(vec![
+ SelectExpr::Wildcard(WildcardOptions::default()),
+ lit(42).into(),
+ ])?
+ .limit(0, None)?;
+
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 14));
+
+ // SQL:
+ // SELECT aggregate_test_100.*, 42 FROM t
+ let res = df.clone().select(vec![
+ SelectExpr::QualifiedWildcard(
+ "aggregate_test_100".into(),
+ WildcardOptions::default(),
+ ),
+ lit(42).into(),
+ ])?;
+
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 14));
+
+ // ----------------------------------------------------------------------
+ // Window functions
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT
+ // c1,
+ // COUNT(c9) OVER (PARTITION BY c1) AS cnt,
+ // SUM(c9) OVER (PARTITION BY c1) AS sum_c9,
+ // AVG(c9) OVER (PARTITION BY c1) AS avg_c9
+ // FROM t
+ // ORDER BY c1
+ let count_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::count::count_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let sum_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::sum::sum_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let avg_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::average::avg_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let res = df
+ .clone()
+ .select(vec![
+ col("c1"),
+ count_window_function.alias("cnt"),
+ sum_window_function.alias("sum_c9"),
+ avg_window_function.alias("avg_c9"),
+ ])?
+ .sort(vec![col("c1").sort(true, true)])?;
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 4));
+ let batch = &batches[0];
+ assert_batches_eq!(
+ &[
+ "+----+-----+-------------+--------------------+",
+ "| c1 | cnt | sum_c9 | avg_c9 |",
+ "+----+-----+-------------+--------------------+",
+ "| a | 21 | 42619217323 | 2029486539.1904762 |",
+ "| b | 19 | 42365566310 | 2229766647.894737 |",
+ "| c | 21 | 46381998762 | 2208666607.714286 |",
+ "| d | 18 | 39910269981 | 2217237221.1666665 |",
+ "| e | 21 | 50812717684 | 2419653223.047619 |",
+ "+----+-----+-------------+--------------------+",
+ ],
+ &[
+ batch.slice(0, 1), // a
+ batch.slice(21, 1), // b
+ batch.slice(40, 1), // c
+ batch.slice(61, 1), // d
+ batch.slice(79, 1) // e
+ ]
+ );
+
+ // ----------------------------------------------------------------------
+ // Window with ORDER BY
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT
+ // c1,
+ // ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c9) AS rn
+ // FROM t
+ // ORDER BY c1
+ let row_number_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::WindowUDF(
+ datafusion::functions_window::row_number::row_number_udwf(),
+ ),
+ vec![],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![col("c9").sort(true, true)])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let res = df
+ .clone()
+ .select(vec![col("c1"), row_number_window_function.alias("rn")])?
+ .sort(vec![col("c1").sort(true, true)])?;
+
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 2));
+ assert_batches_eq!(
+ &[
+ "+----+----+",
+ "| c1 | rn |",
+ "+----+----+",
+ "| a | 1 |",
+ "| a | 2 |",
+ "| a | 3 |",
+ "| a | 4 |",
+ "| a | 5 |",
+ "| a | 6 |",
+ "| a | 7 |",
+ "| a | 8 |",
+ "| a | 9 |",
+ "| a | 10 |",
+ "| a | 11 |",
+ "| a | 12 |",
+ "| a | 13 |",
+ "| a | 14 |",
+ "| a | 15 |",
+ "| a | 16 |",
+ "| a | 17 |",
+ "| a | 18 |",
+ "| a | 19 |",
+ "| a | 20 |",
+ "| a | 21 |",
+ "| b | 1 |",
+ "| b | 2 |",
+ "| b | 3 |",
+ "| b | 4 |",
+ "| b | 5 |",
+ "| b | 6 |",
+ "| b | 7 |",
+ "| b | 8 |",
+ "| b | 9 |",
+ "| b | 10 |",
+ "| b | 11 |",
+ "| b | 12 |",
+ "| b | 13 |",
+ "| b | 14 |",
+ "| b | 15 |",
+ "| b | 16 |",
+ "| b | 17 |",
+ "| b | 18 |",
+ "| b | 19 |",
+ "| c | 1 |",
+ "| c | 2 |",
+ "| c | 3 |",
+ "| c | 4 |",
+ "| c | 5 |",
+ "| c | 6 |",
+ "| c | 7 |",
+ "| c | 8 |",
+ "| c | 9 |",
+ "| c | 10 |",
+ "| c | 11 |",
+ "| c | 12 |",
+ "| c | 13 |",
+ "| c | 14 |",
+ "| c | 15 |",
+ "| c | 16 |",
+ "| c | 17 |",
+ "| c | 18 |",
+ "| c | 19 |",
+ "| c | 20 |",
+ "| c | 21 |",
+ "| d | 1 |",
+ "| d | 2 |",
+ "| d | 3 |",
+ "| d | 4 |",
+ "| d | 5 |",
+ "| d | 6 |",
+ "| d | 7 |",
+ "| d | 8 |",
+ "| d | 9 |",
+ "| d | 10 |",
+ "| d | 11 |",
+ "| d | 12 |",
+ "| d | 13 |",
+ "| d | 14 |",
+ "| d | 15 |",
+ "| d | 16 |",
+ "| d | 17 |",
+ "| d | 18 |",
+ "| e | 1 |",
+ "| e | 2 |",
+ "| e | 3 |",
+ "| e | 4 |",
+ "| e | 5 |",
+ "| e | 6 |",
+ "| e | 7 |",
+ "| e | 8 |",
+ "| e | 9 |",
+ "| e | 10 |",
+ "| e | 11 |",
+ "| e | 12 |",
+ "| e | 13 |",
+ "| e | 14 |",
+ "| e | 15 |",
+ "| e | 16 |",
+ "| e | 17 |",
+ "| e | 18 |",
+ "| e | 19 |",
+ "| e | 20 |",
+ "| e | 21 |",
+ "+----+----+",
+ ],
+ &batches
+ );
+
+ // ----------------------------------------------------------------------
+ // Window inside expression
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT
+ // c1,
+ // COUNT(c9) OVER (PARTITION BY c1) + 1 AS cnt_plus
+ // FROM t
+ // ORDER BY c1
+ let count_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::count::count_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let cnt_plus_expr = count_window_function + lit(1);
+
+ let res = df
+ .clone()
+ .select(vec![col("c1"), cnt_plus_expr.alias("cnt_plus")])?
+ .sort(vec![col("c1").sort(true, true)])?;
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 2));
+ let batch = &batches[0];
+ assert_batches_eq!(
+ &[
+ "+----+----------+",
+ "| c1 | cnt_plus |",
+ "+----+----------+",
+ "| a | 22 |",
+ "| b | 20 |",
+ "| c | 22 |",
+ "| d | 19 |",
+ "| e | 22 |",
+ "+----+----------+",
+ ],
+ &[
+ batch.slice(0, 1), // a
+ batch.slice(21, 1), // b
+ batch.slice(40, 1), // c
+ batch.slice(61, 1), // d
+ batch.slice(79, 1) // e
+ ]
+ );
+
+ // ----------------------------------------------------------------------
+ // Window functions with mixed frames
+ // ----------------------------------------------------------------------
+ // SQL:
+ // SELECT
+ // c1,
+ // SUM(c9) OVER () AS total_sum,
+ // COUNT(c9) OVER (PARTITION BY c1) AS cnt
+ // FROM t
+ let sum_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::sum::sum_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let count_window_function = Expr::WindowFunction(Box::new(WindowFunction::new(
+ WindowFunctionDefinition::AggregateUDF(
+ datafusion_functions_aggregate::count::count_udaf(),
+ ),
+ vec![col("c9")],
+ )))
+ .partition_by(vec![col("c1")])
+ .order_by(vec![])
+ .window_frame(WindowFrame::new(None))
+ .build()?;
+
+ let res = df
+ .clone()
+ .select(vec![
+ col("c1"),
+ sum_window_function.alias("total_sum"),
+ count_window_function.alias("cnt"),
+ ])?
+ .sort(vec![col("c1").sort(true, true)])?;
+
+ let batches = res.collect().await?;
+ assert!(!batches.is_empty());
+ assert_eq!(batches.iter().map(|b| b.num_rows()).sum::<usize>(), 100);
+ assert!(batches.iter().all(|b| b.num_columns() == 3));
+ let batch = &batches[0];
+
+ assert_batches_eq!(
+ &[
+ "+----+--------------+-----+",
+ "| c1 | total_sum | cnt |",
+ "+----+--------------+-----+",
+ "| a | 222089770060 | 21 |",
+ "+----+--------------+-----+"
+ ],
+ &[batch.slice(0, 1)]
+ );
+
+ Ok(())
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org