[SPARK-50492][SS] Fix java.util.NoSuchElementException when event tim…

…e column is dropped after dropDuplicatesWithinWatermark ### What changes were proposed in this pull request? Update `DeduplicateWithinWatermark` references to include all attributes that could be the watermarking column. ### Why are the changes needed? Fix `java.util.NoSuchElementException` due to ColumnPruning. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Added unit test ### Was this patch authored or co-authored using generative AI tooling? No. Closes apache#49065 from liviazhu-db/liviazhu-db/dedup-watermark-fix. Authored-by: Livia Zhu <[email protected]> Signed-off-by: Jungtaek Lim <[email protected]> (cherry picked from commit 851f5f2) Signed-off-by: Jungtaek Lim <[email protected]>
changgyoopark-db · Dec 6, 2024 · d01f34f · d01f34f
1 parent 86e29e9
commit d01f34f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/...st/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/...st/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
@@ -1944,6 +1944,9 @@ case class Deduplicate(
 }
 
 case class DeduplicateWithinWatermark(keys: Seq[Attribute], child: LogicalPlan) extends UnaryNode {
+  // Ensure that references include event time columns so they are not pruned away.
+  override def references: AttributeSet = AttributeSet(keys) ++
+    AttributeSet(child.output.filter(_.metadata.contains(EventTimeWatermark.delayKey)))
   override def maxRows: Option[Long] = child.maxRows
   override def output: Seq[Attribute] = child.output
   final override val nodePatterns: Seq[TreePattern] = Seq(DISTINCT_LIKE)

diff --git a/...est/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala b/...est/scala/org/apache/spark/sql/streaming/StreamingDeduplicationWithinWatermarkSuite.scala
@@ -220,4 +220,18 @@ class StreamingDeduplicationWithinWatermarkSuite extends StateStoreMetricsTest {
       )
     }
   }
+
+  test("SPARK-50492: drop event time column after dropDuplicatesWithinWatermark") {
+    val inputData = MemoryStream[(Int, Int)]
+    val result = inputData.toDS()
+      .withColumn("first", timestamp_seconds($"_1"))
+      .withWatermark("first", "10 seconds")
+      .dropDuplicatesWithinWatermark("_2")
+      .select("_2")
+
+    testStream(result, Append)(
+      AddData(inputData, (1, 2)),
+      CheckAnswer(2)
+    )
+  }
 }