How to read data from AWS S3 and Athena in pandas with column validation

This is a step by step tutorial on reading data from AWS S3 and Athena into a pandas DataFrame and doing column validation to assess the quality of the data.

Step 1: Install the python dependencies needed

We’ll need two python libraries for this:

pip install awswrangler
pip install pandas_schema

Step 2: Add your AWS credentials

In order to be able to reach your AWS account we will have to setup the credentials. We should have a file ~/.aws/credentials with the structure:

[default] 
aws_access_key_id = YOUR_ACCESS_KEY
aws_secret_access_key = YOUR_SECRET_KEY
region = YOUR_AWS_REGION

Step 3: Read data from AWS

Now we will be able to read the data from S3 or from an Athena easily using the awswrangler library:

import awswrangler as wr
df_s3 = wr.s3.read_csv(path='s3://bucket/prefix/')
df_athena = wr.athena.read_sql_query("SELECT col1,col2,col3 FROM table", database="db")

If the data in s3 is partitioned for example by year and date, we can use a pushdown filter to read only relevant partitions. For example we have in s3 the paths:

s3://bucket/prefix/year=2022/month=10/
s3://bucket/prefix/year=2022/month=11/
s3://bucket/prefix/year=2023/month=10/

To read only the data from 2022 October we use the logic:

pushdown_filter = lambda x: True if x["year"]== "2022" and x["month"] == "10" else False
df_filter = wr.s3.read_csv("s3://bucket/prefix/", dataset=True, partition_filter=pushdown_filter)

Step 4: Add data validation

After we read the data we can apply column validation. Let’s day we read a DataFrame that contains 3 columns: age, sexand date:

age,sex,date
10,Male,2022-10-10
85,Female,2022-11-11
24,Female,2020-10-10

To be able to validate the data we’ll have to create a schema:

from pandas_schema import Column, Schema
from pandas_schema.validation import InRangeValidation, InListValidation,DateFormatValidation

schema = Schema([
    Column('age', [InRangeValidation(0, 100)]),
    Column('sex', [InListValidation(['Male', 'Female', 'Other'])]),
    Column('date', [DateFormatValidation('%Y-%m-%d')])
])

So the schema will check that the age column is in range 0–100, the sex column contain values from a list, and the date column should have a certain format.

Step 5: Check data validation

We can apply the schema on our DataFrames:

errors = schema.validate(df_s3)
errors_index_rows = [e.row for e in errors]
df_clean = df_s3.drop(index=errors_index_rows)

All the rows that have invalid data will be stored in the errors DataFrame. To eliminate there error rows we can drop them from the df_s3 DataFrame.