🔴 ML / AI

Python + Pandas: Data Wrangling

📖 Lesson 44 ⏱ 55 min 🧪 5 questions 💻 3 exercises

🎯 Learning Objectives

  • Create and inspect DataFrames and Series from dicts, lists, CSV, and JSON
  • Select data with [], .loc[], .iloc[], and boolean masks
  • Handle missing data: detect, drop, fill, and interpolate
  • Transform data: apply, map, assign, rename, type casting
  • Aggregate and group data with groupby, agg, pivot_table
  • Merge, join, and concatenate DataFrames
  • Reshape data with melt, pivot, stack, unstack

1 · DataFrames & Series

Pandas is built on NumPy; the two core types are Series (1-D labelled array) and DataFrame (2-D table with labelled rows and columns).

Install:

pip install pandas
terminal

Create DataFrames and Series from multiple sources:

import pandas as pd
import numpy as np

# ── From dict ──
df = pd.DataFrame({
    "name":   ["Alice", "Bob", "Carol", "Dave"],
    "age":    [30, 25, 28, 35],
    "salary": [95000, 72000, 85000, 110000],
    "dept":   ["Eng", "Sales", "Eng", "Mgmt"],
})

# ── From list of dicts ──
df2 = pd.DataFrame([
    {"name": "Alice", "score": 92},
    {"name": "Bob",   "score": 78},
])

# ── Series ──
s = pd.Series([10, 20, 30, 40], index=["a", "b", "c", "d"], name="values")
print(s["b"])    # 20
print(s[1:3])    # b:20, c:30

# ── Reading files ──
df_csv  = pd.read_csv("data.csv")
df_json = pd.read_json("data.json")
df_xl   = pd.read_excel("data.xlsx", sheet_name="Sheet1")
df_sql  = pd.read_sql("SELECT * FROM users", con=engine)

# ── Writing files ──
df.to_csv("out.csv", index=False)
df.to_json("out.json", orient="records")
df.to_parquet("out.parquet")   # compressed columnar format — recommended for large data
create_dataframes.py

Key inspection methods:

df.shape          # (4, 4)
df.dtypes         # column types
df.info()         # non-null counts + dtypes
df.describe()     # count, mean, std, min, quartiles, max
df.head(3)        # first 3 rows
df.tail(2)        # last 2 rows
df.columns.tolist()
df.index.tolist()
df.value_counts("dept")   # frequency of each value
inspect.py

2 · Selecting Data

import pandas as pd

df = pd.DataFrame({
    "name":   ["Alice", "Bob", "Carol", "Dave"],
    "age":    [30, 25, 28, 35],
    "salary": [95000, 72000, 85000, 110000],
    "dept":   ["Eng", "Sales", "Eng", "Mgmt"],
})

# ── Column selection ──
df["name"]                   # Series
df[["name", "salary"]]       # DataFrame with 2 columns

# ── Row selection ──
df[1:3]                      # rows 1 and 2 (slice, not loc)

# ── .loc[] — label-based ──
df.loc[0]                    # row with index label 0
df.loc[0, "name"]            # single value: "Alice"
df.loc[0:2, "name":"salary"] # rows 0-2, cols name-salary (inclusive)
df.loc[df["age"] > 27]       # boolean mask

# ── .iloc[] — integer position-based ──
df.iloc[0]                   # first row
df.iloc[0, 1]                # row 0, col 1 → 30
df.iloc[1:3, 0:2]            # rows 1-2, cols 0-1

# ── Boolean indexing ──
df[df["salary"] > 80000]
df[(df["dept"] == "Eng") & (df["age"] < 32)]
df[df["dept"].isin(["Eng", "Mgmt"])]
df[~df["name"].str.startswith("A")]   # NOT starting with A

# ── .query() — string expression ──
df.query("dept == 'Eng' and salary > 80000")
df.query("age.between(25, 30)")
selecting.py
Warning: Use .loc[] for label-based and .iloc[] for position-based selection. Never use chained indexing df["col"][0] for setting values — it produces a SettingWithCopyWarning. Always use df.loc[0, "col"] = value.

3 · Missing Data

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "A": [1.0, np.nan, 3.0, np.nan, 5.0],
    "B": [np.nan, 2.0, 3.0, 4.0, 5.0],
    "C": ["x", "y", None, "w", "v"],
})

# ── Detection ──
df.isna()             # boolean DataFrame
df.isna().sum()       # count per column
df.isna().sum() / len(df)   # fraction missing per column
df.notna()

# ── Dropping ──
df.dropna()                        # drop rows with ANY NaN
df.dropna(how="all")               # drop rows with ALL NaN
df.dropna(subset=["A", "B"])       # only check these columns
df.dropna(axis=1, thresh=3)        # drop columns with fewer than 3 non-NaN

# ── Filling ──
df.fillna(0)                              # fill all NaN with 0
df["A"].fillna(df["A"].mean())            # fill with column mean
df.fillna({"A": 0, "B": df["B"].median()})  # per-column fill values
df["A"].ffill()                           # forward-fill (propagate last valid)
df["A"].bfill()                           # backward-fill

# ── Interpolation ──
df["A"].interpolate(method="linear")      # linear interpolation
df["A"].interpolate(method="polynomial", order=2)

# ── Replacing specific values ──
df.replace(-999, np.nan)          # treat -999 as missing
df.replace({"A": {-1: np.nan}})   # column-specific replacement
missing_data.py
Tip: Always check df.isna().sum() as the first step of any data analysis — understanding the missingness pattern guides whether to drop, fill, or impute.

4 · Transforming Data

import pandas as pd

df = pd.DataFrame({
    "name":   ["alice smith", "bob jones", "carol white"],
    "age":    [30, 25, 28],
    "salary": [95000, 72000, 85000],
    "hire_date": ["2020-03-15", "2019-07-22", "2021-01-10"],
})

# ── Type casting ──
df["hire_date"] = pd.to_datetime(df["hire_date"])
df["age"]       = df["age"].astype(np.int32)
df["dept_code"] = pd.Categorical(["Eng", "Sales", "Eng"])   # memory-efficient

# ── String operations ──
df["name"].str.title()              # "Alice Smith", "Bob Jones"
df["name"].str.split(" ")           # split into list
df["name"].str.extract(r"(\w+)$")   # last word (surname)
df["name"].str.contains("alice", case=False)

# ── assign() — add/replace columns without mutation ──
df = df.assign(
    annual_bonus = df["salary"] * 0.10,
    name_upper   = df["name"].str.upper(),
    years_hired  = (pd.Timestamp.now() - df["hire_date"]).dt.days // 365,
)

# ── apply() — row or column-wise custom function ──
def grade(salary):
    if salary >= 90000: return "Senior"
    if salary >= 75000: return "Mid"
    return "Junior"

df["grade"] = df["salary"].apply(grade)

# Row-wise apply (use axis=1)
df["summary"] = df.apply(
    lambda row: f"{row['name']} ({row['grade']})",
    axis=1
)

# ── rename() ──
df = df.rename(columns={"salary": "annual_salary", "age": "age_years"})

# ── map() for Series value substitution ──
dept_map = {"Eng": "Engineering", "Sales": "Sales Dept"}
df["dept_full"] = df["dept_code"].map(dept_map)
transform.py
Tip: Prefer vectorised string methods (str.title(), str.contains()) and assign() over apply() for column transformations — apply() is a Python loop and can be 10–100× slower.

5 · Grouping & Aggregation

import pandas as pd
import numpy as np

df = pd.DataFrame({
    "dept":   ["Eng", "Sales", "Eng", "Mgmt", "Sales", "Eng"],
    "name":   ["Alice", "Bob", "Carol", "Dave", "Eve", "Frank"],
    "salary": [95000, 72000, 85000, 110000, 68000, 91000],
    "yoe":    [5, 3, 4, 10, 2, 6],
})

# ── groupby basics ──
grp = df.groupby("dept")
grp["salary"].mean()       # mean salary per dept
grp["salary"].agg(["mean", "min", "max", "count"])

# ── Multiple aggregations with named output ──
grp.agg(
    avg_salary = ("salary", "mean"),
    total_yoe  = ("yoe",    "sum"),
    headcount  = ("name",   "count"),
).reset_index()

# ── Custom aggregation ──
grp["salary"].agg(lambda s: s.max() - s.min())   # salary range per dept

# ── transform() — broadcast aggregated values back to original shape ──
df["dept_avg_salary"] = df.groupby("dept")["salary"].transform("mean")
df["salary_vs_avg"]   = df["salary"] - df["dept_avg_salary"]

# ── filter() — keep groups satisfying a condition ──
df.groupby("dept").filter(lambda g: len(g) >= 2)   # keep depts with 2+ people

# ── pivot_table ──
pivot = df.pivot_table(
    values="salary",
    index="dept",
    aggfunc={"salary": ["mean", "count"]},
)

# ── crosstab ──
pd.crosstab(df["dept"], df["yoe"] > 4, margins=True)
grouping.py

6 · Merging & Joining DataFrames

import pandas as pd

employees = pd.DataFrame({
    "emp_id": [1, 2, 3, 4],
    "name":   ["Alice", "Bob", "Carol", "Dave"],
    "dept_id": [10, 20, 10, 30],
})

departments = pd.DataFrame({
    "dept_id":   [10, 20, 40],
    "dept_name": ["Engineering", "Sales", "Marketing"],
})

reviews = pd.DataFrame({
    "emp_id": [1, 2, 1, 3],
    "score":  [4.5, 3.8, 4.9, 4.2],
    "year":   [2023, 2023, 2024, 2024],
})

# ── merge() — SQL-style joins ──
pd.merge(employees, departments, on="dept_id", how="inner")  # 3 rows (Dave excluded)
pd.merge(employees, departments, on="dept_id", how="left")   # 4 rows (Dave keeps NaN dept)
pd.merge(employees, departments, on="dept_id", how="right")  # includes Marketing
pd.merge(employees, departments, on="dept_id", how="outer")  # all rows, NaN where no match

# Different column names
pd.merge(employees, departments, left_on="dept_id", right_on="dept_id")

# ── join() — merge on index ──
employees.set_index("dept_id").join(departments.set_index("dept_id"))

# ── concat() — stack DataFrames ──
df_2023 = reviews[reviews["year"] == 2023]
df_2024 = reviews[reviews["year"] == 2024]
pd.concat([df_2023, df_2024], ignore_index=True)    # vertical (row-wise)
pd.concat([employees, reviews], axis=1)              # horizontal (column-wise)

# ── Multi-key merge ──
pd.merge(employees, reviews, on="emp_id").groupby("emp_id")["score"].mean()
merging.py

Join Type Comparison

Join Type Rows Kept SQL Equivalent
inner Only rows with matching keys in both DataFrames SELECT … FROM A INNER JOIN B ON …
left All rows from left DF; matched rows from right (NaN if no match) SELECT … FROM A LEFT JOIN B ON …
right All rows from right DF; matched rows from left (NaN if no match) SELECT … FROM A RIGHT JOIN B ON …
outer All rows from both DFs; NaN where no match on either side SELECT … FROM A FULL OUTER JOIN B ON …

7 · Reshaping: melt, pivot, stack, unstack

import pandas as pd

# ── Wide → Long with melt() ──
wide = pd.DataFrame({
    "name":  ["Alice", "Bob"],
    "math":  [90, 78],
    "english": [85, 92],
    "science": [88, 71],
})

long = wide.melt(
    id_vars=["name"],
    value_vars=["math", "english", "science"],
    var_name="subject",
    value_name="score",
)
#    name  subject  score
# 0  Alice  math      90
# 1  Bob    math      78
# ...

# ── Long → Wide with pivot() ──
long.pivot(index="name", columns="subject", values="score")

# ── pivot_table (handles duplicates, supports aggfunc) ──
sales = pd.DataFrame({
    "month":  ["Jan","Jan","Feb","Feb"],
    "region": ["N","S","N","S"],
    "units":  [100, 150, 120, 130],
})
sales.pivot_table(index="month", columns="region", values="units", aggfunc="sum")

# ── stack() / unstack() ── (works with MultiIndex)
df_mi = long.set_index(["name", "subject"])["score"]
df_mi.unstack("subject")   # subjects become columns  (same as pivot)
df_mi.unstack("subject").stack("subject")  # back to long

# ── explode() — one row per list element ──
df_tags = pd.DataFrame({
    "post_id": [1, 2],
    "tags":    [["python", "pandas"], ["numpy", "ml", "ai"]],
})
df_tags.explode("tags")
#    post_id    tags
# 0        1  python
# 0        1  pandas
# 1        2   numpy
# ...
reshaping.py
Concept: melt (wide-to-long) and pivot (long-to-wide) are inverses. Most tidy-data operations require long format — one observation per row. Visualisation libraries like Seaborn and Plotly Express also prefer long format.

Time Series

Pandas has first-class datetime support — parsing, resampling, rolling windows, and timezone handling are all built in.

import pandas as pd
import numpy as np

# ── Create a datetime index ──
idx = pd.date_range("2024-01-01", periods=365, freq="D")
ts  = pd.Series(np.random.randn(365).cumsum(), index=idx, name="price")

# ── Indexing by date string ──
ts["2024-03"]              # all of March 2024
ts["2024-01-01":"2024-03-31"]  # date range slice

# ── Resampling (change frequency) ──
ts.resample("W").mean()    # weekly averages
ts.resample("ME").last()   # month-end values   ("M" deprecated → "ME")
ts.resample("QE").agg({"price": ["first","last","max","min"]})

# ── Rolling windows ──
ts.rolling(window=7).mean()     # 7-day moving average
ts.rolling(window=30).std()     # 30-day rolling std
ts.ewm(span=7).mean()           # exponentially weighted moving average

# ── Shifting ──
ts.shift(1)           # lag by 1 period (NaN at start)
ts.shift(-1)          # lead by 1 period
ts.pct_change()       # daily percentage return = (t - t-1) / t-1
ts.diff()             # first difference

# ── Timezone handling ──
ts_utc = ts.tz_localize("UTC")
ts_ny  = ts_utc.tz_convert("America/New_York")

# ── Date component extraction ──
df_dt = pd.DataFrame({"date": pd.date_range("2024-01-01", periods=10)})
df_dt["year"]    = df_dt["date"].dt.year
df_dt["month"]   = df_dt["date"].dt.month
df_dt["weekday"] = df_dt["date"].dt.day_name()
df_dt["is_weekend"] = df_dt["date"].dt.dayofweek >= 5
time_series.py
Frequency aliases — common resample codes: "D" (calendar day), "B" (business day), "W" (week), "ME" (month end), "QE" (quarter end), "YE" (year end), "h" (hour), "min" (minute). Use pd.tseries.frequencies.to_offset("2W") for custom offsets.

Performance: Categorical, Chunking & PyArrow

Large DataFrames can exhaust RAM. These techniques reduce memory and speed up operations.

import pandas as pd
import numpy as np

# ── Categorical dtype ──
# If a string column has low cardinality, convert to Categorical
df = pd.DataFrame({"dept": ["Eng"] * 1000 + ["Sales"] * 500 + ["Mgmt"] * 100})
print(df["dept"].memory_usage(deep=True))     # ~96 KB as object (string)
df["dept"] = df["dept"].astype("category")
print(df["dept"].memory_usage(deep=True))     # ~2 KB as category

# Operations on categoricals are faster
df.groupby("dept").size()   # uses integer codes internally

# ── Optimal dtypes at load time ──
df_large = pd.read_csv(
    "large.csv",
    dtype={
        "id":     np.int32,     # vs int64 default
        "score":  np.float32,   # vs float64 default
        "status": "category",   # vs object default
    },
    parse_dates=["created_at"],
)

# ── Chunked reading for files larger than RAM ──
total = 0
for chunk in pd.read_csv("huge.csv", chunksize=100_000):
    total += chunk["amount"].sum()   # process 100k rows at a time
print(f"Total: {total}")

# ── Parquet — columnar, compressed, fast ──
df.to_parquet("data.parquet", index=False)
df2 = pd.read_parquet("data.parquet", columns=["name", "salary"])  # read only needed cols

# ── PyArrow backend (Pandas 2.0+) ──
df_arrow = pd.read_csv("data.csv", dtype_backend="pyarrow")
# nullable dtypes, better performance, smaller memory footprint
performance.py
The Parquet format is the recommended storage format for analytical data: columnar storage (reads only needed columns), compressed (typically 5–10× smaller than CSV), and preserves dtypes (no re-parsing on load). Always prefer Parquet over CSV for datasets larger than a few MB.

A Complete Wrangling Workflow

End-to-end: load raw data → inspect → clean → transform → aggregate → export.

import pandas as pd
import numpy as np

# ── 1. Load ──
df = pd.read_csv("sales.csv", parse_dates=["order_date"])
print(df.shape)          # (10000, 8)
print(df.dtypes)
print(df.isna().sum())

# ── 2. Clean ──
df = df.dropna(subset=["customer_id", "amount"])   # required fields
df["amount"] = df["amount"].clip(lower=0)           # no negative amounts
df["category"] = df["category"].str.strip().str.lower().astype("category")
df = df.drop_duplicates(subset=["order_id"])

# ── 3. Transform ──
df = df.assign(
    year        = df["order_date"].dt.year,
    month       = df["order_date"].dt.month,
    quarter     = df["order_date"].dt.quarter,
    revenue     = df["amount"] * df["quantity"],
    is_large    = df["amount"] > df["amount"].quantile(0.9),
)

# ── 4. Aggregate ──
monthly = (
    df.groupby(["year", "month", "category"])
    .agg(
        orders      = ("order_id",  "count"),
        revenue     = ("revenue",   "sum"),
        avg_amount  = ("amount",    "mean"),
    )
    .reset_index()
    .sort_values(["year", "month"])
)

# ── 5. Top-N analysis ──
top_customers = (
    df.groupby("customer_id")["revenue"]
    .sum()
    .nlargest(10)
    .reset_index()
    .rename(columns={"revenue": "total_revenue"})
)

# ── 6. Export ──
monthly.to_parquet("monthly_summary.parquet", index=False)
top_customers.to_csv("top_customers.csv", index=False)
workflow.py

Best Practices

  • Use .loc[] for all assignmentsdf.loc[mask, "col"] = val avoids SettingWithCopyWarning and ensures you modify the original DataFrame, not a hidden copy.
  • Prefer vectorised operations over apply() — string methods (str.*), arithmetic, pd.to_datetime(), and assign() are all faster than row-wise apply(axis=1).
  • Convert low-cardinality string columns to category — dramatically reduces memory and speeds up groupby and sort_values.
  • Use Parquet instead of CSV for anything larger than a few MB — preserves dtypes, 5–10× smaller, and reads only requested columns.
  • Chain operations with assign and method chaining — avoid intermediate variables; use parentheses to wrap multi-line chains for readability.
  • Always check df.isna().sum() first — understand the missingness pattern before deciding whether to drop, fill, or model-impute.
  • Use reset_index() after groupby.agg() — the group keys become regular columns, making subsequent operations cleaner.
  • Store tidy data in long format — one observation per row. Use melt() to convert wide data; plotting libraries expect long format.

Exercises

Exercise 1 — Employee Analysis Pipeline

Given a CSV with columns emp_id, name, dept, salary, hire_date, is_active:

  • Load it; cast hire_date to datetime and dept to category.
  • Drop rows where salary is null; fill missing is_active with True.
  • Add columns: years_tenure (years since hire), salary_band (Junior/Mid/Senior using pd.cut).
  • Produce a summary table: for each dept show headcount, mean salary, median tenure, and % active.
  • Find the top-3 highest-paid employees per department.
💡 Hint — salary band & top-3
df["salary_band"] = pd.cut(
    df["salary"],
    bins=[0, 60000, 90000, float("inf")],
    labels=["Junior", "Mid", "Senior"],
)

# Top-3 per dept
df.sort_values("salary", ascending=False).groupby("dept").head(3)

Exercise 2 — Sales Time Series

Given a CSV with order_id, order_date, amount, region:

  • Parse order_date; set it as the index.
  • Compute daily total revenue and plot a 7-day rolling average.
  • Resample to monthly totals; pivot so rows = months, columns = regions.
  • Find the month with the highest revenue in each region.
  • Compute month-over-month percentage change per region using pct_change().
💡 Hint — monthly pivot
daily = df.groupby([pd.Grouper(freq="D"), "region"])["amount"].sum().reset_index()
monthly_pivot = daily.set_index("order_date").groupby("region").resample("ME")["amount"].sum().unstack("region")
monthly_pivot.pct_change()   # MoM change per region

Exercise 3 — Data Cleaning Challenge

Download the Titanic dataset (pd.read_csv("https://raw.githubusercontent.com/datasciencedojo/datasets/master/titanic.csv")) and perform a full cleaning pipeline:

  • Report null counts and dtypes.
  • Fill Age with the median age per Pclass + Sex group (use groupby + transform).
  • Drop Cabin (too many nulls) and Ticket (non-informative).
  • Extract the title from Name (Mr, Mrs, Miss, etc.) into a new column; group rare titles as "Other".
  • Encode Sex as 0/1; one-hot-encode Embarked and title with pd.get_dummies.
  • Output the final clean DataFrame shape and confirm zero nulls.
💡 Hint — fill Age by group median
df["Age"] = df.groupby(["Pclass", "Sex"])["Age"].transform(
    lambda s: s.fillna(s.median())
)

# Extract title
df["title"] = df["Name"].str.extract(r",\s*(\w+)\.")
common = df["title"].value_counts().nlargest(4).index
df["title"] = df["title"].where(df["title"].isin(common), other="Other")