diff --git a/paimon-common/src/main/java/org/apache/paimon/fs/FileIO.java b/paimon-common/src/main/java/org/apache/paimon/fs/FileIO.java index 2b0dcec3f760..6a1da31c631f 100644 --- a/paimon-common/src/main/java/org/apache/paimon/fs/FileIO.java +++ b/paimon-common/src/main/java/org/apache/paimon/fs/FileIO.java @@ -294,6 +294,25 @@ default void deleteQuietly(Path file) { } } + /** + * Same as {@link #deleteQuietly(Path)}, but temporarily clears the calling thread's interrupt + * status so Hadoop/RPC deletes are not rejected with {@link java.io.InterruptedIOException}. + * + *

Use this only for abort/cleanup paths (e.g. speculative task kill) where best-effort + * deletion must succeed even though the task thread has been interrupted. The previous + * interrupt status is restored afterwards. + */ + default void deleteQuietlyIgnoringInterrupt(Path file) { + boolean interrupted = Thread.interrupted(); + try { + deleteQuietly(file); + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + default void deleteFilesQuietly(List files) { for (Path file : files) { deleteQuietly(file); diff --git a/paimon-common/src/test/java/org/apache/paimon/fs/FileIOTest.java b/paimon-common/src/test/java/org/apache/paimon/fs/FileIOTest.java index 96e023d1eb46..d6976b1ef3a1 100644 --- a/paimon-common/src/test/java/org/apache/paimon/fs/FileIOTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/fs/FileIOTest.java @@ -29,6 +29,7 @@ import java.io.File; import java.io.FileNotFoundException; import java.io.IOException; +import java.io.InterruptedIOException; import java.nio.file.AccessDeniedException; import java.nio.file.DirectoryNotEmptyException; import java.nio.file.FileAlreadyExistsException; @@ -138,6 +139,27 @@ public static void testOverwriteFileUtf8(Path file, FileIO fileIO) throws Interr assertThat(exception.get()).isNull(); } + @Test + public void testDeleteQuietlyIgnoringInterrupt() throws Exception { + Path file = new Path(tempDir.resolve("interrupted-delete.txt").toUri()); + LocalFileIO local = new LocalFileIO(); + local.writeFile(file, "data", false); + + FileIO interruptSensitive = new InterruptSensitiveFileIO(local); + Thread.currentThread().interrupt(); + try { + // Plain deleteQuietly fails under interrupt with Hadoop-like FileIO. + interruptSensitive.deleteQuietly(file); + assertThat(local.exists(file)).isTrue(); + + interruptSensitive.deleteQuietlyIgnoringInterrupt(file); + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + assertThat(local.exists(file)).isFalse(); + } finally { + Thread.interrupted(); + } + } + @Test public void testListFiles() throws Exception { Path fileA = new Path(tempDir.resolve("a").toUri()); @@ -198,6 +220,84 @@ public void testDefaultArchiveUnsupported() { .hasMessageContaining("unarchive"); } + /** + * Delegates to another {@link FileIO}, but mimics Hadoop RPC behavior by failing with {@link + * InterruptedIOException} when the calling thread is interrupted. + */ + private static class InterruptSensitiveFileIO implements FileIO { + + private final FileIO delegate; + + private InterruptSensitiveFileIO(FileIO delegate) { + this.delegate = delegate; + } + + private void failIfInterrupted() throws InterruptedIOException { + if (Thread.currentThread().isInterrupted()) { + throw new InterruptedIOException("Call interrupted"); + } + } + + @Override + public boolean isObjectStore() { + return delegate.isObjectStore(); + } + + @Override + public void configure(CatalogContext context) { + delegate.configure(context); + } + + @Override + public SeekableInputStream newInputStream(Path path) throws IOException { + failIfInterrupted(); + return delegate.newInputStream(path); + } + + @Override + public PositionOutputStream newOutputStream(Path path, boolean overwrite) + throws IOException { + failIfInterrupted(); + return delegate.newOutputStream(path, overwrite); + } + + @Override + public FileStatus getFileStatus(Path path) throws IOException { + failIfInterrupted(); + return delegate.getFileStatus(path); + } + + @Override + public FileStatus[] listStatus(Path path) throws IOException { + failIfInterrupted(); + return delegate.listStatus(path); + } + + @Override + public boolean exists(Path path) throws IOException { + failIfInterrupted(); + return delegate.exists(path); + } + + @Override + public boolean delete(Path path, boolean recursive) throws IOException { + failIfInterrupted(); + return delegate.delete(path, recursive); + } + + @Override + public boolean mkdirs(Path path) throws IOException { + failIfInterrupted(); + return delegate.mkdirs(path); + } + + @Override + public boolean rename(Path src, Path dst) throws IOException { + failIfInterrupted(); + return delegate.rename(src, dst); + } + } + /** A {@link FileIO} on local filesystem to test various default implementations. */ private static class DummyFileIO implements FileIO { private static final ReentrantLock RENAME_LOCK = new ReentrantLock(); diff --git a/paimon-core/src/main/java/org/apache/paimon/append/AppendOnlyWriter.java b/paimon-core/src/main/java/org/apache/paimon/append/AppendOnlyWriter.java index 5f47789a2619..0a704073dfef 100644 --- a/paimon-core/src/main/java/org/apache/paimon/append/AppendOnlyWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/append/AppendOnlyWriter.java @@ -279,16 +279,31 @@ public void close() throws Exception { sync(); compactManager.close(); + + // Delete newly generated but not committed files (matches RecordWriter.close contract / + // MergeTreeWriter). Includes buffer-flush files that never reached prepareCommit. + List toDelete = new ArrayList<>(newFiles); + newFiles.clear(); + deletedFiles.clear(); for (DataFileMeta file : compactAfter) { // appendOnlyCompactManager will rewrite the file and no file upgrade will occur, so we // can directly delete the file in compactAfter. - fileIO.deleteQuietly(pathFactory.toPath(file)); + toDelete.add(file); } + compactAfter.clear(); + compactBefore.clear(); - sinkWriter.close(); - - if (compactDeletionFile != null) { - compactDeletionFile.clean(); + // Abort in-flight rolling writer even if delete loops fail. + try { + sinkWriter.close(); + } finally { + for (DataFileMeta file : toDelete) { + // Interrupt-tolerant: close runs after speculative task kill. + fileIO.deleteQuietlyIgnoringInterrupt(pathFactory.toPath(file)); + } + if (compactDeletionFile != null) { + compactDeletionFile.clean(); + } } } @@ -309,7 +324,7 @@ public void toBufferedWriter() throws Exception { } finally { // remove small files for (DataFileMeta file : files) { - fileIO.deleteQuietly(pathFactory.toPath(file)); + fileIO.deleteQuietlyIgnoringInterrupt(pathFactory.toPath(file)); } } } diff --git a/paimon-core/src/main/java/org/apache/paimon/append/dataevolution/DataEvolutionNormalCompactTask.java b/paimon-core/src/main/java/org/apache/paimon/append/dataevolution/DataEvolutionNormalCompactTask.java index 1db0786d27cc..9e3196e2033d 100644 --- a/paimon-core/src/main/java/org/apache/paimon/append/dataevolution/DataEvolutionNormalCompactTask.java +++ b/paimon-core/src/main/java/org/apache/paimon/append/dataevolution/DataEvolutionNormalCompactTask.java @@ -33,9 +33,6 @@ import org.apache.paimon.utils.RecordWriter; import org.apache.paimon.utils.SetUtils; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - import java.util.List; import java.util.Set; import java.util.stream.Collectors; @@ -48,8 +45,6 @@ /** Compacts normal structured files of a data evolution table. */ public class DataEvolutionNormalCompactTask extends DataEvolutionCompactTask { - private static final Logger LOG = LoggerFactory.getLogger(DataEvolutionNormalCompactTask.class); - public DataEvolutionNormalCompactTask(BinaryRow partition, List files) { super(partition, files); } @@ -98,32 +93,51 @@ public CommitMessage doCompact(FileStoreTable table, String commitUser) throws E storeWrite.withWriteType(readWriteType); RecordWriter writer = storeWrite.createWriter(partition, 0); - reader.forEachRemaining( - row -> { - try { - writer.write(row); - } catch (Exception e) { - throw new RuntimeException(e); - } - }); - - List writeResult = writer.prepareCommit(false).newFilesIncrement().newFiles(); - checkArgument( - writeResult.size() == 1, "Data evolution compaction should produce one file."); - try { - writer.close(); - storeWrite.close(); - } catch (Exception e) { - LOG.warn("Failed to close reader and writer.", e); + reader.forEachRemaining( + row -> { + try { + writer.write(row); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); + + List writeResult = + writer.prepareCommit(false).newFilesIncrement().newFiles(); + checkArgument( + writeResult.size() == 1, "Data evolution compaction should produce one file."); + + DataFileMeta dataFileMeta = writeResult.get(0).assignFirstRowId(firstRowId); + dataFileMeta = + dataFileMeta.assignSequenceNumber( + minSequenceId(compactBefore), maxSequenceId(compactBefore)); + compactAfter.add(dataFileMeta); + } finally { + // Close on the failure path too: RecordWriter.close deletes flushed-but-unprepared + // files, so a mid-task failure (e.g. speculative kill) does not leave orphans. + // Each close runs independently so a reader failure cannot skip writer cleanup. + boolean interrupted = Thread.interrupted(); + try { + try { + reader.close(); + } catch (Exception ignored) { + } + try { + writer.close(); + } catch (Exception ignored) { + } + try { + storeWrite.close(); + } catch (Exception ignored) { + } + } finally { + if (interrupted) { + Thread.currentThread().interrupt(); + } + } } - DataFileMeta dataFileMeta = writeResult.get(0).assignFirstRowId(firstRowId); - dataFileMeta = - dataFileMeta.assignSequenceNumber( - minSequenceId(compactBefore), maxSequenceId(compactBefore)); - compactAfter.add(dataFileMeta); - return commitMessage(compactBefore, compactAfter); } } diff --git a/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionFileWriter.java index ca2330664dd8..082b9b99b30b 100644 --- a/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionFileWriter.java @@ -52,6 +52,11 @@ public long getPos() { return out.size(); } + /** Path of the in-flight deletion file, used to delete the partial file on failure. */ + Path path() { + return path; + } + public void write(String key, DeletionVector deletionVector) throws IOException { int start = out.size(); int length = deletionVector.serializeTo(out); diff --git a/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionVectorIndexFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionVectorIndexFileWriter.java index 6ce9e1c3875e..dedec85c5742 100644 --- a/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionVectorIndexFileWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/deletionvectors/DeletionVectorIndexFileWriter.java @@ -53,12 +53,21 @@ public DeletionVectorIndexFileWriter( */ public IndexFileMeta writeSingleFile(Map input) throws IOException { DeletionFileWriter writer = new DeletionFileWriter(indexPathFactory, fileIO); + boolean success = false; try { for (Map.Entry entry : input.entrySet()) { writer.write(entry.getKey(), entry.getValue()); } - } finally { writer.close(); + success = true; + } finally { + if (!success) { + // Delete the partial index file (e.g. speculative task kill interrupted the + // write); the caller never sees its meta, so nobody else can reclaim it. + // Interrupt-tolerant because task kill leaves the thread interrupted. + closeQuietly(writer); + fileIO.deleteQuietlyIgnoringInterrupt(writer.path()); + } } return writer.result(); } @@ -70,8 +79,17 @@ public List writeWithRolling(Map input) } List result = new ArrayList<>(); Iterator> iterator = input.entrySet().iterator(); - while (iterator.hasNext()) { - result.add(tryWriter(iterator)); + try { + while (iterator.hasNext()) { + result.add(tryWriter(iterator)); + } + } catch (Throwable t) { + // A failed roll loses the returned metas; delete the already-written files so they + // do not stay as orphans. + for (IndexFileMeta meta : result) { + fileIO.deleteQuietlyIgnoringInterrupt(indexPathFactory.toPath(meta)); + } + throw t; } return result; } @@ -79,6 +97,7 @@ public List writeWithRolling(Map input) private IndexFileMeta tryWriter(Iterator> iterator) throws IOException { DeletionFileWriter writer = new DeletionFileWriter(indexPathFactory, fileIO); + boolean success = false; try { while (iterator.hasNext()) { Map.Entry entry = iterator.next(); @@ -87,9 +106,22 @@ private IndexFileMeta tryWriter(Iterator> iter break; } } - } finally { writer.close(); + success = true; + } finally { + if (!success) { + closeQuietly(writer); + fileIO.deleteQuietlyIgnoringInterrupt(writer.path()); + } } return writer.result(); } + + private static void closeQuietly(DeletionFileWriter writer) { + try { + writer.close(); + } catch (Throwable ignored) { + // best-effort stream close on the failure path + } + } } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/FileWriterAbortExecutor.java b/paimon-core/src/main/java/org/apache/paimon/io/FileWriterAbortExecutor.java index 4c10608c1137..1333bfd43d0f 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/FileWriterAbortExecutor.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/FileWriterAbortExecutor.java @@ -32,6 +32,8 @@ public FileWriterAbortExecutor(FileIO fileIO, Path path) { } public void abort() { - fileIO.deleteQuietly(path); + // Task kill leaves Thread interrupted; use interrupt-tolerant delete so abort + // during prepareCommit (speculative loser) can still remove written files. + fileIO.deleteQuietlyIgnoringInterrupt(path); } } diff --git a/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java b/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java index 7c55360fef4c..8f0b38bbbc35 100644 --- a/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java +++ b/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java @@ -256,7 +256,8 @@ public void deleteFile(DataFileMeta file) { // in WriteFormatKey DataFilePathFactory pathFactory = formatContext.pathFactory(new WriteFormatKey(file.level(), false)); - file.collectFiles(pathFactory).forEach(fileIO::deleteQuietly); + // Interrupt-tolerant: used by MergeTreeWriter.close unpublished cleanup after task kill. + file.collectFiles(pathFactory).forEach(fileIO::deleteQuietlyIgnoringInterrupt); } public FileIO getFileIO() { diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java index e83f5e81c1b7..5488188e3350 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java @@ -217,6 +217,19 @@ public void notifyNewFiles( @Override public List prepareCommit(boolean waitCompaction, long commitIdentifier) throws Exception { + return prepareCommit(waitCompaction, commitIdentifier, null); + } + + /** + * Like {@link #prepareCommit(boolean, long)} but invokes {@code onPrepared} immediately after + * each bucket increment is drained, so callers can register or abort partial results before the + * whole prepare finishes. + */ + public List prepareCommit( + boolean waitCompaction, + long commitIdentifier, + @Nullable java.util.function.Consumer onPrepared) + throws Exception { Function, Boolean> writerCleanChecker; if (writers.values().stream() .map(Map::values) @@ -271,6 +284,9 @@ public List prepareCommit(boolean waitCompaction, long commitIden newFilesIncrement, compactIncrement); result.add(committable); + if (onPrepared != null) { + onPrepared.accept(committable); + } if (committable.isEmpty()) { if (writerCleanChecker.apply(writerContainer) diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreCommitImpl.java b/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreCommitImpl.java index 1eacfa71f78d..86ea9634c173 100644 --- a/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreCommitImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/operation/FileStoreCommitImpl.java @@ -692,6 +692,9 @@ public void abort(List commitMessages) { DataFilePathFactories dataFactories = new DataFilePathFactories(pathFactory); IndexFilePathFactories indexFactories = new IndexFilePathFactories(pathFactory); for (CommitMessage message : commitMessages) { + if (message == null) { + continue; + } DataFilePathFactory dataPathFactory = dataFactories.get(message.partition(), message.bucket()); IndexPathFactory indexPathFactory = @@ -704,14 +707,14 @@ public void abort(List commitMessages) { dataFilesToDelete.addAll(commitMessage.compactIncrement().changelogFiles()); for (DataFileMeta file : dataFilesToDelete) { - fileIO.deleteQuietly(dataPathFactory.toPath(file)); + fileIO.deleteQuietlyIgnoringInterrupt(dataPathFactory.toPath(file)); } List indexFilesToDelete = new ArrayList<>(); indexFilesToDelete.addAll(commitMessage.newFilesIncrement().newIndexFiles()); indexFilesToDelete.addAll(commitMessage.compactIncrement().newIndexFiles()); for (IndexFileMeta file : indexFilesToDelete) { - fileIO.deleteQuietly(indexPathFactory.toPath(file)); + fileIO.deleteQuietlyIgnoringInterrupt(indexPathFactory.toPath(file)); } } } diff --git a/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketWriter.java b/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketWriter.java index 47068513a864..ac5d769ead66 100644 --- a/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketWriter.java +++ b/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketWriter.java @@ -58,6 +58,7 @@ public class PostponeBucketWriter implements RecordWriter, MemoryOwner private final MergeFunction mergeFunction; private final KeyValueFileWriterFactory writerFactory; private final List files; + private final List uncommittedFiles; private final IOFunction, RecordReaderIterator> fileRead; private final @Nullable IOManager ioManager; private final CompressOptions spillCompression; @@ -86,6 +87,7 @@ public PostponeBucketWriter( this.spillCompression = spillCompression; this.maxDiskSize = maxDiskSize; this.files = new ArrayList<>(); + this.uncommittedFiles = new ArrayList<>(); if (restoreIncrement != null) { files.addAll(restoreIncrement.newFilesIncrement().newFiles()); } @@ -142,7 +144,9 @@ private void validateRetract(KeyValue kv) { } private void flush() throws Exception { - files.addAll(sinkWriter.flush()); + List flushedFiles = sinkWriter.flush(); + files.addAll(flushedFiles); + uncommittedFiles.addAll(flushedFiles); } @Override @@ -227,11 +231,16 @@ public CommitIncrement prepareCommit(boolean waitCompaction) throws Exception { flush(); writerFactory.prepareCommit(); List result = new ArrayList<>(files); + CommitIncrement increment = + new CommitIncrement( + new DataIncrement(result, Collections.emptyList(), Collections.emptyList()), + CompactIncrement.emptyIncrement(), + null); files.clear(); - return new CommitIncrement( - new DataIncrement(result, Collections.emptyList(), Collections.emptyList()), - CompactIncrement.emptyIncrement(), - null); + // Ownership transfers to the returned commit increment. Its abort/commit path is now + // responsible for these files; close must not delete them. + uncommittedFiles.clear(); + return increment; } @VisibleForTesting @@ -252,7 +261,17 @@ public void close() throws Exception { try { writerFactory.abortManagedBlobWrites(); } finally { - sinkWriter.close(); + try { + sinkWriter.close(); + } finally { + // Only files produced by this writer and not yet handed to prepareCommit are owned + // here. Files restored from an increment or supplied through addNewFiles are merely + // referenced and must not be deleted. + for (DataFileMeta file : uncommittedFiles) { + writerFactory.deleteFile(file); + } + uncommittedFiles.clear(); + } } } } diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilder.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilder.java index ee0091a7bff9..da2956becaf7 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilder.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilder.java @@ -55,6 +55,18 @@ public interface BatchWriteBuilder extends WriteBuilder { long COMMIT_IDENTIFIER = Long.MAX_VALUE; + /** Get commit user, set by {@link #withCommitUser}. */ + default String commitUser() { + throw new UnsupportedOperationException( + String.format("%s does not expose commit user.", getClass().getName())); + } + + /** Set commit user. */ + default BatchWriteBuilder withCommitUser(String commitUser) { + throw new UnsupportedOperationException( + String.format("%s does not support setting commit user.", getClass().getName())); + } + /** Overwrite writing, same as the 'INSERT OVERWRITE' semantics of SQL. */ default BatchWriteBuilder withOverwrite() { withOverwrite(Collections.emptyMap()); diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java index c06bd9de42d6..03279c48c581 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/BatchWriteBuilderImpl.java @@ -37,7 +37,7 @@ public class BatchWriteBuilderImpl implements BatchWriteBuilder { private static final long serialVersionUID = 1L; private final InnerTable table; - private final String commitUser; + private String commitUser; private Map staticPartition; private boolean appendCommitCheckConflict = false; @@ -70,6 +70,17 @@ public Optional newWriteSelector() { return table.newWriteSelector(); } + @Override + public String commitUser() { + return commitUser; + } + + @Override + public BatchWriteBuilder withCommitUser(String commitUser) { + this.commitUser = commitUser; + return this; + } + @Override public BatchWriteBuilder withOverwrite(@Nullable Map staticPartition) { this.staticPartition = staticPartition; diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java index 101a51e8a551..dfb469eb2431 100644 --- a/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java +++ b/paimon-core/src/main/java/org/apache/paimon/table/sink/TableWriteImpl.java @@ -28,6 +28,7 @@ import org.apache.paimon.io.DataFileMeta; import org.apache.paimon.memory.MemoryPoolFactory; import org.apache.paimon.metrics.MetricRegistry; +import org.apache.paimon.operation.AbstractFileStoreWrite; import org.apache.paimon.operation.BundleFileStoreWriter; import org.apache.paimon.operation.FileStoreWrite; import org.apache.paimon.operation.FileStoreWrite.State; @@ -42,6 +43,7 @@ import java.util.List; import java.util.concurrent.ExecutorService; +import java.util.function.Consumer; import java.util.stream.Collectors; import static org.apache.paimon.utils.Preconditions.checkState; @@ -249,14 +251,41 @@ public void notifyNewFiles( @Override public List prepareCommit(boolean waitCompaction, long commitIdentifier) throws Exception { - return write.prepareCommit(waitCompaction, commitIdentifier); + return prepareCommit(waitCompaction, commitIdentifier, null); } - @Override - public List prepareCommit() throws Exception { + public List prepareCommit( + boolean waitCompaction, + long commitIdentifier, + @Nullable Consumer onPrepared) + throws Exception { + if (write instanceof AbstractFileStoreWrite) { + return ((AbstractFileStoreWrite) write) + .prepareCommit(waitCompaction, commitIdentifier, onPrepared); + } + List messages = write.prepareCommit(waitCompaction, commitIdentifier); + if (onPrepared != null) { + for (CommitMessage message : messages) { + onPrepared.accept(message); + } + } + return messages; + } + + /** + * Prepare commit for batch write and invoke {@code onPrepared} after each drained bucket + * increment so task-side cleanup can abort partial results on speculative kill. + */ + public List prepareCommit(@Nullable Consumer onPrepared) + throws Exception { checkState(!batchCommitted, "BatchTableWrite only support one-time committing."); batchCommitted = true; - return prepareCommit(true, BatchWriteBuilder.COMMIT_IDENTIFIER); + return prepareCommit(true, BatchWriteBuilder.COMMIT_IDENTIFIER, onPrepared); + } + + @Override + public List prepareCommit() throws Exception { + return prepareCommit((Consumer) null); } @Override diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/ObjectsFile.java b/paimon-core/src/main/java/org/apache/paimon/utils/ObjectsFile.java index deb6c9121414..4533d5252706 100644 --- a/paimon-core/src/main/java/org/apache/paimon/utils/ObjectsFile.java +++ b/paimon-core/src/main/java/org/apache/paimon/utils/ObjectsFile.java @@ -192,7 +192,8 @@ protected Pair writeWithoutRolling(Iterator records) { return Pair.of(path.getName(), pos); } } catch (Throwable e) { - fileIO.deleteQuietly(path); + // Interrupt-tolerant cleanup of the unfinished objects/manifest file. + fileIO.deleteQuietlyIgnoringInterrupt(path); throw new RuntimeException( "Exception occurs when writing records to " + path + ". Clean up.", e); } @@ -214,7 +215,8 @@ public long fileSize(Path file) throws IOException { } public void delete(String fileName) { - fileIO.deleteQuietly(pathFactory.toPath(fileName)); + // Interrupt-tolerant: used to clean tmp manifests after failed commit prep. + fileIO.deleteQuietlyIgnoringInterrupt(pathFactory.toPath(fileName)); } public static List readFromIterator( diff --git a/paimon-core/src/test/java/org/apache/paimon/append/AppendOnlyWriterTest.java b/paimon-core/src/test/java/org/apache/paimon/append/AppendOnlyWriterTest.java index fae027b7e081..f5008741bb27 100644 --- a/paimon-core/src/test/java/org/apache/paimon/append/AppendOnlyWriterTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/append/AppendOnlyWriterTest.java @@ -767,6 +767,31 @@ public void testMultipleFlush() throws Exception { Assertions.assertThat(rowCount).isEqualTo(200); } + @Test + public void testCloseDeletesFlushedButUnpreparedFiles() throws Exception { + AppendOnlyWriter writer = createEmptyWriter(Long.MAX_VALUE, true); + writer.setMemoryPool(new HeapMemorySegmentPool(16384L, 1024)); + + char[] s = new char[990]; + Arrays.fill(s, 'a'); + for (int j = 0; j < 100; j++) { + writer.write(row(j, String.valueOf(s), PART)); + } + writer.flush(false, false); + List flushed = new ArrayList<>(writer.getNewFiles()); + Assertions.assertThat(flushed).isNotEmpty(); + for (DataFileMeta meta : flushed) { + Assertions.assertThat(LocalFileIO.create().exists(pathFactory.toPath(meta))).isTrue(); + } + + writer.close(); + + for (DataFileMeta meta : flushed) { + Assertions.assertThat(LocalFileIO.create().exists(pathFactory.toPath(meta))).isFalse(); + } + Assertions.assertThat(writer.getNewFiles()).isEmpty(); + } + @Test public void testClose() throws Exception { AppendOnlyWriter writer = createEmptyWriter(Long.MAX_VALUE, true); diff --git a/paimon-core/src/test/java/org/apache/paimon/deletionvectors/DeletionVectorsIndexFileTest.java b/paimon-core/src/test/java/org/apache/paimon/deletionvectors/DeletionVectorsIndexFileTest.java index 9a243bf8bf09..3c34681930a7 100644 --- a/paimon-core/src/test/java/org/apache/paimon/deletionvectors/DeletionVectorsIndexFileTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/deletionvectors/DeletionVectorsIndexFileTest.java @@ -32,6 +32,7 @@ import org.junit.jupiter.params.provider.ValueSource; import java.io.DataInputStream; +import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; import java.util.ArrayList; @@ -42,10 +43,12 @@ import java.util.Map; import java.util.Random; import java.util.UUID; +import java.util.function.LongConsumer; import java.util.stream.Collectors; import java.util.stream.Stream; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; /** Test for {@link DeletionVectorsIndexFile}. */ public class DeletionVectorsIndexFileTest { @@ -395,6 +398,85 @@ private DeletionVector createEmptyDV(boolean bitmap64) { return bitmap64 ? new Bitmap64DeletionVector() : new BitmapDeletionVector(); } + @Test + public void testWriteSingleFileFailureDeletesPartialFile() { + IndexPathFactory pathFactory = getPathFactory(); + DeletionVectorsIndexFile deletionVectorsIndexFile = + deletionVectorsIndexFile(pathFactory, false); + + Map deleteMap = new LinkedHashMap<>(); + deleteMap.put("file1.parquet", createEmptyDV(false)); + deleteMap.put("file2.parquet", new FailingSerializeDeletionVector(createEmptyDV(false))); + + assertThatThrownBy(() -> deletionVectorsIndexFile.writeSingleFile(deleteMap)) + .hasMessageContaining("Failed to write deletion vectors"); + // The partial index file must be deleted; its meta is never returned to the caller. + assertThat(tempPath.toFile().listFiles()).isEmpty(); + } + + @Test + public void testWriteWithRollingFailureDeletesWrittenFiles() { + IndexPathFactory pathFactory = getPathFactory(); + // Tiny target size so every deletion vector rolls into its own index file. + DeletionVectorsIndexFile deletionVectorsIndexFile = + deletionVectorsIndexFile(pathFactory, MemorySize.ofBytes(1), false); + + Map deleteMap = new LinkedHashMap<>(); + deleteMap.put("file1.parquet", createEmptyDV(false)); + deleteMap.put("file2.parquet", createEmptyDV(false)); + deleteMap.put("file3.parquet", new FailingSerializeDeletionVector(createEmptyDV(false))); + + assertThatThrownBy(() -> deletionVectorsIndexFile.writeWithRolling(deleteMap)) + .hasMessageContaining("Failed to write deletion vectors"); + // Both the in-flight file and the already-rolled files must be deleted. + assertThat(tempPath.toFile().listFiles()).isEmpty(); + } + + /** A {@link DeletionVector} that fails on serialization to simulate an interrupted write. */ + private static class FailingSerializeDeletionVector implements DeletionVector { + + private final DeletionVector delegate; + + private FailingSerializeDeletionVector(DeletionVector delegate) { + this.delegate = delegate; + } + + @Override + public void delete(long position) { + delegate.delete(position); + } + + @Override + public void merge(DeletionVector deletionVector) { + delegate.merge(deletionVector); + } + + @Override + public boolean isEmpty() { + return delegate.isEmpty(); + } + + @Override + public long getCardinality() { + return delegate.getCardinality(); + } + + @Override + public void forEachDeletedPosition(LongConsumer consumer) { + delegate.forEachDeletedPosition(consumer); + } + + @Override + public boolean isDeleted(long position) { + return delegate.isDeleted(position); + } + + @Override + public int serializeTo(DataOutputStream out) throws IOException { + throw new IOException("injected serialize failure"); + } + } + private DeletionVectorsIndexFile deletionVectorsIndexFile( IndexPathFactory pathFactory, boolean bitmap64) { return deletionVectorsIndexFile(pathFactory, MemorySize.ofBytes(Long.MAX_VALUE), bitmap64); diff --git a/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitAbortTest.java b/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitAbortTest.java new file mode 100644 index 000000000000..1fc1455737c4 --- /dev/null +++ b/paimon-core/src/test/java/org/apache/paimon/operation/FileStoreCommitAbortTest.java @@ -0,0 +1,195 @@ +/* + * 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.paimon.operation; + +import org.apache.paimon.CoreOptions; +import org.apache.paimon.TestAppendFileStore; +import org.apache.paimon.TestKeyValueGenerator; +import org.apache.paimon.TestKeyValueGenerator.GeneratorMode; +import org.apache.paimon.data.BinaryRow; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.index.IndexFileMeta; +import org.apache.paimon.index.IndexPathFactory; +import org.apache.paimon.io.CompactIncrement; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.io.DataFilePathFactory; +import org.apache.paimon.io.DataIncrement; +import org.apache.paimon.table.sink.CommitMessageImpl; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThat; + +/** Tests for {@link FileStoreCommitImpl#abort}. */ +public class FileStoreCommitAbortTest { + + @TempDir java.nio.file.Path tempDir; + + private CommitMessageImpl writeDvCommitMessage(TestAppendFileStore store, BinaryRow partition) + throws Exception { + Map> dvs = new HashMap<>(); + dvs.put("f1", Arrays.asList(1, 3, 5)); + return store.writeDVIndexFiles(partition, 0, dvs); + } + + @Test + public void abortDeletesNewIndexFiles() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.DELETION_VECTOR_BITMAP64.key(), "true"); + TestAppendFileStore store = TestAppendFileStore.createAppendStore(tempDir, options); + TestKeyValueGenerator generator = + new TestKeyValueGenerator(GeneratorMode.MULTI_PARTITIONED); + BinaryRow partition = generator.getPartition(generator.next()); + CommitMessageImpl commitMessage = writeDvCommitMessage(store, partition); + + IndexPathFactory indexPathFactory = + store.pathFactory().indexFileFactory(commitMessage.partition(), 0); + FileIO fileIO = store.fileIO(); + assertThat(commitMessage.newFilesIncrement().newIndexFiles()).isNotEmpty(); + for (IndexFileMeta indexFile : commitMessage.newFilesIncrement().newIndexFiles()) { + assertThat(fileIO.exists(indexPathFactory.toPath(indexFile))).isTrue(); + } + + store.newCommit().abort(Collections.singletonList(commitMessage)); + + for (IndexFileMeta indexFile : commitMessage.newFilesIncrement().newIndexFiles()) { + assertThat(fileIO.exists(indexPathFactory.toPath(indexFile))).isFalse(); + } + } + + @Test + public void abortUnderThreadInterruptStillDeletesFiles() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.DELETION_VECTOR_BITMAP64.key(), "true"); + TestAppendFileStore store = TestAppendFileStore.createAppendStore(tempDir, options); + TestKeyValueGenerator generator = + new TestKeyValueGenerator(GeneratorMode.MULTI_PARTITIONED); + BinaryRow partition = generator.getPartition(generator.next()); + CommitMessageImpl commitMessage = writeDvCommitMessage(store, partition); + + IndexPathFactory indexPathFactory = + store.pathFactory().indexFileFactory(commitMessage.partition(), 0); + FileIO fileIO = store.fileIO(); + assertThat(commitMessage.newFilesIncrement().newIndexFiles()).isNotEmpty(); + + Thread.currentThread().interrupt(); + try { + store.newCommit().abort(Collections.singletonList(commitMessage)); + assertThat(Thread.currentThread().isInterrupted()).isTrue(); + } finally { + Thread.interrupted(); + } + + for (IndexFileMeta indexFile : commitMessage.newFilesIncrement().newIndexFiles()) { + assertThat(fileIO.exists(indexPathFactory.toPath(indexFile))).isFalse(); + } + } + + @Test + public void abortSkipsNullMessages() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.DELETION_VECTOR_BITMAP64.key(), "true"); + TestAppendFileStore store = TestAppendFileStore.createAppendStore(tempDir, options); + TestKeyValueGenerator generator = + new TestKeyValueGenerator(GeneratorMode.MULTI_PARTITIONED); + BinaryRow partition = generator.getPartition(generator.next()); + CommitMessageImpl commitMessage = writeDvCommitMessage(store, partition); + + IndexPathFactory indexPathFactory = + store.pathFactory().indexFileFactory(commitMessage.partition(), 0); + FileIO fileIO = store.fileIO(); + assertThat(commitMessage.newFilesIncrement().newIndexFiles()).isNotEmpty(); + + store.newCommit().abort(Arrays.asList(null, commitMessage)); + + for (IndexFileMeta indexFile : commitMessage.newFilesIncrement().newIndexFiles()) { + assertThat(fileIO.exists(indexPathFactory.toPath(indexFile))).isFalse(); + } + } + + @Test + public void abortDoesNotDeleteDeletedOrCompactBeforeFiles() throws Exception { + Map options = new HashMap<>(); + options.put(CoreOptions.DELETION_VECTOR_BITMAP64.key(), "true"); + TestAppendFileStore store = TestAppendFileStore.createAppendStore(tempDir, options); + TestKeyValueGenerator generator = + new TestKeyValueGenerator(GeneratorMode.MULTI_PARTITIONED); + BinaryRow partition = generator.getPartition(generator.next()); + FileIO fileIO = store.fileIO(); + DataFilePathFactory dataPathFactory = + store.pathFactory().createDataFilePathFactory(partition, 0); + IndexPathFactory indexPathFactory = store.pathFactory().indexFileFactory(partition, 0); + + CommitMessageImpl keepData = + store.writeDataFiles( + partition, 0, Arrays.asList("keep-deleted.parquet", "keep-before.parquet")); + DataFileMeta deletedFile = keepData.newFilesIncrement().newFiles().get(0); + DataFileMeta compactBeforeFile = keepData.newFilesIncrement().newFiles().get(1); + + CommitMessageImpl abortData = + store.writeDataFiles( + partition, 0, Arrays.asList("abort-new.parquet", "abort-after.parquet")); + DataFileMeta newFile = abortData.newFilesIncrement().newFiles().get(0); + DataFileMeta compactAfterFile = abortData.newFilesIncrement().newFiles().get(1); + + CommitMessageImpl keepIndexMsg = writeDvCommitMessage(store, partition); + CommitMessageImpl abortIndexMsg = writeDvCommitMessage(store, partition); + IndexFileMeta deletedIndexFile = keepIndexMsg.newFilesIncrement().newIndexFiles().get(0); + IndexFileMeta newIndexFile = abortIndexMsg.newFilesIncrement().newIndexFiles().get(0); + + assertThat(fileIO.exists(dataPathFactory.toPath(deletedFile))).isTrue(); + assertThat(fileIO.exists(dataPathFactory.toPath(compactBeforeFile))).isTrue(); + assertThat(fileIO.exists(dataPathFactory.toPath(newFile))).isTrue(); + assertThat(fileIO.exists(dataPathFactory.toPath(compactAfterFile))).isTrue(); + assertThat(fileIO.exists(indexPathFactory.toPath(deletedIndexFile))).isTrue(); + assertThat(fileIO.exists(indexPathFactory.toPath(newIndexFile))).isTrue(); + + CommitMessageImpl abortMessage = + new CommitMessageImpl( + partition, + 0, + store.options().bucket(), + new DataIncrement( + Collections.singletonList(newFile), + Collections.singletonList(deletedFile), + Collections.emptyList(), + Collections.singletonList(newIndexFile), + Collections.singletonList(deletedIndexFile)), + new CompactIncrement( + Collections.singletonList(compactBeforeFile), + Collections.singletonList(compactAfterFile), + Collections.emptyList())); + + store.newCommit().abort(Collections.singletonList(abortMessage)); + + assertThat(fileIO.exists(dataPathFactory.toPath(deletedFile))).isTrue(); + assertThat(fileIO.exists(dataPathFactory.toPath(compactBeforeFile))).isTrue(); + assertThat(fileIO.exists(indexPathFactory.toPath(deletedIndexFile))).isTrue(); + assertThat(fileIO.exists(dataPathFactory.toPath(newFile))).isFalse(); + assertThat(fileIO.exists(dataPathFactory.toPath(compactAfterFile))).isFalse(); + assertThat(fileIO.exists(indexPathFactory.toPath(newIndexFile))).isFalse(); + } +} diff --git a/paimon-core/src/test/java/org/apache/paimon/table/PrimaryKeySimpleTableTest.java b/paimon-core/src/test/java/org/apache/paimon/table/PrimaryKeySimpleTableTest.java index ea8f096a53b3..a6290200e809 100644 --- a/paimon-core/src/test/java/org/apache/paimon/table/PrimaryKeySimpleTableTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/table/PrimaryKeySimpleTableTest.java @@ -107,6 +107,7 @@ import java.util.Map; import java.util.Objects; import java.util.Random; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; @@ -212,6 +213,54 @@ public void testPostponeBucket() throws Exception { assertThat(file.valueStatsCols()).isEmpty(); } + @Test + public void testPostponeBucketCloseDeletesUnpreparedFlushedFiles() throws Exception { + FileStoreTable table = + createFileStoreTable(options -> options.set(BUCKET, BucketMode.POSTPONE_BUCKET)); + + BatchTableWrite write = table.newBatchWriteBuilder().newWrite(); + List flushedFiles = Collections.emptyList(); + try { + write.write(rowData(0, 0, 0L)); + PostponeBucketFileStoreWrite fileStoreWrite = + (PostponeBucketFileStoreWrite) ((TableWriteImpl) write).getWrite(); + PostponeBucketWriter writer = + (PostponeBucketWriter) + fileStoreWrite + .writers() + .values() + .iterator() + .next() + .values() + .iterator() + .next() + .writer; + + // DirectSinkWriter cannot release memory, so this forces an on-disk flush before + // prepareCommit and reproduces the speculative-attempt cleanup path. + writer.flushMemory(); + flushedFiles = new ArrayList<>(writer.dataFiles()); + assertThat(flushedFiles).isNotEmpty(); + assertThat( + Arrays.stream(table.fileIO().listFiles(table.location(), true)) + .map(status -> status.getPath().getName()) + .collect(Collectors.toSet())) + .contains(flushedFiles.get(0).fileName()); + } finally { + write.close(); + } + + Set remainingFiles = + Arrays.stream(table.fileIO().listFiles(table.location(), true)) + .map(status -> status.getPath().getName()) + .collect(Collectors.toSet()); + assertThat(remainingFiles) + .doesNotContainAnyElementsOf( + flushedFiles.stream() + .map(DataFileMeta::fileName) + .collect(Collectors.toList())); + } + @ParameterizedTest(name = "format-{0}") @ValueSource(strings = {"avro", "parquet"}) public void testStatsModePerLevel(String format) throws Exception { diff --git a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/CompactProcedure.java b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/CompactProcedure.java index 6868bc2bcf95..7889f45e8f3d 100644 --- a/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/CompactProcedure.java +++ b/paimon-spark/paimon-spark-common/src/main/java/org/apache/paimon/spark/procedure/CompactProcedure.java @@ -43,6 +43,8 @@ import org.apache.paimon.spark.sort.TableSorter; import org.apache.paimon.spark.util.ScanPlanHelper$; import org.apache.paimon.spark.utils.SparkProcedureUtils; +import org.apache.paimon.spark.write.SparkAttemptCleanup; +import org.apache.paimon.spark.write.SparkAttemptWrite; import org.apache.paimon.table.BucketMode; import org.apache.paimon.table.FileStoreTable; import org.apache.paimon.table.SpecialFields; @@ -54,6 +56,7 @@ import org.apache.paimon.table.sink.CommitMessageImpl; import org.apache.paimon.table.sink.CommitMessageSerializer; import org.apache.paimon.table.sink.TableCommitImpl; +import org.apache.paimon.table.sink.TableWriteImpl; import org.apache.paimon.table.source.DataSplit; import org.apache.paimon.table.source.EndOfScanException; import org.apache.paimon.table.source.snapshot.SnapshotReader; @@ -352,31 +355,74 @@ private void compactAwareBucketTable( IOManager ioManager = SparkUtils.createIOManager(); BatchTableWrite write = writeBuilder.newWrite(); write.withIOManager(ioManager); - try { - while (pairIterator.hasNext()) { - Pair pair = - pairIterator.next(); - write.compact( - SerializationUtils.deserializeBinaryRow( - pair.getLeft()), - pair.getRight(), - fullCompact); - } - CommitMessageSerializer serializer = - new CommitMessageSerializer(); - List messages = - write.prepareCommit(); - List serializedMessages = - new ArrayList<>(messages.size()); - for (CommitMessage commitMessage : messages) { - serializedMessages.add( - serializer.serialize(commitMessage)); - } - return serializedMessages.iterator(); - } finally { - write.close(); - ioManager.close(); - } + return SparkAttemptWrite.runJava( + table, + writeBuilder, + () -> { + try { + write.close(); + ioManager.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + cleanup -> { + try { + while (pairIterator.hasNext()) { + cleanup + .checkInterruptedPeriodically(); + Pair pair = + pairIterator.next(); + write.compact( + SerializationUtils + .deserializeBinaryRow( + pair + .getLeft()), + pair.getRight(), + fullCompact); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + cleanup -> { + try { + List messages = + new ArrayList<>(); + ((TableWriteImpl) write) + .prepareCommit( + message -> { + cleanup + .addPreparedJava( + Collections + .singletonList( + message)); + messages.add( + message); + }); + return messages; + } catch (Exception e) { + throw new RuntimeException(e); + } + }, + messages -> { + try { + CommitMessageSerializer serializer = + new CommitMessageSerializer(); + List serializedMessages = + new ArrayList<>( + messages.size()); + for (CommitMessage commitMessage : + messages) { + serializedMessages.add( + serializer.serialize( + commitMessage)); + } + return serializedMessages.iterator(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); }); try (BatchTableCommit commit = writeBuilder.newCommit()) { @@ -442,6 +488,8 @@ private void compactUnAwareBucketTable( int readParallelism = readParallelism(serializedTasks, spark()); String commitUser = createCommitUser(table.coreOptions().toConfiguration()); + BatchWriteBuilder appendCompactWriteBuilder = + table.newBatchWriteBuilder().withCommitUser(commitUser); JavaRDD commitMessageJavaRDD = javaSparkContext .parallelize(serializedTasks, readParallelism) @@ -457,6 +505,18 @@ private void compactUnAwareBucketTable( SpecialFields.rowTypeWithRowTracking( table.rowType())); } + SparkAttemptCleanup cleanup = + SparkAttemptCleanup.forJava( + table.fullName(), + commitUser, + appendCompactWriteBuilder, + () -> { + try { + write.close(); + } catch (Exception e) { + throw new RuntimeException(e); + } + }); AppendCompactTaskSerializer ser = new AppendCompactTaskSerializer(); List messages = new ArrayList<>(); @@ -464,17 +524,27 @@ private void compactUnAwareBucketTable( CommitMessageSerializer messageSer = new CommitMessageSerializer(); while (taskIterator.hasNext()) { + cleanup.checkInterruptedPeriodically(); AppendCompactTask task = ser.deserialize( ser.getVersion(), taskIterator.next()); + CommitMessage commitMessage = + task.doCompact(table, write); + // Incremental registration avoids O(n²) + // copying of the full prepared list on every + // compacted file. + cleanup.addPreparedJava( + java.util.Collections.singletonList( + commitMessage)); messages.add( - messageSer.serialize( - task.doCompact(table, write))); + messageSer.serialize(commitMessage)); } + cleanup.checkInterrupted("before return"); + cleanup.markReturned(); return messages.iterator(); } finally { - write.close(); + cleanup.close(); } }); @@ -553,29 +623,55 @@ private void compactDataEvolutionTable( } int readParallelism = readParallelism(serializedTasks, spark()); + BatchWriteBuilder dataEvolutionCompactWriteBuilder = + table.newBatchWriteBuilder().withCommitUser(commitUser); JavaRDD commitMessageJavaRDD = javaSparkContext .parallelize(serializedTasks, readParallelism) .mapPartitions( (FlatMapFunction, byte[]>) taskIterator -> { + // doCompact creates and closes its own + // writers (finally inside the task + // implementations), so there is no shared + // unprepared writer to close here. + SparkAttemptCleanup cleanup = + SparkAttemptCleanup.forJava( + table.fullName(), + commitUser, + dataEvolutionCompactWriteBuilder, + () -> {}); DataEvolutionCompactTaskSerializer ser = new DataEvolutionCompactTaskSerializer(); List messagesBytes = new ArrayList<>(); CommitMessageSerializer messageSer = new CommitMessageSerializer(); - while (taskIterator.hasNext()) { - DataEvolutionCompactTask task = - ser.deserialize( - ser.getVersion(), - taskIterator.next()); - messagesBytes.add( - messageSer.serialize( - task.doCompact( - table, - commitUser))); + try { + while (taskIterator.hasNext()) { + cleanup.checkInterruptedPeriodically(); + DataEvolutionCompactTask task = + ser.deserialize( + ser.getVersion(), + taskIterator.next()); + CommitMessage commitMessage = + task.doCompact( + table, commitUser); + // Incremental registration so a kill + // between tasks can abort the files + // of already-finished tasks. + cleanup.addPreparedJava( + Collections.singletonList( + commitMessage)); + messagesBytes.add( + messageSer.serialize( + commitMessage)); + } + cleanup.checkInterrupted("before return"); + cleanup.markReturned(); + return messagesBytes.iterator(); + } finally { + cleanup.close(); } - return messagesBytes.iterator(); }); List messages = new ArrayList<>(); diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala index aaba0f1ca0bf..782473a4c9de 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/commands/PaimonSparkWriter.scala @@ -36,7 +36,7 @@ import org.apache.paimon.spark.schema.SparkSystemColumns.{BUCKET_COL, ROW_KIND_C import org.apache.paimon.spark.sort.TableSorter import org.apache.paimon.spark.util.OptionUtils.paimonExtensionEnabled import org.apache.paimon.spark.util.SparkRowUtils -import org.apache.paimon.spark.write.{PaimonDataWrite, WriteHelper, WriteTaskResult} +import org.apache.paimon.spark.write.{PaimonDataWrite, SparkAttemptCleanup, WriteHelper, WriteTaskResult} import org.apache.paimon.table.{FileStoreTable, PostponeUtils, SpecialFields} import org.apache.paimon.table.BucketMode._ import org.apache.paimon.table.sink._ @@ -380,49 +380,70 @@ case class PaimonSparkWriter( .groupByKey(_.bucketPath) .mapGroups { (_, iter: Iterator[SparkDeletionVector]) => - val indexHandler = table.store().newIndexFileHandler() - var dvIndexFileMaintainer: BaseAppendDeleteFileMaintainer = null - while (iter.hasNext) { - val sdv: SparkDeletionVector = iter.next() - if (dvIndexFileMaintainer == null) { - val partition = SerializationUtils.deserializeBinaryRow(sdv.partition) - dvIndexFileMaintainer = if (bucketMode == BUCKET_UNAWARE) { - BaseAppendDeleteFileMaintainer.forUnawareAppend(indexHandler, snapshot, partition) - } else { - BaseAppendDeleteFileMaintainer.forBucketedAppend( - indexHandler, - snapshot, - partition, - sdv.bucket) + // DV files are only materialized by dvIndexFileMaintainer.persist(); before that + // everything is in-memory, so there is no unprepared writer to close. A partial + // index file left by a kill inside persist() is deleted on the failure path by + // DeletionVectorIndexFileWriter; files from a completed persist() are registered + // via setPrepared below and reclaimed by abort. + val cleanup = new SparkAttemptCleanup( + table.fullName(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder), + writeBuilder, + () => ()) + try { + val indexHandler = table.store().newIndexFileHandler() + var dvIndexFileMaintainer: BaseAppendDeleteFileMaintainer = null + while (iter.hasNext) { + cleanup.checkInterruptedPeriodically() + val sdv: SparkDeletionVector = iter.next() + if (dvIndexFileMaintainer == null) { + val partition = SerializationUtils.deserializeBinaryRow(sdv.partition) + dvIndexFileMaintainer = if (bucketMode == BUCKET_UNAWARE) { + BaseAppendDeleteFileMaintainer.forUnawareAppend(indexHandler, snapshot, partition) + } else { + BaseAppendDeleteFileMaintainer.forBucketedAppend( + indexHandler, + snapshot, + partition, + sdv.bucket) + } + } + if (dvIndexFileMaintainer == null) { + throw new RuntimeException("can't create the dv maintainer.") } - } - if (dvIndexFileMaintainer == null) { - throw new RuntimeException("can't create the dv maintainer.") - } - dvIndexFileMaintainer.notifyNewDeletionVector( - new Path(sdv.dataFilePath).getName, - DeletionVector.deserializeFromBytes(sdv.deletionVector)) + dvIndexFileMaintainer.notifyNewDeletionVector( + new Path(sdv.dataFilePath).getName, + DeletionVector.deserializeFromBytes(sdv.deletionVector)) + } + cleanup.checkInterrupted("before persist") + val indexEntries = dvIndexFileMaintainer.persist() + + val (added, deleted) = indexEntries.asScala.partition(_.kind() == FileKind.ADD) + + val commitMessage = new CommitMessageImpl( + dvIndexFileMaintainer.getPartition, + dvIndexFileMaintainer.getBucket, + null, + new DataIncrement( + java.util.Collections.emptyList(), + java.util.Collections.emptyList(), + java.util.Collections.emptyList(), + added.map(_.indexFile).asJava, + deleted.map(_.indexFile).asJava + ), + CompactIncrement.emptyIncrement() + ) + cleanup.setPrepared(Seq(commitMessage)) + cleanup.checkInterrupted("after persist") + val serializer = new CommitMessageSerializer + val bytes = serializer.serialize(commitMessage) + cleanup.checkInterrupted("before return") + cleanup.markReturned() + bytes + } finally { + cleanup.close() } - val indexEntries = dvIndexFileMaintainer.persist() - - val (added, deleted) = indexEntries.asScala.partition(_.kind() == FileKind.ADD) - - val commitMessage = new CommitMessageImpl( - dvIndexFileMaintainer.getPartition, - dvIndexFileMaintainer.getBucket, - null, - new DataIncrement( - java.util.Collections.emptyList(), - java.util.Collections.emptyList(), - java.util.Collections.emptyList(), - added.map(_.indexFile).asJava, - deleted.map(_.indexFile).asJava - ), - CompactIncrement.emptyIncrement() - ) - val serializer = new CommitMessageSerializer - serializer.serialize(commitMessage) } serializedCommits .collect() diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/format/FormatTableBatchWriteBase.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/format/FormatTableBatchWriteBase.scala index 853c7f41467b..0b029419093c 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/format/FormatTableBatchWriteBase.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/format/FormatTableBatchWriteBase.scala @@ -19,7 +19,7 @@ package org.apache.paimon.spark.format import org.apache.paimon.spark.SparkInternalRowWrapper -import org.apache.paimon.spark.write.{FormatTableWriteTaskResult, V2DataWrite, WriteTaskResult} +import org.apache.paimon.spark.write.{FormatTableWriteTaskResult, SparkAttemptCleanup, V2DataWrite, WriteTaskResult} import org.apache.paimon.table.FormatTable import org.apache.paimon.table.sink.{BatchTableWrite, BatchWriteBuilder, CommitMessage} @@ -100,13 +100,24 @@ private class FormatTableDataWriter(batchWriteBuilder: BatchWriteBuilder, writeS private val write: BatchTableWrite = batchWriteBuilder.newWrite() + private val cleanup = new SparkAttemptCleanup( + batchWriteBuilder.tableName(), + SparkAttemptCleanup.commitUserOrUnknown(batchWriteBuilder), + batchWriteBuilder, + () => write.close()) + + override protected def attemptCleanup: Option[SparkAttemptCleanup] = Some(cleanup) + override def write(record: InternalRow): Unit = { + cleanup.checkInterruptedPeriodically() val paimonRow = rowConverter.apply(record) write.write(paimonRow) } override def commitImpl(): Seq[CommitMessage] = { - write.prepareCommit().asScala.toSeq + val messages = write.prepareCommit().asScala + registerPrepared(messages) + messages } def buildWriteTaskResult(commitMessages: Seq[CommitMessage]): FormatTableWriteTaskResult = { @@ -119,16 +130,9 @@ private class FormatTableDataWriter(batchWriteBuilder: BatchWriteBuilder, writeS override def abort(): Unit = { logInfo("Aborting FormatTable data writer") - close() + cleanup.abortPrepared() + cleanup.close() } - override def close(): Unit = { - try { - write.close() - } catch { - case e: Exception => - logError("Error closing FormatTableDataWriter", e) - throw new RuntimeException(e) - } - } + override def close(): Unit = cleanup.close() } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataEvolutionTableDataWrite.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataEvolutionTableDataWrite.scala index 1ba681c8134e..d679a42e1bfa 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataEvolutionTableDataWrite.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataEvolutionTableDataWrite.scala @@ -60,6 +60,30 @@ case class DataEvolutionTableDataWrite( private val toPaimonRow = { SparkRowUtils.toPaimonRow(writeType, -1, uriReaderFactory) } + + private val cleanup = new SparkAttemptCleanup( + writeBuilder.tableName(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder), + writeBuilder, + () => closeLocalResources()) + + override protected def attemptCleanup: Option[SparkAttemptCleanup] = Some(cleanup) + + private def closeLocalResources(): Unit = { + try { + if (currentWriter != null) { + try { + // Abort in-flight PerFileWriter so partial data files are deleted. RecordWriter.close + // deletes unprepared files (RecordWriter.close contract). + currentWriter.abortQuietly() + } finally { + currentWriter = null + } + } + } finally { + ioManager.close() + } + } private val rawBlobFallbackFields = rawBlobPlaceholderMarkerIndexes.toSeq.sortBy(_._1).toArray private val rawBlobFallbackFieldIndexes = rawBlobFallbackFields.map(_._1) private val rawBlobFallbackMappings = { @@ -105,6 +129,7 @@ case class DataEvolutionTableDataWrite( } def write(row: Row): Unit = { + cleanup.checkInterruptedPeriodically() val firstRowId = row.getLong(firstRowIdIndex) val rowId = row.getLong(rowIdIndex) @@ -160,9 +185,20 @@ case class DataEvolutionTableDataWrite( private def finishCurrentWriter(): Unit = { if (currentWriter != null) { - commitMessages ++= currentWriter.finish() + // Register each prepared file incrementally so a mid-attempt kill can abort already + // finished PerFileWriters. prepareAndRegister() calls addPrepared inside prepareMessage + // immediately after prepareCommit drains newFiles, before returning to closeQuietly(). + val writer = currentWriter + currentWriter = null + try { + cleanup.checkInterrupted("before per-file prepareCommit") + val messages = writer.prepareAndRegister(cleanup.addPrepared) + commitMessages ++= messages + cleanup.checkInterrupted("after per-file prepareCommit") + } finally { + writer.closeQuietly() + } } - currentWriter = null } def write(row: Row, bucket: Int): Unit = { @@ -175,9 +211,7 @@ case class DataEvolutionTableDataWrite( commitMessages.toSeq } - override def close(): Unit = { - ioManager.close() - } + override def close(): Unit = cleanup.close() private case class PerFileWriter( partition: BinaryRow, @@ -209,38 +243,52 @@ case class DataEvolutionTableDataWrite( recordWriter.write(row) } - def finish(): Seq[CommitMessageImpl] = { + /** + * Prepare commit messages and register them with cleanup in the same stack frame as + * prepareCommit, so a speculative kill cannot land between draining newFiles and addPrepared. + */ + def prepareAndRegister( + registerPrepared: Seq[CommitMessageImpl] => Unit): Seq[CommitMessageImpl] = { + fillGapUntil(firstRowId + numRecords) + + assert( + numRecords == numWritten, + s"Number of written records $numWritten does not match expected number $numRecords for first row ID $firstRowId.") + if (fillerRows > 0) { + DataEvolutionTableDataWrite.LOG.warn( + s"Data evolution merge wrote $fillerRows filler rows out of $numRecords rows " + + s"for row-id range [$firstRowId, ${firstRowId + numRecords}) to preserve " + + "row-id continuity. Raw blob fields in filler rows are written as NULL.") + } + val result = recordWriter.prepareCommit(false) + val dataFiles = result.newFilesIncrement().newFiles() + val dataFileMetas = assignFirstRowIds(dataFiles.asScala.toSeq) + val messages = Seq( + new CommitMessageImpl( + partition, + 0, + null, + new DataIncrement(dataFileMetas.asJava, Collections.emptyList(), Collections.emptyList()), + CompactIncrement.emptyIncrement() + )) + registerPrepared(messages) + messages + } + + def closeQuietly(): Unit = { try { - fillGapUntil(firstRowId + numRecords) - - assert( - numRecords == numWritten, - s"Number of written records $numWritten does not match expected number $numRecords for first row ID $firstRowId.") - if (fillerRows > 0) { - DataEvolutionTableDataWrite.LOG.warn( - s"Data evolution merge wrote $fillerRows filler rows out of $numRecords rows " + - s"for row-id range [$firstRowId, ${firstRowId + numRecords}) to preserve " + - "row-id continuity. Raw blob fields in filler rows are written as NULL.") - } - val result = recordWriter.prepareCommit(false) - val dataFiles = result.newFilesIncrement().newFiles() - val dataFileMetas = assignFirstRowIds(dataFiles.asScala.toSeq) - Seq( - new CommitMessageImpl( - partition, - 0, - null, - new DataIncrement( - dataFileMetas.asJava, - Collections.emptyList(), - Collections.emptyList()), - CompactIncrement.emptyIncrement() - )) - } finally { recordWriter.close() + } catch { + case _: Exception => } } + /** + * Abort an unfinished PerFileWriter. Unlike a normal prepareMessage path, this must delete the + * in-flight data file rather than leave a partial orphan on disk. + */ + def abortQuietly(): Unit = closeQuietly() + private def fillGapUntil(rowId: Long, fillerSourceRow: InternalRow = null): Unit = { if (fillerRow == null && fillerSourceRow != null) { // Cache the first record eagerly because finish may need it to fill a trailing gap. diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataWrite.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataWrite.scala index 603ec7ecff1c..4f86dbbc4b92 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataWrite.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/DataWrite.scala @@ -35,12 +35,35 @@ trait DataWrite[T] extends AutoCloseable { def write(record: T): Unit + /** + * Optional task-side cleanup for Spark speculative execution. When set, prepared/unprepared files + * of a killed (loser) attempt are aborted before the attempt returns a successful result. + * Defaults to None so non-inner-table writers are unaffected. + */ + protected def attemptCleanup: Option[SparkAttemptCleanup] = None + + /** + * Register prepared commit messages incrementally so a mid-attempt kill can abort already drained + * files. Call this immediately after each prepareCommit increment (per bucket / per writer) + * inside {@link #commitImpl()}. + */ + protected def registerPrepared(messages: Seq[CommitMessage]): Unit = { + attemptCleanup.foreach(_.addPrepared(messages)) + } + def commit: WriteTaskResult = { + val cleanup = attemptCleanup try { + cleanup.foreach(_.checkInterrupted("before prepareCommit")) preCommit() val commitMessages = commitImpl() + // commitImpl() must register prepared messages incrementally via registerPrepared(). + cleanup.foreach(_.checkInterrupted("after prepareCommit")) postCommit(commitMessages) - buildWriteTaskResult(commitMessages) + val result = buildWriteTaskResult(commitMessages) + cleanup.foreach(_.checkInterrupted("before return")) + cleanup.foreach(_.markReturned()) + result } finally { close() } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonDataWrite.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonDataWrite.scala index a954bd64b2f9..c324825d2467 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonDataWrite.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonDataWrite.scala @@ -28,8 +28,6 @@ import org.apache.paimon.utils.UriReaderFactory import org.apache.spark.sql.Row -import scala.collection.JavaConverters._ - case class PaimonDataWrite( writeBuilder: BatchWriteBuilder, writeType: RowType, @@ -60,36 +58,57 @@ case class PaimonDataWrite( SparkRowUtils.toPaimonRow(writeType, rowKindColIdx, uriReaderFactory) } + private val cleanup = new SparkAttemptCleanup( + writeBuilder.tableName(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder), + writeBuilder, + () => closeLocalResources()) + + override protected def attemptCleanup: Option[SparkAttemptCleanup] = Some(cleanup) + + private def closeLocalResources(): Unit = { + write.close() + ioManager.close() + } + def write(row: Row): Unit = { + cleanup.checkInterruptedPeriodically() postWrite(write.writeAndReturn(toPaimonRow(row))) } def write(row: Row, bucket: Int): Unit = { + cleanup.checkInterruptedPeriodically() postWrite(write.writeAndReturn(toPaimonRow(row), bucket)) } override def commitImpl(): Seq[CommitMessage] = { - val commitMessages = write.prepareCommit().asScala.toSeq + val messages = scala.collection.mutable.ListBuffer[CommitMessage]() + write.prepareCommit( + (msg: CommitMessage) => { + val transformed = transformCommitMessage(msg) + registerPrepared(Seq(transformed)) + messages += transformed + }) + messages.toSeq + } + private def transformCommitMessage(message: CommitMessage): CommitMessage = { if (postponePartitionBucketComputer.isDefined) { - commitMessages.map { - case message: CommitMessageImpl => + message match { + case m: CommitMessageImpl => new CommitMessageImpl( - message.partition(), - message.bucket(), - postponePartitionBucketComputer.get.apply(message.partition()), - message.newFilesIncrement(), - message.compactIncrement() + m.partition(), + m.bucket(), + postponePartitionBucketComputer.get.apply(m.partition()), + m.newFilesIncrement(), + m.compactIncrement() ) case _ => throw new RuntimeException() } } else { - commitMessages + message } } - override def close(): Unit = { - write.close() - ioManager.close() - } + override def close(): Unit = cleanup.close() } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonV2DataWriter.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonV2DataWriter.scala index 05b6bfd307c4..388fd738a1bc 100644 --- a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonV2DataWriter.scala +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/PaimonV2DataWriter.scala @@ -73,6 +73,23 @@ case class PaimonV2DataWriter( } } + private val cleanup = new SparkAttemptCleanup( + writeBuilder.tableName(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder), + writeBuilder, + () => closeLocalResources()) + + override protected def attemptCleanup: Option[SparkAttemptCleanup] = Some(cleanup) + + private def closeLocalResources(): Unit = { + try { + val closeables = Seq[AutoCloseable](write) ++ plainWrite.toSeq ++ Seq(ioManager) + IOUtils.closeAll(closeables.asJava) + } catch { + case e: Exception => throw new RuntimeException(e) + } + } + private def createRowConverter( writeSchema: StructType, schema: StructType): InternalRow => SparkInternalRowWrapper = { @@ -95,6 +112,7 @@ case class PaimonV2DataWriter( private val joinedRow = new JoinedRow() override def write(record: InternalRow): Unit = { + cleanup.checkInterruptedPeriodically() plainRowConverter match { case Some(converter) => postWrite(getPlainWrite.writeAndReturn(converter.apply(record))) @@ -113,22 +131,33 @@ case class PaimonV2DataWriter( } override def commitImpl(): Seq[CommitMessage] = { - val metadataMessages = write.prepareCommit().asScala.toSeq - val plainMessages = plainWrite.map(_.prepareCommit().asScala.toSeq).getOrElse(Seq.empty) - metadataMessages ++ plainMessages + val result = scala.collection.mutable.ListBuffer[CommitMessage]() + write.prepareCommit( + (msg: CommitMessage) => { + registerPrepared(Seq(msg)) + result += msg + }) + plainWrite.foreach { + plain => + plain.prepareCommit( + (msg: CommitMessage) => { + registerPrepared(Seq(msg)) + result += msg + }) + } + result.toSeq } - override def abort(): Unit = close() - - override def close(): Unit = { - try { - val closeables = Seq[AutoCloseable](write) ++ plainWrite.toSeq ++ Seq(ioManager) - IOUtils.closeAll(closeables.asJava) - } catch { - case e: Exception => throw new RuntimeException(e) - } + override def abort(): Unit = { + // abortPrepared deletes prepared files (or closes the writer if still writing). Delegate + // residual writer/ioManager close to cleanup.close() so failures are swallowed consistently + // and we do not double-close outside the cleanup state machine. + cleanup.abortPrepared() + cleanup.close() } + override def close(): Unit = cleanup.close() + override def currentMetricsValues(): Array[CustomTaskMetric] = { metricRegistry.buildSparkWriteMetrics() } diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptCleanup.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptCleanup.scala new file mode 100644 index 000000000000..1d15e46b9725 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptCleanup.scala @@ -0,0 +1,379 @@ +/* + * 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.paimon.spark.write + +import org.apache.paimon.table.sink.{BatchTableCommit, BatchWriteBuilder, CommitMessage} + +import org.apache.spark.TaskContext +import org.apache.spark.TaskKilledException +import org.apache.spark.internal.Logging + +import java.util.concurrent.atomic.AtomicReference + +import scala.collection.JavaConverters._ + +/** + * Task-side cleanup helper for Spark speculative execution and task interruption. + * + *

When a task attempt is killed by Spark (e.g. speculative execution loser), this helper aborts + * unprepared or prepared files before the attempt returns a successful result. + */ +final class SparkAttemptCleanup( + tableName: String, + commitUser: String, + writeBuilder: BatchWriteBuilder, + closeUnprepared: () => Unit) + extends Logging { + + private val INTERRUPT_CHECK_INTERVAL = 1024 + + sealed private trait State + private object State { + case object Writing extends State + case object Prepared extends State + case object Aborted extends State + case object Returned extends State + case object CloseFailed extends State + case object Closed extends State + } + + private val state = new AtomicReference[State](State.Writing) + + /** + * Incrementally accumulated prepared messages. Uses a mutable ListBuffer so each {@link + * #addPrepared} call is O(1) amortized; the list is only materialized once in {@link + * #abortIfNeeded} when a kill or close actually needs to delete the files. This avoids the O(n²) + * cost of repeatedly exporting an immutable Seq on every incremental prepareCommit. + */ + private val preparedMessages = scala.collection.mutable.ListBuffer.empty[CommitMessage] + @volatile private var preparedSnapshot: Option[Seq[CommitMessage]] = None + private var recordCount: Long = 0L + + registerCompletionListener() + + def checkInterrupted(stage: String): Unit = { + if (isTaskInterrupted) { + abortIfNeeded(s"interrupted at $stage") + throw taskKilledException() + } + } + + def checkInterruptedPeriodically(): Unit = { + recordCount += 1 + if (recordCount % INTERRUPT_CHECK_INTERVAL == 0) { + checkInterrupted("write") + } + } + + /** + * Replace the full prepared list. Convenience for one-shot prepareCommit callers; multi-file + * writers should prefer {@link #addPrepared} to avoid O(n²) copying. + */ + def setPrepared(messages: Seq[CommitMessage]): Unit = { + preparedMessages.clear() + preparedMessages ++= Option(messages).getOrElse(Seq.empty) + preparedSnapshot = None + transitionToPrepared() + } + + /** + * Incrementally register prepared messages. O(1) amortized; call this after each incremental + * prepareCommit so a mid-attempt kill can abort already-finished files. + */ + def addPrepared(messages: Seq[CommitMessage]): Unit = { + preparedMessages ++= Option(messages).getOrElse(Seq.empty) + preparedSnapshot = None + transitionToPrepared() + } + + private def transitionToPrepared(): Unit = { + while (true) { + val current = state.get() + current match { + case State.Writing | State.Prepared => + if (state.compareAndSet(current, State.Prepared)) { + return + } + case State.Aborted | State.Returned | State.CloseFailed | State.Closed => + logWarning( + s"Ignoring prepared registration for table $tableName because cleanup is already " + + s"in state $current.") + return + } + } + } + + /** Java-friendly wrapper for procedure callers. */ + def setPreparedJava(messages: java.util.List[CommitMessage]): Unit = { + setPrepared(Option(messages).map(_.asScala.toSeq).getOrElse(Seq.empty)) + } + + /** Java-friendly incremental registration for procedure callers. */ + def addPreparedJava(messages: java.util.List[CommitMessage]): Unit = { + addPrepared(Option(messages).map(_.asScala.toSeq).getOrElse(Seq.empty)) + } + + def markReturned(): Unit = { + // Force Returned from Writing or Prepared so close() does not treat the attempt as + // "close without return" and abort. Files remain reclaimable: Spark may still kill + // the attempt after commit() returns but before DataWritingSparkTaskResult is + // accepted (see abortIfNeeded / completion listener). Snapshot-published callers + // must use {@link #markCommitted} instead so a later abort cannot delete live files. + while (true) { + val current = state.get() + current match { + case State.Prepared | State.Writing => + if (state.compareAndSet(current, State.Returned)) { + return + } + case State.Returned | State.CloseFailed | State.Aborted | State.Closed => + return + } + } + } + + /** + * Mark that prepared files are already published into a snapshot (procedure / driver-local + * commit). Clears prepared messages so a later abort or interrupted completion listener cannot + * delete live snapshot files, then behaves like {@link #markReturned} for close(). + */ + def markCommitted(): Unit = { + preparedMessages.clear() + preparedSnapshot = Some(Seq.empty) + markReturned() + } + + def abortPrepared(): Unit = abortIfNeeded("manual abort") + + def close(): Unit = { + state.get() match { + case State.Closed => + // Already closed: a previous close() ran the writer/ioManager close. Spark always + // calls DataWriter#close() again after commit()/abort(), so this must be a no-op + // rather than a double-close. + () + case State.Returned => + // Normal success path: a writer/ioManager close failure must propagate so the + // driver does not commit an incomplete result. Only mark the attempt Closed after + // close succeeds. On failure for Spark result-handoff paths, abort prepared files + // immediately — V1 / SparkAttemptWrite only close once and never call abort(). + try { + runWithClearedThreadInterrupt { + closeUnprepared() + } + state.compareAndSet(State.Returned, State.Closed) + } catch { + case e: Exception => + state.compareAndSet(State.Returned, State.CloseFailed) + abortIfNeeded("returned writer close failure") + throw e + } + case State.CloseFailed => + // Follow-up close after a Returned-path close failure. Abort if still needed + // (idempotent when the failure path already aborted). + abortIfNeeded("close after returned writer close failure") + case State.Aborted => + // Already aborted: best-effort residual close, suppress failures (the task is + // already failing). + runWithClearedThreadInterrupt { + safeCloseUnprepared() + } + case State.Prepared => + abortIfNeeded("close without return") + runWithClearedThreadInterrupt { + safeCloseUnprepared() + } + case State.Writing => + if (state.compareAndSet(State.Writing, State.Closed)) { + runWithClearedThreadInterrupt { + safeCloseUnprepared() + } + } else { + // State advanced concurrently (e.g. setPrepared/abort); handle the new state. + close() + } + } + } + + private def registerCompletionListener(): Unit = { + val context = TaskContext.get() + if (context != null) { + context.addTaskCompletionListener[Unit]( + _ => { + val current = state.get() + // Spark kill: reclaim even from Returned/Closed — markReturned/close run before + // Spark accepts DataWritingSparkTaskResult, so a speculative kill in that window + // would otherwise leave prepared files. Only Spark's explicit kill + // (context.isInterrupted) triggers this, NOT a stray thread interrupt flag: + // a successful winner attempt must never see its committed files deleted by a + // false-positive thread interrupt. CloseFailed without interrupt: reclaim when + // the caller never invokes abort()/close again. + if ( + current != State.Aborted && (context.isInterrupted() || current == State.CloseFailed) + ) { + abortIfNeeded("task completion listener") + } + }) + } + } + + private def abortIfNeeded(reason: String): Unit = { + val current = state.get() + if (current == State.Aborted) { + return + } + // Returned/Closed still abortable: commit() calls markReturned()+close() before Spark + // constructs/delivers DataWritingSparkTaskResult. A loser kill in that window sets + // interrupted and/or invokes DataWriter.abort(); both must reclaim prepared files. + // Snapshot-published callers use markCommitted() which clears prepared messages first. + if (!state.compareAndSet(current, State.Aborted)) { + return + } + + val context = TaskContext.get() + val stageId = if (context != null) context.stageId() else -1 + val partitionId = if (context != null) context.partitionId() else -1 + val attemptNumber = if (context != null) context.attemptNumber() else -1 + + // Spark kill sets Thread.interrupt(); clear it around abort deletes so HDFS RPC is + // not rejected with InterruptedIOException, then restore for Spark task teardown. + // Always close the unprepared writer as well: multi-file writers (data evolution, + // append compact) can have both prepared messages and an in-flight file open. + runWithClearedThreadInterrupt { + val messages = preparedSnapshot.getOrElse { + val s = preparedMessages.toSeq + preparedSnapshot = Some(s) + s + } + try { + if (messages.nonEmpty) { + logInfo( + s"Aborting ${messages.size} prepared commit messages for table $tableName " + + s"(commitUser=$commitUser, stage=$stageId, partition=$partitionId, " + + s"attempt=$attemptNumber) due to $reason.") + var tableCommit: BatchTableCommit = null + try { + tableCommit = writeBuilder.newCommit() + tableCommit.abort(messages.asJava) + SparkAttemptCleanup.notifyAbortedMessagesProbe(messages) + } catch { + case e: Exception => + logWarning( + s"Failed to abort prepared commit messages for table $tableName " + + s"(commitUser=$commitUser, stage=$stageId, partition=$partitionId, " + + s"attempt=$attemptNumber).", + e + ) + } finally { + if (tableCommit != null) { + try tableCommit.close() + catch { case _: Exception => } + } + } + } else { + logInfo( + s"Closing unprepared writer for table $tableName (commitUser=$commitUser, " + + s"stage=$stageId, partition=$partitionId, attempt=$attemptNumber) due to $reason.") + } + } finally { + // Writer cleanup must run even when creating the abort commit itself fails. + safeCloseUnprepared() + } + } + } + + private def safeCloseUnprepared(): Unit = { + try closeUnprepared() + catch { + case e: Exception => + logWarning(s"Failed to close unprepared writer for table $tableName.", e) + } + } + + /** + * Temporarily clears the thread interrupt flag for best-effort abort/cleanup IO, then restores + * it. Hadoop RPC deletes fail with InterruptedIOException while the flag is set. + */ + private def runWithClearedThreadInterrupt(block: => Unit): Unit = { + val interrupted = Thread.interrupted() + try { + block + } finally { + if (interrupted) { + Thread.currentThread().interrupt() + } + } + } + + private def isTaskInterrupted: Boolean = { + val context = TaskContext.get() + (context != null && context.isInterrupted()) || Thread.currentThread().isInterrupted + } + + private def taskKilledException(): TaskKilledException = { + new TaskKilledException("Paimon writer interrupted for speculative execution cleanup") + } +} + +object SparkAttemptCleanup { + + /** + * Best-effort commit user for logging. {@link BatchWriteBuilder#commitUser()} is an optional + * extension point whose default implementation throws; builders that do not expose a commit user + * (e.g. format tables or custom implementations) must not break writer construction. The value is + * only used for abort logging — abort itself goes through {@link BatchWriteBuilder#newCommit()}. + */ + def commitUserOrUnknown(writeBuilder: BatchWriteBuilder): String = { + try { + writeBuilder.commitUser() + } catch { + case _: UnsupportedOperationException => "" + } + } + + /** + * Optional test probe invoked after prepared messages are aborted successfully. Not for + * production use; IT cases install a probe to assert synchronous loser cleanup before orphan + * cleaner runs. The probe is JVM-global: ITs that install it must not run concurrently in the + * same JVM (each test must clear it afterwards, see {@code @AfterEach} usage). + */ + @volatile private var abortedMessagesProbe + : java.util.function.Consumer[java.util.List[CommitMessage]] = null + + /** Install or clear the abort probe used by speculative-write IT cases. */ + def setAbortedMessagesProbe( + probe: java.util.function.Consumer[java.util.List[CommitMessage]]): Unit = { + abortedMessagesProbe = probe + } + + private def notifyAbortedMessagesProbe(messages: Seq[CommitMessage]): Unit = { + val probe = abortedMessagesProbe + if (probe != null) { + probe.accept(messages.asJava) + } + } + + def forJava( + tableName: String, + commitUser: String, + writeBuilder: BatchWriteBuilder, + closeUnprepared: Runnable): SparkAttemptCleanup = { + new SparkAttemptCleanup(tableName, commitUser, writeBuilder, () => closeUnprepared.run()) + } +} diff --git a/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptWrite.scala b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptWrite.scala new file mode 100644 index 000000000000..2d8944a148fe --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/main/scala/org/apache/paimon/spark/write/SparkAttemptWrite.scala @@ -0,0 +1,88 @@ +/* + * 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.paimon.spark.write + +import org.apache.paimon.table.FileStoreTable +import org.apache.paimon.table.sink.{BatchWriteBuilder, CommitMessage} + +import scala.collection.JavaConverters._ + +/** + * Shared helper for Spark task-side write attempts used by procedures (compact / rescale / chain + * compact) and other Java/Scala callers that prepare commit messages inside mapPartitions. + * + *

V1 Dataset writes ({@code PaimonDataWrite} / {@code PaimonV2DataWriter}) embed {@link + * SparkAttemptCleanup} directly rather than going through this helper. + */ +object SparkAttemptWrite { + + def run[R]( + table: FileStoreTable, + writeBuilder: BatchWriteBuilder, + closeUnprepared: () => Unit, + write: SparkAttemptCleanup => Unit, + prepareCommit: SparkAttemptCleanup => Seq[CommitMessage], + toResult: Seq[CommitMessage] => R): R = { + val cleanup = + new SparkAttemptCleanup( + table.fullName(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder), + writeBuilder, + closeUnprepared) + try { + write(cleanup) + cleanup.checkInterrupted("before prepareCommit") + val messages = prepareCommit(cleanup) + cleanup.checkInterrupted("after prepareCommit") + val result = toResult(messages) + cleanup.checkInterrupted("before return") + cleanup.markReturned() + result + } finally { + cleanup.close() + } + } + + /** Java-friendly entry point for procedures and other Java callers. */ + def runJava[R]( + table: FileStoreTable, + writeBuilder: BatchWriteBuilder, + closeUnprepared: Runnable, + write: java.util.function.Consumer[SparkAttemptCleanup], + prepareCommit: java.util.function.Function[ + SparkAttemptCleanup, + java.util.List[CommitMessage]], + toResult: java.util.function.Function[java.util.List[CommitMessage], R]): R = { + run( + table, + writeBuilder, + () => runUnchecked(closeUnprepared), + cleanup => write.accept(cleanup), + cleanup => Option(prepareCommit.apply(cleanup)).map(_.asScala.toSeq).getOrElse(Seq.empty), + messages => toResult.apply(messages.asJava) + ) + } + + private def runUnchecked(action: Runnable): Unit = { + try action.run() + catch { + case e: Exception => throw new RuntimeException(e) + } + } +} diff --git a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/DataEvolutionTableDataWriteTest.scala b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/DataEvolutionTableDataWriteTest.scala new file mode 100644 index 000000000000..3d2ee7667358 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/DataEvolutionTableDataWriteTest.scala @@ -0,0 +1,207 @@ +/* + * 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.paimon.spark + +import org.apache.paimon.CoreOptions +import org.apache.paimon.data.{BinaryRow, BinaryString, GenericRow} +import org.apache.paimon.fs.{FileStatus, Path} +import org.apache.paimon.fs.local.LocalFileIO +import org.apache.paimon.schema.{Schema, SchemaManager, TableSchema} +import org.apache.paimon.spark.write.DataEvolutionTableDataWrite +import org.apache.paimon.table.{FileStoreTable, FileStoreTableFactory} +import org.apache.paimon.table.sink.{BatchTableWrite, BatchWriteBuilder, CommitMessage, CommitMessageImpl} +import org.apache.paimon.types.DataTypes +import org.apache.paimon.utils.{SerializationUtils, UriReaderFactory} + +import org.apache.spark.sql.Row +import org.apache.spark.sql.SparkSession +import org.junit.jupiter.api.{AfterAll, Assertions, BeforeAll, Test, TestInstance} +import org.junit.jupiter.api.io.TempDir + +import java.util.{ArrayList => JArrayList, Arrays => JArrays, Collections, List => JList} + +import scala.collection.JavaConverters._ +import scala.collection.mutable + +/** Regression tests for {@link DataEvolutionTableDataWrite} speculative-execution cleanup. */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +class DataEvolutionTableDataWriteTest { + + @BeforeAll + def startSpark(): Unit = { + SparkSession.builder + .master("local[1]") + .appName("DataEvolutionTableDataWriteTest") + .getOrCreate() + } + + @AfterAll + def stopSpark(): Unit = { + SparkSession.getActiveSession.orElse(SparkSession.getDefaultSession).foreach(_.stop()) + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + } + + @TempDir + private var tempDir: java.nio.file.Path = null + + private val fileIO = LocalFileIO.create() + + private def buildTable(): FileStoreTable = { + val tablePath = new Path(tempDir.toString) + val schema = Schema.newBuilder + .column("f0", DataTypes.INT()) + .column("f1", DataTypes.STRING()) + .column("f2", DataTypes.STRING()) + .option(CoreOptions.ROW_TRACKING_ENABLED.key(), "true") + .option(CoreOptions.DATA_EVOLUTION_ENABLED.key(), "true") + .option("file.format", "parquet") + .option("manifest.format", "avro") + .option("metadata.stats-mode", "none") + .build() + val schemaManager = new SchemaManager(fileIO, tablePath) + val tableSchema: TableSchema = schemaManager.createTable(schema) + FileStoreTableFactory.create(fileIO, tablePath, tableSchema) + } + + private def setFirstRowId(messages: JList[CommitMessage], firstRowId: Long): Unit = { + messages.asScala.foreach { + c => + val commitMessage = c.asInstanceOf[CommitMessageImpl] + val newFiles = new JArrayList(commitMessage.newFilesIncrement().newFiles()) + commitMessage.newFilesIncrement().newFiles().clear() + newFiles.asScala.foreach { + file => + commitMessage.newFilesIncrement().newFiles().add(file.assignFirstRowId(firstRowId)) + } + } + } + + /** Two committed data files with firstRowId 0 and 10, one row each. */ + private def seedTwoFileTable(table: FileStoreTable): Unit = { + val builder = table.newBatchWriteBuilder() + val writeType01 = table.rowType().project(JArrays.asList("f0", "f1")) + + val write: BatchTableWrite = builder.newWrite().withWriteType(table.rowType()) + try { + write.write( + GenericRow.of(Int.box(1), BinaryString.fromString("a"), BinaryString.fromString("b"))) + val commit1 = builder.newCommit() + try { + commit1.commit(write.prepareCommit()) + } finally { + commit1.close() + } + } finally { + write.close() + } + + val write01: BatchTableWrite = builder.newWrite().withWriteType(writeType01) + try { + write01.write(GenericRow.of(Int.box(2), BinaryString.fromString("c"))) + val messages = write01.prepareCommit() + setFirstRowId(messages, 10L) + val commit2 = builder.newCommit() + try { + commit2.commit(messages) + } finally { + commit2.close() + } + } finally { + write01.close() + } + } + + private def buildPartitionMap( + table: FileStoreTable): mutable.HashMap[Long, (Array[Byte], Long)] = { + val map = new mutable.HashMap[Long, (Array[Byte], Long)] + table + .store() + .newScan() + .readFileIterator() + .forEachRemaining { + entry => + map.put( + entry.file().firstRowId(), + (SerializationUtils.serializeBinaryRow(entry.partition()), entry.file().rowCount())) + } + map + } + + private def countParquetFiles(table: FileStoreTable): Int = { + val collected = new JArrayList[FileStatus]() + listRecursive(fileIO, table.location(), collected) + collected.asScala.count(_.getPath.getName.endsWith(".parquet")) + } + + private def listRecursive( + io: org.apache.paimon.fs.FileIO, + path: Path, + out: JList[FileStatus]): Unit = { + try { + io.listStatus(path).foreach { + status => + if (status.isDir) listRecursive(io, status.getPath, out) + else out.add(status) + } + } catch { + case _: java.io.FileNotFoundException => + } + } + + private def evolutionRow(f2: String, rowId: Long, firstRowId: Long): Row = { + Row(f2, rowId, firstRowId) + } + + @Test + def closeAfterFinishCurrentWriterAbortsPreparedPerFileMessages(): Unit = { + val table = buildTable() + seedTwoFileTable(table) + val baseline = countParquetFiles(table) + Assertions.assertEquals(2, baseline) + + val writeBuilder: BatchWriteBuilder = table.newBatchWriteBuilder() + val writeType = table.rowType().project(Collections.singletonList("f2")) + val partitionMap = buildPartitionMap(table) + val firstRowIds = partitionMap.keys.toSeq.sorted + Assertions.assertEquals(2, firstRowIds.size) + + val uriReaderFactory = UriReaderFactory.fromFileIO(LocalFileIO.create()) + val writer = + DataEvolutionTableDataWrite( + writeBuilder, + writeType, + partitionMap, + uriReaderFactory, + Map.empty) + try { + writer.write(evolutionRow("new-f2-0", firstRowIds(0), firstRowIds(0))) + // Switching firstRowId finishes the first PerFileWriter (prepare + addPrepared). + writer.write(evolutionRow("new-f2-1", firstRowIds(1), firstRowIds(1))) + writer.close() + } finally { + writer.close() + } + + Assertions.assertEquals( + baseline, + countParquetFiles(table), + "close without return must abort incrementally prepared per-file messages") + } +} diff --git a/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/write/SparkAttemptCleanupTest.scala b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/write/SparkAttemptCleanupTest.scala new file mode 100644 index 000000000000..8c9fa0068c74 --- /dev/null +++ b/paimon-spark/paimon-spark-common/src/test/scala/org/apache/paimon/spark/write/SparkAttemptCleanupTest.scala @@ -0,0 +1,649 @@ +/* + * 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.paimon.spark.write + +import org.apache.paimon.data.{BinaryString, GenericRow} +import org.apache.paimon.disk.IOManagerImpl +import org.apache.paimon.fs.{FileStatus, Path} +import org.apache.paimon.fs.local.LocalFileIO +import org.apache.paimon.schema.{Schema, SchemaManager, TableSchema} +import org.apache.paimon.table.{FileStoreTable, FileStoreTableFactory} +import org.apache.paimon.table.sink.{BatchTableCommit, BatchTableWrite, BatchWriteBuilder, TableWriteImpl, WriteSelector} +import org.apache.paimon.types.{IntType, RowType, VarCharType} + +import org.junit.jupiter.api.{Assertions, Test} +import org.junit.jupiter.api.io.TempDir + +import java.util +import java.util.concurrent.atomic.AtomicInteger +import java.util.function.Consumer + +import scala.collection.JavaConverters._ + +class SparkAttemptCleanupTest { + + @TempDir + private var tempDir: java.nio.file.Path = null + + private val fileIO = LocalFileIO.create() + + private def buildTable(): FileStoreTable = { + val tablePath = new Path(tempDir.toString) + val schema = Schema.newBuilder + .column("id", new IntType()) + .column("v", new VarCharType()) + .option("file.format", "parquet") + .option("manifest.format", "avro") + .option("metadata.stats-mode", "none") + .build() + val schemaManager = new SchemaManager(fileIO, tablePath) + val tableSchema: TableSchema = schemaManager.createTable(schema) + FileStoreTableFactory.create(fileIO, tablePath, tableSchema) + } + + private def countParquetFiles(table: FileStoreTable): Int = { + val collected = new util.ArrayList[FileStatus]() + listRecursive(fileIO, table.location(), collected) + collected.asScala.count(_.getPath.getName.endsWith(".parquet")) + } + + private def listRecursive( + io: org.apache.paimon.fs.FileIO, + path: Path, + out: util.List[FileStatus]): Unit = { + try { + io.listStatus(path).foreach { + status => + if (status.isDir) listRecursive(io, status.getPath, out) + else out.add(status) + } + } catch { + case _: java.io.FileNotFoundException => + } + } + + @Test + def abortPreparedAlsoClosesUnpreparedWriter(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => closeCount.incrementAndGet()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + // Simulate multi-file writer: prepared messages exist AND an unprepared close callback. + cleanup.abortPrepared() + } finally { + write.close() + cleanup.close() + } + + Assertions.assertEquals( + 2, + closeCount.get(), + "abort must close unprepared writer even when prepared messages are present; close() also closes") + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def abortPreparedDeletesUnpublishedFiles(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => closeCount.incrementAndGet()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.abortPrepared() + } finally { + write.close() + cleanup.close() + } + + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def abortPreparedUnderThreadInterruptStillDeletesFiles(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + Thread.currentThread().interrupt() + try { + cleanup.abortPrepared() + Assertions.assertTrue(Thread.currentThread().isInterrupted) + } finally { + Thread.interrupted() + } + } finally { + write.close() + cleanup.close() + } + + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def markReturnedDoesNotAbortPreparedFilesOnClose(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => closeCount.incrementAndGet()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.markReturned() + } finally { + write.close() + cleanup.close() + } + + Assertions.assertEquals(1, closeCount.get()) + Assertions.assertEquals(1, countParquetFiles(table)) + } + + @Test + def abortAfterMarkReturnedAndCloseReclaimsPreparedFiles(): Unit = { + // P1: commit() marks Returned and close() reaches Closed before Spark accepts the + // DataWritingSparkTaskResult. A subsequent DataWriter.abort() must still delete files. + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.markReturned() + cleanup.close() + Assertions.assertEquals(1, countParquetFiles(table)) + + cleanup.abortPrepared() + Assertions.assertEquals( + 0, + countParquetFiles(table), + "Spark DataWriter.abort after markReturned+close must reclaim prepared files") + } finally { + write.close() + } + } + + @Test + def abortAfterMarkReturnedAndCloseUnderInterruptReclaimsPreparedFiles(): Unit = { + // P1: speculative kill after markReturned/close sets the interrupt flag; abort must + // still reclaim (completion listener and DataWriter.abort share abortIfNeeded). + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.markReturned() + cleanup.close() + Assertions.assertEquals(1, countParquetFiles(table)) + + Thread.currentThread().interrupt() + try { + cleanup.abortPrepared() + Assertions.assertTrue(Thread.currentThread().isInterrupted) + } finally { + Thread.interrupted() + } + Assertions.assertEquals(0, countParquetFiles(table)) + } finally { + write.close() + } + } + + @Test + def closeWithoutReturnAbortsPreparedFiles(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + } finally { + write.close() + cleanup.close() + } + + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def prepareCommitWithoutAddPreparedLeavesOrphanAfterClose(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + write.prepareCommit() + } finally { + // PerFileWriter.finish() used to close here before addPrepared; prepareCommit drains + // newFiles so close() no longer deletes them. + write.close() + } + + Assertions.assertEquals(1, countParquetFiles(table)) + cleanup.abortPrepared() + Assertions.assertEquals( + 1, + countParquetFiles(table), + "abort without addPrepared cannot reclaim drained files") + } + + @Test + def prepareCommitWithAddPreparedReclaimsAfterClose(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write = writeBuilder.newWrite().asInstanceOf[TableWriteImpl[_]] + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + write.prepareCommit(new Consumer[org.apache.paimon.table.sink.CommitMessage] { + override def accept(message: org.apache.paimon.table.sink.CommitMessage): Unit = + cleanup.addPreparedJava(java.util.Collections.singletonList(message)) + }) + } finally { + write.close() + } + + Assertions.assertEquals(1, countParquetFiles(table)) + cleanup.abortPrepared() + cleanup.close() + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def incrementalSetPreparedAbortsAllPreparedFiles(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val prepared = new util.ArrayList[org.apache.paimon.table.sink.CommitMessage]() + val write1: BatchTableWrite = writeBuilder.newWrite() + write1.withIOManager(new IOManagerImpl(tempDir.resolve("io1").toString)) + try { + write1.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + prepared.addAll(write1.prepareCommit()) + cleanup.setPreparedJava(prepared) + } finally { + write1.close() + } + + val write2: BatchTableWrite = writeBuilder.newWrite() + write2.withIOManager(new IOManagerImpl(tempDir.resolve("io2").toString)) + try { + write2.write(GenericRow.of(Int.box(2), BinaryString.fromString("b"))) + prepared.addAll(write2.prepareCommit()) + cleanup.setPreparedJava(prepared) + } finally { + write2.close() + } + + Assertions.assertTrue(countParquetFiles(table) >= 2) + cleanup.abortPrepared() + cleanup.close() + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def closeFailureOnReturnedPathAbortsPreparedFilesImmediately(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => { + closeCount.incrementAndGet() + throw new RuntimeException("injected writer close failure") + }) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.markReturned() + + val thrown = + Assertions.assertThrows(classOf[RuntimeException], () => cleanup.close()) + Assertions.assertTrue(thrown.getMessage.contains("injected writer close failure")) + Assertions.assertTrue( + closeCount.get() >= 1, + "Returned path must invoke closeUnprepared (and abort may also close)") + // V1 / SparkAttemptWrite only close once — abort must run on the failing close itself. + Assertions.assertEquals(0, countParquetFiles(table)) + } finally { + write.close() + } + } + + @Test + def newCommitFailureStillClosesUnpreparedWriter(): Unit = { + val table = buildTable() + val realWriteBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + realWriteBuilder.commitUser(), + newCommitFailingBuilder(realWriteBuilder), + () => closeCount.incrementAndGet()) + + val write: BatchTableWrite = realWriteBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + val messages = + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + write.prepareCommit() + } finally { + write.close() + } + + cleanup.setPreparedJava(messages) + cleanup.abortPrepared() + Assertions.assertEquals( + 1, + closeCount.get(), + "writer close must run even when constructing the abort commit fails") + + // The injected builder could not delete the prepared file, so clean it with the real builder. + val commit = realWriteBuilder.newCommit() + try commit.abort(messages) + finally commit.close() + } + + @Test + def closeOnAbortedPathSuppressesCloseFailure(): Unit = { + // Abort/cleanup path: close failures are suppressed because the task is already failing. + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val closeCount = new AtomicInteger() + + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => { + closeCount.incrementAndGet() + throw new RuntimeException("injected writer close failure") + }) + + // abortPrepared transitions Writing -> Aborted and suppresses the close failure. + cleanup.abortPrepared() + // close() on Aborted must also suppress, never propagate. + cleanup.close() + Assertions.assertTrue(closeCount.get() >= 1, "abort path must still attempt close") + } + + @Test + def markCommittedAfterPublishPreventsAbortOfCommittedFiles(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + val messages = + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + write.prepareCommit() + } finally { + write.close() + } + + cleanup.setPreparedJava(messages) + val commit = writeBuilder.newCommit() + try { + commit.commit(messages) + cleanup.markCommitted() + } finally { + commit.close() + } + + // Simulate close/interrupt/abort after successful snapshot publish: must not delete + // snapshot files (prepared messages were cleared by markCommitted). + Thread.currentThread().interrupt() + try { + cleanup.abortPrepared() + cleanup.close() + } finally { + Thread.interrupted() + } + Assertions.assertEquals(1, countParquetFiles(table)) + } + + @Test + def setPreparedIgnoredAfterAbort(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + writeBuilder.commitUser(), + writeBuilder, + () => ()) + + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io").toString)) + val messages = + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + write.prepareCommit() + } finally { + write.close() + } + + cleanup.setPreparedJava(messages) + cleanup.abortPrepared() + Assertions.assertEquals(0, countParquetFiles(table)) + + // Recreate files and try to re-register after abort: state must stay Aborted. + val write2: BatchTableWrite = writeBuilder.newWrite() + write2.withIOManager(new IOManagerImpl(tempDir.resolve("io2").toString)) + val messages2 = + try { + write2.write(GenericRow.of(Int.box(2), BinaryString.fromString("b"))) + write2.prepareCommit() + } finally { + write2.close() + } + Assertions.assertEquals(1, countParquetFiles(table)) + cleanup.setPreparedJava(messages2) + cleanup.close() + // Ignored setPrepared must not resurrect Prepared, so close does not abort messages2. + // Caller owns messages2 cleanup here. + Assertions.assertEquals(1, countParquetFiles(table)) + val abortCommit = writeBuilder.newCommit() + try { + abortCommit.abort(messages2) + } finally { + abortCommit.close() + } + Assertions.assertEquals(0, countParquetFiles(table)) + } + + @Test + def commitUserOrUnknownFallsBackForBuildersWithoutCommitUser(): Unit = { + val table = buildTable() + val writeBuilder = table.newBatchWriteBuilder() + Assertions.assertEquals( + writeBuilder.commitUser(), + SparkAttemptCleanup.commitUserOrUnknown(writeBuilder)) + + // A custom BatchWriteBuilder that does not override commitUser(): the default interface + // implementation throws, and the fallback must keep writer construction working. + val builderWithoutCommitUser = new BatchWriteBuilder { + override def tableName(): String = writeBuilder.tableName() + + override def rowType(): RowType = writeBuilder.rowType() + + override def newWriteSelector(): java.util.Optional[WriteSelector] = + writeBuilder.newWriteSelector() + + override def newWrite(): BatchTableWrite = writeBuilder.newWrite() + + override def newCommit(): BatchTableCommit = writeBuilder.newCommit() + + override def withOverwrite( + staticPartition: java.util.Map[String, String]): BatchWriteBuilder = { + writeBuilder.withOverwrite(staticPartition) + this + } + } + Assertions.assertThrows( + classOf[UnsupportedOperationException], + () => builderWithoutCommitUser.commitUser()) + Assertions.assertEquals( + "", + SparkAttemptCleanup.commitUserOrUnknown(builderWithoutCommitUser)) + + // Construction with such a builder must succeed; abort still goes through newCommit(). + val cleanup = SparkAttemptCleanup.forJava( + table.fullName(), + SparkAttemptCleanup.commitUserOrUnknown(builderWithoutCommitUser), + builderWithoutCommitUser, + () => ()) + val write: BatchTableWrite = writeBuilder.newWrite() + write.withIOManager(new IOManagerImpl(tempDir.resolve("io-unknown").toString)) + try { + write.write(GenericRow.of(Int.box(1), BinaryString.fromString("a"))) + val messages = write.prepareCommit() + cleanup.setPreparedJava(messages) + cleanup.abortPrepared() + } finally { + write.close() + cleanup.close() + } + Assertions.assertEquals(0, countParquetFiles(table)) + } + + private def newCommitFailingBuilder(delegate: BatchWriteBuilder): BatchWriteBuilder = { + new BatchWriteBuilder { + override def tableName(): String = delegate.tableName() + + override def rowType(): RowType = delegate.rowType() + + override def newWriteSelector(): java.util.Optional[WriteSelector] = + delegate.newWriteSelector() + + override def newWrite(): BatchTableWrite = delegate.newWrite() + + override def newCommit(): BatchTableCommit = + throw new RuntimeException("injected newCommit failure") + + override def withOverwrite( + staticPartition: java.util.Map[String, String]): BatchWriteBuilder = { + delegate.withOverwrite(staticPartition) + this + } + + override def commitUser(): String = delegate.commitUser() + } + } +} diff --git a/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkSpeculativeWriteITCase.java b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkSpeculativeWriteITCase.java new file mode 100644 index 000000000000..c72427518fa3 --- /dev/null +++ b/paimon-spark/paimon-spark-ut/src/test/java/org/apache/paimon/spark/SparkSpeculativeWriteITCase.java @@ -0,0 +1,455 @@ +/* + * 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.paimon.spark; + +import org.apache.paimon.data.BinaryString; +import org.apache.paimon.data.GenericRow; +import org.apache.paimon.disk.IOManagerImpl; +import org.apache.paimon.fs.FileIO; +import org.apache.paimon.fs.FileStatus; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.local.LocalFileIO; +import org.apache.paimon.io.DataFileMeta; +import org.apache.paimon.operation.LocalOrphanFilesClean; +import org.apache.paimon.spark.extensions.PaimonSparkSessionExtensions; +import org.apache.paimon.spark.write.SparkAttemptCleanup; +import org.apache.paimon.table.FileStoreTable; +import org.apache.paimon.table.FileStoreTableFactory; +import org.apache.paimon.table.sink.BatchTableWrite; +import org.apache.paimon.table.sink.BatchWriteBuilder; +import org.apache.paimon.table.sink.CommitMessage; +import org.apache.paimon.table.sink.CommitMessageImpl; +import org.apache.paimon.table.source.DataSplit; +import org.apache.paimon.table.source.ReadBuilder; +import org.apache.paimon.table.source.Split; + +import org.apache.spark.TaskContext; +import org.apache.spark.scheduler.SparkListener; +import org.apache.spark.scheduler.SparkListenerSpeculativeTaskSubmitted; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.api.java.UDF1; +import org.apache.spark.sql.types.DataTypes; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.io.TempDir; + +import java.util.HashSet; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.apache.spark.sql.functions.callUDF; +import static org.apache.spark.sql.functions.col; +import static org.apache.spark.sql.functions.lit; +import static org.apache.spark.sql.functions.pmod; +import static org.apache.spark.sql.functions.rpad; +import static org.apache.spark.sql.functions.when; +import static org.assertj.core.api.Assertions.assertThat; + +/** IT cases for Spark speculative execution with Paimon DSv2 writer. */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public class SparkSpeculativeWriteITCase { + + private SparkSession spark; + + private Path warehousePath; + + // Counts speculative task submissions observed during the current test. Reset per test that + // cares about speculation firing, so prior tests do not pollute the count. + private final AtomicInteger speculativeTaskCount = new AtomicInteger(); + + // Records task attempt ids that have already taken the injected slow-path sleep, so each + // attempt sleeps at most once (on its first part-0 row) instead of once per row. Shared + // across tasks because local[8] runs them in one JVM; cleared before each skew test. + private static final Set SLEPT_ATTEMPTS = ConcurrentHashMap.newKeySet(); + + // File names from CommitMessages that SparkAttemptCleanup aborted during the current test. + // Used to assert synchronous loser cleanup before orphan cleaner runs. + private static final Set ABORTED_FILE_NAMES = ConcurrentHashMap.newKeySet(); + + @BeforeAll + public void startSpark(@TempDir java.nio.file.Path tempDir) { + warehousePath = new Path("file:///" + tempDir.toString()); + SparkSession.clearActiveSession(); + SparkSession.clearDefaultSession(); + spark = + SparkSession.builder() + .master("local[8]") + .appName("SparkSpeculativeWriteITCase") + .config( + "spark.sql.extensions", + PaimonSparkSessionExtensions.class.getName()) + .config("spark.sql.catalog.paimon", SparkCatalog.class.getName()) + .config("spark.sql.catalog.paimon.warehouse", warehousePath.toString()) + .config("spark.sql.defaultCatalog", "paimon") + .config("spark.speculation", "true") + .config("spark.speculation.multiplier", "1.0") + .config("spark.speculation.quantile", "0.0") + .config("spark.speculation.interval", "50") + .config("spark.speculation.minTaskRuntime", "1") + .getOrCreate(); + // Record when Spark actually launches a duplicate (speculative) attempt, so the test can + // confirm the loser-cleanup path was exercised rather than silently skipped. Note: some + // Spark builds do not schedule the speculation executor in local mode, so the count may + // stay 0; the safety assertions (data correctness + no orphan files after clean) hold + // regardless and are the primary guarantee. + spark.sparkContext() + .addSparkListener( + new SparkListener() { + @Override + public void onSpeculativeTaskSubmitted( + SparkListenerSpeculativeTaskSubmitted event) { + speculativeTaskCount.incrementAndGet(); + } + }); + spark.sql("CREATE DATABASE db"); + spark.sql("USE db"); + + // UDF that sleeps once per task attempt on part 0. It is applied AFTER a repartition in + // testSpeculativeLoserFilesAreCleanedUp, so it executes inside the write-stage task (the + // task Spark would speculate and kill), maximizing the chance that speculation fires in + // environments where it is supported. + spark.udf() + .register( + "slowOnPart0", + (UDF1) + part -> { + if (part != null && part == 0) { + TaskContext ctx = TaskContext.get(); + if (ctx != null + && SLEPT_ATTEMPTS.add(ctx.taskAttemptId())) { + try { + Thread.sleep(3000L); + } catch (InterruptedException e) { + // Spark kills the loser by interrupting it; + // preserve the interrupt so the writer abort + // path runs. + Thread.currentThread().interrupt(); + } + } + } + return part; + }, + DataTypes.IntegerType); + } + + @AfterAll + public void stopSpark() { + if (spark != null) { + spark.stop(); + spark = null; + } + } + + @AfterEach + public void dropTable() { + SparkAttemptCleanup.setAbortedMessagesProbe(null); + ABORTED_FILE_NAMES.clear(); + spark.sql("DROP TABLE IF EXISTS spec_t"); + } + + @Test + public void testInsertWithSpeculationEnabled() throws Exception { + spark.sql("CREATE TABLE spec_t (id BIGINT, part INT) USING paimon"); + spark.sql( + "INSERT INTO spec_t " + + "SELECT id, CAST(pmod(id, 8) AS INT) " + + "FROM range(0, 2000, 1, 16)") + .collectAsList(); + + List rows = + spark.sql("SELECT count(*), count(DISTINCT id) FROM spec_t").collectAsList(); + assertThat(rows.get(0).getLong(0)).isEqualTo(2000L); + assertThat(rows.get(0).getLong(1)).isEqualTo(2000L); + + FileStoreTable table = loadTable("spec_t"); + Assertions.assertNotNull(table.snapshotManager().latestSnapshot()); + assertReferencedDataFilesExist(table); + assertOrphanCleanerReclaimsUnreferencedFiles(table); + } + + @Test + public void testV1WriteWithSpeculationEnabled() throws Exception { + spark.sql("CREATE TABLE spec_t (id BIGINT, v STRING) USING paimon"); + spark.sql( + "INSERT INTO spec_t " + + "SELECT id, concat('v-', CAST(id AS STRING)) " + + "FROM range(0, 1000, 1, 12)") + .collectAsList(); + + List rows = spark.sql("SELECT count(*) FROM spec_t").collectAsList(); + assertThat(rows.get(0).getLong(0)).isEqualTo(1000L); + + FileStoreTable table = loadTable("spec_t"); + assertReferencedDataFilesExist(table); + assertOrphanCleanerReclaimsUnreferencedFiles(table); + } + + @Test + public void testSpeculativeWriteDoesNotLeaveExcessiveOrphanFiles() throws Exception { + spark.sql("CREATE TABLE spec_t (id BIGINT, part INT) USING paimon"); + spark.sql( + "INSERT INTO spec_t " + + "SELECT id, CAST(pmod(id, 16) AS INT) " + + "FROM range(0, 5000, 1, 40)") + .collectAsList(); + + FileStoreTable table = loadTable("spec_t"); + assertReferencedDataFilesExist(table); + assertOrphanCleanerReclaimsUnreferencedFiles(table); + } + + @Test + public void testSpeculativeLoserFilesAreCleanedUpSynchronously() throws Exception { + // Drive a real Spark task skew so the loser-abort path is exercised when speculation + // fires: + // * part is a real partition key (PARTITIONED BY), and repartition(8, part) + // concentrates all ~50000 part-0 rows into a single write task while parts 1..7 get + // one row each — a genuine per-task skew, not just a skewed data column. + // * the slowOnPart0 UDF is applied AFTER the repartition, so it runs in the write + // stage (not the map stage) and sleeps once per attempt on part 0, giving speculation + // a deterministic trigger in environments that schedule the speculation executor. + // When a duplicate attempt wins, Spark kills the slow original mid-write, exercising the + // abort/close path that must delete the loser's in-flight and prepared files. + // + // Sync cleanup is verified BEFORE LocalOrphanFilesClean runs: aborted file names recorded + // by SparkAttemptCleanup must already be gone from disk. Orphan cleaner fallback is covered + // by testOrphanCleanerFallbackReclaimsUnreferencedFiles, not this test. + spark.sql( + "CREATE TABLE spec_t (id BIGINT, part INT, payload STRING) USING paimon " + + "PARTITIONED BY (part)"); + + SLEPT_ATTEMPTS.clear(); + speculativeTaskCount.set(0); + ABORTED_FILE_NAMES.clear(); + SparkAttemptCleanup.setAbortedMessagesProbe( + messages -> { + for (CommitMessage message : messages) { + CommitMessageImpl impl = (CommitMessageImpl) message; + for (DataFileMeta file : impl.newFilesIncrement().newFiles()) { + ABORTED_FILE_NAMES.add(file.fileName()); + } + for (DataFileMeta file : impl.newFilesIncrement().changelogFiles()) { + ABORTED_FILE_NAMES.add(file.fileName()); + } + for (DataFileMeta file : impl.compactIncrement().compactAfter()) { + ABORTED_FILE_NAMES.add(file.fileName()); + } + for (DataFileMeta file : impl.compactIncrement().changelogFiles()) { + ABORTED_FILE_NAMES.add(file.fileName()); + } + } + }); + + // First select builds the skewed part column + a modest payload (real parquet IO weight + // on top of the injected sleep). repartition(8, part) puts each part value in its own + // write task. The second select applies slowOnPart0 AFTER the shuffle, so the UDF (and + // its sleep) executes inside the write-stage task. + Dataset skewed = + spark.range(0, 50007, 1, 8) + .select( + col("id"), + when(col("id").lt(50000), lit(0)) + .otherwise(pmod(col("id"), lit(7)).plus(lit(1))) + .cast("int") + .as("part"), + rpad(lit("x"), 256, "x").as("payload")) + .repartition(8, col("part")) + .select( + col("id"), + callUDF("slowOnPart0", col("part")).as("part"), + col("payload")); + + skewed.write().format("paimon").mode("append").saveAsTable("spec_t"); + + List rows = + spark.sql("SELECT count(*), count(DISTINCT id) FROM spec_t").collectAsList(); + assertThat(rows.get(0).getLong(0)).isEqualTo(50007L); + assertThat(rows.get(0).getLong(1)).isEqualTo(50007L); + + FileStoreTable table = loadTable("spec_t"); + assertReferencedDataFilesExist(table); + + // When speculation is supported in the runtime, it must actually have launched a + // duplicate attempt for the skewed part-0 task. Use assumeTrue so environments without a + // working speculation scheduler (this local UT build) skip this strict check rather than + // fail, while environments that do support speculation enforce it. + org.junit.jupiter.api.Assumptions.assumeTrue( + speculativeTaskCount.get() > 0, + "Spark speculation did not fire in this environment (speculativeTaskCount=0); " + + "skipping the sync loser-cleanup assertion. Data-correctness assertions " + + "above still hold."); + + // Sync cleanup assertion — must run BEFORE orphan cleaner. Files aborted by + // SparkAttemptCleanup must already be absent from disk; running LocalOrphanFilesClean + // first would hide a broken sync path. + Set referenced = referencedDataFileNames(table); + Set physicalBeforeClean = new HashSet<>(); + collectParquetFileNames(table.fileIO(), table.location(), physicalBeforeClean); + + // Best-effort sync cleanup: when SparkAttemptCleanup did abort prepared files, they must + // already be gone from disk before the orphan cleaner runs. Losers killed before + // prepareCommit or successful losers ignored by the driver may leave no aborted names; + // orphan-cleaner fallback is covered by testOrphanCleanerFallbackReclaimsUnreferencedFiles. + for (String aborted : ABORTED_FILE_NAMES) { + Assertions.assertFalse( + physicalBeforeClean.contains(aborted), + "loser file aborted by SparkAttemptCleanup must be gone before orphan cleaner: " + + aborted); + Assertions.assertFalse( + referenced.contains(aborted), + "aborted loser file must not be snapshot-referenced: " + aborted); + } + } + + @Test + public void testOrphanCleanerFallbackReclaimsUnreferencedFiles() throws Exception { + // Independent of SparkAttemptCleanup: plant an unreferenced parquet file, run the orphan + // cleaner, and assert it disappears. Sync loser cleanup is covered elsewhere; this test + // only validates the fallback path that reclaim successful-loser / hard-kill leftovers. + spark.sql("CREATE TABLE spec_t (id INT, v STRING) USING paimon"); + spark.sql("INSERT INTO spec_t VALUES (1, 'a')").collectAsList(); + + FileStoreTable table = loadTable("spec_t"); + Set referencedBefore = referencedDataFileNames(table); + Assertions.assertFalse(referencedBefore.isEmpty()); + + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + BatchTableWrite write = writeBuilder.newWrite(); + java.nio.file.Path ioDir = java.nio.file.Files.createTempDirectory("spec-orphan-io"); + write.withIOManager(new IOManagerImpl(ioDir.toString())); + Set plantedOrphans = new HashSet<>(); + try { + write.write(GenericRow.of(2, BinaryString.fromString("orphan"))); + List messages = write.prepareCommit(); + for (CommitMessage message : messages) { + CommitMessageImpl impl = (CommitMessageImpl) message; + for (DataFileMeta file : impl.newFilesIncrement().newFiles()) { + plantedOrphans.add(file.fileName()); + } + } + } finally { + write.close(); + } + Assertions.assertFalse(plantedOrphans.isEmpty(), "should plant at least one orphan file"); + + Set physicalBeforeClean = new HashSet<>(); + collectParquetFileNames(table.fileIO(), table.location(), physicalBeforeClean); + for (String orphan : plantedOrphans) { + Assertions.assertTrue( + physicalBeforeClean.contains(orphan), + "planted orphan should exist before cleaner: " + orphan); + Assertions.assertFalse( + referencedBefore.contains(orphan), + "planted orphan must not be snapshot-referenced: " + orphan); + } + + assertOrphanCleanerReclaimsUnreferencedFiles(table); + + Set physicalAfterClean = new HashSet<>(); + collectParquetFileNames(table.fileIO(), table.location(), physicalAfterClean); + for (String orphan : plantedOrphans) { + Assertions.assertFalse( + physicalAfterClean.contains(orphan), + "orphan cleaner must remove planted unreferenced file: " + orphan); + } + } + + private FileStoreTable loadTable(String tableName) { + Path tablePath = new Path(warehousePath, "db.db/" + tableName); + return FileStoreTableFactory.create(LocalFileIO.create(), tablePath); + } + + private void assertReferencedDataFilesExist(FileStoreTable table) { + Set referenced = referencedDataFileNames(table); + Assertions.assertFalse(referenced.isEmpty(), "snapshot should reference data files"); + + FileIO fileIO = table.fileIO(); + Set physical = new HashSet<>(); + collectParquetFileNames(fileIO, table.location(), physical); + + for (String fileName : referenced) { + Assertions.assertTrue( + physical.contains(fileName), + "referenced data file should exist on disk: " + fileName); + } + } + + private void assertOrphanCleanerReclaimsUnreferencedFiles(FileStoreTable table) + throws Exception { + // Fallback path only: run LocalOrphanFilesClean then assert physical == referenced. + // Do not use this to validate SparkAttemptCleanup — that would pass even when sync + // abort is broken. Use a future older-than threshold so freshly written files are + // eligible. + new LocalOrphanFilesClean(table, System.currentTimeMillis() + TimeUnit.HOURS.toMillis(1)) + .clean(); + + Set referenced = referencedDataFileNames(table); + FileIO fileIO = table.fileIO(); + Set physical = new HashSet<>(); + collectParquetFileNames(fileIO, table.location(), physical); + + Assertions.assertTrue( + physical.size() >= referenced.size(), + "physical files should cover all snapshot-referenced files"); + Assertions.assertEquals( + referenced.size(), + physical.size(), + "no orphan parquet files should remain after orphan clean, referenced=" + + referenced.size() + + ", physical=" + + physical.size()); + } + + private Set referencedDataFileNames(FileStoreTable table) { + Set referenced = new HashSet<>(); + ReadBuilder readBuilder = table.newReadBuilder(); + List splits = readBuilder.newScan().plan().splits(); + for (Split split : splits) { + if (split instanceof DataSplit) { + for (DataFileMeta file : ((DataSplit) split).dataFiles()) { + referenced.add(file.fileName()); + } + } + } + return referenced; + } + + private void collectParquetFileNames(FileIO fileIO, Path path, Set out) { + try { + for (FileStatus status : fileIO.listStatus(path)) { + if (status.isDir()) { + collectParquetFileNames(fileIO, status.getPath(), out); + } else if (status.getPath().getName().endsWith(".parquet")) { + out.add(status.getPath().getName()); + } + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } +}