Bug: foreachBatch can cause IllegalStateException if used in a wrong way
Status: Open | Priority: Minor | Resolution: 
Created: 09/Jul/24 00:52
Affects Version/s: 3.5.1

When foreachBatch is used for a stateful query, it could cause IllegalStateException if the function doesn't iterate the passed batch fully.

The following code demonstrates the issue:
{code:scala}
import org.apache.spark.sql.DataFrame
spark.readStream
  .format("rate")
  .option("rowsPerSecond", "10")
  .load()
  .dropDuplicates()
  .writeStream
  .foreachBatch((batch: DataFrame, id: Long) => batch.limit(1).show())
  .start()
{code}
Error:
{noformat}
24/07/09 00:21:03 ERROR Executor: Exception in task 1.0 in stage 17.0 (TID 239)
java.lang.IllegalStateException: Error reading delta file file:/tmp/temporary-db172114-81ec-490b-a667-93047609fcf1/state/0/2/2.delta of HDFSStateStoreProvider[id = (op=0,part=2),dir = file:/tmp/temporary-db172114-81ec-490b-a667-93047609fcf1/state/0/2]: file:/tmp/temporary-db172114-81ec-490b-a667-93047609fcf1/state/0/2/2.delta does not exist
         at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.updateFromDeltaFile(HDFSBackedStateStoreProvider.scala:461)
         at org.apache.spark.sql.execution.streaming.state.HDFSBackedStateStoreProvider.$anonfun$loadMap$4(HDFSBackedStateStoreProvider.scala:417)
{noformat}
This is because every record must be processed to compute states correctly. Specifically, if the output isn't fully exhausted, StateStore.commit() is not called, and StateStore.abort() is called eventually by the TaskCompletionListener set in mapPartitionsWithStateStore.

There are also some related but minor issues:
 # If the function doesn't iterate the batch at all, state store is not updated at all.
 # If the function iterates the batch multiple times, stateful operators run multiple times for the same data. It can be an issue if the user is using arbitrary stateful operators which are non-deterministic.

This is an example for the second point:
{code:scala}
import org.apache.spark.sql._
import org.apache.spark.sql.streaming._

case class Data(timestamp: java.sql.Timestamp, value: Long)

def func(key: Long, values: Iterator[Data], state: GroupState[Int]): Iterator[String] = {
    val counter = new scala.util.Random().nextInt()
    println(s"counter=$counter")
    val len = values.toSeq.length
    if (state.exists) {
        state.update(state.get + len)
    } else {
        state.update(len)
    }
    Seq(s"key=$key,count=${state.get},counter=$counter").toIterator
}

def outputDF(df: Dataset[String], df_id: Long) = {
    println(df_id)
    println(df.collect().toSeq)
    println(df.collect().toSeq) // state store is overwritten here
}

val q = spark.readStream.format("rate").option("rowsPerSecond", "10").load().as[Data]
    .groupByKey(_.value)
    .flatMapGroupsWithState(OutputMode.Update, GroupStateTimeout.NoTimeout)(func)
    .writeStream
    .foreachBatch((ds: Dataset[String], id: Long) => outputDF(ds, id))
    .outputMode("update")
    .start
{code}
It would be helpful if this is explained in [the documentation for foreachBatch|https://spark.apache.org/docs/3.5.1/structured-streaming-programming-guide.html#foreachbatch].