Core: Preserve concurrent snapshots in replace transactions

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

@@ -21,6 +21,7 @@
 import java.io.Closeable;
 import java.io.IOException;
 import java.util.Map;
+import java.util.function.UnaryOperator;
 import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
@@ -237,12 +238,23 @@ private Transaction newReplaceTableTransaction(boolean orCreate) {
         throw new NoSuchTableException("Table does not exist: %s", identifier);
       }
 
-      TableMetadata metadata;
       tableProperties.putAll(tableOverrideProperties());
+      PartitionSpec replaceSpec = spec;
+      SortOrder replaceSortOrder = sortOrder;
+      String replaceLocation = location;
+      Map<String, String> replaceProperties = Maps.newHashMap(tableProperties);
+      UnaryOperator<TableMetadata> replacement =
+          base ->
+              base.buildReplacement(
+                  schema,
+                  replaceSpec,
+                  replaceSortOrder,
+                  replaceLocation != null ? replaceLocation : base.location(),
+                  replaceProperties);
+
+      TableMetadata metadata;
       if (ops.current() != null) {
-        String baseLocation = location != null ? location : ops.current().location();
-        metadata =
-            ops.current().buildReplacement(schema, spec, sortOrder, baseLocation, tableProperties);
+        metadata = replacement.apply(ops.current());
       } else {
         String baseLocation = location != null ? location : defaultWarehouseLocation(identifier);
         metadata =
@@ -251,10 +263,10 @@ private Transaction newReplaceTableTransaction(boolean orCreate) {
 
       if (orCreate) {
         return Transactions.createOrReplaceTableTransaction(
-            identifier.toString(), ops, metadata, metricsReporter());
+            identifier.toString(), ops, metadata, metricsReporter(), replacement);
       } else {
         return Transactions.replaceTableTransaction(
-            identifier.toString(), ops, metadata, metricsReporter());
+            identifier.toString(), ops, metadata, metricsReporter(), replacement);
       }
     }
 

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

@@ -33,6 +33,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.function.Consumer;
+import java.util.function.UnaryOperator;
 import java.util.stream.Collectors;
 import org.apache.iceberg.encryption.EncryptionManager;
 import org.apache.iceberg.exceptions.CleanableFailure;
@@ -71,14 +72,16 @@ enum TransactionType {
       Sets.newHashSet(); // keep track of files deleted in the most recent commit
   private final Consumer<String> enqueueDelete = deletedFiles::add;
   private final TransactionType type;
+  // rebuilds replace metadata on the refreshed table; null to commit the replacement as built
+  private final UnaryOperator<TableMetadata> replacement;
   private TableMetadata base;
   private TableMetadata current;
   private boolean hasLastOpCommitted;
   private final MetricsReporter reporter;
 
   BaseTransaction(
       String tableName, TableOperations ops, TransactionType type, TableMetadata start) {
-    this(tableName, ops, type, start, LoggingMetricsReporter.instance());
+    this(tableName, ops, type, start, null, LoggingMetricsReporter.instance());
   }
 
   BaseTransaction(
@@ -87,6 +90,16 @@ enum TransactionType {
       TransactionType type,
       TableMetadata start,
       MetricsReporter reporter) {
+    this(tableName, ops, type, start, null, reporter);
+  }
+
+  BaseTransaction(
+      String tableName,
+      TableOperations ops,
+      TransactionType type,
+      TableMetadata start,
+      UnaryOperator<TableMetadata> replacement,
+      MetricsReporter reporter) {
     this.tableName = tableName;
     this.ops = ops;
     this.transactionTable = new TransactionTable();
@@ -95,6 +108,7 @@ enum TransactionType {
     this.updates = Lists.newArrayList();
     this.base = ops.current();
     this.type = type;
+    this.replacement = replacement;
     this.hasLastOpCommitted = true;
     this.reporter = reporter;
   }
@@ -311,20 +325,7 @@ private void commitReplaceTransaction(boolean orCreate) {
           .onlyRetryOn(CommitFailedException.class)
           .run(
               underlyingOps -> {
-                try {
-                  underlyingOps.refresh();
-                } catch (NoSuchTableException e) {
-                  if (!orCreate) {
-                    throw e;
-                  }
-                }
-
-                // because this is a replace table, it will always completely replace the table
-                // metadata. even if it was just updated.
-                if (base != underlyingOps.current()) {
-                  this.base = underlyingOps.current(); // just refreshed
-                }
-
+                refreshReplacement(underlyingOps, orCreate);
                 underlyingOps.commit(base, current);
               });
 
@@ -335,15 +336,54 @@ private void commitReplaceTransaction(boolean orCreate) {
       // the commit failed and no files were committed. clean up each update.
       if (!ops.requireStrictCleanup() || e instanceof CleanableFailure) {
         cleanAllUpdates();
+        deleteUncommittedFiles(deletedFiles);
       }
 
       throw e;
+    }
 
-    } finally {
-      // replace table never needs to retry because the table state is completely replaced. because
-      // retries are not
-      // a concern, it is safe to delete all the deleted files from individual operations
-      deleteUncommittedFiles(deletedFiles);
+    if (!deletedFiles.isEmpty()) {
+      cleanUpAfterCommit();
+    }
+  }
+
+  private void refreshReplacement(TableOperations underlyingOps, boolean orCreate) {
+    try {
+      underlyingOps.refresh();
+    } catch (NoSuchTableException e) {
+      if (!orCreate) {
+        throw e;
+      }
+    }
+
+    if (base == underlyingOps.current()) {
+      return;
+    }
+
+    this.base = underlyingOps.current(); // just refreshed
+    if (replacement == null || base == null) {
+      return;
+    }
+
+    TableMetadata replaced = current;
+    try {
+      TableMetadata rebuilt = replacement.apply(base);
+      // the files this transaction already wrote depend on its schema, spec, sort order and format
+      if (rebuilt.schema().sameSchema(current.schema())
+          && rebuilt.spec().equals(current.spec())
+          && rebuilt.sortOrder().equals(current.sortOrder())
+          && rebuilt.formatVersion() == current.formatVersion()) {
+        this.current = rebuilt;
+        for (PendingUpdate update : updates) {
+          update.commit();
+        }
+      }
+    } catch (RuntimeException e) {
+      LOG.warn(
+          "Cannot rebuild replace of {} on refreshed metadata, concurrent changes will be dropped",
+          tableName,
+          e);
+      this.current = replaced;
     }
   }
 
@@ -384,15 +424,20 @@ private void commitSimpleTransaction() {
     }
 
     // the commit succeeded
+    cleanUpAfterCommit();
+  }
 
+  private void cleanUpAfterCommit() {
     try {
       // clean up the data files that were deleted by each operation. first, get the list of
       // committed manifests to ensure that no committed manifest is deleted.
       // A manifest could be deleted in one successful operation commit, but reused in another
       // successful commit of that operation if the whole transaction is retried.
       // only snapshots added on top of the committed base can reference this transaction's files
       Set<Long> startingSnapshots =
-          base.snapshots().stream().map(Snapshot::snapshotId).collect(Collectors.toSet());
+          base != null
+              ? base.snapshots().stream().map(Snapshot::snapshotId).collect(Collectors.toSet())
+              : Sets.newHashSet();
       Set<Snapshot> newSnapshots = Sets.newHashSet();
       for (Snapshot snapshot : current.snapshots()) {
         if (!startingSnapshots.contains(snapshot.snapshotId())) {

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

@@ -18,6 +18,7 @@
  */
 package org.apache.iceberg;
 
+import java.util.function.UnaryOperator;
 import org.apache.iceberg.BaseTransaction.TransactionType;
 import org.apache.iceberg.metrics.MetricsReporter;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
@@ -36,6 +37,16 @@ public static Transaction createOrReplaceTableTransaction(
         tableName, ops, TransactionType.CREATE_OR_REPLACE_TABLE, start, reporter);
   }
 
+  static Transaction createOrReplaceTableTransaction(
+      String tableName,
+      TableOperations ops,
+      TableMetadata start,
+      MetricsReporter reporter,
+      UnaryOperator<TableMetadata> replacement) {
+    return new BaseTransaction(
+        tableName, ops, TransactionType.CREATE_OR_REPLACE_TABLE, start, replacement, reporter);
+  }
+
   public static Transaction replaceTableTransaction(
       String tableName, TableOperations ops, TableMetadata start) {
     return new BaseTransaction(tableName, ops, TransactionType.REPLACE_TABLE, start);
@@ -46,6 +57,16 @@ public static Transaction replaceTableTransaction(
     return new BaseTransaction(tableName, ops, TransactionType.REPLACE_TABLE, start, reporter);
   }
 
+  static Transaction replaceTableTransaction(
+      String tableName,
+      TableOperations ops,
+      TableMetadata start,
+      MetricsReporter reporter,
+      UnaryOperator<TableMetadata> replacement) {
+    return new BaseTransaction(
+        tableName, ops, TransactionType.REPLACE_TABLE, start, replacement, reporter);
+  }
+
   public static Transaction createTableTransaction(
       String tableName, TableOperations ops, TableMetadata start) {
     Preconditions.checkArgument(

core/src/test/java/org/apache/iceberg/TestReplaceTransaction.java

@@ -33,7 +33,9 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.iceberg.exceptions.CommitFailedException;
 import org.apache.iceberg.exceptions.CommitStateUnknownException;
+import org.apache.iceberg.io.CloseableIterable;
 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.Maps;
 import org.apache.iceberg.relocated.com.google.common.collect.Sets;
 import org.apache.iceberg.transforms.Transform;
@@ -315,6 +317,66 @@ public void testReplaceTransactionConflict() {
     assertThat(listManifestFiles()).containsExactlyElementsOf(manifests);
   }
 
+  @TestTemplate
+  void replaceTransactionConcurrentCommitRetainsHistory() {
+    // TestTables assigns sequential snapshot ids, which collide between the two writers
+    table.updateProperties().set("random-snapshot-ids", "true").commit();
+
+    table.newAppend().appendFile(FILE_A).commit();
+
+    validateSnapshot(null, table.currentSnapshot(), FILE_A);
+    long firstSnapshotId = table.currentSnapshot().snapshotId();
+
+    Transaction replace = TestTables.beginReplace(tableDir, "test", table.schema(), table.spec());
+    replace.newAppend().appendFile(FILE_B).commit();
+
+    table.newAppend().appendFile(FILE_C).commit();
+    long concurrentSnapshotId = table.currentSnapshot().snapshotId();
+
+    replace.commitTransaction();
+
+    table.refresh();
+
+    assertThat(table.currentSnapshot()).isNotNull();
+    validateSnapshot(null, table.currentSnapshot(), FILE_B);
+
+    assertThat(table.snapshot(concurrentSnapshotId))
+        .as("Concurrent writer's snapshot should be preserved in history")
+        .isNotNull();
+    assertThat(table.snapshot(firstSnapshotId))
+        .as("Original snapshot should be preserved in history")
+        .isNotNull();
+    assertThat(table.snapshots()).hasSize(3);
+  }
+
+  @TestTemplate
+  void replaceTransactionFilesMatchTheirSortOrder() throws IOException {
+    table.newAppend().appendFile(FILE_A).commit();
+
+    SortOrder requested = SortOrder.builderFor(table.schema()).asc("data").build();
+    Transaction replace =
+        TestTables.beginReplace(
+            tableDir, "test", table.schema(), table.spec(), requested, ImmutableMap.of());
+    DataFile sortedFile =
+        DataFiles.builder(table.spec())
+            .copy(FILE_B)
+            .withSortOrder(replace.table().sortOrder())
+            .build();
+    replace.newAppend().appendFile(sortedFile).commit();
+
+    table.replaceSortOrder().asc("id").commit();
+
+    replace.commitTransaction();
+
+    table.refresh();
+    try (CloseableIterable<FileScanTask> tasks = table.newScan().planFiles()) {
+      DataFile committed = Iterables.getOnlyElement(tasks).file();
+      assertThat(table.sortOrders().get(committed.sortOrderId()).sameOrder(table.sortOrder()))
+          .as("Data file should reference the replace's sort order")
+          .isTrue();
+    }
+  }
+
   @TestTemplate
   public void testReplaceToCreateAndAppend() throws IOException {
     // this table doesn't exist.
@@ -427,6 +489,31 @@ public void testCreateTransactionWithUnknownState() throws IOException {
     validateSnapshot(null, meta.currentSnapshot(), FILE_A, FILE_B);
   }
 
+  @TestTemplate
+  void createOrReplaceOfMissingTableDeletesUnusedManifestList() {
+    String name = "missing";
+    TestTables.TestTableOperations ops = new TestTables.TestTableOperations(name, tableDir);
+    TableMetadata metadata =
+        TableMetadata.newTableMetadata(
+            SCHEMA,
+            SPEC,
+            SortOrder.unsorted(),
+            tableDir.toURI().toString(),
+            ImmutableMap.of(),
+            formatVersion);
+    Transaction createOrReplace = Transactions.createOrReplaceTableTransaction(name, ops, metadata);
+
+    AppendFiles append = createOrReplace.newFastAppend().appendFile(FILE_A);
+    Snapshot applied = append.apply();
+    append.commit();
+
+    createOrReplace.commitTransaction();
+
+    Snapshot committed = TestTables.readMetadata(name).currentSnapshot();
+    assertThat(new File(applied.manifestListLocation())).doesNotExist();
+    assertThat(new File(committed.manifestListLocation())).exists();
+  }
+
   private static Schema assignFreshIds(Schema schema) {
     AtomicInteger lastColumnId = new AtomicInteger(0);
     return TypeUtil.assignFreshIds(schema, lastColumnId::incrementAndGet);

core/src/test/java/org/apache/iceberg/TestTables.java

@@ -22,6 +22,7 @@
 
 import java.io.File;
 import java.util.Map;
+import java.util.function.UnaryOperator;
 import org.apache.iceberg.exceptions.AlreadyExistsException;
 import org.apache.iceberg.exceptions.CommitFailedException;
 import org.apache.iceberg.exceptions.CommitStateUnknownException;
@@ -32,6 +33,7 @@
 import org.apache.iceberg.io.LocationProvider;
 import org.apache.iceberg.io.OutputFile;
 import org.apache.iceberg.io.SupportsBulkOperations;
+import org.apache.iceberg.metrics.LoggingMetricsReporter;
 import org.apache.iceberg.metrics.MetricsReporter;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
@@ -204,8 +206,11 @@ public static Transaction beginReplace(
     TableMetadata current = ops.current();
     TableMetadata metadata;
     if (current != null) {
-      metadata = current.buildReplacement(schema, spec, sortOrder, current.location(), properties);
-      return Transactions.replaceTableTransaction(name, ops, metadata);
+      UnaryOperator<TableMetadata> replacement =
+          base -> base.buildReplacement(schema, spec, sortOrder, base.location(), properties);
+      metadata = replacement.apply(current);
+      return Transactions.replaceTableTransaction(
+          name, ops, metadata, LoggingMetricsReporter.instance(), replacement);
     } else {
       metadata = newTableMetadata(schema, spec, sortOrder, temp.toURI().toString(), properties);
       return Transactions.createTableTransaction(name, ops, metadata);