feat: Add option to select query types enbled for HBO

presto-docs/src/main/sphinx/optimizer/history-based-optimization.rst

@@ -38,6 +38,7 @@ Configuration Property Name                                   Description
 ``optimizer.history-based-optimizer-timeout``                 Timeout for history based optimizer.                                                                                        ``10 seconds``
 ``optimizer.enforce-timeout-for-hbo-query-registration``      Enforce timeout for query registration in HBO optimizer                                                                     ``False``
 ``optimizer.treat-low-confidence-zero-estimation-as-unknown`` Treat ``LOW`` confidence, zero estimations as ``UNKNOWN`` during joins.                                                     ``False``
+``optimizer.query-types-enabled-for-hbo``                     Query types which are enabled for history based optimization.                                                               ``SELECT,INSERT``
 ``optimizer.confidence-based-broadcast``                      Broadcast based on the confidence of the statistics that are being used, by broadcasting the side of a joinNode which       ``False``
                                                               has the highest confidence statistics. If confidence is the same, then the original behavior will be followed.
 ``optimizer.retry-query-with-history-based-optimization``     Retry a failed query automatically if HBO can help change the existing query plan                                           ``False``
@@ -73,6 +74,8 @@ Session property Name                                       Description
                                                             ``optimizer.confidence-based-broadcast`` in the current session.                                     ``optimizer.confidence-based-broadcast``
 ``retry_query_with_history_based_optimization``             Overrides the behavior of the configuration property
                                                             ``optimizer.retry-query-with-history-based-optimization`` in the current session.                    ``optimizer.retry-query-with-history-based-optimization``
+``query_types_enabled_for_history_based_optimization``      Overrides the behavior of the configuration property
+                                                            ``optimizer.query-types-enabled-for-hbo`` in the current session.                                    ``optimizer.query-types-enabled-for-hbo``
 =========================================================== ==================================================================================================== ==============================================================
 
 Example

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

@@ -17,6 +17,7 @@
 import com.facebook.airlift.units.Duration;
 import com.facebook.presto.common.WarningHandlingLevel;
 import com.facebook.presto.common.plan.PlanCanonicalizationStrategy;
+import com.facebook.presto.common.resourceGroups.QueryType;
 import com.facebook.presto.cost.HistoryBasedOptimizationConfig;
 import com.facebook.presto.execution.QueryManagerConfig;
 import com.facebook.presto.execution.QueryManagerConfig.ExchangeMaterializationStrategy;
@@ -76,6 +77,7 @@
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.JoinReorderingStrategy.ELIMINATE_CROSS_JOINS;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.PartialAggregationStrategy.ALWAYS;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.PartialAggregationStrategy.NEVER;
+import static com.facebook.presto.sql.analyzer.FeaturesConfig.parseQueryTypesFromString;
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static java.lang.Boolean.TRUE;
@@ -271,6 +273,7 @@ public final class SystemSessionProperties
     public static final String ENABLE_VERBOSE_HISTORY_BASED_OPTIMIZER_RUNTIME_STATS = "enable_verbose_history_based_optimizer_runtime_stats";
     public static final String LOG_QUERY_PLANS_USED_IN_HISTORY_BASED_OPTIMIZER = "log_query_plans_used_in_history_based_optimizer";
     public static final String ENFORCE_HISTORY_BASED_OPTIMIZER_REGISTRATION_TIMEOUT = "enforce_history_based_optimizer_register_timeout";
+    public static final String QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION = "query_types_enabled_for_history_based_optimization";
     public static final String MAX_LEAF_NODES_IN_PLAN = "max_leaf_nodes_in_plan";
     public static final String LEAF_NODE_LIMIT_ENABLED = "leaf_node_limit_enabled";
     public static final String PUSH_REMOTE_EXCHANGE_THROUGH_GROUP_ID = "push_remote_exchange_through_group_id";
@@ -1554,6 +1557,18 @@ public SystemSessionProperties(
                         "Enforce timeout for query registration in HBO optimizer",
                         featuresConfig.isEnforceTimeoutForHBOQueryRegistration(),
                         false),
+                new PropertyMetadata<>(
+                        QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION,
+                        format("Query types which are enabled for history based optimization. Specify as a comma-separated string of QueryType values. Allowed options: %s",
+                                Stream.of(QueryType.values())
+                                        .map(QueryType::name)
+                                        .collect(joining(","))),
+                        VARCHAR,
+                        (Class<List<QueryType>>) (Class<?>) List.class,
+                        featuresConfig.getQueryTypesEnabledForHbo(),
+                        false,
+                        value -> parseQueryTypesFromString((String) value),
+                        queryTypes -> ((List<QueryType>) queryTypes).stream().map(QueryType::name).collect(joining(","))),
                 new PropertyMetadata<>(
                         MAX_LEAF_NODES_IN_PLAN,
                         "Maximum number of leaf nodes in the logical plan of SQL statement",
@@ -3022,6 +3037,11 @@ public static List<PlanCanonicalizationStrategy> getHistoryOptimizationPlanCanon
         return strategyList;
     }
 
+    public static List<QueryType> getQueryTypesEnabledForHBO(Session session)
+    {
+        return (List<QueryType>) session.getSystemProperty(QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION, List.class);
+    }
+
     public static boolean enableVerboseHistoryBasedOptimizerRuntimeStats(Session session)
     {
         return session.getSystemProperty(ENABLE_VERBOSE_HISTORY_BASED_OPTIMIZER_RUNTIME_STATS, Boolean.class);

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

@@ -61,11 +61,10 @@
 
 import static com.facebook.presto.SystemSessionProperties.estimateSizeUsingVariablesForHBO;
 import static com.facebook.presto.SystemSessionProperties.getHistoryBasedOptimizerTimeoutLimit;
+import static com.facebook.presto.SystemSessionProperties.getQueryTypesEnabledForHBO;
 import static com.facebook.presto.SystemSessionProperties.trackHistoryBasedPlanStatisticsEnabled;
 import static com.facebook.presto.SystemSessionProperties.trackHistoryStatsFromFailedQuery;
 import static com.facebook.presto.SystemSessionProperties.trackPartialAggregationHistory;
-import static com.facebook.presto.common.resourceGroups.QueryType.INSERT;
-import static com.facebook.presto.common.resourceGroups.QueryType.SELECT;
 import static com.facebook.presto.cost.HistoricalPlanStatisticsUtil.updatePlanStatistics;
 import static com.facebook.presto.cost.HistoryBasedPlanStatisticsManager.historyBasedPlanCanonicalizationStrategyList;
 import static com.facebook.presto.sql.planner.SystemPartitioningHandle.SCALED_WRITER_DISTRIBUTION;
@@ -79,7 +78,6 @@
 public class HistoryBasedPlanStatisticsTracker
 {
     private static final Logger LOG = Logger.get(HistoryBasedPlanStatisticsTracker.class);
-    private static final Set<QueryType> ALLOWED_QUERY_TYPES = ImmutableSet.of(SELECT, INSERT);
 
     private final Supplier<HistoryBasedPlanStatisticsProvider> historyBasedPlanStatisticsProvider;
     private final HistoryBasedStatisticsCacheManager historyBasedStatisticsCacheManager;
@@ -140,7 +138,8 @@ public Map<PlanNodeWithHash, PlanStatisticsWithSourceInfo> getQueryStats(QueryIn
         }
 
         // Only update statistics for SELECT/INSERT queries
-        if (!queryInfo.getQueryType().isPresent() || !ALLOWED_QUERY_TYPES.contains(queryInfo.getQueryType().get())) {
+        List<QueryType> queryTypesEnabled = getQueryTypesEnabledForHBO(session);
+        if (!queryInfo.getQueryType().isPresent() || !queryTypesEnabled.contains(queryInfo.getQueryType().get())) {
             return ImmutableMap.of();
         }
 

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

@@ -22,6 +22,8 @@
 import com.facebook.airlift.units.MaxDataSize;
 import com.facebook.presto.CompressionCodec;
 import com.facebook.presto.common.function.OperatorType;
+import com.facebook.presto.common.resourceGroups.QueryType;
+import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.function.FunctionMetadata;
 import com.facebook.presto.sql.tree.CreateView;
 import com.google.common.annotations.VisibleForTesting;
@@ -36,18 +38,22 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.List;
+import java.util.stream.Stream;
 
 import static com.facebook.airlift.units.DataSize.Unit.KILOBYTE;
 import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
+import static com.facebook.presto.spi.StandardErrorCode.INVALID_SESSION_PROPERTY;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.AggregationPartitioningMergingStrategy.LEGACY;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.JoinNotNullInferenceStrategy.NONE;
 import static com.facebook.presto.sql.analyzer.FeaturesConfig.TaskSpillingStrategy.ORDER_BY_CREATE_TIME;
 import static com.facebook.presto.sql.expressions.ExpressionOptimizerManager.DEFAULT_EXPRESSION_OPTIMIZER_NAME;
 import static com.facebook.presto.sql.tree.CreateView.Security.DEFINER;
 import static com.google.common.collect.ImmutableList.toImmutableList;
+import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MINUTES;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static java.util.stream.Collectors.joining;
 
 @DefunctConfig({
         "resource-group-manager",
@@ -106,6 +112,7 @@ public class FeaturesConfig
     private boolean logPlansUsedInHistoryBasedOptimizer;
     private boolean enforceTimeoutForHBOQueryRegistration;
     private boolean historyBasedOptimizerEstimateSizeUsingVariables;
+    private List<QueryType> queryTypesEnabledForHbo = ImmutableList.of(QueryType.SELECT, QueryType.INSERT);
     private boolean redistributeWrites;
     private boolean scaleWriters = true;
     private DataSize writerMinSize = new DataSize(32, MEGABYTE);
@@ -871,6 +878,33 @@ public FeaturesConfig setHistoryBasedOptimizerPlanCanonicalizationStrategies(Str
         return this;
     }
 
+    @NotNull
+    public List<QueryType> getQueryTypesEnabledForHbo()
+    {
+        return queryTypesEnabledForHbo;
+    }
+
+    @Config("optimizer.query-types-enabled-for-hbo")
+    public FeaturesConfig setQueryTypesEnabledForHbo(String queryTypesEnabledForHbo)
+    {
+        this.queryTypesEnabledForHbo = parseQueryTypesFromString(queryTypesEnabledForHbo);
+        return this;
+    }
+
+    public static List<QueryType> parseQueryTypesFromString(String queryTypes)
+    {
+        try {
+            return Splitter.on(",").trimResults().splitToList(queryTypes).stream()
+                    .map(QueryType::valueOf).collect(toImmutableList());
+        }
+        catch (Exception e) {
+            throw new PrestoException(INVALID_SESSION_PROPERTY, format("Allowed options for query_types_enabled_for_history_based_optimization are: %s",
+                    Stream.of(QueryType.values())
+                            .map(QueryType::name)
+                            .collect(joining(","))));
+        }
+    }
+
     public boolean isLogPlansUsedInHistoryBasedOptimizer()
     {
         return logPlansUsedInHistoryBasedOptimizer;

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

@@ -85,6 +85,7 @@ public void testDefaults()
                 .setHistoryCanonicalPlanNodeLimit(1000)
                 .setHistoryBasedOptimizerTimeout(new Duration(10, SECONDS))
                 .setHistoryBasedOptimizerPlanCanonicalizationStrategies("IGNORE_SAFE_CONSTANTS")
+                .setQueryTypesEnabledForHbo("SELECT,INSERT")
                 .setLogPlansUsedInHistoryBasedOptimizer(false)
                 .setEnforceTimeoutForHBOQueryRegistration(false)
                 .setHistoryBasedOptimizerEstimateSizeUsingVariables(false)
@@ -314,6 +315,7 @@ public void testExplicitPropertyMappings()
                 .put("optimizer.use-perfectly-consistent-histories", "true")
                 .put("optimizer.history-canonical-plan-node-limit", "2")
                 .put("optimizer.history-based-optimizer-plan-canonicalization-strategies", "IGNORE_SAFE_CONSTANTS,IGNORE_SCAN_CONSTANTS")
+                .put("optimizer.query-types-enabled-for-hbo", "SELECT,INSERT,DELETE")
                 .put("optimizer.log-plans-used-in-history-based-optimizer", "true")
                 .put("optimizer.enforce-timeout-for-hbo-query-registration", "true")
                 .put("optimizer.history-based-optimizer-estimate-size-using-variables", "true")
@@ -527,6 +529,7 @@ public void testExplicitPropertyMappings()
                 .setHistoryCanonicalPlanNodeLimit(2)
                 .setHistoryBasedOptimizerTimeout(new Duration(1, SECONDS))
                 .setHistoryBasedOptimizerPlanCanonicalizationStrategies("IGNORE_SAFE_CONSTANTS,IGNORE_SCAN_CONSTANTS")
+                .setQueryTypesEnabledForHbo("SELECT,INSERT,DELETE")
                 .setLogPlansUsedInHistoryBasedOptimizer(true)
                 .setEnforceTimeoutForHBOQueryRegistration(true)
                 .setHistoryBasedOptimizerEstimateSizeUsingVariables(true)

presto-tests/src/test/java/com/facebook/presto/execution/TestHistoryBasedStatsTracking.java

@@ -56,6 +56,7 @@
 import static com.facebook.presto.SystemSessionProperties.HISTORY_CANONICAL_PLAN_NODE_LIMIT;
 import static com.facebook.presto.SystemSessionProperties.JOIN_DISTRIBUTION_TYPE;
 import static com.facebook.presto.SystemSessionProperties.JOIN_REORDERING_STRATEGY;
+import static com.facebook.presto.SystemSessionProperties.QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION;
 import static com.facebook.presto.SystemSessionProperties.RESTRICT_HISTORY_BASED_OPTIMIZATION_TO_COMPLEX_QUERY;
 import static com.facebook.presto.SystemSessionProperties.TRACK_HISTORY_BASED_PLAN_STATISTICS;
 import static com.facebook.presto.SystemSessionProperties.TRACK_HISTORY_STATS_FROM_FAILED_QUERIES;
@@ -154,6 +155,33 @@ public void testHistoryBasedStatsCalculatorUsingVariables()
                 anyTree(node(FilterNode.class, any()).withOutputRowCount(2).withOutputSize(256)));
     }
 
+    @Test
+    public void testHistoryBasedStatsCalculatorQueryType()
+    {
+        Session insertAndDelete = Session.builder(createSession())
+                .setSystemProperty(QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION, "INSERT,DELETE")
+                .build();
+        Session select = Session.builder(createSession())
+                .setSystemProperty(QUERY_TYPES_ENABLED_FOR_HISTORY_BASED_OPTIMIZATION, "SELECT")
+                .build();
+        // CBO Statistics
+        assertPlan(
+                "SELECT * FROM nation where substr(name, 1, 1) = 'A'",
+                anyTree(node(FilterNode.class, any()).withOutputRowCount(Double.NaN)));
+
+        // Select query type is not enabled, expect no history to be rewritten
+        executeAndNoHistoryWritten("SELECT * FROM nation where substr(name, 1, 1) = 'A'", insertAndDelete);
+        assertPlan(
+                "SELECT * FROM nation where substr(name, 1, 1) = 'A'",
+                anyTree(node(FilterNode.class, any()).withOutputRowCount(Double.NaN)));
+
+        // Select query type is enabled, expect history to be rewritten
+        executeAndTrackHistory("SELECT * FROM nation where substr(name, 1, 1) = 'A'", select);
+        assertPlan(
+                "SELECT * FROM nation where substr(name, 1, 1) = 'A'",
+                anyTree(node(FilterNode.class, any()).withOutputRowCount(2).withOutputSize(199)));
+    }
+
     @Test
     public void testHistoryBasedStatsCalculatorEnforceTimeOut()
     {