Spark DataFrame examples for Spark & Data Eng developers

Spark DataFrame examples for Spark & Data Eng developers

If you're building data pipelines in production, you'll spend most of your time transforming DataFrames—selecting columns, filtering rows, computing aggregates, and joining datasets together. These operations form the backbone of every real-world Spark job, from ETL workflows to analytics engines. In this guide, we'll walk through the essential DataFrame operations you'll use daily, with practical examples pulled from actual production scenarios.

Select and filter — Pick columns and rows

Selection and filtering are your first line of defense when working with large datasets. In production, you often inherit schemas with dozens of columns you don't need, and raw data containing incomplete or invalid records. The select() and filter() methods let you trim both dimensions efficiently.

The select operation reduces your DataFrame to only the columns you care about. This matters for performance: fewer columns mean less memory pressure and faster I/O. When you're working with Parquet or Delta tables, Spark's predicate pushdown can skip entire column families if you select carefully upfront.

from pyspark.sql import functions as F

# Load a user events table with 40+ columns
df = spark.read.parquet("/data/events")

# Keep only what you need
active = df.filter(F.col("status") == "active").select("id", "email", "country")
active.show(5)

The filter operation removes rows that don't meet your criteria. You can chain multiple conditions using & (AND), | (OR), and ~ (NOT). Avoid using Python's native boolean operators—Spark needs its Column expressions to push the filter down to the storage layer.

# Multiple conditions
recent_active = df.filter(
    (F.col("status") == "active") &
    (F.col("last_seen") > "2025-01-01") &
    (F.col("country").isin("US", "CA", "GB"))
)

# Wrong: uses Python bool, doesn't optimize
# bad = df.filter(df.status == "active" and df.last_seen > "2025-01-01")

In e-commerce pipelines, you might filter out test orders before joining with inventory. In SaaS analytics, you filter by account status before computing monthly active users. The principle is the same: reduce early, compute less.

Add columns — withColumn for derived values

Real-world data rarely arrives in the shape you need. You'll often create new columns by transforming existing ones, combining multiple fields, or calling UDFs. The withColumn() method is your tool for derived values.

Common scenarios include parsing timestamps, computing age ranges, normalizing text, and flagging anomalies. Each withColumn() call creates a new DataFrame (Spark is immutable), so chain them when possible rather than calling separately.

# Parse a timestamp string and extract components
df_parsed = df.withColumn(
    "signup_date",
    F.to_date(F.col("signup_timestamp"))
).withColumn(
    "signup_year",
    F.year(F.col("signup_date"))
).withColumn(
    "email_domain",
    F.split(F.col("email"), "@")[1]
)

df_parsed.select("email", "email_domain", "signup_year").show(3)

You can also use withColumn() to flag records for downstream processing. In fraud detection, you might add a risk score. In ETL, you might add a processing timestamp to track data lineage.

# Add a processing timestamp and risk flag
df_enriched = df.withColumn(
    "processed_at",
    F.current_timestamp()
).withColumn(
    "is_high_value",
    F.col("revenue") > 10000
).withColumn(
    "customer_tier",
    F.when(F.col("revenue") > 50000, "premium")
     .when(F.col("revenue") > 10000, "gold")
     .otherwise("standard")
)

Watch the performance cost of UDFs (user-defined functions) in Python—they serialize data back to the Python runtime, bypassing Catalyst optimization. Use Spark SQL functions whenever possible; they execute in the JVM and are orders of magnitude faster.

Group and aggregate — groupBy with agg

Aggregations power dashboards, reports, and feature engineering. You'll use groupBy() paired with aggregate functions to compute sums, counts, averages, percentiles, and custom metrics across dimensions.

In a real ride-hailing system, you group by driver and day to calculate total earnings and trip count. In e-commerce, you group by product category to find top sellers. In SaaS billing, you group by account to compute monthly recurring revenue.

# Basic aggregation: trips per driver per day
daily_driver_stats = df_trips.groupBy("driver_id", "trip_date").agg(
    F.count("*").alias("total_trips"),
    F.sum("fare_amount").alias("total_revenue"),
    F.avg("rating").alias("avg_rating"),
    F.min("rating").alias("min_rating")
)

daily_driver_stats.show()

Chain multiple aggregations together in a single agg() call—this is far more efficient than separate operations. Spark can compute all metrics in one pass, not multiple passes over the data.

# Multi-column grouping with multiple aggregates
customer_summary = df_orders.groupBy("customer_id", "country").agg(
    F.count("order_id").alias("order_count"),
    F.sum("amount").alias("lifetime_value"),
    F.avg("amount").alias("avg_order_value"),
    F.percentile_approx("amount", 0.95).alias("p95_order_value"),
    F.max("order_date").alias("last_order_date"),
    F.collect_list("product_id").alias("products_purchased")
)

customer_summary.filter(F.col("order_count") > 5).show()

When you need custom logic inside groupBy, use groupBy().apply() with pandas UDFs for better performance than row-at-a-time operations. For window functions (running totals, ranks, lag/lead), use Window specifications instead of groupBy.

Join — df.join(other, on="id", how="left")

Joins combine datasets on shared keys. You'll join user tables with event tables, orders with products, transactions with accounts. Understanding join types and performance implications is critical in production.

Spark supports inner, left, right, outer, left_semi, and left_anti joins. Each has different semantics and cost. A left join preserves all rows from the left DataFrame; an inner join keeps only matching rows.

# Simple left join: keep all users, add their order info
df_users = spark.read.parquet("/data/users")
df_orders = spark.read.parquet("/data/orders")

result = df_users.join(
    df_orders,
    on="user_id",
    how="left"
)

result.select("user_id", "email", "order_id", "amount").show()

When join keys have different names, specify them explicitly:

# Join on mismatched column names
result = df_users.join(
    df_orders,
    df_users.id == df_orders.user_id,
    how="left"
)

For joins across multiple columns, pass a list or combine conditions with &:

# Multi-column join: match on user_id AND country
result = df_users.join(
    df_regional_settings,
    on=["user_id", "country"],
    how="inner"
)

# Or explicitly:
result = df_users.join(
    df_regional_settings,
    (df_users.user_id == df_regional_settings.user_id) &
    (df_users.country == df_regional_settings.country),
    how="inner"
)

Join performance matters. Large table-to-large table joins can be expensive. Spark broadcasts small DataFrames (under a few hundred MB) automatically. For very large joins, partition both sides on the join key before joining, or pre-aggregate to reduce cardinality. Monitor join performance with explain():

# See the execution plan
result.explain()

# Broadcast hint for small tables
from pyspark.sql import functions as F
result = df_large.join(
    F.broadcast(df_small),
    on="id",
    how="left"
)

Common mistakes

We've seen several patterns slow down production pipelines:

• Chaining filter and select inefficiently: Always filter first to reduce data volume, then select columns. Selecting first wastes computation on columns you'll discard.
• Using Python functions instead of Spark SQL functions: Python UDFs force serialization and break optimization. Use F.when(), F.split(), F.date_format() instead of custom Python logic where possible.
• Forgetting to handle nulls: Null comparisons are tricky in SQL. Use F.col("x").isNull() or F.coalesce() to handle missing values explicitly.
• Not partitioning before large joins: Joining two billion-row tables without partitioning will shuffle massive data across the cluster. Partition on the join key first.
• Collecting to driver without limits: .collect() brings everything to the driver's memory. Always use .limit() or .head()` for exploration.
• Ignoring column naming conflicts in joins: When join columns have the same name, you get duplicates. Use select() after the join to keep only one copy.

Key takeaways

• select() and filter() are your first optimization: trim columns early, then rows, to reduce memory and I/O.
• withColumn() creates derived fields; chain multiple calls and prefer Spark SQL functions over Python UDFs.
• groupBy().agg() computes metrics across dimensions; batch all aggregates in one agg() call for efficiency.
• join() combines datasets; choose join type carefully (left vs. inner), broadcast small tables, and partition large joins on the key.
• Always use Column expressions (F.col) for filtering and conditions—they enable Catalyst optimization and predicate pushdown.
• Call .explain() before running expensive operations to see the execution plan and catch join or shuffle bottlenecks.