[GitHub] spark pull request #18355: [SPARK-21145][SS] Added StateStoreProviderId with...

[zsxwing <gi...@git.apache.org>]

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 infrastructure@apache.org or file a JIRA ticket
with INFRA.
---

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