diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7b7f0f051839..187a240198c2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -513,7 +513,7 @@ function(precompile_fixed_width_kernel_fragments) ) set(VARIANT_NAME transform_kernel_${FILE_INDEX}) set(INSTANCE - "cudf::jit::transform_kernel>, cudf::jit::type_list>>" + "cudf::jit::transform_kernel>, cudf::jit::type_list>>" ) add_fragment( cudf_fragments @@ -548,7 +548,7 @@ function(precompile_fixed_width_kernel_fragments) ) set(VARIANT_NAME transform_kernel_${FILE_INDEX}) set(INSTANCE - "cudf::jit::transform_kernel, cudf::jit::column_accessor<1, cudf::column_device_view_core, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list>>" + "cudf::jit::transform_kernel, cudf::jit::column_accessor<1, ${TYPE}, ${RHS_IS_SCALAR}, 0>>, cudf::jit::type_list>>" ) add_fragment( cudf_fragments @@ -578,20 +578,9 @@ endfunction() # amortize the cost of JIT compilation for kernels that have matching signatures function(precompile_string_kernel_fragments) list(APPEND CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES uint8_t uint16_t uint32_t uint64_t - cuda::std::span - ) - list( - APPEND - CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES - cudf::mutable_column_device_view_core - cudf::mutable_column_device_view_core - cudf::mutable_column_device_view_core - cudf::mutable_column_device_view_core - cudf::jit::mutable_strings_column_device_view - ) - foreach(OUTPUT_TYPE OUTPUT_COLUMN_TYPE IN ZIP_LISTS CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES - CUDF_PRECOMPILE_STRING_OUTPUT_COLUMN_TYPES + cudf::mutable_string_view ) + foreach(OUTPUT_TYPE IN LISTS CUDF_PRECOMPILE_STRING_OUTPUT_PHYSICAL_TYPES) # Pre-compile for up to 8 output columns foreach(OUTPUT_INDICES IN ITEMS "0" "0;1" "0;1;2" "0;1;2;3" "0;1;2;3;4" "0;1;2;3;4;5" "0;1;2;3;4;5;6" "0;1;2;3;4;5;6;7" @@ -606,17 +595,15 @@ function(precompile_string_kernel_fragments) set(OUTPUT_ACCESSORS "") set(OUTPUT_POINTERS "") foreach(OUTPUT_INDEX IN LISTS OUTPUT_INDICES) - list( - APPEND - OUTPUT_ACCESSORS - "cudf::jit::column_accessor<${OUTPUT_INDEX}, ${OUTPUT_COLUMN_TYPE}, ${OUTPUT_TYPE}, false, 0>" + list(APPEND OUTPUT_ACCESSORS + "cudf::jit::column_accessor<${OUTPUT_INDEX}, ${OUTPUT_TYPE}, false, 0>" ) list(APPEND OUTPUT_POINTERS "${OUTPUT_TYPE} *") endforeach() list(JOIN OUTPUT_ACCESSORS " ," OUTPUT_ACCESSORS_STR) list(JOIN OUTPUT_POINTERS " ," OUTPUT_POINTERS_STR) set(INSTANCE - "cudf::jit::transform_kernel>, cudf::jit::type_list<${OUTPUT_ACCESSORS_STR}>>" + "cudf::jit::transform_kernel>, cudf::jit::type_list<${OUTPUT_ACCESSORS_STR}>>" ) add_fragment( cudf_fragments diff --git a/cpp/examples/string_transforms/url_logs/fragments.cu b/cpp/examples/string_transforms/url_logs/fragments.cu index d1315b719ed5..e119cb76be5e 100644 --- a/cpp/examples/string_transforms/url_logs/fragments.cu +++ b/cpp/examples/string_transforms/url_logs/fragments.cu @@ -3,11 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include #include -#include struct range32 { int32_t begin{}; @@ -179,18 +179,18 @@ __device__ int compute_url_component_sizes(int32_t* protocol_size, } // Copies the six parsed URL components into their preallocated string buffers. -__device__ int write_url_components(cuda::std::span* protocol, - cuda::std::span* host, - cuda::std::span* port, - cuda::std::span* path, - cuda::std::span* query, - cuda::std::span* fragment, +__device__ int write_url_components(cudf::mutable_string_view* protocol, + cudf::mutable_string_view* host, + cudf::mutable_string_view* port, + cudf::mutable_string_view* path, + cudf::mutable_string_view* query, + cudf::mutable_string_view* fragment, cudf::string_view input) { url_ranges ranges; if (!parse_url(input, &ranges)) { return 0; } - cuda::std::span* outputs[] = {protocol, host, port, path, query, fragment}; - range32 components[] = { + cudf::mutable_string_view* outputs[] = {protocol, host, port, path, query, fragment}; + range32 components[] = { ranges.protocol, ranges.host, ranges.port, ranges.path, ranges.query, ranges.fragment}; for (auto component = 0; component < 6; ++component) { auto range = components[component]; @@ -216,12 +216,12 @@ extern "C" __device__ int transform(int32_t* protocol_size, #else #ifdef UDF_WRITE_OUTPUT // Exposes the component-writing pass through the transform LTO ABI. -extern "C" __device__ int transform(cuda::std::span* protocol, - cuda::std::span* host, - cuda::std::span* port, - cuda::std::span* path, - cuda::std::span* query, - cuda::std::span* fragment, +extern "C" __device__ int transform(cudf::mutable_string_view* protocol, + cudf::mutable_string_view* host, + cudf::mutable_string_view* port, + cudf::mutable_string_view* path, + cudf::mutable_string_view* query, + cudf::mutable_string_view* fragment, cudf::string_view input) { return write_url_components(protocol, host, port, path, query, fragment, input); diff --git a/cpp/examples/string_transforms/url_logs/transforms.cpp b/cpp/examples/string_transforms/url_logs/transforms.cpp index 86e07c204e96..2cf0bb7f826a 100644 --- a/cpp/examples/string_transforms/url_logs/transforms.cpp +++ b/cpp/examples/string_transforms/url_logs/transforms.cpp @@ -205,17 +205,17 @@ __device__ int compute_url_component_sizes(int32_t* protocol_size, // Builds the output UDF from the same parser so both CUDA passes use identical ranges. std::string const url_component_output_udf = std::string{R"***( // Copies the six parsed URL components into their preallocated string buffers. -__device__ int write_url_components(cuda::std::span* protocol, - cuda::std::span* host, - cuda::std::span* port, - cuda::std::span* path, - cuda::std::span* query, - cuda::std::span* fragment, +__device__ int write_url_components(cudf::mutable_string_view* protocol, + cudf::mutable_string_view* host, + cudf::mutable_string_view* port, + cudf::mutable_string_view* path, + cudf::mutable_string_view* query, + cudf::mutable_string_view* fragment, cudf::string_view input) { )***"} + parse_url_udf + R"***( url_ranges ranges; if (!parse_url(&ranges)) { return 0; } - cuda::std::span* outputs[] = {protocol, host, port, path, query, fragment}; + cudf::mutable_string_view* outputs[] = {protocol, host, port, path, query, fragment}; range32 components[] = { ranges.protocol, ranges.host, ranges.port, ranges.path, ranges.query, ranges.fragment}; for (auto component = 0; component < 6; ++component) { diff --git a/cpp/include/cudf/column/column_child_offsets.hpp b/cpp/include/cudf/column/column_child_offsets.hpp index 89baa563e17e..4dd21aad43be 100644 --- a/cpp/include/cudf/column/column_child_offsets.hpp +++ b/cpp/include/cudf/column/column_child_offsets.hpp @@ -14,7 +14,11 @@ namespace CUDF_EXPORT cudf { -static constexpr size_type offsets_column_index = 0; ///< Child index of the offsets column +/** + * @brief Child index of the offsets column. + * @ingroup column_classes + */ +constexpr size_type offsets_column_index = 0; static constexpr size_type dictionary_indices_column_index = 0; ///< Child index of the dictionary indices column @@ -22,4 +26,16 @@ static constexpr size_type dictionary_indices_column_index = static constexpr size_type dictionary_keys_column_index = 1; ///< Child index of the dictionary key column +/** + * @brief Child index of the list offsets column. + * @ingroup lists_classes + */ +constexpr size_type lists_offsets_column_index = offsets_column_index; + +/** + * @brief Child index of the list values column. + * @ingroup lists_classes + */ +constexpr size_type lists_child_column_index = 1; + } // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/column/column_device_view_base.cuh b/cpp/include/cudf/column/column_device_view_base.cuh index 263579199a79..f6f6b9b3c564 100644 --- a/cpp/include/cudf/column/column_device_view_base.cuh +++ b/cpp/include/cudf/column/column_device_view_base.cuh @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -69,7 +70,7 @@ struct nullate { * @tparam KeyType The type of the dictionary keys */ template - requires(is_index_type() && is_relationally_comparable()) + requires(is_index_type()) struct dictionary_element { using index_type = IndexType; ///< The type of the dictionary indices using key_type = KeyType; ///< The type of the dictionary keys @@ -92,6 +93,472 @@ inline constexpr bool is_dictionary_encoded = false; template inline constexpr bool is_dictionary_encoded> = true; +/** + * @brief An element interpreted directly as its stored representation. + * @tparam T The underlying type used to store the value (e.g., `bool` or `uint8_t`, `char`, + * `string_view`, etc.) + * + */ +template +struct repr_element { + using value_type = T; ///< The underlying type used to store the value + + T value{}; ///< The value for this element +}; + +/** + * @brief A type trait to determine if a type is a representation element. + * @tparam T The type to check + */ +template +inline constexpr bool is_repr_element = false; + +/** + * @brief Specialization of `is_repr_element` for `repr_element`. + * @tparam T The underlying type used to store the value + */ +template +inline constexpr bool is_repr_element> = true; + +/** + * @brief A non-owning, type-erased, read-only view of a list row's child elements. + * + * @ingroup lists_classes + * + * The caller must keep the underlying device storage alive. + * List row validity is separate from the child validity retained here. + * Use the column view's nullable_element() to check the list row itself. + * + * Element indices are relative to this list row and must be in `[0, size())`. + * The caller must choose a `T` matching the child column's type. + * This view does not check types at runtime. Representation-compatible types, + * fixed-point types, and `string_view` are supported. + */ +struct list_element { + friend class column_device_view_core; + friend class mutable_column_device_view_core; + + /** + * @brief Returns the number of child elements in this list row, including null elements. + * @return Number of child elements + */ + [[nodiscard]] CUDF_HOST_DEVICE size_type size() const { return size_; } + + /** + * @brief Returns whether this list row contains no child elements. + * @return `true` if size() is zero + */ + [[nodiscard]] CUDF_HOST_DEVICE bool empty() const { return size_ == 0; } + + /** + * @brief Returns whether the child storage has a validity mask. + * @return `true` if a child validity mask is present, even if all children are valid + */ + [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const { return null_mask_ != nullptr; } + + /** + * @brief Returns whether the specified child element is valid. + * @param list_index Child index within this row in `[0, size())` + * @return `true` if the child is valid or no child validity mask is present + */ + [[nodiscard]] CUDF_HOST_DEVICE bool is_valid(size_type list_index) const + { + return bit_value_or(null_mask_, list_index + offset_, true); + } + + /** + * @brief Returns whether the specified child element is null. + * @param list_index Child index within this row in `[0, size())` + * @return `true` if the child is null + */ + [[nodiscard]] CUDF_HOST_DEVICE bool is_null(size_type list_index) const + { + return !is_valid(list_index); + } + + /** + * @brief Reads a boolean child value. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The boolean child type + * @param list_index Child index within this row in `[0, size())` + * @return The child value + */ + template , bool>)> + [[nodiscard]] __device__ T element(size_type list_index) const noexcept + { + return element(list_index) != 0; + } + + /** + * @brief Reads a representation-compatible non-boolean child value. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The representation-compatible non-boolean child type, satisfying + * `is_rep_layout_compatible() && !is_same_v, bool>` + * @param list_index Child index within this row in `[0, size())` + * @return The child value + */ + template () && + !cuda::std::is_same_v, bool>)> + [[nodiscard]] __device__ T const& element(size_type list_index) const noexcept + { + return static_cast(data_)[list_index + offset_]; + } + + /** + * @brief Reads a fixed-point child value using the child column's scale. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The child's fixed-point type + * @param list_index Child index within this row in `[0, size())` + * @return A fixed-point value constructed from the stored representation and scale + */ + template ())> + [[nodiscard]] __device__ T element(size_type list_index) const noexcept + { + using rep = typename T::rep; + return T{numeric::scaled_integer{static_cast(data_)[list_index + offset_], + numeric::scale_type{scale_}}}; + } + + /** + * @brief Returns a read-only view of a string child. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T `string_view` + * @param list_index Child index within this row in `[0, size())` + * @return A view of the child's existing character storage + */ + template , string_view>)> + [[nodiscard]] __device__ T element(size_type list_index) const noexcept + { + auto const offsets = detail::input_offsetalator(offsets_, data_type{offsets_type_}); + auto const begin = offsets[list_index + offset_]; + auto const end = offsets[list_index + offset_ + 1]; + return T{static_cast(data_) + begin, static_cast(end - begin)}; + } + + /** + * @brief Reads a raw child value. + * + * The child must be valid. Use `nullable_element()` to check child validity. + * @tparam T The raw child type, satisfying `is_repr_element` + * @param list_index Child index within this row in `[0, size())` + * @return The child value + */ + template )> + [[nodiscard]] __device__ typename T::value_type const& element( + size_type list_index) const noexcept + { + return static_cast(data_)[list_index + offset_]; + } + + /** + * @brief Reads a child value if it is valid, otherwise returns `nullopt`. + * + * This checks child validity, not the validity of the containing list row. + * @tparam T A child type supported by element() + * @param list_index Child index within this row in `[0, size())` + * @return An optional child value or string view, or `nullopt` for a null child + */ + template + [[nodiscard]] __device__ auto nullable_element(size_type list_index) const noexcept + { + using value_type = cuda::std::remove_cvref_t(list_index))>; + if (is_null(list_index)) { return cuda::std::optional{}; } + return cuda::std::optional{element(list_index)}; + } + + protected: + /** + * @brief Constructs a read-only view from the child storage of a list row. + * @param data Base pointer to child data or string characters, before applying `offset` + * @param offsets Base pointer to string offsets, or `nullptr` for other child types + * @param null_mask Base pointer to the child validity mask, or `nullptr` if absent + * @param size Number of child elements in the row + * @param scale Fixed-point child scale, ignored for other child types + * @param offset Child index at which the row begins, including the child column's offset + * @param offsets_type `INT32` or `INT64` for string offsets, otherwise `EMPTY` + */ + CUDF_HOST_DEVICE list_element(void const* data, + void const* offsets, + bitmask_type const* null_mask, + size_type size, + numeric::scale_type scale, + size_type offset, + type_id offsets_type) + : data_{data}, + offsets_{offsets}, + null_mask_{null_mask}, + size_{size}, + scale_{scale}, + offset_{offset}, + offsets_type_{offsets_type} + { + } + + void const* data_ = nullptr; ///< Child data or string character buffer + void const* offsets_ = nullptr; ///< String offsets storage + bitmask_type const* null_mask_ = nullptr; ///< Child validity mask + size_type size_ = 0; ///< Number of child elements in the row + int32_t scale_ = 0; ///< Fixed-point scale + size_type offset_ = 0; ///< Row start within the child storage + type_id offsets_type_ = type_id::EMPTY; ///< Type of string offsets +}; + +/** + * @brief A non-owning, type-erased view for reading and modifying a list row's child elements. + * + * @ingroup lists_classes + * + * Inherits read access and validity queries from `list_element`. Mutable column device views + * construct this view from writable child storage, which must remain alive while the view + * is used. Assignment and validity updates apply the stored child offset and affect only + * the children, not the containing list row's validity. + * + * Writes must stay within the existing child range. String characters may be modified + * through element(), but string offsets and lengths cannot be changed. + * The caller must choose a type matching the child column. No runtime type check is performed. + */ +struct mutable_list_element { + friend class mutable_column_device_view_core; + + /** + * @brief Returns the number of child elements in this list row, including null elements. + * @return Number of child elements + */ + [[nodiscard]] CUDF_HOST_DEVICE size_type size() const { return size_; } + + /** + * @brief Returns whether this list row contains no child elements. + * @return `true` if size() is zero + */ + [[nodiscard]] CUDF_HOST_DEVICE bool empty() const { return size_ == 0; } + + /** + * @brief Returns whether the child storage has a validity mask. + * @return `true` if a child validity mask is present, even if all children are valid + */ + [[nodiscard]] CUDF_HOST_DEVICE bool nullable() const { return null_mask_ != nullptr; } + + /** + * @brief Returns whether the specified child element is valid. + * @param list_index Child index within this row in `[0, size())` + * @return `true` if the child is valid or no child validity mask is present + */ + [[nodiscard]] CUDF_HOST_DEVICE bool is_valid(size_type list_index) const + { + return bit_value_or(null_mask_, list_index + offset_, true); + } + + /** + * @brief Returns whether the specified child element is null. + * @param list_index Child index within this row in `[0, size())` + * @return `true` if the child is null + */ + [[nodiscard]] CUDF_HOST_DEVICE bool is_null(size_type list_index) const + { + return !is_valid(list_index); + } + + /** + * @brief Reads a representation-compatible child value. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The child type, satisfying `is_rep_layout_compatible()` + * @param list_index Child index within this row in `[0, size())` + * @return The child value + */ + template ())> + [[nodiscard]] __device__ T& element(size_type list_index) const noexcept + { + return static_cast(data_)[list_index + offset_]; + } + + /** + * @brief Reads a fixed-point child value using the child column's scale. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The child's fixed-point type + * @param list_index Child index within this row in `[0, size())` + * @return A fixed-point value constructed from the stored representation and scale + */ + template ())> + [[nodiscard]] __device__ T element(size_type list_index) const noexcept + { + using rep = typename T::rep; + return T{numeric::scaled_integer{static_cast(data_)[list_index + offset_], + numeric::scale_type{scale_}}}; + } + + /** + * @brief Returns a writable view of a valid string child's existing character storage. + * + * Use nullable_element() to check child validity before accessing a nullable child. + * Modifications must stay within the string's existing byte range. + * @tparam T `mutable_string_view` + * @param list_index Child index within this row in `[0, size())` + * @return A mutable view of the child's characters without changing offsets or length + */ + template , mutable_string_view>)> + [[nodiscard]] __device__ T element(size_type list_index) const noexcept + { + auto const offsets = detail::input_offsetalator(offsets_, data_type{offsets_type_}); + auto const begin = offsets[list_index + offset_]; + auto const end = offsets[list_index + offset_ + 1]; + return T{static_cast(data_) + begin, static_cast(end - begin)}; + } + + /** + * @brief Reads a raw child value. + * + * The child must be valid. Use `nullable_element()` to check child validity. + * @tparam T The raw child type, satisfying `is_repr_element` + * @param list_index Child index within this row in `[0, size())` + * @return The child value + */ + template )> + [[nodiscard]] __device__ typename T::value_type& element(size_type list_index) const noexcept + { + return static_cast(data_)[list_index + offset_]; + } + + /** + * @brief Reads a child value or view if it is valid, otherwise returns `nullopt`. + * + * This checks child validity, not the validity of the containing list row. + * @tparam T A child type supported by element(), including `mutable_string_view` + * @param list_index Child index within this row in `[0, size())` + * @return An optional child value or view, or `nullopt` for a null child + */ + template + [[nodiscard]] __device__ auto nullable_element(size_type list_index) const noexcept + { + using value_type = cuda::std::remove_cvref_t(list_index))>; + if (is_null(list_index)) { return cuda::std::optional{}; } + return cuda::std::optional{element(list_index)}; + } + + /** + * @brief Assigns a representation-compatible child value without changing its validity. + * @param list_index Child index within this row in `[0, size())` + * @param value The value to store + */ + template ())> + __device__ void assign(size_type list_index, T value) const noexcept + { + static_cast(data_)[list_index + offset_] = value; + } + + /** + * @brief Assigns a fixed-point child's representation without rescaling or changing validity. + * + * Only `value.value()` is stored. The child column's scale is unchanged; the caller is + * responsible for supplying a representation appropriate for that scale. + * @param list_index Child index within this row in `[0, size())` + * @param value The fixed-point value whose representation is stored + */ + template ())> + __device__ void assign(size_type list_index, T value) const noexcept + { + using rep = typename T::rep; + static_cast(data_)[list_index + offset_] = value.value(); + } + + /** + * @brief Performs no assignment for a mutable string view. + * + * String characters are modified in place through element(). + * This method does not copy characters, change offsets or length, or update validity. + * @tparam T `mutable_string_view` + * @param list_index Child index within this row in `[0, size())` + * @param value The mutable string view to assign + */ + template , mutable_string_view>)> + __device__ void assign([[maybe_unused]] size_type list_index, + [[maybe_unused]] T value) const noexcept + { + } + + /** + * @brief Assigns `value` to the element at `list_index` + * + * @tparam T The element type + * @param list_index Child index within this row in `[0, size())` + * @param value The value to assign + */ + template )> + __device__ void assign(size_type list_index, typename T::value_type value) const noexcept + { + static_cast(data_)[list_index + offset_] = value; + } + +#ifdef __CUDACC__ + /** + * @brief Marks a child element null using an atomic update to the child validity mask. + * + * Requires nullable() to be true. The containing list row's validity is unchanged. + * @param list_index Child index within this row in `[0, size())` + */ + __device__ void set_null(size_type list_index) const noexcept + { + clear_bit(null_mask_, list_index + offset_); + } + + /** + * @brief Marks a child element valid using an atomic update to the child validity mask. + * + * Requires nullable() to be true. The containing list row's validity is unchanged. + * @param list_index Child index within this row in `[0, size())` + */ + __device__ void set_valid(size_type list_index) const noexcept + { + set_bit(null_mask_, list_index + offset_); + } +#endif + + protected: + /** + * @brief Constructs a mutable view from the writable child storage of a list row. + * @param data Base pointer to writable child data or string characters + * @param offsets Base pointer to string offsets, or `nullptr` for other child types + * @param null_mask Base pointer to the writable child validity mask, or `nullptr` if absent + * @param size Number of child elements in the row + * @param scale Fixed-point child scale, ignored for other child types + * @param offset Child index at which the row begins, including the child column's offset + * @param offsets_type `INT32` or `INT64` for string offsets, otherwise `EMPTY` + */ + CUDF_HOST_DEVICE mutable_list_element(void* data, + void* offsets, + bitmask_type* null_mask, + size_type size, + numeric::scale_type scale, + size_type offset, + type_id offsets_type) + : data_{data}, + offsets_{offsets}, + null_mask_{null_mask}, + size_{size}, + scale_{scale}, + offset_{offset}, + offsets_type_{offsets_type} + { + } + + void* data_ = nullptr; ///< Child data or string character buffer + void* offsets_ = nullptr; ///< String offsets storage + bitmask_type* null_mask_ = nullptr; ///< Child validity mask + size_type size_ = 0; ///< Number of child elements in the row + int32_t scale_ = 0; ///< Fixed-point scale + size_type offset_ = 0; ///< Row start within the child storage + type_id offsets_type_ = type_id::EMPTY; ///< Type of string offsets +}; + namespace detail { /** * @brief An immutable, non-owning view of device data as a column of elements @@ -425,9 +892,9 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba } /** - * @brief Returns a copy of the element at the specified index + * @brief Returns a copy of the boolean element at the specified index * - * Boolean values are normalized: zero is false and any nonzero byte is true. + * The boolean value is normalized: zero is false and any nonzero byte is true. * * If the element at the specified index is NULL, i.e., * `is_null(element_index) == true`, then any attempt to use the result will @@ -435,24 +902,43 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba * * This function accounts for the offset. * - * This function does not participate in overload resolution if `is_rep_layout_compatible` is - * false. Specializations of this function may exist for types `T` where - *`is_rep_layout_compatible` is false. + * This function does not participate in overload resolution if + * `is_same_v, bool>` is false. * * @tparam T The element type * @param element_index Position of the desired element * @return The element at the specified index */ - template ())> + template , bool>)> [[nodiscard]] __device__ T element(size_type element_index) const noexcept { - if constexpr (cuda::std::is_same_v, bool>) { - // BOOL8 permits any nonzero byte for true. Read the storage as a byte before converting - // to bool, since loading a noncanonical representation as bool is undefined behavior. - return data()[element_index] != 0; - } else { - return data()[element_index]; - } + // BOOL8 permits any nonzero byte for true. Read the storage as a byte before converting + // to bool, since loading a noncanonical representation as bool is undefined behavior. + return element(element_index) != 0; + } + + /** + * @brief Returns a copy of the element at the specified index + * + * If the element at the specified index is NULL, i.e., + * `is_null(element_index) == true`, then any attempt to use the result will + * lead to undefined behavior. + * + * This function accounts for the offset. + * + * This function does not participate in overload resolution if `is_rep_layout_compatible && + * !is_same_v, bool>` is false. + * + * @tparam T The element type + * @param element_index Position of the desired element + * @return The element at the specified index + */ + template () && + !cuda::std::is_same_v, bool>)> + [[nodiscard]] __device__ T const& element(size_type element_index) const noexcept + { + return data()[element_index]; } /** @@ -466,7 +952,8 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba * @param element_index Position of the desired string element * @return string_view instance representing this element at this index */ - template )> + template , string_view>)> [[nodiscard]] __device__ T element(size_type element_index) const noexcept { size_type index = element_index + offset(); // account for this view's _offset @@ -498,7 +985,7 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba } /** - * @brief Returns a decoded copy of the element at the specified index. + * @brief Returns a decoded element at the specified index. * * If the element at the specified index is NULL, i.e., * `is_null(element_index) == true`, then any attempt to use the result will @@ -524,20 +1011,77 @@ class alignas(16) column_device_view_core : public detail::column_device_view_ba return keys.template element(index); } + /** + * @brief Returns a read-only view of the child elements of the list at the specified index. + * + * This function accounts for the column offset. + * + * The column must have type `LIST` and a supported child type. + * + * Fixed-point children include their scale, and string children retain their offsets and + * character buffer. + * + * @tparam T `list_element` + * @param element_index Position of the desired list row in `[0, size())` + * @return A view of the row's child elements in device storage + */ + template , list_element>)> + [[nodiscard]] __device__ T element(size_type element_index) const noexcept + { + auto const& offsets = child(lists_offsets_column_index); + auto const begin = offsets.template element(element_index + offset()); + auto const end = offsets.template element(element_index + offset() + 1); + auto const& values = child(lists_child_column_index); + if (values.type().id() == type_id::STRING) { + auto const string_offsets = values.child(offsets_column_index); + return T{values.head(), + string_offsets.head(), + values.null_mask(), + end - begin, + numeric::scale_type{0}, + values.offset() + begin, + string_offsets.type().id()}; + } else { + return T{values.head(), + nullptr, + values.null_mask(), + end - begin, + numeric::scale_type{values.type().scale()}, + values.offset() + begin, + type_id::EMPTY}; + } + } + + /** + * @brief Reads a raw child value. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The raw child type, satisfying `is_repr_element` + * @param element_index Position of the desired element in `[0, size())` + * @return The child value + */ + template )> + [[nodiscard]] __device__ typename T::value_type const& element( + size_type element_index) const noexcept + { + return static_cast(_data)[_offset + element_index]; + } + /** * @brief Returns a nullable element at the specified index. If the element is null, returns * `nullopt`. * * @param element_index Position of the desired element - * @return `optional` containing the element at the specified index, or `nullopt` if the element - * is null + * @return `optional` containing the decoded value returned by `element()`, or `nullopt` + * if the element is null. For dictionary tags, the optional contains the decoded key type. */ template - [[nodiscard]] __device__ cuda::std::optional nullable_element( - size_type element_index) const noexcept + [[nodiscard]] __device__ auto nullable_element(size_type element_index) const noexcept { - if (is_null(element_index)) { return cuda::std::nullopt; } - return element(element_index); + using value_type = cuda::std::remove_cvref_t(element_index))>; + if (is_null(element_index)) { return cuda::std::optional{}; } + return cuda::std::optional{element(element_index)}; } /** @@ -699,7 +1243,8 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device * @param element_index Position of the desired string element * @return string_view instance representing this element at this index */ - template )> + template , string_view>)> [[nodiscard]] __device__ T element(size_type element_index) const noexcept { size_type index = element_index + offset(); // account for this view's _offset @@ -729,6 +1274,108 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device return T{scaled_integer{data()[element_index], scale}}; } + /** + * @brief Returns a mutable view into the existing string character buffer. + * + * The column offset is applied. The element must be valid, and mutations must stay + * within the existing string's byte range. + * @tparam T `mutable_string_view` + * @param element_index Position of the desired string element + * @return A mutable view of the string's preallocated character storage + */ + template , mutable_string_view>)> + [[nodiscard]] __device__ T element(size_type element_index) const noexcept + { + auto const index = element_index + offset(); + auto const offsets = child(offsets_column_index); + auto const itr = detail::input_offsetalator(offsets.head(), offsets.type()); + auto const begin = itr[index]; + auto const end = itr[index + 1]; + return T{head() + begin, static_cast(end - begin)}; + } + + /** + * @brief Returns a decoded element at the specified index. + * + * If the element at the specified index is NULL, i.e., + * `is_null(element_index) == true`, then any attempt to use the result will + * lead to undefined behavior. + * + * This function accounts for the offset. + * + * This function does not participate in overload resolution if `is_dictionary_encoded` is + * false. + * + * @tparam T The element type, e.g., `dictionary_element` for a dictionary column + * with `int32_t` indices and `float` keys + * @param element_index Position of the desired element + * @return The element at the specified index + */ + template )> + [[nodiscard]] __device__ decltype(auto) element(size_type element_index) const noexcept + { + auto const& indices = child(dictionary_indices_column_index); + auto const& keys = child(dictionary_keys_column_index); + auto const index = indices.template element( + element_index + offset()); // account for this view's _offset + return keys.template element(index); + } + + /** + * @brief Returns a view of a list row's child elements in their existing storage. + * + * The column offset and child validity mask are retained. The list row must be valid; + * use nullable_element() to check row validity. Child validity is checked separately + * through the returned list view. Request `list_element` for read-only access or + * `mutable_list_element` for in-place writes within the existing child range. + * @tparam T `list_element` or `mutable_list_element` + * @param element_index Position of the desired list row + * @return A non-owning view of the list row's child storage + */ + template , list_element> || + cuda::std::is_same_v, mutable_list_element>)> + [[nodiscard]] __device__ T element(size_type element_index) const noexcept + { + auto const offsets = child(lists_offsets_column_index); + auto const begin = offsets.template element(element_index + offset()); + auto const end = offsets.template element(element_index + offset() + 1); + auto const values = child(lists_child_column_index); + if (values.type().id() == type_id::STRING) { + auto const string_offsets = values.child(offsets_column_index); + return T{values.head(), + string_offsets.head(), + values.null_mask(), + end - begin, + numeric::scale_type{0}, + values.offset() + begin, + string_offsets.type().id()}; + } else { + return T{values.head(), + nullptr, + values.null_mask(), + end - begin, + numeric::scale_type{values.type().scale()}, + values.offset() + begin, + type_id::EMPTY}; + } + } + + /** + * @brief Reads a raw child value. + * + * The child must be valid. Use nullable_element() to check child validity. + * @tparam T The raw child type, satisfying `is_repr_element` + * @param element_index Position of the desired element in `[0, size())` + * @return The child value + */ + template )> + [[nodiscard]] __device__ typename T::value_type& element(size_type element_index) const noexcept + { + return static_cast(const_cast(_data))[_offset + element_index]; + } + /** * @brief Returns a nullable element at the specified index. If the element is null, returns * `nullopt`. @@ -738,11 +1385,11 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device * is null */ template - [[nodiscard]] __device__ cuda::std::optional nullable_element( - size_type element_index) const noexcept + [[nodiscard]] __device__ auto nullable_element(size_type element_index) const noexcept { - if (is_null(element_index)) { return cuda::std::nullopt; } - return element(element_index); + using value_type = cuda::std::remove_cvref_t(element_index))>; + if (is_null(element_index)) { return cuda::std::optional{}; } + return cuda::std::optional{element(element_index)}; } /** @@ -775,6 +1422,35 @@ class alignas(16) mutable_column_device_view_core : public detail::column_device data()[element_index] = value.value(); } + /** + * @brief Assigns `value` to the element at `element_index`. + * This is a no-op for mutable string and list elements. + * + * @tparam T The element type + * @param element_index Position of the desired element + * @param value The value to assign + */ + template , mutable_string_view> || + cuda::std::is_same_v, mutable_list_element>)> + __device__ void assign(size_type element_index, T value) const noexcept + { + // no-op for mutable string and list elements + } + + /** + * @brief Assigns `value` to the element at `element_index` + * + * @tparam T The element type + * @param element_index Position of the desired element + * @param value The value to assign + */ + template )> + __device__ void assign(size_type element_index, typename T::value_type value) const noexcept + { + static_cast(const_cast(_data))[_offset + element_index] = value; + } + /** * @brief Returns raw pointer to the underlying bitmask allocation. * diff --git a/cpp/include/cudf/lists/lists_column_view.hpp b/cpp/include/cudf/lists/lists_column_view.hpp index dce853bb2ca9..8e44ffd4bcd5 100644 --- a/cpp/include/cudf/lists/lists_column_view.hpp +++ b/cpp/include/cudf/lists/lists_column_view.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include @@ -51,8 +52,10 @@ class lists_column_view : private column_view { */ lists_column_view& operator=(lists_column_view&&) = default; - static constexpr size_type offsets_column_index{0}; ///< The index of the offsets column - static constexpr size_type child_column_index{1}; ///< The index of the child column + static constexpr size_type offsets_column_index = + lists_offsets_column_index; ///< The index of the offsets column + static constexpr size_type child_column_index = + lists_child_column_index; ///< The index of the child column using column_view::child_begin; using column_view::child_end; diff --git a/cpp/include/cudf/strings/mutable_string_view.hpp b/cpp/include/cudf/strings/mutable_string_view.hpp new file mode 100644 index 000000000000..9653b4db7100 --- /dev/null +++ b/cpp/include/cudf/strings/mutable_string_view.hpp @@ -0,0 +1,65 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace CUDF_EXPORT cudf { + +/** + * @brief A non-owning, mutable view of device data representing a UTF-8 string. + * + * @ingroup strings_classes + * + * The caller must maintain the device memory for the lifetime of this instance. + * Writes must stay within the existing byte range and preserve valid UTF-8 encoding. + * Access to the device memory must occur in a kernel. + */ +class mutable_string_view { + public: + /** @brief Default constructor represents an empty string. */ + CUDF_HOST_DEVICE mutable_string_view() {} + + /** + * @brief Create a view of an existing device char array. + * + * @param data Device char array encoded in UTF-8 + * @param bytes Number of bytes in the array + */ + CUDF_HOST_DEVICE mutable_string_view(char* data, size_type bytes) : _data{data}, _bytes{bytes} {} + + /** + * @brief Return a pointer to the mutable device array. + * @return A pointer to the mutable device array + */ + [[nodiscard]] CUDF_HOST_DEVICE char* data() const { return _data; } + + /** + * @brief Return the number of bytes in this string. + * @return The number of bytes in this string + */ + [[nodiscard]] CUDF_HOST_DEVICE size_type size_bytes() const { return _bytes; } + + /** + * @brief Return whether this string is empty. + * @return true if the string has no bytes + */ + [[nodiscard]] CUDF_HOST_DEVICE bool empty() const { return size_bytes() == 0; } + + /** + * @brief Return an immutable view of the same string. + * @return A string_view of the same device array and byte range + */ + [[nodiscard]] CUDF_HOST_DEVICE operator string_view() const { return string_view{_data, _bytes}; } + + private: + char* _data{}; ///< Pointer to the device char array + size_type _bytes{}; ///< Number of bytes in the array + [[maybe_unused]] mutable size_type + _length{}; ///< Most recently computed UTF-8 character count. Intentionally left unused for + ///< ABI-compatibility with string_view +}; + +} // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/transform.hpp b/cpp/include/cudf/transform.hpp index a29561b615c3..6b8d36689dba 100644 --- a/cpp/include/cudf/transform.hpp +++ b/cpp/include/cudf/transform.hpp @@ -49,6 +49,8 @@ struct transform_output { output_nullability nullability = output_nullability::PRESERVE; ///< Signifies if a null mask should be created for the output ///< column + + std::vector children = {}; ///< Children specifications for nested types }; /** @@ -85,9 +87,9 @@ struct transform_output { * @brief Describes a transform input independently of a particular column. * * An input specification contains the type information needed to reflect and retrieve a transform - * kernel. Dictionary specifications recursively describe their indices and keys through `children`. - * String specifications retain their offsets child type so `INT32` and `INT64` layouts can be - * distinguished. + * kernel. Dictionary and list specifications recursively describe their children through + * `children`. String specifications retain their offsets child type so `INT32` and `INT64` layouts + * can be distinguished. */ struct transform_input_spec { type_id type = type_id::EMPTY; ///< Logical type of the input @@ -95,15 +97,15 @@ struct transform_input_spec { bool is_scalar = false; ///< Whether the input is presented to the UDF as a scalar std::vector children = - {}; ///< Specifications of dictionary children or string offsets + {}; ///< Specifications of dictionary/list children or string offsets }; /** * @brief Describes a transform output independently of a particular output column. * - * The string-offset setting identifies the device-view representation required by the kernel. The - * nullability setting is retained so inputs to `transform_program::run` can be validated against - * the output policy used to construct the program. + * The offset setting for string and list outputs identifies the device-view representation required + * by the kernel. The nullability setting is retained so inputs to `transform_program::run` can be + * validated against the output policy used to construct the program. */ struct transform_output_spec { type_id type = type_id::EMPTY; ///< Logical type of the output @@ -111,7 +113,8 @@ struct transform_output_spec { output_nullability nullability = output_nullability::PRESERVE; ///< Null-mask policy for the output - bool has_string_offsets = false; ///< Whether a string output uses preallocated offsets + bool has_offsets = false; ///< Whether a string or list output uses supplied offsets + std::vector children = {}; ///< Specifications of string offsets or nested child columns }; @@ -146,7 +149,7 @@ struct transform_program { * UDF by `run` * @param inputs Inputs from which to derive the input specifications * @param outputs Outputs from which to derive the output type and nullability specifications - * @param string_offsets Optional string offsets used to determine each string output + * @param output_offsets Optional string and required list offsets used to determine the output * representation */ transform_program(std::string const& udf, @@ -155,7 +158,7 @@ struct transform_program { std::optional user_data, std::span inputs, std::span outputs, - std::span const> string_offsets); + std::span const> output_offsets); /** * @brief Constructs a reusable program from explicit input and output specifications. @@ -215,13 +218,14 @@ struct transform_program { /** * @brief Runs the transform program on the given inputs and outputs. * - * @throws std::invalid_argument if the inputs, outputs, or string offsets are not compatible with + * @throws std::invalid_argument if the inputs, outputs, or output offsets are not compatible with * the specifications used to construct the program * * @param inputs The inputs to the transform program * @param outputs The outputs of the transform program - * @param string_offsets For string output columns, the offsets can be pre-allocated and passed in - * to prevent overhead of compacting string views into run-end strings column. + * @param output_offsets For string and list output columns, offsets can be preallocated and + * passed in. String offsets avoid compaction; list offsets are required and allocate the child + * storage. * @param row_size The row size of the transform operation. If not provided, it will be inferred * from the inputs. * @param stream CUDA stream used for device memory operations and kernel launches @@ -232,7 +236,7 @@ struct transform_program { std::unique_ptr run( std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); @@ -268,18 +272,17 @@ struct transform_program { * - String-view outputs (`cudf::string_view`) have uninitialized values. * * @throws std::invalid_argument if any of the input columns have different sizes (except scalars) - * @throws std::invalid_argument if `output_type` or any of the inputs are not fixed-width or string - * types + * @throws std::invalid_argument if an output is not fixed-width, STRING, or fixed-width LIST type. * @throws std::invalid_argument if the inputs only have a scalar with no column inputs and * `row_size` is not provided. This is because the row size cannot be inferred from the inputs in * this case. * @throws std::invalid_argument if any of the output or input types are not supported. - * CUDA-supported input types are fixed-width, string, and their dictionary types. PTX-supported - * input types are integrals, floats, and their dictionary types. CUDA-supported output types are - * fixed-width, string, and their dictionary types. PTX-supported output types are integrals, - * floats, and their dictionary types. - * @throws std::invalid_argument if string offsets are provided for non-string output columns, or - * if the number of string offsets does not match the number of output columns. + * CUDA-supported input types include fixed-width, string, dictionary, and lists with fixed-width + * children. PTX-supported input types are integrals, floats, and their dictionary types. + * CUDA-supported output types are fixed-width, STRING, and fixed-width LIST. PTX-supported output + * types are integrals and floats. + * @throws std::invalid_argument if output offsets are provided for outputs other than strings or + * lists, or if the number of output offsets does not match the number of output columns. * @throws cudf::evaluation_error if the UDF produces an error during execution. * * The size of the resulting column is the `row_size` if provided, otherwise it is inferred from @@ -291,8 +294,8 @@ struct transform_program { * @param user_data User-defined device data to pass to the UDF. * @param inputs Immutable views of the inputs to transform (columns and scalar columns) * @param outputs Specification of the output columns to be created - * @param string_offsets For string output columns, the offsets can be pre-allocated and passed in - * to prevent overhead of compacting string views into run-end strings column. + * @param output_offsets For string and list output columns, offsets can be preallocated and passed + * in. String offsets avoid compaction; list offsets are required and allocate the child storage. * @param row_size The row size of the transform operation. If not provided, it is inferred from the * input columns. * @param stream CUDA stream used for device memory operations and kernel launches @@ -308,7 +311,7 @@ std::unique_ptr
transform( std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); @@ -325,7 +328,7 @@ std::unique_ptr
transform( * @param user_data User-defined device data to pass to the UDF * @param inputs Immutable views of the inputs to transform (columns and scalar columns) * @param outputs Specification of the output columns to be created - * @param string_offsets Pre-allocated offsets for string output columns + * @param output_offsets Preallocated offsets for string or list output columns * @param row_size The row size of the transform operation * @param stream CUDA stream used for device memory operations and kernel launches * @param mr Device memory resource used to allocate the returned column's device memory @@ -338,7 +341,7 @@ std::unique_ptr
transform( std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); @@ -348,7 +351,8 @@ std::unique_ptr
transform( */ enum class lto_binary_type : uint8_t { LTO_IR, //< LTO-IR binary - FATBIN //< FATBIN binary + FATBIN, //< FATBIN binary + PTX //< PTX binary }; /** @@ -360,13 +364,13 @@ enum class lto_binary_type : uint8_t { * * * @throws std::invalid_argument if any of the input columns have different sizes (except scalars) - * @throws std::invalid_argument if `output_type` or any of the inputs are not fixed-width or string - * types + * @throws std::invalid_argument if an output is not fixed-width, STRING, or fixed-width LIST, or + * dictionary type. * @throws std::invalid_argument if the inputs only have a scalar with no column inputs and * `row_size` is not provided. This is because the row size cannot be inferred from the inputs in * this case - * @throws std::invalid_argument if string offsets are provided for non-string output columns, or - * if the number of string offsets does not match the number of output columns + * @throws std::invalid_argument if output offsets are provided for outputs other than strings or + * lists, or if the number of output offsets does not match the number of output columns * @throws cudf::evaluation_error if the UDF produces an error during execution * * The size of the resulting column is the `row_size` if provided, otherwise it is inferred from @@ -379,8 +383,8 @@ enum class lto_binary_type : uint8_t { * @param user_data User-defined device data to pass to the UDF * @param inputs Immutable views of the inputs to transform * @param outputs Specification of the output columns to be created - * @param string_offsets For string output columns, the offsets can be pre-allocated and passed in - * to prevent overhead of compacting string views into run-end strings column. + * @param output_offsets For string and list output columns, offsets can be preallocated and passed + * in. String offsets avoid compaction; list offsets are required and allocate the child storage. * @param row_size The row size of the transform operation. If not provided, it is inferred from the * input columns * @param stream CUDA stream used for device memory operations and kernel launches @@ -396,7 +400,7 @@ std::unique_ptr
transform_lto( std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream = cudf::get_default_stream(), rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); diff --git a/cpp/include/cudf/types.hpp b/cpp/include/cudf/types.hpp index cc3b30c98e8b..1a3467eb77cb 100644 --- a/cpp/include/cudf/types.hpp +++ b/cpp/include/cudf/types.hpp @@ -44,6 +44,7 @@ class column; class column_view; class mutable_column_view; class string_view; +class mutable_string_view; class list_view; class struct_view; class scalar; diff --git a/cpp/src/jit/column_accessor.cuh b/cpp/src/jit/column_accessor.cuh index 9851bbe4b868..d2d7d65b7e9d 100644 --- a/cpp/src/jit/column_accessor.cuh +++ b/cpp/src/jit/column_accessor.cuh @@ -8,26 +8,22 @@ #include #include -#include #include -#include #include namespace cudf { namespace jit { -template +template struct column_accessor { static constexpr int32_t index = Index; - static constexpr int32_t table_index = TableIndex; - using column_type = Column; using element_type = Element; - using optional_element_type = cuda::std::optional; - - static constexpr bool as_scalar = AsScalar; + static constexpr bool as_scalar = AsScalar; + static constexpr int32_t table_index = TableIndex; - static constexpr bool is_mutable_string = - cuda::std::is_same_v>; + static constexpr bool is_inplace_output = + cuda::std::is_same_v || + cuda::std::is_same_v; static __device__ constexpr size_type map_index(size_type row) { @@ -38,69 +34,68 @@ struct column_accessor { } } - template - static __device__ auto& column(T const* __restrict__ cols) - requires(sizeof(T) == sizeof(column_type) && alignof(T) == alignof(column_type)) - { - return reinterpret_cast(cols[index]); - } + static __device__ auto& column(auto const* __restrict__ columns) { return columns[index]; } - static __device__ auto element(auto const* __restrict__ cols, size_type row) + static __device__ decltype(auto) element(auto const* __restrict__ columns, size_type row) { - return column(cols).template element(map_index(row)); + return columns[index].template element(map_index(row)); } - static __device__ bool is_null(auto const* __restrict__ cols, size_type row) + static __device__ bool is_null(auto const* __restrict__ columns, size_type row) { - return column(cols).is_null(map_index(row)); + return columns[index].is_null(map_index(row)); } - static __device__ bool is_valid(auto const* __restrict__ cols, size_type row) + static __device__ bool is_valid(auto const* __restrict__ columns, size_type row) { - return column(cols).is_valid(map_index(row)); + return columns[index].is_valid(map_index(row)); } - static __device__ auto nullable_element(auto const* __restrict__ cols, size_type row) + static __device__ decltype(auto) nullable_element(auto const* __restrict__ columns, size_type row) { - return column(cols).template nullable_element(map_index(row)); + return columns[index].template nullable_element(map_index(row)); } - static __device__ void set_null_mask_word(auto const* __restrict__ cols, + static __device__ void set_null_mask_word(auto const* __restrict__ columns, size_type word_index, bitmask_type word) requires(!as_scalar) { - auto* mask = column(cols).null_mask(); + auto* mask = columns[index].null_mask(); if (mask == nullptr) { return; } mask[word_index] = word; } - static __device__ void assign(auto const* __restrict__ cols, size_type row, element_type value) + static __device__ void assign(auto const* __restrict__ columns, size_type row, auto value) requires(!as_scalar) { - column(cols).template assign(row, value); + columns[index].template assign(row, value); } - static __device__ element_type output_arg(auto const* __restrict__ cols, size_type row) + static __device__ decltype(auto) output_arg(auto const* __restrict__ columns, size_type row) requires(!as_scalar) { - if constexpr (is_mutable_string) { - return element(cols, row); + if constexpr (is_inplace_output) { + return element(columns, row); + } else if constexpr (is_repr_element) { + return typename element_type::value_type{}; } else { - return {}; + return element_type{}; } } - static __device__ optional_element_type null_output_arg(auto const* __restrict__ cols, - size_type row) + static __device__ decltype(auto) nullable_output_arg(auto const* __restrict__ columns, + size_type row) requires(!as_scalar) { - if constexpr (is_mutable_string) { - return element(cols, row); + if constexpr (is_inplace_output) { + return cuda::std::optional(element(columns, row)); + } else if constexpr (is_repr_element) { + return cuda::std::optional{}; } else { - return {}; + return cuda::std::optional{}; } } }; diff --git a/cpp/src/jit/column_device_view_wrappers.cuh b/cpp/src/jit/column_device_view_wrappers.cuh deleted file mode 100644 index 2f170e6c2fd5..000000000000 --- a/cpp/src/jit/column_device_view_wrappers.cuh +++ /dev/null @@ -1,126 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include - -#include -#include - -namespace cudf { -namespace jit { - -/** - * @brief A column wrapper type that treats a column as a vector of elements. - * - */ -struct mutable_vector_device_view : private mutable_column_device_view_core { - using base = mutable_column_device_view_core; - - CUDF_HOST_DEVICE constexpr mutable_vector_device_view(base const& src) : base{src} {} - ~mutable_vector_device_view() = default; - mutable_vector_device_view(mutable_vector_device_view const&) = default; - mutable_vector_device_view(mutable_vector_device_view&&) = default; - mutable_vector_device_view& operator=(mutable_vector_device_view const&) = default; - mutable_vector_device_view& operator=(mutable_vector_device_view&&) = default; - - using base::nullable; - using base::offset; - using base::size; - using base::type; - - template - CUDF_HOST_DEVICE T* __restrict__ data() const noexcept - { - return static_cast(const_cast(_data)) + _offset; - } - - using base::is_null; - using base::is_valid; - using base::null_mask; - - template - [[nodiscard]] __device__ decltype(auto) element(size_type element_index) const noexcept - { - return data()[element_index]; - } - - template - [[nodiscard]] __device__ cuda::std::optional nullable_element( - size_type element_index) const noexcept - { - if (is_null(element_index)) { return cuda::std::nullopt; } - return element(element_index); - } - - template - __device__ void assign(size_type row, T value) const noexcept - { - data()[row] = value; - } -}; - -/** - * @brief A column wrapper type that treats a column as a column of mutable strings. - * The offsets will have been pre-initialized and the chars will have been pre-allocated. - */ -struct mutable_strings_column_device_view : private mutable_column_device_view_core { - using base = mutable_column_device_view_core; - - CUDF_HOST_DEVICE constexpr mutable_strings_column_device_view(base const& src) : base{src} {} - - ~mutable_strings_column_device_view() = default; - mutable_strings_column_device_view(mutable_strings_column_device_view const&) = default; - mutable_strings_column_device_view(mutable_strings_column_device_view&&) = default; - mutable_strings_column_device_view& operator=(mutable_strings_column_device_view const&) = - default; - mutable_strings_column_device_view& operator=(mutable_strings_column_device_view&&) = default; - - using base::is_null; - using base::is_valid; - using base::null_mask; - using base::nullable; - using base::offset; - using base::size; - using base::type; - - template > - [[nodiscard]] __device__ cuda::std::span element(size_type element_index) const noexcept - requires(cuda::std::is_same_v>) - { - auto index = element_index + offset(); - auto chars = static_cast(const_cast(_data)); - auto offsets = child(offsets_column_index); - auto itr = cudf::detail::input_offsetalator(offsets.head(), offsets.type()); - auto beg = itr[index]; - auto end = itr[index + 1]; - auto* __restrict__ str = chars + beg; - return cuda::std::span{str, static_cast(end - beg)}; - } - - template > - [[nodiscard]] __device__ cuda::std::optional> nullable_element( - size_type element_index) const noexcept - requires(cuda::std::is_same_v>) - { - if (is_null(element_index)) { return cuda::std::nullopt; } - return element(element_index); - } - - template > - __device__ void assign(size_type row, cuda::std::span value) const noexcept - requires(cuda::std::is_same_v>) - { - // no-op since we assume the chars have already been pre-allocated and they are mutated - // in-place - return; - } -}; - -} // namespace jit -} // namespace cudf diff --git a/cpp/src/jit/output_columns.hpp b/cpp/src/jit/output_columns.hpp index 56c960828216..1d9fbe72d253 100644 --- a/cpp/src/jit/output_columns.hpp +++ b/cpp/src/jit/output_columns.hpp @@ -140,5 +140,31 @@ struct mutable_strings_column { [[nodiscard]] bitmask_type* null_mask() { return _col->mutable_view().null_mask(); } }; -using output_column = std::variant; +struct mutable_lists_column_view { + mutable_column_view _view; + + [[nodiscard]] auto to_device(cuda::stream_ref stream) const + { + return mutable_column_device_view::create(_view, stream); + } +}; + +// Lists retain the supplied offsets and own the directly written child. +struct mutable_lists_column { + std::unique_ptr _col; + + [[nodiscard]] auto mutable_view() const + { + return mutable_lists_column_view{_col->mutable_view()}; + } + + void set_null_count(size_type count) { _col->set_null_count(count); } + + [[nodiscard]] bool nullable() const { return _col->nullable(); } + + [[nodiscard]] bitmask_type* null_mask() { return _col->mutable_view().null_mask(); } +}; + +using output_column = std:: + variant; } // namespace cudf::jit diff --git a/cpp/src/jit/reflect.cpp b/cpp/src/jit/reflect.cpp index 7693635d9acb..7aee60f3240e 100644 --- a/cpp/src/jit/reflect.cpp +++ b/cpp/src/jit/reflect.cpp @@ -4,6 +4,7 @@ */ #include +#include #include #include #include @@ -23,14 +24,15 @@ #include namespace cudf::jit { + namespace { -std::string get_element_type_name(transform_input_spec const& spec, bool use_physical_type); +std::string reflect_input_element(transform_input_spec const& spec, bool use_physical_type); -struct element_type_name_fn { +struct reflect_input_element_fn { template std::string operator()(transform_input_spec const& spec, bool use_physical_type) const - requires(is_fixed_width() || std::same_as) + requires(is_fixed_width() || std::same_as) { auto type = data_type{spec.type}; return type_to_name(use_physical_type ? jit::physical_type_of(type) : type); @@ -38,74 +40,170 @@ struct element_type_name_fn { template std::string operator()(transform_input_spec const& spec, bool use_physical_type) const - requires(std::same_as) + requires(std::same_as) { return std::format( "cudf::dictionary_element<{}, {}>", - get_element_type_name(spec.children.at(dictionary_indices_column_index), use_physical_type), - get_element_type_name(spec.children.at(dictionary_keys_column_index), use_physical_type)); + reflect_input_element(spec.children.at(dictionary_indices_column_index), use_physical_type), + reflect_input_element(spec.children.at(dictionary_keys_column_index), use_physical_type)); + } + + template + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return "cudf::list_element"; } template std::string operator()(transform_input_spec const& spec, bool) const - requires(!is_fixed_width() && !std::same_as && - !std::same_as) + requires(!is_fixed_width() && !std::same_as && + !std::same_as && !std::same_as) { - CUDF_FAIL("Unsupported type for JIT compilation: " + type_to_name(data_type{spec.type})); + CUDF_FAIL(std::format("Unsupported input element type for JIT input: {}", + type_to_name(data_type{spec.type})), + std::invalid_argument); } }; -std::string get_element_type_name(transform_input_spec const& spec, bool use_physical_type) -{ - return cudf::type_dispatcher( - data_type{spec.type}, element_type_name_fn{}, spec, use_physical_type); -} +std::string reflect_output_element(transform_output_spec const& spec, bool use_physical_type); + +struct reflect_output_element_fn { + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(is_fixed_width()) + { + auto type = data_type{spec.type}; + return type_to_name(use_physical_type ? jit::physical_type_of(type) : type); + } + + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return spec.has_offsets ? "cudf::mutable_string_view" : "cudf::repr_element"; + } + + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return std::format( + "cudf::dictionary_element<{}, {}>", + reflect_output_element(spec.children.at(dictionary_indices_column_index), use_physical_type), + reflect_output_element(spec.children.at(dictionary_keys_column_index), use_physical_type)); + } + + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return "cudf::mutable_list_element"; + } + + template + std::string operator()(transform_output_spec const& spec, bool) const + requires(!is_fixed_width() && !std::same_as && + !std::same_as && !std::same_as) + { + CUDF_FAIL(std::format("Unsupported output element type for JIT output: {}", + type_to_name(data_type{spec.type})), + std::invalid_argument); + } +}; + +struct reflect_input_value_fn { + template + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const + requires(is_fixed_width() || std::same_as) + { + auto type = data_type{spec.type}; + return type_to_name(use_physical_type ? jit::physical_type_of(type) : type); + } + + template + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return reflect_input_element(spec.children.at(dictionary_keys_column_index), use_physical_type); + } + + template + std::string operator()(transform_input_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return "cudf::list_element"; + } + + template + std::string operator()(transform_input_spec const& spec, bool) const + requires(!is_fixed_width() && !std::same_as && + !std::same_as && !std::same_as) + { + CUDF_FAIL(std::format("Unsupported input value type for JIT input: {}", + type_to_name(data_type{spec.type})), + std::invalid_argument); + } +}; + +struct reflect_output_value_fn { + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(is_fixed_width()) + { + auto type = data_type{spec.type}; + return type_to_name(use_physical_type ? jit::physical_type_of(type) : type); + } + + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return spec.has_offsets ? "cudf::mutable_string_view" : "cudf::string_view"; + } + + template + std::string operator()(transform_output_spec const& spec, bool use_physical_type) const + requires(std::same_as) + { + return "cudf::mutable_list_element"; + } + + template + std::string operator()(transform_output_spec const& spec, bool) const + requires(!is_fixed_width() && !std::same_as && !std::same_as) + { + CUDF_FAIL(std::format("Unsupported output value type for JIT output: {}", + type_to_name(data_type{spec.type})), + std::invalid_argument); + } +}; std::string reflect_input_element(transform_input_spec const& spec, bool use_physical_type) { - return get_element_type_name(spec, use_physical_type); + return type_dispatcher(data_type{spec.type}, reflect_input_element_fn{}, spec, use_physical_type); } std::string reflect_output_element(transform_output_spec const& spec, bool use_physical_type) { - if (spec.type == type_id::STRING) { - return spec.has_string_offsets ? "cuda::std::span" : "cudf::string_view"; - } - return get_element_type_name(transform_input_spec{.type = spec.type}, use_physical_type); + return type_dispatcher( + data_type{spec.type}, reflect_output_element_fn{}, spec, use_physical_type); } std::string reflect_input_value_type(transform_input_spec const& spec, bool use_physical_type) { - if (spec.type == type_id::DICTIONARY32) { - return reflect_input_value_type(spec.children.at(dictionary_keys_column_index), - use_physical_type); - } - return reflect_input_element(spec, use_physical_type); + return type_dispatcher(data_type{spec.type}, reflect_input_value_fn{}, spec, use_physical_type); } std::string reflect_output_value_type(transform_output_spec const& spec, bool use_physical_type) { - return reflect_output_element(spec, use_physical_type); -} - -std::string reflect_input_column(transform_input_spec const&) -{ - return "cudf::column_device_view_core"; -} - -std::string reflect_output_column(transform_output_spec const& spec) -{ - if (spec.type == type_id::STRING) { - return spec.has_string_offsets ? "cudf::jit::mutable_strings_column_device_view" - : "cudf::jit::mutable_vector_device_view"; - } - return "cudf::mutable_column_device_view_core"; + return type_dispatcher(data_type{spec.type}, reflect_output_value_fn{}, spec, use_physical_type); } transform_input_spec make_input_spec(column_view const& column, bool is_scalar) { transform_input_spec result{.type = column.type().id(), .is_scalar = is_scalar}; - if (is_dictionary(column.type())) { + if (is_dictionary(column.type()) || column.type().id() == type_id::LIST) { for (size_type i = 0; i < column.num_children(); ++i) { result.children.push_back(make_input_spec(column.child(i), false)); } @@ -137,12 +235,19 @@ transform_output_spec make_output_spec(string_views_column const&) return {.type = type_id::STRING}; } +transform_output_spec make_output_spec(mutable_lists_column const& output) +{ + return { + .type = type_id::LIST, + .has_offsets = true, + .children = {{.type = type_id::INT32}, {.type = output._col->view().child(1).type().id()}}}; +} + transform_output_spec make_output_spec(mutable_strings_column const& output) { auto offsets = output._col->view().child(strings_column_view::offsets_column_index); - return {.type = type_id::STRING, - .has_string_offsets = true, - .children = {{.type = offsets.type().id()}}}; + return { + .type = type_id::STRING, .has_offsets = true, .children = {{.type = offsets.type().id()}}}; } } // namespace @@ -172,13 +277,11 @@ transform_reflection reflect(std::variant sour bool use_physical_types = std::holds_alternative(source_type); for (int32_t i = 0; i < std::ssize(inputs); i++) { auto& in = inputs[i]; - auto column = reflect_input_column(in); auto element = reflect_input_element(in, use_physical_types); bool as_scalar = in.is_scalar; auto accessor = rtcx::reflect_template( "cudf::jit::column_accessor", input_table_sources.empty() ? rtcx::reflect(i) : rtcx::reflect(input_table_indices[i]), - column, element, rtcx::reflect(as_scalar), input_table_sources.empty() ? rtcx::reflect(0) : rtcx::reflect(input_table_sources[i])); @@ -188,13 +291,11 @@ transform_reflection reflect(std::variant sour std::vector out_types; for (int32_t i = 0; i < std::ssize(outputs); i++) { auto& out = outputs[i]; - auto column = reflect_output_column(out); auto element = reflect_output_element(out, use_physical_types); bool as_scalar = false; // never scalar auto accessor = rtcx::reflect_template( "cudf::jit::column_accessor", output_table_sources.empty() ? rtcx::reflect(i) : rtcx::reflect(output_table_indices[i]), - column, element, rtcx::reflect(as_scalar), output_table_sources.empty() ? rtcx::reflect(0) : rtcx::reflect(output_table_sources[i])); @@ -241,18 +342,25 @@ std::vector make_output_specs(std::span make_output_specs( std::span outputs, - std::span const> string_offsets) + std::span const> output_offsets) { - CUDF_EXPECTS(string_offsets.empty() || string_offsets.size() == outputs.size(), - "Number of string offsets must be empty or match the number of outputs", + CUDF_EXPECTS(output_offsets.empty() || output_offsets.size() == outputs.size(), + "Number of output offsets must be empty or match the number of outputs", std::invalid_argument); + std::vector result; for (size_t i = 0; i < outputs.size(); ++i) { - auto has_string_offsets = !string_offsets.empty() && string_offsets[i] != nullptr; - transform_output_spec spec{.type = outputs[i].type.id(), - .nullability = outputs[i].nullability, - .has_string_offsets = has_string_offsets}; - if (has_string_offsets) { spec.children.push_back({.type = string_offsets[i]->type().id()}); } + auto has_offsets = !output_offsets.empty() && output_offsets[i] != nullptr; + transform_output_spec spec{.type = outputs[i].type.id(), + .nullability = outputs[i].nullability, + .has_offsets = has_offsets}; + if (spec.type == type_id::LIST) { + for (auto const& child : outputs[i].children) { + spec.children.push_back({.type = child.type.id()}); + } + } else if (spec.type == type_id::STRING && spec.has_offsets) { + spec.children.push_back({.type = output_offsets[i]->type().id()}); + } result.push_back(std::move(spec)); } return result; @@ -286,14 +394,14 @@ std::string reflect_udf_signature(bool is_null_aware, std::vector in_types; for (auto const& input : inputs) { - auto element = reflect_input_element(input, use_physical_types); + auto element = reflect_input_value_type(input, use_physical_types); in_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}>", element) : element); } std::vector out_types; for (auto const& output : outputs) { - auto element = reflect_output_element(output, use_physical_types); + auto element = reflect_output_value_type(output, use_physical_types); out_types.push_back(is_null_aware ? std::format("cuda::std::optional<{}> *", element) : std::format("{} *", element)); } diff --git a/cpp/src/jit/reflect.hpp b/cpp/src/jit/reflect.hpp index 9794696bde46..b0ac0f02191a 100644 --- a/cpp/src/jit/reflect.hpp +++ b/cpp/src/jit/reflect.hpp @@ -12,7 +12,6 @@ #include #include #include -#include #include #include @@ -24,7 +23,7 @@ struct transform_reflection { std::vector ptx_input_value_types; std::vector ptx_output_value_types; - bool operator==(transform_reflection const& other) const = default; + [[nodiscard]] bool operator==(transform_reflection const& other) const = default; }; /** diff --git a/cpp/src/jit/transform_kernel.cuh b/cpp/src/jit/transform_kernel.cuh index 7915c2a4a0a1..b7dc30ca43c6 100644 --- a/cpp/src/jit/transform_kernel.cuh +++ b/cpp/src/jit/transform_kernel.cuh @@ -104,7 +104,7 @@ __device__ void transform_kernel(size_type row_size, ...); thread_error = cuda::std::max(thread_error, error); }); - }(Out::null_output_arg(output_cols, row_index)...); + }(Out::nullable_output_arg(output_cols, row_index)...); }); } } diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_jit.cu b/cpp/src/join/filter_join_indices/filter_join_indices_jit.cu index e4b3be094edf..d972ee227f8f 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_jit.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices_jit.cu @@ -69,7 +69,6 @@ std::vector build_join_filter_template_params( if (auto* col = std::get_if(&input)) { accessors.emplace_back(rtcx::reflect_template("cudf::jit::column_accessor", rtcx::reflect(i), - "cudf::column_device_view_core", cudf::type_to_name(col->type()), rtcx::reflect(false), rtcx::reflect(table_sources[i].value()))); @@ -78,7 +77,6 @@ std::vector build_join_filter_template_params( accessors.emplace_back(rtcx::reflect_template( "cudf::jit::column_accessor", rtcx::reflect(i), - "cudf::column_device_view_core", cudf::type_to_name(scalar.as_column_view().type()), rtcx::reflect(true), rtcx::reflect(0) // scalars don't belong to a table, so just use 0 as placeholder diff --git a/cpp/src/transform/checked_arithmetic.cu b/cpp/src/transform/checked_arithmetic.cu index a19cce749ca3..88b6a6bfc505 100644 --- a/cpp/src/transform/checked_arithmetic.cu +++ b/cpp/src/transform/checked_arithmetic.cu @@ -117,11 +117,9 @@ CUDF_KERNEL void checked_binary_kernel(size_type row_size, binary_operator op, int32_t* max_error) { - using input_accessors = - jit::type_list, - jit::column_accessor<1, column_device_view_core, T, RhsIsScalar, 0>>; - using output_accessors = - jit::type_list>; + using input_accessors = jit::type_list, + jit::column_accessor<1, T, RhsIsScalar, 0>>; + using output_accessors = jit::type_list>; column_device_view_core const inputs[] = {lhs, rhs}; mutable_column_device_view_core const outputs[] = {out}; @@ -136,10 +134,8 @@ CUDF_KERNEL void checked_unary_kernel(size_type row_size, unary_operator op, int32_t* max_error) { - using input_accessors = - jit::type_list>; - using output_accessors = - jit::type_list>; + using input_accessors = jit::type_list>; + using output_accessors = jit::type_list>; column_device_view_core const inputs[] = {input}; mutable_column_device_view_core const outputs[] = {out}; diff --git a/cpp/src/transform/jit/kernel.cu b/cpp/src/transform/jit/kernel.cu index 77db5230ea60..af7430a914ce 100644 --- a/cpp/src/transform/jit/kernel.cu +++ b/cpp/src/transform/jit/kernel.cu @@ -16,7 +16,6 @@ #include #include -#include #include #pragma nv_hdrstop // The above headers are used by the kernel below and need to be included before diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index de630c8e87ad..3aefe9c92c10 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -20,9 +21,14 @@ #include #include +#include + #include #include +#include #include +#include +#include #include #include @@ -35,6 +41,7 @@ #include #include +#include #include #include @@ -42,14 +49,16 @@ namespace cudf { namespace { using jit::fixed_width_column; +using jit::mutable_lists_column; using jit::mutable_strings_column; using jit::output_column; using jit::string_views_column; using input_column_view = transform_input; -using handle = std::variant< - std::unique_ptr>, - std::unique_ptr>>; + +using handle = std::variant< + std::unique_ptr>, + std::unique_ptr>>; namespace jit_transform { @@ -243,6 +252,7 @@ rtcx::binary_type as_rtcx_binary_type(lto_binary_type type) switch (type) { case lto_binary_type::LTO_IR: return rtcx::binary_type::LTO_IR; case lto_binary_type::FATBIN: return rtcx::binary_type::FATBIN; + case lto_binary_type::PTX: return rtcx::binary_type::PTX; default: CUDF_FAIL( std::format("Unrecognized LTO binary type {} for LTO transform", static_cast(type)), @@ -332,9 +342,13 @@ size_type inplace_null_mask_and(bitmask_type* null_mask, if (!is_nullable) { return 0; } + auto init_mask = [&](bool valid) { + CUDF_CUDA_TRY(cudaMemsetAsync( + null_mask, valid ? 0xff : 0, bitmask_allocation_size_bytes(row_size), stream.get())); + }; if (inputs.empty()) { // no input, set all to valid - set_null_mask(null_mask, 0, row_size, true, stream); + init_mask(true); return 0; } @@ -350,7 +364,7 @@ size_type inplace_null_mask_and(bitmask_type* null_mask, if (scalar_is_null) { // scalar is null, all rows will be null - set_null_mask(null_mask, 0, row_size, false, stream); + init_mask(false); return row_size; } } @@ -360,7 +374,7 @@ size_type inplace_null_mask_and(bitmask_type* null_mask, if (!has_cols) { // no non-scalar columns, so all rows are valid - set_null_mask(null_mask, 0, row_size, true, stream); + init_mask(true); return 0; } @@ -381,7 +395,7 @@ size_type inplace_null_mask_and(bitmask_type* null_mask, if (nullable_masks.empty()) { // we only have non-nullable columns, so all rows are valid - set_null_mask(null_mask, 0, row_size, true, stream); + init_mask(true); return 0; } @@ -452,7 +466,7 @@ void perform_checks(std::variant source_type, std::optional in_row_size, std::span inputs, std::span outputs, - std::span const> string_offsets) + std::span const> output_offsets) { if (auto* udf_source = std::get_if(&source_type); udf_source != nullptr && *udf_source == udf_source_type::PTX) { @@ -485,24 +499,48 @@ void perform_checks(std::variant source_type, CUDF_EXPECTS(std::none_of(outputs.begin(), outputs.end(), [](auto& out) { - return !is_fixed_width(out.type) && out.type.id() != type_id::STRING; + return !(is_fixed_width(out.type) || + out.type.id() == type_id::STRING || + out.type.id() == type_id::LIST); }), - "Transforms only support output of fixed-width or string types", + "Transforms only support fixed-width, STRING, or fixed-width LIST outputs", std::invalid_argument); - static constexpr auto is_input_value_supported = [](auto const& c) { - return is_fixed_width(c.type()) || c.type().id() == type_id::STRING || is_dictionary(c.type()); + for (auto const& out : outputs) { + if (out.type.id() == type_id::LIST) { + CUDF_EXPECTS(out.children.size() == 2, + "LIST outputs require offsets and values child specifications", + std::invalid_argument); + auto const& offsets = out.children[lists_offsets_column_index]; + auto const& values = out.children[lists_child_column_index]; + CUDF_EXPECTS(offsets.type.id() == type_id::INT32 && is_fixed_width(values.type), + "LIST outputs require INT32 offsets and a fixed-width values child", + std::invalid_argument); + } else { + CUDF_EXPECTS(out.children.empty(), + "Child specifications are only supported for LIST outputs", + std::invalid_argument); + } + } + + static constexpr auto is_input_value_supported = [](auto const& self, + column_view const& c) -> bool { + if (is_fixed_width(c.type()) || c.type().id() == type_id::STRING) { return true; } + if (is_dictionary(c.type())) { return self(self, c.child(dictionary_keys_column_index)); } + if (c.type().id() == type_id::LIST) { + return is_fixed_width(c.child(lists_child_column_index).type()); + } + + return false; }; static constexpr auto is_supported_input_type = [&](auto const& c) { auto col = std::visit([](auto const& c) { return jit::as_column_view(c); }, c); - return is_input_value_supported(col) || - (is_dictionary(col.type()) && - is_input_value_supported(col.child(dictionary_keys_column_index))); + return is_input_value_supported(is_input_value_supported, col); }; CUDF_EXPECTS( std::none_of( inputs.begin(), inputs.end(), [&](auto const& in) { return !is_supported_input_type(in); }), - "Transforms only support input of fixed-width, string, or dictionary types", + "Transforms only support fixed-width, string, dictionary, or lists of fixed-width inputs", std::invalid_argument); if (!in_row_size.has_value()) { @@ -526,19 +564,33 @@ void perform_checks(std::variant source_type, "All transform input columns must have the same size", std::invalid_argument); - CUDF_EXPECTS(string_offsets.empty() || (string_offsets.size() == outputs.size()), - "Number of string offsets must be empty or match the number of outputs (with nulls " - "for each non-string column)", + CUDF_EXPECTS(output_offsets.empty() || (output_offsets.size() == outputs.size()), + "Number of output offsets must be empty or match the number of outputs (with nulls " + "for each output without offsets)", std::invalid_argument); CUDF_EXPECTS(std::all_of(cuda::counting_iterator(size_t{0}), - cuda::counting_iterator(string_offsets.size()), + cuda::counting_iterator(output_offsets.size()), [&](auto i) { - if (outputs[i].type.id() == type_id::STRING) { return true; } - return string_offsets.empty() || string_offsets[i] == nullptr; + if (outputs[i].type.id() == type_id::STRING || + outputs[i].type.id() == type_id::LIST) { + return true; + } + return output_offsets.empty() || output_offsets[i] == nullptr; }), - "String offsets must only be provided for string outputs", + "Output offsets must only be provided for string or list outputs", std::invalid_argument); + + for (size_t i = 0; i < outputs.size(); ++i) { + if (outputs[i].type.id() != type_id::LIST) { continue; } + CUDF_EXPECTS(!output_offsets.empty() && output_offsets[i] != nullptr, + "List outputs require supplied offsets", + std::invalid_argument); + auto const& off = *output_offsets[i]; + CUDF_EXPECTS(off.type().id() == type_id::INT32 && !off.nullable() && off.size() == row_size + 1, + "List output offsets must be non-nullable INT32 with row_size + 1 elements", + std::invalid_argument); + } } std::optional> make_stencil( @@ -644,12 +696,80 @@ std::unique_ptr make_strings_column(device_span strin size, std::move(offsets), chars.release(), null_count, std::move(null_mask)); } +struct make_output_column_fn { + template + requires(is_fixed_width()) + output_column operator()(transform_output const& output, + size_type row_size, + cuda::device_buffer&& null_mask, + std::optional>&& offsets, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + { + return fixed_width_column::make(output.type, row_size, std::move(null_mask), 0, stream, mr); + } + + template + requires(std::is_same_v) + output_column operator()(transform_output const& output, + size_type row_size, + cuda::device_buffer&& null_mask, + std::optional>&& offsets, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + { + if (!offsets.has_value()) { + return string_views_column::make(row_size, std::move(null_mask), 0, stream, mr); + } else { + auto chars_size = + strings::detail::get_offset_value(offsets.value()->view(), row_size, stream); + auto chars = rmm::device_buffer{static_cast(chars_size), stream, mr}; + return mutable_strings_column::make( + row_size, std::move(chars), std::move(offsets.value()), std::move(null_mask), 0); + } + } + + template + requires(std::is_same_v) + output_column operator()(transform_output const& output, + size_type row_size, + cuda::device_buffer&& null_mask, + std::optional>&& offsets, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + { + auto count = strings::detail::get_offset_value(offsets.value()->view(), row_size, stream); + auto element_type = output.children.at(lists_child_column_index).type; + // TODO: are nulls in the child columns handled correctly? + // TODO: if children column are nullable, we need to allocate null masks for each of the + // children elements + // Must obey propagation rules + // TODO: null-propagation based on attributes + auto child = make_fixed_width_column(element_type, count, mask_state::UNALLOCATED, stream, mr); + return mutable_lists_column{make_lists_column( + row_size, std::move(offsets.value()), std::move(child), 0, std::move(null_mask))}; + } + + template + requires(!(is_fixed_width() || std::is_same_v || + std::is_same_v)) + output_column operator()(transform_output const& output, + size_type row_size, + cuda::device_buffer&& null_mask, + std::optional>&& offsets, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + { + CUDF_FAIL("Unsupported output type for transform"); + } +}; + auto make_outputs(null_aware is_null_aware, size_type row_size, std::span inputs, std::span outputs, std::span is_output_nullable, - std::vector> string_offsets, + std::vector> output_offsets, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { @@ -662,26 +782,22 @@ auto make_outputs(null_aware is_null_aware, auto nullable = is_output_nullable[i]; auto null_state = nullable ? mask_state::UNINITIALIZED : mask_state::UNALLOCATED; auto null_mask = create_null_mask(row_size, null_state, stream, mr); + auto offsets = std::optional>{}; - if (is_fixed_width(output.type)) { - auto col = - fixed_width_column::make(output.type, row_size, std::move(null_mask), 0, stream, mr); - cols.emplace_back(std::move(col)); - } else if (output.type.id() == type_id::STRING) { - if (string_offsets.empty() || string_offsets[i] == nullptr) { - auto col = string_views_column::make(row_size, std::move(null_mask), 0, stream, mr); - cols.emplace_back(std::move(col)); - } else { - auto chars_size = - strings::detail::get_offset_value(string_offsets[i]->view(), row_size, stream); - auto chars = rmm::device_buffer{static_cast(chars_size), stream, mr}; - auto col = mutable_strings_column::make( - row_size, std::move(chars), std::move(string_offsets[i]), std::move(null_mask), 0); - cols.emplace_back(std::move(col)); - } - } else { - CUDF_UNREACHABLE("Unsupported output type for transform"); + if (!output_offsets.empty() && output_offsets[i] != nullptr) { + offsets = std::move(output_offsets[i]); } + + auto output_column = type_dispatcher(output.type, + make_output_column_fn{}, + output, + row_size, + std::move(null_mask), + std::move(offsets), + stream, + mr); + + cols.emplace_back(std::move(output_column)); } auto stencil = make_stencil(is_null_aware, row_size, inputs, cols, stream); @@ -741,6 +857,13 @@ auto finalize_output(string_views_column&& c, mr); } +auto finalize_output(mutable_lists_column&& c, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + return std::move(c._col); +} + auto finalize_outputs(null_aware is_null_aware, size_type row_size, std::vector outputs, @@ -767,7 +890,7 @@ std::unique_ptr
execute_transform(std::string const& udf, std::optional user_data, std::span inputs, std::span outputs, - std::vector> string_offsets, + std::vector> output_offsets, kernel const* compiled_kernel, cuda::stream_ref stream, rmm::device_async_resource_ref mr) @@ -779,7 +902,7 @@ std::unique_ptr
execute_transform(std::string const& udf, inputs, outputs, output_may_be_nullable, - std::move(string_offsets), + std::move(output_offsets), stream, mr); @@ -835,13 +958,13 @@ std::unique_ptr
transform(std::string const& udf, std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); - perform_checks(source_type, is_null_aware, row_size, inputs, outputs, string_offsets); + perform_checks(source_type, is_null_aware, row_size, inputs, outputs, output_offsets); return execute_transform(udf, source_type, is_null_aware, @@ -849,7 +972,7 @@ std::unique_ptr
transform(std::string const& udf, user_data, inputs, outputs, - std::move(string_offsets), + std::move(output_offsets), nullptr, stream, mr); @@ -861,7 +984,7 @@ std::unique_ptr
multi_transform(std::string const& udf, std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream, rmm::device_async_resource_ref mr) @@ -872,7 +995,7 @@ std::unique_ptr
multi_transform(std::string const& udf, user_data, inputs, outputs, - std::move(string_offsets), + std::move(output_offsets), row_size, stream, mr); @@ -990,13 +1113,13 @@ std::unique_ptr
transform_lto(std::span udf, std::optional user_data, std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional in_row_size, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); - perform_checks(binary_type, is_null_aware, in_row_size, inputs, outputs, string_offsets); + perform_checks(binary_type, is_null_aware, in_row_size, inputs, outputs, output_offsets); auto row_size = in_row_size.has_value() ? *in_row_size : jit::get_projection_size(inputs); auto output_may_be_nullable = get_null_transformation(is_null_aware, inputs, outputs); @@ -1005,7 +1128,7 @@ std::unique_ptr
transform_lto(std::span udf, inputs, outputs, output_may_be_nullable, - std::move(string_offsets), + std::move(output_offsets), stream, mr); auto stencil_arg = stencil.has_value() ? stencil->first : nullptr; @@ -1047,10 +1170,10 @@ struct transform_program::impl { void validate(udf_source_type source_type, std::span actual_inputs, std::span actual_outputs, - std::span const> string_offsets) const + std::span const> output_offsets) const { auto actual_input_specs = jit::make_input_specs(actual_inputs); - auto actual_output_specs = jit::make_output_specs(actual_outputs, string_offsets); + auto actual_output_specs = jit::make_output_specs(actual_outputs, output_offsets); auto actual_reflection = jit::reflect(source_type, actual_input_specs, actual_output_specs); CUDF_EXPECTS(reflection_ == actual_reflection, "Transform program specifications do not match the provided inputs and outputs", @@ -1094,13 +1217,13 @@ transform_program::transform_program(std::string const& udf, std::optional user_data, std::span inputs, std::span outputs, - std::span const> string_offsets) + std::span const> output_offsets) : transform_program(udf, source_type, is_null_aware, user_data, jit::make_input_specs(inputs), - jit::make_output_specs(outputs, string_offsets)) + jit::make_output_specs(outputs, output_offsets)) { } @@ -1163,15 +1286,15 @@ transform_program::~transform_program() = default; std::unique_ptr
transform_program::run(std::span inputs, std::span outputs, - std::vector>&& string_offsets, + std::vector>&& output_offsets, std::optional row_size, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); - impl_->validate(impl_->source_type_, inputs, outputs, string_offsets); + impl_->validate(impl_->source_type_, inputs, outputs, output_offsets); perform_checks( - impl_->source_type_, impl_->is_null_aware_, row_size, inputs, outputs, string_offsets); + impl_->source_type_, impl_->is_null_aware_, row_size, inputs, outputs, output_offsets); return execute_transform({}, impl_->source_type_, impl_->is_null_aware_, @@ -1179,7 +1302,7 @@ std::unique_ptr
transform_program::run(std::span i impl_->user_data_, inputs, outputs, - std::move(string_offsets), + std::move(output_offsets), &impl_->kernel_, stream, mr); diff --git a/cpp/tests/transform/integration/unary_transform_test.cpp b/cpp/tests/transform/integration/unary_transform_test.cpp index f1307cdc0a53..520dc4a8f459 100644 --- a/cpp/tests/transform/integration/unary_transform_test.cpp +++ b/cpp/tests/transform/integration/unary_transform_test.cpp @@ -30,8 +30,15 @@ #include #include +#include +#include +#include +#include #include +#include #include +#include +#include #include #include @@ -1183,7 +1190,7 @@ TEST_F(StringOperationTest, OutputOffsetted) auto c = cudf::test::strings_column_wrapper{"1", "2", "3", "4", "5", "6"}; std::string cuda = R"***( - __device__ void concat(cuda::std::span * out, cudf::string_view a, cudf::string_view b, cudf::string_view c){ + __device__ void concat(cudf::mutable_string_view * out, cudf::string_view a, cudf::string_view b, cudf::string_view c){ auto iter = out->data(); memcpy(iter, a.data(), a.size_bytes()); iter += a.size_bytes(); @@ -1241,7 +1248,7 @@ TEST_F(StringOperationTest, OutputOffsettedMixed) auto c = cudf::test::strings_column_wrapper{"1", "2", "3", "4", "5", "67"}; std::string cuda = R"***( - __device__ void concat(int32_t * ab_size, cuda::std::span * out, bool * is_odd, cudf::string_view a, cudf::string_view b, cudf::string_view c){ + __device__ void concat(int32_t * ab_size, cudf::mutable_string_view * out, bool * is_odd, cudf::string_view a, cudf::string_view b, cudf::string_view c){ auto iter = out->data(); auto begin = iter; memcpy(iter, a.data(), a.size_bytes()); @@ -1842,4 +1849,386 @@ __device__ cudf::errc expression ( } } +using offsets_type = cudf::test::fixed_width_column_wrapper; + +// Construct explicit offsets and storage to avoid initializer-list nesting ambiguity. +std::unique_ptr make_lists(std::initializer_list> rows, + std::span validity = {}) +{ + std::vector boundaries{0}, values; + for (auto row : rows) { + values.insert(values.end(), row.begin(), row.end()); + boundaries.push_back(values.size()); + } + auto mask = cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED); + cudf::size_type null_count = 0; + if (!validity.empty()) { + std::vector parent(rows.size(), 0); + offsets_type source(parent.begin(), parent.end(), validity.begin()); + null_count = static_cast(source).null_count(); + mask = cudf::copy_bitmask(source); + } + return cudf::make_lists_column(rows.size(), + offsets_type(boundaries.begin(), boundaries.end()).release(), + offsets_type(values.begin(), values.end()).release(), + null_count, + std::move(mask)); +} + +std::vector> list_offsets(std::initializer_list values) +{ + std::vector> result; + result.push_back(offsets_type(values).release()); + return result; +} + +cudf::transform_output list_output(cudf::data_type type = cudf::data_type{cudf::type_id::INT32}) +{ + return {cudf::data_type{cudf::type_id::LIST}, + cudf::output_nullability::PRESERVE, + {{cudf::data_type{cudf::type_id::INT32}}, {type}}}; +} + +constexpr auto copy_udf = R"( +__device__ void transform(cudf::mutable_list_element* output, cudf::list_element input) +{ + for (size_t index = 0; index < input.size(); ++index) { + output->assign(index, input.element(index)); + } +} +)"; + +std::unique_ptr transform_lists( + std::span inputs, + std::span outputs, + std::vector> output_offsets, + bool null_aware = false, + std::optional row_count = std::nullopt) +{ + std::string udf = null_aware ? R"( +__device__ void transform(cuda::std::optional* output, + cuda::std::optional input) +{ + if (!input || (!input->empty() && input->element(0) == 7)) { + *output = cuda::std::nullopt; + return; + } + for (size_t index = 0; index < input->size(); ++index) { + (*output)->assign(index, input->element(index)); + } +} +)" + : copy_udf; + return cudf::transform(udf, + cudf::udf_source_type::CUDA, + null_aware ? cudf::null_aware::YES : cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + std::move(output_offsets), + row_count); +} + +struct ListOperationTest : cudf::test::BaseFixture {}; + +TEST_F(ListOperationTest, DictionaryListKeys) +{ + auto keys = make_lists({{}, {1, -2}, {7}}); + offsets_type indices{1, 0, 2, 1}; + auto dictionary = cudf::make_dictionary_column(keys->view(), indices); + auto sliced = cudf::slice(dictionary->view(), {1, 4}).front(); + std::array inputs{sliced}; + std::array output{list_output()}; + auto expected = make_lists({{}, {7}, {1, -2}}); + auto result = transform_lists(inputs, output, list_offsets({0, 0, 1, 3})); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected->view()); + + auto scalar = cudf::slice(dictionary->view(), {0, 1}).front(); + inputs[0] = cudf::scalar_column_view{scalar}; + expected = make_lists({{1, -2}, {1, -2}, {1, -2}}); + result = transform_lists(inputs, output, list_offsets({0, 2, 4, 6}), false, 3); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected->view()); +} + +TEST_F(ListOperationTest, NullAwareDictionaryListKeys) +{ + auto keys = make_lists({{}, {1, -2}, {7}}); + offsets_type indices{{1, 0, 2, 1}, {true, true, true, false}}; + auto dictionary = cudf::make_dictionary_column(keys->view(), indices); + std::array inputs{dictionary->view()}; + std::array output{list_output()}; + auto result = transform_lists(inputs, output, list_offsets({0, 2, 2, 3, 5}), true); + std::array validity{true, true, false, false}; + auto expected = make_lists({{1, -2}, {}, {}, {}}, validity); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected->view()); + EXPECT_FALSE(cudf::has_nonempty_nulls(result->get_column(0).view())); +} + +TEST_F(ListOperationTest, DecimalChild) +{ + auto input = cudf::make_lists_column( + 3, + offsets_type{0, 2, 2, 3}.release(), + cudf::test::fixed_point_column_wrapper<__int128_t>{{-2, 3, 11}, numeric::scale_type{-3}} + .release(), + 0, + cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED)); + std::array inputs{input->view()}; + std::array outputs{list_output(input->view().child(1).type())}; + auto result = cudf::transform(R"( +__device__ void transform(cudf::mutable_list_element* output, + cudf::list_element input) +{ + for (size_t i = 0; i < input.size(); ++i) { + output->assign(i, input.element(i)); + } +} +)", + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + list_offsets({0, 2, 2, 3}), + std::nullopt); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), input->view()); +} + +TEST_F(ListOperationTest, StringChild) +{ + cudf::test::lists_column_wrapper input{ + {"skip"}, {"hello", "", "é"}, {}, {"world"}}; + std::array inputs{cudf::slice(input, {1, 4}).front()}; + std::array outputs{cudf::transform_output{cudf::data_type{cudf::type_id::INT32}}}; + EXPECT_THROW(cudf::transform(R"( +__device__ void transform(int32_t* output, cudf::list_element input) +{ + *output = 0; + for (cudf::size_type i = 0; i < input.size(); ++i) { + *output += input.element(i).size_bytes(); + } +} +)", + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + {}, + std::nullopt), + std::invalid_argument); +} + +TEST_F(ListOperationTest, SizeScanThenEncode) +{ + auto input = make_lists({{-3, 2}, {}, {9}, {4, 5, 6}}); + std::array inputs{input->view()}; + std::array size_output{cudf::transform_output{cudf::data_type{cudf::type_id::INT32}}}; + auto sizes = cudf::transform(R"( +__device__ void transform(int32_t* size, cudf::list_element input) { *size = input.size(); } +)", + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + size_output, + {}, + std::nullopt); + auto prefix = cudf::scan(sizes->get_column(0).view(), + *cudf::make_sum_aggregation(), + cudf::scan_type::INCLUSIVE); + offsets_type zero{0}; + auto scanned = cudf::concatenate(std::vector{zero, prefix->view()}); + offsets_type expected_offsets{0, 2, 2, 3, 6}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(scanned->view(), expected_offsets); + std::array output{list_output()}; + std::vector> supplied; + supplied.push_back(std::move(scanned)); + auto result = transform_lists(inputs, output, std::move(supplied)); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), input->view()); + + auto offsets = list_offsets({0, 2, 2, 3, 6}); + cudf::transform_program program{copy_udf, + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + output, + offsets}; + result = program.run(inputs, output, std::move(offsets), std::nullopt); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), input->view()); + + auto next_input = make_lists({{10}, {11, 12}, {}, {-5}}); + inputs[0] = next_input->view(); + result = program.run(inputs, output, list_offsets({0, 1, 3, 3, 4}), std::nullopt); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), next_input->view()); +} + +TEST_F(ListOperationTest, SlicedAndScalarInputs) +{ + auto input = make_lists({{-10}, {2, -3}, {}, {5}, {99}}); + auto sliced = cudf::slice(input->view(), {1, 4}).front(); + std::array inputs{sliced}; + std::array output{list_output()}; + + auto result = transform_lists(inputs, output, list_offsets({0, 2, 2, 3})); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), sliced); + + auto one_row = cudf::slice(input->view(), {1, 2}).front(); + inputs[0] = cudf::scalar_column_view{one_row}; + auto expected = make_lists({{2, -3}, {2, -3}, {2, -3}}); + + result = transform_lists(inputs, output, list_offsets({0, 2, 4, 6}), false, 3); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected->view()); +} + +TEST_F(ListOperationTest, NullAwareAndPreservedParentValidity) +{ + std::array validity{true, true, true, false}; + auto input = make_lists({{1, -2}, {}, {7}, {}}, validity); + std::array inputs{input->view()}; + std::array output{list_output()}; + std::array expected_validity{true, true, false, false}; + auto expected = make_lists({{1, -2}, {}, {}, {}}, expected_validity); + + auto preserved = transform_lists(inputs, output, list_offsets({0, 2, 2, 3, 5})); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(preserved->get_column(0), input->view()); + EXPECT_FALSE(cudf::has_nonempty_nulls(preserved->get_column(0).view())); + auto aware = transform_lists(inputs, output, list_offsets({0, 2, 2, 3, 5}), true); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(aware->get_column(0), expected->view()); + EXPECT_FALSE(cudf::has_nonempty_nulls(aware->get_column(0).view())); + offsets_type sanitized_offsets{0, 2, 2, 2, 2}; + offsets_type initialized_child{1, -2}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(aware->get_column(0).child(0), sanitized_offsets); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(aware->get_column(0).child(1), initialized_child); + EXPECT_FALSE(aware->get_column(0).child(1).nullable()); +} + +TEST_F(ListOperationTest, EmptyAndAllNull) +{ + auto empty = make_lists({}); + std::array inputs{empty->view()}; + std::array output{list_output()}; + + auto result = transform_lists(inputs, output, list_offsets({0})); + EXPECT_EQ(result->num_rows(), 0); + EXPECT_EQ(result->get_column(0).child(1).type(), cudf::data_type{cudf::type_id::INT32}); + + std::array validity{false, false}; + auto nulls = make_lists({{}, {}}, validity); + inputs[0] = nulls->view(); + + result = transform_lists(inputs, output, list_offsets({0, 0, 0})); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), nulls->view()); + + for (bool null_aware : {false, true}) { + result = transform_lists(inputs, output, list_offsets({0, 2, 4}), null_aware); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), nulls->view()); + EXPECT_EQ(result->get_column(0).child(1).size(), 0); + EXPECT_FALSE(cudf::has_nonempty_nulls(result->get_column(0).view())); + } +} + +TEST_F(ListOperationTest, MixedListStringAndFixedWidthOutputs) +{ + auto input = make_lists({{1, 2}, {}, {-3}}); + std::array inputs{input->view()}; + std::array outputs{list_output(), + cudf::transform_output{cudf::data_type{cudf::type_id::STRING}}, + cudf::transform_output{cudf::data_type{cudf::type_id::INT32}}, + list_output()}; + auto supplied = list_offsets({0, 2, 2, 3}); + supplied.push_back(offsets_type{0, 1, 2, 3}.release()); + supplied.push_back(nullptr); + supplied.push_back(offsets_type{0, 2, 2, 3}.release()); + auto result = cudf::transform(R"( +__device__ void transform(cudf::mutable_list_element* first, cudf::mutable_string_view* text, + int32_t* size, cudf::mutable_list_element* second, + cudf::list_element input) +{ + *size = input.size(); + text->data()[0] = 'a' + input.size(); + for (size_t index = 0; index < input.size(); ++index) { + first->assign(index, input.element(index)); + second->assign(index, -input.element(index)); + } +} +)", + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + outputs, + std::move(supplied), + std::nullopt); + auto negative = make_lists({{-1, -2}, {}, {3}}); + cudf::test::strings_column_wrapper text{"c", "a", "b"}; + offsets_type sizes{2, 0, 1}; + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), input->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(1), text); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(2), sizes); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(3), negative->view()); +} + +TEST_F(ListOperationTest, NullableChildren) +{ + auto input = cudf::make_lists_column(1, + offsets_type{0, 2}.release(), + offsets_type{{1, 2}, {true, false}}.release(), + 0, + cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED)); + std::array inputs{input->view()}; + std::array output{list_output()}; + auto expected = make_lists({{1, -1}}); + offsets_type indices{0}; + auto dictionary = cudf::make_dictionary_column(input->view(), indices); + for (auto const& view : {input->view(), dictionary->view()}) { + inputs[0] = view; + auto result = cudf::transform(R"( +__device__ void transform(cudf::mutable_list_element* output, cudf::list_element input) +{ + for (size_t index = 0; index < input.size(); ++index) { + output->assign(index, input.nullable_element(index).value_or(-1)); + } +} +)", + cudf::udf_source_type::CUDA, + cudf::null_aware::NO, + std::nullopt, + inputs, + output, + list_offsets({0, 2}), + std::nullopt); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0), expected->view()); + } +} + +TEST_F(ListOperationTest, RejectsUnsupportedListChildren) +{ + std::array inputs; + std::array output{list_output()}; + auto child = make_lists({{1, 2}}); + auto nested = cudf::make_lists_column(1, + offsets_type{0, 1}.release(), + std::move(child), + 0, + cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED)); + inputs[0] = nested->view(); + EXPECT_THROW(transform_lists(inputs, output, list_offsets({0, 2})), std::invalid_argument); + auto strings = + cudf::make_lists_column(1, + offsets_type{0, 2}.release(), + cudf::test::strings_column_wrapper{{"a", "b"}, {true, false}}.release(), + 0, + cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED)); + inputs[0] = strings->view(); + EXPECT_THROW(transform_lists(inputs, output, list_offsets({0, 2})), std::invalid_argument); + offsets_type indices{0}; + for (auto const& keys : {nested->view(), strings->view()}) { + auto dictionary = cudf::make_dictionary_column(keys, indices); + inputs[0] = dictionary->view(); + EXPECT_THROW(transform_lists(inputs, output, list_offsets({0, 2})), std::invalid_argument); + } +} + } // namespace transformation diff --git a/python/pylibcudf/pylibcudf/libcudf/transform.pxd b/python/pylibcudf/pylibcudf/libcudf/transform.pxd index 8b7ec4a25242..5690cb2a6926 100644 --- a/python/pylibcudf/pylibcudf/libcudf/transform.pxd +++ b/python/pylibcudf/pylibcudf/libcudf/transform.pxd @@ -46,6 +46,7 @@ cdef extern from "cudf/transform.hpp" namespace "cudf" nogil: cdef cppclass transform_output: data_type type output_nullability nullability + vector[transform_output] children ctypedef const transform_output const_transform_output @@ -76,7 +77,7 @@ cdef extern from "cudf/transform.hpp" namespace "cudf" nogil: optional[void*] user_data, span[const_transform_input] inputs, span[const_transform_output] outputs, - vector[unique_ptr[column]]&& string_offsets, + vector[unique_ptr[column]]&& output_offsets, optional[size_type] row_size, cudaStream_t stream, device_async_resource_ref mr