Support contains() on lists of primitives

CHANGELOG.md

@@ -9,6 +9,7 @@
 - PR #6907 Add `replace_null` API with `replace_policy` parameter, `fixed_width` column support
 
 - PR #6775 Implement cudf.DateOffset for months
+- PR #7039 Support contains() on lists of primitives
 
 ## Improvements
 

conda/recipes/libcudf/meta.yaml

@@ -123,6 +123,7 @@ test:
     - test -f $PREFIX/include/cudf/lists/detail/concatenate.hpp
     - test -f $PREFIX/include/cudf/lists/detail/copying.hpp
     - test -f $PREFIX/include/cudf/lists/extract.hpp
+    - test -f $PREFIX/include/cudf/lists/contains.hpp
     - test -f $PREFIX/include/cudf/lists/lists_column_view.hpp
     - test -f $PREFIX/include/cudf/merge.hpp
     - test -f $PREFIX/include/cudf/null_mask.hpp

cpp/include/cudf/lists/contains.hpp

@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cudf/column/column.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+namespace cudf {
+namespace lists {
+
+/**
+ * @brief Create a column of bool values indicating whether the specified scalar
+ * is an element of each row of a list column.
+ *
+ * The output column has as many elements as the input `lists` column.
+ * Output `column[i]` is set to true if the lists row `lists[i]` contains the value
+ * specified in skey. Otherwise, it is set to false.
+ *
+ * Output `column[i]` is set to null if even one of the following holds true:
+ *   1. The search key `skey` is null
+ *   2. The list row `lists[i]` is null
+ *   3. The list row `lists[i]` contains even *one* null
+ *
+ * @param lists Lists column whose `n` rows are to be searched
+ * @param skey The scalar key to be looked up in each list row
+ * @param mr Device memory resource used to allocate the returned column's device memory.
+ * @return std::unique_ptr<column> BOOL8 column of `n` rows with the result of the lookup
+ */
+std::unique_ptr<column> contains(
+  cudf::lists_column_view const& lists,
+  cudf::scalar const& skey,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+std::unique_ptr<column> contains(
+  cudf::lists_column_view const& lists,
+  cudf::column_view const& skeys,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+}  // namespace lists
+}  // namespace cudf

cpp/src/lists/contains.cu

@@ -0,0 +1,160 @@
+/*
+ * Copyright (c) 2020, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <thrust/logical.h>
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/valid_if.cuh>
+#include <cudf/lists/contains.hpp>
+#include <cudf/lists/list_device_view.cuh>
+#include <cudf/lists/lists_column_device_view.cuh>
+#include <cudf/lists/lists_column_view.hpp>
+#include <cudf/scalar/scalar.hpp>
+#include <cudf/scalar/scalar_device_view.cuh>
+#include <cudf/utilities/type_dispatcher.hpp>
+#include <rmm/exec_policy.hpp>
+#include <type_traits>
+
+namespace cudf {
+namespace lists {
+
+namespace {
+
+auto CUDA_HOST_DEVICE_CALLABLE counting_iter(size_type n)
+{
+  return thrust::make_counting_iterator(n);
+}
+
+std::pair<rmm::device_buffer, size_type> construct_null_mask(
+  cudf::detail::lists_column_device_view const& d_lists,
+  cudf::scalar const& skey,
+  bool input_column_has_nulls,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr)
+{
+  using namespace cudf;
+  using namespace cudf::detail;
+
+  auto const skey_is_valid    = skey.is_valid(stream);
+  auto const output_has_nulls = !skey_is_valid || input_column_has_nulls;
+
+  if (!output_has_nulls) { return std::make_pair(rmm::device_buffer{0, stream, mr}, size_type{0}); }
+
+  if (!skey_is_valid) {
+    return std::make_pair(cudf::create_null_mask(d_lists.size(), mask_state::ALL_NULL, mr),
+                          d_lists.size());
+  }
+
+  return cudf::detail::valid_if(
+    counting_iter(0), counting_iter(d_lists.size()), [d_lists] __device__(auto const& row_index) {
+      auto list = cudf::list_device_view(d_lists, row_index);
+      if (list.is_null()) { return false; }
+      return thrust::none_of(thrust::seq,
+                             counting_iter(0),
+                             counting_iter(list.size()),
+                             [&list] __device__(auto const& i) { return list.is_null(i); });
+    });
+}
+
+struct lookup_functor {
+  template <typename T, typename... Args>
+  std::enable_if_t<!cudf::is_numeric<T>() && !std::is_same<T, cudf::string_view>::value, void>
+  operator()(Args&&...) const
+  {
+    CUDF_FAIL("lists::contains() is only supported on numeric types and strings.");
+  }
+
+  template <typename T>
+  std::enable_if_t<cudf::is_numeric<T>() || std::is_same<T, cudf::string_view>::value, void>
+  operator()(cudf::detail::lists_column_device_view const& d_lists,
+             cudf::scalar const& skey,
+             cudf::mutable_column_device_view output_bools,
+             rmm::cuda_stream_view stream,
+             rmm::mr::device_memory_resource* mr) const
+  {
+    assert(skey.is_valid() && "skey should have been checked for nulls by this point.");
+
+    auto h_scalar = static_cast<cudf::scalar_type_t<T> const&>(skey);
+    auto d_scalar = cudf::get_scalar_device_view(h_scalar);
+
+    thrust::transform(rmm::exec_policy(stream),
+                      counting_iter(0),
+                      counting_iter(d_lists.size()),
+                      output_bools.data<bool>(),
+                      [d_lists, d_scalar] __device__(auto row_index) {
+                        auto list = cudf::list_device_view(d_lists, row_index);
+                        if (list.is_null()) { return false; }
+                        for (size_type i{0}; i < list.size(); ++i) {
+                          if (list.is_null(i)) { return false; }
+                          auto list_element = list.template element<T>(i);
+                          if (list_element == d_scalar.value()) { return true; }
+                        }
+                        return false;
+                      });
+  }
+};
+
+}  // namespace
+
+namespace detail {
+
+std::unique_ptr<column> contains(cudf::lists_column_view const& lists,
+                                 cudf::scalar const& skey,
+                                 rmm::cuda_stream_view stream,
+                                 rmm::mr::device_memory_resource* mr)
+{
+  using namespace cudf;
+  using namespace cudf::detail;
+
+  CUDF_EXPECTS(!cudf::is_nested(lists.child().type()),
+               "Nested types not supported in lists::contains()");
+  CUDF_EXPECTS(lists.child().type().id() == skey.type().id(),
+               "Type of search key does not match list column element type.");
+  CUDF_EXPECTS(skey.type().id() != type_id::EMPTY, "Type cannot be empty.");
+
+  auto const device_view = column_device_view::create(lists.parent(), stream);
+  auto const d_lists     = lists_column_device_view(*device_view);
+
+  rmm::device_buffer null_mask;
+  size_type num_nulls;
+
+  std::tie(null_mask, num_nulls) =
+    construct_null_mask(d_lists, skey, lists.has_nulls(), stream, mr);
+
+  auto ret_bools = make_fixed_width_column(
+    data_type{type_id::BOOL8}, lists.size(), std::move(null_mask), num_nulls, stream, mr);
+
+  auto ret_bools_mutable_device_view =
+    mutable_column_device_view::create(ret_bools->mutable_view(), stream);
+
+  if (skey.is_valid()) {
+    cudf::type_dispatcher(
+      skey.type(), lookup_functor{}, d_lists, skey, *ret_bools_mutable_device_view, stream, mr);
+  }
+
+  return ret_bools;
+}
+
+}  // namespace detail
+
+std::unique_ptr<column> contains(cudf::lists_column_view const& lists,
+                                 cudf::scalar const& skey,
+                                 rmm::mr::device_memory_resource* mr)
+{
+  return detail::contains(lists, skey, rmm::cuda_stream_default, mr);
+}
+
+}  // namespace lists
+}  // namespace cudf

cpp/tests/CMakeLists.txt

@@ -652,7 +652,8 @@ ConfigureTest(AST_TEST "${AST_TEST_SRC}")
 # - lists tests ----------------------------------------------------------------------------------
 
 set(LISTS_TEST_SRC
-    "${CMAKE_CURRENT_SOURCE_DIR}/lists/extract_tests.cpp")
+    "${CMAKE_CURRENT_SOURCE_DIR}/lists/extract_tests.cpp"
+    "${CMAKE_CURRENT_SOURCE_DIR}/lists/contains_tests.cpp")
 
 ConfigureTest(LISTS_TEST "${LISTS_TEST_SRC}")