ml-pipeline-nyc-taxi-trip-duration

NYC Taxi Trip Duration: an end-to-end, tested ML pipeline on AWS

CI Python Tests AWS Lambda License

Predicts how long a New York City taxi trip will take, from information known at pickup. Trained on 1.44 million cleaned trips: MAE 171 s (R² 0.80) on a random test split and 187 s on an honest time-based split (train January–May, test June), a 23.6% lower error than the leak-free baseline. Served as a live API and interactive map on AWS Lambda.

▶ Live demo: click two points on a map of NYC and get a predicted trip time · API docs · Interactive speed map · Model card

API demo


Contents

  1. Project overview
  2. Architecture
  3. Key findings
  4. Results
  5. Quickstart
  6. Prediction API
  7. User stories and acceptance criteria
  8. Data quality and testing
  9. SQL layer, Spark and orchestration
  10. Azure version
  11. Experiment tracking
  12. Limitations
  13. Future work
  14. Project structure
  15. Data sources and background

Project overview

The problem. Ride-hailing, dispatch and logistics teams need to know how long a trip will take before it starts, to plan vehicles, quote arrival times and spot unusual delays.

The data. 1,458,644 NYC yellow-taxi trips from January–June 2016 (Kaggle), plus daily Central Park weather from the US National Weather Service.

What this project does, from raw file to deployed service:

Stage What happens Tools
Validate Data contracts reject malformed or truncated files before anything else runs pandera
Clean 4 rules remove 15,704 invalid trips (1.08%); verified identical between SQL and pandas DuckDB SQL, Parquet
Features 16 leak-free features known at pickup: distances, compass bearing, time, location, airports pandas, PySpark (cross-check)
Train LightGBM on log duration; 10 experiments per run, including an honest time-based split LightGBM, scikit-learn, MLflow
Explain What drives each prediction, and where traffic is slowest SHAP, DuckDB, folium
Serve Validated REST API and an interactive map page FastAPI, Docker, AWS Lambda, S3, ECR
Automate Every push linted, tested and Docker-checked; the pipeline runs weekly as a DAG GitHub Actions, pytest, Prefect

Roles this project demonstrates: data science (feature engineering, leakage detection, honest evaluation, explainability), ML engineering (API, Docker, AWS, MLflow, CI/CD) and data engineering (SQL ETL, data contracts, Parquet, Spark, orchestration).


Architecture

flowchart LR
    K[Kaggle: 1.46M trips<br/>+ NWS weather] --> V[Validation<br/>pandera contracts]
    V --> C[Cleaning<br/>DuckDB SQL]
    C --> P[(Parquet)]
    P --> F[16 features]
    F --> T[LightGBM training<br/>MLflow tracking]
    F -. cross-check .-> S[PySpark ETL]
    T --> M[(Model file<br/>Amazon S3)]
    M --> L[FastAPI in Docker<br/>AWS Lambda]
    U([Browser / client]) -- map page, POST /predict --> L
    O[Prefect: weekly schedule] -. runs .-> V
    CI[GitHub Actions] -. lint, 91 tests,<br/>Docker smoke test .-> L

The same Docker image runs locally, in CI and on AWS Lambda (through the Lambda Web Adapter). The model is stored in S3 and loaded at startup; the image is stored in Amazon ECR and runs on arm64 (Graviton) in the Canada (Central) region.

AWS Lambda console: the nyc-taxi-api function, deployed as a container image The nyc-taxi-api function in the AWS Lambda console, deployed as a container image (package type: Image).


Key findings


Results

Model comparison (random split, 72,147 test trips)

Model Features MAE (s) RMSE (s) RMSLE R²
Template LightGBM (leaks the target) 6 17.3 70.9 0.0603 0.989
LightGBM, no leak (honest baseline) 5 223.9 357.9 0.3934 0.708
+ engineered features 16 181.4 311.6 0.3230 0.778
Improved LightGBM (tuned) 16 171.0 296.9 0.3103 0.799
Linear Regression 16 262.9 425.7 0.4332 0.587

Honest evaluation: time-based split

A random split lets the same days appear in training and test. In real use the model always predicts future trips, so it is also trained on January–May and tested on June:

Evaluation MAE What it shows
Random split 171.0 s Same days in training and test
Random split, June trips only 181.6 s June trips are harder: trips got longer every month
Time-based split (train Jan–May, test June) 186.6 s The realistic estimate for future trips

About two-thirds of the gap (+10.6 s) comes from June being a harder month, one-third (+5.0 s) from the model never having seen June.

Weather experiment No weather + real weather + day of year only
Random split 171.0 s 167.2 s 165.8 s
Time split 186.6 s 189.8 s 195.2 s

Where the model is accurate, and where it isn’t

Actual trip length Share of trips Average error Bias
Under 20 minutes 80% 1.2–2.6 min about 0
30–60 minutes 6.9% 7.9 min −5.2 min
Over 1 hour 0.7% 21.6 min −20.6 min (underestimated)
What drives the predictions (SHAP) Where traffic is slowest
SHAP summary Speed map
Distance dominates; the pickup hour comes next. Median trip speed by pickup area: Midtown is slowest, airport trips fastest; Central Park shows as a gap.
Daily traffic cycle the model learned Predicted vs actual (10,000 test trips)
SHAP pickup hour Predicted vs actual
Pickups at 5 a.m. are predicted about a third faster, midday about 20% slower; weekend mornings have no rush hour. Accurate up to about 30 minutes; the longest trips fall below the diagonal.

Full details, including the cleaning summary and known weaknesses: model card (regenerated on every pipeline run).


Quickstart

git clone https://github.com/Fatema-alam43/ml-pipeline-nyc-taxi-trip-duration.git && cd ml-pipeline-nyc-taxi-trip-duration
python -m venv .venv && source .venv/bin/activate && pip install -r requirements.txt
python scripts/download_data.py && python -m src.pipeline

download_data.py needs a Kaggle API token in ~/.kaggle/kaggle.json and the competition rules accepted. The pipeline takes a few minutes and writes the cleaned data, features, model, metrics, figures and model card. Other options:

python -m src.pipeline --fast           # only the main model
python -m src.pipeline --spark-check    # plus the PySpark cross-check (needs Java 17 or 21)
python -m flows.pipeline_flow           # run as a Prefect flow
uvicorn api.main:app --port 8080        # API and map page at http://localhost:8080
pytest                                  # 91 tests

Prediction API

Live: https://ifksrlyyuqede7xad4naehovee0qrlno.lambda-url.ca-central-1.on.aws/ (map page) · /docs (interactive documentation)

curl -X POST https://ifksrlyyuqede7xad4naehovee0qrlno.lambda-url.ca-central-1.on.aws/predict \
  -H "Content-Type: application/json" -d '{
  "pickup_latitude": 40.7580, "pickup_longitude": -73.9855,
  "dropoff_latitude": 40.6413, "dropoff_longitude": -73.7781,
  "pickup_datetime": "2016-06-15T17:30:00", "passenger_count": 1
}'
{"predicted_duration_minutes": 82.8, "predicted_duration_seconds": 4970.5, "straight_line_distance_km": 21.77,
 "model_version": "c1763ba", "warnings": []}

Invalid requests get a clear message for every problem at once:

{"error": "Invalid trip request", "details": [
  {"field": "pickup_latitude", "message": "pickup_latitude 42.5 is outside New York City (allowed 40.5 to 41.0)"},
  {"field": "passenger_count", "message": "Input should be less than or equal to 6"}]}
Endpoint What it does
GET / Interactive map: click pickup and dropoff, see the prediction and the same trip at five times of day
POST /predict Predicted duration (minutes and seconds), straight-line distance, model version, warnings
GET /health Status, model version (git commit), SHA-256 of the model file, where it was loaded from
GET /docs Interactive API documentation
# Run with Docker
docker build --build-arg PYTHON_VERSION=$(cat .python-version) -t nyc-taxi-api .
docker run -p 8080:8080 -v "$(pwd)/models:/models:ro" -e MODEL_PATH=/models/lgbm_improved.txt nyc-taxi-api

User stories and acceptance criteria

Each criterion is checked by an automated test or a measured result.

1. As a dispatcher, I want an estimated trip time before the trip starts, so that I can plan vehicles and tell customers when they will arrive.

2. As an API user, I want clear errors for invalid input, so that I can fix my request instead of receiving a meaningless prediction.

3. As a data engineer, I want bad input data to stop the pipeline, so that a corrupted file never reaches the model.

4. As a data scientist, I want an honest estimate of performance on future trips, so that stakeholders are not misled by optimistic numbers.

5. As an ML engineer, I want reproducible, traceable experiments, so that any result can be explained and re-created.

6. As a team lead, I want every change tested automatically, so that a mistake is caught before it is deployed.

7. As an operations owner, I want the pipeline to run on a schedule and recover from failures, so that the model stays current without manual work.


Data quality and testing

Data contracts (src/schemas.py, pandera):

Stage Contract
Raw file Exact columns and types, unique ids, valid vendor and flag values, plausible coordinates, positive durations; rejects Excel-truncated files and damaged timestamps
Clean data Every row satisfies every cleaning rule
Features The 16 expected features with valid ranges, no rush hour on weekends, and no leaky columns

Cleaning rules (all 1,458,644 raw trips):

Rule Flagged Newly removed
Duration under 60 s or over 6 h 10,656 10,656
Zero distance 5,897 4,219
Coordinates outside NYC 997 813
Zero passengers 60 16
Total   15,704 (1.08%)

91 automated tests (pytest, about a minute), on synthetic test data with planted problems, so no Kaggle data is stored in the repository and every expected count is known exactly:

CI catches real mistakes. Changing a single character in the SQL duration rule (< to <=) makes 5 tests fail, each naming the problem:

CI catching a broken rule

SQL and pandas disagree (959 vs 961 trips kept) The whole pipeline stops: 5 failed, 55 passed
SQL vs pandas Pipeline test failed

The SQL rule wrongly removed the 2 planted trips of exactly 60 seconds. The test for that boundary failed, and so did the independent pandas cross-check and the end-to-end pipeline test, so the mistake could not reach the model.


SQL layer, Spark and orchestration

DuckDB SQL (sql/) runs the cleaning and the analytics directly on Parquet files: CTEs, window functions (running totals, 7-day rolling averages, LAG), RANK() OVER (PARTITION BY ...) with QUALIFY, conditional aggregation with FILTER, and a reusable haversine macro.

Parquet vs CSV:

Dataset CSV Parquet Size pandas load DuckDB query
Raw trips 200.6 MB 48.0 MB 4.2x smaller 62.5x faster 29.0x faster
Engineered features 175.0 MB 59.6 MB 2.9x smaller 55.7x faster 39.4x faster

PySpark (src/spark_etl.py) re-implements the cleaning and features for data too large for one machine. On all 1,458,644 trips it removes exactly the same 15,704 rows per rule, keeps the same 1,442,940 trips and produces bit-identical values for 15 of the 16 features. The compass bearing differs in 2 rows by 0.0000076°, because Java and NumPy trigonometry can differ in the last bit on Apple Silicon, so it is compared within 0.001°. Two details had to match exactly: coordinates are compared as 32-bit floats (a 64-bit limit would wrongly exclude longitude −73.7), and Spark counts weekdays from Sunday = 1 while pandas uses Monday = 0. On a single laptop DuckDB is faster; Spark pays off on a cluster.

Prefect (flows/pipeline_flow.py) runs the pipeline as a scheduled DAG, every Monday at 06:00 Toronto time:

flowchart LR
    D[download<br/>retries: 2] --> V[validate] --> C[clean] --> F[features]
    F --> S[spark-crosscheck<br/>optional]
    F --> T[train] --> A[analytics] --> R[report]

Prefect flow run


Azure version: Data Factory, Databricks and Power BI

The same pipeline also runs on Microsoft Azure, as a bronze → silver → gold lakehouse, and ends in a Power BI report for business users. It removes exactly the same 15,704 rows as the local pipeline.

flowchart LR
    K[Kaggle CSVs] --> L[(ADLS Gen2<br/>landing)]
    L --> A[Data Factory<br/>copy to Parquet<br/>+ row-count check]
    A --> B[(bronze<br/>Parquet)]
    B --> D[Databricks serverless<br/>PySpark notebook]
    D --> S[(silver<br/>Delta: clean trips, features)]
    D --> G[(gold<br/>Delta + CSV: 6 summary tables)]
    G --> P[Power BI<br/>3-page report]
    T[Weekly trigger<br/>Mondays 06:00] -. runs .-> A
Layer Service What happens
landing Azure Data Lake Storage Gen2 Raw Kaggle files, exactly as downloaded
bronze Azure Data Factory Copies CSV to Parquet, then checks the row count: a file with exactly 1,048,575 rows (Excel’s limit) is deleted and the pipeline fails with a clear message
silver Azure Databricks (serverless, PySpark) The same 4 cleaning rules and 16 features as the local pipeline, saved as Delta tables; the notebook fails itself unless exactly 15,704 rows are removed
gold Azure Databricks 6 small summary tables for reporting: daily KPIs with weather, hourly patterns, speed by area, airport vs other trips, weather impact, data quality
report Power BI 3 pages built on the gold tables, with a date table and DAX measures

Same answer in four engines

The cleaning step was implemented four times and gives the same result every time:

Engine Rows removed Clean trips
pandas 15,704 1,442,940
DuckDB SQL 15,704 1,442,940
PySpark (local) 15,704 1,442,940
Azure Databricks (serverless) 15,704 1,442,940

Including the 997 trips on the edge of the NYC boundary, where a 64-bit instead of 32-bit comparison would have changed the count by one.

Databricks cleaning check

Data Factory: ingestion and the data-quality check

Pipeline Excel-truncated file rejected
ADF pipeline ADF failed run
Copy trips and weather to bronze as Parquet, then check the row count Test with a 1,048,575-row file: the bad copy is deleted and the run fails on purpose

The pipeline runs every Monday at 06:00 (Toronto time) through a schedule trigger. Pipeline, dataset, linked service and trigger definitions are in azure/adf/.

Power BI report

Overview When and where
Power BI overview Power BI when and where

Power BI weather and data quality

The report file and a PDF export are in azure/powerbi/; the PDF opens without Power BI.

Security and cost

What is not automated yet

The Data Factory pipeline (ingestion and the row-count check) runs on a weekly schedule; the Databricks notebook is currently run on demand after ingestion. Connecting it as the final pipeline step is listed under future work.

More screenshots: successful run, trigger, lake containers, Databricks results | Successful pipeline run | Weekly trigger | |---|---| | ![ADF success](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_adf_success.png) | ![ADF trigger](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_adf_trigger.png) | | Landing container | Silver container | Gold container | |---|---|---| | ![Landing](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_landing.png) | ![Silver](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_silver_container.png) | ![Gold](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_gold_container.png) | | Bronze row count | Results: airport, weather, slowest days | |---|---| | ![Bronze](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_databricks_bronze.png) | ![Results](/ml-pipeline-nyc-taxi-trip-duration/reports/figures/azure_databricks_results.png) |

Experiment tracking

Every model run is logged to MLflow with its parameters, metrics, feature list, git commit and model file (10 runs per pipeline run: baselines, the tuned model, Linear Regression, and weather experiments on both splits).

mlflow ui --backend-store-uri sqlite:///mlflow.db --port 5001

MLflow charts MLflow runs


Limitations

Future work


Project structure

├── README.md
├── config.yaml                  # paths, validation and cleaning thresholds, model settings
├── requirements.txt             # main packages (requirements-lock.txt: exact versions)
├── requirements-api.txt         # API image packages, pinned to the training versions
├── requirements-dev.txt         # pytest, ruff, httpx
├── pyproject.toml               # pytest and ruff settings
├── Dockerfile                   # API image (runs locally, in CI and on AWS Lambda)
├── .github/workflows/ci.yml     # lint + tests + Docker smoke test on every push
├── src/
│   ├── config.py                # loads config.yaml
│   ├── etl.py                   # load, validate, clean (SQL, verified against pandas)
│   ├── schemas.py               # pandera data contracts (raw, clean, features)
│   ├── features.py              # 16 features, shared by training and the API
│   ├── db.py                    # runs the SQL files with DuckDB
│   ├── train.py                 # random and time-based splits, models, evaluation
│   ├── tracking.py              # MLflow experiment tracking
│   ├── explain.py               # SHAP explanations
│   ├── plots.py                 # figures and speed maps
│   ├── model_card.py            # writes reports/model_card.md
│   ├── spark_etl.py             # PySpark version of cleaning + features
│   ├── steps.py                 # the pipeline steps (each reads and writes files)
│   └── pipeline.py              # runs everything: python -m src.pipeline
├── sql/                         # DuckDB: cleaning, views, 6 analysis queries
├── flows/pipeline_flow.py       # scheduled Prefect flow (the pipeline as a DAG)
├── api/
│   ├── main.py                  # FastAPI app: /, /predict, /health
│   ├── schemas.py               # request validation (Pydantic)
│   ├── model.py                 # loads the model from a file or S3
│   └── static/index.html        # interactive map page
├── deploy/                      # AWS Lambda deploy and teardown scripts
├── azure/
│   ├── adf/                     # Data Factory pipeline, datasets, linked service, trigger (JSON)
│   ├── databricks/              # bronze → silver → gold notebook (PySpark)
│   └── powerbi/                 # report (.pbix) and PDF export
├── docs/deploy_aws.md           # deployment guide
├── scripts/
│   ├── download_data.py         # downloads and verifies the Kaggle data
│   ├── benchmark_storage.py     # CSV vs Parquet benchmark
│   ├── make_test_fixtures.py    # synthetic test data with planted problems
│   ├── make_demo_model.py       # small model for tests and CI
│   └── pin_api_requirements.py  # pins the API image to training versions
├── tests/                       # 91 pytest tests + synthetic fixtures
├── examples/                    # sample API request
└── reports/
    ├── model_card.md            # generated model card
    ├── metrics.json             # generated metrics
    ├── speed_map.html           # interactive speed map
    └── figures/                 # SHAP, map, results, CI, MLflow, Prefect, AWS

Data sources and background

The data is downloaded by scripts/download_data.py and is not stored in this repository.

License: code under MIT. The data is subject to Kaggle’s terms and is not redistributed.