Joblib: Joblib pickling performance for pandas DataFrame

[rth]

Related to https://github.com/joblib/joblib/issues/467

Below are a few fairly limited benchmarks, that illustrate the serialization / de-serialization performance of pandas.DataFrame with joblib,

import pandas as pd
from sklearn.externals import joblib
import pickle

df = pd.read_csv('http://www.gagolewski.com/resources/data/titanic3.csv',
                 comment='#')
print(df.shape)

# writing to tmpfs to ignore disk I/O cost
%timeit joblib.dump(df, '/dev/shm/df_2.pkl')
%timeit df.to_pickle('/dev/shm/df_1.pkl')
%timeit with open('/dev/shm/df_3.pkl', 'wb') as fh: pickle.dump(df, fh)

produces,

3.17 ms ± 275 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.22 ms ± 48.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
855 µs ± 36.8 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

While for read access,

%timeit joblib.load('/dev/shm/df_2.pkl')
%timeit pd.read_pickle('/dev/shm/df_1.pkl')
%timeit with open('/dev/shm/df_3.pkl', 'rb') as fh: pickle.load(fh)

produces,

3.47 ms ± 344 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
1.4 ms ± 38 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.11 ms ± 9.04 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Run on Linux with Python 3.6 and joblib 0.11. I can run more careful / extensive benchmarks if needed.

So in this particular case, it appears that joblib serialization is almost 3x slower. Since pandas has has a pickle serialization function, I'm wondering how hard would be to use it in joblib.

This is in particular relevant when serializing complex python objects (e.g. scikit-learn pipelines) that can contain numpy arrays and DataFrames. cc @lesteve @jorisvandenbossche

[ogrisel]

Could you try on a larger dataframe? We don't really care to optimize joblib.dump for 3ms workloads and I am not sure the 3x difference will hold on larger datasets.

We cannot use df.to_pickle in a nested datastructure because:

• it requires access to the open file object which is not available when customizing a pickler with the custom reducers;
• it would not make it possible to honor the mmap_mode='r' option at joblib.load time.

[rth]

Thanks for the explanations!

Here are the results with a larger dataset (200MB CSV),

df = pd.read_csv('http://download.geonames.org/export/zip/GB_full.csv.zip',
                 compression='zip', sep='\t', header=0)

Write performance
(first joblib, then pandas read/write and finally cpickle from Python 3)

1.99 s ± 166 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
914 ms ± 14.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
910 ms ± 19.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Read performance

1.23 s ± 24.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
615 ms ± 79.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
619 ms ± 10.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

still twice slower. I can try to find a still bigger dataset if needed. The time to write to SSD with pd.to_pickle is 1.32 s so this would matter less when writing to disk as opposed to serializing in RAM.

[ogrisel]

If we get 2x slower on a workload that lasts on the order of 1s, it means that we have a performance issue. This might be caused by the fact we use a Python-based pickler instead of the C-pickler. Can you try with pickle._dump instead of pickle.dump?

[jjakenichol]

Is there an update on this? It would be nice to use one serializer for my numpy arrays and dataframes.

[cottrell]

Why not just use parquet or better yet pyarrow record batch? I can add this in pretty quickly if someone points to where this should go. I think I have a custom backend somewhere that does this but it really should be the default.

I tend to inspect outputs of functions and allow for depth-0 or depth-1. So return pd.DataFrame is detected or return list of dataframes or return dict of dataframes etc.

If this dict_of_things serialization was in joblib I would probably start using it for most small projects.

Would like to get some positive feedback from the joblib comptrollers before doing any work though.

And as of 2020, pyarrow record batches are fastest deserializations for pandas dataframes.

[Lerbytech]

Still waiting for updates. I'm dealing with dataframes of 2-8gbs and load/dump operations are long from 30 seconds to 2 minutes. I also see 100% load for cpu core as it is the only used core. Willing to do more tests and bring you metrics if it will help

[cottrell]

@Lerbytech You can do your own backend pretty quickly. I have a sketch here that used to work on some version of joblib https://github.com/cottrell/superbasic

[RoyalTS]

Writing pandas/numpy arrays to parquet sounds like an extremely useful piece of functionality. Any chance of cleaning this up and getting it into joblib proper?

[abrichr]

Here's my (unsuccessful) attempt to use @cottrell's https://github.com/cottrell/superbasic:

pip install git+https://github.com/cottrell/superbasic.git pyarrow

from joblib import Memory
import superbasic
memory = Memory('__cache__', backend='superbasic')

Unfortunately this resulted in the following error:

pyarrow.lib.ArrowInvalid: ('Could not convert 80907 with type str: tried to convert to int', 'Conversion failed for column LossZip with type object') 

It seems that pyarrow does not support all pandas types as of time of writing (e.g. see https://github.com/apache/arrow/issues/4168, https://issues.apache.org/jira/browse/ARROW-6222). I avoided the error by casting everything to string before calling the cached function:

df = df.astype(str)
my_cached_function(df)

Unfortunately the result was significantly slower. With default backend:

df.shape=(883798, 11)
...
duration=15.042984962463379

With superbasic backend:

df.shape=(883798, 11)
...
duration=55.95361304283142

(The value of duration is the difference in time between immediately before the cached function is called, and immediately after the cached function is entered.)

Unfortunately for now it seems I will need to avoid caching results of functions parameterized by DataFrames.

[rth]

Unfortunately for now it seems I will need to avoid caching results of functions parameterized by DataFrames.

If the DataFrame is an input of cached functions, you don't actually need fast pickling (to then hash it) -- joblib could potentially just switch to faster hashing of DataFrames directly https://github.com/joblib/joblib/issues/343#issuecomment-770459449 So it's an orthogonal issue.