/** * 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)); } }
public static PCollection<?> fromProto( RunnerApi.PCollection pCollection, Pipeline pipeline, RehydratedComponents components) throws IOException { Coder<?> coder = components.getCoder(pCollection.getCoderId()); return PCollection.createPrimitiveOutputInternal( pipeline, components.getWindowingStrategy(pCollection.getWindowingStrategyId()), fromProto(pCollection.getIsBounded()), (Coder) coder); }
public static Optional<Environment> getEnvironment( PTransform ptransform, RehydratedComponents components) { try { String envId = KNOWN_URN_SPEC_EXTRACTORS .getOrDefault(ptransform.getSpec().getUrn(), DEFAULT_SPEC_EXTRACTOR) .getEnvironmentId(ptransform); if (!Strings.isNullOrEmpty(envId)) { // Some PTransform payloads may have an empty (default) Environment ID, for example a // WindowIntoPayload with a known WindowFn. Others will never have an Environment ID, such // as a GroupByKeyPayload, and the Default extractor returns null in this case. return Optional.of(components.getEnvironment(envId)); } else { return Optional.empty(); } } catch (IOException e) { throw new RuntimeException(e); } }
@Test public void testSimpleCoder() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Coder<?> coder = VarIntCoder.of(); String id = sdkComponents.registerCoder(coder); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> rehydratedCoder = rehydratedComponents.getCoder(id); assertThat(rehydratedCoder, equalTo((Coder) coder)); assertThat(rehydratedComponents.getCoder(id), theInstance((Coder) rehydratedCoder)); }
@Test public void testEncodeDecodeFields() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.PCollection protoCollection = PCollectionTranslation.toProto(testCollection, sdkComponents); RehydratedComponents protoComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> decodedCoder = protoComponents.getCoder(protoCollection.getCoderId()); WindowingStrategy<?, ?> decodedStrategy = protoComponents.getWindowingStrategy(protoCollection.getWindowingStrategyId()); IsBounded decodedIsBounded = PCollectionTranslation.isBounded(protoCollection); assertThat(decodedCoder, equalTo(testCollection.getCoder())); assertThat(decodedStrategy, equalTo(testCollection.getWindowingStrategy().fixDefaults())); assertThat(decodedIsBounded, equalTo(testCollection.isBounded())); }
@Test public void testEnvironment() { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Environment env = Environments.createDockerEnvironment("java_test"); String id = sdkComponents.registerEnvironment(env); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Environment rehydratedEnv = rehydratedComponents.getEnvironment(id); assertThat(rehydratedEnv, equalTo(env)); assertThat(rehydratedComponents.getEnvironment(id), theInstance(rehydratedEnv)); } }
@Test public void testWindowingStrategy() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); WindowingStrategy windowingStrategy = WindowingStrategy.of(FixedWindows.of(Duration.millis(1))) .withAllowedLateness(Duration.standardSeconds(4)); String id = sdkComponents.registerWindowingStrategy(windowingStrategy); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); WindowingStrategy<?, ?> rehydratedStrategy = rehydratedComponents.getWindowingStrategy(id); assertThat(rehydratedStrategy, equalTo((WindowingStrategy) windowingStrategy.fixDefaults())); assertThat( rehydratedComponents.getWindowingStrategy(id), theInstance((WindowingStrategy) rehydratedStrategy)); }
private static TestStream<?> testStreamFromProtoPayload( RunnerApi.TestStreamPayload testStreamPayload, RehydratedComponents components) throws IOException { Coder<Object> coder = (Coder<Object>) components.getCoder(testStreamPayload.getCoderId()); List<TestStream.Event<Object>> events = new ArrayList<>(); for (RunnerApi.TestStreamPayload.Event event : testStreamPayload.getEventsList()) { events.add(eventFromProto(event, coder)); } return TestStream.fromRawEvents(coder, events); }
/** Create a new {@link RehydratedComponents} from a Runner API {@link Components}. */ public static RehydratedComponents forComponents(RunnerApi.Components components) { return new RehydratedComponents(components, null); }
@Override public Map<TupleTag<?>, PValue> getAdditionalInputs() { Map<TupleTag<?>, PValue> additionalInputs = new HashMap<>(); for (Map.Entry<String, SideInput> sideInputEntry : payload.getSideInputsMap().entrySet()) { try { additionalInputs.put( new TupleTag<>(sideInputEntry.getKey()), rehydratedComponents.getPCollection( protoTransform.getInputsOrThrow(sideInputEntry.getKey()))); } catch (IOException exc) { throw new IllegalStateException( String.format( "Could not find input with name %s for %s transform", sideInputEntry.getKey(), WriteFiles.class.getSimpleName())); } } return additionalInputs; }
@Test public void testCompoundCoder() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); Coder<?> coder = VarIntCoder.of(); Coder<?> compoundCoder = NullableCoder.of(coder); String compoundCoderId = sdkComponents.registerCoder(compoundCoder); String coderId = sdkComponents.registerCoder(coder); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(sdkComponents.toComponents()); Coder<?> rehydratedCoder = rehydratedComponents.getCoder(coderId); Coder<?> rehydratedCompoundCoder = rehydratedComponents.getCoder(compoundCoderId); assertThat(rehydratedCoder, equalTo((Coder) coder)); assertThat(rehydratedCompoundCoder, equalTo((Coder) compoundCoder)); assertThat(rehydratedComponents.getCoder(coderId), theInstance((Coder) rehydratedCoder)); assertThat( rehydratedComponents.getCoder(compoundCoderId), theInstance((Coder) rehydratedCompoundCoder)); }
private static Coder<?> getAccumulatorCoder( CombinePayload payload, RehydratedComponents components) throws IOException { String id = payload.getAccumulatorCoderId(); return components.getCoder(id); }
/** Create a new {@link RehydratedComponents} with a pipeline attached. */ public RehydratedComponents withPipeline(Pipeline pipeline) { return new RehydratedComponents(components, pipeline); }
private Coder<T> getCoder() throws IOException { if (coder == null) { coder = (Coder) CoderTranslation.fromProto( coderSpec.getCoder(), RehydratedComponents.forComponents(coderSpec.getComponents())); } return coder; }
@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())); } }
switch (stateSpec.getSpecCase()) { case VALUE_SPEC: return StateSpecs.value(components.getCoder(stateSpec.getValueSpec().getCoderId())); case BAG_SPEC: return StateSpecs.bag(components.getCoder(stateSpec.getBagSpec().getElementCoderId())); case COMBINING_SPEC: FunctionSpec combineFnSpec = stateSpec.getCombiningSpec().getCombineFn().getSpec(); (Coder) components.getCoder(stateSpec.getCombiningSpec().getAccumulatorCoderId()), combineFn); components.getCoder(stateSpec.getMapSpec().getKeyCoderId()), components.getCoder(stateSpec.getMapSpec().getValueCoderId())); return StateSpecs.set(components.getCoder(stateSpec.getSetSpec().getElementCoderId()));
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); } } }
RehydratedComponents components = RehydratedComponents.forComponents(componentsProto); RunnerApi.PTransform spkTransform = componentsProto.getTransformsOrThrow(fn); assertEquals( instanceOf(TestSplittableFn.class)); assertThat( components.getCoder(payload.getRestrictionCoderId()), instanceOf(SerializableCoder.class));
private static Coder<?> fromKnownCoder(RunnerApi.Coder coder, RehydratedComponents components) throws IOException { String coderUrn = coder.getSpec().getSpec().getUrn(); List<Coder<?>> coderComponents = new ArrayList<>(); for (String componentId : coder.getComponentCoderIdsList()) { Coder<?> innerCoder = components.getCoder(componentId); coderComponents.add(innerCoder); } Class<? extends Coder> coderType = KNOWN_CODER_URNS.inverse().get(coderUrn); CoderTranslator<?> translator = KNOWN_TRANSLATORS.get(coderType); checkArgument( translator != null, "Unknown Coder URN %s. Known URNs: %s", coderUrn, KNOWN_CODER_URNS.values()); return translator.fromComponents( coderComponents, coder.getSpec().getSpec().getPayload().toByteArray()); }
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); } } }