Examples with TaskId io.prestosql.execution.TaskId used on opensource projects

Example 1 with TaskId

use of io.prestosql.execution.TaskId in project hetu-core by openlookeng.

the class TestDynamicFilterUtil method registerDf.

public static void registerDf(String filterId, Session session, JoinNode.DistributionType joinType, DynamicFilterService dynamicFilterService) {
    JoinNode node = mock(JoinNode.class);
    HashMap<String, Symbol> dfs = new HashMap<>();
    List<JoinNode.EquiJoinClause> criteria = new ArrayList<JoinNode.EquiJoinClause>();
    Symbol right = new Symbol("rightCol");
    Symbol left = new Symbol("leftCol");
    JoinNode.EquiJoinClause clause = new JoinNode.EquiJoinClause(left, right);
    criteria.add(clause);
    dfs.put(filterId, right);
    when(node.getCriteria()).thenReturn(criteria);
    when(node.getDynamicFilters()).thenReturn(dfs);
    when(node.getDistributionType()).thenReturn(Optional.of(joinType));
    RemoteSourceNode leftNode = mock(RemoteSourceNode.class);
    when(node.getLeft()).thenReturn(leftNode);
    HashSet<TaskId> tasks = new HashSet<>();
    tasks.add(new TaskId("task1.0"));
    tasks.add(new TaskId("task1.1"));
    StageStateMachine stateMachine = mock(StageStateMachine.class);
    when(stateMachine.getSession()).thenReturn(session);
    InternalNode worker = mock(InternalNode.class);
    InternalNode worker2 = mock(InternalNode.class);
    HashSet<InternalNode> workers = new HashSet<>();
    when(worker.getNodeIdentifier()).thenReturn("w1");
    when(worker2.getNodeIdentifier()).thenReturn("w2");
    workers.add(worker);
    workers.add(worker2);
    dynamicFilterService.registerTasks(node, tasks, workers, stateMachine);
}

Example 2 with TaskId

use of io.prestosql.execution.TaskId in project hetu-core by openlookeng.

the class TestMemoryTracking method setUpTest.

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

Example 3 with TaskId

use of io.prestosql.execution.TaskId in project hetu-core by openlookeng.

the class TestSnapshotStateId method testFromString.

@Test
public void testFromString() {
    String testQueryId = "testingquery";
    long snapshotId = 18L;
    int stageId = 500;
    int taskInt = 600;
    TaskId taskId = new TaskId(testQueryId, stageId, taskInt);
    SnapshotStateId notString = new SnapshotStateId(snapshotId, taskId);
    String str = notString.toString();
    SnapshotStateId withString = SnapshotStateId.fromString(str);
    Assert.assertTrue(notString.equals(withString));
}

Example 4 with TaskId

use of io.prestosql.execution.TaskId in project hetu-core by openlookeng.

the class TestSnapshotStateId method testWithSnapshotId.

@Test
public void testWithSnapshotId() {
    String queryId = "query1";
    long snapshotId = 1L;
    int stageId = 100;
    int taskId = 200;
    TaskId taskId1 = new TaskId(queryId, stageId, taskId);
    SnapshotStateId stateId = new SnapshotStateId(snapshotId, taskId1, "component");
    stateId = stateId.withSnapshotId(2);
    Assert.assertEquals(stateId.getId(), "query1/2/100/200/component");
}

Example 5 with TaskId

use of io.prestosql.execution.TaskId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testSpilledDeleted.

@Test
public void testSpilledDeleted() throws Exception {
    queryId = new QueryId("spilleddeletedquery");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 0);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    snapshotManager.setTotalComponents(1);
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId1);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // first store
    long firstSnapshotId = 1L;
    File dirs = new File("/tmp/test_snapshot_manager/" + queryId + "/");
    File firstFile = new File("/tmp/test_snapshot_manager/" + queryId + "/firstFile");
    dirs.mkdirs();
    firstFile.createNewFile();
    try (FileWriter fw1 = new FileWriter(firstFile)) {
        String firstStr = "first string";
        fw1.write(firstStr);
        SnapshotStateId stateId = SnapshotStateId.forOperator(firstSnapshotId, taskId1, 3, 4, 5);
        snapshotManager.storeFile(stateId, firstFile.toPath());
        snapshotManager.succeededToCapture(stateId);
    }
    // second store, then deleted
    File secondFile = new File("/tmp/test_snapshot_manager/" + queryId + "/secondFile");
    secondFile.createNewFile();
    try (FileWriter fw2 = new FileWriter(secondFile)) {
        String secondStr = "second string";
        fw2.write(secondStr);
        SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5);
        snapshotManager.storeFile(secondId, secondFile.toPath());
        snapshotManager.succeededToCapture(secondId);
        querySnapshotManager.snapshotInitiated(firstSnapshotId);
        querySnapshotManager.updateQueryCapture(taskId1, Collections.singletonMap(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
        File secondFileOperator = new File("/tmp/test_snapshot_manager/" + queryId + "/2/1/0/3/4/5/secondFile");
        assertTrue(secondFileOperator.exists());
        secondFileOperator.delete();
        assertFalse(secondFileOperator.exists());
        assertFalse(snapshotManager.loadFile(secondId, secondFile.toPath()));
    }
}