diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopN.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopN.java index 1b5f262bc884df..9acec57125a4b2 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopN.java +++ b/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 topN, CascadesContext ct materializedSlots.add(slot); } } - List requiredOutputSlots = new ArrayList<>(); - for (Map.Entry 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 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 topN, CascadesContext ct return result; } + @VisibleForTesting + static List collectRequiredOutputSlots(Map materializeMap, + Set requiredMaterializedSlots, Set materializedSlots) { + List requiredOutputSlots = new ArrayList<>(); + for (Map.Entry 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 requiredMaterializedSlots) { if (plan instanceof PhysicalProject) { PhysicalProject project = (PhysicalProject) plan; diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruning.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruning.java index f2bf5831206334..3273a7d82d3fa2 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruning.java +++ b/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 filter, Context c filter.child().accept(this, contextForScan)); filter = (PhysicalFilter) filter .copyStatsAndGroupIdFrom(filter).resetLogicalProperties(); - List filterOutput = Lists.newArrayList(filter.getOutput()); - filterOutput.removeAll(filter.getInputSlots()); + // Predicate slots that are not lazy can still be required by TopN order keys. + List 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 filter, Context c return visit(filter, context); } + @VisibleForTesting + static List computeFilterOutput(List output, List lazySlots) { + List filterOutput = Lists.newArrayList(output); + filterOutput.removeAll(lazySlots); + return filterOutput; + } + @Override public Plan visitPhysicalOlapScan(PhysicalOlapScan scan, Context context) { if (scan.getOutput().containsAll(context.lazySlots)) { diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopNTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazyMaterializeTopNTest.java new file mode 100644 index 00000000000000..d8763f270820b5 --- /dev/null +++ b/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 materializeMap = ImmutableMap.of( + aliasSlot, new MaterializeSource(relation, baseSlot), + independentAliasSlot, new MaterializeSource(relation, independentBaseSlot)); + + List requiredOutputSlots = LazyMaterializeTopN.collectRequiredOutputSlots( + materializeMap, ImmutableSet.of(), ImmutableSet.of(baseSlot)); + + Assertions.assertEquals(ImmutableList.of(aliasSlot), requiredOutputSlots); + } +} diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruningTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/processor/post/materialize/LazySlotPruningTest.java new file mode 100644 index 00000000000000..84312664242c70 --- /dev/null +++ b/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 filterOutput = LazySlotPruning.computeFilterOutput( + ImmutableList.of(orderKey, lazyPredicateSlot, rowId), + ImmutableList.of(lazyPredicateSlot, lazyOutputSlot)); + + Assertions.assertEquals(ImmutableList.of(orderKey, rowId), filterOutput); + } +} diff --git a/regression-test/data/query_p0/topn_lazy/topn_lazy_on_data_model.out b/regression-test/data/query_p0/topn_lazy/topn_lazy_on_data_model.out index 694a8249cdf0eb..23c468d9bdd4b6 100644 --- a/regression-test/data/query_p0/topn_lazy/topn_lazy_on_data_model.out +++ b/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 diff --git a/regression-test/data/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.out b/regression-test/data/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.out index 10720c0b9f08b9..38a5cba4eff40d 100644 --- a/regression-test/data/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.out +++ b/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 + diff --git a/regression-test/suites/query_p0/topn_lazy/topn_lazy_on_data_model.groovy b/regression-test/suites/query_p0/topn_lazy/topn_lazy_on_data_model.groovy index ce73bcf01da72b..ee121cc25a29d4 100644 --- a/regression-test/suites/query_p0/topn_lazy/topn_lazy_on_data_model.groovy +++ b/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") { -} \ No newline at end of file +} diff --git a/regression-test/suites/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.groovy b/regression-test/suites/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.groovy index 9f51a025981a77..a7f34129f86fe7 100644 --- a/regression-test/suites/query_p0/topn_lazy/usingIndex/topNLazyMaterializationUsingIndex.groovy +++ b/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; + """ + }