[SPARK-44897][SQL] Propagating local properties to subquery broadcast exec

[ChenMichael]

What changes were proposed in this pull request?

https://issues.apache.org/jira/browse/SPARK-32748 previously proposed propagating these local properties to the subquery broadcast exec threads but was then reverted since it was said that local properties would already be propagated to the broadcast threads.
I believe this is not always true. In the scenario where a separate BroadcastExchangeExec is the first to compute the broadcast, this is fine. However, in the scenario where the SubqueryBroadcastExec is the first to compute the broadcast, then the local properties that are propagated to the broadcast threads would not have been propagated correctly. This is because the local properties from the subquery broadcast exec were not propagated to its Future thread.
It is difficult to write a unit test that reproduces this behavior because usually BroadcastExchangeExec is the first computing the broadcast variable. However, by adding a Thread.sleep(10) to SubqueryBroadcastExec.doPrepare after relationFuture is initialized, the added test will consistently fail.

Why are the changes needed?

Local properties are not propagated correctly to SubqueryBroadcastExec

Does this PR introduce any user-facing change?

No

How was this patch tested?

Following test can reproduce the bug and test the solution by adding sleep to SubqueryBroadcastExec.doPrepare

protected override def doPrepare(): Unit = {
    relationFuture
    Thread.sleep(10)
}

    withSQLConf(StaticSQLConf.BROADCAST_EXCHANGE_MAX_THREAD_THRESHOLD.key -> "1") {
      withTable("a", "b") {
        val confKey = "spark.sql.y"
        val confValue1 = UUID.randomUUID().toString()
        val confValue2 = UUID.randomUUID().toString()
        Seq((confValue1, "1")).toDF("key", "value")
          .write
          .format("parquet")
          .partitionBy("key")
          .mode("overwrite")
          .saveAsTable("a")
        val df1 = spark.table("a")

        def generateBroadcastDataFrame(confKey: String, confValue: String): Dataset[String] = {
          val df = spark.range(1).mapPartitions { _ =>
            Iterator(TaskContext.get.getLocalProperty(confKey))
          }.filter($"value".contains(confValue)).as("c")
          df.hint("broadcast")
        }

        // set local property and assert
        val df2 = generateBroadcastDataFrame(confKey, confValue1)
        spark.sparkContext.setLocalProperty(confKey, confValue1)
        val checkDF = df1.join(df2).where($"a.key" === $"c.value").select($"a.key", $"c.value")
        val checks = checkDF.collect()
        assert(checks.forall(_.toSeq == Seq(confValue1, confValue1)))

        // change local property and re-assert
        Seq((confValue2, "1")).toDF("key", "value")
          .write
          .format("parquet")
          .partitionBy("key")
          .mode("overwrite")
          .saveAsTable("b")
        val df3 = spark.table("b")
        val df4 = generateBroadcastDataFrame(confKey, confValue2)
        spark.sparkContext.setLocalProperty(confKey, confValue2)
        val checks2DF = df3.join(df4).where($"b.key" === $"c.value").select($"b.key", $"c.value")
        val checks2 = checks2DF.collect()
        assert(checks2.forall(_.toSeq == Seq(confValue2, confValue2)))
        assert(checks2.nonEmpty)
      }
    }
  }

Was this patch authored or co-authored using generative AI tooling?

No

[ChenMichael]

cc @cloud-fan, @maropu
Is there some mechanism that is guaranteeing the separate BroadcastExchangeExec will be computing the relationFuture instead of the child of the SubqueryBroadcastExec?

[cloud-fan]

if this test can't reproduce the bug directly, we can put it in the PR descriptions with instructions about the extra change to reproduce the bug (add sleep).

[ChenMichael]

Ok. Does it make sense to leave this test in the code still or do you think putting it in the description with instructions is enough

[dongjoon-hyun]

Please remove the test case from the code, @ChenMichael . PR description (with your reproducible instruction is enough).

[ChenMichael]

👍 Done

[cloud-fan]

thanks, merging to master/3.5!