diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d484c2d4a10c..37fda31d8bd9 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -465,6 +465,7 @@ add_library( src/hash/xxhash_64.cu src/interop/dlpack.cpp src/interop/arrow_utilities.cpp + src/interop/arrow_data_structures.cpp src/interop/to_arrow_device.cu src/interop/to_arrow_host.cu src/interop/from_arrow_device.cu diff --git a/cpp/include/cudf/interop.hpp b/cpp/include/cudf/interop.hpp index 276a1ea77e29..475c0bbc3cf9 100644 --- a/cpp/include/cudf/interop.hpp +++ b/cpp/include/cudf/interop.hpp @@ -25,6 +25,8 @@ #include #include +#include + #include struct DLManagedTensor; @@ -37,6 +39,14 @@ struct ArrowArray; struct ArrowArrayStream; +///@cond +// These are types from arrow that we are forward declaring for our API to +// avoid needing to include nanoarrow headers. +typedef int32_t ArrowDeviceType; // NOLINT + +#define ARROW_DEVICE_CUDA 2 // NOLINT +///@endcond + namespace CUDF_EXPORT cudf { /** * @addtogroup interop_dlpack @@ -130,6 +140,315 @@ using unique_schema_t = std::unique_ptr; */ using unique_device_array_t = std::unique_ptr; +/** + * @brief typedef for a vector of owning columns, used for conversion from ArrowDeviceArray + * + */ +using owned_columns_t = std::vector>; + +/** + * @brief functor for a custom deleter to a unique_ptr of table_view + * + * When converting from an ArrowDeviceArray, there are cases where data can't + * be zero-copy (i.e. bools or non-UINT32 dictionary indices). This custom deleter + * is used to maintain ownership over the data allocated since a `cudf::table_view` + * doesn't hold ownership. + */ +template +struct custom_view_deleter { + /** + * @brief Construct a new custom view deleter object + * + * @param owned Vector of owning columns + */ + explicit custom_view_deleter(owned_columns_t&& owned) : owned_mem_{std::move(owned)} {} + + /** + * @brief operator to delete the unique_ptr + * + * @param ptr Pointer to the object to be deleted + */ + void operator()(ViewType* ptr) const { delete ptr; } + + owned_columns_t owned_mem_; ///< Owned columns that must be deleted. +}; + +/** + * @brief typedef for a unique_ptr to a `cudf::table_view` with custom deleter + * + */ +using unique_table_view_t = + std::unique_ptr>; + +/** + * @brief typedef for a unique_ptr to a `cudf::column_view` with custom deleter + * + */ +using unique_column_view_t = + std::unique_ptr>; + +namespace interop { + +struct arrow_array_container; + +/** + * @brief Helper function to generate empty column metadata (column with no + * name) for arrow conversion. + * + * This function is helpful for internal conversions between host and device + * data using existing arrow functions. It is also convenient for external + * usage of the libcudf Arrow APIs to produce the canonical mapping from cudf + * column names to Arrow column names (i.e. empty names with appropriate + * nesting). + * + * @param input The column to generate metadata for + * @return The metadata for the column + */ +cudf::column_metadata get_column_metadata(cudf::column_view const& input); + +/** + * @brief Helper function to generate empty table metadata (all columns with no + * names) for arrow conversion. + * + * This function is helpful for internal conversions between host and device + * data using existing arrow functions. It is also convenient for external + * usage of the libcudf Arrow APIs to produce the canonical mapping from cudf + * column names to Arrow column names (i.e. empty names with appropriate + * nesting). + * + * @param input The table to generate metadata for + * @return The metadata for the table + */ +std::vector get_table_metadata(cudf::table_view const& input); + +/** + * @brief A standard interchange medium for ArrowDeviceArray data in cudf. + * + * This class provides a way to work with ArrowDeviceArray data in cudf without + * sacrificing the APIs expected of a cudf column. On the other end, it + * provides the shared lifetime management expected by arrow consumers rather + * than the single-owner mechanism of cudf::column. + */ +class arrow_column { + public: + /** + * @brief Construct a new arrow column object + * + * The input array will be moved into the arrow_column, so it is no longer + * suitable for use afterwards. For consistency, this is done even if the + * source array points to host data. + * + * @param schema Arrow schema for the column + * @param input ArrowDeviceArray data for the column + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_column(ArrowSchema&& schema, + ArrowDeviceArray&& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Construct a new arrow column object + * + * The input array will be released, so it is no longer suitable for use + * afterwards. This is done for consistency with other constructors of arrow_table even though the + * source data is always host data. + * + * @param schema Arrow schema for the column + * @param input ArrowArray data for the column + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_column(ArrowSchema&& schema, + ArrowArray&& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Construct a new arrow column object + * + * The input column will be moved into the arrow_column, so it is no longer + * suitable for use afterwards. + * + * @param input cudf column to convert to arrow + * @param metadata Column metadata for the column + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_column(cudf::column&& input, + column_metadata const& metadata, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Convert the column to an ArrowSchema + * + * The resulting schema is a deep copy of the arrow_column's schema and is + * not tied to its lifetime. + * + * @param output ArrowSchema to populate with the column's schema + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + void to_arrow_schema( + ArrowSchema* output, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + /** + * @brief Convert the column to an ArrowDeviceArray + * + * @param output ArrowDeviceArray to populate with the column's data + * @param device_type ArrowDeviceType to set on the output + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + void to_arrow(ArrowDeviceArray* output, + ArrowDeviceType device_type = ARROW_DEVICE_CUDA, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + // TODO: mutable_view + /** + * @brief Get a view of the column data + * + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + * @return unique_column_view_t containing a view of the column data + */ + [[nodiscard]] unique_column_view_t view( + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + private: + std::shared_ptr + container; ///< Shared pointer to container for the ArrowDeviceArray data; shared_ptr allows + ///< re-export via to_arrow +}; + +/** + * @brief A standard interchange medium for ArrowDeviceArray data in cudf. + * + * This class provides a way to work with ArrowDeviceArray data in cudf without + * sacrificing the APIs expected of a cudf table. On the other end, it + * provides the shared lifetime management expected by arrow consumers rather + * than the single-owner mechanism of cudf::table. + */ +class arrow_table { + public: + /** + * @brief Construct a new arrow table object + * + * The input array will be moved into the arrow_table, so it is no longer + * suitable for use afterwards. For consistency, this is done even if the + * source array points to host data. + * + * @param schema Arrow schema for the table + * @param input ArrowDeviceArray data for the table + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_table(ArrowSchema&& schema, + ArrowDeviceArray&& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Construct a new arrow table object + * + * The stream will be released after the table is created, so it is no longer + * suitable for use afterwards. This is done for consistency with other constructors of + * arrow_table even though the source data is always host data. + * + * @param input ArrowArrayStream data for the table + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_table(ArrowArrayStream&& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Construct a new arrow table object + * + * The input array will be released, so it is no longer suitable for use + * afterwards. This is done for consistency with other constructors of arrow_table even though the + * source data is always host data. + * + * @param schema Arrow schema for the table + * @param input ArrowArray data for the table + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_table(ArrowSchema&& schema, + ArrowArray&& input, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Construct a new arrow table object + * + * The input table will be moved into the arrow_table, so it is no longer + * suitable for use afterwards. + * + * @param input cudf table to convert to arrow + * @param metadata The hierarchy of names of columns and children + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + arrow_table(cudf::table&& input, + cudf::host_span metadata, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + + /** + * @brief Get a view of the table data + * + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + * @return unique_table_view_t containing a view of the table data + */ + [[nodiscard]] unique_table_view_t view( + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + /** + * @brief Convert the table to an ArrowSchema + * + * The resulting schema is a deep copy of the arrow_column's schema and is + * not tied to its lifetime. + * + * @param output ArrowSchema to populate with the table's schema + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + void to_arrow_schema( + ArrowSchema* output, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + /** + * @brief Convert the table to an ArrowDeviceArray + * + * @param output ArrowDeviceArray to populate with the table's data + * @param device_type ArrowDeviceType to set on the output + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used for any allocations during conversion + */ + void to_arrow(ArrowDeviceArray* output, + ArrowDeviceType device_type = ARROW_DEVICE_CUDA, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + private: + std::shared_ptr + container; ///< Shared pointer to container for the ArrowDeviceArray data; shared_ptr allows + ///< re-export via to_arrow +}; + +} // namespace interop + /** * @brief Create ArrowSchema from cudf table and metadata * @@ -431,46 +750,6 @@ std::unique_ptr from_arrow_host_column( rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); -/** - * @brief typedef for a vector of owning columns, used for conversion from ArrowDeviceArray - * - */ -using owned_columns_t = std::vector>; - -/** - * @brief functor for a custom deleter to a unique_ptr of table_view - * - * When converting from an ArrowDeviceArray, there are cases where data can't - * be zero-copy (i.e. bools or non-UINT32 dictionary indices). This custom deleter - * is used to maintain ownership over the data allocated since a `cudf::table_view` - * doesn't hold ownership. - */ -template -struct custom_view_deleter { - /** - * @brief Construct a new custom view deleter object - * - * @param owned Vector of owning columns - */ - explicit custom_view_deleter(owned_columns_t&& owned) : owned_mem_{std::move(owned)} {} - - /** - * @brief operator to delete the unique_ptr - * - * @param ptr Pointer to the object to be deleted - */ - void operator()(ViewType* ptr) const { delete ptr; } - - owned_columns_t owned_mem_; ///< Owned columns that must be deleted. -}; - -/** - * @brief typedef for a unique_ptr to a `cudf::table_view` with custom deleter - * - */ -using unique_table_view_t = - std::unique_ptr>; - /** * @brief Create `cudf::table_view` from given `ArrowDeviceArray` and `ArrowSchema` * @@ -514,13 +793,6 @@ unique_table_view_t from_arrow_device( rmm::cuda_stream_view stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); -/** - * @brief typedef for a unique_ptr to a `cudf::column_view` with custom deleter - * - */ -using unique_column_view_t = - std::unique_ptr>; - /** * @brief Create `cudf::column_view` from given `ArrowDeviceArray` and `ArrowSchema` * diff --git a/cpp/include/cudf_test/nanoarrow_utils.hpp b/cpp/include/cudf_test/nanoarrow_utils.hpp index a1211a16e103..faeaea9e1d97 100644 --- a/cpp/include/cudf_test/nanoarrow_utils.hpp +++ b/cpp/include/cudf_test/nanoarrow_utils.hpp @@ -17,6 +17,7 @@ #pragma once #include +#include #include #include #include @@ -416,3 +417,42 @@ get_decimal_precision() else return cudf::detail::max_precision(); } + +struct VectorOfArrays { + std::vector arrays; + nanoarrow::UniqueSchema schema; + size_t index{0}; + + static int get_schema(ArrowArrayStream* stream, ArrowSchema* out_schema) + { + auto private_data = static_cast(stream->private_data); + + NANOARROW_THROW_NOT_OK(ArrowSchemaDeepCopy(private_data->schema.get(), out_schema)); + return 0; + } + + static int get_next(ArrowArrayStream* stream, ArrowArray* out_array) + { + auto private_data = static_cast(stream->private_data); + if (private_data->index >= private_data->arrays.size()) { + out_array->release = nullptr; + return 0; + } + ArrowArrayMove(private_data->arrays[private_data->index++].get(), out_array); + return 0; + } + + static const char* get_last_error(ArrowArrayStream* stream) { return nullptr; } + + static void release(ArrowArrayStream* stream) + { + delete static_cast(stream->private_data); + } +}; + +void makeStreamFromArrays(std::vector arrays, + nanoarrow::UniqueSchema schema, + ArrowArrayStream* out); + +std::tuple, nanoarrow::UniqueSchema, ArrowArrayStream> +get_nanoarrow_stream(int num_copies); diff --git a/cpp/src/interop/arrow_data_structures.cpp b/cpp/src/interop/arrow_data_structures.cpp new file mode 100644 index 000000000000..6100a34ce244 --- /dev/null +++ b/cpp/src/interop/arrow_data_structures.cpp @@ -0,0 +1,373 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include + +namespace cudf::interop { + +/** + * @brief A wrapper around ArrowDeviceArray data used for flexible lifetime management. + * + * The arrow_array_container is the core object for storing and managing the + * lifetime of arrow data in libcudf. Ultimately, data is always owned by this + * container type regardless of the source. There are a few important cases to + * consider: + * 1. We are importing third-party Arrow device data: The data is moved + * directly into the container. + * 2. We are converting a cudf arrow/table: We construct an ArrowDeviceArray + * that owns the data formerly owned by the cudf object and then fall back + * to case 1 with the new array. + * 3. We are importing third-party Arrow host data: We construct a cudf + * column/table from the Arrow data and then fall back to case 2. + * + * Any export of arrow_column or arrow_table to an arrow array produces + * an ArrowDeviceArray whose private_data is an instance of a cudf-internal + * type (the ArrowArrayPrivateData struct) that also holds a shared pointer to + * this container, ensuring shared ownership and of the data and compatible + * management of its lifetime. All the release semantics boil down to a simple + * deletion of a shared_ptr, so no actually freeing needs to be done manually. + * The shared_ptr's reference counting is sufficient across all use cases. The + * original array's release callback is called when the container is + * destructed, which in practice means when all shared references to the + * container are gone. + */ +struct arrow_array_container { + arrow_array_container() = default; + + template + arrow_array_container(ArrowSchema&& schema_, + T input_, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + { + auto output = cudf::to_arrow_device(std::move(input_), stream, mr); + ArrowSchemaMove(&schema_, &schema); + ArrowDeviceArrayMove(output.get(), &owner); + } + + arrow_array_container(ArrowSchema&& schema_, + ArrowDeviceArray&& input_, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + { + switch (input_.device_type) { + case ARROW_DEVICE_CUDA: + case ARROW_DEVICE_CUDA_HOST: + case ARROW_DEVICE_CUDA_MANAGED: { + ArrowSchemaMove(&schema_, &schema); + ArrowDeviceArrayMove(&input_, &owner); + break; + } + default: CUDF_FAIL("Unsupported ArrowDeviceArray type", std::runtime_error); + } + } + ~arrow_array_container() + { + if (owner.array.release != nullptr) { ArrowArrayRelease(&owner.array); } + } + + ArrowDeviceArray owner{}; //< ArrowDeviceArray that owns the data + ArrowSchema schema{}; //< ArrowSchema that describes the data +}; + +cudf::column_metadata get_column_metadata(cudf::column_view const& input) +{ + cudf::column_metadata meta{}; + std::transform( + input.child_begin(), input.child_end(), std::back_inserter(meta.children_meta), [](auto& cv) { + return get_column_metadata(cv); + }); + return meta; +} + +std::vector get_table_metadata(cudf::table_view const& input) +{ + auto meta = std::vector{}; + std::transform(input.begin(), input.end(), std::back_inserter(meta), [](auto& cv) { + return get_column_metadata(cv); + }); + return meta; +} + +namespace { + +/** + * @brief Private data for an ArrowArray that contains a struct array. + * + * This struct is used to manage the lifetimes of the children of a struct array. + */ +struct ArrowArrayPrivateData { + std::shared_ptr parent; + std::vector> children; + std::vector children_raw; +}; + +/** + * @brief Release callback for an ArrowArray that contains a struct array. + * + * This function is called when the ArrowArray is released. It releases all of the children of the + * struct array. + * + * @param array The ArrowArray to release + */ +void ArrayReleaseCallback(ArrowArray* array) +{ + auto private_data = reinterpret_cast(array->private_data); + for (auto& child : private_data->children) { + child->release(child.get()); + } + delete private_data; + array->release = nullptr; +} + +/** + * @brief Copy an ArrowArray. + * + * This function shallow copies an ArrowArray and all of its children. It is + * used to export cudf arrow objects to user-provided ArrowDeviceArrays. + * + * The @p input must be the ``owner`` member of the @p container OR a child of + * the ``owner`` member of the @p container. If not, the behavior is undefined. + * + * @param output The ArrowArray to copy to + * @param input The ArrowArray to copy from + * @param container The container that owns the data + */ +void copy_array(ArrowArray* output, + ArrowArray const* input, + std::shared_ptr container) +{ + auto private_data = new ArrowArrayPrivateData{container}; + output->length = input->length; + output->null_count = input->null_count; + output->offset = input->offset; + output->n_buffers = input->n_buffers; + output->n_children = input->n_children; + output->buffers = input->buffers; + + if (input->n_children > 0) { + private_data->children_raw.resize(input->n_children); + for (auto i = 0; i < input->n_children; ++i) { + private_data->children.push_back(std::make_unique()); + private_data->children_raw[i] = private_data->children.back().get(); + copy_array(private_data->children_raw[i], input->children[i], container); + } + } + output->children = private_data->children_raw.data(); + output->dictionary = input->dictionary; + output->release = ArrayReleaseCallback; + output->private_data = private_data; +} + +template +void arrow_obj_to_arrow(T& obj, + std::shared_ptr container, + ArrowDeviceArray* output, + ArrowDeviceType device_type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + switch (device_type) { + case ARROW_DEVICE_CUDA: + case ARROW_DEVICE_CUDA_HOST: + case ARROW_DEVICE_CUDA_MANAGED: { + auto& device_arr = container->owner; + copy_array(&output->array, &device_arr.array, container); + output->device_id = device_arr.device_id; + // We can reuse the sync event by reference from the input. The + // destruction of that event is managed by the destruction of + // the underlying ArrowDeviceArray of this table. + output->sync_event = device_arr.sync_event; + output->device_type = device_type; + break; + } + case ARROW_DEVICE_CPU: { + auto out = cudf::to_arrow_host(*obj.view().get(), stream, mr); + ArrowArrayMove(&out->array, &output->array); + output->device_id = -1; + output->sync_event = nullptr; + output->device_type = ARROW_DEVICE_CPU; + break; + } + default: throw std::runtime_error("Unsupported ArrowDeviceArray type"); + } +} + +} // namespace + +arrow_column::arrow_column(cudf::column&& input, + column_metadata const& metadata, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + : container{[&] { + auto table_meta = std::vector{metadata}; + auto tv = cudf::table_view{{input.view()}}; + auto schema = cudf::to_arrow_schema(tv, table_meta); + return std::make_shared( + std::move(*schema->children[0]), std::move(input), stream, mr); + }()} +{ +} + +arrow_column::arrow_column(ArrowSchema&& schema, + ArrowDeviceArray&& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + switch (input.device_type) { + case ARROW_DEVICE_CPU: { + auto col = from_arrow_host_column(&schema, &input, stream, mr); + auto tmp_column = arrow_column(std::move(*col), get_column_metadata(col->view()), stream, mr); + container = tmp_column.container; + // Should always be non-null unless we're in some odd multithreaded + // context but best to be safe. + if (input.array.release != nullptr) { ArrowArrayRelease(&input.array); } + break; + } + default: + container = + std::make_shared(std::move(schema), std::move(input), stream, mr); + } +} + +arrow_column::arrow_column(ArrowSchema&& schema, + ArrowArray&& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + ArrowDeviceArray arr{.array = {}, .device_id = -1, .device_type = ARROW_DEVICE_CPU}; + ArrowArrayMove(&input, &arr.array); + auto tmp = arrow_column(std::move(schema), std::move(arr), stream, mr); + container = tmp.container; +} + +void arrow_column::to_arrow_schema(ArrowSchema* output, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + ArrowSchemaDeepCopy(&container->schema, output); +} + +void arrow_column::to_arrow(ArrowDeviceArray* output, + ArrowDeviceType device_type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + arrow_obj_to_arrow(*this, container, output, device_type, stream, mr); +} + +// If it proves to be a bottleneck we could do this work on construction of the +// container and store the extra columns in the container. Then the container +// can safely return copies of the view ad infinitum and this call can be +// stream- and mr-free, matching the cudf::column::view method. Also doing this +// on construction would allow us to cache column data for the types where the +// representation is not identical between arrow and cudf (like bools) and +// avoiding constant back-and-forth conversion. +unique_column_view_t arrow_column::view(rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + return from_arrow_device_column(&container->schema, &container->owner, stream, mr); +} + +arrow_table::arrow_table(cudf::table&& input, + cudf::host_span metadata, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + : container{[&]() { + auto schema = cudf::to_arrow_schema(input.view(), metadata); + return std::make_shared( + std::move(*schema), std::move(input), stream, mr); + }()} +{ +} + +unique_table_view_t arrow_table::view(rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + return from_arrow_device(&container->schema, &container->owner, stream, mr); +} + +void arrow_table::to_arrow_schema(ArrowSchema* output, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + ArrowSchemaDeepCopy(&container->schema, output); +} + +void arrow_table::to_arrow(ArrowDeviceArray* output, + ArrowDeviceType device_type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + arrow_obj_to_arrow(*this, container, output, device_type, stream, mr); +} + +arrow_table::arrow_table(ArrowSchema&& schema, + ArrowDeviceArray&& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + switch (input.device_type) { + case ARROW_DEVICE_CPU: { + // I'm not sure if there is a more efficient approach than doing this + // back-and-forth conversion without writing a lot of bespoke logic. I + // suspect that the overhead of the memory copies will dwarf any extra + // work here, but it's worth benchmarking to be sure. + auto tbl = from_arrow_host(&schema, &input, stream, mr); + auto tmp_table = arrow_table(std::move(*tbl), get_table_metadata(tbl->view()), stream, mr); + container = tmp_table.container; + if (input.array.release != nullptr) { ArrowArrayRelease(&input.array); } + break; + } + default: + container = + std::make_shared(std::move(schema), std::move(input), stream, mr); + } +} + +arrow_table::arrow_table(ArrowSchema&& schema, + ArrowArray&& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + ArrowDeviceArray arr{.array = {}, .device_id = -1, .device_type = ARROW_DEVICE_CPU}; + ArrowArrayMove(&input, &arr.array); + container = arrow_table(std::move(schema), std::move(arr), stream, mr).container; +} + +arrow_table::arrow_table(ArrowArrayStream&& input, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto tbl = from_arrow_stream(&input, stream, mr); + auto tmp = arrow_table(std::move(*tbl), get_table_metadata(tbl->view()), stream, mr); + container = tmp.container; +} +} // namespace cudf::interop diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 2a32f7191a59..c4cee114f3b5 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -275,6 +275,7 @@ ConfigureTest( # * interop tests ------------------------------------------------------------------------- ConfigureTest( INTEROP_TEST + interop/arrow_data_structures_test.cpp interop/to_arrow_device_test.cpp interop/to_arrow_test.cpp interop/to_arrow_host_test.cpp diff --git a/cpp/tests/interop/arrow_data_structures_test.cpp b/cpp/tests/interop/arrow_data_structures_test.cpp new file mode 100644 index 000000000000..dc6d941bcf27 --- /dev/null +++ b/cpp/tests/interop/arrow_data_structures_test.cpp @@ -0,0 +1,297 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include + +struct ArrowColumnTest : public cudf::test::BaseFixture {}; + +template +auto export_to_arrow(T& obj, ArrowDeviceType device_type = ARROW_DEVICE_CUDA) +{ + // Now we can extract an ArrowDeviceArray from the arrow_column + auto schema = std::make_unique(); + obj.to_arrow_schema(schema.get()); + auto array = std::make_unique(); + obj.to_arrow(array.get(), device_type); + return std::make_pair(std::move(schema), std::move(array)); +} + +TEST_F(ArrowColumnTest, TwoWayConversion) +{ + cudf::test::fixed_width_column_wrapper int_col{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}; + auto col = cudf::column(int_col); + auto arrow_column_from_cudf_column = + cudf::interop::arrow_column(std::move(col), cudf::interop::get_column_metadata(int_col)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, *arrow_column_from_cudf_column.view()); + + auto [arrow_schema_from_arrow_column, arrow_array_from_arrow_column] = + export_to_arrow(arrow_column_from_cudf_column); + arrow_column_from_cudf_column.to_arrow_schema(arrow_schema_from_arrow_column.get()); + arrow_column_from_cudf_column.to_arrow(arrow_array_from_arrow_column.get(), ARROW_DEVICE_CUDA); + + auto arrow_column_from_arrow_array = cudf::interop::arrow_column( + std::move(*arrow_schema_from_arrow_column), std::move(*arrow_array_from_arrow_column)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, *arrow_column_from_arrow_array.view()); +} + +TEST_F(ArrowColumnTest, LifetimeManagement) +{ + cudf::test::fixed_width_column_wrapper int_col{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}; + auto col = std::make_unique(int_col); + auto arrow_column_from_cudf_column = std::make_unique( + std::move(*col), cudf::interop::get_column_metadata(int_col)); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, *arrow_column_from_cudf_column->view()); + + auto [schema1, array1] = export_to_arrow(*arrow_column_from_cudf_column); + auto [schema2, array2] = export_to_arrow(*arrow_column_from_cudf_column); + + // Delete the original owner of the data, then reimport and ensure that we + // are still referencing the same valid original data. + arrow_column_from_cudf_column.reset(); + auto col1 = + std::make_unique(std::move(*schema1), std::move(*array1)); + auto col2 = + std::make_unique(std::move(*schema2), std::move(*array2)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, *col1->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(*col1->view(), *col2->view()); +} + +TEST_F(ArrowColumnTest, ComplexNanoarrowDeviceTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_tables(100); + for (auto i = 0; i < tbl->num_columns(); i++) { + auto& col = tbl->get_column(i); + + ArrowDeviceArray device_arr{ + .array = {}, + .device_id = 0, + .device_type = ARROW_DEVICE_CUDA, + }; + ArrowArrayMove(arr->children[i], &device_arr.array); + auto arrow_column_from_nanoarrow_array = + cudf::interop::arrow_column(std::move(*schema->children[i]), std::move(device_arr)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_column] = + export_to_arrow(arrow_column_from_nanoarrow_array); + arrow_column_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_column_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_column.get(), + ARROW_DEVICE_CUDA); + + auto arrow_column_from_arrow_array = cudf::interop::arrow_column( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_column)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_arrow_array.view()); + } +} + +TEST_F(ArrowColumnTest, ComplexNanoarrowHostTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_host_tables(100); + for (auto i = 0; i < tbl->num_columns(); i++) { + auto& col = tbl->get_column(i); + + ArrowDeviceArray device_arr{ + .array = {}, + .device_id = -1, + .device_type = ARROW_DEVICE_CPU, + }; + ArrowArrayMove(arr->children[i], &device_arr.array); + auto arrow_column_from_nanoarrow_array = + cudf::interop::arrow_column(std::move(*schema->children[i]), std::move(device_arr)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_column] = + export_to_arrow(arrow_column_from_nanoarrow_array); + arrow_column_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_column_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_column.get(), + ARROW_DEVICE_CUDA); + + auto arrow_column_from_arrow_array = cudf::interop::arrow_column( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_column)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_arrow_array.view()); + } +} + +TEST_F(ArrowColumnTest, ComplexNanoarrowHostArrowArrayTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_host_tables(100); + for (auto i = 0; i < tbl->num_columns(); i++) { + auto& col = tbl->get_column(i); + + auto arrow_column_from_nanoarrow_array = + cudf::interop::arrow_column(std::move(*schema->children[i]), std::move(*arr->children[i])); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_column] = + export_to_arrow(arrow_column_from_nanoarrow_array); + arrow_column_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_column_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_column.get(), + ARROW_DEVICE_CUDA); + + auto arrow_column_from_arrow_array = cudf::interop::arrow_column( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_column)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(col.view(), *arrow_column_from_arrow_array.view()); + } +} + +TEST_F(ArrowColumnTest, ToFromHost) +{ + cudf::test::fixed_width_column_wrapper int_col{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}; + auto col = cudf::column(int_col); + auto arrow_column_from_cudf_column = + cudf::interop::arrow_column(std::move(col), cudf::interop::get_column_metadata(int_col)); + + auto [arrow_schema_from_arrow_column, arrow_array_from_arrow_column] = + export_to_arrow(arrow_column_from_cudf_column, ARROW_DEVICE_CPU); + arrow_column_from_cudf_column.to_arrow_schema(arrow_schema_from_arrow_column.get()); + arrow_column_from_cudf_column.to_arrow(arrow_array_from_arrow_column.get(), ARROW_DEVICE_CPU); + + auto arrow_column_from_arrow_array = cudf::interop::arrow_column( + std::move(*arrow_schema_from_arrow_column), std::move(*arrow_array_from_arrow_column)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(int_col, *arrow_column_from_arrow_array.view()); +} + +struct ArrowTableTest : public cudf::test::BaseFixture {}; + +TEST_F(ArrowTableTest, TwoWayConversion) +{ + cudf::test::fixed_width_column_wrapper int_col{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}; + cudf::test::fixed_width_column_wrapper float_col{ + {1., 2., 3., 4., 5., 6., 7., 8., 9., 10.}}; + auto original_view = cudf::table_view{{int_col, float_col}}; + cudf::table table{cudf::table_view{{int_col, float_col}}}; + auto arrow_table_from_cudf_table = + cudf::interop::arrow_table(std::move(table), cudf::interop::get_table_metadata(original_view)); + + CUDF_TEST_EXPECT_TABLES_EQUAL(original_view, *arrow_table_from_cudf_table.view()); + + auto [arrow_schema_from_arrow_table, arrow_array_from_arrow_table] = + export_to_arrow(arrow_table_from_cudf_table); + arrow_table_from_cudf_table.to_arrow_schema(arrow_schema_from_arrow_table.get()); + arrow_table_from_cudf_table.to_arrow(arrow_array_from_arrow_table.get(), ARROW_DEVICE_CUDA); + + auto arrow_table_from_arrow_array = cudf::interop::arrow_table( + std::move(*arrow_schema_from_arrow_table), std::move(*arrow_array_from_arrow_table)); + CUDF_TEST_EXPECT_TABLES_EQUAL(original_view, *arrow_table_from_arrow_array.view()); +} + +TEST_F(ArrowTableTest, ComplexNanoarrowDeviceTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_tables(100); + ArrowDeviceArray device_arr{ + .array = {}, + .device_id = 0, + .device_type = ARROW_DEVICE_CUDA, + }; + ArrowArrayMove(arr.get(), &device_arr.array); + auto arrow_table_from_nanoarrow_array = + cudf::interop::arrow_table(std::move(*schema.get()), std::move(device_arr)); + + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_table] = + export_to_arrow(arrow_table_from_nanoarrow_array); + arrow_table_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_table_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_table.get(), ARROW_DEVICE_CUDA); + + auto arrow_table_from_arrow_array = cudf::interop::arrow_table( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_table)); + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_arrow_array.view()); +} + +TEST_F(ArrowTableTest, ComplexNanoarrowHostTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_host_tables(100); + ArrowDeviceArray device_arr{ + .array = {}, + .device_id = -1, + .device_type = ARROW_DEVICE_CPU, + }; + ArrowArrayMove(arr.get(), &device_arr.array); + auto arrow_table_from_nanoarrow_array = + cudf::interop::arrow_table(std::move(*schema.get()), std::move(device_arr)); + + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_table] = + export_to_arrow(arrow_table_from_nanoarrow_array); + arrow_table_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_table_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_table.get(), ARROW_DEVICE_CUDA); + + auto arrow_table_from_arrow_array = cudf::interop::arrow_table( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_table.get())); + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_arrow_array.view()); +} + +TEST_F(ArrowTableTest, ComplexNanoarrowHostArrowArrayTables) +{ + auto [tbl, schema, arr] = get_nanoarrow_host_tables(100); + auto arrow_table_from_nanoarrow_array = + cudf::interop::arrow_table(std::move(*schema.get()), std::move(*arr.get())); + + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_nanoarrow_array.view()); + + auto [arrow_schema_from_nanoarrow_array, arrow_array_from_arrow_table] = + export_to_arrow(arrow_table_from_nanoarrow_array); + arrow_table_from_nanoarrow_array.to_arrow_schema(arrow_schema_from_nanoarrow_array.get()); + arrow_table_from_nanoarrow_array.to_arrow(arrow_array_from_arrow_table.get(), ARROW_DEVICE_CUDA); + + auto arrow_table_from_arrow_array = cudf::interop::arrow_table( + std::move(*arrow_schema_from_nanoarrow_array), std::move(*arrow_array_from_arrow_table)); + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *arrow_table_from_arrow_array.view()); +} + +TEST_F(ArrowTableTest, ToFromHost) +{ + cudf::test::fixed_width_column_wrapper int_col{{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}}; + cudf::test::fixed_width_column_wrapper float_col{ + {1., 2., 3., 4., 5., 6., 7., 8., 9., 10.}}; + auto original_view = cudf::table_view{{int_col, float_col}}; + cudf::table table{cudf::table_view{{int_col, float_col}}}; + auto arrow_table_from_cudf_table = + cudf::interop::arrow_table(std::move(table), cudf::interop::get_table_metadata(original_view)); + + auto [arrow_schema_from_arrow_table, arrow_array_from_arrow_table] = + export_to_arrow(arrow_table_from_cudf_table, ARROW_DEVICE_CPU); + arrow_table_from_cudf_table.to_arrow_schema(arrow_schema_from_arrow_table.get()); + arrow_table_from_cudf_table.to_arrow(arrow_array_from_arrow_table.get(), ARROW_DEVICE_CPU); + + auto arrow_table_from_arrow_array = cudf::interop::arrow_table( + std::move(*arrow_schema_from_arrow_table), std::move(*arrow_array_from_arrow_table)); + CUDF_TEST_EXPECT_TABLES_EQUAL(original_view, *arrow_table_from_arrow_array.view()); +} + +TEST_F(ArrowTableTest, FromArrowArrayStream) +{ + auto num_copies = 3; + auto [tbl, sch, stream] = get_nanoarrow_stream(num_copies); + + auto result = cudf::interop::arrow_table(std::move(stream)); + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), *result.view()); +} diff --git a/cpp/tests/interop/from_arrow_stream_test.cpp b/cpp/tests/interop/from_arrow_stream_test.cpp index fe4b311cf7a7..d8e11e375bc2 100644 --- a/cpp/tests/interop/from_arrow_stream_test.cpp +++ b/cpp/tests/interop/from_arrow_stream_test.cpp @@ -24,38 +24,6 @@ #include #include -struct VectorOfArrays { - std::vector arrays; - nanoarrow::UniqueSchema schema; - size_t index{0}; - - static int get_schema(ArrowArrayStream* stream, ArrowSchema* out_schema) - { - auto private_data = static_cast(stream->private_data); - - [[maybe_unused]] auto rc = ArrowSchemaDeepCopy(private_data->schema.get(), out_schema); - return 0; - } - - static int get_next(ArrowArrayStream* stream, ArrowArray* out_array) - { - auto private_data = static_cast(stream->private_data); - if (private_data->index >= private_data->arrays.size()) { - out_array->release = nullptr; - return 0; - } - ArrowArrayMove(private_data->arrays[private_data->index++].get(), out_array); - return 0; - } - - static const char* get_last_error(ArrowArrayStream* stream) { return nullptr; } - - static void release(ArrowArrayStream* stream) - { - delete static_cast(stream->private_data); - } -}; - struct FromArrowStreamTest : public cudf::test::BaseFixture {}; void makeStreamFromArrays(std::vector arrays, @@ -70,15 +38,15 @@ void makeStreamFromArrays(std::vector arrays, out->private_data = private_data; } -TEST_F(FromArrowStreamTest, BasicTest) +std::tuple, nanoarrow::UniqueSchema, ArrowArrayStream> +get_nanoarrow_stream(int num_copies) { - constexpr auto num_copies = 3; std::vector> tables; // The schema is unique across all tables. nanoarrow::UniqueSchema schema; std::vector arrays; for (auto i = 0; i < num_copies; ++i) { - auto [tbl, sch, arr] = get_nanoarrow_host_tables(0); + auto [tbl, sch, arr] = get_nanoarrow_host_tables(3); tables.push_back(std::move(tbl)); arrays.push_back(std::move(arr)); if (i == 0) { sch.move(schema.get()); } @@ -91,8 +59,16 @@ TEST_F(FromArrowStreamTest, BasicTest) ArrowArrayStream stream; makeStreamFromArrays(std::move(arrays), std::move(schema), &stream); + return std::make_tuple(std::move(expected), std::move(schema), stream); +} + +TEST_F(FromArrowStreamTest, BasicTest) +{ + constexpr auto num_copies = 3; + auto [tbl, sch, stream] = get_nanoarrow_stream(num_copies); + auto result = cudf::from_arrow_stream(&stream); - CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), result->view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(tbl->view(), result->view()); } TEST_F(FromArrowStreamTest, EmptyTest)