Java Code Examples for org.apache.beam.sdk.io.Read#Unbounded

private static <T> TransformEvaluator<Read.Unbounded<T>> readUnbounded() {
  return new TransformEvaluator<Read.Unbounded<T>>() {
    @Override
    public void evaluate(Read.Unbounded<T> transform, EvaluationContext context) {
      final String stepName = context.getCurrentTransform().getFullName();
      context.putDataset(
          transform,
          SparkUnboundedSource.read(
              context.getStreamingContext(),
              context.getSerializableOptions(),
              transform.getSource(),
              stepName));
    }

    @Override
    public String toNativeString() {
      return "streamingContext.<readFrom(<source>)>()";
    }
  };
}
 

@Test
public void readSourceTranslatorUnboundedWithMaxParallelism() {

  final int maxParallelism = 6;
  final int parallelism = 2;

  Read.Unbounded transform = Read.from(new TestUnboundedSource());
  StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
  env.setParallelism(parallelism);
  env.setMaxParallelism(maxParallelism);

  OneInputTransformation<?, ?> sourceTransform =
      (OneInputTransformation)
          applyReadSourceTransform(transform, PCollection.IsBounded.UNBOUNDED, env);

  UnboundedSourceWrapper source =
      (UnboundedSourceWrapper)
          ((SourceTransformation) sourceTransform.getInput()).getOperator().getUserFunction();

  assertEquals(maxParallelism, source.getSplitSources().size());
}
 

@Test
public void readSourceTranslatorUnboundedWithoutMaxParallelism() {

  final int parallelism = 2;

  Read.Unbounded transform = Read.from(new TestUnboundedSource());
  StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
  env.setParallelism(parallelism);

  OneInputTransformation<?, ?> sourceTransform =
      (OneInputTransformation)
          applyReadSourceTransform(transform, PCollection.IsBounded.UNBOUNDED, env);

  UnboundedSourceWrapper source =
      (UnboundedSourceWrapper)
          ((SourceTransformation) sourceTransform.getInput()).getOperator().getUserFunction();

  assertEquals(parallelism, source.getSplitSources().size());
}
 

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator(Read.Unbounded.class)
private static void unboundedReadTranslator(final PipelineTranslationContext ctx,
                                            final TransformHierarchy.Node beamNode,
                                            final Read.Unbounded<?> transform) {
  final IRVertex vertex = new BeamUnboundedSourceVertex<>(transform.getSource(), DisplayData.from(transform));
  ctx.addVertex(vertex);
  beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
  beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
}
 

@Override
public FunctionSpec translate(
    AppliedPTransform<?, ?, Read.Unbounded<?>> transform, SdkComponents components) {
  ReadPayload payload = toProto(transform.getTransform(), components);
  return RunnerApi.FunctionSpec.newBuilder()
      .setUrn(getUrn(transform.getTransform()))
      .setPayload(payload.toByteString())
      .build();
}
 

@Test
public void testToFromProtoUnbounded() throws Exception {
  assumeThat(source, instanceOf(UnboundedSource.class));
  UnboundedSource<?, ?> unboundedSource = (UnboundedSource<?, ?>) this.source;
  Read.Unbounded<?> unboundedRead = Read.from(unboundedSource);
  SdkComponents components = SdkComponents.create();
  // No environment set for unbounded sources
  ReadPayload payload = ReadTranslation.toProto(unboundedRead, components);
  assertThat(payload.getIsBounded(), equalTo(RunnerApi.IsBounded.Enum.UNBOUNDED));
  UnboundedSource<?, ?> deserializedSource = ReadTranslation.unboundedSourceFromProto(payload);
  assertThat(deserializedSource, equalTo(source));
}
 

@Override
public String toString() {
  try {
    Class<? extends PTransform> transformClass = transform.getClass();
    if ("KafkaIO.Read".equals(node.getFullName())) {
      return "KafkaUtils.createDirectStream(...)";
    }
    if (composite) {
      return "_.<" + transformClass.getName() + ">";
    }
    String transformString = transformEvaluator.toNativeString();
    if (transformString.contains("<fn>")) {
      transformString = replaceFnString(transformClass, transformString, "fn");
    } else if (transformString.contains("<windowFn>")) {
      transformString = replaceFnString(transformClass, transformString, "windowFn");
    } else if (transformString.contains("<source>")) {
      String sourceName = "...";
      if (transform instanceof Read.Bounded) {
        sourceName = ((Read.Bounded<?>) transform).getSource().getClass().getName();
      } else if (transform instanceof Read.Unbounded) {
        sourceName = ((Read.Unbounded<?>) transform).getSource().getClass().getName();
      }
      transformString = transformString.replace("<source>", sourceName);
    }
    if (transformString.startsWith("sparkContext")
        || transformString.startsWith("streamingContext")) {
      return transformString;
    }
    return "_." + transformString;
  } catch (NoSuchMethodException
      | InvocationTargetException
      | IllegalAccessException
      | NoSuchFieldException e) {
    return "<FailedTranslation>";
  }
}
 

private void processElement(final TestUnboundedSource<String> source) throws Exception {
  final EvaluationContext context =
      EvaluationContext.create(
          MockClock.fromInstant(Instant.now()),
          CloningBundleFactory.create(),
          DirectGraph.create(
              emptyMap(), emptyMap(), LinkedListMultimap.create(), emptySet(), emptyMap()),
          emptySet(),
          Executors.newCachedThreadPool());
  final UnboundedReadEvaluatorFactory factory =
      new UnboundedReadEvaluatorFactory(context, options);

  final Read.Unbounded<String> unbounded = Read.from(source);
  final Pipeline pipeline = Pipeline.create(options);
  final PCollection<String> pCollection = pipeline.apply(unbounded);
  final AppliedPTransform<PBegin, PCollection<String>, Read.Unbounded<String>> application =
      AppliedPTransform.of(
          "test",
          new HashMap<>(),
          singletonMap(new TupleTag(), pCollection),
          unbounded,
          pipeline);
  final TransformEvaluator<UnboundedSourceShard<String, TestCheckpointMark>> evaluator =
      factory.forApplication(application, null);
  final UnboundedSource.UnboundedReader<String> reader = source.createReader(options, null);
  final UnboundedSourceShard<String, TestCheckpointMark> shard =
      UnboundedSourceShard.of(source, new NeverDeduplicator(), reader, null);
  final WindowedValue<UnboundedSourceShard<String, TestCheckpointMark>> value =
      WindowedValue.of(
          shard, BoundedWindow.TIMESTAMP_MAX_VALUE, GlobalWindow.INSTANCE, PaneInfo.NO_FIRING);
  TestUnboundedSource.readerClosedCount = 0;
  evaluator.processElement(value);
}
 

@Override
public String getUrn(Read.Unbounded<?> transform) {
  return PTransformTranslation.READ_TRANSFORM_URN;
}
 

private static AppliedPTransform<?, ?, ?> read(Pipeline pipeline) {
  Read.Unbounded<Long> transform = Read.from(CountingSource.unbounded());
  PCollection<Long> pcollection = pipeline.apply(transform);
  return AppliedPTransform.of(
      "ReadTheCount", pipeline.begin().expand(), pcollection.expand(), transform, pipeline);
}
 

@Override
public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform(
    AppliedPTransform<PBegin, PCollection<T>, Read.Unbounded<T>> transform) {
  return PTransformReplacement.of(
      transform.getPipeline().begin(), new StreamingUnboundedRead<>(transform.getTransform()));
}
 

public StreamingUnboundedRead(Read.Unbounded<T> transform) {
  this.source = transform.getSource();
}
 

private <OutputT> TransformEvaluator<?> createEvaluator(
    AppliedPTransform<PBegin, PCollection<OutputT>, Read.Unbounded<OutputT>> application) {
  return new UnboundedReadEvaluator<>(application, evaluationContext, options, readerReuseChance);
}