title: How to get the correct input_file_name in PySpark after a repartition
body: <p>I'm developing a PySpark job that read text files and write parquet files on a AWS S3 bucket.</p>
  
  <p>I need to add a column to each parquet line that describe the source file of the original text line.</p>
  
  <p>I'm using the <code>input_file_name</code> but i don't see the correct source file name in the output column. Seems that use the same filename for an entire tasks or for the wrong partition.</p>
  
  <p>I don't know if this is caused by the repartition that i need to do or because i'm reading an entire folder (prefix) in S3.</p>
  
  <p>Here an example of my code:</p>
  
  <pre><code>rdd = spark.sparkContext.textFile("s3a://source-bucket/prefix/")
  rdd = rdd.repartition(numPartitions=(120))
  rdd = rdd.map(lambda line: custom_regex_parser)
  
  df = spark.createDataFrame(rdd, schema=custom_source_schema)
  df = df..withColumn("source_file_name", input_file_name)
  df.write.mode('overwrite').parquet("s3://dest-bucket/prefix/")
  </code></pre>
  
  <p>In the S3 bucket i have lots of GZ files, of various dimension (from few KB to 50 MB compressed). I need to do a repartition to avoid random memory issues on some Spark tasks.</p>
  
  <p>Could be this repartition the problem with input_file_name?</p>
  
  <p>After that, do you have any suggestion to better spread data across executors homogeneously ? This is quite time consuming...</p>