Examples with ValuesOperatorFactory io.trino.operator.ValuesOperator.ValuesOperatorFactory used on opensource projects

Example 1 with ValuesOperatorFactory

use of io.trino.operator.ValuesOperator.ValuesOperatorFactory in project trino by trinodb.

the class TestSqlTaskExecution method testComplex.

@Test(dataProvider = "executionStrategies", timeOut = 20_000)
public void testComplex(PipelineExecutionStrategy executionStrategy) throws Exception {
    ScheduledExecutorService taskNotificationExecutor = newScheduledThreadPool(10, threadsNamed("task-notification-%s"));
    ScheduledExecutorService driverYieldExecutor = newScheduledThreadPool(2, threadsNamed("driver-yield-%s"));
    TaskExecutor taskExecutor = new TaskExecutor(5, 10, 3, 4, Ticker.systemTicker());
    taskExecutor.start();
    try {
        TaskStateMachine taskStateMachine = new TaskStateMachine(TASK_ID, taskNotificationExecutor);
        PartitionedOutputBuffer outputBuffer = newTestingOutputBuffer(taskNotificationExecutor);
        OutputBufferConsumer outputBufferConsumer = new OutputBufferConsumer(outputBuffer, OUTPUT_BUFFER_ID);
        // test initialization: complex test with 4 pipelines
        // Take a task with the following set of pipelines for example:
        // 
        // pipeline 0        pipeline 1       pipeline 2    pipeline 3    ... pipeline id
        // partitioned      unpartitioned     partitioned  unpartitioned  ... partitioned/unpartitioned pipeline
        // grouped           grouped          grouped      ungrouped    ... execution strategy (in grouped test)
        // ungrouped         ungrouped        ungrouped     ungrouped    ... execution strategy (in ungrouped test)
        // 
        // TaskOutput-0
        // |
        // CrossJoin-C  ................................... Build-C
        // |                                               |
        // CrossJoin-A  ..... Build-A                       Values-3
        // |                |
        // Scan-0         CrossJoin-B  ....  Build-B
        // (effectively ExchangeSink)    |
        // |               Scan-2
        // Values-1
        // (1 row)
        // 
        // CrossJoin operator here has the same lifecycle behavior as a real cross/hash-join, and produces
        // the correct number of rows, but doesn't actually produce a cross-join for simplicity.
        // 
        // A single task can never have all 4 combinations: partitioned/unpartitioned x grouped/ungrouped.
        // * In the case of ungrouped test, this test covers driver with
        // 1) split lifecycle (partitioned ungrouped)
        // 2) task lifecycle (unpartitioned ungrouped)
        // These are the only 2 possible pipeline execution strategy a task can have if the task has ungrouped execution strategy.
        // * In the case of grouped test, this covers:
        // 1) split lifecycle (partitioned grouped)
        // 2) driver group lifecycle (unpartitioned grouped)
        // 3) task lifecycle (unpartitioned ungrouped)
        // These are the only 3 possible pipeline execution strategy a task can have if the task has grouped execution strategy.
        // 
        // The following behaviors are tested:
        // * DriverFactory are marked as noMoreDriver/Operator for particular lifespans as soon as they can be:
        // * immediately, if the pipeline has task lifecycle (ungrouped and unpartitioned).
        // * when SplitAssignment containing the lifespan is encountered, if the pipeline has driver group lifecycle (grouped and unpartitioned).
        // * when SplitAssignment indicate that no more splits will be produced for the plan node (and plan nodes that schedule before it
        // due to phased scheduling) and lifespan combination, if the pipeline has split lifecycle (partitioned).
        // * DriverFactory are marked as noMoreDriver/Operator as soon as they can be:
        // * immediately, if the pipeline has task lifecycle (ungrouped and unpartitioned).
        // * when SplitAssignment indicate that will no more splits, otherwise.
        // * Driver groups are marked as completed as soon as they should be:
        // * when there are no active driver, and all DriverFactory for the lifespan (across all pipelines) are marked as completed.
        // * Rows are produced as soon as they should be:
        // * streams data through as soon as the build side is ready, for CrossJoin
        // * streams data through, otherwise.
        PlanNodeId scan0NodeId = new PlanNodeId("scan-0");
        PlanNodeId values1NodeId = new PlanNodeId("values-1");
        PlanNodeId scan2NodeId = new PlanNodeId("scan-2");
        PlanNodeId values3NodeId = new PlanNodeId("values-3");
        PlanNodeId joinANodeId = new PlanNodeId("join-a");
        PlanNodeId joinBNodeId = new PlanNodeId("join-b");
        PlanNodeId joinCNodeId = new PlanNodeId("join-c");
        BuildStates buildStatesA = new BuildStates(executionStrategy);
        BuildStates buildStatesB = new BuildStates(executionStrategy);
        BuildStates buildStatesC = new BuildStates(UNGROUPED_EXECUTION);
        TestingScanOperatorFactory scanOperatorFactory0 = new TestingScanOperatorFactory(1, scan0NodeId, ImmutableList.of(VARCHAR));
        ValuesOperatorFactory valuesOperatorFactory1 = new ValuesOperatorFactory(101, values1NodeId, ImmutableList.of(new Page(createStringsBlock("multiplier1"))));
        TestingScanOperatorFactory scanOperatorFactory2 = new TestingScanOperatorFactory(201, scan2NodeId, ImmutableList.of(VARCHAR));
        ValuesOperatorFactory valuesOperatorFactory3 = new ValuesOperatorFactory(301, values3NodeId, ImmutableList.of(new Page(createStringsBlock("x", "y", "multiplier3"))));
        TaskOutputOperatorFactory taskOutputOperatorFactory = new TaskOutputOperatorFactory(4, joinCNodeId, outputBuffer, Function.identity(), new PagesSerdeFactory(new TestingBlockEncodingSerde(), false));
        TestingCrossJoinOperatorFactory joinOperatorFactoryA = new TestingCrossJoinOperatorFactory(2, joinANodeId, buildStatesA);
        TestingCrossJoinOperatorFactory joinOperatorFactoryB = new TestingCrossJoinOperatorFactory(102, joinBNodeId, buildStatesB);
        TestingCrossJoinOperatorFactory joinOperatorFactoryC = new TestingCrossJoinOperatorFactory(3, joinCNodeId, buildStatesC);
        TestingBuildOperatorFactory buildOperatorFactoryA = new TestingBuildOperatorFactory(103, joinANodeId, buildStatesA);
        TestingBuildOperatorFactory buildOperatorFactoryB = new TestingBuildOperatorFactory(202, joinBNodeId, buildStatesB);
        TestingBuildOperatorFactory buildOperatorFactoryC = new TestingBuildOperatorFactory(302, joinCNodeId, buildStatesC);
        LocalExecutionPlan localExecutionPlan = new LocalExecutionPlan(ImmutableList.of(new DriverFactory(0, true, true, ImmutableList.of(scanOperatorFactory0, joinOperatorFactoryA, joinOperatorFactoryC, taskOutputOperatorFactory), OptionalInt.empty(), executionStrategy), new DriverFactory(1, false, false, ImmutableList.of(valuesOperatorFactory1, joinOperatorFactoryB, buildOperatorFactoryA), OptionalInt.empty(), executionStrategy), new DriverFactory(2, true, false, ImmutableList.of(scanOperatorFactory2, buildOperatorFactoryB), OptionalInt.empty(), executionStrategy), new DriverFactory(3, false, false, ImmutableList.of(valuesOperatorFactory3, buildOperatorFactoryC), OptionalInt.empty(), UNGROUPED_EXECUTION)), ImmutableList.of(scan2NodeId, scan0NodeId), executionStrategy == GROUPED_EXECUTION ? StageExecutionDescriptor.fixedLifespanScheduleGroupedExecution(ImmutableList.of(scan0NodeId, scan2NodeId)) : StageExecutionDescriptor.ungroupedExecution());
        TaskContext taskContext = newTestingTaskContext(taskNotificationExecutor, driverYieldExecutor, taskStateMachine);
        SqlTaskExecution sqlTaskExecution = SqlTaskExecution.createSqlTaskExecution(taskStateMachine, taskContext, outputBuffer, localExecutionPlan, taskExecutor, taskNotificationExecutor, createTestSplitMonitor());
        // 
        // test body
        assertEquals(taskStateMachine.getState(), RUNNING);
        switch(executionStrategy) {
            case UNGROUPED_EXECUTION:
                // assert that pipeline 1 and pipeline 3 will have no more drivers
                // (Unpartitioned ungrouped pipelines can have all driver instance created up front.)
                waitUntilEquals(joinOperatorFactoryB::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryA::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryC::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // add assignment for pipeline 2, and mark as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan2NodeId, ImmutableSet.of(newScheduledSplit(0, scan2NodeId, Lifespan.taskWide(), 100000, 1), newScheduledSplit(1, scan2NodeId, Lifespan.taskWide(), 300000, 2)), false)));
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan2NodeId, ImmutableSet.of(newScheduledSplit(2, scan2NodeId, Lifespan.taskWide(), 300000, 2)), true)));
                // assert that pipeline 2 will have no more drivers
                waitUntilEquals(scanOperatorFactory2::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryB::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                scanOperatorFactory0.getPauser().pause();
                // add assignment for pipeline 0, mark as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan0NodeId, ImmutableSet.of(newScheduledSplit(3, scan0NodeId, Lifespan.taskWide(), 400000, 100)), true)));
                // assert that pipeline 0 will have no more drivers
                waitUntilEquals(scanOperatorFactory0::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryA::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryC::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that no DriverGroup is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of());
                // resume operator execution
                scanOperatorFactory0.getPauser().resume();
                // assert that task result is produced
                outputBufferConsumer.consume(100 * 5 * 3, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                break;
            case GROUPED_EXECUTION:
                // assert that pipeline 3 will have no more drivers
                // (Unpartitioned ungrouped pipelines can have all driver instances created up front.)
                waitUntilEquals(buildOperatorFactoryC::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // add assignment for pipeline 2 driver group 3, and mark driver group 3 as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan2NodeId, ImmutableSet.of(newScheduledSplit(0, scan2NodeId, Lifespan.driverGroup(3), 0, 1), newScheduledSplit(1, scan2NodeId, Lifespan.driverGroup(3), 100000, 2)), false)));
                // assert that pipeline 1 driver group [3] will have no more drivers
                waitUntilEquals(joinOperatorFactoryB::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryA::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan2NodeId, ImmutableSet.of(newScheduledSplit(2, scan2NodeId, Lifespan.driverGroup(3), 200000, 2)), ImmutableSet.of(Lifespan.driverGroup(3)), false)));
                // assert that pipeline 2 driver group [3] will have no more drivers
                waitUntilEquals(scanOperatorFactory2::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryB::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                scanOperatorFactory0.getPauser().pause();
                // add assignment for pipeline 0 driver group 3, and mark driver group 3 as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan0NodeId, ImmutableSet.of(newScheduledSplit(3, scan0NodeId, Lifespan.driverGroup(3), 300000, 10)), ImmutableSet.of(Lifespan.driverGroup(3)), false)));
                // assert that pipeline 0 driver group [3] will have no more drivers
                waitUntilEquals(scanOperatorFactory0::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryA::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryC::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                // assert that no DriverGroup is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of());
                // resume operator execution
                scanOperatorFactory0.getPauser().resume();
                // assert that partial task result is produced
                outputBufferConsumer.consume(10 * 5 * 3, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [3] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT);
                // add assignment for pipeline 2 driver group 7, and mark pipeline as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan2NodeId, ImmutableSet.of(newScheduledSplit(4, scan2NodeId, Lifespan.driverGroup(7), 400000, 2)), ImmutableSet.of(Lifespan.driverGroup(7)), true)));
                // assert that pipeline 2 driver group [3, 7] will have no more drivers
                waitUntilEquals(scanOperatorFactory2::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryB::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                // pause operator execution to make sure that
                // * operatorFactory will be closed even though operator can't execute
                // * completedDriverGroups will NOT include the newly scheduled driver group while pause is in place
                scanOperatorFactory0.getPauser().pause();
                // add assignment for pipeline 0 driver group 7, mark pipeline as no more splits
                sqlTaskExecution.addSplitAssignments(ImmutableList.of(new SplitAssignment(scan0NodeId, ImmutableSet.of(newScheduledSplit(5, scan0NodeId, Lifespan.driverGroup(7), 500000, 1000)), ImmutableSet.of(Lifespan.driverGroup(7)), true)));
                // assert that pipeline 0 driver group [3, 7] will have no more drivers
                waitUntilEquals(scanOperatorFactory0::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryA::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryC::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                // assert that pipeline 0 will have no more drivers
                waitUntilEquals(scanOperatorFactory0::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryA::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(joinOperatorFactoryC::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that pipeline 1 driver group [3, 7] will have no more drivers
                waitUntilEquals(joinOperatorFactoryB::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryA::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                // assert that pipeline 1 will have no more drivers
                // (Unpartitioned grouped pipelines will have no more driver instances when there can be no more driver groups.)
                waitUntilEquals(joinOperatorFactoryB::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryA::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that pipeline 2 will have no more drivers
                // note: One could argue that this should have happened as soon as pipeline 2 driver group 7 is marked as noMoreSplits.
                // This is not how SqlTaskExecution is currently implemented. And such a delay in closing DriverFactory does not matter much.
                waitUntilEquals(scanOperatorFactory2::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                waitUntilEquals(buildOperatorFactoryB::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT);
                // assert that driver group [3] (but not 7) is fully completed
                assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(3)));
                // resume operator execution
                scanOperatorFactory0.getPauser().resume();
                // assert that partial task result is produced
                outputBufferConsumer.consume(1000 * 2 * 3, ASSERT_WAIT_TIMEOUT);
                outputBufferConsumer.assertBufferComplete(ASSERT_WAIT_TIMEOUT);
                // assert that driver group [3, 7] is fully completed
                waitUntilEquals(taskContext::getCompletedDriverGroups, ImmutableSet.of(Lifespan.driverGroup(3), Lifespan.driverGroup(7)), ASSERT_WAIT_TIMEOUT);
                break;
            default:
                throw new UnsupportedOperationException();
        }
        assertEquals(taskStateMachine.getStateChange(RUNNING).get(10, SECONDS), FLUSHING);
        // complete the task by calling abort on it
        outputBufferConsumer.abort();
        assertEquals(taskStateMachine.getStateChange(FLUSHING).get(10, SECONDS), FINISHED);
    } finally {
        taskExecutor.stop();
        taskNotificationExecutor.shutdownNow();
        driverYieldExecutor.shutdown();
    }
}

Example 2 with ValuesOperatorFactory

use of io.trino.operator.ValuesOperator.ValuesOperatorFactory 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 3 with ValuesOperatorFactory

use of io.trino.operator.ValuesOperator.ValuesOperatorFactory in project trino by trinodb.

the class TestNestedLoopJoinOperator method newJoinOperatorFactoryWithCompletedBuild.

private static NestedLoopJoinOperatorFactory newJoinOperatorFactoryWithCompletedBuild(TaskContext taskContext, RowPagesBuilder buildPages, List<Integer> probeChannels, List<Integer> buildChannels) {
    DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext();
    ValuesOperatorFactory valuesOperatorFactory = new ValuesOperatorFactory(0, new PlanNodeId("test"), buildPages.build());
    JoinBridgeManager<NestedLoopJoinBridge> nestedLoopJoinBridgeManager = new JoinBridgeManager<>(false, PipelineExecutionStrategy.UNGROUPED_EXECUTION, PipelineExecutionStrategy.UNGROUPED_EXECUTION, lifespan -> new NestedLoopJoinPagesSupplier(), buildPages.getTypes());
    NestedLoopBuildOperatorFactory nestedLoopBuildOperatorFactory = new NestedLoopBuildOperatorFactory(1, new PlanNodeId("test"), nestedLoopJoinBridgeManager);
    NestedLoopJoinOperatorFactory joinOperatorFactory = new NestedLoopJoinOperatorFactory(3, new PlanNodeId("test"), nestedLoopJoinBridgeManager, probeChannels, buildChannels);
    Operator valuesOperator = valuesOperatorFactory.createOperator(driverContext);
    Operator nestedLoopBuildOperator = nestedLoopBuildOperatorFactory.createOperator(driverContext);
    Driver driver = Driver.createDriver(driverContext, valuesOperator, nestedLoopBuildOperator);
    valuesOperatorFactory.noMoreOperators();
    nestedLoopBuildOperatorFactory.noMoreOperators();
    while (nestedLoopBuildOperator.isBlocked().isDone()) {
        driver.process();
    }
    return joinOperatorFactory;
}

Example 4 with ValuesOperatorFactory

use of io.trino.operator.ValuesOperator.ValuesOperatorFactory in project trino by trinodb.

the class HashBuildBenchmark method createDrivers.

@Override
protected List<Driver> createDrivers(TaskContext taskContext) {
    // hash build
    List<Type> ordersTypes = getColumnTypes("orders", "orderkey", "totalprice");
    OperatorFactory ordersTableScan = createTableScanOperator(0, new PlanNodeId("test"), "orders", "orderkey", "totalprice");
    BlockTypeOperators blockTypeOperators = new BlockTypeOperators(new TypeOperators());
    JoinBridgeManager<PartitionedLookupSourceFactory> lookupSourceFactoryManager = JoinBridgeManager.lookupAllAtOnce(new PartitionedLookupSourceFactory(ordersTypes, ImmutableList.of(0, 1).stream().map(ordersTypes::get).collect(toImmutableList()), Ints.asList(0).stream().map(ordersTypes::get).collect(toImmutableList()), 1, false, blockTypeOperators));
    HashBuilderOperatorFactory hashBuilder = new HashBuilderOperatorFactory(1, new PlanNodeId("test"), lookupSourceFactoryManager, ImmutableList.of(0, 1), Ints.asList(0), OptionalInt.empty(), Optional.empty(), Optional.empty(), ImmutableList.of(), 1_500_000, new PagesIndex.TestingFactory(false), false, SingleStreamSpillerFactory.unsupportedSingleStreamSpillerFactory(), incrementalLoadFactorHashArraySizeSupplier(session));
    DriverFactory hashBuildDriverFactory = new DriverFactory(0, true, true, ImmutableList.of(ordersTableScan, hashBuilder), OptionalInt.empty(), UNGROUPED_EXECUTION);
    // empty join so build finishes
    ImmutableList.Builder<OperatorFactory> joinDriversBuilder = ImmutableList.builder();
    joinDriversBuilder.add(new ValuesOperatorFactory(0, new PlanNodeId("values"), ImmutableList.of()));
    OperatorFactory joinOperator = operatorFactories.innerJoin(2, new PlanNodeId("test"), lookupSourceFactoryManager, false, false, false, ImmutableList.of(BIGINT), Ints.asList(0), OptionalInt.empty(), Optional.empty(), OptionalInt.empty(), unsupportedPartitioningSpillerFactory(), blockTypeOperators);
    joinDriversBuilder.add(joinOperator);
    joinDriversBuilder.add(new NullOutputOperatorFactory(3, new PlanNodeId("test")));
    DriverFactory joinDriverFactory = new DriverFactory(1, true, true, joinDriversBuilder.build(), OptionalInt.empty(), UNGROUPED_EXECUTION);
    Driver hashBuildDriver = hashBuildDriverFactory.createDriver(taskContext.addPipelineContext(0, true, true, false).addDriverContext());
    hashBuildDriverFactory.noMoreDrivers();
    Driver joinDriver = joinDriverFactory.createDriver(taskContext.addPipelineContext(1, true, true, false).addDriverContext());
    joinDriverFactory.noMoreDrivers();
    return ImmutableList.of(hashBuildDriver, joinDriver);
}