diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f542bc4b6dc8..ca14107a023d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -881,12 +881,8 @@ add_library( src/join/key_remapping.cu src/join/mark_join.cu src/join/mixed_join.cu - src/join/mixed_join_kernel.cu - src/join/mixed_join_kernel_nulls.cu src/join/mixed_join_kernels_semi.cu src/join/mixed_join_semi.cu - src/join/mixed_join_size_kernel.cu - src/join/mixed_join_size_kernel_nulls.cu src/join/sort_merge_join.cu src/json/json_path.cu src/lists/contains.cu @@ -1420,7 +1416,7 @@ if(CUDF_BUILD_TESTUTIL) ) target_link_libraries( - cudftestutil INTERFACE $ Threads::Threads cudf + cudftestutil INTERFACE $ Threads::Threads cudf rmm::rmm cudftest_default_stream $ ) diff --git a/cpp/include/cudf/detail/join/join.hpp b/cpp/include/cudf/detail/join/join.hpp index bf9fd5d42def..92d43eb055a7 100644 --- a/cpp/include/cudf/detail/join/join.hpp +++ b/cpp/include/cudf/detail/join/join.hpp @@ -1,15 +1,51 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include +#include +#include #include +#include + +#include +#include +#include + +#include +#include +#include +#include namespace cudf { namespace detail { constexpr int DEFAULT_JOIN_CG_SIZE = 2; +/** + * @brief Internal `filter_join_indices` accepting a precomputed output size. + * + * Same semantics as `cudf::filter_join_indices`. When `output_size` is provided it is used directly + * to size the output, skipping the internal size-counting pass. The value must equal the size that + * the function would otherwise compute (for example the result of `filter_join_indices_output_size` + * for the same inputs); behavior is undefined otherwise. + * + * @param output_size Optional precomputed number of output rows; computed internally if not + * provided + */ +std::pair>, + std::unique_ptr>> +filter_join_indices(table_view const& left, + table_view const& right, + device_span left_indices, + device_span right_indices, + ast::expression const& predicate, + join_kind join_kind, + std::optional output_size, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr); + } // namespace detail } // namespace cudf diff --git a/cpp/include/cudf/join/join.hpp b/cpp/include/cudf/join/join.hpp index a0171b1d4833..333bd69b0912 100644 --- a/cpp/include/cudf/join/join.hpp +++ b/cpp/include/cudf/join/join.hpp @@ -17,7 +17,11 @@ #include +#include #include +#include +#include +#include /** * @file @@ -340,6 +344,9 @@ std::unique_ptr cross_join( * @param right_indices Device span of row indices in the right table from hash join. * @param predicate An AST expression that returns a boolean for each pair of rows. * @param join_kind The type of join operation. Must be INNER_JOIN, LEFT_JOIN, or FULL_JOIN. + * @param output_size Optional precomputed number of output rows. When provided, skips the internal + * size-counting pass. Behavior is undefined if it differs from the size the function would + * otherwise produce for the same inputs. * @param stream CUDA stream used for kernel launches and memory operations. * @param mr Device memory resource used to allocate output indices. * @@ -354,7 +361,37 @@ filter_join_indices(cudf::table_view const& left, cudf::device_span right_indices, cudf::ast::expression const& predicate, cudf::join_kind join_kind, - rmm::cuda_stream_view stream = cudf::get_default_stream(), + std::optional output_size = std::nullopt, + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); + +/** + * @brief Filters join result indices based on a conditional predicate and join type. + * + * @deprecated Use the overload that accepts an optional output size instead. + * + * @param left The left table for predicate evaluation (conditional columns only). + * @param right The right table for predicate evaluation (conditional columns only). + * @param left_indices Device span of row indices in the left table from hash join. + * @param right_indices Device span of row indices in the right table from hash join. + * @param predicate An AST expression that returns a boolean for each pair of rows. + * @param join_kind The type of join operation. Must be INNER_JOIN, LEFT_JOIN, or FULL_JOIN. + * @param stream CUDA stream used for kernel launches and memory operations. + * @param mr Device memory resource used to allocate output indices. + * + * @return A pair of device vectors [filtered_left_indices, filtered_right_indices] + * corresponding to rows that satisfy the join semantics and predicate. + */ +[[deprecated("Use the overload that takes an optional output_size parameter.")]] +std::pair>, + std::unique_ptr>> +filter_join_indices(cudf::table_view const& left, + cudf::table_view const& right, + cudf::device_span left_indices, + cudf::device_span right_indices, + cudf::ast::expression const& predicate, + cudf::join_kind join_kind, + rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** @@ -362,15 +399,19 @@ filter_join_indices(cudf::table_view const& left, * the filtered index vectors. * * Runs the same predicate evaluation as `filter_join_indices` but skips the index - * materialization step, returning only the total number of pairs that would be - * emitted. The semantics per `join_kind` match `filter_join_indices`: - * - INNER_JOIN: number of pairs where the predicate evaluates to true. - * - LEFT_JOIN: predicate-passing pairs plus one entry per left row with no passing match. - * - FULL_JOIN: input pairs plus one extra entry per pair whose predicate failed - * (because failed matches split into `(left, JoinNoMatch)` and `(JoinNoMatch, right)`). - * - * The returned size may be passed as a precomputed hint to APIs that compose - * `filter_join_indices` (for example, the mixed join APIs). + * materialization step, returning the total number of pairs that would be emitted along with the + * per-output contribution counts whose sum is that total. The counts are laid out per `join_kind` + * so that each entry records how many output rows the corresponding input contributes: + * - INNER_JOIN: indexed per input pair; entry `i` is `1` if the predicate passes and `0` otherwise. + * - FULL_JOIN: indexed per input pair; entry `i` is `1` for a preserved pair (predicate passes or + * the pair already contains a `JoinNoMatch`) and `2` for a failed valid pair (which splits into + * `(left, JoinNoMatch)` and `(JoinNoMatch, right)`). + * - LEFT_JOIN: indexed per left row; each entry holds the number of passing pairs for that left + * row, floored to `1` to account for the synthetic `(left, JoinNoMatch)` entry. + * + * The returned size and contribution counts may be passed as a precomputed hint to APIs that + * compose `filter_join_indices` (for example, the mixed join APIs). The layout above is an + * implementation detail that callers should treat as opaque rather than rely upon. * * @throw std::invalid_argument if `join_kind` is not INNER_JOIN, LEFT_JOIN, or FULL_JOIN. * @throw std::invalid_argument if `left_indices` and `right_indices` have different sizes. @@ -383,17 +424,21 @@ filter_join_indices(cudf::table_view const& left, * @param predicate An AST expression that returns a boolean for each pair of rows. * @param join_kind The type of join operation. Must be INNER_JOIN, LEFT_JOIN, or FULL_JOIN. * @param stream CUDA stream used for kernel launches and memory operations. + * @param mr Device memory resource used to allocate the returned contribution counts. * - * @return The exact number of pairs that `filter_join_indices` would produce. + * @return A pair containing the exact number of pairs that `filter_join_indices` would produce + * and the per-output contribution counts that sum to that number. */ -[[nodiscard]] std::size_t filter_join_indices_output_size( +[[nodiscard]] std::pair>> +filter_join_indices_output_size( cudf::table_view const& left, cudf::table_view const& right, cudf::device_span left_indices, cudf::device_span right_indices, cudf::ast::expression const& predicate, cudf::join_kind join_kind, - rmm::cuda_stream_view stream = cudf::get_default_stream()); + rmm::cuda_stream_view stream = cudf::get_default_stream(), + rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref()); /** * @brief JIT-based filtering of join result indices using string predicate. diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index b119ebc3b2e5..3fcc8f5ddaf1 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp @@ -19,6 +19,10 @@ namespace cudf_streaming::detail { /** * @brief A bloom filter, used for approximate set membership queries. + * + * @note All methods of this class launch work on the streams provided. It is the caller's + * responsibility to ensure that data are valid to read/write on that stream as + * appropriate. */ struct device_bloom_filter { /** @@ -28,12 +32,8 @@ struct device_bloom_filter { * @param seed Seed used for hashing each value. * @param storage Storage to view as a bloom filter, must be appropriately * initialized. - * @param stream CUDA stream for device operations. */ - device_bloom_filter(std::size_t num_blocks, - std::uint64_t seed, - void* storage, - rmm::cuda_stream_view stream); + device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage); /** * @brief Create a read-only filter. @@ -41,14 +41,12 @@ struct device_bloom_filter { * @param num_blocks Number of blocks in the filter. * @param seed Seed used for hashing each value. * @param storage View of storage, must be appropriately initialized. - * @param stream CUDA stream for device operations. * * @return A const-qualified bloom filter viewing the underlying storage. */ static device_bloom_filter const view(std::size_t num_blocks, std::uint64_t seed, - void const* storage, - rmm::cuda_stream_view stream); + void const* storage); /** * @brief Create uninitialized storage for a filter. @@ -98,11 +96,6 @@ struct device_bloom_filter { rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) const; - /** - * @brief @return The stream the underlying storage is valid on. - */ - [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; - /** * @brief @return Pointer to the underlying storage. */ @@ -127,10 +120,9 @@ struct device_bloom_filter { [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept; private: - std::size_t num_blocks_; ///< Number of blocks used in the filter. - std::uint64_t seed_; ///< Seed used when hashing values. - void* storage_; ///< Backing storage. - rmm::cuda_stream_view stream_; ///< Stream storage is valid on. + std::size_t num_blocks_; ///< Number of blocks used in the filter. + std::uint64_t seed_; ///< Seed used when hashing values. + void* storage_; ///< Backing storage. }; } // namespace cudf_streaming::detail diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 54dbfcf52cea..78acf8431f05 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/cpp/libcudf_streaming/src/bloom_filter.cpp @@ -37,8 +37,8 @@ rapidsmpf::streaming::Actor bloom_filter::build( auto storage = cudf_streaming::detail::device_bloom_filter::storage(num_filter_blocks_, filter_stream, mr); RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream)); - auto filter = cudf_streaming::detail::device_bloom_filter( - num_filter_blocks_, seed_, storage->data(), filter_stream); + auto filter = + cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage->data()); rapidsmpf::CudaEvent build_event; build_event.record(filter_stream); while (!ch_out->is_shutdown()) { @@ -48,6 +48,14 @@ rapidsmpf::streaming::Actor bloom_filter::build( chunk = co_await chunk.make_available( ctx_, -rapidsmpf::safe_cast(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE))); + + // Reservation for the hash values in add. + auto res = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE) + ->reserve_or_wait(rapidsmpf::safe_cast(chunk.table_view().num_rows()) + // TODO: no magic numbers: the hashing algorithm in + // `add` below returns an int64 column. + * sizeof(std::int64_t), + 0); // Filter is allocated on `filter_stream`, but we run the additions on the chunk's // stream. The addition modifies global memory but we can safely launch two // kernels doing that concurrently because the updates are atomic. @@ -67,9 +75,9 @@ rapidsmpf::streaming::Actor bloom_filter::build( [num_blocks = num_filter_blocks_, seed = seed_](rapidsmpf::Buffer const* left, rapidsmpf::Buffer* right) { right->write_access([&](std::byte* out_bytes, rmm::cuda_stream_view stream) { - auto const in = cudf_streaming::detail::device_bloom_filter::view( - num_blocks, seed, left->data(), stream); - cudf_streaming::detail::device_bloom_filter(num_blocks, seed, out_bytes, stream) + auto const in = + cudf_streaming::detail::device_bloom_filter::view(num_blocks, seed, left->data()); + cudf_streaming::detail::device_bloom_filter(num_blocks, seed, out_bytes) .merge(in, stream); }); }); @@ -95,7 +103,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply( auto stream = storage.stream(); rapidsmpf::CudaEvent event; auto filter = - cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage.data(), stream); + cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage.data()); auto meta = co_await ch_in->receive_metadata(); if (!meta.empty()) { co_await ch_out->send_metadata(std::move(meta)); } while (!ch_out->is_shutdown()) { diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 155cc522d320..0525f139babf 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -30,7 +30,6 @@ #include #include -#include #include @@ -72,15 +71,9 @@ using StorageType = BloomFilterRefType::filter_block_type; } // namespace -device_bloom_filter::device_bloom_filter(std::size_t num_blocks, - std::uint64_t seed, - void* storage, - rmm::cuda_stream_view stream) - : num_blocks_{num_blocks}, seed_{seed}, storage_{storage}, stream_{stream} +device_bloom_filter::device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage) + : num_blocks_{num_blocks}, seed_{seed}, storage_{storage} { - // TODO: use an aligned allocator adaptor to ensure this holds. - // Today all RMM device allocators guarantee at least 256 byte alignment, but that is - // an implementation detail. RAPIDSMPF_EXPECTS( reinterpret_cast(storage_) % std::alignment_of_v == 0, "Allocation for bloom filter is not aligned."); @@ -88,19 +81,19 @@ device_bloom_filter::device_bloom_filter(std::size_t num_blocks, device_bloom_filter const device_bloom_filter::view(std::size_t num_blocks, std::uint64_t seed, - void const* storage, - rmm::cuda_stream_view stream) + void const* storage) { // const-cast is safe because the returned object is also const and therefore can't // call methods that throw away constness. - return device_bloom_filter(num_blocks, seed, const_cast(storage), stream); + return device_bloom_filter(num_blocks, seed, const_cast(storage)); } std::unique_ptr device_bloom_filter::storage(std::size_t num_blocks, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - return std::make_unique(num_blocks * sizeof(StorageType), stream, mr); + return std::make_unique( + num_blocks * sizeof(StorageType), std::alignment_of_v, stream, mr); } void device_bloom_filter::add(cudf::table_view const& values_to_hash, @@ -110,8 +103,7 @@ void device_bloom_filter::add(cudf::table_view const& values_to_hash, RAPIDSMPF_NVTX_FUNC_RANGE(); auto filter_ref = BloomFilterRefType{ static_cast(storage_), num_blocks_, cuco::thread_scope_device, {}}; - auto hashes = cudf::hashing::xxhash_64( - values_to_hash, seed_, stream, cudf::get_current_device_resource_ref()); + auto hashes = cudf::hashing::xxhash_64(values_to_hash, seed_, stream, mr); auto hash_view = hashes->view(); RAPIDSMPF_EXPECTS(hash_view.type().id() == cudf::type_to_id(), "Hash values do not have correct type"); @@ -136,9 +128,8 @@ rmm::device_uvector device_bloom_filter::contains(cudf::table_view const& RAPIDSMPF_NVTX_FUNC_RANGE(); auto filter_ref = BloomFilterRefType{ static_cast(storage_), num_blocks_, cuco::thread_scope_device, {}}; - auto hashes = - cudf::hashing::xxhash_64(values, seed_, stream, cudf::get_current_device_resource_ref()); - auto view = hashes->view(); + auto hashes = cudf::hashing::xxhash_64(values, seed_, stream, mr); + auto view = hashes->view(); rmm::device_uvector result{static_cast(view.size()), stream, mr}; filter_ref.contains_async(view.begin(), view.end(), result.begin(), stream); return result; @@ -149,8 +140,6 @@ std::size_t device_bloom_filter::fitting_num_blocks(std::size_t l2size) noexcept return (l2size * 2) / (3 * sizeof(StorageType)); } -rmm::cuda_stream_view device_bloom_filter::stream() const noexcept { return stream_; } - void* device_bloom_filter::data() noexcept { return storage_; } void const* device_bloom_filter::data() const noexcept { return storage_; } diff --git a/cpp/src/join/filter_join_indices/filter_join_indices.cu b/cpp/src/join/filter_join_indices/filter_join_indices.cu index 868e3a4ecec7..bbc7b2e64d73 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices.cu @@ -12,7 +12,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -38,8 +40,11 @@ #include #include #include +#include +#include #include +#include #include namespace cudf { @@ -53,6 +58,7 @@ filter_join_indices(cudf::table_view const& left, cudf::device_span right_indices, ast::expression const& predicate, join_kind join_kind, + std::optional output_size, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { @@ -168,10 +174,13 @@ filter_join_indices(cudf::table_view const& left, auto valid_predicate = [=] __device__(size_type i) -> bool { return predicate_results_ptr[i]; }; auto const num_valid = - cudf::detail::count_if(cuda::counting_iterator{0}, - cuda::counting_iterator{static_cast(left_indices.size())}, - valid_predicate, - stream); + output_size.has_value() + ? *output_size + : cudf::detail::count_if( + cuda::counting_iterator{0}, + cuda::counting_iterator{static_cast(left_indices.size())}, + valid_predicate, + stream); if (num_valid == 0) { return make_empty_result(); } @@ -224,12 +233,12 @@ filter_join_indices(cudf::table_view const& left, auto const num_invalid = left.num_rows() - num_filter_passing; - // Find the number of indices passing the filter i.e. rows that are valid according to the - // predicate CUB APIs are used instead of Thrust to enable 64-bit operations on index vectors of - // size greater than integer limits - cudf::detail::device_scalar d_num_valid(stream, - cudf::get_current_device_resource_ref()); - { + auto const num_valid = [&]() -> std::size_t { + if (output_size.has_value()) { return *output_size - num_invalid; } + // CUB APIs are used instead of Thrust to enable 64-bit operations on index vectors of size + // greater than integer limits + cudf::detail::device_scalar d_num_valid(stream, + cudf::get_current_device_resource_ref()); auto const predicate_it = cuda::transform_iterator{predicate_results_ptr, cuda::proclaim_return_type( @@ -248,12 +257,12 @@ filter_join_indices(cudf::table_view const& left, d_num_valid.data(), left_indices.size(), stream.value()); - } - auto const num_valid = d_num_valid.value(stream); - auto const output_size = num_valid + num_invalid; - if (output_size == 0) { return make_empty_result(); } + return d_num_valid.value(stream); + }(); + auto const result_size = num_valid + num_invalid; + if (result_size == 0) { return make_empty_result(); } - auto [filtered_left_indices, filtered_right_indices] = make_result_vectors(output_size); + auto [filtered_left_indices, filtered_right_indices] = make_result_vectors(result_size); if (num_valid > 0) { auto input_iter = thrust::make_zip_iterator(cuda::std::tuple{left_indices.begin(), right_indices.begin()}); @@ -307,15 +316,18 @@ filter_join_indices(cudf::table_view const& left, // Count failed matches for output sizing auto const failed_matched_count = - cudf::detail::count_if(cuda::counting_iterator{0}, - cuda::counting_iterator{static_cast(left_indices.size())}, - is_failed_matched_pair, - stream); - auto const output_size = left_indices.size() + failed_matched_count; + output_size.has_value() + ? *output_size - left_indices.size() + : cudf::detail::count_if( + cuda::counting_iterator{0}, + cuda::counting_iterator{static_cast(left_indices.size())}, + is_failed_matched_pair, + stream); + auto const result_size = left_indices.size() + failed_matched_count; - if (output_size == 0) { return make_empty_result(); } + if (result_size == 0) { return make_empty_result(); } - auto [filtered_left_indices, filtered_right_indices] = make_result_vectors(output_size); + auto [filtered_left_indices, filtered_right_indices] = make_result_vectors(result_size); // Use two-step approach with optimized memory management // Step 1: Handle primary pairs @@ -361,13 +373,15 @@ filter_join_indices(cudf::table_view const& left, } } -std::size_t filter_join_indices_output_size(cudf::table_view const& left, - cudf::table_view const& right, - cudf::device_span left_indices, - cudf::device_span right_indices, - ast::expression const& predicate, - join_kind join_kind, - rmm::cuda_stream_view stream) +std::pair>> +filter_join_indices_output_size(cudf::table_view const& left, + cudf::table_view const& right, + cudf::device_span left_indices, + cudf::device_span right_indices, + ast::expression const& predicate, + join_kind join_kind, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { // Validate inputs (same constraints as filter_join_indices) CUDF_EXPECTS(left_indices.size() == right_indices.size(), @@ -379,8 +393,12 @@ std::size_t filter_join_indices_output_size(cudf::table_view const& left, "filter_join_indices_output_size only supports INNER_JOIN, LEFT_JOIN, and FULL_JOIN.", std::invalid_argument); - if (left_indices.empty()) { return 0; } - if (join_kind == join_kind::LEFT_JOIN && left.num_rows() == 0) { return 0; } + auto empty_counts = [&]() { + return std::make_unique>(0, stream, mr); + }; + + if (left_indices.empty()) { return {0, empty_counts()}; } + if (join_kind == join_kind::LEFT_JOIN && left.num_rows() == 0) { return {0, empty_counts()}; } auto const has_nulls = predicate.may_evaluate_null(left, right, stream); @@ -399,16 +417,13 @@ std::size_t filter_join_indices_output_size(cudf::table_view const& left, detail::grid_1d const config(left_indices.size(), DEFAULT_JOIN_BLOCK_SIZE); auto const shmem_per_block = parser.shmem_per_thread * DEFAULT_JOIN_BLOCK_SIZE; - // The count kernel uses a single atomic counter. Allocate device_scalar zero-initialized. - cudf::detail::device_scalar d_count( - std::size_t{0}, stream, cudf::get_current_device_resource_ref()); - - // For LEFT_JOIN, allocate a zeroed per-left-row mark buffer; for others, pass nullptr. - auto left_passing_marks = cudf::detail::make_zeroed_device_uvector_async( - join_kind == join_kind::LEFT_JOIN ? static_cast(left.num_rows()) : 0, - stream, - cudf::get_current_device_resource_ref()); - auto* const marks_ptr = join_kind == join_kind::LEFT_JOIN ? left_passing_marks.data() : nullptr; + auto const counts_size = join_kind == join_kind::LEFT_JOIN + ? static_cast(left.num_rows()) + : left_indices.size(); + auto output_counts = + join_kind == join_kind::LEFT_JOIN + ? cudf::detail::make_zeroed_device_uvector_async(counts_size, stream, mr) + : rmm::device_uvector(counts_size, stream, mr); cudf::detail::dispatch_bool(has_nulls, [&](auto has_nulls_c) { cudf::detail::dispatch_bool(has_complex_type, [&](auto has_complex_c) { @@ -422,25 +437,27 @@ std::size_t filter_join_indices_output_size(cudf::table_view const& left, config, shmem_per_block, join_kind, - d_count.data(), - marks_ptr, + output_counts.data(), stream); }); }); - auto const num_predicate_passing = d_count.value(stream); - - switch (join_kind) { - case join_kind::INNER_JOIN: return num_predicate_passing; - case join_kind::FULL_JOIN: return left_indices.size() + num_predicate_passing; - case join_kind::LEFT_JOIN: { - auto const num_filter_passing = cudf::detail::count_if( - left_passing_marks.begin(), left_passing_marks.end(), cuda::std::identity{}, stream); - auto const num_invalid = static_cast(left.num_rows()) - num_filter_passing; - return num_predicate_passing + num_invalid; - } - default: CUDF_FAIL("Unsupported join kind for filter_join_indices_output_size"); + if (join_kind == join_kind::LEFT_JOIN) { + thrust::transform(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + output_counts.begin(), + output_counts.end(), + output_counts.begin(), + cuda::proclaim_return_type( + [] __device__(size_type count) { return count > 0 ? count : 1; })); } + + std::size_t const total = + thrust::reduce(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + output_counts.begin(), + output_counts.end(), + std::size_t{0}); + + return {total, std::make_unique>(std::move(output_counts))}; } } // namespace detail @@ -454,25 +471,44 @@ filter_join_indices(cudf::table_view const& left, cudf::device_span right_indices, ast::expression const& predicate, cudf::join_kind join_kind, + std::optional output_size, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); return detail::filter_join_indices( - left, right, left_indices, right_indices, predicate, join_kind, stream, mr); + left, right, left_indices, right_indices, predicate, join_kind, output_size, stream, mr); +} + +std::pair>, + std::unique_ptr>> +filter_join_indices(cudf::table_view const& left, + cudf::table_view const& right, + cudf::device_span left_indices, + cudf::device_span right_indices, + ast::expression const& predicate, + cudf::join_kind join_kind, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) +{ + CUDF_FUNC_RANGE(); + return detail::filter_join_indices( + left, right, left_indices, right_indices, predicate, join_kind, std::nullopt, stream, mr); } -std::size_t filter_join_indices_output_size(cudf::table_view const& left, - cudf::table_view const& right, - cudf::device_span left_indices, - cudf::device_span right_indices, - ast::expression const& predicate, - cudf::join_kind join_kind, - rmm::cuda_stream_view stream) +std::pair>> +filter_join_indices_output_size(cudf::table_view const& left, + cudf::table_view const& right, + cudf::device_span left_indices, + cudf::device_span right_indices, + ast::expression const& predicate, + cudf::join_kind join_kind, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { CUDF_FUNC_RANGE(); return detail::filter_join_indices_output_size( - left, right, left_indices, right_indices, predicate, join_kind, stream); + left, right, left_indices, right_indices, predicate, join_kind, stream, mr); } } // namespace cudf diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.cuh b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.cuh index e158e3883662..f835add8b016 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.cuh +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -18,23 +18,25 @@ #include -#include #include -#include - -#include namespace cudf::detail { /** - * @brief Counts the per-join-kind output size of `filter_join_indices` without materializing - * a per-pair boolean buffer. + * @brief Fills the per-output contribution counts of `filter_join_indices` without materializing + * the filtered index vectors. * - * Each thread accumulates a private partial count, the block aggregates with CUB, and each - * block adds its block-sum to `*count_out` exactly once via `cuda::atomic_ref`. For LEFT_JOIN, - * `left_passing_marks[left_row_index]` is additionally set to `true` for every left row that - * contributes to the count, which lets the host derive the number of synthetic JoinNoMatch - * entries. + * The total output size is the sum of `output_counts`, which is laid out per join kind so that + * each entry records how many output rows the corresponding input contributes: + * - INNER_JOIN: `output_counts` is indexed per input pair; entry `i` is `1` if the predicate + * passes and `0` otherwise. + * - FULL_JOIN: `output_counts` is indexed per input pair; entry `i` is `1` for a preserved pair + * (predicate passes or the pair already contains a `JoinNoMatch`) and `2` for a failed valid + * pair (which splits into `(left, JoinNoMatch)` and `(JoinNoMatch, right)`). + * - LEFT_JOIN: `output_counts` is indexed per left row; the kernel atomically accumulates the + * number of passing pairs for each left row. Left rows with no passing pair are floored to `1` + * by the host afterwards to account for the synthetic `(left, JoinNoMatch)` entry. The buffer + * must be zero-initialized before the launch. */ template CUDF_KERNEL __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) void filter_join_indices_output_size_kernel( @@ -44,8 +46,7 @@ CUDF_KERNEL __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) void filter_join_indices_ cudf::device_span right_indices, cudf::ast::detail::expression_device_view device_expression_data, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks) + cudf::size_type* output_counts) { extern __shared__ char raw_intermediate_storage[]; auto* intermediate_storage = @@ -53,17 +54,12 @@ CUDF_KERNEL __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) void filter_join_indices_ auto thread_intermediate_storage = &intermediate_storage[threadIdx.x * device_expression_data.num_intermediates]; - using BlockReduce = cub::BlockReduce; - __shared__ typename BlockReduce::TempStorage temp_storage; - auto const tid = cudf::detail::grid_1d::global_thread_id(); auto const stride = cudf::detail::grid_1d::grid_stride(); auto evaluator = cudf::ast::detail::expression_evaluator{ left_table, right_table, device_expression_data}; - cuda::std::size_t thread_local_count = 0; - for (auto i = tid; i < static_cast(left_indices.size()); i += stride) { auto const left_row_index = left_indices[i]; auto const right_row_index = right_indices[i]; @@ -85,35 +81,20 @@ CUDF_KERNEL __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) void filter_join_indices_ } switch (join_kind) { - case cudf::join_kind::INNER_JOIN: - if (predicate_pass) { ++thread_local_count; } + case cudf::join_kind::INNER_JOIN: output_counts[i] = predicate_pass ? 1 : 0; break; + case cudf::join_kind::FULL_JOIN: + output_counts[i] = (both_valid && !predicate_pass) ? 2 : 1; break; case cudf::join_kind::LEFT_JOIN: - if (predicate_pass) { - ++thread_local_count; - // Mark the left row as "passing" so the host can derive how many left rows need a - // synthetic JoinNoMatch entry. For matched-passing pairs and for pre-existing - // (left, JoinNoMatch) entries from upstream hash_join.left_join the left index is a - // valid row index in [0, left_table.num_rows()). - if (left_row_index >= 0 && left_row_index < left_table.num_rows()) { - left_passing_marks[left_row_index] = true; - } + if (predicate_pass && left_row_index >= 0 && left_row_index < left_table.num_rows()) { + cuda::atomic_ref count_ref{ + output_counts[left_row_index]}; + count_ref.fetch_add(1, cuda::memory_order_relaxed); } break; - case cudf::join_kind::FULL_JOIN: - // Count failed matches: predicate false AND both indices valid. - if (both_valid && !predicate_pass) { ++thread_local_count; } - break; default: break; } } - - cuda::std::size_t const block_sum = BlockReduce(temp_storage).Sum(thread_local_count); - - if (threadIdx.x == 0) { - cuda::atomic_ref count_ref{*count_out}; - count_ref.fetch_add(block_sum, cuda::memory_order_relaxed); - } } template @@ -126,8 +107,7 @@ void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream) { filter_join_indices_output_size_kernel @@ -138,8 +118,7 @@ void launch_filter_output_size_kernel( right_indices, device_expression_data, join_kind, - count_out, - left_passing_marks); + output_counts); CUDF_CUDA_TRY(cudaGetLastError()); } diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.hpp b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.hpp index b5b25e5e7a0e..b343694e88af 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.hpp +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -20,19 +20,15 @@ namespace cudf::detail { /** - * @brief Launches a kernel that counts the per-join-kind output size for `filter_join_indices`. + * @brief Launches a kernel that fills the per-output contribution counts for `filter_join_indices`. * - * For INNER_JOIN this is the number of pairs whose predicate evaluates to true. - * For LEFT_JOIN this is the number of input pairs whose predicate evaluates to true - * (including pre-existing unmatched pairs that are preserved); additionally, - * `left_passing_marks[left_row_index]` is set to `true` for every left row that - * contributes to that count (used by the host code to derive the number of left - * rows that need a synthetic JoinNoMatch entry). - * For FULL_JOIN this is the number of failed matched pairs (predicate false and - * both indices valid), which is added on top of `left_indices.size()` host-side. - * - * The kernel avoids materializing a per-pair boolean buffer; it folds the count - * directly into `count_out` via atomic increments. + * The total output size is the sum of `output_counts`. Its layout depends on the join kind: + * - INNER_JOIN: per input pair, `1` if the predicate passes and `0` otherwise. + * - FULL_JOIN: per input pair, `1` for a preserved pair and `2` for a failed valid pair (which + * splits into `(left, JoinNoMatch)` and `(JoinNoMatch, right)`). + * - LEFT_JOIN: per left row, the number of passing pairs (accumulated atomically). The host floors + * empty rows to `1` afterwards to account for the synthetic `(left, JoinNoMatch)` entry, so the + * buffer must be zero-initialized before the launch. * * @tparam has_nulls Indicates whether the expression may evaluate to null * @tparam has_complex_type Indicates whether the expression may contain complex types @@ -45,10 +41,9 @@ namespace cudf::detail { * @param[in] config Grid configuration for kernel launch * @param[in] shmem_per_block Amount of shared memory to allocate per block * @param[in] join_kind The join kind. Must be INNER_JOIN, LEFT_JOIN, or FULL_JOIN. - * @param[out] count_out Atomic counter for the per-kind count described above - * @param[out] left_passing_marks Byte buffer of size `left_table.num_rows()` used by LEFT_JOIN - * to mark left rows whose entries contribute to `count_out`. Must be zero-initialized - * before the kernel launch and may be `nullptr` for INNER_JOIN and FULL_JOIN. + * @param[out] output_counts Per-output contribution counts described above. Sized to + * `left_indices.size()` for INNER_JOIN and FULL_JOIN, and to `left_table.num_rows()` + * (zero-initialized) for LEFT_JOIN. * @param[in] stream CUDA stream on which to launch the kernel */ template @@ -61,8 +56,7 @@ void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream); } // namespace cudf::detail diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_complex.cu b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_complex.cu index 3bb93635d553..b55a7e33d152 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_complex.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_complex.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,7 +16,6 @@ template void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream); } // namespace cudf::detail diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_complex.cu b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_complex.cu index 195feaeeb65a..794b6ae55b48 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_complex.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_complex.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,7 +16,6 @@ template void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream); } // namespace cudf::detail diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_primitive.cu b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_primitive.cu index babeb76a3f82..84358e6fe9d7 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_primitive.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_null_primitive.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,7 +16,6 @@ template void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream); } // namespace cudf::detail diff --git a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_primitive.cu b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_primitive.cu index f695652818fb..f2de1322e301 100644 --- a/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_primitive.cu +++ b/cpp/src/join/filter_join_indices/filter_join_indices_output_size_kernel_primitive.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,7 +16,6 @@ template void launch_filter_output_size_kernel( cudf::detail::grid_1d const& config, std::size_t shmem_per_block, cudf::join_kind join_kind, - std::size_t* count_out, - bool* left_passing_marks, + cudf::size_type* output_counts, rmm::cuda_stream_view stream); } // namespace cudf::detail diff --git a/cpp/src/join/mixed_filter_join_common_utils.cuh b/cpp/src/join/mixed_filter_join_common_utils.cuh index d80638c81be8..829b5584a943 100644 --- a/cpp/src/join/mixed_filter_join_common_utils.cuh +++ b/cpp/src/join/mixed_filter_join_common_utils.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -7,9 +7,14 @@ #include "mixed_join_common_utils.cuh" #include +#include + +#include #include +#include + namespace cudf::detail { /** diff --git a/cpp/src/join/mixed_join.cu b/cpp/src/join/mixed_join.cu index 79d359c7b01a..baeb35670a36 100644 --- a/cpp/src/join/mixed_join.cu +++ b/cpp/src/join/mixed_join.cu @@ -1,532 +1,128 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -#include "join_common_utils.cuh" #include "join_common_utils.hpp" -#include "mixed_join_common_utils.cuh" -#include "mixed_join_kernel.hpp" -#include "mixed_join_size_kernel.hpp" -#include #include -#include -#include +#include #include -#include +#include #include #include -#include #include #include #include +#include #include #include +#include #include -#include -#include -#include -#include -#include +#include #include #include -#include +#include namespace cudf { namespace detail { namespace { -/** - * @brief Builds the hash table based on the given `build_table`. - * - * @tparam HashTable The type of the hash table - * - * @param build Table of columns used to build join hash. - * @param preprocessed_build shared_ptr to cudf::detail::row::equality::preprocessed_table - * for build - * @param hash_table Build hash table. - * @param has_nested_nulls Flag to denote if build or probe tables have nested nulls - * @param nulls_equal Flag to denote nulls are equal or not. - * @param bitmask Bitmask to denote whether a row is valid. - * @param stream CUDA stream used for device memory operations and kernel launches. - */ -template -void build_join_hash_table( - cudf::table_view const& build, - std::shared_ptr const& preprocessed_build, - HashTable& hash_table, - bool has_nested_nulls, - null_equality nulls_equal, - [[maybe_unused]] bitmask_type const* bitmask, - rmm::cuda_stream_view stream) -{ - CUDF_EXPECTS(0 != build.num_columns(), "Selected build dataset is empty", std::invalid_argument); - CUDF_EXPECTS(0 != build.num_rows(), "Build side table has no rows", std::invalid_argument); - - auto insert_rows = [&](auto const& build, auto const& d_hasher) { - auto const iter = cudf::detail::make_counting_transform_iterator(0, pair_fn{d_hasher}); - - if (nulls_equal == cudf::null_equality::EQUAL or not nullable(build)) { - hash_table.insert_async(iter, iter + build.num_rows(), stream.value()); - } else { - auto const stencil = cuda::counting_iterator{0}; - auto const pred = row_is_valid{bitmask}; - - hash_table.insert_if_async(iter, iter + build.num_rows(), stencil, pred, stream.value()); - } - }; - - auto const nulls = nullate::DYNAMIC{has_nested_nulls}; - - auto const row_hash = detail::row::hash::row_hasher{preprocessed_build}; - auto const d_hasher = row_hash.device_hasher(nulls); - - insert_rows(build, d_hasher); -} - -/** - * @brief Precomputes double hashing indices and row hash values for mixed join operations. - * - * This function exists as a performance optimization to work around the register spilling issue - * reported in https://github.com/NVIDIA/cuCollections/issues/761. The new cuco hash table - * implementation suffers from register spilling due to longer register live ranges, which can - * cause up to 20x performance degradation. - * - * By precomputing the double hashing indices (initial slot and step size) and row hash values - * in a separate pass, we reduce register pressure in the subsequent count and retrieve kernels. - * This approach yields approximately 20% speedup compared to the legacy multimap-based - * implementation. - * - * The tradeoff is that we cannot use cuco's device APIs directly in mixed join operations. - * Instead, we must reimplement the entire double hashing probing logic in cudf without relying - * on cuco's device APIs. This should be revisited and potentially removed once issue #761 is - * fully resolved. - * - * @param hash_table The cuco multiset hash table - * @param hash_probe Hash function for computing row hashes - * @param probe_table_num_rows Number of rows in the probe table - * @param stream CUDA stream for operations - * @param mr Memory resource for allocations - * @return A pair of device vectors: (input_pairs, hash_indices) where input_pairs contains - * (row_hash, row_index) pairs and hash_indices contains (initial_slot, step_size) pairs - */ -template -std::pair>, - rmm::device_uvector>> -precompute_mixed_join_data(mixed_multiset_type const& hash_table, - HashProbe const& hash_probe, - size_type probe_table_num_rows, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - auto input_pairs = - rmm::device_uvector>(probe_table_num_rows, stream, mr); - auto hash_indices = rmm::device_uvector>( - probe_table_num_rows, stream, mr); - - auto const capacity = hash_table.capacity(); - auto const probe_hash_fn = hash_table.hash_function(); - static constexpr std::size_t bucket_size = mixed_multiset_type::bucket_size; - - auto const num_buckets = capacity / bucket_size; - auto const num_buckets_minus_one = num_buckets - 1; - - // Functor to pre-compute both input pairs and initial slots and step sizes for double hashing. - auto precompute_fn = [=] __device__(size_type i) { - auto const probe_key = cuco::pair{hash_probe(i), i}; - - // Use the probing scheme's hash functions for proper double hashing - auto const hash1_val = cuda::std::get<0>(probe_hash_fn)(probe_key); - auto const hash2_val = cuda::std::get<1>(probe_hash_fn)(probe_key); - - auto const init_idx = static_cast( - (static_cast(hash1_val) % num_buckets) * bucket_size); - auto const step_val = static_cast( - ((static_cast(hash2_val) % num_buckets_minus_one) + 1) * bucket_size); - - return cuda::std::pair{probe_key, cuda::std::pair{init_idx, step_val}}; - }; - - // Single transform to fill both arrays using zip iterator - thrust::transform( - rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - cuda::counting_iterator{0}, - cuda::counting_iterator{probe_table_num_rows}, - thrust::make_zip_iterator(cuda::std::make_tuple(input_pairs.begin(), hash_indices.begin())), - precompute_fn); - - return std::make_pair(std::move(input_pairs), std::move(hash_indices)); -} - -struct mixed_join_setup_data { - bool swap_tables; - size_type outer_num_rows; - cudf::nullate::DYNAMIC has_nulls; - ast::detail::expression_parser parser; - mixed_multiset_type hash_table; - std::shared_ptr preprocessed_build; - std::shared_ptr preprocessed_probe; - std::unique_ptr> left_conditional_view; - std::unique_ptr> - right_conditional_view; - detail::grid_1d config; - thread_index_type shmem_size_per_block; - row_equality equality_probe; - cudf::device_span> hash_table_storage; - rmm::device_uvector> input_pairs; - rmm::device_uvector> hash_indices; -}; - -mixed_join_setup_data setup_mixed_join_common(table_view const& left_equality, - table_view const& right_equality, - table_view const& left_conditional, - table_view const& right_conditional, - ast::expression const& binary_predicate, - null_equality compare_nulls, - join_kind join_type, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) -{ - CUDF_EXPECTS(left_conditional.num_rows() == left_equality.num_rows(), - "The left conditional and equality tables must have the same number of rows."); - CUDF_EXPECTS(right_conditional.num_rows() == right_equality.num_rows(), - "The right conditional and equality tables must have the same number of rows."); - - auto const right_num_rows = right_conditional.num_rows(); - auto const left_num_rows = left_conditional.num_rows(); - auto const swap_tables = (join_type == join_kind::INNER_JOIN) && (right_num_rows > left_num_rows); - auto const outer_num_rows = swap_tables ? right_num_rows : left_num_rows; - - // If evaluating the expression may produce null outputs we create a nullable - // output column and follow the null-supporting expression evaluation code path. - auto const has_nulls = cudf::nullate::DYNAMIC{ - cudf::has_nulls(left_equality) || cudf::has_nulls(right_equality) || - binary_predicate.may_evaluate_null(left_conditional, right_conditional, stream)}; - - auto parser = ast::detail::expression_parser{ - binary_predicate, left_conditional, right_conditional, has_nulls, stream, mr}; - CUDF_EXPECTS(parser.output_type().id() == type_id::BOOL8, - "The expression must produce a boolean output.", - cudf::data_type_error); - - // TODO: The non-conditional join impls start with a dictionary matching, - // figure out what that is and what it's needed for (and if conditional joins - // need to do the same). - auto& probe = swap_tables ? right_equality : left_equality; - auto& build = swap_tables ? left_equality : right_equality; - - // Create hash table with load factor following hash join pattern - mixed_multiset_type hash_table{ - cuco::extent{static_cast(build.num_rows())}, - cudf::detail::CUCO_DESIRED_LOAD_FACTOR, - cuco::empty_key{cuco::pair{std::numeric_limits::max(), cudf::JoinNoMatch}}, - {}, - {}, - {}, - {}, - rmm::mr::polymorphic_allocator{}, - stream.value()}; - - // TODO: To add support for nested columns we will need to flatten in many - // places. However, this probably isn't worth adding any time soon since we - // won't be able to support AST conditions for those types anyway. - auto const row_bitmask = - cudf::detail::bitmask_and(build, stream, cudf::get_current_device_resource_ref()).first; - auto preprocessed_build = detail::row::equality::preprocessed_table::create(build, stream); - build_join_hash_table(build, - preprocessed_build, - hash_table, - has_nulls, - compare_nulls, - static_cast(row_bitmask.data()), - stream); - - auto left_conditional_view = table_device_view::create(left_conditional, stream); - auto right_conditional_view = table_device_view::create(right_conditional, stream); - - // For inner joins we support optimizing the join by launching one thread for - // whichever table is larger rather than always using the left table. - detail::grid_1d const config(outer_num_rows, DEFAULT_JOIN_BLOCK_SIZE); - auto const shmem_size_per_block = parser.shmem_per_thread * config.num_threads_per_block; - - auto preprocessed_probe = detail::row::equality::preprocessed_table::create(probe, stream); - auto const row_hash = cudf::detail::row::hash::row_hasher{preprocessed_probe}; - auto const hash_probe = row_hash.device_hasher(has_nulls); - auto const row_comparator = - cudf::detail::row::equality::two_table_comparator{preprocessed_probe, preprocessed_build}; - auto const equality_probe = row_comparator.equal_to(has_nulls, compare_nulls); - - // Precompute hash table storage and input data - auto hash_table_storage = cudf::device_span>{ - hash_table.data(), hash_table.capacity()}; - CUDF_EXPECTS(reinterpret_cast(hash_table_storage.data()) % - (2 * sizeof(cuco::pair)) == - 0, - "Hash table storage must be aligned to 2-element boundary"); - auto [input_pairs, hash_indices] = - precompute_mixed_join_data(hash_table, hash_probe, outer_num_rows, stream, mr); - - return {swap_tables, - outer_num_rows, - has_nulls, - std::move(parser), - std::move(hash_table), - std::move(preprocessed_build), - std::move(preprocessed_probe), - std::move(left_conditional_view), - std::move(right_conditional_view), - config, - shmem_size_per_block, - equality_probe, - hash_table_storage, - std::move(input_pairs), - std::move(hash_indices)}; -} /** - * @brief Helper function to compute the output size for mixed joins by launching count kernels. + * @brief Probes the equality hash table for the given join kind. * - * This function encapsulates the common logic needed by both mixed_join and - * compute_mixed_join_output_size to count the number of matches per row. + * The hash table is built on the right equality table and probed with the left equality table, + * yielding the index pairs that the conditional predicate is subsequently applied to. */ -std::pair>> -compute_mixed_join_matches_per_row( - cudf::nullate::DYNAMIC has_nulls, - table_device_view const& left_conditional_view, - table_device_view const& right_conditional_view, - bool is_outer_join, - bool swap_tables, - row_equality const& equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type outer_num_rows, - detail::grid_1d config, - thread_index_type shmem_size_per_block, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +std::pair>, + std::unique_ptr>> +equality_join_indices(cudf::hash_join const& hash_joiner, + table_view const& left_equality, + join_kind join_type, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { - auto matches_per_row = std::make_unique>( - static_cast(outer_num_rows), stream, mr); - auto matches_per_row_span = cudf::device_span{ - matches_per_row->begin(), static_cast(outer_num_rows)}; - - if (has_nulls) { - launch_mixed_join_count(left_conditional_view, - right_conditional_view, - is_outer_join, - swap_tables, - equality_probe, - hash_table_storage, - input_pairs, - hash_indices, - device_expression_data, - matches_per_row_span, - config, - shmem_size_per_block, - stream); - } else { - launch_mixed_join_count(left_conditional_view, - right_conditional_view, - is_outer_join, - swap_tables, - equality_probe, - hash_table_storage, - input_pairs, - hash_indices, - device_expression_data, - matches_per_row_span, - config, - shmem_size_per_block, - stream); + switch (join_type) { + case join_kind::INNER_JOIN: return hash_joiner.inner_join(left_equality, {}, stream, mr); + case join_kind::LEFT_JOIN: return hash_joiner.left_join(left_equality, {}, stream, mr); + case join_kind::FULL_JOIN: return hash_joiner.full_join(left_equality, {}, stream, mr); + default: CUDF_FAIL("Invalid join kind."); } - - std::size_t const size = - thrust::reduce(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - matches_per_row_span.begin(), - matches_per_row_span.end(), - std::size_t{0}); - - return {size, std::move(matches_per_row)}; } + } // anonymous namespace std::pair>, std::unique_ptr>> -mixed_join( - table_view const& left_equality, - table_view const& right_equality, - table_view const& left_conditional, - table_view const& right_conditional, - ast::expression const& binary_predicate, - null_equality compare_nulls, - join_kind join_type, - std::optional>> const& output_size_data, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr) +mixed_join(table_view const& left_equality, + table_view const& right_equality, + table_view const& left_conditional, + table_view const& right_conditional, + ast::expression const& binary_predicate, + null_equality compare_nulls, + join_kind join_type, + output_size_data_type const& output_size_data, + rmm::cuda_stream_view stream, + rmm::device_async_resource_ref mr) { CUDF_EXPECTS((join_type != join_kind::LEFT_SEMI_JOIN) && (join_type != join_kind::LEFT_ANTI_JOIN), "Left semi and anti joins should use mixed_join_semi."); + CUDF_EXPECTS(left_conditional.num_rows() == left_equality.num_rows(), + "The left conditional and equality tables must have the same number of rows."); + CUDF_EXPECTS(right_conditional.num_rows() == right_equality.num_rows(), + "The right conditional and equality tables must have the same number of rows."); - auto const right_num_rows = right_conditional.num_rows(); - auto const left_num_rows = left_conditional.num_rows(); - - // We can immediately filter out cases where the right table is empty. In - // some cases, we return all the rows of the left table with a corresponding - // null index for the right table; in others, we return an empty output. - if (right_num_rows == 0) { + // hash_join requires a non-empty build (right) table. + if (right_conditional.num_rows() == 0) { switch (join_type) { - // Left and full joins all return all the row indices from - // left with a corresponding NULL from the right. case join_kind::LEFT_JOIN: case join_kind::FULL_JOIN: return get_trivial_left_join_indices(left_conditional, stream, mr); - // Inner joins return empty output because no matches can exist. - case join_kind::INNER_JOIN: - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - default: CUDF_FAIL("Invalid join kind."); break; - } - } else if (left_num_rows == 0) { - switch (join_type) { - // Left and inner joins all return empty sets. - case join_kind::LEFT_JOIN: case join_kind::INNER_JOIN: - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - // Full joins need to return the trivial complement. - case join_kind::FULL_JOIN: { - auto ret_flipped = get_trivial_left_join_indices(right_conditional, stream, mr); - return std::pair(std::move(ret_flipped.second), std::move(ret_flipped.first)); - } - default: CUDF_FAIL("Invalid join kind."); break; + return std::pair{std::make_unique>(0, stream, mr), + std::make_unique>(0, stream, mr)}; + default: CUDF_FAIL("Invalid join kind."); } } - auto setup = setup_mixed_join_common(left_equality, - right_equality, - left_conditional, - right_conditional, - binary_predicate, - compare_nulls, - join_type, - stream, - mr); - - bool const is_outer_join = - (join_type == join_kind::LEFT_JOIN || join_type == join_kind::FULL_JOIN); - - // If the join size data was not provided as an input, compute it here. - std::size_t join_size = 0; - // Using an optional because we only need to allocate a new vector if one was - // not passed as input, and rmm::device_uvector is not default constructible - std::optional> matches_per_row{}; - device_span matches_per_row_span{}; - - if (output_size_data.has_value()) { - join_size = output_size_data->first; - matches_per_row_span = output_size_data->second; - } else { - auto [size, matches] = compute_mixed_join_matches_per_row(setup.has_nulls, - *setup.left_conditional_view, - *setup.right_conditional_view, - is_outer_join, - setup.swap_tables, - setup.equality_probe, - setup.hash_table_storage, - setup.input_pairs.data(), - setup.hash_indices.data(), - setup.parser.device_expression_data, - setup.outer_num_rows, - setup.config, - setup.shmem_size_per_block, - stream, - mr); - join_size = size; - matches_per_row = std::move(*matches); - matches_per_row_span = cudf::device_span{ - matches_per_row->begin(), static_cast(setup.outer_num_rows)}; - } - - // Given the number of matches per row, we need to compute the offsets for insertion. - auto join_result_offsets = - rmm::device_uvector{static_cast(setup.outer_num_rows), stream, mr}; - thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - matches_per_row_span.begin(), - matches_per_row_span.end(), - join_result_offsets.begin()); - - // Get total count from scan result: last offset + last matches_per_row - if (setup.outer_num_rows > 0 && !output_size_data.has_value()) { - auto const last_offset = join_result_offsets.element(setup.outer_num_rows - 1, stream); - auto const last_matches = matches_per_row->element(setup.outer_num_rows - 1, stream); - join_size = last_offset + last_matches; - } - - // The initial early exit clauses guarantee that we will not reach this point - // unless both the left and right tables are non-empty. Under that - // constraint, neither left nor full joins can return an empty result since - // at minimum we are guaranteed null matches for all non-matching rows. In - // all other cases (inner, left semi, and left anti joins) if we reach this - // point we can safely return an empty result. - if (join_size == 0) { - return std::pair(std::make_unique>(0, stream, mr), - std::make_unique>(0, stream, mr)); - } - - auto left_indices = std::make_unique>(join_size, stream, mr); - auto right_indices = std::make_unique>(join_size, stream, mr); - - auto const& join_output_l = left_indices->data(); - auto const& join_output_r = right_indices->data(); - - if (setup.has_nulls) { - launch_mixed_join(*setup.left_conditional_view, - *setup.right_conditional_view, - is_outer_join, - setup.swap_tables, - setup.equality_probe, - setup.hash_table_storage, - setup.input_pairs.data(), - setup.hash_indices.data(), - setup.parser.device_expression_data, - join_output_l, - join_output_r, - join_result_offsets.data(), - setup.config, - setup.shmem_size_per_block, - stream); - } else { - launch_mixed_join(*setup.left_conditional_view, - *setup.right_conditional_view, - is_outer_join, - setup.swap_tables, - setup.equality_probe, - setup.hash_table_storage, - setup.input_pairs.data(), - setup.hash_indices.data(), - setup.parser.device_expression_data, - join_output_l, - join_output_r, - join_result_offsets.data(), - setup.config, - setup.shmem_size_per_block, - stream); - } - - auto join_indices = std::pair(std::move(left_indices), std::move(right_indices)); - - // For full joins, get the indices in the right table that were not joined to - // by any row in the left table. + // A full join is a left join plus the unmatched-right complement. Build the left-outer result and + // append the complement with finalize_full_join rather than splitting failed pairs, which would + // emit spurious unmatched rows for keys that also match elsewhere. if (join_type == join_kind::FULL_JOIN) { - join_indices = detail::finalize_full_join( - std::move(join_indices), left_num_rows, right_num_rows, stream, mr); + auto left_outer = mixed_join(left_equality, + right_equality, + left_conditional, + right_conditional, + binary_predicate, + compare_nulls, + join_kind::LEFT_JOIN, + std::nullopt, + stream, + mr); + return finalize_full_join( + std::move(left_outer), left_conditional.num_rows(), right_conditional.num_rows(), stream, mr); } - return join_indices; + + auto const hash_joiner = cudf::hash_join{right_equality, compare_nulls, stream}; + auto const [left_indices, right_indices] = + equality_join_indices(hash_joiner, left_equality, join_type, stream, mr); + + auto const output_size = output_size_data.has_value() + ? std::optional{output_size_data->first} + : std::nullopt; + + return detail::filter_join_indices(left_conditional, + right_conditional, + *left_indices, + *right_indices, + binary_predicate, + join_type, + output_size, + stream, + mr); } std::pair>> @@ -542,67 +138,43 @@ compute_mixed_join_output_size(table_view const& left_equality, { CUDF_EXPECTS(join_type != join_kind::FULL_JOIN, "Size estimation is not available for full joins."); - CUDF_EXPECTS( (join_type != join_kind::LEFT_SEMI_JOIN) && (join_type != join_kind::LEFT_ANTI_JOIN), "Left semi and anti join size estimation should use compute_mixed_join_output_size_semi."); + CUDF_EXPECTS(left_conditional.num_rows() == left_equality.num_rows(), + "The left conditional and equality tables must have the same number of rows."); + CUDF_EXPECTS(right_conditional.num_rows() == right_equality.num_rows(), + "The right conditional and equality tables must have the same number of rows."); - auto const right_num_rows = right_conditional.num_rows(); - auto const left_num_rows = left_conditional.num_rows(); - - // Handle empty table cases early - if (right_num_rows == 0 || left_num_rows == 0) { - auto const outer_num_rows = - ((join_type == join_kind::INNER_JOIN) && (right_num_rows > left_num_rows)) ? right_num_rows - : left_num_rows; - auto matches_per_row = std::make_unique>( - static_cast(outer_num_rows), stream, mr); - auto matches_per_row_span = cudf::device_span{ - matches_per_row->begin(), static_cast(outer_num_rows)}; - - if (right_num_rows == 0 && join_type == join_kind::LEFT_JOIN) { - thrust::fill(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - matches_per_row_span.begin(), - matches_per_row_span.end(), - 1); - return {left_num_rows, std::move(matches_per_row)}; - } else { - thrust::fill(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), - matches_per_row_span.begin(), - matches_per_row_span.end(), - 0); - return {0, std::move(matches_per_row)}; + // hash_join requires a non-empty build (right) table. + if (right_conditional.num_rows() == 0) { + auto const left_num_rows = left_conditional.num_rows(); + if (join_type == join_kind::LEFT_JOIN) { + auto counts = + rmm::device_uvector(static_cast(left_num_rows), stream, mr); + thrust::uninitialized_fill( + rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), + counts.begin(), + counts.end(), + size_type{1}); + return {static_cast(left_num_rows), + std::make_unique>(std::move(counts))}; } + return {0, std::make_unique>(0, stream, mr)}; } - auto setup = setup_mixed_join_common(left_equality, - right_equality, - left_conditional, - right_conditional, - binary_predicate, - compare_nulls, - join_type, - stream, - mr); - - bool const is_outer_join = (join_type == join_kind::LEFT_JOIN); - - // Use the helper function to compute matches per row - return compute_mixed_join_matches_per_row(setup.has_nulls, - *setup.left_conditional_view, - *setup.right_conditional_view, - is_outer_join, - setup.swap_tables, - setup.equality_probe, - setup.hash_table_storage, - setup.input_pairs.data(), - setup.hash_indices.data(), - setup.parser.device_expression_data, - setup.outer_num_rows, - setup.config, - setup.shmem_size_per_block, - stream, - mr); + auto const hash_joiner = cudf::hash_join{right_equality, compare_nulls, stream}; + auto const [left_indices, right_indices] = + equality_join_indices(hash_joiner, left_equality, join_type, stream, mr); + + return cudf::filter_join_indices_output_size(left_conditional, + right_conditional, + *left_indices, + *right_indices, + binary_predicate, + join_type, + stream, + mr); } } // namespace detail diff --git a/cpp/src/join/mixed_join_common_utils.cuh b/cpp/src/join/mixed_join_common_utils.cuh index ecf9360c096d..195ef534d80a 100644 --- a/cpp/src/join/mixed_join_common_utils.cuh +++ b/cpp/src/join/mixed_join_common_utils.cuh @@ -1,85 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once #include -#include #include #include #include -#include #include -#include - -#include -#include -#include - -#include namespace cudf::detail { -using pair_type = cuco::pair; - -using hash_type = cuco::murmurhash3_32; - -/** - * @brief A custom comparator used for the mixed join multiset insertion - */ -struct mixed_join_always_not_equal { - __device__ constexpr bool operator()(pair_type const&, pair_type const&) const noexcept - { - return false; - } -}; - -/** - * @brief Hash functions for double hashing in mixed joins. - * - * These hashers implement a double hashing scheme for the mixed join multiset: - * - * - mixed_join_hasher1: Determines the initial probe slot for a given key. We simply use - * the precomputed row hash value, which is the first element of our (row_hash, row_index) pair. - * - * - mixed_join_hasher2: Determines the step size for the probing sequence. This allows keys - * with the same hash value to have different step sizes, helping to avoid secondary clustering. - * - * Note: Strictly speaking, this setup does not truly avoid secondary clustering because rows with - * the same hash value still receive the same step size. A true secondary clustering avoidance - * method would compute a different hash value for each row. However, based on performance testing, - * this current approach actually delivers better performance than computing row hashes with a - * different hasher. - */ -struct mixed_join_hasher1 { - __device__ constexpr hash_value_type operator()(pair_type const& key) const noexcept - { - return key.first; - } -}; - -struct mixed_join_hasher2 { - mixed_join_hasher2(hash_value_type seed) : _hash{seed} {} - - __device__ constexpr hash_value_type operator()(pair_type const& key) const noexcept - { - return _hash(key.first); - } - - private: - hash_type _hash; -}; - -using mixed_multiset_type = - cuco::static_multiset, - cuda::thread_scope_device, - mixed_join_always_not_equal, - cuco::double_hashing<1, mixed_join_hasher1, mixed_join_hasher2>, - rmm::mr::polymorphic_allocator, - cuco::storage<2>>; - using row_hash = cudf::detail::row::hash::device_row_hasher; @@ -116,141 +48,4 @@ struct expression_equality { row_equality const& equality_probe; }; -/** - * @brief Equality comparator for cuco::static_multiset queries. - * - * This equality comparator is designed for use with cuco::static_multiset's APIs. - * A probe hit indicates that the hashes of the keys are equal, at which point - * this comparator checks whether the keys themselves are equal (using the - * provided row_equality comparator) and then evaluates the conditional expression - */ -template -struct pair_expression_equality : public expression_equality { - using expression_equality::expression_equality; - -#ifndef NDEBUG - __attribute__((noinline)) -#else - __forceinline__ -#endif - __device__ bool - operator()(pair_type const& left_row, pair_type const& right_row) const noexcept - { - using cudf::detail::row::lhs_index_type; - using cudf::detail::row::rhs_index_type; - - auto output_dest = cudf::ast::detail::value_expression_result(); - // Three levels of checks: - // 1. Row hashes of the columns involved in the equality condition are equal. - // 2. The contents of the columns involved in the equality condition are equal. - // 3. The predicate evaluated on the relevant columns (already encoded in the evaluator) - // evaluates to true. - if ((left_row.first == right_row.first) && - this->equality_probe(lhs_index_type{left_row.second}, rhs_index_type{right_row.second})) { - auto const lrow_idx = this->swap_tables ? right_row.second : left_row.second; - auto const rrow_idx = this->swap_tables ? left_row.second : right_row.second; - this->evaluator.evaluate( - output_dest, lrow_idx, rrow_idx, 0, this->thread_intermediate_storage); - return (output_dest.is_valid() && output_dest.value()); - } - return false; - } -}; - -/** - * @brief Common utility for probing a hash table bucket and checking slot equality - * - * This encapsulates the common logic of reading bucket slots and checking for - * empty slots and key equality, used by both count and retrieve operations. - */ -template -struct hash_probe_result { - bool first_slot_is_empty_; - bool second_slot_is_empty_; - bool first_slot_equals_; - bool second_slot_equals_; - - __device__ __forceinline__ hash_probe_result( - pair_expression_equality const& key_equal, - cudf::device_span> hash_table_storage, - cuco::pair const& probe_key, - std::size_t probe_idx) - { - auto const* data = hash_table_storage.data(); - __builtin_assume_aligned(data, 2 * sizeof(cuco::pair)); - auto const first = *(data + probe_idx); - auto const second = *(data + probe_idx + 1); - - first_slot_is_empty_ = first.second == cudf::JoinNoMatch; - second_slot_is_empty_ = second.second == cudf::JoinNoMatch; - first_slot_equals_ = (not first_slot_is_empty_ and key_equal(probe_key, first)); - second_slot_equals_ = (not second_slot_is_empty_ and key_equal(probe_key, second)); - } - - __device__ __forceinline__ bool has_empty_slot() const noexcept - { - return first_slot_is_empty_ or second_slot_is_empty_; - } - - __device__ __forceinline__ cudf::size_type match_count() const noexcept - { - return static_cast(first_slot_equals_) + - static_cast(second_slot_equals_); - } - - __device__ __forceinline__ bool has_match() const noexcept - { - return first_slot_equals_ or second_slot_equals_; - } -}; - -/** - * @brief Iterator-style wrapper for probing through a hash table - * - * This encapsulates the common double hashing probe sequence used by both - * count and retrieve kernels. - */ -template -struct hash_table_prober { - cudf::device_span> hash_table_storage_; - pair_expression_equality const& key_equal_; - cuco::pair const& probe_key_; - std::size_t probe_idx_; - std::size_t step_; - std::size_t extent_; - - __device__ __forceinline__ hash_table_prober( - pair_expression_equality const& key_equal, - cudf::device_span> hash_table_storage, - cuco::pair const& probe_key, - cuda::std::pair const& hash_idx) - : hash_table_storage_{hash_table_storage}, - key_equal_{key_equal}, - probe_key_{probe_key}, - probe_idx_{static_cast(hash_idx.first)}, - step_{static_cast(hash_idx.second)}, - extent_{hash_table_storage.size()} - { - } - - __device__ __forceinline__ hash_probe_result probe_current_bucket() const - { - return hash_probe_result{key_equal_, hash_table_storage_, probe_key_, probe_idx_}; - } - - __device__ __forceinline__ void advance() noexcept - { - probe_idx_ = (probe_idx_ + step_) % extent_; - } - - __device__ __forceinline__ auto get_bucket_slots() const noexcept - { - auto const* data = hash_table_storage_.data(); - __builtin_assume_aligned(data, 2 * sizeof(cuco::pair)); - auto const first = *(data + probe_idx_); - auto const second = *(data + probe_idx_ + 1); - return cuda::std::pair{first, second}; - } -}; - } // namespace cudf::detail diff --git a/cpp/src/join/mixed_join_kernel.cu b/cpp/src/join/mixed_join_kernel.cu deleted file mode 100644 index 982ca557f5b6..000000000000 --- a/cpp/src/join/mixed_join_kernel.cu +++ /dev/null @@ -1,28 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mixed_join_kernel.cuh" -#include "mixed_join_kernel.hpp" - -namespace cudf::detail { - -template void launch_mixed_join( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type* join_output_l, - size_type* join_output_r, - cudf::size_type const* join_result_offsets, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace cudf::detail diff --git a/cpp/src/join/mixed_join_kernel.cuh b/cpp/src/join/mixed_join_kernel.cuh deleted file mode 100644 index 4c865b44d3b4..000000000000 --- a/cpp/src/join/mixed_join_kernel.cuh +++ /dev/null @@ -1,182 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "join_common_utils.hpp" -#include "mixed_join_common_utils.cuh" -#include "mixed_join_kernel.hpp" - -#include -#include -#include -#include -#include - -#include -#include - -namespace cudf { -namespace detail { - -/** - * @brief Optimized retrieve implementation using precomputed matches per row - * - * This implementation uses precomputed match counts to avoid expensive atomic - * operations and directly fills output arrays based on known match positions. - * - * @tparam is_outer Boolean flag indicating whether outer join semantics should be used - * @tparam has_nulls Whether the input tables may contain nulls - */ -template -__device__ __forceinline__ void retrieve_matches( - cudf::device_span> hash_table_storage, - pair_expression_equality const& key_equal, - cuco::pair const& probe_key, - cuda::std::pair const& hash_idx, - cudf::size_type* probe_output, - cudf::size_type* match_output) noexcept -{ - auto const probe_row_index = probe_key.second; - cudf::size_type output_idx = 0; - bool found_match = false; - auto prober = hash_table_prober{key_equal, hash_table_storage, probe_key, hash_idx}; - - while (true) { - auto const result = prober.probe_current_bucket(); - auto const bucket_slots = prober.get_bucket_slots(); - - if (result.first_slot_equals_) { - probe_output[output_idx] = probe_row_index; - match_output[output_idx] = bucket_slots.first.second; - output_idx++; - found_match = true; - } - - if (result.second_slot_equals_) { - probe_output[output_idx] = probe_row_index; - match_output[output_idx] = bucket_slots.second.second; - output_idx++; - found_match = true; - } - - // Exit if we find an empty slot - if (result.has_empty_slot()) { break; } - - prober.advance(); - } - - // Handle outer join logic for non-matching rows - if constexpr (is_outer) { - if (not found_match) { - probe_output[0] = probe_row_index; - match_output[0] = cudf::JoinNoMatch; - } - } -} - -template -CUDF_KERNEL void __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) - mixed_join(table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type* join_output_l, - size_type* join_output_r, - cudf::size_type const* join_result_offsets) -{ - // Normally the casting of a shared memory array is used to create multiple - // arrays of different types from the shared memory buffer, but here it is - // used to circumvent conflicts between arrays of different types between - // different template instantiations due to the extern specifier. - extern __shared__ char raw_intermediate_storage[]; - cudf::ast::detail::IntermediateDataType* intermediate_storage = - reinterpret_cast*>(raw_intermediate_storage); - auto thread_intermediate_storage = - &intermediate_storage[threadIdx.x * device_expression_data.num_intermediates]; - - cudf::size_type const left_num_rows = left_table.num_rows(); - cudf::size_type const right_num_rows = right_table.num_rows(); - auto const outer_num_rows = (swap_tables ? right_num_rows : left_num_rows); - - auto const start_idx = cudf::detail::grid_1d::global_thread_id(); - auto const stride = cudf::detail::grid_1d::grid_stride(); - - auto const evaluator = cudf::ast::detail::expression_evaluator{ - left_table, right_table, device_expression_data}; - - auto const equality = pair_expression_equality{ - evaluator, thread_intermediate_storage, swap_tables, equality_probe}; - - // Process each row and write matches to precomputed output positions - for (auto outer_row_index = start_idx; outer_row_index < outer_num_rows; - outer_row_index += stride) { - auto const& probe_key = input_pairs[outer_row_index]; - auto const& hash_idx = hash_indices[outer_row_index]; - auto const output_offset = join_result_offsets[outer_row_index]; - - if (is_outer_join) { - retrieve_matches( - hash_table_storage, - equality, - probe_key, - hash_idx, - swap_tables ? join_output_r + output_offset : join_output_l + output_offset, - swap_tables ? join_output_l + output_offset : join_output_r + output_offset); - } else { - retrieve_matches( - hash_table_storage, - equality, - probe_key, - hash_idx, - swap_tables ? join_output_r + output_offset : join_output_l + output_offset, - swap_tables ? join_output_l + output_offset : join_output_r + output_offset); - } - } -} - -template -void launch_mixed_join( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type* join_output_l, - size_type* join_output_r, - cudf::size_type const* join_result_offsets, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream) -{ - mixed_join - <<>>( - left_table, - right_table, - is_outer_join, - swap_tables, - equality_probe, - hash_table_storage, - input_pairs, - hash_indices, - device_expression_data, - join_output_l, - join_output_r, - join_result_offsets); - CUDF_CUDA_TRY(cudaGetLastError()); -} - -} // namespace detail - -} // namespace cudf diff --git a/cpp/src/join/mixed_join_kernel.hpp b/cpp/src/join/mixed_join_kernel.hpp deleted file mode 100644 index d465e8b21c60..000000000000 --- a/cpp/src/join/mixed_join_kernel.hpp +++ /dev/null @@ -1,73 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "mixed_join_common_utils.cuh" - -#include -#include -#include -#include -#include - -#include - -#include -#include - -namespace CUDF_EXPORT cudf { -namespace detail { - -/** - * @brief Performs a join using the combination of a hash lookup to identify - * equal rows between one pair of tables and the evaluation of an expression - * containing an arbitrary expression. - * - * This method probes the hash table with each row in the probe table using a - * custom equality comparator that also checks that the conditional expression - * evaluates to true between the left/right tables when a match is found - * between probe and build rows. - * - * @tparam has_nulls Whether or not the inputs may contain nulls. - * - * @param[in] left_table The left table - * @param[in] right_table The right table - * @param[in] is_outer_join Whether this is an outer join - * @param[in] swap_tables If true, the kernel was launched with one thread per right row and - * the kernel needs to internally loop over left rows. Otherwise, loop over right rows. - * @param[in] equality_probe The equality comparator used when probing the hash table. - * @param[in] hash_table_storage Device span of the hash table storage - * @param[in] input_pairs Precomputed input pairs for probing - * @param[in] hash_indices Precomputed hash indices for efficient probing - * @param[in] device_expression_data Container of device data required to evaluate the desired - * expression. - * @param[out] join_output_l The left result of the join operation - * @param[out] join_output_r The right result of the join operation - * @param[in] join_result_offsets Prefix sum of matches_per_row to get output offsets - * @param[in] config Grid configuration for the kernel launch - * @param[in] shmem_size_per_block Shared memory size per block - * @param[in] stream CUDA stream to use - */ -template -void launch_mixed_join( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type* join_output_l, - size_type* join_output_r, - cudf::size_type const* join_result_offsets, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace detail -} // namespace CUDF_EXPORT cudf diff --git a/cpp/src/join/mixed_join_kernel_nulls.cu b/cpp/src/join/mixed_join_kernel_nulls.cu deleted file mode 100644 index d3b382331dc9..000000000000 --- a/cpp/src/join/mixed_join_kernel_nulls.cu +++ /dev/null @@ -1,28 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mixed_join_kernel.cuh" -#include "mixed_join_kernel.hpp" - -namespace cudf::detail { - -template void launch_mixed_join( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - cudf::ast::detail::expression_device_view device_expression_data, - size_type* join_output_l, - size_type* join_output_r, - cudf::size_type const* join_result_offsets, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace cudf::detail diff --git a/cpp/src/join/mixed_join_semi.cu b/cpp/src/join/mixed_join_semi.cu index aa7674fb86dc..8021cdab21bc 100644 --- a/cpp/src/join/mixed_join_semi.cu +++ b/cpp/src/join/mixed_join_semi.cu @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include diff --git a/cpp/src/join/mixed_join_size_kernel.cu b/cpp/src/join/mixed_join_size_kernel.cu deleted file mode 100644 index b594b8a1a334..000000000000 --- a/cpp/src/join/mixed_join_size_kernel.cu +++ /dev/null @@ -1,28 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mixed_join_size_kernel.cuh" -#include "mixed_join_size_kernel.hpp" - -namespace cudf { -namespace detail { - -template void launch_mixed_join_count( - cudf::table_device_view left_table, - cudf::table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - ast::detail::expression_device_view device_expression_data, - cudf::device_span matches_per_row, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace detail -} // namespace cudf diff --git a/cpp/src/join/mixed_join_size_kernel.cuh b/cpp/src/join/mixed_join_size_kernel.cuh deleted file mode 100644 index 9a719170d7a3..000000000000 --- a/cpp/src/join/mixed_join_size_kernel.cuh +++ /dev/null @@ -1,136 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "join_common_utils.hpp" -#include "mixed_join_common_utils.cuh" -#include "mixed_join_size_kernel.hpp" - -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace cudf::detail { - -/** - * @brief Standalone count implementation using precomputed hash indices - * - * This implementation provides essential count functionality for mixed joins - * using precomputed probe indices and step sizes. - */ -template -__device__ __forceinline__ auto standalone_count( - pair_expression_equality const& key_equal, - cudf::device_span> hash_table_storage, - cuco::pair const& probe_key, - cuda::std::pair const& hash_idx, - bool is_outer_join) noexcept -{ - cudf::size_type count = 0; - auto prober = hash_table_prober{key_equal, hash_table_storage, probe_key, hash_idx}; - - while (true) { - auto const result = prober.probe_current_bucket(); - count += result.match_count(); - - // Exit if we find an empty slot - if (result.has_empty_slot()) { - // Handle outer join logic: non-matching rows are counted as 1 match - if (is_outer_join && count == 0) { return 1; } - return count; - } - - prober.advance(); - } -} - -template -CUDF_KERNEL void __launch_bounds__(DEFAULT_JOIN_BLOCK_SIZE) mixed_join_count( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - ast::detail::expression_device_view device_expression_data, - cudf::device_span matches_per_row) -{ - // The (required) extern storage of the shared memory array leads to - // conflicting declarations between different templates. The easiest - // workaround is to declare an arbitrary (here char) array type then cast it - // after the fact to the appropriate type. - extern __shared__ char raw_intermediate_storage[]; - auto intermediate_storage = - reinterpret_cast*>(raw_intermediate_storage); - auto thread_intermediate_storage = - intermediate_storage + (threadIdx.x * device_expression_data.num_intermediates); - - auto const start_idx = cudf::detail::grid_1d::global_thread_id(); - auto const stride = cudf::detail::grid_1d::grid_stride(); - cudf::size_type const left_num_rows = left_table.num_rows(); - cudf::size_type const right_num_rows = right_table.num_rows(); - auto const outer_num_rows = (swap_tables ? right_num_rows : left_num_rows); - - auto const evaluator = cudf::ast::detail::expression_evaluator{ - left_table, right_table, device_expression_data}; - - // Figure out the number of elements for this key. - // TODO: Address asymmetry in operator. - auto count_equality = pair_expression_equality{ - evaluator, thread_intermediate_storage, swap_tables, equality_probe}; - - for (auto outer_row_index = start_idx; outer_row_index < outer_num_rows; - outer_row_index += stride) { - auto const& probe_key = input_pairs[outer_row_index]; - auto const& hash_idx = hash_indices[outer_row_index]; - - auto match_count = - standalone_count(count_equality, hash_table_storage, probe_key, hash_idx, is_outer_join); - - matches_per_row[outer_row_index] = match_count; - } -} - -template -void launch_mixed_join_count( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - ast::detail::expression_device_view device_expression_data, - cudf::device_span matches_per_row, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream) -{ - mixed_join_count - <<>>( - left_table, - right_table, - is_outer_join, - swap_tables, - equality_probe, - hash_table_storage, - input_pairs, - hash_indices, - device_expression_data, - matches_per_row); - CUDF_CUDA_TRY(cudaGetLastError()); -} - -} // namespace cudf::detail diff --git a/cpp/src/join/mixed_join_size_kernel.hpp b/cpp/src/join/mixed_join_size_kernel.hpp deleted file mode 100644 index 16909cee1518..000000000000 --- a/cpp/src/join/mixed_join_size_kernel.hpp +++ /dev/null @@ -1,74 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include "mixed_join_common_utils.cuh" - -#include -#include -#include -#include -#include -#include - -#include - -#include -#include - -namespace CUDF_EXPORT cudf { -namespace detail { - -/** - * @brief Computes the output size of joining the left table to the right table. - * - * This method probes the hash table with each row in the probe table using a - * custom equality comparator that also checks that the conditional expression - * evaluates to true between the left/right tables when a match is found - * between probe and build rows. - * - * @tparam has_nulls Whether or not the inputs may contain nulls. - * - * @param[in] left_table The left table - * @param[in] right_table The right table - * @param[in] is_outer_join Whether this is an outer join - * @param[in] swap_tables If true, the kernel was launched with one thread per right row and - * the kernel needs to internally loop over left rows. Otherwise, loop over right rows. - * @param[in] equality_probe The equality comparator used when probing the hash table. - * @param[in] hash_table_storage Device span of the hash table storage - * @param[in] input_pairs Precomputed input pairs for probing - * @param[in] hash_indices Precomputed hash indices for efficient probing - * @param[in] device_expression_data Container of device data required to evaluate the desired - * expression. - * @param[out] matches_per_row The number of matches in one pair of - * equality/conditional tables for each row in the other pair of tables. If - * swap_tables is true, matches_per_row corresponds to the right_table, - * otherwise it corresponds to the left_table. Note that corresponding swap of - * left/right tables to determine which is the build table and which is the - * probe table has already happened on the host. - * @param[in] config Grid configuration for the kernel launch - * @param[in] shmem_size_per_block Shared memory size per block - * @param[in] stream CUDA stream to use - */ - -template -void launch_mixed_join_count( - table_device_view left_table, - table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - ast::detail::expression_device_view device_expression_data, - cudf::device_span matches_per_row, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace detail -} // namespace CUDF_EXPORT cudf diff --git a/cpp/src/join/mixed_join_size_kernel_nulls.cu b/cpp/src/join/mixed_join_size_kernel_nulls.cu deleted file mode 100644 index 2f94da3bfd03..000000000000 --- a/cpp/src/join/mixed_join_size_kernel_nulls.cu +++ /dev/null @@ -1,28 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mixed_join_size_kernel.cuh" -#include "mixed_join_size_kernel.hpp" - -namespace cudf { -namespace detail { - -template void launch_mixed_join_count( - cudf::table_device_view left_table, - cudf::table_device_view right_table, - bool is_outer_join, - bool swap_tables, - row_equality equality_probe, - cudf::device_span> hash_table_storage, - cuco::pair const* input_pairs, - cuda::std::pair const* hash_indices, - ast::detail::expression_device_view device_expression_data, - cudf::device_span matches_per_row, - detail::grid_1d config, - int64_t shmem_size_per_block, - rmm::cuda_stream_view stream); - -} // namespace detail -} // namespace cudf diff --git a/cpp/src/strings/regex/glushkov_regcomp.cpp b/cpp/src/strings/regex/glushkov_regcomp.cpp index cb740d0cb9af..8d06aa980208 100644 --- a/cpp/src/strings/regex/glushkov_regcomp.cpp +++ b/cpp/src/strings/regex/glushkov_regcomp.cpp @@ -278,9 +278,9 @@ bool positions_chars_overlap(gkprog const& gp, uint32_t const p, uint32_t const * Glushkov's bit-order cannot represent. * * Two rules: - * Rule 1 – END before char: an ACCEPT item appears before the first CHAR_POS - * item in Thompson priority order → the pattern can empty-match in a way - * that priority_kill cannot handle correctly. + * Rule 1 – ACCEPT before later char: an ACCEPT item appears before a CHAR_POS + * item in Thompson priority order → the accepted path has higher + * priority than a continuation that Glushkov cannot kill correctly. * Rule 2 – non-monotone gpos + char overlap: two CHAR_POS items appear with * the higher-priority one at a larger gpos (inverted bit order), AND * they can match a common character → priority_kill picks the wrong @@ -288,19 +288,13 @@ bool positions_chars_overlap(gkprog const& gp, uint32_t const p, uint32_t const */ bool frontier_has_priority_conflict(std::vector const& items, gkprog const& gp) { - // Rule 1: ACCEPT before any CHAR_POS, but only when the frontier also - // contains at least one CHAR_POS. An ACCEPT-only frontier (the normal - // "end of pattern" case) is not a priority conflict. - bool seen_char = false; - bool accept_before_char = false; + // Rule 1: ACCEPT before a later CHAR_POS. A frontier ending in ACCEPT (the + // normal "end of pattern" case) is not a priority conflict. + bool seen_accept = false; for (auto const& item : items) { - if (item.kind == frontier_item::CHAR_POS) { - seen_char = true; - } else if (item.kind == frontier_item::ACCEPT && !seen_char) { - accept_before_char = true; - } + if (item.kind == frontier_item::ACCEPT) { seen_accept = true; } + if (item.kind == frontier_item::CHAR_POS && seen_accept) { return true; } } - if (accept_before_char && seen_char) { return true; } // Rule 2: non-monotone gpos pair with character overlap for (size_t i = 0; i < items.size(); ++i) { diff --git a/cpp/src/strings/regex/glushkov_regcomp.hpp b/cpp/src/strings/regex/glushkov_regcomp.hpp index 32c306520c25..6d7450aef070 100644 --- a/cpp/src/strings/regex/glushkov_regcomp.hpp +++ b/cpp/src/strings/regex/glushkov_regcomp.hpp @@ -102,6 +102,8 @@ struct gkprog { * - Pattern has more than GLUSHKOV_MAX_STATES character-consuming positions. * - Pattern is nullable (matches the empty string): priority semantics cannot * be faithfully represented without an ε-position for the empty match. + * - Pattern has a Thompson-priority frontier that Glushkov's bit ordering + * cannot represent faithfully. * * @param prog Compiled Thompson NFA (after reprog::finalize()). * @return Host-side Glushkov program, or nullptr on failure. diff --git a/cpp/tests/join/mixed_join_tests.cu b/cpp/tests/join/mixed_join_tests.cu index 3300e3b56f17..701800cd5726 100644 --- a/cpp/tests/join/mixed_join_tests.cu +++ b/cpp/tests/join/mixed_join_tests.cu @@ -18,6 +18,8 @@ #include #include +#include + #include #include #include @@ -269,15 +271,14 @@ struct MixedJoinPairReturnTest : public MixedJoinTest { left_equality, right_equality, left_conditional, right_conditional, predicate, compare_nulls); EXPECT_TRUE(result_size == expected_outputs.size()); - cudf::test::fixed_width_column_wrapper expected_counts_cw( - expected_counts.begin(), expected_counts.end()); - auto const actual_counts_view = - cudf::column_view(cudf::data_type{cudf::type_to_id()}, - actual_counts->size(), - actual_counts->data(), - nullptr, - 0); - CUDF_TEST_EXPECT_COLUMNS_EQUAL(expected_counts_cw, actual_counts_view); + auto const expected_total = + std::accumulate(expected_counts.begin(), expected_counts.end(), std::size_t{0}); + EXPECT_EQ(expected_total, result_size); + auto const actual_total = thrust::reduce(rmm::exec_policy_nosync(cudf::get_default_stream()), + actual_counts->begin(), + actual_counts->end(), + std::size_t{0}); + EXPECT_EQ(actual_total, result_size); auto result = this->join(left_equality, right_equality, @@ -434,14 +435,14 @@ struct MixedInnerJoinTest : public MixedJoinPairReturnTest { this->compare_join_results(mixed_result, ast_filter_result); // Verify filter_join_indices_output_size matches the materialized output size. - auto const fji_size = cudf::filter_join_indices_output_size( + auto const filter_output_size_result = cudf::filter_join_indices_output_size( left_conditional, right_conditional, cudf::device_span(*hash_join_result.first), cudf::device_span(*hash_join_result.second), predicate, cudf::join_kind::INNER_JOIN); - EXPECT_EQ(fji_size, ast_filter_result.first->size()); + EXPECT_EQ(filter_output_size_result.first, ast_filter_result.first->size()); // Verify JIT filter_join_indices if provided if (!jit_predicate.empty()) { @@ -1102,14 +1103,14 @@ struct MixedLeftJoinTest : public MixedJoinPairReturnTest { this->compare_join_results(mixed_result, ast_filter_result); // Verify filter_join_indices_output_size matches the materialized output size. - auto const fji_size = cudf::filter_join_indices_output_size( + auto const filter_output_size_result = cudf::filter_join_indices_output_size( left_conditional, right_conditional, cudf::device_span(*hash_join_result.first), cudf::device_span(*hash_join_result.second), predicate, cudf::join_kind::LEFT_JOIN); - EXPECT_EQ(fji_size, ast_filter_result.first->size()); + EXPECT_EQ(filter_output_size_result.first, ast_filter_result.first->size()); // Verify JIT filter_join_indices if provided if (!jit_predicate.empty()) { @@ -1369,60 +1370,8 @@ struct MixedFullJoinTest : public MixedJoinPairReturnTest { cudf::null_equality compare_nulls = cudf::null_equality::EQUAL, std::string const& jit_predicate = "") override { - // Test both approaches and verify they produce the same results - auto mixed_result = cudf::mixed_full_join( + return cudf::mixed_full_join( left_equality, right_equality, left_conditional, right_conditional, predicate, compare_nulls); - - // Alternative approach: hash_join + filter_join_indices - // Skip hash_join approach for empty tables (hash_join doesn't support empty tables) - if (left_equality.num_rows() > 0 && right_equality.num_rows() > 0) { - cudf::hash_join hash_joiner(right_equality, compare_nulls); - auto hash_join_result = hash_joiner.full_join(left_equality); - - // Verify AST filter_join_indices - auto ast_filter_result = cudf::filter_join_indices( - left_conditional, - right_conditional, - cudf::device_span(*hash_join_result.first), - cudf::device_span(*hash_join_result.second), - predicate, - cudf::join_kind::FULL_JOIN); - this->compare_join_results(mixed_result, ast_filter_result); - - // Verify filter_join_indices_output_size matches the materialized output size. - auto const fji_size = cudf::filter_join_indices_output_size( - left_conditional, - right_conditional, - cudf::device_span(*hash_join_result.first), - cudf::device_span(*hash_join_result.second), - predicate, - cudf::join_kind::FULL_JOIN); - EXPECT_EQ(fji_size, ast_filter_result.first->size()); - - // Verify JIT filter_join_indices if provided - if (!jit_predicate.empty()) { - auto jit_filter_result = cudf::filter_join_indices_jit( - left_conditional, - right_conditional, - cudf::device_span(*hash_join_result.first), - cudf::device_span(*hash_join_result.second), - jit_predicate, - cudf::join_kind::FULL_JOIN); - this->compare_join_results(mixed_result, jit_filter_result); - } - - // Verify AST-based JIT filter_join_indices - auto jit_ast_filter_result = cudf::filter_join_indices_jit( - left_conditional, - right_conditional, - cudf::device_span(*hash_join_result.first), - cudf::device_span(*hash_join_result.second), - predicate, - cudf::join_kind::FULL_JOIN); - this->compare_join_results(mixed_result, jit_ast_filter_result); - } - - return mixed_result; } std::pair>> join_size( @@ -1520,6 +1469,24 @@ TYPED_TEST(MixedFullJoinTest, Basic2) {cudf::JoinNoMatch, 2}}); } +TYPED_TEST(MixedFullJoinTest, MultiMatchUnmatchedDedup) +{ + auto const predicate = + cudf::ast::operation(cudf::ast::ast_operator::GREATER, col_ref_left_0, col_ref_right_0); + this->test({{5, 5, 7}, {10, 1, 10}}, + {{5, 5, 9}, {2, 20, 0}}, + {0}, + {1}, + predicate, + {}, + {{0, 0}, + {1, cudf::JoinNoMatch}, + {2, cudf::JoinNoMatch}, + {cudf::JoinNoMatch, 1}, + {cudf::JoinNoMatch, 2}}, + make_jit_comparison(1, 1, 0, 0, ">")); +} + using MixedFullJoinTest_int32 = MixedFullJoinTest; TEST_F(MixedFullJoinTest_int32, NullableColumnsWithModuloFilter) { diff --git a/cpp/tests/streams/join_test.cpp b/cpp/tests/streams/join_test.cpp index abebb6930791..91c77b011e84 100644 --- a/cpp/tests/streams/join_test.cpp +++ b/cpp/tests/streams/join_test.cpp @@ -161,6 +161,7 @@ TEST_F(JoinTest, LeftJoinWithPostFilter) cudf::device_span(*hash_join_result.second), left_zero_eq_right_zero, cudf::join_kind::LEFT_JOIN, + std::nullopt, cudf::test::get_default_stream()); } diff --git a/cpp/tests/strings/split_tests.cpp b/cpp/tests/strings/split_tests.cpp index cacd98328844..ce9fd7a8b10e 100644 --- a/cpp/tests/strings/split_tests.cpp +++ b/cpp/tests/strings/split_tests.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -517,6 +517,35 @@ TEST_F(StringsSplitTest, SplitRecordRegex) } } +TEST_F(StringsSplitTest, SplitRecordRegexLazyQuantifier) +{ + auto const input = cudf::test::strings_column_wrapper({"\rbaab\r\ra"}); + auto const sv = cudf::strings_column_view(input); + using LCW = cudf::test::lists_column_wrapper; + + { + LCW expected({LCW{"\rbaa", "\ra"}}); + auto const prog = + cudf::strings::regex_program::create("[^ \v\n\t\r\f]\\r+?\\n*", + cudf::strings::regex_flags::EXT_NEWLINE, + cudf::strings::capture_groups::NON_CAPTURE); + auto const result = cudf::strings::split_record_re(sv, *prog); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view(), expected); + } + + { + LCW expected({LCW{"\rbaa", "a"}}); + auto const prog = + cudf::strings::regex_program::create("[^ \v\n\t\r\f]\\r+\\n*", + cudf::strings::regex_flags::EXT_NEWLINE, + cudf::strings::capture_groups::NON_CAPTURE); + auto const result = cudf::strings::split_record_re(sv, *prog); + + CUDF_TEST_EXPECT_COLUMNS_EQUAL(result->view(), expected); + } +} + TEST_F(StringsSplitTest, SplitRegexWithMaxSplit) { std::vector h_strings{" Héllo\tthesé", nullptr, "are\nsome ", "tést\rString", ""}; diff --git a/docs/cudf/source/cudf_polars/api.md b/docs/cudf/source/cudf_polars/api.md index a1c6fcb76ff0..5c099b8b7177 100644 --- a/docs/cudf/source/cudf_polars/api.md +++ b/docs/cudf/source/cudf_polars/api.md @@ -76,6 +76,7 @@ Most users interact with them through `StreamingOptions` fields rather than dire .. automodule:: cudf_polars.utils.config :members: DynamicPlanningOptions, + JoinFilterPushdownOptions, MemoryResourceConfig, ParquetOptions, StreamingExecutor, diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index 5d813e74bbe1..6744ad959a93 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -108,6 +108,7 @@ Environment variables follow these patterns: | `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 | +| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `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` diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py index d6e16af4d9d3..79aaa7a47dbd 100644 --- a/python/cudf/cudf/core/groupby/groupby.py +++ b/python/cudf/cudf/core/groupby/groupby.py @@ -1273,6 +1273,17 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): ) elif agg_kind == "NUNIQUE": cast_dtype = np.dtype(np.int64) + elif ( + agg_name in {"cumsum", "cumprod"} + and is_pandas_nullable_extension_dtype(orig_dtype) + and orig_dtype.kind in {"i", "u"} + ): + # libcudf's SUM/PRODUCT scans promote narrow integers + # to 64-bit. pandas does the same for numpy dtypes + # (int8 -> int64, GH#37493) but preserves masked + # extension dtypes (Int16 stays Int16, GH#58811), + # wrapping on overflow. + cast_dtype = orig_dtype elif ( ( isinstance(agg_name, str) @@ -1335,7 +1346,7 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): # RangeIndex) columns. data = ColumnAccessor( data, - multiindex=False, + multiindex=self.obj._data.multiindex, level_names=self.obj._data.level_names, rangeindex=self.obj._data.rangeindex, label_dtype=self.obj._data.label_dtype, @@ -1346,11 +1357,26 @@ def agg(self, func=None, *args, engine=None, engine_kwargs=None, **kwargs): and self.obj.ndim == 2 and self.obj._data.level_names != (None,) ): + mi_kwargs: dict[str, Any] = {} + if self.obj._data.multiindex and all( + isinstance(label, tuple) + and len(label) == self.obj._data.nlevels + for label in data + ): + # the aggregation kept the source's tuple labels: preserve + # the MultiIndex columns and their per-level metadata. + # Relabeling aggregations (``agg(new=(col, func))``) emit + # new flat labels the source's multi-level metadata does + # not describe, so they keep the flat default. + mi_kwargs = { + "multiindex": True, + "level_dtypes": self.obj._data.level_dtypes, + } data = ColumnAccessor( data, - multiindex=False, level_names=self.obj._data.level_names, label_dtype=self.obj._data.label_dtype, + **mi_kwargs, ) else: data = ColumnAccessor(data, multiindex=multilevel) diff --git a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py index f6d058cc1dda..3e219a23e8ec 100644 --- a/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py +++ b/python/cudf/cudf/pandas/scripts/pandas-testing-plugin.py @@ -1763,7 +1763,6 @@ def pytest_unconfigure(config): "tests/groupby/aggregate/test_aggregate.py::test_agg_str_with_kwarg_axis_1_raises[nunique]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_directory[size]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_aggregate.py::test_groupby_aggregate_empty_key_empty_return": "TODO: Add a reason for failure", - "tests/groupby/aggregate/test_aggregate.py::test_multiindex_custom_func[0]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_aggregate.py::test_order_aggregate_multiple_funcs": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-False]": "TODO: Add a reason for failure", "tests/groupby/aggregate/test_cython.py::test_cython_agg_EA_known_dtypes[data1-prod-large_int-True]": "TODO: Add a reason for failure", @@ -1838,20 +1837,13 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Float32-False-val1]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Float64-False-val1]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int16-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int16-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int32-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int32-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int64-False-val1]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int8-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[Int8-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt16-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt16-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt32-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt32-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt64-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt64-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt8-False-val1]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_cumsum_mask[UInt8-True-3]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_skipna_false": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_cumsum_timedelta64": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_groups_in_BaseGrouper": "TODO: Add a reason for failure", @@ -1883,7 +1875,6 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[ns]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[s]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groupby_with_Time_Grouper[us]": "TODO: Add a reason for failure", - "tests/groupby/test_groupby.py::test_groupby_with_hier_columns": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groups_repr_truncates[1-{0: [0], ...}]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groups_repr_truncates[4-{0: [0], 1: [1], 2: [2], 3: [3], ...}]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_groups_repr_truncates[5-{0: [0], 1: [1], 2: [2], 3: [3], 4: [4]}]": "TODO: Add a reason for failure", @@ -1892,7 +1883,6 @@ def pytest_unconfigure(config): "tests/groupby/test_groupby.py::test_ops_not_as_index[idxmin]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_ops_not_as_index[size]": "TODO: Add a reason for failure", "tests/groupby/test_groupby.py::test_single_element_listlike_level_grouping[level_arg0-False]": "AssertionError: assert ['x', 'y'] == [('x',), ('y',)]", - "tests/groupby/test_groupby.py::test_wrap_aggregated_output_multindex": "TODO: Add a reason for failure", "tests/groupby/test_groupby_dropna.py::test_groupby_nan_included": "GroupBy.indices returns cupy arrays nested in a dict that cudf.pandas does not wrap, so assert_numpy_array_equal sees mismatched array classes", "tests/groupby/test_groupby_subclass.py::test_groupby_preserves_metadata": "proxy pd.Index passed to the raw SubclassedSeries constructor loses its name: pandas' maybe_extract_name checks isinstance against the concrete Index class, which proxies cannot satisfy", "tests/groupby/test_grouping.py::TestGetGroup::test_get_group_grouped_by_tuple": "TODO: Add a reason for failure", @@ -1918,7 +1908,6 @@ def pytest_unconfigure(config): "tests/groupby/test_reductions.py::test_sum_skipna_object[False]": "Inherent cudf.pandas None-vs-NaN difference for object-dtype null (skipna logic is correct)", "tests/groupby/test_timegrouper.py::TestGroupBy::test_groupby_with_timegrouper": "TODO: Add a reason for failure", "tests/groupby/test_timegrouper.py::TestGroupBy::test_scalar_call_versus_list_call": "TODO: Add a reason for failure", - "tests/groupby/transform/test_transform.py::test_nan_in_cumsum_group_label": "AssertionError: Attributes of Series are different", "tests/indexes/base_class/test_reshape.py::TestReshape::test_insert_missing[Decimal]": "TODO: Add a reason for failure", "tests/indexes/categorical/test_astype.py::TestAstype::test_categorical_date_roundtrip[False]": "TODO: Add a reason for failure", "tests/indexes/categorical/test_astype.py::TestAstype::test_categorical_date_roundtrip[True]": "TODO: Add a reason for failure", diff --git a/python/cudf/cudf/tests/groupby/test_agg.py b/python/cudf/cudf/tests/groupby/test_agg.py index 43b295e499ec..8e2a5ccfaf7f 100644 --- a/python/cudf/cudf/tests/groupby/test_agg.py +++ b/python/cudf/cudf/tests/groupby/test_agg.py @@ -804,3 +804,32 @@ def test_groupby_idxminmax_all_na_group_raises(op): getattr(pdf.groupby("key"), op)() with pytest.raises(ValueError): getattr(gdf.groupby("key"), op)() + + +def test_agg_multiindex_columns_preserved(): + # aggregating a MultiIndex-column frame keeps hierarchical columns + pdf = pd.DataFrame( + [[1, 2, 3], [1, 5, 6], [2, 8, 9]], + columns=pd.MultiIndex.from_tuples( + [("k", ""), ("x", "a"), ("x", "b")], names=["l0", "l1"] + ), + ) + gdf = cudf.DataFrame(pdf) + + expect = pdf.groupby(("k", "")).agg("sum") + got = gdf.groupby(("k", "")).agg("sum") + assert_eq(expect, got) + + +def test_agg_relabel_flat_columns_from_multiindex(): + # relabeling aggregations emit new flat labels; the source's + # multi-level column metadata must not be attached to them + pdf = pd.DataFrame( + [[1, 2], [1, 5], [2, 8]], + columns=pd.MultiIndex.from_tuples([("k", ""), ("x", "a")]), + ) + gdf = cudf.DataFrame(pdf) + + expect = pdf.groupby(("k", "")).agg(total=(("x", "a"), "sum")) + got = gdf.groupby(("k", "")).agg(total=(("x", "a"), "sum")) + assert_eq(expect, got) diff --git a/python/cudf/cudf/tests/groupby/test_cummulative.py b/python/cudf/cudf/tests/groupby/test_cummulative.py index c0135fe6c264..b44c273e16cd 100644 --- a/python/cudf/cudf/tests/groupby/test_cummulative.py +++ b/python/cudf/cudf/tests/groupby/test_cummulative.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import numpy as np import pandas as pd @@ -105,3 +105,27 @@ def test_scan_int_null_pandas_compatible(op): with cudf.option_context("mode.pandas_compatible", True): result = getattr(df_cudf.groupby("b")["a"], op)() assert_eq(result, expected) + + +@pytest.mark.parametrize("op", ["cumsum", "cumprod"]) +def test_groupby_cumscan_masked_dtype_preserved(op): + # pandas preserves masked extension dtypes for groupby cum-scans + # (Int16 stays Int16, GH#58811) while numpy ints promote to 64-bit + pdf = pd.DataFrame({"a": [1, 1, 2], "b": [1, pd.NA, 2]}, dtype="Int16") + gdf = cudf.DataFrame(pdf) + + expected = getattr(pdf.groupby("a")["b"], op)() + result = getattr(gdf.groupby("a")["b"], op)() + + assert_eq(expected, result) + + +def test_groupby_cumsum_numpy_dtype_promotes(): + # numpy int8 promotes to int64 (pandas GH#37493) + pdf = pd.DataFrame({"a": [1, 1], "b": [111, 111]}, dtype="int8") + gdf = cudf.DataFrame(pdf) + + expected = pdf.groupby("a").cumsum() + result = gdf.groupby("a").cumsum() + + assert_eq(expected, result) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py new file mode 100644 index 000000000000..382d0027f649 --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -0,0 +1,143 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Utilities for tracking column value domains between IR nodes.""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import singledispatch +from typing import TYPE_CHECKING + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) + +if TYPE_CHECKING: + from collections.abc import Mapping + + from cudf_polars.dsl.ir import IR + +__all__ = [ + "ColumnBinding", + "ColumnLineage", + "ColumnRef", + "column_domain_bindings", +] + + +@dataclass(frozen=True) +class ColumnBinding: + """A direct binding to a named column on a specific child edge.""" + + child_index: int + name: str + + +@dataclass(frozen=True) +class ColumnRef: + """A named column produced by an IR node.""" + + node: IR + name: str + + +@dataclass(frozen=True) +class ColumnLineage: + """Persistent value-domain lineage, sharing suffixes across DAG branches.""" + + column: ColumnRef + source: ColumnLineage | None = None + source_child_index: int | None = None + """Child edge leading to ``source``, or None if there is no source.""" + + +@singledispatch +def column_domain_bindings(node: IR) -> Mapping[str, ColumnBinding]: + """ + Map output columns to child columns containing their value domains. + + For every ``output_name -> ColumnBinding(child_index, input_name)`` binding, + every value appearing in ``node[output_name]`` is guaranteed to appear in + ``node.children[child_index][input_name]``. Row order, multiplicity, and + cardinality are not preserved. + + If a name in ``node.schema`` does not appear in the mapping it means + that it was not possible to derive a relationship between the domain of + the output and input values for that column. + """ + return {} + + +@column_domain_bindings.register(Select) +def _(node: Select) -> Mapping[str, ColumnBinding]: + return { + item.name: ColumnBinding(0, item.value.name) + for item in node.exprs + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(HStack) +def _(node: HStack) -> Mapping[str, ColumnBinding]: + child = node.children[0] + replaced = {item.name for item in node.columns} + return { + name: ColumnBinding(0, name) for name in child.schema if name not in replaced + } | { + item.name: ColumnBinding(0, item.value.name) + for item in node.columns + if isinstance(item.value, expr.Col) + } + + +@column_domain_bindings.register(GroupBy) +def _(node: GroupBy) -> Mapping[str, ColumnBinding]: + return { + key.name: ColumnBinding(0, key.value.name) + for key in node.keys + if isinstance(key.value, expr.Col) + } + + +@column_domain_bindings.register(Join) +def _(node: Join) -> Mapping[str, ColumnBinding]: + left, right = node.children + how = node.options[0] + if how in ("Semi", "Anti"): + return { + name: ColumnBinding(0, name) for name in node.schema if name in left.schema + } + if how != "Inner": + return {} + + bindings = {name: ColumnBinding(0, name) for name in left.schema} + suffix = node.options[3] + for name in right.schema: + output_name = f"{name}{suffix}" if name in left.schema else name + if output_name in node.schema: + bindings[output_name] = ColumnBinding(1, name) + return bindings + + +@column_domain_bindings.register(Distinct) +@column_domain_bindings.register(Filter) +@column_domain_bindings.register(Projection) +@column_domain_bindings.register(Slice) +@column_domain_bindings.register(Sort) +def _( + node: Distinct | Filter | Projection | Slice | Sort, +) -> Mapping[str, ColumnBinding]: + child = node.children[0] + return { + name: ColumnBinding(0, name) for name in node.schema if name in child.schema + } diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 80359eab925d..b382f12099a4 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -734,11 +734,17 @@ def evaluate_on_rank( """ stats = allgather_stats(comm, ctx.br(), ir, config_options, py_executor) + lowering, node_map = lower_ir_graph_with_node_map( + ir, config_options, stats, rank=comm.rank, nranks=comm.nranks + ) + optimized = lowering.optimized + ir = lowering.lowered + partition_info = lowering.partition_info if config_options.executor.quent_context is not None: assert local_quent_context is not None - logical_plan_id = ir.get_stable_plan_id() + logical_plan_id = optimized.get_stable_plan_id() plan, ops, ports, logical_op_by_id = build_plan( - ir, + optimized, config_options, query=local_quent_context.context.query, plan_id=logical_plan_id, @@ -752,10 +758,6 @@ def evaluate_on_rank( local_quent_context.logger, plan, ops, ports ) - ir, partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats, rank=comm.rank, nranks=comm.nranks - ) - if comm.rank == 0: log_query_plan(ir, config_options) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index c959f24f98d9..295a081307fa 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -26,6 +26,7 @@ from cudf_polars.quent import QuentContext from cudf_polars.utils.config import ( DynamicPlanningOptions, + JoinFilterPushdownOptions, ParquetOptions, ) @@ -248,6 +249,14 @@ class StreamingOptions: Env: ``CUDF_POLARS__EXECUTOR__DYNAMIC_PLANNING``. Default: enabled. Category: executor. + join_filter_pushdown + Config for join filter pushdown optimizations, dict or + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions`. ``None`` + disables the rewrite. + Env: ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN`` and + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__*``. + Default: enabled. + Category: executor. sink_to_directory Whether multi-partition sink operations should write to a directory rather than a single file. The ``spmd``/``ray``/``dask`` engines @@ -346,6 +355,9 @@ class StreamingOptions: dynamic_planning: dict[str, Any] | DynamicPlanningOptions | None | Unspecified = ( _opt("executor") ) + join_filter_pushdown: ( + dict[str, Any] | JoinFilterPushdownOptions | None | Unspecified + ) = _opt("executor") sink_to_directory: bool | Unspecified = _opt( "executor", "CUDF_POLARS__EXECUTOR__SINK_TO_DIRECTORY", parse_boolean ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index d02077206739..256e45440d8d 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -38,8 +38,9 @@ import cudf_polars.dsl.tracing from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Cast, Col, NamedExpr, TemporalFunction -from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select +from cudf_polars.dsl.ir import Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope +from cudf_polars.dsl.utils.column_domain import column_domain_bindings from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk @@ -446,9 +447,8 @@ def _derived_ordering( def _select_column_targets(select: Select) -> dict[str, dict[str, None]]: old_to_new_names: defaultdict[str, dict[str, None]] = defaultdict(dict) - for ne in select.exprs: - if isinstance(ne.value, Col): - old_to_new_names[ne.value.name][ne.name] = None + for output_name, source in column_domain_bindings(select).items(): + old_to_new_names[source.name][output_name] = None return dict(old_to_new_names) @@ -603,7 +603,7 @@ def maybe_remap_partitioning( ), local=_remap_scheme_simple(ir, partitioning.local, ir.children[0]), ) - if isinstance(ir, (Cache, Join, Projection, Filter)): + if isinstance(ir, (Join, Projection, Filter)): child = child_ir if child_ir is not None else ir.children[0] return Partitioning( inter_rank=_remap_scheme_simple(ir, partitioning.inter_rank, child), diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 8e62e161cd7c..b585138d583a 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -123,8 +123,10 @@ def explain_query( if physical: with cm: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) - return _repr_ir_tree(lowered_ir, partition_info, stats=stats, config=config) + lowered = lower_ir_graph(ir, config, stats) + return _repr_ir_tree( + lowered.lowered, lowered.partition_info, stats=stats, config=config + ) else: if config.executor.name == "streaming": # Include row-count statistics for the logical plan @@ -150,7 +152,9 @@ def collect_partition_plan( with concurrent.futures.ThreadPoolExecutor() as executor: stats = collect_statistics(ir, config, executor) - lowered_ir, partition_info = lower_ir_graph(ir, config, stats) + lowered = lower_ir_graph(ir, config, stats) + lowered_ir = lowered.lowered + partition_info = lowered.partition_info seen: set[tuple] = set() rows: list[PartitionPlanRow] = [] @@ -755,7 +759,9 @@ def from_ir( if lowered: with cm: stats = collect_statistics(ir, config_options, executor) - ir, partition_info_d = lower_ir_graph(ir, config_options, stats) + lowering = lower_ir_graph(ir, config_options, stats) + ir = lowering.lowered + partition_info_d = lowering.partition_info partition_info_dict = {} nodes: dict[str, SerializableIRNode] = {} diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 70b1585ba348..63d76d6c328f 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -75,7 +75,7 @@ def _make_hash_join( partition_info, output_count, ) - # Always reconstruct in case children contain Cache nodes + # Reconstruct with the lowered and possibly shuffled children. ir = ir.reconstruct([left, right]) # Record new partitioning info diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py new file mode 100644 index 000000000000..5ba74515cf1b --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -0,0 +1,844 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Rewrite a plan, inserting prefilters in join DAGs. + +For a supported inner equijoin, this optimization tries to use the join-key +values produced by one input to reduce the size of the other input before +the original join. In relational notation, a simple rewrite is:: + + left join[left.key = right.key] right + + -> + + (left semijoin[left.key = right.key] project(right.key)) + join[left.key = right.key] right + +In this example, the right hand table is selected to pre-filter the left +table before performing the inner join. + +The implementation uses the following terms: + +``column lineage`` + A chain from a named output column towards columns in its input subplan. + Each step guarantees that every value in the output column also appears in + the referenced child column, although row order and multiplicity are not + preserved and the child may contain additional values. +``child edge`` + One particular parent-to-child position in the IR DAG. The same child node + may occur on more than one edge, so a lineage records child indices and a + rewrite follows the resulting edge path to change only the chosen + occurrence. +``target`` + The side of the join to filter. +``domain`` + The side of the join used to provide key values for the filtering of + ``target``. +``producer`` + A node on a column lineage, together with the column name at that node and + its edge path from the join input. So termed because it "produces" the + key values participating in the join. +``constraint domain`` + Selective values of another join key from the target input, used to reduce + the domain before deriving the values that will filter the target. +``simple candidate`` + A rewrite that projects one domain join key and uses it to filter the + corresponding target key directly. +``composite candidate`` + For a multi-key join, a rewrite that first semi-joins the domain using the + constraint domain, then projects the reduced domain's key used to filter + the target. + +Plan rewrite has three stages. ``analyze_plan`` gathers row estimates, +selective nodes, and column value-domain lineages. Candidate selection +consumes those facts and returns a decision. ``apply_candidate`` then +constructs the selected semi-join rewrite. + +Row estimates, selectivity propagation, thresholds, and candidate scores are +only heuristics for deciding whether a safe rewrite is likely to improve +execution. Poor estimates can choose an unprofitable rewrite. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from functools import singledispatch +from typing import TYPE_CHECKING, Any, Literal, TypeAlias, TypedDict + +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + IR, + ConditionalJoin, + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Rolling, + Scan, + Select, + Slice, + Sort, + Union, +) +from cudf_polars.dsl.tracing import Scope, log +from cudf_polars.dsl.traversal import ( + CachingVisitor, + post_traversal, + reuse_if_unchanged, + traversal, +) +from cudf_polars.dsl.utils.column_domain import ( + ColumnLineage, + ColumnRef, + column_domain_bindings, +) + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator, Mapping, Sequence + + from cudf_polars.streaming.base import StatsCollector + from cudf_polars.typing import GenericTransformer + from cudf_polars.utils.config import ConfigOptions, StreamingExecutor + + +@dataclass(frozen=True) +class _Producer: + """A subtree and its bound column names at an insertion point.""" + + node: IR + columns: tuple[str, ...] + rows: int + path: tuple[int, ...] = () + """Child-edge path from the candidate root to ``node``.""" + + @property + def column(self) -> str: + """First bound column in the producer.""" + return self.columns[0] + + +@dataclass(frozen=True) +class SimpleCandidate: + """A direct key-domain prefilter candidate.""" + + mode = "simple" + target_side: Literal["left", "right"] + target: _Producer + target_key: expr.Col + domain: _Producer + domain_key: expr.Col + + @property + def score(self) -> tuple[int, int, int]: + """Rank after composite candidates, then by domain size.""" + return (1, self.domain.rows, self.domain.rows) + + +@dataclass(frozen=True) +class CompositeCandidate: + """A key-domain prefilter constrained by another join key.""" + + mode = "composite" + target_side: Literal["left", "right"] + target: _Producer + target_key: expr.Col + domain: _Producer + domain_key: expr.Col + constraint_domain: _Producer + domain_constraint_key: expr.Col + target_constraint_key: expr.Col + + @property + def score(self) -> tuple[int, int, int]: + """Prefer smaller constraint and domain inputs.""" + return (0, self.constraint_domain.rows, self.domain.rows) + + +Candidate: TypeAlias = SimpleCandidate | CompositeCandidate +DecisionReason: TypeAlias = Literal[ + "applied", + "maintain_order", + "no_selective_domain", + "non_column_join_key", + "not_inner_join", + "sliced_join", +] + + +@dataclass(frozen=True) +class Decision: + """Result of considering a join for a domain prefilter.""" + + reason: DecisionReason + candidate: Candidate | None = None + + +@dataclass(frozen=True) +class PlanFacts: + """Facts derived in one bottom-up traversal of an IR DAG.""" + + row_estimates: Mapping[IR, int | None] + selective_nodes: frozenset[IR] + column_lineages: Mapping[ColumnRef, ColumnLineage] + + +class _RewriteState(TypedDict): + """State shared by the join-domain prefilter DAG rewrite.""" + + threshold: float + trace: bool + stats: StatsCollector + facts: PlanFacts + + +def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: + """ + Derive row, selectivity, and column-domain facts for an IR DAG. + + Parameters + ---------- + ir + Root node to gather facts for. + stats + Pre-populated statistics + + Returns + ------- + Gather facts about the plan. + """ + row_estimates: dict[IR, int | None] = {} + selective_nodes: set[IR] = set() + column_lineages: dict[ColumnRef, ColumnLineage] = {} + + for node in post_traversal([ir]): + if isinstance(node, (Scan, DataFrameScan)): + source_info = stats.scan_stats.get(node) + rows = None if source_info is None else source_info.row_count + if rows is None and isinstance(node, DataFrameScan): + rows = node.df.shape()[0] + elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): + rows = row_estimates[node.children[0]] + elif isinstance(node, Join): + rows = _estimate_join_rows( + node.options[0], + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) + else: + child_estimates = [ + estimate + for child in node.children + if (estimate := row_estimates[child]) is not None + ] + rows = max(child_estimates, default=None) + row_estimates[node] = rows + + if ( + (isinstance(node, Scan) and node.predicate is not None) + or isinstance(node, Filter) + or any(child in selective_nodes for child in node.children) + ): + selective_nodes.add(node) + + bindings = column_domain_bindings(node) + for name in node.schema: + column = ColumnRef(node, name) + binding = bindings.get(name) + if binding is None: + source_lineage = None + source_child_index = None + else: + source_child_index = binding.child_index + source = ColumnRef( + node.children[source_child_index], + binding.name, + ) + source_lineage = column_lineages[source] + column_lineages[column] = ColumnLineage( + column, source_lineage, source_child_index + ) + + return PlanFacts( + row_estimates=row_estimates, + selective_nodes=frozenset(selective_nodes), + column_lineages=column_lineages, + ) + + +def blocks_pushdown(node: IR) -> bool: + """ + Return whether a node blocks filter pushdown. + + Parameters + ---------- + node + Node to check + + Returns + ------- + bool + True if a semijoin cannot be pushed past this node, otherwise False. + """ + return ( + # TODO: Distinct and Rolling only block pushdown in some + # circumstances, but we'd need to make the logic more complicated: + # - We can push through distinct if the filter applies to the columns + # that are being used to determine distinct rows + # - We can push through rolling if the filter applies to the + # groupby keys. + # TODO: We can push through an unsliced Union, but need to + # distribute the filter onto every child. + isinstance(node, (Distinct, Rolling, Slice, Union)) + # Can't push through anything that is sliced. + or (isinstance(node, (GroupBy, Sort)) and node.zlice is not None) + or (isinstance(node, (ConditionalJoin, Join)) and node.options[2] is not None) + ) + + +def semijoin_pushdown_candidates( + facts: PlanFacts, root: IR, column: str +) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: + """ + Yield column domain lineage providing valid locations for semijoin pushdown. + + Parameters + ---------- + facts + Gathered facts about the plan + root + Root node to search from + column + Name of column we're finding the lineage of. + + Returns + ------- + Iterator + Of valid insertion points and their child-edge paths from ``root``. + """ + try: + lineage = facts.column_lineages[ColumnRef(root, column)] + except KeyError: + return + path: tuple[int, ...] = () + while True: + yield lineage.column, path + source = lineage.source + source_child_index = lineage.source_child_index + if blocks_pushdown(lineage.column.node) or source is None: + return + assert source_child_index is not None + path = (*path, source_child_index) + lineage = source + + +def optimize_join_filter_pushdown( + ir: IR, + stats: StatsCollector, + config_options: ConfigOptions[StreamingExecutor], +) -> IR: + """ + Rewrite an IR DAG to apply filter pushdown of keys. + + This optimization pass inspects joins in the DAG and attempts to push a + prefilter obtained from the keys of one side of the join onto the + inputs of the other side. This can be highly beneficial at large scale + since if we have a selective join we can avoid data movement by + prefiltering before performing the actual join. + + Parameters + ---------- + ir + DAG to rewrite. + stats + Pre-populated statistics. + config_options + Configuration options controlling the rewrite. + + Returns + ------- + Rewritten DAG. + """ + options = config_options.executor.join_filter_pushdown + if options is None: + return ir + threshold = options.threshold + trace = options.trace + if threshold == 0: + return ir + + state = _RewriteState( + threshold=threshold, + trace=trace, + stats=stats, + facts=analyze_plan(ir, stats), + ) + mapper: GenericTransformer[IR, IR, _RewriteState] = CachingVisitor( + _rewrite, state=state + ) + return mapper(ir) + + +@singledispatch +def _rewrite(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + raise AssertionError + + +@_rewrite.register(IR) +def _(node: IR, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + return reuse_if_unchanged(node, rec) + + +@_rewrite.register(Join) +def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: + original = node + rewritten = reuse_if_unchanged(node, rec) + assert isinstance(rewritten, Join) + node = rewritten + if node is original: + facts = rec.state["facts"] + else: + # Child rewrites introduce new semi joins and reconstructed ancestors. + # Re-analyze that current subtree so parent joins can use the derived + # selectivity and cardinality when ranking their own candidates. + facts = analyze_plan(node, rec.state["stats"]) + decision = _select_candidate( + node, + rec.state["threshold"], + facts, + ) + if rec.state["trace"]: + _trace_decision(node, rec.state["threshold"], decision) + if decision.candidate is None: + return node + return apply_candidate(node, decision.candidate) + + +def apply_candidate(ir: Join, candidate: Candidate) -> IR: + """Apply a selected join-domain prefilter candidate to a join.""" + left, right = ir.children + domain = _make_domain(candidate, ir) + target = candidate.target + target_filter = _make_semi_join( + target.node, + expr.Col(target.node.schema[target.column], target.column), + domain, + expr.Col(domain.schema[candidate.domain_key.name], candidate.domain_key.name), + nulls_equal=ir.options[1], + suffix=ir.options[3], + ) + if candidate.target_side == "left": + left = replace_at_path(left, target.path, target_filter) + else: + right = replace_at_path(right, target.path, target_filter) + return ir.reconstruct((left, right)) + + +def replace_at_path(root: IR, path: Sequence[int], replacement: IR) -> IR: + """ + Replace a specific child in a DAG starting at root. + + Parameters + ---------- + root + Root of DAG to carry out replacement. + path + Breadcrumb trail selecting which child at every level to recurse + into. + replacement + Replacement node to return when the path becomes empty. + + Returns + ------- + IR + New DAG with the selected child replaced with replacement. + + Notes + ----- + This specifically does not use replacement by equality so that we can + disambiguate between shared children in the DAG where we only want to + replace one. + """ + if not path: + return replacement + index, *path = path + children = list(root.children) + children[index] = replace_at_path(children[index], path, replacement) + return root.reconstruct(children) + + +def _select_candidate( + ir: Join, + threshold: float, + facts: PlanFacts, +) -> Decision: + if ir.options[0] != "Inner": + return Decision(reason="not_inner_join") + if ir.options[2] is not None: + return Decision(reason="sliced_join") + if ir.options[5] != "none": + return Decision(reason="maintain_order") + + left_keys = _simple_keys(ir.left_on) + right_keys = _simple_keys(ir.right_on) + if len(left_keys) != len(ir.left_on) or len(right_keys) != len(ir.right_on): + return Decision(reason="non_column_join_key") + + candidates: list[Candidate] = [] + left: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "left", + ir.children[0], + left_keys, + ) + right: tuple[Literal["left", "right"], IR, tuple[expr.Col, ...]] = ( + "right", + ir.children[1], + right_keys, + ) + for (target_side, target_child, target_keys), ( + _, + domain_child, + domain_keys, + ) in ((left, right), (right, left)): + candidates.extend( + _composite_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + facts, + ) + ) + candidates.extend( + _simple_candidates( + target_side, + target_child, + domain_child, + target_keys, + domain_keys, + threshold, + facts, + ) + ) + + if not candidates: + return Decision(reason="no_selective_domain") + return Decision(reason="applied", candidate=min(candidates, key=lambda c: c.score)) + + +def _simple_keys(keys: Sequence[expr.NamedExpr]) -> tuple[expr.Col, ...]: + return tuple(key.value for key in keys if isinstance(key.value, expr.Col)) + + +def _simple_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], + threshold: float, + facts: PlanFacts, +) -> Iterable[SimpleCandidate]: + for target_key, domain_key in zip(target_keys, domain_keys, strict=True): + target = _largest_key_source(target_child, target_key.name, facts) + if target is None: + continue + domain = _smallest_key_producer( + domain_child, + domain_key.name, + facts, + require_selective=True, + ) + if domain is None: + continue + if domain.rows / target.rows > threshold: + continue + if contains_node(target.node, domain.node): + continue + yield SimpleCandidate( + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + ) + + +def _composite_candidates( + target_side: Literal["left", "right"], + target_child: IR, + domain_child: IR, + target_keys: tuple[expr.Col, ...], + domain_keys: tuple[expr.Col, ...], + threshold: float, + facts: PlanFacts, +) -> Iterable[CompositeCandidate]: + if len(target_keys) < 2: + return + + for filter_index, (target_key, domain_key) in enumerate( + zip(target_keys, domain_keys, strict=True) + ): + target = _largest_key_source(target_child, target_key.name, facts) + if target is None: + continue + + for constraint_index, ( + target_constraint_key, + domain_constraint_key, + ) in enumerate(zip(target_keys, domain_keys, strict=True)): + if constraint_index == filter_index: + continue + domain = _smallest_node_containing_all( + domain_child, + (domain_key.name, domain_constraint_key.name), + facts, + ) + if domain is None: + continue + if domain.rows / target.rows > threshold: + continue + constraint_domain = _smallest_key_producer( + target_child, + target_constraint_key.name, + facts, + require_selective=True, + exclude=target.node, + ) + if constraint_domain is None: + continue + if constraint_domain.rows / domain.rows > threshold: + continue + if contains_node(target.node, domain.node) or contains_node( + target.node, constraint_domain.node + ): + continue + yield CompositeCandidate( + target_side=target_side, + target=target, + target_key=target_key, + domain=domain, + domain_key=domain_key, + constraint_domain=constraint_domain, + domain_constraint_key=domain_constraint_key, + target_constraint_key=target_constraint_key, + ) + + +def _make_domain(candidate: Candidate, ir: Join) -> IR: + if isinstance(candidate, SimpleCandidate): + return _project_bound_key( + candidate.domain.node, + candidate.domain.column, + candidate.domain_key, + ) + + constraint_domain = _project_bound_key( + candidate.constraint_domain.node, + candidate.constraint_domain.column, + candidate.target_constraint_key, + ) + constrained = _make_semi_join( + candidate.domain.node, + expr.Col( + candidate.domain.node.schema[candidate.domain.columns[1]], + candidate.domain.columns[1], + ), + constraint_domain, + expr.Col( + constraint_domain.schema[candidate.target_constraint_key.name], + candidate.target_constraint_key.name, + ), + nulls_equal=ir.options[1], + suffix=ir.options[3], + ) + return _project_bound_key( + constrained, candidate.domain.column, candidate.domain_key + ) + + +def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> Select: + """Project a bound source column under its join-visible key name.""" + dtype = source.schema[bound_column] + assert dtype == output_key.dtype + return Select( + {output_key.name: dtype}, + (expr.NamedExpr(output_key.name, expr.Col(dtype, bound_column)),), + True, # noqa: FBT003 + source, + ) + + +def _make_semi_join( + target: IR, + target_key: expr.Col, + domain: IR, + domain_key: expr.Col, + *, + nulls_equal: bool, + suffix: str, +) -> Join: + return Join( + target.schema, + (expr.NamedExpr(target_key.name, target_key),), + (expr.NamedExpr(domain_key.name, domain_key),), + ("Semi", nulls_equal, None, suffix, False, "none"), + target, + domain, + ) + + +def _smallest_key_producer( + root: IR, + column: str, + facts: PlanFacts, + *, + require_selective: bool, + exclude: IR | None = None, +) -> _Producer | None: + candidates = [] + for reference, path in semijoin_pushdown_candidates(facts, root, column): + node, bound_column = reference.node, reference.name + if node is exclude: + continue + rows = facts.row_estimates.get(node) + if rows is None or rows <= 0: + continue + if require_selective and node not in facts.selective_nodes: + continue + candidates.append( + (rows, len(node.schema), _Producer(node, (bound_column,), rows, path)) + ) + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _smallest_node_containing_all( + root: IR, columns: Sequence[str], facts: PlanFacts +) -> _Producer | None: + candidates = [] + lineages: list[ColumnLineage] = [] + for column in columns: + lineage = facts.column_lineages.get(ColumnRef(root, column)) + if lineage is None: + return None + lineages.append(lineage) + if not lineages: + return None + path: tuple[int, ...] = () + while True: + node = lineages[0].column.node + if any(lineage.column.node != node for lineage in lineages[1:]): + break + bound_columns = tuple(lineage.column.name for lineage in lineages) + rows = facts.row_estimates.get(node) + if rows is not None and rows > 0: + candidates.append( + ( + rows, + len(node.schema), + _Producer(node, bound_columns, rows, path), + ) + ) + if blocks_pushdown(node): + break + source_child_index = lineages[0].source_child_index + if source_child_index is None or any( + lineage.source_child_index != source_child_index for lineage in lineages[1:] + ): + break + sources = [lineage.source for lineage in lineages if lineage.source is not None] + if len(sources) != len(lineages): + # Some sources are None + break + path = (*path, source_child_index) + lineages = sources + if not candidates: + return None + return min(candidates, key=lambda item: (item[0], item[1]))[2] + + +def _largest_key_source(root: IR, column: str, facts: PlanFacts) -> _Producer | None: + source_candidates = [] + fallback_candidates = [] + for reference, path in semijoin_pushdown_candidates(facts, root, column): + node, bound_column = reference.node, reference.name + rows = facts.row_estimates.get(node) + if rows is None or rows <= 0: + continue + item = ( + rows, + len(node.schema), + _Producer(node, (bound_column,), rows, path), + ) + if isinstance(node, (Scan, DataFrameScan)): + source_candidates.append(item) + else: + fallback_candidates.append(item) + candidates = source_candidates or fallback_candidates + if not candidates: + return None + return max(candidates, key=lambda item: (item[0], -item[1]))[2] + + +def _estimate_join_rows( + how: str, left_rows: int | None, right_rows: int | None +) -> int | None: + if left_rows is None: + return right_rows + if right_rows is None: + return left_rows + if how in ("Inner", "Semi", "Anti"): + return min(left_rows, right_rows) + if how == "Left": + return left_rows + if how == "Right": + return right_rows + if how == "Full": + return max(left_rows, right_rows) + return None + + +def contains_node(root: IR, needle: IR) -> bool: + """Return whether an equal node occurs in a DAG rooted at ``root``.""" + return needle in traversal([root]) + + +def _trace_decision(ir: Join, threshold: float, decision: Decision) -> None: + join_filter_pushdown: dict[str, Any] = { + "considered": True, + "threshold": threshold, + "reason": decision.reason, + } + record = { + "scope": Scope.PLAN.value, + "join_filter_pushdown": join_filter_pushdown, + "actor_ir_id": ir.get_stable_id(), + "actor_ir_type": type(ir).__name__, + } + if (candidate := decision.candidate) is not None: + join_filter_pushdown.update( + { + "mode": candidate.mode, + "target_side": candidate.target_side, + "target_key": candidate.target_key.name, + "domain_key": candidate.domain_key.name, + "estimated_target_rows": candidate.target.rows, + "estimated_domain_rows": candidate.domain.rows, + "target_node_type": type(candidate.target.node).__name__, + "domain_node_type": type(candidate.domain.node).__name__, + } + ) + if isinstance(candidate, CompositeCandidate): + join_filter_pushdown.update( + { + "constraint_key": candidate.target_constraint_key.name, + "estimated_constraint_rows": candidate.constraint_domain.rows, + } + ) + log("Join Filter Pushdown", **record) diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index 295798116d26..2ac5c8c2eef7 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -4,8 +4,9 @@ from __future__ import annotations +import dataclasses import operator -from functools import partial, reduce +from functools import reduce from typing import TYPE_CHECKING import polars as pl @@ -35,7 +36,7 @@ Slice, Union, ) -from cudf_polars.dsl.traversal import CachingVisitor, traversal +from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged, traversal from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node @@ -52,6 +53,7 @@ from cudf_polars.streaming.base import StatsCollector from cudf_polars.streaming.dispatch import LowerIRTransformer, State + from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -65,6 +67,64 @@ def _( ) +@lower_ir_node.register(Cache) +def _( + ir: Cache, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: # pragma: no cover + raise AssertionError("Cache nodes should have been removed before lowering") + + +@dataclasses.dataclass +class LoweringInfo: + """Information produced by optimizing and lowering an IR graph.""" + + optimized: IR # IR after optimization + lowered: IR # optimized IR after lowering + partition_info: MutableMapping[ + IR, PartitionInfo + ] # Partition mapping for nodes in the lowered IR. + + +def remove_cache_nodes(ir: IR) -> IR: + """Remove logical cache nodes while preserving shared DAG structure.""" + + def rewrite(node: IR, rec: GenericTransformer[IR, IR, None]) -> IR: + if isinstance(node, Cache): + return rec(node.children[0]) + return reuse_if_unchanged(node, rec) + + mapper: GenericTransformer[IR, IR, None] = CachingVisitor(rewrite, state=None) + return mapper(ir) + + +def optimize_with_stats( + ir: IR, config_options: ConfigOptions[StreamingExecutor], stats: StatsCollector +) -> IR: + """ + Optimize an IR graph given some statistics. + + Parameters + ---------- + ir + Root of the graph to optimize. + config_options + GPUEngine configuration options. + stats + Pre-computed statistics. + + Returns + ------- + IR + The optimized IR graph. + """ + from cudf_polars.streaming.join_filter_pushdown import ( + optimize_join_filter_pushdown, + ) + + ir = remove_cache_nodes(ir) + return optimize_join_filter_pushdown(ir, stats, config_options) + + def _lower_ir_graph_impl( ir: IR, config_options: ConfigOptions[StreamingExecutor], @@ -72,15 +132,19 @@ def _lower_ir_graph_impl( *, rank: int = 0, nranks: int = 1, -) -> tuple[tuple[IR, MutableMapping[IR, PartitionInfo]], LowerIRTransformer]: +) -> tuple[LoweringInfo, LowerIRTransformer]: state: State = { "config_options": config_options, "stats": stats, "rank": rank, "nranks": nranks, } + optimized = optimize_with_stats(ir, config_options, stats) mapper: LowerIRTransformer = CachingVisitor(lower_ir_node, state=state) - return mapper(ir), mapper + lowered, partition_info = mapper(optimized) + return LoweringInfo( + optimized=optimized, lowered=lowered, partition_info=partition_info + ), mapper def lower_ir_graph( @@ -90,7 +154,7 @@ def lower_ir_graph( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: +) -> LoweringInfo: """ Rewrite an IR graph and extract partitioning information. @@ -109,9 +173,7 @@ def lower_ir_graph( Returns ------- - new_ir, partition_info - The rewritten graph and a mapping from unique nodes - in the new graph to associated partitioning information. + LoweringInfo Notes ----- @@ -132,7 +194,7 @@ def lower_ir_graph_with_node_map( *, rank: int = 0, nranks: int = 1, -) -> tuple[IR, MutableMapping[IR, PartitionInfo], dict[str, list[str]]]: +) -> tuple[LoweringInfo, dict[str, list[str]]]: """ Lower an IR graph and return a mapping from physical to logical stable IDs. @@ -155,10 +217,8 @@ def lower_ir_graph_with_node_map( Returns ------- - new_ir - The rewritten IR graph. - partition_info - Mapping from unique nodes in the new graph to partitioning info. + LoweringInfo + Information about the lowered IR graph. node_map Mapping ``{physical_stable_id: [logical_stable_id, ...]}`` built from the internal :class:`CachingVisitor` cache. Nodes inserted @@ -173,7 +233,7 @@ def lower_ir_graph_with_node_map( old_key = str(old_node.get_stable_id()) node_map.setdefault(new_key, []).append(old_key) - return *result, node_map + return result, node_map def evaluate_streaming( @@ -282,8 +342,6 @@ def _lower_ir_pwise( return new_node, partition_info -_lower_ir_pwise_preserve = partial(_lower_ir_pwise, preserve_partitioning=True) -lower_ir_node.register(Cache, _lower_ir_pwise_preserve) lower_ir_node.register(HConcat, _lower_ir_pwise) diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 3451010b464a..8a265d875b39 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -54,6 +54,7 @@ "DaskContext", "DynamicPlanningOptions", "InMemoryExecutor", + "JoinFilterPushdownOptions", "ParquetOptions", "RayContext", "SPMDContext", @@ -362,8 +363,9 @@ class DynamicPlanningOptions: The maximum number of chunks to sample before making dynamic-planning decisions. Default is 2. join_prefilter_threshold - Row-count ratio (small / large) below which a join key prefilter is - applied. Set to 0 to disable join prefiltering. Default is 0.5. + Row-count ratio (small / large) below which one side of a join is + filtered by a bloom filter built from the other side before + performing the join. Set to 0 to disable. Default is 0.5. join_prefilter_max_key_columns Maximum number of columns from the join-key prefix to use for the prefilter. Set to ``None`` to use the full join-key list. Default is 1. @@ -428,6 +430,59 @@ def __post_init__(self) -> None: # noqa: D105 raise TypeError("join_prefilter_trace must be a bool") +@dataclasses.dataclass(frozen=True) +class JoinFilterPushdownOptions: + """ + Configuration options for join filter pushdown in the logical plan. + + When performing a join between two tables, it is often favourable + to pre-filter one side of the join with the keys (full or partial) of + the other side. This can reduce the size of tables that actually + participate in the join. + + cudf-polars supports a form of this where we can rewrite inner joins by + selecting a side to be filtered by the keys of the other side. + + Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to + disable the rewrite. + + These options can be configured via environment variables with the prefix + ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__``. + + Parameters + ---------- + threshold + Row-count ratio (key-provider-rows / to-be-filtered-table-rows) below which a + filter on is inserted on the to-be-filtered table. Default is 0.5. + trace + Whether to emit plan-time trace decisions for filter decisions. Default is False. + """ + + _env_prefix = "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN" + + threshold: float = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__THRESHOLD", float, default=0.5 + ) + ) + trace: bool = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__TRACE", _bool_converter, default=False + ) + ) + + def __post_init__(self) -> None: # noqa: D105 + threshold = self.threshold + if isinstance(threshold, bool) or not isinstance(threshold, (int, float)): + raise TypeError("threshold must be a float or int") + threshold = float(threshold) + object.__setattr__(self, "threshold", threshold) + if not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be between 0 and 1") + if not isinstance(self.trace, bool): + raise TypeError("trace must be a bool") + + @dataclasses.dataclass(frozen=True, eq=True) class MemoryResourceConfig: """ @@ -693,6 +748,10 @@ class StreamingExecutor: dynamic_planning Options controlling dynamic shuffle planning. See :class:`~cudf_polars.utils.config.DynamicPlanningOptions` for more. + join_filter_pushdown + Options controlling the logical join-domain prefilter rewrite. See + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for more. + ``None`` disables the rewrite. max_io_threads Maximum number of IO threads. Default is 4. This controls the parallelism of IO operations when reading data. @@ -762,6 +821,9 @@ class StreamingExecutor: dynamic_planning: DynamicPlanningOptions | None = dataclasses.field( default_factory=DynamicPlanningOptions ) + join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( + default_factory=JoinFilterPushdownOptions + ) max_io_threads: int = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__MAX_IO_THREADS", int, default=4 @@ -823,6 +885,20 @@ def __post_init__(self) -> None: # noqa: D105 DynamicPlanningOptions(**self.dynamic_planning), ) + if isinstance(self.join_filter_pushdown, dict): + object.__setattr__( + self, + "join_filter_pushdown", + JoinFilterPushdownOptions(**self.join_filter_pushdown), + ) + if self.join_filter_pushdown is not None and not isinstance( + self.join_filter_pushdown, JoinFilterPushdownOptions + ): + raise TypeError( + "join_filter_pushdown must be a JoinFilterPushdownOptions " + "instance, dict, or None" + ) + if self.cluster in ("spmd", "ray", "dask"): if self.sink_to_directory is False: raise ValueError( @@ -855,6 +931,7 @@ def __hash__(self) -> int: # noqa: D105 # to json and hash that. d = dataclasses.asdict(self) d["dynamic_planning"] = json.dumps(d["dynamic_planning"]) + d["join_filter_pushdown"] = json.dumps(d["join_filter_pushdown"]) # Hash the quent context UUIDs as ints quent_context = d["quent_context"] @@ -1019,6 +1096,17 @@ def from_polars_engine( if not _bool_converter(env_dynamic_planning): user_executor_options["dynamic_planning"] = None + # Handle join_filter_pushdown: check user config, then env var + user_join_filter_pushdown = user_executor_options.get( + "join_filter_pushdown", None + ) + if user_join_filter_pushdown is None: + env_join_filter_pushdown = os.environ.get( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1" + ) + if not _bool_converter(env_join_filter_pushdown): + user_executor_options["join_filter_pushdown"] = None + executor = StreamingExecutor(**user_executor_options) case _: # pragma: no cover; Unreachable raise ValueError(f"Unsupported executor: {user_executor}") diff --git a/python/cudf_polars/tests/dsl/test_column_domain.py b/python/cudf_polars/tests/dsl/test_column_domain.py new file mode 100644 index 000000000000..999f3f2463d2 --- /dev/null +++ b/python/cudf_polars/tests/dsl/test_column_domain.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import DataType +from cudf_polars.dsl import expr +from cudf_polars.dsl.ir import ( + DataFrameScan, + Distinct, + Filter, + GroupBy, + HStack, + Join, + Projection, + Select, + Slice, + Sort, +) +from cudf_polars.dsl.utils.column_domain import ( + ColumnBinding, + column_domain_bindings, +) + +I64 = DataType(pl.Int64()) +BOOL = DataType(pl.Boolean()) + + +def make_scan(*names: str) -> DataFrameScan: + frame = pl.DataFrame({name: [1] for name in names}) + return DataFrameScan(dict.fromkeys(names, I64), frame._df, None) + + +def col(name: str) -> expr.Col: + return expr.Col(I64, name) + + +def named_col(output: str, source: str) -> expr.NamedExpr: + return expr.NamedExpr(output, col(source)) + + +def test_source_has_no_column_domain_bindings() -> None: + assert column_domain_bindings(make_scan("a")) == {} + + +def test_select_binds_aliases_and_omits_derived_columns() -> None: + child = make_scan("a", "b") + node = Select( + {"renamed": I64, "derived": I64}, + ( + named_col("renamed", "a"), + expr.NamedExpr("derived", expr.Literal(I64, 1)), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "renamed": ColumnBinding(0, "a"), + } + + +def test_hstack_binds_passthrough_alias_and_override() -> None: + child = make_scan("a", "b") + node = HStack( + {"a": I64, "b": I64, "alias": I64}, + ( + expr.NamedExpr("a", expr.Literal(I64, 1)), + named_col("alias", "b"), + ), + True, # noqa: FBT003 + child, + ) + + assert column_domain_bindings(node) == { + "b": ColumnBinding(0, "b"), + "alias": ColumnBinding(0, "b"), + } + + +def test_groupby_binds_only_direct_keys() -> None: + child = make_scan("a", "b") + node = GroupBy( + {"key": I64, "value": I64}, + (named_col("key", "a"),), + (named_col("value", "b"),), + False, # noqa: FBT003 + None, + child, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "a"), + } + + +def test_inner_join_binds_left_right_and_suffixed_columns() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + { + "key": I64, + "left_value": I64, + "key_right": I64, + "right_value": I64, + }, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, (0, 1), "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "key_right": ColumnBinding(1, "key"), + "right_value": ColumnBinding(1, "right_value"), + } + + +def test_inner_join_omits_coalesced_right_key() -> None: + left = make_scan("key", "left_value") + right = make_scan("key", "right_value") + node = Join( + {"key": I64, "left_value": I64, "right_value": I64}, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Inner", False, None, "_right", True, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "left_value": ColumnBinding(0, "left_value"), + "right_value": ColumnBinding(1, "right_value"), + } + + +def test_semi_join_binds_only_left_columns() -> None: + left = make_scan("key", "value") + right = make_scan("key") + node = Join( + left.schema, + (named_col("key", "key"),), + (named_col("key", "key"),), + ("Semi", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == { + "key": ColumnBinding(0, "key"), + "value": ColumnBinding(0, "value"), + } + + +def test_outer_join_has_no_column_domain_bindings() -> None: + left = make_scan("key") + right = make_scan("other") + node = Join( + {**left.schema, **right.schema}, + (named_col("key", "key"),), + (named_col("other", "other"),), + ("Left", False, None, "_right", False, "none"), + left, + right, + ) + + assert column_domain_bindings(node) == {} + + +def test_passthrough_nodes_bind_same_named_columns() -> None: + child = make_scan("a", "b") + mask = expr.NamedExpr("mask", expr.Literal(BOOL, True)) # noqa: FBT003 + nodes = ( + Filter(child.schema, mask, child), + Projection({"b": I64}, child), + Slice(child.schema, 0, 1, child), + Distinct( + child.schema, + plc.stream_compaction.DuplicateKeepOption.KEEP_ANY, + None, + (0, 1), + False, # noqa: FBT003 + child, + ), + Sort( + child.schema, + (named_col("a", "a"),), + (plc.types.Order.ASCENDING,), + (plc.types.NullOrder.AFTER,), + False, # noqa: FBT003 + (0, 1), + child, + ), + ) + + for node in nodes: + assert column_domain_bindings(node) == { + name: ColumnBinding(0, name) for name in node.schema + } diff --git a/python/cudf_polars/tests/quent/test_quent.py b/python/cudf_polars/tests/quent/test_quent.py index be1aedc770b8..2b7217d0689c 100644 --- a/python/cudf_polars/tests/quent/test_quent.py +++ b/python/cudf_polars/tests/quent/test_quent.py @@ -449,10 +449,9 @@ def test_lower_ir_graph_with_node_map() -> None: ir, config_options, concurrent.futures.ThreadPoolExecutor() ) - _lowered_ir, _partition_info, node_map = lower_ir_graph_with_node_map( - ir, config_options, stats - ) + lowering, node_map = lower_ir_graph_with_node_map(ir, config_options, stats) + assert lowering.optimized is ir assert len(node_map) > 0 for physical_sid, logical_sids in node_map.items(): assert isinstance(physical_sid, str) diff --git a/python/cudf_polars/tests/streaming/test_dataframescan.py b/python/cudf_polars/tests/streaming/test_dataframescan.py index 014851f689e7..0fd27f2b57e8 100644 --- a/python/cudf_polars/tests/streaming/test_dataframescan.py +++ b/python/cudf_polars/tests/streaming/test_dataframescan.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -59,7 +59,7 @@ def test_parallel_dataframescan( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -68,6 +68,8 @@ def test_parallel_dataframescan( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if max_rows_per_partition < total_row_count: assert count > 1 @@ -106,7 +108,7 @@ def test_join_in_memory_lazy_stable_id_pickle( right = pl.LazyFrame({"k": [2, 3, 4], "y": [1, 2, 3]}).collect(engine=engine).lazy() qir = Translator(left.join(right, on="k")._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -115,6 +117,7 @@ def test_join_in_memory_lazy_stable_id_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered _assert_stable_ids_match(ir, pickle.loads(pickle.dumps(ir))) @@ -128,7 +131,7 @@ def test_dataframescan_pickle( ) qir = Translator(df._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, _ = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -137,6 +140,7 @@ def test_dataframescan_pickle( parquet_stats_executor, ), ) + ir = lowering.lowered # Pickle and unpickle the IR (which contains DataFrameScan) pickled = pickle.dumps(ir) diff --git a/python/cudf_polars/tests/streaming/test_hstack.py b/python/cudf_polars/tests/streaming/test_hstack.py index f7a3c6213340..232ad33e255e 100644 --- a/python/cudf_polars/tests/streaming/test_hstack.py +++ b/python/cudf_polars/tests/streaming/test_hstack.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 """Tests for CSE HStack handling in the streaming executor.""" @@ -150,14 +150,16 @@ def test_cse_agg_shared_decomposition( assert len(inner_hstacks) == (1 if comm_subexpr_elim else 0) config_options = ConfigOptions.from_polars_engine(engine) - lowered, _ = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) # Both paths must lower to a single Repartition computing one aggregation. - repartitions = [n for n in traversal([lowered]) if isinstance(n, Repartition)] + repartitions = [ + n for n in traversal([lowering.lowered]) if isinstance(n, Repartition) + ] assert len(repartitions) == 1 assert len(repartitions[0].children[0].exprs) == 1 # type: ignore[attr-defined] assert_gpu_result_equal(q, engine=engine, collect_kwargs={"optimizations": opts}) diff --git a/python/cudf_polars/tests/streaming/test_join.py b/python/cudf_polars/tests/streaming/test_join.py index cdc544c43169..2ac2edc88903 100644 --- a/python/cudf_polars/tests/streaming/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -491,18 +491,17 @@ def test_broadcast_limit( q = left.join(right, on="y", how="inner") ir = Translator(q._ldf.visit(), engine).translate_ir() config_options = ConfigOptions.from_polars_engine(engine) - shuffle_nodes = [ - type(node) - for node in lower_ir_graph( + lowering = lower_ir_graph( + ir, + config_options, + collect_statistics( ir, config_options, - collect_statistics( - ir, - config_options, - parquet_stats_executor, - ), - )[1] - if isinstance(node, Shuffle) + parquet_stats_executor, + ), + ) + shuffle_nodes = [ + type(node) for node in lowering.partition_info if isinstance(node, Shuffle) ] # NOTE: Expect small table to have 3 partitions (9 / 3). @@ -516,7 +515,7 @@ def test_broadcast_limit( assert len(shuffle_nodes) == 0 -def test_cache_preserves_partitioning_join( +def test_shared_join_preserves_partitioning( parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, ): engine = pl.GPUEngine( @@ -542,20 +541,24 @@ def test_cache_preserves_partitioning_join( config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() - lowered_ir, partition_info = lower_ir_graph( + lowering = lower_ir_graph( ir, config_options, collect_statistics(ir, config_options, parquet_stats_executor), ) + lowered_ir = lowering.lowered + partition_info = lowering.partition_info + + assert not any(isinstance(node, Cache) for node in traversal([lowered_ir])) - # Cache should preserve partitioning on 'key' - cache_partitioning = [ + # Removing Cache should preserve the shared join's partitioning on 'key'. + join_partitioning = [ [ne.name for ne in partition_info[node].partitioned_on] for node in traversal([lowered_ir]) - if isinstance(node, Cache) + if isinstance(node, Join) ] - assert cache_partitioning == [["key"]], ( - f"Cache should preserve partitioning on 'key', got {cache_partitioning}" + assert join_partitioning == [["key"]], ( + f"Shared join should be partitioned on 'key', got {join_partitioning}" ) # Only 2 shuffles needed (for join sides, not for groupby) diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py new file mode 100644 index 000000000000..3dfe42df7fd3 --- /dev/null +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -0,0 +1,834 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +import polars as pl + +from cudf_polars import Translator +from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice +from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.utils.column_domain import ColumnRef +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.join_filter_pushdown import ( + CompositeCandidate, + Decision, + PlanFacts, + SimpleCandidate, + _select_candidate, + _smallest_node_containing_all, + analyze_plan, + apply_candidate, + contains_node, + optimize_join_filter_pushdown, + semijoin_pushdown_candidates, +) +from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes +from cudf_polars.streaming.statistics import collect_statistics +from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.config import ConfigOptions + +if TYPE_CHECKING: + import concurrent.futures + from typing import Any + + from cudf_polars.dsl.ir import IR + from cudf_polars.engine.spmd import SPMDEngine + + +@pytest.fixture +def engine(spmd_engine_factory) -> SPMDEngine: + """Return an SPMD engine configured for join-domain prefilter tests.""" + return spmd_engine_factory( + StreamingOptions( + join_filter_pushdown={"threshold": 0.5}, + raise_on_fail=True, + ) + ) + + +def make_config( + *, dynamic_planning: bool = True, join_filter_pushdown: bool = True +) -> ConfigOptions: + executor_options: dict[str, Any] = { + "join_filter_pushdown": {"trace": False} if join_filter_pushdown else None + } + if not dynamic_planning: + executor_options["dynamic_planning"] = None + return ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options=executor_options, + ) + ) + + +def find_joins(ir: IR, how: str | None = None) -> list[Join]: + return [ + node + for node in traversal([ir]) + if isinstance(node, Join) and (how is None or node.options[0] == how) + ] + + +def translate_query(query: pl.LazyFrame, engine: SPMDEngine) -> IR: + """Translate a public Polars query and remove logical Cache nodes.""" + t = Translator(query._ldf.visit(), engine) + root = t.translate_ir() + assert not t.errors + return remove_cache_nodes(root) + + +def dataframe_scan(ir: IR, column: str) -> DataFrameScan: + """Return the unique in-memory scan containing ``column``.""" + (match,) = ( + node + for node in traversal([ir]) + if isinstance(node, DataFrameScan) and column in node.schema + ) + return match + + +@pytest.fixture +def simple_query() -> pl.LazyFrame: + """Return a query with a small selective join domain.""" + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_partkey": [i % 10 for i in range(20)], + "l_suppkey": range(20), + } + ) + return part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + + +def test_simple_prefilter_filters_large_side( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) + + assert isinstance(root, Join) + part_ir, _ = root.children + lineitem_ir = dataframe_scan(root, "l_partkey") + facts = analyze_plan(root, StatsCollector()) + decision = _select_candidate(root, 0.5, facts) + + assert decision.reason == "applied" + assert isinstance(decision.candidate, SimpleCandidate) + optimized = apply_candidate(root, decision.candidate) + + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + semis = find_joins(optimized, "Semi") + assert len(semis) == 1 + assert semis[0].children[0] is lineitem_ir + assert not find_joins(part_ir, "Semi") + assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) + + +def test_filter_pushdown_is_independent_of_dynamic_planning( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + make_config(dynamic_planning=False), + ) + + assert find_joins(optimized, "Semi") + + +def test_filter_pushdown_can_be_disabled( + simple_query: pl.LazyFrame, engine: SPMDEngine +) -> None: + root = translate_query(simple_query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + make_config(join_filter_pushdown=False), + ) + + assert optimized is root + + +@pytest.mark.parametrize( + "nulls_equal", [False, True], ids=["nulls_not_equal", "nulls_equal"] +) +def test_nullable_join_keys_preserve_results( + nulls_equal: bool, # noqa: FBT001 + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + domain = pl.LazyFrame( + { + "key": [None, 1, 2, 9], + "active": [True, True, True, False], + } + ).filter("active") + target = pl.LazyFrame( + { + "key": [None, 1, 2, 3] * 10, + "value": range(40), + } + ) + query = domain.join(target, on="key", nulls_equal=nulls_equal) + + ir = Translator(query._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_join_filter_pushdown( + ir, + collect_statistics(ir, config, parquet_stats_executor), + config, + ) + + semi_joins = find_joins(optimized, "Semi") + assert semi_joins + assert all(join.options[1] is nulls_equal for join in semi_joins) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_prefilter_does_not_move_below_distinct_on_non_subset_column( + engine: SPMDEngine, + parquet_stats_executor: concurrent.futures.ThreadPoolExecutor, +) -> None: + target = pl.LazyFrame( + { + "group": [1, 1] * 100, + "key": [0, 1] * 100, + } + ).unique(subset="group", keep="first", maintain_order=True) + domain = pl.LazyFrame( + { + "key": [1, 2], + "active": [True, False], + } + ).filter("active") + query = target.join(domain, on="key") + + ir = Translator(query._ldf.visit(), engine).translate_ir() + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_join_filter_pushdown( + ir, + collect_statistics(ir, config, parquet_stats_executor), + config, + ) + + semis = find_joins(optimized, "Semi") + assert any(isinstance(semi.children[0], Distinct) for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_no_simple_filter_pushdown_when_domain_is_not_selective( + engine: SPMDEngine, +) -> None: + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + lineitem = pl.LazyFrame({"l_suppkey": [i % 3 for i in range(20)]}) + query = supplier.join( + lineitem, + left_on="s_suppkey", + right_on="l_suppkey", + ) + root = translate_query(query, engine) + stats = StatsCollector() + assert isinstance(root, Join) + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) + + optimized = optimize_join_filter_pushdown( + root, + stats, + ConfigOptions.from_polars_engine(engine), + ) + + assert decision == Decision(reason="no_selective_domain") + assert optimized is root + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_composite_filter_pushdown_constrains_domain_first( + engine: SPMDEngine, +) -> None: + nation = ( + pl.LazyFrame( + { + "n_nationkey": range(10), + "active": [True] * 5 + [False] * 5, + } + ) + .filter("active") + .select("n_nationkey") + ) + orders = pl.LazyFrame( + { + "o_orderkey": range(90), + "n_nationkey": [i % 10 for i in range(90)], + } + ) + lineitem = pl.LazyFrame( + { + "l_orderkey": [i % 90 for i in range(180)], + "l_suppkey": [i % 30 for i in range(180)], + } + ) + supplier = pl.LazyFrame( + { + "s_suppkey": range(30), + "s_nationkey": [i % 10 for i in range(30)], + } + ) + query = ( + nation.join(orders, on="n_nationkey") + .join( + lineitem, + left_on="o_orderkey", + right_on="l_orderkey", + maintain_order="left", + ) + .join( + supplier, + left_on=("l_suppkey", "n_nationkey"), + right_on=("s_suppkey", "s_nationkey"), + ) + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + config = ConfigOptions.from_polars_engine(engine) + stats = StatsCollector() + + assert isinstance(root, Join) + order_lineitem, supplier_ir = root.children + assert isinstance(order_lineitem, Join) + lineitem_ir = order_lineitem.children[1] + decision = _select_candidate(root, 0.5, analyze_plan(root, stats)) + optimized = optimize_join_filter_pushdown(root, stats, config) + + assert decision.reason == "applied" + assert isinstance(decision.candidate, CompositeCandidate) + semis = find_joins(optimized, "Semi") + assert isinstance(optimized, Join) + assert optimized.options[0] == "Inner" + assert optimized.children[1] is supplier_ir + assert any(semi.children[0] is supplier_ir for semi in semis) + assert any(semi.children[0] is lineitem_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_derived_selectivity_propagates_through_rewritten_children( + engine: SPMDEngine, +) -> None: + region = ( + pl.LazyFrame( + { + "r_regionkey": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("r_regionkey") + ) + nation = pl.LazyFrame( + { + "n_nationkey": range(10), + "n_regionkey": [i % 2 for i in range(10)], + } + ) + customer = pl.LazyFrame( + { + "c_custkey": range(40), + "c_nationkey": [i % 10 for i in range(40)], + } + ) + orders = pl.LazyFrame( + { + "o_orderkey": range(200), + "o_custkey": [i % 40 for i in range(200)], + } + ) + query = ( + region.join(nation, left_on="r_regionkey", right_on="n_regionkey") + .join(customer, left_on="n_nationkey", right_on="c_nationkey") + .join(orders, left_on="c_custkey", right_on="o_custkey") + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + semis = find_joins(optimized, "Semi") + expected_targets = { + dataframe_scan(root, "n_nationkey"), + dataframe_scan(root, "c_custkey"), + dataframe_scan(root, "o_orderkey"), + } + assert expected_targets <= {semi.children[0] for semi in semis} + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_rewritten_domain_filters_other_side_instead_of_stacking( + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": range(6), + "part_active": [True] * 3 + [False] * 3, + } + ) + .filter("part_active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame( + { + "l_orderkey": [i % 15 for i in range(180)], + "l_partkey": [i % 6 for i in range(180)], + "l_suppkey": [i % 3 for i in range(180)], + } + ) + supplier = pl.LazyFrame({"s_suppkey": range(3)}) + orders = ( + pl.LazyFrame( + { + "o_orderkey": range(15), + "order_active": [True] * 8 + [False] * 7, + } + ) + .filter("order_active") + .select("o_orderkey") + ) + query = ( + part.join(lineitem, left_on="p_partkey", right_on="l_partkey") + .join(supplier, left_on="l_suppkey", right_on="s_suppkey") + .join(orders, left_on="l_orderkey", right_on="o_orderkey") + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + lineitem_ir = dataframe_scan(root, "l_orderkey") + orders_ir = dataframe_scan(root, "o_orderkey") + semis = find_joins(optimized, "Semi") + assert sum(semi.children[0] is lineitem_ir for semi in semis) == 1 + assert any(semi.children[0] is orders_ir for semi in semis) + assert not any( + isinstance(semi.children[0], Join) and semi.children[0].options[0] == "Semi" + for semi in semis + ) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + big = pl.LazyFrame( + { + "left_key": range(20), + "other": [i % 5 for i in range(20)], + } + ) + renamed_big = big.select(pl.col("left_key").alias("foo"), "other") + small = pl.LazyFrame( + { + "left_key": range(10), + "other2": [i % 5 for i in range(10)], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = renamed_big.join( + small, + left_on="other", + right_on="other2", + maintain_order="left", + ).join( + domain, + left_on="left_key", + right_on="domain_key", + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + + assert isinstance(root, Join) + joined = root.children[0] + assert isinstance(joined, Join) + renamed_big_ir, small_ir = joined.children + assert isinstance(renamed_big_ir, Select) + big_ir = renamed_big_ir.children[0] + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + semis = find_joins(optimized, "Semi") + assert any(semi.children[0] is small_ir for semi in semis) + assert not any(semi.children[0] is big_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_domain_source_follows_join_key_through_rename( + engine: SPMDEngine, +) -> None: + target = pl.LazyFrame({"target_key": range(20)}) + unrelated = pl.LazyFrame( + { + "domain_key": [100, 101], + "other": [0, 1], + "active": [True, False], + } + ) + renamed_unrelated = unrelated.filter("active").select( + pl.col("domain_key").alias("foo"), "other" + ) + domain_source = pl.LazyFrame( + { + "domain_key": range(1, 6), + "other2": range(5), + } + ) + domain = renamed_unrelated.join( + domain_source, + left_on="other", + right_on="other2", + maintain_order="left", + ) + query = target.join( + domain, + left_on="target_key", + right_on="domain_key", + ) + root = remove_cache_nodes(Translator(query._ldf.visit(), engine).translate_ir()) + + assert isinstance(root, Join) + target_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + renamed_unrelated_ir, domain_source_ir = domain_ir.children + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + semi = next( + semi for semi in find_joins(optimized, "Semi") if semi.children[0] is target_ir + ) + selected_domain = semi.children[1] + assert isinstance(selected_domain, Select) + rewritten_domain_source = selected_domain.children[0] + assert isinstance(rewritten_domain_source, Join) + assert rewritten_domain_source.options[0] == "Semi" + assert rewritten_domain_source.children[0] is domain_source_ir + assert rewritten_domain_source.children[0] is not renamed_unrelated_ir + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_composite_domain_columns_follow_renames(engine: SPMDEngine) -> None: + query = pl.LazyFrame( + { + "raw_key": range(10), + "raw_constraint": range(10), + } + ).select( + pl.col("raw_key").alias("domain_key"), + pl.col("raw_constraint").alias("domain_constraint"), + ) + renamed = translate_query(query, engine) + source = dataframe_scan(renamed, "raw_key") + + analyzed = analyze_plan(renamed, StatsCollector()) + facts = PlanFacts( + row_estimates={renamed: 20, source: 10}, + selective_nodes=analyzed.selective_nodes, + column_lineages=analyzed.column_lineages, + ) + producer = _smallest_node_containing_all( + renamed, ("domain_key", "domain_constraint"), facts + ) + + assert producer is not None + assert producer.node is source + assert producer.columns == ("raw_key", "raw_constraint") + + +def test_composite_domain_columns_do_not_reconverge_after_join( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + joined = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(joined, Join) + assert isinstance(joined.children[0], Cache) + assert joined.children[0] is joined.children[1] + joined = remove_cache_nodes(joined) + assert isinstance(joined, Join) + assert joined.children[0] is joined.children[1] + + facts = analyze_plan(joined, StatsCollector()) + producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) + + candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + assert candidates[0] == (ColumnRef(joined, "value"), ()) + assert len(candidates) >= 2 + assert all(path == (0,) * len(path) for _, path in candidates[1:]) + assert producer is not None + assert producer.node is joined + assert producer.columns == ("value", "value_right") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_contains_node_uses_dag_equality(engine: SPMDEngine) -> None: + query = pl.LazyFrame({"key": range(3)}).filter(pl.col("key") >= 0).slice(0, 2) + root = translate_query(query, engine) + + assert isinstance(root, Slice) + source = root.children[0] + equal_source = source.reconstruct(source.children) + + assert source is not equal_source + assert source == equal_source + assert contains_node(root, equal_source) + + +def test_plan_facts_share_lineage_suffixes_across_shared_dag( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame({"raw_key": range(10)}) + query = source.select(pl.col("raw_key").alias("left_key")).join( + source.select(pl.col("raw_key").alias("right_key")), + left_on="left_key", + right_on="right_key", + ) + root = translate_query(query, engine) + + assert isinstance(root, Join) + left, right = root.children + source_ir = dataframe_scan(root, "raw_key") + facts = analyze_plan(root, StatsCollector()) + left_lineage = facts.column_lineages[ColumnRef(left, "left_key")] + right_lineage = facts.column_lineages[ColumnRef(right, "right_key")] + source_lineage = facts.column_lineages[ColumnRef(source_ir, "raw_key")] + + assert left_lineage.column == ColumnRef(left, "left_key") + assert right_lineage.column == ColumnRef(right, "right_key") + while left_lineage.source is not source_lineage: + assert left_lineage.source is not None + left_lineage = left_lineage.source + while right_lineage.source is not source_lineage: + assert right_lineage.source is not None + right_lineage = right_lineage.source + assert left_lineage.source is right_lineage.source + assert source_lineage.source is None + + +def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: + target = ( + pl.LazyFrame({"target_key": range(20)}) + .filter(pl.col("target_key") >= 0) + .slice(0, 10) + ) + domain = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key") + ) + query = target.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) + + assert isinstance(root, Join) + sliced = root.children[0] + assert isinstance(sliced, Slice) + target_ir = dataframe_scan(root, "target_key") + stats = StatsCollector() + facts = analyze_plan(root, stats) + lineage = facts.column_lineages[ColumnRef(sliced, "target_key")] + assert lineage.column == ColumnRef(sliced, "target_key") + assert tuple(semijoin_pushdown_candidates(facts, sliced, "target_key")) == ( + (ColumnRef(sliced, "target_key"), ()), + ) + + optimized = optimize_join_filter_pushdown( + root, + stats, + ConfigOptions.from_polars_engine(engine), + ) + + semis = find_joins(optimized, "Semi") + assert any(semi.children[0] is sliced for semi in semis) + assert not any(semi.children[0] is target_ir for semi in semis) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_replacement_does_not_rewrite_shared_domain_side( + engine: SPMDEngine, +) -> None: + shared = pl.LazyFrame( + { + "target_key": range(20), + "other": [i % 2 for i in range(20)], + } + ) + domain_source = ( + pl.LazyFrame( + { + "domain_key": [1, 99], + "other2": [0, 1], + "active": [True, False], + } + ) + .filter("active") + .select("domain_key", "other2") + ) + domain = shared.join( + domain_source, + left_on=pl.col("other").cast(pl.Int32), + right_on=pl.col("other2").cast(pl.Int32), + ) + query = shared.join(domain, left_on="target_key", right_on="domain_key") + root = translate_query(query, engine) + + assert isinstance(root, Join) + shared_ir, domain_ir = root.children + assert isinstance(domain_ir, Join) + assert domain_ir.children[0] is shared_ir + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + assert isinstance(optimized, Join) + filtered, unfiltered_domain = optimized.children + assert unfiltered_domain is domain_ir + assert domain_ir.children[0] is shared_ir + semis = find_joins(filtered, "Semi") + assert len(semis) == 1 + assert semis[0].children[0] is dataframe_scan(root, "target_key") + assert not find_joins(unfiltered_domain, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +def test_target_prefilter_rewrites_only_selected_self_join_edge( + engine: SPMDEngine, +) -> None: + source = pl.LazyFrame( + { + "key": [1, 1, 2, 2], + "value": [10, 20, 30, 40], + } + ) + domain = ( + pl.LazyFrame( + { + "domain_value": [10, 999], + "active": [True, False], + } + ) + .filter("active") + .select("domain_value") + ) + query = source.join(source, on="key", suffix="_right").join( + domain, + left_on="value", + right_on="domain_value", + ) + translated = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(translated, Join) + translated_self_join = translated.children[0] + assert isinstance(translated_self_join, Join) + shared_cache = translated_self_join.children[0] + assert isinstance(shared_cache, Cache) + assert translated_self_join.children[1] is shared_cache + source_ir = shared_cache.children[0] + + optimized = optimize_with_stats( + translated, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + rewritten_self_join = optimized.children[0] + assert isinstance(rewritten_self_join, Join) + filtered, unfiltered = rewritten_self_join.children + assert unfiltered is source_ir + filtered_semis = find_joins(filtered, "Semi") + assert len(filtered_semis) == 1 + assert not find_joins(unfiltered, "Semi") + assert any(filtered_semis[0].children[0] is node for node in traversal([source_ir])) + assert_gpu_result_equal(query, engine=engine, check_row_order=False) + + +@pytest.mark.parametrize("how", ["left", "right", "cross", "full"]) +def test_no_filter_pushdown_for_unsupported_joins( + how: Any, + engine: SPMDEngine, +) -> None: + part = ( + pl.LazyFrame( + { + "p_partkey": [1, 99], + "active": [True, False], + } + ) + .filter("active") + .select("p_partkey") + ) + lineitem = pl.LazyFrame({"l_partkey": [i % 10 for i in range(20)]}) + if how == "cross": + query = part.join(lineitem, how=how) + else: + query = part.join( + lineitem, + left_on="p_partkey", + right_on="l_partkey", + how=how, + ) + root = translate_query(query, engine) + + optimized = optimize_join_filter_pushdown( + root, + StatsCollector(), + ConfigOptions.from_polars_engine(engine), + ) + + assert optimized is root + assert not find_joins(optimized, "Semi") + assert_gpu_result_equal(query, engine=engine, check_row_order=False) diff --git a/python/cudf_polars/tests/streaming/test_options.py b/python/cudf_polars/tests/streaming/test_options.py index c5a42062c9a7..d1a5a54aafa2 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -83,6 +83,11 @@ def test_executor_options_sink_to_directory_absent_when_unspecified() -> None: assert "sink_to_directory" not in StreamingOptions().to_executor_options() +def test_executor_options_join_filter_pushdown_disabled() -> None: + result = StreamingOptions(join_filter_pushdown=None).to_executor_options() + assert result["join_filter_pushdown"] is None + + # --------------------------------------------------------------------------- # to_engine_options # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/streaming/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py index 44d1b4ce0757..e32fab2ce3a3 100644 --- a/python/cudf_polars/tests/streaming/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -12,9 +12,13 @@ from polars.testing import assert_frame_equal from cudf_polars import Translator +from cudf_polars.dsl.ir import Cache, Join from cudf_polars.dsl.traversal import traversal from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.parallel import optimize_with_stats from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.config import ConfigOptions @pytest.mark.parametrize("column", ["a", "b"]) @@ -88,6 +92,30 @@ def test_evaluate_streaming(streaming_engine): assert_frame_equal(expected, got_streaming) +def test_optimize_removes_cache_nodes() -> None: + source = pl.LazyFrame({"key": [1, 1, 2], "value": [10, 20, 30]}) + query = source.join(source, on="key", suffix="_right") + engine = GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": None}, + ) + ir = Translator(query._ldf.visit(), engine).translate_ir() + + assert isinstance(ir, Join) + assert isinstance(ir.children[0], Cache) + assert ir.children[0] is ir.children[1] + + optimized = optimize_with_stats( + ir, + ConfigOptions.from_polars_engine(engine), + StatsCollector(), + ) + + assert isinstance(optimized, Join) + assert optimized.children[0] is optimized.children[1] + assert not any(isinstance(node, Cache) for node in traversal([optimized])) + + # --------------------------------------------------------------------------- # Tests migrated from tests/streaming/test_parallel.py (round 3) # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 15ba081e69fb..9cb44ce6c0a7 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -204,7 +204,7 @@ def test_target_partition_size( ) qir = Translator(q._ldf.visit(), _engine).translate_ir() config_options = ConfigOptions.from_polars_engine(_engine) - ir, info = lower_ir_graph( + lowering = lower_ir_graph( qir, config_options, collect_statistics( @@ -213,6 +213,8 @@ def test_target_partition_size( parquet_stats_executor, ), ) + ir = lowering.lowered + info = lowering.partition_info count = info[ir].count if blocksize <= 12_000: assert count > n_files diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index a0718e9be5bf..1d1395b48f0f 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -31,6 +31,7 @@ ConfigOptions, DynamicPlanningOptions, InMemoryExecutor, + JoinFilterPushdownOptions, MemoryResourceConfig, StreamingExecutor, ) @@ -614,6 +615,9 @@ def test_dynamic_planning_defaults() -> None: assert config.executor.dynamic_planning.join_prefilter_threshold == 0.5 assert config.executor.dynamic_planning.join_prefilter_max_key_columns == 1 assert not config.executor.dynamic_planning.join_prefilter_trace + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace def test_dynamic_planning_disabled_from_env(monkeypatch: pytest.MonkeyPatch) -> None: @@ -653,6 +657,31 @@ def test_join_prefilter_options_from_env(monkeypatch: pytest.MonkeyPatch) -> Non assert config.executor.dynamic_planning.join_prefilter_threshold == 0.25 assert config.executor.dynamic_planning.join_prefilter_max_key_columns is None assert config.executor.dynamic_planning.join_prefilter_trace + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.5 + assert not config.executor.join_filter_pushdown.trace + + +def test_join_filter_pushdown_options_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" + ) + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.join_filter_pushdown is not None + assert config.executor.join_filter_pushdown.threshold == 0.125 + assert config.executor.join_filter_pushdown.trace + + +def test_join_filter_pushdown_disabled_from_env( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "0") + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__TRACE", "1") + config = ConfigOptions.from_polars_engine(pl.GPUEngine()) + assert config.executor.join_filter_pushdown is None @pytest.mark.parametrize("value, expected", [("none", None), ("null", None), ("2", 2)]) @@ -747,6 +776,65 @@ def test_validate_join_prefilter_trace() -> None: ) +def test_validate_join_filter_pushdown_options() -> None: + with pytest.raises(TypeError, match="threshold must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"threshold": "bad"}}, + ) + ) + with pytest.raises(ValueError, match="threshold must be between"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"threshold": 1.5}}, + ) + ) + with pytest.raises(TypeError, match="trace must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": {"trace": "bad"}}, + ) + ) + + +def test_validate_join_filter_pushdown_type() -> None: + with pytest.raises( + TypeError, + match="join_filter_pushdown must be a JoinFilterPushdownOptions instance", + ): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": object()}, + ) + ) + + +def test_join_filter_pushdown_from_instance() -> None: + options = JoinFilterPushdownOptions(threshold=0.25, trace=True) + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": options}, + ) + ) + assert config.executor.join_filter_pushdown is options + + +def test_join_filter_pushdown_disabled_from_options() -> None: + config = ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={"join_filter_pushdown": None}, + ) + ) + assert config.executor.join_filter_pushdown is None + assert hash(config) == hash(config) + + def test_dynamic_planning_from_instance() -> None: config = ConfigOptions.from_polars_engine( pl.GPUEngine( diff --git a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx b/python/cudf_streaming/cudf_streaming/bloom_filter.pyx index c0554d5ddf21..5116a8aa44e2 100644 --- a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx +++ b/python/cudf_streaming/cudf_streaming/bloom_filter.pyx @@ -18,6 +18,7 @@ from rapidsmpf.streaming._detail.libcoro_spawn_task cimport cpp_set_py_future from rapidsmpf.streaming.chunks.utils cimport py_deleter from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel from rapidsmpf.streaming.core.context cimport Context, cpp_Context +from rapidsmpf.streaming.core.cancellation import (await_cpp_future, shutdown_channels) import asyncio @@ -208,7 +209,11 @@ cdef class BloomFilter: cpp_set_py_future, move(cpp_OwningWrapper(ret, py_deleter)), ) - await ret + # Note: multi-rank, if we get an exception we can't cancel the + # in-progress AllReduce, so this might still hang. + await await_cpp_future( + ret, on_cancel=lambda: shutdown_channels(ctx, ch_in, ch_out) + ) async def apply( self, @@ -248,4 +253,6 @@ cdef class BloomFilter: cpp_set_py_future, move(cpp_OwningWrapper(ret, py_deleter)), ) - await ret + await await_cpp_future( + ret, on_cancel=lambda: shutdown_channels(ctx, bloom_filter, ch_in, ch_out) + ) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py index 2f4c055bc308..d00e7926e50e 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py +++ b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py @@ -172,3 +172,60 @@ def test_bloom_filter_empty_build_filters_all( ) result.stream.synchronize() assert_eq(result.table_view(), expected) + + +def test_bloom_filter_build_exception_no_shutdown( + context: Context, comm: Communicator +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + stream = context.br().stream_pool.get_stream() + bloom = BloomFilter( + context, + comm, + seed=42, + num_filter_blocks=BloomFilter.fitting_num_blocks(1 << 20), + ) + ch_in: Channel[TableChunk] = context.create_channel() + ch_out: Channel[BloomFilterChunk] = context.create_channel() + messages = [ + Message( + sequence_number, + make_table( + np.arange(10, dtype=np.int32), + stream=stream, + br=context.br(), + ), + ) + for sequence_number in range(3) + ] + + async def recv_then_raise( + context: Context, ch_in: Channel[BloomFilterChunk] + ): + await ch_in.recv(context) + raise RuntimeError("Raising but didn't shutdown channel") + + # With no consumer for ch_out, the bloom-filter build blocks while draining + # its output channel. pytest-timeout interrupts run_actor_network, which + # must cancel and drain the worker before propagating the timeout failure. + with pytest.RaisesGroup( + pytest.RaisesExc(RuntimeError, match="didn't shutdown channel") + ): + run_actor_network( + context, + actors=[ + push_to_channel(context, ch_in, messages), + bloom.build(context, ch_in=ch_in, ch_out=ch_out, tag=0), + recv_then_raise(context, ch_out), + ], + ) + + async def recv_after_cancellation() -> Message | None: + return await asyncio.wait_for(ch_out.recv(context), timeout=1) + + # Cancellation should close the output channel. Without shutting down the + # channels inside the bloom filter if we get a cancellation, this + # receive picks up the message that is still in the channel. + assert asyncio.run(recv_after_cancellation()) is None