Add rewrite_manifests table procedure to Iceberg

docs/src/main/sphinx/connector/iceberg.md

@@ -758,6 +758,17 @@ EXECUTE <alter-table-execute>`.
 ```{include} optimize.fragment
 ```
 
+(iceberg-rewrite-manifests)=
+##### rewrite_manifests
+
+Rewrite manifests for a table to optimize scan planning.
+
+`rewrite_manifests` can be run as follows:
+
+```sql
+ALTER TABLE test_table EXECUTE rewrite_manifests;
+```
+
 (iceberg-expire-snapshots)=
 ##### expire_snapshots
 

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

@@ -53,6 +53,7 @@
 import io.trino.plugin.iceberg.procedure.IcebergExpireSnapshotsHandle;
 import io.trino.plugin.iceberg.procedure.IcebergOptimizeHandle;
 import io.trino.plugin.iceberg.procedure.IcebergRemoveOrphanFilesHandle;
+import io.trino.plugin.iceberg.procedure.IcebergRewriteManifestsHandle;
 import io.trino.plugin.iceberg.procedure.IcebergRollbackToSnapshotHandle;
 import io.trino.plugin.iceberg.procedure.IcebergTableExecuteHandle;
 import io.trino.plugin.iceberg.procedure.IcebergTableProcedureId;
@@ -147,6 +148,7 @@
 import org.apache.iceberg.PartitionSpecParser;
 import org.apache.iceberg.ReplaceSortOrder;
 import org.apache.iceberg.RewriteFiles;
+import org.apache.iceberg.RewriteManifests;
 import org.apache.iceberg.RowDelta;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.SchemaParser;
@@ -329,6 +331,7 @@
 import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.EXPIRE_SNAPSHOTS;
 import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.OPTIMIZE;
 import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.REMOVE_ORPHAN_FILES;
+import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.REWRITE_MANIFESTS;
 import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.ROLLBACK_TO_SNAPSHOT;
 import static io.trino.plugin.iceberg.procedure.MigrationUtils.addFiles;
 import static io.trino.plugin.iceberg.procedure.MigrationUtils.addFilesFromTable;
@@ -377,6 +380,8 @@
 import static org.apache.iceberg.TableProperties.DELETE_ISOLATION_LEVEL;
 import static org.apache.iceberg.TableProperties.DELETE_ISOLATION_LEVEL_DEFAULT;
 import static org.apache.iceberg.TableProperties.FORMAT_VERSION;
+import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES;
+import static org.apache.iceberg.TableProperties.MANIFEST_TARGET_SIZE_BYTES_DEFAULT;
 import static org.apache.iceberg.TableProperties.OBJECT_STORE_ENABLED;
 import static org.apache.iceberg.TableProperties.ORC_BLOOM_FILTER_COLUMNS;
 import static org.apache.iceberg.TableProperties.PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX;
@@ -1587,6 +1592,7 @@ public Optional<ConnectorTableExecuteHandle> getTableHandleForExecute(
 
         return switch (procedureId) {
             case OPTIMIZE -> getTableHandleForOptimize(tableHandle, icebergTable, executeProperties, retryMode);
+            case REWRITE_MANIFESTS -> getTableHandleForRewriteManifests(session, tableHandle);
             case DROP_EXTENDED_STATS -> getTableHandleForDropExtendedStats(session, tableHandle);
             case ROLLBACK_TO_SNAPSHOT -> getTableHandleForRollbackToSnapshot(session, tableHandle, executeProperties);
             case EXPIRE_SNAPSHOTS -> getTableHandleForExpireSnapshots(session, tableHandle, executeProperties);
@@ -1623,6 +1629,18 @@ private Optional<ConnectorTableExecuteHandle> getTableHandleForOptimize(
                 icebergTable.io().properties()));
     }
 
+    private Optional<ConnectorTableExecuteHandle> getTableHandleForRewriteManifests(ConnectorSession session, IcebergTableHandle tableHandle)
+    {
+        Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName());
+
+        return Optional.of(new IcebergTableExecuteHandle(
+                tableHandle.getSchemaTableName(),
+                REWRITE_MANIFESTS,
+                new IcebergRewriteManifestsHandle(),
+                icebergTable.location(),
+                icebergTable.io().properties()));
+    }
+
     private Optional<ConnectorTableExecuteHandle> getTableHandleForDropExtendedStats(ConnectorSession session, IcebergTableHandle tableHandle)
     {
         Table icebergTable = catalog.loadTable(session, tableHandle.getSchemaTableName());
@@ -1786,6 +1804,7 @@ public Optional<ConnectorTableLayout> getLayoutForTableExecute(ConnectorSession
         switch (executeHandle.procedureId()) {
             case OPTIMIZE:
                 return getLayoutForOptimize(session, executeHandle);
+            case REWRITE_MANIFESTS:
             case DROP_EXTENDED_STATS:
             case ROLLBACK_TO_SNAPSHOT:
             case EXPIRE_SNAPSHOTS:
@@ -1816,6 +1835,7 @@ public BeginTableExecuteResult<ConnectorTableExecuteHandle, ConnectorTableHandle
         switch (executeHandle.procedureId()) {
             case OPTIMIZE:
                 return beginOptimize(session, executeHandle, table);
+            case REWRITE_MANIFESTS:
             case DROP_EXTENDED_STATS:
             case ROLLBACK_TO_SNAPSHOT:
             case EXPIRE_SNAPSHOTS:
@@ -1862,6 +1882,7 @@ public void finishTableExecute(ConnectorSession session, ConnectorTableExecuteHa
             case OPTIMIZE:
                 finishOptimize(session, executeHandle, fragments, splitSourceInfo);
                 return;
+            case REWRITE_MANIFESTS:
             case DROP_EXTENDED_STATS:
             case ROLLBACK_TO_SNAPSHOT:
             case EXPIRE_SNAPSHOTS:
@@ -1994,6 +2015,9 @@ public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteH
     {
         IcebergTableExecuteHandle executeHandle = (IcebergTableExecuteHandle) tableExecuteHandle;
         switch (executeHandle.procedureId()) {
+            case REWRITE_MANIFESTS:
+                executeRewriteManifests(session, executeHandle);
+                return;
             case DROP_EXTENDED_STATS:
                 executeDropExtendedStats(session, executeHandle);
                 return;
@@ -2017,6 +2041,26 @@ public void executeTableExecute(ConnectorSession session, ConnectorTableExecuteH
         }
     }
 
+    private void executeRewriteManifests(ConnectorSession session, IcebergTableExecuteHandle executeHandle)
+    {
+        checkArgument(executeHandle.procedureHandle() instanceof IcebergRewriteManifestsHandle, "Unexpected procedure handle %s", executeHandle.procedureHandle());
+
+        BaseTable icebergTable = (BaseTable) catalog.loadTable(session, executeHandle.schemaTableName());
+        List<ManifestFile> manifests = icebergTable.currentSnapshot().allManifests(icebergTable.io());
+        if (manifests.isEmpty()) {
+            return;
+        }
+        if (manifests.size() == 1 && manifests.getFirst().length() < icebergTable.operations().current().propertyAsLong(MANIFEST_TARGET_SIZE_BYTES, MANIFEST_TARGET_SIZE_BYTES_DEFAULT)) {
+            return;
+        }
+
+        beginTransaction(icebergTable);
+        RewriteManifests rewriteManifests = transaction.rewriteManifests();
+        rewriteManifests.clusterBy(_ -> "file").commit();
+        commitTransaction(transaction, "rewrite manifests");
+        transaction = null;
+    }
+
     private void executeDropExtendedStats(ConnectorSession session, IcebergTableExecuteHandle executeHandle)
     {
         checkArgument(executeHandle.procedureHandle() instanceof IcebergDropExtendedStatsHandle, "Unexpected procedure handle %s", executeHandle.procedureHandle());

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

@@ -50,6 +50,7 @@
 import io.trino.plugin.iceberg.procedure.OptimizeTableProcedure;
 import io.trino.plugin.iceberg.procedure.RegisterTableProcedure;
 import io.trino.plugin.iceberg.procedure.RemoveOrphanFilesTableProcedure;
+import io.trino.plugin.iceberg.procedure.RewriteManifestsTableProcedure;
 import io.trino.plugin.iceberg.procedure.RollbackToSnapshotProcedure;
 import io.trino.plugin.iceberg.procedure.RollbackToSnapshotTableProcedure;
 import io.trino.plugin.iceberg.procedure.UnregisterTableProcedure;
@@ -131,6 +132,7 @@ public void configure(Binder binder)
 
         Multibinder<TableProcedureMetadata> tableProcedures = newSetBinder(binder, TableProcedureMetadata.class);
         tableProcedures.addBinding().toProvider(OptimizeTableProcedure.class).in(Scopes.SINGLETON);
+        tableProcedures.addBinding().toProvider(RewriteManifestsTableProcedure.class).in(Scopes.SINGLETON);
         tableProcedures.addBinding().toProvider(DropExtendedStatsTableProcedure.class).in(Scopes.SINGLETON);
         tableProcedures.addBinding().toProvider(RollbackToSnapshotTableProcedure.class).in(Scopes.SINGLETON);
         tableProcedures.addBinding().toProvider(ExpireSnapshotsTableProcedure.class).in(Scopes.SINGLETON);

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

@@ -146,6 +146,7 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
                         sortingFileWriterMaxOpenFiles,
                         typeManager,
                         pageSorter);
+            case REWRITE_MANIFESTS:
             case DROP_EXTENDED_STATS:
             case ROLLBACK_TO_SNAPSHOT:
             case EXPIRE_SNAPSHOTS:

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergProcedureHandle.java

@@ -24,6 +24,7 @@
         @JsonSubTypes.Type(value = IcebergRollbackToSnapshotHandle.class, name = "rollback_to_snapshot"),
         @JsonSubTypes.Type(value = IcebergExpireSnapshotsHandle.class, name = "expire_snapshots"),
         @JsonSubTypes.Type(value = IcebergOptimizeHandle.class, name = "optimize"),
+        @JsonSubTypes.Type(value = IcebergRewriteManifestsHandle.class, name = "rewrite_manifests"),
         @JsonSubTypes.Type(value = IcebergRemoveOrphanFilesHandle.class, name = "remove_orphan_files"),
         @JsonSubTypes.Type(value = IcebergAddFilesHandle.class, name = "add_files"),
         @JsonSubTypes.Type(value = IcebergAddFilesFromTableHandle.class, name = "add_files_from_table"),

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergRewriteManifestsHandle.java

@@ -0,0 +1,17 @@
+/*
+ * Licensed 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 io.trino.plugin.iceberg.procedure;
+
+public record IcebergRewriteManifestsHandle()
+        implements IcebergProcedureHandle {}

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/IcebergTableProcedureId.java

@@ -16,6 +16,7 @@
 public enum IcebergTableProcedureId
 {
     OPTIMIZE,
+    REWRITE_MANIFESTS,
     DROP_EXTENDED_STATS,
     ROLLBACK_TO_SNAPSHOT,
     EXPIRE_SNAPSHOTS,

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/procedure/RewriteManifestsTableProcedure.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed 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 io.trino.plugin.iceberg.procedure;
+
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Provider;
+import io.trino.spi.connector.TableProcedureMetadata;
+
+import static io.trino.plugin.iceberg.procedure.IcebergTableProcedureId.REWRITE_MANIFESTS;
+import static io.trino.spi.connector.TableProcedureExecutionMode.coordinatorOnly;
+
+public class RewriteManifestsTableProcedure
+        implements Provider<TableProcedureMetadata>
+{
+    @Override
+    public TableProcedureMetadata get()
+    {
+        return new TableProcedureMetadata(
+                REWRITE_MANIFESTS.name(),
+                coordinatorOnly(),
+                ImmutableList.of());
+    }
+}

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

@@ -0,0 +1,127 @@
+/*
+ * Licensed 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 io.trino.plugin.iceberg.procedure;
+
+import io.trino.plugin.iceberg.IcebergQueryRunner;
+import io.trino.testing.AbstractTestQueryFramework;
+import io.trino.testing.QueryRunner;
+import io.trino.testing.sql.TestTable;
+import org.junit.jupiter.api.Test;
+
+import java.util.Set;
+
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
+import static org.assertj.core.api.Assertions.assertThat;
+
+final class TestIcebergRewriteManifestsProcedure
+        extends AbstractTestQueryFramework
+{
+    @Override
+    protected QueryRunner createQueryRunner()
+            throws Exception
+    {
+        return IcebergQueryRunner.builder().build();
+    }
+
+    @Test
+    void testRewriteManifests()
+    {
+        try (TestTable table = newTrinoTable("test_rewrite_manifests", "(x int)")) {
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1);
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES 2", 1);
+
+            Set<String> manifestFiles = manifestFiles(table.getName());
+            assertThat(manifestFiles).hasSize(2);
+
+            assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests");
+            assertThat(manifestFiles(table.getName()))
+                    .hasSize(1)
+                    .doesNotContainAnyElementsOf(manifestFiles);
+
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("VALUES 1, 2");
+        }
+    }
+
+    @Test
+    void testPartitionTable()
+    {
+        try (TestTable table = newTrinoTable("test_partition", "(id int, part int) WITH (partitioning = ARRAY['part'])")) {
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES (1, 10)", 1);
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES (2, 10)", 1);
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES (3, 20)", 1);
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES (4, 20)", 1);
+
+            Set<String> manifestFiles = manifestFiles(table.getName());
+            assertThat(manifestFiles).hasSize(4);
+
+            assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests");
+            assertThat(manifestFiles(table.getName()))
+                    .hasSize(1)
+                    .doesNotContainAnyElementsOf(manifestFiles);
+
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("VALUES (1, 10), (2, 10), (3, 20), (4, 20)");
+        }
+    }
+
+    @Test
+    void testEmptyManifest()
+    {
+        try (TestTable table = newTrinoTable("test_no_rewrite", "(x int)")) {
+            Set<String> manifestFiles = manifestFiles(table.getName());
+            assertThat(manifestFiles).isEmpty();
+
+            assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests");
+            assertThat(manifestFiles(table.getName())).isEmpty();
+
+            assertQueryReturnsEmptyResult("SELECT * FROM " + table.getName());
+        }
+    }
+
+    @Test
+    void testNotRewriteSingleManifest()
+    {
+        try (TestTable table = newTrinoTable("test_no_rewrite", "(x int)")) {
+            assertUpdate("INSERT INTO " + table.getName() + " VALUES 1", 1);
+
+            Set<String> manifestFiles = manifestFiles(table.getName());
+            assertThat(manifestFiles).hasSize(1);
+
+            assertUpdate("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests");
+            assertThat(manifestFiles(table.getName()))
+                    .hasSize(1)
+                    .isEqualTo(manifestFiles);
+
+            assertThat(query("SELECT * FROM " + table.getName()))
+                    .matches("VALUES 1");
+        }
+    }
+
+    @Test
+    void testUnsupportedWhere()
+    {
+        try (TestTable table = newTrinoTable("test_unsupported_where", "WITH (partitioning = ARRAY['part']) AS SELECT 1 id, 1 part")) {
+            assertQueryFails("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests WHERE id = 1", ".* WHERE not supported for procedure REWRITE_MANIFESTS");
+            assertQueryFails("ALTER TABLE " + table.getName() + " EXECUTE rewrite_manifests WHERE part = 10", ".* WHERE not supported for procedure REWRITE_MANIFESTS");
+        }
+    }
+
+    private Set<String> manifestFiles(String tableName)
+    {
+        return computeActual("SELECT path FROM \"" + tableName + "$manifests\"").getOnlyColumnAsSet().stream()
+                .map(path -> (String) path)
+                .collect(toImmutableSet());
+    }
+}