| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.runners.core.construction; |
| |
| import com.google.protobuf.ByteString; |
| import com.google.protobuf.InvalidProtocolBufferException; |
| import com.google.protobuf.util.Durations; |
| import com.google.protobuf.util.Timestamps; |
| import java.io.IOException; |
| import java.io.Serializable; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.Components; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.FunctionSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.OutputTime; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.SdkFunctionSpec; |
| import org.apache.beam.model.pipeline.v1.StandardWindowFns; |
| import org.apache.beam.model.pipeline.v1.StandardWindowFns.FixedWindowsPayload; |
| import org.apache.beam.model.pipeline.v1.StandardWindowFns.SessionsPayload; |
| import org.apache.beam.model.pipeline.v1.StandardWindowFns.SlidingWindowsPayload; |
| import org.apache.beam.sdk.transforms.windowing.FixedWindows; |
| import org.apache.beam.sdk.transforms.windowing.GlobalWindows; |
| import org.apache.beam.sdk.transforms.windowing.Sessions; |
| import org.apache.beam.sdk.transforms.windowing.SlidingWindows; |
| import org.apache.beam.sdk.transforms.windowing.TimestampCombiner; |
| import org.apache.beam.sdk.transforms.windowing.Trigger; |
| import org.apache.beam.sdk.transforms.windowing.Window.ClosingBehavior; |
| import org.apache.beam.sdk.transforms.windowing.Window.OnTimeBehavior; |
| import org.apache.beam.sdk.transforms.windowing.WindowFn; |
| import org.apache.beam.sdk.util.SerializableUtils; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.sdk.values.WindowingStrategy.AccumulationMode; |
| import org.joda.time.Duration; |
| |
| /** Utilities for working with {@link WindowingStrategy WindowingStrategies}. */ |
| public class WindowingStrategyTranslation implements Serializable { |
| |
| public static AccumulationMode fromProto(RunnerApi.AccumulationMode.Enum proto) { |
| switch (proto) { |
| case DISCARDING: |
| return AccumulationMode.DISCARDING_FIRED_PANES; |
| case ACCUMULATING: |
| return AccumulationMode.ACCUMULATING_FIRED_PANES; |
| case UNRECOGNIZED: |
| default: |
| // Whether or not it is proto that cannot recognize it (due to the version of the |
| // generated code we link to) or the switch hasn't been updated to handle it, |
| // the situation is the same: we don't know what this OutputTime means |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| RunnerApi.AccumulationMode.class.getCanonicalName(), |
| AccumulationMode.class.getCanonicalName(), |
| proto)); |
| } |
| } |
| |
| public static RunnerApi.AccumulationMode.Enum toProto(AccumulationMode accumulationMode) { |
| switch (accumulationMode) { |
| case DISCARDING_FIRED_PANES: |
| return RunnerApi.AccumulationMode.Enum.DISCARDING; |
| case ACCUMULATING_FIRED_PANES: |
| return RunnerApi.AccumulationMode.Enum.ACCUMULATING; |
| default: |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| AccumulationMode.class.getCanonicalName(), |
| RunnerApi.AccumulationMode.class.getCanonicalName(), |
| accumulationMode)); |
| } |
| } |
| |
| public static RunnerApi.ClosingBehavior.Enum toProto(ClosingBehavior closingBehavior) { |
| switch (closingBehavior) { |
| case FIRE_ALWAYS: |
| return RunnerApi.ClosingBehavior.Enum.EMIT_ALWAYS; |
| case FIRE_IF_NON_EMPTY: |
| return RunnerApi.ClosingBehavior.Enum.EMIT_IF_NONEMPTY; |
| default: |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| ClosingBehavior.class.getCanonicalName(), |
| RunnerApi.ClosingBehavior.class.getCanonicalName(), |
| closingBehavior)); |
| } |
| } |
| |
| public static ClosingBehavior fromProto(RunnerApi.ClosingBehavior.Enum proto) { |
| switch (proto) { |
| case EMIT_ALWAYS: |
| return ClosingBehavior.FIRE_ALWAYS; |
| case EMIT_IF_NONEMPTY: |
| return ClosingBehavior.FIRE_IF_NON_EMPTY; |
| case UNRECOGNIZED: |
| default: |
| // Whether or not it is proto that cannot recognize it (due to the version of the |
| // generated code we link to) or the switch hasn't been updated to handle it, |
| // the situation is the same: we don't know what this OutputTime means |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| RunnerApi.ClosingBehavior.class.getCanonicalName(), |
| ClosingBehavior.class.getCanonicalName(), |
| proto)); |
| } |
| } |
| |
| public static RunnerApi.OnTimeBehavior.Enum toProto(OnTimeBehavior onTimeBehavior) { |
| switch (onTimeBehavior) { |
| case FIRE_ALWAYS: |
| return RunnerApi.OnTimeBehavior.Enum.FIRE_ALWAYS; |
| case FIRE_IF_NON_EMPTY: |
| return RunnerApi.OnTimeBehavior.Enum.FIRE_IF_NONEMPTY; |
| default: |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| OnTimeBehavior.class.getCanonicalName(), |
| RunnerApi.OnTimeBehavior.class.getCanonicalName(), |
| onTimeBehavior)); |
| } |
| } |
| |
| public static OnTimeBehavior fromProto(RunnerApi.OnTimeBehavior.Enum proto) { |
| switch (proto) { |
| case FIRE_ALWAYS: |
| return OnTimeBehavior.FIRE_ALWAYS; |
| case FIRE_IF_NONEMPTY: |
| return OnTimeBehavior.FIRE_IF_NON_EMPTY; |
| case UNRECOGNIZED: |
| default: |
| // Whether or not it is proto that cannot recognize it (due to the version of the |
| // generated code we link to) or the switch hasn't been updated to handle it, |
| // the situation is the same: we don't know what this OutputTime means |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| RunnerApi.OnTimeBehavior.class.getCanonicalName(), |
| OnTimeBehavior.class.getCanonicalName(), |
| proto)); |
| } |
| } |
| |
| public static RunnerApi.OutputTime.Enum toProto(TimestampCombiner timestampCombiner) { |
| switch(timestampCombiner) { |
| case EARLIEST: |
| return OutputTime.Enum.EARLIEST_IN_PANE; |
| case END_OF_WINDOW: |
| return OutputTime.Enum.END_OF_WINDOW; |
| case LATEST: |
| return OutputTime.Enum.LATEST_IN_PANE; |
| default: |
| throw new IllegalArgumentException( |
| String.format( |
| "Unknown %s: %s", |
| TimestampCombiner.class.getSimpleName(), |
| timestampCombiner)); |
| } |
| } |
| |
| public static TimestampCombiner timestampCombinerFromProto(RunnerApi.OutputTime.Enum proto) { |
| switch (proto) { |
| case EARLIEST_IN_PANE: |
| return TimestampCombiner.EARLIEST; |
| case END_OF_WINDOW: |
| return TimestampCombiner.END_OF_WINDOW; |
| case LATEST_IN_PANE: |
| return TimestampCombiner.LATEST; |
| case UNRECOGNIZED: |
| default: |
| // Whether or not it is proto that cannot recognize it (due to the version of the |
| // generated code we link to) or the switch hasn't been updated to handle it, |
| // the situation is the same: we don't know what this OutputTime means |
| throw new IllegalArgumentException( |
| String.format( |
| "Cannot convert unknown %s to %s: %s", |
| RunnerApi.OutputTime.class.getCanonicalName(), |
| OutputTime.class.getCanonicalName(), |
| proto)); |
| } |
| } |
| |
| public static final String GLOBAL_WINDOWS_FN = "beam:windowfn:global_windows:v0.1"; |
| public static final String FIXED_WINDOWS_FN = "beam:windowfn:fixed_windows:v0.1"; |
| public static final String SLIDING_WINDOWS_FN = "beam:windowfn:sliding_windows:v0.1"; |
| public static final String SESSION_WINDOWS_FN = "beam:windowfn:session_windows:v0.1"; |
| // This URN says that the WindowFn is just a UDF blob the Java SDK understands |
| // TODO: standardize such things |
| public static final String SERIALIZED_JAVA_WINDOWFN_URN = "beam:windowfn:javasdk:v0.1"; |
| public static final String OLD_SERIALIZED_JAVA_WINDOWFN_URN = "urn:beam:windowfn:javasdk:0.1"; |
| // Remove this once the dataflow worker understands all the above formats. |
| private static final boolean USE_OLD_SERIALIZED_JAVA_WINDOWFN_URN = true; |
| |
| /** |
| * Converts a {@link WindowFn} into a {@link RunnerApi.MessageWithComponents} where {@link |
| * RunnerApi.MessageWithComponents#getFunctionSpec()} is a {@link RunnerApi.FunctionSpec} for the |
| * input {@link WindowFn}. |
| */ |
| public static SdkFunctionSpec toProto( |
| WindowFn<?, ?> windowFn, @SuppressWarnings("unused") SdkComponents components) { |
| // TODO: Set environment IDs |
| ByteString serializedFn = ByteString.copyFrom(SerializableUtils.serializeToByteArray(windowFn)); |
| if (USE_OLD_SERIALIZED_JAVA_WINDOWFN_URN) { |
| return SdkFunctionSpec.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(OLD_SERIALIZED_JAVA_WINDOWFN_URN) |
| .setPayload(serializedFn) |
| .build()) |
| .build(); |
| } else if (windowFn instanceof GlobalWindows) { |
| return SdkFunctionSpec.newBuilder() |
| .setSpec(FunctionSpec.newBuilder().setUrn(GLOBAL_WINDOWS_FN)) |
| .build(); |
| } else if (windowFn instanceof FixedWindows) { |
| FixedWindowsPayload fixedWindowsPayload = |
| FixedWindowsPayload.newBuilder() |
| .setSize(Durations.fromMillis(((FixedWindows) windowFn).getSize().getMillis())) |
| .setOffset(Timestamps.fromMillis(((FixedWindows) windowFn).getOffset().getMillis())) |
| .build(); |
| return SdkFunctionSpec.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(FIXED_WINDOWS_FN) |
| .setPayload(fixedWindowsPayload.toByteString())) |
| .build(); |
| } else if (windowFn instanceof SlidingWindows) { |
| SlidingWindowsPayload slidingWindowsPayload = SlidingWindowsPayload.newBuilder() |
| .setSize(Durations.fromMillis(((SlidingWindows) windowFn).getSize().getMillis())) |
| .setOffset(Timestamps.fromMillis(((SlidingWindows) windowFn).getOffset().getMillis())) |
| .setPeriod(Durations.fromMillis(((SlidingWindows) windowFn).getPeriod().getMillis())) |
| .build(); |
| return SdkFunctionSpec.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(SLIDING_WINDOWS_FN) |
| .setPayload(slidingWindowsPayload.toByteString())) |
| .build(); |
| } else if (windowFn instanceof Sessions) { |
| SessionsPayload sessionsPayload = |
| SessionsPayload.newBuilder() |
| .setGapSize(Durations.fromMillis(((Sessions) windowFn).getGapDuration().getMillis())) |
| .build(); |
| return SdkFunctionSpec.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(SESSION_WINDOWS_FN) |
| .setPayload(sessionsPayload.toByteString())) |
| .build(); |
| } else { |
| return SdkFunctionSpec.newBuilder() |
| .setSpec( |
| FunctionSpec.newBuilder() |
| .setUrn(SERIALIZED_JAVA_WINDOWFN_URN) |
| .setPayload(serializedFn)) |
| .build(); |
| } |
| } |
| |
| /** |
| * Converts a {@link WindowingStrategy} into a {@link RunnerApi.MessageWithComponents} where |
| * {@link RunnerApi.MessageWithComponents#getWindowingStrategy()} ()} is a {@link |
| * RunnerApi.WindowingStrategy RunnerApi.WindowingStrategy (proto)} for the input {@link |
| * WindowingStrategy}. |
| */ |
| public static RunnerApi.MessageWithComponents toProto(WindowingStrategy<?, ?> windowingStrategy) |
| throws IOException { |
| SdkComponents components = SdkComponents.create(); |
| RunnerApi.WindowingStrategy windowingStrategyProto = toProto(windowingStrategy, components); |
| |
| return RunnerApi.MessageWithComponents.newBuilder() |
| .setWindowingStrategy(windowingStrategyProto) |
| .setComponents(components.toComponents()) |
| .build(); |
| } |
| |
| /** |
| * Converts a {@link WindowingStrategy} into a {@link RunnerApi.WindowingStrategy}, registering |
| * any components in the provided {@link SdkComponents}. |
| */ |
| public static RunnerApi.WindowingStrategy toProto( |
| WindowingStrategy<?, ?> windowingStrategy, SdkComponents components) throws IOException { |
| SdkFunctionSpec windowFnSpec = toProto(windowingStrategy.getWindowFn(), components); |
| |
| RunnerApi.WindowingStrategy.Builder windowingStrategyProto = |
| RunnerApi.WindowingStrategy.newBuilder() |
| .setOutputTime(toProto(windowingStrategy.getTimestampCombiner())) |
| .setAccumulationMode(toProto(windowingStrategy.getMode())) |
| .setClosingBehavior(toProto(windowingStrategy.getClosingBehavior())) |
| .setAllowedLateness(windowingStrategy.getAllowedLateness().getMillis()) |
| .setTrigger(TriggerTranslation.toProto(windowingStrategy.getTrigger())) |
| .setWindowFn(windowFnSpec) |
| .setAssignsToOneWindow(windowingStrategy.getWindowFn().assignsToOneWindow()) |
| .setOnTimeBehavior(toProto(windowingStrategy.getOnTimeBehavior())) |
| .setWindowCoderId( |
| components.registerCoder(windowingStrategy.getWindowFn().windowCoder())); |
| |
| return windowingStrategyProto.build(); |
| } |
| |
| /** |
| * Converts from a {@link RunnerApi.WindowingStrategy} accompanied by {@link Components} |
| * to the SDK's {@link WindowingStrategy}. |
| */ |
| public static WindowingStrategy<?, ?> fromProto(RunnerApi.MessageWithComponents proto) |
| throws InvalidProtocolBufferException { |
| switch (proto.getRootCase()) { |
| case WINDOWING_STRATEGY: |
| return fromProto( |
| proto.getWindowingStrategy(), |
| RehydratedComponents.forComponents(proto.getComponents())); |
| default: |
| throw new IllegalArgumentException( |
| String.format( |
| "Expected a %s with components but received %s", |
| RunnerApi.WindowingStrategy.class.getCanonicalName(), proto)); |
| } |
| } |
| |
| /** |
| * Converts from {@link RunnerApi.WindowingStrategy} to the SDK's {@link WindowingStrategy} using |
| * the provided components to dereferences identifiers found in the proto. |
| */ |
| public static WindowingStrategy<?, ?> fromProto( |
| RunnerApi.WindowingStrategy proto, RehydratedComponents components) |
| throws InvalidProtocolBufferException { |
| |
| SdkFunctionSpec windowFnSpec = proto.getWindowFn(); |
| WindowFn<?, ?> windowFn = windowFnFromProto(windowFnSpec); |
| TimestampCombiner timestampCombiner = timestampCombinerFromProto(proto.getOutputTime()); |
| AccumulationMode accumulationMode = fromProto(proto.getAccumulationMode()); |
| Trigger trigger = TriggerTranslation.fromProto(proto.getTrigger()); |
| ClosingBehavior closingBehavior = fromProto(proto.getClosingBehavior()); |
| Duration allowedLateness = Duration.millis(proto.getAllowedLateness()); |
| OnTimeBehavior onTimeBehavior = fromProto(proto.getOnTimeBehavior()); |
| |
| return WindowingStrategy.of(windowFn) |
| .withAllowedLateness(allowedLateness) |
| .withMode(accumulationMode) |
| .withTrigger(trigger) |
| .withTimestampCombiner(timestampCombiner) |
| .withClosingBehavior(closingBehavior) |
| .withOnTimeBehavior(onTimeBehavior); |
| } |
| |
| public static WindowFn<?, ?> windowFnFromProto(SdkFunctionSpec windowFnSpec) { |
| try { |
| switch (windowFnSpec.getSpec().getUrn()) { |
| case GLOBAL_WINDOWS_FN: |
| return new GlobalWindows(); |
| case FIXED_WINDOWS_FN: |
| StandardWindowFns.FixedWindowsPayload fixedParams = null; |
| fixedParams = |
| StandardWindowFns.FixedWindowsPayload.parseFrom( |
| windowFnSpec.getSpec().getPayload()); |
| return FixedWindows.of(Duration.millis(Durations.toMillis(fixedParams.getSize()))) |
| .withOffset(Duration.millis(Timestamps.toMillis(fixedParams.getOffset()))); |
| case SLIDING_WINDOWS_FN: |
| StandardWindowFns.SlidingWindowsPayload slidingParams = |
| StandardWindowFns.SlidingWindowsPayload.parseFrom( |
| windowFnSpec.getSpec().getPayload()); |
| return SlidingWindows.of(Duration.millis(Durations.toMillis(slidingParams.getSize()))) |
| .every(Duration.millis(Durations.toMillis(slidingParams.getPeriod()))) |
| .withOffset(Duration.millis(Timestamps.toMillis(slidingParams.getOffset()))); |
| case SESSION_WINDOWS_FN: |
| StandardWindowFns.SessionsPayload sessionParams = |
| StandardWindowFns.SessionsPayload.parseFrom(windowFnSpec.getSpec().getPayload()); |
| return Sessions.withGapDuration( |
| Duration.millis(Durations.toMillis(sessionParams.getGapSize()))); |
| case SERIALIZED_JAVA_WINDOWFN_URN: |
| case OLD_SERIALIZED_JAVA_WINDOWFN_URN: |
| return (WindowFn<?, ?>) |
| SerializableUtils.deserializeFromByteArray( |
| windowFnSpec.getSpec().getPayload().toByteArray(), "WindowFn"); |
| default: |
| throw new IllegalArgumentException( |
| "Unknown or unsupported WindowFn: " + windowFnSpec.getSpec().getUrn()); |
| } |
| } catch (InvalidProtocolBufferException e) { |
| throw new IllegalArgumentException( |
| String.format( |
| "%s for %s with URN %s did not contain expected proto message for payload", |
| FunctionSpec.class.getSimpleName(), |
| WindowFn.class.getSimpleName(), |
| windowFnSpec.getSpec().getUrn()), |
| e); |
| } |
| } |
| } |
