Coder.newBuilder() .setSpec( SdkFunctionSpec.newBuilder()
@Test public void toAndFromProto() throws Exception { SdkComponents sdkComponents = SdkComponents.create(); sdkComponents.registerEnvironment(Environments.createDockerEnvironment("java")); RunnerApi.Coder coderProto = CoderTranslation.toProto(coder, sdkComponents); Components encodedComponents = sdkComponents.toComponents(); Coder<?> decodedCoder = CoderTranslation.fromProto( coderProto, RehydratedComponents.forComponents(encodedComponents)); assertThat(decodedCoder, equalTo(coder)); if (KNOWN_CODERS.contains(coder)) { for (RunnerApi.Coder encodedCoder : encodedComponents.getCodersMap().values()) { assertThat( encodedCoder.getSpec().getSpec().getUrn(), not(equalTo(CoderTranslation.JAVA_SERIALIZED_CODER_URN))); } } }
public static KvCoderComponents getKvCoderComponents(Coder coder) { checkArgument(KV_CODER_URN.equals(coder.getSpec().getSpec().getUrn())); return new AutoValue_ModelCoders_KvCoderComponents( coder.getComponentCoderIds(0), coder.getComponentCoderIds(1)); }
Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build())
Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build())
Components components = Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build())
.setComponents( Components.newBuilder() .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws",
components.getCodersOrThrow(coderId).getComponentCoderIdsList()) { checkArgument( components.containsCoders(componentCoderId),
.putCoders( "coder", Coder.newBuilder().setSpec(SdkFunctionSpec.getDefaultInstance()).build())) .build();
Coder.newBuilder() .setSpec( SdkFunctionSpec.newBuilder()
.putCoders( "coder", Coder.newBuilder().setSpec(SdkFunctionSpec.getDefaultInstance()).build())) .build();
/** This method doesn't do any pruning for reachability, but this may not require a test. */ @Test public void retainOnlyPrimitivesIgnoresUnreachableNodes() { Pipeline p = Pipeline.create(); p.apply( new org.apache.beam.sdk.transforms.PTransform<PBegin, PCollection<Long>>() { @Override public PCollection<Long> expand(PBegin input) { return input .apply(GenerateSequence.from(2L)) .apply(Window.into(FixedWindows.of(Duration.standardMinutes(5L)))) .apply(MapElements.into(TypeDescriptors.longs()).via(l -> l + 1)); } }); Components augmentedComponents = PipelineTranslation.toProto(p) .getComponents() .toBuilder() .putCoders("extra-coder", RunnerApi.Coder.getDefaultInstance()) .putWindowingStrategies( "extra-windowing-strategy", RunnerApi.WindowingStrategy.getDefaultInstance()) .putEnvironments("extra-env", RunnerApi.Environment.getDefaultInstance()) .putPcollections("extra-pc", RunnerApi.PCollection.getDefaultInstance()) .build(); Collection<String> primitiveComponents = QueryablePipeline.getPrimitiveTransformIds(augmentedComponents); } }
@Before public void setup() { partialComponents = Components.newBuilder() .putTransforms( "impulse", PTransform.newBuilder() .setUniqueName("Impulse") .putOutputs("output", "impulse.out") .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.IMPULSE_TRANSFORM_URN)) .build()) .putPcollections("impulse.out", pc("impulse.out")) .putEnvironments("go", Environments.createDockerEnvironment("go")) .putEnvironments("py", Environments.createDockerEnvironment("py")) .putCoders("coder", Coder.newBuilder().build()) .putCoders("windowCoder", Coder.newBuilder().build()) .putWindowingStrategies( "ws", WindowingStrategy.newBuilder().setWindowCoderId("windowCoder").build()) .build(); }
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 RunnerApi.Coder toKnownCoder(Coder<?> coder, SdkComponents components) throws IOException { CoderTranslator translator = KNOWN_TRANSLATORS.get(coder.getClass()); List<String> componentIds = registerComponents(coder, translator, components); return RunnerApi.Coder.newBuilder() .addAllComponentCoderIds(componentIds) .setSpec( SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(KNOWN_CODER_URNS.get(coder.getClass())) .setPayload(ByteString.copyFrom(translator.getPayload(coder))))) .build(); }
@Test public void kvCoderComponentsNoUrn() { thrown.expect(IllegalArgumentException.class); ModelCoders.getKvCoderComponents( Coder.newBuilder().setSpec(SdkFunctionSpec.getDefaultInstance()).build()); } }
private static RunnerApi.Coder toCustomCoder(Coder<?> coder) throws IOException { RunnerApi.Coder.Builder coderBuilder = RunnerApi.Coder.newBuilder(); return coderBuilder .setSpec( SdkFunctionSpec.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(JAVA_SERIALIZED_CODER_URN) .setPayload( ByteString.copyFrom(SerializableUtils.serializeToByteArray(coder))) .build())) .build(); }
@Test public void windowedValueCoderComponentsWrongUrn() { thrown.expect(IllegalArgumentException.class); ModelCoders.getWindowedValueCoderComponents( Coder.newBuilder() .setSpec( SdkFunctionSpec.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.LENGTH_PREFIX_CODER_URN))) .build()); }
public static Coder kvCoder(String keyCoderId, String valueCoderId) { return Coder.newBuilder() .setSpec( SdkFunctionSpec.newBuilder().setSpec(FunctionSpec.newBuilder().setUrn(KV_CODER_URN))) .addComponentCoderIds(keyCoderId) .addComponentCoderIds(valueCoderId) .build(); }
@Test public void kvCoderComponentsWrongUrn() { thrown.expect(IllegalArgumentException.class); ModelCoders.getKvCoderComponents( Coder.newBuilder() .setSpec( SdkFunctionSpec.newBuilder() .setSpec(FunctionSpec.newBuilder().setUrn(ModelCoders.LENGTH_PREFIX_CODER_URN))) .build()); }