HeartSaVioR commented on PR #60073:
URL: https://github.com/apache/spark/pull/60073#issuecomment-6631279692
Thanks for the change - the code change itself looks OK to me.
W.r.t testing, I'd rather see more detailed and controlled unit test(s)
where we can verify that ReusedExchangeExec is taken into account "and" the
watermark works. That detailed verification isn't something e2e golden file
test can provide; golden file test can only compare the final output, not the
executed plan or the state of each operator.
You can use temp view to replicate very similar code with Scala unit test,
or you could probably give a try with below code (Claude took a look at the
test code and converted with pure DataFrame API):
```
import org.apache.spark.sql.execution.exchange.ReusedExchangeExec
Seq(true, false).foreach { reuseEnabled =>
test("SPARK-60862: watermark propagates through reused exchange - " +
s"exchange reuse enabled: $reuseEnabled") {
withSQLConf(SQLConf.EXCHANGE_REUSE_ENABLED.key ->
reuseEnabled.toString) {
val inputData = MemoryStream[Int]
// Both union branches read the same aggregated DataFrame, hence
share the same
// EventTimeWatermark node. With exchange reuse enabled, the shuffle
under the second
// branch's stateful aggregation is replaced by ReusedExchangeExec.
val aggregated = inputData.toDF()
.withColumn("eventTime", timestamp_seconds($"value"))
.withWatermark("eventTime", "0 seconds")
.groupBy(window($"eventTime", "5 seconds").as("window"), ($"value"
% 2).as("key"))
.agg(count("*").as("count"))
.select($"window.start".cast("long").as("windowStart"), $"key",
$"count")
val stream = aggregated.select(lit("left").as("branch"), $"*")
.union(aggregated.select(lit("right").as("branch"), $"*"))
.as[(String, Long, Int, Long)]
testStream(stream)(
AddData(inputData, 0, 1, 2, 3, 4),
// watermark advances to 4s, which does not close [0, 5) yet
CheckNewAnswer(),
AddData(inputData, 10),
// no-data batch with watermark 10s closes [0, 5) in BOTH branches
CheckNewAnswer(
("left", 0L, 0, 3L), ("left", 0L, 1, 2L),
("right", 0L, 0, 3L), ("right", 0L, 1, 2L)),
AssertOnQuery { q =>
val plan = q.lastExecution.executedPlan
val reused = plan.collect { case r: ReusedExchangeExec => r }
assert(reused.nonEmpty === reuseEnabled,
s"Expected ReusedExchangeExec only when reuse is enabled.
Plan: $plan")
// Both stateful operators must see the propagated watermark.
Without the fix,
// the operator above ReusedExchangeExec gets Some(0).
val evictionWatermarks = plan.collect {
case s: StateStoreSaveExec => s.eventTimeWatermarkForEviction
}
assert(evictionWatermarks === Seq(Some(10000L), Some(10000L)))
true
},
// only [10, 15) with key 0 remains in each operator's state
assertNumStateRows(Seq(1, 1))
)
}
}
}
```
Maybe an useful negative test could be following: we create a single
streaming DataFrame and apply withWatermark two times (becomes two different
DataFrames) and join them. They should have different watermark nodes. Arguably
this is more related to reused exchange itself since it doesn't trigger reused
exchange, but since we are here to address the relationship between reused
exchange and watermark, it's still good to address while we are here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
