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.

Sign up for GitHub

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

Fix unsanitized nulls from strings segmented-reduce #14586

Merged
merged 3 commits into from
Dec 8, 2023
+19 −4
Merged

Fix unsanitized nulls from strings segmented-reduce #14586

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
74bbd26
Fix unsanitized nulls from strings segmented-reduce
davidwendt Dec 6, 2023
a7260cc
Merge branch 'branch-24.02' into fix-segred-unsanitized
davidwendt Dec 8, 2023
21e80dc
move comments
davidwendt Dec 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 19 additions & 4 deletions cpp/src/reductions/segmented/simple.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@
#include <cudf/detail/gather.hpp>
#include <cudf/detail/unary.hpp>
#include <cudf/detail/utilities/cuda.cuh>
#include <cudf/detail/utilities/element_argminmax.cuh>
#include <cudf/detail/valid_if.cuh>
#include <cudf/reduction/detail/segmented_reduction.cuh>
#include <cudf/types.hpp>
Expand Down Expand Up @@ -114,6 +113,24 @@ std::unique_ptr<column> simple_segmented_reduction(
return result;
}

template <typename T>
struct reduce_argminmax_fn {
column_device_view const d_col; // column data
bool const arg_min; // true if argmin, otherwise argmax
null_policy null_handler; // include or exclude nulls

__device__ inline auto operator()(size_type const& lhs_idx, size_type const& rhs_idx) const
{
// CUB segmented reduce calls with OOB indices
if (lhs_idx < 0 || lhs_idx >= d_col.size()) { return rhs_idx; }
if (rhs_idx < 0 || rhs_idx >= d_col.size()) { return lhs_idx; }
if (d_col.is_null(lhs_idx)) { return null_handler == null_policy::INCLUDE ? lhs_idx : rhs_idx; }
if (d_col.is_null(rhs_idx)) { return null_handler == null_policy::INCLUDE ? rhs_idx : lhs_idx; }
auto const less = d_col.element<T>(lhs_idx) < d_col.element<T>(rhs_idx);
return less == arg_min ? lhs_idx : rhs_idx;
}
};

/**
* @brief String segmented reduction for 'min', 'max'.
*
Expand All @@ -130,7 +147,6 @@ std::unique_ptr<column> simple_segmented_reduction(
* @param mr Device memory resource used to allocate the returned column's device memory
* @return Output column in device memory
*/

template <typename InputType,
typename Op,
CUDF_ENABLE_IF(std::is_same_v<Op, cudf::reduction::detail::op::min> ||
Expand All @@ -148,8 +164,7 @@ std::unique_ptr<column> string_segmented_reduction(column_view const& col,
auto const num_segments = static_cast<size_type>(offsets.size()) - 1;

bool constexpr is_argmin = std::is_same_v<Op, cudf::reduction::detail::op::min>;
auto string_comparator =
cudf::detail::element_argminmax_fn<InputType>{*device_col, col.has_nulls(), is_argmin};
auto string_comparator = reduce_argminmax_fn<InputType>{*device_col, is_argmin, null_handling};
auto constexpr identity =
is_argmin ? cudf::detail::ARGMIN_SENTINEL : cudf::detail::ARGMAX_SENTINEL;

Expand Down