Examples with ReadRowsResponse com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse used on opensource projects

Example 56 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse in project spark-bigquery-connector by GoogleCloudDataproc.

the class ArrowInputPartitionContext method createPartitionReaderContext.

public InputPartitionReaderContext<ColumnarBatch> createPartitionReaderContext() {
    BigQueryStorageReadRowsTracer tracer = tracerFactory.newReadRowsTracer(Joiner.on(",").join(streamNames));
    List<ReadRowsRequest.Builder> readRowsRequests = streamNames.stream().map(name -> ReadRowsRequest.newBuilder().setReadStream(name)).collect(Collectors.toList());
    ReadRowsHelper readRowsHelper = new ReadRowsHelper(bigQueryReadClientFactory, readRowsRequests, options);
    tracer.startStream();
    Iterator<ReadRowsResponse> readRowsResponses = readRowsHelper.readRows();
    return new ArrowColumnBatchPartitionReaderContext(readRowsResponses, serializedArrowSchema, readRowsHelper, selectedFields, tracer, userProvidedSchema.toJavaUtil(), options.numBackgroundThreads());
}

Example 57 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse in project spark-bigquery-connector by GoogleCloudDataproc.

the class BigQueryInputPartitionReaderContext method next.

@Override
public boolean next() throws IOException {
    while (!rows.hasNext()) {
        if (!readRowsResponses.hasNext()) {
            return false;
        }
        ReadRowsResponse readRowsResponse = readRowsResponses.next();
        rows = converter.convert(readRowsResponse);
    }
    currentRow = rows.next();
    return true;
}

Example 58 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse in project trino by trinodb.

the class BigQueryResultPageSource method getNextPage.

@Override
public Page getNextPage() {
    checkState(pageBuilder.isEmpty(), "PageBuilder is not empty at the beginning of a new page");
    ReadRowsResponse response = responses.next();
    Iterable<GenericRecord> records = parse(response);
    for (GenericRecord record : records) {
        pageBuilder.declarePosition();
        for (int column = 0; column < columnTypes.size(); column++) {
            BlockBuilder output = pageBuilder.getBlockBuilder(column);
            appendTo(columnTypes.get(column), record.get(columnNames.get(column)), output);
        }
    }
    Page page = pageBuilder.build();
    pageBuilder.reset();
    return page;
}

Example 59 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse in project trino by trinodb.

the class TestReadRowsHelper method testRetryOfSingleFailure.

@Test
public void testRetryOfSingleFailure() {
    BigQueryReadClient client = mock(BigQueryReadClient.class);
    MockResponsesBatch batch1 = new MockResponsesBatch();
    batch1.addResponse(ReadRowsResponse.newBuilder().setRowCount(10).build());
    batch1.addException(new StatusRuntimeException(Status.INTERNAL.withDescription("Received unexpected EOS on DATA frame from server.")));
    MockResponsesBatch batch2 = new MockResponsesBatch();
    batch2.addResponse(ReadRowsResponse.newBuilder().setRowCount(11).build());
    List<ReadRowsResponse> responses = ImmutableList.copyOf(new MockReadRowsHelper(client, "test", 3, ImmutableList.of(batch1, batch2)).readRows());
    assertThat(responses.size()).isEqualTo(2);
    assertThat(responses.stream().mapToLong(ReadRowsResponse::getRowCount).sum()).isEqualTo(21);
}

Example 60 with ReadRowsResponse

use of com.google.cloud.bigquery.storage.v1beta2.ReadRowsResponse in project java-bigquerystorage by googleapis.

the class EnhancedBigQueryReadStub method readRowsCallable.

public ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> readRowsCallable() {
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> innerCallable = GrpcRawCallableFactory.createServerStreamingCallable(GrpcCallSettings.<ReadRowsRequest, ReadRowsResponse>newBuilder().setMethodDescriptor(BigQueryReadGrpc.getReadRowsMethod()).setParamsExtractor(new RequestParamsExtractor<ReadRowsRequest>() {

        @Override
        public Map<String, String> extract(ReadRowsRequest request) {
            return ImmutableMap.of("read_stream", String.valueOf(request.getReadStream()));
        }
    }).build(), stubSettings.readRowsSettings().getRetryableCodes());
    ServerStreamingCallSettings<ReadRowsRequest, ReadRowsResponse> callSettings = stubSettings.readRowsSettings();
    StreamingRetryAlgorithm<Void> retryAlgorithm = new StreamingRetryAlgorithm<>(new ApiResultRetryAlgorithm<Void>(readRowsRetryAttemptListener), new ExponentialRetryAlgorithm(callSettings.getRetrySettings(), context.getClock()));
    ScheduledRetryingExecutor<Void> retryingExecutor = new ScheduledRetryingExecutor<>(retryAlgorithm, context.getExecutor());
    if (context.getStreamWatchdog() != null) {
        innerCallable = Callables.watched(innerCallable, callSettings, context);
    }
    ReadRowsRetryingCallable outerCallable = new ReadRowsRetryingCallable(context.getDefaultCallContext(), innerCallable, retryingExecutor, callSettings.getResumptionStrategy());
    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> traced = new TracedServerStreamingCallable<>(outerCallable, context.getTracerFactory(), SpanName.of(TRACING_OUTER_CLIENT_NAME, "ReadRows"));
    return traced.withDefaultCallContext(context.getDefaultCallContext());
}