branch-4.2: [fix](fe) Lock table metadata while constructing UPDATE plans (#69974)

cherry-pick #69880

Conflicts in UpdateCommand.java / UpdateCommandTest.java resolved: kept
branch-4.2 UPDATE planning APIs and wrapped metadata reads + plan
construction in readLock/readUnlock; tests use ImmutableList.of instead
of List.of (Java 8 target).

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java

@@ -137,94 +137,101 @@ public class UpdateCommand extends Command implements ForwardWithSync, Explainab
     public LogicalPlan completeQueryPlan(ConnectContext ctx, LogicalPlan logicalQuery) {
         checkTable(ctx);
 
-        Map<String, Expression> colNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER);
-        Map<String, Expression> partialUpdateColNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER);
-        for (EqualTo equalTo : assignments) {
-            List<String> colNameParts = ((UnboundSlot) equalTo.left()).getNameParts();
-            checkAssignmentColumn(ctx, colNameParts, this.nameParts, this.tableAlias);
-            colNameToExpression.put(colNameParts.get(colNameParts.size() - 1), equalTo.right());
-            partialUpdateColNameToExpression.put(colNameParts.get(colNameParts.size() - 1), equalTo.right());
-        }
-        // check if any key in update clause
-        if (targetTable.getFullSchema().stream().filter(Column::isKey)
-                .anyMatch(column -> partialUpdateColNameToExpression.containsKey(column.getName()))) {
-            throw new AnalysisException("Only value columns of unique table could be updated");
-        }
-        List<NamedExpression> selectItems = Lists.newArrayList();
-        String tableName = tableAlias != null ? tableAlias : Util.getTempTableDisplayName(targetTable.getName());
-        Expression setExpr = null;
-        for (Column column : targetTable.getFullSchema()) {
-            // if it sets sequence column in stream load phase, the sequence map column is null, we query it.
-            if (!column.isVisible() && !column.isSequenceColumn()) {
-                continue;
+        // Light schema changes rebuild the column metadata in place under the table write lock.
+        // Keep all schema reads consistent until the UPDATE logical plan has been constructed.
+        targetTable.readLock();
+        try {
+            Map<String, Expression> colNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER);
+            Map<String, Expression> partialUpdateColNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER);
+            for (EqualTo equalTo : assignments) {
+                List<String> colNameParts = ((UnboundSlot) equalTo.left()).getNameParts();
+                checkAssignmentColumn(ctx, colNameParts, this.nameParts, this.tableAlias);
+                colNameToExpression.put(colNameParts.get(colNameParts.size() - 1), equalTo.right());
+                partialUpdateColNameToExpression.put(colNameParts.get(colNameParts.size() - 1), equalTo.right());
             }
-            if (colNameToExpression.containsKey(column.getName())) {
-                Expression expr = colNameToExpression.get(column.getName());
-                // when updating the sequence map column, the real sequence column need to set with the same value.
-                boolean isSequenceMapColumn = targetTable.hasSequenceCol()
-                        && targetTable.getSequenceMapCol() != null
-                        && column.getName().equalsIgnoreCase(targetTable.getSequenceMapCol());
-                if (setExpr == null && isSequenceMapColumn) {
-                    setExpr = expr;
-                }
-                selectItems.add(expr instanceof UnboundSlot
-                        ? ((NamedExpression) expr)
-                        : new UnboundAlias(expr));
-                colNameToExpression.remove(column.getName());
-            } else {
-                if (column.isSequenceColumn() && setExpr != null) {
-                    selectItems.add(new UnboundAlias(setExpr, column.getName()));
-                } else if (column.hasOnUpdateDefaultValue()) {
-                    Expression defualtValueExpression =
-                            new NereidsParser().parseExpression(column.getOnUpdateDefaultValueSql());
-                    selectItems.add(new UnboundAlias(defualtValueExpression, column.getName()));
-                } else {
-                    selectItems.add(new UnboundSlot(tableName, column.getName()));
-                }
+            // check if any key in update clause
+            if (targetTable.getFullSchema().stream().filter(Column::isKey)
+                    .anyMatch(column -> partialUpdateColNameToExpression.containsKey(column.getName()))) {
+                throw new AnalysisException("Only value columns of unique table could be updated");
             }
-        }
-        if (!colNameToExpression.isEmpty()) {
-            throw new AnalysisException("unknown column in assignment list: "
-                    + String.join(", ", colNameToExpression.keySet()));
-        }
-
-        boolean isPartialUpdate = targetTable.getEnableUniqueKeyMergeOnWrite()
-                && selectItems.size() < targetTable.getColumns().size()
-                && targetTable.getSequenceCol() == null
-                && partialUpdateColNameToExpression.size() <= targetTable.getFullSchema().size() * 3 / 10
-                && !targetTable.isUniqKeyMergeOnWriteWithClusterKeys();
-
-        List<String> partialUpdateColNames = new ArrayList<>();
-        List<NamedExpression> partialUpdateSelectItems = new ArrayList<>();
-        if (isPartialUpdate) {
+            List<NamedExpression> selectItems = Lists.newArrayList();
+            String tableName = tableAlias != null ? tableAlias : Util.getTempTableDisplayName(targetTable.getName());
+            Expression setExpr = null;
             for (Column column : targetTable.getFullSchema()) {
-                Expression expr = new UnboundSlot(tableName, column.getName());
-                boolean existInExpr = false;
-                for (String colName : partialUpdateColNameToExpression.keySet()) {
-                    if (colName.equalsIgnoreCase(column.getName())) {
-                        expr = partialUpdateColNameToExpression.get(column.getName());
-                        existInExpr = true;
-                        break;
-                    }
+                // if it sets sequence column in stream load phase, the sequence map column is null, we query it.
+                if (!column.isVisible() && !column.isSequenceColumn()) {
+                    continue;
                 }
-                if (column.isKey() || existInExpr) {
-                    partialUpdateSelectItems.add(expr instanceof UnboundSlot
+                if (colNameToExpression.containsKey(column.getName())) {
+                    Expression expr = colNameToExpression.get(column.getName());
+                    // when updating the sequence map column, the real sequence column need to set with the same value.
+                    boolean isSequenceMapColumn = targetTable.hasSequenceCol()
+                            && targetTable.getSequenceMapCol() != null
+                            && column.getName().equalsIgnoreCase(targetTable.getSequenceMapCol());
+                    if (setExpr == null && isSequenceMapColumn) {
+                        setExpr = expr;
+                    }
+                    selectItems.add(expr instanceof UnboundSlot
                             ? ((NamedExpression) expr)
                             : new UnboundAlias(expr));
-                    partialUpdateColNames.add(column.getName());
+                    colNameToExpression.remove(column.getName());
+                } else {
+                    if (column.isSequenceColumn() && setExpr != null) {
+                        selectItems.add(new UnboundAlias(setExpr, column.getName()));
+                    } else if (column.hasOnUpdateDefaultValue()) {
+                        Expression defualtValueExpression =
+                                new NereidsParser().parseExpression(column.getOnUpdateDefaultValueSql());
+                        selectItems.add(new UnboundAlias(defualtValueExpression, column.getName()));
+                    } else {
+                        selectItems.add(new UnboundSlot(tableName, column.getName()));
+                    }
                 }
             }
-        }
+            if (!colNameToExpression.isEmpty()) {
+                throw new AnalysisException("unknown column in assignment list: "
+                        + String.join(", ", colNameToExpression.keySet()));
+            }
 
-        logicalQuery = new LogicalProject<>(isPartialUpdate ? partialUpdateSelectItems : selectItems, logicalQuery);
-        if (cte.isPresent()) {
-            logicalQuery = ((LogicalPlan) cte.get().withChildren(logicalQuery));
+            boolean isPartialUpdate = targetTable.getEnableUniqueKeyMergeOnWrite()
+                    && selectItems.size() < targetTable.getColumns().size()
+                    && targetTable.getSequenceCol() == null
+                    && partialUpdateColNameToExpression.size() <= targetTable.getFullSchema().size() * 3 / 10
+                    && !targetTable.isUniqKeyMergeOnWriteWithClusterKeys();
+
+            List<String> partialUpdateColNames = new ArrayList<>();
+            List<NamedExpression> partialUpdateSelectItems = new ArrayList<>();
+            if (isPartialUpdate) {
+                for (Column column : targetTable.getFullSchema()) {
+                    Expression expr = new UnboundSlot(tableName, column.getName());
+                    boolean existInExpr = false;
+                    for (String colName : partialUpdateColNameToExpression.keySet()) {
+                        if (colName.equalsIgnoreCase(column.getName())) {
+                            expr = partialUpdateColNameToExpression.get(column.getName());
+                            existInExpr = true;
+                            break;
+                        }
+                    }
+                    if (column.isKey() || existInExpr) {
+                        partialUpdateSelectItems.add(expr instanceof UnboundSlot
+                                ? ((NamedExpression) expr)
+                                : new UnboundAlias(expr));
+                        partialUpdateColNames.add(column.getName());
+                    }
+                }
+            }
+
+            logicalQuery = new LogicalProject<>(isPartialUpdate ? partialUpdateSelectItems : selectItems, logicalQuery);
+            if (cte.isPresent()) {
+                logicalQuery = ((LogicalPlan) cte.get().withChildren(logicalQuery));
+            }
+            // make UnboundTableSink
+            return UnboundTableSinkCreator.createUnboundTableSink(nameParts,
+                    isPartialUpdate ? partialUpdateColNames : ImmutableList.of(), ImmutableList.of(),
+                    false, ImmutableList.of(), isPartialUpdate, TPartialUpdateNewRowPolicy.APPEND,
+                    DMLCommandType.UPDATE, logicalQuery);
+        } finally {
+            targetTable.readUnlock();
         }
-        // make UnboundTableSink
-        return UnboundTableSinkCreator.createUnboundTableSink(nameParts,
-                isPartialUpdate ? partialUpdateColNames : ImmutableList.of(), ImmutableList.of(),
-                false, ImmutableList.of(), isPartialUpdate, TPartialUpdateNewRowPolicy.APPEND,
-                DMLCommandType.UPDATE, logicalQuery);
     }
 
     /**

fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/UpdateCommandTest.java

@@ -17,16 +17,34 @@
 
 package org.apache.doris.nereids.trees.plans;
 
+import org.apache.doris.catalog.Column;
+import org.apache.doris.catalog.OlapTable;
+import org.apache.doris.common.jmockit.Deencapsulation;
+import org.apache.doris.nereids.analyzer.UnboundTableSink;
+import org.apache.doris.nereids.exceptions.AnalysisException;
 import org.apache.doris.nereids.parser.NereidsParser;
 import org.apache.doris.nereids.trees.plans.commands.UpdateCommand;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 import org.apache.doris.nereids.util.PlanChecker;
 import org.apache.doris.nereids.util.PlanPatternMatchSupported;
+import org.apache.doris.nereids.util.RelationUtil;
 import org.apache.doris.utframe.TestWithFeService;
 
+import com.google.common.collect.ForwardingMap;
+import com.google.common.collect.ImmutableList;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 
+import java.util.AbstractCollection;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 public class UpdateCommandTest extends TestWithFeService implements PlanPatternMatchSupported {
     @Override
     public void runBeforeAll() throws Exception {
@@ -89,6 +107,104 @@ public class UpdateCommandTest extends TestWithFeService implements PlanPatternM
                 );
     }
 
+    @Test
+    public void testUpdatePlanWithConcurrentSchemaRebuild() throws Exception {
+        OlapTable table = (OlapTable) RelationUtil.getTable(
+                RelationUtil.getQualifierName(connectContext, ImmutableList.of("t1")),
+                connectContext.getEnv(), Optional.empty());
+        Assertions.assertTrue(table.getEnableUniqueKeyMergeOnWrite());
+        Map<String, Column> columns = Deencapsulation.getField(table, "nameToColumn");
+        ExecutorService schemaChange = Executors.newSingleThreadExecutor();
+        AtomicBoolean schemaRebuiltDuringCopy = new AtomicBoolean();
+        AtomicBoolean copyStarted = new AtomicBoolean();
+        // Keep the real TreeMap iterator, but try a schema rebuild after its first entry is read.
+        // Without the UPDATE read lock, clear/put in rebuildFullSchema invalidates this iterator.
+        Map<String, Column> instrumentedColumns = new ForwardingMap<String, Column>() {
+            @Override
+            protected Map<String, Column> delegate() {
+                return columns;
+            }
+
+            @Override
+            public Collection<Column> values() {
+                Collection<Column> values = columns.values();
+                return new AbstractCollection<Column>() {
+                    @Override
+                    public int size() {
+                        return values.size();
+                    }
+
+                    @Override
+                    public Iterator<Column> iterator() {
+                        Iterator<Column> iterator = values.iterator();
+                        return new Iterator<Column>() {
+                            @Override
+                            public boolean hasNext() {
+                                return iterator.hasNext();
+                            }
+
+                            @Override
+                            public Column next() {
+                                Column column = iterator.next();
+                                if (copyStarted.compareAndSet(false, true)) {
+                                    schemaRebuiltDuringCopy.set(Assertions.assertDoesNotThrow(() ->
+                                            schemaChange.submit(() -> tryRebuildSchema(table))
+                                                    .get(10, TimeUnit.SECONDS)));
+                                }
+                                return column;
+                            }
+                        };
+                    }
+                };
+            }
+        };
+        Deencapsulation.setField(table, "nameToColumn", instrumentedColumns);
+        try {
+            UpdateCommand command = (UpdateCommand) new NereidsParser()
+                    .parseSingle("update t1 set v1 = v1 + 1 where k1 = 1");
+            LogicalPlan plan = command.completeQueryPlan(connectContext, command.getLogicalQuery());
+            Assertions.assertTrue(((UnboundTableSink<?>) plan).isPartialUpdate());
+            Assertions.assertTrue(copyStarted.get(), "The test must reach the column-map copy");
+            Assertions.assertFalse(schemaRebuiltDuringCopy.get(), "Schema changes must wait for UPDATE planning");
+            Assertions.assertTrue(schemaChange.submit(() -> tryRebuildSchema(table)).get(10, TimeUnit.SECONDS),
+                    "UPDATE planning must release its read lock before returning");
+            PlanChecker.from(connectContext, plan).analyze(plan).rewrite();
+        } finally {
+            schemaChange.shutdownNow();
+            Deencapsulation.setField(table, "nameToColumn", columns);
+        }
+    }
+
+    @Test
+    public void testUpdatePlanReleasesReadLockOnAnalysisFailure() throws Exception {
+        OlapTable table = (OlapTable) RelationUtil.getTable(
+                RelationUtil.getQualifierName(connectContext, ImmutableList.of("t1")),
+                connectContext.getEnv(), Optional.empty());
+        ExecutorService schemaChange = Executors.newSingleThreadExecutor();
+        try {
+            UpdateCommand command = (UpdateCommand) new NereidsParser()
+                    .parseSingle("update t1 set k1 = 2 where k1 = 1");
+            Assertions.assertThrows(AnalysisException.class,
+                    () -> command.completeQueryPlan(connectContext, command.getLogicalQuery()));
+            Assertions.assertTrue(schemaChange.submit(() -> tryRebuildSchema(table)).get(10, TimeUnit.SECONDS),
+                    "Failed UPDATE planning must release its read lock");
+        } finally {
+            schemaChange.shutdownNow();
+        }
+    }
+
+    private static boolean tryRebuildSchema(OlapTable table) {
+        if (!table.tryWriteLock(0, TimeUnit.SECONDS)) {
+            return false;
+        }
+        try {
+            table.rebuildFullSchema();
+            return true;
+        } finally {
+            table.writeUnlock();
+        }
+    }
+
     @Test
     public void testFromClauseUpdate() {
         String sql = "update t1 a set v1 = t2.v1 + 2, v2 = a.v1 * 2 "