Addressed comments

apache · May 29, 2018 · 2ea9cbc · 2ea9cbc
1 parent 9d95c12
commit 2ea9cbc
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/python/pyspark/taskcontext.py b/python/pyspark/taskcontext.py
@@ -94,4 +94,4 @@ def getLocalProperty(self, key):
         """
         Get a local property set upstream in the driver, or None if it is missing.
         """
-        return self._localProperties[key]
+        return self._localProperties.get(key, default=None)
diff --git a/python/pyspark/tests.py b/python/pyspark/tests.py
@@ -545,15 +545,17 @@ def test_tc_on_driver(self):
 
     def test_get_local_property(self):
         """Verify that local properties set on the driver are available in TaskContext."""
+        key = "testkey"
+        value = "testvalue"
+        self.sc.setLocalProperty(key, value)
         try:
-            self.sc.setLocalProperty("testkey", "testvalue")
             rdd = self.sc.parallelize(range(1), 1)
-            prop1 = rdd.map(lambda x: TaskContext.get().getLocalProperty("testkey")).collect()[0]
-            self.assertEqual(prop1, "testvalue")
+            prop1 = rdd.map(lambda x: TaskContext.get().getLocalProperty(key)).collect()[0]
+            self.assertEqual(prop1, value)
             prop2 = rdd.map(lambda x: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
             self.assertTrue(prop2 is None)
         finally:
-            self.sc.setLocalProperty("testkey", None)
+            self.sc.setLocalProperty(key, None)
 
 
 class RDDTests(ReusedPySparkTestCase):