Learn / Lesson 2 · 6 min

Selecting & Filtering

Ninety percent of DataFrame work is two moves: pick the columns you want and keep the rows you care about. Master these and you can already do useful things to real data.

Throughout this lesson, imagine orders — the order log of a coffee chain:

order_iddrinksizepricetip
1001oat lattelarge6.501.00
1002drip coffeesmall3.000.00
1003caramel macchiatolarge7.252.00

select — choose your columns

result = orders.select("order_id", "drink", "price")

Strings work for simple picks. For anything fancier, use column expressions:

from pyspark.sql import functions as F

result = orders.select(
    F.col("order_id"),
    F.col("price") + F.col("tip"),                     # arithmetic on columns
    (F.col("tip") / F.col("price")).alias("tip_ratio") # rename with .alias()
)

F.col("price") isn't a value — it's an expression that means "the price column." Expressions combine with +, -, *, /, comparisons, and functions from pyspark.sql.functions.

filter — choose your rows

large_orders = orders.filter(F.col("size") == "large")

where is a synonym for filter; use whichever reads better. Combine conditions with & (and), | (or), ~ (not) — and wrap each condition in parentheses, because Python's operator precedence will betray you otherwise:

fancy = orders.filter((F.col("size") == "large") & (F.col("price") > 6))

Other row-picking tools you'll reach for constantly:

orders.filter(F.col("drink").isin("oat latte", "flat white"))
orders.filter(F.col("tip").isNotNull())
orders.filter(F.col("drink").startswith("caramel"))

withColumn — add or replace a column

orders_with_total = orders.withColumn("total", F.col("price") + F.col("tip"))

Remember the mental model: this returns a new DataFrame with the extra column. The original orders is untouched.

For conditional logic, when/otherwise is your if/else:

labeled = orders.withColumn(
    "vibe",
    F.when(F.col("price") > 6, "bougie").otherwise("sensible")
)

Chaining it all together

Real answers are usually a pipeline:

result = (
    orders
    .withColumn("total", F.col("price") + F.col("tip"))
    .filter(F.col("size") == "large")
    .select("order_id", "drink", "total")
)

Read it top to bottom: add a column, keep large orders, pick three columns. Every step is a lazy transformation — nothing runs until the grader calls the action.

Time to prove it. The practice problems for this lesson start with a warm-up filter and escalate to null-wrangling with coalesce and when. Go get 'em.

Practice what you learned