Matthew Rocklin: Pandas Categoricals

tl;dr: Pandas Categoricals efficiently encode and dramatically improve performance on data with text categories

Disclaimer: Categoricals were created by the Pandas development team and not by me.

There is More to Speed Than Parallelism

I usually write about parallelism. As a result people ask me how to parallelize their slow computations. The answer is usually just use pandas in a better way

• Q: How do I make my pandas code faster with parallelism?
• A: You don’t need parallelism, you can use Pandas better

This is almost always simpler and more effective than using multiple cores or multiple machines. You should look towards parallelism only after you’ve made sane choices about storage format, compression, data representation, etc..

Today we’ll talk about how Pandas can represent categorical text data numerically. This is a cheap and underused trick to get an order of magnitude speedup on common queries.

Categoricals

Often our data includes text columns with many repeated elements. Examples:

• Stock symbols – GOOG, APPL, MSFT, ...
• Gender – Female, Male, ...
• Experiment outcomes – Healthy, Sick, No Change, ...
• States – California, Texas, New York, ...

We usually represent these as text. Pandas represents text with the object dtype which holds a normal Python string. This is a common culprit for slow code because object dtypes run at Python speeds, not at Pandas’ normal C speeds.

Pandas categoricals are a new and powerful feature that encodes categorical data numerically so that we can leverage Pandas’ fast C code on this kind of text data.

>>># Example dataframe with names, balances, and genders as object dtypes>>>df=pd.DataFrame({'name':['Alice','Bob','Charlie','Danielle'],...'balance':[100.0,200.0,300.0,400.0],...'gender':['Female','Male','Male','Female']},...columns=['name','balance','gender'])>>>df.dtypes# Oh no!  Slow object dtypes!nameobjectbalancefloat64genderobjectdtype:object

We can represent columns with many repeats, like gender, more efficiently by using categoricals. This stores our original data in two pieces

• Original data

   Female, Male, Male, Female
  

1. Index mapping each category to an integer

   Female: 0
   Male: 1
   ...
   

2. Normal array of integers

   0, 1, 1, 0
   

This integer array is more compact and is now a normal C array. This allows for normal C-speeds on previously slow object dtype columns. Categorizing a column is easy:

In[5]:df['gender']=df['gender'].astype('category')# Categorize!

Lets look at the result

In[6]:df# DataFrame looks the sameOut[6]:namebalancegender0Alice100Female1Bob200Male2Charlie300Male3Danielle400FemaleIn[7]:df.dtypes# But dtypes have changedOut[7]:nameobjectbalancefloat64gendercategorydtype:objectIn[8]:df.gender# Note Categories at the bottomOut[8]:0Female1Male2Male3FemaleName:gender,dtype:categoryCategories(2,object):[Female,Male]In[9]:df.gender.cat.categories# Category indexOut[9]:Index([u'Female',u'Male'],dtype='object')In[10]:df.gender.cat.codes# Numerical valuesOut[10]:00112130dtype:int8# Stored in single bytes!

Notice that we can store our genders much more compactly as single bytes. We can continue to add genders (there are more than just two) and Pandas will use new values (2, 3, …) as necessary.

Our dataframe looks and feels just like it did before. Pandas internals will smooth out the user experience so that you don’t notice that you’re actually using a compact array of integers.

Performance

Lets look at a slightly larger example to see the performance difference.

We take a small subset of the NYC Taxi dataset and group by medallion ID to find the taxi drivers who drove the longest distance during a certain period.

In[1]:importpandasaspdIn[2]:df=pd.read_csv('trip_data_1_00.csv')In[3]:%timedf.groupby(df.medallion).trip_distance.sum().sort(ascending=False,inplace=False).head()CPUtimes:user161ms,sys:0ns,total:161msWalltime:175msOut[3]:medallion1E76B5DCA3A19D03B0FB39BCF2A2F534870.836945300E90C69061B463CCDA370DE5D6832.914F4BEA1914E323156BE0B24EF8205B73811.99191115180C29B1E2AF8BE0FD0ABD138F787.33B83044D63E9421B76011917CE280C137782.78Name:trip_distance,dtype:float64

That took around 170ms. We categorize in about the same time.

In[4]:%timedf['medallion']=df['medallion'].astype('category')CPUtimes:user168ms,sys:12.1ms,total:180msWalltime:197ms

Now that we have numerical categories our computaion runs 20ms, improving by about an order of magnitude.

In[5]:%timedf.groupby(df.medallion).trip_distance.sum().sort(ascending=False,inplace=False).head()CPUtimes:user16.4ms,sys:3.89ms,total:20.3msWalltime:20.3msOut[5]:medallion1E76B5DCA3A19D03B0FB39BCF2A2F534870.836945300E90C69061B463CCDA370DE5D6832.914F4BEA1914E323156BE0B24EF8205B73811.99191115180C29B1E2AF8BE0FD0ABD138F787.33B83044D63E9421B76011917CE280C137782.78Name:trip_distance,dtype:float64

We see almost an order of magnitude speedup after we do the one-time-operation of replacing object dtypes with categories. Most other computations on this column will be similarly fast. Our memory use drops dramatically as well.

Conclusion

Pandas Categoricals efficiently encode repetitive text data. Categoricals are useful for data like stock symbols, gender, experiment outcomes, cities, states, etc.. Categoricals are easy to use and greatly improve performance on this data.

We have several options to increase performance when dealing with inconveniently large or slow data. Good choices in storage format, compression, column layout, and data representation can dramatically improve query times and memory use. Each of these choices is as important as parallelism but isn’t overly hyped and so is often overlooked.

Jeff Reback gave a nice talk on categoricals (and other featuress in Pandas) at PyData NYC 2014 and is giving another this weekend at PyData London.