gaborkaszab commented on code in PR #16958: URL: https://github.com/apache/iceberg/pull/16958#discussion_r3506521442
########## core/src/test/java/org/apache/iceberg/TestV4ManifestReader.java: ########## @@ -0,0 +1,551 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg; + +import static org.apache.iceberg.types.Types.NestedField.optional; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.FileAppender; +import org.apache.iceberg.io.FileIO; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.metrics.DefaultMetricsContext; +import org.apache.iceberg.metrics.ScanMetrics; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.transforms.Transforms; +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.TestTemplate; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.io.TempDir; + +@ExtendWith(ParameterizedTestExtension.class) +public class TestV4ManifestReader { + private static final long SNAPSHOT_ID = 42L; + private static final int FORMAT_VERSION_V4 = 4; + + private static final Schema TABLE_SCHEMA = + new Schema( + optional(1, "id", Types.IntegerType.get()), optional(2, "data", Types.StringType.get())); + private static final PartitionSpec SPEC = + PartitionSpec.builderFor(TABLE_SCHEMA).identity("id").build(); + private static final Types.StructType PARTITION_TYPE = SPEC.partitionType(); + private static final Types.StructType EMPTY_PARTITION = Types.StructType.of(); + private static final PartitionData EMPTY_PARTITION_DATA = new PartitionData(EMPTY_PARTITION); + private static final Map<Integer, PartitionSpec> PARTITIONED_SPECS = + ImmutableMap.of(SPEC.specId(), SPEC); + private static final Map<Integer, PartitionSpec> UNPARTITIONED_SPECS = + ImmutableMap.of(PartitionSpec.unpartitioned().specId(), PartitionSpec.unpartitioned()); + + private static final List<Types.NestedField> SCHEMA_FIELDS = + TrackedFile.schemaWithContentStats(Types.StructType.of(), Types.StructType.of()).fields(); + + @Parameter private FileFormat format; + + @Parameters(name = "format = {0}") + protected static List<FileFormat> parameters() { + return Arrays.asList(FileFormat.AVRO, FileFormat.PARQUET); + } + + @TempDir private Path tempDir; + + private final FileIO fileIO = new TestTables.LocalFileIO(); + + @TestTemplate + public void testRoundTrip() throws IOException { + DeletionVector dv = + DeletionVectorStruct.builder() + .location("s3://bucket/dv.puffin") + .offset(100L) + .sizeInBytes(50L) + .cardinality(5L) + .build(); + + TrackedFile file = + dataFileBuilder("s3://bucket/data/file.parquet", partition(7)) + .sortOrderId(1) + .deletionVector(dv) + .keyMetadata(ByteBuffer.wrap(new byte[] {1, 2, 3})) + .splitOffsets(ImmutableList.of(50L, 100L)) + .build(); + + InputFile manifest = writeManifest(PARTITION_TYPE, ImmutableList.of(file)); + + List<TrackedFile> read = read(manifest, PARTITIONED_SPECS); + assertThat(read).hasSize(1); + TrackedFile actual = read.get(0); + + assertThat(actual.contentType()).isEqualTo(file.contentType()); + assertThat(actual.formatVersion()).isEqualTo(file.formatVersion()); + assertThat(actual.location()).isEqualTo(file.location()); + assertThat(actual.fileFormat()).isEqualTo(file.fileFormat()); + assertThat(actual.recordCount()).isEqualTo(file.recordCount()); + assertThat(actual.fileSizeInBytes()).isEqualTo(file.fileSizeInBytes()); + assertThat(actual.specId()).isEqualTo(file.specId()); + assertThat(actual.sortOrderId()).isEqualTo(file.sortOrderId()); + assertThat(actual.keyMetadata()).isEqualTo(file.keyMetadata()); + assertThat(actual.splitOffsets()).isEqualTo(file.splitOffsets()); + assertThat(actual.partition().get(0, Integer.class)) + .isEqualTo(file.partition().get(0, Integer.class)); + + assertThat(actual.tracking()).isNotNull(); + assertThat(actual.tracking().status()).isEqualTo(file.tracking().status()); + assertThat(actual.tracking().snapshotId()).isEqualTo(file.tracking().snapshotId()); + + assertThat(actual.deletionVector()).isNotNull(); + assertThat(actual.deletionVector().location()).isEqualTo(file.deletionVector().location()); + assertThat(actual.deletionVector().offset()).isEqualTo(file.deletionVector().offset()); + assertThat(actual.deletionVector().sizeInBytes()) + .isEqualTo(file.deletionVector().sizeInBytes()); + assertThat(actual.deletionVector().cardinality()) + .isEqualTo(file.deletionVector().cardinality()); + } + + @TestTemplate + public void testEqualityDeleteRoundTrip() throws IOException { + TrackedFile delete = + TrackedFileBuilder.equalityDelete(SNAPSHOT_ID) + .formatVersion(FORMAT_VERSION_V4) + .location("s3://bucket/eq-delete.parquet") + .fileFormat(FileFormat.PARQUET) + .recordCount(10L) + .fileSizeInBytes(128L) + .partition(EMPTY_PARTITION_DATA) + .specId(0) + .equalityIds(ImmutableList.of(1, 2)) + .build(); + + InputFile manifest = writeManifest(EMPTY_PARTITION, ImmutableList.of(delete)); + + TrackedFile actual = read(manifest, UNPARTITIONED_SPECS).get(0); + assertThat(actual.contentType()).isEqualTo(FileContent.EQUALITY_DELETES); + assertThat(actual.equalityIds()).containsExactly(1, 2); + } + + @TestTemplate + public void testLiveFilesExcludesDeletedAndReplaced() throws IOException { + List<TrackedFile> files = + ImmutableList.of( + fileWithStatus(EntryStatus.ADDED, "s3://bucket/added.parquet"), + fileWithStatus(EntryStatus.EXISTING, "s3://bucket/existing.parquet"), + fileWithStatus(EntryStatus.MODIFIED, "s3://bucket/modified.parquet"), + fileWithStatus(EntryStatus.DELETED, "s3://bucket/deleted.parquet"), + fileWithStatus(EntryStatus.REPLACED, "s3://bucket/replaced.parquet")); + + InputFile manifest = writeManifest(EMPTY_PARTITION, files); + + try (V4ManifestReader reader = newReader(manifest, UNPARTITIONED_SPECS).build()) { + assertThat(reader.allFiles()) + .extracting(file -> file.tracking().status()) + .containsExactly( + EntryStatus.ADDED, + EntryStatus.EXISTING, + EntryStatus.MODIFIED, + EntryStatus.DELETED, + EntryStatus.REPLACED); + + assertThat(reader.liveFiles()) + .extracting(file -> file.tracking().status()) + .containsExactly(EntryStatus.ADDED, EntryStatus.EXISTING, EntryStatus.MODIFIED); + } + } + + @TestTemplate + public void testManifestLocationAndPosition() throws IOException { + List<TrackedFile> files = + ImmutableList.of( + dataFile("s3://bucket/a.parquet", EMPTY_PARTITION_DATA), + dataFile("s3://bucket/b.parquet", EMPTY_PARTITION_DATA), + dataFile("s3://bucket/c.parquet", EMPTY_PARTITION_DATA)); + + InputFile manifest = writeManifest(EMPTY_PARTITION, files); + + List<TrackedFile> read = read(manifest, UNPARTITIONED_SPECS); + assertThat(read) + .allSatisfy( + file -> assertThat(file.tracking().manifestLocation()).isEqualTo(manifest.location())); + assertThat(read).extracting(file -> file.tracking().manifestPos()).containsExactly(0L, 1L, 2L); + } + + @TestTemplate + public void testProjectionRestrictsFields() throws IOException { + TrackedFile file = + dataFileBuilder("s3://bucket/file.parquet", EMPTY_PARTITION_DATA).sortOrderId(7).build(); + + InputFile manifest = writeManifest(EMPTY_PARTITION, ImmutableList.of(file)); + + Schema projection = new Schema(TrackedFile.LOCATION); + try (V4ManifestReader reader = + newReader(manifest, UNPARTITIONED_SPECS).project(projection).build()) { + TrackedFile actual = Lists.newArrayList(reader.allFiles()).get(0); + assertThat(actual.location()).isEqualTo(file.location()); + // a minimal status-only tracking is force-added when the caller omits tracking + assertThat(actual.tracking()).isNotNull(); + assertThat(actual.tracking().status()).isEqualTo(EntryStatus.ADDED); + // sort_order_id, file_format, and spec_id are null because they were not projected + assertThat(actual.sortOrderId()).isNull(); + assertThat(actual.fileFormat()).isNull(); + assertThat(actual.specId()).isNull(); + } + } + + @TestTemplate + public void testUnpartitioned() throws IOException { + TrackedFile file = dataFile("s3://bucket/file.parquet", EMPTY_PARTITION_DATA); + Review Comment: Is there a test where we do a projection that doesn't include partition, and partition added to the projection because of the partition filtering? ########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,276 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.iceberg.expressions.Evaluator; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.expressions.Projections; +import org.apache.iceberg.io.CloseableGroup; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.CloseableIterator; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.metrics.ScanMetrics; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +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.util.StructProjection; + +/** Reader that reads a v4+ manifest file as {@link TrackedFile}s. */ +class V4ManifestReader extends CloseableGroup implements CloseableIterable<TrackedFile> { + private final InputFile file; + private final Types.StructType partitionType; + private final Schema fileProjection; + private final ScanMetrics scanMetrics; + + // partition pruning state, keyed by spec ID; empty when no filtering is required + private final Map<Integer, Evaluator> partitionEvaluators; + private final Map<Integer, StructProjection> partitionProjections; + + private V4ManifestReader( + InputFile file, + Types.StructType partitionType, + Map<Integer, Evaluator> partitionEvaluators, + Map<Integer, StructProjection> partitionProjections, + Schema fileProjection, + ScanMetrics scanMetrics) { + this.file = file; + this.partitionType = partitionType; + this.partitionEvaluators = partitionEvaluators; + this.partitionProjections = partitionProjections; + this.fileProjection = fileProjection; + this.scanMetrics = scanMetrics; + } + + static Builder builder( + InputFile file, Schema tableSchema, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, tableSchema, specsById); + } + + /** Returns all tracked files in this manifest, regardless of status. */ + CloseableIterable<TrackedFile> allFiles() { + return files(false /* all files */); + } + + /** Returns tracked files whose tracking {@link Tracking#isLive() is live}. */ + CloseableIterable<TrackedFile> liveFiles() { + return files(true /* only live files */); + } + + /** Returns live tracked files, each as an independent copy. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + return CloseableIterable.transform(liveFiles(), TrackedFile::copy).iterator(); + } + + private CloseableIterable<TrackedFile> files(boolean onlyLive) { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionEvaluators.isEmpty()) { + entries = CloseableIterable.filter(entries, this::matchesPartition); + } + + if (onlyLive) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return entries; + } + + private boolean matchesPartition(TrackedFile trackedFile) { + FileContent content = trackedFile.contentType(); + if (content == FileContent.DATA_MANIFEST || content == FileContent.DELETE_MANIFEST) { + // manifest references are expanded later and are not pruned by the partition filter + return true; + } + + Integer specId = trackedFile.specId(); + Evaluator evaluator = specId != null ? partitionEvaluators.get(specId) : null; + StructProjection projection = specId != null ? partitionProjections.get(specId) : null; + Preconditions.checkState( + evaluator != null && projection != null, + "Cannot apply partition filter: file %s has spec ID %s, not one of the known specs %s " + + "in manifest %s", + trackedFile.location(), + specId, + partitionEvaluators.keySet(), + file.location()); + + boolean matches = evaluator.eval(projection.wrap(trackedFile.partition())); + if (!matches) { + if (content == FileContent.DATA) { + scanMetrics.skippedDataFiles().increment(); + } else { + scanMetrics.skippedDeleteFiles().increment(); + } + } + + return matches; + } + + private CloseableIterable<TrackedFile> open() { + FileFormat format = FileFormat.fromFileName(file.location()); + Preconditions.checkArgument( + format != null, "Unable to determine format of manifest: %s", file.location()); + + CloseableIterable<TrackedFile> reader = + InternalData.read(format, file) + .project(readSchema()) + .setRootType(TrackedFileStruct.class) + .setCustomType(TrackedFile.TRACKING.fieldId(), TrackingStruct.class) + .setCustomType(TrackedFile.DELETION_VECTOR.fieldId(), DeletionVectorStruct.class) + .setCustomType(TrackedFile.MANIFEST_INFO.fieldId(), ManifestInfoStruct.class) + .setCustomType(TrackedFile.PARTITION_ID, PartitionData.class) + .reuseContainers() + .build(); + addCloseable(reader); + return reader; + } + + private TrackedFile prepare(TrackedFile trackedFile) { + Tracking tracking = trackedFile.tracking(); + // manifestLocation is not stored in the manifest; the reader fills it from the file location. + // manifestPos is filled from ROW_POSITION while reading the tracking struct. + if (tracking instanceof TrackingStruct) { + ((TrackingStruct) tracking).setManifestLocation(file.location()); + } + + return trackedFile; + } + + private Schema readSchema() { + // content_stats is not projected yet, so build the schema with an empty stats type + Types.StructType fullType = + TrackedFile.schemaWithContentStats(partitionType, Types.StructType.of()); + boolean unpartitioned = partitionType.fields().isEmpty(); + + Set<Integer> projectedIds = null; + if (fileProjection != null) { + projectedIds = Sets.newHashSet(); + for (Types.NestedField field : fileProjection.asStruct().fields()) { + projectedIds.add(field.fieldId()); + } + + // tracking carries the status used to filter live files and is always projected + projectedIds.add(TrackedFile.TRACKING.fieldId()); + + // spec_id is required to resolve each entry's partition spec when pruning + if (!partitionEvaluators.isEmpty()) { + projectedIds.add(TrackedFile.SPEC_ID.fieldId()); + } + } + + List<Types.NestedField> fields = Lists.newArrayList(); + for (Types.NestedField field : fullType.fields()) { + if (projectedIds != null && !projectedIds.contains(field.fieldId())) { + continue; + } + + if (field.fieldId() == TrackedFile.TRACKING.fieldId()) { + fields.add(trackingWithRowPosition()); + } else if (field.fieldId() == TrackedFile.CONTENT_STATS_ID) { Review Comment: Thanks for the explanation! Projecting stats makes total sense. I think most of my comment on this area are for a single reason: for me this readSchema() function seems a bit more complicated than what it is for. What I have in mind is like this: ``` if (fileProjection == null) { return TrackedFileStruct.READ_SCHEMA; // this could be prepared with the desired Tracking schema, omitting stats, including partition unconditionally. Preferable a static field in TrackedFileStruct. } // construct the projected schema including the mandatory fields similarly as now ``` ########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,276 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.iceberg.expressions.Evaluator; +import org.apache.iceberg.expressions.Expression; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.expressions.Projections; +import org.apache.iceberg.io.CloseableGroup; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.CloseableIterator; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.metrics.ScanMetrics; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +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.util.StructProjection; + +/** Reader that reads a v4+ manifest file as {@link TrackedFile}s. */ +class V4ManifestReader extends CloseableGroup implements CloseableIterable<TrackedFile> { + private final InputFile file; + private final Types.StructType partitionType; + private final Schema fileProjection; + private final ScanMetrics scanMetrics; + + // partition pruning state, keyed by spec ID; empty when no filtering is required + private final Map<Integer, Evaluator> partitionEvaluators; + private final Map<Integer, StructProjection> partitionProjections; + + private V4ManifestReader( + InputFile file, + Types.StructType partitionType, + Map<Integer, Evaluator> partitionEvaluators, + Map<Integer, StructProjection> partitionProjections, + Schema fileProjection, + ScanMetrics scanMetrics) { + this.file = file; + this.partitionType = partitionType; + this.partitionEvaluators = partitionEvaluators; + this.partitionProjections = partitionProjections; + this.fileProjection = fileProjection; + this.scanMetrics = scanMetrics; + } + + static Builder builder( + InputFile file, Schema tableSchema, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, tableSchema, specsById); + } + + /** Returns all tracked files in this manifest, regardless of status. */ + CloseableIterable<TrackedFile> allFiles() { + return files(false /* all files */); + } + + /** Returns tracked files whose tracking {@link Tracking#isLive() is live}. */ + CloseableIterable<TrackedFile> liveFiles() { + return files(true /* only live files */); + } + + /** Returns live tracked files, each as an independent copy. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + return CloseableIterable.transform(liveFiles(), TrackedFile::copy).iterator(); + } + + private CloseableIterable<TrackedFile> files(boolean onlyLive) { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionEvaluators.isEmpty()) { + entries = CloseableIterable.filter(entries, this::matchesPartition); + } + + if (onlyLive) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return entries; + } + + private boolean matchesPartition(TrackedFile trackedFile) { + FileContent content = trackedFile.contentType(); + if (content == FileContent.DATA_MANIFEST || content == FileContent.DELETE_MANIFEST) { + // manifest references are expanded later and are not pruned by the partition filter + return true; + } + + Integer specId = trackedFile.specId(); + Evaluator evaluator = specId != null ? partitionEvaluators.get(specId) : null; + StructProjection projection = specId != null ? partitionProjections.get(specId) : null; + Preconditions.checkState( Review Comment: My point was that specId = null is a valid state, but we throw an exception here because of that because both the evaluator and the projection will be null. -- 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]
