HIVE-29788: Support for rewrite_manifests iceberg procedure

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/HiveIcebergStorageHandler.java

@@ -1249,6 +1249,9 @@ public void executeOperation(org.apache.hadoop.hive.ql.metadata.Table hmsTable,
             (AlterTableExecuteSpec.CherryPickSpec) executeSpec.getOperationParams();
         IcebergTableUtil.cherryPick(icebergTable, cherryPickSpec.getSnapshotId());
         break;
+      case REWRITE_MANIFESTS:
+        IcebergTableUtil.rewriteManifests(icebergTable);
+        break;
       case DELETE_METADATA:
         AlterTableExecuteSpec.DeleteMetadataSpec deleteMetadataSpec =
             (AlterTableExecuteSpec.DeleteMetadataSpec) executeSpec.getOperationParams();

iceberg/iceberg-handler/src/main/java/org/apache/iceberg/mr/hive/IcebergTableUtil.java

@@ -31,6 +31,7 @@
 import java.util.Optional;
 import java.util.Properties;
 import java.util.Set;
+import java.util.TreeSet;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -67,11 +68,14 @@
 import org.apache.hadoop.hive.ql.session.SessionStateUtil;
 import org.apache.hadoop.util.Sets;
 import org.apache.iceberg.ContentFile;
+import org.apache.iceberg.DataFile;
 import org.apache.iceberg.DeleteFiles;
 import org.apache.iceberg.FileFormat;
 import org.apache.iceberg.FileScanTask;
 import org.apache.iceberg.ManageSnapshots;
 import org.apache.iceberg.ManifestFile;
+import org.apache.iceberg.ManifestFiles;
+import org.apache.iceberg.ManifestReader;
 import org.apache.iceberg.MetadataTableType;
 import org.apache.iceberg.MetadataTableUtils;
 import org.apache.iceberg.PartitionData;
@@ -110,11 +114,13 @@
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.apache.iceberg.transforms.Transform;
+import org.apache.iceberg.types.Comparators;
 import org.apache.iceberg.types.Conversions;
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.ByteBuffers;
 import org.apache.iceberg.util.Pair;
+import org.apache.iceberg.util.PartitionUtil;
 import org.apache.iceberg.util.SnapshotUtil;
 import org.apache.iceberg.util.StructProjection;
 import org.slf4j.Logger;
@@ -880,6 +886,96 @@
     });
   }
 
+  public static void rewriteManifests(Table table) {
+    if (!table.spec().isPartitioned()) {
+      table.rewriteManifests().clusterBy(file -> "").commit();
+    } else {
+      // Determine the target size for each new manifest file (defaults to 8MB)
+      long manifestTargetSizeBytes = TableProperties.MANIFEST_TARGET_SIZE_BYTES_DEFAULT;
+      if (table.properties().containsKey(TableProperties.MANIFEST_TARGET_SIZE_BYTES)) {
+        manifestTargetSizeBytes =
+            Long.parseLong(table.properties().get(TableProperties.MANIFEST_TARGET_SIZE_BYTES));
+      }
+
+      List<ManifestFile> dataManifests = table.currentSnapshot().dataManifests(table.io());
+      if (dataManifests.isEmpty()) {
+        return;
+      }
+
+      // Calculate ideal number of manifest files based on current total metadata size.
+      // We hard-cap the maximum number of clusters at 200 to prevent the JVM from
+      // opening too many concurrent file writers and causing OOM or OS ulimit (Too Many Open Files).
+      long totalManifestsSize = dataManifests.stream().mapToLong(ManifestFile::length).sum();
+      int targetClusters =
+          (int)
+              Math.min(
+                  (totalManifestsSize + manifestTargetSizeBytes - 1) / manifestTargetSizeBytes,
+                  200);
+
+      if (targetClusters <= 1) {
+        table.rewriteManifests().clusterBy(file -> 0).commit();
+        return;
+      }
+
+      // To cluster files efficiently, we want to group them naturally.
+      // We extract the native Type of the first partition column (e.g. Timestamp, String)
+      // and use Iceberg's native Comparators to maintain a sorted TreeSet of all unique partition values.
+      Type.PrimitiveType firstPartitionFieldType =
+          table.spec().partitionType().fields().getFirst().type().asPrimitiveType();
+      Set<Object> uniqueValues = new TreeSet<>(Comparators.forType(firstPartitionFieldType));
+
+      for (ManifestFile manifestFile : dataManifests) {
+        try (ManifestReader<DataFile> reader =
+            ManifestFiles.read(manifestFile, table.io(), table.specs())
+                .select(List.of(DataFile.PARTITION_NAME))) {
+          for (DataFile dataFile : reader) {
+            // Coerce partition struct in case of partition evolution
+            StructLike partition =
+                PartitionUtil.coercePartition(
+                    table.spec().partitionType(),
+                    table.specs().get(dataFile.specId()),
+                    dataFile.partition());
+            // Only extract and sort by the FIRST partition column for read optimization
+            Object value = partition.get(0, Object.class);
+            if (value != null) {
+              uniqueValues.add(value);
+            }
+          }
+        } catch (IOException e) {
+          throw new RuntimeException("Failed to read manifest file", e);
+        }
+      }
+
+      if (uniqueValues.isEmpty()) {
+        table.rewriteManifests().clusterBy(file -> 0).commit();
+        return;
+      }
+
+      // Divide the naturally sorted unique values evenly into our calculated `targetClusters`
+      Object[] sortedValues = uniqueValues.toArray();
+      Map<Object, Integer> valueToBucket = Maps.newHashMap();
+      for (int i = 0; i < sortedValues.length; i++) {
+        // e.g. If we have 1000 sorted partition values and 200 clusters, this groups 5 values per bucket ID
+        valueToBucket.put(sortedValues[i], i * targetClusters / sortedValues.length);
+      }
+
+      // Rewrite manifests, telling Iceberg to group data files based on our pre-calculated bucket mapping
+      table
+          .rewriteManifests()
+          .clusterBy(
+              file -> {
+                StructLike partition =
+                    PartitionUtil.coercePartition(
+                        table.spec().partitionType(),
+                        table.specs().get(file.specId()),
+                        file.partition());
+                Object value = partition.get(0, Object.class);
+                return value != null ? valueToBucket.getOrDefault(value, 0) : 0;
+              })
+          .commit();
+    }
+  }
+
   public static boolean hasUndergonePartitionEvolution(Table table) {
     return table.specs().size() > 1;
   }

iceberg/iceberg-handler/src/test/java/org/apache/iceberg/mr/hive/TestHiveIcebergSnapshotOperations.java

@@ -119,4 +119,76 @@ public void testReplaceBranchWithSnapshot() {
     result = shell.executeStatement("SELECT COUNT(*) FROM default.testReplaceBranchWithSnapshot.branch_branch1");
     assertEquals(6L, result.get(0)[0]);
   }
+
+  @Test
+  public void testRewriteManifests() {
+    TableIdentifier identifier = TableIdentifier.of("default", "testRewriteManifests");
+    shell.executeStatement(
+        String.format(
+            "CREATE EXTERNAL TABLE %s (id INT, data STRING) STORED BY iceberg %s %s",
+            identifier.name(),
+            testTables.locationForCreateTableSQL(identifier),
+            testTables.propertiesForCreateTableSQL(ImmutableMap.of("commit.manifest.min-count-to-compact", "2"))));
+
+    // Create 5 manifests by executing 5 separate INSERT operations
+    for (int i = 1; i <= 5; i++) {
+      shell.executeStatement(
+          String.format("INSERT INTO TABLE %s VALUES(%d, 'val')", identifier.name(), i));
+    }
+
+    org.apache.iceberg.Table icebergTable = testTables.loadTable(identifier);
+    icebergTable.refresh();
+
+    // After 5 inserts, Iceberg will have generated 5 separate manifest files
+    int manifestCountBefore = icebergTable.currentSnapshot().allManifests(icebergTable.io()).size();
+    assertEquals("Manifests keep accumulating for each insert", 5, manifestCountBefore);
+
+    // Execute REWRITE_MANIFESTS procedure
+    shell.executeStatement(
+        String.format("ALTER TABLE %s EXECUTE REWRITE_MANIFESTS", identifier.name()));
+
+    icebergTable.refresh();
+    int manifestCountAfterRewrite =
+        icebergTable.currentSnapshot().allManifests(icebergTable.io()).size();
+
+    // Rewrite manifests should confidently compact all of them into exactly 1 manifest
+    assertEquals(
+        "Should have exactly 1 manifest after REWRITE_MANIFESTS", 1, manifestCountAfterRewrite);
+  }
+
+  @Test
+  public void testRewriteManifestsPartitioned() {
+    TableIdentifier identifier = TableIdentifier.of("default", "testRewriteManifestsPartitioned");
+    shell.executeStatement(
+        String.format(
+            "CREATE EXTERNAL TABLE %s (id INT, data STRING) PARTITIONED BY (part STRING) STORED BY iceberg %s %s",
+            identifier.name(),
+            testTables.locationForCreateTableSQL(identifier),
+            testTables.propertiesForCreateTableSQL(
+                ImmutableMap.of("commit.manifest.min-count-to-compact", "2"))));
+
+    // Create 5 manifests by executing 5 separate INSERT operations across 2 partitions
+    for (int i = 1; i <= 5; i++) {
+      String partitionVal = (i % 2 == 0) ? "p2" : "p1";
+      shell.executeStatement(
+          String.format(
+              "INSERT INTO TABLE %s VALUES(%d, 'val', '%s')", identifier.name(), i, partitionVal));
+    }
+
+    org.apache.iceberg.Table icebergTable = testTables.loadTable(identifier);
+    icebergTable.refresh();
+
+    int manifestCountBefore = icebergTable.currentSnapshot().allManifests(icebergTable.io()).size();
+    assertEquals("Manifests keep accumulating for each insert", 5, manifestCountBefore);
+
+    shell.executeStatement(
+        String.format("ALTER TABLE %s EXECUTE REWRITE_MANIFESTS", identifier.name()));
+
+    icebergTable.refresh();
+    int manifestCountAfterRewrite =
+        icebergTable.currentSnapshot().allManifests(icebergTable.io()).size();
+
+    assertEquals(
+        "Should have exactly 1 manifest after REWRITE_MANIFESTS", 1, manifestCountAfterRewrite);
+  }
 }

parser/src/java/org/apache/hadoop/hive/ql/parse/AlterClauseParser.g

@@ -534,6 +534,8 @@ alterStatementSuffixExecute
     -> ^(TOK_ALTERTABLE_EXECUTE KW_ROLLBACK $rollbackParam)
     | KW_EXECUTE KW_EXPIRE_SNAPSHOTS (LPAREN (expireParam=expression) RPAREN)?
     -> ^(TOK_ALTERTABLE_EXECUTE KW_EXPIRE_SNAPSHOTS $expireParam?)
+    | KW_EXECUTE KW_REWRITE_MANIFESTS (LPAREN (rewriteManifestsParam=expression) RPAREN)?
+    -> ^(TOK_ALTERTABLE_EXECUTE KW_REWRITE_MANIFESTS $rewriteManifestsParam?)
     | KW_EXECUTE KW_SET_CURRENT_SNAPSHOT LPAREN (snapshotParam=expression) RPAREN
     -> ^(TOK_ALTERTABLE_EXECUTE KW_SET_CURRENT_SNAPSHOT $snapshotParam)
     | KW_EXECUTE KW_FAST_FORWARD sourceBranch=StringLiteral (targetBranch=StringLiteral)?

parser/src/java/org/apache/hadoop/hive/ql/parse/HiveLexerParent.g

@@ -398,6 +398,7 @@ KW_SPEC: 'SPEC';
 KW_SYSTEM_TIME: 'SYSTEM_TIME';
 KW_SYSTEM_VERSION: 'SYSTEM_VERSION';
 KW_EXPIRE_SNAPSHOTS: 'EXPIRE_SNAPSHOTS';
+KW_REWRITE_MANIFESTS: 'REWRITE_MANIFESTS';
 KW_SET_CURRENT_SNAPSHOT: 'SET_CURRENT_SNAPSHOT';
 KW_BRANCH: 'BRANCH';
 KW_SNAPSHOTS: 'SNAPSHOTS';

parser/src/java/org/apache/hadoop/hive/ql/parse/IdentifiersParser.g

@@ -1027,6 +1027,7 @@ nonReserved
     | KW_SPEC
     | KW_SYSTEM_TIME | KW_SYSTEM_VERSION
     | KW_EXPIRE_SNAPSHOTS
+    | KW_REWRITE_MANIFESTS
     | KW_SET_CURRENT_SNAPSHOT
     | KW_BRANCH | KW_SNAPSHOTS | KW_RETAIN | KW_RETENTION
     | KW_TAG

ql/src/java/org/apache/hadoop/hive/ql/ddl/table/execute/AlterTableExecuteAnalyzer.java

@@ -103,6 +103,10 @@
       case HiveParser.KW_ORPHAN_FILES:
         desc = getDeleteOrphanFilesDesc(tableName, partitionSpec,  command.getChildren());
         break;
+      case HiveParser.KW_REWRITE_MANIFESTS:
+        desc = new AlterTableExecuteDesc(tableName, partitionSpec,
+            new AlterTableExecuteSpec(AlterTableExecuteSpec.ExecuteOperationType.REWRITE_MANIFESTS, null));
+        break;
     }
 
     rootTasks.add(TaskFactory.get(new DDLWork(getInputs(), getOutputs(), desc)));

ql/src/java/org/apache/hadoop/hive/ql/parse/AlterTableExecuteSpec.java

@@ -43,7 +43,8 @@ public enum ExecuteOperationType {
     FAST_FORWARD,
     CHERRY_PICK,
     DELETE_METADATA,
-    DELETE_ORPHAN_FILES;
+    DELETE_ORPHAN_FILES,
+    REWRITE_MANIFESTS;
   }
 
   private final ExecuteOperationType operationType;