Github user HeartSaVioR commented on a diff in the pull request:
https://github.com/apache/spark/pull/21428#discussion_r191629272
--- Diff:
sql/core/src/test/scala/org/apache/spark/sql/streaming/continuous/shuffle/ContinuousShuffleSuite.scala
---
@@ -288,4 +267,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()
+ }
+
+ test("multiple writer partitions") {
--- End diff --
Would we want to have another test which covers out-of-order epoch between
writers (if that's valid case for us), or rely on the test in
ContinuousShuffleReadRDD?
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
