branch-3.1: [fix](nereids) partition topn optimization requires all window expressions are in the same order #56622

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalWindow.java

@@ -23,6 +23,7 @@
 import org.apache.doris.nereids.properties.DataTrait.Builder;
 import org.apache.doris.nereids.properties.FdItem;
 import org.apache.doris.nereids.properties.LogicalProperties;
+import org.apache.doris.nereids.trees.expressions.Alias;
 import org.apache.doris.nereids.trees.expressions.BinaryOperator;
 import org.apache.doris.nereids.trees.expressions.EqualTo;
 import org.apache.doris.nereids.trees.expressions.ExprId;
@@ -211,6 +212,7 @@ && child(0).child(0) instanceof LogicalPartitionTopN)) {
         // 1. The window function should be one of the 'row_number()', 'rank()', 'dense_rank()'.
         // 2. The window frame should be 'UNBOUNDED' to 'CURRENT'.
         // 3. The 'PARTITION' key and 'ORDER' key can not be empty at the same time.
+        // 4. order of window expressions should be compatible.
         WindowExpression chosenWindowFunc = null;
         long chosenPartitionLimit = Long.MAX_VALUE;
         long chosenRowNumberPartitionLimit = Long.MAX_VALUE;
@@ -303,10 +305,24 @@ && child(0).child(0) instanceof LogicalPartitionTopN)) {
                 }
             }
         }
+
         if (chosenWindowFunc == null || (chosenPartitionLimit == Long.MAX_VALUE
                 && chosenRowNumberPartitionLimit == Long.MAX_VALUE)) {
             return null;
         } else {
+            // 4. check all windowExpression's order key is empty or is the same as chosenWindowFunc's order key
+            for (NamedExpression windowExpr : windowExpressions) {
+                if (windowExpr != null && windowExpr instanceof Alias
+                        && windowExpr.child(0) instanceof WindowExpression) {
+                    WindowExpression windowFunc = (WindowExpression) windowExpr.child(0);
+                    if (windowFunc.getOrderKeys().isEmpty()
+                            || windowFunc.getOrderKeys().equals(chosenWindowFunc.getOrderKeys())) {
+                        continue;
+                    } else {
+                        return null;
+                    }
+                }
+            }
             return Pair.of(chosenWindowFunc, hasRowNumber ? chosenRowNumberPartitionLimit : chosenPartitionLimit);
         }
     }

fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/GeneratePartitionTopnFromWindowTest.java

@@ -117,4 +117,44 @@ public void testGeneratePartitionTopnFromWindow() {
                     )
                 );
     }
+
+    @Test
+    public void testMultipleWindowsWithDifferentOrders() {
+        ConnectContext context = MemoTestUtils.createConnectContext();
+        context.getSessionVariable().setEnablePartitionTopN(true);
+        NamedExpression gender = scan.getOutput().get(1).toSlot();
+        NamedExpression age = scan.getOutput().get(3).toSlot();
+
+        List<Expression> partitionKeyList = ImmutableList.of(gender);
+        List<OrderExpression> orderKeyList = ImmutableList.of(new OrderExpression(
+                new OrderKey(age, true, true)));
+        WindowFrame windowFrame = new WindowFrame(WindowFrame.FrameUnitsType.ROWS,
+                WindowFrame.FrameBoundary.newPrecedingBoundary(),
+                WindowFrame.FrameBoundary.newCurrentRowBoundary());
+        WindowExpression window1 = new WindowExpression(new RowNumber(), partitionKeyList, orderKeyList, windowFrame);
+        Alias windowAlias1 = new Alias(window1, window1.toSql());
+
+        List<OrderExpression> orderKeyList2 = ImmutableList.of(new OrderExpression(
+                new OrderKey(age, false, true)));
+        WindowExpression window2 = new WindowExpression(new RowNumber(), partitionKeyList, orderKeyList2, windowFrame);
+        Alias windowAlias2 = new Alias(window2, window2.toSql());
+
+        List<NamedExpression> expressions = Lists.newArrayList(windowAlias1, windowAlias2);
+        LogicalWindow<LogicalOlapScan> window = new LogicalWindow<>(expressions, scan);
+        Expression filterPredicate = new LessThanEqual(window.getOutput().get(4).toSlot(), Literal.of(100));
+
+        LogicalPlan plan = new LogicalPlanBuilder(window)
+                .filter(filterPredicate)
+                .project(ImmutableList.of(0))
+                .build();
+
+        PlanChecker.from(context, plan)
+                .applyTopDown(new CreatePartitionTopNFromWindow())
+                .matches(
+                        logicalProject(
+                                logicalFilter(
+                                        logicalWindow(
+                                                logicalOlapScan()
+                                        ))));
+    }
 }

regression-test/data/nereids_p0/partition_topn/check_orderkey.out

@@ -0,0 +1,4 @@
+-- This file is automatically generated. You should know what you did if you want to edit this
+-- !exe --
+1	c	3	1
+

regression-test/suites/nereids_p0/partition_topn/check_orderkey.groovy

@@ -0,0 +1,61 @@
+// 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.
+
+suite("check_orderkey") {
+    sql """
+    create table students02 (
+        id int, age int, name varchar
+    )properties ('replication_num'='1');
+
+    insert into students02 values (1, 1, 'a'), (1, 2, 'b'), (1, 3, 'c');
+    """
+    //"ORDER BY age desc" and "ORDER BY age" is different, do not use partition topn
+    explain {
+        sql """
+            select
+            t.id,
+            t.name,
+            t.age,
+            t.desc_age
+            from (
+            SELECT
+                ROW_NUMBER() OVER (partition by id ORDER BY age desc) AS desc_age,
+                ROW_NUMBER() OVER (partition by id ORDER BY age) AS age,
+                id,
+                name
+            FROM students02) t
+            where t.desc_age=1;
+            """
+        notContains("VPartitionTopN")
+    }
+
+    qt_exe """
+    select
+        t.id,
+        t.name,
+        t.age,
+        t.desc_age
+    from (
+    SELECT
+        ROW_NUMBER() OVER (partition by id ORDER BY age desc) AS desc_age,
+        ROW_NUMBER() OVER (partition by id ORDER BY age) AS age,
+        id,
+        name
+    FROM students02) t
+    where t.desc_age=1;
+        """
+}