| /* |
| * Licensed to the Apache Software Foundation (ASF) under one |
| * or more contributor license agreements. See the NOTICE file |
| * distributed with this work for additional information |
| * regarding copyright ownership. The ASF licenses this file |
| * to you under the Apache License, Version 2.0 (the |
| * "License"); you may not use this file except in compliance |
| * with the License. You may obtain a copy of the License at |
| * |
| * http://www.apache.org/licenses/LICENSE-2.0 |
| * |
| * Unless required by applicable law or agreed to in writing, software |
| * distributed under the License is distributed on an "AS IS" BASIS, |
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| * See the License for the specific language governing permissions and |
| * limitations under the License. |
| */ |
| package org.apache.beam.fn.harness; |
| |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; |
| |
| import java.io.IOException; |
| import java.util.Map; |
| import java.util.function.Supplier; |
| import org.apache.beam.fn.harness.control.BundleSplitListener; |
| import org.apache.beam.fn.harness.data.BeamFnDataClient; |
| import org.apache.beam.fn.harness.data.PCollectionConsumerRegistry; |
| import org.apache.beam.fn.harness.data.PTransformFunctionRegistry; |
| import org.apache.beam.fn.harness.state.BeamFnStateClient; |
| import org.apache.beam.fn.harness.state.SideInputSpec; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PCollection; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.PTransform; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.ParDoPayload; |
| import org.apache.beam.runners.core.construction.PCollectionViewTranslation; |
| import org.apache.beam.runners.core.construction.ParDoTranslation; |
| import org.apache.beam.runners.core.construction.RehydratedComponents; |
| import org.apache.beam.runners.core.construction.Timer; |
| import org.apache.beam.sdk.Pipeline; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.fn.data.FnDataReceiver; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.schemas.SchemaCoder; |
| import org.apache.beam.sdk.state.TimeDomain; |
| import org.apache.beam.sdk.transforms.DoFn; |
| import org.apache.beam.sdk.transforms.Materializations; |
| import org.apache.beam.sdk.transforms.reflect.DoFnSignature; |
| import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; |
| import org.apache.beam.sdk.transforms.windowing.BoundedWindow; |
| import org.apache.beam.sdk.util.WindowedValue; |
| import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder; |
| import org.apache.beam.sdk.values.KV; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableListMultimap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableMap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ListMultimap; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Maps; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Sets; |
| |
| /** A {@link PTransformRunnerFactory} for transforms invoking a {@link DoFn}. */ |
| abstract class DoFnPTransformRunnerFactory< |
| TransformInputT, |
| FnInputT, |
| OutputT, |
| RunnerT extends DoFnPTransformRunnerFactory.DoFnPTransformRunner<TransformInputT>> |
| implements PTransformRunnerFactory<RunnerT> { |
| interface DoFnPTransformRunner<T> { |
| void startBundle() throws Exception; |
| |
| void processElement(WindowedValue<T> input) throws Exception; |
| |
| void processTimer( |
| String timerId, TimeDomain timeDomain, WindowedValue<KV<Object, Timer>> input); |
| |
| void finishBundle() throws Exception; |
| } |
| |
| @Override |
| public final RunnerT createRunnerForPTransform( |
| PipelineOptions pipelineOptions, |
| BeamFnDataClient beamFnDataClient, |
| BeamFnStateClient beamFnStateClient, |
| String pTransformId, |
| PTransform pTransform, |
| Supplier<String> processBundleInstructionId, |
| Map<String, PCollection> pCollections, |
| Map<String, RunnerApi.Coder> coders, |
| Map<String, RunnerApi.WindowingStrategy> windowingStrategies, |
| PCollectionConsumerRegistry pCollectionConsumerRegistry, |
| PTransformFunctionRegistry startFunctionRegistry, |
| PTransformFunctionRegistry finishFunctionRegistry, |
| BundleSplitListener splitListener) { |
| Context<FnInputT, OutputT> context = |
| new Context<>( |
| pipelineOptions, |
| beamFnStateClient, |
| pTransformId, |
| pTransform, |
| processBundleInstructionId, |
| pCollections, |
| coders, |
| windowingStrategies, |
| pCollectionConsumerRegistry, |
| splitListener); |
| |
| RunnerT runner = createRunner(context); |
| |
| // Register the appropriate handlers. |
| startFunctionRegistry.register(pTransformId, runner::startBundle); |
| Iterable<String> mainInput = |
| Sets.difference( |
| pTransform.getInputsMap().keySet(), |
| Sets.union( |
| context.parDoPayload.getSideInputsMap().keySet(), |
| context.parDoPayload.getTimerSpecsMap().keySet())); |
| for (String localInputName : mainInput) { |
| pCollectionConsumerRegistry.register( |
| pTransform.getInputsOrThrow(localInputName), |
| pTransformId, |
| (FnDataReceiver) (FnDataReceiver<WindowedValue<TransformInputT>>) runner::processElement); |
| } |
| |
| // Register as a consumer for each timer PCollection. |
| for (String localName : context.parDoPayload.getTimerSpecsMap().keySet()) { |
| TimeDomain timeDomain = |
| DoFnSignatures.getTimerSpecOrThrow( |
| context.doFnSignature.timerDeclarations().get(localName), context.doFn) |
| .getTimeDomain(); |
| pCollectionConsumerRegistry.register( |
| pTransform.getInputsOrThrow(localName), |
| pTransformId, |
| (FnDataReceiver) |
| timer -> |
| runner.processTimer( |
| localName, timeDomain, (WindowedValue<KV<Object, Timer>>) timer)); |
| } |
| |
| finishFunctionRegistry.register(pTransformId, runner::finishBundle); |
| return runner; |
| } |
| |
| abstract RunnerT createRunner(Context<FnInputT, OutputT> context); |
| |
| static class Context<InputT, OutputT> { |
| final PipelineOptions pipelineOptions; |
| final BeamFnStateClient beamFnStateClient; |
| final String ptransformId; |
| final PTransform pTransform; |
| final Supplier<String> processBundleInstructionId; |
| final RehydratedComponents rehydratedComponents; |
| final DoFn<InputT, OutputT> doFn; |
| final DoFnSignature doFnSignature; |
| final TupleTag<OutputT> mainOutputTag; |
| final Coder<?> inputCoder; |
| final SchemaCoder<InputT> schemaCoder; |
| final Coder<?> keyCoder; |
| final SchemaCoder<OutputT> mainOutputSchemaCoder; |
| final Coder<? extends BoundedWindow> windowCoder; |
| final WindowingStrategy<InputT, ?> windowingStrategy; |
| final Map<TupleTag<?>, SideInputSpec> tagToSideInputSpecMap; |
| Map<TupleTag<?>, Coder<?>> outputCoders; |
| final ParDoPayload parDoPayload; |
| final ListMultimap<String, FnDataReceiver<WindowedValue<?>>> localNameToConsumer; |
| final BundleSplitListener splitListener; |
| |
| Context( |
| PipelineOptions pipelineOptions, |
| BeamFnStateClient beamFnStateClient, |
| String ptransformId, |
| PTransform pTransform, |
| Supplier<String> processBundleInstructionId, |
| Map<String, PCollection> pCollections, |
| Map<String, RunnerApi.Coder> coders, |
| Map<String, RunnerApi.WindowingStrategy> windowingStrategies, |
| PCollectionConsumerRegistry pCollectionConsumerRegistry, |
| BundleSplitListener splitListener) { |
| this.pipelineOptions = pipelineOptions; |
| this.beamFnStateClient = beamFnStateClient; |
| this.ptransformId = ptransformId; |
| this.pTransform = pTransform; |
| this.processBundleInstructionId = processBundleInstructionId; |
| ImmutableMap.Builder<TupleTag<?>, SideInputSpec> tagToSideInputSpecMapBuilder = |
| ImmutableMap.builder(); |
| try { |
| rehydratedComponents = |
| RehydratedComponents.forComponents( |
| RunnerApi.Components.newBuilder() |
| .putAllCoders(coders) |
| .putAllPcollections(pCollections) |
| .putAllWindowingStrategies(windowingStrategies) |
| .build()) |
| .withPipeline(Pipeline.create()); |
| parDoPayload = ParDoPayload.parseFrom(pTransform.getSpec().getPayload()); |
| doFn = (DoFn) ParDoTranslation.getDoFn(parDoPayload); |
| doFnSignature = DoFnSignatures.signatureForDoFn(doFn); |
| mainOutputTag = (TupleTag) ParDoTranslation.getMainOutputTag(parDoPayload); |
| String mainInputTag = |
| Iterables.getOnlyElement( |
| Sets.difference( |
| pTransform.getInputsMap().keySet(), |
| Sets.union( |
| parDoPayload.getSideInputsMap().keySet(), |
| parDoPayload.getTimerSpecsMap().keySet()))); |
| PCollection mainInput = pCollections.get(pTransform.getInputsOrThrow(mainInputTag)); |
| inputCoder = rehydratedComponents.getCoder(mainInput.getCoderId()); |
| if (inputCoder instanceof KvCoder |
| // TODO: Stop passing windowed value coders within PCollections. |
| || (inputCoder instanceof WindowedValue.WindowedValueCoder |
| && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof KvCoder))) { |
| this.keyCoder = |
| inputCoder instanceof WindowedValueCoder |
| ? ((KvCoder) ((WindowedValueCoder) inputCoder).getValueCoder()).getKeyCoder() |
| : ((KvCoder) inputCoder).getKeyCoder(); |
| } else { |
| this.keyCoder = null; |
| } |
| if (inputCoder instanceof SchemaCoder |
| // TODO: Stop passing windowed value coders within PCollections. |
| || (inputCoder instanceof WindowedValue.WindowedValueCoder |
| && (((WindowedValueCoder) inputCoder).getValueCoder() instanceof SchemaCoder))) { |
| this.schemaCoder = |
| inputCoder instanceof WindowedValueCoder |
| ? (SchemaCoder<InputT>) ((WindowedValueCoder) inputCoder).getValueCoder() |
| : ((SchemaCoder<InputT>) inputCoder); |
| } else { |
| this.schemaCoder = null; |
| } |
| |
| windowingStrategy = |
| (WindowingStrategy) |
| rehydratedComponents.getWindowingStrategy(mainInput.getWindowingStrategyId()); |
| windowCoder = windowingStrategy.getWindowFn().windowCoder(); |
| |
| outputCoders = Maps.newHashMap(); |
| for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) { |
| TupleTag<?> outputTag = new TupleTag<>(entry.getKey()); |
| RunnerApi.PCollection outputPCollection = pCollections.get(entry.getValue()); |
| Coder<?> outputCoder = rehydratedComponents.getCoder(outputPCollection.getCoderId()); |
| if (outputCoder instanceof WindowedValueCoder) { |
| outputCoder = ((WindowedValueCoder) outputCoder).getValueCoder(); |
| } |
| outputCoders.put(outputTag, outputCoder); |
| } |
| Coder<OutputT> outputCoder = (Coder<OutputT>) outputCoders.get(mainOutputTag); |
| mainOutputSchemaCoder = |
| (outputCoder instanceof SchemaCoder) ? (SchemaCoder<OutputT>) outputCoder : null; |
| |
| // Build the map from tag id to side input specification |
| for (Map.Entry<String, RunnerApi.SideInput> entry : |
| parDoPayload.getSideInputsMap().entrySet()) { |
| String sideInputTag = entry.getKey(); |
| RunnerApi.SideInput sideInput = entry.getValue(); |
| checkArgument( |
| Materializations.MULTIMAP_MATERIALIZATION_URN.equals( |
| sideInput.getAccessPattern().getUrn()), |
| "This SDK is only capable of dealing with %s materializations " |
| + "but was asked to handle %s for PCollectionView with tag %s.", |
| Materializations.MULTIMAP_MATERIALIZATION_URN, |
| sideInput.getAccessPattern().getUrn(), |
| sideInputTag); |
| |
| PCollection sideInputPCollection = |
| pCollections.get(pTransform.getInputsOrThrow(sideInputTag)); |
| WindowingStrategy sideInputWindowingStrategy = |
| rehydratedComponents.getWindowingStrategy( |
| sideInputPCollection.getWindowingStrategyId()); |
| tagToSideInputSpecMapBuilder.put( |
| new TupleTag<>(entry.getKey()), |
| SideInputSpec.create( |
| rehydratedComponents.getCoder(sideInputPCollection.getCoderId()), |
| sideInputWindowingStrategy.getWindowFn().windowCoder(), |
| PCollectionViewTranslation.viewFnFromProto(entry.getValue().getViewFn()), |
| PCollectionViewTranslation.windowMappingFnFromProto( |
| entry.getValue().getWindowMappingFn()))); |
| } |
| } catch (IOException exn) { |
| throw new IllegalArgumentException("Malformed ParDoPayload", exn); |
| } |
| |
| ImmutableListMultimap.Builder<String, FnDataReceiver<WindowedValue<?>>> |
| localNameToConsumerBuilder = ImmutableListMultimap.builder(); |
| for (Map.Entry<String, String> entry : pTransform.getOutputsMap().entrySet()) { |
| localNameToConsumerBuilder.putAll( |
| entry.getKey(), pCollectionConsumerRegistry.getMultiplexingConsumer(entry.getValue())); |
| } |
| localNameToConsumer = localNameToConsumerBuilder.build(); |
| tagToSideInputSpecMap = tagToSideInputSpecMapBuilder.build(); |
| this.splitListener = splitListener; |
| } |
| } |
| } |
