diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java index 29fd9e9e4c5c34..5cb069b5aaba12 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopier.java @@ -20,7 +20,6 @@ import org.apache.doris.nereids.exceptions.AnalysisException; import org.apache.doris.nereids.hint.DistributeHint; import org.apache.doris.nereids.properties.OrderKey; -import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.ExprId; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.MarkJoinSlotReference; @@ -46,6 +45,7 @@ import org.apache.doris.nereids.trees.plans.logical.LogicalIntersect; import org.apache.doris.nereids.trees.plans.logical.LogicalJoin; import org.apache.doris.nereids.trees.plans.logical.LogicalLimit; +import org.apache.doris.nereids.trees.plans.logical.LogicalOlapScan; import org.apache.doris.nereids.trees.plans.logical.LogicalOneRowRelation; import org.apache.doris.nereids.trees.plans.logical.LogicalPartitionTopN; import org.apache.doris.nereids.trees.plans.logical.LogicalPlan; @@ -108,22 +108,24 @@ public Plan visitLogicalCatalogRelation(LogicalCatalogRelation catalogRelation, LogicalCatalogRelation newRelation = catalogRelation.withRelationId(StatementScopeIdGenerator.newRelationId()); updateReplaceMapWithOutput(catalogRelation, newRelation, context.exprIdReplaceMap); - List virtualColumns = catalogRelation.getVirtualColumns().stream() - .map(e -> { - if (e instanceof Alias) { - return new Alias(((Alias) e).child(), e.getName()); - } - return e; - }) - .collect(ImmutableList.toImmutableList()); - for (int i = 0; i < virtualColumns.size(); i++) { - context.exprIdReplaceMap.put(catalogRelation.getVirtualColumns().get(i).getExprId(), - virtualColumns.get(i).getExprId()); + if (!catalogRelation.getVirtualColumns().isEmpty()) { + // A virtual column can refer to the others, so each one gets its new id before any is copied. + for (NamedExpression column : catalogRelation.getVirtualColumns()) { + context.exprIdReplaceMap.put(column.getExprId(), StatementScopeIdGenerator.newExprId()); + } + List virtualColumns = catalogRelation.getVirtualColumns().stream() + .map(column -> (NamedExpression) ExpressionDeepCopier.INSTANCE.deepCopy(column, context)) + .collect(ImmutableList.toImmutableList()); + newRelation = newRelation.withVirtualColumns(virtualColumns); + if (newRelation instanceof LogicalOlapScan) { + // The score and ANN order keys sort by virtual columns. + LogicalOlapScan olapScan = (LogicalOlapScan) newRelation; + newRelation = olapScan.appendVirtualColumnsAndTopN(ImmutableList.of(), + copyOrderKeys(olapScan.getAnnOrderKeys(), context), olapScan.getAnnLimit(), + copyOrderKeys(olapScan.getScoreOrderKeys(), context), olapScan.getScoreLimit(), + olapScan.getScoreRangeInfo()); + } } - virtualColumns = virtualColumns.stream() - .map(o -> (NamedExpression) ExpressionDeepCopier.INSTANCE.deepCopy(o, context)) - .collect(ImmutableList.toImmutableList()); - newRelation = newRelation.withVirtualColumns(virtualColumns); context.putRelation(catalogRelation.getRelationId(), newRelation); return updateOperativeSlots(catalogRelation, newRelation); } @@ -258,11 +260,15 @@ public Plan visitLogicalSort(LogicalSort sort, DeepCopierContext @Override public Plan visitLogicalTopN(LogicalTopN topN, DeepCopierContext context) { Plan child = topN.child().accept(this, context); - List orderKeys = topN.getOrderKeys().stream() + List orderKeys = copyOrderKeys(topN.getOrderKeys(), context); + return new LogicalTopN<>(orderKeys, topN.getLimit(), topN.getOffset(), child); + } + + private static List copyOrderKeys(List orderKeys, DeepCopierContext context) { + return orderKeys.stream() .map(o -> new OrderKey(ExpressionDeepCopier.INSTANCE.deepCopy(o.getExpr(), context), o.isAsc(), o.isNullFirst())) .collect(ImmutableList.toImmutableList()); - return new LogicalTopN<>(orderKeys, topN.getLimit(), topN.getOffset(), child); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java index 56e8015d9147ac..d212e361c7e8ad 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalOlapScan.java @@ -559,16 +559,14 @@ public LogicalOlapScan withTableAlias(String tableAlias) { } /** - * add virtual column to olap scan. + * replace the virtual columns of olap scan, also in its output. * @param virtualColumns generated virtual columns * @return scan with virtual columns */ @Override public LogicalOlapScan withVirtualColumns(List virtualColumns) { - LogicalProperties logicalProperties = getLogicalProperties(); - List output = Lists.newArrayList(logicalProperties.getOutput()); - output.addAll(virtualColumns.stream().map(NamedExpression::toSlot).collect(Collectors.toList())); - logicalProperties = new LogicalProperties(() -> output, this::computeDataTrait); + List output = outputWithVirtualColumns(virtualColumns); + LogicalProperties logicalProperties = new LogicalProperties(() -> output, this::computeDataTrait); return new LogicalOlapScan(relationId, (Table) table, qualifier, groupExpression, Optional.of(logicalProperties), selectedPartitionIds, partitionPruned, hasPartitionPredicate, selectedTabletIds, @@ -578,6 +576,21 @@ public LogicalOlapScan withVirtualColumns(List virtualColumns) scoreRangeInfo, annOrderKeys, annLimit, partitionPrunablePredicates, tableAlias); } + /** + * The output with the slots of the given virtual columns instead of the slots of the current ones. The other + * slots keep their ids, which the slots of the parent plan may already refer to. + */ + protected List outputWithVirtualColumns(List newVirtualColumns) { + Set oldVirtualColumns = virtualColumns.stream() + .map(NamedExpression::getExprId) + .collect(Collectors.toSet()); + List output = getLogicalProperties().getOutput().stream() + .filter(slot -> !oldVirtualColumns.contains(slot.getExprId())) + .collect(Collectors.toCollection(ArrayList::new)); + newVirtualColumns.forEach(column -> output.add(column.toSlot())); + return output; + } + /** * Append additional virtual columns to existing ones. * Unlike {@link #withVirtualColumns} which replaces, this merges existing + new. diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CTEInlineTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CTEInlineTest.java index 7ae9aa1e99569f..1acebe8e3f8259 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CTEInlineTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/rules/rewrite/CTEInlineTest.java @@ -49,6 +49,24 @@ protected void runBeforeAll() throws Exception { + "DUPLICATE KEY(id)\n" + "DISTRIBUTED BY HASH(id) BUCKETS 1\n" + "PROPERTIES (\"replication_num\" = \"1\")"); + createTable("CREATE TABLE cte_inline_text (\n" + + " id int NULL,\n" + + " c1 text NULL,\n" + + " v variant NULL,\n" + + " INDEX idx_c1 (c1) USING INVERTED PROPERTIES (\"parser\" = \"english\")\n" + + ") ENGINE=OLAP\n" + + "DUPLICATE KEY(id)\n" + + "DISTRIBUTED BY HASH(id) BUCKETS 1\n" + + "PROPERTIES (\"replication_num\" = \"1\", \"inverted_index_storage_format\" = \"V3\")"); + createTable("CREATE TABLE cte_inline_ann (\n" + + " id int NOT NULL,\n" + + " embedding array NOT NULL,\n" + + " INDEX idx_ann (embedding) USING ANN PROPERTIES (\n" + + " \"index_type\" = \"hnsw\", \"metric_type\" = \"l2_distance\", \"dim\" = \"2\")\n" + + ") ENGINE=OLAP\n" + + "DUPLICATE KEY(id)\n" + + "DISTRIBUTED BY HASH(id) BUCKETS 1\n" + + "PROPERTIES (\"replication_num\" = \"1\")"); } @Test @@ -119,4 +137,28 @@ public void refreshCteConsumersAfterNormalizeEliminatesEmptyBranch() { connectContext.getSessionVariable().setDisableNereidsRules(""); } } + + @Test + public void inlineCteWhoseScanHasVirtualColumns() { + // The score and ANN TopN and the MATCH projection add virtual columns to the scan of the CTE, + // before the optimizer inlines it. + connectContext.getSessionVariable().setDisableNereidsRules("PRUNE_EMPTY_PARTITION"); + try { + for (String body : new String[] { + "select id, score() as s from cte_inline_text where c1 match 'apple' order by s desc limit 2", + "select id from cte_inline_ann order by l2_distance_approximate(embedding, [1.0, 2.0]) limit 2", + "select id, c1 match 'apple' as m from cte_inline_text", + // The copied scan also outputs the VARIANT subpath, whose slot the copy has mapped. + "select id, v['x'] as x, c1 match 'apple' as m from cte_inline_text"}) { + for (String query : new String[] {"with cte as (%s) select * from cte", + "with cte as (%s) select * from cte union all select * from cte where 1 = 0"}) { + String sql = String.format(query, body); + Assertions.assertDoesNotThrow(() -> PlanChecker.from(connectContext).checkPlannerResult(sql), + sql); + } + } + } finally { + connectContext.getSessionVariable().setDisableNereidsRules(""); + } + } } diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java index 98f07698a9ef29..b022d27ea5c31a 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/copier/LogicalPlanDeepCopierTest.java @@ -17,10 +17,15 @@ package org.apache.doris.nereids.trees.copier; +import org.apache.doris.nereids.properties.OrderKey; +import org.apache.doris.nereids.trees.expressions.Add; +import org.apache.doris.nereids.trees.expressions.Alias; import org.apache.doris.nereids.trees.expressions.Expression; +import org.apache.doris.nereids.trees.expressions.Multiply; import org.apache.doris.nereids.trees.expressions.NamedExpression; import org.apache.doris.nereids.trees.expressions.Slot; import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.functions.scalar.Score; import org.apache.doris.nereids.trees.plans.Plan; import org.apache.doris.nereids.trees.plans.algebra.Repeat.RepeatType; import org.apache.doris.nereids.trees.plans.logical.LogicalAggregate; @@ -50,6 +55,47 @@ public void testDeepCopyOlapScan() { } } + @Test + public void testDeepCopyOlapScanWithVirtualColumn() { + LogicalOlapScan scan = PlanConstructor.newLogicalOlapScan(0, "a", 0); + Alias scoreColumn = new Alias(new Score(), "__score"); + scan = scan.appendVirtualColumnsAndTopN(ImmutableList.of(scoreColumn), ImmutableList.of(), Optional.empty(), + ImmutableList.of(new OrderKey(scoreColumn.toSlot(), false, false)), Optional.of(0L), + Optional.empty()); + DeepCopierContext context = new DeepCopierContext(); + LogicalOlapScan copy = (LogicalOlapScan) scan.accept(LogicalPlanDeepCopier.INSTANCE, context); + + // The copy outputs one new score column, which its order key and the slots above it use. + Slot copiedScore = copy.getVirtualColumns().get(0).toSlot(); + Assertions.assertNotEquals(scoreColumn.getExprId(), copiedScore.getExprId()); + Assertions.assertEquals(scan.getOutput().size(), copy.getOutput().size()); + Assertions.assertTrue(copy.getOutput().contains(copiedScore)); + Assertions.assertFalse(copy.getOutput().contains(scoreColumn.toSlot())); + Assertions.assertEquals(copiedScore, copy.getScoreOrderKeys().get(0).getExpr()); + Assertions.assertEquals(copiedScore.getExprId(), context.exprIdReplaceMap.get(scoreColumn.getExprId())); + for (Slot slot : scan.getOutput()) { + Assertions.assertTrue(copy.getOutputExprIdSet().contains(context.exprIdReplaceMap.get(slot.getExprId())), + slot.toString()); + } + } + + @Test + public void testDeepCopyOlapScanWithDependentVirtualColumns() { + LogicalOlapScan scan = PlanConstructor.newLogicalOlapScan(0, "a", 0); + Slot id = scan.getOutput().get(0); + Alias sum = new Alias(new Add(id, id), "v1"); + Alias product = new Alias(new Multiply(sum.toSlot(), sum.toSlot()), "v2"); + // A virtual column can refer to one that comes after it. + for (List virtualColumns : ImmutableList.>of( + ImmutableList.of(sum, product), ImmutableList.of(product, sum))) { + LogicalOlapScan copy = (LogicalOlapScan) scan.withVirtualColumns(virtualColumns) + .accept(LogicalPlanDeepCopier.INSTANCE, new DeepCopierContext()); + for (NamedExpression column : copy.getVirtualColumns()) { + Assertions.assertTrue(copy.getOutputSet().containsAll(column.getInputSlots()), column.toString()); + } + } + } + @Test public void testDeepCopyOlapScanWithNonFirstOperativeSlot() { LogicalOlapScan relationPlan = PlanConstructor.newLogicalOlapScan(0, "a", 0); diff --git a/regression-test/data/nereids_rules_p0/cte/test_cte_inline_virtual_column.out b/regression-test/data/nereids_rules_p0/cte/test_cte_inline_virtual_column.out new file mode 100644 index 00000000000000..3770fa376c362c --- /dev/null +++ b/regression-test/data/nereids_rules_p0/cte/test_cte_inline_virtual_column.out @@ -0,0 +1,14 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !score_topn -- +1 0.69596696 +2 0.49917626 + +-- !ann_topn -- +1 +2 + +-- !match_projection -- +1 true +2 true +3 false + diff --git a/regression-test/suites/nereids_rules_p0/cte/test_cte_inline_virtual_column.groovy b/regression-test/suites/nereids_rules_p0/cte/test_cte_inline_virtual_column.groovy new file mode 100644 index 00000000000000..88f95cb7ae5c29 --- /dev/null +++ b/regression-test/suites/nereids_rules_p0/cte/test_cte_inline_virtual_column.groovy @@ -0,0 +1,68 @@ +// 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. + +// The score and ANN TopN and the MATCH projection add virtual columns to the scan of a CTE, +// which is copied when the optimizer inlines the CTE. +suite("test_cte_inline_virtual_column") { + sql "DROP TABLE IF EXISTS test_cte_inline_virtual_column_text" + sql """ + CREATE TABLE test_cte_inline_virtual_column_text ( + id INT, + c1 TEXT, + INDEX idx_c1 (c1) USING INVERTED PROPERTIES ("parser" = "english") + ) DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ("replication_num" = "1", "inverted_index_storage_format" = "V3") + """ + sql """ + INSERT INTO test_cte_inline_virtual_column_text VALUES + (1, 'apple apple apple'), (2, 'apple banana'), (3, 'banana cherry') + """ + + sql "DROP TABLE IF EXISTS test_cte_inline_virtual_column_ann" + sql """ + CREATE TABLE test_cte_inline_virtual_column_ann ( + id INT NOT NULL, + embedding ARRAY NOT NULL, + INDEX idx_ann (embedding) USING ANN PROPERTIES ( + "index_type" = "hnsw", "metric_type" = "l2_distance", "dim" = "2") + ) DUPLICATE KEY(id) DISTRIBUTED BY HASH(id) BUCKETS 1 + PROPERTIES ("replication_num" = "1") + """ + sql """ + INSERT INTO test_cte_inline_virtual_column_ann VALUES + (1, [1.0, 2.0]), (2, [2.0, 3.0]), (3, [10.0, 10.0]) + """ + + order_qt_score_topn """ + WITH c AS ( + SELECT id, score() AS s FROM test_cte_inline_virtual_column_text + WHERE c1 MATCH 'apple' ORDER BY s DESC LIMIT 2 + ) + SELECT * FROM c + """ + order_qt_ann_topn """ + WITH c AS ( + SELECT id FROM test_cte_inline_virtual_column_ann + ORDER BY l2_distance_approximate(embedding, [1.0, 2.0]) LIMIT 2 + ) + SELECT * FROM c + """ + order_qt_match_projection """ + WITH c AS (SELECT id, c1 MATCH 'apple' AS m FROM test_cte_inline_virtual_column_text) + SELECT * FROM c UNION ALL SELECT * FROM c WHERE 1 = 0 + """ +}