[fix](fe) Preserve required slots in TopN lazy materialization

fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopN.java

@@ -42,6 +42,7 @@
 import org.apache.doris.nereids.trees.plans.physical.PhysicalTopN;
 import org.apache.doris.qe.SessionVariable;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.collect.ImmutableList;
@@ -133,13 +134,10 @@ private Plan doComputeTopN(PhysicalTopN<? extends Plan> topN, CascadesContext ct
                 materializedSlots.add(slot);
             }
         }
-        List<Slot> requiredOutputSlots = new ArrayList<>();
-        for (Map.Entry<Slot, MaterializeSource> entry : materializeMap.entrySet()) {
-            if (requiredMaterializedSlots.contains(entry.getKey())
-                    || requiredMaterializedSlots.contains(entry.getValue().baseSlot)) {
-                requiredOutputSlots.add(entry.getKey());
-            }
-        }
+        // A lazy alias can share its base slot with another output that must be materialized for TopN.
+        // Keep the alias materialized too, otherwise LazySlotPruning removes the base slot needed by that output.
+        List<Slot> requiredOutputSlots = collectRequiredOutputSlots(
+                materializeMap, requiredMaterializedSlots, new HashSet<>(materializedSlots));
         for (Slot slot : requiredOutputSlots) {
             if (materializeMap.remove(slot) != null) {
                 materializedSlots.add(slot);
@@ -235,6 +233,20 @@ private Plan doComputeTopN(PhysicalTopN<? extends Plan> topN, CascadesContext ct
         return result;
     }
 
+    @VisibleForTesting
+    static List<Slot> collectRequiredOutputSlots(Map<Slot, MaterializeSource> materializeMap,
+            Set<Slot> requiredMaterializedSlots, Set<Slot> materializedSlots) {
+        List<Slot> requiredOutputSlots = new ArrayList<>();
+        for (Map.Entry<Slot, MaterializeSource> entry : materializeMap.entrySet()) {
+            if (requiredMaterializedSlots.contains(entry.getKey())
+                    || requiredMaterializedSlots.contains(entry.getValue().baseSlot)
+                    || materializedSlots.contains(entry.getValue().baseSlot)) {
+                requiredOutputSlots.add(entry.getKey());
+            }
+        }
+        return requiredOutputSlots;
+    }
+
     private void collectProjectExprInputSlots(Plan plan, Set<Slot> requiredMaterializedSlots) {
         if (plan instanceof PhysicalProject) {
             PhysicalProject<?> project = (PhysicalProject<?>) plan;

fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruning.java

@@ -43,6 +43,7 @@
 import org.apache.doris.nereids.trees.plans.visitor.DefaultPlanRewriter;
 import org.apache.doris.qe.SessionVariable;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Lists;
 
@@ -146,8 +147,8 @@ public Plan visitPhysicalFilter(PhysicalFilter<? extends Plan> filter, Context c
                         filter.child().accept(this, contextForScan));
                 filter = (PhysicalFilter<? extends Plan>) filter
                         .copyStatsAndGroupIdFrom(filter).resetLogicalProperties();
-                List<Slot> filterOutput = Lists.newArrayList(filter.getOutput());
-                filterOutput.removeAll(filter.getInputSlots());
+                // Predicate slots that are not lazy can still be required by TopN order keys.
+                List<Slot> filterOutput = computeFilterOutput(filter.getOutput(), context.lazySlots);
                 return new PhysicalProject<>(
                         filterOutput.stream().map(s -> (SlotReference) s).collect(Collectors.toList()),
                         Optional.empty(), null,
@@ -157,6 +158,13 @@ public Plan visitPhysicalFilter(PhysicalFilter<? extends Plan> filter, Context c
         return visit(filter, context);
     }
 
+    @VisibleForTesting
+    static List<Slot> computeFilterOutput(List<Slot> output, List<Slot> lazySlots) {
+        List<Slot> filterOutput = Lists.newArrayList(output);
+        filterOutput.removeAll(lazySlots);
+        return filterOutput;
+    }
+
     @Override
     public Plan visitPhysicalOlapScan(PhysicalOlapScan scan, Context context) {
         if (scan.getOutput().containsAll(context.lazySlots)) {

fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopNTest.java

@@ -0,0 +1,54 @@
+// 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.processor.post.materialize;
+
+import org.apache.doris.nereids.trees.expressions.Alias;
+import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.expressions.SlotReference;
+import org.apache.doris.nereids.trees.plans.algebra.Relation;
+import org.apache.doris.nereids.types.IntegerType;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+import java.util.List;
+import java.util.Map;
+
+public class LazyMaterializeTopNTest {
+
+    @Test
+    public void testAliasIsRequiredWhenItsBaseSlotIsMaterialized() {
+        SlotReference baseSlot = new SlotReference("base", IntegerType.INSTANCE);
+        Slot aliasSlot = new Alias(baseSlot, "alias").toSlot();
+        SlotReference independentBaseSlot = new SlotReference("independent", IntegerType.INSTANCE);
+        Slot independentAliasSlot = new Alias(independentBaseSlot, "independent_alias").toSlot();
+        Relation relation = Mockito.mock(Relation.class);
+        Map<Slot, MaterializeSource> materializeMap = ImmutableMap.of(
+                aliasSlot, new MaterializeSource(relation, baseSlot),
+                independentAliasSlot, new MaterializeSource(relation, independentBaseSlot));
+
+        List<Slot> requiredOutputSlots = LazyMaterializeTopN.collectRequiredOutputSlots(
+                materializeMap, ImmutableSet.of(), ImmutableSet.of(baseSlot));
+
+        Assertions.assertEquals(ImmutableList.of(aliasSlot), requiredOutputSlots);
+    }
+}

fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruningTest.java

@@ -0,0 +1,45 @@
+// 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.processor.post.materialize;
+
+import org.apache.doris.nereids.trees.expressions.Slot;
+import org.apache.doris.nereids.trees.expressions.SlotReference;
+import org.apache.doris.nereids.types.IntegerType;
+
+import com.google.common.collect.ImmutableList;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+public class LazySlotPruningTest {
+
+    @Test
+    public void testFilterKeepsMaterializedPredicateSlot() {
+        SlotReference orderKey = new SlotReference("order_key", IntegerType.INSTANCE);
+        SlotReference lazyPredicateSlot = new SlotReference("lazy_predicate", IntegerType.INSTANCE);
+        SlotReference lazyOutputSlot = new SlotReference("lazy_output", IntegerType.INSTANCE);
+        SlotReference rowId = new SlotReference("row_id", IntegerType.INSTANCE);
+
+        List<Slot> filterOutput = LazySlotPruning.computeFilterOutput(
+                ImmutableList.of(orderKey, lazyPredicateSlot, rowId),
+                ImmutableList.of(lazyPredicateSlot, lazyOutputSlot));
+
+        Assertions.assertEquals(ImmutableList.of(orderKey, rowId), filterOutput);
+    }
+}

regression-test/data/query_p0/topn_lazy/topn_lazy_on_data_model.out

@@ -42,6 +42,10 @@ PhysicalResultSink
 --------------filter((mow.__DORIS_DELETE_SIGN__ = 0))
 ----------------PhysicalLazyMaterializeOlapScan[mow lazySlots:(mow.age)]
 
+-- !mow_alias_with_materialized_base_slot --
+a	\N
+b	\N
+
 -- !shape_aggkey_not_lazy --
 PhysicalResultSink
 --PhysicalProject

regression-test/data/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.out

@@ -41,3 +41,6 @@ PhysicalResultSink
 ------------filter((t1.user_id > 0))
 --------------PhysicalLazyMaterializeOlapScan[t1 lazySlots:(t1.username,t1.age,t1.addr)]
 
+-- !filter_order_key_not_pruned --
+1	10	\N	a
+

regression-test/suites/query_p0/topn_lazy/topn_lazy_on_data_model.groovy

@@ -61,6 +61,14 @@ suite("topn_lazy_on_data_model") {
     qt_shape_mow_key_lazy "explain shape plan select user_id  from mow order by username limit 1"
     qt_shape_mow_value_lazy "explain shape plan select age  from mow order by username limit 1"
 
+    order_qt_mow_alias_with_materialized_base_slot """
+        SELECT username AS username_alias, NULL
+        FROM mow
+        QUALIFY ROW_NUMBER() OVER (PARTITION BY username, user_id) = 1
+        ORDER BY username
+        LIMIT 5
+    """
+
     // agg key user_id is lazy materialized
     sql """
         drop table if exists agg; 
@@ -85,4 +93,4 @@ suite("topn_lazy_on_data_model") {
 
 
 
-}
+}

regression-test/suites/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.groovy

@@ -50,7 +50,30 @@ suite("topNLazyMaterializationUsingIndex") {
 
         insert into t2 values ( 1, 'a', 10, 'cd'),(1,'b', 20, 'cq');
 
+        drop table if exists topn_lazy_filter_order_key;
+        CREATE TABLE topn_lazy_filter_order_key
+        (
+        `k1` INT NOT NULL,
+        `k2` INT NOT NULL,
+        `v` INT NULL,
+        `pad` STRING NULL,
+        INDEX idx_v(v) USING INVERTED
+        )
+        DUPLICATE KEY(k1, k2)
+        DISTRIBUTED BY HASH(k1) BUCKETS 1
+        PROPERTIES (
+        "replication_allocation" = "tag.location.default: 1");
+
+        insert into topn_lazy_filter_order_key values
+            (1, 10, NULL, 'a'),
+            (2, 20, 1, 'b'),
+            (3, 30, 2, 'c'),
+            (4, 40, NULL, 'd'),
+            (5, 50, 3, 'e');
+
         set topn_lazy_materialization_using_index = true;
+        set topn_lazy_materialization_threshold = 1;
+        set enable_segment_limit_pushdown = false;
         SET detail_shape_nodes='PhysicalProject';
         """
         qt_plan """
@@ -76,4 +99,12 @@ suite("topNLazyMaterializationUsingIndex") {
             user_id > 0 order by user_id limit 1;
             """
 
+        order_qt_filter_order_key_not_pruned """
+        select k1, k2, v, pad
+        from topn_lazy_filter_order_key
+        where (v is null or k2 < 40)
+        order by k1, k2
+        limit 1;
+        """
+
 }