Examples with ReadSession com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession used on opensource projects

Example 1 with ReadSession

use of com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession in project presto by prestodb.

the class ReadSessionCreator method create.

public Storage.ReadSession create(TableId table, ImmutableList<String> selectedFields, Optional<String> filter, int parallelism) {
    TableInfo tableDetails = bigQueryClient.getTable(table);
    TableInfo actualTable = getActualTable(tableDetails, selectedFields, new String[] {});
    try (BigQueryStorageClient bigQueryStorageClient = bigQueryStorageClientFactory.createBigQueryStorageClient()) {
        ReadOptions.TableReadOptions.Builder readOptions = ReadOptions.TableReadOptions.newBuilder().addAllSelectedFields(selectedFields);
        filter.ifPresent(readOptions::setRowRestriction);
        TableReferenceProto.TableReference tableReference = toTableReference(actualTable.getTableId());
        Storage.ReadSession readSession = bigQueryStorageClient.createReadSession(Storage.CreateReadSessionRequest.newBuilder().setParent("projects/" + bigQueryClient.getProjectId()).setFormat(Storage.DataFormat.AVRO).setRequestedStreams(parallelism).setReadOptions(readOptions).setTableReference(tableReference).setShardingStrategy(Storage.ShardingStrategy.BALANCED).build());
        return readSession;
    }
}

Example 2 with ReadSession

use of com.google.cloud.bigquery.storage.v1beta1.Storage.ReadSession in project presto by prestodb.

the class BigQuerySplitManager method readFromBigQuery.

private ImmutableList<BigQuerySplit> readFromBigQuery(TableId tableId, Optional<List<ColumnHandle>> projectedColumns, int actualParallelism, Optional<String> filter) {
    List<ColumnHandle> columns = projectedColumns.orElse(ImmutableList.of());
    ImmutableList<String> projectedColumnsNames = columns.stream().map(column -> ((BigQueryColumnHandle) column).getName()).collect(toImmutableList());
    ReadSession readSession = new ReadSessionCreator(readSessionCreatorConfig, bigQueryClient, bigQueryStorageClientFactory).create(tableId, projectedColumnsNames, filter, actualParallelism);
    return readSession.getStreamsList().stream().map(stream -> BigQuerySplit.forStream(stream.getName(), readSession.getAvroSchema().getSchema(), columns)).collect(toImmutableList());
}