Fix stats/cost caching of group-referenced plan nodes

core/trino-main/src/main/java/io/trino/connector/ConnectorManager.java

@@ -33,7 +33,7 @@
 import io.trino.metadata.HandleResolver;
 import io.trino.metadata.InternalNodeManager;
 import io.trino.metadata.MaterializedViewPropertyManager;
-import io.trino.metadata.MetadataManager;
+import io.trino.metadata.Metadata;
 import io.trino.metadata.ProcedureRegistry;
 import io.trino.metadata.SchemaPropertyManager;
 import io.trino.metadata.SessionPropertyManager;
@@ -99,7 +99,7 @@ public class ConnectorManager
 {
     private static final Logger log = Logger.get(ConnectorManager.class);
 
-    private final MetadataManager metadataManager;
+    private final Metadata metadata;
     private final CatalogManager catalogManager;
     private final AccessControlManager accessControlManager;
     private final SplitManager splitManager;
@@ -139,7 +139,7 @@ public class ConnectorManager
 
     @Inject
     public ConnectorManager(
-            MetadataManager metadataManager,
+            Metadata metadata,
             CatalogManager catalogManager,
             AccessControlManager accessControlManager,
             SplitManager splitManager,
@@ -167,7 +167,7 @@ public ConnectorManager(
             TableProceduresPropertyManager tableProceduresPropertyManager,
             NodeSchedulerConfig nodeSchedulerConfig)
     {
-        this.metadataManager = metadataManager;
+        this.metadata = metadata;
         this.catalogManager = catalogManager;
         this.accessControlManager = accessControlManager;
         this.splitManager = splitManager;
@@ -254,7 +254,7 @@ private synchronized void createCatalog(CatalogName catalogName, String connecto
 
         MaterializedConnector informationSchemaConnector = new MaterializedConnector(
                 createInformationSchemaCatalogName(catalogName),
-                new InformationSchemaConnector(catalogName.getCatalogName(), nodeManager, metadataManager, accessControlManager),
+                new InformationSchemaConnector(catalogName.getCatalogName(), nodeManager, metadata, accessControlManager),
                 () -> {});
 
         CatalogName systemId = createSystemTablesCatalogName(catalogName);
@@ -263,7 +263,7 @@ private synchronized void createCatalog(CatalogName catalogName, String connecto
         if (nodeManager.getCurrentNode().isCoordinator()) {
             systemTablesProvider = new CoordinatorSystemTablesProvider(
                     transactionManager,
-                    metadataManager,
+                    metadata,
                     catalogName.getCatalogName(),
                     new StaticSystemTablesProvider(connector.getSystemTables()));
         }
@@ -394,7 +394,7 @@ private Connector createConnector(
                 new ConnectorAwareNodeManager(nodeManager, nodeInfo.getEnvironment(), catalogName, schedulerIncludeCoordinator),
                 versionEmbedder,
                 typeManager,
-                new InternalMetadataProvider(metadataManager, typeManager),
+                new InternalMetadataProvider(metadata, typeManager),
                 pageSorter,
                 pageIndexerFactory,
                 duplicatePluginClassLoaderFactory);

core/trino-main/src/main/java/io/trino/connector/InternalMetadataProvider.java

@@ -16,7 +16,7 @@
 import io.trino.FullConnectorSession;
 import io.trino.Session;
 import io.trino.metadata.MaterializedViewDefinition;
-import io.trino.metadata.MetadataManager;
+import io.trino.metadata.Metadata;
 import io.trino.metadata.QualifiedObjectName;
 import io.trino.metadata.TableHandle;
 import io.trino.metadata.ViewColumn;
@@ -37,12 +37,12 @@
 public class InternalMetadataProvider
         implements MetadataProvider
 {
-    private final MetadataManager metadataManager;
+    private final Metadata metadata;
     private final TypeManager typeManager;
 
-    public InternalMetadataProvider(MetadataManager metadataManager, TypeManager typeManager)
+    public InternalMetadataProvider(Metadata metadata, TypeManager typeManager)
     {
-        this.metadataManager = requireNonNull(metadataManager, "metadataManager is null");
+        this.metadata = requireNonNull(metadata, "metadata is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
     }
 
@@ -52,19 +52,19 @@ public Optional<ConnectorTableSchema> getRelationMetadata(ConnectorSession conne
         Session session = ((FullConnectorSession) connectorSession).getSession();
         QualifiedObjectName qualifiedName = new QualifiedObjectName(tableName.getCatalogName(), tableName.getSchemaTableName().getSchemaName(), tableName.getSchemaTableName().getTableName());
 
-        Optional<MaterializedViewDefinition> materializedView = metadataManager.getMaterializedView(session, qualifiedName);
+        Optional<MaterializedViewDefinition> materializedView = metadata.getMaterializedView(session, qualifiedName);
         if (materializedView.isPresent()) {
             return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(materializedView.get().getColumns())));
         }
 
-        Optional<ViewDefinition> view = metadataManager.getView(session, qualifiedName);
+        Optional<ViewDefinition> view = metadata.getView(session, qualifiedName);
         if (view.isPresent()) {
             return Optional.of(new ConnectorTableSchema(tableName.getSchemaTableName(), toColumnSchema(view.get().getColumns())));
         }
 
-        Optional<TableHandle> tableHandle = metadataManager.getTableHandle(session, qualifiedName);
+        Optional<TableHandle> tableHandle = metadata.getTableHandle(session, qualifiedName);
         if (tableHandle.isPresent()) {
-            return Optional.of(metadataManager.getTableSchema(session, tableHandle.get()).getTableSchema());
+            return Optional.of(metadata.getTableSchema(session, tableHandle.get()).getTableSchema());
         }
 
         return Optional.empty();

core/trino-main/src/main/java/io/trino/cost/CachingCostProvider.java

@@ -98,7 +98,7 @@ private PlanCostEstimate getGroupCost(GroupReference groupReference)
             return knownCost.get();
         }
 
-        PlanCostEstimate cost = calculateCost(memo.getNode(group));
+        PlanCostEstimate cost = getCost(memo.getNode(group));
         verify(memo.getCost(group).isEmpty(), "Group cost already set");
         memo.storeCost(group, cost);
         return cost;

core/trino-main/src/main/java/io/trino/cost/CachingStatsProvider.java

@@ -100,7 +100,7 @@ private PlanNodeStatsEstimate getGroupStats(GroupReference groupReference)
             return stats.get();
         }
 
-        PlanNodeStatsEstimate groupStats = statsCalculator.calculateStats(memo.getNode(group), this, lookup, session, types);
+        PlanNodeStatsEstimate groupStats = getStats(memo.getNode(group));
         verify(memo.getStats(group).isEmpty(), "Group stats already set");
         memo.storeStats(group, groupStats);
         return groupStats;

core/trino-main/src/main/java/io/trino/testing/LocalQueryRunner.java

@@ -95,6 +95,7 @@
 import io.trino.metadata.SessionPropertyManager;
 import io.trino.metadata.Split;
 import io.trino.metadata.SystemFunctionBundle;
+import io.trino.metadata.SystemSecurityMetadata;
 import io.trino.metadata.TableHandle;
 import io.trino.metadata.TableProceduresPropertyManager;
 import io.trino.metadata.TableProceduresRegistry;
@@ -307,6 +308,7 @@ private LocalQueryRunner(
             int nodeCountForStats,
             Map<String, List<PropertyMetadata<?>>> defaultSessionProperties,
             PlanOptimizersProvider planOptimizersProvider,
+            MetadataProvider metadataProvider,
             OperatorFactories operatorFactories,
             Set<SystemSessionPropertiesProvider> extraSessionProperties)
     {
@@ -352,7 +354,7 @@ private LocalQueryRunner(
         globalFunctionCatalog.addFunctions(new InternalFunctionBundle(new LiteralFunction(blockEncodingSerde)));
         globalFunctionCatalog.addFunctions(SystemFunctionBundle.create(featuresConfig, typeOperators, blockTypeOperators, nodeManager.getCurrentNode().getNodeVersion()));
         this.functionManager = new FunctionManager(globalFunctionCatalog);
-        MetadataManager metadata = new MetadataManager(
+        Metadata metadata = metadataProvider.getMetadata(
                 featuresConfig,
                 new DisabledSystemSecurityMetadata(),
                 transactionManager,
@@ -1139,6 +1141,16 @@ PlanOptimizersFactory getPlanOptimizers(
                 RuleStatsRecorder ruleStats);
     }
 
+    public interface MetadataProvider
+    {
+        Metadata getMetadata(
+                FeaturesConfig featuresConfig,
+                SystemSecurityMetadata systemSecurityMetadata,
+                TransactionManager transactionManager,
+                GlobalFunctionCatalog globalFunctionCatalog,
+                TypeManager typeManager);
+    }
+
     public static class Builder
     {
         private final Session defaultSession;
@@ -1150,6 +1162,7 @@ public static class Builder
         private Set<SystemSessionPropertiesProvider> extraSessionProperties = ImmutableSet.of();
         private int nodeCountForStats;
         private PlanOptimizersProvider planOptimizersProvider = PlanOptimizers::new;
+        private MetadataProvider metadataProvider = MetadataManager::new;
         private OperatorFactories operatorFactories = new TrinoOperatorFactories();
 
         private Builder(Session defaultSession)
@@ -1199,6 +1212,12 @@ public Builder withPlanOptimizersProvider(PlanOptimizersProvider planOptimizersP
             return this;
         }
 
+        public Builder withMetadataProvider(MetadataProvider metadataProvider)
+        {
+            this.metadataProvider = metadataProvider;
+            return this;
+        }
+
         public Builder withOperatorFactories(OperatorFactories operatorFactories)
         {
             this.operatorFactories = requireNonNull(operatorFactories, "operatorFactories is null");
@@ -1227,6 +1246,7 @@ public LocalQueryRunner build()
                     nodeCountForStats,
                     defaultSessionProperties,
                     planOptimizersProvider,
+                    metadataProvider,
                     operatorFactories,
                     extraSessionProperties);
         }