Use multiple workers in MSQ WF drill test suite (#16949)

This commit is contained in:
Akshat Jain 2024-08-26 11:34:40 +05:30 committed by GitHub
parent 418da92228
commit 72f8e79a42
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
15 changed files with 174 additions and 148 deletions

View File

@ -20,6 +20,7 @@
package org.apache.druid.msq.exec;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.ImmutableMap;
import com.google.inject.Injector;
import com.google.inject.Module;
import org.apache.druid.guice.DruidInjectorBuilder;
@ -29,17 +30,32 @@ import org.apache.druid.msq.test.CalciteMSQTestsHelper;
import org.apache.druid.msq.test.ExtractResultsFactory;
import org.apache.druid.msq.test.MSQTestOverlordServiceClient;
import org.apache.druid.msq.test.VerifyMSQSupportedNativeQueriesPredicate;
import org.apache.druid.msq.util.MultiStageQueryContext;
import org.apache.druid.query.QueryContexts;
import org.apache.druid.query.groupby.TestGroupByBuffers;
import org.apache.druid.server.QueryLifecycleFactory;
import org.apache.druid.sql.calcite.DrillWindowQueryTest;
import org.apache.druid.sql.calcite.QueryTestBuilder;
import org.apache.druid.sql.calcite.SqlTestFrameworkConfig;
import org.apache.druid.sql.calcite.TempDirProducer;
import org.apache.druid.sql.calcite.planner.PlannerCaptureHook;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import org.apache.druid.sql.calcite.run.SqlEngine;
import org.junit.jupiter.api.Test;
import java.util.HashMap;
import java.util.Map;
@SqlTestFrameworkConfig.ComponentSupplier(DrillWindowQueryMSQComponentSupplier.class)
public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
{
private final Map<String, Object> queryContext = new HashMap<>(ImmutableMap.of(
PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
QueryContexts.ENABLE_DEBUG, true,
MultiStageQueryContext.CTX_MAX_NUM_TASKS, 5
));
public static class DrillWindowQueryMSQComponentSupplier extends DrillComponentSupplier
{
public DrillWindowQueryMSQComponentSupplier(TempDirProducer tempFolderProducer)
@ -74,4 +90,91 @@ public class MSQDrillWindowQueryTest extends DrillWindowQueryTest
.skipVectorize(true)
.verifyNativeQueries(new VerifyMSQSupportedNativeQueriesPredicate());
}
@Override
protected Map<String, Object> getQueryContext()
{
return queryContext;
}
@Override
@DrillTest("druid_queries/empty_over_clause/multiple_empty_over_1")
@Test
public void test_empty_over_multiple_empty_over_1()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/empty_over_clause/single_empty_over_1")
@Test
public void test_empty_over_single_empty_over_1()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/empty_over_clause/single_empty_over_2")
@Test
public void test_empty_over_single_empty_over_2()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/empty_and_non_empty_over/wikipedia_query_1")
@Test
public void test_empty_and_non_empty_over_wikipedia_query_1()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/empty_and_non_empty_over/wikipedia_query_2")
@Test
public void test_empty_and_non_empty_over_wikipedia_query_2()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/empty_and_non_empty_over/wikipedia_query_3")
@Test
public void test_empty_and_non_empty_over_wikipedia_query_3()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/over_clause_only_partitioning/multiple_over_multiple_partition_columns_2")
@Test
public void test_over_clause_with_only_partitioning_multiple_over_multiple_partition_columns_2()
{
useSingleWorker();
windowQueryTest();
}
@Override
@DrillTest("druid_queries/over_clause_only_partitioning/multiple_over_different_partition_column")
@Test
public void test_over_clause_with_only_partitioning_multiple_over_different_partition_column()
{
useSingleWorker();
windowQueryTest();
}
/*
Queries having window functions can give multiple correct results because of using MixShuffleSpec in the previous stage.
So we want to use a single worker to get the same result everytime for such test cases.
*/
private void useSingleWorker()
{
queryContext.put(MultiStageQueryContext.CTX_MAX_NUM_TASKS, 2);
}
}

View File

@ -69,8 +69,10 @@ import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
@ -385,14 +387,12 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
DrillTestCase testCase = drillTestCaseRule.testCase;
thread.setName("drillWindowQuery-" + testCase.filename);
final Map<String, Object> queryContext = new HashMap<>(testBuilder().getQueryContext());
queryContext.putAll(getQueryContext());
testBuilder()
.skipVectorize(true)
.queryContext(ImmutableMap.of(
PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
QueryContexts.ENABLE_DEBUG, true
)
)
.queryContext(queryContext)
.sql(testCase.getQueryString())
.expectedResults(new TextualResultsVerifier(testCase.getExpectedResults(), null))
.run();
@ -404,6 +404,14 @@ public class DrillWindowQueryTest extends BaseCalciteQueryTest
}
}
protected Map<String, Object> getQueryContext()
{
return ImmutableMap.of(
PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
PlannerCaptureHook.NEED_CAPTURE_HOOK, true,
QueryContexts.ENABLE_DEBUG, true
);
}
// testcases_start
@DrillTest("aggregates/aggOWnFn_11")

View File

@ -1,13 +1,6 @@
Austria null #de.wikipedia 1 null
Austria Horsching #de.wikipedia 2 null
Austria Vienna #de.wikipedia 3 Horsching
Austria Vienna #es.wikipedia 5 Vienna
Austria Vienna #tr.wikipedia 13 Vienna
Republic of Korea null #en.wikipedia 4 Vienna
Republic of Korea null #ja.wikipedia 6 null
Republic of Korea null #ko.wikipedia 7 null
Republic of Korea Jeonju #ko.wikipedia 8 null
Republic of Korea Seongnam-si #ko.wikipedia 9 Jeonju
Republic of Korea Seoul #ko.wikipedia 10 Seongnam-si
Republic of Korea Suwon-si #ko.wikipedia 11 Seoul
Republic of Korea Yongsan-dong #ko.wikipedia 12 Suwon-si
Guatemala null Sport Lisboa e Benfica -167 5 null
El Salvador San Salvador Wendigo -1 6 null
Guatemala El Salvador Fuerzas Armadas de Honduras 1 3 San Salvador
El Salvador San Salvador Pedro de Valdivia 9 4 El Salvador
El Salvador San Salvador Amor de barrio 106 1 San Salvador
Guatemala Guatemala City Anexo:Segunda temporada de Gravity Falls 173 2 San Salvador

View File

@ -1,7 +1,7 @@
select
countryName, cityName, channel,
row_number() over (order by channel) as c1,
lag(cityName) over (order by countryName) as c2
countryName, cityName, page, delta,
row_number() over (order by page) as c1,
lag(cityName) over (order by delta) as c2
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel
where countryName in ('Guatemala', 'El Salvador')
group by countryName, cityName, page, delta

View File

@ -1,74 +1,6 @@
Austria null #de.wikipedia 0 1 null
Austria Horsching #de.wikipedia 0 7 null
Austria Vienna #de.wikipedia 0 8 Horsching
Austria null #de.wikipedia 14 2 Vienna
Austria Vienna #de.wikipedia 72 9 null
Austria null #de.wikipedia 94 3 Vienna
Austria null #de.wikipedia 272 4 null
Austria null #de.wikipedia 4685 5 null
Austria null #de.wikipedia 6979 6 null
Republic of Korea null #en.wikipedia 0 12 null
Republic of Korea null #en.wikipedia 5 13 null
Republic of Korea null #en.wikipedia 75 14 null
Austria Vienna #es.wikipedia 0 10 null
Republic of Korea null #ja.wikipedia 0 15 Vienna
Republic of Korea null #ja.wikipedia 43 16 null
Republic of Korea null #ja.wikipedia 46 17 null
Republic of Korea null #ja.wikipedia 356 18 null
Republic of Korea null #ko.wikipedia 0 19 null
Republic of Korea Seoul #ko.wikipedia 0 59 null
Republic of Korea Yongsan-dong #ko.wikipedia 0 74 Seoul
Republic of Korea null #ko.wikipedia 1 20 Yongsan-dong
Republic of Korea Jeonju #ko.wikipedia 1 57 null
Republic of Korea null #ko.wikipedia 2 21 Jeonju
Republic of Korea Seoul #ko.wikipedia 2 60 null
Republic of Korea null #ko.wikipedia 3 22 Seoul
Republic of Korea null #ko.wikipedia 4 23 null
Republic of Korea Seoul #ko.wikipedia 4 61 null
Republic of Korea Seoul #ko.wikipedia 5 62 Seoul
Republic of Korea Seoul #ko.wikipedia 6 63 Seoul
Republic of Korea null #ko.wikipedia 8 24 Seoul
Republic of Korea null #ko.wikipedia 11 25 null
Republic of Korea Seoul #ko.wikipedia 12 64 null
Republic of Korea null #ko.wikipedia 13 26 Seoul
Republic of Korea null #ko.wikipedia 14 27 null
Republic of Korea Seoul #ko.wikipedia 15 65 null
Republic of Korea null #ko.wikipedia 19 28 Seoul
Republic of Korea Seoul #ko.wikipedia 19 66 null
Republic of Korea null #ko.wikipedia 22 29 Seoul
Republic of Korea Seoul #ko.wikipedia 22 67 null
Republic of Korea null #ko.wikipedia 23 30 Seoul
Republic of Korea null #ko.wikipedia 24 31 null
Republic of Korea Seoul #ko.wikipedia 24 68 null
Republic of Korea null #ko.wikipedia 26 32 Seoul
Republic of Korea null #ko.wikipedia 29 33 null
Republic of Korea null #ko.wikipedia 30 34 null
Republic of Korea null #ko.wikipedia 34 35 null
Republic of Korea Suwon-si #ko.wikipedia 40 72 null
Republic of Korea null #ko.wikipedia 41 36 Suwon-si
Republic of Korea null #ko.wikipedia 42 37 null
Republic of Korea null #ko.wikipedia 51 38 null
Republic of Korea null #ko.wikipedia 52 39 null
Republic of Korea null #ko.wikipedia 56 40 null
Republic of Korea null #ko.wikipedia 59 41 null
Republic of Korea null #ko.wikipedia 62 42 null
Republic of Korea null #ko.wikipedia 65 43 null
Republic of Korea null #ko.wikipedia 73 44 null
Republic of Korea null #ko.wikipedia 159 45 null
Republic of Korea Seoul #ko.wikipedia 243 69 null
Republic of Korea Seongnam-si #ko.wikipedia 254 58 Seoul
Republic of Korea null #ko.wikipedia 357 46 Seongnam-si
Republic of Korea null #ko.wikipedia 390 47 null
Republic of Korea null #ko.wikipedia 400 48 null
Republic of Korea null #ko.wikipedia 495 49 null
Republic of Korea null #ko.wikipedia 524 50 null
Republic of Korea null #ko.wikipedia 566 51 null
Republic of Korea Seoul #ko.wikipedia 663 70 null
Republic of Korea null #ko.wikipedia 757 52 Seoul
Republic of Korea null #ko.wikipedia 827 53 null
Republic of Korea null #ko.wikipedia 832 54 null
Republic of Korea null #ko.wikipedia 1006 55 null
Republic of Korea Seoul #ko.wikipedia 1036 71 null
Republic of Korea null #ko.wikipedia 2539 56 Seoul
Republic of Korea Suwon-si #ko.wikipedia 3234 73 null
Austria Vienna #tr.wikipedia 93 11 Suwon-si
Guatemala null #es.wikipedia 0 Sport Lisboa e Benfica -167 4 null
El Salvador San Salvador #es.wikipedia 0 Wendigo -1 3 null
Guatemala El Salvador #es.wikipedia 1 Fuerzas Armadas de Honduras 1 5 San Salvador
El Salvador San Salvador #es.wikipedia 9 Pedro de Valdivia 9 2 El Salvador
El Salvador San Salvador #es.wikipedia 106 Amor de barrio 106 1 San Salvador
Guatemala Guatemala City #es.wikipedia 173 Anexo:Segunda temporada de Gravity Falls 173 6 San Salvador

View File

@ -1,7 +1,7 @@
select
countryName, cityName, channel, added,
row_number() over (order by countryName, cityName) as c1,
lag(cityName) over (order by channel, added) as c2
countryName, cityName, channel, added, page, delta,
row_number() over (order by countryName, cityName, page) as c1,
lag(cityName) over (order by channel, added, delta) as c2
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel, added
where countryName in ('Guatemala', 'El Salvador')
group by countryName, cityName, channel, added, page, delta

View File

@ -1,13 +1,8 @@
Austria null #de.wikipedia 1 null
Austria Horsching #de.wikipedia 2 null
Austria Vienna #de.wikipedia 3 Horsching
Republic of Korea null #en.wikipedia 4 Vienna
Austria Vienna #es.wikipedia 5 null
Republic of Korea null #ja.wikipedia 6 Vienna
Republic of Korea null #ko.wikipedia 7 null
Republic of Korea Jeonju #ko.wikipedia 8 null
Republic of Korea Seongnam-si #ko.wikipedia 9 Jeonju
Republic of Korea Seoul #ko.wikipedia 10 Seongnam-si
Republic of Korea Suwon-si #ko.wikipedia 11 Seoul
Republic of Korea Yongsan-dong #ko.wikipedia 12 Suwon-si
Austria Vienna #tr.wikipedia 13 Yongsan-dong
El Salvador San Salvador 0 1 null
Egypt null 1 2 San Salvador
El Salvador San Salvador 9 3 null
Egypt Cairo 14 4 San Salvador
Egypt null 16 5 Cairo
Egypt null 27 6 null
El Salvador San Salvador 106 7 null
Egypt null 112 8 San Salvador

View File

@ -1,7 +1,7 @@
select
countryName, cityName, channel,
row_number() over (order by channel) as c1,
lag(cityName) over (order by channel) as c2
countryName, cityName, added,
row_number() over (order by added) as c1,
lag(cityName) over (order by added) as c2
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel
where countryName in ('Egypt', 'El Salvador')
group by countryName, cityName, added

View File

@ -1,13 +1,8 @@
Austria null 1
Republic of Korea null 2
Republic of Korea null 3
Republic of Korea null 4
Austria Horsching 5
Republic of Korea Jeonju 6
Republic of Korea Seongnam-si 7
Republic of Korea Seoul 8
Republic of Korea Suwon-si 9
Austria Vienna 10
Austria Vienna 11
Austria Vienna 12
Republic of Korea Yongsan-dong 13
El Salvador 0 1
Egypt 1 2
El Salvador 9 3
Egypt 14 4
Egypt 16 5
Egypt 27 6
El Salvador 106 7
Egypt 112 8

View File

@ -1,6 +1,6 @@
select
countryName, cityName,
row_number() over (order by cityName) as c1
countryName, added,
row_number() over (order by added) as c1
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel
where countryName in ('Egypt', 'El Salvador')
group by countryName, cityName, added

View File

@ -1,6 +1,6 @@
select
countryName, cityName, channel,
lag(cityName) over (order by channel, countryName) as c1
lag(cityName) over (order by channel, countryName, cityName) as c1
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel

View File

@ -1,3 +1,7 @@
Republic of Korea Yongsan-dong #ko.wikipedia 1
Republic of Korea Suwon-si #ko.wikipedia 1
Austria Horsching #de.wikipedia 1
Republic of Korea Seongnam-si #ko.wikipedia 1
Austria null #de.wikipedia 1
Republic of Korea null #en.wikipedia 2
Republic of Korea null #ja.wikipedia 3
@ -5,9 +9,5 @@ Republic of Korea null #ko.wikipedia 4
Republic of Korea Seoul #ko.wikipedia 1
Austria Vienna #de.wikipedia 1
Austria Vienna #es.wikipedia 2
Austria Vienna #tr.wikipedia 3
Republic of Korea Jeonju #ko.wikipedia 4
Republic of Korea Suwon-si #ko.wikipedia 1
Austria Horsching #de.wikipedia 1
Republic of Korea Seongnam-si #ko.wikipedia 1
Republic of Korea Yongsan-dong #ko.wikipedia 1
Republic of Korea Jeonju #ko.wikipedia 3
Austria Vienna #tr.wikipedia 4

View File

@ -1,6 +1,6 @@
select
countryName, cityName, channel,
row_number() over (partition by array[1,2,length(cityName)] order by countryName) as c
row_number() over (partition by array[1,2,length(cityName)] order by channel) as c
from wikipedia
where countryName in ('Austria', 'Republic of Korea')
group by countryName, cityName, channel

View File

@ -1,7 +1,7 @@
SELECT
countryName,
cityName,
ROW_NUMBER() OVER(PARTITION BY countryName),
ROW_NUMBER() OVER(PARTITION BY countryName order by countryName, cityName, channel),
channel,
COUNT(channel) over (PARTITION BY cityName order by countryName, cityName, channel)
FROM wikipedia

View File

@ -1,5 +1,5 @@
SELECT
ROW_NUMBER() OVER(PARTITION BY countryName),
ROW_NUMBER() OVER(PARTITION BY countryName order by countryName, cityName, channel),
countryName,
cityName,
COUNT(channel) over (PARTITION BY cityName order by countryName, cityName, channel),