Why Every Data Engineer Must Learn NumPy This Year

# pip install numpy
import time
import numpy as np

prices = np.random.default_rng(0).uniform(100, 5000, size=1_000_000)

# Python loop: add 18% tax one value at a time
start = time.perf_counter()
with_tax_loop = [p * 1.18 for p in prices]
loop_s = time.perf_counter() - start

# NumPy: one vectorised operation on the whole array
start = time.perf_counter()
with_tax_vec = prices * 1.18
vec_s = time.perf_counter() - start

print(f"loop: {loop_s:.3f}s  vectorised: {vec_s:.4f}s")
print("same result:", np.allclose(with_tax_loop, with_tax_vec))

Why vectorised code wins:

Python loop
- each number is a separate Python object
- the interpreter handles every step
- overhead repeated a million times

NumPy vectorised
- numbers stored side by side in memory
- the loop runs inside compiled code
- far less interpreter overhead

Rule of thumb for pipelines:
If you are looping over rows of numbers,
there is probably a vectorised way.

import numpy as np

n = 5_000_000
ids_64 = np.arange(n, dtype=np.int64)
ids_32 = ids_64.astype(np.int32)

print("int64 MB:", ids_64.nbytes / 1_000_000)
print("int32 MB:", ids_32.nbytes / 1_000_000)

# Always check the range fits before downcasting
info = np.iinfo(np.int32)
print("fits in int32:", ids_64.max() <= info.max and ids_64.min() >= info.min)

import numpy as np

readings = np.array([21.5, np.nan, 23.0, 22.1, np.nan])

print("plain mean:", readings.mean())          # nan - one NaN spoils it
print("nan-safe mean:", np.nanmean(readings))   # ignores NaN
print("missing count:", np.isnan(readings).sum())

# Fill gaps with the mean of valid readings
filled = np.where(np.isnan(readings), np.nanmean(readings), readings)
print("filled:", filled)

# Note: np.nan is a float, so an int array cannot store it.

# pip install pandas numpy
import numpy as np
import pandas as pd

orders = pd.DataFrame({
    "order_id": [101, 102, 103, 104],
    "amount": [450, 2200, 15000, -50],
    "delivery_days": [2, 5, 1, 3],
})

# Clean impossible values
orders["amount"] = np.clip(orders["amount"], 0, None)

# Multi-condition label without a slow apply()
conditions = [
    orders["amount"] >= 10000,
    orders["amount"] >= 1000,
]
choices = ["high", "medium"]
orders["value_band"] = np.select(conditions, choices, default="low")

orders["late_flag"] = np.where(orders["delivery_days"] > 4, 1, 0)
print(orders)

NumPy code review checklist for pipelines

[ ] No Python loops over large numeric columns
[ ] apply() replaced with vectorised logic where possible
[ ] dtypes checked after reading files
[ ] Downcasting only after range checks
[ ] NaN handled explicitly (nan-safe functions)
[ ] np.select used for multi-condition labels
[ ] Results compared with a small known sample

NumPy cheat sheet for pipelines

Create      np.array, np.arange, np.zeros, np.full
Inspect     arr.shape, arr.dtype, arr.nbytes
Convert     arr.astype(np.int32)
Filter      arr[arr > 0]
Condition   np.where(cond, a, b), np.select(conds, choices)
Missing     np.isnan, np.nanmean, np.nansum
Limits      np.clip(arr, low, high)
Bins        np.digitize(arr, bins)
Compare     np.allclose(a, b)
Stats       arr.mean(), arr.std(), np.percentile(arr, 95)

Benchmark Log

Transformation: _______________
Rows tested: _____

Approach            Time       Memory
1. Loop / apply     ____ s     ____ MB
2. Vectorised       ____ s     ____ MB

Results identical? (np.allclose)  Y / N
Chosen approach: _______________

Dtype Plan for a new dataset

Column         Current   Min    Max    Planned dtype
order_id       int64     ___    ___    int32
quantity       int64     ___    ___    int16
price          float64   ___    ___    float64
rating         int64     ___    ___    int8

Columns with missing values: _______________
Memory before: ____ MB   after: ____ MB