[FEA] Function to "pack" a table into a single buffer

[jrhemstad]

Is your feature request related to a problem? Please describe.

I would like to be able to create a copy of a table_view where the storage for the constituent column is "packed" into a single allocation.

Describe the solution you'd like

An API something like this:

std::pair<rmm::device_buffer, table_view> pack(table_view const& t);

Where the returned device_buffer contains a copy of the contents of t in a single allocation and the returned table_view points into that buffer.

Additional context

This functionality is effectively already developed for contiguous_split in the alloc_and_copy function. The ask is just to wrap and expose this as a public API.

[felipeblazing]

We would love to have this. Shouldn't we have somethin that is alse like

std::unique_ptr<table> unpack(const rmm::device_buffer & buffer, unsigned long long size); 

[jrhemstad]

We would love to have this. Shouldn't we have somethin that is alse like

std::unique_ptr<table> unpack(const rmm::device_buffer & buffer, unsigned long long size); 

I think it'd be better to return a table_view that points into the returned device_buffer. I updated the issue.

Then if you want to create a non-packed table, from a packed one, you could do this:

table t;

[buffer, view] = pack(t);

table t(view); // copy ctor table from view will create a new `table` that copies from the `view`

[jakirkham]

This would also be useful from the UCX side as we could serialize a larger contiguous piece of memory and then unpack it later so that individual columns could be returned to RMM when no longer needed.

cc @quasiben

[jrhemstad]

@jakirkham there was some offline conversation with @jlowe @kkraus14 @nvdbaranec about how we'd serialize a "packed" table for communication.

The gist of it is, when you "pack" a table, you get a device_buffer and a table_view that points into that device_buffer. If you were to simply communicate the device_buffer and table_view as-is, the pointers in the table_view on the receiving side wouldn't be correct because they'd refer to memory locations on the sending side.

It's pretty easy to update the pointers on the receiving side if you know what the base pointers were on the sending side. It's just a bunch of pointer offset arithmetic.

We're trying to figure out a common set of serialize/deserialize primitives for libcudf that both Dask and Spark can use. Could you weigh in and help to define that the serialized layout should look like?

[quasiben]

@madsbk do you have opinions here as well on things ucx-py could be/should be doing ?

[madsbk]

@madsbk do you have opinions here as well on things ucx-py could be/should be doing ?

I don't think that ucx-py should support any serialization directly.
But it would be great if cuDF could serialize/deserialize multiple columns to/from a single device_buffer.

[kkraus14]

@jrhemstad I don't think we have any special requirements for the serialized layout other than using the typical standard containers of things like rmm::device_buffer, and std::vector so that we can easily send them over the wire without having to play games.

From our perspective we don't have a lot of opinions as our path would look something like:
Python Table --> cudf::table_view --> cudf::pack --> (rmm::device_buffer, std::vector) --> UCX Send(s) --> UCX Receive(s) --> (rmm::device_buffer, std::vector) --> cudf::unpack --> cudf::table_view --> Python Table

[jlowe]

I don't think we have any special requirements for the serialized layout other than using the typical standard containers of things

+1 from the Spark plugin perspective. We should be able to work with just the device buffer representing the contiguous GPU data and an opaque blob of host bytes representing the serialized table_view.

[jakirkham]

@madsbk do you have opinions here as well on things ucx-py could be/should be doing ?

I don't think that ucx-py should support any serialization directly.
But it would be great if cuDF could serialize/deserialize multiple columns to/from a single device_buffer.

Yeah I was thinking we might do this in Distributed before it got to UCX-Py. More-or-less the same as PR ( dask/distributed#3453 ).

[jakirkham]

FWIW I think from the Dask side w.r.t. cuDF we would be well served for most cases by having something we could plug into cuDF "cuda" serialization. We could then reuse this for transmission with TCP or UCX and spilling since they all go through that code path now.

There are some other objects we would like similar functionality (like CuPy sparse arrays). Though I think other Numba and CuPy objects are already working with contiguous memory chunks to the extent possible.