Add Segmented sort

conda/recipes/libcudf/meta.yaml

@@ -128,6 +128,7 @@ test:
     - test -f $PREFIX/include/cudf/lists/contains.hpp
     - test -f $PREFIX/include/cudf/lists/gather.hpp
     - test -f $PREFIX/include/cudf/lists/lists_column_view.hpp
+    - test -f $PREFIX/include/cudf/lists/sorting.hpp
     - test -f $PREFIX/include/cudf/merge.hpp
     - test -f $PREFIX/include/cudf/null_mask.hpp
     - test -f $PREFIX/include/cudf/partitioning.hpp

cpp/include/cudf/detail/sorting.hpp

@@ -63,5 +63,18 @@ std::unique_ptr<table> 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::sort_lists
+ *
+ * @param[in] stream CUDA stream used for device memory operations and kernel launches.
+ */
+std::unique_ptr<table> sort_lists(
+  table_view const& values,
+  table_view const& keys,
+  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());
+
 }  // namespace detail
 }  // namespace cudf

cpp/include/cudf/lists/sorting.hpp

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2021, 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/column/column_view.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+
+namespace cudf {
+namespace lists {
+/**
+ * @addtogroup lists_sort
+ * @{
+ * @file
+ */
+
+/**
+ * @brief Segmented sort of the elements within a list in each row of a list column.
+ *
+ * `source_column` with depth 1 is only supported.
+ *
+ * * @code{.pseudo}
+ * source_column            : [{4, 2, 3, 1}, {1, 2, NULL, 4}, {-10, 10, 0}]
+ *
+ * Ascending,  Null After   : [{1, 2, 3, 4}, {1, 2, 4, NULL}, {-10, 0, 10}]
+ * Ascending,  Null Before  : [{1, 2, 3, 4}, {NULL, 1, 2, 4}, {-10, 0, 10}]
+ * Descending, Null After   : [{4, 3, 2, 1}, {NULL, 4, 2, 1}, {10, 0, -10}]
+ * Descending, Null Before  : [{4, 3, 2, 1}, {4, 2, 1, NULL}, {10, 0, -10}]
+ * @endcode
+ *
+ * @param source_column View of the list column of numeric types to sort
+ * @param column_order The desired sort order
+ * @param null_precedence The desired order of null compared to other elements in the list
+ * @param mr Device memory resource to allocate any returned objects
+ * @return list column with elements in each list sorted.
+ *
+ */
+std::unique_ptr<column> 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

@@ -128,7 +128,7 @@ std::unique_ptr<table> sort(
  * @param values The table to reorder
  * @param keys The table that determines the ordering
  * @param column_order The desired order for each column in `keys`. Size must be
- * equal to `input.num_columns()` or empty. If empty, all columns are sorted in
+ * equal to `keys.num_columns()` or empty. If empty, all columns are sorted in
  * ascending order.
  * @param null_precedence The desired order of a null element compared to other
  * elements for each column in `keys`. Size must be equal to
@@ -184,5 +184,33 @@ std::unique_ptr<column> rank(
   bool percentage,
   rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource());
 
+/**
+ * @brief Performs a lexicographic segmented sort of the list in each row of a table of list columns
+ *
+ * `keys` with list columns of depth 1 is only supported.
+ * @throws cudf::logic_error if `values.num_rows() != keys.num_rows()`.
+ * @throws cudf::logic_error if any list sizes of corresponding row in each column are not equal.
+ * @throws cudf::logic_error if any column of `keys` or `values` is not a list column.
+ *
+ * @param values The table with list columns to reorder
+ * @param keys The table with list coumns that determines the ordering of elements in each list
+ * @param column_order The desired order for each column in `keys`. Size must be
+ * equal to `keys.num_columns()` or empty. If empty, all columns are sorted in
+ * ascending order.
+ * @param null_precedence The desired order of a null element compared to other
+ * elements for each column in `keys`. Size must be equal to
+ * `keys.num_columns()` or empty. If empty, all columns will be sorted with
+ * `null_order::BEFORE`.
+ * @param mr Device memory resource to allocate any returned objects
+ * @return table with list columns with elements in each list sorted.
+ *
+ */
+std::unique_ptr<table> sort_lists(
+  table_view const& values,
+  table_view const& keys,
+  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/src/lists/segmented_sort.cu

@@ -0,0 +1,246 @@
+/*
+ * Copyright (c) 2021, 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.hpp>
+#include <cudf/column/column_device_view.cuh>
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/copy.hpp>
+#include <cudf/detail/gather.cuh>
+#include <cudf/detail/iterator.cuh>
+#include <cudf/detail/null_mask.hpp>
+#include <cudf/detail/nvtx/ranges.hpp>
+#include <cudf/lists/lists_column_view.hpp>
+#include <cudf/lists/sorting.hpp>
+#include <cudf/utilities/error.hpp>
+#include <cudf/utilities/type_dispatcher.hpp>
+
+#include <thrust/copy.h>
+#include <cub/device/device_segmented_radix_sort.cuh>
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/device_buffer.hpp>
+#include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
+
+namespace cudf {
+namespace lists {
+namespace detail {
+
+struct SortPairs {
+  template <typename KeyT, typename ValueT, typename OffsetIteratorT>
+  void SortPairsAscending(KeyT const* keys_in,
+                          KeyT* keys_out,
+                          ValueT const* values_in,
+                          ValueT* values_out,
+                          int num_items,
+                          int num_segments,
+                          OffsetIteratorT begin_offsets,
+                          OffsetIteratorT end_offsets,
+                          rmm::cuda_stream_view stream)
+  {
+    rmm::device_buffer d_temp_storage;
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairs(d_temp_storage.data(),
+                                             temp_storage_bytes,
+                                             keys_in,
+                                             keys_out,
+                                             values_in,
+                                             values_out,
+                                             num_items,
+                                             num_segments,
+                                             begin_offsets,
+                                             end_offsets,
+                                             0,
+                                             sizeof(KeyT) * 8,
+                                             stream.value());
+    d_temp_storage = rmm::device_buffer{temp_storage_bytes, stream};
+
+    cub::DeviceSegmentedRadixSort::SortPairs(d_temp_storage.data(),
+                                             temp_storage_bytes,
+                                             keys_in,
+                                             keys_out,
+                                             values_in,
+                                             values_out,
+                                             num_items,
+                                             num_segments,
+                                             begin_offsets,
+                                             end_offsets,
+                                             0,
+                                             sizeof(KeyT) * 8,
+                                             stream.value());
+  }
+
+  template <typename KeyT, typename ValueT, typename OffsetIteratorT>
+  void SortPairsDescending(KeyT const* keys_in,
+                           KeyT* keys_out,
+                           ValueT const* values_in,
+                           ValueT* values_out,
+                           int num_items,
+                           int num_segments,
+                           OffsetIteratorT begin_offsets,
+                           OffsetIteratorT end_offsets,
+                           rmm::cuda_stream_view stream)
+  {
+    rmm::device_buffer d_temp_storage;
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage.data(),
+                                                       temp_storage_bytes,
+                                                       keys_in,
+                                                       keys_out,
+                                                       values_in,
+                                                       values_out,
+                                                       num_items,
+                                                       num_segments,
+                                                       begin_offsets,
+                                                       end_offsets,
+                                                       0,
+                                                       sizeof(KeyT) * 8,
+                                                       stream.value());
+    d_temp_storage = rmm::device_buffer{temp_storage_bytes, stream};
+
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage.data(),
+                                                       temp_storage_bytes,
+                                                       keys_in,
+                                                       keys_out,
+                                                       values_in,
+                                                       values_out,
+                                                       num_items,
+                                                       num_segments,
+                                                       begin_offsets,
+                                                       end_offsets,
+                                                       0,
+                                                       sizeof(KeyT) * 8,
+                                                       stream.value());
+  }
+  template <typename T>
+  std::enable_if_t<not is_numeric<T>(), std::unique_ptr<column>> operator()(
+    column_view const& child,
+    column_view const& offsets,
+    order column_order,
+    null_order null_precedence,
+    rmm::cuda_stream_view stream,
+    rmm::mr::device_memory_resource* mr)
+  {
+    CUDF_FAIL("segmented sort is not supported for non-numeric list types");
+  }
+  template <typename T>
+  std::enable_if_t<is_numeric<T>(), std::unique_ptr<column>> operator()(
+    column_view const& child,
+    column_view const& offsets,
+    order column_order,
+    null_order null_precedence,
+    rmm::cuda_stream_view stream,
+    rmm::mr::device_memory_resource* mr)
+  {
+    auto output =
+      cudf::detail::allocate_like(child, child.size(), mask_allocation_policy::NEVER, stream, mr);
+    mutable_column_view mutable_output_view = output->mutable_view();
+
+    auto keys = [&]() {
+      if (child.nullable()) {
+        rmm::device_uvector<T> keys(child.size(), stream);
+        auto null_replace_T = null_precedence == null_order::AFTER ? std::numeric_limits<T>::max()
+                                                                   : std::numeric_limits<T>::min();
+        auto device_child = column_device_view::create(child, stream);
+        auto keys_in =
+          cudf::detail::make_null_replacement_iterator<T>(*device_child, null_replace_T);
+        thrust::copy_n(rmm::exec_policy(stream), keys_in, child.size(), keys.begin());
+        return keys;
+      }
+      return rmm::device_uvector<T>{0, stream};
+    }();
+
+    std::unique_ptr<column> sorted_indices = cudf::make_numeric_column(
+      data_type(type_to_id<size_type>()), child.size(), mask_state::UNALLOCATED, stream, mr);
+    mutable_column_view mutable_indices_view = sorted_indices->mutable_view();
+    thrust::sequence(rmm::exec_policy(stream),
+                     mutable_indices_view.begin<size_type>(),
+                     mutable_indices_view.end<size_type>(),
+                     0);
+
+    if (column_order == order::ASCENDING)
+      SortPairsAscending(child.nullable() ? keys.data() : child.begin<T>(),
+                         mutable_output_view.begin<T>(),
+                         mutable_indices_view.begin<size_type>(),
+                         mutable_indices_view.begin<size_type>(),
+                         child.size(),
+                         offsets.size() - 1,
+                         offsets.begin<size_type>(),
+                         offsets.begin<size_type>() + 1,
+                         stream);
+    else
+      SortPairsDescending(child.nullable() ? keys.data() : child.begin<T>(),
+                          mutable_output_view.begin<T>(),
+                          mutable_indices_view.begin<size_type>(),
+                          mutable_indices_view.begin<size_type>(),
+                          child.size(),
+                          offsets.size() - 1,
+                          offsets.begin<size_type>(),
+                          offsets.begin<size_type>() + 1,
+                          stream);
+    std::vector<std::unique_ptr<column>> output_cols;
+    output_cols.push_back(std::move(output));
+    // rearrange the null_mask.
+    cudf::detail::gather_bitmask(cudf::table_view{{child}},
+                                 mutable_indices_view.begin<size_type>(),
+                                 output_cols,
+                                 cudf::detail::gather_bitmask_op::DONT_CHECK,
+                                 stream,
+                                 mr);
+    return std::move(output_cols.front());
+  }
+};
+
+std::unique_ptr<column> 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)
+{
+  if (input.size() == 0) return {};
+
+  auto output_child = type_dispatcher(input.child().type(),
+                                      SortPairs{},
+                                      input.get_sliced_child(stream),
+                                      input.offsets(),
+                                      column_order,
+                                      null_precedence,
+                                      stream,
+                                      mr);
+
+  // Copy list offsets.
+  auto output_offset = std::make_unique<column>(input.offsets(), stream, mr);
+  auto null_mask     = cudf::detail::copy_bitmask(input.parent(), stream, mr);
+
+  // Assemble list column & return
+  return make_lists_column(input.size(),
+                           std::move(output_offset),
+                           std::move(output_child),
+                           input.null_count(),
+                           std::move(null_mask));
+}
+}  // namespace detail
+
+std::unique_ptr<column> sort_lists(lists_column_view const& input,
+                                   order column_order,
+                                   null_order null_precedence,
+                                   rmm::mr::device_memory_resource* mr)
+{
+  CUDF_FUNC_RANGE();
+  return detail::sort_lists(input, column_order, null_precedence, rmm::cuda_stream_default, mr);
+}
+
+}  // namespace lists
+}  // namespace cudf