io.prestosql.spi.connector.ConnectorSplit Java Examples

The following examples show how to use io.prestosql.spi.connector.ConnectorSplit. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: JmxSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableHandle table, SplitSchedulingStrategy splitSchedulingStrategy)
{
    JmxTableHandle tableHandle = (JmxTableHandle) table;

    //TODO is there a better way to get the node column?
    Optional<JmxColumnHandle> nodeColumnHandle = tableHandle.getColumnHandles().stream()
            .filter(jmxColumnHandle -> jmxColumnHandle.getColumnName().equals(NODE_COLUMN_NAME))
            .findFirst();
    checkState(nodeColumnHandle.isPresent(), "Failed to find %s column", NODE_COLUMN_NAME);

    TupleDomain<ColumnHandle> nodeFilter = tableHandle.getNodeFilter();

    List<ConnectorSplit> splits = nodeManager.getAllNodes().stream()
            .filter(node -> {
                NullableValue value = NullableValue.of(createUnboundedVarcharType(), utf8Slice(node.getNodeIdentifier()));
                return nodeFilter.overlaps(fromFixedValues(ImmutableMap.of(nodeColumnHandle.get(), value)));
            })
            .map(node -> new JmxSplit(ImmutableList.of(node.getHostAndPort())))
            .collect(toList());

    return new FixedSplitSource(splits);
}
 
Example #2
Source File: ExampleSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorTableHandle connectorTableHandle,
        SplitSchedulingStrategy splitSchedulingStrategy)
{
    ExampleTableHandle tableHandle = (ExampleTableHandle) connectorTableHandle;
    ExampleTable table = exampleClient.getTable(tableHandle.getSchemaName(), tableHandle.getTableName());

    // this can happen if table is removed during a query
    if (table == null) {
        throw new TableNotFoundException(tableHandle.toSchemaTableName());
    }

    List<ConnectorSplit> splits = new ArrayList<>();
    for (URI uri : table.getSources()) {
        splits.add(new ExampleSplit(uri));
    }
    Collections.shuffle(splits);

    return new FixedSplitSource(splits);
}
 
Example #3
Source File: TpchSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableHandle tableHandle, SplitSchedulingStrategy splitSchedulingStrategy)
{
    Set<Node> nodes = nodeManager.getRequiredWorkerNodes();

    int totalParts = nodes.size() * splitsPerNode;
    int partNumber = 0;

    // Split the data using split and skew by the number of nodes available.
    ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
    for (Node node : nodes) {
        for (int i = 0; i < splitsPerNode; i++) {
            splits.add(new TpchSplit(partNumber, totalParts, ImmutableList.of(node.getHostAndPort())));
            partNumber++;
        }
    }
    return new FixedSplitSource(splits.build());
}
 
Example #4
Source File: RaptorSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
private Supplier<ConnectorSplitBatch> batchSupplier(int maxSize)
{
    return () -> {
        ImmutableList.Builder<ConnectorSplit> list = ImmutableList.builder();
        for (int i = 0; i < maxSize; i++) {
            if (Thread.currentThread().isInterrupted()) {
                throw new RuntimeException("Split batch fetch was interrupted");
            }
            if (!iterator.hasNext()) {
                break;
            }
            list.add(createSplit(iterator.next()));
        }
        return new ConnectorSplitBatch(list.build(), isFinished());
    };
}
 
Example #5
Source File: PinotSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
protected void generateSegmentSplits(
        List<ConnectorSplit> splits,
        Map<String, Map<String, List<String>>> routingTable,
        String tableName,
        String tableNameSuffix,
        ConnectorSession session,
        Optional<String> timePredicate)
{
    final String finalTableName = tableName + tableNameSuffix;
    int segmentsPerSplitConfigured = PinotSessionProperties.getSegmentsPerSplit(session);
    for (String routingTableName : routingTable.keySet()) {
        if (!routingTableName.equalsIgnoreCase(finalTableName)) {
            continue;
        }

        Map<String, List<String>> hostToSegmentsMap = routingTable.get(routingTableName);
        hostToSegmentsMap.forEach((host, segments) -> {
            int numSegmentsInThisSplit = Math.min(segments.size(), segmentsPerSplitConfigured);
            // segments is already shuffled
            Iterables.partition(segments, numSegmentsInThisSplit).forEach(
                    segmentsForThisSplit -> splits.add(
                            createSegmentSplit(tableNameSuffix, segmentsForThisSplit, host, timePredicate)));
        });
    }
}
 
Example #6
Source File: MongoPageSourceProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle table,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    MongoTableHandle tableHandle = (MongoTableHandle) table;

    ImmutableList.Builder<MongoColumnHandle> handles = ImmutableList.builder();
    for (ColumnHandle handle : requireNonNull(columns, "columns is null")) {
        handles.add((MongoColumnHandle) handle);
    }

    return new MongoPageSource(mongoSession, tableHandle, handles.build());
}
 
Example #7
Source File: RaptorSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
private ConnectorSplit createBucketSplit(int bucketNumber, Set<ShardNodes> shards)
{
    // Bucket splits contain all the shards for the bucket
    // and run on the node assigned to the bucket.

    String nodeId = bucketToNode.get().get(bucketNumber);
    Node node = nodesById.get(nodeId);
    if (node == null) {
        throw new PrestoException(NO_NODES_AVAILABLE, "Node for bucket is offline: " + nodeId);
    }

    Set<UUID> shardUuids = shards.stream()
            .map(ShardNodes::getShardUuid)
            .collect(toSet());
    HostAddress address = node.getHostAndPort();

    return new RaptorSplit(shardUuids, bucketNumber, address, transactionId);
}
 
Example #8
Source File: TestPrometheusSplit.java    From presto with Apache License 2.0 6 votes vote down vote up
@Test
public void testQueryDividedIntoSplitsLastSplitHasRightTime()
        throws URISyntaxException
{
    LocalDateTime now = LocalDateTime.of(2019, 10, 2, 7, 26, 56, 00);
    PrometheusConnectorConfig config = getCommonConfig(prometheusHttpServer.resolve("/prometheus-data/prometheus-metrics.json"), now);
    PrometheusClient client = new PrometheusClient(config, METRIC_CODEC, TYPE_MANAGER);
    PrometheusTable table = client.getTable("default", "up");
    PrometheusSplitManager splitManager = new PrometheusSplitManager(client, config);
    ConnectorSplitSource splitsMaybe = splitManager.getSplits(
            null,
            null,
            (ConnectorTableHandle) new PrometheusTableHandle("default", table.getName()),
            null);
    List<ConnectorSplit> splits = splitsMaybe.getNextBatch(NOT_PARTITIONED, NUMBER_MORE_THAN_EXPECTED_NUMBER_SPLITS).getNow(null).getSplits();
    int lastSplitIndex = splits.size() - 1;
    PrometheusSplit lastSplit = (PrometheusSplit) splits.get(lastSplitIndex);
    String queryInSplit = lastSplit.getUri().getQuery();
    String timeShouldBe = decimalSecondString(now.atZone(ZoneId.systemDefault()).toInstant().toEpochMilli());
    URI uriAsFormed = new URI("http://doesnotmatter:9090/api/v1/query?query=up[" +
            getQueryChunkSizeDurationAsPrometheusCompatibleDurationString(config) + "]" +
            "&time=" + timeShouldBe);
    assertEquals(queryInSplit, uriAsFormed.getQuery());
}
 
Example #9
Source File: SheetsSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorTableHandle connectorTableHandle,
        SplitSchedulingStrategy splitSchedulingStrategy)
{
    SheetsTableHandle tableHandle = (SheetsTableHandle) connectorTableHandle;
    Optional<SheetsTable> table = sheetsClient.getTable(tableHandle.getTableName());

    // this can happen if table is removed during a query
    if (table.isEmpty()) {
        throw new TableNotFoundException(tableHandle.toSchemaTableName());
    }

    List<ConnectorSplit> splits = new ArrayList<>();
    splits.add(new SheetsSplit(tableHandle.getSchemaName(), tableHandle.getTableName(), table.get().getValues()));
    Collections.shuffle(splits);
    return new FixedSplitSource(splits);
}
 
Example #10
Source File: BlackHolePageSourceProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transactionHandle,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle tableHandle,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    BlackHoleTableHandle table = (BlackHoleTableHandle) tableHandle;

    ImmutableList.Builder<Type> builder = ImmutableList.builder();

    for (ColumnHandle column : columns) {
        builder.add(((BlackHoleColumnHandle) column).getColumnType());
    }
    List<Type> types = builder.build();

    Page page = generateZeroPage(types, table.getRowsPerPage(), table.getFieldsLength());
    return new BlackHolePageSource(page, table.getPagesPerSplit(), executorService, table.getPageProcessingDelay());
}
 
Example #11
Source File: TestBeginQuery.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSourceProvider getPageSourceProvider()
{
    return new ConnectorPageSourceProvider()
    {
        @Override
        public ConnectorPageSource createPageSource(
                ConnectorTransactionHandle transaction,
                ConnectorSession session,
                ConnectorSplit split,
                ConnectorTableHandle table,
                List<ColumnHandle> columns,
                TupleDomain<ColumnHandle> dynamicFilter)
        {
            return new FixedPageSource(ImmutableList.of());
        }
    };
}
 
Example #12
Source File: TestJmxSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@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))));
        JmxTableHandle tableHandle = new JmxTableHandle(new SchemaTableName("schema", "tableName"), ImmutableList.of("objectName"), ImmutableList.of(columnHandle), true, nodeTupleDomain);

        ConnectorSplitSource splitSource = splitManager.getSplits(JmxTransactionHandle.INSTANCE, SESSION, tableHandle, UNGROUPED_SCHEDULING);
        List<ConnectorSplit> allSplits = getAllSplits(splitSource);

        assertEquals(allSplits.size(), 1);
        assertEquals(allSplits.get(0).getAddresses().size(), 1);
        assertEquals(allSplits.get(0).getAddresses().get(0).getHostText(), nodeIdentifier);
    }
}
 
Example #13
Source File: BigQueryPageSourceProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle table,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    log.debug("createPageSource(transaction=%s, session=%s, split=%s, table=%s, columns=%s)", transaction, session, split, table, columns);
    BigQuerySplit bigQuerySplit = (BigQuerySplit) split;
    if (bigQuerySplit.representsEmptyProjection()) {
        return new BigQueryEmptyProjectionPageSource(bigQuerySplit.getEmptyRowsToGenerate());
    }

    // not empty projection
    List<BigQueryColumnHandle> bigQueryColumnHandles = columns.stream()
            .map(BigQueryColumnHandle.class::cast)
            .collect(toImmutableList());

    return new BigQueryResultPageSource(bigQueryStorageClientFactory, maxReadRowsRetries, bigQuerySplit, bigQueryColumnHandles);
}
 
Example #14
Source File: TestSourcePartitionedScheduler.java    From presto with Apache License 2.0 6 votes vote down vote up
private synchronized List<ConnectorSplit> getBatch(int maxSize)
{
    // take up to maxSize elements from the queue
    List<ConnectorSplit> elements = new ArrayList<>(maxSize);
    queue.drainTo(elements, maxSize);

    // if the queue is empty and the current future is finished, create a new one so
    // a new readers can be notified when the queue has elements to read
    if (queue.isEmpty() && !closed) {
        if (notEmptyFuture.isDone()) {
            notEmptyFuture = new CompletableFuture<>();
        }
    }

    return ImmutableList.copyOf(elements);
}
 
Example #15
Source File: TestingPageSourceProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle table,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    requireNonNull(columns, "columns is null");

    ImmutableList<Block> blocks = columns.stream()
            .map(column -> new ByteArrayBlock(1, Optional.of(new boolean[] {true}), new byte[1]))
            .collect(toImmutableList());

    return new FixedPageSource(ImmutableList.of(new Page(blocks.toArray(new Block[blocks.size()]))));
}
 
Example #16
Source File: CassandraRecordSetProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    CassandraSplit cassandraSplit = (CassandraSplit) split;
    CassandraTableHandle cassandraTable = (CassandraTableHandle) table;

    List<CassandraColumnHandle> cassandraColumns = columns.stream()
            .map(column -> (CassandraColumnHandle) column)
            .collect(toList());

    String selectCql = CassandraCqlUtils.selectFrom(cassandraTable, cassandraColumns).getQueryString();
    StringBuilder sb = new StringBuilder(selectCql);
    if (sb.charAt(sb.length() - 1) == ';') {
        sb.setLength(sb.length() - 1);
    }
    sb.append(cassandraSplit.getWhereClause());
    String cql = sb.toString();
    log.debug("Creating record set: %s", cql);

    return new CassandraRecordSet(cassandraSession, cql, cassandraColumns);
}
 
Example #17
Source File: JdbcRecordSetProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    JdbcSplit jdbcSplit = (JdbcSplit) split;
    JdbcTableHandle jdbcTable = (JdbcTableHandle) table;

    // In the current API, the columns (and order) needed by the engine are provided via an argument to this method. Make sure that
    // any columns that were recorded in the table handle match the requested set.
    // If no columns are recorded, it means that applyProjection never got called (e.g., in the case all columns are being used) and all
    // table columns should be returned. TODO: this is something that should be addressed once the getRecordSet API is revamped
    jdbcTable.getColumns()
            .ifPresent(tableColumns -> verify(columns.equals(tableColumns)));

    ImmutableList.Builder<JdbcColumnHandle> handles = ImmutableList.builder();
    for (ColumnHandle handle : columns) {
        handles.add((JdbcColumnHandle) handle);
    }

    return new JdbcRecordSet(jdbcClient, session, jdbcSplit, jdbcTable, handles.build());
}
 
Example #18
Source File: NodePartitioningManager.java    From presto with Apache License 2.0 6 votes vote down vote up
private ToIntFunction<Split> getSplitToBucket(Session session, PartitioningHandle partitioningHandle)
{
    ConnectorNodePartitioningProvider partitioningProvider = partitioningProviders.get(partitioningHandle.getConnectorId().get());
    checkArgument(partitioningProvider != null, "No partitioning provider for connector %s", partitioningHandle.getConnectorId().get());

    ToIntFunction<ConnectorSplit> splitBucketFunction = partitioningProvider.getSplitBucketFunction(
            partitioningHandle.getTransactionHandle().orElse(null),
            session.toConnectorSession(),
            partitioningHandle.getConnectorHandle());
    checkArgument(splitBucketFunction != null, "No partitioning %s", partitioningHandle);

    return split -> {
        int bucket;
        if (split.getConnectorSplit() instanceof EmptySplit) {
            bucket = split.getLifespan().isTaskWide() ? 0 : split.getLifespan().getId();
        }
        else {
            bucket = splitBucketFunction.applyAsInt(split.getConnectorSplit());
        }
        if (!split.getLifespan().isTaskWide()) {
            checkArgument(split.getLifespan().getId() == bucket);
        }
        return bucket;
    };
}
 
Example #19
Source File: TpcdsSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorTableLayoutHandle layout, SplitSchedulingStrategy splitSchedulingStrategy)
{
    Set<Node> nodes = nodeManager.getRequiredWorkerNodes();
    checkState(!nodes.isEmpty(), "No TPCDS nodes available");

    int totalParts = nodes.size() * splitsPerNode;
    int partNumber = 0;

    // Split the data using split and skew by the number of nodes available.
    ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
    for (Node node : nodes) {
        for (int i = 0; i < splitsPerNode; i++) {
            splits.add(new TpcdsSplit(partNumber, totalParts, ImmutableList.of(node.getHostAndPort()), noSexism));
            partNumber++;
        }
    }
    return new FixedSplitSource(splits.build());
}
 
Example #20
Source File: AtopPageSourceProvider.java    From presto with Apache License 2.0 6 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle table,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    AtopTableHandle tableHandle = (AtopTableHandle) table;
    AtopSplit atopSplit = (AtopSplit) split;

    ImmutableList.Builder<Type> types = ImmutableList.builder();
    ImmutableList.Builder<AtopColumn> atopColumns = ImmutableList.builder();

    for (ColumnHandle column : columns) {
        AtopColumnHandle atopColumnHandle = (AtopColumnHandle) column;
        AtopColumn atopColumn = tableHandle.getTable().getColumn(atopColumnHandle.getName());
        atopColumns.add(atopColumn);
        types.add(typeManager.getType(atopColumn.getType()));
    }

    ZonedDateTime date = atopSplit.getDate();
    checkArgument(date.equals(date.withHour(0).withMinute(0).withSecond(0).withNano(0)), "Expected date to be at beginning of day");
    return new AtopPageSource(readerPermits, atopFactory, session, utf8Slice(atopSplit.getHost().getHostText()), tableHandle.getTable(), date, atopColumns.build(), types.build());
}
 
Example #21
Source File: CassandraSplitManager.java    From presto with Apache License 2.0 6 votes vote down vote up
private List<ConnectorSplit> getSplitsByTokenRange(CassandraTable table, String partitionId, Optional<Long> sessionSplitsPerNode)
{
    String schema = table.getTableHandle().getSchemaName();
    String tableName = table.getTableHandle().getTableName();
    String tokenExpression = table.getTokenExpression();

    ImmutableList.Builder<ConnectorSplit> builder = ImmutableList.builder();
    List<CassandraTokenSplitManager.TokenSplit> tokenSplits = tokenSplitMgr.getSplits(schema, tableName, sessionSplitsPerNode);
    for (CassandraTokenSplitManager.TokenSplit tokenSplit : tokenSplits) {
        String condition = buildTokenCondition(tokenExpression, tokenSplit.getStartToken(), tokenSplit.getEndToken());
        List<HostAddress> addresses = new HostAddressFactory().hostAddressNamesToHostAddressList(tokenSplit.getHosts());
        CassandraSplit split = new CassandraSplit(partitionId, condition, addresses);
        builder.add(split);
    }

    return builder.build();
}
 
Example #22
Source File: ClassLoaderSafeConnectorRecordSetProvider.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
        return new ClassLoaderSafeRecordSet(delegate.getRecordSet(transaction, session, split, table, columns), classLoader);
    }
}
 
Example #23
Source File: ConnectorAwareSplitSource.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public ListenableFuture<SplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, Lifespan lifespan, int maxSize)
{
    ListenableFuture<ConnectorSplitBatch> nextBatch = toListenableFuture(source.getNextBatch(partitionHandle, maxSize));
    return Futures.transform(nextBatch, splitBatch -> {
        ImmutableList.Builder<Split> result = ImmutableList.builder();
        for (ConnectorSplit connectorSplit : splitBatch.getSplits()) {
            result.add(new Split(catalogName, connectorSplit, lifespan));
        }
        return new SplitBatch(result.build(), splitBatch.isNoMoreSplits());
    }, directExecutor());
}
 
Example #24
Source File: TpchRecordSetProvider.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<? extends ColumnHandle> columns)
{
    TpchSplit tpchSplit = (TpchSplit) split;
    TpchTableHandle tpchTable = (TpchTableHandle) table;

    return getRecordSet(
            TpchTable.getTable(tpchTable.getTableName()),
            columns,
            tpchTable.getScaleFactor(),
            tpchSplit.getPartNumber(),
            tpchSplit.getTotalParts(),
            tpchTable.getConstraint());
}
 
Example #25
Source File: Split.java    From presto with Apache License 2.0 5 votes vote down vote up
@JsonCreator
public Split(
        @JsonProperty("catalogName") CatalogName catalogName,
        @JsonProperty("connectorSplit") ConnectorSplit connectorSplit,
        @JsonProperty("lifespan") Lifespan lifespan)
{
    this.catalogName = requireNonNull(catalogName, "catalogName is null");
    this.connectorSplit = requireNonNull(connectorSplit, "connectorSplit is null");
    this.lifespan = requireNonNull(lifespan, "lifespan is null");
}
 
Example #26
Source File: LocalFileSplitManager.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorTableHandle table, SplitSchedulingStrategy splitSchedulingStrategy)
{
    List<ConnectorSplit> splits = nodeManager.getAllNodes().stream()
            .map(node -> new LocalFileSplit(node.getHostAndPort()))
            .collect(Collectors.toList());

    return new FixedSplitSource(splits);
}
 
Example #27
Source File: AbstractTestHiveFileSystem.java    From presto with Apache License 2.0 5 votes vote down vote up
protected MaterializedResult readTable(SchemaTableName tableName)
        throws IOException
{
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession();

        ConnectorTableHandle table = getTableHandle(metadata, tableName);
        List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, table).values());

        metadata.beginQuery(session);
        ConnectorSplitSource splitSource = splitManager.getSplits(transaction.getTransactionHandle(), session, table, UNGROUPED_SCHEDULING);

        List<Type> allTypes = getTypes(columnHandles);
        List<Type> dataTypes = getTypes(columnHandles.stream()
                .filter(columnHandle -> !((HiveColumnHandle) columnHandle).isHidden())
                .collect(toImmutableList()));
        MaterializedResult.Builder result = MaterializedResult.resultBuilder(session, dataTypes);

        List<ConnectorSplit> splits = getAllSplits(splitSource);
        for (ConnectorSplit split : splits) {
            try (ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, split, table, columnHandles, TupleDomain.all())) {
                MaterializedResult pageSourceResult = materializeSourceDataStream(session, pageSource, allTypes);
                for (MaterializedRow row : pageSourceResult.getMaterializedRows()) {
                    Object[] dataValues = IntStream.range(0, row.getFieldCount())
                            .filter(channel -> !((HiveColumnHandle) columnHandles.get(channel)).isHidden())
                            .mapToObj(row::getField)
                            .toArray();
                    result.row(dataValues);
                }
            }
        }

        metadata.cleanupQuery(session);
        return result.build();
    }
}
 
Example #28
Source File: ThriftSplitManager.java    From presto with Apache License 2.0 5 votes vote down vote up
/**
 * Returns a future with a list of splits.
 * This method is assumed to be called in a single-threaded way.
 * It can be called by multiple threads, but only if the previous call finished.
 */
@Override
public CompletableFuture<ConnectorSplitBatch> getNextBatch(ConnectorPartitionHandle partitionHandle, int maxSize)
{
    checkState(future.get() == null || future.get().isDone(), "previous batch not completed");
    checkState(hasMoreData.get(), "this method cannot be invoked when there's no more data");
    PrestoThriftId currentToken = nextToken.get();
    ListenableFuture<PrestoThriftSplitBatch> splitsFuture = client.getSplits(
            schemaTableName,
            new PrestoThriftNullableColumnSet(columnNames.orElse(null)),
            constraint,
            maxSize,
            new PrestoThriftNullableToken(currentToken));
    ListenableFuture<ConnectorSplitBatch> resultFuture = Futures.transform(
            splitsFuture,
            batch -> {
                requireNonNull(batch, "batch is null");
                List<ConnectorSplit> splits = batch.getSplits().stream()
                        .map(ThriftSplitSource::toConnectorSplit)
                        .collect(toImmutableList());
                checkState(nextToken.compareAndSet(currentToken, batch.getNextToken()));
                checkState(hasMoreData.compareAndSet(true, nextToken.get() != null));
                return new ConnectorSplitBatch(splits, isFinished());
            }, directExecutor());
    resultFuture = catchingThriftException(resultFuture);
    future.set(resultFuture);
    return toCompletableFuture(resultFuture);
}
 
Example #29
Source File: KuduPageSourceProvider.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public ConnectorPageSource createPageSource(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorSplit split,
        ConnectorTableHandle table,
        List<ColumnHandle> columns,
        TupleDomain<ColumnHandle> dynamicFilter)
{
    KuduRecordSet recordSet = (KuduRecordSet) recordSetProvider.getRecordSet(transaction, session, split, table, columns);
    if (columns.contains(KuduColumnHandle.ROW_ID_HANDLE)) {
        return new KuduUpdatablePageSource(recordSet);
    }
    return new RecordPageSource(recordSet);
}
 
Example #30
Source File: SystemSplitManager.java    From presto with Apache License 2.0 5 votes vote down vote up
@Override
public ConnectorSplitSource getSplits(
        ConnectorTransactionHandle transaction,
        ConnectorSession session,
        ConnectorTableHandle tableHandle,
        SplitSchedulingStrategy splitSchedulingStrategy)
{
    SystemTableHandle table = (SystemTableHandle) tableHandle;
    TupleDomain<ColumnHandle> constraint = table.getConstraint();

    SystemTable systemTable = tables.getSystemTable(session, table.getSchemaTableName())
            // table might disappear in the meantime
            .orElseThrow(() -> new TableNotFoundException(table.getSchemaTableName()));

    Distribution tableDistributionMode = systemTable.getDistribution();
    if (tableDistributionMode == SINGLE_COORDINATOR) {
        HostAddress address = nodeManager.getCurrentNode().getHostAndPort();
        ConnectorSplit split = new SystemSplit(address, constraint);
        return new FixedSplitSource(ImmutableList.of(split));
    }

    ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
    ImmutableSet.Builder<InternalNode> nodes = ImmutableSet.builder();
    if (tableDistributionMode == ALL_COORDINATORS) {
        nodes.addAll(nodeManager.getCoordinators());
    }
    else if (tableDistributionMode == ALL_NODES) {
        nodes.addAll(nodeManager.getNodes(ACTIVE));
    }
    Set<InternalNode> nodeSet = nodes.build();
    for (InternalNode node : nodeSet) {
        splits.add(new SystemSplit(node.getHostAndPort(), constraint));
    }
    return new FixedSplitSource(splits.build());
}