Spark Observation API Bug, reproduce

TestObservationWithForeach.scala

import org.apache.spark.sql.Observation
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions._
import org.scalatest.FunSuite

class TestObservationWithForeach extends FunSuite {

  test("observation") {
    implicit val spark: SparkSession = SparkSession.builder()
      .master("local")
      .getOrCreate()
    import spark.implicits._

    spark.listenerManager.register(new QueryExecutionListener {

      override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit =
        println(qe)

      override def onFailure(funcName: String, qe: QueryExecution, exception: Exception): Unit =
        println(qe)

    })

    val observation = Observation.apply("test")
    spark
      .range(200)
      .repartition(5)
      .observe(observation, count($"id"))
      .foreach(println(_))
//      .foreachPartition { it: Iterator[lang.Long] => it.foreach(println) }
//      .count()
    println(observation.get)
  }

}

uncomment count and test will pass, but with foreach and foreachPartition it will stuck.
Because these two actions do not trigger QueryExecutionListener .