Examples with PlanNodeId io.trino.sql.planner.plan.PlanNodeId used on opensource projects

Example 96 with PlanNodeId

use of io.trino.sql.planner.plan.PlanNodeId in project trino by trinodb.

the class CountAggregationBenchmark method createOperatorFactories.

@Override
protected List<? extends OperatorFactory> createOperatorFactories() {
    OperatorFactory tableScanOperator = createTableScanOperator(0, new PlanNodeId("test"), "orders", "orderkey");
    AggregationOperatorFactory aggregationOperator = new AggregationOperatorFactory(1, new PlanNodeId("test"), ImmutableList.of(countFunction.bind(ImmutableList.of(0))));
    return ImmutableList.of(tableScanOperator, aggregationOperator);
}

Example 97 with PlanNodeId

use of io.trino.sql.planner.plan.PlanNodeId in project trino by trinodb.

the class SqlTaskExecution method schedulePartitionedSource.

private synchronized void schedulePartitionedSource(SplitAssignment splitAssignmentUpdate) {
    mergeIntoPendingSplits(splitAssignmentUpdate.getPlanNodeId(), splitAssignmentUpdate.getSplits(), splitAssignmentUpdate.getNoMoreSplitsForLifespan(), splitAssignmentUpdate.isNoMoreSplits());
    while (true) {
        // SchedulingLifespanManager tracks how far each Lifespan has been scheduled. Here is an example.
        // Let's say there are 4 source pipelines/nodes: A, B, C, and D, in scheduling order.
        // And we're processing 3 concurrent lifespans at a time. In this case, we could have
        // 
        // * Lifespan 10:  A   B  [C]  D; i.e. Pipeline A and B has finished scheduling (but not necessarily finished running).
        // * Lifespan 20: [A]  B   C   D
        // * Lifespan 30:  A  [B]  C   D
        // 
        // To recap, SchedulingLifespanManager records the next scheduling source node for each lifespan.
        Iterator<SchedulingLifespan> activeLifespans = schedulingLifespanManager.getActiveLifespans();
        boolean madeProgress = false;
        while (activeLifespans.hasNext()) {
            SchedulingLifespan schedulingLifespan = activeLifespans.next();
            Lifespan lifespan = schedulingLifespan.getLifespan();
            // This is why getSchedulingPlanNode returns an Optional.
            while (true) {
                Optional<PlanNodeId> optionalSchedulingPlanNode = schedulingLifespan.getSchedulingPlanNode();
                if (optionalSchedulingPlanNode.isEmpty()) {
                    break;
                }
                PlanNodeId schedulingPlanNode = optionalSchedulingPlanNode.get();
                DriverSplitRunnerFactory partitionedDriverRunnerFactory = driverRunnerFactoriesWithSplitLifeCycle.get(schedulingPlanNode);
                PendingSplits pendingSplits = pendingSplitsByPlanNode.get(schedulingPlanNode).getLifespan(lifespan);
                // Enqueue driver runners with driver group lifecycle for this driver life cycle, if not already enqueued.
                if (!lifespan.isTaskWide() && !schedulingLifespan.getAndSetDriversForDriverGroupLifeCycleScheduled()) {
                    scheduleDriversForDriverGroupLifeCycle(lifespan);
                }
                // Enqueue driver runners with split lifecycle for this plan node and driver life cycle combination.
                ImmutableList.Builder<DriverSplitRunner> runners = ImmutableList.builder();
                for (ScheduledSplit scheduledSplit : pendingSplits.removeAllSplits()) {
                    // create a new driver for the split
                    runners.add(partitionedDriverRunnerFactory.createDriverRunner(scheduledSplit, lifespan));
                }
                enqueueDriverSplitRunner(false, runners.build());
                // move on to the next plan node.
                if (pendingSplits.getState() != NO_MORE_SPLITS) {
                    break;
                }
                partitionedDriverRunnerFactory.noMoreDriverRunner(ImmutableList.of(lifespan));
                pendingSplits.markAsCleanedUp();
                schedulingLifespan.nextPlanNode();
                madeProgress = true;
                if (schedulingLifespan.isDone()) {
                    break;
                }
            }
        }
        if (!madeProgress) {
            break;
        }
    }
    if (splitAssignmentUpdate.isNoMoreSplits()) {
        schedulingLifespanManager.noMoreSplits(splitAssignmentUpdate.getPlanNodeId());
    }
}

Example 98 with PlanNodeId

use of io.trino.sql.planner.plan.PlanNodeId in project trino by trinodb.

the class ExtractSpatialJoins method tryCreateSpatialJoin.

private static Result tryCreateSpatialJoin(Context context, JoinNode joinNode, Expression filter, PlanNodeId nodeId, List<Symbol> outputSymbols, ComparisonExpression spatialComparison, PlannerContext plannerContext, SplitManager splitManager, PageSourceManager pageSourceManager, TypeAnalyzer typeAnalyzer) {
    PlanNode leftNode = joinNode.getLeft();
    PlanNode rightNode = joinNode.getRight();
    List<Symbol> leftSymbols = leftNode.getOutputSymbols();
    List<Symbol> rightSymbols = rightNode.getOutputSymbols();
    Expression radius;
    Optional<Symbol> newRadiusSymbol;
    ComparisonExpression newComparison;
    if (spatialComparison.getOperator() == LESS_THAN || spatialComparison.getOperator() == LESS_THAN_OR_EQUAL) {
        // ST_Distance(a, b) <= r
        radius = spatialComparison.getRight();
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new ComparisonExpression(spatialComparison.getOperator(), spatialComparison.getLeft(), toExpression(newRadiusSymbol, radius));
        } else {
            return Result.empty();
        }
    } else {
        // r >= ST_Distance(a, b)
        radius = spatialComparison.getLeft();
        Set<Symbol> radiusSymbols = extractUnique(radius);
        if (radiusSymbols.isEmpty() || (rightSymbols.containsAll(radiusSymbols) && containsNone(leftSymbols, radiusSymbols))) {
            newRadiusSymbol = newRadiusSymbol(context, radius);
            newComparison = new ComparisonExpression(spatialComparison.getOperator().flip(), spatialComparison.getRight(), toExpression(newRadiusSymbol, radius));
        } else {
            return Result.empty();
        }
    }
    Expression newFilter = replaceExpression(filter, ImmutableMap.of(spatialComparison, newComparison));
    PlanNode newRightNode = newRadiusSymbol.map(symbol -> addProjection(context, rightNode, symbol, radius)).orElse(rightNode);
    JoinNode newJoinNode = new JoinNode(joinNode.getId(), joinNode.getType(), leftNode, newRightNode, joinNode.getCriteria(), joinNode.getLeftOutputSymbols(), joinNode.getRightOutputSymbols(), joinNode.isMaySkipOutputDuplicates(), Optional.of(newFilter), joinNode.getLeftHashSymbol(), joinNode.getRightHashSymbol(), joinNode.getDistributionType(), joinNode.isSpillable(), joinNode.getDynamicFilters(), joinNode.getReorderJoinStatsAndCost());
    return tryCreateSpatialJoin(context, newJoinNode, newFilter, nodeId, outputSymbols, (FunctionCall) newComparison.getLeft(), Optional.of(newComparison.getRight()), plannerContext, splitManager, pageSourceManager, typeAnalyzer);
}

Example 99 with PlanNodeId

use of io.trino.sql.planner.plan.PlanNodeId in project trino by trinodb.

the class EliminateCrossJoins method buildJoinTree.

public static PlanNode buildJoinTree(List<Symbol> expectedOutputSymbols, JoinGraph graph, List<Integer> joinOrder, PlanNodeIdAllocator idAllocator) {
    requireNonNull(expectedOutputSymbols, "expectedOutputSymbols is null");
    requireNonNull(idAllocator, "idAllocator is null");
    requireNonNull(graph, "graph is null");
    joinOrder = ImmutableList.copyOf(requireNonNull(joinOrder, "joinOrder is null"));
    checkArgument(joinOrder.size() >= 2);
    PlanNode result = graph.getNode(joinOrder.get(0));
    Set<PlanNodeId> alreadyJoinedNodes = new HashSet<>();
    alreadyJoinedNodes.add(result.getId());
    for (int i = 1; i < joinOrder.size(); i++) {
        PlanNode rightNode = graph.getNode(joinOrder.get(i));
        alreadyJoinedNodes.add(rightNode.getId());
        ImmutableList.Builder<JoinNode.EquiJoinClause> criteria = ImmutableList.builder();
        for (JoinGraph.Edge edge : graph.getEdges(rightNode)) {
            PlanNode targetNode = edge.getTargetNode();
            if (alreadyJoinedNodes.contains(targetNode.getId())) {
                criteria.add(new JoinNode.EquiJoinClause(edge.getTargetSymbol(), edge.getSourceSymbol()));
            }
        }
        result = new JoinNode(idAllocator.getNextId(), JoinNode.Type.INNER, result, rightNode, criteria.build(), result.getOutputSymbols(), rightNode.getOutputSymbols(), false, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(), Optional.empty());
    }
    List<Expression> filters = graph.getFilters();
    for (Expression filter : filters) {
        result = new FilterNode(idAllocator.getNextId(), result, filter);
    }
    // Some nodes are sensitive to what's produced (e.g., DistinctLimit node)
    return restrictOutputs(idAllocator, result, ImmutableSet.copyOf(expectedOutputSymbols)).orElse(result);
}

Example 100 with PlanNodeId

use of io.trino.sql.planner.plan.PlanNodeId in project trino by trinodb.

the class TestNodeScheduler method testTopologyAwareScheduling.

@Test(timeOut = 60 * 1000)
public void testTopologyAwareScheduling() {
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    InMemoryNodeManager nodeManager = new InMemoryNodeManager();
    ImmutableList.Builder<InternalNode> nodeBuilder = ImmutableList.builder();
    nodeBuilder.add(new InternalNode("node1", URI.create("http://host1.rack1:11"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new InternalNode("node2", URI.create("http://host2.rack1:12"), NodeVersion.UNKNOWN, false));
    nodeBuilder.add(new InternalNode("node3", URI.create("http://host3.rack2:13"), NodeVersion.UNKNOWN, false));
    ImmutableList<InternalNode> nodes = nodeBuilder.build();
    nodeManager.addNode(CONNECTOR_ID, nodes);
    // contents of taskMap indicate the node-task map for the current stage
    Map<InternalNode, RemoteTask> taskMap = new HashMap<>();
    NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig().setMaxSplitsPerNode(25).setIncludeCoordinator(false).setMaxPendingSplitsPerTask(20);
    TestNetworkTopology topology = new TestNetworkTopology();
    NodeSelectorFactory nodeSelectorFactory = new TopologyAwareNodeSelectorFactory(topology, nodeManager, nodeSchedulerConfig, nodeTaskMap, getNetworkTopologyConfig());
    NodeScheduler nodeScheduler = new NodeScheduler(nodeSelectorFactory);
    NodeSelector nodeSelector = nodeScheduler.createNodeSelector(session, Optional.of(CONNECTOR_ID));
    // Fill up the nodes with non-local data
    ImmutableSet.Builder<Split> nonRackLocalBuilder = ImmutableSet.builder();
    for (int i = 0; i < (25 + 11) * 3; i++) {
        nonRackLocalBuilder.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("data.other_rack", 1)), Lifespan.taskWide()));
    }
    Set<Split> nonRackLocalSplits = nonRackLocalBuilder.build();
    Multimap<InternalNode, Split> assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int task = 0;
    for (InternalNode node : assignments.keySet()) {
        TaskId taskId = new TaskId(new StageId("test", 1), task, 0);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(25);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    // Continue assigning to fill up part of the queue
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    assignments = nodeSelector.computeAssignments(nonRackLocalSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    nonRackLocalSplits = Sets.difference(nonRackLocalSplits, new HashSet<>(assignments.values()));
    // Check that 3 of the splits were rejected, since they're non-local
    assertEquals(nonRackLocalSplits.size(), 3);
    // Assign rack-local splits
    ImmutableSet.Builder<Split> rackLocalSplits = ImmutableSet.builder();
    HostAddress dataHost1 = HostAddress.fromParts("data.rack1", 1);
    HostAddress dataHost2 = HostAddress.fromParts("data.rack2", 1);
    for (int i = 0; i < 6 * 2; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(dataHost1), Lifespan.taskWide()));
    }
    for (int i = 0; i < 6; i++) {
        rackLocalSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(dataHost2), Lifespan.taskWide()));
    }
    assignments = nodeSelector.computeAssignments(rackLocalSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    Set<Split> unassigned = Sets.difference(rackLocalSplits.build(), new HashSet<>(assignments.values()));
    // Compute the assignments a second time to account for the fact that some splits may not have been assigned due to asynchronous
    // loading of the NetworkLocationCache
    assignments = nodeSelector.computeAssignments(unassigned, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments.get(node)).build());
    }
    unassigned = Sets.difference(unassigned, new HashSet<>(assignments.values()));
    assertEquals(unassigned.size(), 3);
    int rack1 = 0;
    int rack2 = 0;
    for (Split split : unassigned) {
        String rack = topology.locate(split.getAddresses().get(0)).getSegments().get(0);
        switch(rack) {
            case "rack1":
                rack1++;
                break;
            case "rack2":
                rack2++;
                break;
            default:
                throw new AssertionError("Unexpected rack: " + rack);
        }
    }
    assertEquals(rack1, 2);
    assertEquals(rack2, 1);
    // Assign local splits
    ImmutableSet.Builder<Split> localSplits = ImmutableSet.builder();
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host1.rack1", 1)), Lifespan.taskWide()));
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host2.rack1", 1)), Lifespan.taskWide()));
    localSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromParts("host3.rack2", 1)), Lifespan.taskWide()));
    assignments = nodeSelector.computeAssignments(localSplits.build(), ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertEquals(assignments.size(), 3);
    assertEquals(assignments.keySet().size(), 3);
}