Collect NDV stats on all columns when replacing Iceberg tables

core/trino-main/src/main/java/io/trino/metadata/Metadata.java

@@ -353,7 +353,7 @@ Optional<TableExecuteHandle> getTableHandleForExecute(
     /**
      * Describes statistics that must be collected during a write.
      */
-    TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata);
+    TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata, boolean tableReplace);
 
     /**
      * Describe statistics that must be collected during a statistics collection

core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java

@@ -1067,11 +1067,11 @@ public Optional<TableLayout> getInsertLayout(Session session, TableHandle table)
     }
 
     @Override
-    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata)
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata, boolean tableReplace)
     {
         CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, catalogHandle);
         ConnectorMetadata metadata = catalogMetadata.getMetadata(session);
-        return metadata.getStatisticsCollectionMetadataForWrite(session.toConnectorSession(catalogHandle), tableMetadata);
+        return metadata.getStatisticsCollectionMetadataForWrite(session.toConnectorSession(catalogHandle), tableMetadata, tableReplace);
     }
 
     @Override

core/trino-main/src/main/java/io/trino/sql/planner/LogicalPlanner.java

@@ -489,7 +489,7 @@ private RelationPlan createTableCreationPlan(Analysis analysis, Query query)
                 .map(ColumnMetadata::getName)
                 .collect(toImmutableList());
 
-        TableStatisticsMetadata statisticsMetadata = metadata.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata);
+        TableStatisticsMetadata statisticsMetadata = metadata.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata, create.isReplace());
 
         return createTableWriterPlan(
                 analysis,
@@ -600,7 +600,7 @@ private RelationPlan getInsertPlan(
                 .map(ColumnMetadata::getName)
                 .collect(toImmutableList());
 
-        TableStatisticsMetadata statisticsMetadata = metadata.getStatisticsCollectionMetadataForWrite(session, tableHandle.catalogHandle(), tableMetadata.metadata());
+        TableStatisticsMetadata statisticsMetadata = metadata.getStatisticsCollectionMetadataForWrite(session, tableHandle.catalogHandle(), tableMetadata.metadata(), false);
 
         if (materializedViewRefreshWriterTarget.isPresent()) {
             RefreshType refreshType = IncrementalRefreshVisitor.canIncrementallyRefresh(plan.getRoot());

core/trino-main/src/main/java/io/trino/tracing/TracingConnectorMetadata.java

@@ -578,6 +578,15 @@ public Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
         }
     }
 
+    @Override
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean tableReplace)
+    {
+        Span span = startSpan("getStatisticsCollectionMetadataForWrite", tableMetadata.getTable());
+        try (var _ = scopedSpan(span)) {
+            return delegate.getStatisticsCollectionMetadataForWrite(session, tableMetadata, tableReplace);
+        }
+    }
+
     @Override
     public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {

core/trino-main/src/main/java/io/trino/tracing/TracingMetadata.java

@@ -630,11 +630,11 @@ public Optional<TableLayout> getInsertLayout(Session session, TableHandle target
     }
 
     @Override
-    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata)
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata, boolean tableReplace)
     {
         Span span = startSpan("getStatisticsCollectionMetadataForWrite", catalogHandle.getCatalogName().toString(), tableMetadata);
         try (var _ = scopedSpan(span)) {
-            return delegate.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata);
+            return delegate.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata, tableReplace);
         }
     }
 

core/trino-main/src/test/java/io/trino/metadata/AbstractMockMetadata.java

@@ -436,7 +436,7 @@ public Optional<TableLayout> getInsertLayout(Session session, TableHandle target
     }
 
     @Override
-    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata)
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Session session, CatalogHandle catalogHandle, ConnectorTableMetadata tableMetadata, boolean tableReplace)
     {
         throw new UnsupportedOperationException();
     }

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorMetadata.java

@@ -684,7 +684,20 @@ default Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
 
     /**
      * Describes statistics that must be collected during a write.
+     *
+     * @param tableReplace if true, replace old statistics
      */
+    default TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean tableReplace)
+    {
+        return getStatisticsCollectionMetadataForWrite(session, tableMetadata);
+    }
+
+    /**
+     * Describes statistics that must be collected during a write.
+     *
+     * @deprecated use {@link #getStatisticsCollectionMetadataForWrite(ConnectorSession, ConnectorTableMetadata, boolean)}
+     */
+    @Deprecated
     default TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {
         return TableStatisticsMetadata.empty();

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeConnectorMetadata.java

@@ -154,6 +154,14 @@ public Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
         }
     }
 
+    @Override
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean tableReplace)
+    {
+        try (ThreadContextClassLoader _ = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getStatisticsCollectionMetadataForWrite(session, tableMetadata, tableReplace);
+        }
+    }
+
     @Override
     public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {

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 tableReplace)
     {
         if (!isExtendedStatisticsEnabled(session) || !isCollectExtendedStatisticsOnWrite(session)) {
             return TableStatisticsMetadata.empty();
         }
 
+        if (tableReplace) {
+            return getStatisticsCollectionMetadata(tableMetadata, Optional.empty(), availableColumnNames -> {});
+        }
+
         ConnectorTableHandle tableHandle = getTableHandle(session, tableMetadata.getTable(), Optional.empty(), Optional.empty());
         if (tableHandle == null) {
             // Assume new table (CTAS), collect NDV stats on all columns
@@ -2874,6 +2878,12 @@ public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(Connector
         return getStatisticsCollectionMetadata(tableMetadata, Optional.of(columnsWithExtendedStatistics), availableColumnNames -> {});
     }
 
+    @Override
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
+    {
+        throw new UnsupportedOperationException("This variant of getStatisticsCollectionMetadataForWrite is unsupported");
+    }
+
     @Override
     public ConnectorAnalyzeMetadata getStatisticsCollectionMetadata(ConnectorSession session, ConnectorTableHandle tableHandle, Map<String, Object> analyzeProperties)
     {

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

@@ -476,10 +476,16 @@ public Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
         return forHandle(tableHandle).getInsertLayout(session, tableHandle);
     }
 
+    @Override
+    public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean tableReplace)
+    {
+        return forProperties(tableMetadata.getProperties()).getStatisticsCollectionMetadataForWrite(session, unwrapTableMetadata(tableMetadata), tableReplace);
+    }
+
     @Override
     public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {
-        return forProperties(tableMetadata.getProperties()).getStatisticsCollectionMetadataForWrite(session, unwrapTableMetadata(tableMetadata));
+        throw new UnsupportedOperationException("This variant of getStatisticsCollectionMetadataForWrite is unsupported");
     }
 
     @Override