Github user zsxwing commented on a diff in the pull request:

    https://github.com/apache/spark/pull/18355#discussion_r123357731
  
    --- Diff: 
sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/state/StateStoreSuite.scala
 ---
    @@ -408,12 +413,60 @@ class StateStoreSuite extends 
StateStoreSuiteBase[HDFSBackedStateStoreProvider]
         assert(numDeltaFiles === 3)
       }
     
    +  test("SPARK-21145: Restarted queries create new provider instances") {
    +    try {
    +      val checkpointLocation = Utils.createTempDir().getAbsoluteFile
    +      val spark = SparkSession.builder().master("local[2]").getOrCreate()
    +      SparkSession.setActiveSession(spark)
    +      implicit val sqlContext = spark.sqlContext
    +      spark.conf.set("spark.sql.shuffle.partitions", "1")
    +      import spark.implicits._
    +      val inputData = MemoryStream[Int]
    +
    +      def runQueryAndGetLoadedProviders(): Seq[StateStoreProvider] = {
    +        val aggregated = inputData.toDF().groupBy("value").agg(count("*"))
    +        // stateful query
    +        val query = aggregated.writeStream
    +          .format("memory")
    +          .outputMode("complete")
    +          .queryName("query")
    +          .option("checkpointLocation", checkpointLocation.toString)
    +          .start()
    +        inputData.addData(1, 2, 3)
    +        query.processAllAvailable()
    +        require(query.lastProgress != null) // at least one batch 
processed after start
    +        val loadedProvidersMethod =
    +          PrivateMethod[mutable.HashMap[StateStoreProviderId, 
StateStoreProvider]]('loadedProviders)
    +        val loadedProvidersMap = StateStore invokePrivate 
loadedProvidersMethod()
    +        val loadedProviders = loadedProvidersMap.synchronized { 
loadedProvidersMap.values.toSeq }
    +        query.stop()
    +        loadedProviders
    +      }
    +
    +      val loadedProvidersAfterRun1 = runQueryAndGetLoadedProviders()
    +      require(loadedProvidersAfterRun1.length === 1)
    +
    +      val loadedProvidersAfterRun2 = runQueryAndGetLoadedProviders()
    +      assert(loadedProvidersAfterRun2.length === 2)   // two providers 
loaded for 2 runs
    +
    +      // Both providers should have the same StateStoreId, but the should 
be different objects
    +      assert(loadedProvidersAfterRun2(0).stateStoreId === 
loadedProvidersAfterRun2(1).stateStoreId)
    +      assert(loadedProvidersAfterRun2(0).hashCode !== 
loadedProvidersAfterRun2(1).hashCode)
    --- End diff --
    
    nit:  `assert(loadedProvidersAfterRun2(0) ne loadedProvidersAfterRun2(1))`



---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org

Reply via email to