Abacn commented on code in PR #39793: URL: https://github.com/apache/beam/pull/39793#discussion_r3809794788
########## runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java: ########## @@ -0,0 +1,316 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.oneOfEncoder; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; +import static org.apache.spark.sql.functions.col; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.spark.api.java.function.FlatMapGroupsFunction; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.Column; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.KeyValueGroupedDataset; +import org.apache.spark.sql.TypedColumn; +import org.apache.spark.storage.StorageLevel; +import scala.Tuple2; + +/** + * Translator for a stateful {@link ParDo.MultiOutput}, or one requiring time sorted input. + * + * <p>Selected by {@link PipelineTranslatorBatch} in place of {@link ParDoTranslatorBatch} when the + * {@link DoFn} uses state, uses timers, or is annotated with {@link DoFn.RequiresTimeSortedInput}; + * see {@link #appliesTo}. + * + * <p>Unlike {@link ParDoTranslatorBatch} this translator never produces an {@code + * UnresolvedTranslation}: a stateful {@link DoFn} must not be fused with neighbouring {@link ParDo + * ParDos}, because the fused runner cannot drive timers. Resolving the input dataset via {@code + * Context#getDataset} breaks any pending fusion chain. + * + * <p>Additional (tagged) outputs are encoded as one column per tag, as in {@link + * ParDoTranslatorBatch}. + */ +@SuppressWarnings({"rawtypes", "unchecked"}) Review Comment: New codes should avoid suppressWarnings ########## runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulDoFnGroupFunction.java: ########## @@ -0,0 +1,392 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.utils.ScalaInterop.tuple; + +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.Iterator; +import java.util.Map; +import java.util.function.Supplier; +import javax.annotation.CheckForNull; +import org.apache.beam.runners.core.InMemoryStateInternals; +import org.apache.beam.runners.core.InMemoryTimerInternals; +import org.apache.beam.runners.core.StateInternals; +import org.apache.beam.runners.core.StateNamespaces; +import org.apache.beam.runners.core.StepContext; +import org.apache.beam.runners.core.TimerInternals; +import org.apache.beam.runners.core.TimerInternals.TimerData; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.DoFnRunnerFactory.DoFnRunnerWithTeardown; +import org.apache.beam.sdk.options.PipelineOptions; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.windowing.BoundedWindow; +import org.apache.beam.sdk.util.WindowedValueMultiReceiver; +import org.apache.beam.sdk.values.CausedByDrain; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.AbstractIterator; +import org.apache.spark.TaskContext; +import org.apache.spark.api.java.function.FlatMapGroupsFunction; +import org.apache.spark.util.TaskCompletionListener; +import org.checkerframework.checker.nullness.qual.NonNull; Review Comment: NonNull is default, no need to annotate with it in codes. ########## runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/PipelineTranslatorBatch.java: ########## @@ -82,11 +82,27 @@ public class PipelineTranslatorBatch extends PipelineTranslator { SplittableParDo.PrimitiveBoundedRead.class, new ReadSourceTranslatorBatch<>()); } + /** + * Translators that shadow the {@link #TRANSFORM_TRANSLATORS} entry for their transform class when + * a predicate matches, so that a single transform class can be translated in more than one way + * depending on the transform instance. + * + * <p>Currently only {@link ParDo.MultiOutput} needs this, to route stateful and time sorted + * {@link org.apache.beam.sdk.transforms.DoFn DoFns} away from {@link ParDoTranslatorBatch}. + */ + @SuppressWarnings("rawtypes") + private static final TransformTranslator STATEFUL_PARDO_TRANSLATOR = + new StatefulParDoTranslatorBatch<>(); + /** Returns a {@link TransformTranslator} for the given {@link PTransform} if known. */ @Override @Nullable protected <InT extends PInput, OutT extends POutput, TransformT extends PTransform<InT, OutT>> TransformTranslator<InT, OutT, TransformT> getTransformTranslator(TransformT transform) { + if (transform instanceof ParDo.MultiOutput Review Comment: We introduced a special case. Consider adding a brief comment note why breaking the current pattern is necessary: ParDo.MultiOutput may map to different translator depends on DoFn signature ########## runners/spark/src/main/java/org/apache/beam/runners/spark/structuredstreaming/translation/batch/StatefulParDoTranslatorBatch.java: ########## @@ -0,0 +1,316 @@ +/* + * 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.spark.structuredstreaming.translation.batch; + +import static org.apache.beam.runners.spark.structuredstreaming.translation.helpers.EncoderHelpers.oneOfEncoder; +import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull; +import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState; +import static org.apache.spark.sql.functions.col; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import java.util.Map; +import org.apache.beam.runners.core.SideInputReader; +import org.apache.beam.runners.spark.SparkCommonPipelineOptions; +import org.apache.beam.runners.spark.structuredstreaming.metrics.MetricsAccumulator; +import org.apache.beam.runners.spark.structuredstreaming.translation.TransformTranslator; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SideInputValues; +import org.apache.beam.runners.spark.structuredstreaming.translation.batch.functions.SparkSideInputReader; +import org.apache.beam.sdk.coders.Coder; +import org.apache.beam.sdk.coders.KvCoder; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.transforms.reflect.DoFnSignature; +import org.apache.beam.sdk.transforms.reflect.DoFnSignatures; +import org.apache.beam.sdk.values.KV; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionTuple; +import org.apache.beam.sdk.values.PCollectionView; +import org.apache.beam.sdk.values.TupleTag; +import org.apache.beam.sdk.values.WindowedValue; +import org.apache.beam.sdk.values.WindowingStrategy; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Maps; +import org.apache.spark.api.java.function.FlatMapGroupsFunction; +import org.apache.spark.broadcast.Broadcast; +import org.apache.spark.sql.Column; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.KeyValueGroupedDataset; +import org.apache.spark.sql.TypedColumn; +import org.apache.spark.storage.StorageLevel; +import scala.Tuple2; + +/** + * Translator for a stateful {@link ParDo.MultiOutput}, or one requiring time sorted input. + * + * <p>Selected by {@link PipelineTranslatorBatch} in place of {@link ParDoTranslatorBatch} when the + * {@link DoFn} uses state, uses timers, or is annotated with {@link DoFn.RequiresTimeSortedInput}; + * see {@link #appliesTo}. + * + * <p>Unlike {@link ParDoTranslatorBatch} this translator never produces an {@code + * UnresolvedTranslation}: a stateful {@link DoFn} must not be fused with neighbouring {@link ParDo + * ParDos}, because the fused runner cannot drive timers. Resolving the input dataset via {@code + * Context#getDataset} breaks any pending fusion chain. + * + * <p>Additional (tagged) outputs are encoded as one column per tag, as in {@link + * ParDoTranslatorBatch}. + */ +@SuppressWarnings({"rawtypes", "unchecked"}) +class StatefulParDoTranslatorBatch<K, V, OutputT> + extends TransformTranslator< + PCollection<? extends KV<K, V>>, PCollectionTuple, ParDo.MultiOutput<KV<K, V>, OutputT>> { + + StatefulParDoTranslatorBatch() { + // A stateful ParDo introduces a shuffle to co-locate and order each key, so it contributes to + // plan complexity much like GroupByKey rather than like a plain ParDo. + super(0.2f); + } + + /** + * Whether {@code transform} must be translated by this translator rather than {@link + * ParDoTranslatorBatch}. + * + * <p>Note {@link DoFn.RequiresTimeSortedInput} is tested independently of state: the SDK only + * treats state and timers as making a {@link DoFn} stateful, so a {@code DoFn} carrying only that + * annotation reaches the runner with neither signature flag set. + */ + static boolean appliesTo(ParDo.MultiOutput<?, ?> transform) { + DoFnSignature signature = DoFnSignatures.signatureForDoFn(transform.getFn()); + return signature.usesState() + || signature.usesTimers() + || signature.processElement().requiresTimeSortedInput(); + } + + @Override + protected boolean canTranslate(ParDo.MultiOutput<KV<K, V>, OutputT> transform) { + DoFn<KV<K, V>, OutputT> doFn = transform.getFn(); + DoFnSignature signature = DoFnSignatures.signatureForDoFn(doFn); + + checkState( + appliesTo(transform), + "Not a stateful or time sorted DoFn, should have been translated by %s: %s", + ParDoTranslatorBatch.class.getSimpleName(), + doFn); + + checkState( + isSupported(), + "Stateful and time sorted ParDo require Spark 3.4+ " + + "(KeyValueGroupedDataset#flatMapSortedGroups): %s", + doFn); + + checkState( + !signature.processElement().isSplittable(), + "Not expected to directly translate splittable DoFn, should have been overridden: %s", + doFn); + + // Not implemented: firing @OnWindowExpiration requires tracking the windows observed per key + // and a dedicated firing pass at the end of each key, see + // https://github.com/apache/beam/issues/22524 + checkState( + signature.onWindowExpiration() == null, "onWindowExpiration is not supported: %s", doFn); + + SparkSideInputReader.validateMaterializations(transform.getSideInputs().values()); + return true; + } + + @Override + protected void translate(ParDo.MultiOutput<KV<K, V>, OutputT> transform, Context cxt) + throws IOException { + PCollection<KV<K, V>> input = (PCollection<KV<K, V>>) cxt.getInput(); + + validateKeyCoder(input.getCoder(), transform.getFn()); + validateWindowingStrategy(input.getWindowingStrategy(), transform.getFn()); + + TupleTag<OutputT> mainOut = transform.getMainOutputTag(); + // Filter out obsolete PCollections to only cache when absolutely necessary + Map<TupleTag<?>, PCollection<?>> outputs = + ParDoTranslatorBatch.skipUnconsumedOutputs( + cxt.getOutputs(), mainOut, transform.getAdditionalOutputTags(), cxt); + + KvCoder<K, V> inputCoder = (KvCoder<K, V>) input.getCoder(); + Encoder<K> keyEnc = cxt.keyEncoderOf(inputCoder); + MetricsAccumulator metrics = MetricsAccumulator.getInstance(cxt.getSparkSession()); + SideInputReader sideInputReader = createSideInputReader(transform, cxt); + + // Group by key, then order each group by event time before handing it to the DoFn. The + // timestamp is a top level LongType column of the WindowedValue encoder (epoch millis), so + // ordering is plain signed numeric ordering; no composite sort key is needed. Nulls sort + // last: a null timestamp encodes END_OF_WINDOW (see GroupByKeyTranslatorBatch), which no + // concrete timestamp of the same window can exceed. Only null and concrete timestamps of + // different windows mixed into one key group may still order imprecisely; deriving the + // timestamp from the window column is not portable across Spark versions. + Column[] sortCols = new Column[] {col(TIMESTAMP_COLUMN).asc_nulls_last()}; + + if (outputs.size() > 1) { + // In case of multiple outputs / tags, map each tag to a column by index. + // At the end split the result into multiple datasets selecting one column each. + Map<String, Integer> tagColIdx = ParDoTranslatorBatch.tagsColumnIndex(outputs.keySet()); + List<Encoder<WindowedValue<Object>>> encoders = createEncoders(outputs, tagColIdx, cxt); + + DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory = + DoFnRunnerFactory.simple(cxt.getCurrentTransform(), input, sideInputReader, false); + StatefulDoFnGroupFunction<K, KV<K, V>, Tuple2<Integer, WindowedValue<Object>>> groupFn = + StatefulDoFnGroupFunction.multiOutput( + cxt.getOptionsSupplier(), metrics, runnerFactory, tagColIdx); + + SparkCommonPipelineOptions opts = cxt.getOptions().as(SparkCommonPipelineOptions.class); + StorageLevel storageLevel = StorageLevel.fromString(opts.getStorageLevel()); + + // Persist as wide rows with one column per TupleTag to support different schemas + Dataset<Tuple2<Integer, WindowedValue<Object>>> allTagsDS = + cxt.getDataset(input) + .groupByKey(GroupByKeyHelpers.valueKey(), keyEnc) + .flatMapSortedGroups(sortCols, groupFn, oneOfEncoder(encoders)); + allTagsDS.persist(storageLevel); + + // divide into separate output datasets per tag + for (TupleTag<?> tag : outputs.keySet()) { + int colIdx = checkStateNotNull(tagColIdx.get(tag.getId()), "Unknown tag"); + // Resolve specific column matching the tuple tag (by id) + TypedColumn<Tuple2<Integer, WindowedValue<Object>>, WindowedValue<Object>> col = + (TypedColumn) col(Integer.toString(colIdx)).as(encoders.get(colIdx)); + + // Caching of the returned outputs is disabled to avoid caching the same data twice. + cxt.putDataset( + cxt.getOutput((TupleTag) tag), allTagsDS.filter(col.isNotNull()).select(col), false); + } + } else { + PCollection<OutputT> output = cxt.getOutput(mainOut); + // Obsolete outputs might have to be filtered out + boolean filterMainOutput = cxt.getOutputs().size() > 1; + DoFnRunnerFactory<KV<K, V>, OutputT> runnerFactory = + DoFnRunnerFactory.simple( + cxt.getCurrentTransform(), input, sideInputReader, filterMainOutput); + StatefulDoFnGroupFunction<K, KV<K, V>, WindowedValue<OutputT>> groupFn = + StatefulDoFnGroupFunction.singleOutput(cxt.getOptionsSupplier(), metrics, runnerFactory); + + Dataset<WindowedValue<OutputT>> result = + cxt.getDataset(input) + .groupByKey(GroupByKeyHelpers.valueKey(), keyEnc) + .flatMapSortedGroups(sortCols, groupFn, cxt.windowedEncoder(output.getCoder())); + + cxt.putDataset(output, result); + } + } + + /** List of encoders matching the order of tagIds. */ + private List<Encoder<WindowedValue<Object>>> createEncoders( + Map<TupleTag<?>, PCollection<?>> outputs, Map<String, Integer> tagIdColIdx, Context ctx) { + ArrayList<Encoder<WindowedValue<Object>>> encoders = new ArrayList<>(outputs.size()); + for (Map.Entry<TupleTag<?>, PCollection<?>> e : outputs.entrySet()) { + Encoder<WindowedValue<Object>> enc = ctx.windowedEncoder((Coder) e.getValue().getCoder()); + int colIdx = checkStateNotNull(tagIdColIdx.get(e.getKey().getId())); + encoders.add(colIdx, enc); Review Comment: IndexOutOfBoundsException if colIdx is out of order. new ArrayList(size) creates a list of zero size and initial capacity of "size" Consider re-allocate an array or fixed-size list: ``` Encoder<WindowedValue<Object>>[] encoders = new Encoder[outputs.size()]; for (Map.Entry<TupleTag<?>, PCollection<?>> e : outputs.entrySet()) { Encoder<WindowedValue<Object>> enc = ctx.windowedEncoder((Coder) e.getValue().getCoder()); int colIdx = checkStateNotNull(tagIdColIdx.get(e.getKey().getId())); encoders[colIdx] = enc; } return Arrays.asList(encoders); ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
