Handling Frequency-Based vs Timetable Schedules

pip install pandas>=1.5.0 pyarrow numpy

import pandas as pd
import numpy as np

# Explicit dtypes prevent silent coercion of trip_id integers and stop_id strings
stop_times = pd.read_csv(
    "stop_times.txt",
    dtype={
        "trip_id": str,
        "stop_id": str,
        "stop_sequence": "int32",
        "arrival_time": str,
        "departure_time": str,
    },
    usecols=["trip_id", "stop_id", "stop_sequence", "arrival_time", "departure_time"],
)

frequencies = pd.read_csv(
    "frequencies.txt",
    dtype={
        "trip_id": str,
        "start_time": str,
        "end_time": str,
        "headway_secs": "int32",
        "exact_times": "Int8",  # nullable integer; field is optional in some feeds
    },
)

# Identify frequency-governed trip_ids
freq_trip_ids = set(frequencies["trip_id"].unique())

# Tag every stop_times row by schedule paradigm
stop_times["schedule_type"] = np.where(
    stop_times["trip_id"].isin(freq_trip_ids), "frequency", "timetable"
)

# Split for separate processing paths
timetable_stop_times = stop_times[stop_times["schedule_type"] == "timetable"].copy()

# For frequency trips: keep only one stop_times row per (trip_id, stop_sequence)
# to preserve the relative offset template; duplicates arise from multi-agency merges
freq_stop_times = (
    stop_times[stop_times["schedule_type"] == "frequency"]
    .drop_duplicates(subset=["trip_id", "stop_sequence"])
    .copy()
)

def gtfs_time_to_timedelta(time_str: str) -> pd.Timedelta:
    """Convert a GTFS HH:MM:SS string (which may exceed 23:59) to a pd.Timedelta."""
    parts = time_str.strip().split(":")
    total_seconds = int(parts[0]) * 3600 + int(parts[1]) * 60 + int(parts[2])
    return pd.Timedelta(seconds=total_seconds)

def parse_time_column(series: pd.Series) -> pd.Series:
    """Vectorized GTFS time parsing. Returns a Series of pd.Timedelta."""
    # Split once and compute without a Python loop
    split = series.str.strip().str.split(":", expand=True).astype("int32")
    return pd.to_timedelta(split[0] * 3600 + split[1] * 60 + split[2], unit="s")

def expand_frequencies(
    freq_df: pd.DataFrame,
    freq_stop_times_df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Expand each frequencies.txt row into discrete per-stop departure times.
    Returns a DataFrame with the same schema as timetable_stop_times plus
    'exact_times' and 'departure_utc_offset' columns.
    """
    freq_df = freq_df.copy()
    freq_df["start_td"] = parse_time_column(freq_df["start_time"])
    freq_df["end_td"] = parse_time_column(freq_df["end_time"])
    freq_df["headway_td"] = pd.to_timedelta(freq_df["headway_secs"], unit="s")
    # Fill missing exact_times with 0 (approximate) per spec default
    freq_df["exact_times"] = freq_df["exact_times"].fillna(0).astype("int8")

    # Merge frequency windows onto the stop sequence template
    merged = freq_stop_times_df.merge(
        freq_df[["trip_id", "start_td", "end_td", "headway_td", "exact_times"]],
        on="trip_id",
        how="inner",
    )

    # For each row, calculate the first stop's relative offset from trip start
    # (arrival_time in stop_times is the offset when trips are frequency-governed)
    merged["stop_offset_td"] = parse_time_column(merged["arrival_time"])
    # The offset of stop 0 (the trip origin); we subtract it to get a relative offset
    first_stop_offset = (
        merged.groupby("trip_id")["stop_offset_td"].transform("min")
    )
    merged["relative_offset_td"] = merged["stop_offset_td"] - first_stop_offset

    # Generate departure sequences without a Python loop using groupby + apply
    expanded_chunks = []
    for (trip_id, exact_t), group in merged.groupby(["trip_id", "exact_times"]):
        start_td = group["start_td"].iloc[0]
        end_td = group["end_td"].iloc[0]
        headway_td = group["headway_td"].iloc[0]

        # Sequence of trip-start times within the service window
        trip_starts = pd.timedelta_range(start=start_td, end=end_td, freq=headway_td)

        # Cross-join trip_starts with the stop sequence
        for stop_row in group.itertuples(index=False):
            departure_offsets = trip_starts + stop_row.relative_offset_td
            chunk = pd.DataFrame({
                "trip_id": trip_id,
                "stop_id": stop_row.stop_id,
                "stop_sequence": stop_row.stop_sequence,
                "departure_time": departure_offsets,
                "schedule_type": "frequency",
                "exact_times": exact_t,
            })
            expanded_chunks.append(chunk)

    if not expanded_chunks:
        return pd.DataFrame(columns=[
            "trip_id", "stop_id", "stop_sequence",
            "departure_time", "schedule_type", "exact_times",
        ])

    return pd.concat(expanded_chunks, ignore_index=True)

# Add exact_times column for schema consistency (timetable trips are always exact)
timetable_stop_times["exact_times"] = 1
timetable_stop_times["departure_time"] = parse_time_column(
    timetable_stop_times["departure_time"]
)

def merge_and_deduplicate(
    timetable_df: pd.DataFrame,
    expanded_freq_df: pd.DataFrame,
) -> pd.DataFrame:
    """
    Combine timetable and expanded frequency departures.
    Timetable rows take precedence over frequency-expanded rows for the same
    (trip_id, stop_id, departure_time) combination.
    """
    # Timetable rows first so drop_duplicates keeps them over frequency rows
    combined = pd.concat(
        [timetable_df, expanded_freq_df],
        ignore_index=True,
        sort=False,
    )

    combined = combined.sort_values(
        ["trip_id", "stop_sequence", "departure_time"]
    ).reset_index(drop=True)

    # Round to 1-second precision before dedup to absorb floating-point drift
    combined["departure_time"] = combined["departure_time"].dt.round("s")

    before_count = len(combined)
    combined = combined.drop_duplicates(
        subset=["trip_id", "stop_id", "departure_time"],
        keep="first",  # first = timetable row, since we concatenated it first
    )
    duplicates_removed = before_count - len(combined)
    if duplicates_removed > 0:
        print(f"[INFO] Removed {duplicates_removed} duplicate departure records.")

    return combined

def normalize_to_utc(
    df: pd.DataFrame,
    service_date_str: str,
    agency_tz: str,
) -> pd.DataFrame:
    """
    Convert GTFS timedelta departure_time values to UTC wall-clock timestamps.

    Parameters
    ----------
    df : DataFrame with a 'departure_time' column of pd.Timedelta
    service_date_str : ISO date string for the service day, e.g. '2024-01-15'
    agency_tz : IANA timezone name from agency.txt agency_timezone field
    """
    # Build a timezone-aware base date representing midnight of the service day
    base_date = pd.Timestamp(service_date_str).tz_localize(agency_tz)

    # Adding a timedelta to a tz-aware Timestamp handles day-wrap automatically.
    # A departure_time of 25:30:00 correctly maps to 01:30:00 the next calendar day.
    df = df.copy()
    df["departure_utc"] = (base_date + df["departure_time"]).dt.tz_convert("UTC")

    df = df.drop(columns=["departure_time"])
    return df

def validate_normalized_schedule(df: pd.DataFrame) -> bool:
    """
    Run structural checks on a normalized departure DataFrame.
    Returns True if all checks pass; prints warnings for any violations found.
    """
    checks = {
        "missing_departure_utc": df["departure_utc"].isna().sum(),
        "missing_stop_id": df["stop_id"].isna().sum(),
        "missing_trip_id": df["trip_id"].isna().sum(),
        "duplicate_trip_stop_departure": df.duplicated(
            subset=["trip_id", "stop_id", "departure_utc"]
        ).sum(),
        "negative_headway_within_trip": (
            df.sort_values(["trip_id", "stop_sequence", "departure_utc"])
            .groupby(["trip_id", "stop_sequence"])["departure_utc"]
            .diff()
            .dropna()
            .lt(pd.Timedelta(0))
            .sum()
        ),
    }

    all_pass = True
    for metric, count in checks.items():
        if count > 0:
            print(f"[WARN] {metric}: {count} records flagged")
            all_pass = False

    # Assert minimum expected row count (basic sanity: at least one departure per trip)
    assert len(df) >= df["trip_id"].nunique(), (
        "Normalized table has fewer rows than unique trip_ids — expansion likely failed."
    )

    if all_pass:
        print(f"[OK] Validation passed: {len(df):,} departure records, "
              f"{df['trip_id'].nunique():,} unique trips.")
    return all_pass

import pyarrow as pa
import pyarrow.parquet as pq

def write_normalized_partition(
    df: pd.DataFrame,
    output_dir: str,
    service_date: str,
    agency_id: str,
) -> None:
    """Write a normalized departure batch to a Parquet partition."""
    table = pa.Table.from_pandas(df, preserve_index=False)
    pq.write_to_dataset(
        table,
        root_path=output_dir,
        partition_cols=["agency_id", "service_date"],
    )

# Vectorized alternative to the itertuples inner loop
trip_starts_frames = []
for (trip_id,), group in frequencies.groupby(["trip_id"]):
    starts = pd.timedelta_range(
        start=group["start_td"].iloc[0],
        end=group["end_td"].iloc[0],
        freq=group["headway_td"].iloc[0],
    )
    trip_starts_frames.append(
        pd.DataFrame({"trip_id": trip_id, "trip_start_td": starts})
    )

all_trip_starts = pd.concat(trip_starts_frames, ignore_index=True)

# Now merge the full stop sequence against all trip starts in a single join
expanded = freq_stop_times.merge(all_trip_starts, on="trip_id", how="inner")
expanded["departure_time"] = expanded["trip_start_td"] + expanded["relative_offset_td"]