[IcebergIO] Serialize using json partition

.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": 1
+    "modification": 2
 }

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

@@ -514,7 +514,7 @@ private Callable<ProcessResult> createProcessTask(
                 .withPartitionPath(partitionPath)
                 .build();
         return new ProcessResult(
-            SerializableDataFile.from(df, partitionPath), null, timestamp, window, paneInfo);
+            SerializableDataFile.from(df, table.spec()), null, timestamp, window, paneInfo);
       };
     }
 

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

@@ -154,7 +154,7 @@ static Term toIcebergTerm(String field) {
    * {@link ContentScanTask}s.
    */
   public static Map<Integer, ?> constantsMap(
-      PartitionSpec spec, ContentFile<?> file, @Nullable Long fileSequenceNumber) {
+      PartitionSpec spec, ContentFile<?> file, @Nullable Long dataSequenceNumber) {
     Preconditions.checkState(
         spec.specId() == file.specId(),
         "File spec ID (%s) does not match PartitionSpec ID (%s)",
@@ -172,13 +172,13 @@ static Term toIcebergTerm(String field) {
           convertConstant(Types.LongType.get(), file.firstRowId()));
     }
 
-    // When reconstructing a DataFile, we lose the ability to attach its fileSequenceNumber,
+    // When reconstructing a DataFile, we lose the ability to attach its dataSequenceNumber,
     // so we pipe it along the util methods to include it here.
-    fileSequenceNumber =
-        fileSequenceNumber != null ? fileSequenceNumber : file.fileSequenceNumber();
+    dataSequenceNumber =
+        dataSequenceNumber != null ? dataSequenceNumber : file.dataSequenceNumber();
     idToConstant.put(
         MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId(),
-        convertConstant(Types.LongType.get(), fileSequenceNumber));
+        convertConstant(Types.LongType.get(), dataSequenceNumber));
 
     // add _file
     idToConstant.put(

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

@@ -104,7 +104,6 @@ class DestinationState {
     final Cache<PartitionKey, RecordWriter> writers;
     private final List<SerializableDataFile> dataFiles = Lists.newArrayList();
     @VisibleForTesting final Map<PartitionKey, Integer> writerCounts = Maps.newHashMap();
-    private final Map<String, PartitionField> partitionFieldMap = Maps.newHashMap();
     private final List<Exception> exceptions = Lists.newArrayList();
     private final InternalRecordWrapper wrapper; // wrapper that facilitates partitioning
 
@@ -115,9 +114,6 @@ class DestinationState {
       this.routingPartitionKey = new PartitionKey(spec, schema);
       this.wrapper = new InternalRecordWrapper(schema.asStruct());
       this.table = table;
-      for (PartitionField partitionField : spec.fields()) {
-        partitionFieldMap.put(partitionField.name(), partitionField);
-      }
 
       // build a cache of RecordWriters.
       // writers will expire after 1 min of idle time.
@@ -127,7 +123,6 @@ class DestinationState {
               .expireAfterAccess(1, TimeUnit.MINUTES)
               .removalListener(
                   (RemovalNotification<PartitionKey, RecordWriter> removal) -> {
-                    final PartitionKey pk = Preconditions.checkStateNotNull(removal.getKey());
                     final RecordWriter recordWriter =
                         Preconditions.checkStateNotNull(removal.getValue());
                     try {
@@ -144,9 +139,9 @@ class DestinationState {
                       throw rethrow;
                     }
                     openWriters--;
-                    String partitionPath = getPartitionDataPath(pk.toPath(), partitionFieldMap);
+                    // Serialize against the file's own spec (looked up by its spec id)
                     dataFiles.add(
-                        SerializableDataFile.from(recordWriter.getDataFile(), partitionPath));
+                        SerializableDataFile.from(recordWriter.getDataFile(), table.specs()));
                   })
               .build();
     }

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

@@ -26,9 +26,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import org.apache.beam.sdk.annotations.Internal;
 import org.apache.beam.sdk.schemas.AutoValueSchema;
 import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
 import org.apache.beam.sdk.schemas.annotations.SchemaFieldNumber;
+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.Equivalence;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps;
 import org.apache.iceberg.DataFile;
@@ -37,6 +39,8 @@
 import org.apache.iceberg.Metrics;
 import org.apache.iceberg.PartitionKey;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.SingleValueParser;
+import org.apache.iceberg.StructLike;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 /**
@@ -49,11 +53,12 @@
  * <p>NOTE: If you add any new fields here, you need to also update the {@link #equals} and {@link
  * #hashCode()} methods.
  *
- * <p>Use {@link #from(DataFile, String)} to create a {@link SerializableDataFile} and {@link
+ * <p>Use {@link #from(DataFile, PartitionSpec)} to create a {@link SerializableDataFile} and {@link
  * #createDataFile(Map)} to reconstruct the original {@link DataFile}.
  */
 @DefaultSchema(AutoValueSchema.class)
 @AutoValue
+@Internal
 public abstract class SerializableDataFile {
   public static Builder builder() {
     return new AutoValue_SerializableDataFile.Builder();
@@ -71,7 +76,9 @@ public static Builder builder() {
   @SchemaFieldNumber("3")
   public abstract long getFileSizeInBytes();
 
+  /** @deprecated Use {@link #getJsonPartition()} instead. */
   @SchemaFieldNumber("4")
+  @Deprecated
   public abstract String getPartitionPath();
 
   @SchemaFieldNumber("5")
@@ -110,6 +117,9 @@ public static Builder builder() {
   @SchemaFieldNumber("16")
   public abstract @Nullable Long getFirstRowId();
 
+  @SchemaFieldNumber("17")
+  abstract @Nullable String getJsonPartition();
+
   @AutoValue.Builder
   public abstract static class Builder {
     abstract Builder setPath(String path);
@@ -122,6 +132,8 @@ public abstract static class Builder {
 
     abstract Builder setPartitionPath(String partitionPath);
 
+    abstract Builder setJsonPartition(String jsonPartition);
+
     abstract Builder setPartitionSpecId(int partitionSpec);
 
     abstract Builder setKeyMetadata(ByteBuffer keyMetadata);
@@ -149,23 +161,46 @@ public abstract static class Builder {
     abstract SerializableDataFile build();
   }
 
-  public static SerializableDataFile from(DataFile f, String partitionPath) {
-    return from(f, partitionPath, true);
+  public static SerializableDataFile from(DataFile f, Map<Integer, PartitionSpec> specs) {
+    return from(
+        f,
+        checkStateNotNull(
+            specs.get(f.specId()),
+            "Could not create a SerializableDataFile because DataFile is written using a partition spec id '%s' that is not found in the provided specs: %s",
+            f.specId(),
+            specs.keySet()),
+        true);
+  }
+
+  public static SerializableDataFile from(DataFile f, PartitionSpec spec) {
+    return from(f, spec, true);
   }
 
   /**
    * Create a {@link SerializableDataFile} from a {@link DataFile} and its associated {@link
    * PartitionKey}.
    */
-  public static SerializableDataFile from(
-      DataFile f, String partitionPath, boolean includeMetrics) {
+  public static SerializableDataFile from(DataFile f, PartitionSpec spec, boolean includeMetrics) {
+    if (spec.specId() != f.specId()) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Cannot serialize DataFile: its partition spec id %s does not match the provided "
+                  + "spec id %s. Serialize the file with the exact spec it was written with.",
+              f.specId(), spec.specId()));
+    }
+    // jsonPartition is the primary (handles evolved specs, special characters).
+    // partitionPath is the fallback for values that don't round-trip through JSON.
+    String jsonPartition = SingleValueParser.toJson(spec.partitionType(), f.partition());
+    String partitionPath = spec.partitionToPath(f.partition());
+
     SerializableDataFile.Builder builder =
         SerializableDataFile.builder()
             .setPath(f.location())
             .setFileFormat(f.format().toString())
             .setRecordCount(f.recordCount())
             .setFileSizeInBytes(f.fileSizeInBytes())
             .setPartitionPath(partitionPath)
+            .setJsonPartition(jsonPartition)
             .setPartitionSpecId(f.specId())
             .setKeyMetadata(f.keyMetadata())
             .setSplitOffsets(f.splitOffsets())
@@ -211,16 +246,36 @@ public DataFile createDataFile(Map<Integer, PartitionSpec> partitionSpecs) {
             toByteBufferMap(getLowerBounds()),
             toByteBufferMap(getUpperBounds()));
 
-    return DataFiles.builder(partitionSpec)
-        .withFormat(FileFormat.fromString(getFileFormat()))
-        .withPath(getPath())
-        .withPartitionPath(getPartitionPath())
-        .withEncryptionKeyMetadata(getKeyMetadata())
-        .withFileSizeInBytes(getFileSizeInBytes())
-        .withMetrics(dataFileMetrics)
-        .withSplitOffsets(getSplitOffsets())
-        .withFirstRowId(getFirstRowId())
-        .build();
+    DataFiles.Builder builder =
+        DataFiles.builder(partitionSpec)
+            .withFormat(FileFormat.fromString(getFileFormat()))
+            .withPath(getPath())
+            .withEncryptionKeyMetadata(getKeyMetadata())
+            .withFileSizeInBytes(getFileSizeInBytes())
+            .withMetrics(dataFileMetrics)
+            .withSplitOffsets(getSplitOffsets())
+            .withFirstRowId(getFirstRowId());
+
+    @Nullable String jsonPartition = getJsonPartition();
+    if (jsonPartition != null) {
+      try {
+        builder = builder.withPartition(partition(partitionSpec));
+      } catch (RuntimeException e) {
+        // Some partition values (e.g. NaN / Infinity floating-point) don't round-trip through the
+        // JSON representation; fall back to the partition-path string, which handles them.
+        builder = builder.withPartitionPath(getPartitionPath());
+      }
+    } else {
+      // Elements decoded from a pre-jsonPartition release carry only the partition path.
+      builder = builder.withPartitionPath(getPartitionPath());
+    }
+    return builder.build();
+  }
+
+  @VisibleForTesting
+  StructLike partition(PartitionSpec spec) {
+    return (StructLike)
+        SingleValueParser.fromJson(spec.partitionType(), checkStateNotNull(getJsonPartition()));
   }
 
   // ByteBuddyUtils has trouble converting Map value type ByteBuffer
@@ -275,6 +330,8 @@ && getRecordCount() == that.getRecordCount()
         && getFileSizeInBytes() == that.getFileSizeInBytes()
         && getPartitionPath().equals(that.getPartitionPath())
         && getPartitionSpecId() == that.getPartitionSpecId()
+        && Objects.equals(getPartitionPath(), that.getPartitionPath())
+        && Objects.equals(getJsonPartition(), that.getJsonPartition())
         && Objects.equals(getKeyMetadata(), that.getKeyMetadata())
         && Objects.equals(getSplitOffsets(), that.getSplitOffsets())
         && Objects.equals(getColumnSizes(), that.getColumnSizes())
@@ -320,6 +377,7 @@ public final int hashCode() {
             getRecordCount(),
             getFileSizeInBytes(),
             getPartitionPath(),
+            getJsonPartition(),
             getPartitionSpecId(),
             getKeyMetadata(),
             getSplitOffsets(),

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

@@ -152,7 +152,8 @@ public void processElement(
         writer.close();
       }
 
-      SerializableDataFile sdf = SerializableDataFile.from(writer.getDataFile(), partitionPath);
+      // Serialize against the file's own spec
+      SerializableDataFile sdf = SerializableDataFile.from(writer.getDataFile(), table.specs());
       out.output(
           FileWriteResult.builder()
               .setTableIdentifier(destination.getTableIdentifier())

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

@@ -135,7 +135,7 @@ public static CloseableIterable<Record> createReader(
         outputSchema,
         checkStateNotNull(table.specs().get(task.getSpecId())),
         task.getDataFile().createDataFile(table.specs()),
-        task.getDataFile().getFileSequenceNumber(),
+        task.getDataFile().getDataSequenceNumber(),
         start,
         length,
         combined);

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

@@ -118,8 +118,8 @@ public abstract static class Builder {
     abstract Builder setDataFile(SerializableDataFile dataFile);
 
     @SchemaIgnore
-    public Builder setDataFile(DataFile df, String partitionPath, boolean includeMetrics) {
-      return setDataFile(SerializableDataFile.from(df, partitionPath, includeMetrics));
+    public Builder setDataFile(DataFile df, PartitionSpec spec, boolean includeMetrics) {
+      return setDataFile(SerializableDataFile.from(df, spec, includeMetrics));
     }
 
     abstract Builder setExistingDeletes(List<SerializableDeleteFile> existingDeletes);
@@ -159,10 +159,7 @@ public static SerializableChangelogTask from(
             .setOperation(task.operation())
             .setOrdinal(task.changeOrdinal())
             .setCommitSnapshotId(task.commitSnapshotId())
-            .setDataFile(
-                contentScanTask.file(),
-                spec.partitionToPath(contentScanTask.partition()),
-                includeMetrics)
+            .setDataFile(contentScanTask.file(), spec, includeMetrics)
             .setSpecId(spec.specId())
             .setStart(contentScanTask.start())
             .setLength(contentScanTask.length())

sdks/java/io/iceberg/src/test/java/org/apache/beam/sdk/io/iceberg/PartitionUtilsTest.java

@@ -190,7 +190,7 @@ public void testConstantsMapIncludesCdcMetadataAndIdentityConstants() throws Exc
             .withRecordCount(2L)
             .withFirstRowId(99L)
             .build();
-    setFileSequenceNumber(file, 42L);
+    setDataSequenceNumber(file, 42L);
 
     Map<Integer, ?> constants = PartitionUtils.constantsMap(spec, file, null);
 
@@ -202,7 +202,7 @@ public void testConstantsMapIncludesCdcMetadataAndIdentityConstants() throws Exc
   }
 
   @Test
-  public void testConstantsMapUsesExplicitSequenceNumberWhenFileSequenceIsUnavailable() {
+  public void testConstantsMapUsesExplicitSequenceNumberWhenDataSequenceIsUnavailable() {
     org.apache.iceberg.Schema icebergSchema =
         new org.apache.iceberg.Schema(
             Types.NestedField.required(1, "id", Types.IntegerType.get()),
@@ -223,12 +223,12 @@ public void testConstantsMapUsesExplicitSequenceNumberWhenFileSequenceIsUnavaila
     assertEquals("B", constants.get(2));
   }
 
-  private static void setFileSequenceNumber(DataFile dataFile, long fileSequenceNumber)
+  private static void setDataSequenceNumber(DataFile dataFile, long dataSequenceNumber)
       throws Exception {
-    Method method = dataFile.getClass().getMethod("setFileSequenceNumber", Long.class);
+    Method method = dataFile.getClass().getMethod("setDataSequenceNumber", Long.class);
     method.setAccessible(true);
     try {
-      method.invoke(dataFile, fileSequenceNumber);
+      method.invoke(dataFile, dataSequenceNumber);
     } catch (InvocationTargetException e) {
       throw (Exception) e.getCause();
     }