diff --git a/Cargo.lock b/Cargo.lock index f00c931f15032..5b47b5dd094db 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2366,6 +2366,7 @@ dependencies = [ "datafusion-common", "datafusion-expr", "datafusion-expr-common", + "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-window", "datafusion-functions-window-common", diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index 81f573fc2a23e..375ed89597086 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -1390,6 +1390,14 @@ config_namespace! { /// cause regressions in both memory usage and runtime. pub enable_window_topn: bool, default = false + /// When set to true, coalesce peer `first_value` / `last_value` + /// expressions that share the same `ORDER BY` key — e.g. + /// `first_value(a ORDER BY o), first_value(b ORDER BY o), ... GROUP BY p` + /// — into a single `first_value(named_struct(a, b, ...) ORDER BY o)`, + /// cutting N argmax scans to one pass over the input. + /// Particularly beneficial for wide payloads. + pub enable_coalesce_first_last: bool, default = false + /// When set to true, the optimizer will push TopK (Sort with fetch) /// below hash repartition when the partition key is a prefix of the /// sort key, reducing data volume before the shuffle. diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md index 7652c2dcae984..40ddb681bdd90 100644 --- a/datafusion/core/src/optimizer_rule_reference.md +++ b/datafusion/core/src/optimizer_rule_reference.md @@ -58,10 +58,11 @@ Rule order matters. The default pipeline may change between releases. | 19 | `push_down_filter` | Moves filters as early as possible through filter-commutative operators. | | 20 | `single_distinct_aggregation_to_group_by` | Rewrites single-column `DISTINCT` aggregations into two-stage `GROUP BY` plans. | | 21 | `eliminate_group_by_constant` | Removes constant or functionally redundant expressions from `GROUP BY`. | -| 22 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. | -| 23 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. | -| 24 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. | -| 25 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. | +| 22 | `coalesce_first_last` | Coalesces peer `first_value` / `last_value` aggregates that share an `ORDER BY` into a single struct aggregate. | +| 23 | `common_sub_expression_eliminate` | Computes repeated subexpressions once and reuses the result. | +| 24 | `extract_leaf_expressions` | Pulls cheap leaf expressions closer to data sources so later pruning and filter rules can act earlier. | +| 25 | `push_down_leaf_projections` | Pushes the helper projections created by leaf extraction toward leaf inputs. | +| 26 | `optimize_projections` | Prunes unused columns and removes unnecessary logical projections. | ### Physical Optimizer Rules diff --git a/datafusion/optimizer/Cargo.toml b/datafusion/optimizer/Cargo.toml index 0822a17f24f16..fa705c7a9f4c1 100644 --- a/datafusion/optimizer/Cargo.toml +++ b/datafusion/optimizer/Cargo.toml @@ -69,6 +69,7 @@ regex-syntax = "0.8.9" async-trait = { workspace = true } criterion = { workspace = true } ctor = { workspace = true } +datafusion-functions = { workspace = true } datafusion-functions-aggregate = { workspace = true } datafusion-functions-window = { workspace = true } datafusion-functions-window-common = { workspace = true } diff --git a/datafusion/optimizer/src/coalesce_first_last.rs b/datafusion/optimizer/src/coalesce_first_last.rs new file mode 100644 index 0000000000000..c1168984dbe21 --- /dev/null +++ b/datafusion/optimizer/src/coalesce_first_last.rs @@ -0,0 +1,517 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`CoalesceFirstLast`] coalesces peer `first_value` / `last_value` aggregate +//! expressions that share the same `ORDER BY` key into a single struct-valued +//! aggregate. + +use std::collections::HashMap; +use std::sync::Arc; + +use crate::optimizer::ApplyOrder; +use crate::{OptimizerConfig, OptimizerRule}; + +use datafusion_common::Result; +use datafusion_common::tree_node::Transformed; +use datafusion_expr::expr::{ + AggregateFunction, AggregateFunctionParams, NullTreatment, Sort, +}; +use datafusion_expr::{ + Aggregate, AggregateUDF, Expr, LogicalPlan, LogicalPlanBuilder, col, lit, +}; + +use indexmap::IndexMap; + +const FIRST_VALUE: &str = "first_value"; +const LAST_VALUE: &str = "last_value"; +const NAMED_STRUCT: &str = "named_struct"; +const GET_FIELD: &str = "get_field"; + +/// Coalesces peer `first_value` / `last_value` aggregates that share one +/// `ORDER BY` key into a single struct-valued aggregate, with a projection on +/// top to unpack the struct back into the original columns: +/// +/// ```text +/// Aggregate: groupBy=[[p]], aggr=[[first_value(a ORDER BY o DESC), +/// first_value(b ORDER BY o DESC)]] +/// ``` +/// +/// becomes +/// +/// ```text +/// Projection: p, get_field(wrapped, 'c0'), get_field(wrapped, 'c1') +/// Aggregate: groupBy=[[p]], +/// aggr=[[first_value(named_struct('c0', a, 'c1', b) ORDER BY o DESC) AS wrapped]] +/// ``` +/// +/// The input is scanned once, not once per expression, and holds one per-group +/// state slot instead of N. +/// +/// Off by default (`optimizer.enable_coalesce_first_last`); no-ops if +/// `named_struct` / `get_field` are not registered. +#[derive(Default, Debug)] +pub struct CoalesceFirstLast {} + +impl CoalesceFirstLast { + pub fn new() -> Self { + Self {} + } +} + +/// `(function name, ORDER BY key, null treatment)` — peers may be coalesced only +/// when all three match. +type BucketKey = (String, Vec, Option); + +struct Coalesceable { + key: BucketKey, + func: Arc, + value: Expr, +} + +fn classify(expr: &Expr) -> Option { + let Expr::AggregateFunction(AggregateFunction { func, params }) = expr else { + return None; + }; + let name = func.name(); + if name != FIRST_VALUE && name != LAST_VALUE { + return None; + } + let AggregateFunctionParams { + args, + distinct, + filter, + order_by, + null_treatment, + } = params; + + // DISTINCT / FILTER would break the shared scan, and IGNORE NULLS cannot be + // reproduced through a single struct (it would skip rows where the whole + // struct is null rather than where the individual value is null). + if *distinct + || filter.is_some() + || args.len() != 1 + || order_by.is_empty() + || *null_treatment == Some(NullTreatment::IgnoreNulls) + { + return None; + } + + Some(Coalesceable { + key: (name.to_string(), order_by.clone(), *null_treatment), + func: Arc::clone(func), + value: args[0].clone(), + }) +} + +impl OptimizerRule for CoalesceFirstLast { + fn name(&self) -> &str { + "coalesce_first_last" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::BottomUp) + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result> { + if !config.options().optimizer.enable_coalesce_first_last { + return Ok(Transformed::no(plan)); + } + + let LogicalPlan::Aggregate(aggregate) = plan else { + return Ok(Transformed::no(plan)); + }; + + // Grouping sets expand the group columns (plus an internal grouping id) + // beyond `group_expr.len()`, which the schema slicing below assumes; skip + // them (consistent with the other cases this rule declines to coalesce). + if matches!(aggregate.group_expr.as_slice(), [Expr::GroupingSet(_)]) { + return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate))); + } + + let classified: Vec> = + aggregate.aggr_expr.iter().map(classify).collect(); + + let mut buckets: IndexMap> = IndexMap::new(); + for (i, c) in classified.iter().enumerate() { + if let Some(c) = c { + buckets.entry(c.key.clone()).or_default().push(i); + } + } + buckets.retain(|_, idxs| idxs.len() >= 2); + if buckets.is_empty() { + return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate))); + } + + // No-op rather than fail if the required scalar functions are missing + // (e.g. registry not wired, or a reduced function library). + let Some(registry) = config.function_registry() else { + return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate))); + }; + let (Ok(named_struct), Ok(get_field)) = + (registry.udf(NAMED_STRUCT), registry.udf(GET_FIELD)) + else { + return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate))); + }; + + let group_len = aggregate.group_expr.len(); + let orig_columns = aggregate.schema.columns(); + + // Original aggregate index -> (struct alias, struct field position). + let mut coalesced_at: HashMap = HashMap::new(); + let mut coalesced_aggr_exprs: Vec = Vec::with_capacity(buckets.len()); + + for (_key, idxs) in &buckets { + let alias = config.alias_generator().next("__coalesce_first_last"); + + let mut struct_args = Vec::with_capacity(idxs.len() * 2); + for (pos, &i) in idxs.iter().enumerate() { + let c = classified[i].as_ref().expect("bucket member is classified"); + struct_args.push(lit(format!("c{pos}"))); + struct_args.push(c.value.clone()); + coalesced_at.insert(i, (alias.clone(), pos)); + } + + let first = classified[idxs[0]].as_ref().expect("non-empty bucket"); + let (_, order_by, null_treatment) = &first.key; + let coalesced = Expr::AggregateFunction(AggregateFunction::new_udf( + Arc::clone(&first.func), + vec![named_struct.call(struct_args)], + /* distinct */ false, + /* filter */ None, + order_by.clone(), + *null_treatment, + )) + .alias(alias); + coalesced_aggr_exprs.push(coalesced); + } + + // Untouched aggregates keep their original order, then the coalesced ones. + let mut new_aggr_exprs: Vec = aggregate + .aggr_expr + .iter() + .enumerate() + .filter(|(i, _)| !coalesced_at.contains_key(i)) + .map(|(_, e)| e.clone()) + .collect(); + new_aggr_exprs.extend(coalesced_aggr_exprs); + + let new_aggregate = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::clone(&aggregate.input), + aggregate.group_expr.clone(), + new_aggr_exprs, + )?); + + // Rebuild the original output schema column-for-column: group columns, + // then each aggregate column, unpacking the struct where coalesced. + let mut projection_exprs: Vec = Vec::with_capacity(orig_columns.len()); + for column in &orig_columns[..group_len] { + projection_exprs.push(Expr::Column(column.clone())); + } + for (i, column) in orig_columns[group_len..].iter().enumerate() { + if let Some((alias, pos)) = coalesced_at.get(&i) { + let field = + get_field.call(vec![col(alias.as_str()), lit(format!("c{pos}"))]); + projection_exprs.push(field.alias(column.name().to_string())); + } else { + projection_exprs.push(Expr::Column(column.clone())); + } + } + + let projection = LogicalPlanBuilder::from(new_aggregate) + .project(projection_exprs)? + .build()?; + + Ok(Transformed::yes(projection)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::OptimizerContext; + use crate::assert_optimized_plan_eq_snapshot; + + use arrow::datatypes::{DataType, Field, Schema}; + use chrono::{DateTime, Utc}; + use datafusion_common::alias::AliasGenerator; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::expr::GroupingSet; + use datafusion_expr::logical_plan::table_scan; + use datafusion_expr::registry::{FunctionRegistry, MemoryFunctionRegistry}; + use datafusion_functions_aggregate::expr_fn::count; + use datafusion_functions_aggregate::first_last::{first_value_udaf, last_value_udaf}; + + /// An [`OptimizerConfig`] that enables the rule and exposes a registry with + /// `named_struct` / `get_field`, so the rewrite can be exercised in isolation. + #[derive(Debug)] + struct TestConfig { + options: Arc, + alias_generator: Arc, + registry: MemoryFunctionRegistry, + has_registry: bool, + } + + impl TestConfig { + fn new() -> Self { + Self::build(true) + } + + fn without_registry() -> Self { + Self::build(false) + } + + fn build(has_registry: bool) -> Self { + let mut registry = MemoryFunctionRegistry::new(); + datafusion_functions::register_all(&mut registry).unwrap(); + let mut options = ConfigOptions::default(); + options.optimizer.enable_coalesce_first_last = true; + Self { + options: Arc::new(options), + alias_generator: Arc::new(AliasGenerator::new()), + registry, + has_registry, + } + } + } + + impl OptimizerConfig for TestConfig { + fn query_execution_start_time(&self) -> Option> { + None + } + + fn alias_generator(&self) -> &Arc { + &self.alias_generator + } + + fn options(&self) -> Arc { + Arc::clone(&self.options) + } + + fn function_registry(&self) -> Option<&dyn FunctionRegistry> { + self.has_registry + .then_some(&self.registry as &dyn FunctionRegistry) + } + } + + macro_rules! assert_coalesced { + ($config:expr, $plan:expr, @ $expected:literal $(,)?) => {{ + let rules: Vec> = + vec![Arc::new(CoalesceFirstLast::new())]; + assert_optimized_plan_eq_snapshot!($config, rules, $plan, @ $expected,) + }}; + } + + /// Scan with a partition key `p`, two value columns `a` / `b`, and an + /// ordering column `o`. + fn scan() -> Result { + let schema = Schema::new(vec![ + Field::new("p", DataType::UInt32, false), + Field::new("a", DataType::UInt32, true), + Field::new("b", DataType::UInt32, true), + Field::new("o", DataType::UInt32, true), + ]); + table_scan(Some("t"), &schema, None)?.build() + } + + fn order_by_o() -> Vec { + vec![Sort::new(col("o"), false, false)] + } + + fn first_value(arg: Expr, order_by: Vec, distinct: bool) -> Expr { + Expr::AggregateFunction(AggregateFunction::new_udf( + first_value_udaf(), + vec![arg], + distinct, + None, + order_by, + None, + )) + } + + fn last_value(arg: Expr, order_by: Vec) -> Expr { + Expr::AggregateFunction(AggregateFunction::new_udf( + last_value_udaf(), + vec![arg], + false, + None, + order_by, + None, + )) + } + + #[test] + fn coalesces_peer_first_values() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + first_value(col("a"), order_by_o(), false), + first_value(col("b"), order_by_o(), false), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ r#" + Projection: t.p, get_field(__coalesce_first_last_1, Utf8("c0")) AS first_value(t.a) ORDER BY [t.o DESC NULLS LAST], get_field(__coalesce_first_last_1, Utf8("c1")) AS first_value(t.b) ORDER BY [t.o DESC NULLS LAST] + Aggregate: groupBy=[[t.p]], aggr=[[first_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o DESC NULLS LAST] AS __coalesce_first_last_1]] + TableScan: t + "#) + } + + #[test] + fn preserves_non_coalesced_aggregate() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + count(col("a")), + first_value(col("a"), order_by_o(), false), + first_value(col("b"), order_by_o(), false), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ r#" + Projection: t.p, count(t.a), get_field(__coalesce_first_last_1, Utf8("c0")) AS first_value(t.a) ORDER BY [t.o DESC NULLS LAST], get_field(__coalesce_first_last_1, Utf8("c1")) AS first_value(t.b) ORDER BY [t.o DESC NULLS LAST] + Aggregate: groupBy=[[t.p]], aggr=[[count(t.a), first_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o DESC NULLS LAST] AS __coalesce_first_last_1]] + TableScan: t + "#) + } + + #[test] + fn coalesces_peer_last_values() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + last_value(col("a"), order_by_o()), + last_value(col("b"), order_by_o()), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ r#" + Projection: t.p, get_field(__coalesce_first_last_1, Utf8("c0")) AS last_value(t.a) ORDER BY [t.o DESC NULLS LAST], get_field(__coalesce_first_last_1, Utf8("c1")) AS last_value(t.b) ORDER BY [t.o DESC NULLS LAST] + Aggregate: groupBy=[[t.p]], aggr=[[last_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o DESC NULLS LAST] AS __coalesce_first_last_1]] + TableScan: t + "#) + } + + #[test] + fn no_op_single_first_value() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![first_value(col("a"), order_by_o(), false)], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ " + Aggregate: groupBy=[[t.p]], aggr=[[first_value(t.a) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } + + #[test] + fn no_op_mixed_first_and_last() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + first_value(col("a"), order_by_o(), false), + last_value(col("b"), order_by_o()), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ " + Aggregate: groupBy=[[t.p]], aggr=[[first_value(t.a) ORDER BY [t.o DESC NULLS LAST], last_value(t.b) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } + + #[test] + fn no_op_distinct() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + first_value(col("a"), order_by_o(), true), + first_value(col("b"), order_by_o(), true), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ " + Aggregate: groupBy=[[t.p]], aggr=[[first_value(DISTINCT t.a) ORDER BY [t.o DESC NULLS LAST], first_value(DISTINCT t.b) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } + + #[test] + fn no_op_when_disabled() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + first_value(col("a"), order_by_o(), false), + first_value(col("b"), order_by_o(), false), + ], + )? + .build()?; + let config = OptimizerContext::new(); + let rules: Vec> = + vec![Arc::new(CoalesceFirstLast::new())]; + assert_optimized_plan_eq_snapshot!(config, rules, plan, @ " + Aggregate: groupBy=[[t.p]], aggr=[[first_value(t.a) ORDER BY [t.o DESC NULLS LAST], first_value(t.b) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } + + #[test] + fn no_op_without_registry() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![col("p")], + vec![ + first_value(col("a"), order_by_o(), false), + first_value(col("b"), order_by_o(), false), + ], + )? + .build()?; + assert_coalesced!(TestConfig::without_registry(), plan, @ " + Aggregate: groupBy=[[t.p]], aggr=[[first_value(t.a) ORDER BY [t.o DESC NULLS LAST], first_value(t.b) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } + + #[test] + fn no_op_grouping_set() -> Result<()> { + let plan = LogicalPlanBuilder::from(scan()?) + .aggregate( + vec![Expr::GroupingSet(GroupingSet::Cube(vec![col("p")]))], + vec![ + first_value(col("a"), order_by_o(), false), + first_value(col("b"), order_by_o(), false), + ], + )? + .build()?; + assert_coalesced!(TestConfig::new(), plan, @ " + Aggregate: groupBy=[[CUBE (t.p)]], aggr=[[first_value(t.a) ORDER BY [t.o DESC NULLS LAST], first_value(t.b) ORDER BY [t.o DESC NULLS LAST]]] + TableScan: t + ") + } +} diff --git a/datafusion/optimizer/src/lib.rs b/datafusion/optimizer/src/lib.rs index fbe7ad2f4d327..416a6d8e4396e 100644 --- a/datafusion/optimizer/src/lib.rs +++ b/datafusion/optimizer/src/lib.rs @@ -39,6 +39,7 @@ //! [`LogicalPlan`]: datafusion_expr::LogicalPlan //! [`TypeCoercion`]: analyzer::type_coercion::TypeCoercion pub mod analyzer; +pub mod coalesce_first_last; pub mod common_subexpr_eliminate; pub mod decorrelate; pub mod decorrelate_lateral_join; diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index db7ad8475273a..53b7ef600623c 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -41,6 +41,7 @@ use datafusion_expr::{ Window, }; +use crate::coalesce_first_last::CoalesceFirstLast; use crate::common_subexpr_eliminate::CommonSubexprEliminate; use crate::decorrelate_lateral_join::DecorrelateLateralJoin; use crate::decorrelate_predicate_subquery::DecorrelatePredicateSubquery; @@ -312,6 +313,7 @@ impl Optimizer { // The previous optimizations added expressions and projections, // that might benefit from the following rules Arc::new(EliminateGroupByConstant::new()), + Arc::new(CoalesceFirstLast::new()), Arc::new(CommonSubexprEliminate::new()), Arc::new(ExtractLeafExpressions::new()), Arc::new(PushDownLeafProjections::new()), diff --git a/datafusion/sqllogictest/test_files/coalesce_first_last.slt b/datafusion/sqllogictest/test_files/coalesce_first_last.slt new file mode 100644 index 0000000000000..d07ea079c54a1 --- /dev/null +++ b/datafusion/sqllogictest/test_files/coalesce_first_last.slt @@ -0,0 +1,84 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Data with multiple groups, NULL value cells, and a unique ORDER BY key. +statement ok +CREATE TABLE t(p int, a int, b int, o int) AS VALUES +(1, 10, 100, 1), +(1, NULL, 200, 2), +(1, 30, NULL, 3), +(2, 40, 400, 4), +(2, 50, 500, 5), +(3, NULL, NULL, 1); + +# Baseline: rule disabled +statement ok +set datafusion.optimizer.enable_coalesce_first_last = false; + +query IIIII +SELECT p, + first_value(a ORDER BY o ASC) AS fa, + first_value(b ORDER BY o ASC) AS fb, + last_value(a ORDER BY o ASC) AS la, + last_value(b ORDER BY o ASC) AS lb +FROM t GROUP BY p ORDER BY p; +---- +1 10 100 30 NULL +2 40 400 50 500 +3 NULL NULL NULL NULL + +# Same query with the rule enabled: results must be identical to the baseline. +statement ok +set datafusion.optimizer.enable_coalesce_first_last = true; + +query IIIII +SELECT p, + first_value(a ORDER BY o ASC) AS fa, + first_value(b ORDER BY o ASC) AS fb, + last_value(a ORDER BY o ASC) AS la, + last_value(b ORDER BY o ASC) AS lb +FROM t GROUP BY p ORDER BY p; +---- +1 10 100 30 NULL +2 40 400 50 500 +3 NULL NULL NULL NULL + +# The coalesced plan scans once: a single struct-valued first_value feeds a +# projection that unpacks the fields with get_field. +query TT +EXPLAIN SELECT p, + first_value(a ORDER BY o ASC) AS fa, + first_value(b ORDER BY o ASC) AS fb +FROM t GROUP BY p; +---- +logical_plan +01)Projection: t.p, get_field(__coalesce_first_last_1, Utf8("c0")) AS fa, get_field(__coalesce_first_last_1, Utf8("c1")) AS fb +02)--Aggregate: groupBy=[[t.p]], aggr=[[first_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o ASC NULLS LAST] AS __coalesce_first_last_1]] +03)----TableScan: t projection=[p, a, b, o] +physical_plan +01)ProjectionExec: expr=[p@0 as p, get_field(__coalesce_first_last_1@1, c0) as fa, get_field(__coalesce_first_last_1@1, c1) as fb] +02)--AggregateExec: mode=FinalPartitioned, gby=[p@0 as p], aggr=[first_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o ASC NULLS LAST] as __coalesce_first_last_1] +03)----RepartitionExec: partitioning=Hash([p@0], 4), input_partitions=1 +04)------AggregateExec: mode=Partial, gby=[p@0 as p], aggr=[first_value(named_struct(Utf8("c0"), t.a, Utf8("c1"), t.b)) ORDER BY [t.o ASC NULLS LAST] as __coalesce_first_last_1] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Restore the default so the harness sees no leaked configuration. +statement ok +set datafusion.optimizer.enable_coalesce_first_last = false; + +statement ok +DROP TABLE t; diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 5405c7ce0e779..64e3eab5dace2 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -197,6 +197,7 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE +logical_plan after coalesce_first_last SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after extract_leaf_expressions SAME TEXT AS ABOVE logical_plan after push_down_leaf_projections SAME TEXT AS ABOVE @@ -222,6 +223,7 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE +logical_plan after coalesce_first_last SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after extract_leaf_expressions SAME TEXT AS ABOVE logical_plan after push_down_leaf_projections SAME TEXT AS ABOVE @@ -576,6 +578,7 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE +logical_plan after coalesce_first_last SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after extract_leaf_expressions SAME TEXT AS ABOVE logical_plan after push_down_leaf_projections SAME TEXT AS ABOVE @@ -601,6 +604,7 @@ logical_plan after push_down_limit SAME TEXT AS ABOVE logical_plan after push_down_filter SAME TEXT AS ABOVE logical_plan after single_distinct_aggregation_to_group_by SAME TEXT AS ABOVE logical_plan after eliminate_group_by_constant SAME TEXT AS ABOVE +logical_plan after coalesce_first_last SAME TEXT AS ABOVE logical_plan after common_sub_expression_eliminate SAME TEXT AS ABOVE logical_plan after extract_leaf_expressions SAME TEXT AS ABOVE logical_plan after push_down_leaf_projections SAME TEXT AS ABOVE diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 12306b4529c46..6b95583161da4 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -305,6 +305,7 @@ datafusion.format.types_info false datafusion.optimizer.allow_symmetric_joins_without_pruning true datafusion.optimizer.default_filter_selectivity 20 datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown true +datafusion.optimizer.enable_coalesce_first_last false datafusion.optimizer.enable_distinct_aggregation_soft_limit true datafusion.optimizer.enable_dynamic_filter_pushdown true datafusion.optimizer.enable_join_dynamic_filter_pushdown true @@ -464,6 +465,7 @@ datafusion.format.types_info false Show types in visual representation batches datafusion.optimizer.allow_symmetric_joins_without_pruning true Should DataFusion allow symmetric hash joins for unbounded data sources even when its inputs do not have any ordering or filtering If the flag is not enabled, the SymmetricHashJoin operator will be unable to prune its internal buffers, resulting in certain join types - such as Full, Left, LeftAnti, LeftSemi, Right, RightAnti, and RightSemi - being produced only at the end of the execution. This is not typical in stream processing. Additionally, without proper design for long runner execution, all types of joins may encounter out-of-memory errors. datafusion.optimizer.default_filter_selectivity 20 The default filter selectivity used by Filter Statistics when an exact selectivity cannot be determined. Valid values are between 0 (no selectivity) and 100 (all rows are selected). datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown true When set to true, the optimizer will attempt to push down Aggregate dynamic filters into the file scan phase. +datafusion.optimizer.enable_coalesce_first_last false When set to true, coalesce peer `first_value` / `last_value` expressions that share the same `ORDER BY` key — e.g. `first_value(a ORDER BY o), first_value(b ORDER BY o), ... GROUP BY p` — into a single `first_value(named_struct(a, b, ...) ORDER BY o)`, cutting N argmax scans to one pass over the input. Particularly beneficial for wide payloads. datafusion.optimizer.enable_distinct_aggregation_soft_limit true When set to true, the optimizer will push a limit operation into grouped aggregations which have no aggregate expressions, as a soft limit, emitting groups once the limit is reached, before all rows in the group are read. datafusion.optimizer.enable_dynamic_filter_pushdown true When set to true attempts to push down dynamic filters generated by operators (TopK, Join & Aggregate) into the file scan phase. For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans. This means that if we already have 10 timestamps in the year 2025 any files that only have timestamps in the year 2024 can be skipped / pruned at various stages in the scan. The config will suppress `enable_join_dynamic_filter_pushdown`, `enable_topk_dynamic_filter_pushdown` & `enable_aggregate_dynamic_filter_pushdown` So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden. datafusion.optimizer.enable_join_dynamic_filter_pushdown true When set to true, the optimizer will attempt to push down Join dynamic filters into the file scan phase. diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index f6e072b59bceb..0e8b11146b190 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -148,6 +148,7 @@ The following configuration settings are available: | datafusion.optimizer.enable_topk_aggregation | true | When set to true, the optimizer will attempt to perform limit operations during aggregations, if possible | | datafusion.optimizer.enable_window_limits | true | When set to true, the optimizer will attempt to push limit operations past window functions, if possible | | datafusion.optimizer.enable_window_topn | false | When set to true, the optimizer will replace Filter(rn<=K) → Window(ROW_NUMBER) → Sort patterns with a PartitionedTopKExec that maintains per-partition heaps, avoiding a full sort of the input. When the window partition key has low cardinality, enabling this optimization can improve performance. However, for high cardinality keys, it may cause regressions in both memory usage and runtime. | +| datafusion.optimizer.enable_coalesce_first_last | false | When set to true, coalesce peer `first_value` / `last_value` expressions that share the same `ORDER BY` key — e.g. `first_value(a ORDER BY o), first_value(b ORDER BY o), ... GROUP BY p` — into a single `first_value(named_struct(a, b, ...) ORDER BY o)`, cutting N argmax scans to one pass over the input. Particularly beneficial for wide payloads. | | datafusion.optimizer.enable_topk_repartition | true | When set to true, the optimizer will push TopK (Sort with fetch) below hash repartition when the partition key is a prefix of the sort key, reducing data volume before the shuffle. | | datafusion.optimizer.enable_topk_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down TopK dynamic filters into the file scan phase. | | datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery | true | When set to true, uncorrelated scalar subqueries are left in the logical plan and executed by `ScalarSubqueryExec` during physical execution. When set to false, all scalar subqueries (including uncorrelated ones) are rewritten to left joins by the `ScalarSubqueryToJoin` optimizer rule. Note disabling this option is not recommended. It restores pre behavior, which silently produces incorrect results for multi-row subqueries and does not support scalar subqueries in ORDER BY / JOIN ON / aggregate-function arguments. This option is intended as a temporary escape hatch for distributed execution frameworks and is planned to be removed in a future DataFusion release. |