From a4d40dfc9729354255443837637adfec763a3c1e Mon Sep 17 00:00:00 2001
From: Lyublena Antova <lyublena@fb.com>
Date: Tue, 31 Oct 2023 11:47:20 -0700
Subject: [PATCH] Use in-memory history provider with HiveQueryRunner

To enable adhoc local testing, install the in-memory history provider plugin when creating a HiveQueryRunner

> explain (type distributed) select partkey, count(*) from lineitem group by 1;
                                                                                                                           Query Plan
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Fragment 0 [SINGLE]
     Output layout: [partkey, count]
     Output partitioning: SINGLE []
     Stage Execution Strategy: UNGROUPED_EXECUTION
     - Output[PlanNodeId 8][partkey, _col1] => [partkey:bigint, count:bigint]
             Estimates: {source: HistoryBasedSourceInfo, rows: 2000 (35.16kB), cpu: ?, memory: ?, network: ?}
             _col1 := count (1:44)
         - RemoteSource[1] => [partkey:bigint, count:bigint]

 Fragment 1 [HASH]
     Output layout: [partkey, count]
     Output partitioning: SINGLE []
     Stage Execution Strategy: UNGROUPED_EXECUTION
     - Aggregate(FINAL)[partkey][PlanNodeId 3] => [partkey:bigint, count:bigint]
             Estimates: {source: HistoryBasedSourceInfo, rows: 2000 (35.16kB), cpu: ?, memory: ?, network: ?}
             count := "presto.default.count"((count_8)) (1:44)
         - LocalExchange[PlanNodeId 194][HASH][$hashvalue] (partkey) => [partkey:bigint, count_8:bigint, $hashvalue:bigint]
             - RemoteSource[2] => [partkey:bigint, count_8:bigint, $hashvalue_9:bigint]

 Fragment 2 [SOURCE]
     Output layout: [partkey, count_8, $hashvalue_10]
     Output partitioning: HASH [partkey][$hashvalue_10]
     Stage Execution Strategy: UNGROUPED_EXECUTION
     - Project[PlanNodeId 226][projectLocality = LOCAL] => [partkey:bigint, count_8:bigint, $hashvalue_10:bigint]
             $hashvalue_10 := combine_hash(BIGINT'0', COALESCE($operator$hash_code(partkey), BIGINT'0')) (1:35)
         - Aggregate(PARTIAL)[partkey][PlanNodeId 198] => [partkey:bigint, count_8:bigint]
                 count_8 := "presto.default.count"(*) (1:44)
             - TableScan[PlanNodeId 0][TableHandle {connectorId='hive', connectorHandle='HiveTableHandle{schemaName=tpch, tableName=lineitem, analyzePartitionValues=Optional.empty}', layout='Optional[tpch.lineitem{}]'}, grouped = false] => [partkey:bigint]
                     Estimates: {source: HistoryBasedSourceInfo, rows: 60175 (528.88kB), cpu: 541572.00, memory: 0.00, network: 0.00}
                     LAYOUT: tpch.lineitem{}
                     partkey := partkey:bigint:1:REGULAR (1:58)
---
 .../facebook/presto/hive/HiveQueryRunner.java  | 18 ++++++++++++++++++
 1 file changed, 18 insertions(+)

diff --git a/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java b/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java
index fb3cfc0742255..1541c76fa02bd 100644
--- a/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java
+++ b/presto-hive/src/test/java/com/facebook/presto/hive/HiveQueryRunner.java
@@ -24,9 +24,12 @@
 import com.facebook.presto.hive.metastore.ExtendedHiveMetastore;
 import com.facebook.presto.hive.metastore.MetastoreContext;
 import com.facebook.presto.hive.metastore.file.FileHiveMetastore;
+import com.facebook.presto.spi.Plugin;
 import com.facebook.presto.spi.security.Identity;
 import com.facebook.presto.spi.security.PrincipalType;
 import com.facebook.presto.spi.security.SelectedRole;
+import com.facebook.presto.spi.statistics.HistoryBasedPlanStatisticsProvider;
+import com.facebook.presto.testing.InMemoryHistoryBasedPlanStatisticsProvider;
 import com.facebook.presto.tests.DistributedQueryRunner;
 import com.facebook.presto.tests.tpcds.TpcdsTableName;
 import com.facebook.presto.tpcds.TpcdsPlugin;
@@ -470,6 +473,21 @@ else if (!dataDirectoryFile.canRead() || !dataDirectoryFile.canWrite()) {
         }
 
         DistributedQueryRunner queryRunner = createQueryRunner(TpchTable.getTables(), getAllTpcdsTableNames(), ImmutableMap.of("http-server.http.port", "8080"), dataDirectory);
+
+        try {
+            queryRunner.installPlugin(new Plugin()
+            {
+                @Override
+                public Iterable<HistoryBasedPlanStatisticsProvider> getHistoryBasedPlanStatisticsProviders()
+                {
+                    return ImmutableList.of(new InMemoryHistoryBasedPlanStatisticsProvider());
+                }
+            });
+        }
+        catch (Exception e) {
+            queryRunner.close();
+            throw e;
+        }
         Thread.sleep(10);
         Logger log = Logger.get(DistributedQueryRunner.class);
         log.info("======== SERVER STARTED ========");