From 0d26b3aa55f9cc75096b0e2b309f64fe3270b9a5 Mon Sep 17 00:00:00 2001
From: Shixiong Zhu <shixiong@databricks.com>
Date: Wed, 2 Aug 2017 14:02:13 -0700
Subject: [PATCH] [SPARK-21546][SS] dropDuplicates should ignore watermark when
it's not a key
## What changes were proposed in this pull request?
When the watermark is not a column of `dropDuplicates`, right now it will crash. This PR fixed this issue.
## How was this patch tested?
The new unit test.
Author: Shixiong Zhu <shixiong@databricks.com>
Closes #18822 from zsxwing/SPARK-21546.
---
.../sql/execution/streaming/statefulOperators.scala | 9 +++++++--
.../spark/sql/streaming/DeduplicateSuite.scala | 13 +++++++++++++
2 files changed, 20 insertions(+), 2 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
index 6addab69f1..e46356392c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/statefulOperators.scala
@@ -156,8 +156,13 @@ trait WatermarkSupport extends UnaryExecNode {
}
/** Predicate based on keys that matches data older than the watermark */
- lazy val watermarkPredicateForKeys: Option[Predicate] =
- watermarkExpression.map(newPredicate(_, keyExpressions))
+ lazy val watermarkPredicateForKeys: Option[Predicate] = watermarkExpression.flatMap { e =>
+ if (keyExpressions.exists(_.metadata.contains(EventTimeWatermark.delayKey))) {
+ Some(newPredicate(e, keyExpressions))
+ } else {
+ None
+ }
+ }
/** Predicate based on the child output that matches data older than the watermark. */
lazy val watermarkPredicateForData: Option[Predicate] =
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/streaming/DeduplicateSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/streaming/DeduplicateSuite.scala
index a15c2cff93..e858b7d999 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/streaming/DeduplicateSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/streaming/DeduplicateSuite.scala
@@ -268,4 +268,17 @@ class DeduplicateSuite extends StateStoreMetricsTest with BeforeAndAfterAll {
CheckLastBatch(7)
)
}
+
+ test("SPARK-21546: dropDuplicates should ignore watermark when it's not a key") {
+ val input = MemoryStream[(Int, Int)]
+ val df = input.toDS.toDF("id", "time")
+ .withColumn("time", $"time".cast("timestamp"))
+ .withWatermark("time", "1 second")
+ .dropDuplicates("id")
+ .select($"id", $"time".cast("long"))
+ testStream(df)(
+ AddData(input, 1 -> 1, 1 -> 2, 2 -> 2),
+ CheckLastBatch(1 -> 1, 2 -> 2)
+ )
+ }
}
--
GitLab