diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 5f4d7877582d..b2aa109e98ba 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -32,7 +32,7 @@ conda install cudf -c rapidsai-nightly -c conda-forge ``` 3. Build and view the docs locally following the instructions in the [Building -documentation docs](https://docs.rapids.ai/api/cudf/stable/developer_guide/documentation/#building-documentation) +documentation docs](https://docs.nvidia.com/cudf/latest/cudf/developer_guide/documentation/#building-documentation) 4. Follow steps 7-10 in the section [Your first issue](#your-first-issue) ## Code contributions @@ -325,9 +325,9 @@ This will bring up an interactive prompt to select which spelling fixes to apply ## Developer Guidelines -The [C++ Developer Guide](cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md) includes details on contributing to libcudf C++ code. +The [C++ Developer Guide](https://docs.nvidia.com/cudf/latest/libcudf/developer_guide/) includes details on contributing to libcudf C++ code. -The [Python Developer Guide](https://docs.rapids.ai/api/cudf/stable/cudf/developer_guide/) includes details on contributing to cuDF Python code. +The [Python Developer Guide](https://docs.nvidia.com/cudf/latest/developer_guide/) includes details on contributing to cuDF Python code. ## Attribution diff --git a/README.md b/README.md index 60ef0be89fa9..1c4477c2a4e4 100644 --- a/README.md +++ b/README.md @@ -1,21 +1,22 @@ -#
 cuDF - A GPU-accelerated DataFrame library for tabular data processing
+# NVIDIA cuDF: A GPU-accelerated DataFrame library for tabular data processing -cuDF (pronounced "KOO-dee-eff") is an [Apache 2.0 licensed](LICENSE), GPU-accelerated DataFrame library -for tabular data processing. The cuDF library is one part of the [RAPIDS](https://rapids.ai/) GPU -Accelerated Data Science suite of libraries. +NVIDIA cuDF (pronounced "KOO-dee-eff") is an [Apache 2.0 licensed](LICENSE), GPU-accelerated DataFrame library +for tabular data processing. The cuDF library is one part of the [NVIDIA +CUDA-X](https://developer.nvidia.com/cuda/cuda-x-libraries) suite of GPU +Accelerated libraries. ## About cuDF is composed of multiple libraries including: -* [libcudf](https://docs.rapids.ai/api/libcudf/stable/): A CUDA C++ library with [Apache Arrow](https://arrow.apache.org/) compliant +* [libcudf](https://docs.nvidia.com/cudf/latest/libcudf/): A CUDA C++ library with [Apache Arrow](https://arrow.apache.org/) compliant data structures and fundamental algorithms for tabular data. -* [pylibcudf](https://docs.rapids.ai/api/cudf/stable/pylibcudf/): A Python library providing [Cython](https://cython.org/) bindings for libcudf. -* [cudf](https://docs.rapids.ai/api/cudf/stable/cudf/): A Python library providing +* [pylibcudf](https://docs.nvidia.com/cudf/latest/pylibcudf/): A Python library providing [Cython](https://cython.org/) bindings for libcudf. +* [cudf](https://docs.nvidia.com/cudf/latest/cudf/): A Python library providing - A DataFrame library mirroring the [pandas](https://pandas.pydata.org/) API - - A zero-code change accelerator, [cudf.pandas](https://docs.rapids.ai/api/cudf/stable/cudf_pandas/), for existing pandas code. -* [cudf-polars](https://docs.rapids.ai/api/cudf/stable/cudf_polars/): A Python library providing a GPU engine for [Polars](https://pola.rs/) -* [dask-cudf](https://docs.rapids.ai/api/dask-cudf/stable/): A Python library providing a GPU backend for [Dask](https://www.dask.org/) DataFrames + - A zero-code change accelerator, [cudf.pandas](https://docs.nvidia.com/cudf/latest/cudf_pandas/), for existing pandas code. +* [cudf-polars](https://docs.nvidia.com/cudf/latest/cudf_polars/): A Python library providing a GPU engine for [Polars](https://pola.rs/) +* [dask-cudf](https://docs.nvidia.com/dask-cudf/latest/): A Python library providing a GPU backend for [Dask](https://www.dask.org/) DataFrames Notable projects that use cuDF include: diff --git a/conda/recipes/libcudf/recipe.yaml b/conda/recipes/libcudf/recipe.yaml index 304bfd9df0c2..59b80c8a241e 100644 --- a/conda/recipes/libcudf/recipe.yaml +++ b/conda/recipes/libcudf/recipe.yaml @@ -209,7 +209,7 @@ outputs: - libzlib - libnvcomp about: - homepage: https://rapids.ai/ + homepage: https://docs.nvidia.com/cudf/ license: Apache-2.0 summary: libcudf_kafka library @@ -253,7 +253,7 @@ outputs: - libzlib - libnvcomp about: - homepage: https://rapids.ai/ + homepage: https://docs.nvidia.com/cudf/ license: Apache-2.0 summary: libcudf-streaming library @@ -431,6 +431,6 @@ outputs: - librmm - libnvcomp about: - homepage: https://rapids.ai/ + homepage: https://docs.nvidia.com/cudf/ license: Apache-2.0 summary: libcudf-streaming test & benchmark executables diff --git a/cpp/cmake/Modules/ConfigureCUDA.cmake b/cpp/cmake/Modules/ConfigureCUDA.cmake index c2635ec68ae5..40b4aea434db 100644 --- a/cpp/cmake/Modules/ConfigureCUDA.cmake +++ b/cpp/cmake/Modules/ConfigureCUDA.cmake @@ -1,12 +1,12 @@ # ============================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ============================================================================= if(CMAKE_COMPILER_IS_GNUCXX) - list(APPEND CUDF_CXX_FLAGS -Wall -Werror -Wno-unknown-pragmas -Wno-error=deprecated-declarations) + list(APPEND CUDF_CXX_FLAGS -Wall -Werror -Wno-unknown-pragmas) endif() list(APPEND CUDF_CUDA_FLAGS --expt-extended-lambda) @@ -17,7 +17,7 @@ if(CUDA_WARNINGS_AS_ERRORS) else() list(APPEND CUDF_CUDA_FLAGS -Werror=cross-execution-space-call) endif() -list(APPEND CUDF_CUDA_FLAGS -Xcompiler=-Wall,-Werror,-Wno-error=deprecated-declarations) +list(APPEND CUDF_CUDA_FLAGS -Xcompiler=-Wall,-Werror) # This warning needs to be suppressed because some parts of cudf instantiate templated CCCL # functions in contexts where the resulting instantiations would have internal linkage (e.g. in # anonymous namespaces). In such contexts, the visibility attribute on the template is ignored, and diff --git a/cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md b/cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md index 84a0e37fd5c3..e3b9260a570a 100644 --- a/cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md +++ b/cpp/doxygen/developer_guide/DEVELOPER_GUIDE.md @@ -26,7 +26,7 @@ A column is an array of data of a single type. Along with Tables, columns are th structures used in libcudf. Most libcudf algorithms operate on columns. Columns may have a validity mask representing whether each element is valid or null (invalid). Columns of nested types are supported, meaning that a column may have child columns. A column is the C++ equivalent to a cuDF -Python [Series](https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.series/). +Python [Series](https://docs.nvidia.com/cudf/latest/cudf/api_docs/series/). ### Element @@ -41,7 +41,7 @@ A type representing a single element of a data type. A table is a collection of columns that all have the same number of elements (rows). A table may also have zero columns while still carrying a row count, mirroring an `(N, 0)` DataFrame. A table is the C++ equivalent to a cuDF Python -[DataFrame](https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.dataframe/). +[DataFrame](https://docs.nvidia.com/cudf/latest/cudf/api_docs/dataframe/). ### View @@ -674,7 +674,7 @@ cudf::detail::copy_if( ## Memory Allocation -Device [memory resources](#rmmdevice_memory_resource) are used in libcudf to abstract and control +Device [memory resources](#memory-resources) are used in libcudf to abstract and control how device memory is allocated. ### Output Memory diff --git a/cpp/doxygen/developer_guide/DOCUMENTATION.md b/cpp/doxygen/developer_guide/DOCUMENTATION.md index 14ac0fa247bb..83224fb16748 100644 --- a/cpp/doxygen/developer_guide/DOCUMENTATION.md +++ b/cpp/doxygen/developer_guide/DOCUMENTATION.md @@ -1,6 +1,6 @@ # libcudf C++ Documentation Guide -These guidelines apply to documenting all libcudf C++ source files using doxygen style formatting although only public APIs and classes are actually [published](https://docs.rapids.ai/api/libcudf/stable/index.html). +These guidelines apply to documenting all libcudf C++ source files using doxygen style formatting although only public APIs and classes are actually [published](https://docs.nvidia.com/cudf/latest/libcudf/api_docs/). ## Copyright License @@ -438,4 +438,5 @@ Then open `:8000` in your local web browser, inserting the IP addres The doxygen output is intended for building documentation only for the public APIs and classes. For example, the output should not include documentation for `detail` or `/src` files, and these directories are excluded in the `Doxyfile` configuration. -When published by the build/CI system, the doxygen output will appear on our external [RAPIDS web site](https://docs.rapids.ai/api/libcudf/stable/index.html). +When published by the build/CI system, the doxygen output will appear as +part of the [cuDF documentation](https://docs.nvidia.com/cudf/latest/libcudf/). diff --git a/cpp/include/cudf/groupby.hpp b/cpp/include/cudf/groupby.hpp index a3e21b58093f..da98e4e97142 100644 --- a/cpp/include/cudf/groupby.hpp +++ b/cpp/include/cudf/groupby.hpp @@ -519,6 +519,13 @@ class streaming_groupby { * are updated atomically. The input `data` table is not referenced after this * call returns. * + * This function may be called concurrently from multiple host threads on the same object, + * and each call may supply a different stream. Callers do not need to serialize the calls or + * synchronize between them. Key insertion is serialized internally, on the host and across + * streams, because a batch's newly discovered keys are held in a transient encoding that is + * only valid while that one insertion is in flight. The aggregation that follows each + * insertion updates every group atomically, so those phases overlap freely across streams. + * * @param data Table containing both key and value columns * @param stream CUDA stream used for device memory operations and kernel launches * @@ -532,6 +539,10 @@ class streaming_groupby { * * Extracts the other object's accumulated intermediate state and merges it into this * object's persistent hash table. The other object is not modified. + * + * This function shares the insertion path with `aggregate()` and is serialized against it, so + * it is safe to call while other host threads are calling `aggregate()` on this object. The + * source object must not be mutated concurrently. * Both objects must have been constructed with compatible aggregation requests, * and this object must have had at least one `aggregate()` call. * diff --git a/cpp/src/groupby/streaming_groupby/aggregate.cu b/cpp/src/groupby/streaming_groupby/aggregate.cu index c024565720d1..575f24cd84fc 100644 --- a/cpp/src/groupby/streaming_groupby/aggregate.cu +++ b/cpp/src/groupby/streaming_groupby/aggregate.cu @@ -17,15 +17,14 @@ #include #include +#include #include namespace cudf::groupby { void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_ref stream) { - CUDF_EXPECTS(!_invalidated, - "streaming_groupby is in an invalidated state from a prior failure; " - "no further aggregate()/merge() is allowed. finalize() may still be called."); + ensure_not_invalidated(); auto const batch_size = data.num_rows(); if (batch_size == 0) { return; } @@ -40,15 +39,29 @@ void streaming_groupby::impl::do_aggregate(table_view const& data, cuda::stream_ "Transient key encoding (max_distinct_keys + batch_size) would overflow size_type.", std::invalid_argument); - if (!_initialized) { initialize(data, stream); } + // The transient key encoding is only valid while a single insertion is in flight, so + // insertion is serialized across concurrent callers on the host and, via the event, on the + // device. The aggregation below is per-group atomic and runs unserialized. + auto const result = [&] { + std::lock_guard const lock{_insert_mutex}; - auto const batch_keys = data.select(_key_indices); + // Re-check under the lock: another caller may have invalidated the object since the + // fail-fast check above. + ensure_not_invalidated(); - update_nullable_state(batch_keys); + if (!_initialized) { initialize(data, stream); } - if (!_key_set) { create_key_set(stream); } + auto const batch_keys = data.select(_key_indices); - auto result = probe_and_insert(batch_keys, stream); + update_nullable_state(batch_keys); + + if (!_key_set) { create_key_set(stream); } + + _insert_done.wait(stream); + auto inserted = probe_and_insert(batch_keys, stream); + _insert_done.record(stream); + return inserted; + }(); auto const values_view = data.select(_value_col_indices); auto const d_values = table_device_view::create(values_view, stream); diff --git a/cpp/src/groupby/streaming_groupby/common.cuh b/cpp/src/groupby/streaming_groupby/common.cuh index 100837863792..5f68ceb35866 100644 --- a/cpp/src/groupby/streaming_groupby/common.cuh +++ b/cpp/src/groupby/streaming_groupby/common.cuh @@ -26,11 +26,40 @@ #include #include +#include #include +#include #include namespace cudf::groupby { +/* + * Minimal owning wrapper around a CUDA event, used to order the insertion phase of + * `aggregate()` / `merge()` calls that overlap on different streams. Only the insertion + * phase needs this ordering; the aggregation phase updates each group with atomics and is + * safe to overlap. + */ +class insert_order_event { + public: + insert_order_event() { CUDF_CUDA_TRY(cudaEventCreateWithFlags(&_event, cudaEventDisableTiming)); } + ~insert_order_event() { cudaEventDestroy(_event); } + insert_order_event(insert_order_event const&) = delete; + insert_order_event& operator=(insert_order_event const&) = delete; + + /// Makes `stream` wait for the most recently recorded insertion. No-op before the first + /// `record()`, which is exactly the behavior the first call needs. + void wait(cuda::stream_ref stream) const + { + CUDF_CUDA_TRY(cudaStreamWaitEvent(stream.get(), _event)); + } + + /// Records completion of the insertion just enqueued on `stream`. + void record(cuda::stream_ref stream) { CUDF_CUDA_TRY(cudaEventRecord(_event, stream.get())); } + + private: + cudaEvent_t _event{}; +}; + /* * Companion location for a stored dense ID: which compacted batch table the key * lives in (`first`) and the row index within that table (`second`). Packed into @@ -250,17 +279,27 @@ struct streaming_groupby::impl { null_policy _null_handling; cuda::mr::any_resource _mr; + /* + * Serializes the insertion phase of `aggregate()` and `merge()`. Callers may invoke those + * from multiple host threads; everything they mutate on the host, and the transient key + * encoding they place in the hash set, is guarded here. + */ + std::mutex _insert_mutex; + /// Orders the insertion phase across calls that supply different streams. + insert_order_event _insert_done; + bool _initialized{false}; /// Set true once an `aggregate()` / `merge()` call has thrown after touching the /// hash set. Subsequent `aggregate()` / `merge()` calls fail fast; only - /// `finalize()` may still be called to recover partial results. - bool _invalidated{false}; + /// `finalize()` may still be called to recover partial results. Atomic so the + /// fail-fast check in `do_aggregate` can run ahead of `_insert_mutex`. + std::atomic _invalidated{false}; /* * Number of distinct keys accumulated so far. Also serves as the high-water * mark of dense IDs in the persistent hash set: stored slot values are in * [0, _distinct_keys). */ - size_type _distinct_keys{0}; + std::atomic _distinct_keys{0}; bool _has_nullable_keys{false}; bool _has_nested_keys{false}; @@ -301,7 +340,12 @@ struct streaming_groupby::impl { std::unique_ptr _key_set; [[nodiscard]] size_type num_keys() const { return static_cast(_key_indices.size()); } - [[nodiscard]] bool has_state() const { return _initialized && _distinct_keys > 0; } + void ensure_not_invalidated() const + { + CUDF_EXPECTS(!_invalidated.load(std::memory_order_relaxed), + "streaming_groupby is in an invalidated state from a prior failure; " + "no further aggregate()/merge() is allowed. finalize() may still be called."); + } impl(host_span key_indices, host_span requests, diff --git a/cpp/src/groupby/streaming_groupby/impl.cu b/cpp/src/groupby/streaming_groupby/impl.cu index de8ddcbb3254..19b94b7ce132 100644 --- a/cpp/src/groupby/streaming_groupby/impl.cu +++ b/cpp/src/groupby/streaming_groupby/impl.cu @@ -243,7 +243,9 @@ std::unique_ptr streaming_groupby::impl::gather_agg_results( // The results we care about are dense in `[0, _distinct_keys)` and can be extracted by // slice+copy. auto const sliced = - cudf::detail::slice(_agg_results->view(), {0, _distinct_keys}, stream).front(); + cudf::detail::slice( + _agg_results->view(), {0, _distinct_keys.load(std::memory_order_relaxed)}, stream) + .front(); return std::make_unique
(sliced, stream, mr); } @@ -378,7 +380,10 @@ std::pair, std::vector> streaming_gro return _impl->do_finalize(stream, mr); } -size_type streaming_groupby::distinct_keys() const noexcept { return _impl->_distinct_keys; } +size_type streaming_groupby::distinct_keys() const noexcept +{ + return _impl->_distinct_keys.load(std::memory_order_relaxed); +} bool is_streaming_groupby_supported(data_type values_type, aggregation::Kind kind) { diff --git a/cpp/src/groupby/streaming_groupby/insert.cuh b/cpp/src/groupby/streaming_groupby/insert.cuh index 9df90ecfcb93..d2c51acd5667 100644 --- a/cpp/src/groupby/streaming_groupby/insert.cuh +++ b/cpp/src/groupby/streaming_groupby/insert.cuh @@ -96,9 +96,10 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ // Bound check: the hash set has already been written above (transient slot values), // so on failure the object is left invalidated; further aggregate()/merge() calls // will throw immediately while finalize() can still recover partial results. - if (_distinct_keys + new_distinct_keys > _max_distinct_keys) { + auto const distinct_so_far = _distinct_keys.load(std::memory_order_relaxed); + if (distinct_so_far + new_distinct_keys > _max_distinct_keys) { _invalidated = true; - CUDF_FAIL("Distinct key count (" + std::to_string(_distinct_keys + new_distinct_keys) + + CUDF_FAIL("Distinct key count (" + std::to_string(distinct_so_far + new_distinct_keys) + ") would exceed max_distinct_keys (" + std::to_string(_max_distinct_keys) + ")."); } @@ -115,7 +116,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ // Store the compacted batch. auto const new_batch_id = static_cast(_compacted_batches.size()); - auto const dense_id_offset = _distinct_keys; + auto const dense_id_offset = distinct_so_far; _compacted_batches.push_back(std::move(compacted)); _preprocessed_batches.push_back(preprocessed_compacted); @@ -138,7 +139,7 @@ streaming_groupby::impl::batch_insert_result streaming_groupby::impl::probe_and_ update_transient_target_indices_fn{ base, slot_offsets.data(), _max_distinct_keys, target_indices.data()}); - _distinct_keys += new_distinct_keys; + _distinct_keys.fetch_add(new_distinct_keys, std::memory_order_relaxed); } // If new_distinct_keys == 0, target_indices is already final from Pass 1 — every // slot held a dense ID at probe time, so *iter was already the correct dense ID. diff --git a/cpp/src/groupby/streaming_groupby/merge.cu b/cpp/src/groupby/streaming_groupby/merge.cu index 6d79fdd7c18b..7294fa295e41 100644 --- a/cpp/src/groupby/streaming_groupby/merge.cu +++ b/cpp/src/groupby/streaming_groupby/merge.cu @@ -18,6 +18,7 @@ #include #include +#include #include namespace cudf::groupby { @@ -90,17 +91,20 @@ struct merge_single_pass_aggs_fn { void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref stream) { - CUDF_EXPECTS(!_invalidated, - "streaming_groupby is in an invalidated state from a prior failure; " - "no further aggregate()/merge() is allowed. finalize() may still be called."); - CUDF_EXPECTS(!other._invalidated, "Cannot merge from an invalidated streaming_groupby."); + // `other` is only read from, so a single lock on this object's insertion state is enough. + std::lock_guard const lock{_insert_mutex}; - if (!other._initialized || !other.has_state()) { return; } + ensure_not_invalidated(); + CUDF_EXPECTS(!other._invalidated.load(std::memory_order_relaxed), + "Cannot merge from an invalidated streaming_groupby."); + + auto const other_distinct_keys = other._distinct_keys.load(std::memory_order_relaxed); + if (!other._initialized || other_distinct_keys == 0) { return; } CUDF_EXPECTS(_initialized, "Cannot merge into an uninitialized streaming_groupby. " "Call aggregate() at least once before merge()."); - CUDF_EXPECTS(other._distinct_keys <= _max_distinct_keys, - "Merge source distinct keys (" + std::to_string(other._distinct_keys) + + CUDF_EXPECTS(other_distinct_keys <= _max_distinct_keys, + "Merge source distinct keys (" + std::to_string(other_distinct_keys) + ") exceeds max_distinct_keys (" + std::to_string(_max_distinct_keys) + ").", std::invalid_argument); CUDF_EXPECTS(other._agg_kinds == _agg_kinds, @@ -115,16 +119,16 @@ void streaming_groupby::impl::do_merge(impl const& other, cuda::stream_ref strea auto const mr = cudf::get_current_device_resource_ref(); - auto other_keys = other.gather_distinct_keys(stream, mr); - auto const other_key_view = other_keys->view(); - auto const other_distinct_keys = other._distinct_keys; - if (other_distinct_keys == 0) { return; } + auto other_keys = other.gather_distinct_keys(stream, mr); + auto const other_key_view = other_keys->view(); update_nullable_state(other_key_view); if (!_key_set) { create_key_set(stream); } + _insert_done.wait(stream); auto result = probe_and_insert(other_key_view, stream); + _insert_done.record(stream); // Merge aggregation values using dense target indices. We only read from // `other._agg_results`; no need to deep-copy the source rows like keys. diff --git a/cpp/tests/groupby/streaming_groupby_test.cpp b/cpp/tests/groupby/streaming_groupby_test.cpp index 885e704ad94b..343ca6663ef2 100644 --- a/cpp/tests/groupby/streaming_groupby_test.cpp +++ b/cpp/tests/groupby/streaming_groupby_test.cpp @@ -19,8 +19,12 @@ #include #include +#include +#include #include +#include +#include #include static std::vector const KEY_COL{0}; @@ -325,6 +329,79 @@ TEST_F(StreamingGroupbyTest, MergeTwoObjects) check(keys, results, cudf::table_view{{ek}}, {ev}); } +TEST_F(StreamingGroupbyTest, ConcurrentAggregate) +{ + using K = int32_t; + using V = int32_t; + + constexpr int num_batches = 8; + + // Every batch re-hits keys 0 and 1 and introduces one key of its own, so concurrent calls + // both collide on existing groups and discover new keys at the same time. + std::vector> keys; + std::vector> vals; + keys.reserve(num_batches); + vals.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + keys.emplace_back(std::initializer_list{0, 1, static_cast(i + 2)}); + vals.emplace_back(std::initializer_list{1, 10, 100}); + } + + std::vector batches; + batches.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + batches.push_back(cudf::table_view{{keys[i], vals[i]}}); + } + + std::vector> streams; + streams.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + streams.push_back(std::make_unique()); + } + + auto reqs = single_agg_req(1, cudf::make_sum_aggregation()); + cudf::groupby::streaming_groupby streaming_agg(KEY_COL, reqs, DEFAULT_MAX_DISTINCT_KEYS); + + auto const device = rmm::get_current_cuda_device(); + std::vector threads; + std::vector errors(num_batches); + // `ready` lets the main thread wait until every worker is spinning, and `start` then releases + // them together, so the aggregate() calls actually overlap. + std::atomic ready{0}; + std::atomic start{false}; + threads.reserve(num_batches); + for (int i = 0; i < num_batches; ++i) { + threads.emplace_back([&, i] { + rmm::cuda_set_device_raii const device_guard{device}; + ready.fetch_add(1, std::memory_order_relaxed); + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + try { + streaming_agg.aggregate(batches[i], streams[i]->view()); + } catch (...) { + errors[i] = std::current_exception(); + } + }); + } + while (ready.load(std::memory_order_relaxed) != num_batches) { + std::this_thread::yield(); + } + start.store(true, std::memory_order_release); + for (auto& thread : threads) { + thread.join(); + } + for (auto const& error : errors) { + EXPECT_FALSE(error); + } + for (auto const& stream : streams) { + stream->synchronize(); + } + + auto [out_keys, results] = streaming_agg.finalize(); + verify_against_groupby(out_keys, results, batches, KEY_COL, reqs); +} + TEST_F(StreamingGroupbyTest, EmptyBatch) { using K = int32_t; @@ -849,6 +926,14 @@ TEST_F(StreamingGroupbyTest, ExceedsDistinctKeyCapacityThrows) cudf::test::fixed_width_column_wrapper k2{4}; cudf::test::fixed_width_column_wrapper v2{50}; EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k2, v2}}), cudf::logic_error); + + // The object is now invalidated: even an empty batch is rejected, while finalize() still + // recovers the groups inserted before the failure. + cudf::test::fixed_width_column_wrapper k_empty{}; + cudf::test::fixed_width_column_wrapper v_empty{}; + EXPECT_THROW(streaming_agg.aggregate(cudf::table_view{{k_empty, v_empty}}), cudf::logic_error); + auto [keys, results] = streaming_agg.finalize(); + EXPECT_EQ(keys->num_rows(), 4); } // Test that sliced input columns with non-zero offsets work correctly. diff --git a/docs/cudf/source/_static/RAPIDS-logo-purple.png b/docs/cudf/source/_static/RAPIDS-logo-purple.png deleted file mode 100644 index d884e01374dc..000000000000 Binary files a/docs/cudf/source/_static/RAPIDS-logo-purple.png and /dev/null differ diff --git a/docs/cudf/source/conf.py b/docs/cudf/source/conf.py index bf89f727e793..7159393f4950 100644 --- a/docs/cudf/source/conf.py +++ b/docs/cudf/source/conf.py @@ -31,19 +31,73 @@ from typing import Any import cudf +from breathe.directives.content_block import DoxygenPageDirective +from breathe.renderer.sphinxrenderer import SphinxRenderer from docutils import nodes from docutils.nodes import Node, Text from packaging.version import Version from pygments.lexer import RegexLexer from pygments.token import Text as PText +from sphinx import addnodes from sphinx.addnodes import pending_xref from sphinx.application import Sphinx from sphinx.ext import intersphinx from sphinx.ext.autodoc import ClassDocumenter from sphinx.highlighting import lexers +from sphinx.util import logging from sphinx.util.nodes import clean_astext, make_refnode +# TODO: ship this upstream in breathe. Today breathe doesn't render +# docsect4 headers or the contents at all. +def visit_docsect4(self, node): + (title,) = (item for item in node.content_ if item.name == "title") + + section = nodes.section(ids=[self.get_refid(node.id)]) + section += nodes.title("", "", *self.render(title)) + section += self.create_doxygen_target(node) + section += self.render_iterable( + [item for item in node.content_ if item.name != "title"] + ) + return [section] + + +SphinxRenderer.methods["docsect4"] = visit_docsect4 + + +class FlatDoxygenPageDirective(DoxygenPageDirective): + """ + Render a Doxygen page body as direct children of its RST section. + + This is basically the same as breathe's doxygenpage directive but + allows us to splice a page in to the toctree at the right level such + that it is rendered with an "on this page" sidebar by the theme we use. + """ + + def run(self): + self.options["content-only"] = None + rendered = super().run() + + if len(rendered) != 1 or not isinstance(rendered[0], nodes.container): + return rendered + + content = next( + ( + child + for child in rendered[0].children + if isinstance(child, addnodes.desc_content) + ), + None, + ) + if content is None: + return rendered + + children = list(content.children) + for child in children: + content.remove(child) + return children + + class PseudoLexer(RegexLexer): """Trivial lexer for pseudocode.""" @@ -166,6 +220,9 @@ def clean_all_xml_files(path): copybutton_prompt_text = ">>> " autosummary_generate = True +toc_object_entries_show_parents = "hide" +maximum_signature_line_length = 70 + # Enable automatic generation of systematic, namespaced labels for sections myst_heading_anchors = 2 @@ -176,7 +233,7 @@ def clean_all_xml_files(path): # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = {".rst": "restructuredtext"} +source_suffix = {".rst": "restructuredtext", ".md": "myst-nb"} # The master toctree document. master_doc = "index" @@ -324,22 +381,33 @@ def clean_all_xml_files(path): ) ] +with open("../../../RAPIDS_BRANCH", "r") as f: + branch = f.read().strip() +intersphinx_version = "latest" if branch == "main" else version -# Example configuration for intersphinx: refer to the Python standard library. intersphinx_mapping = { "cupy": ("https://docs.cupy.dev/en/stable/", None), + "dask-cuda": ( + f"https://docs.nvidia.com/dask-cuda/{intersphinx_version}/", + None, + ), + "dask-cudf": ( + f"https://docs.nvidia.com/dask-cudf/{intersphinx_version}/", + None, + ), "dlpack": ("https://dmlc.github.io/dlpack/latest/", None), + "kvikio": (f"https://docs.nvidia.com/kvikio/{intersphinx_version}/", None), "nanoarrow": ("https://arrow.apache.org/nanoarrow/latest/", None), "numpy": ("https://numpy.org/doc/stable/", None), - # Temporarily disable nitpick warnings for pandas: https://github.com/pandas-dev/pandas/issues/64584 - # "pandas": ( - # "https://pandas.pydata.org/pandas-docs/stable/", - # None, - # ), + "pandas": ("https://pandas.pydata.org/pandas-docs/stable/", None), "polars": ("https://docs.pola.rs/api/python/stable/", None), "pyarrow": ("https://arrow.apache.org/docs/", None), "python": ("https://docs.python.org/3/", None), - "rmm": ("https://docs.rapids.ai/api/rmm/nightly/", None), + "rmm": (f"https://docs.nvidia.com/rmm/{intersphinx_version}/", None), + "rapidsmpf": ( + f"https://docs.nvidia.com/rapidsmpf/{intersphinx_version}/", + None, + ), "typing_extensions": ( "https://typing-extensions.readthedocs.io/en/stable/", None, @@ -456,6 +524,8 @@ def _generate_namespaces(namespaces): "type_to_scalar_type_impl", "type_to_scalar_type_impl", "detail", + # Test-only helper types are intentionally not published as API pages. + "classcudf_1_1test_1_1", # kafka objects "python_callable_type", "kafka_oauth_callback_wrapper_type", @@ -479,7 +549,7 @@ def _generate_namespaces(namespaces): _intersphinx_extra_prefixes = ("rmm", "rmm::mr", "mr") _external_intersphinx_aliases = { - # "pandas": "pd", + "pandas": "pd", "pyarrow": "pa", "numpy": "np", "cupy": "cp", @@ -653,6 +723,12 @@ def on_missing_reference(app, env, node, contnode): nitpick_ignore = [ ("py:class", "Dtype"), ("py:class", "pandas.core.indexes.frozen.FrozenList"), + # pandas does not publish these implementation types in its inventory. + ("py:class", "pandas.api.typing.FrozenList"), + ( + "py:class", + "pandas.core.arrays.arrow.extension_types.ArrowIntervalType", + ), ("py:class", "ScalarLike"), ("py:class", "StringColumn"), ("py:class", "ColumnLike"), @@ -670,6 +746,11 @@ def on_missing_reference(app, env, node, contnode): ("py:class", "Statistics"), ("py:class", "Communicator"), ("py:class", "Options"), + # Not yet published in API docs. + ("py:class", "rapidsmpf.streaming.core.context.Context"), + ("py:func", "rapidsmpf.rrun.rrun.bind"), + # kvikio aliases that don't match the public intersphinx targets. + ("py:class", "kvikio.Summary"), # polars aliases that don't match the public intersphinx targets. ("py:class", "pl.DataFrame"), ("py:class", "pl.DataType"), @@ -694,17 +775,11 @@ def on_missing_reference(app, env, node, contnode): ("py:class", "SupportsCudaArrayInterface"), ("py:class", "T"), ] -# Temporarily disable nitpick warnings for pandas: https://github.com/pandas-dev/pandas/issues/64584 + nitpick_ignore_regex = [ - ("py:.*", "pandas.*"), - ("py:.*", "pd.*"), - ("ref.*", ".*pandas.*"), # External libs without configured intersphinx inventories. - ("py:.*", r"rapidsmpf(\..*)?"), - ("py:.*", r"kvikio(\..*)?"), ("py:.*", r"ray(\..*)?"), ("py:.*", r"distributed(\..*)?"), - ("py:.*", r"dask_cuda(\..*)?"), ] @@ -854,15 +929,201 @@ def register_sections_as_label(app: Sphinx, document: Node) -> None: domain.labels[name] = docname, labelid, title +_libcudf_developer_guide_documents = { + "DEVELOPER_GUIDE.md": "libcudf/developer_guide/DEVELOPER_GUIDE", + "DOCUMENTATION.md": "libcudf/developer_guide/DOCUMENTATION", + "TESTING.md": "libcudf/developer_guide/TESTING", + "BENCHMARKING.md": "libcudf/developer_guide/BENCHMARKING", + "PROFILING.md": "libcudf/developer_guide/PROFILING", +} +_libcudf_developer_guide_source_files = { + docname: filename + for filename, docname in _libcudf_developer_guide_documents.items() +} +_libcudf_developer_guide_xref_prefix = "libcudf-guide-md:" +_libcudf_developer_guide_logger = logging.getLogger(__name__) +_libcudf_developer_guide_source_dir = os.path.abspath( + os.path.join( + os.path.dirname(__file__), "../../../cpp/doxygen/developer_guide" + ) +) + + +def _markdown_heading_slug(title: str) -> str: + """Return the GitHub-style fragment generated for a Markdown heading.""" + slug = title.lower() + slug = re.sub(r"[^\w\s-]", "", slug) + return re.sub(r"-+", "-", re.sub(r"\s+", "-", slug)).strip("-") + + +def _find_libcudf_developer_guide_section(env, docname, fragment): + """Return the guide section identified by a Markdown or Doxygen fragment.""" + slugs: dict[str, int] = {} + for section in env.get_doctree(docname).findall(nodes.section): + title = clean_astext(section[0]) + base_slug = _markdown_heading_slug(title) + occurrence = slugs.get(base_slug, 0) + slugs[base_slug] = occurrence + 1 + heading_slug = ( + base_slug if occurrence == 0 else f"{base_slug}-{occurrence}" + ) + if fragment in (heading_slug, section["ids"][0]): + return section, title + for target in section.findall(nodes.target): + refid = target.get("refid", "") + if fragment in target["ids"] or fragment in ( + refid, + refid.rsplit("_1", maxsplit=1)[-1], + ): + return section, title + return None + + +def _source_markdown_fragment(app, docname, link_text): + """Find the unique same-page fragment for a link in its source Markdown.""" + source_filename = _libcudf_developer_guide_source_files[docname] + source_path = os.path.join( + _libcudf_developer_guide_source_dir, source_filename + ) + with open(source_path) as source: + markdown = source.read() + + normalized_text = " ".join(link_text.replace("`", "").split()) + fragments = { + match.group("fragment") + for match in re.finditer( + r"\[([^]]+)\]\((?P[^#)]*)#(?P[^)\s]+)\)", + markdown, + ) + if " ".join(match.group(1).replace("`", "").split()) == normalized_text + and match.group("path").removeprefix("./") in ("", source_filename) + } + return fragments.pop() if len(fragments) == 1 else None + + +def rewrite_libcudf_developer_guide_references( + app: Sphinx, document: Node +) -> None: + """Rewrite Doxygen assets and local Markdown links in the developer guide.""" + if not app.env.docname.startswith("libcudf/developer_guide/"): + return + + for image in document.findall(nodes.image): + if image["uri"].endswith("cpp/doxygen/xml/strings.png"): + image["uri"] = "strings.png" + + for reference in list(document.findall(nodes.reference)): + refuri = reference.get("refuri") + if not refuri or "#" not in refuri or not reference.children: + continue + + path, fragment = refuri.split("#", maxsplit=1) + # No anchor, nothing to do. + if not fragment: + continue + if not path: + path = _libcudf_developer_guide_source_files[app.env.docname] + else: + path = path.removeprefix("./") + # Paths outside this directory remain ordinary source/external links. + if "/" in path: + continue + + target = _libcudf_developer_guide_documents.get(path) + reftarget = f"{_libcudf_developer_guide_xref_prefix}{path}#{fragment}" + xref = pending_xref( + "", + refdomain="std", + reftype="ref", + reftarget=reftarget, + refexplicit=True, + refwarn=True, + ) + xref["refdoc"] = app.env.docname + xref["libcudf_guide_target_docname"] = target + xref.source = reference.source + xref.line = reference.line + xref.extend(child.deepcopy() for child in reference.children) + reference.replace_self(xref) + + +def resolve_libcudf_developer_guide_markdown_link(app, env, node, contnode): + """Resolve a rewritten guide Markdown link or warn with its original URL.""" + reftarget = node.get("reftarget", "") + if not reftarget.startswith(_libcudf_developer_guide_xref_prefix): + return None + + markdown_target = reftarget.removeprefix( + _libcudf_developer_guide_xref_prefix + ) + target_docname = node.get("libcudf_guide_target_docname") + if target_docname is not None: + _, fragment = markdown_target.split("#", maxsplit=1) + section_data = _find_libcudf_developer_guide_section( + env, target_docname, fragment + ) + if section_data is None and target_docname == node["refdoc"]: + # Doxygen may replace a Markdown heading fragment with an + # unrelated Doxygen anchor. Recover only an unambiguous fragment + # from the original Markdown so invalid source links still warn. + source_fragment = _source_markdown_fragment( + app, target_docname, contnode.astext() + ) + if source_fragment is not None: + section_data = _find_libcudf_developer_guide_section( + env, target_docname, source_fragment + ) + if section_data is not None: + section, title = section_data + return make_refnode( + app.builder, + node["refdoc"], + target_docname, + section["ids"][0], + contnode, + title, + ) + + _libcudf_developer_guide_logger.warning( + "libcudf developer-guide Markdown link target not found: %s", + markdown_target, + location=node, + type="ref", + subtype="libcudf_guide_markdown", + ) + return contnode + + def use_slugged_duplicate_ids(app): # Use default docutils deduplication scheme for duplicate node ids. app.env.settings["auto_id_prefix"] = "%" -def setup(app): +def setup(app: Sphinx): + app.add_directive("flatdoxygenpage", FlatDoxygenPageDirective) app.connect("builder-inited", use_slugged_duplicate_ids) - app.connect("doctree-read", resolve_aliases) - app.connect("doctree-read", register_sections_as_label) - app.connect("missing-reference", on_missing_reference) + + # Do some rewrite passes on the doctrees. Lower priority hooks run + # earlier, equal priority in registration order. + # First rewrite Doxygen assets and Markdown links in the libcudf dev guide. + app.connect( + "doctree-read", + rewrite_libcudf_developer_guide_references, + priority=100, + ) + # Then rewrite xrefs in all documents for aliases. + app.connect("doctree-read", resolve_aliases, priority=200) + # Finally add std:label labels to all section headers for intersphinx + app.connect("doctree-read", register_sections_as_label, priority=300) + + # Now hook up missing-reference rewrites. First handle libcudf dev + # guide links. + app.connect( + "missing-reference", + resolve_libcudf_developer_guide_markdown_link, + priority=100, + ) + # Let intersphinx and other default-priority resolvers run first. + app.connect("missing-reference", on_missing_reference, priority=501) app.setup_extension("sphinx.ext.autodoc") app.add_autodocumenter(PLCIntEnumDocumenter) diff --git a/docs/cudf/source/cudf/10min.ipynb b/docs/cudf/source/cudf/10min.ipynb index f43b657869b8..db178e74df65 100644 --- a/docs/cudf/source/cudf/10min.ipynb +++ b/docs/cudf/source/cudf/10min.ipynb @@ -15,7 +15,7 @@ "\n", "[Dask](https://www.dask.org/) is a flexible library for parallel computing in Python that makes scaling out your workflow smooth and simple. On the CPU, Dask uses Pandas to execute operations in parallel on DataFrame partitions.\n", "\n", - "[Dask cuDF](https://github.com/NVIDIA/cudf/tree/main/python/dask_cudf) extends Dask where necessary to allow its DataFrame partitions to be processed using cuDF GPU DataFrames instead of Pandas DataFrames. For instance, when you call `dask_cudf.read_csv(...)`, your cluster's GPUs do the work of parsing the CSV file(s) by calling [`cudf.read_csv()`](https://docs.rapids.ai/api/cudf/stable/cudf/api_docs/api/cudf.read_csv/).\n", + "[Dask cuDF](https://docs.nvidia.com/dask-cudf/) extends Dask where necessary to allow its DataFrame partitions to be processed using cuDF GPU DataFrames instead of Pandas DataFrames. For instance, when you call `dask_cudf.read_csv(...)`, your cluster's GPUs do the work of parsing the CSV file(s) by calling [`cudf.read_csv()`](https://docs.nvidia.com/cudf/latest/cudf/api_docs/api/cudf.read_csv/).\n", "\n", "\n", "
\n", @@ -2570,7 +2570,7 @@ "id": "fd3fc4f3", "metadata": {}, "source": [ - "Like pandas, cuDF provides string processing methods in the `str` attribute of `Series`. Full documentation of string methods is a work in progress. Please see the [cuDF API documentation](https://docs.rapids.ai/api/cudf/stable/cudf/api_docs/series/#string-handling) for more information." + "Like pandas, cuDF provides string processing methods in the `str` attribute of `Series`. Full documentation of string methods is a work in progress. Please see the [cuDF API documentation](https://docs.nvidia.com/cudf/latest/cudf/api_docs/series/#string-handling) for more information." ] }, { @@ -2635,7 +2635,7 @@ "id": "44fe1243", "metadata": {}, "source": [ - "As well as simple manipulation, We can also match strings using [regular expressions](https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.core.accessors.string.StringMethods.match.html)." + "As well as simple manipulation, We can also match strings using [regular expressions](https://docs.nvidia.com/cudf/latest/cudf/api_docs/api/cudf.core.accessors.string.StringMethods.match/)." ] }, { diff --git a/docs/cudf/source/cudf/cupy-interop.ipynb b/docs/cudf/source/cudf/cupy-interop.ipynb index 4c09d23c5c7e..ca9cd228de53 100644 --- a/docs/cudf/source/cudf/cupy-interop.ipynb +++ b/docs/cudf/source/cudf/cupy-interop.ipynb @@ -1399,7 +1399,7 @@ "source": [ "From here, we could continue our workflow with a CuPy sparse matrix.\n", "\n", - "For a full list of the functionality built into these libraries, we encourage you to check out the API docs for [cuDF](https://docs.rapids.ai/api/cudf/nightly/) and [CuPy](https://docs.cupy.dev/en/stable/index.html)." + "For a full list of the functionality built into these libraries, we encourage you to check out the API docs for [cuDF](https://docs.nvidia.com/cudf/) and [CuPy](https://docs.cupy.dev/en/stable/index.html)." ] } ], diff --git a/docs/cudf/source/cudf/developer_guide/udf_memory_management.md b/docs/cudf/source/cudf/developer_guide/udf_memory_management.md index 264c36ee02b2..afede91051de 100644 --- a/docs/cudf/source/cudf/developer_guide/udf_memory_management.md +++ b/docs/cudf/source/cudf/developer_guide/udf_memory_management.md @@ -155,8 +155,8 @@ that view the strings owned by ``udf_string`` instances. The cuDF extensions to Numba generate code to manipulate instances of these classes, so we outline the members of these classes to aid in understanding -them. These classes also have various methods; consult the [cuDF C++ Developer -Documentation for further details of these structures.](https://docs.rapids.ai/api/libcudf/stable/developer_guide) +them. These classes also have various methods; consult the {ref}`cuDF C++ Developer +Documentation ` for further details of these structures. ```c++ class string_view { diff --git a/docs/cudf/source/cudf/guide-to-udfs.ipynb b/docs/cudf/source/cudf/guide-to-udfs.ipynb index 589c76529052..9e2483c011cb 100644 --- a/docs/cudf/source/cudf/guide-to-udfs.ipynb +++ b/docs/cudf/source/cudf/guide-to-udfs.ipynb @@ -1700,7 +1700,7 @@ "\n", "For time-series data, we may need to operate on a small \\\"window\\\" of our column at a time, processing each portion independently. We could slide (\\\"roll\\\") this window over the entire column to answer questions like \\\"What is the 3-day moving average of a stock price over the past year?\"\n", "\n", - "We can apply more complex functions to rolling windows to `rolling` Series and DataFrames using `apply`. This example is adapted from cuDF's [API documentation](https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/api/cudf.dataframe.rolling/). First, we'll create an example Series and then create a `rolling` object from the Series." + "We can apply more complex functions to rolling windows to `rolling` Series and DataFrames using `apply`. This example is adapted from cuDF's [API documentation](https://docs.nvidia.com/cudf/latest/cudf/api_docs/api/cudf.DataFrame.rolling/). First, we'll create an example Series and then create a `rolling` object from the Series." ] }, { @@ -2160,7 +2160,7 @@ "- String UDFs\n", "\n", "\n", - "For more information please see the [cuDF](https://docs.rapids.ai/api/cudf/nightly/), [Numba.cuda](https://numba.readthedocs.io/en/stable/cuda/index.html), and [CuPy](https://docs.cupy.dev/en/stable/) documentation." + "For more information please see the [cuDF](https://docs.nvidia.com/cudf/), [Numba.cuda](https://numba.readthedocs.io/en/stable/cuda/index.html), and [CuPy](https://docs.cupy.dev/en/stable/) documentation." ] } ], diff --git a/docs/cudf/source/cudf/io/io.md b/docs/cudf/source/cudf/io/io.md index ff1f96c8cd49..5d550bb62944 100644 --- a/docs/cudf/source/cudf/io/io.md +++ b/docs/cudf/source/cudf/io/io.md @@ -200,6 +200,6 @@ By default, cuDF's parquet and json readers will try to read the entire file in To better support low memory systems, cuDF provides a "low-memory" reader for parquet and json files. This low memory reader processes data in chunks, leading to lower peak memory usage due to the smaller size of intermediate allocations. -To read a parquet or json file in low memory mode, there are [cuDF options](https://docs.rapids.ai/api/cudf/nightly/cudf/api_docs/options/#api-options) that must be set globally prior to calling the reader. To set those options, call: +To read a parquet or json file in low memory mode, there are {doc}`cuDF options <../api_docs/options>` that must be set globally prior to calling the reader. To set those options, call: - `cudf.set_option("io.parquet.low_memory", True)` for parquet files, or - `cudf.set_option("io.json.low_memory", True)` for json files. diff --git a/docs/cudf/source/cudf/memory-profiling.md b/docs/cudf/source/cudf/memory-profiling.md index f69965c180a6..d984caaf1875 100644 --- a/docs/cudf/source/cudf/memory-profiling.md +++ b/docs/cudf/source/cudf/memory-profiling.md @@ -6,7 +6,7 @@ Peak memory usage is a common concern in GPU programming because GPU memory is t ## Enabling Memory Profiling -First, enable memory profiling in RMM by calling {py:func}`rmm.statistics.enable_statistics()`. This adds a statistics resource adaptor to the current RMM memory resource, which enables cuDF to access memory profiling information. See the [RMM documentation](https://docs.rapids.ai/api/rmm/stable/user_guide/guide/#memory-statistics-and-profiling) for more details. +First, enable memory profiling in RMM by calling {py:func}`rmm.statistics.enable_statistics()`. This adds a statistics resource adaptor to the current RMM memory resource, which enables cuDF to access memory profiling information. See the [RMM documentation](inv:rmm:std:label:#user_guide/guide:memory-statistics-and-profiling) for more details. Second, enable memory profiling in cuDF by setting the `memory_profiling` option to `True`. Use {py:func}`cudf.set_option` or set the environment variable ``CUDF_MEMORY_PROFILING=1`` prior to the launch of the Python interpreter. diff --git a/docs/cudf/source/cudf_pandas/faq.md b/docs/cudf/source/cudf_pandas/faq.md index 8df5f76a84aa..8241ff6507a1 100644 --- a/docs/cudf/source/cudf_pandas/faq.md +++ b/docs/cudf/source/cudf_pandas/faq.md @@ -12,8 +12,7 @@ the cuDF library directly should be considered. from increased performance by using cuDF directly. - cuDF does offer some functions and methods that pandas does not. For - example, cuDF has a [`.list` - accessor](https://docs.rapids.ai/api/cudf/stable/user_guide/api_docs/series/#list-handling) + example, cuDF has a {ref}`.list accessor ` for working with list-like data. If you need access to the additional functionality in cuDF, you will need to use the cuDF package directly. @@ -139,7 +138,7 @@ Both Dask and Apache Spark support accelerated computing through configuration based interfaces. Dask allows you to [configure the dataframe backend](https://docs.dask.org/en/latest/how-to/selecting-the-collection-backend.html) to use cuDF (learn more in [this -blog](https://medium.com/rapids-ai/easy-cpu-gpu-arrays-and-dataframes-run-your-dask-code-where-youd-like-e349d92351d)) and the [RAPIDS Accelerator for Apache Spark](https://nvidia.github.io/spark-rapids/) +blog](https://medium.com/rapids-ai/easy-cpu-gpu-arrays-and-dataframes-run-your-dask-code-where-youd-like-e349d92351d)) and the [RAPIDS Accelerator for Apache Spark](https://docs.nvidia.com/spark-rapids/) provides a similar configuration-based plugin for Spark. ## How do I know if an object is a `cudf.pandas` proxy object? diff --git a/docs/cudf/source/cudf_pandas/index.rst b/docs/cudf/source/cudf_pandas/index.rst index 964dae75ebf0..f9639012d269 100644 --- a/docs/cudf/source/cudf_pandas/index.rst +++ b/docs/cudf/source/cudf_pandas/index.rst @@ -34,8 +34,10 @@ automatically **falling back to pandas** for other operations. | Nothing changes, not even your `import` statements, when going from CPU to GPU. | Combines the full flexibility of Pandas with blazing fast performance of cuDF | +---------------------------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------+ -``cudf.pandas`` is now Generally Available (GA) as part of the ``cudf`` package. See `RAPIDS -Quick Start `_ to get up-and-running with ``cudf``. +``cudf.pandas`` is available as part of the ``cudf`` package. See the +`installation and deployment guide +_` +to get up-and-running with cuDF. .. toctree:: :maxdepth: 1 diff --git a/docs/cudf/source/cudf_polars/benchmarks.md b/docs/cudf/source/cudf_polars/benchmarks.md index 969be5a0b5a6..0dc38068fc56 100644 --- a/docs/cudf/source/cudf_polars/benchmarks.md +++ b/docs/cudf/source/cudf_polars/benchmarks.md @@ -9,7 +9,7 @@ The steps below reproduce the PDS-H benchmark results using the Polars GPU engin ### Setup Install `cudf-polars` following the -[RAPIDS installation guide](https://docs.rapids.ai/install/). For nightly wheels, install with +[NVIDIA CUDA-X installation guide](https://docs.rapids.ai/install/#install-rapids). For nightly wheels, install with the `ray` extra (required for multi-GPU benchmarking): ```bash diff --git a/docs/cudf/source/cudf_polars/dask_engine.md b/docs/cudf/source/cudf_polars/dask_engine.md index 1cfa4d9f7ad4..6112ba45bb57 100644 --- a/docs/cudf/source/cudf_polars/dask_engine.md +++ b/docs/cudf/source/cudf_polars/dask_engine.md @@ -194,5 +194,5 @@ created inside an `rrun` cluster. [dask-distributed]: https://distributed.dask.org/en/stable/ [dask-cli]: https://docs.dask.org/en/latest/deploying-cli.html -[dask-cuda]: https://docs.rapids.ai/api/dask-cuda/nightly/ -[dask-cuda-worker]: https://docs.rapids.ai/api/dask-cuda/nightly/quickstart/#dask-cuda-worker +[dask-cuda]: inv:dask-cuda:std:doc:#index +[dask-cuda-worker]: diff --git a/docs/cudf/source/cudf_polars/developer_docs.md b/docs/cudf/source/cudf_polars/developer_docs.md index 069ee07ce00b..12c4ea511c7c 100644 --- a/docs/cudf/source/cudf_polars/developer_docs.md +++ b/docs/cudf/source/cudf_polars/developer_docs.md @@ -2,7 +2,7 @@ You will need: -1. Rust development environment. If you use the rapids [combined +1. Rust development environment. If you use the [combined devcontainer](https://github.com/rapidsai/devcontainers/), add `"./features/src/rust": {"version": "latest", "profile": "default"},` to your preferred configuration. Or else, use @@ -625,25 +625,21 @@ another `nvtx` range (e.g. `Scan.do_evaluate`, `GroupBy.do_evaluate`, etc.). These provide a higher-level grouping over the lower-level libcudf calls (e.g. `read_chunk`, `aggregate`). -Finally, if using [rapidsmpf](https://docs.rapids.ai/api/rapidsmpf/nightly/) -for shuffling, the methods inserting and extracting partitions to shuffle are -annotated with nvtx ranges. - # Query Plans -The module `cudf_polars.experimental.explain` contains functions for dumping +The module `cudf_polars.streaming.explain` contains functions for dumping the query for a given `LazyFrame`. ## Structured Output -`cudf_polars.experimental.explain.serialize_query` can be used to output +`cudf_polars.streaming.explain.serialize_query` can be used to output the query plan in a structured format. ```python >>> import dataclasses >>> import polars as pl ->>> from cudf_polars.experimental.explain import serialize_query +>>> from cudf_polars.streaming.explain import serialize_query >>> q = pl.LazyFrame({"a": ['a', 'b', 'a'], "b": [1, 2, 3]}).group_by("a").agg(pl.len()) >>> dataclasses.asdict(serialize_query(q, engine=pl.GPUEngine())) {'roots': ['526964741'], diff --git a/docs/cudf/source/cudf_polars/index.md b/docs/cudf/source/cudf_polars/index.md index 9f641f33a024..e076634f7ce0 100644 --- a/docs/cudf/source/cudf_polars/index.md +++ b/docs/cudf/source/cudf_polars/index.md @@ -9,8 +9,10 @@ and runs on the CPU. ## Install -Follow the [RAPIDS installation guide](https://docs.rapids.ai/install/) and pick the -`cudf-polars` package for your CUDA and Python versions. For example, with conda: +Follow the [NVIDIA CUDA-X installation +guide](https://docs.rapids.ai/install/#install-rapids) +and pick the `cudf-polars` package for your CUDA and Python versions. For +example, with conda: ```bash conda install -c rapidsai -c conda-forge -c nvidia cudf-polars diff --git a/docs/cudf/source/cudf_polars/memory_errors.md b/docs/cudf/source/cudf_polars/memory_errors.md index f39a4938d4a9..693a2e5699f8 100644 --- a/docs/cudf/source/cudf_polars/memory_errors.md +++ b/docs/cudf/source/cudf_polars/memory_errors.md @@ -110,4 +110,6 @@ constructing the GPU engine for queries. For the full list of engine configuration options, including `target_partition_size` and `max_concurrent_io_tasks`, see {doc}`options`. For the full list of memory -and spill configuration options see the [RapidsMPF configuration reference](https://docs.rapids.ai/api/rapidsmpf/stable/configuration/#general). +and spill configuration options see the [RapidsMPF configuration reference][rapidsmpf-config]. + +[rapidsmpf-config]: inv:rapidsmpf:std:label:#configuration:general diff --git a/docs/cudf/source/cudf_polars/options.md b/docs/cudf/source/cudf_polars/options.md index d4af20497ff9..72db94eb7499 100644 --- a/docs/cudf/source/cudf_polars/options.md +++ b/docs/cudf/source/cudf_polars/options.md @@ -109,7 +109,7 @@ Environment variables follow these patterns: | `target_partition_size` | Target partition size in bytes. Used for IO and dynamic planning. `0` means auto. | auto | | `max_concurrent_io_tasks` | Number of concurrent IO producer tasks for each scan node. Tune with an integer or a `{"local": ..., "remote": ...}` dict. | 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 | +| `join_filter_pushdown` | Configuration for join filter pushdown plan rewrites, dict or {class}`~cudf_polars.utils.config.JoinFilterPushdownOptions`. `None` disables. | disabled | | `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` @@ -136,4 +136,4 @@ These environment variables are intended for library developers and advanced use | `CUDF_POLARS_WARN_UNSTABLE` | Raises a `cudf_polars.UnstableWarning` whenever an unstable cudf-polars feature is used. Set to `1` to enable. | `0` | -[rapidsmpf-config]: https://docs.rapids.ai/api/rapidsmpf/nightly/configuration/ +[rapidsmpf-config]: inv:rapidsmpf:std:doc:#configuration diff --git a/docs/cudf/source/cudf_polars/profiling.md b/docs/cudf/source/cudf_polars/profiling.md index 35e0baa7ffe6..526ab6c473e1 100644 --- a/docs/cudf/source/cudf_polars/profiling.md +++ b/docs/cudf/source/cudf_polars/profiling.md @@ -103,8 +103,8 @@ KvikIO I/O summary ``` Every row is also an attribute, `s.bytes_read`, `s.busy_ns` and so on. See the -[KvikIO reference][kvikio-stats] for the full set, and [busy time and bandwidth][kvikio-busy] -for how the busy figures are measured. +[KvikIO statistics reference][kvikio-stats] for the full set, and [busy time and +bandwidth][kvikio-busy] for how the busy figures are measured. ### What is and is not counted @@ -250,9 +250,9 @@ shape: (2, 3) [nsight]: https://developer.nvidia.com/nsight-systems [nvtx]: https://nvidia.github.io/NVTX/ -[kvikio-stats]: https://docs.rapids.ai/api/kvikio/nightly/statistics/ -[kvikio-busy]: https://docs.rapids.ai/api/kvikio/nightly/statistics/#busy-time-and-bandwidth -[rapidsmpf-stats]: https://docs.rapids.ai/api/rapidsmpf/nightly/statistics/ +[kvikio-stats]: inv:kvikio:std:doc:#statistics +[kvikio-busy]: +[rapidsmpf-stats]: inv:rapidsmpf:std:doc:#statistics [structlog]: https://www.structlog.org/en/stable/ [structlog-configure]: https://www.structlog.org/en/stable/configuration.html [structlog-context]: https://www.structlog.org/en/stable/contextvars.html diff --git a/docs/cudf/source/index.rst b/docs/cudf/source/index.rst index 373fa95c54ef..c74f32f8bf00 100644 --- a/docs/cudf/source/index.rst +++ b/docs/cudf/source/index.rst @@ -1,9 +1,10 @@ NVIDIA cuDF Documentation ========================= -**NVIDIA cuDF** (pronounced "KOO-dee-eff") is a GPU-accelerated library for tabular -data processing. It is part of the `RAPIDS `_ suite of -libraries and is composed of multiple sub-projects: +**NVIDIA cuDF** (pronounced "KOO-dee-eff") is a GPU-accelerated library for +tabular data processing. It is part of `NVIDIA CUDA-X for Data Science +`_ +suite, and is composed of multiple sub-projects: .. list-table:: :header-rows: 1 @@ -11,15 +12,15 @@ libraries and is composed of multiple sub-projects: * - Library - Description - * - `cudf `_ - - A Python library providing a `pandas `_-like DataFrame API and a zero-code change accelerator, `cudf.pandas `_, for existing pandas code. - * - `cudf-polars `_ + * - :doc:`cudf ` + - A Python library providing a `pandas `_-like DataFrame API and a zero-code change accelerator, :doc:`cudf.pandas `, for existing pandas code. + * - :doc:`cudf-polars ` - A Python library providing a GPU engine for `Polars `_. - * - `dask-cudf `_ + * - :doc:`dask-cudf ` - A Python library providing a GPU backend for `Dask `_ DataFrames. - * - `libcudf `_ + * - :doc:`libcudf ` - A CUDA C++ library with `Apache Arrow `_ compliant data structures and fundamental algorithms for tabular data. - * - `pylibcudf `_ + * - :doc:`pylibcudf ` - A Python library providing `Cython `_ bindings for libcudf. Accelerated Data Engines and Tools @@ -42,10 +43,10 @@ The following data engines and tools integrate with cuDF: - `Sirius documentation `_ * - pandas - cudf.pandas - - `cudf.pandas documentation `_ + - :doc:`cudf.pandas documentation ` * - Polars - Polars GPU engine - - `Polars GPU engine documentation `_ + - :doc:`Polars GPU engine documentation ` * - Presto - Presto-GPU - `Presto on GPU tutorial `_ @@ -53,6 +54,10 @@ The following data engines and tools integrate with cuDF: - Velox on GPU (experimental) - `Velox-cuDF documentation `_ +See the `installation and deployment guide +`_ +to get up-and-running with cuDF. + .. toctree:: :maxdepth: 1 :caption: Libraries diff --git a/docs/cudf/source/libcudf/api_docs/lists_classes.rst b/docs/cudf/source/libcudf/api_docs/lists_classes.rst index 9b89c1647466..9444b202a123 100644 --- a/docs/cudf/source/libcudf/api_docs/lists_classes.rst +++ b/docs/cudf/source/libcudf/api_docs/lists_classes.rst @@ -3,3 +3,6 @@ Lists Classes .. doxygengroup:: lists_classes :members: + +.. doxygenclass:: cudf::list_view + :project: libcudf diff --git a/docs/cudf/source/libcudf/api_docs/structs_classes.rst b/docs/cudf/source/libcudf/api_docs/structs_classes.rst index 2669c2884d63..2f6e2be7c02c 100644 --- a/docs/cudf/source/libcudf/api_docs/structs_classes.rst +++ b/docs/cudf/source/libcudf/api_docs/structs_classes.rst @@ -3,3 +3,6 @@ Structs Classes .. doxygengroup:: structs_classes :members: + +.. doxygenclass:: cudf::struct_view + :project: libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/BENCHMARKING.rst b/docs/cudf/source/libcudf/developer_guide/BENCHMARKING.rst new file mode 100644 index 000000000000..15b860a28146 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/BENCHMARKING.rst @@ -0,0 +1,8 @@ +.. _md_developer_guide_benchmarking: +.. _md_doxygen_developer_guide_BENCHMARKING: + +Unit Benchmarking in libcudf +============================ + +.. flatdoxygenpage:: md_doxygen_developer_guide_BENCHMARKING + :project: libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/DEVELOPER_GUIDE.rst b/docs/cudf/source/libcudf/developer_guide/DEVELOPER_GUIDE.rst new file mode 100644 index 000000000000..8fe5f11ac7b9 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/DEVELOPER_GUIDE.rst @@ -0,0 +1,17 @@ +.. _md_developer_guide: +.. _DEVELOPER_GUIDE: + +libcudf C++ Developer Guide +=========================== + +.. flatdoxygenpage:: DEVELOPER_GUIDE + :project: libcudf + +.. toctree:: + :hidden: + :maxdepth: 1 + + libcudf C++ Documentation Guide + Unit Testing in libcudf + Unit Benchmarking in libcudf + Profiling libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/DOCUMENTATION.rst b/docs/cudf/source/libcudf/developer_guide/DOCUMENTATION.rst new file mode 100644 index 000000000000..159c40a37ef2 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/DOCUMENTATION.rst @@ -0,0 +1,8 @@ +.. _md_developer_guide_documentation: +.. _md_doxygen_developer_guide_DOCUMENTATION: + +libcudf C++ Documentation Guide +=============================== + +.. flatdoxygenpage:: md_doxygen_developer_guide_DOCUMENTATION + :project: libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/PROFILING.rst b/docs/cudf/source/libcudf/developer_guide/PROFILING.rst new file mode 100644 index 000000000000..2047dd2062d4 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/PROFILING.rst @@ -0,0 +1,8 @@ +.. _md_developer_guide_profiling: +.. _md_doxygen_developer_guide_PROFILING: + +Profiling libcudf +================= + +.. flatdoxygenpage:: md_doxygen_developer_guide_PROFILING + :project: libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/TESTING.rst b/docs/cudf/source/libcudf/developer_guide/TESTING.rst new file mode 100644 index 000000000000..93910cbfb024 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/TESTING.rst @@ -0,0 +1,8 @@ +.. _md_developer_guide_testing: +.. _md_doxygen_developer_guide_TESTING: + +Unit Testing in libcudf +======================= + +.. flatdoxygenpage:: md_doxygen_developer_guide_TESTING + :project: libcudf diff --git a/docs/cudf/source/libcudf/developer_guide/strings.png b/docs/cudf/source/libcudf/developer_guide/strings.png new file mode 120000 index 000000000000..acde8bb2dc62 --- /dev/null +++ b/docs/cudf/source/libcudf/developer_guide/strings.png @@ -0,0 +1 @@ +../../../../../cpp/doxygen/developer_guide/strings.png \ No newline at end of file diff --git a/docs/cudf/source/libcudf/index.rst b/docs/cudf/source/libcudf/index.rst index 9f390f647ef0..7aa61fe3209f 100644 --- a/docs/cudf/source/libcudf/index.rst +++ b/docs/cudf/source/libcudf/index.rst @@ -6,5 +6,6 @@ libcudf :caption: Contents: api_docs/index.rst + developer_guide/DEVELOPER_GUIDE md_regex unicode_limitations diff --git a/docs/cudf/source/libcudf/md_regex.rst b/docs/cudf/source/libcudf/md_regex.rst index 0eb0f464063a..3dff835a2280 100644 --- a/docs/cudf/source/libcudf/md_regex.rst +++ b/docs/cudf/source/libcudf/md_regex.rst @@ -1,4 +1,7 @@ -.. _md_regex: +.. _mr::md_regex: -.. include:: ../../../../cpp/doxygen/regex.md - :parser: myst_parser.sphinx_ +Regex Features +============== + +.. flatdoxygenpage:: md_regex + :project: libcudf diff --git a/docs/cudf/source/libcudf/unicode_limitations.rst b/docs/cudf/source/libcudf/unicode_limitations.rst index 1f0690881606..1a53e08f7f5a 100644 --- a/docs/cudf/source/libcudf/unicode_limitations.rst +++ b/docs/cudf/source/libcudf/unicode_limitations.rst @@ -1,4 +1,7 @@ -.. _unicode_limitations: +.. _mr::md_doxygen_unicode: -.. include:: ../../../../cpp/doxygen/unicode.md - :parser: myst_parser.sphinx_ +Unicode Limitations +=================== + +.. flatdoxygenpage:: md_doxygen_unicode + :project: libcudf diff --git a/docs/dask_cudf/source/_static/RAPIDS-logo-purple.png b/docs/dask_cudf/source/_static/RAPIDS-logo-purple.png deleted file mode 100644 index d884e01374dc..000000000000 Binary files a/docs/dask_cudf/source/_static/RAPIDS-logo-purple.png and /dev/null differ diff --git a/docs/dask_cudf/source/best_practices.rst b/docs/dask_cudf/source/best_practices.rst index 675e5fc1c11a..037de0e1c5cb 100644 --- a/docs/dask_cudf/source/best_practices.rst +++ b/docs/dask_cudf/source/best_practices.rst @@ -3,8 +3,8 @@ Dask cuDF Best Practices ======================== -This page outlines several important guidelines for using `Dask cuDF -`__ effectively. +This page outlines several important guidelines for using +:doc:`Dask cuDF ` effectively. .. note:: Since Dask cuDF is a backend extension for @@ -22,24 +22,23 @@ Use Dask-CUDA ~~~~~~~~~~~~~ To execute a Dask workflow on multiple GPUs, a Dask cluster must -be deployed with `Dask-CUDA `__ +be deployed with :doc:`Dask-CUDA ` and `Dask.distributed `__. -When running on a single machine, the `LocalCUDACluster `__ +When running on a single machine, the :class:`~dask_cuda.LocalCUDACluster` convenience function is strongly recommended. No matter how many GPUs are -available on the machine (even one!), using `Dask-CUDA has many advantages -`__ +available on the machine (even one!), using :ref:`Dask-CUDA has many advantages +` over default (threaded) execution. Just to list a few: * Dask-CUDA makes it easy to pin workers to specific devices. * Dask-CUDA makes it easy to configure memory-spilling options. * The distributed scheduler collects useful diagnostic information that can be viewed on a dashboard in real time. -Please see `Dask-CUDA's API `__ -and `Best Practices `__ +Please see :doc:`Dask-CUDA's API ` +and :doc:`Best Practices ` documentation for detailed information. Typical ``LocalCUDACluster`` usage -is also illustrated within the multi-GPU section of `Dask cuDF's -`__ documentation. +is also shown in :ref:`multiple_gpus`. .. note:: When running on cloud infrastructure or HPC systems, it is usually best to @@ -47,7 +46,7 @@ is also illustrated within the multi-GPU section of `Dask cuDF's `__ and `Dask-Jobqueue `__. - Please see `the RAPIDS deployment documentation `__ + Please see `the cloud deployment documentation `__ for further details and examples. @@ -71,23 +70,23 @@ Enable cuDF spilling ~~~~~~~~~~~~~~~~~~~~ When using Dask cuDF for classic ETL workloads, it is usually best -to enable `native spilling support in cuDF -`__. -When using :class:`dask_cuda.LocalCUDACluster`, this is easily accomplished by +to enable :ref:`native spilling support in cuDF +`. +When using :class:`~dask_cuda.LocalCUDACluster`, this is easily accomplished by setting ``enable_cudf_spill=True``. Use RMM ~~~~~~~ Memory allocations in cuDF are significantly faster and more efficient when -the `RAPIDS Memory Manager (RMM) `__ -library is configured appropriately on worker processes. In most cases, the best way to manage +:doc:`NVIDIA RMM ` +is configured appropriately on worker processes. In most cases, the best way to manage memory is by initializing an RMM pool on each worker before executing a -workflow. When using :class:`dask_cuda.LocalCUDACluster`, this is easily accomplished +workflow. When using :class:`~dask_cuda.LocalCUDACluster`, this is easily accomplished by setting ``rmm_pool_size`` to a large fraction (e.g. ``0.9``). -See the `Dask-CUDA memory-management documentation -`__ +See the :ref:`Dask-CUDA memory-management documentation +` for more details. Use the Dask DataFrame API @@ -289,15 +288,15 @@ bottleneck is typically device-to-host memory spilling. Although every workflow is different, the following guidelines are often recommended: -* Use a distributed cluster with `Dask-CUDA `__ workers +* Use a distributed cluster with :doc:`Dask-CUDA ` workers -* Use native cuDF spilling whenever possible (`Dask-CUDA spilling documentation `__) +* Use native cuDF spilling whenever possible (:doc:`Dask-CUDA spilling documentation `) * Avoid shuffling whenever possible * Use ``split_out=1`` for low-cardinality groupby aggregations * Use ``broadcast=True`` for joins when at least one collection comprises a small number of partitions (e.g. ``<=5``) -* `Use UCX `__ if communication is a bottleneck. +* :doc:`Use UCX ` if communication is a bottleneck. .. note:: UCX enables Dask-CUDA workers to communicate using high-performance diff --git a/docs/dask_cudf/source/conf.py b/docs/dask_cudf/source/conf.py index 635ba31f5754..15f9f98cb496 100644 --- a/docs/dask_cudf/source/conf.py +++ b/docs/dask_cudf/source/conf.py @@ -59,7 +59,7 @@ htmlhelp_basename = "dask-cudfdoc" html_use_modindex = True -html_static_path = ["_static"] +html_static_path = [] pygments_style = "sphinx" @@ -78,16 +78,23 @@ } include_pandas_compat = True +with open("../../../RAPIDS_BRANCH", "r") as f: + branch = f.read().strip() +intersphinx_version = "latest" if branch == "main" else version + intersphinx_mapping = { "python": ("https://docs.python.org/3/", None), "cupy": ("https://docs.cupy.dev/en/stable/", None), "numpy": ("https://numpy.org/doc/stable/", None), "pyarrow": ("https://arrow.apache.org/docs/", None), - "cudf": ("https://docs.rapids.ai/api/cudf/stable/", None), + "cudf": (f"https://docs.nvidia.com/cudf/{intersphinx_version}/", None), "dask": ("https://docs.dask.org/en/stable/", None), - # Temporarily disable pandas intersphinx: https://github.com/pandas-dev/pandas/issues/64584 - # "pandas": ("https://pandas.pydata.org/docs/", None), - "dask-cuda": ("https://docs.rapids.ai/api/dask-cuda/stable/", None), + "pandas": ("https://pandas.pydata.org/docs/", None), + "dask-cuda": ( + f"https://docs.nvidia.com/dask-cuda/{intersphinx_version}/", + None, + ), + "rmm": (f"https://docs.nvidia.com/rmm/{intersphinx_version}/", None), } numpydoc_show_inherited_class_members = True diff --git a/docs/dask_cudf/source/index.rst b/docs/dask_cudf/source/index.rst index eee1bc39fc4b..0bf4f5a4f50b 100644 --- a/docs/dask_cudf/source/index.rst +++ b/docs/dask_cudf/source/index.rst @@ -21,12 +21,11 @@ as the ``"cudf"`` dataframe backend for of the GPU and networking hardware. If you are familiar with Dask and `pandas `__ or -`cuDF `__, then Dask cuDF +:doc:`cuDF `, then Dask cuDF should feel familiar to you. If not, we recommend starting with `10 minutes to Dask `__ followed -by `10 minutes to cuDF and Dask cuDF -`__. +by :doc:`10 minutes to cuDF and Dask cuDF `. After reviewing the sections below, please see the :ref:`Best Practices ` page for further guidance on @@ -120,7 +119,7 @@ automatic query planning (see the next section). Query Planning ~~~~~~~~~~~~~~ -Dask cuDF now provides automatic query planning by default (RAPIDS 24.06+). +Since version 24.06, Dask cuDF provides automatic query planning by default. As long as the ``"dataframe.query-planning"`` configuration is set to ``True`` (the default) when ``dask.dataframe`` is first imported, `Dask Expressions `__ will be used under the hood. @@ -149,6 +148,8 @@ Simplified expression graph (``df.simplify().pprint()``):: (via :func:`dask.compute` or :func:`dask.persist`). You do not need to optimize or simplify the graph yourself. +.. _multiple_gpus: + Using Multiple GPUs and Multiple Nodes ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -193,8 +194,7 @@ to define a client object. For example:: Please see the :doc:`dask-cuda:index` documentation for more information about deploying GPU-aware clusters -(including `best practices -`__). +(including :doc:`best practices `). API Reference @@ -202,7 +202,7 @@ API Reference Generally speaking, Dask cuDF tries to offer exactly the same API as Dask DataFrame. There are, however, some minor differences mostly because -cuDF does not `perfectly mirror `__ +cuDF does not :doc:`perfectly mirror ` the pandas API, or because cuDF provides additional configuration flags (these mostly occur in data reading and writing interfaces). diff --git a/java/pom.xml b/java/pom.xml index c72f0445d5c8..a6df10620007 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -17,7 +17,7 @@ This project provides java bindings for cudf, to be able to process large amounts of data on a GPU. This is still a work in progress so some APIs may change until the 1.0 release. - https://rapids.ai/ + https://docs.nvidia.com/cudf/ diff --git a/java/src/main/java/ai/rapids/cudf/ColumnView.java b/java/src/main/java/ai/rapids/cudf/ColumnView.java index 269a28ad8549..4c00bf3c4e5d 100644 --- a/java/src/main/java/ai/rapids/cudf/ColumnView.java +++ b/java/src/main/java/ai/rapids/cudf/ColumnView.java @@ -3876,7 +3876,7 @@ public final ColumnVector clamp(Scalar lo, Scalar loReplace, Scalar hi, Scalar h * ``` * Any null string entries return corresponding null output column entries. * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * * @param pattern Regex pattern to match to each string. * @return New ColumnVector of boolean results for each string. @@ -3898,7 +3898,7 @@ public final ColumnVector matchesRe(String pattern) { * ``` * Any null string entries return corresponding null output column entries. * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * * @param regexProg Regex program to match to each string. * @return New ColumnVector of boolean results for each string. @@ -3922,7 +3922,7 @@ public final ColumnVector matchesRe(RegexProgram regexProg) { * ``` * Any null string entries return corresponding null output column entries. * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * * @param pattern Regex pattern to match to each string. * @return New ColumnVector of boolean results for each string. @@ -3944,7 +3944,7 @@ public final ColumnVector containsRe(String pattern) { * ``` * Any null string entries return corresponding null output column entries. * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * * @param regexProg Regex program to match to each string. * @return New ColumnVector of boolean results for each string. @@ -3963,7 +3963,7 @@ public final ColumnVector containsRe(RegexProgram regexProg) { * does not match. Any null inputs also result in null output entries. * * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * @param pattern the pattern to use * @return the table of extracted matches * @throws CudfException if any error happens including if the RE does @@ -3980,7 +3980,7 @@ public final Table extractRe(String pattern) throws CudfException { * does not match. Any null inputs also result in null output entries. * * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * @param regexProg the regex program to use * @return the table of extracted matches * @throws CudfException if any error happens including if the regex @@ -3998,7 +3998,7 @@ public final Table extractRe(RegexProgram regexProg) throws CudfException { * regular expression group index. Any null inputs also result in null output entries. * * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * @param pattern The regex pattern * @param idx The regex group index * @return A new column vector of extracted matches @@ -4016,7 +4016,7 @@ public final ColumnVector extractAllRecord(String pattern, int idx) { * regular expression group index. Any null inputs also result in null output entries. * * For supported regex patterns refer to: - * @link https://docs.rapids.ai/api/libcudf/nightly/md_regex.html + * @link https://docs.nvidia.com/cudf/latest/libcudf/md_regex/ * @param regexProg The regex program * @param idx The regex group index * @return A new column vector of extracted matches diff --git a/java/src/main/native/src/RmmJni.cpp b/java/src/main/native/src/RmmJni.cpp index e5f03fba9946..771975d992e2 100644 --- a/java/src/main/native/src/RmmJni.cpp +++ b/java/src/main/native/src/RmmJni.cpp @@ -840,7 +840,7 @@ class pinned_fallback_host_memory_resource { // If the pool is exhausted, fall back to the upstream memory resource } } - return prior_cudf_pinned_mr().allocate(stream, bytes); + return prior_cudf_pinned_mr().allocate(stream, bytes, alignment); } void deallocate(cuda::stream_ref stream, @@ -851,7 +851,7 @@ class pinned_fallback_host_memory_resource { if (bytes <= pool.pool_size() && ptr >= pool_begin && ptr < pool_end) { pool.deallocate(stream, ptr, bytes, alignment); } else { - prior_cudf_pinned_mr().deallocate(stream, ptr, bytes); + prior_cudf_pinned_mr().deallocate(stream, ptr, bytes, alignment); } } @@ -938,7 +938,7 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocInternal(JNIEnv* env, cudf::jni::auto_set_device(env); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); auto c_stream = cuda::stream_ref(reinterpret_cast(stream)); - void* ret = mr.allocate(c_stream, size); + void* ret = mr.allocate(c_stream, size, rmm::CUDA_ALLOCATION_ALIGNMENT); return reinterpret_cast(ret); } JNI_CATCH(env, 0); @@ -953,7 +953,7 @@ Java_ai_rapids_cudf_Rmm_free(JNIEnv* env, jclass clazz, jlong ptr, jlong size, j rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); void* cptr = reinterpret_cast(ptr); auto c_stream = cuda::stream_ref(reinterpret_cast(stream)); - mr.deallocate(c_stream, cptr, size); + mr.deallocate(c_stream, cptr, size, rmm::CUDA_ALLOCATION_ALIGNMENT); } JNI_CATCH(env, ); } @@ -1449,7 +1449,8 @@ JNIEXPORT jlong JNICALL Java_ai_rapids_cudf_Rmm_allocFromFallbackPinnedPool(JNIE JNI_TRY { cudf::jni::auto_set_device(env); - void* ret = cudf::get_pinned_memory_resource().allocate(cudf::get_default_stream(), size); + void* ret = cudf::get_pinned_memory_resource().allocate( + cudf::get_default_stream(), size, rmm::CUDA_ALLOCATION_ALIGNMENT); return reinterpret_cast(ret); } JNI_CATCH(env, 0); @@ -1465,7 +1466,8 @@ JNIEXPORT void JNICALL Java_ai_rapids_cudf_Rmm_freeFromFallbackPinnedPool(JNIEnv { cudf::jni::auto_set_device(env); void* cptr = reinterpret_cast(ptr); - cudf::get_pinned_memory_resource().deallocate(cudf::get_default_stream(), cptr, size); + cudf::get_pinned_memory_resource().deallocate( + cudf::get_default_stream(), cptr, size, rmm::CUDA_ALLOCATION_ALIGNMENT); } JNI_CATCH(env, ); } diff --git a/python/cudf/cudf/core/dataframe.py b/python/cudf/cudf/core/dataframe.py index a08ed067c5ee..feae07611dfb 100644 --- a/python/cudf/cudf/core/dataframe.py +++ b/python/cudf/cudf/core/dataframe.py @@ -5500,8 +5500,8 @@ def apply( Thus the allowed operations within ``func`` are limited to `those supported by the CUDA Python Numba target `__. - For more information, see the `cuDF guide to user defined functions - `__. + For more information, see the :doc:`cuDF guide to user defined functions + `. Some string functions and methods are supported. Refer to the guide to UDFs for details. @@ -5684,8 +5684,8 @@ def apply( >>> df.apply(f, axis=1) # doctest: +SKIP For a complete list of supported functions and methods that may be - used to manipulate string data, see the UDF guide, - + used to manipulate string data, see the :doc:`UDF guide + ` """ if axis != 1: raise NotImplementedError( diff --git a/python/cudf/cudf/core/groupby/groupby.py b/python/cudf/cudf/core/groupby/groupby.py index ce3bb1dcd28e..eed40a6d9d76 100644 --- a/python/cudf/cudf/core/groupby/groupby.py +++ b/python/cudf/cudf/core/groupby/groupby.py @@ -519,7 +519,7 @@ def _collect_series_key_column_names(obj, by) -> dict[int, Hashable]: class GroupByNthSelector: - """Mirror of :class:`pandas.core.groupby.indexing.GroupByNthSelector`. + """Mirror of ``pandas.core.groupby.indexing.GroupByNthSelector``. ``GroupBy.nth`` supports both the call form ``gb.nth(n, dropna=...)`` and the index form ``gb.nth[n]``. @@ -1505,8 +1505,8 @@ def _reduce( Computed {op} of values within each group. .. pandas-compat:: - :meth:`pandas.core.groupby.DataFrameGroupBy.{op}`, - :meth:`pandas.core.groupby.SeriesGroupBy.{op}` + :meth:`pandas.api.typing.DataFrameGroupBy.{op}`, + :meth:`pandas.api.typing.SeriesGroupBy.{op}` The numeric_only, min_count """ @@ -2594,8 +2594,8 @@ def apply( std, idxmax, and idxmin and any arithmetic formula involving them are allowed. Binary operations are not yet supported, so syntax like `df['x'] * 2` is not yet allowed. - For more information, see the `cuDF guide to user defined functions - `__. + For more information, see the :doc:`cuDF guide to user defined functions + `. Use `cudf` to select the iterative groupby apply algorithm which aims to provide maximum flexibility at the expense of performance. The default value `auto` will attempt to use the numba JIT pipeline @@ -2640,8 +2640,8 @@ def mult(df): 6 2 6 12 .. pandas-compat:: - :meth:`pandas.core.groupby.DataFrameGroupBy.apply`, - :meth:`pandas.core.groupby.SeriesGroupBy.apply` + :meth:`pandas.api.typing.DataFrameGroupBy.apply`, + :meth:`pandas.api.typing.SeriesGroupBy.apply` cuDF's ``groupby.apply`` is limited compared to pandas. In some situations, Pandas returns the grouped keys as part of @@ -3593,8 +3593,8 @@ def shift( Object shifted within each group. .. pandas-compat:: - :meth:`pandas.core.groupby.DataFrameGroupBy.shift`, - :meth:`pandas.core.groupby.SeriesGroupBy.shift` + :meth:`pandas.api.typing.DataFrameGroupBy.shift`, + :meth:`pandas.api.typing.SeriesGroupBy.shift` Parameter ``freq`` is unsupported. """ diff --git a/python/cudf/cudf/core/indexed_frame.py b/python/cudf/cudf/core/indexed_frame.py index 37b88e55c19a..53ce03856b43 100644 --- a/python/cudf/cudf/core/indexed_frame.py +++ b/python/cudf/cudf/core/indexed_frame.py @@ -434,7 +434,7 @@ def flags(self) -> pd.Flags: The available flags are - * :attr:`pandas.Flags.allows_duplicate_labels` + * ``allows_duplicate_labels`` See Also -------- diff --git a/python/cudf/cudf/core/multiindex.py b/python/cudf/cudf/core/multiindex.py index 5be6335601cd..cbfa76c83055 100644 --- a/python/cudf/cudf/core/multiindex.py +++ b/python/cudf/cudf/core/multiindex.py @@ -594,7 +594,7 @@ def __repr__(self) -> str: @property @_external_only_api("Use ._codes instead") @_performance_tracking - def codes(self) -> pd.core.indexes.frozen.FrozenList: + def codes(self) -> pd.api.typing.FrozenList: """ Returns the codes of the underlying MultiIndex. diff --git a/python/cudf/cudf/core/series.py b/python/cudf/cudf/core/series.py index dcc6f0df164b..33637be6f0b2 100644 --- a/python/cudf/cudf/core/series.py +++ b/python/cudf/cudf/core/series.py @@ -2618,8 +2618,8 @@ def apply( Thus the allowed operations within ``func`` are limited to `those supported by the CUDA Python Numba target `__. - For more information, see the `cuDF guide to user defined functions - `__. + For more information, see the :doc:`cuDF guide to user defined functions + `. Some string functions and methods are supported. Refer to the guide to UDFs for details. @@ -2751,8 +2751,8 @@ def apply( >>> sr.apply(f) # doctest: +SKIP For a complete list of supported functions and methods that may be - used to manipulate string data, see the UDF guide, - + used to manipulate string data, see the :doc:`UDF guide + ` """ if convert_dtype is not True: diff --git a/python/cudf/cudf/core/udf/groupby_typing.py b/python/cudf/cudf/core/udf/groupby_typing.py index 42e3fa2a5194..e81f816f1708 100644 --- a/python/cudf/cudf/core/udf/groupby_typing.py +++ b/python/cudf/cudf/core/udf/groupby_typing.py @@ -33,7 +33,7 @@ numpy_support.as_dtype(dt) for dt in SUPPORTED_GROUPBY_NUMBA_TYPES ] -_UDF_DOC_URL = "https://docs.rapids.ai/api/cudf/stable/cudf/guide-to-udfs/" +_UDF_DOC_URL = "https://docs.nvidia.com/cudf/latest/cudf/guide-to-udfs/" class Group: diff --git a/python/cudf/cudf/utils/ioutils.py b/python/cudf/cudf/utils/ioutils.py index 4ae2471c5e9f..be0f2b6602ec 100644 --- a/python/cudf/cudf/utils/ioutils.py +++ b/python/cudf/cudf/utils/ioutils.py @@ -228,8 +228,8 @@ - Setting the cudf option `io.parquet.low_memory=True` will result in the chunked low memory parquet reader being used. This can make it easier to read large - parquet datasets on systems with limited GPU memory. See all `available options - `_. + parquet datasets on systems with limited GPU memory. See all :ref:`available options + `. Examples -------- @@ -809,8 +809,8 @@ - Setting the cudf option `io.json.low_memory=True` will result in the chunked low memory json reader being used. This can make it easier to read large - json datasets on systems with limited GPU memory. See all `available options - `_. + json datasets on systems with limited GPU memory. See all :ref:`available options + `. See Also -------- diff --git a/python/cudf/pyproject.toml b/python/cudf/pyproject.toml index b6eafde33d73..f58821f87e15 100644 --- a/python/cudf/pyproject.toml +++ b/python/cudf/pyproject.toml @@ -85,7 +85,7 @@ cudf-pandas-tests = [ [project.urls] Homepage = "https://github.com/NVIDIA/cudf" -Documentation = "https://docs.rapids.ai/api/cudf/stable/" +Documentation = "https://docs.nvidia.com/cudf/" [tool.pydistcheck] select = [ diff --git a/python/cudf_kafka/pyproject.toml b/python/cudf_kafka/pyproject.toml index 3e5ff58eb98f..ea0ae81635a2 100644 --- a/python/cudf_kafka/pyproject.toml +++ b/python/cudf_kafka/pyproject.toml @@ -31,7 +31,7 @@ test = [ [project.urls] Homepage = "https://github.com/NVIDIA/cudf" -Documentation = "https://docs.rapids.ai/api/cudf/stable/" +Documentation = "https://docs.nvidia.com/cudf/" [tool.ruff] extend = "../../pyproject.toml" diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 8dc366275852..4941391032ac 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -2756,7 +2756,12 @@ class Join(IR): """A join of two dataframes.""" __slots__ = ("left_on", "options", "right_on") - _non_child = ("schema", "left_on", "right_on", "options") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "left_on", + "right_on", + "options", + ) _n_non_child_args = 3 left_on: tuple[expr.NamedExpr, ...] """List of expressions used as keys in the left frame.""" diff --git a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py index 382d0027f649..f15f5f01a0b9 100644 --- a/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py +++ b/python/cudf_polars/cudf_polars/dsl/utils/column_domain.py @@ -21,6 +21,7 @@ Slice, Sort, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Mapping @@ -141,3 +142,11 @@ def _( return { name: ColumnBinding(0, name) for name in node.schema if name in child.schema } + + +@column_domain_bindings.register(PushdownFilterHint) +def _(node: PushdownFilterHint) -> Mapping[str, ColumnBinding]: + target = node.children[0] + return { + name: ColumnBinding(0, name) for name in node.schema if name in target.schema + } diff --git a/python/cudf_polars/cudf_polars/engine/core.py b/python/cudf_polars/cudf_polars/engine/core.py index 4b96f939396b..ba3937982119 100644 --- a/python/cudf_polars/cudf_polars/engine/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -211,6 +211,8 @@ def resolve_rapidsmpf_options(rapidsmpf_options: Options | None) -> Options: - ``num_streaming_threads=4``: moderate worker count for the rapidsmpf streaming runtime, shared across frontends. + - ``pinned_memory=true``, ``pinned_initial_pool_size=0``: pinned host + memory enabled by default. Parameters ---------- @@ -226,7 +228,13 @@ def resolve_rapidsmpf_options(rapidsmpf_options: Options | None) -> Options: if rapidsmpf_options is None: rapidsmpf_options = Options(get_environment_variables()) - rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) + rapidsmpf_options.insert_if_absent( + { + "num_streaming_threads": "4", + "pinned_memory": "true", + "pinned_initial_pool_size": "0", + } + ) return rapidsmpf_options @@ -649,7 +657,7 @@ def execute_ir_on_rank( hint = ( f"Try lowering `target_partition_size` (current {target_partition_size}) " f"and/or RAPIDSMPF_SPILL_DEVICE_LIMIT (default '80%') to reduce peak memory." - f"\nSee https://docs.rapids.ai/api/cudf/stable/cudf_polars/memory_errors/ " + f"\nSee https://docs.nvidia.com/cudf/latest/cudf_polars/memory_errors/ " f"for troubleshooting guidance." f"\nOriginal error:\n{mem_error}" ) diff --git a/python/cudf_polars/cudf_polars/engine/options.py b/python/cudf_polars/cudf_polars/engine/options.py index e1272475f941..cae2a26202ad 100644 --- a/python/cudf_polars/cudf_polars/engine/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -46,6 +46,7 @@ def _opt( category: str, env_var: str | None = None, coerce: Callable[[str], Any] = str, + default: Any = UNSPECIFIED, ) -> Any: """ Factory for ``StreamingOptions`` fields with category and env-var metadata. @@ -60,10 +61,14 @@ def _opt( :class:`StreamingOptions` is instantiated without an explicit value for this field, the factory reads the environment variable (if set) on the constructing process. ``None`` means no environment variable; the field defaults to - :data:`UNSPECIFIED`. + *default*. coerce Callable used to convert the raw env-var string to the field's type. Defaults to ``str`` (no conversion). + default + Value used when neither an explicit value nor the environment variable + is set. Defaults to :data:`UNSPECIFIED`, which defers to rapidsmpf's + built-in default. """ def _default() -> Any: @@ -71,7 +76,7 @@ def _default() -> Any: raw = os.environ.get(env_var) if raw is not None: return coerce(raw) - return UNSPECIFIED + return default return dataclasses.field( default_factory=_default, @@ -161,7 +166,7 @@ class StreamingOptions: pinned_memory Enable pinned host memory. Env: ``RAPIDSMPF_PINNED_MEMORY``. - Default: ``False``. + Default: ``True``. Category: rapidsmpf. pinned_initial_pool_size Initial pinned memory pool size (bytes). @@ -248,7 +253,7 @@ class StreamingOptions: disables the rewrite. Env: ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN`` and ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__*``. - Default: enabled. + Default: disabled. Category: executor. sink_to_directory Whether multi-partition sink operations should write to a directory @@ -433,8 +438,8 @@ def to_dict(self) -> dict[str, Any]: Examples -------- - >>> StreamingOptions(fallback_mode="silent").to_dict() - {'fallback_mode': 'silent'} + >>> StreamingOptions(fallback_mode="silent").to_dict() # doctest: +ELLIPSIS + {..., 'fallback_mode': 'silent'} >>> StreamingOptions.from_dict( ... StreamingOptions(fallback_mode="silent").to_dict() ... ) # doctest: +ELLIPSIS @@ -639,7 +644,7 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None: action=argparse.BooleanOptionalAction, help=textwrap.dedent("""\ Enable pinned host memory if available on the system. - Env: RAPIDSMPF_PINNED_MEMORY. Built-in default: false."""), + Env: RAPIDSMPF_PINNED_MEMORY. Default: true."""), ) g.add_argument( "--pinned-initial-pool-size", @@ -648,7 +653,7 @@ def _add_cli_args(parser: argparse.ArgumentParser) -> None: type=int, help=textwrap.dedent("""\ Starting allocation for the pinned memory pool in bytes. - Env: RAPIDSMPF_PINNED_INITIAL_POOL_SIZE. Built-in default: 0."""), + Env: RAPIDSMPF_PINNED_INITIAL_POOL_SIZE. Default: 0."""), ) g.add_argument( "--pinned-max-pool-size", diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py index c78c8084ae54..c3733c08eec2 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py @@ -16,6 +16,7 @@ import cudf_polars.streaming.actor_graph.io import cudf_polars.streaming.actor_graph.join import cudf_polars.streaming.actor_graph.over +import cudf_polars.streaming.actor_graph.prefilter_actor import cudf_polars.streaming.actor_graph.repartition import cudf_polars.streaming.actor_graph.union # noqa: F401 diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py index bae39d58cc08..7cb34f02d107 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py @@ -12,6 +12,7 @@ from cudf_polars.dsl.ir import Distinct, GroupBy, Sort from cudf_polars.dsl.traversal import traversal +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.io import StreamingSink from cudf_polars.streaming.join import Join from cudf_polars.streaming.over import Over @@ -107,6 +108,7 @@ def __init__( GroupBy, Distinct, Over, + PushdownFilterHint, ) self.collective_nodes: list[IR] = [ diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index fbec424d3104..82688495f859 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -22,6 +22,7 @@ generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext @@ -176,11 +177,12 @@ def _mark_children_unbounded(node: IR) -> None: for node in traversal([ir]): if node in unbounded: _mark_children_unbounded(node) - elif isinstance(node, (Union, Join, Over)): + elif isinstance(node, (Union, Join, Over, PushdownFilterHint)): # Union processes children sequentially; Join may broadcast one # side; Over buffers (or samples-then-replays) its input before - # producing output. In every case the input source needs - # unbounded fanout so other consumers don't block it. + # producing output; PushdownFilterHint similarly might buffer + # then replay. In every case the input source needs unbounded + # fanout so other consumers don't block it. _mark_children_unbounded(node) elif len(node.children) > 1: # Check if this node is doing any broadcasting. diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index d0c10e750b81..951ba3f4c0b2 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -7,6 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Literal, TypeAlias, assert_never +from cudf_streaming import CardinalityEstimator from cudf_streaming.channel_metadata import ( ChannelMetadata, HashScheme, @@ -26,7 +27,7 @@ ) from cudf_polars.containers import DataFrame -from cudf_polars.dsl.ir import IR, Join +from cudf_polars.dsl.ir import IR, Join, Projection from cudf_polars.dsl.utils.naming import names_to_indices from cudf_polars.streaming.actor_graph.collectives.allgather import ( AllGatherManager, @@ -42,18 +43,23 @@ from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) +from cudf_polars.streaming.actor_graph.join_planning import JoinPlanningState from cudf_polars.streaming.actor_graph.nodes import default_node_multi -from cudf_polars.streaming.actor_graph.tracing import send_chunk +from cudf_polars.streaming.actor_graph.prefilter import ( + JoinPrefilterExecution, + add_bloom_prefilter, + choose_prefilter, +) +from cudf_polars.streaming.actor_graph.tracing import LOG_TRACES, send_chunk from cudf_polars.streaming.actor_graph.utils import ( CUDF_ROW_LIMIT, MAX_ROWS_PER_PARTITION, ChannelManager, + ChunkSampler, ChunkStore, NormalizedPartitioning, TableSizeStats, - _sample_chunks, _update_ordering_indices, - allgather_reduce, chunk_to_frame, clear_local_ordering, empty_table_chunk, @@ -63,9 +69,15 @@ process_children, recv_metadata, replay_buffered_channel, + sample_inputs, send_metadata, shutdown_on_error, ) +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.utils import _concat @@ -80,8 +92,18 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR, IRExecutionContext from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.join_planning import JoinInput + from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterDecision, + PrefilterExecution, + ) from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.filter_hint import ( + JoinSide, + Prefilter, + PushdownFilterHint, + ) from cudf_polars.utils.config import StreamingExecutor @@ -147,6 +169,53 @@ class OrderedJoinStrategy: ] +@dataclass(frozen=True, slots=True) +class JoinCollectiveIds: + """Named collective-ID slots reserved for a dynamic join.""" + + size_estimate: int + left_redistribution: int + right_redistribution: int + + @classmethod + def from_reserved(cls, collective_ids: list[int]) -> JoinCollectiveIds: + """Construct the named slots from IDs reserved for a dynamic join.""" + if len(collective_ids) < 3: + raise ValueError( + "Dynamic join requires 3 reserved collective IDs " + "(allgather + left shuffle + right shuffle); got " + f"{len(collective_ids)} for this Join. " + "Ensure ReserveOpIDs is run with dynamic_planning enabled." + ) + return cls(*collective_ids[:3]) + + @property + def cardinality_tags(self) -> tuple[int, int]: + """Tags available for concurrent prefilter cardinality estimates.""" + return (self.size_estimate, self.left_redistribution) + + @property + def broadcast(self) -> int: + """ID used by a broadcast join after size estimation completes.""" + return self.left_redistribution + + def shuffle(self, side: JoinSide) -> int: + """Return the collective ID for one shuffle input.""" + if side == "left": + return self.left_redistribution + return self.right_redistribution + + def prefilter(self, strategy: JoinStrategy, target_side: JoinSide) -> int: + """Return the subsequent join collective reused by a prefilter.""" + if isinstance(strategy, BroadcastJoinStrategy): + if target_side != strategy.side: + raise ValueError( + "Only the broadcast input can have an active prefilter" + ) + return self.broadcast + return self.shuffle(target_side) + + @define_actor() async def broadcast_join_actor( context: Context, @@ -194,7 +263,7 @@ async def broadcast_join_actor( trace_ir=ir, ir_context=ir_context, ) as tracer: - await _broadcast_join( + await broadcast_join( context, comm, ir, @@ -203,7 +272,7 @@ async def broadcast_join_actor( ch_left, ch_right, BroadcastJoinStrategy(side=broadcast_side), - [collective_id], + collective_id, target_partition_size, tracer=tracer, ) @@ -298,7 +367,7 @@ async def _broadcast_join_large_chunk( broadcast_side: Literal["left", "right"], *, tracer: ActorTracer | None, -) -> None: +) -> int: """Join one large-side chunk with the small DataFrame(s) and send the result.""" large_df = chunk_to_frame(large_chunk, large_child) large_chunk_size = large_chunk.data_alloc_size() @@ -329,11 +398,13 @@ async def _broadcast_join_large_chunk( output_chunk = TableChunk.from_pylibcudf_table( df.table, df.stream, exclusive_view=True, br=context.br() ) + output_rows = output_chunk.shape[0] await send_chunk(context, ch_out, output_chunk, seq_num, tracer=tracer) del df, large_df + return output_rows -async def _broadcast_join( +async def broadcast_join( context: Context, comm: Communicator, ir: Join, @@ -342,26 +413,26 @@ async def _broadcast_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: BroadcastJoinStrategy, - collective_ids: list[int], - target_partition_size: int, + collective_id: int, + target_partition_size: int | None, *, tracer: ActorTracer | None, + trace_stats: dict[str, Any] | None = None, ) -> None: """ Execute a broadcast join after initial sampling. The small side is gathered (if not already duplicated) and concatenated into a single DataFrame, then joined with each chunk from the large side. - Pops one collective ID from collective_ids for allgather when needed. + Uses ``collective_id`` for the allgather when needed. """ left_metadata, right_metadata = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), ) - collective_id = collective_ids.pop(0) if collective_ids else 0 broadcast_side = strategy.side - left, right = ir.children + left, right = ir.children[:2] if tracer is not None: tracer.decision = f"broadcast_{broadcast_side}" @@ -407,8 +478,6 @@ async def _broadcast_join( partitioning=partitioning, duplicated=output_duplicated, ) - await send_metadata(ch_out, context, metadata_out) - small_dfs, small_size = await _collect_small_side_for_broadcast( context, comm, @@ -420,6 +489,13 @@ async def _broadcast_join( concat_size_limit=(target_partition_size if ir.options[0] == "Inner" else None), ) + # Publish output metadata only once the broadcast-side collective has + # completed. Besides making the data channel ready when advertised, this + # permits a consumer to reuse the collective ID after receiving metadata. + await send_metadata(ch_out, context, metadata_out) + + input_rows = 0 + output_rows = 0 while (msg := await large_ch.recv(context)) is not None: # Unknown: the large chunk is freed but the join output replaces # it, and its size depends on selectivity we cannot estimate here. @@ -429,7 +505,8 @@ async def _broadcast_join( reserve_extra=0, net_memory_delta=missing_net_memory_delta, ) - await _broadcast_join_large_chunk( + input_rows += large_chunk.shape[0] + output_rows += await _broadcast_join_large_chunk( context, ir, ir_context, @@ -444,9 +521,157 @@ async def _broadcast_join( tracer=tracer, ) + if trace_stats is not None: + trace_stats["input_rows"] = input_rows + trace_stats["output_rows"] = output_rows await ch_out.drain(context) +def add_prefilter( + execution: PrefilterExecution, + comm: Communicator, + *, + spec: Prefilter | PushdownFilterHint, + decision: PrefilterDecision, + target: IR, + domain: IR, + ch_target: Channel[TableChunk], + ch_domain_keys: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + ir_context: IRExecutionContext, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the actors and channels that apply one selected prefilter.""" + context = execution.context + if decision.method == "bloom": + if decision.bloom_bytes is None: + raise ValueError("Bloom prefilter decision has no filter size") + add_bloom_prefilter( + context, + comm, + decision.bloom_bytes, + execution, + names_to_indices(spec.target_on, target.schema), + ch_domain_keys, + ch_target, + ch_filtered, + collective_id, + trace_stats, + ) + elif decision.method == "broadcast_semi_join": + domain_schema = {key.name: key.value.dtype for key in spec.domain_on} + if len(domain_schema) != len(spec.domain_on): + raise ValueError("Broadcast semi-join keys must have unique names") + semi_join = Join( + target.schema, + spec.target_on, + spec.domain_on, + ("Semi", spec.nulls_equal, None, "", False, "none"), + target, + Projection(domain_schema, domain), + ) + execution.add_task( + broadcast_join( + context, + comm, + semi_join, + ir_context, + ch_filtered, + ch_target, + ch_domain_keys, + BroadcastJoinStrategy(side="right"), + collective_id, + target_partition_size=None, + tracer=None, + trace_stats=trace_stats, + ) + ) + else: + raise ValueError(f"Cannot apply prefilter method {decision.method!r}") + + +def make_prefilter_execution( + context: Context, + comm: Communicator, + ir: Join, + ir_context: IRExecutionContext, + strategy: JoinStrategy, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + join_state: JoinPlanningState, + collective_ids: JoinCollectiveIds, +) -> JoinPrefilterExecution: + """Create the actors and channels that realize selected prefilters.""" + execution = JoinPrefilterExecution(context, ch_left, ch_right) + + # Prepare every required domain before connecting target-side filters. This + # is important for opposing direct filters: each filter must consume the + # replay produced while the same input's keys are copied for the other one. + for candidate in join_state.candidates: + decision = candidate.decision + if decision is None: + raise ValueError("Join prefilter has no runtime decision") + spec = candidate.spec + if decision.method == "skip": + continue + + if isinstance(spec.domain, JoinInputDomain): + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + candidate.key_channel = execution.buffer_domain(spec.domain.side, indices) + else: + sample = candidate.domain.sample + if sample is None: + raise ValueError("Active external prefilter has no domain sample") + indices = names_to_indices(spec.domain_on, candidate.domain.node.schema) + if indices != tuple(range(len(candidate.domain.node.schema))): + raise ValueError("External prefilter domains must contain only keys") + candidate.key_channel = context.create_channel() + execution.add_channel(candidate.key_channel) + execution.add_task( + replay_buffered_channel( + context, + candidate.key_channel, + candidate.domain.channel, + sample.chunks, + candidate.domain.metadata, + trace_ir=ir, + ) + ) + + for candidate in join_state.candidates: + decision = candidate.decision + assert decision is not None + if decision.method == "skip": + continue + spec = candidate.spec + ch_domain_keys = candidate.key_channel + assert ch_domain_keys is not None + target_side = spec.target_side + target = candidate.target.node + ch_target = execution.join_inputs[target_side] + ch_filtered: Channel[TableChunk] = context.create_channel() + trace_stats = candidate.trace + + add_prefilter( + execution, + comm, + spec=spec, + decision=decision, + target=target, + domain=candidate.domain.node, + ch_target=ch_target, + ch_domain_keys=ch_domain_keys, + ch_filtered=ch_filtered, + collective_id=collective_ids.prefilter(strategy, target_side), + ir_context=ir_context, + trace_stats=trace_stats, + ) + execution.replace_join_input(target_side, ch_filtered) + + return execution + + def _get_key_indices( ir: Join, n_partitioned_keys: int | None, @@ -457,7 +682,7 @@ def _get_key_indices( tuple[NamedExpr, ...], tuple[NamedExpr, ...], ]: - left, right = ir.children + left, right = ir.children[:2] n_keys = n_partitioned_keys if n_partitioned_keys is not None else len(ir.left_on) left_keys = ir.left_on[:n_keys] right_keys = ir.right_on[:n_keys] @@ -586,7 +811,7 @@ async def _join_chunks( recv_metadata(ch_right, context), ) - left, right = ir.children + left, right = ir.children[:2] while True: left_msg, right_msg = await gather_in_task_group( ch_left.recv(context), ch_right.recv(context) @@ -687,7 +912,8 @@ async def _shuffle_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: ShuffleJoinStrategy, - collective_ids: list[int], + left_collective_id: int, + right_collective_id: int, *, tracer: ActorTracer | None, ) -> None: @@ -730,7 +956,7 @@ async def _shuffle_join( strategy.left_keys, ir.children[0].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + left_collective_id, ), _global_shuffle( context, @@ -741,7 +967,7 @@ async def _shuffle_join( strategy.right_keys, ir.children[1].schema, strategy.shuffle_modulus, - collective_ids.pop(0), + right_collective_id, ), _join_chunks( context, @@ -803,7 +1029,7 @@ async def _adjust_ordered_join_side( context, ch_out, ch_in, - (), + ChunkStore(context), output_metadata, trace_ir=schema_ir, ) @@ -832,7 +1058,7 @@ async def _ordered_join( ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], strategy: OrderedJoinStrategy, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, *, tracer: ActorTracer | None, ) -> None: @@ -892,8 +1118,8 @@ async def _ordered_join( ch_left, strategy.left_input_ordering, strategy.left_output_ordering, + collective_id=collective_ids.shuffle("left"), already_aligned=left_aligned, - collective_id=collective_ids.pop(0), ), _adjust_ordered_join_side( context, @@ -904,8 +1130,8 @@ async def _ordered_join( ch_right, strategy.right_input_ordering, strategy.right_output_ordering, + collective_id=collective_ids.shuffle("right"), already_aligned=right_aligned, - collective_id=collective_ids.pop(0), ), _join_chunks( context, @@ -962,49 +1188,6 @@ def _num_indices(partitioning: NormalizedPartitioning) -> int: ) -async def _aggregate_estimates( - context: Context, - comm: Communicator, - left_sample: TableSizeStats, - right_sample: TableSizeStats, - collective_ids: list[int], -) -> tuple[TableSizeStats, TableSizeStats]: - """Aggregate table-size and row estimates across ranks.""" - # AllGather size, row, and chunk count estimates across ranks - ( - left_total, - right_total, - left_total_rows, - right_total_rows, - left_total_chunks, - right_total_chunks, - ) = await allgather_reduce( - context, - comm, - collective_ids.pop(0), - left_sample.total_size, - right_sample.total_size, - left_sample.total_rows, - right_sample.total_rows, - left_sample.total_chunks, - right_sample.total_chunks, - ) - - new_left_sample = TableSizeStats( - chunks=left_sample.chunks, - total_size=left_total, - total_rows=left_total_rows, - total_chunks=left_total_chunks, - ) - new_right_sample = TableSizeStats( - chunks=right_sample.chunks, - total_size=right_total, - total_rows=right_total_rows, - total_chunks=right_total_chunks, - ) - return new_left_sample, new_right_sample - - def _choose_strategy_from_samples( comm: Communicator, ir: Join, @@ -1148,45 +1331,241 @@ def _modulus(partitioning: NormalizedPartitioning) -> int | None: return max(large, min_shuffle_modulus) -async def _choose_strategy( +def join_input_requires_redistribution( + strategy: JoinStrategy, + side: Literal["left", "right"], + partitioning: NormalizedPartitioning, + metadata: ChannelMetadata, +) -> bool: + """Return whether the join strategy redistributes an input side.""" + if isinstance(strategy, BroadcastJoinStrategy): + return side == strategy.side and not metadata.duplicated + if isinstance(strategy, OrderedJoinStrategy): + # Ordered inputs already have a viable join strategy without adaptive + # sampling. Keep that strategy and avoid introducing a sampled + # prefilter pipeline solely to optimize boundary alignment. + return False + + assert isinstance(strategy, ShuffleJoinStrategy) + indices = strategy.left_indices if side == "left" else strategy.right_indices + if not indices: + return True + desired = HashScheme(indices, strategy.shuffle_modulus) + return not ( + partitioning.inter_rank_scheme == desired + and partitioning.local_scheme == "inherit" + ) + + +def choose_prefilters( + join_state: JoinPlanningState, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> None: + """Choose strategies for prefilters with sufficient available statistics.""" + partitionings = { + "left": left_partitioning, + "right": right_partitioning, + } + for candidate in join_state.candidates: + if candidate.decision is not None: + continue + target = candidate.target.sample + if target is None: + raise ValueError("Join target has not been sampled") + target_side = candidate.spec.target_side + target_requires_redistribution = join_input_requires_redistribution( + strategy, + target_side, + partitionings[target_side], + candidate.target.metadata, + ) + if ( + isinstance(candidate.spec.domain, ExternalDomain) + and candidate.domain.sample is None + and target_requires_redistribution + ): + continue + candidate.decision = choose_prefilter( + candidate.spec, + target, + candidate.domain.sample, + target_requires_redistribution=target_requires_redistribution, + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +async def collect_samples( + context: Context, + comm: Communicator, + join_state: JoinPlanningState, + inputs: tuple[JoinInput, ...], + sample_chunk_count: int, + target_partition_size: int, + collective_id: int, +) -> None: + """Sample inputs and attach aggregate estimates to their planning state.""" + if not inputs: + return + sampling_inputs = [] + for input_ in inputs: + candidates = [ + candidate + for candidate in join_state.candidates + if candidate.domain is input_ + ] + if len(candidates) > 1: + raise ValueError("One join input cannot provide multiple prefilter domains") + sampling_inputs.append((input_, candidates[0] if candidates else None)) + samplers = [] + for input_, candidate in sampling_inputs: + if candidate is None: + cardinality_estimator = None + cardinality_columns: tuple[int, ...] = () + else: + cardinality_estimator = CardinalityEstimator( + context, + comm, + tag=candidate.cardinality_tag, + ) + cardinality_columns = names_to_indices( + candidate.spec.domain_on, + input_.node.schema, + ) + assert len(cardinality_columns) == len(candidate.spec.domain_on), ( + "Prefilter domain keys must be columns" + ) + samplers.append( + ChunkSampler( + context=context, + ch_in=input_.channel, + max_chunks=sample_chunk_count, + max_bytes=target_partition_size, + ch_in_chunk_count=input_.metadata.local_count, + cardinality_estimator=cardinality_estimator, + cardinality_columns=cardinality_columns, + ) + ) + samples = await sample_inputs( + context, + comm, + samplers, + collective_id, + ) + for (input_, _), sample in zip(sampling_inputs, samples, strict=True): + input_.sample = sample + + +async def release_skipped_external_domains( + context: Context, join_state: JoinPlanningState +) -> None: + """Release buffered data and stop external domains rejected by planning.""" + channels = [] + for candidate in join_state.candidates: + if not isinstance(candidate.spec.domain, ExternalDomain): + continue + if candidate.decision is None: + raise ValueError("Join prefilter has no runtime decision") + if candidate.decision.method != "skip": + continue + if candidate.domain.sample is not None: + candidate.domain.sample.chunks.clear() + channels.append(candidate.domain.channel) + if channels: + await gather_in_task_group(*(channel.shutdown(context) for channel in channels)) + + +async def resolve_prefilters( + context: Context, + comm: Communicator, + join_state: JoinPlanningState, + strategy: JoinStrategy, + left_partitioning: NormalizedPartitioning, + right_partitioning: NormalizedPartitioning, + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Resolve optional prefilters after selecting the join strategy.""" + config = executor.join_filter_pushdown + if config is None or not join_state.candidates: + return + + choose_prefilters( + join_state, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + assert executor.dynamic_planning is not None + await collect_samples( + context, + comm, + join_state, + tuple( + candidate.domain + for candidate in join_state.candidates + if isinstance(candidate.spec.domain, ExternalDomain) + and candidate.decision is None + ), + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_id, + ) + choose_prefilters( + join_state, + strategy, + left_partitioning, + right_partitioning, + executor.broadcast_limit, + config.bloom_filter_max_size, + ) + await release_skipped_external_domains(context, join_state) + + +async def choose_strategy( context: Context, comm: Communicator, ir: Join, - ch_left: Channel[TableChunk], - ch_right: Channel[TableChunk], - left_metadata: ChannelMetadata, - right_metadata: ChannelMetadata, + join_state: JoinPlanningState, executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, *, tracer: ActorTracer | None, -) -> tuple[TableSizeStats, TableSizeStats, JoinStrategy]: - """Sample both sides, aggregate estimates, and choose broadcast vs shuffle.""" +) -> JoinStrategy: + """Collect any required samples and choose broadcast vs shuffle.""" + left, right = ir.children[:2] + left_metadata = join_state.left.metadata + right_metadata = join_state.right.metadata nranks = comm.nranks left_partitioning = NormalizedPartitioning.from_keys( left_metadata.partitioning, nranks, - keys=names_to_indices(ir.left_on, ir.children[0].schema, concrete_prefix=True), + keys=names_to_indices(ir.left_on, left.schema, concrete_prefix=True), ) right_partitioning = NormalizedPartitioning.from_keys( right_metadata.partitioning, nranks, - keys=names_to_indices(ir.right_on, ir.children[1].schema, concrete_prefix=True), + keys=names_to_indices(ir.right_on, right.schema, concrete_prefix=True), ) - hash_chunkwise = isinstance( left_partitioning.inter_rank_scheme, HashScheme ) and isinstance(right_partitioning.inter_rank_scheme, HashScheme) - if hash_chunkwise and left_partitioning.is_aligned_with( + chunkwise = hash_chunkwise and left_partitioning.is_aligned_with( right_partitioning, context.br() - ): - # We can use a chunkwise join - chunkwise = True - left_sample = TableSizeStats( + ) + + if chunkwise: + join_state.left.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=left_metadata.local_count, ) - right_sample = TableSizeStats( + join_state.right.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=right_metadata.local_count, ) @@ -1204,45 +1583,41 @@ async def _choose_strategy( ): if tracer is not None: tracer.decision = "ordered" - left_sample = TableSizeStats( + join_state.left.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=left_metadata.local_count, ) - right_sample = TableSizeStats( + join_state.right.sample = TableSizeStats( chunks=ChunkStore(context), total_chunks=right_metadata.local_count, ) - return left_sample, right_sample, ordered_strategy + await resolve_prefilters( + context, + comm, + join_state, + ordered_strategy, + left_partitioning, + right_partitioning, + executor, + collective_ids.size_estimate, + ) + return ordered_strategy else: - # Need to shuffle or broadcast - Use sampled data to choose a strategy - chunkwise = False assert executor.dynamic_planning is not None - sample_chunk_count = executor.dynamic_planning.sample_chunk_count - target_partition_size = executor.target_partition_size - left_sample, right_sample = await gather_in_task_group( - _sample_chunks( - context, - ch_left, - sample_chunk_count, - target_partition_size, - left_metadata.local_count, - ), - _sample_chunks( - context, - ch_right, - sample_chunk_count, - target_partition_size, - right_metadata.local_count, - ), - ) - left_sample, right_sample = await _aggregate_estimates( + await collect_samples( context, comm, - left_sample, - right_sample, - collective_ids, + join_state, + (join_state.left, join_state.right), + executor.dynamic_planning.sample_chunk_count, + executor.target_partition_size, + collective_ids.size_estimate, ) + left_sample = join_state.left.sample + right_sample = join_state.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") strategy = _choose_strategy_from_samples( comm, ir, @@ -1256,8 +1631,17 @@ async def _choose_strategy( chunkwise=chunkwise, tracer=tracer, ) - - return left_sample, right_sample, strategy + await resolve_prefilters( + context, + comm, + join_state, + strategy, + left_partitioning, + right_partitioning, + executor, + collective_ids.size_estimate, + ) + return strategy @define_actor() @@ -1269,8 +1653,9 @@ async def join_actor( ch_out: Channel[TableChunk], ch_left: Channel[TableChunk], ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], executor: StreamingExecutor, - collective_ids: list[int], + collective_ids: JoinCollectiveIds, ) -> None: """ Dynamic Join actor that selects the best strategy at runtime. @@ -1295,6 +1680,8 @@ async def join_actor( Input channel for the left side. ch_right Input channel for the right side. + ch_prefilter_domains + Input channels providing the prefilter key domains. executor Streaming executor configuration. collective_ids @@ -1305,32 +1692,72 @@ async def join_actor( ch_out, ch_left, ch_right, + *ch_prefilter_domains, trace_ir=ir, ir_context=ir_context, ) as tracer: - left_metadata, right_metadata = await gather_in_task_group( + ( + left_metadata, + right_metadata, + *prefilter_domain_metadata, + ) = await gather_in_task_group( recv_metadata(ch_left, context), recv_metadata(ch_right, context), + *(recv_metadata(ch, context) for ch in ch_prefilter_domains), ) - left_sample, right_sample, strategy = await _choose_strategy( - context, - comm, + join_state = JoinPlanningState.create( ir, ch_left, ch_right, + ch_prefilter_domains, left_metadata, right_metadata, + tuple(prefilter_domain_metadata), + collective_ids.cardinality_tags, + ) + + strategy = await choose_strategy( + context, + comm, + ir, + join_state, executor, collective_ids, tracer=tracer, ) + prefilter_traces = [] + for candidate in join_state.candidates: + if candidate.decision is None: + raise ValueError("Join prefilter has no runtime decision") + trace = candidate.decision.trace(candidate.spec) + prefilter_traces.append(trace) + if LOG_TRACES: + candidate.trace = trace + if tracer is not None and prefilter_traces: + tracer.set_extra("join_prefilters", prefilter_traces) + left_sample = join_state.left.sample + right_sample = join_state.right.sample + if left_sample is None or right_sample is None: + raise ValueError("Join inputs have not been sampled") ch_left_replay = context.create_channel() ch_right_replay = context.create_channel() + prefilter_execution = make_prefilter_execution( + context, + comm, + ir, + ir_context, + strategy, + ch_left_replay, + ch_right_replay, + join_state, + collective_ids, + ) async with shutdown_on_error( context, ch_left_replay, ch_right_replay, + *prefilter_execution.channels, trace_ir=ir, ir_context=ir_context, ): @@ -1351,13 +1778,14 @@ async def join_actor( right_metadata, trace_ir=ir, ), + *prefilter_execution.tasks, ] - ch_left = ch_left_replay - ch_right = ch_right_replay + ch_left = prefilter_execution.left + ch_right = prefilter_execution.right if isinstance(strategy, BroadcastJoinStrategy): actor_tasks.append( - _broadcast_join( + broadcast_join( context, comm, ir, @@ -1366,7 +1794,7 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.broadcast, executor.target_partition_size, tracer=tracer, ) @@ -1397,7 +1825,8 @@ async def join_actor( ch_left, ch_right, strategy, - collective_ids, + collective_ids.shuffle("left"), + collective_ids.shuffle("right"), tracer=tracer, ) ) @@ -1412,7 +1841,7 @@ def _use_pwise_join( ir: Join, ) -> bool: """Whether to use a static-planning partition-wise join.""" - left, right = ir.children + left, right = ir.children[:2] output_count = partition_info[ir].count if ( output_count == 1 @@ -1438,18 +1867,24 @@ def _use_pwise_join( @generate_ir_sub_network.register(Join) +@generate_ir_sub_network.register(JoinWithPrefilter) def _( - ir: Join, rec: SubNetGenerator + ir: Join | JoinWithPrefilter, rec: SubNetGenerator ) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: # Join operation. - left, right = ir.children + left, right, *prefilter_domains = ir.children partition_info = rec.state["partition_info"] left_count = partition_info[left].count right_count = partition_info[right].count executor = rec.state["config_options"].executor pwise_join = _use_pwise_join(executor, partition_info, ir) - # Process children + if pwise_join and isinstance(ir, JoinWithPrefilter): + raise AssertionError( + "Partition-wise JoinWithPrefilter should have been simplified " + "during IR lowering" + ) + actors, channels = process_children(ir, rec) # Create output ChannelManager @@ -1479,16 +1914,13 @@ def _( and ir.options[0] in ("Inner", "Left", "Right", "Full", "Semi", "Anti") ): # Dynamic join - decide strategy at runtime - collective_ids = list(rec.state["collective_id_map"].get(ir, [])) - # Join uses up to 3 collective IDs: allgather, left shuffle, and - # right shuffle. - if len(collective_ids) < 3: - raise ValueError( - "Dynamic join requires 3 reserved collective IDs " - "(allgather + left shuffle + right shuffle); got " - f"{len(collective_ids)} for this Join. " - "Ensure ReserveOpIDs is run with dynamic_planning enabled." - ) + collective_ids = JoinCollectiveIds.from_reserved( + rec.state["collective_id_map"].get(ir, []) + ) + # Join uses up to 3 collective IDs. Cardinality allreduces complete + # before the size allgather and join collectives. Runtime prefilters + # reuse the collective ID of the target-side join redistribution, with + # their filtered output channel providing the ordering barrier. actors[ir] = [ join_actor( rec.state["context"], @@ -1498,6 +1930,10 @@ def _( channels[ir].reserve_input_slot(), channels[left].reserve_output_slot(), channels[right].reserve_output_slot(), + tuple( + channels[domain].reserve_output_slot() + for domain in prefilter_domains + ), executor, collective_ids, ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py new file mode 100644 index 000000000000..1b0972bb762e --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join_planning.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Actor-local planning state for dynamic joins and optional prefilters.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, +) + +if TYPE_CHECKING: + from typing import Any, Self + + from cudf_streaming.channel_metadata import ChannelMetadata + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.streaming.core.channel import Channel + + from cudf_polars.dsl.ir import IR, Join + from cudf_polars.streaming.actor_graph.prefilter import PrefilterDecision + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import Prefilter + + +@dataclass(slots=True) +class JoinInput: + """Concrete runtime resources for one input to a dynamic join.""" + + node: IR + channel: Channel[TableChunk] + metadata: ChannelMetadata + sample: TableSizeStats | None = None + + +@dataclass(slots=True) +class PrefilterCandidate: + """An optional prefilter and the runtime inputs needed to evaluate it.""" + + spec: Prefilter + target: JoinInput + domain: JoinInput + cardinality_tag: int + decision: PrefilterDecision | None = None + key_channel: Channel[TableChunk] | None = None + trace: dict[str, Any] | None = None + + +@dataclass(frozen=True, slots=True) +class JoinPlanningState: + """Actor-local input and prefilter state for planning a dynamic join.""" + + left: JoinInput + right: JoinInput + candidates: tuple[PrefilterCandidate, ...] = () + + @classmethod + def create( + cls, + ir: Join, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ch_prefilter_domains: tuple[Channel[TableChunk], ...], + left_metadata: ChannelMetadata, + right_metadata: ChannelMetadata, + prefilter_domain_metadata: tuple[ChannelMetadata, ...], + cardinality_tags: tuple[int, ...], + ) -> Self: + """Create actor-local planning state from a join and its runtime inputs.""" + left = JoinInput(ir.children[0], ch_left, left_metadata) + right = JoinInput(ir.children[1], ch_right, right_metadata) + if not isinstance(ir, JoinWithPrefilter): + if ch_prefilter_domains or prefilter_domain_metadata: + raise ValueError("A plain Join cannot have prefilter domain inputs") + return cls(left, right) + + external_inputs = tuple( + JoinInput(node, channel, metadata) + for node, channel, metadata in zip( + ir.children[2:], + ch_prefilter_domains, + prefilter_domain_metadata, + strict=True, + ) + ) + external_prefilter_count = sum( + isinstance(prefilter.domain, ExternalDomain) for prefilter in ir.prefilters + ) + if external_prefilter_count != len(external_inputs): + raise ValueError("Join prefilters and external domain inputs must align") + if len(cardinality_tags) < len(ir.prefilters): + raise ValueError("Each join prefilter requires a cardinality collective ID") + + sides = {"left": left, "right": right} + external_inputs_iter = iter(external_inputs) + cardinality_tags_iter = iter(cardinality_tags) + candidates = [] + for spec in ir.prefilters: + target = sides[spec.target_side] + if isinstance(spec.domain, JoinInputDomain): + domain = sides[spec.domain.side] + else: + domain = next(external_inputs_iter) + candidates.append( + PrefilterCandidate( + spec, + target, + domain, + cardinality_tag=next(cardinality_tags_iter), + ) + ) + return cls(left, right, tuple(candidates)) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py new file mode 100644 index 000000000000..0a518e12f3ed --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter.py @@ -0,0 +1,442 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Runtime planning helpers for optional prefilters.""" + +from __future__ import annotations + +import math +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Any, Literal + +import pylibcudf as plc +from cudf_streaming import BloomFilter +from cudf_streaming.channel_metadata import ChannelMetadata +from cudf_streaming.table_chunk import TableChunk +from pylibcudf.hashing import LIBCUDF_DEFAULT_HASH_SEED +from rapidsmpf.memory.memory_reservation import opaque_memory_usage +from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory +from rapidsmpf.streaming.core.message import Message + +from cudf_polars.streaming.actor_graph.utils import ( + ChunkStore, + recv_metadata, + send_metadata, + shutdown_channels_on_error, +) +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, +) + +if TYPE_CHECKING: + from collections.abc import Coroutine, Iterable, Sequence + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.containers import DataType + from cudf_polars.dsl.expr import NamedExpr + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.streaming.filter_hint import JoinSide, Prefilter + + +def estimate_bytes(dtypes: Sequence[DataType], row_count: int) -> int | None: + """ + Estimate the byte count of a table containing the given datatypes. + + Parameters + ---------- + dtypes + Types of columns in the table. + row_count + Estimated total number of rows. + + Returns + ------- + Estimated table size in bytes, or ``None`` if any dtype is not fixed width. + """ + if not all(plc.traits.is_fixed_width(dtype.plc_type) for dtype in dtypes): + return None + + return int( + # Just assume everything has a validity mask + row_count * sum(plc.types.size_of(dtype.plc_type) + 1 / 8 for dtype in dtypes) + ) + + +@dataclass(frozen=True, slots=True) +class PrefilterDecision: + """Runtime decision for one optional prefilter.""" + + method: Literal["skip", "bloom", "broadcast_semi_join"] + reason: str + target_bytes: int + domain_rows: int | None + estimated_cardinality: int | None = None + bloom_bytes: int | None = None + exact_bytes: int | None = None + + def trace(self, prefilter: Prefilter) -> dict[str, str | int | None]: + """Return serializable actor-trace information.""" + result = asdict(self) + result["target_side"] = prefilter.target_side + if isinstance(prefilter.domain, JoinInputDomain): + result["domain_side"] = prefilter.domain.side + else: + assert isinstance(prefilter.domain, ExternalDomain) + result["domain"] = "external" + return result + + +async def project_key_chunk( + context: Context, chunk: TableChunk, indices: Iterable[int] +) -> TableChunk: + """Copy selected columns into an owning key chunk.""" + columns = tuple(chunk.table_view().columns()[index] for index in indices) + bytes = sum(column.device_buffer_size() for column in columns) + with opaque_memory_usage( + await reserve_memory(context, size=bytes, net_memory_delta=0) + ): + table = plc.Table(columns).copy(stream=chunk.stream, mr=context.br().device_mr) + return TableChunk.from_pylibcudf_table( + table, + chunk.stream, + exclusive_view=True, + br=context.br(), + ) + + +async def buffer_and_project_keys( + context: Context, + ch_in: Channel[TableChunk], + ch_keys: Channel[TableChunk], + ch_replay: Channel[TableChunk], + indices: Iterable[int], +) -> None: + """ + Project owning key chunks while spill-buffering an input for replay. + + The key channel is produced in full before replay begins. Its consumer must + therefore run concurrently with this coroutine. + """ + chunks = ChunkStore(context) + try: + async with shutdown_channels_on_error(context, ch_in, ch_keys, ch_replay): + metadata = await recv_metadata(ch_in, context) + key_metadata = ChannelMetadata( + local_count=metadata.local_count, + partitioning=None, + duplicated=metadata.duplicated, + ) + await send_metadata(ch_replay, context, metadata) + await send_metadata(ch_keys, context, key_metadata) + indices = tuple(indices) + while (msg := await ch_in.recv(context)) is not None: + sequence_number = msg.sequence_number + chunk = await TableChunk.from_message( + msg, br=context.br() + ).make_available_or_wait(context, net_memory_delta=0) + key_chunk = await project_key_chunk(context, chunk, indices) + chunks.insert(Message(sequence_number, chunk)) + await ch_keys.send(context, Message(sequence_number, key_chunk)) + + await ch_keys.drain(context) + for msg in chunks: + await ch_replay.send(context, msg) + await ch_replay.drain(context) + finally: + chunks.clear() + + +async def count_rows_passthrough( + context: Context, + ch_in: Channel[TableChunk], + ch_out: Channel[TableChunk], + trace_stats: dict[str, Any], + row_count_key: str, +) -> None: + """Forward a table-chunk channel while recording its row count.""" + async with shutdown_channels_on_error(context, ch_in, ch_out): + metadata = await recv_metadata(ch_in, context) + await send_metadata(ch_out, context, metadata) + row_count = 0 + while (msg := await ch_in.recv(context)) is not None: + chunk = TableChunk.from_message(msg, br=context.br()) + row_count += chunk.shape[0] + await ch_out.send(context, Message(msg.sequence_number, chunk)) + trace_stats[row_count_key] = row_count + await ch_out.drain(context) + + +class PrefilterExecution: + """Channels and actor tasks used to apply one or more prefilters.""" + + def __init__(self, context: Context) -> None: + self.context = context + self.tasks: list[Coroutine[Any, Any, None]] = [] + self.channels: list[Channel[Any]] = [] + + def add_task(self, task: Coroutine[Any, Any, None]) -> None: + """Add an actor task to the prefilter execution.""" + self.tasks.append(task) + + def add_channel(self, channel: Channel[Any]) -> None: + """Register an auxiliary channel for shutdown on failure.""" + self.channels.append(channel) + + +class JoinPrefilterExecution(PrefilterExecution): + """Channels and actor tasks used to apply prefilters before a join.""" + + def __init__( + self, + context: Context, + ch_left: Channel[TableChunk], + ch_right: Channel[TableChunk], + ) -> None: + super().__init__(context) + self.source_inputs = {"left": ch_left, "right": ch_right} + self.join_inputs = dict(self.source_inputs) + self.buffered_domains: set[JoinSide] = set() + + def buffer_domain( + self, + side: JoinSide, + indices: Iterable[int], + ) -> Channel[TableChunk]: + """Buffer one original input and return its owning key channel.""" + if side in self.buffered_domains: + raise ValueError(f"Join input {side!r} is already a prefilter domain") + + ch_keys: Channel[TableChunk] = self.context.create_channel() + ch_replay: Channel[TableChunk] = self.context.create_channel() + self.tasks.append( + buffer_and_project_keys( + self.context, + self.source_inputs[side], + ch_keys, + ch_replay, + indices, + ) + ) + self.channels.extend((ch_keys, ch_replay)) + self.join_inputs[side] = ch_replay + self.buffered_domains.add(side) + return ch_keys + + def replace_join_input( + self, + side: JoinSide, + channel: Channel[TableChunk], + ) -> None: + """Replace one join-facing input with a prefilter output channel.""" + self.join_inputs[side] = channel + self.channels.append(channel) + + @property + def left(self) -> Channel[TableChunk]: + """Current left join input.""" + return self.join_inputs["left"] + + @property + def right(self) -> Channel[TableChunk]: + """Current right join input.""" + return self.join_inputs["right"] + + +def add_bloom_prefilter( + context: Context, + comm: Communicator, + bloom_bytes: int, + execution: PrefilterExecution, + target_indices: Iterable[int], + ch_domain_keys: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_filtered: Channel[TableChunk], + collective_id: int, + trace_stats: dict[str, Any] | None, +) -> None: + """Add the channels and actors for an approximate Bloom prefilter.""" + bloom = BloomFilter( + context, + comm, + LIBCUDF_DEFAULT_HASH_SEED, + bloom_bytes, + ) + ch_filter = context.create_channel() + execution.add_channel(ch_filter) + execution.add_task( + bloom.build( + context, + ch_domain_keys, + ch_filter, + collective_id, + ) + ) + ch_apply_input = ch_target + ch_apply_output = ch_filtered + if trace_stats is not None: + ch_counted_input: Channel[TableChunk] = context.create_channel() + ch_raw_output: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_counted_input) + execution.add_channel(ch_raw_output) + execution.add_task( + count_rows_passthrough( + context, + ch_target, + ch_counted_input, + trace_stats, + "input_rows", + ) + ) + execution.add_task( + count_rows_passthrough( + context, + ch_raw_output, + ch_filtered, + trace_stats, + "output_rows", + ) + ) + ch_apply_input = ch_counted_input + ch_apply_output = ch_raw_output + execution.add_task( + bloom.apply( + context, + ch_filter, + ch_apply_input, + ch_apply_output, + target_indices, + ) + ) + + +def estimate_bloom_filter_bytes( + cardinality: int, + desired_false_positive_rate: float = 0.1, +) -> int: + """Estimate Bloom-filter bytes for the block-split policy.""" + if cardinality < 0: + raise ValueError("cardinality must be non-negative") + if not 0 < desired_false_positive_rate < 1: + raise ValueError("false_positive_rate must be between zero and one") + if cardinality == 0: + return 0 + # TODO: cuco could offer this as a static utility on the policy + # Then we wouldn't have to hardcode these magic numbers. + bits = ( + -8 # number of fingerprint bits + * cardinality + / math.log(1 - desired_false_positive_rate ** (1 / 8)) + ) + return math.ceil(bits / 8) + + +def choose_prefilter( + prefilter: Prefilter, + target: TableSizeStats, + domain: TableSizeStats | None, + *, + target_requires_redistribution: bool, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose whether one join prefilter is eligible to be applied.""" + domain_rows = None if domain is None else domain.total_rows + if not target_requires_redistribution: + return PrefilterDecision( + "skip", + "target_not_redistributed", + target.total_size, + domain_rows, + ) + if domain is None: + raise ValueError("A redistributed target requires domain statistics") + if ( + isinstance(prefilter.domain, JoinInputDomain) + and prefilter.target_side == prefilter.domain.side + ): + return PrefilterDecision( + "skip", + "same_input", + target.total_size, + domain_rows, + ) + + return choose_prefilter_method( + prefilter.domain_on, + target, + domain, + broadcast_limit=broadcast_limit, + bloom_filter_max_size=bloom_filter_max_size, + ) + + +def choose_prefilter_method( + domain_on: Sequence[NamedExpr], + target: TableSizeStats, + domain: TableSizeStats, + *, + broadcast_limit: int, + bloom_filter_max_size: int, +) -> PrefilterDecision: + """Choose the implementation for an eligible prefilter.""" + distinct_count = domain.distinct_count() + if distinct_count is None: + return PrefilterDecision( + "skip", + "missing_cardinality", + target.total_size, + domain.total_rows, + ) + if distinct_count == 0: + return PrefilterDecision( + "skip", + "zero_cardinality", + target.total_size, + domain.total_rows, + estimated_cardinality=0, + bloom_bytes=0, + exact_bytes=0, + ) + + bloom_bytes = max( + 32, + BloomFilter.aligned_size(estimate_bloom_filter_bytes(distinct_count)), + ) + exact_bytes = estimate_bytes( + tuple(key.value.dtype for key in domain_on), + domain.total_rows, + ) + if bloom_bytes <= min(bloom_filter_max_size, target.total_size): + return PrefilterDecision( + "bloom", + "bloom_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + if exact_bytes is not None and exact_bytes <= min( + broadcast_limit, target.total_size + ): + return PrefilterDecision( + "broadcast_semi_join", + "exact_domain_fits", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) + return PrefilterDecision( + "skip", + "no_viable_filter", + target.total_size, + domain.total_rows, + estimated_cardinality=distinct_count, + bloom_bytes=bloom_bytes, + exact_bytes=exact_bytes, + ) diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py new file mode 100644 index 000000000000..76ecee442821 --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/prefilter_actor.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone execution of optional pushdown-filter hints.""" + +from __future__ import annotations + +from dataclasses import asdict +from typing import TYPE_CHECKING, Any + +from cudf_streaming import CardinalityEstimator +from rapidsmpf.streaming.core.actor import define_actor + +from cudf_polars.dsl.utils.naming import names_to_indices +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.join import add_prefilter +from cudf_polars.streaming.actor_graph.prefilter import ( + PrefilterExecution, + choose_prefilter_method, +) +from cudf_polars.streaming.actor_graph.utils import ( + ChannelManager, + ChunkSampler, + gather_in_task_group, + process_children, + recv_metadata, + replay_buffered_channel, + sample_inputs, + shutdown_on_error, +) +from cudf_polars.streaming.filter_hint import PushdownFilterHint + +if TYPE_CHECKING: + from collections.abc import Sequence + + from cudf_streaming.table_chunk import TableChunk + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + + from cudf_polars.dsl.ir import IR, IRExecutionContext + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.utils import TableSizeStats + from cudf_polars.utils.config import StreamingExecutor + + +@define_actor() +async def pushdown_filter_actor( + context: Context, + comm: Communicator, + ir: PushdownFilterHint, + ir_context: IRExecutionContext, + ch_out: Channel[TableChunk], + ch_target: Channel[TableChunk], + ch_domain: Channel[TableChunk], + executor: StreamingExecutor, + collective_id: int, +) -> None: + """Choose and optionally execute one standalone pushdown-filter hint.""" + collected_samples: Sequence[TableSizeStats] = [] + async with shutdown_on_error( + context, + ch_out, + ch_target, + ch_domain, + trace_ir=ir, + ir_context=ir_context, + ) as tracer: + try: + target_metadata, domain_metadata = await gather_in_task_group( + recv_metadata(ch_target, context), + recv_metadata(ch_domain, context), + ) + dynamic_planning = executor.dynamic_planning + if dynamic_planning is None: + raise ValueError("Standalone prefilters require dynamic planning") + collected_samples = await sample_inputs( + context, + comm, + ( + ChunkSampler( + context=context, + ch_in=ch_target, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=target_metadata.local_count, + ), + ChunkSampler( + context=context, + ch_in=ch_domain, + max_chunks=dynamic_planning.sample_chunk_count, + max_bytes=executor.target_partition_size, + ch_in_chunk_count=domain_metadata.local_count, + cardinality_estimator=CardinalityEstimator( + context, comm, tag=collective_id + ), + cardinality_columns=names_to_indices( + ir.domain_on, ir.children[1].schema + ), + ), + ), + collective_id, + ) + if len(collected_samples) != 2: + raise ValueError("Standalone prefilters require two input samples") + target_sample, domain_sample = collected_samples + config = executor.join_filter_pushdown + if config is None: + raise ValueError("Standalone prefilter has no runtime configuration") + decision = choose_prefilter_method( + ir.domain_on, + target_sample, + domain_sample, + broadcast_limit=executor.broadcast_limit, + bloom_filter_max_size=config.bloom_filter_max_size, + ) + trace = asdict(decision) + trace["placement"] = "standalone" + trace["target_on"] = [key.name for key in ir.target_on] + trace["domain_on"] = [key.name for key in ir.domain_on] + trace_stats = trace if tracer is not None else None + if tracer is not None: + tracer.decision = decision.method + tracer.set_extra("prefilter", trace) + + if decision.method == "skip": + domain_sample.chunks.clear() + await gather_in_task_group( + ch_domain.shutdown(context), + replay_buffered_channel( + context, + ch_out, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ), + ) + else: + target, domain = ir.children + domain_indices = names_to_indices(ir.domain_on, domain.schema) + if domain_indices != tuple(range(len(domain.schema))): + raise ValueError("Pushdown filter domains must contain only keys") + + execution = PrefilterExecution(context) + ch_target_replay: Channel[TableChunk] = context.create_channel() + ch_domain_replay: Channel[TableChunk] = context.create_channel() + execution.add_channel(ch_target_replay) + execution.add_channel(ch_domain_replay) + execution.add_task( + replay_buffered_channel( + context, + ch_target_replay, + ch_target, + target_sample.chunks, + target_metadata, + trace_ir=ir, + ) + ) + execution.add_task( + replay_buffered_channel( + context, + ch_domain_replay, + ch_domain, + domain_sample.chunks, + domain_metadata, + trace_ir=ir, + ) + ) + add_prefilter( + execution, + comm, + spec=ir, + decision=decision, + target=target, + domain=domain, + ch_target=ch_target_replay, + ch_domain_keys=ch_domain_replay, + ch_filtered=ch_out, + collective_id=collective_id, + ir_context=ir_context, + trace_stats=trace_stats, + ) + async with shutdown_on_error( + context, + *execution.channels, + trace_ir=ir, + ir_context=ir_context, + ): + await gather_in_task_group(*execution.tasks) + finally: + for sample in collected_samples: + sample.chunks.clear() + + +@generate_ir_sub_network.register(PushdownFilterHint) +def generate_pushdown_filter_subnetwork( + ir: PushdownFilterHint, rec: SubNetGenerator +) -> tuple[dict[IR, list[Any]], dict[IR, ChannelManager]]: + """Generate the actor subnetwork for a standalone filter hint.""" + target, domain = ir.children + actors, channels = process_children(ir, rec) + channels[ir] = ChannelManager(rec.state["context"]) + (collective_id,) = rec.state["collective_id_map"][ir] + actors[ir] = [ + pushdown_filter_actor( + rec.state["context"], + rec.state["comm"], + ir, + rec.state["ir_context"], + channels[ir].reserve_input_slot(), + channels[target].reserve_output_slot(), + channels[domain].reserve_output_slot(), + rec.state["config_options"].executor, + collective_id, + ) + ] + return actors, channels 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 c39f1d1c3283..7472bdecb30e 100644 --- a/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -7,6 +7,7 @@ import asyncio import contextlib import itertools +import math import operator import struct import time @@ -54,7 +55,6 @@ Callable, Coroutine, Generator, - Iterable, Iterator, Sequence, ) @@ -176,7 +176,7 @@ def _keys_match( class ChunkStore: - """Ordered spillable buffer for TableChunk messages.""" + """Ordered spillable buffer for Messages.""" def __init__(self, ctx: Context) -> None: self._mids: deque[int] = deque() @@ -186,6 +186,12 @@ def __len__(self) -> int: """Return the number of messages in the store.""" return len(self._mids) + def clear(self) -> None: + """Discard all messages in the store.""" + for mid in self._mids: + self._store.extract(mid=mid) + self._mids.clear() + def insert(self, msg: Message) -> None: """Insert a message into the store.""" self._mids.append(self._store.insert(msg)) @@ -325,6 +331,7 @@ async def shutdown_on_error( record["row_count"] = tracer.row_count if tracer.decision is not None: record["decision"] = tracer.decision + record.update(tracer.extra) cudf_polars.dsl.tracing.log( "Streaming Actor", start=start, stop=stop, **record ) @@ -1099,6 +1106,57 @@ class TableSizeStats: cardinality: CardinalityEstimate | None = None """Global cardinality statistics for the sampled rows, when requested.""" + def distinct_count(self) -> int | None: + """Extrapolate sampled distinct count to the estimated full row count.""" + if self.total_rows == 0: + return 0 + if self.cardinality is None or self.cardinality.row_count == 0: + return None + return min( + self.total_rows, + math.ceil( + self.cardinality.distinct_count + * self.total_rows + / self.cardinality.row_count + ), + ) + + +async def aggregate_table_size_stats( + context: Context, + comm: Communicator, + samples: tuple[TableSizeStats, ...], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Aggregate table-size and row estimates across ranks.""" + totals = await allgather_reduce( + context, + comm, + collective_id, + *( + value + for sample in samples + for value in ( + sample.total_size, + sample.total_rows, + sample.total_chunks, + int(sample.is_complete), + ) + ), + ) + totals_iter = iter(totals) + return tuple( + TableSizeStats( + chunks=sample.chunks, + total_size=next(totals_iter), + total_rows=next(totals_iter), + total_chunks=next(totals_iter), + is_complete=next(totals_iter) == comm.nranks, + cardinality=sample.cardinality, + ) + for sample in samples + ) + @dataclass(frozen=True) class ChunkSampler: @@ -1228,6 +1286,26 @@ async def sample(self) -> TableSizeStats: ) +async def sample_inputs( + context: Context, + comm: Communicator, + samplers: Sequence[ChunkSampler], + collective_id: int, +) -> tuple[TableSizeStats, ...]: + """Sample input channels concurrently and aggregate their statistics.""" + if not samplers: + return () + local_samples = await gather_in_task_group( + *(sampler.sample() for sampler in samplers) + ) + return await aggregate_table_size_stats( + context, + comm, + tuple(local_samples), + collective_id, + ) + + async def _sample_chunks( context: Context, ch: Channel[TableChunk], @@ -1279,7 +1357,7 @@ async def replay_buffered_channel( context: Context, ch_out: Channel[TableChunk], ch_in: Channel[TableChunk], - buffered_chunks: Iterable[Message], + buffered_chunks: ChunkStore, metadata: ChannelMetadata, *, trace_ir: IR, @@ -1296,19 +1374,23 @@ async def replay_buffered_channel( ch_in The buffered input channel. buffered_chunks - Buffered messages to yield first. May be empty. + The buffered chunks to yield first. The store is empty when this + coroutine exits, including on cancellation or error. metadata The metadata to send to the output channel. trace_ir The IR node to trace. Passed through to shutdown_on_error. """ - async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): - await send_metadata(ch_out, context, metadata) - for msg in buffered_chunks: - await ch_out.send(context, msg) - while (msg := await ch_in.recv(context)) is not None: - await ch_out.send(context, msg) - await ch_out.drain(context) + try: + async with shutdown_on_error(context, ch_out, ch_in, trace_ir=trace_ir): + await send_metadata(ch_out, context, metadata) + for msg in buffered_chunks: + await ch_out.send(context, msg) + while (msg := await ch_in.recv(context)) is not None: + await ch_out.send(context, msg) + await ch_out.drain(context) + finally: + buffered_chunks.clear() @dataclass(frozen=True) diff --git a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 54116be53657..29ff22ca30e9 100644 --- a/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -874,7 +874,12 @@ def print_query_plan( elif CUDF_POLARS_AVAILABLE: assert isinstance(engine, pl.GPUEngine) if args.explain_logical: - logical_plan = explain_query(q, engine, physical=False) + logical_plan = explain_query( + q, + engine, + optimized=run_config.frontend in _STREAMING_FRONTENDS, + physical=False, + ) if args.explain and run_config.frontend in _STREAMING_FRONTENDS: plan = explain_query(q, engine) else: diff --git a/python/cudf_polars/cudf_polars/streaming/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py index 011e3be232a7..fda95175eda6 100644 --- a/python/cudf_polars/cudf_polars/streaming/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -37,8 +37,14 @@ from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import IOPartitionFlavor +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, + PushdownFilterHint, +) from cudf_polars.streaming.io import StreamingScan, scan_partition_plan -from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.parallel import lower_ir_graph, optimize_with_stats from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.statistics import ( collect_statistics, @@ -53,6 +59,7 @@ from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IR from cudf_polars.streaming.base import PartitionInfo, StatsCollector + from cudf_polars.streaming.filter_hint import Prefilter @dataclasses.dataclass @@ -84,6 +91,7 @@ def explain_query( q: pl.LazyFrame, engine: pl.GPUEngine, *, + optimized: bool = True, physical: bool = True, executor: concurrent.futures.Executor | None = None, ) -> str: @@ -96,6 +104,9 @@ def explain_query( The LazyFrame to explain. engine : pl.GPUEngine The configured GPU engine to use. + optimized + If True and showing the logical plan, run cudf-polars specific + query optimization. physical : bool, default True If True, show the physical (lowered) plan. If False, show the logical (pre-lowering) plan. @@ -134,6 +145,8 @@ def explain_query( # Include row-count statistics for the logical plan with cm: stats = collect_statistics(ir, config, executor) + if optimized: + ir = optimize_with_stats(ir, config, stats) return _repr_ir_tree(ir, stats=stats) else: return _repr_ir_tree(ir) @@ -469,6 +482,29 @@ def _(ir: Join, *, offset: str = "") -> str: return _repr_header(offset, f"JOIN {ir.options[0]} {left_on} {right_on}", ir.schema) +@_repr_ir.register +def _(ir: JoinWithPrefilter, *, offset: str = "") -> str: + left_on = tuple(ne.name for ne in ir.left_on) + right_on = tuple(ne.name for ne in ir.right_on) + prefilters = tuple(type(prefilter.domain).__name__ for prefilter in ir.prefilters) + return _repr_header( + offset, + f"JOIN {ir.options[0]} {left_on} {right_on} {prefilters=}", + ir.schema, + ) + + +@_repr_ir.register +def _(ir: PushdownFilterHint, *, offset: str = "") -> str: + target_on = tuple(ne.name for ne in ir.target_on) + domain_on = tuple(ne.name for ne in ir.domain_on) + return _repr_header( + offset, + f"PUSHDOWN FILTER HINT {target_on} {domain_on} {ir.placement}", + ir.schema, + ) + + _BinaryOperator = plc.binaryop.BinaryOperator _BINOP_SYMBOLS: dict[_BinaryOperator, str] = { _BinaryOperator.EQUAL: "==", @@ -580,6 +616,45 @@ def _(ir: Join) -> dict[str, Serializable]: } +def _serialize_prefilter(prefilter: Prefilter) -> dict[str, Serializable]: + """Serialize a normalized join prefilter descriptor.""" + properties: dict[str, Serializable] = { + "type": type(prefilter).__name__, + "target_side": prefilter.target_side, + "target_on": [ne.name for ne in prefilter.target_on], + "domain_on": [ne.name for ne in prefilter.domain_on], + "nulls_equal": prefilter.nulls_equal, + } + if isinstance(prefilter.domain, JoinInputDomain): + properties["domain"] = { + "type": type(prefilter.domain).__name__, + "side": prefilter.domain.side, + } + elif isinstance(prefilter.domain, ExternalDomain): + properties["domain"] = {"type": type(prefilter.domain).__name__} + return properties + + +@_serialize_properties.register +def _(ir: JoinWithPrefilter) -> dict[str, Serializable]: + return { + "how": ir.options[0], + "left_on": [ne.name for ne in ir.left_on], + "right_on": [ne.name for ne in ir.right_on], + "prefilters": [_serialize_prefilter(prefilter) for prefilter in ir.prefilters], + } + + +@_serialize_properties.register +def _(ir: PushdownFilterHint) -> dict[str, Serializable]: + return { + "target_on": [ne.name for ne in ir.target_on], + "domain_on": [ne.name for ne in ir.domain_on], + "nulls_equal": ir.nulls_equal, + "placement": ir.placement, + } + + @_serialize_properties.register def _(ir: GroupBy) -> dict[str, Serializable]: return { @@ -815,4 +890,10 @@ def from_query( """ config_options = ConfigOptions.from_polars_engine(engine) ir = Translator(q._ldf.visit(), engine).translate_ir() + if not lowered and config_options.executor.name == "streaming": + with concurrent.futures.ThreadPoolExecutor( + thread_name_prefix="cudf-polars-explain" + ) as executor: + stats = collect_statistics(ir, config_options, executor) + ir = optimize_with_stats(ir, config_options, stats) return cls.from_ir(ir, config_options=config_options, lowered=lowered) diff --git a/python/cudf_polars/cudf_polars/streaming/filter_hint.py b/python/cudf_polars/cudf_polars/streaming/filter_hint.py new file mode 100644 index 000000000000..de5f59ad3cdd --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/filter_hint.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Logical filter hints for the streaming runtime.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, ClassVar, Literal, TypeAlias + +from cudf_polars.dsl.ir import IR, Join + +if TYPE_CHECKING: + from collections.abc import Sequence + + from cudf_polars.containers import DataFrame + from cudf_polars.dsl.expr import NamedExpr + from cudf_polars.dsl.ir import IRExecutionContext + from cudf_polars.typing import Schema + + +JoinSide: TypeAlias = Literal["left", "right"] +HintPlacement: TypeAlias = Literal["join_input", "pushed_down"] + + +@dataclass(frozen=True, slots=True) +class JoinInputDomain: + """A prefilter domain provided by an input of its owning join.""" + + side: JoinSide + + +@dataclass(frozen=True, slots=True) +class ExternalDomain: + """A prefilter domain provided by an additional join input.""" + + +PrefilterDomain: TypeAlias = JoinInputDomain | ExternalDomain + + +@dataclass(frozen=True, slots=True) +class Prefilter: + """Description of an optional join prefilter.""" + + target_side: JoinSide + target_on: tuple[NamedExpr, ...] + domain: PrefilterDomain + domain_on: tuple[NamedExpr, ...] + nulls_equal: bool + + +class JoinWithPrefilter(Join): + """Lowered join with normalized prefilter descriptors.""" + + __slots__ = ("prefilters",) + _non_child = ("schema", "left_on", "right_on", "options", "prefilters") + _n_non_child_args = 4 + + prefilters: tuple[Prefilter, ...] + + def __init__( + self, + schema: Schema, + left_on: Sequence[NamedExpr], + right_on: Sequence[NamedExpr], + options: Any, + prefilters: Sequence[Prefilter], + left: IR, + right: IR, + *external_domains: IR, + ): + self.schema = schema + self.left_on = tuple(left_on) + self.right_on = tuple(right_on) + self.options = options + self.prefilters = tuple(prefilters) + self.children = (left, right, *external_domains) + self._non_child_args = ( + self.left_on, + self.right_on, + self.options, + self.prefilters, + ) + + if not self.prefilters: + raise ValueError("JoinWithPrefilter requires at least one prefilter") + external_domain_count = sum( + isinstance(prefilter.domain, ExternalDomain) + for prefilter in self.prefilters + ) + if external_domain_count != len(external_domains): + raise ValueError( + "External prefilters and additional JoinWithPrefilter children " + "must align" + ) + + @classmethod + def do_evaluate( + cls, + left_on: tuple[NamedExpr, ...], + right_on: tuple[NamedExpr, ...], + options: Any, + prefilters: tuple[Prefilter, ...], + left: DataFrame, + right: DataFrame, + *external_domains: DataFrame, + context: IRExecutionContext, + ) -> DataFrame: + """Evaluate the join while ignoring its optional prefilters.""" + del prefilters, external_domains + return Join.do_evaluate( + left_on, + right_on, + options, + left, + right, + context=context, + ) + + +class PushdownFilterHint(IR): + """ + Optional join-key filter placed in a logical plan. + + The first child is the target to filter and the second child, the + domain, provides the keys to filter against. Applying the filter is + optional. + """ + + __slots__ = ("domain_on", "nulls_equal", "placement", "target_on") + _non_child: ClassVar[tuple[str, ...]] = ( + "schema", + "target_on", + "domain_on", + "nulls_equal", + "placement", + ) + _n_non_child_args: ClassVar[int] = 4 + + target_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the target.""" + domain_on: tuple[NamedExpr, ...] + """Expressions selecting filter keys from the domain.""" + nulls_equal: bool + """Whether null key values compare equal.""" + placement: HintPlacement + """Whether the hint remains at the motivating join input.""" + + def __init__( + self, + schema: Schema, + target_on: Sequence[NamedExpr], + domain_on: Sequence[NamedExpr], + nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, + target: IR, + domain: IR, + ): + self.schema = schema + self.target_on = tuple(target_on) + self.domain_on = tuple(domain_on) + self.nulls_equal = nulls_equal + self.placement = placement + self._non_child_args = ( + self.target_on, + self.domain_on, + self.nulls_equal, + self.placement, + ) + self.children = (target, domain) + + @classmethod + def do_evaluate( + cls, + target_on: tuple[NamedExpr, ...], + domain_on: tuple[NamedExpr, ...], + nulls_equal: bool, # noqa: FBT001 + placement: HintPlacement, + target: DataFrame, + domain: DataFrame, + *, + context: IRExecutionContext, + ) -> DataFrame: + """Ignore the optional filter and return the target.""" + del placement + return target diff --git a/python/cudf_polars/cudf_polars/streaming/join.py b/python/cudf_polars/cudf_polars/streaming/join.py index 63d76d6c328f..a9a266a5c166 100644 --- a/python/cudf_polars/cudf_polars/streaming/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -8,10 +8,17 @@ from functools import reduce from typing import TYPE_CHECKING -from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice +from cudf_polars.dsl.ir import ConditionalJoin, Join, Projection, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.streaming.base import PartitionInfo from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.filter_hint import ( + ExternalDomain, + JoinInputDomain, + JoinWithPrefilter, + Prefilter, + PushdownFilterHint, +) from cudf_polars.streaming.repartition import Repartition from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.streaming.utils import ( @@ -25,6 +32,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR + from cudf_polars.streaming.filter_hint import JoinSide from cudf_polars.streaming.parallel import LowerIRTransformer @@ -149,6 +157,144 @@ def _has_non_pointwise_keys(ir: Join) -> bool: return not all(expr.is_pointwise for expr in traversal(keys)) +def is_direct_join_prefilter(ir: IR) -> bool: + """Return whether a hint belongs to its immediately enclosing join.""" + return isinstance(ir, PushdownFilterHint) and ir.placement == "join_input" + + +def lower_join_with_prefilters( + ir: Join, + rec: LowerIRTransformer, +) -> tuple[Join, MutableMapping[IR, PartitionInfo]]: + """Lower a join and normalize its adjacent filter hints.""" + targets = tuple( + child.children[0] if is_direct_join_prefilter(child) else child + for child in ir.children + ) + lowered_targets, target_partition_info = zip( + *(rec(target) for target in targets), + strict=True, + ) + partition_info: MutableMapping[IR, PartitionInfo] = reduce( + operator.or_, target_partition_info + ) + + if all( + isinstance(target, Repartition) and partition_info[target].count == 1 + for target in lowered_targets + ): + # This join will execute partition-wise, so its optional prefilters + # are unnecessary. Moreover, the piecewise join special case + # execution at runtime never has a chance to shut down prefilter + # channels that would be produced, which would leave an actor graph + # in a deadlocked state. Since they are unnecessary, drop them + # before lowering their domains and before the actor graph derives + # fanout from the lowered DAG. + return ( + Join( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + *lowered_targets, + ), + partition_info, + ) + + prefilters: list[Prefilter] = [] + external_domains: list[IR] = [] + claimed_sides: set[JoinSide] = set() + for target_index, child in enumerate(ir.children): + if not is_direct_join_prefilter(child): + continue + assert isinstance(child, PushdownFilterHint) + + _target, domain = child.children + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) + + # A key-only Projection retains an explicit edge to its source. If that + # source is a join input and contains every requested key, the join can + # project those keys itself rather than execute a separate domain input. + left, right = lowered_targets + direct_domain = domain + while True: + if direct_domain == left and direct_domain == right: + domain_side: JoinSide | None = "right" if target_index == 0 else "left" + break + if direct_domain == left: + domain_side = "left" + break + if direct_domain == right: + domain_side = "right" + break + if isinstance(direct_domain, Projection) and all( + key.name in direct_domain.children[0].schema for key in child.domain_on + ): + (direct_domain,) = direct_domain.children + continue + domain_side = None + break + + target_side: JoinSide = "left" if target_index == 0 else "right" + if domain_side in claimed_sides: + domain_side = None + elif domain_side is not None: + claimed_sides.add(domain_side) + + if domain_side is not None: + prefilters.append( + Prefilter( + target_side, + child.target_on, + JoinInputDomain(domain_side), + child.domain_on, + child.nulls_equal, + ) + ) + else: + external_domains.append(domain) + prefilters.append( + Prefilter( + target_side, + child.target_on, + ExternalDomain(), + child.domain_on, + child.nulls_equal, + ) + ) + + return ( + JoinWithPrefilter( + ir.schema, + ir.left_on, + ir.right_on, + ir.options, + prefilters, + *lowered_targets, + *external_domains, + ), + partition_info, + ) + + +@lower_ir_node.register(PushdownFilterHint) +def _( + ir: PushdownFilterHint, rec: LowerIRTransformer +) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: + """Preserve optional filters for dynamic execution, otherwise discard them.""" + target, domain = ir.children + target, partition_info = rec(target) + if not _dynamic_planning_on(rec.state["config_options"]): + return target, partition_info + + domain, domain_partition_info = rec(domain) + partition_info.update(domain_partition_info) + lowered = ir.reconstruct((target, domain)) + partition_info[lowered] = partition_info[target] + return lowered, partition_info + + @lower_ir_node.register(ConditionalJoin) def _( ir: ConditionalJoin, rec: LowerIRTransformer @@ -214,17 +360,33 @@ def _( ) return rec(Slice(ir.schema, offset, length, new_join)) - # Lower children - children, _partition_info = zip(*(rec(c) for c in ir.children), strict=True) - partition_info = reduce(operator.or_, _partition_info) - - # Check for dynamic planning - may have more partitions at runtime config_options = rec.state["config_options"] dynamic_planning = _dynamic_planning_on(config_options) + has_non_pointwise_keys = _has_non_pointwise_keys(ir) + if ( + dynamic_planning + and ir.options[0] != "Cross" + and ir.options[5] == "none" + and not has_non_pointwise_keys + and any(is_direct_join_prefilter(child) for child in ir.children) + ): + preserve_prefilters = True + else: + preserve_prefilters = False - left, right = children + if preserve_prefilters: + ir, partition_info = lower_join_with_prefilters(ir, rec) + children = ir.children + else: + # Hints not owned by an adaptive join use the generic identity lowering. + children, _partition_info = zip( + *(rec(child) for child in ir.children), + strict=True, + ) + partition_info = reduce(operator.or_, _partition_info) + + left, right = children[:2] output_count = max(partition_info[left].count, partition_info[right].count) - has_non_pointwise_keys = _has_non_pointwise_keys(ir) if output_count == 1 and not dynamic_planning: new_node = ir.reconstruct(children) partition_info[new_node] = PartitionInfo(count=1) diff --git a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py index b099e1584ea6..a82c49ef9d82 100644 --- a/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py +++ b/python/cudf_polars/cudf_polars/streaming/join_filter_pushdown.py @@ -5,13 +5,14 @@ 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:: +the original join. It records that opportunity with a logical +``PushdownFilterHint``:: left join[left.key = right.key] right -> - (left semijoin[left.key = right.key] project(right.key)) + PushdownFilterHint(left, left.key, project(right.key), right.key) join[left.key = right.key] right In this example, the right hand table is selected to pre-filter the left @@ -49,14 +50,14 @@ 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. + For a multi-key join, a rewrite that first hints that the domain should be + filtered 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, source scan facts, 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. +``apply_candidate`` then constructs the selected filter-hint rewrite. Row estimates, selectivity propagation, thresholds, and candidate scores are only heuristics for deciding whether a safe rewrite is likely to improve @@ -100,11 +101,13 @@ ColumnRef, column_domain_bindings, ) +from cudf_polars.streaming.filter_hint import PushdownFilterHint if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Mapping, Sequence from cudf_polars.streaming.base import StatsCollector + from cudf_polars.streaming.filter_hint import HintPlacement from cudf_polars.typing import GenericTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -249,6 +252,12 @@ def analyze_plan(ir: IR, stats: StatsCollector) -> PlanFacts: rows = node.df.shape()[0] elif isinstance(node, (Select, Projection, HStack, Filter, Distinct, GroupBy)): rows = row_estimates[node.children[0]] + elif isinstance(node, PushdownFilterHint): + rows = _estimate_join_rows( + "Semi", + row_estimates[node.children[0]], + row_estimates[node.children[1]], + ) elif isinstance(node, Join): rows = _estimate_join_rows( node.options[0], @@ -329,7 +338,7 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: Returns ------- bool - True if a semijoin cannot be pushed past this node, otherwise False. + True if a filter hint cannot be pushed past this node, otherwise False. """ # TODO: Need better cost model to handle nodes that are shared. Pushing # a filter into a shared node will typically mean that it is no longer @@ -350,11 +359,11 @@ def blocks_pushdown(node: IR, facts: PlanFacts) -> bool: ) -def semijoin_pushdown_candidates( +def filter_hint_pushdown_candidates( facts: PlanFacts, root: IR, column: str ) -> Iterator[tuple[ColumnRef, tuple[int, ...]]]: """ - Yield column domain lineage providing valid locations for semijoin pushdown. + Yield column domain lineage providing valid locations for a filter hint. Parameters ---------- @@ -452,7 +461,7 @@ def _(node: Join, rec: GenericTransformer[IR, IR, _RewriteState]) -> IR: if node is original: facts = rec.state["facts"] else: - # Child rewrites introduce new semi joins and reconstructed ancestors. + # Child rewrites introduce new filter hints 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"]) @@ -473,13 +482,13 @@ def apply_candidate(ir: Join, candidate: Candidate) -> IR: left, right = ir.children domain = _make_domain(candidate, ir) target = candidate.target - target_filter = _make_semi_join( + target_filter = _make_filter_hint( 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], + placement="join_input" if not target.path else "pushed_down", ) if candidate.target_side == "left": left = replace_at_path(left, target.path, target_filter) @@ -611,7 +620,7 @@ def _simple_candidates( continue if contains_node(target.node, domain.node): continue - if domain.is_single_source and has_filtering_semi_ancestor( + if domain.is_single_source and has_filtering_hint_ancestor( target_child, target.path ): continue @@ -704,7 +713,7 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: candidate.constraint_domain.column, candidate.target_constraint_key, ) - constrained = _make_semi_join( + constrained = _make_filter_hint( candidate.domain.node, expr.Col( candidate.domain.node.schema[candidate.domain.columns[1]], @@ -716,7 +725,6 @@ def _make_domain(candidate: Candidate, ir: Join) -> IR: 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 @@ -735,20 +743,21 @@ def _project_bound_key(source: IR, bound_column: str, output_key: expr.Col) -> S ) -def _make_semi_join( +def _make_filter_hint( target: IR, target_key: expr.Col, domain: IR, domain_key: expr.Col, *, nulls_equal: bool, - suffix: str, -) -> Join: - return Join( + placement: HintPlacement = "pushed_down", +) -> PushdownFilterHint: + return PushdownFilterHint( target.schema, (expr.NamedExpr(target_key.name, target_key),), (expr.NamedExpr(domain_key.name, domain_key),), - ("Semi", nulls_equal, None, suffix, False, "none"), + nulls_equal, + placement, target, domain, ) @@ -784,7 +793,7 @@ def _smallest_key_producer( exclude: IR | None = None, ) -> _Producer | None: producers = [] - for reference, path in semijoin_pushdown_candidates(facts, root, column): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name if node is exclude: continue @@ -840,7 +849,7 @@ def _smallest_node_containing_all( 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): + for reference, path in filter_hint_pushdown_candidates(facts, root, column): node, bound_column = reference.node, reference.name producer = make_producer(node, (bound_column,), path, facts) if producer is None: @@ -890,11 +899,11 @@ def domain_cost_is_small( return domain.cost / target.rows <= threshold -def has_filtering_semi_ancestor(root: IR, path: Sequence[int]) -> bool: - """Return whether a selected child edge is below a filtering semi join.""" +def has_filtering_hint_ancestor(root: IR, path: Sequence[int]) -> bool: + """Return whether a selected child edge is below a pushdown-filter hint.""" node = root for child_index in path: - if isinstance(node, Join) and node.options[0] == "Semi" and child_index == 0: + if isinstance(node, PushdownFilterHint) and child_index == 0: return True node = node.children[child_index] return False diff --git a/python/cudf_polars/cudf_polars/streaming/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py index d2098b20cb7b..1f2dd6bcae33 100644 --- a/python/cudf_polars/cudf_polars/streaming/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -15,6 +15,7 @@ # handlers at import time so the dispatch table is populated before any query # is lowered. import cudf_polars.streaming.distinct +import cudf_polars.streaming.filter_hint import cudf_polars.streaming.groupby import cudf_polars.streaming.io import cudf_polars.streaming.join diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 47c8a9f6dfdb..03724c5010a3 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -509,7 +509,7 @@ def __post_init__(self) -> None: # noqa: D105 @dataclasses.dataclass(frozen=True) class JoinFilterPushdownOptions: """ - Configuration options for join filter pushdown in the logical plan. + Configuration options for join filter pushdown. 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 @@ -517,7 +517,8 @@ class JoinFilterPushdownOptions: 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. + selecting a side to be filtered by the keys of the other side. At execution + time, these options also control how optional filters are applied. Pass ``None`` to ``StreamingExecutor(join_filter_pushdown=...)`` to disable the rewrite. @@ -530,6 +531,10 @@ class JoinFilterPushdownOptions: 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. + bloom_filter_max_size + Maximum Bloom-filter size in bytes. If the estimated Bloom filter exceeds + this size, an exact semi-join is preferred when its projected keys fit the + broadcast limit. Set to 0 to disable Bloom filters. Default is 32 MiB. trace Whether to emit plan-time trace decisions for filter decisions. Default is False. """ @@ -541,6 +546,13 @@ class JoinFilterPushdownOptions: f"{_env_prefix}__THRESHOLD", float, default=0.5 ) ) + bloom_filter_max_size: int = dataclasses.field( + default_factory=_make_default_factory( + f"{_env_prefix}__BLOOM_FILTER_MAX_SIZE", + int, + default=32 * 1024 * 1024, + ) + ) trace: bool = dataclasses.field( default_factory=_make_default_factory( f"{_env_prefix}__TRACE", _bool_converter, default=False @@ -555,6 +567,12 @@ def __post_init__(self) -> None: # noqa: D105 object.__setattr__(self, "threshold", threshold) if not 0.0 <= threshold <= 1.0: raise ValueError("threshold must be between 0 and 1") + if isinstance(self.bloom_filter_max_size, bool) or not isinstance( + self.bloom_filter_max_size, int + ): + raise TypeError("bloom_filter_max_size must be an int") + if self.bloom_filter_max_size < 0: + raise ValueError("bloom_filter_max_size must be non-negative") if not isinstance(self.trace, bool): raise TypeError("trace must be a bool") @@ -826,8 +844,8 @@ class StreamingExecutor: :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. + :class:`~cudf_polars.utils.config.JoinFilterPushdownOptions` for + more. Disabled by default (or by explicitly providing ``None``). Enable through environment variables with ``CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN=1``. @@ -926,7 +944,7 @@ class StreamingExecutor: default_factory=DynamicPlanningOptions ) join_filter_pushdown: JoinFilterPushdownOptions | None = dataclasses.field( - default_factory=JoinFilterPushdownOptions + default=None ) max_concurrent_io_tasks: MaxConcurrentIOTasks = dataclasses.field( default_factory=_make_default_factory( @@ -1259,15 +1277,14 @@ def from_polars_engine( 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: + if "join_filter_pushdown" not in user_executor_options: env_join_filter_pushdown = os.environ.get( "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "0" ) - if not _bool_converter(env_join_filter_pushdown): - user_executor_options["join_filter_pushdown"] = None + if _bool_converter(env_join_filter_pushdown): + user_executor_options["join_filter_pushdown"] = ( + JoinFilterPushdownOptions() + ) executor = StreamingExecutor(**user_executor_options) case _: # pragma: no cover; Unreachable diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 0aaf27d073d8..e0dc29917dd9 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -779,6 +779,6 @@ argument; user-supplied keys are merged with reserved entries set by `SPMDEngine [spmd-wiki]: https://en.wikipedia.org/wiki/Single_program,_multiple_data [ray-docs]: https://docs.ray.io/en/latest/ [ray-actors]: https://docs.ray.io/en/latest/ray-core/actors.html -[rapidsmpf-communicator]: https://docs.rapids.ai/api/rapidsmpf/stable/glossary/#term-Communicator -[rapidsmpf-context]: https://docs.rapids.ai/api/rapidsmpf/stable/glossary/#term-Context +[rapidsmpf-communicator]: https://docs.nvidia.com/rapidsmpf/latest/glossary/#term-Communicator +[rapidsmpf-context]: https://docs.nvidia.com/rapidsmpf/latest/glossary/#term-Context [polars-gpuengine]: https://docs.pola.rs/api/python/stable/reference/api/polars.GPUEngine.html diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index 7fd92c53ebbb..c95499462e13 100644 --- a/python/cudf_polars/docs/overview.md +++ b/python/cudf_polars/docs/overview.md @@ -647,10 +647,6 @@ another `nvtx` range (e.g. `Scan.do_evaluate`, `GroupBy.do_evaluate`, etc.). These provide a higher-level grouping over the lower-level libcudf calls (e.g. `read_chunk`, `aggregate`). -Finally, if using [rapidsmpf](https://docs.rapids.ai/api/rapidsmpf/nightly/) -for shuffling, the methods inserting and extracting partitions to shuffle are -annotated with nvtx ranges. - # Query Plans The module `cudf_polars.streaming.explain` contains functions for dumping diff --git a/python/cudf_polars/tests/streaming/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py index 5ea6be578ae4..9458c507e623 100644 --- a/python/cudf_polars/tests/streaming/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -137,6 +137,54 @@ def test_explain_logical_plan_with_join(tmp_path, df): assert "JOIN Inner ('x',) ('x',)" in plan +def test_explain_pushdown_filter_hint_in_dynamic_physical_plan(): + domain = ( + pl.LazyFrame({"key": [1, 99], "active": [True, False]}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame({"key": [i % 10 for i in range(20)]}) + query = domain.join(target, on="key") + engine = pl.GPUEngine( + executor="streaming", + raise_on_fail=True, + executor_options={"join_filter_pushdown": {"threshold": 0.5}}, + ) + + logical = explain_query(query, engine, physical=False) + physical = explain_query(query, engine, physical=True) + logical_serialized = serialize_query(query, engine, physical=False) + physical_serialized = serialize_query(query, engine, physical=True) + + assert "PUSHDOWN FILTER HINT ('key',) ('key',)" in logical + assert "prefilters=('JoinInputDomain',)" in physical + expected_properties = { + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + "placement": "join_input", + } + assert any( + node.type == "PushdownFilterHint" and node.properties == expected_properties + for node in logical_serialized.nodes.values() + ) + assert any( + node.type == "JoinWithPrefilter" + and node.properties["prefilters"] + == [ + { + "type": "Prefilter", + "target_side": "right", + "target_on": ["key"], + "domain_on": ["key"], + "nulls_equal": False, + "domain": {"type": "JoinInputDomain", "side": "left"}, + } + ] + for node in physical_serialized.nodes.values() + ) + + def test_explain_logical_plan_with_sort(tmp_path, df): make_partitioned_source(df, tmp_path, fmt="parquet", n_files=2) diff --git a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py index 7601788400db..90e5d52fa92e 100644 --- a/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py +++ b/python/cudf_polars/tests/streaming/test_join_filter_pushdown.py @@ -3,7 +3,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest @@ -11,11 +11,30 @@ from cudf_polars import Translator from cudf_polars.dsl.expr import Col -from cudf_polars.dsl.ir import Cache, DataFrameScan, Distinct, Join, Select, Slice -from cudf_polars.dsl.traversal import traversal +from cudf_polars.dsl.ir import ( + IR, + Cache, + DataFrameScan, + Distinct, + Join, + Projection, + Select, + Slice, +) +from cudf_polars.dsl.traversal import CachingVisitor, 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.base import PartitionInfo, StatsCollector +from cudf_polars.streaming.filter_hint import ( + JoinInputDomain, + JoinWithPrefilter, + Prefilter, + PushdownFilterHint, +) +from cudf_polars.streaming.join import ( + is_direct_join_prefilter, + lower_join_with_prefilters, +) from cudf_polars.streaming.join_filter_pushdown import ( CompositeCandidate, Decision, @@ -26,10 +45,15 @@ analyze_plan, apply_candidate, contains_node, + filter_hint_pushdown_candidates, optimize_join_filter_pushdown, - semijoin_pushdown_candidates, ) -from cudf_polars.streaming.parallel import optimize_with_stats, remove_cache_nodes +from cudf_polars.streaming.parallel import ( + lower_ir_graph, + optimize_with_stats, + remove_cache_nodes, +) +from cudf_polars.streaming.repartition import Repartition 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 @@ -77,6 +101,10 @@ def find_joins(ir: IR, how: str | None = None) -> list[Join]: ] +def find_hints(ir: IR) -> list[PushdownFilterHint]: + return [node for node in traversal([ir]) if isinstance(node, PushdownFilterHint)] + + 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) @@ -95,10 +123,12 @@ def dataframe_scan(ir: IR, column: str) -> DataFrameScan: return match -def join_key_names(join: Join) -> tuple[str, ...]: - """Return the column names used on the left of a simple-column join.""" - names = tuple(key.value.name for key in join.left_on if isinstance(key.value, Col)) - assert len(names) == len(join.left_on) +def hint_key_names(hint: PushdownFilterHint) -> tuple[str, ...]: + """Return the target column names used by a filter hint.""" + names = tuple( + key.value.name for key in hint.target_on if isinstance(key.value, Col) + ) + assert len(names) == len(hint.target_on) return names @@ -141,10 +171,11 @@ def test_simple_prefilter_filters_large_side( 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 not find_joins(optimized, "Semi") + hints = find_hints(optimized) + assert len(hints) == 1 + assert hints[0].children[0] is lineitem_ir + assert not find_hints(part_ir) assert_gpu_result_equal(simple_query, engine=engine, check_row_order=False) @@ -153,14 +184,90 @@ def test_filter_pushdown_is_independent_of_dynamic_planning( engine: SPMDEngine, ) -> None: root = translate_query(simple_query, engine) + config = make_config(dynamic_planning=False) optimized = optimize_join_filter_pushdown( root, StatsCollector(), - make_config(dynamic_planning=False), + config, + ) + + assert find_hints(optimized) + lowering = lower_ir_graph(root, config, StatsCollector()) + assert not any( + isinstance(node, JoinWithPrefilter) for node in traversal([lowering.lowered]) + ) + + +def test_adjacent_filter_hint_is_recorded_on_lowered_join( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + + lowering = lower_ir_graph(root, config, StatsCollector()) + + assert find_hints(lowering.optimized) + assert isinstance(lowering.lowered, JoinWithPrefilter) + left, right = lowering.lowered.children + assert not isinstance(left, PushdownFilterHint) + assert not isinstance(right, PushdownFilterHint) + (prefilter,) = lowering.lowered.prefilters + assert isinstance(prefilter, Prefilter) + assert isinstance(prefilter.domain, JoinInputDomain) + assert find_hints(lowering.optimized)[0].placement == "join_input" + assert prefilter.target_side == "right" + assert prefilter.domain.side == "left" + assert tuple(right.schema) == ("l_partkey", "l_suppkey") + assert tuple(ne.name for ne in prefilter.target_on) == ("l_partkey",) + assert tuple(ne.name for ne in prefilter.domain_on) == ("p_partkey",) + assert not prefilter.nulls_equal + assert tuple(left.schema) == ("p_partkey",) + assert not find_joins(lowering.lowered, "Semi") + + +def test_partition_wise_join_discards_prefilters_before_lowering_domains( + simple_query: pl.LazyFrame, + engine: SPMDEngine, +) -> None: + """Partition-wise joins must not retain optional prefilter inputs.""" + root = translate_query(simple_query, engine) + config = ConfigOptions.from_polars_engine(engine) + optimized = optimize_with_stats(root, config, StatsCollector()) + assert isinstance(optimized, Join) + + children = list(optimized.children) + (hint_index,) = ( + index for index, child in enumerate(children) if is_direct_join_prefilter(child) ) + hint = children[hint_index] + assert isinstance(hint, PushdownFilterHint) + domain = Projection(hint.children[1].schema, hint.children[1]) + children[hint_index] = hint.reconstruct((hint.children[0], domain)) + optimized = optimized.reconstruct(children) - assert find_joins(optimized, "Semi") + targets = tuple( + child.children[0] if is_direct_join_prefilter(child) else child + for child in optimized.children + ) + repartitions = tuple(Repartition(target.schema, target) for target in targets) + lowered_targets = dict(zip(targets, repartitions, strict=True)) + + def lower_target(child: IR, rec: Any) -> tuple[IR, dict[IR, PartitionInfo]]: + assert child in lowered_targets, "prefilter domain was lowered" + lowered = lowered_targets[child] + return lowered, {lowered: PartitionInfo(count=1)} + + rec: Any = CachingVisitor( + lower_target, + state={"config_options": config}, + ) + lowered, partition_info = lower_join_with_prefilters(optimized, rec) + + assert type(lowered) is Join + assert lowered.children == repartitions + assert domain not in partition_info def test_filter_pushdown_can_be_disabled( @@ -207,9 +314,9 @@ def test_nullable_join_keys_preserve_results( config, ) - semi_joins = find_joins(optimized, "Semi") - assert semi_joins - assert all(join.options[1] is nulls_equal for join in semi_joins) + hints = find_hints(optimized) + assert hints + assert all(hint.nulls_equal is nulls_equal for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -239,8 +346,8 @@ def test_prefilter_does_not_move_below_distinct_on_non_subset_column( config, ) - semis = find_joins(optimized, "Semi") - assert any(isinstance(semi.children[0], Distinct) for semi in semis) + hints = find_hints(optimized) + assert any(isinstance(hint.children[0], Distinct) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -267,7 +374,7 @@ def test_no_simple_filter_pushdown_when_domain_is_not_selective( assert decision == Decision(reason="no_profitable_domain") assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -329,12 +436,12 @@ def test_composite_filter_pushdown_constrains_domain_first( assert decision.reason == "applied" assert isinstance(decision.candidate, CompositeCandidate) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) 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 any(hint.children[0] is supplier_ir for hint in hints) + assert any(hint.children[0] is lineitem_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -388,16 +495,16 @@ def test_prefilter_uses_cheaper_source_domain_and_skips_expensive_domain( supplier_ir = dataframe_scan(root, "s_suppkey") lineitem_ir = dataframe_scan(root, "l_orderkey") orders_ir = dataframe_scan(root, "o_orderkey") - semis = find_joins(optimized, "Semi") - partkey_semis = [ - semi - for semi in semis - if semi.children[0] is lineitem_ir and join_key_names(semi) == ("l_partkey",) + hints = find_hints(optimized) + partkey_hints = [ + hint + for hint in hints + if hint.children[0] is lineitem_ir and hint_key_names(hint) == ("l_partkey",) ] - assert partkey_semis - assert not any(semi.children[0] is orders_ir for semi in semis) - assert contains_node(partkey_semis[0].children[1], part_ir) - assert not contains_node(partkey_semis[0].children[1], supplier_ir) + assert partkey_hints + assert not any(hint.children[0] is orders_ir for hint in hints) + assert contains_node(partkey_hints[0].children[1], part_ir) + assert not contains_node(partkey_hints[0].children[1], supplier_ir) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -442,13 +549,11 @@ def test_source_only_domain_does_not_stack_on_prefiltered_source( ) lineitem_ir = dataframe_scan(root, "l_partkey") - lineitem_semis = [ - semi - for semi in find_joins(optimized, "Semi") - if semi.children[0] is lineitem_ir + lineitem_hints = [ + hint for hint in find_hints(optimized) if hint.children[0] is lineitem_ir ] - assert any(join_key_names(semi) == ("l_partkey",) for semi in lineitem_semis) - assert not any(join_key_names(semi) == ("l_orderkey",) for semi in lineitem_semis) + assert any(hint_key_names(hint) == ("l_partkey",) for hint in lineitem_hints) + assert not any(hint_key_names(hint) == ("l_orderkey",) for hint in lineitem_hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -496,13 +601,13 @@ def test_derived_selectivity_propagates_through_rewritten_children( ConfigOptions.from_polars_engine(engine), ) - semis = find_joins(optimized, "Semi") + hints = find_hints(optimized) 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 expected_targets <= {hint.children[0] for hint in hints} assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -552,13 +657,10 @@ def test_rewritten_domain_filters_other_side_instead_of_stacking( 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 not 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 - ) + hints = find_hints(optimized) + assert sum(hint.children[0] is lineitem_ir for hint in hints) == 1 + assert not any(hint.children[0] is orders_ir for hint in hints) + assert not any(isinstance(hint.children[0], PushdownFilterHint) for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -612,9 +714,9 @@ def test_target_source_follows_join_key_through_rename( 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) + hints = find_hints(optimized) + assert any(hint.children[0] is small_ir for hint in hints) + assert not any(hint.children[0] is big_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -661,14 +763,11 @@ def test_domain_source_follows_join_key_through_rename( 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] + hint = next(hint for hint in find_hints(optimized) if hint.children[0] is target_ir) + selected_domain = hint.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 isinstance(rewritten_domain_source, PushdownFilterHint) 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) @@ -721,7 +820,7 @@ def test_composite_domain_columns_do_not_reconverge_after_join( facts = analyze_plan(joined, StatsCollector()) producer = _smallest_node_containing_all(joined, ("value", "value_right"), facts) - candidates = tuple(semijoin_pushdown_candidates(facts, joined, "value")) + candidates = tuple(filter_hint_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:]) @@ -802,7 +901,7 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: 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")) == ( + assert tuple(filter_hint_pushdown_candidates(facts, sliced, "target_key")) == ( (ColumnRef(sliced, "target_key"), ()), ) @@ -812,9 +911,9 @@ def test_target_prefilter_does_not_move_below_slice(engine: SPMDEngine) -> None: 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) + hints = find_hints(optimized) + assert any(hint.children[0] is sliced for hint in hints) + assert not any(hint.children[0] is target_ir for hint in hints) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -861,10 +960,10 @@ def test_target_replacement_does_not_rewrite_shared_domain_side( 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 shared_ir - assert not find_joins(unfiltered_domain, "Semi") + hints = find_hints(filtered) + assert len(hints) == 1 + assert hints[0].children[0] is shared_ir + assert not find_hints(unfiltered_domain) assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -913,12 +1012,12 @@ def test_target_prefilter_rewrites_only_selected_self_join_edge( 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") + filtered_hints = find_hints(filtered) + assert len(filtered_hints) == 1 + assert not find_hints(unfiltered) # The shared node is a valid insertion point, but its children are not: - # Only this consumer should be wrapped by the semi-join. - assert filtered_semis[0].children[0] is source_ir + # Only this consumer should be wrapped by the filter hint. + assert filtered_hints[0].children[0] is source_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -967,8 +1066,8 @@ def test_internal_prefilter_rewrites_shared_subplan_once( rewritten_left, rewritten_right = optimized.children assert rewritten_left is rewritten_right assert rewritten_left is not original_shared - (internal_semi,) = find_joins(rewritten_left, "Semi") - assert internal_semi.children[0] is target_ir + (internal_hint,) = find_hints(rewritten_left) + assert internal_hint.children[0] is target_ir assert_gpu_result_equal(query, engine=engine, check_row_order=False) @@ -1006,5 +1105,5 @@ def test_no_filter_pushdown_for_unsupported_joins( ) assert optimized is root - assert not find_joins(optimized, "Semi") + assert not find_hints(optimized) 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 781b41668161..a2f573fb33ac 100644 --- a/python/cudf_polars/tests/streaming/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -418,10 +418,6 @@ def test_from_argparse_omitted_flag_still_picks_up_env_var( # --------------------------------------------------------------------------- -def test_to_dict_empty_when_all_unspecified() -> None: - assert StreamingOptions().to_dict() == {} - - def test_to_dict_contains_only_set_fields() -> None: opts = StreamingOptions( fallback_mode="silent", diff --git a/python/cudf_polars/tests/streaming/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py index 6eb99da2b0bd..99ec3b52af55 100644 --- a/python/cudf_polars/tests/streaming/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -23,6 +23,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.streaming.actor_graph.io import Lineariser from cudf_polars.streaming.actor_graph.tracing import ActorTracer, send_chunk +from cudf_polars.utils.versions import POLARS_VERSION_LT_138 if TYPE_CHECKING: import pathlib @@ -257,6 +258,407 @@ def test_io_tasks_wait_for_memory_admission( assert second["admitted"] >= first["stop"] +@pytest.mark.parametrize( + "ordered,broadcast_limit,bloom_filter_max_size,join_strategy,method,reason,domain_rows,output_rows", + [ + (False, 1, 32 * 1024 * 1024, "shuffle", "bloom", "bloom_fits", 1, 10), + (False, 64, 0, "shuffle", "broadcast_semi_join", "exact_domain_fits", 1, 10), + ( + False, + 1_000_000, + 32 * 1024 * 1024, + "broadcast_left", + "skip", + "target_not_redistributed", + 1, + None, + ), + ( + True, + 1, + 32 * 1024 * 1024, + "ordered_aligned", + "skip", + "target_not_redistributed", + None, + None, + ), + ], + ids=["bloom", "exact", "broadcast-skip", "ordered-skip"], +) +def test_local_join_prefilter_trace_records_decision_and_effect( + request: pytest.FixtureRequest, + tmp_path: pathlib.Path, + timeout_seconds: int, + ordered: bool, # noqa: FBT001 + broadcast_limit: int, + bloom_filter_max_size: int, + join_strategy: str, + method: str, + reason: str, + domain_rows: int | None, + output_rows: int | None, +) -> None: + """Trace a direct-input join prefilter selected through the public engine.""" + pytest.importorskip("structlog") + if ordered and POLARS_VERSION_LT_138: + request.applymarker( + pytest.mark.xfail(reason="set_sorted lowers to unsupported hint ir") + ) + + domain_path = tmp_path / "domain.parquet" + target_path = tmp_path / "target.parquet" + pl.DataFrame( + { + "key": range(100), + "active": [i % 10 == 0 for i in range(100)], + } + ).write_parquet(domain_path) + pl.DataFrame( + { + "key": range(1_000), + "value": range(1_000), + } + ).write_parquet(target_path) + code = textwrap.dedent(f"""\ + import json + import os + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + ordered = {ordered!r} + if ordered: + domain = ( + pl.scan_parquet({str(domain_path)!r}) + .filter("active") + .select("key") + .set_sorted("key") + ) + target = pl.scan_parquet({str(target_path)!r}).set_sorted("key") + else: + domain = ( + pl.LazyFrame({{"key": [1, 99], "active": [True, False]}}) + .filter("active") + .select("key") + ) + target = pl.LazyFrame( + {{"key": [i % 100 for i in range(1_000)], "value": range(1_000)}} + ) + query = domain.join(target, on="key") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 1 << 30 if ordered else 64, + "max_rows_per_partition": 1_000_000 if ordered else 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" and "join_prefilters" in log + ) + record = {{ + "result_rows": result.height, + "join_strategy": event["decision"], + "prefilter": event["join_prefilters"][0], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 10 + assert record["join_strategy"] == join_strategy + expected_prefilter: dict[str, str | int] = { + "target_side": "right", + "domain_side": "left", + "method": method, + "reason": reason, + } + if domain_rows is not None: + expected_prefilter["domain_rows"] = domain_rows + assert record["prefilter"].items() >= expected_prefilter.items() + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 1 + assert record["prefilter"]["input_rows"] == 1_000 + assert record["prefilter"]["output_rows"] == output_rows + + +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,method,reason,output_rows", + [ + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 20), + (64, 0, "broadcast_semi_join", "exact_domain_fits", 20), + (1, 0, "skip", "no_viable_filter", None), + ], + ids=["bloom", "exact", "skip"], +) +def test_standalone_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + method: str, + reason: str, + output_rows: int | None, +) -> None: + """Trace a prefilter pushed below an intervening join.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + domain = ( + pl.LazyFrame( + {{"p_partkey": range(10), "active": [True] * 2 + [False] * 8}} + ) + .filter("active") + .select("p_partkey") + ) + target = ( + pl.LazyFrame( + {{ + "l_partkey": [i % 10 for i in range(100)], + "bridge_key": range(100), + "value": range(100), + }} + ) + .join(pl.LazyFrame({{"bridge_key": range(100)}}), on="bridge_key") + .with_columns((pl.col("value") + 1).alias("derived")) + ) + query = domain.join(target, left_on="p_partkey", right_on="l_partkey") + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 10, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and log.get("prefilter", {{}}).get("placement") == "standalone" + ) + record = {{ + "result_rows": result.height, + "decision": event["decision"], + "prefilter": event["prefilter"], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 20 + assert record["decision"] == method + assert ( + record["prefilter"].items() + >= { + "placement": "standalone", + "method": method, + "reason": reason, + "domain_rows": 2, + }.items() + ) + if output_rows is None: + assert "input_rows" not in record["prefilter"] + assert "output_rows" not in record["prefilter"] + else: + assert record["prefilter"]["estimated_cardinality"] == 2 + assert record["prefilter"]["input_rows"] == 100 + assert record["prefilter"]["output_rows"] == output_rows + + +@pytest.mark.parametrize( + "broadcast_limit,bloom_filter_max_size,method,reason,domain_rows", + [ + (1, 32 * 1024 * 1024, "bloom", "bloom_fits", 15), + (512, 0, "broadcast_semi_join", "exact_domain_fits", 15), + ( + 1_000_000, + 32 * 1024 * 1024, + "bloom", + "bloom_fits", + 15, + ), + ], + ids=["bloom", "exact", "bloom_despite_intervening_broadcast"], +) +def test_indirect_prefilter_trace_records_decision_and_effect( + timeout_seconds: int, + broadcast_limit: int, + bloom_filter_max_size: int, + method: str, + reason: str, + domain_rows: int, +) -> None: + """Trace a composite prefilter pushed below an intervening join.""" + pytest.importorskip("structlog") + code = textwrap.dedent(f"""\ + import json + + import polars as pl + import rmm + import structlog + + rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) + + from cudf_polars.engine.spmd import SPMDEngine + + 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 % 60 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"), + ) + ) + options = {{ + "join_filter_pushdown": {{ + "threshold": 0.5, + "bloom_filter_max_size": {bloom_filter_max_size}, + }}, + "broadcast_limit": {broadcast_limit}, + "target_partition_size": 64, + "max_rows_per_partition": 100, + }} + with SPMDEngine(executor_options=options) as engine: + with structlog.testing.capture_logs() as logs: + result = query.collect(engine=engine) + + (event,) = ( + log + for log in logs + if log.get("scope") == "actor" + and log.get("prefilter", {{}}).get("placement") == "standalone" + and log.get("prefilter", {{}}).get("target_on") == ["l_suppkey"] + ) + record = {{ + "result_rows": result.height, + "prefilter": event["prefilter"], + }} + print("PREFILTER_TRACE=" + json.dumps(record)) + """) + + env = os.environ.copy() + env["CUDF_POLARS_LOG_TRACES"] = "1" + result = subprocess.check_output( + [sys.executable, "-c", code], + env=env, + stderr=subprocess.STDOUT, + timeout=timeout_seconds, + ) + (payload,) = ( + line.removeprefix(b"PREFILTER_TRACE=") + for line in result.splitlines() + if line.startswith(b"PREFILTER_TRACE=") + ) + record = json.loads(payload) + + assert record["result_rows"] == 45 + assert record["prefilter"]["target_on"] == ["l_suppkey"] + assert ( + record["prefilter"].items() + >= { + "placement": "standalone", + "method": method, + "reason": reason, + "domain_rows": domain_rows, + }.items() + ) + assert record["prefilter"]["estimated_cardinality"] == domain_rows + assert record["prefilter"]["input_rows"] == 180 + if method == "broadcast_semi_join": + assert record["prefilter"]["output_rows"] == 45 + else: + assert 45 <= record["prefilter"]["output_rows"] < 180 + + def test_structlog_disabled_by_default(timeout_seconds: int): """Test that structlog does NOT emit events when CUDF_POLARS_LOG_TRACES is not set.""" pytest.importorskip("structlog") diff --git a/python/cudf_polars/tests/test_config.py b/python/cudf_polars/tests/test_config.py index 4f0a953dd6cc..15108a42d074 100644 --- a/python/cudf_polars/tests/test_config.py +++ b/python/cudf_polars/tests/test_config.py @@ -856,10 +856,15 @@ def test_join_filter_pushdown_options_from_env( monkeypatch.setenv( "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__THRESHOLD", "0.125" ) + monkeypatch.setenv( + "CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN__BLOOM_FILTER_MAX_SIZE", + "1024", + ) 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.bloom_filter_max_size == 1024 assert config.executor.join_filter_pushdown.trace @@ -894,6 +899,24 @@ def test_validate_join_filter_pushdown_options() -> None: executor_options={"join_filter_pushdown": {"trace": "bad"}}, ) ) + with pytest.raises(TypeError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": "bad"} + }, + ) + ) + with pytest.raises(ValueError, match="bloom_filter_max_size must be"): + ConfigOptions.from_polars_engine( + pl.GPUEngine( + executor="streaming", + executor_options={ + "join_filter_pushdown": {"bloom_filter_max_size": -1} + }, + ) + ) def test_validate_join_filter_pushdown_type() -> None: @@ -910,7 +933,9 @@ def test_validate_join_filter_pushdown_type() -> None: def test_join_filter_pushdown_from_instance() -> None: - options = JoinFilterPushdownOptions(threshold=0.25, trace=True) + options = JoinFilterPushdownOptions( + threshold=0.25, bloom_filter_max_size=1024, trace=True + ) config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", @@ -920,7 +945,10 @@ def test_join_filter_pushdown_from_instance() -> None: assert config.executor.join_filter_pushdown is options -def test_join_filter_pushdown_disabled_from_options() -> None: +def test_join_filter_pushdown_disabled_from_options( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("CUDF_POLARS__EXECUTOR__JOIN_FILTER_PUSHDOWN", "1") config = ConfigOptions.from_polars_engine( pl.GPUEngine( executor="streaming", diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index f7cf81f83999..1e44a6eb6bd5 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -36,7 +36,7 @@ test = [ [project.urls] Homepage = "https://github.com/NVIDIA/cudf" -Documentation = "https://docs.rapids.ai/api/cudf/stable/" +Documentation = "https://docs.nvidia.com/cudf/" [tool.ruff] extend = "../../pyproject.toml" diff --git a/python/dask_cudf/README.md b/python/dask_cudf/README.md index d0d4eee7db62..2cdce1b97096 100644 --- a/python/dask_cudf/README.md +++ b/python/dask_cudf/README.md @@ -1,13 +1,13 @@ -#
 Dask cuDF - A GPU Backend for Dask DataFrame
+# Dask cuDF - A GPU Backend for Dask DataFrame Dask cuDF (a.k.a. dask-cudf or `dask_cudf`) is an extension library for [Dask DataFrame](https://docs.dask.org/en/stable/dataframe.html) that provides a Pandas-like API for parallel and larger-than-memory DataFrame computing on GPUs. When installed, Dask cuDF is automatically registered as the `"cudf"` [dataframe backend](https://docs.dask.org/en/stable/how-to/selecting-the-collection-backend.html) for Dask DataFrame. > [!IMPORTANT] -> Dask cuDF does not provide support for multi-GPU or multi-node execution on its own. You must also deploy a distributed cluster (ideally with [Dask-CUDA](https://docs.rapids.ai/api/dask-cuda/stable/)) to leverage multiple GPUs efficiently. +> Dask cuDF does not provide support for multi-GPU or multi-node execution on its own. You must also deploy a distributed cluster (ideally with [Dask-CUDA](https://docs.nvidia.com/dask-cuda/)) to leverage multiple GPUs efficiently. ## Using Dask cuDF -Please visit [the official documentation page](https://docs.rapids.ai/api/dask-cudf/stable/) for detailed information about using Dask cuDF. +Please visit [the official documentation page](https://docs.nvidia.com/dask-cudf/) for detailed information about using Dask cuDF. ## Installation @@ -15,11 +15,11 @@ See the [RAPIDS install page](https://docs.rapids.ai/install/) for the most up-t ## Resources -- [Dask cuDF documentation](https://docs.rapids.ai/api/dask-cudf/stable/) -- [Best practices](https://docs.rapids.ai/api/dask-cudf/stable/best_practices/) -- [cuDF documentation](https://docs.rapids.ai/api/cudf/stable/) -- [10 Minutes to cuDF and Dask cuDF](https://docs.rapids.ai/api/cudf/latest/user_guide/10min/) -- [Dask-CUDA documentation](https://docs.rapids.ai/api/dask-cuda/stable/) +- [Dask cuDF documentation](https://docs.nvidia.com/dask-cudf) +- [Best practices](https://docs.nvidia.com/dask-cudf/latest/best_practices/) +- [cuDF documentation](https://docs.nvidia.com/cudf/) +- [10 Minutes to cuDF and Dask cuDF](https://docs.nvidia.com/cudf/latest/cudf/10min/) +- [Dask-CUDA documentation](https://docs.nvidia.com/dask-cuda/) - [Deployment](https://docs.rapids.ai/deployment/stable/) - [RAPIDS Community](https://rapids.ai/learn-more/#get-involved): Get help, contribute, and collaborate. @@ -59,6 +59,6 @@ if __name__ == "__main__": query.head() ``` -If you do not have multiple GPUs available, using `LocalCUDACluster` is optional. However, it is still a good idea to [enable cuDF spilling](https://docs.rapids.ai/api/cudf/stable/cudf/developer_guide/library_design/#spilling-to-host-memory). +If you do not have multiple GPUs available, using `LocalCUDACluster` is optional. However, it is still a good idea to [enable cuDF spilling](https://docs.nvidia.com/cudf/latest/cudf/developer_guide/library_design/#spilling-to-host-memory). If you wish to scale across multiple nodes, you will need to use a different mechanism to deploy your Dask-CUDA workers. Please see [the RAPIDS deployment documentation](https://docs.rapids.ai/deployment/stable/) for more instructions. diff --git a/python/pylibcudf/pyproject.toml b/python/pylibcudf/pyproject.toml index 695e106b7066..ab0b2fd97c02 100644 --- a/python/pylibcudf/pyproject.toml +++ b/python/pylibcudf/pyproject.toml @@ -59,7 +59,7 @@ numpy = [ [project.urls] Homepage = "https://github.com/NVIDIA/cudf" -Documentation = "https://docs.rapids.ai/api/cudf/stable/" +Documentation = "https://docs.nvidia.com/cudf/" [tool.ruff] extend = "../../pyproject.toml"