Skip to content

Commit 4f0c792

Browse files
kevintang2022kevintang2022
committed
Add unit test
1 parent af7425b commit 4f0c792

File tree

5 files changed

+216
-3
lines changed

5 files changed

+216
-3
lines changed

‎presto-main/src/main/java/com/facebook/presto/server/testing/TestingPrestoServer.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@
5151
import com.facebook.presto.memory.LocalMemoryManager;
5252
import com.facebook.presto.metadata.AllNodes;
5353
import com.facebook.presto.metadata.CatalogManager;
54+
import com.facebook.presto.metadata.FunctionAndTypeManager;
5455
import com.facebook.presto.metadata.InternalNode;
5556
import com.facebook.presto.metadata.InternalNodeManager;
5657
import com.facebook.presto.metadata.Metadata;
@@ -147,6 +148,7 @@ public class TestingPrestoServer
147148
private final boolean preserveData;
148149
private final LifeCycleManager lifeCycleManager;
149150
private final PluginManager pluginManager;
151+
private final FunctionAndTypeManager functionAndTypeManager;
150152
private final ConnectorManager connectorManager;
151153
private final TestingHttpServer server;
152154
private final CatalogManager catalogManager;
@@ -367,6 +369,8 @@ public TestingPrestoServer(
367369

368370
connectorManager = injector.getInstance(ConnectorManager.class);
369371

372+
functionAndTypeManager = injector.getInstance(FunctionAndTypeManager.class);
373+
370374
server = injector.getInstance(TestingHttpServer.class);
371375
catalogManager = injector.getInstance(CatalogManager.class);
372376
transactionManager = injector.getInstance(TransactionManager.class);
@@ -501,6 +505,11 @@ private Map<String, String> getServerProperties(
501505
return ImmutableMap.copyOf(serverProperties);
502506
}
503507

508+
public void registerNativeFunctions()
509+
{
510+
functionAndTypeManager.registerNativeFunctions();
511+
}
512+
504513
@Override
505514
public void close()
506515
throws IOException

‎presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,7 @@ public static class HiveQueryRunnerBuilder
123123
private String security;
124124
private boolean addStorageFormatToPath;
125125
private boolean coordinatorSidecarEnabled;
126+
private boolean workerSidecarEnabled;
126127
private boolean enableRuntimeMetricsCollection;
127128
private boolean enableSsdCache;
128129
private boolean failOnNestedLoopJoin;
@@ -218,6 +219,13 @@ public HiveQueryRunnerBuilder setCoordinatorSidecarEnabled(boolean coordinatorSi
218219
return this;
219220
}
220221

222+
223+
public HiveQueryRunnerBuilder setWorkerSidecarEnabled(boolean workerSidecarEnabled)
224+
{
225+
this.workerSidecarEnabled = workerSidecarEnabled;
226+
return this;
227+
}
228+
221229
public HiveQueryRunnerBuilder setStorageFormat(String storageFormat)
222230
{
223231
this.storageFormat = storageFormat;
@@ -261,7 +269,7 @@ public QueryRunner build()
261269
Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher = Optional.empty();
262270
if (this.useExternalWorkerLauncher) {
263271
externalWorkerLauncher = getExternalWorkerLauncher("hive", serverBinary, cacheMaxSize, remoteFunctionServerUds,
264-
failOnNestedLoopJoin, coordinatorSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache);
272+
failOnNestedLoopJoin, coordinatorSidecarEnabled, workerSidecarEnabled, enableRuntimeMetricsCollection, enableSsdCache);
265273
}
266274
return HiveQueryRunner.createQueryRunner(
267275
ImmutableList.of(),
@@ -350,7 +358,7 @@ public QueryRunner build()
350358
Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher = Optional.empty();
351359
if (this.useExternalWorkerLauncher) {
352360
externalWorkerLauncher = getExternalWorkerLauncher("iceberg", serverBinary, cacheMaxSize, remoteFunctionServerUds,
353-
false, false, false, false);
361+
false, false, false,false, false);
354362
}
355363
return IcebergQueryRunner.builder()
356364
.setExtraProperties(extraProperties)
@@ -446,6 +454,7 @@ public static Optional<BiFunction<Integer, URI, Process>> getExternalWorkerLaunc
446454
Optional<String> remoteFunctionServerUds,
447455
Boolean failOnNestedLoopJoin,
448456
boolean isCoordinatorSidecarEnabled,
457+
boolean isWorkerSidecarEnabled,
449458
boolean enableRuntimeMetricsCollection,
450459
boolean enableSsdCache)
451460
{
@@ -468,6 +477,10 @@ public static Optional<BiFunction<Integer, URI, Process>> getExternalWorkerLaunc
468477
"native-sidecar=true%n" +
469478
"presto.default-namespace=native.default%n", configProperties);
470479
}
480+
else if (isWorkerSidecarEnabled) {
481+
configProperties = format("%s%n" +
482+
"native-sidecar=true%n", configProperties);
483+
}
471484

472485
if (enableRuntimeMetricsCollection) {
473486
configProperties = format("%s%n" +

‎presto-native-execution/src/test/java/com/facebook/presto/nativeworker/TestBuiltInNativeFunctions.java‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
15+
package com.facebook.presto.nativeworker;
16+
17+
import com.facebook.presto.cost.CostCalculator;
18+
import com.facebook.presto.cost.CostCalculatorUsingExchanges;
19+
import com.facebook.presto.cost.CostCalculatorWithEstimatedExchanges;
20+
import com.facebook.presto.cost.CostComparator;
21+
import com.facebook.presto.cost.TaskCountEstimator;
22+
import com.facebook.presto.execution.QueryManagerConfig;
23+
import com.facebook.presto.execution.TaskManagerConfig;
24+
import com.facebook.presto.metadata.InMemoryNodeManager;
25+
import com.facebook.presto.metadata.Metadata;
26+
import com.facebook.presto.nodeManager.PluginNodeManager;
27+
import com.facebook.presto.spi.WarningCollector;
28+
import com.facebook.presto.sql.analyzer.FeaturesConfig;
29+
import com.facebook.presto.sql.analyzer.QueryExplainer;
30+
import com.facebook.presto.sql.expressions.ExpressionOptimizerManager;
31+
import com.facebook.presto.sql.planner.PartitioningProviderManager;
32+
import com.facebook.presto.sql.planner.PlanFragmenter;
33+
import com.facebook.presto.sql.planner.PlanOptimizers;
34+
import com.facebook.presto.sql.planner.optimizations.PlanOptimizer;
35+
import com.facebook.presto.sql.planner.sanity.PlanChecker;
36+
import com.facebook.presto.sql.tree.ExplainType;
37+
import com.facebook.presto.testing.QueryRunner;
38+
import com.facebook.presto.tests.AbstractTestQueryFramework;
39+
import com.facebook.presto.tests.DistributedQueryRunner;
40+
import com.google.common.collect.ImmutableMap;
41+
import org.intellij.lang.annotations.Language;
42+
import org.testng.annotations.Test;
43+
import org.weakref.jmx.MBeanExporter;
44+
import org.weakref.jmx.testing.TestingMBeanServer;
45+
46+
import java.util.List;
47+
import java.util.regex.Pattern;
48+
49+
import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createLineitem;
50+
import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createNation;
51+
import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrders;
52+
import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createOrdersEx;
53+
import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.createRegion;
54+
import static com.facebook.presto.transaction.TransactionBuilder.transaction;
55+
import static com.facebook.presto.util.AnalyzerUtil.createParsingOptions;
56+
import static java.util.Collections.emptyList;
57+
import static org.testng.Assert.fail;
58+
59+
public class TestBuiltInNativeFunctions
60+
extends AbstractTestQueryFramework
61+
{
62+
@Override
63+
protected void createTables()
64+
{
65+
QueryRunner queryRunner = (QueryRunner) getExpectedQueryRunner();
66+
createLineitem(queryRunner);
67+
createNation(queryRunner);
68+
createOrders(queryRunner);
69+
createOrdersEx(queryRunner);
70+
createRegion(queryRunner);
71+
}
72+
73+
@Override
74+
protected QueryRunner createQueryRunner()
75+
throws Exception
76+
{
77+
DistributedQueryRunner queryRunner = (DistributedQueryRunner) PrestoNativeQueryRunnerUtils.nativeHiveQueryRunnerBuilder()
78+
.setExtraProperties(ImmutableMap.of("built-in-sidecar-functions-enabled", "true"))
79+
.setAddStorageFormatToPath(true)
80+
.setWorkerSidecarEnabled(true)
81+
.build();
82+
83+
queryRunner.registerNativeFunctions();
84+
85+
return queryRunner;
86+
}
87+
88+
@Override
89+
protected QueryRunner createExpectedQueryRunner()
90+
throws Exception
91+
{
92+
return PrestoNativeQueryRunnerUtils.javaHiveQueryRunnerBuilder()
93+
.setAddStorageFormatToPath(true)
94+
.build();
95+
}
96+
97+
private void assertJsonPlan(@Language("SQL") String query, boolean withBuiltInSidecarEnabled, @Language("RegExp") String jsonPlanRegex, boolean shouldContainRegex)
98+
{
99+
QueryRunner queryRunner;
100+
if (withBuiltInSidecarEnabled) {
101+
queryRunner = getQueryRunner();
102+
}
103+
else {
104+
queryRunner = (QueryRunner) getExpectedQueryRunner();
105+
}
106+
107+
QueryExplainer explainer = getQueryExplainerFromProvidedQueryRunner(queryRunner);
108+
transaction(queryRunner.getTransactionManager(), queryRunner.getAccessControl())
109+
.singleStatement()
110+
.execute(queryRunner.getDefaultSession(), transactionSession -> {
111+
String actualPlan = explainer.getJsonPlan(transactionSession, getSqlParser().createStatement(query, createParsingOptions(transactionSession)), ExplainType.Type.LOGICAL, emptyList(), WarningCollector.NOOP, query);
112+
Pattern p = Pattern.compile(jsonPlanRegex, Pattern.MULTILINE);
113+
if (shouldContainRegex) {
114+
if (!p.matcher(actualPlan).find()) {
115+
fail("Query plan text does not contain regex");
116+
}
117+
}
118+
else {
119+
if (p.matcher(actualPlan).find()) {
120+
fail("Query plan text contains bad pattern");
121+
}
122+
}
123+
124+
return null;
125+
});
126+
}
127+
128+
private QueryExplainer getQueryExplainerFromProvidedQueryRunner(QueryRunner queryRunner)
129+
{
130+
Metadata metadata = queryRunner.getMetadata();
131+
FeaturesConfig featuresConfig = createFeaturesConfig();
132+
boolean noExchange = queryRunner.getNodeCount() == 1;
133+
TaskCountEstimator taskCountEstimator = new TaskCountEstimator(queryRunner::getNodeCount);
134+
CostCalculator costCalculator = new CostCalculatorUsingExchanges(taskCountEstimator);
135+
List<PlanOptimizer> optimizers = new PlanOptimizers(
136+
metadata,
137+
getSqlParser(),
138+
noExchange,
139+
new MBeanExporter(new TestingMBeanServer()),
140+
queryRunner.getSplitManager(),
141+
queryRunner.getPlanOptimizerManager(),
142+
queryRunner.getPageSourceManager(),
143+
queryRunner.getStatsCalculator(),
144+
costCalculator,
145+
new CostCalculatorWithEstimatedExchanges(costCalculator, taskCountEstimator),
146+
new CostComparator(featuresConfig),
147+
taskCountEstimator,
148+
new PartitioningProviderManager(),
149+
featuresConfig,
150+
new ExpressionOptimizerManager(
151+
new PluginNodeManager(new InMemoryNodeManager()),
152+
queryRunner.getMetadata().getFunctionAndTypeManager()),
153+
new TaskManagerConfig())
154+
.getPlanningTimeOptimizers();
155+
return new QueryExplainer(
156+
optimizers,
157+
new PlanFragmenter(metadata, queryRunner.getNodePartitioningManager(), new QueryManagerConfig(), featuresConfig, queryRunner.getPlanCheckerProviderManager()),
158+
metadata,
159+
queryRunner.getAccessControl(),
160+
getSqlParser(),
161+
queryRunner.getStatsCalculator(),
162+
costCalculator,
163+
ImmutableMap.of(),
164+
new PlanChecker(featuresConfig, false, queryRunner.getPlanCheckerProviderManager()));
165+
}
166+
167+
@Test
168+
public void testUdfQueries()
169+
{
170+
assertQuery("SELECT ARRAY['abc']");
171+
assertQuery("SELECT ARRAY[1, 2, 3]");
172+
assertQuery("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
173+
assertQuery("SELECT presto.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )");
174+
assertQueryFails("SELECT native.default.map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", ".*Function native.default.map_remove_null_values not registered.*");
175+
assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", true, "lambda", false);
176+
assertJsonPlan("SELECT map_remove_null_values( MAP( ARRAY['a', 'b', 'c'], ARRAY[1, NULL, 3] ) )", false, "lambda", true);
177+
}
178+
}

‎presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueryFramework.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -565,7 +565,7 @@ protected SubPlan subplan(String sql, Session session)
565565
}
566566
}
567567

568-
private QueryExplainer getQueryExplainer()
568+
protected QueryExplainer getQueryExplainer()
569569
{
570570
Metadata metadata = queryRunner.getMetadata();
571571
FeaturesConfig featuresConfig = createFeaturesConfig();
@@ -630,6 +630,12 @@ protected ExpectedQueryRunner getExpectedQueryRunner()
630630
return expectedQueryRunner;
631631
}
632632

633+
protected SqlParser getSqlParser()
634+
{
635+
checkState(sqlParser != null, "sqlParser not set");
636+
return sqlParser;
637+
}
638+
633639
public interface QueryRunnerSupplier
634640
{
635641
QueryRunner get()

‎presto-tests/src/main/java/com/facebook/presto/tests/DistributedQueryRunner.java‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1040,6 +1040,13 @@ public void loadPlanCheckerProviderManager(String planCheckerProviderName, Map<S
10401040
}
10411041
}
10421042

1043+
public void registerNativeFunctions()
1044+
{
1045+
for (TestingPrestoServer server : servers) {
1046+
server.registerNativeFunctions();
1047+
}
1048+
}
1049+
10431050
private static void closeUnchecked(AutoCloseable closeable)
10441051
{
10451052
try {

0 commit comments

Comments
 (0)