Java Code Examples for org.apache.beam.sdk.runners.AppliedPTransform#getTransform()

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
 * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and {@link
 * DoFnSignature#usesTimers()}.
 */
public static PTransformMatcher stateOrTimerParDoSingle() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      PTransform<?, ?> transform = application.getTransform();
      if (transform instanceof ParDo.SingleOutput) {
        DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn();
        DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
        return signature.usesState() || signature.usesTimers();
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("StateOrTimerParDoSingleMatcher").toString();
    }
  };
}
 

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
 * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
 */
public static PTransformMatcher splittableParDoSingle() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      PTransform<?, ?> transform = application.getTransform();
      if (transform instanceof ParDo.SingleOutput) {
        DoFn<?, ?> fn = ((ParDo.SingleOutput<?, ?>) transform).getFn();
        DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
        return signature.processElement().isSplittable();
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("SplittableParDoSingleMatcher").toString();
    }
  };
}
 

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
 * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
 */
public static PTransformMatcher splittableProcessKeyedBounded() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      PTransform<?, ?> transform = application.getTransform();
      if (transform instanceof SplittableParDo.ProcessKeyedElements) {
        DoFn<?, ?> fn = ((SplittableParDo.ProcessKeyedElements) transform).getFn();
        DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
        return signature.processElement().isSplittable()
            && signature.isBoundedPerElement() == IsBounded.BOUNDED;
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("SplittableProcessKeyedBoundedMatcher").toString();
    }
  };
}
 

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.SingleOutput} containing a {@link DoFn}
 * that is splittable, as signified by {@link ProcessElementMethod#isSplittable()}.
 */
public static PTransformMatcher splittableProcessKeyedUnbounded() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      PTransform<?, ?> transform = application.getTransform();
      if (transform instanceof SplittableParDo.ProcessKeyedElements) {
        DoFn<?, ?> fn = ((SplittableParDo.ProcessKeyedElements) transform).getFn();
        DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
        return signature.processElement().isSplittable()
            && signature.isBoundedPerElement() == IsBounded.UNBOUNDED;
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("SplittableProcessKeyedUnboundedMatcher").toString();
    }
  };
}
 

@Override
public Collection<CommittedBundle<TestStreamIndex<T>>> getInitialInputs(
    AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>> transform,
    int targetParallelism) {

  // This will always be run on an execution-time transform, so it can be downcast
  DirectTestStreamFactory.DirectTestStream<T> testStream =
      (DirectTestStreamFactory.DirectTestStream<T>) transform.getTransform();

  CommittedBundle<TestStreamIndex<T>> initialBundle =
      evaluationContext
          .<TestStreamIndex<T>>createRootBundle()
          .add(WindowedValue.valueInGlobalWindow(TestStreamIndex.of(testStream.original)))
          .commit(BoundedWindow.TIMESTAMP_MAX_VALUE);
  return Collections.singleton(initialBundle);
}
 

/**
 * A {@link PTransformMatcher} which matches a {@link Flatten.PCollections} which consumes a
 * single input {@link PCollection} multiple times.
 */
public static PTransformMatcher flattenWithDuplicateInputs() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      if (application.getTransform() instanceof Flatten.PCollections) {
        Set<PValue> observed = new HashSet<>();
        for (PValue pvalue : application.getInputs().values()) {
          boolean firstInstance = observed.add(pvalue);
          if (!firstInstance) {
            return true;
          }
        }
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("FlattenWithDuplicateInputsMatcher").toString();
    }
  };
}
 

public static TupleTagList getAdditionalOutputTags(AppliedPTransform<?, ?, ?> application)
    throws IOException {
  PTransform<?, ?> transform = application.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    return ((ParDo.MultiOutput<?, ?>) transform).getAdditionalOutputTags();
  }

  RunnerApi.PTransform protoTransform =
      PTransformTranslation.toProto(
          application, SdkComponents.create(application.getPipeline().getOptions()));

  ParDoPayload payload = ParDoPayload.parseFrom(protoTransform.getSpec().getPayload());
  TupleTag<?> mainOutputTag = getMainOutputTag(payload);
  Set<String> outputTags =
      Sets.difference(
          protoTransform.getOutputsMap().keySet(), Collections.singleton(mainOutputTag.getId()));

  ArrayList<TupleTag<?>> additionalOutputTags = new ArrayList<>();
  for (String outputTag : outputTags) {
    additionalOutputTags.add(new TupleTag<>(outputTag));
  }
  return TupleTagList.of(additionalOutputTags);
}
 

@Override
public Vertex translate(
    Pipeline pipeline,
    AppliedPTransform<?, ?, ?> appliedTransform,
    TransformHierarchy.Node node,
    JetTranslationContext context) {
  String transformName = appliedTransform.getFullName();
  DAGBuilder dagBuilder = context.getDagBuilder();
  String vertexId = dagBuilder.newVertexId(transformName);

  TestStream<T> testStream = (TestStream<T>) appliedTransform.getTransform();

  // events in the transform are not serializable, we have to translate them. We'll also flatten
  // the collection.
  Map.Entry<TupleTag<?>, PValue> output = Utils.getOutput(appliedTransform);
  Coder outputCoder = Utils.getCoder((PCollection) output.getValue());
  TestStream.TestStreamCoder<T> payloadCoder =
      TestStream.TestStreamCoder.of(testStream.getValueCoder());
  byte[] encodedPayload = getEncodedPayload(testStream, payloadCoder);
  Vertex vertex =
      dagBuilder.addVertex(
          vertexId, TestStreamP.supplier(encodedPayload, payloadCoder, outputCoder));

  String outputEdgeId = Utils.getTupleTagId(output.getValue());
  dagBuilder.registerCollectionOfEdge(outputEdgeId, output.getKey().getId());
  dagBuilder.registerEdgeStartPoint(outputEdgeId, vertex, outputCoder);
  return vertex;
}
 

static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> appliedTransform) {
  PTransform<?, ?> transform = appliedTransform.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    ParDo.MultiOutput multiParDo = (ParDo.MultiOutput) transform;
    return (List) multiParDo.getSideInputs().values().stream().collect(Collectors.toList());
  } else if (transform instanceof ParDo.SingleOutput) {
    ParDo.SingleOutput singleParDo = (ParDo.SingleOutput) transform;
    return (List) singleParDo.getSideInputs().values().stream().collect(Collectors.toList());
  }
  return Collections.emptyList();
}
 

@SuppressWarnings("unchecked")
static PValue getInput(AppliedPTransform<?, ?, ?> appliedTransform) {
  if (appliedTransform.getTransform() == null) {
    return null;
  }
  return Iterables.getOnlyElement(TransformInputs.nonAdditionalInputs(appliedTransform));
}
 

@SuppressWarnings("unchecked")
private PTransform<PCollection<? extends InputT>, PCollectionTuple> getReplacementForApplication(
    AppliedPTransform<
            PCollection<? extends InputT>,
            PCollectionTuple,
            PTransform<PCollection<? extends InputT>, PCollectionTuple>>
        application)
    throws IOException {

  DoFn<InputT, OutputT> fn = (DoFn<InputT, OutputT>) ParDoTranslation.getDoFn(application);

  DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass());

  if (signature.processElement().isSplittable()) {
    return SplittableParDo.forAppliedParDo((AppliedPTransform) application);
  } else if (signature.stateDeclarations().size() > 0
      || signature.timerDeclarations().size() > 0
      || signature.timerFamilyDeclarations().size() > 0) {
    return new GbkThenStatefulParDo(
        fn,
        ParDoTranslation.getMainOutputTag(application),
        ParDoTranslation.getAdditionalOutputTags(application),
        ParDoTranslation.getSideInputs(application),
        ParDoTranslation.getSchemaInformation(application),
        ParDoTranslation.getSideInputMapping(application));
  } else {
    return application.getTransform();
  }
}
 

public static List<PCollectionView<?>> getSideInputs(AppliedPTransform<?, ?, ?> application)
    throws IOException {
  PTransform<?, ?> transform = application.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    return ((ParDo.MultiOutput<?, ?>) transform)
        .getSideInputs().values().stream().collect(Collectors.toList());
  }

  SdkComponents sdkComponents = SdkComponents.create(application.getPipeline().getOptions());
  RunnerApi.PTransform parDoProto = PTransformTranslation.toProto(application, sdkComponents);
  ParDoPayload payload = ParDoPayload.parseFrom(parDoProto.getSpec().getPayload());

  List<PCollectionView<?>> views = new ArrayList<>();
  RehydratedComponents components =
      RehydratedComponents.forComponents(sdkComponents.toComponents());
  for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) {
    String sideInputTag = sideInputEntry.getKey();
    RunnerApi.SideInput sideInput = sideInputEntry.getValue();
    PCollection<?> originalPCollection =
        checkNotNull(
            (PCollection<?>) application.getInputs().get(new TupleTag<>(sideInputTag)),
            "no input with tag %s",
            sideInputTag);
    views.add(
        PCollectionViewTranslation.viewFromProto(
            sideInput, sideInputTag, originalPCollection, parDoProto, components));
  }
  return views;
}
 

public static TupleTag<?> getMainOutputTag(AppliedPTransform<?, ?, ?> application)
    throws IOException {
  PTransform<?, ?> transform = application.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    return ((ParDo.MultiOutput<?, ?>) transform).getMainOutputTag();
  }

  return getMainOutputTag(getParDoPayload(application));
}
 

public static DoFn<?, ?> getDoFn(AppliedPTransform<?, ?, ?> application) throws IOException {
  PTransform<?, ?> transform = application.getTransform();
  if (transform instanceof ParDo.MultiOutput) {
    return ((ParDo.MultiOutput<?, ?>) transform).getFn();
  }

  return getDoFn(getParDoPayload(application));
}
 

@Override
public PTransformReplacement<PCollection<InputT>, PValue> getReplacementTransform(
    AppliedPTransform<PCollection<InputT>, PValue, PTransform<PCollection<InputT>, PValue>>
        transform) {
  Combine.GloballyAsSingletonView<?, ?> combineTransform =
      (Combine.GloballyAsSingletonView) transform.getTransform();
  return PTransformReplacement.of(
      PTransformReplacements.getSingletonMainInput(transform),
      new BatchViewOverrides.BatchViewAsSingleton(
          runner,
          findCreatePCollectionView(transform),
          (CombineFn) combineTransform.getCombineFn(),
          combineTransform.getFanout()));
}
 

/**
 * Gets all inputs of the {@link AppliedPTransform} that are not returned by {@link
 * PTransform#getAdditionalInputs()}.
 */
public static Collection<PValue> nonAdditionalInputs(AppliedPTransform<?, ?, ?> application) {
  ImmutableList.Builder<PValue> mainInputs = ImmutableList.builder();
  PTransform<?, ?> transform = application.getTransform();
  for (Map.Entry<TupleTag<?>, PValue> input : application.getInputs().entrySet()) {
    if (!transform.getAdditionalInputs().containsKey(input.getKey())) {
      mainInputs.add(input.getValue());
    }
  }
  checkArgument(
      !mainInputs.build().isEmpty() || application.getInputs().isEmpty(),
      "Expected at least one main input if any inputs exist");
  return mainInputs.build();
}
 

/**
 * A {@link PTransformMatcher} which matches a {@link Flatten.PCollections} which consumes no
 * input {@link PCollection PCollections}.
 */
public static PTransformMatcher emptyFlatten() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      return (application.getTransform() instanceof Flatten.PCollections)
          && application.getInputs().isEmpty();
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("EmptyFlattenMatcher").toString();
    }
  };
}
 

/**
 * A {@link PTransformMatcher} that matches a {@link ParDo.MultiOutput} containing a {@link DoFn}
 * that uses state or timers, as specified by {@link DoFnSignature#usesState()} and {@link
 * DoFnSignature#usesTimers()}.
 */
public static PTransformMatcher stateOrTimerParDoMulti() {
  return new PTransformMatcher() {
    @Override
    public boolean matches(AppliedPTransform<?, ?, ?> application) {
      PTransform<?, ?> transform = application.getTransform();
      if (transform instanceof ParDo.MultiOutput) {
        DoFn<?, ?> fn = ((ParDo.MultiOutput<?, ?>) transform).getFn();
        DoFnSignature signature = DoFnSignatures.signatureForDoFn(fn);
        return signature.usesState() || signature.usesTimers();
      }
      return false;
    }

    @Override
    public String toString() {
      return MoreObjects.toStringHelper("StateOrTimerParDoMultiMatcher").toString();
    }
  };
}
 

@Override
public PTransformReplacement<PBegin, PCollection<T>> getReplacementTransform(
    AppliedPTransform<PBegin, PCollection<T>, Create.Values<T>> transform) {
  Create.Values<T> original = transform.getTransform();
  PCollection<T> output =
      (PCollection) Iterables.getOnlyElement(transform.getOutputs().values());
  return PTransformReplacement.of(
      transform.getPipeline().begin(), new StreamingFnApiCreate<>(original, output));
}
 

static Map<TupleTag<?>, PValue> getOutputs(AppliedPTransform<?, ?, ?> appliedTransform) {
  if (appliedTransform.getTransform() == null) {
    return null;
  }
  return appliedTransform.getOutputs();
}