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.

AlenkaF · 2026-10-30T16:56:48Z

@tadeja @rok I have checked typed.pyi and decided to push a commit (4b97c69) after all =) Instead of suggestions I will post the reasoning behind the changes here:

NumPy is also an optional dependency (next to pandas) so I added # type: ignore[..] there too.
Locally mypy and pyright both pass with only # type: ignore[..] so I removed pyright ignores
ty doesn't honor mypy's # type: ignore[...] and # ty: ignore[unresolved-import] fails if numpy/pandas are installed, so I disabled ty unresolved-import in pyproject.toml instead (mypy and pyright still catch broken imports).
I removed noqa: F401 on the lib import and ruff didn't complain
I have changed some classmethods in stubs as they are defined as static methods in source (types.pxi).
Example: _import_from_c in DataType, Field and Schema.
I propose to use typing_extensions.CapsuleType instead of Any in __arrow_c_schema__.
Claude added a test for this case, I have reviewed and changed it a bit. typing_extensions dependency had to be added.
Note: we should leave the Protocols in _stubs_typing.pyi at Any. Polars' public pl.DataFrame and pl.Series, for example, annotate __arrow_c_stream__ -> object so narrowing the Protocol would reject them
- https://github.com/pola-rs/polars/blob/fe540d641165534f2a5e571234c3ee78d51587ee/py-polars/src/polars/dataframe/frame.py#L1587
- https://github.com/pola-rs/polars/blob/fe540d641165534f2a5e571234c3ee78d51587ee/py-polars/src/polars/series/series.py#L1632
Claude suggested to add a test for Protocols also (under the Table stubs tests, to make sure they do not become to narrow) but I think these should be added later.

KeyValueMetadata stubs do not include keys(), values() and items() stubs because they inherit them from Mapping[bytes, bytes]:
https://github.com/python/typeshed/blob/17cbf254205549ccdcf886fadf8c265e634ad596/stdlib/typing.pyi#L823-L825
but our code overrides all three as generators, not Mapping views (probably so that duplicate keys are kept):

arrow/python/pyarrow/types.pxi

Lines 2409 to 2419 in 9b9fad6

    
           def keys(self): 
        
               for i in range(self.metadata.size()): 
        
                   yield self.metadata.key(i) 
        
           def values(self): 
        
               for i in range(self.metadata.size()): 
        
                   yield self.metadata.value(i) 
        
           def items(self): 
        
               for i in range(self.metadata.size()): 
        
                   yield (self.metadata.key(i), self.metadata.value(i))

I have added overrides to the stubs so they are in sync with runtime and added ignores because of

Return type "Iterator[bytes]" of "keys" incompatible with return type "KeysView[bytes]" in supertype "typing.Mapping" [override]

I have also added a test test_key_value_metadata.
I think code overrides need a separate issue as we might want to decide on changing the overrides in types.pxi KeyValueMetadata class to return Mapping views instead of generators (and to somehow keep support of duplicate keys).

keys_sorted in map_ take bool, None is by truthiness converted to False in Cython but I think it should not be suggested as a possible input so I changed the default from None to False. Same for dictionary().
Included the use of np.typing.DTypeLike for parameter type in from_numpy_dtype https://numpy.org/devdocs/reference/typing.html#numpy.typing.DTypeLike
I have changed the Field methods that returned Self to return Field as the runtime always builds a plain Field via pyarrow_wrap_field.

Future TODO:

string_to_tzinfo, tzinfo_to_string and ensure_metadata are not in the stubs currently (not public). We should add them to the stubtest allowlist at some point or think of adding them to the stubs.

I have run

all 3 type checkers on this file with and without numpy/pandas installed
the tests
ruff linter check

AlenkaF

Adding some extra comments for things I didn't change in types.pyi but have suggestions/questions about.

AlenkaF · 2026-10-30T17:02:08Z

+class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]):
+    def __init__(
+        self,
+        __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None,


Would it make sense to put None out of the alias (_KeyValueMetadataInput) and add | None = None explicitly in all the parameters connected to metadata, so it's clear which parameters accept it?

(just a scratch note, went back to compare with: https://github.com/zen-xu/pyarrow-stubs/blob/148368c532f96f953cf39f34120da9d1043bb4dc/pyarrow-stubs/__lib_pxi/types.pyi#L1535)

AlenkaF · 2026-10-30T17:06:19Z

+    def __contains__(self, /, __key: object) -> bool: ...
+    def __getitem__(self, /, __key: str | bytes) -> bytes: ...


typeshed puts "/" at the end:
https://github.com/python/typeshed/blob/17cbf254205549ccdcf886fadf8c265e634ad596/stdlib/typing.pyi#L826
Maybe worth a change on our side or maybe removing it altogether? Not sure about this though ...

Suggested change

def __contains__(self, /, __key: object) -> bool: ...

def __getitem__(self, /, __key: str | bytes) -> bytes: ...

def __contains__(self, __key: object, /) -> bool: ...

def __getitem__(self, __key: str | bytes, /) -> bytes: ...

AlenkaF · 2026-10-30T17:13:44Z

+def int32() -> Int32Type: ...
+def int64() -> Int64Type: ...
+def uint64() -> UInt64Type: ...
+def timestamp(unit: TimeUnit | str, tz: str | None = None) -> TimestampType: ...


This comment regarding the unit type is a curiosity from my side as the tradeoff between being strict and permissive here is not super clear to me.

If we are strict we can warn users in advance that only strings from the Literal are accepted but then type checkers will error on the use of str variables and users will have to update the code with casts.
If we accept any string (and give a hint of the options) we can accept str variables but with the cost of getting runtime errors as type checkers will catch nothing.

What do others think about this?

AlenkaF · 2026-10-30T17:15:18Z

+def binary(
+    length: Literal[-1] | int = ...,
+) -> BinaryType | FixedSizeBinaryType: ...


I think binary() should have the same handling as list_() so we add overloads or drop them for both.

tadeja · 2026-10-31T17:56:53Z

I propose to use typing_extensions.CapsuleType instead of Any in __arrow_c_schema__.
Claude added a test for this case, I have reviewed and changed it a bit. typing_extensions dependency had to be added.

Ok! I've added 'typing_extensions>=4.12.0' to dependency-groups with commit 110511a as without it or with version before 'CapsuleType' in 4.12.0 this test would fail:

python -m pytest -q python/pyarrow-stubs/tests/test_types.py

python/pyarrow-stubs/tests/test_types.py:24: in <module>
from typing_extensions import CapsuleType
E ImportError: cannot import name 'CapsuleType' from 'typing_extensions'

(Edit: Dev / Lint job failure is on main, opened pr to fix it: #55353)

AlenkaF · 2026-11-02T09:17:39Z

I added two additional changes connected to __hash__ in __lib/types.pyi (4b70969):

ExtensionType and KeyValueMetadata are unhashable at runtime, but the stubs treated them as hashable.
I added __hash__: ClassVar[None] # type: ignore[assignment] to both classes, following typeshed's pattern for list, dict and set. I also added KeyValueMetadata case to test_invalid_usage.

Note: pyright checks hashability (reportUnhashable) but mypy doesn't, Claude suggested the use of Hashable so both checkers will error.
__hash__ from DataType, Field and Schema (and Schema.__sizeof__) seem redundant.
They have the same signature as object in typeshed plus the stubs don't declare __eq__. Added test_data_type_hashable for this.

tadeja · 2026-11-02T14:26:58Z

 pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil

 # Install type checkers
 pip install mypy pyright ty


scratch note: should we add "typing_extensions>=4.12.0" also here to be certain to have it? Although CI job log of Conda Python 3.12 Sphinx & Numpydoc already shows the package in the conda env where pyarrow-stubs/tests currently run OK; typing_extensions 4.16.0 pyhcf101f3_0 conda-forge

I think mypy (at least) installs typing_extensions as transitive dependency. Not sure if we need to install typing_extensions separately (yet).

… test

tadeja · 2026-11-03T13:08:00Z

A note while I'm still reviewing array.pyi. With 3ddf085:
__lib/array.pyi - removed duplicating NullableCollection, added optional import ignores (like in types.pyi already), and corrected __dlpack__ signature
tests/test_array.py a new regression test for DLPack
( Optional import ignores not yet added for _stubs_typing.pyi, __lib/tensor.pyi or __lib/table.pyi )

Type checking without NumPy, pandas, SciPy

uv run --isolated --no-project --with mypy==2.4.0 --with pyright==1.1.414 --with ty==0.0.92 bash -c 'mypy --no-incremental; pyright; ty check'

Resolved:
pyarrow-stubs/pyarrow/__lib/array.pyi:29: error: Cannot find implementation or library stub for module named "numpy" [import-not-found]
pyarrow-stubs/pyarrow/__lib/array.pyi:29: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
pyarrow-stubs/pyarrow/__lib/array.pyi:30: error: Library stubs not installed for "pandas" [import-untyped]
pyarrow-stubs/pyarrow/__lib/array.pyi:30: note: Hint: "python3 -m pip install pandas-stubs"
pyarrow-stubs/pyarrow/__lib/array.pyi:30: note: (or run "mypy --install-types" to install all missing stub packages)

Remaining:
pyarrow-stubs/pyarrow/_stubs_typing.pyi:30: error: Cannot find implementation or library stub for module named "numpy" [import-not-found]
pyarrow-stubs/pyarrow/_stubs_typing.pyi:32: error: Cannot find implementation or library stub for module named "numpy.typing" [import-not-found]
pyarrow-stubs/pyarrow/__lib/tensor.pyi:26: error: Cannot find implementation or library stub for module named "numpy" [import-not-found]
pyarrow-stubs/pyarrow/__lib/tensor.pyi:30: error: Library stubs not installed for "scipy.sparse" [import-untyped]
pyarrow-stubs/pyarrow/__lib/tensor.pyi:30: note: Hint: "python3 -m pip install scipy-stubs"
pyarrow-stubs/pyarrow/__lib/table.pyi:32: error: Cannot find implementation or library stub for module named "numpy" [import-not-found]
pyarrow-stubs/pyarrow/__lib/table.pyi:33: error: Library stubs not installed for "pandas" [import-untyped]
pyarrow-stubs/pyarrow/__lib/table.pyi:33: note: Hint: "python3 -m pip install pandas-stubs"
pyarrow-stubs/pyarrow/__lib/table.pyi:33: note: (or run "mypy --install-types" to install all missing stub packages)
pyarrow-stubs/pyarrow/__lib/table.pyi:35: error: Cannot find implementation or library stub for module named "numpy.typing" [import-not-found]
pyarrow-stubs/pyarrow/__lib/table.pyi:35: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
Found 7 errors in 3 files (checked 22 source files)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi:32:8 - error: Import "numpy" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi:33:8 - error: Import "pandas" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/table.pyi:35:6 - error: Import "numpy.typing" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi:26:8 - error: Import "numpy" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi:30:6 - error: Import "scipy.sparse" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi:30:8 - error: Import "numpy" could not be resolved (reportMissingImports)
/Users/tadeja/Documents/arrow/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi:32:6 - error: Import "numpy.typing" could not be resolved (reportMissingImports)
7 errors, 0 warnings, 0 informations
All checks passed!

AlenkaF · 2026-11-03T13:12:59Z

I am reviewing __lib/table.pyi at the moment and will add the import ignores there too 👍

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.

github-actions Bot added the awaiting changes Awaiting changes label Oct 29, 2026

Move remaining ipc/types.pyi to __lib/

fd19ff4

github-actions Bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Oct 29, 2026

tadeja added a commit to rok/arrow that referenced this pull request Oct 29, 2026

Move compat.pyi and pandas_shim.pyi to __lib/ like in apache#55264

abc4724

tadeja mentioned this pull request Oct 29, 2026

GH-49198: [Python][Annotations] Flight, CUDA and other modules #49205

Draft

Review type.pyi

4b97c69

AlenkaF reviewed Oct 30, 2026

View reviewed changes

github-actions Bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Oct 31, 2026

Add typing_extensions to dependency-groups test

110511a

github-actions Bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Oct 31, 2026

tadeja and others added 2 commits October 31, 2026 23:47

Array staticmethod _import_from_c_capsule/device and test

aee2040

Add hash stubs to ExtensionType and KeyValueMetadata

4b70969

github-actions Bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Nov 2, 2026

tadeja reviewed Nov 2, 2026

View reviewed changes

Simplified array inputs, ignore opt imports, fix DLPack signature and…

3ddf085

… test

github-actions Bot added awaiting change review Awaiting change review and removed awaiting changes Awaiting changes labels Nov 2, 2026

Change Array.equals Any to Array and add tests

aac6c4b

tadeja force-pushed the pyarrow-stubs-pr3-core-data-sans-generic branch from 8e8a811 to aac6c4b Compare November 3, 2026 18:46

github-actions Bot added awaiting changes Awaiting changes and removed awaiting change review Awaiting change review labels Nov 3, 2026

		def __contains__(self, /, __key: object) -> bool: ...
		def __getitem__(self, /, __key: str \| bytes) -> bytes: ...

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 30, 2026

Uh oh!

AlenkaF left a comment

Choose a reason for hiding this comment

Uh oh!

AlenkaF Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

tadeja Oct 31, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

AlenkaF Oct 30, 2026

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Uh oh!

AlenkaF commented Nov 2, 2026

Uh oh!

tadeja Nov 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlenkaF Nov 2, 2026

Choose a reason for hiding this comment

Uh oh!

tadeja commented Nov 3, 2026

Uh oh!

AlenkaF commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

rok commented Oct 26, 2026 •

edited by github-actions Bot

Loading

tadeja commented Oct 31, 2026 •

edited

Loading

tadeja Nov 2, 2026 •

edited

Loading