nastra commented on code in PR #16958: URL: https://github.com/apache/iceberg/pull/16958#discussion_r3622644314
########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,315 @@ +/* + * 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.Collection; +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.Maps; +import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.types.TypeUtil; +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 Schema readSchema; + private final boolean includeTombstones; + private final ScanMetrics scanMetrics; + + // partition pruning state, keyed by spec ID + private final Map<Integer, Evaluator> partitionEvaluators; + private final Map<Integer, StructProjection> partitionProjections; + + private V4ManifestReader( + InputFile file, + Schema readSchema, + Map<Integer, Evaluator> partitionEvaluators, + Map<Integer, StructProjection> partitionProjections, + boolean includeTombstones, + ScanMetrics scanMetrics) { + this.file = file; + this.readSchema = readSchema; + this.partitionEvaluators = partitionEvaluators; + this.partitionProjections = partitionProjections; + this.includeTombstones = includeTombstones; + this.scanMetrics = scanMetrics; + } + + static Builder builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, specsById); + } + + /** Returns copies of the tracked files that match this reader's configured filters. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionEvaluators.isEmpty()) { + // manifest references are expanded later and are not pruned by the partition filter + entries = + CloseableIterable.filter(entries, entry -> isManifest(entry) || matchesPartition(entry)); + } + + if (!includeTombstones) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return CloseableIterable.transform(entries, TrackedFile::copy).iterator(); + } + + private boolean matchesPartition(TrackedFile trackedFile) { + Integer specId = trackedFile.specId(); + if (specId == null) { + // a file without a spec is not partitioned and may match the filter + return true; + } + + Evaluator evaluator = partitionEvaluators.get(specId); + if (evaluator == null) { + // the row filter does not project to a partition filter for this spec + return true; + } + + StructProjection projection = partitionProjections.get(specId); + Preconditions.checkState( + projection != null, + "Cannot produce partition tuple for spec ID %s in manifest %s", + specId, + file.location()); + + boolean matches = evaluator.eval(projection.wrap(trackedFile.partition())); + if (!matches) { + incrementSkipCount(trackedFile.contentType()); + } + + return matches; + } + + private void incrementSkipCount(FileContent content) { + switch (content) { + case DATA: + scanMetrics.skippedDataFiles().increment(); + break; + case EQUALITY_DELETES: + scanMetrics.skippedDeleteFiles().increment(); + break; + case DATA_MANIFEST: + scanMetrics.skippedDataManifests().increment(); + break; + case DELETE_MANIFEST: + scanMetrics.skippedDeleteManifests().increment(); + break; + default: + throw new UnsupportedOperationException("Unsupported content type: " + content); + } + } + + private CloseableIterable<TrackedFile> open() { + FileFormat format = FileFormat.fromFileName(file.location()); + Preconditions.checkArgument( + format != null, "Cannot 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 in + if (tracking instanceof TrackingStruct) { + ((TrackingStruct) tracking).setManifestLocation(file.location()); + } + + return trackedFile; + } + + private static boolean isManifest(TrackedFile trackedFile) { + FileContent content = trackedFile.contentType(); + return content == FileContent.DATA_MANIFEST || content == FileContent.DELETE_MANIFEST; + } + + static class Builder { + private final InputFile file; + private final Types.StructType unionPartitionType; + private final Map<Integer, PartitionSpec> specsById; + private Expression rowFilter = Expressions.alwaysTrue(); + private boolean caseSensitive = true; + private boolean includeTombstones = false; + private boolean scanPlanning = false; + private Collection<String> columns = null; + private Schema fileProjection = null; + private ScanMetrics scanMetrics = ScanMetrics.noop(); + + private Builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + this.file = file; + this.unionPartitionType = Partitioning.unionPartitionTypes(specsById.values()); + this.specsById = specsById; + } + + /** Sets a row filter; files that cannot match the expression are skipped. */ + Builder filterRows(Expression expr) { + Preconditions.checkArgument(expr != null, "Invalid row filter: null"); + this.rowFilter = expr; + return this; + } + + Builder caseSensitive(boolean isCaseSensitive) { + this.caseSensitive = isCaseSensitive; + return this; + } + + /** Returns deleted and replaced files in addition to {@link Tracking#isLive() live} files. */ + Builder includeTombstones() { + this.includeTombstones = true; + return this; + } + + /** Configures the reader to select the minimal fields needed for scan planning. */ + Builder forScanPlanning() { + Preconditions.checkState( + columns == null && fileProjection == null, + "Cannot use forScanPlanning() with select(Collection<String>) or project(Schema)"); + this.scanPlanning = true; + return this; + } + + /** Selects columns to read by name; fields needed by the reader are always read. */ + Builder select(Collection<String> newColumns) { + Preconditions.checkArgument(newColumns != null, "Invalid columns: null"); + Preconditions.checkState( + !scanPlanning, "Cannot use select(Collection<String>) with forScanPlanning()"); + Preconditions.checkState( + fileProjection == null, + "Cannot select columns using both select(Collection<String>) and project(Schema)"); Review Comment: nit: `Cannot use select(Collection<String>) with project(Schema)` ########## api/src/test/java/org/apache/iceberg/types/TestTypeUtil.java: ########## @@ -1037,4 +1038,41 @@ public void testIndexStatsNames() { .containsEntry(24, "addresses_value") // the leaf takes precedence .hasSize(22); } + + @Test + public void testReplaceFieldTypes() { + Types.StructType replacement = Types.StructType.of(required(10, "x", IntegerType.get())); + Schema schema = + new Schema( + required(1, "id", IntegerType.get()), + required(2, "s", Types.StructType.of(required(3, "a", Types.LongType.get())))); + + Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, (Type) replacement)); + + assertThat(result.findField(1).type()).isEqualTo(IntegerType.get()); + assertThat(result.findField(2).type()).isEqualTo(replacement); + } + + @Test + public void testReplaceFieldTypesListElement() { + Schema schema = + new Schema(required(1, "list", Types.ListType.ofRequired(2, Types.StructType.of()))); + + Schema result = + TypeUtil.replaceFieldTypes( + schema, + ImmutableMap.of(2, (Type) Types.StructType.of(required(3, "x", IntegerType.get())))); Review Comment: ```suggestion ImmutableMap.of(2, Types.StructType.of(required(3, "x", IntegerType.get())))); ``` ########## api/src/test/java/org/apache/iceberg/types/TestTypeUtil.java: ########## @@ -1037,4 +1038,41 @@ public void testIndexStatsNames() { .containsEntry(24, "addresses_value") // the leaf takes precedence .hasSize(22); } + + @Test + public void testReplaceFieldTypes() { + Types.StructType replacement = Types.StructType.of(required(10, "x", IntegerType.get())); + Schema schema = + new Schema( + required(1, "id", IntegerType.get()), + required(2, "s", Types.StructType.of(required(3, "a", Types.LongType.get())))); + + Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, (Type) replacement)); Review Comment: ```suggestion Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, replacement)); ``` unnecessary cast here and in the other tests ########## api/src/main/java/org/apache/iceberg/types/ReplaceTypeById.java: ########## @@ -0,0 +1,107 @@ +/* + * 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.types; + +import java.util.List; +import java.util.Map; +import org.apache.iceberg.Schema; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; + +class ReplaceTypeById extends TypeUtil.SchemaVisitor<Type> { + private final Map<Integer, Type> replacementsById; + + ReplaceTypeById(Map<Integer, Type> replacementsById) { + this.replacementsById = replacementsById; + } + + @Override + public Type schema(Schema schema, Type structResult) { + return structResult; + } + + @Override + public Type struct(Types.StructType struct, List<Type> fieldResults) { + List<Types.NestedField> fields = struct.fields(); + List<Types.NestedField> newFields = Lists.newArrayListWithExpectedSize(fields.size()); + boolean hasChanged = false; + + for (int i = 0; i < fields.size(); i += 1) { + Type fieldReplacement = fieldResults.get(i); + Types.NestedField field = fields.get(i); + if (field.type() != fieldReplacement) { + hasChanged = true; + newFields.add(Types.NestedField.from(field).ofType(fieldReplacement).build()); Review Comment: do we need to handle initial/write defaults here when we're changing the type? What happens if the original type is int with some defaults but we change the type to string? I think we should have a test for this scenario ########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,310 @@ +/* + * 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.Arrays; +import java.util.Collection; +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.ImmutableMap; +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.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Pair; +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 Schema readSchema; + private final boolean includeAll; + private final ScanMetrics scanMetrics; + + // partition filters keyed by spec ID; empty when no partition filter applies + private final Map<Integer, Pair<Evaluator, StructProjection>> partitionFilters; + + private V4ManifestReader( + InputFile file, + Schema readSchema, + Map<Integer, Pair<Evaluator, StructProjection>> partitionFilters, + boolean includeAll, + ScanMetrics scanMetrics) { + this.file = file; + this.readSchema = readSchema; + this.partitionFilters = partitionFilters; + this.includeAll = includeAll; + this.scanMetrics = scanMetrics; + } + + static Builder builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, specsById); + } + + /** Returns copies of the tracked files that match this reader's configured filters. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionFilters.isEmpty()) { + // manifests have no partition, so the partition filter cannot apply to them + entries = + CloseableIterable.filter(entries, entry -> isManifest(entry) || matchesPartition(entry)); + } + + if (!includeAll) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return CloseableIterable.transform(entries, TrackedFile::copy).iterator(); + } + + private boolean matchesPartition(TrackedFile trackedFile) { + Integer specId = trackedFile.specId(); + if (specId == null) { + // a file without a spec is not partitioned and may match the filter + return true; + } + + Pair<Evaluator, StructProjection> partitionFilter = partitionFilters.get(specId); + if (partitionFilter == null) { + // the row filter does not project to a partition filter for this spec + return true; + } + + Evaluator evaluator = partitionFilter.first(); + StructProjection projection = partitionFilter.second(); + boolean matches = evaluator.eval(projection.wrap(trackedFile.partition())); + if (!matches) { + incrementSkipCount(trackedFile.contentType()); + } + + return matches; + } + + private void incrementSkipCount(FileContent content) { + switch (content) { + case DATA: + scanMetrics.skippedDataFiles().increment(); + break; + case EQUALITY_DELETES: + scanMetrics.skippedDeleteFiles().increment(); + break; + case DATA_MANIFEST: + scanMetrics.skippedDataManifests().increment(); + break; + case DELETE_MANIFEST: + scanMetrics.skippedDeleteManifests().increment(); + break; + default: + throw new UnsupportedOperationException("Unsupported content type: " + content); + } + } + + private CloseableIterable<TrackedFile> open() { + FileFormat format = FileFormat.fromFileName(file.location()); + Preconditions.checkArgument( + format != null, "Cannot 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 in + if (tracking instanceof TrackingStruct) { + ((TrackingStruct) tracking).setManifestLocation(file.location()); + } + + return trackedFile; + } + + private static boolean isManifest(TrackedFile trackedFile) { + FileContent content = trackedFile.contentType(); + return content == FileContent.DATA_MANIFEST || content == FileContent.DELETE_MANIFEST; + } + + static class Builder { + private final InputFile file; + private final Types.StructType unionPartitionType; + private final Map<Integer, PartitionSpec> specsById; + private final Schema fullSchema; + private Expression rowFilter = Expressions.alwaysTrue(); + private boolean caseSensitive = true; + private boolean includeAll = false; + private boolean scanPlanning = false; + private Collection<String> columns = null; + private Schema requestedProjection = null; + private ScanMetrics scanMetrics = ScanMetrics.noop(); + + private Builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + this.file = file; + this.specsById = specsById; + this.unionPartitionType = Partitioning.unionPartitionTypes(specsById.values()); + Schema base = TrackedFile.schema(unionPartitionType, Types.StructType.of()); + // the read schema carries row_position (via BASE_TYPE) so the reader can fill manifestPos + this.fullSchema = + TypeUtil.replaceFieldTypes( + base, ImmutableMap.of(TrackedFile.TRACKING.fieldId(), TrackingStruct.BASE_TYPE)); + } + + /** Sets a filter; files that cannot match the expression are skipped. */ + Builder filter(Expression expr) { + Preconditions.checkArgument(expr != null, "Invalid filter: null"); + this.rowFilter = expr; + return this; + } + + Builder caseSensitive(boolean isCaseSensitive) { + this.caseSensitive = isCaseSensitive; + return this; + } + + /** Returns all entries without filtering by {@link Tracking#isLive() liveness}. */ + Builder includeAll() { + this.includeAll = true; + return this; + } + + /** Configures the reader to select the minimal fields needed for scan planning. */ + Builder forScanPlanning() { + Preconditions.checkState( + columns == null && requestedProjection == null, + "Cannot use forScanPlanning() with select(Collection<String>) or project(Schema)"); + this.scanPlanning = true; + return this; + } + + /** Selects columns to read by name; fields needed by the reader are always read. */ + Builder select(String... newColumns) { + return select(Arrays.asList(newColumns)); + } + + /** Selects columns to read by name; fields needed by the reader are always read. */ + Builder select(Collection<String> newColumns) { + Preconditions.checkArgument(newColumns != null, "Invalid columns: null"); + Preconditions.checkState( + !scanPlanning, "Cannot use select(Collection<String>) with forScanPlanning()"); + Preconditions.checkState( + requestedProjection == null, + "Cannot select columns using both select(Collection<String>) and project(Schema)"); + this.columns = newColumns; + return this; + } + + /** Sets the exact schema to read; used in place of {@link #select(Collection)}. */ + Builder project(Schema newProjection) { Review Comment: nit: `project(null)` would read everything. If that is a supported case, then we should have a test for it or add a `newProjection != null` check ########## api/src/main/java/org/apache/iceberg/types/ReplaceTypeById.java: ########## @@ -0,0 +1,107 @@ +/* + * 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.types; + +import java.util.List; +import java.util.Map; +import org.apache.iceberg.Schema; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; + +class ReplaceTypeById extends TypeUtil.SchemaVisitor<Type> { Review Comment: minor: I think it might be worth pulling this out into a separate PR in order to reduce complexity in this one and harden the impl + have more extensive tests (especially around default value handling) ########## api/src/test/java/org/apache/iceberg/types/TestTypeUtil.java: ########## @@ -1037,4 +1038,41 @@ public void testIndexStatsNames() { .containsEntry(24, "addresses_value") // the leaf takes precedence .hasSize(22); } + + @Test + public void testReplaceFieldTypes() { + Types.StructType replacement = Types.StructType.of(required(10, "x", IntegerType.get())); + Schema schema = + new Schema( + required(1, "id", IntegerType.get()), + required(2, "s", Types.StructType.of(required(3, "a", Types.LongType.get())))); + + Schema result = TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(2, (Type) replacement)); + + assertThat(result.findField(1).type()).isEqualTo(IntegerType.get()); + assertThat(result.findField(2).type()).isEqualTo(replacement); + } + + @Test + public void testReplaceFieldTypesListElement() { + Schema schema = + new Schema(required(1, "list", Types.ListType.ofRequired(2, Types.StructType.of()))); + + Schema result = + TypeUtil.replaceFieldTypes( + schema, + ImmutableMap.of(2, (Type) Types.StructType.of(required(3, "x", IntegerType.get())))); + + Types.ListType list = (Types.ListType) result.findField(1).type(); + assertThat(list.elementType().asStructType().field(3).name()).isEqualTo("x"); + assertThat(list.isElementRequired()).isTrue(); + } + + @Test + public void testReplaceFieldTypesNoMatchReturnsSameSchema() { + Schema schema = new Schema(required(1, "id", IntegerType.get())); + Schema result = + TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(99, (Type) Types.LongType.get())); Review Comment: ```suggestion TypeUtil.replaceFieldTypes(schema, ImmutableMap.of(99, Types.LongType.get())); ``` ########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,315 @@ +/* + * 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.Collection; +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.Maps; +import org.apache.iceberg.relocated.com.google.common.collect.Sets; +import org.apache.iceberg.types.TypeUtil; +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 Schema readSchema; + private final boolean includeTombstones; + private final ScanMetrics scanMetrics; + + // partition pruning state, keyed by spec ID + private final Map<Integer, Evaluator> partitionEvaluators; + private final Map<Integer, StructProjection> partitionProjections; + + private V4ManifestReader( + InputFile file, + Schema readSchema, + Map<Integer, Evaluator> partitionEvaluators, + Map<Integer, StructProjection> partitionProjections, + boolean includeTombstones, + ScanMetrics scanMetrics) { + this.file = file; + this.readSchema = readSchema; + this.partitionEvaluators = partitionEvaluators; + this.partitionProjections = partitionProjections; + this.includeTombstones = includeTombstones; + this.scanMetrics = scanMetrics; + } + + static Builder builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, specsById); + } + + /** Returns copies of the tracked files that match this reader's configured filters. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionEvaluators.isEmpty()) { + // manifest references are expanded later and are not pruned by the partition filter + entries = + CloseableIterable.filter(entries, entry -> isManifest(entry) || matchesPartition(entry)); + } + + if (!includeTombstones) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return CloseableIterable.transform(entries, TrackedFile::copy).iterator(); + } + + private boolean matchesPartition(TrackedFile trackedFile) { + Integer specId = trackedFile.specId(); + if (specId == null) { + // a file without a spec is not partitioned and may match the filter + return true; + } + + Evaluator evaluator = partitionEvaluators.get(specId); + if (evaluator == null) { + // the row filter does not project to a partition filter for this spec + return true; + } + + StructProjection projection = partitionProjections.get(specId); + Preconditions.checkState( + projection != null, + "Cannot produce partition tuple for spec ID %s in manifest %s", + specId, + file.location()); + + boolean matches = evaluator.eval(projection.wrap(trackedFile.partition())); + if (!matches) { + incrementSkipCount(trackedFile.contentType()); + } + + return matches; + } + + private void incrementSkipCount(FileContent content) { + switch (content) { + case DATA: + scanMetrics.skippedDataFiles().increment(); + break; + case EQUALITY_DELETES: + scanMetrics.skippedDeleteFiles().increment(); + break; + case DATA_MANIFEST: + scanMetrics.skippedDataManifests().increment(); + break; + case DELETE_MANIFEST: + scanMetrics.skippedDeleteManifests().increment(); + break; + default: + throw new UnsupportedOperationException("Unsupported content type: " + content); + } + } + + private CloseableIterable<TrackedFile> open() { + FileFormat format = FileFormat.fromFileName(file.location()); + Preconditions.checkArgument( + format != null, "Cannot 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 in + if (tracking instanceof TrackingStruct) { + ((TrackingStruct) tracking).setManifestLocation(file.location()); + } + + return trackedFile; + } + + private static boolean isManifest(TrackedFile trackedFile) { + FileContent content = trackedFile.contentType(); + return content == FileContent.DATA_MANIFEST || content == FileContent.DELETE_MANIFEST; + } + + static class Builder { + private final InputFile file; + private final Types.StructType unionPartitionType; + private final Map<Integer, PartitionSpec> specsById; + private Expression rowFilter = Expressions.alwaysTrue(); + private boolean caseSensitive = true; + private boolean includeTombstones = false; + private boolean scanPlanning = false; + private Collection<String> columns = null; + private Schema fileProjection = null; + private ScanMetrics scanMetrics = ScanMetrics.noop(); + + private Builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + this.file = file; + this.unionPartitionType = Partitioning.unionPartitionTypes(specsById.values()); + this.specsById = specsById; + } + + /** Sets a row filter; files that cannot match the expression are skipped. */ + Builder filterRows(Expression expr) { + Preconditions.checkArgument(expr != null, "Invalid row filter: null"); + this.rowFilter = expr; + return this; + } + + Builder caseSensitive(boolean isCaseSensitive) { + this.caseSensitive = isCaseSensitive; + return this; + } + + /** Returns deleted and replaced files in addition to {@link Tracking#isLive() live} files. */ + Builder includeTombstones() { + this.includeTombstones = true; + return this; + } + + /** Configures the reader to select the minimal fields needed for scan planning. */ + Builder forScanPlanning() { + Preconditions.checkState( + columns == null && fileProjection == null, + "Cannot use forScanPlanning() with select(Collection<String>) or project(Schema)"); + this.scanPlanning = true; + return this; + } + + /** Selects columns to read by name; fields needed by the reader are always read. */ + Builder select(Collection<String> newColumns) { + Preconditions.checkArgument(newColumns != null, "Invalid columns: null"); + Preconditions.checkState( + !scanPlanning, "Cannot use select(Collection<String>) with forScanPlanning()"); + Preconditions.checkState( + fileProjection == null, + "Cannot select columns using both select(Collection<String>) and project(Schema)"); + this.columns = newColumns; + return this; + } + + /** Sets the exact schema to read; used in place of {@link #select(Collection)}. */ + Builder project(Schema newFileProjection) { + Preconditions.checkState(!scanPlanning, "Cannot use project(Schema) with forScanPlanning()"); + Preconditions.checkState( + columns == null, + "Cannot select columns using both select(Collection<String>) and project(Schema)"); Review Comment: nit: since we're inside `project` maybe the error msg should be `Cannot use project(Schema) with select(...)` to be slightly more consistent around what gets mentioned first ########## core/src/main/java/org/apache/iceberg/V4ManifestReader.java: ########## @@ -0,0 +1,310 @@ +/* + * 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.Arrays; +import java.util.Collection; +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.ImmutableMap; +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.TypeUtil; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Pair; +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 Schema readSchema; + private final boolean includeAll; + private final ScanMetrics scanMetrics; + + // partition filters keyed by spec ID; empty when no partition filter applies + private final Map<Integer, Pair<Evaluator, StructProjection>> partitionFilters; + + private V4ManifestReader( + InputFile file, + Schema readSchema, + Map<Integer, Pair<Evaluator, StructProjection>> partitionFilters, + boolean includeAll, + ScanMetrics scanMetrics) { + this.file = file; + this.readSchema = readSchema; + this.partitionFilters = partitionFilters; + this.includeAll = includeAll; + this.scanMetrics = scanMetrics; + } + + static Builder builder(InputFile file, Map<Integer, PartitionSpec> specsById) { + return new Builder(file, specsById); + } + + /** Returns copies of the tracked files that match this reader's configured filters. */ + @Override + public CloseableIterator<TrackedFile> iterator() { + CloseableIterable<TrackedFile> entries = CloseableIterable.transform(open(), this::prepare); + if (!partitionFilters.isEmpty()) { + // manifests have no partition, so the partition filter cannot apply to them + entries = + CloseableIterable.filter(entries, entry -> isManifest(entry) || matchesPartition(entry)); + } + + if (!includeAll) { + entries = CloseableIterable.filter(entries, entry -> entry.tracking().isLive()); + } + + return CloseableIterable.transform(entries, TrackedFile::copy).iterator(); + } + + private boolean matchesPartition(TrackedFile trackedFile) { + Integer specId = trackedFile.specId(); + if (specId == null) { + // a file without a spec is not partitioned and may match the filter + return true; + } + + Pair<Evaluator, StructProjection> partitionFilter = partitionFilters.get(specId); + if (partitionFilter == null) { + // the row filter does not project to a partition filter for this spec + return true; + } + + Evaluator evaluator = partitionFilter.first(); + StructProjection projection = partitionFilter.second(); + boolean matches = evaluator.eval(projection.wrap(trackedFile.partition())); + if (!matches) { + incrementSkipCount(trackedFile.contentType()); + } + + return matches; + } + + private void incrementSkipCount(FileContent content) { Review Comment: I think this can be shortened/simplified via ``` private void incrementSkipCount(FileContent content) { switch (content) { case DATA -> scanMetrics.skippedDataFiles().increment(); case EQUALITY_DELETES -> scanMetrics.skippedDeleteFiles().increment(); case DATA_MANIFEST -> scanMetrics.skippedDataManifests().increment(); case DELETE_MANIFEST -> scanMetrics.skippedDeleteManifests().increment(); default -> throw new UnsupportedOperationException("Unsupported content type: " + content); } } ``` ########## api/src/main/java/org/apache/iceberg/types/ReplaceTypeById.java: ########## @@ -0,0 +1,107 @@ +/* + * 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.types; + +import java.util.List; +import java.util.Map; +import org.apache.iceberg.Schema; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; + +class ReplaceTypeById extends TypeUtil.SchemaVisitor<Type> { + private final Map<Integer, Type> replacementsById; + + ReplaceTypeById(Map<Integer, Type> replacementsById) { + this.replacementsById = replacementsById; + } + + @Override + public Type schema(Schema schema, Type structResult) { + return structResult; + } + + @Override + public Type struct(Types.StructType struct, List<Type> fieldResults) { + List<Types.NestedField> fields = struct.fields(); + List<Types.NestedField> newFields = Lists.newArrayListWithExpectedSize(fields.size()); + boolean hasChanged = false; + + for (int i = 0; i < fields.size(); i += 1) { + Type fieldReplacement = fieldResults.get(i); + Types.NestedField field = fields.get(i); + if (field.type() != fieldReplacement) { + hasChanged = true; + newFields.add(Types.NestedField.from(field).ofType(fieldReplacement).build()); + } else { + newFields.add(field); + } + } + + if (hasChanged) { + return Types.StructType.of(newFields); + } + + return struct; + } + + @Override + public Type field(Types.NestedField field, Type fieldResult) { + return replacementsById.getOrDefault(field.fieldId(), fieldResult); + } + + @Override + public Type list(Types.ListType list, Type elementResult) { + Type elementReplacement = replacementsById.getOrDefault(list.elementId(), elementResult); + if (list.elementType() != elementReplacement) { + if (list.isElementRequired()) { + return Types.ListType.ofRequired(list.elementId(), elementReplacement); + } else { + return Types.ListType.ofOptional(list.elementId(), elementReplacement); + } + } + + return list; + } + + @Override + public Type map(Types.MapType map, Type keyResult, Type valueResult) { Review Comment: do we have a test that verifies replacement on maps? ########## api/src/main/java/org/apache/iceberg/types/TypeUtil.java: ########## @@ -157,6 +157,24 @@ public static Schema selectNot(Schema schema, Set<Integer> fieldIds) { return project(schema, projectedIds); } + /** + * Returns a copy of the schema with the type of each field in {@code replacementsById} replaced + * by its mapped type. Fields not in the map are unchanged. + * + * @param schema a schema + * @param replacementsById a map from field ID to the type that should replace the field's type + * @return a schema with the replaced field types + */ + public static Schema replaceFieldTypes(Schema schema, Map<Integer, Type> replacementsById) { + Types.StructType struct = visit(schema, new ReplaceTypeById(replacementsById)).asStructType(); + if (struct.equals(schema.asStruct())) { + return schema; + } + + return new Schema( + schema.schemaId(), struct.fields(), schema.getAliases(), schema.identifierFieldIds()); Review Comment: what happens if the replacement drops an identifier field or makes it optional/non-primitive? I think we should have some tests for this scenario -- 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]
