diff --git a/presto-benchmark/src/main/java/io/prestosql/benchmark/AbstractOperatorBenchmark.java b/presto-benchmark/src/main/java/io/prestosql/benchmark/AbstractOperatorBenchmark.java index adda4cb51..3b81b2810 100644 --- a/presto-benchmark/src/main/java/io/prestosql/benchmark/AbstractOperatorBenchmark.java +++ b/presto-benchmark/src/main/java/io/prestosql/benchmark/AbstractOperatorBenchmark.java @@ -300,7 +300,8 @@ public abstract class AbstractOperatorBenchmark session, false, false, - OptionalInt.empty()); + OptionalInt.empty(), + Optional.empty()); CpuTimer cpuTimer = new CpuTimer(); Map executionStats = execute(taskContext); diff --git a/presto-benchmark/src/test/java/io/prestosql/benchmark/MemoryLocalQueryRunner.java b/presto-benchmark/src/test/java/io/prestosql/benchmark/MemoryLocalQueryRunner.java index 8c9f1b65b..29bd68a51 100644 --- a/presto-benchmark/src/test/java/io/prestosql/benchmark/MemoryLocalQueryRunner.java +++ b/presto-benchmark/src/test/java/io/prestosql/benchmark/MemoryLocalQueryRunner.java @@ -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 output = ImmutableList.builder(); diff --git a/presto-main/src/main/java/io/prestosql/execution/MemoryTrackingRemoteTaskFactory.java b/presto-main/src/main/java/io/prestosql/execution/MemoryTrackingRemoteTaskFactory.java index d54675542..55af66d48 100644 --- a/presto-main/src/main/java/io/prestosql/execution/MemoryTrackingRemoteTaskFactory.java +++ b/presto-main/src/main/java/io/prestosql/execution/MemoryTrackingRemoteTaskFactory.java @@ -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 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; diff --git a/presto-main/src/main/java/io/prestosql/execution/RemoteTaskFactory.java b/presto-main/src/main/java/io/prestosql/execution/RemoteTaskFactory.java index 75cf51172..100dce4f9 100644 --- a/presto-main/src/main/java/io/prestosql/execution/RemoteTaskFactory.java +++ b/presto-main/src/main/java/io/prestosql/execution/RemoteTaskFactory.java @@ -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 initialSplits, - OptionalInt totalPartitions, - OutputBuffers outputBuffers, - PartitionedSplitCountTracker partitionedSplitCountTracker, - boolean summarizeTaskInfo); + TaskId taskId, + InternalNode node, + PlanFragment fragment, + Multimap initialSplits, + OptionalInt totalPartitions, + OutputBuffers outputBuffers, + PartitionedSplitCountTracker partitionedSplitCountTracker, + boolean summarizeTaskInfo, Optional parent); } diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java b/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java index 852083426..0e7119cad 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java @@ -119,6 +119,8 @@ public final class SqlStageExecution @GuardedBy("SqlStageExecution.class") public static Map> 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; + } } diff --git a/presto-main/src/main/java/io/prestosql/execution/TaskManager.java b/presto-main/src/main/java/io/prestosql/execution/TaskManager.java index 61b8920fb..3745573ee 100644 --- a/presto-main/src/main/java/io/prestosql/execution/TaskManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/TaskManager.java @@ -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 fragment, List sources, OutputBuffers outputBuffers, OptionalInt totalPartitions); + TaskInfo updateTask(Session session, TaskId taskId, Optional fragment, List sources, OutputBuffers outputBuffers, OptionalInt totalPartitions, Optional consumer); /** * Cancels a task. If the task does not already exist, is is created and then diff --git a/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java b/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java index 60cc9e5c6..51b641a51 100644 --- a/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java +++ b/presto-main/src/main/java/io/prestosql/execution/scheduler/SqlQueryScheduler.java @@ -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 visitedPlanFrags = new HashSet<>(); + public static SqlQueryScheduler createSqlQueryScheduler( QueryStateMachine queryStateMachine, LocationFactory locationFactory, @@ -206,6 +209,7 @@ public class SqlQueryScheduler Map partitioningCache = new HashMap<>(); OutputBufferId rootBufferId = Iterables.getOnlyElement(rootOutputBuffers.getBuffers().keySet()); + visitedPlanFrags.add(plan.getFragment().getId()); List stages = createStages( (fragmentId, tasks, noMoreExchangeLocations) -> updateQueryOutputLocations(queryStateMachine, rootBufferId, tasks, noMoreExchangeLocations), new AtomicInteger(), @@ -425,6 +429,11 @@ public class SqlQueryScheduler ImmutableSet.Builder childStagesBuilder = ImmutableSet.builder(); for (StageExecutionPlan subStagePlan : plan.getSubStages()) { + if (visitedPlanFrags.contains(subStagePlan.getFragment().getId())) { + continue; + } + + visitedPlanFrags.add(subStagePlan.getFragment().getId()); List subTree = createStages( stage::addExchangeLocations, nextStageId, @@ -446,6 +455,10 @@ public class SqlQueryScheduler SqlStageExecution childStage = subTree.get(0); childStagesBuilder.add(childStage); + Optional 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 childStages = childStagesBuilder.build(); stage.addStateChangeListener(newState -> { diff --git a/presto-main/src/main/java/io/prestosql/memory/QueryContext.java b/presto-main/src/main/java/io/prestosql/memory/QueryContext.java index 242ef2539..7d7b1fa36 100644 --- a/presto-main/src/main/java/io/prestosql/memory/QueryContext.java +++ b/presto-main/src/main/java/io/prestosql/memory/QueryContext.java @@ -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 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(); + } } diff --git a/presto-main/src/main/java/io/prestosql/operator/CommonTableExpressionOperator.java b/presto-main/src/main/java/io/prestosql/operator/CommonTableExpressionOperator.java new file mode 100644 index 000000000..7d9ac8190 --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/operator/CommonTableExpressionOperator.java @@ -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 types; + private final DataSize minOutputPageSize; + private final int minOutputPageRowCount; + private boolean closed; + private Set parents = new HashSet<>(); + private CommonTableExecutionContext cteCtx; + private final AtomicInteger operatorCounter = new AtomicInteger(0); + + public CommonTableExpressionOperatorFactory( + int operatorId, + PlanNodeId planNodeId, + CommonTableExecutionContext cteCtx, + List 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. + *

+ * Spawned threads can not modify OperatorContext because it's not thread safe. + * For this purpose implement {@link #finishMemoryRevoke()} + *

+ * 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. + *

+ * 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"); + } +} diff --git a/presto-main/src/main/java/io/prestosql/server/HttpRemoteTaskFactory.java b/presto-main/src/main/java/io/prestosql/server/HttpRemoteTaskFactory.java index 825a836e7..35932add5 100644 --- a/presto-main/src/main/java/io/prestosql/server/HttpRemoteTaskFactory.java +++ b/presto-main/src/main/java/io/prestosql/server/HttpRemoteTaskFactory.java @@ -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 initialSplits, - OptionalInt totalPartitions, - OutputBuffers outputBuffers, - PartitionedSplitCountTracker partitionedSplitCountTracker, - boolean summarizeTaskInfo) + TaskId taskId, + InternalNode node, + PlanFragment fragment, + Multimap initialSplits, + OptionalInt totalPartitions, + OutputBuffers outputBuffers, + PartitionedSplitCountTracker partitionedSplitCountTracker, + boolean summarizeTaskInfo, Optional parent) { return new HttpRemoteTask(session, taskId, @@ -164,6 +165,7 @@ public class HttpRemoteTaskFactory taskUpdateRequestCodec, partitionedSplitCountTracker, stats, - isBinaryEncoding); + isBinaryEncoding, + parent); } } diff --git a/presto-main/src/main/java/io/prestosql/server/TaskResource.java b/presto-main/src/main/java/io/prestosql/server/TaskResource.java index a5abeca18..62095ddf2 100644 --- a/presto-main/src/main/java/io/prestosql/server/TaskResource.java +++ b/presto-main/src/main/java/io/prestosql/server/TaskResource.java @@ -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(); diff --git a/presto-main/src/main/java/io/prestosql/server/TaskUpdateRequest.java b/presto-main/src/main/java/io/prestosql/server/TaskUpdateRequest.java index f81157e9d..8e3c9542e 100644 --- a/presto-main/src/main/java/io/prestosql/server/TaskUpdateRequest.java +++ b/presto-main/src/main/java/io/prestosql/server/TaskUpdateRequest.java @@ -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 sources; private final OutputBuffers outputIds; private final OptionalInt totalPartitions; + private final Optional consumerId; @JsonCreator public TaskUpdateRequest( @@ -46,7 +48,8 @@ public class TaskUpdateRequest @JsonProperty("fragment") Optional fragment, @JsonProperty("sources") List sources, @JsonProperty("outputIds") OutputBuffers outputIds, - @JsonProperty("totalPartitions") OptionalInt totalPartitions) + @JsonProperty("totalPartitions") OptionalInt totalPartitions, + @JsonProperty("consumerId")Optional 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 getConsumerId() + { + return consumerId; + } + @Override public String toString() { diff --git a/presto-main/src/main/java/io/prestosql/server/remotetask/HttpRemoteTask.java b/presto-main/src/main/java/io/prestosql/server/remotetask/HttpRemoteTask.java index 32b9ef326..bfe8d8731 100644 --- a/presto-main/src/main/java/io/prestosql/server/remotetask/HttpRemoteTask.java +++ b/presto-main/src/main/java/io/prestosql/server/remotetask/HttpRemoteTask.java @@ -159,28 +159,30 @@ public final class HttpRemoteTask private final AtomicBoolean aborting = new AtomicBoolean(false); private final boolean isBinaryEncoding; + private Optional parent; public HttpRemoteTask(Session session, - TaskId taskId, - String nodeId, - URI location, - PlanFragment planFragment, - Multimap initialSplits, - OptionalInt totalPartitions, - OutputBuffers outputBuffers, - HttpClient httpClient, - Executor executor, - ScheduledExecutorService updateScheduledExecutor, - ScheduledExecutorService errorScheduledExecutor, - Duration maxErrorDuration, - Duration taskStatusRefreshMaxWait, - Duration taskInfoUpdateInterval, - boolean summarizeTaskInfo, - Codec taskStatusCodec, - Codec taskInfoCodec, - Codec taskUpdateRequestCodec, - PartitionedSplitCountTracker partitionedSplitCountTracker, - RemoteTaskStats stats, boolean isBinaryEncoding) + TaskId taskId, + String nodeId, + URI location, + PlanFragment planFragment, + Multimap initialSplits, + OptionalInt totalPartitions, + OutputBuffers outputBuffers, + HttpClient httpClient, + Executor executor, + ScheduledExecutorService updateScheduledExecutor, + ScheduledExecutorService errorScheduledExecutor, + Duration maxErrorDuration, + Duration taskStatusRefreshMaxWait, + Duration taskInfoUpdateInterval, + boolean summarizeTaskInfo, + Codec taskStatusCodec, + Codec taskInfoCodec, + Codec taskUpdateRequestCodec, + PartitionedSplitCountTracker partitionedSplitCountTracker, + RemoteTaskStats stats, boolean isBinaryEncoding, + Optional 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 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); diff --git a/presto-main/src/main/java/io/prestosql/testing/TestingTaskContext.java b/presto-main/src/main/java/io/prestosql/testing/TestingTaskContext.java index 514575f98..faf46ad2c 100644 --- a/presto-main/src/main/java/io/prestosql/testing/TestingTaskContext.java +++ b/presto-main/src/main/java/io/prestosql/testing/TestingTaskContext.java @@ -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) diff --git a/presto-main/src/test/java/io/prestosql/execution/MockRemoteTaskFactory.java b/presto-main/src/test/java/io/prestosql/execution/MockRemoteTaskFactory.java index 4e1b8b7ad..2edb3ce9f 100644 --- a/presto-main/src/test/java/io/prestosql/execution/MockRemoteTaskFactory.java +++ b/presto-main/src/test/java/io/prestosql/execution/MockRemoteTaskFactory.java @@ -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 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 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); diff --git a/presto-main/src/test/java/io/prestosql/execution/TaskTestUtils.java b/presto-main/src/test/java/io/prestosql/execution/TaskTestUtils.java index 4d421e252..7551e1f1b 100644 --- a/presto-main/src/test/java/io/prestosql/execution/TaskTestUtils.java +++ b/presto-main/src/test/java/io/prestosql/execution/TaskTestUtils.java @@ -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 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() diff --git a/presto-main/src/test/java/io/prestosql/execution/TestMemoryRevokingScheduler.java b/presto-main/src/test/java/io/prestosql/execution/TestMemoryRevokingScheduler.java index 7dcf650ed..85662dfdf 100644 --- a/presto-main/src/test/java/io/prestosql/execution/TestMemoryRevokingScheduler.java +++ b/presto-main/src/test/java/io/prestosql/execution/TestMemoryRevokingScheduler.java @@ -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"); diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java b/presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java index e6f9fc1e8..c377bb1bc 100644 --- a/presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java +++ b/presto-main/src/test/java/io/prestosql/execution/TestSqlStageExecution.java @@ -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()); } } diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSqlTask.java b/presto-main/src/test/java/io/prestosql/execution/TestSqlTask.java index 3e0bb5164..8595a3f3d 100644 --- a/presto-main/src/test/java/io/prestosql/execution/TestSqlTask.java +++ b/presto-main/src/test/java/io/prestosql/execution/TestSqlTask.java @@ -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, diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSqlTaskExecution.java b/presto-main/src/test/java/io/prestosql/execution/TestSqlTaskExecution.java index df9247f07..6a524efb9 100644 --- a/presto-main/src/test/java/io/prestosql/execution/TestSqlTaskExecution.java +++ b/presto-main/src/test/java/io/prestosql/execution/TestSqlTaskExecution.java @@ -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) diff --git a/presto-main/src/test/java/io/prestosql/operator/TestCommonTableExpressionOperator.java b/presto-main/src/test/java/io/prestosql/operator/TestCommonTableExpressionOperator.java new file mode 100644 index 000000000..17e44d4f3 --- /dev/null +++ b/presto-main/src/test/java/io/prestosql/operator/TestCommonTableExpressionOperator.java @@ -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 toPages(Operator operator) + { + ImmutableList.Builder 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(); + } +}