Repository: dask/distributed

Issue:
Serialization problem with arrow Tables?
I think the below code should work (but it doesn't) is this a bug or am I missing something?

Setup code
```python 
import pandas as pd
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

# create parquet files
filenames = []
for n in range(10):
    df = n*pd.DataFrame(np.ones((10, 5)))
    tbl = pa.Table.from_pandas(df)
    filename = f'part{n}.pq'
    pq.write_table(tbl, filename, version='2.0')
    filenames.append(filename)
```
Demo of problem:
```python
def read_table(filename):
    import pyarrow.parquet as pq
    tbl = pq.read_table(filename)
    return tbl


from distributed import LocalCluster, Client

cluster = LocalCluster(2, threads_per_worker=1)
client = Client(cluster.scheduler.address)
futs = client.map(read_table, filenames)
tbl = client.submit(pa.concat_tables, futs).result()
```
Exception:
```python-traceback
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-67-74f644c68e47> in <module>
      1 futs = client.map(read_table, filenames)
----> 2 tbl = client.submit(pa.concat_tables, futs).result()

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\client.py in result(self, timeout)
    191         # shorten error traceback
    192         result = self.client.sync(self._result, callback_timeout=timeout,
--> 193                                   raiseit=False)
    194         if self.status == 'error':
    195             six.reraise(*result)

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\client.py in sync(self, func, *args, **kwargs)
    673             return future
    674         else:
--> 675             return sync(self.loop, func, *args, **kwargs)
    676 
    677     def __repr__(self):

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\utils.py in sync(loop, func, *args, **kwargs)
    275             e.wait(10)
    276     if error[0]:
--> 277         six.reraise(*error[0])
    278     else:
    279         return result[0]

C:\Miniconda3\envs\py-dev\lib\site-packages\six.py in reraise(tp, value, tb)
    691             if value.__traceback__ is not tb:
    692                 raise value.with_traceback(tb)
--> 693             raise value
    694         finally:
    695             value = None

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\utils.py in f()
    260             if timeout is not None:
    261                 future = gen.with_timeout(timedelta(seconds=timeout), future)
--> 262             result[0] = yield future
    263         except Exception as exc:
    264             error[0] = sys.exc_info()

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\client.py in _result(self, raiseit)
    216                 raise gen.Return(exception)
    217         else:
--> 218             result = yield self.client._gather([self])
    219             raise gen.Return(result[0])
    220 

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\client.py in _gather(self, futures, errors, direct, local_worker)
   1520                 else:
   1521                     self._gather_future = future
-> 1522                 response = yield future
   1523 
   1524             if response['status'] == 'error':

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\client.py in _gather_remote(self, direct, local_worker)
   1569 
   1570             else:  # ask scheduler to gather data for us
-> 1571                 response = yield self.scheduler.gather(keys=keys)
   1572         finally:
   1573             self._gather_semaphore.release()

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\core.py in send_recv_from_rpc(**kwargs)
    578             try:
    579                 comm = yield self.live_comm()
--> 580                 result = yield send_recv(comm=comm, op=key, **kwargs)
    581             except (RPCClosed, CommClosedError) as e:
    582                 raise e.__class__("%s: while trying to call remote method %r"

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run(self)
   1145                             exc_info = None
   1146                     else:
-> 1147                         yielded = self.gen.send(value)
   1148 
   1149                     if stack_context._state.contexts is not orig_stack_contexts:

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\core.py in send_recv(comm, reply, serializers, deserializers, **kwargs)
    470     if isinstance(response, dict) and response.get('status') == 'uncaught-error':
    471         if comm.deserialize:
--> 472             six.reraise(*clean_exception(**response))
    473         else:
    474             raise Exception(response['text'])

C:\Miniconda3\envs\py-dev\lib\site-packages\six.py in reraise(tp, value, tb)
    690                 value = tp()
    691             if value.__traceback__ is not tb:
--> 692                 raise value.with_traceback(tb)
    693             raise value
    694         finally:

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\core.py in handle_comm()
    344                         if type(result) is gen.Future:
    345                             self._ongoing_coroutines.add(result)
--> 346                             result = yield result
    347                     except (CommClosedError, CancelledError) as e:
    348                         if self.status == 'running':

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\scheduler.py in gather()
   2315 
   2316         data, missing_keys, missing_workers = yield gather_from_workers(
-> 2317             who_has, rpc=self.rpc, close=False, serializers=serializers)
   2318         if not missing_keys:
   2319             result = {'status': 'OK', 'data': data}

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\utils_comm.py in gather_from_workers()
     65             for worker, c in coroutines.items():
     66                 try:
---> 67                     r = yield c
     68                 except EnvironmentError:
     69                     missing_workers.add(worker)

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1139                     if exc_info is not None:
   1140                         try:
-> 1141                             yielded = self.gen.throw(*exc_info)
   1142                         finally:
   1143                             # Break up a reference to itself

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\worker.py in get_data_from_worker()
   2669                                    deserializers=deserializers,
   2670                                    op='get_data', keys=keys, who=who,
-> 2671                                    max_connections=max_connections)
   2672         try:
   2673             status = response['status']

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1131 
   1132                     try:
-> 1133                         value = future.result()
   1134                     except Exception:
   1135                         self.had_exception = True

C:\Miniconda3\envs\py-dev\lib\site-packages\tornado\gen.py in run()
   1145                             exc_info = None
   1146                     else:
-> 1147                         yielded = self.gen.send(value)
   1148 
   1149                     if stack_context._state.contexts is not orig_stack_contexts:

C:\Miniconda3\envs\py-dev\lib\site-packages\distributed\core.py in send_recv()
    472             six.reraise(*clean_exception(**response))
    473         else:
--> 474             raise Exception(response['text'])
    475     raise gen.Return(response)
    476 

Exception: Object <pyarrow.lib.Buffer object at 0x000002A6C3FF03E8> is neither a bytes object nor has an encode method
```

I'll try updating and reproducing but won't get to that for a while so any confirmation appreciated...
```
print(pa.__version__)
print(dask.__version__)
print(distributed.__version__)

0.12.0
1.1.1
1.25.3
```

Discussion:
dhirschfeld: The returned `header = [buf]` is actually picklable so it's somewhat icky that it would blow up elsewhere:
```python
In [21]: sink
Out[21]: <pyarrow.lib.BufferOutputStream at 0x2c85aee33b0>

In [22]: len(pickle.dumps([sink.getvalue()]))
Out[22]: 1961

In [24]: len(distributed.protocol.dumps([sink.getvalue()]))
distributed.protocol.core - CRITICAL - Failed to Serialize
<snip>
TypeError: can not serialize 'pyarrow.lib.Buffer' object
```

That said, the workaround is straightforward (don't return complex objects in the header) so I I'm :+1: for this fix.

dhirschfeld: > *For pyarrow, this would be something like:*
```python
if hasattr(s, 'to_pybytes'):
    return s.to_pybytes()
```
Adding an arrow-specific workaround is a slippery slope so I think fixing it in the `serialize_table` 
function as is done here is better.
TomAugspurger: > Adding an arrow-specific workaround is a slippery slope

Yeah, we certainly don't want that. But my hope is that `bytes(thing)` in `ensure_bytes` is generic enough. It should work for anything implementing the buffer protocol.
TomAugspurger: Updated to fix the actual problem in `ensure_bytes` I think. I'll followup with a PR to Dask's version if this seems reasonable.

Patch:
diff --git a/distributed/__init__.py b/distributed/__init__.py
index ca36613c..d79993df 100644
--- a/distributed/__init__.py
+++ b/distributed/__init__.py
@@ -3,7 +3,7 @@ from dask.config import config
 from .actor import Actor, ActorFuture
 from .core import connect, rpc
 from .deploy import LocalCluster, Adaptive, SpecCluster
-from .diagnostics import progress
+from .diagnostics.progressbar import progress
 from .client import (
     Client,
     Executor,
diff --git a/distributed/utils.py b/distributed/utils.py
index c8ea8d64..65f5c188 100644
--- a/distributed/utils.py
+++ b/distributed/utils.py
@@ -904,22 +904,43 @@ def tmpfile(extension=""):
 
 
 def ensure_bytes(s):
-    """ Turn string or bytes to bytes
+    """Attempt to turn `s` into bytes.
+
+    Parameters
+    ----------
+    s : Any
+        The object to be converted. Will correctly handled
+
+        * str
+        * bytes
+        * objects implementing the buffer protocol (memoryview, ndarray, etc.)
+
+    Returns
+    -------
+    b : bytes
+
+    Raises
+    ------
+    TypeError
+        When `s` cannot be converted
+
+    Examples
+    --------
 
     >>> ensure_bytes('123')
     b'123'
     >>> ensure_bytes(b'123')
     b'123'
     """
-    if isinstance(s, bytes):
-        return s
-    if isinstance(s, memoryview):
-        return s.tobytes()
-    if isinstance(s, bytearray):  # noqa: F821
-        return bytes(s)
     if hasattr(s, "encode"):
         return s.encode()
-    raise TypeError("Object %s is neither a bytes object nor has an encode method" % s)
+    else:
+        try:
+            return bytes(s)
+        except Exception as e:
+            raise TypeError(
+                "Object %s is neither a bytes object nor has an encode method" % s
+            ) from e
 
 
 def divide_n_among_bins(n, bins):
diff --git a/docs/source/limitations.rst b/docs/source/limitations.rst
index 3c64d445..e272359e 100644
--- a/docs/source/limitations.rst
+++ b/docs/source/limitations.rst
@@ -35,7 +35,6 @@ Dask assumes the following about your functions and your data:
 -  Dask may run your functions multiple times,
    such as if a worker holding an intermediate result dies.  Any side effects
    should be `idempotent <https://en.wikipedia.org/wiki/Idempotence>`_.
--
 
 Security
 --------


Test patch:
diff --git a/distributed/protocol/tests/test_arrow.py b/distributed/protocol/tests/test_arrow.py
index a363ee95..37aff3a2 100644
--- a/distributed/protocol/tests/test_arrow.py
+++ b/distributed/protocol/tests/test_arrow.py
@@ -3,9 +3,9 @@ import pytest
 
 pa = pytest.importorskip("pyarrow")
 
+import distributed
 from distributed.utils_test import gen_cluster
-from distributed.protocol import deserialize, serialize
-
+from distributed.protocol import deserialize, serialize, to_serialize
 
 df = pd.DataFrame({"A": list("abc"), "B": [1, 2, 3]})
 tbl = pa.Table.from_pandas(df, preserve_index=False)
@@ -35,3 +35,12 @@ def test_scatter(obj):
         assert obj.equals(result)
 
     run_test()
+
+
+def test_dumps_compression():
+    # https://github.com/dask/distributed/issues/2966
+    # large enough to trigger compression
+    t = pa.Table.from_pandas(pd.DataFrame({"A": [1] * 10000}))
+    msg = {"op": "update", "data": to_serialize(t)}
+    result = distributed.protocol.loads(distributed.protocol.dumps(msg))
+    assert result["data"].equals(t)
diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py
index e5e18eb3..81541686 100644
--- a/distributed/tests/test_utils.py
+++ b/distributed/tests/test_utils.py
@@ -1,3 +1,4 @@
+import array
 import datetime
 from functools import partial
 import io
@@ -275,13 +276,25 @@ def test_funcname():
 
 
 def test_ensure_bytes():
-    data = [b"1", "1", memoryview(b"1"), bytearray(b"1")]
+    data = [b"1", "1", memoryview(b"1"), bytearray(b"1"), array.array("b", [49])]
     for d in data:
         result = ensure_bytes(d)
         assert isinstance(result, bytes)
         assert result == b"1"
 
 
+def test_ensure_bytes_ndarray():
+    result = ensure_bytes(np.arange(12))
+    assert isinstance(result, bytes)
+
+
+def test_ensure_bytes_pyarrow_buffer():
+    pa = pytest.importorskip("pyarrow")
+    buf = pa.py_buffer(b"123")
+    result = ensure_bytes(buf)
+    assert isinstance(result, bytes)
+
+
 def test_nbytes():
     def check(obj, expected):
         assert nbytes(obj) == expected