slight more refactor

HeartSaVioR · Sep 30, 2024 · d50a906 · d50a906
1 parent 4f14646
commit d50a906
Showing 1 changed file with 16 additions and 12 deletions.
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/StreamingJoinSuite.scala
@@ -1602,6 +1602,18 @@ class StreamingOuterJoinSuite extends StreamingJoinSuite {
         val joinedDf = data1.join(data2, Seq("key", "eventTime"), "leftOuter")
           .selectExpr("key", "CAST(eventTime AS long) AS eventTime")
 
+        def assertLeftRows(expected: Seq[Row]): AssertOnQuery = {
+          assertStateStoreRows(0L, "left", expected) { df =>
+            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
+          }
+        }
+
+        def assertRightRows(expected: Seq[Row]): AssertOnQuery = {
+          assertStateStoreRows(0L, "right", expected) { df =>
+            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
+          }
+        }
+
         testStream(joinedDf)(
           StartStream(checkpointLocation = checkpoint.getCanonicalPath),
           // batch 0
@@ -1611,12 +1623,8 @@ class StreamingOuterJoinSuite extends StreamingJoinSuite {
             (memoryStream2, Seq(("b", 2), ("c", 1)))
           ),
           CheckNewAnswer(("b", 2)),
-          assertStateStoreRows(0L, "left", Seq(Row("a", 1), Row("b", 2))) { df =>
-            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
-          },
-          assertStateStoreRows(0L, "right", Seq(Row("b", 2), Row("c", 1))) { df =>
-            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
-          },
+          assertLeftRows(Seq(Row("a", 1), Row("b", 2))),
+          assertRightRows(Seq(Row("b", 1), Row("c", 1))),
           // batch 1
           // WM: late record = 0, eviction = 2
           // Before Spark introduces multiple stateful operator, WM for late record was same as
@@ -1625,12 +1633,8 @@ class StreamingOuterJoinSuite extends StreamingJoinSuite {
           // batch 1. Before SPARK-49829, this wasn't producing unmatched row, and it is fixed.
           AddData(memoryStream1, ("d", 1)),
           CheckNewAnswer(("a", 1), ("d", 1)),
-          assertStateStoreRows(0L, "left", Seq()) { df =>
-            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
-          },
-          assertStateStoreRows(0L, "right", Seq()) { df =>
-            df.selectExpr("value.key", "CAST(value.eventTime AS long)")
-          }
+          assertLeftRows(Seq()),
+          assertRightRows(Seq()),
         )
       }
     }