This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 8debe32663ad [SPARK-58508][CORE][SS] Activate pipelined shuffles for
Real-Time Mode streaming queries to support Stateful Queries
8debe32663ad is described below
commit 8debe32663ad988684a2ecfb11351a0c9c6641e0
Author: Boyang Jerry Peng <[email protected]>
AuthorDate: Wed Aug 5 14:26:10 2026 +0800
[SPARK-58508][CORE][SS] Activate pipelined shuffles for Real-Time Mode
streaming queries to support Stateful Queries
### What changes were proposed in this pull request?
This PR activates **pipelined (streaming) shuffles for Real-Time Mode (RTM)
streaming queries**. It is the payload PR of the pipelined-shuffle stack:
earlier PRs built the transient-shuffle machinery (group-atomic scheduling,
MapOutputTracker decoupling) as dormant infrastructure with no production
caller; this PR is the first to actually **construct and use** a
`PipelinedShuffleDependency`.
The `core/` dependency it builds on, #57646 ("Fully decouple pipelined
shuffles from the MapOutputTracker"), has merged, so this PR now applies
directly to `master` and its diff is SQL/streaming only.
Concretely:
- **Build the dependency (`ShuffleExchangeExec`)**: a new `pipelined:
Boolean` flag; when set, it constructs a `PipelinedShuffleDependency` instead
of a regular `ShuffleDependency`. The dependency type is the entire opt-in --
routing to the streaming shuffle manager and pipelined-group co-scheduling in
the DAGScheduler both follow from it. Plumbed through `EnsureRequirements`,
`AQEUtils`, and `InsertSortForLimitAndOffset`.
- **Mark shuffles on the streaming path (`IncrementalExecution`)**: a
`MarkPipelinedShuffleForRealTimeMode` rule marks the `ShuffleExchangeExec`s
whose subtree reaches a `RealTimeStreamScanExec`. Inert for a non-RTM batch. A
chain of several pipelined shuffles (e.g. two repartitions, or a repartition
feeding a keyed stateful operator) is handled -- the whole all-pipelined job is
co-scheduled as one pipelined group. Shuffles in a *static* subtree (for
instance the static side of a broa [...]
- **Extend the RTM allowlist (`RealTimeModeAllowlist`)**: admit
`ShuffleExchangeExec` plus the stateful operators it enables --
`StateStoreRestoreExec`, `StateStoreSaveExec`, `StreamingDeduplicateExec` -- as
supported members of a pipelined group rather than materialization barriers.
Admitting shuffles is qualified by partitioning: a **range-partitioned**
shuffle is still rejected, because building its `RangePartitioner` runs a
separate job that samples the input to compute range boun [...]
- **Make `repartition(n)` usable in RTM (`MicroBatchExecution`,
`ShuffleExchangeExec`)**: `repartition(n)` produces a `RoundRobinPartitioning`,
and by default (SPARK-23207) Spark inserts a deterministic local sort before a
round-robin shuffle so a retried task cannot reassign rows and lose data. That
sort is **fully blocking** -- it drains its entire input before emitting a row
-- so on an unbounded RTM stream the producer never emits and the query returns
no rows at all, with no erro [...]
The net effect: a stateful streaming query runs in Real-Time Mode with its
input repartitioned through a transient, incrementally-read shuffle, instead of
the consumer waiting for the producer to fully materialize.
### Why are the changes needed?
Real-Time Mode aims for low end-to-end latency by co-scheduling a query's
producer and consumer stages so records stream through rather than
materializing at a shuffle boundary. Until now RTM rejected shuffles and the
stateful operators behind them (they were not in the allowlist), so a keyed
stateful query could not run in RTM at all. This PR turns on the
pipelined-shuffle path -- built and made safe by the earlier PRs in the stack
-- for exactly those queries.
### Does this PR introduce _any_ user-facing change?
Yes, for Real-Time Mode streaming queries: a query that repartitions by key
(e.g. streaming deduplication) is now admitted and runs over a pipelined
shuffle, where it was previously rejected by the RTM operator allowlist. RTM
also now forces `spark.sql.execution.sortBeforeRepartition=false` for its own
(isolated) session, overriding an explicit `true` and logging a warning --
honouring it would hang the query forever, and it guards against a task retry
that does not happen in RTM. Bat [...]
### How was this patch tested?
New end-to-end RTM tests in `StreamRealTimeModeSuite` /
`StreamRealTimeModeWithManualClockSuite` (the suites shared with the Databricks
runtime, rather than a new standalone file):
- stateful dedup runs in RTM over a pipelined shuffle and genuinely
co-schedules its stages (asserted from a `SparkListener`, counting only the
stages of the query under test);
- multi-key dedup, and an explicit repartition-by-key;
- a chain of two shuffles (round-robin repartition feeding a keyed dedup)
forms one co-scheduled pipelined group;
- a chain of two round-robin repartitions -- the shape that the
order-sensitivity exemption exists for;
- an explicit `sortBeforeRepartition=true` does not stall the query;
- a shuffle on the static side of a broadcast stream-static join is not
marked pipelined;
- multiple broadcast joins against the same static table, asserting the
plan really does reuse the broadcast exchange (SC-209926);
- recovery from a task failure via checkpoint restart, and from a
commit-log write failure.
`StreamRealTimeModeAllowlistSuite` also gains a rejection test: a
range-partitioned shuffle (`repartitionByRange`) fails fast with the allowlist
error rather than stalling.
Updated for the new allowlist and the `pipelined` field:
`StreamRealTimeModeAllowlistSuite`, `EnsureRequirementsSuite`, `PlannerSuite`,
`DataFrameWindowFunctionsSuite`, `DatasetSuite`, `StreamingJoinSuite`.
### Was this patch authored or co-authored using generative AI tooling?
Co-authored with Claude Code (Opus 4.8)
Closes #57692 from jerrypeng/stack/pipelined-shuffle-pr8-rtm.
Authored-by: Boyang Jerry Peng <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
.../execution/InsertSortForLimitAndOffset.scala | 2 +-
.../spark/sql/execution/adaptive/AQEUtils.scala | 2 +-
.../execution/exchange/EnsureRequirements.scala | 6 +-
.../execution/exchange/ShuffleExchangeExec.scala | 65 +++-
.../streaming/runtime/IncrementalExecution.scala | 93 ++++-
.../streaming/runtime/MicroBatchExecution.scala | 25 ++
.../streaming/runtime/RealTimeModeAllowlist.scala | 39 +-
.../spark/sql/DataFrameWindowFunctionsSuite.scala | 2 +-
.../scala/org/apache/spark/sql/DatasetSuite.scala | 2 +-
.../apache/spark/sql/execution/ExchangeSuite.scala | 28 ++
.../apache/spark/sql/execution/PlannerSuite.scala | 5 +-
.../exchange/EnsureRequirementsSuite.scala | 110 +++---
.../StreamRealTimeModeAllowlistSuite.scala | 52 ++-
.../sql/streaming/StreamRealTimeModeSuite.scala | 418 ++++++++++++++++++++-
.../spark/sql/streaming/StreamingJoinSuite.scala | 4 +-
15 files changed, 737 insertions(+), 116 deletions(-)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
index aa29128cda7e..88f0de2555f7 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/InsertSortForLimitAndOffset.scala
@@ -44,7 +44,7 @@ object InsertSortForLimitAndOffset extends Rule[SparkPlan] {
_,
// Should not match AQE shuffle stage because we only target
un-submitted stages which
// we can still rewrite the query plan.
- s @ ShuffleExchangeExec(SinglePartition, child, _, _),
+ s @ ShuffleExchangeExec(SinglePartition, child, _, _, _),
_) if child.logicalLink.isDefined =>
extractOrderingAndPropagateOrderingColumns(child) match {
case Some((ordering, newChild)) =>
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala
index 578e0acd8052..cc3ca9073909 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/adaptive/AQEUtils.scala
@@ -30,7 +30,7 @@ object AQEUtils {
// Project/Filter/LocalSort/CollectMetrics.
// Note: we only care about `HashPartitioning` as `EnsureRequirements` can
only optimize out
// user-specified repartition with `HashPartitioning`.
- case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _)
+ case ShuffleExchangeExec(h: HashPartitioning, _, shuffleOrigin, _, _)
if shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin ==
REPARTITION_BY_NUM =>
val numPartitions = if (shuffleOrigin == REPARTITION_BY_NUM) {
Some(h.numPartitions)
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
index c632b3d841e6..8ca344e51194 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
@@ -272,8 +272,8 @@ case class EnsureRequirements(
}
child match {
- case ShuffleExchangeExec(_, c, so, ps) =>
- ShuffleExchangeExec(newPartitioning, c, so, ps)
+ case s: ShuffleExchangeExec =>
+ s.copy(outputPartitioning = newPartitioning)
case gpe: GroupPartitionsExec =>
ShuffleExchangeExec(newPartitioning, gpe.child)
case _ => ShuffleExchangeExec(newPartitioning, child)
}
@@ -896,7 +896,7 @@ case class EnsureRequirements(
def apply(plan: SparkPlan): SparkPlan = {
val newPlan = plan.transformUp {
- case operator @ ShuffleExchangeExec(upper: HashPartitioning, child,
shuffleOrigin, _)
+ case operator @ ShuffleExchangeExec(upper: HashPartitioning, child,
shuffleOrigin, _, _)
if optimizeOutRepartition &&
(shuffleOrigin == REPARTITION_BY_COL || shuffleOrigin ==
REPARTITION_BY_NUM) =>
def hasSemanticEqualPartitioning(partitioning: Partitioning): Boolean
= {
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
index dd829e697df6..0059473b3c65 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/ShuffleExchangeExec.scala
@@ -191,7 +191,8 @@ case class ShuffleExchangeExec(
override val outputPartitioning: Partitioning,
child: SparkPlan,
shuffleOrigin: ShuffleOrigin = ENSURE_REQUIREMENTS,
- advisoryPartitionSize: Option[Long] = None)
+ advisoryPartitionSize: Option[Long] = None,
+ pipelined: Boolean = false)
extends ShuffleExchangeLike {
private lazy val writeMetrics =
@@ -205,6 +206,17 @@ case class ShuffleExchangeExec(
override def nodeName: String = "Exchange"
+ // `pipelined` is only meaningful for a Real-Time Mode plan, and the default
arg string is
+ // positional, so printing it unconditionally would add a bare `false` to
every shuffle in every
+ // plan. Show it only when set, and name it when shown.
+ override def stringArgs: Iterator[Any] = {
+ // `pipelined` is the last field, so drop it positionally rather than by
value; argString drops
+ // the child on its own. Exchange's `[plan_id=...]` suffix is re-appended
here.
+ val argsWithoutPipelined = productIterator.toSeq.dropRight(1).iterator
+ val pipelinedArg = if (pipelined) Iterator("isPipelined=true") else
Iterator.empty
+ argsWithoutPipelined ++ pipelinedArg ++ Iterator(s"[plan_id=$id]")
+ }
+
private lazy val serializer: Serializer =
new UnsafeRowSerializer(child.output.size, longMetric("dataSize"))
@@ -252,7 +264,8 @@ case class ShuffleExchangeExec(
child.output,
outputPartitioning,
serializer,
- writeMetrics)
+ writeMetrics,
+ pipelined)
metrics("numPartitions").set(dep.partitioner.numPartitions)
val executionId =
sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
SQLMetrics.postDriverMetricUpdates(
@@ -346,7 +359,8 @@ object ShuffleExchangeExec {
outputAttributes: Seq[Attribute],
newPartitioning: Partitioning,
serializer: Serializer,
- writeMetrics: Map[String, SQLMetric])
+ writeMetrics: Map[String, SQLMetric],
+ pipelined: Boolean = false)
: ShuffleDependency[Int, InternalRow, InternalRow] = {
val part: Partitioner = newPartitioning match {
case RoundRobinPartitioning(numPartitions) => new
HashPartitioner(numPartitions)
@@ -509,8 +523,16 @@ object ShuffleExchangeExec {
// round-robin function is order sensitive if we don't sort the input.
// Stateful partition assignment is order-sensitive when it depends on
row visitation order.
- val isOrderSensitive =
- (isRoundRobin || isNullAwareHashPartitioning) &&
!SQLConf.get.sortBeforeRepartition
+ //
+ // A pipelined shuffle is exempt. Marking the map RDD order-sensitive
only serves to make it
+ // INDETERMINATE when its own input is UNORDERED (see
+ // MapPartitionsRDD.getOutputDeterministicLevel), which tells the
scheduler a retry cannot be
+ // trusted and the stage must be rolled back and recomputed. A pipelined
stage is never
+ // retried (a pipelined task set gets a single attempt) and never
recomputed, and the
+ // DAGScheduler rejects an indeterminate pipelined producer outright --
so keeping the flag
+ // would reject a chain of round-robin repartitions rather than protect
anything.
+ val isOrderSensitive = (isRoundRobin || isNullAwareHashPartitioning) &&
+ !SQLConf.get.sortBeforeRepartition && !pipelined
if (needToCopyObjectsBeforeShuffle(part)) {
newRdd.mapPartitionsWithIndexInternal((_, iter) => {
val getPartitionKey = getPartitionKeyExtractor()
@@ -537,15 +559,30 @@ object ShuffleExchangeExec {
}
}
val dependency =
- new ShuffleDependency[Int, InternalRow, InternalRow](
- rddWithPartitionIds,
- new PartitionIdPassthrough(part.numPartitions),
- serializer,
- shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics),
- rowBasedChecksums =
UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize),
- _checksumMismatchFullRetryEnabled =
SQLConf.get.shuffleChecksumMismatchFullRetryEnabled,
- checksumMismatchQueryLevelRollbackEnabled =
- SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled)
+ if (pipelined) {
+ // A pipelined shuffle is transient and incrementally readable: the
DAGScheduler
+ // co-schedules its producer and consumer stages instead of
materializing the shuffle
+ // first. The PipelinedShuffleDependency type is the entire opt-in --
routing to the
+ // streaming shuffle manager and pipelined-group co-scheduling both
follow from it. The
+ // checksum-mismatch retry knobs are intentionally not carried over: a
transient shuffle is
+ // never recomputed, so PipelinedShuffleDependency does not expose
them (they stay off).
+ new PipelinedShuffleDependency[Int, InternalRow, InternalRow](
+ rddWithPartitionIds,
+ new PartitionIdPassthrough(part.numPartitions),
+ serializer,
+ shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics),
+ rowBasedChecksums =
UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize))
+ } else {
+ new ShuffleDependency[Int, InternalRow, InternalRow](
+ rddWithPartitionIds,
+ new PartitionIdPassthrough(part.numPartitions),
+ serializer,
+ shuffleWriterProcessor = createShuffleWriteProcessor(writeMetrics),
+ rowBasedChecksums =
UnsafeRowChecksum.createUnsafeRowChecksums(checksumSize),
+ _checksumMismatchFullRetryEnabled =
SQLConf.get.shuffleChecksumMismatchFullRetryEnabled,
+ checksumMismatchQueryLevelRollbackEnabled =
+ SQLConf.get.shuffleChecksumMismatchQueryLevelRollbackEnabled)
+ }
dependency
}
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
index 0d2e4a6941a0..f7b64af45a31 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/IncrementalExecution.scala
@@ -36,8 +36,9 @@ import org.apache.spark.sql.classic.SparkSession
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.{CommandExecutionMode, LocalLimitExec,
QueryExecution, SerializeFromObjectExec, SparkPlan, SparkPlanner, SparkStrategy
=> Strategy, UnaryExecNode}
import org.apache.spark.sql.execution.aggregate.{HashAggregateExec,
MergingSessionsExec, ObjectHashAggregateExec, SortAggregateExec,
UpdatingSessionsExec}
+import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec
import
org.apache.spark.sql.execution.datasources.v2.state.metadata.StateMetadataPartitionReader
-import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike
+import org.apache.spark.sql.execution.exchange.{ShuffleExchangeExec,
ShuffleExchangeLike}
import
org.apache.spark.sql.execution.python.streaming.{FlatMapGroupsInPandasWithStateExec,
TransformWithStateInPySparkExec}
import org.apache.spark.sql.execution.streaming.{StreamingErrors,
StreamingQueryPlanTraverseHelper}
import
org.apache.spark.sql.execution.streaming.checkpointing.{CheckpointFileManager,
OffsetSeqMetadata, OffsetSeqMetadataBase}
@@ -663,7 +664,95 @@ class IncrementalExecution(
}
}
- override def preparations: Seq[Rule[SparkPlan]] = state +: super.preparations
+ /**
+ * For a Real-Time Mode batch, mark the shuffle exchanges as pipelined so
the DAGScheduler
+ * co-schedules a stateful query's producer (source scan) and consumer
(stateful operator) stages
+ * as one pipelined group -- records stream through a transient shuffle
instead of the consumer
+ * waiting for the producer to fully materialize. The exchange carries the
decision as a field
+ * (see ShuffleExchangeExec.pipelined); the PipelinedShuffleDependency it
then builds is the whole
+ * opt-in -- routing to the streaming shuffle manager and pipelined-group
scheduling both follow
+ * from that dependency type.
+ *
+ * Real-Time Mode is detected structurally by a RealTimeStreamScanExec leaf
(there is no
+ * RTM-specific plan flag). Inert for a non-RTM batch, so the ordinary
microbatch path is
+ * unchanged.
+ *
+ * Marks every eligible shuffle exchange on the streaming path, so a plan
with several pipelined
+ * shuffles in a chain (e.g. two repartitions, or a repartition feeding a
keyed stateful operator)
+ * is handled: each becomes a PipelinedShuffleDependency and the whole
all-pipelined job is
+ * co-scheduled as one pipelined group (the DAGScheduler treats an
all-pipelined job's stage graph
+ * as a single group). There is no shuffle-count restriction. An exchange
whose subtree does not
+ * reach the real-time scan is skipped -- the static side of a broadcast
stream-static join must
+ * materialize, because it runs to completion rather than streaming. A
partitioning the pipelined
+ * path cannot serve, such as range partitioning, is rejected up front by
RealTimeModeAllowlist
+ * rather than being handled here.
+ *
+ * The walk does not descend into a ReusedExchangeExec (a leaf whose wrapped
exchange is a field,
+ * not a tree child), so a REUSED shuffle exchange would keep
pipelined=false while its standalone
+ * twin flips to true. That divergence is not reachable: a reused shuffle
requires
+ * referencing the same streaming source more than once (self-join /
self-union / CTE read twice),
+ * which Real-Time Mode rejects when the query starts (MicroBatchExecution,
+ * IDENTICAL_SOURCES_IN_UNION_NOT_SUPPORTED) before this rule runs. The only
ReusedExchangeExec
+ * that reaches an RTM plan wraps a BROADCAST exchange (multiple broadcast
joins on the same
+ * static table, SC-209926), which this rule does not match.
+ *
+ * Fan-out -- one shuffle read by more than one consumer -- is a limitation
that marking a shuffle
+ * here introduces, not one that was already there. A regular materialized
shuffle serves any
+ * number of consumers; a pipelined one is transient and read once, so the
DAGScheduler rejects
+ * fan-out for a PipelinedShuffleDependency specifically
(checkPipelinedGroupsSupportedInRDDGraph,
+ * inert for a regular ShuffleDependency). So a fan-out query that runs fine
unmarked is rejected
+ * once this rule marks it. The check runs in handleJobSubmitted, so it
rejects the batch's job
+ * rather than the query: such a query fails the same way on every batch
instead of failing once
+ * when it is planned. A plan-time guard here would improve only where the
failure is reported,
+ * not whether the query can run.
+ */
+ object MarkPipelinedShuffleForRealTimeMode extends Rule[SparkPlan] {
+ override def apply(plan: SparkPlan): SparkPlan = {
+ val isRealTimeMode = plan.exists(_.isInstanceOf[RealTimeStreamScanExec])
+ if (!isRealTimeMode) {
+ plan
+ } else {
+ markStreamingPath(plan)._1
+ }
+ }
+
+ /**
+ * Marks the shuffles that are on the streaming path -- those whose
subtree reaches a
+ * [[RealTimeStreamScanExec]] -- and returns the rewritten plan along with
whether this
+ * subtree reaches one.
+ *
+ * A plan can hold a static subtree alongside the streaming one: the
static side of a
+ * broadcast stream-static join is planned in the same physical plan and
may contain its own
+ * shuffle. That shuffle materializes normally and is not part of the
pipelined group -- and
+ * cannot be, since a static side runs to completion rather than
streaming. Marking it
+ * pipelined would pull it into the group and demand slots for stages that
must instead
+ * finish, which fails admission (CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT).
This mirrors the
+ * streaming-path detection the operator allowlist uses
(RealTimeModeAllowlist), which only
+ * inspects nodes whose subtree reaches the real-time scan; marking a
wider set than the
+ * allowlist checks would flip shuffles it never validated.
+ */
+ private def markStreamingPath(plan: SparkPlan): (SparkPlan, Boolean) =
plan match {
+ case rts: RealTimeStreamScanExec => (rts, true)
+ case p if p.children.isEmpty => (p, false)
+ case p =>
+ val results = p.children.map(markStreamingPath)
+ val onStreamingPath = results.exists(_._2)
+ val newPlan = p.withNewChildren(results.map(_._1))
+ newPlan match {
+ case s: ShuffleExchangeExec if onStreamingPath && !s.pipelined =>
+ // A bare case-class copy does not carry the node's tags, which is
where the logical
+ // link lives, so copy them over the way the tree transforms do.
+ val marked = s.copy(pipelined = true)
+ marked.copyTagsFrom(s)
+ (marked, true)
+ case other => (other, onStreamingPath)
+ }
+ }
+
+ }
+
+ override def preparations: Seq[Rule[SparkPlan]] =
+ state +: (super.preparations :+ MarkPipelinedShuffleForRealTimeMode)
/** no need to try-catch again as this is already done once */
override def assertAnalyzed(): Unit = analyzed
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
index de84c9d15a6e..39fb0038ad4a 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/MicroBatchExecution.scala
@@ -371,6 +371,31 @@ class MicroBatchExecution(
messageParameters = Map("className" -> s.getClass.getName)
)
}
+
+ // `repartition(n)` uses RoundRobinPartitioning, and by default
(SPARK-23207) Spark inserts a
+ // local sort before a round-robin shuffle so the row-to-partition
assignment is deterministic
+ // -- otherwise a retried task could emit rows in a different order and
lose data. That sort
+ // is fully blocking: it drains its whole input before emitting a row. A
Real-Time Mode task
+ // reads an unbounded stream, so the input never ends, the producer
never emits, and the query
+ // makes no progress at all. Note this sort is not a SortExec node (it
lives inside
+ // ShuffleExchangeExec's RDD), so the Real-Time Mode operator allowlist
cannot catch it.
+ //
+ // Determinism is not needed here: Real-Time Mode does not retry tasks
(TaskSetManager caps a
+ // pipelined task set at one attempt, and its group is aborted as a unit
rather than
+ // recomputed), so the hazard the sort guards against does not arise.
+ //
+ // This overrides an explicit `true` as well. Honouring it would
reintroduce the hang above,
+ // and a query that never produces a row is worse than silently ignoring
a config that only
+ // guards against a retry that cannot happen here. Log at warning level
when overriding an
+ // explicit setting so the operator can see their config is not the one
in force.
+ if
(sparkSessionForStream.conf.contains(SQLConf.SORT_BEFORE_REPARTITION.key) &&
+ sparkSessionForStream.conf.get(SQLConf.SORT_BEFORE_REPARTITION)) {
+ logWarning(log"Ignoring ${MDC(LogKeys.CONFIG,
SQLConf.SORT_BEFORE_REPARTITION.key)}=true " +
+ log"for this Real-Time Mode query: the local sort it inserts before
a round-robin " +
+ log"repartition never completes on an unbounded stream. It is not
needed because a " +
+ log"Real-Time Mode task is not retried.")
+ }
+ sparkSessionForStream.conf.set(SQLConf.SORT_BEFORE_REPARTITION.key,
"false")
}
// Initializing TriggerExecutor relies on `sources`, hence calling this
after initializing
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
index 7ae557797f36..2cae2f7ee2ad 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/runtime/RealTimeModeAllowlist.scala
@@ -19,9 +19,11 @@ package org.apache.spark.sql.execution.streaming.runtime
import org.apache.spark.SparkIllegalArgumentException
import org.apache.spark.internal.{Logging, LogKeys, MessageWithContext}
+import org.apache.spark.sql.catalyst.plans.physical.{Partitioning,
RangePartitioning}
import org.apache.spark.sql.connector.catalog.Table
import org.apache.spark.sql.execution.SparkPlan
import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec
+import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
import org.apache.spark.sql.execution.streaming.operators.stateful._
object RealTimeModeAllowlist extends Logging {
@@ -55,7 +57,17 @@ object RealTimeModeAllowlist extends Logging {
"org.apache.spark.sql.execution.datasources.v2.WriteToDataSourceV2Exec",
"org.apache.spark.sql.execution.exchange.BroadcastExchangeExec",
"org.apache.spark.sql.execution.exchange.ReusedExchangeExec",
+ // A pipelined (streaming) shuffle repartitions a stateful query's input
by key. In Real-Time
+ // Mode the DAGScheduler co-schedules the shuffle's producer and consumer
stages via a
+ // PipelinedShuffleDependency (see IncrementalExecution's
pipelined-shuffle rule), so the
+ // exchange is a supported member of a pipelined group rather than a
materialization barrier.
+ "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec",
"org.apache.spark.sql.execution.joins.BroadcastHashJoinExec",
+ // Streaming deduplication and the state-store access operators it plans
into. These run in the
+ // pipelined-shuffle consumer stage, keyed by the same columns the shuffle
repartitions on.
+
"org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreRestoreExec",
+
"org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec",
+
"org.apache.spark.sql.execution.streaming.operators.stateful.StreamingDeduplicateExec",
classOf[EventTimeWatermarkExec].getName
)
@@ -120,12 +132,37 @@ object RealTimeModeAllowlist extends Logging {
collectNodesWhoseSubtreeHasRTS(root)._2
}
+ /**
+ * Whether this operator is supported in Real-Time Mode.
+ *
+ * Membership in `allowedOperators` is the general rule, but a shuffle needs
a second look: only
+ * some partitionings can run in Real-Time Mode. A `RangePartitioning`
cannot, because building
+ * its `RangePartitioner` runs a separate job that samples the input to
compute range bounds (see
+ * `ShuffleExchangeExec.prepareShuffleDependency`), and that job cannot
complete while the source
+ * keeps producing. Rejecting it here fails the query fast with the usual
allowlist error instead
+ * of letting it stall on a sampling job that never finishes.
+ */
+ private def isAllowed(node: SparkPlan): Boolean = {
+ allowedOperators.contains(node.getClass.getName) && (node match {
+ case e: ShuffleExchangeExec => supportsPartitioning(e.outputPartitioning)
+ case _ => true
+ })
+ }
+
+ /**
+ * Whether a shuffle with this partitioning is supported in Real-Time Mode.
See `isAllowed` for
+ * why range partitioning is not.
+ */
+ private def supportsPartitioning(partitioning: Partitioning): Boolean = {
+ !partitioning.isInstanceOf[RangePartitioning]
+ }
+
def checkAllowedPhysicalOperator(operator: SparkPlan, throwException:
Boolean): Unit = {
val nodesToCheck = collectRealtimeNodes(operator)
val violations = nodesToCheck
.collect {
case node =>
- if (allowedOperators.contains(node.getClass.getName)) {
+ if (isAllowed(node)) {
None
} else {
Some(node.getClass.getName)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
index c1d1f13915a6..136dc57fdfcf 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
@@ -1360,7 +1360,7 @@ class DataFrameWindowFunctionsSuite extends
SharedSparkSession
def isShuffleExecByRequirement(
plan: ShuffleExchangeExec,
desiredClusterColumns: Seq[String]): Boolean = plan match {
- case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS,
_) =>
+ case ShuffleExchangeExec(op: HashPartitioning, _, ENSURE_REQUIREMENTS,
_, _) =>
partitionExpressionsColumns(op.expressions) === desiredClusterColumns
case _ => false
}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
index 3b28cae31a13..096eee137b32 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DatasetSuite.scala
@@ -1942,7 +1942,7 @@ class DatasetSuite extends SharedSparkSession
val agg = cp.groupBy($"id" % 2).agg(count($"id"))
agg.queryExecution.executedPlan.collectFirst {
- case ShuffleExchangeExec(_, _: RDDScanExec, _, _) =>
+ case ShuffleExchangeExec(_, _: RDDScanExec, _, _, _) =>
case BroadcastExchangeExec(_, _: RDDScanExec) =>
}.foreach { _ =>
fail(
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala
index b7798b0bde5d..c6c262379d58 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ExchangeSuite.scala
@@ -238,4 +238,32 @@ class ExchangeSuite extends SharedSparkSession {
"ReusedExchangeExec should reuse an existing exchange")
}
}
+
+ test("ShuffleExchangeExec string args show `pipelined` only when it is set")
{
+ val plan = spark.range(10).selectExpr("id % 4 AS
key").queryExecution.executedPlan
+ val partitioning = HashPartitioning(Seq(Literal(1)), 4)
+
+ // The default: nothing about pipelining is printed, so plan output (and
every golden file that
+ // captures it) is unchanged by the existence of the field.
+ val nonPipelined = ShuffleExchangeExec(partitioning, plan).simpleString(10)
+ assert(!nonPipelined.contains("pipelined") &&
!nonPipelined.contains("false"),
+ s"a non-pipelined exchange should not mention pipelining, but was:
$nonPipelined")
+
+ // When set, it is printed with a name rather than as a bare positional
`true`.
+ val pipelined = ShuffleExchangeExec(partitioning, plan, pipelined =
true).simpleString(10)
+ assert(pipelined.contains("isPipelined=true"),
+ s"a pipelined exchange should report it, but was: $pipelined")
+
+ // `stringArgs` drops `pipelined` by position, so it must stay the last
constructor field --
+ // otherwise a newly added field would be dropped instead and the bare
`false` would come back.
+ val exchange = ShuffleExchangeExec(partitioning, plan)
+ assert(exchange.productElementName(exchange.productArity - 1) ==
"pipelined",
+ "`pipelined` must remain the last constructor field of
ShuffleExchangeExec, because " +
+ "stringArgs drops it positionally")
+
+ // Exchange's plan_id suffix is re-appended by the override and must
survive.
+ Seq(nonPipelined, pipelined).foreach { s =>
+ assert(s.contains("[plan_id="), s"the plan_id suffix should be
preserved, but was: $s")
+ }
+ }
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
index 71192e92af0d..41b74c6d4ff2 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala
@@ -733,10 +733,11 @@ class PlannerSuite extends SharedSparkSession with
AdaptiveSparkPlanHelper {
outputPlan match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _,
-
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _), _, _, _),
_),
+
ShuffleExchangeExec(HashPartitioning(leftPartitioningExpressions, _),
+ _, _, _, _), _),
SortExec(_, _,
ShuffleExchangeExec(HashPartitioning(rightPartitioningExpressions, _),
- _, _, _), _), _) =>
+ _, _, _, _), _), _) =>
assert(leftKeys === smjExec.leftKeys)
assert(rightKeys === smjExec.rightKeys)
assert(leftKeys === leftPartitioningExpressions)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
index 17d00ec055e0..95085b71bc25 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala
@@ -60,7 +60,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec1) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
_) =>
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_), _) =>
assert(leftKeys === Seq(exprA, exprB))
assert(rightKeys === Seq(exprB, exprA))
case other => fail(other.toString)
@@ -71,7 +71,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprB :: exprA :: Nil, Inner, None, plan2, plan1)
EnsureRequirements.apply(smjExec2) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _),
_), _) =>
assert(leftKeys === Seq(exprB, exprA))
assert(rightKeys === Seq(exprA, exprB))
@@ -84,7 +84,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprD :: exprC :: Nil, exprB :: exprA :: Nil, Inner, None, plan1, plan1)
EnsureRequirements.apply(smjExec3) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _),
_), _) =>
assert(leftKeys === Seq(exprC, exprD))
assert(rightKeys === Seq(exprA, exprB))
@@ -126,8 +126,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
)
EnsureRequirements.apply(smjExec2) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _)
=>
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _), _),
_) =>
assert(leftKeys === Seq(exprC, exprB, exprD))
assert(rightKeys === Seq(exprD, exprA, exprC))
case other => fail(other.toString)
@@ -145,7 +145,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan2)
EnsureRequirements.apply(smjExec1) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _)
=>
assert(leftKeys === Seq(exprB, exprA))
assert(rightKeys === Seq(exprB, exprC))
@@ -159,7 +159,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprB :: Nil, Inner, None, plan1, plan3)
EnsureRequirements.apply(smjExec2) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _),
_), _) =>
assert(leftKeys === Seq(exprB, exprA))
assert(rightKeys === Seq(exprB, exprC))
@@ -173,7 +173,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec3) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: PartitioningCollection, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
_) =>
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_), _) =>
assert(leftKeys === Seq(exprB, exprC))
assert(rightKeys === Seq(exprB, exprA))
case other => fail(other.toString)
@@ -319,7 +319,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _),
_) =>
+ SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _),
_), _) =>
assert(leftKeys === Seq(exprA, exprB))
assert(rightKeys === Seq(exprC, exprD))
assert(p.expressions == Seq(exprC))
@@ -335,7 +335,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _), _)
=>
assert(leftKeys === Seq(exprA, exprB))
assert(rightKeys === Seq(exprC, exprD))
@@ -353,7 +353,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _),
_) =>
+ SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _),
_), _) =>
assert(leftKeys === Seq(exprA, exprB))
assert(rightKeys === Seq(exprC, exprD))
assert(p.expressions == Seq(exprC))
@@ -372,7 +372,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _)
=>
+ SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _),
_) =>
assert(leftKeys === Seq(exprA, exprB, exprB))
assert(rightKeys === Seq(exprA, exprC, exprC))
assert(p.expressions == Seq(exprA, exprC, exprA))
@@ -388,7 +388,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _), _), _)
=>
+ SortExec(_, _, ShuffleExchangeExec(p: HashPartitioning, _, _, _, _), _),
_) =>
assert(leftKeys === Seq(exprA, exprB, exprB))
assert(rightKeys === Seq(exprA, exprC, exprD))
assert(p.expressions == Seq(exprA, exprC, exprA))
@@ -444,8 +444,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == 5)
assert(right.numPartitions == 5)
case other => fail(other.toString)
@@ -462,7 +462,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == 10)
assert(right.numPartitions == 10)
assert(right.expressions == Seq(exprC, exprD))
@@ -481,8 +481,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == 5)
assert(left.expressions == Seq(exprA, exprB))
assert(right.numPartitions == 5)
@@ -492,7 +492,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == 1)
assert(right.numPartitions == 1)
assert(right.expressions == Seq(exprC))
@@ -510,8 +510,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == conf.numShufflePartitions)
assert(left.expressions == Seq(exprA, exprB))
assert(right.numPartitions == conf.numShufflePartitions)
@@ -529,7 +529,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
SortExec(_, _, DummySparkPlan(_, _, left: PartitioningCollection, _,
_), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.numPartitions == 10)
assert(right.numPartitions == 10)
assert(right.expressions == Seq(exprA))
@@ -545,7 +545,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
Inner, None, plan1, plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
SortExec(_, _, DummySparkPlan(_, _, right: PartitioningCollection, _,
_), _), _) =>
assert(left.numPartitions == 20)
assert(left.expressions == Seq(exprC))
@@ -587,7 +587,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.expressions === Seq(exprA, exprB))
assert(right.expressions === Seq(exprC, exprD))
assert(left.numPartitions == 6)
@@ -606,7 +606,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB))
assert(right.expressions === Seq(exprC, exprD))
@@ -624,7 +624,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB))
assert(right.expressions === Seq(exprC, exprD))
@@ -644,7 +644,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
SortExec(_, _, DummySparkPlan(_, _, right: HashPartitioning, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB))
assert(right.expressions === Seq(exprC, exprD))
@@ -666,8 +666,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: Nil, exprC :: exprD :: Nil, Inner, None, plan1,
plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.expressions === Seq(exprA, exprB))
assert(right.expressions === Seq(exprC, exprD))
assert(left.numPartitions == conf.numShufflePartitions)
@@ -691,7 +691,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, left: HashPartitioning, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _,
_), _), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _,
_, _), _), _) =>
assert(leftKeys === Seq(exprA, exprB))
assert(rightKeys === Seq(exprC, exprD))
assert(left.numPartitions == 9)
@@ -714,8 +714,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
var smjExec = SortMergeJoinExec(exprA :: Nil, exprC :: Nil, Inner, None,
plan1, plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(leftKeys === Seq(exprA))
assert(rightKeys === Seq(exprC))
assert(left.numPartitions == 20)
@@ -733,7 +733,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
SortExec(_, _, DummySparkPlan(_, _, _: HashPartitioning, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(leftKeys === Seq(exprA))
assert(rightKeys === Seq(exprC))
assert(right.numPartitions == 10)
@@ -765,8 +765,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
} else {
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(leftKeys, rightKeys, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _,
_), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _,
_), _), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _,
_, _), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _,
_, _), _), _) =>
assert(leftKeys === Seq(exprA))
assert(rightKeys === Seq(exprC))
assert(left.numPartitions == 5)
@@ -875,8 +875,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner,
None, plan1, plan2)
EnsureRequirements.apply(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _,
_), _), _) =>
assert(left.expressions === Seq(exprA, exprB, exprB))
assert(right.expressions === Seq(exprA, exprC, exprC))
case other => fail(other.toString)
@@ -951,8 +951,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: exprC :: Nil, exprA :: exprB :: exprC :: Nil, Inner,
None, plan1, plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB, exprC))
assert(right.expressions === Seq(exprA, exprB, exprC))
case other => fail(other.toString)
@@ -969,8 +969,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner,
None, plan1, plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB, exprB))
assert(right.expressions === Seq(exprA, exprC, exprC))
case other => fail(other.toString)
@@ -987,8 +987,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner,
None, plan1, plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB, exprB))
assert(right.expressions === Seq(exprA, exprC, exprC))
case other => fail(other.toString)
@@ -1008,8 +1008,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
exprA :: exprB :: exprB :: Nil, exprA :: exprC :: exprC :: Nil, Inner,
None, plan1, plan2)
applyEnsureRequirementsWithSubsetKeys(smjExec) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(left: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(right: HashPartitioning, _, _, _, _),
_), _) =>
assert(left.expressions === Seq(exprA, exprB, exprB))
assert(right.expressions === Seq(exprA, exprC, exprC))
case other => fail(other.toString)
@@ -1130,7 +1130,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
case ShuffledHashJoinExec(_, _, _, _, _,
DummySparkPlan(_, _, left: KeyedPartitioning, _, _),
ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _),
- DummySparkPlan(_, _, SinglePartition, _, _), _, _), _) =>
+ DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) =>
assert(left.expressions == a1 :: Nil)
assert(attrs == a1 :: Nil)
assert(partitionKeys == pks.map(_.row))
@@ -1234,8 +1234,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(join) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _,
_), _), _) =>
// Both sides should be shuffled to default partitions
assert(p1.numPartitions == 10)
assert(p2.numPartitions == 10)
@@ -1259,8 +1259,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(join) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_), _) =>
// Both sides shuffled due to key mismatch
case other => fail(s"Expected shuffles on both sides, but got: $other")
}
@@ -1278,7 +1278,7 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(join) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
SortExec(_, _, _: DummySparkPlan, _), _) =>
// Left side shuffled, right side kept as-is
case other => fail(s"Expected shuffle on the left side, but got:
$other")
@@ -1296,8 +1296,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(join) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _),
- SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_),
+ SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _, _),
_), _) =>
// Both sides shuffled due to canCreatePartitioning = false
case other => fail(s"Expected shuffles on both sides, but got: $other")
}
@@ -1371,8 +1371,8 @@ class EnsureRequirementsSuite extends SharedSparkSession {
EnsureRequirements.apply(join) match {
case SortMergeJoinExec(_, _, _, _,
- SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _),
_),
- SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _),
_), _) =>
+ SortExec(_, _, ShuffleExchangeExec(p1: HashPartitioning, _, _, _,
_), _),
+ SortExec(_, _, ShuffleExchangeExec(p2: HashPartitioning, _, _, _,
_), _), _) =>
// Both sides should be shuffled because partition keys not in join
keys
assert(p1.numPartitions == 10)
assert(p2.numPartitions == 10)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
index a9b8d54a6c61..de4af97656a6 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeAllowlistSuite.scala
@@ -70,7 +70,6 @@ class StreamRealTimeModeAllowlistSuite extends
StreamRealTimeModeE2ESuiteBase {
"errorType" -> "operator",
"message" -> (
"org.apache.spark.sql.execution.SortExec, " +
- "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec,
" +
"org.apache.spark.sql.execution.joins.SortMergeJoinExec are"
)
)
@@ -107,31 +106,34 @@ class StreamRealTimeModeAllowlistSuite extends
StreamRealTimeModeE2ESuiteBase {
}
}
- // TODO(SPARK-54237) : Remove this test after RTM can shuffle to multiple
stages
- test("repartition not allowed") {
- val inputData = LowLatencyMemoryStream[Int](2)
+ // A repartitionByRange produces a range-partitioned shuffle. Building its
RangePartitioner runs a
+ // separate job that samples the input to compute range bounds, which cannot
complete while the
+ // source keeps producing, so it must be rejected up front rather than
stalling the query.
+ test("range-partitioned shuffle not allowed") {
+ val inputData = LowLatencyMemoryStream[Int](2)
- val df = inputData.toDF()
- .select(col("value").as("key"))
- .repartition(4, col("key"))
+ val df = inputData.toDF()
+ .select(col("value").as("key"))
+ .repartitionByRange(3, col("key"))
+ .select(col("key"))
- val query = runStreamingQuery("repartition_allowlist", df)
+ val query = runStreamingQuery("range_shuffle_allowlist", df)
- eventually(timeout(60.seconds)) {
- checkError(
- exception =
query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException],
- condition =
"STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST",
- parameters = Map(
- "errorType" -> "operator",
- "message" -> (
- "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec
is"
- )
- )
+ eventually(timeout(60.seconds)) {
+ checkError(
+ exception =
query.exception.get.getCause.asInstanceOf[SparkIllegalArgumentException],
+ condition =
"STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST",
+ parameters = Map(
+ "errorType" -> "operator",
+ "message" ->
"org.apache.spark.sql.execution.exchange.ShuffleExchangeExec is"
)
- }
+ )
+ }
}
- // TODO(SPARK-54236) : Remove this test after RTM supports stateful queries
+ // A stateful aggregation is still rejected under RTM because
HashAggregateExec is not in the
+ // allowlist, even though the pipelined shuffle and the StateStore operators
now are. When RTM
+ // supports aggregation (SPARK-54236), remove this test.
test("stateful queries not allowed") {
val inputData = LowLatencyMemoryStream[Int](2)
@@ -141,7 +143,7 @@ class StreamRealTimeModeAllowlistSuite extends
StreamRealTimeModeE2ESuiteBase {
.count()
.select(concat(col("key"), lit("-"), col("count")))
- val query = runStreamingQuery("repartition_allowlist", df)
+ val query = runStreamingQuery("stateful_allowlist", df)
eventually(timeout(60.seconds)) {
checkError(
@@ -149,13 +151,7 @@ class StreamRealTimeModeAllowlistSuite extends
StreamRealTimeModeE2ESuiteBase {
condition =
"STREAMING_REAL_TIME_MODE.OPERATOR_OR_SINK_NOT_IN_ALLOWLIST",
parameters = Map(
"errorType" -> "operator",
- "message" -> (
- "org.apache.spark.sql.execution.aggregate.HashAggregateExec, " +
- "org.apache.spark.sql.execution.exchange.ShuffleExchangeExec, " +
- "org.apache.spark.sql.execution.streaming" +
- ".operators.stateful.StateStoreRestoreExec, " +
-
"org.apache.spark.sql.execution.streaming.operators.stateful.StateStoreSaveExec
are"
- )
+ "message" ->
"org.apache.spark.sql.execution.aggregate.HashAggregateExec is"
)
)
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala
index 0f584d818c06..689d89e19e00 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala
@@ -17,18 +17,23 @@
package org.apache.spark.sql.streaming
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicInteger
+import java.io.IOException
+import java.util.concurrent.{ConcurrentHashMap, TimeUnit}
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
import scala.concurrent.duration.Duration
import org.scalatest.concurrent.PatienceConfiguration.Timeout
-import org.apache.spark.{SparkIllegalArgumentException,
SparkIllegalStateException, TaskContext}
+import org.apache.spark.{SparkException, SparkIllegalArgumentException,
SparkIllegalStateException, TaskContext}
+import org.apache.spark.scheduler.{SparkListener, SparkListenerJobStart,
SparkListenerStageCompleted, SparkListenerStageSubmitted}
+import org.apache.spark.sql.execution.datasources.v2.RealTimeStreamScanExec
+import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec,
ShuffleExchangeExec}
import org.apache.spark.sql.execution.streaming.RealTimeTrigger
-import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
+import org.apache.spark.sql.execution.streaming.runtime.{MemoryStream,
StreamExecution}
import org.apache.spark.sql.execution.streaming.sources.{ContinuousMemorySink,
LowLatencyMemoryStream}
-import org.apache.spark.sql.functions.udf
+import
org.apache.spark.sql.execution.streaming.state.{FailureInjectionCheckpointFileManager,
FailureInjectionFileSystem}
+import org.apache.spark.sql.functions.{broadcast, concat, lit, udf}
import org.apache.spark.sql.internal.SQLConf
class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase {
@@ -209,6 +214,74 @@ class StreamRealTimeModeSuite extends
StreamRealTimeModeSuiteBase {
StopStream
)
}
+
+ test("pipelined shuffle: a static-side shuffle is not marked pipelined") {
+ // A broadcast stream-static join can carry a shuffle on its STATIC side,
in a subtree with no
+ // RealTimeStreamScanExec. That shuffle must materialize normally: a
static side runs to
+ // completion rather than streaming, so pulling it into the pipelined
group would demand
+ // concurrent slots for a stage that must instead finish, failing
admission. Only shuffles on
+ // the streaming path are marked, matching what the operator allowlist
actually validates.
+ withSQLConf(
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "10485760") {
+ val staticData = spark.range(0, 3).selectExpr("id as sk", "id * 10 as
sv")
+ .repartition(3, $"sk")
+ val inputData = LowLatencyMemoryStream[(String, Int)](2)
+ val streamDf = inputData.toDF().select($"_1".as("key"),
$"_2".cast("long").as("sk"))
+ val result = streamDf
+ .repartition(2, $"sk")
+ .join(broadcast(staticData), Seq("sk"), "left")
+ .select($"key")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 2)),
+ StartStream(),
+ CheckAnswerWithTimeout(60000, "a", "b"),
+ Execute { q =>
+ val exchanges = q.lastExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }
+ // The streaming-side repartition is pipelined; a static-side
shuffle, if it survives
+ // into this plan, must not be.
+ val streamingSide = exchanges.filter(_.exists {
+ case _: RealTimeStreamScanExec => true
+ case _ => false
+ })
+ assert(streamingSide.nonEmpty, "expected a shuffle on the streaming
path")
+ assert(streamingSide.forall(_.pipelined),
+ "a streaming-path shuffle must be pipelined")
+
assert(exchanges.filterNot(streamingSide.contains).forall(!_.pipelined),
+ "a static-side shuffle must not be marked pipelined")
+ },
+ StopStream
+ )
+ }
+ }
+
+ test("multiple broadcast joins with the same static table run in Real-Time
Mode") {
+ // Two broadcast joins against the SAME static table -- exchange reuse
collapses the second
+ // broadcast into a ReusedExchangeExec. This is the supported reuse shape
in Real-Time Mode: a
+ // reused BROADCAST exchange (allowlisted), NOT a reused shuffle. The RTM
marking rule only
+ // marks ShuffleExchangeExec, so it correctly leaves the broadcast reuse
alone.
+ val inputData = LowLatencyMemoryStream.singlePartition[Int]
+ val staticData = Seq((1, "a"), (2, "b"), (3, "c")).toDF("key", "value")
+ val df = inputData.toDS().toDF("key").join(broadcast(staticData),
Seq("key"), "left")
+ val df2 = df.join(broadcast(staticData), Seq("key"), "left")
+ testStream(df2, OutputMode.Update, Map.empty, new ContinuousMemorySink())(
+ AddData(inputData, 1, 2, 3),
+ StartStream(),
+ CheckAnswerWithTimeout(30000, (1, "a", "a"), (2, "b", "b"), (3, "c",
"c")),
+ Execute { q =>
+ val plan = q.lastExecution.executedPlan
+ // Verify the shape this test is about, not just the answer: the
second broadcast is reused,
+ // and the marking rule left it alone because it only marks shuffle
exchanges.
+ assert(plan.exists(_.isInstanceOf[ReusedExchangeExec]),
+ s"expected the second broadcast to be reused, got:\n$plan")
+ assert(plan.collect { case s: ShuffleExchangeExec => s }.isEmpty,
+ "a broadcast stream-static join should introduce no shuffle
exchange")
+ },
+ StopStream
+ )
+ }
}
class StreamRealTimeModeWithManualClockSuite extends
StreamRealTimeModeManualClockSuiteBase {
@@ -393,4 +466,339 @@ class StreamRealTimeModeWithManualClockSuite extends
StreamRealTimeModeManualClo
StopStream
)
}
+
+ //
========================================================================================
+ // Pipelined (streaming) shuffle: a stateful/repartition Real-Time Mode
query whose shuffle is a
+ // PipelinedShuffleDependency, so the producer (source scan) and consumer
stages are co-scheduled
+ // and stream records through a transient shuffle instead of the consumer
waiting for the producer
+ // to fully materialize.
+ //
========================================================================================
+
+ override def beforeEach(): Unit = {
+ super.beforeEach()
+ StreamRealTimeModeSuite.failTasks = false
+ }
+
+ /** Assert every shuffle exchange in the query's last executed plan is
pipelined. */
+ private def assertAllExchangesPipelined(q: StreamExecution): Unit = {
+ val exchanges = q.lastExecution.executedPlan.collect { case s:
ShuffleExchangeExec => s }
+ assert(exchanges.nonEmpty, "expected at least one shuffle exchange in the
plan")
+ assert(exchanges.forall(_.pipelined),
+ "expected all Real-Time Mode shuffle exchanges to be pipelined, got: " +
+ exchanges.map(e => s"pipelined=${e.pipelined}").mkString(", "))
+ }
+
+ test("pipelined shuffle: stateful dedup runs in Real-Time Mode and
co-schedules its stages") {
+ // Track, from the driver, whether the producer (source scan) and consumer
(dedup) stages of the
+ // pipelined group were ever RUNNING simultaneously. A sequential
producer-then-consumer
+ // schedule never exceeds one running stage at a time; >= 2 proves genuine
co-scheduling.
+ val runningStages = ConcurrentHashMap.newKeySet[Int]()
+ val maxConcurrentStages = new AtomicInteger(0)
+ val queryStageIds = ConcurrentHashMap.newKeySet[Int]()
+ // Count only stages belonging to the query under test. The suite shares
one SparkContext, so a
+ // stage from any other streaming query would otherwise satisfy the
co-scheduling assertion
+ // below even if this query's producer and consumer actually ran one after
the other. The id is
+ // captured from the query once it is running, and every job is matched
against it.
+ val queryId = new AtomicReference[String](null)
+ val listener = new SparkListener {
+ override def onJobStart(e: SparkListenerJobStart): Unit = {
+ // StreamExecution tags every streaming job with its query id.
+ val id = queryId.get()
+ if (id != null &&
e.properties.getProperty(StreamExecution.QUERY_ID_KEY) == id) {
+ e.stageIds.foreach(queryStageIds.add(_))
+ }
+ }
+ override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = {
+ if (queryStageIds.contains(e.stageInfo.stageId)) {
+ runningStages.add(e.stageInfo.stageId)
+ maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max)
+ }
+ }
+ override def onStageCompleted(e: SparkListenerStageCompleted): Unit = {
+ runningStages.remove(e.stageInfo.stageId)
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ try {
+ val inputData = LowLatencyMemoryStream[(String, Int)]
+ // scan --shuffle(repartition by key)--> streaming dropDuplicates -->
sink.
+ val result =
inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2), ("b", 2),
("c", 2)),
+ StartStream(),
+ // StartStream only launches the query -- the stream runs on its own
thread and can submit
+ // the first batch's jobs before the id is recorded here, in which
case those jobs are not
+ // attributed. Rather than race, the assertion below relies on the
SECOND batch, which is
+ // driven after this point and so is always observed in full.
+ Execute(q => queryId.set(q.id.toString)),
+ CheckAnswerWithTimeout(60000, "a", "b", "c"),
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)),
+ CheckAnswerWithTimeout(60000, "a", "b", "c", "d"),
+ Execute { q =>
+ assertAllExchangesPipelined(q)
+ assert(maxConcurrentStages.get() >= 2,
+ s"expected >= 2 stages running concurrently, saw max
${maxConcurrentStages.get()}")
+ },
+ StopStream
+ )
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+
+ test("pipelined shuffle: multi-key dedup runs in Real-Time Mode over a
pipelined shuffle") {
+ // dropDuplicates on more than one column: the hash-partitioning is by the
composite key.
+ // Confirms the pipelined path is not specific to a single dedup column.
+ val inputData = LowLatencyMemoryStream[(String, Int)]
+ val result = inputData.toDF().select($"_1".as("k1"), $"_2".as("k2"))
+ .dropDuplicates("k1", "k2")
+ .select(concat($"k1", lit("-"), $"k2").as("out"))
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ // (a,1) twice -> once; (a,2) is a distinct composite key -> also
emitted.
+ AddData(inputData, ("a", 1), ("a", 1), ("a", 2), ("b", 1)),
+ StartStream(),
+ CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1"),
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ AddData(inputData, ("a", 1), ("a", 2), ("b", 2)),
+ CheckAnswerWithTimeout(60000, "a-1", "a-2", "b-1", "b-2"),
+ Execute { q => assertAllExchangesPipelined(q) },
+ StopStream
+ )
+ }
+
+ test("pipelined shuffle: an explicit repartition-by-key runs in Real-Time
Mode") {
+ // A stateless repartition (no dedup): the ShuffleExchangeExec is still
marked pipelined in RTM.
+ val inputData = LowLatencyMemoryStream[Int]
+ val result = inputData.toDF().repartition(4, $"value").select($"value")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, 1, 2, 3, 4, 5),
+ StartStream(),
+ CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5),
+ Execute { q => assertAllExchangesPipelined(q) },
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ AddData(inputData, 6, 7),
+ CheckAnswerWithTimeout(60000, 1, 2, 3, 4, 5, 6, 7),
+ StopStream
+ )
+ }
+
+ test("pipelined shuffle: a chain of two shuffles is a single co-scheduled
pipelined group") {
+ // A round-robin repartition (RoundRobinPartitioning) feeding a
dropDuplicates (HashPartitioning
+ // on the dedup key) produces TWO distinct shuffle exchanges -- neither
distribution satisfies
+ // the other, so they are not collapsed. BOTH must be marked pipelined and
the all-pipelined job
+ // co-schedules as one group (>= 2 stages running at once). This is the
multi-shuffle-per-
+ // group scenario: more than one pipelined shuffle in a single Real-Time
Mode job.
+ val runningStages = ConcurrentHashMap.newKeySet[Int]()
+ val maxConcurrentStages = new AtomicInteger(0)
+ val queryStageIds = ConcurrentHashMap.newKeySet[Int]()
+ // Count only stages belonging to the query under test. The suite shares
one SparkContext, so a
+ // stage from any other streaming query would otherwise satisfy the
co-scheduling assertion
+ // below even if this query's producer and consumer actually ran one after
the other. The id is
+ // captured from the query once it is running, and every job is matched
against it.
+ val queryId = new AtomicReference[String](null)
+ val listener = new SparkListener {
+ override def onJobStart(e: SparkListenerJobStart): Unit = {
+ // StreamExecution tags every streaming job with its query id.
+ val id = queryId.get()
+ if (id != null &&
e.properties.getProperty(StreamExecution.QUERY_ID_KEY) == id) {
+ e.stageIds.foreach(queryStageIds.add(_))
+ }
+ }
+ override def onStageSubmitted(e: SparkListenerStageSubmitted): Unit = {
+ if (queryStageIds.contains(e.stageInfo.stageId)) {
+ runningStages.add(e.stageInfo.stageId)
+ maxConcurrentStages.accumulateAndGet(runningStages.size(), Math.max)
+ }
+ }
+ override def onStageCompleted(e: SparkListenerStageCompleted): Unit = {
+ runningStages.remove(e.stageInfo.stageId)
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ // Keep the shuffle partition count small: the whole group (scan + 2
shuffles) must fit in the
+ // test cluster's slots at once (gang admission), so a 2-shuffle chain at
the default 200
+ // partitions would fail with CONCURRENT_SCHEDULER_INSUFFICIENT_SLOT.
+ try {
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val inputData = LowLatencyMemoryStream[(String, Int)]
+ val result = inputData.toDF().select($"_1".as("key"))
+ .repartition(4) // RoundRobinPartitioning -> shuffle #1
+ .dropDuplicates("key") // HashPartitioning(key) -> shuffle #2
+ .select($"key")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)),
+ StartStream(),
+ // StartStream only launches the query -- the stream runs on its own
thread and can submit
+ // the first batch's jobs before the id is recorded here, in which
case those jobs are not
+ // attributed and contribute nothing to maxConcurrentStages. So do
not assert on the first
+ // batch: record the id, let batch 0 finish, then drive a SECOND
batch that is guaranteed
+ // to run entirely after this point, and assert on that one.
+ Execute(q => queryId.set(q.id.toString)),
+ CheckAnswerWithTimeout(60000, "a", "b", "c"),
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ AddData(inputData, ("d", 1), ("e", 1), ("d", 2)),
+ CheckAnswerWithTimeout(60000, "a", "b", "c", "d", "e"),
+ Execute { q =>
+ val exchanges = q.lastExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got
${exchanges.size}")
+ assert(exchanges.forall(_.pipelined),
+ "every exchange in the chain must be pipelined, got: " +
+ exchanges.map(_.pipelined).mkString(", "))
+ assert(maxConcurrentStages.get() >= 2,
+ s"a 2-shuffle chain must co-schedule >= 2 stages, saw
${maxConcurrentStages.get()}")
+ },
+ StopStream
+ )
+ }
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+
+ test("pipelined shuffle: a chain of two round-robin repartitions runs in
Real-Time Mode") {
+ withSQLConf(SQLConf.SHUFFLE_PARTITIONS.key -> "2") {
+ val inputData = LowLatencyMemoryStream[(String, Int)](2)
+ // Two round-robin repartitions with a map between them; consecutive
repartitions collapse to
+ // the last one, so the map is what keeps both shuffles in the plan.
Both are pipelined, and
+ // the second one's producer reads the first one's output -- an
UNORDERED input. A round-robin
+ // shuffle is order-sensitive once the deterministic local sort is off,
which would escalate
+ // that producer to INDETERMINATE and get the group rejected; a
pipelined shuffle is exempt
+ // from the order-sensitive marking precisely so this shape runs.
+ val result = inputData.toDS()
+ .repartition(3)
+ .map(row => row)
+ .repartition(2)
+ .map(row => row)
+ .toDF().select($"_1".as("key"))
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 2), ("c", 3)),
+ StartStream(),
+ CheckAnswerWithTimeout(60000, "a", "b", "c"),
+ Execute { q =>
+ val exchanges = q.lastExecution.executedPlan.collect {
+ case s: ShuffleExchangeExec => s
+ }
+ assert(exchanges.size >= 2, s"expected >= 2 shuffle exchanges, got
${exchanges.size}")
+ assert(exchanges.forall(_.pipelined),
+ "every exchange in the chain must be pipelined, got: " +
+ exchanges.map(_.pipelined).mkString(", "))
+ },
+ StopStream
+ )
+ }
+ }
+
+ test("pipelined shuffle: an explicit sortBeforeRepartition=true does not
stall the query") {
+ // The deterministic local sort before a round-robin repartition never
drains an unbounded
+ // stream, so Real-Time Mode overrides this config even when it is set
explicitly -- honouring
+ // it would hang the query forever. See MicroBatchExecution.
+ withSQLConf(
+ SQLConf.SHUFFLE_PARTITIONS.key -> "2",
+ SQLConf.SORT_BEFORE_REPARTITION.key -> "true") {
+ val inputData = LowLatencyMemoryStream[(String, Int)](2)
+ val result = inputData.toDF().select($"_1".as("key"))
+ .repartition(4)
+ .dropDuplicates("key")
+ .select($"key")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 1), ("c", 1), ("a", 2)),
+ StartStream(),
+ CheckAnswerWithTimeout(60000, "a", "b", "c"),
+ StopStream
+ )
+ }
+ }
+
+ test("pipelined shuffle: dedup recovers from a task failure via checkpoint
restart") {
+ withTempDir { checkpointDir =>
+ // A UDF that throws on demand, placed after the dedup so the failure
lands in the pipelined
+ // consumer stage while the query runs. RTM does not retry tasks, so one
failure fails it.
+ val failUDF = udf { (key: String) =>
+ if (StreamRealTimeModeSuite.failTasks) {
+ throw new RuntimeException(s"forced task failure on $key")
+ }
+ key
+ }
+ val inputData = LowLatencyMemoryStream[(String, Int)]
+ val result =
inputData.toDF().select($"_1".as("key")).dropDuplicates("key")
+ .select(failUDF($"key").as("key"))
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 1), ("a", 2)),
+ StartStream(checkpointLocation = checkpointDir.getAbsolutePath),
+ CheckAnswerWithTimeout(60000, "a", "b"),
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ Execute { _ => StreamRealTimeModeSuite.failTasks = true },
+ AddData(inputData, ("c", 1)),
+ advanceRealTimeClock,
+ ExpectFailure[SparkException] { ex =>
+ val msg =
Option(ex.getCause).map(_.getMessage).getOrElse(ex.getMessage)
+ assert(msg != null && msg.contains("forced task failure"),
+ s"expected a forced task failure, got: $msg")
+ }
+ )
+ // Restart from the same checkpoint: batch-0 dedup state must survive
(a, b not re-emitted),
+ // only the genuinely-new c, d appear -- recovery to the last committed
batch.
+ StreamRealTimeModeSuite.failTasks = false
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ StartStream(checkpointLocation = checkpointDir.getAbsolutePath),
+ AddData(inputData, ("a", 3), ("b", 3), ("c", 3), ("d", 1)),
+ CheckAnswerWithTimeout(60000, "c", "d"),
+ advanceRealTimeClock,
+ StopStream
+ )
+ }
+ }
+
+ test("pipelined shuffle: dedup recovers from a commit-log write failure") {
+ withSQLConf(
+ SQLConf.STREAMING_CHECKPOINT_FILE_MANAGER_CLASS.parent.key ->
+ classOf[FailureInjectionCheckpointFileManager].getName) {
+ withTempDir { checkpointDir =>
+ val injectionState =
FailureInjectionFileSystem.registerTempPath(checkpointDir.getPath)
+ try {
+ val inputData = LowLatencyMemoryStream[(String, Int)]
+ val result =
inputData.toDF().select($"_1".as("key")).dropDuplicates("key").select($"key")
+ // Batch 0 dedups a, b and commits. Then fail the close() of batch
1's commit-log write so
+ // batch 1 cannot commit and the query fails after processing.
+ injectionState.createAtomicDelayCloseRegex = Seq(".*/commits/1")
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ AddData(inputData, ("a", 1), ("b", 1)),
+ StartStream(checkpointLocation = checkpointDir.getAbsolutePath),
+ CheckAnswerWithTimeout(60000, "a", "b"),
+ advanceRealTimeClock,
+ WaitUntilBatchProcessed(0),
+ AddData(inputData, ("c", 1)),
+ CheckAnswerWithTimeout(60000, "a", "b", "c"),
+ advanceRealTimeClock,
+ ExpectFailure[IOException]()
+ )
+ // Clear injection, restart: batch-0 state survives (a, b seen);
uncommitted batch 1
+ // re-runs so its new key c is still emitted, plus a further new key
d.
+ injectionState.createAtomicDelayCloseRegex = Seq.empty
+ testStream(result, OutputMode.Update, Map.empty, new
ContinuousMemorySink())(
+ StartStream(checkpointLocation = checkpointDir.getAbsolutePath),
+ AddData(inputData, ("a", 2), ("b", 2), ("c", 2), ("d", 1)),
+ CheckAnswerWithTimeout(60000, "c", "d"),
+ advanceRealTimeClock,
+ StopStream
+ )
+ } finally {
+
FailureInjectionFileSystem.removePathFromTempToInjectionState(checkpointDir.getPath)
+ }
+ }
+ }
+ }
+}
+
+/** Driver-side switch a UDF reads on executors to fail a task on demand
(fault-tolerance tests). */
+object StreamRealTimeModeSuite {
+ @volatile var failTasks: Boolean = false
}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
index c46f0076721b..c0983f338abe 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
@@ -677,8 +677,8 @@ abstract class StreamingInnerJoinBase extends
StreamingJoinSuite {
assert(query.lastExecution.executedPlan.collect {
case j @ StreamingSymmetricHashJoinExec(_, _, _, _, _, _, _, _, _,
- ShuffleExchangeExec(opA: HashPartitioning, _, _, _),
- ShuffleExchangeExec(opB: HashPartitioning, _, _, _), _)
+ ShuffleExchangeExec(opA: HashPartitioning, _, _, _, _),
+ ShuffleExchangeExec(opB: HashPartitioning, _, _, _, _), _)
if partitionExpressionsColumns(opA.expressions) === Seq("a", "b")
&& partitionExpressionsColumns(opB.expressions) === Seq("a",
"b")
&& opA.numPartitions == numPartitions && opB.numPartitions ==
numPartitions => j
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
