rahulsmahadev commented on code in PR #17622: URL: https://github.com/apache/iceberg/pull/17622#discussion_r4140266753
########## spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/actions/TestRepairTableAction.java: ########## @@ -0,0 +1,1085 @@ +/* + * 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.spark.actions; + +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.Assumptions.assumeThat; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Files; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.ManifestFiles; +import org.apache.iceberg.ManifestWriter; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.Parameter; +import org.apache.iceberg.ParameterizedTestExtension; +import org.apache.iceberg.Parameters; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.actions.RepairTable; +import org.apache.iceberg.data.FileHelpers; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.BaseDVFileWriter; +import org.apache.iceberg.deletes.DVFileWriter; +import org.apache.iceberg.exceptions.CommitFailedException; +import org.apache.iceberg.exceptions.CommitStateUnknownException; +import org.apache.iceberg.exceptions.ValidationException; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.io.OutputFileFactory; +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.spark.TestBase; +import org.apache.iceberg.spark.source.ThreeColumnRecord; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Pair; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.junit.jupiter.api.BeforeEach; +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 TestRepairTableAction extends TestBase { + + private static final HadoopTables TABLES = new HadoopTables(new Configuration()); + private static final Schema SCHEMA = + new Schema( + optional(1, "c1", Types.IntegerType.get()), + optional(2, "c2", Types.StringType.get()), + optional(3, "c3", Types.StringType.get())); + + @Parameters(name = "formatVersion = {0}") + public static Object[] parameters() { + return new Object[][] {new Object[] {1}, new Object[] {2}, new Object[] {3}}; + } + + @Parameter private int formatVersion; + + private String tableLocation = null; + + @TempDir private Path temp; + @TempDir private File tableDir; + + @BeforeEach + public void setupTableLocation() { + this.tableLocation = tableDir.toURI().toString(); + } + + @TestTemplate + public void testRepairEmptyTable() { + Table table = createTable(PartitionSpec.unpartitioned()); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + } + + @TestTemplate + public void testRepairTableWithCorrectStats() { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + Snapshot before = table.currentSnapshot(); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + + table.refresh(); + assertThat(table.currentSnapshot().snapshotId()) + .as("should not commit a snapshot when nothing is repaired") + .isEqualTo(before.snapshotId()); + } + + @TestTemplate + public void testNoRepairSelectedIsNoOp() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + replaceManifestWithCorruptStats(table, original); + + table.refresh(); + Snapshot before = table.currentSnapshot(); + DataFile corrupt = onlyDataFile(table); + + // no repair was selected, so execute() must do nothing even though the stats are incorrect + RepairTable.Result result = SparkActions.get().repairTable(table).execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + + table.refresh(); + assertThat(table.currentSnapshot().snapshotId()) + .as("a repair with nothing selected must not commit") + .isEqualTo(before.snapshotId()); + assertThat(onlyDataFile(table).recordCount()) + .as("a repair with nothing selected must leave the incorrect stats in place") + .isEqualTo(corrupt.recordCount()); + } + + @TestTemplate + public void testRepairIncorrectRecordCountAndFileSize() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + List<Object[]> expectedRows = currentRows(); + DataFile original = onlyDataFile(table); + + // replace the manifest with one whose entry records a wrong record count and file size + replaceManifestWithCorruptStats(table, original); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedEntryCount()).isEqualTo(1); + assertThat(result.repairedManifests()).hasSize(1); + + table.refresh(); + DataFile repaired = onlyDataFile(table); + assertThat(repaired.recordCount()).isEqualTo(original.recordCount()); + assertThat(repaired.fileSizeInBytes()).isEqualTo(original.fileSizeInBytes()); + assertThat(repaired.location()).isEqualTo(original.location()); + + assertThat(currentRows()) + .as("table contents must be unchanged by the repair") + .containsExactlyInAnyOrderElementsOf(expectedRows); + } + + @TestTemplate + public void testRepairPreservesEntryLineage() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + List<Row> lineageBefore = entryLineage(); + + replaceManifestWithCorruptStats(table, original); + + SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + table.refresh(); + assertThat(entryLineage()) + .as("snapshot id and sequence numbers must be carried through the repair") + .containsExactlyInAnyOrderElementsOf(lineageBefore); + } + + @TestTemplate + public void testDryRunDoesNotCommit() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + replaceManifestWithCorruptStats(table, original); + + table.refresh(); + Snapshot before = table.currentSnapshot(); + DataFile corrupt = onlyDataFile(table); + + RepairTable.Result result = + SparkActions.get().repairTable(table).repairFileMetrics().dryRun().execute(); + + assertThat(result.repairedEntryCount()).isEqualTo(1); + assertThat(result.repairedManifests()).hasSize(1); + + table.refresh(); + assertThat(table.currentSnapshot().snapshotId()) + .as("dry run must not commit") + .isEqualTo(before.snapshotId()); + assertThat(onlyDataFile(table).recordCount()) + .as("dry run must leave the incorrect stats in place") + .isEqualTo(corrupt.recordCount()); + } + + @TestTemplate + public void testRepairOnlyRewritesAffectedManifests() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(2)); + appendRecords(table, records(2)); + + table.refresh(); + assertThat(table.currentSnapshot().dataManifests(table.io())).hasSize(2); + + List<ManifestFile> manifests = table.currentSnapshot().dataManifests(table.io()); + ManifestFile untouched = manifests.get(1); + + // corrupt the entry of one manifest only + DataFile fileToCorrupt = readDataFiles(table, manifests.get(0)).get(0); + corruptStats(table, manifests.get(0), fileToCorrupt.location()); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedManifests()).hasSize(1); + assertThat(result.repairedEntryCount()).isEqualTo(1); + + table.refresh(); + assertThat(table.currentSnapshot().dataManifests(table.io())) + .as("the manifest without incorrect entries must be left in place") + .anyMatch(manifest -> manifest.path().equals(untouched.path())); + } + + @TestTemplate + public void testRepairPartitionedTable() throws IOException { + Table table = createTable(PartitionSpec.builderFor(SCHEMA).identity("c1").build()); + + Dataset<Row> df = + spark + .createDataFrame( + Lists.newArrayList( + new ThreeColumnRecord(1, "AAAA", "A"), new ThreeColumnRecord(2, "BBBB", "B")), + ThreeColumnRecord.class) + .coalesce(1); + df.select("c1", "c2", "c3").write().format("iceberg").mode("append").save(tableLocation); + + table.refresh(); + List<Object[]> expectedRows = currentRows(); + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + List<DataFile> files = readDataFiles(table, manifest); + + corruptStats(table, manifest, files.get(0).location()); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedEntryCount()).isEqualTo(1); + assertThat(currentRows()).containsExactlyInAnyOrderElementsOf(expectedRows); + } + + @TestTemplate + public void testRepairSkipsColumnMetricsByDefault() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + + // only the column level statistics are wrong, the record count and the file size are correct + ManifestFile manifest = table.currentSnapshot().dataManifests(table.io()).get(0); + corruptStats(table, manifest, original.location(), false); + + RepairTable.Result skipped = + SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(skipped.repairedEntryCount()) + .as("column metrics must not be compared by default") + .isEqualTo(0); + + RepairTable.Result repaired = + SparkActions.get() + .repairTable(table) + .repairFileMetrics() + .option(RepairTableSparkAction.REPAIR_COLUMN_METRICS, "true") + .execute(); + + assertThat(repaired.repairedEntryCount()) + .as("column metrics are compared when enabled") + .isEqualTo(1); + } + + @TestTemplate + void repairWithCustomColumnMetrics() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + table + .updateProperties() + .set(TableProperties.METRICS_MODE_COLUMN_CONF_PREFIX + "c2", "none") + .commit(); + appendRecords(table, records(4)); + + List<Object[]> expectedRows = currentRows(); + DataFile original = onlyDataFile(table); + int columnId = table.schema().findField("c2").fieldId(); + assertThat(original.valueCounts()).doesNotContainKey(columnId); + assertThat(original.lowerBounds()).doesNotContainKey(columnId); + assertThat(original.upperBounds()).doesNotContainKey(columnId); + + replaceManifestWithCorruptStats(table, original); + + RepairTable.Result result = + SparkActions.get() + .repairTable(table) + .repairFileMetrics() + .option(RepairTableSparkAction.REPAIR_COLUMN_METRICS, "true") + .execute(); + + assertThat(result.repairedEntryCount()).isEqualTo(1); + assertThat(result.repairedManifests()).hasSize(1); + + DataFile repaired = onlyDataFile(table); + assertThat(repaired.recordCount()).isEqualTo(original.recordCount()); + assertThat(repaired.fileSizeInBytes()).isEqualTo(original.fileSizeInBytes()); + assertThat(repaired.columnSizes()).isEqualTo(original.columnSizes()); + assertThat(repaired.valueCounts()).isEqualTo(original.valueCounts()); + assertThat(repaired.nullValueCounts()).isEqualTo(original.nullValueCounts()); + assertThat(repaired.lowerBounds()).isEqualTo(original.lowerBounds()); + assertThat(repaired.upperBounds()).isEqualTo(original.upperBounds()); + assertThat(currentRows()).containsExactlyInAnyOrderElementsOf(expectedRows); Review Comment: cool added an idempotenc check ########## spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/actions/TestRepairTableAction.java: ########## @@ -0,0 +1,1085 @@ +/* + * 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.spark.actions; + +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.Assumptions.assumeThat; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.ContentFile; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileContent; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.FileMetadata; +import org.apache.iceberg.Files; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.ManifestFiles; +import org.apache.iceberg.ManifestWriter; +import org.apache.iceberg.Metrics; +import org.apache.iceberg.Parameter; +import org.apache.iceberg.ParameterizedTestExtension; +import org.apache.iceberg.Parameters; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Snapshot; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; +import org.apache.iceberg.actions.RepairTable; +import org.apache.iceberg.data.FileHelpers; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.deletes.BaseDVFileWriter; +import org.apache.iceberg.deletes.DVFileWriter; +import org.apache.iceberg.exceptions.CommitFailedException; +import org.apache.iceberg.exceptions.CommitStateUnknownException; +import org.apache.iceberg.exceptions.ValidationException; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.io.OutputFileFactory; +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.spark.TestBase; +import org.apache.iceberg.spark.source.ThreeColumnRecord; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Pair; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.junit.jupiter.api.BeforeEach; +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 TestRepairTableAction extends TestBase { + + private static final HadoopTables TABLES = new HadoopTables(new Configuration()); + private static final Schema SCHEMA = + new Schema( + optional(1, "c1", Types.IntegerType.get()), + optional(2, "c2", Types.StringType.get()), + optional(3, "c3", Types.StringType.get())); + + @Parameters(name = "formatVersion = {0}") + public static Object[] parameters() { + return new Object[][] {new Object[] {1}, new Object[] {2}, new Object[] {3}}; + } + + @Parameter private int formatVersion; + + private String tableLocation = null; + + @TempDir private Path temp; + @TempDir private File tableDir; + + @BeforeEach + public void setupTableLocation() { + this.tableLocation = tableDir.toURI().toString(); + } + + @TestTemplate + public void testRepairEmptyTable() { + Table table = createTable(PartitionSpec.unpartitioned()); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + } + + @TestTemplate + public void testRepairTableWithCorrectStats() { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + Snapshot before = table.currentSnapshot(); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + + table.refresh(); + assertThat(table.currentSnapshot().snapshotId()) + .as("should not commit a snapshot when nothing is repaired") + .isEqualTo(before.snapshotId()); + } + + @TestTemplate + public void testNoRepairSelectedIsNoOp() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + replaceManifestWithCorruptStats(table, original); + + table.refresh(); + Snapshot before = table.currentSnapshot(); + DataFile corrupt = onlyDataFile(table); + + // no repair was selected, so execute() must do nothing even though the stats are incorrect + RepairTable.Result result = SparkActions.get().repairTable(table).execute(); + + assertThat(result.repairedManifests()).isEmpty(); + assertThat(result.repairedEntryCount()).isEqualTo(0); + + table.refresh(); + assertThat(table.currentSnapshot().snapshotId()) + .as("a repair with nothing selected must not commit") + .isEqualTo(before.snapshotId()); + assertThat(onlyDataFile(table).recordCount()) + .as("a repair with nothing selected must leave the incorrect stats in place") + .isEqualTo(corrupt.recordCount()); + } + + @TestTemplate + public void testRepairIncorrectRecordCountAndFileSize() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + List<Object[]> expectedRows = currentRows(); + DataFile original = onlyDataFile(table); + + // replace the manifest with one whose entry records a wrong record count and file size + replaceManifestWithCorruptStats(table, original); + + RepairTable.Result result = SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + assertThat(result.repairedEntryCount()).isEqualTo(1); + assertThat(result.repairedManifests()).hasSize(1); + + table.refresh(); + DataFile repaired = onlyDataFile(table); + assertThat(repaired.recordCount()).isEqualTo(original.recordCount()); + assertThat(repaired.fileSizeInBytes()).isEqualTo(original.fileSizeInBytes()); + assertThat(repaired.location()).isEqualTo(original.location()); + + assertThat(currentRows()) + .as("table contents must be unchanged by the repair") + .containsExactlyInAnyOrderElementsOf(expectedRows); + } + + @TestTemplate + public void testRepairPreservesEntryLineage() throws IOException { + Table table = createTable(PartitionSpec.unpartitioned()); + appendRecords(table, records(4)); + + DataFile original = onlyDataFile(table); + List<Row> lineageBefore = entryLineage(); + + replaceManifestWithCorruptStats(table, original); + + SparkActions.get().repairTable(table).repairFileMetrics().execute(); + + table.refresh(); + assertThat(entryLineage()) Review Comment: added -- 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]
