[spark] branch branch-2.4 updated: [SPARK-27234][SS][PYTHON][BRANCH-2.4] Use InheritableThreadLocal for current epoch in EpochTracker (to support Python UDFs)

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dongjoon pushed a commit to branch branch-2.4
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-2.4 by this push:
     new 5a06584  [SPARK-27234][SS][PYTHON][BRANCH-2.4] Use InheritableThreadLocal for current epoch in EpochTracker (to support Python UDFs)
5a06584 is described below

commit 5a06584f8ea316484d138081d26eb4b10d02ac32
Author: HyukjinKwon <gu...@apache.org>
AuthorDate: Wed Aug 14 21:36:20 2019 -0700

    [SPARK-27234][SS][PYTHON][BRANCH-2.4] Use InheritableThreadLocal for current epoch in EpochTracker (to support Python UDFs)
    
    ## 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?
    
    Manually tested:
    
    ```python
    from pyspark.sql.functions import col, udf
    fooUDF = udf(lambda p: "foo")
    
    spark \
        .readStream \
        .format("rate") \
        .load()\
        .withColumn("foo", fooUDF(col("value")))\
        .writeStream\
        .format("console")\
        .trigger(continuous="1 second").start()
    ```
    
    Note that test was not ported because:
    
    1. `IntegratedUDFTestUtils` only exists in master.
    2. Missing SS testing utils in PySpark code base.
    3. Writing new test for branch-2.4 specifically might bring even more overhead due to mismatch against master.
    
    Closes #25457 from HyukjinKwon/SPARK-27234-backport.
    
    Authored-by: HyukjinKwon <gu...@apache.org>
    Signed-off-by: Dongjoon Hyun <dh...@apple.com>
---
 .../streaming/continuous/ContinuousCoalesceRDD.scala          |  5 +++--
 .../sql/execution/streaming/continuous/EpochTracker.scala     | 11 +++++++++--
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousCoalesceRDD.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousCoalesceRDD.scala
index aec756c..14046f6 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousCoalesceRDD.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/ContinuousCoalesceRDD.scala
@@ -110,8 +110,9 @@ class ContinuousCoalesceRDD(
               context.getLocalProperty(ContinuousExecution.START_EPOCH_KEY).toLong)
             while (!context.isInterrupted() && !context.isCompleted()) {
               writer.write(prev.compute(prevSplit, context).asInstanceOf[Iterator[UnsafeRow]])
-              // Note that current epoch is a non-inheritable thread local, so each writer thread
-              // can properly increment its own epoch without affecting the main task thread.
+              // Note that current epoch is a inheritable thread local but makes another instance,
+              // so each writer thread can properly increment its own epoch without affecting
+              // the main task thread.
               EpochTracker.incrementCurrentEpoch()
             }
           }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochTracker.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochTracker.scala
index bc0ae42..631ae48 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochTracker.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/continuous/EpochTracker.scala
@@ -26,8 +26,15 @@ import java.util.concurrent.atomic.AtomicLong
 object EpochTracker {
   // The current epoch. Note that this is a shared reference; ContinuousWriteRDD.compute() will
   // update the underlying AtomicLong as it finishes epochs. Other code should only read the value.
-  private val currentEpoch: ThreadLocal[AtomicLong] = new ThreadLocal[AtomicLong] {
-    override def initialValue() = new AtomicLong(-1)
+  private val currentEpoch: InheritableThreadLocal[AtomicLong] = {
+    new InheritableThreadLocal[AtomicLong] {
+      override protected def childValue(parent: AtomicLong): AtomicLong = {
+        // Note: make another instance so that changes in the parent epoch aren't reflected in
+        // those in the children threads. This is required at `ContinuousCoalesceRDD`.
+        new AtomicLong(parent.get)
+      }
+      override def initialValue() = new AtomicLong(-1)
+    }
   }
 
   /**


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