Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Rework cudf::strings::pad and zfill to use make_strings_children #12238

Merged
204 changes: 97 additions & 107 deletions cpp/src/strings/padding.cu
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,6 @@

#include <cudf/column/column.hpp>
#include <cudf/column/column_device_view.cuh>
#include <cudf/column/column_factories.hpp>
#include <cudf/detail/get_value.cuh>
#include <cudf/detail/iterator.cuh>
#include <cudf/detail/null_mask.hpp>
#include <cudf/detail/nvtx/ranges.hpp>
#include <cudf/strings/detail/pad_impl.cuh>
Expand All @@ -29,142 +26,135 @@
#include <cudf/utilities/default_stream.hpp>

#include <rmm/cuda_stream_view.hpp>
#include <rmm/exec_policy.hpp>

#include <thrust/for_each.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/iterator/transform_iterator.h>

namespace cudf {
namespace strings {
namespace detail {
namespace {

struct compute_pad_output_length_fn {
column_device_view d_strings;
size_type width;
size_type fill_char_size;
/**
* @brief Base class for pad_fn and zfill_fn functors
*
* This handles the output size calculation while delegating the
* pad operation to Derived.
*
* @tparam Derived class uses the CRTP pattern to reuse code logic
* and must include a `pad(string_view,char*)` member function.
*/
template <typename Derived>
struct base_fn {
column_device_view const d_column;
size_type const width;
size_type const fill_char_size;
offset_type* d_offsets{};
char* d_chars{};

base_fn(column_device_view const& d_column, size_type width, size_type fill_char_size)
: d_column(d_column), width(width), fill_char_size(fill_char_size)
{
}

__device__ size_type operator()(size_type idx)
__device__ void operator()(size_type idx) const
{
if (d_column.is_null(idx)) {
if (!d_chars) d_offsets[idx] = 0;
}
davidwendt marked this conversation as resolved.
Show resolved Hide resolved

auto const d_str = d_column.element<string_view>(idx);
auto const& derived = static_cast<Derived const&>(*this);
if (d_chars) {
derived.pad(d_str, d_chars + d_offsets[idx]);
} else {
d_offsets[idx] = compute_padded_size(d_str, width, fill_char_size);
}
};
};

/**
* @brief Pads each string to specified width
*
* @tparam side Side of the string to pad
*/
template <side_type side>
struct pad_fn : base_fn<pad_fn<side>> {
using Base = base_fn<pad_fn<side>>;

cudf::char_utf8 const d_fill_char;

pad_fn(column_device_view const& d_column,
size_type width,
size_type fill_char_size,
char_utf8 fill_char)
: Base(d_column, width, fill_char_size), d_fill_char(fill_char)
{
if (d_strings.is_null(idx)) return 0;
string_view d_str = d_strings.element<string_view>(idx);
return compute_padded_size(d_str, width, fill_char_size);
}

__device__ void pad(string_view d_str, char* output) const
{
pad_impl<side>(d_str, Base::width, d_fill_char, output);
}
};

} // namespace

std::unique_ptr<column> pad(strings_column_view const& strings,
std::unique_ptr<column> pad(strings_column_view const& input,
size_type width,
side_type side,
std::string_view fill_char,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
{
size_type strings_count = strings.size();
if (strings_count == 0) return make_empty_column(type_id::STRING);
if (input.is_empty()) return make_empty_column(type_id::STRING);
CUDF_EXPECTS(!fill_char.empty(), "fill_char parameter must not be empty");
char_utf8 d_fill_char = 0;
size_type fill_char_size = to_char_utf8(fill_char.data(), d_fill_char);

auto strings_column = column_device_view::create(strings.parent(), stream);
auto d_strings = *strings_column;

// create null_mask
rmm::device_buffer null_mask = cudf::detail::copy_bitmask(strings.parent(), stream, mr);

// build offsets column
auto offsets_transformer_itr = cudf::detail::make_counting_transform_iterator(
0, compute_pad_output_length_fn{d_strings, width, fill_char_size});
auto offsets_column = make_offsets_child_column(
offsets_transformer_itr, offsets_transformer_itr + strings_count, stream, mr);
auto d_offsets = offsets_column->view().data<int32_t>();

// build chars column
auto const bytes =
cudf::detail::get_value<int32_t>(offsets_column->view(), strings_count, stream);
auto chars_column = strings::detail::create_chars_child_column(bytes, stream, mr);
auto d_chars = chars_column->mutable_view().data<char>();

if (side == side_type::LEFT) {
thrust::for_each_n(
rmm::exec_policy(stream),
thrust::make_counting_iterator<cudf::size_type>(0),
strings_count,
[d_strings, width, d_fill_char, d_offsets, d_chars] __device__(size_type idx) {
if (d_strings.is_valid(idx)) {
pad_impl<side_type::LEFT>(
d_strings.element<string_view>(idx), width, d_fill_char, d_chars + d_offsets[idx]);
}
});
} else if (side == side_type::RIGHT) {
thrust::for_each_n(
rmm::exec_policy(stream),
thrust::make_counting_iterator<cudf::size_type>(0),
strings_count,
[d_strings, width, d_fill_char, d_offsets, d_chars] __device__(size_type idx) {
if (d_strings.is_valid(idx)) {
pad_impl<side_type::RIGHT>(
d_strings.element<string_view>(idx), width, d_fill_char, d_chars + d_offsets[idx]);
}
});
} else if (side == side_type::BOTH) {
thrust::for_each_n(
rmm::exec_policy(stream),
thrust::make_counting_iterator<cudf::size_type>(0),
strings_count,
[d_strings, width, d_fill_char, d_offsets, d_chars] __device__(size_type idx) {
if (d_strings.is_valid(idx)) {
pad_impl<side_type::BOTH>(
d_strings.element<string_view>(idx), width, d_fill_char, d_chars + d_offsets[idx]);
}
});
}
auto d_fill_char = char_utf8{0};
auto const fill_char_size = to_char_utf8(fill_char.data(), d_fill_char);

auto d_strings = column_device_view::create(input.parent(), stream);

auto children = [&] {
if (side == side_type::LEFT) {
auto fn = pad_fn<side_type::LEFT>{*d_strings, width, fill_char_size, d_fill_char};
return make_strings_children(fn, input.size(), stream, mr);
} else if (side == side_type::RIGHT) {
auto fn = pad_fn<side_type::RIGHT>{*d_strings, width, fill_char_size, d_fill_char};
return make_strings_children(fn, input.size(), stream, mr);
}
auto fn = pad_fn<side_type::BOTH>{*d_strings, width, fill_char_size, d_fill_char};
return make_strings_children(fn, input.size(), stream, mr);
davidwendt marked this conversation as resolved.
Show resolved Hide resolved
}();

return make_strings_column(strings_count,
std::move(offsets_column),
std::move(chars_column),
strings.null_count(),
std::move(null_mask));
return make_strings_column(input.size(),
std::move(children.first),
std::move(children.second),
input.null_count(),
cudf::detail::copy_bitmask(input.parent(), stream, mr));
}

namespace {

/**
* @brief Zero-fill each string to specified width
*/
struct zfill_fn : base_fn<zfill_fn> {
zfill_fn(column_device_view const& d_column, size_type width) : base_fn(d_column, width, 1) {}

__device__ void pad(string_view d_str, char* output) const { zfill_impl(d_str, width, output); }
};
} // namespace

std::unique_ptr<column> zfill(strings_column_view const& input,
size_type width,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr)
{
if (input.is_empty()) return make_empty_column(type_id::STRING);

auto strings_column = column_device_view::create(input.parent(), stream);
auto d_strings = *strings_column;

// build offsets column
auto offsets_transformer_itr = thrust::make_transform_iterator(
thrust::make_counting_iterator<int32_t>(0),
compute_pad_output_length_fn{d_strings, width, 1}); // fillchar is 1 byte
auto offsets_column = make_offsets_child_column(
offsets_transformer_itr, offsets_transformer_itr + input.size(), stream, mr);
auto const d_offsets = offsets_column->view().data<int32_t>();

// build chars column
auto const bytes = cudf::detail::get_value<int32_t>(offsets_column->view(), input.size(), stream);
auto chars_column = strings::detail::create_chars_child_column(bytes, stream, mr);
auto d_chars = chars_column->mutable_view().data<char>();

thrust::for_each_n(rmm::exec_policy(stream),
thrust::make_counting_iterator<cudf::size_type>(0),
input.size(),
[d_strings, width, d_offsets, d_chars] __device__(size_type idx) {
if (d_strings.is_valid(idx)) {
zfill_impl(
d_strings.element<string_view>(idx), width, d_chars + d_offsets[idx]);
}
});
auto d_strings = column_device_view::create(input.parent(), stream);
auto children = make_strings_children(zfill_fn{*d_strings, width}, input.size(), stream, mr);

return make_strings_column(input.size(),
std::move(offsets_column),
std::move(chars_column),
std::move(children.first),
std::move(children.second),
input.null_count(),
cudf::detail::copy_bitmask(input.parent(), stream, mr));
}
Expand Down