cloud-fan commented on a change in pull request #29589:
URL: https://github.com/apache/spark/pull/29589#discussion_r481642185



##########
File path: 
sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
##########
@@ -1342,6 +1345,50 @@ abstract class DynamicPartitionPruningSuiteBase
       }
     }
   }
+
+  test("SPARK-32748: propagate local properties to dynamic pruning thread") {
+    def checkPropertyValueByUdfResult(propKey: String, propValue: String): 
Unit = {
+      spark.sparkContext.setLocalProperty(propKey, propValue)
+      val df = sql(
+        s"""
+           |SELECT compare_property_value(f.date_id, '$propKey', '$propValue') 
as col
+           |FROM fact_sk f
+           |INNER JOIN dim_store s
+           |ON f.store_id = s.store_id AND s.country = 'NL'
+          """.stripMargin)
+
+      checkPartitionPruningPredicate(df, false, true)
+      assert(df.collect().forall(_.toSeq == Seq(true)))
+    }
+
+    try {
+      SQLConf.get.setConf(StaticSQLConf.SUBQUERY_MAX_THREAD_THRESHOLD, 1)
+
+      withSQLConf(
+        SQLConf.DYNAMIC_PARTITION_PRUNING_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
+
+        spark.udf.register(
+          "compare_property_value",
+          (input: Int, propKey: String, propValue: String) =>
+            TaskContext.get().getLocalProperty(propKey) == propValue
+        )
+        val propKey = "spark.sql.subquery.broadcast.prop.key"
+
+        // set local property and assert
+        val propValue1 = UUID.randomUUID().toString()
+        checkPropertyValueByUdfResult(propKey, propValue1)
+
+        // change local property and re-assert
+        val propValue2 = UUID.randomUUID().toString()
+        checkPropertyValueByUdfResult(propKey, propValue2)
+
+        spark.sessionState.catalog.dropTempFunction("compare_property_value", 
false)

Review comment:
       shall we put it in `finally`




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]



---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to