Examples with TableReference com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference used on opensource projects

Example 11 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class ITBigQueryStorageLongRunningTest method testLongRunningReadSession.

@Test
public void testLongRunningReadSession() throws InterruptedException, ExecutionException {
    // This test reads a larger table with the goal of doing a simple validation of timeout settings
    // for a longer running session.
    TableReference tableReference = TableReference.newBuilder().setProjectId("bigquery-public-data").setDatasetId("samples").setTableId("wikipedia").build();
    ReadSession session = client.createReadSession(/* tableReference = */
    tableReference, /* parent = */
    parentProjectId, /* requestedStreams = */
    5);
    assertEquals(String.format("Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s", TextFormat.shortDebugString(tableReference), session.toString()), 5, session.getStreamsCount());
    List<Callable<Long>> tasks = new ArrayList<>(session.getStreamsCount());
    for (final Stream stream : session.getStreamsList()) {
        tasks.add(new Callable<Long>() {

            @Override
            public Long call() throws Exception {
                return readAllRowsFromStream(stream);
            }
        });
    }
    ExecutorService executor = Executors.newFixedThreadPool(tasks.size());
    List<Future<Long>> results = executor.invokeAll(tasks);
    long rowCount = 0;
    for (Future<Long> result : results) {
        rowCount += result.get();
    }
    assertEquals(313_797_035, rowCount);
}

Example 12 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference 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 13 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project urban-eureka by errir503.

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 14 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class BigQueryStorageClientTest method createReadSessionExceptionTest.

@Test
@SuppressWarnings("all")
public void createReadSessionExceptionTest() throws Exception {
    StatusRuntimeException exception = new StatusRuntimeException(Status.INVALID_ARGUMENT);
    mockBigQueryStorage.addException(exception);
    try {
        TableReference tableReference = TableReference.newBuilder().build();
        String parent = "parent-995424086";
        int requestedStreams = 1017221410;
        client.createReadSession(tableReference, parent, requestedStreams);
        Assert.fail("No exception raised");
    } catch (InvalidArgumentException e) {
    // Expected exception
    }
}

Example 15 with TableReference

use of com.google.cloud.bigquery.storage.v1beta1.TableReferenceProto.TableReference in project java-bigquerystorage by googleapis.

the class ITBigQueryStorageTest method testFilter.

@Test
public void testFilter() throws IOException {
    TableReference tableReference = TableReference.newBuilder().setProjectId("bigquery-public-data").setDatasetId("samples").setTableId("shakespeare").build();
    TableReadOptions options = TableReadOptions.newBuilder().setRowRestriction("word_count > 100").build();
    CreateReadSessionRequest request = CreateReadSessionRequest.newBuilder().setParent(parentProjectId).setRequestedStreams(1).setTableReference(tableReference).setReadOptions(options).setFormat(DataFormat.AVRO).build();
    ReadSession session = client.createReadSession(request);
    assertEquals(String.format("Did not receive expected number of streams for table reference '%s' CreateReadSession response:%n%s", TextFormat.shortDebugString(tableReference), session.toString()), 1, session.getStreamsCount());
    StreamPosition readPosition = StreamPosition.newBuilder().setStream(session.getStreams(0)).build();
    ReadRowsRequest readRowsRequest = ReadRowsRequest.newBuilder().setReadPosition(readPosition).build();
    SimpleRowReader reader = new SimpleRowReader(new Schema.Parser().parse(session.getAvroSchema().getSchema()));
    long rowCount = 0;
    ServerStream<ReadRowsResponse> stream = client.readRowsCallable().call(readRowsRequest);
    for (ReadRowsResponse response : stream) {
        rowCount += response.getRowCount();
        reader.processRows(response.getAvroRows(), new SimpleRowReader.AvroRowConsumer() {

            @Override
            public void accept(GenericData.Record record) {
                Long wordCount = (Long) record.get("word_count");
                assertWithMessage("Row not matching expectations: %s", record.toString()).that(wordCount).isGreaterThan(100L);
            }
        });
    }
    assertEquals(1_333, rowCount);
}