File size: 1,387 Bytes
d95db82 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import pandas as pd
COLUMN_MAP = {
"tpep_pickup_datetime": "pickup",
"tpep_dropoff_datetime": "dropoff",
"passenger_count": "passengers",
"trip_distance": "distance",
"fare_amount": "fare",
"tip_amount": "tip",
"tolls_amount": "tolls",
"total_amount": "total",
"color": "color",
}
PAYMENT_TYPES = {
1: "credit card",
2: "cash",
}
MAX_TRIP_DURATION = 8000
if __name__ == "__main__":
raw = pd.read_csv(
"raw/taxis.csv",
parse_dates=["tpep_pickup_datetime", "tpep_dropoff_datetime"]
).rename(columns=str.lower)
loc = pd.read_csv("raw/taxi_zones.csv").set_index("LocationID").drop_duplicates()
clean = (
raw[list(COLUMN_MAP)]
.rename(columns=COLUMN_MAP)
.assign(payment=raw["payment_type"].map(PAYMENT_TYPES))
.assign(pickup_zone=raw["pulocationid"].map(loc["zone"]))
.assign(dropoff_zone=raw["dolocationid"].map(loc["zone"]))
.assign(pickup_borough=raw["pulocationid"].map(loc["borough"]))
.assign(dropoff_borough=raw["dolocationid"].map(loc["borough"]))
.loc[lambda x: x["dropoff_borough"] != "EWR"]
.loc[lambda x: x.eval("dropoff - pickup").dt.seconds < MAX_TRIP_DURATION]
.loc[lambda x: (x["fare"] > 0) & (x["fare"] < 200)]
.loc[lambda x: (x["tip"] / x["fare"]) < 1]
)
clean.to_csv("taxis.csv", index=False) |