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); }
private static Coder<?> getAccumulatorCoder( CombinePayload payload, RehydratedComponents components) throws IOException { String id = payload.getAccumulatorCoderId(); return components.getCoder(id); }
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 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); }
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()); }
private static <T> void verifyTestStreamEncoding( TestStream<T> testStream, RunnerApi.TestStreamPayload payload, RehydratedComponents protoComponents) throws Exception { // This reverse direction is only valid for Java-based coders assertThat(protoComponents.getCoder(payload.getCoderId()), equalTo(testStream.getValueCoder())); assertThat(payload.getEventsList().size(), equalTo(testStream.getEvents().size())); for (int i = 0; i < payload.getEventsList().size(); ++i) { assertThat( TestStreamTranslation.eventFromProto(payload.getEvents(i), testStream.getValueCoder()), equalTo(testStream.getEvents().get(i))); } } }
@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 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)); }
@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 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())); } }
instanceOf(TestSplittableFn.class)); assertThat( components.getCoder(payload.getRestrictionCoderId()), instanceOf(SerializableCoder.class));