Skip to content

PySpark DataFrames

PySpark is the Python API for Apache Spark. This sheet covers DataFrame operations, transformations, actions, joins, aggregations, UDFs, and performance optimization.

Table of Contents

Installation / Setup

pip install pyspark

# Create SparkSession
from pyspark.sql import SparkSession
from pyspark.sql import functions as F
from pyspark.sql.types import *

spark = SparkSession.builder \
    .appName("PySparkApp") \
    .master("local[*]") \
    .getOrCreate()

sc = spark.sparkContext

Core Syntax

Creating DataFrames

# From list of tuples
data = [("Alice", 25, "Engineering"), ("Bob", 30, "Marketing")]
df = spark.createDataFrame(data, ["name", "age", "department"])

# From Pandas
import pandas as pd
pdf = pd.DataFrame({"name": ["Alice", "Bob"], "age": [25, 30]})
df = spark.createDataFrame(pdf)

# From files
df = spark.read.csv("data.csv", header=True, inferSchema=True)
df = spark.read.parquet("data.parquet")
df = spark.read.json("data.json")
df = spark.read.text("data.txt")
df = spark.read.load("data.avro", format="avro")

# S3 / cloud storage
df = spark.read.parquet("s3://bucket/path/to/data.parquet")

Schema Definition

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

schema = StructType([
    StructField("name", StringType(), True),
    StructField("age", IntegerType(), True),
    StructField("department", StringType(), True),
])

df = spark.createDataFrame(data, schema)

# Programmatic schema from CSV
df = spark.read.csv("data.csv", header=True, schema=schema)

# Print schema
df.printSchema()

Common Operations

Selecting and Filtering

# Select columns
df.select("name", "age")
df.select(F.col("name"), F.col("age"))

# Select with expressions
df.selectExpr("name", "age + 1 as next_age")

# Filter / Where
df.filter(F.col("age") > 25)
df.filter("age > 25")
df.where((F.col("age") > 25) & (F.col("department") == "Engineering"))
df.where(F.col("name").isin(["Alice", "Bob"]))
df.where(F.col("name").like("%lic%"))
df.where(F.col("name").startswith("A"))
df.where(F.col("name").rlike("^[A-C]"))

Adding and Modifying Columns

# Add column
df.withColumn("country", F.lit("US"))

# Modify column
df.withColumn("age", F.col("age") + 1)

# Conditional column
df.withColumn("seniority",
    F.when(F.col("age") > 30, "Senior")
     .when(F.col("age") > 20, "Mid")
     .otherwise("Junior")
)

# Rename column
df.withColumnRenamed("name", "first_name")

# Drop column
df.drop("country")

Aggregations

from pyspark.sql.functions import sum, avg, count, min, max, collect_list, collect_set

# Single aggregation
df.groupBy("department").agg(
    F.avg("salary").alias("avg_salary"),
    F.count("*").alias("count")
)

# Multiple aggregations
df.groupBy("department", "level").agg(
    F.sum("salary").alias("total"),
    F.countDistinct("employee_id").alias("num_employees"),
    F.collect_list("name").alias("employees")
)

# Pivot
df.groupBy("department").pivot("year").sum("salary")

Joins

# Inner join
df1.join(df2, df1.id == df2.id, "inner")

# Join on multiple columns
df1.join(df2, ["id", "year"], "inner")

# Left join
df1.join(df2, "id", "left")

# Right join
df1.join(df2, "id", "right")

# Full outer join
df1.join(df2, "id", "outer")

# Cross join
df1.crossJoin(df2)

# Semi join (returns left rows that have a match)
df1.join(df2, "id", "left_semi")

# Anti join (returns left rows with no match)
df1.join(df2, "id", "left_anti")

# Broadcast join (for small tables)
df1.join(F.broadcast(df2), "id")

Sorting

# Sort ascending
df.orderBy("age")

# Sort descending
df.orderBy(F.col("age").desc())

# Multiple sort keys
df.orderBy(F.col("department"), F.col("age").desc())

Window Functions

from pyspark.sql.window import Window

window = Window.partitionBy("department").orderBy("salary")

df.withColumn("rank", F.rank().over(window))
df.withColumn("dense_rank", F.dense_rank().over(window))
df.withColumn("row_number", F.row_number().over(window))
df.withColumn("percent_rank", F.percent_rank().over(window))
df.withColumn("ntile", F.ntile(4).over(window))

# Lag / Lead
df.withColumn("prev_salary", F.lag("salary", 1).over(window))
df.withColumn("next_salary", F.lead("salary", 1).over(window))

# Rolling window
window_7d = Window.orderBy("date").rowsBetween(-6, 0)
df.withColumn("ma_7d", F.avg("amount").over(window_7d))

Handling Nulls

# Drop nulls
df.dropna()
df.dropna(subset=["name", "email"])

# Fill nulls
df.fillna(0)
df.fillna({"name": "Unknown", "age": 0})

# Conditional null handling
df.withColumn("age", F.coalesce(F.col("age"), F.lit(0)))

# Check for nulls
df.where(F.col("name").isNull())
df.where(F.col("name").isNotNull())

SQL Queries

# Register temp view
df.createOrReplaceTempView("employees")

# Run SQL
spark.sql("""
    SELECT department, AVG(salary) as avg_salary
    FROM employees
    WHERE age > 25
    GROUP BY department
    ORDER BY avg_salary DESC
""").show()

UDFs (User Defined Functions)

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType

# Regular UDF
@udf(returnType=StringType())
def uppercase(name):
    return name.upper()

df.withColumn("name_upper", uppercase("name"))

# Pandas UDF (vectorized, faster)
from pyspark.sql.functions import pandas_udf
import pandas as pd

@pandas_udf(IntegerType())
def add_one(s: pd.Series) -> pd.Series:
    return s + 1

df.withColumn("age_plus_one", add_one("age"))

Partitions and Caching

# Check partitions
df.rdd.getNumPartitions()

# Repartition (full shuffle)
df.repartition(10)
df.repartition("department")

# Coalesce (reduce partitions, no shuffle)
df.coalesce(5)

# Cache / Persist
df.cache()
df.persist(StorageLevel.MEMORY_AND_DISK)
df.unpersist()

# Write partitioned
df.write.partitionBy("year", "month").parquet("output/")
df.write.mode("overwrite").partitionBy("department").parquet("output/")

Practical Examples

Example 1 — Data cleaning pipeline

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("Cleaning").getOrCreate()

# Read raw data
df = spark.read.csv("raw_data.csv", header=True, inferSchema=True)

# Clean
cleaned = df \
    .dropDuplicates(["id"]) \
    .filter(F.col("email").isNotNull()) \
    .withColumn("email", F.lower(F.trim(F.col("email")))) \
    .withColumn("age", F.coalesce(F.col("age"), F.lit(0))) \
    .withColumn("name", F.initcap(F.col("name"))) \
    .withColumn("processed_at", F.current_timestamp())

# Write
cleaned.write.mode("overwrite").parquet("cleaned_data.parquet")

Example 2 — Aggregation with pivot

# Monthly revenue by region
monthly = df \
    .withColumn("month", F.month("date")) \
    .groupBy("region", "month") \
    .agg(F.sum("amount").alias("revenue"))

# Pivot to wide format
pivot = monthly.groupBy("region").pivot("month").sum("revenue")
pivot.show()

Example 3 — Complex join with aggregation

# Top customer per region
from pyspark.sql.window import Window

window = Window.partitionBy("region").orderBy(F.col("total_spent").desc())

result = customers.join(orders, "customer_id") \
    .groupBy("customer_id", "name", "region") \
    .agg(F.sum("amount").alias("total_spent")) \
    .withColumn("rank", F.row_number().over(window)) \
    .filter(F.col("rank") == 1) \
    .drop("rank")

Common Patterns

Pattern: deduplication

# Remove exact duplicates
df.dropDuplicates()

# Remove duplicates based on specific columns
df.dropDuplicates(["id"])

# Keep first occurrence
window = Window.partitionBy("id").orderBy(F.col("created_at").desc())
df.withColumn("rn", F.row_number().over(window)) \
  .filter(F.col("rn") == 1) \
  .drop("rn")

Pattern: string operations

# Common string operations
df.withColumn("name_upper", F.upper("name"))
df.withColumn("name_lower", F.lower("name"))
df.withColumn("name_init", F.initcap("name"))
df.withColumn("name_trim", F.trim("name"))
df.withColumn("name_len", F.length("name"))
df.withColumn("name_sub", F.substring("name", 1, 3))
df.withColumn("name_replace", F.regexp_replace("name", " ", "_"))
df.withColumn("name_split", F.split("name", " "))

Pattern: date operations

# Current date/time
df.withColumn("today", F.current_date())
df.withColumn("now", F.current_timestamp())

# Extract components
df.withColumn("year", F.year("date"))
df.withColumn("month", F.month("date"))
df.withColumn("day", F.dayofmonth("date"))
df.withColumn("weekday", F.dayofweek("date"))

# Date arithmetic
df.withColumn("next_week", F.date_add("date", 7))
df.withColumn("prev_week", F.date_sub("date", 7))
df.withColumn("days_diff", F.datediff(F.current_date(), "date"))

Best Practices

  • Use DataFrames over RDDs — Catalyst optimizer makes DataFrames faster.
  • Prefer built-in functions over Python UDFs — native functions are serialized efficiently.
  • Use Pandas UDFs when you must use custom logic — vectorized operations are faster.
  • Filter early — reduce data volume before expensive operations.
  • Cache only when needed — cache only if you reuse a DataFrame 3+ times.
  • Use broadcast joins for small tables (< 10MB) — avoids shuffle.
  • Write to Parquet — columnar, compressed, schema-preserving.
  • Partition output by commonly filtered columns.

Common Mistakes

Mistake Fix
collect() on large data Use show(), take(), or write output
Not caching reused DataFrames Call .cache() before multiple actions
Using Python UDFs over native functions Prefer pyspark.sql.functions
Too many partitions Aim for 128MB–256MB per partition
Filtering after aggregations Filter before groupBy when possible
Ignoring data skew Use salting or broadcast joins
Not repartitioning before write Repartition for parallel write

Troubleshooting

Symptom Cause / Fix
OutOfMemoryError Reduce partition size; increase executor memory
Slow performance Check data skew; repartition; cache
Job failed with NPE Null handling — check for null values
Shuffle overflow Increase spark.sql.shuffle.partitions
Serialization error Use mapPartitions for heavy libraries
Data skew Use salting technique or broadcast joins

Official Documentation

Last Updated

2026-09-06 — Initial version.