Core: Support explicit row ID assignment in manifest rewrites

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

@@ -48,6 +48,28 @@ public interface RewriteManifests extends SnapshotUpdate<RewriteManifests> {
    */
   RewriteManifests clusterBy(Function<DataFile, Object> func);
 
+  /**
+   * Reserves the row ID range assigned by replacement manifests.
+   *
+   * <p>This is a low-level operation for distributed manifest rewrites during an upgrade to format
+   * version 3. All current data manifests without assigned row IDs must be replaced. Added
+   * manifests must contain the same live data files as the deleted manifests and preserve each
+   * manifest entry's snapshot ID, data sequence number, and file sequence number.
+   *
+   * <p>As with direct manifest replacement through {@link #addManifest(ManifestFile)}, the caller
+   * is responsible for the contents of the supplied manifests. In particular, the committing
+   * process does not load all manifest entries and does not validate that row ID ranges are
+   * globally disjoint. Every data file's row ID range must be contained in {@code
+   * [expectedFirstRowId, nextRowIdExclusive)} and must not overlap another data file's range.
+   *
+   * @param expectedFirstRowId expected table {@code next-row-id} before this operation
+   * @param nextRowIdExclusive first row ID available after this operation
+   * @return this for method chaining
+   */
+  default RewriteManifests reserveRowIdRange(long expectedFirstRowId, long nextRowIdExclusive) {
+    throw new UnsupportedOperationException("Explicit row ID ranges are not supported");
+  }
+
   /**
    * Determines which existing {@link ManifestFile} for the table should be rewritten. Manifests
    * that do not match the predicate are kept as-is. If this is not called and no predicate is set,

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

@@ -26,6 +26,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
@@ -65,6 +66,9 @@ public class BaseRewriteManifests extends SnapshotProducer<RewriteManifests>
 
   private Function<DataFile, Object> clusterByFunc;
   private Predicate<ManifestFile> predicate;
+  private Long reservedRangeStart;
+  private Long reservedRangeEndExclusive;
+  private Long expectedSnapshotId;
 
   private final SnapshotSummary.Builder summaryBuilder = SnapshotSummary.builder();
 
@@ -113,12 +117,43 @@ protected Map<String, String> summary() {
 
   @Override
   public RewriteManifests clusterBy(Function<DataFile, Object> func) {
+    Preconditions.checkState(
+        !hasAssignedRowIdRange(),
+        "Manifest clustering cannot be combined with an assigned row ID range");
     this.clusterByFunc = func;
     return this;
   }
 
+  @Override
+  public RewriteManifests reserveRowIdRange(long expectedFirstRowId, long nextRowIdExclusive) {
+    Preconditions.checkArgument(
+        expectedFirstRowId >= 0,
+        "Expected first row ID must be non-negative: %s",
+        expectedFirstRowId);
+    Preconditions.checkArgument(
+        nextRowIdExclusive >= expectedFirstRowId,
+        "Next row ID %s must not be less than expected first row ID %s",
+        nextRowIdExclusive,
+        expectedFirstRowId);
+    Preconditions.checkState(
+        !hasAssignedRowIdRange(), "Assigned row ID range is already configured");
+    Preconditions.checkState(
+        predicate == null && clusterByFunc == null,
+        "Assigned row ID ranges can only be used with direct manifest replacement");
+
+    this.reservedRangeStart = expectedFirstRowId;
+    this.reservedRangeEndExclusive = nextRowIdExclusive;
+    Snapshot currentSnapshot = ops().current().currentSnapshot();
+    this.expectedSnapshotId = currentSnapshot != null ? currentSnapshot.snapshotId() : null;
+
+    return this;
+  }
+
   @Override
   public RewriteManifests rewriteIf(Predicate<ManifestFile> pred) {
+    Preconditions.checkState(
+        !hasAssignedRowIdRange(),
+        "Manifest selection cannot be combined with an assigned row ID range");
     this.predicate = pred;
     return this;
   }
@@ -180,6 +215,7 @@ public List<ManifestFile> apply(TableMetadata base, Snapshot snapshot) {
     }
 
     validateFilesCounts();
+    validateAssignedRowIdRange();
 
     Iterable<ManifestFile> newManifestsWithMetadata =
         Iterables.transform(
@@ -337,6 +373,87 @@ private void appendEntry(ManifestEntry<DataFile> entry, Object key, int partitio
     entryCount.incrementAndGet();
   }
 
+  private void validateAssignedRowIdRange() {
+    if (!hasAssignedRowIdRange()) {
+      return;
+    }
+
+    ValidationException.check(
+        !deletedManifests.isEmpty()
+            && (!addedManifests.isEmpty() || !rewrittenAddedManifests.isEmpty()),
+        "Assigned row ID ranges require direct manifest replacement");
+
+    deletedManifests.forEach(
+        manifest ->
+            ValidationException.check(
+                manifest.content() == ManifestContent.DATA,
+                "Cannot assign row IDs by replacing a delete manifest: %s",
+                manifest.path()));
+    Iterables.concat(addedManifests, rewrittenAddedManifests)
+        .forEach(
+            manifest -> {
+              ValidationException.check(
+                  manifest.content() == ManifestContent.DATA,
+                  "Cannot assign row IDs with a replacement delete manifest: %s",
+                  manifest.path());
+              ValidationException.check(
+                  manifest.firstRowId() != null,
+                  "Replacement data manifest must have first row ID: %s",
+                  manifest.path());
+              ValidationException.check(
+                  manifest.firstRowId() >= reservedRangeStart
+                      && manifest.firstRowId() < reservedRangeEndExclusive,
+                  "Replacement data manifest first row ID %s is outside reserved range [%s, %s]: %s",
+                  manifest.firstRowId(),
+                  reservedRangeStart,
+                  reservedRangeEndExclusive,
+                  manifest.path());
+            });
+
+    keptManifests.forEach(
+        manifest ->
+            ValidationException.check(
+                manifest.content() != ManifestContent.DATA || manifest.firstRowId() != null,
+                "All data manifests without assigned row IDs must be replaced: %s",
+                manifest.path()));
+  }
+
+  @Override
+  protected void validate(TableMetadata currentMetadata, Snapshot snapshot) {
+    if (!hasAssignedRowIdRange()) {
+      return;
+    }
+
+    ValidationException.check(
+        currentMetadata.formatVersion() == 3,
+        "Can only assign row IDs to a table with format version 3, found %s",
+        currentMetadata.formatVersion());
+    Long currentSnapshotId = snapshot != null ? snapshot.snapshotId() : null;
+    ValidationException.check(
+        Objects.equals(expectedSnapshotId, currentSnapshotId),
+        "Cannot assign first row IDs after the table snapshot changed: expected %s, found %s",
+        expectedSnapshotId,
+        currentSnapshotId);
+    ValidationException.check(
+        currentMetadata.nextRowId() == reservedRangeStart,
+        "Cannot assign first row IDs after table next-row-id changed: expected %s, found %s",
+        reservedRangeStart,
+        currentMetadata.nextRowId());
+  }
+
+  @Override
+  protected long assignedRows(TableMetadata base, long manifestListNextRowId) {
+    if (!hasAssignedRowIdRange()) {
+      return super.assignedRows(base, manifestListNextRowId);
+    }
+
+    return Math.subtractExact(reservedRangeEndExclusive, reservedRangeStart);
+  }
+
+  private boolean hasAssignedRowIdRange() {
+    return reservedRangeStart != null;
+  }
+
   private WriterWrapper getWriter(Object key, int partitionSpecId) {
     return writers.computeIfAbsent(
         Pair.of(key, partitionSpecId), k -> new WriterWrapper(specsById.get(partitionSpecId)));

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

@@ -24,6 +24,7 @@
 import java.math.RoundingMode;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ExecutorService;
@@ -35,6 +36,7 @@
 import org.apache.iceberg.encryption.EncryptedFiles;
 import org.apache.iceberg.encryption.EncryptedOutputFile;
 import org.apache.iceberg.exceptions.RuntimeIOException;
+import org.apache.iceberg.exceptions.ValidationException;
 import org.apache.iceberg.io.CloseableIterable;
 import org.apache.iceberg.io.ContentCache;
 import org.apache.iceberg.io.FileIO;
@@ -295,6 +297,173 @@ public static ManifestWriter<DataFile> write(
         writerProperties);
   }
 
+  /**
+   * Rewrites the live entries in a data manifest with explicit first row IDs.
+   *
+   * <p>The replacement contains only {@link ManifestEntry.Status#EXISTING EXISTING} entries and
+   * preserves each entry's snapshot ID, data sequence number, and file sequence number. The caller
+   * can therefore build replacement manifests in parallel and commit them with {@link
+   * RewriteManifests#addManifest(ManifestFile)}.
+   *
+   * <p>This method validates that assigned row ID ranges do not overlap within the replacement
+   * manifest. The caller must validate that ranges are globally disjoint across all replacement
+   * manifests.
+   *
+   * <p>The caller owns the output file until it is successfully committed. The output file may
+   * exist when this method throws and must be cleaned up by the caller.
+   *
+   * @param manifest data manifest to rewrite
+   * @param io file IO used to read the manifest
+   * @param specsById table partition specs by ID
+   * @param outputFile output file for the replacement manifest
+   * @param firstRowIdForFile function that returns the first row ID for each live data file
+   * @param rangeStart inclusive start of the reserved row ID range
+   * @param rangeEnd exclusive end of the reserved row ID range
+   * @return the replacement manifest
+   */
+  public static ManifestFile rewriteDataManifestWithFirstRowIds(
+      ManifestFile manifest,
+      FileIO io,
+      Map<Integer, PartitionSpec> specsById,
+      OutputFile outputFile,
+      Function<DataFile, Long> firstRowIdForFile,
+      long rangeStart,
+      long rangeEnd) {
+    return rewriteDataManifestWithFirstRowIds(
+        manifest,
+        io,
+        specsById,
+        EncryptedFiles.plainAsEncryptedOutput(outputFile),
+        firstRowIdForFile,
+        rangeStart,
+        rangeEnd);
+  }
+
+  /**
+   * Rewrites the live entries in a data manifest with explicit first row IDs.
+   *
+   * @see #rewriteDataManifestWithFirstRowIds(ManifestFile, FileIO, Map, OutputFile, Function, long,
+   *     long)
+   */
+  public static ManifestFile rewriteDataManifestWithFirstRowIds(
+      ManifestFile manifest,
+      FileIO io,
+      Map<Integer, PartitionSpec> specsById,
+      EncryptedOutputFile outputFile,
+      Function<DataFile, Long> firstRowIdForFile,
+      long rangeStart,
+      long rangeEnd) {
+    Preconditions.checkArgument(
+        manifest.content() == ManifestContent.DATA,
+        "Cannot assign row IDs in a delete manifest: %s",
+        manifest.path());
+    Preconditions.checkArgument(firstRowIdForFile != null, "First row ID function cannot be null");
+    Preconditions.checkArgument(
+        rangeStart >= 0, "Row ID range start must be non-negative: %s", rangeStart);
+    Preconditions.checkArgument(
+        rangeEnd >= rangeStart,
+        "Row ID range end %s must not be less than range start %s",
+        rangeEnd,
+        rangeStart);
+
+    PartitionSpec spec =
+        Preconditions.checkNotNull(
+            specsById.get(manifest.partitionSpecId()),
+            "Cannot find partition spec: %s",
+            manifest.partitionSpecId());
+    ManifestWriter<DataFile> writer = newWriter(3, spec, outputFile, null, rangeStart);
+    List<RowIdRange> assignedRanges = Lists.newArrayList();
+    boolean threw = true;
+    try (ManifestReader<DataFile> reader = read(manifest, io, specsById)) {
+      for (ManifestEntry<DataFile> entry : reader.liveEntries()) {
+        DataFile file = entry.file();
+        ValidationException.check(
+            file.firstRowId() == null, "Data file already has a first row ID: %s", file.location());
+        Long assignedFirstRowId = firstRowIdForFile.apply(file);
+        ValidationException.check(
+            assignedFirstRowId != null, "No first row ID for data file: %s", file.location());
+
+        long lastRowIdExclusive;
+        try {
+          lastRowIdExclusive = Math.addExact(assignedFirstRowId, file.recordCount());
+        } catch (ArithmeticException e) {
+          throw new ValidationException(
+              e, "Row ID range overflows for data file: %s", file.location());
+        }
+        ValidationException.check(
+            assignedFirstRowId >= rangeStart && lastRowIdExclusive <= rangeEnd,
+            "Row ID range [%s, %s) for data file %s is outside reserved range [%s, %s)",
+            assignedFirstRowId,
+            lastRowIdExclusive,
+            file.location(),
+            rangeStart,
+            rangeEnd);
+        assignedRanges.add(new RowIdRange(assignedFirstRowId, lastRowIdExclusive, file.location()));
+
+        DataFile withFirstRowId =
+            DataFiles.builder(spec).copy(file).withFirstRowId(assignedFirstRowId).build();
+        writer.existing(
+            withFirstRowId,
+            Preconditions.checkNotNull(
+                entry.snapshotId(), "Missing snapshot ID for data file: %s", file.location()),
+            Preconditions.checkNotNull(
+                entry.dataSequenceNumber(),
+                "Missing data sequence number for data file: %s",
+                file.location()),
+            entry.fileSequenceNumber());
+      }
+
+      validateNonOverlappingRowIdRanges(assignedRanges);
+      threw = false;
+    } catch (IOException e) {
+      throw new RuntimeIOException(e, "Failed to close manifest: %s", manifest.path());
+    } finally {
+      try {
+        writer.close();
+      } catch (IOException e) {
+        if (!threw) {
+          throw new RuntimeIOException(e, "Failed to close manifest: %s", outputFile);
+        }
+      }
+    }
+
+    return writer.toManifestFile();
+  }
+
+  private static void validateNonOverlappingRowIdRanges(List<RowIdRange> ranges) {
+    ranges.sort(Comparator.comparingLong(range -> range.firstRowId));
+
+    RowIdRange previous = null;
+    for (RowIdRange current : ranges) {
+      if (previous != null && current.firstRowId < previous.lastRowIdExclusive) {
+        throw new ValidationException(
+            "Overlapping row ID ranges for data files %s [%s, %s) and %s [%s, %s)",
+            previous.filePath,
+            previous.firstRowId,
+            previous.lastRowIdExclusive,
+            current.filePath,
+            current.firstRowId,
+            current.lastRowIdExclusive);
+      }
+
+      if (previous == null || current.lastRowIdExclusive > previous.lastRowIdExclusive) {
+        previous = current;
+      }
+    }
+  }
+
+  private static class RowIdRange {
+    private final long firstRowId;
+    private final long lastRowIdExclusive;
+    private final String filePath;
+
+    private RowIdRange(long firstRowId, long lastRowIdExclusive, String filePath) {
+      this.firstRowId = firstRowId;
+      this.lastRowIdExclusive = lastRowIdExclusive;
+      this.filePath = filePath;
+    }
+  }
+
   @VisibleForTesting
   static ManifestWriter<DataFile> newWriter(
       int formatVersion,