Java Code Examples for org.apache.beam.model.pipeline.v1.RunnerApi#ReadPayload

The following examples show how to use org.apache.beam.model.pipeline.v1.RunnerApi#ReadPayload . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.

Example 1

Source File: FlinkStreamingPortablePipelineTranslator.java From beam with Apache License 2.0

5 votes

@SuppressWarnings("unchecked")
private <T> void translateUnboundedRead(
    String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext context) {
  RunnerApi.PTransform transform = pipeline.getComponents().getTransformsOrThrow(id);
  String outputCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values());

  RunnerApi.ReadPayload payload;
  try {
    payload = RunnerApi.ReadPayload.parseFrom(transform.getSpec().getPayload());
  } catch (IOException e) {
    throw new RuntimeException("Failed to parse ReadPayload from transform", e);
  }

  Preconditions.checkState(
      payload.getIsBounded() != RunnerApi.IsBounded.Enum.BOUNDED,
      "Bounded reads should run inside an environment instead of being translated by the Runner.");

  DataStream<WindowedValue<T>> source =
      translateUnboundedSource(
          transform.getUniqueName(),
          outputCollectionId,
          payload,
          pipeline,
          context.getPipelineOptions(),
          context.getExecutionEnvironment());

  context.addDataStream(outputCollectionId, source);
}

Example 2

Source File: PubsubIOExternalTest.java From beam with Apache License 2.0

4 votes

@Test
public void testConstructPubsubRead() throws Exception {
  String topic = "projects/project-1234/topics/topic_name";
  String idAttribute = "id_foo";
  Boolean needsAttributes = true;

  ExternalTransforms.ExternalConfigurationPayload payload =
      ExternalTransforms.ExternalConfigurationPayload.newBuilder()
          .putConfiguration(
              "topic",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(topic)))
                  .build())
          .putConfiguration(
              "id_label",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(idAttribute)))
                  .build())
          .putConfiguration(
              "with_attributes",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:bool:v1")
                  .setPayload(ByteString.copyFrom(encodeBoolean(needsAttributes)))
                  .build())
          .build();

  RunnerApi.Components defaultInstance = RunnerApi.Components.getDefaultInstance();
  ExpansionApi.ExpansionRequest request =
      ExpansionApi.ExpansionRequest.newBuilder()
          .setComponents(defaultInstance)
          .setTransform(
              RunnerApi.PTransform.newBuilder()
                  .setUniqueName("test")
                  .setSpec(
                      RunnerApi.FunctionSpec.newBuilder()
                          .setUrn("beam:external:java:pubsub:read:v1")
                          .setPayload(payload.toByteString())))
          .setNamespace("test_namespace")
          .build();

  ExpansionService expansionService = new ExpansionService();
  TestStreamObserver<ExpansionApi.ExpansionResponse> observer = new TestStreamObserver<>();
  expansionService.expand(request, observer);

  ExpansionApi.ExpansionResponse result = observer.result;
  RunnerApi.PTransform transform = result.getTransform();
  assertThat(
      transform.getSubtransformsList(),
      Matchers.contains(
          "test_namespacetest/PubsubUnboundedSource", "test_namespacetest/MapElements"));
  assertThat(transform.getInputsCount(), Matchers.is(0));
  assertThat(transform.getOutputsCount(), Matchers.is(1));

  RunnerApi.PTransform pubsubComposite =
      result.getComponents().getTransformsOrThrow(transform.getSubtransforms(0));
  RunnerApi.PTransform pubsubRead =
      result.getComponents().getTransformsOrThrow(pubsubComposite.getSubtransforms(0));
  RunnerApi.ReadPayload readPayload =
      RunnerApi.ReadPayload.parseFrom(pubsubRead.getSpec().getPayload());
  PubsubUnboundedSource.PubsubSource source =
      (PubsubUnboundedSource.PubsubSource) ReadTranslation.unboundedSourceFromProto(readPayload);
  PubsubUnboundedSource spec = source.outer;

  assertThat(
      spec.getTopicProvider() == null ? null : String.valueOf(spec.getTopicProvider()),
      Matchers.is(topic));
  assertThat(spec.getIdAttribute(), Matchers.is(idAttribute));
  assertThat(spec.getNeedsAttributes(), Matchers.is(true));
}

Example 3

Source File: KafkaIOExternalTest.java From beam with Apache License 2.0

4 votes

@Test
public void testConstructKafkaRead() throws Exception {
  List<String> topics = ImmutableList.of("topic1", "topic2");
  String keyDeserializer = "org.apache.kafka.common.serialization.ByteArrayDeserializer";
  String valueDeserializer = "org.apache.kafka.common.serialization.LongDeserializer";
  ImmutableMap<String, String> consumerConfig =
      ImmutableMap.<String, String>builder()
          .put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "server1:port,server2:port")
          .put("key2", "value2")
          .put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, keyDeserializer)
          .put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, valueDeserializer)
          .build();

  ExternalTransforms.ExternalConfigurationPayload payload =
      ExternalTransforms.ExternalConfigurationPayload.newBuilder()
          .putConfiguration(
              "topics",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:iterable:v1")
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(listAsBytes(topics)))
                  .build())
          .putConfiguration(
              "consumer_config",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:iterable:v1")
                  .addCoderUrn("beam:coder:kv:v1")
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(mapAsBytes(consumerConfig)))
                  .build())
          .putConfiguration(
              "key_deserializer",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(keyDeserializer)))
                  .build())
          .putConfiguration(
              "value_deserializer",
              ExternalTransforms.ConfigValue.newBuilder()
                  .addCoderUrn("beam:coder:string_utf8:v1")
                  .setPayload(ByteString.copyFrom(encodeString(valueDeserializer)))
                  .build())
          .build();

  RunnerApi.Components defaultInstance = RunnerApi.Components.getDefaultInstance();
  ExpansionApi.ExpansionRequest request =
      ExpansionApi.ExpansionRequest.newBuilder()
          .setComponents(defaultInstance)
          .setTransform(
              RunnerApi.PTransform.newBuilder()
                  .setUniqueName("test")
                  .setSpec(
                      RunnerApi.FunctionSpec.newBuilder()
                          .setUrn("beam:external:java:kafka:read:v1")
                          .setPayload(payload.toByteString())))
          .setNamespace("test_namespace")
          .build();

  ExpansionService expansionService = new ExpansionService();
  TestStreamObserver<ExpansionApi.ExpansionResponse> observer = new TestStreamObserver<>();
  expansionService.expand(request, observer);

  ExpansionApi.ExpansionResponse result = observer.result;
  RunnerApi.PTransform transform = result.getTransform();
  assertThat(
      transform.getSubtransformsList(),
      Matchers.contains(
          "test_namespacetest/KafkaIO.Read", "test_namespacetest/Remove Kafka Metadata"));
  assertThat(transform.getInputsCount(), Matchers.is(0));
  assertThat(transform.getOutputsCount(), Matchers.is(1));

  RunnerApi.PTransform kafkaComposite =
      result.getComponents().getTransformsOrThrow(transform.getSubtransforms(0));
  RunnerApi.PTransform kafkaRead =
      result.getComponents().getTransformsOrThrow(kafkaComposite.getSubtransforms(0));
  RunnerApi.ReadPayload readPayload =
      RunnerApi.ReadPayload.parseFrom(kafkaRead.getSpec().getPayload());
  KafkaUnboundedSource source =
      (KafkaUnboundedSource) ReadTranslation.unboundedSourceFromProto(readPayload);
  KafkaIO.Read spec = source.getSpec();

  assertThat(spec.getConsumerConfig(), Matchers.is(consumerConfig));
  assertThat(spec.getTopics(), Matchers.is(topics));
  assertThat(
      spec.getKeyDeserializerProvider()
          .getDeserializer(spec.getConsumerConfig(), true)
          .getClass()
          .getName(),
      Matchers.is(keyDeserializer));
  assertThat(
      spec.getValueDeserializerProvider()
          .getDeserializer(spec.getConsumerConfig(), false)
          .getClass()
          .getName(),
      Matchers.is(valueDeserializer));
}

Example 4

Source File: FlinkStreamingPortablePipelineTranslator.java From beam with Apache License 2.0

4 votes

private static <T> DataStream<WindowedValue<T>> translateUnboundedSource(
    String transformName,
    String outputCollectionId,
    RunnerApi.ReadPayload payload,
    RunnerApi.Pipeline pipeline,
    PipelineOptions pipelineOptions,
    StreamExecutionEnvironment env) {

  final DataStream<WindowedValue<T>> source;
  final DataStream<WindowedValue<ValueWithRecordId<T>>> nonDedupSource;
  Coder<WindowedValue<T>> windowCoder =
      instantiateCoder(outputCollectionId, pipeline.getComponents());

  TypeInformation<WindowedValue<T>> outputTypeInfo = new CoderTypeInformation<>(windowCoder);

  WindowingStrategy windowStrategy =
      getWindowingStrategy(outputCollectionId, pipeline.getComponents());
  TypeInformation<WindowedValue<ValueWithRecordId<T>>> withIdTypeInfo =
      new CoderTypeInformation<>(
          WindowedValue.getFullCoder(
              ValueWithRecordId.ValueWithRecordIdCoder.of(
                  ((WindowedValueCoder) windowCoder).getValueCoder()),
              windowStrategy.getWindowFn().windowCoder()));

  UnboundedSource unboundedSource = ReadTranslation.unboundedSourceFromProto(payload);

  try {
    int parallelism =
        env.getMaxParallelism() > 0 ? env.getMaxParallelism() : env.getParallelism();
    UnboundedSourceWrapper sourceWrapper =
        new UnboundedSourceWrapper<>(
            transformName, pipelineOptions, unboundedSource, parallelism);
    nonDedupSource =
        env.addSource(sourceWrapper)
            .name(transformName)
            .uid(transformName)
            .returns(withIdTypeInfo);

    if (unboundedSource.requiresDeduping()) {
      source =
          nonDedupSource
              .keyBy(new FlinkStreamingTransformTranslators.ValueWithRecordIdKeySelector<>())
              .transform("deduping", outputTypeInfo, new DedupingOperator<>(pipelineOptions))
              .uid(format("%s/__deduplicated__", transformName));
    } else {
      source =
          nonDedupSource
              .flatMap(new FlinkStreamingTransformTranslators.StripIdsMap<>(pipelineOptions))
              .returns(outputTypeInfo);
    }
  } catch (Exception e) {
    throw new RuntimeException("Error while translating UnboundedSource: " + unboundedSource, e);
  }

  return source;
}