Support arrow Table/RecordBatch types

[dhirschfeld]

xref: #614

I'm sure this is possible with #606 so this issue is mostly just to document my attempt to get this working. Out-of-the box (1.22.0+9.gdb758d0f), attempting to pass an arrow Table or RecordBatch results in a TypeError:

TypeError: no default __reduce__ due to non-trivial __cinit__

from distributed import Client
import pandas as pd
import pyarrow as pa

client = Client()

df = pd.DataFrame({'A': list('abc'), 'B': [1,2,3]})
tbl = pa.Table.from_pandas(df, preserve_index=False)

def echo(arg):
    return arg

>>> client.submit(echo, df).result().equals(df)
True
>>> client.submit(echo, tbl).result()
distributed.protocol.pickle - INFO - Failed to serialize (pyarrow.Table
A: string
B: int64
metadata
--------
{b'pandas': b'{"index_columns": [], "column_indexes": [], "columns": [{"name":'
            b' "A", "field_name": "A", "pandas_type": "unicode", "numpy_type":'
            b' "object", "metadata": null}, {"name": "B", "field_name": "B", "'
            b'pandas_type": "int64", "numpy_type": "int64", "metadata": null}]'
            b', "pandas_version": "0.23.1"}'},). Exception: no default __reduce__ due to non-trivial __cinit__
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py in dumps(x)
     37     try:
---> 38         result = pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
     39         if len(result) < 1000:

C:\Miniconda3\lib\site-packages\pyarrow\lib.cp36-win_amd64.pyd in pyarrow.lib.RecordBatch.__reduce_cython__()

TypeError: no default __reduce__ due to non-trivial __cinit__

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-48-4e8e2ea90e79> in <module>()
----> 1 client.submit(echo, tbl).result()

C:\Miniconda3\lib\site-packages\distributed\client.py in submit(self, func, *args, **kwargs)
   1236                                          resources={skey: resources} if resources else None,
   1237                                          retries=retries,
-> 1238                                          fifo_timeout=fifo_timeout)
   1239 
   1240         logger.debug("Submit %s(...), %s", funcname(func), key)

C:\Miniconda3\lib\site-packages\distributed\client.py in _graph_to_futures(self, dsk, keys, restrictions, loose_restrictions, priority, user_priority, resources, retries, fifo_timeout)
   2093 
   2094             self._send_to_scheduler({'op': 'update-graph',
-> 2095                                      'tasks': valmap(dumps_task, dsk3),
   2096                                      'dependencies': dependencies,
   2097                                      'keys': list(flatkeys),

C:\Miniconda3\lib\site-packages\cytoolz\dicttoolz.pyx in cytoolz.dicttoolz.valmap()

C:\Miniconda3\lib\site-packages\cytoolz\dicttoolz.pyx in cytoolz.dicttoolz.valmap()

C:\Miniconda3\lib\site-packages\distributed\worker.py in dumps_task(task)
    799         elif not any(map(_maybe_complex, task[1:])):
    800             return {'function': dumps_function(task[0]),
--> 801                     'args': warn_dumps(task[1:])}
    802     return to_serialize(task)
    803 

C:\Miniconda3\lib\site-packages\distributed\worker.py in warn_dumps(obj, dumps, limit)
    808 def warn_dumps(obj, dumps=pickle.dumps, limit=1e6):
    809     """ Dump an object to bytes, warn if those bytes are large """
--> 810     b = dumps(obj)
    811     if not _warn_dumps_warned[0] and len(b) > limit:
    812         _warn_dumps_warned[0] = True

C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py in dumps(x)
     49     except Exception:
     50         try:
---> 51             return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
     52         except Exception as e:
     53             logger.info("Failed to serialize %s. Exception: %s", x, e)

C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py in dumps(obj, protocol)
    893     try:
    894         cp = CloudPickler(file, protocol=protocol)
--> 895         cp.dump(obj)
    896         return file.getvalue()
    897     finally:

C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py in dump(self, obj)
    266         self.inject_addons()
    267         try:
--> 268             return Pickler.dump(self, obj)
    269         except RuntimeError as e:
    270             if 'recursion' in e.args[0]:

C:\Miniconda3\lib\pickle.py in dump(self, obj)
    407         if self.proto >= 4:
    408             self.framer.start_framing()
--> 409         self.save(obj)
    410         self.write(STOP)
    411         self.framer.end_framing()

C:\Miniconda3\lib\pickle.py in save(self, obj, save_persistent_id)
    474         f = self.dispatch.get(t)
    475         if f is not None:
--> 476             f(self, obj) # Call unbound method with explicit self
    477             return
    478 

C:\Miniconda3\lib\pickle.py in save_tuple(self, obj)
    734         if n <= 3 and self.proto >= 2:
    735             for element in obj:
--> 736                 save(element)
    737             # Subtle.  Same as in the big comment below.
    738             if id(obj) in memo:

C:\Miniconda3\lib\pickle.py in save(self, obj, save_persistent_id)
    494             reduce = getattr(obj, "__reduce_ex__", None)
    495             if reduce is not None:
--> 496                 rv = reduce(self.proto)
    497             else:
    498                 reduce = getattr(obj, "__reduce__", None)

C:\Miniconda3\lib\site-packages\pyarrow\lib.cp36-win_amd64.pyd in pyarrow.lib.RecordBatch.__reduce_cython__()

TypeError: no default __reduce__ due to non-trivial __cinit__

[mrocklin]

Interesting.

I can reproduce this with just Pickle. I recommend raising this as an issue with Arrow.

In [1]: import pandas as pd
   ...: import pyarrow as pa
   ...: df = pd.DataFrame({'A': list('abc'), 'B': [1,2,3]})
   ...: tbl = pa.Table.from_pandas(df, preserve_index=False)
   ...: 
   ...: 

In [2]: import pickle

In [3]: b = pickle.dumps(tbl)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-264061ce888e> in <module>()
----> 1 b = pickle.dumps(tbl)

~/Software/anaconda/envs/arrow/lib/python3.6/site-packages/pyarrow/lib.cpython-36m-x86_64-linux-gnu.so in pyarrow.lib.Table.__reduce_cython__()

TypeError: no default __reduce__ due to non-trivial __cinit__

Alternatively if you wanted you could also help to implement a custom serialization solution for Arrow for Dask (see docs here) which would be useful to help avoid memory copies during transfer, but this is probably a minor performance concern in the common case. Getting other projects to implement the pickle protocol is probably the first step.

[dhirschfeld]

The docs seem pretty clear/straightforward, however I haven't (yet) gotten it to work:

from distributed.protocol.serialize import register_serialization
import pyarrow as arw


def serialize(batch):
    sink = arw.BufferOutputStream()
    writer = arw.RecordBatchStreamWriter(sink, batch.schema)
    writer.write_batch(batch)
    writer.close()
    buf = sink.get_result()
    return buf.to_pybytes()

def deserialize(blob):
    reader = arw.RecordBatchStreamReader(arw.BufferReader(blob))
    return reader.read_next_batch()

register_serialization(arw.RecordBatch, serialize, deserialize)

>>> batch = arw.RecordBatch.from_pandas(df)
>>> batch
<pyarrow.lib.RecordBatch at 0x9fd9ae8>
>>> isinstance(batch, arw.RecordBatch)
True
>>> batch.equals(deserialize(serialize(batch)))
True
>>> client.submit(echo, batch).result()
distributed.protocol.pickle - INFO - Failed to serialize (<pyarrow.lib.RecordBatch object at 0x0000000009FD9AE8>,). Exception: no default __reduce__ due to non-trivial __cinit__
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py in dumps(x)
     37     try:
---> 38         result = pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
     39         if len(result) < 1000:

C:\Miniconda3\lib\site-packages\pyarrow\lib.cp36-win_amd64.pyd in pyarrow.lib.RecordBatch.__reduce_cython__()

TypeError: no default __reduce__ due to non-trivial __cinit__

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
<ipython-input-111-39ca69d5a2b4> in <module>()
----> 1 client.submit(echo, batch).result()

C:\Miniconda3\lib\site-packages\distributed\client.py in submit(self, func, *args, **kwargs)
   1236                                          resources={skey: resources} if resources else None,
   1237                                          retries=retries,
-> 1238                                          fifo_timeout=fifo_timeout)
   1239 
   1240         logger.debug("Submit %s(...), %s", funcname(func), key)

C:\Miniconda3\lib\site-packages\distributed\client.py in _graph_to_futures(self, dsk, keys, restrictions, loose_restrictions, priority, user_priority, resources, retries, fifo_timeout)
   2093 
   2094             self._send_to_scheduler({'op': 'update-graph',
-> 2095                                      'tasks': valmap(dumps_task, dsk3),
   2096                                      'dependencies': dependencies,
   2097                                      'keys': list(flatkeys),

C:\Miniconda3\lib\site-packages\cytoolz\dicttoolz.pyx in cytoolz.dicttoolz.valmap()

C:\Miniconda3\lib\site-packages\cytoolz\dicttoolz.pyx in cytoolz.dicttoolz.valmap()

C:\Miniconda3\lib\site-packages\distributed\worker.py in dumps_task(task)
    799         elif not any(map(_maybe_complex, task[1:])):
    800             return {'function': dumps_function(task[0]),
--> 801                     'args': warn_dumps(task[1:])}
    802     return to_serialize(task)
    803 

C:\Miniconda3\lib\site-packages\distributed\worker.py in warn_dumps(obj, dumps, limit)
    808 def warn_dumps(obj, dumps=pickle.dumps, limit=1e6):
    809     """ Dump an object to bytes, warn if those bytes are large """
--> 810     b = dumps(obj)
    811     if not _warn_dumps_warned[0] and len(b) > limit:
    812         _warn_dumps_warned[0] = True

C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py in dumps(x)
     49     except Exception:
     50         try:
---> 51             return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
     52         except Exception as e:
     53             logger.info("Failed to serialize %s. Exception: %s", x, e)

C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py in dumps(obj, protocol)
    893     try:
    894         cp = CloudPickler(file, protocol=protocol)
--> 895         cp.dump(obj)
    896         return file.getvalue()
    897     finally:

C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py in dump(self, obj)
    266         self.inject_addons()
    267         try:
--> 268             return Pickler.dump(self, obj)
    269         except RuntimeError as e:
    270             if 'recursion' in e.args[0]:

C:\Miniconda3\lib\pickle.py in dump(self, obj)
    407         if self.proto >= 4:
    408             self.framer.start_framing()
--> 409         self.save(obj)
    410         self.write(STOP)
    411         self.framer.end_framing()

C:\Miniconda3\lib\pickle.py in save(self, obj, save_persistent_id)
    474         f = self.dispatch.get(t)
    475         if f is not None:
--> 476             f(self, obj) # Call unbound method with explicit self
    477             return
    478 

C:\Miniconda3\lib\pickle.py in save_tuple(self, obj)
    734         if n <= 3 and self.proto >= 2:
    735             for element in obj:
--> 736                 save(element)
    737             # Subtle.  Same as in the big comment below.
    738             if id(obj) in memo:

C:\Miniconda3\lib\pickle.py in save(self, obj, save_persistent_id)
    494             reduce = getattr(obj, "__reduce_ex__", None)
    495             if reduce is not None:
--> 496                 rv = reduce(self.proto)
    497             else:
    498                 reduce = getattr(obj, "__reduce__", None)

C:\Miniconda3\lib\site-packages\pyarrow\lib.cp36-win_amd64.pyd in pyarrow.lib.RecordBatch.__reduce_cython__()

TypeError: no default __reduce__ due to non-trivial __cinit__

...so it seems like my custom serialization functions are being ignored?

Am I missing something or is this a bug? Maybe I need to somehow register the functions with the remote workers?

[dhirschfeld]

I recommend raising this as an issue with Arrow

Looks like this is already being tracked:
https://issues.apache.org/jira/browse/ARROW-1715

[dhirschfeld]

With the scipy sprints coming up I just thought I'd clarify the status of this issue which is that I believe it is highlighting a bug in the dask custom serialization.

Irrespective of the fact that arrow doesn't implement the pickle protocol IIUC I should still be able to register custom serialization/deserialisation functions however:

register_serialization(arw.RecordBatch, serialize, deserialize)

...appears to have no effect.

I'd be glad to learn that I'm just doing something wrong in which case I'd be happy to put in a PR to clarify the documentation.

[mrocklin]

Do you have an implementation for serialize and deserialize then? Are you able to publish them? Perhaps either here as a PR?

[dhirschfeld]

Published above, but it might have got lost in the traceback noise:

from distributed.protocol.serialize import register_serialization
import pyarrow as arw


def serialize(batch):
    sink = arw.BufferOutputStream()
    writer = arw.RecordBatchStreamWriter(sink, batch.schema)
    writer.write_batch(batch)
    writer.close()
    buf = sink.get_result()
    return buf.to_pybytes()

def deserialize(blob):
    reader = arw.RecordBatchStreamReader(arw.BufferReader(blob))
    return reader.read_next_batch()

register_serialization(arw.RecordBatch, serialize, deserialize)

[mrocklin]

See http://distributed.readthedocs.io/en/latest/serialization.html#id2 . You also need to return a header, which in your case is likely an empty dictionary.

…

On Wed, Jul 11, 2018 at 10:56 PM, Dave Hirschfeld ***@***.***> wrote: Published above, but it might have got lost in the traceback noise: from distributed.protocol.serialize import register_serializationimport pyarrow as arw def serialize(batch): sink = arw.BufferOutputStream() writer = arw.RecordBatchStreamWriter(sink, batch.schema) writer.write_batch(batch) writer.close() buf = sink.get_result() return buf.to_pybytes() def deserialize(blob): reader = arw.RecordBatchStreamReader(arw.BufferReader(blob)) return reader.read_next_batch() register_serialization(arw.RecordBatch, serialize, deserialize) — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#2103 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszFFuFM7MoaXI1jP2oh8hkKvvuqlNks5uFskGgaJpZM4VGi9I> .

[dhirschfeld]

Right - user error! Sorry for the noise.

It is right there in both the docs and the example :|

[mrocklin]

A pull request would be welcome!

…

On Wed, Jul 11, 2018 at 11:08 PM, Dave Hirschfeld ***@***.***> wrote: Right - user error! Sorry for the noise. It is right there in both the docs and the example :| — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#2103 (comment)>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AASszJ6u50HLy2nyXuwFmLgNBPNKplM5ks5uFsvFgaJpZM4VGi9I> .

[dhirschfeld]

You mean add a file distributed/protocol/arrow.py and add it to the __init__.py?

Can do, if that's what you're after...

[dhirschfeld]

Hmmm, still having problems:

import pyarrow as arw
import pandas as pd

df = pd.DataFrame({'A': list('abc'), 'B': [1,2,3]})
batch = arw.RecordBatch.from_pandas(df, preserve_index=False)

def echo(arg):
    return arg

from distributed.protocol.serialize import register_serialization

def serialize(batch):
    sink = arw.BufferOutputStream()
    writer = arw.RecordBatchStreamWriter(sink, batch.schema)
    writer.write_batch(batch)
    writer.close()
    buf = sink.get_result()
    header = {}
    frames = [buf.to_pybytes()]
    return header, frames

def deserialize(header, frames):
    blob = frames[0]
    reader = arw.RecordBatchStreamReader(arw.BufferReader(blob))
    return reader.read_next_batch()

register_serialization(arw.RecordBatch, serialize, deserialize)

The functions still work (roundtrip):

>>> batch.equals(deserialize(*serialize(batch)))
True

...but the client.submit call still seems to ignore the custom functions?

>>> client.submit(echo, batch, pure=False)
distributed.protocol.pickle - INFO - Failed to serialize (<pyarrow.lib.RecordBatch object at 0x000000000CDE1EF8>,). Exception: no default __reduce__ due to non-trivial __cinit__
<snip>
TypeError: no default __reduce__ due to non-trivial __cinit__

[dhirschfeld]

Just as a sanity check I tried the Human example and it works fine:

>>> human = Human('Alice')
>>> client.submit(echo, human, pure=False).result().name
'Alice'

I thought maybe it's something to do with different environments between my desktop and the workers but we're all running arrow=0.9.0

>>> def version():
...     import pyarrow as awr
...     return arw.__version__

>>> version()
'0.9.0'

>>> client.run(version)
{'tcp://192.168.224.180:64230': '0.9.0',
 'tcp://192.168.224.196:54674': '0.9.0',
 'tcp://192.168.224.34:63531': '0.9.0',
 'tcp://192.168.226.190:62284': '0.9.0',
 'tcp://192.168.228.225:55086': '0.9.0',
 'tcp://192.168.230.94:60170': '0.9.0',
 'tcp://192.168.231.121:58951': '0.9.0',
 'tcp://192.168.232.202:55595': '0.9.0',
 'tcp://192.168.235.98:59804': '0.9.0',
 'tcp://192.168.236.232:60676': '0.9.0'}

[dhirschfeld]

Note:

The Human example also fails if it isn't picklable:

class Human(object):
    def __init__(self, name):
        self.name = name

    def __getstate__(self):
        1/0

>>> client.submit(echo, human, pure=False)
distributed.protocol.pickle - INFO - Failed to serialize (<__main__.Human object at 0x00000000054AD6A0>,). Exception: division by zero
Traceback (most recent call last):
  File "C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py", line 38, in dumps
    result = pickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
  File "C:/Users/dhirschf/.PyCharm2018.1/config/scratches/scratch_35.py", line 38, in __getstate__
    1/0
ZeroDivisionError: division by zero
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
  File "C:\Miniconda3\lib\site-packages\IPython\core\interactiveshell.py", line 2963, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-4-c602a524231a>", line 1, in <module>
    client.submit(echo, human, pure=False)
  File "C:\Miniconda3\lib\site-packages\distributed\client.py", line 1238, in submit
    fifo_timeout=fifo_timeout)
  File "C:\Miniconda3\lib\site-packages\distributed\client.py", line 2095, in _graph_to_futures
    'tasks': valmap(dumps_task, dsk3),
  File "cytoolz/dicttoolz.pyx", line 165, in cytoolz.dicttoolz.valmap
  File "cytoolz/dicttoolz.pyx", line 190, in cytoolz.dicttoolz.valmap
  File "C:\Miniconda3\lib\site-packages\distributed\worker.py", line 801, in dumps_task
    'args': warn_dumps(task[1:])}
  File "C:\Miniconda3\lib\site-packages\distributed\worker.py", line 810, in warn_dumps
    b = dumps(obj)
  File "C:\Miniconda3\lib\site-packages\distributed\protocol\pickle.py", line 51, in dumps
    return cloudpickle.dumps(x, protocol=pickle.HIGHEST_PROTOCOL)
  File "C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py", line 895, in dumps
    cp.dump(obj)
  File "C:\Miniconda3\lib\site-packages\cloudpickle\cloudpickle.py", line 268, in dump
    return Pickler.dump(self, obj)
  File "C:\Miniconda3\lib\pickle.py", line 409, in dump
    self.save(obj)
  File "C:\Miniconda3\lib\pickle.py", line 476, in save
    f(self, obj) # Call unbound method with explicit self
  File "C:\Miniconda3\lib\pickle.py", line 736, in save_tuple
    save(element)
  File "C:\Miniconda3\lib\pickle.py", line 496, in save
    rv = reduce(self.proto)
  File "C:/Users/dhirschf/.PyCharm2018.1/config/scratches/scratch_35.py", line 38, in __getstate__
    1/0
ZeroDivisionError: division by zero