Efficient Pandas serialization

[mrocklin]

We need an efficient way to serialize Pandas Dataframes. As of #606 we can now customize this beyond just pickle. There are a number of different data regimes here including pure numeric data, very compressible numeric data like time series, text data with repeats, categoricals, long text data etc.. We care both about fast encoding and about fast compression for larger results. We want something that has minimal overhead on small dataframes (important for shuffling).

Several options come to mind:

1. Pickle + lz4: what we use now. Actually not that bad. Pickle deduplicates text data, handles numeric data without too much overhead, etc.. It's a shame that we can't compress per type.
2. NumPy serialization + blosc on blocks + pickle or msgpack + lz4 on text: This is similar to what we do in partd. A custom approach can get significant speed boosts in some cases.
3. Arrow: we could give arrow a shot. I suspect that it handles encoding intelligently. It would be nice if it were to provide different compressions per column but we could probably push a bit upstream if this went into use.
4. Other thoughts?

@jreback @wesm @shoyer

It would be good to make this decision with a benchmark in hand. It would be good both to get people's opinions on solutions we should consider as well as some benchmarks that are representative of data that they care about.

[wesm]

I've spent a lot of time lately working on DataFrame serialization with as little overhead as possible. I'm happy to assist as best I can with development to support this use case and benchmarking to support the decision making. I recently did some rudimentary benchmarking on uncompressed numeric data, see slide 24 in

http://www.slideshare.net/wesm/python-data-wrangling-preparing-for-the-future

the benchmarking code (which I can clean up and publish if so desired) looks like:

import pyarrow as pa
import pyarrow.io as io
import pyarrow.ipc as ipc

def write_arrow(df):
    batch = pa.RecordBatch.from_pandas(df)

    buf = BytesIO()
    writer = ipc.ArrowFileWriter(buf, batch.schema)
    writer.write_record_batch(batch)
    writer.close()
    return buf.getvalue()


def read_arrow(buf):
    reader = ipc.ArrowFileReader(buf)
    return [reader.get_record_batch(i).to_pandas()
            for i in range(reader.num_record_batches)]

Here, on the read path, this is all zero-copy up until converting to pandas.DataFrame, which introduces some overhead. This overhead could be reduced through the slightly messy process of deserializing directly into a 2D ndarray and then constructing pandas's internal BlockManager data structure. this will become much more favorable with pandas 2.0 (can be fully zero-copy if we use bitmaps for all null handling, see discussion in wesm/pandas2#46).

For compressible data you can obviously add buffer-level compression -- see https://issues.apache.org/jira/browse/ARROW-300. So in other words, if you have an array with a null bitmap buffer and a contiguous buffer, these buffers would be compressed independently on going into the IPC blob.

Dictionary encoding isn't yet implemented, but it's provided for in the IPC metadata.

Maybe this is pie in the sky, but at some point it may be in dask's interests to adopt a language independent serialization protocol so there is the option to run task worker code written in C++, Julia, or other languages.

[mrocklin]

Hi @wesm . Thanks for the read and write functions. These should be easy to benchmark and easy to incorporate as an optional dependency (hooray for Arrow being on conda-forge by the way). @jreback any interest in trying these out in the wild? You have some communication heavy computations.

Zero copy sounds nice. So too does compression on different components. It'll be nice to try some workloads in the wild though to see how much impact these have. That might help us to prioritize or deprioritize them. I have not yet run into workloads that are limited by single memory copies. Network bandwidth is more commonly a larger issue for me, at least on Amazon's hardware.

Maybe this is pie in the sky, but at some point it may be in dask's interests to adopt a language independent serialization protocol so there is the option to run task worker code written in C++, Julia, or other languages.

It's decently possible and there is some slow momentum towards that direction. The custom serialization system is a step in that direction. There are other users clamoring for destination-specific serialization (e.g. gpu-gpu transfers) that would also benefit from this. Presumably you would say "OK, this data of type T has to move from worker A to B, what is the best serialization that we have for this case?" If there were a use case today the serialization work could happen in about a week. I put the cost of making a decent Julia/R/whatever worker/client at "a couple of weeks" by someone familiar with the networking stack in that language. See #586

[jreback]

sure, i'll give this a go. IIRC you are going to have this settable via dask.set_options()?
(or someway completely hacky ok), just so I can go back and forth to try things out?

though since on win-64.... will wait for @wesm to kindly make packages available for pyarrow :>

[mrocklin]

Ah, forgot about the Windows thing. You should wait to try this until I give you the go ahead. I'm fairly comfortable putting Arrow into an optional plugin after #606 gets merged. There will be a few of these for other formats as well.

[mrocklin]

Initial timings comparing Pickle and Arrow. Corrections or counter-benchmarks welcome. Based on this very narrow benchmark custom serialization + blosc is winning out.

In [1]: import pyarrow as pa
   ...: import pyarrow.io as io
   ...: import pyarrow.ipc as ipc
   ...: 
   ...: def write_arrow(df):
   ...:     batch = pa.RecordBatch.from_pandas(df)
   ...: 
   ...:     buf = BytesIO()
   ...:     writer = ipc.ArrowFileWriter(buf, batch.schema)
   ...:     writer.write_record_batch(batch)
   ...:     writer.close()
   ...:     return buf.getvalue()
   ...: 
   ...: 
   ...: def read_arrow(buf):
   ...:     reader = ipc.ArrowFileReader(buf)
   ...:     return [reader.get_record_batch(i).to_pandas()
   ...:             for i in range(reader.num_record_batches)]
   ...: 

In [2]: from io import BytesIO

In [3]: import string

In [4]: import numpy as np

In [5]: import pandas as pd

In [6]: n = 1000000

In [7]: df = pd.DataFrame({'x': np.random.randint(0, 10000, size=n, dtype='i4'), 'y': np.random.random(n), 
   ...:  't': np.random.choice(list(string.ascii_lowercase), size=n)})

In [8]: import pickle

In [9]: %time len(write_arrow(df))
CPU times: user 120 ms, sys: 4 ms, total: 124 ms
Wall time: 121 ms
Out[9]: 17000630

In [10]: %time len(pickle.dumps(df))
CPU times: user 36 ms, sys: 8 ms, total: 44 ms
Wall time: 42.7 ms
Out[10]: 14003232

In [11]: %time len(write_arrow(df[['x']]))
CPU times: user 4 ms, sys: 0 ns, total: 4 ms
Wall time: 3.4 ms
Out[11]: 4000334

In [12]: %time len(pickle.dumps(df[['x']]))
CPU times: user 0 ns, sys: 0 ns, total: 0 ns
Wall time: 2.85 ms
Out[12]: 4000781

In [13]: import blosc

In [14]: x = df.x.values  # There is an easier way to do this

In [15]: %time len(blosc.compress(x.data, typesize=x.dtype.itemsize, cname='lz4', clevel=5))
CPU times: user 8 ms, sys: 0 ns, total: 8 ms
Wall time: 3.72 ms
Out[15]: 2020787

In [16]: %time len(pickle.dumps(df[['t']]))
CPU times: user 44 ms, sys: 0 ns, total: 44 ms
Wall time: 45.5 ms
Out[16]: 2002878

In [17]: %time len(write_arrow(df[['t']]))
CPU times: user 96 ms, sys: 0 ns, total: 96 ms
Wall time: 95 ms
Out[17]: 5000406

[wesm]

How did you install pyarrow? I believe you are using a debug build (the conda-forge artifact is a debug build, I can fix this to allow you to do a more apples-to-apples comparison)

[mrocklin]

conda create -n arrow-test -c conda-forge python=3 pyarrow numpy pandas ipython -y

I may have pip installed some thing in between

(arrow-test) mrocklin@carbon:~$ conda list -e
# This file may be used to create an environment using:
# $ conda create --name <env> --file <this file>
# platform: linux-64
Using Anaconda Cloud api site https://api.anaconda.org
arrow-cpp=0.1.post=1
backports.shutil_get_terminal_size=1.0.0=py35_0
blas=1.1=openblas
ca-certificates=2016.8.31=0
certifi=2016.8.31=py35_0
decorator=4.0.10=py35_0
ipython=5.1.0=py35_1
ipython_genutils=0.1.0=py35_0
ncurses=5.9=9
numpy=1.11.2=py35_blas_openblas_200
openblas=0.2.18=5
openssl=1.0.2h=2
pandas=0.19.0=np111py35_0
parquet-cpp=0.1.pre=3
pexpect=4.2.1=py35_0
pickleshare=0.7.3=py35_0
pip=8.1.2=py35_0
prompt_toolkit=1.0.8=py35_0
ptyprocess=0.5.1=py35_0
pyarrow=0.1.post=0
pygments=2.1.3=py35_1
python=3.5.2=2
python-dateutil=2.5.3=py35_0
pytz=2016.7=py35_0
readline=6.2=0
setuptools=26.1.1=py35_0
simplegeneric=0.8.1=py35_0
six=1.10.0=py35_1
sqlite=3.13.0=1
tk=8.5.19=0
traitlets=4.3.0=py35_0
wcwidth=0.1.7=py35_0
wheel=0.29.0=py35_0
xz=5.2.2=0
zlib=1.2.8=3
libgfortran=3.0.0=1
lz4=0.8.2=py35_0

[wesm]

Yes, you're on a debug build of pyarrow. I'll let you know when there's a new build available built with optimizations turned on (this build is gcc -O0)

[wesm]

Digging into a memory leak bug (ARROW-362) but will get a release build up on conda-forge after that's fixed

[wesm]

PyArrow release builds are up. You ideally should to disable multithreading in blosc to make the benchmarks comparable -- we could add multithreaded writes to Arrow, too, and get similar speedups

tmp = df[['x']]  # this is a copy -- should not be included in benchmark

import blosc

blosc.set_nthreads(1)  # multithreading is not apples-to-apples
x = df.x.values

df_t = df[['t']]  # this is a copy -- should not be included in benchmark

%timeit len(write_arrow(df))

%timeit len(pickle.dumps(df))
%timeit len(write_arrow(tmp))
%timeit len(pickle.dumps(tmp))
%timeit len(blosc.compress(x.data, typesize=x.dtype.itemsize, cname='lz4', clevel=5))
%time len(pickle.dumps(df_t))
%time len(write_arrow(df_t))

In [27]: %timeit len(write_arrow(df))
    ...: 
10 loops, best of 3: 27.3 ms per loop

In [28]: %timeit len(pickle.dumps(df))
    ...: 
10 loops, best of 3: 24.4 ms per loop

In [29]: %timeit len(pickle.dumps(tmp))
    ...: 
The slowest run took 5.50 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 626 µs per loop

In [30]: %timeit len(write_arrow(tmp))
    ...: 
The slowest run took 6.02 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 689 µs per loop

In [31]: %timeit len(pickle.dumps(tmp))
^[[41;1R
1000 loops, best of 3: 627 µs per loop

In [34]: %timeit len(blosc.compress(x.data, typesize=x.dtype.itemsize, cname='lz4', clevel=5))
    ...: 
The slowest run took 5.77 times longer than the fastest. This could mean that an intermediate result is being cached.
1000 loops, best of 3: 683 µs per loop

In [36]: %timeit len(pickle.dumps(df_t))
10 loops, best of 3: 26.1 ms per loop

In [37]: %timeit len(write_arrow(df_t))
100 loops, best of 3: 17.7 ms per loop

if you change the characteristics of the string data, you get different results:

In [96]: paste
from pandas.util.testing import rands
string_choices = [rands(10) for i in range(1000)]
df = pd.DataFrame({'x': np.random.randint(0, 10000, size=n, dtype='i4'), 'y': np.random.random(n), 
   't': np.random.choice(string_choices, size=n)})
## -- End pasted text --

In [97]: %timeit len(write_arrow(df))
10 loops, best of 3: 47.1 ms per loop

In [98]: %timeit len(pickle.dumps(df))
10 loops, best of 3: 164 ms per loop

In [99]: df_t = df[['t']]  # this is a copy -- should not be included in benchmark

In [100]: %timeit len(pickle.dumps(df_t))
10 loops, best of 3: 152 ms per loop

In [101]: %timeit len(write_arrow(df_t))
10 loops, best of 3: 40.1 ms per loop

there's some performance loss on the Arrow read path right now because we aren't constructing the precise pandas BlockManager. If we wanted to hyperoptimize for the pandas 1.x memory layout this would be a nice project for someone to tackle (multithreaded blockmanager read/write)

note you can also go through arrow's native IO subsystem and avoid extra pybytes interactions -- this lets you interact with memory in C/C++ with zero copy

def write_arrow(df):
    batch = pa.RecordBatch.from_pandas(df)

    buf = io.InMemoryOutputStream()
    writer = ipc.ArrowFileWriter(buf, batch.schema)
    writer.write_record_batch(batch)
    writer.close()
    return buf.get_result()

YMMV, worth double checking my work in case I made some mistakes

[mrocklin]

OK, there is a trivial implementation using Arrow on the Dask side here: #643

Once questions there get resolved we should have enough in place to start doing more integrative benchmarks.

In my mind there are two approaches, Blosc-and-custom-code or Arrow. Each has pros and cons.

The Blosc solution can be implemented today and is likely to be near optimal in speed, at least for numeric data. The main cons here are that we need to write code around the block manager etc. and we need to maintain this going forward. We're likely to miss things like new dtypes, categorical index, etc.. The good news here is that we've already done a lot of this a couple times before and can probably steal ideas from previous implementations in systems like partd.

The Arrow solution has two main things going for it:

1. Other people are in charge of making sure the entire Pandas abstraction is well covered.
2. It keeps open the cross-language door

However in the short term it is missing some things like per-column or dtype-specific compression, comprehensive support for the Pandas abstraction (see issues in #643, I'd love to be wrong about this) etc..

I would like to write up and tune some performance benchmarks for dask.dataframe sometime in the next month. Serialization and communication costs do contribute significantly in shuffling so this is likely to end up on my critical path soon-ish.

[wesm]

It seems reasonable to have as many options available as possible (assuming the development burden is not too much), and as long as we have reproducible benchmarks to evaluate performance that can help drive performance engineering work. The Arrow<->Pandas conversion has a number of immaturities so it's not an turnkey solution at the moment -- it is an accessible codebase, though, so I invite others to get involved and I can help review patches and where help reach consensus w/ the other Arrow devs where design changes are needed (e.g. adding compression options).