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

### What problem does this PR solve? Issue Number: N/A Related PR: #69292 (depends on this fix) Problem Summary: A CTE fails to plan when the optimizer inlines it after one of its scans got a virtual column: ```sql WITH c AS ( SELECT id, score() AS s FROM t WHERE c1 MATCH 'apple' ORDER BY s DESC LIMIT 2 ) SELECT * FROM c; -- ERROR: A expression contains slot not from children ``` A score or ANN TopN pushed into a scan, and a MATCH in the select list, add a virtual column to the scan. After rewrite, the optimizer inlines a CTE that has one consumer left (`Optimizer.normalizeCtePlan` → `CTEInliner`) by deep copying its rewritten body, so the copy includes these scans: - `score()` and `l2_distance_approximate()` are nondeterministic, so the rewrite does not inline a CTE that calls them, even with one consumer. Once the TopN pushdown has replaced them with a virtual column of the scan, the optimizer inlines the CTE. So every CTE with `ORDER BY score() LIMIT n` or an ANN TopN fails, with default settings. - A CTE with MATCH in its select list reaches the optimizer when the rewrite removes all but one of its consumers, e.g. `... UNION ALL SELECT * FROM c WHERE 1 = 0`. `LogicalPlanDeepCopier` copied a scan with virtual columns wrongly: - it gave each virtual column two new ids: the slots above the scan were mapped to the first one, and the copied scan output the second one; - it copied the virtual columns in list order, so a virtual column that refers to a later one (CSE with `experimental_enable_virtual_slot_for_cse=true`, e.g. `[v2 := v1 * v1, v1 := x + y]`) kept the old id of the later one; - the copied scan still output the original virtual column, with its original id, next to the new one; - the score and ANN order keys of the copied scan still used the original slots. So the plan above the copied scan read a slot that the scan does not output. The fix: - The copier gives every virtual column its new id first and then copies their expressions, which reuse those ids, so virtual columns may refer to each other in any order. It also maps the score and ANN order keys to the new slots. Relations without virtual columns are copied as before. - `LogicalOlapScan.withVirtualColumns` replaces the virtual columns of the scan, but kept their old slots in the output next to the new ones. It now replaces them there too and keeps the other output slots, which the copier has already mapped. Computing the output of the copy again would not work: it gives VARIANT subpath slots new ids. `AdjustNullable`, the other caller that replaces existing virtual columns, no longer keeps two slots for a virtual column whose nullability changes. #69292 binds `score()` to its scan before CTEs are inlined, so it needs this fix; it carries the same change until this PR is merged.
2026-11-04 07:31:17 +00:00 · 2026-10-22 10:35:35 +08:00
1 parent 1052fc31b1
commit 88511083a6
@@ -20,7 +20,6 @@ package org.apache.doris.nereids.trees.copier;
 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;
@@ -113,22 +112,24 @@ public class LogicalPlanDeepCopier extends DefaultPlanRewriter<DeepCopierContext
            newRelation = ((LogicalOlapScan) newRelation).withPartitionPruned(false);
        }
        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);
        if (catalogRelation instanceof LogicalSchemaScan
                && ((LogicalSchemaScan) catalogRelation).isFilterPushed()) {
            LogicalSchemaScan oldSchemaScan = (LogicalSchemaScan) catalogRelation;
@@ -273,11 +274,15 @@ public class LogicalPlanDeepCopier extends DefaultPlanRewriter<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
@@ -687,15 +687,13 @@ public class LogicalOlapScan extends LogicalCatalogRelation implements OlapScan,
    }

    /**
-     * 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()));
+        List<Slot> output = outputWithVirtualColumns(virtualColumns);
        LogicalProperties finalLogicalProperties = new LogicalProperties(() -> output, this::computeDataTrait);
        return AbstractPlan.copyWithSameId(this, () ->
                new LogicalOlapScan(relationId, (Table) table, qualifier,
@@ -707,6 +705,21 @@ public class LogicalOlapScan extends LogicalCatalogRelation implements OlapScan,
                scoreRangeInfo, annOrderKeys, annLimit, tableAlias, partitionPrunablePredicates, scanParams));
    }

+    /**
+     * 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.
@@ -445,9 +445,7 @@ public class LogicalOlapTableStreamScan extends LogicalOlapScan {

    @Override
    public LogicalOlapTableStreamScan withVirtualColumns(List<NamedExpression> virtualColumns) {
-        LogicalProperties logicalProperties = getLogicalProperties();
-        List<Slot> output = Lists.newArrayList(logicalProperties.getOutput());
-        output.addAll(virtualColumns.stream().map(NamedExpression::toSlot).collect(Collectors.toList()));
+        List<Slot> output = outputWithVirtualColumns(virtualColumns);
        LogicalProperties finalLogicalProperties = new LogicalProperties(() -> output, this::computeDataTrait);
        return AbstractPlan.copyWithSameId(this, () ->
                new LogicalOlapTableStreamScan(relationId, (Table) table, qualifier,
@@ -50,6 +50,24 @@ public class CTEInlineTest extends TestWithFeService implements MemoPatternMatch
                + "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
@@ -158,4 +176,28 @@ public class CTEInlineTest extends TestWithFeService implements MemoPatternMatch
            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("");
+        }
+    }
 }
@@ -18,11 +18,16 @@
 package org.apache.doris.nereids.trees.copier;

 import org.apache.doris.catalog.SchemaTable;
+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.EqualTo;
 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.expressions.literal.VarcharLiteral;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.algebra.Repeat.RepeatType;
@@ -84,6 +89,47 @@ public class LogicalPlanDeepCopierTest {
        Assertions.assertTrue(copiedPlan.hasPartitionPredicate());
    }

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