@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); } }
return RunnerApi.CombinePayload.newBuilder() .setAccumulatorCoderId( components.registerCoder(
@Test public void getEnvironmentCombine() throws IOException { SdkComponents components = SdkComponents.create(); components.registerEnvironment(Environments.createDockerEnvironment("java")); CombinePayload payload = CombinePayload.newBuilder() .setCombineFn(CombineTranslation.toProto(Sum.ofLongs(), components)) .build(); RehydratedComponents rehydratedComponents = RehydratedComponents.forComponents(components.toComponents()); PTransform builder = PTransform.newBuilder() .setSpec( FunctionSpec.newBuilder() .setUrn(PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payload.toByteString()) .build()) .build(); Environment env = Environments.getEnvironment(builder, rehydratedComponents).get(); assertThat( env, equalTo( components .toComponents() .getEnvironmentsOrThrow(payload.getCombineFn().getEnvironmentId()))); } }