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