| /* |
| * 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.runners.dataflow.worker; |
| |
| import static org.apache.beam.runners.dataflow.util.Structs.getBytes; |
| import static org.apache.beam.runners.dataflow.util.Structs.getObject; |
| import static org.apache.beam.runners.dataflow.util.Structs.getString; |
| import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; |
| |
| import com.google.api.services.dataflow.model.SideInputInfo; |
| import java.util.ArrayList; |
| import java.util.List; |
| import java.util.Map; |
| import javax.annotation.Nullable; |
| import org.apache.beam.model.pipeline.v1.RunnerApi; |
| import org.apache.beam.model.pipeline.v1.RunnerApi.MessageWithComponents.RootCase; |
| import org.apache.beam.runners.core.NullSideInputReader; |
| import org.apache.beam.runners.core.SideInputReader; |
| import org.apache.beam.runners.core.StateInternalsFactory; |
| import org.apache.beam.runners.core.construction.RehydratedComponents; |
| import org.apache.beam.runners.core.construction.WindowingStrategyTranslation; |
| import org.apache.beam.runners.dataflow.DataflowRunner; |
| import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions; |
| import org.apache.beam.runners.dataflow.util.CloudObject; |
| import org.apache.beam.runners.dataflow.util.CloudObjects; |
| import org.apache.beam.runners.dataflow.util.PropertyNames; |
| import org.apache.beam.runners.dataflow.worker.util.BatchGroupAlsoByWindowsDoFns; |
| import org.apache.beam.runners.dataflow.worker.util.WorkerPropertyNames; |
| import org.apache.beam.runners.dataflow.worker.util.common.worker.ParDoFn; |
| import org.apache.beam.sdk.coders.CannotProvideCoderException; |
| import org.apache.beam.sdk.coders.Coder; |
| import org.apache.beam.sdk.coders.CoderRegistry; |
| import org.apache.beam.sdk.coders.KvCoder; |
| import org.apache.beam.sdk.coders.ListCoder; |
| import org.apache.beam.sdk.options.PipelineOptions; |
| import org.apache.beam.sdk.options.StreamingOptions; |
| import org.apache.beam.sdk.transforms.Combine.CombineFn; |
| import org.apache.beam.sdk.transforms.CombineFnBase.GlobalCombineFn; |
| import org.apache.beam.sdk.transforms.CombineWithContext.CombineFnWithContext; |
| import org.apache.beam.sdk.transforms.CombineWithContext.Context; |
| import org.apache.beam.sdk.util.AppliedCombineFn; |
| import org.apache.beam.sdk.util.SerializableUtils; |
| import org.apache.beam.sdk.util.WindowedValue.WindowedValueCoder; |
| import org.apache.beam.sdk.values.TupleTag; |
| import org.apache.beam.sdk.values.WindowingStrategy; |
| import org.apache.beam.vendor.grpc.v1p21p0.com.google.protobuf.InvalidProtocolBufferException; |
| import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables; |
| import org.slf4j.Logger; |
| import org.slf4j.LoggerFactory; |
| |
| /** |
| * A {@link ParDoFnFactory} to create GroupAlsoByWindowsDoFn instances according to specifications |
| * from the Dataflow service. |
| */ |
| class GroupAlsoByWindowParDoFnFactory implements ParDoFnFactory { |
| |
| private static final Logger LOG = LoggerFactory.getLogger(GroupAlsoByWindowParDoFnFactory.class); |
| |
| @Override |
| public ParDoFn create( |
| PipelineOptions options, |
| CloudObject cloudUserFn, |
| @Nullable List<SideInputInfo> sideInputInfos, |
| TupleTag<?> mainOutputTag, |
| Map<TupleTag<?>, Integer> outputTupleTagsToReceiverIndices, |
| final DataflowExecutionContext<?> executionContext, |
| DataflowOperationContext operationContext) |
| throws Exception { |
| Map.Entry<TupleTag<?>, Integer> entry = |
| Iterables.getOnlyElement(outputTupleTagsToReceiverIndices.entrySet()); |
| checkArgument( |
| entry.getKey().equals(mainOutputTag), |
| "Output tags should reference only the main output tag: %s vs %s", |
| entry.getKey(), |
| mainOutputTag); |
| checkArgument( |
| entry.getValue() == 0, |
| "There should be a single receiver, but using receiver index %s", |
| entry.getValue()); |
| |
| byte[] encodedWindowingStrategy = getBytes(cloudUserFn, PropertyNames.SERIALIZED_FN); |
| WindowingStrategy windowingStrategy; |
| try { |
| windowingStrategy = deserializeWindowingStrategy(encodedWindowingStrategy); |
| } catch (Exception e) { |
| // Temporarily choose default windowing strategy if fn API is enabled. |
| // TODO: Catch block disappears, becoming an error once Python SDK is compliant. |
| if (DataflowRunner.hasExperiment( |
| options.as(DataflowPipelineDebugOptions.class), "beam_fn_api")) { |
| LOG.info("FnAPI: Unable to deserialize windowing strategy, assuming default", e); |
| windowingStrategy = WindowingStrategy.globalDefault(); |
| } else { |
| throw e; |
| } |
| } |
| |
| byte[] serializedCombineFn = getBytes(cloudUserFn, WorkerPropertyNames.COMBINE_FN, null); |
| AppliedCombineFn<?, ?, ?, ?> combineFn = null; |
| if (serializedCombineFn != null) { |
| Object combineFnObj = |
| SerializableUtils.deserializeFromByteArray(serializedCombineFn, "serialized combine fn"); |
| checkArgument( |
| combineFnObj instanceof AppliedCombineFn, |
| "unexpected kind of AppliedCombineFn: " + combineFnObj.getClass().getName()); |
| combineFn = (AppliedCombineFn<?, ?, ?, ?>) combineFnObj; |
| } |
| |
| Map<String, Object> inputCoderObject = getObject(cloudUserFn, WorkerPropertyNames.INPUT_CODER); |
| |
| Coder<?> inputCoder = CloudObjects.coderFromCloudObject(CloudObject.fromSpec(inputCoderObject)); |
| checkArgument( |
| inputCoder instanceof WindowedValueCoder, |
| "Expected WindowedValueCoder for inputCoder, got: " + inputCoder.getClass().getName()); |
| @SuppressWarnings("unchecked") |
| WindowedValueCoder<?> windowedValueCoder = (WindowedValueCoder<?>) inputCoder; |
| |
| Coder<?> elemCoder = windowedValueCoder.getValueCoder(); |
| checkArgument( |
| elemCoder instanceof KvCoder, |
| "Expected KvCoder for inputCoder, got: " + elemCoder.getClass().getName()); |
| @SuppressWarnings("unchecked") |
| KvCoder<?, ?> kvCoder = (KvCoder<?, ?>) elemCoder; |
| |
| boolean isStreamingPipeline = options.as(StreamingOptions.class).isStreaming(); |
| |
| SideInputReader sideInputReader = NullSideInputReader.empty(); |
| @Nullable AppliedCombineFn<?, ?, ?, ?> maybeMergingCombineFn = null; |
| if (combineFn != null) { |
| sideInputReader = |
| executionContext.getSideInputReader( |
| sideInputInfos, combineFn.getSideInputViews(), operationContext); |
| |
| String phase = getString(cloudUserFn, WorkerPropertyNames.PHASE, CombinePhase.ALL); |
| checkArgument( |
| phase.equals(CombinePhase.ALL) || phase.equals(CombinePhase.MERGE), |
| "Unexpected phase: %s", |
| phase); |
| if (phase.equals(CombinePhase.MERGE)) { |
| maybeMergingCombineFn = makeAppliedMergingFunction(combineFn); |
| } else { |
| maybeMergingCombineFn = combineFn; |
| } |
| } |
| |
| StateInternalsFactory<?> stateInternalsFactory = |
| key -> executionContext.getStepContext(operationContext).stateInternals(); |
| |
| // This will be a GABW Fn for either batch or streaming, with combiner in it or not |
| GroupAlsoByWindowFn<?, ?> fn; |
| |
| // This will be a FakeKeyedWorkItemCoder for streaming or null for batch |
| Coder<?> gabwInputCoder; |
| |
| // TODO: do not do this with mess of "if" |
| if (isStreamingPipeline) { |
| if (maybeMergingCombineFn == null) { |
| fn = |
| StreamingGroupAlsoByWindowsDoFns.createForIterable( |
| windowingStrategy, stateInternalsFactory, ((KvCoder) kvCoder).getValueCoder()); |
| gabwInputCoder = WindmillKeyedWorkItem.FakeKeyedWorkItemCoder.of(kvCoder); |
| } else { |
| fn = |
| StreamingGroupAlsoByWindowsDoFns.create( |
| windowingStrategy, |
| stateInternalsFactory, |
| (AppliedCombineFn) maybeMergingCombineFn, |
| ((KvCoder) kvCoder).getKeyCoder()); |
| gabwInputCoder = |
| WindmillKeyedWorkItem.FakeKeyedWorkItemCoder.of( |
| ((AppliedCombineFn) maybeMergingCombineFn).getKvCoder()); |
| } |
| |
| } else { |
| if (maybeMergingCombineFn == null) { |
| fn = |
| BatchGroupAlsoByWindowsDoFns.createForIterable( |
| windowingStrategy, stateInternalsFactory, ((KvCoder) kvCoder).getValueCoder()); |
| gabwInputCoder = null; |
| } else { |
| fn = |
| BatchGroupAlsoByWindowsDoFns.create( |
| windowingStrategy, (AppliedCombineFn) maybeMergingCombineFn); |
| gabwInputCoder = null; |
| } |
| } |
| |
| // TODO: or anyhow related to it, do not do this with mess of "if" |
| if (maybeMergingCombineFn != null) { |
| return new GroupAlsoByWindowsParDoFn( |
| options, |
| fn, |
| windowingStrategy, |
| ((AppliedCombineFn) maybeMergingCombineFn).getSideInputViews(), |
| gabwInputCoder, |
| sideInputReader, |
| mainOutputTag, |
| executionContext.getStepContext(operationContext)); |
| } else { |
| return new GroupAlsoByWindowsParDoFn( |
| options, |
| fn, |
| windowingStrategy, |
| null, |
| gabwInputCoder, |
| sideInputReader, |
| mainOutputTag, |
| executionContext.getStepContext(operationContext)); |
| } |
| } |
| |
| static WindowingStrategy deserializeWindowingStrategy(byte[] encodedWindowingStrategy) |
| throws InvalidProtocolBufferException { |
| RunnerApi.MessageWithComponents strategyProto = |
| RunnerApi.MessageWithComponents.parseFrom(encodedWindowingStrategy); |
| checkArgument( |
| strategyProto.getRootCase() == RootCase.WINDOWING_STRATEGY, |
| "Invalid windowing strategy: %s", |
| strategyProto); |
| return WindowingStrategyTranslation.fromProto( |
| strategyProto.getWindowingStrategy(), |
| RehydratedComponents.forComponents(strategyProto.getComponents())); |
| } |
| |
| private static <K, AccumT> |
| AppliedCombineFn<K, AccumT, List<AccumT>, AccumT> makeAppliedMergingFunction( |
| AppliedCombineFn<K, ?, AccumT, ?> appliedFn) { |
| GlobalCombineFn<AccumT, List<AccumT>, AccumT> mergingCombineFn; |
| if (appliedFn.getFn() instanceof CombineFnWithContext) { |
| mergingCombineFn = |
| new MergingKeyedCombineFnWithContext<>( |
| (CombineFnWithContext<?, AccumT, ?>) appliedFn.getFn(), |
| appliedFn.getAccumulatorCoder()); |
| } else { |
| mergingCombineFn = |
| new MergingCombineFn<>( |
| (CombineFn<?, AccumT, ?>) appliedFn.getFn(), appliedFn.getAccumulatorCoder()); |
| } |
| return AppliedCombineFn.<K, AccumT, List<AccumT>, AccumT>withAccumulatorCoder( |
| mergingCombineFn, |
| ListCoder.of(appliedFn.getAccumulatorCoder()), |
| appliedFn.getSideInputViews(), |
| KvCoder.of(appliedFn.getKvCoder().getKeyCoder(), appliedFn.getAccumulatorCoder()), |
| appliedFn.getWindowingStrategy()); |
| } |
| |
| private static int MAX_ACCUMULATOR_BUFFER_SIZE = 10; |
| |
| static class MergingCombineFn<K, AccumT> extends CombineFn<AccumT, List<AccumT>, AccumT> { |
| |
| private final CombineFn<?, AccumT, ?> combineFn; |
| private final Coder<AccumT> accumCoder; |
| |
| MergingCombineFn(CombineFn<?, AccumT, ?> combineFn, Coder<AccumT> accumCoder) { |
| this.combineFn = combineFn; |
| this.accumCoder = accumCoder; |
| } |
| |
| @Override |
| public List<AccumT> createAccumulator() { |
| ArrayList<AccumT> result = new ArrayList<>(); |
| result.add(this.combineFn.createAccumulator()); |
| return result; |
| } |
| |
| @Override |
| public List<AccumT> addInput(List<AccumT> accumulator, AccumT input) { |
| accumulator.add(input); |
| if (accumulator.size() < MAX_ACCUMULATOR_BUFFER_SIZE) { |
| return accumulator; |
| } else { |
| return mergeToSingleton(accumulator); |
| } |
| } |
| |
| @Override |
| public List<AccumT> mergeAccumulators(Iterable<List<AccumT>> accumulators) { |
| return mergeToSingleton(Iterables.concat(accumulators)); |
| } |
| |
| @Override |
| public List<AccumT> compact(List<AccumT> accumulator) { |
| return mergeToSingleton(accumulator); |
| } |
| |
| @Override |
| public AccumT extractOutput(List<AccumT> accumulator) { |
| if (accumulator.isEmpty()) { |
| return combineFn.createAccumulator(); |
| } else { |
| return combineFn.mergeAccumulators(accumulator); |
| } |
| } |
| |
| private List<AccumT> mergeToSingleton(Iterable<AccumT> accumulators) { |
| List<AccumT> singleton = new ArrayList<>(); |
| singleton.add(combineFn.mergeAccumulators(accumulators)); |
| return singleton; |
| } |
| |
| @Override |
| public Coder<List<AccumT>> getAccumulatorCoder(CoderRegistry registry, Coder<AccumT> inputCoder) |
| throws CannotProvideCoderException { |
| return ListCoder.of(accumCoder); |
| } |
| } |
| |
| static class MergingKeyedCombineFnWithContext<K, AccumT> |
| extends CombineFnWithContext<AccumT, List<AccumT>, AccumT> { |
| |
| private final CombineFnWithContext<?, AccumT, ?> combineFnWithContext; |
| private final Coder<AccumT> accumCoder; |
| |
| MergingKeyedCombineFnWithContext( |
| CombineFnWithContext<?, AccumT, ?> combineFnWithContext, Coder<AccumT> accumCoder) { |
| this.combineFnWithContext = combineFnWithContext; |
| this.accumCoder = accumCoder; |
| } |
| |
| @Override |
| public List<AccumT> createAccumulator(Context c) { |
| ArrayList<AccumT> result = new ArrayList<>(); |
| result.add(this.combineFnWithContext.createAccumulator(c)); |
| return result; |
| } |
| |
| @Override |
| public List<AccumT> addInput(List<AccumT> accumulator, AccumT input, Context c) { |
| accumulator.add(input); |
| if (accumulator.size() < MAX_ACCUMULATOR_BUFFER_SIZE) { |
| return accumulator; |
| } else { |
| return mergeToSingleton(accumulator, c); |
| } |
| } |
| |
| @Override |
| public List<AccumT> mergeAccumulators(Iterable<List<AccumT>> accumulators, Context c) { |
| return mergeToSingleton(Iterables.concat(accumulators), c); |
| } |
| |
| @Override |
| public List<AccumT> compact(List<AccumT> accumulator, Context c) { |
| return mergeToSingleton(accumulator, c); |
| } |
| |
| @Override |
| public AccumT extractOutput(List<AccumT> accumulator, Context c) { |
| if (accumulator.isEmpty()) { |
| return combineFnWithContext.createAccumulator(c); |
| } else { |
| return combineFnWithContext.mergeAccumulators(accumulator, c); |
| } |
| } |
| |
| private List<AccumT> mergeToSingleton(Iterable<AccumT> accumulators, Context c) { |
| List<AccumT> singleton = new ArrayList<>(); |
| singleton.add(combineFnWithContext.mergeAccumulators(accumulators, c)); |
| return singleton; |
| } |
| |
| @Override |
| public Coder<List<AccumT>> getAccumulatorCoder(CoderRegistry registry, Coder<AccumT> inputCoder) |
| throws CannotProvideCoderException { |
| return ListCoder.of(accumCoder); |
| } |
| } |
| } |
