Added

apache · May 26, 2018 · ac0270a · ac0270a
1 parent 1b1528a
commit ac0270a
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 0 deletions.
diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
@@ -183,6 +183,13 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
         dataOut.writeInt(context.partitionId())
         dataOut.writeInt(context.attemptNumber())
         dataOut.writeLong(context.taskAttemptId())
+        val localProps = context.asInstanceOf[TaskContextImpl].getLocalProperties.asScala
+        dataOut.writeInt(localProps.size)
+        localProps.foreach { case (k, v) =>
+          PythonRDD.writeUTF(k, dataOut)
+          PythonRDD.writeUTF(v, dataOut)
+        }
+
         // sparkFilesDir
         PythonRDD.writeUTF(SparkFiles.getRootDirectory(), dataOut)
         // Python includes (*.zip and *.egg files)

diff --git a/python/pyspark/taskcontext.py b/python/pyspark/taskcontext.py
@@ -34,6 +34,7 @@ class TaskContext(object):
     _partitionId = None
     _stageId = None
     _taskAttemptId = None
+    _localProperties = None
 
     def __new__(cls):
         """Even if users construct TaskContext instead of using get, give them the singleton."""
@@ -88,3 +89,9 @@ def taskAttemptId(self):
         TaskAttemptID.
         """
         return self._taskAttemptId
+
+    def getLocalProperty(self, key):
+        """
+        Get a local property set upstream in the driver, or None if it is missing.
+        """
+        return self._localProperties[key]
diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py
@@ -543,6 +543,14 @@ def test_tc_on_driver(self):
         tc = TaskContext.get()
         self.assertTrue(tc is None)
 
+    def test_get_local_property(self):
+        """Verify that local properties set on the driver are available in TaskContext."""
+        self.sc.setLocalProperty("testkey", "testvalue")
+        rdd = self.sc.parallelize(range(1), 1)
+        prop1 = rdd1.map(lambda x: TaskContext.get().getLocalProperty("testkey")).collect()[0]
+        self.assertEqual(prop1, "testkey")
+        prop2 = rdd1.map(lambda x: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
+        self.assertTrue(prop2 is None)
 
 class RDDTests(ReusedPySparkTestCase):
 

diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
@@ -222,6 +222,12 @@ def main(infile, outfile):
         taskContext._partitionId = read_int(infile)
         taskContext._attemptNumber = read_int(infile)
         taskContext._taskAttemptId = read_long(infile)
+        taskContext._localProperties = dict()
+        for i in range(read_int(infile)):
+            k = utf8_deserializer.loads(infile)
+            v = utf8_deserializer.loads(infile)
+            taskContext._localProperties[k] = v
+
         shuffle.MemoryBytesSpilled = 0
         shuffle.DiskBytesSpilled = 0
         _accumulatorRegistry.clear()