[Python] Pickling a sliced array serializes all the buffers

[asfimport]

If a large array is sliced, and pickled, it seems the full buffer is serialized, this leads to excessive memory usage and data transfer when using multiprocessing or dask.

>>> import pyarrow as pa
>>> ar = pa.array(['foo'] * 100_000)
>>> ar.nbytes
700004
>>> import pickle
>>> len(pickle.dumps(ar.slice(10, 1)))
700165

NumPy for instance
>>> import numpy as np
>>> ar_np = np.array(ar)
>>> ar_np
array(['foo', 'foo', 'foo', ..., 'foo', 'foo', 'foo'], dtype=object)
>>> import pickle
>>> len(pickle.dumps(ar_np[10:11]))
165

I think this makes sense if you know arrow, but kind of unexpected as a user.

Is there a workaround for this? For instance copy an arrow array to get rid of the offset, and trim the buffers?

Reporter: Maarten Breddels / @maartenbreddels
Assignee: Clark Zinzow

Related issues:

• [C++][Python] Create utils for deep-copying an Array/ ArrayData #30503 (is related to)

_{Note: This issue was originally created as ARROW-10739. Please see the migration documentation for further details.}

[asfimport]

Wes McKinney / @wesm:
We truncate the buffers on sliced arrays when writing record batches to the IPC protocol, so the buffers should be similarly truncated in the case of pickling. 

[asfimport]

Maarten Breddels / @maartenbreddels:
Ok, good to know.

Two workarounds I came up with

 

%%timeit
s = pa.serialize(ar.slice(10, 1))
ar2 = pa.deserialize(s.to_buffer())
790 µs ± 578 ns per loop (mean ± std. dev. of 7 runs, 1000 loops each)

 

 

 

import vaex.arrow.convert

----

def trim_buffers(ar):
    '''
    >>> ar = pa.array([1, 2, 3, 4], pa.int8())
    >>> ar.nbytes
    4
    >>> ar.slice(2, 2) #doctest: +ELLIPSIS
    <pyarrow.lib.Int8Array object at 0x...>
    [
      3,
      4
    ]
    >>> ar.slice(2, 2).nbytes
    4
    >>> trim_buffers(ar.slice(2, 2)).nbytes
    2
    >>> trim_buffers(ar.slice(2, 2))#doctest: +ELLIPSIS
    <pyarrow.lib.Int8Array object at 0x...>
    [
      3,
      4
    ]
    '''
    schema = pa.schema({'x': ar.type})
    with pa.BufferOutputStream() as sink:
        with pa.ipc.new_stream(sink, schema) as writer:
            writer.write_table(pa.table({'x': ar}))
    with pa.BufferReader(sink.getvalue()) as source:
        with pa.ipc.open_stream(source) as reader:
            table = reader.read_all()
            assert table.num_columns == 1
            assert table.num_rows == len(ar)
            trimmed_ar = table.column(0)
    if isinstance(trimmed_ar, pa.ChunkedArray):
        assert len(trimmed_ar.chunks) == 1
        trimmed_ar = trimmed_ar.chunks[0]


    return trimmed_ar
----

%%timeit
vaex.arrow.convert.trim_buffers(ar.slice(10, 1))
202 µs ± 2.31 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

 

 

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Note that pyarrow.serialize is deprecated, so best not use that as a workaround

[asfimport]

Maarten Breddels / @maartenbreddels:
Thanks Joris!

I cannot reproduce the previous timings (I guess I had an debug install without optimization), but this one seems fastest:

%%timeit
pa.concat_arrays([ar.slice(10, 1)])
2.16 µs ± 9.22 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

(vs 8 and 125 us using ipc and (de)serialize respectively)

[asfimport]

Wes McKinney / @wesm:
Seems like this would be good to fix?

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
That would indeed be good.

For a moment, I naively thought this would just be adding a SliceBuffer call when wrapping the buffer in _reduce_array_data (

arrow/python/pyarrow/array.pxi

Line 597 in c43fab3

buffers.append(pyarrow_wrap_buffer(buf)

). But of course, the offset and length to slice the buffer with depends on the array type and bit width, or whether it's a bitmap, etc.
In the IPC code, the truncation is handled in the Visit methods of RecordBatchSerializer (eg for primitive arrays:

arrow/cpp/src/arrow/ipc/writer.cc

Line 331 in c43fab3

data = SliceBuffer(data, byte_offset, buffer_length);

), and this is quite a lot of code for doing this correctly for all the different data types. Something we shouldn't start replicating in cython, I think.

Are there other utilities in C++ that can be reused to do this truncation? Or could we "just" use the IPC serialization under the hood for pickling in Python?

[asfimport]

Jim Crist-Harif:
We're running into this in Dask right now when attempting to integrate Pandas string[pyarrow], since pickling pyarrow string arrays ends up pickling all the data even if the result only includes a small slice. I'm willing to hack on this if no one else has the bandwidth, but upon initial inspection this looks a bit more complicated than I'd like to bite off for a new-ish arrow contribute. With some guidance on the best path forward though I could possibly get something working though? @jorisvandenbossche any further thoughts on a solution here?

[asfimport]

Krisztian Szucs / @kszucs:
Postponing to 10.0 since there is no PR available at the moment.

[asfimport]

Philipp Moritz / @pcmoritz:
In Ray we are also planning to work around this ray-project/ray#22891 – it would be wonderful to see this fixed in Arrow :)

[asfimport]

Clark Zinzow:
Hey folks, I'm the author of the Ray PR that @pcmoritz linked to, which essentially ports Arrow's buffer truncation in the IPC serialization path to Python as a custom pickle serializer. I'd be happy to help push on getting this fixed upstream for Arrow 10.0.0.

First, is there any in-progress work by [~jcrist] or others?

If not, I could take this on in the next month or so; the two implementation routes that I've thought of when looking at the IPC serialization code (these are basically the same routes that @jorisvandenbossche pointed out a year ago) are:

1. Refactor the IPC writer's per-type buffer truncation logic into utilities that can be shared by the IPC serialization path and the pickle serialization path.

2. Directly use the Arrow IPC format in its pickle serialization, where the pickle reducer is a light wrapper around the IPC serialization and deserialization hooks.

   Do either of these routes sound appealing? (2) has the added benefits of consolidating the serialization schemes on the IPC format and pushing all expensive serialization code into C++ land, but is a larger change and would involve otherwise-unnecessary wrapping of plain Arrow (chunked) arrays in record batches in order to match the IPC format, so maybe (1) is the better option.

[asfimport]

Clark Zinzow:
Ping on this, @amol- [~jcrist] are either of y'all actively working on this for Arrow 10.0.0? And if not, does option (1) that I gave in my previous comment sound like a good path forward?

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
[~clarkzinzow] sorry for the slow reply, several of us were on holidays. As far as I know, nobody is actively working on this, so a PR is certainly welcome! I think option (1) is a good path forward.

[asfimport]

Clark Zinzow:
@jorisvandenbossche No worries! And sounds good, I should be able to start working on this in a few weeks, I will update this thread once I've started working on it and again once I have a PR out.

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
That sounds good!

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
[~clarkzinzow] were you able to make some progress on this?

(I am removing the "In progress" label, just to make it a bit easier to keep track of this as open issue in JIRA, but can change that again once there is an open PR)