Examples with TaskStateMachine io.trino.execution.TaskStateMachine used on opensource projects

Example 1 with TaskStateMachine

use of io.trino.execution.TaskStateMachine in project trino by trinodb.

the class TestHashJoinOperator method innerJoinWithSpill.

private void innerJoinWithSpill(boolean probeHashEnabled, List<WhenSpill> whenSpill, SingleStreamSpillerFactory buildSpillerFactory, PartitioningSpillerFactory joinSpillerFactory) throws Exception {
    TaskStateMachine taskStateMachine = new TaskStateMachine(new TaskId(new StageId("query", 0), 0, 0), executor);
    TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TEST_SESSION, taskStateMachine);
    DriverContext joinDriverContext = taskContext.addPipelineContext(2, true, true, false).addDriverContext();
    // force a yield for every match in LookupJoinOperator, set called to true after first
    AtomicBoolean called = new AtomicBoolean(false);
    InternalJoinFilterFunction filterFunction = new TestInternalJoinFilterFunction((leftPosition, leftPage, rightPosition, rightPage) -> {
        called.set(true);
        joinDriverContext.getYieldSignal().forceYieldForTesting();
        return true;
    });
    // build factory
    RowPagesBuilder buildPages = rowPagesBuilder(false, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).addSequencePage(4, 20, 200).addSequencePage(4, 20, 200).addSequencePage(4, 30, 300).addSequencePage(4, 40, 400);
    BuildSideSetup buildSideSetup = setupBuildSide(nodePartitioningManager, true, taskContext, buildPages, Optional.of(filterFunction), true, buildSpillerFactory);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = buildSideSetup.getLookupSourceFactoryManager();
    // probe factory
    RowPagesBuilder probePages = rowPagesBuilder(probeHashEnabled, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).row("20", 123_000L).row("20", 123_000L).pageBreak().addSequencePage(20, 0, 123_000).addSequencePage(10, 30, 123_000);
    OperatorFactory joinOperatorFactory = innerJoinOperatorFactory(operatorFactories, lookupSourceFactoryManager, probePages, joinSpillerFactory, true);
    // build drivers and operators
    instantiateBuildDrivers(buildSideSetup, taskContext);
    List<Driver> buildDrivers = buildSideSetup.getBuildDrivers();
    int buildOperatorCount = buildDrivers.size();
    checkState(buildOperatorCount == whenSpill.size());
    LookupSourceFactory lookupSourceFactory = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide());
    try (Operator joinOperator = joinOperatorFactory.createOperator(joinDriverContext)) {
        // build lookup source
        ListenableFuture<LookupSourceProvider> lookupSourceProvider = lookupSourceFactory.createLookupSourceProvider();
        List<Boolean> revoked = new ArrayList<>(nCopies(buildOperatorCount, false));
        while (!lookupSourceProvider.isDone()) {
            for (int i = 0; i < buildOperatorCount; i++) {
                checkErrors(taskStateMachine);
                buildDrivers.get(i).process();
                HashBuilderOperator buildOperator = buildSideSetup.getBuildOperators().get(i);
                if (whenSpill.get(i) == WhenSpill.DURING_BUILD && buildOperator.getOperatorContext().getReservedRevocableBytes() > 0) {
                    checkState(!lookupSourceProvider.isDone(), "Too late, LookupSource already done");
                    revokeMemory(buildOperator);
                    revoked.set(i, true);
                }
            }
        }
        getFutureValue(lookupSourceProvider).close();
        assertEquals(revoked, whenSpill.stream().map(WhenSpill.DURING_BUILD::equals).collect(toImmutableList()), "Some operators not spilled before LookupSource built");
        for (int i = 0; i < buildOperatorCount; i++) {
            if (whenSpill.get(i) == WhenSpill.AFTER_BUILD) {
                revokeMemory(buildSideSetup.getBuildOperators().get(i));
            }
        }
        for (Driver buildDriver : buildDrivers) {
            runDriverInThread(executor, buildDriver);
        }
        ValuesOperatorFactory valuesOperatorFactory = new ValuesOperatorFactory(17, new PlanNodeId("values"), probePages.build());
        PageBuffer pageBuffer = new PageBuffer(10);
        PageBufferOperatorFactory pageBufferOperatorFactory = new PageBufferOperatorFactory(18, new PlanNodeId("pageBuffer"), pageBuffer);
        Driver joinDriver = Driver.createDriver(joinDriverContext, valuesOperatorFactory.createOperator(joinDriverContext), joinOperator, pageBufferOperatorFactory.createOperator(joinDriverContext));
        while (!called.get()) {
            // process first row of first page of LookupJoinOperator
            processRow(joinDriver, taskStateMachine);
        }
        for (int i = 0; i < buildOperatorCount; i++) {
            if (whenSpill.get(i) == WhenSpill.DURING_USAGE) {
                triggerMemoryRevokingAndWait(buildSideSetup.getBuildOperators().get(i), taskStateMachine);
            }
        }
        // process remaining LookupJoinOperator pages
        while (!joinDriver.isFinished()) {
            checkErrors(taskStateMachine);
            processRow(joinDriver, taskStateMachine);
        }
        checkErrors(taskStateMachine);
        List<Page> actualPages = getPages(pageBuffer);
        MaterializedResult expected = MaterializedResult.resultBuilder(taskContext.getSession(), concat(probePages.getTypesWithoutHash(), buildPages.getTypesWithoutHash())).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("20", 123_000L, "20", 200L).row("30", 123_000L, "30", 300L).row("31", 123_001L, "31", 301L).row("32", 123_002L, "32", 302L).row("33", 123_003L, "33", 303L).build();
        assertEqualsIgnoreOrder(getProperColumns(joinOperator, concat(probePages.getTypes(), buildPages.getTypes()), probePages, actualPages).getMaterializedRows(), expected.getMaterializedRows());
    } finally {
        joinOperatorFactory.noMoreOperators();
    }
}

Example 2 with TaskStateMachine

use of io.trino.execution.TaskStateMachine in project trino by trinodb.

the class TestMemoryTracking method setUpTest.

@BeforeMethod
public void setUpTest() {
    memoryPool = new MemoryPool(memoryPoolSize);
    queryContext = new QueryContext(new QueryId("test_query"), queryMaxMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker);
    taskContext = queryContext.addTaskContext(new TaskStateMachine(new TaskId(new StageId("query", 0), 0, 0), notificationExecutor), testSessionBuilder().build(), () -> {
    }, true, true);
    pipelineContext = taskContext.addPipelineContext(0, true, true, false);
    driverContext = pipelineContext.addDriverContext();
    operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("a"), "test-operator");
}

Example 3 with TaskStateMachine

use of io.trino.execution.TaskStateMachine in project trino by trinodb.

the class TestHashJoinOperator method testBuildGracefulSpill.

@Test(timeOut = 30_000)
public void testBuildGracefulSpill() throws Exception {
    TaskStateMachine taskStateMachine = new TaskStateMachine(new TaskId(new StageId("query", 0), 0, 0), executor);
    TaskContext taskContext = TestingTaskContext.createTaskContext(executor, scheduledExecutor, TEST_SESSION, taskStateMachine);
    // build factory
    RowPagesBuilder buildPages = rowPagesBuilder(false, Ints.asList(0), ImmutableList.of(VARCHAR, BIGINT)).addSequencePage(4, 20, 200);
    DummySpillerFactory buildSpillerFactory = new DummySpillerFactory();
    BuildSideSetup buildSideSetup = setupBuildSide(nodePartitioningManager, true, taskContext, buildPages, Optional.empty(), true, buildSpillerFactory);
    instantiateBuildDrivers(buildSideSetup, taskContext);
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = buildSideSetup.getLookupSourceFactoryManager();
    PartitionedLookupSourceFactory lookupSourceFactory = lookupSourceFactoryManager.getJoinBridge(Lifespan.taskWide());
    // finish probe before any build partition is spilled
    lookupSourceFactory.finishProbeOperator(OptionalInt.of(1));
    // spill build partition after probe is finished
    HashBuilderOperator hashBuilderOperator = buildSideSetup.getBuildOperators().get(0);
    hashBuilderOperator.startMemoryRevoke().get();
    hashBuilderOperator.finishMemoryRevoke();
    hashBuilderOperator.finish();
    // hash builder operator should not deadlock waiting for spilled lookup source to be disposed
    hashBuilderOperator.isBlocked().get();
    lookupSourceFactory.destroy();
    assertTrue(hashBuilderOperator.isFinished());
}

Example 4 with TaskStateMachine

use of io.trino.execution.TaskStateMachine in project trino by trinodb.

the class AbstractOperatorBenchmark method runOnce.

@Override
protected Map<String, Long> runOnce() {
    Session session = testSessionBuilder().setSystemProperty("optimizer.optimize-hash-generation", "true").setTransactionId(this.session.getRequiredTransactionId()).build();
    MemoryPool memoryPool = new MemoryPool(DataSize.of(1, GIGABYTE));
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(DataSize.of(1, GIGABYTE));
    TaskContext taskContext = new QueryContext(new QueryId("test"), DataSize.of(256, MEGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), DataSize.of(256, MEGABYTE), spillSpaceTracker).addTaskContext(new TaskStateMachine(new TaskId(new StageId("query", 0), 0, 0), localQueryRunner.getExecutor()), session, () -> {
    }, false, false);
    CpuTimer cpuTimer = new CpuTimer();
    Map<String, Long> executionStats = execute(taskContext);
    CpuDuration executionTime = cpuTimer.elapsedTime();
    TaskStats taskStats = taskContext.getTaskStats();
    long inputRows = taskStats.getRawInputPositions();
    long inputBytes = taskStats.getRawInputDataSize().toBytes();
    long outputRows = taskStats.getOutputPositions();
    long outputBytes = taskStats.getOutputDataSize().toBytes();
    double inputMegaBytes = ((double) inputBytes) / MEGABYTE.inBytes();
    return ImmutableMap.<String, Long>builder().putAll(executionStats).put("elapsed_millis", executionTime.getWall().toMillis()).put("input_rows_per_second", (long) (inputRows / executionTime.getWall().getValue(SECONDS))).put("output_rows_per_second", (long) (outputRows / executionTime.getWall().getValue(SECONDS))).put("input_megabytes", (long) inputMegaBytes).put("input_megabytes_per_second", (long) (inputMegaBytes / executionTime.getWall().getValue(SECONDS))).put("wall_nanos", executionTime.getWall().roundTo(NANOSECONDS)).put("cpu_nanos", executionTime.getCpu().roundTo(NANOSECONDS)).put("user_nanos", executionTime.getUser().roundTo(NANOSECONDS)).put("input_rows", inputRows).put("input_bytes", inputBytes).put("output_rows", outputRows).put("output_bytes", outputBytes).buildOrThrow();
}

Example 5 with TaskStateMachine

use of io.trino.execution.TaskStateMachine in project trino by trinodb.

the class MemoryLocalQueryRunner method execute.

public List<Page> execute(@Language("SQL") String query) {
    MemoryPool memoryPool = new MemoryPool(DataSize.of(2, GIGABYTE));
    SpillSpaceTracker spillSpaceTracker = new SpillSpaceTracker(DataSize.of(1, GIGABYTE));
    QueryContext queryContext = new QueryContext(new QueryId("test"), DataSize.of(1, GIGABYTE), memoryPool, new TestingGcMonitor(), localQueryRunner.getExecutor(), localQueryRunner.getScheduler(), DataSize.of(4, GIGABYTE), spillSpaceTracker);
    TaskContext taskContext = queryContext.addTaskContext(new TaskStateMachine(new TaskId(new StageId("query", 0), 0, 0), localQueryRunner.getExecutor()), localQueryRunner.getDefaultSession(), () -> {
    }, false, false);
    // Use NullOutputFactory to avoid coping out results to avoid affecting benchmark results
    ImmutableList.Builder<Page> output = ImmutableList.builder();
    List<Driver> drivers = localQueryRunner.createDrivers(query, new PageConsumerOperator.PageConsumerOutputFactory(types -> output::add), taskContext);
    boolean done = false;
    while (!done) {
        boolean processed = false;
        for (Driver driver : drivers) {
            if (!driver.isFinished()) {
                driver.process();
                processed = true;
            }
        }
        done = !processed;
    }
    return output.build();
}