You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by as...@apache.org on 2025/09/16 22:04:32 UTC

(spark) branch branch-3.5 updated: [SPARK-53560][SS][SQL] Crash looping when retrying uncommitted batch in Kafka source and AvailableNow trigger

This is an automated email from the ASF dual-hosted git repository.

ashrigondekar pushed a commit to branch branch-3.5
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-3.5 by this push:
     new a9e81fb277d1 [SPARK-53560][SS][SQL] Crash looping when retrying uncommitted batch in Kafka source and AvailableNow trigger
a9e81fb277d1 is described below

commit a9e81fb277d133a381cc022a3d6136620725fc1e
Author: Yuchen Liu <17...@users.noreply.github.com>
AuthorDate: Tue Sep 16 15:04:17 2025 -0700

    [SPARK-53560][SS][SQL] Crash looping when retrying uncommitted batch in Kafka source and AvailableNow trigger
    
    ### What changes were proposed in this pull request?
    
    This PR solves the issue by not fetching the latest topic partitions in function `prepareForTriggerAvailableNow`. Instead, it fetches in `latestOffset`, because Spark will not call `latestOffset` when retrying the uncommitted batch.
    
    ### Why are the changes needed?
    
    There is a crash loop when using Kafka source with AvailableNow. It will be triggered deterministically when the query fails / terminates in the middle of the run, and the user changes the Kafka topic partition, then restarts the query. This is because the topic partitions in uncommitted batch and the latest topic partitions aren't the same, but trigger.AvailableNow requires it anyway.
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Created a new unit test.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #52326 from eason-yuchen-liu/kafkaAvailableNow3.5.
    
    Authored-by: Yuchen Liu <17...@users.noreply.github.com>
    Signed-off-by: Anish Shrigondekar <an...@databricks.com>
---
 .../spark/sql/kafka010/KafkaMicroBatchStream.scala | 15 +++++---
 .../apache/spark/sql/kafka010/KafkaSource.scala    | 13 +++++--
 .../sql/kafka010/KafkaMicroBatchSourceSuite.scala  | 40 +++++++++++++++++++++-
 .../streaming/SupportsTriggerAvailableNow.java     |  7 ++++
 4 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
index 24244c6de63b..473fa5b1523e 100644
--- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
+++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchStream.scala
@@ -89,6 +89,8 @@ private[kafka010] class KafkaMicroBatchStream(
 
   private var allDataForTriggerAvailableNow: PartitionOffsetMap = _
 
+  private var isTriggerAvailableNow: Boolean = false
+
   /**
    * Lazily initialize `initialPartitionOffsets` to make sure that `KafkaConsumer.poll` is only
    * called in StreamExecutionThread. Otherwise, interrupting a thread while running
@@ -124,8 +126,14 @@ private[kafka010] class KafkaMicroBatchStream(
     val startPartitionOffsets = start.asInstanceOf[KafkaSourceOffset].partitionToOffsets
 
     // Use the pre-fetched list of partition offsets when Trigger.AvailableNow is enabled.
-    latestPartitionOffsets = if (allDataForTriggerAvailableNow != null) {
-      allDataForTriggerAvailableNow
+    latestPartitionOffsets = if (isTriggerAvailableNow) {
+      if (allDataForTriggerAvailableNow != null) {
+        allDataForTriggerAvailableNow
+      } else {
+        allDataForTriggerAvailableNow =
+          kafkaOffsetReader.fetchLatestOffsets(Some(startPartitionOffsets))
+        allDataForTriggerAvailableNow
+      }
     } else {
       kafkaOffsetReader.fetchLatestOffsets(Some(startPartitionOffsets))
     }
@@ -357,8 +365,7 @@ private[kafka010] class KafkaMicroBatchStream(
   }
 
   override def prepareForTriggerAvailableNow(): Unit = {
-    allDataForTriggerAvailableNow = kafkaOffsetReader.fetchLatestOffsets(
-      Some(getOrCreateInitialPartitionOffsets()))
+    isTriggerAvailableNow = true
   }
 }
 
diff --git a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
index f5d4abb569a3..ab94b3964ee5 100644
--- a/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
+++ b/connector/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaSource.scala
@@ -108,6 +108,8 @@ private[kafka010] class KafkaSource(
 
   private var allDataForTriggerAvailableNow: PartitionOffsetMap = _
 
+  private var isTriggerAvailableNow = false
+
   /**
    * Lazily initialize `initialPartitionOffsets` to make sure that `KafkaConsumer.poll` is only
    * called in StreamExecutionThread. Otherwise, interrupting a thread while running
@@ -171,8 +173,13 @@ private[kafka010] class KafkaSource(
     val currentOffsets = currentPartitionOffsets.orElse(Some(initialPartitionOffsets))
 
     // Use the pre-fetched list of partition offsets when Trigger.AvailableNow is enabled.
-    val latest = if (allDataForTriggerAvailableNow != null) {
-      allDataForTriggerAvailableNow
+    val latest = if (isTriggerAvailableNow) {
+      if (allDataForTriggerAvailableNow != null) {
+        allDataForTriggerAvailableNow
+      } else {
+        allDataForTriggerAvailableNow = kafkaReader.fetchLatestOffsets(currentOffsets)
+        allDataForTriggerAvailableNow
+      }
     } else {
       kafkaReader.fetchLatestOffsets(currentOffsets)
     }
@@ -399,7 +406,7 @@ private[kafka010] class KafkaSource(
   }
 
   override def prepareForTriggerAvailableNow(): Unit = {
-    allDataForTriggerAvailableNow = kafkaReader.fetchLatestOffsets(Some(initialPartitionOffsets))
+    isTriggerAvailableNow = true
   }
 }
 
diff --git a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
index d63b9805e553..6493b8e29f98 100644
--- a/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
+++ b/connector/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/KafkaMicroBatchSourceSuite.scala
@@ -36,7 +36,7 @@ import org.scalatest.concurrent.PatienceConfiguration.Timeout
 import org.scalatest.matchers.should._
 import org.scalatest.time.SpanSugar._
 
-import org.apache.spark.TestUtils
+import org.apache.spark.{SparkException, TestUtils}
 import org.apache.spark.sql.{Dataset, ForeachWriter, Row, SparkSession}
 import org.apache.spark.sql.catalyst.util.CaseInsensitiveMap
 import org.apache.spark.sql.connector.read.streaming.SparkDataStream
@@ -2056,6 +2056,44 @@ abstract class KafkaSourceSuiteBase extends KafkaSourceTest {
       "subscribePattern" -> s"$topicPrefix-.*")
   }
 
+  test("SPARK-53560: no crash looping during uncommitted batch retry in AvailableNow trigger") {
+    val topic = newTopic()
+    testUtils.createTopic(topic, partitions = 1)
+    testUtils.sendMessages(topic, (1 to 7).map(_.toString).toArray, Some(0))
+    def udfFailOn7(x: Int): Int = {
+      if (x == 7) throw new RuntimeException("error for 7")
+      x
+    }
+    val kafka =
+      spark.readStream.format("kafka")
+        .option("kafka.bootstrap.servers", testUtils.brokerAddress)
+        .option("subscribe", topic)
+        .option("startingOffsets", "earliest")
+        .load()
+        .select(expr("CAST(CAST(value AS STRING) AS INT)").as("value"))
+        .as[Int]
+        .map(udfFailOn7)
+
+    withTempDir { dir =>
+      testStream(kafka)(
+        StartStream(Trigger.AvailableNow, checkpointLocation = dir.getAbsolutePath),
+        ExpectFailure[SparkException] { e =>
+          assert(e.getMessage.contains("error for 7"))
+        },
+        AssertOnQuery { q =>
+          testUtils.addPartitions(topic, 2)
+          !q.isActive
+        },
+        StartStream(Trigger.AvailableNow, checkpointLocation = dir.getAbsolutePath),
+        // Getting this error means the query has passed the planning stage, so
+        // verifyEndOffsetForTriggerAvailableNow succeeds.
+        ExpectFailure[SparkException] { e =>
+          assert(e.getMessage.contains("error for 7"))
+        }
+      )
+    }
+  }
+
   private def testFromSpecificTimestampsWithNoMatchingStartingOffset(
       topic: String,
       options: (String, String)*): Unit = {
diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsTriggerAvailableNow.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsTriggerAvailableNow.java
index 47662dc97cc1..268fa577b29e 100644
--- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsTriggerAvailableNow.java
+++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/read/streaming/SupportsTriggerAvailableNow.java
@@ -36,6 +36,13 @@ public interface SupportsTriggerAvailableNow extends SupportsAdmissionControl {
    * the query). The source will behave as if there is no new data coming in after the target
    * offset, i.e., the source will not return an offset higher than the target offset when
    * {@link #latestOffset(Offset, ReadLimit) latestOffset} is called.
+   * <p>
+   * Note that there is an exception on the first uncommitted batch after a restart, where the end
+   * offset is not derived from the current latest offset. Sources need to take special
+   * considerations if wanting to assert such relation. One possible way is to have an internal
+   * flag in the source to indicate whether it is Trigger.AvailableNow, set the flag in this method,
+   * and record the target offset in the first call of
+   * {@link #latestOffset(Offset, ReadLimit) latestOffset}.
    */
   void prepareForTriggerAvailableNow();
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org