Add annotations to pylibcudf source files - #23643
Open
vyasr wants to merge 13 commits into
Open
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
github-actions
Bot
added
Python
vyasr
force-pushed
the
codex/pylibcudf-source-annotations-stubgen
branch
from
August 13, 2026 00:49
f617fb7 to
93e1691
Compare
vyasr
changed the title
vyasr
added
improvement
vyasr
marked this pull request as ready for review
August 13, 2026 00:53
mroeschke
reviewed
Aug 13, 2026
| from functools import cache | ||
| from typing import Any, TypeAlias | ||
|
|
||
| PyarrowDataType: TypeAlias = type[Any] |
Contributor
There was a problem hiding this comment.
Could we import pyarrow in a if TYPE_CHECKING block and use -> pa.DataType?
Contributor
Author
There was a problem hiding this comment.
Done in 8839ca3: this now uses if TYPE_CHECKING: import pyarrow as pa and annotates with pa.DataType.
Comment on lines
+23
to
+30
| from typing import Protocol | ||
|
|
||
|
|
||
| class HasCudaStream(Protocol): | ||
| def __cuda_stream__(self) -> object: ... | ||
|
|
||
|
|
||
| CudaStreamLike = Stream | HasCudaStream |
Contributor
There was a problem hiding this comment.
Should we consider adding this to a pylibcudf.typing module and expose CudaStreamLike (and other aliases) just in case another library wants to use this annotation in their library with pylibcudf?
Contributor
Author
There was a problem hiding this comment.
Done in 8839ca3: added pylibcudf.typing and exposed the shared aliases/protocols there.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe PR adds static type annotations across pylibcudf Cython APIs. It defines shared CUDA stream, array-interface, metadata, and recursive structure types. Existing defaults, control flow, native calls, and runtime behavior remain unchanged. ChangesPublic API typing
Estimated code review effort: 3 (Moderate) | ~25 minutes Mergeability Score: 🟡 Moderate · up to The annotation changes may pass an unnormalized stream through an unchecked cast and may narrow valid stream-like inputs in generated stubs. This could cause incorrect stream handling or reject supported calls, so owner follow-up is needed before merging. 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: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
python/pylibcudf/pylibcudf/scalar.pyx (1)
160-164: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winAlign all public
Scalarstream annotations.
scalar.pyideclares five methods withCudaStreamLike | None, butscalar.pyxstill declaresStream | None. Update the five source annotations before regenerating the stub. Regeneration should also preserveis_valid’s defaultNone.🤖 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/pylibcudf/pylibcudf/scalar.pyx` around lines 160 - 164, Update the five public Scalar method annotations in scalar.pyx, including Scalar.is_valid, from Stream | None to CudaStreamLike | None; retain None as the default value for is_valid, then regenerate scalar.pyi so its declarations remain aligned.Sources: Coding guidelines, Learnings, MCP tools
🤖 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/pylibcudf/pylibcudf/column.pyx`:
- Around line 88-104: Update the module’s __all__ declaration to include the
public typing symbols ArrayInterfaceBase, SupportsCudaArrayInterface, and
SupportsArrayInterface, preserving all existing exports.
In `@python/pylibcudf/pylibcudf/concatenate.pyx`:
- Around line 28-32: Added typing-contract tests and representative unit
benchmarks across the cohort: in
python/pylibcudf/pylibcudf/concatenate.pyx:28-32 cover correlated Column/Table
typing and concatenate benchmarking; contiguous_split.pyx:169-174 cover
ChunkedPack.create’s return type and packing; copying.pyx:458-462 and :519-523
cover slice and split element typing; datetime.pyx:58 cover generated stream
annotations and a representative benchmark; experimental/_join_streams.pyx:23-25
cover stream-like inputs and stream joining; expressions.pyx:479 cover the
Expression return contract and construction; and filling.pyx:49 cover generated
stream annotations and a representative operation. Use the existing test and
benchmark conventions, including generated-stub/runtime validation.
Apply the same fix in `@python/pylibcudf/pylibcudf/nvtext/replace.pyx` at line 36:
Representative stream-like typing and benchmark coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/utils.pyx` at line 58: Shared
coverage for stream, alias, and collection annotations.
Apply the same fix in `@python/pylibcudf/pylibcudf/aggregation.pyx` at line 431:
Generated quantile stub and representative benchmark coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/datetime.pyx` at line 58:
Generated datetime stubs and representative operation coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/filling.pyx` at line 49:
Generated filling stubs and representative operation coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyx` at
line 37: Stream-like generated signature and runtime coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/io/json.pyx` around lines 72 -
73: Recursive alias and accepted input contract coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/join.pyx` around lines 24 - 27:
Compiled import and generated stream signature coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/merge.pyx` around lines 27 -
32: Generated stream and ordering signature coverage.
Apply the same fix in `@python/pylibcudf/pylibcudf/strings/char_types.pyx` around
lines 14 - 17: Representative string API typing coverage.
In `@python/pylibcudf/pylibcudf/io/parquet.pyx`:
- Around line 689-694: Add return annotations to read_parquet in
python/pylibcudf/pylibcudf/io/parquet.pyx at lines 689-694,
ChunkedParquetWriter.from_options in python/pylibcudf/pylibcudf/io/parquet.pyx
at lines 791-793, and OrcChunkedWriter.from_options in
python/pylibcudf/pylibcudf/io/orc.pyx at lines 732-734. Match each annotation to
the corresponding declared return type in parquet.pyi or orc.pyi, preserving the
existing implementations.
In `@python/pylibcudf/pylibcudf/io/types.pyx`:
- Around line 688-695: Update the SinkInfo.__init__ sinks parameter and its
corresponding types.pyi declaration so both accept the advertised Sequence
variants, including tuples, by removing the list-only Cython annotation;
preserve the existing runtime validation and keep documentation consistent with
the shared contract.
In `@python/pylibcudf/pylibcudf/strings/translate.pyx`:
- Around line 53-56: Align the Cython runtime boundary with the public Mapping
annotation by changing the chars_table parameter from dict to object in both
public function declarations and _table_to_c_table, then ensure conversion
accepts Mapping implementations such as UserDict. Add coverage using UserDict
for both public translation functions while preserving existing dict behavior.
In `@python/pylibcudf/pylibcudf/table.pyx`:
- Line 135: Update both return paths that call from_table_view_of_arbitrary to
pass the normalized _stream value instead of the original stream argument,
ensuring objects implementing __cuda_stream__ are converted before the callee’s
unchecked cast.
In `@python/pylibcudf/pylibcudf/types.pyx`:
- Around line 29-31: Update the PyarrowDataType alias in types.pyx to represent
PyArrow DataType instances rather than class objects: use pa.DataType under
TYPE_CHECKING and provide a runtime fallback that avoids importing optional
PyArrow. Ensure the checked-in types.pyi exposes the same instance-based alias
for from_arrow parameters and to_arrow returns.
---
Outside diff comments:
In `@python/pylibcudf/pylibcudf/scalar.pyx`:
- Around line 160-164: Update the five public Scalar method annotations in
scalar.pyx, including Scalar.is_valid, from Stream | None to CudaStreamLike |
None; retain None as the default value for is_valid, then regenerate scalar.pyi
so its declarations remain aligned.
🪄 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: 11b7070b-b8d7-43c9-a793-7869d626652c
📒 Files selected for processing (91)
python/pylibcudf/pylibcudf/aggregation.pyxpython/pylibcudf/pylibcudf/binaryop.pyxpython/pylibcudf/pylibcudf/column.pyxpython/pylibcudf/pylibcudf/column_factories.pyxpython/pylibcudf/pylibcudf/concatenate.pyxpython/pylibcudf/pylibcudf/contiguous_split.pyxpython/pylibcudf/pylibcudf/copying.pyxpython/pylibcudf/pylibcudf/datetime.pyxpython/pylibcudf/pylibcudf/experimental/_join_streams.pyxpython/pylibcudf/pylibcudf/expressions.pyxpython/pylibcudf/pylibcudf/filling.pyxpython/pylibcudf/pylibcudf/gpumemoryview.pyxpython/pylibcudf/pylibcudf/groupby.pyxpython/pylibcudf/pylibcudf/hashing.pyxpython/pylibcudf/pylibcudf/interop.pyxpython/pylibcudf/pylibcudf/io/avro.pyxpython/pylibcudf/pylibcudf/io/csv.pyxpython/pylibcudf/pylibcudf/io/experimental/hybrid_scan.pyxpython/pylibcudf/pylibcudf/io/json.pyxpython/pylibcudf/pylibcudf/io/orc.pyxpython/pylibcudf/pylibcudf/io/parquet.pyxpython/pylibcudf/pylibcudf/io/parquet_metadata.pyxpython/pylibcudf/pylibcudf/io/text.pyxpython/pylibcudf/pylibcudf/io/timezone.pyxpython/pylibcudf/pylibcudf/io/types.pyxpython/pylibcudf/pylibcudf/join.pyxpython/pylibcudf/pylibcudf/json.pyxpython/pylibcudf/pylibcudf/labeling.pyxpython/pylibcudf/pylibcudf/lists.pyxpython/pylibcudf/pylibcudf/merge.pyxpython/pylibcudf/pylibcudf/null_mask.pyxpython/pylibcudf/pylibcudf/nvtext/byte_pair_encode.pyxpython/pylibcudf/pylibcudf/nvtext/deduplicate.pyxpython/pylibcudf/pylibcudf/nvtext/edit_distance.pyxpython/pylibcudf/pylibcudf/nvtext/generate_ngrams.pyxpython/pylibcudf/pylibcudf/nvtext/jaccard.pyxpython/pylibcudf/pylibcudf/nvtext/minhash.pyxpython/pylibcudf/pylibcudf/nvtext/ngrams_tokenize.pyxpython/pylibcudf/pylibcudf/nvtext/normalize.pyxpython/pylibcudf/pylibcudf/nvtext/replace.pyxpython/pylibcudf/pylibcudf/nvtext/stemmer.pyxpython/pylibcudf/pylibcudf/nvtext/tokenize.pyxpython/pylibcudf/pylibcudf/nvtext/wordpiece_tokenize.pyxpython/pylibcudf/pylibcudf/partitioning.pyxpython/pylibcudf/pylibcudf/quantiles.pyxpython/pylibcudf/pylibcudf/reduce.pyxpython/pylibcudf/pylibcudf/replace.pyxpython/pylibcudf/pylibcudf/reshape.pyxpython/pylibcudf/pylibcudf/rolling.pyxpython/pylibcudf/pylibcudf/round.pyxpython/pylibcudf/pylibcudf/scalar.pyxpython/pylibcudf/pylibcudf/search.pyxpython/pylibcudf/pylibcudf/sorting.pyxpython/pylibcudf/pylibcudf/stream_compaction.pyxpython/pylibcudf/pylibcudf/strings/attributes.pyxpython/pylibcudf/pylibcudf/strings/capitalize.pyxpython/pylibcudf/pylibcudf/strings/case.pyxpython/pylibcudf/pylibcudf/strings/char_types.pyxpython/pylibcudf/pylibcudf/strings/combine.pyxpython/pylibcudf/pylibcudf/strings/contains.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_booleans.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_datetime.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_durations.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_fixed_point.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_floats.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_integers.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_ipv4.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_lists.pyxpython/pylibcudf/pylibcudf/strings/convert/convert_urls.pyxpython/pylibcudf/pylibcudf/strings/extract.pyxpython/pylibcudf/pylibcudf/strings/find.pyxpython/pylibcudf/pylibcudf/strings/find_multiple.pyxpython/pylibcudf/pylibcudf/strings/findall.pyxpython/pylibcudf/pylibcudf/strings/padding.pyxpython/pylibcudf/pylibcudf/strings/repeat.pyxpython/pylibcudf/pylibcudf/strings/replace.pyxpython/pylibcudf/pylibcudf/strings/replace_re.pyxpython/pylibcudf/pylibcudf/strings/reverse.pyxpython/pylibcudf/pylibcudf/strings/slice.pyxpython/pylibcudf/pylibcudf/strings/split/partition.pyxpython/pylibcudf/pylibcudf/strings/split/split.pyxpython/pylibcudf/pylibcudf/strings/strip.pyxpython/pylibcudf/pylibcudf/strings/translate.pyxpython/pylibcudf/pylibcudf/strings/wrap.pyxpython/pylibcudf/pylibcudf/table.pyxpython/pylibcudf/pylibcudf/table_equality.pyxpython/pylibcudf/pylibcudf/transform.pyxpython/pylibcudf/pylibcudf/transpose.pyxpython/pylibcudf/pylibcudf/types.pyxpython/pylibcudf/pylibcudf/unary.pyxpython/pylibcudf/pylibcudf/utils.pyx
| obj: ArrowLike, | ||
| dtype: DataType | None = None, | ||
| object stream=None, | ||
| object stream: CudaStreamLike | None = None, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate file structure ---'
ast-grep outline python/pylibcudf/pylibcudf/table.pyx
printf '%s\n' '--- relevant implementation ---'
sed -n '120,245p' python/pylibcudf/pylibcudf/table.pyx
sed -n '300,345p' python/pylibcudf/pylibcudf/table.pyx
printf '%s\n' '--- stream definitions and call sites ---'
rg -n -C 4 'CudaStreamLike|from_table_view_of_arbitrary|def from_arrow|_stream\s*=' python/pylibcudf python -g '*.pyx' -g '*.pxd' -g '*.py' | head -n 500Repository: NVIDIA/cudf
Length of output: 41079
🏁 Script executed:
#!/bin/bash
set -eu
python3 - <<'PY'
from pathlib import Path
path = Path("python/pylibcudf/pylibcudf/table.pyx")
text = path.read_text()
lines = text.splitlines()
for i, line in enumerate(lines, 1):
if "def from_arrow" in line or "from_table_view_of_arbitrary" in line or "_stream" in line:
start = max(1, i - 8)
end = min(len(lines), i + 14)
print(f"--- {path}:{start}-{end} ---")
for n in range(start, end + 1):
print(f"{n}: {lines[n-1]}")
PY
printf '%s\n' '--- all stream-like definitions ---'
rg -n -C 5 'CudaStreamLike|class Stream|cdef class Stream|def.*stream|stream_from' python -g '*.pyx' -g '*.pxd' -g '*.py' | head -n 500Repository: NVIDIA/cudf
Length of output: 50367
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- utility files ---'
fd -t f -i 'utils' python/pylibcudf | head -n 50
printf '%s\n' '--- exact stream definitions and normalization ---'
rg -n -C 12 'CudaStreamLike|def _get_stream|cpdef _get_stream|cdef _get_stream|_get_stream\(' python/pylibcudf/pylibcudf -g '*.py' -g '*.pyx' -g '*.pxd' | head -n 300
printf '%s\n' '--- Stream declarations ---'
rg -n -C 8 'cdef class Stream|class Stream|cdef Stream|from .*Stream' python/pylibcudf/pylibcudf -g '*.pyx' -g '*.pxd' -g '*.py' | head -n 300Repository: NVIDIA/cudf
Length of output: 41629
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- focused diff ---'
git diff --unified=12 -- python/pylibcudf/pylibcudf/table.pyx | sed -n '1,220p'
printf '%s\n' '--- direct arbitrary-owner stream conversions ---'
rg -n -C 5 'cdef Stream _stream = <Stream>stream|from_column_view_of_arbitrary' python/pylibcudf/pylibcudf -g '*.pyx' -g '*.pxd' | head -n 250
printf '%s\n' '--- read-only structural verifier ---'
python3 - <<'PY'
from pathlib import Path
text = Path("python/pylibcudf/pylibcudf/table.pyx").read_text()
required = [
"cdef Stream _stream = _get_stream(stream)",
"cdef Stream _stream = <Stream>stream",
]
assert required[0] in text
assert required[1] in text
returns = text[text.index("def from_arrow("):text.index("cdef table_view view(self):")]
calls = returns.count("Table.from_table_view_of_arbitrary(")
raw_streams = returns.count(" stream,")
normalized_streams = returns.count(" _stream,")
print({
"from_arrow_arbitrary_table_calls": calls,
"raw_stream_arguments": raw_streams,
"normalized_stream_arguments": normalized_streams,
"normalizer": "_get_stream(stream)",
"callee_conversion": "<Stream>stream",
})
assert calls == 2
assert raw_streams == 2
assert normalized_streams == 0
PYRepository: NVIDIA/cudf
Length of output: 8909
Pass _stream to from_table_view_of_arbitrary.
CudaStreamLike includes objects that implement __cuda_stream__, but the callee performs an unchecked <Stream>stream cast. Both return paths currently pass the unnormalized stream instead of _stream.
🤖 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/pylibcudf/pylibcudf/table.pyx` at line 135, Update both return paths
that call from_table_view_of_arbitrary to pass the normalized _stream value
instead of the original stream argument, ensuring objects implementing
__cuda_stream__ are converted before the callee’s unchecked cast.
Source: Learnings
Contributor
Author
There was a problem hiding this comment.
This is a runtime behavior fix rather than a typing/source-annotation change, so I am leaving it out of scope for this PR. It looks reasonable as a separate follow-up if we want Table.from_arrow to accept the full CudaStreamLike contract along that path.
There was a problem hiding this comment.
@vyasr, agreed. This is a valid runtime behavior issue, but it is outside the scope of this typing-only pull request.
Would you like me to open a follow-up GitHub issue for the Table.from_arrow CudaStreamLike normalization path?
🧠 Learnings used
Learnt from: CR
Repo: NVIDIA/cudf PR: 0
File: python/REVIEW_GUIDELINES.md:0-0
Timestamp: 2026-08-12T20:45:50.605Z
Learning: Applies to python/**/pylibcudf/**/*.{pyx,pxd} : - Incorrect CUDA stream handling in Cython bindings
You are interacting with an AI system.
Contributor
Author
There was a problem hiding this comment.
No issue from this PR, thanks. Since this branch is intentionally scoped to source annotations for generated stubs, I will leave follow-up tracking for that runtime stream-normalization path to maintainers if they want it.
There was a problem hiding this comment.
@vyasr, understood. No follow-up issue will be created from this PR.
🧠 Learnings used
Learnt from: CR
Repo: NVIDIA/cudf PR: 0
File: python/REVIEW_GUIDELINES.md:0-0
Timestamp: 2026-08-12T20:45:50.605Z
Learning: Applies to python/**/pylibcudf/**/*.{pyx,pxd} : - Incorrect CUDA stream handling in Cython bindings
Learnt from: vyasr
Repo: NVIDIA/cudf PR: 23643
File: python/pylibcudf/pylibcudf/concatenate.pyx:28-32
Timestamp: 2026-08-13T01:27:58.502Z
Learning: For source-annotation-only pull requests affecting Python-facing pylibcudf Cython files, validate changes with `build-pylibcudf-python`, focused `stubgen-pyx` generation, runtime import and type-hint checks for the public typing module, and relevant pre-commit hooks. Do not require generated-stub tests or benchmarks unless the pull request changes runtime behavior.
You are interacting with an AI system.
wence-
approved these changes
Aug 13, 2026
| ) | ||
|
|
||
| def all_row_groups(self, ParquetReaderOptions options): | ||
| def all_row_groups(self, ParquetReaderOptions options) -> list[int]: |
Contributor
There was a problem hiding this comment.
question: Some of these changes don't introduce a python-like type-annotation for method arguments. Is that because it is not necessary because the type stub generation already knows how to turn ParquetReaderOptions options into options: ParquetReaderOptions?
madsbk
approved these changes
Aug 13, 2026
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Currently pylibcudf maintains handwritten
.pyiannotations. We want to generate these annotations from source files. In order to do so, the source files must contain all annotation information that cannot be inferred from the Cython typing. Currently all of that information is only maintained in the handwritten.pyifiles. This PR migrates all of that typing information to annotations in the.pyxsource files. By doing so, we can generate fully descriptive type stubs.Some specific changes in this PR:
def/property return annotations where Cython accepts them.