Add replace flag to getStatisticsCollectionMetadataForWrite

Author: ebyhr

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 replace);
 
     /**
      * 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 replace)
     {
         CatalogMetadata catalogMetadata = getCatalogMetadataForWrite(session, catalogHandle);
         ConnectorMetadata metadata = catalogMetadata.getMetadata(session);
-        return metadata.getStatisticsCollectionMetadataForWrite(session.toConnectorSession(catalogHandle), tableMetadata);
+        return metadata.getStatisticsCollectionMetadataForWrite(session.toConnectorSession(catalogHandle), tableMetadata, replace);
     }
 
     @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 replace)
+    {
+        Span span = startSpan("getStatisticsCollectionMetadataForWrite", tableMetadata.getTable());
+        try (var _ = scopedSpan(span)) {
+            return delegate.getStatisticsCollectionMetadataForWrite(session, tableMetadata, replace);
+        }
+    }
+
     @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 replace)
     {
         Span span = startSpan("getStatisticsCollectionMetadataForWrite", catalogHandle.getCatalogName().toString(), tableMetadata);
         try (var _ = scopedSpan(span)) {
-            return delegate.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata);
+            return delegate.getStatisticsCollectionMetadataForWrite(session, catalogHandle, tableMetadata, replace);
         }
     }
 

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 replace)
     {
         throw new UnsupportedOperationException();
     }

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

@@ -685,6 +685,17 @@ default Optional<ConnectorTableLayout> getInsertLayout(ConnectorSession session,
     /**
      * Describes statistics that must be collected during a write.
      */
+    default TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata, boolean replace)
+    {
+        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 replace)
+    {
+        try (ThreadContextClassLoader _ = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getStatisticsCollectionMetadataForWrite(session, tableMetadata, replace);
+        }
+    }
+
     @Override
     public TableStatisticsMetadata getStatisticsCollectionMetadataForWrite(ConnectorSession session, ConnectorTableMetadata tableMetadata)
     {

plugin/trino-lakehouse/src/test/java/io/trino/plugin/lakehouse/TestLakehouseMetadata.java

@@ -17,6 +17,7 @@
 import io.trino.spi.connector.ConnectorMetadata;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorTableHandle;
+import io.trino.spi.connector.ConnectorTableMetadata;
 import io.trino.spi.connector.JoinStatistics;
 import io.trino.spi.connector.JoinType;
 import io.trino.spi.connector.SaveMode;
@@ -48,6 +49,7 @@ public class TestLakehouseMetadata
                     .add(ConnectorMetadata.class.getMethod("supportsMissingColumnsOnInsert"))
                     .add(ConnectorMetadata.class.getMethod("refreshMaterializedView", ConnectorSession.class, SchemaTableName.class))
                     .add(ConnectorMetadata.class.getMethod("resolveIndex", ConnectorSession.class, ConnectorTableHandle.class, Set.class, Set.class, TupleDomain.class))
+                    .add(ConnectorMetadata.class.getMethod("getStatisticsCollectionMetadataForWrite", ConnectorSession.class, ConnectorTableMetadata.class))
                     .add(ConnectorMetadata.class.getMethod("listFunctions", ConnectorSession.class, String.class))
                     .add(ConnectorMetadata.class.getMethod("getFunctions", ConnectorSession.class, SchemaFunctionName.class))
                     .add(ConnectorMetadata.class.getMethod("getFunctionMetadata", ConnectorSession.class, FunctionId.class))