Examples with SourceSplitOptions com.google.api.services.dataflow.model.SourceSplitOptions used on opensource projects

Example 1 with SourceSplitOptions

use of com.google.api.services.dataflow.model.SourceSplitOptions in project beam by apache.

the class WorkerCustomSources method performSplitWithApiLimit.

/**
 * A helper method like {@link #performSplit(SourceSplitRequest, PipelineOptions)} but that allows
 * overriding the API size limit for testing.
 */
static SourceOperationResponse performSplitWithApiLimit(SourceSplitRequest request, PipelineOptions options, int numBundlesLimit, long apiByteLimit) throws Exception {
    // Compute the desired bundle size given by the service, or default if none was provided.
    long desiredBundleSizeBytes = DEFAULT_DESIRED_BUNDLE_SIZE_BYTES;
    SourceSplitOptions splitOptions = request.getOptions();
    if (splitOptions != null && splitOptions.getDesiredBundleSizeBytes() != null) {
        desiredBundleSizeBytes = splitOptions.getDesiredBundleSizeBytes();
    }
    Source<?> anySource = deserializeFromCloudSource(request.getSource().getSpec());
    checkArgument(anySource instanceof BoundedSource, "Cannot split a non-Bounded source: %s", anySource);
    return performSplitTyped(options, (BoundedSource<?>) anySource, desiredBundleSizeBytes, numBundlesLimit, apiByteLimit);
}

Example 2 with SourceSplitOptions

use of com.google.api.services.dataflow.model.SourceSplitOptions in project beam by apache.

the class WorkerCustomSourcesTest method performSplit.

static SourceSplitResponse performSplit(com.google.api.services.dataflow.model.Source source, PipelineOptions options, @Nullable Long desiredBundleSizeBytes, @Nullable Integer numBundlesLimitForTest, @Nullable Long apiByteLimitForTest) throws Exception {
    SourceSplitRequest splitRequest = new SourceSplitRequest();
    splitRequest.setSource(source);
    if (desiredBundleSizeBytes != null) {
        splitRequest.setOptions(new SourceSplitOptions().setDesiredBundleSizeBytes(desiredBundleSizeBytes));
    }
    SourceOperationResponse response = WorkerCustomSources.performSplitWithApiLimit(splitRequest, options, MoreObjects.firstNonNull(numBundlesLimitForTest, WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT), MoreObjects.firstNonNull(apiByteLimitForTest, WorkerCustomSources.DATAFLOW_SPLIT_RESPONSE_API_SIZE_LIMIT));
    return response.getSplit();
}