mxm commented on code in PR #18668:
URL: https://github.com/apache/iceberg/pull/18668#discussion_r4426212007
##########
flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestHashKeyGenerator.java:
##########
@@ -754,6 +754,172 @@ void testCacheMissOnDistributionModeChange() throws
Exception {
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 =
+ distributionFieldsWriteKey(
+ generator, unpartitioned, writeParallelism, distributionFields,
row1);
+ int writeKey2 =
+ distributionFieldsWriteKey(
+ generator, unpartitioned, writeParallelism, distributionFields,
row2);
+ int writeKey3 =
+ distributionFieldsWriteKey(
+ generator, unpartitioned, writeParallelism, 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"));
+
+ assertThat(
+ distributionFieldsWriteKey(
+ generator, partitioned, writeParallelism, distributionFields,
row1))
+ .isEqualTo(
+ distributionFieldsWriteKey(
+ generator, partitioned, writeParallelism, distributionFields,
row2));
+
+ // Distinct 'data' values must not all collapse onto a single writer
+ Set<Integer> writeKeys = Sets.newHashSet();
+ for (int i = 0; i < 8; i++) {
+ writeKeys.add(
+ distributionFieldsWriteKey(
+ generator,
+ partitioned,
+ writeParallelism,
+ distributionFields,
+ GenericRowData.of(i, StringData.fromString("value-" + i))));
+ }
+
+ assertThat(writeKeys).hasSizeGreaterThan(1);
+ }
+
+ @Test
+ void testDistributionFieldsMissingFromSchema() {
+ HashKeyGenerator generator = new HashKeyGenerator(16, 8);
+ DynamicRecord record =
+ distributionFieldsRecord(
+ PartitionSpec.unpartitioned(),
+ DistributionMode.HASH,
+ Collections.singleton("does_not_exist"));
+
+ assertThatThrownBy(() -> generator.generateKey(record))
+ .isInstanceOf(NullPointerException.class)
+ .hasMessageContaining(
+ "Missing distribution field 'does_not_exist'. "
+ + "Expected distribution field 'does_not_exist' to exist in
the table schema.");
+ }
+
+ @Test
+ void testDistributionFieldsNotSubsetOfEqualityFields() {
Review Comment:
Could we also test the accepted case, e.g. equality fields `{id, data}` and
distribution field `{id}`? Rows `(1, "a")` and `(1, "b")` should then get the
same writer key.
##########
flink/v2.3/flink/src/main/java/org/apache/iceberg/flink/sink/dynamic/HashKeyGenerator.java:
##########
@@ -117,23 +122,52 @@ private KeySelector<RowData, Integer> getKeySelector(
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 '%s' to be a subset of "
+ + "equality fields '%s'.",
Review Comment:
NIT: This prints the distribution fields twice. Could we shorten it, e.g.
`"%s: Distribution fields %s must be a subset of equality fields %s"`?
##########
flink/v2.3/flink/src/test/java/org/apache/iceberg/flink/sink/dynamic/TestDynamicIcebergSink.java:
##########
@@ -1184,6 +1215,43 @@ void testUpsertV3(boolean immediateUpdate) throws
Exception {
}
}
+ @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();
+ for (Snapshot snapshot : dataSnapshots) {
+ assertThat(snapshot.addedDataFiles(table.io())).hasSize(1);
+ }
Review Comment:
NIT: I'd prefer not to use loops for assertions. Makes test failures easier
to reason about.
```suggestion
assertThat(dataSnapshots)
.isNotEmpty()
.allSatisfy(snapshot ->
assertThat(snapshot.addedDataFiles(table.io())).hasSize(1));
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
