Skip to content

[SPARK-56023][SS] Better load balance in LowLatencyMemoryStream #54848

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
eason-yuchen-liu wants to merge 5 commits into from
+30 −5
Closed

[SPARK-56023][SS] Better load balance in LowLatencyMemoryStream #54848

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
69061a7
Improve load balance in LowLatencyMemoryStream addData
eason-yuchen-liu Mar 17, 2026
cde4a5c
address comments
eason-yuchen-liu Mar 18, 2026
be1b3f4
address comment
eason-yuchen-liu Mar 18, 2026
6041b7e
compile
eason-yuchen-liu Mar 19, 2026
5ef5079
style
eason-yuchen-liu Mar 19, 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
12 changes: 8 additions & 4 deletions .../main/scala/org/apache/spark/sql/execution/streaming/sources/LowLatencyMemoryStream.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,16 +76,19 @@ class LowLatencyMemoryStream[A: Encoder](
private val records =
Seq.fill(numPartitions)(new ListBuffer[(UnsafeRow, Long)]())

@GuardedBy("this")
private var numRecords: Long = 0L

private val recordEndpoint = new LowLatencyMemoryStreamEndpoint(records, this)
@volatile private var endpointRef: RpcEndpointRef = _

override def addData(data: IterableOnce[A]): Offset = synchronized {
// Distribute data evenly among partition lists.
val timestamp = clock.getTimeMillis()
data.iterator.to(Seq).zipWithIndex.map {
case (item, index) =>
val partitionId = index % numPartitions
records(partitionId) += ((toRow(item).copy().asInstanceOf[UnsafeRow], timestamp))
data.iterator.foreach { item =>
val partitionId: Int = (numRecords % numPartitions).toInt
records(partitionId) += ((toRow(item).copy().asInstanceOf[UnsafeRow], timestamp))
numRecords += 1
}

// The new target offset is the offset where all records in all partitions have been processed.
Expand All @@ -102,6 +105,7 @@ class LowLatencyMemoryStream[A: Encoder](
val timestamp = clock.getTimeMillis()
data.iterator.foreach { item =>
records(partitionId) += ((toRow(item).copy().asInstanceOf[UnsafeRow], timestamp))
numRecords += 1
Copy link
Contributor

@HeartSaVioR HeartSaVioR Mar 18, 2026

Choose a reason for hiding this comment

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

Beyond the PR so this comment doesn't block this PR to merge.

Do we have a pattern of mix-up between writing to specific partition & writing without specifying partition? We probably need to be smarter if we want to keep the balance for that pattern, but I agree this is sorta over engineering, and I don't know we ever have that pattern.

}

// The new target offset is the offset where all records in all partitions have been processed.
Expand Down
23 changes: 22 additions & 1 deletion sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamRealTimeModeSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,11 @@ import scala.concurrent.duration.Duration

import org.scalatest.concurrent.PatienceConfiguration.Timeout

import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException}
import org.apache.spark.{SparkIllegalArgumentException, SparkIllegalStateException, TaskContext}
import org.apache.spark.sql.execution.streaming.{LowLatencyMemoryStream, RealTimeTrigger}
import org.apache.spark.sql.execution.streaming.runtime.MemoryStream
import org.apache.spark.sql.execution.streaming.sources.ContinuousMemorySink
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.internal.SQLConf

class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase {
Expand Down Expand Up @@ -204,6 +205,26 @@ class StreamRealTimeModeSuite extends StreamRealTimeModeSuiteBase {
}
)
}

test("LowLatencyMemoryStream load balance among all partitions") {
val numPartitions = 3
val inputData = LowLatencyMemoryStream[Int](numPartitions)

val getPartitionId = udf(() => TaskContext.getPartitionId())

val mapped = inputData.toDS().select($"value", getPartitionId()).as[(Int, Int)]

testStream(mapped, OutputMode.Update, Map.empty, new ContinuousMemorySink())(
StartStream(),
// 6 items round-robin across 3 partitions: item i goes to partition (i-1) % 3
AddData(inputData, 1, 2, 3),
AddData(inputData, 4),
AddData(inputData, 5),
AddData(inputData, 6),
CheckAnswerWithTimeout(10000, (1, 0), (2, 1), (3, 2), (4, 0), (5, 1), (6, 2)),
StopStream
)
}
}

class StreamRealTimeModeWithManualClockSuite extends StreamRealTimeModeManualClockSuiteBase {
Expand Down