branch-4.2: [fix](nereids) Fix deep copy of scans with virtual columns #69399

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<NamedExpression> 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<NamedExpression> 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<? extends Plan> sort, DeepCopierContext
     @Override
     public Plan visitLogicalTopN(LogicalTopN<? extends Plan> topN, DeepCopierContext context) {
         Plan child = topN.child().accept(this, context);
-        List<OrderKey> orderKeys = topN.getOrderKeys().stream()
+        List<OrderKey> orderKeys = copyOrderKeys(topN.getOrderKeys(), context);
+        return new LogicalTopN<>(orderKeys, topN.getLimit(), topN.getOffset(), child);
+    }
+
+    private static List<OrderKey> copyOrderKeys(List<OrderKey> 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

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<NamedExpression> virtualColumns) {
-        LogicalProperties logicalProperties = getLogicalProperties();
-        List<Slot> output = Lists.newArrayList(logicalProperties.getOutput());
-        output.addAll(virtualColumns.stream().map(NamedExpression::toSlot).collect(Collectors.toList()));
-        logicalProperties = new LogicalProperties(() -> output, this::computeDataTrait);
+        List<Slot> 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<NamedExpression> 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<Slot> outputWithVirtualColumns(List<NamedExpression> newVirtualColumns) {
+        Set<ExprId> oldVirtualColumns = virtualColumns.stream()
+                .map(NamedExpression::getExprId)
+                .collect(Collectors.toSet());
+        List<Slot> 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.

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<float> 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("");
+        }
+    }
 }

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<NamedExpression> virtualColumns : ImmutableList.<List<NamedExpression>>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);

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
+

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<FLOAT> 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
+    """
+}