Merge branch 'branch-21.12' of github.com:rapidsai/cudf into branch-2…

…1.12

conda/recipes/libcudf/meta.yaml

@@ -93,6 +93,7 @@ test:
     - test -f $PREFIX/include/cudf/detail/sequence.hpp
     - test -f $PREFIX/include/cudf/detail/sorting.hpp
     - test -f $PREFIX/include/cudf/detail/stream_compaction.hpp
+    - test -f $PREFIX/include/cudf/detail/structs/utilities.hpp
     - test -f $PREFIX/include/cudf/detail/tdigest/tdigest.hpp
     - test -f $PREFIX/include/cudf/detail/transform.hpp
     - test -f $PREFIX/include/cudf/detail/transpose.hpp

cpp/CMakeLists.txt

@@ -97,6 +97,11 @@ rapids_find_package(CUDAToolkit REQUIRED
     INSTALL_EXPORT_SET cudf-exports)
 include(cmake/Modules/ConfigureCUDA.cmake) # set other CUDA compilation flags
 
+# ctest cuda memcheck
+find_program(CUDA_SANITIZER compute-sanitizer)
+set(MEMORYCHECK_COMMAND ${CUDA_SANITIZER})
+set(MEMORYCHECK_TYPE CudaSanitizer)
+set(CUDA_SANITIZER_COMMAND_OPTIONS "--tool memcheck")
 
 ###################################################################################################
 # - dependencies ----------------------------------------------------------------------------------

cpp/include/cudf/detail/groupby/sort_helper.hpp

@@ -18,6 +18,7 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_view.hpp>
+#include <cudf/detail/structs/utilities.hpp>
 #include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
 
@@ -218,7 +219,7 @@ struct sort_groupby_helper {
   column_ptr _keys_bitmask_column;   ///< Column representing rows with one or more nulls values
   table_view _keys;                  ///< Input keys to sort by
   table_view _unflattened_keys;      ///< Input keys, unflattened and possibly nested
-  std::vector<column_ptr> _struct_null_vectors;  ///< Null vectors for struct columns in _keys
+  structs::detail::flattened_table _flattened;  ///< Support datastructures for _keys
 
   index_vector_ptr
     _group_offsets;  ///< Indices into sorted _keys indicating starting index of each groups

cpp/include/cudf/detail/sorting.hpp

@@ -76,6 +76,19 @@ std::unique_ptr<column> segmented_sorted_order(
   rmm::cuda_stream_view stream                   = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
 
+/**
+ * @copydoc cudf::stable_segmented_sorted_order
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<column> stable_segmented_sorted_order(
+  table_view const& keys,
+  column_view const& segment_offsets,
+  std::vector<order> const& column_order         = {},
+  std::vector<null_order> const& null_precedence = {},
+  rmm::cuda_stream_view stream                   = rmm::cuda_stream_default,
+  rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
+
 /**
  * @copydoc cudf::segmented_sort_by_key
  *
@@ -90,6 +103,20 @@ std::unique_ptr<table> segmented_sort_by_key(
   rmm::cuda_stream_view stream                   = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
 
+/**
+ * @copydoc cudf::stable_segmented_sort_by_key
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<table> stable_segmented_sort_by_key(
+  table_view const& values,
+  table_view const& keys,
+  column_view const& segment_offsets,
+  std::vector<order> const& column_order         = {},
+  std::vector<null_order> const& null_precedence = {},
+  rmm::cuda_stream_view stream                   = rmm::cuda_stream_default,
+  rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
+
 /**
  * @copydoc cudf::sort
  *

cpp/src/structs/utilities.hpp → cpp/include/cudf/detail/structs/utilities.hpp

@@ -16,10 +16,12 @@
 #pragma once
 
 #include <cudf/structs/structs_column_view.hpp>
+#include <cudf/table/table_view.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_buffer.hpp>
 
 namespace cudf {
 namespace structs {
@@ -64,6 +66,71 @@ std::vector<std::vector<column_view>> extract_ordered_struct_children(
  */
 bool is_or_has_nested_lists(cudf::column_view const& col);
 
+/**
+ * @brief Result of `flatten_nested_columns()`, where all `STRUCT` columns are replaced with
+ * their non-nested member columns, and `BOOL8` columns for their null masks.
+ *
+ * `flatten_nested_columns()` produces a "flattened" table_view with all `STRUCT` columns
+ * replaced with their child column_views, preceded by their null masks.
+ * All newly allocated columns and device_buffers that back the returned table_view
+ * are also encapsulated in `flatten_result`.
+ *
+ * Objects of `flatten_result` need to kept alive while its table_view is accessed.
+ */
+class flattened_table {
+ public:
+  /**
+   * @brief Constructor, to be used from `flatten_nested_columns()`.
+   *
+   * @param flattened_columns_ table_view resulting from `flatten_nested_columns()`
+   * @param orders_ Per-column ordering of the table_view
+   * @param null_orders_ Per-column null_order of the table_view
+   * @param columns_ Newly allocated columns to back the table_view
+   * @param null_masks_ Newly allocated null masks to back the table_view
+   */
+  flattened_table(table_view const& flattened_columns_,
+                  std::vector<order> const& orders_,
+                  std::vector<null_order> const& null_orders_,
+                  std::vector<std::unique_ptr<column>>&& columns_,
+                  std::vector<rmm::device_buffer>&& null_masks_)
+    : _flattened_columns{flattened_columns_},
+      _orders{orders_},
+      _null_orders{null_orders_},
+      _columns{std::move(columns_)},
+      _superimposed_nullmasks{std::move(null_masks_)}
+  {
+  }
+
+  flattened_table() = default;
+
+  /**
+   * @brief Getter for the flattened columns, as a `table_view`.
+   */
+  table_view flattened_columns() const { return _flattened_columns; }
+
+  /**
+   * @brief Getter for the cudf::order of the table_view's columns.
+   */
+  std::vector<order> orders() const { return _orders; }
+
+  /**
+   * @brief Getter for the cudf::null_order of the table_view's columns.
+   */
+  std::vector<null_order> null_orders() const { return _null_orders; }
+
+  /**
+   * @brief Conversion to `table_view`, to fetch flattened columns.
+   */
+  operator table_view() const { return flattened_columns(); }
+
+ private:
+  table_view _flattened_columns;
+  std::vector<order> _orders;
+  std::vector<null_order> _null_orders;
+  std::vector<std::unique_ptr<column>> _columns;
+  std::vector<rmm::device_buffer> _superimposed_nullmasks;
+};
+
 /**
  * @brief Flatten table with struct columns to table with constituent columns of struct columns.
  *
@@ -74,17 +141,14 @@ bool is_or_has_nested_lists(cudf::column_view const& col);
  * @param null_precedence null order for input table
  * @param nullability force output to have nullability columns even if input columns
  * are all valid
- * @return tuple with flattened table, flattened column order, flattened null precedence,
- * vector of boolean columns (struct validity).
+ * @return `flatten_result` with flattened table, flattened column order, flattened null precedence,
+ * alongside the supporting columns and device_buffers for the flattened table.
  */
-std::tuple<table_view,
-           std::vector<order>,
-           std::vector<null_order>,
-           std::vector<std::unique_ptr<column>>>
-flatten_nested_columns(table_view const& input,
-                       std::vector<order> const& column_order,
-                       std::vector<null_order> const& null_precedence,
-                       column_nullability nullability = column_nullability::MATCH_INCOMING);
+flattened_table flatten_nested_columns(
+  table_view const& input,
+  std::vector<order> const& column_order,
+  std::vector<null_order> const& null_precedence,
+  column_nullability nullability = column_nullability::MATCH_INCOMING);
 
 /**
  * @brief Unflatten columns flattened as by `flatten_nested_columns()`,
@@ -156,6 +220,31 @@ std::tuple<cudf::column_view, std::vector<rmm::device_buffer>> superimpose_paren
   rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Push down nulls from a parent mask into child columns, using bitwise AND,
+ * for all columns in the specified table.
+ *
+ * This function constructs a table_view containing a new column_view instance equivalent to
+ * every column_view in the specified table. Each column_view might contain possibly new
+ * child column_views, all with possibly new null mask values reflecting null rows from the
+ * parent column:
+ * 1. If the column is not STRUCT, the column is returned unmodified, with no new
+ *    supporting device_buffer instances.
+ * 2. If the column is STRUCT, the null masks of the parent and child are bitwise-ANDed, and a
+ *    modified column_view is returned. This applies recursively.
+ *
+ * @param table The table_view of (possibly STRUCT) columns whose nulls need to be pushed to its
+ * members.
+ * @param stream CUDA stream used for device memory operations and kernel launches.
+ * @param mr     Device memory resource used to allocate new device memory.
+ * @return A pair of:
+ *         1. table_view of columns with nulls pushed down to child columns, as appropriate.
+ *         2. Supporting device_buffer instances, for any newly constructed null masks.
+ */
+std::tuple<cudf::table_view, std::vector<rmm::device_buffer>> superimpose_parent_nulls(
+  table_view const& table,
+  rmm::cuda_stream_view stream        = rmm::cuda_stream_default,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 }  // namespace detail
 }  // namespace structs
 }  // namespace cudf

cpp/include/cudf/lists/detail/sorting.hpp

@@ -34,6 +34,19 @@ std::unique_ptr<column> sort_lists(
   null_order null_precedence,
   rmm::cuda_stream_view stream,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
+/**
+ * @copydoc cudf::lists::stable_sort_lists
+ *
+ * @param stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<column> stable_sort_lists(
+  lists_column_view const& input,
+  order column_order,
+  null_order null_precedence,
+  rmm::cuda_stream_view stream,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 }  // namespace detail
 }  // namespace lists
 }  // namespace cudf

cpp/include/cudf/lists/sorting.hpp

@@ -54,6 +54,18 @@ std::unique_ptr<column> sort_lists(
   null_order null_precedence,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Segmented sort of the elements within a list in each row of a list column using stable
+ * sort.
+ *
+ * @copydoc cudf::lists::sort_lists
+ */
+std::unique_ptr<column> stable_sort_lists(
+  lists_column_view const& source_column,
+  order column_order,
+  null_order null_precedence,
+  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
+
 /** @} */  // end of group
 }  // namespace lists
 }  // namespace cudf

cpp/include/cudf/sorting.hpp

@@ -211,6 +211,18 @@ std::unique_ptr<column> segmented_sorted_order(
   std::vector<null_order> const& null_precedence = {},
   rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Returns sorted order after stably sorting each segment in the table.
+ *
+ * @copydoc cudf::segmented_sorted_order
+ */
+std::unique_ptr<column> stable_segmented_sorted_order(
+  table_view const& keys,
+  column_view const& segment_offsets,
+  std::vector<order> const& column_order         = {},
+  std::vector<null_order> const& null_precedence = {},
+  rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
+
 /**
  * @brief Performs a lexicographic segmented sort of a table
  *
@@ -241,5 +253,18 @@ std::unique_ptr<table> segmented_sort_by_key(
   std::vector<null_order> const& null_precedence = {},
   rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Performs a stably lexicographic segmented sort of a table
+ *
+ * @copydoc cudf::segmented_sort_by_key
+ */
+std::unique_ptr<table> stable_segmented_sort_by_key(
+  table_view const& values,
+  table_view const& keys,
+  column_view const& segment_offsets,
+  std::vector<order> const& column_order         = {},
+  std::vector<null_order> const& null_precedence = {},
+  rmm::mr::device_memory_resource* mr            = rmm::mr::get_current_device_resource());
+
 /** @} */  // end of group
 }  // namespace cudf

cpp/include/cudf_test/base_fixture.hpp

@@ -266,16 +266,22 @@ inline std::shared_ptr<rmm::mr::device_memory_resource> create_memory_resource(
  *
  * Currently only supports 'rmm_mode' string parameter, which set the rmm
  * allocation mode. The default value of the parameter is 'pool'.
+ * Environment variable 'CUDF_TEST_RMM_MODE' can also be used to set the rmm
+ * allocation mode. If both are set, the value of 'rmm_mode' string parameter
+ * takes precedence.
  *
  * @return Parsing results in the form of unordered map
  */
 inline auto parse_cudf_test_opts(int argc, char** argv)
 {
   try {
     cxxopts::Options options(argv[0], " - cuDF tests command line options");
+    const char* env_rmm_mode = std::getenv("GTEST_CUDF_RMM_MODE");  // Overridden by CLI options
+    auto default_rmm_mode    = env_rmm_mode ? env_rmm_mode : "pool";
     options.allow_unrecognised_options().add_options()(
-      "rmm_mode", "RMM allocation mode", cxxopts::value<std::string>()->default_value("pool"));
-
+      "rmm_mode",
+      "RMM allocation mode",
+      cxxopts::value<std::string>()->default_value(default_rmm_mode));
     return options.parse(argc, argv);
   } catch (const cxxopts::OptionException& e) {
     CUDF_FAIL("Error parsing command line options");

cpp/src/groupby/groupby.cu

@@ -25,6 +25,7 @@
 #include <cudf/detail/groupby/group_replace_nulls.hpp>
 #include <cudf/detail/groupby/sort_helper.hpp>
 #include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/detail/structs/utilities.hpp>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/groupby.hpp>
 #include <cudf/strings/string_view.hpp>
@@ -33,7 +34,6 @@
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/traits.hpp>
-#include <structs/utilities.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 
@@ -76,8 +76,8 @@ std::pair<std::unique_ptr<table>, std::vector<aggregation_result>> groupby::disp
   if (_keys_are_sorted == sorted::NO and not _helper and
       detail::hash::can_use_hash_groupby(_keys, requests)) {
     // Optionally flatten nested key columns.
-    auto [flattened_keys, _, __, ___] =
-      flatten_nested_columns(_keys, {}, {}, column_nullability::FORCE);
+    auto flattened             = flatten_nested_columns(_keys, {}, {}, column_nullability::FORCE);
+    auto flattened_keys        = flattened.flattened_columns();
     auto is_supported_key_type = [](auto col) { return cudf::is_equality_comparable(col.type()); };
     CUDF_EXPECTS(std::all_of(flattened_keys.begin(), flattened_keys.end(), is_supported_key_type),
                  "Unsupported groupby key type does not support equality comparison");

cpp/src/groupby/sort/group_rank_scan.cu

@@ -16,15 +16,14 @@
 
 #include <cudf/column/column.hpp>
 #include <cudf/column/column_factories.hpp>
+#include <cudf/detail/structs/utilities.hpp>
 #include <cudf/detail/utilities/device_operators.cuh>
 #include <cudf/table/row_operators.cuh>
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <structs/utilities.hpp>
-
 namespace cudf {
 namespace groupby {
 namespace detail {
@@ -55,9 +54,9 @@ std::unique_ptr<column> rank_generator(column_view const& order_by,
 {
   auto const superimposed = structs::detail::superimpose_parent_nulls(order_by, stream, mr);
   table_view const order_table{{std::get<0>(superimposed)}};
-  auto const flattener = cudf::structs::detail::flatten_nested_columns(
+  auto const flattened = cudf::structs::detail::flatten_nested_columns(
     order_table, {}, {}, structs::detail::column_nullability::MATCH_INCOMING);
-  auto const d_flat_order = table_device_view::create(std::get<0>(flattener), stream);
+  auto const d_flat_order = table_device_view::create(flattened, stream);
   row_equality_comparator<has_nulls> comparator(*d_flat_order, *d_flat_order, true);
   auto ranks         = make_fixed_width_column(data_type{type_to_id<size_type>()},
                                        order_table.num_rows(),

cpp/src/groupby/sort/scan.cpp

@@ -22,15 +22,14 @@
 #include <cudf/column/column_view.hpp>
 #include <cudf/detail/aggregation/aggregation.hpp>
 #include <cudf/detail/aggregation/result_cache.hpp>
+#include <cudf/detail/structs/utilities.hpp>
 #include <cudf/groupby.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 
-#include <structs/utilities.hpp>
-
 #include <memory>
 
 namespace cudf {