[IcebergIO] Spec aware file reserialization

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

@@ -28,6 +28,7 @@
 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;
@@ -37,11 +38,14 @@
 import org.apache.iceberg.FileMetadata;
 import org.apache.iceberg.Metrics;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.SingleValueParser;
 import org.apache.iceberg.SortOrder;
+import org.apache.iceberg.StructLike;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 @DefaultSchema(AutoValueSchema.class)
 @AutoValue
+@Internal
 public abstract class SerializableDeleteFile {
   public static SerializableDeleteFile.Builder builder() {
     return new AutoValue_SerializableDeleteFile.Builder();
@@ -62,7 +66,11 @@ public static SerializableDeleteFile.Builder builder() {
   @SchemaFieldNumber("4")
   public abstract long getFileSizeInBytes();
 
+  /**
+   * @deprecated Use {@link #getJsonPartition()} instead.
+   */
   @SchemaFieldNumber("5")
+  @Deprecated
   public abstract String getPartitionPath();
 
   @SchemaFieldNumber("6")
@@ -113,6 +121,9 @@ public static SerializableDeleteFile.Builder builder() {
   @SchemaFieldNumber("21")
   public abstract @Nullable Long getFileSequenceNumber();
 
+  @SchemaFieldNumber("22")
+  abstract @Nullable String getJsonPartition();
+
   @AutoValue.Builder
   abstract static class Builder {
     abstract Builder setContentType(FileContent content);
@@ -127,6 +138,8 @@ abstract static class Builder {
 
     abstract Builder setPartitionPath(String partitionPath);
 
+    abstract Builder setJsonPartition(String jsonPartition);
+
     abstract Builder setPartitionSpecId(int partitionSpec);
 
     abstract Builder setSortOrderId(@Nullable Integer sortOrderId);
@@ -163,14 +176,55 @@ abstract static class Builder {
   }
 
   public static SerializableDeleteFile from(
-      DeleteFile deleteFile, String partitionPath, boolean includeMetrics) {
+      DeleteFile deleteFile, Map<Integer, PartitionSpec> specs) {
+    return from(deleteFile, specs, true);
+  }
+
+  /**
+   * Creates a {@link SerializableDeleteFile}, resolving the file's {@link PartitionSpec} by its own
+   * spec id.
+   *
+   * <p>Delete files reached from a scan task may carry a spec id that differs from the spec of the
+   * data file they apply to, so the lookup has to be per delete file rather than against a single
+   * "current" spec.
+   */
+  public static SerializableDeleteFile from(
+      DeleteFile deleteFile, Map<Integer, PartitionSpec> specs, boolean includeMetrics) {
+    return from(
+        deleteFile,
+        checkStateNotNull(
+            specs.get(deleteFile.specId()),
+            "Could not create a SerializableDeleteFile because DeleteFile is written using a partition spec id '%s' that is not found in the provided specs: %s",
+            deleteFile.specId(),
+            specs.keySet()),
+        includeMetrics);
+  }
+
+  public static SerializableDeleteFile from(DeleteFile deleteFile, PartitionSpec spec) {
+    return from(deleteFile, spec, true);
+  }
+
+  public static SerializableDeleteFile from(
+      DeleteFile deleteFile, PartitionSpec spec, boolean includeMetrics) {
+    if (spec.specId() != deleteFile.specId()) {
+      throw new IllegalArgumentException(
+          String.format(
+              "Cannot serialize DeleteFile: its partition spec id %s does not match the provided "
+                  + "spec id %s.",
+              deleteFile.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(), deleteFile.partition());
+    String partitionPath = spec.partitionToPath(deleteFile.partition());
 
     SerializableDeleteFile.Builder builder =
         SerializableDeleteFile.builder()
             .setLocation(deleteFile.location())
             .setFileFormat(deleteFile.format().name())
             .setFileSizeInBytes(deleteFile.fileSizeInBytes())
             .setPartitionPath(partitionPath)
+            .setJsonPartition(jsonPartition)
             .setPartitionSpecId(deleteFile.specId())
             .setRecordCount(deleteFile.recordCount())
             .setColumnSizes(deleteFile.columnSizes())
@@ -228,7 +282,21 @@ public DeleteFile createDeleteFile(
             .withMetrics(metrics)
             .withSplitOffsets(getSplitOffsets())
             .withEncryptionKeyMetadata(getKeyMetadata())
-            .withPartitionPath(getPartitionPath());
+            .withReferencedDataFile(getReferencedDataFile());
+
+    @Nullable String jsonPartition = getJsonPartition();
+    if (jsonPartition != null) {
+      try {
+        deleteFileBuilder = deleteFileBuilder.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
+        deleteFileBuilder = deleteFileBuilder.withPartitionPath(getPartitionPath());
+      }
+    } else {
+      // Elements decoded from a pre-jsonPartition release carry only the partition path.
+      deleteFileBuilder = deleteFileBuilder.withPartitionPath(getPartitionPath());
+    }
 
     switch (getContentType()) {
       case POSITION_DELETES:
@@ -260,17 +328,22 @@ public DeleteFile createDeleteFile(
             "Unexpected content type for DeleteFile: " + getContentType());
     }
 
-    // needed for puffin files
+    // contentOffset / contentSizeInBytes really are Puffin-only: build() rejects a non-null value
+    // for either on any other format, and requires both (plus referencedDataFile) on Puffin.
     if (getFileFormat().equalsIgnoreCase(FileFormat.PUFFIN.name())) {
       deleteFileBuilder =
           deleteFileBuilder
               .withContentOffset(checkStateNotNull(getContentOffset()))
-              .withContentSizeInBytes(checkStateNotNull(getContentSizeInBytes()))
-              .withReferencedDataFile(checkStateNotNull(getReferencedDataFile()));
+              .withContentSizeInBytes(checkStateNotNull(getContentSizeInBytes()));
     }
     return deleteFileBuilder.build();
   }
 
+  private StructLike partition(PartitionSpec spec) {
+    return (StructLike)
+        SingleValueParser.fromJson(spec.partitionType(), checkStateNotNull(getJsonPartition()));
+  }
+
   @Override
   public final boolean equals(@Nullable Object o) {
     if (this == o) {
@@ -287,6 +360,7 @@ && getRecordCount() == that.getRecordCount()
         && getFileSizeInBytes() == that.getFileSizeInBytes()
         && getPartitionPath().equals(that.getPartitionPath())
         && getPartitionSpecId() == that.getPartitionSpecId()
+        && Objects.equals(getJsonPartition(), that.getJsonPartition())
         && Objects.equals(getSortOrderId(), that.getSortOrderId())
         && Objects.equals(getEqualityFieldIds(), that.getEqualityFieldIds())
         && Objects.equals(getKeyMetadata(), that.getKeyMetadata())
@@ -314,6 +388,7 @@ public final int hashCode() {
             getRecordCount(),
             getFileSizeInBytes(),
             getPartitionPath(),
+            getJsonPartition(),
             getPartitionSpecId(),
             getSortOrderId(),
             getEqualityFieldIds(),

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

@@ -17,7 +17,6 @@
  */
 package org.apache.beam.sdk.io.iceberg.cdc;
 
-import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import com.google.auto.value.AutoValue;
@@ -267,13 +266,10 @@ static List<DeleteFile> getAddedDeleteFiles(ChangelogScanTask task) {
 
   private static List<SerializableDeleteFile> toSerializableDeletes(
       List<DeleteFile> dfs, Map<Integer, PartitionSpec> specs, boolean includeMetrics) {
+    // Serialize each delete file against its own spec (looked up by its spec id): a delete file may
+    // carry a different spec id than the data file it applies to.
     return dfs.stream()
-        .map(
-            df ->
-                SerializableDeleteFile.from(
-                    df,
-                    checkStateNotNull(specs.get(df.specId())).partitionToPath(df.partition()),
-                    includeMetrics))
+        .map(df -> SerializableDeleteFile.from(df, specs, includeMetrics))
         .collect(Collectors.toList());
   }
 }