PySpark Tutorial: DataFrames, Transformations and Actions

PySpark is Spark's Python API, and almost all everyday PySpark is the DataFrame API: read a table, transform it, aggregate it, write or show it. This tutorial works through that API on a real dataset, January 2024 New York yellow taxi trips, nearly three million rows, and every result below is real output. If Spark's architecture is new to you, read Apache Spark explained for beginners first; this article assumes you know that transformations are lazy and actions run them.

Setup and reading data with a schema

pip install pyspark is enough to run Spark locally. A SparkSession is the entry point to everything (Spark docs: getting started). Parquet files carry their own schema; CSV files do not, so the zone lookup gets an explicit StructType instead of a costly inferSchema pass. The data comes from the city's Taxi and Limousine Commission (NYC TLC: trip record data).

import urllib.request
from pathlib import Path
from pyspark.sql import SparkSession, Window, functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, StringType

base = "https://d37ci6vzurychx.cloudfront.net"
for name, url in {"yellow_tripdata_2024-01.parquet": f"{base}/trip-data/yellow_tripdata_2024-01.parquet",
                  "taxi_zone_lookup.csv": f"{base}/misc/taxi_zone_lookup.csv"}.items():
    if not Path(name).exists():
        urllib.request.urlretrieve(url, name)

spark = SparkSession.builder.master("local[4]").appName("pyspark-tutorial").getOrCreate()
spark.sparkContext.setLogLevel("ERROR")

zone_schema = StructType([StructField("LocationID", IntegerType()), StructField("Borough", StringType()),
                          StructField("Zone", StringType()), StructField("service_zone", StringType())])
zones = spark.read.csv("taxi_zone_lookup.csv", header=True, schema=zone_schema)
trips = spark.read.parquet("yellow_tripdata_2024-01.parquet")

Filtering and deriving columns

Most transformation work is where to keep rows, withColumn to add or replace a column, and select to keep only what you need. Conditional logic uses when and otherwise, the DataFrame form of SQL's CASE. The payment codes come from the TLC's data dictionary (TLC: yellow taxi data dictionary):

payment = (F.when(F.col("payment_type") == 1, "credit card").when(F.col("payment_type") == 2, "cash")
            .when(F.col("payment_type") == 0, "flex fare").otherwise("other"))
clean = (trips.where((F.col("fare_amount") > 0) & (F.col("tpep_dropoff_datetime") > F.col("tpep_pickup_datetime")))
              .withColumn("payment", payment)
              .withColumn("minutes", (F.unix_timestamp("tpep_dropoff_datetime") - F.unix_timestamp("tpep_pickup_datetime")) / 60)
              .select("PULocationID", "payment", "fare_amount", "tip_amount", "minutes"))

Two habits worth copying: use functions from pyspark.sql.functions rather than Python UDFs, because built-ins run inside Spark's engine, and select the columns you need early, so less data moves through the plan.

Grouping and aggregating

groupBy followed by agg computes several aggregates in one pass:

(clean.groupBy("payment")
      .agg(F.count("*").alias("trips"),
           F.round(F.avg("fare_amount"), 2).alias("avg_fare"),
           F.round(100 * F.sum("tip_amount") / F.sum("fare_amount"), 1).alias("tip_pct_of_fare"))
      .orderBy(F.desc("trips"))
      .show())

+-----------+-------+--------+---------------+
|    payment|  trips|avg_fare|tip_pct_of_fare|
+-----------+-------+--------+---------------+
|credit card|2318753|   18.55|           22.5|
|       cash| 430058|   18.63|            0.0|
|  flex fare| 137935|   20.34|            7.7|
|      other|  38737|   18.47|            0.0|
+-----------+-------+--------+---------------+

Cash passengers apparently tip nothing, which is not a finding about passengers. The data dictionary explains it: tip_amount is populated automatically for credit card tips, and cash tips are not included. An analysis of tipping over all trips would understate tips badly; restrict it to card payments. Code that runs correctly can still answer the wrong question, so read the dictionary before trusting an aggregate.

Joining and ranking with a window function

Joins look like SQL joins. Window functions compute a value for each row from a group of related rows, defined by a partition and an ordering (Spark SQL reference: window functions). Here, the two busiest pickup zones in each borough:

by_zone = clean.groupBy("PULocationID").count().join(zones, F.col("PULocationID") == F.col("LocationID"))
rank = Window.partitionBy("Borough").orderBy(F.desc("count"), "Zone")
(by_zone.withColumn("rank", F.row_number().over(rank))
        .where((F.col("rank") <= 2) & ~F.col("Borough").isin("Unknown", "N/A"))
        .select("Borough", "rank", "Zone", F.col("count").alias("pickups"))
        .orderBy("Borough", "rank")
        .show(truncate=False))

+-------------+----+---------------------------+-------+
|Borough      |rank|Zone                       |pickups|
+-------------+----+---------------------------+-------+
|Bronx        |1   |Mott Haven/Port Morris     |649    |
|Bronx        |2   |Co-Op City                 |408    |
|Brooklyn     |1   |Downtown Brooklyn/MetroTech|1369   |
|Brooklyn     |2   |Brooklyn Heights           |1237   |
|EWR          |1   |Newark Airport             |283    |
|Manhattan    |1   |Midtown Center             |141699 |
|Manhattan    |2   |Upper East Side South      |141225 |
|Queens       |1   |JFK Airport                |141060 |
|Queens       |2   |LaGuardia Airport          |88540  |
|Staten Island|1   |Arrochar/Fort Wadsworth    |21     |
|Staten Island|2   |Heartland Village/Todt Hill|9      |
+-------------+----+---------------------------+-------+

Two details matter. row_number gives each row a unique rank, so adding Zone as a second ordering column makes ties break the same way on every run; rank would give tied rows the same number. And the result shows how concentrated yellow taxi pickups are: Manhattan's top zone has more than two hundred times as many pickups as the Bronx's.

The same logic in Spark SQL

Any DataFrame can be registered as a temporary view, which lasts for the session, and queried with SQL; the result is another DataFrame:

clean.createOrReplaceTempView("trips")
spark.sql("""SELECT payment, ROUND(PERCENTILE_APPROX(minutes, 0.5), 1) AS median_minutes
             FROM trips GROUP BY payment ORDER BY payment""").show()
print("clean trips:", clean.count())
spark.stop()

+-----------+--------------+
|    payment|median_minutes|
+-----------+--------------+
|       cash|          11.2|
|credit card|          11.7|
|  flex fare|          14.0|
|      other|           8.1|
+-----------+--------------+

clean trips: 2925483

SQL and the DataFrame API compile to the same plans, so choose whichever reads better for each step. Teams often write joins and aggregations in SQL and programmatic logic, such as loops over columns, in Python. PERCENTILE_APPROX computes an approximate median, far cheaper than an exact one on large data.

Transformations versus actions in this tutorial

Everything above except four calls was a transformation. The actions were the three show() calls and count(), and each ran the plan from the source files. In a real job that would reread the Parquet file four times; caching clean with clean.cache() before the first action would keep it in memory for the others. Spark performance tuning covers when caching helps and when it hurts.

Errors you will meet, and what they mean

• AnalysisException about an unresolved column: a misspelt column name, or one dropped by an earlier select. Spark checks names when you define the transformation, so this error appears early.
• Ambiguous column references after a join: both sides had a column with the same name. Rename one side first, or join on a list of column names so Spark keeps a single copy.
• Errors that surface only at an action: problems in the data itself, such as a failed type conversion, appear when show() or write runs, possibly several lines after the step that caused them.
• Out-of-memory on the driver: usually a collect() or toPandas() on too much data. Aggregate first, or write the result instead.

df.explain() prints the plan Spark will run, which is the first place to look when a result or a runtime surprises you.

Writing results

In production the last step is usually a write, for example df.write.mode("overwrite").partitionBy("payment").parquet(path). Partitioning output by a column lets later queries skip whole folders; partitioning and bucketing in Spark explains the trade-offs. On Windows, writing files from local Spark needs extra Hadoop setup, which is one reason many people learn on Linux, in a container or on a managed platform.

Where next

With these calls, select, where, withColumn, groupBy with agg, join, Window, spark.sql and the actions, you can write most batch transformations. Next, compare the API with pandas in PySpark vs pandas, and practise window functions in SQL, where the same ideas apply, with window functions for data engineers.

PySpark basics, DataFrames and RDDs, Spark SQL and performance optimization are part of the Apache Spark module of the Data Engineering course.

Sources

• Spark SQL getting started (Apache Spark documentation) — Apache Spark
• Window functions (Spark SQL reference) — Apache Spark
• Data dictionary: yellow taxi trip records (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission
• TLC Trip Record Data (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission