[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState - #39576
[Spark 4] Streaming runner proof of concept, unbounded sources, windowed GBK and stateful ParDo on transformWithState#39576tkaymak wants to merge 21 commits into
Conversation
Prepares the structured streaming runner for a Spark 4 streaming translator without changing batch behavior. - Add PipelineTranslatorFactory, the seam the Spark 4 module shadows to dispatch streaming pipelines. The shared base rejects streaming with a clear message instead of the previous generic checkArgument. - Open up EvaluationContext (non final, protected ctor, leaves()) and add a no-op stop() so a streaming context can override evaluation. - Add a createEvaluationContext hook to PipelineTranslator and skip the persist and lineage breaking optimizations for streaming datasets. - Plumb the EvaluationContext into SparkStructuredStreamingPipelineResult so cancel() can stop a running streaming query. - Default the state store provider to RocksDB, required by Spark 4 transformWithState and inert for batch. - Add watermarkDelayMillis, maxRecordsPerMicroBatch, maxBatchDurationMillis and streamingStopAfterIdleBatches options. - Narrow the Spark 4 test source override exclude to the legacy DStream package. The previous glob also matched structuredstreaming and would have silently dropped the structured streaming tests.
Wraps any Beam UnboundedSource as a Spark 4 DataSourceV2 micro-batch streaming source. The rows have exactly two columns, payload of type BINARY holding the element encoded with a FullWindowedValueCoder, and eventTimestamp of type TIMESTAMP holding the Beam event timestamp, so no Catalyst encoder has to be generated for Beam types. Offsets are opaque epoch counters. The driver never reads from the source and never inspects its progress, latestOffset simply returns the previous epoch plus one so Spark keeps planning micro-batches. Termination is therefore owned by the lifecycle, not by the offsets. The source is split once on the driver and the sub sources are memoized, so every micro-batch plans the same stable set of partitions. On the executor a static cache keyed by checkpoint location, source id and split id keeps the Beam UnboundedReader alive across micro-batches, mirroring the legacy MicrobatchSource. Checkpoint marks are cached in executor memory only and are never committed durably, which is documented POC scope with weak failure recovery. UnboundedSourceDataset is the translator facing entry point. It applies withWatermark exactly once, since Spark 4 forbids a second declaration further down the plan. The test confirms that the EventTimeWatermark logical node survives one and two typed maps in both the logical and the analyzed plan, and that the running query really tracks the watermark afterwards. It also reads a finite set of elements end to end and checks the decoded payloads and timestamps.
Adds a single generic transformWithState operator that can host any keyed Beam transform, both the stateful ParDo stack and the group also by window stack that implements a windowed GroupByKey. The mode is selected by BeamStatefulProcessorConfig, everything else is shared. Keys, inputs and outputs are raw Beam coder bytes and every Spark encoder involved is BINARY or STRING, so Catalyst never has to derive a schema for a Beam type. TwsTransformFactory owns the row layout and documents it. Beam state is bridged by TwsStateInternals on top of BytesKV, a tiny string to bytes store backed by exactly one Spark MapState. One map is a requirement rather than a preference, ReduceFnRunner invents state tags at runtime so the set of state addresses is not known when init has to declare its state variables. Timers are bridged by TwsTimerInternals, which keeps the full TimerData in a second map and registers only bare wake ups with Spark, reconciling them against listTimers so an expiry can never be registered twice. Two boundary conditions needed care. Spark deletes the expiry it is firing after the callback returns, so a wake up re armed at that same millisecond is nudged one millisecond forward. Spark also expires a wake up as soon as the expiry is at or before the batch watermark, while Beam only fires an event time timer once the watermark is strictly past it, so timers sitting exactly on the watermark are withheld and re armed instead of being handed to Beam, which would decline to fire them and silently drop the on time pane. Processing time timers are rejected with a clear message, transformWithState runs in a single TimeMode and this operator uses event time. The state and timer bridges are unit tested against an in memory store, and BeamStatefulProcessorTest runs both modes through a real streaming query on a live SparkSession.
Shadows PipelineTranslatorFactory so streaming pipelines dispatch to a new PipelineTranslatorStreaming, which extends PipelineTranslatorBatch to reuse Impulse, Window.Assign, Flatten, Reshuffle, the bounded read and stateless ParDo unchanged, while registering placeholders for the unbounded read, GroupByKey and stateful ParDo that WS-D2 will replace. Combine.PerKey is deliberately left unregistered so it auto-expands into GroupByKey plus ParDo. Adds StreamingEvaluationContext, which starts one noop-sink streaming query per leaf dataset with a checkpoint directory derived from the pipeline options, registers an idle-stop listener for test termination, and makes stop() idempotent and safe to call concurrently with evaluate() so cancel() can interrupt a running streaming pipeline from another thread.
…etons Adds StreamingTestUtils, the ListBackedUnboundedSource plus static collector and pipeline option factories the streaming translators will be tested against, and five skeleton test classes for the stateless ParDo, windowed GroupByKey, stateful ParDo, chained stateful, and lifecycle scenarios. Every skeleton compiles and runs, reporting as skipped: each is Ignore'd with a comment naming the exact WS-D2 translator it needs once that lands.
Replaces the WS-D2 placeholders in PipelineTranslatorStreaming with real translators for the unbounded read, GroupByKey and stateful ParDo. ReadUnboundedTranslator wraps the Beam UnboundedSource with UnboundedSourceDataset, which already declares the single withWatermark for the whole query, and decodes the binary payload column into the Dataset<WindowedValue<T>> shape every other translator consumes. The event timestamp column is projected away, the EventTimeWatermark plan node survives below the projection and transformWithState reads the query wide watermark rather than a column. GroupByKeyStreamingTranslator and StatefulParDoStreamingTranslator both encode their input into the byte[] row layout of TwsTransformFactory, hand it to the generic transformWithState operator in GROUP_ALSO_BY_WINDOW and STATEFUL_PARDO mode respectively, and decode the tagged output rows back out. The stateful ParDo splits the output by tag index and skips additional outputs that nothing consumes, so an unused tag does not start a second streaming query. StreamingTranslationHelpers holds the shared guards and row conversions. Pipelines outside the POC scope are rejected at translation time with a message naming the feature, specifically merging windows, custom triggers, accumulating panes, processing time timers, non deterministic key coders, non KV input to a stateful ParDo, side inputs on a stateful ParDo, @OnWindowExpiration and @RequiresTimeSortedInput.
Turns the five Ignore'd streaming test skeletons into real, asserting tests against the WS-D2 translators. Every skeleton declared its pipeline as a local TestPipeline variable, which TestPipeline itself rejects at run() because it is not a Rule field. These tests cannot use PAssert anyway, a streaming pipeline here never gets a final watermark so panes never finalize, and each test needs its own options, so TestPipeline buys nothing. They now build a plain Pipeline.create(options) and assert after waitUntilFinish against the static collector in StreamingTestUtils. The unused TestPipeline factory is dropped from StreamingTestUtils and the reasoning is recorded in its javadoc. StatelessParDoStreamingTest and StreamingPipelineLifecycleTest gain the same relaxed Kryo SparkSessionRule as the rest of the suite, so that the whole package shares one session and a cancelled or idle stopped query cannot take that session down with it. Asserted results, all derived from each test's own input list. Stateless ParDo, exactly 0, 2, 4 up to 18. Fixed 10s window count per key, a=2 and b=1. Sliding 10s every 5s, two a=2 panes, one per shared window. Late data dropped, a=1 and sentinel=1, never a=2. Stateful ParDo, a, b, c once each plus two timer sentinels. Chained dedup into windowed sum, a=8 and b=10. The chained case is the one that matters, a=13 would mean the dedup state was lost and an empty result would mean the watermark got stuck between the two operators, so exactly a=8 and b=10 is the observable proof of cross operator watermark propagation. The late data test is the only one that depends on how the source round robins elements across splits, so it pins maxRecordsPerMicroBatch to 1, spells out the resulting per batch schedule in a comment, and asserts the split count assumption up front rather than leaving it silent. The lifecycle tests wait for a query to actually be active before cancelling, since translation is asynchronous and an early cancel would find a null evaluation context, and they check the query really stopped rather than only that the state was relabelled.
Spark 4 broadcasts a StateSchemaMetadata to the executors for every transformWithState query, using the user Kryo instance. Nothing registered that class, so any Beam streaming pipeline with state or timers died on its first micro-batch as soon as spark.kryo.registrationRequired was on, which is the default for this module's tests and a perfectly reasonable production setting. The same applies to MemoryWriterCommitMessage, the commit message of Spark's memory sink, which is nested inside the already registered DataWritingSparkTaskResult. Both are registered by name, because the shared runner base also compiles against Spark 3 where neither class exists, and both are registered with a JavaSerializer rather than Kryo's default field serializer. Both are Scala case classes holding further Scala and Spark types, immutable.Map, StructType, avro Schema and Row, none of which are registered either, so going through Java serialization covers the whole object graph at once instead of forcing this list to track Spark's internal field layout. Neither object is on a hot path. Every streaming test now runs with the module default of spark.kryo.registrationRequired=true, the per test relaxations are gone. SparkKryoRegistratorStreamingTest names both classes at compile time against the Spark 4 classpath, so a future rename becomes a compile error rather than a silently dropped registration.
Two javadoc paragraphs in the shared base drifted from google-java-format, one from the streaming note added to SparkStructuredStreamingRunner and one from the new Kryo registrations. The shared base sources under runners/spark/src are checked by the :runners:spark project rather than by :runners:spark:3 or :runners:spark:4, whose own spotless targets only see their per version override directories, so it is easy to miss.
ChainedStatefulStreamingTest asserts what the chained pipeline computes. It cannot distinguish a single query holding two transformWithState operators from two queries that happen to add up to the same numbers, and it cannot show whether the watermark moved at all or whether the whole input simply landed in one micro-batch where everything is trivially on time. ChainedStatefulStreamingEvidenceTest asserts the run itself, against Spark's own StreamingQueryProgress. It pins one record per split per micro-batch so the watermark has to climb in steps, records every progress event through a listener, and then asserts that all of them carry one query id, that one micro-batch reports exactly two transformWithStateExec state operators, and that the watermark takes at least three strictly increasing values. The second test adds a tap between the two operators, so the late record is observed leaving the dedup operator and absent from the windowed sum. That is the difference between a record excluded for lateness and a record that never arrived. Both tests print the raw per batch progress they recorded, since that output is the evidence the phase gate report quotes. Also drops four javadoc references to a Kryo relaxation that no longer exists anywhere in the suite.
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Two checks fire under -Werror in CI but not in a local build. Passing -Pjava17Home makes the compile fork to Java 17, and the Beam build adds -XepDisableAllChecks whenever it forks, so ErrorProne is silently off locally. Reproduce with JAVA_HOME set to a Java 17 JDK instead. TypeParameterUnusedInFormals on BeamStreamingSource.decode is suppressed. The decoded types include Coder<WindowedValue<T>>, which no Class token can express, so a checked variant is not available. The unused windowingStrategy local in BeamStatefulProcessor.process is removed, every later use reads config.windowingStrategy() directly.
|
Hi @tkaymak, I found this PR while looking into Spark streaming work. I had been planning to work on Spark portable streaming state and side inputs via #20396 and I’m reading through this POC now so I do not duplicate work or build against the wrong runner My initial interest was state/side-input support, but I see this POC already covers stateful |
|
Hi @Eliaaazzz, happy to share where this is heading. The thinking behind the POC: Spark 4 added transformWithState, and it turned out flexible enough to host the Beam streaming model directly. The POC proves the core claim, two chained Beam stateful operators run inside a single Structured Streaming query with a watermark that propagates correctly between them. So yes, this is the intended direction, and I would not build new work on the DStream path. Once there is agreement on the approach, the plan is to split this branch into five smaller PRs: the dispatch seam and options, the Kryo registrations, the DataSourceV2 unbounded source, the state and timer bridge, and the translators with the end to end tests. Concrete places where help would be useful right now:
Streaming side inputs, your original interest, are deliberately out of scope for the POC, but they are a natural follow up phase once this foundation is agreed. Session windows and full trigger support are planned as later phases too. If you want to start somewhere, I would suggest number 1, it is contained and it teaches you the test setup you would need for anything bigger. |
|
R: @Abacn for the runner and build mechanics, you reviewed the Spark 4 batch runner PRs this builds on. R: @kennknowles for the Beam model side, watermark propagation, trigger semantics and the state and timer bridge on transformWithState. The most interesting finding for you is probably the timer off by one described in the PR, Spark fires a wake up at expiry equal to the watermark while the Beam model requires the watermark to be strictly past the timer, which silently dropped the on time pane until we bounded the fire set. (This is a draft to get agreement on the direction first, the plan is to split it into five smaller PRs afterwards.) |
|
Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control. If you'd like to restart, comment |
…istrator The conditional, by-name registrations only resolve on a Spark 4 classpath. Registering them last keeps the auto assigned Kryo ids of every unconditional registration identical on Spark 3 and Spark 4.
When no checkpointDir is configured the streaming evaluation context creates a temporary directory and never removed it. It is now deleted recursively, best effort, once evaluate() finishes. A user configured checkpointDir is never touched.
…slator canTranslate is a predicate in every batch translator, but the streaming stateful translator deliberately throws instead of returning false, a silent decline would route the DoFn to the batch translator and lose the streaming semantics. The throwing checks now live in rejectUnsupported and the override documents the contract. No error message changed.
awaitTermination previously awaited the leaf queries in list order, so a failure in a later query stayed hidden until every earlier query terminated. The queries are now polled round robin with a short timeout, the first failure stops the siblings and is rethrown. Adds a lifecycle test with a healthy leaf that never idles and a poisoned leaf, which hangs under the old sequential await.
Restarting a streaming pipeline previously lost all source progress, the source id was a fresh UUID per translation, checkpoint marks lived only in executor memory, and the split list was recomputed per run. The source id is now derived deterministically from the read transform's full name. The first run pins its split list under the checkpoint location, Beam sources do not guarantee deterministic splitting and the split index keys all per split state. Each partition reader persists its checkpoint mark per batch epoch, atomically via write and rename through the Hadoop filesystem of the checkpoint location, retaining the two newest epochs. A reader created without an in memory mark restores the newest durable mark at or before the epoch its batch starts at. The driver epoch counter fast forwards past every offset replayed from Spark's log so a restarted driver never emits an offset that goes backwards. Semantics are at least once, a mark is written when a batch finishes reading rather than transactionally with Spark's commit.
Runs two pipelines with identical transform names against one checkpoint directory, clearing the in memory reader cache in between to simulate a fresh JVM. The second run must resume from the durable checkpoint marks, re reading the whole range would mean recovery regressed. ListBackedUnboundedSource gains a positional checkpoint mark so a resumed reader continues after the last emitted element.
|
Update for reviewers, two rounds of changes since the review request. First, the four remarks from my own review pass are addressed: the Kryo registrations for Spark's streaming internals moved to the end of the registrator so auto assigned ids stay identical across Spark 3 and 4, the fallback temporary checkpoint directory is cleaned up after evaluation, the throwing checks in the stateful streaming translator now live in a named rejectUnsupported method with the contract documented, and awaitTermination polls the leaf queries round robin so a failing query surfaces immediately and stops its siblings. Second, the biggest known gap is closed: checkpoint recovery is now durable. The source id is derived deterministically from the read transform's full name, the first run pins its split list under the checkpoint location (Beam sources do not guarantee deterministic splitting), each partition reader persists its checkpoint mark per batch epoch atomically through the Hadoop filesystem, and a restarted run resumes from the newest mark at or before its batch start epoch. Semantics are at least once, a mark is written when a batch finishes reading rather than transactionally with Spark's commit, that caveat is documented on BeamReaderCache. A new end to end test restarts a pipeline against the same checkpoint location with a cleared reader cache and proves it resumes instead of re reading. |
…vancement and streaming Impulse
aa1c70e to
be7e3ce
Compare
…ide expiry PAssert on a non globally windowed collection rewindows into the global window and chains two GroupByKeys. The end of stream sentinel pushes the watermark to TIMESTAMP_MAX_VALUE so end of global window timers fire, but LateDataUtils truncates every garbage collection time to the end of the global window, one day earlier. Judged against the raw watermark the global window itself counts as expired, so the late data filter in front of the second GroupByKey dropped the entire pane flushed by the first and PAssert compared against an empty iterable. Arrival side expiry decisions now see the watermark clamped to the end of the global window through a delegating TimerInternals view installed in the step context. The timer firing path and the ReduceFnRunner keep the real watermark, clamping there would withhold the end of global window timers forever. Restores the fixed windows PAssert test and adds a negative variant proving a wrong expectation still fails on the rewindowed path. Late data drop semantics for finite windows are unchanged, the clamp only takes effect once the watermark has passed the end of the global window.
|
PAssert update. Streaming PAssert now works in non global windows too, commit e991c5d. The earlier commit made PAssert resolve in the global window by emitting an end of stream sentinel that advances the watermark to TIMESTAMP_MAX_VALUE. Non global windows still failed with an empty assertion iterable. The cause was not the fixed windows themselves, PAssert rewindows everything into the global window and chains two GroupByKeys, and the late data filter in front of the second one judged the global window expired because LateDataUtils truncates garbage collection times to the end of the global window while the sentinel watermark sits one day past it. The pane flushed by the first GroupByKey was dropped on arrival at the second. The fix clamps the watermark used for arrival side expiry decisions to the end of the global window, through a delegating TimerInternals view in the step context. Timer firing and the ReduceFnRunner triggers keep the real watermark. Late data semantics for finite windows are unchanged, the clamp only takes effect once the watermark is already past the end of the global window, which only the sentinel can cause. Tests: the fixed windows PAssert case is restored, and both window modes now have negative tests proving a wrong expectation genuinely fails the pipeline. Known remaining gap: chained groupings in finite windows that are flushed only by the final sentinel would still drop their in flight panes. PAssert never builds that shape, it always rewindows to global first. The clean general solution I would choose is per operator output watermark holds, which is on the roadmap for the productionization phase. |
… same-batch folding tests
|
Run Java PreCommit |
|
The runner now covers allowed lateness, late firings via AfterPane.elementCountAtLeast(1), accumulating panes, durable checkpoint restart and PAssert in global and non global windows, each guarded by tests that prove the unsupported shapes still fail loudly. Next step is the split announced in the description. The first slice, the streaming dispatch seam, changes nothing for Spark 3 and is coming as a regular PR shortly. I will link it here. This draft stays open as the end to end evidence for the discussion. |
|
The first slice is out: #39906, the streaming dispatch seam. No behavior change for either Spark version, full local gates green. The remaining slices follow in the order listed there once it lands. |
2 participants
Draft, opened for design discussion rather than for merging as one lump. See "How I would like to split this" at the bottom.
Addresses #36841.
What this is
The Spark 4 runner merged in #38255 is batch only.
SparkStructuredStreamingRunnerrejects streaming outright. This branch is a proof of concept that Spark 4'stransformWithStatecan host the Beam streaming model, and it turns out it can, including several chained stateful operators inside a single Structured Streaming query with a watermark that propagates correctly across them.I would like to agree on the approach before splitting this into reviewable pieces.
Approach
Dispatch seam. The shared base gains
translation/PipelineTranslatorFactory.create(boolean streaming), which throws for streaming and names the Spark 4 module.runners/spark/4/srcshadows that one file through the existing source override copy and returns a streaming translator instead. Every piece oftransformWithState,StreamingQueryand DataSourceV2 streaming code lives only underrunners/spark/4/src, because the shared base still compiles against Spark 3.5 wheretransformWithStatedoes not exist. All shared base edits are behaviour preserving for Spark 3.Source. A DataSourceV2
TableProviderwrapping anyUnboundedSource. Offsets are opaque epoch counters, the driver never reads them. Executor sidePartitionReaders holdUnboundedReaders in a static cache and resume from a locally cachedCheckpointMark, in the style of the legacyMicrobatchSource. The row schema ispayload BINARYpluseventTimestamp TIMESTAMP, so no new Catalyst encoder work was needed.Watermark. Declared exactly once, in the read translator, on the raw rows before the typed decode. Spark forbids re-declaring it. The
EventTimeWatermarkplan node survives the projection and the typed map, which was the main thing I was unsure about going in. There is a test that asserts this at both the plan level and at runtime.Stateful execution. One generic
transformWithStateoperator,StatefulProcessor<byte[], byte[], byte[]>withEncoders.BINARY()throughout, hosting anyDoFnthroughDoFnRunners. Stateful ParDo goes throughsimpleRunnerplusdefaultStatefulDoFnRunner, GBK throughGroupAlsoByWindowViaWindowSetNewDoFnplusSystemReduceFn.bufferingandKeyedWorkItems, which follows the FlinkWindowDoFnOperatorrecipe. State is a port of the legacySparkStateInternalsonto a singleMapState<String, byte[]>keyed by namespace plus tag. That layout is the only one I found that can hostReduceFnRunner's dynamically created system tags. Per@StateIdcolumn families would be faster and are an obvious follow up.Lifecycle.
StreamingEvaluationContextstarts one query per leaf against thenoopsink and blocks until all are terminal.cancel()reaches it through anAtomicReferenceset by the async translate task. Tests terminate through an idle stop listener that gracefully stops a query after N consecutive empty micro batches.Results
:runners:spark:3:test:runners:spark:4:test121 of the Spark 4 tests are structured streaming. The streaming suite was run four separate times with identical results.
Asserted end to end: stateless ParDo, fixed and sliding window GBK, late data dropping, stateful dedup with an event time timer, chained stateful then windowed sum, and pipeline lifecycle through RUNNING, DONE and CANCELLED.
The chained case is the one that matters, and it is asserted against Spark's own
StreamingQueryProgressrather than only against the computed values: one query id across every batch, exactly twotransformWithStateExecoperators per batch, and four strictly increasing watermarks. That rules out two separate queries that happen to sum correctly, and it rules out everything landing in one batch where all data is trivially on time.One finding worth flagging on its own
Spark expires a
transformWithStatewake up atexpiry <= batchWatermark. Beam fires an event time timer only once the watermark is strictly past it, andAfterWatermark.pastEndOfWindowis strict as well. Handing Beam a timer one millisecond early is therefore not merely early, it is destructive: the trigger declines, the runner is entitled to assume the timer will not be redelivered, and the on time pane is silently lost with no error anywhere. The end of window timer of a fixed window sits at exactlywindow.maxTimestamp(), so this hits most windows.Fixed by bounding the fire set at
min(firedExpiry, watermark - 1)and re-arming the withheld timer atfiredExpiry + 1. Anyone else bridging Beam timers ontotransformWithStatewill hit this.Rejected at translation time, with a named message
Merging and session windows, custom triggers including
AfterWatermarkwith early or late firings, accumulating panes, processing time timers,@OnWindowExpiration,@RequiresTimeSortedInput, side inputs on a stateful ParDo, non deterministic key coders.Known limitations
PAsserton the streaming path. No source emits a final positive infinity watermark, so panes never finalise. Streaming tests assert against a static collector afterwaitUntilFinish. I would like input on what reviewers expect here, since this is the biggest departure from how the rest of Beam tests runners.stop()does not drain. An in flight micro batch finishes, anything not yet pulled is left unprocessed.Combine.PerKey.Two smaller things in here
SparkSessionFactorynow registersStateSchemaMetadataandMemoryWriterCommitMessagewith Kryo, by name since the shared base also compiles against Spark 3. Spark 4 broadcasts aStateSchemaMetadatafor everytransformWithStatequery, so without this any stateful streaming pipeline dies on its first micro batch underspark.kryo.registrationRequired=true.runners/spark/4/build.gradlenarrowsexclude "**/translation/streaming/**"toexclude "**/runners/spark/translation/streaming/**". The old glob also matchedstructuredstreaming/translation/streaming, so it would have silently dropped every new test in this branch from the Spark 4 module.Not done yet, deliberately
@Ignore("TODO: Reactivate with streaming."). They cannot be un-ignored in place, since the shared base test tree is compiled by both modules and Spark 3 still correctly throws. Reactivating them means moving them into the Spark 4 override tree. Happy to do that, it just needs a decision on where they should live.How I would like to split this
Assuming the approach is acceptable, five PRs, each independently green:
Slices 4 and 5 would benefit from a reviewer who knows Spark's Structured Streaming internals, not only Beam.