Skip to content

Commit e119109

Browse files
wuwenchiwuwenchi
authored
[opt](iceberg) support insert to branch (#52751)
### What problem does this PR solve? Related Issue: #48285 Problem Summary: support insert to branch, like : ```sql insert into tb@branch(br1); ```
1 parent b8a4915 commit e119109

15 files changed

Lines changed: 910 additions & 40 deletions

File tree

‎fe/fe-core/src/main/antlr4/org/apache/doris/nereids/DorisParser.g4‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,9 +128,13 @@ constraintStatement
128128
| SHOW CONSTRAINTS FROM table=multipartIdentifier #showConstraint
129129
;
130130

131+
optSpecBranch
132+
: ATSIGN BRANCH LEFT_PAREN name=identifier RIGHT_PAREN
133+
;
134+
131135
supportedDmlStatement
132136
: explain? cte? INSERT (INTO | OVERWRITE TABLE)
133-
(tableName=multipartIdentifier | DORIS_INTERNAL_TABLE_ID LEFT_PAREN tableId=INTEGER_VALUE RIGHT_PAREN)
137+
(tableName=multipartIdentifier (optSpecBranch)? | DORIS_INTERNAL_TABLE_ID LEFT_PAREN tableId=INTEGER_VALUE RIGHT_PAREN)
134138
partitionSpec? // partition define
135139
(WITH LABEL labelName=identifier)? cols=identifierList? // label and columns define
136140
(LEFT_BRACKET hints=identifierSeq RIGHT_BRACKET)? // hint define

‎fe/fe-core/src/main/java/org/apache/doris/datasource/iceberg/IcebergTransaction.java‎

Lines changed: 39 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
import org.apache.doris.datasource.ExternalTable;
2525
import org.apache.doris.datasource.NameMapping;
2626
import org.apache.doris.datasource.iceberg.helper.IcebergWriterHelper;
27-
import org.apache.doris.nereids.trees.plans.commands.insert.BaseExternalTableInsertCommandContext;
27+
import org.apache.doris.nereids.trees.plans.commands.insert.IcebergInsertCommandContext;
2828
import org.apache.doris.nereids.trees.plans.commands.insert.InsertCommandContext;
2929
import org.apache.doris.thrift.TIcebergCommitData;
3030
import org.apache.doris.thrift.TUpdateMode;
@@ -38,6 +38,7 @@
3838
import org.apache.iceberg.OverwriteFiles;
3939
import org.apache.iceberg.PartitionSpec;
4040
import org.apache.iceberg.ReplacePartitions;
41+
import org.apache.iceberg.SnapshotRef;
4142
import org.apache.iceberg.Table;
4243
import org.apache.iceberg.io.CloseableIterable;
4344
import org.apache.iceberg.io.WriteResult;
@@ -61,6 +62,9 @@ public class IcebergTransaction implements Transaction {
6162
private org.apache.iceberg.Transaction transaction;
6263
private final List<TIcebergCommitData> commitDataList = Lists.newArrayList();
6364

65+
private IcebergInsertCommandContext insertCtx;
66+
private String branchName;
67+
6468
public IcebergTransaction(IcebergMetadataOps ops) {
6569
this.ops = ops;
6670
}
@@ -71,29 +75,43 @@ public void updateIcebergCommitData(List<TIcebergCommitData> commitDataList) {
7175
}
7276
}
7377

74-
public void beginInsert(ExternalTable dorisTable) throws UserException {
78+
public void beginInsert(ExternalTable dorisTable, Optional<InsertCommandContext> ctx) throws UserException {
79+
ctx.ifPresent(c -> this.insertCtx = (IcebergInsertCommandContext) c);
7580
try {
7681
ops.getPreExecutionAuthenticator().execute(() -> {
7782
// create and start the iceberg transaction
7883
this.table = IcebergUtils.getIcebergTable(dorisTable);
84+
// check branch
85+
if (insertCtx != null && insertCtx.getBranchName().isPresent()) {
86+
this.branchName = insertCtx.getBranchName().get();
87+
SnapshotRef branchRef = table.refs().get(branchName);
88+
if (branchRef == null) {
89+
throw new RuntimeException(branchName + " is not founded in " + dorisTable.getName());
90+
} else if (!branchRef.isBranch()) {
91+
throw new RuntimeException(
92+
branchName
93+
+ " is a tag, not a branch. Tags cannot be targets for producing snapshots");
94+
}
95+
}
7996
this.transaction = table.newTransaction();
8097
});
8198
} catch (Exception e) {
82-
throw new UserException("Failed to begin insert for iceberg table " + dorisTable.getName(), e);
99+
throw new UserException("Failed to begin insert for iceberg table " + dorisTable.getName()
100+
+ "because: " + e.getMessage(), e);
83101
}
84102

85103
}
86104

87-
public void finishInsert(NameMapping nameMapping, Optional<InsertCommandContext> insertCtx) {
105+
public void finishInsert(NameMapping nameMapping) {
88106
if (LOG.isDebugEnabled()) {
89107
LOG.info("iceberg table {} insert table finished!", nameMapping.getFullLocalName());
90108
}
91109
try {
92110
ops.getPreExecutionAuthenticator().execute(() -> {
93111
//create and start the iceberg transaction
94112
TUpdateMode updateMode = TUpdateMode.APPEND;
95-
if (insertCtx.isPresent()) {
96-
updateMode = ((BaseExternalTableInsertCommandContext) insertCtx.get()).isOverwrite()
113+
if (insertCtx != null) {
114+
updateMode = insertCtx.isOverwrite()
97115
? TUpdateMode.OVERWRITE
98116
: TUpdateMode.APPEND;
99117
}
@@ -146,6 +164,9 @@ public long getUpdateCnt() {
146164
private void commitAppendTxn(List<WriteResult> pendingResults) {
147165
// commit append files.
148166
AppendFiles appendFiles = transaction.newAppend().scanManifestsWith(ops.getThreadPoolWithPreAuth());
167+
if (branchName != null) {
168+
appendFiles = appendFiles.toBranch(branchName);
169+
}
149170
for (WriteResult result : pendingResults) {
150171
Preconditions.checkState(result.referencedDataFiles().length == 0,
151172
"Should have no referenced data files for append.");
@@ -161,11 +182,14 @@ private void commitReplaceTxn(List<WriteResult> pendingResults) {
161182
// 1. if dst_tb is a partitioned table, it will return directly.
162183
// 2. if dst_tb is an unpartitioned table, the `dst_tb` table will be emptied.
163184
if (!transaction.table().spec().isPartitioned()) {
164-
OverwriteFiles overwriteFiles = transaction
165-
.newOverwrite()
166-
.scanManifestsWith(ops.getThreadPoolWithPreAuth());
185+
OverwriteFiles overwriteFiles = transaction.newOverwrite();
186+
if (branchName != null) {
187+
overwriteFiles = overwriteFiles.toBranch(branchName);
188+
}
189+
overwriteFiles = overwriteFiles.scanManifestsWith(ops.getThreadPoolWithPreAuth());
167190
try (CloseableIterable<FileScanTask> fileScanTasks = table.newScan().planFiles()) {
168-
fileScanTasks.forEach(f -> overwriteFiles.deleteFile(f.file()));
191+
OverwriteFiles finalOverwriteFiles = overwriteFiles;
192+
fileScanTasks.forEach(f -> finalOverwriteFiles.deleteFile(f.file()));
169193
} catch (IOException e) {
170194
throw new RuntimeException(e);
171195
}
@@ -175,9 +199,11 @@ private void commitReplaceTxn(List<WriteResult> pendingResults) {
175199
}
176200

177201
// commit replace partitions
178-
ReplacePartitions appendPartitionOp = transaction
179-
.newReplacePartitions()
180-
.scanManifestsWith(ops.getThreadPoolWithPreAuth());
202+
ReplacePartitions appendPartitionOp = transaction.newReplacePartitions();
203+
if (branchName != null) {
204+
appendPartitionOp = appendPartitionOp.toBranch(branchName);
205+
}
206+
appendPartitionOp = appendPartitionOp.scanManifestsWith(ops.getThreadPoolWithPreAuth());
181207
for (WriteResult result : pendingResults) {
182208
Preconditions.checkState(result.referencedDataFiles().length == 0,
183209
"Should have no referenced data files.");

‎fe/fe-core/src/main/java/org/apache/doris/nereids/parser/LogicalPlanBuilder.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1273,6 +1273,10 @@ public LogicalPlan visitInsertTable(InsertTableContext ctx) {
12731273
throw new ParseException("tableName and tableId cannot both be null");
12741274
}
12751275
Optional<String> labelName = (ctx.labelName == null) ? Optional.empty() : Optional.of(ctx.labelName.getText());
1276+
Optional<String> branchName = Optional.empty();
1277+
if (ctx.optSpecBranch() != null) {
1278+
branchName = Optional.of(ctx.optSpecBranch().name.getText());
1279+
}
12761280
List<String> colNames = ctx.cols == null ? ImmutableList.of() : visitIdentifierList(ctx.cols);
12771281
// TODO visit partitionSpecCtx
12781282
LogicalPlan plan = visitQuery(ctx.query());
@@ -1301,7 +1305,7 @@ public LogicalPlan visitInsertTable(InsertTableContext ctx) {
13011305
}
13021306
LogicalPlan command;
13031307
if (isOverwrite) {
1304-
command = new InsertOverwriteTableCommand(sink, labelName, cte);
1308+
command = new InsertOverwriteTableCommand(sink, labelName, cte, branchName);
13051309
} else {
13061310
if (ConnectContext.get() != null && ConnectContext.get().isTxnModel()
13071311
&& sink.child() instanceof InlineTable
@@ -1311,7 +1315,7 @@ public LogicalPlan visitInsertTable(InsertTableContext ctx) {
13111315
// Now handle it as `insert into select`(a separate load job), should fix it as the legacy.
13121316
command = new BatchInsertIntoTableCommand(sink);
13131317
} else {
1314-
command = new InsertIntoTableCommand(sink, labelName, Optional.empty(), cte);
1318+
command = new InsertIntoTableCommand(sink, labelName, Optional.empty(), cte, true, branchName);
13151319
}
13161320
}
13171321
return withExplain(command, ctx.explain());

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,8 +183,8 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
183183
ImmutableList.of(), ImmutableList.of(), ImmutableList.of(), query);
184184
try {
185185
if (!FeConstants.runningUnitTest) {
186-
new InsertIntoTableCommand(query, Optional.empty(), Optional.empty(), Optional.empty()).run(
187-
ctx, executor);
186+
new InsertIntoTableCommand(query, Optional.empty(), Optional.empty(),
187+
Optional.empty(), true, Optional.empty()).run(ctx, executor);
188188
}
189189
if (ctx.getState().getStateType() == MysqlStateType.ERR) {
190190
handleFallbackFailedCtas(ctx);

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
5454
}
5555
// NOTE: delete from using command is executed as insert command, so txn insert can support it
5656
new InsertIntoTableCommand(completeQueryPlan(ctx, logicalQuery), Optional.empty(), Optional.empty(),
57-
Optional.empty()).run(ctx, executor);
57+
Optional.empty(), true, Optional.empty()).run(ctx, executor);
5858
}
5959

6060
@Override

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ public UpdateCommand(List<String> nameParts, @Nullable String tableAlias, List<E
9797
public void run(ConnectContext ctx, StmtExecutor executor) throws Exception {
9898
// NOTE: update command is executed as insert command, so txn insert can support it
9999
new InsertIntoTableCommand(completeQueryPlan(ctx, logicalQuery), Optional.empty(), Optional.empty(),
100-
Optional.empty()).run(ctx, executor);
100+
Optional.empty(), true, Optional.empty()).run(ctx, executor);
101101
}
102102

103103
/**

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ public class UpdateMvByPartitionCommand extends InsertOverwriteTableCommand {
8585
private static final Logger LOG = LogManager.getLogger(UpdateMvByPartitionCommand.class);
8686

8787
private UpdateMvByPartitionCommand(LogicalPlan logicalQuery) {
88-
super(logicalQuery, Optional.empty(), Optional.empty());
88+
super(logicalQuery, Optional.empty(), Optional.empty(), Optional.empty());
8989
}
9090

9191
@Override

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/IcebergInsertCommandContext.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,19 @@
1717

1818
package org.apache.doris.nereids.trees.plans.commands.insert;
1919

20+
import java.util.Optional;
21+
2022
/**
2123
* For iceberg External Table
2224
*/
2325
public class IcebergInsertCommandContext extends BaseExternalTableInsertCommandContext {
26+
private Optional<String> branchName = Optional.empty();
27+
28+
public Optional<String> getBranchName() {
29+
return branchName;
30+
}
31+
32+
public void setBranchName(Optional<String> branchName) {
33+
this.branchName = branchName;
34+
}
2435
}

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/IcebergInsertExecutor.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ public IcebergInsertExecutor(ConnectContext ctx, IcebergExternalTable table,
4949
@Override
5050
protected void beforeExec() throws UserException {
5151
IcebergTransaction transaction = (IcebergTransaction) transactionManager.getTransaction(txnId);
52-
transaction.beginInsert((IcebergExternalTable) table);
52+
transaction.beginInsert((IcebergExternalTable) table, insertCtx);
5353
}
5454

5555
@Override
@@ -60,7 +60,7 @@ protected void doBeforeCommit() throws UserException {
6060
dorisTable.getRemoteDbName(), dorisTable.getRemoteName());
6161
IcebergTransaction transaction = (IcebergTransaction) transactionManager.getTransaction(txnId);
6262
this.loadedRows = transaction.getUpdateCnt();
63-
transaction.finishInsert(nameMapping, insertCtx);
63+
transaction.finishInsert(nameMapping);
6464
}
6565

6666
@Override

‎fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/commands/insert/InsertIntoTableCommand.java‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ public class InsertIntoTableCommand extends Command implements NeedAuditEncrypti
9999
private LogicalPlan originLogicalQuery;
100100
private Optional<LogicalPlan> logicalQuery;
101101
private Optional<String> labelName;
102+
private Optional<String> branchName;
102103
/**
103104
* When source it's from job scheduler,it will be set.
104105
*/
@@ -111,22 +112,23 @@ public class InsertIntoTableCommand extends Command implements NeedAuditEncrypti
111112

112113
public InsertIntoTableCommand(LogicalPlan logicalQuery, Optional<String> labelName,
113114
Optional<InsertCommandContext> insertCtx, Optional<LogicalPlan> cte) {
114-
this(logicalQuery, labelName, insertCtx, cte, true);
115+
this(logicalQuery, labelName, insertCtx, cte, true, Optional.empty());
115116
}
116117

117118
/**
118119
* constructor
119120
*/
120121
public InsertIntoTableCommand(LogicalPlan logicalQuery, Optional<String> labelName,
121122
Optional<InsertCommandContext> insertCtx, Optional<LogicalPlan> cte,
122-
boolean needNormalizePlan) {
123+
boolean needNormalizePlan, Optional<String> branchName) {
123124
super(PlanType.INSERT_INTO_TABLE_COMMAND);
124125
this.originLogicalQuery = Objects.requireNonNull(logicalQuery, "logicalQuery should not be null");
125126
this.labelName = Objects.requireNonNull(labelName, "labelName should not be null");
126127
this.logicalQuery = Optional.empty();
127128
this.insertCtx = insertCtx;
128129
this.cte = cte;
129130
this.needNormalizePlan = needNormalizePlan;
131+
this.branchName = branchName;
130132
}
131133

132134
/**
@@ -140,7 +142,8 @@ public InsertIntoTableCommand(InsertIntoTableCommand command, PlanType planType)
140142
this.insertCtx = command.insertCtx;
141143
this.cte = command.cte;
142144
this.jobId = command.jobId;
143-
this. needNormalizePlan = true;
145+
this.needNormalizePlan = true;
146+
this.branchName = command.branchName;
144147
}
145148

146149
public LogicalPlan getLogicalQuery() {
@@ -327,6 +330,11 @@ private ExecutorFactory selectInsertExecutorFactory(
327330
String label = this.labelName.orElse(
328331
ctx.isTxnModel() ? null : String.format("label_%x_%x", ctx.queryId().hi, ctx.queryId().lo));
329332

333+
// check branch
334+
if (branchName.isPresent() && !(physicalSink instanceof PhysicalIcebergTableSink)) {
335+
throw new AnalysisException("Only support insert data into iceberg table's branch");
336+
}
337+
330338
if (physicalSink instanceof PhysicalOlapTableSink) {
331339
boolean emptyInsert = childIsEmptyRelation(physicalSink);
332340
OlapTable olapTable = (OlapTable) targetTableIf;
@@ -387,12 +395,16 @@ private ExecutorFactory selectInsertExecutorFactory(
387395
} else if (physicalSink instanceof PhysicalIcebergTableSink) {
388396
boolean emptyInsert = childIsEmptyRelation(physicalSink);
389397
IcebergExternalTable icebergExternalTable = (IcebergExternalTable) targetTableIf;
398+
IcebergInsertCommandContext icebergInsertCtx = insertCtx
399+
.map(insertCommandContext -> (IcebergInsertCommandContext) insertCommandContext)
400+
.orElseGet(IcebergInsertCommandContext::new);
401+
branchName.ifPresent(notUsed -> icebergInsertCtx.setBranchName(branchName));
390402
return ExecutorFactory.from(
391403
planner,
392404
dataSink,
393405
physicalSink,
394406
() -> new IcebergInsertExecutor(ctx, icebergExternalTable, label, planner,
395-
Optional.of(insertCtx.orElse((new BaseExternalTableInsertCommandContext()))),
407+
Optional.of(icebergInsertCtx),
396408
emptyInsert
397409
)
398410
);

0 commit comments

Comments
 (0)