I came across a similar issue. 

Please detail how you are executing the python script. Are you calling it from the web terminal? or from a notebook?

Note: If you are calling it from the web terminal - your spark session wont be passed. You could create a local variable and pass it in if youd like. I have never gotten to that point yet tho