{ "cells": [ { "cell_type": "markdown", "id": "3204c9d6", "metadata": {}, "source": [ "# Spark on a Laptop: 22.6 Million Taxi Trips\n", "\n", "Companion notebook to the blog post. Everything runs locally on one machine.\n", "\n", "The data is public NYC Taxi & Limousine Commission trip records. The notebook downloads it\n", "for you (315 MB, three months of 2019).\n", "\n", "## Setup\n", "\n", "You need Python and about 800 MB of disk. No Java installation required, `jdk4py` handles it.\n", "\n", "```bash\n", "python3 -m venv ~/sparkenv\n", "~/sparkenv/bin/python -m pip install pyspark==3.5.3 jdk4py ipykernel\n", "~/sparkenv/bin/python -m ipykernel install --user --name pyspark-env --display-name \"Python (PySpark)\"\n", "```\n", "\n", "Then pick the `Python (PySpark)` kernel (in VS Code: kernel picker -> Jupyter Kernel)." ] }, { "cell_type": "markdown", "id": "1fdb0346", "metadata": {}, "source": [ "## 1. Point Python at Java\n", "\n", "Spark runs on the JVM, so it needs `JAVA_HOME` set before the session starts." ] }, { "cell_type": "code", "execution_count": 1, "id": "424b6bb1", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:01.295939Z", "iopub.status.busy": "2026-08-05T00:24:01.295212Z", "iopub.status.idle": "2026-08-05T00:24:01.445366Z", "shell.execute_reply": "2026-08-05T00:24:01.444857Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "pyspark 3.5.3\n" ] } ], "source": [ "import os, glob, time, urllib.request\n", "import jdk4py\n", "\n", "os.environ[\"JAVA_HOME\"] = str(jdk4py.JAVA_HOME)\n", "os.environ[\"PATH\"] = str(jdk4py.JAVA_HOME / \"bin\") + os.pathsep + os.environ[\"PATH\"]\n", "\n", "import pyspark\n", "from pyspark.sql import SparkSession\n", "from pyspark.sql.functions import *\n", "print(\"pyspark\", pyspark.__version__)" ] }, { "cell_type": "markdown", "id": "872d7449", "metadata": {}, "source": [ "## 2. Get the data\n", "\n", "Three months of yellow taxi trips, straight from the TLC. Skips anything already downloaded." ] }, { "cell_type": "code", "execution_count": 2, "id": "fbd7c4bf", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:01.447008Z", "iopub.status.busy": "2026-08-05T00:24:01.446876Z", "iopub.status.idle": "2026-08-05T00:24:01.462961Z", "shell.execute_reply": "2026-08-05T00:24:01.462649Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "have yellow_tripdata_2019-01.parquet\n", "have yellow_tripdata_2019-02.parquet\n", "have yellow_tripdata_2019-03.parquet\n" ] } ], "source": [ "DATA = os.path.expanduser(\"~/nyctaxi/\")\n", "os.makedirs(DATA, exist_ok=True)\n", "BASE = \"https://d37ci6vzurychx.cloudfront.net/trip-data/\"\n", "\n", "for month in [\"2019-01\", \"2019-02\", \"2019-03\"]:\n", " fname = f\"yellow_tripdata_{month}.parquet\"\n", " dest = os.path.join(DATA, fname)\n", " if os.path.exists(dest):\n", " print(f\"have {fname}\")\n", " else:\n", " print(f\"downloading {fname} ...\", end=\" \", flush=True)\n", " urllib.request.urlretrieve(BASE + fname, dest)\n", " print(f\"{os.path.getsize(dest)/1048576:.0f} MB\")" ] }, { "cell_type": "markdown", "id": "b062deaa", "metadata": {}, "source": [ "## 3. Start Spark\n", "\n", "`local[4]` runs four worker threads on this machine. On a cluster you would name the cluster here\n", "instead, and nothing else in this notebook would change." ] }, { "cell_type": "code", "execution_count": 3, "id": "468f4821", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:01.464321Z", "iopub.status.busy": "2026-08-05T00:24:01.464223Z", "iopub.status.idle": "2026-08-05T00:24:03.592305Z", "shell.execute_reply": "2026-08-05T00:24:03.591984Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Setting default log level to \"WARN\".\n", "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n", "26/08/04 17:24:02 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "26/08/04 17:24:03 WARN Utils: Service 'SparkUI' could not bind on port 4040. Attempting port 4041.\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Spark 3.5.3\n" ] } ], "source": [ "spark = (SparkSession.builder\n", " .master(\"local[4]\")\n", " .appName(\"taxi\")\n", " .config(\"spark.driver.memory\", \"3g\")\n", " .config(\"spark.sql.shuffle.partitions\", \"8\")\n", " .getOrCreate())\n", "spark.sparkContext.setLogLevel(\"ERROR\")\n", "print(\"Spark\", spark.version)" ] }, { "cell_type": "markdown", "id": "937aab81", "metadata": {}, "source": [ "## 4. Load\n", "\n", "`.parquet()` only builds a plan. `.count()` is what forces Spark to actually do work." ] }, { "cell_type": "code", "execution_count": 4, "id": "6a7c5d1d", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:03.594528Z", "iopub.status.busy": "2026-08-05T00:24:03.594399Z", "iopub.status.idle": "2026-08-05T00:24:05.703633Z", "shell.execute_reply": "2026-08-05T00:24:05.703233Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "22,612,607 rows x 19 columns in 2.1s\n" ] } ], "source": [ "t = time.time()\n", "df_raw = spark.read.parquet(DATA + \"*.parquet\")\n", "n = df_raw.count()\n", "print(f\"{n:,} rows x {len(df_raw.columns)} columns in {time.time()-t:.1f}s\")" ] }, { "cell_type": "markdown", "id": "9db2b327", "metadata": {}, "source": [ "## 5. The analysis functions\n", "\n", "Each takes a dataframe and returns a dataframe." ] }, { "cell_type": "code", "execution_count": 5, "id": "e6afaa73", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.705382Z", "iopub.status.busy": "2026-08-05T00:24:05.705260Z", "iopub.status.idle": "2026-08-05T00:24:05.708802Z", "shell.execute_reply": "2026-08-05T00:24:05.708447Z" } }, "outputs": [], "source": [ "def clean_data(df):\n", " '''\n", " input: df a dataframe\n", " output: df a dataframe with the all the original columns\n", " '''\n", " \n", " # START YOUR CODE HERE ---------\n", " df2 = df.withColumn(\"passenger_count\",df.passenger_count.cast('int'))\\\n", " .withColumn(\"total_amount\",df.total_amount.cast('float'))\\\n", " .withColumn(\"tip_amount\",df.tip_amount.cast('float'))\\\n", " .withColumn(\"trip_distance\",df.trip_distance.cast('float'))\\\n", " .withColumn(\"fare_amount\",df.fare_amount.cast('float'))\\\n", " .withColumn(\"tpep_pickup_datetime\",df.tpep_pickup_datetime.cast('timestamp'))\\\n", " .withColumn(\"tpep_dropoff_datetime\",df.tpep_dropoff_datetime.cast('timestamp'))\n", " # END YOUR CODE HERE -----------\n", " return df2" ] }, { "cell_type": "code", "execution_count": 6, "id": "58a32bfb", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.710101Z", "iopub.status.busy": "2026-08-05T00:24:05.710008Z", "iopub.status.idle": "2026-08-05T00:24:05.713070Z", "shell.execute_reply": "2026-08-05T00:24:05.712819Z" } }, "outputs": [], "source": [ "def common_pair(df):\n", " '''\n", " input: df a dataframe\n", " output: df a dataframe with following columns:\n", " - PULocationID\n", " - DOLocationID\n", " - count\n", " - trip_rate\n", " \n", " trip_rate is the average amount (total_amount) per distance (trip_distance)\n", " \n", " '''\n", " \n", " # START YOUR CODE HERE ---------\n", " \n", " df = df.groupBy(\"PULocationID\", \"DOLocationID\")\\\n", " .agg(sum(\"total_amount\").alias(\"total_amount\"), \\\n", " sum(\"trip_distance\").alias(\"trip_distance\"), \\\n", " count(\"PULocationID\").alias(\"count\")\\\n", " ) \n", " df = df.withColumn(\"trip_rate\", df.total_amount / df.trip_distance)\n", " \n", " df = df.sort(desc(\"count\"),desc(\"trip_rate\")).limit(10)\n", " # END YOUR CODE HERE -----------\n", " \n", " return df" ] }, { "cell_type": "code", "execution_count": 7, "id": "1a504930", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.714312Z", "iopub.status.busy": "2026-08-05T00:24:05.714230Z", "iopub.status.idle": "2026-08-05T00:24:05.716992Z", "shell.execute_reply": "2026-08-05T00:24:05.716750Z" } }, "outputs": [], "source": [ "def time_of_cheapest_fare(df):\n", " '''\n", " input: df a dataframe\n", " output: df a dataframe with following columns:\n", " - day_night\n", " - trip_rate\n", " \n", " day_night will have 'Day' or 'Night' based on following conditions:\n", " - From 9am to 8:59:59pm - Day\n", " - From 9pm to 8:59:59am - Night\n", " \n", " trip_rate is the average amount (total_amount) per distance\n", " \n", " '''\n", " \n", " # START YOUR CODE HERE ---------\n", " df = df.withColumn(\"day_night\", \\\n", " when((hour(df.tpep_pickup_datetime) >= 9) \\\n", " & (hour(df.tpep_pickup_datetime) <= 20) \\\n", " & (minute(df.tpep_pickup_datetime) <= 59) \\\n", " & (second(df.tpep_pickup_datetime) <= 59), 'Day') \\\n", " .otherwise('Night'))\n", " df = df.groupBy(\"day_night\")\\\n", " .agg(sum(\"total_amount\").alias(\"total_amount\"), \\\n", " sum(\"trip_distance\").alias(\"trip_distance\"), \\\n", " ) \n", " df = df.withColumn(\"trip_rate\", df.total_amount / df.trip_distance) \n", " # END YOUR CODE HERE -----------\n", " \n", " return df" ] }, { "cell_type": "code", "execution_count": 8, "id": "7a707f65", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.718125Z", "iopub.status.busy": "2026-08-05T00:24:05.718046Z", "iopub.status.idle": "2026-08-05T00:24:05.720490Z", "shell.execute_reply": "2026-08-05T00:24:05.720226Z" } }, "outputs": [], "source": [ "def passenger_count_for_most_tip(df):\n", " '''\n", " input: df a dataframe\n", " output: df a dataframe with following columns:\n", " - passenger_count\n", " - tip_percent\n", " \n", " trip_percent is the percent of tip out of fare_amount\n", " \n", " '''\n", " \n", " # START YOUR CODE HERE ---------\n", " df = df.filter((df.fare_amount > 2) & (df.passenger_count > 0))\n", " df = df.groupBy(\"passenger_count\")\\\n", " .agg(sum(\"tip_amount\").alias(\"tip_amount\"), \\\n", " sum(\"fare_amount\").alias(\"fare_amount\"), \\\n", " ) \n", " df = df.withColumn(\"tip_percent\", \\\n", " df.tip_amount * 100 / df.fare_amount)\n", " df = df.sort(desc(\"tip_percent\"))\n", " # END YOUR CODE HERE -----------\n", " \n", " return df" ] }, { "cell_type": "code", "execution_count": 9, "id": "979a73bb", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.721660Z", "iopub.status.busy": "2026-08-05T00:24:05.721573Z", "iopub.status.idle": "2026-08-05T00:24:05.724277Z", "shell.execute_reply": "2026-08-05T00:24:05.723986Z" } }, "outputs": [], "source": [ "def day_with_traffic(df):\n", " '''\n", " input: df a dataframe\n", " output: df a dataframe with following columns:\n", " - day_of_week\n", " - average_speed\n", " \n", " day_of_week should be day of week e.g.) Mon, Tue, Wed, ...\n", " average_speed (miles/hour) is calculated as distance / time (in hours)\n", " \n", " '''\n", " \n", " # START YOUR CODE HERE ---------\n", " \n", " df = df.withColumn(\"day_of_week\",date_format(df.tpep_pickup_datetime, 'E'))\n", " df = df.withColumn(\"time\", (col(\"tpep_dropoff_datetime\").cast(\"long\") - col('tpep_pickup_datetime').cast(\"long\")) / 3600)\n", " \n", " \n", " df = df.groupBy(\"day_of_week\")\\\n", " .agg(sum(\"time\").alias(\"time\"), \\\n", " sum(\"trip_distance\").alias(\"trip_distance\"))\n", " \n", " df = df.withColumn(\"average_speed\",df.trip_distance / df.time)\n", "\n", " \n", " # END YOUR CODE HERE -----------\n", " df = df.sort(asc(\"average_speed\"), asc(\"day_of_week\"))\n", " return df" ] }, { "cell_type": "markdown", "id": "014b064e", "metadata": {}, "source": [ "## 6. Lazy evaluation\n", "\n", "`clean_data` casts seven columns across 22.6 million rows. Time it. It returns instantly,\n", "because it converts nothing — it only records what you intend to do." ] }, { "cell_type": "code", "execution_count": 10, "id": "2c313aec", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.725702Z", "iopub.status.busy": "2026-08-05T00:24:05.725613Z", "iopub.status.idle": "2026-08-05T00:24:05.899150Z", "shell.execute_reply": "2026-08-05T00:24:05.898796Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "clean_data returned in 0.071s\n", "...but .count() took 0.1s\n" ] } ], "source": [ "t = time.time(); df = clean_data(df_raw); print(f\"clean_data returned in {time.time()-t:.3f}s\")\n", "t = time.time(); df.count(); print(f\"...but .count() took {time.time()-t:.1f}s\")" ] }, { "cell_type": "markdown", "id": "75217ed7", "metadata": {}, "source": [ "## 7. Results\n", "\n", "Cost per mile, day versus night:" ] }, { "cell_type": "code", "execution_count": 11, "id": "a4b0c9c4", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:05.900681Z", "iopub.status.busy": "2026-08-05T00:24:05.900540Z", "iopub.status.idle": "2026-08-05T00:24:09.372429Z", "shell.execute_reply": "2026-08-05T00:24:09.371350Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 7:==============> (1 + 3) / 4]\r" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "+---------+-----------------+\n", "|day_night|trip_rate |\n", "+---------+-----------------+\n", "|Night |5.655919205504151|\n", "|Day |6.414801278986028|\n", "+---------+-----------------+\n", "\n" ] } ], "source": [ "time_of_cheapest_fare(df).select(\"day_night\",\"trip_rate\").show(truncate=False)" ] }, { "cell_type": "markdown", "id": "81cdaff1", "metadata": {}, "source": [ "Tip percentage by number of passengers:" ] }, { "cell_type": "code", "execution_count": 12, "id": "9b6dfeb0", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:09.375026Z", "iopub.status.busy": "2026-08-05T00:24:09.374898Z", "iopub.status.idle": "2026-08-05T00:24:10.272843Z", "shell.execute_reply": "2026-08-05T00:24:10.272432Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "+---------------+------------------+\n", "|passenger_count|tip_percent |\n", "+---------------+------------------+\n", "|5 |16.86834785679893 |\n", "|6 |16.75340968679461 |\n", "|1 |16.254314786129637|\n", "|2 |16.02281391894225 |\n", "|3 |15.937851499950584|\n", "|4 |15.080371520817177|\n", "|7 |12.879898566907126|\n", "|9 |12.13987210073056 |\n", "|8 |11.735485947863436|\n", "+---------------+------------------+\n", "\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 10:==============> (1 + 3) / 4]\r", "\r", " \r" ] } ], "source": [ "passenger_count_for_most_tip(df).select(\"passenger_count\",\"tip_percent\").show(10, False)" ] }, { "cell_type": "markdown", "id": "fe55e26b", "metadata": {}, "source": [ "Average speed by day of week — a proxy for traffic:" ] }, { "cell_type": "code", "execution_count": 13, "id": "179c7472", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:10.277226Z", "iopub.status.busy": "2026-08-05T00:24:10.277098Z", "iopub.status.idle": "2026-08-05T00:24:14.585680Z", "shell.execute_reply": "2026-08-05T00:24:14.585062Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 13:==============> (1 + 3) / 4]\r" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 13:=============================> (2 + 2) / 4]\r" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 13:===========================================> (3 + 1) / 4]\r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "+-----------+------------------+\n", "|day_of_week|average_speed |\n", "+-----------+------------------+\n", "|Thu |9.567853436366102 |\n", "|Fri |9.729021129695456 |\n", "|Sat |9.820350026501735 |\n", "|Wed |10.030324181882959|\n", "|Tue |10.169196504009705|\n", "|Mon |11.270629967612495|\n", "|Sun |11.655404282705545|\n", "+-----------+------------------+\n", "\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " \r" ] } ], "source": [ "day_with_traffic(df).select(\"day_of_week\",\"average_speed\").show(10, False)" ] }, { "cell_type": "markdown", "id": "12ccf1c3", "metadata": {}, "source": [ "The busiest pickup/dropoff pairs, with the average fare per mile on each:" ] }, { "cell_type": "code", "execution_count": 14, "id": "4b99df6a", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:14.587962Z", "iopub.status.busy": "2026-08-05T00:24:14.587826Z", "iopub.status.idle": "2026-08-05T00:24:15.561130Z", "shell.execute_reply": "2026-08-05T00:24:15.560580Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 16:==============> (1 + 3) / 4]\r", "\r", " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "+------------+------------+------------------+------------------+------+------------------+\n", "|PULocationID|DOLocationID|total_amount |trip_distance |count |trip_rate |\n", "+------------+------------+------------------+------------------+------+------------------+\n", "|264 |264 |4106117.585838042 |575267.4800102618 |230177|7.137753703311712 |\n", "|237 |236 |1403653.1713214517|145293.49008958042|139066|9.660812541952376 |\n", "|236 |236 |1102619.0395024419|79208.99017350748 |127131|13.920377435530389|\n", "|236 |237 |1277749.348745823 |124500.87005715072|118377|10.262975255990472|\n", "|237 |237 |1006906.8190342784|72624.73015023954 |110962|13.864517182387868|\n", "|239 |238 |626669.5557575226 |57126.08006504178 |68588 |10.969937987063322|\n", "|239 |142 |619390.8252599835 |55703.33004354313 |64321 |11.119457755502365|\n", "|142 |239 |620815.7451167107 |60728.67002433352 |62180 |10.22277854706771 |\n", "|238 |239 |540148.3849024773 |47453.8200550247 |59495 |11.38261122658097 |\n", "|161 |237 |655966.584115088 |61868.60002200492 |56876 |10.602576814115386|\n", "+------------+------------+------------------+------------------+------+------------------+\n", "\n" ] } ], "source": [ "common_pair(df).show(10, False)" ] }, { "cell_type": "markdown", "id": "8682a3ab", "metadata": {}, "source": [ "## 8. Column pruning\n", "\n", "Parquet stores data column by column, so Spark reads only the columns you ask for. Compare\n", "reading all 19 against the 9 this analysis actually uses." ] }, { "cell_type": "code", "execution_count": 15, "id": "aa8106bc", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:15.562667Z", "iopub.status.busy": "2026-08-05T00:24:15.562545Z", "iopub.status.idle": "2026-08-05T00:24:19.401526Z", "shell.execute_reply": "2026-08-05T00:24:19.400241Z" } }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 20:==============> (1 + 3) / 4]\r" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 20:=============================> (2 + 2) / 4]\r", "\r", " \r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "all 19 columns : 2.2s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", "[Stage 22:==============> (1 + 3) / 4]\r" ] }, { "name": "stdout", "output_type": "stream", "text": [ "only 9 columns : 1.6s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " \r" ] } ], "source": [ "cols = [\"passenger_count\",\"total_amount\",\"tip_amount\",\"trip_distance\",\"fare_amount\",\n", " \"tpep_pickup_datetime\",\"tpep_dropoff_datetime\",\"PULocationID\",\"DOLocationID\"]\n", "\n", "t = time.time(); spark.read.parquet(DATA + \"*.parquet\").write.format(\"noop\").mode(\"overwrite\").save()\n", "print(f\"all 19 columns : {time.time()-t:.1f}s\")\n", "t = time.time(); spark.read.parquet(DATA + \"*.parquet\").select(*cols).write.format(\"noop\").mode(\"overwrite\").save()\n", "print(f\"only 9 columns : {time.time()-t:.1f}s\")" ] }, { "cell_type": "markdown", "id": "ea870a54", "metadata": {}, "source": [ "## 9. Things to try\n", "\n", "- Change `local[4]` to `local[1]` and rerun. Watch the timings.\n", "- Add `df = df.cache()` after `clean_data` and rerun section 7. Spark stops re-reading from disk.\n", "- Add more months to the download list. The code does not change." ] }, { "cell_type": "code", "execution_count": 16, "id": "589c81d7", "metadata": { "execution": { "iopub.execute_input": "2026-08-05T00:24:19.404095Z", "iopub.status.busy": "2026-08-05T00:24:19.403953Z", "iopub.status.idle": "2026-08-05T00:24:19.545620Z", "shell.execute_reply": "2026-08-05T00:24:19.545157Z" } }, "outputs": [], "source": [ "spark.stop()" ] } ], "metadata": { "kernelspec": { "display_name": "Python (PySpark)", "language": "python", "name": "pyspark-env" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.6" } }, "nbformat": 4, "nbformat_minor": 5 }