[REVIEW] Add dictionary support to cudf::reduce

CHANGELOG.md

@@ -75,6 +75,7 @@
 - PR #6644 Cover different CSV reader/writer options in benchmarks
 - PR #6741 Cover different ORC and Parquet reader/writer options in benchmarks
 - PR #6651 Add cudf::dictionary::make_dictionary_pair_iterator
+- PR #6666 Add dictionary support to `cudf::reduce`
 - PR #6635 Add cudf::test::dictionary_column_wrapper class
 - PR #6702 Fix orc read corruption on boolean column
 - PR #6676 Add dictionary support to `cudf::quantile`

conda/recipes/libcudf/meta.yaml

@@ -71,12 +71,12 @@ test:
     - test -f $PREFIX/include/cudf/detail/null_mask.hpp
     - test -f $PREFIX/include/cudf/detail/nvtx/nvtx3.hpp
     - test -f $PREFIX/include/cudf/detail/nvtx/ranges.hpp
+    - test -f $PREFIX/include/cudf/detail/quantiles.hpp
     - test -f $PREFIX/include/cudf/detail/reduction_functions.hpp
     - test -f $PREFIX/include/cudf/detail/repeat.hpp
     - test -f $PREFIX/include/cudf/detail/replace.hpp
     - test -f $PREFIX/include/cudf/detail/reshape.hpp
     - test -f $PREFIX/include/cudf/detail/round.hpp
-    - test -f $PREFIX/include/cudf/detail/quantiles.hpp
     - test -f $PREFIX/include/cudf/detail/scatter.hpp
     - test -f $PREFIX/include/cudf/detail/search.hpp
     - test -f $PREFIX/include/cudf/detail/sequence.hpp

cpp/include/cudf/column/column_device_view.cuh

@@ -705,14 +705,17 @@ __device__ inline string_view const column_device_view::element<string_view>(
  * The basic dictionary elements are the indices which can be any index type.
  */
 struct index_element_fn {
-  template <typename IndexType, std::enable_if_t<is_index_type<IndexType>()>* = nullptr>
+  template <
+    typename IndexType,
+    std::enable_if_t<is_index_type<IndexType>() and std::is_unsigned<IndexType>::value>* = nullptr>
   __device__ size_type operator()(column_device_view const& input, size_type index)
   {
     return static_cast<size_type>(input.element<IndexType>(index));
   }
   template <typename IndexType,
             typename... Args,
-            std::enable_if_t<not is_index_type<IndexType>()>* = nullptr>
+            std::enable_if_t<not(is_index_type<IndexType>() and
+                                 std::is_unsigned<IndexType>::value)>* = nullptr>
   __device__ size_type operator()(Args&&... args)
   {
     release_assert(false and "indices must be an integral type");

cpp/include/cudf/detail/reduction.cuh

@@ -131,20 +131,6 @@ std::unique_ptr<scalar> reduce(InputIterator d_in,
   return std::unique_ptr<scalar>(s);
 }
 
-// @brief dictionary specialization of simple reduction
-template <typename Op,
-          typename InputIterator,
-          typename OutputType = typename thrust::iterator_value<InputIterator>::type,
-          typename std::enable_if_t<std::is_same<OutputType, dictionary32>::value>* = nullptr>
-std::unique_ptr<scalar> reduce(InputIterator d_in,
-                               cudf::size_type num_items,
-                               op::simple_op<Op> sop,
-                               rmm::cuda_stream_view stream,
-                               rmm::mr::device_memory_resource* mr)
-{
-  CUDF_FAIL("dictionary type not supported");
-}
-
 /** --------------------------------------------------------------------------*
  * @brief compute reduction by the compound operator (reduce and transform)
  *

cpp/src/copying/get_element.cu

@@ -16,6 +16,7 @@
 
 #include <cudf/column/column_device_view.cuh>
 #include <cudf/copying.hpp>
+#include <cudf/detail/indexalator.cuh>
 #include <cudf/dictionary/dictionary_column_view.hpp>
 #include <cudf/scalar/scalar_device_view.cuh>
 #include <cudf/scalar/scalar_factories.hpp>
@@ -84,25 +85,33 @@ struct get_element_functor {
     cudaStream_t stream                 = 0,
     rmm::mr::device_memory_resource *mr = rmm::mr::get_current_device_resource())
   {
-    auto dict_view = dictionary_column_view(input);
-    auto key_index_scalar =
-      get_element_functor{}.operator()<int32_t>(dict_view.indices(), index, stream);
-
-    size_type key_index =
-      static_cast<numeric_scalar<int32_t> const *>(key_index_scalar.get())->value(stream);
-    auto result = type_dispatcher(
-      dict_view.keys().type(), get_element_functor{}, dict_view.keys(), key_index, stream, mr);
-
-    auto result_validity = result->validity_data();
-    auto device_col      = column_device_view::create(input, stream);
+    auto dict_view    = dictionary_column_view(input);
+    auto indices_iter = detail::indexalator_factory::make_input_iterator(dict_view.indices());
+    numeric_scalar<size_type> key_index_scalar{index, true, stream};
+    auto d_key_index = get_scalar_device_view(key_index_scalar);
+    auto d_col       = column_device_view::create(input, stream);
 
+    // retrieve the indices value at index
     device_single_thread(
-      [result_validity, d_col = *device_col, index] __device__() mutable {
-        *result_validity = d_col.is_valid(index);
+      [d_key_index, d_col = *d_col, indices_iter, index] __device__() mutable {
+        d_key_index.set_value(indices_iter[index]);
+        d_key_index.set_valid(d_col.is_valid(index));
       },
       stream);
 
-    return result;
+    if (!key_index_scalar.is_valid(stream)) {
+      auto null_result = make_default_constructed_scalar(dict_view.keys().type());
+      null_result->set_valid(false, stream);
+      return null_result;
+    }
+
+    // retrieve the key element using the key-index
+    return type_dispatcher(dict_view.keys().type(),
+                           get_element_functor{},
+                           dict_view.keys(),
+                           key_index_scalar.value(stream),
+                           stream,
+                           mr);
   }
 
   template <typename T, std::enable_if_t<std::is_same<T, list_view>::value> *p = nullptr>

cpp/src/quantiles/quantile.cu

@@ -140,54 +140,28 @@ std::unique_ptr<column> quantile(column_view const& input,
 std::unique_ptr<column> quantile(column_view const& input,
                                  std::vector<double> const& q,
                                  interpolation interp,
-                                 column_view const& ordered_indices,
+                                 column_view const& indices,
                                  bool exact,
                                  rmm::cuda_stream_view stream,
                                  rmm::mr::device_memory_resource* mr)
 {
-  if (ordered_indices.is_empty()) {
+  if (indices.is_empty()) {
+    auto begin = thrust::make_counting_iterator<size_type>(0);
     if (exact) {
-      return detail::quantile<true>(input,
-                                    thrust::make_counting_iterator<size_type>(0),
-                                    input.size(),
-                                    q,
-                                    interp,
-                                    exact,
-                                    stream,
-                                    mr);
+      return quantile<true>(input, begin, input.size(), q, interp, exact, stream, mr);
     } else {
-      return detail::quantile<false>(input,
-                                     thrust::make_counting_iterator<size_type>(0),
-                                     input.size(),
-                                     q,
-                                     interp,
-                                     exact,
-                                     stream,
-                                     mr);
+      return quantile<false>(input, begin, input.size(), q, interp, exact, stream, mr);
     }
 
   } else {
-    CUDF_EXPECTS(ordered_indices.type() == data_type{type_to_id<size_type>()},
-                 "`ordered_indicies` type must be `INT32`.");
-
+    CUDF_EXPECTS(indices.type() == data_type{type_to_id<size_type>()},
+                 "`indicies` type must be `INT32`.");
     if (exact) {
-      return detail::quantile<true>(input,
-                                    ordered_indices.data<size_type>(),
-                                    ordered_indices.size(),
-                                    q,
-                                    interp,
-                                    exact,
-                                    stream,
-                                    mr);
+      return quantile<true>(
+        input, indices.begin<size_type>(), indices.size(), q, interp, exact, stream, mr);
     } else {
-      return detail::quantile<false>(input,
-                                     ordered_indices.data<size_type>(),
-                                     ordered_indices.size(),
-                                     q,
-                                     interp,
-                                     exact,
-                                     stream,
-                                     mr);
+      return quantile<false>(
+        input, indices.begin<size_type>(), indices.size(), q, interp, exact, stream, mr);
     }
   }
 }

cpp/src/reductions/all.cu

@@ -14,11 +14,9 @@
  * limitations under the License.
  */
 
-#include <reductions/simple.cuh>
-
 #include <cudf/detail/reduction_functions.hpp>
-
-#include <rmm/cuda_stream_view.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
+#include <reductions/simple.cuh>
 
 namespace cudf {
 namespace reduction {
@@ -30,7 +28,9 @@ std::unique_ptr<cudf::scalar> all(column_view const& col,
 {
   CUDF_EXPECTS(output_dtype == cudf::data_type(cudf::type_id::BOOL8),
                "all() operation can be applied with output type `BOOL8` only");
-  return cudf::type_dispatcher(col.type(),
+  auto const dispatch_type =
+    cudf::is_dictionary(col.type()) ? dictionary_column_view(col).keys().type() : col.type();
+  return cudf::type_dispatcher(dispatch_type,
                                simple::bool_result_element_dispatcher<cudf::reduction::op::min>{},
                                col,
                                stream,

cpp/src/reductions/any.cu

@@ -14,11 +14,9 @@
  * limitations under the License.
  */
 
-#include <reductions/simple.cuh>
-
 #include <cudf/detail/reduction_functions.hpp>
-
-#include <rmm/cuda_stream_view.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
+#include <reductions/simple.cuh>
 
 namespace cudf {
 namespace reduction {
@@ -30,7 +28,9 @@ std::unique_ptr<cudf::scalar> any(column_view const& col,
 {
   CUDF_EXPECTS(output_dtype == cudf::data_type(cudf::type_id::BOOL8),
                "any() operation can be applied with output type `bool8` only");
-  return cudf::type_dispatcher(col.type(),
+  auto const dispatch_type =
+    cudf::is_dictionary(col.type()) ? dictionary_column_view(col).keys().type() : col.type();
+  return cudf::type_dispatcher(dispatch_type,
                                simple::bool_result_element_dispatcher<cudf::reduction::op::max>{},
                                col,
                                stream,

cpp/src/reductions/compound.cuh

@@ -17,13 +17,11 @@
 #pragma once
 
 #include <cudf/detail/reduction.cuh>
-
+#include <cudf/dictionary/detail/iterator.cuh>
 #include <cudf/scalar/scalar_factories.hpp>
+#include <cudf/utilities/traits.hpp>
 #include <cudf/utilities/type_dispatcher.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
-#include <rmm/device_scalar.hpp>
-
 namespace cudf {
 namespace reduction {
 namespace compound {
@@ -58,23 +56,37 @@ std::unique_ptr<scalar> compound_reduction(column_view const& col,
   std::unique_ptr<scalar> result;
   Op compound_op{};
 
-  if (col.has_nulls()) {
-    auto it = thrust::make_transform_iterator(
-      dcol->pair_begin<ElementType, true>(),
-      compound_op.template get_null_replacing_element_transformer<ResultType>());
-    result = detail::reduce<Op, decltype(it), ResultType>(
-      it, col.size(), compound_op, valid_count, ddof, stream, mr);
+  if (!cudf::is_dictionary(col.type())) {
+    if (col.has_nulls()) {
+      auto it = thrust::make_transform_iterator(
+        dcol->pair_begin<ElementType, true>(),
+        compound_op.template get_null_replacing_element_transformer<ResultType>());
+      result = detail::reduce<Op, decltype(it), ResultType>(
+        it, col.size(), compound_op, valid_count, ddof, stream, mr);
+    } else {
+      auto it = thrust::make_transform_iterator(
+        dcol->begin<ElementType>(), compound_op.template get_element_transformer<ResultType>());
+      result = detail::reduce<Op, decltype(it), ResultType>(
+        it, col.size(), compound_op, valid_count, ddof, stream, mr);
+    }
   } else {
-    auto it = thrust::make_transform_iterator(
-      dcol->begin<ElementType>(), compound_op.template get_element_transformer<ResultType>());
-    result = detail::reduce<Op, decltype(it), ResultType>(
-      it, col.size(), compound_op, valid_count, ddof, stream, mr);
+    if (col.has_nulls()) {
+      auto it = thrust::make_transform_iterator(
+        cudf::dictionary::detail::make_dictionary_pair_iterator<ElementType, true>(*dcol),
+        compound_op.template get_null_replacing_element_transformer<ResultType>());
+      result = detail::reduce<Op, decltype(it), ResultType>(
+        it, col.size(), compound_op, valid_count, ddof, stream, mr);
+    } else {
+      auto it = thrust::make_transform_iterator(
+        cudf::dictionary::detail::make_dictionary_iterator<ElementType>(*dcol),
+        compound_op.template get_element_transformer<ResultType>());
+      result = detail::reduce<Op, decltype(it), ResultType>(
+        it, col.size(), compound_op, valid_count, ddof, stream, mr);
+    }
   }
+
   // set scalar is valid
-  if (col.null_count() < col.size())
-    result->set_valid(true, stream);
-  else
-    result->set_valid(false, stream);
+  result->set_valid(col.null_count() < col.size(), stream);
   return result;
 };
 

cpp/src/reductions/max.cu

@@ -14,9 +14,9 @@
  * limitations under the License.
  */
 
-#include <reductions/simple.cuh>
-
 #include <cudf/detail/reduction_functions.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
+#include <reductions/simple.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 
@@ -28,9 +28,17 @@ std::unique_ptr<cudf::scalar> max(column_view const& col,
                                   rmm::cuda_stream_view stream,
                                   rmm::mr::device_memory_resource* mr)
 {
-  CUDF_EXPECTS(col.type() == output_dtype, "max() operation requires matching output type");
-  return cudf::type_dispatcher(
-    col.type(), simple::same_element_type_dispatcher<cudf::reduction::op::max>{}, col, stream, mr);
+  auto const input_type =
+    cudf::is_dictionary(col.type()) ? cudf::dictionary_column_view(col).keys().type() : col.type();
+  CUDF_EXPECTS(input_type == output_dtype, "max() operation requires matching output type");
+  auto const dispatch_type = cudf::is_dictionary(col.type())
+                               ? cudf::dictionary_column_view(col).indices().type()
+                               : col.type();
+  return cudf::type_dispatcher(dispatch_type,
+                               simple::same_element_type_dispatcher<cudf::reduction::op::max>{},
+                               col,
+                               stream,
+                               mr);
 }
 
 }  // namespace reduction

cpp/src/reductions/mean.cu

@@ -15,9 +15,9 @@
  */
 // The translation unit for reduction `mean`
 
-#include "compound.cuh"
-
 #include <cudf/detail/reduction_functions.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
+#include <reductions/compound.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 
@@ -27,6 +27,8 @@ std::unique_ptr<cudf::scalar> cudf::reduction::mean(column_view const& col,
                                                     rmm::mr::device_memory_resource* mr)
 {
   using reducer = cudf::reduction::compound::element_type_dispatcher<cudf::reduction::op::mean>;
+  auto col_type =
+    cudf::is_dictionary(col.type()) ? dictionary_column_view(col).keys().type() : col.type();
   return cudf::type_dispatcher(
-    col.type(), reducer(), col, output_dtype, /* ddof is not used for mean*/ 1, stream, mr);
+    col_type, reducer(), col, output_dtype, /* ddof is not used for mean*/ 1, stream, mr);
 }

cpp/src/reductions/min.cu

@@ -14,11 +14,9 @@
  * limitations under the License.
  */
 
-#include <reductions/simple.cuh>
-
 #include <cudf/detail/reduction_functions.hpp>
-
-#include <rmm/cuda_stream_view.hpp>
+#include <cudf/dictionary/dictionary_column_view.hpp>
+#include <reductions/simple.cuh>
 
 namespace cudf {
 namespace reduction {
@@ -28,9 +26,17 @@ std::unique_ptr<cudf::scalar> min(column_view const& col,
                                   rmm::cuda_stream_view stream,
                                   rmm::mr::device_memory_resource* mr)
 {
-  CUDF_EXPECTS(col.type() == output_dtype, "min() operation requires matching output type");
-  return cudf::type_dispatcher(
-    col.type(), simple::same_element_type_dispatcher<cudf::reduction::op::min>{}, col, stream, mr);
+  auto const input_type =
+    cudf::is_dictionary(col.type()) ? cudf::dictionary_column_view(col).keys().type() : col.type();
+  CUDF_EXPECTS(input_type == output_dtype, "min() operation requires matching output type");
+  auto const dispatch_type = cudf::is_dictionary(col.type())
+                               ? cudf::dictionary_column_view(col).indices().type()
+                               : col.type();
+  return cudf::type_dispatcher(dispatch_type,
+                               simple::same_element_type_dispatcher<cudf::reduction::op::min>{},
+                               col,
+                               stream,
+                               mr);
 }
 
 }  // namespace reduction