Fix cudf::hash_partition kernel launch error with decimal128 types

cpp/src/partitioning/partitioning.cu

@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-#include <cub/cub.cuh>
 #include <cudf/column/column_factories.hpp>
 #include <cudf/copying.hpp>
 #include <cudf/detail/gather.hpp>
@@ -36,6 +35,9 @@
 #include <thrust/scan.h>
 #include <thrust/transform.h>
 
+#include <cub/block/block_scan.cuh>
+#include <cub/device/device_histogram.cuh>
+
 namespace cudf {
 namespace {
 // Launch configuration for optimized hash partition
@@ -389,7 +391,15 @@ rmm::device_uvector<size_type> compute_gather_map(size_type num_rows,
 }
 
 struct copy_block_partitions_dispatcher {
-  template <typename DataType, std::enable_if_t<is_fixed_width<DataType>()>* = nullptr>
+  template <typename DataType>
+  constexpr static bool is_copy_block_supported()
+  {
+    // The shared-memory used for fixed-width types in the copy_block_partitions_impl function
+    // will be too large for any DataType greater than int64_t.
+    return is_fixed_width<DataType>() && (sizeof(DataType) <= sizeof(int64_t));
+  }
+
+  template <typename DataType, CUDF_ENABLE_IF(is_copy_block_supported<DataType>())>
   std::unique_ptr<column> operator()(column_view const& input,
                                      const size_type num_partitions,
                                      size_type const* row_partition_numbers,
@@ -416,7 +426,7 @@ struct copy_block_partitions_dispatcher {
     return std::make_unique<column>(input.type(), input.size(), std::move(output));
   }
 
-  template <typename DataType, std::enable_if_t<not is_fixed_width<DataType>()>* = nullptr>
+  template <typename DataType, CUDF_ENABLE_IF(not is_copy_block_supported<DataType>())>
   std::unique_ptr<column> operator()(column_view const& input,
                                      const size_type num_partitions,
                                      size_type const* row_partition_numbers,
@@ -713,7 +723,7 @@ struct dispatch_map_type {
 }  // namespace
 
 namespace detail {
-namespace local {
+namespace {
 template <template <typename> class hash_function>
 std::pair<std::unique_ptr<table>, std::vector<size_type>> hash_partition(
   table_view const& input,
@@ -738,7 +748,7 @@ std::pair<std::unique_ptr<table>, std::vector<size_type>> hash_partition(
       input, table_to_hash, num_partitions, seed, stream, mr);
   }
 }
-}  // namespace local
+}  // namespace
 
 std::pair<std::unique_ptr<table>, std::vector<size_type>> partition(
   table_view const& t,
@@ -779,10 +789,10 @@ std::pair<std::unique_ptr<table>, std::vector<size_type>> hash_partition(
         if (!is_numeric(input.column(column_id).type()))
           CUDF_FAIL("IdentityHash does not support this data type");
       }
-      return detail::local::hash_partition<detail::IdentityHash>(
+      return detail::hash_partition<detail::IdentityHash>(
         input, columns_to_hash, num_partitions, seed, stream, mr);
     case (hash_id::HASH_MURMUR3):
-      return detail::local::hash_partition<detail::MurmurHash3_32>(
+      return detail::hash_partition<detail::MurmurHash3_32>(
         input, columns_to_hash, num_partitions, seed, stream, mr);
     default: CUDF_FAIL("Unsupported hash function in hash_partition");
   }

cpp/tests/partitioning/hash_partition_test.cpp

@@ -349,21 +349,21 @@ TEST_F(HashPartition, FixedPointColumnsToHash)
 {
   fixed_width_column_wrapper<int32_t> to_hash({1});
   cudf::test::fixed_point_column_wrapper<int64_t> first_col({7}, numeric::scale_type{-1});
+  cudf::test::fixed_point_column_wrapper<__int128_t> second_col({77}, numeric::scale_type{0});
 
-  auto first_input = cudf::table_view({to_hash, first_col});
+  auto input = cudf::table_view({to_hash, first_col, second_col});
 
   auto columns_to_hash = std::vector<cudf::size_type>({0});
 
   cudf::size_type const num_partitions = 1;
-  auto [first_result, first_offsets] =
-    cudf::hash_partition(first_input, columns_to_hash, num_partitions);
+  auto [result, offsets] = cudf::hash_partition(input, columns_to_hash, num_partitions);
 
   // Expect offsets to be equal and num_partitions in length
-  EXPECT_EQ(static_cast<size_t>(num_partitions), first_offsets.size());
-
-  CUDF_TEST_EXPECT_COLUMNS_EQUAL(first_result->get_column(0).view(), first_input.column(0));
+  EXPECT_EQ(static_cast<size_t>(num_partitions), offsets.size());
 
-  CUDF_TEST_EXPECT_COLUMNS_EQUAL(first_result->get_column(1).view(), first_input.column(1));
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(0).view(), input.column(0));
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(1).view(), input.column(1));
+  CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->get_column(2).view(), input.column(2));
 }
 
 TEST_F(HashPartition, ListWithNulls)