Spark: Adding simple custom partition sort order option to RewriteManifests Spark Action

.palantir/revapi.yml

@@ -873,6 +873,10 @@ acceptedBreaks:
       new: "method void org.apache.iceberg.encryption.Ciphers::<init>()"
       justification: "Static utility class - should not have public constructor"
   "1.4.0":
+    org.apache.iceberg:iceberg-api:
+    - code: "java.method.addedToInterface"
+      new: "method org.apache.iceberg.actions.RewriteManifests org.apache.iceberg.actions.RewriteManifests::sort(java.util.List<java.lang.String>)"
+      justification: "New method added to Spark RewriteManifests"
     org.apache.iceberg:iceberg-core:
     - code: "java.class.defaultSerializationChanged"
       old: "class org.apache.iceberg.PartitionData"

api/src/main/java/org/apache/iceberg/actions/RewriteManifests.java

@@ -18,6 +18,7 @@
  */
 package org.apache.iceberg.actions;
 
+import java.util.List;
 import java.util.function.Predicate;
 import org.apache.iceberg.ManifestFile;
 
@@ -44,6 +45,16 @@ public interface RewriteManifests
    */
   RewriteManifests rewriteIf(Predicate<ManifestFile> predicate);
 
+  /**
+   * Rewrite manifests in a given order, based on partition columns
+   *
+   * <p>If not set, manifests will be rewritten in the order of the table's spec.
+   *
+   * @param partitionSortOrder a list of partition field names
+   * @return this for method chaining
+   */
+  RewriteManifests sort(List<String> partitionSortOrder);
+
   /**
    * Passes a location where the staged manifests should be written.
    *

spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/RewriteManifestsSparkAction.java

@@ -19,10 +19,13 @@
 package org.apache.iceberg.spark.actions;
 
 import static org.apache.iceberg.MetadataTableType.ENTRIES;
+import static org.apache.spark.sql.functions.col;
+import static org.apache.spark.sql.functions.concat_ws;
 
 import java.io.Serializable;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -90,13 +93,45 @@ public class RewriteManifestsSparkAction
   public static final String USE_CACHING = "use-caching";
   public static final boolean USE_CACHING_DEFAULT = false;
 
+  private List<String> partitionSortColumns = null;
   private static final Logger LOG = LoggerFactory.getLogger(RewriteManifestsSparkAction.class);
   private static final RewriteManifests.Result EMPTY_RESULT =
       ImmutableRewriteManifests.Result.builder()
           .rewrittenManifests(ImmutableList.of())
           .addedManifests(ImmutableList.of())
           .build();
 
+  /**
+   * Supply an optional set of partition columns to sort the rewritten manifests by. Expects real
+   * column names used for partitioning; will resolve to the proper hidden partition names.
+   *
+   * @param partitionSortOrder - Order of partitions to sort manifests by
+   * @return this action
+   */
+  @Override
+  public RewriteManifestsSparkAction sort(List<String> partitionSortOrder) {
+    // Build up a mapping of input column -> partition column name (AKA x -> x_bucket_1000)
+    Map<String, String> partitionFieldMap =
+        spec.fields().stream()
+            .map(
+                partitionField ->
+                    Map.entry(
+                        spec.schema().findField(partitionField.sourceId()).name(),
+                        partitionField.name()))
+            .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+    // Check if these partitions are included in the spec
+    Preconditions.checkArgument(
+        partitionSortOrder.stream().allMatch(partitionFieldMap::containsKey),
+        "Invalid sorting columns. Expected to match column names used for partitioning in spec: %s."
+            + " Available partitionable columns: %s",
+        this.spec.specId(),
+        partitionFieldMap.keySet());
+    this.partitionSortColumns =
+        partitionSortOrder.stream().map(partitionFieldMap::get).collect(Collectors.toList());
+    return this;
+  }
+
   private final Table table;
   private final int formatVersion;
   private final long targetManifestSizeBytes;
@@ -208,7 +243,15 @@ private RewriteManifests.Result rewriteManifests(ManifestContent content) {
     if (spec.isUnpartitioned()) {
       newManifests = writeUnpartitionedManifests(content, manifestEntryDF, targetNumManifests);
     } else {
-      newManifests = writePartitionedManifests(content, manifestEntryDF, targetNumManifests);
+      if (partitionSortColumns != null) {
+        LOG.info(
+            "Sorting manifests for specId {} by partition columns in order of {} ",
+            spec.specId(),
+            partitionSortColumns);
+        newManifests = writeSortedManifests(content, manifestEntryDF, targetNumManifests);
+      } else {
+        newManifests = writePartitionedManifests(content, manifestEntryDF, targetNumManifests);
+      }
     }
 
     return ImmutableRewriteManifests.Result.builder()
@@ -247,6 +290,33 @@ private List<ManifestFile> writeUnpartitionedManifests(
     return writeFunc.apply(transformedManifestEntryDF).collectAsList();
   }
 
+  private List<ManifestFile> writeSortedManifests(
+      ManifestContent content, Dataset<Row> manifestEntryDF, int numManifests) {
+
+    // Map the top level partition column names to the column name referenced within the manifest
+    // entry dataframe
+    Column[] actualPartitionColumns =
+        partitionSortColumns.stream()
+            .map(p -> col("data_file.partition." + p))
+            .toArray(Column[]::new);
+
+    // Form a new temporary column to sort/cluster manifests on, based on the custom sort
+    // order provided
+    Dataset<Row> clusteredManifestEntryDF =
+        manifestEntryDF.withColumn(
+            "__clustering_column__", concat_ws("::", actualPartitionColumns));
+
+    return withReusableDS(
+        clusteredManifestEntryDF,
+        df -> {
+          WriteManifests<?> writeFunc = newWriteManifestsFunc(content, df.schema());
+          Column partitionColumn = df.col("__clustering_column__");
+          Dataset<Row> transformedDF =
+              repartitionAndSort(df, partitionColumn, numManifests).drop("__clustering_column__");
+          return writeFunc.apply(transformedDF).collectAsList();
+        });
+  }
+
   private List<ManifestFile> writePartitionedManifests(
       ManifestContent content, Dataset<Row> manifestEntryDF, int numManifests) {
 

spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestRewriteManifestsAction.java

@@ -34,9 +34,12 @@
 import java.io.File;
 import java.io.IOException;
 import java.nio.file.Path;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
+import java.util.Random;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.DataFiles;
@@ -65,6 +68,7 @@
 import org.apache.iceberg.exceptions.CommitStateUnknownException;
 import org.apache.iceberg.hadoop.HadoopTables;
 import org.apache.iceberg.io.OutputFile;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Iterables;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
@@ -73,6 +77,7 @@
 import org.apache.iceberg.spark.SparkWriteOptions;
 import org.apache.iceberg.spark.TestBase;
 import org.apache.iceberg.spark.source.ThreeColumnRecord;
+import org.apache.iceberg.types.Conversions;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.CharSequenceSet;
 import org.apache.iceberg.util.Pair;
@@ -84,6 +89,7 @@
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.junit.jupiter.api.io.TempDir;
+import scala.Tuple2;
 
 @ExtendWith(ParameterizedTestExtension.class)
 public class TestRewriteManifestsAction extends TestBase {
@@ -466,6 +472,164 @@ public void testRewriteLargeManifestsPartitionedTable() throws IOException {
     assertThat(newManifests).hasSizeGreaterThanOrEqualTo(2);
   }
 
+  @TestTemplate
+  public void testRewriteManifestsPartitionedTableWithInvalidSortColumnsThowsException()
+      throws IOException {
+    PartitionSpec spec = PartitionSpec.builderFor(SCHEMA).identity("c1").bucket("c3", 10).build();
+    Map<String, String> options = Maps.newHashMap();
+    options.put(TableProperties.FORMAT_VERSION, String.valueOf(formatVersion));
+    options.put(TableProperties.SNAPSHOT_ID_INHERITANCE_ENABLED, snapshotIdInheritanceEnabled);
+    Table table = TABLES.create(SCHEMA, spec, options, tableLocation);
+
+    SparkActions actions = SparkActions.get();
+
+    //  c2 is not a partition column, cannot use for sorting
+    List<String> badSortKeys1 = ImmutableList.of("c1", "c2");
+    assertThatThrownBy(
+            () ->
+                actions
+                    .rewriteManifests(table)
+                    .rewriteIf(manifest -> true)
+                    .sort(badSortKeys1)
+                    .execute())
+        .isInstanceOf(IllegalArgumentException.class)
+        .message()
+        .contains("Invalid sorting columns");
+
+    //  c3_bucket is the hidden internal partition name, sort() expects the un-hidden partition
+    // column names
+    List<String> badSortKeys2 = ImmutableList.of("c1", "c3_bucket");
+    assertThatThrownBy(
+            () ->
+                actions
+                    .rewriteManifests(table)
+                    .rewriteIf(manifest -> true)
+                    .sort(badSortKeys2)
+                    .execute())
+        .isInstanceOf(IllegalArgumentException.class)
+        .message()
+        .contains("Invalid sorting columns");
+  }
+
+  @TestTemplate
+  public void testRewriteManifestsPartitionedTableWithCustomSorting() throws IOException {
+    Random random = new Random();
+
+    PartitionSpec spec =
+        PartitionSpec.builderFor(SCHEMA).identity("c1").truncate("c2", 3).bucket("c3", 10).build();
+    Map<String, String> options = Maps.newHashMap();
+    options.put(TableProperties.FORMAT_VERSION, String.valueOf(formatVersion));
+    options.put(TableProperties.SNAPSHOT_ID_INHERITANCE_ENABLED, snapshotIdInheritanceEnabled);
+    Table table = TABLES.create(SCHEMA, spec, options, tableLocation);
+
+    List<DataFile> dataFiles = Lists.newArrayList();
+    for (int fileOrdinal = 0; fileOrdinal < 1000; fileOrdinal++) {
+      dataFiles.add(
+          newDataFile(
+              table,
+              TestHelpers.Row.of(
+                  new Object[] {
+                    fileOrdinal, String.valueOf(random.nextInt() * 100), random.nextInt(10)
+                  })));
+    }
+    ManifestFile appendManifest = writeManifest(table, dataFiles);
+    table.newFastAppend().appendManifest(appendManifest).commit();
+
+    List<ManifestFile> manifests = table.currentSnapshot().allManifests(table.io());
+    assertThat(manifests).as("Should have 1 manifests before rewrite").hasSize(1);
+
+    // Capture the c3 partition's lower and upper bounds - used for later test assertions
+    Integer c3PartitionMin =
+        Conversions.fromByteBuffer(
+            Types.IntegerType.get(), manifests.get(0).partitions().get(2).lowerBound());
+    Integer c3PartitionMax =
+        Conversions.fromByteBuffer(
+            Types.IntegerType.get(), manifests.get(0).partitions().get(2).upperBound());
+
+    // Set the target manifest size to a small value to force splitting records into multiple files
+    table
+        .updateProperties()
+        .set(
+            TableProperties.MANIFEST_TARGET_SIZE_BYTES,
+            String.valueOf(manifests.get(0).length() / 2))
+        .commit();
+
+    SparkActions actions = SparkActions.get();
+
+    String stagingLocation = java.nio.file.Files.createTempDirectory(temp, "junit").toString();
+
+    List<String> manifestSortKeys = ImmutableList.of("c3", "c2", "c1");
+    RewriteManifests.Result result =
+        actions
+            .rewriteManifests(table)
+            .rewriteIf(manifest -> true)
+            .sort(manifestSortKeys)
+            .option(RewriteManifestsSparkAction.USE_CACHING, useCaching)
+            .stagingLocation(stagingLocation)
+            .execute();
+
+    table.refresh();
+    List<ManifestFile> newManifests = table.currentSnapshot().allManifests(table.io());
+
+    assertThat(result.rewrittenManifests()).hasSize(1);
+    assertThat(result.addedManifests()).hasSizeGreaterThanOrEqualTo(2);
+    assertManifestsLocation(result.addedManifests(), stagingLocation);
+
+    assertThat(newManifests).hasSizeGreaterThanOrEqualTo(2);
+
+    // Rewritten manifests are clustered by c3_bucket - each should contain only a subset of the
+    // lower and upper bounds
+    // of the partition 'c3'.
+    List<Tuple2<Integer, Integer>> c3Boundaries =
+        newManifests.stream()
+            .map(manifest -> manifest.partitions().get(2))
+            .sorted(
+                Comparator.comparing(
+                    ptn -> Conversions.fromByteBuffer(Types.IntegerType.get(), ptn.lowerBound())))
+            .map(
+                p ->
+                    new Tuple2<Integer, Integer>(
+                        Conversions.fromByteBuffer(Types.IntegerType.get(), p.lowerBound()),
+                        Conversions.fromByteBuffer(Types.IntegerType.get(), p.upperBound())))
+            .collect(Collectors.toList());
+
+    List<Integer> lowers = c3Boundaries.stream().map(t -> t._1).collect(Collectors.toList());
+    List<Integer> uppers = c3Boundaries.stream().map(t -> t._2).collect(Collectors.toList());
+
+    // With custom sorting, this looks like
+    // - manifest 1 -> [lower bound = 0, upper bound = 4]
+    // - manifest 2 -> [lower bound = 4, upper bound = 9]
+    // Without the custom sorting, each manifest tracks the full range of c3 upper/lower bounds.
+    // AKA they look like
+    // - manifest 1 -> [lower bound = 0, upper bound = 9]
+    // - manifest 2 -> [lower bound = 0, upper bound = 9]
+    // So the upper bound of the partitions tracked by the first file should be LEQ the lower bounds
+    // of the second. Etc
+    assertThat(uppers.get(0))
+        .withFailMessage(
+            "Upper bound of first manifest partition should be LEQ lower bound of second")
+        .isLessThanOrEqualTo(lowers.get(1));
+
+    // Each file should contain less than the full c3 partition span
+    c3Boundaries.forEach(
+        boundary -> {
+          assertThat(boundary._2 - boundary._1)
+              .withFailMessage(
+                  "Manifest should contain less than the full range of c3 bucket partitions")
+              .isLessThanOrEqualTo(c3PartitionMax - c3PartitionMin);
+        });
+
+    // c3's Bucket(10) partition means our true lower bound = 0 and true upper bound is 9. The first
+    // manifest should
+    // include the lower bound of 0, and the last should have the upper bound of 9
+    assertThat(lowers.get(0))
+        .withFailMessage("Lower bound of first manifest partition should be 0")
+        .isEqualTo(c3PartitionMin);
+    assertThat(uppers.get(uppers.size() - 1))
+        .withFailMessage("Lower bound of first manifest partition should be 0")
+        .isEqualTo(c3PartitionMax);
+  }
+
   @TestTemplate
   public void testRewriteManifestsWithPredicate() throws IOException {
     PartitionSpec spec = PartitionSpec.builderFor(SCHEMA).identity("c1").truncate("c2", 2).build();