@Override protected List<Driver> createDrivers(TaskContext taskContext) { DriverFactory driverFactory = createDriverFactory(); DriverContext driverContext = taskContext.addPipelineContext(0, true, true, false).addDriverContext(); Driver driver = driverFactory.createDriver(driverContext); return ImmutableList.of(driver); } }
protected Map<String, Long> execute(TaskContext taskContext) { List<Driver> drivers = createDrivers(taskContext); long peakMemory = 0; boolean done = false; while (!done) { boolean processed = false; for (Driver driver : drivers) { if (!driver.isFinished()) { driver.process(); long lastPeakMemory = peakMemory; peakMemory = (long) taskContext.getTaskStats().getUserMemoryReservation().getValue(BYTE); if (peakMemory <= lastPeakMemory) { peakMemory = lastPeakMemory; } processed = true; } } done = !processed; } return ImmutableMap.of("peak_memory", peakMemory); }
public Session getSession() { return taskContext.getSession(); }
public static TaskContext createTaskContext( QueryContext queryContext, TaskStateMachine taskStateMachine, GcMonitor gcMonitor, Executor notificationExecutor, ScheduledExecutorService yieldExecutor, Session session, MemoryTrackingContext taskMemoryContext, boolean perOperatorCpuTimerEnabled, boolean cpuTimerEnabled, OptionalInt totalPartitions) { TaskContext taskContext = new TaskContext(queryContext, taskStateMachine, gcMonitor, notificationExecutor, yieldExecutor, session, taskMemoryContext, perOperatorCpuTimerEnabled, cpuTimerEnabled, totalPartitions); taskContext.initialize(); return taskContext; }
public SqlTaskIoStats getIoStats() { // if we are finished, return the final IoStats if (finalIoStats != null) { return finalIoStats; } // if we haven't started yet, return an empty IoStats if (taskExecution == null) { return new SqlTaskIoStats(); } // get IoStats from the current task execution TaskContext taskContext = taskExecution.getTaskContext(); return new SqlTaskIoStats(taskContext.getInputDataSize(), taskContext.getInputPositions(), taskContext.getOutputDataSize(), taskContext.getOutputPositions()); } }
long physicalWrittenBytes = 0; TaskContext taskContext = taskHolder.getTaskExecution().getTaskContext(); for (PipelineContext pipelineContext : taskContext.getPipelineContexts()) { PipelineStatus pipelineStatus = pipelineContext.getPipelineStatus(); queuedPartitionedDrivers += pipelineStatus.getQueuedPartitionedDrivers(); userMemoryReservation = taskContext.getMemoryReservation(); systemMemoryReservation = taskContext.getSystemMemoryReservation(); completedDriverGroups = taskContext.getCompletedDriverGroups(); fullGcCount = taskContext.getFullGcCount(); fullGcTime = taskContext.getFullGcTime();
@Test public void testEmptyBuildPage() { TaskContext taskContext = createTaskContext(); // build List<Type> buildTypes = ImmutableList.of(VARCHAR); RowPagesBuilder buildPages = rowPagesBuilder(buildTypes) .pageBreak(); // probe List<Type> probeTypes = ImmutableList.of(VARCHAR); RowPagesBuilder probePages = rowPagesBuilder(probeTypes); List<Page> probeInput = probePages .row("A") .row("B") .pageBreak() .build(); NestedLoopJoinOperatorFactory joinOperatorFactory = newJoinOperatorFactoryWithCompletedBuild(taskContext, buildPages); // expected MaterializedResult expected = resultBuilder(taskContext.getSession(), concat(probeTypes, buildPages.getTypes())) .build(); assertOperatorEquals(joinOperatorFactory, taskContext.addPipelineContext(0, true, true, false).addDriverContext(), probeInput, expected); }
@BeforeMethod public void setUp() { executor = newCachedThreadPool(daemonThreadsNamed("test-executor-%s")); scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-scheduledExecutor-%s")); TaskContext taskContext = TestingTaskContext.builder(executor, scheduledExecutor, TEST_SESSION) .setQueryMaxMemory(DataSize.valueOf("100MB")) .setMemoryPoolSize(DataSize.valueOf("10B")) .setQueryId(QUERY_ID) .build(); memoryPool = taskContext.getQueryContext().getMemoryPool(); driverContext = taskContext .addPipelineContext(0, true, true, false) .addDriverContext(); }
public void failed(Throwable cause) { taskContext.failed(cause); }
assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of()); assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1))); assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(1), Lifespan.driverGroup(5)));
public TaskContext addTaskContext(TaskStateMachine taskStateMachine, Session session, DataSize operatorPreAllocatedMemory, boolean verboseStats, boolean cpuTimerEnabled) { TaskContext taskContext = new TaskContext(this, taskStateMachine, executor, session, operatorPreAllocatedMemory, verboseStats, cpuTimerEnabled); taskContexts.add(taskContext); return taskContext; } }
@Test public void testEmptyProbePage() { TaskContext taskContext = createTaskContext(); // build List<Type> buildTypes = ImmutableList.of(VARCHAR); RowPagesBuilder buildPages = rowPagesBuilder(buildTypes) .row("A") .row("B") .pageBreak() .row("C"); // probe List<Type> probeTypes = ImmutableList.of(VARCHAR); RowPagesBuilder probePages = rowPagesBuilder(probeTypes); List<Page> probeInput = probePages .pageBreak() .build(); NestedLoopJoinOperatorFactory joinOperatorFactory = newJoinOperatorFactoryWithCompletedBuild(taskContext, buildPages); // expected MaterializedResult expected = resultBuilder(taskContext.getSession(), concat(probeTypes, buildPages.getTypes())) .build(); assertOperatorEquals(joinOperatorFactory, taskContext.addPipelineContext(0, true, true, false).addDriverContext(), probeInput, expected); }
public SqlTaskIoStats getIoStats() { // if we are finished, return the final IoStats if (finalIoStats != null) { return finalIoStats; } // if we haven't started yet, return an empty IoStats if (taskExecution == null) { return new SqlTaskIoStats(); } // get IoStats from the current task execution TaskContext taskContext = taskExecution.getTaskContext(); return new SqlTaskIoStats(taskContext.getInputDataSize(), taskContext.getInputPositions(), taskContext.getOutputDataSize(), taskContext.getOutputPositions()); } }
public void failed(Throwable cause) { taskContext.failed(cause); }
waitUntilEquals(joinOperatorFactoryC::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT); assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of()); waitUntilEquals(joinOperatorFactoryC::getDriverGroupsWithNoMoreOperators, ImmutableSet.of(Lifespan.driverGroup(3)), ASSERT_WAIT_TIMEOUT); assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of()); waitUntilEquals(buildOperatorFactoryB::isOverallNoMoreOperators, true, ASSERT_WAIT_TIMEOUT); assertEquals(taskContext.getCompletedDriverGroups(), ImmutableSet.of(Lifespan.driverGroup(3)));
@Test public void testDistributedSpatialSelfJoin() { TaskContext taskContext = createTaskContext(); DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext(); RowPagesBuilder pages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)) .row(POLYGON_A, "A", 1) .row(POLYGON_A, "A", 2) .row(null, "null", null) .pageBreak() .row(POLYGON_B, "B", 0) .row(POLYGON_B, "B", 2); MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)) .row("A", "A") .row("A", "B") .row("B", "A") .row("B", "B") .build(); PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.intersects(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), pages); OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, pages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory); assertOperatorEquals(joinOperatorFactory, driverContext, pages.build(), expected); }
private synchronized void initializeStateIfNecessary() { if (pipelineContext == null) { TaskContext taskContext = taskContextReference.get(); checkState(taskContext != null, "Task context must be set before index can be built"); pipelineContext = taskContext.addPipelineContext(indexBuildDriverFactoryProvider.getPipelineId(), true, true, false); } if (indexSnapshotLoader == null) { indexSnapshotLoader = new IndexSnapshotLoader( indexBuildDriverFactoryProvider, pipelineContext, indexSnapshotReference, lookupSourceInputChannels, keyTypes, keyOutputChannels, keyOutputHashChannel, expectedPositions, maxIndexMemorySize, pagesIndexFactory, joinCompiler); } }
private TaskStats getTaskStats(TaskHolder taskHolder) { TaskInfo finalTaskInfo = taskHolder.getFinalTaskInfo(); if (finalTaskInfo != null) { return finalTaskInfo.getStats(); } SqlTaskExecution taskExecution = taskHolder.getTaskExecution(); if (taskExecution != null) { return taskExecution.getTaskContext().getTaskStats(); } // if the task completed without creation, set end time DateTime endTime = taskStateMachine.getState().isDone() ? DateTime.now() : null; return new TaskStats(taskStateMachine.getCreatedTime(), endTime); }
public Session getSession() { return taskContext.getSession(); }
@Test public void testDistributedSpatialJoin() { TaskContext taskContext = createTaskContext(); DriverContext driverContext = taskContext.addPipelineContext(0, true, true, true).addDriverContext(); RowPagesBuilder buildPages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)) .row(POLYGON_A, "A", 1) .row(POLYGON_A, "A", 2) .row(null, "null", null) .pageBreak() .row(POLYGON_B, "B", 0) .row(POLYGON_B, "B", 2); RowPagesBuilder probePages = rowPagesBuilder(ImmutableList.of(GEOMETRY, VARCHAR, INTEGER)) .row(POINT_X, "x", 2) .row(null, "null", null) .row(POINT_Y, "y", 2) .pageBreak() .row(POINT_Z, "z", 0); MaterializedResult expected = resultBuilder(taskContext.getSession(), ImmutableList.of(VARCHAR, VARCHAR)) .row("x", "A") .row("y", "A") .row("y", "B") .row("z", "B") .build(); PagesSpatialIndexFactory pagesSpatialIndexFactory = buildIndex(driverContext, (build, probe, r) -> build.contains(probe), Optional.empty(), Optional.of(2), Optional.of(KDB_TREE_JSON), Optional.empty(), buildPages); OperatorFactory joinOperatorFactory = new SpatialJoinOperatorFactory(2, new PlanNodeId("test"), INNER, probePages.getTypes(), Ints.asList(1), 0, Optional.of(2), pagesSpatialIndexFactory); assertOperatorEquals(joinOperatorFactory, driverContext, probePages.build(), expected); }