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

Example 16 with PlanFragmentId

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

the class TestDynamicFilterService method createPlan.

private static PlanFragment createPlan(DynamicFilterId consumedDynamicFilterId, DynamicFilterId producedDynamicFilterId, PartitioningHandle stagePartitioning, ExchangeNode.Type exchangeType) {
    Symbol symbol = new Symbol("column");
    Symbol buildSymbol = new Symbol("buildColumn");
    PlanNodeId tableScanNodeId = new PlanNodeId("plan_id");
    TableScanNode tableScan = TableScanNode.newInstance(tableScanNodeId, TEST_TABLE_HANDLE, ImmutableList.of(symbol), ImmutableMap.of(symbol, new TestingMetadata.TestingColumnHandle("column")), false, Optional.empty());
    FilterNode filterNode = new FilterNode(new PlanNodeId("filter_node_id"), tableScan, createDynamicFilterExpression(session, createTestMetadataManager(), consumedDynamicFilterId, VARCHAR, symbol.toSymbolReference()));
    RemoteSourceNode remote = new RemoteSourceNode(new PlanNodeId("remote_id"), new PlanFragmentId("plan_fragment_id"), ImmutableList.of(buildSymbol), Optional.empty(), exchangeType, RetryPolicy.NONE);
    return new PlanFragment(new PlanFragmentId("plan_id"), new JoinNode(new PlanNodeId("join_id"), INNER, filterNode, remote, ImmutableList.of(), tableScan.getOutputSymbols(), remote.getOutputSymbols(), false, Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), Optional.empty(), ImmutableMap.of(producedDynamicFilterId, buildSymbol), Optional.empty()), ImmutableMap.of(symbol, VARCHAR), stagePartitioning, ImmutableList.of(tableScanNodeId), new PartitioningScheme(Partitioning.create(SINGLE_DISTRIBUTION, ImmutableList.of()), ImmutableList.of(symbol)), ungroupedExecution(), StatsAndCosts.empty(), Optional.empty());
}

Example 17 with PlanFragmentId

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

the class TestPhasedExecutionSchedule method testAggregation.

@Test
public void testAggregation() {
    PlanFragment sourceFragment = createTableScanPlanFragment("probe");
    PlanFragment aggregationFragment = createAggregationFragment("aggregation", sourceFragment);
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, REPLICATED, "join", buildFragment, aggregationFragment);
    TestingStageExecution sourceStage = new TestingStageExecution(sourceFragment);
    TestingStageExecution aggregationStage = new TestingStageExecution(aggregationFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(sourceStage, aggregationStage, buildStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId(), joinFragment.getId());
    // aggregation and source stage should start immediately, join stage should wait for build stage to complete
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactly(new FragmentsEdge(buildFragment.getId(), joinFragment.getId()));
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
}

Example 18 with PlanFragmentId

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

the class TestPhasedExecutionSchedule method testDependentStageAbortedBeforeStarted.

@Test
public void testDependentStageAbortedBeforeStarted() {
    PlanFragment sourceFragment = createTableScanPlanFragment("probe");
    PlanFragment aggregationFragment = createAggregationFragment("aggregation", sourceFragment);
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, REPLICATED, "join", buildFragment, aggregationFragment);
    TestingStageExecution sourceStage = new TestingStageExecution(sourceFragment);
    TestingStageExecution aggregationStage = new TestingStageExecution(aggregationFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(sourceStage, aggregationStage, buildStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId(), joinFragment.getId());
    // aggregation and source stage should start immediately, join stage should wait for build stage to complete
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactly(new FragmentsEdge(buildFragment.getId(), joinFragment.getId()));
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), sourceFragment.getId(), aggregationFragment.getId());
    // abort non-active join stage
    joinStage.setState(ABORTED);
    // dependencies finish
    buildStage.setState(FINISHED);
    aggregationStage.setState(FINISHED);
    sourceStage.setState(FINISHED);
    // join stage already aborted. Whole schedule should be marked as finished
    schedule.schedule();
    assertThat(schedule.isFinished()).isTrue();
}

Example 19 with PlanFragmentId

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

the class TestPhasedExecutionSchedule method testPartitionedJoin.

@Test
public void testPartitionedJoin() {
    PlanFragment buildFragment = createTableScanPlanFragment("build");
    PlanFragment probeFragment = createTableScanPlanFragment("probe");
    PlanFragment joinFragment = createJoinPlanFragment(INNER, PARTITIONED, "join", buildFragment, probeFragment);
    TestingStageExecution buildStage = new TestingStageExecution(buildFragment);
    TestingStageExecution probeStage = new TestingStageExecution(probeFragment);
    TestingStageExecution joinStage = new TestingStageExecution(joinFragment);
    PhasedExecutionSchedule schedule = PhasedExecutionSchedule.forStages(ImmutableSet.of(buildStage, probeStage, joinStage), dynamicFilterService);
    assertThat(schedule.getSortedFragments()).containsExactly(buildFragment.getId(), probeFragment.getId(), joinFragment.getId());
    // single dependency between build and probe stages
    DirectedGraph<PlanFragmentId, FragmentsEdge> dependencies = schedule.getFragmentDependency();
    assertThat(dependencies.edgeSet()).containsExactlyInAnyOrder(new FragmentsEdge(buildFragment.getId(), probeFragment.getId()));
    // build and join stage should start immediately
    assertThat(getActiveFragments(schedule)).containsExactly(buildFragment.getId(), joinFragment.getId());
    // probe stage should start after build stage is completed
    ListenableFuture<Void> rescheduleFuture = schedule.getRescheduleFuture().orElseThrow();
    assertThat(rescheduleFuture).isNotDone();
    buildStage.setState(FLUSHING);
    assertThat(rescheduleFuture).isDone();
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).containsExactly(joinFragment.getId(), probeFragment.getId());
    // make sure scheduler finishes
    rescheduleFuture = schedule.getRescheduleFuture().orElseThrow();
    assertThat(rescheduleFuture).isNotDone();
    probeStage.setState(FINISHED);
    assertThat(rescheduleFuture).isNotDone();
    joinStage.setState(FINISHED);
    schedule.schedule();
    assertThat(getActiveFragments(schedule)).isEmpty();
    assertThat(schedule.isFinished()).isTrue();
}

Example 20 with PlanFragmentId

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

the class StageTaskSourceFactory method getInputsForRemoteSources.

private static ListMultimap<PlanNodeId, ExchangeSourceHandle> getInputsForRemoteSources(List<RemoteSourceNode> remoteSources, Multimap<PlanFragmentId, ExchangeSourceHandle> exchangeSourceHandles) {
    ImmutableListMultimap.Builder<PlanNodeId, ExchangeSourceHandle> result = ImmutableListMultimap.builder();
    for (RemoteSourceNode remoteSource : remoteSources) {
        for (PlanFragmentId fragmentId : remoteSource.getSourceFragmentIds()) {
            Collection<ExchangeSourceHandle> handles = requireNonNull(exchangeSourceHandles.get(fragmentId), () -> "exchange source handle is missing for fragment: " + fragmentId);
            if (remoteSource.getExchangeType() == GATHER || remoteSource.getExchangeType() == REPLICATE) {
                checkArgument(handles.size() <= 1, "at most 1 exchange source handle is expected, got: %s", handles);
            }
            result.putAll(remoteSource.getId(), handles);
        }
    }
    return result.build();
}