[feature](Nereids) eliminate sort that is not directly below result sink (#22550)

eliminate sort that is not directly below result sink.
TODO:
handle select c1 + c2 from (select c1, c2 from t order by c1) v;

fe/fe-core/src/main/java/org/apache/doris/nereids/jobs/executor/Rewriter.java

@@ -56,6 +56,7 @@ import org.apache.doris.nereids.rules.rewrite.EliminateLimit;
 import org.apache.doris.nereids.rules.rewrite.EliminateNotNull;
 import org.apache.doris.nereids.rules.rewrite.EliminateNullAwareLeftAntiJoin;
 import org.apache.doris.nereids.rules.rewrite.EliminateOrderByConstant;
+import org.apache.doris.nereids.rules.rewrite.EliminateSort;
 import org.apache.doris.nereids.rules.rewrite.EliminateUnnecessaryProject;
 import org.apache.doris.nereids.rules.rewrite.EnsureProjectOnTopJoin;
 import org.apache.doris.nereids.rules.rewrite.ExtractAndNormalizeWindowExpression;
@@ -296,6 +297,8 @@ public class Rewriter extends AbstractBatchJobExecutor {
                             new PushdownFilterThroughProject(),
                             new MergeProjects()
                     ),
+                    // SORT_PRUNING should be applied after mergeLimit
+                    custom(RuleType.ELIMINATE_SORT, EliminateSort::new),
                     custom(RuleType.ADJUST_CONJUNCTS_RETURN_TYPE, AdjustConjunctsReturnType::new),
                     bottomUp(
                             new ExpressionRewrite(CheckLegalityAfterRewrite.INSTANCE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/RuleType.java

@@ -158,6 +158,7 @@ public enum RuleType {
     PUSHDOWN_DISTINCT_THROUGH_JOIN(RuleTypeClass.REWRITE),
 
     COLUMN_PRUNING(RuleTypeClass.REWRITE),
+    ELIMINATE_SORT(RuleTypeClass.REWRITE),
 
     PUSHDOWN_TOP_N_THROUGH_PROJECTION_WINDOW(RuleTypeClass.REWRITE),
     PUSHDOWN_TOP_N_THROUGH_WINDOW(RuleTypeClass.REWRITE),

fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/EliminateSort.java

@@ -0,0 +1,87 @@
+// 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.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.jobs.JobContext;
+import org.apache.doris.nereids.trees.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalProject;
+import org.apache.doris.nereids.trees.plans.logical.LogicalSink;
+import org.apache.doris.nereids.trees.plans.logical.LogicalSort;
+import org.apache.doris.nereids.trees.plans.visitor.CustomRewriter;
+import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
+
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * Eliminate sort that is not directly below result sink
+ * Note we have put limit in sort node so that we don't need to consider limit
+ */
+
+public class EliminateSort extends DefaultPlanRewriter<Boolean> implements CustomRewriter {
+    @Override
+    public Plan rewriteRoot(Plan plan, JobContext jobContext) {
+        Boolean eliminateSort = false;
+        return plan.accept(this, eliminateSort);
+    }
+
+    @Override
+    public Plan visit(Plan plan, Boolean pruneSort) {
+        List<Plan> newChildren = new ArrayList<>();
+        boolean hasNewChildren = false;
+        for (Plan child : plan.children()) {
+            Plan newChild = child.accept(this, true);
+            if (newChild != child) {
+                hasNewChildren = true;
+            }
+            newChildren.add(newChild);
+        }
+        return hasNewChildren ? plan.withChildren(newChildren) : plan;
+    }
+
+    @Override
+    public Plan visitLogicalSort(LogicalSort<? extends Plan> sort, Boolean eliminateSort) {
+        if (eliminateSort) {
+            return visit(sort.child(), true);
+        }
+        return visit(sort, true);
+    }
+
+    @Override
+    public Plan visitLogicalProject(LogicalProject<? extends Plan> project, Boolean eliminateSort) {
+        return skipEliminateSort(project, eliminateSort);
+    }
+
+    @Override
+    public Plan visitLogicalSink(LogicalSink<? extends Plan> sink, Boolean eliminateSort) {
+        return skipEliminateSort(sink, eliminateSort);
+    }
+
+    private Plan skipEliminateSort(Plan plan, Boolean eliminateSort) {
+        List<Plan> newChildren = new ArrayList<>();
+        boolean hasNewChildren = false;
+        for (Plan child : plan.children()) {
+            Plan newChild = child.accept(this, eliminateSort);
+            if (newChild != child) {
+                hasNewChildren = true;
+            }
+            newChildren.add(newChild);
+        }
+        return hasNewChildren ? plan.withChildren(newChildren) : plan;
+    }
+}

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

@@ -0,0 +1,50 @@
+// 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.
+
+package org.apache.doris.nereids.rules.rewrite;
+
+import org.apache.doris.nereids.util.MemoPatternMatchSupported;
+import org.apache.doris.nereids.util.PlanChecker;
+import org.apache.doris.utframe.TestWithFeService;
+
+import org.junit.jupiter.api.Test;
+
+/**
+ * column prune ut.
+ */
+public class EliminateSortTest extends TestWithFeService implements MemoPatternMatchSupported {
+    @Override
+    protected void runBeforeAll() throws Exception {
+        createDatabase("test");
+        createTable("create table test.student (\n" + "id int not null,\n" + "name varchar(128),\n"
+                + "age int,sex int)\n" + "distributed by hash(id) buckets 10\n"
+                + "properties('replication_num' = '1');");
+        connectContext.setDatabase("default_cluster:test");
+    }
+
+    @Test
+    public void test() {
+        PlanChecker.from(connectContext)
+                .analyze("select * from student order by id")
+                .rewrite()
+                .matches(logicalSort());
+        PlanChecker.from(connectContext)
+                .analyze("select count(*) from (select * from student order by id) t")
+                .rewrite()
+                .nonMatch(logicalSort());
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/nereids/util/PlanChecker.java

@@ -464,6 +464,13 @@ public class PlanChecker {
         return this;
     }
 
+    public PlanChecker nonMatch(PatternDescriptor<? extends Plan> patternDesc) {
+        Memo memo = cascadesContext.getMemo();
+        checkSlotFromChildren(memo);
+        assertMatches(memo, () -> !MatchingUtils.topDownFindMatching(memo.getRoot(), patternDesc.pattern));
+        return this;
+    }
+
     // TODO: remove it.
     public PlanChecker matchesNotCheck(PatternDescriptor<? extends Plan> patternDesc) {
         Memo memo = cascadesContext.getMemo();