Java Code Examples for org.apache.beam.sdk.io.BoundedSource#getEstimatedSizeBytes()

private static <T> List<BoundedSource<T>> split(
    BoundedSource<T> source, SamzaPipelineOptions pipelineOptions) throws Exception {
  final int numSplits = pipelineOptions.getMaxSourceParallelism();
  if (numSplits > 1) {
    final long estimatedSize = source.getEstimatedSizeBytes(pipelineOptions);
    // calculate the size of each split, rounded up to the ceiling.
    final long bundleSize = (estimatedSize + numSplits - 1) / numSplits;
    @SuppressWarnings("unchecked")
    final List<BoundedSource<T>> splits =
        (List<BoundedSource<T>>) source.split(bundleSize, pipelineOptions);
    // Need the empty check here because Samza doesn't handle empty partition well
    if (!splits.isEmpty()) {
      return splits;
    }
  }
  return Collections.singletonList(source);
}
 

/**
 * Version of {@link CustomSources#serializeToCloudSource(Source, PipelineOptions)} intended for
 * use on splits of {@link BoundedSource}.
 */
private static com.google.api.services.dataflow.model.Source serializeSplitToCloudSource(
    BoundedSource<?> source) throws Exception {
  com.google.api.services.dataflow.model.Source cloudSource =
      new com.google.api.services.dataflow.model.Source();
  cloudSource.setSpec(CloudObject.forClass(CustomSources.class));
  addString(
      cloudSource.getSpec(), SERIALIZED_SOURCE, encodeBase64String(serializeToByteArray(source)));
  SourceMetadata metadata = new SourceMetadata();
  // Size estimation is best effort so we continue even if it fails here.
  try {
    long estimatedSize = source.getEstimatedSizeBytes(PipelineOptionsFactory.create());
    if (estimatedSize >= 0) {
      metadata.setEstimatedSizeBytes(estimatedSize);
    } else {
      LOG.warn(
          "Ignoring negative estimated size {} produced by source {}", estimatedSize, source);
    }
  } catch (Exception e) {
    LOG.warn("Size estimation of the source failed: " + source, e);
  }
  cloudSource.setMetadata(metadata);
  return cloudSource;
}
 

@Override
public Collection<CommittedBundle<BoundedSourceShard<T>>> getInitialInputs(
    AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform,
    int targetParallelism)
    throws Exception {
  BoundedSource<T> source = ReadTranslation.boundedSourceFromTransform(transform);
  long estimatedBytes = source.getEstimatedSizeBytes(options);
  long bytesPerBundle = estimatedBytes / targetParallelism;
  List<? extends BoundedSource<T>> bundles = source.split(bytesPerBundle, options);
  ImmutableList.Builder<CommittedBundle<BoundedSourceShard<T>>> shards =
      ImmutableList.builder();
  for (BoundedSource<T> bundle : bundles) {
    CommittedBundle<BoundedSourceShard<T>> inputShard =
        evaluationContext
            .<BoundedSourceShard<T>>createRootBundle()
            .add(WindowedValue.valueInGlobalWindow(BoundedSourceShard.of(bundle)))
            .commit(BoundedWindow.TIMESTAMP_MAX_VALUE);
    shards.add(inputShard);
  }
  return shards.build();
}
 

/**
 * Constructor of BeamBoundedSourceVertex.
 *
 * @param source      BoundedSource to read from.
 * @param displayData data to display.
 */
public BeamBoundedSourceVertex(final BoundedSource<O> source, final DisplayData displayData) {
  this.source = source;
  this.displayData = displayData;
  try {
    this.estimatedSizeBytes = source.getEstimatedSizeBytes(null);
  } catch (Exception e) {
    throw new MetricException(e);
  }
}
 

private Future<BoundedSource<OutputT>> startDynamicSplitThread(
    BoundedSource<OutputT> source, BoundedReader<OutputT> reader) throws Exception {
  if (source.getEstimatedSizeBytes(options) > minimumDynamicSplitSize) {
    return produceSplitExecutor.submit(new GenerateSplitAtHalfwayPoint<>(reader));
  } else {
    SettableFuture<BoundedSource<OutputT>> emptyFuture = SettableFuture.create();
    emptyFuture.set(null);
    return emptyFuture;
  }
}