Python Essentials for Data Engineers

Python knowledge is essential for data engineers. But what does knowing Python mean?

If you have

Wondered what aspects of Python you’d need to know to become a data engineer

Questioned your ability to learn Python, especially since Python keeps on adding new things

Want to learn how to use Python practically for data engineering

Then, this post is for you.

Python is the glue that holds data pipelines together.

By the end of this post, you will understand Python’s role in data engineering and how to use it to build data pipelines.

Python instructs data systems to process data

Python is not designed for data processing (think joins, group by, etc), and the amount of data it can process is limited by the machine’s memory.

In Memory Python v Out of Memory Processing

Most large-scale data processing is done by systems such as Spark, Snowflake, BigQuery, DuckDB, and Polars.

Python has APIs/librarues to instruct these systems on how to process data.

For example, in the code below, we use the PySpark API to specify how Spark should process data.

spark\
.sql(""" 
Select orderdate
, sum(price) as total_amount 
From orders 
Left Join customer 
on customer_id 
Group by orderdate
""")\
.show(5)

Project code along

As we progress through this post, we will build a 3-stage pipeline and see how Python supports each stage.

1. Weather data: Extracted from public S3, transformed with Spark, and loaded to DuckDB
2. Pokemon data: Extracted from PokeAPI, transformed with Spark, and loaded to DuckDB
3. Orders data: Extracted from the file system, transformed with Spark, DQ checked with pointblank, and loaded to DuckDB

Python Pipelines (click to enlarge)

Prerequisite:

1. Python basics
2. SQL basics.

Run the code below to get started.

git clone https://github.com/josephmachado/python_essentials_for_data_engineers.git 
cd python_essentials_for_data_engineers
uv sync
uv run jupyter lab

1

Clone and cd into the repo with code for this post

2

Install all dependencies at pyproject.toml

3

Start a jupyter lab instance at http://localhost:8888

Open Jupyter Lab at http://localhost:8888. Copy and paste the code into the workshop notebook.

Solutions are available at solutions notebook.

Video walkthrough

Python can Extract data from any system and format

Python has multiple libraries that enable reading and writing to various systems. In addition, Python can read/write to any data format.

Note

Data processing systems like Spark need their own libraries to connect directly to systems.

E.g., Spark needs JDBC to interact with databases.

Let’s extract data from our sources.

API data extraction with requests

Let’s extract Pokémon information from pokeapi as shown below.

import requests
url = "https://pokeapi.co/api/v2/pokemon/"

response = requests.get(url)
pokemon = response.json() 
pokemon_data = pokemon['results']

pokemon_data[:2]

1

requests library to interact with APIs

2

HTTP GET to get data from the API

3

Explore the pokemon variable.

4

Print a sample of pokemon_data.

S3 data extraction with boto3

Let’s extract weather data from a publicly open S3 bucket as shown below.

import csv
import gzip
from io import StringIO

import boto3
from botocore import UNSIGNED
from botocore.client import Config

bucket_name = "noaa-ghcn-pds"
file_key = "csv.gz/by_station/ASN00002022.csv.gz"

s3_client = boto3.client("s3", config=Config(signature_version=UNSIGNED))

response = s3_client.get_object(Bucket=bucket_name, Key=file_key)
compressed_data = response["Body"].read()

csv_data = gzip.decompress(compressed_data).decode("utf-8")

csv_reader = csv.reader(StringIO(csv_data))
weather_data = list(csv_reader)

1

We use boto3 to connect to AWS services.

2

Since this is a public bucket, we specify that we can just download this (don’t in prod)

3

We specify the bucket name and file path.

4

We create an S3 client.

5

We pull the specific file.

6

We unzip the downloaded file.

7

We convert the data into a list of dicts

Data documentation for NOAA

Extracting data from a local file with PySpark

Let’s see how we can use the PySpark API to instruct Spark to read data from a file.

from sqlframe import activate
activate("duckdb")

from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
customer_df = spark.read.csv("../data/customer.csv", header=True, inferSchema=True)

customer_df.show(2)

1

We use sqlframe to simulate Spark

2

We use spark.read.csv to read data from a local folder.

Exercise: Extract data from orders.csv

Use the PySpark API to read data from ../data/orders.csv similar to customer_df.

TipSolution

orders_df = spark\
  .read\
  .csv(
  "../data/orders.csv", 
  header=True, 
  inferSchema=True
  )

orders_df.show(2)

NoteTakeaways

When extracting data from a source system, Python will need 2 capabilities:

1. Connecting to the source system. E.g., boto3 to connect to S3
2. Ability to parse the data format. Python natively supports multiple formats: csv, XML, json. However, you will need to install libraries to read data formats like parquet, avro, etc

Video walkthrough

Python transforms data with SQL/Dataframe interface

Most transformation in data engineering is SQL-like, meaning it is often a mix of

1. Modifying/Adding new columns
2. Joins
3. Group By with aggregates
4. Window functions

Data processing Python libraries (PySpark, Pandas, Polars, DuckDB, etc.) make it easy to implement these operations.

Let’s transform our weather data.

weather_header = [
  'id', 
  'date', 
  'element', 
  'value', 
  'm_flag', 
  'q_flag', 
  's_flag', 
  'obs_time'
]
weather_df = spark\
.createDataFrame(
  weather_data, 
  schema=weather_header
)

1

Create a Spark dataframe from a list of dicts. No-op transformation.

NoteDataFrame

When working with PySpark DataFrame, there are 2 main patterns to remember.

1. Pyspark Functions: Functions that operate on column(s) Pyspark Functions
2. Dataframe Functions: Changing the shape of the dataframe as a whole, e.g., group by, join, sort, filter, drop rows, union, etc

# Examples
from pyspark.sql import functions as F

F.upper("your-column")
# etc

df.groupBy("date").agg({"price": "sum"}).show()

1

Pyspark Functions

2

Dataframe Functions

Let’s join orders and customer data and enrich them with an age column.

from pyspark.sql import functions as F

enriched_orders_df = orders_df.join(
    customer_df,
    orders_df["customer_id"] == customer_df["id"],
    how="left",
)

enriched_orders_df = enriched_orders_df\
.withColumn(
    "age",
    (
      F.year(F.current_date()) - 
      F.year(F.to_date("date_of_birth"))
    ).cast("int"),
)

1

We enrich order data with customer information.

2

Computing customers’ age as of today.

Exercise: Create age_bucket column

Use when & otherwise to create a age_bucket column, with the following categories: lt18, 18-30, 31-50, 51-70, 70+.

TipSolution

enriched_orders_df = enriched_orders_df.withColumn(
    "age_bucket",
    F.when(F.col("age") < 18, "lt18")
     .when(F.col("age") <= 30, "18-30")
     .when(F.col("age") <= 50, "31-50")
     .when(F.col("age") <= 70, "51-70")
     .otherwise("70+"),
)
enriched_orders_df.show(2)

1

Similar to case .. when in SQL

There are cases where we need to work directly in Python.

Assume we need to get the ID, name, base_experience, height, and weight for the first 151 Pokémon. From the extract section, we only got the first few Pokémon (because we did not paginate).

Recommended reading: Extracting data from an API

Let’s take a swing at getting all 151 Pokémon.

import requests
 
url = "https://pokeapi.co/api/v2/pokemon/"
results = []
while url and len(results) < 151:
    resp = requests.get(url).json()
    results.extend(resp["results"])
    url = resp["next"]  # follow pagination link
results = results[:151]

len(results)

1

Logic to loop and follow the next URL

However, there is a better way. Reading the documents, we can see that there is a parameter we can use to get all 151 without looping.

Documentation describing limit in the pokeapi

url = "https://pokeapi.co/api/v2/pokemon/?limit=151"
results = requests.get(url).json()["results"]
len(results)

Tip

Always read the documentation.

Video walkthrough

Exercise: Enrich Pokemon data

Enrich the 151 Pokémon with ID, name, base_experience, height, and weight details.

Here is a sample of an enrichment of a single Pokémon.

url = 'https://pokeapi.co/api/v2/pokemon/59/'
arcanine_data = requests.get(url).json()
print(f"id: {arcanine_data['id']}, name: {arcanine_data['name']}, base_experience: {arcanine_data['base_experience']}, height: {arcanine_data['height']}, weight: {arcanine_data['weight']}")

1

We hit the URL from the results variable.

TipSolution

We can loop through each URL; however, that would be very slow.

pokemon_data = []
i = 0
for p in results: # 
    print(f'Hitting {p["url"]}')
    detail = requests.get(p["url"]).json()
    pokemon_data.append({
        "id": detail["id"],
        "name": detail["name"],
        "base_experience": detail["base_experience"],
        "height": detail["height"],
        "weight": detail["weight"],
    })
    i += 1
    if i > 5:
        break

1

Requesting data 151 times will take a while

A better way is to read the docs. If we look at the docs, we can see a graphql endpoint.

TL;DR: GraphQL enables us to get the exact data we need using the GraphQL API. Reference docs.

Let’s take a look at what that would look like.

query = """
query samplePokeAPIquery {
  pokemon(limit: 151, order_by: {id: asc}) {
    id
    name
    height
    base_experience
    weight
  }
}
"""

resp = requests.post(
    "https://graphql.pokeapi.co/v1beta2",
    json={"query": query},
)
pokemon_data = resp.json()["data"]["pokemon"]
print(len(pokemon_data), pokemon_data[58])

1

This is the GraphQL query

2

We make a POST request to get this data

Video walkthrough

Load data with Python into your destination system

As with Extract, Python can load data into every data system.

import duckdb

con = duckdb.connect("my_database.duckdb")

pokemon_df = spark.createDataFrame(pokemon_data)

for name, sdf in [
    ("pokemon", pokemon_df),
    ("weather", weather_df),
    ("enriched_orders", enriched_orders_df),
]:
    pdf = sdf.toPandas()
    con.register(f"{name}_tmp", pdf)
    con.execute(f"CREATE OR REPLACE TABLE {name} AS SELECT * FROM {name}_tmp")
    con.unregister(f"{name}_tmp")

con.close()

1

Creates a duckdb database

2

We convert Spark to a pandas df; don’t do this with large data

3

Loads weather, Pokémon, & orders data into duckdb tables

Now assume a stakeholder accesses the data & finds duplicate data.

Warning

con = duckdb.connect("my_database.duckdb")
report = """
  select order_id,
  count(*) 
  from enriched_orders 
  group by 1 
  having count(*) > 1
  """
print(con.execute(report).fetchdf())
con.close()

Quality check data before end-user access

We missed a key part of any data pipeline, the data quality check. The quality check must be completed before the data is made available to stakeholders.

Recommended reading: WAP pattern

Based on the issues above, we need to check for duplicates and ensure uniqueness.

We use the pointblank library for a data quality check.

con = duckdb.connect("my_database.duckdb")
enriched_orders_query = "select * from enriched_orders"
enriched_orders_data = con.execute(enriched_orders_query).fetchdf()
con.close()

import pointblank as pb

validation = (
    pb.Validate(
      data=enriched_orders_data,
      tbl_name="enriched_orders_data",
      thresholds=pb.Thresholds(error=1, critical=2)
    )
    .rows_distinct(columns_subset=["order_id"])
    .col_vals_not_null(columns="first_name")
    .col_vals_not_null(columns="price")
    .interrogate()
)
validation.get_tabular_report().show()
print("All passed:", validation.all_passed())

1

Read the bad data

2

If the number of failing rows is >= 1, it’s considered an error.

3

Check for order_id distinctness.

4

Check for first_name not null.

5

Check for price not null.

6

Print output as a table

Note the failures:

Data Quality Failures

Now let’s clean our data and check again.

cleaned_enriched_orders_df = enriched_orders_df\
  .dropDuplicates(["id"])\
  .fillna({"first_name": "UNKNOWN", "price": 0})

import pointblank as pb

validation = (
    pb.Validate(
      data=cleaned_enriched_orders_df.toPandas(),
      tbl_name="cleaned_enriched_orders_data",
      thresholds=pb.Thresholds(error=1, critical=2)
    )
    .rows_distinct(columns_subset=["order_id"])
    .col_vals_not_null(columns="first_name")
    .col_vals_not_null(columns="price")
    .interrogate()
)

validation.get_tabular_report().show()
print("All passed:", validation.all_passed())

1

Remove duplicates and fill nulls.

2

If the number of failing rows is >= 1, it’s considered an error.

3

Check for order_id distinctness.

4

Check for first_name not null.

5

Check for price not null.

6

Print output as a table

We can be certain now that our data doesn’t have duplicates.

Data Quality Pass

Our data is now ready to be loaded into the destination table.

Video walkthrough

Test code, be sure it does what you think it does

Testing your code ensures it does exactly what it is supposed to do.

Tests (different from data quality checks) are run before your code is deployed to production.

When you write code, you and your colleagues must maintain it over time. Bugs are inevitable with code changes/updates; tests will help you catch these before your code goes live!

Let’s use pytest to check that the logic in this function works as expected.

! uv run pytest ./

Recommended reading: How to use pytest

Define pipeline run order with an orchestrator.

In the examples above, we’ve seen ETL written in Python. But you may need to run tasks across multiple systems, e.g., PySpark for ETL and then send an email.

We may also have to implement time-based logic (e.g., run this pipeline only at the end of the month).

These dependencies can be defined with an orchestrator.

Learn the industry standard orchestrators:

1. dbt core
2. Airflow

Define pipeline run time with a scheduler.

Schedulers are used to run data pipelines at specified times (once a day, once an hour, once a week, etc).

Scheduler systems run continuously (more accurately, run every n seconds, sleep, then run again) to check whether there are any data pipelines (or tasks) to run and, if so, start them.

Learn the industry-standard scheduler:

1. Airflow

Other such tools are: cron, APScheduler, Windows Scheduler.

Conclusion

To recap, we saw

1. How python can extract data from any source
2. Using SQL/Dataframe API to tranform data
3. Loading data into any destination system
4. Checking data quality with tools
5. Testing code
6. Scheduling and orchestration tools

We saw how Python can be used to glue together the parts of a data pipeline and how Python’s vast libraries allow it to work with any system and any data format.

I want to hear about your favorite Python libraries. Let me know in the comments below.