Optimize compaction operations

cpp/benchmarks/CMakeLists.txt

@@ -123,7 +123,7 @@ ConfigureBench(APPLY_BOOLEAN_MASK_BENCH stream_compaction/apply_boolean_mask_ben
 
 # ##################################################################################################
 # * stream_compaction benchmark -------------------------------------------------------------------
-ConfigureBench(STREAM_COMPACTION_BENCH stream_compaction/drop_duplicates_benchmark.cpp)
+ConfigureNVBench(STREAM_COMPACTION_NVBENCH stream_compaction/drop_duplicates_nvbench.cpp)
 
 # ##################################################################################################
 # * join benchmark --------------------------------------------------------------------------------

cpp/benchmarks/stream_compaction/drop_duplicates_nvbench.cpp

@@ -0,0 +1,135 @@
+/*
+ * Copyright (c) 2022, 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 <cudf/column/column_view.hpp>
+#include <cudf/detail/stream_compaction.hpp>
+#include <cudf/types.hpp>
+#include <cudf_test/base_fixture.hpp>
+#include <cudf_test/column_wrapper.hpp>
+
+#include <fixture/rmm_pool_raii.hpp>
+
+#include <nvbench/nvbench.cuh>
+
+#include <memory>
+#include <random>
+
+enum class algorithm { SORT_BASED, HASH_BASED };
+
+// mandatory for enum types
+NVBENCH_DECLARE_ENUM_TYPE_STRINGS(
+  // Enum type:
+  algorithm,
+  // Callable to generate input strings:
+  // Short identifier used for tables, command-line args, etc.
+  // Used when context is available to figure out the enum type.
+  [](algorithm algo) {
+    switch (algo) {
+      case algorithm::SORT_BASED: return "SORT_BASED";
+      case algorithm::HASH_BASED: return "HASH_BASED";
+      default: return "ERROR";
+    }
+  },
+  // Callable to generate descriptions:
+  // If non-empty, these are used in `--list` to describe values.
+  // Used when context may not be available to figure out the type from the
+  // input string.
+  // Just use `[](auto) { return std::string{}; }` if you don't want these.
+  [](auto) { return std::string{}; })
+
+NVBENCH_DECLARE_ENUM_TYPE_STRINGS(
+  // Enum type:
+  cudf::duplicate_keep_option,
+  // Callable to generate input strings:
+  // Short identifier used for tables, command-line args, etc.
+  // Used when context is available to figure out the enum type.
+  [](cudf::duplicate_keep_option option) {
+    switch (option) {
+      case cudf::duplicate_keep_option::KEEP_FIRST: return "KEEP_FIRST";
+      case cudf::duplicate_keep_option::KEEP_LAST: return "KEEP_LAST";
+      case cudf::duplicate_keep_option::KEEP_NONE: return "KEEP_NONE";
+      default: return "ERROR";
+    }
+  },
+  // Callable to generate descriptions:
+  // If non-empty, these are used in `--list` to describe values.
+  // Used when context may not be available to figure out the type from the
+  // input string.
+  // Just use `[](auto) { return std::string{}; }` if you don't want these.
+  [](auto) { return std::string{}; })
+
+NVBENCH_DECLARE_TYPE_STRINGS(cudf::timestamp_ms, "cudf::timestamp_ms", "cudf::timestamp_ms");
+
+template <typename Type, algorithm Algo, cudf::duplicate_keep_option Keep>
+void nvbench_compaction(
+  nvbench::state& state,
+  nvbench::type_list<Type, nvbench::enum_type<Algo>, nvbench::enum_type<Keep>>)
+{
+  if constexpr ((not std::is_same_v<Type, int32_t> and
+                 Keep != cudf::duplicate_keep_option::KEEP_FIRST and
+                 Algo == algorithm::SORT_BASED) or
+                (Algo == algorithm::HASH_BASED and
+                 Keep != cudf::duplicate_keep_option::KEEP_FIRST)) {
+    state.skip("Skip unwanted benchmarks.");
+  }
+
+  cudf::rmm_pool_raii pool_raii;
+
+  auto const num_rows = state.get_int64("NumRows");
+
+  cudf::test::UniformRandomGenerator<long> rand_gen(0, 100);
+  auto elements = cudf::detail::make_counting_transform_iterator(
+    0, [&rand_gen](auto row) { return rand_gen.generate(); });
+  auto valids = cudf::detail::make_counting_transform_iterator(
+    0, [](auto i) { return i % 100 == 0 ? false : true; });
+  cudf::test::fixed_width_column_wrapper<Type, long> values(elements, elements + num_rows, valids);
+
+  auto input_column = cudf::column_view(values);
+  auto input_table  = cudf::table_view({input_column, input_column, input_column, input_column});
+
+  state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) {
+    rmm::cuda_stream_view stream_view{launch.get_stream()};
+    auto const result = [&]() {
+      if constexpr (Algo == algorithm::HASH_BASED) {
+        return cudf::detail::unordered_drop_duplicates(
+          input_table, {0}, cudf::null_equality::EQUAL, stream_view);
+      } else {
+        return cudf::detail::sort_and_drop_duplicates(input_table,
+                                                      {0},
+                                                      Keep,
+                                                      cudf::null_equality::EQUAL,
+                                                      cudf::null_order::BEFORE,
+                                                      stream_view);
+      }
+    }();
+  });
+}
+
+using data_type   = nvbench::type_list<bool,
+                                     nvbench::int8_t,
+                                     nvbench::int32_t,
+                                     nvbench::int64_t,
+                                     nvbench::float32_t,
+                                     cudf::timestamp_ms>;
+using algo        = nvbench::enum_type_list<algorithm::SORT_BASED, algorithm::HASH_BASED>;
+using keep_option = nvbench::enum_type_list<cudf::duplicate_keep_option::KEEP_FIRST,
+                                            cudf::duplicate_keep_option::KEEP_LAST,
+                                            cudf::duplicate_keep_option::KEEP_NONE>;
+
+NVBENCH_BENCH_TYPES(nvbench_compaction, NVBENCH_TYPE_AXES(data_type, algo, keep_option))
+  .set_name("drop_duplicates")
+  .set_type_axes_names({"Type", "Algorithm", "KeepOption"})
+  .add_int64_axis("NumRows", {10'000, 100'000, 1'000'000, 10'000'000});

cpp/cmake/thirdparty/get_cucollections.cmake

@@ -1,5 +1,5 @@
 # =============================================================================
-# Copyright (c) 2021, NVIDIA CORPORATION.
+# Copyright (c) 2021-2022, 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
@@ -21,7 +21,7 @@ function(find_and_configure_cucollections)
     cuco 0.0
     GLOBAL_TARGETS cuco::cuco
     CPM_ARGS GITHUB_REPOSITORY NVIDIA/cuCollections
-    GIT_TAG 193de1aa74f5721717f991ca757dc610c852bb17
+    GIT_TAG 922a87856aac17742fb964eeaf1b9bbc5d7a916e
     OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF"
   )
 

cpp/include/cudf/detail/stream_compaction.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2020, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2022, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -67,6 +67,19 @@ std::unique_ptr<table> apply_boolean_mask(
  * @param[in] stream CUDA stream used for device memory operations and kernel launches.
  */
 std::unique_ptr<table> drop_duplicates(
+  table_view const& input,
+  std::vector<size_type> const& keys,
+  duplicate_keep_option keep,
+  null_equality nulls_equal           = null_equality::EQUAL,
+  rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+/**
+ * @copydoc cudf::sort_and_drop_duplicates
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<table> sort_and_drop_duplicates(
   table_view const& input,
   std::vector<size_type> const& keys,
   duplicate_keep_option keep,
@@ -75,6 +88,18 @@ std::unique_ptr<table> drop_duplicates(
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @copydoc cudf::unordered_drop_duplicates
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<table> unordered_drop_duplicates(
+  table_view const& input,
+  std::vector<size_type> const& keys,
+  null_equality nulls_equal           = null_equality::EQUAL,
+  rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /**
  * @copydoc cudf::distinct_count(column_view const&, null_policy, nan_policy)
  *
@@ -94,5 +119,24 @@ cudf::size_type distinct_count(table_view const& input,
                                null_equality nulls_equal    = null_equality::EQUAL,
                                rmm::cuda_stream_view stream = rmm::cuda_stream_default);
 
+/**
+ * @copydoc cudf::unordered_distinct_count(column_view const&, null_policy, nan_policy)
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+cudf::size_type unordered_distinct_count(column_view const& input,
+                                         null_policy null_handling,
+                                         nan_policy nan_handling,
+                                         rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+
+/**
+ * @copydoc cudf::unordered_distinct_count(table_view const&, null_equality)
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+cudf::size_type unordered_distinct_count(table_view const& input,
+                                         null_equality nulls_equal    = null_equality::EQUAL,
+                                         rmm::cuda_stream_view stream = rmm::cuda_stream_default);
+
 }  // namespace detail
 }  // namespace cudf