A bit of progress cleaning up the stream stream join examples: Open: @holden needs to update STB to handle the strongly typed Datasets that are not just valid type aliases to DataFrame.

Author: holdenk

core/src/main/scala/com/high-performance-spark-examples/streaming/structuredstreaming/RateSourceStressExample.scala

@@ -1,12 +1,11 @@
 package com.highperformancespark.examples.structuredstreaming
 
-// tag::stream_stream_join_basic_both_side_watermark[]
 // Stream-stream join with watermark on both sides
 // State can be cleaned up
 
-import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
-import org.apache.spark.sql.streaming.Trigger
+import org.apache.spark.sql.streaming._
 
 object StreamStreamJoinBothSideWatermark {
   def main(args: Array[String]): Unit = {
@@ -15,33 +14,51 @@ object StreamStreamJoinBothSideWatermark {
       .appName("StreamStreamJoinBothSideWatermark")
       .master("local[2]")
       .getOrCreate()
-    import spark.implicits._
+  }
 
+  def run(spark: SparkSession): Unit = {
     val left = spark.readStream
       .format("memory")
       .load()
-      .withWatermark("timestamp", "10 minutes")
+
     val right = spark.readStream
       .format("memory")
       .load()
-      .withWatermark("timestamp", "10 minutes")
+
+    val query = streamStreamJoin(spark, left, right)
+    query.awaitTermination()
+  }
+
+  def streamStreamJoinDF(spark: SparkSession, stream1: DataFrame, stream2: DataFrame): Dataset[Row] = {
+    // Note the watermarks don't need to be the same, by default Spark will pick the min.
+    // tag::stream_stream_join_basic_both_side_watermark[]
+    val left = stream1.withWatermark("timestamp", "10 minutes")
+    val right = stream2.withWatermark("timestamp", "5 minutes")
 
     val joined = left.join(
       right,
       expr(
-        "left.timestamp >= right.timestamp - interval 5 minutes AND left.timestamp <= right.timestamp + interval 5 minutes AND left.key = right.key"
+        "left.timestamp >= right.timestamp - interval 5 minutes " +
+         " AND left.timestamp <= right.timestamp + interval 5 minutes " +
+         " AND left.key = right.key"
       )
     )
+    // end::stream_stream_join_basic_both_side_watermark[]
+    joined
+  }
 
-    val query = joined.writeStream
+  def streamStreamJoin(spark: SparkSession, stream1: DataFrame, stream2: DataFrame): StreamingQuery = {
+    val joined = streamStreamJoinDF(spark, stream1, stream2)
+    // tag::ex_with_checkpoin_at_writet[]
+    val writer = joined.writeStream
       .outputMode("append")
       .format("console")
       .option(
         "checkpointLocation",
         "./tmp/checkpoints/stream_stream_join_both_side_watermark"
       )
-      .start()
-    query.awaitTermination()
+    // end::ex_with_checkpoin_at_writet[]
+    val query = writer.start()
+    query
   }
 }
-// end::stream_stream_join_basic_both_side_watermark[]

core/src/main/scala/com/high-performance-spark-examples/streaming/structuredstreaming/StreamStreamJoinBothSideWatermark.scala

@@ -1,60 +1,130 @@
 package com.highperformancespark.examples.structuredstreaming
 
-// tag::stream_stream_join_basic_both_side_watermark_test[]
-// Test for stream-stream join with watermark on both sides
-// Verifies bounded state and correct join results
+import java.sql.Timestamp
+import java.nio.file.Files
 
-import org.scalatest.funsuite.AnyFunSuite
-import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.streaming.Trigger
+import org.apache.spark.sql._
 import org.apache.spark.sql.functions._
-import java.sql.Timestamp
+import org.apache.spark.sql.streaming._
+import org.apache.spark.sql.execution.streaming.MemoryStream
+import org.scalatest.funsuite.AnyFunSuite
+
+// spark-testing-base
+import com.holdenkarau.spark.testing.DatasetSuiteBase
 
-class StreamStreamJoinBothSideWatermarkSuite extends AnyFunSuite {
-  test("join with both-side watermark yields bounded state and correct results") {
-  val spark = SparkSession.builder()
-      .master("local[2]")
-      .appName("StreamStreamJoinBothSideWatermarkSuite")
-      .getOrCreate()
-    import spark.implicits._
-
-    import org.apache.spark.sql.execution.streaming.MemoryStream
-    val now = System.currentTimeMillis()
-    val leftStream = MemoryStream[(Timestamp, String)](1, spark.sqlContext)
-    val rightStream = MemoryStream[(Timestamp, String)](2, spark.sqlContext)
-    val leftRows = Seq(
-      (new Timestamp(now - 1000 * 60 * 5), "k1"), // within window
-      (new Timestamp(now - 1000 * 60 * 20), "k2") // late, beyond watermark
-    )
-    val rightRows = Seq(
-      (new Timestamp(now - 1000 * 60 * 5), "k1"), // within window
-      (new Timestamp(now - 1000 * 60 * 20), "k2") // late, beyond watermark
-    )
-    leftStream.addData(leftRows: _*)
-    rightStream.addData(rightRows: _*)
-    val leftDF = leftStream.toDF().toDF("timestamp", "key").withWatermark("timestamp", "10 minutes")
-    val rightDF = rightStream.toDF().toDF("timestamp", "key").withWatermark("timestamp", "10 minutes")
-
-    val joined = leftDF.join(
-      rightDF,
-      leftDF("key") === rightDF("key") &&
-        leftDF("timestamp") >= rightDF("timestamp") - expr("interval 5 minutes") &&
-        leftDF("timestamp") <= rightDF("timestamp") + expr("interval 5 minutes")
-    )
-
-    val query = joined.writeStream
+final case class Ev(key: String, timestamp: Timestamp, v: Int)
+
+class StreamStreamJoinBothSideWatermarkSTBSpec
+    extends AnyFunSuite
+    with DatasetSuiteBase {
+
+  import spark.implicits._
+
+  private def ts(mins: Long): Timestamp =
+    new Timestamp(mins * 60L * 1000L) // epoch + minutes
+
+  private def joinedDF(leftIn: DataFrame, rightIn: DataFrame): DataFrame = {
+    StreamStreamJoinBothSideWatermark.streamStreamJoinDF(spark, leftIn, rightIn)
+  }
+
+  test("joins rows with same key within ±5 minutes") {
+    val leftMem  = MemoryStream[Ev](1, spark.sqlContext)
+    val rightMem = MemoryStream[Ev](2, spark.sqlContext)
+
+    val outName = "stb_out_basic"
+    val q = joinedDF(leftMem.toDF(), rightMem.toDF())
+      .writeStream
+      .format("memory")
+      .queryName(outName)
       .outputMode("append")
+      .option("checkpointLocation", Files.createTempDirectory("chk-basic").toString)
+      .start()
+
+    // Left @ 10, Right @ 12 -> within window and same key
+    leftMem.addData(Ev("A", ts(10), 1))
+    rightMem.addData(Ev("A", ts(12), 2))
+    q.processAllAvailable()
+
+    // Select a stable set of columns to compare
+    val actual = spark.table(outName)
+      .selectExpr("left.key as key", "left.timestamp as lt", "right.timestamp as rt")
+      .as[(String, Timestamp, Timestamp)]
+
+    val expected = Seq(("A", ts(10), ts(12))).toDS()
+
+    assertDataFrameEquals(actual, expected)
+
+    q.stop()
+  }
+
+  test("does not join when outside tolerance or key mismatch") {
+    val leftMem  = MemoryStream[Ev](3, spark.sqlContext)
+    val rightMem = MemoryStream[Ev](4, spark.sqlContext)
+
+    val outName = "stb_out_filtering"
+    val q = joinedDF(leftMem.toDF(), rightMem.toDF())
+      .writeStream
+      .format("memory")
+      .queryName(outName)
+      .outputMode("append")
+      .option("checkpointLocation", Files.createTempDirectory("chk-filter").toString)
+      .start()
+
+    // Outside ±5 minutes (0 vs 7 -> 7 minutes apart)
+    leftMem.addData(Ev("A", ts(0), 1))
+    rightMem.addData(Ev("A", ts(7), 2))
+    q.processAllAvailable()
+    assert(spark.table(outName).isEmpty)
+
+    // Within time but different keys
+    rightMem.addData(Ev("B", ts(2), 9))
+    q.processAllAvailable()
+    assert(spark.table(outName).isEmpty)
+
+    q.stop()
+  }
+
+  test("late data are dropped after both watermarks advance") {
+    val leftMem  = MemoryStream[Ev](5, spark.sqlContext)
+    val rightMem = MemoryStream[Ev](6, spark.sqlContext)
+
+    val outName = "stb_out_late"
+    val q = joinedDF(leftMem.toDF(), rightMem.toDF())
+      .writeStream
       .format("memory")
-      .queryName("stream_stream_join_both_side_watermark")
-      .trigger(Trigger.Once())
-      .option("checkpointLocation", "./tmp/checkpoints/stream_stream_join_both_side_watermark_test")
+      .queryName(outName)
+      .outputMode("append")
+      .option("checkpointLocation", Files.createTempDirectory("chk-late").toString)
       .start()
-    query.processAllAvailable()
-    query.awaitTermination()
 
-    val result = spark.sql("select key from stream_stream_join_both_side_watermark").collect().map(_.getString(0)).toSet
-    assert(result == Set("k1"), "Only non-late key should join")
-    spark.stop()
+    // 1) Valid pair near t ~ 10..12
+    leftMem.addData(Ev("A", ts(10), 1))
+    rightMem.addData(Ev("A", ts(12), 2))
+    q.processAllAvailable()
+    assert(spark.table(outName).count() == 1)
+
+    // 2) Advance BOTH watermarks far ahead:
+    //    left WM delay 10m -> add t=100 -> WM ~ 90
+    //    right WM delay 5m -> add t=100 -> WM ~ 95
+    leftMem.addData(Ev("A", ts(100), 3))
+    rightMem.addData(Ev("A", ts(100), 4))
+    q.processAllAvailable()
+
+    // 3) Inject events that would have joined in the past (t=20..22)
+    //    but are now far older than both watermarks -> should be dropped.
+    leftMem.addData(Ev("A", ts(20), 5))
+    rightMem.addData(Ev("A", ts(22), 6))
+    q.processAllAvailable()
+
+    // Still only the first result
+    assert(spark.table(outName).count() == 1)
+
+    // Optional sanity: state metrics shouldn't balloon
+    Option(q.lastProgress).foreach { p =>
+      assert(p.stateOperators != null && p.stateOperators.nonEmpty)
+      assert(p.stateOperators.head.numRowsTotal >= 0)
+    }
+
+    q.stop()
   }
 }
-// end::stream_stream_join_basic_both_side_watermark_test[]