Skip to content

Commit

Permalink
use more efficient make_strings_children utility
Browse files Browse the repository at this point in the history
  • Loading branch information
davidwendt committed Mar 15, 2021
1 parent 2458329 commit b60ab0c
Showing 1 changed file with 61 additions and 75 deletions.
136 changes: 61 additions & 75 deletions cpp/src/strings/substring.cu
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
/*
* Copyright (c) 2019-2020, NVIDIA CORPORATION.
* Copyright (c) 2019-2021, NVIDIA CORPORATION.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
Expand Down Expand Up @@ -43,17 +43,25 @@ namespace {
* using the provided start, stop, and step parameters.
*/
struct substring_fn {
const column_device_view d_column;
numeric_scalar_device_view<size_type> d_start, d_stop, d_step;
const int32_t* d_offsets{};
column_device_view const d_column;
numeric_scalar_device_view<size_type> const d_start;
numeric_scalar_device_view<size_type> const d_stop;
numeric_scalar_device_view<size_type> const d_step;
int32_t* d_offsets{};
char* d_chars{};

__device__ cudf::size_type operator()(size_type idx)
__device__ void operator()(size_type idx)
{
if (d_column.is_null(idx)) return 0; // null string
string_view d_str = d_column.template element<string_view>(idx);
if (d_column.is_null(idx)) {
if (!d_chars) d_offsets[idx] = 0;
return;
}
auto const d_str = d_column.template element<string_view>(idx);
auto const length = d_str.length();
if (length == 0) return 0; // empty string
if (length == 0) {
if (!d_chars) d_offsets[idx] = 0;
return;
}
size_type const step = d_step.is_valid() ? d_step.value() : 1;
auto const begin = [&] { // always inclusive
// when invalid, default depends on step
Expand Down Expand Up @@ -88,7 +96,7 @@ struct substring_fn {
if (d_buffer) d_buffer += from_char_utf8(*itr, d_buffer);
itr += step;
}
return bytes;
if (!d_chars) d_offsets[idx] = bytes;
}
};

Expand All @@ -103,42 +111,26 @@ std::unique_ptr<column> slice_strings(
rmm::cuda_stream_view stream = rmm::cuda_stream_default,
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
{
size_type strings_count = strings.size();
if (strings_count == 0) return make_empty_strings_column(stream, mr);
if (strings.is_empty()) return make_empty_strings_column(stream, mr);

if (step.is_valid()) CUDF_EXPECTS(step.value(stream) != 0, "Step parameter must not be 0");

auto strings_column = column_device_view::create(strings.parent(), stream);
auto d_column = *strings_column;
auto d_start = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(start));
auto d_stop = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(stop));
auto d_step = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(step));

// copy the null mask
rmm::device_buffer null_mask = cudf::detail::copy_bitmask(strings.parent(), stream, mr);

// build offsets column
auto offsets_transformer_itr = thrust::make_transform_iterator(
thrust::make_counting_iterator<int32_t>(0), substring_fn{d_column, d_start, d_stop, d_step});
auto offsets_column = make_offsets_child_column(
offsets_transformer_itr, offsets_transformer_itr + strings_count, stream, mr);
auto d_new_offsets = offsets_column->view().data<int32_t>();

// build chars column
auto bytes = cudf::detail::get_value<int32_t>(offsets_column->view(), strings_count, stream);
auto chars_column = strings::detail::create_chars_child_column(
strings_count, strings.null_count(), bytes, stream, mr);
auto d_chars = chars_column->mutable_view().data<char>();
thrust::for_each_n(rmm::exec_policy(stream),
thrust::make_counting_iterator<size_type>(0),
strings_count,
substring_fn{d_column, d_start, d_stop, d_step, d_new_offsets, d_chars});
auto const d_column = column_device_view::create(strings.parent(), stream);
auto const d_start = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(start));
auto const d_stop = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(stop));
auto const d_step = get_scalar_device_view(const_cast<numeric_scalar<size_type>&>(step));

return make_strings_column(strings_count,
std::move(offsets_column),
std::move(chars_column),
auto children = make_strings_children(substring_fn{*d_column, d_start, d_stop, d_step},
strings.size(),
strings.null_count(),
stream,
mr);

return make_strings_column(strings.size(),
std::move(children.first),
std::move(children.second),
strings.null_count(),
std::move(null_mask),
cudf::detail::copy_bitmask(strings.parent(), stream, mr),
stream,
mr);
}
Expand Down Expand Up @@ -166,25 +158,33 @@ namespace {
* This both calculates the output size and executes the substring.
*/
struct substring_from_fn {
const column_device_view d_column;
const cudf::detail::input_indexalator starts;
const cudf::detail::input_indexalator stops;
const int32_t* d_offsets{};
column_device_view const d_column;
cudf::detail::input_indexalator const starts;
cudf::detail::input_indexalator const stops;
int32_t* d_offsets{};
char* d_chars{};

__device__ size_type operator()(size_type idx)
__device__ void operator()(size_type idx)
{
if (d_column.is_null(idx)) return 0; // null string
string_view d_str = d_column.template element<string_view>(idx);
if (d_column.is_null(idx)) {
if (!d_chars) d_offsets[idx] = 0;
return;
}
auto const d_str = d_column.template element<string_view>(idx);
auto const length = d_str.length();
auto const start = starts[idx];
if (start >= length) return 0; // empty string
if (start >= length) {
if (!d_chars) d_offsets[idx] = 0;
return;
}
auto const stop = stops[idx];
auto const end = (((stop < 0) || (stop > length)) ? length : stop);

string_view d_substr = d_str.substr(start, end - start);
if (d_chars) memcpy(d_chars + d_offsets[idx], d_substr.data(), d_substr.size_bytes());
return d_substr.size_bytes();
auto const d_substr = d_str.substr(start, end - start);
if (d_chars)
memcpy(d_chars + d_offsets[idx], d_substr.data(), d_substr.size_bytes());
else
d_offsets[idx] = d_substr.size_bytes();
}
};

Expand Down Expand Up @@ -212,32 +212,18 @@ std::unique_ptr<column> compute_substrings_from_fn(column_device_view const& d_c
auto strings_count = d_column.size();

// Copy the null mask
rmm::device_buffer null_mask{0, stream, mr};
if (d_column.nullable())
null_mask = rmm::device_buffer(
d_column.null_mask(), cudf::bitmask_allocation_size_bytes(strings_count), stream, mr);

// Build offsets column
auto offsets_transformer_itr = thrust::make_transform_iterator(
thrust::make_counting_iterator<size_type>(0), substring_from_fn{d_column, starts, stops});
auto offsets_column = cudf::strings::detail::make_offsets_child_column(
offsets_transformer_itr, offsets_transformer_itr + strings_count, stream, mr);
auto d_new_offsets = offsets_column->view().data<int32_t>();

// Build chars column
auto bytes = cudf::detail::get_value<int32_t>(offsets_column->view(), strings_count, stream);
auto chars_column =
cudf::strings::detail::create_chars_child_column(strings_count, null_count, bytes, stream, mr);
auto chars_view = chars_column->mutable_view();
auto d_chars = chars_view.template data<char>();
thrust::for_each_n(rmm::exec_policy(stream),
thrust::make_counting_iterator<cudf::size_type>(0),
strings_count,
substring_from_fn{d_column, starts, stops, d_new_offsets, d_chars});
rmm::device_buffer null_mask =
!d_column.nullable()
? rmm::device_buffer{0, stream, mr}
: rmm::device_buffer(
d_column.null_mask(), cudf::bitmask_allocation_size_bytes(strings_count), stream, mr);

auto children = make_strings_children(
substring_from_fn{d_column, starts, stops}, strings_count, null_count, stream, mr);

return make_strings_column(strings_count,
std::move(offsets_column),
std::move(chars_column),
std::move(children.first),
std::move(children.second),
null_count,
std::move(null_mask),
stream,
Expand Down

0 comments on commit b60ab0c

Please sign in to comment.