Skip to content

[SPARK-57638][SQL] Avoid busy-waiting in Declarative Pipelines flow resolution #56700

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
LuciferYang wants to merge 4 commits into from
Closed

[SPARK-57638][SQL] Avoid busy-waiting in Declarative Pipelines flow resolution #56700

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e348463
[SPARK-57638][SQL] Avoid busy-waiting in Declarative Pipelines flow r…
LuciferYang Jun 23, 2026
3e138a4
[SPARK-57638][SQL][FOLLOWUP] Explain why the take() branch guards on …
LuciferYang Jun 24, 2026
4b86d97
[SPARK-57638][SQL][FOLLOWUP] Add deterministic tests for resolution w…
LuciferYang Jun 30, 2026
70268c3
[SPARK-57638][SQL][FOLLOWUP] Make the re-queue test deterministic and…
LuciferYang Jun 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
263 changes: 136 additions & 127 deletions ...elines/src/main/scala/org/apache/spark/sql/pipelines/graph/DataflowGraphTransformer.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,10 @@ import java.util.concurrent.{
ConcurrentLinkedDeque,
ConcurrentLinkedQueue,
ExecutionException,
ExecutorCompletionService,
Future
}

import scala.collection.mutable.ArrayBuffer
import scala.jdk.CollectionConverters._
import scala.util.control.NoStackTrace

Expand Down Expand Up @@ -134,159 +134,168 @@ class DataflowGraphTransformer(graph: DataflowGraph) extends AutoCloseable {
val failedFlowsQueue = new ConcurrentLinkedQueue[ResolutionFailedFlow]()
val failedDependentFlows = new ConcurrentHashMap[TableIdentifier, Seq[ResolutionFailedFlow]]()

var futures = ArrayBuffer[Future[Unit]]()
val completionService = new ExecutorCompletionService[Unit](executor)
var outstanding = 0
val toBeResolvedFlows = new ConcurrentLinkedDeque[Flow]()
toBeResolvedFlows.addAll(flows.asJava)

while (futures.nonEmpty || toBeResolvedFlows.peekFirst() != null) {
val (done, notDone) = futures.partition(_.isDone)
// Explicitly call future.get() to propagate exceptions one by one if any
// Waits on a finished resolution task and propagates its exception, if any.
def reap(finished: Future[Unit]): Unit = {
try {
done.foreach(_.get())
finished.get()
} catch {
case exn: ExecutionException =>
// Computation threw the exception that is the cause of exn
throw exn.getCause
}
futures = notDone
val flowOpt = {
// We only schedule [[batchSize]] number of flows in parallel.
if (futures.size < batchSize) {
Option(toBeResolvedFlows.pollFirst())
} else {
None
}
outstanding -= 1
}

while (outstanding > 0 || toBeResolvedFlows.peekFirst() != null) {
// Reap every resolution task that has already finished, without blocking.
var finished = completionService.poll()
while (finished != null) {
reap(finished)
finished = completionService.poll()
}
flowOpt.foreach { flow =>
futures.append(
executor.submit(
() =>
// We only schedule [[batchSize]] number of flows in parallel.
if (outstanding < batchSize && toBeResolvedFlows.peekFirst() != null) {
val flow = toBeResolvedFlows.pollFirst()
outstanding += 1
completionService.submit(
() =>
try {
try {
try {
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
val result = transformer(flow, Seq.empty)
require(
result.forall(_.isInstanceOf[ResolvedFlow]),
"transformer must return a Seq[Flow]"
)
// Note: Flow don't need their inputs passed, so for now we send empty Seq.
val result = transformer(flow, Seq.empty)
require(
result.forall(_.isInstanceOf[ResolvedFlow]),
"transformer must return a Seq[Flow]"
)

val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
resolvedFlowsMap.put(flow.identifier, transformedFlows)
resolvedFlows.addAll(transformedFlows.asJava)
} catch {
case e: TransformNodeRetryableException =>
val datasetIdentifier = e.datasetIdentifier
failedDependentFlows.compute(
datasetIdentifier,
(_, flows) => {
// Don't add the input flow back but the failed flow object
// back which has relevant failure information.
val failedFlow = e.failedNode
if (flows == null) {
Seq(failedFlow)
} else {
flows :+ failedFlow
}
val transformedFlows = result.map(_.asInstanceOf[ResolvedFlow])
resolvedFlowsMap.put(flow.identifier, transformedFlows)
resolvedFlows.addAll(transformedFlows.asJava)
} catch {
case e: TransformNodeRetryableException =>
val datasetIdentifier = e.datasetIdentifier
failedDependentFlows.compute(
datasetIdentifier,
(_, flows) => {
// Don't add the input flow back but the failed flow object
// back which has relevant failure information.
val failedFlow = e.failedNode
if (flows == null) {
Seq(failedFlow)
} else {
flows :+ failedFlow
}
)
// Between the time the flow started and finished resolving, perhaps the
// dependent dataset was resolved
resolvedFlowDestinationsMap.computeIfPresent(
datasetIdentifier,
(_, resolved) => {
if (resolved) {
// Check if the dataset that the flow is dependent on has been resolved
// and if so, remove all dependent flows from the failedDependentFlows and
// add them to the toBeResolvedFlows queue for retry.
failedDependentFlows.computeIfPresent(
datasetIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
}
resolved
}
)
// Between the time the flow started and finished resolving, perhaps the
// dependent dataset was resolved
resolvedFlowDestinationsMap.computeIfPresent(
datasetIdentifier,
(_, resolved) => {
if (resolved) {
// Check if the dataset that the flow is dependent on has been resolved
// and if so, remove all dependent flows from the failedDependentFlows and
// add them to the toBeResolvedFlows queue for retry.
failedDependentFlows.computeIfPresent(
datasetIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
}
resolved
}
)
case other: Throwable => throw other
}
// If all flows to this particular destination are resolved, move to the destination
// node transformer
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
resolvedFlowsMap.containsKey(flowToDestination.identifier)
})) {
// If multiple flows completed in parallel, ensure we resolve the destination only
// once by electing a leader via computeIfAbsent
var isCurrentThreadLeader = false
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
isCurrentThreadLeader = true
// Set initial value as false as flow destination is not resolved yet.
false
})
if (isCurrentThreadLeader) {
if (tableMap.contains(flow.destinationIdentifier)) {
val transformed =
transformer(
tableMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
resolvedTables.addAll(
transformed.collect { case t: Table => t }.asJava
)
case other: Throwable => throw other
}
// If all flows to this particular destination are resolved, move to the destination
// node transformer
if (flowsTo(flow.destinationIdentifier).forall({ flowToDestination =>
resolvedFlowsMap.containsKey(flowToDestination.identifier)
})) {
// If multiple flows completed in parallel, ensure we resolve the destination only
// once by electing a leader via computeIfAbsent
var isCurrentThreadLeader = false
resolvedFlowDestinationsMap.computeIfAbsent(flow.destinationIdentifier, _ => {
isCurrentThreadLeader = true
// Set initial value as false as flow destination is not resolved yet.
false
})
if (isCurrentThreadLeader) {
if (tableMap.contains(flow.destinationIdentifier)) {
resolvedFlows.addAll(
transformed.collect { case f: ResolvedFlow => f }.asJava
)
} else if (viewMap.contains(flow.destinationIdentifier)) {
resolvedViews.addAll {
val transformed =
transformer(
tableMap(flow.destinationIdentifier),
viewMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
resolvedTables.addAll(
transformed.collect { case t: Table => t }.asJava
)
resolvedFlows.addAll(
transformed.collect { case f: ResolvedFlow => f }.asJava
)
} else if (viewMap.contains(flow.destinationIdentifier)) {
resolvedViews.addAll {
val transformed =
transformer(
viewMap(flow.destinationIdentifier),
flowsTo(flow.destinationIdentifier)
)
transformed.map(_.asInstanceOf[View]).asJava
}
} else if (sinkMap.contains(flow.destinationIdentifier)) {
resolvedSinks.addAll {
val transformed =
transformer(
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
)
require(
transformed.forall(_.isInstanceOf[Sink]),
"transformer must return a Seq[Sink]"
transformed.map(_.asInstanceOf[View]).asJava
}
} else if (sinkMap.contains(flow.destinationIdentifier)) {
resolvedSinks.addAll {
val transformed =
transformer(
sinkMap(flow.destinationIdentifier), flowsTo(flow.destinationIdentifier)
)
transformed.map(_.asInstanceOf[Sink]).asJava
}
} else {
throw new IllegalArgumentException(
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
s" in flow: ${flow.displayName} at transformDownNodes"
require(
transformed.forall(_.isInstanceOf[Sink]),
"transformer must return a Seq[Sink]"
)
transformed.map(_.asInstanceOf[Sink]).asJava
}
// Set flow destination as resolved now.
resolvedFlowDestinationsMap.computeIfPresent(
flow.destinationIdentifier,
(_, _) => {
// If there are any other node failures dependent on this destination, retry
// them
failedDependentFlows.computeIfPresent(
flow.destinationIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
true
}
} else {
throw new IllegalArgumentException(
s"Unsupported destination ${flow.destinationIdentifier.unquotedString}" +
s" in flow: ${flow.displayName} at transformDownNodes"
)
}
// Set flow destination as resolved now.
resolvedFlowDestinationsMap.computeIfPresent(
flow.destinationIdentifier,
(_, _) => {
// If there are any other node failures dependent on this destination, retry
// them
failedDependentFlows.computeIfPresent(
flow.destinationIdentifier,
(_, toRetryFlows) => {
toRetryFlows.foreach(toBeResolvedFlows.addFirst(_))
null
}
)
true
}
)
}
} catch {
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
}
)
} catch {
case ex: TransformNodeFailedException => failedFlowsQueue.add(ex.failedNode)
}
)
} else if (outstanding > 0) {
// Nothing could be scheduled (slots full, or the queue is drained) but tasks are still
// running: block until the next finishes instead of busy-spinning on Future.isDone. The
// outstanding > 0 guard is required, not redundant: the poll() drain above can take
// outstanding to 0 with an empty queue, and then there is nothing to wait for - the loop
// should just exit rather than block forever in take().
reap(completionService.take())
Comment on lines +292 to +298

@uros-b uros-b Jun 23, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: the else if (outstanding > 0) guard is effectively always true when reached. Control only reaches it when the if is false, i.e. outstanding >= batchSize || queue.isEmpty. Since batchSize >= 1, the first disjunct implies outstanding > 0; and if the queue is empty the loop invariant guarantees outstanding > 0. So it could just be a plain else?

@LuciferYang LuciferYang Jun 24, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good question, but the guard is needed - a plain else would deadlock the last iteration. The reasoning misses that the poll() drain loop runs after the loop-condition check and mutates outstanding. The invariant outstanding > 0 || queue.nonEmpty only holds at the top of the loop; by the time we reach this branch, the drain may have reaped the last in-flight tasks and taken outstanding to 0 while the queue is already empty. In that case the if is false (empty queue) and there is nothing left to wait for - take() would block forever. The outstanding > 0 guard lets the loop fall through so the next condition check (0 > 0 || empty) exits cleanly. I added a comment spelling this out.

}
}

Expand Down
58 changes: 58 additions & 0 deletions ...lines/src/test/scala/org/apache/spark/sql/pipelines/graph/ConnectValidPipelineSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -542,6 +542,64 @@ class ConnectValidPipelineSuite extends PipelineTest with SharedSparkSession {
)
}

test("resolution terminates and resolves all flows when flow count exceeds parallelism") {
val session = spark
import session.implicits._

// DataflowGraphTransformer caps in-flight resolutions at `parallelism` (10). With more
// independent flows than that, the slots fill and the scheduler blocks on a finished task
// (the `take()` branch) once `parallelism` tasks are outstanding - the path the busy-wait
// rewrite changes. Small graphs in the other suites never reach this regime. This asserts the
// outcome only (all flows resolve and the call returns), so it is deterministic and has no
// timing dependence; a regression that deadlocked would hang here until the suite times out.
val numFlows = 25
class P extends TestGraphRegistrationContext(spark) {
(0 until numFlows).foreach { i =>
registerPersistedView(s"v$i", query = dfFlowFunc(Seq(i).toDF("x")))
}
}
val p = new P().resolveToDataflowGraph()

assert(p.resolved, "all flows should resolve when their count exceeds parallelism")
(0 until numFlows).foreach { i =>
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"v$i")),
s"flow v$i was not resolved")
}
}

test("resolution re-queues retryable flows under load when consumers exceed parallelism") {
val session = spark
import session.implicits._

// A wide fan-out: many consumers reading from one source view, with the consumer count above
// `parallelism` (10), so slots fill and the loop blocks on take(). The consumers are
// registered (and therefore scheduled) before `src`, so the first batch resolves consumers
// whose `src` input is not yet available: each throws TransformNodeRetryableException and is
// parked as a dependent of `src`; once `src` resolves they are re-queued onto the deque and
// retried, re-driving the loop until every consumer resolves. This deterministically exercises
// the retryable re-queue path together with the blocking branch. Asserts only that everything
// resolves and the call returns - no timing assertions.
val numConsumers = 20
class P extends TestGraphRegistrationContext(spark) {
(0 until numConsumers).foreach { i =>
registerPersistedView(s"c$i", query = sqlFlowFunc(spark, "SELECT x FROM src"))
}
registerPersistedView("src", query = dfFlowFunc(Seq(1, 2, 3).toDF("x")))
}
val p = new P().resolveToDataflowGraph()

assert(p.resolved, "source and all consumers should resolve under load")
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier("src")),
"source flow was not resolved")
(0 until numConsumers).foreach { i =>
assert(
p.resolvedFlow.contains(fullyQualifiedIdentifier(s"c$i")),
s"consumer flow c$i was not resolved")
}
}

/** Verifies the [[DataflowGraph]] has the specified [[Flow]] with the specified schema. */
private def verifyFlowSchema(
pipeline: DataflowGraph,
Expand Down