Optimize Python ETL by extending Pandas with AWS Data Wrangler

April 2024: This post was reviewed for accuracy.

Developing extract, transform, and load (ETL) data pipelines is one of the most time-consuming steps to keep data lakes, data warehouses, and databases up to date and ready to provide business insights. You can categorize these pipelines into distributed and non-distributed, and the choice of one or the other depends on the amount of data you need to process.

Apache Spark is widely used to build distributed pipelines, whereas Pandas is preferred for lightweight, non-distributed pipelines. With the second use case in mind, the AWS Professional Service team created AWS Data Wrangler, aiming to fill the integration gap between Pandas and several AWS services, such as Amazon Simple Storage Service (Amazon S3), Amazon Redshift, AWS Glue, Amazon Athena, Amazon Aurora, Amazon QuickSight, and Amazon CloudWatch Log Insights.

AWS Data Wrangler is an open-source Python library that enables you to focus on the transformation step of ETL by using familiar Pandas transformation commands and relying on abstracted functions to handle the extraction and load steps.

You can use AWS Data Wrangler in different environments on AWS and on premises (for more information, see Install). This post focuses on data preparation for a data science project on Jupyter. By the end of this walkthrough, you will be able to set up AWS Data Wrangler on your Amazon SageMaker notebook.

Use case overview

In the following walkthrough, you use data stored in the NOAA public S3 bucket. For more information, see NOAA Global Historical Climatology Network Daily. The objective is to convert 10 CSV files (approximately 240 MB total) to a partitioned Parquet dataset, store its related metadata into the AWS Glue Data Catalog, and query the data using Athena to create a data analysis.

Configuring Amazon S3

Your first step is to create an S3 bucket to store the Parquet dataset.

1. On the Amazon S3 console, choose Create bucket.
2. For Bucket name, enter a name for your bucket.

3. Choose Create.

Creating a new database in the Data Catalog

The Data Catalog is an Apache Hive-compatible managed metadata storage that lets you store, annotate, and share metadata on AWS.

For this use case, you use it to store the metadata associated with your Parquet dataset. The Data Catalog is integrated with many analytics services, including Athena, Amazon Redshift Spectrum, and Amazon EMR (Apache Spark, Apache Hive, and Presto).

1. On the AWS Glue console, choose Databases.
2. Choose Add database.
3. For Database name, enter awswrangler_test.
4. Choose Create.

Launching an Amazon SageMaker notebook

An Amazon SageMaker notebook is a managed instance running the Jupyter Notebook app. For this use case, you use it to write and run your code.

1. On the Amazon SageMaker console, choose Notebook instance.
2. Choose Create a notebook instance.
3. For Notebook instance name, enter a name.
4. For IAM role, choose an existing AWS Identity and Access Management (IAM) role or create a role that allows you to run Amazon SageMaker and grants access to Amazon S3, Athena, and AWS Glue for the related resources.

5. Wait for the notebook status to show as InService.
6. Choose Open Jupyter from the notebook instance you created.

Exploring the data

This section walks you through several notebook paragraphs to expose how to install and use AWS Data Wrangler.

1. On Jupyter console, under New, choose conda_python3.
2. To install AWS Data Wrangler, enter the following code:

   !pip install awswrangler

3. To avoid dependency conflicts, restart the notebook kernel by choosing kernel -> Restart.
4. Import the library given the usual alias wr:

   import awswrangler as wr

5. List all files in the NOAA public bucket from the decade of 1880:

   wr.s3.list_objects("s3://noaa-ghcn-pds/csv/188")

The following screenshot shows the output.

6. Load the whole decade (10 files) into a Pandas DataFrame using the Amazon S3 prefix s3://noaa-ghcn-pds/csv/188:

   col_names = ["id", "dt", "element", "value", "m_flag", "q_flag", "s_flag", "obs_time"]
   
   df = wr.s3.read_csv(
       path="s3://noaa-ghcn-pds/csv/188",
       names=col_names,
       parse_dates=["dt", "obs_time"]  # Hint to parse these columns as date instead of strings
   )

   The following screenshot shows the output.

7. Create a new column extracting the year from the dt column (the new column is useful for creating partitions in the Parquet dataset):

   df["year"] = df["dt"].dt.year

The following screenshot shows the output.

8. Store the Pandas DataFrame in the S3 bucket you created in the beginning of this post (replace the [BUCKET] placeholder in the code with your bucket name):

   wr.s3.to_parquet(
       df=df,
       path="s3://[BUCKET]/noaa/",
       dataset=True,
       database="awswrangler_test",
       table="noaa",
       partition_cols=["year"]
   );

The preceding code creates the table noaa in the awswrangler_test database in the Data Catalog.

9. After processing this, you can confirm the Parquet files exist in Amazon S3 and the table noaa is in AWS Glue data catalog. See the following code:

   wr.s3.list_objects("s3://[BUCKET]/noaa/")
   wr.catalog.table(database="awswrangler_test", table="noaa")

The following screenshot shows the output.

10. Run a SQL query from Athena that filters only the US maximum temperature measurements of the last 3 years (1887–1889) and receive the result as a Pandas DataFrame:

    sql = """
    SELECT
        dt,
        (value / 10.0) AS temperature  -- Converting tenths of degrees C to regular degrees C
    FROM noaa
    WHERE year BETWEEN 1887 AND 1889  -- Only last 3 years (PARTITION filter)
    AND substr(id, 1, 2)='US'  -- Only U.S. stations
    AND element='TMAX'  -- Only Maximum temperature elements
    AND q_flag is NULL  -- Only HIGH quality measurement
    """
    
    df = wr.athena.read_sql_query(sql, database="awswrangler_test")

The following screenshot shows the output.

The following two queries illustrate how you can visualize the data.

11. To plot the average maximum temperature measured in the tracked station, enter the following code:

    %matplotlib inline
    df.groupby("dt").mean().plot();

The following screenshot shows the output.

12. To plot a moving average of the previous metric with a 30-day window, enter the following code:

    %matplotlib inline
    df.groupby("dt").mean().rolling(window=30, center=True).mean().plot();

The following screenshot shows the output.

Cleaning up

To avoid incurring future charges, delete the resources from the following services:

1. AWS Glue database
   • On the AWS Glue console, choose the database you created.
   • From the Actions drop-down menu, choose Delete database.
   • Choose Delete.
2. Amazon SageMaker notebook
   • On the Amazon SageMaker console, choose the notebook instance you created.
   • From the Actions drop-down menu, choose Stop.
   • When the status shows as Stopped, choose Database.
   • Choose Delete.
3. S3 bucket
   • On the Amazon S3 console, choose Buckets.
   • Choose the bucket you created.
   • Choose Empty and enter your bucket name.
   • Choose Confirm.
   • Choose Delete and enter your bucket name.
   • Choose Delete bucket.
4. IAM Role
   • On the IAM console, choose Roles.
   • Choose the role you attached to Amazon SageMaker.
   • Choose Delete role.
   • Choose Yes.

Conclusion

Installing AWS Data Wrangler is a breeze. With a single command, you can connect ETL tasks to multiple data sources and different data services. The library is a work in progress, with new features and enhancements added regularly. For more tutorials, see the GitHub repo.

About the Authors

Satoshi Kuramitsu is a Solutions Architect in AWS. His favorite AWS services are AWS Glue, Amazon Kinesis, and Amazon S3.

Igor Tavares is a Data & Machine Learning Engineer in the AWS Professional Services team and the original creator of AWS Data Wrangler.

Audit History

Last reviewed and updated in April 2024 by Priyanka Chaudhary | Sr. Solutions Architect