stevenzwu commented on code in PR #18853:
URL: https://github.com/apache/iceberg/pull/18853#discussion_r4540268343
##########
core/src/main/java/org/apache/iceberg/V4Metadata.java:
##########
@@ -541,4 +546,607 @@ public Long pos() {
return null;
}
}
+
+ /** Base wrapper used to write a {@link TrackedFile} using the v4 tracked
file schema. */
+ private abstract static class TrackedFileWrapper implements StructLike {
+ private static final int SIZE =
+ TrackedFile.schema(Types.StructType.of(),
Types.StructType.of()).columns().size();
+
+ private final Map<Integer, PartitionSpec> specsById;
+ private final Map<Integer, StructProjection> partitionProjections;
+ private final ContentStatsWrapper contentStats;
+ private final DeletionVectorWrapper deletionVector = new
DeletionVectorWrapper();
+ private final List<ColumnFileWrapper> columnFileWrappers =
Lists.newArrayList();
+ private final TrackingWrapper tracking;
+ private final ManifestInfoWrapper manifestInfo;
+
+ private TrackedFile wrapped = null;
+ private Integer specId = null;
+
+ private TrackedFileWrapper(
+ Map<Integer, PartitionSpec> specsById,
+ Types.StructType partitionType,
+ Types.StructType contentStatsType,
+ TrackingWrapper tracking,
+ ManifestInfoWrapper manifestInfo) {
+ validatePartitionType(specsById.values(), partitionType);
+ this.specsById = specsById;
+ this.partitionProjections = partitionProjections(specsById.values(),
partitionType);
+ this.contentStats = new ContentStatsWrapper(contentStatsType);
+ this.tracking = tracking;
+ this.manifestInfo = manifestInfo;
+ }
+
+ protected TrackingWrapper tracking() {
+ return tracking;
+ }
+
+ protected void wrapInternal(TrackedFile file, Tracking newTracking) {
+ Preconditions.checkArgument(file != null, "Invalid file: null");
+ Preconditions.checkArgument(newTracking != null, "Invalid tracking:
null");
+
+ FileContent content = file.contentType();
+ if (content == FileContent.DATA) {
+ Preconditions.checkArgument(
+ file.manifestInfo() == null,
+ "Invalid data file, manifest info must be null: %s",
+ file.location());
+ } else if (content == FileContent.DATA_MANIFEST || content ==
FileContent.DELETE_MANIFEST) {
+ Preconditions.checkArgument(
+ file.manifestInfo() != null,
+ "Invalid manifest, missing manifest info: %s",
+ file.location());
+ Preconditions.checkArgument(
+ file.deletionVector() == null,
+ "Invalid manifest, deletion vector must be null: %s",
+ file.location());
+ } else {
+ throw new IllegalArgumentException(
+ String.format("Cannot write %s to a v4 manifest: %s", content,
file.location()));
+ }
+
+ Integer fileSpecId = file.specId();
+ PartitionSpec spec = fileSpecId != null ? specsById.get(fileSpecId) :
null;
+ Preconditions.checkArgument(
+ fileSpecId == null || spec != null,
+ "Cannot find partition spec for spec ID: %s",
+ fileSpecId);
+
+ this.wrapped = file;
+ // v4 metadata stores a null spec ID and partition for unpartitioned
files
+ this.specId = spec != null && !spec.isUnpartitioned() ? fileSpecId :
null;
+ }
+
+ @Override
+ public int size() {
+ return SIZE;
+ }
+
+ @Override
+ public <T> void set(int pos, T value) {
+ throw new UnsupportedOperationException("Cannot modify
TrackedFileWrapper wrapper via set");
+ }
+
+ @Override
+ public <T> T get(int pos, Class<T> javaClass) {
+ return javaClass.cast(get(pos));
+ }
+
+ private Object get(int pos) {
+ return switch (pos) {
+ case 0 -> tracking;
+ case 1 -> wrapped.contentType().id();
+ case 2 -> wrapped.location();
+ case 3 -> wrapped.fileFormat().toString();
+ case 4 -> wrapped.recordCount();
+ case 5 -> wrapped.fileSizeInBytes();
+ case 6 -> specId;
+ case 7 -> unionPartition();
+ case 8 -> contentStats.wrap(wrapped.contentStats());
+ case 9 -> wrapped.tightStats();
+ case 10 -> wrapped.sortOrderId();
+ case 11 -> deletionVector.wrap(wrapped.deletionVector());
+ case 12 ->
+ wrapped.manifestInfo() != null ?
manifestInfo.wrap(wrapped.manifestInfo()) : null;
+ case 13 -> wrapped.keyMetadata();
+ case 14 -> wrapped.splitOffsets();
+ case 15 -> wrapColumnFiles(wrapped.columnFiles());
+ default -> throw new UnsupportedOperationException("Unknown field
ordinal: " + pos);
+ };
+ }
+
+ private List<ColumnFileWrapper> wrapColumnFiles(List<ColumnFile>
columnFiles) {
+ if (columnFiles == null) {
+ return null;
+ }
+
+ for (int pos = columnFileWrappers.size(); pos < columnFiles.size(); pos
+= 1) {
+ columnFileWrappers.add(new ColumnFileWrapper());
+ }
+
+ for (int pos = 0; pos < columnFiles.size(); pos += 1) {
+ columnFileWrappers.get(pos).wrap(columnFiles.get(pos));
+ }
+
+ return columnFileWrappers.subList(0, columnFiles.size());
+ }
+
+ private StructLike unionPartition() {
+ if (specId == null) {
+ return null;
+ }
+
+ StructProjection projection = partitionProjections.get(specId);
+ return projection != null ? projection.wrap(wrapped.partition()) :
wrapped.partition();
+ }
+
+ private static void validatePartitionType(
+ Iterable<PartitionSpec> specs, Types.StructType unionType) {
+ for (PartitionSpec spec : specs) {
+ for (Types.NestedField field : spec.partitionType().fields()) {
+ Preconditions.checkArgument(
+ unionType.field(field.fieldId()) != null,
+ "Invalid union partition type, missing field %s from spec %s:
%s",
+ field.fieldId(),
+ spec.specId(),
+ unionType);
+ }
+ }
+ }
+
+ private static Map<Integer, StructProjection> partitionProjections(
+ Iterable<PartitionSpec> specs, Types.StructType unionType) {
+ Map<Integer, StructProjection> projections = Maps.newHashMap();
+ for (PartitionSpec spec : specs) {
+ Types.StructType partitionType = spec.partitionType();
+ StructProjection projection =
+ partitionType.equals(unionType)
+ ? null
+ : StructProjection.createAllowMissing(partitionType,
unionType);
+ projections.put(spec.specId(), projection);
+ }
+
+ return projections;
+ }
+ }
+
+ /** Wrapper used to write a data file to a v4 leaf manifest. */
+ static class LeafTrackedFileWrapper extends TrackedFileWrapper {
+ LeafTrackedFileWrapper(
+ Long snapshotId,
+ Map<Integer, PartitionSpec> specsById,
+ Types.StructType partitionType,
+ Types.StructType contentStatsType) {
+ super(
+ specsById,
+ partitionType,
+ contentStatsType,
+ new TrackingWrapper(snapshotId, null /* sequence numbers are
inherited */),
+ null /* leaf manifests only contain data files */);
+ }
+
+ LeafTrackedFileWrapper wrap(TrackedFile file, Tracking newTracking) {
+ wrapInternal(file, newTracking);
Review Comment:
I tried a few structure change to split the validation properly. this is the
version that I settled in.
The leaf content check runs after `wrapFile` so that a null `file` is
rejected by validation inside `wrapFile`, instead of failing with an NPE on
`file.contentType()`.
##########
core/src/main/java/org/apache/iceberg/V4Metadata.java:
##########
@@ -541,4 +546,607 @@ public Long pos() {
return null;
}
}
+
+ /** Base wrapper used to write a {@link TrackedFile} using the v4 tracked
file schema. */
+ private abstract static class TrackedFileWrapper implements StructLike {
+ private static final int SIZE =
+ TrackedFile.schema(Types.StructType.of(),
Types.StructType.of()).columns().size();
+
+ private final Map<Integer, PartitionSpec> specsById;
+ private final Map<Integer, StructProjection> partitionProjections;
+ private final ContentStatsWrapper contentStats;
+ private final DeletionVectorWrapper deletionVector = new
DeletionVectorWrapper();
+ private final List<ColumnFileWrapper> columnFileWrappers =
Lists.newArrayList();
+ private final TrackingWrapper tracking;
+ private final ManifestInfoWrapper manifestInfo;
+
+ private TrackedFile wrapped = null;
+ private Integer specId = null;
+
+ private TrackedFileWrapper(
+ Map<Integer, PartitionSpec> specsById,
+ Types.StructType partitionType,
+ Types.StructType contentStatsType,
+ TrackingWrapper tracking,
+ ManifestInfoWrapper manifestInfo) {
+ validatePartitionType(specsById.values(), partitionType);
+ this.specsById = specsById;
+ this.partitionProjections = partitionProjections(specsById.values(),
partitionType);
+ this.contentStats = new ContentStatsWrapper(contentStatsType);
+ this.tracking = tracking;
+ this.manifestInfo = manifestInfo;
+ }
+
+ protected TrackingWrapper tracking() {
+ return tracking;
+ }
+
+ protected void wrapInternal(TrackedFile file, Tracking newTracking) {
+ Preconditions.checkArgument(file != null, "Invalid file: null");
+ Preconditions.checkArgument(newTracking != null, "Invalid tracking:
null");
+
+ FileContent content = file.contentType();
+ if (content == FileContent.DATA) {
+ Preconditions.checkArgument(
+ file.manifestInfo() == null,
+ "Invalid data file, manifest info must be null: %s",
+ file.location());
+ } else if (content == FileContent.DATA_MANIFEST || content ==
FileContent.DELETE_MANIFEST) {
+ Preconditions.checkArgument(
+ file.manifestInfo() != null,
+ "Invalid manifest, missing manifest info: %s",
+ file.location());
+ Preconditions.checkArgument(
+ file.deletionVector() == null,
+ "Invalid manifest, deletion vector must be null: %s",
+ file.location());
+ } else {
+ throw new IllegalArgumentException(
+ String.format("Cannot write %s to a v4 manifest: %s", content,
file.location()));
+ }
+
+ Integer fileSpecId = file.specId();
+ PartitionSpec spec = fileSpecId != null ? specsById.get(fileSpecId) :
null;
+ Preconditions.checkArgument(
+ fileSpecId == null || spec != null,
+ "Cannot find partition spec for spec ID: %s",
+ fileSpecId);
+
+ this.wrapped = file;
+ // v4 metadata stores a null spec ID and partition for unpartitioned
files
+ this.specId = spec != null && !spec.isUnpartitioned() ? fileSpecId :
null;
+ }
+
+ @Override
+ public int size() {
+ return SIZE;
+ }
+
+ @Override
+ public <T> void set(int pos, T value) {
+ throw new UnsupportedOperationException("Cannot modify
TrackedFileWrapper wrapper via set");
+ }
+
+ @Override
+ public <T> T get(int pos, Class<T> javaClass) {
+ return javaClass.cast(get(pos));
+ }
+
+ private Object get(int pos) {
+ return switch (pos) {
+ case 0 -> tracking;
+ case 1 -> wrapped.contentType().id();
+ case 2 -> wrapped.location();
+ case 3 -> wrapped.fileFormat().toString();
+ case 4 -> wrapped.recordCount();
+ case 5 -> wrapped.fileSizeInBytes();
+ case 6 -> specId;
+ case 7 -> unionPartition();
+ case 8 -> contentStats.wrap(wrapped.contentStats());
+ case 9 -> wrapped.tightStats();
+ case 10 -> wrapped.sortOrderId();
+ case 11 -> deletionVector.wrap(wrapped.deletionVector());
+ case 12 ->
+ wrapped.manifestInfo() != null ?
manifestInfo.wrap(wrapped.manifestInfo()) : null;
+ case 13 -> wrapped.keyMetadata();
+ case 14 -> wrapped.splitOffsets();
+ case 15 -> wrapColumnFiles(wrapped.columnFiles());
+ default -> throw new UnsupportedOperationException("Unknown field
ordinal: " + pos);
+ };
+ }
+
+ private List<ColumnFileWrapper> wrapColumnFiles(List<ColumnFile>
columnFiles) {
+ if (columnFiles == null) {
+ return null;
+ }
+
+ for (int pos = columnFileWrappers.size(); pos < columnFiles.size(); pos
+= 1) {
+ columnFileWrappers.add(new ColumnFileWrapper());
+ }
+
+ for (int pos = 0; pos < columnFiles.size(); pos += 1) {
+ columnFileWrappers.get(pos).wrap(columnFiles.get(pos));
+ }
+
+ return columnFileWrappers.subList(0, columnFiles.size());
+ }
+
+ private StructLike unionPartition() {
+ if (specId == null) {
+ return null;
+ }
+
+ StructProjection projection = partitionProjections.get(specId);
+ return projection != null ? projection.wrap(wrapped.partition()) :
wrapped.partition();
+ }
+
+ private static void validatePartitionType(
+ Iterable<PartitionSpec> specs, Types.StructType unionType) {
+ for (PartitionSpec spec : specs) {
+ for (Types.NestedField field : spec.partitionType().fields()) {
+ Preconditions.checkArgument(
+ unionType.field(field.fieldId()) != null,
+ "Invalid union partition type, missing field %s from spec %s:
%s",
+ field.fieldId(),
+ spec.specId(),
+ unionType);
+ }
+ }
+ }
+
+ private static Map<Integer, StructProjection> partitionProjections(
+ Iterable<PartitionSpec> specs, Types.StructType unionType) {
+ Map<Integer, StructProjection> projections = Maps.newHashMap();
+ for (PartitionSpec spec : specs) {
+ Types.StructType partitionType = spec.partitionType();
+ StructProjection projection =
+ partitionType.equals(unionType)
+ ? null
+ : StructProjection.createAllowMissing(partitionType,
unionType);
+ projections.put(spec.specId(), projection);
+ }
+
+ return projections;
+ }
+ }
+
+ /** Wrapper used to write a data file to a v4 leaf manifest. */
+ static class LeafTrackedFileWrapper extends TrackedFileWrapper {
+ LeafTrackedFileWrapper(
+ Long snapshotId,
+ Map<Integer, PartitionSpec> specsById,
+ Types.StructType partitionType,
+ Types.StructType contentStatsType) {
+ super(
+ specsById,
+ partitionType,
+ contentStatsType,
+ new TrackingWrapper(snapshotId, null /* sequence numbers are
inherited */),
+ null /* leaf manifests only contain data files */);
+ }
+
+ LeafTrackedFileWrapper wrap(TrackedFile file, Tracking newTracking) {
+ wrapInternal(file, newTracking);
+ Preconditions.checkArgument(
+ file.contentType() == FileContent.DATA,
+ "Invalid content type for leaf manifest: %s",
+ file.contentType());
+
+ tracking().wrapLeaf(file, newTracking);
+ return this;
+ }
+ }
+
+ /** Wrapper used to write a data file or leaf manifest to a v4 root
manifest. */
+ static class RootTrackedFileWrapper extends TrackedFileWrapper {
+ RootTrackedFileWrapper(
+ long commitSnapshotId,
+ long commitSequenceNumber,
+ Map<Integer, PartitionSpec> specsById,
+ Types.StructType partitionType,
+ Types.StructType contentStatsType) {
+ super(
+ specsById,
+ partitionType,
+ contentStatsType,
+ new TrackingWrapper(commitSnapshotId, commitSequenceNumber),
+ new ManifestInfoWrapper(commitSequenceNumber));
+ }
+
+ RootTrackedFileWrapper wrap(TrackedFile file, Tracking newTracking, Long
assignedFirstRowId) {
+ wrapInternal(file, newTracking);
+ tracking().wrapRoot(file, newTracking, assignedFirstRowId);
+ return this;
+ }
+ }
+
+ /** Wrapper used to write {@link Tracking} using the v4 tracking schema. */
+ private static class TrackingWrapper implements StructLike {
+ private static final int SIZE = Tracking.schema().fields().size();
+
+ private final Long commitSnapshotId;
+ private final Long commitSequenceNumber;
+
+ private FileContent content = null;
+ private Tracking wrapped = null;
+ private Long assignedFirstRowId = null;
+
+ private TrackingWrapper(Long commitSnapshotId, Long commitSequenceNumber) {
+ this.commitSnapshotId = commitSnapshotId;
+ this.commitSequenceNumber = commitSequenceNumber;
+ }
+
+ private TrackingWrapper wrapLeaf(TrackedFile file, Tracking tracking) {
+ validateSequenceNumbers(file, tracking);
+
+ if (tracking.status() == EntryStatus.ADDED) {
Review Comment:
I can extract this first row id validation block to a small validation
method to share with wrapRoot. but it might affect readability. kept the
duplicates for now
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
