anoopj commented on code in PR #17541: URL: https://github.com/apache/iceberg/pull/17541#discussion_r4020227000
########## core/src/test/java/org/apache/iceberg/TestFilePlanner.java: ########## @@ -0,0 +1,569 @@ +/* + * 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 static org.assertj.core.api.Assertions.tuple; + +import java.io.IOException; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.function.UnaryOperator; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.inmemory.InMemoryFileIO; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.FileAppender; +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.apache.iceberg.util.LocationUtil; +import org.apache.iceberg.util.ThreadPools; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.FieldSource; + +class TestFilePlanner { + private static final long SNAPSHOT_ID = 42L; + private static final int WRITER_FORMAT_VERSION = 4; + private static final long RECORD_COUNT = 100L; + private static final long FILE_SIZE_IN_BYTES = 1024L; + private static final String TABLE_LOCATION = "s3://bucket/db/table"; + private static final String DV_LOCATION = "s3://bucket/db/table/dv.puffin"; + private static final long DV_OFFSET = 100L; + private static final long DV_SIZE_IN_BYTES = 50L; + private static final long DV_CARDINALITY = 5L; + + 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<FileFormat> MANIFEST_FORMATS = + ImmutableList.of(FileFormat.AVRO, FileFormat.PARQUET); + + private final InMemoryFileIO fileIO = new InMemoryFileIO(); + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void rootWithDirectDataEntries(FileFormat format) throws IOException { + InputFile root = + writeManifest( + format, + EMPTY_PARTITION, + ImmutableList.of( + dataFile("a.parquet", EMPTY_PARTITION_DATA), + dataFile("b.parquet", EMPTY_PARTITION_DATA))); + + List<FileScanTask> tasks = plan(root, UNPARTITIONED_SPECS); + + assertThat(tasks) + .hasSize(2) + .extracting(task -> task.file().location()) + .containsExactlyInAnyOrder(resolved("a.parquet"), resolved("b.parquet")); + assertThat(tasks).allSatisfy(task -> assertThat(task.deletes()).isEmpty()); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void rootWithLeafManifestEntriesOnly(FileFormat format) throws IOException { + InputFile leaf = + writeManifest( + format, + EMPTY_PARTITION, + ImmutableList.of( + dataFile("leaf-a.parquet", EMPTY_PARTITION_DATA), + dataFile("leaf-b.parquet", EMPTY_PARTITION_DATA))); + InputFile root = + writeManifest(format, EMPTY_PARTITION, ImmutableList.of(dataManifest(leaf.location()))); + + ScanMetrics metrics = ScanMetrics.of(new DefaultMetricsContext()); + List<FileScanTask> tasks = + plan(root, UNPARTITIONED_SPECS, planner -> planner.scanMetrics(metrics)); + + assertThat(tasks) + .hasSize(2) + .extracting(task -> task.file().location()) + .containsExactlyInAnyOrder(resolved("leaf-a.parquet"), resolved("leaf-b.parquet")); + assertThat(metrics.scannedDataManifests().value()) + .as("the root and the expanded leaf are each counted as a scanned manifest") + .isEqualTo(2L); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void mixedRootDataAndLeafManifest(FileFormat format) throws IOException { + InputFile leaf = + writeManifest( + format, + EMPTY_PARTITION, + ImmutableList.of(dataFile("leaf-data-file.parquet", EMPTY_PARTITION_DATA))); + InputFile root = + writeManifest( + format, + EMPTY_PARTITION, + ImmutableList.of( + dataFile("root-data-file.parquet", EMPTY_PARTITION_DATA), + dataManifest(leaf.location()))); + + List<FileScanTask> tasks = plan(root, UNPARTITIONED_SPECS); + + assertThat(tasks) + .hasSize(2) + .extracting(task -> task.file().location()) + .containsExactlyInAnyOrder( + resolved("root-data-file.parquet"), resolved("leaf-data-file.parquet")); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void deletionVectorAttachedToTask(FileFormat format) throws IOException { + TrackedFile fileWithDv = + dataFile( + "with-dv.parquet", + EMPTY_PARTITION_DATA, + deletionVector(DV_LOCATION, DV_OFFSET, DV_SIZE_IN_BYTES, DV_CARDINALITY)); + InputFile root = writeManifest(format, EMPTY_PARTITION, ImmutableList.of(fileWithDv)); + + List<FileScanTask> tasks = plan(root, UNPARTITIONED_SPECS); + + assertThat(tasks).hasSize(1); + assertThat(tasks.get(0).deletes()) + .hasSize(1) + .allSatisfy( + delete -> { + assertThat(delete.content()).isEqualTo(FileContent.POSITION_DELETES); + assertThat(delete.location()).isEqualTo(DV_LOCATION); + assertThat(delete.referencedDataFile()).isEqualTo(resolved("with-dv.parquet")); + assertThat(delete.recordCount()).isEqualTo(DV_CARDINALITY); + assertThat(delete.contentOffset()).isEqualTo(DV_OFFSET); + assertThat(delete.contentSizeInBytes()).isEqualTo(DV_SIZE_IN_BYTES); + }); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void residualAttachedFromFilter(FileFormat format) throws IOException { + InputFile root = + writeManifest( + format, PARTITION_TYPE, ImmutableList.of(dataFile("keep.parquet", partition(1)))); + + List<FileScanTask> withResidual = + plan( + root, PARTITIONED_SPECS, planner -> planner.filterData(Expressions.equal("data", "x"))); + assertThat(withResidual.get(0).residual()) + .hasToString(Expressions.equal("data", "x").toString()); + + List<FileScanTask> ignored = + plan( + root, + PARTITIONED_SPECS, + planner -> planner.filterData(Expressions.equal("data", "x")).ignoreResiduals()); + assertThat(ignored.get(0).residual()).isEqualTo(Expressions.alwaysTrue()); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void sameSpecFilesShareContext(FileFormat format) throws IOException { + InputFile root = + writeManifest( + format, + PARTITION_TYPE, + ImmutableList.of( + dataFile("a.parquet", partition(1)), dataFile("b.parquet", partition(1)))); + + List<FileScanTask> tasks = + plan( + root, PARTITIONED_SPECS, planner -> planner.filterData(Expressions.equal("data", "x"))); + + // both files share spec 0, so both tasks resolve from the same shared context + assertThat(tasks).hasSize(2); + FileScanTask first = tasks.get(0); + FileScanTask second = tasks.get(1); + assertThat(second.schema()).isSameAs(first.schema()); + assertThat(second.spec()).isSameAs(first.spec()); + assertThat(second.residual()).isSameAs(first.residual()); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void buildingAndClosingWithoutIteratingDoesNotScanLeaves(FileFormat format) throws IOException { + InputFile leaf = + writeManifest( + format, + EMPTY_PARTITION, + ImmutableList.of(dataFile("leaf.parquet", EMPTY_PARTITION_DATA))); + InputFile root = + writeManifest(format, EMPTY_PARTITION, ImmutableList.of(dataManifest(leaf.location()))); + + ScanMetrics metrics = ScanMetrics.of(new DefaultMetricsContext()); + FilePlanner planner = + FilePlanner.builder(fileIO, asManifest(root), UNPARTITIONED_SPECS) + .tableLocation(TABLE_LOCATION) + .scanMetrics(metrics) + .build(); + planner.planFiles().close(); + + assertThat(metrics.scannedDataManifests().value()) + .as( + "the root is read eagerly but leaf readers open lazily, so closing " + + "without iterating scans only the root, not the leaf") + .isEqualTo(1L); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void parallelPlanningMatchesSequential(FileFormat format) throws IOException { + TrackedFile withDv = + dataFile( + "leaf1-data.parquet", + partition(1), + deletionVector(DV_LOCATION, DV_OFFSET, DV_SIZE_IN_BYTES, DV_CARDINALITY)); + InputFile leaf1 = writeManifest(format, PARTITION_TYPE, ImmutableList.of(withDv)); + InputFile leaf2 = + writeManifest( + format, PARTITION_TYPE, ImmutableList.of(dataFile("leaf2-data.parquet", partition(1)))); + InputFile root = + writeManifest( + format, + PARTITION_TYPE, + ImmutableList.of(dataManifest(leaf1.location()), dataManifest(leaf2.location()))); + + List<FileScanTask> sequential = + plan( + root, PARTITIONED_SPECS, planner -> planner.filterData(Expressions.equal("data", "x"))); + + ExecutorService pool = ThreadPools.newFixedThreadPool("test-scan-task-planner", 2); + try { + List<FileScanTask> parallel = + plan( + root, + PARTITIONED_SPECS, + planner -> planner.filterData(Expressions.equal("data", "x")).planWith(pool)); + assertThat(parallel) + .extracting( + task -> task.file().location(), + task -> task.residual().toString(), + task -> task.deletes().size()) + .containsExactlyInAnyOrderElementsOf( + Lists.transform( + sequential, + task -> + tuple( + task.file().location(), + task.residual().toString(), + task.deletes().size()))); + } finally { + pool.shutdownNow(); + } + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void deleteContentInRootIsUnsupported(FileFormat format) throws IOException { + InputFile root = + writeManifest(format, EMPTY_PARTITION, ImmutableList.of(deleteManifest("deletes.avro"))); + + FilePlanner planner = + FilePlanner.builder(fileIO, asManifest(root), UNPARTITIONED_SPECS) + .tableLocation(TABLE_LOCATION) + .build(); + assertThatThrownBy(() -> Lists.newArrayList(planner.planFiles())) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessage("v3 and earlier deletes are not yet supported"); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void nonDataEntryInLeafFailsPlanning(FileFormat format) throws IOException { + InputFile leaf = + writeManifest( + format, EMPTY_PARTITION, ImmutableList.of(dataManifest("nested-leaf.parquet"))); + InputFile root = + writeManifest(format, EMPTY_PARTITION, ImmutableList.of(dataManifest(leaf.location()))); + + FilePlanner planner = + FilePlanner.builder(fileIO, asManifest(root), UNPARTITIONED_SPECS) + .tableLocation(TABLE_LOCATION) + .build(); + assertThatThrownBy(() -> Lists.newArrayList(planner.planFiles())) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("Invalid content type for DataFile: DATA_MANIFEST"); + } + + @Test + void emptyRootYieldsNoTasks() throws IOException { + // Avro is used because the Parquet writer does not materialize a file when no records are + // appended. + InputFile root = writeManifest(FileFormat.AVRO, EMPTY_PARTITION, ImmutableList.of()); + + assertThat(plan(root, UNPARTITIONED_SPECS)).isEmpty(); + } + + @ParameterizedTest + @FieldSource("MANIFEST_FORMATS") + void planFilesAcrossMultipleSpecs(FileFormat format) throws IOException { Review Comment: This test will fail until the bug fix https://github.com/apache/iceberg/pull/18108/ is merged. -- 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]
