Skip to content

Flink: Fix HashKeyGenerator SelectorKey cache ignoring writeParallelism and distributionMode #15740

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pvary merged 1 commit into from
Mar 25, 2026
Merged

Flink: Fix HashKeyGenerator SelectorKey cache ignoring writeParallelism and distributionMode #15740

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 25 additions & 4 deletions flink/v2.1/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,9 @@ int generateKey(
tableSpec != null ? tableSpec.specId() : null,
dynamicRecord.schema(),
dynamicRecord.spec(),
dynamicRecord.equalityFields());
dynamicRecord.equalityFields(),
MoreObjects.firstNonNull(dynamicRecord.distributionMode(), DistributionMode.NONE),
Math.min(dynamicRecord.writeParallelism(), maxWriteParallelism));
Comment on lines +91 to +92

@mxm mxm Mar 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see you added the same logic as for creating the actual KeySelector below. Should we consolidate the SelectorKey and the getKeySelector (below) parameters?

@mxm mxm Mar 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can do that in a follow-up.

KeySelector<RowData, Integer> keySelector =
keySelectorCache.computeIfAbsent(
cacheKey,
Expand Down Expand Up @@ -322,6 +324,8 @@ static class SelectorKey {
private final Schema schema;
private final PartitionSpec spec;
private final Set<String> equalityFields;
private final DistributionMode distributionMode;
private final int writeParallelism;

SelectorKey(
String tableName,
Expand All @@ -330,14 +334,18 @@ static class SelectorKey {
@Nullable Integer tableSpecId,
Schema schema,
PartitionSpec spec,
Set<String> equalityFields) {
Set<String> equalityFields,
DistributionMode distributionMode,
int writeParallelism) {
this.tableName = tableName;
this.branch = branch;
this.schemaId = tableSchemaId;
this.specId = tableSpecId;
this.schema = tableSchemaId == null ? schema : null;
this.spec = tableSpecId == null ? spec : null;
this.equalityFields = equalityFields;
this.distributionMode = distributionMode;
this.writeParallelism = writeParallelism;
}

@Override
Expand All @@ -357,12 +365,23 @@ public boolean equals(Object other) {
&& Objects.equals(specId, that.specId)
&& Objects.equals(schema, that.schema)
&& Objects.equals(spec, that.spec)
&& Objects.equals(equalityFields, that.equalityFields);
&& Objects.equals(equalityFields, that.equalityFields)
&& distributionMode == that.distributionMode
&& writeParallelism == that.writeParallelism;
}

@Override
public int hashCode() {
return Objects.hash(tableName, branch, schemaId, specId, schema, spec, equalityFields);
return Objects.hash(
tableName,
branch,
schemaId,
specId,
schema,
spec,
equalityFields,
distributionMode,
writeParallelism);
}

@Override
Expand All @@ -375,6 +394,8 @@ public String toString() {
.add("schema", schema)
.add("spec", spec)
.add("equalityFields", equalityFields)
.add("distributionMode", distributionMode)
.add("writeParallelism", writeParallelism)
.toString();
}
}
Expand Down
68 changes: 68 additions & 0 deletions .../v2.1/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,74 @@ void testCaching() throws Exception {
assertThat(writeKey1).isEqualTo(writeKey3);
}

@Test
void testCacheMissOnWriteParallelismChange() throws Exception {
int maxWriteParallelism = 8;
HashKeyGenerator generator = new HashKeyGenerator(10, maxWriteParallelism);
Map<HashKeyGenerator.SelectorKey, KeySelector<RowData, Integer>> keySelectorCache =
generator.getKeySelectorCache();

PartitionSpec unpartitioned = PartitionSpec.unpartitioned();
DynamicRecord record1 =
new DynamicRecord(
TABLE_IDENTIFIER,
BRANCH,
SCHEMA,
GenericRowData.of(1, StringData.fromString("foo")),
unpartitioned,
DistributionMode.NONE,
2);
DynamicRecord record2 =
new DynamicRecord(
TABLE_IDENTIFIER,
BRANCH,
SCHEMA,
GenericRowData.of(1, StringData.fromString("foo")),
unpartitioned,
DistributionMode.NONE,
4);

generator.generateKey(record1);
assertThat(keySelectorCache).hasSize(1);

generator.generateKey(record2);
assertThat(keySelectorCache).hasSize(2);
}

@Test
void testCacheMissOnDistributionModeChange() throws Exception {
int maxWriteParallelism = 8;
HashKeyGenerator generator = new HashKeyGenerator(10, maxWriteParallelism);
Map<HashKeyGenerator.SelectorKey, KeySelector<RowData, Integer>> keySelectorCache =
generator.getKeySelectorCache();

PartitionSpec partitioned = PartitionSpec.builderFor(SCHEMA).identity("id").build();
DynamicRecord record1 =
new DynamicRecord(
TABLE_IDENTIFIER,
BRANCH,
SCHEMA,
GenericRowData.of(1, StringData.fromString("foo")),
partitioned,
DistributionMode.NONE,
2);
DynamicRecord record2 =
new DynamicRecord(
TABLE_IDENTIFIER,
BRANCH,
SCHEMA,
GenericRowData.of(1, StringData.fromString("foo")),
partitioned,
DistributionMode.HASH,
2);

generator.generateKey(record1);
assertThat(keySelectorCache).hasSize(1);

generator.generateKey(record2);
assertThat(keySelectorCache).hasSize(2);
}

private static int getWriteKey(
HashKeyGenerator generator,
PartitionSpec spec,
Expand Down
Loading