@Override public WindowingStrategy<?, ?> load(String id) throws Exception { @Nullable RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrDefault(id, null); checkState( windowingStrategyProto != null, "No WindowingStrategy with id '%s' in serialized components", id); return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.this); } });
/** * 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); }
private static byte[] serializeWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) { try { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.JAVA_SDK_HARNESS_ENVIRONMENT); return WindowingStrategyTranslation.toMessageProto(windowingStrategy, sdkComponents) .toByteArray(); } catch (Exception e) { throw new RuntimeException( String.format("Unable to format windowing strategy %s as bytes", windowingStrategy), e); } }
public static @Nullable WindowFn<?, ?> getWindowFn(AppliedPTransform<?, ?, ?> application) { return WindowingStrategyTranslation.windowFnFromProto( getWindowIntoPayload(application).getWindowFn()); }
@Test public void testToProtoAndBack() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy<?, ?> toProtoAndBackWindowingStrategy = WindowingStrategyTranslation.fromProto( WindowingStrategyTranslation.toMessageProto(windowingStrategy, components)); assertThat( toProtoAndBackWindowingStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); }
@Test public void testToProtoAndBackWithComponents() throws Exception { WindowingStrategy<?, ?> windowingStrategy = toProtoAndBackSpec.getWindowingStrategy(); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.WindowingStrategy proto = WindowingStrategyTranslation.toProto(windowingStrategy, components); RehydratedComponents protoComponents = RehydratedComponents.forComponents(components.toComponents()); assertThat( WindowingStrategyTranslation.fromProto(proto, protoComponents).fixDefaults(), equalTo(windowingStrategy.fixDefaults())); protoComponents.getCoder( components.registerCoder(windowingStrategy.getWindowFn().windowCoder())); assertThat( proto.getAssignsToOneWindow(), equalTo(windowingStrategy.getWindowFn().assignsToOneWindow())); } }
public static WindowIntoPayload toProto(Window.Assign<?> transform, SdkComponents components) { return WindowIntoPayload.newBuilder() .setWindowFn(WindowingStrategyTranslation.toProto(transform.getWindowFn(), components)) .build(); }
private <T> void translateAssignWindows( String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext context) { RunnerApi.Components components = pipeline.getComponents(); RunnerApi.PTransform transform = components.getTransformsOrThrow(id); RunnerApi.WindowIntoPayload payload; try { payload = RunnerApi.WindowIntoPayload.parseFrom(transform.getSpec().getPayload()); } catch (InvalidProtocolBufferException e) { throw new IllegalArgumentException(e); } //TODO: https://issues.apache.org/jira/browse/BEAM-4296 // This only works for well known window fns, we should defer this execution to the SDK // if the WindowFn can't be parsed or just defer it all the time. WindowFn<T, ? extends BoundedWindow> windowFn = (WindowFn<T, ? extends BoundedWindow>) WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn()); String inputCollectionId = Iterables.getOnlyElement(transform.getInputsMap().values()); String outputCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); Coder<WindowedValue<T>> outputCoder = instantiateCoder(outputCollectionId, components); TypeInformation<WindowedValue<T>> resultTypeInfo = new CoderTypeInformation<>(outputCoder); DataStream<WindowedValue<T>> inputDataStream = context.getDataStreamOrThrow(inputCollectionId); FlinkAssignWindows<T, ? extends BoundedWindow> assignWindowsFunction = new FlinkAssignWindows<>(windowFn); DataStream<WindowedValue<T>> resultDataStream = inputDataStream .flatMap(assignWindowsFunction) .name(transform.getUniqueName()) .returns(resultTypeInfo); context.addDataStream(outputCollectionId, resultDataStream); }
/** * 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 toMessageProto( WindowingStrategy<?, ?> windowingStrategy, SdkComponents components) throws IOException { RunnerApi.WindowingStrategy windowingStrategyProto = toProto(windowingStrategy, components); return RunnerApi.MessageWithComponents.newBuilder() .setWindowingStrategy(windowingStrategyProto) .setComponents(components.toComponents()) .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)); } }
private <T> void translateAssignWindows( String id, RunnerApi.Pipeline pipeline, StreamingTranslationContext context) { RunnerApi.Components components = pipeline.getComponents(); RunnerApi.PTransform transform = components.getTransformsOrThrow(id); RunnerApi.WindowIntoPayload payload; try { payload = RunnerApi.WindowIntoPayload.parseFrom(transform.getSpec().getPayload()); } catch (InvalidProtocolBufferException e) { throw new IllegalArgumentException(e); } //TODO: https://issues.apache.org/jira/browse/BEAM-4296 // This only works for well known window fns, we should defer this execution to the SDK // if the WindowFn can't be parsed or just defer it all the time. WindowFn<T, ? extends BoundedWindow> windowFn = (WindowFn<T, ? extends BoundedWindow>) WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn()); String inputCollectionId = Iterables.getOnlyElement(transform.getInputsMap().values()); String outputCollectionId = Iterables.getOnlyElement(transform.getOutputsMap().values()); Coder<WindowedValue<T>> outputCoder = instantiateCoder(outputCollectionId, components); TypeInformation<WindowedValue<T>> resultTypeInfo = new CoderTypeInformation<>(outputCoder); DataStream<WindowedValue<T>> inputDataStream = context.getDataStreamOrThrow(inputCollectionId); FlinkAssignWindows<T, ? extends BoundedWindow> assignWindowsFunction = new FlinkAssignWindows<>(windowFn); DataStream<WindowedValue<T>> resultDataStream = inputDataStream .flatMap(assignWindowsFunction) .name(transform.getUniqueName()) .returns(resultTypeInfo); context.addDataStream(outputCollectionId, resultDataStream); }
/** * 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(); }
public static WindowingStrategy getWindowingStrategy( String pCollectionId, RunnerApi.Components components) { RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow( components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId()); final WindowingStrategy<?, ?> windowingStrategy; try { return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.forComponents(components)); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException( String.format( "Unable to hydrate windowing strategy %s for %s.", windowingStrategyProto, pCollectionId), e); } } }
WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn());
/** * Registers the provided {@link WindowingStrategy} into this {@link SdkComponents}, returning a * unique ID for the {@link WindowingStrategy}. Multiple registrations of the same {@link * WindowingStrategy} will return the same unique ID. */ public String registerWindowingStrategy(WindowingStrategy<?, ?> windowingStrategy) throws IOException { String existing = windowingStrategyIds.get(windowingStrategy); if (existing != null) { return existing; } String baseName = String.format( "%s(%s)", NameUtils.approximateSimpleName(windowingStrategy), NameUtils.approximateSimpleName(windowingStrategy.getWindowFn())); String name = uniqify(baseName, windowingStrategyIds.values()); windowingStrategyIds.put(windowingStrategy, name); RunnerApi.WindowingStrategy windowingStrategyProto = WindowingStrategyTranslation.toProto(windowingStrategy, this); componentsBuilder.putWindowingStrategies(name, windowingStrategyProto); return name; }
public static WindowingStrategy getWindowingStrategy( String pCollectionId, RunnerApi.Components components) { RunnerApi.WindowingStrategy windowingStrategyProto = components.getWindowingStrategiesOrThrow( components.getPcollectionsOrThrow(pCollectionId).getWindowingStrategyId()); final WindowingStrategy<?, ?> windowingStrategy; try { return WindowingStrategyTranslation.fromProto( windowingStrategyProto, RehydratedComponents.forComponents(components)); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException( String.format( "Unable to hydrate windowing strategy %s for %s.", windowingStrategyProto, pCollectionId), e); } } }
WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn());
@Test public void getEnvironmentWindowIntoKnown() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoPayload.newBuilder() .setWindowFn( WindowingStrategyTranslation.toProto( FixedWindows.of(Duration.standardMinutes(5L)), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getWindowFn().getEnvironmentId()))); }
try { windowingStrategy = WindowingStrategyTranslation.fromProto(windowingStrategyProto, rehydratedComponents); } catch (InvalidProtocolBufferException e) { throw new IllegalStateException(
@Test public void testToFromProto() throws InvalidProtocolBufferException { pipeline.apply(GenerateSequence.from(0)).apply(Window.<Long>into((WindowFn) windowFn)); final AtomicReference<AppliedPTransform<?, ?, Assign<?>>> assign = new AtomicReference<>(null); pipeline.traverseTopologically( new PipelineVisitor.Defaults() { @Override public void visitPrimitiveTransform(Node node) { if (node.getTransform() instanceof Window.Assign) { checkState(assign.get() == null); assign.set( (AppliedPTransform<?, ?, Assign<?>>) node.toAppliedPTransform(getPipeline())); } } }); checkState(assign.get() != null); SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); WindowIntoPayload payload = WindowIntoTranslation.toProto(assign.get().getTransform(), components); assertEquals(windowFn, WindowingStrategyTranslation.windowFnFromProto(payload.getWindowFn())); }