AWS Glue: How to add a column with the source filename in the output?

Question

Does anyone know of a way to add the source filename as a column in a Glue job?

We created a flow where we crawled some files in S3 to create a schema. We then wrote a job that transforms the files to a new format, and the writes those files back to another S3 bucket as CSV, to be used by the rest of our pipeline. What we would like to do is get access to some sort of job meta properties so we can add a new column to the output file that contains the original filename.

I looked through the AWS documentation and the aws-glue-libs source, but didn't see anything that jumped out. Ideally there would be some way to get metadata from the awsglue.job package (we're using the python flavor).

I'm still learning Glue, so apologies if I'm using the wrong terminology. I tagged this with the spark tag as well, because I believe that's what Glue is using under the covers.

Solution 1

You can do it with spark in your etl job:

var df = glueContext.getCatalogSource(
  database = database,
  tableName = table,
  transformationContext = s"source-$database.$table"
).getDynamicFrame()
 .toDF()
 .withColumn("input_file_name", input_file_name())

glueContext.getSinkWithFormat(
  connectionType = "s3",
  options = JsonOptions(Map(
    "path" -> args("DST_S3_PATH")
  )),
  transformationContext = "",
  format = "parquet"
).writeDynamicFrame(DynamicFrame(df, glueContext))

Remember it works with getCatalogSource() API only and not with create_dynamic_frame_from_options()

Solution 2

With an AWS Glue Python auto-generated script, I've added the following lines:

from pyspark.sql.functions import input_file_name

## Add the input file name column
datasource1 = datasource0.toDF().withColumn("input_file_name", input_file_name())

## Convert DataFrame back to DynamicFrame
datasource2 = datasource0.fromDF(datasource1, glueContext, "datasource2")

Then, in the ApplyMapping or datasink portions of the code, you reference datasource2.

Solution 3

I am using an AWS Glue Python auto-generated script. I tried using the solution from JcMaco as this is exactly what I needed and it is a very simple solution to use input_file_name().

However, I could not get that to work, my column always came back empty aside from the header of the column, but I was able to get the name of the Glue job and use that as a constant in a new column and it serves the same purpose as input_file_name() in this particular use case for me.

If you look at the top left of your script, you will see where the args variable gets created. Use args to get access to JOB_NAME as shown below.

How I did it:

from pyspark.sql.functions import *

job_name = args['JOB_NAME'] # define new variable

(The JOB_NAME is passed in as a command line argument.)

Then, after the datasource0 definition in your script, use job_name along with the lit function:

applymapping1 = ApplyMapping.apply(frame = datasource0, mappings = […] , transformation_ctx = "applymapping1") 
applymapping2 = applymapping1.toDF().withColumn("job_name", lit(job_name))
applymapping3 = applymapping1.fromDF(applymapping2, glueContext, "applymapping3")

In the example above, you would change the assignment of your frame parameter in your datasink definition to applymapping3.

Comments

[Roberto]

Works great! I imported input_file_name with from pyspark.sql.functions import input_file_name.

[Hugo]

doesn't work for me