@Override public ConnectorTransactionHandle beginTransaction(IsolationLevel isolationLevel, boolean readOnly) { return TestingTransactionHandle.create(); }
public static TestingTransactionHandle create() { return new TestingTransactionHandle(UUID.randomUUID()); }
@Test public void testScheduleRemote() { Set<Split> splits = new HashSet<>(); splits.add(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())); Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments(); assertEquals(assignments.size(), 1); }
public static TestingTransactionHandle create(String connectorId) { return new TestingTransactionHandle(connectorId, UUID.randomUUID()); }
@Test public void testBasicAssignment() { // One split for each node Set<Split> splits = new HashSet<>(); for (int i = 0; i < 3; i++) { splits.add(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())); } Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments(); assertEquals(assignments.entries().size(), 3); for (Node node : nodeManager.getActiveConnectorNodes(CONNECTOR_ID)) { assertTrue(assignments.keySet().contains(node)); } }
@Test public void testScheduleLocal() { Split split = new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitLocal()); Set<Split> splits = ImmutableSet.of(split); Map.Entry<Node, Split> assignment = Iterables.getOnlyElement(nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments().entries()); assertEquals(assignment.getKey().getHostAndPort(), split.getAddresses().get(0)); assertEquals(assignment.getValue(), split); }
public SourceOperator newTableScanOperator(DriverContext driverContext) { ConnectorPageSource pageSource = newPageSource(); SourceOperatorFactory sourceOperatorFactory = new TableScanOperatorFactory( 0, new PlanNodeId("0"), (session, split, columnHandles) -> pageSource, columns.stream().map(columnHandle -> (ColumnHandle) columnHandle).collect(toList())); SourceOperator operator = sourceOperatorFactory.createOperator(driverContext); operator.addSplit(new Split(new ConnectorId("test"), TestingTransactionHandle.create(), TestingSplit.createLocalSplit())); return operator; }
public SourceOperator newScanFilterAndProjectOperator(DriverContext driverContext) { ConnectorPageSource pageSource = newPageSource(); ImmutableList.Builder<RowExpression> projectionsBuilder = ImmutableList.builder(); for (int i = 0; i < types.size(); i++) { projectionsBuilder.add(field(i, types.get(i))); } Supplier<CursorProcessor> cursorProcessor = EXPRESSION_COMPILER.compileCursorProcessor(Optional.empty(), projectionsBuilder.build(), "key"); Supplier<PageProcessor> pageProcessor = EXPRESSION_COMPILER.compilePageProcessor(Optional.empty(), projectionsBuilder.build()); SourceOperatorFactory sourceOperatorFactory = new ScanFilterAndProjectOperatorFactory( 0, new PlanNodeId("test"), new PlanNodeId("0"), (session, split, columnHandles) -> pageSource, cursorProcessor, pageProcessor, columns.stream().map(columnHandle -> (ColumnHandle) columnHandle).collect(toList()), types, new DataSize(0, BYTE), 0); SourceOperator operator = sourceOperatorFactory.createOperator(driverContext); operator.addSplit(new Split(new ConnectorId("test"), TestingTransactionHandle.create(), TestingSplit.createLocalSplit())); return operator; }
@Test public void testMaxSplitsPerNodePerTask() TestingTransactionHandle transactionHandle = TestingTransactionHandle.create();
@Test public void testTaskCompletion() throws Exception { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0); TaskId taskId = new TaskId("test", 1, 1); RemoteTask remoteTask = remoteTaskFactory.createTableScanTask( taskId, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId)); nodeTaskMap.addTask(chosenNode, remoteTask); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 1); remoteTask.abort(); MILLISECONDS.sleep(100); // Sleep until cache expires assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); remoteTask.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); }
private static Split createRecordSetSplit() { return new Split(new ConnectorId("test"), TestingTransactionHandle.create(), new TestSplit(true)); }
private static Split newMockSplit() { return new Split(new ConnectorId("test"), TestingTransactionHandle.create(), new MockSplit()); }
private static Split createNormalSplit() { return new Split(new ConnectorId("test"), TestingTransactionHandle.create(), new TestSplit(false)); }
@Test public void testMaxSplitsPerNode() { TestingTransactionHandle transactionHandle = TestingTransactionHandle.create(); Node newNode = new PrestoNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false); nodeManager.addNode(CONNECTOR_ID, newNode); ImmutableList.Builder<Split> initialSplits = ImmutableList.builder(); for (int i = 0; i < 10; i++) { initialSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote())); } MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); // Max out number of splits on node TaskId taskId1 = new TaskId("test", 1, 1); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId1)); nodeTaskMap.addTask(newNode, remoteTask1); TaskId taskId2 = new TaskId("test", 1, 2); RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask(taskId2, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId2)); nodeTaskMap.addTask(newNode, remoteTask2); Set<Split> splits = new HashSet<>(); for (int i = 0; i < 5; i++) { splits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote())); } Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments(); // no split should be assigned to the newNode, as it already has maxNodeSplits assigned to it assertFalse(assignments.keySet().contains(newNode)); remoteTask1.abort(); remoteTask2.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(newNode), 0); }
@Setup public void setup() TestingTransactionHandle transactionHandle = TestingTransactionHandle.create();
@Test public void testSplitCount() { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0); TaskId taskId1 = new TaskId("test", 1, 1); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, chosenNode, ImmutableList.of( new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote()), new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId1)); TaskId taskId2 = new TaskId("test", 1, 2); RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask( taskId2, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId2)); nodeTaskMap.addTask(chosenNode, remoteTask1); nodeTaskMap.addTask(chosenNode, remoteTask2); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 3); remoteTask1.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 1); remoteTask2.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); }
ImmutableMap.of(tableScanNodeId, new ConnectorAwareSplitSource(CONNECTOR_ID, TestingTransactionHandle.create(), splitSource)), ImmutableList.of());
throws Exception TestingTransactionHandle transactionHandle = TestingTransactionHandle.create(); NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); InMemoryNodeManager nodeManager = new InMemoryNodeManager();
@BeforeClass public void setup() { metadata = getQueryRunner().getMetadata(); builder = new PlanBuilder(new PlanNodeIdAllocator(), metadata); ConnectorId connectorId = getCurrentConnectorId(); TableHandle nationTableHandle = new TableHandle( connectorId, new TpchTableHandle("nation", 1.0)); TableLayoutHandle nationTableLayoutHandle = new TableLayoutHandle(connectorId, TestingTransactionHandle.create(), new TpchTableLayoutHandle((TpchTableHandle) nationTableHandle.getConnectorHandle(), TupleDomain.all())); TpchColumnHandle nationkeyColumnHandle = new TpchColumnHandle("nationkey", BIGINT); symbol = new Symbol("nationkey"); tableScanNode = builder.tableScan(nationTableHandle, ImmutableList.of(symbol), ImmutableMap.of(symbol, nationkeyColumnHandle), Optional.of(nationTableLayoutHandle)); }
@BeforeClass public void setup() { metadata = getQueryRunner().getMetadata(); sqlParser = getQueryRunner().getSqlParser(); ConnectorId connectorId = getCurrentConnectorId(); nationTableHandle = new TableHandle( connectorId, new TpchTableHandle("nation", 1.0)); nationTableLayoutHandle = new TableLayoutHandle(connectorId, TestingTransactionHandle.create(), new TpchTableLayoutHandle((TpchTableHandle) nationTableHandle.getConnectorHandle(), TupleDomain.all())); }