[core] Introduce withBlobConsumer in TableWrite

Author: JingsongLi

paimon-common/src/main/java/org/apache/paimon/data/BlobConsumer.java

@@ -0,0 +1,30 @@
+/*
+ * 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.data;
+
+/** A consumer for blob. */
+public interface BlobConsumer {
+
+    /**
+     * Accept a blob descriptor.
+     *
+     * @return Whether to flush to output stream.
+     */
+    boolean accept(BlobDescriptor blobDescriptor);
+}

paimon-common/src/main/java/org/apache/paimon/format/FileAwareFormatWriter.java

@@ -0,0 +1,31 @@
+/*
+ * 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.format;
+
+import org.apache.paimon.fs.Path;
+
+/** A {@link FormatWriter} that can be aware of the file. */
+public interface FileAwareFormatWriter extends FormatWriter {
+
+    /** Set the path of the file. */
+    void setFile(Path file);
+
+    /** Whether to delete the file upon abort. */
+    boolean deleteFileUponAbort();
+}

paimon-core/src/main/java/org/apache/paimon/append/AppendOnlyWriter.java

@@ -22,6 +22,7 @@
 import org.apache.paimon.compact.CompactDeletionFile;
 import org.apache.paimon.compact.CompactManager;
 import org.apache.paimon.compression.CompressOptions;
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.disk.IOManager;
 import org.apache.paimon.disk.RowBuffer;
@@ -92,6 +93,7 @@ public class AppendOnlyWriter implements BatchRecordWriter, MemoryOwner {
     @Nullable private final IOManager ioManager;
     private final FileIndexOptions fileIndexOptions;
     private final MemorySize maxDiskSize;
+    @Nullable private final BlobConsumer blobConsumer;
 
     @Nullable private CompactDeletionFile compactDeletionFile;
     private SinkWriter<InternalRow> sinkWriter;
@@ -120,7 +122,8 @@ public AppendOnlyWriter(
             MemorySize maxDiskSize,
             FileIndexOptions fileIndexOptions,
             boolean asyncFileWrite,
-            boolean statsDenseStore) {
+            boolean statsDenseStore,
+            @Nullable BlobConsumer blobConsumer) {
         this.fileIO = fileIO;
         this.schemaId = schemaId;
         this.fileFormat = fileFormat;
@@ -134,6 +137,7 @@ public AppendOnlyWriter(
         this.forceCompact = forceCompact;
         this.asyncFileWrite = asyncFileWrite;
         this.statsDenseStore = statsDenseStore;
+        this.blobConsumer = blobConsumer;
         this.newFiles = new ArrayList<>();
         this.deletedFiles = new ArrayList<>();
         this.compactBefore = new ArrayList<>();
@@ -311,7 +315,8 @@ private RollingFileWriter<InternalRow, DataFileMeta> createRollingRowWriter() {
                     fileIndexOptions,
                     FileSource.APPEND,
                     asyncFileWrite,
-                    statsDenseStore);
+                    statsDenseStore,
+                    blobConsumer);
         }
         return new RowDataRollingFileWriter(
                 fileIO,

paimon-core/src/main/java/org/apache/paimon/append/RollingBlobFileWriter.java

@@ -18,6 +18,7 @@
 
 package org.apache.paimon.append;
 
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.fileindex.FileIndexOptions;
 import org.apache.paimon.format.FileFormat;
@@ -44,6 +45,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.annotation.Nullable;
+
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -91,7 +94,6 @@ public class RollingBlobFileWriter implements RollingFileWriter<InternalRow, Dat
                             RollingFileWriterImpl<InternalRow, DataFileMeta>, List<DataFileMeta>>>
             blobWriterFactory;
     private final long targetFileSize;
-    private final long blobTargetFileSize;
 
     // State management
     private final List<FileWriterAbortExecutor> closedWriters;
@@ -118,11 +120,10 @@ public RollingBlobFileWriter(
             FileIndexOptions fileIndexOptions,
             FileSource fileSource,
             boolean asyncFileWrite,
-            boolean statsDenseStore) {
-
+            boolean statsDenseStore,
+            @Nullable BlobConsumer blobConsumer) {
         // Initialize basic fields
         this.targetFileSize = targetFileSize;
-        this.blobTargetFileSize = blobTargetFileSize;
         this.results = new ArrayList<>();
         this.closedWriters = new ArrayList<>();
 
@@ -161,7 +162,8 @@ public RollingBlobFileWriter(
                                 fileSource,
                                 asyncFileWrite,
                                 statsDenseStore,
-                                blobTargetFileSize);
+                                blobTargetFileSize,
+                                blobConsumer);
     }
 
     /** Creates a factory for normal data writers. */
@@ -221,9 +223,10 @@ public RollingBlobFileWriter(
                     FileSource fileSource,
                     boolean asyncFileWrite,
                     boolean statsDenseStore,
-                    long targetFileSize) {
-
+                    long targetFileSize,
+                    @Nullable BlobConsumer blobConsumer) {
         BlobFileFormat blobFileFormat = new BlobFileFormat();
+        blobFileFormat.setWriteConsumer(blobConsumer);
         List<String> blobNames = blobType.getFieldNames();
 
         // Validate blob field count
@@ -375,7 +378,7 @@ private void closeCurrentWriter() throws IOException {
     /** Closes the main writer and returns its metadata. */
     private DataFileMeta closeMainWriter() throws IOException {
         currentWriter.close();
-        closedWriters.add(currentWriter.writer().abortExecutor());
+        currentWriter.writer().abortExecutor().ifPresent(closedWriters::add);
         return currentWriter.result();
     }
 

paimon-core/src/main/java/org/apache/paimon/io/RollingFileWriterImpl.java

@@ -129,7 +129,7 @@ protected void closeCurrentWriter() throws IOException {
         // only store abort executor in memory
         // cannot store whole writer, it includes lots of memory for example column vectors to read
         // and write
-        closedWriters.add(currentWriter.abortExecutor());
+        currentWriter.abortExecutor().ifPresent(closedWriters::add);
         results.add(currentWriter.result());
         currentWriter = null;
     }

paimon-core/src/main/java/org/apache/paimon/io/SingleFileWriter.java

@@ -20,6 +20,7 @@
 
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.format.BundleFormatWriter;
+import org.apache.paimon.format.FileAwareFormatWriter;
 import org.apache.paimon.format.FormatWriter;
 import org.apache.paimon.format.FormatWriterFactory;
 import org.apache.paimon.format.SupportsDirectWrite;
@@ -36,6 +37,7 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.util.Optional;
 import java.util.function.Function;
 
 /**
@@ -51,9 +53,10 @@ public abstract class SingleFileWriter<T, R> implements FileWriter<T, R> {
     protected final FileIO fileIO;
     protected final Path path;
     private final Function<T, InternalRow> converter;
+    private final boolean deleteFileUponAbort;
 
     private FormatWriter writer;
-    private PositionOutputStream out;
+    @Nullable private PositionOutputStream out;
 
     @Nullable private Long outputBytes;
     private long recordCount;
@@ -72,14 +75,23 @@ public SingleFileWriter(
 
         try {
             if (factory instanceof SupportsDirectWrite) {
-                writer = ((SupportsDirectWrite) factory).create(fileIO, path, compression);
+                SupportsDirectWrite supportsDirectWrite = (SupportsDirectWrite) factory;
+                writer = supportsDirectWrite.create(fileIO, path, compression);
             } else {
                 out = fileIO.newOutputStream(path, false);
                 if (asyncWrite) {
                     out = new AsyncPositionOutputStream(out);
                 }
                 writer = factory.create(out, compression);
             }
+
+            if (writer instanceof FileAwareFormatWriter) {
+                FileAwareFormatWriter fileAwareFormatWriter = (FileAwareFormatWriter) writer;
+                fileAwareFormatWriter.setFile(path);
+                deleteFileUponAbort = fileAwareFormatWriter.deleteFileUponAbort();
+            } else {
+                deleteFileUponAbort = true;
+            }
         } catch (IOException e) {
             LOG.warn(
                     "Failed to open the bulk writer, closing the output stream and throw the error.",
@@ -118,7 +130,7 @@ public void writeBundle(BundleRecords bundle) throws IOException {
             }
             recordCount += bundle.rowCount();
         } catch (Throwable e) {
-            LOG.warn("Exception occurs when writing file " + path + ". Cleaning up.", e);
+            LOG.warn("Exception occurs when writing file {}. Cleaning up.", path, e);
             abort();
             throw e;
         }
@@ -135,7 +147,7 @@ protected InternalRow writeImpl(T record) throws IOException {
             recordCount++;
             return rowData;
         } catch (Throwable e) {
-            LOG.warn("Exception occurs when writing file " + path + ". Cleaning up.", e);
+            LOG.warn("Exception occurs when writing file {}. Cleaning up.", path, e);
             abort();
             throw e;
         }
@@ -160,15 +172,13 @@ public void abort() {
             IOUtils.closeQuietly(out);
             out = null;
         }
-        fileIO.deleteQuietly(path);
+        abortExecutor().ifPresent(FileWriterAbortExecutor::abort);
     }
 
-    public FileWriterAbortExecutor abortExecutor() {
-        if (!closed) {
-            throw new RuntimeException("Writer should be closed!");
-        }
-
-        return new FileWriterAbortExecutor(fileIO, path);
+    public Optional<FileWriterAbortExecutor> abortExecutor() {
+        return deleteFileUponAbort
+                ? Optional.of(new FileWriterAbortExecutor(fileIO, path))
+                : Optional.empty();
     }
 
     @Override

paimon-core/src/main/java/org/apache/paimon/operation/AbstractFileStoreWrite.java

@@ -24,6 +24,7 @@
 import org.apache.paimon.annotation.VisibleForTesting;
 import org.apache.paimon.compact.CompactDeletionFile;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.deletionvectors.BucketedDvMaintainer;
 import org.apache.paimon.disk.IOManager;
 import org.apache.paimon.index.DynamicBucketIndexMaintainer;
@@ -133,6 +134,11 @@ public FileStoreWrite<T> withMemoryPoolFactory(MemoryPoolFactory memoryPoolFacto
         return this;
     }
 
+    @Override
+    public FileStoreWrite<T> withBlobConsumer(BlobConsumer blobConsumer) {
+        return this;
+    }
+
     @Override
     public void withIgnorePreviousFiles(boolean ignorePreviousFiles) {
         this.ignorePreviousFiles = ignorePreviousFiles;

paimon-core/src/main/java/org/apache/paimon/operation/BaseAppendFileStoreWrite.java

@@ -24,6 +24,7 @@
 import org.apache.paimon.append.cluster.Sorter;
 import org.apache.paimon.compact.CompactManager;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.deletionvectors.BucketedDvMaintainer;
 import org.apache.paimon.deletionvectors.DeletionVector;
@@ -82,6 +83,7 @@ public abstract class BaseAppendFileStoreWrite extends MemoryFileStoreWrite<Inte
     private @Nullable List<String> writeCols;
     private boolean forceBufferSpill = false;
     private boolean withBlob;
+    private @Nullable BlobConsumer blobConsumer;
 
     public BaseAppendFileStoreWrite(
             FileIO fileIO,
@@ -109,6 +111,12 @@ public BaseAppendFileStoreWrite(
         this.fileIndexOptions = options.indexColumnsOptions();
     }
 
+    @Override
+    public BaseAppendFileStoreWrite withBlobConsumer(BlobConsumer blobConsumer) {
+        this.blobConsumer = blobConsumer;
+        return this;
+    }
+
     @Override
     protected RecordWriter<InternalRow> createWriter(
             BinaryRow partition,
@@ -142,7 +150,8 @@ protected RecordWriter<InternalRow> createWriter(
                 options.writeBufferSpillDiskSize(),
                 fileIndexOptions,
                 options.asyncFileWrite(),
-                options.statsDenseStore());
+                options.statsDenseStore(),
+                blobConsumer);
     }
 
     @Override

paimon-core/src/main/java/org/apache/paimon/operation/FileStoreWrite.java

@@ -20,6 +20,7 @@
 
 import org.apache.paimon.FileStore;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.deletionvectors.BucketedDvMaintainer;
 import org.apache.paimon.disk.IOManager;
 import org.apache.paimon.index.DynamicBucketIndexMaintainer;
@@ -64,6 +65,8 @@ default void withWriteType(RowType writeType) {
      */
     FileStoreWrite<T> withMemoryPoolFactory(MemoryPoolFactory memoryPoolFactory);
 
+    FileStoreWrite<T> withBlobConsumer(BlobConsumer blobConsumer);
+
     /**
      * Set whether the write operation should ignore previously stored files.
      *

paimon-core/src/main/java/org/apache/paimon/table/format/FormatTableWrite.java

@@ -21,6 +21,7 @@
 import org.apache.paimon.CoreOptions;
 import org.apache.paimon.casting.DefaultValueRow;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.BlobConsumer;
 import org.apache.paimon.data.InternalRow;
 import org.apache.paimon.disk.IOManager;
 import org.apache.paimon.fs.FileIO;
@@ -43,14 +44,13 @@
 /** {@link TableWrite} implementation for format table. */
 public class FormatTableWrite implements BatchTableWrite {
 
-    private RowType rowType;
+    private final RowType rowType;
     private final FormatTableFileWriter write;
     private final FormatTableRowPartitionKeyExtractor partitionKeyExtractor;
 
     private final int[] notNullFieldIndex;
     private final @Nullable DefaultValueRow defaultValueRow;
     private final ProjectedRow projectedRow;
-    private final RowType writeRowType;
 
     public FormatTableWrite(
             FileIO fileIO,
@@ -68,7 +68,7 @@ public FormatTableWrite(
                         .collect(Collectors.toList());
         this.notNullFieldIndex = rowType.getFieldIndices(notNullColumnNames);
         this.defaultValueRow = DefaultValueRow.create(rowType);
-        this.writeRowType =
+        RowType writeRowType =
                 rowType.project(
                         rowType.getFieldNames().stream()
                                 .filter(name -> !partitionType.getFieldNames().contains(name))
@@ -117,6 +117,11 @@ public TableWrite withMemoryPoolFactory(MemoryPoolFactory memoryPoolFactory) {
         return this;
     }
 
+    @Override
+    public TableWrite withBlobConsumer(BlobConsumer blobConsumer) {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public BatchTableWrite withIOManager(IOManager ioManager) {
         return this;