diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 7c3607bd1df3..8c6701a90a88 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -67,6 +67,12 @@ jobs: - name: Run pre-commit run: | pre-commit run --all-files --color=always --show-diff-on-failure + - name: Run pre-commit for type stubs + env: + # The previous step already ran the ruff hooks. + SKIP: ruff-check,ruff-format + run: | + pre-commit run --hook-stage manual python-stubs --all-files --color=always --show-diff-on-failure release: name: Source Release and Merge Script on ${{ matrix.runs-on }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 50293d291f4d..f70cfbf73923 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -145,6 +145,70 @@ repos: ( ?^python/pyarrow/vendored/| ) + # Linting, formatting and type checking of the type stubs. The type checkers + # don't run on commit, only on request (the Lint CI job runs them too): + # pre-commit run --hook-stage manual python-stubs --all-files + # The type checkers run with uvx, in the environment of + # python/requirements-stubs.txt. They don't need PyArrow installed: the + # modules that aren't annotated yet, e.g. pyarrow.compute, are treated as Any + # (see python/pyproject.toml). + - 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) + args: + - "--fix" + files: >- + ^python/pyarrow-stubs/ + types: + - pyi + - id: ruff-format + alias: python-stubs + name: Python Stubs Format (ruff) + 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 uvx --with-requirements requirements-stubs.txt mypy" + pass_filenames: false + files: >- + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: &uv + - uv==0.13.7 + - id: pyright + alias: python-stubs + name: Python Stubs Type Check (pyright) + stages: + - manual + language: python + entry: uvx --with-requirements python/requirements-stubs.txt pyright -p python/pyproject.toml + pass_filenames: false + files: >- + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: *uv + - id: ty + alias: python-stubs + name: Python Stubs Type Check (ty) + stages: + - manual + language: python + # Unset the VIRTUAL_ENV of the hook, so that ty uses the environment + # of uvx. + entry: env -u VIRTUAL_ENV uvx --with-requirements python/requirements-stubs.txt ty check --project python + pass_filenames: false + files: >- + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: *uv - repo: https://github.com/MarcoGorelli/cython-lint rev: v0.21.1 hooks: diff --git a/ci/scripts/python_test_type_annotations.sh b/ci/scripts/python_test_type_annotations.sh index 092bedf3f5ea..107c5ccf4f61 100755 --- a/ci/scripts/python_test_type_annotations.sh +++ b/ci/scripts/python_test_type_annotations.sh @@ -25,14 +25,20 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then . "${ARROW_PYTHON_VENV}/bin/activate" fi -# Install library stubs. Note some libraries contain their own type hints so they need to be installed. -pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil +cd "${pyarrow_dir}" -# Install type checkers -pip install mypy pyright ty +# Install type checkers and library stubs. Note some libraries contain their +# own type hints so they need to be installed. +pip install -r requirements-stubs.txt # Run type checkers -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 expected annotated types against the runtime types. +pytest -r s pyarrow-stubs-tests + +# Check that the stubs match the signatures of the installed PyArrow. +python scripts/run_stubtest.py diff --git a/ci/scripts/python_wheel_validate_contents.py b/ci/scripts/python_wheel_validate_contents.py index 36956cf0b294..f11b24c11e0a 100644 --- a/ci/scripts/python_wheel_validate_contents.py +++ b/ci/scripts/python_wheel_validate_contents.py @@ -60,7 +60,7 @@ def validate_wheel(path): ), "pyarrow/py.typed is present in the wheel." source_root = Path(__file__).resolve().parents[2] - stubs_dir = source_root / "python" / "pyarrow-stubs" / "pyarrow" + stubs_dir = source_root / "python" / "pyarrow-stubs" assert stubs_dir.exists(), f"Stub source directory not found: {stubs_dir}" expected_stub_files = { @@ -75,7 +75,7 @@ def validate_wheel(path): } assert not (wheel_stub_files == expected_stub_files), ( - "Wheel .pyi files do not differ from python/pyarrow-stubs/pyarrow.\n" + "Wheel .pyi files do not differ from python/pyarrow-stubs.\n" f"Missing in wheel: {sorted(expected_stub_files - wheel_stub_files)}\n" f"Unexpected in wheel: {sorted(wheel_stub_files - expected_stub_files)}" ) diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index 91f98cab060b..ad46671b1cc4 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -105,8 +105,11 @@ Type Checking ============= PyArrow provides type stubs (``*.pyi`` files) for static type checking. These -stubs are located in the ``pyarrow-stubs/`` directory and are automatically -included in the distributed wheel packages. +stubs are located in the ``pyarrow-stubs/`` directory, which is laid out as a +:pep:`561` stub-only package for ``pyarrow``, so type checkers and IDEs pick +it up from the ``arrow/python`` directory. The stubs aren't included in the +distributed wheel packages until they are complete (see +`GH-49831 `_). Running Type Checkers --------------------- @@ -114,6 +117,29 @@ 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 type checkers and stub packages pinned in ``requirements-stubs.txt`` (with +uv). This also lints and formats the stubs with ruff: + +.. code-block:: + + $ pre-commit run --hook-stage manual python-stubs --all-files + +The type checkers only run on request, while ruff also runs on commit. The +Lint CI job runs all of them as well. CI also runs the type checkers, the type +stub tests and stubtest (see below) against the PyArrow it builds, with +``ci/scripts/python_test_type_annotations.sh``. + +To run a type checker directly instead, install the pinned type checkers and +stub packages in your PyArrow development environment: + +.. code-block:: + + $ pip install -r requirements-stubs.txt + +Checking the PyArrow tests also needs their dependencies from +``requirements-test.txt``. + **mypy** To run mypy on the PyArrow codebase: @@ -124,6 +150,12 @@ To run mypy on the PyArrow codebase: $ mypy The mypy configuration is in the ``[tool.mypy]`` section of ``pyproject.toml``. +mypy only uses the stubs when it checks them too, so pass ``pyarrow-stubs`` +along with any other files you check: + +.. code-block:: + + $ mypy pyarrow-stubs pyarrow/tests/test_array.py **pyright** @@ -145,6 +177,20 @@ To run ty (note: currently only partially configured): $ cd arrow/python $ ty check +**stubtest** + +mypy's stubtest compares the stubs with the signatures of the installed +PyArrow. Run it with: + +.. code-block:: + + $ cd arrow/python + $ python scripts/run_stubtest.py + +Stubs that aren't written yet are skipped. The known mismatches are listed in +``pyarrow-stubs-tests/stubtest_allowlist.txt``; remove the entries of the +mismatches you fix, as stubtest fails on unused entries. + Maintaining Type Stubs ----------------------- @@ -169,6 +215,44 @@ 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 + + from typing import TYPE_CHECKING + + import pyarrow as pa + + from _check import assert_type, check + + def test_list_type() -> None: + lst = pa.list_(pa.int32()) + check(assert_type(lst, pa.ListType), pa.ListType) + check(assert_type(lst.value_type, pa.DataType), pa.DataType) + + if TYPE_CHECKING: + pa.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 + +Code under ``if TYPE_CHECKING:`` is never executed. It holds the usage that +should fail type checking, which is asserted with a +``# type: ignore[]`` comment: unused ignores are reported as errors +in these tests, so the test fails once the usage is no longer rejected. ty +doesn't honor mypy error codes, so it doesn't check invalid usage. + Doctest ======= diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 5bff8bcc6499..28c163328ff3 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -1077,7 +1077,7 @@ endif() # # TODO(GH-32609): Reintroduce type stubs into PyArrow wheels. # GH-49831: Temporarily do not install pyarrow-stubs into wheels. -# Stubs live in pyarrow-stubs/pyarrow/ during development but are installed +# Stubs live in pyarrow-stubs/ during development but are installed # alongside the package so type checkers can find them (PEP 561). # The stubs are currently incomplete, and some type checkers consume .pyi files # even without the py.typed marker. Re-enable this when the stubs are complete. @@ -1085,7 +1085,7 @@ if(PYARROW_REQUIRE_STUB_DOCSTRINGS) message(FATAL_ERROR "PYARROW_REQUIRE_STUB_DOCSTRINGS cannot be used while " "pyarrow-stubs are omitted from wheels (GH-49831).") endif() -# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs/pyarrow") +# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs") # if(EXISTS "${PYARROW_STUBS_SOURCE_DIR}") # install(DIRECTORY "${PYARROW_STUBS_SOURCE_DIR}/" # DESTINATION "." diff --git a/python/pyarrow-stubs-tests/_check.py b/python/pyarrow-stubs-tests/_check.py new file mode 100644 index 000000000000..4ebd47608b52 --- /dev/null +++ b/python/pyarrow-stubs-tests/_check.py @@ -0,0 +1,48 @@ +# 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. + +"""Helpers for the type annotation tests in this directory. + +Each test pairs a static and a runtime check of the same expression:: + + check(assert_type(pa.list_(pa.int32()), pa.ListType), pa.ListType) + +``assert_type`` is verified by every type checker we run in CI (mypy, +pyright and ty) and is a no-op at runtime. ``check`` runs under pytest and +verifies that the value really is an instance of the given class, so the +stubs can't drift away from what pyarrow actually returns. + +Code that type checkers should *reject* goes under ``if TYPE_CHECKING:`` +with a ``# type: ignore[]`` comment. The block never runs, and +mypy and pyright are configured to report unused ignores in these tests, so +the test fails if the usage stops being an error. +""" + +from typing import Any, TypeVar, assert_type + +__all__ = ["assert_type", "check"] + +_T = TypeVar("_T") + + +def check(actual: _T, klass: type[Any] | tuple[type[Any], ...]) -> _T: + """Assert that ``actual`` is an instance of ``klass`` and return it.""" + if not isinstance(actual, klass): + raise AssertionError( + f"Expected an instance of {klass}, got {type(actual)}: {actual!r}" + ) + return actual diff --git a/python/pyarrow-stubs-tests/conftest.py b/python/pyarrow-stubs-tests/conftest.py new file mode 100644 index 000000000000..fa4936a64ede --- /dev/null +++ b/python/pyarrow-stubs-tests/conftest.py @@ -0,0 +1,37 @@ +# 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 sys +import types + +import pyarrow as pa + + +def _stub_only_type(name: str) -> type[pa.DataType]: + # The factories of the stub-only DataType subclasses (Int32Type, ...) + # return plain DataType instances at runtime. + if name.endswith("Type") and not name.startswith("_"): + return pa.DataType + raise AttributeError(name) + + +# pyarrow._stubs_typing only exists in the type stubs. Provide it at runtime, +# so that the tests can import the stub-only types without a TYPE_CHECKING +# block. +_stubs_typing = types.ModuleType("pyarrow._stubs_typing") +setattr(_stubs_typing, "__getattr__", _stub_only_type) +sys.modules["pyarrow._stubs_typing"] = _stubs_typing diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt new file mode 100644 index 000000000000..b361f3157c92 --- /dev/null +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -0,0 +1,30 @@ +# 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. + +# Mismatches between the type stubs and the runtime that stubtest skips, one +# regex per line; see scripts/run_stubtest.py. Unused entries fail stubtest, +# so remove the entries of the mismatches you fix. Print the entries for the +# current mismatches with +# python scripts/run_stubtest.py --generate-allowlist + +# Stub-only modules. The pyarrow.__lib modules are checked against +# pyarrow.lib, which has no __all__. +pyarrow\._stubs_typing +pyarrow\.__lib\.[a-z]+\.__all__ + +# The __init__ of Cython extension types takes (*args, **kwargs) at runtime. +pyarrow\.__lib\.[a-z]+\.[A-Za-z_]+\.__init__ diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py new file mode 100644 index 000000000000..57fabfa8f4e5 --- /dev/null +++ b/python/pyarrow-stubs-tests/test_array.py @@ -0,0 +1,252 @@ +# 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. + +from collections.abc import Hashable +from typing import Any, cast, TYPE_CHECKING + +import pytest +from typing_extensions import CapsuleType + +import pyarrow as pa +from pyarrow._compute import CastOptions + +from _check import assert_type, check + + +def test_array() -> None: + # pa.array() isn't annotated with a specific return type yet, so tell + # the type checker which class to expect. + arr = check(cast(pa.Int64Array, pa.array([1, 2, None])), pa.Int64Array) + + check(assert_type(arr.type, pa.DataType), pa.DataType) + check(assert_type(arr.null_count, int), int) + check(assert_type(len(arr), int), int) + check(assert_type(arr.to_pylist(), list[Any]), list) + # Scalars aren't annotated yet, so indexing returns Any | Int64Array. + check(assert_type(arr[0], Any | pa.Int64Array), pa.Int64Scalar) + check(assert_type(arr[1:], Any | pa.Int64Array), pa.Int64Array) + + # Methods returning the same kind of array keep the subclass. + check(assert_type(arr.slice(1), pa.Int64Array), pa.Int64Array) + check(assert_type(arr.fill_null(0), pa.Int64Array), pa.Int64Array) + check(assert_type(arr.drop_null(), pa.Int64Array), pa.Int64Array) + indices: pa.Array = pa.array([1, 0]) + check(assert_type(arr.take(indices), pa.Int64Array), pa.Int64Array) + + check(assert_type(arr.is_null(), pa.BooleanArray), pa.BooleanArray) + check(assert_type(arr.cast(pa.float64()), pa.Array), pa.DoubleArray) + check(assert_type(arr.cast("float64"), pa.Array), pa.DoubleArray) + check(assert_type(arr.cast(pa.int32(), safe=None), pa.Array), pa.Int32Array) + # without target_type argument, as CastOptions already specifies it + check(assert_type(arr.cast(options=CastOptions.safe(pa.int32())), pa.Array), + pa.Int32Array) + check(assert_type(arr.view("uint64"), pa.Array), pa.UInt64Array) + check(assert_type(arr.to_string(element_size_limit=5), str), str) + check(assert_type(arr.to_tensor(allow_nulls=True), pa.Tensor), pa.Tensor) + check(assert_type(pa.concat_arrays([arr, arr]), pa.Array), + pa.Int64Array) + check(assert_type(pa.nulls(2, pa.int8()), pa.Array), pa.Int8Array) + check(assert_type(pa.repeat("a", 2), pa.Array), pa.StringArray) + check(assert_type(pa.Array.from_pandas([1, 2], None, "int64"), + pa.Array), + pa.Int64Array) + + other = check(cast(pa.StringArray, pa.array(["a", "b", None])), pa.StringArray) + check(assert_type(arr.equals(arr), bool), bool) + check(assert_type(arr.equals(other), bool), bool) + # Arrays of any type can diff() + check(assert_type(arr.diff(other), str), str) + + bool_mask = cast(pa.BooleanArray, pa.array([True, False, True])) + check(assert_type(arr.filter(bool_mask), pa.Int64Array), pa.Int64Array) + + if TYPE_CHECKING: + pa.array([1, 2], mask=[True, None]) # type: ignore[list-item] + # TODO: when scalars get stubs + # pa.array([1, 2], mask=pa.chunked_array([[True, False]])) + + +def test_array_numpy() -> None: + pytest.importorskip("numpy") + import numpy as np + + arr = cast(pa.Int64Array, pa.array([1, 2])) + check(assert_type(arr[np.int64(0)], Any | pa.Int64Array), pa.Int64Scalar) + check(assert_type(pa.infer_type(np.array([1])), pa.DataType), pa.DataType) + offsets = np.array([0, 1, 2], dtype="int32") + check(assert_type(pa.ListArray.from_arrays(offsets, [1, 2]), pa.ListArray), + pa.ListArray) + + np_mask = np.array([True, False]) + check(assert_type(arr.filter(np_mask), pa.Int64Array), pa.Int64Array) + + if TYPE_CHECKING: + dtype = pa.dictionary(pa.int8(), pa.string()) + strs = np.array(["a"]) + pa.DictionaryArray.from_buffers(dtype, 1, [], strs) # type: ignore[arg-type] + + +def test_array_dlpack() -> None: + arr = cast(pa.Int64Array, pa.array([1, 2])) + capsule = arr.__dlpack__( + max_version=(1, 0), + dl_device=arr.__dlpack_device__(), + copy=False, + ) + check(assert_type(capsule, CapsuleType), CapsuleType) + + +def test_import_from_c_capsule() -> None: + arr = cast(pa.Int64Array, pa.array([1, 2])) + capsules = arr.__arrow_c_array__() + check(assert_type(capsules, tuple[CapsuleType, CapsuleType]), tuple) + imported = pa.BooleanArray._import_from_c_capsule(*capsules) + check(assert_type(imported, pa.Array), pa.Int64Array) + + check(assert_type(arr.__arrow_c_device_array__(), + tuple[CapsuleType, CapsuleType]), + tuple) + + +def test_from_buffers() -> None: + values = cast(pa.Int64Array, pa.array([1, 2])) + offsets = pa.array([0, 2], pa.int32()).buffers()[1] + # Buffers and children can be any collection, like dict values or a tuple + buffers = {"validity": None, "offsets": offsets} + list_arr = pa.Array.from_buffers(pa.list_(pa.int64()), 1, buffers.values(), + children=(values,)) + check(assert_type(list_arr, pa.Array), pa.ListArray) + + +def test_list_arrays() -> None: + lst = pa.ListArray.from_arrays([0, 1, 3], [1, 2, 3]) + check(assert_type(lst.value_lengths(), pa.Int32Array), pa.Int32Array) + + large = pa.LargeListArray.from_arrays([0, 1, 3], pa.array([1, 2, 3])) + check(assert_type(large.value_lengths(), pa.Int64Array), pa.Int64Array) + + list_type = pa.list_(pa.int64()) + check(assert_type(pa.ListArray.from_arrays([0, 2], [1, 2], list_type), + pa.ListArray), + pa.ListArray) + check(assert_type(pa.LargeListArray.from_arrays([0, 2], [1, 2]), + pa.LargeListArray), + pa.LargeListArray) + check(assert_type(pa.ListViewArray.from_arrays([0], [2], [1, 2]), + pa.ListViewArray), + pa.ListViewArray) + check(assert_type(pa.LargeListViewArray.from_arrays([0], [2], [1, 2]), + pa.LargeListViewArray), + pa.LargeListViewArray) + check(assert_type(pa.FixedSizeListArray.from_arrays([1, 2], 2), + pa.FixedSizeListArray), + pa.FixedSizeListArray) + + # MapArray is a ListArray (of key-item structs) at runtime. + map_arr = pa.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"]) + check(assert_type(map_arr, pa.MapArray), pa.ListArray) + check(assert_type(map_arr.offsets, pa.Int32Array), pa.Int32Array) + check(assert_type(map_arr.values, pa.Array), pa.StructArray) + check(assert_type(map_arr.keys, pa.Array), pa.Int64Array) + map_type = pa.map_(pa.int64(), pa.string()) + check(assert_type(pa.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"], map_type), + pa.MapArray), + pa.ListArray) + + +def test_dictionary_array() -> None: + arr = check(cast(pa.StringArray, pa.array(["a", "b", "a"])), + pa.StringArray) + dict_arr = arr.dictionary_encode() + check(assert_type(dict_arr, pa.DictionaryArray), pa.DictionaryArray) + check(assert_type(dict_arr.indices, pa.IntegerArray), pa.Int32Array) + check(assert_type(dict_arr.dictionary, pa.Array), pa.StringArray) + check(assert_type(dict_arr.dictionary_decode(), pa.Array), + pa.StringArray) + + rebuilt = pa.DictionaryArray.from_buffers( + dict_arr.type, len(dict_arr), dict_arr.indices.buffers(), dict_arr.dictionary + ) + check(assert_type(rebuilt, pa.DictionaryArray), pa.DictionaryArray) + + +def test_run_end_encoded_array() -> None: + ree = pa.RunEndEncodedArray.from_arrays([2, 3], ["a", "b"]) + check(assert_type(ree, pa.RunEndEncodedArray), pa.RunEndEncodedArray) + check( + assert_type(ree.run_ends, + pa.Int16Array | pa.Int32Array | pa.Int64Array), + pa.Int64Array, + ) + check(assert_type(ree.values, pa.Array), pa.StringArray) + + rebuilt = pa.RunEndEncodedArray.from_buffers( + ree.type, len(ree), [None], children=(ree.run_ends, ree.values) + ) + check(assert_type(rebuilt, pa.RunEndEncodedArray), pa.RunEndEncodedArray) + + +def test_chunked_array() -> None: + chunked = pa.chunked_array([[1, 2], [3]]) + check(assert_type(chunked, pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.type, pa.DataType), pa.DataType) + check(assert_type(chunked.num_chunks, int), int) + check(assert_type(chunked.chunk(0), pa.Array), pa.Int64Array) + # Scalars aren't annotated yet, so indexing returns Any | ChunkedArray. + check(assert_type(chunked[0], Any | pa.ChunkedArray), pa.Int64Scalar) + check(assert_type(chunked[1:], Any | pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.chunks, list[pa.Array]), list) + check(assert_type(chunked.combine_chunks(), pa.Array), pa.Int64Array) + check(assert_type(chunked.unique(), pa.Array), pa.Int64Array) + check(assert_type(chunked.slice(1), pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.cast(pa.float64()), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(chunked.is_null(), pa.ChunkedArray), pa.ChunkedArray) + + +def test_invalid_usage() -> None: + arr = cast(pa.Int64Array, pa.array([1, 2])) + if TYPE_CHECKING: + arr.cast(1) # type: ignore[arg-type] + arr.slice("1") # type: ignore[arg-type] + arr.equals(object()) # type: ignore[arg-type] + arr.equals(None) # type: ignore[arg-type] + arr.format() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] # noqa: E501 + # Arrays define __eq__ but not __hash__, hence unhashable at runtime + _: Hashable = arr # type: ignore[assignment] # ty: ignore[invalid-assignment] + arr.sort(null_placement="first") # type: ignore[arg-type] + arr.filter([True], null_selection_behavior="keep") # type: ignore[arg-type] + arr.dictionary_encode("drop") # type: ignore[arg-type] + arr.to_pylist(maps_as_pydicts="loose") # type: ignore[arg-type] + # Generators aren't accepted, only Sequences + pa.infer_type(x for x in [1]) # type: ignore[arg-type] + + +def test_buffer_constructors_invalid_usage() -> None: + types = cast(pa.Int8Array, pa.array([0], pa.int8())) + children: list[pa.Array] = [pa.array([1])] + strings = cast(pa.StringArray, pa.array(["a"])) + off, data = cast(list[pa.Buffer], strings.buffers()[1:]) + if TYPE_CHECKING: + # Only lists are accepted for children and type codes + pa.UnionArray.from_sparse(types, tuple(children)) # type: ignore[arg-type] + pa.UnionArray.from_sparse(types, children, None, + types) # type: ignore[arg-type] + pa.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] + # Fields are passed as `fields`, `names` only takes strings + pa.StructArray.from_arrays( + children, names=[pa.field("a", pa.int64())]) diff --git a/python/pyarrow-stubs-tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py new file mode 100644 index 000000000000..31732b991436 --- /dev/null +++ b/python/pyarrow-stubs-tests/test_table.py @@ -0,0 +1,227 @@ +# 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. + +from collections.abc import Hashable +from types import MappingProxyType +from typing import Any, TYPE_CHECKING + +import pytest +from typing_extensions import CapsuleType + +import pyarrow as pa + +from _check import assert_type, check + + +class _PyArrowArrayProvider: + def __arrow_array__(self, type: pa.DataType | None = None) -> pa.Array: + return pa.array([1, 2], type=type) + + +class _ArrowArrayProvider: + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: + return pa.array([1, 2]).__arrow_c_array__(requested_schema) + + +class _ArrowDeviceArrayProvider: + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: + return pa.array([1, 2]).__arrow_c_device_array__(requested_schema, **kwargs) + + +class _ArrowStreamProvider: + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: + return pa.chunked_array([[1, 2]]).__arrow_c_stream__(requested_schema) + + +def test_table() -> None: + table = pa.table({"a": [1, 2], "b": ["x", "y"]}) + check(assert_type(table, pa.Table), pa.Table) + check(assert_type(table.num_rows, int), int) + check(assert_type(table.num_columns, int), int) + check(assert_type(table.shape, tuple[int, int]), tuple) + check(assert_type(table.column_names, list[str]), list) + check(assert_type(table.schema, pa.Schema), pa.Schema) + check(assert_type(table.field("a"), pa.Field), pa.Field) + + # Columns aren't typed per tabular class, so Table and RecordBatch both + # return the union here. + check(assert_type(table.column("a"), pa.Array | pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(table["a"], pa.Array | pa.ChunkedArray | pa.Table), + pa.ChunkedArray) + check(assert_type(table[0], pa.Array | pa.ChunkedArray | pa.Table), + pa.ChunkedArray) + check(assert_type(table[:1], pa.Array | pa.ChunkedArray | pa.Table), + pa.Table) + + check(assert_type(table.set_column(0, "a", [_ArrowDeviceArrayProvider()]), + pa.Table), + pa.Table) + check(assert_type(table.append_column("c", ([3], [4])), pa.Table), pa.Table) + + check(assert_type(table.slice(0, 1), pa.Table), pa.Table) + check(assert_type(table.select(["a"]), pa.Table), pa.Table) + check(assert_type(table.select(["a", 1]), pa.Table), pa.Table) + check(assert_type(table.equals(None), bool), bool) + check(assert_type(table.combine_chunks(), pa.Table), pa.Table) + check(assert_type(table.to_batches(), list[pa.RecordBatch]), list) + check(assert_type(table.to_pydict(), dict[str, list]), dict) + check(assert_type(pa.Table.from_arrays([(1, 2)], names=["a"]), pa.Table), + pa.Table) + + chunked = pa.chunked_array([[1, 2]]) + check(assert_type(table.add_column(0, "c", chunked), pa.Table), pa.Table) + check(assert_type(table.add_column(0, "c", pa.array([1, 2])), pa.Table), + pa.Table) + check(assert_type(table.add_column(0, "c", [[1], [2]]), pa.Table), + pa.Table) + check(assert_type(table.add_column(0, "c", [pa.array([1, 2])]), pa.Table), + pa.Table) + + table = pa.table(([1],), names=["a"]) + check(assert_type(table, pa.Table), pa.Table) + + chunked_struct = pa.chunked_array([[{"x": 1}]]) + check(assert_type(pa.Table.from_struct_array(chunked_struct), pa.Table), + pa.Table) + + # The reader isn't annotated yet, but must not reject valid usage. + check(table.to_reader().read_all(), pa.Table) + for batch in table.to_reader(): + check(batch, pa.RecordBatch) + + +def test_table_from_arrays_protocols() -> None: + check(assert_type(pa.Table.from_arrays([_PyArrowArrayProvider()], names=["a"]), + pa.Table), + pa.Table) + check(assert_type(pa.Table.from_arrays([_ArrowArrayProvider()], names=["a"]), + pa.Table), + pa.Table) + check(assert_type(pa.Table.from_arrays([_ArrowDeviceArrayProvider()], names=["a"]), + pa.Table), + pa.Table) + + +def test_table_from_arrays_numpy() -> None: + pytest.importorskip("numpy") + import numpy as np + + columns = np.array([[1, 2], [3, 4]]) + check(assert_type(pa.Table.from_arrays(columns, names=["a", "b"]), pa.Table), + pa.Table) + + +def test_record_batch() -> None: + batch = pa.record_batch({"a": [1, 2]}) # from a dict + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) + check(assert_type(batch.num_rows, int), int) + check(assert_type(batch.column(0), pa.Array | pa.ChunkedArray), + pa.Array) + check(assert_type(batch.slice(1), pa.RecordBatch), pa.RecordBatch) + check(assert_type(pa.Table.from_batches([batch]), pa.Table), pa.Table) + check(assert_type(pa.concat_batches([batch, batch]), pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.add_column(0, "b", pa.array([3, 4])), + pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.add_column(0, "b", [3, 4]), pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.append_column("b", [3, 4]), pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.equals(batch), bool), bool) + # Any object can be compared, a non-batch returns False + check(assert_type(batch.equals(42), bool), bool) + + batch = pa.record_batch([pa.array([1])], names=["a"]) # from a list + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) + + struct_array = batch.to_struct_array() + check(assert_type(struct_array, pa.StructArray), pa.StructArray) + check(assert_type(pa.RecordBatch.from_struct_array(struct_array), + pa.RecordBatch), + pa.RecordBatch) + + +def test_record_batch_from_pandas() -> None: + pd = pytest.importorskip("pandas") + df = pd.DataFrame({"a": [1], "b": [2]}) + # Columns only need to be iterable, e.g. a generator. + batch = pa.RecordBatch.from_pandas(df, columns=(c for c in ["a"])) + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) + + +def test_chunked_array() -> None: + chunked_array = pa.chunked_array([[1, 2], [None]]) + check(assert_type(chunked_array, pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked_array.cast(pa.int32(), safe=None), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(chunked_array.fill_null(0), pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked_array.sort(null_placement="at_start"), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(chunked_array.equals(None), bool), bool) + + arr = pa.array([1, 2]) + check(assert_type(pa.chunked_array(arr), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([arr, [3]]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_PyArrowArrayProvider()]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_ArrowDeviceArrayProvider()]), + pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array(_ArrowArrayProvider()), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_ArrowArrayProvider()]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array(_ArrowStreamProvider()), pa.ChunkedArray), + pa.ChunkedArray) + + +def test_invalid_usage() -> None: + table = pa.table({"a": [1, 2]}) + streams = [_ArrowStreamProvider()] + columns = {"a": pa.array([1])}.values() # dict as non-sequence example + + if TYPE_CHECKING: + table.column(1.5) # type: ignore[arg-type] + table.slice("0") # type: ignore[arg-type] + pa.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] + # Stream providers are accepted as the whole input, not individual chunks + pa.chunked_array(streams) # type: ignore[arg-type] + table.add_column(0, "b", streams) # type: ignore[arg-type] + # Tables define __eq__ but not __hash__, hence unhashable at runtime + _: Hashable = table + + +def test_record_batch_invalid_usage() -> None: + arr = pa.array([1]) + chunked = pa.chunked_array([[{"a": 1}]]) + if TYPE_CHECKING: + # Tuples are failing at runtime. Recheck after the fix of GH-55418 + pa.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] + pa.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] + pa.record_batch((arr,), names=["a"]) # type: ignore[arg-type] + # Only a `dict` is passed on to `from_pydict`, other mappings raise + pa.record_batch(MappingProxyType({"a": [1]})) # type: ignore[arg-type] diff --git a/python/pyarrow-stubs-tests/test_tensor.py b/python/pyarrow-stubs-tests/test_tensor.py new file mode 100644 index 000000000000..917736e7b3f0 --- /dev/null +++ b/python/pyarrow-stubs-tests/test_tensor.py @@ -0,0 +1,103 @@ +# 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. +# mypy: warn-unused-ignores + +from collections.abc import Hashable +from typing import TYPE_CHECKING + +import numpy as np + +from scipy.sparse import coo_array, csc_array, csr_array +from typing_extensions import CapsuleType + +import pyarrow as pa + +from _check import assert_type, check + +DENSE = np.array([[1, 0, 2], [0, 3, 0]]) + + +def test_tensor() -> None: + tensor = pa.Tensor.from_numpy(DENSE, dim_names=("x", "y")) + check(assert_type(tensor, pa.Tensor), pa.Tensor) + check(assert_type(tensor.dim_names, list[str]), list) + + from_dlpack = pa.Tensor.from_dlpack(DENSE) + check(assert_type(from_dlpack, pa.Tensor), pa.Tensor) + check(assert_type(tensor.__dlpack__(max_version=(1, 0)), CapsuleType), + CapsuleType) + check(assert_type(tensor.__dlpack_device__(), tuple[int, int]), tuple) + # __dlpack_device__() output must be accepted back as `device` + round_trip = pa.Tensor.from_dlpack( + tensor, device=tensor.__dlpack_device__()) + check(assert_type(round_trip, pa.Tensor), pa.Tensor) + + +def test_sparse_from_numpy() -> None: + dim_names = ("x", "y") + coo = pa.SparseCOOTensor.from_numpy( + np.array([[1, 2, 3]]).T, np.array([[0, 0], [0, 2], [1, 1]]), + (2, 3), dim_names) + check(assert_type(coo, pa.SparseCOOTensor), pa.SparseCOOTensor) + check(assert_type(coo.dim_names, tuple[str, ...]), tuple) + + coo = pa.SparseCOOTensor.from_numpy( + np.array([[1, 2, 3]]).T, np.array([[0, 0], [0, 2], [1, 1]]), + [2, 3], dim_names) + check(assert_type(coo, pa.SparseCOOTensor), pa.SparseCOOTensor) + + csr = pa.SparseCSRMatrix.from_numpy( + np.array([[1, 2, 3]]).T, np.array([0, 2, 3]), np.array([0, 2, 1]), + (2, 3), dim_names) + check(assert_type(csr, pa.SparseCSRMatrix), pa.SparseCSRMatrix) + check(assert_type(csr.dim_names, tuple[str, ...]), tuple) + + csc = pa.SparseCSCMatrix.from_numpy( + np.array([[1, 3, 2]]).T, np.array([0, 1, 2, 3]), + np.array([0, 1, 0]), (2, 3), dim_names) + check(assert_type(csc, pa.SparseCSCMatrix), pa.SparseCSCMatrix) + check(assert_type(csc.dim_names, tuple[str, ...]), tuple) + + csf = pa.SparseCSFTensor.from_numpy( + np.array([[1, 2, 3]]).T, [np.array([0, 2, 3])], + [np.array([0, 1]), np.array([0, 2, 1])], (2, 3), + axis_order=(0, 1), dim_names=dim_names) + check(assert_type(csf, pa.SparseCSFTensor), pa.SparseCSFTensor) + check(assert_type(csf.dim_names, tuple[str, ...]), tuple) + + csf = pa.SparseCSFTensor.from_numpy( + np.array([[1, 2, 3]]).T, (np.array([0, 2, 3]),), + (np.array([0, 1]), np.array([0, 2, 1])), (2, 3)) + check(assert_type(csf, pa.SparseCSFTensor), pa.SparseCSFTensor) + + +def test_sparse_scipy() -> None: + coo = pa.SparseCOOTensor.from_scipy(coo_array(DENSE)) + check(assert_type(coo.to_scipy(), coo_array), coo_array) + + csr = pa.SparseCSRMatrix.from_scipy(csr_array(DENSE)) + check(assert_type(csr.to_scipy(), csr_array), csr_array) + + csc = pa.SparseCSCMatrix.from_scipy(csc_array(DENSE)) + check(assert_type(csc.to_scipy(), csc_array), csc_array) + + +def test_invalid_usage() -> None: + tensor = pa.Tensor.from_numpy(DENSE) + if TYPE_CHECKING: + pa.Tensor.from_numpy(DENSE, dim_names=[1, 2]) # type: ignore[list-item] + _: Hashable = tensor # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs-tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py new file mode 100644 index 000000000000..f081cd62d8df --- /dev/null +++ b/python/pyarrow-stubs-tests/test_types.py @@ -0,0 +1,215 @@ +# 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. + +from collections.abc import Hashable, Iterator +from typing import Literal, TYPE_CHECKING + +from typing_extensions import CapsuleType + +import pyarrow as pa +from pyarrow._stubs_typing import ( + BoolType, + Int16Type, + Int32Type, + Int64Type, + IntegerType, + StringType, +) +import pyarrow.types as types + +from _check import assert_type, check + + +def _erase(t: pa.DataType) -> pa.DataType: + # Hide the precise type, so we can test narrowing by the is_* functions. + return t + + +def test_primitive_factories() -> None: + check(assert_type(pa.bool_(), BoolType), pa.DataType) + check(assert_type(pa.int32(), Int32Type), pa.DataType) + check(assert_type(pa.string(), StringType), pa.DataType) + check(assert_type(pa.utf8(), StringType), pa.DataType) + + +def test_temporal_types() -> None: + ts = pa.timestamp("ms", tz="UTC") + check(assert_type(ts, pa.TimestampType), pa.TimestampType) + check(assert_type(ts.unit, Literal["s", "ms", "us", "ns"]), str) + check(assert_type(ts.tz, str | None), str) + + check(assert_type(pa.time32("s").unit, Literal["s", "ms"]), str) + check(assert_type(pa.time64("ns").unit, Literal["us", "ns"]), str) + check(assert_type(pa.duration("us"), pa.DurationType), pa.DurationType) + + +def test_decimal_types() -> None: + dec = pa.decimal128(10, 2) + check(assert_type(dec, pa.Decimal128Type), pa.Decimal128Type) + check(assert_type(dec.precision, int), int) + check(assert_type(dec.scale, int), int) + + +def test_list_types() -> None: + lst = pa.list_(pa.int32()) + check(assert_type(lst, pa.ListType), pa.ListType) + check(assert_type(lst.value_type, pa.DataType), pa.DataType) + check(assert_type(lst.value_field, pa.Field), pa.Field) + + variable = pa.list_(pa.int32(), -1) + check(assert_type(variable, pa.ListType), pa.ListType) + fixed = pa.list_(pa.int32(), 2) + check(assert_type(fixed, pa.ListType | pa.FixedSizeListType), pa.FixedSizeListType) + + check(assert_type(pa.large_list(pa.int32()), pa.LargeListType), pa.LargeListType) + check(assert_type(pa.list_view(pa.int32()), pa.ListViewType), pa.ListViewType) + + +def test_dictionary_type() -> None: + dict_type = pa.dictionary(pa.int8(), pa.string()) + check(assert_type(dict_type, pa.DictionaryType), pa.DictionaryType) + check(assert_type(dict_type.ordered, bool), bool) + check(assert_type(dict_type.value_type, pa.DataType), pa.DataType) + check(assert_type(dict_type.index_type, IntegerType), pa.DataType) + + +def test_map_and_run_end_encoded_types() -> None: + map_type = pa.map_(pa.string(), pa.int64()) + check(assert_type(map_type, pa.MapType), pa.MapType) + check(assert_type(map_type.key_type, pa.DataType), pa.DataType) + check(assert_type(map_type.keys_sorted, bool), bool) + + ree = pa.run_end_encoded(pa.int16(), pa.string()) + check(assert_type(ree, pa.RunEndEncodedType), pa.RunEndEncodedType) + check(assert_type(ree.run_end_type, Int16Type | Int32Type | Int64Type), pa.DataType) + + +def test_data_type_hashable() -> None: + data_type = pa.int32() + check(assert_type(hash(data_type), int), int) + + +def test_field() -> None: + field = pa.field("a", pa.int32()) + check(assert_type(field, pa.Field), pa.Field) + check(assert_type(field.name, str), str) + check(assert_type(field.metadata, dict[bytes, bytes] | None), type(None)) + check(assert_type(field.type, pa.DataType), pa.DataType) + check(assert_type(field.nullable, bool), bool) + check(assert_type(field.with_nullable(False), pa.Field), pa.Field) + check(assert_type(pa.field("b", "string"), pa.Field), pa.Field) + + field = pa.field("a", pa.int32(), metadata={"k": "v"}) + check(assert_type(field.metadata, dict[bytes, bytes] | None), dict) + check(assert_type(field.with_metadata(None), pa.Field), pa.Field) + check( + assert_type(field.with_metadata(None).metadata, dict[bytes, bytes] | None), + dict # Empty dict + ) + + +def test_schema() -> None: + schema = pa.schema([("a", pa.int32()), pa.field("b", pa.string())]) + check(assert_type(schema, pa.Schema), pa.Schema) + check(assert_type(schema.names, list[str]), list) + check(assert_type(schema.types, list[pa.DataType]), list) + check(assert_type(schema.field("a"), pa.Field), pa.Field) + check(assert_type(schema[0], pa.Field), pa.Field) + + +def test_key_value_metadata() -> None: + # The runtime yields generators instead of Mapping views (duplicated keys) + metadata = pa.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) + check(assert_type(metadata.keys(), Iterator[bytes]), Iterator) + check(assert_type(metadata.values(), Iterator[bytes]), Iterator) + check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), Iterator) + + +def test_type_predicates_narrow() -> None: + lst = _erase(pa.list_(pa.int32())) + assert types.is_list(lst) + check(assert_type(lst, pa.ListType), pa.ListType) + + ts = _erase(pa.timestamp("s")) + assert types.is_timestamp(ts) + check(assert_type(ts, pa.TimestampType), pa.TimestampType) + + dec = _erase(pa.decimal64(5)) + assert types.is_decimal(dec) + check( + assert_type( + dec, + pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type | pa.Decimal256Type, + ), + pa.Decimal64Type, + ) + + # Decimal types subclass FixedSizeBinaryType, but don't match it, so + # they must not be narrowed away when the check fails. + dec128 = _erase(pa.decimal128(5, 2)) + if not types.is_fixed_size_binary(dec128): + assert types.is_decimal(dec128) + check( + assert_type( + dec128, + pa.Decimal32Type + | pa.Decimal64Type + | pa.Decimal128Type + | pa.Decimal256Type, + ), + pa.Decimal128Type, + ) + + # Run-end encoded types are nested, so this doesn't narrow to Never. + ree = pa.run_end_encoded(pa.int32(), pa.string()) + assert types.is_nested(ree) + check(assert_type(ree, pa.RunEndEncodedType), pa.RunEndEncodedType) + + +def test_import_from_c_capsule() -> None: + capsule = pa.int32().__arrow_c_schema__() + check(assert_type(capsule, CapsuleType), CapsuleType) + + # The import methods are static and return whatever type the C schema + # holds, not the class they are called on. + imported = pa.ListType._import_from_c_capsule(capsule) + check(assert_type(imported, pa.DataType), pa.DataType) + assert not isinstance(imported, pa.ListType) + + field = pa.field("a", pa.int32()) + field_capsule = field.__arrow_c_schema__() + check(assert_type(field_capsule, CapsuleType), CapsuleType) + + imported_field = pa.Field._import_from_c_capsule(field_capsule) + check(assert_type(imported_field, pa.Field), pa.Field) + + schema = pa.schema([field]) + schema_capsule = schema.__arrow_c_schema__() + check(assert_type(schema_capsule, CapsuleType), CapsuleType) + + imported_schema = pa.Schema._import_from_c_capsule(schema_capsule) + check(assert_type(imported_schema, pa.Schema), pa.Schema) + + +def test_invalid_usage() -> None: + if TYPE_CHECKING: + # A Field needs a type unless it's created from an Arrow C schema. + pa.field("a") # type: ignore[call-overload] + pa.list_(1) # type: ignore[call-overload] + pa.decimal128("10") # type: ignore[arg-type] + pa.MonthDayNano() # type: ignore[call-arg] # ty: ignore[missing-argument] + _: Hashable = pa.KeyValueMetadata({b"a": b"1"}) # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs/__init__.pyi b/python/pyarrow-stubs/__init__.pyi new file mode 100644 index 000000000000..68e7c457e65f --- /dev/null +++ b/python/pyarrow-stubs/__init__.pyi @@ -0,0 +1,229 @@ +# 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. + +from typing import Any + +# The names that pyarrow/__init__.py imports from pyarrow.lib and that have +# stubs already. Type checkers only fall back to __getattr__ below for the +# names that aren't defined here. +from pyarrow.lib import ( + Array as Array, + ArrowCancelled as ArrowCancelled, + ArrowCapacityError as ArrowCapacityError, + ArrowException as ArrowException, + ArrowIOError as ArrowIOError, + ArrowIndexError as ArrowIndexError, + ArrowInvalid as ArrowInvalid, + ArrowKeyError as ArrowKeyError, + ArrowMemoryError as ArrowMemoryError, + ArrowNotImplementedError as ArrowNotImplementedError, + ArrowSerializationError as ArrowSerializationError, + ArrowTypeError as ArrowTypeError, + BaseExtensionType as BaseExtensionType, + BinaryArray as BinaryArray, + BinaryViewArray as BinaryViewArray, + Bool8Array as Bool8Array, + Bool8Type as Bool8Type, + BooleanArray as BooleanArray, + BuildInfo as BuildInfo, + ChunkedArray as ChunkedArray, + CppBuildInfo as CppBuildInfo, + DataType as DataType, + Date32Array as Date32Array, + Date64Array as Date64Array, + Decimal128Array as Decimal128Array, + Decimal128Type as Decimal128Type, + Decimal256Array as Decimal256Array, + Decimal256Type as Decimal256Type, + Decimal32Array as Decimal32Array, + Decimal32Type as Decimal32Type, + Decimal64Array as Decimal64Array, + Decimal64Type as Decimal64Type, + DenseUnionType as DenseUnionType, + Device as Device, + DeviceAllocationType as DeviceAllocationType, + DictionaryArray as DictionaryArray, + DictionaryMemo as DictionaryMemo, + DictionaryType as DictionaryType, + DoubleArray as DoubleArray, + DurationArray as DurationArray, + DurationType as DurationType, + ExtensionArray as ExtensionArray, + ExtensionType as ExtensionType, + Field as Field, + FixedShapeTensorArray as FixedShapeTensorArray, + FixedShapeTensorType as FixedShapeTensorType, + FixedSizeBinaryArray as FixedSizeBinaryArray, + FixedSizeBinaryType as FixedSizeBinaryType, + FixedSizeListArray as FixedSizeListArray, + FixedSizeListType as FixedSizeListType, + FloatArray as FloatArray, + FloatingPointArray as FloatingPointArray, + HalfFloatArray as HalfFloatArray, + Int16Array as Int16Array, + Int32Array as Int32Array, + Int64Array as Int64Array, + Int8Array as Int8Array, + IntegerArray as IntegerArray, + JsonArray as JsonArray, + JsonType as JsonType, + KeyValueMetadata as KeyValueMetadata, + LargeBinaryArray as LargeBinaryArray, + LargeListArray as LargeListArray, + LargeListType as LargeListType, + LargeListViewArray as LargeListViewArray, + LargeListViewType as LargeListViewType, + LargeStringArray as LargeStringArray, + ListArray as ListArray, + ListType as ListType, + ListViewArray as ListViewArray, + ListViewType as ListViewType, + LoggingMemoryPool as LoggingMemoryPool, + MapArray as MapArray, + MapType as MapType, + MemoryManager as MemoryManager, + MemoryPool as MemoryPool, + MonthDayNano as MonthDayNano, + MonthDayNanoIntervalArray as MonthDayNanoIntervalArray, + NullArray as NullArray, + NumericArray as NumericArray, + OpaqueArray as OpaqueArray, + OpaqueType as OpaqueType, + ProxyMemoryPool as ProxyMemoryPool, + RecordBatch as RecordBatch, + RunEndEncodedArray as RunEndEncodedArray, + RunEndEncodedType as RunEndEncodedType, + RuntimeInfo as RuntimeInfo, + Schema as Schema, + SparseCOOTensor as SparseCOOTensor, + SparseCSCMatrix as SparseCSCMatrix, + SparseCSFTensor as SparseCSFTensor, + SparseCSRMatrix as SparseCSRMatrix, + SparseUnionType as SparseUnionType, + StringArray as StringArray, + StringViewArray as StringViewArray, + StructArray as StructArray, + StructType as StructType, + Table as Table, + TableGroupBy as TableGroupBy, + Tensor as Tensor, + Time32Array as Time32Array, + Time32Type as Time32Type, + Time64Array as Time64Array, + Time64Type as Time64Type, + TimestampArray as TimestampArray, + TimestampType as TimestampType, + UInt16Array as UInt16Array, + UInt32Array as UInt32Array, + UInt64Array as UInt64Array, + UInt8Array as UInt8Array, + UnionArray as UnionArray, + UnionType as UnionType, + UnknownExtensionType as UnknownExtensionType, + UuidArray as UuidArray, + UuidType as UuidType, + VersionInfo as VersionInfo, + array as array, + binary as binary, + binary_view as binary_view, + bool8 as bool8, + bool_ as bool_, + build_info as build_info, + chunked_array as chunked_array, + concat_arrays as concat_arrays, + concat_batches as concat_batches, + concat_tables as concat_tables, + cpp_build_info as cpp_build_info, + cpp_version as cpp_version, + cpp_version_info as cpp_version_info, + date32 as date32, + date64 as date64, + decimal128 as decimal128, + decimal256 as decimal256, + decimal32 as decimal32, + decimal64 as decimal64, + default_cpu_memory_manager as default_cpu_memory_manager, + default_memory_pool as default_memory_pool, + dense_union as dense_union, + dictionary as dictionary, + duration as duration, + enable_signal_handlers as enable_signal_handlers, + field as field, + fixed_shape_tensor as fixed_shape_tensor, + float16 as float16, + float32 as float32, + float64 as float64, + from_numpy_dtype as from_numpy_dtype, + infer_type as infer_type, + int16 as int16, + int32 as int32, + int64 as int64, + int8 as int8, + jemalloc_memory_pool as jemalloc_memory_pool, + jemalloc_set_decay_ms as jemalloc_set_decay_ms, + json_ as json_, + large_binary as large_binary, + large_list as large_list, + large_list_view as large_list_view, + large_string as large_string, + large_utf8 as large_utf8, + list_ as list_, + list_view as list_view, + log_memory_allocations as log_memory_allocations, + logging_memory_pool as logging_memory_pool, + map_ as map_, + mimalloc_memory_pool as mimalloc_memory_pool, + month_day_nano_interval as month_day_nano_interval, + null as null, + nulls as nulls, + opaque as opaque, + proxy_memory_pool as proxy_memory_pool, + record_batch as record_batch, + register_extension_type as register_extension_type, + repeat as repeat, + run_end_encoded as run_end_encoded, + runtime_info as runtime_info, + schema as schema, + set_memory_pool as set_memory_pool, + set_timezone_db_path as set_timezone_db_path, + sparse_union as sparse_union, + string as string, + string_view as string_view, + struct as struct, + supported_memory_backends as supported_memory_backends, + system_memory_pool as system_memory_pool, + table as table, + time32 as time32, + time64 as time64, + timestamp as timestamp, + total_allocated_bytes as total_allocated_bytes, + type_for_alias as type_for_alias, + uint16 as uint16, + uint32 as uint32, + uint64 as uint64, + uint8 as uint8, + unify_schemas as unify_schemas, + union as union, + unregister_extension_type as unregister_extension_type, + utf8 as utf8, + uuid as uuid, +) + +# TODO(GH-48970): remove __getattr__ before release as this +# will annotate non-existing attributes as Any. +# https://github.com/apache/arrow/issues/48970 +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/__lib/__init__.pyi b/python/pyarrow-stubs/__lib/__init__.pyi new file mode 100644 index 000000000000..50d570015931 --- /dev/null +++ b/python/pyarrow-stubs/__lib/__init__.pyi @@ -0,0 +1,18 @@ +# 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. + +# Separate subpackage so that e.g. array.pyi doesn't shadow pyarrow.array() diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi new file mode 100644 index 000000000000..21d8d732cd56 --- /dev/null +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -0,0 +1,620 @@ +# 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 builtins +from collections.abc import Collection, Iterable, Iterator, Sequence +from typing import Any, ClassVar, Self, SupportsIndex + +import numpy as np +import pandas as pd +from numpy.typing import NDArray + +# Placeholders (Any) until scalar.pyi is annotated; mypy rejects them as +# types when imported through pyarrow.lib. +from pyarrow.__lib.scalar import Int64Scalar, Scalar +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + MapsAsPydicts, + Mask, + NullEncoding, + NullPlacement, + NullSelectionBehavior, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + BaseExtensionType, + Buffer, + DataType, + Device, + DeviceAllocationType, + Field, + MemoryManager, + MemoryPool, + StructType, + Tensor, + _Weakrefable, +) +from typing_extensions import CapsuleType + +def array( + obj: Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray, + type: Any | None = None, + mask: BooleanArray | Iterable[bool] | NDArray[np.bool_] | None = None, + size: int | None = None, + from_pandas: bool | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... +def asarray( + values: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, + type: Any | None = None, +) -> ArrayLike: ... +def nulls( + size: int, + type: Any | None = None, + memory_pool: MemoryPool | None = None, +) -> Array: ... +def repeat( + value: Any, + size: int, + memory_pool: MemoryPool | None = None, +) -> Array: ... +def infer_type( + values: Sequence[Any] | np.ndarray, + mask: Mask | None = None, + from_pandas: bool = False, +) -> DataType: ... + +class ArrayStatistics(_Weakrefable): + @property + def null_count(self) -> int | None: ... + @property + def distinct_count(self) -> int | None: ... + @property + def is_null_count_exact(self) -> bool | None: ... + @property + def is_distinct_count_exact(self) -> bool | None: ... + @property + def min(self) -> Any | None: ... + @property + def is_min_exact(self) -> bool | None: ... + @property + def max(self) -> Any | None: ... + @property + def is_max_exact(self) -> bool | None: ... + +class _PandasConvertible(_Weakrefable): + def to_pandas( + self, + memory_pool: MemoryPool | None = None, + categories: list | tuple | None = None, + strings_to_categorical: bool = False, + zero_copy_only: bool = False, + integer_object_nulls: bool = False, + date_as_object: bool = True, + timestamp_as_object: bool = False, + use_threads: bool = True, + deduplicate_objects: bool = True, + ignore_metadata: bool = False, + safe: bool = True, + split_blocks: bool = False, + self_destruct: bool = False, + maps_as_pydicts: MapsAsPydicts | None = None, + types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None + coerce_temporal_nanoseconds: bool = False, + ) -> pd.Series | pd.DataFrame: ... + +class Array(_PandasConvertible): + __hash__: ClassVar[None] # type: ignore[assignment] + def diff(self, other: Array) -> str: ... + def cast( + self, + target_type: DataType | str | None = None, + safe: bool | None = None, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + def view(self, target_type: DataType | str) -> Array: ... + def sum(self, **kwargs: Any) -> Scalar: ... + @property + def type(self) -> DataType: ... + def unique(self) -> Self: ... + def dictionary_encode( + self, null_encoding: NullEncoding = "mask" + ) -> DictionaryArray: ... + def value_counts(self) -> StructArray: ... + @staticmethod + def from_pandas( + obj: pd.Series | np.ndarray | ArrayLike, + mask: Mask | None = None, + type: DataType | str | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + @staticmethod + def from_buffers( + type: DataType, + length: int, + buffers: Collection[Buffer | None], + null_count: int = -1, + offset: int = 0, + children: Collection[Array] | None = None, + ) -> Array: ... + @property + def null_count(self) -> int: ... + @property + def nbytes(self) -> int: ... + def get_total_buffer_size(self) -> int: ... + def __sizeof__(self) -> int: ... + def __iter__(self) -> Iterator[Scalar]: ... + def to_string( + self, + *, + indent: int = 2, + top_level_indent: int = 0, + window: int = 10, + container_window: int = 2, + skip_new_lines: bool = False, + element_size_limit: int = 100, + ) -> str: ... + def equals(self, other: Array) -> bool: ... + def __len__(self) -> int: ... + def is_null(self, *, nan_is_null: bool = False) -> BooleanArray: ... + def is_nan(self) -> BooleanArray: ... + def is_valid(self) -> BooleanArray: ... + def fill_null(self, fill_value: Any) -> Self: ... + def __getitem__(self, key: SupportsIndex | builtins.slice) -> Scalar | Self: ... + def slice(self, offset: int = 0, length: int | None = None) -> Self: ... + def take(self, indices: Indices) -> Self: ... + def drop_null(self) -> Self: ... + def filter( + self, + mask: Mask, + *, + null_selection_behavior: NullSelectionBehavior = "drop", + ) -> Self: ... + def index( + self, + value: Any, + start: int | None = None, + end: int | None = None, + *, + memory_pool: MemoryPool | None = None, + ) -> Int64Scalar: ... + def sort( + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + **kwargs: Any, + ) -> Self: ... + def __array__( + self, dtype: np.dtype | None = None, copy: bool | None = None + ) -> np.ndarray: ... + def to_numpy( + self, zero_copy_only: bool = True, writable: bool = False + ) -> np.ndarray: ... + def to_pylist( + self, + *, + maps_as_pydicts: MapsAsPydicts | None = None, + ) -> list[Any]: ... + + tolist = to_pylist + def to_tensor(self, *, allow_nulls: bool = False) -> Tensor: ... + def validate(self, *, full: bool = False) -> None: ... + @property + def offset(self) -> int: ... + def buffers(self) -> list[Buffer | None]: ... + def copy_to(self, destination: MemoryManager | Device) -> Self: ... + def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int, type: int | DataType) -> Array: ... + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: ... + @staticmethod + def _import_from_c_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> Array: ... + def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c_device(in_ptr: int, type: DataType | int) -> Array: ... + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: ... + @staticmethod + def _import_from_c_device_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> Array: ... + def __dlpack__( + self, + *, + stream: int | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> CapsuleType: ... + def __dlpack_device__(self) -> tuple[int, int]: ... + @property + def device_type(self) -> DeviceAllocationType: ... + @property + def is_cpu(self) -> bool: ... + @property + def statistics(self) -> ArrayStatistics | None: ... + +class NullArray(Array): ... + +class BooleanArray(Array): + @property + def false_count(self) -> int: ... + @property + def true_count(self) -> int: ... + +class NumericArray(Array): ... +class IntegerArray(NumericArray): ... +class FloatingPointArray(NumericArray): ... +class Int8Array(IntegerArray): ... +class UInt8Array(IntegerArray): ... +class Int16Array(IntegerArray): ... +class UInt16Array(IntegerArray): ... +class Int32Array(IntegerArray): ... +class UInt32Array(IntegerArray): ... +class Int64Array(IntegerArray): ... +class UInt64Array(IntegerArray): ... +class Date32Array(NumericArray): ... +class Date64Array(NumericArray): ... +class TimestampArray(NumericArray): ... +class Time32Array(NumericArray): ... +class Time64Array(NumericArray): ... +class DurationArray(NumericArray): ... +class MonthDayNanoIntervalArray(Array): ... +class HalfFloatArray(FloatingPointArray): ... +class FloatArray(FloatingPointArray): ... +class DoubleArray(FloatingPointArray): ... +class FixedSizeBinaryArray(Array): ... +class Decimal32Array(FixedSizeBinaryArray): ... +class Decimal64Array(FixedSizeBinaryArray): ... +class Decimal128Array(FixedSizeBinaryArray): ... +class Decimal256Array(FixedSizeBinaryArray): ... + +class BaseListArray(Array): + def flatten(self, recursive: bool = False) -> Array: ... + def value_parent_indices(self) -> Int64Array: ... + def value_lengths(self) -> Int32Array | Int64Array: ... + +class ListArray(BaseListArray): + @staticmethod + def from_arrays( + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> ListArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + +class LargeListArray(BaseListArray): + @staticmethod + def from_arrays( + offsets: Int64Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + +class ListViewArray(BaseListArray): + @staticmethod + def from_arrays( + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, + sizes: Int32Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> ListViewArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int32Array: ... + @property + def sizes(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + +class LargeListViewArray(BaseListArray): + @staticmethod + def from_arrays( + offsets: Int64Array | list[int] | list[int | None] | np.ndarray, + sizes: Int64Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListViewArray: ... + @property + def values(self) -> Array: ... + @property + def offsets(self) -> Int64Array: ... + @property + def sizes(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + +class FixedSizeListArray(BaseListArray): + @staticmethod + def from_arrays( + values: Array | list[Any], + list_size: int | None = None, + type: DataType | None = None, + mask: Mask | None = None, + ) -> FixedSizeListArray: ... + @property + def values(self) -> Array: ... + def value_lengths(self) -> Int32Array: ... + +class MapArray(ListArray): + @staticmethod + def from_arrays( # type: ignore[override] + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, + keys: Array | np.ndarray | list, + items: Array | np.ndarray | list, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> MapArray: ... + @property + def keys(self) -> Array: ... + @property + def items(self) -> Array: ... + +class UnionArray(Array): + def field(self, pos: int) -> Array: ... + @property + def type_codes(self) -> Int8Array: ... + @property + def offsets(self) -> Int32Array: ... + @staticmethod + def from_dense( + types: Int8Array, + value_offsets: Int32Array, + children: list[Array], + field_names: list[str] | None = None, + type_codes: list[int] | None = None, + ) -> UnionArray: ... + @staticmethod + def from_sparse( + types: Int8Array, + children: list[Array], + field_names: list[str] | None = None, + type_codes: list[int] | None = None, + ) -> UnionArray: ... + +class StringArray(Array): + @staticmethod + def from_buffers( # type: ignore[override] + length: int, + value_offsets: Buffer, + data: Buffer, + null_bitmap: Buffer | None = None, + null_count: int = -1, + offset: int = 0, + ) -> StringArray: ... + +class LargeStringArray(Array): + @staticmethod + def from_buffers( # type: ignore[override] + length: int, + value_offsets: Buffer, + data: Buffer, + null_bitmap: Buffer | None = None, + null_count: int = -1, + offset: int = 0, + ) -> LargeStringArray: ... + +class StringViewArray(Array): ... + +class BinaryArray(Array): + @property + def total_values_length(self) -> int: ... + +class LargeBinaryArray(Array): + @property + def total_values_length(self) -> int: ... + +class BinaryViewArray(Array): ... + +class DictionaryArray(Array): + def dictionary_encode(self) -> Self: ... # type: ignore[override] + def dictionary_decode(self) -> Array: ... + @property + def indices(self) -> IntegerArray: ... + @property + def dictionary(self) -> Array: ... + @staticmethod + def from_buffers( # type: ignore[override] + type: DataType, + length: int, + buffers: Collection[Buffer | None], + dictionary: Array, + null_count: int = -1, + offset: int = 0, + ) -> DictionaryArray: ... + @staticmethod + def from_arrays( + indices: Indices, + dictionary: Array | np.ndarray | pd.Series | list[Any], + mask: np.ndarray | pd.Series | BooleanArray | None = None, + ordered: bool = False, + from_pandas: bool = False, + safe: bool = True, + memory_pool: MemoryPool | None = None, + ) -> DictionaryArray: ... + +class StructArray(Array): + def field(self, index: int | str) -> Array: ... + def flatten(self, memory_pool: MemoryPool | None = None) -> list[Array]: ... + @staticmethod + def from_arrays( + arrays: Iterable[Array | np.ndarray | list], + names: Sequence[str] | list[Field] | None = None, + fields: list[Field] | None = None, + mask: BooleanArray | None = None, + memory_pool: MemoryPool | None = None, + type: StructType | None = None, + ) -> StructArray: ... + def sort( + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + by: str | None = None, + **kwargs: Any, + ) -> StructArray: ... + +class RunEndEncodedArray(Array): + @staticmethod + def from_arrays( + run_ends: Int16Array | Int32Array | Int64Array | list[int], + values: Array | list[Any], + type: DataType | None = None, + ) -> RunEndEncodedArray: ... + @staticmethod + def from_buffers( # type: ignore[override] + type: DataType, + length: int, + buffers: list[Buffer] | list[None], + null_count: int = -1, + offset: int = 0, + children: tuple[Array, Array] | list[list[int]] | None = None, + ) -> RunEndEncodedArray: ... + @property + def run_ends(self) -> Int16Array | Int32Array | Int64Array: ... + @property + def values(self) -> Array: ... + def find_physical_offset(self) -> int: ... + def find_physical_length(self) -> int: ... + +class ExtensionArray(Array): + @property + def storage(self) -> Any: ... + @staticmethod + def from_storage(typ: BaseExtensionType, storage: Array) -> ExtensionArray: ... + +class JsonArray(ExtensionArray): ... +class UuidArray(ExtensionArray): ... + +class FixedShapeTensorArray(ExtensionArray): + def to_numpy_ndarray(self) -> np.ndarray: ... + @staticmethod + def from_numpy_ndarray( + obj: np.ndarray, dim_names: list[str] | tuple[str, ...] | None = None + ) -> FixedShapeTensorArray: ... + +class OpaqueArray(ExtensionArray): ... + +class Bool8Array(ExtensionArray): + def to_numpy( + self, zero_copy_only: bool = ..., writable: bool = ... + ) -> np.ndarray: ... + @staticmethod + def from_storage(storage: Int8Array) -> Bool8Array: ... # type: ignore[override] + @staticmethod + def from_numpy(obj: np.ndarray) -> Bool8Array: ... + +def concat_arrays( + arrays: Iterable[Array], memory_pool: MemoryPool | None = None +) -> Array: ... +def _empty_array(type: DataType) -> Array: ... + +__all__ = [ + "array", + "asarray", + "nulls", + "repeat", + "infer_type", + "ArrayStatistics", + "_PandasConvertible", + "Array", + "NullArray", + "BooleanArray", + "NumericArray", + "IntegerArray", + "FloatingPointArray", + "Int8Array", + "UInt8Array", + "Int16Array", + "UInt16Array", + "Int32Array", + "UInt32Array", + "Int64Array", + "UInt64Array", + "Date32Array", + "Date64Array", + "TimestampArray", + "Time32Array", + "Time64Array", + "DurationArray", + "MonthDayNanoIntervalArray", + "HalfFloatArray", + "FloatArray", + "DoubleArray", + "FixedSizeBinaryArray", + "Decimal32Array", + "Decimal64Array", + "Decimal128Array", + "Decimal256Array", + "BaseListArray", + "ListArray", + "LargeListArray", + "ListViewArray", + "LargeListViewArray", + "FixedSizeListArray", + "MapArray", + "UnionArray", + "StringArray", + "LargeStringArray", + "StringViewArray", + "BinaryArray", + "LargeBinaryArray", + "BinaryViewArray", + "DictionaryArray", + "StructArray", + "RunEndEncodedArray", + "ExtensionArray", + "Bool8Array", + "UuidArray", + "JsonArray", + "OpaqueArray", + "FixedShapeTensorArray", + "concat_arrays", + "_empty_array", +] diff --git a/python/pyarrow-stubs/__lib/builder.pyi b/python/pyarrow-stubs/__lib/builder.pyi new file mode 100644 index 000000000000..869e9c27727f --- /dev/null +++ b/python/pyarrow-stubs/__lib/builder.pyi @@ -0,0 +1,45 @@ +# 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. + +from collections.abc import Iterable + +from pyarrow.lib import ( + MemoryPool, + StringArray, + StringViewArray, + _Weakrefable, +) + +class StringBuilder(_Weakrefable): + def __init__(self, memory_pool: MemoryPool | None = None) -> None: ... + def append(self, value: str | bytes | float | None) -> None: ... + def append_values(self, values: Iterable[str | bytes | float | None]) -> None: ... + def finish(self) -> StringArray: ... + @property + def null_count(self) -> int: ... + def __len__(self) -> int: ... + +class StringViewBuilder(_Weakrefable): + def __init__(self, memory_pool: MemoryPool | None = None) -> None: ... + def append(self, value: str | bytes | float | None) -> None: ... + def append_values(self, values: Iterable[str | bytes | float | None]) -> None: ... + def finish(self) -> StringViewArray: ... + @property + def null_count(self) -> int: ... + def __len__(self) -> int: ... + +__all__ = ["StringBuilder", "StringViewBuilder"] diff --git a/python/pyarrow-stubs/__lib/config.pyi b/python/pyarrow-stubs/__lib/config.pyi new file mode 100644 index 000000000000..5f2fbc5ee863 --- /dev/null +++ b/python/pyarrow-stubs/__lib/config.pyi @@ -0,0 +1,65 @@ +# 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. + +from typing import NamedTuple + +class VersionInfo(NamedTuple): + major: int + minor: int + patch: int + +class CppBuildInfo(NamedTuple): + version: str + version_info: VersionInfo + so_version: str + full_so_version: str + compiler_id: str + compiler_version: str + compiler_flags: str + git_id: str + git_description: str + package_kind: str + build_type: str + +class BuildInfo(NamedTuple): + build_type: str + cpp_build_info: CppBuildInfo + +class RuntimeInfo(NamedTuple): + simd_level: str + detected_simd_level: str + +build_info: BuildInfo +cpp_build_info: CppBuildInfo +cpp_version: str +cpp_version_info: VersionInfo + +def runtime_info() -> RuntimeInfo: ... +def set_timezone_db_path(path: str) -> None: ... + +__all__ = [ + "VersionInfo", + "BuildInfo", + "CppBuildInfo", + "RuntimeInfo", + "build_info", + "cpp_build_info", + "cpp_version", + "cpp_version_info", + "runtime_info", + "set_timezone_db_path", +] diff --git a/python/pyarrow-stubs/__lib/device.pyi b/python/pyarrow-stubs/__lib/device.pyi new file mode 100644 index 000000000000..7bc20ffb7ec2 --- /dev/null +++ b/python/pyarrow-stubs/__lib/device.pyi @@ -0,0 +1,61 @@ +# 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 enum + +from pyarrow.lib import _Weakrefable + +class DeviceAllocationType(enum.IntEnum): + CPU = 1 + CUDA = 2 + CUDA_HOST = 3 + OPENCL = 4 + VULKAN = 7 + METAL = 8 + VPI = 9 + ROCM = 10 + ROCM_HOST = 11 + EXT_DEV = 12 + CUDA_MANAGED = 13 + ONEAPI = 14 + WEBGPU = 15 + HEXAGON = 16 + +class Device(_Weakrefable): + @property + def type_name(self) -> str: ... + @property + def device_id(self) -> int: ... + @property + def is_cpu(self) -> bool: ... + @property + def device_type(self) -> DeviceAllocationType: ... + +class MemoryManager(_Weakrefable): + @property + def device(self) -> Device: ... + @property + def is_cpu(self) -> bool: ... + +def default_cpu_memory_manager() -> MemoryManager: ... + +__all__ = [ + "DeviceAllocationType", + "Device", + "MemoryManager", + "default_cpu_memory_manager", +] diff --git a/python/pyarrow-stubs/__lib/error.pyi b/python/pyarrow-stubs/__lib/error.pyi new file mode 100644 index 000000000000..fa67de04535c --- /dev/null +++ b/python/pyarrow-stubs/__lib/error.pyi @@ -0,0 +1,70 @@ +# 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. + +from types import TracebackType +from typing import Self + +class ArrowException(Exception): ... +class ArrowInvalid(ValueError, ArrowException): ... +class ArrowMemoryError(MemoryError, ArrowException): ... +class ArrowKeyError(KeyError, ArrowException): ... +class ArrowTypeError(TypeError, ArrowException): ... +class ArrowNotImplementedError(NotImplementedError, ArrowException): ... +class ArrowCapacityError(ArrowException): ... +class ArrowIndexError(IndexError, ArrowException): ... +class ArrowSerializationError(ArrowException): ... + +class ArrowCancelled(ArrowException): + signum: int | None + def __init__(self, message: str, signum: int | None = None) -> None: ... + +ArrowIOError = IOError + +class StopToken: ... + +def enable_signal_handlers(enable: bool) -> None: ... + +have_signal_refcycle: bool + +class SignalStopHandler: + def __enter__(self) -> Self: ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: ... + @property + def stop_token(self) -> StopToken: ... + +__all__ = [ + "ArrowException", + "ArrowInvalid", + "ArrowMemoryError", + "ArrowKeyError", + "ArrowTypeError", + "ArrowNotImplementedError", + "ArrowCapacityError", + "ArrowIndexError", + "ArrowSerializationError", + "ArrowCancelled", + "ArrowIOError", + "StopToken", + "enable_signal_handlers", + "have_signal_refcycle", + "SignalStopHandler", +] diff --git a/python/pyarrow-stubs/pyarrow/__init__.pyi b/python/pyarrow-stubs/__lib/io.pyi similarity index 75% rename from python/pyarrow-stubs/pyarrow/__init__.pyi rename to python/pyarrow-stubs/__lib/io.pyi index ccec8d5abc07..8a1da39a1cf2 100644 --- a/python/pyarrow-stubs/pyarrow/__init__.pyi +++ b/python/pyarrow-stubs/__lib/io.pyi @@ -15,15 +15,11 @@ # 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. -""" +# Placeholder stub - complete annotations in future PR. +# At runtime, these symbols are provided by the pyarrow.lib C extension. from typing import Any -# TODO(GH-48970): remove __getattr__ before release as this -# will annotate non-existing attributes as Any. -# https://github.com/apache/arrow/issues/48970 +class Buffer: ... + def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/__lib/ipc.pyi b/python/pyarrow-stubs/__lib/ipc.pyi new file mode 100644 index 000000000000..c84964043e4a --- /dev/null +++ b/python/pyarrow-stubs/__lib/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: ... diff --git a/python/pyarrow-stubs/__lib/memory.pyi b/python/pyarrow-stubs/__lib/memory.pyi new file mode 100644 index 000000000000..2645f2fab067 --- /dev/null +++ b/python/pyarrow-stubs/__lib/memory.pyi @@ -0,0 +1,60 @@ +# 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. + +from pyarrow.lib import _Weakrefable + +class MemoryPool(_Weakrefable): + def release_unused(self) -> None: ... + def bytes_allocated(self) -> int: ... + def total_bytes_allocated(self) -> int: ... + def max_memory(self) -> int | None: ... + def num_allocations(self) -> int: ... + def print_stats(self) -> None: ... + @property + def backend_name(self) -> str: ... + +class LoggingMemoryPool(MemoryPool): ... +class ProxyMemoryPool(MemoryPool): ... + +def default_memory_pool() -> MemoryPool: ... +def proxy_memory_pool(parent: MemoryPool) -> ProxyMemoryPool: ... +def logging_memory_pool(parent: MemoryPool) -> LoggingMemoryPool: ... +def system_memory_pool() -> MemoryPool: ... +def jemalloc_memory_pool() -> MemoryPool: ... +def mimalloc_memory_pool() -> MemoryPool: ... +def set_memory_pool(pool: MemoryPool) -> None: ... +def log_memory_allocations(enable: bool = True) -> None: ... +def total_allocated_bytes() -> int: ... +def jemalloc_set_decay_ms(decay_ms: int) -> None: ... +def supported_memory_backends() -> list[str]: ... + +__all__ = [ + "MemoryPool", + "LoggingMemoryPool", + "ProxyMemoryPool", + "default_memory_pool", + "proxy_memory_pool", + "logging_memory_pool", + "system_memory_pool", + "jemalloc_memory_pool", + "mimalloc_memory_pool", + "set_memory_pool", + "log_memory_allocations", + "total_allocated_bytes", + "jemalloc_set_decay_ms", + "supported_memory_backends", +] diff --git a/python/pyarrow-stubs/__lib/scalar.pyi b/python/pyarrow-stubs/__lib/scalar.pyi new file mode 100644 index 000000000000..307b3e689a25 --- /dev/null +++ b/python/pyarrow-stubs/__lib/scalar.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 + +class ExtensionScalar: ... + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi new file mode 100644 index 000000000000..3237d51d3e44 --- /dev/null +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -0,0 +1,619 @@ +# 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 builtins +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, ClassVar, Literal, Self, TypeAlias + +import numpy as np +import pandas as pd +from numpy.typing import NDArray + +# Placeholders (Any) until ipc.pyi and scalar.pyi are annotated; mypy +# rejects them as types when imported through pyarrow.lib. +from pyarrow.__lib.ipc import RecordBatchReader +from pyarrow.__lib.scalar import Int64Scalar, Scalar +from pyarrow._compute import ( + CastOptions, + CountOptions, + FunctionOptions, + ScalarAggregateOptions, + TDigestOptions, + VarianceOptions, +) +from pyarrow._stubs_typing import ( + Indices, + MapsAsPydicts, + Mask, + NullEncoding, + NullPlacement, + NullSelectionBehavior, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsArrowStream, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + Array, + Buffer, + DataType, + Device, + DeviceAllocationType, + Field, + MemoryManager, + MemoryPool, + Schema, + StructArray, + Tensor, + _PandasConvertible, +) +from typing_extensions import CapsuleType + +from pyarrow.compute import Expression +from pyarrow.interchange.dataframe import _PyArrowDataFrame + +_ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray + +_Aggregation: TypeAlias = Literal[ + "all", + "any", + "approximate_median", + "count", + "count_all", + "count_distinct", + "distinct", + "first", + "first_last", + "kurtosis", + "last", + "list", + "max", + "mean", + "min", + "min_max", + "one", + "pivot_wider", + "product", + "skew", + "stddev", + "sum", + "tdigest", + "variance", +] +_AggregationPrefixed: TypeAlias = Literal[ + "hash_all", + "hash_any", + "hash_approximate_median", + "hash_count", + "hash_count_all", + "hash_count_distinct", + "hash_distinct", + "hash_first", + "hash_first_last", + "hash_kurtosis", + "hash_last", + "hash_list", + "hash_max", + "hash_mean", + "hash_min", + "hash_min_max", + "hash_one", + "hash_pivot_wider", + "hash_product", + "hash_skew", + "hash_stddev", + "hash_sum", + "hash_tdigest", + "hash_variance", +] +_AggregationName: TypeAlias = _Aggregation | _AggregationPrefixed | str +_AggregateOptions: TypeAlias = ( + ScalarAggregateOptions + | CountOptions + | TDigestOptions + | VarianceOptions + | FunctionOptions +) + +_ChunkedArrayConstructorInput: TypeAlias = ( + Iterable[ + Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray + ] + | Array + | SupportsArrowArray + | SupportsArrowStream +) + +_UnarySelector: TypeAlias = str +_NullarySelector: TypeAlias = tuple[()] +_NarySelector: TypeAlias = list[str] | tuple[str, ...] +_ColumnSelector: TypeAlias = _UnarySelector | _NullarySelector | _NarySelector + +class ChunkedArray(_PandasConvertible): + @property + def type(self) -> DataType: ... + def length(self) -> int: ... + + __len__ = length + + def to_string( + self, + *, + indent: int = 0, + window: int = 5, + container_window: int = 2, + skip_new_lines: bool = False, + element_size_limit: int = 100, + ) -> str: ... + def validate(self, *, full: bool = False) -> None: ... + @property + def null_count(self) -> int: ... + @property + def nbytes(self) -> int: ... + def get_total_buffer_size(self) -> int: ... + def __sizeof__(self) -> int: ... + def __getitem__( + self, key: int | np.integer | builtins.slice + ) -> Scalar | ChunkedArray: ... + __hash__: ClassVar[None] # type: ignore[assignment] + def is_null(self, *, nan_is_null: bool = False) -> ChunkedArray: ... + def is_nan(self) -> ChunkedArray: ... + def is_valid(self) -> ChunkedArray: ... + def cast( + self, + target_type: DataType | str | None = None, + safe: bool | None = None, + options: CastOptions | None = None, + ) -> ChunkedArray: ... + def fill_null(self, fill_value: Any) -> ChunkedArray: ... + def equals(self, other: ChunkedArray | None) -> bool: ... + def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray: ... + def __array__( + self, dtype: np.dtype | None = None, copy: bool | None = None + ) -> np.ndarray: ... + def dictionary_encode( + self, null_encoding: NullEncoding = "mask" + ) -> ChunkedArray: ... + def flatten(self, memory_pool: MemoryPool | None = None) -> list[ChunkedArray]: ... + def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Array: ... + def unique(self) -> Array: ... + def value_counts(self) -> StructArray: ... + def slice(self, offset: int = 0, length: int | None = None) -> ChunkedArray: ... + def filter( + self, mask: Mask, null_selection_behavior: NullSelectionBehavior = "drop" + ) -> ChunkedArray: ... + def index( + self, + value: Any, + start: int | None = None, + end: int | None = None, + *, + memory_pool: MemoryPool | None = None, + ) -> Int64Scalar: ... + def take(self, indices: Indices) -> ChunkedArray: ... + def drop_null(self) -> ChunkedArray: ... + def sort( + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + **kwargs: Any, + ) -> ChunkedArray: ... + def unify_dictionaries( + self, memory_pool: MemoryPool | None = None + ) -> ChunkedArray: ... + @property + def num_chunks(self) -> int: ... + def chunk(self, i: int) -> Array: ... + @property + def chunks(self) -> list[Array]: ... + def iterchunks(self) -> Iterator[Array]: ... + def __iter__(self) -> Iterator[Scalar]: ... + def to_pylist( + self, + *, + maps_as_pydicts: MapsAsPydicts | None = None, + ) -> list[Any]: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(stream: CapsuleType) -> ChunkedArray: ... + @property + def is_cpu(self) -> bool: ... + +def chunked_array( + arrays: _ChunkedArrayConstructorInput, + type: DataType | str | None = None, +) -> ChunkedArray: ... + +class _Tabular(_PandasConvertible): + def __array__( + self, dtype: np.dtype | None = None, copy: bool | None = None + ) -> np.ndarray: ... + def __dataframe__( + self, nan_as_null: bool = False, allow_copy: bool = True + ) -> _PyArrowDataFrame: ... + def __getitem__(self, key: int | str | slice) -> _ArrayOrChunkedArray | Self: ... + def __len__(self) -> int: ... + def column(self, i: int | str) -> _ArrayOrChunkedArray: ... + @property + def column_names(self) -> list[str]: ... + @property + def columns(self) -> list[Array] | list[ChunkedArray]: ... + def drop_null(self) -> Self: ... + def field(self, i: int | str) -> Field: ... + @classmethod + def from_pydict( + cls, + mapping: Mapping[ + Any, + _ArrayOrChunkedArray + | Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray, + ], + schema: Schema | None = None, + metadata: Mapping[str | bytes, str | bytes] | None = None, + ) -> Self: ... + @classmethod + def from_pylist( + cls, + mapping: Sequence[Mapping[str, Any]], + schema: Schema | None = None, + metadata: Mapping[str | bytes, str | bytes] | None = None, + ) -> Self: ... + def itercolumns(self) -> Iterator[_ArrayOrChunkedArray]: ... + @property + def num_columns(self) -> int: ... + @property + def num_rows(self) -> int: ... + @property + def shape(self) -> tuple[int, int]: ... + @property + def schema(self) -> Schema: ... + def sort_by( + self, + sorting: str | Sequence[tuple[str, Order] | tuple[str, Order, NullPlacement]], + **kwargs: Any, + ) -> Self: ... + def take(self, indices: Indices) -> Self: ... + def filter( + self, + mask: Mask | Expression, + null_selection_behavior: NullSelectionBehavior = "drop", + ) -> Self: ... + def to_pydict( + self, *, maps_as_pydicts: MapsAsPydicts | None = None + ) -> dict[str, list[Any]]: ... + def to_pylist( + self, *, maps_as_pydicts: MapsAsPydicts | None = None + ) -> list[dict[str, Any]]: ... + def to_string( + self, *, show_metadata: bool = False, preview_cols: int = 0 + ) -> str: ... + def remove_column(self, i: int) -> Self: ... + def drop_columns(self, columns: str | list[str]) -> Self: ... + def add_column( + self, + i: int, + field_: str | Field, + column: _ArrayOrChunkedArray | list[list[Any]], + ) -> Self: ... + def append_column( + self, field_: str | Field, column: _ArrayOrChunkedArray | list[list[Any]] + ) -> Self: ... + +class RecordBatch(_Tabular): + def validate(self, *, full: bool = False) -> None: ... + def replace_schema_metadata( + self, + metadata: dict[str, str] + | dict[bytes, bytes] + | dict[bytes, str] + | dict[str, bytes] + | None = None, + ) -> RecordBatch: ... + @property + def num_columns(self) -> int: ... + @property + def num_rows(self) -> int: ... + @property + def schema(self) -> Schema: ... + @property + def nbytes(self) -> int: ... + def get_total_buffer_size(self) -> int: ... + def __sizeof__(self) -> int: ... + def add_column( + self, i: int, field_: str | Field, column: _ArrayOrChunkedArray | list[Any] + ) -> RecordBatch: ... + def append_column( + self, field_: str | Field, column: _ArrayOrChunkedArray | list[Any] + ) -> RecordBatch: ... + def remove_column(self, i: int) -> RecordBatch: ... + def set_column( + self, i: int, field_: str | Field, column: Array | list + ) -> RecordBatch: ... + def rename_columns(self, names: list[str] | dict[str, str]) -> RecordBatch: ... + def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... + def slice(self, offset: int = 0, length: int | None = None) -> RecordBatch: ... + def equals(self, other: object, check_metadata: bool = False) -> bool: ... + def select( + self, columns: Iterable[str] | Iterable[int] | NDArray[np.str_] + ) -> RecordBatch: ... + def cast( + self, + target_schema: Schema, + safe: bool | None = None, + options: CastOptions | None = None, + ) -> RecordBatch: ... + @staticmethod + def from_arrays( + arrays: list[Any], + names: Sequence[str] | None = None, + schema: Schema | None = None, + metadata: Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] + | None = None, + ) -> RecordBatch: ... + @classmethod + def from_pandas( + cls, + df: pd.DataFrame, + schema: Schema | None = None, + preserve_index: bool | None = None, + nthreads: int | None = None, + columns: Iterable[str | int] | None = None, + ) -> RecordBatch: ... + @staticmethod + def from_struct_array(struct_array: StructArray) -> RecordBatch: ... + def to_struct_array(self) -> StructArray: ... + def to_tensor( + self, + null_to_nan: bool = False, + row_major: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Tensor: ... + def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int, schema: Schema | int) -> RecordBatch: ... + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> RecordBatch: ... + def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + @staticmethod + def _import_from_c_device(in_ptr: int, schema: Schema | int) -> RecordBatch: ... + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: ... + @staticmethod + def _import_from_c_device_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> RecordBatch: ... + @property + def device_type(self) -> DeviceAllocationType: ... + @property + def is_cpu(self) -> bool: ... + def copy_to(self, destination: MemoryManager | Device) -> RecordBatch: ... + +_JoinType: TypeAlias = Literal[ + "left semi", + "right semi", + "left anti", + "right anti", + "inner", + "left outer", + "right outer", + "full outer", +] + +class Table(_Tabular): + def validate(self, *, full: bool = False) -> None: ... + def slice(self, offset: int = 0, length: int | None = None) -> Table: ... + def select(self, columns: Iterable[str | int]) -> Table: ... + def replace_schema_metadata( + self, + metadata: dict[str, str] + | dict[bytes, bytes] + | dict[bytes, str] + | dict[str, bytes] + | None = None, + ) -> Table: ... + def flatten(self, memory_pool: MemoryPool | None = None) -> Table: ... + def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Table: ... + def unify_dictionaries(self, memory_pool: MemoryPool | None = None) -> Table: ... + def equals(self, other: Table | None, check_metadata: bool = False) -> bool: ... + def cast( + self, + target_schema: Schema, + safe: bool | None = None, + options: CastOptions | None = None, + ) -> Table: ... + @classmethod + def from_pandas( + cls, + df: pd.DataFrame, + schema: Schema | None = None, + preserve_index: bool | None = None, + nthreads: int | None = None, + columns: Iterable[str | int] | None = None, + safe: bool = True, + ) -> Table: ... + # When no schema, arrays should also support indexing + @staticmethod + def from_arrays( + arrays: Sequence[ + Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray + ] + | NDArray[Any], + names: Sequence[str] | None = None, + schema: Schema | None = None, + metadata: Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] + | None = None, + ) -> Table: ... + @staticmethod + def from_struct_array(struct_array: StructArray | ChunkedArray) -> Table: ... + def to_struct_array(self, max_chunksize: int | None = None) -> ChunkedArray: ... + @staticmethod + def from_batches( + batches: Iterable[RecordBatch], schema: Schema | None = None + ) -> Table: ... + def to_batches(self, max_chunksize: int | None = None) -> list[RecordBatch]: ... + def to_tensor( + self, + null_to_nan: bool = False, + row_major: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Tensor: ... + def to_reader(self, max_chunksize: int | None = None) -> RecordBatchReader: ... + @property + def schema(self) -> Schema: ... + @property + def num_columns(self) -> int: ... + @property + def num_rows(self) -> int: ... + @property + def nbytes(self) -> int: ... + def get_total_buffer_size(self) -> int: ... + def __sizeof__(self) -> int: ... + def add_column( + self, + i: int, + field_: str | Field, + column: ChunkedArray | _ChunkedArrayConstructorInput, + ) -> Table: ... + def append_column( + self, + field_: str | Field, + column: ChunkedArray | _ChunkedArrayConstructorInput, + ) -> Table: ... + def remove_column(self, i: int) -> Table: ... + def set_column( + self, + i: int, + field_: str | Field, + column: ChunkedArray | _ChunkedArrayConstructorInput, + ) -> Table: ... + def rename_columns(self, names: list[str] | dict[str, str]) -> Table: ... + def drop(self, columns: str | list[str]) -> Table: ... + def group_by( + self, keys: str | list[str], use_threads: bool = True + ) -> TableGroupBy: ... + def join( + self, + right_table: Table, + keys: str | list[str], + right_keys: str | list[str] | None = None, + join_type: _JoinType = "left outer", + left_suffix: str | None = None, + right_suffix: str | None = None, + coalesce_keys: bool = True, + use_threads: bool = True, + filter_expression: Expression | None = None, + ) -> Table: ... + def join_asof( + self, + right_table: Table, + on: str, + by: str | list[str], + tolerance: int, + right_on: str | list[str] | None = None, + right_by: str | list[str] | None = None, + ) -> Table: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... + @property + def is_cpu(self) -> bool: ... + +def record_batch( + data: dict[str, Any] + | list[Any] + | pd.DataFrame + | SupportsArrowArray + | SupportsArrowDeviceArray, + names: list[str] | Schema | None = None, + schema: Schema | None = None, + metadata: Mapping[str | bytes, str | bytes] | None = None, +) -> RecordBatch: ... +def table( + data: dict[str, Any] + | list[Any] + | tuple[Any, ...] + | pd.DataFrame + | SupportsArrowArray + | SupportsArrowStream + | SupportsArrowDeviceArray, + names: list[str] | Schema | None = None, + schema: Schema | None = None, + metadata: Mapping[str | bytes, str | bytes] | None = None, + nthreads: int | None = None, +) -> Table: ... +def concat_tables( + tables: Iterable[Table], + memory_pool: MemoryPool | None = None, + promote_options: Literal["none", "default", "permissive"] = "none", +) -> Table: ... + +class TableGroupBy: + keys: str | list[str] + + def __init__( + self, table: Table, keys: str | list[str], use_threads: bool = True + ): ... + def aggregate( + self, + aggregations: Iterable[ + tuple[_ColumnSelector, _AggregationName] + | tuple[_ColumnSelector, _AggregationName, _AggregateOptions | None] + ], + ) -> Table: ... + +def concat_batches( + recordbatches: Iterable[RecordBatch], memory_pool: MemoryPool | None = None +) -> RecordBatch: ... + +__all__ = [ + "ChunkedArray", + "chunked_array", + "_Tabular", + "RecordBatch", + "Table", + "record_batch", + "table", + "concat_tables", + "TableGroupBy", + "concat_batches", +] diff --git a/python/pyarrow-stubs/__lib/tensor.pyi b/python/pyarrow-stubs/__lib/tensor.pyi new file mode 100644 index 000000000000..281efe578595 --- /dev/null +++ b/python/pyarrow-stubs/__lib/tensor.pyi @@ -0,0 +1,245 @@ +# 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. + +from collections.abc import Iterable, Sequence +from typing import ClassVar + +import numpy as np +from pyarrow.lib import DataType, _Weakrefable +from scipy.sparse import ( + coo_array, + coo_matrix, + csc_array, + csc_matrix, + csr_array, + csr_matrix, +) +from sparse import COO +from typing_extensions import CapsuleType + +class Tensor(_Weakrefable): + @staticmethod + def from_numpy( + obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> Tensor: ... + def to_numpy(self) -> np.ndarray: ... + def equals(self, other: Tensor) -> bool: ... + def dim_name(self, i: int) -> str: ... + @property + def dim_names(self) -> list[str]: ... + @property + def is_mutable(self) -> bool: ... + @property + def is_contiguous(self) -> bool: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def shape(self) -> tuple[int, ...]: ... + @property + def strides(self) -> tuple[int, ...]: ... + @staticmethod + def from_dlpack( + x: object, + /, + *, + device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> Tensor: ... + def __dlpack__( + self, + *, + stream: int | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> CapsuleType: ... + def __dlpack_device__(self) -> tuple[int, int]: ... + @property + def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +class SparseCOOTensor(_Weakrefable): + @classmethod + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod + def from_numpy( + data: np.ndarray, + coords: np.ndarray, + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCOOTensor: ... + @staticmethod + def from_scipy( + obj: coo_array | coo_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod + def from_pydata_sparse( + obj: COO, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCOOTensor: ... + def to_numpy(self) -> tuple[np.ndarray, np.ndarray]: ... + def to_scipy(self) -> coo_array: ... + def to_pydata_sparse(self) -> COO: ... + def to_tensor(self) -> Tensor: ... + def equals(self, other: SparseCOOTensor) -> bool: ... + @property + def is_mutable(self) -> bool: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def shape(self) -> tuple[int, ...]: ... + def dim_name(self, i: int) -> str: ... + @property + def dim_names(self) -> tuple[str, ...]: ... + @property + def non_zero_length(self) -> int: ... + @property + def has_canonical_format(self) -> bool: ... + @property + def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +class SparseCSRMatrix(_Weakrefable): + @classmethod + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSRMatrix: ... + @staticmethod + def from_numpy( + data: np.ndarray, + indptr: np.ndarray, + indices: np.ndarray, + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCSRMatrix: ... + @staticmethod + def from_scipy( + obj: csr_array | csr_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCSRMatrix: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSRMatrix: ... + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + def to_scipy(self) -> csr_array: ... + def to_tensor(self) -> Tensor: ... + def equals(self, other: SparseCSRMatrix) -> bool: ... + @property + def is_mutable(self) -> bool: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def shape(self) -> tuple[int, ...]: ... + def dim_name(self, i: int) -> str: ... + @property + def dim_names(self) -> tuple[str, ...]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +class SparseCSCMatrix(_Weakrefable): + @classmethod + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSCMatrix: ... + @staticmethod + def from_numpy( + data: np.ndarray, + indptr: np.ndarray, + indices: np.ndarray, + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCSCMatrix: ... + @staticmethod + def from_scipy( + obj: csc_array | csc_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCSCMatrix: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSCMatrix: ... + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + def to_scipy(self) -> csc_array: ... + def to_tensor(self) -> Tensor: ... + def equals(self, other: SparseCSCMatrix) -> bool: ... + @property + def is_mutable(self) -> bool: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def shape(self) -> tuple[int, ...]: ... + def dim_name(self, i: int) -> str: ... + @property + def dim_names(self) -> tuple[str, ...]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +class SparseCSFTensor(_Weakrefable): + @classmethod + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSFTensor: ... + @staticmethod + def from_numpy( + data: np.ndarray, + indptr: list[np.ndarray] | tuple[np.ndarray, ...], + indices: list[np.ndarray] | tuple[np.ndarray, ...], + shape: Sequence[int], + axis_order: Iterable[int] | None = None, + dim_names: Iterable[str] | None = None, + ) -> SparseCSFTensor: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSFTensor: ... + def to_numpy(self) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]: ... + def to_tensor(self) -> Tensor: ... + def equals(self, other: SparseCSFTensor) -> bool: ... + @property + def is_mutable(self) -> bool: ... + @property + def ndim(self) -> int: ... + @property + def size(self) -> int: ... + @property + def shape(self) -> tuple[int, ...]: ... + def dim_name(self, i: int) -> str: ... + @property + def dim_names(self) -> tuple[str, ...]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +__all__ = [ + "Tensor", + "SparseCOOTensor", + "SparseCSRMatrix", + "SparseCSCMatrix", + "SparseCSFTensor", +] diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi new file mode 100644 index 000000000000..d3f98a94c585 --- /dev/null +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -0,0 +1,662 @@ +# 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. + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, ClassVar, Literal, Self, TypeAlias, overload + +import numpy as np +import pandas as pd +from pyarrow._stubs_typing import ( + BinaryType, + BinaryViewType, + BoolType, + Date32Type, + Date64Type, + Float16Type, + Float32Type, + Float64Type, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + IntegerType, + LargeBinaryType, + LargeStringType, + MonthDayNanoIntervalType, + NullType, + StringType, + StringViewType, + SupportsArrowSchema, + TimeUnit, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, +) +from pyarrow.lib import ( + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) +from typing_extensions import CapsuleType + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field: ... + @property + def id(self) -> int: ... + @property + def bit_width(self) -> int: ... + @property + def byte_width(self) -> int: ... + @property + def num_fields(self) -> int: ... + @property + def num_buffers(self) -> int: ... + @property + def has_variadic_buffers(self) -> bool: ... + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... + def to_pandas_dtype(self) -> _PandasDtype: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +_Time32Unit: TypeAlias = Literal["s", "ms"] +_Time64Unit: TypeAlias = Literal["us", "ns"] + +class TimestampType(DataType): + @property + def unit(self) -> TimeUnit: ... + @property + def tz(self) -> str | None: ... + +class Time32Type(DataType): + @property + def unit(self) -> _Time32Unit: ... + +class Time64Type(DataType): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(DataType): + @property + def unit(self) -> TimeUnit: ... + +class FixedSizeBinaryType(DataType): ... + +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal128Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... + +class ListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class ListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class LargeListViewType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + +class FixedSizeListType(DataType): + @property + def value_field(self) -> Field: ... + @property + def value_type(self) -> DataType: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> IntegerType: ... + @property + def value_type(self) -> DataType: ... + +class MapType(DataType): + @property + def key_field(self) -> Field: ... + @property + def key_type(self) -> DataType: ... + @property + def item_field(self) -> Field: ... + @property + def item_type(self) -> DataType: ... + @property + def keys_sorted(self) -> bool: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field]: ... + +class UnionType(DataType): + @property + def mode(self) -> Literal["sparse", "dense"]: ... + @property + def type_codes(self) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type + +class RunEndEncodedType(DataType): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> DataType: ... + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... + @property + def extension_name(self) -> str: ... + @property + def storage_type(self) -> DataType: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... + +class ExtensionType(BaseExtensionType): + def __init__(self, storage_type: DataType, extension_name: str) -> None: ... + def __arrow_ext_serialize__(self) -> bytes: ... + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... + __hash__: ClassVar[None] # type: ignore[assignment] + +class FixedShapeTensorType(BaseExtensionType): + @property + def value_type(self) -> DataType: ... + @property + def shape(self) -> list[int]: ... + @property + def dim_names(self) -> list[str] | None: ... + @property + def permutation(self) -> list[int] | None: ... + +class Bool8Type(BaseExtensionType): ... +class UuidType(BaseExtensionType): ... +class JsonType(BaseExtensionType): ... + +class OpaqueType(BaseExtensionType): + @property + def type_name(self) -> str: ... + @property + def vendor_name(self) -> str: ... + +class UnknownExtensionType(ExtensionType): + def __init__(self, storage_type: DataType, serialized: bytes) -> None: ... + +def register_extension_type(ext_type: BaseExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_FieldName: TypeAlias = str | bytes +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable +_DataTypeAlias: TypeAlias = Literal[ + "null", + "bool", + "boolean", + "i1", + "int8", + "i2", + "int16", + "i4", + "int32", + "i8", + "int64", + "u1", + "uint8", + "u2", + "uint16", + "u4", + "uint32", + "u8", + "uint64", + "f2", + "halffloat", + "float16", + "f4", + "float", + "float32", + "f8", + "double", + "float64", + "string", + "str", + "utf8", + "binary", + "large_string", + "large_str", + "large_utf8", + "large_binary", + "binary_view", + "string_view", + "date32", + "date64", + "date32[day]", + "date64[ms]", + "time32[s]", + "time32[ms]", + "time64[us]", + "time64[ns]", + "timestamp[s]", + "timestamp[ms]", + "timestamp[us]", + "timestamp[ns]", + "duration[s]", + "duration[ms]", + "duration[us]", + "duration[ns]", + "month_day_nano_interval", +] +_DataTypeAliasInput: TypeAlias = _DataTypeAlias | str +_DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata | None = None, + **kwargs: str | bytes, + ) -> None: ... + __hash__: ClassVar[None] # type: ignore[assignment] + def equals(self, other: KeyValueMetadata) -> bool: ... + def __len__(self) -> int: ... + def __contains__(self, key: object, /) -> bool: ... + def __getitem__(self, key: str | bytes, /) -> bytes: ... + def __iter__(self) -> Iterator[bytes]: ... + def key(self, i: int) -> bytes: ... + def value(self, i: int) -> bytes: ... + # The runtime yields generators instead of Mapping views + def keys(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def values(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def items(self) -> Iterator[tuple[bytes, bytes]]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def get_all(self, key: str | bytes) -> list[bytes]: ... + def to_dict(self) -> dict[bytes, bytes]: ... + +class Field(_Weakrefable): + def equals(self, other: Field, check_metadata: bool = False) -> bool: ... + @property + def nullable(self) -> bool: ... + @property + def name(self) -> str: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + @property + def type(self) -> DataType: ... + def with_metadata( + self, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None, + ) -> Field: ... + def remove_metadata(self) -> Field: ... + def with_type(self, new_type: DataType) -> Field: ... + def with_name(self, name: str | bytes) -> Field: ... + def with_nullable(self, nullable: bool) -> Field: ... + def flatten(self) -> list[Field]: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Field: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Field: ... + +_FieldTuple: TypeAlias = ( + tuple[_FieldName, _DataTypeLike] + | tuple[_FieldName, _DataTypeLike, bool | None] + | tuple[ + _FieldName, + _DataTypeLike, + bool | None, + _KeyValueMetadataInput | KeyValueMetadata | None, + ] +) +_StructFieldsInput: TypeAlias = ( + Iterable[Field | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] +) +_SchemaFieldsInput: TypeAlias = _StructFieldsInput + +class Schema(_Weakrefable): + def __len__(self) -> int: ... + def __getitem__(self, key: int, /) -> Field: ... + + _field = __getitem__ + def __iter__(self) -> Iterator[Field]: ... + def __sizeof__(self) -> int: ... + @property + def pandas_metadata(self) -> dict[str, Any] | None: ... + @property + def names(self) -> list[str]: ... + @property + def types(self) -> list[DataType]: ... + @property + def metadata(self) -> dict[bytes, bytes] | None: ... + def empty_table(self) -> Table: ... + def equals(self, other: Schema, check_metadata: bool = False) -> bool: ... + @classmethod + def from_pandas( + cls, df: pd.DataFrame, preserve_index: bool | None = None + ) -> Schema: ... + def field(self, i: int | str | bytes) -> Field: ... + def get_field_index(self, name: str | bytes) -> int: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def append(self, field: Field) -> Schema: ... + def insert(self, i: int, field: Field) -> Schema: ... + def remove(self, i: int) -> Schema: ... + def set(self, i: int, field: Field) -> Schema: ... + def with_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata | None + ) -> Schema: ... + def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... + def remove_metadata(self) -> Schema: ... + def to_string( + self, + truncate_metadata: bool = True, + show_field_metadata: bool = True, + show_schema_metadata: bool = True, + element_size_limit: int = 100, + ) -> str: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Schema: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Schema: ... + +def unify_schemas( + schemas: Iterable[Schema], + *, + promote_options: Literal["default", "permissive"] = "default", +) -> Schema: ... +@overload +def field( + name: SupportsArrowSchema, + type: None = None, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, +) -> Field: ... +@overload +def field( + name: _FieldName, + type: _DataTypeLike, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, +) -> Field: ... +def null() -> NullType: ... +def bool_() -> BoolType: ... +def uint8() -> UInt8Type: ... +def int8() -> Int8Type: ... +def uint16() -> UInt16Type: ... +def int16() -> Int16Type: ... +def uint32() -> UInt32Type: ... +def int32() -> Int32Type: ... +def int64() -> Int64Type: ... +def uint64() -> UInt64Type: ... +def timestamp(unit: TimeUnit | str, tz: str | None = None) -> TimestampType: ... +def time32(unit: _Time32Unit | str) -> Time32Type: ... +def time64(unit: _Time64Unit | str) -> Time64Type: ... +def duration(unit: TimeUnit | str) -> DurationType: ... +def month_day_nano_interval() -> MonthDayNanoIntervalType: ... +def date32() -> Date32Type: ... +def date64() -> Date64Type: ... +def float16() -> Float16Type: ... +def float32() -> Float32Type: ... +def float64() -> Float64Type: ... +def decimal32(precision: int, scale: int = 0) -> Decimal32Type: ... +def decimal64(precision: int, scale: int = 0) -> Decimal64Type: ... +def decimal128(precision: int, scale: int = 0) -> Decimal128Type: ... +def decimal256(precision: int, scale: int = 0) -> Decimal256Type: ... +def string() -> StringType: ... + +utf8 = string + +def binary( + length: Literal[-1] | int = ..., +) -> BinaryType | FixedSizeBinaryType: ... +def large_binary() -> LargeBinaryType: ... +def large_string() -> LargeStringType: ... + +large_utf8 = large_string + +def binary_view() -> BinaryViewType: ... +def string_view() -> StringViewType: ... +@overload +def list_(value_type: DataType | Field, list_size: Literal[-1] = -1) -> ListType: ... +@overload +def list_( + value_type: DataType | Field, list_size: int +) -> ListType | FixedSizeListType: ... +def large_list(value_type: DataType | Field) -> LargeListType: ... +def list_view(value_type: DataType | Field) -> ListViewType: ... +def large_list_view(value_type: DataType | Field) -> LargeListViewType: ... +def map_( + key_type: _DataTypeLike | Field, + item_type: _DataTypeLike | Field, + keys_sorted: bool = False, +) -> MapType: ... +def dictionary( + index_type: IntegerType | _DataTypeAliasInput, + value_type: _DataTypeLike, + ordered: bool = False, +) -> DictionaryType: ... +def struct( + fields: _StructFieldsInput, +) -> StructType: ... +def sparse_union( + child_fields: Iterable[Field], type_codes: Sequence[int] | None = None +) -> SparseUnionType: ... +def dense_union( + child_fields: Iterable[Field], type_codes: Sequence[int] | None = None +) -> DenseUnionType: ... +def union( + child_fields: Iterable[Field], + mode: Literal["sparse", "dense"] | int, + type_codes: Sequence[int] | None = None, +) -> SparseUnionType | DenseUnionType: ... +def run_end_encoded( + run_end_type: _RunEndType | _DataTypeAliasInput, + value_type: _DataTypeLike, +) -> RunEndEncodedType: ... +def json_(storage_type: DataType = ...) -> JsonType: ... +def uuid() -> UuidType: ... +def fixed_shape_tensor( + value_type: DataType, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + permutation: Sequence[int] | None = None, +) -> FixedShapeTensorType: ... +def bool8() -> Bool8Type: ... +def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... +def type_for_alias(name: str) -> DataType: ... +def schema( + fields: SupportsArrowSchema | _SchemaFieldsInput, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, +) -> Schema: ... +def from_numpy_dtype(dtype: np.typing.DTypeLike) -> DataType: ... + +__all__ = [ + "_Weakrefable", + "_Metadata", + "DataType", + "TimestampType", + "Time32Type", + "Time64Type", + "DurationType", + "FixedSizeBinaryType", + "Decimal32Type", + "Decimal64Type", + "Decimal128Type", + "Decimal256Type", + "ListType", + "LargeListType", + "ListViewType", + "LargeListViewType", + "FixedSizeListType", + "DictionaryMemo", + "DictionaryType", + "MapType", + "StructType", + "UnionType", + "SparseUnionType", + "DenseUnionType", + "RunEndEncodedType", + "BaseExtensionType", + "ExtensionType", + "FixedShapeTensorType", + "Bool8Type", + "UuidType", + "JsonType", + "OpaqueType", + "UnknownExtensionType", + "register_extension_type", + "unregister_extension_type", + "KeyValueMetadata", + "Field", + "Schema", + "unify_schemas", + "field", + "null", + "bool_", + "uint8", + "int8", + "uint16", + "int16", + "uint32", + "int32", + "int64", + "uint64", + "timestamp", + "time32", + "time64", + "duration", + "month_day_nano_interval", + "date32", + "date64", + "float16", + "float32", + "float64", + "decimal32", + "decimal64", + "decimal128", + "decimal256", + "string", + "utf8", + "binary", + "large_binary", + "large_string", + "large_utf8", + "binary_view", + "string_view", + "list_", + "large_list", + "list_view", + "large_list_view", + "map_", + "dictionary", + "struct", + "sparse_union", + "dense_union", + "union", + "run_end_encoded", + "json_", + "uuid", + "fixed_shape_tensor", + "bool8", + "opaque", + "type_for_alias", + "schema", + "from_numpy_dtype", +] diff --git a/python/pyarrow-stubs/_compute.pyi b/python/pyarrow-stubs/_compute.pyi new file mode 100644 index 000000000000..3108cb063747 --- /dev/null +++ b/python/pyarrow-stubs/_compute.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 + +Expression: TypeAlias = Any + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi new file mode 100644 index 000000000000..e7b58eb63f75 --- /dev/null +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -0,0 +1,93 @@ +# 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. + +from collections.abc import Sequence +from typing import Any, Literal, Protocol, TypeAlias + +import numpy as np +from numpy.typing import NDArray +from pyarrow.lib import Array, ChunkedArray + +from pyarrow import lib + +ArrayLike: TypeAlias = Any +Order: TypeAlias = Literal["ascending", "descending"] +MapsAsPydicts: TypeAlias = Literal["lossy", "strict"] +NullEncoding: TypeAlias = Literal["mask", "encode"] +NullPlacement: TypeAlias = Literal["at_start", "at_end"] +NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] +TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] + +Mask: TypeAlias = Sequence[bool | None] | NDArray[np.bool_] | Array | ChunkedArray +Indices: TypeAlias = ( + Sequence[int | None] | NDArray[np.integer[Any]] | Array | ChunkedArray +) + +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: ... + +# Data types without a class of their own at runtime: their factories, e.g. +# pyarrow.int32(), return plain DataType instances. These classes only exist +# for type checking, where they distinguish the results of the factories and +# the narrowing of the pyarrow.types.is_* checks. +class NullType(lib.DataType): ... +class BoolType(lib.DataType): ... +class UInt8Type(lib.DataType): ... +class Int8Type(lib.DataType): ... +class UInt16Type(lib.DataType): ... +class Int16Type(lib.DataType): ... +class UInt32Type(lib.DataType): ... +class Int32Type(lib.DataType): ... +class UInt64Type(lib.DataType): ... +class Int64Type(lib.DataType): ... +class Float16Type(lib.DataType): ... +class Float32Type(lib.DataType): ... +class Float64Type(lib.DataType): ... +class Date32Type(lib.DataType): ... +class Date64Type(lib.DataType): ... +class MonthDayNanoIntervalType(lib.DataType): ... +class StringType(lib.DataType): ... +class LargeStringType(lib.DataType): ... +class StringViewType(lib.DataType): ... +class BinaryType(lib.DataType): ... +class LargeBinaryType(lib.DataType): ... +class BinaryViewType(lib.DataType): ... + +IntegerType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi new file mode 100644 index 000000000000..c585531eed4a --- /dev/null +++ b/python/pyarrow-stubs/lib.pyi @@ -0,0 +1,44 @@ +# 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. + +# Stub for the pyarrow.lib C extension module. Re-exports the modules that +# have been annotated so far; everything else falls back to Any via +# __getattr__ until it gets annotated in subsequent PRs. + +from typing import Any, Self, final + +from pyarrow.__lib.array import * # noqa: F401,F403 +from pyarrow.__lib.builder import * # noqa: F401,F403 +from pyarrow.__lib.config import * # noqa: F401,F403 +from pyarrow.__lib.device import * # noqa: F401,F403 +from pyarrow.__lib.error import * # noqa: F401,F403 +from pyarrow.__lib.io import * # noqa: F401,F403 +from pyarrow.__lib.memory import * # noqa: F401,F403 +from pyarrow.__lib.scalar import * # noqa: F401,F403 +from pyarrow.__lib.table import * # noqa: F401,F403 +from pyarrow.__lib.tensor import * # noqa: F401,F403 +from pyarrow.__lib.types import * # noqa: F401,F403 + +@final +class MonthDayNano(tuple): + months: int + days: int + nanoseconds: int + + def __new__(cls, sequence: tuple[int, int, int] | list[int]) -> Self: ... + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/types.pyi b/python/pyarrow-stubs/types.pyi new file mode 100644 index 000000000000..8f8773b08d56 --- /dev/null +++ b/python/pyarrow-stubs/types.pyi @@ -0,0 +1,200 @@ +# 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 sys +from enum import IntEnum +from typing import Any, TypeAlias, TypeGuard + +if sys.version_info >= (3, 13): + from typing import TypeIs +else: + from typing_extensions import TypeIs + +from pyarrow._stubs_typing import ( + BinaryType, + BinaryViewType, + BoolType, + Date32Type, + Date64Type, + Float16Type, + Float32Type, + Float64Type, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + IntegerType, + LargeBinaryType, + LargeStringType, + MonthDayNanoIntervalType, + NullType, + StringType, + StringViewType, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, +) +from pyarrow.lib import ( + DataType, + Decimal32Type, + Decimal64Type, + Decimal128Type, + Decimal256Type, + DenseUnionType, + DictionaryType, + DurationType, + FixedSizeBinaryType, + FixedSizeListType, + LargeListType, + LargeListViewType, + ListType, + ListViewType, + MapType, + RunEndEncodedType, + SparseUnionType, + StructType, + Time32Type, + Time64Type, + TimestampType, +) + +_SignedInteger: TypeAlias = Int8Type | Int16Type | Int32Type | Int64Type +_UnsignedInteger: TypeAlias = UInt8Type | UInt16Type | UInt32Type | UInt64Type +_Floating: TypeAlias = Float16Type | Float32Type | Float64Type +_Decimal: TypeAlias = Decimal32Type | Decimal64Type | Decimal128Type | Decimal256Type +_Date: TypeAlias = Date32Type | Date64Type +_Time: TypeAlias = Time32Type | Time64Type +_Interval: TypeAlias = MonthDayNanoIntervalType +_Temporal: TypeAlias = TimestampType | DurationType | _Time | _Date | _Interval +_Union: TypeAlias = SparseUnionType | DenseUnionType +_Nested: TypeAlias = ( + ListType + | FixedSizeListType + | LargeListType + | ListViewType + | LargeListViewType + | StructType + | MapType + | _Union + | RunEndEncodedType +) + +def is_null(t: DataType) -> TypeIs[NullType]: ... +def is_boolean(t: DataType) -> TypeIs[BoolType]: ... +def is_integer(t: DataType) -> TypeIs[IntegerType]: ... +def is_signed_integer(t: DataType) -> TypeIs[_SignedInteger]: ... +def is_unsigned_integer(t: DataType) -> TypeIs[_UnsignedInteger]: ... +def is_int8(t: DataType) -> TypeIs[Int8Type]: ... +def is_int16(t: DataType) -> TypeIs[Int16Type]: ... +def is_int32(t: DataType) -> TypeIs[Int32Type]: ... +def is_int64(t: DataType) -> TypeIs[Int64Type]: ... +def is_uint8(t: DataType) -> TypeIs[UInt8Type]: ... +def is_uint16(t: DataType) -> TypeIs[UInt16Type]: ... +def is_uint32(t: DataType) -> TypeIs[UInt32Type]: ... +def is_uint64(t: DataType) -> TypeIs[UInt64Type]: ... +def is_floating(t: DataType) -> TypeIs[_Floating]: ... +def is_float16(t: DataType) -> TypeIs[Float16Type]: ... +def is_float32(t: DataType) -> TypeIs[Float32Type]: ... +def is_float64(t: DataType) -> TypeIs[Float64Type]: ... +def is_list(t: DataType) -> TypeIs[ListType]: ... +def is_large_list(t: DataType) -> TypeIs[LargeListType]: ... +def is_fixed_size_list(t: DataType) -> TypeIs[FixedSizeListType]: ... +def is_list_view(t: DataType) -> TypeIs[ListViewType]: ... +def is_large_list_view(t: DataType) -> TypeIs[LargeListViewType]: ... +def is_struct(t: DataType) -> TypeIs[StructType]: ... +def is_union(t: DataType) -> TypeIs[_Union]: ... +def is_nested(t: DataType) -> TypeIs[_Nested]: ... +def is_run_end_encoded(t: DataType) -> TypeIs[RunEndEncodedType]: ... +def is_temporal(t: DataType) -> TypeIs[_Temporal]: ... +def is_timestamp(t: DataType) -> TypeIs[TimestampType]: ... +def is_duration(t: DataType) -> TypeIs[DurationType]: ... +def is_time(t: DataType) -> TypeIs[_Time]: ... +def is_time32(t: DataType) -> TypeIs[Time32Type]: ... +def is_time64(t: DataType) -> TypeIs[Time64Type]: ... +def is_binary(t: DataType) -> TypeIs[BinaryType]: ... +def is_large_binary(t: DataType) -> TypeIs[LargeBinaryType]: ... +def is_unicode(t: DataType) -> TypeIs[StringType]: ... +def is_string(t: DataType) -> TypeIs[StringType]: ... +def is_large_unicode(t: DataType) -> TypeIs[LargeStringType]: ... +def is_large_string(t: DataType) -> TypeIs[LargeStringType]: ... + +# TypeGuard, as decimal types subclass FixedSizeBinaryType but aren't matched. +def is_fixed_size_binary(t: DataType) -> TypeGuard[FixedSizeBinaryType]: ... +def is_binary_view(t: DataType) -> TypeIs[BinaryViewType]: ... +def is_string_view(t: DataType) -> TypeIs[StringViewType]: ... +def is_date(t: DataType) -> TypeIs[_Date]: ... +def is_date32(t: DataType) -> TypeIs[Date32Type]: ... +def is_date64(t: DataType) -> TypeIs[Date64Type]: ... +def is_map(t: DataType) -> TypeIs[MapType]: ... +def is_decimal(t: DataType) -> TypeIs[_Decimal]: ... +def is_decimal32(t: DataType) -> TypeIs[Decimal32Type]: ... +def is_decimal64(t: DataType) -> TypeIs[Decimal64Type]: ... +def is_decimal128(t: DataType) -> TypeIs[Decimal128Type]: ... +def is_decimal256(t: DataType) -> TypeIs[Decimal256Type]: ... +def is_dictionary(t: DataType) -> TypeIs[DictionaryType]: ... +def is_interval(t: DataType) -> TypeIs[_Interval]: ... +def is_primitive(t: DataType) -> bool: ... +def is_boolean_value(obj: Any) -> bool: ... +def is_integer_value(obj: Any) -> bool: ... +def is_float_value(obj: Any) -> bool: ... + +class TypesEnum(IntEnum): + NA = 0 + BOOL = 1 + INT8 = 3 + INT16 = 5 + INT32 = 7 + INT64 = 9 + UINT8 = 2 + UINT16 = 4 + UINT32 = 6 + UINT64 = 8 + HALF_FLOAT = 10 + FLOAT = 11 + DOUBLE = 12 + BINARY = 14 + BINARY_VIEW = 40 + LARGE_BINARY = 35 + STRING = 13 + STRING_VIEW = 39 + LARGE_STRING = 34 + FIXED_SIZE_BINARY = 15 + DECIMAL32 = 43 + DECIMAL64 = 44 + DECIMAL128 = 23 + DECIMAL256 = 24 + LIST = 25 + LARGE_LIST = 36 + LIST_VIEW = 41 + LARGE_LIST_VIEW = 42 + MAP = 30 + FIXED_SIZE_LIST = 32 + STRUCT = 26 + SPARSE_UNION = 27 + DENSE_UNION = 28 + RUN_END_ENCODED = 38 + DATE32 = 16 + DATE64 = 17 + TIME32 = 19 + TIME64 = 20 + TIMESTAMP = 18 + INTERVAL_MONTHS = 21 + INTERVAL_DAY_TIME = 22 + INTERVAL_MONTH_DAY_NANO = 37 + DURATION = 33 + DICTIONARY = 29 diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index e090b1d0b3db..702f9eff94e7 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -3936,7 +3936,7 @@ cdef class FixedSizeListArray(BaseListArray): Or create from a values array, list size and matching type: >>> typ = pa.list_(pa.field("values", pa.int64()), 2) - >>> arr = pa.FixedSizeListArray.from_arrays(values,type=typ) + >>> arr = pa.FixedSizeListArray.from_arrays(values, type=typ) >>> arr [ diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index f96fb40572b9..ad562a0dbd5c 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -1108,7 +1108,7 @@ cdef class StructScalar(Scalar, Mapping): Parameters ---------- - index : Union[int, str] + key : Union[int, str] Index / position or name of the field. Returns diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi index 7f006f294175..5ddcb02543b4 100644 --- a/python/pyarrow/tensor.pxi +++ b/python/pyarrow/tensor.pxi @@ -1376,14 +1376,14 @@ shape: {self.shape}""" ---------- data : numpy.ndarray Data used to populate the sparse tensor. - indptr : numpy.ndarray + indptr : list or tuple of numpy.ndarray The sparsity structure. Each two consecutive dimensions in a tensor correspond to a buffer in indices. A pair of consecutive values at `indptr[dim][i]` `indptr[dim][i + 1]` signify a range of nodes in `indices[dim + 1]` who are children of `indices[dim][i]` node. - indices : numpy.ndarray + indices : list or tuple of numpy.ndarray Stores values of nodes. Each tensor dimension corresponds to a buffer in indptr. shape : tuple diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 63c0eb0bea4e..7619e65109b7 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -31,7 +31,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.tests.strategies as past @@ -323,11 +323,11 @@ def test_asarray(): arr = pa.array(range(4)) - # The iterator interface gives back an array of Int64Value's + # The iterator interface gives back an array of Int64Scalar's np_arr = np.asarray([_ for _ in arr]) assert np_arr.tolist() == [0, 1, 2, 3] assert np_arr.dtype == np.dtype('O') - assert isinstance(np_arr[0], pa.lib.Int64Value) + assert isinstance(np_arr[0], pa.Int64Scalar) # Calling with the arrow array gives back an array with 'int64' dtype np_arr = np.asarray(arr) @@ -3188,7 +3188,7 @@ def test_buffers_primitive(): # Slicing does not affect the buffers but the offset a_sliced = a[1:] buffers = a_sliced.buffers() - a_sliced.offset == 1 + assert a_sliced.offset == 1 assert len(buffers) == 2 null_bitmap = buffers[0].to_pybytes() assert 1 <= len(null_bitmap) <= 64 # XXX this is varying diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index 3ac7be88eaa0..72b418715b4e 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -27,7 +27,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] from pyarrow.pandas_compat import _pandas_api # noqa import pyarrow as pa diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index 01bf3c31a09b..0ed3cec1408c 100644 --- a/python/pyarrow/tests/test_schema.py +++ b/python/pyarrow/tests/test_schema.py @@ -23,7 +23,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.tests.util as test_util diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index eca8090d77a9..a5b031ba4658 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -28,10 +28,10 @@ try: from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix except ImportError: - coo_matrix = None - csr_matrix = None - csr_array = None - coo_array = None + coo_matrix = None # type: ignore[assignment, misc] + csr_matrix = None # type: ignore[assignment, misc] + csr_array = None # type: ignore[assignment, misc] + coo_array = None # type: ignore[assignment, misc] try: import sparse @@ -401,7 +401,7 @@ def test_dense_to_sparse_tensor(dtype_str, arrow_type, sparse_tensor_type): assert np.array_equal(array, result_array) -@pytest.mark.skipif(not coo_matrix, reason="requires scipy") +@pytest.mark.skipif(coo_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (coo_array, coo_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, @@ -443,7 +443,7 @@ def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, assert out_scipy_matrix.has_canonical_format -@pytest.mark.skipif(not csr_matrix, reason="requires scipy") +@pytest.mark.skipif(csr_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (csr_array, csr_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_csr_matrix_scipy_roundtrip(dtype_str, arrow_type, diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index 0fefabf02abe..9a2a99b4e535 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -23,7 +23,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pytest import pyarrow as pa import pyarrow.compute as pc diff --git a/python/pyarrow/tests/test_types.py b/python/pyarrow/tests/test_types.py index 4f251d143343..064eeb6e7a0e 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -28,13 +28,13 @@ try: import hypothesis.extra.pytz as tzst except ImportError: - tzst = None + tzst = None # type: ignore[assignment] import weakref try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.types as types import pyarrow.tests.strategies as past @@ -415,7 +415,7 @@ def test_tzinfo_to_string_errors(): if tzst: timezones = tzst.timezones() else: - timezones = st.none() + timezones = st.none() # type: ignore[assignment] @h.given(timezones) @@ -752,7 +752,7 @@ def test_struct_type(): ty.field(None) for a, b in zip(ty, fields): - a == b + assert a == b # Construct from list of tuples ty = pa.struct([('a', pa.int64()), @@ -760,7 +760,7 @@ def test_struct_type(): ('b', pa.int32())]) assert list(ty) == fields for a, b in zip(ty, fields): - a == b + assert a == b # Construct from mapping fields = [pa.field('a', pa.int64()), @@ -769,7 +769,7 @@ def test_struct_type(): ('b', pa.int32())])) assert list(ty) == fields for a, b in zip(ty, fields): - a == b + assert a == b # Invalid args with pytest.raises(TypeError): @@ -1098,12 +1098,14 @@ def test_timedelta_overflow(): pa.scalar(d, type=pa.duration('ns')) # microsecond resolution, not overflow - pa.scalar(d, type=pa.duration('us')).as_py() == d + assert pa.scalar(d, type=pa.duration('us')).as_py() == d - # second/millisecond resolution, not overflow + # second/millisecond resolution, not overflow (sub-unit part is truncated) for d in [datetime.timedelta.min, datetime.timedelta.max]: - pa.scalar(d, type=pa.duration('ms')).as_py() == d - pa.scalar(d, type=pa.duration('s')).as_py() == d + ms = datetime.timedelta(milliseconds=1) + assert pa.scalar(d, type=pa.duration('ms')).as_py() == d // ms * ms + s = datetime.timedelta(seconds=1) + assert pa.scalar(d, type=pa.duration('s')).as_py() == d // s * s def test_type_equality_operators(): diff --git a/python/pyproject.toml b/python/pyproject.toml index 53b6abb4e02a..db8aba207f39 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -111,41 +111,78 @@ version_scheme = 'guess-next-dev' git_describe_command = 'git describe --dirty --tags --long --match "apache-arrow-[0-9]*.*"' fallback_version = '27.0.0a0' -# TODO: Enable type checking once stubs are merged [tool.mypy] -files = ["pyarrow-stubs"] -mypy_path = "$MYPY_CONFIG_FILE_DIR/pyarrow-stubs" -exclude = [ - "^pyarrow/", - "^benchmarks/", - "^examples/", - "^scripts/", -] +# mypy can't find a "-stubs" package on mypy_path, but it maps pyarrow-stubs/ +# to the pyarrow module when checking it. Pass pyarrow-stubs as well when +# checking other files, e.g. `mypy pyarrow-stubs pyarrow/tests/test_array.py`. +files = ["pyarrow-stubs", "pyarrow-stubs-tests"] +# The type annotation tests use `# type: ignore[...]` to assert that some +# usage is rejected, so an ignore that is no longer needed is an error there. +warn_unused_ignores = true + +# Modules that will be annotated in later PRs (GH-32609). Treat them as Any +# rather than type checking their sources. +[[tool.mypy.overrides]] +module = ["pyarrow._compute", "pyarrow.compute", "pyarrow.interchange.*"] +follow_imports = "skip" + +# Optional dependencies without type hints. +[[tool.mypy.overrides]] +module = ["fsspec", "sparse"] +ignore_missing_imports = true -# TODO: Enable type checking once stubs are merged [tool.pyright] pythonPlatform = "All" pythonVersion = "3.11" -include = ["pyarrow-stubs"] -exclude = [ - "pyarrow", - "benchmarks", - "examples", - "scripts", - "build", -] -stubPath = "pyarrow-stubs" +include = ["pyarrow-stubs", "pyarrow-stubs-tests"] typeCheckingMode = "basic" +# The type annotation tests use `# type: ignore[...]` to assert that some +# usage is rejected, so an ignore that is no longer needed is an error there. +executionEnvironments = [ + { root = "pyarrow-stubs-tests", extraPaths = ["."], reportUnnecessaryTypeIgnoreComment = "error" }, + # The stubs import modules that will be annotated in later PRs (GH-32609), + # which only resolve when PyArrow is installed, and the optional sparse + # dependency. mypy reports the other imports that can't be resolved. + { root = ".", reportMissingImports = false }, +] +# Don't warn about imports that resolve to a stub without a Python source: +# pyarrow._stubs_typing only exists as a stub, and compiled modules such as +# pyarrow._compute have no source that pyright can read. +reportMissingModuleSource = false -# TODO: Enable type checking once stubs are merged -[tool.ty.environment] -extra-paths = ["pyarrow-stubs"] +[tool.ty.analysis] +# Modules that will be annotated in later PRs (GH-32609), and the optional +# sparse dependency, which has no type hints. Treat them as Any. +replace-imports-with-any = [ + "pyarrow._compute", + "pyarrow.compute", + "pyarrow.interchange.**", + "sparse", +] [tool.ty.src] -include = ["pyarrow-stubs"] -exclude = [ - "pyarrow", - "benchmarks", - "examples", - "scripts", -] +include = ["pyarrow-stubs", "pyarrow-stubs-tests"] + +# ty doesn't honor `# type: ignore[]` comments, so disable the +# rules that those comments suppress for mypy. +[[tool.ty.overrides]] +include = ["pyarrow-stubs/__lib/array.pyi"] +[tool.ty.overrides.rules] +# Some array classes override methods with a different signature at runtime. +invalid-method-override = "ignore" + +[[tool.ty.overrides]] +include = ["pyarrow-stubs-tests/**"] +[tool.ty.overrides.rules] +# mypy and pyright check that the invalid usage in the tests is rejected. +invalid-argument-type = "ignore" + +# Lints and formats the type stubs (*.pyi) only; see .pre-commit-config.yaml. +# The other Python files are checked with autopep8 and flake8. +[tool.ruff] +include = ["pyarrow-stubs/**/*.pyi"] +line-length = 88 + +[tool.ruff.lint] +# PYI: the flake8-pyi rules for type stubs. +select = ["E", "F", "W", "PYI"] diff --git a/python/requirements-stubs.txt b/python/requirements-stubs.txt new file mode 100644 index 000000000000..ba6f48eec2fc --- /dev/null +++ b/python/requirements-stubs.txt @@ -0,0 +1,20 @@ +# Type checkers and the type stubs of PyArrow's dependencies, pinned so that +# everyone gets the same results; see the "Type Checking" section of +# docs/source/developers/python/development.rst. The python-stubs hooks in +# .pre-commit-config.yaml and ci/scripts/python_test_type_annotations.sh +# install them from this file. +fsspec==2026.9.0 +mypy==2.4.0 +pandas-stubs==3.0.5.260914 +pyright==1.1.414 +# The scipy extra installs the matching scipy, which the type stub tests import. +scipy-stubs[scipy]==1.17.1.5 +ty==0.0.99 +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 +# Annotated dependencies of the type stub tests, usually installed already. +numpy +pytest +typing_extensions>=4.12.0 diff --git a/python/scripts/run_stubtest.py b/python/scripts/run_stubtest.py new file mode 100644 index 000000000000..e89e2a6e6db9 --- /dev/null +++ b/python/scripts/run_stubtest.py @@ -0,0 +1,75 @@ +# 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. + +""" +Check the type stubs in pyarrow-stubs/ against the installed PyArrow with +mypy's stubtest. + +Usage: + python scripts/run_stubtest.py [stubtest options] + +Stubs that aren't written yet are skipped (--ignore-missing-stub), and the +known mismatches are listed in pyarrow-stubs-tests/stubtest_allowlist.txt. +Pass --generate-allowlist to print the entries for the current mismatches. +""" + +import contextlib +import os +import sys +import tempfile +from pathlib import Path + +from mypy import stubtest + +import pyarrow.lib + +PYTHON_DIR = Path(__file__).resolve().parents[1] +STUBS_DIR = PYTHON_DIR / "pyarrow-stubs" +ALLOWLIST = PYTHON_DIR / "pyarrow-stubs-tests" / "stubtest_allowlist.txt" + + +def main(argv): + # The stubs of pyarrow.lib are split into the stub-only pyarrow.__lib + # package, one module per .pxi file. Alias those modules to pyarrow.lib, + # otherwise stubtest can't import them and skips everything they define. + for stub in (STUBS_DIR / "__lib").glob("*.pyi"): + name = "pyarrow.__lib" + if stub.stem != "__init__": + name += f".{stub.stem}" + sys.modules[name] = pyarrow.lib + + # mypy doesn't look for "-stubs" packages in its search path, so give it + # the stubs under the name of the package. Run it from the directory with + # the stubs: mypy searches the working directory too, so in python/ it + # would read the sources of the modules without stubs, e.g. pyarrow/fs.py, + # and stubtest would check them as if they were stubs. + with tempfile.TemporaryDirectory() as tmpdir, contextlib.chdir(tmpdir): + os.symlink(STUBS_DIR, Path(tmpdir) / "pyarrow") + options = stubtest.parse_options([ + "pyarrow", + "--ignore-missing-stub", + # Every Cython extension type is a disjoint base (PEP 800). + "--ignore-disjoint-bases", + "--allowlist", str(ALLOWLIST), + "--mypy-config-file", str(PYTHON_DIR / "pyproject.toml"), + *argv, + ]) + return stubtest.test_stubs(options) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/python/scripts/update_stub_docstrings.py b/python/scripts/update_stub_docstrings.py index 5ac424f138a0..982445bc3a56 100644 --- a/python/scripts/update_stub_docstrings.py +++ b/python/scripts/update_stub_docstrings.py @@ -170,10 +170,6 @@ def leave_FunctionDef(self, original_node, updated_node): return updated_node -LIB_MODULES = {"array", "builder", "compat", "config", "device", "error", "io", - "_ipc", "memory", "pandas_shim", "scalar", "table", "tensor", "_types"} - - def add_docstrings_to_stubs(stubs_dir): """Update all stub files in stubs_dir with docstrings from pyarrow runtime.""" stubs_dir = Path(stubs_dir) @@ -186,7 +182,7 @@ def add_docstrings_to_stubs(stubs_dir): continue module_name = stub_file.stem - if module_name in LIB_MODULES: + if stub_file.parent.name == "__lib": namespace = "lib" elif stub_file.parent.name in ("parquet", "interchange"): namespace = (stub_file.parent.name if module_name == "__init__"