A Practical Introduction to Polars for Python Data Processing

Polars is a "lightning-fast DataFrame library for Rust and Python" - if you're familiar with Pandas, Polars is quite similar but significantly faster and with a cleaner API. This post is a brief tour of how to set up and use Polars, working through a pragmatic example of the things you might do in a typical data exploration project.

The Setup

Install polars and jupyter:

python3 -m venv venv
source venv/bin/activate
pip install polars notebook

This sets up a python virtual environment, then installs the Polars library as well as jupyter. The page you're looking at was created as a jupyter notebook.

The Data

I'm using the New York Taxi Dataset parquet files in this notebook. I downloaded the October 2022 Yellow Taxi Trip Records (PARQUET) file, feel free to grab that as well.

The data is in parquet format, which is quite a nice format for data processing, significantly better than CSV.

Polars

Let's import polars and read the file:

import polars as pl

fname = "yellow_tripdata_2022-10.parquet"
data = pl.read_parquet(fname)

Let's see what it looks like

data.head(5)

What are the columns in this file?

data.columns

['VendorID',
 'tpep_pickup_datetime',
 'tpep_dropoff_datetime',
 'passenger_count',
 'trip_distance',
 'RatecodeID',
 'store_and_fwd_flag',
 'PULocationID',
 'DOLocationID',
 'payment_type',
 'fare_amount',
 'extra',
 'mta_tax',
 'tip_amount',
 'tolls_amount',
 'improvement_surcharge',
 'total_amount',
 'congestion_surcharge',
 'airport_fee']

Pretty easy to guess what the columns mean.

Let's poke around a little bit - what's the average taxi fare in NYC?

data["total_amount"].mean()

22.247540895418666

I wonder how many passengers there are in each ride?

data["passenger_count"].mean()

1.3846936672698376

How about the average fare from the airport? I noticed the airport_fee column, let's look at average total_amount when the airport fee is greater than zero:

data.filter( pl.col('airport_fee') > 0 )["total_amount"].mean()

55.83249747837022

How about the average non-airport ride cost?

data.filter( pl.col('airport_fee') == 0 )["total_amount"].mean()

18.957161390687517

Most expensive rides

Let's see what the most expensive fares in this data set are. We'll sort the data by total_amount and look at the top 10.

most_expensive = data\
    .sort("total_amount", reverse=True)\
    .head(10)\
    .select(["total_amount", "PULocationID", "DOLocationID"])

most_expensive

I'm selecting just the 3 columns I'm interested in so it's easier to look at the data - the amount and pick up and drop off locations.

Note the coding style above: I split the query into separate lines so it's easier to read. I've become a stickler when it comes to query / code style - you'll thank yourself for keeping it clean when you come back months later and try to figure out what you were doing.

We can make it even cleaner if we use the paranthesis syntax:

most_expensive = (
    data
    .sort("total_amount", reverse=True)
    .head(10)
    .select(["total_amount", "PULocationID", "DOLocationID"])
)

I wasn't crazy about the above syntax at first, but once I started using it I really liked it.

Also notice the style of computation here - we're chaining the actions, one after another, javascript style: sort, then take 10 rows, then select the columns we want. This I also like.

Looking at the data, apparently somebody paid \$1,301.85 for a single ride. I wonder where they went to? Let's find out.

PULocationID and DOLocationID are the pick up and drop off location IDs. They're just numbers right now, not too informative. Let's convert them to names.

Zone Names

We need a mapping from the Zone IDs to the Zone names. From the NYC TLC Data page, look for the Taxi Zone Lookup Table (CSV) link and download the CSV file.

zone_data = pl.read_csv("taxi+_zone_lookup.csv")
zone_data.head(10)

Ok, that's helpful, but I'd rather have the Zone and Borough together in a single column so it has all the info I need.

zone_data = (
    zone_data
    .with_column(
        (pl.col("Borough") + ": " + pl.col("Zone")).alias("Location"),
    )
    .select(["LocationID", "Location"])
)

zone_data.head(10)

Ok, that's looking pretty good. Now let's look at those expensive rides again, but add in the pick up and drop off location names.

To do this I'll join the most_expensive dataframe I created above, which just had the fare and pick up and drop off IDs, with the zone_data dataframe which has the zone names. First I'll join the PULocationID with the zone data, add the location name, and rename the column from Location to PickUpLocation, then I'll do the same with the DOLocationID.

The left join below means use the rows from the left dataframe, which is the most_expensive dataframe, and injects into those rows the Location names from the zone_data dataframe, joining on LocationID.

most_expensive_rides = (
    most_expensive
    .join(zone_data, how="left", left_on="PULocationID", right_on="LocationID").rename({"Location": "PickUpLocation"})
    .join(zone_data, how="left", left_on="DOLocationID", right_on="LocationID").rename({"Location": "DropOffLocation"})
)
most_expensive_rides

The location names are getting cut off when we view the dataframe - let's configure the output to show wider columns

pl.Config. set_fmt_str_lengths(100)

most_expensive_rides

The Unknown: NA locations are not very useful, let's get rid of them. I'll filter the data so they're excluded.

Location ID of 256 means Unknown: NA, so we need to get rid of any row that has 256 in the PULocationID or DOLocationID columns.

clean_data = (
    data
    .filter( 
        (pl.col("PULocationID") != 265) &
        (pl.col("DOLocationID") != 265)
    )
)
clean_data.head(5)

Ok, now we can look at the most expensive rides again, with the Unknown locations removed

most_expensive = (
    clean_data
    .sort("total_amount", reverse=True)
    .head(10)
    .select(["total_amount", "PULocationID", "DOLocationID"])
)

most_expensive

Let's add the location names back in. I'll create a function to do this so we can use it whenever we need.

def add_location_names(data):
    return (
        data
        .join(zone_data, how="left", left_on="PULocationID", right_on="LocationID").rename({"Location": "PickUpLocation"})
        .join(zone_data, how="left", left_on="DOLocationID", right_on="LocationID").rename({"Location": "DropOffLocation"})
    )

add_location_names(most_expensive)

Hmm. Somebody paid \$950.3 to go from Queens: Sunnyside to Queens: Sunnyside. Same for Woodside. There was also an \$800 ride from Lenox Hill East to Sutton Place, which according to Google is a 4 minute car ride.

I'm starting to think our data might have junk in it. Exciting - finding junk in you data means you're doing Data Science™. Now we just have to throw in some Excel and we'll have Machine Learning™.

There's a lot we could do from here. For example, we could calculate the average fare by pick up and drop off location. Or we can look at the impact of pick up time on fares (there's a tpep_pickup_datetime column in the data).

For some reason I decided it'd be interesting to see which pickup-dropoff location pairs have the highest variation in fare. We can figure this out by grouping the data by ("PULocationID", "DOLocationID"), then calculating the standard deviation for each grouping, and looking at the locations with the highest deviation.

amount_stddevs = (
    clean_data
    .groupby(["PULocationID", "DOLocationID"])
    .agg(
        [
            pl.col("total_amount").std().alias("total_amount_stddev")
        ]
    )
    .sort("total_amount_stddev", reverse=True)
    .head(20)

)

amount_stddevs

Let's add the location names:

add_location_names(amount_stddevs)

That's... less interesting than I thought it would be. I suppose the absolute values are pretty interesting - our top hit, Queens: Bayside to Queens: Oakland Gardens, is about a 5 minute car ride and has a standard deviatiion of \$220. That seems suspect.

This much standard deviation means we have bad data.

Data Exploration and Cleansing

A lot of real-world data work involves cleaning data and finding outliers, so let's try that with this data.

There are a number of techniques for finding outliers, but first let's take a better look at our data. Polars has a built-in describe function that calculates basic statistics for your dataframe:

data.describe()

That's pretty helpful - for example, the minimum pick up date is 2008-12-31, which means this data starts all the way back to 2009. Because of the way the TLC data download page is organized I thought I only had data from October 2023.

The largest tip amount is \$500 (max of the tip_amount column). The minimum tip is \$-100.01. Must've been a pretty rough ride.

Looking at the total_amount column we see that the maximum ride amount was \$1301.85 and the minimum was -\$1301.85. I don't know much about this data, but I'm going to guess that's an incorrectly charged ride and a refund. That's unfortunate - having large invalid data (\$1301.85 and -\$1301.85) at either extreme of values will wreak havoc on our standard deviations and such.

The longest ride (max of the trip_distance field) is 389678.46. Is that in miles? Let's look at the data dictionary provided by the TLC.

It is in fact in miles. The length of the United States is about 3000 miles, so this enterprising taxi travelled the entire US about 130 times.

The data dictionary tells us a lot of other useful things - for example, the Payment_type column should help us find disputed, voided, and other invalid charges.

Basic Cleaning

So what should we do? I'm thinking we look at more recent data, say anything 2020 and after. Let's get rid of anything with more than 3000 miles, and let's also get rid of rows with unreasonable negative numbers - eg. negative total amounts, tolls, etc. We can get rid of rows where the Payment_type is No Charge, Dispute, Voided, etc.

The trouble is if we just get rid of the negative amounts and disputes then we're leaving the corresponding equally invalid positive amount in the data - eg. if we drop the row that has the -\$1301.85 amount we'll still have the \$1301.85 row in our data.

We'll have to figure out how to drop both: for each of the negative numbers, we'll need to find the corresponding positive amount and get rid of those rows too.

Let's take a look at the outlier data and see if spot any clues for how to get rid of them.

(
    data
    .filter(
        (pl.col("total_amount") > 1300) |
        (pl.col("total_amount") < -1300)
    )
)

The payment type is 2, which is cash. Somebody paid $1300 in cash, then got it handed back to him. I was hoping the payment type would be Voided or something reasonable.

Getting rid of these invalid pairs is going to be a bit painful - we'll have to find each invalid ride and also find its pair.

Before we do that, let's see if the problems still exists in data we care about, which is data from 2020 onwards.

data.filter( pl.col("tpep_pickup_datetime") > "2020" ).describe()

Warning: Comparing date/datetime/time column to string value, this will lead to string comparison and is unlikely what you want.
If this is intended, consider using an explicit cast to silence this warning.

That didn't work - our min tpep_pickup_datetime is 2022-09-30 (I was expecting 2020). I guess I can't filter a datetime column (tpep_pickup_datetime) using a string ("2020"). Makes sense, let's fix it.

from datetime import datetime
start_date = datetime(2020,1,1,0,0,0)
print(f"Filtering data to start at {start_date}")
data.filter( pl.col("tpep_pickup_datetime") > start_date ).describe()

Filtering data to start at 2020-01-01 00:00:00

My min tpep_pickup_datetime is still 2022-09-30 (I was expecting 2020). What's going on? Maybe our data is only for Octover 2022 after all, and the 2008 date I was seeing was spurious? Note that we had 3675411 data points before we filtered by date and now we have 3675407 data points (I'm getting this by looking at the count field for tpep_pickup_datetime). Our date filtering only got rid of 4 data points. So we're looking only at October 2022 data with a couple of spurious data points thrown in.

I'm curious what the data with bad dates look like.

(
    data
    .filter(
        (pl.col("tpep_pickup_datetime") < start_date)
    )
)

Finding bad pairs

I don't see any obvious ways to correlate the bad pairs - the rows don't have IDs, there is no cab ID, and the payment method field doesn't give us much either. We'll have to data science.

Let's try this: - Find negative data points - For each negative data point, find a data point with the same amount but positive that also matches other parameters such as date and pick up / drop off location

(
    data
    .filter( pl.col("fare_amount") < 0)
)

There's probably a clever way to do this, but my initial thought for getting rid of the corresponding positives is to create a new dataframe with the negative rows, then do an anti-join back to the original dataframe. An anti-join should find matching rows and remove them.

negative_data = (
    data
    .filter( pl.col("fare_amount") < 0)
).select(["tpep_pickup_datetime", "PULocationID", "fare_amount"])

negative_data

I'm only grabbing a few columns for my negative_data dataframe, since all I'm going to use it for is to remove the corresponding positive rows from the original data. I'll try matching on the pick up date and pick up location.

Let's convert the negative amounts to positive so we can directly match them to their corresponding positive rows in the original data.

negative_data = negative_data.with_columns([
    (pl.col("fare_amount") * -1).alias("amount_positive")
])

negative_data

Now that we have a dataframe with the negative amounts converted to positive, we can join back to the original data, matching on the (now positive) fare amount and the pick up location. We'll do an anti-join, which gets rids of the rows that match. If Polars didn't have anti-join as an option we could have merged the new dataframe with the original data and dropped duplicate rows.

data.join(negative_data, on=("PULocationID", "fare_amount"), how="anti")

Did that work? Let's see how many rows that got rid of. After the anti-join we end up with 3649512 rows. We had and 25899 rows with negative values.

n_rows_original = data.shape[0]
n_rows_negative = negative_data.shape[0]
print(f"We are expecting {n_rows_original - n_rows_negative} rows if our removal worked perfectly")

We are expecting 3649512 rows if our removal worked perfectly

Sure enough the result of the anti-join has 3649512 rows (you can see that from the shape output a few cells above), so our method is working well. We should sanity check to make sure it's working as expected, but for now let's assume all is well.

We can now do our final cleaning: we'll get rid of the negative and corresponding positive rows, get rid of anything where the mileage is over 3000 miles, remove dates that are out of range, and get rid of rows with unknown pick up and drop off locations:

clean_data = (
    data
    .join(negative_data, on=("PULocationID", "fare_amount"), how="anti")
    .filter( pl.col("fare_amount") >= 0)
    .filter( pl.col("trip_distance") <= 3000 )
    .filter( pl.col("tpep_pickup_datetime") >= start_date )
    .filter( 
        (pl.col("PULocationID") != 265) &
        (pl.col("DOLocationID") != 265)
    )
)

clean_data

How many rows of data did we get rid of?

data.shape[0] - clean_data.shape[0]

45288

clean_data.describe()

That's looking better - we still have some negatives and other strange looking things, but you get a sense of how we'd approach getting rid of those. For next steps I'd probably filter more things that look incorrect, then I'd plot the data as a histogram and eyeball it to see if there's anything strange in there.

Let's stop there.

This was my first foray into Polars, so it's likely there are more polarsy ways of doing things - if you know of improvements please let me know.

Overall I really like Polars. It's fast and clean and the resulting code is easy on the eyes. I'm likely to use it in place of Pandas going forward.

Links

Modern Polars is quite nice with many examples.

Towards Data Science has a good article called Polars: Pandas DataFrame but Much Faster.

This article from Code Magazine has a lot of good Polars examples.

The Polars documentation page on Coming from Pandas is very useful if you have Pandas experience.

While we're at it, read this piece from a former cab driver on the ride he'll never forget.