private static String registerCoderOrThrow(SdkComponents components, Coder coder) { try { return components.registerCoder(coder); } catch (IOException exc) { throw new RuntimeException("Failure to register coder", exc); } }
private static <T extends Coder<?>> List<String> registerComponents( T coder, CoderTranslator<T> translator, SdkComponents components) throws IOException { List<String> componentIds = new ArrayList<>(); for (Coder<?> component : translator.getComponents(coder)) { componentIds.add(components.registerCoder(component)); } return componentIds; }
/** Convert the provided {@link Coder} into a {@link CloudObject}. */ public static CloudObject asCloudObject(Coder<?> coder, @Nullable SdkComponents sdkComponents) { CloudObjectTranslator<Coder> translator = (CloudObjectTranslator<Coder>) CODER_TRANSLATORS.get(coder.getClass()); CloudObject encoding; if (translator != null) { encoding = translator.toCloudObject(coder, sdkComponents); } else { CloudObjectTranslator customCoderTranslator = CODER_TRANSLATORS.get(CustomCoder.class); checkNotNull( customCoderTranslator, "No %s registered for %s, but it is in the %s", CloudObjectTranslator.class.getSimpleName(), CustomCoder.class.getSimpleName(), DefaultCoderCloudObjectTranslatorRegistrar.class.getSimpleName()); encoding = customCoderTranslator.toCloudObject(coder, sdkComponents); } if (sdkComponents != null && !ModelCoderRegistrar.isKnownCoder(coder)) { try { String coderId = sdkComponents.registerCoder(coder); Structs.addString(encoding, PropertyNames.PIPELINE_PROTO_CODER_ID, coderId); } catch (Exception e) { throw new RuntimeException("Unable to register coder " + coder, e); } } return encoding; }
/** Produces a {@link RunnerApi.TestStreamPayload} from a {@link TestStream}. */ static <T> RunnerApi.TestStreamPayload payloadForTestStream( final TestStream<T> transform, SdkComponents components) throws IOException { List<RunnerApi.TestStreamPayload.Event> protoEvents = new ArrayList<>(); try { for (TestStream.Event<T> event : transform.getEvents()) { protoEvents.add(eventToProto(event, transform.getValueCoder())); } } catch (IOException e) { throw new RuntimeException(e); } return RunnerApi.TestStreamPayload.newBuilder() .setCoderId(components.registerCoder(transform.getValueCoder())) .addAllEvents(protoEvents) .build(); } }
@VisibleForTesting static CombinePayload toProto( AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>> combine, SdkComponents sdkComponents) throws IOException { checkArgument( combine.getTransform().getSideInputs().isEmpty(), "CombineTranslation.toProto cannot translate Combines with side inputs."); GlobalCombineFn<?, ?, ?> combineFn = combine.getTransform().getFn(); try { Coder<?> accumulatorCoder = extractAccumulatorCoder(combineFn, (AppliedPTransform) combine); return RunnerApi.CombinePayload.newBuilder() .setAccumulatorCoderId(sdkComponents.registerCoder(accumulatorCoder)) .setCombineFn(toProto(combineFn, sdkComponents)) .build(); } catch (CannotProvideCoderException e) { throw new IllegalArgumentException(e); } }
@Test public void registerCoder() throws IOException { Coder<?> coder = KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(SetCoder.of(ByteArrayCoder.of()))); String id = components.registerCoder(coder); assertThat(components.registerCoder(coder), equalTo(id)); assertThat(id, not(isEmptyOrNullString())); VarLongCoder otherCoder = VarLongCoder.of(); assertThat(components.registerCoder(otherCoder), not(equalTo(id))); components.toComponents().getCodersOrThrow(id); components.toComponents().getCodersOrThrow(components.registerCoder(otherCoder)); }
final Coder<?> restrictionCoder = DoFnInvokers.invokerFor(doFn).invokeGetRestrictionCoder(pipeline.getCoderRegistry()); restrictionCoderId = components.registerCoder(restrictionCoder); } else { restrictionCoderId = "";
public static RunnerApi.PCollection toProto(PCollection<?> pCollection, SdkComponents components) throws IOException { String coderId = components.registerCoder(pCollection.getCoder()); String windowingStrategyId = components.registerWindowingStrategy(pCollection.getWindowingStrategy()); // TODO: Display Data return RunnerApi.PCollection.newBuilder() .setUniqueName(pCollection.getName()) .setCoderId(coderId) .setIsBounded(toProto(pCollection.isBounded())) .setWindowingStrategyId(windowingStrategyId) .build(); }
private static void checkPipelineProtoCoderIds( Coder<?> coder, CloudObject cloudObject, SdkComponents sdkComponents) throws Exception { if (ModelCoderRegistrar.isKnownCoder(coder)) { assertFalse(cloudObject.containsKey(PropertyNames.PIPELINE_PROTO_CODER_ID)); } else { assertTrue(cloudObject.containsKey(PropertyNames.PIPELINE_PROTO_CODER_ID)); assertEquals( sdkComponents.registerCoder(coder), cloudObject.get(PropertyNames.PIPELINE_PROTO_CODER_ID)); } List<? extends Coder<?>> coderArguments = coder.getCoderArguments(); Object cloudComponentsObject = cloudObject.get(PropertyNames.COMPONENT_ENCODINGS); assertTrue(cloudComponentsObject instanceof List); List<CloudObject> cloudComponents = (List<CloudObject>) cloudComponentsObject; assertEquals(coderArguments.size(), cloudComponents.size()); for (int i = 0; i < coderArguments.size(); i++) { checkPipelineProtoCoderIds(coderArguments.get(i), cloudComponents.get(i), sdkComponents); } } }
@Test public void registerCoderEqualsNotSame() throws IOException { Coder<?> coder = KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(SetCoder.of(ByteArrayCoder.of()))); Coder<?> otherCoder = KvCoder.of(StringUtf8Coder.of(), IterableCoder.of(SetCoder.of(ByteArrayCoder.of()))); assertThat(coder, equalTo(otherCoder)); String id = components.registerCoder(coder); String otherId = components.registerCoder(otherCoder); assertThat(otherId, not(equalTo(id))); components.toComponents().getCodersOrThrow(id); components.toComponents().getCodersOrThrow(otherId); }
final DoFn<?, ?> fn = pke.getFn(); final DoFnSignature signature = DoFnSignatures.getSignature(fn.getClass()); final String restrictionCoderId = components.registerCoder(pke.getRestrictionCoder());
@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 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)); }
return RunnerApi.CombinePayload.newBuilder() .setAccumulatorCoderId( components.registerCoder( extractAccumulatorCoder(combineFn, (AppliedPTransform) combine))) .setCombineFn(
/** * 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(); }
@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())); } }