Spark:Add branch support to rewrite_data_files procedure

core/src/main/java/org/apache/iceberg/actions/RewriteDataFilesCommitManager.java

@@ -40,6 +40,7 @@ public class RewriteDataFilesCommitManager {
   private final long startingSnapshotId;
   private final boolean useStartingSequenceNumber;
   private final Map<String, String> snapshotProperties;
+  private final String branch;
 
   // constructor used for testing
   public RewriteDataFilesCommitManager(Table table) {
@@ -60,10 +61,20 @@ public RewriteDataFilesCommitManager(
       long startingSnapshotId,
       boolean useStartingSequenceNumber,
       Map<String, String> snapshotProperties) {
+    this(table, startingSnapshotId, useStartingSequenceNumber, snapshotProperties, null);
+  }
+
+  public RewriteDataFilesCommitManager(
+      Table table,
+      long startingSnapshotId,
+      boolean useStartingSequenceNumber,
+      Map<String, String> snapshotProperties,
+      String branch) {
     this.table = table;
     this.startingSnapshotId = startingSnapshotId;
     this.useStartingSequenceNumber = useStartingSequenceNumber;
     this.snapshotProperties = snapshotProperties;
+    this.branch = branch;
   }
 
   /**
@@ -94,6 +105,10 @@ public void commitFileGroups(Set<RewriteFileGroup> fileGroups) {
 
     snapshotProperties.forEach(rewrite::set);
 
+    if (branch != null) {
+      rewrite.toBranch(branch);
+    }
+
     rewrite.commit();
   }
 

spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestRewriteDataFilesProcedure.java

@@ -54,6 +54,7 @@
 import org.apache.iceberg.spark.SparkCatalogConfig;
 import org.apache.iceberg.spark.SparkTableCache;
 import org.apache.iceberg.spark.SystemFunctionPushDownHelper;
+import org.apache.iceberg.spark.actions.SparkActions;
 import org.apache.iceberg.spark.source.ThreeColumnRecord;
 import org.apache.spark.sql.AnalysisException;
 import org.apache.spark.sql.Dataset;
@@ -1105,4 +1106,209 @@ private List<Object[]> currentData() {
   private List<Object[]> currentData(String table) {
     return rowsToJava(spark.sql("SELECT * FROM " + table + " order by c1, c2, c3").collectAsList());
   }
+
+  @TestTemplate
+  public void testRewriteDataFilesOnBranch() {
+    createTable();
+    insertData(10);
+
+    String branchName = "testBranch";
+    sql("ALTER TABLE %s CREATE BRANCH %s", tableName, branchName);
+
+    List<Object[]> expectedRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+
+    // Get snapshot IDs before rewrite
+    Table table = validationCatalog.loadTable(tableIdent);
+    long mainSnapshotId = table.currentSnapshot().snapshotId();
+    long branchSnapshotId = table.refs().get(branchName).snapshotId();
+
+    // Call rewrite_data_files on the branch
+    List<Object[]> output =
+        sql(
+            "CALL %s.system.rewrite_data_files(table => '%s', branch => '%s')",
+            catalogName, tableName, branchName);
+
+    assertEquals(
+        "Action should rewrite 10 data files and add 1 data file",
+        row(10, 1),
+        Arrays.copyOf(output.get(0), 2));
+
+    // Verify branch data is preserved after compaction
+    List<Object[]> actualRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+
+    // Verify branch snapshot changed
+    table.refresh();
+    assertThat(table.refs().get(branchName).snapshotId())
+        .as("Branch snapshot should be updated when files are rewritten")
+        .isNotEqualTo(branchSnapshotId);
+
+    // Verify main snapshot unchanged
+    assertThat(table.currentSnapshot().snapshotId())
+        .as("Main snapshot should remain unchanged")
+        .isEqualTo(mainSnapshotId);
+  }
+
+  @TestTemplate
+  public void testRewriteDataFilesToNullBranchFails() {
+    createTable();
+    insertData(10);
+
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    assertThatThrownBy(() -> SparkActions.get(spark).rewriteDataFiles(table).toBranch(null))
+        .as("Invalid branch")
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessage("Invalid branch name: null");
+  }
+
+  @TestTemplate
+  public void testRewriteDataFilesOnBranchWithFilter() {
+    createPartitionTable();
+    insertData(10);
+
+    String branchName = "filteredBranch";
+    sql("ALTER TABLE %s CREATE BRANCH %s", tableName, branchName);
+
+    List<Object[]> expectedRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+
+    // Get snapshot IDs before rewrite
+    Table table = validationCatalog.loadTable(tableIdent);
+    long mainSnapshotId = table.currentSnapshot().snapshotId();
+    long branchSnapshotId = table.refs().get(branchName).snapshotId();
+
+    // Call rewrite_data_files on the branch with filter (select only partition c2 = 'bar')
+    List<Object[]> output =
+        sql(
+            "CALL %s.system.rewrite_data_files(table => '%s', branch => '%s', where => 'c2 = \"bar\"')",
+            catalogName, tableName, branchName);
+
+    assertEquals(
+        "Action should rewrite 5 data files from single matching partition"
+            + "(containing c2 = bar) and add 1 data file",
+        row(5, 1),
+        Arrays.copyOf(output.get(0), 2));
+
+    // Verify branch data is preserved after compaction
+    List<Object[]> actualRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+    assertEquals("Data after compaction should not change", expectedRecords, actualRecords);
+
+    // Verify branch snapshot changed after rewrite
+    table.refresh();
+    assertThat(table.refs().get(branchName).snapshotId())
+        .as("Branch snapshot should be updated when files are rewritten")
+        .isNotEqualTo(branchSnapshotId);
+
+    // Verify main snapshot unchanged
+    assertThat(table.currentSnapshot().snapshotId())
+        .as("Main snapshot should remain unchanged")
+        .isEqualTo(mainSnapshotId);
+  }
+
+  @TestTemplate
+  public void testBranchCompactionDoesNotAffectMain() {
+    createTable();
+    // create 10 files under non-partitioned table
+    insertData(10);
+
+    Table table = validationCatalog.loadTable(tableIdent);
+
+    // Create branch from current main state
+    String branchName = "compactionBranch";
+    sql("ALTER TABLE %s CREATE BRANCH %s", tableName, branchName);
+
+    // Add more data to MAIN to make it diverge from branch
+    insertData(tableName, 10);
+
+    // Refresh to get new main snapshot after divergence
+    table.refresh();
+    long mainSnapshotAfterDivergence = table.currentSnapshot().snapshotId();
+    List<Object[]> expectedMainRecords = currentData();
+
+    // Get branch data before compaction
+    List<Object[]> expectedBranchRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+
+    long branchSnapshotBeforeCompaction = table.refs().get(branchName).snapshotId();
+
+    // Verify that branch and main have diverged
+    assertThat(branchSnapshotBeforeCompaction)
+        .as("Branch and main should have different snapshots")
+        .isNotEqualTo(mainSnapshotAfterDivergence);
+
+    List<Object[]> output =
+        sql(
+            "CALL %s.system.rewrite_data_files(table => '%s', branch => '%s')",
+            catalogName, tableName, branchName);
+
+    assertEquals(
+        "Action should rewrite 10 data files and add 1 data file",
+        row(10, 1),
+        Arrays.copyOf(output.get(0), 2));
+
+    table.refresh();
+
+    // Verify main snapshot unchanged
+    assertThat(table.currentSnapshot().snapshotId())
+        .as("Main snapshot ID must remain unchanged after branch compaction")
+        .isEqualTo(mainSnapshotAfterDivergence);
+
+    // Verify main data unchanged
+    List<Object[]> actualMainRecords = currentData();
+    assertEquals(
+        "Main data after compaction should not change", expectedMainRecords, actualMainRecords);
+
+    // Verify branch data unchanged
+    List<Object[]> actualBranchRecords =
+        rowsToJava(
+            spark
+                .sql(
+                    String.format(
+                        "SELECT * FROM %s.branch_%s ORDER BY c1, c2, c3", tableName, branchName))
+                .collectAsList());
+    assertEquals(
+        "Branch data after compaction should not change",
+        expectedBranchRecords,
+        actualBranchRecords);
+
+    // Verify branch snapshot changed
+    long branchSnapshotAfterCompaction = table.refs().get(branchName).snapshotId();
+    assertThat(branchSnapshotAfterCompaction)
+        .as("Branch snapshot must be updated after compaction")
+        .isNotEqualTo(branchSnapshotBeforeCompaction);
+
+    // Verify the new branch snapshot is a child of the previous branch snapshot
+    assertThat(table.snapshot(branchSnapshotAfterCompaction).parentId())
+        .as("New branch snapshot must be a child of the previous branch snapshot")
+        .isEqualTo(branchSnapshotBeforeCompaction);
+  }
 }

spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/actions/RewriteDataFilesSparkAction.java

@@ -31,6 +31,7 @@
 import java.util.stream.Collectors;
 import org.apache.iceberg.DataFile;
 import org.apache.iceberg.FileScanTask;
+import org.apache.iceberg.SnapshotRef;
 import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.Table;
@@ -96,6 +97,7 @@ public class RewriteDataFilesSparkAction
   private boolean removeDanglingDeletes;
   private boolean useStartingSequenceNumber;
   private boolean caseSensitive;
+  private String branch = SnapshotRef.MAIN_BRANCH;
   private BinPackRewriteFilePlanner planner = null;
   private FileRewriteRunner<FileGroupInfo, FileScanTask, DataFile, RewriteFileGroup> runner = null;
 
@@ -157,13 +159,24 @@ public RewriteDataFilesSparkAction filter(Expression expression) {
     return this;
   }
 
+  public RewriteDataFilesSparkAction toBranch(String targetBranch) {
+    Preconditions.checkArgument(targetBranch != null, "Invalid branch name: null");
+    this.branch = targetBranch;
+    return this;
+  }
+
   @Override
   public RewriteDataFiles.Result execute() {
     if (table.currentSnapshot() == null) {
       return EMPTY_RESULT;
     }
 
-    long startingSnapshotId = table.currentSnapshot().snapshotId();
+    Preconditions.checkArgument(
+        table.snapshot(branch) != null,
+        "Cannot rewrite data files for branch %s: branch does not exist",
+        branch);
+
+    long startingSnapshotId = table.snapshot(branch).snapshotId();
 
     init(startingSnapshotId);
 
@@ -230,7 +243,7 @@ private ExecutorService rewriteService() {
   @VisibleForTesting
   RewriteDataFilesCommitManager commitManager(long startingSnapshotId) {
     return new RewriteDataFilesCommitManager(
-        table, startingSnapshotId, useStartingSequenceNumber, commitSummary());
+        table, startingSnapshotId, useStartingSequenceNumber, commitSummary(), branch);
   }
 
   private Builder doExecute(