diff --git a/ci/validate_wheel.sh b/ci/validate_wheel.sh index 3668749316a4..14d7121260d2 100755 --- a/ci/validate_wheel.sh +++ b/ci/validate_wheel.sh @@ -21,11 +21,11 @@ PYDISTCHECK_ARGS=( if [[ "${package_dir}" == "python/libcudf" ]]; then if [[ "${RAPIDS_CUDA_MAJOR}" == "12" ]]; then PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '675M' + --max-allowed-size-compressed '700M' ) else PYDISTCHECK_ARGS+=( - --max-allowed-size-compressed '325M' + --max-allowed-size-compressed '350M' ) fi elif [[ "${package_dir}" != "python/cudf" ]] && \ diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f89ecf862ebf..91ff5c193fd9 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -496,6 +496,7 @@ add_library( src/filling/repeat.cu src/filling/sequence.cu src/groupby/common/m2_var_std.cu + src/groupby/common/utils.cpp src/groupby/groupby.cu src/groupby/hash/compute_global_memory_aggs.cu src/groupby/hash/compute_global_memory_aggs_null.cu @@ -539,6 +540,16 @@ add_library( src/groupby/sort/host_udf_aggregation.cpp src/groupby/sort/scan.cpp src/groupby/sort/sort_helper.cu + src/groupby/streaming_groupby.cpp + src/groupby/streaming_groupby/aggregate.cu + src/groupby/streaming_groupby/impl.cu + src/groupby/streaming_groupby/insert.cu + src/groupby/streaming_groupby/insert_first.cu + src/groupby/streaming_groupby/insert_first_nested.cu + src/groupby/streaming_groupby/insert_nested.cu + src/groupby/streaming_groupby/insert_subsequent.cu + src/groupby/streaming_groupby/insert_subsequent_nested.cu + src/groupby/streaming_groupby/merge.cu src/hash/md5_hash.cu src/hash/murmurhash3_x86_32.cu src/hash/murmurhash3_x64_128.cu diff --git a/cpp/benchmarks/groupby/group_max.cpp b/cpp/benchmarks/groupby/group_max.cpp index d837cfac6659..706299b61041 100644 --- a/cpp/benchmarks/groupby/group_max.cpp +++ b/cpp/benchmarks/groupby/group_max.cpp @@ -84,11 +84,85 @@ void bench_groupby_max(nvbench::state& state, nvbench::type_list) template void bench_groupby_max_cardinality(nvbench::state& state, nvbench::type_list) { - auto constexpr num_rows = 20'000'000; - auto constexpr null_probability = 0.; - auto const cardinality = static_cast(state.get_int64("cardinality")); + auto const num_rows = static_cast(state.get_int64("num_rows")); + auto const cardinality = static_cast(state.get_int64("cardinality")); + auto const num_aggregations = state.get_int64("num_aggregations"); + auto const is_streaming = state.get_string("api") == "streaming"; + + // TODO: streaming groupby reuses the cudf hash element_aggregator, which has + // no decimal128 MIN/MAX/SUM specialization (no native 128-bit atomics). The + // stateless `normal` path falls back to sort-based aggregation, but streaming + // has no fallback and rejects the request. Re-enable once streaming has a + // non-atomic aggregator path or 128-bit atomics gain hardware support. + if (is_streaming && std::is_same_v) { + state.skip("streaming groupby does not support decimal128 MAX/MIN/SUM"); + return; + } - groupby_max_helper(state, num_rows, cardinality, null_probability); + auto const keys = [&] { + data_profile const profile = + data_profile_builder() + .cardinality(cardinality) + .no_validity() + .distribution(cudf::type_to_id(), distribution_id::UNIFORM, 0, num_rows); + return create_random_column(cudf::type_to_id(), row_count{num_rows}, profile); + }(); + + auto const make_values = [&]() { + auto builder = data_profile_builder().cardinality(0).no_validity().distribution( + cudf::type_to_id(), distribution_id::UNIFORM, 0, num_rows); + return create_random_column( + cudf::type_to_id(), row_count{num_rows}, data_profile{builder}); + }; + + std::vector> val_cols; + for (int64_t i = 0; i < num_aggregations; i++) { + val_cols.emplace_back(make_values()); + } + + auto keys_view = keys->view(); + + auto const mem_stats_logger = cudf::memory_stats_logger(); + state.set_cuda_stream(nvbench::make_cuda_stream_view(cudf::get_default_stream().value())); + + if (is_streaming) { + std::vector all_columns = {keys_view, keys_view, keys_view}; + for (auto const& vc : val_cols) { + all_columns.push_back(vc->view()); + } + auto const full_table = cudf::table_view(all_columns); + + std::vector key_indices = {0, 1, 2}; + std::vector requests; + for (int64_t i = 0; i < num_aggregations; i++) { + cudf::groupby::streaming_aggregation_request req; + req.column_index = static_cast(3 + i); + req.aggregation = cudf::make_max_aggregation(); + requests.push_back(std::move(req)); + } + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { + auto sgb = cudf::groupby::streaming_groupby(key_indices, requests, num_rows); + sgb.aggregate(full_table); + auto const result = sgb.finalize(); + }); + } else { + std::vector requests; + for (int64_t i = 0; i < num_aggregations; i++) { + requests.emplace_back(); + requests[i].values = val_cols[i]->view(); + requests[i].aggregations.push_back(cudf::make_max_aggregation()); + } + state.exec(nvbench::exec_tag::sync, [&](nvbench::launch& launch) { + auto gb_obj = cudf::groupby::groupby(cudf::table_view({keys_view, keys_view, keys_view})); + auto const result = gb_obj.aggregate(requests); + }); + } + + auto const elapsed_time = state.get_summary("nv/cold/time/gpu/mean").get_float64("value"); + state.add_element_count( + static_cast(num_rows * num_aggregations) / elapsed_time / 1'000'000., "Mrows/s"); + state.add_buffer_size( + mem_stats_logger.peak_memory_usage(), "peak_memory_usage", "peak_memory_usage"); } NVBENCH_BENCH_TYPES(bench_groupby_max, @@ -102,5 +176,7 @@ NVBENCH_BENCH_TYPES(bench_groupby_max, NVBENCH_BENCH_TYPES(bench_groupby_max_cardinality, NVBENCH_TYPE_AXES(nvbench::type_list)) .set_name("groupby_max_cardinality") + .add_int64_axis("num_rows", {20'000'000}) .add_int64_axis("num_aggregations", {1, 2, 3, 4, 5, 6, 7, 8}) - .add_int64_axis("cardinality", {20, 50, 100, 1'000, 10'000, 100'000, 1'000'000}); + .add_int64_axis("cardinality", {20, 50, 100, 1'000, 10'000, 100'000, 1'000'000}) + .add_string_axis("api", {"normal", "streaming"}); diff --git a/cpp/include/cudf/detail/aggregation/aggregation.hpp b/cpp/include/cudf/detail/aggregation/aggregation.hpp index cd94625c2dcd..b848f4417b57 100644 --- a/cpp/include/cudf/detail/aggregation/aggregation.hpp +++ b/cpp/include/cudf/detail/aggregation/aggregation.hpp @@ -14,6 +14,8 @@ #include #include +#include + #include #include #include @@ -1341,9 +1343,9 @@ data_type target_type(data_type source_type, aggregation::Kind k); * @tparam k The aggregation to perform */ template -constexpr inline bool is_valid_aggregation() +CUDF_HOST_DEVICE constexpr inline bool is_valid_aggregation() { - return (not std::is_void_v>); + return (not cuda::std::is_void_v>); } /** diff --git a/cpp/include/cudf/groupby.hpp b/cpp/include/cudf/groupby.hpp index e814f7ec76bb..21a039f85690 100644 --- a/cpp/include/cudf/groupby.hpp +++ b/cpp/include/cudf/groupby.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -137,7 +137,7 @@ class groupby { * result is the same order as was specified in the request. * * The returned `table` contains the group labels for each group, i.e., the - * unique rows from `keys`. Element `i` across all aggregation results + * distinct rows from `keys`. Element `i` across all aggregation results * belongs to the group at row `i` in the group labels table. * * The order of the rows in the group labels is arbitrary. Furthermore, @@ -169,7 +169,7 @@ class groupby { * perform * @param stream CUDA stream used for device memory operations and kernel launches. * @param mr Device memory resource used to allocate the returned table and columns' device memory - * @return Pair containing the table with each group's unique key and + * @return Pair containing the table with each group's distinct key and * a vector of aggregation_results for each request in the same order as * specified in `requests`. */ @@ -407,6 +407,189 @@ class groupby { rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr); }; + +/** + * @brief Request for a single streaming groupby aggregation on a column. + * + * Analogous to `aggregation_request` but identifies the value column by index rather than + * by `column_view`, since data arrives in batches after construction, and carries exactly + * one aggregation per request. + * + * `column_index` refers to the position of the value column in the `table_view` passed + * to `streaming_groupby::aggregate()`. Multiple aggregations on the same column are + * expressed as separate requests (e.g., `[{col, sum}, {col, mean}]`). Internal + * deduplication ensures redundant computations are shared automatically. + */ +struct streaming_aggregation_request { + size_type column_index; ///< Index of the value column + std::unique_ptr aggregation; ///< Desired aggregation +}; + +/** + * @brief Stateful streaming groupby that accumulates partial aggregates across batches. + * + * `streaming_groupby` and the stateless `groupby` serve different use cases. Use + * the stateless `groupby` for single-shot aggregation when all input fits in memory + * at once. Use `streaming_groupby` when input arrives over multiple batches and + * memory efficiency matters: peak memory does not scale with the number of *input* + * rows, because only the distinct keys seen so far and one aggregation slot per + * group are kept across batches. Arbitrarily long high-duplicate streams therefore + * accumulate without running out of memory. + * + * If memory is not a concern, concatenating all batches and calling the stateless + * `groupby` once is also a valid choice. Reach for `streaming_groupby` when + * (a) the cumulative input does not fit in memory, or (b) partial-state aggregation + * across distributed workers (`merge()`) is part of the workload. + * + * Per-batch cost is O(batch_size): each batch does direct hash table insertion + * and in-place aggregation updates against the persistent state. Partial states + * can be combined via `merge()`, and final results are produced via `finalize()`. + * + * The `max_distinct_keys` parameter sets the upper bound on the number of distinct key + * combinations across the lifetime of this object. The persistent state is sized to + * `max_distinct_keys` (constant for the lifetime of the object); the stored distinct + * keys grow with the number of distinct keys actually seen, so the incremental key + * storage is O(`distinct_keys()` × key_size) and does not scale with cumulative input + * rows. + * + * Cumulative input rows are not bounded — only cumulative distinct keys. A single + * batch may also not exceed `max_distinct_keys` rows; this is an implementation + * limit because each in-flight batch row is encoded as `max_distinct_keys + row_idx` + * inside the hash set, which must fit in `cudf::size_type`. + * + * All column types (including variable-width types such as strings, lists, and structs) + * are supported for key columns. Only hash-based aggregation kinds are supported; use + * `is_streaming_groupby_supported()` to query a specific (value type, aggregation kind) + * combination. + * + * Supported aggregation kinds: + * SUM, SUM_OF_SQUARES, PRODUCT, MIN, MAX, COUNT_VALID, COUNT_ALL, + * MEAN, M2, VARIANCE, STD + * + * @throws std::invalid_argument for unsupported aggregation kinds + * @throws std::invalid_argument if a single batch exceeds `max_distinct_keys` rows + * @throws cudf::logic_error if cumulative distinct keys exceed `max_distinct_keys` + */ +class streaming_groupby { + public: + streaming_groupby() = delete; + ~streaming_groupby(); + streaming_groupby(streaming_groupby const&) = delete; + streaming_groupby& operator=(streaming_groupby const&) = delete; + + /** @brief Move constructor. */ + streaming_groupby(streaming_groupby&&) noexcept; + + /** + * @brief Move assignment operator. + * @return Reference to this object. + */ + streaming_groupby& operator=(streaming_groupby&&) noexcept; + + /** + * @brief Construct a streaming groupby object with a persistent hash table. + * + * @param key_indices Indices of columns in the data table that serve as groupby keys + * @param requests The aggregations to perform and which columns to aggregate + * @param max_distinct_keys Upper bound on distinct key combinations. The hash set, + * companion vectors, and aggregation results table are all sized to this + * capacity. Cumulative input rows are not bounded. + * @param null_handling Indicates whether rows in keys that contain NULL values should be included + * + * @throws std::invalid_argument if `max_distinct_keys <= 0` + * @throws std::invalid_argument if any requested aggregation kind is unsupported + */ + explicit streaming_groupby(host_span key_indices, + host_span requests, + size_type max_distinct_keys, + null_policy null_handling = null_policy::EXCLUDE); + + /** + * @brief Feed a batch of data into the streaming aggregation. + * + * Batch keys are inserted into the persistent hash set and aggregation results + * are updated atomically. The input `data` table is not referenced after this + * call returns. + * + * @param data Table containing both key and value columns + * @param stream CUDA stream used for device memory operations and kernel launches + * + * @throws std::invalid_argument if `data.num_rows()` exceeds `max_distinct_keys` + * @throws cudf::logic_error if cumulative distinct keys exceed `max_distinct_keys` + */ + void aggregate(table_view const& data, rmm::cuda_stream_view stream = cudf::get_default_stream()); + + /** + * @brief Merge another streaming_groupby's accumulated partial state into this one. + * + * Extracts the other object's accumulated intermediate state and merges it into this + * object's persistent hash table. The other object is not modified. + * Both objects must have been constructed with compatible aggregation requests, + * and this object must have had at least one `aggregate()` call. + * + * @param other The streaming_groupby whose partial state to merge + * @param stream CUDA stream used for device memory operations and kernel launches + * + * @throws std::invalid_argument if the other object has more distinct keys than + * `max_distinct_keys` + * @throws cudf::logic_error if this object has not been initialized via `aggregate()` + * @throws cudf::logic_error if distinct keys exceed `max_distinct_keys` after merge + */ + void merge(streaming_groupby const& other, + rmm::cuda_stream_view stream = cudf::get_default_stream()); + + /** + * @brief Finalize the accumulated partial aggregates into final results. + * + * For most aggregation kinds the partial state is the final result. For kinds like + * MEAN, VARIANCE, or STD, a finalization step converts the internal partial representation + * (e.g., sum+count) into the user-facing result. + * + * This does not modify the internal state; `aggregate()` may be called again afterward. + * + * @param stream CUDA stream used for device memory operations and kernel launches + * @param mr Device memory resource used to allocate the returned table and columns + * @return Pair of distinct keys table and a vector of aggregation_results (one per request) + * + * @throws cudf::logic_error if no data has been accumulated + */ + [[nodiscard]] std::pair, std::vector> finalize( + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()) const; + + /** + * @brief Returns the number of distinct keys accumulated so far. + * + * Returns 0 before any successful `aggregate()` or `merge()` call. + * + * @return The current count of distinct keys in the persistent hash table + */ + [[nodiscard]] size_type distinct_keys() const noexcept; + + private: + struct impl; + std::unique_ptr _impl; + + void do_aggregate(table_view const& data, rmm::cuda_stream_view stream); + void do_merge(streaming_groupby const& other, rmm::cuda_stream_view stream); + [[nodiscard]] std::pair, std::vector> do_finalize( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const; +}; + +/** + * @brief Returns true if `streaming_groupby` supports the given value type and + * aggregation kind combination. + * + * Use this to query support without constructing a `streaming_groupby`. A `true` + * return implies that an `aggregate()` call with a value column of `values_type` and + * an aggregation of `kind` will not be rejected on type/kind grounds. + * + * @param values_type Type of the value column the aggregation would run on + * @param kind Aggregation kind + * @return True if the combination is supported, false otherwise + */ +[[nodiscard]] bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind); + /** @} */ } // namespace groupby } // namespace CUDF_EXPORT cudf diff --git a/cpp/src/groupby/common/utils.cpp b/cpp/src/groupby/common/utils.cpp new file mode 100644 index 000000000000..93eb7029e28f --- /dev/null +++ b/cpp/src/groupby/common/utils.cpp @@ -0,0 +1,34 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "groupby/common/utils.hpp" + +#include +#include +#include + +namespace cudf::groupby::detail { + +std::pair compute_row_bitmask(table_view const& keys, + rmm::cuda_stream_view stream) +{ + auto const mr = cudf::get_current_device_resource_ref(); + if (keys.num_columns() == 0 || !cudf::has_nulls(keys)) { + return {rmm::device_buffer{0, stream, mr}, nullptr}; + } + // Single-column fast path: reuse the column's null mask directly. + if (keys.num_columns() == 1) { + auto const& col = keys.column(0); + if (col.offset() == 0) { return {rmm::device_buffer{0, stream, mr}, col.null_mask()}; } + auto buf = cudf::copy_bitmask(col, stream, mr); + auto ptr = static_cast(buf.data()); + return {std::move(buf), ptr}; + } + auto [buf, null_count] = cudf::bitmask_and(keys, stream, mr); + if (null_count == 0) { return {rmm::device_buffer{0, stream, mr}, nullptr}; } + return {std::move(buf), static_cast(buf.data())}; +} + +} // namespace cudf::groupby::detail diff --git a/cpp/src/groupby/common/utils.hpp b/cpp/src/groupby/common/utils.hpp index 7c232e098520..e856ae38ecd0 100644 --- a/cpp/src/groupby/common/utils.hpp +++ b/cpp/src/groupby/common/utils.hpp @@ -1,21 +1,26 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include #include #include +#include +#include #include #include +#include +#include + #include +#include #include -namespace cudf { -namespace groupby { -namespace detail { +namespace cudf::groupby::detail { template inline std::vector extract_results(host_span requests, @@ -47,6 +52,34 @@ inline std::vector extract_results(host_span compute_row_bitmask( + table_view const& keys, rmm::cuda_stream_view stream); + +/// Whether the given aggregation kind is supported by hash-based groupby. +constexpr bool is_hash_aggregation(aggregation::Kind k) +{ + switch (k) { + case aggregation::SUM: + case aggregation::SUM_WITH_OVERFLOW: + case aggregation::SUM_OF_SQUARES: + case aggregation::PRODUCT: + case aggregation::MIN: + case aggregation::MAX: + case aggregation::COUNT_VALID: + case aggregation::COUNT_ALL: + case aggregation::ARGMIN: + case aggregation::ARGMAX: + case aggregation::MEAN: + case aggregation::M2: + case aggregation::STD: + case aggregation::VARIANCE: return true; + default: return false; + } +} + +} // namespace cudf::groupby::detail diff --git a/cpp/src/groupby/hash/compute_groupby.cu b/cpp/src/groupby/hash/compute_groupby.cu index 44f4fd33c86a..3e94dd7b7c9b 100644 --- a/cpp/src/groupby/hash/compute_groupby.cu +++ b/cpp/src/groupby/hash/compute_groupby.cu @@ -5,6 +5,7 @@ #include "compute_groupby.hpp" #include "compute_single_pass_aggs.hpp" +#include "groupby/common/utils.hpp" #include "hash_compound_agg_finalizer.hpp" #include "helpers.cuh" #include "output_utils.hpp" @@ -58,36 +59,11 @@ std::unique_ptr compute_groupby(table_view const& keys, { auto const num_keys = keys.num_rows(); - [[maybe_unused]] auto const [row_bitmask_data, row_bitmask] = - [&]() -> std::pair { - if (!skip_rows_with_nulls) { - return {rmm::device_buffer{0, stream, cudf::get_current_device_resource_ref()}, nullptr}; - } - - if (keys.num_columns() == 1) { - auto const& keys_col = keys.column(0); - // Only use the input null mask directly if the keys table was not sliced. - if (keys_col.offset() == 0) { - return {rmm::device_buffer{0, stream, cudf::get_current_device_resource_ref()}, - keys_col.null_mask()}; - } - // If the keys table was sliced, we need to copy the null mask to ensure its first bit aligns - // with the first row of the keys table. - auto null_mask_data = - cudf::copy_bitmask(keys_col, stream, cudf::get_current_device_resource_ref()); - auto const null_mask = static_cast(null_mask_data.data()); - return {std::move(null_mask_data), null_mask}; - } - - auto [null_mask_data, null_count] = - cudf::bitmask_and(keys, stream, cudf::get_current_device_resource_ref()); - if (null_count == 0) { - return {rmm::device_buffer{0, stream, cudf::get_current_device_resource_ref()}, nullptr}; - } - - auto const null_mask = static_cast(null_mask_data.data()); - return {std::move(null_mask_data), null_mask}; - }(); + [[maybe_unused]] auto [row_bitmask_data, row_bitmask] = + skip_rows_with_nulls + ? cudf::groupby::detail::compute_row_bitmask(keys, stream) + : std::pair{ + rmm::device_buffer{0, stream, cudf::get_current_device_resource_ref()}, nullptr}; auto const cached_hashes = [&]() -> rmm::device_uvector { auto const num_columns = diff --git a/cpp/src/groupby/hash/groupby.cu b/cpp/src/groupby/hash/groupby.cu index a49446abdafb..7a53f0cb0016 100644 --- a/cpp/src/groupby/hash/groupby.cu +++ b/cpp/src/groupby/hash/groupby.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -24,47 +24,11 @@ #include #include -#include #include #include namespace cudf::groupby::detail::hash { namespace { -/** - * @brief List of aggregation operations that can be computed with a hash-based implementation. - * - * For single pass aggregations, the supported operations are the ones that can be atomically - * updated: SUM, SUM_WITH_OVERFLOW, SUM_OF_SQUARES, PRODUCT, MIN, MAX, COUNT_VALID, COUNT_ALL. - * For compound aggregations, the supported operations are the ones that depends on the single pass - * aggregations above: ARGMIN(MIN), ARGMAX(MAX), MEAN(SUM, COUNT_VALID), M2/STD/VARIANCE(M2, - * COUNT_VALID). - */ -const auto hash_aggregations = std::unordered_set{// Single pass aggregations: - aggregation::SUM, - aggregation::SUM_WITH_OVERFLOW, - aggregation::SUM_OF_SQUARES, - aggregation::PRODUCT, - aggregation::MIN, - aggregation::MAX, - aggregation::COUNT_VALID, - aggregation::COUNT_ALL, - // Compound aggregations: - aggregation::ARGMIN, - aggregation::ARGMAX, - aggregation::MEAN, - aggregation::M2, - aggregation::STD, - aggregation::VARIANCE}; - -/** - * @brief Indicates whether the specified aggregation operation can be computed - * with a hash-based implementation. - * - * @param t The aggregation operation to verify - * @return true `t` is valid for a hash based groupby - * @return false `t` is invalid for a hash based groupby - */ -bool is_hash_aggregation(aggregation::Kind t) { return hash_aggregations.contains(t); } std::unique_ptr
dispatch_groupby(table_view const& keys, host_span requests, diff --git a/cpp/src/groupby/streaming_groupby.cpp b/cpp/src/groupby/streaming_groupby.cpp new file mode 100644 index 000000000000..0aea2baf1a25 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby.cpp @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace cudf::groupby { + +void streaming_groupby::aggregate(table_view const& data, rmm::cuda_stream_view stream) +{ + CUDF_FUNC_RANGE(); + do_aggregate(data, stream); +} + +void streaming_groupby::merge(streaming_groupby const& other, rmm::cuda_stream_view stream) +{ + CUDF_FUNC_RANGE(); + do_merge(other, stream); +} + +std::pair, std::vector> streaming_groupby::finalize( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const +{ + CUDF_FUNC_RANGE(); + return do_finalize(stream, mr); +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/aggregate.cu b/cpp/src/groupby/streaming_groupby/aggregate.cu new file mode 100644 index 000000000000..a8568bbff4f5 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/aggregate.cu @@ -0,0 +1,66 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "common.cuh" +#include "groupby/hash/single_pass_functors.cuh" + +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include + +namespace cudf::groupby { + +void streaming_groupby::impl::do_aggregate(table_view const& data, rmm::cuda_stream_view stream) +{ + CUDF_EXPECTS(!_invalidated, + "streaming_groupby is in an invalidated state from a prior failure; " + "no further aggregate()/merge() is allowed. finalize() may still be called."); + + auto const batch_size = data.num_rows(); + if (batch_size == 0) { return; } + + CUDF_EXPECTS(batch_size <= _max_distinct_keys, + "Batch size (" + std::to_string(batch_size) + ") exceeds max_distinct_keys (" + + std::to_string(_max_distinct_keys) + ").", + std::invalid_argument); + + CUDF_EXPECTS(static_cast(_max_distinct_keys) + static_cast(batch_size) <= + static_cast(std::numeric_limits::max()), + "Transient key encoding (max_distinct_keys + batch_size) would overflow size_type.", + std::invalid_argument); + + if (!_initialized) { initialize(data, stream); } + + auto const batch_keys = data.select(_key_indices); + + update_nullable_state(batch_keys); + + if (!_key_set) { create_key_set(stream); } + + auto result = probe_and_insert(batch_keys, stream); + + auto const values_view = data.select(_value_col_indices); + auto const d_values = table_device_view::create(values_view, stream); + + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const num_agg_cols = static_cast(_agg_kinds.size()); + thrust::for_each_n( + rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + static_cast(batch_size) * num_agg_cols, + detail::hash::compute_single_pass_aggs_dense_output_fn{ + result.target_indices.begin(), _d_agg_kinds.data(), *d_values, *_d_agg_results}); +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/common.cuh b/cpp/src/groupby/streaming_groupby/common.cuh new file mode 100644 index 000000000000..5ff5b1c2078e --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/common.cuh @@ -0,0 +1,376 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "groupby/hash/helpers.cuh" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include + +#include +#include + +namespace cudf::groupby { + +/* + * Companion location for a stored dense ID: which compacted batch table the key + * lives in (`first`) and the row index within that table (`second`). Packed into + * one 8-byte pair so the comparator does a single load instead of two. + */ +using key_location_t = cuda::std::pair; + +using streaming_probing_scheme_t = + cuco::linear_probing>; + +using streaming_set_t = cuco::static_set, + cuda::thread_scope_device, + cuda::std::equal_to, + streaming_probing_scheme_t, + rmm::mr::polymorphic_allocator, + cuco::storage>; + +/* + * Comparator for the first batch only. All slot values are transient (encoded as + * `max_distinct_keys + row_idx`) since no dense IDs exist yet; this wrapper subtracts + * the offset and delegates to the batch self-equality. + */ +template +struct first_batch_comparator { + RowEqT batch_self_eq; + size_type max_distinct_keys; + + __device__ bool operator()(size_type lhs, size_type rhs) const noexcept + { + return batch_self_eq(lhs - max_distinct_keys, rhs - max_distinct_keys); + } +}; + +/* + * N-table comparator for the persistent hash set. + * + * Slot values < max_distinct_keys are "stored" dense IDs resolved via the companion + * vector key_loc[id] = {batch_id, row_within_batch} to a (compacted_batch_table, + * row) location. Slot values >= max_distinct_keys are transient batch values: row index + * = value - max_distinct_keys in the current batch table. The transient encoding lives + * only for the duration of one probe_and_insert call; new keys are rewritten to + * dense IDs before the next batch's insertion. + * + * Cross-comparators are pre-built as device_row_comparator(batch, compacted[k]) + * and stored in a device array. Self-comparisons use batch_self_eq. + */ +template +struct n_table_comparator { + RowEqT batch_self_eq; ///< Self-comparator on the current batch table + RowEqT const* cross_eqs; ///< Device array [num_compacted_batches]: batch vs compacted[k] + key_location_t const* key_loc; ///< {batch_id, row_in_compacted} per dense ID + size_type max_distinct_keys; ///< Threshold: idx >= max_distinct_keys is a transient batch value + + __device__ bool operator()(size_type lhs, size_type rhs) const noexcept + { + bool const lhs_is_batch = (lhs >= max_distinct_keys); + bool const rhs_is_batch = (rhs >= max_distinct_keys); + + if (lhs_is_batch && rhs_is_batch) { + return batch_self_eq(lhs - max_distinct_keys, rhs - max_distinct_keys); + } + if (lhs_is_batch) { + auto const loc = key_loc[rhs]; + return cross_eqs[loc.first](lhs - max_distinct_keys, loc.second); + } + if (rhs_is_batch) { + auto const loc = key_loc[lhs]; + return cross_eqs[loc.first](rhs - max_distinct_keys, loc.second); + } + // During probe_and_insert, at least one operand is always the batch row being + // inserted (transient-encoded), so two dense IDs cannot be compared here. + CUDF_UNREACHABLE("n_table_comparator received two dense-ID operands"); + } +}; + +/* + * Predicate used by `thrust::copy_if` to compact the batch row indices of newly + * inserted keys in a single pass. Each invocation calls `set_ref.insert_and_find` + * for one batch row, records the resident slot's value in `target_indices[row_idx]` + * and the slot offset in `slot_offsets[row_idx]` (or `CUDF_SIZE_TYPE_SENTINEL` for + * null/excluded rows), and returns true iff this row inserted a new key. When + * the return is true `copy_if` writes `row_idx` to `batch_local_indices`; the + * total insert count is the iterator distance returned by `copy_if`. + * + * If the batch produces no new keys, target_indices is already final. If new + * keys exist, Pass 2 rewrites the affected slots and the caller re-reads + * target_indices via slot_offsets in a single transform. + */ +template +struct insert_and_check_fn { + mutable SetRef set_ref; + bitmask_type const* row_bitmask; + size_type max_distinct_keys; + size_type const* base; + size_type* target_indices; + size_type* slot_offsets; + + __device__ bool operator()(size_type row_idx) const + { + if (row_bitmask && !cudf::bit_is_set(row_bitmask, row_idx)) { + target_indices[row_idx] = cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + slot_offsets[row_idx] = cudf::detail::CUDF_SIZE_TYPE_SENTINEL; + return false; + } + auto const [iter, inserted] = set_ref.insert_and_find(max_distinct_keys + row_idx); + target_indices[row_idx] = *iter; + slot_offsets[row_idx] = static_cast(iter - base); + return inserted; + } +}; + +/* + * Per-row hash producer used to populate the precomputed batch hash cache. + * Returns 0 (a dummy value never read) for rows excluded by the null bitmask + * under `null_policy::EXCLUDE`, avoiding wasted hash work for rows that won't + * probe the set. + */ +template +struct conditional_hash_fn { + RowHasher row_hasher; + bitmask_type const* row_bitmask; + + __device__ hash_value_type operator()(size_type i) const noexcept + { + if (row_bitmask && !cudf::bit_is_set(row_bitmask, i)) { return hash_value_type{0}; } + return row_hasher(i); + } +}; + +/* + * Hasher backed by a precomputed cache, indexed by `idx - offset`. + * In streaming_groupby `offset = max_distinct_keys`, so transient batch values + * `max_distinct_keys + row_idx` resolve to `cache[row_idx]`. + */ +struct offset_cache_hasher { + hash_value_type const* cache; + size_type offset; + __device__ hash_value_type operator()(size_type idx) const noexcept + { + return cache[idx - offset]; + } +}; + +/* + * For each newly discovered key (dense rank `r` within this batch): + * 1. Rewrite its slot from transient encoding to its dense ID. Pass 2 runs + * after Pass 1's kernel completes and each slot has a single writer, so a + * plain store is sufficient. + * 2. Write the (batch_id, row) pair to the companion vector at the dense ID. + */ +struct finalize_new_key_fn { + size_type const* + batch_local_indices; ///< batch-local row indices of new keys [new_distinct_keys] + size_type* base; ///< hash set storage base + size_type const* slot_offsets; ///< slot offset per batch row [batch_size] + key_location_t* key_loc; ///< {batch_id, row_in_compacted} per dense ID + size_type batch_id; ///< the index of this batch in _compacted_batches + size_type dense_id_offset; ///< first dense ID assigned to this batch's new keys + + __device__ void operator()(size_type r) const + { + auto const dense_id = dense_id_offset + r; + auto const batch_local = batch_local_indices[r]; + + *(base + slot_offsets[batch_local]) = dense_id; + *(key_loc + dense_id) = key_location_t{batch_id, r}; + } +}; + +struct update_transient_target_indices_fn { + size_type const* base; + size_type const* slot_offsets; + size_type max_distinct_keys; + size_type* target_indices; + + __device__ void operator()(size_type i) const + { + if (target_indices[i] >= max_distinct_keys) { target_indices[i] = *(base + slot_offsets[i]); } + } +}; + +template +auto build_cross_comparators( + std::shared_ptr const& preprocessed_batch, + std::vector> const& + preprocessed_batches, + cudf::nullate::DYNAMIC has_null, + rmm::cuda_stream_view stream) +{ + using eq_t = cudf::detail::row::equality::device_row_comparator< + has_nested_columns, + cudf::nullate::DYNAMIC, + cudf::detail::row::equality::nan_equal_physical_equality_comparator>; + + auto const n = static_cast(preprocessed_batches.size()); + auto const temp_mr = cudf::get_current_device_resource_ref(); + + std::vector h_eqs; + h_eqs.reserve(n); + for (size_type k = 0; k < n; ++k) { + auto const cross_cmp = cudf::detail::row::equality::two_table_comparator{ + preprocessed_batch, preprocessed_batches[k]}; + auto const adapter = cross_cmp.equal_to(has_null, null_equality::EQUAL); + h_eqs.push_back(adapter.comparator); + } + + return cudf::detail::make_device_uvector_async(h_eqs, stream, temp_mr); +} + +/// The impl struct for streaming_groupby. Defined in impl.cu. +struct streaming_groupby::impl { + std::vector _key_indices; + std::vector _requests_clone; + size_type _max_distinct_keys; + null_policy _null_handling; + + bool _initialized{false}; + /// Set true once an `aggregate()` / `merge()` call has thrown after touching the + /// hash set. Subsequent `aggregate()` / `merge()` calls fail fast; only + /// `finalize()` may still be called to recover partial results. + bool _invalidated{false}; + /* + * Number of distinct keys accumulated so far. Also serves as the high-water + * mark of dense IDs in the persistent hash set: stored slot values are in + * [0, _distinct_keys). + */ + size_type _distinct_keys{0}; + bool _has_nullable_keys{false}; + bool _has_nested_keys{false}; + + // -- Compacted batch storage -- + std::vector> _compacted_batches; + std::vector> + _preprocessed_batches; + + /// Empty (0-row) table preserving the key schema, used by gather_distinct_keys when + /// no batches produced any groups (e.g. first batch empty or fully null-excluded). + std::unique_ptr
_empty_key_schema; + + /// Companion vector indexed by dense ID, sized to max_distinct_keys. + /// Each entry is {batch_id, row_in_compacted_batch}. + std::unique_ptr> _key_loc; + + std::vector _request_first_agg_offset; + std::vector _agg_kinds; + std::vector> _agg_objects; + std::vector _is_agg_intermediate; + bool _has_compound_aggs{false}; + + /* + * Aggregation results table, pre-allocated to max_distinct_keys rows. + * Indexed by dense ID (== row index). + */ + std::unique_ptr
_agg_results; + /* + * Cached mutable_table_device_view of `_agg_results`. `_agg_results` is allocated + * once at initialize() and never resized, so this device-side descriptor can be + * built once and reused on every aggregate() / merge() call rather than rebuilt + * (which requires a host-to-device copy of the column metadata). + */ + std::unique_ptr _d_agg_results; + std::vector _value_col_indices; + rmm::device_uvector _d_agg_kinds; + + std::unique_ptr _key_set; + + [[nodiscard]] size_type num_keys() const { return static_cast(_key_indices.size()); } + [[nodiscard]] bool has_state() const { return _initialized && _distinct_keys > 0; } + + impl(host_span key_indices, + host_span requests, + size_type max_distinct_keys, + null_policy null_handling); + + void initialize(table_view const& data, rmm::cuda_stream_view stream); + void create_key_set(rmm::cuda_stream_view stream); + void update_nullable_state(table_view const& batch_keys); + + struct batch_insert_result { + rmm::device_uvector target_indices; + size_type new_insertions; + rmm::device_buffer bitmask_buffer; + }; + + batch_insert_result probe_and_insert(table_view const& batch_keys, rmm::cuda_stream_view stream); + + /* + * Template implementation of probe_and_insert, split by has_nested. + * Defined in insert.cuh, instantiated in insert.cu and insert_nested.cu. + */ + template + batch_insert_result probe_and_insert_impl(table_view const& batch_keys, + rmm::cuda_stream_view stream); + + /* + * Two helpers split off probe_and_insert_impl for compile-time parallelism. + * Each builds the cuco set_ref + comparator and runs the fused insert/compact + * via thrust::copy_if. Split into separate TUs because both cuco and thrust + * algorithm template instantiations dominate compile time. + * + * Defined in insert_first.cuh / insert_subsequent.cuh, instantiated in + * insert_first{,_nested}.cu and insert_subsequent{,_nested}.cu respectively. + */ + template + size_type probe_and_insert_first_batch( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + + template + size_type probe_and_insert_subsequent( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + + void do_aggregate(table_view const& data, rmm::cuda_stream_view stream); + + [[nodiscard]] std::unique_ptr
gather_agg_results(rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const; + [[nodiscard]] std::unique_ptr
gather_distinct_keys( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const; + + [[nodiscard]] std::pair, std::vector> do_finalize( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const; + + void do_merge(impl const& other, rmm::cuda_stream_view stream); +}; + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/impl.cu b/cpp/src/groupby/streaming_groupby/impl.cu new file mode 100644 index 000000000000..bd442e36d60d --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/impl.cu @@ -0,0 +1,405 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "common.cuh" +#include "groupby/common/utils.hpp" +#include "groupby/hash/extract_single_pass_aggs.hpp" +#include "groupby/hash/hash_compound_agg_finalizer.hpp" +#include "groupby/hash/output_utils.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace cudf::groupby { + +namespace { + +void validate_requests(host_span requests) +{ + for (auto const& req : requests) { + CUDF_EXPECTS(req.aggregation != nullptr, + "streaming_aggregation_request must have a non-null aggregation.", + std::invalid_argument); + CUDF_EXPECTS(detail::is_hash_aggregation(req.aggregation->kind) && + req.aggregation->kind != aggregation::ARGMIN && + req.aggregation->kind != aggregation::ARGMAX, + "Unsupported aggregation kind for streaming groupby. " + "ARGMIN/ARGMAX are not supported because row indices are batch-local.", + std::invalid_argument); + } +} + +// Group streaming requests by `column_index` so multiple aggregations on the same +// column become a single `aggregation_request{values, [aggs...]}`. `extract_single_pass_aggs` +// then dedups repeated simple kinds within the group (e.g. SUM and MEAN both want SUM). +std::vector build_aggregation_requests( + host_span requests, table_view const& data) +{ + std::vector result; + std::unordered_map col_to_idx; + col_to_idx.reserve(requests.size()); + for (auto const& req : requests) { + auto cloned = std::unique_ptr{ + dynamic_cast(req.aggregation->clone().release())}; + auto const [it, inserted] = + col_to_idx.try_emplace(req.column_index, static_cast(result.size())); + if (inserted) { + aggregation_request ar; + ar.values = data.column(req.column_index); + ar.aggregations.push_back(std::move(cloned)); + result.push_back(std::move(ar)); + } else { + result[it->second].aggregations.push_back(std::move(cloned)); + } + } + return result; +} + +} // namespace + +streaming_groupby::impl::impl(host_span key_indices, + host_span requests, + size_type max_distinct_keys, + null_policy null_handling) + : _max_distinct_keys{max_distinct_keys}, + _null_handling{null_handling}, + _d_agg_kinds{0, rmm::cuda_stream_default, cudf::get_current_device_resource_ref()}, + _d_agg_results{nullptr, +[](mutable_table_device_view*) {}} +{ + CUDF_EXPECTS(max_distinct_keys > 0, "max_distinct_keys must be positive.", std::invalid_argument); + if (!key_indices.empty()) { _key_indices.assign(key_indices.begin(), key_indices.end()); } + validate_requests(requests); + + for (auto const& req : requests) { + streaming_aggregation_request clone; + clone.column_index = req.column_index; + clone.aggregation = std::unique_ptr{ + dynamic_cast(req.aggregation->clone().release())}; + _requests_clone.push_back(std::move(clone)); + } +} + +void streaming_groupby::impl::initialize(table_view const& data, rmm::cuda_stream_view stream) +{ + auto const mr = cudf::get_current_device_resource_ref(); + + // Detect nested key columns (struct, list) for comparator template dispatch. + std::vector key_cols; + key_cols.reserve(_key_indices.size()); + for (auto idx : _key_indices) { + key_cols.push_back(data.column(idx)); + } + _has_nested_keys = cudf::detail::has_nested_columns(table_view{key_cols}); + + std::vector> empty_key_cols; + empty_key_cols.reserve(key_cols.size()); + for (auto const& kc : key_cols) { + empty_key_cols.push_back(cudf::empty_like(kc)); + } + _empty_key_schema = std::make_unique
(std::move(empty_key_cols)); + + auto agg_requests = build_aggregation_requests(_requests_clone, data); + + // TODO: streaming aggregation reuses the cudf hash-groupby element_aggregator, + // so it inherits the same atomic-support requirement. In particular, decimal128 + // MIN/MAX/SUM falls through to CUDF_UNREACHABLE because __int128 is not + // lock-free atomic. Stateless cudf::groupby falls back to sort-based groupby + // in that case; streaming has no such fallback. Until streaming has a + // non-atomic aggregator path (or 128-bit atomics gain hardware support), gate + // by the same predicate to fail loudly instead of silently producing garbage. + CUDF_EXPECTS(detail::hash::can_use_hash_groupby(agg_requests), + "streaming_groupby does not support this combination of value type and " + "aggregation kind (e.g. decimal128 MIN/MAX/SUM require 128-bit atomics).", + std::invalid_argument); + + auto [values_view, agg_kinds_hv, agg_objects, is_intermediate, has_compound] = + detail::hash::extract_single_pass_aggs(agg_requests, stream); + + _agg_kinds.assign(agg_kinds_hv.begin(), agg_kinds_hv.end()); + _agg_objects = std::move(agg_objects); + _is_agg_intermediate = std::move(is_intermediate); + _has_compound_aggs = has_compound; + + // Reject aggregation kinds that are unsupported in streaming after decomposition. + for (auto k : _agg_kinds) { + CUDF_EXPECTS(k != aggregation::ARGMIN && k != aggregation::ARGMAX, + "Streaming groupby does not support MIN/MAX on variable-width types " + "(internally decomposed to ARGMIN/ARGMAX).", + std::invalid_argument); + CUDF_EXPECTS(k != aggregation::SUM_WITH_OVERFLOW, + "Streaming groupby does not support SUM_WITH_OVERFLOW " + "(struct intermediate cannot be merged across batches).", + std::invalid_argument); + } + + _agg_results = detail::hash::create_results_table( + _max_distinct_keys, values_view, _agg_kinds, _is_agg_intermediate, stream, mr); + + // Cache the mutable_table_device_view once; the underlying table is fixed-size and + // never reallocated, so the device-side descriptor stays valid for the whole + // lifetime of this impl. + { + auto raii = mutable_table_device_view::create(*_agg_results, stream); + _d_agg_results = + decltype(_d_agg_results){raii.release(), +[](mutable_table_device_view* t) { t->destroy(); }}; + } + + _d_agg_kinds = cudf::detail::make_device_uvector_async(_agg_kinds, stream, mr); + + // Map each column in `values_view` back to its index in `data`. + _value_col_indices.reserve(values_view.num_columns()); + for (size_type i = 0; i < values_view.num_columns(); ++i) { + auto const& col = values_view.column(i); + bool found = false; + for (size_type c = 0; c < data.num_columns(); ++c) { + if (cudf::detail::is_shallow_equivalent(data.column(c), col)) { + _value_col_indices.push_back(c); + found = true; + break; + } + } + CUDF_EXPECTS(found, "Internal error: agg column not found in input data."); + } + + // For each user streaming request, locate the offset in the dedup'd `_agg_kinds` + // where its first decomposed simple agg lives (matched by kind + column identity). + _request_first_agg_offset.reserve(_requests_clone.size()); + for (auto const& req : _requests_clone) { + auto const& target_col = data.column(req.column_index); + auto const first_kind = + detail::hash::get_simple_aggregations(*req.aggregation, target_col.type()).front(); + bool found = false; + for (size_type k = 0; k < static_cast(_agg_kinds.size()); ++k) { + if (_agg_kinds[k] == first_kind && + cudf::detail::is_shallow_equivalent(values_view.column(k), target_col)) { + _request_first_agg_offset.push_back(k); + found = true; + break; + } + } + CUDF_EXPECTS(found, "Internal error: request's first simple agg not found."); + } + + // Companion vector: indexed by dense ID, one {batch_id, row} entry per distinct key. + _key_loc = std::make_unique>(_max_distinct_keys, stream, mr); + + _initialized = true; +} + +void streaming_groupby::impl::create_key_set(rmm::cuda_stream_view stream) +{ + _key_set = std::make_unique( + cuco::extent{static_cast(_max_distinct_keys)}, + cudf::detail::CUCO_DESIRED_LOAD_FACTOR, + cuco::empty_key{cudf::detail::CUDF_SIZE_TYPE_SENTINEL}, + cuda::std::equal_to{}, + streaming_probing_scheme_t{cudf::hashing::detail::default_hash{}}, + cuco::thread_scope_device, + cuco::storage{}, + rmm::mr::polymorphic_allocator{}, + stream.value()); +} + +void streaming_groupby::impl::update_nullable_state(table_view const& batch_keys) +{ + if (_has_nullable_keys) return; + for (size_type c = 0; c < num_keys(); ++c) { + if (batch_keys.column(c).nullable()) { + _has_nullable_keys = true; + return; + } + } +} + +std::unique_ptr
streaming_groupby::impl::gather_agg_results( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const +{ + // The results we care about are dense in `[0, _distinct_keys)` and can be extracted by + // slice+copy. + auto const sliced = + cudf::detail::slice(_agg_results->view(), {0, _distinct_keys}, stream).front(); + return std::make_unique
(sliced, stream, mr); +} + +std::unique_ptr
streaming_groupby::impl::gather_distinct_keys( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const +{ + if (_compacted_batches.empty()) { + return std::make_unique
(_empty_key_schema->view(), stream, mr); + } + if (_compacted_batches.size() == 1) { + return std::make_unique
(_compacted_batches[0]->view(), stream, mr); + } + std::vector distinct_keys(_compacted_batches.size()); + std::transform(_compacted_batches.begin(), + _compacted_batches.end(), + distinct_keys.begin(), + [](auto const& batch) { return batch->view(); }); + return cudf::concatenate(distinct_keys, stream, mr); +} + +std::pair, std::vector> +streaming_groupby::impl::do_finalize(rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) const +{ + CUDF_EXPECTS(_initialized, "Cannot finalize streaming_groupby with no accumulated data."); + + auto keys = gather_distinct_keys(stream, mr); + auto agg_gathered = gather_agg_results(stream, mr); + + // Group user requests by their target column in `agg_gathered` so the cache layout + // produced by `extract_single_pass_aggs` matches the dedup'd `agg_gathered`. Uses + // linear search on a small `group_offsets` vector since the number of distinct + // columns is typically small. + auto const agg_gathered_view = agg_gathered->view(); + std::vector column_grouped; + std::vector group_offsets; + for (size_t i = 0; i < _requests_clone.size(); ++i) { + auto const offset = _request_first_agg_offset[i]; + auto cloned = std::unique_ptr{ + dynamic_cast(_requests_clone[i].aggregation->clone().release())}; + auto const it = std::find(group_offsets.begin(), group_offsets.end(), offset); + if (it == group_offsets.end()) { + aggregation_request ar; + ar.values = agg_gathered_view.column(offset); + ar.aggregations.push_back(std::move(cloned)); + column_grouped.push_back(std::move(ar)); + group_offsets.push_back(offset); + } else { + column_grouped[std::distance(group_offsets.begin(), it)].aggregations.push_back( + std::move(cloned)); + } + } + + auto [values_view_fin, agg_kinds_fin, agg_objects_fin, is_intermediate_fin, has_compound_fin] = + detail::hash::extract_single_pass_aggs(column_grouped, stream); + + cudf::detail::result_cache cache(_agg_kinds.size()); + detail::hash::finalize_output(values_view_fin, agg_objects_fin, agg_gathered, &cache, stream); + + if (_has_compound_aggs) { + // Compute compound aggs (MEAN/STD/VARIANCE/M2/...) into the cache. The cache itself + // dedupes: skip if (column, kind) is already there from a prior agg in the group. + for (auto const& req : column_grouped) { + auto const finalizer = + detail::hash::hash_compound_agg_finalizer(req.values, &cache, nullptr, stream, mr); + for (auto const& agg : req.aggregations) { + if (cache.has_result(req.values, *agg)) continue; + cudf::detail::aggregation_dispatcher(agg->kind, finalizer, *agg); + } + } + } + + // User-1:1 lookup keys so `extract_results` returns results in the order of the + // original user requests, sharing cache entries across requests on the same column. + std::vector user_requests; + user_requests.reserve(_requests_clone.size()); + for (size_t i = 0; i < _requests_clone.size(); ++i) { + aggregation_request ar; + ar.values = agg_gathered_view.column(_request_first_agg_offset[i]); + ar.aggregations.push_back(std::unique_ptr{ + dynamic_cast(_requests_clone[i].aggregation->clone().release())}); + user_requests.push_back(std::move(ar)); + } + + return {std::move(keys), + detail::extract_results( + host_span{user_requests}, cache, stream, mr)}; +} + +streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_insert( + table_view const& batch_keys, rmm::cuda_stream_view stream) +{ + if (_has_nested_keys) { + return probe_and_insert_impl(batch_keys, stream); + } else { + return probe_and_insert_impl(batch_keys, stream); + } +} + +// Constructor, destructor, and move ops require full impl definition. +streaming_groupby::streaming_groupby(host_span key_indices, + host_span requests, + size_type max_distinct_keys, + null_policy null_handling) + : _impl{std::make_unique(key_indices, requests, max_distinct_keys, null_handling)} +{ +} + +streaming_groupby::~streaming_groupby() = default; + +streaming_groupby::streaming_groupby(streaming_groupby&&) noexcept = default; + +streaming_groupby& streaming_groupby::operator=(streaming_groupby&&) noexcept = default; + +// Private member functions defined here (requires full impl definition). +// The public API wrappers in streaming_groupby.cpp call these. +void streaming_groupby::do_aggregate(table_view const& data, rmm::cuda_stream_view stream) +{ + _impl->do_aggregate(data, stream); +} + +void streaming_groupby::do_merge(streaming_groupby const& other, rmm::cuda_stream_view stream) +{ + _impl->do_merge(*other._impl, stream); +} + +std::pair, std::vector> streaming_groupby::do_finalize( + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const +{ + return _impl->do_finalize(stream, mr); +} + +size_type streaming_groupby::distinct_keys() const noexcept { return _impl->_distinct_keys; } + +bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind) +{ + switch (kind) { + case aggregation::SUM: + case aggregation::PRODUCT: + case aggregation::COUNT_VALID: + case aggregation::COUNT_ALL: + case aggregation::MEAN: + case aggregation::M2: + case aggregation::VARIANCE: + case aggregation::STD: + case aggregation::SUM_OF_SQUARES: break; + case aggregation::MIN: + case aggregation::MAX: + // Variable-width / compound types decompose to ARGMIN/ARGMAX (unsupported). + if (!cudf::is_fixed_width(values_type)) { return false; } + break; + default: return false; + } + // decimal128 SUM/MIN/MAX needs 128-bit atomics, which aren't supported. + if ((kind == aggregation::SUM || kind == aggregation::MIN || kind == aggregation::MAX) && + values_type.id() == type_id::DECIMAL128) { + return false; + } + return true; +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert.cu b/cpp/src/groupby/streaming_groupby/insert.cu new file mode 100644 index 000000000000..621b42e69917 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert.cu @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert.cuh" + +namespace cudf::groupby { + +template streaming_groupby::impl::batch_insert_result +streaming_groupby::impl::probe_and_insert_impl(table_view const& batch_keys, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert.cuh b/cpp/src/groupby/streaming_groupby/insert.cuh new file mode 100644 index 000000000000..39be36b4af7e --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert.cuh @@ -0,0 +1,149 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common.cuh" +#include "groupby/common/utils.hpp" + +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include + +#include + +namespace cudf::groupby { + +template +streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_insert_impl( + table_view const& batch_keys, rmm::cuda_stream_view stream) +{ + auto const batch_size = batch_keys.num_rows(); + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const has_null = cudf::nullate::DYNAMIC{_has_nullable_keys}; + + // Preprocess batch for row operators. + auto preprocessed_batch = cudf::detail::row::hash::preprocessed_table::create(batch_keys, stream); + auto const batch_hasher_obj = cudf::detail::row::hash::row_hasher{preprocessed_batch}; + auto const d_batch_hash = batch_hasher_obj.device_hasher(has_null); + + // Compute the null-exclusion bitmask first so the hash cache pass can skip + // hashing rows that will be excluded + auto const skip_rows_with_nulls = _has_nullable_keys && _null_handling == null_policy::EXCLUDE; + auto [bitmask_buffer, batch_bitmask] = + skip_rows_with_nulls + ? detail::compute_row_bitmask(batch_keys, stream) + : std::pair{rmm::device_buffer{0, stream}, nullptr}; + + // Precompute batch hash values. Caching is faster than inlining the row hasher + rmm::device_uvector batch_hash_cache(batch_size, stream, temp_mr); + thrust::transform(rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + cuda::counting_iterator(batch_size), + batch_hash_cache.begin(), + conditional_hash_fn{d_batch_hash, batch_bitmask}); + + // Pass 1 — fused insert_and_find + compact via `thrust::copy_if`. + // Per-row work in the predicate: writes target_indices and slot_offsets. + // Output: batch row indices of newly inserted keys, compacted into batch_local_indices. + // Count: iterator distance returned by copy_if inside the helper. + // slot_offsets stores 4-byte slot offsets (vs. 8-byte raw pointers) to halve temp memory. + rmm::device_uvector target_indices(batch_size, stream, temp_mr); + rmm::device_uvector slot_offsets(batch_size, stream, temp_mr); + rmm::device_uvector batch_local_indices(batch_size, stream, temp_mr); + + // First batch has no compacted batches yet, so all slot values are transient (>= + // _max_distinct_keys) and only the batch-self equality branch of n_table_comparator + // can fire. Dispatch to a dedicated helper that uses first_batch_comparator to skip + // the cross-table dispatch, the cross-comparator build, and the dense-ID branches. + // The two helpers live in separate TUs to parallelize the heavy cuco/thrust template + // instantiations. + size_type const new_distinct_keys = + _compacted_batches.empty() + ? probe_and_insert_first_batch(preprocessed_batch, + has_null, + batch_bitmask, + batch_hash_cache.data(), + batch_size, + target_indices.data(), + slot_offsets.data(), + batch_local_indices.data(), + stream) + : probe_and_insert_subsequent(preprocessed_batch, + has_null, + batch_bitmask, + batch_hash_cache.data(), + batch_size, + target_indices.data(), + slot_offsets.data(), + batch_local_indices.data(), + stream); + batch_local_indices.resize(new_distinct_keys, stream); + + if (new_distinct_keys > 0) { + // Bound check: the hash set has already been written above (transient slot values), + // so on failure the object is left invalidated; further aggregate()/merge() calls + // will throw immediately while finalize() can still recover partial results. + if (_distinct_keys + new_distinct_keys > _max_distinct_keys) { + _invalidated = true; + CUDF_FAIL("Distinct key count (" + std::to_string(_distinct_keys + new_distinct_keys) + + ") would exceed max_distinct_keys (" + std::to_string(_max_distinct_keys) + ")."); + } + + // Gather compacted distinct keys from the batch. + auto compacted = cudf::detail::gather(batch_keys, + batch_local_indices, + out_of_bounds_policy::DONT_CHECK, + cudf::negative_index_policy::NOT_ALLOWED, + stream, + temp_mr); + + auto preprocessed_compacted = + cudf::detail::row::hash::preprocessed_table::create(compacted->view(), stream); + + // Store the compacted batch. + auto const new_batch_id = static_cast(_compacted_batches.size()); + auto const dense_id_offset = _distinct_keys; + _compacted_batches.push_back(std::move(compacted)); + _preprocessed_batches.push_back(preprocessed_compacted); + + // Pass 2 — fused: rewrites slots from transient encoding to dense IDs and + // writes the {batch_id, row} companion entry. + auto* const base = _key_set->data(); + thrust::for_each_n(rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + new_distinct_keys, + finalize_new_key_fn{batch_local_indices.data(), + base, + slot_offsets.data(), + _key_loc->data(), + new_batch_id, + dense_id_offset}); + + thrust::for_each_n(rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + batch_size, + update_transient_target_indices_fn{ + base, slot_offsets.data(), _max_distinct_keys, target_indices.data()}); + + _distinct_keys += new_distinct_keys; + } + // If new_distinct_keys == 0, target_indices is already final from Pass 1 — every + // slot held a dense ID at probe time, so *iter was already the correct dense ID. + + return batch_insert_result{ + std::move(target_indices), new_distinct_keys, std::move(bitmask_buffer)}; +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_first.cu b/cpp/src/groupby/streaming_groupby/insert_first.cu new file mode 100644 index 000000000000..baf7b48dfb51 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_first.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert_first.cuh" + +namespace cudf::groupby { + +template size_type streaming_groupby::impl::probe_and_insert_first_batch( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_first.cuh b/cpp/src/groupby/streaming_groupby/insert_first.cuh new file mode 100644 index 000000000000..74f26ce14a60 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_first.cuh @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common.cuh" + +#include + +#include +#include + +#include +#include + +namespace cudf::groupby { + +template +size_type streaming_groupby::impl::probe_and_insert_first_batch( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream) +{ + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const batch_self_cmp = cudf::detail::row::equality::self_comparator{preprocessed_batch}; + auto const batch_self_eq = batch_self_cmp.equal_to(has_null, null_equality::EQUAL); + auto const hasher = offset_cache_hasher{batch_hash_cache, _max_distinct_keys}; + auto const set_ref_base = _key_set->ref(cuco::op::insert_and_find).rebind_hash_function(hasher); + auto const first_batch_cmp = first_batch_comparator{batch_self_eq, _max_distinct_keys}; + auto* const base = _key_set->data(); + + auto const out_end = + thrust::copy_if(rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + cuda::counting_iterator(batch_size), + batch_local_indices, + insert_and_check_fn{set_ref_base.rebind_key_eq(first_batch_cmp), + batch_bitmask, + _max_distinct_keys, + base, + target_indices, + slot_offsets}); + return static_cast(out_end - batch_local_indices); +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_first_nested.cu b/cpp/src/groupby/streaming_groupby/insert_first_nested.cu new file mode 100644 index 000000000000..98225875f76d --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_first_nested.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert_first.cuh" + +namespace cudf::groupby { + +template size_type streaming_groupby::impl::probe_and_insert_first_batch( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_nested.cu b/cpp/src/groupby/streaming_groupby/insert_nested.cu new file mode 100644 index 000000000000..33117248ee8c --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_nested.cu @@ -0,0 +1,14 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert.cuh" + +namespace cudf::groupby { + +template streaming_groupby::impl::batch_insert_result +streaming_groupby::impl::probe_and_insert_impl(table_view const& batch_keys, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_subsequent.cu b/cpp/src/groupby/streaming_groupby/insert_subsequent.cu new file mode 100644 index 000000000000..a5ea95b7572a --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_subsequent.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert_subsequent.cuh" + +namespace cudf::groupby { + +template size_type streaming_groupby::impl::probe_and_insert_subsequent( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_subsequent.cuh b/cpp/src/groupby/streaming_groupby/insert_subsequent.cuh new file mode 100644 index 000000000000..4f14691d3f9d --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_subsequent.cuh @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "common.cuh" + +#include + +#include +#include + +#include +#include + +namespace cudf::groupby { + +template +size_type streaming_groupby::impl::probe_and_insert_subsequent( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream) +{ + auto const temp_mr = cudf::get_current_device_resource_ref(); + auto const batch_self_cmp = cudf::detail::row::equality::self_comparator{preprocessed_batch}; + auto const batch_self_eq = batch_self_cmp.equal_to(has_null, null_equality::EQUAL); + auto const hasher = offset_cache_hasher{batch_hash_cache, _max_distinct_keys}; + auto const set_ref_base = _key_set->ref(cuco::op::insert_and_find).rebind_hash_function(hasher); + auto* const base = _key_set->data(); + + auto d_cross_eqs = build_cross_comparators( + preprocessed_batch, _preprocessed_batches, has_null, stream); + auto const comparator = + n_table_comparator{batch_self_eq, d_cross_eqs.data(), _key_loc->data(), _max_distinct_keys}; + + auto const out_end = thrust::copy_if(rmm::exec_policy_nosync(stream, temp_mr), + cuda::counting_iterator(0), + cuda::counting_iterator(batch_size), + batch_local_indices, + insert_and_check_fn{set_ref_base.rebind_key_eq(comparator), + batch_bitmask, + _max_distinct_keys, + base, + target_indices, + slot_offsets}); + return static_cast(out_end - batch_local_indices); +} + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/insert_subsequent_nested.cu b/cpp/src/groupby/streaming_groupby/insert_subsequent_nested.cu new file mode 100644 index 000000000000..8c0a6faab895 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/insert_subsequent_nested.cu @@ -0,0 +1,21 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "insert_subsequent.cuh" + +namespace cudf::groupby { + +template size_type streaming_groupby::impl::probe_and_insert_subsequent( + std::shared_ptr const& preprocessed_batch, + cudf::nullate::DYNAMIC has_null, + bitmask_type const* batch_bitmask, + hash_value_type const* batch_hash_cache, + size_type batch_size, + size_type* target_indices, + size_type* slot_offsets, + size_type* batch_local_indices, + rmm::cuda_stream_view stream); + +} // namespace cudf::groupby diff --git a/cpp/src/groupby/streaming_groupby/merge.cu b/cpp/src/groupby/streaming_groupby/merge.cu new file mode 100644 index 000000000000..34cd7abe39a6 --- /dev/null +++ b/cpp/src/groupby/streaming_groupby/merge.cu @@ -0,0 +1,144 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "common.cuh" + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include + +namespace cudf::groupby { + +namespace { + +/** + * @brief Element aggregator for merging intermediate results. + */ +struct merge_element_aggregator { + template + __device__ void operator()(mutable_column_device_view target, + size_type target_index, + column_device_view source, + size_type source_index) const noexcept + { + if constexpr (!cudf::detail::is_valid_aggregation()) { + return; + } else { + if constexpr (k != aggregation::COUNT_ALL) { + if (source.is_null(source_index)) { return; } + } + if constexpr (!(k == aggregation::COUNT_VALID || k == aggregation::COUNT_ALL)) { + if (target.is_null(target_index)) { target.set_valid(target_index); } + } + + if constexpr (k == aggregation::COUNT_VALID || k == aggregation::COUNT_ALL) { + using Target = cudf::detail::target_type_t; + cudf::detail::atomic_add(&target.element(target_index), + source.element(source_index)); + } else if constexpr (k == aggregation::SUM_OF_SQUARES) { + using Target = cudf::detail::target_type_t; + cudf::detail::atomic_add(&target.element(target_index), + static_cast(source.element(source_index))); + } else { + cudf::detail::update_target_element{}( + target, target_index, source, source_index); + } + } + } +}; + +struct merge_single_pass_aggs_fn { + size_type const* target_indices; + aggregation::Kind const* aggs; + table_device_view source_values; + mutable_table_device_view target_values; + + __device__ void operator()(int64_t idx) const + { + auto const num_rows = source_values.num_rows(); + auto const source_row_idx = static_cast(idx % num_rows); + if (auto const target_row_idx = target_indices[source_row_idx]; + target_row_idx != cudf::detail::CUDF_SIZE_TYPE_SENTINEL) { + auto const col_idx = static_cast(idx / num_rows); + auto const& source_col = source_values.column(col_idx); + auto const& target_col = target_values.column(col_idx); + cudf::detail::dispatch_type_and_aggregation(source_col.type(), + aggs[col_idx], + merge_element_aggregator{}, + target_col, + target_row_idx, + source_col, + source_row_idx); + } + } +}; + +} // namespace + +void streaming_groupby::impl::do_merge(impl const& other, rmm::cuda_stream_view stream) +{ + CUDF_EXPECTS(!_invalidated, + "streaming_groupby is in an invalidated state from a prior failure; " + "no further aggregate()/merge() is allowed. finalize() may still be called."); + CUDF_EXPECTS(!other._invalidated, "Cannot merge from an invalidated streaming_groupby."); + + if (!other._initialized || !other.has_state()) { return; } + CUDF_EXPECTS(_initialized, + "Cannot merge into an uninitialized streaming_groupby. " + "Call aggregate() at least once before merge()."); + CUDF_EXPECTS(other._distinct_keys <= _max_distinct_keys, + "Merge source distinct keys (" + std::to_string(other._distinct_keys) + + ") exceeds max_distinct_keys (" + std::to_string(_max_distinct_keys) + ").", + std::invalid_argument); + CUDF_EXPECTS(other._agg_kinds == _agg_kinds, + "Cannot merge streaming_groupby objects with different aggregation schemas.", + std::invalid_argument); + CUDF_EXPECTS(other._key_indices == _key_indices, + "Cannot merge streaming_groupby objects with different key column indices.", + std::invalid_argument); + CUDF_EXPECTS(other._null_handling == _null_handling, + "Cannot merge streaming_groupby objects with different null handling policies.", + std::invalid_argument); + + auto const mr = cudf::get_current_device_resource_ref(); + + auto other_keys = other.gather_distinct_keys(stream, mr); + auto const other_key_view = other_keys->view(); + auto const other_distinct_keys = other._distinct_keys; + if (other_distinct_keys == 0) { return; } + + update_nullable_state(other_key_view); + + if (!_key_set) { create_key_set(stream); } + + auto result = probe_and_insert(other_key_view, stream); + + // Merge aggregation values using dense target indices. We only read from + // `other._agg_results`; no need to deep-copy the source rows like keys. + auto const other_aggs_view = + cudf::detail::slice(other._agg_results->view(), {0, other_distinct_keys}, stream).front(); + auto const d_source = table_device_view::create(other_aggs_view, stream); + + auto const num_agg_cols = static_cast(_agg_kinds.size()); + thrust::for_each_n( + rmm::exec_policy_nosync(stream, mr), + cuda::counting_iterator(0), + static_cast(other_distinct_keys) * num_agg_cols, + merge_single_pass_aggs_fn{ + result.target_indices.begin(), _d_agg_kinds.data(), *d_source, *_d_agg_results}); +} + +} // namespace cudf::groupby diff --git a/cpp/src/io/json/nested_json_gpu.cu b/cpp/src/io/json/nested_json_gpu.cu index 7329e86975e5..f22c6cbd7e4a 100644 --- a/cpp/src/io/json/nested_json_gpu.cu +++ b/cpp/src/io/json/nested_json_gpu.cu @@ -1637,15 +1637,6 @@ std::pair, rmm::device_uvector> ge (format == tokenizer_pda::json_format_cfg_t::JSON_LINES_RECOVER ? 1 : 0); // Perform a PDA-transducer pass - // Compute the maximum amount of tokens that can possibly be emitted for a given input size - // Worst case ratio of tokens per input char is given for a struct with an empty field name, that - // may be arbitrarily deeply nested: {"":_}, where '_' is a placeholder for any JSON value, - // possibly another such struct. That is, 6 tokens for 5 chars (plus chars and tokens of '_') - std::size_t constexpr min_chars_per_struct = 5; - std::size_t constexpr max_tokens_per_struct = 6; - auto const max_token_out_count = - cudf::util::div_rounding_up_safe(json_in.size(), min_chars_per_struct) * max_tokens_per_struct + - delimiter_offset; cudf::detail::device_scalar num_written_tokens{ stream, cudf::get_current_device_resource_ref()}; @@ -1680,9 +1671,6 @@ std::pair, rmm::device_uvector> ge tokens_indices = std::move(filtered_tokens_indices); } - CUDF_EXPECTS(num_total_tokens <= max_token_out_count, - "Generated token count exceeds the expected token count"); - return std::make_pair(std::move(tokens), std::move(tokens_indices)); } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index bae781f36bda..ce7bfdbed0f0 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -154,6 +154,7 @@ ConfigureTest( groupby/replace_nulls_tests.cpp groupby/shift_tests.cpp groupby/std_tests.cpp + groupby/streaming_groupby_test.cpp groupby/structs_tests.cpp groupby/sum_of_squares_tests.cpp groupby/sum_scan_tests.cpp diff --git a/cpp/tests/groupby/groupby_test_util.cpp b/cpp/tests/groupby/groupby_test_util.cpp index d6a0608c3e9b..7c0599be8b5b 100644 --- a/cpp/tests/groupby/groupby_test_util.cpp +++ b/cpp/tests/groupby/groupby_test_util.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,7 @@ void test_single_agg(cudf::column_view const& keys, std::vector const& column_order, std::vector const& null_precedence, cudf::sorted reference_keys_are_sorted, + test_streaming use_streaming, std::source_location const& location) { SCOPED_TRACE("Original failure location: " + std::string{location.file_name()} + ":" + @@ -47,39 +48,85 @@ void test_single_agg(cudf::column_view const& keys, } }(); - std::vector requests; - requests.emplace_back(); - requests[0].values = values; + // --- Standard groupby path --- + { + std::vector requests; + requests.emplace_back(); + requests[0].values = values; - requests[0].aggregations.push_back(std::move(agg)); + requests[0].aggregations.push_back(std::unique_ptr{ + dynamic_cast(agg->clone().release())}); + + if (use_sort == force_use_sort_impl::YES) { + // WAR to force cudf::groupby to use sort implementation + requests[0].aggregations.push_back( + cudf::make_nth_element_aggregation(0)); + } + + // since the default behavior of cudf::groupby(...) for an empty null_precedence vector is + // null_order::AFTER whereas for cudf::sorted_order(...) it's null_order::BEFORE + auto const precedence = null_precedence.empty() + ? std::vector(1, cudf::null_order::BEFORE) + : null_precedence; + + cudf::groupby::groupby gb_obj( + cudf::table_view({keys}), include_null_keys, keys_are_sorted, column_order, precedence); + + auto result = gb_obj.aggregate(requests, cudf::test::get_default_stream()); - if (use_sort == force_use_sort_impl::YES) { - // WAR to force cudf::groupby to use sort implementation - requests[0].aggregations.push_back( - cudf::make_nth_element_aggregation(0)); + if (use_sort == force_use_sort_impl::YES && keys_are_sorted == cudf::sorted::NO) { + CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_expect_keys, result.first->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(sorted_expect_vals->get_column(0), + *result.second[0].results[0]); + + } else { + auto const sort_order = cudf::sorted_order(result.first->view(), column_order, precedence); + auto const sorted_keys = cudf::gather(result.first->view(), *sort_order); + auto const sorted_vals = + cudf::gather(cudf::table_view({result.second[0].results[0]->view()}), *sort_order); + + CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_expect_keys, *sorted_keys); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(sorted_expect_vals->get_column(0), + sorted_vals->get_column(0)); + } } - // since the default behavior of cudf::groupby(...) for an empty null_precedence vector is - // null_order::AFTER whereas for cudf::sorted_order(...) it's null_order::BEFORE - auto const precedence = null_precedence.empty() - ? std::vector(1, cudf::null_order::BEFORE) - : null_precedence; + // --- Streaming groupby path (single-batch, validates against same expected output) --- + // Skip streaming for: empty input (finalize throws), dictionary values (unsupported by + // streaming's row operators), ARGMIN/ARGMAX (batch-local row indices), or pre-sorted + // keys (streaming doesn't support sorted mode). + auto const skip_streaming = + keys.size() == 0 || expect_keys.size() == 0 || + values.type().id() == cudf::type_id::DICTIONARY32 || agg->kind == cudf::aggregation::ARGMIN || + agg->kind == cudf::aggregation::ARGMAX || keys_are_sorted == cudf::sorted::YES; - cudf::groupby::groupby gb_obj( - cudf::table_view({keys}), include_null_keys, keys_are_sorted, column_order, precedence); + if (use_streaming == test_streaming::YES && !skip_streaming) { + SCOPED_TRACE("streaming groupby path"); - auto result = gb_obj.aggregate(requests, cudf::test::get_default_stream()); + cudf::table_view data{{keys, values}}; + std::vector key_indices{0}; - if (use_sort == force_use_sort_impl::YES && keys_are_sorted == cudf::sorted::NO) { - CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_expect_keys, result.first->view()); - CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(sorted_expect_vals->get_column(0), - *result.second[0].results[0]); + cudf::groupby::streaming_aggregation_request sreq; + sreq.column_index = 1; + sreq.aggregation = std::unique_ptr{ + dynamic_cast(agg->clone().release())}; + + std::vector sreqs; + sreqs.push_back(std::move(sreq)); + + auto const max_distinct_keys = std::max(keys.size(), cudf::size_type{64}); + cudf::groupby::streaming_groupby sgb(key_indices, sreqs, max_distinct_keys, include_null_keys); + sgb.aggregate(data, cudf::test::get_default_stream()); + auto [skeys, sresults] = sgb.finalize(cudf::test::get_default_stream()); + + auto const precedence = null_precedence.empty() + ? std::vector(1, cudf::null_order::BEFORE) + : null_precedence; - } else { - auto const sort_order = cudf::sorted_order(result.first->view(), column_order, precedence); - auto const sorted_keys = cudf::gather(result.first->view(), *sort_order); + auto const sort_order = cudf::sorted_order(skeys->view(), column_order, precedence); + auto const sorted_keys = cudf::gather(skeys->view(), *sort_order); auto const sorted_vals = - cudf::gather(cudf::table_view({result.second[0].results[0]->view()}), *sort_order); + cudf::gather(cudf::table_view({sresults[0].results[0]->view()}), *sort_order); CUDF_TEST_EXPECT_TABLES_EQUAL(*sorted_expect_keys, *sorted_keys); CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(sorted_expect_vals->get_column(0), @@ -105,6 +152,7 @@ void test_sum_agg(cudf::column_view const& keys, {}, {}, cudf::sorted::NO, + test_streaming::NO, location); }; do_test(force_use_sort_impl::YES); diff --git a/cpp/tests/groupby/groupby_test_util.hpp b/cpp/tests/groupby/groupby_test_util.hpp index f581db75ab77..4d56342cc7f6 100644 --- a/cpp/tests/groupby/groupby_test_util.hpp +++ b/cpp/tests/groupby/groupby_test_util.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -11,6 +11,7 @@ #include enum class force_use_sort_impl : bool { NO, YES }; +enum class test_streaming : bool { NO, YES }; void test_single_agg(cudf::column_view const& keys, cudf::column_view const& values, @@ -23,6 +24,7 @@ void test_single_agg(cudf::column_view const& keys, std::vector const& column_order = {}, std::vector const& null_precedence = {}, cudf::sorted reference_keys_are_sorted = cudf::sorted::NO, + test_streaming use_streaming = test_streaming::NO, std::source_location const& location = std::source_location::current()); void test_sum_agg(cudf::column_view const& keys, cudf::column_view const& values, diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp new file mode 100644 index 000000000000..9711cce4add0 --- /dev/null +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -0,0 +1,1326 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +static std::vector const KEY_COL{0}; +static cudf::size_type constexpr DEFAULT_MAX_DISTINCT_KEYS = 1024; + +namespace { + +void sort_and_compare(std::unique_ptr& lhs_keys, + std::vector& lhs_results, + std::unique_ptr& rhs_keys, + std::vector& rhs_results, + std::vector const& null_prec = {}) +{ + auto const lhs_order = cudf::sorted_order(lhs_keys->view(), {}, null_prec); + auto const rhs_order = cudf::sorted_order(rhs_keys->view(), {}, null_prec); + + EXPECT_EQ(lhs_keys->num_rows(), rhs_keys->num_rows()); + + ASSERT_EQ(lhs_results.size(), rhs_results.size()); + for (size_t r = 0; r < lhs_results.size(); ++r) { + ASSERT_EQ(lhs_results[r].results.size(), rhs_results[r].results.size()); + for (size_t c = 0; c < lhs_results[r].results.size(); ++c) { + auto const lhs_sorted = + cudf::gather(cudf::table_view{{lhs_results[r].results[c]->view()}}, *lhs_order); + auto const rhs_sorted = + cudf::gather(cudf::table_view{{rhs_results[r].results[c]->view()}}, *rhs_order); + + auto const& lhs_col = lhs_sorted->get_column(0); + auto const& rhs_col = rhs_sorted->get_column(0); + + if (lhs_col.type() == rhs_col.type()) { + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(lhs_col, rhs_col); + } else if (cudf::is_fixed_width(lhs_col.type()) && cudf::is_fixed_width(rhs_col.type())) { + auto const lhs_cast = cudf::cast(lhs_col, rhs_col.type()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(*lhs_cast, rhs_col); + } else { + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(lhs_col, rhs_col); + } + } + } +} + +void verify_against_groupby( + std::unique_ptr& streaming_keys, + std::vector& streaming_results, + std::vector const& batches, + std::vector const& key_indices, + std::vector const& requests, + cudf::null_policy null_handling = cudf::null_policy::EXCLUDE) +{ + auto const all_data = cudf::concatenate(batches); + + std::vector key_cols; + for (auto idx : key_indices) { + key_cols.push_back(all_data->view().column(idx)); + } + cudf::groupby::groupby reference_groupby{cudf::table_view{key_cols}, null_handling}; + + std::vector ref_requests; + for (auto const& req : requests) { + cudf::groupby::aggregation_request ref_req; + ref_req.values = all_data->view().column(req.column_index); + ref_req.aggregations.push_back(std::unique_ptr{ + dynamic_cast(req.aggregation->clone().release())}); + ref_requests.push_back(std::move(ref_req)); + } + + auto [ref_keys, ref_results] = reference_groupby.aggregate(ref_requests); + + sort_and_compare(streaming_keys, streaming_results, ref_keys, ref_results); +} + +void check(std::unique_ptr& keys, + std::vector& results, + cudf::table_view expect_keys, + std::vector const& expect_vals) +{ + auto const order = cudf::sorted_order(keys->view()); + auto const sorted_keys = cudf::gather(keys->view(), *order); + CUDF_TEST_EXPECT_TABLES_EQUAL(expect_keys, sorted_keys->view()); + + size_t val_idx = 0; + for (auto& agg_res : results) { + for (auto& col : agg_res.results) { + auto const sorted = cudf::gather(cudf::table_view{{col->view()}}, *order); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(expect_vals[val_idx], sorted->get_column(0)); + ++val_idx; + } + } +} + +cudf::groupby::streaming_aggregation_request make_req( + cudf::size_type col_idx, std::unique_ptr&& agg) +{ + cudf::groupby::streaming_aggregation_request req; + req.column_index = col_idx; + req.aggregation = std::move(agg); + return req; +} + +std::vector single_agg_req( + cudf::size_type col_idx, std::unique_ptr&& agg) +{ + std::vector reqs; + reqs.push_back(make_req(col_idx, std::move(agg))); + return reqs; +} + +} // namespace + +struct StreamingGroupbyTest : public cudf::test::BaseFixture {}; + +TEST_F(StreamingGroupbyTest, SumTwoBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 3, 1}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30, 40}; + + cudf::test::fixed_width_column_wrapper keys2{2, 3, 1, 4}; + cudf::test::fixed_width_column_wrapper vals2{5, 15, 25, 35}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MinMaxTwoBatches) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{5.0, 2.0, 8.0}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2, 3}; + cudf::test::fixed_width_column_wrapper vals2{3.0, 9.0, 1.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_min_aggregation())); + reqs.push_back(make_req(1, cudf::make_max_aggregation())); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, CountValidTwoBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{{10, 20, 30, 40}, {true, false, true, true}}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{{50, 60}, {false, true}}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req( + 1, cudf::make_count_aggregation(cudf::null_policy::EXCLUDE)); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MeanTwoBatches) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{10.0, 20.0, 30.0}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{50.0, 40.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_mean_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, ProductTwoBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{3, 5}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{4, 2}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_product_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MaxMinOnIntegers) +{ + using K = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{0, 1}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{1, 1}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_max_aggregation())); + reqs.push_back(make_req(1, cudf::make_min_aggregation())); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MergeTwoObjects) +{ + using K = int32_t; + using V = int32_t; + using R = int64_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30}; + + cudf::test::fixed_width_column_wrapper keys2{2, 3}; + cudf::test::fixed_width_column_wrapper vals2{40, 50}; + + auto reqs1 = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby worker1(KEY_COL, reqs1, DEFAULT_MAX_DISTINCT_KEYS); + worker1.aggregate(cudf::table_view{{keys1, vals1}}); + + auto reqs2 = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby worker2(KEY_COL, reqs2, DEFAULT_MAX_DISTINCT_KEYS); + worker2.aggregate(cudf::table_view{{keys2, vals2}}); + + worker1.merge(worker2); + auto [keys, results] = worker1.finalize(); + + cudf::test::fixed_width_column_wrapper ek{1, 2, 3}; + cudf::test::fixed_width_column_wrapper ev{40, 60, 50}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, EmptyBatch) +{ + using K = int32_t; + using V = int32_t; + using R = int64_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::fixed_width_column_wrapper keys_empty{}; + cudf::test::fixed_width_column_wrapper vals_empty{}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(cudf::table_view{{keys_empty, vals_empty}}); + streaming_agg.aggregate(cudf::table_view{{keys1, vals1}}); + auto [keys, results] = streaming_agg.finalize(); + + cudf::test::fixed_width_column_wrapper ek{1, 2}; + cudf::test::fixed_width_column_wrapper ev{10, 20}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, SingleBatch) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 3, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30, 40, 50}; + + cudf::table_view batch1{{keys1, vals1}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, NewKeysInLaterBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::fixed_width_column_wrapper keys2{3, 4}; + cudf::test::fixed_width_column_wrapper vals2{30, 40}; + + cudf::test::fixed_width_column_wrapper keys3{1, 4}; + cudf::test::fixed_width_column_wrapper vals3{50, 60}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + cudf::table_view batch3{{keys3, vals3}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + streaming_agg.aggregate(batch3); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2, batch3}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MultipleRequestsOnDifferentColumns) +{ + using K = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper col_a1{10, 20, 30}; + cudf::test::fixed_width_column_wrapper col_b1{1.0, 2.0, 3.0}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1}; + cudf::test::fixed_width_column_wrapper col_a2{40, 50}; + cudf::test::fixed_width_column_wrapper col_b2{4.0, 5.0}; + + cudf::table_view batch1{{keys1, col_a1, col_b1}}; + cudf::table_view batch2{{keys2, col_a2, col_b2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(2, cudf::make_min_aggregation())); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, FinalizeDoesNotModifyState) +{ + using K = int32_t; + using V = int32_t; + using R = int64_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::fixed_width_column_wrapper keys2{1}; + cudf::test::fixed_width_column_wrapper vals2{30}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(cudf::table_view{{keys1, vals1}}); + + { + auto [k1, r1] = streaming_agg.finalize(); + } + + streaming_agg.aggregate(cudf::table_view{{keys2, vals2}}); + auto [keys, results] = streaming_agg.finalize(); + + cudf::test::fixed_width_column_wrapper ek{1, 2}; + cudf::test::fixed_width_column_wrapper ev{40, 20}; + check(keys, results, cudf::table_view{{ek}}, {ev}); +} + +TEST_F(StreamingGroupbyTest, NullKeysExcluded) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{{1, 2, 3}, {true, false, true}}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30}; + + cudf::test::fixed_width_column_wrapper keys2{{1, 2}, {true, false}}; + cudf::test::fixed_width_column_wrapper vals2{40, 50}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg( + KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS, cudf::null_policy::EXCLUDE); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby( + keys, results, {batch1, batch2}, KEY_COL, reqs, cudf::null_policy::EXCLUDE); +} + +TEST_F(StreamingGroupbyTest, NullKeysIncluded) +{ + using K = int32_t; + using V = int32_t; + using R = int64_t; + + cudf::test::fixed_width_column_wrapper keys1{{1, 2, 3}, {true, false, true}}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30}; + + cudf::test::fixed_width_column_wrapper keys2{{1, 2}, {true, false}}; + cudf::test::fixed_width_column_wrapper vals2{40, 50}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg( + KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS, cudf::null_policy::INCLUDE); + streaming_agg.aggregate(cudf::table_view{{keys1, vals1}}); + streaming_agg.aggregate(cudf::table_view{{keys2, vals2}}); + auto [keys, results] = streaming_agg.finalize(); + + EXPECT_EQ(keys->num_rows(), 3); + + auto const order = cudf::sorted_order(keys->view(), {}, {cudf::null_order::AFTER}); + auto const sorted_keys = cudf::gather(keys->view(), *order); + auto const sorted_vals = cudf::gather(cudf::table_view{{results[0].results[0]->view()}}, *order); + + cudf::test::fixed_width_column_wrapper ek{{1, 3, 2}, {true, true, false}}; + cudf::test::fixed_width_column_wrapper ev{50, 30, 70}; + CUDF_TEST_EXPECT_TABLES_EQUAL(cudf::table_view{{ek}}, sorted_keys->view()); + CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(ev, sorted_vals->get_column(0)); +} + +TEST_F(StreamingGroupbyTest, AllNullKeysExcluded) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{{1, 2}, {false, false}}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg( + KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS, cudf::null_policy::EXCLUDE); + streaming_agg.aggregate(cudf::table_view{{keys1, vals1}}); + auto [keys, results] = streaming_agg.finalize(); + + EXPECT_EQ(keys->num_rows(), 0); + EXPECT_EQ(results[0].results[0]->size(), 0); +} + +template +struct StreamingGroupbySumTypedTest : public cudf::test::BaseFixture {}; + +using SumSupportedTypes = + cudf::test::Concat, + cudf::test::DurationTypes>; + +TYPED_TEST_SUITE(StreamingGroupbySumTypedTest, SumSupportedTypes); + +TYPED_TEST(StreamingGroupbySumTypedTest, TwoBatches) +{ + using K = int32_t; + using V = TypeParam; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 3, 1, 2, 2, 1, 3, 3, 2}; + cudf::test::fixed_width_column_wrapper vals1{0, 1, 2, 3, 4, 5, 6, 7, 8, 9}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2, 3}; + cudf::test::fixed_width_column_wrapper vals2{10, 20, 30}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +template +struct StreamingGroupbyMinTypedTest : public cudf::test::BaseFixture {}; + +using MinSupportedTypes = + cudf::test::Concat, + cudf::test::DurationTypes>; + +TYPED_TEST_SUITE(StreamingGroupbyMinTypedTest, MinSupportedTypes); + +TYPED_TEST(StreamingGroupbyMinTypedTest, TwoBatches) +{ + using K = int32_t; + using V = TypeParam; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{5, 2, 8}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{3, 9}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_min_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, VarianceBasic) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 3, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{0, 1, 2, 3, 4}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1, 3, 3, 2}; + cudf::test::fixed_width_column_wrapper vals2{5, 6, 7, 8, 9}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_variance_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StdBasic) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 3, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{0, 1, 2, 3, 4}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1, 3, 3, 2}; + cudf::test::fixed_width_column_wrapper vals2{5, 6, 7, 8, 9}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_std_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, UnsupportedAggThrows) +{ + auto reqs = single_agg_req(1, cudf::make_collect_list_aggregation()); + EXPECT_THROW(cudf::groupby::streaming_groupby(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS), + std::invalid_argument); +} + +TEST_F(StreamingGroupbyTest, BatchExceedsMaxDistinctKeysThrows) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys{1, 2, 3, 4, 5}; + cudf::test::fixed_width_column_wrapper vals{10, 20, 30, 40, 50}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 3); + EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{keys, vals}}), std::invalid_argument); +} + +TEST_F(StreamingGroupbyTest, DisjointKeysAcrossBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::fixed_width_column_wrapper keys2{3, 4}; + cudf::test::fixed_width_column_wrapper vals2{30, 40}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, AllDuplicateKeysAcrossBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{30, 40}; + + cudf::test::fixed_width_column_wrapper keys3{1, 2}; + cudf::test::fixed_width_column_wrapper vals3{50, 60}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + cudf::table_view batch3{{keys3, vals3}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + streaming_agg.aggregate(batch3); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2, batch3}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, SingleRowBatches) +{ + using K = int32_t; + using V = int32_t; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + + std::vector batches; + std::vector> key_owners; + std::vector> val_owners; + + for (int32_t i = 0; i < 10; ++i) { + auto k = std::make_unique(cudf::test::fixed_width_column_wrapper{i % 3}); + auto v = std::make_unique(cudf::test::fixed_width_column_wrapper{i * 10}); + cudf::table_view batch{{k->view(), v->view()}}; + streaming_agg.aggregate(batch); + key_owners.push_back(std::move(k)); + val_owners.push_back(std::move(v)); + } + + auto [keys, results] = streaming_agg.finalize(); + EXPECT_EQ(keys->num_rows(), 3); +} + +// Regression test for staging-offset corruption: when a batch has internal duplicate keys, +// the last distinct key's canonical position in the staging buffer equals the distinct-key +// count, not the end of the written range. Without the fix, the next batch writes at the +// wrong offset and corrupts the canonical entry for that key. +TEST_F(StreamingGroupbyTest, InternalDuplicatesDoNotCorruptStaging) +{ + using K = int32_t; + using V = int32_t; + + // Batch 1: key 1 is a duplicate at position 1, so key 3 ends up at position 3 + // (_num_unique_keys == 3). Without the fix, batch 2 writes at offset 3, overwriting + // the canonical slot for key 3 with key 4, making them appear identical. + cudf::test::fixed_width_column_wrapper keys1{1, 1, 2, 3}; + cudf::test::fixed_width_column_wrapper vals1{10, 10, 20, 30}; + + cudf::test::fixed_width_column_wrapper keys2{4, 5}; + cudf::test::fixed_width_column_wrapper vals2{40, 50}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, SumAndMeanOnSameColumn) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{10.0, 20.0, 30.0}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1}; + cudf::test::fixed_width_column_wrapper vals2{40.0, 50.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_sum_aggregation())); + reqs.push_back(make_req(1, cudf::make_mean_aggregation())); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +// Test that many small batches accumulate correctly within the fixed-capacity key table. +TEST_F(StreamingGroupbyTest, ManySmallBatches) +{ + using K = int32_t; + using V = int32_t; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + // 10 batches of 4 rows = 40 total rows; set max_distinct_keys=40 to fit all rows. + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 40); + + std::vector batches; + std::vector> key_owners; + std::vector> val_owners; + + for (int32_t b = 0; b < 10; ++b) { + auto k = std::make_unique( + cudf::test::fixed_width_column_wrapper{b % 8, (b + 1) % 8, (b + 2) % 8, (b + 3) % 8}); + auto v = std::make_unique( + cudf::test::fixed_width_column_wrapper{b * 10, b * 10 + 1, b * 10 + 2, b * 10 + 3}); + cudf::table_view batch{{k->view(), v->view()}}; + streaming_agg.aggregate(batch); + batches.push_back(batch); + key_owners.push_back(std::move(k)); + val_owners.push_back(std::move(v)); + } + + auto [keys, results] = streaming_agg.finalize(); + verify_against_groupby(keys, results, batches, KEY_COL, reqs); +} + +// Test that exceeding distinct-key capacity throws. +TEST_F(StreamingGroupbyTest, ExceedsDistinctKeyCapacityThrows) +{ + using K = int32_t; + using V = int32_t; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + // max_distinct_keys=4: can hold at most 4 distinct keys. + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 4); + + // Batch with 4 distinct keys fills distinct-key capacity. + cudf::test::fixed_width_column_wrapper k1{0, 1, 2, 3}; + cudf::test::fixed_width_column_wrapper v1{10, 20, 30, 40}; + streaming_agg.aggregate(cudf::table_view{{k1, v1}}); + + // Any further batch with a new distinct key exceeds distinct-key capacity (4 + 1 > 4). + cudf::test::fixed_width_column_wrapper k2{4}; + cudf::test::fixed_width_column_wrapper v2{50}; + EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k2, v2}}), cudf::logic_error); +} + +// Test that sliced input columns with non-zero offsets work correctly. +TEST_F(StreamingGroupbyTest, SlicedInputColumns) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper full_keys{0, 1, 2, 3, 1, 2}; + cudf::test::fixed_width_column_wrapper full_vals{10, 20, 30, 40, 50, 60}; + + // Slice to get a view with offset=2: keys={2,3,1,2}, vals={30,40,50,60} + auto sliced = cudf::slice(cudf::table_view{{full_keys, full_vals}}, {2, 6}); + ASSERT_EQ(sliced[0].num_rows(), 4); + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(sliced[0]); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {sliced[0]}, KEY_COL, reqs); +} + +// Test that finalize() before any aggregate() throws. +TEST_F(StreamingGroupbyTest, FinalizeBeforeAggregateThrows) +{ + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + EXPECT_THROW(static_cast(streaming_agg.finalize()), cudf::logic_error); +} + +// Test merge with MEAN aggregation (compound: SUM + COUNT intermediates). +TEST_F(StreamingGroupbyTest, MergeMeanTwoBatches) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{10.0, 20.0, 30.0}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1, 3}; + cudf::test::fixed_width_column_wrapper vals2{40.0, 50.0, 60.0}; + + auto reqs1 = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby worker1(KEY_COL, reqs1, DEFAULT_MAX_DISTINCT_KEYS); + worker1.aggregate(cudf::table_view{{keys1, vals1}}); + + auto reqs2 = single_agg_req(1, cudf::make_mean_aggregation()); + cudf::groupby::streaming_groupby worker2(KEY_COL, reqs2, DEFAULT_MAX_DISTINCT_KEYS); + worker2.aggregate(cudf::table_view{{keys2, vals2}}); + + worker1.merge(worker2); + auto [keys, results] = worker1.finalize(); + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs1); +} + +// Test merge with COUNT_VALID aggregation (counts must be summed, not incremented). +TEST_F(StreamingGroupbyTest, MergeCountTwoBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 1}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30, 40}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1, 2}; + cudf::test::fixed_width_column_wrapper vals2{50, 60, 70}; + + auto reqs1 = single_agg_req( + 1, cudf::make_count_aggregation(cudf::null_policy::EXCLUDE)); + cudf::groupby::streaming_groupby worker1(KEY_COL, reqs1, DEFAULT_MAX_DISTINCT_KEYS); + worker1.aggregate(cudf::table_view{{keys1, vals1}}); + + auto reqs2 = single_agg_req( + 1, cudf::make_count_aggregation(cudf::null_policy::EXCLUDE)); + cudf::groupby::streaming_groupby worker2(KEY_COL, reqs2, DEFAULT_MAX_DISTINCT_KEYS); + worker2.aggregate(cudf::table_view{{keys2, vals2}}); + + worker1.merge(worker2); + auto [keys, results] = worker1.finalize(); + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs1); +} + +// Test merge with VARIANCE aggregation (compound: SUM_OF_SQUARES + SUM + COUNT intermediates). +TEST_F(StreamingGroupbyTest, MergeVarianceTwoBatches) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 3, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{0, 1, 2, 3, 4}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1, 3, 3, 2}; + cudf::test::fixed_width_column_wrapper vals2{5, 6, 7, 8, 9}; + + auto reqs1 = single_agg_req(1, cudf::make_variance_aggregation()); + cudf::groupby::streaming_groupby worker1(KEY_COL, reqs1, DEFAULT_MAX_DISTINCT_KEYS); + worker1.aggregate(cudf::table_view{{keys1, vals1}}); + + auto reqs2 = single_agg_req(1, cudf::make_variance_aggregation()); + cudf::groupby::streaming_groupby worker2(KEY_COL, reqs2, DEFAULT_MAX_DISTINCT_KEYS); + worker2.aggregate(cudf::table_view{{keys2, vals2}}); + + worker1.merge(worker2); + auto [keys, results] = worker1.finalize(); + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs1); +} + +TEST_F(StreamingGroupbyTest, SumOfSquaresBasic) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{3.0, 4.0, 5.0}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1}; + cudf::test::fixed_width_column_wrapper vals2{6.0, 7.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_of_squares_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, M2Basic) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{1.0, 2.0, 3.0, 4.0}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2}; + cudf::test::fixed_width_column_wrapper vals2{5.0, 6.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_m2_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StdWithNullValues) +{ + using K = int32_t; + using V = double; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1, 2}; + cudf::test::fixed_width_column_wrapper vals1{{1.0, 2.0, 3.0, 4.0}, {true, true, false, true}}; + + cudf::test::fixed_width_column_wrapper keys2{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals2{{5.0, 6.0, 7.0}, {true, false, true}}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_std_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +// ===== String key tests ===== + +TEST_F(StreamingGroupbyTest, StringKeySumTwoBatches) +{ + using V = int32_t; + + cudf::test::strings_column_wrapper keys1{"a", "b", "c", "a"}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30, 40}; + + cudf::test::strings_column_wrapper keys2{"b", "c", "a", "d"}; + cudf::test::fixed_width_column_wrapper vals2{5, 15, 25, 35}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StringKeyNonAsciiUtf8) +{ + using V = int32_t; + + cudf::test::strings_column_wrapper keys1{"αλφα", "βητα", "γαμμα", "αλφα"}; + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30, 40}; + + cudf::test::strings_column_wrapper keys2{"βητα", "δελτα", "αλφα", "🙂"}; + cudf::test::fixed_width_column_wrapper vals2{5, 15, 25, 35}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StringKeyMinMaxTwoBatches) +{ + cudf::test::strings_column_wrapper keys1{"cat", "dog", "cat"}; + cudf::test::fixed_width_column_wrapper vals1{5.0, 2.0, 8.0}; + + cudf::test::strings_column_wrapper keys2{"cat", "dog", "bird"}; + cudf::test::fixed_width_column_wrapper vals2{3.0, 9.0, 1.0}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + std::vector reqs; + reqs.push_back(make_req(1, cudf::make_min_aggregation())); + reqs.push_back(make_req(1, cudf::make_max_aggregation())); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StringKeyManySmallBatches) +{ + using V = int32_t; + + std::vector key_universe{"alpha", "beta", "gamma", "delta"}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + + std::vector batches; + std::vector> key_owners; + std::vector> val_owners; + + for (int32_t b = 0; b < 8; ++b) { + auto k = std::make_unique( + cudf::test::strings_column_wrapper{key_universe[b % 4], key_universe[(b + 1) % 4]}); + auto v = + std::make_unique(cudf::test::fixed_width_column_wrapper{b * 10, b * 10 + 1}); + cudf::table_view batch{{k->view(), v->view()}}; + streaming_agg.aggregate(batch); + batches.push_back(batch); + key_owners.push_back(std::move(k)); + val_owners.push_back(std::move(v)); + } + + auto [keys, results] = streaming_agg.finalize(); + verify_against_groupby(keys, results, batches, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StringKeyDisjointBatches) +{ + using V = int32_t; + + cudf::test::strings_column_wrapper keys1{"x", "y"}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::strings_column_wrapper keys2{"z", "w"}; + cudf::test::fixed_width_column_wrapper vals2{30, 40}; + + cudf::test::strings_column_wrapper keys3{"x", "w"}; + cudf::test::fixed_width_column_wrapper vals3{50, 60}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + cudf::table_view batch3{{keys3, vals3}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + streaming_agg.aggregate(batch3); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2, batch3}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StringKeyNullKeysExcluded) +{ + using V = int32_t; + + cudf::test::strings_column_wrapper keys1(std::initializer_list{"a", "b", "c"}, + std::initializer_list{true, false, true}); + cudf::test::fixed_width_column_wrapper vals1{10, 20, 30}; + + cudf::test::strings_column_wrapper keys2(std::initializer_list{"a", "b"}, + std::initializer_list{true, false}); + cudf::test::fixed_width_column_wrapper vals2{40, 50}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg( + KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS, cudf::null_policy::EXCLUDE); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby( + keys, results, {batch1, batch2}, KEY_COL, reqs, cudf::null_policy::EXCLUDE); +} + +TEST_F(StreamingGroupbyTest, StringKeyMerge) +{ + using V = int32_t; + + cudf::test::strings_column_wrapper keys1{"a", "b"}; + cudf::test::fixed_width_column_wrapper vals1{10, 20}; + + cudf::test::strings_column_wrapper keys2{"b", "c"}; + cudf::test::fixed_width_column_wrapper vals2{30, 40}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby obj1(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + obj1.aggregate(cudf::table_view{{keys1, vals1}}); + + cudf::groupby::streaming_groupby obj2(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + obj2.aggregate(cudf::table_view{{keys2, vals2}}); + + obj1.merge(obj2); + auto [keys, results] = obj1.finalize(); + + verify_against_groupby(keys, + results, + {cudf::table_view{{keys1, vals1}}, cudf::table_view{{keys2, vals2}}}, + KEY_COL, + reqs); +} + +TEST_F(StreamingGroupbyTest, CountAllTwoBatches) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper keys1{1, 2, 1}; + cudf::test::fixed_width_column_wrapper vals1{{10, 20, 30}, {true, false, true}}; + + cudf::test::fixed_width_column_wrapper keys2{2, 1}; + cudf::test::fixed_width_column_wrapper vals2{{40, 50}, {false, true}}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_count_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, MultiColumnKeys) +{ + using K = int32_t; + using V = int32_t; + + cudf::test::fixed_width_column_wrapper k1a{1, 1, 2}; + cudf::test::fixed_width_column_wrapper k1b{10, 20, 10}; + cudf::test::fixed_width_column_wrapper v1{100, 200, 300}; + + cudf::test::fixed_width_column_wrapper k2a{1, 2}; + cudf::test::fixed_width_column_wrapper k2b{10, 10}; + cudf::test::fixed_width_column_wrapper v2{400, 500}; + + cudf::table_view batch1{{k1a, k1b, v1}}; + cudf::table_view batch2{{k2a, k2b, v2}}; + + std::vector key_cols{0, 1}; + auto reqs = single_agg_req(2, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(key_cols, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, key_cols, reqs); +} + +// A single batch larger than `max_distinct_keys` cannot be encoded because transient +// batch values (`max_distinct_keys + row_idx`) would collide with stored dense IDs. +TEST_F(StreamingGroupbyTest, BatchExceedingMaxDistinctKeysThrows) +{ + using K = int32_t; + using V = int32_t; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 3); + + cudf::test::fixed_width_column_wrapper k{0, 1, 2, 3}; + cudf::test::fixed_width_column_wrapper v{10, 20, 30, 40}; + EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k, v}}), std::invalid_argument); +} + +// Cumulative input rows are not bounded by `max_distinct_keys` — only cumulative distinct +// keys are. Re-feeding the same batch many times keeps distinct_keys constant and +// must never throw, regardless of how many cumulative rows have been processed. +TEST_F(StreamingGroupbyTest, CumulativeRowsCanExceedMaxDistinctKeys) +{ + using K = int32_t; + using V = int32_t; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, 3); + + cudf::test::fixed_width_column_wrapper k{0, 1, 2}; + cudf::test::fixed_width_column_wrapper v{10, 20, 30}; + cudf::table_view batch{{k, v}}; + + // Five repeats: 15 cumulative rows >> max_distinct_keys=3, distinct_keys stays at 3. + streaming_agg.aggregate(batch); + streaming_agg.aggregate(batch); + streaming_agg.aggregate(batch); + streaming_agg.aggregate(batch); + streaming_agg.aggregate(batch); + + EXPECT_EQ(streaming_agg.distinct_keys(), 3); + + auto [keys, results] = streaming_agg.finalize(); + verify_against_groupby(keys, results, {batch, batch, batch, batch, batch}, KEY_COL, reqs); +} + +TEST_F(StreamingGroupbyTest, StructKeySumTwoBatches) +{ + using V = int32_t; + + // Struct key: {int, int} + cudf::test::fixed_width_column_wrapper s1a{1, 1, 2}; + cudf::test::fixed_width_column_wrapper s1b{10, 20, 10}; + auto keys1 = cudf::test::structs_column_wrapper{{s1a, s1b}}; + cudf::test::fixed_width_column_wrapper vals1{100, 200, 300}; + + cudf::test::fixed_width_column_wrapper s2a{1, 2}; + cudf::test::fixed_width_column_wrapper s2b{10, 10}; + auto keys2 = cudf::test::structs_column_wrapper{{s2a, s2b}}; + cudf::test::fixed_width_column_wrapper vals2{400, 500}; + + cudf::table_view batch1{{keys1, vals1}}; + cudf::table_view batch2{{keys2, vals2}}; + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + streaming_agg.aggregate(batch1); + streaming_agg.aggregate(batch2); + auto [keys, results] = streaming_agg.finalize(); + + verify_against_groupby(keys, results, {batch1, batch2}, KEY_COL, reqs); +} diff --git a/cpp/tests/groupby/sum_tests.cpp b/cpp/tests/groupby/sum_tests.cpp index 0d6b8422a7c6..b6b16669a866 100644 --- a/cpp/tests/groupby/sum_tests.cpp +++ b/cpp/tests/groupby/sum_tests.cpp @@ -18,6 +18,43 @@ using namespace cudf::test::iterators; +namespace { +// Run SUM aggregation through hash, sort, AND streaming groupby paths. +void test_sum_all_paths(cudf::column_view const& keys, + cudf::column_view const& values, + cudf::column_view const& expect_keys, + cudf::column_view const& expect_vals, + std::source_location const& loc = std::source_location::current()) +{ + test_single_agg(keys, + values, + expect_keys, + expect_vals, + cudf::make_sum_aggregation(), + force_use_sort_impl::NO, + cudf::null_policy::EXCLUDE, + cudf::sorted::NO, + {}, + {}, + cudf::sorted::NO, + test_streaming::YES, + loc); + test_single_agg(keys, + values, + expect_keys, + expect_vals, + cudf::make_sum_aggregation(), + force_use_sort_impl::YES, + cudf::null_policy::EXCLUDE, + cudf::sorted::NO, + {}, + {}, + cudf::sorted::NO, + test_streaming::NO, + loc); +} +} // namespace + template struct groupby_sum_test : public cudf::test::BaseFixture {}; @@ -39,11 +76,7 @@ TYPED_TEST(groupby_sum_test, basic) cudf::test::fixed_width_column_wrapper expect_keys{1, 2, 3}; cudf::test::fixed_width_column_wrapper expect_vals{9, 19, 17}; - auto agg = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg)); - - auto agg2 = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg2), force_use_sort_impl::YES); + test_sum_all_paths(keys, vals, expect_keys, expect_vals); } TYPED_TEST(groupby_sum_test, empty_cols) @@ -57,11 +90,7 @@ TYPED_TEST(groupby_sum_test, empty_cols) cudf::test::fixed_width_column_wrapper expect_keys{}; cudf::test::fixed_width_column_wrapper expect_vals{}; - auto agg = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg)); - - auto agg2 = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg2), force_use_sort_impl::YES); + test_sum_all_paths(keys, vals, expect_keys, expect_vals); } TYPED_TEST(groupby_sum_test, zero_valid_keys) @@ -75,11 +104,7 @@ TYPED_TEST(groupby_sum_test, zero_valid_keys) cudf::test::fixed_width_column_wrapper expect_keys{}; cudf::test::fixed_width_column_wrapper expect_vals{}; - auto agg = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg)); - - auto agg2 = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg2), force_use_sort_impl::YES); + test_sum_all_paths(keys, vals, expect_keys, expect_vals); } TYPED_TEST(groupby_sum_test, zero_valid_values) @@ -93,11 +118,7 @@ TYPED_TEST(groupby_sum_test, zero_valid_values) cudf::test::fixed_width_column_wrapper expect_keys{1}; cudf::test::fixed_width_column_wrapper expect_vals({0}, cudf::test::iterators::all_nulls()); - auto agg = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg)); - - auto agg2 = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg2), force_use_sort_impl::YES); + test_sum_all_paths(keys, vals, expect_keys, expect_vals); } TYPED_TEST(groupby_sum_test, null_keys_and_values) @@ -117,13 +138,11 @@ TYPED_TEST(groupby_sum_test, null_keys_and_values) // { 3, 6, 1, 4, 9, 2, 8, -} cudf::test::fixed_width_column_wrapper expect_vals({9, 14, 10, 0}, {1, 1, 1, 0}); - auto agg = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg)); - - auto agg2 = cudf::make_sum_aggregation(); - test_single_agg(keys, vals, expect_keys, expect_vals, std::move(agg2), force_use_sort_impl::YES); + test_sum_all_paths(keys, vals, expect_keys, expect_vals); } +// streaming_groupby does not accept dictionary-typed value columns, so this case +// runs only the stateless hash and sort paths via test_single_agg. TYPED_TEST(groupby_sum_test, dictionary) { using V = TypeParam; diff --git a/cpp/tests/io/json/json_test.cpp b/cpp/tests/io/json/json_test.cpp index 1e67cd6c9f9e..23fc1943d27d 100644 --- a/cpp/tests/io/json/json_test.cpp +++ b/cpp/tests/io/json/json_test.cpp @@ -2062,22 +2062,47 @@ TEST_F(JsonReaderTest, JSONLinesRecovering) TEST_F(JsonReaderTest, JSONLinesRecoveringMalformedOpenBraces) { - std::string data = - // Two lines with just an open brace (malformed JSON) - "{\n" - "{"; + // Test recovery mode with various malformed JSON patterns + for (int num_lines : {2, 4, 8, 16}) { + // Test "{\n" pattern + { + std::string data; + for (int i = 0; i < num_lines - 1; ++i) { + data += "{\n"; + } + data += "{"; - cudf::io::json_reader_options in_options = - cudf::io::json_reader_options::builder( - cudf::io::source_info{cudf::host_span{ - reinterpret_cast(data.data()), data.size()}}) - .lines(true) - .recovery_mode(cudf::io::json_recovery_mode_t::RECOVER_WITH_NULL); + cudf::io::json_reader_options in_options = + cudf::io::json_reader_options::builder( + cudf::io::source_info{cudf::host_span{ + reinterpret_cast(data.data()), data.size()}}) + .lines(true) + .recovery_mode(cudf::io::json_recovery_mode_t::RECOVER_WITH_NULL); - cudf::io::table_with_metadata result = cudf::io::read_json(in_options); + cudf::io::table_with_metadata result = cudf::io::read_json(in_options); + EXPECT_EQ(result.tbl->num_rows(), 0) << "Failed {\\n pattern with " << num_lines << " lines"; + } + + // Test {"\n pattern + { + std::string data; + for (int i = 0; i < num_lines - 1; ++i) { + data += "{\"\n"; + } + data += "{\""; - // All rows are invalid with no schema to infer, so we expect 0 rows - EXPECT_EQ(result.tbl->num_rows(), 0); + cudf::io::json_reader_options in_options = + cudf::io::json_reader_options::builder( + cudf::io::source_info{cudf::host_span{ + reinterpret_cast(data.data()), data.size()}}) + .lines(true) + .recovery_mode(cudf::io::json_recovery_mode_t::RECOVER_WITH_NULL); + + cudf::io::table_with_metadata result = cudf::io::read_json(in_options); + EXPECT_EQ(result.tbl->num_rows(), 0) + << "Failed {\"\\n pattern with " << num_lines << " lines"; + } + } } TEST_F(JsonReaderTest, JSONLinesRecoveringIgnoreExcessChars) diff --git a/dependencies.yaml b/dependencies.yaml index b9ef783e44a1..0780cd61eef8 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -161,6 +161,7 @@ files: - depends_on_pylibcudf - depends_on_libcudf - depends_on_cudf_polars + - depends_on_ray - docs - py_version py_build_cudf: diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index 24afd929b948..0b4fd3e809e6 100644 --- a/docs/cudf/source/conf.py +++ b/docs/cudf/source/conf.py @@ -625,12 +625,30 @@ def on_missing_reference(app, env, node, contnode): ("py:class", "Axis"), ("py:class", "ArrowLike"), ("py:class", "ExecutorType"), + # cudf-polars: bare rapidsmpf type names appear in autodoc'd signatures + # because they are imported under ``if TYPE_CHECKING:`` and rendered as + # unqualified strings in type annotations. The ``rapidsmpf.*`` regex below + # only matches fully qualified targets, so the bare leaf names are listed + # explicitly here. + ("py:class", "Statistics"), + ("py:class", "Communicator"), + ("py:class", "Options"), + # polars aliases that don't match the public intersphinx targets. + ("py:class", "pl.DataFrame"), + ("py:class", "polars.LazyFrame"), + ("py:class", "polars.DataFrame"), + ("py:class", "polars.dataframe.frame.DataFrame"), ] # Temporarily disable nitpick warnings for pandas: https://github.com/pandas-dev/pandas/issues/64584 nitpick_ignore_regex = [ ("py:.*", "pandas.*"), ("py:.*", "pd.*"), ("ref.*", ".*pandas.*"), + # External libs without configured intersphinx inventories. + ("py:.*", r"rapidsmpf(\..*)?"), + ("py:.*", r"ray(\..*)?"), + ("py:.*", r"distributed(\..*)?"), + ("py:.*", r"dask_cuda(\..*)?"), ] diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index 823954a3b087..6acff73d4f03 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -1,17 +1,73 @@ (cudf-polars-api)= -# API +# API Reference -For the most part, the public API of `cudf-polars` is the polars API. +For the most part, the public API of `cudf-polars` is the Polars API itself. This page +documents the additional classes and functions that `cudf-polars` exposes for the streaming +multi-GPU engines. + +## Streaming engines + +```{eval-rst} +.. autoclass:: cudf_polars.engine.ray.RayEngine + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :show-inheritance: + +.. autoclass:: cudf_polars.engine.dask.DaskEngine + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :show-inheritance: + +.. autoclass:: cudf_polars.engine.spmd.SPMDEngine + :members: from_options, gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks, rank, comm, context + :show-inheritance: + +.. autoclass:: cudf_polars.engine.default_singleton_engine.DefaultSingletonEngine + :members: get_or_create, shutdown + :show-inheritance: +``` + +The engine classes share a common base class: + +```{eval-rst} +.. autoclass:: cudf_polars.engine.core.StreamingEngine + :members: gather_cluster_info, gather_statistics, global_statistics, shutdown, nranks + :show-inheritance: + +.. autoclass:: cudf_polars.engine.core.ClusterInfo + :members: +``` + +## Configuration + +```{eval-rst} +.. autoclass:: cudf_polars.engine.options.StreamingOptions + :members: from_dict, to_dict, to_rapidsmpf_options, to_executor_options, to_engine_options + +.. autodata:: cudf_polars.engine.options.UNSPECIFIED + +.. autoclass:: cudf_polars.engine.hardware_binding.HardwareBindingPolicy + :members: + +.. autofunction:: cudf_polars.engine.hardware_binding.bind_to_gpu +``` + +## SPMD helpers + +```{eval-rst} +.. autofunction:: cudf_polars.engine.spmd.allgather_polars_dataframe + +.. autofunction:: cudf_polars.streaming.actor_graph.collectives.common.reserve_op_id +``` + +## Internal configuration objects + +These dataclasses back the `engine_options` surfaced by `pl.GPUEngine` and `StreamingOptions`. +Most users interact with them through `StreamingOptions` fields rather than directly. ```{eval-rst} .. automodule:: cudf_polars.utils.config :members: - Cluster, - ConfigOptions, - CUDAStreamPoolConfig, DynamicPlanningOptions, - ExecutorType, - InMemoryExecutor, + MemoryResourceConfig, ParquetOptions, StreamingExecutor, StreamingFallbackMode, diff --git a/docs/cudf/source/cudf_polars/dask_engine.md b/docs/cudf/source/cudf_polars/dask_engine.md new file mode 100644 index 000000000000..190968cc07a7 --- /dev/null +++ b/docs/cudf/source/cudf_polars/dask_engine.md @@ -0,0 +1,198 @@ +(cudf-polars-dask-engine)= +# Dask + +{class}`~cudf_polars.engine.dask.DaskEngine` runs the streaming executor +on a [Dask distributed][dask-distributed] cluster: one Dask worker per GPU, coordinated by a +single client process. Partitions are streamed through the query plan and collective operations +(shuffles, allgathers, joins) run across workers over a shared UCXX communicator. On startup, +each worker is pinned to the CPU cores and NUMA node closest to its GPU (see +[Pre-configured GPU clusters](#pre-configured-gpu-clusters) below). + +```python +import polars as pl +from cudf_polars.engine.dask import DaskEngine + +with DaskEngine() as engine: + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + print(result) +``` + +With no arguments, {class}`~cudf_polars.engine.dask.DaskEngine` creates a +`distributed.LocalCluster` with one worker per visible GPU, a `distributed.Client`, and +bootstraps a UCXX communicator across all workers. On exit, everything it created is torn down. + +```{note} +`.collect()` pulls the full result back to the client process. For large distributed outputs, +prefer `.sink_*()` or aggregate/sample inside the query before `.collect()`. See +[Result collection](engines.md#result-collection). +``` + +## Configuring `DaskEngine` + +For custom configuration, build +{class}`~cudf_polars.engine.options.StreamingOptions` and use +`DaskEngine.from_options()`: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.dask import DaskEngine + +opts = StreamingOptions(num_streaming_threads=8, fallback_mode="silent") + +with DaskEngine.from_options(opts) as engine: + result = pl.scan_parquet("/data/dataset/*.parquet").collect(engine=engine) +``` + +See {doc}`options` for the available fields. + +## Bring your own Dask client + +Pass an existing `distributed.Client` via `dask_client=` to attach to an already-running +scheduler: + +```python +from distributed import Client +import polars as pl +from cudf_polars.engine.dask import DaskEngine + +with Client("scheduler-address:8786") as dc: + with DaskEngine(dask_client=dc) as engine: + result = pl.scan_parquet("/data/*.parquet").collect(engine=engine) +``` + +When you supply the client, {class}`~cudf_polars.engine.dask.DaskEngine` +leaves it (and the cluster) alone on exit. + +(pre-configured-gpu-clusters)= +### Pre-configured GPU clusters + +Some Dask launchers, notably `dask_cuda.LocalCUDACluster`, already pin CPU affinity and set +`CUDA_VISIBLE_DEVICES` per worker. Disable the built-in hardware binding via +{class}`~cudf_polars.engine.hardware_binding.HardwareBindingPolicy` +to avoid having both layers fight over each worker's affinity (the second to run wins, which +makes the resulting placement non-deterministic): + +```python +from dask_cuda import LocalCUDACluster +from distributed import Client +from cudf_polars.engine.dask import DaskEngine +from cudf_polars.engine.hardware_binding import ( + HardwareBindingPolicy, +) + +with Client(LocalCUDACluster()) as dc, DaskEngine( + dask_client=dc, + engine_options={ + "hardware_binding": HardwareBindingPolicy(enabled=False), + }, +) as engine: + ... +``` + +### Manually launched workers + +When launching workers yourself (for example on a multi-node HPC cluster), use the built-in nanny +preload to assign one GPU per worker. The preload sets `CUDA_VISIBLE_DEVICES` on each worker +before the process spawns: + +```bash +# On each node, launch one worker per GPU with a single thread each: +dask worker SCHEDULER_ADDRESS:8786 --nworkers N --nthreads 1 \ + --preload-nanny cudf_polars.engine.dask +``` + +Then connect from the client: + +```python +import polars as pl +from distributed import Client +from cudf_polars.engine.dask import DaskEngine + +with Client("SCHEDULER_ADDRESS:8786") as dc: + with DaskEngine(dask_client=dc) as engine: + result = pl.scan_parquet("/data/*.parquet").collect(engine=engine) +``` + +Hardware binding (CPU affinity, NUMA, network) is handled automatically by +{class}`~cudf_polars.engine.dask.DaskEngine`; the nanny preload only +deals with GPU assignment. + +See the [Dask CLI deployment guide][dask-cli] for more on `dask worker` options. + +#### Using `dask-cuda-worker` + +As an alternative to the built-in nanny preload, you can launch workers with +[`dask-cuda-worker`][dask-cuda-worker] from the [dask-cuda][dask-cuda] project. It launches one +worker per visible GPU and installs a set of plugins on every worker: a `CPUAffinity` plugin +that pins the worker to the NUMA node of its GPU, an `RMMSetup` plugin, and a nanny preload that +configures UCX. + +`DaskEngine` sets up the same things for its own streaming runtime, so the two need to be +coordinated or they will fight: + +* **CPU affinity is unconditional in `dask-cuda-worker`**, the `CPUAffinity` plugin is always + installed and there is no CLI flag to turn it off. Pass `hardware_binding=HardwareBindingPolicy(enabled=False)` + to `DaskEngine` so it does not try to re-pin affinity on top of dask-cuda's binding. +* **Do not pass `--rmm-pool-size`, `--rmm-managed-memory`, or similar RMM flags** to + `dask-cuda-worker`. Let `DaskEngine` own the memory resource via its `memory_resource_config` + (see {doc}`options`) otherwise two different memory resources will be installed on the same + worker. +* **Do not pass `--enable-tcp-over-ucx`, `--enable-infiniband`, `--enable-nvlink`, or + `--enable-rdmacm`** to `dask-cuda-worker`. `DaskEngine` bootstraps its own UCXX communicator + and will select transports itself. Enabling them on both sides can produce inconsistent UCX + configuration across the cluster. + +```bash +# On each node, GPU assignment + CPU affinity only (no RMM, no UCX flags): +dask-cuda-worker SCHEDULER_ADDRESS:8786 +``` + +```python +import polars as pl +from distributed import Client +from cudf_polars.engine.dask import DaskEngine +from cudf_polars.engine.hardware_binding import ( + HardwareBindingPolicy, +) + +with Client("SCHEDULER_ADDRESS:8786") as dc: + with DaskEngine( + dask_client=dc, + engine_options={ + # dask-cuda-worker always pins CPU affinity; disable DaskEngine's + # binding so the two don't conflict. + "hardware_binding": HardwareBindingPolicy(enabled=False), + }, + ) as engine: + result = pl.scan_parquet("/data/*.parquet").collect(engine=engine) +``` + +## Cluster diagnostics + +{meth}`~cudf_polars.engine.dask.DaskEngine.gather_cluster_info` returns +placement information for every worker: + +```python +with DaskEngine() as engine: + print(f"cluster has {engine.nranks} workers") + for info in engine.gather_cluster_info(): + print( + f"hostname={info['hostname']}, pid={info['pid']}, " + f"CUDA_VISIBLE_DEVICES={info['cuda_visible_devices']}" + ) +``` + +{class}`~cudf_polars.engine.dask.DaskEngine` raises `RuntimeError` if +created inside an `rrun` cluster. + +[dask-distributed]: https://distributed.dask.org/ +[dask-cli]: https://docs.dask.org/en/latest/deploying-cli.html +[dask-cuda]: https://docs.rapids.ai/api/dask-cuda/nightly/ +[dask-cuda-worker]: https://docs.rapids.ai/api/dask-cuda/nightly/quickstart/#dask-cuda-worker diff --git a/docs/cudf/source/cudf_polars/default_singleton_engine.md b/docs/cudf/source/cudf_polars/default_singleton_engine.md new file mode 100644 index 000000000000..05d3017d7575 --- /dev/null +++ b/docs/cudf/source/cudf_polars/default_singleton_engine.md @@ -0,0 +1,109 @@ +(cudf-polars-default-singleton-engine)= +# Default `engine="gpu"` + +`.collect(engine="gpu")` (and `engine=pl.GPUEngine()`) is the API you invoke when you don't +construct a streaming engine explicitly. It runs the same streaming executor as the explicit +engines (Ray, Dask, SPMD), conceptually similar to +[Polars' own streaming engine](https://docs.pola.rs/user-guide/concepts/streaming/) but on the +GPU. Under the hood it's backed by {class}`~cudf_polars.engine.default_singleton_engine.DefaultSingletonEngine`, +a process-wide singleton specialization of {class}`~cudf_polars.engine.spmd.SPMDEngine`. At most one live +instance exists per process, which is created lazily on first use and torn down at interpreter +exit. Ray is the showcased explicit engine (see {doc}`usage`); this page documents what +`engine="gpu"` does *without* you having to construct anything. + +```{important} +`engine="gpu"` is meant for trivial setup: single-GPU execution with no +configuration or engine object to manage. +For any non-trivial workflow, construct an engine explicitly. To tune +options, use +{meth}`RayEngine.from_options(...) `. +`engine="gpu"` accepts no options, so settings such as +`spill_to_pinned_memory=True` for spill-heavy workloads require an +explicit engine. See {doc}`usage` and {doc}`options`. +``` + +## What you get without an explicit engine + +When you just write: + +```python +import polars as pl + +result = ( + pl.scan_parquet("/data/*.parquet") + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine="gpu") +) +``` + +cudf-polars uses +{class}`~cudf_polars.engine.default_singleton_engine.DefaultSingletonEngine` +under the hood. No cluster is set up, the rapidsmpf `Context` is bootstrapped on first use, +and subsequent `.collect()` calls in the same process reuse it. + +## Explicit handle + +If you genuinely want the singleton (for example in tests or scripts that need to call +`.shutdown()` deterministically) you can obtain it via the factory: + +```python +from cudf_polars.engine.default_singleton_engine import ( + DefaultSingletonEngine, +) + +engine = DefaultSingletonEngine.get_or_create() +result = query.collect(engine=engine) +``` + +`get_or_create()` is idempotent: calling it again returns the same instance. + +For anything beyond defaults, prefer an explicit engine. See {doc}`usage`. + +## Lifecycle + +The singleton is bootstrapped once per process. The rapidsmpf `Context`, RMM adaptor, and +Python thread-pool executor are reused across every `.collect()` call. + +Shutdown is automatic: the engine registers an `atexit` hook that tears it down at interpreter +exit. To shut it down explicitly (for example to release resources before constructing a +multi-GPU engine), call the static method: + +```python +from cudf_polars.engine.default_singleton_engine import ( + DefaultSingletonEngine, +) + +DefaultSingletonEngine.shutdown() +``` + +`shutdown()` is idempotent (calling it twice is safe) and a no-op if no live engine exists. + +## Mutual exclusion with explicit engines + +`DefaultSingletonEngine`, {class}`~cudf_polars.engine.ray.RayEngine`, +{class}`~cudf_polars.engine.dask.DaskEngine`, and +{class}`~cudf_polars.engine.spmd.SPMDEngine` cannot coexist in the same +process. Concretely: + +- Constructing `RayEngine` / `DaskEngine` / `SPMDEngine` while the singleton is alive raises + `RuntimeError`. +- `DefaultSingletonEngine.get_or_create()` raises `RuntimeError` if any explicit streaming + engine is alive. + +Recommended pattern: pick one engine for the lifetime of the program. If you need to switch, +shut down the active engine first: + +```python +DefaultSingletonEngine.shutdown() +explicit_engine = SPMDEngine.from_options(opts) +``` + +## No options + +`DefaultSingletonEngine.get_or_create()` takes no arguments. To tune `StreamingOptions` such +as `spill_to_pinned_memory`, `fallback_mode`, `max_rows_per_partition`, or any rapidsmpf +runtime knob, construct an explicit +{class}`~cudf_polars.engine.ray.RayEngine` via +{meth}`RayEngine.from_options(...) `. +See {doc}`options` for the available fields. diff --git a/docs/cudf/source/cudf_polars/engine_options.md b/docs/cudf/source/cudf_polars/engine_options.md deleted file mode 100644 index ba6085275b87..000000000000 --- a/docs/cudf/source/cudf_polars/engine_options.md +++ /dev/null @@ -1,176 +0,0 @@ -# GPUEngine Configuration Options - -The `polars.GPUEngine` object may be configured in several different ways. - -## Executor - -`cudf-polars` includes multiple *executors*, backends that take a Polars query and execute it to produce the result (either an in-memory `polars.DataFrame` from `.collect()` or one or more files with `.sink_`). These can be specified with the `executor` option when you create the `GPUEngine`. - -```python -import polars as pl - -engine = pl.GPUEngine(executor="streaming") -query = ... - -result = query.collect(engine=engine) -``` - -The `streaming` executor is the default executor as of RAPIDS 25.08, and is -equivalent to passing `engine="gpu"` or `engine=pl.GPUEngine()` to `collect`. At -a high-level, the `streaming` executor works by breaking inputs (in-memory -DataFrames or parquet files) into multiple pieces and streaming those pieces -through the series of operations needed to produce the final result. - -We also provide an `in-memory` executor. This executor is often faster when the -underlying data fits comfortably in device memory, because the overhead of splitting -inputs and executing them in batches is less beneficial at this scale. With that said, -this executor must rely on [Unified Virtual Memory] (UVM) if the input and intermediate -data do not fit in device memory. The `in-memory` executor can be used with - -```python -engine = pl.GPUEngine(executor="in-memory") -``` - -In general, we recommend starting with the default `streaming` executor, because -it scales significantly better than `in-memory`. The `streaming` executor includes -several configuration options, which can be provided with the `executor_options` -key when constructing the `GPUEngine`: - -```python -engine = pl.GPUEngine( - executor="streaming", # the default - executor_options={ - "max_rows_per_partition": 500_000, - } -) -``` - -You can configure the default value for configuration options through -environment variables with the prefix `CUDF_POLARS__EXECUTOR__{option_name}`. -For example, the environment variable -`CUDF_POLARS__EXECUTOR__MAX_ROWS_PER_PARTITION` will set the default -`max_rows_per_partition` to use if it isn't overridden through -`executor_options`. - -For boolean options, like `sink_to_directory`, the values `{"1", "true", "yes", "y"}` -are considered `True` and `{"0", "false", "no", "n"}` are considered `False`. - -See [Configuration Reference](#cudf-polars-api) for a full list of options, and -[Streaming Execution](#cudf-polars-streaming) for more on the streaming executor, -including multi-GPU execution. - -## Parquet Reader Options - -Reading large parquet files can use a large amount of memory, especially when the files are compressed. This may lead to out of memory errors for some workflows. To mitigate this, the "chunked" parquet reader may be selected. When enabled, parquet files are read in chunks, limiting the peak memory usage at the cost of a small drop in performance. - -To configure the parquet reader, we provide a dictionary of options to the `parquet_options` keyword of the `GPUEngine` object. Valid keys and values are: -- `chunked` indicates that chunked parquet reading is to be used. By default, chunked reading is turned on. -- [`chunk_read_limit`](https://docs.rapids.ai/api/libcudf/legacy/classcudf_1_1io_1_1chunked__parquet__reader#aad118178b7536b7966e3325ae1143a1a) controls the maximum size per chunk. By default, the maximum chunk size is unlimited. -- [`pass_read_limit`](https://docs.rapids.ai/api/libcudf/legacy/classcudf_1_1io_1_1chunked__parquet__reader#aad118178b7536b7966e3325ae1143a1a) controls the maximum memory used for decompression. The default pass read limit is 16GiB. - -For example, to select the chunked reader with custom values for `pass_read_limit` and `chunk_read_limit`: -```python -engine = GPUEngine( - parquet_options={ - 'chunked': True, - 'chunk_read_limit': int(1e9), - 'pass_read_limit': int(4e9) - } -) -result = query.collect(engine=engine) -``` -Note that passing `chunked: False` disables chunked reading entirely, and thus `chunk_read_limit` and `pass_read_limit` will have no effect. - -You can configure the default value for configuration options through -environment variables with the prefix -`CUDF_POLARS__PARQUET_OPTIONS__{option_name}`. For example, the environment -variable `CUDF_POLARS__PARQUET_OPTIONS__CHUNKED=0` will set the default -`chunked` to `False`. - -## CUDA Stream Policy - -By default, all CUDA operations in `cudf-polars` are launched on the default -stream. You can configure `cudf-polars` to use multiple CUDA streams, which may -improve performance by overlapping data transfers and kernel launches. - -This behavior is configured by the `cuda_stream_policy` keyword or -`CUDF_POLARS__CUDA_STREAM_POLICY` environment variable. The valid options are - -* `default`: use the default CUDA stream for all kernel launches and memory operations -* `new`: create a new CUDA stream when necessary (e.g. when reading from a file or loading an in-memory `polars.LazyFrame` object, - or when performing a join where the inputs might be on different streams) -* `pool`: use an RMM stream pool (only supported with the rapidsmpf runtime) - -The ``rapidsmpf`` runtime for the streaming executor also supports using a CUDA Stream Pool. - -```python -engine = GPUEngine( - executor="streaming", - executor_options={ - "runtime": "rapidsmpf", - }, - cuda_stream_policy="pool", -) -``` - -Or provide a dictionary with configuration options for the pool, like `cuda_stream_pool={"pool_size": 16}`. -You can also set the `CUDF_POLARS__CUDA_STREAM_POLICY` environment variable the JSON encoded configuration dictionary. - -This stream pool is shared between cudf-polars and rapidsmpf. - -## Memory Resource - -All GPU memory allocations made by cudf-polars use an RMM Memory Resource object from {mod}`rmm.mr`. You can specify -the memory resource to use by: - -1. Passing a concrete `MemoryResource` instance to {class}`~polars.lazyframe.engine_config.GPUEngine`. -2. Passing the configuration options for a Memory Resource as the `memory_resource_config` keyword argument to {class}`~polars.lazyframe.engine_config.GPUEngine`. -3. Relying on the default behavior, which creates a memory resource for you (details below). - -By default, cudf-polars will create a new RMM Memory Resource for you, which is cached and reused -for each query. The type of that memory resource is hardware-dependent. GPUs that support [Unified Virtual Memory] memory, -use a {class}`rmm.mr.ManagedMemoryResource` wrapped in a {class}`rmm.mr.PoolMemoryResource` and {class}`rmm.mr.PrefetchResourceAdaptor`. -Otherwise, {class}`rmm.mr.CudaAsyncMemoryResource` is used. - -Set `POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY=0` to disabled managed memory and use {class}`rmm.mr.CudaAsyncMemoryResource` instead. - -Alternatively, you can customize the pool by passing the configuration for an RMM Memory Resource object as `memory_resource_config` -when creating your {class}`~polars.lazyframe.engine_config.GPUEngine`: - -```python -memory_resource_config = { - "qualname": "rmm.mr.CudaAsyncMemoryResource", - "options": { - "initial_pool_size": "100 MiB", - } -} - -engine = pl.GPUEngine(memory_resource_config=memory_resource_config) -``` - -This lets you control things like the initial pool size or release threshold. - -Finally, for maximum flexibility, you can create your own memory resource object and pass it into the {class}`~polars.lazyframe.engine_config.GPUEngine`: - -```python -import polars as pl -import rmm - -mr = rmm.mr.CudaAsyncMemoryResource() -engine = pl.GPUEngine(memory_resource=mr) -``` - -Passing a concrete memory resource takes precedence over passing the `memory_resource_config` options, -which takes precedence over the default memory resource. - -Note that providing a concrete memory resource isn't an option with the distributed scheduler, -because the concrete memory resource is only valid for the process in which it was created. - -## Disabling CUDA Managed Memory - -By default the `in-memory` executor will use [CUDA managed memory](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#unified-memory-introduction) with RMM's pool allocator. On systems that don't support managed memory, a non-managed asynchronous pool -allocator is used. -Managed memory can be turned off by setting `POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY` to `0`. System requirements for managed memory can be found [here]( -https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#system-requirements-for-unified-memory). - -[Unified Virtual Memory]: https://developer.nvidia.com/blog/unified-memory-cuda-beginners/ diff --git a/docs/cudf/source/cudf_polars/engines.md b/docs/cudf/source/cudf_polars/engines.md new file mode 100644 index 000000000000..3bb83bfcbea4 --- /dev/null +++ b/docs/cudf/source/cudf_polars/engines.md @@ -0,0 +1,94 @@ +(cudf-polars-engines)= +# Engines + +## What is an engine? + +`cudf-polars` executes Polars `LazyFrame` queries on GPU. You select GPU execution by passing an +`engine=` argument to `.collect()` or `.sink_*()`. The `engine` you pass decides *how* the +query runs: whether it streams through partitioned inputs or fits everything in device memory, +whether it runs in-process or distributes work across a cluster of GPU workers, and which +cluster backend coordinates those workers. + +## Execution modes + +### Streaming + +Streaming engines partition their inputs (Parquet files or in-memory `DataFrame`s) and process +those partitions through the query graph in chunks. This lets queries scale past device memory +and (on Ray, Dask, and SPMD) across multiple GPUs and multiple nodes. cudf-polars' streaming +executor is its own GPU implementation, but conceptually parallels +[Polars' CPU streaming engine](https://docs.pola.rs/user-guide/concepts/streaming/): the same +partition-and-stream model, just on the GPU. + +All four ways of running cudf-polars use this same streaming executor: +{class}`~cudf_polars.engine.ray.RayEngine`, +{class}`~cudf_polars.engine.dask.DaskEngine`, +{class}`~cudf_polars.engine.spmd.SPMDEngine`, and the default +`engine="gpu"` (backed internally by +{class}`~cudf_polars.engine.default_singleton_engine.DefaultSingletonEngine`). +They differ only in how their GPU worker(s) are provisioned. +{class}`~cudf_polars.engine.ray.RayEngine` with no arguments uses every +GPU visible to the process, so on a single node with N GPUs it runs the query on all N of them +without any extra configuration. Launching a multi-node cluster simply means pointing the +engine at that cluster; the user-facing code is the same. + +### In-memory + +The in-memory engine (`engine=pl.GPUEngine(executor="in-memory")`) is +the only non-streaming path. It runs the query on a single GPU, materializing intermediates in +device memory. Use it for small queries (data that fits in device memory), debugging, or when +you specifically need `LazyFrame.profile` support (see {doc}`profiling`). For production +workloads on any nontrivial dataset, use a streaming engine. See {doc}`in_memory_engine` for +details. + +## Cluster backends + +| Engine | Cluster model | Extra runtime dependency | Typical use | +| --------------------------------------------- | --------------------------------------------------------| ------------------------ | ------------------------------------------------------------------------------- | +| {class}`~cudf_polars.engine.ray.RayEngine` | Single client, one Ray actor per GPU | [Ray][ray-docs] | Works from a laptop to a cloud cluster. No separate cluster setup needed. | +| {class}`~cudf_polars.engine.dask.DaskEngine` | Single client, one Dask worker per GPU | [Dask distributed][dask] | Teams with an existing Dask deployment or a preferred Dask launcher. | +| {class}`~cudf_polars.engine.spmd.SPMDEngine` | Same script runs once per GPU, joined by a communicator | UCXX (under `rrun`) | HPC / SPMD launchers such as `rrun`. Single-rank mode needs no cluster at all. | +| [`engine="gpu"`](default_singleton_engine.md) | Implicit process-wide singleton on one GPU; no cluster | None | Default when no engine is constructed. Short scripts and notebooks. No options. | + +All four approaches use the same execution model under the hood, so which to select depends +on your preferred deployment method, not performance tradeoffs. For any non-trivial workflow, +construct one of the first three engines explicitly (see {doc}`usage`). `engine="gpu"` is a +convenience and accepts no options, so it cannot be tuned. See {doc}`default_singleton_engine` +for details on the implementation that backs it. + +## Result collection + +`.collect()` returns a single `pl.DataFrame` on the **caller's process**. On the streaming +engines that has two flavors: + +- **`RayEngine` / `DaskEngine`** (single client): every partition is pulled from the + cluster workers back to the client and concatenated there. This is convenient for small + results but does not scale to large queries. E.g., calling `.collect()` on a 1 TB query + result sends 1 TB through your client. Sink the result + (`.sink_parquet("path/")`, `.sink_csv(...)`, …) so each rank writes its own partition + directly, or reduce/sample the data inside the query before `.collect()`. +- **`SPMDEngine`** (one process per GPU): each rank's `.collect()` returns *that rank's* + local fragment. There is no client to gather to. If you need a single concatenated + `pl.DataFrame` across ranks, call + {func}`~cudf_polars.engine.spmd.allgather_polars_dataframe` explicitly (see + [Collecting distributed results](spmd_engine.md#collecting-distributed-results)). If you + want to keep processing the data rank-by-rank, just stay in `SPMDEngine` and use its + MPI-style model: each rank already owns its fragment. +- **`engine="gpu"`**: single GPU, no cluster, so `.collect()` is the only sensible option. + +Rules of thumb for multi-machine `RayEngine` / `DaskEngine` runs: + +- For exports: prefer `.sink_*()` over `.collect()`. +- For analysis: aggregate, sample, or `limit()` the result inside the lazy query before + `.collect()` so the client only sees a small DataFrame. +- For further distributed processing in Python: switch to `SPMDEngine` so each rank keeps + its fragment. + +## Where to go next + +- {doc}`usage`: tutorial that walks through running your first GPU query end-to-end. +- {doc}`other_engines`: per-engine reference pages for DaskEngine and SPMDEngine. +- {doc}`options`: the `StreamingOptions` configuration object and every field it surfaces. + +[ray-docs]: https://docs.ray.io/ +[dask]: https://distributed.dask.org/ diff --git a/docs/cudf/source/cudf_polars/in_memory_engine.md b/docs/cudf/source/cudf_polars/in_memory_engine.md new file mode 100644 index 000000000000..8ee82af3ae41 --- /dev/null +++ b/docs/cudf/source/cudf_polars/in_memory_engine.md @@ -0,0 +1,37 @@ +(cudf-polars-in-memory-engine)= +# In-memory engine + +The in-memory engine (`engine=pl.GPUEngine(executor="in-memory")`) is +the only non-streaming path in cudf-polars. It materializes the whole query in device memory +on a single GPU. + +For most workflows, prefer a streaming engine. Use the in-memory engine when: + +- The data comfortably fits in device memory and you want minimum setup. +- You need `LazyFrame.profile` (see {doc}`profiling`). +- You are debugging and want the simpler, non-streaming execution path. + +```python +result = query.collect(engine=pl.GPUEngine(executor="in-memory")) +``` + +This is the path documented in Polars' own [GPU support guide][polars-gpu]. By contrast, +`engine="gpu"` (or `engine=pl.GPUEngine()`) selects the default streaming path on a single GPU +(see {doc}`default_singleton_engine`). That default accepts no options, so for anything beyond +a quick script, construct an explicit engine. + +## Configuration + +The in-memory engine does not accept +{class}`~cudf_polars.engine.options.StreamingOptions`. Pass keyword +arguments to `pl.GPUEngine(...)` directly: + +```python +import polars as pl + +engine = pl.GPUEngine(executor="in-memory", parquet_options={"chunked": True}) +``` + +See the [Polars GPU support guide][polars-gpu] for the full in-memory usage story. + +[polars-gpu]: https://docs.pola.rs/user-guide/gpu-support/ diff --git a/docs/cudf/source/cudf_polars/index.md b/docs/cudf/source/cudf_polars/index.md new file mode 100644 index 000000000000..f98d3e2ecb44 --- /dev/null +++ b/docs/cudf/source/cudf_polars/index.md @@ -0,0 +1,104 @@ +# Polars GPU engine + +cuDF provides GPU-accelerated execution engines for Python users of the Polars Lazy API. The +engines support most of the core expressions and data types as well as a growing set of more +advanced dataframe manipulations and data file formats. When a GPU engine is selected, Polars +converts expressions into an optimized query plan and determines whether the plan is supported +on the GPU. If it is not, the execution transparently falls back to the standard Polars engine +and runs on the CPU. + +## Install + +Follow the [RAPIDS installation guide](https://docs.rapids.ai/install) and pick the +`cudf-polars` package for your CUDA and Python versions. For example, with conda: + +```bash +conda install -c rapidsai -c conda-forge -c nvidia cudf-polars +``` + +Or with pip (CUDA 13 wheels; use `cudf-polars-cu12` for CUDA 12): + +```bash +pip install cudf-polars-cu13 +``` + +## Quick start + +{class}`~cudf_polars.engine.ray.RayEngine` with no arguments uses +every GPU visible to the process, so the same code runs on one GPU and scales to multi-GPU / +multi-node setups automatically: + +```python +import polars as pl +from cudf_polars.engine.ray import RayEngine + +query = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) +) + +with RayEngine() as engine: + result = query.collect(engine=engine) +``` + +See {doc}`usage` for the full tutorial, {doc}`engines` for a conceptual overview of the +available engines, and {doc}`options` for the +{class}`~cudf_polars.engine.options.StreamingOptions` configuration. + +## Benchmark + +```{note} +The following benchmarks were performed with the `POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY` +environment variable set to `"0"`. Using managed memory (the default) imposes a performance cost +in order to avoid out of memory errors. Peak performance can still be attained by setting the +environment variable to `0`. +``` + +We reproduced the [Polars Decision Support (PDS)](https://github.com/pola-rs/polars-benchmark) +benchmark to compare Polars GPU engine with the default CPU settings across several dataset sizes. +Here are the results: + +```{figure} ../_static/pds_benchmark_polars.png +:width: 600px +``` + +You can see up to 13x speedup using the GPU engine on the compute-heavy PDS queries involving +complex aggregation and join operations. Below are the speedups for the top performing queries: + +```{figure} ../_static/compute_heavy_queries_polars.png +:width: 1000px +``` + +*PDS-H benchmark | GPU: NVIDIA H100 PCIe | CPU: Intel Xeon W9-3495X (Sapphire Rapids) | Storage: +Local NVMe* + +You can reproduce the results by visiting the [Polars Decision Support (PDS) GitHub repository](https://github.com/pola-rs/polars-benchmark). + +## Learn More + +The GPU engine for Polars is now available in Open Beta and the engine is undergoing rapid development. +To learn more, visit the [GPU Support page](https://docs.pola.rs/user-guide/gpu-support/) on the Polars website. + +```{toctree} +:maxdepth: 1 +:caption: Contents: + +usage +engines +options +profiling +other_engines +api +``` + +## Launch on Google Colab + +```{figure} ../_static/colab.png +:width: 200px +:target: https://nvda.ws/4eKlWZW + +Try out the GPU engine for Polars in a free GPU notebook environment. +Sign in with your Google account and [launch the demo on Colab](https://nvda.ws/4eKlWZW). +``` diff --git a/docs/cudf/source/cudf_polars/index.rst b/docs/cudf/source/cudf_polars/index.rst deleted file mode 100644 index 40759a72fa66..000000000000 --- a/docs/cudf/source/cudf_polars/index.rst +++ /dev/null @@ -1,56 +0,0 @@ -Polars GPU engine -================= - -cuDF provides an in-memory, GPU-accelerated execution engine for Python users of the Polars Lazy API. -The engine supports most of the core expressions and data types as well as a growing set of more advanced dataframe manipulations -and data file formats. When using the GPU engine, Polars will convert expressions into an optimized query plan and determine -whether the plan is supported on the GPU. If it is not, the execution will transparently fall back to the standard Polars engine -and run on the CPU. This functionality is available in Open Beta, is undergoing rapid development, and is currently a single GPU implementation. - -Benchmark ---------- - -.. note:: - The following benchmarks were performed with the ``POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY`` environment variable set to ``"0"``. - Using managed memory (the default) imposes a performance cost in order to avoid out of memory errors. - Peak performance can still be attained by setting the environment variable to ``0``. - -We reproduced the `Polars Decision Support (PDS) `__ benchmark to compare Polars GPU engine with the default CPU settings across several dataset sizes. Here are the results: - -.. figure:: ../_static/pds_benchmark_polars.png - :width: 600px - - - -You can see up to 13x speedup using the GPU engine on the compute-heavy PDS queries involving complex aggregation and join operations. Below are the speedups for the top performing queries: - - -.. figure:: ../_static/compute_heavy_queries_polars.png - :width: 1000px - -:emphasis:`PDS-H benchmark | GPU: NVIDIA H100 PCIe | CPU: Intel Xeon W9-3495X (Sapphire Rapids) | Storage: Local NVMe` - -You can reproduce the results by visiting the `Polars Decision Support (PDS) GitHub repository `__. - -Learn More ----------- - -The GPU engine for Polars is now available in Open Beta and the engine is undergoing rapid development. To learn more, visit the `GPU Support page `__ on the Polars website. - -.. toctree:: - :maxdepth: 1 - :caption: Contents: - - usage - streaming_execution - engine_options - api - -Launch on Google Colab ----------------------- - -.. figure:: ../_static/colab.png - :width: 200px - :target: https://nvda.ws/4eKlWZW - - Try out the GPU engine for Polars in a free GPU notebook environment. Sign in with your Google account and `launch the demo on Colab `__. diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md new file mode 100644 index 000000000000..20ff7ce1b4ab --- /dev/null +++ b/docs/cudf/source/cudf_polars/options.md @@ -0,0 +1,130 @@ +(cudf-polars-options)= +# Configuration Options + +{class}`~cudf_polars.engine.options.StreamingOptions` is the recommended +way to configure the streaming engines (Ray, Dask, SPMD. The default `engine="gpu"` accepts no +options, see the note below). Build one and pass it to `RayEngine.from_options()` +to construct a {class}`~cudf_polars.engine.ray.RayEngine`: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine + +opts = StreamingOptions( + num_streaming_threads=8, + fallback_mode="silent", + spill_device_limit="70%", +) + +with RayEngine.from_options(opts) as engine: + result = ( + pl.scan_parquet("/data/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) +``` + +```{note} +`engine="gpu"` (the default when no engine is constructed) accepts no +{class}`~cudf_polars.engine.options.StreamingOptions`. Many of the +fields below have a noticeable runtime impact (for example `spill_to_pinned_memory=True` +significantly speeds up spill-heavy workflows), so to use any non-default value construct one +of the engines listed below. +``` + +{class}`~cudf_polars.engine.options.StreamingOptions` covers three +categories of fields: + +| Category | Scope | Env var prefix | +| ----------- | -------------------------------------------------------------------------------------- | ------------------------- | +| `executor` | Query execution and partitioning, e.g. `max_rows_per_partition`, `fallback_mode`, ... | `CUDF_POLARS__EXECUTOR__` | +| `engine` | `pl.GPUEngine` kwargs, e.g. Parquet, memory resource, CUDA streams, hardware binding | `CUDF_POLARS__` | +| `rapidsmpf` | Streaming runtime, e.g. threads, CUDA streams, spilling, pinned memory, log level | `RAPIDSMPF_` | + +The `engine` category surfaces the same tuning knobs as plain `pl.GPUEngine(...)`. For example, +`parquet_options` and `memory_resource_config`. Configure these settings through +{class}`~cudf_polars.engine.options.StreamingOptions` rather than +passing them to `pl.GPUEngine(...)` directly. + +The `rapidsmpf` category adds lower-level configuration for the streaming runtime that has no equivalent on +the plain `pl.GPUEngine`. Most users will not need to touch these directly. See the +[streaming runtime configuration reference][rapidsmpf-config] for the underlying meaning of each +`RAPIDSMPF_*` field. + +Every option has a corresponding environment variable. When an option is not set explicitly, its +value is read from the environment variable if present; otherwise the underlying library applies +its built-in default. Boolean environment variables accept `{"1", "true", "yes", "y"}` as true +and `{"0", "false", "no", "n"}` as false. + + +## Building from a dictionary + +{meth}`~cudf_polars.engine.options.StreamingOptions.from_dict` accepts a flat dict of field names. +Unknown keys raise `TypeError` and `None` values leave the field unspecified: + +```python +opts = StreamingOptions.from_dict({ + "num_streaming_threads": 8, + "fallback_mode": "silent", +}) +``` + +This is convenient when options come from a config file or CLI. + +## Engine keyword arguments + +Each engine ({class}`~cudf_polars.engine.ray.RayEngine`, +{class}`~cudf_polars.engine.dask.DaskEngine`, or +{class}`~cudf_polars.engine.spmd.SPMDEngine`) accepts +`rapidsmpf_options`, `executor_options`, and `engine_options` as raw keyword arguments. +We recommend using this only when you need fine-grained control that doesn't fit the +{class}`~cudf_polars.engine.options.StreamingOptions` schema. +Otherwise, prefer the engine's `from_options` constructor with +{class}`~cudf_polars.engine.options.StreamingOptions`. + +For the in-memory engine, +{class}`~cudf_polars.engine.options.StreamingOptions` does not apply. +See {doc}`in_memory_engine` for how to configure it. + + +## Options Reference + +Environment variables follow these patterns: + +* `executor`: `CUDF_POLARS__EXECUTOR__` (e.g. `CUDF_POLARS__EXECUTOR__FALLBACK_MODE`) +* `engine`: `CUDF_POLARS__` (e.g. `CUDF_POLARS__RAISE_ON_FAIL`; nested prefixes for structured options) +* `rapidsmpf`: `RAPIDSMPF_` (e.g. `RAPIDSMPF_NUM_STREAMING_THREADS`) + +### Category: `executor` + +| Field | Description | Default | +|--------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------|-------------| +| `num_py_executors` | Workers for the internal Python `ThreadPoolExecutor`. | `8` | +| `fallback_mode` | When an unsupported operation forces a fallback to CPU execution: `"warn"`, `"raise"`, `"silent"`. | `"warn"` | +| `max_rows_per_partition` | Maximum number of rows per partition. Only used for in-memory `DataFrame` sources, never for disk IO or dynamic planning. | `1_000_000` | +| `broadcast_limit` | Maximum number of bytes for broadcast joins. | auto | +| `target_partition_size` | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto. | auto | +| `dynamic_planning` | Dynamic planning configuration, dict or {class}`~cudf_polars.utils.config.DynamicPlanningOptions`. `None` disables. | enabled | +| `sink_to_directory` | Whether `.sink_*()` writes its output as a directory. The `spmd`, `ray`, and `dask` engines always use `True`; passing `False` raises `ValueError`. | `True` | + +### Category: `engine` + +| Field | Description | Default | +|--------------------------|-------------------------------------------------------------------------------------------------------------------------------|---------------------------| +| `raise_on_fail` | Raise an error instead of falling back to CPU execution. | `False` | +| `parquet_options` | Parquet configuration, dict or {class}`~cudf_polars.utils.config.ParquetOptions`. | — | +| `memory_resource_config` | RMM configuration, dict or {class}`~cudf_polars.utils.config.MemoryResourceConfig`. | — | +| `cuda_stream_policy` | CUDA stream policy (`"default"`, `"pool"`, or a configuration dict). | — | +| `hardware_binding` | Hardware binding policy. Pass a {class}`~cudf_polars.engine.hardware_binding.HardwareBindingPolicy` for fine-grained control. | `HardwareBindingPolicy()` | +| `allow_gpu_sharing` | When `False` (default), the engine raises if multiple ranks share the same physical GPU. | `False` | + +### Category: `rapidsmpf` + +Lower-level streaming runtime knobs. Most users will not need to touch these directly. See the +[streaming runtime configuration reference][rapidsmpf-config] for the full list of fields and defaults. + + +[rapidsmpf-config]: https://docs.rapids.ai/api/rapidsmpf/nightly/configuration/ diff --git a/docs/cudf/source/cudf_polars/other_engines.md b/docs/cudf/source/cudf_polars/other_engines.md new file mode 100644 index 000000000000..d09b031aa189 --- /dev/null +++ b/docs/cudf/source/cudf_polars/other_engines.md @@ -0,0 +1,34 @@ +(cudf-polars-other-engines)= +# Other Engines + +The examples in {doc}`usage` use +{class}`~cudf_polars.engine.ray.RayEngine`. The pages below cover +other ways to run cudf-polars: + +* **{doc}`dask_engine`** runs on a [Dask distributed][dask] cluster with one Dask worker per + GPU. Use this when you already have a Dask deployment or a preferred Dask launcher. +* **{doc}`spmd_engine`** is single program, multiple data: the same script runs once per GPU, + typically launched with `rrun`. Single-rank mode needs no external cluster at all. +* **{doc}`default_singleton_engine`** documents what `engine="gpu"` does under the hood when no + engine is constructed explicitly. Useful to *understand*, but for any non-trivial workflow we + recommend constructing an explicit engine so you can pass {class}`~cudf_polars.engine.options.StreamingOptions`. +* **{doc}`in_memory_engine`** (`engine=pl.GPUEngine(executor="in-memory")`) is the only non-streaming + path. Suitable for small queries (data that fits in device memory), debugging, or when you specifically + need `LazyFrame.profile`. + +See {doc}`engines` for the conceptual comparison with `RayEngine` (cluster model, runtime +dependencies, typical use), and {doc}`options` for the shared +{class}`~cudf_polars.engine.options.StreamingOptions` configuration +(the in-memory engine does not accept `StreamingOptions`). + +```{toctree} +:maxdepth: 1 +:hidden: + +dask_engine +spmd_engine +default_singleton_engine +in_memory_engine +``` + +[dask]: https://distributed.dask.org/ diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md new file mode 100644 index 000000000000..ed19c8dd178e --- /dev/null +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -0,0 +1,190 @@ +(cudf-polars-profiling)= +# Profiling and Tracing + +## Streaming Statistics + +When a query runs on a streaming engine +({class}`~cudf_polars.engine.ray.RayEngine`, +{class}`~cudf_polars.engine.dask.DaskEngine`, +{class}`~cudf_polars.engine.spmd.SPMDEngine`, or the default +`engine="gpu"`), the underlying streaming runtime can record detailed per-rank statistics: +shuffle byte counts, allgather participation, memory-pool high-water marks, and more. See the +[underlying statistics reference][rapidsmpf-stats] for the full list of metrics. + +Statistics collection is off by default. Enable it by setting `statistics=True` on +{class}`~cudf_polars.engine.options.StreamingOptions` (or exporting +`RAPIDSMPF_STATISTICS=1`), then call `gather_statistics()` on the engine to pull the per-rank +records: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine + +opts = StreamingOptions(statistics=True) + +with RayEngine.from_options(opts) as engine: + result = ( + pl.scan_parquet("/data/*.parquet") + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) + + per_rank = engine.gather_statistics(clear=True) + for rank, stats in enumerate(per_rank): + print(f"rank {rank}:\n{stats}") +``` + +`gather_statistics(*, clear=False)` returns a list of `rapidsmpf.statistics.Statistics` objects, +one per rank, in rank order. Passing `clear=True` resets each rank's counters after the gather — +useful when you want to scope statistics to a single query. + +Use `global_statistics(*, clear=False)` when you only need the cluster-wide picture. It gathers +and merges the per-rank statistics into a single `Statistics` (counts and values summed, maxima +reduced with `max`). Capture it inside the engine context, then print after exit: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine + +opts = StreamingOptions(statistics=True) + +with RayEngine.from_options(opts) as engine: + result = pl.scan_parquet("/data/*.parquet").collect(engine=engine) + total = engine.global_statistics(clear=True) +print(total) +``` + + +## GPU Profiling + +For streaming queries, we recommend profiling with [NVIDIA NSight Systems][nsight]. `cudf-polars` +includes [nvtx][nvtx] annotations to help you understand where time is being spent. Streaming +engines do not support `LazyFrame.profile`, since `profile` requires a single in-memory pass. + +If you specifically need [`LazyFrame.profile`](https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.profile.html), +the in-memory engine supports it. This is useful for small queries during development: + +```python +import polars as pl +q = pl.scan_parquet("ny-taxi/2024/*.parquet").filter(pl.col("total_amount") > 15.0) +profile = q.profile(engine=pl.GPUEngine(executor="in-memory")) +``` + +The result is `(result_df, timings_df)`, see the Polars docs link above for the schema. + +## Tracing + +cudf-polars can optionally trace execution of each node in the query plan. To enable tracing, set +the environment variable ``CUDF_POLARS_LOG_TRACES`` to a true value ("1", "true", "y", "yes") +before starting your process. + +cudf-polars logs traces at three scopes (levels): + +1. `plan`: These generally happen once per query. This will include things like the (serialized) + query plan. +2. `actor`: (streaming engines only). There will be roughly one `actor` trace per node in the + logical plan. +3. `evaluate_ir_node`: Logs the evaluation of a physical node in the query plan. Note that one + logical node might expand to more than one physical nodes. + +Each trace includes a `scope` key indicating which level that trace belongs to. `actor`-scoped +nodes will be nested under a `plan`-scoped node. When using a streaming engine, +`evaluate_ir_node`-scoped nodes will be nested under an `actor`-scoped node. + +### Schemas + +The different scopes have different schemas. Fields in **bold** are required / always present. + +#### scope=plan + +| Field Name | Type | Description | +| ---------- | ----- | ----------- | +| **scope** | Literal["plan"] | The string literal `"plan"`. Useful for distinguishing from other types of traces. | +| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | +| **plan** | `PlanObject` | A serialized representation of the query plan. | +| **event** | String | A message like "Query Plan" | + +#### scope=actor + +`actor`-scoped traces only appear when running on a streaming engine. + +| Field Name | Type | Description | +| ---------- | ----- | ----------- | +| **scope** | Literal["actor"] | The string literal `"actor"`. Useful for distinguishing from other types of traces. | +| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | +| **start** | int | A nanosecond-resolution counter indicating when the actor started. Note: actors generally start early in the query and suspend waiting for data. | +| **stop** | int | A nanosecond-resolution counter indicating when the actor completed. | +| **event** | String | A message like "Streaming Actor". | +| **actor_ir_type** | String | The type of the actor, like `"Scan"`. | +| **actor_ir_id** | int | A unique identifier for the actor. All traces logged under this actor will include this value. | +| chunk_count | int | A counter for how many table chunks have been processed by this actor at the time of logging. | +| duplicated | bool | Whether the output rows are duplicated across ranks (e.g. after an allgather). | +| row_count | int | Total row count produced by this node during execution. | + +#### scope=evaluate_ir_node + +| Field Name | Type | Description | +| ---------- | ----- | ----------- | +| **scope** | `Literal["evaluate_ir_node"]` | The string literal `"evaluate_ir_node"`. Useful for distinguishing from other types of traces. | +| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | +| **type** | string | The name of the IR node | +| **start** | int | A nanosecond-precision counter indicating when this node started executing | +| **stop** | int | A nanosecond-precision counter indicating when this node finished executing | +| **overhead_duration** | int | The overhead, in nanoseconds, added by tracing | +| `count_frames_{phase}` | int | The number of dataframes for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_DATAFRAMES=0`. | +| `frames_{phase}` | `list[dict]` | A list with dictionaries with "shape" and "size" fields, one per input dataframe, for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_DATAFRAMES=0`. | +| `total_bytes_{phase}` | int | The sum of the size (in bytes) of the dataframes for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_current_bytes_{phase}` | int | The current number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_current_count_{phase}` | int | The current number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_peak_bytes_{phase}` | int | The peak number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_peak_count_{phase}` | int | The peak number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_total_bytes_{phase}` | int | The total number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `rmm_total_count_{phase}` | int | The total number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| `nvml_current_bytes_{phase}` | int | The device memory usage of this process, as reported by NVML, for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | +| actor_ir_id | int | A unique identifier for the parent actor (streaming engines only). | + +Setting `CUDF_POLARS_LOG_TRACES=1` enables all the metrics. Depending on the query, the overhead +from collecting the memory or dataframe metrics can be measurable. You can disable some metrics +through additional environment variables. For example, to disable the memory-related metrics, set: + +```bash +CUDF_POLARS_LOG_TRACES=1 CUDF_POLARS_LOG_TRACES_MEMORY=0 +``` + +And to disable the memory and dataframe metrics, which essentially leaves just the duration +metrics, set +```bash +CUDF_POLARS_LOG_TRACES=1 CUDF_POLARS_LOG_TRACES_MEMORY=0 CUDF_POLARS_LOG_TRACES_DATAFRAMES=0 +``` + +Note that tracing still needs to be enabled with `CUDF_POLARS_LOG_TRACES=1`. + +The implementation uses [structlog] to build log records. You can configure the output using +structlog's [configuration][structlog-configure] and enrich the records with +[context variables][structlog-context]. + +```python +>>> df = pl.DataFrame({"a": ["a", "a", "b"], "b": [1, 2, 3]}).lazy() +>>> df.group_by("a").agg(pl.col("b").min().alias("min"), pl.col("b").max().alias("max")).collect(engine=pl.GPUEngine(executor="in-memory")) +2025-09-10 07:44:01 [info ] Execute IR count_frames_input=0 count_frames_output=1 ... type=DataFrameScan +2025-09-10 07:44:01 [info ] Execute IR count_frames_input=1 count_frames_output=1 ... type=GroupBy +shape: (2, 3) +┌─────┬─────┬─────┐ +│ a ┆ min ┆ max │ +│ --- ┆ --- ┆ --- │ +│ str ┆ i64 ┆ i64 │ +╞═════╪═════╪═════╡ +│ b ┆ 3 ┆ 3 │ +│ a ┆ 1 ┆ 2 │ +└─────┴─────┴─────┘ +``` + +[nsight]: https://developer.nvidia.com/nsight-systems +[nvtx]: https://nvidia.github.io/NVTX/ +[rapidsmpf-stats]: https://docs.rapids.ai/api/rapidsmpf/nightly/statistics/ +[structlog]: https://www.structlog.org/ +[structlog-configure]: https://www.structlog.org/en/stable/configuration.html +[structlog-context]: https://www.structlog.org/en/stable/contextvars.html diff --git a/docs/cudf/source/cudf_polars/spmd_engine.md b/docs/cudf/source/cudf_polars/spmd_engine.md new file mode 100644 index 000000000000..a15e727ab53c --- /dev/null +++ b/docs/cudf/source/cudf_polars/spmd_engine.md @@ -0,0 +1,210 @@ +(cudf-polars-spmd-engine)= +# SPMD + +{class}`~cudf_polars.engine.spmd.SPMDEngine` runs the streaming executor +in [SPMD][spmd-wiki] mode: the same Python script runs once per GPU, and each process owns its +local data fragment. Collective operations (shuffles, allgathers, joins) coordinate across +processes to produce a globally consistent result. + +On startup, `SPMDEngine` pins the process to the CPU cores and NUMA node closest to its GPU. +Under `rrun` this binding is delegated to the launcher; outside `rrun` (single-process mode) +`SPMDEngine` performs it itself. See +{class}`~cudf_polars.engine.hardware_binding.HardwareBindingPolicy` +to override this behaviour. + +## Single-GPU setup + +To use {class}`~cudf_polars.engine.spmd.SPMDEngine` on a single GPU, create the engine and +run your Python script as normal. You still get the full streaming executor (partitioned inputs, +spilling, scaling past device memory), you just don't need any multi-process coordination: + +```python +# python my_script.py +import polars as pl +from cudf_polars.engine.spmd import SPMDEngine + +with SPMDEngine() as engine: + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) +``` + +With a single rank, the [Query symmetry requirement](#query-symmetry-requirement) and +[Collecting distributed results](#collecting-distributed-results) steps below do not apply, +`collect()` returns the full result directly. + +## Multi-GPU with `rrun` + +To run on more than one GPU, the same Python script must be launched collectively, and all +processes must be informed that they are participating in the cluster. This is the role of the +`rrun` launcher: it starts one process per GPU, +{class}`~cudf_polars.engine.spmd.SPMDEngine` detects this and bootstraps +a UCXX communicator across all ranks. + +When the same script is launched without `rrun`, `SPMDEngine` falls back to a single-process, +single-GPU communicator that requires no external communication library. This mode is useful +for local development, unit tests, and single-GPU pipelines (see [Single-GPU setup](#single-gpu-setup) above). + +```python +# multi-GPU launch: rrun -n 4 python my_script.py +# single-GPU: python my_script.py +import polars as pl +from cudf_polars.engine.spmd import SPMDEngine + +with SPMDEngine() as engine: + result = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) +``` + +File-based sources (`scan_parquet`, `scan_csv`, …) are automatically partitioned so that each +rank reads a different file or row-group range. In-memory `DataFrame` objects are already +rank-local, so each rank processes its own copy. + +## Configuring `SPMDEngine` + +For custom configuration, build a +{class}`~cudf_polars.engine.options.StreamingOptions` and use +`SPMDEngine.from_options()`: + +```python +import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import SPMDEngine + +opts = StreamingOptions(num_streaming_threads=8, fallback_mode="silent") + +with SPMDEngine.from_options(opts) as engine: + result = pl.scan_parquet("/data/dataset/*.parquet").collect(engine=engine) +``` + +See {doc}`options` for the available fields. + +{class}`~cudf_polars.engine.spmd.SPMDEngine` exposes a few properties +that are useful in SPMD code: + +* `engine.nranks` / `engine.rank`: cluster size and local rank index. +* `engine.comm`: the active `rapidsmpf.communicator.Communicator`. +* `engine.context`: the active `rapidsmpf.streaming.core.context.Context`. + +## Query symmetry requirement + +All ranks must execute the **same sequence of queries in the same order**. Collective operations +are matched using internal operation IDs. If one rank executes a collective that another rank +does not, the program will deadlock. + +In practice: + +* Avoid rank-conditional `collect()` or `sink*()` calls. +* Avoid branches that change the query graph. +* Keep the client script deterministic. + +```python +# OK: every rank runs the same query in the same order. +with SPMDEngine() as engine: + result = ( + pl.scan_parquet("/data/*.parquet") + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) +``` + +```python +# DEADLOCKS: rank 0 issues a group_by collective the other ranks never see. +with SPMDEngine() as engine: + df = pl.scan_parquet("/data/*.parquet") + if engine.rank == 0: # don't do this + df = df.group_by("customer_id").agg(pl.col("amount").sum()) + result = df.collect(engine=engine) +``` + +## Collecting distributed results + +Unlike `RayEngine` / `DaskEngine`, where `.collect()` gathers every partition to the client, +here each rank's `.collect()` returns *its own* fragment. If you want to keep processing the +data rank-by-rank, just use that fragment directly. If you need a single concatenated view, +use the helper below. + +`collect()` returns a rank-local result. Use +{func}`~cudf_polars.engine.spmd.allgather_polars_dataframe` to assemble +the full dataset on every rank: + +```python +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id +from cudf_polars.engine.spmd import ( + SPMDEngine, + allgather_polars_dataframe, +) + +with SPMDEngine() as engine: + result = pl.scan_parquet("/data/*.parquet").collect(engine=engine) + + with reserve_op_id() as op_id: + full = allgather_polars_dataframe( + engine=engine, + local_df=result, + op_id=op_id, + ) +``` + +`op_id` identifies the collective across ranks. All ranks must pass the same value. +{func}`~cudf_polars.streaming.actor_graph.collectives.common.reserve_op_id` draws from the same +pool that cudf-polars uses internally for shuffle and join collectives, so there is no risk of +collision. Do not pass hardcoded integers: they may silently collide with an ID reserved by an +active collective inside `collect()`. + +The result is a `pl.DataFrame` containing rows from all ranks in rank order (rank 0 first, then +rank 1, …, rank N-1). + +## Reusing a communicator + +By default {class}`~cudf_polars.engine.spmd.SPMDEngine` bootstraps a new +UCXX communicator on every construction. When running multiple engines in sequence (for example +in a test suite or interactive session), repeated bootstrapping is unnecessary and can race on +the file-based coordination layer shared by all ranks. + +Pass a pre-created communicator via `comm=` to skip the bootstrap entirely. The engine does +**not** close the communicator on shutdown. The caller retains ownership and can reuse it +across multiple {class}`~cudf_polars.engine.spmd.SPMDEngine` lifetimes: + +```python +from rapidsmpf import bootstrap +from rapidsmpf.progress_thread import ProgressThread +from cudf_polars.engine.spmd import SPMDEngine + +# Bootstrap once. +comm = bootstrap.create_ucxx_comm(progress_thread=ProgressThread()) + +# Reuse across multiple engine lifetimes, no re-bootstrap between them. +with SPMDEngine(comm=comm) as engine: + result1 = df1.lazy().collect(engine=engine) + +with SPMDEngine(comm=comm) as engine: + result2 = df2.lazy().collect(engine=engine) +``` + +## Cluster diagnostics + +{meth}`~cudf_polars.engine.spmd.SPMDEngine.gather_cluster_info` returns +placement information for every rank: + +```python +with SPMDEngine() as engine: + if engine.rank == 0: + for i, info in enumerate(engine.gather_cluster_info()): + print( + f"rank {i}: hostname={info['hostname']}, pid={info['pid']}, " + f"CUDA_VISIBLE_DEVICES={info['cuda_visible_devices']}" + ) +``` + +[spmd-wiki]: https://en.wikipedia.org/wiki/Single_program,_multiple_data diff --git a/docs/cudf/source/cudf_polars/streaming_execution.md b/docs/cudf/source/cudf_polars/streaming_execution.md deleted file mode 100644 index b599401a8749..000000000000 --- a/docs/cudf/source/cudf_polars/streaming_execution.md +++ /dev/null @@ -1,140 +0,0 @@ -(cudf-polars-streaming)= -# Streaming Execution - -The streaming executors work best when the inputs to your query come -from parquet files. That is, start with `scan_parquet`, not existing -Polars `DataFrame`s or CSV files. - -## Single GPU streaming - -The simplest case, requiring no additional dependencies, is the -`single` cluster option. An appropriate engine is: - -```python -engine = pl.GPUEngine() -``` - -This uses the default single-GPU *cluster* and is equivalent to -`pl.GPUEngine(executor="streaming", executor_options={"cluster": "single"})`, -or simply passing `engine="gpu"` to `.collect()`. - -When executed with this engine, any parquet inputs are split into -"partitions" that are streamed through the query graph. We try to -pick a good default for the typical partition size (based on the -amount of GPU memory available), however, it might not be optimal. You -can configure the execution by providing more options to the executor. -For example, to split input parquet files into 125 MB chunks: - -```python -engine = pl.GPUEngine( - executor="streaming", - executor_options={ - "target_partition_size": 125_000_000 # 125 MB - } -) -``` - -Use the executor option `max_rows_per_partition` to control how in-memory -``DataFrame`` inputs are split into multiple partitions. - -You may find, at the cost of higher memory footprint, that a larger value gives -better performance. - -````{note} -If part of a query does not run in streaming mode, but _does_ run -using the in-memory GPU engine, then we automatically concatenate the -inputs for that operation into a single partition, and effectively -fall back to the in-memory engine. - -The `fallback_mode` option can be used to raise an exception when -this fallback occurs or silence the warning instead: - - - engine = pl.GPUEngine( - executor="streaming", - executor_options={ - "fallback_mode": "raise", - } - ) -```` - -## Multi GPU streaming - -```{note} -The distributed cluster is considered experimental and might change without warning. -``` - -Streaming utilising multiple GPUs simultaneously is supported by -setting the `"cluster"` to `"distributed"`: -```python -engine = pl.GPUEngine( - executor="streaming", - executor_options={"cluster": "distributed"}, -) -``` - -Unlike the single GPU executor, this does require a number of -additional dependencies. We currently require -[Dask](https://www.dask.org/) and -[Dask-CUDA](https://docs.rapids.ai/api/dask-cuda/nightly/) to be -installed. In addition, we recommend that Dask Distributed plugin of -[UCXX](https://github.com/rapidsai/ucxx) and -[RapidsMPF](https://github.com/rapidsai/rapidsmpf) are installed to -take advantage of any high-performance networking. - -To quickly install all of these dependencies into a conda environment, -you can run: - -``` -conda install -c rapidsai -c conda-forge \ - cudf-polars rapidsmpf dask-cuda distributed-ucxx -``` - - -````{note} -Identically to the single-GPU streaming case, if part of a query does -not support execution with multiple partitions, but is supported by -the in-memory GPU engine, we concatenate the inputs and execute using -a single partition. -```` - -The multi-GPU engine uses the currently active Dask client to carry -out the partitioned execution, so for multi-GPU we would use something -like - -```python -from dask_cuda import LocalCUDACluster - -... - -client = LocalCUDACluster(...).get_client() - -q = ... -engine = pl.GPUEngine( - executor="streaming", - executor_options={"cluster": "distributed"}, -) -result = q.collect(engine=engine) -``` - -````{warning} -If you request a `"distributed"` cluster but do not have a cluster -deployed, `collect`ing the query will fail. -```` - -### Streaming sink operations - -When the `"distributed"` cluster option is active, sink operations like -`df.sink_parquet("my_path")` will always produce a directory containing -one or more files. It is not currently possible to disable this behavior. - -When the `"single"` cluster option is active, sink operations will -generate a single file by default. However, you may opt into the -distributed sink behavior by adding `{"sink_to_directory": True}` -to your `executor_options` dictionary. - -## Get Started - -The experimental streaming GPU executor is now available. For a quick -walkthrough of a multi-GPU example workflow and performance on a real dataset, -check out the [multi-GPU Polars demo](https://github.com/rapidsai-community/showcase/blob/main/accelerated_data_processing_examples/multi_gpu_polars_demo.ipynb). diff --git a/docs/cudf/source/cudf_polars/usage.md b/docs/cudf/source/cudf_polars/usage.md index 934cce9a9f0d..6bcd4220774d 100644 --- a/docs/cudf/source/cudf_polars/usage.md +++ b/docs/cudf/source/cudf_polars/usage.md @@ -1,230 +1,170 @@ +(cudf-polars-usage)= # Usage -`cudf-polars` enables GPU acceleration for Polars' LazyFrame API by executing logical plans with cuDF and pylibcudf. It requires minimal code changes and works by specifying a GPU engine during execution. +`cudf-polars` runs your Polars `LazyFrame` queries on GPU. You select GPU execution by passing +an `engine=` argument to `.collect()` or `.sink_*()`. See {doc}`engines` for the conceptual +picture, this page walks through running your first query. -For a high-level overview of GPU support in Polars, see the [Polars GPU support guide](https://docs.pola.rs/user-guide/gpu-support/). +We always recommend constructing an engine object and using it in a context manager to ensure proper +resource cleanup. The engine constructor is where you specify {class}`~cudf_polars.engine.options.StreamingOptions` +such as `spill_to_pinned_memory` or `fallback_mode`. Ray is the showcased example below, see also +{doc}`other_engines`. -## Getting Started - -Use `cudf-polars` by calling `.collect(engine="gpu")` or `.sink_(engine="gpu")` on a LazyFrame: +## Your first GPU query ```python import polars as pl +from cudf_polars.engine.ray import RayEngine + +query = ( + pl.scan_parquet("/data/dataset/*.parquet") + .filter(pl.col("amount") > 100) + .group_by("customer_id") + .agg(pl.col("amount").sum()) +) + +with RayEngine() as engine: + result = query.collect(engine=engine) +print(result) +``` + +{class}`~cudf_polars.engine.ray.RayEngine` with no arguments uses every +GPU visible to the process, so the example above runs on one GPU if that's all that's available +and scales automatically to every GPU on the node otherwise. It also attaches to an existing +Ray cluster if one is already running (see [Attaching to an existing Ray cluster](#attaching-to-an-existing-ray-cluster)). + +```{note} +The examples on this page use {class}`~cudf_polars.engine.ray.RayEngine`. `cudf-polars` supports +multiple engines for GPU execution. See {doc}`other_engines` for alternatives, or {doc}`engines` for a conceptual overview of when to pick which. +``` -q = pl.scan_parquet("ny-taxi/2024/*.parquet").filter(pl.col("total_amount") > 15.0) -result = q.collect(engine="gpu") +```{note} +`.collect()` pulls the full result back to the client process. For large distributed outputs, +prefer `.sink_*()` or aggregate/sample inside the query before `.collect()`. See +[Result collection](engines.md#result-collection). ``` -Alternatively, you can create a `GPUEngine` instance with custom configuration: +## Configuring `RayEngine` + +{class}`~cudf_polars.engine.ray.RayEngine` with no arguments starts a +local [Ray][ray-docs] cluster and creates one GPU worker per visible GPU. + +For custom configuration, build a +{class}`~cudf_polars.engine.options.StreamingOptions` and use +`RayEngine.from_options()`: ```python import polars as pl +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine -engine = pl.GPUEngine(raise_on_fail=True) +opts = StreamingOptions(num_streaming_threads=8, fallback_mode="silent") -q = pl.scan_parquet("ny-taxi/2024/*.parquet").filter(pl.col("total_amount") > 15.0) -result = q.collect(engine=engine) +with RayEngine.from_options(opts) as engine: + result = pl.scan_parquet("/data/dataset/*.parquet").collect(engine=engine) ``` -With `raise_on_fail=True`, the query will raise an exception if it cannot be run on the GPU instead of transparently falling back to polars CPU. See more [engine options](engine_options.md). +See {doc}`options` for the available fields. -## GPU Profiling +```{note} +`RayEngine` is an object you create and pass to `.collect(engine=engine)`. Prefer the +context-manager form so the Ray cluster and GPU workers are torn down automatically. +``` -The `streaming` executor does not support profiling query execution through the `LazyFrame.profile` method. With the default `synchronous` scheduler for the `streaming` executor, we recommend using [NVIDIA NSight Systems](https://developer.nvidia.com/nsight-systems) to profile your queries. -cudf-polars includes [nvtx](https://nvidia.github.io/NVTX/) annotations to help you understand where time is being spent. +The same `from_options()` / `StreamingOptions` pattern shown here works for every streaming +engine. See {doc}`other_engines` for the DaskEngine and SPMDEngine variants. -With the `distributed` scheduler for the `streaming` executor, we recommend using Dask's [built-in diagnostics](https://docs.dask.org/en/stable/diagnostics-distributed.html). +## Attaching to an existing Ray cluster -Finally, the `"in-memory"` *does* support [`LazyFrame.profile`](https://docs.pola.rs/api/python/stable/reference/lazyframe/api/polars.LazyFrame.profile.html). +For multi-node runs, start a Ray cluster separately (for example with `ray start` on each +node) and attach to it from your client script. When Ray is already initialized, +{class}`~cudf_polars.engine.ray.RayEngine` connects to the running +cluster and leaves it untouched on exit: ```python +import ray import polars as pl -q = pl.scan_parquet("ny-taxi/2024/*.parquet").filter(pl.col("total_amount") > 15.0) -profile = q.profile(engine=pl.GPUEngine(executor="in-memory")) +from cudf_polars.engine.ray import RayEngine + +ray.init(address="auto") # attach to a running cluster +with RayEngine() as engine: + result = ( + pl.scan_parquet("s3://bucket/*.parquet") + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) + ) ``` -The result is a tuple containing 2 materialized DataFrames - the first with the query result and the second with profiling information of each node that is executed. -```python -print(profile[0]) -``` -``` -shape: (32_439_327, 19) -┌──────────┬──────────────────────┬───────────────────────┬─────────────────┬───┬───────────────────────┬──────────────┬──────────────────────┬─────────────┐ -│ VendorID ┆ tpep_pickup_datetime ┆ tpep_dropoff_datetime ┆ passenger_count ┆ … ┆ improvement_surcharge ┆ total_amount ┆ congestion_surcharge ┆ Airport_fee │ -│ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │ -│ i32 ┆ datetime[μs] ┆ datetime[μs] ┆ i64 ┆ ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ -╞══════════╪══════════════════════╪═══════════════════════╪═════════════════╪═══╪═══════════════════════╪══════════════╪══════════════════════╪═════════════╡ -│ 2 ┆ 2024-01-01 00:57:55 ┆ 2024-01-01 01:17:43 ┆ 1 ┆ … ┆ 1.0 ┆ 22.7 ┆ 2.5 ┆ 0.0 │ -│ 1 ┆ 2024-01-01 00:03:00 ┆ 2024-01-01 00:09:36 ┆ 1 ┆ … ┆ 1.0 ┆ 18.75 ┆ 2.5 ┆ 0.0 │ -│ 1 ┆ 2024-01-01 00:17:06 ┆ 2024-01-01 00:35:01 ┆ 1 ┆ … ┆ 1.0 ┆ 31.3 ┆ 2.5 ┆ 0.0 │ -│ 1 ┆ 2024-01-01 00:36:38 ┆ 2024-01-01 00:44:56 ┆ 1 ┆ … ┆ 1.0 ┆ 17.0 ┆ 2.5 ┆ 0.0 │ -│ 1 ┆ 2024-01-01 00:46:51 ┆ 2024-01-01 00:52:57 ┆ 1 ┆ … ┆ 1.0 ┆ 16.1 ┆ 2.5 ┆ 0.0 │ -│ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … │ -│ 2 ┆ 2024-12-31 23:05:43 ┆ 2024-12-31 23:18:15 ┆ null ┆ … ┆ 1.0 ┆ 24.67 ┆ null ┆ null │ -│ 2 ┆ 2024-12-31 23:02:00 ┆ 2024-12-31 23:22:14 ┆ null ┆ … ┆ 1.0 ┆ 15.25 ┆ null ┆ null │ -│ 2 ┆ 2024-12-31 23:17:15 ┆ 2024-12-31 23:17:34 ┆ null ┆ … ┆ 1.0 ┆ 24.46 ┆ null ┆ null │ -│ 1 ┆ 2024-12-31 23:14:53 ┆ 2024-12-31 23:35:13 ┆ null ┆ … ┆ 1.0 ┆ 32.88 ┆ null ┆ null │ -│ 1 ┆ 2024-12-31 23:15:33 ┆ 2024-12-31 23:36:29 ┆ null ┆ … ┆ 1.0 ┆ 28.57 ┆ null ┆ null │ -└──────────┴──────────────────────┴───────────────────────┴─────────────────┴───┴───────────────────────┴──────────────┴──────────────────────┴─────────────┘ -``` +{class}`~cudf_polars.engine.ray.RayEngine` creates one rank per GPU in the Ray cluster. +It raises `RuntimeError` if no GPUs are available. -```python -print(profile[1]) -``` -``` -shape: (3, 3) -┌────────────────────┬───────┬────────┐ -│ node ┆ start ┆ end │ -│ --- ┆ --- ┆ --- │ -│ str ┆ u64 ┆ u64 │ -╞════════════════════╪═══════╪════════╡ -│ optimization ┆ 0 ┆ 416 │ -│ gpu-ir-translation ┆ 416 ┆ 741 │ -│ Scan ┆ 813 ┆ 233993 │ -└────────────────────┴───────┴────────┘ -``` +## Manual Engine Lifetime Control -## Tracing - -cudf-polars can optionally trace execution of each node in the query plan. -To enable tracing, set the environment variable ``CUDF_POLARS_LOG_TRACES`` to a -true value ("1", "true", "y", "yes") before starting your process. - -cudf-polars logs traces at three scopes (levels): - -1. `plan`: These generally happen once per query. This will include things - like the (serialized) query plan. -2. `actor`: (rapidsmpf runtime only). There will be roughly one `actor` - trace per node in the logical plan. -3. `evaluate_ir_node`: Logs the evaluation of a physical node in the query plan. - Note that one logical node might expand to more than one physical nodes. - -Each trace includes a `scope` key indicating which level that trace belongs to. -`actor`-scoped nodes will be nested under a `plan`-scoped node. When using the -rapidsmpf runtime, `evaluate_ir_node`-scoped nodes will -be nested under an `actor`-scoped node. - -### Schemas - -The different scopes have different schemas. Fields in **bold** are required / always present. - -#### scope=plan - -| Field Name | Type | Description | -| ---------- | ----- | ----------- | -| **scope** | Literal["plan"] | The string literal `"plan"`. Useful for distinguishing from other types of traces. | -| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | -| **plan** | `PlanObject` | A serialized representation of the query plan. See #TODO below | -| **event** | String | A message like "Query Plan" | - -#### scope=actor - -`actor`-scoped traces will only appear with the rapidsmpf runtime. - -| Field Name | Type | Description | -| ---------- | ----- | ----------- | -| **scope** | Literal["actor"] | The string literal `"actor"`. Useful for distinguishing from other types of traces. | -| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | -| **start** | int | A nanosecond-resolution counter indicating when the actor started. Note: actors generally start early in the query and suspend waiting for data. | -| **stop** | int | A nanosecond-resolution counter indicating when the actor completed. | -| **event** | String | A message like "Streaming Actor". | -| **actor_ir_type** | String | The type of the actor, like `"Scan"`. | -| **actor_ir_id** | int | A unique identifier for the actor. All traces logged under this actor will include this value. | -| chunk_count | int | A counter for how many table chunks have been processed by this actor at the time of logging. | -| duplicated | bool | Whether the output rows are duplicated across ranks (e.g. after an allgather). | -| row_count | int | Total row count produced by this node during execution. | - -#### scope=evaluate_ir_node - -| Field Name | Type | Description | -| ---------- | ----- | ----------- | -| **scope** | `Literal["evaluate_ir_node"]` | The string literal `"evaluate_ir_node"`. Useful for distinguishing from other types of traces. | -| **cudf_polars_query_id** | UUID4 | A unique identifier for the polars query being executed. All traces logged as part of this query use this ID. | -| **type** | string | The name of the IR node | -| **start** | int | A nanosecond-precision counter indicating when this node started executing | -| **stop** | int | A nanosecond-precision counter indicating when this node finished executing | -| **overhead_duration** | int | The overhead, in nanoseconds, added by tracing | -| `count_frames_{phase}` | int | The number of dataframes for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_DATAFRAMES=0`. | -| `frames_{phase}` | `list[dict]` | A list with dictionaries with "shape" and "size" fields, one per input dataframe, for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_DATAFRAMES=0`. | -| `total_bytes_{phase}` | int | The sum of the size (in bytes) of the dataframes for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_current_bytes_{phase}` | int | The current number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_current_count_{phase}` | int | The current number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_peak_bytes_{phase}` | int | The peak number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_peak_count_{phase}` | int | The peak number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_total_bytes_{phase}` | int | The total number of bytes allocated by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `rmm_total_count_{phase}` | int | The total number of allocations made by RMM Memory Resource used by cudf-polars for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| `nvml_current_bytes_{phase}` | int | The device memory usage of this process, as reported by NVML, for the input / output `phase`. This metric can be disabled by setting `CUDF_POLARS_LOG_TRACES_MEMORY=0`. | -| actor_ir_id | int | A unique identifier for the parent actor (rapidsmpf runtime only). | - -Setting `CUDF_POLARS_LOG_TRACES=1` enables all the metrics. Depending on the query, the overhead -from collecting the memory or dataframe metrics can be measurable. You can disable some metrics -through additional environment variables. For example, do disable the memory related metrics, set: +When you need to control the engine lifetime explicitly, for example in a Jupyter notebook +where a `with` block cannot span multiple cells, construct a `RayEngine` once and reuse it, +then call `engine.shutdown()` when you are done: -``` -CUDF_POLARS_LOG_TRACES=1 CUDF_POLARS_LOG_TRACES_MEMORY=0 -``` +```python +# Cell 1: start the engine +from cudf_polars.engine.ray import RayEngine -And to disable the memory and dataframe metrics, which essentially leaves just -the duration metrics, set +engine = RayEngine() ``` -CUDF_POLARS_LOG_TRACES=1 CUDF_POLARS_LOG_TRACES_MEMORY=0 CUDF_POLARS_LOG_TRACES_DATAFRAMES=0 -``` - -Note that tracing still needs to be enabled with `CUDF_POLARS_LOG_TRACES=1`. -The implementation uses [structlog] to build log records. You can configure the -output using structlog's [configuration][structlog-configure] and enrich the -records with [context variables][structlog-context]. +```python +# Cell 2: run a query +import polars as pl -``` ->>> df = pl.DataFrame({"a": ["a", "a", "b"], "b": [1, 2, 3]}).lazy() ->>> df.group_by("a").agg(pl.col("b").min().alias("min"), pl.col("b").max().alias("max")).collect(engine="gpu") -2025-09-10 07:44:01 [info ] Execute IR count_frames_input=0 count_frames_output=1 ... type=DataFrameScan -2025-09-10 07:44:01 [info ] Execute IR count_frames_input=1 count_frames_output=1 ... type=GroupBy -shape: (2, 3) -┌─────┬─────┬─────┐ -│ a ┆ min ┆ max │ -│ --- ┆ --- ┆ --- │ -│ str ┆ i64 ┆ i64 │ -╞═════╪═════╪═════╡ -│ b ┆ 3 ┆ 3 │ -│ a ┆ 1 ┆ 2 │ -└─────┴─────┴─────┘ +result = ( + pl.scan_parquet("/data/*.parquet") + .group_by("customer_id") + .agg(pl.col("amount").sum()) + .collect(engine=engine) +) +result ``` -### Serialized Query Plan +```python +# Cell 3: run another query reusing the same engine +other = pl.scan_parquet("/data/other/*.parquet").collect(engine=engine) +``` -The query plan is serialized with the following schema: +```python +# Final cell: tear everything down +engine.shutdown() +``` -| Field Name | Type | Description | -| ---------- | ---- | ----------- | -| roots | `list` | A list of string node IDs for the root nodes | -| nodes | `Mapping` | A mapping from string node ID to the Node | -| partition_info | `Mapping` | A mapping from string node ID to the Node | +`engine.shutdown()` stops the GPU worker processes (rank actors) and, if the engine started Ray itself, +also calls `ray.shutdown()`. It is idempotent, so calling it twice is safe. -`nodes` and `partition_info` are flat: they contain every node in the query plan. +## Sink behavior -`IRNode` objects have the following schema: +When a streaming engine is used, sink operations such as `df.sink_parquet("my_path")` always produce +a directory containing one or more files. It is not currently possible to disable this behavior, and +setting `sink_to_directory=False` raises a `ValueError`. -| Field Name | Type | Description | -| ---------- | ---- | ----------- | -| id | `str` | The string node ID. This is unique within the query plan | -| children | `list` | The node IDs of this node's children nodes. Each child node ID is also available in the plan's `nodes` field. | -| schema | `Mapping` | A mapping from column name to (string) data type identifier. | -| properties | `Mapping` | Additional properties, unique to each node type. | -| type | `str` | The name of the IR node. | +The in-memory engine, by contrast, follows standard Polars semantics and writes to a single file at +the specified path. -`PartitionInfo` objects have the following schema: +## Cluster diagnostics -| Field Name | Type | Description | -| ---------- | ---- | ----------- | -| count | int | The number of partitions for this node | -| partitioned_on | `list[str]` | The columns this node is partitioned on. | +{meth}`~cudf_polars.engine.ray.RayEngine.gather_cluster_info` returns +a list of {class}`~cudf_polars.engine.core.ClusterInfo`, one per rank +actor, with fields `hostname`, `pid`, `cuda_visible_devices`, and `gpu_uuid`: +```python +with RayEngine() as engine: + print(f"cluster has {engine.nranks} ranks") + for i, info in enumerate(engine.gather_cluster_info()): + print( + f"rank {i}: hostname={info.hostname}, pid={info.pid}, " + f"cuda_visible_devices={info.cuda_visible_devices}, " + f"gpu_uuid={info.gpu_uuid}" + ) +# rank 0: hostname=node-0, pid=12345, cuda_visible_devices=0, gpu_uuid=GPU-abc123... +# rank 1: hostname=node-0, pid=12346, cuda_visible_devices=1, gpu_uuid=GPU-def456... +``` -[nvml]: https://developer.nvidia.com/management-library-nvml -[rmm-stats]: https://docs.rapids.ai/api/rmm/stable/guide/#memory-statistics-and-profiling -[structlog]: https://www.structlog.org/ -[structlog-configure]: https://www.structlog.org/en/stable/configuration.html -[structlog-context]: https://www.structlog.org/en/stable/contextvars.html +[ray-docs]: https://docs.ray.io/ diff --git a/python/cudf_polars/cudf_polars/engine/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py index 509980fc2c2a..fcb0b9ce280a 100644 --- a/python/cudf_polars/cudf_polars/engine/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -761,7 +761,7 @@ def from_options( dask_client: distributed.Client | None = None, ) -> DaskEngine: """ - Create a :class:`DaskEngine` from a :class:`StreamingOptions` object. + Create a :class:`DaskEngine` from a :class:`~cudf_polars.engine.options.StreamingOptions` object. This is the recommended way to construct a ``DaskEngine`` for typical use. All RapidsMPF, executor, and engine options are read from @@ -806,7 +806,7 @@ def gather_cluster_info(self) -> list[ClusterInfo]: Returns ------- - List of :class:`ClusterInfo`, one per rank. + List of :class:`~cudf_polars.engine.core.ClusterInfo`, one per rank. """ return list(self._dask_ctx.client.run(ClusterInfo.local).values()) diff --git a/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py b/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py index 73f08730df63..fe3f03698103 100644 --- a/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py +++ b/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -"""Single-GPU, single-instance specialization of :class:`SPMDEngine`.""" +"""Single-GPU, single-instance specialization of :class:`~cudf_polars.engine.spmd.SPMDEngine`.""" from __future__ import annotations @@ -221,7 +221,7 @@ def check_no_live_default_singleton(self_engine: Any) -> None: class DefaultSingletonEngine(SPMDEngine): """ - Process-wide single-GPU singleton specialization of :class:`SPMDEngine`. + Process-wide single-GPU singleton specialization of :class:`~cudf_polars.engine.spmd.SPMDEngine`. At most one live instance exists per process. Use :meth:`get_or_create` to obtain it and :meth:`shutdown` to tear it down. @@ -230,7 +230,8 @@ class DefaultSingletonEngine(SPMDEngine): executor, and engine settings from the environment. Users needing custom configuration should construct an engine explicitly. - See :class:`RayEngine`, :class:`DaskEngine`, and :class:`SPMDEngine`. + See :class:`~cudf_polars.engine.ray.RayEngine`, :class:`~cudf_polars.engine.dask.DaskEngine`, + and :class:`~cudf_polars.engine.spmd.SPMDEngine`. Examples -------- @@ -275,7 +276,7 @@ def get_or_create(cls) -> DefaultSingletonEngine: Raises ------ RuntimeError - If any other :class:`StreamingEngine` is currently alive. + If any other :class:`~cudf_polars.engine.core.StreamingEngine` is currently alive. """ with _state.lock: if _state.instance is not None: @@ -292,7 +293,7 @@ def shutdown() -> None: Submits teardown to the dedicated worker thread, the same thread that constructed the rapidsmpf ``Context``, and waits up to - :data:`SHUTDOWN_TIMEOUT_SECONDS`. + ``SHUTDOWN_TIMEOUT_SECONDS`` seconds. """ with _state.lock: instance = _state.instance diff --git a/python/cudf_polars/cudf_polars/engine/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py index 532e863bdf72..25b5206b514c 100644 --- a/python/cudf_polars/cudf_polars/engine/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -470,15 +470,15 @@ class RayEngine(StreamingEngine): Hardware binding is disabled implicitly but the caller must pass ``engine_options={"allow_gpu_sharing": True}`` explicitly to acknowledge the multi-tenant GPU semantics. - .. note:: - Oversubscription does not increase throughput. When multiple - ranks share a GPU, they compete for the same compute and - memory resources, which may increase memory pressure and - reduce overall performance. This option is primarily useful - for testing multi-rank code paths on machines with fewer - GPUs than ranks, and for downstream projects that need to - validate distributed logic in resource-constrained CI - environments. + + Note, oversubscription does not increase throughput. When multiple + ranks share a GPU, they compete for the same compute and + memory resources, which may increase memory pressure and + reduce overall performance. This option is primarily useful + for testing multi-rank code paths on machines with fewer + GPUs than ranks, and for downstream projects that need to + validate distributed logic in resource-constrained CI + environments. Raises ------ @@ -671,7 +671,7 @@ def from_options( ray_init_options: dict[str, object] | None = None, ) -> RayEngine: """ - Create a :class:`RayEngine` from a :class:`StreamingOptions` object. + Create a :class:`RayEngine` from a :class:`~cudf_polars.engine.options.StreamingOptions` object. This is the recommended way to construct a ``RayEngine`` for typical use. All RapidsMPF, executor, and engine options are read from @@ -725,7 +725,7 @@ def gather_cluster_info(self) -> list[ClusterInfo]: Returns ------- - List of :class:`ClusterInfo`, one per rank. + List of :class:`~cudf_polars.engine.core.ClusterInfo`, one per rank. """ return ray.get([rank.get_info.remote() for rank in self.rank_actors]) diff --git a/python/cudf_polars/cudf_polars/engine/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py index 64e7b1893501..18804f410ef8 100644 --- a/python/cudf_polars/cudf_polars/engine/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -141,7 +141,7 @@ def allgather_polars_dataframe( Rank-local DataFrame to contribute. op_id Operation ID for this AllGather collective. Must be identical on every - rank. For example, use :func:`reserve_op_id` to obtain a collision-free + rank. For example, use :func:`~cudf_polars.streaming.actor_graph.collectives.common.reserve_op_id` to obtain a collision-free ID from the same pool used internally by cudf-polars. Avoid passing hardcoded integers. @@ -309,7 +309,7 @@ class SPMDEngine(StreamingEngine): time, before RMM and communicator initialisation, so that CPU affinity, NUMA memory policy, and ``UCX_NET_DEVICES`` are set as early as possible. By default, binding is skipped under ``rrun`` (which already performs its own binding), - see :attr:`HardwareBindingPolicy.skip_under_rrun`. + see ``HardwareBindingPolicy.skip_under_rrun``. Examples -------- @@ -425,7 +425,7 @@ def _cleanup_ctx(self) -> None: @classmethod def from_options(cls, options: StreamingOptions) -> SPMDEngine: """ - Create an :class:`SPMDEngine` from a :class:`StreamingOptions` object. + Create an :class:`SPMDEngine` from a :class:`~cudf_polars.engine.options.StreamingOptions` object. This is the recommended way to construct an ``SPMDEngine`` for typical use. All RapidsMPF, executor, and engine options are read from @@ -590,7 +590,7 @@ def gather_cluster_info(self) -> list[ClusterInfo]: Returns ------- - List of :class:`ClusterInfo`, one per rank. + List of :class:`~cudf_polars.engine.core.ClusterInfo`, one per rank. """ data = json.dumps(dataclasses.asdict(ClusterInfo.local())).encode() with reserve_op_id() as op_id: