diff --git a/api/src/main/java/org/apache/iceberg/RewriteManifests.java b/api/src/main/java/org/apache/iceberg/RewriteManifests.java index 32a9011ad6f6..5f8e0fdbcb4f 100644 --- a/api/src/main/java/org/apache/iceberg/RewriteManifests.java +++ b/api/src/main/java/org/apache/iceberg/RewriteManifests.java @@ -48,6 +48,28 @@ public interface RewriteManifests extends SnapshotUpdate { */ RewriteManifests clusterBy(Function func); + /** + * Reserves the row ID range assigned by replacement manifests. + * + *

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. + * + *

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, diff --git a/core/src/main/java/org/apache/iceberg/BaseRewriteManifests.java b/core/src/main/java/org/apache/iceberg/BaseRewriteManifests.java index e98027ec4a0a..2aed7df43b60 100644 --- a/core/src/main/java/org/apache/iceberg/BaseRewriteManifests.java +++ b/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 private Function clusterByFunc; private Predicate predicate; + private Long reservedRangeStart; + private Long reservedRangeEndExclusive; + private Long expectedSnapshotId; private final SnapshotSummary.Builder summaryBuilder = SnapshotSummary.builder(); @@ -113,12 +117,43 @@ protected Map summary() { @Override public RewriteManifests clusterBy(Function 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 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 apply(TableMetadata base, Snapshot snapshot) { } validateFilesCounts(); + validateAssignedRowIdRange(); Iterable newManifestsWithMetadata = Iterables.transform( @@ -337,6 +373,87 @@ private void appendEntry(ManifestEntry 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))); diff --git a/core/src/main/java/org/apache/iceberg/ManifestFiles.java b/core/src/main/java/org/apache/iceberg/ManifestFiles.java index dae46e5ec49e..e76f870e84da 100644 --- a/core/src/main/java/org/apache/iceberg/ManifestFiles.java +++ b/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 write( writerProperties); } + /** + * Rewrites the live entries in a data manifest with explicit first row IDs. + * + *

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)}. + * + *

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. + * + *

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 specsById, + OutputFile outputFile, + Function 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 specsById, + EncryptedOutputFile outputFile, + Function 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 writer = newWriter(3, spec, outputFile, null, rangeStart); + List assignedRanges = Lists.newArrayList(); + boolean threw = true; + try (ManifestReader reader = read(manifest, io, specsById)) { + for (ManifestEntry 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 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 newWriter( int formatVersion, diff --git a/core/src/main/java/org/apache/iceberg/SnapshotProducer.java b/core/src/main/java/org/apache/iceberg/SnapshotProducer.java index d97c63b61608..20ebfc7a31f4 100644 --- a/core/src/main/java/org/apache/iceberg/SnapshotProducer.java +++ b/core/src/main/java/org/apache/iceberg/SnapshotProducer.java @@ -334,7 +334,7 @@ public Snapshot apply() { Long assignedRows = null; if (base.formatVersion() >= 3) { nextRowId = base.nextRowId(); - assignedRows = writer.nextRowId() - base.nextRowId(); + assignedRows = assignedRows(base, writer.nextRowId()); } Map summary = summary(); @@ -368,6 +368,16 @@ public Snapshot apply() { writer.toManifestListFile().encryptionKeyID()); } + /** + * Returns the number of row IDs reserved by this snapshot. + * + *

Snapshot producers that assign explicit row IDs may override this method when the reserved + * range cannot be inferred from manifest-level row ID inheritance. + */ + protected long assignedRows(TableMetadata tableMetadata, long manifestListNextRowId) { + return Math.subtractExact(manifestListNextRowId, tableMetadata.nextRowId()); + } + private void runValidations(Snapshot parentSnapshot) { validate(base, parentSnapshot); diff --git a/core/src/test/java/org/apache/iceberg/TestRowLineageAssignment.java b/core/src/test/java/org/apache/iceberg/TestRowLineageAssignment.java index 6dab3c0dad84..123ce4d5e031 100644 --- a/core/src/test/java/org/apache/iceberg/TestRowLineageAssignment.java +++ b/core/src/test/java/org/apache/iceberg/TestRowLineageAssignment.java @@ -19,6 +19,7 @@ package org.apache.iceberg; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.fail; import java.io.File; @@ -27,10 +28,15 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.UUID; +import java.util.stream.Collectors; import org.apache.iceberg.data.Record; +import org.apache.iceberg.exceptions.ValidationException; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.InputFile; 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.types.Types; import org.apache.iceberg.types.Types.NestedField; import org.apache.iceberg.util.PartitionUtil; @@ -531,6 +537,331 @@ public void testTableUpgrade(@TempDir File altLocation) { checkDataFileAssignment(upgradeTable, manifests.get(1), (Long) null); } + @Test + public void commitDistributedRowIdManifestRewrite(@TempDir File altLocation) { + BaseTable upgradeTable = + TestTables.create( + altLocation, + "test_upgrade", + SCHEMA, + PartitionSpec.unpartitioned(), + 2, + Map.of("random-snapshot-ids", "true")); + + upgradeTable.newAppend().appendFile(FILE_A).commit(); + upgradeTable.newAppend().appendFile(FILE_B).commit(); + upgradeTable.newRowDelta().addDeletes(FILE_A_DELETES).commit(); + + Map entriesBefore = liveDataEntriesByPath(upgradeTable); + List deleteManifestsBefore = + upgradeTable.currentSnapshot().deleteManifests(upgradeTable.io()).stream() + .map(ManifestFile::path) + .collect(Collectors.toList()); + long currentSnapshotIdBefore = upgradeTable.currentSnapshot().snapshotId(); + Set snapshotIdsBefore = Sets.newHashSet(); + upgradeTable.snapshots().forEach(snapshot -> snapshotIdsBefore.add(snapshot.snapshotId())); + + Map firstRowIds = Map.of(FILE_A.location(), 0L, FILE_B.location(), 1_000L); + long nextRowIdExclusive = 2_000L; + List oldDataManifests = + upgradeTable.currentSnapshot().dataManifests(upgradeTable.io()); + List replacementManifests = + oldDataManifests.stream() + .map( + manifest -> + rewriteManifestWithFirstRowIds( + upgradeTable, manifest, firstRowIds, 0L, nextRowIdExclusive)) + .collect(Collectors.toList()); + + Transaction upgradeTransaction = upgradeTable.newTransaction(); + upgradeTransaction.updateProperties().set(TableProperties.FORMAT_VERSION, "3").commit(); + TableMetadata stagedMetadata = ((BaseTransaction) upgradeTransaction).currentMetadata(); + assertThat(stagedMetadata.formatVersion()).isEqualTo(3); + assertThat(stagedMetadata.nextRowId()).isEqualTo(0L); + RewriteManifests rewrite = + upgradeTransaction.rewriteManifests().reserveRowIdRange(0L, nextRowIdExclusive); + oldDataManifests.forEach(rewrite::deleteManifest); + replacementManifests.forEach(rewrite::addManifest); + rewrite.commit(); + upgradeTransaction.commitTransaction(); + upgradeTable.refresh(); + + Snapshot assignmentSnapshot = upgradeTable.currentSnapshot(); + assertThat(assignmentSnapshot.parentId()).isEqualTo(currentSnapshotIdBefore); + assertThat(assignmentSnapshot.operation()).isEqualTo(DataOperations.REPLACE); + assertThat(assignmentSnapshot.firstRowId()).isEqualTo(0L); + assertThat(assignmentSnapshot.addedRows()).isEqualTo(nextRowIdExclusive); + assertThat(upgradeTable.operations().current().nextRowId()).isEqualTo(nextRowIdExclusive); + + assertThat(snapshotIdsBefore) + .allSatisfy(snapshotId -> assertThat(upgradeTable.snapshot(snapshotId)).isNotNull()); + assertThat( + assignmentSnapshot.deleteManifests(upgradeTable.io()).stream() + .map(ManifestFile::path) + .collect(Collectors.toList())) + .containsExactlyElementsOf(deleteManifestsBefore); + + Map entriesAfter = liveDataEntriesByPath(upgradeTable); + assertThat(entriesAfter.keySet()).containsExactlyInAnyOrderElementsOf(entriesBefore.keySet()); + entriesBefore.forEach( + (path, before) -> { + EntryMetadata after = entriesAfter.get(path); + assertThat(after.status).isEqualTo(ManifestEntry.Status.EXISTING); + assertThat(after.snapshotId).isEqualTo(before.snapshotId); + assertThat(after.dataSequenceNumber).isEqualTo(before.dataSequenceNumber); + assertThat(after.fileSequenceNumber).isEqualTo(before.fileSequenceNumber); + assertThat(after.recordCount).isEqualTo(before.recordCount); + assertThat(after.firstRowId).isEqualTo(firstRowIds.get(path)); + }); + + upgradeTable.newAppend().appendFile(FILE_C).commit(); + assertThat(liveDataEntriesByPath(upgradeTable).get(FILE_C.location()).firstRowId) + .isEqualTo(nextRowIdExclusive); + } + + @Test + public void rejectInvalidDistributedRowIdManifestRewrite(@TempDir File altLocation) { + BaseTable upgradeTable = upgradedTableWithTwoFiles(altLocation); + ManifestFile oldManifest = + Iterables.getOnlyElement(upgradeTable.currentSnapshot().dataManifests(upgradeTable.io())); + + assertThatThrownBy( + () -> + ManifestFiles.rewriteDataManifestWithFirstRowIds( + oldManifest, + upgradeTable.io(), + upgradeTable.specs(), + upgradeTable + .io() + .newOutputFile(new File(altLocation, "overlapping-range.avro").toString()), + file -> file.location().equals(FILE_A.location()) ? 0L : 100L, + 125L, + 225L)) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("outside reserved range"); + + assertThatThrownBy( + () -> + ManifestFiles.rewriteDataManifestWithFirstRowIds( + oldManifest, + upgradeTable.io(), + upgradeTable.specs(), + upgradeTable + .io() + .newOutputFile(new File(altLocation, "overlapping-ids.avro").toString()), + ignored -> 0L, + 0L, + FILE_A.recordCount())) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("Overlapping row ID ranges"); + + assertThatThrownBy(() -> upgradeTable.rewriteManifests().reserveRowIdRange(10L, 9L)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("must not be less than"); + + RewriteManifests selectedRewrite = upgradeTable.rewriteManifests().rewriteIf(ignored -> true); + assertThatThrownBy(() -> selectedRewrite.reserveRowIdRange(0L, 225L)) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("direct manifest replacement"); + } + + @Test + public void rejectPartialLegacyManifestReplacement(@TempDir File altLocation) { + BaseTable upgradeTable = + TestTables.create( + altLocation, + "test_upgrade", + SCHEMA, + PartitionSpec.unpartitioned(), + 2, + Map.of("random-snapshot-ids", "true")); + upgradeTable.newFastAppend().appendFile(FILE_A).commit(); + upgradeTable.newFastAppend().appendFile(FILE_B).commit(); + TestTables.upgrade(altLocation, "test_upgrade", 3); + upgradeTable.refresh(); + + List oldManifests = + upgradeTable.currentSnapshot().dataManifests(upgradeTable.io()); + assertThat(oldManifests).hasSize(2); + ManifestFile replacedManifest = oldManifests.get(0); + ManifestFile replacement = + rewriteManifestWithFirstRowIds( + upgradeTable, + replacedManifest, + Map.of(FILE_A.location(), 0L, FILE_B.location(), 0L), + 0L, + FILE_A.recordCount()); + + RewriteManifests rewrite = upgradeTable.rewriteManifests(); + rewrite.deleteManifest(replacedManifest); + rewrite.addManifest(replacement); + rewrite.reserveRowIdRange(0L, FILE_A.recordCount()); + + assertThatThrownBy(rewrite::commit) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("All data manifests without assigned row IDs must be replaced"); + } + + @Test + public void rejectConcurrentSnapshotDuringDistributedRowIdRewrite(@TempDir File altLocation) { + BaseTable upgradeTable = + TestTables.create( + altLocation, + "test_upgrade", + SCHEMA, + PartitionSpec.unpartitioned(), + 2, + Map.of("random-snapshot-ids", "true")); + upgradeTable.newAppend().appendFile(FILE_A).commit(); + TestTables.upgrade(altLocation, "test_upgrade", 3); + upgradeTable.refresh(); + + ManifestFile oldManifest = + Iterables.getOnlyElement(upgradeTable.currentSnapshot().dataManifests(upgradeTable.io())); + ManifestFile replacement = + rewriteManifestWithFirstRowIds( + upgradeTable, oldManifest, Map.of(FILE_A.location(), 0L), 0L, FILE_A.recordCount()); + RewriteManifests assignment = upgradeTable.rewriteManifests(); + assignment.deleteManifest(oldManifest); + assignment.addManifest(replacement); + assignment.reserveRowIdRange(0L, FILE_A.recordCount()); + upgradeTable.newAppend().appendFile(FILE_B).commit(); + + assertThatThrownBy(assignment::commit) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("table snapshot changed"); + } + + @Test + public void rejectConcurrentNextRowIdChange(@TempDir File altLocation) { + BaseTable upgradeTable = + TestTables.create( + altLocation, + "test_upgrade", + SCHEMA, + PartitionSpec.unpartitioned(), + 2, + Map.of("random-snapshot-ids", "true")); + upgradeTable.newAppend().appendFile(FILE_A).commit(); + TestTables.upgrade(altLocation, "test_upgrade", 3); + upgradeTable.refresh(); + + long mainSnapshotId = upgradeTable.currentSnapshot().snapshotId(); + ManifestFile oldManifest = + Iterables.getOnlyElement(upgradeTable.currentSnapshot().dataManifests(upgradeTable.io())); + ManifestFile replacement = + rewriteManifestWithFirstRowIds( + upgradeTable, oldManifest, Map.of(FILE_A.location(), 0L), 0L, FILE_A.recordCount()); + RewriteManifests assignment = upgradeTable.rewriteManifests(); + assignment.deleteManifest(oldManifest); + assignment.addManifest(replacement); + assignment.reserveRowIdRange(0L, FILE_A.recordCount()); + + upgradeTable.manageSnapshots().createBranch("audit", mainSnapshotId).commit(); + upgradeTable.newAppend().appendFile(FILE_B).toBranch("audit").commit(); + assertThat(upgradeTable.currentSnapshot().snapshotId()).isEqualTo(mainSnapshotId); + + assertThatThrownBy(assignment::commit) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("table next-row-id changed"); + } + + @Test + public void rejectRowIdReservationForFormatV4(@TempDir File altLocation) { + BaseTable upgradeTable = upgradedTableWithTwoFiles(altLocation); + TestTables.upgrade(altLocation, "test_upgrade", 4); + upgradeTable.refresh(); + + ManifestFile oldManifest = + Iterables.getOnlyElement(upgradeTable.currentSnapshot().dataManifests(upgradeTable.io())); + long rangeEnd = Math.addExact(FILE_A.recordCount(), FILE_B.recordCount()); + ManifestFile replacement = + rewriteManifestWithFirstRowIds( + upgradeTable, + oldManifest, + Map.of(FILE_A.location(), 0L, FILE_B.location(), FILE_A.recordCount()), + 0L, + rangeEnd); + RewriteManifests assignment = upgradeTable.rewriteManifests(); + assignment.deleteManifest(oldManifest); + assignment.addManifest(replacement); + assignment.reserveRowIdRange(0L, rangeEnd); + + assertThatThrownBy(assignment::commit) + .isInstanceOf(ValidationException.class) + .hasMessageContaining("format version 3, found 4"); + } + + private BaseTable upgradedTableWithTwoFiles(File altLocation) { + BaseTable upgradeTable = + TestTables.create( + altLocation, + "test_upgrade", + SCHEMA, + PartitionSpec.unpartitioned(), + 2, + Map.of("random-snapshot-ids", "true")); + upgradeTable.newAppend().appendFile(FILE_A).appendFile(FILE_B).commit(); + TestTables.upgrade(altLocation, "test_upgrade", 3); + upgradeTable.refresh(); + return upgradeTable; + } + + private static ManifestFile rewriteManifestWithFirstRowIds( + Table table, + ManifestFile manifest, + Map firstRowIds, + long rangeStart, + long rangeEnd) { + String outputPath = table.location() + "/metadata/row-id-" + UUID.randomUUID() + ".avro"; + return ManifestFiles.rewriteDataManifestWithFirstRowIds( + manifest, + table.io(), + table.specs(), + table.io().newOutputFile(outputPath), + file -> firstRowIds.get(file.location()), + rangeStart, + rangeEnd); + } + + private static Map liveDataEntriesByPath(Table table) { + Map entriesByPath = Maps.newHashMap(); + for (ManifestFile manifest : table.currentSnapshot().dataManifests(table.io())) { + try (ManifestReader reader = + ManifestFiles.read(manifest, table.io(), table.specs())) { + for (ManifestEntry entry : reader.liveEntries()) { + EntryMetadata previous = + entriesByPath.put(entry.file().location(), new EntryMetadata(entry)); + assertThat(previous) + .as("Active data file path must be unique: %s", entry.file().location()) + .isNull(); + } + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + return entriesByPath; + } + + private static class EntryMetadata { + private final ManifestEntry.Status status; + private final Long snapshotId; + private final Long dataSequenceNumber; + private final Long fileSequenceNumber; + private final long recordCount; + private final Long firstRowId; + + private EntryMetadata(ManifestEntry entry) { + this.status = entry.status(); + this.snapshotId = entry.snapshotId(); + this.dataSequenceNumber = entry.dataSequenceNumber(); + this.fileSequenceNumber = entry.fileSequenceNumber(); + this.recordCount = entry.file().recordCount(); + this.firstRowId = entry.file().firstRowId(); + } + } + @Test public void testAssignmentAfterUpgrade(@TempDir File altLocation) { // data manifests: [added(FILE_C)], [existing(FILE_A), deleted(FILE_B)] diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestStructuredStreamingRead3.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestStructuredStreamingRead3.java index 2cf728830bc8..de89aceb09e2 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestStructuredStreamingRead3.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestStructuredStreamingRead3.java @@ -26,8 +26,10 @@ import java.io.IOException; import java.time.Duration; import java.util.Collections; +import java.util.Comparator; import java.util.List; import java.util.Map; +import java.util.UUID; import java.util.concurrent.TimeoutException; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.IntStream; @@ -40,17 +42,22 @@ import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.Files; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.ManifestFiles; import org.apache.iceberg.Parameter; import org.apache.iceberg.ParameterizedTestExtension; import org.apache.iceberg.Parameters; import org.apache.iceberg.RewriteFiles; +import org.apache.iceberg.RewriteManifests; import org.apache.iceberg.Schema; import org.apache.iceberg.Snapshot; import org.apache.iceberg.SnapshotChanges; import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; import org.apache.iceberg.TableOperations; +import org.apache.iceberg.TableProperties; import org.apache.iceberg.TestHelpers; +import org.apache.iceberg.Transaction; import org.apache.iceberg.data.FileHelpers; import org.apache.iceberg.data.GenericRecord; import org.apache.iceberg.data.Record; @@ -64,6 +71,7 @@ import org.apache.iceberg.spark.SparkCatalogConfig; import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkReadOptions; +import org.apache.iceberg.util.SnapshotUtil; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.api.java.function.VoidFunction2; import org.apache.spark.sql.Dataset; @@ -760,6 +768,104 @@ public void testResumingStreamReadFromCheckpoint() throws Exception { } } + @TestTemplate + public void resumeCheckpointAfterV2ToV3RowIdAssignment() throws Exception { + File writerCheckpoint = temp.resolve("upgrade-checkpoint").toFile(); + File output = temp.resolve("upgrade-output").toFile(); + List checkpointed = Lists.newArrayList(new SimpleRecord(1, "checkpointed")); + List pendingBeforeUpgrade = + Lists.newArrayList(new SimpleRecord(2, "pending-before-upgrade")); + List afterUpgrade = Lists.newArrayList(new SimpleRecord(3, "after-upgrade")); + + appendData(checkpointed); + table.refresh(); + long checkpointSnapshotId = table.currentSnapshot().snapshotId(); + + DataStreamWriter querySource = + spark + .readStream() + .format("iceberg") + .load(tableName) + .writeStream() + .option("checkpointLocation", writerCheckpoint.toString()) + .format("parquet") + .queryName("v2_v3_upgrade_checkpoint_test") + .option("path", output.getPath()); + + StreamingQuery initialQuery = querySource.start(); + initialQuery.processAllAvailable(); + initialQuery.stop(); + + // This append is deliberately newer than the checkpoint. A resumed stream must still process + // it after walking through the v3 manifest-rewrite snapshot. + appendData(pendingBeforeUpgrade); + table.refresh(); + + List activeFiles = Lists.newArrayList(); + try (CloseableIterable tasks = table.newScan().planFiles()) { + tasks.forEach(task -> activeFiles.add(task.file())); + } + activeFiles.sort(Comparator.comparing(DataFile::location)); + + Map firstRowIds = Maps.newHashMap(); + long nextRowId = 0L; + for (DataFile file : activeFiles) { + firstRowIds.put(file.location(), nextRowId); + nextRowId = Math.addExact(nextRowId, file.recordCount()); + } + long nextRowIdExclusive = nextRowId; + + List oldDataManifests = table.currentSnapshot().dataManifests(table.io()); + List replacementManifests = + oldDataManifests.stream() + .map( + manifest -> + rewriteManifestWithFirstRowIds(manifest, firstRowIds, 0L, nextRowIdExclusive)) + .collect(java.util.stream.Collectors.toList()); + + Transaction upgrade = table.newTransaction(); + upgrade.updateProperties().set(TableProperties.FORMAT_VERSION, "3").commit(); + RewriteManifests rewrite = upgrade.rewriteManifests().reserveRowIdRange(0L, nextRowIdExclusive); + oldDataManifests.forEach(rewrite::deleteManifest); + replacementManifests.forEach(rewrite::addManifest); + rewrite.commit(); + upgrade.commitTransaction(); + table.refresh(); + + assertThat(table.snapshot(checkpointSnapshotId)).isNotNull(); + assertThat(SnapshotUtil.isAncestorOf(table, checkpointSnapshotId)).isTrue(); + assertThat(table.currentSnapshot().operation()).isEqualTo(DataOperations.REPLACE); + + appendData(afterUpgrade); + + StreamingQuery resumedQuery = querySource.start(); + resumedQuery.processAllAvailable(); + resumedQuery.stop(); + + StreamingQuery restartedQuery = querySource.start(); + restartedQuery.processAllAvailable(); + restartedQuery.stop(); + + List actual = + spark.read().load(output.getPath()).as(Encoders.bean(SimpleRecord.class)).collectAsList(); + assertThat(actual) + .containsExactlyInAnyOrderElementsOf( + Iterables.concat(checkpointed, pendingBeforeUpgrade, afterUpgrade)); + } + + private ManifestFile rewriteManifestWithFirstRowIds( + ManifestFile manifest, Map firstRowIds, long rangeStart, long rangeEnd) { + String outputPath = table.location() + "/metadata/row-id-" + UUID.randomUUID() + ".avro"; + return ManifestFiles.rewriteDataManifestWithFirstRowIds( + manifest, + table.io(), + table.specs(), + table.io().newOutputFile(outputPath), + file -> firstRowIds.get(file.location()), + rangeStart, + rangeEnd); + } + @TestTemplate public void testFailReadingCheckpointInvalidSnapshot() throws IOException, TimeoutException { File writerCheckpointFolder = temp.resolve("writer-checkpoint-folder").toFile();