Java Code Examples for org.apache.beam.sdk.transforms.windowing.Window#Assign

/**
 * @param ctx       provides translation context
 * @param beamNode  the beam node to be translated
 * @param transform transform which can be obtained from {@code beamNode}
 */
@PrimitiveTransformTranslator({Window.class, Window.Assign.class})
private static void windowTranslator(final PipelineTranslationContext ctx,
                                     final TransformHierarchy.Node beamNode,
                                     final PTransform<?, ?> transform) {
  final WindowFn windowFn;
  if (transform instanceof Window) {
    windowFn = ((Window) transform).getWindowFn();
  } else if (transform instanceof Window.Assign) {
    windowFn = ((Window.Assign) transform).getWindowFn();
  } else {
    throw new UnsupportedOperationException(String.format("%s is not supported", transform));
  }
  final IRVertex vertex = new OperatorVertex(
    new WindowFnTransform(windowFn, DisplayData.from(beamNode.getTransform())));
  ctx.addVertex(vertex);
  beamNode.getInputs().values().forEach(input -> ctx.addEdgeTo(vertex, input));
  beamNode.getOutputs().values().forEach(output -> ctx.registerMainOutputFrom(beamNode, vertex, output));
}
 

private static <T, W extends BoundedWindow> TransformEvaluator<Window.Assign<T>> window() {
  return new TransformEvaluator<Window.Assign<T>>() {
    @Override
    public void evaluate(final Window.Assign<T> transform, EvaluationContext context) {
      @SuppressWarnings("unchecked")
      UnboundedDataset<T> unboundedDataset =
          (UnboundedDataset<T>) context.borrowDataset(transform);
      JavaDStream<WindowedValue<T>> dStream = unboundedDataset.getDStream();
      JavaDStream<WindowedValue<T>> outputStream;
      if (TranslationUtils.skipAssignWindows(transform, context)) {
        // do nothing.
        outputStream = dStream;
      } else {
        outputStream =
            dStream.transform(rdd -> rdd.map(new SparkAssignWindowFn<>(transform.getWindowFn())));
      }
      context.putDataset(
          transform, new UnboundedDataset<>(outputStream, unboundedDataset.getStreamSources()));
    }

    @Override
    public String toNativeString() {
      return "map(new <windowFn>())";
    }
  };
}
 

private static <T, W extends BoundedWindow> TransformEvaluator<Window.Assign<T>> window() {
  return new TransformEvaluator<Window.Assign<T>>() {
    @Override
    public void evaluate(Window.Assign<T> transform, EvaluationContext context) {
      @SuppressWarnings("unchecked")
      JavaRDD<WindowedValue<T>> inRDD =
          ((BoundedDataset<T>) context.borrowDataset(transform)).getRDD();

      if (TranslationUtils.skipAssignWindows(transform, context)) {
        context.putDataset(transform, new BoundedDataset<>(inRDD));
      } else {
        context.putDataset(
            transform,
            new BoundedDataset<>(inRDD.map(new SparkAssignWindowFn<>(transform.getWindowFn()))));
      }
    }

    @Override
    public String toNativeString() {
      return "map(new <windowFn>())";
    }
  };
}
 

@Override
public void translateTransform(
    PTransform<PCollection<T>, PCollection<T>> transform, TranslationContext context) {

  Window.Assign<T> assignTransform = (Window.Assign<T>) transform;
  @SuppressWarnings("unchecked")
  final PCollection<T> input = (PCollection<T>) context.getInput();
  @SuppressWarnings("unchecked")
  final PCollection<T> output = (PCollection<T>) context.getOutput();

  Dataset<WindowedValue<T>> inputDataset = context.getDataset(input);
  if (WindowingHelpers.skipAssignWindows(assignTransform, context)) {
    context.putDataset(output, inputDataset);
  } else {
    WindowFn<T, ?> windowFn = assignTransform.getWindowFn();
    WindowedValue.FullWindowedValueCoder<T> windowedValueCoder =
        WindowedValue.FullWindowedValueCoder.of(input.getCoder(), windowFn.windowCoder());
    Dataset<WindowedValue<T>> outputDataset =
        inputDataset.map(
            WindowingHelpers.assignWindowsMapFunction(windowFn),
            EncoderHelpers.fromBeamCoder(windowedValueCoder));
    context.putDataset(output, outputDataset);
  }
}
 

@SuppressWarnings("unchecked")
public WindowIntoEvaluator(
    AppliedPTransform<PCollection<InputT>, PCollection<InputT>, Window.Assign<InputT>>
        transform,
    WindowFn<? super InputT, ?> windowFn,
    UncommittedBundle<InputT> outputBundle) {
  this.outputBundle = outputBundle;
  this.transform = transform;
  // Safe contravariant cast
  this.windowFn = (WindowFn<InputT, ?>) windowFn;
}
 

private <InputT> TransformEvaluator<InputT> createTransformEvaluator(
    AppliedPTransform<PCollection<InputT>, PCollection<InputT>, Window.Assign<InputT>>
        transform) {

  WindowFn<? super InputT, ?> fn = (WindowFn) WindowIntoTranslation.getWindowFn(transform);

  UncommittedBundle<InputT> outputBundle =
      evaluationContext.createBundle(
          (PCollection<InputT>) Iterables.getOnlyElement(transform.getOutputs().values()));
  if (fn == null) {
    return PassthroughTransformEvaluator.create(transform, outputBundle);
  }
  return new WindowIntoEvaluator<>(transform, fn, outputBundle);
}
 

private <T> void translateHelper(Window.Assign<T> transform, TranslationContext context) {
  StepTranslationContext stepContext = context.addStep(transform, "Bucket");
  PCollection<T> input = context.getInput(transform);
  stepContext.addInput(PropertyNames.PARALLEL_INPUT, input);
  stepContext.addOutput(PropertyNames.OUTPUT, context.getOutput(transform));

  WindowingStrategy<?, ?> strategy = context.getOutput(transform).getWindowingStrategy();
  byte[] serializedBytes =
      serializeWindowingStrategy(strategy, context.getPipelineOptions());
  String serializedJson = byteArrayToJsonString(serializedBytes);
  stepContext.addInput(PropertyNames.SERIALIZED_FN, serializedJson);
}
 

@Override
public void translate(
    Window.Assign<T> transform, TransformHierarchy.Node node, TranslationContext ctx) {
  final PCollection<T> output = ctx.getOutput(transform);

  @SuppressWarnings("unchecked")
  final WindowFn<T, ?> windowFn = (WindowFn<T, ?>) output.getWindowingStrategy().getWindowFn();

  final MessageStream<OpMessage<T>> inputStream = ctx.getMessageStream(ctx.getInput(transform));

  final MessageStream<OpMessage<T>> outputStream =
      inputStream.flatMap(OpAdapter.adapt(new WindowAssignOp<>(windowFn)));

  ctx.registerMessageStream(output, outputStream);
}
 

/**
 * Checks if the window transformation should be applied or skipped.
 *
 * <p>Avoid running assign windows if both source and destination are global window or if the user
 * has not specified the WindowFn (meaning they are just messing with triggering or allowed
 * lateness).
 */
@SuppressWarnings("unchecked")
public static <T, W extends BoundedWindow> boolean skipAssignWindows(
    Window.Assign<T> transform, TranslationContext context) {
  WindowFn<? super T, W> windowFnToApply = (WindowFn<? super T, W>) transform.getWindowFn();
  PCollection<T> input = (PCollection<T>) context.getInput();
  WindowFn<?, ?> windowFnOfInput = input.getWindowingStrategy().getWindowFn();
  return windowFnToApply == null
      || (windowFnOfInput instanceof GlobalWindows && windowFnToApply instanceof GlobalWindows);
}
 

@Override
public FunctionSpec translate(
    AppliedPTransform<?, ?, Window.Assign<?>> transform, SdkComponents components) {
  WindowIntoPayload payload = toProto(transform.getTransform(), components);
  return RunnerApi.FunctionSpec.newBuilder()
      .setUrn(getUrn(transform.getTransform()))
      .setPayload(payload.toByteString())
      .build();
}
 

@Override
public FunctionSpec translate(
    AppliedPTransform<?, ?, Window.Assign<?>> transform, SdkComponents components) {
  return FunctionSpec.newBuilder()
      .setUrn("beam:transform:window:v1")
      .setPayload(
          WindowIntoTranslation.toProto(transform.getTransform(), components).toByteString())
      .build();
}
 

@Override
public void translateNode(Window.Assign<T> transform, Twister2BatchTranslationContext context) {
  BatchTSetImpl<WindowedValue<T>> inputTTSet =
      context.getInputDataSet(context.getInput(transform));

  final WindowingStrategy<T, BoundedWindow> windowingStrategy =
      (WindowingStrategy<T, BoundedWindow>) context.getOutput(transform).getWindowingStrategy();

  WindowFn<T, BoundedWindow> windowFn = windowingStrategy.getWindowFn();
  ComputeTSet<WindowedValue<T>, Iterator<WindowedValue<T>>> outputTSet =
      inputTTSet.direct().compute(new AssignWindowsFunction(windowFn, context.getOptions()));
  context.setOutputDataSet(context.getOutput(transform), outputTSet);
}
 

@Override
public void translateNode(Window.Assign<T> transform, Twister2BatchTranslationContext context) {
  BatchTSetImpl<WindowedValue<T>> inputTTSet = context.getInputDataSet(context.getInput(
      transform));

  final WindowingStrategy<T, BoundedWindow> windowingStrategy =
      (WindowingStrategy<T, BoundedWindow>) context.getOutput(transform).getWindowingStrategy();

  WindowFn<T, BoundedWindow> windowFn = windowingStrategy.getWindowFn();
  ComputeTSet<WindowedValue<T>, Iterator<WindowedValue<T>>> outputTSet =
      inputTTSet.direct().compute(new AssignWindowsFunction(windowFn));
  context.setOutputDataSet(context.getOutput(transform), outputTSet);
}
 

@Override
public String getUrn(Window.Assign<?> transform) {
  return PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN;
}
 

public static WindowIntoPayload toProto(Window.Assign<?> transform, SdkComponents components) {
  return WindowIntoPayload.newBuilder()
      .setWindowFn(WindowingStrategyTranslation.toProto(transform.getWindowFn(), components))
      .build();
}
 

@Override
public void translate(Window.Assign transform, TranslationContext context) {
  translateHelper(transform, context);
}
 

/**
 * Checks if the window transformation should be applied or skipped.
 *
 * <p>Avoid running assign windows if both source and destination are global window or if the user
 * has not specified the WindowFn (meaning they are just messing with triggering or allowed
 * lateness).
 *
 * @param transform The {@link Window.Assign} transformation.
 * @param context The {@link EvaluationContext}.
 * @param <T> PCollection type.
 * @param <W> {@link BoundedWindow} type.
 * @return if to apply the transformation.
 */
public static <T, W extends BoundedWindow> boolean skipAssignWindows(
    Window.Assign<T> transform, EvaluationContext context) {
  @SuppressWarnings("unchecked")
  WindowFn<? super T, W> windowFn = (WindowFn<? super T, W>) transform.getWindowFn();
  return windowFn == null
      || (context.getInput(transform).getWindowingStrategy().getWindowFn()
              instanceof GlobalWindows
          && windowFn instanceof GlobalWindows);
}