[FEA] Produce and Consume ArrowDeviceArray struct from cudf::table / cudf::column

[zeroshade]

Is your feature request related to a problem? Please describe.
I would like to generate Arrow IPC payloads from a cudf::table without copying the data off of the GPU device. Currently the to_arrow and from_arrow functions explicitly perform copies to and from the GPU device. There is not currently any efficient way to generate Arrow IPC payloads from libcudf without copying all of the data off of the device.

Describe the solution you'd like
In addition to the existing to_arrow and from_arrow functions, we could have a to_arrow_device_arr function that populates an ArrowDeviceArray struct from a cudf::table or cudf::column. We'd also create a from_arrow_device_arr function that could construct a cudf::table / cudf::column from an ArrowDeviceArray that describes Arrow data which is already on the device. Once you have the ArrowDeviceArray struct, the Arrow C++ library itself can be used to generate the IPC payloads without needing to copy the data off the device. This would also increase the interoperability options that libcudf has, as anything which produces or consumes ArrowDeviceArray structs could hand data off to libcudf and vice versa.

Describe alternatives you've considered
An alternative would be to implement Arrow IPC creating inside of the libcudf library, but I saw that this was explicitly removed from libcudf due to the requirement of linking against libarrow_cuda.so. (#10994). Implementing conversions to and from ArrowDeviceArray wouldn't require linking against libarrow_cuda.so at all and would provide an easy way to allow any consumers to create Arrow IPC payloads, or whatever else they want to do with the resulting Arrow data. Such as leveraging CUDA IPC with the data.

Additional context
When designing the ArrowDeviceArray struct, I created https://github.com/zeroshade/arrow-non-cpu as a POC which used Python numba to generate and operate on some GPU data before handing it off to libcudf, and then getting it back without copying off the device. Using ArrowDeviceArray as the way it handed the data off.

More recently I've been working on creating a protocol for sending Arrow IPC data that is located on GPUs across high-performance transports like UCX. To this end, I created a POC using libcudf to pass the data. As a result I have a partial implementation of the to_arrow_device_arr which can be found here. There's likely better ways than what I'm doing in there, but at least for my POC it was working.

The contribution guidelines say I should file this issue first for discussion rather than just submitting a PR, so that's where I'm at. I plan on trying to create a full implementation that I can contribute but wanted to have this discussion and get feedback here first.

Thanks for hearing me out everyone!

[davidwendt]

I don't think this requires new APIs to libcudf.

The cudf::column and cudf::table are data-owning structs in libcudf.
For zero-copy you should be able wrap the arrow data in device memory with a cudf::column_view (and cudf::table_view) which are non-owning data structures.

All of the libcudf APIs accept cudf::column_view objects and so do not require an owning object so there should be no need to copy the arrow data in order to call a libcudf function.

Generally, libcudf APIs will return new cudf::column object since they are modifying or creating new column/table data. You can take ownership of this data (which should be in arrow format in device memory) using the cudf::column::release() method and then place the data in an appropriate arrow structure.

[zeroshade]

For zero-copy you should be able wrap the arrow data in device memory with a cudf::column_view (and cudf::table_view) which are non-owning data structures.

The issue is that it's not clear-cut how to perform that wrapping since libcudf's memory representation still differs from Arrow in some cases, in addition to differences in how the buffers are handled (such as with string columns using children for their offsets/data and Arrow string the offsets and data buffers as plain buffers, not children). There's a significant amount of code required to correctly wrap cudf::column_views around Arrow data in device memory (note the significant amount of code in https://github.com/rapidsai/cudf/blob/branch-24.04/cpp/src/interop/from_arrow.cu and https://github.com/rapidsai/cudf/blob/branch-24.04/cpp/src/interop/to_arrow.cu), which makes it worthwhile to have functions in libcudf itself to encapsulate this logic. Rather than have consumers replicate the logic in their own libraries.

You can take ownership of this data (which should be in arrow format in device memory) using the cudf::column::release() method and then place the data in an appropriate arrow structure.

Sure, but like I mentioned above it's not necessarily as simple as placing it in the appropriate arrow structure. It requires a significant amount of code to do it correctly and properly, and it makes sense for that to exist within libcudf. Particularly because it can then remain updated as libcudf adds support for more Arrow data types.

[davidwendt]

Ok. That makes sense. This appears to be just a wrapper around the cudf::column_view constructor. But resolving the data-type and the other components needed for the parameters would involve significant code. And the reverse as well.

[zeroshade]

Before I go further working on it, could you take a look at my partial implementation in https://github.com/zeroshade/cudf-flight-ucx/blob/main/to_arrow.cc and let me know if you think that's a good direction to go in as opposed to a different approach?

If it's a good approach then i'll work on creating a PR for this

[davidwendt]

Who owns the data after this call?

I expected that this signature

arrow::Status to_arrow_device_arr(std::shared_ptr<cudf::table> input,
                                  struct ArrowDeviceArray* out,
                                  rmm::cuda_stream_view stream)

would be more like

std::unique_ptr<struct struct ArrowDeviceArray> to_arrow_device_array(cudf::table_view const& input,
                                                                      rmm::cuda_stream_view stream)

And throw an exception with an error message instead of returning a status.

Also, I don't understand what the cudaEventRecord objects are for and why they should be created here.
It seems a fragile piece of logic subject to fleeting changes in your arrow struct implementation.
I wonder if Arrow could have an API to build an ArrowDeviceArray from simple native elements like pointers and integers similar to the parameters for cudf::column_view and cudf::table_view constructors.

[zeroshade]

Who owns the data after this call?

Technically the data has shared ownership. The ArrowDeviceArray maintains a reference to the passed in std::shared_ptr<cudf::table> to keep it alive until the release callback in the struct is called. Because the ArrowDeviceArray is intended to be a C ABI, it uses a release callback to control the lifetime of the underlying data.

The problem with this

std::unique_ptr<struct ArrowDeviceArray> to_arrow_device_array(cudf::table_view const& input,
                                                                      rmm::cuda_stream_view stream)

Is that because you're only passing in a cudf::table_view which doesn't own its data, there's no way to guarantee that the data stays alive until the release callback on the ArrowDeviceArray is called. We need to ensure that the underlying data stays valid and alive until the release callback is used.

If we don't like the idea of the shared ownership, then the interface could take a cudf::table and give the ownership to the ArrowDeviceArray entirely rather than sharing ownership.

Also, I don't understand what the cudaEventRecord objects are for and why they should be created here.
It seems a fragile piece of logic subject to fleeting changes in your arrow struct implementation.

We don't manually synchronize on the stream during the creation of the ArrowDeviceArray, instead we create an event and record it on the stream provided. The event is then part of the ArrowDeviceArray struct, so that a consumer can have their own stream wait on that event to synchronize before attempting to access the data. This lets the consumer of the struct choose when they synchronize, and on whatever stream they want to synchronize on. This allows the struct to be passed across C boundaries to different libraries and/or runtimes (such as python or Go or Rust etc.) and allow the consumer to synchronize the GPU however they want. @kkraus14 might be better able to explain the reasoning for the event in the struct than I.

That said, the ArrowDeviceArray struct is intended to be ABI stable and will not change. You can find the full definition, documentation and reasoning behind the structure of the ArrowDeviceArray here: https://arrow.apache.org/docs/format/CDeviceDataInterface.html

I wonder if Arrow could have an API to build an ArrowDeviceArray from simple native elements like pointers and integers similar to the parameters for cudf::column_view and cudf::table_view constructors.

That's exactly what the ArrowDeviceArray struct is. It's the collection of "simple native elements like pointers and integers" that describe the array and its children. Arrow provides APIs which take the struct and construct Arrow Arrays or RecordBatches from the struct (https://github.com/apache/arrow/blob/main/cpp/src/arrow/c/bridge.h#L232)

In this scenario: arrow::Array and arrow::RecordBatch are equivalent to cudf::column and cudf::table. ArrowDeviceArray is a struct used to encapsulate all of the pointers/length/null counts/etc. to zero-copy send the data across a C ABI boundary.

[kkraus14]

Hey @davidwendt 😃

The CUDA event in the struct that @zeroshade mentioned is the responsibility of the producer of the struct. It should have all relevant work related to allocating and populating the memory that is being handed / shared to the struct captured so that a downstream user of the struct can wait on the event to guarantee that whatever stream they're working on doesn't cause a race condition with the relevant allocations / kernels that produced the memory.

The reason behind using a CUDA event as opposed to a CUDA stream is that often frameworks don't have a mechanism to share or extend the lifetime of their streams outside of their framework.

As far as the lifetime management of the actual memory, the struct's release callback is designed to be flexible to allow accommodating both owning and non-owning situations. I.E. if someone had a cudf::column_view, we could basically just have an empty release callback to have it function as a view as opposed to having any form of ownership. In the case of cudf having unique ownership in something like a std::unique_ptr<cudf::column>, then it would likely make sense to transfer ownership to the struct.

[davidwendt]

Hey Keith. Thanks but it seems this kind of Arrow-specific logic for an Arrow-specific struct does not belong in libcudf. It seems a bit fragile in that changing how Arrow manages objects would require changes in a non-Arrow repo (like cudf). For example, if in the future Arrow decided the cudaEventCreate was not sufficient and now relied on cudaEventCreateWithFlags instead, a new cudf PR would be required to make this compliant again.

I was picturing more of a arrow::ArrowArray::Make() factory function that would handle these kinds of details.
Something like this perhaps (likely needs tweaking):

std::unique_ptr<ArrowArray> Make( int length, int null_count, int offset, void* buffer, std::vector<ArrowArray> children);

(And similar one for ArrowDeviceArray)
And this function would handle all the Arrow-specific things including the release mechanism and whatever CUDA objects it needs. It also allows the Arrow code complete control on how it is created and destroyed.

Then the libcudf function could simply call this with the appropriate counts and device pointers.

[davidwendt]

Is that because you're only passing in a cudf::table_view which doesn't own its data, there's no way to guarantee that the data stays alive until the release callback on the ArrowDeviceArray is called. We need to ensure that the underlying data stays valid and alive until the release callback is used.

Generally, libcudf is based on accepting views that are non-owning as per our developer guidelines.
https://github.com/rapidsai/cudf/blob/branch-24.04/cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md#views-and-ownership
The caller must ensure proper ownership scope and lifetime. This also provides a great deal of flexibility since there is no guarantee the original data is owned by a cudf::table in the first place.

[kkraus14]

Thanks but it seems this kind of Arrow-specific logic for an Arrow-specific struct does not belong in libcudf. It seems a bit fragile in that changing how Arrow manages objects would require changes in a non-Arrow repo (like cudf).

This is Arrow format specific, but not Arrow library specific. Libcudf already has to_arrow and from_arrow functions to go from device memory in libcudf containers to host memory in arrow containers, so it's already been exposed to fragility in Arrow for years, but the memory layout and ABI has been stable for years.

What is proposed here doesn't use Arrow containers and is designed to be a vendorable single header with a stable ABI so there really isn't additional exposure to Arrow that isn't already there.

I was picturing more of a arrow::ArrowArray::Make() factory function that would handle these kinds of details. Something like this perhaps (likely needs tweaking):

std::unique_ptr<ArrowArray> Make( int length, int null_count, int offset, void* buffer, std::vector<ArrowArray> children);

(And similar one for ArrowDeviceArray) And this function would handle all the Arrow-specific things including the release mechanism and whatever CUDA objects it needs. It also allows the Arrow code complete control on how it is created and destroyed.

Then the libcudf function could simply call this with the appropriate counts and device pointers.

In theory something like this could be added as a free function in the vendorable header, but you'd need to handle all the nesting structure that columns can have where you'd ultimately end up likely recreating a healthy chunk of what this struct describes in itself. No matter what there's some translation that needs to happen from how libcudf organizes its device pointers into some type of interface, and that's basically what this struct is.

[kkraus14]

Also, supporting this interface could be used to replace the existing to_arrow and from_arrow functions and remove the need to actually depend on the arrow library for supporting this functionality in the future. You could return host memory via this interface and there would be functions in the arrow library that could be called against the returned struct to get arrow containers similar to what the current to_arrow / from_arrow functions do.

[davidwendt]

In theory something like this could be added as a free function in the vendorable header, but you'd need to handle all the nesting structure that columns can have where you'd ultimately end up likely recreating a healthy chunk of what this struct describes in itself. No matter what there's some translation that needs to happen from how libcudf organizes its device pointers into some type of interface, and that's basically what this struct is.

No, I would not expect Arrow to unwind libcudf data structures. My suggestion leaves most of the proposed logic intact (type-dispatch, etc) but just replaces the pieces that create the ArrowArray and ArrowDeviceArray with factory functions implemented in the Arrow source.

I will work on a counter-proposal.

[kkraus14]

I will work on a counter-proposal.

Thank you! We'll more than happily review and iterate on it with you! 😃