diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 50293d291f4d..3b6ea465b5de 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -136,15 +136,100 @@ repos: - id: flake8 alias: python name: Python Lint + additional_dependencies: + - flake8-pyi args: - "--config" - "python/setup.cfg" + - "--extend-select" + - "Y" + - "--per-file-ignores" + - "python/pyarrow-stubs/pyarrow/*.pyi:E301,E302,E305,E701,F821" files: >- ^(c_glib|dev|python)/ + types: + - file + types_or: + - python + - pyi exclude: >- ( ?^python/pyarrow/vendored/| ) + # Linting, formatting and type checking of the type stubs. These hooks only + # run on request, not on commit or in the existing CI jobs: + # pre-commit run --hook-stage manual python-stubs --all-files + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.17.2 + hooks: + - id: ruff-check + alias: python-stubs + name: Python Stubs Lint (ruff) + stages: + - manual + args: + - "--fix" + files: >- + ^python/pyarrow-stubs/ + types: + - pyi + - id: ruff-format + alias: python-stubs + name: Python Stubs Format (ruff) + stages: + - manual + files: >- + ^python/pyarrow-stubs/ + types: + - pyi + - repo: local + hooks: + - id: mypy + alias: python-stubs + name: Python Stubs Type Check (mypy) + stages: + - manual + language: python + entry: bash -c "cd python && exec mypy" + pass_filenames: false + files: >- + ^python/(pyarrow-stubs/|pyproject\.toml$) + additional_dependencies: &python-type-check-dependencies + - fsspec==2026.9.0 + - mypy==2.4.0 + - numpy==2.4.6 + - pandas-stubs==3.0.5.260914 + # Provides the modules that aren't annotated yet, e.g. pyarrow.compute. + - pyarrow==26.0.0 + - pyright==1.1.414 + - scipy-stubs==1.17.1.5 + - ty==0.0.92 + - types-cffi==2.1.0.20260827 + - types-psutil==7.2.2.20260906 + - types-python-dateutil==2.9.0.20261021 + - types-requests==2.33.0.20261006 + - id: pyright + alias: python-stubs + name: Python Stubs Type Check (pyright) + stages: + - manual + language: python + entry: pyright -p python/pyproject.toml + pass_filenames: false + files: >- + ^python/(pyarrow-stubs/|pyproject\.toml$) + additional_dependencies: *python-type-check-dependencies + - id: ty + alias: python-stubs + name: Python Stubs Type Check (ty) + stages: + - manual + language: python + entry: ty check --project python + pass_filenames: false + files: >- + ^python/(pyarrow-stubs/|pyproject\.toml$) + additional_dependencies: *python-type-check-dependencies - repo: https://github.com/MarcoGorelli/cython-lint rev: v0.21.1 hooks: diff --git a/ci/scripts/python_test_type_annotations.sh b/ci/scripts/python_test_type_annotations.sh index 092bedf3f5ea..96fe9a3135de 100755 --- a/ci/scripts/python_test_type_annotations.sh +++ b/ci/scripts/python_test_type_annotations.sh @@ -36,3 +36,7 @@ cd "${pyarrow_dir}" mypy pyright --stats ty check --verbose --output-format concise + +# The type annotation tests are type checked above; also run them to verify +# the annotated types against the installed pyarrow. +pytest -r s pyarrow-stubs/tests diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index 91f98cab060b..e50a78096ca8 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -114,6 +114,18 @@ Running Type Checkers We support multiple type checkers. Their configurations are in ``pyproject.toml``. +The easiest way to run all of them is with ``pre-commit``, which installs +the same versions of the type checkers and stub packages for everyone. This +also lints and formats the stubs with ruff: + +.. code-block:: + + $ pre-commit run --hook-stage manual python-stubs --all-files + +These hooks only run on request, not on commit. + +To run a type checker directly instead: + **mypy** To run mypy on the PyArrow codebase: @@ -169,6 +181,48 @@ in downstream projects and for users' IDEs. Note: ``py.typed`` marker file in the ``pyarrow/`` directory indicates to type checkers that PyArrow supports type checking according to :pep:`561`. +Testing Type Stubs +------------------ + +The tests in ``pyarrow-stubs/tests/`` pin down the annotated types of public +APIs. Each check pairs :func:`typing.assert_type`, which every type checker +verifies statically, with a runtime ``isinstance`` check: + +.. code-block:: python + + import pyarrow.lib as lib + + from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + def test_list_type() -> None: + lst = lib.list_(lib.int32()) + check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst.value_type, lib.DataType), lib.DataType) + + if TYPE_CHECKING_INVALID_USAGE: + lib.list_(1) # type: ignore[call-overload] + +The type checkers check these files as part of the regular runs above, and +pytest runs them against the installed PyArrow to make sure the stubs agree +with what PyArrow actually returns: + +.. code-block:: + + $ cd arrow/python + $ pytest pyarrow-stubs/tests + +Run ``pytest`` directly and pass it the directory: ``python -m pytest`` puts +``arrow/python`` on ``sys.path``, so the source tree would be imported instead +of the installed PyArrow. + +Code under ``if TYPE_CHECKING_INVALID_USAGE:`` is never executed. It should +fail type checking, which is asserted with a ``# type: ignore[]`` +comment: unused ignores are reported as errors in these tests, so the test +fails once the usage is no longer rejected. Start new test files with a +``# mypy: warn-unused-ignores`` comment to enable this for mypy (pyright is +configured in ``pyproject.toml``). ty doesn't honor mypy error codes, so it +doesn't check invalid usage. + Doctest ======= diff --git a/python/pyarrow-stubs/pyarrow/__init__.pyi b/python/pyarrow-stubs/pyarrow/__init__.pyi index ccec8d5abc07..4b2a8da206c2 100644 --- a/python/pyarrow-stubs/pyarrow/__init__.pyi +++ b/python/pyarrow-stubs/pyarrow/__init__.pyi @@ -15,12 +15,6 @@ # specific language governing permissions and limitations # under the License. -"""Type stubs for PyArrow. - -This is a placeholder stub file. -Complete type annotations will be added in subsequent PRs. -""" - from typing import Any # TODO(GH-48970): remove __getattr__ before release as this diff --git a/python/pyarrow-stubs/pyarrow/_compute.pyi b/python/pyarrow-stubs/pyarrow/_compute.pyi new file mode 100644 index 000000000000..3108cb063747 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_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/pyarrow/_ipc.pyi b/python/pyarrow-stubs/pyarrow/_ipc.pyi new file mode 100644 index 000000000000..c84964043e4a --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_ipc.pyi @@ -0,0 +1,25 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Placeholder stub - complete annotations in future PR. +# At runtime, these symbols are provided by the pyarrow.lib C extension. + +from typing import Any, TypeAlias + +RecordBatchReader: TypeAlias = Any + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi new file mode 100644 index 000000000000..89c36cba98a0 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -0,0 +1,117 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import datetime as dt + +from _typeshed import ( + SupportsDunderGE, + SupportsDunderGT, + SupportsDunderLE, + SupportsDunderLT, +) +from collections.abc import Collection, Container, Iterator, Sequence, Sized +from decimal import Decimal +from typing import Any, Literal, Protocol, TypeAlias + +import numpy as np + +from numpy.typing import NDArray + +from pyarrow import lib +from pyarrow.lib import ChunkedArray + +ArrayLike: TypeAlias = Any +ScalarLike: TypeAlias = Any +Order: TypeAlias = Literal["ascending", "descending"] +JoinType: TypeAlias = Literal[ + "left semi", + "right semi", + "left anti", + "right anti", + "inner", + "left outer", + "right outer", + "full outer", +] +Compression: TypeAlias = Literal[ + "gzip", "bz2", "brotli", "lz4", "lz4_frame", "lz4_raw", "zstd", "snappy" +] +NullEncoding: TypeAlias = Literal["mask", "encode"] +NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] +TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] + +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) +NumpyScalar: TypeAlias = np.generic[Any] + +Mask: TypeAlias = ( + Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray +) +Indices: TypeAlias = ( + Sequence[int | None] + | NDArray[np.integer[Any]] + | lib.IntegerArray + | ChunkedArray +) + +class SupportsDunderEQ(Protocol): + def __eq__(self, other: object, /) -> bool: ... + +FilterTuple: TypeAlias = ( + tuple[str, Literal["=", "==", "!="], SupportsDunderEQ] + | tuple[str, Literal["<"], SupportsDunderLT[Any]] + | tuple[str, Literal[">"], SupportsDunderGT[Any]] + | tuple[str, Literal["<="], SupportsDunderLE[Any]] + | tuple[str, Literal[">="], SupportsDunderGE[Any]] + | tuple[str, Literal["in", "not in"], Collection[Any]] + | tuple[str, str, Any] # Allow general str for operator to avoid type errors +) + +class Buffer(Protocol): ... +class SupportsPyBuffer(Protocol): ... + +class SupportsArrowStream(Protocol): + def __arrow_c_stream__(self, requested_schema: Any = None, /) -> Any: ... + +class SupportsPyArrowArray(Protocol): + def __arrow_array__(self, type: Any = None, /) -> Any: ... + +class SupportsArrowArray(Protocol): + def __arrow_c_array__(self, requested_schema: Any = None, /) -> Any: ... + +class SupportsArrowDeviceArray(Protocol): + def __arrow_c_device_array__( + self, requested_schema: Any = None, /, **kwargs: Any + ) -> Any: ... + +class SupportsArrowSchema(Protocol): + def __arrow_c_schema__(self) -> Any: ... + +class NullableCollection(Sized, Container[Any], Protocol): + def __iter__(self) -> Iterator[Any]: ... + def __len__(self) -> int: ... + def __contains__(self, item: Any, /) -> bool: ... diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi new file mode 100644 index 000000000000..a46f72394cd4 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -0,0 +1,698 @@ +# 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, Literal, TypeAlias, overload + +import numpy as np +import pandas as pd # type: ignore[import-not-found,import-untyped] # pyright: ignore[reportMissingImports,reportMissingTypeStubs] # noqa: E501 +from typing_extensions import Self, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( # noqa: F401 + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + Table, +) + +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 __hash__(self) -> int: ... + 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: ... + @classmethod + def _import_from_c(cls, in_ptr: int) -> Self: ... + def __arrow_c_schema__(self) -> Any: ... + @classmethod + def _import_from_c_capsule(cls, schema: Any) -> Self: ... + +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class NullType(DataType): ... +class BoolType(DataType): ... +class UInt8Type(DataType): ... +class Int8Type(DataType): ... +class UInt16Type(DataType): ... +class Int16Type(DataType): ... +class UInt32Type(DataType): ... +class Int32Type(DataType): ... +class UInt64Type(DataType): ... +class Int64Type(DataType): ... +class Float16Type(DataType): ... +class Float32Type(DataType): ... +class Float64Type(DataType): ... +class Date32Type(DataType): ... +class Date64Type(DataType): ... +class MonthDayNanoIntervalType(DataType): ... +class StringType(DataType): ... +class LargeStringType(DataType): ... +class StringViewType(DataType): ... +class BinaryType(DataType): ... +class LargeBinaryType(DataType): ... +class BinaryViewType(DataType): ... + +_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): ... + +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) + +class DictionaryType(DataType): + @property + def ordered(self) -> bool: ... + @property + def index_type(self) -> _IndexType: ... + @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: ... + +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 | None +_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, + **kwargs: str | bytes, + ) -> None: ... + 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: ... + 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: ... + def __hash__(self) -> int: ... + @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, + ) -> Self: ... + def remove_metadata(self) -> Self: ... + def with_type(self, new_type: DataType) -> Field: ... + def with_name(self, name: str | bytes) -> Self: ... + def with_nullable(self, nullable: bool) -> Self: ... + def flatten(self) -> list[Field]: ... + def _export_to_c(self, out_ptr: int) -> None: ... + @classmethod + def _import_from_c(cls, in_ptr: int) -> Self: ... + def __arrow_c_schema__(self) -> Any: ... + @classmethod + def _import_from_c_capsule(cls, schema: Any) -> Self: ... + +_FieldTuple: TypeAlias = ( + tuple[_FieldName, _DataTypeLike] + | tuple[_FieldName, _DataTypeLike, bool | None] + | tuple[ + _FieldName, + _DataTypeLike, + bool | None, + _KeyValueMetadataInput | KeyValueMetadata, + ] +) +_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 __hash__(self) -> int: ... + 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: ... + @deprecated("Use 'field' instead") + def field_by_name(self, name: str | bytes) -> Field | None: ... + 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: ... + @deprecated("Use 'with_metadata' instead") + def add_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> Schema: ... + def with_metadata( + self, metadata: _KeyValueMetadataInput | KeyValueMetadata + ) -> 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: ... + @classmethod + def _import_from_c(cls, in_ptr: int) -> Schema: ... + def __arrow_c_schema__(self) -> Any: ... + @staticmethod + def _import_from_c_capsule(schema: Any) -> 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, +) -> Field: ... +@overload +def field( + name: _FieldName, + type: _DataTypeLike, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = 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 | None = None, +) -> MapType: ... +def dictionary( + index_type: _IndexType | _DataTypeAliasInput, + value_type: _DataTypeLike, + ordered: bool | None = None, +) -> 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, +) -> Schema: ... +def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... + +__all__ = [ + "_Weakrefable", + "_Metadata", + "DataType", + "NullType", + "BoolType", + "UInt8Type", + "Int8Type", + "UInt16Type", + "Int16Type", + "UInt32Type", + "Int32Type", + "UInt64Type", + "Int64Type", + "Float16Type", + "Float32Type", + "Float64Type", + "Date32Type", + "Date64Type", + "MonthDayNanoIntervalType", + "StringType", + "LargeStringType", + "StringViewType", + "BinaryType", + "LargeBinaryType", + "BinaryViewType", + "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/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi new file mode 100644 index 000000000000..bca6cdbf37fd --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -0,0 +1,816 @@ +# 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 collections.abc import Iterable, Iterator, Sequence + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +from typing import Any, Literal + +import numpy as np +import pandas as pd + +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + Mask, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, +) +from pyarrow.lib import ( + Buffer, + Device, + MemoryManager, + MemoryPool, + Tensor, + _Weakrefable, +) +from typing_extensions import deprecated +import builtins + +from .scalar import Int64Scalar, Scalar +from .device import DeviceAllocationType +from ._types import BaseExtensionType, DataType, Field, StructType +from ._stubs_typing import NullableCollection + + +def array( + values: NullableCollection | Iterable[Any] | SupportsArrowArray + | SupportsArrowDeviceArray | SupportsPyArrowArray, + type: Any | None = None, + mask: Mask | pd.Series[bool] | None = None, + size: int | None = None, + from_pandas: bool | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... + + +def asarray( + values: NullableCollection | Iterable[Any] | SupportsArrowArray + | SupportsArrowDeviceArray, + type: Any | None = None, +) -> ArrayLike: ... + + +def nulls( + size: int, + type: Any | None = None, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... + + +def repeat( + value: Any, + size: int, + memory_pool: MemoryPool | None = None, +) -> ArrayLike: ... + + +def infer_type(values: Iterable[Any], 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: Literal["None", "lossy", "strict"] | None = None, + types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None + coerce_temporal_nanoseconds: bool = False, + ) -> pd.Series | pd.DataFrame: ... + + +class Array(_PandasConvertible): + def diff(self, other: Self) -> str: ... + + # Private attribute used internally (e.g., for column names in batches) + _name: str | None + + def cast( + self, + target_type: DataType | str, + safe: bool = True, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + + def view(self, target_type: DataType) -> Array: ... + + def sum(self, **kwargs) -> Scalar: ... + + @property + def type(self) -> DataType: ... + def unique(self) -> Self: ... + + def dictionary_encode(self, null_encoding: str = "mask") -> DictionaryArray: ... + + def value_counts(self) -> StructArray: ... + + @staticmethod + def from_pandas( + obj: pd.Series | np.ndarray | ArrayLike, + mask: Mask | None = None, + type: DataType | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Array: ... + + @staticmethod + def from_buffers( + type: DataType, + length: int, + buffers: Sequence[Buffer | None], + null_count: int = -1, + offset=0, + children: NullableCollection | 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, + ) -> str: ... + + format = to_string + def equals(self, other: Array | Any) -> 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: int | 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: Literal["drop", "emit_null"] = "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", **kwargs) -> 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: Literal["lossy", "strict"] | None = None, + ) -> list[Any]: ... + + tolist = to_pylist + 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: ... + + @classmethod + def _import_from_c(cls, in_ptr: int, type: int | DataType) -> Self: ... + + def __arrow_c_array__(self, requested_schema=None) -> Any: ... + + @classmethod + def _import_from_c_capsule(cls, schema_capsule, array_capsule) -> Self: ... + def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + + @classmethod + def _import_from_c_device(cls, in_ptr: int, type: DataType | int) -> Self: ... + + def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... + + @classmethod + def _import_from_c_device_capsule(cls, schema_capsule, array_capsule) -> Self: ... + def __dlpack__(self, stream: int | None = None) -> Any: ... + + 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): + @classmethod + def from_arrays( + cls, + offsets: Int32Array | list[int] | list[int | None], + 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): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | list[int | None], + values: Array, + *, + 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): + @classmethod + def from_arrays( + cls, + offsets: Int32Array | list[int] | list[int | None], + sizes: Int32Array | list[int] | list[int | None], + values: Array, + *, + 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): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | list[int | None], + sizes: Int64Array | list[int] | list[int | None], + values: Array, + *, + 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): + @classmethod + def from_arrays( + cls, + values: Array, + 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): + @classmethod + def from_arrays( # type: ignore[override] + cls, + offsets: Int32Array | list[int] | list[int | None], + 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): + @deprecated("Use fields() instead") + def child(self, pos: int) -> Field: ... + + 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: NullableCollection, + field_names: list[str] | None = None, + type_codes: Int8Array | list[int] | None = None, + ) -> UnionArray: ... + + @staticmethod + def from_sparse( + types: Int8Array, + children: NullableCollection, + field_names: list[str] | None = None, + type_codes: Int8Array | 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 | None = -1, + offset: int | None = 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 | None = -1, + offset: int | None = 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: list[Buffer], + dictionary: Array | np.ndarray | pd.Series, + null_count: int = -1, + offset: int = 0, + ) -> DictionaryArray: ... + + @staticmethod + def from_arrays( + indices: Indices | Sequence[int | None], + 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=None, + memory_pool: MemoryPool | None = None, + type: StructType | None = None, + ) -> StructArray: ... + + def sort(self, order: Order = "ascending", by: str | + None = None, **kwargs) -> 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=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: ... + + def to_tensor(self) -> Tensor: ... + + @classmethod + def from_numpy_ndarray( + cls, obj: np.ndarray, + dim_names: list[str] | tuple[str, ...] | None = None + ) -> Self: ... + + +class OpaqueArray(ExtensionArray): + ... + + +class Bool8Array(ExtensionArray): + def to_numpy(self, zero_copy_only: bool = ..., + writable: bool = ...) -> np.ndarray: ... + + @classmethod + def from_storage(cls, storage: Int8Array) -> Self: ... # type: ignore[override] + + @classmethod + def from_numpy(cls, obj: np.ndarray) -> Self: ... + + +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", + "_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/pyarrow/builder.pyi b/python/pyarrow-stubs/pyarrow/builder.pyi new file mode 100644 index 000000000000..9001d9835b6c --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/builder.pyi @@ -0,0 +1,51 @@ +# 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, _Weakrefable + +from .array import StringArray, StringViewArray + + +class StringBuilder(_Weakrefable): + def __init__(self, memory_pool: MemoryPool | None = None) -> None: ... + def append(self, value: str | bytes | float | None): ... + + def append_values(self, values: Iterable[str | bytes | float | 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): ... + + def append_values(self, values: Iterable[str | bytes | float | None]): ... + + def finish(self) -> StringViewArray: ... + + @property + def null_count(self) -> int: ... + def __len__(self) -> int: ... + + +__all__ = ["StringBuilder", "StringViewBuilder"] diff --git a/python/pyarrow-stubs/pyarrow/config.pyi b/python/pyarrow-stubs/pyarrow/config.pyi new file mode 100644 index 000000000000..069b70e553ac --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/config.pyi @@ -0,0 +1,72 @@ +# 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/pyarrow/device.pyi b/python/pyarrow-stubs/pyarrow/device.pyi new file mode 100644 index 000000000000..7787ac44deb8 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/device.pyi @@ -0,0 +1,66 @@ +# 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.Enum): + CPU = enum.auto() + CUDA = enum.auto() + CUDA_HOST = enum.auto() + OPENCL = enum.auto() + VULKAN = enum.auto() + METAL = enum.auto() + VPI = enum.auto() + ROCM = enum.auto() + ROCM_HOST = enum.auto() + EXT_DEV = enum.auto() + CUDA_MANAGED = enum.auto() + ONEAPI = enum.auto() + WEBGPU = enum.auto() + HEXAGON = enum.auto() + + +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/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/error.pyi new file mode 100644 index 000000000000..67be73eadb6d --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/error.pyi @@ -0,0 +1,71 @@ +# 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_extensions 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/io.pyi b/python/pyarrow-stubs/pyarrow/io.pyi new file mode 100644 index 000000000000..8a1da39a1cf2 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/io.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 Buffer: ... + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi new file mode 100644 index 000000000000..e1c64d028bcf --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -0,0 +1,36 @@ +# 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 + +from pyarrow._types import * # noqa: F401,F403 +from pyarrow.array import * # noqa: F401,F403 +from pyarrow.builder import * # noqa: F401,F403 +from pyarrow.config import * # noqa: F401,F403 +from pyarrow.device import * # noqa: F401,F403 +from pyarrow.error import * # noqa: F401,F403 +from pyarrow.io import * # noqa: F401,F403 +from pyarrow.memory import * # noqa: F401,F403 +from pyarrow.scalar import * # noqa: F401,F403 +from pyarrow.table import * # noqa: F401,F403 +from pyarrow.tensor import * # noqa: F401,F403 + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/pyarrow/memory.pyi b/python/pyarrow-stubs/pyarrow/memory.pyi new file mode 100644 index 000000000000..f80e01ab21c0 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/memory.pyi @@ -0,0 +1,94 @@ +# 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/pyarrow/scalar.pyi b/python/pyarrow-stubs/pyarrow/scalar.pyi new file mode 100644 index 000000000000..307b3e689a25 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/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/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi new file mode 100644 index 000000000000..e60218af9424 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -0,0 +1,675 @@ +# 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 + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias +from collections.abc import ( + Collection, Generator, Iterable, Iterator, Sequence, Mapping) +from typing import Any, Literal +import builtins + +import numpy as np +import pandas as pd + +from numpy.typing import NDArray +from pyarrow._compute import ( + CastOptions, + CountOptions, + FunctionOptions, + ScalarAggregateOptions, + TDigestOptions, + VarianceOptions, +) +from pyarrow._stubs_typing import ( + Indices, + Mask, + NullEncoding, + NullSelectionBehavior, + Order, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsArrowStream, +) +from pyarrow.compute import Expression +from pyarrow.interchange.dataframe import _PyArrowDataFrame +from pyarrow.lib import Device, Field, MemoryManager, MemoryPool, Schema + +from .array import Array, StructArray, _PandasConvertible +from .device import DeviceAllocationType +from .io import Buffer +from ._ipc import RecordBatchReader +from .scalar import Int64Scalar, Scalar +from .tensor import Tensor +from ._stubs_typing import NullableCollection +from ._types import DataType + +ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray + +_Aggregation: TypeAlias = Literal[ + "all", + "any", + "approximate_median", + "count", + "count_all", + "count_distinct", + "distinct", + "first", + "first_last", + "last", + "list", + "max", + "mean", + "min", + "min_max", + "one", + "product", + "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_last", + "hash_list", + "hash_max", + "hash_mean", + "hash_min", + "hash_min_max", + "hash_one", + "hash_product", + "hash_stddev", + "hash_sum", + "hash_tdigest", + "hash_variance", +] +Aggregation: TypeAlias = _Aggregation | _AggregationPrefixed | str +AggregateOptions: TypeAlias = (ScalarAggregateOptions | CountOptions + | TDigestOptions | VarianceOptions | FunctionOptions) + +UnarySelector: TypeAlias = str +NullarySelector: TypeAlias = tuple[()] +NarySelector: TypeAlias = list[str] | tuple[str, ...] +ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector + + +class ChunkedArray(_PandasConvertible): + + @property + def data(self) -> Self: ... + @property + def type(self) -> DataType: ... + + # Private attribute used internally for column names + _name: str | None + + 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, + ) -> str: ... + + format = to_string + 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 | Self: ... + + def getitem(self, i: int) -> Scalar: ... + 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, safe: bool = True, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None + ) -> ChunkedArray: ... + + def fill_null(self, fill_value: Any) -> Self: ... + + def equals(self, other: Self | Any) -> 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") -> Self: ... + + 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) -> 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 take(self, indices: Indices) -> Self: ... + + def drop_null(self) -> Self: ... + + def sort(self, order: Order = "ascending", **kwargs) -> Self: ... + + def unify_dictionaries(self, memory_pool: MemoryPool | None = None) -> Self: ... + + @property + def num_chunks(self) -> int: ... + + def chunk(self, i: int) -> Array: ... + + @property + def chunks(self) -> list[Array]: ... + + def iterchunks(self) -> Generator[Array, None, None]: ... + + def __iter__(self) -> Iterator[Scalar]: ... + + def to_pylist( + self, + *, + maps_as_pydicts: Literal["lossy", "strict"] | None = None, + ) -> list[Any]: ... + + def __arrow_c_stream__(self, requested_schema=None) -> Any: ... + + @classmethod + def _import_from_c_capsule(cls, stream) -> Self: ... + + @property + def is_cpu(self) -> bool: ... + + +def chunked_array( + arrays: Iterable[NullableCollection] + | Iterable[Iterable[Any] | SupportsArrowStream | SupportsArrowArray] + | Iterable[Array] | Array + | SupportsArrowArray | SupportsArrowStream, + 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 | list[Any] | np.ndarray | range], + 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) -> Generator[ArrayOrChunkedArray, None, None]: ... + + @property + def num_columns(self) -> int: ... + @property + def num_rows(self) -> int: ... + @property + def shape(self) -> tuple[int, int]: ... + + @property + def schema(self) -> Schema: ... + @property + def nbytes(self) -> int: ... + def sort_by(self, sorting: str | list[tuple[str, Order]], **kwargs) -> 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: Literal["lossy", "strict"] | None = None + ) -> dict[str, list[Any]]: ... + + def to_pylist( + self, *, maps_as_pydicts: Literal["lossy", "strict"] | 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 + ) -> Self: ... + + @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 + ) -> Self: ... + + def remove_column(self, i: int) -> Self: ... + + def set_column(self, i: int, field_: str | Field, column: Array | list) -> Self: ... + + def rename_columns(self, names: list[str] | dict[str, str]) -> Self: ... + + def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... + + def slice(self, offset: int = 0, length: int | None = None) -> Self: ... + + def equals(self, other: Self | Any, check_metadata: bool = False) -> bool: ... + + def select(self, columns: Iterable[str] | + Iterable[int] | NDArray[np.str_]) -> Self: ... + + def cast(self, target_schema: Schema, safe: bool | None = None, + options: CastOptions | None = None) -> Self: ... + + @classmethod + def from_arrays( + cls, + arrays: Iterable[Any], + names: list[str] | tuple[str, ...] | None = None, + schema: Schema | None = None, + metadata: Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] + | None = None, + ) -> Self: ... + + @classmethod + def from_pandas( + cls, + df: pd.DataFrame, + schema: Schema | None = None, + preserve_index: bool | None = None, + nthreads: int | None = None, + columns: Sequence[str | int] | None = None, + ) -> Self: ... + + @classmethod + def from_struct_array( + cls, struct_array: StructArray | ChunkedArray + ) -> Self: ... + + 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): ... + + @classmethod + def _import_from_c(cls, in_ptr: int, schema: Schema) -> Self: ... + + def __arrow_c_array__(self, requested_schema=None): ... + + def __arrow_c_stream__(self, requested_schema=None): ... + + @classmethod + def _import_from_c_capsule(cls, schema_capsule, array_capsule) -> Self: ... + + def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... + + @classmethod + def _import_from_c_device(cls, in_ptr: int, schema: Schema) -> Self: ... + + def __arrow_c_device_array__(self, requested_schema=None, **kwargs): ... + + @classmethod + def _import_from_c_device_capsule(cls, schema_capsule, array_capsule) -> Self: ... + + @property + def device_type(self) -> DeviceAllocationType: ... + + @property + def is_cpu(self) -> bool: ... + + def copy_to(self, destination: MemoryManager | Device) -> Self: ... + + +def table_to_blocks(options, table: Table, categories, extension_columns): ... + + +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) -> Self: ... + + def select(self, columns: Iterable[str] | + Iterable[int] | NDArray[np.str_]) -> Self: ... + + def replace_schema_metadata( + self, metadata: dict[str, str] + | dict[bytes, bytes] + | dict[bytes, str] + | dict[str, bytes] + | None = None + ) -> Self: ... + + def flatten(self, memory_pool: MemoryPool | None = None) -> Self: ... + + def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Self: ... + + def unify_dictionaries(self, memory_pool: MemoryPool | None = None) -> Self: ... + + def equals(self, other: Self | Any, check_metadata: bool = False) -> bool: ... + + def cast(self, target_schema: Schema, safe: bool | None = None, + options: CastOptions | None = None) -> Self: ... + + @classmethod + def from_pandas( + cls, + df: pd.DataFrame, + schema: Schema | None = None, + preserve_index: bool | None = None, + nthreads: int | None = None, + columns: Sequence[str | int] | None = None, + safe: bool = True, + ) -> Self: ... + + @classmethod + def from_arrays( + cls, + arrays: + Collection[ArrayOrChunkedArray | Collection[NDArray[Any]] | list[Any]], + names: list[str] | tuple[str, ...] | None = None, + schema: Schema | None = None, + metadata: Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] | None = None, + ) -> Self: ... + + @classmethod + def from_struct_array( + cls, struct_array: StructArray | ChunkedArray + ) -> Self: ... + + def to_struct_array( + self, max_chunksize: int | None = None + ) -> ChunkedArray: ... + + @classmethod + def from_batches(cls, batches: Iterable[RecordBatch], + schema: Schema | None = None) -> Self: ... + + def to_batches(self, max_chunksize: int | None = None) -> list[RecordBatch]: ... + + 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: ArrayOrChunkedArray | list[list[Any]]) -> Self: ... + + def remove_column(self, i: int) -> Self: ... + + def set_column(self, i: int, field_: str | Field, + column: ArrayOrChunkedArray | list[list[Any]]) -> Self: ... + + def rename_columns(self, names: list[str] | dict[str, str]) -> Self: ... + + def drop(self, columns: str | list[str]) -> Self: ... + + def group_by(self, keys: str | list[str], + use_threads: bool = True) -> TableGroupBy: ... + + def join( + self, + right_table: Self, + 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, + ) -> Self: ... + + def join_asof( + self, + right_table: Self, + on: str, + by: str | list[str], + tolerance: int, + right_on: str | list[str] | None = None, + right_by: str | list[str] | None = None, + ) -> Self: ... + + def __arrow_c_stream__(self, requested_schema=None): ... + + @property + def is_cpu(self) -> bool: ... + + +def record_batch( + data: Mapping[str, list[Any] | Array] + | Collection[Array | ChunkedArray | 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: Collection[ArrayOrChunkedArray | list[Any] | range | str] + | pd.DataFrame + | SupportsArrowArray + | SupportsArrowStream + | SupportsArrowDeviceArray + | Mapping[str, list[Any] | Array | ChunkedArray | range] + | Mapping[str, Any], + 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", + **kwargs: Any, +) -> 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, Aggregation] + | tuple[ColumnSelector, Aggregation, AggregateOptions | None] + ], + ) -> Table: ... + + def _table(self) -> Table: ... + @property + def _use_threads(self) -> bool: ... + + +def concat_batches( + recordbatches: Iterable[RecordBatch], memory_pool: MemoryPool | None = None +) -> RecordBatch: ... + + +__all__ = [ + "ChunkedArray", + "chunked_array", + "_Tabular", + "RecordBatch", + "table_to_blocks", + "Table", + "record_batch", + "table", + "concat_tables", + "TableGroupBy", + "concat_batches", + "Aggregation", + "AggregateOptions", +] diff --git a/python/pyarrow-stubs/pyarrow/tensor.pyi b/python/pyarrow-stubs/pyarrow/tensor.pyi new file mode 100644 index 000000000000..b4bd52c832f0 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/tensor.pyi @@ -0,0 +1,270 @@ +# 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 + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +from collections.abc import Sequence +import numpy as np + +from pyarrow.lib import _Weakrefable +from pyarrow._types import DataType +from scipy.sparse import coo_matrix, csc_matrix, csr_matrix +from sparse import COO # type: ignore[import-untyped, import-not-found] + + +class Tensor(_Weakrefable): + @classmethod + def from_numpy(cls, obj: np.ndarray, + dim_names: Sequence[str] | None = None) -> Self: ... + + 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, ...]: ... + + @property + def type(self) -> DataType: ... + + +class SparseCOOTensor(_Weakrefable): + @classmethod + def from_dense_numpy(cls, obj: np.ndarray, + dim_names: list[str] | None = None) -> Self: ... + + @classmethod + def from_numpy( + cls, + data: np.ndarray, + coords: np.ndarray, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + ) -> Self: ... + + @classmethod + def from_scipy(cls, obj: coo_matrix, + dim_names: Sequence[str] | None = None) -> Self: ... + + @classmethod + def from_pydata_sparse( + cls, obj: COO, dim_names: Sequence[str] | None = None) -> Self: ... + + @classmethod + def from_tensor(cls, obj: Tensor) -> Self: ... + + def to_numpy(self) -> tuple[np.ndarray, np.ndarray]: ... + + def to_scipy(self) -> coo_matrix: ... + + def to_pydata_sparse(self) -> COO: ... + + def to_tensor(self) -> Tensor: ... + + def equals(self, other: Self) -> 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) -> list[str]: ... + @property + def non_zero_length(self) -> int: ... + @property + def has_canonical_format(self) -> bool: ... + @property + def type(self) -> DataType: ... + + +class SparseCSRMatrix(_Weakrefable): + @classmethod + def from_dense_numpy(cls, obj: np.ndarray, + dim_names: list[str] | None = None) -> Self: ... + + @classmethod + def from_numpy( + cls, + data: np.ndarray, + indptr: np.ndarray, + indices: np.ndarray, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + ) -> Self: ... + + @classmethod + def from_scipy(cls, obj: csr_matrix, + dim_names: Sequence[str] | None = None) -> Self: ... + + @classmethod + def from_tensor(cls, obj: Tensor) -> Self: ... + + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + + def to_scipy(self) -> csr_matrix: ... + + def to_tensor(self) -> Tensor: ... + + def equals(self, other: Self) -> 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) -> list[str]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + + +class SparseCSCMatrix(_Weakrefable): + @classmethod + def from_dense_numpy(cls, obj: np.ndarray, + dim_names: list[str] | None = None) -> Self: ... + + @classmethod + def from_numpy( + cls, + data: np.ndarray, + indptr: np.ndarray, + indices: np.ndarray, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + ) -> Self: ... + + @classmethod + def from_scipy(cls, obj: csc_matrix, + dim_names: Sequence[str] | None = None) -> Self: ... + + @classmethod + def from_tensor(cls, obj: Tensor) -> Self: ... + + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + + def to_scipy(self) -> csc_matrix: ... + + def to_tensor(self) -> Tensor: ... + + def equals(self, other: Self) -> 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) -> list[str]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + + +class SparseCSFTensor(_Weakrefable): + @classmethod + def from_dense_numpy(cls, obj: np.ndarray, + dim_names: Sequence[str] | None = None) -> Self: ... + + @classmethod + def from_numpy( + cls, + data: np.ndarray, + indptr: Sequence[np.ndarray], + indices: Sequence[np.ndarray], + shape: tuple[int, ...], + axis_order: Sequence[int] | None = None, + dim_names: Sequence[str] | None = None, + ) -> Self: ... + + @classmethod + def from_tensor(cls, obj: Tensor) -> Self: ... + + def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... + + def to_tensor(self) -> Tensor: ... + + def equals(self, other: Self) -> 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) -> list[str]: ... + @property + def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... + + +__all__ = [ + "Tensor", + "SparseCOOTensor", + "SparseCSRMatrix", + "SparseCSCMatrix", + "SparseCSFTensor", +] diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi new file mode 100644 index 000000000000..1bb66bd1209a --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -0,0 +1,264 @@ +# 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, TypeGuard + +if sys.version_info >= (3, 13): + from typing import TypeIs +else: + from typing_extensions import TypeIs +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias + +import pyarrow.lib as lib + +from pyarrow._types import ( + BinaryType, + BinaryViewType, + BoolType, + DataType, + Date32Type, + Date64Type, + Decimal32Type, + Decimal64Type, + Decimal128Type, + Decimal256Type, + DenseUnionType, + DictionaryType, + DurationType, + FixedSizeBinaryType, + FixedSizeListType, + Float16Type, + Float32Type, + Float64Type, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + LargeBinaryType, + LargeListType, + LargeListViewType, + LargeStringType, + ListType, + ListViewType, + MapType, + MonthDayNanoIntervalType, + NullType, + RunEndEncodedType, + SparseUnionType, + StringType, + StringViewType, + StructType, + Time32Type, + Time64Type, + TimestampType, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, +) + +_SignedInteger: TypeAlias = Int8Type | Int16Type | Int32Type | Int64Type +_UnsignedInteger: TypeAlias = UInt8Type | UInt16Type | UInt32Type | UInt64Type +_Integer: TypeAlias = _SignedInteger | _UnsignedInteger +_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[_Integer]: ... +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: ... + + +__all__ = [ + "lib", + "is_binary", + "is_binary_view", + "is_boolean", + "is_date", + "is_date32", + "is_date64", + "is_decimal", + "is_decimal128", + "is_decimal256", + "is_decimal32", + "is_decimal64", + "is_dictionary", + "is_duration", + "is_fixed_size_binary", + "is_fixed_size_list", + "is_float16", + "is_float32", + "is_float64", + "is_floating", + "is_int16", + "is_int32", + "is_int64", + "is_int8", + "is_integer", + "is_interval", + "is_large_binary", + "is_large_list", + "is_large_list_view", + "is_large_string", + "is_large_unicode", + "is_list", + "is_list_view", + "is_map", + "is_nested", + "is_null", + "is_primitive", + "is_run_end_encoded", + "is_signed_integer", + "is_string", + "is_string_view", + "is_struct", + "is_temporal", + "is_time", + "is_time32", + "is_time64", + "is_timestamp", + "is_uint16", + "is_uint32", + "is_uint64", + "is_uint8", + "is_unicode", + "is_union", + "is_unsigned_integer", +] + + +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-stubs/tests/_check.py b/python/pyarrow-stubs/tests/_check.py new file mode 100644 index 000000000000..596129874188 --- /dev/null +++ b/python/pyarrow-stubs/tests/_check.py @@ -0,0 +1,58 @@ +# 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.lib.list_(pa.lib.int32()), pa.lib.ListType), + pa.lib.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_INVALID_USAGE:`` 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. +""" + +import sys +from typing import Any, Final, TYPE_CHECKING, TypeVar + +if sys.version_info >= (3, 11): + from typing import assert_type +else: + from typing_extensions import assert_type + +__all__ = ["TYPE_CHECKING_INVALID_USAGE", "assert_type", "check"] + +TYPE_CHECKING_INVALID_USAGE: Final = TYPE_CHECKING + +_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/test_array.py b/python/pyarrow-stubs/tests/test_array.py new file mode 100644 index 000000000000..d41826eed8dc --- /dev/null +++ b/python/pyarrow-stubs/tests/test_array.py @@ -0,0 +1,113 @@ +# 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 typing import Any, cast + +import pyarrow.lib as lib + +from _check import TYPE_CHECKING_INVALID_USAGE, 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(lib.Int64Array, lib.array([1, 2, None])), lib.Int64Array) + + check(assert_type(arr.type, lib.DataType), lib.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 | lib.Int64Array), lib.Int64Scalar) + check(assert_type(arr[1:], Any | lib.Int64Array), lib.Int64Array) + + # Methods returning the same kind of array keep the subclass. + check(assert_type(arr.slice(1), lib.Int64Array), lib.Int64Array) + check(assert_type(arr.fill_null(0), lib.Int64Array), lib.Int64Array) + check(assert_type(arr.drop_null(), lib.Int64Array), lib.Int64Array) + + check(assert_type(arr.is_null(), lib.BooleanArray), lib.BooleanArray) + check(assert_type(arr.cast(lib.float64()), lib.Array), lib.DoubleArray) + check(assert_type(arr.cast("float64"), lib.Array), lib.DoubleArray) + check(assert_type(lib.concat_arrays([arr, arr]), lib.Array), + lib.Int64Array) + check(assert_type(lib.Array.from_pandas([1, 2], None, lib.int64()), + lib.Array), + lib.Int64Array) + + +def test_list_arrays() -> None: + lst = lib.ListArray.from_arrays([0, 1, 3], [1, 2, 3]) + check(assert_type(lst.value_lengths(), lib.Int32Array), lib.Int32Array) + + large = lib.LargeListArray.from_arrays([0, 1, 3], lib.array([1, 2, 3])) + check(assert_type(large.value_lengths(), lib.Int64Array), lib.Int64Array) + + # MapArray is a ListArray (of key-item structs) at runtime. + map_arr = lib.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"]) + check(assert_type(map_arr, lib.MapArray), lib.ListArray) + check(assert_type(map_arr.offsets, lib.Int32Array), lib.Int32Array) + check(assert_type(map_arr.values, lib.Array), lib.StructArray) + check(assert_type(map_arr.keys, lib.Array), lib.Int64Array) + + +def test_dictionary_array() -> None: + arr = check(cast(lib.StringArray, lib.array(["a", "b", "a"])), + lib.StringArray) + dict_arr = arr.dictionary_encode() + check(assert_type(dict_arr, lib.DictionaryArray), lib.DictionaryArray) + check(assert_type(dict_arr.indices, lib.IntegerArray), lib.Int32Array) + check(assert_type(dict_arr.dictionary, lib.Array), lib.StringArray) + check(assert_type(dict_arr.dictionary_decode(), lib.Array), + lib.StringArray) + + +def test_run_end_encoded_array() -> None: + ree = lib.RunEndEncodedArray.from_arrays([2, 3], ["a", "b"]) + check(assert_type(ree, lib.RunEndEncodedArray), lib.RunEndEncodedArray) + check( + assert_type(ree.run_ends, + lib.Int16Array | lib.Int32Array | lib.Int64Array), + lib.Int64Array, + ) + check(assert_type(ree.values, lib.Array), lib.StringArray) + + +def test_chunked_array() -> None: + chunked = lib.chunked_array([[1, 2], [3]]) + check(assert_type(chunked, lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(chunked.type, lib.DataType), lib.DataType) + check(assert_type(chunked.num_chunks, int), int) + check(assert_type(chunked.chunk(0), lib.Array), lib.Int64Array) + # Scalars aren't annotated yet, so indexing returns Any | ChunkedArray. + check(assert_type(chunked[0], Any | lib.ChunkedArray), lib.Int64Scalar) + check(assert_type(chunked[1:], Any | lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(chunked.chunks, list[lib.Array]), list) + check(assert_type(chunked.combine_chunks(), lib.Array), lib.Int64Array) + check(assert_type(chunked.unique(), lib.Array), lib.Int64Array) + check(assert_type(chunked.slice(1), lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(chunked.cast(lib.float64()), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(chunked.is_null(), lib.ChunkedArray), lib.ChunkedArray) + + +def test_invalid_usage() -> None: + arr = cast(lib.Int64Array, lib.array([1, 2])) + if TYPE_CHECKING_INVALID_USAGE: + arr.cast(1) # type: ignore[arg-type] + arr.slice("1") # type: ignore[arg-type] diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py new file mode 100644 index 000000000000..9544f44be794 --- /dev/null +++ b/python/pyarrow-stubs/tests/test_table.py @@ -0,0 +1,73 @@ +# 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 + +import pyarrow.lib as lib + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + + +def test_table() -> None: + table = lib.table({"a": [1, 2], "b": ["x", "y"]}) + check(assert_type(table, lib.Table), lib.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, lib.Schema), lib.Schema) + check(assert_type(table.field("a"), lib.Field), lib.Field) + + # Columns aren't typed per tabular class, so Table and RecordBatch both + # return the union here. + check(assert_type(table.column("a"), lib.Array | lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(table["a"], lib.Array | lib.ChunkedArray | lib.Table), + lib.ChunkedArray) + check(assert_type(table[0], lib.Array | lib.ChunkedArray | lib.Table), + lib.ChunkedArray) + check(assert_type(table[:1], lib.Array | lib.ChunkedArray | lib.Table), + lib.Table) + + check(assert_type(table.slice(0, 1), lib.Table), lib.Table) + check(assert_type(table.select(["a"]), lib.Table), lib.Table) + check(assert_type(table.combine_chunks(), lib.Table), lib.Table) + check(assert_type(table.to_batches(), list[lib.RecordBatch]), list) + check(assert_type(table.to_pydict(), dict[str, list]), dict) + + # The reader isn't annotated yet, but must not reject valid usage. + check(table.to_reader().read_all(), lib.Table) + for batch in table.to_reader(): + check(batch, lib.RecordBatch) + + +def test_record_batch() -> None: + batch = lib.record_batch({"a": [1, 2]}) + check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + check(assert_type(batch.num_rows, int), int) + check(assert_type(batch.column(0), lib.Array | lib.ChunkedArray), + lib.Array) + check(assert_type(batch.slice(1), lib.RecordBatch), lib.RecordBatch) + check(assert_type(lib.Table.from_batches([batch]), lib.Table), lib.Table) + check(assert_type(lib.concat_batches([batch, batch]), lib.RecordBatch), + lib.RecordBatch) + + +def test_invalid_usage() -> None: + table = lib.table({"a": [1, 2]}) + if TYPE_CHECKING_INVALID_USAGE: + table.column(1.5) # type: ignore[arg-type] + table.slice("0") # type: ignore[arg-type] diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py new file mode 100644 index 000000000000..2133cb4de874 --- /dev/null +++ b/python/pyarrow-stubs/tests/test_types.py @@ -0,0 +1,187 @@ +# 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 typing import TYPE_CHECKING, Literal + +import pyarrow.lib as lib +import pyarrow.types as types + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + +if TYPE_CHECKING: + # The stubs have a class per primitive type, but at runtime the + # factories return plain DataType instances, so these can only be + # referred to from type checking code. + from pyarrow.lib import ( # noqa: F401 + BoolType, + Int8Type, + Int16Type, + Int32Type, + Int64Type, + StringType, + UInt8Type, + UInt16Type, + UInt32Type, + UInt64Type, + ) + + IntegerType = ( + UInt8Type | Int8Type | UInt16Type | Int16Type + | UInt32Type | Int32Type | UInt64Type | Int64Type + ) + + +def _erase(t: lib.DataType) -> lib.DataType: + # Hide the precise type, so we can test narrowing by the is_* functions. + return t + + +def test_primitive_factories() -> None: + check(assert_type(lib.bool_(), "BoolType"), lib.DataType) + check(assert_type(lib.int32(), "Int32Type"), lib.DataType) + check(assert_type(lib.string(), "StringType"), lib.DataType) + check(assert_type(lib.utf8(), "StringType"), lib.DataType) + + +def test_temporal_types() -> None: + ts = lib.timestamp("ms", tz="UTC") + check(assert_type(ts, lib.TimestampType), lib.TimestampType) + check(assert_type(ts.unit, Literal["s", "ms", "us", "ns"]), str) + check(assert_type(ts.tz, str | None), str) + + check(assert_type(lib.time32("s").unit, Literal["s", "ms"]), str) + check(assert_type(lib.time64("ns").unit, Literal["us", "ns"]), str) + check(assert_type(lib.duration("us"), lib.DurationType), lib.DurationType) + + +def test_decimal_types() -> None: + dec = lib.decimal128(10, 2) + check(assert_type(dec, lib.Decimal128Type), lib.Decimal128Type) + check(assert_type(dec.precision, int), int) + check(assert_type(dec.scale, int), int) + + +def test_list_types() -> None: + lst = lib.list_(lib.int32()) + check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst.value_type, lib.DataType), lib.DataType) + check(assert_type(lst.value_field, lib.Field), lib.Field) + + # With an explicit list_size we can't tell statically whether it is -1. + fixed = lib.list_(lib.int32(), 2) + check(assert_type(fixed, lib.ListType | lib.FixedSizeListType), + lib.FixedSizeListType) + + check(assert_type(lib.large_list(lib.int32()), lib.LargeListType), + lib.LargeListType) + check(assert_type(lib.list_view(lib.int32()), lib.ListViewType), + lib.ListViewType) + + +def test_dictionary_type() -> None: + dict_type = lib.dictionary(lib.int8(), lib.string()) + check(assert_type(dict_type, lib.DictionaryType), lib.DictionaryType) + check(assert_type(dict_type.ordered, bool), bool) + check(assert_type(dict_type.value_type, lib.DataType), lib.DataType) + check(assert_type(dict_type.index_type, "IntegerType"), lib.DataType) + + +def test_map_and_run_end_encoded_types() -> None: + map_type = lib.map_(lib.string(), lib.int64()) + check(assert_type(map_type, lib.MapType), lib.MapType) + check(assert_type(map_type.key_type, lib.DataType), lib.DataType) + check(assert_type(map_type.keys_sorted, bool), bool) + + ree = lib.run_end_encoded(lib.int16(), lib.string()) + check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) + check(assert_type(ree.run_end_type, "Int16Type | Int32Type | Int64Type"), + lib.DataType) + + +def test_types_enum() -> None: + assert lib.int64().id == types.TypesEnum.INT64 + assert lib.list_(lib.int8()).id == types.TypesEnum.LIST + assert (lib.month_day_nano_interval().id + == types.TypesEnum.INTERVAL_MONTH_DAY_NANO) + + +def test_field() -> None: + field = lib.field("a", lib.int32()) + check(assert_type(field, lib.Field), lib.Field) + check(assert_type(field.name, str), str) + check(assert_type(field.type, lib.DataType), lib.DataType) + check(assert_type(field.nullable, bool), bool) + check(assert_type(field.with_nullable(False), lib.Field), lib.Field) + check(assert_type(lib.field("b", "string"), lib.Field), lib.Field) + + +def test_schema() -> None: + schema = lib.schema([("a", lib.int32()), lib.field("b", lib.string())]) + check(assert_type(schema, lib.Schema), lib.Schema) + check(assert_type(schema.names, list[str]), list) + check(assert_type(schema.types, list[lib.DataType]), list) + check(assert_type(schema.field("a"), lib.Field), lib.Field) + check(assert_type(schema[0], lib.Field), lib.Field) + + +def test_type_predicates_narrow() -> None: + lst = _erase(lib.list_(lib.int32())) + assert types.is_list(lst) + check(assert_type(lst, lib.ListType), lib.ListType) + + ts = _erase(lib.timestamp("s")) + assert types.is_timestamp(ts) + check(assert_type(ts, lib.TimestampType), lib.TimestampType) + + dec = _erase(lib.decimal64(5)) + assert types.is_decimal(dec) + check( + assert_type( + dec, + lib.Decimal32Type | lib.Decimal64Type | lib.Decimal128Type + | lib.Decimal256Type, + ), + lib.Decimal64Type, + ) + + # Decimal types subclass FixedSizeBinaryType, but don't match it, so + # they must not be narrowed away when the check fails. + dec128 = _erase(lib.decimal128(5, 2)) + if not types.is_fixed_size_binary(dec128): + assert types.is_decimal(dec128) + check( + assert_type( + dec128, + lib.Decimal32Type | lib.Decimal64Type | lib.Decimal128Type + | lib.Decimal256Type, + ), + lib.Decimal128Type, + ) + + # Run-end encoded types are nested, so this doesn't narrow to Never. + ree = lib.run_end_encoded(lib.int32(), lib.string()) + assert types.is_nested(ree) + check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) + + +def test_invalid_usage() -> None: + if TYPE_CHECKING_INVALID_USAGE: + # A Field needs a type unless it's created from an Arrow C schema. + lib.field("a") # type: ignore[call-overload] + lib.list_(1) # type: ignore[call-overload] + lib.decimal128("10") # type: ignore[arg-type] diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index bf9a6d672280..3b5dce1eda3a 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -2787,7 +2787,7 @@ cdef class FixedSizeBinaryArray(Array): Concrete class for Arrow arrays of a fixed-size binary data type. """ -cdef class Decima32Array(FixedSizeBinaryArray): +cdef class Decimal32Array(FixedSizeBinaryArray): """ Concrete class for Arrow arrays of decimal32 data type. """ @@ -3954,7 +3954,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/fs.py b/python/pyarrow/fs.py index 670ccaaf2455..482650699ea5 100644 --- a/python/pyarrow/fs.py +++ b/python/pyarrow/fs.py @@ -111,7 +111,7 @@ def _ensure_filesystem(filesystem, *, use_mmap=False): else: # handle fsspec-compatible filesystems try: - import fsspec + import fsspec # type: ignore[import-not-found,import-untyped] except ImportError: pass else: 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/tests/test_array.py b/python/pyarrow/tests/test_array.py index 74ff837e16e8..cc42c81835e7 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -18,20 +18,24 @@ from collections.abc import Iterable import datetime import decimal -import hypothesis as h -import hypothesis.strategies as st +import hypothesis as h # type: ignore[import-not-found] +import hypothesis.strategies as st # type: ignore[import-not-found] import itertools import operator -import pytest +import pytest # type: ignore[import-not-found] import struct import subprocess import sys import weakref +from typing import TYPE_CHECKING -try: +if TYPE_CHECKING: import numpy as np -except ImportError: - np = None +else: + try: + import numpy as np + except ImportError: + np = None import pyarrow as pa import pyarrow.tests.strategies as past @@ -73,7 +77,7 @@ def test_constructor_raises(): # This could happen by wrong capitalization. # ARROW-2638: prevent calling extension class constructors directly with pytest.raises(TypeError): - pa.Array([1, 2]) + pa.Array([1, 2]) # type: ignore[reportCallIssue] def test_list_format(): @@ -323,11 +327,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) @@ -766,8 +770,8 @@ def test_array_eq(): @pytest.mark.numpy def test_array_from_buffers(): - values_buf = pa.py_buffer(np.int16([4, 5, 6, 7])) - nulls_buf = pa.py_buffer(np.uint8([0b00001101])) + values_buf = pa.py_buffer(np.array([4, 5, 6, 7], dtype=np.int16())) + nulls_buf = pa.py_buffer(np.array([0b00001101], dtype=np.uint8())) arr = pa.Array.from_buffers(pa.int16(), 4, [nulls_buf, values_buf]) assert arr.type == pa.int16() assert arr.to_pylist() == [4, None, 6, 7] @@ -782,7 +786,9 @@ def test_array_from_buffers(): assert arr.to_pylist() == [None, 6, 7] with pytest.raises(TypeError): - pa.Array.from_buffers(pa.int16(), 3, ['', ''], offset=1) + pa.Array.from_buffers( + pa.int16(), 3, ['', ''], offset=1 # type: ignore[reportArgumentType] + ) def test_string_binary_from_buffers(): @@ -976,7 +982,8 @@ def test_struct_array_from_chunked(): chunked_arr = pa.chunked_array([[1, 2, 3], [4, 5, 6]]) with pytest.raises(TypeError, match="Expected Array"): - pa.StructArray.from_arrays([chunked_arr], ["foo"]) + pa.StructArray.from_arrays( + [chunked_arr], ["foo"]) # type: ignore[reportArgumentType] @pytest.mark.parametrize("offset", (0, 1)) @@ -1296,24 +1303,24 @@ def test_map_from_arrays(): keys = pa.array(pykeys, type='binary') items = pa.array(pyitems, type='i4') - result = pa.MapArray.from_arrays(offsets, keys, items) + result = pa.MapArray.from_arrays(offsets, keys, items) # type: ignore[arg-type] expected = pa.array(pyentries, type=pa.map_(pa.binary(), pa.int32())) assert result.equals(expected) # pass in the type explicitly - result = pa.MapArray.from_arrays(offsets, keys, items, pa.map_( - keys.type, - items.type - )) + result = pa.MapArray.from_arrays(offsets, keys, items, # type: ignore[arg-type] + pa.map_(keys.type, items.type)) assert result.equals(expected) # pass in invalid types with pytest.raises(pa.ArrowTypeError, match='Expected map type, got string'): - pa.MapArray.from_arrays(offsets, keys, items, pa.string()) + pa.MapArray.from_arrays( + offsets, keys, items, pa.string() # type: ignore[arg-type] + ) with pytest.raises(pa.ArrowTypeError, match='Mismatching map items type'): - pa.MapArray.from_arrays(offsets, keys, items, pa.map_( + pa.MapArray.from_arrays(offsets, keys, items, pa.map_( # type: ignore[arg-type] keys.type, # Larger than the original i4 pa.int64() @@ -1351,7 +1358,7 @@ def test_map_from_arrays(): # error if null bitmap and offsets with nulls passed msg1 = 'Ambiguous to specify both validity map and offsets with nulls' with pytest.raises(pa.ArrowInvalid, match=msg1): - pa.MapArray.from_arrays(offsets, keys, items, pa.map_( + pa.MapArray.from_arrays(offsets, keys, items, pa.map_( # type: ignore[arg-type] keys.type, items.type), mask=pa.array([False, True, False], type=pa.bool_()) @@ -2874,7 +2881,7 @@ def test_interval_array_from_relativedelta(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, + pa.MonthDayNano([13, 8, # type: ignore[arg-type] (datetime.timedelta(seconds=1, microseconds=1, minutes=1, hours=1) // datetime.timedelta(microseconds=1)) * 1000])] @@ -2906,7 +2913,7 @@ def test_interval_array_from_tuple(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([1, 2, -3])] + pa.MonthDayNano([1, 2, -3])] # type: ignore[arg-type] expected = pa.array(expected_list) assert arr.equals(expected) assert arr.to_pylist() == expected_list @@ -2927,8 +2934,8 @@ def test_interval_array_from_dateoffset(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, 3661000001001]), - pa.MonthDayNano([0, 0, 0])] + pa.MonthDayNano([13, 8, 3661000001001]), # type: ignore[arg-type] + pa.MonthDayNano([0, 0, 0])] # type: ignore[arg-type] expected = pa.array(expected_list) assert arr.equals(expected) expected_from_pandas = [ @@ -3194,7 +3201,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 @@ -3202,7 +3209,7 @@ def test_buffers_primitive(): assert struct.unpack('hhxxh', buffers[1].to_pybytes()) == (1, 2, 4) - a = pa.array(np.int8([4, 5, 6])) + a = pa.array(np.array([4, 5, 6], dtype=np.int8)) buffers = a.buffers() assert len(buffers) == 2 # No null bitmap from Numpy int array @@ -3288,7 +3295,7 @@ def test_nbytes_size(): def test_invalid_tensor_constructor_repr(): # ARROW-2638: prevent calling extension class constructors directly with pytest.raises(TypeError): - repr(pa.Tensor([1])) + repr(pa.Tensor([1])) # type: ignore[reportCallIssue] def test_invalid_tensor_construction(): @@ -3806,7 +3813,7 @@ def test_array_supported_masks(): with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], - mask=[1.0, 2.0, 3.0, 4.0]) + mask=[1.0, 2.0, 3.0, 4.0]) # type: ignore[reportArgumentType] with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], @@ -4084,11 +4091,11 @@ def test_concat_array_invalid_type(): # ARROW-9920 - do not segfault on non-array input with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays([None]) + pa.concat_arrays([None]) # type: ignore[reportArgumentType] arr = pa.chunked_array([[0, 1], [3, 4]]) with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays(arr) + pa.concat_arrays(arr) # type: ignore[reportArgumentType] @pytest.mark.pandas @@ -4617,7 +4624,7 @@ def test_non_cpu_array(): with pytest.raises(NotImplementedError): [i for i in iter(arr)] with pytest.raises(NotImplementedError): - arr == arr2 + _ = arr == arr2 with pytest.raises(NotImplementedError): arr.is_null() with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index 3ac7be88eaa0..19d33a6e2e05 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -21,13 +21,18 @@ import itertools import math import re +from typing import TYPE_CHECKING, cast import hypothesis as h import pytest -try: + +if TYPE_CHECKING: import numpy as np -except ImportError: - np = None +else: + try: + import numpy as np + except ImportError: + np = None from pyarrow.pandas_compat import _pandas_api # noqa import pyarrow as pa @@ -66,7 +71,7 @@ def __int__(self): class MyBrokenInt: def __int__(self): - 1/0 # MARKER + _ = 1/0 # MARKER def test_iterable_types(): @@ -160,7 +165,7 @@ def test_object_with_getitem(): # https://github.com/apache/arrow/issues/34944 # considered as sequence because of __getitem__, but has no length with pytest.raises(TypeError, match="has no len()"): - pa.array(ObjectWithOnlyGetitem()) + pa.array(ObjectWithOnlyGetitem()) # type: ignore[arg-type] def _as_list(xs): @@ -868,7 +873,7 @@ def test_large_binary_value(ty): assert isinstance(arr, pa.Array) assert arr.type == ty assert len(arr) == 4 - buf = arr[1].as_buffer() + buf = cast(pa.FixedSizeBinaryScalar, arr[1]).as_buffer() assert len(buf) == len(s) * nrepeats @@ -1114,11 +1119,11 @@ def expected_datetime_value(dt): ), ] utcdata = [ - pytz.utc.localize(data[0]), + pytz.utc.localize(cast(datetime.datetime, data[0])), data[1], None, - data[3].astimezone(pytz.utc), - data[4].astimezone(pytz.utc), + cast(datetime.datetime, data[3]).astimezone(pytz.utc), + cast(datetime.datetime, data[4]).astimezone(pytz.utc), ] ty = pa.timestamp(unit, tz=timezone) @@ -1246,9 +1251,9 @@ def test_sequence_timestamp_from_mixed_builtin_and_pandas_datetimes(): None, ] utcdata = [ - data[0].astimezone(pytz.utc), - pytz.utc.localize(data[1]), - data[2].astimezone(pytz.utc), + cast(datetime.datetime, data[0]).astimezone(pytz.utc), + pytz.utc.localize(cast(datetime.datetime, data[1])), + cast(datetime.datetime, data[2]).astimezone(pytz.utc), None, ] @@ -2126,8 +2131,8 @@ def test_map_from_dicts(): assert arr.to_pylist() == expected # With omitted values - data[1] = None - expected[1] = None + data[1] = None # type: ignore[call-overload] + expected[1] = None # type: ignore[call-overload] arr = pa.array(expected, type=pa.map_(pa.binary(), pa.int32())) @@ -2452,6 +2457,7 @@ def test_nested_auto_chunking(ty, char): } +@pytest.mark.numpy @pytest.mark.large_memory def test_array_from_pylist_data_overflow(): # Regression test for ARROW-12983 @@ -2474,6 +2480,7 @@ def test_array_from_pylist_data_overflow(): assert len(arr.chunks) > 1 +@pytest.mark.numpy @pytest.mark.slow @pytest.mark.large_memory def test_array_from_pylist_offset_overflow(): @@ -2498,6 +2505,7 @@ def test_array_from_pylist_offset_overflow(): assert len(arr.chunks) > 1 +@pytest.mark.numpy @parametrize_with_collections_types @pytest.mark.parametrize(('data', 'scalar_data', 'value_type'), [ ([True, False, None], [pa.scalar(True), pa.scalar(False), None], pa.bool_()), @@ -2535,8 +2543,10 @@ def test_array_from_pylist_offset_overflow(): pa.timestamp('us') ), ( - [pa.MonthDayNano([1, -1, -10100])], - [pa.scalar(pa.MonthDayNano([1, -1, -10100]))], + [pa.MonthDayNano([1, -1, -10100])], # type: ignore[call-arg, arg-type] + [pa.scalar( + pa.MonthDayNano([1, -1, -10100]) # type: ignore[call-arg, arg-type] + )], pa.month_day_nano_interval() ), (["a", "b"], [pa.scalar("a"), pa.scalar("b")], pa.string()), diff --git a/python/pyarrow/tests/test_device.py b/python/pyarrow/tests/test_device.py index dc1a51e6d009..00f8bbf720da 100644 --- a/python/pyarrow/tests/test_device.py +++ b/python/pyarrow/tests/test_device.py @@ -59,11 +59,15 @@ def test_copy_to(): batch_copied = batch.copy_to(dest) assert batch_copied.equals(batch) - assert batch_copied["col"].buffers()[1].device == mm.device - assert batch_copied["col"].buffers()[1].address != arr.buffers()[1].address + buffer = batch_copied.column("col").buffers()[1] + assert buffer is not None + assert buffer.device == mm.device + buffer_orig = arr.buffers()[1] + assert buffer_orig is not None + assert buffer.address != buffer_orig.address with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - arr.copy_to(mm.device.device_type) + arr.copy_to(mm.device.device_type) # type: ignore[arg-type] with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - batch.copy_to(mm.device.device_type) + batch.copy_to(mm.device.device_type) # type: ignore[arg-type] diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index b35b5350a814..15f61b1b4002 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 + pass import pyarrow as pa import pyarrow.tests.util as test_util @@ -253,7 +253,7 @@ def test_schema(): child 0, item: int8""" with pytest.raises(TypeError): - pa.schema([None]) + pa.schema([None]) # type: ignore[list-item] def test_schema_weakref(): @@ -542,7 +542,7 @@ def test_schema_equals_invalid_type(): for val in [None, 'string', pa.array([1, 2])]: with pytest.raises(TypeError): - schema.equals(val) + schema.equals(val) # type: ignore[invalid-argument-type] def test_schema_equality_operators(): @@ -588,7 +588,7 @@ def test_schema_get_fields(): with pytest.raises(KeyError): schema.field('other') with pytest.raises(TypeError): - schema.field(0.0) + schema.field(0.0) # type: ignore[arg-type] with pytest.raises(IndexError): schema.field(4) @@ -700,6 +700,7 @@ def test_empty_table(): assert table.schema == schema +@pytest.mark.numpy @pytest.mark.pandas def test_schema_from_pandas(): import pandas as pd @@ -784,7 +785,7 @@ def test_schema_merge(): # raise proper error when passing a non-Schema value with pytest.raises(TypeError): - pa.unify_schemas([a, 1]) + pa.unify_schemas([a, 1]) # type: ignore[list-item] def test_undecodable_metadata(): diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index eca8090d77a9..2ce48b651b14 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -26,15 +26,16 @@ import pyarrow as pa try: - from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix + from scipy.sparse import ( # type: ignore[reportMissingModuleSource] + 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 + import sparse # type: ignore[import-untyped, import-not-found] except ImportError: sparse = None @@ -401,7 +402,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 +444,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, @@ -483,7 +484,8 @@ def test_pydata_sparse_sparse_coo_tensor_roundtrip(dtype_str, arrow_type): shape = (4, 6) dim_names = ("x", "y") - sparse_array = sparse.COO(data=data, coords=coords, shape=shape) + sparse_array = sparse.COO( # type: ignore[reportOptionalMemberAccess] + data=data, coords=coords, shape=shape) sparse_tensor = pa.SparseCOOTensor.from_pydata_sparse(sparse_array, dim_names=dim_names) out_sparse_array = sparse_tensor.to_pydata_sparse() diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index eb02051693f2..60bb0e51902e 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -18,12 +18,13 @@ from collections import OrderedDict from collections.abc import Iterable import sys +from typing import cast import weakref try: import numpy as np except ImportError: - np = None + pass import pytest import pyarrow as pa import pyarrow.compute as pc @@ -417,7 +418,8 @@ def test_to_pandas_empty_table(): table = pa.table(df) result = table.schema.empty_table().to_pandas() assert result.shape == (0, 2) - tm.assert_frame_equal(result, df.iloc[:0]) + expected = df.iloc[:0] + tm.assert_frame_equal(result, expected) @pytest.mark.pandas @@ -485,12 +487,25 @@ def test_chunked_array_unify_dictionaries(): pa.array(["foo", "bar", None, "foo"]).dictionary_encode(), pa.array(["quux", None, "foo"]).dictionary_encode(), ]) - assert arr.chunk(0).dictionary.equals(pa.array(["foo", "bar"])) - assert arr.chunk(1).dictionary.equals(pa.array(["quux", "foo"])) + chunk_0 = arr.chunk(0) + assert isinstance(chunk_0, pa.DictionaryArray) + assert chunk_0.dictionary.equals(pa.array(["foo", "bar"])) + + chunk_1 = arr.chunk(1) + assert isinstance(chunk_1, pa.DictionaryArray) + assert chunk_1.dictionary.equals(pa.array(["quux", "foo"])) + arr = arr.unify_dictionaries() expected_dict = pa.array(["foo", "bar", "quux"]) - assert arr.chunk(0).dictionary.equals(expected_dict) - assert arr.chunk(1).dictionary.equals(expected_dict) + + chunk_0 = arr.chunk(0) + assert isinstance(chunk_0, pa.DictionaryArray) + assert chunk_0.dictionary.equals(expected_dict) + + chunk_1 = arr.chunk(1) + assert isinstance(chunk_1, pa.DictionaryArray) + assert chunk_1.dictionary.equals(expected_dict) + assert arr.to_pylist() == ["foo", "bar", None, "foo", "quux", None, "foo"] @@ -715,7 +730,7 @@ def test_recordbatch_take(): def test_recordbatch_column_sets_private_name(): # ARROW-6429 rb = pa.record_batch([pa.array([1, 2, 3, 4])], names=['a0']) - assert rb[0]._name == 'a0' + assert rb.column(0)._name == 'a0' def test_recordbatch_from_arrays_validate_schema(): @@ -797,7 +812,7 @@ def test_recordbatch_get_field(): batch.field('d') with pytest.raises(TypeError): - batch.field(None) + batch.field(None) # type: ignore[arg-type] with pytest.raises(IndexError): batch.field(4) @@ -818,7 +833,7 @@ def test_recordbatch_select_column(): batch.column('d') with pytest.raises(TypeError): - batch.column(None) + batch.column(None) # type: ignore[arg-type] with pytest.raises(IndexError): batch.column(4) @@ -937,7 +952,10 @@ def test_table_from_struct_array_chunked_array(): [[{"ints": 1}, {"floats": 1.0}]], type=pa.struct([("ints", pa.int32()), ("floats", pa.float32())]), ) - result = pa.Table.from_struct_array(chunked_struct_array) + assert isinstance(chunked_struct_array.type, pa.StructType) + # Cast to the proper type for type checker + struct_chunked_array = cast(pa.ChunkedArray, chunked_struct_array) + result = pa.Table.from_struct_array(struct_chunked_array) assert result.equals(pa.Table.from_arrays( [ pa.array([1, None], type=pa.int32()), @@ -1191,7 +1209,8 @@ def test_recordbatch_to_tensor_null(): batch.to_tensor() result = batch.to_tensor(null_to_nan=True, row_major=False) - x = np.column_stack([arr1, arr2]).astype(np.float64, order="F") + x = np.column_stack([arr1, arr2]).astype( + np.float64, order="F") # type: ignore[no-matching-overload] expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1225,7 +1244,8 @@ def test_recordbatch_to_tensor_null(): ) result = batch.to_tensor(null_to_nan=True, row_major=False) - x = np.column_stack([arr1, arr2]).astype(np.float32, order="F") + x = np.column_stack([arr1, arr2]).astype( + np.float32, order="F") # type: ignore[no-matching-overload] expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1447,7 +1467,7 @@ def test_recordbatchlist_schema_equals(): def test_table_column_sets_private_name(): # ARROW-6429 t = pa.table([pa.array([1, 2, 3, 4])], names=['a0']) - assert t[0]._name == 'a0' + assert t.column(0)._name == 'a0' def test_table_equals(): @@ -1608,7 +1628,8 @@ def test_table_from_arrays_preserves_column_metadata(): field1 = pa.field('field2', pa.int64(), nullable=False) table = pa.Table.from_arrays([arr0, arr1], schema=pa.schema([field0, field1])) - assert b"a" in table.field(0).metadata + field0_metadata = table.field(0).metadata + assert field0_metadata is not None and b"a" in field0_metadata assert table.field(1).nullable is False @@ -1673,7 +1694,7 @@ def test_table_get_field(): table.field('d') with pytest.raises(TypeError): - table.field(None) + table.field(None) # type: ignore[arg-type] with pytest.raises(IndexError): table.field(4) @@ -1694,7 +1715,7 @@ def test_table_select_column(): table.column('d') with pytest.raises(TypeError): - table.column(None) + table.column(None) # type: ignore[arg-type] with pytest.raises(IndexError): table.column(4) @@ -1987,22 +2008,41 @@ def test_table_unify_dictionaries(): table = pa.Table.from_batches([batch1, batch2]) table = table.replace_schema_metadata({b"key1": b"value1"}) - assert table.column(0).chunk(0).dictionary.equals( - pa.array(["foo", "bar"])) - assert table.column(0).chunk(1).dictionary.equals( - pa.array(["quux", "foo"])) - assert table.column(1).chunk(0).dictionary.equals( - pa.array([123, 456, 789])) - assert table.column(1).chunk(1).dictionary.equals( - pa.array([456, 789])) + chunk_0_0 = table.column(0).chunk(0) + assert isinstance(chunk_0_0, pa.DictionaryArray) + assert chunk_0_0.dictionary.equals(pa.array(["foo", "bar"])) + + chunk_0_1 = table.column(0).chunk(1) + assert isinstance(chunk_0_1, pa.DictionaryArray) + assert chunk_0_1.dictionary.equals(pa.array(["quux", "foo"])) + + chunk_1_0 = table.column(1).chunk(0) + assert isinstance(chunk_1_0, pa.DictionaryArray) + assert chunk_1_0.dictionary.equals(pa.array([123, 456, 789])) + + chunk_1_1 = table.column(1).chunk(1) + assert isinstance(chunk_1_1, pa.DictionaryArray) + assert chunk_1_1.dictionary.equals(pa.array([456, 789])) table = table.unify_dictionaries(pa.default_memory_pool()) expected_dict_0 = pa.array(["foo", "bar", "quux"]) expected_dict_1 = pa.array([123, 456, 789]) - assert table.column(0).chunk(0).dictionary.equals(expected_dict_0) - assert table.column(0).chunk(1).dictionary.equals(expected_dict_0) - assert table.column(1).chunk(0).dictionary.equals(expected_dict_1) - assert table.column(1).chunk(1).dictionary.equals(expected_dict_1) + + chunk_0_0 = table.column(0).chunk(0) + assert isinstance(chunk_0_0, pa.DictionaryArray) + assert chunk_0_0.dictionary.equals(expected_dict_0) + + chunk_0_1 = table.column(0).chunk(1) + assert isinstance(chunk_0_1, pa.DictionaryArray) + assert chunk_0_1.dictionary.equals(expected_dict_0) + + chunk_1_0 = table.column(1).chunk(0) + assert isinstance(chunk_1_0, pa.DictionaryArray) + assert chunk_1_0.dictionary.equals(expected_dict_1) + + chunk_1_1 = table.column(1).chunk(1) + assert isinstance(chunk_1_1, pa.DictionaryArray) + assert chunk_1_1.dictionary.equals(expected_dict_1) assert table.to_pydict() == { 'a': ["foo", "bar", None, "foo", "quux", "foo", None, "quux"], @@ -2072,13 +2112,13 @@ def test_concat_tables_invalid_option(): t = pa.Table.from_arrays([list(range(10))], names=('a',)) with pytest.raises(ValueError, match="Invalid promote_options: invalid"): - pa.concat_tables([t, t], promote_options="invalid") + pa.concat_tables([t, t], promote_options="invalid") # type: ignore[arg-type] def test_concat_tables_none_table(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_tables([None]) + pa.concat_tables([None]) # type: ignore[arg-type] @pytest.mark.pandas @@ -2221,7 +2261,7 @@ def test_concat_batches_different_schema(): def test_concat_batches_none_batches(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_batches([None]) + pa.concat_batches([None]) # type: ignore[arg-type] @pytest.mark.parametrize( @@ -2372,7 +2412,7 @@ def test_from_arrays_schema(data, klass): # with different and incompatible schema schema = pa.schema([('strs', pa.utf8()), ('floats', pa.timestamp('s'))]) with pytest.raises((NotImplementedError, TypeError)): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] # Cannot pass both schema and metadata / names with pytest.raises(ValueError): @@ -2477,7 +2517,7 @@ def test_table_from_pydict_arrow_arrays(data, klass): # with different and incompatible schema schema = pa.schema([('strs', pa.utf8()), ('floats', pa.timestamp('s'))]) with pytest.raises((NotImplementedError, TypeError)): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] @pytest.mark.parametrize('data, klass', [ @@ -2494,7 +2534,7 @@ def test_table_from_pydict_schema(data, klass): schema = pa.schema([('strs', pa.utf8()), ('floats', pa.float64()), ('ints', pa.int64())]) with pytest.raises(KeyError, match='ints'): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] # data has columns not present in schema -> ignored schema = pa.schema([('strs', pa.utf8())]) @@ -2698,10 +2738,10 @@ def test_table_factory_function_args_pandas(): def test_factory_functions_invalid_input(): with pytest.raises(TypeError, match="Expected pandas DataFrame, python"): - pa.table("invalid input") + pa.table("invalid input") # type: ignore[arg-type] with pytest.raises(TypeError, match="Expected pandas DataFrame"): - pa.record_batch("invalid input") + pa.record_batch("invalid input") # type: ignore[arg-type] def test_table_repr_to_string(): @@ -2835,8 +2875,8 @@ def test_table_function_unicode_schema(): schema = pa.schema([(col_a, pa.int32()), (col_b, pa.string())]) result = pa.table(d, schema=schema) - assert result[0].chunk(0).equals(pa.array([1, 2, 3], type='int32')) - assert result[1].chunk(0).equals(pa.array(['a', 'b', 'c'], type='string')) + assert result.column(0).chunk(0).equals(pa.array([1, 2, 3], type='int32')) + assert result.column(1).chunk(0).equals(pa.array(['a', 'b', 'c'], type='string')) def test_table_take_vanilla_functionality(): @@ -3768,7 +3808,7 @@ def test_chunked_array_non_cpu(cuda_context, cpu_chunked_array, cuda_chunked_arr # equals() test with pytest.raises(NotImplementedError): - cuda_chunked_array == cuda_chunked_array + cuda_chunked_array == cuda_chunked_array # type: ignore[reportUnusedExpression] # to_pandas() test with pytest.raises(NotImplementedError): @@ -4025,7 +4065,7 @@ def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch, # __dataframe__() test with pytest.raises(NotImplementedError): - from_dataframe(cuda_recordbatch.__dataframe__()) + from_dataframe(cuda_recordbatch.__dataframe__()) # type: ignore[misc] def verify_cuda_table(table, expected_schema): @@ -4224,7 +4264,7 @@ def test_table_non_cpu(cuda_context, cpu_table, cuda_table, # __dataframe__() test with pytest.raises(NotImplementedError): - from_dataframe(cuda_table.__dataframe__()) + from_dataframe(cuda_table.__dataframe__()) # type: ignore[misc] # __reduce__() test with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_tensor.py b/python/pyarrow/tests/test_tensor.py index debb1066280c..c3726fdbbf41 100644 --- a/python/pyarrow/tests/test_tensor.py +++ b/python/pyarrow/tests/test_tensor.py @@ -213,7 +213,7 @@ def test_tensor_memoryview(): dtype = data.dtype lst = data.tolist() tensor = pa.Tensor.from_numpy(data) - m = memoryview(tensor) + m = memoryview(tensor) # type: ignore[reportArgumentType] assert m.format == expected_format assert m.shape == data.shape assert m.strides == data.strides diff --git a/python/pyarrow/tests/test_types.py b/python/pyarrow/tests/test_types.py index 4f251d143343..04a842443989 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -25,16 +25,22 @@ import pytest import hypothesis as h import hypothesis.strategies as st -try: - import hypothesis.extra.pytz as tzst -except ImportError: - tzst = None +from typing import Any, TYPE_CHECKING import weakref -try: +if TYPE_CHECKING: import numpy as np -except ImportError: - np = None + import hypothesis.extra.pytz as tzst +else: + try: + import numpy as np + except ImportError: + np = None + try: + import hypothesis.extra.pytz as tzst + except ImportError: + tzst = None + import pyarrow as pa import pyarrow.types as types import pyarrow.tests.strategies as past @@ -415,7 +421,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) @@ -469,7 +475,7 @@ class BuggyTimezone2(datetime.tzinfo): def tzname(self, dt): return None - def utcoffset(self, dt): + def utcoffset(self, dt): # type: ignore[override] return "one hour" class BuggyTimezone3(datetime.tzinfo): @@ -477,7 +483,7 @@ class BuggyTimezone3(datetime.tzinfo): Wrong timezone name type """ - def tzname(self, dt): + def tzname(self, dt): # type: ignore[override] return 240 def utcoffset(self, dt): @@ -746,13 +752,13 @@ def test_struct_type(): # Neither integer nor string with pytest.raises(TypeError): - ty[None] + ty[None] # type: ignore[reportArgumentType] with pytest.raises(TypeError): - ty.field(None) + ty.field(None) # type: ignore[reportArgumentType] 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 +766,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 +775,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): @@ -876,7 +882,7 @@ def test_dictionary_type(): # invalid index type raises with pytest.raises(TypeError): - pa.dictionary(pa.string(), pa.int64()) + pa.dictionary(pa.string(), pa.int64()) # type: ignore[reportArgumentType] def test_dictionary_ordered_equals(): @@ -965,7 +971,7 @@ def test_run_end_encoded_type(): pa.run_end_encoded(None, pa.utf8()) with pytest.raises(ValueError): - pa.run_end_encoded(pa.int8(), pa.utf8()) + pa.run_end_encoded(pa.int8(), pa.utf8()) # type: ignore[reportArgumentType] @pytest.mark.parametrize('t,check_func', [ @@ -1098,12 +1104,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(): @@ -1141,11 +1149,11 @@ def test_key_value_metadata(): assert m1 != {'a': 'A', 'b': 'C'} with pytest.raises(TypeError): - pa.KeyValueMetadata({'a': 1}) + pa.KeyValueMetadata({'a': 1}) # type: ignore[reportArgumentType] with pytest.raises(TypeError): - pa.KeyValueMetadata({1: 'a'}) + pa.KeyValueMetadata({1: 'a'}) # type: ignore[reportArgumentType] with pytest.raises(TypeError): - pa.KeyValueMetadata(a=1) + pa.KeyValueMetadata(a=1) # type: ignore[reportArgumentType] expected = [(b'a', b'A'), (b'b', b'B')] result = [(k, v) for k, v in m3.items()] @@ -1272,6 +1280,7 @@ def test_field_metadata(): assert f1.metadata is None assert f2.metadata == {} + assert f3.metadata is not None assert f3.metadata[b'bizz'] == b'bazz' @@ -1408,7 +1417,7 @@ def __arrow_c_schema__(self): return self.schema.__arrow_c_schema__() -class SchemaMapping(Mapping): +class SchemaMapping(Mapping[Any, Any]): def __init__(self, schema): self.schema = schema diff --git a/python/pyproject.toml b/python/pyproject.toml index 53b6abb4e02a..2b713a39032a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -122,6 +122,12 @@ exclude = [ "^scripts/", ] +# 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" + # TODO: Enable type checking once stubs are merged [tool.pyright] pythonPlatform = "All" @@ -136,11 +142,29 @@ exclude = [ ] stubPath = "pyarrow-stubs" 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", reportUnnecessaryTypeIgnoreComment = "error" }, + { root = "." }, +] +# In editable install, pyright finds the stubs but compiled modules (e.g. pyarrow.lib) +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 = [ @@ -149,3 +173,25 @@ exclude = [ "examples", "scripts", ] + +# ty doesn't honor `# type: ignore[]` comments, so disable the +# rules that those comments suppress for mypy. +[[tool.ty.overrides]] +include = ["pyarrow-stubs/pyarrow/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); see .pre-commit-config.yaml. +[tool.ruff] +line-length = 88 + +[tool.ruff.lint] +# PYI: the flake8-pyi rules for type stubs. +select = ["E", "F", "W", "PYI"]