Polars vs Pandas

data_folder = '/Users/alexlee/Desktop/Data'
filename_gnaf = f'{data_folder}/geo/gnaf_feb2024/GNAF_CORE.psv'
filename_gnaf_parquet = f'{data_folder}/geo/gnaf_feb2024/GNAF_CORE.parquet'

gnaf = pd.read_csv(filename_gnaf, sep='|')
gnaf.shape

# polars
t1 = time()
gnaf = pl.read_csv(filename_gnaf, separator='|', infer_schema_length=10000)
t2 = time()
print(t2-t1)

# pandas
t1 = time()
gnaf = pd.read_csv(filename_gnaf, sep='|')
t2 = time()
print(t2-t1)

gnaf_polars = pl.read_csv(filename_gnaf, separator='|', infer_schema_length=10000)
gnaf_pandas = pd.read_csv(filename_gnaf, sep='|')

# polars
t1 = time()
(
    gnaf_polars
    .filter(pl.col('STATE') == "NSW")
    .select(pl.col("STREET_NAME"))
    .with_columns(pl.col("STREET_NAME").str.slice(0, 1).alias("letter"))
    .group_by('letter')
    .count()
    .sort('count', descending=True)
)
t2 = time()
print(t2-t1)

# pandas
t1 = time()
(
    gnaf_pandas
    .query('STATE == "NSW"')
    .STREET_NAME
    .str.slice(0, 1)
    .value_counts()
)
t2 = time()
print(t2-t1)

# polars
t1 = time()
(
    gnaf_polars
    .filter(pl.col("POSTCODE") == 3000)
    .group_by("STREET_NAME")
    .len()
    .sort('len', descending=True)
)
t2 = time()
print(t2 - t1)

# pandas
t1 = time()
results = (
    gnaf_pandas
    .query('POSTCODE == 3000')
    .groupby('STREET_NAME')
    .size()
    .sort_values(ascending=False)
)
t2 = time()
print(t2-t1)

# load in pandas
t1 = time()
gnaf = pd.read_parquet(filename_gnaf_parquet)
t2 = time()
print(t2-t1)

# load in polars
t1 = time()
gnaf = pl.read_parquet(filename_gnaf_parquet)
t2 = time()
print(t2-t1)

gnaf = pl.scan_parquet(filename_gnaf_parquet)
t1 = time()
(
    gnaf
    .filter(pl.col("POSTCODE") == 3000)
    .group_by("STREET_NAME")
    .len()
    .sort('len', descending=True)
    .collect()
)
t2 = time()
print(t2 - t1)

t1 = time()
(
    gnaf
    .filter(pl.col('STATE') == "NSW")
    .select(pl.col("STREET_NAME"))
    .with_columns(pl.col("STREET_NAME").str.slice(0, 1).alias("letter"))
    .group_by('letter')
    .len()
    .sort('len', descending=True)
    .collect()
)
t2 = time()
print(t2-t1)

from os import listdir

nyc_folder = '/Users/alexlee/Desktop/Data'

filenames_nyc = f'{data_folder}/test_data/nyc_taxi'

all_files = listdir(filenames_nyc)

all_files = sorted(all_files)

# column names of 2023 data so that they match the previous years
for filename in all_files[-11:]:
    print(filename)
    nyc_data = pl.read_parquet(f'{filenames_nyc}/{filename}').rename({'Airport_fee': 'airport_fee'})
    nyc_data.write_parquet(f'{filenames_nyc}/{filename}')

# consider only data from 2011 as the format changed

# ensure that the columns of all files are the data types
# probably not the best way of doing this
for filename in all_files:
    nyc_data = pl.scan_parquet(f'{filenames_nyc}/{filename}')
    nyc_data.select(
        pl.col('VendorID').cast(pl.Int64),
        pl.col('tpep_pickup_datetime').cast(pl.Datetime),
        pl.col('tpep_dropoff_datetime').cast(pl.Datetime),
        pl.col('passenger_count').cast(pl.Int64),
        pl.col('trip_distance').cast(pl.Float64),
        pl.col('RatecodeID').cast(pl.Int64),
        pl.col('store_and_fwd_flag').cast(pl.String),
        pl.col('PULocationID').cast(pl.Int64),
        pl.col('DOLocationID').cast(pl.Int64),
        pl.col('payment_type').cast(pl.Int64),
        pl.col('fare_amount').cast(pl.Float64),
        pl.col('extra').cast(pl.Float64),
        pl.col('mta_tax').cast(pl.Float64),
        pl.col('tip_amount').cast(pl.Float64),
        pl.col('tolls_amount').cast(pl.Float64),
        pl.col('improvement_surcharge').cast(pl.Float64),
        pl.col('total_amount').cast(pl.Float64),
        pl.col('congestion_surcharge').cast(pl.Float64),
        pl.col('airport_fee').cast(pl.Float64)
    ).collect().write_parquet(f'{filenames_nyc}/{filename}')

# load the data folder
nyc_taxis = pl.scan_parquet(f'{filenames_nyc}/*.parquet')

# how many rows in the data?
nyc_taxis.select(pl.len()).collect()

# what is the max fare?
nyc_taxis.select(pl.max('fare_amount')).collect()

# what is the average fare?
nyc_taxis.select(pl.mean('fare_amount')).collect()

# what is the median fare?
nyc_taxis.select(pl.median('fare_amount')).collect()

# what is the distribution of passenger counts?
nyc_taxis.group_by('passenger_count').len().sort('len', descending=True).collect().head(10)

nyc_taxis.columns

# how many different pickup locations are there?
nyc_taxis.group_by('PULocationID').len().collect()

# let's see how people are paying over time
payment_types_year = (
    nyc_taxis
    .with_columns(pl.col('tpep_pickup_datetime').dt.year().alias('pickup_year'))
    .group_by(pl.col(['pickup_year', 'payment_type']))
    .len()
    .filter(pl.col('pickup_year').is_in(list(range(2011,2024))))
    .filter(pl.col('payment_type').is_in([1,2]))
    .sort(pl.col('pickup_year'))
    .collect()
)

plt.figure(figsize=(6, 4))
sns.lineplot(data=payment_types_year.to_pandas(), x='pickup_year', y='len', hue='payment_type')
plt.xlabel('Year')
plt.ylabel('Number of trips')
plt.tight_layout()
plt.savefig('payments_over_time_nyc_taxis.png')