Thanks a lot for your response. I'll give it a try and get back to you.

Not sure I understand that input_file_name() is an UDF - I didn't write it myself, it's imported from pyspark. I guess what you are saying is that it still is an UDF, please correct me if I'm wrong.