GH-55263: [Python][Annotations] Add internal types and core data structures

.pre-commit-config.yaml

@@ -136,15 +136,100 @@ repos:
       - id: flake8
         alias: python
         name: Python Lint
+        additional_dependencies:
+          - flake8-pyi
         args:
           - "--config"
           - "python/setup.cfg"
+          - "--extend-select"
+          - "Y"
+          - "--per-file-ignores"
+          - "python/pyarrow-stubs/pyarrow/*.pyi:E301,E302,E305,E701,F821"
         files: >-
           ^(c_glib|dev|python)/
+        types:
+          - file
+        types_or:
+          - python
+          - pyi
         exclude: >-
           (
           ?^python/pyarrow/vendored/|
           )
+  # Linting, formatting and type checking of the type stubs. These hooks only
+  # run on request, not on commit or in the existing CI jobs:
+  #   pre-commit run --hook-stage manual python-stubs --all-files
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.17.2
+    hooks:
+      - id: ruff-check
+        alias: python-stubs
+        name: Python Stubs Lint (ruff)
+        stages:
+          - manual
+        args:
+          - "--fix"
+        files: >-
+          ^python/pyarrow-stubs/
+        types:
+          - pyi
+      - id: ruff-format
+        alias: python-stubs
+        name: Python Stubs Format (ruff)
+        stages:
+          - manual
+        files: >-
+          ^python/pyarrow-stubs/
+        types:
+          - pyi
+  - repo: local
+    hooks:
+      - id: mypy
+        alias: python-stubs
+        name: Python Stubs Type Check (mypy)
+        stages:
+          - manual
+        language: python
+        entry: bash -c "cd python && exec mypy"
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs/|pyproject\.toml$)
+        additional_dependencies: &python-type-check-dependencies
+          - fsspec==2026.9.0
+          - mypy==2.4.0
+          - numpy==2.4.6
+          - pandas-stubs==3.0.5.260914
+          # Provides the modules that aren't annotated yet, e.g. pyarrow.compute.
+          - pyarrow==26.0.0
+          - pyright==1.1.414
+          - scipy-stubs==1.17.1.5
+          - ty==0.0.92
+          - types-cffi==2.1.0.20260827
+          - types-psutil==7.2.2.20260906
+          - types-python-dateutil==2.9.0.20261021
+          - types-requests==2.33.0.20261006
+      - id: pyright
+        alias: python-stubs
+        name: Python Stubs Type Check (pyright)
+        stages:
+          - manual
+        language: python
+        entry: pyright -p python/pyproject.toml
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs/|pyproject\.toml$)
+        additional_dependencies: *python-type-check-dependencies
+      - id: ty
+        alias: python-stubs
+        name: Python Stubs Type Check (ty)
+        stages:
+          - manual
+        language: python
+        entry: ty check --project python
+        pass_filenames: false
+        files: >-
+          ^python/(pyarrow-stubs/|pyproject\.toml$)
+        additional_dependencies: *python-type-check-dependencies
   - repo: https://github.com/MarcoGorelli/cython-lint
     rev: v0.21.1
     hooks:

ci/scripts/python_test_type_annotations.sh

@@ -36,3 +36,7 @@ cd "${pyarrow_dir}"
 mypy
 pyright --stats
 ty check --verbose --output-format concise
+
+# The type annotation tests are type checked above; also run them to verify
+# the annotated types against the installed pyarrow.
+pytest -r s pyarrow-stubs/tests

docs/source/developers/python/development.rst

@@ -114,6 +114,18 @@ Running Type Checkers
 We support multiple type checkers. Their configurations are in
 ``pyproject.toml``.
 
+The easiest way to run all of them is with ``pre-commit``, which installs
+the same versions of the type checkers and stub packages for everyone. This
+also lints and formats the stubs with ruff:
+
+.. code-block::
+
+   $ pre-commit run --hook-stage manual python-stubs --all-files
+
+These hooks only run on request, not on commit.
+
+To run a type checker directly instead:
+
 **mypy**
 
 To run mypy on the PyArrow codebase:
@@ -169,6 +181,48 @@ in downstream projects and for users' IDEs.
 Note: ``py.typed`` marker file in the ``pyarrow/`` directory indicates to type
 checkers that PyArrow supports type checking according to :pep:`561`.
 
+Testing Type Stubs
+------------------
+
+The tests in ``pyarrow-stubs/tests/`` pin down the annotated types of public
+APIs. Each check pairs :func:`typing.assert_type`, which every type checker
+verifies statically, with a runtime ``isinstance`` check:
+
+.. code-block:: python
+
+   import pyarrow.lib as lib
+
+   from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check
+
+   def test_list_type() -> None:
+       lst = lib.list_(lib.int32())
+       check(assert_type(lst, lib.ListType), lib.ListType)
+       check(assert_type(lst.value_type, lib.DataType), lib.DataType)
+
+       if TYPE_CHECKING_INVALID_USAGE:
+           lib.list_(1)  # type: ignore[call-overload]
+
+The type checkers check these files as part of the regular runs above, and
+pytest runs them against the installed PyArrow to make sure the stubs agree
+with what PyArrow actually returns:
+
+.. code-block::
+
+   $ cd arrow/python
+   $ pytest pyarrow-stubs/tests
+
+Run ``pytest`` directly and pass it the directory: ``python -m pytest`` puts
+``arrow/python`` on ``sys.path``, so the source tree would be imported instead
+of the installed PyArrow.
+
+Code under ``if TYPE_CHECKING_INVALID_USAGE:`` is never executed. It should
+fail type checking, which is asserted with a ``# type: ignore[<mypy code>]``
+comment: unused ignores are reported as errors in these tests, so the test
+fails once the usage is no longer rejected. Start new test files with a
+``# mypy: warn-unused-ignores`` comment to enable this for mypy (pyright is
+configured in ``pyproject.toml``). ty doesn't honor mypy error codes, so it
+doesn't check invalid usage.
+
 Doctest
 =======
 

python/pyarrow-stubs/pyarrow/__init__.pyi

@@ -15,12 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-"""Type stubs for PyArrow.
-
-This is a placeholder stub file.
-Complete type annotations will be added in subsequent PRs.
-"""
-
 from typing import Any
 
 # TODO(GH-48970): remove __getattr__ before release as this

python/pyarrow-stubs/pyarrow/_ipc.pyi

@@ -0,0 +1,25 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# Placeholder stub - complete annotations in future PR.
+# At runtime, these symbols are provided by the pyarrow.lib C extension.
+
+from typing import Any, TypeAlias
+
+RecordBatchReader: TypeAlias = Any
+
+def __getattr__(name: str) -> Any: ...

python/pyarrow-stubs/pyarrow/_stubs_typing.pyi

@@ -0,0 +1,117 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import datetime as dt
+
+from _typeshed import (
+    SupportsDunderGE,
+    SupportsDunderGT,
+    SupportsDunderLE,
+    SupportsDunderLT,
+)
+from collections.abc import Collection, Container, Iterator, Sequence, Sized
+from decimal import Decimal
+from typing import Any, Literal, Protocol, TypeAlias
+
+import numpy as np
+
+from numpy.typing import NDArray
+
+from pyarrow import lib
+from pyarrow.lib import ChunkedArray
+
+ArrayLike: TypeAlias = Any
+ScalarLike: TypeAlias = Any
+Order: TypeAlias = Literal["ascending", "descending"]
+JoinType: TypeAlias = Literal[
+    "left semi",
+    "right semi",
+    "left anti",
+    "right anti",
+    "inner",
+    "left outer",
+    "right outer",
+    "full outer",
+]
+Compression: TypeAlias = Literal[
+    "gzip", "bz2", "brotli", "lz4", "lz4_frame", "lz4_raw", "zstd", "snappy"
+]
+NullEncoding: TypeAlias = Literal["mask", "encode"]
+NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"]
+TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"]
+
+PyScalar: TypeAlias = (
+    bool
+    | int
+    | float
+    | Decimal
+    | str
+    | bytes
+    | dt.date
+    | dt.datetime
+    | dt.time
+    | dt.timedelta
+)
+NumpyScalar: TypeAlias = np.generic[Any]
+
+Mask: TypeAlias = (
+    Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray
+)
+Indices: TypeAlias = (
+    Sequence[int | None]
+    | NDArray[np.integer[Any]]
+    | lib.IntegerArray
+    | ChunkedArray
+)
+
+class SupportsDunderEQ(Protocol):
+    def __eq__(self, other: object, /) -> bool: ...
+
+FilterTuple: TypeAlias = (
+    tuple[str, Literal["=", "==", "!="], SupportsDunderEQ]
+    | tuple[str, Literal["<"], SupportsDunderLT[Any]]
+    | tuple[str, Literal[">"], SupportsDunderGT[Any]]
+    | tuple[str, Literal["<="], SupportsDunderLE[Any]]
+    | tuple[str, Literal[">="], SupportsDunderGE[Any]]
+    | tuple[str, Literal["in", "not in"], Collection[Any]]
+    | tuple[str, str, Any]  # Allow general str for operator to avoid type errors
+)
+
+class Buffer(Protocol): ...
+class SupportsPyBuffer(Protocol): ...
+
+class SupportsArrowStream(Protocol):
+    def __arrow_c_stream__(self, requested_schema: Any = None, /) -> Any: ...
+
+class SupportsPyArrowArray(Protocol):
+    def __arrow_array__(self, type: Any = None, /) -> Any: ...
+
+class SupportsArrowArray(Protocol):
+    def __arrow_c_array__(self, requested_schema: Any = None, /) -> Any: ...
+
+class SupportsArrowDeviceArray(Protocol):
+    def __arrow_c_device_array__(
+        self, requested_schema: Any = None, /, **kwargs: Any
+    ) -> Any: ...
+
+class SupportsArrowSchema(Protocol):
+    def __arrow_c_schema__(self) -> Any: ...
+
+class NullableCollection(Sized, Container[Any], Protocol):
+    def __iter__(self) -> Iterator[Any]: ...
+    def __len__(self) -> int: ...
+    def __contains__(self, item: Any, /) -> bool: ...