Skip to content

[SPARK-24397][PYSPARK] Added TaskContext.getLocalProperty(key) in Python #21437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
tdas wants to merge 7 commits into from
Closed

[SPARK-24397][PYSPARK] Added TaskContext.getLocalProperty(key) in Python #21437

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
ac0270a
Added
tdas May 26, 2018
52610e3
Fixed test
tdas May 26, 2018
b9d8dd3
Fixed lint
tdas May 28, 2018
e197460
fixed
tdas May 29, 2018
9d95c12
fixed style
tdas May 29, 2018
2ea9cbc
Addressed comments
tdas May 29, 2018
6fef2f1
Removed keyword
tdas May 30, 2018
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
7 changes: 7 additions & 0 deletions python/pyspark/taskcontext.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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.

@BryanCutler BryanCutler May 29, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If it's missing it will result in a KeyError, maybe you want return self._localProperties.get(key) which returns None as the default? That seems better to me too, although you might want to add an optional default value.

@tdas tdas May 29, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right.

@tdas tdas May 30, 2018 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks @BryanCutler for catching this stupid stuff. not at all comfortable with python.

@HyukjinKwon HyukjinKwon May 30, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

get(key) would return None tho if it's missing.

@tdas tdas May 30, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yeah. I added the default=None to make it super obvious. i see no harm in making the code more intuitive.

@JoshRosen JoshRosen May 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The Java / Scala equivalents of this API return null for missing keys, so on the one hand returning None is kinda consistent with that.

On the other hand, consider a case where you want to specify an alternative in case a key is not set:

With this API, you might think of doing something like tc.getLocalProperty('key') or 'defaultValue', which potentially could be a problem in case a non-None key could have a False-y value. I suppose we're only dealing with strings here, though, and that'd only happen for empty strings. If we allowed non-strings to be returned here, though, then we'd have problems if we're returning values like 0. For that case, having a getLocalProperty('key', 'defaultValue') is a bit more useful.

@tdas tdas May 30, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am trying to match the Java TaskContext API, which only has a TaskContext.getLocalProperty(key)

@JoshRosen JoshRosen May 30, 2018

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FWIW we can always evolve what you have here towards the getLocalProperty(key, default=None) case; that evolution would maintain behavior and source compatibility. Therefore maybe it's fine to defer that until later if we're not sure whether that variant is useful.

@HyukjinKwon HyukjinKwon May 31, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense but +1 for leaving it out since I either don't know how commonly it will be used for now.

@BryanCutler BryanCutler May 31, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good to me since this matches the Scala API

"""
return self._localProperties[key]
9 changes: 9 additions & 0 deletions python/pyspark/tests.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -543,6 +543,15 @@ 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")

@ueshin ueshin May 29, 2018

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should cleanup the property after the test to avoid affecting to other tests?

@tdas tdas May 29, 2018

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

there isnt a clear way to clear local property in the spark context? what do you propose is the right approach here? set the local property as "None"?

rdd = self.sc.parallelize(range(1), 1)
prop1 = rdd.map(lambda x: TaskContext.get().getLocalProperty("testkey")).collect()[0]
self.assertEqual(prop1, "testkey")
prop2 = rdd.map(lambda x: TaskContext.get().getLocalProperty("otherkey")).collect()[0]
self.assertTrue(prop2 is None)


class RDDTests(ReusedPySparkTestCase):

Expand Down
6 changes: 6 additions & 0 deletions python/pyspark/worker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down