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

The following examples show how to use org.apache.beam.sdk.transforms.windowing.Window#Assign . These examples are extracted from open source projects. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
/**
 * @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));
}
 
Example 2
@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);
  }
}
 
Example 3
Source Project: beam   File: TransformTranslator.java    License: Apache License 2.0 6 votes vote down vote up
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>())";
    }
  };
}
 
Example 4
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>())";
    }
  };
}
 
Example 5
@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);
}
 
Example 6
@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);
}
 
Example 7
Source Project: beam   File: WindowIntoTranslation.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 8
Source Project: beam   File: WindowIntoTranslation.java    License: Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example 9
Source Project: beam   File: WindowingHelpers.java    License: Apache License 2.0 5 votes vote down vote up
/**
 * 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);
}
 
Example 10
Source Project: beam   File: WindowAssignTranslator.java    License: Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example 11
Source Project: beam   File: DataflowPipelineTranslator.java    License: Apache License 2.0 5 votes vote down vote up
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);
}
 
Example 12
Source Project: beam   File: WindowEvaluatorFactory.java    License: Apache License 2.0 5 votes vote down vote up
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);
}
 
Example 13
Source Project: beam   File: WindowEvaluatorFactory.java    License: Apache License 2.0 5 votes vote down vote up
@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;
}
 
Example 14
Source Project: beam   File: WindowIntoTranslation.java    License: Apache License 2.0 4 votes vote down vote up
public static WindowIntoPayload toProto(Window.Assign<?> transform, SdkComponents components) {
  return WindowIntoPayload.newBuilder()
      .setWindowFn(WindowingStrategyTranslation.toProto(transform.getWindowFn(), components))
      .build();
}
 
Example 15
Source Project: beam   File: WindowIntoTranslation.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public String getUrn(Window.Assign<?> transform) {
  return PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN;
}
 
Example 16
Source Project: beam   File: DataflowPipelineTranslator.java    License: Apache License 2.0 4 votes vote down vote up
@Override
public void translate(Window.Assign transform, TranslationContext context) {
  translateHelper(transform, context);
}
 
Example 17
Source Project: beam   File: TranslationUtils.java    License: Apache License 2.0 3 votes vote down vote up
/**
 * 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);
}