Examples with QueryId io.prestosql.spi.QueryId used on opensource projects

Example 26 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testConsolidatedOverlap.

@Test
public void testConsolidatedOverlap() throws Exception {
    queryId = new QueryId("consolidateoverlapquery");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 0);
    TaskId taskId2 = new TaskId(queryId.getId(), 1, 1);
    TaskSnapshotManager snapshotManager1 = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    TaskSnapshotManager snapshotManager2 = new TaskSnapshotManager(taskId2, 0, snapshotUtils);
    snapshotManager1.setTotalComponents(1);
    snapshotManager2.setTotalComponents(1);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // statexy means task x snapshot y
    MockState state11 = new MockState("mockstate1.1");
    MockState state12 = new MockState("mockstate1.2");
    MockState state21 = new MockState("mockstate2.1");
    MockState state22 = new MockState("mockstate2.2");
    // similar naming as above
    SnapshotStateId stateId11 = SnapshotStateId.forOperator(1L, taskId1, 3, 4, 5);
    SnapshotStateId stateId12 = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5);
    SnapshotStateId stateId21 = SnapshotStateId.forOperator(1L, taskId2, 3, 4, 5);
    SnapshotStateId stateId22 = SnapshotStateId.forOperator(2L, taskId2, 3, 4, 5);
    // save 11 (part of snapshot 1), then 12 and 22 (snapshot 2), then finish snapshot 1 (21)
    snapshotManager1.storeConsolidatedState(stateId11, state11, 0);
    snapshotManager2.storeConsolidatedState(stateId22, state22, 0);
    snapshotManager1.storeConsolidatedState(stateId12, state12, 0);
    snapshotManager1.succeededToCapture(stateId12);
    snapshotManager2.succeededToCapture(stateId22);
    snapshotManager2.storeConsolidatedState(stateId21, state21, 0);
    snapshotManager1.succeededToCapture(stateId11);
    snapshotManager2.succeededToCapture(stateId21);
    // similar naming as above
    MockState newState11 = (MockState) snapshotManager1.loadConsolidatedState(stateId11).get();
    MockState newState12 = (MockState) snapshotManager1.loadConsolidatedState(stateId12).get();
    MockState newState21 = (MockState) snapshotManager2.loadConsolidatedState(stateId21).get();
    MockState newState22 = (MockState) snapshotManager2.loadConsolidatedState(stateId22).get();
    Assert.assertEquals(state11.getState(), newState11.getState());
    Assert.assertEquals(state12.getState(), newState12.getState());
    Assert.assertEquals(state21.getState(), newState21.getState());
    Assert.assertEquals(state22.getState(), newState22.getState());
}

Example 27 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testCapture.

@Test
public void testCapture() throws Exception {
    queryId = new QueryId("query");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 1);
    TaskId taskId2 = new TaskId(queryId.getId(), 1, 2);
    TaskSnapshotManager snapshotManager1 = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    TaskSnapshotManager snapshotManager2 = new TaskSnapshotManager(taskId2, 0, snapshotUtils);
    snapshotManager1.setTotalComponents(2);
    snapshotManager2.setTotalComponents(2);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // Test capture successfully
    snapshotManager1.succeededToCapture(new SnapshotStateId(1, taskId1, "component1"));
    Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.IN_PROGRESS);
    snapshotManager1.succeededToCapture(new SnapshotStateId(1, taskId1, "component2"));
    Assert.assertEquals(snapshotManager1.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.SUCCESSFUL);
    // Test capture failed
    snapshotManager2.failedToCapture(new SnapshotStateId(1, taskId2, "component1"));
    Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.IN_PROGRESS_FAILED);
    snapshotManager2.failedToCapture(new SnapshotStateId(1, taskId2, "component2"));
    Assert.assertEquals(snapshotManager2.getSnapshotCaptureResult().get(1L).getSnapshotResult(), SnapshotResult.FAILED);
}

Example 28 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method TestStoreAndLoad.

@Test
public void TestStoreAndLoad() throws Exception {
    queryId = new QueryId("saveandload");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 2);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // Test operator state
    MockState operatorState = new MockState("operator-state");
    SnapshotStateId operatorStateId = SnapshotStateId.forOperator(1L, taskId1, 3, 4, 5);
    snapshotManager.storeState(operatorStateId, operatorState, 0);
    snapshotManager.setTotalComponents(1);
    snapshotManager.succeededToCapture(operatorStateId);
    MockState newOperatorState = (MockState) snapshotManager.loadState(operatorStateId).get();
    Assert.assertEquals(operatorState.getState(), newOperatorState.getState());
    // Test task state
    MockState taskState = new MockState("task-state");
    TaskId taskId2 = new TaskId(queryId.getId(), 3, 4);
    SnapshotStateId taskStateId = SnapshotStateId.forOperator(3L, taskId2, 5, 6, 7);
    snapshotManager.storeState(taskStateId, taskState, 0);
    snapshotManager.succeededToCapture(taskStateId);
    MockState newTaskState = (MockState) snapshotManager.loadState(taskStateId).get();
    Assert.assertEquals(taskState.getState(), newTaskState.getState());
}

Example 29 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method TestLoadBacktrack.

@Test
public void TestLoadBacktrack() throws Exception {
    queryId = new QueryId("loadbacktrack");
    StageId stageId = new StageId(queryId, 0);
    TaskId taskId = new TaskId(stageId, 0);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId, 0, snapshotUtils);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    // Save operator state
    MockState state = new MockState("state");
    SnapshotStateId stateId = SnapshotStateId.forOperator(1L, taskId, 3, 4, 5);
    snapshotManager.storeState(stateId, state, 0);
    snapshotManager.setTotalComponents(1);
    snapshotManager.succeededToCapture(stateId);
    // Make snapshot manager think that snapshot #1 is complete
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId);
    querySnapshotManager.snapshotInitiated(1L);
    querySnapshotManager.updateQueryCapture(taskId, Collections.singletonMap(1L, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL)));
    querySnapshotManager.getResumeSnapshotId();
    SnapshotStateId newStateId = stateId.withSnapshotId(2);
    Optional<Object> loadedState = snapshotManager.loadState(newStateId);
    // Should return state saved for snapshot #1
    assertTrue(loadedState.isPresent());
    Assert.assertEquals(((MockState) loadedState.get()).getState(), state.getState());
}

Example 30 with QueryId

use of io.prestosql.spi.QueryId in project hetu-core by openlookeng.

the class TestTaskSnapshotManager method testConsolidatedBacktrack.

@Test
public void testConsolidatedBacktrack() throws Exception {
    queryId = new QueryId("consolidatebacktrackquery");
    TaskId taskId1 = new TaskId(queryId.getId(), 1, 0);
    TaskSnapshotManager snapshotManager = new TaskSnapshotManager(taskId1, 0, snapshotUtils);
    snapshotManager.setTotalComponents(1);
    snapshotUtils.getOrCreateQuerySnapshotManager(queryId, TEST_SNAPSHOT_SESSION);
    QuerySnapshotManager querySnapshotManager = new QuerySnapshotManager(queryId, snapshotUtils, TEST_SNAPSHOT_SESSION);
    querySnapshotManager.addNewTask(taskId1);
    // first store
    long firstSnapshotId = 1L;
    MockState state = new MockState("mockstate");
    SnapshotStateId stateId = SnapshotStateId.forOperator(firstSnapshotId, taskId1, 3, 4, 5);
    snapshotManager.storeConsolidatedState(stateId, state, 0);
    snapshotManager.succeededToCapture(stateId);
    // second store (null)
    SnapshotStateId secondId = SnapshotStateId.forOperator(2L, taskId1, 3, 4, 5);
    snapshotManager.storeConsolidatedState(secondId, null, 0);
    snapshotManager.succeededToCapture(secondId);
    Map<Long, SnapshotInfo> queryCaptureResult = new LinkedHashMap<>();
    queryCaptureResult.put(firstSnapshotId, SnapshotInfo.withStatus(SnapshotResult.SUCCESSFUL));
    snapshotUtils.storeSnapshotResult(queryId.getId(), queryCaptureResult);
    MockState newState = (MockState) snapshotManager.loadConsolidatedState(secondId).get();
    Assert.assertEquals(state.getState(), newState.getState());
}