Limit number of manifest files written by OPTIMIZE_MANIFESTS

Avoids coordinator OOM and too many small manifests

Author: raunaqmorarka

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java

@@ -176,7 +176,6 @@
 import org.apache.iceberg.SortField;
 import org.apache.iceberg.SortOrder;
 import org.apache.iceberg.StatisticsFile;
-import org.apache.iceberg.StructLike;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.TableProperties;
 import org.apache.iceberg.TableScan;
@@ -200,6 +199,7 @@
 import org.apache.iceberg.types.Types.NestedField;
 import org.apache.iceberg.types.Types.StringType;
 import org.apache.iceberg.types.Types.StructType;
+import org.apache.iceberg.util.StructLikeWrapper;
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
@@ -219,6 +219,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.OptionalLong;
 import java.util.Set;
@@ -2116,17 +2117,23 @@ private void executeOptimizeManifests(ConnectorSession session, IcebergTableExec
         if (manifests.isEmpty()) {
             return;
         }
-        if (manifests.size() == 1 && manifests.getFirst().length() < icebergTable.operations().current().propertyAsLong(MANIFEST_TARGET_SIZE_BYTES, MANIFEST_TARGET_SIZE_BYTES_DEFAULT)) {
+        long manifestTargetSizeBytes = icebergTable.operations().current().propertyAsLong(MANIFEST_TARGET_SIZE_BYTES, MANIFEST_TARGET_SIZE_BYTES_DEFAULT);
+        if (manifests.size() == 1 && manifests.getFirst().length() < manifestTargetSizeBytes) {
             return;
         }
+        long totalManifestsSize = manifests.stream().mapToLong(ManifestFile::length).sum();
+        // Having too many open manifest writers can potentially cause OOM on the coordinator
+        long targetManifestClusters = Math.min(((totalManifestsSize + manifestTargetSizeBytes - 1) / manifestTargetSizeBytes), 100);
 
         beginTransaction(icebergTable);
         RewriteManifests rewriteManifests = transaction.rewriteManifests();
+        Types.StructType structType = icebergTable.spec().partitionType();
         rewriteManifests
                 .clusterBy(file -> {
-                    // Use the first partition field as the clustering key
-                    StructLike partition = file.partition();
-                    return partition.size() > 1 ? Optional.ofNullable(partition.get(0, Object.class)) : partition;
+                    // Cluster by partitions for better locality when reading data files
+                    StructLikeWrapper partitionWrapper = StructLikeWrapper.forType(structType).set(file.partition());
+                    // Limit the number of clustering buckets to avoid creating too many small manifest files
+                    return Objects.hash(partitionWrapper) % targetManifestClusters;
                 })
                 .scanManifestsWith(icebergScanExecutor)
                 .commit();

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/procedure/TestIcebergOptimizeManifestsProcedure.java

@@ -66,7 +66,7 @@ void testOptimizeManifestWithNullPartitions()
             assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests");
 
             assertThat(manifestFiles(table.getName()))
-                    .hasSize(2)
+                    .hasSize(1)
                     .doesNotContainAnyElementsOf(manifestFiles);
 
             assertThat(query("SELECT * FROM " + table.getName()))
@@ -138,7 +138,7 @@ void testPartitionTable()
 
             assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests");
             assertThat(manifestFiles(table.getName()))
-                    .hasSize(2)
+                    .hasSize(1)
                     .doesNotContainAnyElementsOf(manifestFiles);
 
             assertThat(query("SELECT * FROM " + table.getName()))
@@ -147,24 +147,40 @@ void testPartitionTable()
     }
 
     @Test
-    void testFirstPartitionField()
+    void testMultiplePartitioningColumns()
     {
         try (TestTable table = newTrinoTable("test_partition", "(id int, part int, nested int) WITH (partitioning = ARRAY['part', 'nested'])")) {
-            assertUpdate("INSERT INTO " + table.getName() + " VALUES (1, 10, 100)", 1);
-            assertUpdate("INSERT INTO " + table.getName() + " VALUES (2, 10, 200)", 1);
-            assertUpdate("INSERT INTO " + table.getName() + " VALUES (3, 20, 300)", 1);
-            assertUpdate("INSERT INTO " + table.getName() + " VALUES (4, 20, 400)", 1);
+            for (int i = 0; i < 30; i++) {
+                assertUpdate("INSERT INTO " + table.getName() + " VALUES (%d, %d, %d)".formatted(i, i % 10, i % 3), 1);
+            }
 
             Set<String> manifestFiles = manifestFiles(table.getName());
-            assertThat(manifestFiles).hasSize(4);
+            assertThat(manifestFiles).hasSize(30);
 
             assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests");
-            assertThat(manifestFiles(table.getName()))
+            Set<String> currentManifestFiles = manifestFiles(table.getName());
+            assertThat(currentManifestFiles)
+                    .hasSize(1)
+                    .doesNotContainAnyElementsOf(manifestFiles);
+
+            assertThat(query("SELECT COUNT(*) FROM " + table.getName()))
+                    .matches("VALUES BIGINT '30'");
+
+            // Set small target size to force split
+            BaseTable icebergTable = loadTable(table.getName());
+            icebergTable.updateProperties()
+                    .set("commit.manifest.target-size-bytes", "8000")
+                    .commit();
+            manifestFiles = currentManifestFiles;
+            assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE optimize_manifests");
+
+            currentManifestFiles = manifestFiles(table.getName());
+            assertThat(currentManifestFiles)
                     .hasSize(2)
                     .doesNotContainAnyElementsOf(manifestFiles);
 
-            assertThat(query("SELECT * FROM " + table.getName()))
-                    .matches("VALUES (1, 10, 100), (2, 10, 200), (3, 20, 300), (4, 20, 400)");
+            assertThat(query("SELECT COUNT(*) FROM " + table.getName()))
+                    .matches("VALUES BIGINT '30'");
         }
     }