Github user zsxwing commented on a diff in the pull request:
https://github.com/apache/spark/pull/21428#discussion_r194914332
--- Diff:
sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/shuffle/ContinuousShuffleSuite.scala
---
@@ -288,4 +264,153 @@ class ContinuousShuffleReadSuite extends StreamTest {
val thirdEpoch = rdd.compute(rdd.partitions(0),
ctx).map(_.getUTF8String(0).toString).toSet
assert(thirdEpoch == Set("writer1-row1", "writer2-row0"))
}
+
+ test("one epoch") {
+ val reader = new ContinuousShuffleReadRDD(sparkContext, numPartitions
= 1)
+ val writer = new RPCContinuousShuffleWriter(
+ 0, new HashPartitioner(1), Array(readRDDEndpoint(reader)))
+
+ writer.write(Iterator(1, 2, 3))
+
+ assert(readEpoch(reader) == Seq(1, 2, 3))
+ }
+
+ test("multiple epochs") {
+ val reader = new ContinuousShuffleReadRDD(sparkContext, numPartitions
= 1)
+ val writer = new RPCContinuousShuffleWriter(
+ 0, new HashPartitioner(1), Array(readRDDEndpoint(reader)))
+
+ writer.write(Iterator(1, 2, 3))
+ writer.write(Iterator(4, 5, 6))
+
+ assert(readEpoch(reader) == Seq(1, 2, 3))
+ assert(readEpoch(reader) == Seq(4, 5, 6))
+ }
+
+ test("empty epochs") {
+ val reader = new ContinuousShuffleReadRDD(sparkContext, numPartitions
= 1)
+ val writer = new RPCContinuousShuffleWriter(
+ 0, new HashPartitioner(1), Array(readRDDEndpoint(reader)))
+
+ writer.write(Iterator())
+ writer.write(Iterator(1, 2))
+ writer.write(Iterator())
+ writer.write(Iterator())
+ writer.write(Iterator(3, 4))
+ writer.write(Iterator())
+
+ assert(readEpoch(reader) == Seq())
+ assert(readEpoch(reader) == Seq(1, 2))
+ assert(readEpoch(reader) == Seq())
+ assert(readEpoch(reader) == Seq())
+ assert(readEpoch(reader) == Seq(3, 4))
+ assert(readEpoch(reader) == Seq())
+ }
+
+ test("blocks waiting for writer") {
+ val reader = new ContinuousShuffleReadRDD(sparkContext, numPartitions
= 1)
+ val writer = new RPCContinuousShuffleWriter(
+ 0, new HashPartitioner(1), Array(readRDDEndpoint(reader)))
+
+ val readerEpoch = reader.compute(reader.partitions(0), ctx)
+
+ val readRowThread = new Thread {
+ override def run(): Unit = {
+ assert(readerEpoch.toSeq.map(_.getInt(0)) == Seq(1))
+ }
+ }
+ readRowThread.start()
+
+ eventually(timeout(streamingTimeout)) {
+ assert(readRowThread.getState == Thread.State.TIMED_WAITING)
+ }
+
+ // Once we write the epoch the thread should stop waiting and succeed.
+ writer.write(Iterator(1))
+ readRowThread.join()
--- End diff --
nit: it's better to add a timeout here, such as
`readRowThread.join(streamingTimeout.toMillis)`. Without a timeout, if there is
a bug causing this hang, we will need to wait until the jenkins build timeout,
which is much longer.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
