Writing GTFS to Partitioned Parquet
Define an explicit pyarrow schema so string IDs, integer sequences, and text time fields are pinned regardless of the DataFrame’s inferred dtypes, attach agency_id and feed_date as partition columns, then write with pyarrow.dataset.write_dataset using existing_data_behavior="delete_matching" so re-runs overwrite exactly one partition. Read the data back with a filter on the partition columns and pyarrow prunes untouched directories before opening a single file — predicate pushdown that turns a full-archive scan into a single-partition read.
Root Cause Analysis
A quick DataFrame.to_parquet("stop_times.parquet") per feed works until you have a second feed version, at which point three problems surface — and each is what the script below is built to avoid.
Schema drift across versions. Letting pyarrow infer the schema from each DataFrame means the types follow whatever pandas happened to produce. A stop_id that was read without dtype=str writes as int64 in one feed and string in another; an all-null optional column writes as double one week and string the next. When a query engine tries to read the two files as one dataset, the incompatible schemas collide. Pinning an explicit schema, the same way the exporting GTFS to databases and warehouses guide pins typed DDL, makes every version mergeable.
No partitioning, no pruning. A single flat Parquet file per table forces every query to open the whole file even when it only wants one agency or one date. Without partition columns the reader cannot skip anything, so an archive of a hundred feed versions is scanned in full for a question about one day.
Non-idempotent reloads. Writing with a timestamped filename each run silently accumulates duplicate rows for the same feed version; overwriting a single file loses history. Neither is what you want. Partitioning by feed_date and overwriting only the matching partition directory makes a reload replace exactly one feed version and leave the rest of the archive untouched.
Production-Ready Python Implementation
The script below writes a GTFS stop_times table to a partitioned Parquet dataset with an explicit schema, then reads it back with a predicate that prunes to one partition. The same pattern applies to any GTFS table; stop_times is the one where columnar storage and pushdown matter most.
import logging
from pathlib import Path
import pandas as pd
import pyarrow as pa
import pyarrow.dataset as ds
import pyarrow.compute as pc
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
# Explicit schema: IDs are strings, sequence is int32, time fields stay strings
# (GTFS permits times past 24:00:00), partition keys are strings.
STOP_TIMES_SCHEMA = pa.schema([
("trip_id", pa.string()),
("stop_id", pa.string()),
("stop_sequence", pa.int32()),
("arrival_time", pa.string()),
("departure_time", pa.string()),
("shape_dist_traveled", pa.float64()),
("agency_id", pa.string()),
("feed_date", pa.string()),
])
# The read dtypes mirror the write schema so pandas never coerces an ID.
READ_DTYPES = {
"trip_id": str,
"stop_id": str,
"stop_sequence": "Int32",
"arrival_time": str,
"departure_time": str,
"shape_dist_traveled": float,
}
def write_stop_times(
feed_dir: str, base_dir: str, agency_id: str, feed_date: str
) -> int:
"""Write one feed's stop_times as a partition under agency_id / feed_date."""
# dtype keys that name an absent column are ignored by pandas, so a feed
# missing the optional shape_dist_traveled still reads cleanly.
frame = pd.read_csv(Path(feed_dir) / "stop_times.txt", dtype=READ_DTYPES)
# Fill the optional column if the feed omitted it, then tag partitions.
if "shape_dist_traveled" not in frame.columns:
frame["shape_dist_traveled"] = pd.NA
frame = frame.assign(agency_id=agency_id, feed_date=feed_date)
table = pa.Table.from_pandas(
frame, schema=STOP_TIMES_SCHEMA, preserve_index=False
)
ds.write_dataset(
table,
base_dir=base_dir,
format="parquet",
partitioning=ds.partitioning(
pa.schema([("agency_id", pa.string()), ("feed_date", pa.string())]),
flavor="hive",
),
# Overwrite only the partition(s) present in this write — idempotent reload
existing_data_behavior="delete_matching",
basename_template="stop_times-{i}.parquet",
)
logging.info(
"Wrote %d rows to %s (agency_id=%s, feed_date=%s)",
len(frame), base_dir, agency_id, feed_date,
)
return len(frame)
def read_stop_times(base_dir: str, agency_id: str, feed_date: str) -> pd.DataFrame:
"""Read back one partition using predicate pushdown on the partition keys."""
dataset = ds.dataset(base_dir, format="parquet", partitioning="hive")
predicate = (
(pc.field("agency_id") == agency_id)
& (pc.field("feed_date") == feed_date)
)
table = dataset.to_table(filter=predicate)
logging.info(
"Read %d rows for agency_id=%s feed_date=%s (pruned to one partition)",
table.num_rows, agency_id, feed_date,
)
return table.to_pandas()
# --- Entry point ---
# write_stop_times("/data/gtfs/mta", "warehouse/stop_times", "mta", "2024-04-20")
# frame = read_stop_times("warehouse/stop_times", "mta", "2024-04-20")
Step-by-Step Walkthrough
Explicit schema pins the types. STOP_TIMES_SCHEMA declares trip_id and stop_id as string, stop_sequence as int32, and the two time columns as string. Passing this schema to Table.from_pandas forces the conversion to match, so even if pandas inferred an int64 stop_id upstream, the write fails loudly rather than silently baking a bad type into the file. Every feed version then produces an identical, mergeable schema — the same discipline the step-by-step partridge parsing guide applies at read time.
Partition columns are attached, not inferred. assign(agency_id=..., feed_date=...) adds the two low-cardinality keys the dataset is partitioned on. write_dataset reads those columns out of the table and writes them into the directory path (agency_id=mta/feed_date=2024-04-20/) rather than into the Parquet files, so they cost no storage and become prunable at read time.
Hive-flavored partitioning. Declaring ds.partitioning(..., flavor="hive") writes the key=value directory names that every major engine — DuckDB, Spark, Polars, and cloud warehouses — recognizes automatically, so the same dataset is portable across the whole ecosystem.
Idempotent reloads via delete_matching. existing_data_behavior="delete_matching" tells write_dataset to clear any existing files in exactly the partitions this write touches before writing the new ones. Re-running the load for mta / 2024-04-20 replaces that one partition and leaves every other agency and date intact — structural idempotency without an upsert. This mirrors the version-keyed reload model in the exporting to databases and warehouses guide.
Predicate pushdown on read. In read_stop_times, the pc.field(...) filter on the partition columns lets pyarrow prune whole directories before opening a single file. Filters on non-partition columns are additionally pushed down to per-row-group statistics, so unmatched row groups are skipped inside the files that do get opened. A query for one agency and date reads only that partition.
Verification and Output
Confirm the written dataset has the expected schema, partitions, and row counts before trusting it downstream:
import pyarrow.dataset as ds
import pyarrow as pa
dataset = ds.dataset("warehouse/stop_times", format="parquet", partitioning="hive")
# Schema must match what we declared, including partition columns
assert dataset.schema.field("stop_id").type == pa.string()
assert dataset.schema.field("stop_sequence").type == pa.int32()
assert dataset.schema.field("arrival_time").type == pa.string()
assert "agency_id" in dataset.schema.names
assert "feed_date" in dataset.schema.names
# Enumerate the partition directories that were written
fragments = sorted({f.path.rsplit("/", 1)[0] for f in dataset.get_fragments()})
print("Partitions:")
for path in fragments:
print(" ", path)
# A filtered read must return only the requested partition's rows
one = dataset.to_table(
filter=(ds.field("agency_id") == "mta") & (ds.field("feed_date") == "2024-04-20")
)
assert set(one.column("agency_id").to_pylist()) == {"mta"}
assert set(one.column("feed_date").to_pylist()) == {"2024-04-20"}
print(f"Single-partition read returned {one.num_rows} rows")
# A reload of the same partition must not change the total row count
total_before = dataset.count_rows()
# write_stop_times(...) re-run here would leave this unchanged:
print(f"Total rows across all partitions: {total_before}")
The schema assertions guarantee that IDs stayed strings and sequences stayed int32, so a later engine will not choke on a drifted type. The partition listing shows the agency_id=.../feed_date=.../ directory layout, and the filtered read returns rows from a single feed version — proof that pushdown is pruning correctly.
Gotchas and Edge Cases
-
Over-partitioning creates a small-file storm. Partitioning by
trip_idorstop_idyields one directory per value — potentially millions of tiny files that make both writes and reads pathologically slow. Restrict partitioning toagency_idandfeed_date, and let row groups within each Parquet file handle finer-grained pruning. -
feed_dateas an integer sorts and prunes wrong. Store the partition date as an ISO string ("2024-04-20") rather than an integer like20240420. String ISO dates sort lexicographically in the correct chronological order and read back unambiguously, which matters when a query filters a date range across the archive. Anchor the value to your agency metadata and feed versioning convention. -
Optional columns absent from some feeds break schema unification.
shape_dist_traveledis optional in GTFS; if one feed omits it and you do not add it back, that partition’s file lacks the column and a cross-partition read raises a schema-mismatch error. Always reindex each frame to the full declared schema before writing, as the script does. -
Large
stop_timesneeds bounded memory on write. A national feed’sstop_times.txtcan exceed RAM. Read and write it in chunks per partition rather than materializing the whole table; the chunking techniques are in optimizing pandas memory usage for transit feeds.
Frequently Asked Questions
Why define an explicit pyarrow schema instead of letting it infer?
Inference reads dtypes from the DataFrame, so a stop_id of "007" that pandas already coerced to an integer, or an all-null optional column typed as float, gets locked into the Parquet file. An explicit schema pins string IDs, int32 sequences, and text time fields so every feed version writes an identical, mergeable schema.
What columns should I partition GTFS Parquet by?
Partition by low-cardinality keys that queries filter on — agency_id and feed_date. That keeps the number of partition directories small and lets a reader skip whole feed versions or agencies. Never partition by a high-cardinality column like trip_id, which creates millions of tiny files.
How does predicate pushdown speed up GTFS reads?
When you pass a filter on a partition column, pyarrow prunes entire directories before opening any file, and filters on other columns are pushed to row-group statistics so unmatched row groups are skipped. A query for one agency and one feed_date reads only that partition instead of the whole dataset.
Related
- Loading GTFS into PostGIS with Python — the transactional, spatial counterpart to this columnar-warehouse write, from the same DataFrames
- Memory-Efficient Processing for Large Feeds — chunked reads and writes so a national
stop_times.txtnever overruns RAM - Optimizing Pandas Memory Usage for Transit Feeds — dtype downcasting that keeps each partition’s frame small before it becomes an Arrow table
- Up: Exporting GTFS to Databases and Warehouses — the schema and idempotent-load model this Parquet workflow implements
- Section: Python Parsing & Data Normalization — the full GTFS parsing and normalization pipeline · Home