Optimize partition

cpp/benchmarks/hashing/partition.cpp

@@ -8,6 +8,7 @@
 
 #include <cudf/copying.hpp>
 #include <cudf/partitioning.hpp>
+#include <cudf/reduction.hpp>
 #include <cudf/scalar/scalar.hpp>
 #include <cudf/table/table.hpp>
 #include <cudf/utilities/default_stream.hpp>
@@ -402,8 +403,142 @@ void bench_hash_partition_all_keys_stress(nvbench::state& state)
   run_hash_partition(state, *input, num_columns, num_partitions);
 }
 
+// Keep map generation outside the timed region and record output plus temporary allocation peaks.
+void run_map_partition(nvbench::state& state,
+                       cudf::table const& input,
+                       cudf::column_view map,
+                       cudf::size_type partitions)
+{
+  // Invalid partition ids have undefined behavior; check generated INT32 maps before timing.
+  {
+    auto const [minimum, maximum] = cudf::minmax(map);
+    CUDF_EXPECTS(
+      static_cast<cudf::numeric_scalar<cudf::size_type> const&>(*minimum).value() >= 0 &&
+        static_cast<cudf::numeric_scalar<cudf::size_type> const&>(*maximum).value() < partitions,
+      "Generated partition map is out of range");
+  }
+  state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().get()));
+  state.add_element_count(input.num_rows(), "rows");
+  state.add_buffer_size(input.alloc_size(), "input_size", "input_size");
+  auto const mem_stats_logger = cudf::memory_stats_logger();
+  state.exec(nvbench::exec_tag::sync, [&](nvbench::launch&) {
+    [[maybe_unused]] auto output = cudf::partition(input.view(), map, partitions);
+  });
+  state.add_buffer_size(
+    mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage");
+}
+
+// Measure partition-map materialization separately from hash computation.
+void bench_partition_map(nvbench::state& state)
+{
+  auto const rows       = get_size_type_axis(state, "num_rows");
+  auto const partitions = get_size_type_axis(state, "num_partitions");
+  auto const hot        = state.get_int64("hot_partition") != 0;
+  data_profile const profile =
+    data_profile_builder()
+      .distribution(cudf::type_id::INT32, distribution_id::UNIFORM, 0, hot ? 0 : partitions - 1)
+      .cardinality(0)
+      .avg_run_length(1)
+      .no_validity();
+  auto map          = create_random_column(cudf::type_id::INT32, row_count{rows}, profile);
+  auto const layout = state.get_string("layout");
+  std::unique_ptr<cudf::table> input;
+  if (layout == "int8") {
+    input = make_input_table(std::array{cudf::type_id::INT8}, {}, rows);
+  } else if (layout == "int8x2") {
+    input = make_input_table(std::array{cudf::type_id::INT8, cudf::type_id::INT8}, {}, rows);
+  } else if (layout == "int16") {
+    input = make_input_table(std::array{cudf::type_id::INT16}, {}, rows);
+  } else if (layout == "int32") {
+    input = make_input_table(std::array{cudf::type_id::INT32}, {}, rows);
+  } else if (layout == "int64") {
+    input = make_input_table(std::array{cudf::type_id::INT64}, {}, rows);
+  } else if (layout == "decimal128") {
+    input = make_input_table(std::array{cudf::type_id::DECIMAL128}, {}, rows);
+  } else if (layout == "strings10") {
+    input = make_input_table({}, std::vector<cudf::type_id>(10, cudf::type_id::STRING), rows);
+  } else if (layout == "mixed_nullable") {
+    input = make_input_table(std::array{cudf::type_id::INT64},
+                             std::array{cudf::type_id::STRING, cudf::type_id::INT32},
+                             rows,
+                             0.1,
+                             0.1);
+  } else {
+    state.skip("Unknown partition layout: " + layout);
+    return;
+  }
+  run_map_partition(state, *input, map->view(), partitions);
+}
+
+// Vary byte lengths independently from the 10-column scaling workload.
+void run_string_partition(nvbench::state& state,
+                          cudf::size_type rows,
+                          cudf::size_type min_length,
+                          cudf::size_type max_length,
+                          std::optional<double> null_probability)
+{
+  auto const partitions  = get_size_type_axis(state, "num_partitions");
+  auto const num_columns = get_size_type_axis(state, "num_columns");
+  data_profile const map_profile =
+    data_profile_builder()
+      .distribution(cudf::type_id::INT32, distribution_id::UNIFORM, 0, partitions - 1)
+      .cardinality(0)
+      .avg_run_length(1)
+      .no_validity();
+  data_profile const string_profile =
+    data_profile_builder()
+      .distribution(cudf::type_id::STRING, distribution_id::UNIFORM, min_length, max_length)
+      .cardinality(0)
+      .avg_run_length(1)
+      .null_probability(null_probability);
+  auto map = create_random_column(cudf::type_id::INT32, row_count{rows}, map_profile);
+  std::vector<std::unique_ptr<cudf::column>> columns;
+  for (cudf::size_type seed = 0; seed < num_columns; ++seed) {
+    columns.push_back(create_ascii_string_column(string_profile, rows, seed));
+  }
+  auto input = cudf::table(std::move(columns));
+  run_map_partition(state, input, map->view(), partitions);
+}
+
+void bench_partition_strings(nvbench::state& state)
+{
+  run_string_partition(state, get_size_type_axis(state, "num_rows"), 16, 48, std::nullopt);
+}
+
+void bench_partition_string_lengths(nvbench::state& state)
+{
+  auto const nulls = state.get_int64("null_percent");
+  run_string_partition(state,
+                       1'000'000,
+                       0,
+                       get_size_type_axis(state, "max_length"),
+                       nulls == 0 ? std::nullopt : std::optional<double>{nulls / 100.0});
+}
+
 }  // namespace
 
+NVBENCH_BENCH(bench_partition_strings)
+  .set_name("partition_strings")
+  .add_int64_axis("num_rows", {1'000'000, 10'000'000})
+  .add_int64_axis("num_columns", {2, 10})
+  .add_int64_axis("num_partitions", {100, 200});
+
+NVBENCH_BENCH(bench_partition_string_lengths)
+  .set_name("partition_string_lengths")
+  .add_int64_axis("max_length", {0, 8, 16, 24, 32, 128, 256})
+  .add_int64_axis("num_columns", {2, 10})
+  .add_int64_axis("null_percent", {0, 10})
+  .add_int64_axis("num_partitions", {100, 200});
+
+NVBENCH_BENCH(bench_partition_map)
+  .set_name("partition_map")
+  .add_int64_axis("num_rows", {1024, 1 << 20, 1 << 23, 10'000'000})
+  .add_int64_axis("num_partitions", {2, 100, 200, 256, 257})
+  .add_int64_axis("hot_partition", {0, 1})
+  .add_string_axis(
+    "layout",
+    {"int8", "int8x2", "int16", "int32", "int64", "decimal128", "strings10", "mixed_nullable"});
+
 NVBENCH_BENCH_TYPES(bench_hash_partition_analytic,
                     NVBENCH_TYPE_AXES(nvbench::type_list<std::int32_t, std::int64_t>))
   .set_name("hash_partition_analytic")

cpp/src/partitioning/partition_fixed_width.cuh

@@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+#include <cudf/column/column_factories.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
+#include <cudf/detail/utilities/vector_factories.hpp>
+#include <cudf/table/table.hpp>
+
+#include <cuda/std/bit>
+#include <cuda/std/cstring>
+#include <cuda/std/memory>
+
+namespace cudf::detail::partition_fixed {
+template <int Width>
+__device__ void copy_value(unsigned char* out, unsigned char const* in)
+{
+  cuda::std::memcpy(
+    cuda::std::assume_aligned<Width>(out), cuda::std::assume_aligned<Width>(in), Width);
+}
+template <int Width>
+struct single_output {
+  unsigned char const* input;
+  unsigned char* output;
+  __device__ void operator()(size_type source, size_type destination) const
+  {
+    copy_value<Width>(output + std::size_t(destination) * Width,
+                      input + std::size_t(source) * Width);
+  }
+};
+struct data_column {
+  unsigned char const* input;
+  unsigned char* output;
+  bitmask_type const* input_mask;
+  bitmask_type* output_mask;
+  size_type offset;
+  size_type width;
+};
+template <bool Nullable>
+struct table_output {
+  data_column const* columns;
+  size_type count;
+  __device__ void operator()(size_type source, size_type destination) const
+  {
+    for (size_type c = 0; c < count; ++c) {
+      auto const col = columns[c];
+      auto const row = source + col.offset;
+      bool valid     = true;
+      if constexpr (Nullable)
+        valid = !col.input_mask || (col.input_mask[row / 32] & (1u << (row % 32)));
+      if (valid) {
+        auto const* in = col.input + std::size_t(row) * col.width;
+        auto* out      = col.output + std::size_t(destination) * col.width;
+        switch (col.width) {
+          case 1: copy_value<1>(out, in); break;
+          case 2: copy_value<2>(out, in); break;
+          case 4: copy_value<4>(out, in); break;
+          case 8: copy_value<8>(out, in); break;
+          case 16: copy_value<16>(out, in); break;
+        }
+        if constexpr (Nullable)
+          if (col.output_mask) {
+            atomicOr(col.output_mask + destination / 32, 1u << (destination % 32));
+          }
+      }
+    }
+  }
+};
+
+template <typename Launch>
+std::pair<std::unique_ptr<table>, std::vector<size_type>> partition(
+  table_view const& input,
+  Launch&& launch,
+  cuda::stream_ref stream,
+  rmm::device_async_resource_ref mr)
+{
+  auto const rows = input.num_rows();
+  if (input.num_columns() == 1 && input.column(0).null_count() == 0) {
+    auto const& col = input.column(0);
+    auto output = make_fixed_width_column(col.type(), rows, mask_state::UNALLOCATED, stream, mr);
+    auto run    = [&]<int Width>() {
+      return launch(
+        single_output<Width>{col.head<unsigned char>() + std::size_t(col.offset()) * Width,
+                             output->mutable_view().head<unsigned char>()});
+    };
+    std::vector<size_type> offsets;
+    switch (cudf::size_of(col.type())) {
+      case 1: offsets = run.template operator()<1>(); break;
+      case 2: offsets = run.template operator()<2>(); break;
+      case 4: offsets = run.template operator()<4>(); break;
+      case 8: offsets = run.template operator()<8>(); break;
+      case 16: offsets = run.template operator()<16>(); break;
+      default: CUDF_FAIL("Unsupported fixed-width storage size");
+    }
+    std::vector<std::unique_ptr<column>> columns;
+    columns.push_back(std::move(output));
+    return {std::make_unique<table>(std::move(columns), rows), std::move(offsets)};
+  }
+  std::vector<std::unique_ptr<column>> columns;
+  columns.reserve(input.num_columns());
+  std::vector<data_column> descriptors;
+  bool nullable = false;
+  for (auto const& col : input) {
+    auto output =
+      make_fixed_width_column(col.type(),
+                              rows,
+                              col.null_count() ? mask_state::ALL_NULL : mask_state::UNALLOCATED,
+                              stream,
+                              mr);
+    output->set_null_count(col.null_count());
+    if (col.null_count() != rows) {
+      nullable |= col.has_nulls();
+      descriptors.push_back({col.head<unsigned char>(),
+                             output->mutable_view().head<unsigned char>(),
+                             col.has_nulls() ? col.null_mask() : nullptr,
+                             output->mutable_view().null_mask(),
+                             col.offset(),
+                             static_cast<size_type>(cudf::size_of(col.type()))});
+    }
+    columns.push_back(std::move(output));
+  }
+  auto device_columns =
+    cudf::detail::make_device_buffer(descriptors, stream, cudf::get_current_device_resource_ref());
+  std::vector<size_type> offsets;
+  if (!nullable)
+    offsets = launch(
+      table_output<false>{device_columns.data(), static_cast<size_type>(descriptors.size())});
+  else
+    offsets =
+      launch(table_output<true>{device_columns.data(), static_cast<size_type>(descriptors.size())});
+  return {std::make_unique<table>(std::move(columns), rows), std::move(offsets)};
+}
+}  // namespace cudf::detail::partition_fixed