Examples with SplitSchedulingContext com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext used on opensource projects

Example 1 with SplitSchedulingContext

use of com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext in project presto by prestodb.

the class TestCassandraConnector method testGetRecords.

@Test
public void testGetRecords() {
    ConnectorTableHandle tableHandle = getTableHandle(table);
    ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(SESSION, tableHandle);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
    Map<String, Integer> columnIndex = indexColumns(columnHandles);
    ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE;
    List<ConnectorTableLayoutResult> layouts = metadata.getTableLayouts(SESSION, tableHandle, Constraint.alwaysTrue(), Optional.empty());
    ConnectorTableLayoutHandle layout = getOnlyElement(layouts).getTableLayout().getHandle();
    List<ConnectorSplit> splits = getAllSplits(splitManager.getSplits(transaction, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP)));
    long rowNumber = 0;
    for (ConnectorSplit split : splits) {
        CassandraSplit cassandraSplit = (CassandraSplit) split;
        long completedBytes = 0;
        try (RecordCursor cursor = recordSetProvider.getRecordSet(transaction, SESSION, cassandraSplit, columnHandles).cursor()) {
            while (cursor.advanceNextPosition()) {
                try {
                    assertReadFields(cursor, tableMetadata.getColumns());
                } catch (RuntimeException e) {
                    throw new RuntimeException("row " + rowNumber, e);
                }
                rowNumber++;
                String keyValue = cursor.getSlice(columnIndex.get("key")).toStringUtf8();
                assertTrue(keyValue.startsWith("key "));
                int rowId = Integer.parseInt(keyValue.substring(4));
                assertEquals(keyValue, String.format("key %d", rowId));
                assertEquals(Bytes.toHexString(cursor.getSlice(columnIndex.get("typebytes")).getBytes()), String.format("0x%08X", rowId));
                // VARINT is returned as a string
                assertEquals(cursor.getSlice(columnIndex.get("typeinteger")).toStringUtf8(), String.valueOf(rowId));
                assertEquals(cursor.getLong(columnIndex.get("typelong")), 1000 + rowId);
                assertEquals(cursor.getSlice(columnIndex.get("typeuuid")).toStringUtf8(), String.format("00000000-0000-0000-0000-%012d", rowId));
                assertEquals(cursor.getSlice(columnIndex.get("typetimestamp")).toStringUtf8(), Long.valueOf(DATE.getTime()).toString());
                long newCompletedBytes = cursor.getCompletedBytes();
                assertTrue(newCompletedBytes >= completedBytes);
                completedBytes = newCompletedBytes;
            }
        }
    }
    assertEquals(rowNumber, 9);
}

Example 2 with SplitSchedulingContext

use of com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext in project presto by prestodb.

the class TestJmxSplitManager method getRecordSet.

private RecordSet getRecordSet(SchemaTableName schemaTableName) throws Exception {
    JmxTableHandle tableHandle = metadata.getTableHandle(SESSION, schemaTableName);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
    ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, TupleDomain.all());
    ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP));
    List<ConnectorSplit> allSplits = getAllSplits(splitSource);
    assertEquals(allSplits.size(), nodes.size());
    ConnectorSplit split = allSplits.get(0);
    return recordSetProvider.getRecordSet(JmxTransactionHandle.INSTANCE, SESSION, split, columnHandles);
}

Example 3 with SplitSchedulingContext

use of com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext in project presto by prestodb.

the class SplitManager method getSplits.

public SplitSource getSplits(Session session, TableHandle table, SplitSchedulingStrategy splitSchedulingStrategy, WarningCollector warningCollector) {
    long startTime = System.nanoTime();
    ConnectorId connectorId = table.getConnectorId();
    ConnectorSplitManager splitManager = getConnectorSplitManager(connectorId);
    ConnectorSession connectorSession = session.toConnectorSession(connectorId);
    // Now we will fetch the layout handle if it's not presented in TableHandle.
    // In the future, ConnectorTableHandle will be used to fetch splits since it will contain layout information.
    ConnectorTableLayoutHandle layout;
    if (!table.getLayout().isPresent()) {
        TableLayoutResult result = metadata.getLayout(session, table, Constraint.alwaysTrue(), Optional.empty());
        layout = result.getLayout().getLayoutHandle();
    } else {
        layout = table.getLayout().get();
    }
    ConnectorSplitSource source = splitManager.getSplits(table.getTransaction(), connectorSession, layout, new SplitSchedulingContext(splitSchedulingStrategy, preferSplitHostAddresses, warningCollector));
    SplitSource splitSource = new ConnectorAwareSplitSource(connectorId, table.getTransaction(), source);
    if (minScheduleSplitBatchSize > 1) {
        splitSource = new BufferingSplitSource(splitSource, minScheduleSplitBatchSize);
    }
    return splitSource;
}

Example 4 with SplitSchedulingContext

use of com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext in project presto by prestodb.

the class TestJmxSplitManager method testNoPredicate.

@Test
public void testNoPredicate() throws Exception {
    ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, TupleDomain.all());
    ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP));
    List<ConnectorSplit> allSplits = getAllSplits(splitSource);
    assertEquals(allSplits.size(), nodes.size());
    Set<String> actualNodes = nodes.stream().map(Node::getNodeIdentifier).collect(toSet());
    Set<String> expectedNodes = new HashSet<>();
    for (ConnectorSplit split : allSplits) {
        List<HostAddress> addresses = ((JmxSplit) split).getAddresses();
        assertEquals(addresses.size(), 1);
        expectedNodes.add(addresses.get(0).getHostText());
    }
    assertEquals(actualNodes, expectedNodes);
}

Example 5 with SplitSchedulingContext

use of com.facebook.presto.spi.connector.ConnectorSplitManager.SplitSchedulingContext in project presto by prestodb.

the class TestJmxSplitManager method testPredicatePushdown.

@Test
public void testPredicatePushdown() throws Exception {
    for (Node node : nodes) {
        String nodeIdentifier = node.getNodeIdentifier();
        TupleDomain<ColumnHandle> nodeTupleDomain = TupleDomain.fromFixedValues(ImmutableMap.of(columnHandle, NullableValue.of(createUnboundedVarcharType(), utf8Slice(nodeIdentifier))));
        ConnectorTableLayoutHandle layout = new JmxTableLayoutHandle(tableHandle, nodeTupleDomain);
        ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, layout, new SplitSchedulingContext(UNGROUPED_SCHEDULING, false, WarningCollector.NOOP));
        List<ConnectorSplit> allSplits = getAllSplits(splitSource);
        assertEquals(allSplits.size(), 1);
        assertEquals(((JmxSplit) allSplits.get(0)).getAddresses().size(), 1);
        assertEquals(((JmxSplit) allSplits.get(0)).getAddresses().get(0).getHostText(), nodeIdentifier);
    }
}