Skip to content

Commit

Permalink
Fix unsanitized nulls from strings segmented-reduce (#14586)
Browse files Browse the repository at this point in the history
Fixes the string specialization logic in `cudf::segmented_reduce` to not produce unsanitized null entries. The functor used to build a gather map for argmin/argmax was corrected to handle include/exclude nulls correctly.

Reference: #14559

Authors:
  - David Wendt (https://github.com/davidwendt)

Approvers:
  - Nghia Truong (https://github.com/ttnghia)
  - Bradley Dice (https://github.com/bdice)

URL: #14586
  • Loading branch information
davidwendt authored Dec 8, 2023
1 parent a8f071d commit 8b23436
Showing 1 changed file with 19 additions and 4 deletions.
23 changes: 19 additions & 4 deletions cpp/src/reductions/segmented/simple.cuh
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

0 comments on commit 8b23436

Please sign in to comment.