Skip to content

Commit

Permalink
Fully support nested types in cudf::contains (#10656)
Browse files Browse the repository at this point in the history 
This extends the `cudf::contains` API to support nested types (lists + structs) with arbitrarily nested levels. As such, `cudf::contains` will work with literally any type of input data.

In addition, this fixes null handling of `cudf::contains` with structs column + struct scalar input when the structs column contains null rows at the top level while the scalar key is valid but all nulls at children levels.

Closes: #8965
Depends on:
 * #10730
 * #10883
 * #10802
 * #10997
 * NVIDIA/cuCollections#172
 * NVIDIA/cuCollections#173
 * #11037
 * #11356

Authors:
  - Nghia Truong (https://github.com/ttnghia)
  - Devavret Makkar (https://github.com/devavret)
  - Bradley Dice (https://github.com/bdice)
  - Karthikeyan (https://github.com/karthikeyann)

Approvers:
  - AJ Schmidt (https://github.com/ajschmidt8)
  - Bradley Dice (https://github.com/bdice)
  - Yunsong Wang (https://github.com/PointKernel)

URL: #10656
  • Loading branch information
@ttnghia
ttnghia authored Aug 17, 2022
1 parent 89fa003 commit 127d574
Show file tree
Hide file tree
Showing 8 changed files with 894 additions and 393 deletions.
2 changes: 1 addition & 1 deletion cpp/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -449,8 +449,8 @@ add_library(
src/scalar/scalar.cpp
src/scalar/scalar_factories.cpp
src/search/contains_column.cu
src/search/contains_scalar.cu
src/search/contains_table.cu
src/search/contains_nested.cu
src/search/search_ordered.cu
src/sort/is_sorted.cu
src/sort/rank.cu
Expand Down
16 changes: 0 additions & 16 deletions cpp/include/cudf/detail/search.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,20 +97,4 @@ rmm::device_uvector<bool> contains(
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());

/**
* @brief Check if the (unique) row of the `needle` column is contained in the `haystack` column.
*
* If the input `needle` column has more than one row, only the first row will be considered.
*
* This function is designed for nested types only. It can also work with non-nested types
* but with lower performance due to the complexity of the implementation.
*
* @param haystack The column containing search space.
* @param needle A scalar value to check for existence in the search space.
* @return true if the given `needle` value exists in the `haystack` column.
*/
bool contains_nested_element(column_view const& haystack,
column_view const& needle,
rmm::cuda_stream_view stream);

} // namespace cudf::detail
202 changes: 47 additions & 155 deletions cpp/src/search/contains_column.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,22 +17,18 @@
#include <hash/unordered_multiset.cuh>

#include <cudf/column/column_factories.hpp>
#include <cudf/detail/null_mask.hpp>
#include <cudf/detail/nvtx/ranges.hpp>
#include <cudf/detail/search.hpp>
#include <cudf/dictionary/detail/search.hpp>
#include <cudf/dictionary/detail/update_keys.hpp>
#include <cudf/lists/list_view.hpp>
#include <cudf/scalar/scalar.hpp>
#include <cudf/scalar/scalar_device_view.cuh>
#include <cudf/structs/struct_view.hpp>
#include <cudf/table/table_view.hpp>
#include <cudf/utilities/default_stream.hpp>

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

#include <thrust/count.h>
#include <thrust/iterator/counting_iterator.h>
#include <thrust/pair.h>
#include <thrust/transform.h>
#include <thrust/uninitialized_fill.h>

Expand All @@ -41,104 +37,33 @@ namespace detail {

namespace {

/**
* @brief Get the underlying value of a scalar through a scalar device view.
*
* @tparam Type The scalar's value type
* @tparam ScalarDView Type of the input scalar device view
* @param d_scalar The input scalar device view
*/
template <typename Type, typename ScalarDView>
__device__ auto inline get_scalar_value(ScalarDView d_scalar)
{
if constexpr (cudf::is_fixed_point<Type>()) {
return d_scalar.rep();
} else {
return d_scalar.value();
}
}

struct contains_scalar_dispatch {
template <typename Type>
bool operator()(column_view const& haystack,
scalar const& needle,
rmm::cuda_stream_view stream) const
{
CUDF_EXPECTS(haystack.type() == needle.type(), "scalar and column types must match");

using DType = device_storage_type_t<Type>;
auto const d_haystack = column_device_view::create(haystack, stream);
auto const d_needle =
get_scalar_device_view(static_cast<cudf::scalar_type_t<Type>&>(const_cast<scalar&>(needle)));

if (haystack.has_nulls()) {
auto const begin = d_haystack->pair_begin<DType, true>();
auto const end = d_haystack->pair_end<DType, true>();

return thrust::count_if(
rmm::exec_policy(stream), begin, end, [d_needle] __device__(auto const val_pair) {
auto const needle_pair = thrust::make_pair(get_scalar_value<Type>(d_needle), true);
return val_pair == needle_pair;
}) > 0;
} else {
auto const begin = d_haystack->begin<DType>();
auto const end = d_haystack->end<DType>();

return thrust::count_if(
rmm::exec_policy(stream), begin, end, [d_needle] __device__(auto const val) {
return val == get_scalar_value<Type>(d_needle);
}) > 0;
struct contains_column_dispatch {
template <typename Element, typename Haystack>
struct contains_fn {
bool __device__ operator()(size_type const idx) const
{
if (needles_have_nulls && needles.is_null_nocheck(idx)) {
// Exit early. The value doesn't matter, and will be masked as a null element.
return true;
}

return haystack.contains(needles.template element<Element>(idx));
}
}
};

template <>
bool contains_scalar_dispatch::operator()<cudf::list_view>(column_view const&,
scalar const&,
rmm::cuda_stream_view) const
{
CUDF_FAIL("list_view type not supported yet");
}
Haystack const haystack;
column_device_view const needles;
bool const needles_have_nulls;
};

template <>
bool contains_scalar_dispatch::operator()<cudf::struct_view>(column_view const& haystack,
scalar const& needle,
rmm::cuda_stream_view stream) const
{
CUDF_EXPECTS(haystack.type() == needle.type(), "scalar and column types must match");
// Haystack and needle structure compatibility will be checked by the table comparator
// constructor during call to `contains_nested_element`.

auto const needle_as_col = make_column_from_scalar(needle, 1, stream);
return contains_nested_element(haystack, needle_as_col->view(), stream);
}

template <>
bool contains_scalar_dispatch::operator()<cudf::dictionary32>(column_view const& haystack,
scalar const& needle,
rmm::cuda_stream_view stream) const
{
auto const dict_col = cudf::dictionary_column_view(haystack);
// first, find the needle in the dictionary's key set
auto const index = cudf::dictionary::detail::get_index(dict_col, needle, stream);
// if found, check the index is actually in the indices column
return index->is_valid(stream) && cudf::type_dispatcher(dict_col.indices().type(),
contains_scalar_dispatch{},
dict_col.indices(),
*index,
stream);
}

struct multi_contains_dispatch {
template <typename Type>
template <typename Element, CUDF_ENABLE_IF(!is_nested<Element>())>
std::unique_ptr<column> operator()(column_view const& haystack,
column_view const& needles,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const
{
auto result = make_numeric_column(data_type{type_to_id<bool>()},
needles.size(),
copy_bitmask(needles),
copy_bitmask(needles, stream, mr),
needles.null_count(),
stream,
mr);
Expand All @@ -151,57 +76,38 @@ struct multi_contains_dispatch {
return result;
}

auto const haystack_set = cudf::detail::unordered_multiset<Type>::create(haystack, stream);
auto const haystack_set = cudf::detail::unordered_multiset<Element>::create(haystack, stream);
auto const haystack_set_dv = haystack_set.to_device();
auto const needles_cdv_ptr = column_device_view::create(needles, stream);
auto const needles_it = thrust::make_counting_iterator<size_type>(0);

if (needles.has_nulls()) {
thrust::transform(rmm::exec_policy(stream),
needles_it,
needles_it + needles.size(),
out_begin,
[haystack = haystack_set.to_device(),
needles = *needles_cdv_ptr] __device__(size_type const idx) {
return needles.is_null_nocheck(idx) ||
haystack.contains(needles.template element<Type>(idx));
});
} else {
thrust::transform(rmm::exec_policy(stream),
needles_it,
needles_it + needles.size(),
out_begin,
[haystack = haystack_set.to_device(),
needles = *needles_cdv_ptr] __device__(size_type const index) {
return haystack.contains(needles.template element<Type>(index));
});
}

thrust::transform(rmm::exec_policy(stream),
thrust::make_counting_iterator<size_type>(0),
thrust::make_counting_iterator<size_type>(needles.size()),
out_begin,
contains_fn<Element, decltype(haystack_set_dv)>{
haystack_set_dv, *needles_cdv_ptr, needles.has_nulls()});
return result;
}
};

template <>
std::unique_ptr<column> multi_contains_dispatch::operator()<list_view>(
column_view const&,
column_view const&,
rmm::cuda_stream_view,
rmm::mr::device_memory_resource*) const
{
CUDF_FAIL("list_view type not supported");
}

template <>
std::unique_ptr<column> multi_contains_dispatch::operator()<struct_view>(
column_view const&,
column_view const&,
rmm::cuda_stream_view,
rmm::mr::device_memory_resource*) const
{
CUDF_FAIL("struct_view type not supported");
}
template <typename Element, CUDF_ENABLE_IF(is_nested<Element>())>
std::unique_ptr<column> operator()(column_view const& haystack,
column_view const& needles,
rmm::cuda_stream_view stream,
rmm::mr::device_memory_resource* mr) const
{
auto result_v = detail::contains(table_view{{haystack}},
table_view{{needles}},
null_equality::EQUAL,
nan_equality::ALL_EQUAL,
stream,
mr);
return std::make_unique<column>(
std::move(result_v), copy_bitmask(needles, stream, mr), needles.null_count());
}
};

template <>
std::unique_ptr<column> multi_contains_dispatch::operator()<dictionary32>(
std::unique_ptr<column> contains_column_dispatch::operator()<dictionary32>(
column_view const& haystack_in,
column_view const& needles_in,
rmm::cuda_stream_view stream,
Expand All @@ -219,22 +125,14 @@ std::unique_ptr<column> multi_contains_dispatch::operator()<dictionary32>(
column_view const haystack_indices = haystack_view.get_indices_annotated();
column_view const needles_indices = needles_view.get_indices_annotated();
return cudf::type_dispatcher(haystack_indices.type(),
multi_contains_dispatch{},
contains_column_dispatch{},
haystack_indices,
needles_indices,
stream,
mr);
}
} // namespace

bool contains(column_view const& haystack, scalar const& needle, rmm::cuda_stream_view stream)
{
if (haystack.is_empty()) { return false; }
if (not needle.is_valid(stream)) { return haystack.has_nulls(); }

return cudf::type_dispatcher(
haystack.type(), contains_scalar_dispatch{}, haystack, needle, stream);
}
} // namespace

std::unique_ptr<column> contains(column_view const& haystack,
column_view const& needles,
Expand All @@ -244,17 +142,11 @@ std::unique_ptr<column> contains(column_view const& haystack,
CUDF_EXPECTS(haystack.type() == needles.type(), "DTYPE mismatch");

return cudf::type_dispatcher(
haystack.type(), multi_contains_dispatch{}, haystack, needles, stream, mr);
haystack.type(), contains_column_dispatch{}, haystack, needles, stream, mr);
}

} // namespace detail

bool contains(column_view const& haystack, scalar const& needle)
{
CUDF_FUNC_RANGE();
return detail::contains(haystack, needle, cudf::default_stream_value);
}

std::unique_ptr<column> contains(column_view const& haystack,
column_view const& needles,
rmm::mr::device_memory_resource* mr)
Expand Down
66 changes: 0 additions & 66 deletions cpp/src/search/contains_nested.cu
View file

This file was deleted.

Loading

0 comments on commit 127d574

Please sign in to comment.