zhuqi-lucas commented on code in PR #23824: URL: https://github.com/apache/datafusion/pull/23824#discussion_r3694914536
########## datafusion/optimizer/src/replace_filter_top1.rs: ########## @@ -0,0 +1,627 @@ +// 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. + +//! [`ReplaceFilterTop1`] rewrites a "top-1 per group" `row_number()` filter +//! into a `first_value` aggregate. + +use crate::optimizer::{ApplyOrder, ApplyOrder::BottomUp}; +use crate::{OptimizerConfig, OptimizerRule}; +use std::sync::Arc; + +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::Transformed; +use datafusion_common::{Column, Result}; +use datafusion_expr::{ + Aggregate, BinaryExpr, Expr, Filter, LogicalPlan, Operator, Projection, SortExpr, +}; +use datafusion_expr::{ExprFunctionExt, LogicalPlanBuilder, lit}; + +/// Optimizer that replaces a logical `Filter` with a top-1 predicate over a +/// `row_number()` window with an aggregate. +/// +/// ```text +/// SELECT * FROM ( +/// SELECT *, +/// ROW_NUMBER() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// FROM t +/// ) WHERE rn = 1 +/// ``` +/// +/// Input plan: +/// ```text +/// Filter: rn = 1 -- or rn <= 1, or rn < 2 +/// [Projection: ...] -- optional passthrough projection +/// WindowAggr: row_number() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// child +/// ``` +/// +/// Rewritten plan: +/// ```text +/// Projection: <original Filter output schema; rn -> literal 1> +/// Aggregate: +/// group_by=[p] +/// aggr=[first_value(col_i ORDER BY o DESC) for each non-partition col_i] +/// child +/// ``` +/// +/// Notes: +/// - the window function must be `row_number` +/// - filter predicate must be "top-1" (rn = 1, <= 1, < 2) +/// - window has a `PARTITION BY` clause +/// - gated behind the `optimizer.enable_row_number_to_aggregate` config option (off by default) +#[derive(Default, Debug)] +pub struct ReplaceFilterTop1 {} + +impl ReplaceFilterTop1 { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl OptimizerRule for ReplaceFilterTop1 { + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result<Transformed<LogicalPlan>> { + if !config.options().optimizer.enable_row_number_to_aggregate { + return Ok(Transformed::no(plan)); + } + + let LogicalPlan::Filter(Filter { + ref predicate, + ref input, + .. + }) = plan + else { + return Ok(Transformed::no(plan)); + }; + + let Some(WindowTop1 { + order_by, + partition_by, + rn_col, + input_cols, + child, + projection, + }) = validate_window_input(input) + else { + return Ok(Transformed::no(plan)); + }; + + // Resolve the rn name the filter references (projection alias if present) + let rn_ref_name = match projection { + None => rn_col.name.clone(), + Some(p) => match rn_passthrough_name(p, &rn_col) { + Some(name) => name, + None => return Ok(Transformed::no(plan)), + }, + }; + if !has_valid_predicate(predicate, &rn_ref_name) { + return Ok(Transformed::no(plan)); + } + + let is_partition = |c: &Column| { + partition_by.iter().any(|e| matches!(e, Expr::Column(p) if p.name == c.name && p.relation == c.relation)) + }; + + // Group by the partition keys and take `first_value(col ORDER BY ...)` for every other + // input column, aliased back to that column's qualifier+name + let Some(registry) = config.function_registry() else { + return Ok(Transformed::no(plan)); + }; + let first_value = registry.udaf("first_value")?; + let aggr_expr = input_cols + .iter() + .filter(|c| !is_partition(c)) + .map(|c| { + first_value + .call(vec![Expr::Column(c.clone())]) + .order_by(order_by.to_vec()) + .build() + .map(|e| e.alias_qualified(c.relation.clone(), c.name.clone())) + }) + .collect::<Result<Vec<_>>>()?; + + let aggregate = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::clone(child), + partition_by.to_vec(), + aggr_expr, + )?); + + // Restoring the Filter's exact output schema on top of the aggregate + let proj_exprs = match projection { + None => input + .schema() + .iter() + .map(|(qualifier, field)| { + if qualifier.is_none() && field.name() == &rn_col.name { + lit(1u64).alias(field.name()) + } else { + Expr::Column(Column::new(qualifier.cloned(), field.name())) + } + }) + .collect::<Vec<_>>(), + Some(p) => p + .expr + .iter() + .zip(p.schema.iter()) + .map(|(expr, (qualifier, field))| { + if matches!(unalias(expr), Expr::Column(c) if *c == rn_col) { + lit(1u64).alias_qualified(qualifier.cloned(), field.name()) + } else { + expr.clone() + } + }) + .collect::<Vec<_>>(), + }; + + let new_plan = LogicalPlanBuilder::from(aggregate) + .project(proj_exprs)? + .build()?; + Ok(Transformed::yes(new_plan)) + } + + fn name(&self) -> &str { + "replace_filter_top1" + } + + fn apply_order(&self) -> Option<ApplyOrder> { + Some(BottomUp) + } +} + +/// Validating that the filter predicate is `rn_name` == 1 (or equivalent) +fn has_valid_predicate(predicate: &Expr, rn_name: &str) -> bool { + let Expr::BinaryExpr(BinaryExpr { left, right, op }) = predicate else { + return false; + }; + + let (name, op, val) = match (&**left, &**right) { + ( + Expr::Column(Column { name, .. }), + Expr::Literal(ScalarValue::UInt64(Some(val)), _), + ) => (name, *op, *val), + ( + Expr::Literal(ScalarValue::UInt64(Some(val)), _), + Expr::Column(Column { name, .. }), + ) => { + let Some(op) = op.swap() else { return false }; + (name, op, *val) + } + _ => return false, + }; + + name.as_str() == rn_name Review Comment: **This name-only match can rewrite away a filter on a payload column.** Both `rn_passthrough_name` (returns the bare `field.name()`, dropping the qualifier) and this comparison ignore the column's relation, so a payload column that shares the rn alias's *name* under a different qualifier matches too. Verified reproducer (buildable with the existing test helpers in this file): ```rust // projection: [test.a, test.b, test.c, rn AS w.b] — rn aliased under qualifier "w" // filter: test.b = 1 — a *payload* filter, not the window rn let plan = top1_plan_with_projection( vec![col("a")], |rn| vec![col("test.a"), col("test.b"), col("test.c"), rn.alias_qualified(Some("w"), "b")], |_| col("test.b").eq(lit(1u64)), )?; ``` The rule fires and produces: ```text Projection: test.a, test.b, test.c, UInt64(1) AS b Aggregate: groupBy=[[test.a]], aggr=[[first_value(test.b) ORDER BY [...], ...]] ``` — the `test.b = 1` filter is silently discarded and replaced by top-1-per-partition semantics: **wrong results**. (The unqualified-alias variant of this shape is unreachable — the builder rejects it as ambiguous — but the qualified-alias form constructs fine, and optimizer rules themselves produce qualified aliases.) Fix is to carry the qualifier through both spots — I applied and verified this locally (all 11 existing tests + the reproducer pass): ```rust // 1. rn_passthrough_name → return the full output column fn rn_passthrough_column(projection: &Projection, rn_col: &Column) -> Option<Column> { let mut rn_output: Option<Column> = None; for (expr, (qualifier, field)) in projection.expr.iter().zip(projection.schema.iter()) { let is_passthrough = matches!(unalias(expr), Expr::Column(c) if c == rn_col); if is_passthrough { if rn_output.is_some() { return None; } rn_output = Some(Column::new(qualifier.cloned(), field.name().clone())); } else if expr.column_refs().iter().any(|c| *c == rn_col) { return None; } } rn_output } // 2. has_valid_predicate: take `rn_ref: &Column`, match `Expr::Column(c)` whole, // and compare `column == rn_ref` instead of name-only. // 3. call site: `None => rn_col.clone()` for the no-projection path. ``` Happy to share the reproducer as a test to keep in the suite. ########## datafusion/optimizer/src/replace_filter_top1.rs: ########## @@ -0,0 +1,627 @@ +// 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. + +//! [`ReplaceFilterTop1`] rewrites a "top-1 per group" `row_number()` filter +//! into a `first_value` aggregate. + +use crate::optimizer::{ApplyOrder, ApplyOrder::BottomUp}; +use crate::{OptimizerConfig, OptimizerRule}; +use std::sync::Arc; + +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::Transformed; +use datafusion_common::{Column, Result}; +use datafusion_expr::{ + Aggregate, BinaryExpr, Expr, Filter, LogicalPlan, Operator, Projection, SortExpr, +}; +use datafusion_expr::{ExprFunctionExt, LogicalPlanBuilder, lit}; + +/// Optimizer that replaces a logical `Filter` with a top-1 predicate over a +/// `row_number()` window with an aggregate. +/// +/// ```text +/// SELECT * FROM ( +/// SELECT *, +/// ROW_NUMBER() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// FROM t +/// ) WHERE rn = 1 +/// ``` +/// +/// Input plan: +/// ```text +/// Filter: rn = 1 -- or rn <= 1, or rn < 2 +/// [Projection: ...] -- optional passthrough projection +/// WindowAggr: row_number() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// child +/// ``` +/// +/// Rewritten plan: +/// ```text +/// Projection: <original Filter output schema; rn -> literal 1> +/// Aggregate: +/// group_by=[p] +/// aggr=[first_value(col_i ORDER BY o DESC) for each non-partition col_i] +/// child +/// ``` +/// +/// Notes: +/// - the window function must be `row_number` +/// - filter predicate must be "top-1" (rn = 1, <= 1, < 2) Review Comment: Worth a note here (same discussion as on #23682, this epic's sibling): with duplicate `ORDER BY` keys, `row_number()` and the rewritten `first_value` aggregate may pick **different (equally valid) tied rows** — each path breaks ties deterministically on its own, but the two don't guarantee the same choice as each other. Someone toggling the flag on tied data could observe a result change and report it as a bug; one sentence here preempts that. ########## datafusion/optimizer/src/replace_filter_top1.rs: ########## @@ -0,0 +1,627 @@ +// 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. + +//! [`ReplaceFilterTop1`] rewrites a "top-1 per group" `row_number()` filter +//! into a `first_value` aggregate. + +use crate::optimizer::{ApplyOrder, ApplyOrder::BottomUp}; +use crate::{OptimizerConfig, OptimizerRule}; +use std::sync::Arc; + +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::Transformed; +use datafusion_common::{Column, Result}; +use datafusion_expr::{ + Aggregate, BinaryExpr, Expr, Filter, LogicalPlan, Operator, Projection, SortExpr, +}; +use datafusion_expr::{ExprFunctionExt, LogicalPlanBuilder, lit}; + +/// Optimizer that replaces a logical `Filter` with a top-1 predicate over a +/// `row_number()` window with an aggregate. +/// +/// ```text +/// SELECT * FROM ( +/// SELECT *, +/// ROW_NUMBER() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// FROM t +/// ) WHERE rn = 1 +/// ``` +/// +/// Input plan: +/// ```text +/// Filter: rn = 1 -- or rn <= 1, or rn < 2 +/// [Projection: ...] -- optional passthrough projection +/// WindowAggr: row_number() OVER (PARTITION BY p ORDER BY o DESC) AS rn +/// child +/// ``` +/// +/// Rewritten plan: +/// ```text +/// Projection: <original Filter output schema; rn -> literal 1> +/// Aggregate: +/// group_by=[p] +/// aggr=[first_value(col_i ORDER BY o DESC) for each non-partition col_i] +/// child +/// ``` +/// +/// Notes: +/// - the window function must be `row_number` +/// - filter predicate must be "top-1" (rn = 1, <= 1, < 2) +/// - window has a `PARTITION BY` clause +/// - gated behind the `optimizer.enable_row_number_to_aggregate` config option (off by default) +#[derive(Default, Debug)] +pub struct ReplaceFilterTop1 {} + +impl ReplaceFilterTop1 { + #[expect(missing_docs)] + pub fn new() -> Self { + Self {} + } +} + +impl OptimizerRule for ReplaceFilterTop1 { + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + config: &dyn OptimizerConfig, + ) -> Result<Transformed<LogicalPlan>> { + if !config.options().optimizer.enable_row_number_to_aggregate { + return Ok(Transformed::no(plan)); + } + + let LogicalPlan::Filter(Filter { + ref predicate, + ref input, + .. + }) = plan + else { + return Ok(Transformed::no(plan)); + }; + + let Some(WindowTop1 { + order_by, + partition_by, + rn_col, + input_cols, + child, + projection, + }) = validate_window_input(input) + else { + return Ok(Transformed::no(plan)); + }; + + // Resolve the rn name the filter references (projection alias if present) + let rn_ref_name = match projection { + None => rn_col.name.clone(), + Some(p) => match rn_passthrough_name(p, &rn_col) { + Some(name) => name, + None => return Ok(Transformed::no(plan)), + }, + }; + if !has_valid_predicate(predicate, &rn_ref_name) { + return Ok(Transformed::no(plan)); + } + + let is_partition = |c: &Column| { + partition_by.iter().any(|e| matches!(e, Expr::Column(p) if p.name == c.name && p.relation == c.relation)) + }; + + // Group by the partition keys and take `first_value(col ORDER BY ...)` for every other + // input column, aliased back to that column's qualifier+name + let Some(registry) = config.function_registry() else { + return Ok(Transformed::no(plan)); + }; + let first_value = registry.udaf("first_value")?; + let aggr_expr = input_cols + .iter() + .filter(|c| !is_partition(c)) + .map(|c| { + first_value + .call(vec![Expr::Column(c.clone())]) + .order_by(order_by.to_vec()) + .build() + .map(|e| e.alias_qualified(c.relation.clone(), c.name.clone())) + }) + .collect::<Result<Vec<_>>>()?; + + let aggregate = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::clone(child), + partition_by.to_vec(), + aggr_expr, + )?); + + // Restoring the Filter's exact output schema on top of the aggregate + let proj_exprs = match projection { + None => input + .schema() + .iter() + .map(|(qualifier, field)| { + if qualifier.is_none() && field.name() == &rn_col.name { + lit(1u64).alias(field.name()) + } else { + Expr::Column(Column::new(qualifier.cloned(), field.name())) + } + }) + .collect::<Vec<_>>(), + Some(p) => p + .expr + .iter() + .zip(p.schema.iter()) + .map(|(expr, (qualifier, field))| { + if matches!(unalias(expr), Expr::Column(c) if *c == rn_col) { + lit(1u64).alias_qualified(qualifier.cloned(), field.name()) + } else { + expr.clone() + } + }) + .collect::<Vec<_>>(), + }; + + let new_plan = LogicalPlanBuilder::from(aggregate) + .project(proj_exprs)? + .build()?; + Ok(Transformed::yes(new_plan)) + } + + fn name(&self) -> &str { + "replace_filter_top1" + } + + fn apply_order(&self) -> Option<ApplyOrder> { + Some(BottomUp) + } +} + +/// Validating that the filter predicate is `rn_name` == 1 (or equivalent) +fn has_valid_predicate(predicate: &Expr, rn_name: &str) -> bool { + let Expr::BinaryExpr(BinaryExpr { left, right, op }) = predicate else { + return false; + }; + + let (name, op, val) = match (&**left, &**right) { + ( + Expr::Column(Column { name, .. }), + Expr::Literal(ScalarValue::UInt64(Some(val)), _), + ) => (name, *op, *val), + ( + Expr::Literal(ScalarValue::UInt64(Some(val)), _), + Expr::Column(Column { name, .. }), + ) => { + let Some(op) = op.swap() else { return false }; + (name, op, *val) + } + _ => return false, + }; + + name.as_str() == rn_name + && match op { + Operator::Lt => val == 2, + Operator::Eq | Operator::LtEq => val == 1, + _ => false, + } +} + +/// A `row_number()` window recognized beneath a `Filter`, decomposed into the +/// pieces the rewrite needs. Borrows from the matched [`LogicalPlan`]. +struct WindowTop1<'a> { + /// `ORDER BY` of the window (becomes the `first_value` ordering). + order_by: &'a [SortExpr], + /// `PARTITION BY` of the window (becomes the aggregate group-by). + partition_by: &'a [Expr], + /// The `row_number()` output column produced by the window. + rn_col: Column, + /// Columns produced by the window's input (candidate payload columns). + input_cols: Vec<Column>, + /// The window's input, which becomes the aggregate's input. + child: &'a Arc<LogicalPlan>, + /// Optional passthrough `Projection` sitting between the `Filter` and the + /// `Window`. When present, its expressions are re-applied on top of the + /// rewritten aggregate. + projection: Option<&'a Projection>, +} + +/// Recognize either `Filter -> Window` or `Filter -> Projection -> Window`, +/// where the window is a single `row_number()` with a non-empty `PARTITION BY`. +fn validate_window_input(input: &Arc<LogicalPlan>) -> Option<WindowTop1<'_>> { + let (projection, window) = match &**input { + LogicalPlan::Window(w) => (None, w), + LogicalPlan::Projection(p) => match &*p.input { + LogicalPlan::Window(w) => (Some(p), w), + _ => return None, + }, + _ => return None, + }; + + if window.window_expr.len() != 1 { + return None; + } + + let window_expr = window.window_expr.first()?; + let Expr::WindowFunction(window_function) = window_expr else { + return None; + }; + + if window_function.params.partition_by.is_empty() { + return None; + } + + if window_function.fun.name() != "row_number" { Review Comment: Matching the window function by name assumes the session's `"row_number"` is the builtin with builtin semantics — a session that registers a custom UDWF under that name would get mis-rewritten. #23682 documents the same assumption on its bucket key ("a session resolves one canonical UDF per name"); suggest either adding the same one-line justification here, or tightening the check by comparing against the registry's canonical `row_number` (you already have `config.function_registry()` in scope at the call site). -- 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]
