feat(optimizer): simplify filter predicate before converting row_number + filter to topn (#22295)

stdrc · web-flow · commit 12da00c3d112 · 2025-06-20T06:39:21.000Z
Signed-off-by: Richard Chien &lt;stdrc@outlook.com&gt;
diff --git a/src/frontend/planner_test/tests/testdata/input/over_window_function.yaml b/src/frontend/planner_test/tests/testdata/input/over_window_function.yaml
@@ -393,6 +393,68 @@
     - stream_plan
     - batch_plan
 
+# TopN with arithmetic on row_number/rank function
+- name: test ibis window function optimization
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) - 1 AS rn
+      FROM t
+    )
+    WHERE rn < 10;
+  expected_outputs:
+    - logical_plan
+    - optimized_logical_plan_for_batch
+    - optimized_logical_plan_for_stream
+- name: test arithmetic with addition
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) + 5 AS rn_plus_five
+      FROM t
+    )
+    WHERE rn_plus_five = 6;
+  expected_outputs:
+    - logical_plan
+    - optimized_logical_plan_for_batch
+    - optimized_logical_plan_for_stream
+- name: test optimization when rn column is kept
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c, rn_plus_one
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) + 1 AS rn_plus_one
+      FROM t
+    )
+    WHERE rn_plus_one = 2;
+  expected_outputs:
+    - logical_plan
+    - optimized_logical_plan_for_batch
+    - optimized_logical_plan_for_stream
+- name: test complex arithmetic not optimized
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c, rn
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) * 2 + 1 AS rn
+      FROM t
+    )
+    WHERE rn = 3;
+  expected_outputs:
+    - logical_plan
+    - optimized_logical_plan_for_batch
+    - optimized_logical_plan_for_stream
+
 # TopN on nexmark schema
 - id: create_bid
   sql: |
diff --git a/src/frontend/planner_test/tests/testdata/output/over_window_function.yaml b/src/frontend/planner_test/tests/testdata/output/over_window_function.yaml
@@ -855,6 +855,110 @@
           └─StreamOverWindow { window_functions: [row_number() OVER(PARTITION BY t.x ORDER BY t.y ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
             └─StreamExchange { dist: HashShard(t.x) }
               └─StreamTableScan { table: t, columns: [t.x, t.y, t.z, t._row_id], stream_scan_type: ArrangementBackfill, stream_key: [t._row_id], pk: [_row_id], dist: UpstreamHashShard(t._row_id) }
+- name: test ibis window function optimization
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) - 1 AS rn
+      FROM t
+    )
+    WHERE rn < 10;
+  logical_plan: |-
+    LogicalProject { exprs: [t.a, t.b, t.c] }
+    └─LogicalFilter { predicate: ($expr1 < 10:Int32) }
+      └─LogicalProject { exprs: [t.a, t.b, t.c, (row_number - 1:Int32) as $expr1] }
+        └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+          └─LogicalProject { exprs: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+            └─LogicalScan { table: t, columns: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+  optimized_logical_plan_for_batch: |-
+    LogicalTopN { order: [t.b ASC], limit: 10, offset: 0, group_key: [t.a] }
+    └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+  optimized_logical_plan_for_stream: |-
+    LogicalTopN { order: [t.b ASC], limit: 10, offset: 0, group_key: [t.a] }
+    └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+- name: test arithmetic with addition
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) + 5 AS rn_plus_five
+      FROM t
+    )
+    WHERE rn_plus_five = 6;
+  logical_plan: |-
+    LogicalProject { exprs: [t.a, t.b, t.c] }
+    └─LogicalFilter { predicate: ($expr1 = 6:Int32) }
+      └─LogicalProject { exprs: [t.a, t.b, t.c, (row_number + 5:Int32) as $expr1] }
+        └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+          └─LogicalProject { exprs: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+            └─LogicalScan { table: t, columns: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+  optimized_logical_plan_for_batch: |-
+    LogicalTopN { order: [t.b ASC], limit: 1, offset: 0, group_key: [t.a] }
+    └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+  optimized_logical_plan_for_stream: |-
+    LogicalTopN { order: [t.b ASC], limit: 1, offset: 0, group_key: [t.a] }
+    └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+- name: test optimization when rn column is kept
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c, rn_plus_one
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) + 1 AS rn_plus_one
+      FROM t
+    )
+    WHERE rn_plus_one = 2;
+  logical_plan: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, $expr1] }
+    └─LogicalFilter { predicate: ($expr1 = 2:Int32) }
+      └─LogicalProject { exprs: [t.a, t.b, t.c, (row_number + 1:Int32) as $expr1] }
+        └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+          └─LogicalProject { exprs: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+            └─LogicalScan { table: t, columns: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+  optimized_logical_plan_for_batch: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, (row_number + 1:Int32) as $expr1] }
+    └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+      └─LogicalTopN { order: [t.b ASC], limit: 1, offset: 0, group_key: [t.a] }
+        └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+  optimized_logical_plan_for_stream: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, (row_number + 1:Int32) as $expr1] }
+    └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+      └─LogicalTopN { order: [t.b ASC], limit: 1, offset: 0, group_key: [t.a] }
+        └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+- name: test complex arithmetic not optimized
+  sql: |
+    CREATE TABLE t (a INT, b INT, c INT);
+    SELECT a, b, c, rn
+    FROM (
+      SELECT
+        a, b, c,
+        ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) * 2 + 1 AS rn
+      FROM t
+    )
+    WHERE rn = 3;
+  logical_plan: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, $expr1] }
+    └─LogicalFilter { predicate: ($expr1 = 3:Int32) }
+      └─LogicalProject { exprs: [t.a, t.b, t.c, ((row_number * 2:Int32) + 1:Int32) as $expr1] }
+        └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+          └─LogicalProject { exprs: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+            └─LogicalScan { table: t, columns: [t.a, t.b, t.c, t._row_id, t._rw_timestamp] }
+  optimized_logical_plan_for_batch: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, ((row_number * 2:Int32) + 1:Int32) as $expr1] }
+    └─LogicalFilter { predicate: (((row_number * 2:Int32) + 1:Int32) = 3:Int32) }
+      └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+        └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
+  optimized_logical_plan_for_stream: |-
+    LogicalProject { exprs: [t.a, t.b, t.c, ((row_number * 2:Int32) + 1:Int32) as $expr1] }
+    └─LogicalFilter { predicate: (((row_number * 2:Int32) + 1:Int32) = 3:Int32) }
+      └─LogicalOverWindow { window_functions: [row_number() OVER(PARTITION BY t.a ORDER BY t.b ASC ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)] }
+        └─LogicalScan { table: t, columns: [t.a, t.b, t.c] }
 - id: create_bid
   sql: |
     /*
diff --git a/src/frontend/src/optimizer/rule/over_window_to_topn_rule.rs b/src/frontend/src/optimizer/rule/over_window_to_topn_rule.rs
@@ -18,11 +18,14 @@ use risingwave_expr::window_function::WindowFuncKind;
 
 use super::{BoxedRule, Rule};
 use crate::PlanRef;
-use crate::expr::{ExprImpl, ExprType, collect_input_refs};
+use crate::expr::{
+    Expr, ExprImpl, ExprRewriter, ExprType, FunctionCall, Literal, collect_input_refs,
+};
 use crate::optimizer::plan_node::generic::GenericPlanRef;
 use crate::optimizer::plan_node::{LogicalFilter, LogicalTopN, PlanTreeNodeUnary};
 use crate::optimizer::property::Order;
 use crate::planner::LIMIT_ALL_COUNT;
+use crate::utils::Condition;
 
 /// Transforms the following pattern to group `TopN` (No Ranking Output).
 ///
@@ -42,6 +45,9 @@ use crate::planner::LIMIT_ALL_COUNT;
 /// FROM ..
 /// WHERE rank [ < | <= | > | >= | = ] ..;
 /// ```
+///
+/// Also optimizes filter arithmetic expressions in the `Project <- Filter <- OverWindow` pattern,
+/// such as simplifying `(row_number - 1) = 0` to `row_number = 1`.
 pub struct OverWindowToTopNRule;
 
 impl OverWindowToTopNRule {
@@ -65,6 +71,13 @@ impl Rule for OverWindowToTopNRule {
         // The filter is directly on top of the over window after predicate pushdown.
         let over_window = plan.as_logical_over_window()?;
 
+        // First try to simplify filter arithmetic expressions
+        let filter = if let Some(simplified) = self.simplify_filter_arithmetic(filter) {
+            simplified
+        } else {
+            filter.clone()
+        };
+
         if over_window.window_functions().len() != 1 {
             // Queries with multiple window function calls are not supported yet.
             return None;
@@ -137,6 +150,131 @@ impl Rule for OverWindowToTopNRule {
     }
 }
 
+impl OverWindowToTopNRule {
+    /// Simplify arithmetic expressions in filter conditions before TopN optimization
+    /// For example: `(row_number - 1) = 0` -> `row_number = 1`
+    fn simplify_filter_arithmetic(&self, filter: &LogicalFilter) -> Option<LogicalFilter> {
+        let new_predicate = self.simplify_filter_arithmetic_condition(filter.predicate())?;
+        Some(LogicalFilter::new(filter.input(), new_predicate))
+    }
+
+    /// Simplify arithmetic expressions in the filter condition
+    fn simplify_filter_arithmetic_condition(&self, predicate: &Condition) -> Option<Condition> {
+        let expr = predicate.as_expr_unless_true()?;
+        let mut rewriter = FilterArithmeticRewriter {};
+        let new_expr = rewriter.rewrite_expr(expr.clone());
+
+        if new_expr != expr {
+            Some(Condition::with_expr(new_expr))
+        } else {
+            None
+        }
+    }
+}
+
+/// Filter arithmetic simplification rewriter: simplifies `(col op const) = const2` to `col = (const2 reverse_op const)`
+struct FilterArithmeticRewriter {}
+
+impl ExprRewriter for FilterArithmeticRewriter {
+    fn rewrite_function_call(&mut self, func_call: FunctionCall) -> ExprImpl {
+        use ExprType::{
+            Equal, GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual, NotEqual,
+        };
+
+        // Check if this is a comparison operation
+        match func_call.func_type() {
+            Equal | NotEqual | LessThan | LessThanOrEqual | GreaterThan | GreaterThanOrEqual => {
+                let inputs = func_call.inputs();
+                if inputs.len() == 2 {
+                    // Check if left operand is an arithmetic expression and right operand is a constant
+                    if let ExprImpl::FunctionCall(left_func) = &inputs[0] {
+                        if inputs[1].is_const() {
+                            if let Some(simplified) = self.simplify_arithmetic_comparison(
+                                left_func,
+                                &inputs[1],
+                                func_call.func_type(),
+                            ) {
+                                return simplified;
+                            }
+                        }
+                    }
+                }
+            }
+            _ => {}
+        }
+
+        // Recursively handle sub-expressions
+        let (func_type, inputs, ret_type) = func_call.decompose();
+        let new_inputs: Vec<_> = inputs
+            .into_iter()
+            .map(|input| self.rewrite_expr(input))
+            .collect();
+
+        FunctionCall::new_unchecked(func_type, new_inputs, ret_type).into()
+    }
+}
+
+impl FilterArithmeticRewriter {
+    /// Simplify arithmetic comparison: `(col op const1) comp const2` -> `col comp (const2 reverse_op const1)`
+    fn simplify_arithmetic_comparison(
+        &self,
+        arithmetic_func: &FunctionCall,
+        comparison_const: &ExprImpl,
+        comparison_op: ExprType,
+    ) -> Option<ExprImpl> {
+        use ExprType::{Add, Subtract};
+
+        // Check arithmetic operation
+        match arithmetic_func.func_type() {
+            Add | Subtract => {
+                let inputs = arithmetic_func.inputs();
+                if inputs.len() == 2 {
+                    // Find column reference and constant
+                    let (column_ref, arith_const, reverse_op) = if inputs[1].is_const() {
+                        // col op const
+                        let reverse_op = match arithmetic_func.func_type() {
+                            Add => Subtract,
+                            Subtract => Add,
+                            _ => unreachable!(),
+                        };
+                        (&inputs[0], &inputs[1], reverse_op)
+                    } else if inputs[0].is_const() && arithmetic_func.func_type() == Add {
+                        // const + col
+                        (&inputs[1], &inputs[0], Subtract)
+                    } else {
+                        return None;
+                    };
+
+                    // Calculate new constant value
+                    if let Ok(new_const_func) = FunctionCall::new(
+                        reverse_op,
+                        vec![comparison_const.clone(), arith_const.clone()],
+                    ) {
+                        let new_const_expr: ExprImpl = new_const_func.into();
+                        // Try constant folding
+                        if let Some(Ok(Some(folded_value))) = new_const_expr.try_fold_const() {
+                            let new_const =
+                                Literal::new(Some(folded_value), new_const_expr.return_type())
+                                    .into();
+
+                            // Construct new comparison expression
+                            if let Ok(new_comparison) = FunctionCall::new(
+                                comparison_op,
+                                vec![column_ref.clone(), new_const],
+                            ) {
+                                return Some(new_comparison.into());
+                            }
+                        }
+                    }
+                }
+            }
+            _ => {}
+        }
+
+        None
+    }
+}
+
 /// Returns `None` if the conditions are too complex or invalid. `Some((limit, offset))` otherwise.
 fn handle_rank_preds(rank_preds: &[ExprImpl], window_func_pos: usize) -> Option<(u64, u64)> {
     if rank_preds.is_empty() {
