SQL Optimization for Data Engineers

A slow query is usually slow for a visible reason, and the database will tell you what it is if you ask. Query optimisation is not a bag of tricks; it is a loop: read the plan, find where the work goes, change the query or the indexes, and measure again. This guide runs that loop on 2.9 million real rows, measuring the work each version of a query does.

The optimiser and the plan

You describe what you want in SQL; the engine decides how to get it. In SQL Server, the Query Optimizer takes the query, the schema and the database statistics and builds an execution plan: the order tables are accessed, the method for each, such as an index or a table scan, and how rows are filtered, aggregated and sorted (Microsoft Learn: execution plans). Every serious engine has an equivalent: EXPLAIN in PostgreSQL and MySQL, the graphical plan in SQL Server Management Studio, EXPLAIN QUERY PLAN in SQLite, which reports whether each table is read by a full SCAN or an indexed SEARCH (SQLite docs: EXPLAIN QUERY PLAN).

One query, four ways

The test loads January 2024 New York taxi trips (NYC TLC: trip record data) into SQLite and asks for the trips and fares from one pickup zone, JFK Airport, on one day. Timings vary from run to run, so the script counts the steps SQLite's virtual machine executes, a repeatable measure of work, and prints the plan beside it:

import sqlite3
import urllib.request
from pathlib import Path
import pandas as pd

name = "yellow_tripdata_2024-01.parquet"
if not Path(name).exists():
    urllib.request.urlretrieve(f"https://d37ci6vzurychx.cloudfront.net/trip-data/{name}", name)
trips = pd.read_parquet(name, columns=["tpep_pickup_datetime", "PULocationID", "payment_type", "trip_distance", "fare_amount"])
trips["tpep_pickup_datetime"] = trips["tpep_pickup_datetime"].dt.strftime("%Y-%m-%d %H:%M:%S")
trips.columns = ["pickup", "zone", "payment", "distance", "fare"]
con = sqlite3.connect(":memory:")
trips.to_sql("trips", con, index=False)
print(f"{len(trips)} rows loaded")

def measure(label, sql):
    steps = 0
    def tick():
        nonlocal steps
        steps += 1
    plan = " / ".join(row[3] for row in con.execute("EXPLAIN QUERY PLAN " + sql))
    con.set_progress_handler(tick, 1000)                 # called every 1,000 virtual-machine steps
    result = con.execute(sql).fetchall()
    con.set_progress_handler(None, 0)
    print(f"{label:<34} {steps:>7}k steps | {plan}")
    return result

by_zone = "SELECT COUNT(*), ROUND(SUM(fare), 2) FROM trips WHERE zone = 132 AND pickup >= '2024-01-15' AND pickup < '2024-01-16'"
a = measure("1. no index", by_zone)
con.execute("CREATE INDEX ix_zone_pickup ON trips (zone, pickup)")
b = measure("2. index on (zone, pickup)", by_zone)
con.execute("CREATE INDEX ix_zone_pickup_fare ON trips (zone, pickup, fare)")
c = measure("3. covering index adds fare", by_zone)
d = measure("4. function on the column", by_zone.replace("pickup >= '2024-01-15' AND pickup < '2024-01-16'", "date(pickup) = '2024-01-15'"))
print("same answer every time:", a == b == c == d, a)

2964624 rows loaded
1. no index                           9354k steps | SCAN trips
2. index on (zone, pickup)              39k steps | SEARCH trips USING INDEX ix_zone_pickup (zone=? AND pickup>? AND pickup<?)
3. covering index adds fare             34k steps | SEARCH trips USING COVERING INDEX ix_zone_pickup_fare (zone=? AND pickup>? AND pickup<?)
4. function on the column              748k steps | SEARCH trips USING COVERING INDEX ix_zone_pickup_fare (zone=?)
same answer every time: True [(5651, 345688.95)]

What each version shows

1. No index: a full scan. The plan says SCAN: every one of the 2,964,624 rows is read to find 5,651. Scan cost grows in proportion to the table (SQLite docs: query planning), so this query gets slower every month.

2. A composite index on zone and pickup time. The plan becomes SEARCH, and the work drops from 9,354 thousand steps to 39 thousand, about 240 times less. The column order matters: the left-most column orders the index and later columns break ties, so putting the equality column, zone, first and the range column, pickup, second lets the engine jump to JFK and read one contiguous day.

3. A covering index. Adding fare to the index means the query is answered from the index alone, USING COVERING INDEX, with no visits to the table. The saving here is modest, 39 to 34 thousand steps, because only 5,651 rows match; it grows with the number of matching rows.

4. A function on the filtered column. Writing date(pickup) = '2024-01-15' asks the same question and returns the same answer, but the plan now uses the index only for the zone. The engine cannot use an index on pickup to find values of a function of pickup, so it reads every JFK row for the month and computes the date for each: 748 thousand steps, about twenty times the range version. Predicates the engine can match to an index are called sargable; keeping the bare column on one side of the comparison keeps them that way.

All four return identical results. Optimisation changes the work, never the answer, and checking that is part of the method.

Habits that follow

• Index for your WHERE and JOIN columns, equality columns first, then the range column.
• Keep predicates sargable. Compare the column to a range instead of wrapping it: pickup >= '2024-01-15' AND pickup < '2024-01-16', not date(pickup) = ...; name LIKE 'Sm%', not LIKE '%son'; and avoid implicit type conversions between text and numbers.
• Select only the columns you need. SELECT * defeats covering indexes and moves more data.
• Filter early, before joins and window functions, so later steps handle fewer rows; window functions for data engineers notes the same for sorts.
• Keep statistics current. The optimiser chooses plans from statistics; stale ones produce bad plans.
• Do not over-index. Every index must be updated on each insert and takes space; SQLite's documentation advises dropping an index that is a prefix of another.

Beyond row stores

The same thinking applies in analytical engines, with different tools. Columnar warehouses and Spark rarely use B-tree indexes; they prune with partitioning, clustering or sort order and per-file statistics, and the plan shows whether a filter was pushed down to the scan. A function wrapped around a partition column defeats pruning exactly as it defeated the index above. Spark performance tuning shows plan reading there, and designing fact and dimension tables covers the modelling choices that make queries simple in the first place.

A tuning checklist

1. Get the actual plan, not a guess.
2. Find the largest scan, sort or join in it.
3. Check the predicates on that table are sargable and indexed or prunable.
4. Change one thing, confirm the result is unchanged, and measure again.
5. Stop when the query is fast enough; indexes and rewrites have costs too.

Query optimization, indexing techniques and execution plans are part of the performance optimization module of the SQL Server course, and query optimization for pipelines is covered in the Data Engineering course.

Sources

• Query planning (SQLite documentation) — SQLite
• EXPLAIN QUERY PLAN (SQLite documentation) — SQLite
• Execution plan overview (SQL Server, Microsoft Learn) — Microsoft Learn
• TLC Trip Record Data (NYC Taxi and Limousine Commission) — NYC Taxi and Limousine Commission