[core] Support writing into postpone bucket tables with table API (#5437)

tsreaper · web-flow · commit 9422995a3dc8 · 2025-04-10T10:49:02.000+08:00
diff --git a/paimon-core/src/main/java/org/apache/paimon/AbstractFileStore.java b/paimon-core/src/main/java/org/apache/paimon/AbstractFileStore.java
@@ -121,10 +121,10 @@ protected AbstractFileStore(
 
     @Override
     public FileStorePathFactory pathFactory() {
-        return pathFactory(options.fileFormatString());
+        return pathFactory(options, options.fileFormatString());
     }
 
-    protected FileStorePathFactory pathFactory(String format) {
+    protected FileStorePathFactory pathFactory(CoreOptions options, String format) {
         return new FileStorePathFactory(
                 options.path(),
                 partitionType,
diff --git a/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java b/paimon-core/src/main/java/org/apache/paimon/KeyValueFileStore.java
@@ -170,6 +170,7 @@ public AbstractFileStoreWrite<KeyValue> newWrite(
             return new PostponeBucketFileStoreWrite(
                     fileIO,
                     schema,
+                    commitUser,
                     partitionType,
                     keyType,
                     valueType,
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
@@ -99,10 +99,7 @@ protected AbstractFileStoreWrite(
             @Nullable DeletionVectorsMaintainer.Factory dvMaintainerFactory,
             String tableName,
             CoreOptions options,
-            int numBuckets,
-            RowType partitionType,
-            int writerNumberMax,
-            boolean legacyPartitionName) {
+            RowType partitionType) {
         this.snapshotManager = snapshotManager;
         this.scan = scan;
         // Statistic is useless in writer
@@ -113,12 +110,12 @@ protected AbstractFileStoreWrite(
         }
         this.indexFactory = indexFactory;
         this.dvMaintainerFactory = dvMaintainerFactory;
-        this.numBuckets = numBuckets;
+        this.numBuckets = options.bucket();
         this.partitionType = partitionType;
         this.writers = new HashMap<>();
         this.tableName = tableName;
-        this.writerNumberMax = writerNumberMax;
-        this.legacyPartitionName = legacyPartitionName;
+        this.writerNumberMax = options.writeMaxWritersToSpill();
+        this.legacyPartitionName = options.legacyPartitionName();
     }
 
     @Override
diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/KeyValueFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/KeyValueFileStoreWrite.java
@@ -82,6 +82,7 @@
 import java.util.Comparator;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
@@ -124,7 +125,7 @@ public KeyValueFileStoreWrite(
             Supplier<RecordEqualiser> logDedupEqualSupplier,
             MergeFunctionFactory<KeyValue> mfFactory,
             FileStorePathFactory pathFactory,
-            Function<String, FileStorePathFactory> formatPathFactory,
+            BiFunction<CoreOptions, String, FileStorePathFactory> formatPathFactory,
             SnapshotManager snapshotManager,
             FileStoreScan scan,
             @Nullable IndexMaintainer.Factory<KeyValue> indexFactory,
diff --git a/paimon-core/src/main/java/org/apache/paimon/operation/MemoryFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/operation/MemoryFileStoreWrite.java
@@ -74,10 +74,7 @@ public MemoryFileStoreWrite(
                 dvMaintainerFactory,
                 tableName,
                 options,
-                options.bucket(),
-                partitionType,
-                options.writeMaxWritersToSpill(),
-                options.legacyPartitionName());
+                partitionType);
         this.options = options;
         this.cacheManager =
                 new CacheManager(
diff --git a/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketFileStoreWrite.java b/paimon-core/src/main/java/org/apache/paimon/postpone/PostponeBucketFileStoreWrite.java
@@ -28,6 +28,7 @@
 import org.apache.paimon.operation.AbstractFileStoreWrite;
 import org.apache.paimon.operation.FileStoreScan;
 import org.apache.paimon.operation.FileStoreWrite;
+import org.apache.paimon.options.Options;
 import org.apache.paimon.schema.TableSchema;
 import org.apache.paimon.table.BucketMode;
 import org.apache.paimon.types.RowType;
@@ -40,6 +41,8 @@
 
 import java.util.List;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.function.BiFunction;
 import java.util.function.Function;
 
 import static org.apache.paimon.utils.FileStorePathFactory.createFormatPathFactories;
@@ -53,32 +56,31 @@ public class PostponeBucketFileStoreWrite extends AbstractFileStoreWrite<KeyValu
     public PostponeBucketFileStoreWrite(
             FileIO fileIO,
             TableSchema schema,
+            String commitUser,
             RowType partitionType,
             RowType keyType,
             RowType valueType,
-            Function<String, FileStorePathFactory> formatPathFactory,
+            BiFunction<CoreOptions, String, FileStorePathFactory> formatPathFactory,
             SnapshotManager snapshotManager,
             FileStoreScan scan,
             CoreOptions options,
             String tableName) {
-        super(
-                snapshotManager,
-                scan,
-                null,
-                null,
-                tableName,
-                options,
-                options.bucket(),
-                partitionType,
-                options.writeMaxWritersToSpill(),
-                options.legacyPartitionName());
-
-        // copy options for postpone bucket
-        this.options = new CoreOptions(options.toConfiguration().toMap());
+        super(snapshotManager, scan, null, null, tableName, options, partitionType);
 
+        Options newOptions = new Options(options.toMap());
         // use avro for postpone bucket
-        this.options.toConfiguration().set(CoreOptions.FILE_FORMAT, "avro");
-        this.options.toConfiguration().set(CoreOptions.METADATA_STATS_MODE, "none");
+        newOptions.set(CoreOptions.FILE_FORMAT, "avro");
+        newOptions.set(CoreOptions.METADATA_STATS_MODE, "none");
+        // each writer should have its unique prefix, so files from the same writer can be consumed
+        // by the same compaction reader to keep the input order
+        newOptions.set(
+                CoreOptions.DATA_FILE_PREFIX,
+                String.format(
+                        "%s-u-%s-s-%d-w-",
+                        options.dataFilePrefix(),
+                        commitUser,
+                        ThreadLocalRandom.current().nextInt(Integer.MAX_VALUE)));
+        this.options = new CoreOptions(newOptions);
 
         this.writerFactoryBuilder =
                 KeyValueFileWriterFactory.builder(
@@ -88,7 +90,7 @@ public PostponeBucketFileStoreWrite(
                         valueType,
                         this.options.fileFormat(),
                         createFormatPathFactories(this.options, formatPathFactory),
-                        options.targetFileSize(true));
+                        this.options.targetFileSize(true));
 
         // Ignoring previous files saves scanning time.
         //
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java b/paimon-core/src/main/java/org/apache/paimon/table/PrimaryKeyFileStoreTable.java
@@ -21,6 +21,7 @@
 import org.apache.paimon.CoreOptions;
 import org.apache.paimon.KeyValue;
 import org.apache.paimon.KeyValueFileStore;
+import org.apache.paimon.annotation.VisibleForTesting;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
 import org.apache.paimon.manifest.ManifestCacheFilter;
@@ -32,6 +33,8 @@
 import org.apache.paimon.schema.KeyValueFieldsExtractor;
 import org.apache.paimon.schema.TableSchema;
 import org.apache.paimon.table.query.LocalTableQuery;
+import org.apache.paimon.table.sink.PostponeBucketRowKeyExtractor;
+import org.apache.paimon.table.sink.RowKeyExtractor;
 import org.apache.paimon.table.sink.TableWriteImpl;
 import org.apache.paimon.table.source.InnerTableRead;
 import org.apache.paimon.table.source.KeyValueTableRead;
@@ -55,6 +58,7 @@ public class PrimaryKeyFileStoreTable extends AbstractFileStoreTable {
 
     private transient KeyValueFileStore lazyStore;
 
+    @VisibleForTesting
     PrimaryKeyFileStoreTable(FileIO fileIO, Path path, TableSchema tableSchema) {
         this(fileIO, path, tableSchema, CatalogEnvironment.empty());
     }
@@ -133,7 +137,7 @@ protected BiConsumer<FileStoreScan, Predicate> nonPartitionFilterConsumer() {
                             splitAnd(predicate),
                             tableSchema.fieldNames(),
                             tableSchema.trimmedPrimaryKeys());
-            if (keyFilters.size() > 0) {
+            if (!keyFilters.isEmpty()) {
                 ((KeyValueFileStoreScan) scan).withKeyFilter(and(keyFilters));
             }
 
@@ -185,4 +189,13 @@ protected Runnable newExpireRunnable() {
             return super.newExpireRunnable();
         }
     }
+
+    @Override
+    public RowKeyExtractor createRowKeyExtractor() {
+        if (coreOptions().bucket() == BucketMode.POSTPONE_BUCKET) {
+            return new PostponeBucketRowKeyExtractor(schema());
+        } else {
+            return super.createRowKeyExtractor();
+        }
+    }
 }
diff --git a/paimon-core/src/main/java/org/apache/paimon/table/sink/PostponeBucketRowKeyExtractor.java b/paimon-core/src/main/java/org/apache/paimon/table/sink/PostponeBucketRowKeyExtractor.java
@@ -0,0 +1,38 @@
+/*
+ * 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.table.sink;
+
+import org.apache.paimon.schema.TableSchema;
+import org.apache.paimon.table.BucketMode;
+
+/**
+ * {@link RowKeyExtractor} for postpone bucket tables. The bucket is always {@link
+ * BucketMode#POSTPONE_BUCKET}.
+ */
+public class PostponeBucketRowKeyExtractor extends RowKeyExtractor {
+
+    public PostponeBucketRowKeyExtractor(TableSchema schema) {
+        super(schema);
+    }
+
+    @Override
+    public int bucket() {
+        return BucketMode.POSTPONE_BUCKET;
+    }
+}
diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/FileStorePathFactory.java b/paimon-core/src/main/java/org/apache/paimon/utils/FileStorePathFactory.java
@@ -37,7 +37,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.atomic.AtomicInteger;
-import java.util.function.Function;
+import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 
 /** Factory which produces {@link Path}s for manifest files. */
@@ -298,11 +298,13 @@ public Path toPath(String fileName) {
     }
 
     public static Map<String, FileStorePathFactory> createFormatPathFactories(
-            CoreOptions options, Function<String, FileStorePathFactory> formatPathFactory) {
+            CoreOptions options,
+            BiFunction<CoreOptions, String, FileStorePathFactory> formatPathFactory) {
         Map<String, FileStorePathFactory> pathFactoryMap = new HashMap<>();
         Set<String> formats = new HashSet<>(options.fileFormatPerLevel().values());
         formats.add(options.fileFormatString());
-        formats.forEach(format -> pathFactoryMap.put(format, formatPathFactory.apply(format)));
+        formats.forEach(
+                format -> pathFactoryMap.put(format, formatPathFactory.apply(options, format)));
         return pathFactoryMap;
     }
 }
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
@@ -32,6 +32,7 @@
 import org.apache.paimon.io.BundleRecords;
 import org.apache.paimon.io.DataFileMeta;
 import org.apache.paimon.manifest.FileKind;
+import org.apache.paimon.manifest.ManifestEntry;
 import org.apache.paimon.manifest.ManifestFileMeta;
 import org.apache.paimon.operation.FileStoreScan;
 import org.apache.paimon.options.MemorySize;
@@ -130,19 +131,23 @@ public class PrimaryKeySimpleTableTest extends SimpleTableTestBase {
 
     @Test
     public void testPostponeBucket() throws Exception {
-        FileStoreTable table = createFileStoreTable(options -> options.set(BUCKET, -2));
+        FileStoreTable table =
+                createFileStoreTable(options -> options.set(BUCKET, BucketMode.POSTPONE_BUCKET));
 
         BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder();
         try (BatchTableWrite write = writeBuilder.newWrite();
                 BatchTableCommit commit = writeBuilder.newCommit()) {
-            write.write(rowData(0, 0, 0L), BucketMode.POSTPONE_BUCKET);
+            write.write(rowData(0, 0, 0L));
             commit.commit(write.prepareCommit());
         }
 
         Snapshot snapshot = table.latestSnapshot().get();
         ManifestFileMeta manifest =
                 table.manifestListReader().read(snapshot.deltaManifestList()).get(0);
-        DataFileMeta file = table.manifestFileReader().read(manifest.fileName()).get(0).file();
+        ManifestEntry entry = table.manifestFileReader().read(manifest.fileName()).get(0);
+        assertThat(entry.bucket()).isEqualTo(BucketMode.POSTPONE_BUCKET);
+
+        DataFileMeta file = entry.file();
         assertThat(file.fileName()).endsWith(".avro");
         assertThat(file.level()).isEqualTo(0);
         assertThat(file.valueStatsCols()).isEmpty();
diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/postpone/PostponeBucketCompactSplitSource.java
@@ -179,12 +179,15 @@ public int channel(Split record) {
             Matcher matcher = pattern.matcher(fileName);
             Preconditions.checkState(
                     matcher.find(),
-                    "Data file name does not match the pattern. This is unexpected.");
-            int subtaskId = Integer.parseInt(matcher.group(1));
+                    "Data file name %s does not match the pattern. This is unexpected.",
+                    fileName);
 
+            // use long to avoid overflow
+            long subtaskId = Long.parseLong(matcher.group(1));
             // send records written by the same subtask to the same subtask
             // to make sure we replay the written records in the exact order
-            return (Math.abs(dataSplit.partition().hashCode()) + subtaskId) % numChannels;
+            long channel = (Math.abs(dataSplit.partition().hashCode()) + subtaskId) % numChannels;
+            return (int) channel;
         }
     }
 }
diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/FlinkSinkBuilder.java
@@ -56,6 +56,7 @@
 import static org.apache.paimon.CoreOptions.OrderType.HILBERT;
 import static org.apache.paimon.CoreOptions.OrderType.ORDER;
 import static org.apache.paimon.CoreOptions.OrderType.ZORDER;
+import static org.apache.paimon.CoreOptions.createCommitUser;
 import static org.apache.paimon.flink.FlinkConnectorOptions.CLUSTERING_SAMPLE_FACTOR;
 import static org.apache.paimon.flink.FlinkConnectorOptions.CLUSTERING_STRATEGY;
 import static org.apache.paimon.flink.FlinkConnectorOptions.MIN_CLUSTERING_SAMPLE_FACTOR;
@@ -297,7 +298,10 @@ protected DataStreamSink<?> buildForFixedBucket(DataStream<InternalRow> input) {
     }
 
     private DataStreamSink<?> buildPostponeBucketSink(DataStream<InternalRow> input) {
-        return new PostponeBucketWriteSink(table, overwritePartition).sinkFrom(input, parallelism);
+        FixedBucketSink sink = new FixedBucketSink(table, overwritePartition, null);
+        String commitUser = createCommitUser(table.coreOptions().toConfiguration());
+        DataStream<Committable> written = sink.doWrite(input, commitUser, parallelism);
+        return sink.doCommit(written, commitUser);
     }
 
     private DataStreamSink<?> buildUnawareBucketSink(DataStream<InternalRow> input) {
diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/PostponeBucketTableWriteOperator.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/PostponeBucketTableWriteOperator.java
diff --git a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/PostponeBucketWriteSink.java b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/PostponeBucketWriteSink.java
diff --git a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/TestChangelogDataReadWrite.java b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/TestChangelogDataReadWrite.java
