Collect NDV stats on all columns when replacing Iceberg tables

Author: ebyhr

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

@@ -2834,12 +2834,16 @@ public void dropNotNullConstraint(ConnectorSession session, ConnectorTableHandle
     }
 
     @Override
-    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean replace)
     {
         if (!isExtendedStatisticsEnabled(session) || !isCollectExtendedStatisticsOnWrite(session)) {
             return TableStatisticsMetadata.empty();
         }
 
+        if (replace) {
+            return getStatisticsCollectionMetadata(tableMetadata, Optional.empty(), _ -> {});
+        }
+
         ConnectorTableHandle tableHandle = getTableHandle(session, tableMetadata.getTable(), Optional.empty(), Optional.empty());
         if (tableHandle == null) {
             // Assume new table (CTAS), collect NDV stats on all columns

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

@@ -18,6 +18,7 @@
 import io.trino.Session;
 import io.trino.testing.AbstractTestQueryFramework;
 import io.trino.testing.QueryRunner;
+import io.trino.testing.sql.TestTable;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -972,6 +973,31 @@ public void testShowStatsAsOf()
         assertUpdate("DROP TABLE show_stats_as_of");
     }
 
+    @Test
+    public void testShowStatsReplaceTable()
+    {
+        try (TestTable table = newTrinoTable("show_stats_after_replace_table_", "AS SELECT 1 a, 2 b")) {
+            assertThat(query("SHOW STATS FOR " + table.getName()))
+                    .skippingTypesCheck()
+                    .matches("""
+                        VALUES
+                        ('a', null, 1e0, 0e0, NULL, '1', '1'),
+                        ('b', null, 1e0, 0e0, NULL, '2', '2'),
+                        (NULL, NULL, NULL, NULL, 1e0, NULL, NULL)
+                        """);
+
+            assertUpdate("CREATE OR REPLACE TABLE " + table.getName() + " AS SELECT 3 x, 4 y", 1);
+            assertThat(query("SHOW STATS FOR " + table.getName()))
+                    .skippingTypesCheck()
+                    .matches("""
+                        VALUES
+                        ('x', null, 1e0, 0e0, NULL, '3', '3'),
+                        ('y', null, 1e0, 0e0, NULL, '4', '4'),
+                        (NULL, NULL, NULL, NULL, 1e0, NULL, NULL)
+                        """);
+        }
+    }
+
     @Test
     public void testShowStatsAfterExpiration()
     {

plugin/trino-lakehouse/src/main/java/io/trino/plugin/lakehouse/LakehouseMetadata.java

@@ -477,9 +477,9 @@ public Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
     }
 
     @Override
-    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean replace)
     {
-        return forProperties(tableMetadata.getProperties()).getStatisticsCollectionMetadataForWrite(session, unwrapTableMetadata(tableMetadata));
+        return forProperties(tableMetadata.getProperties()).getStatisticsCollectionMetadataForWrite(session, unwrapTableMetadata(tableMetadata), replace);
     }
 
     @Override