Spark 4.0: Add Support for PartitionStatistics Files in RewriteTablePath.

core/src/main/java/org/apache/iceberg/RewriteTablePathUtil.java

@@ -129,8 +129,8 @@ public static TableMetadata replacePaths(
         metadataLogEntries,
         metadata.refs(),
         updatePathInStatisticsFiles(metadata.statisticsFiles(), sourcePrefix, targetPrefix),
-        // TODO: update partition statistics file paths
-        metadata.partitionStatisticsFiles(),
+        updatePathInPartitionStatisticsFiles(
+            metadata.partitionStatisticsFiles(), sourcePrefix, targetPrefix),
         metadata.nextRowId(),
         metadata.encryptionKeys(),
         metadata.changes());
@@ -175,6 +175,31 @@ private static List<StatisticsFile> updatePathInStatisticsFiles(
         .collect(Collectors.toList());
   }
 
+  /**
+   * This method updates the file paths in a list of PartitionStatisticsFile. It replaces the
+   * sourcePrefix in the file paths with the targetPrefix.
+   *
+   * @param partitionStatisticsFiles The list of PartitionStatisticsFile to update.
+   * @param sourcePrefix The prefix to be replaced in the file paths.
+   * @param targetPrefix The new prefix to replace the sourcePrefix in the file paths.
+   * @return A new list of PartitionStatisticsFile with updated file paths.
+   */
+  private static List<PartitionStatisticsFile> updatePathInPartitionStatisticsFiles(
+      List<PartitionStatisticsFile> partitionStatisticsFiles,
+      String sourcePrefix,
+      String targetPrefix) {
+
+    return partitionStatisticsFiles.stream()
+        .map(
+            existing ->
+                ImmutableGenericPartitionStatisticsFile.builder()
+                    .snapshotId(existing.snapshotId())
+                    .path(newPath(existing.path(), sourcePrefix, targetPrefix))
+                    .fileSizeInBytes(existing.fileSizeInBytes())
+                    .build())
+        .collect(Collectors.toList());
+  }
+
   private static List<TableMetadata.MetadataLogEntry> updatePathInMetadataLogs(
       TableMetadata metadata, String sourcePrefix, String targetPrefix) {
     List<TableMetadata.MetadataLogEntry> metadataLogEntries =

spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/actions/RewriteTablePathSparkAction.java

@@ -34,6 +34,7 @@
 import org.apache.iceberg.HasTableOperations;
 import org.apache.iceberg.ManifestFile;
 import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.PartitionStatisticsFile;
 import org.apache.iceberg.RewriteTablePathUtil;
 import org.apache.iceberg.RewriteTablePathUtil.PositionDeleteReaderWriter;
 import org.apache.iceberg.RewriteTablePathUtil.RewriteResult;
@@ -273,11 +274,6 @@ private String rebuildMetadata() {
     TableMetadata endMetadata =
         ((HasTableOperations) newStaticTable(endVersionName, table.io())).operations().current();
 
-    Preconditions.checkArgument(
-        endMetadata.partitionStatisticsFiles() == null
-            || endMetadata.partitionStatisticsFiles().isEmpty(),
-        "Partition statistics files are not supported yet.");
-
     // rebuild version files
     RewriteResult<Snapshot> rewriteVersionResult = rewriteVersionFiles(endMetadata);
     Set<Snapshot> deltaSnapshots = deltaSnapshots(startMetadata, rewriteVersionResult.toRewrite());
@@ -385,6 +381,9 @@ private Set<Pair<String, String>> rewriteVersionFile(
     // include statistics files in copy plan
     result.addAll(
         statsFileCopyPlan(metadata.statisticsFiles(), newTableMetadata.statisticsFiles()));
+    result.addAll(
+        partitionStatsFileCopyPlan(
+            metadata.partitionStatisticsFiles(), newTableMetadata.partitionStatisticsFiles()));
     return result;
   }
 
@@ -409,6 +408,27 @@ private Set<Pair<String, String>> statsFileCopyPlan(
     return result;
   }
 
+  private Set<Pair<String, String>> partitionStatsFileCopyPlan(
+      List<PartitionStatisticsFile> beforeStats, List<PartitionStatisticsFile> afterStats) {
+    Set<Pair<String, String>> result = Sets.newHashSet();
+    if (beforeStats.isEmpty()) {
+      return result;
+    }
+
+    Preconditions.checkArgument(
+        beforeStats.size() == afterStats.size(),
+        "Before and after path rewrite, partition statistic files count should be same");
+    for (int i = 0; i < beforeStats.size(); i++) {
+      PartitionStatisticsFile before = beforeStats.get(i);
+      PartitionStatisticsFile after = afterStats.get(i);
+      Preconditions.checkArgument(
+          before.fileSizeInBytes() == after.fileSizeInBytes(),
+          "Before and after path rewrite, partition statistic file size should be same");
+      result.add(Pair.of(before.path(), after.path()));
+    }
+    return result;
+  }
+
   /**
    * Rewrite a manifest list representing a snapshot.
    *

spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteTablePathsAction.java

@@ -902,36 +902,37 @@ public void testInvalidArgs() {
   }
 
   @Test
-  public void testPartitionStatisticFile() throws IOException {
+  public void testTableWithPartitionStatisticFile() throws IOException {
     String sourceTableLocation = newTableLocation();
     Map<String, String> properties = Maps.newHashMap();
     properties.put("format-version", "2");
     String tableName = "v2tblwithPartStats";
     Table sourceTable =
         createMetastoreTable(sourceTableLocation, properties, "default", tableName, 0);
 
+    sql("insert into hive.default.%s values (%s, 'AAAAAAAAAA', 'AAAA')", tableName, 0);
+    sourceTable.refresh();
+
     TableMetadata metadata = currentMetadata(sourceTable);
+    File statFile = new File(removePrefix(sourceTableLocation + "/stats/file.parquet"));
     TableMetadata withPartStatistics =
         TableMetadata.buildFrom(metadata)
             .setPartitionStatistics(
                 ImmutableGenericPartitionStatisticsFile.builder()
                     .snapshotId(11L)
-                    .path("/some/partition/stats/file.parquet")
+                    .path(statFile.toURI().toString())
                     .fileSizeInBytes(42L)
                     .build())
             .build();
-
     OutputFile file = sourceTable.io().newOutputFile(metadata.metadataFileLocation());
     TableMetadataParser.overwrite(withPartStatistics, file);
 
-    assertThatThrownBy(
-            () ->
-                actions()
-                    .rewriteTablePath(sourceTable)
-                    .rewriteLocationPrefix(sourceTableLocation, targetTableLocation())
-                    .execute())
-        .isInstanceOf(IllegalArgumentException.class)
-        .hasMessageContaining("Partition statistics files are not supported yet");
+    RewriteTablePath.Result result =
+        actions()
+            .rewriteTablePath(sourceTable)
+            .rewriteLocationPrefix(sourceTableLocation, targetTableLocation())
+            .execute();
+    checkFileNum(2, 1, 1, 6, result);
   }
 
   @Test