Spark: Add flag to handle missing files for importSparkTable

data/src/test/java/org/apache/iceberg/data/TestTableMigrationUtil.java

@@ -0,0 +1,110 @@
+/*
+ * 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.data;
+
+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.File;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.generic.GenericData;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.MetricsConfig;
+import org.apache.iceberg.PartitionSpec;
+import org.apache.iceberg.Schema;
+import org.apache.iceberg.avro.RandomAvroData;
+import org.apache.iceberg.io.FileAppender;
+import org.apache.iceberg.parquet.Parquet;
+import org.apache.iceberg.parquet.ParquetAvroWriter;
+import org.apache.iceberg.types.Types;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+class TestTableMigrationUtil {
+  private static final Schema SCHEMA =
+      new Schema(
+          optional(1, "id", Types.IntegerType.get()), optional(2, "data", Types.StringType.get()));
+  private static final PartitionSpec SPEC = PartitionSpec.builderFor(SCHEMA).identity("id").build();
+  private static final Map<String, String> PARTITION = Map.of("id", "1");
+  private static final String FORMAT = "parquet";
+  private static final Configuration CONF = new Configuration();
+
+  @TempDir protected Path tempTableLocation;
+
+  @Test
+  void testListPartition() throws IOException {
+    Path partitionPath = tempTableLocation.resolve("id=1");
+    String partitionUri = partitionPath.toUri().toString();
+    java.nio.file.Files.createDirectories(partitionPath);
+    writePartitionFile(partitionPath.toFile());
+
+    List<DataFile> dataFiles =
+        TableMigrationUtil.listPartition(
+            PARTITION, partitionUri, FORMAT, SPEC, CONF, MetricsConfig.getDefault(), null);
+    assertThat(dataFiles)
+        .as("List partition with 1 Parquet file should return 1 DataFile")
+        .hasSize(1);
+  }
+
+  @Test
+  void testListEmptyPartition() throws IOException {
+    Path partitionPath = tempTableLocation.resolve("id=1");
+    String partitionUri = partitionPath.toUri().toString();
+    java.nio.file.Files.createDirectories(partitionPath);
+
+    List<DataFile> dataFiles =
+        TableMigrationUtil.listPartition(
+            PARTITION, partitionUri, FORMAT, SPEC, CONF, MetricsConfig.getDefault(), null);
+    assertThat(dataFiles).as("List partition with 0 file should return 0 DataFile").isEmpty();
+  }
+
+  @Test
+  void testListPartitionMissingFilesFailure() {
+    String partitionUri = tempTableLocation.resolve("id=1").toUri().toString();
+
+    assertThatThrownBy(
+            () ->
+                TableMigrationUtil.listPartition(
+                    PARTITION, partitionUri, FORMAT, SPEC, CONF, MetricsConfig.getDefault(), null))
+        .hasMessageContaining("Unable to list files in partition: " + partitionUri)
+        .isInstanceOf(RuntimeException.class)
+        .hasRootCauseInstanceOf(FileNotFoundException.class);
+  }
+
+  private static void writePartitionFile(File outputDir) throws IOException {
+    Iterable<GenericData.Record> records = RandomAvroData.generate(SCHEMA, 1000, 54310);
+    File testFile = File.createTempFile("junit", ".parquet", outputDir);
+    assertThat(testFile.delete()).as("Delete should succeed").isTrue();
+
+    try (FileAppender<GenericData.Record> writer =
+        Parquet.write(Files.localOutput(testFile))
+            .schema(SCHEMA)
+            .createWriterFunc(ParquetAvroWriter::buildWriter)
+            .build()) {
+      writer.addAll(records);
+    }
+  }
+}

spark/v3.4/spark/src/main/java/org/apache/iceberg/spark/SparkTableUtil.java

@@ -20,6 +20,7 @@
 
 import static org.apache.spark.sql.functions.col;
 
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.Serializable;
 import java.net.URI;
@@ -105,6 +106,8 @@
 import org.apache.spark.sql.util.CaseInsensitiveStringMap;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import scala.Function2;
 import scala.Option;
 import scala.Some;
@@ -120,6 +123,7 @@
  * https://github.com/apache/iceberg/blob/apache-iceberg-0.8.0-incubating/spark/src/main/scala/org/apache/iceberg/spark/SparkTableUtil.scala
  */
 public class SparkTableUtil {
+  private static final Logger LOG = LoggerFactory.getLogger(SparkTableUtil.class);
 
   private static final String DUPLICATE_FILE_MESSAGE =
       "Cannot complete import because data files "
@@ -292,34 +296,26 @@ private static List<DataFile> listPartition(
       SerializableConfiguration conf,
       MetricsConfig metricsConfig,
       NameMapping mapping,
-      int parallelism) {
-    return TableMigrationUtil.listPartition(
-        partition.values,
-        partition.uri,
-        partition.format,
-        spec,
-        conf.get(),
-        metricsConfig,
-        mapping,
-        parallelism);
-  }
-
-  private static List<DataFile> listPartition(
-      SparkPartition partition,
-      PartitionSpec spec,
-      SerializableConfiguration conf,
-      MetricsConfig metricsConfig,
-      NameMapping mapping,
+      boolean ignoreMissingFiles,
       ExecutorService service) {
-    return TableMigrationUtil.listPartition(
-        partition.values,
-        partition.uri,
-        partition.format,
-        spec,
-        conf.get(),
-        metricsConfig,
-        mapping,
-        service);
+    try {
+      return TableMigrationUtil.listPartition(
+          partition.values,
+          partition.uri,
+          partition.format,
+          spec,
+          conf.get(),
+          metricsConfig,
+          mapping,
+          service);
+    } catch (RuntimeException e) {
+      if (ignoreMissingFiles && e.getCause() instanceof FileNotFoundException) {
+        LOG.warn("Ignoring FileNotFoundException when listing partition of {}", partition, e);
+        return Collections.emptyList();
+      } else {
+        throw e;
+      }
+    }
   }
 
   private static SparkPartition toSparkPartition(
@@ -524,6 +520,43 @@ public static void importSparkTable(
       Map<String, String> partitionFilter,
       boolean checkDuplicateFiles,
       ExecutorService service) {
+    importSparkTable(
+        spark,
+        sourceTableIdent,
+        targetTable,
+        stagingDir,
+        partitionFilter,
+        checkDuplicateFiles,
+        false,
+        service);
+  }
+
+  /**
+   * Import files from an existing Spark table to an Iceberg table.
+   *
+   * <p>The import uses the Spark session to get table metadata. It assumes no operation is going on
+   * the original and target table and thus is not thread-safe.
+   *
+   * @param spark a Spark session
+   * @param sourceTableIdent an identifier of the source Spark table
+   * @param targetTable an Iceberg table where to import the data
+   * @param stagingDir a staging directory to store temporary manifest files
+   * @param partitionFilter only import partitions whose values match those in the map, can be
+   *     partially defined
+   * @param checkDuplicateFiles if true, throw exception if import results in a duplicate data file
+   * @param ignoreMissingFiles if true, ignore {@link FileNotFoundException} when running {@link
+   *     #listPartition} for the Spark partitions
+   * @param service executor service to use for file reading
+   */
+  public static void importSparkTable(
+      SparkSession spark,
+      TableIdentifier sourceTableIdent,
+      Table targetTable,
+      String stagingDir,
+      Map<String, String> partitionFilter,
+      boolean checkDuplicateFiles,
+      boolean ignoreMissingFiles,
+      ExecutorService service) {
     SessionCatalog catalog = spark.sessionState().catalog();
 
     String db =
@@ -560,6 +593,7 @@ public static void importSparkTable(
               spec,
               stagingDir,
               checkDuplicateFiles,
+              ignoreMissingFiles,
               service);
         }
       }
@@ -744,6 +778,32 @@ public static void importSparkPartitions(
       String stagingDir,
       boolean checkDuplicateFiles,
       ExecutorService service) {
+    importSparkPartitions(
+        spark, partitions, targetTable, spec, stagingDir, checkDuplicateFiles, false, service);
+  }
+
+  /**
+   * Import files from given partitions to an Iceberg table.
+   *
+   * @param spark a Spark session
+   * @param partitions partitions to import
+   * @param targetTable an Iceberg table where to import the data
+   * @param spec a partition spec
+   * @param stagingDir a staging directory to store temporary manifest files
+   * @param checkDuplicateFiles if true, throw exception if import results in a duplicate data file
+   * @param ignoreMissingFiles if true, ignore {@link FileNotFoundException} when running {@link
+   *     #listPartition} for the Spark partitions
+   * @param service executor service to use for file reading
+   */
+  public static void importSparkPartitions(
+      SparkSession spark,
+      List<SparkPartition> partitions,
+      Table targetTable,
+      PartitionSpec spec,
+      String stagingDir,
+      boolean checkDuplicateFiles,
+      boolean ignoreMissingFiles,
+      ExecutorService service) {
     Configuration conf = spark.sessionState().newHadoopConf();
     SerializableConfiguration serializableConf = new SerializableConfiguration(conf);
     int listingParallelism =
@@ -771,6 +831,7 @@ public static void importSparkPartitions(
                             serializableConf,
                             metricsConfig,
                             nameMapping,
+                            ignoreMissingFiles,
                             service)
                         .iterator(),
             Encoders.javaSerialization(DataFile.class));

spark/v3.4/spark/src/test/java/org/apache/iceberg/spark/source/TestIcebergSourceTablesBase.java

@@ -26,11 +26,14 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import java.io.File;
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.nio.file.Path;
 import java.time.LocalDateTime;
 import java.time.ZoneOffset;
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
@@ -2170,6 +2173,90 @@ public void testTableWithInt96Timestamp() throws IOException {
     }
   }
 
+  @Test
+  public void testImportSparkTableWithMissingFilesFailure() throws IOException {
+    TableIdentifier tableIdentifier = TableIdentifier.of("db", "missing_files_test");
+    Table table = createTable(tableIdentifier, SCHEMA, SPEC);
+
+    File parquetTablePath = temp.newFolder("table_missing_files");
+    String parquetTableLocation = parquetTablePath.toURI().toString();
+    spark.sql(
+        String.format(
+            "CREATE TABLE parquet_table (data string, id int) "
+                + "USING parquet PARTITIONED BY (id) LOCATION '%s'",
+            parquetTableLocation));
+
+    List<SimpleRecord> records =
+        Lists.newArrayList(new SimpleRecord(1, "a"), new SimpleRecord(2, "b"));
+
+    Dataset<Row> inputDF = spark.createDataFrame(records, SimpleRecord.class);
+    inputDF.write().mode("overwrite").insertInto("parquet_table");
+
+    // Add a Spark partition of which location is missing
+    spark.sql("ALTER TABLE parquet_table ADD PARTITION (id = 1234)");
+    Path partitionLocationPath = parquetTablePath.toPath().resolve("id=1234");
+    java.nio.file.Files.delete(partitionLocationPath);
+
+    String stagingLocation = table.location() + "/metadata";
+
+    assertThatThrownBy(
+            () ->
+                SparkTableUtil.importSparkTable(
+                    spark,
+                    new org.apache.spark.sql.catalyst.TableIdentifier("parquet_table"),
+                    table,
+                    stagingLocation))
+        .hasMessageContaining(
+            "Unable to list files in partition: " + partitionLocationPath.toFile().toURI())
+        .isInstanceOf(SparkException.class)
+        .hasRootCauseInstanceOf(FileNotFoundException.class);
+  }
+
+  @Test
+  public void testImportSparkTableWithIgnoreMissingFilesEnabled() throws IOException {
+    TableIdentifier tableIdentifier = TableIdentifier.of("db", "missing_files_test");
+    Table table = createTable(tableIdentifier, SCHEMA, SPEC);
+
+    File parquetTableDir = temp.newFolder("table_missing_files");
+    String parquetTableLocation = parquetTableDir.toURI().toString();
+    spark.sql(
+        String.format(
+            "CREATE TABLE parquet_table (data string, id int) "
+                + "USING parquet PARTITIONED BY (id) LOCATION '%s'",
+            parquetTableLocation));
+
+    List<SimpleRecord> records =
+        Lists.newArrayList(new SimpleRecord(1, "a"), new SimpleRecord(2, "b"));
+
+    Dataset<Row> inputDF = spark.createDataFrame(records, SimpleRecord.class);
+    inputDF.write().mode("overwrite").insertInto("parquet_table");
+
+    // Add a Spark partition of which location is missing
+    spark.sql("ALTER TABLE parquet_table ADD PARTITION (id = 1234)");
+    Path partitionLocationPath = parquetTableDir.toPath().resolve("id=1234");
+    java.nio.file.Files.delete(partitionLocationPath);
+
+    String stagingLocation = table.location() + "/metadata";
+
+    SparkTableUtil.importSparkTable(
+        spark,
+        new org.apache.spark.sql.catalyst.TableIdentifier("parquet_table"),
+        table,
+        stagingLocation,
+        Collections.emptyMap(),
+        false,
+        true,
+        SparkTableUtil.migrationService(1));
+
+    List<Row> partitionsTableRows =
+        spark
+            .read()
+            .format("iceberg")
+            .load(loadLocation(tableIdentifier, "partitions"))
+            .collectAsList();
+    assertThat(partitionsTableRows).as("Partitions table should have 2 rows").hasSize(2);
+  }
+
   private void testWithFilter(String filterExpr, TableIdentifier tableIdentifier) {
     List<Row> expected =
         spark.table("parquet_table").select("tmp_col").filter(filterExpr).collectAsList();