Add optimizer information collector for optimizer logging

Add optimization information collector to session, so that it can be used to collect information we want to log from optimizers.
prestodb · Nov 3, 2022 · e2f5ad5 · e2f5ad5
1 parent c32f1b8
commit e2f5ad5
Show file tree

Hide file tree

Showing 11 changed files with 132 additions and 3 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/Session.java b/presto-main/src/main/java/com/facebook/presto/Session.java
@@ -31,6 +31,7 @@
 import com.facebook.presto.spi.session.ResourceEstimates;
 import com.facebook.presto.spi.session.SessionPropertyConfigurationManager.SystemSessionPropertyConfiguration;
 import com.facebook.presto.spi.tracing.Tracer;
+import com.facebook.presto.sql.planner.optimizations.OptimizerInformationCollector;
 import com.facebook.presto.transaction.TransactionId;
 import com.facebook.presto.transaction.TransactionManager;
 import com.google.common.collect.ImmutableMap;
@@ -92,6 +93,7 @@ public final class Session
     private final WarningCollector warningCollector;
 
     private final RuntimeStats runtimeStats = new RuntimeStats();
+    private final OptimizerInformationCollector optimizerInformationCollector = new OptimizerInformationCollector();
 
     public Session(
             QueryId queryId,
@@ -313,6 +315,11 @@ public WarningCollector getWarningCollector()
         return warningCollector;
     }
 
+    public OptimizerInformationCollector getOptimizerInformationCollector()
+    {
+        return optimizerInformationCollector;
+    }
+
     public Session beginTransactionId(TransactionId transactionId, TransactionManager transactionManager, AccessControl accessControl)
     {
         requireNonNull(transactionId, "transactionId is null");

diff --git a/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java b/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java
@@ -214,7 +214,8 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                 ImmutableList.of(),
                 ImmutableList.of(),
                 ImmutableList.of(),
-                Optional.empty()));
+                Optional.empty(),
+                ImmutableList.of()));
 
         logQueryTimeline(queryInfo);
     }
@@ -246,7 +247,8 @@ public void queryCompletedEvent(QueryInfo queryInfo)
                         createOperatorStatistics(queryInfo),
                         createPlanStatistics(queryInfo.getPlanStatsAndCosts()),
                         historyBasedPlanStatisticsTracker.getQueryStats(queryInfo).values().stream().collect(toImmutableList()),
-                        queryInfo.getExpandedQuery()));
+                        queryInfo.getExpandedQuery(),
+                        queryInfo.getOptimizerInformation()));
 
         logQueryTimeline(queryInfo);
     }

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryInfo.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.cost.StatsAndCosts;
 import com.facebook.presto.spi.PrestoWarning;
 import com.facebook.presto.spi.QueryId;
+import com.facebook.presto.spi.eventlistener.PlanOptimizerInformation;
 import com.facebook.presto.spi.function.SqlFunctionId;
 import com.facebook.presto.spi.function.SqlInvokedFunction;
 import com.facebook.presto.spi.memory.MemoryPoolId;
@@ -90,6 +91,7 @@ public class QueryInfo
     private final Map<SqlFunctionId, SqlInvokedFunction> addedSessionFunctions;
     private final Set<SqlFunctionId> removedSessionFunctions;
     private final StatsAndCosts planStatsAndCosts;
+    private final List<PlanOptimizerInformation> optimizerInformation;
     // Using a list rather than map, to avoid implementing map key deserializer
     private final List<CanonicalPlanWithInfo> planCanonicalInfo;
 
@@ -130,6 +132,7 @@ public QueryInfo(
             @JsonProperty("addedSessionFunctions") Map<SqlFunctionId, SqlInvokedFunction> addedSessionFunctions,
             @JsonProperty("removedSessionFunctions") Set<SqlFunctionId> removedSessionFunctions,
             @JsonProperty("planStatsAndCosts") StatsAndCosts planStatsAndCosts,
+            @JsonProperty("optimizerInformation") List<PlanOptimizerInformation> optimizerInformation,
             List<CanonicalPlanWithInfo> planCanonicalInfo)
     {
         requireNonNull(queryId, "queryId is null");
@@ -159,6 +162,7 @@ public QueryInfo(
         requireNonNull(addedSessionFunctions, "addedSessionFunctions is null");
         requireNonNull(removedSessionFunctions, "removedSessionFunctions is null");
         requireNonNull(planStatsAndCosts, "planStatsAndCosts is null");
+        requireNonNull(optimizerInformation, "optimizerInformation is null");
 
         this.queryId = queryId;
         this.session = session;
@@ -200,6 +204,7 @@ public QueryInfo(
         this.addedSessionFunctions = ImmutableMap.copyOf(addedSessionFunctions);
         this.removedSessionFunctions = ImmutableSet.copyOf(removedSessionFunctions);
         this.planStatsAndCosts = planStatsAndCosts;
+        this.optimizerInformation = optimizerInformation;
         this.planCanonicalInfo = planCanonicalInfo == null ? ImmutableList.of() : planCanonicalInfo;
     }
 
@@ -429,6 +434,12 @@ public StatsAndCosts getPlanStatsAndCosts()
         return planStatsAndCosts;
     }
 
+    @JsonProperty
+    public List<PlanOptimizerInformation> getOptimizerInformation()
+    {
+        return optimizerInformation;
+    }
+
     // Don't serialize this field because it can be big
     public List<CanonicalPlanWithInfo> getPlanCanonicalInfo()
     {

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryStateMachine.java
@@ -482,6 +482,7 @@ public QueryInfo getQueryInfo(Optional<StageInfo> rootStage)
                 addedSessionFunctions,
                 removedSessionFunctions,
                 Optional.ofNullable(planStatsAndCosts.get()).orElseGet(StatsAndCosts::empty),
+                session.getOptimizerInformationCollector().getOptimizationInfo(),
                 Optional.ofNullable(planCanonicalInfo.get()).orElseGet(ImmutableList::of));
     }
 
@@ -1057,6 +1058,7 @@ public void pruneQueryInfo()
                 queryInfo.getAddedSessionFunctions(),
                 queryInfo.getRemovedSessionFunctions(),
                 StatsAndCosts.empty(),
+                queryInfo.getOptimizerInformation(),
                 ImmutableList.of());
         finalQueryInfo.compareAndSet(finalInfo, Optional.of(prunedQueryInfo));
     }

diff --git a/...ain/java/com/facebook/presto/sql/planner/optimizations/OptimizerInformationCollector.java b/...ain/java/com/facebook/presto/sql/planner/optimizations/OptimizerInformationCollector.java
@@ -0,0 +1,34 @@
+/*
+ * 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 com.facebook.presto.sql.planner.optimizations;
+
+import com.facebook.presto.spi.eventlistener.PlanOptimizerInformation;
+
+import java.util.LinkedList;
+import java.util.List;
+
+public class OptimizerInformationCollector
+{
+    private final List<PlanOptimizerInformation> optimizationInfo = new LinkedList<PlanOptimizerInformation>();
+
+    public void addInformation(PlanOptimizerInformation optimizerInformation)
+    {
+        this.optimizationInfo.add(optimizerInformation);
+    }
+
+    public List<PlanOptimizerInformation> getOptimizationInfo()
+    {
+        return optimizationInfo;
+    }
+}
diff --git a/presto-main/src/test/java/com/facebook/presto/execution/TestQueryInfo.java b/presto-main/src/test/java/com/facebook/presto/execution/TestQueryInfo.java
@@ -191,6 +191,7 @@ private static QueryInfo createQueryInfo()
                 ImmutableMap.of(),
                 ImmutableSet.of(),
                 StatsAndCosts.empty(),
+                ImmutableList.of(),
                 ImmutableList.of(new CanonicalPlanWithInfo(
                         new CanonicalPlan(
                                 new ValuesNode(Optional.empty(), new PlanNodeId("0"), ImmutableList.of(), ImmutableList.of(), Optional.empty()),

diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestBasicQueryInfo.java
@@ -145,6 +145,7 @@ public void testConstructor()
                         ImmutableMap.of(),
                         ImmutableSet.of(),
                         StatsAndCosts.empty(),
+                        ImmutableList.of(),
                         ImmutableList.of()));
 
         assertEquals(basicInfo.getQueryId().getId(), "0");

diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java
@@ -299,6 +299,7 @@ private QueryInfo createQueryInfo(String queryId, ResourceGroupId resourceGroupI
                 ImmutableMap.of(),
                 ImmutableSet.of(),
                 StatsAndCosts.empty(),
+                ImmutableList.of(),
                 ImmutableList.of());
     }
 }
diff --git a/...-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java b/...-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkQueryExecutionFactory.java
@@ -345,6 +345,7 @@ public static QueryInfo createQueryInfo(
                 ImmutableMap.of(),
                 ImmutableSet.of(),
                 StatsAndCosts.empty(),
+                ImmutableList.of(),
                 ImmutableList.of());
     }
 

diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/PlanOptimizerInformation.java b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/PlanOptimizerInformation.java
@@ -0,0 +1,61 @@
+/*
+ * 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 com.facebook.presto.spi.eventlistener;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Optional;
+
+import static java.util.Objects.requireNonNull;
+
+public class PlanOptimizerInformation
+{
+    // Name of the optimizer, we use the class name of the optimizer here.
+    private final String optimizerName;
+    // True when the optimizer makes changes to query plan, otherwise false
+    private final boolean optimizerTriggered;
+    // For optimizers which are not enabled. True if the query matches the pattern of the optimizer and could be applied.
+    // False if cannot be applied. Empty if information not available.
+    private final Optional<Boolean> optimizerApplicable;
+
+    @JsonCreator
+    public PlanOptimizerInformation(
+            @JsonProperty("optimizerName") String optimizerName,
+            @JsonProperty("optimizerTriggered") boolean optimizerTriggered,
+            @JsonProperty("optimizerApplicable") Optional<Boolean> optimizerApplicable)
+    {
+        this.optimizerName = requireNonNull(optimizerName, "optimizerName is null");
+        this.optimizerTriggered = requireNonNull(optimizerTriggered, "optimizerTriggered is null");
+        this.optimizerApplicable = requireNonNull(optimizerApplicable, "optimizerApplicable is null");
+    }
+
+    @JsonProperty
+    public String getOptimizerName()
+    {
+        return optimizerName;
+    }
+
+    @JsonProperty
+    public boolean getOptimizerTriggered()
+    {
+        return optimizerTriggered;
+    }
+
+    @JsonProperty
+    public Optional<Boolean> getOptimizerApplicable()
+    {
+        return optimizerApplicable;
+    }
+}
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryCompletedEvent.java b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryCompletedEvent.java
@@ -42,6 +42,7 @@ public class QueryCompletedEvent
     private final Instant executionStartTime;
     private final Instant endTime;
     private final Optional<String> expandedQuery;
+    private final List<PlanOptimizerInformation> optimizerInformation;
 
     public QueryCompletedEvent(
             QueryMetadata metadata,
@@ -59,7 +60,8 @@ public QueryCompletedEvent(
             List<OperatorStatistics> operatorStatistics,
             List<PlanStatisticsWithSourceInfo> planStatisticsRead,
             List<PlanStatisticsWithSourceInfo> planStatisticsWritten,
-            Optional<String> expandedQuery)
+            Optional<String> expandedQuery,
+            List<PlanOptimizerInformation> optimizerInformation)
     {
         this.metadata = requireNonNull(metadata, "metadata is null");
         this.statistics = requireNonNull(statistics, "statistics is null");
@@ -77,6 +79,7 @@ public QueryCompletedEvent(
         this.planStatisticsRead = requireNonNull(planStatisticsRead, "planStatisticsRead is null");
         this.planStatisticsWritten = requireNonNull(planStatisticsWritten, "planStatisticsWritten is null");
         this.expandedQuery = requireNonNull(expandedQuery, "expandedQuery is null");
+        this.optimizerInformation = requireNonNull(optimizerInformation, "optimizerInformation is null");
     }
 
     public QueryMetadata getMetadata()
@@ -158,4 +161,9 @@ public Optional<String> getExpandedQuery()
     {
         return expandedQuery;
     }
+
+    public List<PlanOptimizerInformation> getOptimizerInformation()
+    {
+        return optimizerInformation;
+    }
 }