GH-55263: [Python][Annotations] Add internal types and core data structures by rok · Pull Request #55264 · apache/arrow

rok · 2026-10-26T11:25:15Z

Rationale for this change

This is the next step in adding type annotations to PyArrow (#32609). It merges #49103 (internal type system stubs) and #49193 (core data structures) because the two depend on each other and can't be type checked separately.

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Was AI used for this PR?

In accordance to the AI generation guidelines, please disclose below whether and how AI was used in this PR.

PR code and description written by:

Human
AI

Reviewed before submission by:

Human
AI
Not reviewed

GitHub Issue: [Python][Annotations] Add internal types and Core Data Structures #55263

Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Add type annotation stubs for array, table, tensor, builder, memory, device, config, and types modules. Includes type-ignore annotations in related tests.

Drop the TypeVar/Generic parameters from the DataType subclasses, Field, Array, ChunkedArray and the tabular classes. Where a property can only return one of a few types (dictionary index type, run end type, time units) it now returns a union or Literal instead. Most of these parameters weren't buying much yet: scalar.pyi and lib.pyi are still placeholders, so e.g. Array[Scalar[...]] collapsed to Any anyway, and several annotations were being flagged by mypy and pyright (TypeVars used only once, type arguments on subclasses that weren't actually generic). Plain classes are a lot easier to read and review, and we can bring parameters back later where they clearly pay off. This also removes the stub-only _BasicDataType and _HasPrecisionScale helpers, which have no runtime counterpart, plus the IntoArray and SingleOrList aliases.

Add pyarrow-stubs/tests with pandas-stubs style checks: check(assert_type(lib.list_(lib.int32()), lib.ListType), lib.ListType) assert_type is verified by mypy, pyright and ty, which already include pyarrow-stubs/ in their runs, and check() verifies the runtime class when the tests run under pytest, so the stubs can't silently drift from what pyarrow returns. Usage that should be rejected goes under `if TYPE_CHECKING_INVALID_USAGE:` with a `# type: ignore[...]`; unused ignores are errors in these tests, so those fail if the usage starts type checking. To have something to test, lib.pyi now re-exports the stub modules that are annotated so far, still falling back to Any for everything else. The CI type annotation script runs the tests with pytest after the type checkers.

github-actions · 2026-10-26T11:26:09Z

⚠️ GitHub issue #55263 has been automatically assigned in GitHub to PR creator.

table.pyi imports RecordBatchReader from ._ipc, which had no stub, so pyright and ty failed to resolve the import. Add a placeholder in the same style as io.pyi and scalar.pyi until _ipc gets fully annotated.

compute and interchange will be annotated in later PRs, so treat imports of them as Any instead of type checking their untyped sources (mypy) or failing to resolve them (ty). Recent ty versions also no longer honor `# type: ignore[<mypy code>]` comments, so disable the rules those comments suppress: method overrides in array.pyi, and invalid arguments in the tests, where mypy and pyright still check that the invalid usage is rejected.

- TypesEnum: list all members, with the runtime values. - is_nested() also matches run-end encoded types. - Arrays and chunked arrays have no as_py() method. - ChunkedArray.unique() returns an Array. - MapArray is a ListArray, so it has offsets and values. - value_lengths() of large list (view) arrays returns an Int64Array. - Array.from_pandas() arguments can be passed positionally. Add annotation tests for these.

Converting timedelta.max to a millisecond or second duration truncates the sub-unit part, so it doesn't round trip. The comparisons used to be missing their assert, which hid this.

An empty placeholder class rejects valid usage such as table.to_reader().read_all(), as it has no methods.

Decimal types subclass FixedSizeBinaryType, but is_fixed_size_binary() returns False for them. With TypeIs, type checkers would wrongly narrow decimal types away when the check fails.

Run ruff, mypy, pyright and ty on the stubs with pinned versions: pre-commit run --hook-stage manual python-stubs --all-files The hooks only run on request, so they don't affect the existing hooks or CI jobs.

AlenkaF · 2026-10-27T05:54:04Z

Thanks for continuing to push this forward Rok! I think this PR is a great next step and is quite clear to review.

I have checked out the PR locally and started with running type checkers (mypy, pyright and ty). mypy and ty emit no errors or warnings, but pyright does emit reportMissingImports and reportMissingModuleSource which get fixed with adding _compute.pyi placeholder (similar to existing io.pyi) and reportMissingModuleSource = false under [tool.pyright] in pyproject.toml.

The reason for these warnings and errors in pyright is that locally I have editable pyarrow install where scikit-build-core puts the .so files in a separate folder (not python/pyarrow) without __init__.py and pyright can't find them. As we do not need the compiled modules for the type checkers, I think it is fine to use reportMissingModuleSource = false.

I propose to push the changes directly if we agree the additions makes sense.

Meanwhile I will continue with running the pyarrow-stubs tests.

AlenkaF · 2026-10-27T12:56:51Z

I have run pytest pyarrow-stubs/tests, no errors found. I have looked through the tests also and tried to break some, everything looks good. I think the tests are easy to follow while still cover a lot of different angles.

I have also run mypy's stubtest with stubtest pyarrow.lib --concise --mypy-config-file pyproject.toml. All of the reported discrepancies are cases of "... is not present in stub" with majority of them being stubs that are not yet included in this PR (scalars, io, etc.). Part of these are also Type_* stubs, for example pyarrow.lib.Type_UINT8. As we already have TypesEnum defined here, might it make sense to also include Type_* stubs now? (I am not inclined towards adding them, just wondering).

Besides the reports mentioned above which will be fixed in subsequent PRs, stubtest did found one bug in pyarrow:

pyarrow.lib.Decima32Array is not present in stub

Note missing "L" in "Decimal". I will work on the fix for this now and would suggest to add it as part of this PR. Can commit directly, if that is ok.

tadeja · 2026-10-27T12:59:01Z

I've run pre-commit run --hook-stage manual python-stubs --all-files which currently uses pyarrow==26.0.0 wheel modules in site-packages/pyarrow, and got all passing (Python Stubs Lint (ruff), Python Stubs Type Check (mypy/pyright/ty)) - except for ruff-format which requires 10 files to be reformatted, fyi @rok
Python Stubs Format (ruff)...............................................Failed

Running on editable dev build 27, with pyright==1.1.414 in my conda env, I also hit errors reportMissingImports and warnings reportMissingModuleSource - Suggested changes by @AlenkaF sound good and seem to work fine for me locally.

rok · 2026-10-27T13:40:54Z

Thanks for continuing to push this forward Rok! I think this PR is a great next step and is quite clear to review.

I have checked out the PR locally and started with running type checkers (mypy, pyright and ty). mypy and ty emit no errors or warnings, but pyright does emit reportMissingImports and reportMissingModuleSource which get fixed with adding _compute.pyi placeholder (similar to existing io.pyi) and reportMissingModuleSource = false under [tool.pyright] in pyproject.toml.

The reason for these warnings and errors in pyright is that locally I have editable pyarrow install where scikit-build-core puts the .so files in a separate folder (not python/pyarrow) without __init__.py and pyright can't find them. As we do not need the compiled modules for the type checkers, I think it is fine to use reportMissingModuleSource = false.

I propose to push the changes directly if we agree the additions makes sense.

That makes sense and please go for it. I think we should aim for pyproject.toml setting to be as close to the state we want to maintain as possible. So far I've tried to exclude modules that are not covered and ignore as few things as possible. reportMissingModuleSource = false seems good, we will want to challenge it when we have complete coverage.

Part of these are also Type_* stubs, for example pyarrow.lib.Type_UINT8. As we already have TypesEnum defined here, might it make sense to also include Type_* stubs now? (I am not inclined towards adding them, just wondering).

Huh. I've never noticed these were missing. Given we're not really using them as types I'd skip them if we can.

Besides the reports mentioned above which will be fixed in subsequent PRs, stubtest did found one bug in pyarrow:
pyarrow.lib.Decima32Array is not present in stub
Note missing "L" in "Decimal". I will work on the fix for this now and would suggest to add it as part of this PR. Can commit directly, if that is ok.

Huh. Let's do the fix as separate PR, but also include a fix here? (And we'll just rebase over the actual fix on main)

I've run pre-commit run --hook-stage manual python-stubs --all-files which currently uses pyarrow==26.0.0 wheel modules in site-packages/pyarrow, and got all passing (Python Stubs Lint (ruff), Python Stubs Type Check (mypy/pyright/ty)) - except for ruff-format which requires 10 files to be reformatted, fyi @rok Python Stubs Format (ruff)...............................................Failed

I'm not sure we want (can) bring in ruff as a CI dependency. However we should autoformat before we merge, so we don't do it later and change history for no good reason. So I propose one of you runs the ruff formatter and commits the change into this branch.

AlenkaF · 2026-10-27T14:07:40Z

reportMissingModuleSource = false seems good, we will want to challenge it when we have complete coverage.

I think reportMissingModuleSource will have to stay due to local editable installs. Otherwise I agree, we should ignore as few things as possible.
Will push the commit 👍

Huh. Let's do the fix as separate PR, but also include a fix here? (And we'll just rebase over the actual fix on main)

Sounds good. Will open an issue and a PR and add a fix here also.

AlenkaF · 2026-10-27T14:17:29Z

Part of these are also Type_* stubs, for example pyarrow.lib.Type_UINT8. As we already have TypesEnum defined here, might it make sense to also include Type_* stubs now? (I am not inclined towards adding them, just wondering).

Huh. I've never noticed these were missing. Given we're not really using them as types I'd skip them if we can.

Yeah, they are internal and exposed through TypesEnum which are added to the stubs here. For stubtest we can later add them to allowlist. For the type checkers we will need to work around them once lib stubs loose their def __getattr__(name: str) -> Any: .... Maybe we can revisit this topic later (maybe we see it's easier to simply add them to the stubs)

tadeja · 2026-10-27T15:37:40Z

... ruff-format which requires 10 files to be reformatted, fyi @rok Python Stubs Format (ruff)...............................................Failed

I'm not sure we want (can) bring in ruff as a CI dependency. However we should autoformat before we merge, so we don't do it later and change history for no good reason. So I propose one of you runs the ruff formatter and commits the change into this branch.

Ok, I'll do that reformatting in a bit.

AlenkaF · 2026-10-28T08:05:06Z

I have started with reviewing pyarrow tests with running mypy --check-untyped-defs --warn-unused-ignores pyarrow/tests/test_* and looking at the ignores. The stubs seem to have changed since the last time type checkers were run on pyarrow tests, example:

click for output

$ mypy --check-untyped-defs --warn-unused-ignores pyarrow/tests/test_array.py 
pyarrow/tests/test_array.py:21: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:22: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:25: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:80: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:790: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:986: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:1300: error: List item 3 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:1302: error: List item 1 has incompatible type "None"; expected "list[tuple[bytes, int]]"  [list-item]
pyarrow/tests/test_array.py:1313: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:1320: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:1324: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:1348: error: List item 0 has incompatible type "None"; expected "list[tuple[bytes, int]]"  [list-item]
pyarrow/tests/test_array.py:1348: error: List item 1 has incompatible type "None"; expected "list[tuple[bytes, int]]"  [list-item]
pyarrow/tests/test_array.py:1362: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:2513: error: Need type annotation for "buffers" (hint: "buffers: list[<type>] = ...")  [var-annotated]
pyarrow/tests/test_array.py:2539: error: "object" has no attribute "dtype"  [attr-defined]
pyarrow/tests/test_array.py:2541: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:2542: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:2543: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:2885: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:2917: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:2938: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:2939: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:3183: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:3299: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:3720: error: Incompatible types in assignment (expression has type "None", variable has type "ndarray[tuple[Any, ...], dtype[Any]]")  [assignment]
pyarrow/tests/test_array.py:3729: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:3740: error: Value of type "object" is not indexable  [index]
pyarrow/tests/test_array.py:3817: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:4004: error: Incompatible types in assignment (expression has type "MyArrayInvalid", variable has type "MyArray")  [assignment]
pyarrow/tests/test_array.py:4017: error: Incompatible types in assignment (expression has type "MyArray2", variable has type "MyArray")  [assignment]
pyarrow/tests/test_array.py:4031: error: Incompatible types in assignment (expression has type "MyArray3", variable has type "MyArray")  [assignment]
pyarrow/tests/test_array.py:4095: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:4099: error: Unused "type: ignore" comment  [unused-ignore]
pyarrow/tests/test_array.py:4315: error: List item 3 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:4319: error: List item 5 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:4324: error: List item 2 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:4324: error: List item 4 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:4395: error: List item 1 has incompatible type "None"; expected "int"  [list-item]
pyarrow/tests/test_array.py:4407: error: List item 0 has incompatible type "None"; expected "int"  [list-item]
Found 40 errors in 1 file (checked 1 source file)

Type checkers are currently configured to skip python/pyarrow folder (too much noise) so the changes are not caught. Maybe we could start with adding mypy --check-untyped-defs --warn-unused-ignores pyarrow/tests/test_* for specific files that can already be type checked.

I propose to temporarily add the mypy command explicitly naming files which bypasses the exclude setting to ci/scripts/python_test_type_annotations.sh. Will also push the updates of ignores in the tests once I go through the ones connected to this PR.

AlenkaF · 2026-10-28T08:27:50Z

Another thing: I can't run pyright on specific files inside an excluded dir. I guess I would need to be running the checks with pyright only locally (with the pyarrow exclude temporarily removed) until the pyarrow/ exclude in pyproject.toml is removed.

AlenkaF · 2026-10-28T08:53:25Z

Oh, I was too fast here. Unfortunately def __getattr__(name: str) -> Any: ... in python/pyarrow-stubs/pyarrow/__init__.pyi makes anything in form of pa.* be Any making lots of correct ignores (in with pytest.raises ... blocks) falsely reported as unused.

Would it make sense to remove # type: ignores from the pyarrow tests for now and add type checking of the tests (with the ignores) after the stubs are in?

I will leave this for discussion and start with reading the stubs and comparing them with the APIs, starting with _types.pyi first.

rok · 2026-10-28T10:34:08Z

Ignore statements not being needed is more likely due to changes to typecheckers than due to changes in stubs. Or they might be needed once we import other parts of the library. Also note camel cased (e.g. reportArgumentType) ignores are likely coming from pyright. So if mypy thinks it's a redundant ignore you might want to check pyright too.

rok · 2026-10-28T10:36:19Z

Oh, I was too fast here. Unfortunately def __getattr__(name: str) -> Any: ... in python/pyarrow-stubs/pyarrow/__init__.pyi makes anything in form of pa.* be Any making lots of correct ignores (in with pytest.raises ... blocks) falsely reported as unused.

Would it make sense to remove # type: ignores from the pyarrow tests for now and add type checking of the tests (with the ignores) after the stubs are in?

Let's rather remove def __getattr__(name: str) -> Any: ... at least for development work? Otherwise we'll just need to add ignores again.

AlenkaF · 2026-10-28T10:43:03Z

OK, agree, when reviewing pyarrow/tests we should remove def __getattr__(name: str) -> Any: ... from python/pyarrow-stubs/pyarrow/__init__.pyi so we keep the ignores and run mypy and pyright (with updated pyproject.toml) so both pass 👍

I plan to do that later and am focusing on reading the stubs currently.

tadeja · 2026-10-28T16:29:14Z

Ehhh... I've fallen into this interesting rabbit hole - There are over 1.5k mypy errors Module not callable in pyarrow/tests because of the change from upstream's zen-xu/pyarrow-stubs having subfolder __lib_pxi to current PR's flat layout without a subfolder (Past discussion rok#45 (comment) was about the naming, but not really why removing subfolder __lib_pxi)

% mypy --check-untyped-defs --warn-unused-ignores pyarrow/tests/test_convert_builtin.py 
pyarrow/tests/test_convert_builtin.py:78: error: Module not callable  [operator]
pyarrow/tests/test_convert_builtin.py:79: error: Module not callable  [operator]
...
pyarrow/tests/test_convert_builtin.py:168: error: Module not callable  [operator]
...
Found 321 errors in 1 file (checked 1 source file)

% mypy --check-untyped-defs --warn-unused-ignores pyarrow/tests/test_*.py
Found 1866 errors in 35 files (checked 44 source files)
% ... | grep -c  "Module not callable"
1562

test_convert_builtin.py:168 : pa.array(...

A stub pyarrow/array.pyi declares pyarrow.array submodule that doesn't exist at runtime because array.pxi gets compiled into lib. Then pa.array, pa.table or pa.scalar can resolve to stub modules instead of functions.
The count of mypy errors may be even higher, depending on module processing order and mypy cache. (optionally run rm -rf .mypy_cache)

So I propose we reintroduce subfolder __lib_pxi/ (or renamed subfolder _lib_pxi/ or just _lib/) in this PR (for files: array.pyi, builder.pyi, config.pyi, device.pyi, error.pyi, io.pyi, memory.pyi, scalar.pyi, table.pyi, tensor.pyi).
Locally that seems to resolve all 1.5k Module not callable errors and still passes type checkers' tests for me.

rok and others added 25 commits October 26, 2026 11:30

Add internal types and helpers

ee419b0

Update python/pyarrow-stubs/pyarrow/_stubs_typing.pyi

b09050f

Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>

post rebase changes

28a44c7

linting and formatting

3d7a713

add pyi check to pre-commit

fb50721

ruff format

fc48948

Apply suggestions from code review

9b03d17

Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>

review feedback

3874966

Decimals should be covariant

32e439d

fix editable install collisions

93d807e

review feedback

558321c

move placeholder

6a1d410

post rebase

de68c89

Apply suggestions from code review

e3105e3

Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>

Update python/pyarrow-stubs/pyarrow/_types.pyi

de7a5b6

Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>

update

5d920e9

lint

f7573ac

more covariant types

a0c8c7a

Apply suggestions from code review

a05d894

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

lint

821e914

Add type stubs for core data structures

3003457

Add type annotation stubs for array, table, tensor, builder, memory, device, config, and types modules. Includes type-ignore annotations in related tests.

temporary

8886730

post rebase

5ee863c

github-actions Bot added the awaiting committer review Awaiting committer review label Oct 26, 2026

github-actions Bot added Component: Python Component: Documentation labels Oct 26, 2026

Add placeholder stub for pyarrow._ipc

168c7c4

table.pyi imports RecordBatchReader from ._ipc, which had no stub, so pyright and ty failed to resolve the import. Add a placeholder in the same style as io.pyi and scalar.pyi until _ipc gets fully annotated.

rok added 7 commits October 26, 2026 16:31

Fix test_timedelta_overflow for sub-unit truncation

b096b3b

Converting timedelta.max to a millisecond or second duration truncates the sub-unit part, so it doesn't round trip. The comparisons used to be missing their assert, which hid this.

Keep RecordBatchReader as Any until it gets annotated

147e021

An empty placeholder class rejects valid usage such as table.to_reader().read_all(), as it has no methods.

Use TypeGuard for is_fixed_size_binary()

67aeb33

Decimal types subclass FixedSizeBinaryType, but is_fixed_size_binary() returns False for them. With TypeIs, type checkers would wrongly narrow decimal types away when the check fails.

Test indexing of arrays, chunked arrays and tables

440fbee

Add manual pre-commit hooks for linting and type checking the stubs

a72e35b

Run ruff, mypy, pyright and ty on the stubs with pinned versions: pre-commit run --hook-stage manual python-stubs --all-files The hooks only run on request, so they don't affect the existing hooks or CI jobs.

Add _compute.pyi and reportMissingModuleSource

120659e

Fix Decima32Array typo in array.pxi

f93d912

AlenkaF mentioned this pull request Oct 27, 2026

[Python] Typo in Decimal32Array definition in array.pxi #55285

Closed

AlenkaF and others added 2 commits October 27, 2026 15:58

Fix linter error

1add31b

Stubs ruff-format

c570fa9

Move to __lib/ to fix Module errors for pa.array/table/scalar

eda3425

Conversation

rok commented Oct 26, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Was AI used for this PR?

Uh oh!

github-actions Bot commented Oct 26, 2026

Uh oh!

AlenkaF commented Oct 27, 2026

Uh oh!

AlenkaF commented Oct 27, 2026

Uh oh!

tadeja commented Oct 27, 2026

Uh oh!

rok commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AlenkaF commented Oct 27, 2026

Uh oh!

AlenkaF commented Oct 27, 2026

Uh oh!

tadeja commented Oct 27, 2026

Uh oh!

AlenkaF commented Oct 28, 2026

Uh oh!

AlenkaF commented Oct 28, 2026

Uh oh!

AlenkaF commented Oct 28, 2026

Uh oh!

rok commented Oct 28, 2026

Uh oh!

rok commented Oct 28, 2026

Uh oh!

AlenkaF commented Oct 28, 2026

Uh oh!

tadeja commented Oct 28, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rok commented Oct 26, 2026 •

edited by github-actions Bot

Loading

rok commented Oct 27, 2026 •

edited

Loading