Pandas 2.0 and Arrow

[bryevdv]

Pandas 2.0 is coming out and is increasing support for Apache Arrow as a backend:

https://datapythonista.me/blog/pandas-20-and-the-arrow-revolution-part-i

I wanted to start this discussion just to get some preliminary idea of what this might mean for us: will everything work out of the box, do we need any new explicit codepaths, are there longer term opportunities (e.g. to leverage pyarrow and arrow.js?)

More immediately: if things don't work out of the box (or nearly) do we need a pandas upper bound for version 3.1?

cc @bokeh/dev @benrussell80

[benrussell80]

I definitely need to look into this some more (and get a better understanding of Bokeh internals, for that matter), but Bokeh 3.1 should have no problem handling pandas pyarrow arrays since they sublcass pandas ExtensionArrays (see below). However, anytime a pandas array is encoded it is just converted to a numpy array. This means that some of the main benefits of arrow (i.e., minimizing array copies) are being lost.

Working code example for

• pyarrow==11.0.0
• pandas==1.5.3
• bokeh==3.1.0.dev3+44.gfbba8622.dirty

import pandas as pd
from bokeh.plotting import figure, show
from bokeh.models import ColumnDataSource, LabelSet


def main():
    source = ColumnDataSource({
        'x': pd.array([0, 1, 2], dtype='int64[pyarrow]'),
        'y': pd.array([0, 1, 2], dtype='int64[pyarrow]'),
        'text': pd.array(['hello', None, 'world'], dtype='string[pyarrow]')
    })

    p = figure()
    p.circle(source=source)
    p.add_layout(LabelSet(source=source))
    show(p)


if __name__ == '__main__':
    main()

[benrussell80]

No problem, happy to help. Arrow seems super interesting.

So pyarrow structs seem to be encoding as expected (similar to a dataclass). However, I think I need to understand a few things better about how arrays should be encoded. Why are multi-dimensional numpy arrays flattened during encoding?

Example. Console shows a 1-D Int32Array

import numpy as np

from bokeh.models import ColumnDataSource, CustomJS
from bokeh.io import show
from bokeh.plotting import figure

source = ColumnDataSource({
    'arr': np.array([
        [1, 2, 3],
        [4, 5, 6],
        [7, 8, 9],
    ])
})

p = figure()
p.js_on_event('tap', CustomJS(code='console.log(source.data["arr"])', args={'source': source}))

show(p)

[mattpap]

It actually shows Int32NDArray correctly in Chromium. Though depending on your browser it may potentially differ. You can explore the structure of the object in JS console, and it should be like this:

Symbol(__ndarray__) will clearly indicate that this is a NDArray derived from a built-in typed array.

[benrussell80]

Ooh I am in Firefox

[mattpap]

The string is indeed wrong, but it's still an NDArray:

[benrussell80]

I see now, thanks for bearing with me @mattpap

[philippjfr]

For reference here's a Gist demonstrating how to serialize an arrow table across a Jupyter Comm: https://gist.github.com/manzt/5c5e3de6d2cea65d2bb68af03db88249

This should be pretty similar to how we would handle this.

[NoAnyLove]

Hopefully the team could bring issue #10464 into consideration.

If bokeh support arrow format or use arrow format as internal data representation, I guess we could get better performance when user provides data to ColumnDataSource (less data copy) and when bokeh server pushes data to client (push the binary data without extra serialization/deserialization cost in JSON)

[bryevdv]

(push the binary data without extra serialization/deserialization cost in JSON)

FYI @NoAnyLove arrays are already not serialized to JSON for the bokeh server case. Raw memory buffers are (individually, independently) sent over directly a tornado websocket and then those buffers are immediately attached to JS typed arrays. I don't actually think there are any more copies to get rid of, except possibly inside the tornado ws/stream codepaths (but I think those may have been resolved in the last few years too).

[bryevdv]

I should also elaborate: one of the core premises of CDS is that the data you put inside it remains in the format you originally provide. So while I definitely think CDS need to support Arrow arrays, I do not think we can reasonably require or convert to arrow arrays.

[NoAnyLove]

Hi @bryevdv, thanks for elaborating on this. It's nice to know arrays are already optimized.

I haven't dive deep into bokeh/bokehjs code, so I might be completely wrong here. I think the server calls .tobytes method in serialization, which is not zero-copy. If we use Apache arrow as the internal data representation, we could potentially improve the serialization with its zero-copy APIs. We could also potentially improve the data transmission by using its IPC protocol (over WebSocket) and record batches.

Besides, if ColumnDataSource supports arrow representation, we could potentially improve the performance on updating its internal data (with record batches or something else) and also improve the data streaming to client side.

[bryevdv]

I think the server calls .tobytes method in serialization

This was my comment about any further optimizations needing to happen on the Tornado side—Tornado WebsocketHandler does not accept memoryviews (or at least it didn't last time I checked, it's possible the situation has improved and I am not aware). If/when it does, we already pass memoryviews around internally and can just drop the call to tobytes.

Perhaps we could pass memoryviews to Arrow and let arrow send them (zero-copy). At least I assume something along those line would be feasible. But introducing and entirely new separate IPC channel seems like a very heavy lift. We definitely would not want to require two websockets.

We could also potentially improve the data transmission by using its IPC protocol (over WebSocket) and record batches.

I'm not sure there is much improvement here either, our protocol already sends memory buffers individually over the websocket, we do not assemble one big message or anything like that.

I'll just reiterate what I said earlier: we definitely can't require or convert to arrow arrays, we can only support them alongside numpy arrays, plain python lists, etc. I'd love to see a more concrete proposal about what that could look like, or even (ideally) a limited POC.