jffwang
/
hetu-core
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

initial executor part

This commit is contained in:
nitin.kashyap 2021-02-25 13:57:10 +05:30 committed by Neerajunnikrishnan
parent cfd6fe41a0
commit 296b9fdc09
21 changed files with 533 additions and 65 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
presto-benchmark/src/main/java/io/prestosql/benchmark/AbstractOperatorBenchmark.java
View File

@ -300,7 +300,8 @@ public abstract class AbstractOperatorBenchmark
session,
false,
false,
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty());
CpuTimer cpuTimer = new CpuTimer();
Map<String, Long> executionStats = execute(taskContext);

3
presto-benchmark/src/test/java/io/prestosql/benchmark/MemoryLocalQueryRunner.java
View File

@ -92,7 +92,8 @@ public class MemoryLocalQueryRunner
localQueryRunner.getDefaultSession(),
false,
false,
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty());
// Use NullOutputFactory to avoid coping out results to avoid affecting benchmark results
ImmutableList.Builder<Page> output = ImmutableList.builder();

7
presto-main/src/main/java/io/prestosql/execution/MemoryTrackingRemoteTaskFactory.java
View File

@ -23,6 +23,7 @@ import io.prestosql.metadata.Split;
import io.prestosql.spi.plan.PlanNodeId;
import io.prestosql.sql.planner.PlanFragment;
import java.util.Optional;
import java.util.OptionalInt;
import static java.util.Objects.requireNonNull;
@ -48,7 +49,8 @@ public class MemoryTrackingRemoteTaskFactory
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo)
boolean summarizeTaskInfo,
Optional<PlanNodeId> parent)
{
RemoteTask task = remoteTaskFactory.createRemoteTask(session,
taskId,
@ -58,7 +60,8 @@ public class MemoryTrackingRemoteTaskFactory
totalPartitions,
outputBuffers,
partitionedSplitCountTracker,
summarizeTaskInfo);
summarizeTaskInfo,
parent);
task.addStateChangeListener(new UpdatePeakMemory(stateMachine));
return task;

17
presto-main/src/main/java/io/prestosql/execution/RemoteTaskFactory.java
View File

@ -22,17 +22,18 @@ import io.prestosql.metadata.Split;
import io.prestosql.spi.plan.PlanNodeId;
import io.prestosql.sql.planner.PlanFragment;
import java.util.Optional;
import java.util.OptionalInt;
public interface RemoteTaskFactory
{
RemoteTask createRemoteTask(Session session,
TaskId taskId,
InternalNode node,
PlanFragment fragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo);
TaskId taskId,
InternalNode node,
PlanFragment fragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo, Optional<PlanNodeId> parent);
}

15
presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java
View File

@ -119,6 +119,8 @@ public final class SqlStageExecution
@GuardedBy("SqlStageExecution.class")
public static Map<QueryId, List<Integer>> queryIdReuseTableScanMappingIdFinishedMap = new ConcurrentHashMap<>();
private PlanNodeId parentId;
public static SqlStageExecution createSqlStageExecution(
StageId stageId,
URI location,
@ -485,7 +487,8 @@ public final class SqlStageExecution
totalPartitions,
outputBuffers,
nodeTaskMap.createPartitionedSplitCountTracker(node, taskId),
summarizeTaskInfo);
summarizeTaskInfo,
Optional.ofNullable(parentId));
completeSources.forEach(task::noMoreSplits);
@ -705,4 +708,14 @@ public final class SqlStageExecution
}
}
}
public PlanNodeId getParentId()
{
return parentId;
}
public void setParentId(PlanNodeId parentId)
{
this.parentId = parentId;
}
}

3
presto-main/src/main/java/io/prestosql/execution/TaskManager.java
View File

@ -22,6 +22,7 @@ import io.prestosql.execution.buffer.OutputBuffers;
import io.prestosql.execution.buffer.OutputBuffers.OutputBufferId;
import io.prestosql.memory.MemoryPoolAssignmentsRequest;
import io.prestosql.sql.planner.PlanFragment;
import io.prestosql.sql.planner.plan.PlanNodeId;
import java.util.List;
import java.util.Optional;
@ -83,7 +84,7 @@ public interface TaskManager
* Updates the task plan, sources and output buffers. If the task does not
* already exist, is is created and then updated.
*/
TaskInfo updateTask(Session session, TaskId taskId, Optional<PlanFragment> fragment, List<TaskSource> sources, OutputBuffers outputBuffers, OptionalInt totalPartitions);
TaskInfo updateTask(Session session, TaskId taskId, Optional<PlanFragment> fragment, List<TaskSource> sources, OutputBuffers outputBuffers, OptionalInt totalPartitions, Optional<PlanNodeId> consumer);
/**
* Cancels a task. If the task does not already exist, is is created and then

13
presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java
View File

@ -52,6 +52,7 @@ import io.prestosql.sql.planner.NodePartitioningManager;
import io.prestosql.sql.planner.PartitioningHandle;
import io.prestosql.sql.planner.StageExecutionPlan;
import io.prestosql.sql.planner.plan.PlanFragmentId;
import io.prestosql.sql.planner.plan.RemoteSourceNode;
import java.net.URI;
import java.util.ArrayList;
@ -128,6 +129,8 @@ public class SqlQueryScheduler
private final HeuristicIndexerManager heuristicIndexerManager;
private final Session session;
private final Set<PlanFragmentId> visitedPlanFrags = new HashSet<>();
public static SqlQueryScheduler createSqlQueryScheduler(
QueryStateMachine queryStateMachine,
LocationFactory locationFactory,
@ -206,6 +209,7 @@ public class SqlQueryScheduler
Map<PartitioningHandle, NodePartitionMap> partitioningCache = new HashMap<>();
OutputBufferId rootBufferId = Iterables.getOnlyElement(rootOutputBuffers.getBuffers().keySet());
visitedPlanFrags.add(plan.getFragment().getId());
List<SqlStageExecution> stages = createStages(
(fragmentId, tasks, noMoreExchangeLocations) -> updateQueryOutputLocations(queryStateMachine, rootBufferId, tasks, noMoreExchangeLocations),
new AtomicInteger(),
@ -425,6 +429,11 @@ public class SqlQueryScheduler
ImmutableSet.Builder<SqlStageExecution> childStagesBuilder = ImmutableSet.builder();
for (StageExecutionPlan subStagePlan : plan.getSubStages()) {
if (visitedPlanFrags.contains(subStagePlan.getFragment().getId())) {
continue;
}
visitedPlanFrags.add(subStagePlan.getFragment().getId());
List<SqlStageExecution> subTree = createStages(
stage::addExchangeLocations,
nextStageId,
@ -446,6 +455,10 @@ public class SqlQueryScheduler
SqlStageExecution childStage = subTree.get(0);
childStagesBuilder.add(childStage);
Optional<RemoteSourceNode> parentNode = plan.getFragment().getRemoteSourceNodes().stream().filter(x -> x.getSourceFragmentIds().contains(childStage.getFragment().getId())).findAny();
checkArgument(parentNode.isPresent(), "Couldn't find parent of a CTE node");
childStage.setParentId(parentNode.get().getId());
}
Set<SqlStageExecution> childStages = childStagesBuilder.build();
stage.addStateChangeListener(newState -> {

17
presto-main/src/main/java/io/prestosql/memory/QueryContext.java
View File

@ -25,6 +25,7 @@ import io.prestosql.memory.context.MemoryTrackingContext;
import io.prestosql.operator.TaskContext;
import io.prestosql.spi.QueryId;
import io.prestosql.spiller.SpillSpaceTracker;
import io.prestosql.sql.planner.plan.PlanNodeId;
import javax.annotation.concurrent.GuardedBy;
import javax.annotation.concurrent.ThreadSafe;
@ -33,6 +34,7 @@ import java.util.Comparator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
@ -246,7 +248,7 @@ public class QueryContext
return memoryPool;
}
public TaskContext addTaskContext(TaskStateMachine taskStateMachine, Session session, boolean perOperatorCpuTimerEnabled, boolean cpuTimerEnabled, OptionalInt totalPartitions)
public TaskContext addTaskContext(TaskStateMachine taskStateMachine, Session session, boolean perOperatorCpuTimerEnabled, boolean cpuTimerEnabled, OptionalInt totalPartitions, Optional<PlanNodeId> parent)
{
TaskContext taskContext = TaskContext.createTaskContext(
this,
@ -258,7 +260,8 @@ public class QueryContext
queryMemoryContext.newMemoryTrackingContext(),
perOperatorCpuTimerEnabled,
cpuTimerEnabled,
totalPartitions);
totalPartitions,
parent.orElse(null));
taskContexts.put(taskStateMachine.getTaskId(), taskContext);
return taskContext;
}
@ -353,4 +356,14 @@ public class QueryContext
return format("%s, Top Consumers: %s", additionalInfo, topConsumers);
}
public QueryId getQueryId()
{
return queryId;
}
public int getTaskCount()
{
return taskContexts.size();
}
}

259
presto-main/src/main/java/io/prestosql/operator/CommonTableExpressionOperator.java Normal file
View File

@ -0,0 +1,259 @@
/*
* Copyright (C) 2018-2021. Huawei Technologies Co., Ltd. All rights reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.prestosql.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ListenableFuture;
import io.airlift.log.Logger;
import io.airlift.units.DataSize;
import io.prestosql.spi.Page;
import io.prestosql.spi.type.Type;
import io.prestosql.sql.planner.plan.PlanNodeId;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static java.util.Objects.requireNonNull;
public class CommonTableExpressionOperator
implements Operator, Closeable
{
private static final Logger LOG = Logger.get(CommonTableExpressionOperator.class);
private final PlanNodeId self;
private final OperatorContext operatorContext;
private final PlanNodeId consumer;
private final CommonTableExecutionContext cteContext;
private final int operatorInstaceId;
private boolean finish;
private boolean isProducer;
public CommonTableExpressionOperator(
PlanNodeId self,
PlanNodeId consumer,
OperatorContext operatorContext,
CommonTableExecutionContext cteContext,
int operatorInstaceId)
{
this.self = requireNonNull(self, "PlanNode Id is null");
this.consumer = requireNonNull(consumer, "consumer cannot be null");
this.operatorContext = requireNonNull(operatorContext, "operatorContext is null");
this.cteContext = requireNonNull(cteContext, "CTE context is null");
this.operatorInstaceId = operatorInstaceId;
synchronized (cteContext) {
if (cteContext.isProducer(consumer)) {
this.isProducer = true;
cteContext.setProducerState(consumer, operatorInstaceId, true);
}
}
LOG.debug("CTE(" + cteContext.getName() + ")[" + consumer + "-" + operatorInstaceId + "] Operator Initialized (Producer: " + this.isProducer + ")");
}
public static class CommonTableExpressionOperatorFactory
implements OperatorFactory
{
private final int operatorId;
private final PlanNodeId planNodeId;
private final List<Type> types;
private final DataSize minOutputPageSize;
private final int minOutputPageRowCount;
private boolean closed;
private Set<PlanNodeId> parents = new HashSet<>();
private CommonTableExecutionContext cteCtx;
private final AtomicInteger operatorCounter = new AtomicInteger(0);
public CommonTableExpressionOperatorFactory(
int operatorId,
PlanNodeId planNodeId,
CommonTableExecutionContext cteCtx,
List<Type> types,
DataSize minOutputPageSize,
int minOutputPageRowCount)
{
this.operatorId = operatorId;
this.planNodeId = requireNonNull(planNodeId, "planNodeId is null");
this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
this.minOutputPageSize = requireNonNull(minOutputPageSize, "minOutputPageSize is null");
this.minOutputPageRowCount = minOutputPageRowCount;
this.cteCtx = cteCtx;
}
@Override
public Operator createOperator(DriverContext driverContext)
{
checkState(!closed, "Factory is already closed");
checkArgument(parents.size() > 0, "No parent assigned for CTE");
OperatorContext operatorContext = driverContext.addOperatorContext(operatorId, planNodeId, CommonTableExpressionOperator.class.getSimpleName());
return new CommonTableExpressionOperator(
planNodeId,
parents.stream().findAny().get(),
operatorContext,
cteCtx,
operatorCounter.incrementAndGet());
}
@Override
public void noMoreOperators()
{
closed = true;
}
@Override
public OperatorFactory duplicate()
{
return new CommonTableExpressionOperatorFactory(operatorId, planNodeId, cteCtx, types, minOutputPageSize, minOutputPageRowCount);
}
public void addConsumer(PlanNodeId id)
{
parents.add(id);
}
}
@Override
public OperatorContext getOperatorContext()
{
return operatorContext;
}
/**
* Returns a future that will be completed when the operator becomes
* unblocked. If the operator is not blocked, this method should return
* {@code NOT_BLOCKED}.
*/
@Override
public ListenableFuture<?> isBlocked()
{
return cteContext.isBlocked(consumer);
}
/**
* Returns true if and only if this operator can accept an input page.
*/
@Override
public boolean needsInput()
{
return isProducer && !finish;
}
/**
* Adds an input page to the operator. This method will only be called if
* {@code needsInput()} returns true.
*
* @param page
*/
@Override
public void addInput(Page page)
{
/* Got a new page... Place it in the Queue! */
cteContext.addPage(page);
LOG.debug("CTE(" + cteContext.getName() + ")" + "[" + consumer + "-" + operatorInstaceId + "] Page added with " + page.getPositionCount() + " rows");
}
/**
* Gets an output page from the operator. If no output data is currently
* available, return null.
*/
@Override
public Page getOutput()
{
try {
Page page = cteContext.getPage(consumer);
if (page != null) {
LOG.debug("CTE(" + cteContext.getName() + ")" + "[" + consumer + "-" + operatorInstaceId + "] got a page with " + page.getPositionCount() + " rows");
}
return page;
}
catch (CommonTableExecutionContext.CTEDoneException e) {
if (!finish) {
finish = true;
LOG.debug("CTE(" + cteContext.getName() + ")" + "[" + consumer + "-" + operatorInstaceId + "] Done(empty) directed");
}
}
return null;
}
/**
* After calling this method operator should revoke all reserved revocable memory.
* As soon as memory is revoked returned future should be marked as done.
* <p>
* Spawned threads can not modify OperatorContext because it's not thread safe.
* For this purpose implement {@link #finishMemoryRevoke()}
* <p>
* Since memory revoking signal is delivered asynchronously to the Operator, implementation
* must gracefully handle the case when there no longer is any revocable memory allocated.
* <p>
* After this method is called on Operator the Driver is disallowed to call any
* processing methods on it (isBlocked/needsInput/addInput/getOutput) until
* {@link #finishMemoryRevoke()} is called.
*/
@Override
public ListenableFuture<?> startMemoryRevoke()
{
return NOT_BLOCKED;
}
/**
* Clean up and release resources after completed memory revoking. Called by driver
* once future returned by startMemoryRevoke is completed.
*/
@Override
public void finishMemoryRevoke()
{
}
/**
* Notifies the operator that no more pages will be added and the
* operator should finish processing and flush results. This method
* will not be called if the Task is already failed or canceled.
*/
@Override
public void finish()
{
if (isProducer) {
cteContext.setProducerState(consumer, operatorInstaceId, false);
}
LOG.debug("CTE(" + cteContext.getName() + ")[" + consumer + "-" + operatorInstaceId + "] Operator Finished (deferred)");
}
/**
* Is this operator completely finished processing and no more
* output pages will be produced.
*/
@Override
public boolean isFinished()
{
return finish;
}
/**
* This method will always be called before releasing the Operator reference.
*/
@Override
public void close() throws IOException
{
LOG.debug("CTE(" + cteContext.getName() + ")[" + consumer + "-" + operatorInstaceId + "] Operator Closed");
}
}

20
presto-main/src/main/java/io/prestosql/server/HttpRemoteTaskFactory.java
View File

@ -45,6 +45,7 @@ import org.weakref.jmx.Nested;
import javax.annotation.PreDestroy;
import javax.inject.Inject;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
@ -134,14 +135,14 @@ public class HttpRemoteTaskFactory
@Override
public RemoteTask createRemoteTask(Session session,
TaskId taskId,
InternalNode node,
PlanFragment fragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo)
TaskId taskId,
InternalNode node,
PlanFragment fragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo, Optional<PlanNodeId> parent)
{
return new HttpRemoteTask(session,
taskId,
@ -164,6 +165,7 @@ public class HttpRemoteTaskFactory
taskUpdateRequestCodec,
partitionedSplitCountTracker,
stats,
isBinaryEncoding);
isBinaryEncoding,
parent);
}
}

3
presto-main/src/main/java/io/prestosql/server/TaskResource.java
View File

@ -137,7 +137,8 @@ public class TaskResource
taskUpdateRequest.getFragment(),
taskUpdateRequest.getSources(),
taskUpdateRequest.getOutputIds(),
taskUpdateRequest.getTotalPartitions());
taskUpdateRequest.getTotalPartitions(),
taskUpdateRequest.getConsumerId());
if (shouldSummarize(uriInfo)) {
taskInfo = taskInfo.summarize();

12
presto-main/src/main/java/io/prestosql/server/TaskUpdateRequest.java
View File

@ -20,6 +20,7 @@ import io.prestosql.SessionRepresentation;
import io.prestosql.execution.TaskSource;
import io.prestosql.execution.buffer.OutputBuffers;
import io.prestosql.sql.planner.PlanFragment;
import io.prestosql.sql.planner.plan.PlanNodeId;
import java.util.List;
import java.util.Map;
@ -38,6 +39,7 @@ public class TaskUpdateRequest
private final List<TaskSource> sources;
private final OutputBuffers outputIds;
private final OptionalInt totalPartitions;
private final Optional<PlanNodeId> consumerId;
@JsonCreator
public TaskUpdateRequest(
@ -46,7 +48,8 @@ public class TaskUpdateRequest
@JsonProperty("fragment") Optional<PlanFragment> fragment,
@JsonProperty("sources") List<TaskSource> sources,
@JsonProperty("outputIds") OutputBuffers outputIds,
@JsonProperty("totalPartitions") OptionalInt totalPartitions)
@JsonProperty("totalPartitions") OptionalInt totalPartitions,
@JsonProperty("consumerId")Optional<PlanNodeId> consumerPlanNodeId)
{
requireNonNull(session, "session is null");
requireNonNull(extraCredentials, "credentials is null");
@ -61,6 +64,7 @@ public class TaskUpdateRequest
this.sources = ImmutableList.copyOf(sources);
this.outputIds = outputIds;
this.totalPartitions = totalPartitions;
this.consumerId = consumerPlanNodeId;
}
@JsonProperty
@ -99,6 +103,12 @@ public class TaskUpdateRequest
return totalPartitions;
}
@JsonProperty
public Optional<PlanNodeId> getConsumerId()
{
return consumerId;
}
@Override
public String toString()
{

47
presto-main/src/main/java/io/prestosql/server/remotetask/HttpRemoteTask.java
View File

@ -159,28 +159,30 @@ public final class HttpRemoteTask
private final AtomicBoolean aborting = new AtomicBoolean(false);
private final boolean isBinaryEncoding;
private Optional<PlanNodeId> parent;
public HttpRemoteTask(Session session,
TaskId taskId,
String nodeId,
URI location,
PlanFragment planFragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
HttpClient httpClient,
Executor executor,
ScheduledExecutorService updateScheduledExecutor,
ScheduledExecutorService errorScheduledExecutor,
Duration maxErrorDuration,
Duration taskStatusRefreshMaxWait,
Duration taskInfoUpdateInterval,
boolean summarizeTaskInfo,
Codec<TaskStatus> taskStatusCodec,
Codec<TaskInfo> taskInfoCodec,
Codec<TaskUpdateRequest> taskUpdateRequestCodec,
PartitionedSplitCountTracker partitionedSplitCountTracker,
RemoteTaskStats stats, boolean isBinaryEncoding)
TaskId taskId,
String nodeId,
URI location,
PlanFragment planFragment,
Multimap<PlanNodeId, Split> initialSplits,
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
HttpClient httpClient,
Executor executor,
ScheduledExecutorService updateScheduledExecutor,
ScheduledExecutorService errorScheduledExecutor,
Duration maxErrorDuration,
Duration taskStatusRefreshMaxWait,
Duration taskInfoUpdateInterval,
boolean summarizeTaskInfo,
Codec<TaskStatus> taskStatusCodec,
Codec<TaskInfo> taskInfoCodec,
Codec<TaskUpdateRequest> taskUpdateRequestCodec,
PartitionedSplitCountTracker partitionedSplitCountTracker,
RemoteTaskStats stats, boolean isBinaryEncoding,
Optional<PlanNodeId> parent)
{
requireNonNull(session, "session is null");
requireNonNull(taskId, "taskId is null");
@ -196,6 +198,7 @@ public final class HttpRemoteTask
requireNonNull(taskUpdateRequestCodec, "taskUpdateRequestCodec is null");
requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null");
requireNonNull(stats, "stats is null");
requireNonNull(parent, "parent is null");
try (SetThreadName ignored = new SetThreadName("HttpRemoteTask-%s", taskId)) {
this.taskId = taskId;
@ -214,6 +217,7 @@ public final class HttpRemoteTask
this.partitionedSplitCountTracker = requireNonNull(partitionedSplitCountTracker, "partitionedSplitCountTracker is null");
this.stats = stats;
this.isBinaryEncoding = isBinaryEncoding;
this.parent = parent;
for (Entry<PlanNodeId, Split> entry : requireNonNull(initialSplits, "initialSplits is null").entries()) {
ScheduledSplit scheduledSplit = new ScheduledSplit(nextSplitId.getAndIncrement(), entry.getKey(), entry.getValue());
@ -515,7 +519,8 @@ public final class HttpRemoteTask
fragment,
sources,
outputBuffers.get(),
totalPartitions);
totalPartitions,
parent);
byte[] taskUpdateRequestJson = taskUpdateRequestCodec.toBytes(updateRequest);
if (fragment.isPresent()) {
stats.updateWithPlanBytes(taskUpdateRequestJson.length);

4
presto-main/src/main/java/io/prestosql/testing/TestingTaskContext.java
View File

@ -26,6 +26,7 @@ import io.prestosql.spi.QueryId;
import io.prestosql.spi.memory.MemoryPoolId;
import io.prestosql.spiller.SpillSpaceTracker;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
@ -71,7 +72,8 @@ public final class TestingTaskContext
session,
true,
true,
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty());
}
public static Builder builder(Executor notificationExecutor, ScheduledExecutorService yieldExecutor, Session session)

8
presto-main/src/test/java/io/prestosql/execution/MockRemoteTaskFactory.java
View File

@ -122,13 +122,15 @@ public class MockRemoteTaskFactory
new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)),
ungroupedExecution(),
StatsAndCosts.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty());
ImmutableMultimap.Builder<PlanNodeId, Split> initialSplits = ImmutableMultimap.builder();
for (Split sourceSplit : splits) {
initialSplits.put(sourceId, sourceSplit);
}
return createRemoteTask(TEST_SESSION, taskId, newNode, testFragment, initialSplits.build(), OptionalInt.empty(), createInitialEmptyOutputBuffers(BROADCAST), partitionedSplitCountTracker, true);
return createRemoteTask(TEST_SESSION, taskId, newNode, testFragment, initialSplits.build(), OptionalInt.empty(), createInitialEmptyOutputBuffers(BROADCAST), partitionedSplitCountTracker, true, Optional.empty());
}
@Override
@ -141,7 +143,7 @@ public class MockRemoteTaskFactory
OptionalInt totalPartitions,
OutputBuffers outputBuffers,
PartitionedSplitCountTracker partitionedSplitCountTracker,
boolean summarizeTaskInfo)
boolean summarizeTaskInfo, Optional<PlanNodeId> parent)
{
return new MockRemoteTask(taskId, fragment, node.getNodeIdentifier(), executor, scheduledExecutor, initialSplits, totalPartitions, partitionedSplitCountTracker);
}
@ -195,7 +197,7 @@ public class MockRemoteTaskFactory
scheduledExecutor,
new DataSize(1, MEGABYTE),
spillSpaceTracker);
this.taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, true, true, totalPartitions);
this.taskContext = queryContext.addTaskContext(taskStateMachine, TEST_SESSION, true, true, totalPartitions, Optional.empty());
this.location = URI.create("fake://task/" + taskId);

4
presto-main/src/test/java/io/prestosql/execution/TaskTestUtils.java
View File

@ -112,6 +112,8 @@ public final class TaskTestUtils
.withBucketToPartition(Optional.of(new int[1])),
ungroupedExecution(),
StatsAndCosts.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty());
public static LocalExecutionPlanner createTestingPlanner()
@ -174,7 +176,7 @@ public final class TaskTestUtils
public static TaskInfo updateTask(SqlTask sqlTask, List<TaskSource> taskSources, OutputBuffers outputBuffers)
{
return sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), taskSources, outputBuffers, OptionalInt.empty());
return sqlTask.updateTask(TEST_SESSION, Optional.of(PLAN_FRAGMENT), taskSources, outputBuffers, OptionalInt.empty(), Optional.empty(), null);
}
public static SplitMonitor createTestSplitMonitor()

7
presto-main/src/test/java/io/prestosql/execution/TestMemoryRevokingScheduler.java
View File

@ -44,6 +44,7 @@ import org.testng.annotations.Test;
import java.net.URI;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.Set;
import java.util.concurrent.Callable;
@ -118,7 +119,7 @@ public class TestMemoryRevokingScheduler
SqlTask sqlTask1 = newSqlTask();
SqlTask sqlTask2 = newSqlTask();
TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty());
TaskContext taskContext1 = sqlTask1.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q1", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty());
PipelineContext pipelineContext11 = taskContext1.addPipelineContext(0, false, false, false);
DriverContext driverContext111 = pipelineContext11.addDriverContext();
OperatorContext operatorContext1 = driverContext111.addOperatorContext(1, new PlanNodeId("na"), "na");
@ -126,7 +127,7 @@ public class TestMemoryRevokingScheduler
DriverContext driverContext112 = pipelineContext11.addDriverContext();
OperatorContext operatorContext3 = driverContext112.addOperatorContext(3, new PlanNodeId("na"), "na");
TaskContext taskContext2 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty());
TaskContext taskContext2 = sqlTask2.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q2", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty());
PipelineContext pipelineContext21 = taskContext2.addPipelineContext(1, false, false, false);
DriverContext driverContext211 = pipelineContext21.addDriverContext();
OperatorContext operatorContext4 = driverContext211.addOperatorContext(4, new PlanNodeId("na"), "na");
@ -248,7 +249,7 @@ public class TestMemoryRevokingScheduler
private OperatorContext createContexts(SqlTask sqlTask)
{
TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty());
TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty(), Optional.empty());
PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false);
DriverContext driverContext = pipelineContext.addDriverContext();
OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na");

2
presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java
View File

@ -186,6 +186,8 @@ public class TestSqlStageExecution
new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), planNode.getOutputSymbols()),
ungroupedExecution(),
StatsAndCosts.empty(),
Optional.empty(),
Optional.empty(),
Optional.empty());
}
}

23
presto-main/src/test/java/io/prestosql/execution/TestSqlTask.java
View File

@ -113,7 +113,9 @@ public class TestSqlTask
ImmutableList.of(),
createInitialEmptyOutputBuffers(PARTITIONED)
.withNoMoreBufferIds(),
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty(),
null);
assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
taskInfo = sqlTask.getTaskInfo();
@ -124,7 +126,9 @@ public class TestSqlTask
ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(), true)),
createInitialEmptyOutputBuffers(PARTITIONED)
.withNoMoreBufferIds(),
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty(),
null);
assertEquals(taskInfo.getTaskStatus().getState(), TaskState.FINISHED);
taskInfo = sqlTask.getTaskInfo();
@ -141,7 +145,9 @@ public class TestSqlTask
Optional.of(PLAN_FRAGMENT),
ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)),
createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(),
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty(),
null);
assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
taskInfo = sqlTask.getTaskInfo();
@ -179,7 +185,9 @@ public class TestSqlTask
createInitialEmptyOutputBuffers(PARTITIONED)
.withBuffer(OUT, 0)
.withNoMoreBufferIds(),
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty(),
null);
assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
assertNull(taskInfo.getStats().getEndTime());
@ -206,7 +214,9 @@ public class TestSqlTask
Optional.of(PLAN_FRAGMENT),
ImmutableList.of(new TaskSource(TABLE_SCAN_NODE_ID, ImmutableSet.of(SPLIT), true)),
createInitialEmptyOutputBuffers(PARTITIONED).withBuffer(OUT, 0).withNoMoreBufferIds(),
OptionalInt.empty());
OptionalInt.empty(),
Optional.empty(),
null);
assertEquals(taskInfo.getTaskStatus().getState(), TaskState.RUNNING);
taskInfo = sqlTask.getTaskInfo();
@ -311,7 +321,8 @@ public class TestSqlTask
new DataSize(1, MEGABYTE),
new SpillSpaceTracker(new DataSize(1, GIGABYTE)));
queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty());
queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty(),
Optional.empty());
return createSqlTask(
taskId,

8
presto-main/src/test/java/io/prestosql/execution/TestSqlTaskExecution.java
View File

@ -162,7 +162,8 @@ public class TestSqlTaskExecution
OptionalInt.empty(),
executionStrategy)),
ImmutableList.of(TABLE_SCAN_NODE_ID),
executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(TABLE_SCAN_NODE_ID)) : StageExecutionDescriptor.ungroupedExecution());
executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(TABLE_SCAN_NODE_ID)) : StageExecutionDescriptor.ungroupedExecution(),
Optional.empty());
TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine);
SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(
taskStateMachine,
@ -413,7 +414,8 @@ public class TestSqlTaskExecution
OptionalInt.empty(),
UNGROUPED_EXECUTION)),
ImmutableList.of(scan2NodeId, scan0NodeId),
executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(scan0NodeId, scan2NodeId)) : StageExecutionDescriptor.ungroupedExecution());
executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(scan0NodeId, scan2NodeId)) : StageExecutionDescriptor.ungroupedExecution(),
Optional.empty());
TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine);
SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(
taskStateMachine,
@ -601,7 +603,7 @@ public class TestSqlTaskExecution
driverYieldExecutor,
new DataSize(1, MEGABYTE),
new SpillSpaceTracker(new DataSize(1, GIGABYTE)));
return queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty());
return queryContext.addTaskContext(taskStateMachine, TEST_SESSION, false, false, OptionalInt.empty(), Optional.empty());
}
private PartitionedOutputBuffer newTestingOutputBuffer(ScheduledExecutorService taskNotificationExecutor)

123
presto-main/src/test/java/io/prestosql/operator/TestCommonTableExpressionOperator.java Normal file
View File

@ -0,0 +1,123 @@
/*
* Copyright (C) 2018-2020. Huawei Technologies Co., Ltd. All rights reserved.
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.prestosql.operator;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import io.airlift.units.DataSize;
import io.prestosql.SequencePageBuilder;
import io.prestosql.metadata.Metadata;
import io.prestosql.operator.scalar.AbstractTestFunctions;
import io.prestosql.spi.Page;
import io.prestosql.sql.gen.ExpressionCompiler;
import io.prestosql.sql.gen.PageFunctionCompiler;
import io.prestosql.sql.planner.plan.PlanNodeId;
import io.prestosql.testing.MaterializedResult;
import org.testng.annotations.Test;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.prestosql.SessionTestUtils.TEST_SESSION;
import static io.prestosql.metadata.MetadataManager.createTestMetadataManager;
import static io.prestosql.operator.OperatorAssertion.assertOperatorEquals;
import static io.prestosql.spi.type.VarcharType.VARCHAR;
import static io.prestosql.testing.TestingTaskContext.createTaskContext;
import static java.util.concurrent.Executors.newCachedThreadPool;
import static java.util.concurrent.Executors.newScheduledThreadPool;
import static org.testng.Assert.assertTrue;
public class TestCommonTableExpressionOperator
extends AbstractTestFunctions
{
private final Metadata metadata = createTestMetadataManager();
private final ExpressionCompiler expressionCompiler = new ExpressionCompiler(metadata, new PageFunctionCompiler(metadata, 0));
private ExecutorService executor;
private ScheduledExecutorService scheduledExecutor;
public TestCommonTableExpressionOperator()
{
executor = newCachedThreadPool(daemonThreadsNamed("test-executor-%s"));
scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s"));
}
@Test
public void testOperatorSource()
{
final Page input = SequencePageBuilder.createSequencePage(ImmutableList.of(VARCHAR), 10_000, 0);
DriverContext driverContext = newDriverContext();
CommonTableExecutionContext cteContext = new CommonTableExecutionContext("test_cte_prod_1",
ImmutableSet.of(new PlanNodeId("consumer_1"), new PlanNodeId("consumer_2")), new PlanNodeId("consumer_1"),
driverContext.getNotificationExecutor(), 0, 1024, 512);
CommonTableExpressionOperator.CommonTableExpressionOperatorFactory parent1 = new CommonTableExpressionOperator.CommonTableExpressionOperatorFactory(
0,
new PlanNodeId("test"),
cteContext,
ImmutableList.of(VARCHAR),
new DataSize(0, DataSize.Unit.BYTE),
0);
parent1.addConsumer(new PlanNodeId("consumer_1"));
CommonTableExpressionOperator.CommonTableExpressionOperatorFactory parent2 = new CommonTableExpressionOperator.CommonTableExpressionOperatorFactory(
1,
new PlanNodeId("test"),
cteContext,
ImmutableList.of(VARCHAR),
new DataSize(0, DataSize.Unit.BYTE),
0);
parent2.addConsumer(new PlanNodeId("consumer_2"));
//Operator operator = factory.createOperator(driverContext);
MaterializedResult result = MaterializedResult.resultBuilder(driverContext.getSession(), VARCHAR)
.page(input)
.build();
assertOperatorEquals(parent1, driverContext, ImmutableList.of(input), result);
assertOperatorEquals(parent2, driverContext, ImmutableList.of(input), result);
}
private static List<Page> toPages(Operator operator)
{
ImmutableList.Builder<Page> outputPages = ImmutableList.builder();
// read output until input is needed or operator is finished
int nullPages = 0;
while (!operator.isFinished()) {
Page outputPage = operator.getOutput();
if (outputPage == null) {
// break infinite loop due to null pages
assertTrue(nullPages < 1_000_000, "Too many null pages; infinite loop?");
nullPages++;
}
else {
outputPages.add(outputPage);
nullPages = 0;
}
}
return outputPages.build();
}
private DriverContext newDriverContext()
{
return createTaskContext(executor, scheduledExecutor, TEST_SESSION)
.addPipelineContext(0, true, true, false)
.addDriverContext();
}
}