/** * Update all composites present in the {@code originalPipeline} with an URN equal to the provided * {@code urn} using the provided {@link TransformReplacement}. */ public static Pipeline updateTransform( String urn, Pipeline originalPipeline, TransformReplacement compositeBuilder) { Components.Builder resultComponents = originalPipeline.getComponents().toBuilder(); for (Map.Entry<String, PTransform> pt : originalPipeline.getComponents().getTransformsMap().entrySet()) { if (pt.getValue().getSpec() != null && urn.equals(pt.getValue().getSpec().getUrn())) { MessageWithComponents updated = compositeBuilder.getReplacement(pt.getKey(), originalPipeline.getComponents()); checkArgument( updated.getPtransform().getOutputsMap().equals(pt.getValue().getOutputsMap()), "A %s must produce all of the outputs of the original %s", TransformReplacement.class.getSimpleName(), PTransform.class.getSimpleName()); removeSubtransforms(pt.getValue(), resultComponents); resultComponents .mergeFrom(updated.getComponents()) .putTransforms(pt.getKey(), updated.getPtransform()); } } return originalPipeline.toBuilder().setComponents(resultComponents).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 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); } }
return MessageWithComponents.newBuilder() .setPtransform(newPTransform.build()) .setComponents(newComponents)
(transformId, existingComponents) -> { String subtransform = String.format("%s_sub", transformId); return MessageWithComponents.newBuilder() .setPtransform( PTransform.newBuilder()
return MessageWithComponents.newBuilder() .setPtransform(newTransform) .setComponents(newComponents)
/** * 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(); }
@Test public void windowedValueCoderComponentsToConstructor() throws IOException { FullWindowedValueCoder<Iterable<KV<String, Integer>>> javaCoder = FullWindowedValueCoder.of( IterableCoder.of(KvCoder.of(StringUtf8Coder.of(), VarIntCoder.of())), IntervalWindowCoder.of()); MessageWithComponents coderAndComponents = CoderTranslation.toProto(javaCoder); WindowedValueCoderComponents windowedValueCoderComponents = ModelCoders.getWindowedValueCoderComponents(coderAndComponents.getCoder()); Coder windowedCoder = ModelCoders.windowedValueCoder( windowedValueCoderComponents.elementCoderId(), windowedValueCoderComponents.windowCoderId()); assertThat(windowedCoder, equalTo(coderAndComponents.getCoder())); }
@Test public void kvCoderComponentsToConstructor() throws IOException { KvCoder<byte[], Iterable<Long>> javaCoder = KvCoder.of(ByteArrayCoder.of(), IterableCoder.of(LengthPrefixCoder.of(VarLongCoder.of()))); MessageWithComponents coderAndComponents = CoderTranslation.toProto(javaCoder); KvCoderComponents kvCoderComponents = ModelCoders.getKvCoderComponents(coderAndComponents.getCoder()); Coder kvCoder = ModelCoders.kvCoder(kvCoderComponents.keyCoderId(), kvCoderComponents.valueCoderId()); assertThat(kvCoder, equalTo(coderAndComponents.getCoder())); }
public static RunnerApi.MessageWithComponents toProto(Coder<?> coder) throws IOException { SdkComponents components = SdkComponents.create(); RunnerApi.Coder coderProto = toProto(coder, components); return RunnerApi.MessageWithComponents.newBuilder() .setCoder(coderProto) .setComponents(components.toComponents()) .build(); }
@Override public MessageWithComponents getReplacement( String transformId, ComponentsOrBuilder existingComponents) { return MessageWithComponents.newBuilder() .setPtransform(extraTransform) .setComponents(extraComponents) .build(); } }
private Coder<T> getCoder() throws IOException { if (coder == null) { coder = (Coder) CoderTranslation.fromProto( coderSpec.getCoder(), RehydratedComponents.forComponents(coderSpec.getComponents())); } return coder; }