diff --git a/cpp/benchmarks/hashing/partition.cpp b/cpp/benchmarks/hashing/partition.cpp index af75e5aacd7b..956df325e647 100644 --- a/cpp/benchmarks/hashing/partition.cpp +++ b/cpp/benchmarks/hashing/partition.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include #include @@ -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 const&>(*minimum).value() >= 0 && + static_cast 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 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(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 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> 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{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)) .set_name("hash_partition_analytic") diff --git a/cpp/src/partitioning/partition_fixed_width.cuh b/cpp/src/partitioning/partition_fixed_width.cuh new file mode 100644 index 000000000000..680dbc273c78 --- /dev/null +++ b/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 +#include +#include +#include + +#include +#include +#include + +namespace cudf::detail::partition_fixed { +template +__device__ void copy_value(unsigned char* out, unsigned char const* in) +{ + cuda::std::memcpy( + cuda::std::assume_aligned(out), cuda::std::assume_aligned(in), Width); +} +template +struct single_output { + unsigned char const* input; + unsigned char* output; + __device__ void operator()(size_type source, size_type destination) const + { + copy_value(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 +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 +std::pair, std::vector> 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 = [&]() { + return launch( + single_output{col.head() + std::size_t(col.offset()) * Width, + output->mutable_view().head()}); + }; + std::vector 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> columns; + columns.push_back(std::move(output)); + return {std::make_unique(std::move(columns), rows), std::move(offsets)}; + } + std::vector> columns; + columns.reserve(input.num_columns()); + std::vector 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(), + output->mutable_view().head(), + col.has_nulls() ? col.null_mask() : nullptr, + output->mutable_view().null_mask(), + col.offset(), + static_cast(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 offsets; + if (!nullable) + offsets = launch( + table_output{device_columns.data(), static_cast(descriptors.size())}); + else + offsets = + launch(table_output{device_columns.data(), static_cast(descriptors.size())}); + return {std::make_unique
(std::move(columns), rows), std::move(offsets)}; +} +} // namespace cudf::detail::partition_fixed diff --git a/cpp/src/partitioning/partition_identity.cuh b/cpp/src/partitioning/partition_identity.cuh new file mode 100644 index 000000000000..d1404342b000 --- /dev/null +++ b/cpp/src/partitioning/partition_identity.cuh @@ -0,0 +1,199 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace cudf::detail::partition_identity { +struct string_source { + void const* offsets; + size_type row_offset; + bool wide; +}; +struct string_bounds { + int64_t first; + int64_t bytes; +}; + +template +CUDF_KERNEL void read_bounds(string_source const* sources, + size_type count, + size_type rows, + string_bounds* output) +{ + auto const i = static_cast(blockIdx.x) * BlockSize + threadIdx.x; + if (i >= static_cast(count)) return; + auto const source = sources[i]; + auto const first = source.wide ? static_cast(source.offsets)[source.row_offset] + : static_cast(source.offsets)[source.row_offset]; + auto const last = source.wide + ? static_cast(source.offsets)[source.row_offset + rows] + : static_cast(source.offsets)[source.row_offset + rows]; + output[i] = {first, last - first}; +} + +/** Cached host metadata proves the payload is invariant under a permutation. */ +inline bool is_uniform(column_view const& input) +{ + bool const flat = is_fixed_width(input.type()) || input.type().id() == type_id::STRING; + return flat && (input.null_count() == input.size() || + (input.type().id() == type_id::STRING && input.null_count() == 0 && + input.head() == nullptr)); +} + +inline std::unique_ptr uniform_column(column_view const& input, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + auto const rows = input.size(); + if (input.type().id() != type_id::STRING) + return make_fixed_width_column(input.type(), rows, mask_state::ALL_NULL, stream, mr); + auto offsets = + make_numeric_column(data_type{type_id::INT32}, rows + 1, mask_state::UNALLOCATED, stream, mr); + CUDF_CUDA_TRY(cudaMemsetAsync(offsets->mutable_view().data(), + 0, + (static_cast(rows) + 1) * sizeof(size_type), + stream.get())); + return make_strings_column( + rows, + std::move(offsets), + rmm::device_buffer(0, stream, mr), + input.null_count(), + cudf::detail::create_null_mask( + rows, input.null_count() ? mask_state::ALL_NULL : mask_state::UNALLOCATED, stream, mr)); +} + +/** + * Identity materialization for CPU-known one-row/one-partition cases. + * + * String byte bounds are read only for output allocation and offset rebasing. + * They never choose the partition algorithm. Nullable strings and nested values + * use the normal gather with an implicit identity iterator; no N-entry map. + */ +inline std::unique_ptr
copy(table_view const& input, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + auto const rows = input.num_rows(); + std::vector sources; + for (auto const& col : input) + if (col.type().id() == type_id::STRING && col.null_count() == 0) { + auto const& offsets = col.child(strings_column_view::offsets_column_index); + bool const wide = offsets.type().id() == type_id::INT64; + auto const* data = wide ? static_cast(offsets.data()) + : static_cast(offsets.data()); + sources.push_back({data, col.offset(), wide}); + } + auto bounds = [&] { + if (sources.empty()) return std::vector{}; + auto descriptors = + cudf::detail::make_device_buffer(sources, stream, cudf::get_current_device_resource_ref()); + rmm::device_uvector device_bounds(sources.size(), stream); + read_bounds<256><<<(sources.size() + 255) / 256, 256, 0, stream.get()>>>( + descriptors.data(), sources.size(), rows, device_bounds.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + return cudf::detail::make_std_vector(device_bounds, stream); + }(); + auto const threshold = strings::get_offset64_threshold(); + for (auto const& b : bounds) + CUDF_EXPECTS(b.bytes < threshold || strings::is_large_strings_enabled(), + "Size of output exceeds the column size limit", + std::overflow_error); + + std::vector> columns; + columns.reserve(input.num_columns()); + std::size_t string_index = 0; + auto const indices = cuda::counting_iterator{0}; + for (auto const& col : input) { + if (is_uniform(col)) { + if (col.type().id() == type_id::STRING && col.null_count() == 0) { ++string_index; } + columns.push_back(uniform_column(col, stream, mr)); + continue; + } + if (is_fixed_width(col.type())) { + auto const width = cudf::size_of(col.type()); + rmm::device_buffer data(col.head() + static_cast(col.offset()) * width, + static_cast(rows) * width, + stream, + mr); + auto mask = col.null_count() > 0 + ? cudf::detail::copy_bitmask(col, stream, mr) + : cudf::detail::create_null_mask(0, mask_state::UNALLOCATED, stream, mr); + columns.push_back(std::make_unique( + col.type(), rows, std::move(data), std::move(mask), col.null_count())); + } else if (col.type().id() == type_id::STRING && col.null_count() == 0) { + auto const b = bounds[string_index++]; + auto const& source_offsets = col.child(strings_column_view::offsets_column_index); + auto const type = data_type{b.bytes >= threshold ? type_id::INT64 : type_id::INT32}; + std::unique_ptr offsets; + if (type == source_offsets.type() && b.first == 0) { + auto const width = cudf::size_of(type); + rmm::device_buffer data( + source_offsets.head() + + (static_cast(source_offsets.offset()) + col.offset()) * width, + (static_cast(rows) + 1) * width, + stream, + mr); + offsets = std::make_unique( + type, + rows + 1, + std::move(data), + cudf::detail::create_null_mask(0, mask_state::UNALLOCATED, stream, mr), + 0); + } else { + offsets = make_numeric_column(type, rows + 1, mask_state::UNALLOCATED, stream, mr); + auto src = + cudf::detail::offsetalator_factory::make_input_iterator(source_offsets, col.offset()); + auto dst = + cudf::detail::offsetalator_factory::make_output_iterator(offsets->mutable_view()); + thrust::transform( + rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + indices, + indices + rows + 1, + dst, + [src, first = b.first] __device__(size_type i) { return src[i] - first; }); + } + auto const* chars_begin = b.bytes ? col.head() + b.first : nullptr; + rmm::device_buffer chars(chars_begin, b.bytes, stream, mr); + columns.push_back(make_strings_column( + rows, + std::move(offsets), + std::move(chars), + 0, + cudf::detail::create_null_mask(0, mask_state::UNALLOCATED, stream, mr))); + } else if (col.type().id() == type_id::STRING) { + auto gathered = strings::detail::gather( + strings_column_view{col}, indices, indices + rows, stream, cudf::memory_resources{mr}); + gathered->set_null_mask(cudf::detail::copy_bitmask(col, stream, mr), col.null_count()); + columns.push_back(std::move(gathered)); + } else { + auto gathered = cudf::detail::gather(table_view{{col}}, + indices, + indices + rows, + out_of_bounds_policy::DONT_CHECK, + stream, + cudf::memory_resources{mr}); + auto owned = gathered->release(); + columns.push_back(std::move(owned.front())); + } + } + return std::make_unique
(std::move(columns), rows); +} +} // namespace cudf::detail::partition_identity diff --git a/cpp/src/partitioning/partition_map_plan.cuh b/cpp/src/partitioning/partition_map_plan.cuh new file mode 100644 index 000000000000..7b209ba79fe1 --- /dev/null +++ b/cpp/src/partitioning/partition_map_plan.cuh @@ -0,0 +1,349 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cudf::detail::partition_map_plan { +enum class algorithm { saved, recompute, super2, reserve, binary }; +constexpr int threads = 256, tile_rows = 1024, warps = 8, bins = 256; + +// The same contiguous warp ordering as PR1: four 32-row rounds per warp. +// Histogram-only instantiation discards ranks; rank instantiation retains them in registers. +template +__device__ void local_tile(MapType const* map, + int n, + int start, + int p, + int (&hist)[warps][bins], + int (&rank)[4], + int (&keys)[4]) +{ + int lane = threadIdx.x % 32, w = threadIdx.x / 32; + for (int q = lane; q < p; q += 32) + hist[w][q] = 0; + __syncwarp(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int row = start + w * 128 + j * 32 + lane; + bool valid = row < n; + unsigned active = __ballot_sync(0xffffffffu, valid); + if (valid) { + int q = map[row]; + unsigned peers = __match_any_sync(active, q); + int base = hist[w][q]; + if constexpr (Rank) { + keys[j] = q; + rank[j] = base + cuda::std::popcount(peers & ((1u << lane) - 1u)); + } + __syncwarp(active); + if (lane == __ffs(peers) - 1) hist[w][q] = base + cuda::std::popcount(peers); + } + __syncwarp(); + } + __syncthreads(); +} + +template +CUDF_KERNEL void count_tiles(MapType const* map, int n, int p, uint16_t* saved, int32_t* counts) +{ + static_assert(!PackBin || tile_rows <= 32'768, "Binary packed ranks must fit in uint16_t"); + __shared__ int hist[warps][bins]; + int total = 0; + for (int t = 0; t < Tiles; ++t) { + int rank[4], keys[4]; + int start = (blockIdx.x * Tiles + t) * tile_rows; + local_tile(map, n, start, p, hist, rank, keys); + if (threadIdx.x < p) { + int q = threadIdx.x, sum = 0; +#pragma unroll + for (int w = 0; w < warps; ++w) { + int c = hist[w][q]; + if constexpr (SaveRanks) hist[w][q] = sum; + sum += c; + } + total += sum; + } + if constexpr (SaveRanks) { + __syncthreads(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int i = start + (threadIdx.x / 32) * 128 + j * 32 + threadIdx.x % 32; + if (i < n) { + auto const local = hist[threadIdx.x / 32][keys[j]] + rank[j]; + // PackBin is used by the binary STRING plan: its 10-bit tile rank and + // one-bit bin fit in uint16_t, avoiding a map read during byte copying. + saved[i] = local + (PackBin ? keys[j] * tile_rows : 0); + } + } + } + __syncthreads(); + } + if (threadIdx.x < p) counts[size_t(threadIdx.x) * gridDim.x + blockIdx.x] = total; +} + +CUDF_KERNEL void get_offsets(int32_t const* counts, int blocks, int p, int n, int32_t* offsets) +{ + if (threadIdx.x < p) offsets[threadIdx.x] = counts[size_t(threadIdx.x) * blocks]; + if (threadIdx.x == 0) offsets[p] = n; +} + +template +CUDF_KERNEL void scatter_saved(MapType const* map, + uint16_t const* ranks, + int32_t const* offsets, + int n, + int blocks, + int p, + Sink sink) +{ + __shared__ int base[bins]; + if (threadIdx.x < p) base[threadIdx.x] = offsets[size_t(threadIdx.x) * blocks + blockIdx.x]; + __syncthreads(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int i = blockIdx.x * tile_rows + j * threads + threadIdx.x; + if (i < n) sink(i, base[map[i]] + ranks[i]); + } +} + +template +__device__ void emit_reordered(int const (&keys)[4], + int const (&rank)[4], + int (&hist)[warps][bins], + int const* base, + int next, + int p, + int n, + int start, + Sink sink) +{ + __shared__ typename cub::BlockScan::TempStorage scan_temp; + __shared__ int bucket_begin[bins], source_rows[tile_rows], destinations[tile_rows]; + int count = threadIdx.x < p ? next - base[threadIdx.x] : 0, prefix = 0; + cub::BlockScan(scan_temp).ExclusiveSum(count, prefix); + bucket_begin[threadIdx.x] = prefix; + __syncthreads(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int row = start + (threadIdx.x / 32) * 128 + j * 32 + threadIdx.x % 32; + if (row < n) { + int q = keys[j], local = hist[threadIdx.x / 32][q] + rank[j], + position = bucket_begin[q] + local; + source_rows[position] = row; + destinations[position] = base[q] + local; + } + } + __syncthreads(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int position = threadIdx.x + j * threads; + if (position < min(tile_rows, n - start)) sink(source_rows[position], destinations[position]); + } + // The caller synchronizes before reusing the shared storage. +} + +template +CUDF_KERNEL void scatter_recompute(MapType const* map, Offsets offsets, int n, int p, Sink sink) +{ + __shared__ int hist[warps][bins], base[bins]; + if (threadIdx.x < p) base[threadIdx.x] = offsets[size_t(threadIdx.x) * gridDim.x + blockIdx.x]; + __syncthreads(); + for (int t = 0; t < Tiles; ++t) { + int rank[4], keys[4]; + int start = (blockIdx.x * Tiles + t) * tile_rows; + local_tile(map, n, start, p, hist, rank, keys); + int next = 0; + if (threadIdx.x < p) { + int q = threadIdx.x, sum = 0; +#pragma unroll + for (int w = 0; w < warps; ++w) { + int c = hist[w][q]; + hist[w][q] = sum; + sum += c; + } + next = base[q] + sum; + } + __syncthreads(); + if constexpr (Reorder) + emit_reordered(keys, rank, hist, base, next, p, n, start, sink); + else { +#pragma unroll + for (int j = 0; j < 4; ++j) { + int i = start + (threadIdx.x / 32) * 128 + j * 32 + threadIdx.x % 32; + if (i < n) { + int q = keys[j]; + sink(i, base[q] + hist[threadIdx.x / 32][q] + rank[j]); + } + } + } + __syncthreads(); + if (threadIdx.x < p) base[threadIdx.x] = next; + __syncthreads(); + } +} + +template +CUDF_KERNEL void scatter_reserve(MapType const* map, int n, int p, int32_t* cursors, Sink sink) +{ + __shared__ int hist[warps][bins], base[bins]; + int rank[4], keys[4]; + int start = blockIdx.x * tile_rows; + local_tile(map, n, start, p, hist, rank, keys); + if (threadIdx.x < p) { + int q = threadIdx.x, sum = 0; +#pragma unroll + for (int w = 0; w < warps; ++w) { + int c = hist[w][q]; + hist[w][q] = sum; + sum += c; + } + if (sum) base[q] = atomicAdd(cursors + q, sum); + } + __syncthreads(); +#pragma unroll + for (int j = 0; j < 4; ++j) { + int i = start + (threadIdx.x / 32) * 128 + j * 32 + threadIdx.x % 32; + if (i < n) sink(i, base[keys[j]] + hist[threadIdx.x / 32][keys[j]] + rank[j]); + } +} + +template +struct WriteProxy { + Emit sink; + int destination; + __device__ void operator=(int source) const { sink(source, destination); } +}; +template +struct MakeWriteProxy { + Emit sink; + __host__ __device__ WriteProxy operator()(int destination) const + { + return {sink, destination}; + } +}; +template +struct IsZero { + Map const* map; + __device__ bool operator()(int row) const { return map[row] == 0; } +}; + +template +std::vector binary(Map const* map, int n, Emit sink, cuda::stream_ref stream) +{ + rmm::device_uvector selected(1, stream); + auto source = cuda::counting_iterator{0}; + auto output = cuda::transform_iterator(source, MakeWriteProxy{sink}); + size_t bytes = 0; + CUDF_CUDA_TRY(cub::DevicePartition::If( + nullptr, bytes, source, output, selected.data(), n, IsZero{map}, stream.get())); + rmm::device_buffer temporary(bytes, stream); + CUDF_CUDA_TRY(cub::DevicePartition::If( + temporary.data(), bytes, source, output, selected.data(), n, IsZero{map}, stream.get())); + auto const count = cudf::detail::make_std_vector(selected, stream)[0]; + return {0, count, n}; +} + +inline void scan(int32_t* ptr, size_t count, cuda::stream_ref stream) +{ + // Same Thrust in-place scan policy as PR1. + thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + ptr, + ptr + count, + ptr); +} + +template +inline void global_hist(MapType const* map, int n, int p, int32_t* offsets, cuda::stream_ref stream) +{ + size_t bytes = 0; + CUDF_CUDA_TRY(cudaMemsetAsync(offsets + p, 0, sizeof(int), stream.get())); + CUDF_CUDA_TRY(cub::DeviceHistogram::HistogramEven( + nullptr, bytes, map, offsets, p + 1, 0, p, n, stream.get())); + rmm::device_buffer tmp(bytes, stream); + CUDF_CUDA_TRY(cub::DeviceHistogram::HistogramEven( + tmp.data(), bytes, map, offsets, p + 1, 0, p, n, stream.get())); + scan(offsets, p + 1, stream); +} + +template +void recompute( + MapType const* map, int n, int p, Sink sink, int32_t* offsets, cuda::stream_ref stream) +{ + int blocks = (n - 1) / (tile_rows * Tiles) + 1; + rmm::device_uvector counts(size_t(p) * blocks, stream); + count_tiles + <<>>(map, n, p, nullptr, counts.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + scan(counts.data(), counts.size(), stream); + get_offsets<<<1, threads, 0, stream.get()>>>(counts.data(), blocks, p, n, offsets); + scatter_recompute + <<>>(map, counts.data(), n, p, sink); + CUDF_CUDA_TRY(cudaPeekAtLastError()); +} + +template +std::vector run(MapType const* map, + size_type n, + size_type p, + algorithm method, + Sink sink, + bool reorder, + cuda::stream_ref stream) +{ + CUDF_EXPECTS(n > 0 && p > 0 && p <= bins, "Invalid block partition dimensions"); + if (method == algorithm::binary) { + CUDF_EXPECTS(p == 2, "Binary plan requires two partitions"); + return binary(map, n, sink, stream); + } + int const blocks = (n - 1) / tile_rows + 1; + rmm::device_uvector offsets(p + 1, stream); + if (reorder) { + if (method == algorithm::recompute) + recompute<1, true>(map, n, p, sink, offsets.data(), stream); + else if (method == algorithm::super2) + recompute<2, true>(map, n, p, sink, offsets.data(), stream); + else + CUDF_FAIL("Unsupported reordered partition plan"); + } else if (method == algorithm::saved) { + rmm::device_uvector ranks(n, stream); + rmm::device_uvector counts(size_t(p) * blocks, stream); + count_tiles<1, true> + <<>>(map, n, p, ranks.data(), counts.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + scan(counts.data(), counts.size(), stream); + get_offsets<<<1, threads, 0, stream.get()>>>(counts.data(), blocks, p, n, offsets.data()); + scatter_saved<<>>( + map, ranks.data(), counts.data(), n, blocks, p, sink); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + } else if (method == algorithm::recompute) + recompute<1>(map, n, p, sink, offsets.data(), stream); + else if (method == algorithm::super2) + recompute<2>(map, n, p, sink, offsets.data(), stream); + else if (method == algorithm::reserve) { + global_hist(map, n, p, offsets.data(), stream); + rmm::device_uvector cursors(p + 1, stream); + CUDF_CUDA_TRY(cudf::detail::memcpy_async( + cursors.data(), offsets.data(), (p + 1) * sizeof(int32_t), stream)); + scatter_reserve<<>>(map, n, p, cursors.data(), sink); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + } else + CUDF_FAIL("Unknown block partition plan"); + return cudf::detail::make_std_vector(offsets, stream); +} +} // namespace cudf::detail::partition_map_plan diff --git a/cpp/src/partitioning/partition_string_copy.cuh b/cpp/src/partitioning/partition_string_copy.cuh new file mode 100644 index 000000000000..a66f22aaad47 --- /dev/null +++ b/cpp/src/partitioning/partition_string_copy.cuh @@ -0,0 +1,111 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once +#include + +#include +#include +namespace cudf::detail { +struct partition_string_column { + size_type const* input_offsets; + char const* input_chars; + size_type* output_offsets; + char* output_chars; +}; +namespace partition_strings_plan { +// All addresses loaded by load_uint4 stay within this string, including its +// unaligned loads. Destination stores are aligned and the ends use byte copies. +template +__device__ void copy_vector( + partition_string_column col, size_type source, size_type target, size_type length, int lane) +{ + int64_t const first = (int64_t{target} + 4 + 15) / 16 * 16; + int64_t const end = (int64_t{target} + length - 4) / 16 * 16; + for (auto byte = first + lane * 16; byte < end; byte += Lanes * 16) { + auto const value = strings::detail::load_uint4(col.input_chars + source + byte - target); + cuda::std::memcpy(cuda::std::assume_aligned<16>(col.output_chars + byte), &value, 16); + } + for (size_type byte = lane; int64_t{target} + byte < first; byte += Lanes) { + col.output_chars[int64_t{target} + byte] = col.input_chars[int64_t{source} + byte]; + } + for (auto byte = end + lane; byte < int64_t{target} + length; byte += Lanes) { + col.output_chars[byte] = col.input_chars[int64_t{source} + byte - target]; + } +} + +/** All warp lanes participate, including groups beyond the last row. */ +template +__device__ inline void copy_group(partition_string_column col, + std::size_t source_row, + size_type destination, + size_type rows) +{ + constexpr int lanes = 4; + auto const lane = threadIdx.x % lanes; + size_type begin = 0, length = 0, target = 0; + if (source_row < static_cast(rows)) { + begin = col.input_offsets[source_row]; + length = col.input_offsets[source_row + 1] - begin; + target = col.output_offsets[destination]; + } + // This device-side split balances individual strings. It does not change the + // CPU-selected partition method or require a size/skew sampling pass. + auto long_rows = __ballot_sync(0xffffffffu, lane == 0 && length > 128); + if constexpr (RebalanceMedium) { + auto medium_rows = __ballot_sync(0xffffffffu, lane == 0 && length > 8 && length <= 128); + if (length <= 8) { + for (size_type byte = lane; byte < length; byte += lanes) { + col.output_chars[int64_t{target} + byte] = col.input_chars[int64_t{begin} + byte]; + } + } + // Two 16-lane groups copy medium strings; every lane also participates + // when the last queue entry occupies only one half warp. + auto const warp_lane = threadIdx.x % 32; + while (medium_rows) { + auto const first = __ffs(medium_rows) - 1; + medium_rows &= medium_rows - 1; + bool const has_second = medium_rows != 0; + auto const second = has_second ? __ffs(medium_rows) - 1 : first; + if (has_second) { medium_rows &= medium_rows - 1; } + auto const owner = warp_lane < 16 ? first : second; + auto const source = __shfl_sync(0xffffffffu, begin, owner); + auto const output = __shfl_sync(0xffffffffu, target, owner); + auto count = __shfl_sync(0xffffffffu, length, owner); + if (warp_lane >= 16 && !has_second) { count = 0; } + for (size_type byte = warp_lane % 16; byte < count; byte += 16) { + col.output_chars[int64_t{output} + byte] = col.input_chars[int64_t{source} + byte]; + } + } + } else { + if (length <= 128) { + for (size_type byte = lane; byte < length; byte += lanes) { + col.output_chars[int64_t{target} + byte] = col.input_chars[int64_t{begin} + byte]; + } + } + } + while (long_rows) { + auto const leader = __ffs(long_rows) - 1; + auto const source = __shfl_sync(0xffffffffu, begin, leader); + auto const output = __shfl_sync(0xffffffffu, target, leader); + auto const count = __shfl_sync(0xffffffffu, length, leader); + copy_vector<32>(col, source, output, count, threadIdx.x % 32); + long_rows &= long_rows - 1; + } +} + +template +CUDF_KERNEL void copy_characters(partition_string_column const* columns, + size_type const* scatter_map, + size_type rows) +{ + auto const col = columns[blockIdx.y]; + for (auto base = std::size_t{blockIdx.x} * 64; base < static_cast(rows); + base += std::size_t{gridDim.x} * 64) { + auto const source = base + threadIdx.x / 4; + auto const target = source < static_cast(rows) ? scatter_map[source] : 0; + copy_group(col, source, target, rows); + } +} + +} // namespace partition_strings_plan +} // namespace cudf::detail diff --git a/cpp/src/partitioning/partition_string_plan.cuh b/cpp/src/partitioning/partition_string_plan.cuh new file mode 100644 index 000000000000..f7bffd49ab6b --- /dev/null +++ b/cpp/src/partitioning/partition_string_plan.cuh @@ -0,0 +1,328 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 +#pragma once + +#include "partition_map_plan.cuh" +#include "partition_strings.cuh" + +#include + +#include +#include + +#include + +namespace cudf::detail::partition_strings_plan { + +struct sizes_output { + partition_string_column const* columns; + size_type count; + size_type* scatter; + __device__ void operator()(size_type source, size_type destination) const + { + if (scatter) { scatter[source] = destination; } + for (size_type c = 0; c < count; ++c) { + auto const col = columns[c]; + col.output_offsets[destination] = col.input_offsets[source + 1] - col.input_offsets[source]; + } + } +}; + +/** Recompute into an on-chip tile map instead of storing N global indices. */ +template +CUDF_KERNEL void copy_recomputed(MapType const* map, + size_type const* prefixes, + size_type rows, + size_type partitions, + partition_string_column const* columns, + size_type count) +{ + namespace plans = partition_map_plan; + __shared__ int histogram[plans::warps][plans::bins], base[plans::bins]; + __shared__ int destinations[plans::tile_rows]; + if (threadIdx.x < partitions) { + base[threadIdx.x] = prefixes[std::size_t{threadIdx.x} * gridDim.x + blockIdx.x]; + } + __syncthreads(); + int ranks[4], keys[4]; + int const start = blockIdx.x * plans::tile_rows; + plans::local_tile(map, rows, start, partitions, histogram, ranks, keys); + if (threadIdx.x < partitions) { + int sum = 0; + for (int w = 0; w < plans::warps; ++w) { + auto const value = histogram[w][threadIdx.x]; + histogram[w][threadIdx.x] = sum; + sum += value; + } + } + __syncthreads(); +#pragma unroll + for (int item = 0; item < 4; ++item) { + int const local = (threadIdx.x / 32) * 128 + item * 32 + threadIdx.x % 32; + destinations[local] = + start + local < rows + ? base[keys[item]] + histogram[threadIdx.x / 32][keys[item]] + ranks[item] + : 0; + } + __syncthreads(); + // Shared targets let register ranks/keys die before copying variable-length + // data. Keep this loop rolled to avoid retaining all targets in registers. + for (size_type c = 0; c < count; ++c) { + auto const col = columns[c]; +#pragma unroll 1 + for (int batch = 0; batch < plans::tile_rows / 64; ++batch) { + int const local = batch * 64 + threadIdx.x / 4; + copy_group(col, start + local, destinations[local], rows); + } + } +} + +CUDF_KERNEL void scatter_ranked_sizes(uint16_t const* ranks, + size_type const* prefixes, + size_type rows, + size_type tiles, + partition_string_column const* columns, + size_type count) +{ + auto const zero = prefixes[blockIdx.x]; + auto const one = prefixes[tiles + blockIdx.x]; +#pragma unroll + for (int item = 0; item < 4; ++item) { + auto const source = + std::size_t{blockIdx.x} * partition_map_plan::tile_rows + item * blockDim.x + threadIdx.x; + if (source < static_cast(rows)) { + auto const packed = ranks[source]; + auto const rank = packed % partition_map_plan::tile_rows; + auto const target = (packed < partition_map_plan::tile_rows ? zero : one) + rank; + sizes_output{columns, count, nullptr}(source, target); + } + } +} + +/** Copy from packed binary-bin/tile-rank pairs without a full-width row permutation. */ +CUDF_KERNEL void copy_ranked(uint16_t const* ranks, + size_type const* prefixes, + size_type rows, + size_type tiles, + partition_string_column const* columns) +{ + static_assert(partition_map_plan::tile_rows % 64 == 0); + auto const col = columns[blockIdx.y]; + for (auto base = std::size_t{blockIdx.x} * 64; base < static_cast(rows); + base += std::size_t{gridDim.x} * 64) { + // Each 64-row batch stays inside one tile. Load both bin bases independently + // of the packed rank so prefix loads can overlap the rank load. + auto const tile = base / partition_map_plan::tile_rows; + auto const zero = prefixes[tile]; + auto const one = prefixes[tiles + tile]; + auto const source = base + threadIdx.x / 4; + size_type target = 0; + if (source < static_cast(rows)) { + auto const packed = ranks[source]; + auto const rank = packed % partition_map_plan::tile_rows; + target = (packed < partition_map_plan::tile_rows ? zero : one) + rank; + } + copy_group(col, source, target, rows); + } +} + +class output_buffers { + public: + std::vector columns; + std::vector> offsets; + std::vector characters; + cuda::device_buffer device_columns; + rmm::device_buffer scan_storage; + size_type rows; + + output_buffers(table_view const& input, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) + : device_columns(stream, cudf::get_current_device_resource_ref(), 0, cuda::no_init), + scan_storage(0, stream), + rows(input.num_rows()) + { + for (auto const& col : input) { + columns.push_back({strings_column_view(col).offsets().data() + col.offset(), + col.head(), + nullptr, + nullptr}); + } + auto const temporary_mr = cudf::get_current_device_resource_ref(); + device_columns = cudf::detail::make_device_buffer(columns, stream, temporary_mr); + rmm::device_uvector device_sizes(columns.size(), stream); + partition_string_sizes<<<(columns.size() + 255) / 256, 256, 0, stream.get()>>>( + device_columns.data(), columns.size(), rows, device_sizes.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + auto const sizes = cudf::detail::make_pinned_vector(device_sizes, stream); + // Mandatory allocation bounds also enforce the standard output offset-type + // limit. They do not choose the map algorithm or byte-copy strategy. + if (*std::max_element(sizes.begin(), sizes.end()) >= strings::get_offset64_threshold()) { + return; + } + for (std::size_t c = 0; c < columns.size(); ++c) { + offsets.push_back(make_numeric_column( + data_type{type_id::INT32}, rows + 1, mask_state::UNALLOCATED, stream, mr)); + characters.emplace_back(sizes[c], stream, mr); + columns[c].output_offsets = offsets.back()->mutable_view().data(); + columns[c].output_chars = static_cast(characters.back().data()); + CUDF_CUDA_TRY( + cudaMemsetAsync(columns[c].output_offsets + rows, 0, sizeof(size_type), stream.get())); + } + device_columns = cudf::detail::make_device_buffer(columns, stream, temporary_mr); + std::size_t bytes = 0; + CUDF_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + bytes, + static_cast(nullptr), + static_cast(nullptr), + rows + 1, + stream.get())); + scan_storage = rmm::device_buffer(bytes, stream); + } + + bool admitted() const { return offsets.size() == columns.size(); } + + void scan(cuda::stream_ref stream) + { + for (auto const& col : columns) { + auto bytes = scan_storage.size(); + CUDF_CUDA_TRY(cub::DeviceScan::ExclusiveSum(scan_storage.data(), + bytes, + col.output_offsets, + col.output_offsets, + rows + 1, + stream.get())); + } + } + + std::unique_ptr
release(cuda::stream_ref stream, rmm::device_async_resource_ref mr) + { + std::vector> output; + for (std::size_t c = 0; c < columns.size(); ++c) { + output.push_back(make_strings_column( + rows, + std::move(offsets[c]), + std::move(characters[c]), + 0, + cudf::detail::create_null_mask(0, mask_state::UNALLOCATED, stream, mr))); + } + return std::make_unique
(std::move(output), rows); + } +}; + +template +std::optional, std::vector>> partition( + table_view const& input, + column_view const& map, + size_type partitions, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + if (partitions < 2 || partitions > 256 || (partitions != 2 && input.num_columns() < 2) || + input.num_rows() < 1024 || input.num_columns() == 0 || input.num_columns() > 65'535 || + !std::all_of(input.begin(), input.end(), [](auto const& col) { + return col.type().id() == type_id::STRING && col.null_count() == 0 && + strings_column_view(col).offsets().type().id() == type_id::INT32; + })) { + return std::nullopt; + } + namespace plans = partition_map_plan; + auto const rows = input.num_rows(); + // Preserve fused materialization for multi-column strings at larger partition + // counts. With two partitions, small schemas can eliminate one or both maps. + if (input.num_columns() >= 4 || partitions != 2) { + rmm::device_uvector gather_map(rows, stream); + auto offsets = plans::run( + map.begin(), + rows, + partitions, + partitions == 2 ? plans::algorithm::binary : plans::algorithm::super2, + [indices = gather_map.data()] __device__(size_type source, size_type destination) { + indices[destination] = source; + }, + false, + stream); + if (auto output = gather_partition_strings( + input, device_span{gather_map.data(), gather_map.size()}, stream, mr)) { + return std::pair{std::move(output), std::move(offsets)}; + } + return std::nullopt; + } + output_buffers output(input, stream, mr); + if (!output.admitted()) { return std::nullopt; } + std::vector offsets; + constexpr size_type large_input_crossover = 4'000'000; + if (rows >= large_input_crossover) { + auto const blocks = (rows - 1) / plans::tile_rows + 1; + rmm::device_uvector prefixes(static_cast(partitions) * blocks, stream); + rmm::device_uvector boundaries(partitions + 1, stream); + // A single column keeps the mapless plan. For multiple columns, binary bins + // and tile-local ranks share two bytes per row to avoid repeated map reads. + bool const compact = input.num_columns() > 1; + std::optional> ranks; + if (compact) { + ranks.emplace(rows, stream); + plans::count_tiles<1, true, true><<>>( + map.begin(), rows, partitions, ranks->data(), prefixes.data()); + } else { + plans::count_tiles<1, false><<>>( + map.begin(), rows, partitions, nullptr, prefixes.data()); + } + plans::scan(prefixes.data(), prefixes.size(), stream); + plans::get_offsets<<<1, 256, 0, stream.get()>>>( + prefixes.data(), blocks, partitions, rows, boundaries.data()); + if (compact) { + scatter_ranked_sizes<<>>(ranks->data(), + prefixes.data(), + rows, + blocks, + output.device_columns.data(), + input.num_columns()); + } else { + plans::scatter_recompute<1><<>>( + map.begin(), + prefixes.data(), + rows, + partitions, + sizes_output{output.device_columns.data(), input.num_columns(), nullptr}); + } + output.scan(stream); + if (compact) { + auto const copy_blocks = + std::min(65'535, (static_cast(rows) + 63) / 64); + copy_ranked<<>>( + ranks->data(), prefixes.data(), rows, blocks, output.device_columns.data()); + } else { + copy_recomputed<<>>(map.begin(), + prefixes.data(), + rows, + partitions, + output.device_columns.data(), + input.num_columns()); + } + CUDF_CUDA_TRY(cudaPeekAtLastError()); + offsets = cudf::detail::make_std_vector(boundaries, stream); + } else { + rmm::device_uvector scatter_map(rows, stream); + offsets = plans::run( + map.begin(), + rows, + partitions, + plans::algorithm::binary, + sizes_output{output.device_columns.data(), input.num_columns(), scatter_map.data()}, + false, + stream); + output.scan(stream); + auto const blocks = std::min(65'535, (static_cast(rows) + 63) / 64); + if (input.num_columns() == 1) { + copy_characters<<>>( + output.device_columns.data(), scatter_map.data(), rows); + } else { + copy_characters<<>>( + output.device_columns.data(), scatter_map.data(), rows); + } + CUDF_CUDA_TRY(cudaPeekAtLastError()); + } + return std::pair{output.release(stream, mr), std::move(offsets)}; +} +} // namespace cudf::detail::partition_strings_plan diff --git a/cpp/src/partitioning/partition_strings.cuh b/cpp/src/partitioning/partition_strings.cuh new file mode 100644 index 000000000000..849ee993d98e --- /dev/null +++ b/cpp/src/partitioning/partition_strings.cuh @@ -0,0 +1,179 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "partition_string_copy.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cudf::detail { + +CUDF_KERNEL void invert_partition_map(size_type const* gather_map, + size_type* scatter_map, + size_type num_rows) +{ + auto const row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row < num_rows) { scatter_map[gather_map[row]] = static_cast(row); } +} + +CUDF_KERNEL void partition_string_sizes(partition_string_column const* columns, + size_type num_columns, + size_type num_rows, + size_type* sizes) +{ + auto const column = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (column < num_columns) { + sizes[column] = columns[column].input_offsets[num_rows] - columns[column].input_offsets[0]; + } +} + +/** + * @brief Materialize a complete permutation of non-nullable INT32-offset string columns. + * + * Batch the column-size readback and copy all character buffers in one kernel. The gather map + * must be a permutation, as guaranteed by partition's rank/scan construction. Repeated or missing + * indices are not supported. Other schemas return nullptr for the caller's generic gather path. + */ +[[nodiscard]] inline std::unique_ptr
gather_partition_strings( + table_view const& input, + device_span gather_map, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + if (input.num_columns() < 2 || input.num_columns() > 65'535 || + !std::all_of(input.begin(), input.end(), [](auto const& column) { + return column.type().id() == type_id::STRING && column.null_count() == 0 && + strings_column_view(column).offsets().type().id() == type_id::INT32; + })) { + return nullptr; + } + + auto const rows = input.num_rows(); + auto const columns = input.num_columns(); + auto const temp_mr = cudf::get_current_device_resource_ref(); + std::vector descriptors; + std::vector> offsets; + descriptors.reserve(columns); + offsets.reserve(columns); + + // A complete permutation preserves each sliced column's total character count. + // Read the input boundaries in one batch before allocating or scanning the output columns. + for (auto const& column : input) { + auto const strings = strings_column_view(column); + descriptors.push_back({strings.offsets().data() + column.offset(), + strings.chars_begin(stream), + nullptr, + nullptr}); + } + auto device_columns = cudf::detail::make_device_buffer(descriptors, stream, temp_mr); + cuda::device_buffer device_sizes(stream, temp_mr, columns, cuda::no_init); + partition_string_sizes<<<(columns + 255) / 256, 256, 0, stream.get()>>>( + device_columns.data(), columns, rows, device_sizes.data()); + auto const char_sizes = cudf::detail::make_pinned_vector(device_sizes, stream); + auto const max_chars = *std::max_element(char_sizes.begin(), char_sizes.end()); + auto const offset64_threshold = cudf::strings::get_offset64_threshold(); + if (max_chars >= offset64_threshold) { return nullptr; } + + auto const sizes = [&](size_type const* input_offsets) { + return cuda::transform_iterator( + cuda::counting_iterator{0}, + [map = gather_map.data(), input_offsets, rows] __device__(size_type i) -> size_type { + if (i == rows) { return size_type{0}; } + auto const row = map[i]; + return input_offsets[row + 1] - input_offsets[row]; + }); + }; + std::size_t temp_bytes = 0; + if (max_chars > 0) { + CUDF_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + temp_bytes, + sizes(nullptr), + static_cast(nullptr), + rows + 1, + stream.get())); + } + cuda::device_buffer scan_storage(stream, temp_mr, temp_bytes, cuda::no_init); + std::vector chars; + std::vector copy_columns; + chars.reserve(columns); + copy_columns.reserve(columns); + for (size_type i = 0; i < columns; ++i) { + offsets.push_back(cudf::make_numeric_column( + data_type{type_id::INT32}, rows + 1, mask_state::UNALLOCATED, stream, mr)); + chars.emplace_back(char_sizes[i], stream, mr); + descriptors[i].output_offsets = offsets.back()->mutable_view().data(); + descriptors[i].output_chars = static_cast(chars.back().data()); + if (char_sizes[i] == 0) { + CUDF_CUDA_TRY(cudaMemsetAsync(descriptors[i].output_offsets, + 0, + (static_cast(rows) + 1) * sizeof(size_type), + stream.get())); + } else { + CUDF_CUDA_TRY(cub::DeviceScan::ExclusiveSum(scan_storage.data(), + temp_bytes, + sizes(descriptors[i].input_offsets), + descriptors[i].output_offsets, + rows + 1, + stream.get())); + copy_columns.push_back(descriptors[i]); + } + } + + if (!copy_columns.empty()) { + // Complete this upload before releasing or reusing its pageable host source on every path. + device_columns = cudf::detail::make_device_buffer(copy_columns, stream, temp_mr); + cuda::device_buffer scatter_map(stream, temp_mr, rows, cuda::no_init); + invert_partition_map<<<(static_cast(rows) + 255) / 256, 256, 0, stream.get()>>>( + gather_map.data(), scatter_map.data(), rows); + auto const blocks = std::min(65'535, (static_cast(rows) + 63) / 64); + partition_strings_plan::copy_characters + <<>>( + device_columns.data(), scatter_map.data(), rows); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + } + + std::vector> output; + output.reserve(columns); + for (size_type i = 0; i < columns; ++i) { + output.push_back( + cudf::make_strings_column(rows, + std::move(offsets[i]), + std::move(chars[i]), + 0, + cudf::create_null_mask(0, mask_state::UNALLOCATED, stream, mr))); + } + return std::make_unique
(std::move(output)); +} + +} // namespace cudf::detail diff --git a/cpp/src/partitioning/partitioning.cu b/cpp/src/partitioning/partitioning.cu index ec153e0951f2..b5b98256d666 100644 --- a/cpp/src/partitioning/partitioning.cu +++ b/cpp/src/partitioning/partitioning.cu @@ -3,6 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "partition_fixed_width.cuh" +#include "partition_identity.cuh" +#include "partition_map_plan.cuh" +#include "partition_string_plan.cuh" + #include #include #include @@ -15,10 +20,14 @@ #include #include #include +#include #include #include #include #include +#include +#include +#include #include #include @@ -29,10 +38,15 @@ #include #include #include +#include +#include #include #include #include +#include +#include +#include #include namespace cudf { @@ -754,25 +768,569 @@ std::pair, std::vector> hash_partition_table( } } +constexpr size_type MAP_PARTITION_BLOCK_SIZE = 256; +constexpr size_type MAP_PARTITION_ITEMS_PER_THREAD = 4; +constexpr size_type MAP_PARTITION_TILE_SIZE = + MAP_PARTITION_BLOCK_SIZE * MAP_PARTITION_ITEMS_PER_THREAD; +constexpr size_type MAP_PARTITION_MAX_BINS = 256; +constexpr size_type MAP_PARTITION_WARP_SIZE = 32; +constexpr size_type MAP_PARTITION_NUM_WARPS = MAP_PARTITION_BLOCK_SIZE / MAP_PARTITION_WARP_SIZE; +constexpr size_type MAP_PARTITION_GATHER_TILES = 2; +static_assert(MAP_PARTITION_MAX_BINS <= MAP_PARTITION_BLOCK_SIZE, + "Partition-base initialization requires one thread per partition"); + +/** + * @brief Build warp-local histograms and optionally compute each row's warp-local rank. + * + * Each warp owns a histogram. Matching lanes elect one writer for each bin, avoiding + * atomic updates even when many rows belong to the same partition. Contiguous tiles + * per warp keep the resulting gather map in input order within each partition. + */ +template +__device__ void partition_map_warp_histogram( + MapType const* partition_map, + size_type num_rows, + size_type num_partitions, + cuda::std::size_t tile_begin, + size_type (&histogram)[MAP_PARTITION_NUM_WARPS][MAP_PARTITION_MAX_BINS], + size_type (&ranks)[MAP_PARTITION_ITEMS_PER_THREAD], + size_type (&partitions)[MAP_PARTITION_ITEMS_PER_THREAD]) +{ + constexpr auto warp_size = MAP_PARTITION_WARP_SIZE; + auto const lane = threadIdx.x % warp_size; + auto const warp = threadIdx.x / warp_size; + for (auto p = lane; p < num_partitions; p += warp_size) { + histogram[warp][p] = 0; + } + __syncwarp(); + +#pragma unroll + for (size_type item = 0; item < MAP_PARTITION_ITEMS_PER_THREAD; ++item) { + auto const row = + tile_begin + warp * warp_size * MAP_PARTITION_ITEMS_PER_THREAD + item * warp_size + lane; + auto const valid = row < num_rows; + auto const active = __ballot_sync(0xffffffffu, valid); + if (valid) { + auto const partition = static_cast(partition_map[row]); + auto const peers = __match_any_sync(active, partition); + auto const base = histogram[warp][partition]; + if constexpr (ComputeRanks) { + partitions[item] = partition; + ranks[item] = base + cuda::std::popcount(peers & ((1u << lane) - 1u)); + } + // All readers of a bin finish before its elected writer updates it. + __syncwarp(active); + if (lane == __ffs(peers) - 1) { + histogram[warp][partition] = base + cuda::std::popcount(peers); + } + } + __syncwarp(); + } + __syncthreads(); +} + +template +CUDF_KERNEL void partition_map_local_ranks(MapType const* partition_map, + size_type num_rows, + size_type num_partitions, + uint16_t* local_ranks, + size_type* block_counts) +{ + constexpr size_type warp_size = 32; + constexpr size_type num_warps = MAP_PARTITION_BLOCK_SIZE / warp_size; + __shared__ size_type histogram[num_warps][MAP_PARTITION_MAX_BINS]; + auto const lane = threadIdx.x % warp_size; + auto const warp = threadIdx.x / warp_size; + size_type ranks[MAP_PARTITION_ITEMS_PER_THREAD]; + size_type partitions[MAP_PARTITION_ITEMS_PER_THREAD]; + partition_map_warp_histogram( + partition_map, + num_rows, + num_partitions, + static_cast(blockIdx.x) * MAP_PARTITION_TILE_SIZE, + histogram, + ranks, + partitions); + + for (auto partition = threadIdx.x; partition < num_partitions; + partition += MAP_PARTITION_BLOCK_SIZE) { + size_type total = 0; +#pragma unroll + for (size_type w = 0; w < num_warps; ++w) { + auto const count = histogram[w][partition]; + histogram[w][partition] = total; + total += count; + } + block_counts[static_cast(partition) * gridDim.x + blockIdx.x] = total; + } + __syncthreads(); + +#pragma unroll + for (size_type item = 0; item < MAP_PARTITION_ITEMS_PER_THREAD; ++item) { + auto const row = static_cast(blockIdx.x) * MAP_PARTITION_TILE_SIZE + + warp * warp_size * MAP_PARTITION_ITEMS_PER_THREAD + item * warp_size + lane; + if (row < num_rows) { + // A rank is at most 1023, so two bytes per row are sufficient. + local_ranks[row] = static_cast(histogram[warp][partitions[item]] + ranks[item]); + } + } +} + +/** + * @brief Count two consecutive microtiles without materializing per-row ranks. + */ +template +CUDF_KERNEL void partition_map_gather_counts(MapType const* partition_map, + size_type num_rows, + size_type num_partitions, + size_type* block_counts) +{ + __shared__ size_type histogram[MAP_PARTITION_NUM_WARPS][MAP_PARTITION_MAX_BINS]; + size_type total = 0; + for (size_type tile = 0; tile < MAP_PARTITION_GATHER_TILES; ++tile) { + size_type ranks[MAP_PARTITION_ITEMS_PER_THREAD]; + size_type partitions[MAP_PARTITION_ITEMS_PER_THREAD]; + auto const tile_begin = + (static_cast(blockIdx.x) * MAP_PARTITION_GATHER_TILES + tile) * + MAP_PARTITION_TILE_SIZE; + partition_map_warp_histogram( + partition_map, num_rows, num_partitions, tile_begin, histogram, ranks, partitions); + if (threadIdx.x < num_partitions) { +#pragma unroll + for (size_type warp = 0; warp < MAP_PARTITION_NUM_WARPS; ++warp) { + total += histogram[warp][threadIdx.x]; + } + } + // Every histogram reader must finish before the next microtile resets it. + __syncthreads(); + } + if (threadIdx.x < num_partitions) { + block_counts[static_cast(threadIdx.x) * gridDim.x + blockIdx.x] = total; + } +} + +/** + * @brief Recompute microtile ranks and emit an output-to-input gather map. + * + * Each block carries its partition bases from the first microtile into the second, + * preserving source locality without a global per-row rank buffer. + */ +template +CUDF_KERNEL void partition_map_gather_scatter(MapType const* partition_map, + size_type const* block_offsets, + size_type num_rows, + size_type num_partitions, + size_type* gather_map, + size_type* partition_offsets) +{ + __shared__ size_type histogram[MAP_PARTITION_NUM_WARPS][MAP_PARTITION_MAX_BINS]; + __shared__ size_type partition_bases[MAP_PARTITION_MAX_BINS]; + if (threadIdx.x < num_partitions) { + auto const partition_base = static_cast(threadIdx.x) * gridDim.x; + partition_bases[threadIdx.x] = block_offsets[partition_base + blockIdx.x]; + if (blockIdx.x == 0) { partition_offsets[threadIdx.x] = block_offsets[partition_base]; } + } + if (blockIdx.x == 0 && threadIdx.x == 0) { partition_offsets[num_partitions] = num_rows; } + __syncthreads(); + + auto const lane = threadIdx.x % MAP_PARTITION_WARP_SIZE; + auto const warp = threadIdx.x / MAP_PARTITION_WARP_SIZE; + for (size_type tile = 0; tile < MAP_PARTITION_GATHER_TILES; ++tile) { + size_type ranks[MAP_PARTITION_ITEMS_PER_THREAD]; + size_type partitions[MAP_PARTITION_ITEMS_PER_THREAD]; + auto const tile_begin = + (static_cast(blockIdx.x) * MAP_PARTITION_GATHER_TILES + tile) * + MAP_PARTITION_TILE_SIZE; + partition_map_warp_histogram( + partition_map, num_rows, num_partitions, tile_begin, histogram, ranks, partitions); + + size_type next_base = 0; + if (threadIdx.x < num_partitions) { + size_type total = 0; +#pragma unroll + for (size_type w = 0; w < MAP_PARTITION_NUM_WARPS; ++w) { + auto const count = histogram[w][threadIdx.x]; + histogram[w][threadIdx.x] = total; + total += count; + } + next_base = partition_bases[threadIdx.x] + total; + } + __syncthreads(); + +#pragma unroll + for (size_type item = 0; item < MAP_PARTITION_ITEMS_PER_THREAD; ++item) { + auto const row = tile_begin + + warp * MAP_PARTITION_WARP_SIZE * MAP_PARTITION_ITEMS_PER_THREAD + + item * MAP_PARTITION_WARP_SIZE + lane; + if (row < num_rows) { + auto const partition = partitions[item]; + auto const destination = + partition_bases[partition] + histogram[warp][partition] + ranks[item]; + gather_map[destination] = static_cast(row); + } + } + __syncthreads(); + if (threadIdx.x < num_partitions) { partition_bases[threadIdx.x] = next_base; } + __syncthreads(); + } +} + +template +std::vector make_partition_gather_map(column_view const& partition_map, + size_type num_partitions, + size_type* gather_map, + cuda::stream_ref stream) +{ + auto const num_rows = partition_map.size(); + auto const num_blocks = + (num_rows - 1) / (MAP_PARTITION_GATHER_TILES * MAP_PARTITION_TILE_SIZE) + 1; + auto const temp_mr = cudf::get_current_device_resource_ref(); + cuda::device_buffer block_offsets( + stream, temp_mr, static_cast(num_blocks) * num_partitions, cuda::no_init); + cuda::device_buffer offsets(stream, temp_mr, num_partitions + 1, cuda::no_init); + partition_map_gather_counts<<>>( + partition_map.begin(), num_rows, num_partitions, block_offsets.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + thrust::exclusive_scan(rmm::exec_policy_nosync(stream, temp_mr), + block_offsets.begin(), + block_offsets.end(), + block_offsets.begin()); + partition_map_gather_scatter<<>>( + partition_map.begin(), + block_offsets.data(), + num_rows, + num_partitions, + gather_map, + offsets.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + auto const pinned_offsets = cudf::detail::make_pinned_vector(offsets, stream); + return {pinned_offsets.begin(), pinned_offsets.end()}; +} + +template +CUDF_KERNEL void partition_map_scatter(MapType const* partition_map, + uint16_t const* local_ranks, + size_type const* block_offsets, + size_type num_rows, + size_type num_blocks, + size_type num_partitions, + InputIterator input, + OutputIterator output, + size_type* partition_offsets) +{ + __shared__ size_type partition_bases[MAP_PARTITION_MAX_BINS]; + if (threadIdx.x < num_partitions) { + auto const partition_base = static_cast(threadIdx.x) * num_blocks; + partition_bases[threadIdx.x] = block_offsets[partition_base + blockIdx.x]; + if (blockIdx.x == 0) { partition_offsets[threadIdx.x] = block_offsets[partition_base]; } + } + if (blockIdx.x == 0 && threadIdx.x == 0) { partition_offsets[num_partitions] = num_rows; } + __syncthreads(); + + // Reuse each block's partition bases across all rows in its tile. +#pragma unroll + for (size_type item = 0; item < MAP_PARTITION_ITEMS_PER_THREAD; ++item) { + auto const row = static_cast(blockIdx.x) * MAP_PARTITION_TILE_SIZE + + item * MAP_PARTITION_BLOCK_SIZE + threadIdx.x; + if (row < num_rows) { + auto const partition = static_cast(partition_map[row]); + auto const destination = partition_bases[partition] + local_ranks[row]; + output[destination] = input[row]; + } + } +} + +template +std::vector scatter_with_block_ranks(column_view const& partition_map, + size_type num_partitions, + InputIterator input, + OutputIterator output, + cuda::stream_ref stream) +{ + auto const num_rows = partition_map.size(); + auto const num_blocks = (num_rows - 1) / MAP_PARTITION_TILE_SIZE + 1; + auto const temp_mr = cudf::get_current_device_resource_ref(); + std::vector partition_offsets; + { + cuda::device_buffer local_ranks(stream, temp_mr, num_rows, cuda::no_init); + cuda::device_buffer block_offsets( + stream, temp_mr, static_cast(num_blocks) * num_partitions, cuda::no_init); + cuda::device_buffer offsets(stream, temp_mr, num_partitions + 1, cuda::no_init); + partition_map_local_ranks<<>>( + partition_map.begin(), + num_rows, + num_partitions, + local_ranks.data(), + block_offsets.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + // Partition-major layout gives both the partition base and preceding blocks' + // row counts in one scan, without an atomic cursor for each partition. + thrust::exclusive_scan(rmm::exec_policy_nosync(stream, temp_mr), + block_offsets.begin(), + block_offsets.end(), + block_offsets.begin()); + partition_map_scatter<<>>( + partition_map.begin(), + local_ranks.data(), + block_offsets.data(), + num_rows, + num_blocks, + num_partitions, + input, + output, + offsets.data()); + CUDF_CUDA_TRY(cudaPeekAtLastError()); + auto const pinned_offsets = cudf::detail::make_pinned_vector(offsets, stream); + partition_offsets.assign(pinned_offsets.begin(), pinned_offsets.end()); + } + return partition_offsets; +} + +template +struct partition_single_column { + template + std::pair, std::vector> operator()( + column_view const& input, + column_view const& partition_map, + size_type num_partitions, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) const + { + if constexpr (cudf::is_fixed_width()) { + auto output = cudf::allocate_like(input, mask_allocation_policy::NEVER, stream, mr); + auto offsets = scatter_with_block_ranks( + partition_map, num_partitions, input.begin(), output->mutable_view().begin(), stream); + std::vector> columns; + columns.push_back(std::move(output)); + return {std::make_unique
(std::move(columns)), std::move(offsets)}; + } else { + CUDF_FAIL("Expected a fixed-width column"); + } + } +}; + +template +std::pair, std::vector> partition_with_block_ranks( + table_view const& input, + column_view const& partition_map, + size_type num_partitions, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + // A single fixed-width column with no nulls can be scattered directly, avoiding the + // gather map and the subsequent gather of its payload. + if (input.num_columns() == 1 && is_fixed_width(input.column(0).type()) && + (!input.column(0).nullable() || input.column(0).null_count() == 0)) { + return type_dispatcher(input.column(0).type(), + partition_single_column{}, + input.column(0), + partition_map, + num_partitions, + stream, + mr); + } + cuda::device_buffer gather_map( + stream, cudf::get_current_device_resource_ref(), input.num_rows(), cuda::no_init); + auto partition_offsets = + make_partition_gather_map(partition_map, num_partitions, gather_map.data(), stream); + auto output = + cudf::detail::gather(input, + device_span{gather_map.data(), gather_map.size()}, + out_of_bounds_policy::DONT_CHECK, + negative_index_policy::NOT_ALLOWED, + stream, + mr); + return {std::move(output), std::move(partition_offsets)}; +} + +/** + * @brief Compute partition boundaries without materializing a row permutation. + */ +template +std::vector make_map_partition_offsets(column_view const& partition_map, + size_type num_partitions, + cuda::stream_ref stream) +{ + rmm::device_uvector offsets(static_cast(num_partitions) + 1, stream); + CUDF_CUDA_TRY( + cudaMemsetAsync(offsets.data() + num_partitions, 0, sizeof(size_type), stream.get())); + std::size_t bytes = 0; + CUDF_CUDA_TRY(cub::DeviceHistogram::HistogramEven(nullptr, + bytes, + partition_map.begin(), + offsets.data(), + num_partitions + 1, + size_type{0}, + num_partitions, + partition_map.size(), + stream.get())); + { + rmm::device_buffer scratch(bytes, stream); + CUDF_CUDA_TRY(cub::DeviceHistogram::HistogramEven(scratch.data(), + bytes, + partition_map.begin(), + offsets.data(), + num_partitions + 1, + size_type{0}, + num_partitions, + partition_map.size(), + stream.get())); + } + thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + offsets.begin(), + offsets.end(), + offsets.begin()); + auto const pinned_offsets = cudf::detail::make_pinned_vector(offsets, stream); + return {pinned_offsets.begin(), pinned_offsets.end()}; +} + +/** + * @brief Handle cases whose output values do not require a permutation. + * + * Admission uses CPU metadata only. For a single row, the one map value is read + * solely to construct the host partition boundaries required by the API. + */ +template +std::optional, std::vector>> try_partition_trivial( + table_view const& input, + column_view const& partition_map, + size_type num_partitions, + cuda::stream_ref stream, + rmm::device_async_resource_ref mr) +{ + if (num_partitions <= 0 || num_partitions == std::numeric_limits::max()) { + return std::nullopt; + } + auto const rows = input.num_rows(); + auto const uniform = std::all_of(input.begin(), input.end(), [](auto const& col) { + return cudf::detail::partition_identity::is_uniform(col); + }); + if (num_partitions != 1 && rows != 1 && !uniform) { return std::nullopt; } + + std::vector offsets; + if (num_partitions == 1) { + offsets = {0, rows}; + } else if (rows == 1) { + auto const value = cudf::detail::make_std_vector( + device_span{partition_map.begin(), 1}, stream)[0]; + offsets.resize(static_cast(num_partitions) + 1); + for (std::size_t p = 0; p < offsets.size(); ++p) { + offsets[p] = p > static_cast(value) ? rows : 0; + } + } else { + offsets = make_map_partition_offsets(partition_map, num_partitions, stream); + } + if (!uniform) { + return std::pair{cudf::detail::partition_identity::copy(input, stream, mr), std::move(offsets)}; + } + + std::vector> columns; + columns.reserve(input.num_columns()); + for (auto const& col : input) { + columns.push_back(cudf::detail::partition_identity::uniform_column(col, stream, mr)); + } + return std::pair{std::make_unique
(std::move(columns), rows), std::move(offsets)}; +} + +struct partition_gather_output { + size_type* indices; + __device__ void operator()(size_type source, size_type destination) const + { + indices[destination] = source; + } +}; + +struct partition_plan { + cudf::detail::partition_map_plan::algorithm method; + bool direct; + bool reordered; +}; + +/** + * Choose a plan using only cached host metadata. Null-mask presence is not a + * discriminator: a nullable column with null_count() == 0 uses the non-null path. + * The large-input crossover is measured over different partition skews and widths. + */ +std::optional select_partition_plan(table_view const& input, size_type partitions) +{ + namespace plans = cudf::detail::partition_map_plan; + if (partitions <= 1 || partitions > MAP_PARTITION_MAX_BINS || + input.num_rows() < MAP_PARTITION_TILE_SIZE) { + return std::nullopt; + } + auto const fixed = std::all_of( + input.begin(), input.end(), [](auto const& col) { return is_fixed_width(col.type()); }); + auto const nullable = + std::any_of(input.begin(), input.end(), [](auto const& col) { return col.null_count() > 0; }); + std::size_t bytes_per_row = 0; + if (fixed) { + for (auto const& col : input) { + bytes_per_row += cudf::size_of(col.type()); + } + } + constexpr size_type large_input = 1'000'000; + bool const large = input.num_rows() >= large_input; + bool const dense_nulls = + fixed && nullable && large && std::all_of(input.begin(), input.end(), [&](auto const& col) { + return int64_t{col.null_count()} * 4 >= int64_t{input.num_rows()} * 3; + }); + // With few valid rows, skip payload reads and reorder writes without a row map. + if (dense_nulls && input.num_columns() > 2) { + return partition_plan{plans::algorithm::recompute, true, true}; + } + if (partitions == 2) { + bool const direct = fixed && input.num_columns() <= 2 && + (input.num_columns() == 1 || !nullable || !large || bytes_per_row >= 16); + partition_plan plan{plans::algorithm::binary, direct, false}; + if (fixed && !nullable && large) { + plan.direct = true; + if (input.num_columns() > 2) { + // Reordered output pays off for a few wide columns and removes saved ranks. + bool const reorder = input.num_columns() <= 4 && bytes_per_row >= 32; + plan.method = reorder ? plans::algorithm::recompute : plans::algorithm::saved; + plan.reordered = reorder; + } else if (bytes_per_row > 4) { + plan.method = plans::algorithm::recompute; + plan.reordered = input.num_columns() == 1 && bytes_per_row == 16; + } + } + return plan; + } + if (fixed && large) { + if (input.num_columns() == 1 || + (input.num_columns() == 2 && (!nullable || bytes_per_row >= 16))) { + // Saved ranks avoid recomputing sparse nullable two-column tables; the + // direct sink still removes the larger N-entry gather map. + auto const method = nullable && input.num_columns() == 2 && !dense_nulls + ? plans::algorithm::saved + : partitions >= 16 ? plans::algorithm::reserve + : plans::algorithm::recompute; + return partition_plan{method, true, false}; + } + if (!nullable) { + // A one-tile prefix is more robust than supertiles for narrow multi-column + // output with uniform partition maps, while retaining direct-output savings. + return partition_plan{plans::algorithm::recompute, true, true}; + } + } + return std::nullopt; +} + struct dispatch_map_type { /** * @brief Partitions the table `t` according to the `partition_map`. * - * Algorithm: + * For a small partition count and at least one full tile, compute warp-local + * ranks and scan block histograms to scatter a single fixed-width column or build a + * gather map without per-row global atomics. + * + * Otherwise: * - Compute the histogram of the size each partition * - Compute the exclusive scan of the histogram to get the offset for each * partition in the final partitioned output * - Use a transform iterator to materialize the scatter map of the rows from * `t` into the final output. * - * @note JH: It would likely be more efficient to avoid the atomic increments - * in the transform iterator. It would probably be faster to compute a - * per-thread block histogram and compute an exclusive scan of all of the - * per-block histograms (like in hash partition). But I'm purposefully trying - * to reduce memory pressure by avoiding intermediate materializations. Plus, - * atomics resolve in L2 and should be pretty fast since all the offsets will - * fit in L2. - * */ template std::pair, std::vector> operator()( @@ -783,6 +1341,64 @@ struct dispatch_map_type { rmm::device_async_resource_ref mr) const requires(is_index_type()) { + if (auto result = + try_partition_trivial(t, partition_map, num_partitions, stream, mr)) { + return std::move(*result); + } + // All-null and CPU-known empty-string columns are invariant. Build them + // after the reduced call has released its row map and other temporaries. + auto const rows = t.num_rows(); + auto const omit = [](column_view const& col) { + return cudf::detail::partition_identity::is_uniform(col); + }; + if (std::any_of(t.begin(), t.end(), omit)) { + std::vector reduced; + for (auto const& col : t) { + if (!omit(col)) { reduced.push_back(col); } + } + auto [result, offsets] = this->template operator()( + table_view{reduced}, partition_map, num_partitions, stream, mr); + auto owned = result->release(); + std::vector> columns; + columns.reserve(t.num_columns()); + auto next = owned.begin(); + for (auto const& col : t) { + columns.push_back(omit(col) + ? cudf::detail::partition_identity::uniform_column(col, stream, mr) + : std::move(*next++)); + } + return {std::make_unique
(std::move(columns), rows), std::move(offsets)}; + } + if (auto result = cudf::detail::partition_strings_plan::partition( + t, partition_map, num_partitions, stream, mr)) { + return std::move(*result); + } + if (auto plan = select_partition_plan(t, num_partitions)) { + auto launch = [&](auto sink) { + return cudf::detail::partition_map_plan::run(partition_map.begin(), + t.num_rows(), + num_partitions, + plan->method, + sink, + plan->reordered, + stream); + }; + if (plan->direct) { return cudf::detail::partition_fixed::partition(t, launch, stream, mr); } + rmm::device_uvector map(t.num_rows(), stream); + auto offsets = launch(partition_gather_output{map.data()}); + auto output = cudf::detail::gather(t, + device_span{map.data(), map.size()}, + out_of_bounds_policy::DONT_CHECK, + negative_index_policy::NOT_ALLOWED, + stream, + mr); + return {std::move(output), std::move(offsets)}; + } + if (num_partitions > 1 && num_partitions <= MAP_PARTITION_MAX_BINS && + t.num_rows() >= MAP_PARTITION_TILE_SIZE) { + return partition_with_block_ranks(t, partition_map, num_partitions, stream, mr); + } + // Build a histogram of the number of rows in each partition rmm::device_uvector histogram(num_partitions + 1, stream); std::size_t temp_storage_bytes{}; diff --git a/cpp/tests/partitioning/partition_test.cpp b/cpp/tests/partitioning/partition_test.cpp index 824331513fd3..d349e66aaf77 100644 --- a/cpp/tests/partitioning/partition_test.cpp +++ b/cpp/tests/partitioning/partition_test.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include #include @@ -12,12 +13,27 @@ #include #include #include +#include #include +#include +#include + +#include +#include #include #include +#include #include +#include +#include +#include +#include +#include +#include +#include +#include template class PartitionTest : public cudf::test::BaseFixture { @@ -123,6 +139,37 @@ void run_partition_test(cudf::table_view table_to_partition, } // Normal cases +TYPED_TEST(PartitionTest, SingleFixedWidthTileBoundaries) +{ + using value_type = cudf::test::GetType; + using map_type = cudf::test::GetType; + + for (auto const rows : {1024, 1025}) { + for (auto const hot : {false, true}) { + SCOPED_TRACE(::testing::Message() << rows << "/" << hot); + auto const values = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return (i * 37 + 11) % 97; }); + fixed_width_column_wrapper data(values, values + rows + 7); + auto const input = cudf::slice(cudf::table_view{{data}}, {7, rows + 7})[0]; + std::vector map(rows + 3, 0); + std::vector expected_values; + std::vector offsets{0}; + for (cudf::size_type p = 0; p < 2; ++p) { + for (cudf::size_type i = 0; i < rows; ++i) { + map[i + 3] = hot ? 1 : i % 2; + if (map[i + 3] == static_cast(p)) { expected_values.push_back(values[i + 7]); } + } + offsets.push_back(expected_values.size()); + } + fixed_width_column_wrapper map_column(map.begin(), map.end()); + auto const map_view = cudf::slice(map_column, {3, rows + 3})[0]; + fixed_width_column_wrapper expected(expected_values.begin(), + expected_values.end()); + run_partition_test(input, map_view, 2, cudf::table_view{{expected}}, offsets); + } + } +} + TYPED_TEST(PartitionTest, Identity) { using value_type = cudf::test::GetType; @@ -280,6 +327,469 @@ TYPED_TEST(PartitionTestFixedPoint, Partition2) struct PartitionTestNotTyped : public cudf::test::BaseFixture {}; +template +struct PartitionMapRanksTest : public cudf::test::BaseFixture {}; + +TYPED_TEST_SUITE(PartitionMapRanksTest, cudf::test::IntegralTypesNotBool); + +TYPED_TEST(PartitionMapRanksTest, SingleRowWithSlicedMap) +{ + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + std::vector values(15, "value"); + values[7] = std::string{"x\0y", 3}; + for (auto const null_mode : {0, 1, 2}) { + auto const numeric_valid = cuda::make_constant_iterator(null_mode == 0); + auto const string_valid = cuda::make_constant_iterator(null_mode != 2); + fixed_width_column_wrapper numeric(ids, ids + 15, numeric_valid); + strings_column_wrapper strings(values.begin(), values.end(), string_valid); + auto const input = cudf::slice(cudf::table_view{{numeric, strings}}, {7, 8})[0]; + for (auto const partitions : {2, 100}) { + SCOPED_TRACE(::testing::Message() << null_mode << "/" << partitions); + auto const last = cuda::make_constant_iterator(partitions - 1); + fixed_width_column_wrapper map(last, last + 7); + auto const map_view = cudf::slice(cudf::column_view{map}, {3, 4})[0]; + cudf::get_default_stream().sync(); + auto [output, offsets] = cudf::partition(input, map_view, partitions, stream, this->mr()); + std::vector expected(partitions + 1, 0); + expected.back() = 1; + EXPECT_EQ(offsets, expected); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(input, output->view(), stream); + } + } +} + +TEST_F(PartitionTestNotTyped, PruneNullColumnsPreservesSchemaAndRows) +{ + constexpr cudf::size_type rows = 4097; + auto const ids = cuda::counting_iterator{0}; + auto const invalid = cuda::make_constant_iterator(false); + auto const valid = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 3 != 0; }); + fixed_width_column_wrapper row_ids(ids, ids + rows + 14); + fixed_width_column_wrapper null_numbers(ids, ids + rows + 14, invalid); + fixed_width_column_wrapper partial(ids, ids + rows + 14, valid); + std::vector values(rows + 14, "discarded"); + strings_column_wrapper null_strings(values.begin(), values.end(), invalid); + auto const input = + cudf::slice(cudf::table_view{{null_numbers, row_ids, null_strings, partial}}, {7, rows + 7})[0]; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto const partitions : {2, 100, 257}) { + SCOPED_TRACE(partitions); + std::vector assignments(rows); + std::vector expected(partitions + 1, 0); + for (auto i = 0; i < rows; ++i) { + assignments[i] = (i * 17 + 3) % partitions; + ++expected[assignments[i] + 1]; + } + std::partial_sum(expected.begin(), expected.end(), expected.begin()); + fixed_width_column_wrapper map(assignments.begin(), assignments.end()); + cudf::get_default_stream().sync(); + auto [output, offsets] = cudf::partition(input, map, partitions, stream, mr()); + EXPECT_EQ(offsets, expected); + auto const [actual_ids, mask] = cudf::test::to_host(output->get_column(1), stream); + for (auto p = 0; p < partitions; ++p) { + for (auto i = offsets[p]; i < offsets[p + 1]; ++i) { + ASSERT_GE(actual_ids[i], 7); + ASSERT_LT(actual_ids[i], rows + 7); + EXPECT_EQ(assignments[actual_ids[i] - 7], p); + } + } + auto const order = std::vector{1, 0, 2, 3}; + auto expected_sorted = cudf::sort(input.select(order), {}, {}, stream, mr()); + auto actual_sorted = cudf::sort(output->view().select(order), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_sorted->view(), actual_sorted->view(), stream); + } +} + +TEST_F(PartitionTestNotTyped, UniformEmptyStringsAndMixedColumns) +{ + constexpr cudf::size_type rows = 4097; + auto const ids = cuda::counting_iterator{0}; + auto const valid = cuda::make_constant_iterator(true); + auto const invalid = cuda::make_constant_iterator(false); + auto const partial = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 3 != 0; }); + std::vector values(rows + 14, ""); + fixed_width_column_wrapper row_ids(ids, ids + rows + 14); + fixed_width_column_wrapper null_numbers(ids, ids + rows + 14, invalid); + strings_column_wrapper partial_empty(values.begin(), values.end(), partial); + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto wide : {false, true}) { + auto empty = strings_column_wrapper(values.begin(), values.end(), valid).release(); + if (wide) { + auto contents = empty->release(); + auto offsets = + cudf::cast(contents.children[0]->view(), cudf::data_type{cudf::type_id::INT64}); + empty = cudf::make_strings_column(rows + 14, + std::move(offsets), + std::move(*contents.data), + 0, + std::move(*contents.null_mask)); + } + ASSERT_EQ(empty->view().head(), nullptr); + for (auto const& unsliced : + {cudf::table_view{{*empty, null_numbers}}, + cudf::table_view{{*empty, partial_empty, row_ids, null_numbers}}}) { + auto input = cudf::slice(unsliced, {7, rows + 7})[0]; + for (auto partitions : {1, 2, 100, 257}) { + SCOPED_TRACE(::testing::Message() + << wide << "/" << input.num_columns() << "/" << partitions); + std::vector assignments(rows); + std::vector expected(partitions + 1, 0); + for (auto i = 0; i < rows; ++i) { + assignments[i] = (i * 17 + 3) % partitions; + ++expected[assignments[i] + 1]; + } + std::partial_sum(expected.begin(), expected.end(), expected.begin()); + fixed_width_column_wrapper map(assignments.begin(), assignments.end()); + cudf::get_default_stream().sync(); + auto [output, offsets] = cudf::partition(input, map, partitions, stream, mr()); + EXPECT_EQ(offsets, expected); + EXPECT_EQ(output->get_column(0).null_count(), 0); + auto expected_sorted = cudf::sort(input, {}, {}, stream, mr()); + auto actual_sorted = cudf::sort(output->view(), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_sorted->view(), actual_sorted->view(), stream); + if (input.num_columns() == 4) { + auto const [actual_ids, mask] = + cudf::test::to_host(output->get_column(2), stream); + for (auto p = 0; p < partitions; ++p) + for (auto i = offsets[p]; i < offsets[p + 1]; ++i) { + ASSERT_GE(actual_ids[i], 7); + ASSERT_LT(actual_ids[i], rows + 7); + EXPECT_EQ(assignments[actual_ids[i] - 7], p); + } + } + } + } + } +} + +TEST_F(PartitionTestNotTyped, AllNullColumnsWithoutRowMap) +{ + constexpr cudf::size_type rows = 4097; + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + auto const invalid = cuda::make_constant_iterator(false); + fixed_width_column_wrapper narrow(ids, ids + rows + 14, invalid); + fixed_width_column_wrapper wide(ids, ids + rows + 14, invalid); + std::vector values(rows + 14, "unused"); + strings_column_wrapper strings(values.begin(), values.end(), invalid); + auto const input = cudf::slice(cudf::table_view{{narrow, wide, strings}}, {7, rows + 7})[0]; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto const partitions : {1, 2, 100, 256, 257}) { + SCOPED_TRACE(partitions); + std::vector partition_ids(rows + 3, 0); + std::vector expected(partitions + 1, 0); + for (cudf::size_type i = 0; i < rows; ++i) { + auto const p = (i * 17 + 3) % partitions; + partition_ids[i + 3] = p; + ++expected[p + 1]; + } + std::partial_sum(expected.begin(), expected.end(), expected.begin()); + fixed_width_column_wrapper map(partition_ids.begin(), partition_ids.end()); + auto const map_view = cudf::slice(cudf::column_view{map}, {3, rows + 3})[0]; + cudf::get_default_stream().sync(); + auto [output, offsets] = cudf::partition(input, map_view, partitions, stream, mr()); + EXPECT_EQ(offsets, expected); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(input, output->view(), stream); + } +} + +TEST_F(PartitionTestNotTyped, SinglePartitionNestedColumn) +{ + cudf::test::lists_column_wrapper lists{{1, 2}, {}, {3}}; + fixed_width_column_wrapper map{0, 0, 0}; + auto const input = cudf::table_view{{lists}}; + auto [output, offsets] = cudf::partition(input, map, 1); + EXPECT_EQ(offsets, (std::vector{0, 3})); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, output->view()); +} + +TEST_F(PartitionTestNotTyped, ZeroColumnsWithRows) +{ + for (auto const rows : {1, 1025}) { + auto const input = cudf::table_view{std::vector{}, rows}; + for (auto const partitions : {1, 2, 257}) { + SCOPED_TRACE(::testing::Message() << rows << "/" << partitions); + auto const ids = cudf::detail::make_counting_transform_iterator( + 0, [partitions](auto i) { return (i * 17 + 3) % partitions; }); + fixed_width_column_wrapper map(ids, ids + rows); + std::vector expected(partitions + 1, 0); + for (auto i = 0; i < rows; ++i) { + ++expected[ids[i] + 1]; + } + std::partial_sum(expected.begin(), expected.end(), expected.begin()); + auto [output, offsets] = cudf::partition(input, map, partitions); + EXPECT_EQ(offsets, expected); + EXPECT_EQ(output->num_rows(), rows); + EXPECT_EQ(output->num_columns(), 0); + } + } +} + +TEST_F(PartitionTestNotTyped, TrivialStringOutputOffsetThreshold) +{ + struct restore_setting { + char const* name; + std::optional previous; + explicit restore_setting(char const* key) : name(key) + { + if (auto const value = std::getenv(key)) { previous = value; } + } + ~restore_setting() + { + if (previous) { + setenv(name, previous->c_str(), 1); + } else { + unsetenv(name); + } + } + }; + restore_setting restore_threshold{"LIBCUDF_LARGE_STRINGS_THRESHOLD"}; + restore_setting restore_enabled{"LIBCUDF_LARGE_STRINGS_ENABLED"}; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto const rows : {1, 1025}) { + for (auto const wide : {false, true}) { + for (auto const nullable : {false, true}) { + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_THRESHOLD", wide ? "1" : "65536", 1), 0); + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_ENABLED", "1", 1), 0); + auto const length = rows == 1 ? 4 : 1; + std::vector values(rows + 14, std::string(length, 'a')); + auto const valid = cudf::detail::make_counting_transform_iterator( + 0, [nullable](auto i) { return !nullable || i % 10 != 0; }); + auto strings = strings_column_wrapper(values.begin(), values.end(), valid).release(); + if (wide) { + auto const nulls = strings->null_count(); + auto contents = strings->release(); + auto wide_offsets = + cudf::cast(contents.children[0]->view(), cudf::data_type{cudf::type_id::INT64}); + strings = cudf::make_strings_column(rows + 14, + std::move(wide_offsets), + std::move(*contents.data), + nulls, + std::move(*contents.null_mask)); + } + auto const input = cudf::slice(cudf::table_view{{*strings, *strings}}, {7, rows + 7})[0]; + auto const expected_input_type = wide ? cudf::type_id::INT64 : cudf::type_id::INT32; + ASSERT_EQ(cudf::strings_column_view(input.column(0)).offsets().type().id(), + expected_input_type); + auto const partitions = rows == 1 ? 2 : 1; + auto const last = cuda::make_constant_iterator(partitions - 1); + fixed_width_column_wrapper map(last, last + rows); + auto const bytes = (rows - input.column(0).null_count()) * length; + cudf::get_default_stream().sync(); + for (auto const threshold : {bytes - 1, bytes, bytes + 1}) { + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_THRESHOLD", std::to_string(threshold).c_str(), 1), + 0); + for (auto const enabled : {false, true}) { + SCOPED_TRACE(::testing::Message() << rows << "/" << wide << "/" << nullable << "/" + << threshold << "/" << enabled); + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_ENABLED", enabled ? "1" : "0", 1), 0); + if (!enabled && bytes >= threshold) { + EXPECT_THROW(cudf::partition(input, map, partitions, stream, mr()), + std::overflow_error); + } else { + auto [output, offsets] = cudf::partition(input, map, partitions, stream, mr()); + auto const expected_offsets = partitions == 1 ? std::vector{0, rows} + : std::vector{0, 0, 1}; + EXPECT_EQ(offsets, expected_offsets); + auto const type = bytes >= threshold ? cudf::type_id::INT64 : cudf::type_id::INT32; + for (auto const& col : output->view()) { + EXPECT_EQ(cudf::strings_column_view(col).offsets().type().id(), type); + } + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(input, output->view(), stream); + } + } + } + } + } + } +} + +TEST_F(PartitionTestNotTyped, SingleFixedWidthNullMask) +{ + constexpr cudf::size_type rows = 1025; + constexpr cudf::size_type partitions = 100; + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + auto const last_partition = cuda::make_constant_iterator(partitions - 1); + fixed_width_column_wrapper map(last_partition, last_partition + rows); + std::vector expected_offsets(partitions + 1, 0); + expected_offsets.back() = rows; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + + for (auto const null_mode : {0, 1, 2}) { + SCOPED_TRACE(null_mode); + // Cover an all-valid mask, nulls outside the slice, and one null inside the slice. + auto const valid = cudf::detail::make_counting_transform_iterator(0, [null_mode](auto i) { + return null_mode == 0 || (i >= 7 && i < rows + 7 && (null_mode != 2 || i != rows / 2 + 7)); + }); + fixed_width_column_wrapper data(ids, ids + rows + 14, valid); + auto const input = cudf::slice(cudf::column_view{data}, {7, rows + 7})[0]; + ASSERT_TRUE(input.nullable()); + ASSERT_EQ(input.null_count(), null_mode == 2 ? 1 : 0); + + cudf::get_default_stream().sync(); + + auto [output, offsets] = + cudf::partition(cudf::table_view{{input}}, map, partitions, stream, mr()); + EXPECT_EQ(offsets, expected_offsets); + EXPECT_EQ(output->get_column(0).nullable(), null_mode == 2); + auto sorted_input = cudf::sort(cudf::table_view{{input}}, {}, {}, stream, mr()); + auto sorted_output = cudf::sort(output->view(), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(sorted_input->view(), sorted_output->view(), stream); + } +} + +TEST_F(PartitionTestNotTyped, SingleFixedWidthManyPartitions) +{ + constexpr cudf::size_type rows = 4099; + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + fixed_width_column_wrapper data(ids, ids + rows); + for (auto const partitions : {200, 256, 257}) { + for (auto const hot : {false, true}) { + SCOPED_TRACE(::testing::Message() << partitions << "/" << hot); + std::vector map(rows); + std::vector expected_values; + std::vector offsets{0}; + for (cudf::size_type p = 0; p < partitions; ++p) { + for (cudf::size_type i = 0; i < rows; ++i) { + map[i] = hot ? partitions - 1 : (i * 17 + 3) % partitions; + if (map[i] == p) { expected_values.push_back(i); } + } + offsets.push_back(expected_values.size()); + } + fixed_width_column_wrapper map_column(map.begin(), map.end()); + fixed_width_column_wrapper expected(expected_values.begin(), expected_values.end()); + run_partition_test( + cudf::table_view{{data}}, map_column, partitions, cudf::table_view{{expected}}, offsets); + } + } +} + +TYPED_TEST(PartitionMapRanksTest, TileBoundariesAndSkew) +{ + using size_type = cudf::size_type; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto const rows : {1023, 1024, 1025, 2047, 2048, 2049, 4095, 4096, 4097, 4099}) { + for (auto const partitions : {1, 2, 100, 200, 256, 257}) { + if (static_cast(partitions - 1) > + static_cast(std::numeric_limits::max())) { + continue; + } + for (auto const distribution : {0, 1, 2}) { + SCOPED_TRACE(::testing::Message() << rows << "/" << partitions << "/" << distribution); + // Slice both the data and the map, with a different starting offset for each. + auto const ids = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + fixed_width_column_wrapper row_ids(ids, ids + rows + 7); + auto const valid = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 7 != 0; }); + fixed_width_column_wrapper nullable(ids, ids + rows + 7, valid); + std::vector strings(rows + 7); + std::transform(ids, ids + rows + 7, strings.begin(), [](auto i) { + return i % 11 == 0 ? std::string{} : std::to_string(i) + "-\xE4\xB8\xAD"; + }); + strings_column_wrapper text(strings.begin(), strings.end(), valid); + auto input = cudf::slice(cudf::table_view{{row_ids, nullable, text}}, {7, rows + 7})[0]; + + std::vector map(rows + 3, 0); + std::vector expected_offsets(partitions + 1, 0); + for (size_type i = 0; i < rows; ++i) { + auto const partition = distribution == 1 ? partitions - 1 + : distribution == 2 && i % 10 != 0 ? 0 + : (i * 17 + 3) % partitions; + map[i + 3] = static_cast(partition); + ++expected_offsets[partition + 1]; + } + std::partial_sum( + expected_offsets.begin(), expected_offsets.end(), expected_offsets.begin()); + fixed_width_column_wrapper map_column(map.begin(), map.end()); + auto map_view = cudf::slice(map_column, {3, rows + 3})[0]; + cudf::get_default_stream().sync(); + + auto [output, offsets] = cudf::partition(input, map_view, partitions, stream, this->mr()); + EXPECT_EQ(offsets, expected_offsets); + auto const [actual_ids, mask] = + cudf::test::to_host(output->get_column(0), stream); + for (size_type p = 0; p < partitions; ++p) { + for (auto i = offsets[p]; i < offsets[p + 1]; ++i) { + ASSERT_GE(actual_ids[i], 7); + ASSERT_LT(actual_ids[i], rows + 7); + EXPECT_EQ(map[actual_ids[i] - 7 + 3], p); + } + } + // Unique row ids catch dropped/duplicated rows without requiring a stable partition. + auto sorted = cudf::sort(output->view(), {}, {}, stream, this->mr()); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, sorted->view(), stream); + } + } + } +} + +TEST_F(PartitionTestNotTyped, PartitionMemoryResources) +{ + std::vector values(1025); + std::vector partitions(1025, 0); + for (int32_t i = 0; i < 1025; ++i) { + values[i] = "value-" + std::to_string(i); + } + strings_column_wrapper strings(values.begin(), values.end()); + fixed_width_column_wrapper map(partitions.begin(), partitions.end()); + auto const input = cudf::table_view{{strings, strings}}; + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + auto output_mr = rmm::mr::statistics_resource_adaptor{mr()}; + auto temporary_mr = rmm::mr::statistics_resource_adaptor{mr()}; + cudf::get_default_stream().sync(); + std::unique_ptr output; + { + auto current = cudf::test::scoped_current_device_resource{temporary_mr}; + auto result = cudf::partition(input, map, 100, stream, output_mr); + output = std::move(result.first); + stream.sync(); + } + EXPECT_GT(output_mr.get_bytes_counter().value, 0); + EXPECT_GT(temporary_mr.get_bytes_counter().total, 0); + EXPECT_EQ(temporary_mr.get_bytes_counter().value, 0); + auto expected = cudf::sort(input); + auto actual = cudf::sort(output->view()); + CUDF_TEST_EXPECT_TABLES_EQUAL(expected->view(), actual->view()); + output.reset(); + stream.sync(); + EXPECT_EQ(output_mr.get_bytes_counter().value, 0); +} + +TEST_F(PartitionTestNotTyped, PartitionFixedWidthMemoryResources) +{ + constexpr cudf::size_type rows = 1025; + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + fixed_width_column_wrapper data(ids, ids + rows); + auto const ones = cuda::make_constant_iterator(1); + fixed_width_column_wrapper map(ones, ones + rows); + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + cudf::get_default_stream().sync(); + + for (auto const& input : {cudf::table_view{{data}}, cudf::table_view{{data, data}}}) { + SCOPED_TRACE(input.num_columns()); + cudf::test::memory_resource_test_harness resources{mr()}; + std::unique_ptr output; + { + auto current = cudf::test::scoped_current_device_resource{resources.temporary_mr()}; + auto result = cudf::partition(input, map, 2, stream, resources.output_mr()); + output = std::move(result.first); + EXPECT_EQ(result.second, (std::vector{0, 0, rows})); + stream.sync(); + } + resources.expect_resource_usage(output->alloc_size(), + {cudf::test::output_allocation_expectation::EXACT, + cudf::test::temporary_allocation_expectation::SOME}, + stream); + auto sorted = cudf::sort(output->view(), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, sorted->view(), stream); + output.reset(); + resources.expect_no_live_allocations(stream); + } +} + TEST_F(PartitionTestNotTyped, ListOfStringsEmpty) { cudf::test::lists_column_wrapper list{{}, {}}; @@ -354,3 +864,374 @@ TEST_F(PartitionTestNotTyped, NoIntegerOverflow) run_partition_test(table_to_partition, map, 2, expected_table, expected_offsets); } + +TEST_F(PartitionTestNotTyped, StringPermutationAndFallbacks) +{ + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + // Two partitions exercise the specialized string paths. Schemas 1/2/3 fall + // back for actual nulls, mixed types, and INT64 input offsets. Schema 5 checks + // fused descriptor compaction with empty columns between nonempty columns. + for (auto const schema : {0, 1, 2, 3, 4, 5}) { + for (auto const rows : {1025, 4099}) { + auto const total = rows + 5; + std::vector ids(total), values(total), empty(total); + std::vector valid(total); + std::vector numbers(total), partition_ids(total); + for (int32_t i = 0; i < total; ++i) { + ids[i] = "row-" + std::to_string(i); + values[i] = i % 13 == 0 ? std::string{"a\0b", 3} + : i % 7 == 0 + ? std::string{} + : std::string(i % 67, 'a' + i % 26) + "\xE4\xB8\xAD\xF0\x9F\x8C\x8D"; + if (schema == 4) { + values[i] = i % 13 == 0 ? std::string{"a\0b", 3} + : i % 7 == 0 ? "\xE4\xB8\xAD" + : std::string(i % 17, 'a' + i % 26); + } + valid[i] = i % 11 != 0; + numbers[i] = i; + } + strings_column_wrapper id_column(ids.begin(), ids.end()); + strings_column_wrapper empty_column(empty.begin(), empty.end()); + auto value_column = + schema == 1 ? strings_column_wrapper(values.begin(), values.end(), valid.begin()).release() + : strings_column_wrapper(values.begin(), values.end()).release(); + if (schema == 3) { + auto contents = value_column->release(); + auto wide_offsets = + cudf::cast(contents.children[0]->view(), cudf::data_type{cudf::type_id::INT64}); + value_column = + cudf::make_strings_column(total, + std::move(wide_offsets), + std::move(*contents.data), + 0, + cudf::create_null_mask(0, cudf::mask_state::UNALLOCATED)); + } + fixed_width_column_wrapper numeric(numbers.begin(), numbers.end()); + std::vector columns{id_column, *value_column, empty_column}; + if (schema == 2) { columns.push_back(numeric); } + if (schema == 5) { columns = {id_column, empty_column, *value_column, empty_column}; } + auto input = cudf::slice(cudf::table_view{columns}, {5, total})[0]; + for (auto const partitions : {2, 100, 200}) { + for (auto const skew : {false, true}) { + SCOPED_TRACE(::testing::Message() + << schema << "/" << rows << "/" << partitions << "/" << skew); + std::vector expected_offsets(partitions + 1, 0); + for (int32_t i = 0; i < rows; ++i) { + partition_ids[i + 5] = skew && i % 10 != 0 ? partitions - 1 : (i * 37 + 7) % partitions; + ++expected_offsets[partition_ids[i + 5] + 1]; + } + std::partial_sum( + expected_offsets.begin(), expected_offsets.end(), expected_offsets.begin()); + fixed_width_column_wrapper map(partition_ids.begin(), partition_ids.end()); + auto map_view = cudf::slice(map, {5, total})[0]; + cudf::get_default_stream().sync(); + auto [output, offsets] = cudf::partition(input, map_view, partitions, stream, mr()); + EXPECT_EQ(offsets, expected_offsets); + auto const [actual_ids, mask] = + cudf::test::to_host(output->get_column(0), stream); + std::vector seen(total, false); + for (int32_t p = 0; p < partitions; ++p) { + for (auto row = offsets[p]; row < offsets[p + 1]; ++row) { + auto const id = std::stoi(actual_ids[row].substr(4)); + ASSERT_GE(id, 5); + ASSERT_LT(id, total); + EXPECT_FALSE(seen[id]); + seen[id] = true; + EXPECT_EQ(partition_ids[id], p); + } + } + auto expected = cudf::sort(input, {}, {}, stream, mr()); + auto actual = cudf::sort(output->view(), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected->view(), actual->view(), stream); + } + } + } + } +} + +TEST_F(PartitionTestNotTyped, AllEmptyStringPermutation) +{ + std::vector empty(1025); + strings_column_wrapper strings(empty.begin(), empty.end()); + auto input = cudf::table_view{{strings, strings}}; + std::vector partition_ids(1025, 199); + fixed_width_column_wrapper map(partition_ids.begin(), partition_ids.end()); + auto [output, offsets] = cudf::partition(input, map, 200); + EXPECT_TRUE(std::all_of(offsets.begin(), offsets.end() - 1, [](auto x) { return x == 0; })); + EXPECT_EQ(offsets.back(), 1025); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, output->view()); +} + +TEST_F(PartitionTestNotTyped, StringPermutationOutputOffsetThreshold) +{ + struct restore_environment { + char const* name; + std::optional previous; + + explicit restore_environment(char const* variable) : name(variable) + { + if (auto const value = std::getenv(name)) { previous = value; } + } + + ~restore_environment() + { + if (previous) { + setenv(name, previous->c_str(), 1); + } else { + unsetenv(name); + } + } + }; + restore_environment threshold_restore{"LIBCUDF_LARGE_STRINGS_THRESHOLD"}; + restore_environment enabled_restore{"LIBCUDF_LARGE_STRINGS_ENABLED"}; + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_THRESHOLD", "2048", 1), 0); + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_ENABLED", "1", 1), 0); + + constexpr cudf::size_type rows = 1024; + std::vector values(rows, "a"); + strings_column_wrapper strings(values.begin(), values.end()); + auto const input = cudf::table_view{{strings, strings}}; + auto const ones = cuda::make_constant_iterator(1); + fixed_width_column_wrapper map(ones, ones + rows); + ASSERT_EQ(cudf::strings_column_view(input.column(0)).offsets().type().id(), cudf::type_id::INT32); + + // Change the threshold after construction to keep the input offsets INT32. + for (auto const threshold : {1025, 1024, 1023}) { + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_THRESHOLD", std::to_string(threshold).c_str(), 1), 0); + for (auto const enabled : {true, false}) { + SCOPED_TRACE(::testing::Message() << threshold << "/" << enabled); + ASSERT_EQ(setenv("LIBCUDF_LARGE_STRINGS_ENABLED", enabled ? "1" : "0", 1), 0); + if (!enabled && rows >= threshold) { + EXPECT_THROW(cudf::partition(input, map, 2), std::overflow_error); + } else { + auto [output, offsets] = cudf::partition(input, map, 2); + EXPECT_EQ(offsets, (std::vector{0, 0, rows})); + auto const expected_type = rows >= threshold ? cudf::type_id::INT64 : cudf::type_id::INT32; + for (auto const& column : output->view()) { + EXPECT_EQ(cudf::strings_column_view(column).offsets().type().id(), expected_type); + } + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(input, output->view()); + } + } + } +} + +TEST_F(PartitionTestNotTyped, StringPermutationCopyGridStride) +{ + // One row beyond the capped four-lane character-copy grid. Two columns use + // rank recomputation; four columns also exercise the fused copy's second pass. + constexpr cudf::size_type rows = 65'535 * 64 + 1; + for (auto const width : {4, 129}) { + auto const value = std::string(width, 'x'); + auto const values = cuda::make_constant_iterator(value.c_str()); + strings_column_wrapper strings(values, values + rows); + auto const ones = cuda::make_constant_iterator(1); + fixed_width_column_wrapper map(ones, ones + rows); + for (auto const count : {2, 4}) { + SCOPED_TRACE(::testing::Message() << width << "/" << count); + std::vector columns(count, strings); + auto const input = cudf::table_view{columns}; + auto [output, offsets] = cudf::partition(input, map, 2); + EXPECT_EQ(offsets, (std::vector{0, 0, rows})); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, output->view()); + } + } +} + +class PartitionLargeFixedWidthTest + : public cudf::test::BaseFixture, + public ::testing::WithParamInterface> {}; + +TEST_P(PartitionLargeFixedWidthTest, HostOracle) +{ + // Exceed the large-plan threshold and leave the last tile and mask word incomplete. + constexpr cudf::size_type rows = 1'000'003; + constexpr cudf::size_type slice_offset = 7; + auto const [count, partitions, null_percent, hot] = GetParam(); + auto stream_owner = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + cuda::stream_ref stream = stream_owner; + auto const valid = [null_percent](int row, int col) { + return (row * 17 + col * 13) % 100 >= null_percent; + }; + auto const value = [](int row, int col) { return int64_t{row} * (37 + col * 2) + col * 11; }; + std::vector> columns; + for (int c = 0; c < count; ++c) { + auto const values = + cudf::detail::make_counting_transform_iterator(0, [=](auto i) { return value(i, c); }); + auto const validity = + cudf::detail::make_counting_transform_iterator(0, [=](auto i) { return valid(i, c); }); + if (null_percent == 0) { + columns.push_back( + fixed_width_column_wrapper(values, values + rows + slice_offset, stream) + .release()); + } else { + columns.push_back( + fixed_width_column_wrapper(values, values + rows + slice_offset, validity, stream) + .release()); + } + } + auto const owner = cudf::table{std::move(columns)}; + auto const input = cudf::slice(owner.view(), {slice_offset, rows + slice_offset}, stream)[0]; + std::vector assignments(rows + 3, 0); + std::vector expected_offsets(partitions + 1, 0); + std::vector expected_all_null(partitions, 0); + for (int i = 0; i < rows; ++i) { + auto const p = hot && i % 10 != 0 ? partitions - 1 : (i * 17 + 3) % partitions; + assignments[i + 3] = p; + ++expected_offsets[p + 1]; + bool any_valid = false; + for (int c = 0; c < count; ++c) + any_valid |= valid(i + slice_offset, c); + if (!any_valid) ++expected_all_null[p]; + } + std::partial_sum(expected_offsets.begin(), expected_offsets.end(), expected_offsets.begin()); + fixed_width_column_wrapper map(assignments.begin(), assignments.end(), stream); + auto const map_view = cudf::slice(map, {3, rows + 3}, stream)[0]; + auto [output, offsets] = cudf::partition(input, map_view, partitions, stream, mr()); + ASSERT_EQ(offsets, expected_offsets); + ASSERT_EQ(output->num_rows(), rows); + ASSERT_EQ(output->num_columns(), count); + // This oracle uses host values and validity, independent of gather and GPU sorting. + using host_column = decltype(cudf::test::to_host(output->get_column(0), stream)); + std::vector actual; + for (int c = 0; c < count; ++c) { + EXPECT_EQ(output->get_column(c).type(), input.column(c).type()); + EXPECT_EQ(output->get_column(c).null_count(), input.column(c).null_count()); + actual.push_back(cudf::test::to_host(output->get_column(c), stream)); + } + auto const is_valid = [&](int c, int i) { + auto const& mask = actual[c].second; + return mask.empty() || (mask[i / 32] & (1u << (i % 32))) != 0; + }; + std::vector seen(rows, false); + for (int p = 0; p < partitions; ++p) { + cudf::size_type all_null = 0; + for (int i = offsets[p]; i < offsets[p + 1]; ++i) { + int first_valid = 0; + while (first_valid < count && !is_valid(first_valid, i)) + ++first_valid; + if (first_valid == count) { + ++all_null; + continue; + } + auto const source = + (actual[first_valid].first[i] - first_valid * 11) / (37 + first_valid * 2) - slice_offset; + ASSERT_GE(source, 0); + ASSERT_LT(source, rows); + ASSERT_EQ(assignments[source + 3], p); + ASSERT_FALSE(seen[source]); + seen[source] = true; + for (int c = 0; c < count; ++c) { + ASSERT_EQ(is_valid(c, i), valid(source + slice_offset, c)); + if (is_valid(c, i)) { ASSERT_EQ(actual[c].first[i], value(source + slice_offset, c)); } + } + } + ASSERT_EQ(all_null, expected_all_null[p]); + } + for (int i = 0; i < rows; ++i) { + bool any_valid = false; + for (int c = 0; c < count; ++c) + any_valid |= valid(i + slice_offset, c); + ASSERT_EQ(seen[i], any_valid); + } +} + +INSTANTIATE_TEST_SUITE_P(LargePlans, + PartitionLargeFixedWidthTest, + ::testing::Combine(::testing::Values(1, 2, 4, 8), + ::testing::Values(2, 100), + ::testing::Values(0, 10, 90), + ::testing::Bool())); + +TEST_F(PartitionTestNotTyped, LargeSparseStringCopy) +{ + constexpr cudf::size_type rows = 4'194'307; + auto const stream = cudf::get_default_stream(); + auto const first_values = cudf::detail::make_counting_transform_iterator(0, [](auto i) { + auto const slot = i % 1024; + auto const length = slot == 0 ? 129 : slot == 1 ? 65 : slot == 2 ? 7 : 0; + return std::string(length, static_cast('a' + (i / 1024) % 23)); + }); + auto const second_values = cudf::detail::make_counting_transform_iterator(0, [](auto i) { + auto const slot = i % 1024; + auto const length = slot == 0 ? 9 : slot == 1 ? 128 : slot == 3 ? 513 : 0; + return std::string(length, static_cast('A' + (i / 1024) % 23)); + }); + strings_column_wrapper first(first_values, first_values + rows, stream); + strings_column_wrapper second(second_values, second_values + rows, stream); + auto const assignments = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 2; }); + fixed_width_column_wrapper map(assignments, assignments + rows, stream); + auto const all_columns = cudf::table_view{{first, second}}; + constexpr auto half = (rows + 1) / 2; + auto const ordered = cudf::detail::make_counting_transform_iterator( + 0, [](auto i) { return i < half ? i * 2 : (i - half) * 2 + 1; }); + fixed_width_column_wrapper order(ordered, ordered + rows, stream); + for (auto const count : {1, 2}) { + SCOPED_TRACE(count); + auto const input = count == 1 ? all_columns.select({0}) : all_columns; + auto [output, offsets] = cudf::partition(input, map, 2, stream, mr()); + ASSERT_EQ(offsets, (std::vector{0, half, rows})); + auto expected = cudf::gather(input, + order, + cudf::out_of_bounds_policy::DONT_CHECK, + cudf::negative_index_policy::NOT_ALLOWED, + stream, + mr()); + for (auto p = 0; p < 2; ++p) { + auto expected_part = cudf::slice(expected->view(), {offsets[p], offsets[p + 1]}, stream)[0]; + auto actual_part = cudf::slice(output->view(), {offsets[p], offsets[p + 1]}, stream)[0]; + auto expected_sorted = cudf::sort(expected_part, {}, {}, stream, mr()); + auto actual_sorted = cudf::sort(actual_part, {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expected_sorted->view(), actual_sorted->view(), stream); + } + } +} + +TEST_F(PartitionTestNotTyped, StringVectorCopyBoundaries) +{ + constexpr cudf::size_type rows = 4099; + std::vector values(rows + 14); + std::vector assignments(rows); + std::vector gather; + std::vector offsets{0}; + std::vector lengths{0, 1, 3, 15, 16, 17, 127, 128, 129, 130, 143, 144, 255, 256, 257, 1024}; + for (std::size_t i = 0; i < values.size(); ++i) { + values[i] = std::string(lengths[i % lengths.size()], static_cast('a' + i % 26)); + if (values[i].size() > 3) { values[i][2] = '\0'; } + } + auto const valid = cuda::make_constant_iterator(true); + strings_column_wrapper strings(values.begin(), values.end(), valid); + for (auto i = 0; i < rows; ++i) { + assignments[i] = i % 10 == 0 ? 0 : 1; + } + for (auto p = 0; p < 2; ++p) { + for (auto i = 0; i < rows; ++i) { + if (assignments[i] == p) { gather.push_back(i); } + } + offsets.push_back(gather.size()); + } + fixed_width_column_wrapper map(assignments.begin(), assignments.end()); + fixed_width_column_wrapper indices(gather.begin(), gather.end()); + auto stream = cuda::stream{cuda::device_ref{rmm::get_current_cuda_device().value()}}; + for (auto const count : {1, 2, 4}) { + SCOPED_TRACE(count); + std::vector columns(count, strings); + auto const input = cudf::slice(cudf::table_view{columns}, {7, rows + 7})[0]; + auto expected = cudf::gather(input, indices); + cudf::get_default_stream().sync(); + auto [output, boundaries] = cudf::partition(input, map, 2, stream, mr()); + EXPECT_EQ(boundaries, offsets); + stream.sync(); + // An all-valid mask may be discarded; compare partition values and null semantics. + auto const split_points = std::vector{offsets[1]}; + auto const expected_parts = cudf::split(expected->view(), split_points); + auto const actual_parts = cudf::split(output->view(), split_points); + for (std::size_t i = 0; i < expected_parts.size(); ++i) { + auto sorted_expected = cudf::sort(expected_parts[i]); + auto sorted_actual = cudf::sort(actual_parts[i]); + CUDF_TEST_EXPECT_TABLES_EQUIVALENT(sorted_expected->view(), sorted_actual->view()); + } + } +} diff --git a/cpp/tests/streams/partitioning_test.cpp b/cpp/tests/streams/partitioning_test.cpp index cb7da2fe6b5b..0e7cd8448c40 100644 --- a/cpp/tests/streams/partitioning_test.cpp +++ b/cpp/tests/streams/partitioning_test.cpp @@ -1,11 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include +#include #include #include +#include #include #include @@ -15,11 +17,112 @@ #include #include +#include + +#include +#include + using cudf::test::fixed_width_column_wrapper; using cudf::test::strings_column_wrapper; class PartitionTest : public cudf::test::BaseFixture {}; +TEST_F(PartitionTest, MapRanks) +{ + constexpr cudf::size_type rows = 1025; + auto const stream = cudf::test::get_default_stream(); + auto const ids = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); + auto const partitions = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 2; }); + fixed_width_column_wrapper data(ids, ids + rows, stream); + fixed_width_column_wrapper map(partitions, partitions + rows, stream); + for (auto const& input : {cudf::table_view{{data}}, cudf::table_view{{data, data}}}) { + SCOPED_TRACE(input.num_columns()); + auto [output, offsets] = cudf::partition(input, map, 2, stream, mr()); + ASSERT_EQ(offsets, (std::vector{0, (rows + 1) / 2, rows})); + auto const [actual, mask] = cudf::test::to_host(output->get_column(0), stream); + for (cudf::size_type p = 0; p < 2; ++p) { + for (auto row = offsets[p]; row < offsets[p + 1]; ++row) { + EXPECT_EQ(actual[row] % 2, p); + } + } + auto sorted = cudf::sort(output->view(), {}, {}, stream, mr()); + CUDF_TEST_EXPECT_TABLES_EQUAL(input, sorted->view(), stream); + } +} + +TEST_F(PartitionTest, HeuristicFixedWidth) +{ + constexpr cudf::size_type rows = 1'000'003; + auto const stream = cudf::test::get_default_stream(); + auto const ids = cuda::counting_iterator{0}; + for (auto const null_mode : {0, 1, 2}) { + auto const valid = cudf::detail::make_counting_transform_iterator(0, [null_mode](auto i) { + return null_mode == 0 || (null_mode == 1 ? i % 100 != 0 : i % 10 == 0); + }); + fixed_width_column_wrapper data(ids, ids + rows, valid, stream); + for (auto const partitions : {2, 100, 256}) { + auto const assignments = cudf::detail::make_counting_transform_iterator( + 0, [partitions](auto i) { return i % partitions; }); + fixed_width_column_wrapper map(assignments, assignments + rows, stream); + std::vector expected(partitions + 1, 0); + for (auto p = 0; p < partitions; ++p) { + expected[p + 1] = expected[p] + rows / partitions + (p < rows % partitions); + } + for (auto const count : {1, 2, 8}) { + SCOPED_TRACE(::testing::Message() << null_mode << "/" << partitions << "/" << count); + std::vector columns(count, data); + auto [output, offsets] = + cudf::partition(cudf::table_view{columns}, map, partitions, stream, mr()); + EXPECT_EQ(offsets, expected); + EXPECT_EQ(output->num_rows(), rows); + EXPECT_EQ(output->num_columns(), count); + for (auto const& col : output->view()) { + EXPECT_EQ(col.null_count(), cudf::column_view{data}.null_count()); + } + } + } + } + fixed_width_column_wrapper narrow(ids, ids + rows, stream); + auto const partitions = + cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i % 256; }); + fixed_width_column_wrapper map(partitions, partitions + rows, stream); + std::vector columns(8, narrow); + auto result = cudf::partition(cudf::table_view{columns}, map, 256, stream, mr()); + EXPECT_EQ(result.second.back(), rows); +} + +TEST_F(PartitionTest, TrivialMapValues) +{ + constexpr cudf::size_type rows = 1025; + auto const stream = cudf::test::get_default_stream(); + auto const ids = cuda::counting_iterator{0}; + auto const invalid = cuda::make_constant_iterator(false); + auto const zero = cuda::make_constant_iterator(0); + fixed_width_column_wrapper nulls(ids, ids + rows, invalid, stream); + fixed_width_column_wrapper map(zero, zero + rows, stream); + auto const text = cuda::make_constant_iterator("value"); + strings_column_wrapper strings(text, text + rows, stream); + auto const empty_value = cuda::make_constant_iterator(""); + strings_column_wrapper empty(empty_value, empty_value + rows, stream); + for (auto const& input : {cudf::table_view{{nulls}}, + cudf::table_view{{nulls, strings}}, + cudf::table_view{std::vector{}, rows}}) { + auto [output, offsets] = cudf::partition(input, map, 1, stream, mr()); + EXPECT_EQ(offsets, (std::vector{0, rows})); + EXPECT_EQ(output->num_rows(), rows); + } + auto empty_result = cudf::partition(cudf::table_view{{empty, nulls}}, map, 256, stream, mr()); + EXPECT_EQ(empty_result.second.back(), rows); + EXPECT_EQ(empty_result.first->get_column(0).null_count(), 0); + auto mixed_result = cudf::partition(cudf::table_view{{empty, strings}}, map, 256, stream, mr()); + EXPECT_EQ(mixed_result.second.back(), rows); + auto [output, offsets] = cudf::partition(cudf::table_view{{nulls}}, map, 256, stream, mr()); + EXPECT_EQ(offsets.front(), 0); + EXPECT_EQ(offsets.back(), rows); + EXPECT_EQ(output->get_column(0).null_count(), rows); +} + TEST_F(PartitionTest, Struct) { fixed_width_column_wrapper A({1, 2}, {0, 1}); @@ -62,4 +165,32 @@ TEST_F(PartitionTest, ZeroPartitions) cudf::test::get_default_stream()); } +TEST_F(PartitionTest, HeuristicStrings) +{ + auto const stream = cudf::test::get_default_stream(); + for (auto const partitions : {2, 256}) { + for (auto const rows : {4099, 4'000'003}) { + auto const value = cuda::make_constant_iterator("a string with enough characters"); + auto const valid = cuda::make_constant_iterator(true); + strings_column_wrapper strings(value, value + rows, valid, stream); + auto const values = cudf::detail::make_counting_transform_iterator( + 0, [partitions](auto i) { return i % partitions; }); + fixed_width_column_wrapper map(values, values + rows, stream); + for (auto const count : {1, 2, 4}) { + SCOPED_TRACE(::testing::Message() << rows << "/" << count); + std::vector columns(count, strings); + auto [output, offsets] = + cudf::partition(cudf::table_view{columns}, map, partitions, stream, mr()); + std::vector expected(partitions + 1, 0); + for (auto p = 0; p < partitions; ++p) { + expected[p + 1] = expected[p] + rows / partitions + (p < rows % partitions); + } + EXPECT_EQ(offsets, expected); + EXPECT_EQ(output->num_rows(), rows); + EXPECT_EQ(output->num_columns(), count); + } + } + } +} + CUDF_TEST_PROGRAM_MAIN()