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

private SideInputHandler createHandler(SideInputSpec<?, ?> cacheKey) {
  switch (cacheKey.accessPattern().getUrn()) {
    case Materializations.ITERABLE_MATERIALIZATION_URN:
      return sideInputHandlerFactory.forIterableSideInput(
          cacheKey.transformId(),
          cacheKey.sideInputId(),
          cacheKey.elementCoder(),
          cacheKey.windowCoder());

    case Materializations.MULTIMAP_MATERIALIZATION_URN:
      return sideInputHandlerFactory.forMultimapSideInput(
          cacheKey.transformId(),
          cacheKey.sideInputId(),
          (KvCoder) cacheKey.elementCoder(),
          cacheKey.windowCoder());

    default:
      throw new IllegalStateException(
          String.format("Unsupported access pattern for side input %s", cacheKey));
  }
}
 

public Twister2SideInputReader(Map<PCollectionView<?>, WindowingStrategy<?, ?>> indexByView,
                               TSetContext context) {
  this.sideInputs = new HashMap<>();

  for (PCollectionView<?> view : indexByView.keySet()) {
    checkArgument(
        Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
            view.getViewFn().getMaterialization().getUrn()),
        "This handler is only capable of dealing with %s materializations "
            + "but was asked to handle %s for PCollectionView with tag %s.",
        Materializations.MULTIMAP_MATERIALIZATION_URN,
        view.getViewFn().getMaterialization().getUrn(),
        view.getTagInternal().getId());
  }
  for (Map.Entry<PCollectionView<?>, WindowingStrategy<?, ?>> entry : indexByView.entrySet()) {
    sideInputs.put(entry.getKey().getTagInternal(), entry.getValue());
  }
  this.runtimeContext = context;
}
 

@Override
public Integer apply(final Materializations.MultimapView<Void, String> view) {
  int sum = 0;
  // MultimapView.get is Nullable
  for (String s : view.get(null)) {
    sum += 1;
  }
  return sum;
}
 

@Test
public void testMaterializationUrn() {
  assertEquals(
      Materializations.MULTIMAP_MATERIALIZATION_URN,
      DataflowPortabilityPCollectionView.with(TAG, CODER)
          .getViewFn()
          .getMaterialization()
          .getUrn());
}
 

/**
 * Returns an artificial PCollectionView that can be used to fulfill API requirements of a {@link
 * SideInputReader} when used inside the Dataflow runner harness.
 *
 * <p>Generates length prefixed coder variants suitable to be used within the Dataflow Runner
 * harness so that encoding and decoding values matches the length prefixing that occurred when
 * materializing the side input.
 */
public static final PCollectionView<?> transformSideInputForRunner(
    RunnerApi.Pipeline pipeline,
    RunnerApi.PTransform parDoPTransform,
    String sideInputTag,
    RunnerApi.SideInput sideInput) {
  checkArgument(
      Materializations.MULTIMAP_MATERIALIZATION_URN.equals(sideInput.getAccessPattern().getUrn()),
      "This handler is only capable of dealing with %s materializations "
          + "but was asked to handle %s for PCollectionView with tag %s.",
      Materializations.MULTIMAP_MATERIALIZATION_URN,
      sideInput.getAccessPattern().getUrn(),
      sideInputTag);
  String sideInputPCollectionId = parDoPTransform.getInputsOrThrow(sideInputTag);
  RunnerApi.PCollection sideInputPCollection =
      pipeline.getComponents().getPcollectionsOrThrow(sideInputPCollectionId);
  try {
    FullWindowedValueCoder<KV<Object, Object>> runnerSideInputCoder =
        (FullWindowedValueCoder)
            WireCoders.instantiateRunnerWireCoder(
                PipelineNode.pCollection(sideInputPCollectionId, sideInputPCollection),
                pipeline.getComponents());

    return DataflowPortabilityPCollectionView.with(
        new TupleTag<>(sideInputTag), runnerSideInputCoder);
  } catch (IOException e) {
    throw new IllegalStateException("Unable to translate proto to coder", e);
  }
}
 

private Iterable<PCollectionView<?>> buildPCollectionViewsWithSdkSupportedWindowMappingFn(
    IdGenerator idGenerator,
    InstructionRequestHandler instructionRequestHandler,
    FnDataService beamFnDataService,
    ApiServiceDescriptor dataServiceApiServiceDescriptor,
    Coder<BoundedWindow> mainInputWindowCoder,
    Map<PCollectionView<?>, RunnerApi.FunctionSpec> pCollectionViewsToWindowMappingFns) {
  ImmutableList.Builder<PCollectionView<?>> wrappedViews = ImmutableList.builder();
  for (Map.Entry<PCollectionView<?>, RunnerApi.FunctionSpec> entry :
      pCollectionViewsToWindowMappingFns.entrySet()) {
    WindowMappingFn windowMappingFn =
        new FnApiWindowMappingFn(
            idGenerator,
            instructionRequestHandler,
            dataServiceApiServiceDescriptor,
            beamFnDataService,
            entry.getValue(),
            mainInputWindowCoder,
            entry.getKey().getWindowingStrategyInternal().getWindowFn().windowCoder());
    wrappedViews.add(
        new ForwardingPCollectionView<Materializations.MultimapView>(
            (PCollectionView) entry.getKey()) {
          @Override
          public WindowMappingFn<?> getWindowMappingFn() {
            return windowMappingFn;
          }
        });
  }
  return wrappedViews.build();
}
 

@Override
public Iterable<V> get(K key, W window) {
  Materializations.MultimapView<K, V> sideInput =
      (Materializations.MultimapView<K, V>)
          sideInputReader.get(view, (BoundedWindow) windowCoder.structuralValue(window));

  return sideInput.get(key);
}
 

@Override
public Iterable<K> get(W window) {
  Materializations.MultimapView<K, V> sideInput =
      (Materializations.MultimapView<K, V>)
          sideInputReader.get(view, (BoundedWindow) windowCoder.structuralValue(window));
  return sideInput.get();
}
 

private DataflowMultimapSideInputHandler(
    SideInputReader sideInputReader,
    PCollectionView<Materializations.MultimapView<Object, Object>> view,
    Coder<K> keyCoder,
    Coder<V> valueCoder,
    Coder<W> windowCoder) {
  this.sideInputReader = sideInputReader;
  this.view = view;
  this.keyCoder = keyCoder;
  this.valueCoder = valueCoder;
  this.windowCoder = windowCoder;
}
 

@Override
public Materialization<MultimapView<Void, KV<K, V>>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<MultimapView<K, V>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public <K, V, W extends BoundedWindow> MultimapSideInputHandler<K, V, W> forMultimapSideInput(
    String pTransformId, String sideInputId, KvCoder<K, V> elementCoder, Coder<W> windowCoder) {
  checkArgument(
      pTransformId != null && pTransformId.length() > 0, "Expect a valid PTransform ID.");

  SideInputReader sideInputReader = ptransformIdToSideInputReader.get(pTransformId);
  checkState(sideInputReader != null, String.format("Unknown PTransform '%s'", pTransformId));

  PCollectionView<Materializations.MultimapView<Object, Object>> view =
      (PCollectionView<Materializations.MultimapView<Object, Object>>)
          sideInputIdToPCollectionViewMap.get(
              RunnerApi.ExecutableStagePayload.SideInputId.newBuilder()
                  .setTransformId(pTransformId)
                  .setLocalName(sideInputId)
                  .build());
  checkState(
      view != null,
      String.format("Unknown side input '%s' on PTransform '%s'", sideInputId, pTransformId));

  checkState(
      Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
          view.getViewFn().getMaterialization().getUrn()),
      String.format(
          "Unknown materialization for side input '%s' on PTransform '%s' with urn '%s'",
          sideInputId, pTransformId, view.getViewFn().getMaterialization().getUrn()));

  checkState(
      view.getCoderInternal() instanceof KvCoder,
      String.format(
          "Materialization of side input '%s' on PTransform '%s' expects %s but received %s.",
          sideInputId,
          pTransformId,
          KvCoder.class.getSimpleName(),
          view.getCoderInternal().getClass().getSimpleName()));

  KvCoder<K, V> kvCoder = elementCoder;

  return new DataflowMultimapSideInputHandler<>(
      sideInputReader, view, kvCoder.getKeyCoder(), kvCoder.getValueCoder(), windowCoder);
}
 

/**
 * Adds an output with the given name to the previously added Dataflow step, producing the
 * specified output {@code PValue} with the given {@code Coder} (if not {@code null}).
 */
private void addOutput(String name, PValue value, Coder<?> valueCoder) {
  translator.registerOutputName(value, name);

  Map<String, Object> properties = getProperties();
  @Nullable List<Map<String, Object>> outputInfoList = null;
  try {
    // TODO: This should be done via a Structs accessor.
    outputInfoList = (List<Map<String, Object>>) properties.get(PropertyNames.OUTPUT_INFO);
  } catch (Exception e) {
    throw new RuntimeException("Inconsistent dataflow pipeline translation", e);
  }
  if (outputInfoList == null) {
    outputInfoList = new ArrayList<>();
    // TODO: This should be done via a Structs accessor.
    properties.put(PropertyNames.OUTPUT_INFO, outputInfoList);
  }

  Map<String, Object> outputInfo = new HashMap<>();
  addString(outputInfo, PropertyNames.OUTPUT_NAME, name);

  String stepName = getString(properties, PropertyNames.USER_NAME);
  String generatedName = String.format("%s.out%d", stepName, outputInfoList.size());

  addString(outputInfo, PropertyNames.USER_NAME, generatedName);
  if ((value instanceof PCollection
          && translator.runner.doesPCollectionRequireIndexedFormat((PCollection<?>) value))
      || ((value instanceof PCollectionView)
          && (Materializations.MULTIMAP_MATERIALIZATION_URN.equals(
              ((PCollectionView) value).getViewFn().getMaterialization().getUrn())))) {
    addBoolean(outputInfo, PropertyNames.USE_INDEXED_FORMAT, true);
  }
  if (valueCoder != null) {
    // Verify that encoding can be decoded, in order to catch serialization
    // failures as early as possible.
    CloudObject encoding = translateCoder(valueCoder, translator);
    addObject(outputInfo, PropertyNames.ENCODING, encoding);
    translator.outputCoders.put(value, valueCoder);
  }

  outputInfoList.add(outputInfo);
}
 

@Override
public Materialization<MultimapView<Void, KV<K, V>>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<MultimapView<K, V>> getMaterialization() {
  return Materializations.multimap();
}
 

/**
 * 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 Materialization<MultimapView<K, V>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<MultimapView<Void, T>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<MultimapView<Long, ValueOrMetadata<T, OffsetRange>>>
    getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<MultimapView<Void, T>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<IterableView<T>> getMaterialization() {
  return Materializations.iterable();
}
 

@Override
public Materialization<MultimapView<Void, T>> getMaterialization() {
  return Materializations.multimap();
}
 

@Override
public Materialization<IterableView<T>> getMaterialization() {
  return Materializations.iterable();
}
 

@Override
@Nullable
public <T> T get(PCollectionView<T> view, BoundedWindow window) {
  TupleTag<?> tag = view.getTagInternal();

  SideInputSpec sideInputSpec = sideInputSpecMap.get(tag);
  checkArgument(sideInputSpec != null, "Attempting to access unknown side input %s.", view);

  ByteString.Output encodedWindowOut = ByteString.newOutput();
  try {
    sideInputSpec
        .getWindowCoder()
        .encode(sideInputSpec.getWindowMappingFn().getSideInputWindow(window), encodedWindowOut);
  } catch (IOException e) {
    throw new IllegalStateException(e);
  }
  ByteString encodedWindow = encodedWindowOut.toByteString();
  StateKey.Builder cacheKeyBuilder = StateKey.newBuilder();
  Object sideInputAccessor;

  switch (sideInputSpec.getAccessPattern()) {
    case Materializations.ITERABLE_MATERIALIZATION_URN:
      cacheKeyBuilder
          .getIterableSideInputBuilder()
          .setTransformId(ptransformId)
          .setSideInputId(tag.getId())
          .setWindow(encodedWindow);
      sideInputAccessor =
          new IterableSideInput<>(
              beamFnStateClient,
              processBundleInstructionId.get(),
              ptransformId,
              tag.getId(),
              encodedWindow,
              sideInputSpec.getCoder());
      break;

    case Materializations.MULTIMAP_MATERIALIZATION_URN:
      checkState(
          sideInputSpec.getCoder() instanceof KvCoder,
          "Expected %s but received %s.",
          KvCoder.class,
          sideInputSpec.getCoder().getClass());
      KvCoder<?, ?> kvCoder = (KvCoder) sideInputSpec.getCoder();
      cacheKeyBuilder
          .getMultimapSideInputBuilder()
          .setTransformId(ptransformId)
          .setSideInputId(tag.getId())
          .setWindow(encodedWindow);
      sideInputAccessor =
          new MultimapSideInput<>(
              beamFnStateClient,
              processBundleInstructionId.get(),
              ptransformId,
              tag.getId(),
              encodedWindow,
              kvCoder.getKeyCoder(),
              kvCoder.getValueCoder());
      break;

    default:
      throw new IllegalStateException(
          String.format(
              "This SDK is only capable of dealing with %s materializations "
                  + "but was asked to handle %s for PCollectionView with tag %s.",
              ImmutableList.of(
                  Materializations.ITERABLE_MATERIALIZATION_URN,
                  Materializations.MULTIMAP_MATERIALIZATION_URN),
              sideInputSpec.getAccessPattern(),
              tag));
  }

  return (T)
      stateKeyObjectCache.computeIfAbsent(
          cacheKeyBuilder.build(), key -> sideInputSpec.getViewFn().apply(sideInputAccessor));
}
 

@Override
public Materialization<Materializations.MultimapView<Void, String>> getMaterialization() {
  throw new UnsupportedOperationException();
}