[IcebergIO] Fix partition values written by hash distribution

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run.",
-    "modification": 2
+    "modification": 3
 }

CHANGES.md

@@ -149,6 +149,7 @@
 * (Java/Python) The SQL `SqlTransformSchemaTransformProvider` configuration schema (`schematransform:org.apache.beam:sql_transform:v1`) now uses only portable types so it can be discovered from Python: `dialect` is a string, `tableproviders` is an array of strings, and the unused `parameters` field was removed. Invalid dialects and unknown table providers are rejected with a clear error ([#34613](https://github.com/apache/beam/issues/34613)).
 * (Python) An unknown pipeline option is no longer given the value of a later argument that does not directly follow it, which made `Pipeline.run()` and `get_all_options(retain_unknown_options=True)` fail with `expected one argument` for flags like `--unknown_flag --job_name=foo bar` ([#22434](https://github.com/apache/beam/issues/22434)).
 * (Python) Fixed `WriteToJson` silently rounding floats to 10 decimal places: it now defaults pandas' `double_precision` to 15, so values such as `0.1234567890123` and `1e-11` are written correctly. Floats that need 16-17 significant digits may still be rounded ([#41102](https://github.com/apache/beam/issues/41102)).
+* (Java) Makes some fixes to IcebergIO's hash distribution mode, which sometimes wrote incorrect partition values. ([#41012](https://github.com/apache/beam/issues/41012)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/BeamRowWrapper.java

@@ -35,6 +35,7 @@
 import org.apache.beam.sdk.schemas.logicaltypes.FixedPrecisionNumeric;
 import org.apache.beam.sdk.schemas.logicaltypes.MicrosInstant;
 import org.apache.beam.sdk.schemas.logicaltypes.Time;
+import org.apache.beam.sdk.schemas.logicaltypes.Timestamp;
 import org.apache.beam.sdk.values.Row;
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.types.Type;
@@ -158,7 +159,8 @@ private interface PositionalGetter<T> {
         BeamRowWrapper nestedWrapper = new BeamRowWrapper(beamSchema, structType);
         return (row, pos) -> nestedWrapper.wrap(row.getRow(pos));
       case LOGICAL_TYPE:
-        if (beamType.isLogicalType(MicrosInstant.IDENTIFIER)) {
+        if (beamType.isLogicalType(MicrosInstant.IDENTIFIER)
+            || beamType.isLogicalType(Timestamp.IDENTIFIER)) {
           return (row, pos) -> {
             Instant instant = checkStateNotNull(row.getLogicalTypeValue(pos, Instant.class));
             return TimeUnit.SECONDS.toMicros(instant.getEpochSecond()) + instant.getNano() / 1000;

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/RecordWriterManager.java

@@ -21,11 +21,6 @@
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.io.IOException;
-import java.time.LocalDateTime;
-import java.time.YearMonth;
-import java.time.ZoneOffset;
-import java.time.format.DateTimeFormatter;
-import java.time.temporal.ChronoUnit;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
@@ -36,15 +31,13 @@
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.sdk.values.WindowedValue;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
-import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.Cache;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.CacheBuilder;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.cache.RemovalNotification;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.ManifestFile;
-import org.apache.iceberg.PartitionField;
 import org.apache.iceberg.PartitionKey;
 import org.apache.iceberg.PartitionSpec;
 import org.apache.iceberg.SortOrder;
@@ -57,7 +50,6 @@
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.NoSuchTableException;
-import org.apache.iceberg.transforms.Transforms;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -104,7 +96,7 @@ final class DestinationState {
     private final String stateToken = UUID.randomUUID().toString();
     final Cache<PartitionKey, RecordWriter> writers;
     private final List<SerializableDataFile> dataFiles = Lists.newArrayList();
-    @VisibleForTesting final Map<PartitionKey, Integer> writerCounts = Maps.newHashMap();
+    @VisibleForTesting final Map<String, Integer> writerCounts = Maps.newHashMap();
     private final List<Exception> exceptions = Lists.newArrayList();
     private final InternalRecordWrapper wrapper; // wrapper that facilitates partitioning
 
@@ -191,13 +183,13 @@ private RecordWriter createWriter(PartitionKey partitionKey) {
       // keep track of how many writers we opened for each destination-partition path
       // use this as a prefix to differentiate the new path.
       // this avoids overwriting a data file written by a previous writer in this destination state.
-      int recordIndex = writerCounts.merge(partitionKey, 1, Integer::sum);
+      int writerIndex = writerCounts.merge(partitionKey.toPath(), 1, Integer::sum);
       try {
         RecordWriter writer =
             new RecordWriter(
                 table,
                 icebergDestination.getFileFormat(),
-                filePrefix + "_" + stateToken + "_" + recordIndex,
+                filePrefix + "_" + stateToken + "_" + writerIndex,
                 partitionKey,
                 writeProperties);
         openWriters++;
@@ -212,40 +204,6 @@ private RecordWriter createWriter(PartitionKey partitionKey) {
     }
   }
 
-  /**
-   * Returns an equivalent partition path that is made up of partition data. Needed to reconstruct a
-   * {@link DataFile}.
-   */
-  @VisibleForTesting
-  static String getPartitionDataPath(
-      String partitionPath, Map<String, PartitionField> partitionFieldMap) {
-    if (partitionPath.isEmpty() || partitionFieldMap.isEmpty()) {
-      return partitionPath;
-    }
-    List<String> resolved = new ArrayList<>();
-    for (String partition : Splitter.on('/').splitToList(partitionPath)) {
-      List<String> nameAndValue = Splitter.on('=').splitToList(partition);
-      String name = nameAndValue.get(0);
-      String value = nameAndValue.get(1);
-      String transformName =
-          Preconditions.checkArgumentNotNull(partitionFieldMap.get(name)).transform().toString();
-      if (Transforms.month().toString().equals(transformName)) {
-        long months =
-            ChronoUnit.MONTHS.between(EPOCH, YearMonth.parse(value).atDay(1).atStartOfDay());
-        value = String.valueOf(months);
-      } else if (Transforms.hour().toString().equals(transformName)) {
-        long hour = ChronoUnit.HOURS.between(EPOCH, LocalDateTime.parse(value, HOUR_FORMATTER));
-        value = String.valueOf(hour);
-      }
-      resolved.add(name + "=" + value);
-    }
-    return String.join("/", resolved);
-  }
-
-  private static final DateTimeFormatter HOUR_FORMATTER =
-      DateTimeFormatter.ofPattern("yyyy-MM-dd-HH");
-  private static final LocalDateTime EPOCH = LocalDateTime.ofEpochSecond(0, 0, ZoneOffset.UTC);
-
   private final IcebergCatalogConfig catalogConfig;
   private final String filePrefix;
   private final long maxFileSize;