org.apache.beam.sdk.transforms.ViewFn Java Examples

/**
 * Call this constructor to initialize the fields for which this base class provides boilerplate
 * accessors.
 */
private SimplePCollectionView(
    PCollection<ElemT> pCollection,
    TupleTag<PrimitiveViewT> tag,
    ViewFn<PrimitiveViewT, ViewT> viewFn,
    WindowMappingFn<W> windowMappingFn,
    WindowingStrategy<?, W> windowingStrategy) {
  super(pCollection.getPipeline());
  this.pCollection = pCollection;
  if (windowingStrategy.getWindowFn() instanceof InvalidWindows) {
    throw new IllegalArgumentException("WindowFn of PCollectionView cannot be InvalidWindows");
  }
  this.windowMappingFn = windowMappingFn;
  this.tag = tag;
  this.windowingStrategy = windowingStrategy;
  this.viewFn = viewFn;
  this.coder = pCollection.getCoder();
}
 

@Nullable
@Override
public <T> T get(final PCollectionView<T> view, final BoundedWindow window) {
  // This gets called after isReady()
  final T sideInputData = (T) inMemorySideInputs.get(Pair.of(view, window));
  return sideInputData == null
    // The upstream gave us an empty sideInput
    ? ((ViewFn<Object, T>) view.getViewFn()).apply(new CreateViewTransform.MultiView<T>(Collections.emptyList()))
    // The upstream gave us a concrete sideInput
    : sideInputData;
}
 

public static <W extends BoundedWindow> SideInputSpec create(
    String accessPattern,
    Coder<?> coder,
    Coder<W> windowCoder,
    ViewFn<?, ?> viewFn,
    WindowMappingFn<W> windowMappingFn) {
  return new AutoValue_SideInputSpec<>(
      accessPattern, coder, windowCoder, viewFn, windowMappingFn);
}
 

/**
 * Call this constructor to initialize the fields for which this base class provides boilerplate
 * accessors, with an auto-generated tag.
 */
private SimplePCollectionView(
    PCollection<ElemT> pCollection,
    ViewFn<PrimitiveViewT, ViewT> viewFn,
    WindowMappingFn<W> windowMappingFn,
    WindowingStrategy<?, W> windowingStrategy) {
  this(pCollection, new TupleTag<>(), viewFn, windowMappingFn, windowingStrategy);
}
 

/** Create a new {@link RunnerPCollectionView} from the provided components. */
public RunnerPCollectionView(
    @Nullable PCollection<?> pCollection,
    TupleTag<Iterable<WindowedValue<?>>> tag,
    ViewFn<Iterable<WindowedValue<?>>, T> viewFn,
    WindowMappingFn<?> windowMappingFn,
    @Nullable WindowingStrategy<?, ?> windowingStrategy,
    @Nullable Coder<?> coder) {
  this.pCollection = pCollection;
  this.tag = tag;
  this.viewFn = viewFn;
  this.windowMappingFn = windowMappingFn;
  this.windowingStrategy = windowingStrategy;
  this.coder = coder;
}
 

@Test
public void createViewWithViewFn() {
  PCollection<Integer> input = p.apply(Create.of(1));
  PCollectionView<Iterable<Integer>> view = input.apply(View.asIterable());
  ViewFn<?, ?> viewFn = view.getViewFn();
  CreatePCollectionView<?, ?> createView = CreatePCollectionView.of(view);

  PTransformMatcher matcher = PTransformMatchers.createViewWithViewFn(viewFn.getClass());
  assertThat(matcher.matches(getAppliedTransform(createView)), is(true));
}
 

public static PTransformMatcher createViewWithViewFn(final Class<? extends ViewFn> viewFnType) {
  return application -> {
    if (!(application.getTransform() instanceof CreatePCollectionView)) {
      return false;
    }
    CreatePCollectionView<?, ?> createView =
        (CreatePCollectionView<?, ?>) application.getTransform();
    ViewFn<?, ?> viewFn = createView.getView().getViewFn();
    return viewFn.getClass().equals(viewFnType);
  };
}
 

/**
 * Converts a {@link org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec} into a {@link
 * ViewFn} using the URN.
 */
public static ViewFn<?, ?> viewFnFromProto(RunnerApi.FunctionSpec viewFn)
    throws InvalidProtocolBufferException {
  RunnerApi.FunctionSpec spec = viewFn;
  checkArgument(
      spec.getUrn().equals(ParDoTranslation.CUSTOM_JAVA_VIEW_FN_URN),
      "Can't deserialize unknown %s type %s",
      ViewFn.class.getSimpleName(),
      spec.getUrn());
  return (ViewFn<?, ?>)
      SerializableUtils.deserializeFromByteArray(
          spec.getPayload().toByteArray(), "Custom ViewFn");
}
 

/**
 * Create a {@link PCollectionView} from a side input spec and an already-deserialized {@link
 * PCollection} that should be wired up.
 */
public static PCollectionView<?> viewFromProto(
    RunnerApi.SideInput sideInput,
    String localName,
    PCollection<?> pCollection,
    RunnerApi.PTransform parDoTransform,
    RehydratedComponents components)
    throws IOException {
  checkArgument(
      localName != null,
      "%s.viewFromProto: localName must not be null",
      ParDoTranslation.class.getSimpleName());
  TupleTag<?> tag = new TupleTag<>(localName);
  WindowMappingFn<?> windowMappingFn = windowMappingFnFromProto(sideInput.getWindowMappingFn());
  ViewFn<?, ?> viewFn = viewFnFromProto(sideInput.getViewFn());
  WindowingStrategy<?, ?> windowingStrategy = pCollection.getWindowingStrategy().fixDefaults();

  PCollectionView<?> view =
      new RunnerPCollectionView<>(
          pCollection,
          (TupleTag) tag,
          (ViewFn) viewFn,
          windowMappingFn,
          windowingStrategy,
          (Coder) pCollection.getCoder());
  return view;
}
 

@Override
public ViewFn<?, MultimapView<K, V>> getViewFn() {
  return (ViewFn) PortabilityViewFn.INSTANCE;
}
 

/**
 * Constructor of CreateViewTransform.
 *
 * @param viewFn the viewFn that materializes data.
 */
public CreateViewTransform(final ViewFn<Materializations.MultimapView<Void, ?>, O> viewFn) {
  this.viewFn = viewFn;
  this.windowListMap = new HashMap<>();
  this.currentOutputWatermark = Long.MIN_VALUE;
}
 

@Override
public ViewFn<?, T> getViewFn() {
  return delegate.getViewFn();
}
 

private static LinkedHashMap<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>>
    getSideInputIdToPCollectionViewMap(
        RunnerApi.ExecutableStagePayload stagePayload, RunnerApi.Components components) {

  RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components);

  LinkedHashMap<RunnerApi.ExecutableStagePayload.SideInputId, PCollectionView<?>> sideInputs =
      new LinkedHashMap<>();
  // for PCollectionView compatibility, not used to transform materialization
  ViewFn<Iterable<WindowedValue<?>>, ?> viewFn =
      (ViewFn)
          new PCollectionViews.MultimapViewFn<>(
              (PCollectionViews.TypeDescriptorSupplier<Iterable<WindowedValue<Void>>>)
                  () -> TypeDescriptors.iterables(new TypeDescriptor<WindowedValue<Void>>() {}),
              (PCollectionViews.TypeDescriptorSupplier<Void>) TypeDescriptors::voids);

  for (RunnerApi.ExecutableStagePayload.SideInputId sideInputId :
      stagePayload.getSideInputsList()) {

    // TODO: local name is unique as long as only one transform with side input can be within a
    // stage
    String sideInputTag = sideInputId.getLocalName();
    String collectionId =
        components
            .getTransformsOrThrow(sideInputId.getTransformId())
            .getInputsOrThrow(sideInputId.getLocalName());
    RunnerApi.WindowingStrategy windowingStrategyProto =
        components.getWindowingStrategiesOrThrow(
            components.getPcollectionsOrThrow(collectionId).getWindowingStrategyId());

    final WindowingStrategy<?, ?> windowingStrategy;
    try {
      windowingStrategy =
          WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents);
    } catch (InvalidProtocolBufferException e) {
      throw new IllegalStateException(
          String.format(
              "Unable to hydrate side input windowing strategy %s.", windowingStrategyProto),
          e);
    }

    Coder<WindowedValue<Object>> coder = instantiateCoder(collectionId, components);
    // side input materialization via GBK (T -> Iterable<T>)
    WindowedValueCoder wvCoder = (WindowedValueCoder) coder;
    coder = wvCoder.withValueCoder(IterableCoder.of(wvCoder.getValueCoder()));

    sideInputs.put(
        sideInputId,
        new RunnerPCollectionView<>(
            null,
            new TupleTag<>(sideInputTag),
            viewFn,
            // TODO: support custom mapping fn
            windowingStrategy.getWindowFn().getDefaultWindowMappingFn(),
            windowingStrategy,
            coder));
  }
  return sideInputs;
}
 

public static FunctionSpec translateViewFn(ViewFn<?, ?> viewFn, SdkComponents components) {
  return FunctionSpec.newBuilder()
      .setUrn(CUSTOM_JAVA_VIEW_FN_URN)
      .setPayload(ByteString.copyFrom(SerializableUtils.serializeToByteArray(viewFn)))
      .build();
}
 

@Override
public ViewFn<Iterable<WindowedValue<?>>, T> getViewFn() {
  return viewFn;
}
 

@Override
public ViewFn<PrimitiveViewT, ViewT> getViewFn() {
  return viewFn;
}
 

/**
 * <b>For internal use only.</b>
 *
 * @deprecated this method will be removed entirely. The {@link ViewFn} for a side input is an
 *     attribute of the side input's specification with a {@link ParDo} transform, which will
 *     obtain this specification via a package-private channel.
 */
@Deprecated
@Internal
ViewFn<?, T> getViewFn();
 

abstract ViewFn<?, ?> getViewFn();