Populate snapshot summary in iceberg ConnectorOutputMetadata

Author: raunaqmorarka

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

@@ -0,0 +1,40 @@
+/*
+ * 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;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.common.collect.ImmutableMap;
+import io.trino.spi.connector.ConnectorOutputMetadata;
+
+import java.util.Map;
+
+final class IcebergCommitMetadata
+        implements ConnectorOutputMetadata
+{
+    private final Map<String, String> commitMetrics;
+
+    @JsonCreator
+    public IcebergCommitMetadata(@JsonProperty("commitMetrics") Map<String, String> commitMetrics)
+    {
+        this.commitMetrics = ImmutableMap.copyOf(commitMetrics);
+    }
+
+    @Override
+    @JsonProperty
+    public Map<String, String> getInfo()
+    {
+        return commitMetrics;
+    }
+}

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

@@ -46,7 +46,6 @@
 import io.trino.plugin.base.projection.ApplyProjectionUtil.ProjectedColumnRepresentation;
 import io.trino.plugin.hive.HiveCompressionCodec;
 import io.trino.plugin.hive.HiveStorageFormat;
-import io.trino.plugin.hive.HiveWrittenPartitions;
 import io.trino.plugin.iceberg.aggregation.DataSketchStateSerializer;
 import io.trino.plugin.iceberg.aggregation.IcebergThetaSketchForStats;
 import io.trino.plugin.iceberg.catalog.TrinoCatalog;
@@ -1612,9 +1611,12 @@ else if (!computedStatistics.isEmpty()) {
         commitTransaction(transaction, "insert");
         transaction = null;
 
-        return Optional.of(new HiveWrittenPartitions(commitTasks.stream()
-                .map(CommitTaskData::path)
-                .collect(toImmutableList())));
+        long insertSnapshotId = icebergTable.currentSnapshot().snapshotId();
+        Map<String, String> summary = icebergTable.snapshot(insertSnapshotId).summary();
+        if (summary == null) {
+            return Optional.empty();
+        }
+        return Optional.of(new IcebergCommitMetadata(summary));
     }
 
     @Override
@@ -3992,9 +3994,12 @@ public Optional<ConnectorOutputMetadata> finishRefreshMaterializedView(
         commitUpdateAndTransaction(appendFiles, session, transaction, "refresh materialized view");
         transaction = null;
         fromSnapshotForRefresh = Optional.empty();
-        return Optional.of(new HiveWrittenPartitions(commitTasks.stream()
-                .map(CommitTaskData::path)
-                .collect(toImmutableList())));
+        long refreshedSnapshotId = icebergTable.currentSnapshot().snapshotId();
+        Map<String, String> summary = icebergTable.snapshot(refreshedSnapshotId).summary();
+        if (summary == null) {
+            return Optional.empty();
+        }
+        return Optional.of(new IcebergCommitMetadata(summary));
     }
 
     @Override

plugin/trino-iceberg/src/test/java/io/trino/plugin/iceberg/BaseIcebergConnectorTest.java

@@ -25,6 +25,7 @@
 import io.trino.Session;
 import io.trino.connector.MockConnectorFactory;
 import io.trino.connector.MockConnectorPlugin;
+import io.trino.execution.QueryStats;
 import io.trino.execution.StageId;
 import io.trino.execution.StageInfo;
 import io.trino.execution.StagesInfo;
@@ -36,6 +37,7 @@
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.metadata.TableHandle;
 import io.trino.operator.OperatorStats;
+import io.trino.operator.TableFinishInfo;
 import io.trino.plugin.hive.HiveCompressionCodec;
 import io.trino.server.DynamicFilterService;
 import io.trino.spi.QueryId;
@@ -9253,6 +9255,55 @@ void testExplainAnalyzeSplitSourceMetrics()
                 "splits generation metrics");
     }
 
+    @Test
+    void testCommitMetrics()
+    {
+        try (TestTable table = newTrinoTable(
+                "test_commit_metrics",
+                "AS SELECT * FROM nation WITH NO DATA")) {
+            assertQueryStats(
+                    getSession(),
+                    "INSERT INTO " + table.getName() + " SELECT * FROM nation",
+                    queryStats -> {
+                        TableFinishInfo info = getTableFinishInfo(queryStats);
+                        assertThat(info.getConnectorOutputMetadata()).contains("\"added-records\" : \"25\"");
+                    },
+                    _ -> {});
+        }
+
+        String tableName = "test_commit_metrics_ctas_" + randomNameSuffix();
+        assertQueryStats(
+                getSession(),
+                "CREATE TABLE " + tableName + " AS SELECT * FROM nation",
+                queryStats -> {
+                    TableFinishInfo info = getTableFinishInfo(queryStats);
+                    assertThat(info.getConnectorOutputMetadata()).contains("\"added-records\" : \"25\"");
+                },
+                _ -> {});
+        assertUpdate("DROP TABLE " + tableName);
+
+        String materializedViewName = "test_commit_metrics_mv_" + randomNameSuffix();
+        assertUpdate("CREATE MATERIALIZED VIEW " + materializedViewName + " AS SELECT * FROM nation");
+        assertQueryStats(
+                getSession(),
+                "REFRESH MATERIALIZED VIEW " + materializedViewName,
+                queryStats -> {
+                    TableFinishInfo info = getTableFinishInfo(queryStats);
+                    assertThat(info.getConnectorOutputMetadata()).contains("\"added-records\" : \"25\"");
+                },
+                _ -> {});
+        assertUpdate("DROP MATERIALIZED VIEW " + materializedViewName);
+    }
+
+    private static TableFinishInfo getTableFinishInfo(QueryStats queryStats)
+    {
+        OperatorStats finishOperatorStats = queryStats.getOperatorSummaries()
+                .stream()
+                .filter(summary -> summary.getOperatorType().startsWith("TableFinish"))
+                .collect(onlyElement());
+        return (TableFinishInfo) finishOperatorStats.getInfo();
+    }
+
     // regression test for https://github.com/trinodb/trino/issues/22922
     @Test
     void testArrayElementChange()