You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@beam.apache.org by "f-loris (via GitHub)" <gi...@apache.org> on 2026/08/17 19:26:55 UTC
[PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
f-loris opened a new pull request, #39793:
URL: https://github.com/apache/beam/pull/39793
Adds support for stateful ParDo — state, timers and `@RequiresTimeSortedInput` — to the Spark Structured Streaming (Dataset-based) runner in batch mode, including additional (tagged) outputs.
Fixes #39779
### Approach
Stateful, timer-using and `@RequiresTimeSortedInput` `DoFn`s are dispatched to a new `StatefulParDoTranslatorBatch`, which groups by key via `KeyValueGroupedDataset#flatMapSortedGroups` and orders each group by event time. `StatefulDoFnGroupFunction` then runs the `DoFn` per key with `InMemoryStateInternals`/`InMemoryTimerInternals` from `runners-core`.
Because batch has no state store to bridge onto — state is a heap object scoped to a single key — every Beam state type is already implemented by `InMemoryStateInternals`, so the state categories are enabled together rather than incrementally. Memory is bounded by the state one key holds rather than by that key's element count: groups stream off Spark's spillable sort instead of being materialized, unlike the classic runner's RDD `groupByKey`.
Sorting uses the `timestamp` column of the `WindowedValue` encoder, which is `LongType` epoch millis, with `asc_nulls_last`: a null timestamp encodes `END_OF_WINDOW`, which no concrete timestamp of the same window can exceed.
The `DoFn` is set up once per task and torn down from a `TaskCompletionListener`, with state, timers and a bundle per key.
### ValidatesRunner
Enabled for `validatesStructuredStreamingRunnerBatch`: `UsesStatefulParDo`, `UsesKeyInParDo`, `UsesTimersInParDo`, `UsesMapState`, `UsesMultimapState`, `UsesSetState`, `UsesOrderedListState`, `UsesTimerMap`.
Results on this branch:
| Suite | Result |
|---|---|
| `:runners:spark:3:validatesStructuredStreamingRunnerBatch` | 304 tests, 0 failures |
| `:runners:spark:4:validatesStructuredStreamingRunnerBatch` | 303 tests, 0 failures |
Also added: unit tests for translator dispatch and preconditions, and execution tests covering per-key state, state isolation across many keys in one partition, `@RequiresTimeSortedInput` ordering, event-time timers, looping timers, `@FinishBundle` ordering relative to timers, and tagged outputs.
### Notes for reviewers
- **No CI covers this suite on Spark 4.** `beam_PostCommit_Java_ValidatesRunner_Spark4` runs `:runners:spark:4:validatesRunner` (the classic runner), and `beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming` is pinned to `:runners:spark:3:...`. This is pre-existing — the suite has never run on Spark 4 — so the 303/0 above was verified locally. Happy to add a `beam_PostCommit_Java_ValidatesRunner_SparkStructuredStreaming4` workflow here or as a follow-up if you'd like that coverage; it would be a near-copy of the existing one with `java-version: '17'`.
- **Spark 3.4+ requirement.** `flatMapSortedGroups` was added in Spark 3.4, so stateful pipelines are rejected at translation time on 3.1–3.3 with a clear message rather than failing with `NoSuchMethodError` mid-job. The stateful tests skip on those versions so `sparkVersion31Test`/`sparkVersion33Test` stay green.
- **`PerKeyOrderingTest#testMultipleStatefulOrdering*` are excluded by name**, not by category. They build on `PeriodicImpulse` and so are unbounded, but are not categorized as `UsesUnboundedPCollections`; adding that category upstream would also stop Dataflow and Flink running them, so the exclusion is kept local to this runner.
- **The `ParDoTest` OrderedListState range tests gain `UsesOnWindowExpiration`**, which their `DoFn`s declare but which they were not categorized for. This is the one change outside `runners/spark`. It is a no-op for other runners: every runner excluding `UsesOnWindowExpiration` also excludes `UsesOrderedListState`, so those tests did not run there.
- **Deliberate mirroring of the stateless path.** `createEncoders`, `createSideInputReader`, the output receivers and the persist/split-by-tag block intentionally follow `ParDoTranslatorBatch`/`DoFnPartitionIteratorFactory`. Happy to consolidate into shared helpers if preferred.
- **`@OnWindowExpiration` remains unsupported** (#22524) and is rejected explicitly; `UsesOnWindowExpiration` stays excluded. Processing-time timers re-armed from `@OnTimer` do not re-fire in batch, matching the classic runner; this is documented on `fireNextTimer`.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn merged PR #39793:
URL: https://github.com/apache/beam/pull/39793
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on PR #39793:
URL: https://github.com/apache/beam/pull/39793#issuecomment-5346641356
Thanks! Left a few comments.
Also, we can update https://beam.apache.org/documentation/runners/capability-matrix/ after this change is released
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "f-loris (via GitHub)" <gi...@apache.org>.
f-loris commented on PR #39793:
URL: https://github.com/apache/beam/pull/39793#issuecomment-5354423138
@Abacn Thanks, all four addressed in my recent commit.
- @SuppressWarnings: removed. The only raw declarations were in the private createSideInputReader, so I dropped it and made ParDoTranslatorBatch's version static on TranslationState, like skipUnconsumedOutputs/tagsColumnIndex which this translator already reuses.
-createEncoders: used the array form, now a single static shared by both translators (ParDoTranslatorBatch's copy had the same issue). new Encoder[n] does emit [rawtypes] and the module builds with -Xlint:all -Werror, so it carries a method-scoped suppression instead of the class-level one.
- Special case: comment added at the dispatch site.
- @NonNull: dropped.
Happy to do the capability-matrix update as a follow-up once released.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39793:
URL: https://github.com/apache/beam/pull/39793#issuecomment-5321088097
Assigning reviewers:
R: @Abacn for label java.
Note: If you would like to opt out of this review, comment `assign to next reviewer`.
Available commands:
- `stop reviewer notifications` - opt out of the automated review tooling
- `remind me after tests pass` - tag the comment author after tests pass
- `waiting on author` - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)
The PR bot will only process comments in the main thread (not review comments).
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "f-loris (via GitHub)" <gi...@apache.org>.
f-loris commented on PR #39793:
URL: https://github.com/apache/beam/pull/39793#issuecomment-5334402960
@Abacn Thanks for your feedback. Updated the PR based an you comments and triggered the PostCommit runs you mentioned.
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
Abacn commented on code in PR #39793:
URL: https://github.com/apache/beam/pull/39793#discussion_r3806612330
##########
sdks/java/core/src/test/java/org/apache/beam/sdk/transforms/ParDoTest.java:
##########
@@ -3414,13 +3414,23 @@ public void onWindowExpiration(
}
@Test
- @Category({ValidatesRunner.class, UsesStatefulParDo.class, UsesOrderedListState.class})
+ @Category({
Review Comment:
UsesOrderedListState tests always useUsesOnWindowExpiration. Instead of these change, consider add back ` excludeCategories 'org.apache.beam.sdk.testing.UsesOrderedListState'` in spark_runner.gradle
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #39793:
URL: https://github.com/apache/beam/pull/39793#issuecomment-5320011854
Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment `assign set of reviewers`
--
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] [runners-spark] Support stateful ParDo in the Structured Streaming batch runner [beam]
Posted by "Abacn (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@beam.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org