Async Deletion of Previous Metadata and Statistics Files

polaris-service/src/main/java/org/apache/polaris/service/task/ManifestFileCleanupTaskHandler.java

@@ -42,9 +42,11 @@
 import org.slf4j.LoggerFactory;
 
 /**
- * {@link TaskHandler} responsible for deleting all of the files in a manifest and the manifest
- * itself. Since data files may be present in multiple manifests across different snapshots, we
- * assume a data file that doesn't exist is missing because it was already deleted by another task.
+ * {@link TaskHandler} responsible for deleting table files: 1. Manifest files: It contains all the
+ * files in a manifest and the manifest itself. Since data files may be present in multiple
+ * manifests across different snapshots, we assume a data file that doesn't exist is missing because
+ * it was already deleted by another task. 2. Table content files: It contains previous metadata and
+ * statistics files, which are grouped and deleted in batch
  */
 public class ManifestFileCleanupTaskHandler implements TaskHandler {
   public static final int MAX_ATTEMPTS = 3;
@@ -68,58 +70,107 @@ public boolean canHandleTask(TaskEntity task) {
   @Override
   public boolean handleTask(TaskEntity task) {
     ManifestCleanupTask cleanupTask = task.readData(ManifestCleanupTask.class);
-    ManifestFile manifestFile = decodeManifestData(cleanupTask.getManifestFileData());
     TableIdentifier tableId = cleanupTask.getTableId();
     try (FileIO authorizedFileIO = fileIOSupplier.apply(task)) {
-
-      // if the file doesn't exist, we assume that another task execution was successful, but failed
-      // to drop the task entity. Log a warning and return success
-      if (!TaskUtils.exists(manifestFile.path(), authorizedFileIO)) {
+      if (cleanupTask.getManifestFileData() != null) {
+        ManifestFile manifestFile = decodeManifestData(cleanupTask.getManifestFileData());
+        return cleanUpManifestFile(manifestFile, authorizedFileIO, tableId);
+      } else if (cleanupTask.getContentFileBatch() != null) {
+        return cleanUpContentFiles(cleanupTask.getContentFileBatch(), authorizedFileIO, tableId);
+      } else {
         LOGGER
             .atWarn()
-            .addKeyValue("manifestFile", manifestFile.path())
             .addKeyValue("tableId", tableId)
-            .log("Manifest cleanup task scheduled, but manifest file doesn't exist");
+            .log("Cleanup task scheduled, but input file doesn't exist");
         return true;
       }
+    }
+  }
 
-      ManifestReader<DataFile> dataFiles = ManifestFiles.read(manifestFile, authorizedFileIO);
-      List<CompletableFuture<Void>> dataFileDeletes =
-          StreamSupport.stream(
-                  Spliterators.spliteratorUnknownSize(dataFiles.iterator(), Spliterator.IMMUTABLE),
-                  false)
-              .map(
-                  file ->
-                      tryDelete(
-                          tableId, authorizedFileIO, manifestFile, file.path().toString(), null, 1))
-              .toList();
-      LOGGER.debug(
-          "Scheduled {} data files to be deleted from manifest {}",
-          dataFileDeletes.size(),
-          manifestFile.path());
-      try {
-        // wait for all data files to be deleted, then wait for the manifest itself to be deleted
-        CompletableFuture.allOf(dataFileDeletes.toArray(CompletableFuture[]::new))
-            .thenCompose(
-                (v) -> {
-                  LOGGER
-                      .atInfo()
-                      .addKeyValue("manifestFile", manifestFile.path())
-                      .log("All data files in manifest deleted - deleting manifest");
-                  return tryDelete(
-                      tableId, authorizedFileIO, manifestFile, manifestFile.path(), null, 1);
-                })
-            .get();
-        return true;
-      } catch (InterruptedException e) {
-        LOGGER.error(
-            "Interrupted exception deleting data files from manifest {}", manifestFile.path(), e);
-        throw new RuntimeException(e);
-      } catch (ExecutionException e) {
-        LOGGER.error("Unable to delete data files from manifest {}", manifestFile.path(), e);
-        return false;
-      }
+  private boolean cleanUpManifestFile(
+      ManifestFile manifestFile, FileIO fileIO, TableIdentifier tableId) {
+    // if the file doesn't exist, we assume that another task execution was successful, but
+    // failed to drop the task entity. Log a warning and return success
+    if (!TaskUtils.exists(manifestFile.path(), fileIO)) {
+      LOGGER
+          .atWarn()
+          .addKeyValue("manifestFile", manifestFile.path())
+          .addKeyValue("tableId", tableId)
+          .log("Manifest cleanup task scheduled, but manifest file doesn't exist");
+      return true;
+    }
+
+    ManifestReader<DataFile> dataFiles = ManifestFiles.read(manifestFile, fileIO);
+    List<CompletableFuture<Void>> dataFileDeletes =
+        StreamSupport.stream(
+                Spliterators.spliteratorUnknownSize(dataFiles.iterator(), Spliterator.IMMUTABLE),
+                false)
+            .map(file -> tryDelete(tableId, fileIO, manifestFile, file.path().toString(), null, 1))
+            .toList();
+    LOGGER.debug(
+        "Scheduled {} data files to be deleted from manifest {}",
+        dataFileDeletes.size(),
+        manifestFile.path());
+    try {
+      // wait for all data files to be deleted, then wait for the manifest itself to be deleted
+      CompletableFuture.allOf(dataFileDeletes.toArray(CompletableFuture[]::new))
+          .thenCompose(
+              (v) -> {
+                LOGGER
+                    .atInfo()
+                    .addKeyValue("manifestFile", manifestFile.path())
+                    .log("All data files in manifest deleted - deleting manifest");
+                return tryDelete(tableId, fileIO, manifestFile, manifestFile.path(), null, 1);
+              })
+          .get();
+      return true;
+    } catch (InterruptedException e) {
+      LOGGER.error(
+          "Interrupted exception deleting data files from manifest {}", manifestFile.path(), e);
+      throw new RuntimeException(e);
+    } catch (ExecutionException e) {
+      LOGGER.error("Unable to delete data files from manifest {}", manifestFile.path(), e);
+      return false;
+    }
+  }
+
+  private boolean cleanUpContentFiles(
+      List<String> contentFileBatch, FileIO fileIO, TableIdentifier tableId) {
+    List<String> validFiles =
+        contentFileBatch.stream().filter(file -> TaskUtils.exists(file, fileIO)).toList();
+    if (validFiles.isEmpty()) {
+      LOGGER
+          .atWarn()
+          .addKeyValue("contentFileBatch", contentFileBatch.toString())
+          .addKeyValue("tableId", tableId)
+          .log("Table content cleanup task scheduled, but the none of the file in batch exists");
+      return true;
     }
+
+    // Schedule the deletion for each file asynchronously
+    List<CompletableFuture<Void>> deleteFutures =
+        validFiles.stream().map(file -> tryDelete(tableId, fileIO, null, file, null, 1)).toList();
+
+    // Wait for all delete operations to finish
+    try {
+      CompletableFuture<Void> allDeletes =
+          CompletableFuture.allOf(deleteFutures.toArray(new CompletableFuture[0]));
+      allDeletes.join();
+    } catch (Exception e) {
+      LOGGER
+          .atWarn()
+          .addKeyValue("contentFileBatch", contentFileBatch.toString())
+          .addKeyValue("tableId", tableId)
+          .log("Exception detected during content file batch deletion", e);
+    }
+
+    LOGGER
+        .atInfo()
+        .addKeyValue("contentFileBatch", contentFileBatch.toString())
+        .addKeyValue("tableId", tableId)
+        .log("Content file batch deletion has completed");
+
+    return true;
   }
 
   private static ManifestFile decodeManifestData(String manifestFileData) {
@@ -134,16 +185,16 @@ private CompletableFuture<Void> tryDelete(
       TableIdentifier tableId,
       FileIO fileIO,
       ManifestFile manifestFile,
-      String dataFile,
+      String file,
       Throwable e,
       int attempt) {
     if (e != null && attempt <= MAX_ATTEMPTS) {
       LOGGER
           .atWarn()
-          .addKeyValue("dataFile", dataFile)
+          .addKeyValue("file", file)
           .addKeyValue("attempt", attempt)
           .addKeyValue("error", e.getMessage())
-          .log("Error encountered attempting to delete data file");
+          .log("Error encountered attempting to delete file");
     }
     if (attempt > MAX_ATTEMPTS && e != null) {
       return CompletableFuture.failedFuture(e);
@@ -155,27 +206,27 @@ private CompletableFuture<Void> tryDelete(
               // file's existence, but then it is deleted before we have a chance to
               // send the delete request. In such a case, we <i>should</i> retry
               // and find
-              if (TaskUtils.exists(dataFile, fileIO)) {
-                fileIO.deleteFile(dataFile);
+              if (TaskUtils.exists(file, fileIO)) {
+                fileIO.deleteFile(file);
               } else {
                 LOGGER
                     .atInfo()
-                    .addKeyValue("dataFile", dataFile)
-                    .addKeyValue("manifestFile", manifestFile.path())
+                    .addKeyValue("file", file)
+                    .addKeyValue("manifestFile", manifestFile != null ? manifestFile.path() : "")
                     .addKeyValue("tableId", tableId)
-                    .log("Manifest cleanup task scheduled, but data file doesn't exist");
+                    .log("table file cleanup task scheduled, but data file doesn't exist");
               }
             },
             executorService)
         .exceptionallyComposeAsync(
             newEx -> {
               LOGGER
                   .atWarn()
-                  .addKeyValue("dataFile", dataFile)
-                  .addKeyValue("tableIdentifer", tableId)
-                  .addKeyValue("manifestFile", manifestFile.path())
+                  .addKeyValue("dataFile", file)
+                  .addKeyValue("tableIdentifier", tableId)
+                  .addKeyValue("manifestFile", manifestFile != null ? manifestFile.path() : "")
                   .log("Exception caught deleting data file from manifest", newEx);
-              return tryDelete(tableId, fileIO, manifestFile, dataFile, newEx, attempt + 1);
+              return tryDelete(tableId, fileIO, manifestFile, file, newEx, attempt + 1);
             },
             CompletableFuture.delayedExecutor(
                 FILE_DELETION_RETRY_MILLIS, TimeUnit.MILLISECONDS, executorService));
@@ -185,12 +236,18 @@ private CompletableFuture<Void> tryDelete(
   public static final class ManifestCleanupTask {
     private TableIdentifier tableId;
     private String manifestFileData;
+    private List<String> contentFileBatch;
 
     public ManifestCleanupTask(TableIdentifier tableId, String manifestFileData) {
       this.tableId = tableId;
       this.manifestFileData = manifestFileData;
     }
 
+    public ManifestCleanupTask(TableIdentifier tableId, List<String> contentFileBatch) {
+      this.tableId = tableId;
+      this.contentFileBatch = contentFileBatch;
+    }
+
     public ManifestCleanupTask() {}
 
     public TableIdentifier getTableId() {
@@ -209,17 +266,26 @@ public void setManifestFileData(String manifestFileData) {
       this.manifestFileData = manifestFileData;
     }
 
+    public List<String> getContentFileBatch() {
+      return contentFileBatch;
+    }
+
+    public void setContentFileBatch(List<String> contentFileBatch) {
+      this.contentFileBatch = contentFileBatch;
+    }
+
     @Override
     public boolean equals(Object object) {
       if (this == object) return true;
       if (!(object instanceof ManifestCleanupTask that)) return false;
       return Objects.equals(tableId, that.tableId)
-          && Objects.equals(manifestFileData, that.manifestFileData);
+          && Objects.equals(manifestFileData, that.manifestFileData)
+          && Objects.equals(contentFileBatch, that.contentFileBatch);
     }
 
     @Override
     public int hashCode() {
-      return Objects.hash(tableId, manifestFileData);
+      return Objects.hash(tableId, manifestFileData, contentFileBatch);
     }
   }
 }

polaris-service/src/main/java/org/apache/polaris/service/task/TableCleanupTaskHandler.java

@@ -18,11 +18,14 @@
  */
 package org.apache.polaris.service.task;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
 import java.util.function.Function;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 import org.apache.iceberg.ManifestFile;
+import org.apache.iceberg.StatisticsFile;
 import org.apache.iceberg.TableMetadata;
 import org.apache.iceberg.TableMetadataParser;
 import org.apache.iceberg.io.FileIO;
@@ -49,6 +52,7 @@ public class TableCleanupTaskHandler implements TaskHandler {
   private final TaskExecutor taskExecutor;
   private final MetaStoreManagerFactory metaStoreManagerFactory;
   private final Function<TaskEntity, FileIO> fileIOSupplier;
+  private static final int BATCH_SIZE = 10;
 
   public TableCleanupTaskHandler(
       TaskExecutor taskExecutor,
@@ -102,7 +106,7 @@ public boolean handleTask(TaskEntity cleanupTask) {
       // read the manifest list for each snapshot. dedupe the manifest files and schedule a
       // cleanupTask
       // for each manifest file and its data files to be deleted
-      List<TaskEntity> taskEntities =
+      Stream<TaskEntity> manifestCleanupTasks =
           tableMetadata.snapshots().stream()
               .flatMap(sn -> sn.allManifests(fileIO).stream())
               // distinct by manifest path, since multiple snapshots will contain the same
@@ -142,8 +146,40 @@ public boolean handleTask(TaskEntity cleanupTask) {
                         // copy the internal properties, which will have storage info
                         .setInternalProperties(cleanupTask.getInternalPropertiesAsMap())
                         .build();
-                  })
-              .toList();
+                  });
+
+      Stream<TaskEntity> contentFileCleanupTasks =
+          getContentFileBatch(tableMetadata, fileIO).stream()
+              .map(
+                  fileBatch -> {
+                    String taskName =
+                        String.join(
+                            "_",
+                            cleanupTask.getName(),
+                            fileBatch.toString(),
+                            UUID.randomUUID().toString());
+                    LOGGER
+                        .atDebug()
+                        .addKeyValue("taskName", taskName)
+                        .addKeyValue("tableIdentifier", tableEntity.getTableIdentifier())
+                        .addKeyValue("fileBatch", fileBatch)
+                        .log(
+                            "Queueing task to delete content file (prev metadata and statistics files)");
+                    return new TaskEntity.Builder()
+                        .setName(taskName)
+                        .setId(metaStoreManager.generateNewEntityId(polarisCallContext).getId())
+                        .setCreateTimestamp(polarisCallContext.getClock().millis())
+                        .withTaskType(AsyncTaskType.FILE_CLEANUP)
+                        .withData(
+                            new ManifestFileCleanupTaskHandler.ManifestCleanupTask(
+                                tableEntity.getTableIdentifier(), fileBatch))
+                        .setInternalProperties(cleanupTask.getInternalPropertiesAsMap())
+                        .build();
+                  });
+
+      List<TaskEntity> taskEntities =
+          Stream.concat(manifestCleanupTasks, contentFileCleanupTasks).toList();
+
       List<PolarisBaseEntity> createdTasks =
           metaStoreManager
               .createEntitiesIfNotExist(polarisCallContext, null, taskEntities)
@@ -154,15 +190,32 @@ public boolean handleTask(TaskEntity cleanupTask) {
             .addKeyValue("tableIdentifier", tableEntity.getTableIdentifier())
             .addKeyValue("metadataLocation", tableEntity.getMetadataLocation())
             .addKeyValue("taskCount", taskEntities.size())
-            .log("Successfully queued tasks to delete manifests - deleting table metadata file");
+            .log(
+                "Successfully queued tasks to delete manifests, previous metadata, and statistics files - deleting table metadata file");
         for (PolarisBaseEntity createdTask : createdTasks) {
           taskExecutor.addTaskHandlerContext(createdTask.getId(), CallContext.getCurrentContext());
         }
+
         fileIO.deleteFile(tableEntity.getMetadataLocation());
 
         return true;
       }
     }
     return false;
   }
+
+  private List<List<String>> getContentFileBatch(TableMetadata tableMetadata, FileIO fileIO) {
+    List<List<String>> result = new ArrayList<>();
+    List<String> contentFiles =
+        Stream.concat(
+                tableMetadata.previousFiles().stream().map(TableMetadata.MetadataLogEntry::file),
+                tableMetadata.statisticsFiles().stream().map(StatisticsFile::path))
+            .filter(file -> TaskUtils.exists(file, fileIO))
+            .toList();
+
+    for (int i = 0; i < contentFiles.size(); i += BATCH_SIZE) {
+      result.add(contentFiles.subList(i, Math.min(i + BATCH_SIZE, contentFiles.size())));
+    }
+    return result;
+  }
 }