Pythonの高速データフレームライブラリFireDucksを使ってみた

import pandas as pd
import os

dir = "xxx"
df_list = []
for year in [2022, 2023]:
for i in range(12):
month = str(i+1).zfill(2)
fn = f"yellow_tripdata_{year}-{month}.parquet"
file = os.path.join(dir, fn)
if not os.path.exists(file):
continue
df = pd.read_parquet(fn)
df_list.append(df)

all_df = pd.concat(df_list)
all_df.to_csv("taxi_all.csv")

from time import time
from functools import wraps

def timer(func):
@wraps(func)
def wp(*args, **kargs):
t = time()
ret = func(*args, **kargs)
print(f"{func.__name__} : {(time() - t):.5g} [sec]")
return ret
return wp

def evaluate(df):
if hasattr(df, "_evaluate"):
df._evaluate()

import pandas as pd

@timer
def file_read(fn, args={}):
df = pd.read_csv(fn, **args)
evaluate(df)
print(df.shape)
return df

@timer
def drop_na(df):
df.dropna(how="all", inplace=True)
evaluate(df)
return df

@timer
def txt_to_date(df, low):
df[low] = pd.to_datetime(df[low])
evaluate(df)
return df

# 一人以上乗車している
@timer
def check_passenger_c(df):
df_ = df.groupby("passenger_count").size()
evaluate(df_)
df = df[df["passenger_count"] > 0]
evaluate(df)
return df

# 乗車年/月が正しい
@timer
def check_pu_date(df):
df['year'] =  df['tpep_pickup_datetime'].dt.year
df['month'] = df['tpep_pickup_datetime'].dt.month
df['date'] =  df['tpep_pickup_datetime'].dt.day
df['hour'] =  df['tpep_pickup_datetime'].dt.hour

df_ = df.groupby("year").size()
evaluate(df_)
df = df[(df['year'] == 2022) | (df['year'] == 2023)]

df_ = df.groupby("month").size()
evaluate(df_)
df = df[(df['month'] >= 1) & df['month'] <= 12]

df_ = df.groupby(["year", "month"]).size()
evaluate(df_)
evaluate(df)
return df

# 現実的な乗車時間（分）
@timer
def check_ride_time(df):
df["ride_time"] = (df["tpep_dropoff_datetime"] - df["tpep_pickup_datetime"]).dt.seconds / 60
df = df[(df["ride_time"] > 0) & (df["ride_time"] <= 180)]
evaluate(df)
return df

# 現実的な距離
@timer
def check_trip_distance(df):
df = df[(df["trip_distance"] > 0) & (df["trip_distance"] <= 250)]
evaluate(df)
return df

# 現実的な運賃
@timer
def check_total_amount(df):
df = df[(df["total_amount"] > 0) & (df["total_amount"] <= 1000)]
evaluate(df)
return df

# IDから緯度経度を割り出す
@timer
def add_coordinate(df, ID_df):
df = df.merge(ID_df.rename(columns={"longitude": "start_lon", "latitude": "start_lat"}),
left_on="PULocationID", right_on="LocationID", how="left").drop("LocationID", axis=1)
df = df.merge(ID_df.rename(columns={"longitude": "end_lon", "latitude": "end_lat"}),
left_on="DOLocationID", right_on="LocationID", how="left").drop("LocationID", axis=1)
evaluate(df)
return df

# NY内かをcheckする
@timer
def in_NY(df):
df = df[(df["start_lon"] <= -71.47) & (df["start_lon"] >= -79.45)]
df = df[(df["start_lat"] >= 40.29) & (df["start_lat"] <= 45)]
df = df[(df["end_lon"] <= -71.47) & (df["end_lon"] >= -79.45)]
df = df[(df["end_lat"] >= 40.29) & (df["end_lat"] <= 45)]
evaluate(df)
return df

def check_in_NY(df, ID_df):
df = add_coordinate(df, ID_df)
df = in_NY(df)
return df

import os
data_path = "data_sets"
fn = os.path.join(data_path, "taxi_all.csv")
ID_file = os.path.join(data_path, "ID_to_coordinate.csv")

need_cols = ['tpep_pickup_datetime','tpep_dropoff_datetime', 'passenger_count', 'trip_distance',
'PULocationID', 'DOLocationID', 'total_amount', 'improvement_surcharge', 'extra', 'fare_amount', 'RatecodeID']

@timer
def Preprocessing():
df = file_read(fn, {"usecols": need_cols})
ID_df = file_read(ID_file)

df = txt_to_date(df, "tpep_pickup_datetime")
df = txt_to_date(df, "tpep_dropoff_datetime")

df = drop_na(df)
df = check_passenger_c(df)
df = check_pu_date(df)
df = check_ride_time(df)
df = check_trip_distance(df)
df = check_total_amount(df)
df = check_in_NY(df, ID_df)
evaluate(df)
return df

df = Preprocessing()

pip install fireducks

import fireducks.pandas as pd