This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new d4a493d733 [flink] size aware partitioner for full compaction job
(#8583)
d4a493d733 is described below
commit d4a493d7332b20f7a77556efe8f58e187e419d43
Author: dwangatt <[email protected]>
AuthorDate: Sat Jul 18 15:35:23 2026 +1000
[flink] size aware partitioner for full compaction job (#8583)
---
docs/generated/flink_connector_configuration.html | 6 +
.../apache/paimon/flink/FlinkConnectorOptions.java | 38 +++++
.../apache/paimon/flink/action/CompactAction.java | 24 +++-
.../paimon/flink/action/CompactDatabaseAction.java | 11 +-
.../paimon/flink/sink/CompactorSinkBuilder.java | 29 +++-
.../flink/source/CompactorSourceBuilder.java | 103 +++++++++++++-
.../paimon/flink/source/StaticFileStoreSource.java | 85 ++++++++++-
.../source/assigners/PreAssignSplitAssigner.java | 142 +++++++++++++++++--
.../CompactionBucketDistributionStrategyTest.java | 70 +++++++++
.../flink/source/CompactorSourceBuilderTest.java | 157 +++++++++++++++++++++
.../paimon/flink/source/CompactorSourceITCase.java | 61 ++++++++
.../paimon/flink/source/FairAssignModeTest.java | 102 +++++++++++++
12 files changed, 805 insertions(+), 23 deletions(-)
diff --git a/docs/generated/flink_connector_configuration.html
b/docs/generated/flink_connector_configuration.html
index 83e9c0bb25..8f5b2eb9c3 100644
--- a/docs/generated/flink_connector_configuration.html
+++ b/docs/generated/flink_connector_configuration.html
@@ -26,6 +26,12 @@ under the License.
</tr>
</thead>
<tbody>
+ <tr>
+ <td><h5>compaction.bucket-distribution-strategy</h5></td>
+ <td style="word-wrap: break-word;">linear</td>
+ <td><p>Enum</p></td>
+ <td>Defines how dedicated bucket compaction jobs distribute
compact buckets to writers. 'linear' uses the existing stable
partition-plus-bucket mapping. 'size-aware-batch' assigns bounded
full-compaction bucket splits by total data file size and forwards them to
writers to reduce compaction long tail.<br /><br />Possible
values:<ul><li>"linear": Distribute compact buckets by the existing stable
partition-plus-bucket channel mapping.</li><li>"size-aware-batch": For bounded
ful [...]
+ </tr>
<tr>
<td><h5>changelog.precommit-compact.thread-num</h5></td>
<td style="word-wrap: break-word;">(none)</td>
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/FlinkConnectorOptions.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/FlinkConnectorOptions.java
index bc61b8bbab..1f0bf32c09 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/FlinkConnectorOptions.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/FlinkConnectorOptions.java
@@ -84,6 +84,16 @@ public class FlinkConnectorOptions {
+ "This eliminates data transfer overhead
when the source already "
+ "provides suitable data distribution
(e.g., Kafka partitions).");
+ public static final ConfigOption<CompactionBucketDistributionStrategy>
+ COMPACTION_BUCKET_DISTRIBUTION_STRATEGY =
+
ConfigOptions.key("compaction.bucket-distribution-strategy")
+
.enumType(CompactionBucketDistributionStrategy.class)
+
.defaultValue(CompactionBucketDistributionStrategy.LINEAR)
+ .withDescription(
+ "Defines how dedicated bucket compaction
jobs distribute compact buckets to writers. "
+ + "'linear' uses the existing
stable partition-plus-bucket mapping. "
+ + "'size-aware-batch' assigns
bounded full-compaction bucket splits by total data file size and forwards them
to writers to reduce compaction long tail.");
+
public static final ConfigOption<Boolean> INFER_SCAN_PARALLELISM =
ConfigOptions.key("scan.infer-parallelism")
.booleanType()
@@ -644,6 +654,34 @@ public class FlinkConnectorOptions {
}
}
+ /** Bucket distribution strategy for dedicated compaction jobs. */
+ public enum CompactionBucketDistributionStrategy implements DescribedEnum {
+ LINEAR(
+ "linear",
+ "Distribute compact buckets by the existing stable
partition-plus-bucket channel mapping."),
+ SIZE_AWARE_BATCH(
+ "size-aware-batch",
+ "For bounded full compaction, assign compact bucket splits by
total data file size and forward them to writers to reduce long-tail compaction
tasks.");
+
+ private final String value;
+ private final String description;
+
+ CompactionBucketDistributionStrategy(String value, String description)
{
+ this.value = value;
+ this.description = description;
+ }
+
+ @Override
+ public String toString() {
+ return value;
+ }
+
+ @Override
+ public InlineElement getDescription() {
+ return text(description);
+ }
+ }
+
/**
* Split assign mode for {@link
org.apache.paimon.flink.source.StaticFileStoreSplitEnumerator}.
*/
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
index 146c7b7a62..501b6aed16 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactAction.java
@@ -23,6 +23,7 @@ import org.apache.paimon.Snapshot;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.InternalRow;
import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
import org.apache.paimon.flink.compact.AppendTableCompact;
import org.apache.paimon.flink.compact.DataEvolutionTableCompact;
import org.apache.paimon.flink.compact.IncrementalClusterCompact;
@@ -166,6 +167,20 @@ public class CompactAction extends TableActionBase {
return true;
}
+ static CompactionBucketDistributionStrategy
compactionBucketDistributionStrategy(
+ FileStoreTable table, boolean fullCompaction, boolean isStreaming)
{
+ return compactionBucketDistributionStrategy(
+ table.coreOptions().toConfiguration(), fullCompaction,
isStreaming);
+ }
+
+ static CompactionBucketDistributionStrategy
compactionBucketDistributionStrategy(
+ Options options, boolean fullCompaction, boolean isStreaming) {
+ if (!fullCompaction || isStreaming) {
+ return CompactionBucketDistributionStrategy.LINEAR;
+ }
+ return
options.get(FlinkConnectorOptions.COMPACTION_BUCKET_DISTRIBUTION_STRATEGY);
+ }
+
protected void buildForBucketedTableCompact(
StreamExecutionEnvironment env, FileStoreTable table, boolean
isStreaming)
throws Exception {
@@ -192,9 +207,14 @@ public class CompactAction extends TableActionBase {
};
table = table.copy(dynamicOptions);
}
+ CompactionBucketDistributionStrategy bucketDistributionStrategy =
+ compactionBucketDistributionStrategy(table, fullCompaction,
isStreaming);
CompactorSourceBuilder sourceBuilder =
- new CompactorSourceBuilder(identifier.getFullName(), table);
- CompactorSinkBuilder sinkBuilder = new CompactorSinkBuilder(table,
fullCompaction);
+ new CompactorSourceBuilder(identifier.getFullName(), table)
+
.withBucketDistributionStrategy(bucketDistributionStrategy);
+ CompactorSinkBuilder sinkBuilder =
+ new CompactorSinkBuilder(table, fullCompaction)
+
.withBucketDistributionStrategy(bucketDistributionStrategy);
sourceBuilder.withPartitionPredicate(getPartitionPredicate());
DataStreamSource<RowData> source =
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactDatabaseAction.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactDatabaseAction.java
index 1a65b8d1bd..3000c319b5 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactDatabaseAction.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/action/CompactDatabaseAction.java
@@ -23,6 +23,7 @@ import org.apache.paimon.append.MultiTableAppendCompactTask;
import org.apache.paimon.catalog.Catalog;
import org.apache.paimon.catalog.Identifier;
import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
import org.apache.paimon.flink.compact.AppendTableCompact;
import org.apache.paimon.flink.sink.BucketsRowChannelComputer;
import org.apache.paimon.flink.sink.CombinedTableCompactorSink;
@@ -266,10 +267,16 @@ public class CompactDatabaseAction extends ActionBase {
table = table.copy(dynamicOptions);
}
+ CompactionBucketDistributionStrategy bucketDistributionStrategy =
+ CompactAction.compactionBucketDistributionStrategy(
+ table, fullCompaction, isStreaming);
CompactorSourceBuilder sourceBuilder =
new CompactorSourceBuilder(fullName, table)
- .withPartitionIdleTime(partitionIdleTime);
- CompactorSinkBuilder sinkBuilder = new CompactorSinkBuilder(table,
fullCompaction);
+ .withPartitionIdleTime(partitionIdleTime)
+
.withBucketDistributionStrategy(bucketDistributionStrategy);
+ CompactorSinkBuilder sinkBuilder =
+ new CompactorSinkBuilder(table, fullCompaction)
+
.withBucketDistributionStrategy(bucketDistributionStrategy);
DataStreamSource<RowData> source =
sourceBuilder.withEnv(env).withContinuousMode(isStreaming).build();
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/CompactorSinkBuilder.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/CompactorSinkBuilder.java
index 8e1321a9f7..778b2fd551 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/CompactorSinkBuilder.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/sink/CompactorSinkBuilder.java
@@ -19,15 +19,18 @@
package org.apache.paimon.flink.sink;
import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
import org.apache.paimon.table.BucketMode;
import org.apache.paimon.table.FileStoreTable;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSink;
+import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.table.data.RowData;
import java.util.Optional;
+import static org.apache.paimon.CoreOptions.createCommitUser;
import static org.apache.paimon.flink.sink.FlinkStreamPartitioner.partition;
/** Builder for {@link CompactorSink}. */
@@ -39,6 +42,9 @@ public class CompactorSinkBuilder {
private final boolean fullCompaction;
+ private CompactionBucketDistributionStrategy bucketDistributionStrategy =
+ CompactionBucketDistributionStrategy.LINEAR;
+
public CompactorSinkBuilder(FileStoreTable table, boolean fullCompaction) {
this.table = table;
this.fullCompaction = fullCompaction;
@@ -49,6 +55,12 @@ public class CompactorSinkBuilder {
return this;
}
+ public CompactorSinkBuilder withBucketDistributionStrategy(
+ CompactionBucketDistributionStrategy bucketDistributionStrategy) {
+ this.bucketDistributionStrategy = bucketDistributionStrategy;
+ return this;
+ }
+
public DataStreamSink<?> build() {
BucketMode bucketMode = table.bucketMode();
switch (bucketMode) {
@@ -66,8 +78,19 @@ public class CompactorSinkBuilder {
table.options().get(FlinkConnectorOptions.SINK_PARALLELISM.key()))
.map(Integer::valueOf)
.orElse(null);
- DataStream<RowData> partitioned =
- partition(input, new BucketsRowChannelComputer(), parallelism);
- return new CompactorSink(table, fullCompaction).sinkFrom(partitioned);
+ switch (bucketDistributionStrategy) {
+ case SIZE_AWARE_BATCH:
+ CompactorSink sink = new CompactorSink(table, fullCompaction);
+ String commitUser =
createCommitUser(table.coreOptions().toConfiguration());
+ DataStream<Committable> written = sink.doWrite(input,
commitUser, null);
+ return sink.doCommit(
+ ((SingleOutputStreamOperator<Committable>)
written).startNewChain(),
+ commitUser);
+ case LINEAR:
+ default:
+ DataStream<RowData> partitioned =
+ partition(input, new BucketsRowChannelComputer(),
parallelism);
+ return new CompactorSink(table,
fullCompaction).sinkFrom(partitioned);
+ }
}
}
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/CompactorSourceBuilder.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/CompactorSourceBuilder.java
index 83486a8624..96961272f1 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/CompactorSourceBuilder.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/CompactorSourceBuilder.java
@@ -21,12 +21,15 @@ package org.apache.paimon.flink.source;
import org.apache.paimon.CoreOptions;
import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
+import org.apache.paimon.flink.FlinkConnectorOptions.SplitAssignMode;
import org.apache.paimon.flink.LogicalTypeConversion;
import org.apache.paimon.manifest.PartitionEntry;
import org.apache.paimon.options.Options;
import org.apache.paimon.partition.PartitionPredicate;
import org.apache.paimon.partition.PartitionValuesTimeExpireStrategy;
import org.apache.paimon.table.FileStoreTable;
+import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.table.source.ReadBuilder;
import org.apache.paimon.table.system.CompactBucketsTable;
import org.apache.paimon.types.RowType;
@@ -66,6 +69,9 @@ public class CompactorSourceBuilder {
@Nullable private PartitionPredicate partitionPredicate = null;
@Nullable private Duration partitionIdleTime = null;
+ private CompactionBucketDistributionStrategy bucketDistributionStrategy =
+ CompactionBucketDistributionStrategy.LINEAR;
+
public CompactorSourceBuilder(String tableIdentifier, FileStoreTable
table) {
this.tableIdentifier = tableIdentifier;
this.table = table;
@@ -101,11 +107,22 @@ public class CompactorSourceBuilder {
return new ContinuousFileStoreSource(readBuilder,
compactBucketsTable.options(), null);
} else {
Options options =
compactBucketsTable.coreOptions().toConfiguration();
+ SplitAssignMode splitAssignMode =
+
options.get(FlinkConnectorOptions.SCAN_SPLIT_ENUMERATOR_ASSIGN_MODE);
+ validateBucketDistributionStrategy(bucketDistributionStrategy,
splitAssignMode);
return new StaticFileStoreSource(
readBuilder,
null,
options.get(FlinkConnectorOptions.SCAN_SPLIT_ENUMERATOR_BATCH_SIZE),
-
options.get(FlinkConnectorOptions.SCAN_SPLIT_ENUMERATOR_ASSIGN_MODE),
+ splitAssignMode,
+ bucketDistributionStrategy
+ ==
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH
+ ? split -> bucketFileSize((DataSplit)
split.split())
+ : null,
+ bucketDistributionStrategy
+ ==
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH
+ ? split -> bucketKey((DataSplit) split.split())
+ : null,
options.get(CoreOptions.BLOB_AS_DESCRIPTOR));
}
}
@@ -117,12 +134,21 @@ public class CompactorSourceBuilder {
CompactBucketsTable compactBucketsTable = new
CompactBucketsTable(table, isContinuous);
RowType produceType = compactBucketsTable.rowType();
+ Integer parallelism =
+ sourceParallelism(Options.fromMap(table.options()),
bucketDistributionStrategy);
DataStreamSource<RowData> dataStream =
env.fromSource(
buildSource(compactBucketsTable),
WatermarkStrategy.noWatermarks(),
tableIdentifier + "-compact-source",
InternalTypeInfo.of(LogicalTypeConversion.toLogicalType(produceType)));
+ if (parallelism != null) {
+ // Size-aware assignment depends on source-reader to writer
alignment. Set the
+ // parallelism on the original source before adding downstream
filters; otherwise the
+ // configured parallelism may only apply to a filter and
source-to-filter redistribution
+ // could break grouped bucket assignment.
+ dataStream.setParallelism(parallelism);
+ }
if (isContinuous) {
Preconditions.checkArgument(
partitionIdleTime == null, "Streaming mode does not
support partitionIdleTime");
@@ -140,6 +166,9 @@ public class CompactorSourceBuilder {
BinaryRow partition =
deserializeBinaryRow(rowData.getBinary(1));
return partitionInfo.get(partition) <=
historyMilli;
});
+ if (parallelism != null) {
+ filterStream.setParallelism(parallelism);
+ }
dataStream = new DataStreamSource<>(filterStream);
}
CoreOptions coreOptions = table.coreOptions();
@@ -160,13 +189,11 @@ public class CompactorSourceBuilder {
BinaryRow partition =
deserializeBinaryRow(rowData.getBinary(1));
return
!expireStrategy.isExpired(expireDateTime, partition);
});
+ if (parallelism != null) {
+ filterStream.setParallelism(parallelism);
+ }
dataStream = new DataStreamSource<>(filterStream);
}
- Integer parallelism =
-
Options.fromMap(table.options()).get(FlinkConnectorOptions.SCAN_PARALLELISM);
- if (parallelism != null) {
- dataStream.setParallelism(parallelism);
- }
return dataStream;
}
@@ -204,6 +231,70 @@ public class CompactorSourceBuilder {
return this;
}
+ public CompactorSourceBuilder withBucketDistributionStrategy(
+ CompactionBucketDistributionStrategy bucketDistributionStrategy) {
+ this.bucketDistributionStrategy = bucketDistributionStrategy;
+ return this;
+ }
+
+ static void validateBucketDistributionStrategy(
+ CompactionBucketDistributionStrategy bucketDistributionStrategy,
+ SplitAssignMode splitAssignMode) {
+ if (bucketDistributionStrategy ==
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH
+ && splitAssignMode == SplitAssignMode.PREEMPTIVE) {
+ throw new IllegalArgumentException(
+ "compaction.bucket-distribution-strategy=size-aware-batch
requires "
+ + "scan.split-enumerator.mode=fair because it
relies on grouped "
+ + "bucket assignment. Preemptive split assignment
can split the "
+ + "same bucket across different writers while the
sink skips "
+ + "bucket shuffle.");
+ }
+ }
+
+ static Integer sourceParallelism(
+ Options tableOptions, CompactionBucketDistributionStrategy
bucketDistributionStrategy) {
+ Integer parallelism =
tableOptions.get(FlinkConnectorOptions.SCAN_PARALLELISM);
+ if (bucketDistributionStrategy ==
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH) {
+ Integer sinkParallelism =
tableOptions.get(FlinkConnectorOptions.SINK_PARALLELISM);
+ if (sinkParallelism != null) {
+ parallelism = sinkParallelism;
+ }
+ }
+ return parallelism;
+ }
+
+ static long bucketFileSize(DataSplit split) {
+ return split.dataFiles().stream().mapToLong(dataFile ->
dataFile.fileSize()).sum();
+ }
+
+ static BucketKey bucketKey(DataSplit split) {
+ return new BucketKey(split.partition(), split.bucket());
+ }
+
+ static class BucketKey {
+ private final BinaryRow partition;
+ private final int bucket;
+
+ private BucketKey(BinaryRow partition, int bucket) {
+ this.partition = partition.copy();
+ this.bucket = bucket;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (!(o instanceof BucketKey)) {
+ return false;
+ }
+ BucketKey that = (BucketKey) o;
+ return bucket == that.bucket && partition.equals(that.partition);
+ }
+
+ @Override
+ public int hashCode() {
+ return 31 * partition.hashCode() + bucket;
+ }
+ }
+
private Map<BinaryRow, Long> getPartitionInfo(CompactBucketsTable table) {
List<PartitionEntry> partitions =
table.newSnapshotReader().partitionEntries();
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSource.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSource.java
index 940feb9d31..04facbd49e 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSource.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/StaticFileStoreSource.java
@@ -27,6 +27,7 @@ import org.apache.paimon.table.ChainGroupReadTable;
import org.apache.paimon.table.source.InnerTableScan;
import org.apache.paimon.table.source.ReadBuilder;
import org.apache.paimon.table.source.TableScan;
+import org.apache.paimon.utils.SerializableFunction;
import org.apache.flink.api.connector.source.Boundedness;
import org.apache.flink.api.connector.source.SplitEnumerator;
@@ -50,6 +51,10 @@ public class StaticFileStoreSource extends FlinkSource {
@Nullable private final DynamicPartitionFilteringInfo
dynamicPartitionFilteringInfo;
+ @Nullable private final SerializableFunction<FileStoreSourceSplit, Long>
splitWeightFunc;
+
+ @Nullable private final SerializableFunction<FileStoreSourceSplit, ?>
splitGroupFunc;
+
private final boolean skipPreloadTargetSnapshot;
public StaticFileStoreSource(
@@ -65,6 +70,51 @@ public class StaticFileStoreSource extends FlinkSource {
splitAssignMode,
null,
null,
+ null,
+ null,
+ blobAsDescriptor,
+ false);
+ }
+
+ public StaticFileStoreSource(
+ ReadBuilder readBuilder,
+ @Nullable Long limit,
+ int splitBatchSize,
+ SplitAssignMode splitAssignMode,
+ @Nullable DynamicPartitionFilteringInfo
dynamicPartitionFilteringInfo,
+ @Nullable NestedProjectedRowData rowData,
+ boolean blobAsDescriptor,
+ boolean skipPreloadTargetSnapshot) {
+ this(
+ readBuilder,
+ limit,
+ splitBatchSize,
+ splitAssignMode,
+ dynamicPartitionFilteringInfo,
+ rowData,
+ null,
+ null,
+ blobAsDescriptor,
+ skipPreloadTargetSnapshot);
+ }
+
+ public StaticFileStoreSource(
+ ReadBuilder readBuilder,
+ @Nullable Long limit,
+ int splitBatchSize,
+ SplitAssignMode splitAssignMode,
+ @Nullable SerializableFunction<FileStoreSourceSplit, Long>
splitWeightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?>
splitGroupFunc,
+ boolean blobAsDescriptor) {
+ this(
+ readBuilder,
+ limit,
+ splitBatchSize,
+ splitAssignMode,
+ null,
+ null,
+ splitWeightFunc,
+ splitGroupFunc,
blobAsDescriptor,
false);
}
@@ -76,12 +126,16 @@ public class StaticFileStoreSource extends FlinkSource {
SplitAssignMode splitAssignMode,
@Nullable DynamicPartitionFilteringInfo
dynamicPartitionFilteringInfo,
@Nullable NestedProjectedRowData rowData,
+ @Nullable SerializableFunction<FileStoreSourceSplit, Long>
splitWeightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?>
splitGroupFunc,
boolean blobAsDescriptor,
boolean skipPreloadTargetSnapshot) {
super(readBuilder, limit, rowData, blobAsDescriptor);
this.splitBatchSize = splitBatchSize;
this.splitAssignMode = splitAssignMode;
this.dynamicPartitionFilteringInfo = dynamicPartitionFilteringInfo;
+ this.splitWeightFunc = splitWeightFunc;
+ this.splitGroupFunc = splitGroupFunc;
this.skipPreloadTargetSnapshot = skipPreloadTargetSnapshot;
}
@@ -97,7 +151,13 @@ public class StaticFileStoreSource extends FlinkSource {
Collection<FileStoreSourceSplit> splits =
checkpoint == null ? getSplits(context) : checkpoint.splits();
SplitAssigner splitAssigner =
- createSplitAssigner(context, splitBatchSize, splitAssignMode,
splits);
+ createSplitAssigner(
+ context,
+ splitBatchSize,
+ splitAssignMode,
+ splits,
+ splitWeightFunc,
+ splitGroupFunc);
return new StaticFileStoreSplitEnumerator(
context, null, splitAssigner, dynamicPartitionFilteringInfo);
}
@@ -121,9 +181,30 @@ public class StaticFileStoreSource extends FlinkSource {
int splitBatchSize,
SplitAssignMode splitAssignMode,
Collection<FileStoreSourceSplit> splits) {
+ return createSplitAssigner(context, splitBatchSize, splitAssignMode,
splits, null, null);
+ }
+
+ public static SplitAssigner createSplitAssigner(
+ SplitEnumeratorContext<FileStoreSourceSplit> context,
+ int splitBatchSize,
+ SplitAssignMode splitAssignMode,
+ Collection<FileStoreSourceSplit> splits,
+ @Nullable SerializableFunction<FileStoreSourceSplit, Long>
splitWeightFunc) {
+ return createSplitAssigner(
+ context, splitBatchSize, splitAssignMode, splits,
splitWeightFunc, null);
+ }
+
+ public static SplitAssigner createSplitAssigner(
+ SplitEnumeratorContext<FileStoreSourceSplit> context,
+ int splitBatchSize,
+ SplitAssignMode splitAssignMode,
+ Collection<FileStoreSourceSplit> splits,
+ @Nullable SerializableFunction<FileStoreSourceSplit, Long>
splitWeightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?>
splitGroupFunc) {
switch (splitAssignMode) {
case FAIR:
- return new PreAssignSplitAssigner(splitBatchSize, context,
splits);
+ return new PreAssignSplitAssigner(
+ splitBatchSize, context, splits, splitWeightFunc,
splitGroupFunc);
case PREEMPTIVE:
return new FIFOSplitAssigner(splits);
default:
diff --git
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
index fbb31bd108..24ac4a2911 100644
---
a/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
+++
b/paimon-flink/paimon-flink-common/src/main/java/org/apache/paimon/flink/source/assigners/PreAssignSplitAssigner.java
@@ -24,6 +24,7 @@ import org.apache.paimon.flink.FlinkRowData;
import org.apache.paimon.flink.source.FileStoreSourceSplit;
import org.apache.paimon.table.source.DataSplit;
import org.apache.paimon.utils.BinPacking;
+import org.apache.paimon.utils.SerializableFunction;
import org.apache.flink.api.connector.source.SplitEnumeratorContext;
import org.apache.flink.table.connector.source.DynamicFilteringData;
@@ -37,6 +38,7 @@ import java.util.LinkedList;
import java.util.List;
import java.util.ListIterator;
import java.util.Map;
+import java.util.Objects;
import java.util.Optional;
import java.util.Queue;
import java.util.concurrent.atomic.AtomicInteger;
@@ -59,6 +61,8 @@ public class PreAssignSplitAssigner implements SplitAssigner {
private final AtomicInteger numberOfPendingSplits;
private final Collection<FileStoreSourceSplit> splits;
+ private final SerializableFunction<FileStoreSourceSplit, Long> weightFunc;
+ @Nullable private final SerializableFunction<FileStoreSourceSplit, ?>
groupFunc;
public PreAssignSplitAssigner(
int splitBatchSize,
@@ -67,26 +71,80 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
this(splitBatchSize, context.currentParallelism(), splits);
}
+ public PreAssignSplitAssigner(
+ int splitBatchSize,
+ SplitEnumeratorContext<FileStoreSourceSplit> context,
+ Collection<FileStoreSourceSplit> splits,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc) {
+ this(splitBatchSize, context.currentParallelism(), splits, weightFunc);
+ }
+
+ public PreAssignSplitAssigner(
+ int splitBatchSize,
+ SplitEnumeratorContext<FileStoreSourceSplit> context,
+ Collection<FileStoreSourceSplit> splits,
+ @Nullable SerializableFunction<FileStoreSourceSplit, Long>
weightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?> groupFunc)
{
+ this(splitBatchSize, context.currentParallelism(), splits, weightFunc,
groupFunc);
+ }
+
public PreAssignSplitAssigner(
int splitBatchSize,
int parallelism,
Collection<FileStoreSourceSplit> splits,
Projection partitionRowProjection,
DynamicFilteringData dynamicFilteringData) {
+ this(
+ splitBatchSize,
+ parallelism,
+ splits,
+ partitionRowProjection,
+ dynamicFilteringData,
+ split -> split.split().rowCount());
+ }
+
+ public PreAssignSplitAssigner(
+ int splitBatchSize,
+ int parallelism,
+ Collection<FileStoreSourceSplit> splits,
+ Projection partitionRowProjection,
+ DynamicFilteringData dynamicFilteringData,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc) {
this(
splitBatchSize,
parallelism,
splits.stream()
.filter(s -> filter(partitionRowProjection,
dynamicFilteringData, s))
- .collect(Collectors.toList()));
+ .collect(Collectors.toList()),
+ weightFunc);
}
public PreAssignSplitAssigner(
int splitBatchSize, int parallelism,
Collection<FileStoreSourceSplit> splits) {
+ this(splitBatchSize, parallelism, splits, split ->
split.split().rowCount());
+ }
+
+ public PreAssignSplitAssigner(
+ int splitBatchSize,
+ int parallelism,
+ Collection<FileStoreSourceSplit> splits,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc) {
+ this(splitBatchSize, parallelism, splits, weightFunc, null);
+ }
+
+ public PreAssignSplitAssigner(
+ int splitBatchSize,
+ int parallelism,
+ Collection<FileStoreSourceSplit> splits,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?> groupFunc)
{
this.splitBatchSize = splitBatchSize;
this.parallelism = parallelism;
this.splits = splits;
- this.pendingSplitAssignment = createBatchFairSplitAssignment(splits,
parallelism);
+ this.weightFunc = weightFunc == null ? split ->
split.split().rowCount() : weightFunc;
+ this.groupFunc = groupFunc;
+ this.pendingSplitAssignment =
+ createBatchFairSplitAssignment(splits, parallelism,
this.weightFunc, groupFunc);
this.numberOfPendingSplits = new AtomicInteger(splits.size());
}
@@ -132,11 +190,37 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
* this method only reload restore for batch execute, because in streaming
mode, we need to
* assign certain bucket to certain task.
*/
- private static Map<Integer, LinkedList<FileStoreSourceSplit>>
createBatchFairSplitAssignment(
- Collection<FileStoreSourceSplit> splits, int numReaders) {
- List<List<FileStoreSourceSplit>> assignmentList =
- BinPacking.packForFixedBinNumber(
- splits, split -> split.split().rowCount(), numReaders);
+ public static Map<Integer, LinkedList<FileStoreSourceSplit>>
createBatchFairSplitAssignment(
+ Collection<FileStoreSourceSplit> splits,
+ int numReaders,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc) {
+ return createBatchFairSplitAssignment(splits, numReaders, weightFunc,
null);
+ }
+
+ public static Map<Integer, LinkedList<FileStoreSourceSplit>>
createBatchFairSplitAssignment(
+ Collection<FileStoreSourceSplit> splits,
+ int numReaders,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc,
+ @Nullable SerializableFunction<FileStoreSourceSplit, ?> groupFunc)
{
+ List<List<FileStoreSourceSplit>> assignmentList;
+ if (groupFunc == null) {
+ assignmentList = BinPacking.packForFixedBinNumber(splits,
weightFunc, numReaders);
+ } else {
+ List<GroupedSplit> groupedSplits = groupSplits(splits, weightFunc,
groupFunc);
+ List<List<GroupedSplit>> groupedAssignment =
+ BinPacking.packForFixedBinNumber(
+ groupedSplits, GroupedSplit::weight, numReaders);
+ assignmentList =
+ groupedAssignment.stream()
+ .map(
+ group ->
+ group.stream()
+ .flatMap(
+ groupedSplit ->
+
groupedSplit.splits().stream())
+
.collect(Collectors.toList()))
+ .collect(Collectors.toList());
+ }
Map<Integer, LinkedList<FileStoreSourceSplit>> assignment = new
HashMap<>();
for (int i = 0; i < assignmentList.size(); i++) {
assignment.put(i, new LinkedList<>(assignmentList.get(i)));
@@ -144,6 +228,43 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
return assignment;
}
+ private static List<GroupedSplit> groupSplits(
+ Collection<FileStoreSourceSplit> splits,
+ SerializableFunction<FileStoreSourceSplit, Long> weightFunc,
+ SerializableFunction<FileStoreSourceSplit, ?> groupFunc) {
+ Map<Object, List<FileStoreSourceSplit>> grouped = new HashMap<>();
+ for (FileStoreSourceSplit split : splits) {
+ grouped.computeIfAbsent(
+ Objects.requireNonNull(groupFunc.apply(split)),
+ ignored -> new ArrayList<>())
+ .add(split);
+ }
+ return grouped.values().stream()
+ .map(
+ group ->
+ new GroupedSplit(
+ group,
group.stream().mapToLong(weightFunc::apply).sum()))
+ .collect(Collectors.toList());
+ }
+
+ private static class GroupedSplit {
+ private final List<FileStoreSourceSplit> splits;
+ private final long weight;
+
+ private GroupedSplit(List<FileStoreSourceSplit> splits, long weight) {
+ this.splits = splits;
+ this.weight = weight;
+ }
+
+ private List<FileStoreSourceSplit> splits() {
+ return splits;
+ }
+
+ private long weight() {
+ return weight;
+ }
+ }
+
@Override
public Optional<Long> getNextSnapshotId(int subtask) {
LinkedList<FileStoreSourceSplit> pendingSplits =
pendingSplitAssignment.get(subtask);
@@ -160,7 +281,12 @@ public class PreAssignSplitAssigner implements
SplitAssigner {
public SplitAssigner ofDynamicPartitionPruning(
Projection partitionRowProjection, DynamicFilteringData
dynamicFilteringData) {
return new PreAssignSplitAssigner(
- splitBatchSize, parallelism, splits, partitionRowProjection,
dynamicFilteringData);
+ splitBatchSize,
+ parallelism,
+ splits,
+ partitionRowProjection,
+ dynamicFilteringData,
+ weightFunc);
}
private static boolean filter(
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/action/CompactionBucketDistributionStrategyTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/action/CompactionBucketDistributionStrategyTest.java
new file mode 100644
index 0000000000..3ea093f345
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/action/CompactionBucketDistributionStrategyTest.java
@@ -0,0 +1,70 @@
+/*
+ * 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.flink.action;
+
+import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
+import org.apache.paimon.options.Options;
+
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for compaction bucket distribution strategy resolution. */
+public class CompactionBucketDistributionStrategyTest {
+
+ @Test
+ public void testDefaultStrategyIsLinear() {
+ assertThat(
+ Options.fromMap(java.util.Collections.emptyMap())
+
.get(FlinkConnectorOptions.COMPACTION_BUCKET_DISTRIBUTION_STRATEGY))
+ .isEqualTo(CompactionBucketDistributionStrategy.LINEAR);
+ }
+
+ @Test
+ public void testStrategyOptionParsing() {
+
assertThat(strategy("linear")).isEqualTo(CompactionBucketDistributionStrategy.LINEAR);
+ assertThat(strategy("size-aware-batch"))
+
.isEqualTo(CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH);
+ }
+
+ @Test
+ public void testStrategyOnlyAppliesToBatchFullCompaction() {
+ Options options =
+ Options.fromMap(
+ java.util.Collections.singletonMap(
+
FlinkConnectorOptions.COMPACTION_BUCKET_DISTRIBUTION_STRATEGY.key(),
+ "size-aware-batch"));
+
+ assertThat(CompactAction.compactionBucketDistributionStrategy(options,
true, false))
+
.isEqualTo(CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH);
+ assertThat(CompactAction.compactionBucketDistributionStrategy(options,
false, false))
+ .isEqualTo(CompactionBucketDistributionStrategy.LINEAR);
+ assertThat(CompactAction.compactionBucketDistributionStrategy(options,
true, true))
+ .isEqualTo(CompactionBucketDistributionStrategy.LINEAR);
+ }
+
+ private static CompactionBucketDistributionStrategy strategy(String value)
{
+ return Options.fromMap(
+ java.util.Collections.singletonMap(
+
FlinkConnectorOptions.COMPACTION_BUCKET_DISTRIBUTION_STRATEGY.key(),
+ value))
+
.get(FlinkConnectorOptions.COMPACTION_BUCKET_DISTRIBUTION_STRATEGY);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceBuilderTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceBuilderTest.java
new file mode 100644
index 0000000000..4746db671b
--- /dev/null
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceBuilderTest.java
@@ -0,0 +1,157 @@
+/*
+ * 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.flink.source;
+
+import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.flink.FlinkConnectorOptions;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
+import org.apache.paimon.flink.FlinkConnectorOptions.SplitAssignMode;
+import org.apache.paimon.io.DataFileMeta;
+import org.apache.paimon.options.Options;
+import org.apache.paimon.table.source.DataSplit;
+
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Tests for {@link CompactorSourceBuilder}. */
+public class CompactorSourceBuilderTest {
+
+ @Test
+ public void testSizeAwareBatchUsesSinkParallelismForSource() {
+ Options options = optionsWithParallelism(4, 16);
+
+ assertThat(
+ CompactorSourceBuilder.sourceParallelism(
+ options,
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH))
+ .isEqualTo(16);
+ }
+
+ @Test
+ public void testNonSizeAwareBatchUsesScanParallelismForSource() {
+ Options options = optionsWithParallelism(4, 16);
+
+ assertThat(
+ CompactorSourceBuilder.sourceParallelism(
+ options,
CompactionBucketDistributionStrategy.LINEAR))
+ .isEqualTo(4);
+ }
+
+ @Test
+ public void
testSizeAwareBatchFallsBackToScanParallelismWithoutSinkParallelism() {
+ Map<String, String> map = new HashMap<>();
+ map.put(FlinkConnectorOptions.SCAN_PARALLELISM.key(), "4");
+ Options options = Options.fromMap(map);
+
+ assertThat(
+ CompactorSourceBuilder.sourceParallelism(
+ options,
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH))
+ .isEqualTo(4);
+ }
+
+ @Test
+ public void testSizeAwareBatchRejectsPreemptiveSplitAssignment() {
+ assertThatThrownBy(
+ () ->
+
CompactorSourceBuilder.validateBucketDistributionStrategy(
+
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH,
+ SplitAssignMode.PREEMPTIVE))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("size-aware-batch requires")
+ .hasMessageContaining("scan.split-enumerator.mode=fair");
+ }
+
+ @Test
+ public void testSizeAwareBatchAllowsFairSplitAssignment() {
+ assertThatCode(
+ () ->
+
CompactorSourceBuilder.validateBucketDistributionStrategy(
+
CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH,
+ SplitAssignMode.FAIR))
+ .doesNotThrowAnyException();
+ }
+
+ @Test
+ public void testBucketKeyUsesPartitionAndBucket() {
+ DataSplit split1 = dataSplit(1L, BinaryRow.EMPTY_ROW, 0,
dataFile("file-1", 10L));
+ DataSplit split2 = dataSplit(2L, BinaryRow.EMPTY_ROW, 0,
dataFile("file-2", 25L));
+ DataSplit split3 = dataSplit(3L, BinaryRow.EMPTY_ROW, 1,
dataFile("file-3", 30L));
+
+ assertThat(CompactorSourceBuilder.bucketKey(split1))
+ .isEqualTo(CompactorSourceBuilder.bucketKey(split2));
+ assertThat(CompactorSourceBuilder.bucketKey(split1))
+ .isNotEqualTo(CompactorSourceBuilder.bucketKey(split3));
+ }
+
+ @Test
+ public void testBucketFileSizeUsesTotalDataFileSize() {
+ DataSplit split =
+ dataSplit(
+ 1L,
+ BinaryRow.EMPTY_ROW,
+ 0,
+ dataFile("file-1", 10L),
+ dataFile("file-2", 25L));
+
+
assertThat(CompactorSourceBuilder.bucketFileSize(split)).isEqualTo(35L);
+ }
+
+ private static Options optionsWithParallelism(int scanParallelism, int
sinkParallelism) {
+ Map<String, String> map = new HashMap<>();
+ map.put(FlinkConnectorOptions.SCAN_PARALLELISM.key(),
String.valueOf(scanParallelism));
+ map.put(FlinkConnectorOptions.SINK_PARALLELISM.key(),
String.valueOf(sinkParallelism));
+ return Options.fromMap(map);
+ }
+
+ private static DataSplit dataSplit(
+ long snapshot, BinaryRow partition, int bucket, DataFileMeta...
dataFiles) {
+ return DataSplit.builder()
+ .withSnapshot(snapshot)
+ .withPartition(partition)
+ .withBucket(bucket)
+ .withBucketPath("bucket-" + bucket)
+ .withDataFiles(Arrays.asList(dataFiles))
+ .build();
+ }
+
+ private static DataFileMeta dataFile(String fileName, long fileSize) {
+ return DataFileMeta.forAppend(
+ fileName,
+ fileSize,
+ 1L,
+ null,
+ 0L,
+ 0L,
+ 0L,
+ Collections.emptyList(),
+ null,
+ null,
+ null,
+ null,
+ null,
+ null);
+ }
+}
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceITCase.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceITCase.java
index d5224670d4..8f7686dbff 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceITCase.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/CompactorSourceITCase.java
@@ -23,6 +23,7 @@ import org.apache.paimon.data.BinaryRow;
import org.apache.paimon.data.BinaryRowWriter;
import org.apache.paimon.data.BinaryString;
import org.apache.paimon.data.GenericRow;
+import
org.apache.paimon.flink.FlinkConnectorOptions.CompactionBucketDistributionStrategy;
import org.apache.paimon.flink.util.AbstractTestBase;
import org.apache.paimon.fs.Path;
import org.apache.paimon.fs.local.LocalFileIO;
@@ -42,6 +43,10 @@ import org.apache.paimon.types.RowType;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.streaming.api.graph.StreamEdge;
+import org.apache.flink.streaming.api.graph.StreamGraph;
+import org.apache.flink.streaming.api.graph.StreamNode;
+import org.apache.flink.streaming.runtime.partitioner.ForwardPartitioner;
import org.apache.flink.table.data.RowData;
import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.BeforeEach;
@@ -286,6 +291,54 @@ public class CompactorSourceITCase extends
AbstractTestBase {
it.close();
}
+ @Test
+ public void
testSizeAwareSourceAndFilterKeepSameParallelismWithoutShuffle() throws
Exception {
+ FileStoreTable table =
+ createFileStoreTable()
+ .copy(
+ new HashMap<String, String>() {
+ {
+ put("sink.parallelism", "4");
+ put("scan.parallelism", "2");
+ put(
+
"compaction.bucket-distribution-strategy",
+ "size-aware-batch");
+ }
+ });
+ StreamWriteBuilder streamWriteBuilder =
+ table.newStreamWriteBuilder().withCommitUser(commitUser);
+ StreamTableWrite write = streamWriteBuilder.newWrite();
+ StreamTableCommit commit = streamWriteBuilder.newCommit();
+ write.write(rowData(1, 1510, BinaryString.fromString("20221208"), 15));
+ commit.commit(0, write.prepareCommit(true, 0));
+
+ StreamExecutionEnvironment env =
+ streamExecutionEnvironmentBuilder().streamingMode().build();
+ new CompactorSourceBuilder("test", table)
+ .withContinuousMode(false)
+ .withBucketDistributionStrategy(
+ CompactionBucketDistributionStrategy.SIZE_AWARE_BATCH)
+ .withPartitionIdleTime(Duration.ofMillis(1))
+ .withEnv(env)
+ .build();
+
+ StreamGraph streamGraph = env.getStreamGraph(false);
+ StreamNode sourceNode = findStreamNode(streamGraph,
"test-compact-source");
+ StreamNode filterNode = findStreamNode(streamGraph, "Filter");
+ assertThat(sourceNode.getParallelism()).isEqualTo(4);
+ assertThat(filterNode.getParallelism()).isEqualTo(4);
+
+ StreamEdge sourceToFilterEdge =
+ sourceNode.getOutEdges().stream()
+ .filter(edge -> edge.getTargetId() ==
filterNode.getId())
+ .findFirst()
+ .orElseThrow(() -> new AssertionError("Source to
filter edge not found"));
+
assertThat(sourceToFilterEdge.getPartitioner()).isInstanceOf(ForwardPartitioner.class);
+
+ write.close();
+ commit.close();
+ }
+
@ParameterizedTest(name = "defaultOptions = {0}")
@ValueSource(booleans = {true, false})
public void testHistoryPartitionRead(boolean defaultOptions) throws
Exception {
@@ -334,6 +387,14 @@ public class CompactorSourceITCase extends
AbstractTestBase {
it.close();
}
+ private StreamNode findStreamNode(StreamGraph streamGraph, String
operatorNameContains) {
+ return streamGraph.getStreamNodes().stream()
+ .filter(node ->
node.getOperatorName().contains(operatorNameContains))
+ .findFirst()
+ .orElseThrow(
+ () -> new AssertionError("Stream node not found: " +
operatorNameContains));
+ }
+
private String toString(RowData rowData) {
int numFiles;
try {
diff --git
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FairAssignModeTest.java
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FairAssignModeTest.java
index 8a7b36b27d..1f4fcf465d 100644
---
a/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FairAssignModeTest.java
+++
b/paimon-flink/paimon-flink-common/src/test/java/org/apache/paimon/flink/source/FairAssignModeTest.java
@@ -18,6 +18,8 @@
package org.apache.paimon.flink.source;
+import org.apache.paimon.table.source.DataSplit;
+
import
org.apache.flink.connector.testutils.source.reader.TestingSplitEnumeratorContext;
import org.junit.jupiter.api.Test;
@@ -64,6 +66,90 @@ public class FairAssignModeTest extends
StaticFileStoreSplitEnumeratorTestBase {
.containsExactly(splits.get(0), splits.get(2));
}
+ @Test
+ public void testSplitAllocationWithCustomWeight() {
+ final TestingSplitEnumeratorContext<FileStoreSourceSplit> context =
+ getSplitEnumeratorContext(2);
+
+ List<FileStoreSourceSplit> splits = new ArrayList<>();
+ for (int i = 1; i <= 4; i++) {
+ splits.add(createSnapshotSplit(i, 0, Collections.emptyList()));
+ }
+ StaticFileStoreSplitEnumerator enumerator =
+ new StaticFileStoreSplitEnumerator(
+ context,
+ null,
+ StaticFileStoreSource.createSplitAssigner(
+ context,
+ 10,
+ SplitAssignMode.FAIR,
+ splits,
+ split -> {
+ int snapshotId = (int) ((DataSplit)
split.split()).snapshotId();
+ return snapshotId == 1 || snapshotId == 2
? 100L : 1L;
+ }));
+
+ enumerator.handleSplitRequest(0, "test-host");
+ enumerator.handleSplitRequest(1, "test-host");
+ Map<Integer, SplitAssignmentState<FileStoreSourceSplit>> assignments =
+ context.getSplitAssignments();
+
+ assertThat(assignments).containsOnlyKeys(0, 1);
+
assertThat(totalWeight(assignments.get(0).getAssignedSplits())).isEqualTo(101L);
+
assertThat(totalWeight(assignments.get(1).getAssignedSplits())).isEqualTo(101L);
+ assertThat(assignments.get(0).getAssignedSplits()).hasSize(2);
+ assertThat(assignments.get(1).getAssignedSplits()).hasSize(2);
+ }
+
+ @Test
+ public void testGroupedSplitAllocationBalancesByBucketTotalWeight() {
+ final TestingSplitEnumeratorContext<FileStoreSourceSplit> context =
+ getSplitEnumeratorContext(2);
+
+ List<FileStoreSourceSplit> splits = new ArrayList<>();
+ // Bucket 0 has two splits, each with weight 50. Bucket 1 has one
split with weight 100.
+ // The assigner should first group by bucket, calculate bucket-level
total weight, and then
+ // bin-pack the two bucket groups evenly across readers. This also
guarantees all splits of
+ // the same bucket go to the same downstream writer.
+ splits.add(createSnapshotSplit(1, 0, Collections.emptyList()));
+ splits.add(createSnapshotSplit(2, 0, Collections.emptyList()));
+ splits.add(createSnapshotSplit(3, 1, Collections.emptyList()));
+
+ StaticFileStoreSplitEnumerator enumerator =
+ new StaticFileStoreSplitEnumerator(
+ context,
+ null,
+ StaticFileStoreSource.createSplitAssigner(
+ context,
+ 10,
+ SplitAssignMode.FAIR,
+ splits,
+ split -> ((DataSplit) split.split()).bucket()
== 0 ? 50L : 100L,
+ split -> ((DataSplit)
split.split()).bucket()));
+
+ enumerator.handleSplitRequest(0, "test-host");
+ enumerator.handleSplitRequest(1, "test-host");
+ Map<Integer, SplitAssignmentState<FileStoreSourceSplit>> assignments =
+ context.getSplitAssignments();
+
+ assertThat(assignments).containsOnlyKeys(0, 1);
+ assertThat(assignments.values())
+ .anySatisfy(
+ assignment ->
+ assertThat(assignment.getAssignedSplits())
+
.containsExactlyInAnyOrder(splits.get(0), splits.get(1)));
+ assertThat(assignments.values())
+ .anySatisfy(
+ assignment ->
+ assertThat(assignment.getAssignedSplits())
+ .containsExactly(splits.get(2)));
+ assertThat(assignments.values())
+ .allSatisfy(
+ assignment ->
+
assertThat(groupedTotalWeight(assignment.getAssignedSplits()))
+ .isEqualTo(100L));
+ }
+
@Test
public void testSplitBatch() {
final TestingSplitEnumeratorContext<FileStoreSourceSplit> context =
@@ -146,6 +232,22 @@ public class FairAssignModeTest extends
StaticFileStoreSplitEnumeratorTestBase {
assertThat(assignments.get(2).getAssignedSplits()).isEmpty();
}
+ private long totalWeight(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .mapToLong(
+ split -> {
+ int snapshotId = (int) ((DataSplit)
split.split()).snapshotId();
+ return snapshotId == 1 || snapshotId == 2 ? 100L :
1L;
+ })
+ .sum();
+ }
+
+ private long groupedTotalWeight(List<FileStoreSourceSplit> splits) {
+ return splits.stream()
+ .mapToLong(split -> ((DataSplit) split.split()).bucket() == 0
? 50L : 100L)
+ .sum();
+ }
+
@Override
protected SplitAssignMode splitAssignMode() {
return SplitAssignMode.FAIR;
