ConnectorRewriteDriver.java
// 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.
package org.apache.doris.nereids.trees.plans.commands.execute;
import org.apache.doris.catalog.Env;
import org.apache.doris.common.Status;
import org.apache.doris.common.UserException;
import org.apache.doris.connector.spi.ConnectorMetadata;
import org.apache.doris.connector.spi.ConnectorSession;
import org.apache.doris.connector.spi.DorisConnectorException;
import org.apache.doris.connector.spi.handle.ConnectorTableHandle;
import org.apache.doris.connector.spi.handle.ConnectorTransaction;
import org.apache.doris.connector.spi.handle.RewriteCapableTransaction;
import org.apache.doris.connector.spi.procedure.ConnectorProcedureOps;
import org.apache.doris.connector.spi.procedure.ConnectorProcedureResult;
import org.apache.doris.connector.spi.procedure.ConnectorRewriteGroup;
import org.apache.doris.connector.spi.procedure.ConnectorRewriteStatistics;
import org.apache.doris.connector.spi.pushdown.ConnectorPredicate;
import org.apache.doris.datasource.ExternalTable;
import org.apache.doris.datasource.plugin.PluginDrivenExternalCatalog;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.scheduler.exception.JobException;
import org.apache.doris.transaction.PluginDrivenTransactionManager;
import com.google.common.collect.Lists;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
/**
* Engine-neutral driver for a distributed {@code rewrite_data_files} (compaction), the post-flip
* counterpart of the legacy per-source rewrite executor. It orchestrates the read/write distribution; the
* connector owns the file-selection / bin-pack / commit decisions AND the shape of the result row, behind
* neutral SPIs (no {@code instanceof} on a connector type, no source-specific types).
*
* <p>Flow: (0) ask the connector to plan N bin-packed groups ({@link ConnectorProcedureOps#planRewrite}); (1)
* open ONE shared connector transaction; (2) run one {@code INSERT-SELECT} per group concurrently, each
* scoped to its files and bound to the shared transaction; (3) register the union of source files to remove
* (AFTER the groups began the transaction so the table + OCC snapshot are loaded); (4) commit once; (5) read
* the added-files count post-commit and ask the connector to render the result row.</p>
*
* <p><b>R6 scope.</b> No-WHERE rewrite only (WHERE lowering is a later step). Output file SIZING ��� the
* per-group {@code target-file-size}/parallelism that the legacy task threaded via an iceberg session var ���
* is deferred (see the rewrite-output-sizing follow-up): every group GATHERs to a single writer via the
* rewrite sink flag, which is correct but not size-tuned. This only affects real BE writes (exercised at the
* flip rehearsal), not the dormant pre-flip / mock path.</p>
*/
public class ConnectorRewriteDriver {
private static final Logger LOG = LogManager.getLogger(ConnectorRewriteDriver.class);
private final ConnectContext ctx;
private final ExternalTable table;
private final PluginDrivenExternalCatalog catalog;
private final ConnectorMetadata metadata;
private final ConnectorProcedureOps procedureOps;
private final ConnectorSession session;
private final ConnectorTableHandle tableHandle;
private final String procedureName;
private final Map<String, String> properties;
private final List<String> partitionNames;
// The engine-lowered WHERE restricting which files to rewrite, or null when there is no WHERE. Passed
// straight through to the connector's planRewrite (the connector scopes the rewrite to the matching files).
private final ConnectorPredicate whereCondition;
// The outer statement executor, or null in tests/standalone use. Cancellation of the outer statement is
// handed to this driver so the distributed rewrite does not commit after the statement is terminal.
private final StmtExecutor owner;
// Serializes the outer cancellation handoff with the source-registration/final-commit decision.
private final Object cancelLock = new Object();
private volatile Status outerCancelReason;
private volatile List<ConnectorRewriteGroupTask> submittedGroups = Collections.emptyList();
private final AtomicBoolean cancelHandoffInstalled = new AtomicBoolean(false);
/**
* Builds a driver bound to one {@code ALTER TABLE ... EXECUTE rewrite_data_files} invocation; all of
* these are already resolved by {@code ConnectorExecuteAction} (the only caller).
*/
public ConnectorRewriteDriver(ConnectContext ctx, ExternalTable table, PluginDrivenExternalCatalog catalog,
ConnectorMetadata metadata, ConnectorProcedureOps procedureOps, ConnectorSession session,
ConnectorTableHandle tableHandle, String procedureName, Map<String, String> properties,
List<String> partitionNames, ConnectorPredicate whereCondition, StmtExecutor owner) {
this.ctx = ctx;
this.table = table;
this.catalog = catalog;
this.metadata = metadata;
this.procedureOps = procedureOps;
this.session = session;
this.tableHandle = tableHandle;
this.procedureName = procedureName;
this.properties = properties;
this.partitionNames = partitionNames;
this.whereCondition = whereCondition;
this.owner = owner;
}
/**
* Installs the sticky outer-executor cancellation handoff: pick up a cancellation that landed before the
* driver existed, then register for the ones that arrive while the rewrite runs.
*/
private void installCancelHandoff() {
StmtExecutor outer = this.owner;
if (outer == null) {
return;
}
Status prior = outer.getPendingCancelReason();
if (prior != null && !prior.ok()) {
synchronized (cancelLock) {
if (outerCancelReason == null) {
outerCancelReason = prior;
}
}
}
outer.setCancelDelegate(this::onOuterCancel);
cancelHandoffInstalled.set(true);
}
private void clearCancelHandoff() {
if (cancelHandoffInstalled.compareAndSet(true, false) && owner != null) {
owner.clearCancelDelegate();
}
}
/**
* Runs on the cancelling thread. Records the sticky reason and stops the live groups; the owner thread
* drains them (with a shared budget) before it rolls the shared transaction back.
*/
private void onOuterCancel(Status reason) {
synchronized (cancelLock) {
if (outerCancelReason == null) {
outerCancelReason = reason;
}
}
for (ConnectorRewriteGroupTask task : submittedGroups) {
try {
task.cancel();
} catch (Exception e) {
LOG.warn("Failed to cancel rewrite task {}: {}", task.getId(), e.getMessage());
}
}
}
/**
* Sticky: picks up the owner's first terminal reason the first time it is observed. Polled by the wait
* loop and re-checked at the register/commit decision, so it also covers a cancellation whose delegate
* notification was missed.
*/
private boolean isOuterCancelled() {
StmtExecutor outer = this.owner;
if (outer != null && outerCancelReason == null) {
Status current = outer.getPendingCancelReason();
if (current != null && !current.ok()) {
synchronized (cancelLock) {
if (outerCancelReason == null) {
outerCancelReason = current;
}
}
}
}
return outerCancelReason != null;
}
private UserException cancelledException() {
Status reason = outerCancelReason;
return new UserException("Rewrite is cancelled: "
+ (reason == null ? "statement terminated" : reason.getErrorMsg()));
}
/**
* Runs the distributed rewrite and returns the single-row result the engine wraps into a ResultSet.
*/
public ConnectorProcedureResult run() throws UserException {
// STEP 0: ask the connector to plan the bin-packed groups, scoped by the lowered WHERE (null = no WHERE).
List<ConnectorRewriteGroup> groups;
try {
groups = procedureOps.planRewrite(session, tableHandle, procedureName, properties, whereCondition,
partitionNames);
} catch (DorisConnectorException e) {
throw new UserException(e.getMessage(), e);
}
if (groups == null || groups.isEmpty()) {
// Nothing to rewrite: skip the transaction entirely and let the connector render its all-zero
// row (legacy parity). There is no transaction on this path, which is why buildRewriteResult
// must render locally.
return procedureOps.buildRewriteResult(procedureName,
new ConnectorRewriteStatistics(0, 0, 0L, 0));
}
// STEP 1: open ONE shared connector transaction for all groups.
PluginDrivenTransactionManager txnManager =
(PluginDrivenTransactionManager) catalog.getTransactionManager();
ConnectorTransaction connectorTx;
long txnId;
try {
connectorTx = metadata.beginTransaction(session, tableHandle);
txnId = txnManager.begin(connectorTx);
} catch (DorisConnectorException e) {
throw new UserException(e.getMessage(), e);
}
// rewrite_data_files is a rewrite-capable-connector-only procedure; the transaction MUST carry the
// narrow RewriteCapableTransaction capability. Fail loud with a type mismatch here rather than
// discovering it as a runtime UnsupportedOperationException mid-rewrite (only iceberg qualifies today).
if (!(connectorTx instanceof RewriteCapableTransaction)) {
txnManager.rollback(txnId);
throw new UserException("Connector transaction does not support rewrite_data_files: "
+ connectorTx.getClass().getSimpleName());
}
RewriteCapableTransaction rewriteTx = (RewriteCapableTransaction) connectorTx;
installCancelHandoff();
try {
try {
// STEP 2: run one INSERT-SELECT per group concurrently, all sharing the transaction.
runGroups(groups, txnId, connectorTx);
// STEP 3: register the UNION of every group's source data files in a SINGLE call. The connector
// re-derives them from the table at the pinned OCC snapshot with ONE planFiles() scan; the former
// per-group loop repeated that full-table scan once per group (G groups = G+1 scans). Ordering is
// unchanged ��� still AFTER the groups ran (so the first group's write loaded the table + pinned the
// OCC snapshot that the connector re-derives against) and BEFORE commit (which consumes the
// registered files in the RewriteFiles op). Every per-group call scanned the SAME pinned snapshot, so
// one union scan is equivalent; the connector's registration accumulates and dedups by path, and the
// planner emits path-DISJOINT groups (iceberg: planFiles() yields one task per data file, bin-packed
// into disjoint groups), so the union reconstructs exactly the per-group calls' accumulated file set.
synchronized (cancelLock) {
if (isOuterCancelled()) {
throw cancelledException();
}
rewriteTx.registerRewriteSourceFiles(unionSourceFilePaths(groups));
}
} catch (Exception e) {
txnManager.rollback(txnId);
if (e instanceof UserException) {
throw (UserException) e;
}
throw new UserException("Failed to rewrite data files: " + e.getMessage(), e);
}
// STEP 4: commit once, serialized on cancelLock with the outer cancellation handoff. Cancellation
// that wins the lock first rolls back instead of committing; cancellation that arrives while the
// critical section runs is linearized after the commit. The manager deregisters the transaction on
// both success and failure, so a failed commit needs no rollback ��� surface it directly.
synchronized (cancelLock) {
if (isOuterCancelled()) {
txnManager.rollback(txnId);
throw cancelledException();
}
txnManager.commit(txnId);
}
} finally {
clearCancelHandoff();
}
// The rewrite is committed. Persist follower replay identity and refresh leader caches before the
// post-commit statistics and result construction below, which can fail independently of the mutation.
Env.getCurrentEnv().getRefreshManager().refreshTableAfterExternalMutation(table);
// STEP 5: post-commit statistics. The added-files count is only valid after commit (it is
// materialized from the BE commit fragments during commit); the other three are summed from the
// planning groups (the connector exposes them on each ConnectorRewriteGroup).
int addedDataFilesCount = rewriteTx.getRewriteAddedDataFilesCount();
int rewrittenDataFilesCount = groups.stream().mapToInt(ConnectorRewriteGroup::getDataFileCount).sum();
long rewrittenBytesCount = groups.stream().mapToLong(ConnectorRewriteGroup::getTotalSizeBytes).sum();
int removedDeleteFilesCount = groups.stream().mapToInt(ConnectorRewriteGroup::getDeleteFileCount).sum();
// The connector names and types its own result columns; the engine only reports what it ran.
return procedureOps.buildRewriteResult(procedureName, new ConnectorRewriteStatistics(
rewrittenDataFilesCount, addedDataFilesCount, rewrittenBytesCount, removedDeleteFilesCount));
}
/**
* Unions every group's source data-file paths into one dedup'd set, so the connector re-derives them all in
* a single {@code planFiles()} scan (STEP 3) instead of one scan per group. Bin-packed groups are
* path-disjoint so this is a straight union; the connector's own per-path dedup keeps the registered file set
* exact regardless. Package-visible for unit testing (the full distributed STEP 3 needs a live cluster).
*/
static Set<String> unionSourceFilePaths(List<ConnectorRewriteGroup> groups) {
Set<String> sourceFilePaths = new HashSet<>();
for (ConnectorRewriteGroup group : groups) {
sourceFilePaths.addAll(group.getDataFilePaths());
}
return sourceFilePaths;
}
private void runGroups(List<ConnectorRewriteGroup> groups, long txnId, ConnectorTransaction connectorTx)
throws UserException {
List<ConnectorRewriteGroupTask> tasks = Lists.newArrayList();
RewriteResultCollector collector = new RewriteResultCollector(groups.size(), tasks);
for (ConnectorRewriteGroup group : groups) {
ConnectorRewriteGroupTask task = new ConnectorRewriteGroupTask(group, txnId, connectorTx, table, ctx,
new ConnectorRewriteGroupTask.RewriteResultCallback() {
@Override
public void onTaskCompleted(Long taskId) {
collector.onTaskCompleted(taskId);
}
@Override
public void onTaskFailed(Long taskId, Exception error) {
collector.onTaskFailed(taskId, error);
}
});
tasks.add(task);
}
List<ConnectorRewriteGroupTask> submitted = Lists.newArrayList();
try {
for (ConnectorRewriteGroupTask task : tasks) {
Env.getCurrentEnv().getTransientTaskManager().addMemoryTask(task);
submitted.add(task);
}
} catch (JobException e) {
// Groups submitted before the failing call already bind the shared transaction; drain them so the
// caller never rolls that transaction back while a live group still reports into it.
drain(submitted, drainBudgetNanos());
throw new UserException("Failed to submit rewrite tasks: " + e.getMessage(), e);
}
submittedGroups = submitted;
long maxWaitTime = ctx.getSessionVariable().getInsertTimeoutS();
long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(maxWaitTime);
boolean completed = false;
while (!completed) {
// A TIMEOUT/KILL on the outer statement only reaches this driver through the handoff; poll it so
// live groups are stopped instead of finishing and committing after the statement is terminal.
if (isOuterCancelled()) {
drain(submitted, drainBudgetNanos());
throw cancelledException();
}
long remainingMs = TimeUnit.NANOSECONDS.toMillis(deadlineNanos - System.nanoTime());
if (remainingMs <= 0) {
completed = collector.isDone();
break;
}
try {
completed = collector.await(Math.min(200L, remainingMs), TimeUnit.MILLISECONDS);
} catch (InterruptedException e) {
// The interrupt flag was cleared by the exception, so the drain below can still wait.
drain(submitted, drainBudgetNanos());
Thread.currentThread().interrupt();
throw new UserException("Wait for rewrite tasks completion was interrupted", e);
}
}
if (!completed) {
// The owner gave up waiting: stop every live group and wait for its terminal callback so no group
// is still reporting into the shared transaction when the caller rolls it back.
drain(submitted, drainBudgetNanos());
throw new UserException("Rewrite tasks did not complete within timeout");
}
if (collector.getFirstError() != null) {
throw new UserException("Some rewrite tasks failed: " + collector.getFirstError().getMessage(),
collector.getFirstError());
}
}
private long drainBudgetNanos() {
return TimeUnit.SECONDS.toNanos(Math.max(1, ctx.getSessionVariable().getInsertTimeoutS()));
}
/**
* Cancels every submitted group and waits (bounded) for its terminal callback within ONE shared budget,
* so a shared transaction is never rolled back while a live group still has commit data flowing into it,
* and G never-terminal groups cannot multiply the drain deadline into G * insert_timeout.
*/
private void drain(List<ConnectorRewriteGroupTask> submitted, long budgetNanos) {
for (ConnectorRewriteGroupTask task : submitted) {
try {
task.cancel();
} catch (Exception e) {
LOG.warn("Failed to cancel rewrite task {}: {}", task.getId(), e.getMessage());
}
}
long deadline = System.nanoTime() + budgetNanos;
for (ConnectorRewriteGroupTask task : submitted) {
long remaining = deadline - System.nanoTime();
if (remaining <= 0) {
return;
}
try {
task.awaitTerminal(remaining, TimeUnit.NANOSECONDS);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
return;
}
}
}
/**
* Collects concurrent group-task completions and cancels the rest on the first failure (ported verbatim
* from the result collector of the pre-SPI rewrite executor, which no longer exists in the tree).
*/
private static class RewriteResultCollector {
private final int expectedTasks;
private final AtomicInteger completedTasks = new AtomicInteger(0);
private final AtomicInteger failedTasks = new AtomicInteger(0);
private volatile Exception firstError = null;
private final CountDownLatch completionLatch;
private final List<ConnectorRewriteGroupTask> allTasks;
RewriteResultCollector(int expectedTasks, List<ConnectorRewriteGroupTask> tasks) {
this.expectedTasks = expectedTasks;
this.completionLatch = new CountDownLatch(expectedTasks);
this.allTasks = tasks;
}
public synchronized void onTaskCompleted(Long taskId) {
int completed = completedTasks.incrementAndGet();
LOG.info("Connector rewrite task {} completed ({}/{})", taskId, completed, expectedTasks);
completionLatch.countDown();
}
public synchronized void onTaskFailed(Long taskId, Exception error) {
int failed = failedTasks.incrementAndGet();
if (firstError == null) {
firstError = error;
cancelAllOtherTasks(taskId);
}
LOG.warn("Connector rewrite task {} failed ({}/{}): {}", taskId, failed, expectedTasks,
error.getMessage());
completionLatch.countDown();
}
private void cancelAllOtherTasks(Long failedTaskId) {
for (ConnectorRewriteGroupTask task : allTasks) {
if (!task.getId().equals(failedTaskId)) {
try {
task.cancel();
} catch (Exception e) {
LOG.warn("Failed to cancel rewrite task {}: {}", task.getId(), e.getMessage());
}
}
}
}
public boolean await(long timeout, TimeUnit unit) throws InterruptedException {
return completionLatch.await(timeout, unit);
}
public boolean isDone() {
return completionLatch.getCount() == 0;
}
public Exception getFirstError() {
return firstError;
}
}
}