This is an automated email from the ASF dual-hosted git repository.
mxm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new c8d68218df Flink: Distribute rows by custom fields to DynamicWriter in
DynamicSink (#18668)
c8d68218df is described below
commit c8d68218df68f96000e39387e6491736b23a25c2
Author: Swapna Marru <[email protected]>
AuthorDate: Mon Nov 9 07:49:27 2026 -0800
Flink: Distribute rows by custom fields to DynamicWriter in DynamicSink
(#18668)
---
docs/docs/flink-writes.md | 3 +-
.../iceberg/flink/sink/dynamic/DynamicRecord.java | 9 +
.../sink/dynamic/DynamicRecordWithConfig.java | 5 +
.../flink/sink/dynamic/DynamicSinkUtil.java | 15 ++
.../flink/sink/dynamic/HashKeyGenerator.java | 83 +++++--
.../flink/sink/dynamic/TestDynamicIcebergSink.java | 72 +++++-
.../flink/sink/dynamic/TestHashKeyGenerator.java | 261 +++++++++++++++++++++
7 files changed, 427 insertions(+), 21 deletions(-)
diff --git a/docs/docs/flink-writes.md b/docs/docs/flink-writes.md
index deece55bb8..e9b7c4dcc5 100644
--- a/docs/docs/flink-writes.md
+++ b/docs/docs/flink-writes.md
@@ -516,6 +516,7 @@ We need the following information (DynamicRecord) for every
record:
| `Parallelism` | The maximum number of parallel writers for a given
table/branch/schema/spec (WriteTarget). |
| `UpsertMode` | Overrides this table's write.upsert.enabled (optional).
|
| `EqualityFields` | The equality fields for the table (optional). When
unset, the dynamic sink falls back to the schema's identifier fields for both
distribution (records sharing a key route to the same writer) and
equality-delete inference. Identifier fields do not auto-enable upsert mode —
that flag remains opt-in. |
+| `DistributionFields` | The fields to hash records by when the distribution
mode is `HASH` (optional). When set, they replace the partition key and
equality fields as the hash key. If equality fields are set, the distribution
fields must be a subset of them. The same applies to the schema's identifier
fields if equality fields are unset. Ignored for other distribution modes. |
### Schema Evolution
@@ -589,7 +590,7 @@ The `DistributionMode` set on each `DynamicRecord` controls
how that record is r
| Mode | Behavior
|
|--------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| `NONE` | Records are distributed across writer subtasks in a round-robin
fashion (or by equality fields if set).
|
-| `HASH` | Records are distributed by partition key (partitioned tables) or
equality fields (unpartitioned tables). Ensures that records for the same
partition are handled by the same writer subtask. |
+| `HASH` | Records are distributed by the distribution fields if set,
otherwise by partition key (partitioned tables) or equality fields
(unpartitioned tables). Without distribution fields, records for the same
partition are handled by the same writer subtask. |
| `null` | Forward mode: bypasses distribution entirely and sends records
directly via a forward edge (see below).
|
#### Forward Mode
diff --git
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecord.java
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecord.java
index 6507a575c2..60bb3a272a 100644
---
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecord.java
+++
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecord.java
@@ -39,6 +39,7 @@ public class DynamicRecord {
private int writeParallelism;
private boolean upsertMode;
@Nullable private Set<String> equalityFields;
+ @Nullable private Set<String> distributionFields;
@Internal
DynamicRecord() {}
@@ -165,4 +166,12 @@ public class DynamicRecord {
public void setEqualityFields(Set<String> equalityFields) {
this.equalityFields = equalityFields;
}
+
+ public Set<String> distributionFields() {
+ return distributionFields;
+ }
+
+ public void setDistributionFields(Set<String> distributionFields) {
+ this.distributionFields = distributionFields;
+ }
}
diff --git
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecordWithConfig.java
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecordWithConfig.java
index 32716c3e4a..569bcf5c9a 100644
---
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecordWithConfig.java
+++
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicRecordWithConfig.java
@@ -91,4 +91,9 @@ class DynamicRecordWithConfig extends DynamicRecord {
public Set<String> equalityFields() {
return wrapped.equalityFields();
}
+
+ @Override
+ public Set<String> distributionFields() {
+ return wrapped.distributionFields();
+ }
}
diff --git
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicSinkUtil.java
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicSinkUtil.java
index dea135ddca..e278e13612 100644
---
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicSinkUtil.java
+++
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/DynamicSinkUtil.java
@@ -71,6 +71,21 @@ class DynamicSinkUtil {
return schema.identifierFieldNames();
}
+ static Set<Integer> distributionFieldIds(Set<String> distributionFields,
Schema schema) {
+ Set<Integer> distributionFieldIds =
Sets.newHashSetWithExpectedSize(distributionFields.size());
+ for (String distributionField : distributionFields) {
+ Types.NestedField field = schema.findField(distributionField);
+ Preconditions.checkNotNull(
+ field,
+ "Missing distribution field '%s'. Expected distribution field '%s'
to exist in the table schema.",
+ distributionField,
+ distributionField);
+ distributionFieldIds.add(field.fieldId());
+ }
+
+ return distributionFieldIds;
+ }
+
static int safeAbs(int input) {
if (input >= 0) {
return input;
diff --git
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java
index c60dd9b0ce..9cfc66f1fd 100644
---
a/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java
+++
b/flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java
@@ -20,6 +20,7 @@ package org.apache.iceberg.flink.sink.dynamic;
import static org.apache.iceberg.TableProperties.WRITE_DISTRIBUTION_MODE;
+import java.util.Collections;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
@@ -82,6 +83,8 @@ class HashKeyGenerator {
DynamicSinkUtil.resolveEqualityFieldNames(dynamicRecord.equalityFields(),
effectiveSchema);
DistributionMode effectiveDistributionMode =
MoreObjects.firstNonNull(dynamicRecord.distributionMode(),
DistributionMode.NONE);
+ Set<String> effectiveDistributionFields =
+ MoreObjects.firstNonNull(dynamicRecord.distributionFields(),
Collections.emptySet());
SelectorKey cacheKey =
new SelectorKey(
tableIdent,
@@ -92,7 +95,8 @@ class HashKeyGenerator {
dynamicRecord.spec(),
effectiveEqualityFields,
effectiveDistributionMode,
- Math.min(dynamicRecord.writeParallelism(), maxWriteParallelism));
+ Math.min(dynamicRecord.writeParallelism(), maxWriteParallelism),
+ effectiveDistributionFields);
KeySelector<RowData, Integer> keySelector =
keySelectorCache.computeIfAbsent(
cacheKey,
@@ -103,7 +107,8 @@ class HashKeyGenerator {
MoreObjects.firstNonNull(tableSpec, dynamicRecord.spec()),
effectiveDistributionMode,
effectiveEqualityFields,
- Math.min(dynamicRecord.writeParallelism(),
maxWriteParallelism)));
+ Math.min(dynamicRecord.writeParallelism(),
maxWriteParallelism),
+ effectiveDistributionFields));
try {
return keySelector.getKey(
overrideRowData != null ? overrideRowData : dynamicRecord.rowData());
@@ -118,23 +123,51 @@ class HashKeyGenerator {
PartitionSpec spec,
DistributionMode mode,
Set<String> equalityFields,
- int writeParallelism) {
+ int writeParallelism,
+ Set<String> distributionFields) {
LOG.debug(
"Creating new KeySelector for table '{}' with distribution mode '{}'",
tableName, mode);
+
switch (mode) {
case NONE:
+ if (!distributionFields.isEmpty()) {
+ LOG.warn(
+ "{}: Distribution fields {} are ignored when distribution mode
is set to NONE",
+ tableName,
+ distributionFields);
+ }
+
if (equalityFields.isEmpty()) {
return tableKeySelector(tableName, writeParallelism,
maxWriteParallelism);
} else {
LOG.info(
"{}: Distribute rows by equality fields, because there are
equality fields set",
tableName);
- return equalityFieldKeySelector(
- tableName, schema, equalityFields, writeParallelism,
maxWriteParallelism);
+ return fieldKeySelector(
+ tableName,
+ schema,
+ DynamicSinkUtil.getEqualityFieldIds(equalityFields, schema),
+ writeParallelism,
+ maxWriteParallelism);
}
case HASH:
- if (equalityFields.isEmpty()) {
+ if (!distributionFields.isEmpty()) {
+ Preconditions.checkArgument(
+ equalityFields.isEmpty() ||
equalityFields.containsAll(distributionFields),
+ "%s: Invalid distribution fields '%s'. Expected distribution
fields to be a subset of "
+ + "equality fields '%s'.",
+ tableName,
+ distributionFields,
+ equalityFields);
+ LOG.info("{}: Distribute rows by distribution fields {}", tableName,
distributionFields);
+ return fieldKeySelector(
+ tableName,
+ schema,
+ DynamicSinkUtil.distributionFieldIds(distributionFields, schema),
+ writeParallelism,
+ maxWriteParallelism);
+ } else if (equalityFields.isEmpty()) {
if (spec.isUnpartitioned()) {
LOG.warn(
"{}: Fallback to use 'none' distribution mode, because there
are no equality fields set "
@@ -151,8 +184,12 @@ class HashKeyGenerator {
"{}: Distribute rows by equality fields, because there are
equality fields set "
+ "and table is unpartitioned",
tableName);
- return equalityFieldKeySelector(
- tableName, schema, equalityFields, writeParallelism,
maxWriteParallelism);
+ return fieldKeySelector(
+ tableName,
+ schema,
+ DynamicSinkUtil.getEqualityFieldIds(equalityFields, schema),
+ writeParallelism,
+ maxWriteParallelism);
} else {
// Compare by fully-qualified name: a partition source that is a
nested field has a
// simple name (e.g. "name") that differs from the dotted
equality-field name (e.g.
@@ -188,8 +225,12 @@ class HashKeyGenerator {
+ "and {}='range' is not supported yet in flink",
tableName,
WRITE_DISTRIBUTION_MODE);
- return equalityFieldKeySelector(
- tableName, schema, equalityFields, writeParallelism,
maxWriteParallelism);
+ return fieldKeySelector(
+ tableName,
+ schema,
+ DynamicSinkUtil.getEqualityFieldIds(equalityFields, schema),
+ writeParallelism,
+ maxWriteParallelism);
}
default:
@@ -198,17 +239,14 @@ class HashKeyGenerator {
}
}
- private static KeySelector<RowData, Integer> equalityFieldKeySelector(
+ private static KeySelector<RowData, Integer> fieldKeySelector(
String tableName,
Schema schema,
- Set<String> equalityFields,
+ Set<Integer> fieldIds,
int writeParallelism,
int maxWriteParallelism) {
return new TargetLimitedKeySelector(
- new EqualityFieldKeySelector(
- schema,
- FlinkSchemaUtil.convert(schema),
- DynamicSinkUtil.getEqualityFieldIds(equalityFields, schema)),
+ new EqualityFieldKeySelector(schema, FlinkSchemaUtil.convert(schema),
fieldIds),
tableName,
writeParallelism,
maxWriteParallelism);
@@ -331,6 +369,7 @@ class HashKeyGenerator {
private final Set<String> equalityFields;
private final DistributionMode distributionMode;
private final int writeParallelism;
+ private final Set<String> distributionFields;
SelectorKey(
String tableName,
@@ -341,7 +380,8 @@ class HashKeyGenerator {
PartitionSpec spec,
Set<String> equalityFields,
DistributionMode distributionMode,
- int writeParallelism) {
+ int writeParallelism,
+ Set<String> distributionFields) {
this.tableName = tableName;
this.branch = branch;
this.schemaId = tableSchemaId;
@@ -351,9 +391,11 @@ class HashKeyGenerator {
this.equalityFields = equalityFields;
this.distributionMode = distributionMode;
this.writeParallelism = writeParallelism;
+ this.distributionFields = distributionFields;
}
@Override
+ @SuppressWarnings("checkstyle:CyclomaticComplexity")
public boolean equals(Object other) {
if (this == other) {
return true;
@@ -372,7 +414,8 @@ class HashKeyGenerator {
&& Objects.equals(spec, that.spec)
&& Objects.equals(equalityFields, that.equalityFields)
&& distributionMode == that.distributionMode
- && writeParallelism == that.writeParallelism;
+ && writeParallelism == that.writeParallelism
+ && Objects.equals(distributionFields, that.distributionFields);
}
@Override
@@ -386,7 +429,8 @@ class HashKeyGenerator {
spec,
equalityFields,
distributionMode,
- writeParallelism);
+ writeParallelism,
+ distributionFields);
}
@Override
@@ -401,6 +445,7 @@ class HashKeyGenerator {
.add("equalityFields", equalityFields)
.add("distributionMode", distributionMode)
.add("writeParallelism", writeParallelism)
+ .add("distributionFields", distributionFields)
.toString();
}
}
diff --git
a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java
b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java
index 69730d9adb..1e7b1945ff 100644
---
a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java
+++
b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java
@@ -136,6 +136,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
PartitionSpec partitionSpec;
boolean upsertMode;
Set<String> equalityFields;
+ Set<String> distributionFields;
int writeParallelism;
private DynamicIcebergDataImpl(
@@ -149,6 +150,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
false,
Collections.emptySet(),
false,
+ Collections.emptySet(),
10);
}
@@ -167,6 +169,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
false,
Collections.emptySet(),
false,
+ Collections.emptySet(),
10);
}
@@ -187,6 +190,26 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
upsertMode,
equalityFields,
isDuplicate,
+ Collections.emptySet(),
+ 10);
+ }
+
+ private DynamicIcebergDataImpl(
+ Schema schemaProvided,
+ String tableName,
+ String branch,
+ PartitionSpec partitionSpec,
+ Set<String> distributionFields) {
+ this(
+ schemaProvided,
+ schemaProvided,
+ tableName,
+ branch,
+ partitionSpec,
+ false,
+ Collections.emptySet(),
+ false,
+ distributionFields,
10);
}
@@ -199,6 +222,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
boolean upsertMode,
Set<String> equalityFields,
boolean isDuplicate,
+ Set<String> distributionFields,
int writeParallelism) {
this.rowProvided = randomRow(schemaProvided, isDuplicate ? seed :
++seed);
this.rowExpected = isDuplicate ? null : rowProvided;
@@ -209,6 +233,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
this.partitionSpec = partitionSpec;
this.upsertMode = upsertMode;
this.equalityFields = equalityFields;
+ this.distributionFields = distributionFields;
this.writeParallelism = writeParallelism;
}
}
@@ -221,6 +246,11 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
String branch = row.branch;
Schema schema = row.schemaProvided;
PartitionSpec spec = row.partitionSpec;
+ // Distribution fields are only read in HASH mode
+ DistributionMode distributionMode =
+ spec.isPartitioned() || !row.distributionFields.isEmpty()
+ ? DistributionMode.HASH
+ : DistributionMode.NONE;
DynamicRecord dynamicRecord =
new DynamicRecord(
tableIdentifier,
@@ -228,10 +258,11 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
schema,
converter(schema).toInternal(row.rowProvided),
spec,
- spec.isPartitioned() ? DistributionMode.HASH :
DistributionMode.NONE,
+ distributionMode,
row.writeParallelism);
dynamicRecord.setUpsertMode(row.upsertMode);
dynamicRecord.setEqualityFields(row.equalityFields);
+ dynamicRecord.setDistributionFields(row.distributionFields);
out.collect(dynamicRecord);
}
}
@@ -1366,6 +1397,43 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
}
}
+ @Test
+ void testDistributionFieldsWriteSameKeyToSingleDataFile() throws Exception {
+ int parallelism = 2;
+ List<DynamicIcebergDataImpl> rows = Lists.newArrayList();
+ for (int i = 0; i < 50; i++) {
+ DynamicIcebergDataImpl row =
+ new DynamicIcebergDataImpl(
+ SimpleDataUtil.SCHEMA,
+ "t1",
+ SnapshotRef.MAIN_BRANCH,
+ PartitionSpec.unpartitioned(),
+ Collections.singleton("id"));
+
+ row.rowProvided = Row.of(1, "data-" + i);
+ row.rowExpected = row.rowProvided;
+ rows.add(row);
+ }
+
+ executeDynamicSink(rows, this.env, true, parallelism, null);
+ verifyResults(rows);
+
+ Table table =
+ CATALOG_EXTENSION
+ .catalogLoader()
+ .loadCatalog()
+ .loadTable(TableIdentifier.of(DATABASE, "t1"));
+
+ List<Snapshot> dataSnapshots =
+ Lists.newArrayList(table.snapshots()).stream()
+ .filter(snapshot ->
snapshot.addedDataFiles(table.io()).iterator().hasNext())
+ .collect(Collectors.toList());
+
+ assertThat(dataSnapshots)
+ .isNotEmpty()
+ .allSatisfy(snapshot ->
assertThat(snapshot.addedDataFiles(table.io())).hasSize(1));
+ }
+
@Test
void testMultiFormatVersion() throws Exception {
ImmutableMap<String, String> properties =
ImmutableMap.of(TableProperties.FORMAT_VERSION, "3");
@@ -1942,6 +2010,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
false,
Collections.emptySet(),
false,
+ Collections.emptySet(),
-1),
new DynamicIcebergDataImpl(
SimpleDataUtil.SCHEMA,
@@ -1952,6 +2021,7 @@ class TestDynamicIcebergSink extends
TestFlinkIcebergSinkBase {
false,
Collections.emptySet(),
false,
+ Collections.emptySet(),
0));
runTest(
diff --git
a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java
b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java
index 837de70d3d..de6fbddb82 100644
---
a/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java
+++
b/flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java
@@ -38,6 +38,7 @@ import org.apache.iceberg.catalog.TableIdentifier;
import org.apache.iceberg.flink.FlinkWriteConf;
import org.apache.iceberg.flink.FlinkWriteOptions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import org.apache.iceberg.types.Types;
@@ -800,6 +801,224 @@ class TestHashKeyGenerator {
assertThat(keySelectorCache).hasSize(2);
}
+ @Test
+ void testDistributionFieldsWithDistributionModeHash() throws Exception {
+ int writeParallelism = 2;
+ int maxWriteParallelism = 8;
+ HashKeyGenerator generator = new HashKeyGenerator(16, maxWriteParallelism);
+ PartitionSpec unpartitioned = PartitionSpec.unpartitioned();
+ Set<String> distributionFields = Collections.singleton("id");
+
+ GenericRowData row1 = GenericRowData.of(1, StringData.fromString("foo"));
+ GenericRowData row2 = GenericRowData.of(1, StringData.fromString("bar"));
+ GenericRowData row3 = GenericRowData.of(2, StringData.fromString("baz"));
+
+ int writeKey1 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ writeParallelism,
+ Collections.emptySet(),
+ distributionFields,
+ row1);
+ int writeKey2 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ writeParallelism,
+ Collections.emptySet(),
+ distributionFields,
+ row2);
+ int writeKey3 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ writeParallelism,
+ Collections.emptySet(),
+ distributionFields,
+ row3);
+
+ assertThat(writeKey1).isEqualTo(writeKey2);
+ assertThat(writeKey2).isNotEqualTo(writeKey3);
+
+ assertThat(getSubTaskId(writeKey1, writeParallelism,
maxWriteParallelism)).isEqualTo(1);
+ assertThat(getSubTaskId(writeKey2, writeParallelism,
maxWriteParallelism)).isEqualTo(1);
+ assertThat(getSubTaskId(writeKey3, writeParallelism,
maxWriteParallelism)).isEqualTo(0);
+ }
+
+ @Test
+ void testDistributionFieldsIgnorePartitionSpec() throws Exception {
+ int writeParallelism = 4;
+ int maxWriteParallelism = 8;
+ HashKeyGenerator generator = new HashKeyGenerator(16, maxWriteParallelism);
+ PartitionSpec partitioned =
PartitionSpec.builderFor(SCHEMA).identity("id").build();
+ Set<String> distributionFields = Collections.singleton("data");
+
+ // Same 'data' but different partition values must still land on the same
writer
+ GenericRowData row1 = GenericRowData.of(1, StringData.fromString("x"));
+ GenericRowData row2 = GenericRowData.of(2, StringData.fromString("x"));
+
+ int writeKey1 =
+ getWriteKey(
+ generator,
+ partitioned,
+ DistributionMode.HASH,
+ writeParallelism,
+ Collections.emptySet(),
+ distributionFields,
+ row1);
+ int writeKey2 =
+ getWriteKey(
+ generator,
+ partitioned,
+ DistributionMode.HASH,
+ writeParallelism,
+ Collections.emptySet(),
+ distributionFields,
+ row2);
+ assertThat(writeKey1).isEqualTo(writeKey2);
+ }
+
+ @Test
+ void testDistributionFieldsMissingFromSchema() {
+ HashKeyGenerator generator = new HashKeyGenerator(16, 8);
+
+ assertThatThrownBy(
+ () ->
+ getWriteKey(
+ generator,
+ PartitionSpec.unpartitioned(),
+ DistributionMode.HASH,
+ 2,
+ Collections.emptySet(),
+ Collections.singleton("does_not_exist"),
+ GenericRowData.of(1, StringData.fromString("foo"))))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining(
+ "Missing distribution field 'does_not_exist'. "
+ + "Expected distribution field 'does_not_exist' to exist in
the table schema.");
+ }
+
+ @Test
+ void testDistributionFieldsSubsetOfEqualityFields() throws Exception {
+ HashKeyGenerator generator = new HashKeyGenerator(16, 8);
+ PartitionSpec unpartitioned = PartitionSpec.unpartitioned();
+ Set<String> equalityFields = ImmutableSet.of("id", "data");
+ Set<String> distributionFields = Collections.singleton("id");
+
+ int writeKey1 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ 2,
+ equalityFields,
+ distributionFields,
+ GenericRowData.of(1, StringData.fromString("a")));
+ int writeKey2 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ 2,
+ equalityFields,
+ distributionFields,
+ GenericRowData.of(1, StringData.fromString("b")));
+ int writeKey3 =
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ 2,
+ equalityFields,
+ distributionFields,
+ GenericRowData.of(2, StringData.fromString("a")));
+
+ assertThat(writeKey1).isEqualTo(writeKey2);
+ assertThat(writeKey1).isNotEqualTo(writeKey3);
+ }
+
+ @Test
+ void testDistributionFieldsNotSubsetOfEqualityFields() {
+ HashKeyGenerator generator = new HashKeyGenerator(16, 8);
+
+ assertThatThrownBy(
+ () ->
+ getWriteKey(
+ generator,
+ PartitionSpec.unpartitioned(),
+ DistributionMode.HASH,
+ 2,
+ Collections.singleton("id"),
+ Collections.singleton("data"),
+ GenericRowData.of(1, StringData.fromString("foo"))))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining(
+ "Invalid distribution fields '[data]'. Expected distribution
fields "
+ + "to be a subset of equality fields '[id]'.");
+ }
+
+ @Test
+ void testDistributionFieldsNotSubsetOfIdentifierFields() {
+ HashKeyGenerator generator = new HashKeyGenerator(16, 8);
+ Schema schemaWithIdentifier =
+ new Schema(
+ Lists.newArrayList(
+ Types.NestedField.required(1, "id", Types.IntegerType.get()),
+ Types.NestedField.required(2, "data", Types.StringType.get())),
+ Sets.newHashSet(1));
+
+ DynamicRecord record =
+ new DynamicRecord(
+ TABLE_IDENTIFIER,
+ BRANCH,
+ schemaWithIdentifier,
+ GenericRowData.of(1, StringData.fromString("foo")),
+ PartitionSpec.unpartitioned(),
+ DistributionMode.HASH,
+ 2);
+ record.setDistributionFields(Collections.singleton("data"));
+
+ assertThatThrownBy(() -> generator.generateKey(record))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining(
+ "Invalid distribution fields '[data]'. Expected distribution
fields "
+ + "to be a subset of equality fields '[id]'.");
+ }
+
+ @Test
+ void testCacheMissOnDistributionFieldsChange() throws Exception {
+ HashKeyGenerator generator = new HashKeyGenerator(10, 8);
+ Map<HashKeyGenerator.SelectorKey, KeySelector<RowData, Integer>>
keySelectorCache =
+ generator.getKeySelectorCache();
+
+ PartitionSpec unpartitioned = PartitionSpec.unpartitioned();
+ GenericRowData row = GenericRowData.of(1, StringData.fromString("foo"));
+
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ 2,
+ Collections.emptySet(),
+ Collections.singleton("id"),
+ row);
+ assertThat(keySelectorCache).hasSize(1);
+
+ getWriteKey(
+ generator,
+ unpartitioned,
+ DistributionMode.HASH,
+ 2,
+ Collections.emptySet(),
+ Collections.singleton("data"),
+ row);
+ assertThat(keySelectorCache).hasSize(2);
+ }
+
private static int getWriteKey(
HashKeyGenerator generator,
Schema schema,
@@ -843,12 +1062,54 @@ class TestHashKeyGenerator {
DistributionMode mode,
int writeParallelism,
Set<String> equalityFields,
+ Set<String> distributionFields,
+ GenericRowData row)
+ throws Exception {
+ return getWriteKey(
+ generator,
+ spec,
+ mode,
+ writeParallelism,
+ equalityFields,
+ distributionFields,
+ row,
+ new FlinkWriteConf(Collections.emptyMap(), new Configuration()));
+ }
+
+ private static int getWriteKey(
+ HashKeyGenerator generator,
+ PartitionSpec spec,
+ DistributionMode mode,
+ int writeParallelism,
+ Set<String> equalityFields,
+ GenericRowData row,
+ FlinkWriteConf flinkWriteConf)
+ throws Exception {
+ return getWriteKey(
+ generator,
+ spec,
+ mode,
+ writeParallelism,
+ equalityFields,
+ Collections.emptySet(),
+ row,
+ flinkWriteConf);
+ }
+
+ private static int getWriteKey(
+ HashKeyGenerator generator,
+ PartitionSpec spec,
+ DistributionMode mode,
+ int writeParallelism,
+ Set<String> equalityFields,
+ Set<String> distributionFields,
GenericRowData row,
FlinkWriteConf flinkWriteConf)
throws Exception {
DynamicRecord inputRecord =
new DynamicRecord(TABLE_IDENTIFIER, BRANCH, SCHEMA, row, spec, mode,
writeParallelism);
inputRecord.setEqualityFields(equalityFields);
+ inputRecord.setDistributionFields(distributionFields);
DynamicRecordWithConfig dynamicRecordWithConfig = new
DynamicRecordWithConfig(flinkWriteConf);
return generator.generateKey(dynamicRecordWithConfig.wrap(inputRecord));
