Add a flag to disable tracking of partial histories

presto-main/src/main/java/com/facebook/presto/SystemSessionProperties.java

@@ -297,6 +297,7 @@ public final class SystemSessionProperties
     public static final String INFER_INEQUALITY_PREDICATES = "infer_inequality_predicates";
     public static final String ENABLE_HISTORY_BASED_SCALED_WRITER = "enable_history_based_scaled_writer";
     public static final String USE_PARTIAL_AGGREGATION_HISTORY = "use_partial_aggregation_history";
+    public static final String TRACK_PARTIAL_AGGREGATION_HISTORY = "track_partial_aggregation_history";
     public static final String REMOVE_REDUNDANT_CAST_TO_VARCHAR_IN_JOIN = "remove_redundant_cast_to_varchar_in_join";
     public static final String HANDLE_COMPLEX_EQUI_JOINS = "handle_complex_equi_joins";
 
@@ -1790,6 +1791,11 @@ public SystemSessionProperties(
                         "Use collected partial aggregation statistics from HBO",
                         featuresConfig.isUsePartialAggregationHistory(),
                         false),
+                booleanProperty(
+                        TRACK_PARTIAL_AGGREGATION_HISTORY,
+                        "Track partial aggregation statistics in HBO",
+                        featuresConfig.isTrackPartialAggregationHistory(),
+                        false),
                 booleanProperty(
                         REMOVE_REDUNDANT_CAST_TO_VARCHAR_IN_JOIN,
                         "If both left and right side of join clause are varchar cast from int/bigint, remove the cast here",
@@ -2994,6 +3000,11 @@ public static boolean usePartialAggregationHistory(Session session)
         return session.getSystemProperty(USE_PARTIAL_AGGREGATION_HISTORY, Boolean.class);
     }
 
+    public static boolean trackPartialAggregationHistory(Session session)
+    {
+        return session.getSystemProperty(TRACK_PARTIAL_AGGREGATION_HISTORY, Boolean.class);
+    }
+
     public static boolean isRemoveRedundantCastToVarcharInJoinEnabled(Session session)
     {
         return session.getSystemProperty(REMOVE_REDUNDANT_CAST_TO_VARCHAR_IN_JOIN, Boolean.class);

presto-main/src/main/java/com/facebook/presto/cost/HistoryBasedPlanStatisticsTracker.java

@@ -54,6 +54,7 @@
 
 import static com.facebook.presto.SystemSessionProperties.getHistoryBasedOptimizerTimeoutLimit;
 import static com.facebook.presto.SystemSessionProperties.trackHistoryBasedPlanStatisticsEnabled;
+import static com.facebook.presto.SystemSessionProperties.trackPartialAggregationHistory;
 import static com.facebook.presto.common.plan.PlanCanonicalizationStrategy.historyBasedPlanCanonicalizationStrategyList;
 import static com.facebook.presto.common.resourceGroups.QueryType.INSERT;
 import static com.facebook.presto.common.resourceGroups.QueryType.SELECT;
@@ -158,7 +159,7 @@ public Map<PlanNodeWithHash, PlanStatisticsWithSourceInfo> getQueryStats(QueryIn
                 double joinProbeKeyCount = planNodeStats.getPlanNodeJoinProbeKeyCount();
                 PartialAggregationStatistics partialAggregationStatistics = PartialAggregationStatistics.empty();
 
-                if (isAggregation(planNode, AggregationNode.Step.PARTIAL)) {
+                if (isAggregation(planNode, AggregationNode.Step.PARTIAL) && trackPartialAggregationHistory(session)) {
                     // we're doing a depth-first traversal of the plan tree so we must have seen the corresponding final agg already:
                     // find it and update its partial agg stats
                     partialAggregationStatistics = constructAggregationNodeStatistics(planNode, planNodeStatsMap, outputBytes, outputPositions);
@@ -205,7 +206,7 @@ public Map<PlanNodeWithHash, PlanStatisticsWithSourceInfo> getQueryStats(QueryIn
                                 new HistoryBasedSourceInfo(Optional.of(hash), Optional.of(inputTableStatistics)));
                         planStatisticsMap.put(planNodeWithHash, planStatsWithSourceInfo);
 
-                        if (isAggregation(planNode, AggregationNode.Step.FINAL) && ((AggregationNode) planNode).getAggregationId().isPresent()) {
+                        if (isAggregation(planNode, AggregationNode.Step.FINAL) && ((AggregationNode) planNode).getAggregationId().isPresent() && trackPartialAggregationHistory(session)) {
                             // we're doing a depth-first traversal of the plan tree: cache the final agg so that when we encounter the partial agg we can come back
                             // and update the partial agg statistics
                             aggregationNodeMap.put(((AggregationNode) planNode).getAggregationId().get(), new FinalAggregationStatsInfo(planNodeWithHash, planStatsWithSourceInfo));

presto-main/src/main/java/com/facebook/presto/sql/analyzer/FeaturesConfig.java

@@ -288,6 +288,8 @@ public class FeaturesConfig
 
     private boolean usePartialAggregationHistory;
 
+    private boolean trackPartialAggregationHistory = true;
+
     private boolean removeRedundantCastToVarcharInJoin = true;
 
     public enum PartitioningPrecisionStrategy
@@ -2861,6 +2863,19 @@ public FeaturesConfig setUsePartialAggregationHistory(boolean usePartialAggregat
         return this;
     }
 
+    public boolean isTrackPartialAggregationHistory()
+    {
+        return this.trackPartialAggregationHistory;
+    }
+
+    @Config("optimizer.track-partial-aggregation-history")
+    @ConfigDescription("Track partial aggregation histories")
+    public FeaturesConfig setTrackPartialAggregationHistory(boolean trackPartialAggregationHistory)
+    {
+        this.trackPartialAggregationHistory = trackPartialAggregationHistory;
+        return this;
+    }
+
     public boolean isRemoveRedundantCastToVarcharInJoin()
     {
         return removeRedundantCastToVarcharInJoin;

presto-main/src/test/java/com/facebook/presto/sql/analyzer/TestFeaturesConfig.java

@@ -80,6 +80,7 @@ public void testDefaults()
                 .setUseHistoryBasedPlanStatistics(false)
                 .setTrackHistoryBasedPlanStatistics(false)
                 .setUsePartialAggregationHistory(false)
+                .setTrackPartialAggregationHistory(true)
                 .setUsePerfectlyConsistentHistories(false)
                 .setHistoryCanonicalPlanNodeLimit(1000)
                 .setHistoryBasedOptimizerTimeout(new Duration(10, SECONDS))
@@ -301,6 +302,7 @@ public void testExplicitPropertyMappings()
                 .put("optimizer.use-history-based-plan-statistics", "true")
                 .put("optimizer.track-history-based-plan-statistics", "true")
                 .put("optimizer.use-partial-aggregation-history", "true")
+                .put("optimizer.track-partial-aggregation-history", "false")
                 .put("optimizer.use-perfectly-consistent-histories", "true")
                 .put("optimizer.history-canonical-plan-node-limit", "2")
                 .put("optimizer.history-based-optimizer-timeout", "1s")
@@ -489,6 +491,7 @@ public void testExplicitPropertyMappings()
                 .setUseHistoryBasedPlanStatistics(true)
                 .setTrackHistoryBasedPlanStatistics(true)
                 .setUsePartialAggregationHistory(true)
+                .setTrackPartialAggregationHistory(false)
                 .setUsePerfectlyConsistentHistories(true)
                 .setHistoryCanonicalPlanNodeLimit(2)
                 .setHistoryBasedOptimizerTimeout(new Duration(1, SECONDS))

presto-spark-base/src/main/java/com/facebook/presto/spark/PrestoSparkSettingsRequirements.java

@@ -96,6 +96,7 @@ public static void setDefaults(FeaturesConfig config)
         config.setInlineSqlFunctions(true);
         config.setEnforceFixedDistributionForOutputOperator(true);
         config.setPrestoSparkAssignBucketToPartitionForPartitionedTableWriteEnabled(true);
+        config.setTrackPartialAggregationHistory(false);
     }
 
     public static void setDefaults(QueryManagerConfig config)