Skip to content

Run engine-only tests with task level retries #11224

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
losipiuk merged 9 commits into from
Mar 2, 2022
Merged

Run engine-only tests with task level retries #11224

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
ec687ab
Rename test class
losipiuk Feb 28, 2022
21699d5
Reintroduce test class extending from abstract
losipiuk Feb 28, 2022
373a1c0
Move Abstract*TestEngineOnlyQueries to trino-testing
losipiuk Feb 28, 2022
e71757b
Extract dynamic filter explain analyze test method
losipiuk Feb 28, 2022
4981671
Move exchange setup to top-level DistributedQueryRunner
losipiuk Feb 28, 2022
cbb3793
Add builder to MemoryQueryRunner
losipiuk Feb 28, 2022
a3a007a
Support retries in Memory connector
losipiuk Feb 28, 2022
a422074
Run engine-only tests with task level retries
losipiuk Feb 28, 2022
6603cdb
Support replicated sources for arbitrary exchange
losipiuk Feb 28, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 40 additions & 23 deletions core/trino-main/src/main/java/io/trino/execution/scheduler/StageTaskSourceFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand Down Expand Up @@ -209,9 +210,9 @@ public void close()
public static class ArbitraryDistributionTaskSource
implements TaskSource
{
private final Map<PlanFragmentId, PlanNodeId> sourceFragmentToRemoteSourceNodeIdMap;
private final Map<PlanFragmentId, Exchange> sourceExchanges;
private final Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles;
private final IdentityHashMap<ExchangeSourceHandle, Exchange> sourceExchanges;
private final Multimap<PlanNodeId, ExchangeSourceHandle> partitionedExchangeSourceHandles;
private final Multimap<PlanNodeId, ExchangeSourceHandle> replicatedExchangeSourceHandles;
private final long targetPartitionSizeInBytes;

private boolean finished;
Expand All @@ -223,33 +224,33 @@ public static ArbitraryDistributionTaskSource create(
DataSize targetPartitionSize)
{
checkArgument(fragment.getPartitionedSources().isEmpty(), "no partitioned sources (table scans) expected, got: %s", fragment.getPartitionedSources());
checkArgument(fragment.getRemoteSourceNodes().stream().noneMatch(node -> node.getExchangeType() == REPLICATE), "replicated exchanges are not expected in source distributed stage, got: %s", fragment.getRemoteSourceNodes());
IdentityHashMap<ExchangeSourceHandle, Exchange> exchangeForHandleMap = getExchangeForHandleMap(sourceExchanges, exchangeSourceHandles);

return new ArbitraryDistributionTaskSource(
getSourceFragmentToRemoteSourceNodeIdMap(fragment.getRemoteSourceNodes()),
sourceExchanges,
exchangeSourceHandles,
exchangeForHandleMap,
getPartitionedExchangeSourceHandles(fragment, exchangeSourceHandles),
getReplicatedExchangeSourceHandles(fragment, exchangeSourceHandles),
targetPartitionSize);
}

public ArbitraryDistributionTaskSource(
Map<PlanFragmentId, PlanNodeId> sourceFragmentToRemoteSourceNodeIdMap,
Map<PlanFragmentId, Exchange> sourceExchanges,
Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles,
IdentityHashMap<ExchangeSourceHandle, Exchange> sourceExchanges,
Multimap<PlanNodeId, ExchangeSourceHandle> partitionedExchangeSourceHandles,
Multimap<PlanNodeId, ExchangeSourceHandle> replicatedExchangeSourceHandles,
DataSize targetPartitionSize)
{
this.sourceFragmentToRemoteSourceNodeIdMap = ImmutableMap.copyOf(requireNonNull(sourceFragmentToRemoteSourceNodeIdMap, "sourceFragmentToRemoteSourceNodeIdMap is null"));
this.sourceExchanges = ImmutableMap.copyOf(requireNonNull(sourceExchanges, "sourceExchanges is null"));
this.sourceExchanges = new IdentityHashMap<>(requireNonNull(sourceExchanges, "sourceExchanges is null"));
this.partitionedExchangeSourceHandles = ImmutableListMultimap.copyOf(requireNonNull(partitionedExchangeSourceHandles, "partitionedExchangeSourceHandles is null"));
this.replicatedExchangeSourceHandles = ImmutableListMultimap.copyOf(requireNonNull(replicatedExchangeSourceHandles, "replicatedExchangeSourceHandles is null"));
checkArgument(
sourceFragmentToRemoteSourceNodeIdMap.keySet().equals(sourceExchanges.keySet()),
"sourceFragmentToRemoteSourceNodeIdMap and sourceExchanges are expected to have the same set of keys: %s != %s",
sourceFragmentToRemoteSourceNodeIdMap.keySet(),
sourceExchanges.keySet().containsAll(partitionedExchangeSourceHandles.values()),
"Unexpected entries in partitionedExchangeSourceHandles map: %s; allowed keys: %s",
partitionedExchangeSourceHandles.values(),
sourceExchanges.keySet());
this.exchangeSourceHandles = ImmutableListMultimap.copyOf(requireNonNull(exchangeSourceHandles, "exchangeSourceHandles is null"));
checkArgument(
sourceExchanges.keySet().containsAll(exchangeSourceHandles.keySet()),
"Unexpected keys in exchangeSourceHandles map: %s; allowed keys: %s",
exchangeSourceHandles.keySet(),
sourceExchanges.keySet().containsAll(replicatedExchangeSourceHandles.values()),
"Unexpected entries in replicatedExchangeSourceHandles map: %s; allowed keys: %s",
replicatedExchangeSourceHandles.values(),
sourceExchanges.keySet());
this.targetPartitionSizeInBytes = requireNonNull(targetPartitionSize, "targetPartitionSize is null").toBytes();
}
Expand All @@ -266,11 +267,10 @@ public List<TaskDescriptor> getMoreTasks()
long assignedExchangeDataSize = 0;
int assignedExchangeSourceHandleCount = 0;

for (Map.Entry<PlanFragmentId, ExchangeSourceHandle> entry : exchangeSourceHandles.entries()) {
PlanFragmentId sourceFragmentId = entry.getKey();
PlanNodeId remoteSourcePlanNodeId = sourceFragmentToRemoteSourceNodeIdMap.get(sourceFragmentId);
for (Map.Entry<PlanNodeId, ExchangeSourceHandle> entry : partitionedExchangeSourceHandles.entries()) {
PlanNodeId remoteSourcePlanNodeId = entry.getKey();
ExchangeSourceHandle originalExchangeSourceHandle = entry.getValue();
Exchange sourceExchange = sourceExchanges.get(sourceFragmentId);
Exchange sourceExchange = sourceExchanges.get(originalExchangeSourceHandle);

ExchangeSourceSplitter splitter = sourceExchange.split(originalExchangeSourceHandle, targetPartitionSizeInBytes);
ImmutableList.Builder<ExchangeSourceHandle> sourceHandles = ImmutableList.builder();
Expand All @@ -286,6 +286,7 @@ public List<TaskDescriptor> getMoreTasks()
for (ExchangeSourceHandle handle : sourceHandles.build()) {
ExchangeSourceStatistics statistics = sourceExchange.getExchangeSourceStatistics(handle);
if (assignedExchangeDataSize != 0 && assignedExchangeDataSize + statistics.getSizeInBytes() > targetPartitionSizeInBytes) {
assignedExchangeSourceHandles.putAll(replicatedExchangeSourceHandles);
result.add(new TaskDescriptor(currentPartitionId++, ImmutableListMultimap.of(), assignedExchangeSourceHandles.build(), nodeRequirements));
assignedExchangeSourceHandles = ImmutableListMultimap.builder();
assignedExchangeDataSize = 0;
Expand All @@ -299,6 +300,7 @@ public List<TaskDescriptor> getMoreTasks()
}

if (assignedExchangeSourceHandleCount > 0) {
assignedExchangeSourceHandles.putAll(replicatedExchangeSourceHandles);
result.add(new TaskDescriptor(currentPartitionId, ImmutableListMultimap.of(), assignedExchangeSourceHandles.build(), nodeRequirements));
}

Expand Down Expand Up @@ -704,6 +706,21 @@ public void close()
}
}

private static IdentityHashMap<ExchangeSourceHandle, Exchange> getExchangeForHandleMap(
Map<PlanFragmentId, Exchange> sourceExchanges,
Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles)
{
IdentityHashMap<ExchangeSourceHandle, Exchange> exchangeForHandle = new IdentityHashMap<>();
for (Map.Entry<PlanFragmentId, ExchangeSourceHandle> entry : exchangeSourceHandles.entries()) {
PlanFragmentId fragmentId = entry.getKey();
ExchangeSourceHandle handle = entry.getValue();
Exchange exchange = sourceExchanges.get(fragmentId);
requireNonNull(exchange, "Exchange not found for fragment " + fragmentId);
exchangeForHandle.put(handle, exchange);
}
return exchangeForHandle;
}

private static ListMultimap<PlanNodeId, ExchangeSourceHandle> getReplicatedExchangeSourceHandles(PlanFragment fragment, Multimap<PlanFragmentId, ExchangeSourceHandle> handles)
{
return getInputsForRemoteSources(
Expand Down
Loading