Replace rolling.apply implementation with numba-cuda-mlir - #23598
Replace rolling.apply implementation with numba-cuda-mlir#23598mroeschke wants to merge 2 commits into
Conversation
📝 WalkthroughSummary by CodeRabbit
WalkthroughRolling UDF execution now uses cached CUDA-MLIR kernels over precomputed window bounds. Legacy UDF aggregation compilation was removed. Rolling results are cast to ChangesRolling UDF migration
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@python/cudf/cudf/core/window/rolling.py`:
- Around line 393-399: Update the callable branch in the rolling aggregation
path to cast the result of jit_rolling_apply to float64 before returning it,
matching the dtype behavior applied later in the method. Add or update the
relevant assertion for integer-returning UDFs while preserving the existing
Python-float result expectations.
- Around line 393-399: Update the Rolling.apply docstring to document the
numba_cuda_mlir UDF execution path used by the callable branch and its currently
supported features. Remove outdated libcudf and PTX-specific limitations, and
explicitly state that inputs containing nulls and passing args or kwargs are
unsupported.
- Line 391: The default assignment in _apply_agg_column should use
self.window.window_size when self.min_periods is None, preserving explicit
min_periods values. Add a test for FixedForwardWindowIndexer(window_size=3)
without min_periods that verifies the final one- and two-row windows are null.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 81b6b25c-aedd-40c7-b5aa-1dbdf34da058
📒 Files selected for processing (4)
python/cudf/cudf/core/_internals/aggregation.pypython/cudf/cudf/core/udf/rolling_utils.pypython/cudf/cudf/core/window/rolling.pypython/cudf/cudf/tests/window/test_rolling.py
| # not the codes | ||
| source_column = source_column._get_decategorized_column() # type: ignore[attr-defined] | ||
|
|
||
| min_periods = 1 if self.min_periods is None else self.min_periods |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate files ---'
fd -i 'rolling|window|indexer' python/cudf tests 2>/dev/null | head -120 || true
printf '%s\n' '--- symbols and call sites ---'
rg -n --glob '*.py' --glob '*.pyx' \
'class (BaseIndexer|FixedForwardWindowIndexer)|FixedForwardWindowIndexer|min_periods|jit_rolling_apply|_window_start_end' \
python/cudf tests 2>/dev/null | head -300
printf '%s\n' '--- outline ---'
ast-grep outline python/cudf/cudf/core/window/rolling.py 2>/dev/null | head -160 || true
printf '%s\n' '--- implementation context ---'
sed -n '330,455p' python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- indexer definitions ---'
rg -l --glob '*.py' --glob '*.pyx' 'class (BaseIndexer|FixedForwardWindowIndexer)' python/cudf | while read -r f; do
echo "FILE: $f"
sed -n '1,240p' "$f"
doneRepository: rapidsai/cudf
Length of output: 19833
🏁 Script executed:
#!/bin/bash
set -eu
sed -n '330,455p' python/cudf/cudf/core/window/rolling.py
rg -n --glob '*.py' --glob '*.pyx' 'class (BaseIndexer|FixedForwardWindowIndexer)|FixedForwardWindowIndexer|min_periods|jit_rolling_apply|_window_start_end' python/cudf tests 2>/dev/null | head -300Repository: rapidsai/cudf
Length of output: 18510
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- rolling.py relevant lines ---'
nl -ba python/cudf/cudf/core/window/rolling.py | sed -n '300,455p'
printf '%s\n' '--- indexer and test locations ---'
rg -n --glob '*.py' --glob '*.pyx' --glob '*.md' \
'BaseIndexer|FixedForwardWindowIndexer|min_periods' . | head -400
printf '%s\n' '--- rolling tests ---'
fd -i 'rolling' tests python | head -100Repository: rapidsai/cudf
Length of output: 22475
🏁 Script executed:
#!/bin/bash
set -eu
nl -ba python/cudf/cudf/core/window/rolling.py | sed -n '330,455p'
rg -n --glob '*.py' --glob '*.pyx' --glob '*.md' \
'BaseIndexer|FixedForwardWindowIndexer|min_periods|jit_rolling_apply' . | head -400
fd -i 'rolling' tests python | head -100Repository: rapidsai/cudf
Length of output: 22655
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- documented default and normalization ---'
sed -n '60,95p;640,695p' python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- existing BaseIndexer tests ---'
sed -n '385,425p;475,545p;620,650p' python/cudf/cudf/tests/window/test_rolling.py
printf '%s\n' '--- UDF min_periods behavior ---'
sed -n '35,125p' python/cudf/cudf/core/udf/rolling_utils.py
printf '%s\n' '--- pandas behavioral probe ---'
python3 - <<'PY'
import pandas as pd
indexer = pd.api.indexers.FixedForwardWindowIndexer(window_size=3)
s = pd.Series([10, 20, 30, 40])
for operation in ("sum", "apply"):
rolling = s.rolling(window=indexer)
result = getattr(rolling, operation)(lambda x: x.sum() if len(x) else 0, raw=True) if operation == "apply" else rolling.sum()
print(operation, result.tolist(), result.isna().tolist())
print("rolling_min_periods", rolling.min_periods)
PYRepository: rapidsai/cudf
Length of output: 10542
🌐 Web query:
pandas Rolling min_periods BaseIndexer FixedForwardWindowIndexer default window_size source
💡 Result:
In pandas, the BaseIndexer and its subclass FixedForwardWindowIndexer have a default window_size of 0 [1][2][3]. This default value indicates that, by default, the indexer is configured to consider all observations within the window [4][2][3]. Regarding min_periods in the context of rolling operations: 1. For standard integer-based windows, min_periods defaults to the size of the window [5][6]. 2. For offset-based windows, min_periods defaults to 1 [5][6]. 3. When using a BaseIndexer subclass like FixedForwardWindowIndexer as the window argument in rolling(), the min_periods parameter is explicitly passed to the indexer's get_window_bounds method [5][1]. If you do not specify min_periods in the rolling() call, it defaults to None, and the behavior of the indexer depends on the implementation of its get_window_bounds method [5][1][7]. The source code for BaseIndexer and FixedForwardWindowIndexer is located in the pandas repository within pandas/core/indexers/objects.py [1][7]. The logic for how rolling() handles these indexers can be found in pandas/core/window/rolling.py [5].
Citations:
- 1: https://github.com/pandas-dev/pandas/blob/main/pandas/core/indexers/objects.py
- 2: https://pandas.pydata.org/docs/dev/reference/api/pandas.api.indexers.BaseIndexer.html
- 3: http://pandas.pydata.org/docs/reference/api/pandas.api.indexers.BaseIndexer.html
- 4: https://pandas.pydata.org/docs/reference/api/pandas.api.indexers.FixedForwardWindowIndexer.html
- 5: https://github.com/pandas-dev/pandas/blob/main/pandas/core/window/rolling.py
- 6: https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.rolling.html
- 7: https://github.com/pandas-dev/pandas/blob/v3.0.5/pandas/core/indexers/objects.py
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from urllib.request import urlopen
urls = [
"https://raw.githubusercontent.com/pandas-dev/pandas/main/pandas/core/window/rolling.py",
"https://raw.githubusercontent.com/pandas-dev/pandas/main/pandas/core/indexers/objects.py",
]
for url in urls:
print(f"--- {url} ---")
text = urlopen(url, timeout=15).read().decode()
lines = text.splitlines()
needles = (
"min_periods =",
"class FixedForwardWindowIndexer",
"def get_window_bounds",
"window_size",
)
for i, line in enumerate(lines):
if any(needle in line for needle in needles):
lo = max(0, i - 5)
hi = min(len(lines), i + 18)
print("\n".join(f"{j + 1}: {lines[j]}" for j in range(lo, hi)))
print()
PYRepository: rapidsai/cudf
Length of output: 2533
🏁 Script executed:
#!/bin/bash
set -eu
for spec in \
"https://raw.githubusercontent.com/pandas-dev/pandas/main/pandas/core/window/rolling.py|rolling.py" \
"https://raw.githubusercontent.com/pandas-dev/pandas/main/pandas/core/indexers/objects.py|objects.py"
do
url=${spec%%|*}
name=${spec##*|}
curl -fsSLk "$url" -o "/tmp/$name"
echo "--- $name ---"
rg -n -A16 -B6 \
'min_periods = \(|class FixedForwardWindowIndexer|def get_window_bounds|window_size' \
"/tmp/$name" | head -260
doneRepository: rapidsai/cudf
Length of output: 16300
Preserve the BaseIndexer window-size default for min_periods.
When min_periods is omitted for FixedForwardWindowIndexer(window_size=3), _apply_agg_column passes 1 instead of 3. The final one- and two-row windows therefore produce values instead of nulls. Derive the default from self.window.window_size and add a forward-indexer test without an explicit min_periods.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf/cudf/core/window/rolling.py` at line 391, The default assignment
in _apply_agg_column should use self.window.window_size when self.min_periods is
None, preserving explicit min_periods values. Add a test for
FixedForwardWindowIndexer(window_size=3) without min_periods that verifies the
final one- and two-row windows are null.
Source: Coding guidelines
| if callable(agg_name): | ||
| from cudf.core.udf.rolling_utils import jit_rolling_apply | ||
|
|
||
| start, end = self._window_start_end() | ||
| return jit_rolling_apply( | ||
| source_column, start, end, min_periods, agg_name | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Cast callable rolling results to float64.
The early return bypasses line 428. A UDF that returns an integer or float32 therefore exposes its inferred dtype, while this change requires rolling results to use float64.
Cast the jit_rolling_apply result before returning it. Add an integer-returning UDF assertion because both new tests return a Python float.
Proposed fix
- return jit_rolling_apply(
- source_column, start, end, min_periods, agg_name
- )
+ return jit_rolling_apply(
+ source_column, start, end, min_periods, agg_name
+ ).astype(np.dtype("float64"))📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if callable(agg_name): | |
| from cudf.core.udf.rolling_utils import jit_rolling_apply | |
| start, end = self._window_start_end() | |
| return jit_rolling_apply( | |
| source_column, start, end, min_periods, agg_name | |
| ) | |
| if callable(agg_name): | |
| from cudf.core.udf.rolling_utils import jit_rolling_apply | |
| start, end = self._window_start_end() | |
| return jit_rolling_apply( | |
| source_column, start, end, min_periods, agg_name | |
| ).astype(np.dtype("float64")) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf/cudf/core/window/rolling.py` around lines 393 - 399, Update the
callable branch in the rolling aggregation path to cast the result of
jit_rolling_apply to float64 before returning it, matching the dtype behavior
applied later in the method. Add or update the relevant assertion for
integer-returning UDFs while preserving the existing Python-float result
expectations.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- rolling.py outline ---'
ast-grep outline python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- relevant implementation and docstring references ---'
rg -n -C 8 'def apply|Rolling\\.apply|numba_cuda_mlir|libcudf|PTX|rolling_apply|jit_rolling_apply' python/cudf/cudf/core/window/rolling.py python/cudf/cudf/core/udf python/cudf -g '*.py' -g '*.pyx'Repository: rapidsai/cudf
Length of output: 50370
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- rolling.py outline ---'
ast-grep outline python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- relevant implementation and docstring references ---'
rg -n -C 8 'def apply|Rolling\.apply|numba_cuda_mlir|libcudf|PTX|rolling_apply|jit_rolling_apply' \
python/cudf/cudf/core/window/rolling.py \
python/cudf/cudf/core/udf \
python/cudf -g '*.py' -g '*.pyx'Repository: rapidsai/cudf
Length of output: 50370
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- Rolling.apply docstring ---'
sed -n '560,635p' python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- rolling UDF implementation ---'
wc -l python/cudf/cudf/core/udf/rolling_utils.py
ast-grep outline python/cudf/cudf/core/udf/rolling_utils.py
sed -n '1,260p' python/cudf/cudf/core/udf/rolling_utils.py
printf '%s\n' '--- UDF backend and feature documentation references ---'
rg -n -C 4 'numba_cuda_mlir|supported.*feature|unsupported|Rolling\.apply|rolling.*apply|Masked|NA' \
python/cudf/cudf/core/udf \
python/cudf/cudf/core/window/rolling.py \
docs -g '*.py' -g '*.rst' -g '*.md' -g '*.yaml' 2>/dev/null | head -n 500Repository: rapidsai/cudf
Length of output: 44606
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' '--- remainder of Rolling.apply ---'
sed -n '630,700p' python/cudf/cudf/core/window/rolling.py
printf '%s\n' '--- rolling apply tests and documentation ---'
rg -n -C 6 'rolling.*apply|apply\(.*count_if|jit_rolling_apply|cmath|math\.sin|global variables|PTX|numba_cuda_mlir' \
python/cudf tests docs -g '*.py' -g '*.rst' -g '*.md' 2>/dev/null | head -n 800
printf '%s\n' '--- repository changes touching the relevant files ---'
git diff --stat -- python/cudf/cudf/core/window/rolling.py python/cudf/cudf/core/udf/rolling_utils.py
git status --short -- python/cudf/cudf/core/window/rolling.py python/cudf/cudf/core/udf/rolling_utils.pyRepository: rapidsai/cudf
Length of output: 50369
Update the Rolling.apply docstring. Describe the numba_cuda_mlir UDF path and its current supported features. Remove stale libcudf and PTX-specific restrictions, and document that null-containing inputs and args/kwargs are unsupported.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@python/cudf/cudf/core/window/rolling.py` around lines 393 - 399, Update the
Rolling.apply docstring to document the numba_cuda_mlir UDF execution path used
by the callable branch and its currently supported features. Remove outdated
libcudf and PTX-specific limitations, and explicitly state that inputs
containing nulls and passing args or kwargs are unsupported.
Source: Coding guidelines
2 participants
Description
closes #23555
Primarily agent generated implementation of replacing the prior
rolling.apply(PTX UDF aggregation via libcudf) to a pure numba cuda mlir implementation (dedicated numba cuda kernel that jits the users UDF)For a "simple UDF" w/ a small window size e.g.
The new implementation is ~5.5x slower
For a "complex UDF" w/ a larger window size e.g.
The new implementation is about equivalent to the old implementation
cc @brandon-b-miller
Checklist