@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); } }
private static Optional<CombinePayload> getCombinePayload( AppliedPTransform<?, ?, ?> transform, SdkComponents components) throws IOException { RunnerApi.PTransform proto = PTransformTranslation.toProto(transform, Collections.emptyList(), components); // Even if the proto has no spec, calling getSpec still returns a blank spec, which we want to // avoid. It should be clear to the caller whether or not there was a spec in the transform. if (proto.hasSpec()) { return Optional.of(CombinePayload.parseFrom(proto.getSpec().getPayload())); } else { return Optional.empty(); } } }
@Override public FunctionSpec translate( AppliedPTransform<?, ?, Combine.PerKey<?, ?, ?>> transform, SdkComponents components) throws IOException { if (transform.getTransform().getSideInputs().isEmpty()) { return FunctionSpec.newBuilder() .setUrn(COMBINE_PER_KEY_TRANSFORM_URN) .setPayload(payloadForCombine((AppliedPTransform) transform, components).toByteString()) .build(); } else { // Combines with side inputs are translated as generic composites, which have a blank // FunctionSpec. return null; } }
combineFn, SerializableUtils.deserializeFromByteArray( combineProto.getCombineFn().getSpec().getPayload().toByteArray(), "CombineFn"));
combineFn, SerializableUtils.deserializeFromByteArray( combineProto.getCombineFn().getSpec().getPayload().toByteArray(), "CombineFn"));
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()))); } }
private static Coder<?> getAccumulatorCoder( CombinePayload payload, RehydratedComponents components) throws IOException { String id = payload.getAccumulatorCoderId(); return components.getCoder(id); }
private static void validateCombine(String id, PTransform transform, Components components) throws Exception { CombinePayload payload = CombinePayload.parseFrom(transform.getSpec().getPayload()); checkArgument( components.containsCoders(payload.getAccumulatorCoderId()), "Transform %s uses unknown accumulator coder id %s", payload.getAccumulatorCoderId()); }
private static String combineExtractor(PTransform pTransform) throws InvalidProtocolBufferException { return CombinePayload.parseFrom(pTransform.getSpec().getPayload()) .getCombineFn() .getEnvironmentId(); }