Core: Add deleteFiles to FileIO interface. For S3FileIo implement batch deletion via S3#RemoveObjects API

Author: amogh-jahagirdar

api/src/main/java/org/apache/iceberg/io/FileIO.java

@@ -68,6 +68,15 @@ default void deleteFile(OutputFile file) {
   default void initialize(Map<String, String> properties) {
   }
 
+  /**
+   * Delete the files at the given path.
+   */
+  default void deleteFiles(Iterable<String> pathsToDelete) {
+    for (String path : pathsToDelete) {
+      deleteFile(path);
+    }
+  }
+
   /**
    * Close File IO to release underlying resources.
    * <p>

aws/src/integration/java/org/apache/iceberg/aws/s3/TestS3FileIOIntegration.java

@@ -26,6 +26,7 @@
 import java.security.MessageDigest;
 import java.security.SecureRandom;
 import java.util.Base64;
+import java.util.List;
 import java.util.UUID;
 import javax.crypto.KeyGenerator;
 import javax.crypto.SecretKey;
@@ -36,6 +37,7 @@
 import org.apache.iceberg.aws.AwsProperties;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.OutputFile;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.junit.AfterClass;
 import org.junit.Assert;
 import org.junit.Before;
@@ -67,6 +69,7 @@ public class TestS3FileIOIntegration {
   private static byte[] contentBytes;
   private static String content;
   private static String kmsKeyArn;
+  private static int deletionBatchSize;
   private String objectKey;
   private String objectUri;
 
@@ -78,6 +81,7 @@ public static void beforeClass() {
     bucketName = AwsIntegTestUtil.testBucketName();
     prefix = UUID.randomUUID().toString();
     contentBytes = new byte[1024 * 1024 * 10];
+    deletionBatchSize = 3;
     content = new String(contentBytes, StandardCharsets.UTF_8);
     kmsKeyArn = kms.createKey().keyMetadata().arn();
   }
@@ -208,8 +212,49 @@ public void testClientFactorySerialization() throws Exception {
     validateRead(fileIO2);
   }
 
+  @Test
+  public void testDeleteMultipleBatches() throws Exception {
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3, getDeletionTestProperties());
+    testBatchDelete(6, s3FileIO);
+  }
+
+  @Test
+  public void testDeleteLessThanBatchSize() throws Exception {
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3, getDeletionTestProperties());
+    testBatchDelete(2, s3FileIO);
+  }
+
+  @Test
+  public void testDeleteSingleBatchWithRemainder() throws Exception {
+    S3FileIO s3FileIO = new S3FileIO(clientFactory::s3, getDeletionTestProperties());
+    testBatchDelete(5, s3FileIO);
+  }
+
+  private AwsProperties getDeletionTestProperties() {
+    AwsProperties properties = new AwsProperties();
+    properties.setS3FileIoDeleteBatchSize(deletionBatchSize);
+    return properties;
+  }
+
+  private void testBatchDelete(int numObjects, S3FileIO s3FileIO) throws Exception {
+    List<String> paths = Lists.newArrayList();
+    for (int i = 1; i <= numObjects; i++) {
+      String deletionKey = objectKey + "-deletion-" + i;
+      write(s3FileIO, String.format("s3://%s/%s", bucketName, deletionKey));
+      paths.add(String.format("s3://%s/%s", bucketName, deletionKey));
+    }
+    s3FileIO.deleteFiles(paths);
+    for (String path : paths) {
+      Assert.assertFalse(s3FileIO.newInputFile(path).exists());
+    }
+  }
+
   private void write(S3FileIO s3FileIO) throws Exception {
-    OutputFile outputFile = s3FileIO.newOutputFile(objectUri);
+    write(s3FileIO, objectUri);
+  }
+
+  private void write(S3FileIO s3FileIO, String uri) throws Exception {
+    OutputFile outputFile = s3FileIO.newOutputFile(uri);
     OutputStream outputStream = outputFile.create();
     IoUtils.copy(new ByteArrayInputStream(contentBytes), outputStream);
     outputStream.close();

aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java

@@ -171,6 +171,11 @@ public class AwsProperties implements Serializable {
   public static final String S3_CHECKSUM_ENABLED = "s3.checksum-enabled";
   public static final boolean S3_CHECKSUM_ENABLED_DEFAULT = false;
 
+  /**
+   * Configure the batch size used when deleting multiple files from a given S3 bucket
+   */
+  public static final String S3FILEIO_DELETE_BATCH_SIZE = "s3.delete.batch-size";
+
   /**
    * DynamoDB table name for {@link DynamoDbCatalog}
    */
@@ -221,11 +226,17 @@ public class AwsProperties implements Serializable {
   @Deprecated
   public static final boolean CLIENT_ENABLE_ETAG_CHECK_DEFAULT = false;
 
+  /**
+   * Sets the size for batch deletion.
+   */
+  public static final int S3FILEIO_DELETE_BATCH_SIZE_DEFAULT = 250;
+
   private String s3FileIoSseType;
   private String s3FileIoSseKey;
   private String s3FileIoSseMd5;
   private int s3FileIoMultipartUploadThreads;
   private int s3FileIoMultiPartSize;
+  private int s3FileIoDeleteBatchSize;
   private double s3FileIoMultipartThresholdFactor;
   private String s3fileIoStagingDirectory;
   private ObjectCannedACL s3FileIoAcl;
@@ -246,6 +257,7 @@ public AwsProperties() {
     this.s3FileIoMultipartUploadThreads = Runtime.getRuntime().availableProcessors();
     this.s3FileIoMultiPartSize = S3FILEIO_MULTIPART_SIZE_DEFAULT;
     this.s3FileIoMultipartThresholdFactor = S3FILEIO_MULTIPART_THRESHOLD_FACTOR_DEFAULT;
+    this.s3FileIoDeleteBatchSize = S3FILEIO_DELETE_BATCH_SIZE_DEFAULT;
     this.s3fileIoStagingDirectory = System.getProperty("java.io.tmpdir");
     this.isS3ChecksumEnabled = S3_CHECKSUM_ENABLED_DEFAULT;
 
@@ -300,6 +312,11 @@ public AwsProperties(Map<String, String> properties) {
     this.isS3ChecksumEnabled = PropertyUtil.propertyAsBoolean(properties, S3_CHECKSUM_ENABLED,
         S3_CHECKSUM_ENABLED_DEFAULT);
 
+    this.s3FileIoDeleteBatchSize = PropertyUtil.propertyAsInt(properties, S3FILEIO_DELETE_BATCH_SIZE,
+        S3FILEIO_DELETE_BATCH_SIZE_DEFAULT);
+    Preconditions.checkArgument(s3FileIoDeleteBatchSize > 0 && s3FileIoDeleteBatchSize < 1000,
+        "Deletion batch size must be between 1 and 1000");
+
     this.dynamoDbTableName = PropertyUtil.propertyAsString(properties, DYNAMODB_TABLE_NAME,
         DYNAMODB_TABLE_NAME_DEFAULT);
   }
@@ -316,6 +333,14 @@ public String s3FileIoSseKey() {
     return s3FileIoSseKey;
   }
 
+  public int s3FileIoDeleteBatchSize() {
+    return s3FileIoDeleteBatchSize;
+  }
+
+  public void setS3FileIoDeleteBatchSize(int deleteBatchSize) {
+    this.s3FileIoDeleteBatchSize = deleteBatchSize;
+  }
+
   public void setS3FileIoSseKey(String sseKey) {
     this.s3FileIoSseKey = sseKey;
   }

aws/src/main/java/org/apache/iceberg/aws/s3/S3BatchDeletionException.java

@@ -0,0 +1,26 @@
+/*
+ * 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.aws.s3;
+
+public class S3BatchDeletionException extends RuntimeException {
+  public S3BatchDeletionException(String message) {
+    super(message);
+  }
+}

aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java

@@ -19,17 +19,31 @@
 
 package org.apache.iceberg.aws.s3;
 
+import java.util.Collection;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
 import org.apache.iceberg.aws.AwsClientFactories;
 import org.apache.iceberg.aws.AwsClientFactory;
 import org.apache.iceberg.aws.AwsProperties;
 import org.apache.iceberg.io.FileIO;
 import org.apache.iceberg.io.InputFile;
 import org.apache.iceberg.io.OutputFile;
+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.Multimaps;
+import org.apache.iceberg.relocated.com.google.common.collect.SetMultimap;
+import org.apache.iceberg.relocated.com.google.common.collect.Sets;
 import org.apache.iceberg.util.SerializableSupplier;
 import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.Delete;
 import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
+import software.amazon.awssdk.services.s3.model.DeleteObjectsRequest;
+import software.amazon.awssdk.services.s3.model.DeleteObjectsResponse;
+import software.amazon.awssdk.services.s3.model.ObjectIdentifier;
+import software.amazon.awssdk.services.s3.model.S3Error;
 
 /**
  * FileIO implementation backed by S3.
@@ -94,6 +108,66 @@ public void deleteFile(String path) {
     client().deleteObject(deleteRequest);
   }
 
+  /**
+   * Deletes the given paths in a batched manner.
+   * <p>
+   * The paths are grouped by bucket, and deletion is triggered when we either reach the configured batch size
+   * or have a final remainder batch for each bucket.
+   *
+   * @param paths paths to delete
+   */
+  @Override
+  public void deleteFiles(Iterable<String> paths) {
+    SetMultimap<String, String> bucketToObjects = Multimaps.newSetMultimap(Maps.newHashMap(), Sets::newHashSet);
+    List<String> failedDeletions = Lists.newArrayList();
+    for (String path : paths) {
+      S3URI location = new S3URI(path);
+      String bucket = location.bucket();
+      String objectKey = location.key();
+      Set<String> objectsInBucket = bucketToObjects.get(bucket);
+      if (objectsInBucket.size() == awsProperties.s3FileIoDeleteBatchSize()) {
+        List<String> failedDeletionsForBatch = deleteObjectsInBucket(bucket, objectsInBucket);
+        failedDeletions.addAll(failedDeletionsForBatch);
+        bucketToObjects.removeAll(bucket);
+      }
+      bucketToObjects.get(bucket).add(objectKey);
+    }
+    // Delete the remainder
+    List<List<String>> remainderFailedObjects = bucketToObjects
+        .asMap()
+        .entrySet()
+        .stream()
+        .map(entry -> deleteObjectsInBucket(entry.getKey(), entry.getValue()))
+        .collect(Collectors.toList());
+
+    remainderFailedObjects.forEach(failedDeletions::addAll);
+    if (!failedDeletions.isEmpty()) {
+      throw new S3BatchDeletionException(String.format("Failed to delete %d objects. Failed objects: %s",
+              failedDeletions.size(), failedDeletions));
+    }
+  }
+
+  private List<String> deleteObjectsInBucket(String bucket, Collection<String> objects) {
+    if (!objects.isEmpty()) {
+      List<ObjectIdentifier> objectIds = objects
+          .stream()
+          .map(objectKey -> ObjectIdentifier.builder().key(objectKey).build())
+          .collect(Collectors.toList());
+      DeleteObjectsRequest deleteObjectsRequest = DeleteObjectsRequest.builder()
+          .bucket(bucket)
+          .delete(Delete.builder().objects(objectIds).build())
+          .build();
+      DeleteObjectsResponse response = client().deleteObjects(deleteObjectsRequest);
+      if (response.hasErrors()) {
+        return response.errors()
+            .stream()
+            .map(S3Error::key)
+            .collect(Collectors.toList());
+      }
+    }
+    return Lists.newArrayList();
+  }
+
   private S3Client client() {
     if (client == null) {
       client = s3.get();

aws/src/test/java/org/apache/iceberg/aws/TestAwsProperties.java

@@ -97,4 +97,24 @@ public void testS3MultipartThresholdFactorLessThanOne() {
         () -> new AwsProperties(map));
   }
 
+  @Test
+  public void testS3FileIoDeleteBatchSizeTooLarge() {
+    Map<String, String> map = Maps.newHashMap();
+    map.put(AwsProperties.S3FILEIO_DELETE_BATCH_SIZE, "2000");
+    AssertHelpers.assertThrows("should not accept batch size greater than 1000",
+            IllegalArgumentException.class,
+            "Deletion batch size must be between 1 and 1000",
+            () -> new AwsProperties(map));
+  }
+
+  @Test
+  public void testS3FileIoDeleteBatchSizeTooSmall() {
+    Map<String, String> map = Maps.newHashMap();
+    map.put(AwsProperties.S3FILEIO_DELETE_BATCH_SIZE, "0");
+    AssertHelpers.assertThrows("should not accept batch size less than 1",
+            IllegalArgumentException.class,
+            "Deletion batch size must be between 1 and 1000",
+            () -> new AwsProperties(map));
+  }
+
 }