diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java index 06237ec19e7..adee7e5bb3d 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/UpdateCommand.java +++ b/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 colNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); - Map partialUpdateColNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); - for (EqualTo equalTo : assignments) { - List 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 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 colNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + Map partialUpdateColNameToExpression = Maps.newTreeMap(String.CASE_INSENSITIVE_ORDER); + for (EqualTo equalTo : assignments) { + List 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 partialUpdateColNames = new ArrayList<>(); - List partialUpdateSelectItems = new ArrayList<>(); - if (isPartialUpdate) { + List 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 partialUpdateColNames = new ArrayList<>(); + List 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); } /** diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/UpdateCommandTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/UpdateCommandTest.java index 4151115a2c8..2a38b488bf7 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/trees/plans/UpdateCommandTest.java +++ b/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 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 instrumentedColumns = new ForwardingMap() { + @Override + protected Map delegate() { + return columns; + } + + @Override + public Collection values() { + Collection values = columns.values(); + return new AbstractCollection() { + @Override + public int size() { + return values.size(); + } + + @Override + public Iterator iterator() { + Iterator iterator = values.iterator(); + return new Iterator() { + @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 "