This is an automated email from the ASF dual-hosted git repository. echauchot pushed a commit to branch spark-runner_structured-streaming in repository https://gitbox.apache.org/repos/asf/beam.git
commit fb3aa34b55bdb21afc030ce618a27dcf2623f40e Author: Etienne Chauchot <echauc...@apache.org> AuthorDate: Fri Oct 11 16:18:11 2019 +0200 Add missing windowedValue Encoder call in Pardo --- .../translation/batch/ParDoTranslatorBatch.java | 37 ++++++++++++---------- 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java index e73d38e..f5a109e 100644 --- a/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java +++ b/runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/ParDoTranslatorBatch.java @@ -93,6 +93,8 @@ class ParDoTranslatorBatch<InputT, OutputT> List<TupleTag<?>> outputTags = new ArrayList<>(outputs.keySet()); WindowingStrategy<?, ?> windowingStrategy = ((PCollection<InputT>) input).getWindowingStrategy(); + Coder<InputT> inputCoder = ((PCollection<InputT>) input).getCoder(); + Coder<? extends BoundedWindow> windowCoder = windowingStrategy.getWindowFn().windowCoder(); // construct a map from side input to WindowingStrategy so that // the DoFn runner can map main-input windows to side input windows @@ -105,7 +107,6 @@ class ParDoTranslatorBatch<InputT, OutputT> SideInputBroadcast broadcastStateData = createBroadcastSideInputs(sideInputs, context); Map<TupleTag<?>, Coder<?>> outputCoderMap = context.getOutputCoders(); - Coder<InputT> inputCoder = ((PCollection<InputT>) input).getCoder(); MetricsContainerStepMapAccumulator metricsAccum = MetricsAccumulator.getInstance(); List<TupleTag<?>> additionalOutputTags = new ArrayList<>(); @@ -131,24 +132,25 @@ class ParDoTranslatorBatch<InputT, OutputT> broadcastStateData, doFnSchemaInformation); - MultiOuputCoder multipleOutputCoder = - MultiOuputCoder.of( - SerializableCoder.of(TupleTag.class), - outputCoderMap, - windowingStrategy.getWindowFn().windowCoder()); + MultiOuputCoder multipleOutputCoder = MultiOuputCoder + .of(SerializableCoder.of(TupleTag.class), outputCoderMap, windowCoder); Dataset<Tuple2<TupleTag<?>, WindowedValue<?>>> allOutputs = inputDataSet.mapPartitions(doFnWrapper, EncoderHelpers.fromBeamCoder(multipleOutputCoder)); if (outputs.entrySet().size() > 1) { allOutputs.persist(); for (Map.Entry<TupleTag<?>, PValue> output : outputs.entrySet()) { - pruneOutputFilteredByTag(context, allOutputs, output); + pruneOutputFilteredByTag(context, allOutputs, output, windowCoder); } } else { + Coder<OutputT> outputCoder = ((PCollection<OutputT>) outputs.get(mainOutputTag)).getCoder(); + Coder<WindowedValue<?>> windowedValueCoder = + (Coder<WindowedValue<?>>) + (Coder<?>) WindowedValue.getFullCoder(outputCoder, windowCoder); Dataset<WindowedValue<?>> outputDataset = allOutputs.map( (MapFunction<Tuple2<TupleTag<?>, WindowedValue<?>>, WindowedValue<?>>) value -> value._2, - EncoderHelpers.windowedValueEncoder()); + EncoderHelpers.fromBeamCoder(windowedValueCoder)); context.putDatasetWildcard(outputs.entrySet().iterator().next().getValue(), outputDataset); } } @@ -159,14 +161,14 @@ class ParDoTranslatorBatch<InputT, OutputT> JavaSparkContext.fromSparkContext(context.getSparkSession().sparkContext()); SideInputBroadcast sideInputBroadcast = new SideInputBroadcast(); - for (PCollectionView<?> input : sideInputs) { + for (PCollectionView<?> sideInput : sideInputs) { Coder<? extends BoundedWindow> windowCoder = - input.getPCollection().getWindowingStrategy().getWindowFn().windowCoder(); + sideInput.getPCollection().getWindowingStrategy().getWindowFn().windowCoder(); + Coder<WindowedValue<?>> windowedValueCoder = (Coder<WindowedValue<?>>) - (Coder<?>) WindowedValue.getFullCoder(input.getPCollection().getCoder(), windowCoder); - - Dataset<WindowedValue<?>> broadcastSet = context.getSideInputDataSet(input); + (Coder<?>) WindowedValue.getFullCoder(sideInput.getPCollection().getCoder(), windowCoder); + Dataset<WindowedValue<?>> broadcastSet = context.getSideInputDataSet(sideInput); List<WindowedValue<?>> valuesList = broadcastSet.collectAsList(); List<byte[]> codedValues = new ArrayList<>(); for (WindowedValue<?> v : valuesList) { @@ -174,7 +176,7 @@ class ParDoTranslatorBatch<InputT, OutputT> } sideInputBroadcast.add( - input.getTagInternal().getId(), jsc.broadcast(codedValues), windowedValueCoder); + sideInput.getTagInternal().getId(), jsc.broadcast(codedValues), windowedValueCoder); } return sideInputBroadcast; } @@ -213,14 +215,17 @@ class ParDoTranslatorBatch<InputT, OutputT> private void pruneOutputFilteredByTag( TranslationContext context, Dataset<Tuple2<TupleTag<?>, WindowedValue<?>>> allOutputs, - Map.Entry<TupleTag<?>, PValue> output) { + Map.Entry<TupleTag<?>, PValue> output, Coder<? extends BoundedWindow> windowCoder) { Dataset<Tuple2<TupleTag<?>, WindowedValue<?>>> filteredDataset = allOutputs.filter(new DoFnFilterFunction(output.getKey())); + Coder<WindowedValue<?>> windowedValueCoder = + (Coder<WindowedValue<?>>) + (Coder<?>) WindowedValue.getFullCoder(((PCollection<OutputT>)output.getValue()).getCoder(), windowCoder); Dataset<WindowedValue<?>> outputDataset = filteredDataset.map( (MapFunction<Tuple2<TupleTag<?>, WindowedValue<?>>, WindowedValue<?>>) value -> value._2, - EncoderHelpers.windowedValueEncoder()); + EncoderHelpers.fromBeamCoder(windowedValueCoder)); context.putDatasetWildcard(output.getValue(), outputDataset); }