Accelerating Machine Learning with AWS Data Wrangler and AWS Glue 2.0

As machine learning (ML) becomes an integral part of modern applications, data preparation and feature engineering have emerged as key challenges. Data scientists often spend a significant portion of their time cleaning, transforming, and preprocessing raw data before it can be used to train ML models. This is where AWS Data Wrangler comes in.

AWS Data Wrangler is an open-source Python library that extends the power of Pandas to AWS, making it easy to perform data transformations on AWS services like S3, Athena, Glue, and Redshift. With its intuitive API and optimized performance, Data Wrangler has quickly become a go-to tool for data scientists and engineers building ML pipelines on AWS.

In this post, we‘ll explore how you can use AWS Data Wrangler in conjunction with AWS Glue 2.0 jobs to streamline your ML data preparation workflows. We‘ll discuss key benefits, walk through a practical example, and share best practices for productionizing Data Wrangler pipelines. Let‘s dive in!

Why Use AWS Data Wrangler for Machine Learning?

Before we look at how to use Data Wrangler with Glue, let‘s discuss some of the key benefits it offers for ML workflows:

1. Simplified feature engineering: Data Wrangler provides a wide range of built-in functions for common data transformations, such as one-hot encoding, string manipulations, date/time conversions, and more. This allows data scientists to quickly experiment with different feature engineering techniques without writing complex code.

2. Integration with ML frameworks: Data Wrangler seamlessly integrates with popular ML frameworks like scikit-learn, TensorFlow, and PyTorch. You can easily convert Data Wrangler DataFrames to NumPy arrays or TensorFlow Datasets, making it straightforward to feed preprocessed data into your ML models.

3. Scalability and performance: Data Wrangler is optimized for performance and can handle large-scale datasets with ease. It leverages Pandas-on-Spark to parallelize computations across a cluster of machines, enabling efficient processing of big data. Data Wrangler also integrates with Glue‘s DynamicFrame API, allowing you to take advantage of Glue‘s serverless Spark infrastructure.

To quantify the performance benefits of using Data Wrangler, let‘s look at some statistics from a benchmark comparing Data Wrangler to raw PySpark for common ETL tasks:

Source: AWS Data Wrangler benchmarks (link)

As we can see, Data Wrangler significantly outperforms raw PySpark for these ETL tasks, with speedups ranging from 2x to 4x. This is because Data Wrangler is optimized for AWS services and can automatically apply best practices for data layout and partitioning.

Example: Preprocessing Data for Machine Learning with Data Wrangler and Glue

Now that we understand the benefits of using Data Wrangler for ML, let‘s walk through an example of preprocessing data for training an ML model using Data Wrangler in a Glue job.

Suppose we have a dataset of customer reviews stored in S3 as JSON files. Our goal is to build a sentiment analysis model that can predict whether a given review is positive or negative. To prepare the data for training, we need to:

1. Read the JSON files from S3 into a DataFrame
2. Extract relevant features from the JSON data
3. Perform text preprocessing on the review text (lowercase, remove punctuation, etc.)
4. Convert the preprocessed text to numerical features using TF-IDF vectorization
5. Write the preprocessed features and labels to S3 in Parquet format

Here‘s how we can accomplish this using AWS Data Wrangler in a Glue 2.0 job:

import sys
from awsglue.utils import getResolvedOptions
from awsglue.context import GlueContext
from awsglue.job import Job
import awswrangler as wr

from sklearn.feature_extraction.text import TfidfVectorizer

args = getResolvedOptions(sys.argv, [‘JOB_NAME‘])
glueContext = GlueContext(SparkContext.getOrCreate())
job = Job(glueContext)
job.init(args[‘JOB_NAME‘], args)

# Read JSON data from S3
path = "s3://my-bucket/reviews/*.json"
df = wr.s3.read_json(path)

# Extract relevant features
df = df[[‘review_id‘, ‘review_body‘, ‘star_rating‘]]

# Perform text preprocessing
df[‘review_body‘] = df[‘review_body‘].str.lower()
df[‘review_body‘] = df[‘review_body‘].str.replace(r‘[^\w\s]‘, ‘‘)

# Convert star rating to sentiment label
df[‘sentiment‘] = df[‘star_rating‘].apply(lambda x: 1 if x >= 4 else 0)

# Convert preprocessed text to TF-IDF features
vectorizer = TfidfVectorizer(stop_words=‘english‘, max_features=1000)
tfidf_matrix = vectorizer.fit_transform(df[‘review_body‘])
features_df = pd.DataFrame(tfidf_matrix.toarray(), columns=vectorizer.get_feature_names())

# Combine features and labels into single DataFrame
preprocessed_df = pd.concat([df[[‘review_id‘, ‘sentiment‘]], features_df], axis=1)

# Write preprocessed data to S3 in Parquet format
wr.s3.to_parquet(
    df=preprocessed_df,
    path=‘s3://my-bucket/preprocessed_reviews/‘,
    dataset=True,
    partition_cols=[‘sentiment‘]
)

job.commit()

Let‘s break down the key steps:

1. We start by reading the JSON review data from S3 using wr.s3.read_json(). This automatically infers the schema and returns a DataFrame.

2. Next, we extract the relevant columns (review_id, review_body, and star_rating) and perform text preprocessing on the review_body column using Pandas string functions. We lowercase the text and remove punctuation.

3. We convert the star_rating to a binary sentiment label, where reviews with a rating >= 4 are considered positive (1) and reviews with a rating < 4 are considered negative (0).

4. We then use scikit-learn‘s TfidfVectorizer to convert the preprocessed text to numerical features. The vectorizer calculates the TF-IDF scores for each word in the corpus, which capture the importance of a word to a document. We limit the number of features to the top 1000 most important words.

5. Finally, we combine the sentiment labels and TF-IDF features into a single DataFrame and write the preprocessed data to S3 in Parquet format using wr.s3.to_parquet(). We partition the data by the sentiment label to enable efficient querying later.

With just a few lines of code, we were able to preprocess a large dataset of reviews and create a feature matrix ready for training an ML model. Data Wrangler abstracted away the complexities of working with JSON data, text preprocessing, and feature extraction, allowing us to focus on the higher-level ML workflow.

Productionizing Data Wrangler Pipelines

While the previous example showed how to use Data Wrangler in a one-off Glue job, in practice you‘ll want to productionize your Data Wrangler pipelines to run on a scheduled basis or in response to data updates. Here are some best practices to keep in mind:

1. Modularize your code: Break your Data Wrangler pipeline into modular functions or classes that encapsulate specific data transformations. This makes your code more reusable and maintainable.

2. Use Glue Workflows: Instead of running Glue jobs directly, use AWS Glue Workflows to orchestrate multiple jobs and create dependencies between them. This allows you to build complex pipelines that can handle data preprocessing, model training, and deployment.

3. Implement data quality checks: Add data quality checks to your Data Wrangler pipeline to ensure that the input data meets expected schemas and distributions. This can help catch data issues early and prevent downstream errors.

4. Monitor and log pipeline runs: Use AWS services like CloudWatch and X-Ray to monitor your Data Wrangler pipeline runs and log key metrics like data volume, processing time, and error rates. This can help you identify performance bottlenecks and troubleshoot issues.

5. Version control your code: Use a version control system like Git to track changes to your Data Wrangler code over time. This allows you to rollback to previous versions if needed and collaborate with other team members.

By following these best practices, you can build robust and scalable Data Wrangler pipelines that can reliably preprocess data for your ML workflows.

Case Study: Autodesk‘s ML Platform with Data Wrangler and Glue

To illustrate the impact of using Data Wrangler and Glue for ML workloads, let‘s look at a case study from Autodesk, a leading 3D design and engineering software company.

Autodesk built an ML platform on AWS to enable data scientists to easily develop and deploy ML models for a variety of use cases, such as product recommendation and customer churn prediction. A key component of this platform is a feature store that provides reusable, preprocessed features for training ML models.

To build the feature store, Autodesk used AWS Data Wrangler and AWS Glue to preprocess and transform raw data from various sources, including S3, Redshift, and DynamoDB. Data Wrangler allowed data scientists to easily experiment with different feature engineering techniques and create reusable feature sets that could be shared across teams.

By using Data Wrangler and Glue, Autodesk was able to:

• Reduce the time required to onboard new datasets from weeks to days
• Increase the productivity of data scientists by 10x by providing reusable, preprocessed features
• Scale the feature store to handle petabytes of data and serve features to hundreds of ML models

This case study demonstrates the tangible benefits of using Data Wrangler and Glue for ML workflows at enterprise scale. By abstracting away the complexities of data preprocessing and enabling feature reuse, these tools can significantly accelerate the development and deployment of ML models.

The Future of Data Preprocessing with AWS Data Wrangler

As we‘ve seen in this post, AWS Data Wrangler is a powerful tool for simplifying and accelerating data preprocessing for ML workflows on AWS. Its integration with Glue and other AWS services makes it easy to build end-to-end ML pipelines that can handle large-scale datasets.

Looking ahead, we can expect AWS to continue investing in Data Wrangler and expanding its capabilities. Some exciting developments on the roadmap include:

• Integration with AWS SageMaker Feature Store, allowing you to easily create, store, and share reusable feature sets for ML
• Support for more data sources and file formats, such as Apache Iceberg and Delta Lake
• Improved performance and scalability through tighter integration with Spark 3.x and AWS Glue 3.0
• Enhanced data profiling and quality checks, leveraging techniques like anomaly detection and schema inference

As an AI/ML expert, I‘m excited to see how these capabilities will further streamline and democratize the process of building ML applications on AWS. With tools like Data Wrangler, data scientists and engineers can focus on the high-level ML problem at hand, while still leveraging the power and scale of AWS infrastructure.

If you‘re building ML workflows on AWS, I highly recommend giving AWS Data Wrangler a try. You can find more examples, tutorials, and best practices in the official documentation and GitHub repository.

What are your experiences with using Data Wrangler or other data preprocessing tools for ML on AWS? Let me know in the comments below!