diff --git a/cpp/include/cudf/context.hpp b/cpp/include/cudf/context.hpp index e9172c91cdbf..7eb6f5d273fc 100644 --- a/cpp/include/cudf/context.hpp +++ b/cpp/include/cudf/context.hpp @@ -78,4 +78,28 @@ void initialize(init_flags flags = init_flags::INIT_JIT_CACHE); /// teardown and that only one thread calls teardown at a time. void teardown(); +/** + * @brief Enable or disable the JIT program cache + * + * When disabled, the cache will not be used for + * storing or retrieving compiled programs, effectively bypassing the cache. When enabled, the + * cache will be used as normal. This can be used to temporarily disable caching without clearing + * the existing cache contents, allowing for easy re-enabling of the cache later. + * + * @param enable If `true`, the JIT program cache is enabled; if `false`, it is disabled. + */ +void enable_jit_cache(bool enable); + +/** + * @brief Clear the JIT program cache, removing all cached programs from memory and disk. + * + * This is a more expensive operation than simply disabling the cache, as it involves deleting + * cached files from disk, but it also frees up any memory used by the cached programs. Use + * `enable_jit_cache(false)` if you want to temporarily disable caching without clearing existing + * cache contents. + * + * @warning For benchmarking or testing purposes, prefer `enable_jit_cache`. + */ +void clear_jit_cache(); + } // namespace CUDF_EXPORT cudf diff --git a/cpp/include/cudf/detail/aggregation/aggregation.hpp b/cpp/include/cudf/detail/aggregation/aggregation.hpp index eb9cb08c4477..714c6128e984 100644 --- a/cpp/include/cudf/detail/aggregation/aggregation.hpp +++ b/cpp/include/cudf/detail/aggregation/aggregation.hpp @@ -136,7 +136,8 @@ class count_aggregation final : public clonable::derived_from { + reduce_aggregation, + scan_aggregation> { public: count_aggregation(aggregation::Kind kind) : aggregation(kind) {} }; diff --git a/cpp/include/cudf/detail/unary.hpp b/cpp/include/cudf/detail/unary.hpp index 285a798a4787..e55ccfb16aea 100644 --- a/cpp/include/cudf/detail/unary.hpp +++ b/cpp/include/cudf/detail/unary.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -8,50 +8,12 @@ #include #include #include -#include #include #include -#include -#include - -namespace CUDF_EXPORT cudf { +namespace cudf { namespace detail { -/** - * @brief Creates a column of `type_id::BOOL8` elements by applying a predicate to every element - * between - * [`begin, `end`) `true` indicates the value is satisfies the predicate and `false` indicates it - * doesn't. - * - * @tparam InputIterator Iterator type for `begin` and `end` - * @tparam Predicate A predicator type which will be evaluated - * @param begin Beginning of the sequence of elements - * @param end End of the sequence of elements - * @param p Predicate to be applied to each element in `[begin,end)` - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param mr Device memory resource used to allocate the returned column's device memory - * - * @returns A column of type `type_id::BOOL8,` with `true` representing predicate is satisfied. - */ - -template -std::unique_ptr true_if(InputIterator begin, - InputIterator end, - size_type size, - Predicate p, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - auto output = - make_numeric_column(data_type(type_id::BOOL8), size, mask_state::UNALLOCATED, stream, mr); - auto output_mutable_view = output->mutable_view(); - auto output_data = output_mutable_view.data(); - - thrust::transform(rmm::exec_policy_nosync(stream), begin, end, output_data, p); - - return output; -} /** * @copydoc cudf::unary_operation @@ -91,4 +53,4 @@ std::unique_ptr is_not_nan(cudf::column_view const& input, rmm::device_async_resource_ref mr); } // namespace detail -} // namespace CUDF_EXPORT cudf +} // namespace cudf diff --git a/cpp/src/aggregation/aggregation.cpp b/cpp/src/aggregation/aggregation.cpp index 6a4273a28534..ab871abf94ed 100644 --- a/cpp/src/aggregation/aggregation.cpp +++ b/cpp/src/aggregation/aggregation.cpp @@ -121,6 +121,8 @@ template CUDF_EXPORT std::unique_ptr make_count_aggregation(null_policy null_handling); template CUDF_EXPORT std::unique_ptr make_count_aggregation( null_policy null_handling); +template CUDF_EXPORT std::unique_ptr make_count_aggregation( + null_policy null_handling); /// Factory to create a HISTOGRAM aggregation template diff --git a/cpp/src/io/parquet/experimental/page_index_filter.cu b/cpp/src/io/parquet/experimental/page_index_filter.cu index 1c63b65e24c2..51bb397a6bab 100644 --- a/cpp/src/io/parquet/experimental/page_index_filter.cu +++ b/cpp/src/io/parquet/experimental/page_index_filter.cu @@ -38,6 +38,7 @@ #include #include +#include namespace cudf::io::parquet::experimental::detail { @@ -911,7 +912,7 @@ std::unique_ptr aggregate_reader_metadata::build_row_mask_with_pag [&](auto col_idx) { auto const schema_idx = output_column_schemas[col_idx]; auto const& dtype = output_dtypes[col_idx]; - // Only participating columns and comparable types except fixed point are supported + // Only participating columns and comparable types are supported if (not stats_columns_mask[col_idx] or (cudf::is_compound(dtype) && dtype.id() != cudf::type_id::STRING)) { // Placeholder for unsupported types and non-participating columns diff --git a/cpp/src/io/parquet/predicate_pushdown.cpp b/cpp/src/io/parquet/predicate_pushdown.cpp index 03de07c5bf9a..fa6bf42bf8ac 100644 --- a/cpp/src/io/parquet/predicate_pushdown.cpp +++ b/cpp/src/io/parquet/predicate_pushdown.cpp @@ -146,7 +146,7 @@ std::optional>> aggregate_reader_metadata::ap for (size_t col_idx = 0; col_idx < output_dtypes.size(); col_idx++) { auto const schema_idx = output_column_schemas[col_idx]; auto const& dtype = output_dtypes[col_idx]; - // Only participating columns and comparable types except fixed point are supported + // Only participating columns and comparable types are supported if (not stats_columns_mask[col_idx] or (cudf::is_compound(dtype) && dtype.id() != cudf::type_id::STRING)) { // Placeholder for unsupported types and non-participating columns diff --git a/cpp/src/io/parquet/stats_filter_helpers.hpp b/cpp/src/io/parquet/stats_filter_helpers.hpp index ec46ed558111..89128de03acb 100644 --- a/cpp/src/io/parquet/stats_filter_helpers.hpp +++ b/cpp/src/io/parquet/stats_filter_helpers.hpp @@ -18,8 +18,9 @@ #include #include -#include -#include +#include +#include +#include namespace cudf::io::parquet::detail { @@ -38,8 +39,34 @@ constexpr size_t initial_chars_capacity = 1024; */ class stats_caster_base { protected: + static inline numeric::decimal128::rep decode_flba_decimal128(uint8_t const* stats_val) + { + auto constexpr endianness = std::endian::native; + static_assert(endianness == std::endian::little or endianness == std::endian::big, + "Encountered unsupported endianness while decoding decimal128 from FLBA"); + using RepType = numeric::decimal128::rep; + auto value = RepType{}; + std::memcpy(&value, stats_val, sizeof(RepType)); + auto value_rep = std::bit_cast>(value); + // byte-swap to native representation on little-endian platforms + if constexpr (endianness == std::endian::little) { std::ranges::reverse(value_rep); } + return std::bit_cast(value_rep); + } + + template + static inline T decode_fixed_width_value(uint8_t const* stats_val, size_t stats_size) + requires((cudf::is_integral() and !cudf::is_boolean()) or cudf::is_fixed_point() or + cudf::is_chrono()) + { + CUDF_EXPECTS(stats_size == sizeof(T), + "Parquet reader encountered a statistics vector larger than the type's size"); + auto value = T{}; + std::memcpy(&value, stats_val, std::min(stats_size, sizeof(T))); + return value; + } + template - static inline ToType targetType(FromType const value) + static inline ToType target_type(FromType const value) { if constexpr (cudf::is_timestamp()) { return static_cast( @@ -52,43 +79,53 @@ class stats_caster_base { } // uses storage type as T - template () or cudf::is_nested())> + template static inline T convert(uint8_t const* stats_val, size_t stats_size, Type const type) + requires(cudf::is_dictionary() or cudf::is_nested()) { CUDF_FAIL("unsupported type for stats casting"); } - template ())> + template static inline T convert(uint8_t const* stats_val, size_t stats_size, Type const type) + requires(cudf::is_boolean()) { CUDF_EXPECTS(type == Type::BOOLEAN, "Invalid type and stats combination"); - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + return stats_caster_base::target_type(*reinterpret_cast(stats_val)); } // integral but not boolean, and fixed_point, and chrono. - template () and !cudf::is_boolean()) or - cudf::is_fixed_point() or cudf::is_chrono())> + template static inline T convert(uint8_t const* stats_val, size_t stats_size, Type const type) + requires((cudf::is_integral() and !cudf::is_boolean()) or cudf::is_fixed_point() or + cudf::is_chrono()) { switch (type) { case Type::INT32: - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + return stats_caster_base::target_type( + decode_fixed_width_value(stats_val, stats_size)); case Type::INT64: - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + return stats_caster_base::target_type( + decode_fixed_width_value(stats_val, stats_size)); case Type::INT96: // Deprecated in parquet specification - return stats_caster_base::targetType( - static_cast<__int128_t>(reinterpret_cast(stats_val)[0]) << 32 | - reinterpret_cast(stats_val)[2]); + return stats_caster_base::target_type( + static_cast<__int128_t>(decode_fixed_width_value(stats_val, stats_size)) << 32 | + decode_fixed_width_value(stats_val + sizeof(int64_t), stats_size)); case Type::BYTE_ARRAY: [[fallthrough]]; case Type::FIXED_LEN_BYTE_ARRAY: if (stats_size == sizeof(T)) { - // if type size == length of stats_val. then typecast and return. if constexpr (cudf::is_chrono()) { - return stats_caster_base::targetType( - *reinterpret_cast(stats_val)); + return stats_caster_base::target_type( + decode_fixed_width_value(stats_val, stats_size)); + } else if constexpr (std::is_same_v) { + // Decimals with physical type FLBA/BYTE_ARRAY are stored as two's complement using + // big-endian. + return stats_caster_base::target_type(decode_flba_decimal128(stats_val)); } else { - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + // TODO(mh): We may need to add support for `decimal256` (two's complement using + // big-endian) and `UUID` types (big-endian) + return stats_caster_base::target_type( + decode_fixed_width_value(stats_val, stats_size)); } } // unsupported type @@ -96,20 +133,22 @@ class stats_caster_base { } } - template ())> + template static inline T convert(uint8_t const* stats_val, size_t stats_size, Type const type) + requires(cudf::is_floating_point()) { switch (type) { case Type::FLOAT: - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + return stats_caster_base::target_type(*reinterpret_cast(stats_val)); case Type::DOUBLE: - return stats_caster_base::targetType(*reinterpret_cast(stats_val)); + return stats_caster_base::target_type(*reinterpret_cast(stats_val)); default: CUDF_FAIL("Invalid type and stats combination"); } } - template )> + template static inline T convert(uint8_t const* stats_val, size_t stats_size, Type const type) + requires(std::is_same_v) { switch (type) { case Type::BYTE_ARRAY: [[fallthrough]]; diff --git a/cpp/src/jit/cache.cpp b/cpp/src/jit/cache.cpp index 7d03ed8fdbca..b988296b0b7e 100644 --- a/cpp/src/jit/cache.cpp +++ b/cpp/src/jit/cache.cpp @@ -1,8 +1,9 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ +#include "io/utilities/getenv_or.hpp" #include "runtime/context.hpp" #include @@ -95,40 +96,63 @@ std::string get_program_cache_dir() #endif } -std::size_t try_parse_numeric_env_var(char const* const env_name, std::size_t default_val) -{ - auto const value = std::getenv(env_name); - return value != nullptr ? std::stoull(value) : default_val; -} } // namespace jitify2::ProgramCache<>& jit::program_cache::get(jitify2::PreprocessedProgramData const& preprog) { CUDF_FUNC_RANGE(); - std::lock_guard const caches_lock(_caches_mutex); + std::lock_guard caches_lock(_caches_mutex); auto existing_cache = _caches.find(preprog.name()); - if (existing_cache == _caches.end()) { - auto const kernel_limit_proc = - try_parse_numeric_env_var("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 10'000); - auto const kernel_limit_disk = - try_parse_numeric_env_var("LIBCUDF_KERNEL_CACHE_LIMIT_DISK", 100'000); - - // if kernel_limit_disk is zero, jitify will assign it the value of kernel_limit_proc. - // to avoid this, we treat zero as "disable disk caching" by not providing the cache dir. - auto const cache_dir = kernel_limit_disk == 0 ? std::string{} : get_program_cache_dir(); - - auto const res = - _caches.insert({preprog.name(), + if (existing_cache == _caches.end() || _disabled.load(std::memory_order_seq_cst)) { + auto res = + _caches.emplace(preprog.name(), std::make_unique>( - kernel_limit_proc, preprog, nullptr, cache_dir, kernel_limit_disk)}); + _kernel_limit_proc, preprog, nullptr, _cache_dir, _kernel_limit_disk)); existing_cache = res.first; } return *(existing_cache->second); } +void jit::program_cache::clear() +{ + CUDF_FUNC_RANGE(); + std::lock_guard caches_lock(_caches_mutex); + + _caches.clear(); + + // non-atomic + std::filesystem::remove_all(_cache_dir); +} + +void jit::program_cache::enable(bool enable) +{ + _disabled.store(!enable, std::memory_order_seq_cst); +} + +bool jit::program_cache::is_enabled() const { return !_disabled.load(std::memory_order_seq_cst); } + +std::unique_ptr jit::program_cache::create() +{ + auto const kernel_limit_proc = getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_PER_PROCESS", 10'000); + auto const kernel_limit_disk = getenv_or("LIBCUDF_KERNEL_CACHE_LIMIT_DISK", 100'000); + auto const disabled = get_bool_env_or("LIBCUDF_KERNEL_CACHE_DISABLED", false); + auto const clear_cache = get_bool_env_or("LIBCUDF_KERNEL_CACHE_CLEAR", false); + + // if kernel_limit_disk is zero, jitify will assign it the value of kernel_limit_proc. + // to avoid this, we treat zero as "disable disk caching" by not providing the cache dir. + auto cache_dir = kernel_limit_disk == 0 ? std::string{} : get_program_cache_dir(); + + auto cache = + std::make_unique(kernel_limit_proc, kernel_limit_disk, cache_dir, disabled); + + if (clear_cache) { cache->clear(); } + + return cache; +} + jitify2::ProgramCache<>& jit::get_program_cache(jitify2::PreprocessedProgramData const& preprog) { return cudf::get_context().program_cache().get(preprog); diff --git a/cpp/src/jit/cache.hpp b/cpp/src/jit/cache.hpp index c130b953fc9f..0e6738fed176 100644 --- a/cpp/src/jit/cache.hpp +++ b/cpp/src/jit/cache.hpp @@ -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 */ @@ -11,6 +11,8 @@ #include +#include +#include #include #include #include @@ -21,9 +23,23 @@ namespace jit { class program_cache { std::mutex _caches_mutex; std::unordered_map>> _caches; + int32_t _kernel_limit_proc; + int32_t _kernel_limit_disk; + std::filesystem::path _cache_dir; + std::atomic _disabled; public: - program_cache() = default; + program_cache(int32_t kernel_limit_proc, + int32_t kernel_limit_disk, + std::filesystem::path cache_dir, + bool disabled) + : _kernel_limit_proc{kernel_limit_proc}, + _kernel_limit_disk{kernel_limit_disk}, + _cache_dir{std::move(cache_dir)}, + _disabled{disabled} + { + } + program_cache(program_cache const&) = delete; program_cache(program_cache&&) = delete; program_cache& operator=(program_cache const&) = delete; @@ -31,6 +47,14 @@ class program_cache { ~program_cache() = default; jitify2::ProgramCache<>& get(jitify2::PreprocessedProgramData const& preprog); + + void clear(); + + void enable(bool enable); + + bool is_enabled() const; + + static std::unique_ptr create(); }; jitify2::ProgramCache<>& get_program_cache(jitify2::PreprocessedProgramData const& preprog); diff --git a/cpp/src/reductions/scan/scan.cuh b/cpp/src/reductions/scan/scan.cuh index 22ba5ce5ecbe..fe27ff96626f 100644 --- a/cpp/src/reductions/scan/scan.cuh +++ b/cpp/src/reductions/scan/scan.cuh @@ -55,6 +55,12 @@ std::unique_ptr scan_agg_dispatch(column_view const& input, return type_dispatcher( input.type(), DispatchFn(), input, output_mask, stream, mr); case aggregation::EWMA: return exponentially_weighted_moving_average(input, agg, stream, mr); + case aggregation::COUNT_VALID: + return type_dispatcher( + input.type(), DispatchFn(), input, output_mask, stream, mr); + case aggregation::COUNT_ALL: + return type_dispatcher( + input.type(), DispatchFn(), input, nullptr, stream, mr); default: CUDF_FAIL("Unsupported aggregation operator for scan"); } } diff --git a/cpp/src/reductions/scan/scan_inclusive.cu b/cpp/src/reductions/scan/scan_inclusive.cu index 0e203c2d5e3d..b209b64ded73 100644 --- a/cpp/src/reductions/scan/scan_inclusive.cu +++ b/cpp/src/reductions/scan/scan_inclusive.cu @@ -90,6 +90,7 @@ struct scan_functor { }; template + requires(not std::is_same_v) struct scan_functor { static std::unique_ptr invoke(column_view const& input_view, bitmask_type const* mask, @@ -101,6 +102,7 @@ struct scan_functor { }; template + requires(not std::is_same_v) struct scan_functor { static std::unique_ptr invoke(column_view const& input, bitmask_type const*, @@ -111,6 +113,34 @@ struct scan_functor { } }; +template + requires(std::is_same_v) +struct scan_functor { + static std::unique_ptr invoke(column_view const& input_view, + bitmask_type const* mask, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) + { + auto output_column = make_numeric_column(data_type{type_to_id()}, + input_view.size(), + cudf::mask_state::UNALLOCATED, + stream, + mr); + auto result = output_column->mutable_view(); + + auto const begin = make_counting_transform_iterator( + 0, cuda::proclaim_return_type([mask] __device__(auto idx) -> size_type { + return static_cast(mask == nullptr || bit_is_set(mask, idx)); + })); + + thrust::inclusive_scan( + rmm::exec_policy_nosync(stream), begin, begin + input_view.size(), result.data()); + + CUDF_CHECK_CUDA(stream.value()); + return output_column; + } +}; + /** * @brief Dispatcher for running a Scan operation on an input column * @@ -122,6 +152,7 @@ struct scan_dispatcher { template static constexpr bool is_supported() { + if constexpr (std::is_same_v) { return true; } if constexpr (std::is_same_v) { return std::is_same_v || std::is_same_v; } else { diff --git a/cpp/src/runtime/context.cpp b/cpp/src/runtime/context.cpp index 561840b9a587..6e78060e8792 100644 --- a/cpp/src/runtime/context.cpp +++ b/cpp/src/runtime/context.cpp @@ -27,7 +27,7 @@ void context::ensure_nvcomp_loaded() { io::detail::nvcomp::load_nvcomp_library() void context::ensure_jit_cache_initialized() { std::call_once(_program_cache_init_flag, - [&]() { _program_cache = std::make_unique(); }); + [&]() { _program_cache = jit::program_cache::create(); }); } jit::program_cache& context::program_cache() @@ -85,6 +85,10 @@ void teardown() }); } +void enable_jit_cache(bool enable) { get_context().program_cache().enable(enable); } + +void clear_jit_cache() { get_context().program_cache().clear(); } + context& get_context() { cudf::initialize(); diff --git a/cpp/src/transform/compute_column.cu b/cpp/src/transform/compute_column.cu index b512aafeb83d..1b1e02ccdb6c 100644 --- a/cpp/src/transform/compute_column.cu +++ b/cpp/src/transform/compute_column.cu @@ -1,9 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #include "compute_column_kernel.hpp" +#include "runtime/context.hpp" #include #include @@ -31,6 +32,8 @@ std::unique_ptr compute_column(table_view const& table, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { + if (get_context().use_jit()) { return compute_column_jit(table, expr, stream, mr); } + // If evaluating the expression may produce null outputs we create a nullable // output column and follow the null-supporting expression evaluation code // path. diff --git a/cpp/src/transform/transform.cu b/cpp/src/transform/transform.cu index c909b94efe32..55ab90913bdc 100644 --- a/cpp/src/transform/transform.cu +++ b/cpp/src/transform/transform.cu @@ -148,6 +148,7 @@ auto to_device_input_arg(InputsView inputs, rmm::device_async_resource_ref mr) { std::vector columns; + for (auto const& input : inputs) { columns.emplace_back(std::visit([](auto const& col) { return to_column_view(col); }, input)); } diff --git a/cpp/src/unary/math_ops.cu b/cpp/src/unary/math_ops.cu index 742735e848d5..8aae1c01d85b 100644 --- a/cpp/src/unary/math_ops.cu +++ b/cpp/src/unary/math_ops.cu @@ -15,6 +15,7 @@ #include #include +#include #include #include diff --git a/cpp/src/unary/nan_ops.cu b/cpp/src/unary/nan_ops.cu index c5be3d038342..5cfaa0e170d7 100644 --- a/cpp/src/unary/nan_ops.cu +++ b/cpp/src/unary/nan_ops.cu @@ -1,8 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ +#include "true_if.cuh" + #include #include #include diff --git a/cpp/src/unary/null_ops.cu b/cpp/src/unary/null_ops.cu index d38427cc0d28..cff98f234146 100644 --- a/cpp/src/unary/null_ops.cu +++ b/cpp/src/unary/null_ops.cu @@ -1,8 +1,10 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ +#include "true_if.cuh" + #include #include #include diff --git a/cpp/src/unary/true_if.cuh b/cpp/src/unary/true_if.cuh new file mode 100644 index 000000000000..90bd4290e3a0 --- /dev/null +++ b/cpp/src/unary/true_if.cuh @@ -0,0 +1,55 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include + +#include +#include + +#include + +namespace cudf { +namespace detail { +/** + * @brief Creates a column of `type_id::BOOL8` elements by applying a predicate to every element + * between + * [`begin, `end`) `true` indicates the value is satisfies the predicate and `false` indicates it + * doesn't. + * + * @tparam InputIterator Iterator type for `begin` and `end` + * @tparam Predicate A predicator type which will be evaluated + * @param begin Beginning of the sequence of elements + * @param end End of the sequence of elements + * @param size Size of the output column + * @param p Predicate to be applied to each element in `[begin,end)` + * @param stream CUDA stream used for device memory operations and kernel launches. + * @param mr Device memory resource used to allocate the returned column's device memory + * + * @returns A column of type `type_id::BOOL8,` with `true` representing predicate is satisfied. + */ + +template +std::unique_ptr true_if(InputIterator begin, + InputIterator end, + size_type size, + Predicate p, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + auto output = + make_numeric_column(data_type(type_id::BOOL8), size, mask_state::UNALLOCATED, stream, mr); + auto output_mutable_view = output->mutable_view(); + auto output_data = output_mutable_view.data(); + + thrust::transform(rmm::exec_policy_nosync(stream), begin, end, output_data, p); + + return output; +} + +} // namespace detail +} // namespace cudf diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 6d01bc80af83..f13a3af71b69 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -401,7 +401,6 @@ ConfigureTest( copying/copy_if_else_nested_tests.cpp copying/copy_range_tests.cpp copying/copy_tests.cpp - copying/detail_gather_tests.cu copying/gather_list_tests.cpp copying/gather_str_tests.cpp copying/gather_struct_tests.cpp diff --git a/cpp/tests/copying/detail_gather_tests.cu b/cpp/tests/copying/detail_gather_tests.cu deleted file mode 100644 index 92c78550c83c..000000000000 --- a/cpp/tests/copying/detail_gather_tests.cu +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2020-2024, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include - -template -class GatherTest : public cudf::test::BaseFixture {}; - -TYPED_TEST_SUITE(GatherTest, cudf::test::NumericTypes); - -// This test exercises using different iterator types as gather map inputs -// to cudf::detail::gather -- device_uvector and raw pointers. -TYPED_TEST(GatherTest, GatherDetailDeviceVectorTest) -{ - constexpr cudf::size_type source_size{1000}; - rmm::device_uvector gather_map(source_size, cudf::get_default_stream()); - thrust::sequence( - rmm::exec_policy_nosync(cudf::get_default_stream()), gather_map.begin(), gather_map.end()); - - auto data = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); - cudf::test::fixed_width_column_wrapper source_column(data, data + source_size); - - cudf::table_view source_table({source_column}); - - // test with device vector iterators - { - std::unique_ptr result = - cudf::detail::gather(source_table, - gather_map.begin(), - gather_map.end(), - cudf::out_of_bounds_policy::DONT_CHECK, - cudf::get_default_stream(), - cudf::get_current_device_resource_ref()); - - for (auto i = 0; i < source_table.num_columns(); ++i) { - CUDF_TEST_EXPECT_COLUMNS_EQUAL(source_table.column(i), result->view().column(i)); - } - - CUDF_TEST_EXPECT_TABLES_EQUAL(source_table, result->view()); - } - - // test with raw pointers - { - std::unique_ptr result = - cudf::detail::gather(source_table, - gather_map.begin(), - gather_map.data() + gather_map.size(), - cudf::out_of_bounds_policy::DONT_CHECK, - cudf::get_default_stream(), - cudf::get_current_device_resource_ref()); - - for (auto i = 0; i < source_table.num_columns(); ++i) { - CUDF_TEST_EXPECT_COLUMNS_EQUAL(source_table.column(i), result->view().column(i)); - } - - CUDF_TEST_EXPECT_TABLES_EQUAL(source_table, result->view()); - } -} - -TYPED_TEST(GatherTest, GatherDetailInvalidIndexTest) -{ - constexpr cudf::size_type source_size{1000}; - - auto data = cudf::detail::make_counting_transform_iterator(0, [](auto i) { return i; }); - cudf::test::fixed_width_column_wrapper source_column(data, data + source_size); - auto gather_map_data = - cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 2) ? -1 : i; }); - cudf::test::fixed_width_column_wrapper gather_map(gather_map_data, - gather_map_data + (source_size * 2)); - - cudf::table_view source_table({source_column}); - std::unique_ptr result = - cudf::detail::gather(source_table, - gather_map, - cudf::out_of_bounds_policy::NULLIFY, - cudf::detail::negative_index_policy::NOT_ALLOWED, - cudf::get_default_stream(), - cudf::get_current_device_resource_ref()); - - auto expect_data = - cudf::detail::make_counting_transform_iterator(0, [](auto i) { return (i % 2) ? 0 : i; }); - auto expect_valid = cudf::detail::make_counting_transform_iterator( - 0, [](auto i) { return (i % 2) || (i >= source_size) ? 0 : 1; }); - cudf::test::fixed_width_column_wrapper expect_column( - expect_data, expect_data + (source_size * 2), expect_valid); - - for (auto i = 0; i < source_table.num_columns(); ++i) { - CUDF_TEST_EXPECT_COLUMNS_EQUAL(expect_column, result->view().column(i)); - } -} diff --git a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp index fa3d25b862e0..40adbf0e35ff 100644 --- a/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp +++ b/cpp/tests/io/experimental/hybrid_scan_filters_test.cpp @@ -1193,9 +1193,9 @@ TEST_F(HybridScanFiltersTest, FilterRowGroupsWithDictionary) template struct RowGroupFilteringWithDictTest : public HybridScanFiltersTest {}; -// Booleans are not supported for dictionary based filtering +// Booleans and fixed-point types are not supported for dictionary based filtering using DictionaryTestTypes = - cudf::test::RemoveIf>, SupportedTestTypes>; + cudf::test::RemoveIf>, SupportedTestTypesJIT>; TYPED_TEST_SUITE(RowGroupFilteringWithDictTest, DictionaryTestTypes); diff --git a/cpp/tests/io/parquet_common.hpp b/cpp/tests/io/parquet_common.hpp index ff698b9be70d..5f1a85aea697 100644 --- a/cpp/tests/io/parquet_common.hpp +++ b/cpp/tests/io/parquet_common.hpp @@ -45,9 +45,14 @@ using ByteLikeTypes = cudf::test::Types; -// Also fixed point types unsupported, because AST does not support them yet. -using SupportedTestTypes = cudf::test::RemoveIf, - cudf::test::ComparableTypes>; + +// Support types for AST expression evaluator +using SupportedTestTypesAST = + cudf::test::RemoveIf, ComparableAndFixedTypes>; + +// JIT does not yet support fixed point types +using SupportedTestTypesJIT = + cudf::test::RemoveIf, SupportedTestTypesAST>; // removing duration_D, duration_s, and timestamp_s as they don't appear to be supported properly. // see definition of UnsupportedChronoTypes above. diff --git a/cpp/tests/io/parquet_reader_test.cpp b/cpp/tests/io/parquet_reader_test.cpp index 3968dc656bab..f4c1dad3268e 100644 --- a/cpp/tests/io/parquet_reader_test.cpp +++ b/cpp/tests/io/parquet_reader_test.cpp @@ -3111,12 +3111,6 @@ TYPED_TEST(ParquetReaderSourceTest, BufferSourceArrayTypes) ////////////////////////////// // predicate pushdown tests -// Test for Types - numeric, chrono, string. -template -struct ParquetPredicatePushdownTest : public ParquetReaderTest {}; - -TYPED_TEST_SUITE(ParquetPredicatePushdownTest, SupportedTestTypes); - template void filter_typed_test() { @@ -3170,18 +3164,17 @@ void filter_typed_test() // Filtering AST auto literal_value = []() { if constexpr (cudf::is_timestamp()) { - // table[0] < 10000 timestamp days/seconds/milliseconds/microseconds/nanoseconds - return cudf::timestamp_scalar(T(typename T::duration(10000))); // i (0-20,000) + return cudf::timestamp_scalar(T(typename T::duration(10000))); // i ∈ [0, 20,000) } else if constexpr (cudf::is_duration()) { - // table[0] < 10000 day/seconds/milliseconds/microseconds/nanoseconds - return cudf::duration_scalar(T(10000)); // i (0-20,000) + return cudf::duration_scalar(T(10000)); // i ∈ [0, 20,000) } else if constexpr (std::is_same_v) { - // table[0] < "000010000" - return cudf::string_scalar("000010000"); // i (0-20,000) + return cudf::string_scalar("000010000"); // i ∈ [0-20,000) + } else if constexpr (cudf::is_fixed_point()) { + return cudf::fixed_point_scalar(typename T::rep{0}, + numeric::scale_type{0}); // i ∈ [-10,000, 10,000) } else { - // table[0] < 0 or 100u return cudf::numeric_scalar( - (100 - 100 * std::is_signed_v)); // i/100 (-100-100/ 0-200) + (100 - 100 * std::is_signed_v)); // i/100 ∈ [-100, 100) or [0, 200) } }(); @@ -3206,6 +3199,8 @@ void filter_typed_test() return cudf::duration_scalar(T(20000)); } else if constexpr (std::is_same_v) { return cudf::string_scalar("000020000"); + } else if constexpr (cudf::is_fixed_point()) { + return cudf::fixed_point_scalar(typename T::rep{20000}, numeric::scale_type{0}); } else { return cudf::numeric_scalar(std::numeric_limits::max()); } @@ -3353,18 +3348,17 @@ void filter_unary_operation_typed_test() // Filtering AST auto literal_value = []() { if constexpr (cudf::is_timestamp()) { - // table[0] < 10000 timestamp days/seconds/milliseconds/microseconds/nanoseconds - return cudf::timestamp_scalar(T(typename T::duration(10000))); // i (0-20,000) + return cudf::timestamp_scalar(T(typename T::duration(10000))); // i ∈ [0, 20,000) } else if constexpr (cudf::is_duration()) { - // table[0] < 10000 day/seconds/milliseconds/microseconds/nanoseconds - return cudf::duration_scalar(T(10000)); // i (0-20,000) + return cudf::duration_scalar(T(10000)); // i ∈ [0, 20,000) } else if constexpr (std::is_same_v) { - // table[0] < "000010000" - return cudf::string_scalar("000010000"); // i (0-20,000) + return cudf::string_scalar("000010000"); // i ∈ [0-20,000) + } else if constexpr (cudf::is_fixed_point()) { + return cudf::fixed_point_scalar(typename T::rep{0}, + numeric::scale_type{0}); // i ∈ [-10,000, 10,000) } else { - // table[0] < 0 or 100u return cudf::numeric_scalar( - (100 - 100 * std::is_signed_v)); // i/100 (-100-100/ 0-200) + (100 - 100 * std::is_signed_v)); // i/100 ∈ [-100, 100) or [0, 200) } }(); @@ -3402,18 +3396,205 @@ void filter_unary_operation_typed_test() } } -TYPED_TEST(ParquetPredicatePushdownTest, FilterTyped) +template +void decimal_stats_filter_test() +{ + using RepType = typename DecimalType::rep; + + auto constexpr num_input_row_groups = 3; + + auto const filepath = temp_env->get_temp_filepath("DecimalStatsFilter.parquet"); + + for (auto const scale : + {numeric::scale_type{-5}, numeric::scale_type{0}, numeric::scale_type{3}}) { + { + auto const rg0 = cudf::test::fixed_point_column_wrapper( + {RepType{100}, RepType{0}, RepType{200}}, {true, false, true}, scale); + auto const rg1 = cudf::test::fixed_point_column_wrapper( + {RepType{-50}, RepType{300}, RepType{0}}, {true, true, false}, scale); + auto const rg2 = + cudf::test::fixed_point_column_wrapper({RepType{-600}, RepType{-400}}, scale); + auto const t0 = cudf::table_view{{rg0}}; + auto const t1 = cudf::table_view{{rg1}}; + auto const t2 = cudf::table_view{{rg2}}; + + auto const options = + cudf::io::chunked_parquet_writer_options::builder(cudf::io::sink_info{filepath}) + .metadata(cudf::io::table_input_metadata(t0)) + .build(); + + cudf::io::chunked_parquet_writer writer(options); + writer.write(t0); + writer.write(t1); + writer.write(t2); + writer.close(); + } + + // Verify Parquet physical type for the decimal column is as expected + { + auto const meta = cudf::io::read_parquet_metadata(cudf::io::source_info{filepath}); + auto const& root = meta.schema().root(); + ASSERT_GE(root.num_children(), 1); + auto const& col_schema = root.child(0); + if constexpr (std::is_same_v) { + EXPECT_EQ(col_schema.type(), cudf::io::parquet::Type::INT32); + } else if constexpr (std::is_same_v) { + EXPECT_EQ(col_schema.type(), cudf::io::parquet::Type::INT64); + } else if constexpr (std::is_same_v) { + EXPECT_EQ(col_schema.type(), cudf::io::parquet::Type::FIXED_LEN_BYTE_ARRAY); + } + } + + // Helper function to test predicate pushdown for decimal types + auto const test_predicate_pushdown = [&](cudf::ast::operation const& filter, + cudf::size_type expected_filtered_row_groups, + cudf::size_type expected_num_rows) { + auto const options = + cudf::io::parquet_reader_options::builder(cudf::io::source_info{filepath}) + .filter(filter) + .build(); + + auto const result = cudf::io::read_parquet(options); + + EXPECT_EQ(result.metadata.num_input_row_groups, num_input_row_groups); + EXPECT_TRUE(result.metadata.num_row_groups_after_stats_filter.has_value()); + EXPECT_EQ(result.metadata.num_row_groups_after_stats_filter.value(), + expected_filtered_row_groups); + EXPECT_EQ(result.tbl->num_rows(), expected_num_rows); + }; + + auto const col_ref = cudf::ast::column_reference(0); + + // Filter: col0 >= 100 AND col0 <= 200 + { + auto scalar_100 = cudf::fixed_point_scalar(RepType{100}, scale); + auto scalar_200 = cudf::fixed_point_scalar(RepType{200}, scale); + auto const literal_100 = cudf::ast::literal(scalar_100); + auto const literal_200 = cudf::ast::literal(scalar_200); + auto const col_ge_100 = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal_100); + auto const col_le_200 = + cudf::ast::operation(cudf::ast::ast_operator::LESS_EQUAL, col_ref, literal_200); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, col_ge_100, col_le_200); + + // RGs 0, 1 pass + test_predicate_pushdown(filter, 2, 2); + } + + // Filter: col0 >= -550 AND col0 <= -450 + { + auto scalar_neg_550 = cudf::fixed_point_scalar(RepType{-550}, scale); + auto scalar_neg_450 = cudf::fixed_point_scalar(RepType{-450}, scale); + auto const literal_neg_550 = cudf::ast::literal(scalar_neg_550); + auto const literal_neg_450 = cudf::ast::literal(scalar_neg_450); + auto const col_ge_neg_550 = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal_neg_550); + auto const col_le_neg_450 = + cudf::ast::operation(cudf::ast::ast_operator::LESS_EQUAL, col_ref, literal_neg_450); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, col_ge_neg_550, col_le_neg_450); + + // RG 2 passes + test_predicate_pushdown(filter, 1, 0); + } + + // Filter: col0 >= 700 AND col0 <= 900 — matches no row groups + { + auto scalar_700 = cudf::fixed_point_scalar(RepType{700}, scale); + auto scalar_900 = cudf::fixed_point_scalar(RepType{900}, scale); + auto const literal_700 = cudf::ast::literal(scalar_700); + auto const literal_900 = cudf::ast::literal(scalar_900); + auto const col_ge_700 = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal_700); + auto const col_le_900 = + cudf::ast::operation(cudf::ast::ast_operator::LESS_EQUAL, col_ref, literal_900); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, col_ge_700, col_le_900); + + test_predicate_pushdown(filter, 0, 0); + } + + // Filter: col0 == -400 + { + auto scalar_neg_400 = cudf::fixed_point_scalar(RepType{-400}, scale); + auto const literal_neg_400 = cudf::ast::literal(scalar_neg_400); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::EQUAL, col_ref, literal_neg_400); + + // RG 2 passes + test_predicate_pushdown(filter, 1, 1); + } + + // Filter: col0 >= -100 AND col0 <= 250 + { + auto scalar_neg_100 = cudf::fixed_point_scalar(RepType{-100}, scale); + auto scalar_250 = cudf::fixed_point_scalar(RepType{250}, scale); + auto const literal_neg_100 = cudf::ast::literal(scalar_neg_100); + auto const literal_250 = cudf::ast::literal(scalar_250); + auto const col_ge = + cudf::ast::operation(cudf::ast::ast_operator::GREATER_EQUAL, col_ref, literal_neg_100); + auto const col_le = + cudf::ast::operation(cudf::ast::ast_operator::LESS_EQUAL, col_ref, literal_250); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_AND, col_ge, col_le); + + // RGs 0, 1 pass + test_predicate_pushdown(filter, 2, 3); + } + + // Filter: col0 == 100 OR col0 == -50 + { + auto scalar_100 = cudf::fixed_point_scalar(RepType{100}, scale); + auto scalar_neg_50 = cudf::fixed_point_scalar(RepType{-50}, scale); + auto const literal_100 = cudf::ast::literal(scalar_100); + auto const literal_neg_50 = cudf::ast::literal(scalar_neg_50); + auto const eq_100 = + cudf::ast::operation(cudf::ast::ast_operator::EQUAL, col_ref, literal_100); + auto const eq_neg_50 = + cudf::ast::operation(cudf::ast::ast_operator::EQUAL, col_ref, literal_neg_50); + auto const filter = + cudf::ast::operation(cudf::ast::ast_operator::LOGICAL_OR, eq_100, eq_neg_50); + + // RGs 0, 1 pass + test_predicate_pushdown(filter, 2, 2); + } + + // Large value filter only for decimal128 (overflows smaller rep types) + if constexpr (std::is_same_v) { + auto const big_val = (static_cast<__int128_t>(1) << 70) + 1234; + auto scalar_val = cudf::fixed_point_scalar(big_val, scale); + auto const lit = cudf::ast::literal(scalar_val); + auto const filter = cudf::ast::operation(cudf::ast::ast_operator::GREATER, col_ref, lit); + + // No RGs pass + test_predicate_pushdown(filter, 0, 0); + } + } +} + +template +struct ParquetPredicatePushdownTestAST : public ParquetReaderTest {}; +TYPED_TEST_SUITE(ParquetPredicatePushdownTestAST, SupportedTestTypesAST); + +TYPED_TEST(ParquetPredicatePushdownTestAST, FilterTyped) { filter_typed_test(); filter_unary_operation_typed_test(); + if constexpr (cudf::is_fixed_point()) { decimal_stats_filter_test(); } } -TYPED_TEST(ParquetPredicatePushdownTest, FilterTypedJIT) +template +struct ParquetPredicatePushdownTestJIT : public ParquetReaderTest {}; +TYPED_TEST_SUITE(ParquetPredicatePushdownTestJIT, SupportedTestTypesJIT); + +TYPED_TEST(ParquetPredicatePushdownTestJIT, FilterTyped) { filter_typed_test(); - // JIT does not support nullness-dependent operators such as IS_NULL so we can't call - // `filter_unary_operation_typed_test` - // Ref: https://github.com/rapidsai/cudf/issues/20177 + // JIT does not support decimals and nullness-dependent operators (IS_NULL) so we can't test: + // `filter_unary_operation_typed_test()` and `decimal_stats_filter_test()`. + // Refs: https://github.com/rapidsai/cudf/issues/20177 and + // https://github.com/rapidsai/cudf/issues/21584 } TEST_P(ParquetDecompressionTest, RoundTripBasic) diff --git a/cpp/tests/reductions/scan_tests.cpp b/cpp/tests/reductions/scan_tests.cpp index a2ea07a6225d..64493c818f3f 100644 --- a/cpp/tests/reductions/scan_tests.cpp +++ b/cpp/tests/reductions/scan_tests.cpp @@ -473,7 +473,72 @@ TEST_F(ScanStringTest, MinMaxInclusiveWithNullsExclude) CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_max, result_max->view()); } -TEST_F(ScanStringTest, ExclusiveThrows) +// ============== Scan Count ============== +struct ScanCountTest : public cudf::test::BaseFixture {}; + +TEST_F(ScanCountTest, InclusiveNoNulls) +{ + cudf::test::fixed_width_column_wrapper col({5, 4, 6, 0, 1, 6, 5, 3}); + cudf::test::fixed_width_column_wrapper expected({1, 2, 3, 4, 5, 6, 7, 8}); + + auto result = cudf::scan( + col, *cudf::make_count_aggregation(), cudf::scan_type::INCLUSIVE); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); + + result = + cudf::scan(col, + *cudf::make_count_aggregation(cudf::null_policy::INCLUDE), + cudf::scan_type::INCLUSIVE); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); +} + +TEST_F(ScanCountTest, InclusiveWithNullsExclude) +{ + cudf::test::strings_column_wrapper col({"5", "4", "6", "", "1", "6", "5", "3"}, + {1, 1, 1, 0, 1, 1, 1, 1}); + cudf::test::fixed_width_column_wrapper expected({1, 2, 3, 4, 4, 5, 6, 7}, + {1, 1, 1, 0, 1, 1, 1, 1}); + + auto result = cudf::scan(col, + *cudf::make_count_aggregation(), + cudf::scan_type::INCLUSIVE, + cudf::null_policy::EXCLUDE); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); +} + +TEST_F(ScanCountTest, InclusiveWithNullsInclude) +{ + cudf::test::strings_column_wrapper col({"5", "4", "6", "", "1", "6", "5", "3"}, + {1, 1, 1, 0, 1, 1, 1, 1}); + cudf::test::fixed_width_column_wrapper expected({1, 2, 3, 0, 0, 0, 0, 0}, + {1, 1, 1, 0, 0, 0, 0, 0}); + + auto result = + cudf::scan(col, + *cudf::make_count_aggregation(cudf::null_policy::INCLUDE), + cudf::scan_type::INCLUSIVE, + cudf::null_policy::INCLUDE); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); +} + +TEST_F(ScanCountTest, InclusiveWithOffset) +{ + cudf::test::strings_column_wrapper col({"5", "4", "6", "", "1", "6", "5", "3"}, + {1, 1, 1, 0, 1, 1, 1, 1}); + auto input = cudf::slice(col, {1, 7}).front(); + cudf::test::fixed_width_column_wrapper expected({1, 2, 3, 3, 4, 5}, + {1, 1, 0, 1, 1, 1}); + + auto result = cudf::scan(input, + *cudf::make_count_aggregation(), + cudf::scan_type::INCLUSIVE, + cudf::null_policy::EXCLUDE); + CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected, result->view()); +} + +struct ScanExclusiveErrorTest : public cudf::test::BaseFixture {}; + +TEST_F(ScanExclusiveErrorTest, ExclusiveThrows) { cudf::test::strings_column_wrapper col({"a", "b", "c"}); EXPECT_THROW( @@ -484,6 +549,10 @@ TEST_F(ScanStringTest, ExclusiveThrows) cudf::scan( col, *cudf::make_max_aggregation(), cudf::scan_type::EXCLUSIVE), cudf::logic_error); + EXPECT_THROW( + cudf::scan( + col, *cudf::make_count_aggregation(), cudf::scan_type::EXCLUSIVE), + cudf::logic_error); } // ============== Chrono MinMax ============== diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py index 92d1b32f1cf3..ba8b85f7ecfb 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py @@ -290,6 +290,11 @@ class PDSDSPolarsQueries(PDSDSQueries): pl.col("store_qty").cast(pl.Int64), pl.col("other_chan_qty").cast(pl.Int64), ], + 83: [ + pl.col("sr_item_qty").cast(pl.Int64), + pl.col("cr_item_qty").cast(pl.Int64), + pl.col("wr_item_qty").cast(pl.Int64), + ], 94: [pl.col("order count").cast(COUNT_DTYPE)], 95: [pl.col("order count").cast(COUNT_DTYPE)], 96: [pl.col("count_star()").cast(COUNT_DTYPE)], diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py index dc7e3103e5d2..05858eef99c1 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py @@ -11,7 +11,7 @@ import polars as pl from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import get_data +from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: from cudf_polars.experimental.benchmarks.utils import RunConfig @@ -116,12 +116,23 @@ def q83_segment( ) .join(dates, left_on=returned_date_key, right_on="d_date_sk") .group_by("i_item_id") - .agg(pl.col(qty_col).sum().alias(out_qty_name)) + .agg( + [ + pl.col(qty_col).count().alias(f"{out_qty_name}_count"), + pl.col(qty_col).sum().alias(f"{out_qty_name}_sum"), + ] + ) + .with_columns( + pl.when(pl.col(f"{out_qty_name}_count") > 0) + .then(pl.col(f"{out_qty_name}_sum")) + .otherwise(None) + .alias(out_qty_name) + ) .select(["i_item_id", out_qty_name]) ) -def polars_impl(run_config: RunConfig) -> pl.LazyFrame: +def polars_impl(run_config: RunConfig) -> QueryResult: """Query 83.""" params = load_parameters( int(run_config.scale_factor), @@ -180,40 +191,48 @@ def polars_impl(run_config: RunConfig) -> pl.LazyFrame: out_qty_name="wr_item_qty", ) - return ( - sr_items.join(cr_items, on="i_item_id") - .join(wr_items, on="i_item_id") - .with_columns( - ( - pl.col("sr_item_qty") + pl.col("cr_item_qty") + pl.col("wr_item_qty") - ).alias("total_qty") - ) - .with_columns( - [ - (pl.col("total_qty") / 3.0).cast(pl.Float64).alias("average"), - (pl.col("sr_item_qty") / pl.col("total_qty") / 3.0 * 100) - .cast(pl.Float64) - .alias("sr_dev"), - (pl.col("cr_item_qty") / pl.col("total_qty") / 3.0 * 100) - .cast(pl.Float64) - .alias("cr_dev"), - (pl.col("wr_item_qty") / pl.col("total_qty") / 3.0 * 100) - .cast(pl.Float64) - .alias("wr_dev"), - ] - ) - .select( - [ - pl.col("i_item_id").alias("item_id"), - "sr_item_qty", - "sr_dev", - "cr_item_qty", - "cr_dev", - "wr_item_qty", - "wr_dev", - "average", - ] - ) - .sort(["item_id", "sr_item_qty"]) - .limit(100) + sort_by = {"item_id": False, "sr_item_qty": False} + limit = 100 + return QueryResult( + frame=( + sr_items.join(cr_items, on="i_item_id") + .join(wr_items, on="i_item_id") + .with_columns( + ( + pl.col("sr_item_qty") + + pl.col("cr_item_qty") + + pl.col("wr_item_qty") + ).alias("total_qty") + ) + .with_columns( + [ + (pl.col("total_qty") / 3.0).cast(pl.Float64).alias("average"), + (pl.col("sr_item_qty") / pl.col("total_qty") / 3.0 * 100) + .cast(pl.Float64) + .alias("sr_dev"), + (pl.col("cr_item_qty") / pl.col("total_qty") / 3.0 * 100) + .cast(pl.Float64) + .alias("cr_dev"), + (pl.col("wr_item_qty") / pl.col("total_qty") / 3.0 * 100) + .cast(pl.Float64) + .alias("wr_dev"), + ] + ) + .select( + [ + pl.col("i_item_id").alias("item_id"), + "sr_item_qty", + "sr_dev", + "cr_item_qty", + "cr_dev", + "wr_item_qty", + "wr_dev", + "average", + ] + ) + .sort(["item_id", "sr_item_qty"], nulls_last=True) + .limit(limit) + ), + sort_by=list(sort_by.items()), + limit=limit, ) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py index 3e641798eed1..e8fdaba78420 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py @@ -10,7 +10,7 @@ import polars as pl from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import get_data +from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: from cudf_polars.experimental.benchmarks.utils import RunConfig @@ -40,7 +40,7 @@ def duckdb_impl(run_config: RunConfig) -> str: WHERE ca_city = '{city}' AND c_current_addr_sk = ca_address_sk AND ib_lower_bound >= {income} - AND ib_upper_bound <= 54986 + 50000 + AND ib_upper_bound <= {income} + 50000 AND ib_income_band_sk = hd_income_band_sk AND cd_demo_sk = c_current_cdemo_sk AND hd_demo_sk = c_current_hdemo_sk @@ -50,7 +50,7 @@ def duckdb_impl(run_config: RunConfig) -> str: """ -def polars_impl(run_config: RunConfig) -> pl.LazyFrame: +def polars_impl(run_config: RunConfig) -> QueryResult: """Query 84.""" params = load_parameters( int(run_config.scale_factor), @@ -75,48 +75,52 @@ def polars_impl(run_config: RunConfig) -> pl.LazyFrame: store_returns = get_data( run_config.dataset_path, "store_returns", run_config.suffix ) - return ( - customer.join( - customer_address.filter(pl.col("ca_city") == city), - left_on="c_current_addr_sk", - right_on="ca_address_sk", - how="inner", - ) - .join( - customer_demographics, - left_on="c_current_cdemo_sk", - right_on="cd_demo_sk", - how="inner", - ) - .join( - household_demographics, - left_on="c_current_hdemo_sk", - right_on="hd_demo_sk", - how="inner", - ) - .join( - income_band.filter( - (pl.col("ib_lower_bound") >= income) - & (pl.col("ib_upper_bound") <= income + 50000) - ), - left_on="hd_income_band_sk", - right_on="ib_income_band_sk", - how="inner", - ) - .join( - store_returns, - left_on="c_current_cdemo_sk", - right_on="sr_cdemo_sk", - how="inner", - ) - .select( - [ - pl.col("c_customer_id").alias("customer_id"), - (pl.col("c_last_name") + pl.lit(", ") + pl.col("c_first_name")).alias( - "customername" + return QueryResult( + frame=( + customer.join( + customer_address.filter(pl.col("ca_city") == city), + left_on="c_current_addr_sk", + right_on="ca_address_sk", + how="inner", + ) + .join( + customer_demographics, + left_on="c_current_cdemo_sk", + right_on="cd_demo_sk", + how="inner", + ) + .join( + household_demographics, + left_on="c_current_hdemo_sk", + right_on="hd_demo_sk", + how="inner", + ) + .join( + income_band.filter( + (pl.col("ib_lower_bound") >= income) + & (pl.col("ib_upper_bound") <= income + 50000) ), - ] - ) - .sort("customer_id", nulls_last=True) - .limit(100) + left_on="hd_income_band_sk", + right_on="ib_income_band_sk", + how="inner", + ) + .join( + store_returns, + left_on="c_current_cdemo_sk", + right_on="sr_cdemo_sk", + how="inner", + ) + .select( + [ + pl.col("c_customer_id").alias("customer_id"), + ( + pl.col("c_last_name") + pl.lit(", ") + pl.col("c_first_name") + ).alias("customername"), + ] + ) + .sort("customer_id", nulls_last=True) + .limit(100) + ), + sort_by=[("customer_id", False)], + limit=100, ) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py index 707de394fde8..9fb980d15b11 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py @@ -10,7 +10,7 @@ import polars as pl from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import get_data +from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: from cudf_polars.experimental.benchmarks.utils import RunConfig @@ -86,7 +86,7 @@ def duckdb_impl(run_config: RunConfig) -> str: """ -def polars_impl(run_config: RunConfig) -> pl.LazyFrame: +def polars_impl(run_config: RunConfig) -> QueryResult: """Query 85.""" params = load_parameters( int(run_config.scale_factor), @@ -113,126 +113,132 @@ def polars_impl(run_config: RunConfig) -> pl.LazyFrame: ) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) reason = get_data(run_config.dataset_path, "reason", run_config.suffix) - return ( - web_sales.join( - web_returns, - left_on=["ws_item_sk", "ws_order_number"], - right_on=["wr_item_sk", "wr_order_number"], - how="inner", - ) - .join( - web_page, left_on="ws_web_page_sk", right_on="wp_web_page_sk", how="inner" - ) - .join( - date_dim.filter(pl.col("d_year") == year), - left_on="ws_sold_date_sk", - right_on="d_date_sk", - how="inner", - ) - .join( - customer_demographics.select( - [ - pl.col("cd_demo_sk").alias("cd1_demo_sk"), - pl.col("cd_marital_status").alias("cd1_marital_status"), - pl.col("cd_education_status").alias("cd1_education_status"), - ] - ), - left_on="wr_refunded_cdemo_sk", - right_on="cd1_demo_sk", - how="inner", - ) - .join( - customer_demographics.select( - [ - pl.col("cd_demo_sk").alias("cd2_demo_sk"), - pl.col("cd_marital_status").alias("cd2_marital_status"), - pl.col("cd_education_status").alias("cd2_education_status"), - ] - ), - left_on="wr_returning_cdemo_sk", - right_on="cd2_demo_sk", - how="inner", - ) - .join( - customer_address, - left_on="wr_refunded_addr_sk", - right_on="ca_address_sk", - how="inner", - ) - .join(reason, left_on="wr_reason_sk", right_on="r_reason_sk", how="inner") - .filter( - ( - (pl.col("cd1_marital_status") == ms[0]) - & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) - & (pl.col("cd1_education_status") == es[0]) - & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) - & ( - pl.col("ws_sales_price").is_between( - price_ranges[0][0], price_ranges[0][1] + sort_by = { + "substr(r_reason_desc, 1, 20)": False, + "avg(ws_quantity)": False, + "avg(wr_refunded_cash)": False, + "avg(wr_fee)": False, + } + limit = 100 + return QueryResult( + frame=( + web_sales.join( + web_returns, + left_on=["ws_item_sk", "ws_order_number"], + right_on=["wr_item_sk", "wr_order_number"], + how="inner", + ) + .join( + web_page, + left_on="ws_web_page_sk", + right_on="wp_web_page_sk", + how="inner", + ) + .join( + date_dim.filter(pl.col("d_year") == year), + left_on="ws_sold_date_sk", + right_on="d_date_sk", + how="inner", + ) + .join( + customer_demographics.select( + [ + pl.col("cd_demo_sk").alias("cd1_demo_sk"), + pl.col("cd_marital_status").alias("cd1_marital_status"), + pl.col("cd_education_status").alias("cd1_education_status"), + ] + ), + left_on="wr_refunded_cdemo_sk", + right_on="cd1_demo_sk", + how="inner", + ) + .join( + customer_demographics.select( + [ + pl.col("cd_demo_sk").alias("cd2_demo_sk"), + pl.col("cd_marital_status").alias("cd2_marital_status"), + pl.col("cd_education_status").alias("cd2_education_status"), + ] + ), + left_on="wr_returning_cdemo_sk", + right_on="cd2_demo_sk", + how="inner", + ) + .join( + customer_address, + left_on="wr_refunded_addr_sk", + right_on="ca_address_sk", + how="inner", + ) + .join(reason, left_on="wr_reason_sk", right_on="r_reason_sk", how="inner") + .filter( + ( + (pl.col("cd1_marital_status") == ms[0]) + & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) + & (pl.col("cd1_education_status") == es[0]) + & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) + & ( + pl.col("ws_sales_price").is_between( + price_ranges[0][0], price_ranges[0][1] + ) ) ) - ) - | ( - (pl.col("cd1_marital_status") == ms[1]) - & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) - & (pl.col("cd1_education_status") == es[1]) - & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) - & ( - pl.col("ws_sales_price").is_between( - price_ranges[1][0], price_ranges[1][1] + | ( + (pl.col("cd1_marital_status") == ms[1]) + & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) + & (pl.col("cd1_education_status") == es[1]) + & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) + & ( + pl.col("ws_sales_price").is_between( + price_ranges[1][0], price_ranges[1][1] + ) ) ) - ) - | ( - (pl.col("cd1_marital_status") == ms[2]) - & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) - & (pl.col("cd1_education_status") == es[2]) - & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) - & ( - pl.col("ws_sales_price").is_between( - price_ranges[2][0], price_ranges[2][1] + | ( + (pl.col("cd1_marital_status") == ms[2]) + & (pl.col("cd1_marital_status") == pl.col("cd2_marital_status")) + & (pl.col("cd1_education_status") == es[2]) + & (pl.col("cd1_education_status") == pl.col("cd2_education_status")) + & ( + pl.col("ws_sales_price").is_between( + price_ranges[2][0], price_ranges[2][1] + ) ) ) ) - ) - .filter( - ( - (pl.col("ca_country") == "United States") - & (pl.col("ca_state").is_in(states[0:3])) - & (pl.col("ws_net_profit").is_between(np_min, np_max)) + .filter( + ( + (pl.col("ca_country") == "United States") + & (pl.col("ca_state").is_in(states[0:3])) + & (pl.col("ws_net_profit").is_between(np_min, np_max)) + ) + | ( + (pl.col("ca_country") == "United States") + & (pl.col("ca_state").is_in(states[3:6])) + & (pl.col("ws_net_profit").is_between(np_min, np_max)) + ) ) - | ( - (pl.col("ca_country") == "United States") - & (pl.col("ca_state").is_in(states[3:6])) - & (pl.col("ws_net_profit").is_between(np_min, np_max)) + .group_by("r_reason_desc") + .agg( + [ + pl.col("ws_quantity").mean().alias("avg(ws_quantity)"), + pl.col("wr_refunded_cash").mean().alias("avg(wr_refunded_cash)"), + pl.col("wr_fee").mean().alias("avg(wr_fee)"), + ] + ) + .select( + [ + pl.col("r_reason_desc") + .str.slice(0, 20) + .alias("substr(r_reason_desc, 1, 20)"), + "avg(ws_quantity)", + "avg(wr_refunded_cash)", + "avg(wr_fee)", + ] ) - ) - .group_by("r_reason_desc") - .agg( - [ - pl.col("ws_quantity").mean().alias("avg(ws_quantity)"), - pl.col("wr_refunded_cash").mean().alias("avg(wr_refunded_cash)"), - pl.col("wr_fee").mean().alias("avg(wr_fee)"), - ] - ) - .select( - [ - pl.col("r_reason_desc") - .str.slice(0, 20) - .alias("substr(r_reason_desc, 1, 20)"), - "avg(ws_quantity)", - "avg(wr_refunded_cash)", - "avg(wr_fee)", - ] - ) - .sort( - [ - "substr(r_reason_desc, 1, 20)", - "avg(ws_quantity)", - "avg(wr_refunded_cash)", - "avg(wr_fee)", - ], - nulls_last=True, - ) - .limit(100) + .sort(list(sort_by.keys()), nulls_last=True) + .limit(limit) + ), + sort_by=list(sort_by.items()), + limit=limit, ) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py index 4e80d49d0a01..0ff9ec60e9b8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py @@ -33,6 +33,8 @@ import rmm.statistics +from cudf_polars.experimental.rapidsmpf.spmd import spmd_execution + # The dtype for count() aggregations depends on the presence # of the polars-runtime-64 package (`polars[rt64]`). HAS_POLARS_RT_64 = pl.config.plr.RUNTIME_REPR == "rt64" @@ -440,7 +442,13 @@ def from_args(cls, args: argparse.Namespace) -> RunConfig: ) cluster = "single" if scheduler == "synchronous" else "distributed" elif cluster is not None: - scheduler = "synchronous" if cluster == "single" else "distributed" + match cluster: + case "single": + scheduler = "synchronous" + case "distributed": + scheduler = "distributed" + case "spmd": # launched via rrun, not Dask + scheduler = None else: cluster = "single" scheduler = "synchronous" @@ -957,11 +965,12 @@ def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: "--cluster", default=None, type=str, - choices=["single", "distributed"], + choices=["single", "distributed", "spmd"], help=textwrap.dedent("""\ Cluster type to use with the 'streaming' executor. - single : Run locally in a single process - - distributed : Use Dask for multi-GPU execution"""), + - distributed : Use Dask for multi-GPU execution + - spmd : SPMD execution via rrun launcher"""), ) parser.add_argument( "-s", @@ -1439,10 +1448,14 @@ def run_polars_query_iteration( expected: pl.DataFrame | None, query_result: Any, client: Any, + prepare_validation_result: Callable[[pl.DataFrame], pl.DataFrame] | None = None, ) -> SuccessRecord: """Run a single query iteration. Caller must wrap in try/except.""" result, duration = execute_query(q_id, iteration, q, run_config, args, engine) + if expected is not None and prepare_validation_result is not None: + result = prepare_validation_result(result) + if run_config.shuffle == "rapidsmpf" and run_config.gather_shuffle_stats: from rapidsmpf.integrations.dask.shuffler import ( clear_shuffle_statistics, @@ -1493,6 +1506,7 @@ def run_polars_query( numeric_type: str, date_type: str, validation_files: dict[int, Path] | None, + prepare_validation_result: Callable[[pl.DataFrame], pl.DataFrame] | None = None, ) -> QueryRunResult: """Run all iterations for a single query. Caller must wrap in try/except.""" query_result = getattr(benchmark, f"q{q_id}")(run_config) @@ -1562,6 +1576,7 @@ def run_polars_query( expected=expected, query_result=query_result, client=client, + prepare_validation_result=prepare_validation_result, ) except Exception: print(f"❌ query={q_id} iteration={i} failed!") @@ -1599,53 +1614,27 @@ def run_polars_query( ) -def run_polars( +def _run_query_loop( benchmark: Any, args: argparse.Namespace, -) -> None: - """Run the queries using the given benchmark and executor options.""" - vars(args).update({"query_set": benchmark.name}) - run_config = RunConfig.from_args(args) - validation_failures: list[int] = [] - query_failures: list[tuple[int, int]] = [] - - client = initialize_dask_cluster(run_config, args) - - # Update n_workers from the actual cluster when using scheduler file/address - if client is not None: - actual_n_workers = client.scheduler_info()["n_workers"] - run_config = dataclasses.replace(run_config, n_workers=actual_n_workers) - + run_config: RunConfig, + engine: pl.GPUEngine | None, + client: Any, + numeric_type: str, + date_type: str, + validation_files: dict[int, Path] | None, + prepare_validation_result: Callable[[pl.DataFrame], pl.DataFrame] | None = None, +) -> tuple[ + defaultdict[int, list[SuccessRecord | FailedRecord]], + dict[int, Any], + list[int], + list[tuple[int, int]], +]: + """Execute all queries in ``run_config`` and return accumulated results.""" records: defaultdict[int, list[SuccessRecord | FailedRecord]] = defaultdict(list) plans: dict[int, SerializablePlan] = {} - engine: pl.GPUEngine | None = None - numeric_type, date_type = check_input_data_type(run_config) - - if args.validate_directory is not None: - validation_files = list_validation_files(args.validate_directory) - else: - validation_files = None - - if run_config.executor != "cpu": - executor_options = get_executor_options(run_config, benchmark=benchmark) - if run_config.runtime == "rapidsmpf": - parquet_options = { - "use_rapidsmpf_native": run_config.native_parquet, - } - else: - parquet_options = {} - engine = pl.GPUEngine( - raise_on_fail=True, - memory_resource=rmm.mr.CudaAsyncMemoryResource( - release_threshold=args.rmm_release_threshold - ) - if run_config.rmm_async - else None, - cuda_stream_policy=run_config.stream_policy, - executor=run_config.executor, - executor_options=executor_options, - parquet_options=parquet_options, - ) + validation_failures: list[int] = [] + query_failures: list[tuple[int, int]] = [] for q_id in run_config.queries: try: @@ -1659,6 +1648,7 @@ def run_polars( numeric_type=numeric_type, date_type=date_type, validation_files=validation_files, + prepare_validation_result=prepare_validation_result, ) except Exception: print(f"❌ query={q_id} failed (setup or execution)!") @@ -1683,65 +1673,156 @@ def run_polars( if result.validation_failed: validation_failures.append(q_id) - run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) + return records, plans, validation_failures, query_failures - # consolidate logs - if _HAS_STRUCTLOG and run_config.collect_traces: - def gather_logs() -> str: - logger = logging.getLogger() - return logger.handlers[0].stream.getvalue() # type: ignore[attr-defined] +def _consolidate_logs(run_config: RunConfig, client: Any) -> RunConfig: + """Merge structlog traces from the local process and Dask workers into run_config.""" + if not (_HAS_STRUCTLOG and run_config.collect_traces): + return run_config - if client is not None: - # Gather logs from both client (for Query Plan) and workers - worker_logs = "\n".join(client.run(gather_logs).values()) - client_logs = gather_logs() - all_logs = client_logs + "\n" + worker_logs - else: - all_logs = gather_logs() + def gather_logs() -> str: + logger = logging.getLogger() + return logger.handlers[0].stream.getvalue() # type: ignore[attr-defined] - parsed_logs = [json.loads(log) for log in all_logs.splitlines() if log] - # Some other log records can end up in here. Filter those out. - scope_values = {s.value for s in Scope} - parsed_logs = [log for log in parsed_logs if log.get("scope") in scope_values] - # Now we want to augment the existing Records with the trace data. + if client is not None: + # Gather logs from both client (for Query Plan) and workers + worker_logs = "\n".join(client.run(gather_logs).values()) + client_logs = gather_logs() + all_logs = client_logs + "\n" + worker_logs + else: + all_logs = gather_logs() - def group_key(x: dict) -> int: - return x["query_id"] + parsed_logs = [json.loads(log) for log in all_logs.splitlines() if log] + # Some other log records can end up in here. Filter those out. + scope_values = {s.value for s in Scope} + parsed_logs = [log for log in parsed_logs if log.get("scope") in scope_values] + # Now we want to augment the existing Records with the trace data. - def sort_key(x: dict) -> tuple[int, int]: - return x["query_id"], x["iteration"] + def group_key(x: dict) -> int: + return x["query_id"] - grouped = itertools.groupby( - sorted(parsed_logs, key=sort_key), - key=group_key, - ) + def sort_key(x: dict) -> tuple[int, int]: + return x["query_id"], x["iteration"] - for query_id, run_logs_group in grouped: - run_logs = list(run_logs_group) - by_iteration = [ - list(x) - for _, x in itertools.groupby(run_logs, key=lambda x: x["iteration"]) - ] - run_records = run_config.records[query_id] - assert len(by_iteration) == len(run_records) # same number of iterations - all_traces = [list(iteration) for iteration in by_iteration] - - new_records: list[SuccessRecord | FailedRecord] = [] - for rec, traces in zip(run_records, all_traces, strict=True): - if rec.status == "success": - new_records.append(dataclasses.replace(rec, traces=traces)) - else: - new_records.append(rec) + grouped = itertools.groupby( + sorted(parsed_logs, key=sort_key), + key=group_key, + ) + + for query_id, run_logs_group in grouped: + run_logs = list(run_logs_group) + by_iteration = [ + list(x) + for _, x in itertools.groupby(run_logs, key=lambda x: x["iteration"]) + ] + run_records = run_config.records[query_id] + assert len(by_iteration) == len(run_records) # same number of iterations + all_traces = [list(iteration) for iteration in by_iteration] + + new_records: list[SuccessRecord | FailedRecord] = [] + for rec, traces in zip(run_records, all_traces, strict=True): + if rec.status == "success": + new_records.append(dataclasses.replace(rec, traces=traces)) + else: + new_records.append(rec) - run_config.records[query_id] = new_records + run_config.records[query_id] = new_records - if args.summarize: - run_config.summarize() + return run_config + + +def run_polars( + benchmark: Any, + args: argparse.Namespace, +) -> None: + """Run the queries using the given benchmark and executor options.""" + vars(args).update({"query_set": benchmark.name}) + run_config = RunConfig.from_args(args) + numeric_type, date_type = check_input_data_type(run_config) + validation_files = ( + list_validation_files(args.validate_directory) + if args.validate_directory is not None + else None + ) + parquet_options = ( + {"use_rapidsmpf_native": run_config.native_parquet} + if run_config.runtime == "rapidsmpf" + else {} + ) + match run_config.cluster: + case "spmd": + run_polars_spmd( + benchmark, + args, + run_config, + parquet_options, + numeric_type, + date_type, + validation_files, + ) + case "single" | "distributed": + run_polars_single_or_dask( + benchmark, + args, + run_config, + parquet_options, + numeric_type, + date_type, + validation_files, + ) + +def run_polars_single_or_dask( + benchmark: Any, + args: argparse.Namespace, + run_config: RunConfig, + parquet_options: dict[str, Any], + numeric_type: str, + date_type: str, + validation_files: dict[int, Path] | None, +) -> None: + """Run benchmark queries using Dask or single-process execution.""" + client = initialize_dask_cluster(run_config, args) + if client is not None: + run_config = dataclasses.replace( + run_config, n_workers=client.scheduler_info()["n_workers"] + ) + + engine = None + if run_config.executor != "cpu": + executor_options = get_executor_options(run_config, benchmark=benchmark) + engine = pl.GPUEngine( + raise_on_fail=True, + memory_resource=rmm.mr.CudaAsyncMemoryResource( + release_threshold=args.rmm_release_threshold + ) + if run_config.rmm_async + else None, + cuda_stream_policy=run_config.stream_policy, + executor=run_config.executor, + executor_options=executor_options, + parquet_options=parquet_options, + ) + + records, plans, validation_failures, query_failures = _run_query_loop( + benchmark, + args, + run_config, + engine, + client, + numeric_type, + date_type, + validation_files, + ) + run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) + run_config = _consolidate_logs(run_config, client=client) if client is not None: client.close(timeout=60) + if args.summarize: + run_config.summarize() + if args.validate and run_config.executor != "cpu": print("\nValidation Summary") print("==================") @@ -1755,8 +1836,79 @@ def sort_key(x: dict) -> tuple[int, int]: args.output.write(json.dumps(run_config.serialize(engine=engine))) args.output.write("\n") - exit_code = 1 if (query_failures or validation_failures) else 0 - sys.exit(exit_code) + sys.exit(1 if (query_failures or validation_failures) else 0) + + +def run_polars_spmd( + benchmark: Any, + args: argparse.Namespace, + run_config: RunConfig, + parquet_options: dict[str, Any], + numeric_type: str, + date_type: str, + validation_files: dict[int, Path] | None, +) -> None: + """Run benchmark queries using SPMD execution via the ``rrun`` launcher.""" + if run_config.collect_traces: + raise NotImplementedError( + "--collect-traces is not yet supported with --cluster spmd." + ) + executor_options = get_executor_options(run_config, benchmark=benchmark) + # "runtime" and "cluster" are reserved — spmd_execution sets them + executor_options.pop("runtime", None) + executor_options.pop("cluster", None) + with spmd_execution( + mr=rmm.mr.CudaAsyncMemoryResource(release_threshold=args.rmm_release_threshold) + if run_config.rmm_async + else None, + executor_options=executor_options, + parquet_options=parquet_options, + cuda_stream_policy=run_config.stream_policy, + ) as (comm, ctx, engine): + from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id + from cudf_polars.experimental.rapidsmpf.spmd import allgather_polars_dataframe + + def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: + with reserve_op_id() as op_id: + return allgather_polars_dataframe( + comm=comm, + ctx=ctx, + local_df=df, + op_id=op_id, + ) + + rank = comm.rank + run_config = dataclasses.replace(run_config, n_workers=comm.nranks) + records, plans, validation_failures, query_failures = _run_query_loop( + benchmark, + args, + run_config, + engine, + None, + numeric_type, + date_type, + validation_files, + prepare_validation_result=_allgather_result, + ) + run_config = dataclasses.replace(run_config, records=dict(records), plans=plans) + # Only rank 0 writes output and prints summaries to avoid N duplicate outputs. + if rank == 0: + if args.summarize: + run_config.summarize() + if args.validate and run_config.executor != "cpu": + print("\nValidation Summary") + print("==================") + if validation_failures: + print( + f"{len(validation_failures)} queries failed validation: " + f"{sorted(set(validation_failures))}" + ) + else: + print("✅ All validated queries passed.") + # engine is not JSON-serializable (holds the SPMD Cython context) + args.output.write(json.dumps(run_config.serialize(engine=None))) + args.output.write("\n") + sys.exit(1 if (query_failures or validation_failures) else 0) def setup_logging(query_id: int, iteration: int) -> None: # noqa: D103