rdblue commented on code in PR #18853:
URL: https://github.com/apache/iceberg/pull/18853#discussion_r4543056928


##########
core/src/main/java/org/apache/iceberg/V4Metadata.java:
##########
@@ -541,4 +546,608 @@ 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 wrapFile(TrackedFile file) {
+      Preconditions.checkArgument(file != null, "Invalid file: 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;
+      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);
+      // projection is null when the spec's partition type already matches the 
union type
+      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 /* no manifest info wrapper. leaf manifests only contain data 
files */);
+    }
+
+    LeafTrackedFileWrapper wrap(TrackedFile file, Tracking newTracking) {
+      wrapFile(file);
+      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) {
+      wrapFile(file);
+      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) {
+      Preconditions.checkArgument(tracking != null, "Invalid tracking: null");
+      validateSequenceNumbers(file, tracking);
+
+      if (tracking.status() == EntryStatus.ADDED) {
+        Preconditions.checkState(
+            tracking.firstRowId() == null,
+            "Invalid first-row-id for ADDED file (must be null): %s",
+            file.location());
+      } else {
+        Preconditions.checkState(
+            tracking.firstRowId() != null,
+            "Missing first-row-id for %s file: %s",
+            tracking.status(),
+            file.location());
+      }
+
+      return wrapInternal(file, tracking, null /* assignedFirstRowId */);
+    }
+
+    private TrackingWrapper wrapRoot(TrackedFile file, Tracking tracking, Long 
firstRowId) {
+      Preconditions.checkArgument(tracking != null, "Invalid tracking: null");
+      Preconditions.checkArgument(
+          tracking.snapshotId() != null, "Missing snapshot ID for file: %s", 
file.location());
+      validateSequenceNumbers(file, tracking);
+
+      FileContent fileContent = file.contentType();
+      if (fileContent == FileContent.DATA || fileContent == 
FileContent.DATA_MANIFEST) {
+        if (tracking.status() == EntryStatus.ADDED) {
+          Preconditions.checkState(
+              tracking.firstRowId() == null,
+              "Invalid first-row-id for ADDED file (must be null): %s",
+              file.location());
+          Preconditions.checkState(
+              firstRowId != null,
+              "Missing first-row-id assignment for ADDED file: %s",
+              file.location());
+        } else {
+          Preconditions.checkState(
+              tracking.firstRowId() != null,
+              "Missing first-row-id for %s file: %s",
+              tracking.status(),
+              file.location());
+          Preconditions.checkState(
+              firstRowId == null,
+              "Cannot assign first-row-id to %s file: %s",
+              tracking.status(),
+              file.location());
+        }
+      } else if (fileContent == FileContent.DELETE_MANIFEST) {
+        Preconditions.checkState(
+            firstRowId == null,
+            "Cannot assign first-row-id to delete manifest: %s",
+            file.location());
+      } else {
+        throw new IllegalArgumentException(
+            String.format("Cannot write %s to a root manifest: %s", 
fileContent, file.location()));
+      }
+
+      return wrapInternal(file, tracking, firstRowId);
+    }
+
+    private void validateSequenceNumbers(TrackedFile file, Tracking tracking) {
+      EntryStatus status = tracking.status();
+      if (tracking.dataSequenceNumber() == null) {
+        Preconditions.checkState(
+            status == EntryStatus.ADDED || status == EntryStatus.MODIFIED,
+            "Cannot write null data sequence number for status %s (must be 
ADDED or MODIFIED): %s",
+            status,
+            file.location());
+
+        // a modified entry inherits the data sequence number of the commit 
that modified it
+        Long dataSnapshotId =
+            status == EntryStatus.MODIFIED ? tracking.modifiedSnapshotId() : 
tracking.snapshotId();
+        Preconditions.checkState(
+            Objects.equals(commitSnapshotId, dataSnapshotId),
+            "Cannot assign data sequence number to file from snapshot %s 
(committing snapshot %s): %s",
+            dataSnapshotId,
+            commitSnapshotId,
+            file.location());
+      }
+
+      if (tracking.fileSequenceNumber() == null) {
+        Preconditions.checkState(
+            status == EntryStatus.ADDED,
+            "Cannot write null file sequence number for status %s (must be 
ADDED): %s",
+            status,
+            file.location());
+        Preconditions.checkState(
+            Objects.equals(commitSnapshotId, tracking.snapshotId()),
+            "Cannot assign file sequence number to file from snapshot %s 
(committing snapshot %s): %s",
+            tracking.snapshotId(),
+            commitSnapshotId,
+            file.location());
+      }
+    }
+
+    private TrackingWrapper wrapInternal(TrackedFile file, Tracking tracking, 
Long firstRowId) {
+      this.content = file.contentType();
+      this.wrapped = tracking;
+      this.assignedFirstRowId = firstRowId;
+      return this;
+    }
+
+    @Override
+    public int size() {
+      return SIZE;
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      throw new UnsupportedOperationException("Cannot modify TrackingWrapper 
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 -> wrapped.status().id();
+        case 1 -> wrapped.snapshotId();
+        case 2 -> sequenceNumber(wrapped.dataSequenceNumber());
+        case 3 -> sequenceNumber(wrapped.fileSequenceNumber());
+        case 4 -> wrapped.modifiedSnapshotId();
+        case 5 -> firstRowId();
+        case 6 -> wrapped.deletedPositions();
+        case 7 -> wrapped.replacedPositions();
+        default -> throw new UnsupportedOperationException("Unknown field 
ordinal: " + pos);
+      };
+    }
+
+    private Long sequenceNumber(Long sequenceNumber) {
+      return sequenceNumber != null ? sequenceNumber : commitSequenceNumber;
+    }
+
+    private Long firstRowId() {
+      if (content != FileContent.DATA && content != FileContent.DATA_MANIFEST) 
{
+        return null;
+      }
+
+      return assignedFirstRowId != null ? assignedFirstRowId : 
wrapped.firstRowId();
+    }
+  }
+
+  /** Wrapper used to write {@link ManifestInfo} using the v4 manifest info 
schema. */
+  private static class ManifestInfoWrapper implements StructLike {
+    private static final int SIZE = ManifestInfo.schema().fields().size();
+
+    private final long commitSequenceNumber;
+    private ManifestInfo wrapped = null;
+
+    private ManifestInfoWrapper(long commitSequenceNumber) {
+      this.commitSequenceNumber = commitSequenceNumber;
+    }
+
+    private ManifestInfoWrapper wrap(ManifestInfo info) {
+      this.wrapped = info;
+      return this;
+    }
+
+    @Override
+    public int size() {
+      return SIZE;
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      throw new UnsupportedOperationException("Cannot modify 
ManifestInfoWrapper 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 -> wrapped.addedFilesCount();
+        case 1 -> wrapped.existingFilesCount();
+        case 2 -> wrapped.deletedFilesCount();
+        case 3 -> wrapped.replacedFilesCount();
+        case 4 -> wrapped.modifiedFilesCount();
+        case 5 -> wrapped.addedRowsCount();
+        case 6 -> wrapped.existingRowsCount();
+        case 7 -> wrapped.deletedRowsCount();
+        case 8 -> wrapped.replacedRowsCount();
+        case 9 -> wrapped.modifiedRowsCount();
+        case 10 -> minSequenceNumber();
+        case 11 -> mdvBuffer();
+        case 12 -> wrapped.formatVersion();
+        default -> throw new UnsupportedOperationException("Unknown field 
ordinal: " + pos);
+      };
+    }
+
+    private long minSequenceNumber() {
+      if (wrapped.minSequenceNumber() == ManifestWriter.UNASSIGNED_SEQ) {
+        return commitSequenceNumber;
+      }
+
+      return wrapped.minSequenceNumber();
+    }
+
+    private ByteBuffer mdvBuffer() {
+      ManifestBitmap mdv = wrapped.manifestDeletionVector();
+      return mdv != null ? mdv.buffer() : null;
+    }
+  }
+
+  /** Wrapper used to write {@link DeletionVector} using the v4 deletion 
vector schema. */
+  private static class DeletionVectorWrapper implements StructLike {
+    private static final int SIZE = DeletionVector.schema().fields().size();
+
+    private DeletionVector wrapped = null;
+
+    private DeletionVectorWrapper wrap(DeletionVector dv) {
+      if (dv == null) {
+        return null;
+      }
+
+      this.wrapped = dv;
+      return this;
+    }
+
+    @Override
+    public int size() {
+      return SIZE;
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      throw new UnsupportedOperationException(
+          "Cannot modify DeletionVectorWrapper 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 -> wrapped.location();
+        case 1 -> wrapped.offset();
+        case 2 -> wrapped.sizeInBytes();
+        case 3 -> wrapped.cardinality();
+        case 4 -> wrapped.keyMetadata();
+        default -> throw new UnsupportedOperationException("Unknown field 
ordinal: " + pos);
+      };
+    }
+  }
+
+  /** Wrapper used to write {@link ColumnFile} using the v4 column file 
schema. */
+  private static class ColumnFileWrapper implements StructLike {
+    private static final int SIZE = ColumnFile.schema().fields().size();
+
+    private ColumnFile wrapped = null;
+
+    private ColumnFileWrapper wrap(ColumnFile columnFile) {
+      this.wrapped = columnFile;
+      return this;
+    }
+
+    @Override
+    public int size() {
+      return SIZE;
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      throw new UnsupportedOperationException("Cannot modify ColumnFileWrapper 
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 -> wrapped.location();
+        case 1 -> wrapped.fieldIds();
+        case 2 -> wrapped.fileFormat().toString();
+        case 3 -> wrapped.fileSizeInBytes();
+        case 4 -> wrapped.keyMetadata();
+        default -> throw new UnsupportedOperationException("Unknown field 
ordinal: " + pos);
+      };
+    }
+  }
+
+  /** Wrapper used to write {@link ContentStats} using a manifest's content 
stats type. */
+  private static class ContentStatsWrapper implements StructLike {
+    private final int[] posToFieldId;
+    private final FieldStatsWrapper[] fieldStats;
+    private ContentStats wrapped = null;
+
+    private ContentStatsWrapper(Types.StructType type) {
+      List<Types.NestedField> fields = type.fields();
+      this.posToFieldId = ContentStatsStruct.posToId(type);
+      this.fieldStats = new FieldStatsWrapper[fields.size()];
+      for (int pos = 0; pos < fields.size(); pos += 1) {
+        fieldStats[pos] = new 
FieldStatsWrapper(fields.get(pos).type().asStructType());
+      }
+    }
+
+    private ContentStatsWrapper wrap(ContentStats stats) {
+      if (stats == null) {
+        return null;
+      }
+
+      this.wrapped = stats;
+      return this;
+    }
+
+    @Override
+    public int size() {
+      return posToFieldId.length;
+    }
+
+    @Override
+    public <T> void set(int pos, T value) {
+      throw new UnsupportedOperationException("Cannot modify 
ContentStatsWrapper wrapper via set");
+    }
+
+    @Override
+    public <T> T get(int pos, Class<T> javaClass) {
+      FieldStats<?> stats = wrapped.statsFor(posToFieldId[pos]);
+      return stats != null ? javaClass.cast(fieldStats[pos].wrap(stats)) : 
null;
+    }
+  }
+
+  /** Wrapper used to write {@link FieldStats} using a manifest's field stats 
type. */
+  private static class FieldStatsWrapper implements StructLike {
+    private final int[] posToOffset;
+    private FieldStats<?> wrapped = null;
+
+    private FieldStatsWrapper(Types.StructType type) {
+      this.posToOffset = FieldStatsStruct.posToOffset(type);
+    }
+
+    private FieldStatsWrapper wrap(FieldStats<?> stats) {

Review Comment:
   I think this should be null-tolerant like `DeletionVectorWrapper.wrap` so 
that `ContentStatsWrapper` doesn't need to handle it externally.



-- 
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]

Reply via email to