org.apache.beam.sdk.values.TupleTag Java Examples

public SdkHarnessDoFnRunner(
    DoFn<InputT, OutputT> doFn,
    StageBundleFactory stageBundleFactory,
    StateRequestHandler stateRequestHandler,
    BundleProgressHandler progressHandler,
    BufferedOutputManager<OutputT> outputManager,
    Map<String, TupleTag<?>> outputMap,
    Coder<BoundedWindow> windowCoder,
    BiConsumer<Timer<?>, TimerInternals.TimerData> timerRegistration,
    Supplier<Object> keyForTimer) {

  this.doFn = doFn;
  this.stageBundleFactory = stageBundleFactory;
  this.stateRequestHandler = stateRequestHandler;
  this.progressHandler = progressHandler;
  this.outputManager = outputManager;
  this.outputMap = outputMap;
  this.timerRegistration = timerRegistration;
  this.keyForTimer = keyForTimer;
  this.windowCoder = windowCoder;
  this.outputQueue = new LinkedBlockingQueue<>();
}
 

@Test
public void flattenWithDuplicateInputsWithDuplicates() {
  PCollection<Integer> duplicate =
      PCollection.createPrimitiveOutputInternal(
          p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of());
  AppliedPTransform application =
      AppliedPTransform.of(
          "Flatten",
          ImmutableMap.<TupleTag<?>, PValue>builder()
              .put(new TupleTag<Integer>(), duplicate)
              .put(new TupleTag<Integer>(), duplicate)
              .build(),
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Flatten.pCollections(),
          p);

  assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(true));
}
 

BufferedOutputManager(
    Output<StreamRecord<WindowedValue<OutputT>>> output,
    TupleTag<OutputT> mainTag,
    Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags,
    Map<TupleTag<?>, Integer> tagsToIds,
    Lock bufferLock,
    PushedBackElementsHandler<KV<Integer, WindowedValue<?>>> pushedBackElementsHandler) {
  this.output = output;
  this.mainTag = mainTag;
  this.tagsToOutputTags = tagsToOutputTags;
  this.tagsToIds = tagsToIds;
  this.bufferLock = bufferLock;
  this.idsToTags = new HashMap<>();
  for (Map.Entry<TupleTag<?>, Integer> entry : tagsToIds.entrySet()) {
    idsToTags.put(entry.getValue(), entry.getKey());
  }
  this.pushedBackElementsHandler = pushedBackElementsHandler;
}
 

public static <OutputT> PCollectionTuple createPrimitiveOutputFor(
    PCollection<?> input,
    DoFn<?, OutputT> fn,
    TupleTag<OutputT> mainOutputTag,
    TupleTagList additionalOutputTags,
    Map<TupleTag<?>, Coder<?>> outputTagsToCoders,
    WindowingStrategy<?, ?> windowingStrategy) {
  DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());
  PCollectionTuple outputs =
      PCollectionTuple.ofPrimitiveOutputsInternal(
          input.getPipeline(),
          TupleTagList.of(mainOutputTag).and(additionalOutputTags.getAll()),
          outputTagsToCoders,
          windowingStrategy,
          input.isBounded().and(signature.isBoundedPerElement()));

  // Set output type descriptor similarly to how ParDo.MultiOutput does it.
  outputs.get(mainOutputTag).setTypeDescriptor(fn.getOutputTypeDescriptor());

  return outputs;
}
 

@Test
@Category(NeedsRunner.class)
public void testTaggedOutputUnknownCoder() throws Exception {

  PCollection<Integer> input = pipeline.apply(Create.of(Arrays.asList(1, 2, 3)));

  final TupleTag<Integer> mainOutputTag = new TupleTag<>("main");
  final TupleTag<TestDummy> additionalOutputTag = new TupleTag<>("unknownSide");
  input.apply(
      ParDo.of(new TaggedOutputDummyFn(mainOutputTag, additionalOutputTag))
          .withOutputTags(mainOutputTag, TupleTagList.of(additionalOutputTag)));

  thrown.expect(IllegalStateException.class);
  thrown.expectMessage("Unable to return a default Coder");
  pipeline.run();
}
 

@Test
public void flattenWithDuplicateInputsWithoutDuplicates() {
  AppliedPTransform application =
      AppliedPTransform.of(
          "Flatten",
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Collections.singletonMap(
              new TupleTag<Integer>(),
              PCollection.createPrimitiveOutputInternal(
                  p, WindowingStrategy.globalDefault(), IsBounded.BOUNDED, VarIntCoder.of())),
          Flatten.pCollections(),
          p);

  assertThat(PTransformMatchers.flattenWithDuplicateInputs().matches(application), is(false));
}
 

@Internal
public Node addFinalizedPrimitiveNode(
    String name,
    Map<TupleTag<?>, PValue> inputs,
    PTransform<?, ?> transform,
    Map<TupleTag<?>, PValue> outputs) {
  checkNotNull(
      transform, "A %s must be provided for all Nodes", PTransform.class.getSimpleName());
  checkNotNull(
      name, "A name must be provided for all %s Nodes", PTransform.class.getSimpleName());
  checkNotNull(
      inputs, "Inputs must be provided for all %s Nodes", PTransform.class.getSimpleName());
  checkNotNull(
      outputs, "Outputs must be provided for all %s Nodes", PTransform.class.getSimpleName());
  Node node = new Node(current, transform, name, inputs, outputs);
  node.finishedSpecifying = true;
  for (PValue output : outputs.values()) {
    producers.put(output, node);
  }
  current.addComposite(node);
  return node;
}
 

@Override
public ParDoFn create(
    PipelineOptions options,
    CloudObject cloudUserFn,
    List<SideInputInfo> sideInputInfos,
    TupleTag<?> mainOutputTag,
    Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices,
    DataflowExecutionContext<?> executionContext,
    DataflowOperationContext operationContext)
    throws Exception {
  Coder<?> coder =
      CloudObjects.coderFromCloudObject(
          CloudObject.fromSpec(Structs.getObject(cloudUserFn, PropertyNames.ENCODING)));
  checkState(
      coder instanceof IsmRecordCoder,
      "Expected to received an instanceof an %s but got %s",
      IsmRecordCoder.class.getSimpleName(),
      coder);
  IsmRecordCoder<?> ismRecordCoder = (IsmRecordCoder<?>) coder;
  return new ToIsmRecordForMultimapParDoFn(
      KvCoder.of(
          ismRecordCoder.getCoderArguments().get(0), ismRecordCoder.getCoderArguments().get(1)));
}
 

/**
 * Constructs a PartitionDoFn.
 *
 * @throws IllegalArgumentException if {@code numPartitions <= 0}
 */
private PartitionDoFn(
    int numPartitions,
    Contextful<Contextful.Fn<X, Integer>> ctxFn,
    Object originalFnClassForDisplayData) {
  this.ctxFn = ctxFn;
  this.originalFnClassForDisplayData = originalFnClassForDisplayData;
  if (numPartitions <= 0) {
    throw new IllegalArgumentException("numPartitions must be > 0");
  }

  this.numPartitions = numPartitions;

  TupleTagList buildOutputTags = TupleTagList.empty();
  for (int partition = 0; partition < numPartitions; partition++) {
    buildOutputTags = buildOutputTags.and(new TupleTag<X>());
  }
  outputTags = buildOutputTags;
}
 

@Test
@Category(ValidatesRunner.class)
public void testParDoWithEmptyTaggedOutput() {
  TupleTag<String> mainOutputTag = new TupleTag<String>("main") {};
  TupleTag<String> additionalOutputTag1 = new TupleTag<String>("additional1") {};
  TupleTag<String> additionalOutputTag2 = new TupleTag<String>("additional2") {};

  PCollectionTuple outputs =
      pipeline
          .apply(Create.empty(VarIntCoder.of()))
          .apply(
              ParDo.of(new TestNoOutputDoFn())
                  .withOutputTags(
                      mainOutputTag,
                      TupleTagList.of(additionalOutputTag1).and(additionalOutputTag2)));

  PAssert.that(outputs.get(mainOutputTag)).empty();

  PAssert.that(outputs.get(additionalOutputTag1)).empty();
  PAssert.that(outputs.get(additionalOutputTag2)).empty();

  pipeline.run();
}
 

private static <K, V, OutputT>
    PairFlatMapFunction<Iterator<Tuple2<ByteArray, byte[]>>, TupleTag<?>, WindowedValue<?>>
        wrapDoFnFromSortedRDD(
            MultiDoFnFunction<KV<K, V>, OutputT> doFnFunction,
            Coder<K> keyCoder,
            Coder<WindowedValue<V>> wvCoder) {

  return (Iterator<Tuple2<ByteArray, byte[]>> in) -> {
    Iterator<Iterator<Tuple2<TupleTag<?>, WindowedValue<?>>>> mappedGroups;
    mappedGroups =
        Iterators.transform(
            splitBySameKey(in, keyCoder, wvCoder),
            group -> {
              try {
                return doFnFunction.call(group);
              } catch (Exception ex) {
                throw new RuntimeException(ex);
              }
            });
    return flatten(mappedGroups);
  };
}
 

@Test
@Category(NeedsRunner.class)
public void testMainOutputUnregisteredExplicitCoder() {

  PCollection<Integer> input = pipeline.apply(Create.of(Arrays.asList(1, 2, 3)));

  final TupleTag<TestDummy> mainOutputTag = new TupleTag<>("unregisteredMain");
  final TupleTag<Integer> additionalOutputTag = new TupleTag<Integer>("additionalOutput") {};
  PCollectionTuple outputTuple =
      input.apply(
          ParDo.of(new MainOutputDummyFn(mainOutputTag, additionalOutputTag))
              .withOutputTags(mainOutputTag, TupleTagList.of(additionalOutputTag)));

  outputTuple.get(mainOutputTag).setCoder(new TestDummyCoder());

  pipeline.run();
}
 

WritePartition(
    boolean singletonTable,
    DynamicDestinations<?, DestinationT> dynamicDestinations,
    PCollectionView<String> tempFilePrefix,
    int maxNumFiles,
    long maxSizeBytes,
    TupleTag<KV<ShardedKey<DestinationT>, List<String>>> multiPartitionsTag,
    TupleTag<KV<ShardedKey<DestinationT>, List<String>>> singlePartitionTag,
    RowWriterFactory<?, DestinationT> rowWriterFactory) {
  this.singletonTable = singletonTable;
  this.dynamicDestinations = dynamicDestinations;
  this.tempFilePrefix = tempFilePrefix;
  this.maxNumFiles = maxNumFiles;
  this.maxSizeBytes = maxSizeBytes;
  this.multiPartitionsTag = multiPartitionsTag;
  this.singlePartitionTag = singlePartitionTag;
  this.rowWriterFactory = rowWriterFactory;
}
 

@Test
public void testWithOutputTagsDisplayData() {
  DoFn<String, String> fn =
      new DoFn<String, String>() {
        @ProcessElement
        public void proccessElement(ProcessContext c) {}

        @Override
        public void populateDisplayData(Builder builder) {
          builder.add(DisplayData.item("fnMetadata", "foobar"));
        }
      };

  ParDo.MultiOutput<String, String> parDo =
      ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty());

  DisplayData displayData = DisplayData.from(parDo);
  assertThat(displayData, includesDisplayDataFor("fn", fn));
  assertThat(displayData, hasDisplayItem("fn", fn.getClass()));
}
 

@Test
@Category({ValidatesRunner.class, UsesSideInputs.class})
public void testCoGroupByKeyGetOnly() {
  final TupleTag<String> tag1 = new TupleTag<>();
  final TupleTag<String> tag2 = new TupleTag<>();

  PCollection<KV<Integer, CoGbkResult>> coGbkResults = buildGetOnlyGbk(p, tag1, tag2);

  PAssert.thatMap(coGbkResults)
      .satisfies(
          results -> {
            assertEquals("collection1-1", results.get(1).getOnly(tag1));
            assertEquals("collection1-2", results.get(2).getOnly(tag1));
            assertEquals("collection2-2", results.get(2).getOnly(tag2));
            assertEquals("collection2-3", results.get(3).getOnly(tag2));
            return null;
          });

  p.run();
}
 

@Test
public void nonAdditionalInputsWithOnlyAdditionalInputsThrows() {
  Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>();
  additionalInputs.put(new TupleTag<String>() {}, pipeline.apply(Create.of("1, 2", "3")));
  additionalInputs.put(new TupleTag<Long>() {}, pipeline.apply(GenerateSequence.from(3L)));

  AppliedPTransform<PInput, POutput, TestTransform> transform =
      AppliedPTransform.of(
          "additional-only",
          additionalInputs,
          Collections.emptyMap(),
          new TestTransform(additionalInputs),
          pipeline);

  thrown.expect(IllegalArgumentException.class);
  thrown.expectMessage("at least one");
  TransformInputs.nonAdditionalInputs(transform);
}
 

private static Pipeline createPipeline(
    PipelineOptions options, String singleOutputPrefix, String multiOutputPrefix) {
  Pipeline p = Pipeline.create(options);

  SerializableFunction<Void, Void> firstTime =
      (SerializableFunction<Void, Void>)
          value -> {
            latch.countDown();
            return null;
          };

  PCollection<String> impulse = p.apply("CreatePCollectionOfOneValue", Create.of(VALUE));
  impulse
      .apply(
          "Single-PairWithRandomKey",
          MapElements.via(new RequiresStableInputIT.PairWithRandomKeyFn()))
      .apply(
          "Single-MakeSideEffectAndThenFail",
          ParDo.of(
              new RequiresStableInputIT.MakeSideEffectAndThenFailFn(
                  singleOutputPrefix, firstTime)));
  impulse
      .apply(
          "Multi-PairWithRandomKey",
          MapElements.via(new RequiresStableInputIT.PairWithRandomKeyFn()))
      .apply(
          "Multi-MakeSideEffectAndThenFail",
          ParDo.of(
                  new RequiresStableInputIT.MakeSideEffectAndThenFailFn(
                      multiOutputPrefix, firstTime))
              .withOutputTags(new TupleTag<>(), TupleTagList.empty()));

  return p;
}
 

@Test
public void testBadArgs() throws Exception {
  expectedEx.expect(RuntimeException.class);
  expectedEx.expectMessage(Matchers.containsString("Unable to find factory method"));

  InstanceBuilder.ofType(TupleTag.class)
      .fromClassName(InstanceBuilderTest.class.getName())
      .fromFactoryMethod("createTag")
      .withArg(String.class, "hello")
      .withArg(Integer.class, 42)
      .build();
}
 

@Test
public void parDoWithFnTypeWithNoMatch() {
  DoFn<Object, Object> fn =
      new DoFn<Object, Object>() {
        @ProcessElement
        public void process(ProcessContext ctxt) {}
      };
  AppliedPTransform<?, ?, ?> parDoSingle = getAppliedTransform(ParDo.of(fn));
  AppliedPTransform<?, ?, ?> parDoMulti =
      getAppliedTransform(ParDo.of(fn).withOutputTags(new TupleTag<>(), TupleTagList.empty()));

  PTransformMatcher matcher = PTransformMatchers.parDoWithFnType(doFnWithState.getClass());
  assertThat(matcher.matches(parDoSingle), is(false));
  assertThat(matcher.matches(parDoMulti), is(false));
}
 

MultiOutput(
    DoFn<InputT, OutputT> fn,
    Map<String, PCollectionView<?>> sideInputs,
    TupleTag<OutputT> mainOutputTag,
    TupleTagList additionalOutputTags,
    ItemSpec<? extends Class<?>> fnDisplayData) {
  this.sideInputs = sideInputs;
  this.mainOutputTag = mainOutputTag;
  this.additionalOutputTags = additionalOutputTags;
  this.fn = fn;
  this.fnDisplayData = fnDisplayData;
}
 

@Override
public <T> void outputWithTimestamp(TupleTag<T> tag, T value, Instant timestamp) {
  noteOutput();
  if (watermarkEstimator instanceof TimestampObservingWatermarkEstimator) {
    ((TimestampObservingWatermarkEstimator) watermarkEstimator).observeTimestamp(timestamp);
  }
  output.outputWindowedValue(tag, value, timestamp, element.getWindows(), element.getPane());
}
 

private OperationNode createParDoOperation(
    Network<Node, Edge> network,
    ParallelInstructionNode node,
    PipelineOptions options,
    DataflowExecutionContext<?> executionContext,
    DataflowOperationContext operationContext)
    throws Exception {

  ParallelInstruction instruction = node.getParallelInstruction();
  ParDoInstruction parDo = instruction.getParDo();

  TupleTag<?> mainOutputTag = tupleTag(parDo.getMultiOutputInfos().get(0));
  ImmutableMap.Builder<TupleTag<?>, Integer> outputTagsToReceiverIndicesBuilder =
      ImmutableMap.builder();
  int successorOffset = 0;
  for (Node successor : network.successors(node)) {
    for (Edge edge : network.edgesConnecting(node, successor)) {
      outputTagsToReceiverIndicesBuilder.put(
          tupleTag(((MultiOutputInfoEdge) edge).getMultiOutputInfo()), successorOffset);
    }
    successorOffset += 1;
  }
  ParDoFn fn =
      parDoFnFactory.create(
          options,
          CloudObject.fromSpec(parDo.getUserFn()),
          parDo.getSideInputs(),
          mainOutputTag,
          outputTagsToReceiverIndicesBuilder.build(),
          executionContext,
          operationContext);

  OutputReceiver[] receivers = getOutputReceivers(network, node);
  return OperationNode.create(new ParDoOperation(fn, receivers, operationContext));
}
 

@Override
protected DoFnRunner<InputT, OutputT> getDoFnRunner(
    PipelineOptions pipelineOptions,
    DoFn<InputT, OutputT> doFn,
    SideInputReader sideInputReader,
    JetOutputManager outputManager,
    TupleTag<OutputT> mainOutputTag,
    List<TupleTag<?>> additionalOutputTags,
    Coder<InputT> inputValueCoder,
    Map<TupleTag<?>, Coder<?>> outputValueCoders,
    WindowingStrategy<?, ?> windowingStrategy,
    DoFnSchemaInformation doFnSchemaInformation,
    Map<String, PCollectionView<?>> sideInputMapping) {
  return DoFnRunners.simpleRunner(
      pipelineOptions,
      doFn,
      sideInputReader,
      outputManager,
      mainOutputTag,
      additionalOutputTags,
      new NotImplementedStepContext(),
      inputValueCoder,
      outputValueCoders,
      windowingStrategy,
      doFnSchemaInformation,
      sideInputMapping);
}
 

@Test
public void singletonMultipleOriginalsThrows() {
  thrown.expect(IllegalArgumentException.class);
  ReplacementOutputs.singleton(
      ImmutableMap.<TupleTag<?>, PValue>builder()
          .putAll(ints.expand())
          .putAll(moreInts.expand())
          .build(),
      replacementInts);
}
 

@Override
public Map<TupleTag<?>, PValue> expand() {
  Map<TupleTag<?>, PValue> values = new HashMap<>();
  values.put(failuresTag(), failures());
  if (outputTag() != null && output() instanceof PValue) {
    values.put(outputTag(), (PValue) output());
  }
  return values;
}
 

/**
 * Returns a new {@link CoGbkResult} based on this, with the given tag and given data added to it.
 */
public <V> CoGbkResult and(TupleTag<V> tag, List<V> data) {
  if (nextTestUnionId != schema.size()) {
    throw new IllegalArgumentException(
        "Attempting to call and() on a CoGbkResult apparently not created by" + " of().");
  }
  List<Iterable<?>> valueMap = new ArrayList<>(this.valueMap);
  valueMap.add(data);
  return new CoGbkResult(
      new CoGbkResultSchema(schema.getTupleTagList().and(tag)), valueMap, nextTestUnionId + 1);
}
 

private CoGbkResultSchema createSchema(int size) {
  List<TupleTag<?>> tags = new ArrayList<>();
  for (int i = 0; i < size; i++) {
    tags.add(new TupleTag<Integer>("tag" + i));
  }
  return new CoGbkResultSchema(TupleTagList.of(tags));
}
 

AbstractSupplier(
    String stepId,
    String ownerId,
    DoFn<InputT, OutputT> doFn,
    WindowingStrategy<?, ?> windowingStrategy,
    DoFnSchemaInformation doFnSchemaInformation,
    SerializablePipelineOptions pipelineOptions,
    TupleTag<OutputT> mainOutputTag,
    Set<TupleTag<OutputT>> allOutputTags,
    Coder<InputT> inputCoder,
    Map<PCollectionView<?>, Coder<?>> sideInputCoders,
    Map<TupleTag<?>, Coder<?>> outputCoders,
    Coder<InputT> inputValueCoder,
    Map<TupleTag<?>, Coder<?>> outputValueCoders,
    List<PCollectionView<?>> sideInputs) {
  this.stepId = stepId;
  this.ownerId = ownerId;
  this.pipelineOptions = pipelineOptions;
  this.doFn = doFn;
  this.windowingStrategy = windowingStrategy;
  this.doFnSchemaInformation = doFnSchemaInformation;
  this.outputCollToOrdinals =
      allOutputTags.stream()
          .collect(Collectors.toMap(Function.identity(), t -> new ArrayList<>()));
  this.mainOutputTag = mainOutputTag;
  this.inputCoder = inputCoder;
  this.sideInputCoders = sideInputCoders;
  this.outputCoders = outputCoders;
  this.inputValueCoder = inputValueCoder;
  this.outputValueCoders = outputValueCoders;
  this.sideInputs = sideInputs;
}
 

@Test
@Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesParDoLifecycle.class})
public void testFnCallSequenceStateful() {
  PCollectionList.of(p.apply("Impolite", Create.of(KV.of("a", 1), KV.of("b", 2), KV.of("a", 4))))
      .and(
          p.apply(
              "Polite", Create.of(KV.of("b", 3), KV.of("a", 5), KV.of("c", 6), KV.of("c", 7))))
      .apply(Flatten.pCollections())
      .apply(
          ParDo.of(new CallSequenceEnforcingStatefulFn<String, Integer>())
              .withOutputTags(new TupleTag<KV<String, Integer>>() {}, TupleTagList.empty()));

  p.run();
}
 

@Test
public void sdkErrorsSurfaceOnClose() throws Exception {
  TupleTag<Integer> mainOutput = new TupleTag<>("main-output");
  DoFnOperator.MultiOutputOutputManagerFactory<Integer> outputManagerFactory =
      new DoFnOperator.MultiOutputOutputManagerFactory(mainOutput, VoidCoder.of());
  ExecutableStageDoFnOperator<Integer, Integer> operator =
      getOperator(mainOutput, Collections.emptyList(), outputManagerFactory);

  OneInputStreamOperatorTestHarness<WindowedValue<Integer>, WindowedValue<Integer>> testHarness =
      new OneInputStreamOperatorTestHarness<>(operator);

  testHarness.open();

  @SuppressWarnings("unchecked")
  RemoteBundle bundle = Mockito.mock(RemoteBundle.class);
  when(stageBundleFactory.getBundle(any(), any(), any(), any())).thenReturn(bundle);

  @SuppressWarnings("unchecked")
  FnDataReceiver<WindowedValue<?>> receiver = Mockito.mock(FnDataReceiver.class);
  when(bundle.getInputReceivers()).thenReturn(ImmutableMap.of("input", receiver));

  Exception expected = new RuntimeException(new Exception());
  doThrow(expected).when(bundle).close();
  thrown.expectCause(is(expected));

  operator.processElement(new StreamRecord<>(WindowedValue.valueInGlobalWindow(0)));
  testHarness.close();
}