cudf icon indicating copy to clipboard operation
cudf copied to clipboard

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

Open zeroshade opened this issue 1 year ago • 35 comments

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. (https://github.com/rapidsai/cudf/issues/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!

zeroshade avatar Jan 29 '24 21:01 zeroshade

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.

davidwendt avatar Jan 29 '24 22:01 davidwendt

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.

zeroshade avatar Jan 30 '24 00:01 zeroshade

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.

davidwendt avatar Jan 30 '24 15:01 davidwendt

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

zeroshade avatar Jan 30 '24 15:01 zeroshade

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.

davidwendt avatar Jan 30 '24 15:01 davidwendt

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.

zeroshade avatar Jan 30 '24 16:01 zeroshade

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.

kkraus14 avatar Jan 30 '24 18:01 kkraus14

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 avatar Jan 30 '24 18:01 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.

davidwendt avatar Jan 30 '24 18:01 davidwendt

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 avatar Jan 30 '24 19:01 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.

kkraus14 avatar Jan 30 '24 19:01 kkraus14

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.

davidwendt avatar Jan 31 '24 13:01 davidwendt

I will work on a counter-proposal.

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

kkraus14 avatar Jan 31 '24 13:01 kkraus14

Ok, this is what I'm proposing for the 2 Make functions to go in the Arrow source. We can name that whatever makes sense. And we can change the return type to be a std::shared_ptr if that helps too.

namespace arrow {

/// raw pointer and a function to free it
using OwningBuffer = std::pair<const void*, std::function<void()>>;

namespace {
// generic object deleter functor for ArrowArray instances
struct DeleterFn {
  std::vector<ArrowArray*> children;
  std::vector<OwningBuffer> owners;
  ~DeleterFn()
  {
    for (auto& c : children)
      ArrowArrayRelease(c);
    for (auto& o : owners)
      std::invoke(o.second);
  }
};
}  // namespace

std::unique_ptr<ArrowArray> MakeArrowArray(int64_t length,
                                           int64_t null_count,
                                           int64_t offset,
                                           std::vector<OwningBuffer>&& data    = {},
                                           std::vector<ArrowArray*>&& children = {},
                                           ArrowArray&& dictionary             = {0})
{
  auto result = new ArrowArray{};
  std::memset(result, 0, sizeof(ArrowArray));

  const void** buffers = (const void**)(malloc(sizeof(void*) * data.size()));
  std::transform(data.begin(), data.end(), buffers, [](auto& buffer) { return buffer.first; });

  result->length     = length;
  result->null_count = null_count;
  result->offset     = offset;
  result->n_buffers  = 2;
  result->n_children = static_cast<int64_t>(children.size());
  result->buffers    = buffers;
  result->children   = children.data();
  result->dictionary = dictionary.length == 0 ? nullptr : new ArrowArray(std::move(dictionary));
  result->release    = [](struct ArrowArray* arr) {
    free(arr->buffers);
    auto d = static_cast<DeleterFn*>(arr->private_data);
    delete d;
    if (arr->dictionary) ArrowArrayRelease(arr->dictionary);
    ArrowArrayMarkReleased(arr);
  };
  result->private_data = new DeleterFn{std::move(children), std::move(data)};
  return std::unique_ptr<ArrowArray>(result);
}

std::unique_ptr<ArrowDeviceArray> MakeDeviceArray(ArrowArray&& array)
{
  auto result = new ArrowDeviceArray{std::move(array)};
  cudaEventCreate(reinterpret_cast<cudaEvent_t*>(&(result->sync_event)));
  int dev_id = 0;
  cudaGetDevice(&dev_id);
  result->device_id   = dev_id;
  result->device_type = ARROW_DEVICE_CUDA;
  return std::unique_ptr<ArrowDeviceArray>(result);
}
}  // namespace arrow

The DeleterFn can certainly go in a .cpp file along with the 2 function definitions. No need for these to be declared and be defined in a header file. I believe this should work and all the appropriate objects are managed correctly but I've not tested it.

davidwendt avatar Feb 01 '24 23:02 davidwendt

I spent some time recoded each of the dispatch functions to use these Make factories. Here are few of them for reference.

  // handles most of the fixed-width types
  std::unique_ptr<ArrowArray> operator()(cudf::column_view input, rmm::cuda_stream_view)
  {
    std::vector<arrow::OwningBuffer> data{{input.null_mask(), empty_fn}, {input.head(), empty_fn}};
    return arrow::MakeArrowArray(input.size(), input.null_count(), input.offset(), std::move(data));
  }

// the bool specialization shows passing in a custom 'delete' function for freeing the device_buffer
template <>
std::unique_ptr<ArrowArray> dispatch_to_arrow::operator()<bool>(cudf::column_view input,
                                                                rmm::cuda_stream_view stream)
{
  cudf::column_view view_without_offset =
    input.offset() == 0 ? input
                        : view_without_offset = cudf::column_view{input.type(), input.size() + input.offset(),
                                                                  input.head(), input.null_mask(),  input.null_count()};
  auto bitmask = std::get<0>(cudf::detail::bools_to_mask(
    view_without_offset, stream, rmm::mr::get_current_device_resource()));

  std::vector<arrow::OwningBuffer> data{{input.null_mask(), empty_fn}};
  data.emplace_back(device_buffer_to_arrow(std::move(*bitmask.release())));
  return arrow::MakeArrowArray(input.size(), input.null_count(), input.offset(), std::move(data));
}
...
which uses this utility (to be included in the libcudf source along with these dispatch functions):
...
// utility to transfer a device_buffer to an OwningBuffer
arrow::OwningBuffer device_buffer_to_arrow(rmm::device_buffer&& buffer)
{
  auto dbuf    = new rmm::device_buffer(std::move(buffer));
  auto deleter = [dbuf]() { delete dbuf; };
  return arrow::OwningBuffer{dbuf->data(), deleter};
}

// the main public function that returns the new ArrowDeviceArray
std::unique_ptr<arrow::ArrowDeviceArray> to_arrow_device_array(table_view input_view,
                                                               rmm::cuda_stream_view stream)
{
  std::vector<ArrowArray*> children;
  for (auto& c : input_view) {
    auto col = c.type().id() != cudf::type_id::EMPTY
                 ? cudf::type_dispatcher(c.type(), detail::dispatch_to_arrow{}, c, stream)
                 : detail::create_null_array(c.size());
    children.emplace_back(col.release());
  }
  std::vector<arrow::OwningBuffer> data{{nullptr, detail::empty_fn}};
  auto array =  arrow::MakeArrowArray(input_view.num_rows(), 0, 0, std::move(data), std::move(children));
  return arrow::MakeDeviceArray(std::move(*array.release()));
}
...
The create_null_array() was copied from original the get_null_arr()
...
std::unique_ptr<ArrowArray> create_null_array(int size)
{
  auto arr = std::make_shared<arrow::NullArray>(size);
  auto out = new ArrowArray{};
  ARROW_UNUSED(arrow::ExportArray(*arr, out));
  return std::unique_ptr<ArrowArray>(out);
}

Let me know if you want to see any of the other ones. I didn't realize how different each type is built into an ArrowArray but the Make function seems to handle them all. I was a bit surprised the type-id is not included in the structure.

davidwendt avatar Feb 01 '24 23:02 davidwendt

I think we'd also want to look into nanoarrow (#13678) before we design any new structs ourselves. If I'm reading this discussion right it seems like there should be significant overlap given that nanoarrow has a device-side extension.

vyasr avatar Feb 02 '24 21:02 vyasr

The DeleterFn can certainly go in a .cpp file along with the 2 function definitions. No need for these to be declared and be defined in a header file. I believe this should work and all the appropriate objects are managed correctly but I've not tested it.

This is unfortunately a C API as opposed to a CPP API. I imagine we could make this work regardless, but the bigger question is where would we expect this to live? If this lived in the main Arrow library then it eliminates the goal of being dependency free and requires linking libarrow which has a somewhat non-trivial dependency tree on its own. One of the goals of the interfaces is explicitly to avoid an explicit dependency on Arrow: https://arrow.apache.org/docs/format/CDeviceDataInterface.html#goals.

We could potentially implement something like this in nanoarrow as @vyasr mentioned above, but we'd probably need to take in the cuda event somewhere as opposed to having the make function create and record the event since the buffers coming in could potentially be on different streams or something of the like and I don't think there's a nice general way for something like nanoarrow to introspect and handle things properly. Additionally, the device and subsequent CUDA device extension in nanoarrow is quite new where there isn't interfaces for doing things like stream ordered memory management, stream ordered copying, etc. yet where I'm not sure how helpful it would be in the actual implementation here outside of providing the relevant definitions in headers for the Arrow C Device Data Interface.

kkraus14 avatar Feb 04 '24 03:02 kkraus14

Ok. The link was helpful background.

If this lived in the main Arrow library then it eliminates the goal of being dependency free and requires linking libarrow which has a somewhat non-trivial dependency tree on its own.

Since libcudf is already linking to libarrow.so, I'd like to consider Arrow providing these functions as an alternative to hand building the struct elements as illustrated in the original proposal.

davidwendt avatar Feb 06 '24 15:02 davidwendt

Since libcudf is already linking to libarrow.so

My understanding is that there is a desire for libcudf to no longer link against libarrow.so, where this proposal would enable a path to removing one of the key places it's used, in to_arrow and from_arrow as well as enabling handing GPU memory to other libraries that don't link to libcudf.

I believe from some conversations with @beckernick that he's expressed that Arrow increasing major versions ~quarterly and libcudf being tied to a specific major version has caused some compatibility pain in working with other packages across the ecosystem.

I'd like to consider Arrow providing these functions as an alternative to hand building the struct elements as illustrated in the original proposal.

I don't think this is particularly feasible. There's different ownership models / semantics that Arrow would need to capture / support here. I.E. shared ownership where someone would want to more or less stuff some shared_ptrs into the private_data struct member and handle them appropriately in the release callback.

Additionally, in your proposal above you'd still need to organize your buffers and child columns into a flattened structure, pass the device type, and create + record the CUDA event for synchronization yourself. It seems like the main difference would be moving handling the ownership semantics into Arrow as opposed to handling it in libcudf?

kkraus14 avatar Feb 06 '24 15:02 kkraus14

I'd like to consider Arrow providing these functions as an alternative to hand building the struct elements as illustrated in the original proposal.

I'm not sure what the functions would look like if they lived in libarrow, since libarrow can't use definitions of libcudf classes like table_view. Could you sketch the signatures you were thinking of?

bkietz avatar Feb 06 '24 16:02 bkietz

I'd like to consider Arrow providing these functions as an alternative to hand building the struct elements as illustrated in the original proposal.

I'm not sure what the functions would look like if they lived in libarrow, since libarrow can't use definitions of libcudf classes like table_view. Could you sketch the signatures you were thinking of?

https://github.com/rapidsai/cudf/issues/14926#issuecomment-1922458305

davidwendt avatar Feb 06 '24 16:02 davidwendt

What version of Arrow includes ArrowDeviceArray? I don't see it in the version used by libcudf so upgrading may be a prerequisite for this work.

I'm still puzzled by the lack of a type-id in these structures. What is your proposal for from_arrow_device_array? I believe it should be possible to build a cudf::table_view/cudf::column_views but only if the type-ids are available. Building a cudf::table does not look possible since cudf::column objects expect to own their data and for it to be stored in an rmm::device_buffer which is managed by RMM. There is no mechanism for RMM to manage device memory that it has not allocated.

davidwendt avatar Feb 06 '24 21:02 davidwendt

My understanding is that there is a desire for libcudf to no longer link against libarrow.so

@kkraus14 you're right that we would eventually like to stop linking against libarrow if possible.

We could potentially implement something like this in nanoarrow as @vyasr mentioned above

My understanding is that nanoarrow was intended to provide essentially what we would need to decouple the existing Arrow interop functionality in libcudf from linkage to libarrow itself: a small, easily vendored library that provides an implementation of readers/writers of the Arrow C data interface so that various libraries could produce ABI-equivalent versions of Arrow data structures without linking. Do I have that right?

but [...] I don't think there's a nice general way for something like nanoarrow to introspect and handle things properly.

Assuming my understanding of the goals of nanoarrow above is correct, is the main concern here leaking too much CUDA-specific information into the nanoarrow implementation, which would be a long-term issue? Or are you mostly concerned with the more short-term issue that

the device and subsequent CUDA device extension in nanoarrow is quite new where there isn't interfaces for doing things like stream ordered memory management, stream ordered copying, etc. yet

If it's the latter, then could it make sense to implement this kind of functionality in libcudf (or a separate but associated library) for now but eventually upstream it to nanoarrow?

If it's the former, then I'd like to better understand the inherent limitations you see in nanoarrow and see if we can find a path to upstream this. At a high level I think I understand your concerns but I would like to dig into the details a bit since IMHO something like this really ought to be within the long-term scope of nanoarrow if I've understood its intent properly. I think I agree that we'll always need some functionality in cudf, but in an ideal world I would hope that we'd have something close to as simple as (very rough, not trying to be precise with types etc since I imagine all that could change in nanoarrow):

ArrowDeviceArrayView to_arrow(column_view col) {
    // Note that I'm constructing an ArrowDeviceArrayView, not an ArrowDeviceArray,
    // because I assumed those were the intended semantics of that object.
    // Since it's a view and not a copy stream-ordering concerns seem like they'd be obviated.
    return ArrowDeviceArrayViewInit(col.data(), col.size());
}

but I really haven't looked into nanoarrow enough yet to understand where/why this would be problematic.

vyasr avatar Feb 06 '24 23:02 vyasr

The ArrowDeviceArray was introduced with arrow v13 or v14. You'll find it in the header file <arrow/c/abi.h>

I'm still puzzled by the lack of a type-id in these structures.

The type IDs are managed by a corresponding ArrowSchema object, also in the same header. They are separate to allow for a stream of batches of data with the same schema to not have to duplicate the type information for every batch of records. So libcudf would need two methods: one to fill in an ArrowSchema and it's children based on the type of a column/column_view or a table/table_view (a table is treated as a struct column whose fields are the columns of the table so everything is seamless).

The other issue I see with pushing this upstream is that the ArrowDeviceArray on the arrow side is supposed to be device agnostic. Any init function we provide would also need to have the device type passed in (we can't assume CUDA) which would also require the caller to pass in any synchronization event (if required).

Essentially the only thing a helper function like what you are asking could do is be a wrapper around populating a C struct, which seems a little redundant and unnecessary. At least to me.

Nanoarrow could certainly be used to simplify the creating of the ArrowArray and ArrowStruct objects though. @paleolimbot could comment further on nanoarrow for this

zeroshade avatar Feb 06 '24 23:02 zeroshade

Very interesting read! No pressure to use nanoarrow's implementation for any of this...if it can't help, its source might be useful to review and/or it might give you another endpoint to test against.

The nanoarrow C library (without any CUDA integration) can definitely populate the ArrowSchema for you. It might look like:

#include nanoarrow.hpp

int export_column_schema(const cudf::column& col, ArrowSchema* out) {
  nanoarrow::UniqueSchema tmp;
  ArrowSchemInit(tmp.get());
  // I imagine there is already a mapping to an Arrow type id somewhere in cudf, but for example...
  NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(tmp.get(), NANOARROW_TYPE_INT32);
  
  ArrowSchemaMove(tmp.get(), out);
  return NANOARROW_OK;
}

The nanoarrow C library (also without any CUDA integration) can also populate an ArrowArray for you. If you want export an ArrowArray that's actually non-owning (just pretending to be owning), you could do:

int export_column_view_array(const cudf::column_view& col, ArrowArray* out) {
  nanoarrow::UniqueArray tmp;
  NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromType(tmp.get(), NANOARROW_TYPE_INT32);
  tmp->length = col.length;
  // offset, null_count
  tmp->buffers[1] = col.data_buffer_start_addr;
  // If validity bitmaps are a thing in cudf: tmp->buffers[0] = col.validity_buffer;
  
  ArrowArrayMove(tmp.get(), out);
  return NANOARROW_OK;
}

If you want to export an ArrowArray that fully conforms to the spec (i.e., it is safe to access buffer content until the consumer calls the release callback), you could also use nanoarrow but you would have to explode ownership to the buffer level, which it sounds like might involve some shared pointers or reference counting of some kind. Hypothetically:

static void finalize_buffer(ArrowBufferAllocator* allocator, uint8_t* ptr, int64_t size) {
  auto* shared_col = reinterpret_cast<std::shared_ptr<cuda::column>>(allocator->private_data);
  delete shared_col;
}

int export_column_view_array(const std::shared_ptr<cuda::column> col, ArrowArray* out) {
  nanoarrow::UniqueArray tmp;
  NANOARROW_RETURN_NOT_OK(ArrowArrayInitFromType(tmp.get(), NANOARROW_TYPE_INT32);
  tmp->length = col->length;
  
  ArrowBuffer* data_buffer = ArrowArrayBuffer(tmp.get(), 1);
  ArrowBufferSetAllocator(
    data_buffer, 
    ArrowBufferDeallocator(&finalize_buffer, new std::shared_ptr<cuda::column>(col));
  data_buffer->data = col->data_buffer_start_addr;
  
  NANOARROW_RETURN_NOT_OK(ArrowArrayFinishBuilding(tmp.get(), nullptr. NANOARROW_VALIDATION_LEVEL_MINIMAL);
  ArrowArrayMove(tmp.get(), out);
  return NANOARROW_OK;
}

The only CUDA-specific part would be ensuring that the cudaEvent_t pointer in the ArrowDeviceArray struct is cleaned up when the outermost ArrowArray's release callback is called.

@vyasr is correct that there is an ArrowArrayView (and, in the work-in-progress device helpers, an ArrowDeviceArrayView). It sounds like this is the equivalent of your cudf::column_view; however, it's not ABI stable and so I'm not sure it will be all that useful to use as an interface.

paleolimbot avatar Feb 07 '24 01:02 paleolimbot

I haven't read the full thread in detail, but here's my $0.02.

As far as I'm concerned, the whole reason for Arrow's existence and why RAPIDS built on it in the first place was to enable zero-copy data sharing with a common vocabulary for in-memory, columnar data.

My memory is hazy, but I believe the only reason the original cudf::to/from_arrow perform deep copies is because there wasn't yet a way to describe GPU memory with the Arrow data structures, so we had to always assume we had to make copies to/from host memory.

Now it seems we have a zero-copy way to describe GPU memory with Arrow, so libcudf should definitely enable that.

In my mind, this is equivalent to if you're a C++ library that has your own custom string type, you better provide a std::string_view conversion operator.

jrhemstad avatar Feb 09 '24 18:02 jrhemstad

I agree with that, I just want to make sure that we're leveraging newer Arrow tools and concepts (the C Data Interface, nanoarrow, etc) to the maximum extent possible, which also means making sure that we understand exactly what those tools have to offer and whether there is missing functionality that we should be helping to implement. The questions I'm asking are focused on filling the gaps in my understanding.

Ownership

The questions around proper ownership are ultimately quite similar to, for instance, how cuDF Python works. All objects allocated by libcudf algorithms are immediately forced to relinquish their ownership to Python objects that maintain their lifetime, and downstream algorithms then operate on views anyway so it doesn't matter that libcudf no longer owns the memory. It seems to me then that the proper signature would be ArrowDeviceArray to_arrow_device_array(unique_ptr<column> column) in this context because an ownership transfer would indeed be the only way to get proper interop with other Arrow consumers that are expecting a shared ownership model like Matt indicated above. On the flip side, for consuming arrow objects it seems like we'd want cudf::column_view from_arrow_device_array(ArrowDeviceArray) because we can only ever make a view since we cannot claim sole ownership of the data.

Am I missing anything here? It seems like these are the only ways to provide semantics that are consistent with the goal of minimizing data copies while also producing objects that are consistent with the Arrow spec. There is a fundamental difference between the existing implementations in libcudf and the new ones we're proposing here because the host versions always require copies whereas we want the device ones to never(? or maybe sometimes, in which we'd need different versions of the APIs.) make copies.

Object Creation

This is where I was hoping that nanoarrow could help, and thanks to @paleolimbot we have some good examples. The example @zeroshade linked above looks like it's on the right track, and it seems like it could be written to use nanoarrow instead of arrow APIs based on the examples @paleolimbot showed above. If not, is there missing functionality that we should be helping to add? Creating those structures with nanoarrow seems like exactly what it's intended for and would allow the resulting library to have no direct dependency on libarrow, which would be nice and probably be a template for something we'd try to do with our existing Arrow host-data interop APIs eventually.

Where does the code live

Based on the above I certainly think it makes sense for libcudf to own the logic for mapping our internal representation of Arrow data into Arrow's structs. What I would hope is that it would be possible to use nanoarrow to allocate the necessary Arrow structs and then ideally to use nanoarrow APIs to populate those structs within a libcudf-specific function that knows how to translate between our types and our groupings of (Arrow-compliant) data buffers into Arrow's types and Arrows structs. But Matt brings up a few points regarding that:

The other issue I see with pushing this upstream is that the ArrowDeviceArray on the arrow side is supposed to be device agnostic. Any init function we provide would also need to have the device type passed in (we can't assume CUDA) which would also require the caller to pass in any synchronization event (if required).

I seem to recall discussions around Arrow device data also discussing this and designing for the need to pass around synchronization (CUDA) events. @kkraus14 can probably say more, but isn't Arrow already designing for this in some places? Are you thinking that it's just overkill in this context?

Essentially the only thing a helper function like what you are asking could do is be a wrapper around populating a C struct, which seems a little redundant and unnecessary. At least to me. Nanoarrow could certainly be used to simplify the creating of the ArrowArray and ArrowStruct objects though. @paleolimbot could comment further on nanoarrow for this

It seems like using nanoarrow here would at least be helpful and not redundant as a way to protect against future non-ABI-breaking changes in the spec, e.g. if arrow arrays added fields at the end of the struct (that didn't change the alignment). And that also isn't all that different from what's outline in the example above. Maybe I'm exaggerating the likelihood of meaningful changes like this though and reaching for an external tool rather than adding this code to libcudf is unnecessarily complex.

vyasr avatar Feb 09 '24 19:02 vyasr

If not, is there missing functionality that we should be helping to add?

If you do end up using nanoarrow and find that there is some missing functionality, feel free to open an issue! I'm happy to implement or coordinate implementing anything that helps, or to add anything that cudf had to implement themselves that would be useful to a wider audience.

if arrow arrays added fields at the end of the struct

I'm almost positive that we've stated that we won't do that...the main thing that I think nanoarrow can help you with is populating ArrowSchema and ArrowArray structs for nested things like structs and lists in a way that makes sure they're cleaned up. It's not hard to implement that, necessarily, but the details are fiddly.

paleolimbot avatar Feb 09 '24 20:02 paleolimbot

ArrowDeviceArray to_arrow_device_array(unique_ptr<column> column) in this context because an ownership transfer would indeed be the only way to get proper interop with other Arrow consumers that are expecting a shared ownership model like Matt indicated above

That signature should work fine. libcudf tends to use exceptions rather than explicit status returns right? As far as handling the shared ownership model, it's also possible for the ArrowDeviceArray that is constructed to have a no-op for its release callback meaning that no guarantee is provided that the returned ArrowDeviceArray keeps anything alive. So you could have two options, one that takes a unique_ptr<column> and guarantees it maintains the shared ownership and one that takes a column_view and leaves the responsibility on the consumer to maintain the lifetime as long as is necessary.

On the flip side, for consuming arrow objects it seems like we'd want cudf::column_view from_arrow_device_array(ArrowDeviceArray) because we can only ever make a view since we cannot claim sole ownership of the data.

Well, not exactly. When you import from an ArrowDeviceArray you can claim sole ownership of the arrow data in terms of ensuring the release callback inside of the ArrowArray structs gets called upon destruction of the corresponding columns. Currently I don't think that cudf::column_view provides any way to provide anything custom in the destruction, but you could return an unique_ptr<cudf::column> with a custom deleter that calls the release callback?

There is a fundamental difference between the existing implementations in libcudf and the new ones we're proposing here because the host versions always require copies whereas we want the device ones to never(? or maybe sometimes, in which we'd need different versions of the APIs.) make copies.

You've got it exactly correct here.

It seems like using nanoarrow here would at least be helpful and not redundant as a way to protect against future non-ABI-breaking changes in the spec, e.g. if arrow arrays added fields at the end of the struct (that didn't change the alignment). And that also isn't all that different from what's outline in the example above.

Using nanoarrow should definitely be at least helpful for constructing the C structs. I can put together an example code sample if you'd like or I can just start working on a PR for libcudf using nanoarrow (remember that nanoarrow is intended to be vendored/embedded so that would be part of the PR)

zeroshade avatar Feb 09 '24 20:02 zeroshade

Thank you everyone for this discussion. It seems like we are directionally aligned with a few open questions. At this point I would encourage @zeroshade to prepare a draft PR. RAPIDS is interested in adding zero-copy interop via arrow, and the Awkward Array team is interested in testing the feature (see #14959).

I see two areas where perhaps we should agree before @zeroshade kicks off a PR. Please share your thoughts if there are other areas.

  • Agree on including nanoarrow and using its interop utilities. I'm in favor of refactoring our interop module to use nanoarrow, especially if it gets us closer to dropping the libarrow dependency. Please share any concerns you may have about including nanoarrow.
  • Agree on the design for data lifetime and ownership. Since this is ultimately a performance project, would it be more valuable to build the from or the to first? In general, I would prefer to let the application layer manage data lifetime as much as possible. I'm used to reasoning about libcudf functions which almost always make a copy - so shared ownership is unfamiliar to me.

GregoryKimball avatar Feb 09 '24 21:02 GregoryKimball