[SPARK-58635][SS] Support for streaming aggregation in Real-Time Mode (RTM) - #57812
Open
jerrypeng wants to merge 14 commits into
Open
[SPARK-58635][SS] Support for streaming aggregation in Real-Time Mode (RTM)#57812jerrypeng wants to merge 14 commits into
jerrypeng wants to merge 14 commits into
Conversation
An AggregationIterator that produces a fresh aggregation buffer on demand, choosing an UnsafeRow buffer when every aggregate buffer attribute is mutable and a GenericInternalRow otherwise. This is the shared base for the aggregation processors used by Real-Time Mode streaming aggregation, added in following commits. It is useful where the buffer has to be initialized more often than the cardinality of the grouping keys, which is the case when a stateful operator merges each input row against state rather than accumulating a hash table per batch. Co-authored-by: Isaac
A pass-through aggregate operator that only handles the aggregation buffer of its input according to the aggregate mode: it either initializes the buffer or completes it and produces the result. It performs no actual aggregation. Real-Time Mode streaming aggregation uses it twice -- as the partial stage that initializes buffers and as the final stage that projects the result columns -- around the stateful operator that does the merging. Added in a following commit. Co-authored-by: Isaac
…on state manager Removing state rows older than the watermark is currently a side-effecting `removeKeysOlderThanWatermark`, which reports nothing back. An operator that has to emit the rows it evicts, or count how many it removed, cannot use it. Add EvictionIterator: it walks a state store, removes the rows whose event time is older than a given timestamp, returns exactly those rows, and tracks how many rows it examined and removed. The event time comes from the state store key via the watermark metadata on the key attributes, so a key with no event time column yields an empty iterator. Real-Time Mode streaming aggregation, added in a following commit, needs both properties: append mode outputs a grouping key once the watermark passes it, and incremental cleanup counts real removals rather than rows scanned. Existing operators are unchanged. Co-authored-by: Isaac
A stateful streaming aggregation operator that streamlines the read-process-output loop: for each input row it reads the previous value for the grouping key, merges the two, emits the merged result immediately, and buffers the state write to be flushed at the end of the batch. Microbatch streaming aggregation (StateStoreSaveExec) instead emits only once the batch ends, which is not viable for Real-Time Mode where a batch runs until its duration elapses. All three output modes are supported: complete emits the whole result table per batch, append emits a grouping key once the watermark passes it, and update emits an intermediate result per input row. Also adds the operator name constant and two internal configs it reads: one selecting the operator under a microbatch trigger for testing, and the incremental cleanup factor (default 0, meaning eviction happens only at the end of a batch). Nothing plans this operator yet; that follows. Co-authored-by: Isaac
Plans a streaming aggregation as: initialize the aggregation buffer, shuffle on the grouping keys, merge each input row against the state store emitting as it goes, then project the result columns. Aggregating only between an input row and the state store means an output is available from processing a single row, so in update mode there is no blocking operation and output latency is as low as possible. Nothing calls this yet; the planner change follows. Co-authored-by: Isaac
…ator in Real-Time Mode A Real-Time Mode batch runs until its duration elapses rather than until its input is exhausted, so a streaming aggregation that only emits once the batch ends would hold every result back for the whole batch. Plan the streamline operator for those queries instead, which merges each input row against state and emits immediately. A query is in Real-Time Mode when it reads a source whose relation carries a real-time mode duration -- the same signal that decides whether to plan a RealTimeStreamScanExec for it. The new internal spark.sql.streaming.useStreamlineAggregator config selects the operator under a microbatch trigger as well, so it can be tested there. Microbatch and batch aggregation are otherwise unchanged. Co-authored-by: Isaac
…me Mode Give the new operators what every stateful operator needs from IncrementalExecution and let the Real-Time Mode allowlist admit them: - assign a state operator id, partition count and output mode to StatefulStreamlineAggregateExec, and propagate the per-batch late-event and eviction watermarks to it; - set the shuffle partition count on ProjectAggregationBufferExec, which it needs to derive its required child distribution as a streaming aggregate; - admit both operators in RealTimeModeAllowlist. With this, a streaming aggregation runs end-to-end in Real-Time Mode. Co-authored-by: Isaac
Two tests over both state format versions: - a keyed aggregation runs in Real-Time Mode and is planned as the streamline aggregate operator. Update mode emits an intermediate result per input row, so a key that appears twice is seen twice, once merged -- micro-batch aggregation would emit only the final value per key per batch. - with incremental cleanup enabled, rows that incremental eviction does not reach are still removed before the batch completes, checked through the numRowsIncrementallyRemoved and numRowsRemoved progress metrics. The second test pins the shuffle partition count, since the eviction counts it asserts are per state store. Co-authored-by: Isaac
"stateful queries not allowed" asserted that a Real-Time Mode streaming aggregation was rejected by the operator allowlist, because it planned into the micro-batch aggregation operators and HashAggregateExec is not allowlisted. Aggregation is now planned as the streamline aggregate operator, which is allowlisted, so there is no rejection left to assert -- including for a global aggregation with no grouping keys, which I checked separately. The generic operator-allowlist test in the same suite still guards the operators that remain unsupported, and the new StreamlineStreamingAggregationRealTimeSuite covers the aggregation itself. Co-authored-by: Isaac
…treamingAggregationSuite Five tests, each over both state format versions, selecting the streamline operator through spark.sql.streaming.useStreamlineAggregator so it is covered under a micro-batch trigger alongside the micro-batch versions of the same queries: - update mode emits one output per input row. Four rows for the same key produce four outputs, where micro-batch aggregation emits only the final value for the batch. This is the operator's distinguishing behaviour. - complete mode outputs the whole result table per batch. - sum, min, max and avg over a grouping key. - multiple grouping keys. - state survives a stop and restart, so counts continue rather than restarting. They sit here rather than in a new file because the suite already provides testWithAllStateVersions and the per-test config plumbing, the micro-batch expectations for the same queries are right alongside for comparison, and the existing RocksDB subclass re-runs them against that state store for free. Co-authored-by: Isaac
…amline aggregate groupingProjection hands back the same UnsafeRow on every call, overwriting its bytes, and the deferred-write cache stored that row directly as its key. One shared row therefore became the key of every entry: an entry could only be found again through the hash captured when it was inserted, and the equality check that should discriminate between keys degenerated to a reference match against that one row. Aggregate results happen to come out correct today because the stored hash does the discriminating, but the cache is not behaving as the value semantics it is written against would imply, and any change to how the key is compared would silently produce wrong aggregates. The existing comment at the lookup already called for the copy; this makes the code do it. Co-authored-by: Isaac
… not expressions The final stage's child is the post-shuffle output of the stateful aggregate, whose grouping columns are already resolved attributes, so it should group by those. It was passing the original grouping expressions instead, unlike planStreamingAggregation's final stage and unlike the stateful aggregate in the same function, both of which group by the attributes. Adds a test with a computed grouping key, which is the shape where an expression that has to be re-evaluated against the child's output would diverge from a plain attribute reference. Co-authored-by: Isaac
`iter.map { numOutputRows += 1; aggProcessor.process _ }` passes a block whose
value is the eta-expanded function. The block runs once per partition, so the
counter was incremented once per partition rather than once per row, and the
operator reported a numOutputRows of 1 no matter how many rows it produced.
Take the row as a parameter so the increment happens per element.
Co-authored-by: Isaac
…d silent state-write drop Fixes for issues found reviewing the streamline aggregation operator, each verified against real upstream code and, where behavioral, by a probe. Changes are kept close to the Databricks runtime form so the stack merges back cleanly. 1. IncrementalExecution: allow the stateStoreSave <-> StatefulStreamlineAggregate operator transition in the state-metadata check. The two operators share StreamingAggregationStateManager and the same state format, so their on-disk state is mutually readable, but checkOperatorValidWithMetadata compared operator names and aborted the restart. Verified: a same-stream restart that flips the operator threw STREAMING_STATEFUL_OPERATOR_NOT_MATCH_IN_STATE_METADATA before this change; a streamline-on-streamline restart was unaffected. Mirrors the runtime's isOperatorMetadataConvertible, character for character. The operator name string "StatefulStreamlineAggregate" is kept as-is: it matches the runtime and is persisted into checkpoint metadata, so renaming would break checkpoint compatibility on merge-back. 2. StatefulStreamlineAggregateExec: fail the task when a state-store write fails during flush. The removal listener performs the essential stateManager.put, but Guava logs and swallows any exception a removal listener throws (confirmed by a direct probe: invalidateAll() returned normally despite the listener throwing), so a failed put was dropped and the batch still committed at the next version, silently losing that key's update while the sink already emitted the row. StateStoreSaveExec puts inline, so it fails the task on such a write; this restores that guarantee by capturing the first failure and rethrowing it from flushDirtyWrites. This is the one change ahead of the runtime, which has the same latent swallow. 3. EvictionIterator: fix the scaladoc, which claimed "every next() corresponds to one removal" while removal actually happens in hasNext (as the code and its own inline comment describe). 4. SQLConf: reject a negative incrementalCleanupFactor with checkValue; a negative silently disabled cleanup. (This config is OSS-only -- the runtime's equivalent lives in DatabricksSQLConf under a different key -- so there is no divergence.) 5. Drop the allowlist-deletion comment's unsupported claim of coverage for a grouping-key-less aggregation; no such test exists. Tests: new StreamingAggregationSuite case covering the operator transition in both directions (revert-checked: fails without fix 1). StreamingAggregationSuite (56), StreamlineStreamingAggregationRealTimeSuite + StreamRealTimeModeAllowlistSuite (7) pass. Co-authored-by: Isaac
jerrypeng
changed the title
jerrypeng
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
This PR adds support for streaming aggregation in Real-Time Mode (RTM), building on the pipelined
shuffle and RTM stateful work from SPARK-58508.
RTM cannot use the ordinary streaming aggregation plan. The standard plan materializes a partial
aggregate before the shuffle and holds results until the batch ends, which is incompatible with RTM's
requirement to emit rows continuously within a long-running batch. This PR introduces a streamline
aggregation plan that reads, merges, and emits per row instead.
The plan is three operators:
ProjectAggregationBufferExec(Partial) - projects input rows into aggregation buffers.StatefulStreamlineAggregateExec(PartialMerge) - the stateful core. For each row it reads thecurrent buffer from the state store, merges the new value, and emits the updated buffer
immediately. Buffers written during the batch are held in a bounded cache (Guava
maximumSize)and flushed to the state store at batch end.
ProjectAggregationBufferExec(Final) - projects the merged buffers to the output schema.Supporting pieces:
GenericBufferAggregationIterator- shared buffer-merge iteration used by the operators above.EvictionIterator- incremental state-store cleanup, so eviction cost is spread across the batchrather than paid in one pass at batch end.
AggUtils.planStreamlineStreamingAggregation- the planner entry point, selected byStatefulAggregationStrategywhen the plan is an RTM plan.IncrementalExecutionwires the operators into the RTM path and admits thestateStoreSave<->StatefulStreamlineAggregateoperator transition in the state-metadata check(the two share the same state manager and format, so their on-disk state is mutually readable);
RealTimeModeAllowlistadmits the new operators.spark.sql.streaming.useStreamlineAggregator(defaultfalse) andspark.sql.streaming.statefulOperator.incrementalCleanupFactor(default0, meaning cleanup atbatch end). RTM uses the streamline operator regardless of the first config; it exists so the
operator can also be exercised under a microbatch trigger.
The change is inert outside RTM: the streamline plan is only selected for a plan containing a
RealTimeStreamScanExecleaf (or when the config above is set), so the ordinary microbatchaggregation path is untouched.
Why are the changes needed?
Before this PR, a streaming aggregation was rejected in RTM: it planned into
HashAggregateExec,which is not in
RealTimeModeAllowlist, so the query failed at start withSTREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST. Aggregation is one of the most commonstreaming workloads, so this was a significant gap in RTM's coverage. SPARK-58508 enabled stateful
queries in RTM (deduplication); this extends that to aggregation.
Does this PR introduce any user-facing change?
Yes, additive. A streaming aggregation query can now run under
Trigger.RealTime(...)where itpreviously failed at query start. Behavior outside RTM is unchanged: the two new configs are
internal()and default off, and the streamline plan is only selected for an RTM plan.How was this patch tested?
New and updated tests:
StreamlineStreamingAggregationRealTimeSuite- end-to-end RTM aggregation coverage, includingincremental cleanup.
StreamingAggregationSuite- the shared suite now also exercises the streamline operator (via theuseStreamlineAggregatorconfig), so the new path is covered by the same assertions as theexisting micro-batch path. Adds a test that a query keeps its state when switched between the
micro-batch and streamline operators in either direction, guarding the operator-transition
allowance above.
StreamRealTimeModeAllowlistSuite- updated for the newly admitted operators; the obsoleteaggregation-rejection test is removed.
Results on this branch:
StreamingAggregationSuite55/55,StreamlineStreamingAggregationRealTimeSuite+StreamRealTimeModeAllowlistSuite8/8, and thestreamline cases in
RocksDBStateStoreStreamingAggregationSuite26/26.Several defects were caught by these tests during development and fixed, each in its own commit:
counting every output row in
ProjectAggregationBufferExec; grouping the final aggregate byattributes rather than expressions; copying the grouping key before caching it; allowing the
micro-batch <-> streamline operator transition on restart; and failing the task (rather than silently
dropping the write and still committing) when a state-store write fails during the batch-end flush.
Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code