[SPARK-27234][SS][PYTHON] Use InheritableThreadLocal for current epoch in EpochTracker (to support Python UDFs)

[HyukjinKwon]

What changes were proposed in this pull request?

This PR proposes to use InheritableThreadLocal instead of ThreadLocal for current epoch in EpochTracker. Python UDF needs threads to write out to and read it from Python processes and when there are new threads, previously set epoch is lost.

After this PR, Python UDFs can be used at Structured Streaming with the continuous mode.

How was this patch tested?

The test cases were written on the top of #24945.
Unit tests were added.

Manual tests.

[HyukjinKwon]

cc @WeichenXu123, @mengxr, @tdas, @zsxwing, @HeartSaVioR

[HyukjinKwon]

There's similar case that uses InheritableThreadLocal, input_file_name(). cc @cloud-fan as well.

[mengxr]

@HyukjinKwon How does SCALAR_ITER Pandas UDF work with continuous processing after this PR? Do we only run initialization code once per task?

[HyukjinKwon]

It works after this PR as below:

from pyspark.sql.functions import col, pandas_udf, PandasUDFType

@pandas_udf("int", PandasUDFType.SCALAR_ITER)
def the_udf(iterator):
    for col1_batch in iterator:
        yield col1_batch

spark \
    .readStream \
    .format("rate") \
    .load() \
    .withColumn("foo", the_udf(col("value"))) \
    .writeStream \
    .format("console") \
    .trigger(continuous="5 second").start() 

Before:

...
Caused by: java.util.NoSuchElementException: None.get
	at scala.None$.get(Option.scala:366)
	at scala.None$.get(Option.scala:364)
	at org.apache.spark.sql.execution.streaming.continuous.ContinuousQueuedDataReader.next(ContinuousQueuedDataReader.scala:116)
	at org.apache.spark.sql.execution.streaming.continuous.ContinuousDataSourceRDD$$anon$1.getNext(ContinuousDataSourceRDD.scala:93)
	at org.apache.spark.sql.execution.streaming.continuous.ContinuousDataSourceRDD$$anon$1.getNext(ContinuousDataSourceRDD.scala:91)
	at org.apache.spark.util.NextIterator.hasNext(NextIterator.scala:73)
	at org.apache.spark.sql.catalyst.expressions.GeneratedClass$GeneratedIteratorForCodegenStage1.processNext(Unknown Source)
	at org.apache.spark.sql.execution.BufferedRowIterator.hasNext(BufferedRowIterator.java:43)
	at org.apache.spark.sql.execution.WholeStageCodegenExec$$anon$1.hasNext(WholeStageCodegenExec.scala:701)
	at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:458)
...

After:

...
-------------------------------------------
Batch: 3
-------------------------------------------
+--------------------+-----+---+
|           timestamp|value|foo|
+--------------------+-----+---+
|2019-06-25 09:05:...|    9|  9|
|2019-06-25 09:05:...|    2|  2|
|2019-06-25 09:05:...|   10| 10|
|2019-06-25 09:05:...|   11| 11|
|2019-06-25 09:05:...|    5|  5|
|2019-06-25 09:05:...|    3|  3|
|2019-06-25 09:05:...|    1|  1|
|2019-06-25 09:05:...|    6|  6|
|2019-06-25 09:05:...|    4|  4|
|2019-06-25 09:05:...|    8|  8|
|2019-06-25 09:05:...|    7|  7|
|2019-06-25 09:05:...|    0|  0|
+--------------------+-----+---+
...

Because each epoch couldn't be referred in writer thread (to Python process). Each UDF will be executed each execution per each epoch.

[HyukjinKwon]

#24945 is merged. Let me rebase.

[zsxwing]

How do we handle other thread local variables that are not InheritableThreadLocal, such as org.apache.spark.TaskContext.get?

[HyukjinKwon]

In case of org.apache.spark.TaskContext, looks it's manually set for Python's writer thread:

spark/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala

Line 203 in 6d441dc

TaskContext.setTaskContext(context)

and the information are (de)serialized into Python worker, TaskContext instance is constructed and used in Python worker later.

I think it works too but I thought it's better to isolate this logic out of Python. Python runners are in core and this code is in SQL FWIW. We should move the codes around to mimic this approach.

[HyukjinKwon]

gentle ping .. :-) ..

[SparkQA]

Test build #108000 has finished for PR 24946 at commit b103acc.

• This patch passes all tests.
• This patch merges cleanly.
• This patch adds no public classes.

[zsxwing]

LGTM

[HyukjinKwon]

Thanks for review and approval, @zsxwing.

Just FYI, there's similar fix and discussion going on at #24958.

[HyukjinKwon]

Merged to master.

Thanks all.

[aurorazl]

@HyukjinKwon Was problem fixed in spark 3.0? What can i do to fix it in spark2.4.3?

[HyukjinKwon]

This fix will be included in Apache Spark 3.0. I think you should upgrade it later when this is released.

[dongjoon-hyun]

In the dev mailing list, this issue is discussed for 2.4.4. I'll follow the decision from @HyukjinKwon and @zsxwing .

[HyukjinKwon]

Am fine with backporting but @zsxwing WDYT?

[zsxwing]

I'm fine with backporting this small fix.

[dongjoon-hyun]

Thank you, @zsxwing and @HyukjinKwon .
Could you make a backporting PR against branch-2.4, @HyukjinKwon ?

[HyukjinKwon]

Yup. BTW IntegratedUDFTestUtils doesn't exist in master branch. I will have to manually test and backport.