How to Preserve Directory Structure in PySpark ETL Job on AWS Glue

Question

I have been implementing a small ETL job using Pyspark. 
**I plan to deploy it to AWS Glue and will use an S3 bucket. to read and write my files instead of local file, once it is ready.**

This ETL job does the following:

- It reads JSON files from a directory structure (see the following directory structure)  using `create_dynamic_frame_from_options()` and gets its schema.
- Then, it flattens the schema using  `relationalize()`
- Finally, it writes files in parquet format in a flattened schema. **The limitation of the current version is it writes files to a directory without following the directory structure it reads from.**

I want to retrieve the directory structure and write parquet files according to that exact directory structure when I write  parquet file.

The local directory structure from which I have been reading data (in future it will be S3 folder instead of local directory) 
is in `year`, `month`, `day` and `hr` format. For instance, the structure is as follows:

```
   2024 (year) 
    |
    |--- 07 (month)
    |   |---01 (day)
    |   |    |- 12 (hr)
    |   |          |- file_1.json
    |   |          |_ file_2.json
    |   |          |
    |   |    |- 13
    |   |    |- 14
    |   |    |- ...
    |   |
    |   |--- 02
    |   |    |- 12
    |   |    |- 13
    |   |    |- 14
    |   |    |- ...
    |
    |-- 08 (month)
    |
```

```
import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from awsglue.transforms import Map

sc = SparkContext.getOrCreate()
glueContext = GlueContext(sc)
job = Job(glueContext)

connection_options = {
    "paths": ["file:///home/glue_user/workspace/jupyter_workspace/2024/"],
    "recurse": True
}

dynamic_frame = glueContext.create_dynamic_frame_from_options(
    connection_type="file",
    connection_options=connection_options,
    format = "json"
)

dynamic_frame.toDF().printSchema()

flattened_frames = dynamic_frame.relationalize("root", "/tmp/")
flattened_root_df = flattened_frames.select("root").toDF()
flattened_root_df.printSchema()

glueContext.write_dynamic_frame.from_options(
       frame = flattened_frames,
       connection_type = "file",
       connection_options = {"path": "file:///home/glue_user/workspace/jupyter_workspace/output/"},
       format = "parquet")
       
job.commit()
```

Alexander B · Answer

I understand that you have a requirement that you need to dynamically alter the partitions for each write so that it matches the directory structure of your source data. I suggest a method for allowing to pass the directory structure of the source data to your write job by including the directory structure as metadata columns in your json object, as well as some potential issues that will occurs with changing partitions in your destination.

You can effectively write the 'year', 'month', 'day' and 'hr' format using ‘partitionKeys’ in your connection options. [1][2]
“partitionKeys” is an optional key in the connection_options parameter in the write_dynamic_frame.from_options() function, which allows you to create partitions as S3 folders (directories) on write.

It's generally not advisable to change your partitioning structure frequently, as predicate pushdown and tools leveraging partitions for query optimization require consistent partition formats to function efficiently. Changing the structure could impact query performance and break compatibility with existing processes.

However if your use case requires creating unique file structures for each of your write operations, one could still leverage the partitionKeys to achieve this dynamic creation of unique folder structures for each write so it always matches the existing directory structure of the source data. For example in your script prior to write your code could check for the existence of columns you use for partitioning, and add the list of the existing columns in your actual data that are in that list passing it to the partitionKeys value. In this way you would allow for ‘passing’ the available partitions to the write operation by excluding certain columns to be used for partitions from the source file. Your json objects being flattened would just need to have those partition columns updated when publish to s3 for ingestion by your Spark job.

Some example sudo code illustrating the method above:
```
available_partition_columns = [col for col in ["year", "month", "day", "hour"] if col in dynamic_frame.toDF().columns]

glue_context.write_dynamic_frame.from_options(
    frame = dynamic_frame,
    connection_type = "s3",    
    connection_options = {…, "partitionKeys": available_partition_columns, …},
    format = "parquet")
```

This approach allows for flexibility in your partition structure so that it will change on each write based on the columns you pass it to create partition folders by. However, be cautious with this method, as it may lead to inconsistent partitioning across your dataset, potentially complicating queries and integrating with upstream tools which otherwise would be able to limit the data you need to query leveraging partitions. Many tools that query S3 charge by the amount of data scanned which you can limit leveraging partitions, such as Athena: https://calculator.aws/#/createCalculator/Athena

References
* [1] https://docs.aws.amazon.com/glue/latest/dg/aws-glue-programming-etl-partitions.html#aws-glue-programming-etl-partitions-writing
* [2] https://aws.amazon.com/blogs/big-data/work-with-partitioned-data-in-aws-glue/
* [3] https://calculator.aws/#/createCalculator/Athena

How to Preserve Directory Structure in PySpark ETL Job on AWS Glue

Add your answer

Relevant content