[damccorm]

Note: this PR was auto-generated by tooling. Before moving it out of draft, the author has done a full review of the PR, specifically evaluating it on the following criteria:

• The PR is an appropriate fix for the issue (e.g. we're not changing a test to avoid a bug)
• The PR is a simple fix to the issue (e.g. we're not adding infra when we could just modify a test)
• The PR is not covering for other problems (e.g. if we're bumping timeouts, it is not because we're covering for a performance problem)

Full details of the methodology of generation and review can be found here - https://docs.google.com/document/d/14lj7ECVov4T4EdwmIzNlWrRGE3lLDMOwX5xUMGzaLoc/edit?usp=sharing

If you have concerns with the PR or with the underlying methodology used to generate this, please comment below - they will be addressed by the author not the tooling :)

Summary

Fixes a driver concurrency race in Spark 4 Structured Streaming when multiple leaf StreamingQuery instances share an upstream TransformWithState logical plan node (such as GroupByKey or stateful ParDo consumed by multiple PAssert queries in GroupByKeyTest$WindowTests.testGroupByKeyMultipleWindows).

Production Impact

Affects production Spark 4 Structured Streaming pipelines where a stateful transform (GroupByKey or stateful ParDo) is consumed by multiple downstream outputs/sinks. Because Spark starts a separate StreamingQuery per leaf output that shares the upstream logical plan, sharing a single mutable BeamStatefulProcessor instance across sibling queries causes concurrent driver micro-batch planning threads to overwrite and clear statefulProcessorHandle = null mid-initialization, crashing the pipeline with SparkRuntimeException: [STATE_STORE_HANDLE_NOT_INITIALIZED].

Root Cause

Failed master run: https://github.com/apache/beam/actions/runs/42658100679/job/129301949163 (PostCommit Java ValidatesRunner Spark4).

In Spark 4 (org.apache.spark.sql.execution.streaming.TransformWithStateExec), getDriverProcessorHandle() and closeProcessorHandle() are invoked on the driver by each StreamingQuery's StreamExecution thread during micro-batch planning (IncrementalExecution$StateSchemaAndOperatorMetadataRule -> getColFamilySchemas / getStateVariableInfos):

1. getDriverProcessorHandle() calls statefulProcessor.setHandle(driverProcessorHandle) followed by statefulProcessor.init(outputMode, timeMode).
2. closeProcessorHandle() calls statefulProcessor.close() followed by statefulProcessor.setHandle(null).

Because StatefulProcessor.setHandle and StatefulProcessor.getHandle are public final methods that mutate a single non-thread-safe private StatefulProcessorHandle statefulProcessorHandle field on the StatefulProcessor instance, sharing the same TransformWithState logical plan node across sibling leaf StreamingQuerys causes their StreamExecution threads to concurrently set and clear statefulProcessorHandle = null mid-init(), throwing:

org.apache.spark.SparkRuntimeException: [STATE_STORE_HANDLE_NOT_INITIALIZED] State store handle is not initialized.
  at org.apache.spark.sql.streaming.StatefulProcessor.getHandle(StatefulProcessor.scala:146)
  at org.apache.beam.runners.spark.structuredstreaming.translation.streaming.BeamStatefulProcessor.init(BeamStatefulProcessor.java:110)
  at org.apache.spark.sql.execution.streaming.TransformWithStateExec.getDriverProcessorHandle(TransformWithStateExec.scala:155)

Fix

• In StreamingEvaluationContext.startQuery, clone StatefulProcessor instances inside TransformWithState logical plan nodes for additional leaf queries (leafIndex > 0) via SerializableUtils.clone(tws.statefulProcessor()) so each leaf StreamingQuery has its own isolated StatefulProcessor instance on the driver.
• Add WindowedGroupByKeyStreamingTest.testGroupByKeySharedAcrossMultipleLeaves to verify multiple PAssert leaf queries consuming a shared GroupByKey.

Verification

• Triggered beam_PostCommit_Java_ValidatesRunner_Spark4.yml workflow run on branch: 42696643750

[damccorm]

/gemini review

[github-actions]

📋 Review Summary

This PR resolves a race condition on the driver in Spark 4 Structured Streaming when multiple sibling leaf StreamingQuery instances share an upstream TransformWithState logical plan node. By deep-cloning the StatefulProcessor instances in TransformWithState nodes for subsequent leaf queries, each query's StreamExecution thread operates on an isolated StatefulProcessor handle during planning. The fix is well-targeted, clean, and accompanied by a multi-leaf streaming unit test.

🔍 General Feedback

• The recursive traversal via LogicalPlan.mapChildren and TransformWithState.copy with SerializableUtils.clone cleanly isolates the non-thread-safe StatefulProcessor state per query without affecting single-leaf pipelines (leafIndex == 0).
• WindowedGroupByKeyStreamingTest.testGroupByKeySharedAcrossMultipleLeaves effectively exercises the concurrent multi-leaf query execution path.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 57.53%. Comparing base (a537a4c) to head (a4a896c).
⚠️ Report is 36 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff            @@
##             master   #41137   +/-   ##
=========================================
  Coverage     57.53%   57.53%           
  Complexity     2502     2502           
=========================================
  Files          1207     1207           
  Lines        171557   171557           
  Branches       2121     2121           
=========================================
  Hits          98710    98710           
  Misses        70187    70187           
  Partials       2660     2660           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[damccorm]

R: @tkaymak

[github-actions]

Assigning reviewers:

R: @kennknowles added as fallback since no labels match configuration
Backup reviewer: chamikaramj (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

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).

[github-actions]

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment assign set of reviewers

[tkaymak]

Thanks for chasing this down, the diagnosis is right and one StatefulProcessor per query thread is the invariant we need. Two things before I can approve:

1. SerializableUtils.clone also copies the MetricsAccumulator that ParDoStatefulProcessor holds. Spark only merges the accumulator the driver registered (see the note in MetricsAccumulator.getInstance), so a stateful ParDo in every leaf after the first would stop reporting metrics.
2. The copies for leaves 1..n are taken after query 0 is started on the original instance, and StatefulProcessor keeps its handle in a non transient field that query 0 sets and clears while planning its first batch. That narrows the race, it doesn't remove it.

Both go away if BeamStatefulProcessor gets a copy method that the two subclasses implement with their constructor and the same final fields, instead of a Java serialization round trip. It keeps the registered accumulator and reads no Spark handle, so the order against query 0 no longer matters.

Smaller points: could you align ofRows with the existing Dataset$.MODULE$ call sites (BoundedDatasetFactory, UnboundedSourceDataset), add a short comment on why leaf 0 keeps the original, and mention #41198 as the long term fix (one query per pipeline would remove this code). Does the new test fail without the change? A stateful ParDo with a counter and two leaves would also cover point 1.
And could you please bump the Spark 4 ValidatesRunner trigger file so the PostCommit runs here?

[damccorm]

Thanks for chasing this down, the diagnosis is right and one StatefulProcessor per query thread is the invariant we need. Two things before I can approve:

1. SerializableUtils.clone also copies the MetricsAccumulator that ParDoStatefulProcessor holds. Spark only merges the accumulator the driver registered (see the note in MetricsAccumulator.getInstance), so a stateful ParDo in every leaf after the first would stop reporting metrics.
2. The copies for leaves 1..n are taken after query 0 is started on the original instance, and StatefulProcessor keeps its handle in a non transient field that query 0 sets and clears while planning its first batch. That narrows the race, it doesn't remove it.

Both go away if BeamStatefulProcessor gets a copy method that the two subclasses implement with their constructor and the same final fields, instead of a Java serialization round trip. It keeps the registered accumulator and reads no Spark handle, so the order against query 0 no longer matters.

Smaller points: could you align ofRows with the existing Dataset$.MODULE$ call sites (BoundedDatasetFactory, UnboundedSourceDataset), add a short comment on why leaf 0 keeps the original, and mention #41198 as the long term fix (one query per pipeline would remove this code). Does the new test fail without the change? A stateful ParDo with a counter and two leaves would also cover point 1. And could you please bump the Spark 4 ValidatesRunner trigger file so the PostCommit runs here?

Thanks - I think I addressed the feedback. I am a little out of my wheelhouse here (this is part of a broader effort I'm working on to automate away some of our test flakes/toil around that), so there may be pieces I'm missing. If so, feel free to either help me through it or try to fork/fix yourself as you have time available - whatever is easiest for you depending on your availability. Regardless, it would be good to clean up the failures introduced here.

[damccorm]

Does the new test fail without the change?

Yes, both do

[tkaymak]

Thanks, this addresses everything. The copies keep the registered accumulator and the new stateful test proves both leaves report metrics. LGTM.

One small thing for later: the expected 10 in testStatefulParDoSharedAcrossMultipleLeaves is 5 per leaf because each leaf query runs the stateful ParDo, a short comment pointing at #41198 would help whoever fixes that (I am happy to do that too), since it then drops to 5.

[damccorm]

Thanks, this addresses everything. The copies keep the registered accumulator and the new stateful test proves both leaves report metrics. LGTM.

One small thing for later: the expected 10 in testStatefulParDoSharedAcrossMultipleLeaves is 5 per leaf because each leaf query runs the stateful ParDo, a short comment pointing at #41198 would help whoever fixes that (I am happy to do that too), since it then drops to 5.

Thanks, added