From 94046591fd724001a457606293759f7452b3687d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 22 Dec 2025 00:44:32 +0100 Subject: [PATCH 01/93] Add internal types and helpers --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 133 +++ python/pyarrow-stubs/pyarrow/_types.pyi | 966 ++++++++++++++++++ python/pyarrow-stubs/pyarrow/error.pyi | 104 ++ python/pyarrow-stubs/pyarrow/io.pyi | 22 + python/pyarrow-stubs/pyarrow/lib.pyi | 25 + python/pyarrow-stubs/pyarrow/scalar.pyi | 22 + python/pyarrow/fs.py | 2 +- 7 files changed, 1273 insertions(+), 1 deletion(-) create mode 100644 python/pyarrow-stubs/pyarrow/_stubs_typing.pyi create mode 100644 python/pyarrow-stubs/pyarrow/_types.pyi create mode 100644 python/pyarrow-stubs/pyarrow/error.pyi create mode 100644 python/pyarrow-stubs/pyarrow/io.pyi create mode 100644 python/pyarrow-stubs/pyarrow/lib.pyi create mode 100644 python/pyarrow-stubs/pyarrow/scalar.pyi diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi new file mode 100644 index 000000000000..0715012fddc3 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -0,0 +1,133 @@ +# 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 collections.abc import Collection, Iterator, Sequence +from decimal import Decimal +from typing import Any, Literal, Protocol, TypeAlias, TypeVar + +import numpy as np + +from numpy.typing import NDArray + +from pyarrow.lib import BooleanArray, IntegerArray, 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"] +Mask: TypeAlias = ( + Sequence[bool | None] + | NDArray[np.bool_] + | BooleanArray + | ChunkedArray[Any] +) +Indices: TypeAlias = ( + Sequence[int | None] + | NDArray[np.integer[Any]] + | IntegerArray + | ChunkedArray[Any] +) + +PyScalar: TypeAlias = (bool | int | float | Decimal | str | bytes | + dt.date | dt.datetime | dt.time | dt.timedelta) + +_T = TypeVar("_T") +_V = TypeVar("_V", covariant=True) + +SingleOrList: TypeAlias = list[_T] | _T + + +class SupportEq(Protocol): + def __eq__(self, other) -> bool: ... + + +class SupportLt(Protocol): + def __lt__(self, other) -> bool: ... + + +class SupportGt(Protocol): + def __gt__(self, other) -> bool: ... + + +class SupportLe(Protocol): + def __le__(self, other) -> bool: ... + + +class SupportGe(Protocol): + def __ge__(self, other) -> bool: ... + + +FilterTuple: TypeAlias = ( + tuple[str, Literal["=", "==", "!="], SupportEq] + | tuple[str, Literal["<"], SupportLt] + | tuple[str, Literal[">"], SupportGt] + | tuple[str, Literal["<="], SupportLe] + | tuple[str, Literal[">="], SupportGe] + | tuple[str, Literal["in", "not in"], Collection] + | tuple[str, str, Any] # Allow general str for operator to avoid type errors +) + + +class Buffer(Protocol): + ... + + +class SupportPyBuffer(Protocol): + ... + + +class SupportArrowStream(Protocol): + def __arrow_c_stream__(self, requested_schema=None) -> Any: ... + + +class SupportPyArrowArray(Protocol): + def __arrow_array__(self, type=None) -> Any: ... + + +class SupportArrowArray(Protocol): + def __arrow_c_array__(self, requested_schema=None) -> Any: ... + + +class SupportArrowDeviceArray(Protocol): + def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... + + +class SupportArrowSchema(Protocol): + def __arrow_c_schema__(self) -> Any: ... + + +class NullableCollection(Protocol[_V]): # type: ignore[reportInvalidTypeVarUse] + def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... + 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..3d802382ba1c --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -0,0 +1,966 @@ +# 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 # noqa: F401 +import sys + +from collections.abc import Mapping, Sequence, Iterable, Iterator +from decimal import Decimal # noqa: F401 + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +from typing import Any, Generic, Literal + +import numpy as np +import pandas as pd + +from pyarrow._stubs_typing import SupportArrowSchema +from pyarrow.lib import ( # noqa: F401 + Array, + ChunkedArray, + ExtensionArray, + MemoryPool, + MonthDayNano, + Table, +) +from typing_extensions import TypeVar, deprecated + +from .io import Buffer +from .scalar import ExtensionScalar +from ._stubs_typing import TimeUnit + +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: ... + + # Properties that exist on specific subtypes but accessed generically + @property + def list_size(self) -> int: ... + + def __hash__(self) -> int: ... + + def equals(self, other: DataType | str, *, + check_metadata: bool = False) -> bool: ... + + def to_pandas_dtype(self) -> np.generic: ... + + 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) -> Self: ... + + +_AsPyType = TypeVar("_AsPyType") +_DataTypeT = TypeVar("_DataTypeT", bound=DataType) + + +class _BasicDataType(DataType, Generic[_AsPyType]): + ... + + +class NullType(_BasicDataType[None]): + ... + + +class BoolType(_BasicDataType[bool]): + ... + + +class UInt8Type(_BasicDataType[int]): + ... + + +class Int8Type(_BasicDataType[int]): + ... + + +class UInt16Type(_BasicDataType[int]): + ... + + +class Int16Type(_BasicDataType[int]): + ... + + +class UInt32Type(_BasicDataType[int]): + ... + + +class Int32Type(_BasicDataType[int]): + ... + + +class UInt64Type(_BasicDataType[int]): + ... + + +class Int64Type(_BasicDataType[int]): + ... + + +class Float16Type(_BasicDataType[float]): + ... + + +class Float32Type(_BasicDataType[float]): + ... + + +class Float64Type(_BasicDataType[float]): + ... + + +class Date32Type(_BasicDataType[dt.date]): + ... + + +class Date64Type(_BasicDataType[dt.date]): + ... + + +class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): + ... + + +class StringType(_BasicDataType[str]): + ... + + +class LargeStringType(_BasicDataType[str]): + ... + + +class StringViewType(_BasicDataType[str]): + ... + + +class BinaryType(_BasicDataType[bytes]): + ... + + +class LargeBinaryType(_BasicDataType[bytes]): + ... + + +class BinaryViewType(_BasicDataType[bytes]): + ... + + +_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) +_Tz = TypeVar("_Tz", str, None, default=None) + + +class TimestampType(_BasicDataType[int], Generic[_Unit, _Tz]): + + @property + def unit(self) -> _Unit: ... + + @property + def tz(self) -> _Tz: ... + + +_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) + + +class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): + @property + def unit(self) -> _Time32Unit: ... + + +_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) + + +class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): + @property + def unit(self) -> _Time64Unit: ... + + +class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): + @property + def unit(self) -> _Unit: ... + + +class FixedSizeBinaryType(_BasicDataType[Decimal]): + ... + + +_Precision = TypeVar("_Precision", default=Any) +_Scale = TypeVar("_Scale", default=Any) + + +class Decimal32Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): + @property + def precision(self) -> _Precision: ... + + @property + def scale(self) -> _Scale: ... + + +class Decimal64Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): + @property + def precision(self) -> _Precision: ... + + @property + def scale(self) -> _Scale: ... + + +class Decimal128Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): + @property + def precision(self) -> _Precision: ... + + @property + def scale(self) -> _Scale: ... + + +class Decimal256Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): + @property + def precision(self) -> _Precision: ... + + @property + def scale(self) -> _Scale: ... + + +class ListType(DataType, Generic[_DataTypeT]): + @property + def value_field(self) -> Field[_DataTypeT]: ... + + @property + def value_type(self) -> _DataTypeT: ... + + +class LargeListType(DataType, Generic[_DataTypeT]): + @property + def value_field(self) -> Field[_DataTypeT]: ... + @property + def value_type(self) -> _DataTypeT: ... + + +class ListViewType(DataType, Generic[_DataTypeT]): + @property + def value_field(self) -> Field[_DataTypeT]: ... + + @property + def value_type(self) -> _DataTypeT: ... + + +class LargeListViewType(DataType, Generic[_DataTypeT]): + @property + def value_field(self) -> Field[_DataTypeT]: ... + + @property + def value_type(self) -> _DataTypeT: ... + + +class FixedSizeListType(DataType, Generic[_DataTypeT, _Size]): + @property + def value_field(self) -> Field[_DataTypeT]: ... + + @property + def value_type(self) -> _DataTypeT: ... + + @property + def list_size(self) -> int: ... + + +class DictionaryMemo(_Weakrefable): + ... + + +_IndexT = TypeVar( + "_IndexT", + UInt8Type, + Int8Type, + UInt16Type, + Int16Type, + UInt32Type, + Int32Type, + UInt64Type, + Int64Type, +) +_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType) +_ValueT = TypeVar("_ValueT", bound=DataType) +_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) + + +class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): + @property + def ordered(self) -> _Ordered: ... + + @property + def index_type(self) -> _IndexT: ... + + @property + def value_type(self) -> _BasicValueT: ... + + +_K = TypeVar("_K", bound=DataType) + + +class MapType(DataType, Generic[_K, _ValueT, _Ordered]): + @property + def key_field(self) -> Field[_K]: ... + + @property + def key_type(self) -> _K: ... + + @property + def item_field(self) -> Field[_ValueT]: ... + + @property + def item_type(self) -> _ValueT: ... + + @property + def keys_sorted(self) -> _Ordered: ... + + +_Size = TypeVar("_Size", default=int) + + +class StructType(DataType): + def get_field_index(self, name: str) -> int: ... + + def field(self, i: int | str) -> Field: ... + + def get_all_field_indices(self, name: str) -> 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 = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) + + +class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> _BasicValueT: ... + + +_StorageT = TypeVar("_StorageT", bound=Array | ChunkedArray) + + +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: _StorageT) -> _StorageT: ... + + +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, Generic[_ValueT]): + @property + def value_type(self) -> _ValueT: ... + + @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: ExtensionType) -> None: ... + + +def unregister_extension_type(type_name: str) -> None: ... + + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, __arg0__: Mapping[str | bytes, str | bytes] + | Iterable[tuple[str, str]] + | KeyValueMetadata + | None = None, **kwargs: str + ) -> None: ... + + def equals(self, other: KeyValueMetadata) -> bool: ... + + def __len__(self) -> int: ... + + def __contains__(self, /, __key: object) -> bool: ... # type: ignore[override] + + def __getitem__(self, /, __key: Any) -> Any: ... # type: ignore[override] + + def __iter__(self) -> Iterator[bytes]: ... + + def get_all(self, key: str) -> list[bytes]: ... + + def to_dict(self) -> dict[bytes, bytes]: ... + + +class Field(_Weakrefable, Generic[_DataTypeT]): + 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) -> _DataTypeT: ... + def with_metadata(self, metadata: dict[bytes | str, bytes | str] | + Mapping[bytes | str, bytes | str] | Any) -> Self: ... + + def remove_metadata(self) -> Self: ... + + def with_type(self, new_type: DataType) -> Field: ... + + def with_name(self, name: str) -> Self: ... + + def with_nullable(self, nullable: bool) -> Field[_DataTypeT]: ... + + 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) -> Self: ... + + +class Schema(_Weakrefable): + def __len__(self) -> int: ... + + def __getitem__(self, key: str | int) -> Field: ... + + _field = __getitem__ + def __iter__(self) -> Iterator[Field]: ... + + def __hash__(self) -> int: ... + + def __sizeof__(self) -> int: ... + @property + def pandas_metadata(self) -> dict: ... + + @property + def names(self) -> list[str]: ... + + @property + def types(self) -> list[DataType]: ... + + @property + def metadata(self) -> dict[bytes, bytes]: ... + + 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) -> Field: ... + + def get_field_index(self, name: str) -> int: ... + + def get_all_field_indices(self, name: str) -> 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: dict) -> Schema: ... + + def with_metadata(self, metadata: dict) -> 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 | None = None, + ) -> 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: Sequence[Schema], + *, + promote_options: Literal["default", "permissive"] = "default" +) -> Schema: ... + + +def field( + name: SupportArrowSchema | str | Any, type: _DataTypeT | str | None = None, + nullable: bool = ..., + metadata: dict[Any, Any] | None = None +) -> Field[_DataTypeT] | Field[Any]: ... + + +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: _Unit | str, tz: _Tz | None = None) -> TimestampType[_Unit, _Tz]: ... + + +def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... + + +def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... + + +def duration(unit: _Unit | str) -> DurationType[_Unit]: ... + + +def month_day_nano_interval() -> MonthDayNanoIntervalType: ... + + +def date32() -> Date32Type: ... + + +def date64() -> Date64Type: ... + + +def float16() -> Float16Type: ... + + +def float32() -> Float32Type: ... + + +def float64() -> Float64Type: ... + + +def decimal32(precision: _Precision, scale: _Scale | + None = None) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... + + +def decimal64(precision: _Precision, scale: _Scale | + None = None) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... + + +def decimal128(precision: _Precision, scale: _Scale | + None = None) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... + + +def decimal256(precision: _Precision, scale: _Scale | + None = None) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... + + +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: ... + + +def list_( + value_type: _DataTypeT | Field[_DataTypeT] | None = None, + list_size: Literal[-1] | _Size | None = None +) -> ListType[_DataTypeT] | FixedSizeListType[_DataTypeT, _Size]: ... + + +def large_list(value_type: _DataTypeT | + Field[_DataTypeT] | None = None) -> LargeListType[_DataTypeT]: ... + + +def list_view(value_type: _DataTypeT | + Field[_DataTypeT] | None = None) -> ListViewType[_DataTypeT]: ... + + +def large_list_view( + value_type: _DataTypeT | Field[_DataTypeT] | None = None +) -> LargeListViewType[_DataTypeT]: ... + + +def map_( + key_type: _K | Field | str | None = None, + item_type: _ValueT | Field | str | None = None, + keys_sorted: bool | None = None +) -> MapType[_K, _ValueT, Literal[False]]: ... + + +def dictionary( + index_type: _IndexT | str, + value_type: _BasicValueT | str, + ordered: _Ordered | None = None +) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... + + +def struct( + fields: Iterable[ + Field[Any] + | tuple[str, Field[Any] | None] + | tuple[str, DataType | None] + ] | Mapping[str, Field[Any] | DataType | None], +) -> StructType: ... + + +def sparse_union( + child_fields: list[Field[Any]], type_codes: list[int] | None = None +) -> SparseUnionType: ... + + +def dense_union( + child_fields: list[Field[Any]], type_codes: list[int] | None = None +) -> DenseUnionType: ... + + +def union( + child_fields: list[Field[Any]], mode: Literal["sparse", "dense"] | int | str, + type_codes: list[int] | None = None) -> SparseUnionType | DenseUnionType: ... + + +def run_end_encoded( + run_end_type: _RunEndType | str | None, value_type: _BasicValueT | str | None +) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... + + +def json_(storage_type: DataType = ...) -> JsonType: ... + + +def uuid() -> UuidType: ... + + +def fixed_shape_tensor( + value_type: _ValueT, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + permutation: Sequence[int] | None = None, +) -> FixedShapeTensorType[_ValueT]: ... + + +def bool8() -> Bool8Type: ... + + +def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... + + +def type_for_alias(name: Any) -> DataType: ... + + +def schema( + fields: ( + Iterable[Field[Any]] + | Iterable[tuple[str, DataType | str | None]] + | Mapping[Any, DataType | str | None] + ), + metadata: Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] | None = None, +) -> Schema: ... + + +def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... + + +__all__ = [ + "_Weakrefable", + "_Metadata", + "DataType", + "_BasicDataType", + "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", + "_Unit", + "_Tz", + "_Time32Unit", + "_Time64Unit", + "_DataTypeT", +] diff --git a/python/pyarrow-stubs/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/error.pyi new file mode 100644 index 000000000000..eac936afcb53 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/error.pyi @@ -0,0 +1,104 @@ +# 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 + + +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, exc_value, exc_tb) -> None: ... + def __dealloc__(self) -> 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..467ec48cc76d --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/io.pyi @@ -0,0 +1,22 @@ +# 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.""" + +from typing import Any + +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..775434be2ea4 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/lib.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 for pyarrow.lib C extension module. + +Complete type annotations will be added in subsequent PRs. +""" + +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/pyarrow/scalar.pyi b/python/pyarrow-stubs/pyarrow/scalar.pyi new file mode 100644 index 000000000000..467ec48cc76d --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/scalar.pyi @@ -0,0 +1,22 @@ +# 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.""" + +from typing import Any + +def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow/fs.py b/python/pyarrow/fs.py index 670ccaaf2455..f055c5081826 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-untyped] except ImportError: pass else: From 5f7cdc8b4eb0cf806de912843dd420ddea79c10d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Feb 2026 21:08:58 +0100 Subject: [PATCH 02/93] Update python/pyarrow-stubs/pyarrow/_stubs_typing.pyi Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com> --- python/pyarrow-stubs/pyarrow/_stubs_typing.pyi | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 0715012fddc3..0e4013c2b9e6 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -46,16 +46,29 @@ Compression: TypeAlias = Literal[ NullEncoding: TypeAlias = Literal["mask", "encode"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] +from pyarrow import lib + +IntegerType: TypeAlias = ( + lib.Int8Type + | lib.Int16Type + | lib.Int32Type + | lib.Int64Type + | lib.UInt8Type + | lib.UInt16Type + | lib.UInt32Type + | lib.UInt64Type +) + Mask: TypeAlias = ( Sequence[bool | None] | NDArray[np.bool_] - | BooleanArray + | lib.Array[lib.Scalar[lib.BoolType]] | ChunkedArray[Any] ) Indices: TypeAlias = ( Sequence[int | None] | NDArray[np.integer[Any]] - | IntegerArray + | lib.Array[lib.Scalar[IntegerType]] | ChunkedArray[Any] ) From 610c48d98ae37deefb90b5be4eb6ef91caadf290 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Mar 2026 22:16:26 +0100 Subject: [PATCH 03/93] post rebase changes --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 14 ++--- python/pyarrow-stubs/pyarrow/_types.pyi | 63 +++++++++---------- python/pyarrow-stubs/pyarrow/error.pyi | 8 +-- python/pyarrow-stubs/pyarrow/io.pyi | 10 ++- python/pyarrow-stubs/pyarrow/scalar.pyi | 10 ++- 5 files changed, 56 insertions(+), 49 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 0e4013c2b9e6..2b8230758958 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -25,6 +25,7 @@ import numpy as np from numpy.typing import NDArray +from pyarrow import lib from pyarrow.lib import BooleanArray, IntegerArray, ChunkedArray ArrayLike: TypeAlias = Any @@ -46,7 +47,6 @@ Compression: TypeAlias = Literal[ NullEncoding: TypeAlias = Literal["mask", "encode"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] -from pyarrow import lib IntegerType: TypeAlias = ( lib.Int8Type @@ -82,23 +82,23 @@ SingleOrList: TypeAlias = list[_T] | _T class SupportEq(Protocol): - def __eq__(self, other) -> bool: ... + def __eq__(self, other: object) -> bool: ... class SupportLt(Protocol): - def __lt__(self, other) -> bool: ... + def __lt__(self, other: object) -> bool: ... class SupportGt(Protocol): - def __gt__(self, other) -> bool: ... + def __gt__(self, other: object) -> bool: ... class SupportLe(Protocol): - def __le__(self, other) -> bool: ... + def __le__(self, other: object) -> bool: ... class SupportGe(Protocol): - def __ge__(self, other) -> bool: ... + def __ge__(self, other: object) -> bool: ... FilterTuple: TypeAlias = ( @@ -140,7 +140,7 @@ class SupportArrowSchema(Protocol): def __arrow_c_schema__(self) -> Any: ... -class NullableCollection(Protocol[_V]): # type: ignore[reportInvalidTypeVarUse] +class NullableCollection(Protocol[_V]): def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... 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 index 3d802382ba1c..d646551eecb3 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -16,22 +16,18 @@ # under the License. import datetime as dt # noqa: F401 -import sys from collections.abc import Mapping, Sequence, Iterable, Iterator from decimal import Decimal # noqa: F401 - -if sys.version_info >= (3, 11): - from typing import Self -else: - from typing_extensions import Self - from typing import Any, Generic, Literal import numpy as np import pandas as pd -from pyarrow._stubs_typing import SupportArrowSchema +from typing_extensions import Self, TypeVar, deprecated + +from pyarrow._stubs_typing import SupportArrowSchema, TimeUnit +from pyarrow.io import Buffer from pyarrow.lib import ( # noqa: F401 Array, ChunkedArray, @@ -40,11 +36,7 @@ from pyarrow.lib import ( # noqa: F401 MonthDayNano, Table, ) -from typing_extensions import TypeVar, deprecated - -from .io import Buffer -from .scalar import ExtensionScalar -from ._stubs_typing import TimeUnit +from pyarrow.scalar import ExtensionScalar class _Weakrefable: ... @@ -226,7 +218,12 @@ class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): def unit(self) -> _Unit: ... -class FixedSizeBinaryType(_BasicDataType[Decimal]): +_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) + + +class FixedSizeBinaryType( + _BasicDataType[_FixedSizeBinaryAsPyType], Generic[_FixedSizeBinaryAsPyType] +): ... @@ -234,7 +231,7 @@ _Precision = TypeVar("_Precision", default=Any) _Scale = TypeVar("_Scale", default=Any) -class Decimal32Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): +class Decimal32Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... @@ -242,7 +239,7 @@ class Decimal32Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): def scale(self) -> _Scale: ... -class Decimal64Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): +class Decimal64Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... @@ -250,7 +247,7 @@ class Decimal64Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): def scale(self) -> _Scale: ... -class Decimal128Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): +class Decimal128Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... @@ -258,7 +255,7 @@ class Decimal128Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): def scale(self) -> _Scale: ... -class Decimal256Type(FixedSizeBinaryType, Generic[_Precision, _Scale]): +class Decimal256Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... @@ -491,7 +488,7 @@ def unregister_extension_type(type_name: str) -> None: ... class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( self, __arg0__: Mapping[str | bytes, str | bytes] - | Iterable[tuple[str, str]] + | Iterable[tuple[str | bytes, str | bytes]] | KeyValueMetadata | None = None, **kwargs: str ) -> None: ... @@ -500,9 +497,9 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __len__(self) -> int: ... - def __contains__(self, /, __key: object) -> bool: ... # type: ignore[override] + def __contains__(self, /, __key: object) -> bool: ... - def __getitem__(self, /, __key: Any) -> Any: ... # type: ignore[override] + def __getitem__(self, /, __key: Any) -> Any: ... def __iter__(self) -> Iterator[bytes]: ... @@ -636,7 +633,7 @@ def unify_schemas( def field( name: SupportArrowSchema | str | Any, type: _DataTypeT | str | None = None, - nullable: bool = ..., + nullable: bool = True, metadata: dict[Any, Any] | None = None ) -> Field[_DataTypeT] | Field[Any]: ... @@ -702,20 +699,20 @@ def float32() -> Float32Type: ... def float64() -> Float64Type: ... -def decimal32(precision: _Precision, scale: _Scale | - None = None) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... +def decimal32(precision: _Precision, + scale: _Scale | Literal[0] = 0) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... -def decimal64(precision: _Precision, scale: _Scale | - None = None) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... +def decimal64(precision: _Precision, + scale: _Scale | Literal[0] = 0) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... -def decimal128(precision: _Precision, scale: _Scale | - None = None) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... +def decimal128(precision: _Precision, + scale: _Scale | Literal[0] = 0) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... -def decimal256(precision: _Precision, scale: _Scale | - None = None) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... +def decimal256(precision: _Precision, + scale: _Scale | Literal[0] = 0) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... def string() -> StringType: ... @@ -724,7 +721,7 @@ def string() -> StringType: ... utf8 = string -def binary(length: Literal[-1] | int = ...) -> BinaryType | FixedSizeBinaryType: ... +def binary(length: Literal[-1] | int = ...) -> BinaryType | FixedSizeBinaryType[bytes]: ... def large_binary() -> LargeBinaryType: ... @@ -764,8 +761,8 @@ def large_list_view( def map_( key_type: _K | Field | str | None = None, item_type: _ValueT | Field | str | None = None, - keys_sorted: bool | None = None -) -> MapType[_K, _ValueT, Literal[False]]: ... + keys_sorted: _Ordered | None = None +) -> MapType[_K, _ValueT, _Ordered]: ... def dictionary( diff --git a/python/pyarrow-stubs/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/error.pyi index eac936afcb53..6a010071de08 100644 --- a/python/pyarrow-stubs/pyarrow/error.pyi +++ b/python/pyarrow-stubs/pyarrow/error.pyi @@ -15,12 +15,7 @@ # 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 typing_extensions import Self class ArrowException(Exception): @@ -80,7 +75,6 @@ have_signal_refcycle: bool class SignalStopHandler: def __enter__(self) -> Self: ... def __exit__(self, exc_type, exc_value, exc_tb) -> None: ... - def __dealloc__(self) -> None: ... @property def stop_token(self) -> StopToken: ... diff --git a/python/pyarrow-stubs/pyarrow/io.pyi b/python/pyarrow-stubs/pyarrow/io.pyi index 467ec48cc76d..856fb093a6cd 100644 --- a/python/pyarrow-stubs/pyarrow/io.pyi +++ b/python/pyarrow-stubs/pyarrow/io.pyi @@ -15,8 +15,16 @@ # specific language governing permissions and limitations # under the License. -"""Placeholder stub - complete annotations in future PR.""" +"""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/scalar.pyi b/python/pyarrow-stubs/pyarrow/scalar.pyi index 467ec48cc76d..384507d2a2b6 100644 --- a/python/pyarrow-stubs/pyarrow/scalar.pyi +++ b/python/pyarrow-stubs/pyarrow/scalar.pyi @@ -15,8 +15,16 @@ # specific language governing permissions and limitations # under the License. -"""Placeholder stub - complete annotations in future PR.""" +"""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: ... From 7bd5d7ce38229d6eef6e88c7b1574ff230dc2eea Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:27:18 +0100 Subject: [PATCH 04/93] linting and formatting --- python/pyarrow-stubs/pyarrow/__init__.pyi | 8 +- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 55 +- python/pyarrow-stubs/pyarrow/_types.pyi | 481 ++++-------------- python/pyarrow-stubs/pyarrow/error.pyi | 54 +- python/pyarrow-stubs/pyarrow/io.pyi | 11 +- python/pyarrow-stubs/pyarrow/lib.pyi | 6 +- python/pyarrow-stubs/pyarrow/scalar.pyi | 11 +- 7 files changed, 134 insertions(+), 492 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__init__.pyi b/python/pyarrow-stubs/pyarrow/__init__.pyi index ccec8d5abc07..9c7120c7f807 100644 --- a/python/pyarrow-stubs/pyarrow/__init__.pyi +++ b/python/pyarrow-stubs/pyarrow/__init__.pyi @@ -15,11 +15,9 @@ # 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. -""" +# Type stubs for PyArrow. +# This is a placeholder stub file. +# Complete type annotations will be added in subsequent PRs. from typing import Any diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 2b8230758958..1d5f7600c4b7 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -26,7 +26,7 @@ import numpy as np from numpy.typing import NDArray from pyarrow import lib -from pyarrow.lib import BooleanArray, IntegerArray, ChunkedArray +from pyarrow.lib import ChunkedArray ArrayLike: TypeAlias = Any ScalarLike: TypeAlias = Any @@ -72,34 +72,38 @@ Indices: TypeAlias = ( | ChunkedArray[Any] ) -PyScalar: TypeAlias = (bool | int | float | Decimal | str | bytes | - dt.date | dt.datetime | dt.time | dt.timedelta) +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) _T = TypeVar("_T") _V = TypeVar("_V", covariant=True) SingleOrList: TypeAlias = list[_T] | _T - class SupportEq(Protocol): - def __eq__(self, other: object) -> bool: ... - + def __eq__(self, other: object, /) -> bool: ... class SupportLt(Protocol): - def __lt__(self, other: object) -> bool: ... - + def __lt__(self, other: object, /) -> bool: ... class SupportGt(Protocol): - def __gt__(self, other: object) -> bool: ... - + def __gt__(self, other: object, /) -> bool: ... class SupportLe(Protocol): - def __le__(self, other: object) -> bool: ... - + def __le__(self, other: object, /) -> bool: ... class SupportGe(Protocol): - def __ge__(self, other: object) -> bool: ... - + def __ge__(self, other: object, /) -> bool: ... FilterTuple: TypeAlias = ( tuple[str, Literal["=", "==", "!="], SupportEq] @@ -111,35 +115,24 @@ FilterTuple: TypeAlias = ( | tuple[str, str, Any] # Allow general str for operator to avoid type errors ) - -class Buffer(Protocol): - ... - - -class SupportPyBuffer(Protocol): - ... - +class Buffer(Protocol): ... +class SupportPyBuffer(Protocol): ... class SupportArrowStream(Protocol): - def __arrow_c_stream__(self, requested_schema=None) -> Any: ... - + def __arrow_c_stream__(self, requested_schema=None, /) -> Any: ... class SupportPyArrowArray(Protocol): - def __arrow_array__(self, type=None) -> Any: ... - + def __arrow_array__(self, type=None, /) -> Any: ... class SupportArrowArray(Protocol): - def __arrow_c_array__(self, requested_schema=None) -> Any: ... - + def __arrow_c_array__(self, requested_schema=None, /) -> Any: ... class SupportArrowDeviceArray(Protocol): - def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... - + def __arrow_c_device_array__(self, requested_schema=None, /, **kwargs) -> Any: ... class SupportArrowSchema(Protocol): def __arrow_c_schema__(self) -> Any: ... - class NullableCollection(Protocol[_V]): def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... def __len__(self) -> int: ... diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index d646551eecb3..1fd5c5002d1e 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -38,276 +38,157 @@ from pyarrow.lib import ( # noqa: F401 ) from pyarrow.scalar import ExtensionScalar -class _Weakrefable: - ... - - -class _Metadata(_Weakrefable): - ... - +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: ... # Properties that exist on specific subtypes but accessed generically @property def list_size(self) -> int: ... - def __hash__(self) -> int: ... - - def equals(self, other: DataType | str, *, - check_metadata: bool = False) -> bool: ... - + def equals( + self, other: DataType | str, *, check_metadata: bool = False + ) -> bool: ... def to_pandas_dtype(self) -> np.generic: ... - 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) -> Self: ... - _AsPyType = TypeVar("_AsPyType") _DataTypeT = TypeVar("_DataTypeT", bound=DataType) - -class _BasicDataType(DataType, Generic[_AsPyType]): - ... - - -class NullType(_BasicDataType[None]): - ... - - -class BoolType(_BasicDataType[bool]): - ... - - -class UInt8Type(_BasicDataType[int]): - ... - - -class Int8Type(_BasicDataType[int]): - ... - - -class UInt16Type(_BasicDataType[int]): - ... - - -class Int16Type(_BasicDataType[int]): - ... - - -class UInt32Type(_BasicDataType[int]): - ... - - -class Int32Type(_BasicDataType[int]): - ... - - -class UInt64Type(_BasicDataType[int]): - ... - - -class Int64Type(_BasicDataType[int]): - ... - - -class Float16Type(_BasicDataType[float]): - ... - - -class Float32Type(_BasicDataType[float]): - ... - - -class Float64Type(_BasicDataType[float]): - ... - - -class Date32Type(_BasicDataType[dt.date]): - ... - - -class Date64Type(_BasicDataType[dt.date]): - ... - - -class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): - ... - - -class StringType(_BasicDataType[str]): - ... - - -class LargeStringType(_BasicDataType[str]): - ... - - -class StringViewType(_BasicDataType[str]): - ... - - -class BinaryType(_BasicDataType[bytes]): - ... - - -class LargeBinaryType(_BasicDataType[bytes]): - ... - - -class BinaryViewType(_BasicDataType[bytes]): - ... - +class _BasicDataType(DataType, Generic[_AsPyType]): ... +class NullType(_BasicDataType[None]): ... +class BoolType(_BasicDataType[bool]): ... +class UInt8Type(_BasicDataType[int]): ... +class Int8Type(_BasicDataType[int]): ... +class UInt16Type(_BasicDataType[int]): ... +class Int16Type(_BasicDataType[int]): ... +class UInt32Type(_BasicDataType[int]): ... +class Int32Type(_BasicDataType[int]): ... +class UInt64Type(_BasicDataType[int]): ... +class Int64Type(_BasicDataType[int]): ... +class Float16Type(_BasicDataType[float]): ... +class Float32Type(_BasicDataType[float]): ... +class Float64Type(_BasicDataType[float]): ... +class Date32Type(_BasicDataType[dt.date]): ... +class Date64Type(_BasicDataType[dt.date]): ... +class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... +class StringType(_BasicDataType[str]): ... +class LargeStringType(_BasicDataType[str]): ... +class StringViewType(_BasicDataType[str]): ... +class BinaryType(_BasicDataType[bytes]): ... +class LargeBinaryType(_BasicDataType[bytes]): ... +class BinaryViewType(_BasicDataType[bytes]): ... _Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) _Tz = TypeVar("_Tz", str, None, default=None) - class TimestampType(_BasicDataType[int], Generic[_Unit, _Tz]): - @property def unit(self) -> _Unit: ... - @property def tz(self) -> _Tz: ... - _Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) - class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): @property def unit(self) -> _Time32Unit: ... - _Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) - class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): @property def unit(self) -> _Time64Unit: ... - class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): @property def unit(self) -> _Unit: ... - _FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) - class FixedSizeBinaryType( - _BasicDataType[_FixedSizeBinaryAsPyType], Generic[_FixedSizeBinaryAsPyType] -): - ... - + _BasicDataType[_FixedSizeBinaryAsPyType] +): ... _Precision = TypeVar("_Precision", default=Any) _Scale = TypeVar("_Scale", default=Any) - class Decimal32Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... - @property def scale(self) -> _Scale: ... - class Decimal64Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... - @property def scale(self) -> _Scale: ... - class Decimal128Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... - @property def scale(self) -> _Scale: ... - class Decimal256Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): @property def precision(self) -> _Precision: ... - @property def scale(self) -> _Scale: ... - class ListType(DataType, Generic[_DataTypeT]): @property def value_field(self) -> Field[_DataTypeT]: ... - @property def value_type(self) -> _DataTypeT: ... - class LargeListType(DataType, Generic[_DataTypeT]): @property def value_field(self) -> Field[_DataTypeT]: ... @property def value_type(self) -> _DataTypeT: ... - class ListViewType(DataType, Generic[_DataTypeT]): @property def value_field(self) -> Field[_DataTypeT]: ... - @property def value_type(self) -> _DataTypeT: ... - class LargeListViewType(DataType, Generic[_DataTypeT]): @property def value_field(self) -> Field[_DataTypeT]: ... - @property def value_type(self) -> _DataTypeT: ... - class FixedSizeListType(DataType, Generic[_DataTypeT, _Size]): @property def value_field(self) -> Field[_DataTypeT]: ... - @property def value_type(self) -> _DataTypeT: ... - @property def list_size(self) -> int: ... - -class DictionaryMemo(_Weakrefable): - ... - +class DictionaryMemo(_Weakrefable): ... _IndexT = TypeVar( "_IndexT", @@ -324,287 +205,197 @@ _BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType) _ValueT = TypeVar("_ValueT", bound=DataType) _Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) - class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): @property def ordered(self) -> _Ordered: ... - @property def index_type(self) -> _IndexT: ... - @property def value_type(self) -> _BasicValueT: ... - _K = TypeVar("_K", bound=DataType) - class MapType(DataType, Generic[_K, _ValueT, _Ordered]): @property def key_field(self) -> Field[_K]: ... - @property def key_type(self) -> _K: ... - @property def item_field(self) -> Field[_ValueT]: ... - @property def item_type(self) -> _ValueT: ... - @property def keys_sorted(self) -> _Ordered: ... - _Size = TypeVar("_Size", default=int) - class StructType(DataType): def get_field_index(self, name: str) -> int: ... - def field(self, i: int | str) -> Field: ... - def get_all_field_indices(self, name: str) -> 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 = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) - class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): @property def run_end_type(self) -> _RunEndType: ... @property def value_type(self) -> _BasicValueT: ... - _StorageT = TypeVar("_StorageT", bound=Array | ChunkedArray) - 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: _StorageT) -> _StorageT: ... - 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: ... - + cls, storage_type: DataType, serialized: bytes + ) -> Self: ... class FixedShapeTensorType(BaseExtensionType, Generic[_ValueT]): @property def value_type(self) -> _ValueT: ... - @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 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: ExtensionType) -> None: ... - - def unregister_extension_type(type_name: str) -> None: ... - class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( - self, __arg0__: Mapping[str | bytes, str | bytes] + self, + __arg0__: Mapping[str | bytes, str | bytes] | Iterable[tuple[str | bytes, str | bytes]] | KeyValueMetadata - | None = None, **kwargs: str + | None = None, + **kwargs: str, ) -> None: ... - def equals(self, other: KeyValueMetadata) -> bool: ... - def __len__(self) -> int: ... - def __contains__(self, /, __key: object) -> bool: ... - def __getitem__(self, /, __key: Any) -> Any: ... - def __iter__(self) -> Iterator[bytes]: ... - def get_all(self, key: str) -> list[bytes]: ... - def to_dict(self) -> dict[bytes, bytes]: ... - class Field(_Weakrefable, Generic[_DataTypeT]): 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) -> _DataTypeT: ... - def with_metadata(self, metadata: dict[bytes | str, bytes | str] | - Mapping[bytes | str, bytes | str] | Any) -> Self: ... - + def with_metadata( + self, + metadata: dict[bytes | str, bytes | str] + | Mapping[bytes | str, bytes | str] + | Any, + ) -> Self: ... def remove_metadata(self) -> Self: ... - def with_type(self, new_type: DataType) -> Field: ... - def with_name(self, name: str) -> Self: ... - def with_nullable(self, nullable: bool) -> Field[_DataTypeT]: ... - 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) -> Self: ... - class Schema(_Weakrefable): def __len__(self) -> int: ... - def __getitem__(self, key: str | int) -> Field: ... _field = __getitem__ def __iter__(self) -> Iterator[Field]: ... - def __hash__(self) -> int: ... - def __sizeof__(self) -> int: ... @property def pandas_metadata(self) -> dict: ... - @property def names(self) -> list[str]: ... - @property def types(self) -> list[DataType]: ... - @property def metadata(self) -> dict[bytes, bytes]: ... - 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 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) -> Field: ... - def get_field_index(self, name: str) -> int: ... - def get_all_field_indices(self, name: str) -> 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: dict) -> Schema: ... - def with_metadata(self, metadata: dict) -> Schema: ... - def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... - def remove_metadata(self) -> Schema: ... - def to_string( self, truncate_metadata: bool = True, @@ -612,218 +403,126 @@ class Schema(_Weakrefable): show_schema_metadata: bool = True, element_size_limit: int | None = None, ) -> 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: Sequence[Schema], *, - promote_options: Literal["default", "permissive"] = "default" + promote_options: Literal["default", "permissive"] = "default", ) -> Schema: ... - - def field( - name: SupportArrowSchema | str | Any, type: _DataTypeT | str | None = None, + name: SupportArrowSchema | str | Any, + type: _DataTypeT | str | None = None, nullable: bool = True, - metadata: dict[Any, Any] | None = None + metadata: dict[Any, Any] | None = None, ) -> Field[_DataTypeT] | Field[Any]: ... - - 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: _Unit | str, tz: _Tz | None = None) -> TimestampType[_Unit, _Tz]: ... - - + unit: _Unit | str, tz: _Tz | None = None +) -> TimestampType[_Unit, _Tz]: ... def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... - - def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... - - def duration(unit: _Unit | str) -> DurationType[_Unit]: ... - - def month_day_nano_interval() -> MonthDayNanoIntervalType: ... - - def date32() -> Date32Type: ... - - def date64() -> Date64Type: ... - - def float16() -> Float16Type: ... - - def float32() -> Float32Type: ... - - def float64() -> Float64Type: ... - - -def decimal32(precision: _Precision, - scale: _Scale | Literal[0] = 0) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... - - -def decimal64(precision: _Precision, - scale: _Scale | Literal[0] = 0) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... - - -def decimal128(precision: _Precision, - scale: _Scale | Literal[0] = 0) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... - - -def decimal256(precision: _Precision, - scale: _Scale | Literal[0] = 0) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... - - +def decimal32( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... +def decimal64( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... +def decimal128( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... +def decimal256( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... def string() -> StringType: ... - utf8 = string - -def binary(length: Literal[-1] | int = ...) -> BinaryType | FixedSizeBinaryType[bytes]: ... - - +def binary( + length: Literal[-1] | int = ..., +) -> BinaryType | FixedSizeBinaryType[bytes]: ... def large_binary() -> LargeBinaryType: ... - - def large_string() -> LargeStringType: ... - large_utf8 = large_string - def binary_view() -> BinaryViewType: ... - - def string_view() -> StringViewType: ... - - def list_( value_type: _DataTypeT | Field[_DataTypeT] | None = None, - list_size: Literal[-1] | _Size | None = None + list_size: Literal[-1] | _Size | None = None, ) -> ListType[_DataTypeT] | FixedSizeListType[_DataTypeT, _Size]: ... - - -def large_list(value_type: _DataTypeT | - Field[_DataTypeT] | None = None) -> LargeListType[_DataTypeT]: ... - - -def list_view(value_type: _DataTypeT | - Field[_DataTypeT] | None = None) -> ListViewType[_DataTypeT]: ... - - +def large_list( + value_type: _DataTypeT | Field[_DataTypeT] | None = None, +) -> LargeListType[_DataTypeT]: ... +def list_view( + value_type: _DataTypeT | Field[_DataTypeT] | None = None, +) -> ListViewType[_DataTypeT]: ... def large_list_view( - value_type: _DataTypeT | Field[_DataTypeT] | None = None + value_type: _DataTypeT | Field[_DataTypeT] | None = None, ) -> LargeListViewType[_DataTypeT]: ... - - def map_( key_type: _K | Field | str | None = None, item_type: _ValueT | Field | str | None = None, - keys_sorted: _Ordered | None = None + keys_sorted: _Ordered | None = None, ) -> MapType[_K, _ValueT, _Ordered]: ... - - def dictionary( index_type: _IndexT | str, value_type: _BasicValueT | str, - ordered: _Ordered | None = None + ordered: _Ordered | None = None, ) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... - - def struct( fields: Iterable[ - Field[Any] - | tuple[str, Field[Any] | None] - | tuple[str, DataType | None] - ] | Mapping[str, Field[Any] | DataType | None], + Field[Any] | tuple[str, Field[Any] | None] | tuple[str, DataType | None] + ] + | Mapping[str, Field[Any] | DataType | None], ) -> StructType: ... - - def sparse_union( child_fields: list[Field[Any]], type_codes: list[int] | None = None ) -> SparseUnionType: ... - - def dense_union( child_fields: list[Field[Any]], type_codes: list[int] | None = None ) -> DenseUnionType: ... - - def union( - child_fields: list[Field[Any]], mode: Literal["sparse", "dense"] | int | str, - type_codes: list[int] | None = None) -> SparseUnionType | DenseUnionType: ... - - + child_fields: list[Field[Any]], + mode: Literal["sparse", "dense"] | int, + type_codes: list[int] | None = None, +) -> SparseUnionType | DenseUnionType: ... def run_end_encoded( run_end_type: _RunEndType | str | None, value_type: _BasicValueT | str | None ) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... - - def json_(storage_type: DataType = ...) -> JsonType: ... - - def uuid() -> UuidType: ... - - def fixed_shape_tensor( value_type: _ValueT, shape: Sequence[int], dim_names: Sequence[str] | None = None, permutation: Sequence[int] | None = None, ) -> FixedShapeTensorType[_ValueT]: ... - - def bool8() -> Bool8Type: ... - - def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... - - def type_for_alias(name: Any) -> DataType: ... - - def schema( fields: ( Iterable[Field[Any]] @@ -833,13 +532,11 @@ def schema( metadata: Mapping[bytes, bytes] | Mapping[str, str] | Mapping[bytes, str] - | Mapping[str, bytes] | None = None, + | Mapping[str, bytes] + | None = None, ) -> Schema: ... - - def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... - __all__ = [ "_Weakrefable", "_Metadata", diff --git a/python/pyarrow-stubs/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/error.pyi index 6a010071de08..8ee75c3ec414 100644 --- a/python/pyarrow-stubs/pyarrow/error.pyi +++ b/python/pyarrow-stubs/pyarrow/error.pyi @@ -17,68 +17,34 @@ 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 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: - ... - +class StopToken: ... def enable_signal_handlers(enable: bool) -> None: ... - have_signal_refcycle: bool - class SignalStopHandler: def __enter__(self) -> Self: ... def __exit__(self, exc_type, exc_value, exc_tb) -> None: ... @property def stop_token(self) -> StopToken: ... - __all__ = [ "ArrowException", "ArrowInvalid", diff --git a/python/pyarrow-stubs/pyarrow/io.pyi b/python/pyarrow-stubs/pyarrow/io.pyi index 856fb093a6cd..8a1da39a1cf2 100644 --- a/python/pyarrow-stubs/pyarrow/io.pyi +++ b/python/pyarrow-stubs/pyarrow/io.pyi @@ -15,16 +15,11 @@ # 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. -""" +# 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: - ... - +class Buffer: ... def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi index 775434be2ea4..0c539af067ed 100644 --- a/python/pyarrow-stubs/pyarrow/lib.pyi +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -15,10 +15,8 @@ # specific language governing permissions and limitations # under the License. -"""Placeholder stub for pyarrow.lib C extension module. - -Complete type annotations will be added in subsequent PRs. -""" +# Placeholder stub for pyarrow.lib C extension module. +# Complete type annotations will be added in subsequent PRs. from typing import Any diff --git a/python/pyarrow-stubs/pyarrow/scalar.pyi b/python/pyarrow-stubs/pyarrow/scalar.pyi index 384507d2a2b6..307b3e689a25 100644 --- a/python/pyarrow-stubs/pyarrow/scalar.pyi +++ b/python/pyarrow-stubs/pyarrow/scalar.pyi @@ -15,16 +15,11 @@ # 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. -""" +# 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: - ... - +class ExtensionScalar: ... def __getattr__(name: str) -> Any: ... From 2f1a0ba26360505ac6be8682a4818141c5e68233 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:39:04 +0100 Subject: [PATCH 05/93] add pyi check to pre-commit --- .pre-commit-config.yaml | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 50293d291f4d..67bebe09398b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -136,11 +136,22 @@ 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" files: >- ^(c_glib|dev|python)/ + types: + - file + types_or: + - python + - pyi exclude: >- ( ?^python/pyarrow/vendored/| From 42ab12d1ef812fe086f88e9e4516b28826e5bea7 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:39:47 +0100 Subject: [PATCH 06/93] ruff format --- python/pyarrow-stubs/pyarrow/_types.pyi | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 1fd5c5002d1e..c4d88ffce6b8 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -125,9 +125,7 @@ class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): _FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) -class FixedSizeBinaryType( - _BasicDataType[_FixedSizeBinaryAsPyType] -): ... +class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... _Precision = TypeVar("_Precision", default=Any) _Scale = TypeVar("_Scale", default=Any) From 675ae6441e0fb44e1d3a376f4fcd138c30207c44 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 01:20:19 +0100 Subject: [PATCH 07/93] Apply suggestions from code review Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com> --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 7 ++- python/pyarrow-stubs/pyarrow/_types.pyi | 49 ++++++++----------- 2 files changed, 24 insertions(+), 32 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 1d5f7600c4b7..6d0c15fd6e3e 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -133,7 +133,6 @@ class SupportArrowDeviceArray(Protocol): class SupportArrowSchema(Protocol): def __arrow_c_schema__(self) -> Any: ... -class NullableCollection(Protocol[_V]): - def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... - def __len__(self) -> int: ... - def __contains__(self, item: Any, /) -> bool: ... +from collections.abc import Container, Sized +class NullableCollection(Sized, Container[_V], Protocol[_V]): + def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index c4d88ffce6b8..30f719a739f9 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -42,7 +42,7 @@ class _Weakrefable: ... class _Metadata(_Weakrefable): ... class DataType(_Weakrefable): - def field(self, i: int) -> Field: ... + def field(self, i: int) -> Field[Any]: ... @property def id(self) -> int: ... @property @@ -72,7 +72,8 @@ class DataType(_Weakrefable): def _import_from_c_capsule(cls, schema) -> Self: ... _AsPyType = TypeVar("_AsPyType") -_DataTypeT = TypeVar("_DataTypeT", bound=DataType) +_DataTypeT = TypeVar("_DataTypeT", bound=DataType) +_DataTypeT_co = TypeVar("_DataTypeT", bound=DataType, covariant=True) class _BasicDataType(DataType, Generic[_AsPyType]): ... class NullType(_BasicDataType[None]): ... @@ -127,32 +128,24 @@ _FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... -_Precision = TypeVar("_Precision", default=Any) -_Scale = TypeVar("_Scale", default=Any) - -class Decimal32Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): - @property - def precision(self) -> _Precision: ... - @property - def scale(self) -> _Scale: ... - -class Decimal64Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): - @property - def precision(self) -> _Precision: ... - @property - def scale(self) -> _Scale: ... - -class Decimal128Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): - @property - def precision(self) -> _Precision: ... - @property - def scale(self) -> _Scale: ... - -class Decimal256Type(FixedSizeBinaryType[Decimal], Generic[_Precision, _Scale]): - @property - def precision(self) -> _Precision: ... - @property - def scale(self) -> _Scale: ... +from typing import Protocol + +_Precision = TypeVar("_Precision", default=Any , covariant=True) +_Scale = TypeVar("_Scale", default=Any , covariant=True) + +class _HasPrecisionScale(Protocol[_Precision, _Scale]): + @property + def precision(self) -> _Precision: ... + @property + def scale(self) -> _Scale: ... + +class Decimal32Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... + +class Decimal64Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... + +class Decimal128Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... + +class Decimal256Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... class ListType(DataType, Generic[_DataTypeT]): @property From 14b924bb892f4336f2e89e86bc593e736587c8ae Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 01:50:36 +0100 Subject: [PATCH 08/93] review feedback --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 56 +++++--- python/pyarrow-stubs/pyarrow/_types.pyi | 132 ++++++++++-------- 2 files changed, 109 insertions(+), 79 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 6d0c15fd6e3e..572bc7ba99f2 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -17,7 +17,7 @@ import datetime as dt -from collections.abc import Collection, Iterator, Sequence +from collections.abc import Collection, Container, Iterator, Sequence, Sized from decimal import Decimal from typing import Any, Literal, Protocol, TypeAlias, TypeVar @@ -90,49 +90,65 @@ _V = TypeVar("_V", covariant=True) SingleOrList: TypeAlias = list[_T] | _T -class SupportEq(Protocol): + +class SupportsEq(Protocol): def __eq__(self, other: object, /) -> bool: ... -class SupportLt(Protocol): + +class SupportsLt(Protocol): def __lt__(self, other: object, /) -> bool: ... -class SupportGt(Protocol): + +class SupportsGt(Protocol): def __gt__(self, other: object, /) -> bool: ... -class SupportLe(Protocol): + +class SupportsLe(Protocol): def __le__(self, other: object, /) -> bool: ... -class SupportGe(Protocol): + +class SupportsGe(Protocol): def __ge__(self, other: object, /) -> bool: ... + FilterTuple: TypeAlias = ( - tuple[str, Literal["=", "==", "!="], SupportEq] - | tuple[str, Literal["<"], SupportLt] - | tuple[str, Literal[">"], SupportGt] - | tuple[str, Literal["<="], SupportLe] - | tuple[str, Literal[">="], SupportGe] + tuple[str, Literal["=", "==", "!="], SupportsEq] + | tuple[str, Literal["<"], SupportsLt] + | tuple[str, Literal[">"], SupportsGt] + | tuple[str, Literal["<="], SupportsLe] + | tuple[str, Literal[">="], SupportsGe] | tuple[str, Literal["in", "not in"], Collection] | tuple[str, str, Any] # Allow general str for operator to avoid type errors ) + class Buffer(Protocol): ... -class SupportPyBuffer(Protocol): ... -class SupportArrowStream(Protocol): + +class SupportsPyBuffer(Protocol): ... + + +class SupportsArrowStream(Protocol): def __arrow_c_stream__(self, requested_schema=None, /) -> Any: ... -class SupportPyArrowArray(Protocol): + +class SupportsPyArrowArray(Protocol): def __arrow_array__(self, type=None, /) -> Any: ... -class SupportArrowArray(Protocol): + +class SupportsArrowArray(Protocol): def __arrow_c_array__(self, requested_schema=None, /) -> Any: ... -class SupportArrowDeviceArray(Protocol): + +class SupportsArrowDeviceArray(Protocol): def __arrow_c_device_array__(self, requested_schema=None, /, **kwargs) -> Any: ... -class SupportArrowSchema(Protocol): + +class SupportsArrowSchema(Protocol): def __arrow_c_schema__(self) -> Any: ... -from collections.abc import Container, Sized -class NullableCollection(Sized, Container[_V], Protocol[_V]): - def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... + +class NullableCollection(Sized, Container[_V], Protocol[_V]): + def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... + 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 index 30f719a739f9..acd558203949 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -17,16 +17,16 @@ import datetime as dt # noqa: F401 -from collections.abc import Mapping, Sequence, Iterable, Iterator +from collections.abc import Iterable, Iterator, Mapping, Sequence from decimal import Decimal # noqa: F401 -from typing import Any, Generic, Literal +from typing import Any, Generic, Literal, Protocol, TypeAlias import numpy as np import pandas as pd from typing_extensions import Self, TypeVar, deprecated -from pyarrow._stubs_typing import SupportArrowSchema, TimeUnit +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit from pyarrow.io import Buffer from pyarrow.lib import ( # noqa: F401 Array, @@ -72,8 +72,8 @@ class DataType(_Weakrefable): def _import_from_c_capsule(cls, schema) -> Self: ... _AsPyType = TypeVar("_AsPyType") -_DataTypeT = TypeVar("_DataTypeT", bound=DataType) -_DataTypeT_co = TypeVar("_DataTypeT", bound=DataType, covariant=True) +_DataTypeT = TypeVar("_DataTypeT", bound=DataType) +_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) class _BasicDataType(DataType, Generic[_AsPyType]): ... class NullType(_BasicDataType[None]): ... @@ -128,54 +128,60 @@ _FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... -from typing import Protocol - -_Precision = TypeVar("_Precision", default=Any , covariant=True) -_Scale = TypeVar("_Scale", default=Any , covariant=True) - -class _HasPrecisionScale(Protocol[_Precision, _Scale]): - @property - def precision(self) -> _Precision: ... - @property - def scale(self) -> _Scale: ... - -class Decimal32Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... - -class Decimal64Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... - -class Decimal128Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... - -class Decimal256Type(FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale]): ... +_Precision = TypeVar("_Precision", default=Any, covariant=True) +_Scale = TypeVar("_Scale", default=Any, covariant=True) -class ListType(DataType, Generic[_DataTypeT]): +class _HasPrecisionScale(Protocol[_Precision, _Scale]): @property - def value_field(self) -> Field[_DataTypeT]: ... + def precision(self) -> _Precision: ... @property - def value_type(self) -> _DataTypeT: ... + def scale(self) -> _Scale: ... -class LargeListType(DataType, Generic[_DataTypeT]): +class Decimal32Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] +): ... + +class Decimal64Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] +): ... + +class Decimal128Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] +): ... + +class Decimal256Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] +): ... + +class ListType(DataType, Generic[_DataTypeT_co]): @property - def value_field(self) -> Field[_DataTypeT]: ... + def value_field(self) -> Field[_DataTypeT_co]: ... @property - def value_type(self) -> _DataTypeT: ... + def value_type(self) -> _DataTypeT_co: ... -class ListViewType(DataType, Generic[_DataTypeT]): +class LargeListType(DataType, Generic[_DataTypeT_co]): @property - def value_field(self) -> Field[_DataTypeT]: ... + def value_field(self) -> Field[_DataTypeT_co]: ... @property - def value_type(self) -> _DataTypeT: ... + def value_type(self) -> _DataTypeT_co: ... -class LargeListViewType(DataType, Generic[_DataTypeT]): +class ListViewType(DataType, Generic[_DataTypeT_co]): @property - def value_field(self) -> Field[_DataTypeT]: ... + def value_field(self) -> Field[_DataTypeT_co]: ... @property - def value_type(self) -> _DataTypeT: ... + def value_type(self) -> _DataTypeT_co: ... -class FixedSizeListType(DataType, Generic[_DataTypeT, _Size]): +class LargeListViewType(DataType, Generic[_DataTypeT_co]): @property - def value_field(self) -> Field[_DataTypeT]: ... + def value_field(self) -> Field[_DataTypeT_co]: ... @property - def value_type(self) -> _DataTypeT: ... + def value_type(self) -> _DataTypeT_co: ... + +class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... @property def list_size(self) -> int: ... @@ -305,13 +311,22 @@ class UnknownExtensionType(ExtensionType): def register_extension_type(ext_type: ExtensionType) -> None: ... def unregister_extension_type(type_name: str) -> None: ... +_StrOrBytes: TypeAlias = str | bytes +_MetadataMapping: TypeAlias = Mapping[_StrOrBytes, _StrOrBytes] +_MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_FieldTypeInput: TypeAlias = DataType | str | None +_SchemaMetadataInput: TypeAlias = ( + Mapping[bytes, bytes] + | Mapping[str, str] + | Mapping[bytes, str] + | Mapping[str, bytes] +) + class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( self, - __arg0__: Mapping[str | bytes, str | bytes] - | Iterable[tuple[str | bytes, str | bytes]] - | KeyValueMetadata - | None = None, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, **kwargs: str, ) -> None: ... def equals(self, other: KeyValueMetadata) -> bool: ... @@ -322,7 +337,7 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def get_all(self, key: str) -> list[bytes]: ... def to_dict(self) -> dict[bytes, bytes]: ... -class Field(_Weakrefable, Generic[_DataTypeT]): +class Field(_Weakrefable, Generic[_DataTypeT_co]): def equals(self, other: Field, check_metadata: bool = False) -> bool: ... def __hash__(self) -> int: ... @property @@ -332,17 +347,15 @@ class Field(_Weakrefable, Generic[_DataTypeT]): @property def metadata(self) -> dict[bytes, bytes] | None: ... @property - def type(self) -> _DataTypeT: ... + def type(self) -> _DataTypeT_co: ... def with_metadata( self, - metadata: dict[bytes | str, bytes | str] - | Mapping[bytes | str, bytes | str] - | Any, + metadata: _MetadataMapping | Any, ) -> Self: ... def remove_metadata(self) -> Self: ... def with_type(self, new_type: DataType) -> Field: ... def with_name(self, name: str) -> Self: ... - def with_nullable(self, nullable: bool) -> Field[_DataTypeT]: ... + def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... def flatten(self) -> list[Field]: ... def _export_to_c(self, out_ptr: int) -> None: ... @classmethod @@ -351,6 +364,14 @@ class Field(_Weakrefable, Generic[_DataTypeT]): @classmethod def _import_from_c_capsule(cls, schema) -> Self: ... +_StructFieldTuple: TypeAlias = ( + tuple[str, Field[Any] | None] | tuple[str, _FieldTypeInput] +) +_StructFieldsInput: TypeAlias = ( + Iterable[Field[Any] | _StructFieldTuple] + | Mapping[str, Field[Any] | DataType | str | None] +) + class Schema(_Weakrefable): def __len__(self) -> int: ... def __getitem__(self, key: str | int) -> Field: ... @@ -407,10 +428,10 @@ def unify_schemas( promote_options: Literal["default", "permissive"] = "default", ) -> Schema: ... def field( - name: SupportArrowSchema | str | Any, + name: SupportsArrowSchema | str | bytes, type: _DataTypeT | str | None = None, nullable: bool = True, - metadata: dict[Any, Any] | None = None, + metadata: _MetadataMapping | None = None, ) -> Field[_DataTypeT] | Field[Any]: ... def null() -> NullType: ... def bool_() -> BoolType: ... @@ -484,10 +505,7 @@ def dictionary( ordered: _Ordered | None = None, ) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... def struct( - fields: Iterable[ - Field[Any] | tuple[str, Field[Any] | None] | tuple[str, DataType | None] - ] - | Mapping[str, Field[Any] | DataType | None], + fields: _StructFieldsInput, ) -> StructType: ... def sparse_union( child_fields: list[Field[Any]], type_codes: list[int] | None = None @@ -520,11 +538,7 @@ def schema( | Iterable[tuple[str, DataType | str | None]] | Mapping[Any, DataType | str | None] ), - metadata: Mapping[bytes, bytes] - | Mapping[str, str] - | Mapping[bytes, str] - | Mapping[str, bytes] - | None = None, + metadata: _SchemaMetadataInput | None = None, ) -> Schema: ... def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... From 4ec84e4fc880c6d59f85b54d5fefb3b3b23dc86b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 13:57:38 +0100 Subject: [PATCH 09/93] Decimals should be covariant --- python/pyarrow-stubs/pyarrow/_types.pyi | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index acd558203949..02f93645a456 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -128,29 +128,31 @@ _FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... -_Precision = TypeVar("_Precision", default=Any, covariant=True) -_Scale = TypeVar("_Scale", default=Any, covariant=True) +_Precision = TypeVar("_Precision", default=Any) +_Scale = TypeVar("_Scale", default=Any) +_Precision_co = TypeVar("_Precision_co", default=Any, covariant=True) +_Scale_co = TypeVar("_Scale_co", default=Any, covariant=True) -class _HasPrecisionScale(Protocol[_Precision, _Scale]): +class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): @property - def precision(self) -> _Precision: ... + def precision(self) -> _Precision_co: ... @property - def scale(self) -> _Scale: ... + def scale(self) -> _Scale_co: ... class Decimal32Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] ): ... class Decimal64Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] ): ... class Decimal128Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] ): ... class Decimal256Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision, _Scale] + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] ): ... class ListType(DataType, Generic[_DataTypeT_co]): From 235c90d52f60ce77f08dfd93c84362a5b1b7cdc9 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 02:10:18 +0100 Subject: [PATCH 10/93] fix editable install collisions --- docs/source/developers/python/building.rst | 1 + python/CMakeLists.txt | 9 --------- 2 files changed, 1 insertion(+), 9 deletions(-) diff --git a/docs/source/developers/python/building.rst b/docs/source/developers/python/building.rst index 73a2482ecbfa..bb033281f9b9 100644 --- a/docs/source/developers/python/building.rst +++ b/docs/source/developers/python/building.rst @@ -204,6 +204,7 @@ Windows tab under the :ref:`pyarrow_build_section` section. $ export ARROW_HOME=$(pwd)/dist $ export LD_LIBRARY_PATH=$(pwd)/dist/lib:$LD_LIBRARY_PATH + $ export DYLD_LIBRARY_PATH=$(pwd)/dist/lib:$DYLD_LIBRARY_PATH $ export CMAKE_PREFIX_PATH=$ARROW_HOME:$CMAKE_PREFIX_PATH .. tab-item:: Windows diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 5bff8bcc6499..9f407e8514b8 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -295,15 +295,6 @@ message(STATUS "Found Cython version: ${CYTHON_VERSION}") include(GNUInstallDirs) find_package(Arrow REQUIRED) -# When not bundling Arrow C++ libraries on macOS, add the Arrow library -# directory to the RPATH so that the extensions can find libarrow at runtime. -if(APPLE - AND NOT PYARROW_BUNDLE_ARROW_CPP - AND ARROW_SHARED_LIB) - get_filename_component(_arrow_lib_dir "${ARROW_SHARED_LIB}" DIRECTORY) - list(APPEND CMAKE_INSTALL_RPATH "${_arrow_lib_dir}") -endif() - macro(define_option name description arrow_option) set("PYARROW_${name}" "AUTO" From 046fd112426a8b20546a0a0048b5fd81bbd760be Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 21:39:36 +0100 Subject: [PATCH 11/93] review feedback --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 20 ++-- python/pyarrow-stubs/pyarrow/_types.pyi | 96 +++++++++++++++---- 2 files changed, 86 insertions(+), 30 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 572bc7ba99f2..cfa5cb6e27ea 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -91,32 +91,32 @@ _V = TypeVar("_V", covariant=True) SingleOrList: TypeAlias = list[_T] | _T -class SupportsEq(Protocol): +class SupportsDunderEQ(Protocol): def __eq__(self, other: object, /) -> bool: ... -class SupportsLt(Protocol): +class SupportsDunderLT(Protocol): def __lt__(self, other: object, /) -> bool: ... -class SupportsGt(Protocol): +class SupportsDunderGT(Protocol): def __gt__(self, other: object, /) -> bool: ... -class SupportsLe(Protocol): +class SupportsDunderLE(Protocol): def __le__(self, other: object, /) -> bool: ... -class SupportsGe(Protocol): +class SupportsDunderGE(Protocol): def __ge__(self, other: object, /) -> bool: ... FilterTuple: TypeAlias = ( - tuple[str, Literal["=", "==", "!="], SupportsEq] - | tuple[str, Literal["<"], SupportsLt] - | tuple[str, Literal[">"], SupportsGt] - | tuple[str, Literal["<="], SupportsLe] - | tuple[str, Literal[">="], SupportsGe] + tuple[str, Literal["=", "==", "!="], SupportsDunderEQ] + | tuple[str, Literal["<"], SupportsDunderLT] + | tuple[str, Literal[">"], SupportsDunderGT] + | tuple[str, Literal["<="], SupportsDunderLE] + | tuple[str, Literal[">="], SupportsDunderGE] | tuple[str, Literal["in", "not in"], Collection] | tuple[str, str, Any] # Allow general str for operator to avoid type errors ) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 02f93645a456..fd4a9ea49448 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -42,7 +42,7 @@ class _Weakrefable: ... class _Metadata(_Weakrefable): ... class DataType(_Weakrefable): - def field(self, i: int) -> Field[Any]: ... + def field(self, i: int) -> Field[Any]: ... @property def id(self) -> int: ... @property @@ -56,9 +56,6 @@ class DataType(_Weakrefable): @property def has_variadic_buffers(self) -> bool: ... - # Properties that exist on specific subtypes but accessed generically - @property - def list_size(self) -> int: ... def __hash__(self) -> int: ... def equals( self, other: DataType | str, *, check_metadata: bool = False @@ -200,8 +197,9 @@ _IndexT = TypeVar( UInt64Type, Int64Type, ) -_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType) -_ValueT = TypeVar("_ValueT", bound=DataType) +_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType, default=_BasicDataType) +_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) +_K = TypeVar("_K", bound=DataType, default=DataType) _Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): @@ -212,8 +210,6 @@ class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): @property def value_type(self) -> _BasicValueT: ... -_K = TypeVar("_K", bound=DataType) - class MapType(DataType, Generic[_K, _ValueT, _Ordered]): @property def key_field(self) -> Field[_K]: ... @@ -317,7 +313,66 @@ _StrOrBytes: TypeAlias = str | bytes _MetadataMapping: TypeAlias = Mapping[_StrOrBytes, _StrOrBytes] _MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] _KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None -_FieldTypeInput: TypeAlias = DataType | str | 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 +_FieldTypeInput: TypeAlias = _DataTypeLike | None _SchemaMetadataInput: TypeAlias = ( Mapping[bytes, bytes] | Mapping[str, str] @@ -371,7 +426,7 @@ _StructFieldTuple: TypeAlias = ( ) _StructFieldsInput: TypeAlias = ( Iterable[Field[Any] | _StructFieldTuple] - | Mapping[str, Field[Any] | DataType | str | None] + | Mapping[str, Field[Any] | _FieldTypeInput] ) class Schema(_Weakrefable): @@ -406,8 +461,8 @@ class Schema(_Weakrefable): def remove(self, i: int) -> Schema: ... def set(self, i: int, field: Field) -> Schema: ... @deprecated("Use 'with_metadata' instead") - def add_metadata(self, metadata: dict) -> Schema: ... - def with_metadata(self, metadata: dict) -> Schema: ... + def add_metadata(self, metadata: _MetadataMapping) -> Schema: ... + def with_metadata(self, metadata: _MetadataMapping) -> Schema: ... def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... def remove_metadata(self) -> Schema: ... def to_string( @@ -431,7 +486,7 @@ def unify_schemas( ) -> Schema: ... def field( name: SupportsArrowSchema | str | bytes, - type: _DataTypeT | str | None = None, + type: _DataTypeT | _DataTypeAliasInput | None = None, nullable: bool = True, metadata: _MetadataMapping | None = None, ) -> Field[_DataTypeT] | Field[Any]: ... @@ -497,13 +552,13 @@ def large_list_view( value_type: _DataTypeT | Field[_DataTypeT] | None = None, ) -> LargeListViewType[_DataTypeT]: ... def map_( - key_type: _K | Field | str | None = None, - item_type: _ValueT | Field | str | None = None, + key_type: _K | Field | _FieldTypeInput = None, + item_type: _ValueT | Field | _FieldTypeInput = None, keys_sorted: _Ordered | None = None, ) -> MapType[_K, _ValueT, _Ordered]: ... def dictionary( - index_type: _IndexT | str, - value_type: _BasicValueT | str, + index_type: _IndexT | _DataTypeAliasInput, + value_type: _BasicValueT | _DataTypeAliasInput, ordered: _Ordered | None = None, ) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... def struct( @@ -521,7 +576,8 @@ def union( type_codes: list[int] | None = None, ) -> SparseUnionType | DenseUnionType: ... def run_end_encoded( - run_end_type: _RunEndType | str | None, value_type: _BasicValueT | str | None + run_end_type: _RunEndType | _DataTypeAliasInput | None, + value_type: _BasicValueT | _DataTypeAliasInput | None, ) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... def json_(storage_type: DataType = ...) -> JsonType: ... def uuid() -> UuidType: ... @@ -537,8 +593,8 @@ def type_for_alias(name: Any) -> DataType: ... def schema( fields: ( Iterable[Field[Any]] - | Iterable[tuple[str, DataType | str | None]] - | Mapping[Any, DataType | str | None] + | Iterable[tuple[str, _FieldTypeInput]] + | Mapping[Any, _FieldTypeInput] ), metadata: _SchemaMetadataInput | None = None, ) -> Schema: ... From 0454fb963d65d1f932e7d8689cc5422a2625ff1b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 22:38:22 +0100 Subject: [PATCH 12/93] move placeholder --- python/pyarrow-stubs/pyarrow/_types.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index fd4a9ea49448..a94f3cce32bd 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -27,16 +27,16 @@ import pandas as pd from typing_extensions import Self, TypeVar, deprecated from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit -from pyarrow.io import Buffer from pyarrow.lib import ( # noqa: F401 Array, + Buffer, ChunkedArray, ExtensionArray, + ExtensionScalar, MemoryPool, MonthDayNano, Table, ) -from pyarrow.scalar import ExtensionScalar class _Weakrefable: ... class _Metadata(_Weakrefable): ... From e6bc69be5aae0767277ede20575dee68b59b06ad Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 6 May 2026 17:52:37 +0200 Subject: [PATCH 13/93] post rebase --- docs/source/developers/python/building.rst | 1 - python/CMakeLists.txt | 9 +++++++++ python/pyarrow-stubs/pyarrow/__init__.pyi | 8 +++++--- 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/source/developers/python/building.rst b/docs/source/developers/python/building.rst index bb033281f9b9..73a2482ecbfa 100644 --- a/docs/source/developers/python/building.rst +++ b/docs/source/developers/python/building.rst @@ -204,7 +204,6 @@ Windows tab under the :ref:`pyarrow_build_section` section. $ export ARROW_HOME=$(pwd)/dist $ export LD_LIBRARY_PATH=$(pwd)/dist/lib:$LD_LIBRARY_PATH - $ export DYLD_LIBRARY_PATH=$(pwd)/dist/lib:$DYLD_LIBRARY_PATH $ export CMAKE_PREFIX_PATH=$ARROW_HOME:$CMAKE_PREFIX_PATH .. tab-item:: Windows diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 9f407e8514b8..5bff8bcc6499 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -295,6 +295,15 @@ message(STATUS "Found Cython version: ${CYTHON_VERSION}") include(GNUInstallDirs) find_package(Arrow REQUIRED) +# When not bundling Arrow C++ libraries on macOS, add the Arrow library +# directory to the RPATH so that the extensions can find libarrow at runtime. +if(APPLE + AND NOT PYARROW_BUNDLE_ARROW_CPP + AND ARROW_SHARED_LIB) + get_filename_component(_arrow_lib_dir "${ARROW_SHARED_LIB}" DIRECTORY) + list(APPEND CMAKE_INSTALL_RPATH "${_arrow_lib_dir}") +endif() + macro(define_option name description arrow_option) set("PYARROW_${name}" "AUTO" diff --git a/python/pyarrow-stubs/pyarrow/__init__.pyi b/python/pyarrow-stubs/pyarrow/__init__.pyi index 9c7120c7f807..ccec8d5abc07 100644 --- a/python/pyarrow-stubs/pyarrow/__init__.pyi +++ b/python/pyarrow-stubs/pyarrow/__init__.pyi @@ -15,9 +15,11 @@ # specific language governing permissions and limitations # under the License. -# Type stubs for PyArrow. -# This is a placeholder stub file. -# Complete type annotations will be added in subsequent PRs. +"""Type stubs for PyArrow. + +This is a placeholder stub file. +Complete type annotations will be added in subsequent PRs. +""" from typing import Any From a33f32793da3663c129ed091fed31eb0d8d6b895 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 16:16:48 +0200 Subject: [PATCH 14/93] Apply suggestions from code review Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com> --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 73 ++++++------- python/pyarrow-stubs/pyarrow/_types.pyi | 100 +++++++++--------- 2 files changed, 89 insertions(+), 84 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index cfa5cb6e27ea..b7f22480e608 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -48,42 +48,43 @@ NullEncoding: TypeAlias = Literal["mask", "encode"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] -IntegerType: TypeAlias = ( - lib.Int8Type - | lib.Int16Type - | lib.Int32Type - | lib.Int64Type - | lib.UInt8Type - | lib.UInt16Type - | lib.UInt32Type - | lib.UInt64Type -) - -Mask: TypeAlias = ( - Sequence[bool | None] - | NDArray[np.bool_] - | lib.Array[lib.Scalar[lib.BoolType]] - | ChunkedArray[Any] -) -Indices: TypeAlias = ( - Sequence[int | None] - | NDArray[np.integer[Any]] - | lib.Array[lib.Scalar[IntegerType]] - | ChunkedArray[Any] -) - -PyScalar: TypeAlias = ( - bool - | int - | float - | Decimal - | str - | bytes - | dt.date - | dt.datetime - | dt.time - | dt.timedelta -) +IntegerType: TypeAlias = ( + lib.Int8Type + | lib.Int16Type + | lib.Int32Type + | lib.Int64Type + | lib.UInt8Type + | lib.UInt16Type + | lib.UInt32Type + | lib.UInt64Type +) +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) +NumpyScalar: TypeAlias = "np.generic[Any]" + +PyScalarT_co = TypeVar("PyScalarT_co", bound=PyScalar, covariant=True) +NumpyScalarT_co = TypeVar("NumpyScalarT_co", bound=NumpyScalar, covariant=True) +DataTypeT_co = TypeVar("DataTypeT_co", bound=lib.DataType, covariant=True) + +IntoArray: TypeAlias = ( + Sequence[PyScalarT_co | None] + | NDArray[NumpyScalarT_co] + | lib.Array[lib.Scalar[DataTypeT_co]] + | ChunkedArray[Any] +) + +Mask: TypeAlias = IntoArray[bool, np.bool_, lib.BoolType] +Indices: TypeAlias = IntoArray[int, np.integer[Any], IntegerType] _T = TypeVar("_T") _V = TypeVar("_V", covariant=True) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index a94f3cce32bd..69ff034d99b9 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -197,7 +197,9 @@ _IndexT = TypeVar( UInt64Type, Int64Type, ) -_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType, default=_BasicDataType) +_BasicValueT = TypeVar( + "_BasicValueT", bound=_BasicDataType[Any], default=_BasicDataType[Any] +) _ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) _K = TypeVar("_K", bound=DataType, default=DataType) _Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) @@ -226,16 +228,16 @@ _Size = TypeVar("_Size", default=int) class StructType(DataType): def get_field_index(self, name: str) -> int: ... - def field(self, i: int | str) -> Field: ... + def field(self, i: int | str) -> Field[Any]: ... def get_all_field_indices(self, name: str) -> list[int]: ... def __len__(self) -> int: ... - def __iter__(self) -> Iterator[Field]: ... + def __iter__(self) -> Iterator[Field[Any]]: ... __getitem__ = field @property def names(self) -> list[str]: ... @property - def fields(self) -> list[Field]: ... + def fields(self) -> list[Field[Any]]: ... class UnionType(DataType): @property @@ -243,8 +245,8 @@ class UnionType(DataType): @property def type_codes(self) -> list[int]: ... def __len__(self) -> int: ... - def __iter__(self) -> Iterator[Field]: ... - def field(self, i: int) -> Field: ... + def __iter__(self) -> Iterator[Field[Any]]: ... + def field(self, i: int) -> Field[Any]: ... __getitem__ = field @@ -264,10 +266,10 @@ class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): @property def value_type(self) -> _BasicValueT: ... -_StorageT = TypeVar("_StorageT", bound=Array | ChunkedArray) +_StorageT = TypeVar("_StorageT", bound=Array[Any] | ChunkedArray[Any]) class BaseExtensionType(DataType): - def __arrow_ext_class__(self) -> type[ExtensionArray]: ... + def __arrow_ext_class__(self) -> type[ExtensionArray[Any]]: ... def __arrow_ext_scalar_class__(self) -> type[ExtensionScalar]: ... @property def extension_name(self) -> str: ... @@ -310,7 +312,9 @@ def register_extension_type(ext_type: ExtensionType) -> None: ... def unregister_extension_type(type_name: str) -> None: ... _StrOrBytes: TypeAlias = str | bytes -_MetadataMapping: TypeAlias = Mapping[_StrOrBytes, _StrOrBytes] +_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[ @@ -373,12 +377,6 @@ _DataTypeAlias: TypeAlias = Literal[ _DataTypeAliasInput: TypeAlias = _DataTypeAlias | str _DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput _FieldTypeInput: TypeAlias = _DataTypeLike | None -_SchemaMetadataInput: TypeAlias = ( - Mapping[bytes, bytes] - | Mapping[str, str] - | Mapping[bytes, str] - | Mapping[str, bytes] -) class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( @@ -395,7 +393,7 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def to_dict(self) -> dict[bytes, bytes]: ... class Field(_Weakrefable, Generic[_DataTypeT_co]): - def equals(self, other: Field, check_metadata: bool = False) -> bool: ... + def equals(self, other: Field[Any], check_metadata: bool = False) -> bool: ... def __hash__(self) -> int: ... @property def nullable(self) -> bool: ... @@ -410,16 +408,16 @@ class Field(_Weakrefable, Generic[_DataTypeT_co]): metadata: _MetadataMapping | Any, ) -> Self: ... def remove_metadata(self) -> Self: ... - def with_type(self, new_type: DataType) -> Field: ... + def with_type(self, new_type: _DataTypeT) -> Field[_DataTypeT]: ... def with_name(self, name: str) -> Self: ... def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... - def flatten(self) -> list[Field]: ... + def flatten(self) -> list[Field[Any]]: ... 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) -> Self: ... + def _import_from_c_capsule(cls, schema: Any) -> Self: ... _StructFieldTuple: TypeAlias = ( tuple[str, Field[Any] | None] | tuple[str, _FieldTypeInput] @@ -431,14 +429,14 @@ _StructFieldsInput: TypeAlias = ( class Schema(_Weakrefable): def __len__(self) -> int: ... - def __getitem__(self, key: str | int) -> Field: ... + def __getitem__(self, key: str | int) -> Field[Any]: ... _field = __getitem__ - def __iter__(self) -> Iterator[Field]: ... + def __iter__(self) -> Iterator[Field[Any]]: ... def __hash__(self) -> int: ... def __sizeof__(self) -> int: ... @property - def pandas_metadata(self) -> dict: ... + def pandas_metadata(self) -> dict[bytes, bytes]: ... @property def names(self) -> list[str]: ... @property @@ -451,15 +449,15 @@ class Schema(_Weakrefable): 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) -> Field: ... + def field(self, i: int | str | bytes) -> Field[Any]: ... + @deprecated("Use 'field' instead") + def field_by_name(self, name: str) -> Field[Any]: ... def get_field_index(self, name: str) -> int: ... def get_all_field_indices(self, name: str) -> list[int]: ... - def append(self, field: Field) -> Schema: ... - def insert(self, i: int, field: Field) -> Schema: ... - def remove(self, i: int) -> Schema: ... - def set(self, i: int, field: Field) -> Schema: ... + def append(self, field: Field[Any]) -> Schema: ... + def insert(self, i: int, field: Field[Any]) -> Schema: ... + def remove(self, i: int) -> Schema: ... + def set(self, i: int, field: Field[Any]) -> Schema: ... @deprecated("Use 'with_metadata' instead") def add_metadata(self, metadata: _MetadataMapping) -> Schema: ... def with_metadata(self, metadata: _MetadataMapping) -> Schema: ... @@ -538,24 +536,30 @@ large_utf8 = large_string def binary_view() -> BinaryViewType: ... def string_view() -> StringViewType: ... -def list_( - value_type: _DataTypeT | Field[_DataTypeT] | None = None, - list_size: Literal[-1] | _Size | None = None, -) -> ListType[_DataTypeT] | FixedSizeListType[_DataTypeT, _Size]: ... -def large_list( - value_type: _DataTypeT | Field[_DataTypeT] | None = None, -) -> LargeListType[_DataTypeT]: ... -def list_view( - value_type: _DataTypeT | Field[_DataTypeT] | None = None, -) -> ListViewType[_DataTypeT]: ... -def large_list_view( - value_type: _DataTypeT | Field[_DataTypeT] | None = None, -) -> LargeListViewType[_DataTypeT]: ... -def map_( - key_type: _K | Field | _FieldTypeInput = None, - item_type: _ValueT | Field | _FieldTypeInput = None, - keys_sorted: _Ordered | None = None, -) -> MapType[_K, _ValueT, _Ordered]: ... +from typing import overload + +@overload +def list_( + value_type: _DataTypeT | Field[_DataTypeT], list_size: Literal[-1] | None = None +) -> ListType[_DataTypeT]: ... +@overload +def list_( + value_type: _DataTypeT | Field[_DataTypeT], list_size: _Size +) -> FixedSizeListType[_DataTypeT, _Size]: ... +def large_list( + value_type: _DataTypeT | Field[_DataTypeT], +) -> LargeListType[_DataTypeT]: ... +def list_view( + value_type: _DataTypeT | Field[_DataTypeT], +) -> ListViewType[_DataTypeT]: ... +def large_list_view( + value_type: _DataTypeT | Field[_DataTypeT], +) -> LargeListViewType[_DataTypeT]: ... +def map_( + key_type: _K | Field[_K] | _FieldTypeInput, + item_type: _ValueT | Field[_ValueT] | _FieldTypeInput, + keys_sorted: _Ordered | None = None, +) -> MapType[_K, _ValueT, _Ordered]: ... def dictionary( index_type: _IndexT | _DataTypeAliasInput, value_type: _BasicValueT | _DataTypeAliasInput, @@ -596,7 +600,7 @@ def schema( | Iterable[tuple[str, _FieldTypeInput]] | Mapping[Any, _FieldTypeInput] ), - metadata: _SchemaMetadataInput | None = None, + metadata: _MetadataMapping | None = None, ) -> Schema: ... def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... From ca9441cc645fa36db58f54de7635ebe4a4f41bb9 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 16:42:49 +0200 Subject: [PATCH 15/93] Update python/pyarrow-stubs/pyarrow/_types.pyi Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com> --- python/pyarrow-stubs/pyarrow/_types.pyi | 1350 +++++++++++------------ 1 file changed, 675 insertions(+), 675 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 69ff034d99b9..efc3bd847ca5 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -1,541 +1,541 @@ -# 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 # noqa: F401 - -from collections.abc import Iterable, Iterator, Mapping, Sequence -from decimal import Decimal # noqa: F401 -from typing import Any, Generic, Literal, Protocol, TypeAlias - -import numpy as np -import pandas as pd - -from typing_extensions import Self, TypeVar, deprecated - -from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit -from pyarrow.lib import ( # noqa: F401 - Array, - Buffer, - ChunkedArray, - ExtensionArray, - ExtensionScalar, - MemoryPool, - MonthDayNano, - Table, -) - -class _Weakrefable: ... -class _Metadata(_Weakrefable): ... - -class DataType(_Weakrefable): - def field(self, i: int) -> Field[Any]: ... - @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) -> np.generic: ... - 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) -> Self: ... - -_AsPyType = TypeVar("_AsPyType") -_DataTypeT = TypeVar("_DataTypeT", bound=DataType) -_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) - -class _BasicDataType(DataType, Generic[_AsPyType]): ... -class NullType(_BasicDataType[None]): ... -class BoolType(_BasicDataType[bool]): ... -class UInt8Type(_BasicDataType[int]): ... -class Int8Type(_BasicDataType[int]): ... -class UInt16Type(_BasicDataType[int]): ... -class Int16Type(_BasicDataType[int]): ... -class UInt32Type(_BasicDataType[int]): ... -class Int32Type(_BasicDataType[int]): ... -class UInt64Type(_BasicDataType[int]): ... -class Int64Type(_BasicDataType[int]): ... -class Float16Type(_BasicDataType[float]): ... -class Float32Type(_BasicDataType[float]): ... -class Float64Type(_BasicDataType[float]): ... -class Date32Type(_BasicDataType[dt.date]): ... -class Date64Type(_BasicDataType[dt.date]): ... -class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... -class StringType(_BasicDataType[str]): ... -class LargeStringType(_BasicDataType[str]): ... -class StringViewType(_BasicDataType[str]): ... -class BinaryType(_BasicDataType[bytes]): ... -class LargeBinaryType(_BasicDataType[bytes]): ... -class BinaryViewType(_BasicDataType[bytes]): ... - -_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) -_Tz = TypeVar("_Tz", str, None, default=None) - -class TimestampType(_BasicDataType[int], Generic[_Unit, _Tz]): - @property - def unit(self) -> _Unit: ... - @property - def tz(self) -> _Tz: ... - -_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) - -class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): - @property - def unit(self) -> _Time32Unit: ... - -_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) - -class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): - @property - def unit(self) -> _Time64Unit: ... - -class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): - @property - def unit(self) -> _Unit: ... - -_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) - -class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... - -_Precision = TypeVar("_Precision", default=Any) -_Scale = TypeVar("_Scale", default=Any) -_Precision_co = TypeVar("_Precision_co", default=Any, covariant=True) -_Scale_co = TypeVar("_Scale_co", default=Any, covariant=True) - -class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): - @property - def precision(self) -> _Precision_co: ... - @property - def scale(self) -> _Scale_co: ... - -class Decimal32Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal64Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal128Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal256Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class ListType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class LargeListType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class ListViewType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class LargeListViewType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - @property - def list_size(self) -> int: ... - -class DictionaryMemo(_Weakrefable): ... - -_IndexT = TypeVar( - "_IndexT", - UInt8Type, - Int8Type, - UInt16Type, - Int16Type, - UInt32Type, - Int32Type, - UInt64Type, - Int64Type, -) +# 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 # noqa: F401 + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from decimal import Decimal # noqa: F401 +from typing import Any, Generic, Literal, Protocol, TypeAlias + +import numpy as np +import pandas as pd + +from typing_extensions import Self, TypeVar, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( # noqa: F401 + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + MonthDayNano, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field[Any]: ... + @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) -> np.generic: ... + 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: ... + +_AsPyType = TypeVar("_AsPyType") +_DataTypeT = TypeVar("_DataTypeT", bound=DataType) +_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) + +class _BasicDataType(DataType, Generic[_AsPyType]): ... +class NullType(_BasicDataType[None]): ... +class BoolType(_BasicDataType[bool]): ... +class UInt8Type(_BasicDataType[int]): ... +class Int8Type(_BasicDataType[int]): ... +class UInt16Type(_BasicDataType[int]): ... +class Int16Type(_BasicDataType[int]): ... +class UInt32Type(_BasicDataType[int]): ... +class Int32Type(_BasicDataType[int]): ... +class UInt64Type(_BasicDataType[int]): ... +class Int64Type(_BasicDataType[int]): ... +class Float16Type(_BasicDataType[float]): ... +class Float32Type(_BasicDataType[float]): ... +class Float64Type(_BasicDataType[float]): ... +class Date32Type(_BasicDataType[dt.date]): ... +class Date64Type(_BasicDataType[dt.date]): ... +class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... +class StringType(_BasicDataType[str]): ... +class LargeStringType(_BasicDataType[str]): ... +class StringViewType(_BasicDataType[str]): ... +class BinaryType(_BasicDataType[bytes]): ... +class LargeBinaryType(_BasicDataType[bytes]): ... +class BinaryViewType(_BasicDataType[bytes]): ... + +_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) +_Tz = TypeVar("_Tz", str, None, default=None) + +class TimestampType(_BasicDataType[int], Generic[_Unit, _Tz]): + @property + def unit(self) -> _Unit: ... + @property + def tz(self) -> _Tz: ... + +_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) + +class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): + @property + def unit(self) -> _Time32Unit: ... + +_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) + +class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): + @property + def unit(self) -> _Unit: ... + +_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) + +class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... + +_Precision = TypeVar("_Precision", default=Any) +_Scale = TypeVar("_Scale", default=Any) +_Precision_co = TypeVar("_Precision_co", default=Any, covariant=True) +_Scale_co = TypeVar("_Scale_co", default=Any, covariant=True) + +class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): + @property + def precision(self) -> _Precision_co: ... + @property + def scale(self) -> _Scale_co: ... + +class Decimal32Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... + +class Decimal64Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... + +class Decimal128Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... + +class Decimal256Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... + +class ListType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class LargeListType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class ListViewType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class LargeListViewType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + @property + def list_size(self) -> int: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexT = TypeVar( + "_IndexT", + UInt8Type, + Int8Type, + UInt16Type, + Int16Type, + UInt32Type, + Int32Type, + UInt64Type, + Int64Type, +) _BasicValueT = TypeVar( "_BasicValueT", bound=_BasicDataType[Any], default=_BasicDataType[Any] ) -_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) -_K = TypeVar("_K", bound=DataType, default=DataType) -_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) - -class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): - @property - def ordered(self) -> _Ordered: ... - @property - def index_type(self) -> _IndexT: ... - @property - def value_type(self) -> _BasicValueT: ... - -class MapType(DataType, Generic[_K, _ValueT, _Ordered]): - @property - def key_field(self) -> Field[_K]: ... - @property - def key_type(self) -> _K: ... - @property - def item_field(self) -> Field[_ValueT]: ... - @property - def item_type(self) -> _ValueT: ... - @property - def keys_sorted(self) -> _Ordered: ... - -_Size = TypeVar("_Size", default=int) - -class StructType(DataType): - def get_field_index(self, name: str) -> int: ... +_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) +_K = TypeVar("_K", bound=DataType, default=DataType) +_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) + +class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): + @property + def ordered(self) -> _Ordered: ... + @property + def index_type(self) -> _IndexT: ... + @property + def value_type(self) -> _BasicValueT: ... + +class MapType(DataType, Generic[_K, _ValueT, _Ordered]): + @property + def key_field(self) -> Field[_K]: ... + @property + def key_type(self) -> _K: ... + @property + def item_field(self) -> Field[_ValueT]: ... + @property + def item_type(self) -> _ValueT: ... + @property + def keys_sorted(self) -> _Ordered: ... + +_Size = TypeVar("_Size", default=int) + +class StructType(DataType): + def get_field_index(self, name: str) -> int: ... def field(self, i: int | str) -> Field[Any]: ... - def get_all_field_indices(self, name: str) -> list[int]: ... - def __len__(self) -> int: ... + def get_all_field_indices(self, name: str) -> list[int]: ... + def __len__(self) -> int: ... def __iter__(self) -> Iterator[Field[Any]]: ... - - __getitem__ = field - @property - def names(self) -> list[str]: ... - @property + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property def fields(self) -> list[Field[Any]]: ... - -class UnionType(DataType): - @property - def mode(self) -> Literal["sparse", "dense"]: ... - @property - def type_codes(self) -> list[int]: ... - def __len__(self) -> int: ... + +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[Any]]: ... def field(self, i: int) -> Field[Any]: ... - - __getitem__ = field - -class SparseUnionType(UnionType): - @property - def mode(self) -> Literal["sparse"]: ... - -class DenseUnionType(UnionType): - @property - def mode(self) -> Literal["dense"]: ... - -_RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) - -class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): - @property - def run_end_type(self) -> _RunEndType: ... - @property - def value_type(self) -> _BasicValueT: ... - + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) + +class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> _BasicValueT: ... + _StorageT = TypeVar("_StorageT", bound=Array[Any] | ChunkedArray[Any]) - -class BaseExtensionType(DataType): + +class BaseExtensionType(DataType): def __arrow_ext_class__(self) -> type[ExtensionArray[Any]]: ... - 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: _StorageT) -> _StorageT: ... - -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, Generic[_ValueT]): - @property - def value_type(self) -> _ValueT: ... - @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: ExtensionType) -> None: ... -def unregister_extension_type(type_name: str) -> None: ... - -_StrOrBytes: TypeAlias = str | bytes + 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: _StorageT) -> _StorageT: ... + +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, Generic[_ValueT]): + @property + def value_type(self) -> _ValueT: ... + @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: ExtensionType) -> None: ... +def unregister_extension_type(type_name: str) -> None: ... + +_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 -_FieldTypeInput: TypeAlias = _DataTypeLike | None - -class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): - def __init__( - self, - __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, - **kwargs: str, - ) -> None: ... - def equals(self, other: KeyValueMetadata) -> bool: ... - def __len__(self) -> int: ... - def __contains__(self, /, __key: object) -> bool: ... - def __getitem__(self, /, __key: Any) -> Any: ... - def __iter__(self) -> Iterator[bytes]: ... - def get_all(self, key: str) -> list[bytes]: ... - def to_dict(self) -> dict[bytes, bytes]: ... - -class Field(_Weakrefable, Generic[_DataTypeT_co]): +_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 +_FieldTypeInput: TypeAlias = _DataTypeLike | None + +class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): + def __init__( + self, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, + **kwargs: str, + ) -> None: ... + def equals(self, other: KeyValueMetadata) -> bool: ... + def __len__(self) -> int: ... + def __contains__(self, /, __key: object) -> bool: ... + def __getitem__(self, /, __key: Any) -> Any: ... + def __iter__(self) -> Iterator[bytes]: ... + def get_all(self, key: str) -> list[bytes]: ... + def to_dict(self) -> dict[bytes, bytes]: ... + +class Field(_Weakrefable, Generic[_DataTypeT_co]): def equals(self, other: Field[Any], 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) -> _DataTypeT_co: ... - def with_metadata( - self, - metadata: _MetadataMapping | Any, - ) -> Self: ... - def remove_metadata(self) -> Self: ... + 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) -> _DataTypeT_co: ... + def with_metadata( + self, + metadata: _MetadataMapping | Any, + ) -> Self: ... + def remove_metadata(self) -> Self: ... def with_type(self, new_type: _DataTypeT) -> Field[_DataTypeT]: ... - def with_name(self, name: str) -> Self: ... - def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... + def with_name(self, name: str) -> Self: ... + def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... def flatten(self) -> list[Field[Any]]: ... - 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 _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: ... - -_StructFieldTuple: TypeAlias = ( - tuple[str, Field[Any] | None] | tuple[str, _FieldTypeInput] -) -_StructFieldsInput: TypeAlias = ( - Iterable[Field[Any] | _StructFieldTuple] - | Mapping[str, Field[Any] | _FieldTypeInput] -) - -class Schema(_Weakrefable): - def __len__(self) -> int: ... + +_StructFieldTuple: TypeAlias = ( + tuple[str, Field[Any] | None] | tuple[str, _FieldTypeInput] +) +_StructFieldsInput: TypeAlias = ( + Iterable[Field[Any] | _StructFieldTuple] + | Mapping[str, Field[Any] | _FieldTypeInput] +) + +class Schema(_Weakrefable): + def __len__(self) -> int: ... def __getitem__(self, key: str | int) -> Field[Any]: ... - - _field = __getitem__ + + _field = __getitem__ def __iter__(self) -> Iterator[Field[Any]]: ... - def __hash__(self) -> int: ... - def __sizeof__(self) -> int: ... - @property + def __hash__(self) -> int: ... + def __sizeof__(self) -> int: ... + @property def pandas_metadata(self) -> dict[bytes, bytes]: ... - @property - def names(self) -> list[str]: ... - @property - def types(self) -> list[DataType]: ... - @property - def metadata(self) -> dict[bytes, bytes]: ... - 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: ... + @property + def names(self) -> list[str]: ... + @property + def types(self) -> list[DataType]: ... + @property + def metadata(self) -> dict[bytes, bytes]: ... + 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[Any]: ... @deprecated("Use 'field' instead") def field_by_name(self, name: str) -> Field[Any]: ... - def get_field_index(self, name: str) -> int: ... - def get_all_field_indices(self, name: str) -> list[int]: ... + def get_field_index(self, name: str) -> int: ... + def get_all_field_indices(self, name: str) -> list[int]: ... def append(self, field: Field[Any]) -> Schema: ... def insert(self, i: int, field: Field[Any]) -> Schema: ... def remove(self, i: int) -> Schema: ... def set(self, i: int, field: Field[Any]) -> Schema: ... - @deprecated("Use 'with_metadata' instead") - def add_metadata(self, metadata: _MetadataMapping) -> Schema: ... - def with_metadata(self, metadata: _MetadataMapping) -> 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 | None = None, - ) -> 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: Sequence[Schema], - *, - promote_options: Literal["default", "permissive"] = "default", -) -> Schema: ... -def field( - name: SupportsArrowSchema | str | bytes, - type: _DataTypeT | _DataTypeAliasInput | None = None, - nullable: bool = True, - metadata: _MetadataMapping | None = None, -) -> Field[_DataTypeT] | Field[Any]: ... -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: _Unit | str, tz: _Tz | None = None -) -> TimestampType[_Unit, _Tz]: ... -def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... -def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... -def duration(unit: _Unit | str) -> DurationType[_Unit]: ... -def month_day_nano_interval() -> MonthDayNanoIntervalType: ... -def date32() -> Date32Type: ... -def date64() -> Date64Type: ... -def float16() -> Float16Type: ... -def float32() -> Float32Type: ... -def float64() -> Float64Type: ... -def decimal32( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... -def decimal64( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... -def decimal128( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... -def decimal256( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... -def string() -> StringType: ... - -utf8 = string - -def binary( - length: Literal[-1] | int = ..., -) -> BinaryType | FixedSizeBinaryType[bytes]: ... -def large_binary() -> LargeBinaryType: ... -def large_string() -> LargeStringType: ... - -large_utf8 = large_string - -def binary_view() -> BinaryViewType: ... -def string_view() -> StringViewType: ... + @deprecated("Use 'with_metadata' instead") + def add_metadata(self, metadata: _MetadataMapping) -> Schema: ... + def with_metadata(self, metadata: _MetadataMapping) -> 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 | None = None, + ) -> 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: Sequence[Schema], + *, + promote_options: Literal["default", "permissive"] = "default", +) -> Schema: ... +def field( + name: SupportsArrowSchema | str | bytes, + type: _DataTypeT | _DataTypeAliasInput | None = None, + nullable: bool = True, + metadata: _MetadataMapping | None = None, +) -> Field[_DataTypeT] | Field[Any]: ... +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: _Unit | str, tz: _Tz | None = None +) -> TimestampType[_Unit, _Tz]: ... +def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... +def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... +def duration(unit: _Unit | str) -> DurationType[_Unit]: ... +def month_day_nano_interval() -> MonthDayNanoIntervalType: ... +def date32() -> Date32Type: ... +def date64() -> Date64Type: ... +def float16() -> Float16Type: ... +def float32() -> Float32Type: ... +def float64() -> Float64Type: ... +def decimal32( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... +def decimal64( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... +def decimal128( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... +def decimal256( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... +def string() -> StringType: ... + +utf8 = string + +def binary( + length: Literal[-1] | int = ..., +) -> BinaryType | FixedSizeBinaryType[bytes]: ... +def large_binary() -> LargeBinaryType: ... +def large_string() -> LargeStringType: ... + +large_utf8 = large_string + +def binary_view() -> BinaryViewType: ... +def string_view() -> StringViewType: ... from typing import overload @overload @@ -560,168 +560,168 @@ def map_( item_type: _ValueT | Field[_ValueT] | _FieldTypeInput, keys_sorted: _Ordered | None = None, ) -> MapType[_K, _ValueT, _Ordered]: ... -def dictionary( - index_type: _IndexT | _DataTypeAliasInput, - value_type: _BasicValueT | _DataTypeAliasInput, - ordered: _Ordered | None = None, -) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... -def struct( - fields: _StructFieldsInput, -) -> StructType: ... -def sparse_union( - child_fields: list[Field[Any]], type_codes: list[int] | None = None -) -> SparseUnionType: ... -def dense_union( - child_fields: list[Field[Any]], type_codes: list[int] | None = None -) -> DenseUnionType: ... -def union( - child_fields: list[Field[Any]], - mode: Literal["sparse", "dense"] | int, - type_codes: list[int] | None = None, -) -> SparseUnionType | DenseUnionType: ... -def run_end_encoded( - run_end_type: _RunEndType | _DataTypeAliasInput | None, - value_type: _BasicValueT | _DataTypeAliasInput | None, -) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... -def json_(storage_type: DataType = ...) -> JsonType: ... -def uuid() -> UuidType: ... -def fixed_shape_tensor( - value_type: _ValueT, - shape: Sequence[int], - dim_names: Sequence[str] | None = None, - permutation: Sequence[int] | None = None, -) -> FixedShapeTensorType[_ValueT]: ... -def bool8() -> Bool8Type: ... -def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... -def type_for_alias(name: Any) -> DataType: ... -def schema( - fields: ( - Iterable[Field[Any]] - | Iterable[tuple[str, _FieldTypeInput]] - | Mapping[Any, _FieldTypeInput] - ), +def dictionary( + index_type: _IndexT | _DataTypeAliasInput, + value_type: _BasicValueT | _DataTypeAliasInput, + ordered: _Ordered | None = None, +) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... +def struct( + fields: _StructFieldsInput, +) -> StructType: ... +def sparse_union( + child_fields: list[Field[Any]], type_codes: list[int] | None = None +) -> SparseUnionType: ... +def dense_union( + child_fields: list[Field[Any]], type_codes: list[int] | None = None +) -> DenseUnionType: ... +def union( + child_fields: list[Field[Any]], + mode: Literal["sparse", "dense"] | int, + type_codes: list[int] | None = None, +) -> SparseUnionType | DenseUnionType: ... +def run_end_encoded( + run_end_type: _RunEndType | _DataTypeAliasInput | None, + value_type: _BasicValueT | _DataTypeAliasInput | None, +) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... +def json_(storage_type: DataType = ...) -> JsonType: ... +def uuid() -> UuidType: ... +def fixed_shape_tensor( + value_type: _ValueT, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + permutation: Sequence[int] | None = None, +) -> FixedShapeTensorType[_ValueT]: ... +def bool8() -> Bool8Type: ... +def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... +def type_for_alias(name: Any) -> DataType: ... +def schema( + fields: ( + Iterable[Field[Any]] + | Iterable[tuple[str, _FieldTypeInput]] + | Mapping[Any, _FieldTypeInput] + ), metadata: _MetadataMapping | None = None, -) -> Schema: ... -def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... - -__all__ = [ - "_Weakrefable", - "_Metadata", - "DataType", - "_BasicDataType", - "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", - "_Unit", - "_Tz", - "_Time32Unit", - "_Time64Unit", - "_DataTypeT", -] +) -> Schema: ... +def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... + +__all__ = [ + "_Weakrefable", + "_Metadata", + "DataType", + "_BasicDataType", + "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", + "_Unit", + "_Tz", + "_Time32Unit", + "_Time64Unit", + "_DataTypeT", +] From dbc4e400e17d8c7c59620714d2bb8e34500f3554 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:00:13 -0700 Subject: [PATCH 16/93] update --- python/pyarrow-stubs/pyarrow/__init__.pyi | 6 - .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 131 +- python/pyarrow-stubs/pyarrow/_types.pyi | 1476 +++++++++-------- python/pyarrow-stubs/pyarrow/error.pyi | 9 +- 4 files changed, 814 insertions(+), 808 deletions(-) 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/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index b7f22480e608..7802f0330597 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -17,9 +17,16 @@ 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, TypeVar +from typing_extensions import TypeAliasType import numpy as np @@ -48,108 +55,84 @@ NullEncoding: TypeAlias = Literal["mask", "encode"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] -IntegerType: TypeAlias = ( - lib.Int8Type - | lib.Int16Type - | lib.Int32Type - | lib.Int64Type - | lib.UInt8Type - | lib.UInt16Type - | lib.UInt32Type - | lib.UInt64Type -) -PyScalar: TypeAlias = ( - bool - | int - | float - | Decimal - | str - | bytes - | dt.date - | dt.datetime - | dt.time - | dt.timedelta -) -NumpyScalar: TypeAlias = "np.generic[Any]" - -PyScalarT_co = TypeVar("PyScalarT_co", bound=PyScalar, covariant=True) -NumpyScalarT_co = TypeVar("NumpyScalarT_co", bound=NumpyScalar, covariant=True) -DataTypeT_co = TypeVar("DataTypeT_co", bound=lib.DataType, covariant=True) - -IntoArray: TypeAlias = ( - Sequence[PyScalarT_co | None] - | NDArray[NumpyScalarT_co] - | lib.Array[lib.Scalar[DataTypeT_co]] - | ChunkedArray[Any] -) - -Mask: TypeAlias = IntoArray[bool, np.bool_, lib.BoolType] -Indices: TypeAlias = IntoArray[int, np.integer[Any], IntegerType] +IntegerType: TypeAlias = ( + lib.Int8Type + | lib.Int16Type + | lib.Int32Type + | lib.Int64Type + | lib.UInt8Type + | lib.UInt16Type + | lib.UInt32Type + | lib.UInt64Type +) +PyScalar: TypeAlias = ( + bool + | int + | float + | Decimal + | str + | bytes + | dt.date + | dt.datetime + | dt.time + | dt.timedelta +) +NumpyScalar: TypeAlias = np.generic[Any] + +_PyScalarT_co = TypeVar("_PyScalarT_co", bound=PyScalar, covariant=True) +_NumpyScalarT_co = TypeVar("_NumpyScalarT_co", bound=NumpyScalar, covariant=True) +_DataTypeT_co = TypeVar("_DataTypeT_co", bound=lib.DataType, covariant=True) +IntoArray = TypeAliasType( + "IntoArray", + Sequence[_PyScalarT_co | None] + | NDArray[_NumpyScalarT_co] + | lib.Array[lib.Scalar[_DataTypeT_co]] + | ChunkedArray[Any], + type_params=(_PyScalarT_co, _NumpyScalarT_co, _DataTypeT_co), +) + +Mask: TypeAlias = IntoArray[bool, np.bool_, lib.BoolType] +Indices: TypeAlias = IntoArray[int, np.integer[Any], IntegerType] _T = TypeVar("_T") _V = TypeVar("_V", covariant=True) SingleOrList: TypeAlias = list[_T] | _T - class SupportsDunderEQ(Protocol): def __eq__(self, other: object, /) -> bool: ... - -class SupportsDunderLT(Protocol): - def __lt__(self, other: object, /) -> bool: ... - - -class SupportsDunderGT(Protocol): - def __gt__(self, other: object, /) -> bool: ... - - -class SupportsDunderLE(Protocol): - def __le__(self, other: object, /) -> bool: ... - - -class SupportsDunderGE(Protocol): - def __ge__(self, other: object, /) -> bool: ... - - FilterTuple: TypeAlias = ( tuple[str, Literal["=", "==", "!="], SupportsDunderEQ] - | tuple[str, Literal["<"], SupportsDunderLT] - | tuple[str, Literal[">"], SupportsDunderGT] - | tuple[str, Literal["<="], SupportsDunderLE] - | tuple[str, Literal[">="], SupportsDunderGE] - | tuple[str, Literal["in", "not in"], Collection] + | 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=None, /) -> Any: ... - + def __arrow_c_stream__(self, requested_schema: Any = None, /) -> Any: ... class SupportsPyArrowArray(Protocol): - def __arrow_array__(self, type=None, /) -> Any: ... - + def __arrow_array__(self, type: Any = None, /) -> Any: ... class SupportsArrowArray(Protocol): - def __arrow_c_array__(self, requested_schema=None, /) -> Any: ... - + def __arrow_c_array__(self, requested_schema: Any = None, /) -> Any: ... class SupportsArrowDeviceArray(Protocol): - def __arrow_c_device_array__(self, requested_schema=None, /, **kwargs) -> Any: ... - + 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[_V], Protocol[_V]): +class NullableCollection(Sized, Container[Any], Protocol[_V]): def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... 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 index efc3bd847ca5..218993886531 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -1,727 +1,749 @@ -# 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 # noqa: F401 - -from collections.abc import Iterable, Iterator, Mapping, Sequence -from decimal import Decimal # noqa: F401 -from typing import Any, Generic, Literal, Protocol, TypeAlias - -import numpy as np -import pandas as pd - -from typing_extensions import Self, TypeVar, deprecated - -from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit -from pyarrow.lib import ( # noqa: F401 - Array, - Buffer, - ChunkedArray, - ExtensionArray, - ExtensionScalar, - MemoryPool, - MonthDayNano, - Table, -) - -class _Weakrefable: ... -class _Metadata(_Weakrefable): ... - -class DataType(_Weakrefable): - def field(self, i: int) -> Field[Any]: ... - @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) -> np.generic: ... - 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: ... - -_AsPyType = TypeVar("_AsPyType") -_DataTypeT = TypeVar("_DataTypeT", bound=DataType) -_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) - -class _BasicDataType(DataType, Generic[_AsPyType]): ... -class NullType(_BasicDataType[None]): ... -class BoolType(_BasicDataType[bool]): ... -class UInt8Type(_BasicDataType[int]): ... -class Int8Type(_BasicDataType[int]): ... -class UInt16Type(_BasicDataType[int]): ... -class Int16Type(_BasicDataType[int]): ... -class UInt32Type(_BasicDataType[int]): ... -class Int32Type(_BasicDataType[int]): ... -class UInt64Type(_BasicDataType[int]): ... -class Int64Type(_BasicDataType[int]): ... -class Float16Type(_BasicDataType[float]): ... -class Float32Type(_BasicDataType[float]): ... -class Float64Type(_BasicDataType[float]): ... -class Date32Type(_BasicDataType[dt.date]): ... -class Date64Type(_BasicDataType[dt.date]): ... -class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... -class StringType(_BasicDataType[str]): ... -class LargeStringType(_BasicDataType[str]): ... -class StringViewType(_BasicDataType[str]): ... -class BinaryType(_BasicDataType[bytes]): ... -class LargeBinaryType(_BasicDataType[bytes]): ... -class BinaryViewType(_BasicDataType[bytes]): ... - -_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) -_Tz = TypeVar("_Tz", str, None, default=None) - -class TimestampType(_BasicDataType[int], Generic[_Unit, _Tz]): - @property - def unit(self) -> _Unit: ... - @property - def tz(self) -> _Tz: ... - -_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) - -class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): - @property - def unit(self) -> _Time32Unit: ... - -_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) - -class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): - @property - def unit(self) -> _Time64Unit: ... - -class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): - @property - def unit(self) -> _Unit: ... - -_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) - -class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... - -_Precision = TypeVar("_Precision", default=Any) -_Scale = TypeVar("_Scale", default=Any) -_Precision_co = TypeVar("_Precision_co", default=Any, covariant=True) -_Scale_co = TypeVar("_Scale_co", default=Any, covariant=True) - -class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): - @property - def precision(self) -> _Precision_co: ... - @property - def scale(self) -> _Scale_co: ... - -class Decimal32Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal64Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal128Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class Decimal256Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... - -class ListType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class LargeListType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class ListViewType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class LargeListViewType(DataType, Generic[_DataTypeT_co]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - -class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): - @property - def value_field(self) -> Field[_DataTypeT_co]: ... - @property - def value_type(self) -> _DataTypeT_co: ... - @property - def list_size(self) -> int: ... - -class DictionaryMemo(_Weakrefable): ... - -_IndexT = TypeVar( - "_IndexT", - UInt8Type, - Int8Type, - UInt16Type, - Int16Type, - UInt32Type, - Int32Type, - UInt64Type, - Int64Type, -) -_BasicValueT = TypeVar( - "_BasicValueT", bound=_BasicDataType[Any], default=_BasicDataType[Any] -) -_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) -_K = TypeVar("_K", bound=DataType, default=DataType) -_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) - -class DictionaryType(DataType, Generic[_IndexT, _BasicValueT, _Ordered]): - @property - def ordered(self) -> _Ordered: ... - @property - def index_type(self) -> _IndexT: ... - @property - def value_type(self) -> _BasicValueT: ... - -class MapType(DataType, Generic[_K, _ValueT, _Ordered]): - @property - def key_field(self) -> Field[_K]: ... - @property - def key_type(self) -> _K: ... - @property - def item_field(self) -> Field[_ValueT]: ... - @property - def item_type(self) -> _ValueT: ... - @property - def keys_sorted(self) -> _Ordered: ... - -_Size = TypeVar("_Size", default=int) - -class StructType(DataType): - def get_field_index(self, name: str) -> int: ... - def field(self, i: int | str) -> Field[Any]: ... - def get_all_field_indices(self, name: str) -> list[int]: ... - def __len__(self) -> int: ... - def __iter__(self) -> Iterator[Field[Any]]: ... - - __getitem__ = field - @property - def names(self) -> list[str]: ... - @property - def fields(self) -> list[Field[Any]]: ... - -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[Any]]: ... - def field(self, i: int) -> Field[Any]: ... - - __getitem__ = field - -class SparseUnionType(UnionType): - @property - def mode(self) -> Literal["sparse"]: ... - -class DenseUnionType(UnionType): - @property - def mode(self) -> Literal["dense"]: ... - -_RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) - -class RunEndEncodedType(DataType, Generic[_RunEndType, _BasicValueT]): - @property - def run_end_type(self) -> _RunEndType: ... - @property - def value_type(self) -> _BasicValueT: ... - -_StorageT = TypeVar("_StorageT", bound=Array[Any] | ChunkedArray[Any]) - -class BaseExtensionType(DataType): - def __arrow_ext_class__(self) -> type[ExtensionArray[Any]]: ... - 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: _StorageT) -> _StorageT: ... - -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, Generic[_ValueT]): - @property - def value_type(self) -> _ValueT: ... - @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: ExtensionType) -> None: ... -def unregister_extension_type(type_name: str) -> None: ... - -_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 -_FieldTypeInput: TypeAlias = _DataTypeLike | None - -class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): - def __init__( - self, - __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, - **kwargs: str, - ) -> None: ... - def equals(self, other: KeyValueMetadata) -> bool: ... - def __len__(self) -> int: ... - def __contains__(self, /, __key: object) -> bool: ... - def __getitem__(self, /, __key: Any) -> Any: ... - def __iter__(self) -> Iterator[bytes]: ... - def get_all(self, key: str) -> list[bytes]: ... - def to_dict(self) -> dict[bytes, bytes]: ... - -class Field(_Weakrefable, Generic[_DataTypeT_co]): - def equals(self, other: Field[Any], 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) -> _DataTypeT_co: ... - def with_metadata( - self, - metadata: _MetadataMapping | Any, - ) -> Self: ... - def remove_metadata(self) -> Self: ... - def with_type(self, new_type: _DataTypeT) -> Field[_DataTypeT]: ... - def with_name(self, name: str) -> Self: ... - def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... - def flatten(self) -> list[Field[Any]]: ... - 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: ... - -_StructFieldTuple: TypeAlias = ( - tuple[str, Field[Any] | None] | tuple[str, _FieldTypeInput] -) -_StructFieldsInput: TypeAlias = ( - Iterable[Field[Any] | _StructFieldTuple] - | Mapping[str, Field[Any] | _FieldTypeInput] -) - -class Schema(_Weakrefable): - def __len__(self) -> int: ... - def __getitem__(self, key: str | int) -> Field[Any]: ... - - _field = __getitem__ - def __iter__(self) -> Iterator[Field[Any]]: ... - def __hash__(self) -> int: ... - def __sizeof__(self) -> int: ... - @property - def pandas_metadata(self) -> dict[bytes, bytes]: ... - @property - def names(self) -> list[str]: ... - @property - def types(self) -> list[DataType]: ... - @property - def metadata(self) -> dict[bytes, bytes]: ... - 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[Any]: ... - @deprecated("Use 'field' instead") - def field_by_name(self, name: str) -> Field[Any]: ... - def get_field_index(self, name: str) -> int: ... - def get_all_field_indices(self, name: str) -> list[int]: ... - def append(self, field: Field[Any]) -> Schema: ... - def insert(self, i: int, field: Field[Any]) -> Schema: ... - def remove(self, i: int) -> Schema: ... - def set(self, i: int, field: Field[Any]) -> Schema: ... - @deprecated("Use 'with_metadata' instead") - def add_metadata(self, metadata: _MetadataMapping) -> Schema: ... - def with_metadata(self, metadata: _MetadataMapping) -> 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 | None = None, - ) -> 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: Sequence[Schema], - *, - promote_options: Literal["default", "permissive"] = "default", -) -> Schema: ... -def field( - name: SupportsArrowSchema | str | bytes, - type: _DataTypeT | _DataTypeAliasInput | None = None, - nullable: bool = True, - metadata: _MetadataMapping | None = None, -) -> Field[_DataTypeT] | Field[Any]: ... -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: _Unit | str, tz: _Tz | None = None -) -> TimestampType[_Unit, _Tz]: ... -def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... -def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... -def duration(unit: _Unit | str) -> DurationType[_Unit]: ... -def month_day_nano_interval() -> MonthDayNanoIntervalType: ... -def date32() -> Date32Type: ... -def date64() -> Date64Type: ... -def float16() -> Float16Type: ... -def float32() -> Float32Type: ... -def float64() -> Float64Type: ... -def decimal32( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... -def decimal64( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... -def decimal128( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... -def decimal256( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... -def string() -> StringType: ... - -utf8 = string - -def binary( - length: Literal[-1] | int = ..., -) -> BinaryType | FixedSizeBinaryType[bytes]: ... -def large_binary() -> LargeBinaryType: ... -def large_string() -> LargeStringType: ... - -large_utf8 = large_string - -def binary_view() -> BinaryViewType: ... -def string_view() -> StringViewType: ... -from typing import overload - -@overload -def list_( - value_type: _DataTypeT | Field[_DataTypeT], list_size: Literal[-1] | None = None -) -> ListType[_DataTypeT]: ... -@overload -def list_( - value_type: _DataTypeT | Field[_DataTypeT], list_size: _Size -) -> FixedSizeListType[_DataTypeT, _Size]: ... -def large_list( - value_type: _DataTypeT | Field[_DataTypeT], -) -> LargeListType[_DataTypeT]: ... -def list_view( - value_type: _DataTypeT | Field[_DataTypeT], -) -> ListViewType[_DataTypeT]: ... -def large_list_view( - value_type: _DataTypeT | Field[_DataTypeT], -) -> LargeListViewType[_DataTypeT]: ... -def map_( - key_type: _K | Field[_K] | _FieldTypeInput, - item_type: _ValueT | Field[_ValueT] | _FieldTypeInput, - keys_sorted: _Ordered | None = None, -) -> MapType[_K, _ValueT, _Ordered]: ... -def dictionary( - index_type: _IndexT | _DataTypeAliasInput, - value_type: _BasicValueT | _DataTypeAliasInput, - ordered: _Ordered | None = None, -) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ... -def struct( - fields: _StructFieldsInput, -) -> StructType: ... -def sparse_union( - child_fields: list[Field[Any]], type_codes: list[int] | None = None -) -> SparseUnionType: ... -def dense_union( - child_fields: list[Field[Any]], type_codes: list[int] | None = None -) -> DenseUnionType: ... -def union( - child_fields: list[Field[Any]], - mode: Literal["sparse", "dense"] | int, - type_codes: list[int] | None = None, -) -> SparseUnionType | DenseUnionType: ... -def run_end_encoded( - run_end_type: _RunEndType | _DataTypeAliasInput | None, - value_type: _BasicValueT | _DataTypeAliasInput | None, -) -> RunEndEncodedType[_RunEndType, _BasicValueT]: ... -def json_(storage_type: DataType = ...) -> JsonType: ... -def uuid() -> UuidType: ... -def fixed_shape_tensor( - value_type: _ValueT, - shape: Sequence[int], - dim_names: Sequence[str] | None = None, - permutation: Sequence[int] | None = None, -) -> FixedShapeTensorType[_ValueT]: ... -def bool8() -> Bool8Type: ... -def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... -def type_for_alias(name: Any) -> DataType: ... -def schema( - fields: ( - Iterable[Field[Any]] - | Iterable[tuple[str, _FieldTypeInput]] - | Mapping[Any, _FieldTypeInput] - ), - metadata: _MetadataMapping | None = None, -) -> Schema: ... -def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... - -__all__ = [ - "_Weakrefable", - "_Metadata", - "DataType", - "_BasicDataType", - "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", - "_Unit", - "_Tz", - "_Time32Unit", - "_Time64Unit", - "_DataTypeT", -] +# 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 # noqa: F401 + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from decimal import Decimal # noqa: F401 +from typing import Any, Generic, Literal, Protocol, TypeAlias, overload + +import numpy as np +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import Self, TypeVar, deprecated + +from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow.lib import ( # noqa: F401 + Array, + Buffer, + ChunkedArray, + ExtensionArray, + ExtensionScalar, + MemoryPool, + MonthDayNano, + Table, +) + +class _Weakrefable: ... +class _Metadata(_Weakrefable): ... + +class DataType(_Weakrefable): + def field(self, i: int) -> Field[Any]: ... + @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: ... + +_AsPyType = TypeVar("_AsPyType") +_DataTypeT = TypeVar("_DataTypeT", bound=DataType) +_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) +_Size = TypeVar("_Size", bound=int, default=int) +_PandasDtype: TypeAlias = ( + type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype +) + +class _BasicDataType(DataType, Generic[_AsPyType]): ... +class NullType(_BasicDataType[None]): ... +class BoolType(_BasicDataType[bool]): ... +class UInt8Type(_BasicDataType[int]): ... +class Int8Type(_BasicDataType[int]): ... +class UInt16Type(_BasicDataType[int]): ... +class Int16Type(_BasicDataType[int]): ... +class UInt32Type(_BasicDataType[int]): ... +class Int32Type(_BasicDataType[int]): ... +class UInt64Type(_BasicDataType[int]): ... +class Int64Type(_BasicDataType[int]): ... +class Float16Type(_BasicDataType[float]): ... +class Float32Type(_BasicDataType[float]): ... +class Float64Type(_BasicDataType[float]): ... +class Date32Type(_BasicDataType[dt.date]): ... +class Date64Type(_BasicDataType[dt.date]): ... +class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... +class StringType(_BasicDataType[str]): ... +class LargeStringType(_BasicDataType[str]): ... +class StringViewType(_BasicDataType[str]): ... +class BinaryType(_BasicDataType[bytes]): ... +class LargeBinaryType(_BasicDataType[bytes]): ... +class BinaryViewType(_BasicDataType[bytes]): ... + +_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) +_Tz = TypeVar("_Tz", str, None, default=None) + +class TimestampType(_BasicDataType[dt.datetime], Generic[_Unit, _Tz]): + @property + def unit(self) -> _Unit: ... + @property + def tz(self) -> _Tz: ... + +_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) + +class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): + @property + def unit(self) -> _Time32Unit: ... + +_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) + +class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): + @property + def unit(self) -> _Time64Unit: ... + +class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): + @property + def unit(self) -> _Unit: ... + +_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) + +class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... + +_Precision = TypeVar("_Precision", bound=int, default=int) +_Scale = TypeVar("_Scale", bound=int, default=int) +_Precision_co = TypeVar("_Precision_co", bound=int, default=int, covariant=True) +_Scale_co = TypeVar("_Scale_co", bound=int, default=int, covariant=True) + +class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): + @property + def precision(self) -> _Precision_co: ... + @property + def scale(self) -> _Scale_co: ... + +class Decimal32Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... +class Decimal64Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... +class Decimal128Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... +class Decimal256Type( + FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] +): ... + +class ListType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class LargeListType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class ListViewType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class LargeListViewType(DataType, Generic[_DataTypeT_co]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + +class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): + @property + def value_field(self) -> Field[_DataTypeT_co]: ... + @property + def value_type(self) -> _DataTypeT_co: ... + @property + def list_size(self) -> _Size: ... + +class DictionaryMemo(_Weakrefable): ... + +_IndexT = TypeVar( + "_IndexT", + UInt8Type, + Int8Type, + UInt16Type, + Int16Type, + UInt32Type, + Int32Type, + UInt64Type, + Int64Type, +) +_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) +_ValueT_co = TypeVar("_ValueT_co", bound=DataType, default=DataType, covariant=True) +_K = TypeVar("_K", bound=DataType, default=DataType) +_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) + +class DictionaryType(DataType, Generic[_IndexT, _ValueT_co, _Ordered]): + @property + def ordered(self) -> _Ordered: ... + @property + def index_type(self) -> _IndexT: ... + @property + def value_type(self) -> _ValueT_co: ... + +class MapType(DataType, Generic[_K, _ValueT, _Ordered]): + @property + def key_field(self) -> Field[_K]: ... + @property + def key_type(self) -> _K: ... + @property + def item_field(self) -> Field[_ValueT]: ... + @property + def item_type(self) -> _ValueT: ... + @property + def keys_sorted(self) -> _Ordered: ... + +class StructType(DataType): + def get_field_index(self, name: str | bytes) -> int: ... + def field(self, i: int | str | bytes) -> Field[Any]: ... + def get_all_field_indices(self, name: str | bytes) -> list[int]: ... + def __len__(self) -> int: ... + def __iter__(self) -> Iterator[Field[Any]]: ... + + __getitem__ = field + @property + def names(self) -> list[str]: ... + @property + def fields(self) -> list[Field[Any]]: ... + +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[Any]]: ... + def field(self, i: int) -> Field[Any]: ... + + __getitem__ = field + +class SparseUnionType(UnionType): + @property + def mode(self) -> Literal["sparse"]: ... + +class DenseUnionType(UnionType): + @property + def mode(self) -> Literal["dense"]: ... + +_RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) + +class RunEndEncodedType(DataType, Generic[_RunEndType, _ValueT_co]): + @property + def run_end_type(self) -> _RunEndType: ... + @property + def value_type(self) -> _ValueT_co: ... + +_StorageT = TypeVar("_StorageT", bound=Array[Any] | ChunkedArray[Any]) + +class BaseExtensionType(DataType): + def __arrow_ext_class__(self) -> type[ExtensionArray[Any]]: ... + 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: _StorageT) -> _StorageT: ... + +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, Generic[_ValueT]): + @property + def value_type(self) -> _ValueT: ... + @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, + ) -> 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, Generic[_DataTypeT_co]): + def equals(self, other: Field[Any], 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) -> _DataTypeT_co: ... + def with_metadata( + self, + metadata: _KeyValueMetadataInput | KeyValueMetadata, + ) -> Self: ... + def remove_metadata(self) -> Self: ... + def with_type(self, new_type: _DataTypeT) -> Field[_DataTypeT]: ... + def with_name(self, name: str | bytes) -> Self: ... + def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... + def flatten(self) -> list[Field[Any]]: ... + 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[Any] | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] +) +_SchemaFieldsInput: TypeAlias = _StructFieldsInput + +class Schema(_Weakrefable): + def __len__(self) -> int: ... + def __getitem__(self, key: int) -> Field[Any]: ... + + _field = __getitem__ + def __iter__(self) -> Iterator[Field[Any]]: ... + 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[Any]: ... + @deprecated("Use 'field' instead") + def field_by_name(self, name: str | bytes) -> Field[Any] | 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[Any]) -> Schema: ... + def insert(self, i: int, field: Field[Any]) -> Schema: ... + def remove(self, i: int) -> Schema: ... + def set(self, i: int, field: Field[Any]) -> 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[Any]: ... +@overload +def field( + name: _FieldName, + type: _DataTypeT, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field[_DataTypeT]: ... +@overload +def field( + name: _FieldName, + type: _DataTypeAliasInput, + nullable: bool | None = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata = None, +) -> Field[Any]: ... +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: _Unit | str, tz: _Tz | None = None +) -> TimestampType[_Unit, _Tz]: ... +def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... +def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... +def duration(unit: _Unit | str) -> DurationType[_Unit]: ... +def month_day_nano_interval() -> MonthDayNanoIntervalType: ... +def date32() -> Date32Type: ... +def date64() -> Date64Type: ... +def float16() -> Float16Type: ... +def float32() -> Float32Type: ... +def float64() -> Float64Type: ... +def decimal32( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... +def decimal64( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... +def decimal128( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... +def decimal256( + precision: _Precision, scale: _Scale | Literal[0] = 0 +) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... +def string() -> StringType: ... + +utf8 = string + +def binary( + length: Literal[-1] | int = ..., +) -> BinaryType | FixedSizeBinaryType[bytes]: ... +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: _DataTypeT | Field[_DataTypeT]) -> ListType[_DataTypeT]: ... +@overload +def list_( + value_type: _DataTypeT | Field[_DataTypeT], list_size: Literal[-1] +) -> ListType[_DataTypeT]: ... +@overload +def list_( + value_type: _DataTypeT | Field[_DataTypeT], list_size: _Size +) -> ListType[_DataTypeT] | FixedSizeListType[_DataTypeT, _Size]: ... +def large_list( + value_type: _DataTypeT | Field[_DataTypeT], +) -> LargeListType[_DataTypeT]: ... +def list_view( + value_type: _DataTypeT | Field[_DataTypeT], +) -> ListViewType[_DataTypeT]: ... +def large_list_view( + value_type: _DataTypeT | Field[_DataTypeT], +) -> LargeListViewType[_DataTypeT]: ... +def map_( + key_type: _K | Field[_K] | _DataTypeAliasInput, + item_type: _ValueT | Field[_ValueT] | _DataTypeAliasInput, + keys_sorted: _Ordered | None = None, +) -> MapType[_K, _ValueT, _Ordered]: ... +def dictionary( + index_type: _IndexT | _DataTypeAliasInput, + value_type: _ValueT | _DataTypeAliasInput, + ordered: _Ordered | None = None, +) -> DictionaryType[_IndexT, _ValueT, _Ordered]: ... +def struct( + fields: _StructFieldsInput, +) -> StructType: ... +def sparse_union( + child_fields: Sequence[Field[Any]], type_codes: Sequence[int] | None = None +) -> SparseUnionType: ... +def dense_union( + child_fields: Sequence[Field[Any]], type_codes: Sequence[int] | None = None +) -> DenseUnionType: ... +def union( + child_fields: Sequence[Field[Any]], + mode: Literal["sparse", "dense"] | int, + type_codes: Sequence[int] | None = None, +) -> SparseUnionType | DenseUnionType: ... +def run_end_encoded( + run_end_type: _RunEndType | _DataTypeAliasInput, + value_type: _ValueT | _DataTypeAliasInput, +) -> RunEndEncodedType[_RunEndType, _ValueT]: ... +def json_(storage_type: DataType = ...) -> JsonType: ... +def uuid() -> UuidType: ... +def fixed_shape_tensor( + value_type: _ValueT, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, + permutation: Sequence[int] | None = None, +) -> FixedShapeTensorType[_ValueT]: ... +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", + "_BasicDataType", + "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", + "_Unit", + "_Tz", + "_Time32Unit", + "_Time64Unit", + "_DataTypeT", +] diff --git a/python/pyarrow-stubs/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/error.pyi index 8ee75c3ec414..67be73eadb6d 100644 --- a/python/pyarrow-stubs/pyarrow/error.pyi +++ b/python/pyarrow-stubs/pyarrow/error.pyi @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +from types import TracebackType + from typing_extensions import Self class ArrowException(Exception): ... @@ -41,7 +43,12 @@ have_signal_refcycle: bool class SignalStopHandler: def __enter__(self) -> Self: ... - def __exit__(self, exc_type, exc_value, exc_tb) -> None: ... + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: ... @property def stop_token(self) -> StopToken: ... From d2cc8826f9eba2e025ff747add935bf296103c12 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:22:40 -0700 Subject: [PATCH 17/93] lint --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 67bebe09398b..df3e8b799d77 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -144,7 +144,7 @@ repos: - "--extend-select" - "Y" - "--per-file-ignores" - - "python/pyarrow-stubs/pyarrow/*.pyi:E301,E302,E305,E701" + - "python/pyarrow-stubs/pyarrow/*.pyi:E301,E302,E305,E701,F821" files: >- ^(c_glib|dev|python)/ types: From dce40ae9263c05dacf452dac426bdcf852948653 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:44:05 -0700 Subject: [PATCH 18/93] more covariant types --- python/pyarrow-stubs/pyarrow/_types.pyi | 53 +++++++++++++++++-------- 1 file changed, 36 insertions(+), 17 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 218993886531..1e98aefc8186 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -196,30 +196,46 @@ _IndexT = TypeVar( UInt64Type, Int64Type, ) +_IndexT_co = TypeVar( + "_IndexT_co", + UInt8Type, + Int8Type, + UInt16Type, + Int16Type, + UInt32Type, + Int32Type, + UInt64Type, + Int64Type, + covariant=True, +) _ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) _ValueT_co = TypeVar("_ValueT_co", bound=DataType, default=DataType, covariant=True) _K = TypeVar("_K", bound=DataType, default=DataType) +_K_co = TypeVar("_K_co", bound=DataType, default=DataType, covariant=True) _Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) +_Ordered_co = TypeVar( + "_Ordered_co", Literal[True], Literal[False], default=Literal[False], covariant=True +) -class DictionaryType(DataType, Generic[_IndexT, _ValueT_co, _Ordered]): +class DictionaryType(DataType, Generic[_IndexT_co, _ValueT_co, _Ordered_co]): @property - def ordered(self) -> _Ordered: ... + def ordered(self) -> _Ordered_co: ... @property - def index_type(self) -> _IndexT: ... + def index_type(self) -> _IndexT_co: ... @property def value_type(self) -> _ValueT_co: ... -class MapType(DataType, Generic[_K, _ValueT, _Ordered]): +class MapType(DataType, Generic[_K_co, _ValueT_co, _Ordered_co]): @property - def key_field(self) -> Field[_K]: ... + def key_field(self) -> Field[_K_co]: ... @property - def key_type(self) -> _K: ... + def key_type(self) -> _K_co: ... @property - def item_field(self) -> Field[_ValueT]: ... + def item_field(self) -> Field[_ValueT_co]: ... @property - def item_type(self) -> _ValueT: ... + def item_type(self) -> _ValueT_co: ... @property - def keys_sorted(self) -> _Ordered: ... + def keys_sorted(self) -> _Ordered_co: ... class StructType(DataType): def get_field_index(self, name: str | bytes) -> int: ... @@ -254,10 +270,13 @@ class DenseUnionType(UnionType): def mode(self) -> Literal["dense"]: ... _RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) +_RunEndType_co = TypeVar( + "_RunEndType_co", Int16Type, Int32Type, Int64Type, covariant=True +) -class RunEndEncodedType(DataType, Generic[_RunEndType, _ValueT_co]): +class RunEndEncodedType(DataType, Generic[_RunEndType_co, _ValueT_co]): @property - def run_end_type(self) -> _RunEndType: ... + def run_end_type(self) -> _RunEndType_co: ... @property def value_type(self) -> _ValueT_co: ... @@ -280,9 +299,9 @@ class ExtensionType(BaseExtensionType): cls, storage_type: DataType, serialized: bytes ) -> Self: ... -class FixedShapeTensorType(BaseExtensionType, Generic[_ValueT]): +class FixedShapeTensorType(BaseExtensionType, Generic[_ValueT_co]): @property - def value_type(self) -> _ValueT: ... + def value_type(self) -> _ValueT_co: ... @property def shape(self) -> list[int]: ... @property @@ -380,7 +399,7 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( self, __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, - **kwargs: str, + **kwargs: str | bytes, ) -> None: ... def equals(self, other: KeyValueMetadata) -> bool: ... def __len__(self) -> int: ... @@ -595,13 +614,13 @@ def struct( fields: _StructFieldsInput, ) -> StructType: ... def sparse_union( - child_fields: Sequence[Field[Any]], type_codes: Sequence[int] | None = None + child_fields: Iterable[Field[Any]], type_codes: Sequence[int] | None = None ) -> SparseUnionType: ... def dense_union( - child_fields: Sequence[Field[Any]], type_codes: Sequence[int] | None = None + child_fields: Iterable[Field[Any]], type_codes: Sequence[int] | None = None ) -> DenseUnionType: ... def union( - child_fields: Sequence[Field[Any]], + child_fields: Iterable[Field[Any]], mode: Literal["sparse", "dense"] | int, type_codes: Sequence[int] | None = None, ) -> SparseUnionType | DenseUnionType: ... From f278750d2b6e69acc435434c1cd2bfdbbda3ebb3 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 21:15:15 +0200 Subject: [PATCH 19/93] Apply suggestions from code review Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- python/pyarrow-stubs/pyarrow/_types.pyi | 2 +- python/pyarrow/fs.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 1e98aefc8186..8063ac399219 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -22,7 +22,7 @@ from decimal import Decimal # noqa: F401 from typing import Any, Generic, Literal, Protocol, TypeAlias, overload import numpy as np -import pandas as pd # type: ignore[import-not-found,import-untyped] +import pandas as pd # type: ignore[import-not-found,import-untyped] # pyright: ignore[reportMissingImports,reportMissingTypeStubs] from typing_extensions import Self, TypeVar, deprecated from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit diff --git a/python/pyarrow/fs.py b/python/pyarrow/fs.py index f055c5081826..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 # type: ignore[import-untyped] + import fsspec # type: ignore[import-not-found,import-untyped] except ImportError: pass else: From ca7c54e7fdfe039f0b499a42864e78be85db30fc Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 12:30:43 -0700 Subject: [PATCH 20/93] lint --- python/pyarrow-stubs/pyarrow/_types.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 8063ac399219..941caa761446 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -22,7 +22,7 @@ from decimal import Decimal # noqa: F401 from typing import Any, Generic, Literal, Protocol, TypeAlias, overload import numpy as np -import pandas as pd # type: ignore[import-not-found,import-untyped] # pyright: ignore[reportMissingImports,reportMissingTypeStubs] +import pandas as pd # type: ignore[import-not-found,import-untyped] # pyright: ignore[reportMissingImports,reportMissingTypeStubs] # noqa: E501 from typing_extensions import Self, TypeVar, deprecated from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit From c0325965c29e707889e55ee48776f303c7a8f212 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Feb 2026 20:54:52 +0100 Subject: [PATCH 21/93] Add type stubs for core data structures Add type annotation stubs for array, table, tensor, builder, memory, device, config, and types modules. Includes type-ignore annotations in related tests. --- python/pyarrow-stubs/pyarrow/array.pyi | 894 +++++++++++++++++++ python/pyarrow-stubs/pyarrow/builder.pyi | 51 ++ python/pyarrow-stubs/pyarrow/config.pyi | 72 ++ python/pyarrow-stubs/pyarrow/device.pyi | 66 ++ python/pyarrow-stubs/pyarrow/memory.pyi | 94 ++ python/pyarrow-stubs/pyarrow/table.pyi | 686 ++++++++++++++ python/pyarrow-stubs/pyarrow/tensor.pyi | 268 ++++++ python/pyarrow-stubs/pyarrow/types.pyi | 227 +++++ python/pyarrow/array.pxi | 2 +- python/pyarrow/scalar.pxi | 2 +- python/pyarrow/tests/test_array.py | 71 +- python/pyarrow/tests/test_convert_builtin.py | 42 +- python/pyarrow/tests/test_device.py | 12 +- python/pyarrow/tests/test_schema.py | 11 +- python/pyarrow/tests/test_sparse_tensor.py | 20 +- python/pyarrow/tests/test_table.py | 120 ++- python/pyarrow/tests/test_tensor.py | 2 +- python/pyarrow/tests/test_types.py | 55 +- 18 files changed, 2561 insertions(+), 134 deletions(-) create mode 100644 python/pyarrow-stubs/pyarrow/array.pyi create mode 100644 python/pyarrow-stubs/pyarrow/builder.pyi create mode 100644 python/pyarrow-stubs/pyarrow/config.pyi create mode 100644 python/pyarrow-stubs/pyarrow/device.pyi create mode 100644 python/pyarrow-stubs/pyarrow/memory.pyi create mode 100644 python/pyarrow-stubs/pyarrow/table.pyi create mode 100644 python/pyarrow-stubs/pyarrow/tensor.pyi create mode 100644 python/pyarrow-stubs/pyarrow/types.pyi diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi new file mode 100644 index 000000000000..547e9c949d52 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -0,0 +1,894 @@ +# 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, + Generic, + Literal, + TypeVar, +) + +import numpy as np +import pandas as pd + +from pyarrow._compute import CastOptions +from pyarrow._stubs_typing import ( + ArrayLike, + Indices, + Mask, + Order, + SupportArrowArray, + SupportArrowDeviceArray, + SupportPyArrowArray, +) +from pyarrow.lib import ( + Buffer, + Device, + MemoryManager, + MemoryPool, + Tensor, + _Weakrefable, +) +from typing_extensions import deprecated +import builtins + +from .scalar import ( # noqa: F401 + BinaryScalar, + BinaryViewScalar, + BooleanScalar, + Date32Scalar, + Date64Scalar, + DictionaryScalar, + DoubleScalar, + DurationScalar, + ExtensionScalar, + FixedSizeBinaryScalar, + FixedSizeListScalar, + FloatScalar, + HalfFloatScalar, + Int16Scalar, + Int32Scalar, + Int64Scalar, + Int8Scalar, + LargeBinaryScalar, + LargeListScalar, + LargeStringScalar, + ListScalar, + ListViewScalar, + MapScalar, + MonthDayNanoIntervalScalar, + NullScalar, + RunEndEncodedScalar, + Scalar, + StringScalar, + StringViewScalar, + StructScalar, + Time32Scalar, + Time64Scalar, + TimestampScalar, + UInt16Scalar, + UInt32Scalar, + UInt64Scalar, + UInt8Scalar, + UnionScalar, +) +from .device import DeviceAllocationType +from ._types import ( # noqa: F401 + BaseExtensionType, + BinaryType, + DataType, + Field, + Float64Type, + Int64Type, + MapType, + StringType, + StructType, + _AsPyType, + _BasicDataType, + _BasicValueT, + _DataTypeT, + _IndexT, + _RunEndType, + _Size, + _Time32Unit, + _Time64Unit, + _Tz, + _Unit, +) +from ._stubs_typing import NullableCollection + + +def array( + values: NullableCollection[Any] | Iterable[Any] | SupportArrowArray + | SupportArrowDeviceArray | SupportPyArrowArray, + 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[Any] | Iterable[Any] | SupportArrowArray + | SupportArrowDeviceArray, + type: _DataTypeT | Any | None = None, +) -> Array[Scalar[_DataTypeT]] | 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: ... + + +_ConvertAs = TypeVar("_ConvertAs", pd.DataFrame, pd.Series) + + +class _PandasConvertible(_Weakrefable, Generic[_ConvertAs]): + 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, + ) -> _ConvertAs: ... + + +_CastAs = TypeVar("_CastAs", bound=DataType) +_Scalar_co = TypeVar("_Scalar_co", bound=Scalar, covariant=True) +_ScalarT = TypeVar("_ScalarT", bound=Scalar) + + +class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): + def as_py(self) -> list[Any]: ... + + 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: _CastAs | str, + safe: bool = True, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None, + ) -> Array[Scalar[_CastAs]]: ... + + def view(self, target_type: _CastAs) -> Array[Scalar[_CastAs]]: ... + + def sum(self, **kwargs) -> _Scalar_co: ... + + @property + def type(self: Array[Scalar[_DataTypeT]]) -> _DataTypeT: ... + 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: _DataTypeT | None = None, + safe: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Array[Scalar[_DataTypeT]] | Array[Scalar]: ... + + @staticmethod + def from_buffers( + type: _DataTypeT, + length: int, + buffers: Sequence[Buffer | None], + null_count: int = -1, + offset=0, + children: NullableCollection[Array[Scalar[_DataTypeT]]] | None = None, + ) -> Array[Scalar[_DataTypeT]]: ... + + @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_co]: ... + + 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: Array[Scalar[_BasicDataType[_AsPyType]]], fill_value: _AsPyType + ) -> Array[Scalar[_BasicDataType[_AsPyType]]]: ... + + def __getitem__(self, key: int | builtins.slice) -> _Scalar_co | 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: Array[_ScalarT] | Array[Scalar[_BasicDataType[_AsPyType]]], + value: _ScalarT | _AsPyType, + 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[NullScalar]): + ... + + +class BooleanArray(Array[BooleanScalar]): + @property + def false_count(self) -> int: ... + @property + def true_count(self) -> int: ... + + +class NumericArray(Array[_ScalarT]): + ... + + +class IntegerArray(NumericArray[_ScalarT]): + ... + + +class FloatingPointArray(NumericArray[_ScalarT]): + ... + + +class Int8Array(IntegerArray[Int8Scalar]): + ... + + +class UInt8Array(IntegerArray[UInt8Scalar]): + ... + + +class Int16Array(IntegerArray[Int16Scalar]): + ... + + +class UInt16Array(IntegerArray[UInt16Scalar]): + ... + + +class Int32Array(IntegerArray[Int32Scalar]): + ... + + +class UInt32Array(IntegerArray[UInt32Scalar]): + ... + + +class Int64Array(IntegerArray[Int64Scalar]): + ... + + +class UInt64Array(IntegerArray[UInt64Scalar]): + ... + + +class Date32Array(NumericArray[Date32Scalar]): + ... + + +class Date64Array(NumericArray[Date64Scalar]): + ... + + +class TimestampArray(NumericArray[TimestampScalar[_Unit, _Tz]]): + ... + + +class Time32Array(NumericArray[Time32Scalar[_Time32Unit]]): + ... + + +class Time64Array(NumericArray[Time64Scalar[_Time64Unit]]): + ... + + +class DurationArray(NumericArray[DurationScalar[_Unit]]): + ... + + +class MonthDayNanoIntervalArray(Array[MonthDayNanoIntervalScalar]): + ... + + +class HalfFloatArray(FloatingPointArray[HalfFloatScalar]): + ... + + +class FloatArray(FloatingPointArray[FloatScalar]): + ... + + +class DoubleArray(FloatingPointArray[DoubleScalar]): + ... + + +class FixedSizeBinaryArray(Array[FixedSizeBinaryScalar]): + ... + + +class Decimal32Array(FixedSizeBinaryArray): + ... + + +class Decimal64Array(FixedSizeBinaryArray): + ... + + +class Decimal128Array(FixedSizeBinaryArray): + ... + + +class Decimal256Array(FixedSizeBinaryArray): + ... + + +class BaseListArray(Array[_ScalarT]): + def flatten(self, recursive: bool = False) -> Array: ... + + def value_parent_indices(self) -> Int64Array: ... + + def value_lengths(self) -> Int32Array: ... + + +class ListArray(BaseListArray[_ScalarT]): + @classmethod + def from_arrays( + cls, + offsets: Int32Array | list[int] | list[int | None], + values: Array[Scalar[_DataTypeT]] | list[int] | list[float] | list[str] + | list[bytes] | list, + *, + type: _DataTypeT | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> (ListArray[ListScalar[ + _DataTypeT | Int64Type | Float64Type | StringType | BinaryType + ]] | ListArray): ... + + @property + def values(self) -> Array: ... + + @property + def offsets(self) -> Int32Array: ... + + +class LargeListArray(BaseListArray[LargeListScalar[_DataTypeT]]): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | list[int | None], + values: Array[Scalar[_DataTypeT]] | Array, + *, + type: _DataTypeT | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListArray[_DataTypeT]: ... + + @property + def values(self) -> Array: ... + + @property + def offsets(self) -> Int64Array: ... + + +class ListViewArray(BaseListArray[ListViewScalar[_DataTypeT]]): + @classmethod + def from_arrays( + cls, + offsets: Int32Array, + values: Array[Scalar[_DataTypeT]] | Array, + *, + type: _DataTypeT | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> ListViewArray[_DataTypeT]: ... + + @property + def values(self) -> Array: ... + + @property + def offsets(self) -> Int32Array: ... + + @property + def sizes(self) -> Int32Array: ... + + +class LargeListViewArray(BaseListArray[LargeListScalar[_DataTypeT]]): + @classmethod + def from_arrays( + cls, + offsets: Int64Array, + values: Array[Scalar[_DataTypeT]] | Array, + *, + type: _DataTypeT | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> LargeListViewArray[_DataTypeT]: ... + + @property + def values(self) -> Array: ... + + @property + def offsets(self) -> Int64Array: ... + + @property + def sizes(self) -> Int64Array: ... + + +class FixedSizeListArray(BaseListArray[FixedSizeListScalar[_DataTypeT, _Size]]): + @classmethod + def from_arrays( + cls, + values: Array[Scalar[_DataTypeT]], + list_size: _Size | None = None, + *, + type: DataType | None = None, + mask: Mask | None = None, + ) -> FixedSizeListArray[_DataTypeT, _Size | None]: ... + + @property + def values(self) -> BaseListArray[ListScalar[_DataTypeT]]: ... + + +_MapKeyT = TypeVar("_MapKeyT", bound=_BasicDataType) +_MapItemT = TypeVar("_MapItemT", bound=_BasicDataType) + + +class MapArray(BaseListArray[MapScalar[_MapKeyT, _MapItemT]]): + @classmethod + def from_arrays( + cls, + offsets: Int64Array | list[int] | None, + keys: Array[Scalar[_MapKeyT]] | np.ndarray | list | None = None, + items: Array[Scalar[_MapItemT]] | np.ndarray | list | None = None, + values: Array | DataType | None = None, + *, + type: DataType | None = None, + pool: MemoryPool | None = None, + mask: Mask | None = None, + ) -> MapArray[_MapKeyT, _MapItemT]: ... + + @property + def keys(self) -> Array: ... + + @property + def items(self) -> Array: ... + + +class UnionArray(Array[UnionScalar]): + @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[Array], + field_names: list[str] | None = None, + type_codes: Int8Array | list[int] | None = None, + ) -> UnionArray: ... + + @staticmethod + def from_sparse( + types: Int8Array, + children: NullableCollection[Array], + field_names: list[str] | None = None, + type_codes: Int8Array | list[int] | None = None, + ) -> UnionArray: ... + + +class StringArray(Array[StringScalar]): + @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[LargeStringScalar]): + @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 StringViewArray(Array[StringViewScalar]): + ... + + +class BinaryArray(Array[BinaryScalar]): + @property + def total_values_length(self) -> int: ... + + +class LargeBinaryArray(Array[LargeBinaryScalar]): + @property + def total_values_length(self) -> int: ... + + +class BinaryViewArray(Array[BinaryViewScalar]): + ... + + +class DictionaryArray(Array[DictionaryScalar[_IndexT, _BasicValueT]]): + def dictionary_encode(self) -> Self: ... # type: ignore[override] + def dictionary_decode(self) -> Array[Scalar[_BasicValueT]]: ... + + @property + def indices(self) -> Array[Scalar[_IndexT]]: ... + @property + def dictionary(self) -> Array[Scalar[_BasicValueT]]: ... + + @staticmethod + def from_buffers( # type: ignore[override] + type: _BasicValueT, + length: int, + buffers: list[Buffer], + dictionary: Array | np.ndarray | pd.Series, + null_count: int = -1, + offset: int = 0, + ) -> DictionaryArray[Any, _BasicValueT]: ... + + @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[StructScalar]): + 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[RunEndEncodedScalar[_RunEndType, _BasicValueT]]): + @staticmethod + def from_arrays( + run_ends: Int16Array | Int32Array | Int64Array | list[int], + values: Array | list[Any], type: DataType | None = None, + ) -> RunEndEncodedArray[Any, _BasicValueT]: ... + + @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[Any, _BasicValueT]: ... + + @property + def run_ends(self) -> Array[Scalar[_RunEndType]]: ... + + @property + def values(self) -> Array[Scalar[_BasicValueT]]: ... + + def find_physical_offset(self) -> int: ... + + def find_physical_length(self) -> int: ... + + +_ArrayT = TypeVar("_ArrayT", bound=Array) + + +class ExtensionArray(Array[ExtensionScalar], Generic[_ArrayT]): + @property + def storage(self) -> Any: ... + + @staticmethod + def from_storage(typ: BaseExtensionType, + storage: _ArrayT) -> ExtensionArray[_ArrayT]: ... + + +class JsonArray(ExtensionArray[_ArrayT]): + ... + + +class UuidArray(ExtensionArray[_ArrayT]): + ... + + +class FixedShapeTensorArray(ExtensionArray[_ArrayT]): + 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[_ArrayT]): + ... + + +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[_ArrayT], + memory_pool: MemoryPool | None = None) -> _ArrayT: ... + + +def _empty_array(type: _DataTypeT) -> Array[Scalar[_DataTypeT]]: ... + + +__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", + "_CastAs", +] 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/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/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi new file mode 100644 index 000000000000..6dd61674d40c --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -0,0 +1,686 @@ +# 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, Generic, Literal, TypeVar +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, + SupportArrowArray, + SupportArrowDeviceArray, + SupportArrowStream, +) +from pyarrow.compute import Expression +from pyarrow.interchange.dataframe import _PyArrowDataFrame +from pyarrow.lib import Device, MemoryManager, MemoryPool, Schema +from pyarrow.lib import Field as _Field + +from .array import Array, StructArray, _CastAs, _PandasConvertible +from .device import DeviceAllocationType +from .io import Buffer +from ._ipc import RecordBatchReader +from .scalar import BooleanScalar, Int64Scalar, Scalar, StructScalar +from .tensor import Tensor +from ._stubs_typing import NullableCollection +from ._types import DataType, _AsPyType, _BasicDataType, _DataTypeT + +Field: TypeAlias = _Field[DataType] +_ScalarT = TypeVar("_ScalarT", bound=Scalar) +_Scalar_co = TypeVar("_Scalar_co", bound=Scalar, covariant=True) +ArrayOrChunkedArray: TypeAlias = Array[_Scalar_co] | ChunkedArray[_Scalar_co] + +_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[pd.Series], Generic[_Scalar_co]): + + def as_py(self) -> list[Any]: ... + + @property + def data(self) -> Self: ... + @property + def type(self: ChunkedArray[Scalar[_DataTypeT]]) -> _DataTypeT: ... + + # 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_co | Self: ... + + def getitem(self, i: int) -> Scalar: ... + def is_null(self, *, nan_is_null: bool = False) -> ChunkedArray[BooleanScalar]: ... + + def is_nan(self) -> ChunkedArray[BooleanScalar]: ... + + def is_valid(self) -> ChunkedArray[BooleanScalar]: ... + + def cast( + self, target_type: _CastAs | str | None, safe: bool = True, + options: CastOptions | None = None, + memory_pool: MemoryPool | None = None + ) -> Self | ChunkedArray[Scalar[_CastAs]]: ... + + def fill_null(self, fill_value: Scalar[_DataTypeT] | 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[Any]]: ... + + def combine_chunks(self, memory_pool: MemoryPool | + None = None) -> Array[_Scalar_co]: ... + + def unique(self) -> ChunkedArray[_Scalar_co]: ... + + 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: ChunkedArray[Scalar[_BasicDataType[_AsPyType]]], + value: Scalar[_DataTypeT] | _AsPyType, + 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[_Scalar_co]: ... + + @property + def chunks(self) -> list[Array[_Scalar_co]]: ... + + def iterchunks( + self: ArrayOrChunkedArray[_ScalarT], + ) -> Generator[Array, None, None]: ... + + def __iter__(self) -> Iterator[_Scalar_co]: ... + + def to_pylist( + self: ChunkedArray[Scalar[_BasicDataType[_AsPyType]]], + *, + maps_as_pydicts: Literal["lossy", "strict"] | None = None, + ) -> list[_AsPyType | None]: ... + + 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[Any]] + | Iterable[Iterable[Any] | SupportArrowStream | SupportArrowArray] + | Iterable[Array[_ScalarT]] | Array[_ScalarT] + | SupportArrowArray | SupportArrowStream, + type: DataType | str | None = None, +) -> ChunkedArray[Scalar[Any]] | ChunkedArray[_ScalarT]: ... + + +_ColumnT = TypeVar("_ColumnT", bound=ArrayOrChunkedArray[Any]) + + +class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): + 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) -> _ColumnT | Self: ... + + def __len__(self) -> int: ... + def column(self, i: int | str) -> _ColumnT: ... + + @property + def column_names(self) -> list[str]: ... + + @property + def columns(self) -> list[_ColumnT]: ... + + def drop_null(self) -> Self: ... + + def field(self, i: int | str) -> Field: ... + + @classmethod + def from_pydict( + cls, + mapping: + Mapping[Any, ArrayOrChunkedArray[Any] | 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[_ColumnT, 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[Any] | list[list[Any]]) -> Self: ... + + def append_column( + self, field_: str | Field, column: ArrayOrChunkedArray[Any] | list[list[Any]] + ) -> Self: ... + + +class RecordBatch(_Tabular[Array]): + 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[Any] | 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[StructScalar] + ) -> 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[ChunkedArray[Any]]): + 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[Any] | 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[StructScalar] + ) -> Self: ... + + def to_struct_array( + self, max_chunksize: int | None = None + ) -> ChunkedArray[StructScalar]: ... + + @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[Any] | list[list[Any]]) -> Self: ... + + def remove_column(self, i: int) -> Self: ... + + def set_column(self, i: int, field_: str | Field, + column: ArrayOrChunkedArray[Any] | 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[Any]] + | Collection[Array[Any] | ChunkedArray[Any] | list[Any]] + | pd.DataFrame + | SupportArrowArray + | SupportArrowDeviceArray, + names: list[str] | Schema | None = None, + schema: Schema | None = None, + metadata: Mapping[str | bytes, str | bytes] | None = None, +) -> RecordBatch: ... + + +def table( + data: Collection[ArrayOrChunkedArray[Any] | list[Any] | range | str] + | pd.DataFrame + | SupportArrowArray + | SupportArrowStream + | SupportArrowDeviceArray + | Mapping[str, list[Any] | Array[Any] | ChunkedArray[Any] | 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..ba40c7b299db --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/tensor.pyi @@ -0,0 +1,268 @@ +# 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, 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) -> str: ... + + @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: csr_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) -> str: ... + @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) -> str: ... + @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: tuple[int, ...], + dim_names: list[str] | None = None, + ) -> Self: ... + + @classmethod + def from_scipy(cls, obj: csr_matrix, + dim_names: list[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) -> str: ... + @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: ... + + +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) -> str: ... + @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..9e5a0568db05 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -0,0 +1,227 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import sys +from enum import IntEnum + +from typing import Any + +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.lib 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[Any, Any] + | Decimal64Type[Any, Any] + | Decimal128Type[Any, Any] + | Decimal256Type[Any, Any] +) +_Date: TypeAlias = Date32Type | Date64Type +_Time: TypeAlias = Time32Type[Any] | Time64Type[Any] +_Interval: TypeAlias = MonthDayNanoIntervalType +_Temporal: TypeAlias = (TimestampType[Any, Any] + | DurationType[Any] | _Time | _Date | _Interval) +_Union: TypeAlias = SparseUnionType | DenseUnionType +_Nested: TypeAlias = ( + ListType[Any] + | FixedSizeListType[Any, Any] + | LargeListType[Any] + | ListViewType[Any] + | LargeListViewType[Any] + | StructType + | MapType[Any, Any, Any] + | _Union +) + + +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[Any]]: ... +def is_large_list(t: DataType) -> TypeIs[LargeListType[Any]]: ... +def is_fixed_size_list(t: DataType) -> TypeIs[FixedSizeListType[Any, Any]]: ... +def is_list_view(t: DataType) -> TypeIs[ListViewType[Any]]: ... +def is_large_list_view(t: DataType) -> TypeIs[LargeListViewType[Any]]: ... +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[Any, Any]]: ... +def is_temporal(t: DataType) -> TypeIs[_Temporal]: ... +def is_timestamp(t: DataType) -> TypeIs[TimestampType[Any, Any]]: ... +def is_duration(t: DataType) -> TypeIs[DurationType[Any]]: ... +def is_time(t: DataType) -> TypeIs[_Time]: ... +def is_time32(t: DataType) -> TypeIs[Time32Type[Any]]: ... +def is_time64(t: DataType) -> TypeIs[Time64Type[Any]]: ... +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]: ... +def is_fixed_size_binary(t: DataType) -> TypeIs[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[Any, Any, Any]]: ... +def is_decimal(t: DataType) -> TypeIs[_Decimal]: ... +def is_decimal32(t: DataType) -> TypeIs[Decimal32Type[Any, Any]]: ... +def is_decimal64(t: DataType) -> TypeIs[Decimal64Type[Any, Any]]: ... +def is_decimal128(t: DataType) -> TypeIs[Decimal128Type[Any, Any]]: ... +def is_decimal256(t: DataType) -> TypeIs[Decimal256Type[Any, Any]]: ... +def is_dictionary(t: DataType) -> TypeIs[DictionaryType[Any, Any, Any]]: ... +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): + INTERVAL_MONTHS = 0 + INTERVAL_DAY_TIME = 1 + INTERVAL_MONTH_DAY_NANO = 2 diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index e090b1d0b3db..702f9eff94e7 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -3936,7 +3936,7 @@ cdef class FixedSizeListArray(BaseListArray): Or create from a values array, list size and matching type: >>> typ = pa.list_(pa.field("values", pa.int64()), 2) - >>> arr = pa.FixedSizeListArray.from_arrays(values,type=typ) + >>> arr = pa.FixedSizeListArray.from_arrays(values, type=typ) >>> arr [ diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index f96fb40572b9..ad562a0dbd5c 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -1108,7 +1108,7 @@ cdef class StructScalar(Scalar, Mapping): Parameters ---------- - index : Union[int, str] + key : Union[int, str] Index / position or name of the field. Returns diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 63c0eb0bea4e..d3493aeb462f 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 Int64Type'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.lib.Int64Type) # 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_()) @@ -2868,7 +2875,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])] @@ -2900,7 +2907,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 @@ -2921,8 +2928,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 = [ @@ -3188,7 +3195,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 @@ -3196,7 +3203,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 @@ -3282,7 +3289,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(): @@ -3800,7 +3807,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], @@ -4078,11 +4085,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 @@ -4611,7 +4618,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 01bf3c31a09b..2dbfec5815f4 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(): @@ -538,7 +538,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(): @@ -584,7 +584,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) @@ -696,6 +696,7 @@ def test_empty_table(): assert table.schema == schema +@pytest.mark.numpy @pytest.mark.pandas def test_schema_from_pandas(): import pandas as pd @@ -780,7 +781,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 0fefabf02abe..43b2faa42ede 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) @@ -954,7 +969,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()), @@ -1208,7 +1226,7 @@ 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) @@ -1242,7 +1260,7 @@ 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) @@ -1464,7 +1482,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(): @@ -1625,7 +1643,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 @@ -1690,7 +1709,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) @@ -1711,7 +1730,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) @@ -2004,22 +2023,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"], @@ -2089,13 +2127,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 @@ -2227,7 +2265,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( @@ -2378,7 +2416,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): @@ -2483,7 +2521,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', [ @@ -2500,7 +2538,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())]) @@ -2704,10 +2742,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(): @@ -2841,8 +2879,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(): @@ -3774,7 +3812,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): @@ -4031,7 +4069,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): @@ -4230,7 +4268,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..5b79cba44201 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,12 @@ 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 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 + _ = pa.scalar(d, type=pa.duration('ms')).as_py() == d + _ = pa.scalar(d, type=pa.duration('s')).as_py() == d def test_type_equality_operators(): @@ -1141,11 +1147,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 +1278,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 +1415,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 From 3df0567f12c134b7d661dd7ca1538061ea0ecee1 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Mar 2026 01:07:51 +0100 Subject: [PATCH 22/93] temporary --- python/pyarrow-stubs/pyarrow/array.pyi | 17 ++++++++-------- python/pyarrow-stubs/pyarrow/tensor.pyi | 26 +++++++++++++------------ python/pyarrow/tests/test_array.py | 4 ++-- python/pyarrow/tests/test_table.py | 6 ++++-- python/pyarrow/tests/test_types.py | 4 ++-- 5 files changed, 31 insertions(+), 26 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi index 547e9c949d52..d3a4c2f574b6 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -545,7 +545,8 @@ class ListViewArray(BaseListArray[ListViewScalar[_DataTypeT]]): @classmethod def from_arrays( cls, - offsets: Int32Array, + offsets: Int32Array | list[int] | list[int | None], + sizes: Int32Array | list[int] | list[int | None], values: Array[Scalar[_DataTypeT]] | Array, *, type: _DataTypeT | None = None, @@ -567,7 +568,8 @@ class LargeListViewArray(BaseListArray[LargeListScalar[_DataTypeT]]): @classmethod def from_arrays( cls, - offsets: Int64Array, + offsets: Int64Array | list[int] | list[int | None], + sizes: Int64Array | list[int] | list[int | None], values: Array[Scalar[_DataTypeT]] | Array, *, type: _DataTypeT | None = None, @@ -597,7 +599,7 @@ class FixedSizeListArray(BaseListArray[FixedSizeListScalar[_DataTypeT, _Size]]): ) -> FixedSizeListArray[_DataTypeT, _Size | None]: ... @property - def values(self) -> BaseListArray[ListScalar[_DataTypeT]]: ... + def values(self) -> Array: ... _MapKeyT = TypeVar("_MapKeyT", bound=_BasicDataType) @@ -608,10 +610,9 @@ class MapArray(BaseListArray[MapScalar[_MapKeyT, _MapItemT]]): @classmethod def from_arrays( cls, - offsets: Int64Array | list[int] | None, - keys: Array[Scalar[_MapKeyT]] | np.ndarray | list | None = None, - items: Array[Scalar[_MapItemT]] | np.ndarray | list | None = None, - values: Array | DataType | None = None, + offsets: Int32Array | list[int] | list[int | None], + keys: Array[Scalar[_MapKeyT]] | np.ndarray | list, + items: Array[Scalar[_MapItemT]] | np.ndarray | list, *, type: DataType | None = None, pool: MemoryPool | None = None, @@ -676,7 +677,7 @@ class LargeStringArray(Array[LargeStringScalar]): null_bitmap: Buffer | None = None, null_count: int | None = -1, offset: int | None = 0, - ) -> StringArray: ... + ) -> LargeStringArray: ... class StringViewArray(Array[StringViewScalar]): diff --git a/python/pyarrow-stubs/pyarrow/tensor.pyi b/python/pyarrow-stubs/pyarrow/tensor.pyi index ba40c7b299db..b4bd52c832f0 100644 --- a/python/pyarrow-stubs/pyarrow/tensor.pyi +++ b/python/pyarrow-stubs/pyarrow/tensor.pyi @@ -27,7 +27,7 @@ import numpy as np from pyarrow.lib import _Weakrefable from pyarrow._types import DataType -from scipy.sparse import coo_matrix, csr_matrix +from scipy.sparse import coo_matrix, csc_matrix, csr_matrix from sparse import COO # type: ignore[import-untyped, import-not-found] @@ -55,7 +55,7 @@ class Tensor(_Weakrefable): def ndim(self) -> int: ... @property - def size(self) -> str: ... + def size(self) -> int: ... @property def shape(self) -> tuple[int, ...]: ... @@ -82,7 +82,7 @@ class SparseCOOTensor(_Weakrefable): ) -> Self: ... @classmethod - def from_scipy(cls, obj: csr_matrix, + def from_scipy(cls, obj: coo_matrix, dim_names: Sequence[str] | None = None) -> Self: ... @classmethod @@ -107,7 +107,7 @@ class SparseCOOTensor(_Weakrefable): @property def ndim(self) -> int: ... @property - def size(self) -> str: ... + def size(self) -> int: ... @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @@ -157,7 +157,7 @@ class SparseCSRMatrix(_Weakrefable): @property def ndim(self) -> int: ... @property - def size(self) -> str: ... + def size(self) -> int: ... @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @@ -181,20 +181,20 @@ class SparseCSCMatrix(_Weakrefable): data: np.ndarray, indptr: np.ndarray, indices: np.ndarray, - shape: tuple[int, ...], - dim_names: list[str] | None = None, + shape: Sequence[int], + dim_names: Sequence[str] | None = None, ) -> Self: ... @classmethod - def from_scipy(cls, obj: csr_matrix, - dim_names: list[str] | None = None) -> Self: ... + 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) -> csr_matrix: ... + def to_scipy(self) -> csc_matrix: ... def to_tensor(self) -> Tensor: ... @@ -205,7 +205,7 @@ class SparseCSCMatrix(_Weakrefable): @property def ndim(self) -> int: ... @property - def size(self) -> str: ... + def size(self) -> int: ... @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @@ -214,6 +214,8 @@ class SparseCSCMatrix(_Weakrefable): def dim_names(self) -> list[str]: ... @property def non_zero_length(self) -> int: ... + @property + def type(self) -> DataType: ... class SparseCSFTensor(_Weakrefable): @@ -246,7 +248,7 @@ class SparseCSFTensor(_Weakrefable): @property def ndim(self) -> int: ... @property - def size(self) -> str: ... + def size(self) -> int: ... @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index d3493aeb462f..71aa27a7cbb5 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -327,11 +327,11 @@ def test_asarray(): arr = pa.array(range(4)) - # The iterator interface gives back an array of Int64Type'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.Int64Type) + assert isinstance(np_arr[0], pa.Int64Scalar) # Calling with the arrow array gives back an array with 'int64' dtype np_arr = np.asarray(arr) diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index 43b2faa42ede..e28f00862035 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -1226,7 +1226,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") # type: ignore[no-matching-overload] + 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) @@ -1260,7 +1261,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") # type: ignore[no-matching-overload] + 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) diff --git a/python/pyarrow/tests/test_types.py b/python/pyarrow/tests/test_types.py index 5b79cba44201..8d79674f7f46 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -1108,8 +1108,8 @@ def test_timedelta_overflow(): # second/millisecond resolution, not overflow 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 + assert pa.scalar(d, type=pa.duration('ms')).as_py() == d + assert pa.scalar(d, type=pa.duration('s')).as_py() == d def test_type_equality_operators(): From c990f327743acc29c0515d4d852d2abfe49dbe74 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 10:26:07 -0700 Subject: [PATCH 23/93] post rebase --- python/pyarrow-stubs/pyarrow/_types.pyi | 2 ++ python/pyarrow-stubs/pyarrow/array.pyi | 14 +++++++------- python/pyarrow-stubs/pyarrow/table.pyi | 20 ++++++++++---------- python/pyarrow-stubs/pyarrow/types.pyi | 2 +- 4 files changed, 20 insertions(+), 18 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index 941caa761446..d810fb81ce83 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -210,6 +210,7 @@ _IndexT_co = TypeVar( ) _ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) _ValueT_co = TypeVar("_ValueT_co", bound=DataType, default=DataType, covariant=True) +_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType[Any]) _K = TypeVar("_K", bound=DataType, default=DataType) _K_co = TypeVar("_K_co", bound=DataType, default=DataType, covariant=True) _Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) @@ -765,4 +766,5 @@ __all__ = [ "_Time32Unit", "_Time64Unit", "_DataTypeT", + "_BasicValueT", ] diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi index d3a4c2f574b6..0ffeda8ae5ba 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -40,9 +40,9 @@ from pyarrow._stubs_typing import ( Indices, Mask, Order, - SupportArrowArray, - SupportArrowDeviceArray, - SupportPyArrowArray, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsPyArrowArray, ) from pyarrow.lib import ( Buffer, @@ -122,8 +122,8 @@ from ._stubs_typing import NullableCollection def array( - values: NullableCollection[Any] | Iterable[Any] | SupportArrowArray - | SupportArrowDeviceArray | SupportPyArrowArray, + values: NullableCollection[Any] | Iterable[Any] | SupportsArrowArray + | SupportsArrowDeviceArray | SupportsPyArrowArray, type: Any | None = None, mask: Mask | pd.Series[bool] | None = None, size: int | None = None, @@ -134,8 +134,8 @@ def array( def asarray( - values: NullableCollection[Any] | Iterable[Any] | SupportArrowArray - | SupportArrowDeviceArray, + values: NullableCollection[Any] | Iterable[Any] | SupportsArrowArray + | SupportsArrowDeviceArray, type: _DataTypeT | Any | None = None, ) -> Array[Scalar[_DataTypeT]] | ArrayLike: ... diff --git a/python/pyarrow-stubs/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi index 6dd61674d40c..5e6a9f640591 100644 --- a/python/pyarrow-stubs/pyarrow/table.pyi +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -48,9 +48,9 @@ from pyarrow._stubs_typing import ( NullEncoding, NullSelectionBehavior, Order, - SupportArrowArray, - SupportArrowDeviceArray, - SupportArrowStream, + SupportsArrowArray, + SupportsArrowDeviceArray, + SupportsArrowStream, ) from pyarrow.compute import Expression from pyarrow.interchange.dataframe import _PyArrowDataFrame @@ -255,9 +255,9 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def chunked_array( arrays: Iterable[NullableCollection[Any]] - | Iterable[Iterable[Any] | SupportArrowStream | SupportArrowArray] + | Iterable[Iterable[Any] | SupportsArrowStream | SupportsArrowArray] | Iterable[Array[_ScalarT]] | Array[_ScalarT] - | SupportArrowArray | SupportArrowStream, + | SupportsArrowArray | SupportsArrowStream, type: DataType | str | None = None, ) -> ChunkedArray[Scalar[Any]] | ChunkedArray[_ScalarT]: ... @@ -613,8 +613,8 @@ def record_batch( data: Mapping[str, list[Any] | Array[Any]] | Collection[Array[Any] | ChunkedArray[Any] | list[Any]] | pd.DataFrame - | SupportArrowArray - | SupportArrowDeviceArray, + | SupportsArrowArray + | SupportsArrowDeviceArray, names: list[str] | Schema | None = None, schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, @@ -624,9 +624,9 @@ def record_batch( def table( data: Collection[ArrayOrChunkedArray[Any] | list[Any] | range | str] | pd.DataFrame - | SupportArrowArray - | SupportArrowStream - | SupportArrowDeviceArray + | SupportsArrowArray + | SupportsArrowStream + | SupportsArrowDeviceArray | Mapping[str, list[Any] | Array[Any] | ChunkedArray[Any] | range] | Mapping[str, Any], names: list[str] | Schema | None = None, diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index 9e5a0568db05..99d1fbfe084d 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -31,7 +31,7 @@ else: import pyarrow.lib as lib -from pyarrow.lib import ( +from pyarrow._types import ( BinaryType, BinaryViewType, BoolType, From cffdf249d511cd353d0da2c9d5eee8dd4d376044 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 11:43:30 +0100 Subject: [PATCH 24/93] Use plain classes and unions instead of generics in stubs Drop the TypeVar/Generic parameters from the DataType subclasses, Field, Array, ChunkedArray and the tabular classes. Where a property can only return one of a few types (dictionary index type, run end type, time units) it now returns a union or Literal instead. Most of these parameters weren't buying much yet: scalar.pyi and lib.pyi are still placeholders, so e.g. Array[Scalar[...]] collapsed to Any anyway, and several annotations were being flagged by mypy and pyright (TypeVars used only once, type arguments on subclasses that weren't actually generic). Plain classes are a lot easier to read and review, and we can bring parameters back later where they clearly pay off. This also removes the stub-only _BasicDataType and _HasPrecisionScale helpers, which have no runtime counterpart, plus the IntoArray and SingleOrList aliases. --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 43 +- python/pyarrow-stubs/pyarrow/_types.pyi | 384 +++++++----------- python/pyarrow-stubs/pyarrow/array.pyi | 294 +++++--------- python/pyarrow-stubs/pyarrow/table.pyi | 109 +++-- python/pyarrow-stubs/pyarrow/types.pyi | 56 ++- 5 files changed, 346 insertions(+), 540 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 7802f0330597..89c36cba98a0 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -25,8 +25,7 @@ from _typeshed import ( ) from collections.abc import Collection, Container, Iterator, Sequence, Sized from decimal import Decimal -from typing import Any, Literal, Protocol, TypeAlias, TypeVar -from typing_extensions import TypeAliasType +from typing import Any, Literal, Protocol, TypeAlias import numpy as np @@ -55,16 +54,6 @@ NullEncoding: TypeAlias = Literal["mask", "encode"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] -IntegerType: TypeAlias = ( - lib.Int8Type - | lib.Int16Type - | lib.Int32Type - | lib.Int64Type - | lib.UInt8Type - | lib.UInt16Type - | lib.UInt32Type - | lib.UInt64Type -) PyScalar: TypeAlias = ( bool | int @@ -79,25 +68,15 @@ PyScalar: TypeAlias = ( ) NumpyScalar: TypeAlias = np.generic[Any] -_PyScalarT_co = TypeVar("_PyScalarT_co", bound=PyScalar, covariant=True) -_NumpyScalarT_co = TypeVar("_NumpyScalarT_co", bound=NumpyScalar, covariant=True) -_DataTypeT_co = TypeVar("_DataTypeT_co", bound=lib.DataType, covariant=True) -IntoArray = TypeAliasType( - "IntoArray", - Sequence[_PyScalarT_co | None] - | NDArray[_NumpyScalarT_co] - | lib.Array[lib.Scalar[_DataTypeT_co]] - | ChunkedArray[Any], - type_params=(_PyScalarT_co, _NumpyScalarT_co, _DataTypeT_co), +Mask: TypeAlias = ( + Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray +) +Indices: TypeAlias = ( + Sequence[int | None] + | NDArray[np.integer[Any]] + | lib.IntegerArray + | ChunkedArray ) - -Mask: TypeAlias = IntoArray[bool, np.bool_, lib.BoolType] -Indices: TypeAlias = IntoArray[int, np.integer[Any], IntegerType] - -_T = TypeVar("_T") -_V = TypeVar("_V", covariant=True) - -SingleOrList: TypeAlias = list[_T] | _T class SupportsDunderEQ(Protocol): def __eq__(self, other: object, /) -> bool: ... @@ -132,7 +111,7 @@ class SupportsArrowDeviceArray(Protocol): class SupportsArrowSchema(Protocol): def __arrow_c_schema__(self) -> Any: ... -class NullableCollection(Sized, Container[Any], Protocol[_V]): - def __iter__(self) -> Iterator[_V] | Iterator[_V | None]: ... +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 index d810fb81ce83..a46f72394cd4 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -15,15 +15,12 @@ # specific language governing permissions and limitations # under the License. -import datetime as dt # noqa: F401 - from collections.abc import Iterable, Iterator, Mapping, Sequence -from decimal import Decimal # noqa: F401 -from typing import Any, Generic, Literal, Protocol, TypeAlias, overload +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, TypeVar, deprecated +from typing_extensions import Self, deprecated from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit from pyarrow.lib import ( # noqa: F401 @@ -33,7 +30,6 @@ from pyarrow.lib import ( # noqa: F401 ExtensionArray, ExtensionScalar, MemoryPool, - MonthDayNano, Table, ) @@ -41,7 +37,7 @@ class _Weakrefable: ... class _Metadata(_Weakrefable): ... class DataType(_Weakrefable): - def field(self, i: int) -> Field[Any]: ... + def field(self, i: int) -> Field: ... @property def id(self) -> int: ... @property @@ -66,190 +62,157 @@ class DataType(_Weakrefable): @classmethod def _import_from_c_capsule(cls, schema: Any) -> Self: ... -_AsPyType = TypeVar("_AsPyType") -_DataTypeT = TypeVar("_DataTypeT", bound=DataType) -_DataTypeT_co = TypeVar("_DataTypeT_co", bound=DataType, covariant=True) -_Size = TypeVar("_Size", bound=int, default=int) _PandasDtype: TypeAlias = ( type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype ) -class _BasicDataType(DataType, Generic[_AsPyType]): ... -class NullType(_BasicDataType[None]): ... -class BoolType(_BasicDataType[bool]): ... -class UInt8Type(_BasicDataType[int]): ... -class Int8Type(_BasicDataType[int]): ... -class UInt16Type(_BasicDataType[int]): ... -class Int16Type(_BasicDataType[int]): ... -class UInt32Type(_BasicDataType[int]): ... -class Int32Type(_BasicDataType[int]): ... -class UInt64Type(_BasicDataType[int]): ... -class Int64Type(_BasicDataType[int]): ... -class Float16Type(_BasicDataType[float]): ... -class Float32Type(_BasicDataType[float]): ... -class Float64Type(_BasicDataType[float]): ... -class Date32Type(_BasicDataType[dt.date]): ... -class Date64Type(_BasicDataType[dt.date]): ... -class MonthDayNanoIntervalType(_BasicDataType[MonthDayNano]): ... -class StringType(_BasicDataType[str]): ... -class LargeStringType(_BasicDataType[str]): ... -class StringViewType(_BasicDataType[str]): ... -class BinaryType(_BasicDataType[bytes]): ... -class LargeBinaryType(_BasicDataType[bytes]): ... -class BinaryViewType(_BasicDataType[bytes]): ... - -_Unit = TypeVar("_Unit", bound=TimeUnit, default=Literal["us"]) -_Tz = TypeVar("_Tz", str, None, default=None) - -class TimestampType(_BasicDataType[dt.datetime], Generic[_Unit, _Tz]): - @property - def unit(self) -> _Unit: ... - @property - def tz(self) -> _Tz: ... - -_Time32Unit = TypeVar("_Time32Unit", bound=Literal["s", "ms"]) - -class Time32Type(_BasicDataType[dt.time], Generic[_Time32Unit]): +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: ... -_Time64Unit = TypeVar("_Time64Unit", bound=Literal["us", "ns"]) - -class Time64Type(_BasicDataType[dt.time], Generic[_Time64Unit]): +class Time64Type(DataType): @property def unit(self) -> _Time64Unit: ... -class DurationType(_BasicDataType[dt.timedelta], Generic[_Unit]): +class DurationType(DataType): @property - def unit(self) -> _Unit: ... + def unit(self) -> TimeUnit: ... -_FixedSizeBinaryAsPyType = TypeVar("_FixedSizeBinaryAsPyType", default=bytes) +class FixedSizeBinaryType(DataType): ... -class FixedSizeBinaryType(_BasicDataType[_FixedSizeBinaryAsPyType]): ... +class Decimal32Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... -_Precision = TypeVar("_Precision", bound=int, default=int) -_Scale = TypeVar("_Scale", bound=int, default=int) -_Precision_co = TypeVar("_Precision_co", bound=int, default=int, covariant=True) -_Scale_co = TypeVar("_Scale_co", bound=int, default=int, covariant=True) +class Decimal64Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... -class _HasPrecisionScale(Protocol[_Precision_co, _Scale_co]): +class Decimal128Type(FixedSizeBinaryType): @property - def precision(self) -> _Precision_co: ... + def precision(self) -> int: ... @property - def scale(self) -> _Scale_co: ... + def scale(self) -> int: ... -class Decimal32Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... -class Decimal64Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... -class Decimal128Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... -class Decimal256Type( - FixedSizeBinaryType[Decimal], _HasPrecisionScale[_Precision_co, _Scale_co] -): ... +class Decimal256Type(FixedSizeBinaryType): + @property + def precision(self) -> int: ... + @property + def scale(self) -> int: ... -class ListType(DataType, Generic[_DataTypeT_co]): +class ListType(DataType): @property - def value_field(self) -> Field[_DataTypeT_co]: ... + def value_field(self) -> Field: ... @property - def value_type(self) -> _DataTypeT_co: ... + def value_type(self) -> DataType: ... -class LargeListType(DataType, Generic[_DataTypeT_co]): +class LargeListType(DataType): @property - def value_field(self) -> Field[_DataTypeT_co]: ... + def value_field(self) -> Field: ... @property - def value_type(self) -> _DataTypeT_co: ... + def value_type(self) -> DataType: ... -class ListViewType(DataType, Generic[_DataTypeT_co]): +class ListViewType(DataType): @property - def value_field(self) -> Field[_DataTypeT_co]: ... + def value_field(self) -> Field: ... @property - def value_type(self) -> _DataTypeT_co: ... + def value_type(self) -> DataType: ... -class LargeListViewType(DataType, Generic[_DataTypeT_co]): +class LargeListViewType(DataType): @property - def value_field(self) -> Field[_DataTypeT_co]: ... + def value_field(self) -> Field: ... @property - def value_type(self) -> _DataTypeT_co: ... + def value_type(self) -> DataType: ... -class FixedSizeListType(DataType, Generic[_DataTypeT_co, _Size]): +class FixedSizeListType(DataType): @property - def value_field(self) -> Field[_DataTypeT_co]: ... + def value_field(self) -> Field: ... @property - def value_type(self) -> _DataTypeT_co: ... + def value_type(self) -> DataType: ... @property - def list_size(self) -> _Size: ... + def list_size(self) -> int: ... class DictionaryMemo(_Weakrefable): ... -_IndexT = TypeVar( - "_IndexT", - UInt8Type, - Int8Type, - UInt16Type, - Int16Type, - UInt32Type, - Int32Type, - UInt64Type, - Int64Type, -) -_IndexT_co = TypeVar( - "_IndexT_co", - UInt8Type, - Int8Type, - UInt16Type, - Int16Type, - UInt32Type, - Int32Type, - UInt64Type, - Int64Type, - covariant=True, -) -_ValueT = TypeVar("_ValueT", bound=DataType, default=DataType) -_ValueT_co = TypeVar("_ValueT_co", bound=DataType, default=DataType, covariant=True) -_BasicValueT = TypeVar("_BasicValueT", bound=_BasicDataType[Any]) -_K = TypeVar("_K", bound=DataType, default=DataType) -_K_co = TypeVar("_K_co", bound=DataType, default=DataType, covariant=True) -_Ordered = TypeVar("_Ordered", Literal[True], Literal[False], default=Literal[False]) -_Ordered_co = TypeVar( - "_Ordered_co", Literal[True], Literal[False], default=Literal[False], covariant=True +_IndexType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type ) -class DictionaryType(DataType, Generic[_IndexT_co, _ValueT_co, _Ordered_co]): +class DictionaryType(DataType): @property - def ordered(self) -> _Ordered_co: ... + def ordered(self) -> bool: ... @property - def index_type(self) -> _IndexT_co: ... + def index_type(self) -> _IndexType: ... @property - def value_type(self) -> _ValueT_co: ... + def value_type(self) -> DataType: ... -class MapType(DataType, Generic[_K_co, _ValueT_co, _Ordered_co]): +class MapType(DataType): @property - def key_field(self) -> Field[_K_co]: ... + def key_field(self) -> Field: ... @property - def key_type(self) -> _K_co: ... + def key_type(self) -> DataType: ... @property - def item_field(self) -> Field[_ValueT_co]: ... + def item_field(self) -> Field: ... @property - def item_type(self) -> _ValueT_co: ... + def item_type(self) -> DataType: ... @property - def keys_sorted(self) -> _Ordered_co: ... + 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[Any]: ... + 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[Any]]: ... + def __iter__(self) -> Iterator[Field]: ... __getitem__ = field @property def names(self) -> list[str]: ... @property - def fields(self) -> list[Field[Any]]: ... + def fields(self) -> list[Field]: ... class UnionType(DataType): @property @@ -257,8 +220,8 @@ class UnionType(DataType): @property def type_codes(self) -> list[int]: ... def __len__(self) -> int: ... - def __iter__(self) -> Iterator[Field[Any]]: ... - def field(self, i: int) -> Field[Any]: ... + def __iter__(self) -> Iterator[Field]: ... + def field(self, i: int) -> Field: ... __getitem__ = field @@ -270,27 +233,24 @@ class DenseUnionType(UnionType): @property def mode(self) -> Literal["dense"]: ... -_RunEndType = TypeVar("_RunEndType", Int16Type, Int32Type, Int64Type) -_RunEndType_co = TypeVar( - "_RunEndType_co", Int16Type, Int32Type, Int64Type, covariant=True -) +_RunEndType: TypeAlias = Int16Type | Int32Type | Int64Type -class RunEndEncodedType(DataType, Generic[_RunEndType_co, _ValueT_co]): +class RunEndEncodedType(DataType): @property - def run_end_type(self) -> _RunEndType_co: ... + def run_end_type(self) -> _RunEndType: ... @property - def value_type(self) -> _ValueT_co: ... - -_StorageT = TypeVar("_StorageT", bound=Array[Any] | ChunkedArray[Any]) + def value_type(self) -> DataType: ... class BaseExtensionType(DataType): - def __arrow_ext_class__(self) -> type[ExtensionArray[Any]]: ... + 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: _StorageT) -> _StorageT: ... + def wrap_array( + self, storage: Array | ChunkedArray + ) -> ExtensionArray | ChunkedArray: ... class ExtensionType(BaseExtensionType): def __init__(self, storage_type: DataType, extension_name: str) -> None: ... @@ -300,9 +260,9 @@ class ExtensionType(BaseExtensionType): cls, storage_type: DataType, serialized: bytes ) -> Self: ... -class FixedShapeTensorType(BaseExtensionType, Generic[_ValueT_co]): +class FixedShapeTensorType(BaseExtensionType): @property - def value_type(self) -> _ValueT_co: ... + def value_type(self) -> DataType: ... @property def shape(self) -> list[int]: ... @property @@ -412,8 +372,8 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def get_all(self, key: str | bytes) -> list[bytes]: ... def to_dict(self) -> dict[bytes, bytes]: ... -class Field(_Weakrefable, Generic[_DataTypeT_co]): - def equals(self, other: Field[Any], check_metadata: bool = False) -> bool: ... +class Field(_Weakrefable): + def equals(self, other: Field, check_metadata: bool = False) -> bool: ... def __hash__(self) -> int: ... @property def nullable(self) -> bool: ... @@ -422,16 +382,16 @@ class Field(_Weakrefable, Generic[_DataTypeT_co]): @property def metadata(self) -> dict[bytes, bytes] | None: ... @property - def type(self) -> _DataTypeT_co: ... + def type(self) -> DataType: ... def with_metadata( self, metadata: _KeyValueMetadataInput | KeyValueMetadata, ) -> Self: ... def remove_metadata(self) -> Self: ... - def with_type(self, new_type: _DataTypeT) -> Field[_DataTypeT]: ... + def with_type(self, new_type: DataType) -> Field: ... def with_name(self, name: str | bytes) -> Self: ... - def with_nullable(self, nullable: bool) -> Field[_DataTypeT_co]: ... - def flatten(self) -> list[Field[Any]]: ... + 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: ... @@ -450,16 +410,16 @@ _FieldTuple: TypeAlias = ( ] ) _StructFieldsInput: TypeAlias = ( - Iterable[Field[Any] | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] + Iterable[Field | _FieldTuple] | Mapping[_FieldName, _DataTypeLike] ) _SchemaFieldsInput: TypeAlias = _StructFieldsInput class Schema(_Weakrefable): def __len__(self) -> int: ... - def __getitem__(self, key: int) -> Field[Any]: ... + def __getitem__(self, key: int) -> Field: ... _field = __getitem__ - def __iter__(self) -> Iterator[Field[Any]]: ... + def __iter__(self) -> Iterator[Field]: ... def __hash__(self) -> int: ... def __sizeof__(self) -> int: ... @property @@ -476,15 +436,15 @@ class Schema(_Weakrefable): def from_pandas( cls, df: pd.DataFrame, preserve_index: bool | None = None ) -> Schema: ... - def field(self, i: int | str | bytes) -> Field[Any]: ... + def field(self, i: int | str | bytes) -> Field: ... @deprecated("Use 'field' instead") - def field_by_name(self, name: str | bytes) -> Field[Any] | None: ... + 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[Any]) -> Schema: ... - def insert(self, i: int, field: Field[Any]) -> Schema: ... + 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[Any]) -> Schema: ... + def set(self, i: int, field: Field) -> Schema: ... @deprecated("Use 'with_metadata' instead") def add_metadata( self, metadata: _KeyValueMetadataInput | KeyValueMetadata @@ -519,21 +479,14 @@ def field( type: None = None, nullable: bool | None = None, metadata: _KeyValueMetadataInput | KeyValueMetadata = None, -) -> Field[Any]: ... +) -> Field: ... @overload def field( name: _FieldName, - type: _DataTypeT, + type: _DataTypeLike, nullable: bool | None = None, metadata: _KeyValueMetadataInput | KeyValueMetadata = None, -) -> Field[_DataTypeT]: ... -@overload -def field( - name: _FieldName, - type: _DataTypeAliasInput, - nullable: bool | None = None, - metadata: _KeyValueMetadataInput | KeyValueMetadata = None, -) -> Field[Any]: ... +) -> Field: ... def null() -> NullType: ... def bool_() -> BoolType: ... def uint8() -> UInt8Type: ... @@ -544,37 +497,27 @@ def uint32() -> UInt32Type: ... def int32() -> Int32Type: ... def int64() -> Int64Type: ... def uint64() -> UInt64Type: ... -def timestamp( - unit: _Unit | str, tz: _Tz | None = None -) -> TimestampType[_Unit, _Tz]: ... -def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ... -def time64(unit: _Time64Unit | str) -> Time64Type[_Time64Unit]: ... -def duration(unit: _Unit | str) -> DurationType[_Unit]: ... +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: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal32Type[_Precision, _Scale | Literal[0]]: ... -def decimal64( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal64Type[_Precision, _Scale | Literal[0]]: ... -def decimal128( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal128Type[_Precision, _Scale | Literal[0]]: ... -def decimal256( - precision: _Precision, scale: _Scale | Literal[0] = 0 -) -> Decimal256Type[_Precision, _Scale | Literal[0]]: ... +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[bytes]: ... +) -> BinaryType | FixedSizeBinaryType: ... def large_binary() -> LargeBinaryType: ... def large_string() -> LargeStringType: ... @@ -583,60 +526,52 @@ large_utf8 = large_string def binary_view() -> BinaryViewType: ... def string_view() -> StringViewType: ... @overload -def list_(value_type: _DataTypeT | Field[_DataTypeT]) -> ListType[_DataTypeT]: ... -@overload def list_( - value_type: _DataTypeT | Field[_DataTypeT], list_size: Literal[-1] -) -> ListType[_DataTypeT]: ... + value_type: DataType | Field, list_size: Literal[-1] = -1 +) -> ListType: ... @overload def list_( - value_type: _DataTypeT | Field[_DataTypeT], list_size: _Size -) -> ListType[_DataTypeT] | FixedSizeListType[_DataTypeT, _Size]: ... -def large_list( - value_type: _DataTypeT | Field[_DataTypeT], -) -> LargeListType[_DataTypeT]: ... -def list_view( - value_type: _DataTypeT | Field[_DataTypeT], -) -> ListViewType[_DataTypeT]: ... -def large_list_view( - value_type: _DataTypeT | Field[_DataTypeT], -) -> LargeListViewType[_DataTypeT]: ... + 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: _K | Field[_K] | _DataTypeAliasInput, - item_type: _ValueT | Field[_ValueT] | _DataTypeAliasInput, - keys_sorted: _Ordered | None = None, -) -> MapType[_K, _ValueT, _Ordered]: ... + key_type: _DataTypeLike | Field, + item_type: _DataTypeLike | Field, + keys_sorted: bool | None = None, +) -> MapType: ... def dictionary( - index_type: _IndexT | _DataTypeAliasInput, - value_type: _ValueT | _DataTypeAliasInput, - ordered: _Ordered | None = None, -) -> DictionaryType[_IndexT, _ValueT, _Ordered]: ... + index_type: _IndexType | _DataTypeAliasInput, + value_type: _DataTypeLike, + ordered: bool | None = None, +) -> DictionaryType: ... def struct( fields: _StructFieldsInput, ) -> StructType: ... def sparse_union( - child_fields: Iterable[Field[Any]], type_codes: Sequence[int] | None = None + child_fields: Iterable[Field], type_codes: Sequence[int] | None = None ) -> SparseUnionType: ... def dense_union( - child_fields: Iterable[Field[Any]], type_codes: Sequence[int] | None = None + child_fields: Iterable[Field], type_codes: Sequence[int] | None = None ) -> DenseUnionType: ... def union( - child_fields: Iterable[Field[Any]], + 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: _ValueT | _DataTypeAliasInput, -) -> RunEndEncodedType[_RunEndType, _ValueT]: ... + value_type: _DataTypeLike, +) -> RunEndEncodedType: ... def json_(storage_type: DataType = ...) -> JsonType: ... def uuid() -> UuidType: ... def fixed_shape_tensor( - value_type: _ValueT, + value_type: DataType, shape: Sequence[int], dim_names: Sequence[str] | None = None, permutation: Sequence[int] | None = None, -) -> FixedShapeTensorType[_ValueT]: ... +) -> FixedShapeTensorType: ... def bool8() -> Bool8Type: ... def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueType: ... def type_for_alias(name: str) -> DataType: ... @@ -650,7 +585,6 @@ __all__ = [ "_Weakrefable", "_Metadata", "DataType", - "_BasicDataType", "NullType", "BoolType", "UInt8Type", @@ -761,10 +695,4 @@ __all__ = [ "type_for_alias", "schema", "from_numpy_dtype", - "_Unit", - "_Tz", - "_Time32Unit", - "_Time64Unit", - "_DataTypeT", - "_BasicValueT", ] diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi index 0ffeda8ae5ba..4c06b0c9c520 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -24,12 +24,7 @@ if sys.version_info >= (3, 11): else: from typing_extensions import Self -from typing import ( - Any, - Generic, - Literal, - TypeVar, -) +from typing import Any, Literal import numpy as np import pandas as pd @@ -55,74 +50,14 @@ from pyarrow.lib import ( from typing_extensions import deprecated import builtins -from .scalar import ( # noqa: F401 - BinaryScalar, - BinaryViewScalar, - BooleanScalar, - Date32Scalar, - Date64Scalar, - DictionaryScalar, - DoubleScalar, - DurationScalar, - ExtensionScalar, - FixedSizeBinaryScalar, - FixedSizeListScalar, - FloatScalar, - HalfFloatScalar, - Int16Scalar, - Int32Scalar, - Int64Scalar, - Int8Scalar, - LargeBinaryScalar, - LargeListScalar, - LargeStringScalar, - ListScalar, - ListViewScalar, - MapScalar, - MonthDayNanoIntervalScalar, - NullScalar, - RunEndEncodedScalar, - Scalar, - StringScalar, - StringViewScalar, - StructScalar, - Time32Scalar, - Time64Scalar, - TimestampScalar, - UInt16Scalar, - UInt32Scalar, - UInt64Scalar, - UInt8Scalar, - UnionScalar, -) +from .scalar import Int64Scalar, Scalar from .device import DeviceAllocationType -from ._types import ( # noqa: F401 - BaseExtensionType, - BinaryType, - DataType, - Field, - Float64Type, - Int64Type, - MapType, - StringType, - StructType, - _AsPyType, - _BasicDataType, - _BasicValueT, - _DataTypeT, - _IndexT, - _RunEndType, - _Size, - _Time32Unit, - _Time64Unit, - _Tz, - _Unit, -) +from ._types import BaseExtensionType, DataType, Field, StructType from ._stubs_typing import NullableCollection def array( - values: NullableCollection[Any] | Iterable[Any] | SupportsArrowArray + values: NullableCollection | Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray | SupportsPyArrowArray, type: Any | None = None, mask: Mask | pd.Series[bool] | None = None, @@ -134,10 +69,10 @@ def array( def asarray( - values: NullableCollection[Any] | Iterable[Any] | SupportsArrowArray + values: NullableCollection | Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, - type: _DataTypeT | Any | None = None, -) -> Array[Scalar[_DataTypeT]] | ArrayLike: ... + type: Any | None = None, +) -> ArrayLike: ... def nulls( @@ -184,10 +119,7 @@ class ArrayStatistics(_Weakrefable): def is_max_exact(self) -> bool | None: ... -_ConvertAs = TypeVar("_ConvertAs", pd.DataFrame, pd.Series) - - -class _PandasConvertible(_Weakrefable, Generic[_ConvertAs]): +class _PandasConvertible(_Weakrefable): def to_pandas( self, memory_pool: MemoryPool | None = None, @@ -206,15 +138,10 @@ class _PandasConvertible(_Weakrefable, Generic[_ConvertAs]): maps_as_pydicts: Literal["None", "lossy", "strict"] | None = None, types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None coerce_temporal_nanoseconds: bool = False, - ) -> _ConvertAs: ... - - -_CastAs = TypeVar("_CastAs", bound=DataType) -_Scalar_co = TypeVar("_Scalar_co", bound=Scalar, covariant=True) -_ScalarT = TypeVar("_ScalarT", bound=Scalar) + ) -> pd.Series | pd.DataFrame: ... -class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): +class Array(_PandasConvertible): def as_py(self) -> list[Any]: ... def diff(self, other: Self) -> str: ... @@ -224,18 +151,18 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def cast( self, - target_type: _CastAs | str, + target_type: DataType | str, safe: bool = True, options: CastOptions | None = None, memory_pool: MemoryPool | None = None, - ) -> Array[Scalar[_CastAs]]: ... + ) -> Array: ... - def view(self, target_type: _CastAs) -> Array[Scalar[_CastAs]]: ... + def view(self, target_type: DataType) -> Array: ... - def sum(self, **kwargs) -> _Scalar_co: ... + def sum(self, **kwargs) -> Scalar: ... @property - def type(self: Array[Scalar[_DataTypeT]]) -> _DataTypeT: ... + def type(self) -> DataType: ... def unique(self) -> Self: ... def dictionary_encode(self, null_encoding: str = "mask") -> DictionaryArray: ... @@ -247,20 +174,20 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): obj: pd.Series | np.ndarray | ArrayLike, *, mask: Mask | None = None, - type: _DataTypeT | None = None, + type: DataType | None = None, safe: bool = True, memory_pool: MemoryPool | None = None, - ) -> Array[Scalar[_DataTypeT]] | Array[Scalar]: ... + ) -> Array: ... @staticmethod def from_buffers( - type: _DataTypeT, + type: DataType, length: int, buffers: Sequence[Buffer | None], null_count: int = -1, offset=0, - children: NullableCollection[Array[Scalar[_DataTypeT]]] | None = None, - ) -> Array[Scalar[_DataTypeT]]: ... + children: NullableCollection | None = None, + ) -> Array: ... @property def null_count(self) -> int: ... @@ -270,7 +197,7 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def get_total_buffer_size(self) -> int: ... def __sizeof__(self) -> int: ... - def __iter__(self) -> Iterator[_Scalar_co]: ... + def __iter__(self) -> Iterator[Scalar]: ... def to_string( self, @@ -293,11 +220,9 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def is_valid(self) -> BooleanArray: ... - def fill_null( - self: Array[Scalar[_BasicDataType[_AsPyType]]], fill_value: _AsPyType - ) -> Array[Scalar[_BasicDataType[_AsPyType]]]: ... + def fill_null(self, fill_value: Any) -> Self: ... - def __getitem__(self, key: int | builtins.slice) -> _Scalar_co | Self: ... + def __getitem__(self, key: int | builtins.slice) -> Scalar | Self: ... def slice(self, offset: int = 0, length: int | None = None) -> Self: ... @@ -313,8 +238,8 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): ) -> Self: ... def index( - self: Array[_ScalarT] | Array[Scalar[_BasicDataType[_AsPyType]]], - value: _ScalarT | _AsPyType, + self, + value: Any, start: int | None = None, end: int | None = None, *, @@ -377,102 +302,102 @@ class Array(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def statistics(self) -> ArrayStatistics | None: ... -class NullArray(Array[NullScalar]): +class NullArray(Array): ... -class BooleanArray(Array[BooleanScalar]): +class BooleanArray(Array): @property def false_count(self) -> int: ... @property def true_count(self) -> int: ... -class NumericArray(Array[_ScalarT]): +class NumericArray(Array): ... -class IntegerArray(NumericArray[_ScalarT]): +class IntegerArray(NumericArray): ... -class FloatingPointArray(NumericArray[_ScalarT]): +class FloatingPointArray(NumericArray): ... -class Int8Array(IntegerArray[Int8Scalar]): +class Int8Array(IntegerArray): ... -class UInt8Array(IntegerArray[UInt8Scalar]): +class UInt8Array(IntegerArray): ... -class Int16Array(IntegerArray[Int16Scalar]): +class Int16Array(IntegerArray): ... -class UInt16Array(IntegerArray[UInt16Scalar]): +class UInt16Array(IntegerArray): ... -class Int32Array(IntegerArray[Int32Scalar]): +class Int32Array(IntegerArray): ... -class UInt32Array(IntegerArray[UInt32Scalar]): +class UInt32Array(IntegerArray): ... -class Int64Array(IntegerArray[Int64Scalar]): +class Int64Array(IntegerArray): ... -class UInt64Array(IntegerArray[UInt64Scalar]): +class UInt64Array(IntegerArray): ... -class Date32Array(NumericArray[Date32Scalar]): +class Date32Array(NumericArray): ... -class Date64Array(NumericArray[Date64Scalar]): +class Date64Array(NumericArray): ... -class TimestampArray(NumericArray[TimestampScalar[_Unit, _Tz]]): +class TimestampArray(NumericArray): ... -class Time32Array(NumericArray[Time32Scalar[_Time32Unit]]): +class Time32Array(NumericArray): ... -class Time64Array(NumericArray[Time64Scalar[_Time64Unit]]): +class Time64Array(NumericArray): ... -class DurationArray(NumericArray[DurationScalar[_Unit]]): +class DurationArray(NumericArray): ... -class MonthDayNanoIntervalArray(Array[MonthDayNanoIntervalScalar]): +class MonthDayNanoIntervalArray(Array): ... -class HalfFloatArray(FloatingPointArray[HalfFloatScalar]): +class HalfFloatArray(FloatingPointArray): ... -class FloatArray(FloatingPointArray[FloatScalar]): +class FloatArray(FloatingPointArray): ... -class DoubleArray(FloatingPointArray[DoubleScalar]): +class DoubleArray(FloatingPointArray): ... -class FixedSizeBinaryArray(Array[FixedSizeBinaryScalar]): +class FixedSizeBinaryArray(Array): ... @@ -492,7 +417,7 @@ class Decimal256Array(FixedSizeBinaryArray): ... -class BaseListArray(Array[_ScalarT]): +class BaseListArray(Array): def flatten(self, recursive: bool = False) -> Array: ... def value_parent_indices(self) -> Int64Array: ... @@ -500,20 +425,17 @@ class BaseListArray(Array[_ScalarT]): def value_lengths(self) -> Int32Array: ... -class ListArray(BaseListArray[_ScalarT]): +class ListArray(BaseListArray): @classmethod def from_arrays( cls, offsets: Int32Array | list[int] | list[int | None], - values: Array[Scalar[_DataTypeT]] | list[int] | list[float] | list[str] - | list[bytes] | list, + values: Array | list[Any], *, - type: _DataTypeT | None = None, + type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, - ) -> (ListArray[ListScalar[ - _DataTypeT | Int64Type | Float64Type | StringType | BinaryType - ]] | ListArray): ... + ) -> ListArray: ... @property def values(self) -> Array: ... @@ -522,17 +444,17 @@ class ListArray(BaseListArray[_ScalarT]): def offsets(self) -> Int32Array: ... -class LargeListArray(BaseListArray[LargeListScalar[_DataTypeT]]): +class LargeListArray(BaseListArray): @classmethod def from_arrays( cls, offsets: Int64Array | list[int] | list[int | None], - values: Array[Scalar[_DataTypeT]] | Array, + values: Array, *, - type: _DataTypeT | None = None, + type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, - ) -> LargeListArray[_DataTypeT]: ... + ) -> LargeListArray: ... @property def values(self) -> Array: ... @@ -541,18 +463,18 @@ class LargeListArray(BaseListArray[LargeListScalar[_DataTypeT]]): def offsets(self) -> Int64Array: ... -class ListViewArray(BaseListArray[ListViewScalar[_DataTypeT]]): +class ListViewArray(BaseListArray): @classmethod def from_arrays( cls, offsets: Int32Array | list[int] | list[int | None], sizes: Int32Array | list[int] | list[int | None], - values: Array[Scalar[_DataTypeT]] | Array, + values: Array, *, - type: _DataTypeT | None = None, + type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, - ) -> ListViewArray[_DataTypeT]: ... + ) -> ListViewArray: ... @property def values(self) -> Array: ... @@ -564,18 +486,18 @@ class ListViewArray(BaseListArray[ListViewScalar[_DataTypeT]]): def sizes(self) -> Int32Array: ... -class LargeListViewArray(BaseListArray[LargeListScalar[_DataTypeT]]): +class LargeListViewArray(BaseListArray): @classmethod def from_arrays( cls, offsets: Int64Array | list[int] | list[int | None], sizes: Int64Array | list[int] | list[int | None], - values: Array[Scalar[_DataTypeT]] | Array, + values: Array, *, - type: _DataTypeT | None = None, + type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, - ) -> LargeListViewArray[_DataTypeT]: ... + ) -> LargeListViewArray: ... @property def values(self) -> Array: ... @@ -587,37 +509,33 @@ class LargeListViewArray(BaseListArray[LargeListScalar[_DataTypeT]]): def sizes(self) -> Int64Array: ... -class FixedSizeListArray(BaseListArray[FixedSizeListScalar[_DataTypeT, _Size]]): +class FixedSizeListArray(BaseListArray): @classmethod def from_arrays( cls, - values: Array[Scalar[_DataTypeT]], - list_size: _Size | None = None, + values: Array, + list_size: int | None = None, *, type: DataType | None = None, mask: Mask | None = None, - ) -> FixedSizeListArray[_DataTypeT, _Size | None]: ... + ) -> FixedSizeListArray: ... @property def values(self) -> Array: ... -_MapKeyT = TypeVar("_MapKeyT", bound=_BasicDataType) -_MapItemT = TypeVar("_MapItemT", bound=_BasicDataType) - - -class MapArray(BaseListArray[MapScalar[_MapKeyT, _MapItemT]]): +class MapArray(BaseListArray): @classmethod def from_arrays( cls, offsets: Int32Array | list[int] | list[int | None], - keys: Array[Scalar[_MapKeyT]] | np.ndarray | list, - items: Array[Scalar[_MapItemT]] | np.ndarray | list, + keys: Array | np.ndarray | list, + items: Array | np.ndarray | list, *, type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, - ) -> MapArray[_MapKeyT, _MapItemT]: ... + ) -> MapArray: ... @property def keys(self) -> Array: ... @@ -626,7 +544,7 @@ class MapArray(BaseListArray[MapScalar[_MapKeyT, _MapItemT]]): def items(self) -> Array: ... -class UnionArray(Array[UnionScalar]): +class UnionArray(Array): @deprecated("Use fields() instead") def child(self, pos: int) -> Field: ... @@ -642,7 +560,7 @@ class UnionArray(Array[UnionScalar]): def from_dense( types: Int8Array, value_offsets: Int32Array, - children: NullableCollection[Array], + children: NullableCollection, field_names: list[str] | None = None, type_codes: Int8Array | list[int] | None = None, ) -> UnionArray: ... @@ -650,13 +568,13 @@ class UnionArray(Array[UnionScalar]): @staticmethod def from_sparse( types: Int8Array, - children: NullableCollection[Array], + children: NullableCollection, field_names: list[str] | None = None, type_codes: Int8Array | list[int] | None = None, ) -> UnionArray: ... -class StringArray(Array[StringScalar]): +class StringArray(Array): @staticmethod def from_buffers( # type: ignore[override] length: int, @@ -668,7 +586,7 @@ class StringArray(Array[StringScalar]): ) -> StringArray: ... -class LargeStringArray(Array[LargeStringScalar]): +class LargeStringArray(Array): @staticmethod def from_buffers( # type: ignore[override] length: int, @@ -680,42 +598,42 @@ class LargeStringArray(Array[LargeStringScalar]): ) -> LargeStringArray: ... -class StringViewArray(Array[StringViewScalar]): +class StringViewArray(Array): ... -class BinaryArray(Array[BinaryScalar]): +class BinaryArray(Array): @property def total_values_length(self) -> int: ... -class LargeBinaryArray(Array[LargeBinaryScalar]): +class LargeBinaryArray(Array): @property def total_values_length(self) -> int: ... -class BinaryViewArray(Array[BinaryViewScalar]): +class BinaryViewArray(Array): ... -class DictionaryArray(Array[DictionaryScalar[_IndexT, _BasicValueT]]): +class DictionaryArray(Array): def dictionary_encode(self) -> Self: ... # type: ignore[override] - def dictionary_decode(self) -> Array[Scalar[_BasicValueT]]: ... + def dictionary_decode(self) -> Array: ... @property - def indices(self) -> Array[Scalar[_IndexT]]: ... + def indices(self) -> IntegerArray: ... @property - def dictionary(self) -> Array[Scalar[_BasicValueT]]: ... + def dictionary(self) -> Array: ... @staticmethod def from_buffers( # type: ignore[override] - type: _BasicValueT, + type: DataType, length: int, buffers: list[Buffer], dictionary: Array | np.ndarray | pd.Series, null_count: int = -1, offset: int = 0, - ) -> DictionaryArray[Any, _BasicValueT]: ... + ) -> DictionaryArray: ... @staticmethod def from_arrays( @@ -729,7 +647,7 @@ class DictionaryArray(Array[DictionaryScalar[_IndexT, _BasicValueT]]): ) -> DictionaryArray: ... -class StructArray(Array[StructScalar]): +class StructArray(Array): def field(self, index: int | str) -> Array: ... def flatten(self, memory_pool: MemoryPool | None = None) -> list[Array]: ... @@ -748,12 +666,12 @@ class StructArray(Array[StructScalar]): None = None, **kwargs) -> StructArray: ... -class RunEndEncodedArray(Array[RunEndEncodedScalar[_RunEndType, _BasicValueT]]): +class RunEndEncodedArray(Array): @staticmethod def from_arrays( run_ends: Int16Array | Int32Array | Int64Array | list[int], values: Array | list[Any], type: DataType | None = None, - ) -> RunEndEncodedArray[Any, _BasicValueT]: ... + ) -> RunEndEncodedArray: ... @staticmethod def from_buffers( # type: ignore[override] @@ -763,40 +681,37 @@ class RunEndEncodedArray(Array[RunEndEncodedScalar[_RunEndType, _BasicValueT]]): null_count: int = -1, offset=0, children: tuple[Array, Array] | list[list[int]] | None = None, - ) -> RunEndEncodedArray[Any, _BasicValueT]: ... + ) -> RunEndEncodedArray: ... @property - def run_ends(self) -> Array[Scalar[_RunEndType]]: ... + def run_ends(self) -> Int16Array | Int32Array | Int64Array: ... @property - def values(self) -> Array[Scalar[_BasicValueT]]: ... + def values(self) -> Array: ... def find_physical_offset(self) -> int: ... def find_physical_length(self) -> int: ... -_ArrayT = TypeVar("_ArrayT", bound=Array) - - -class ExtensionArray(Array[ExtensionScalar], Generic[_ArrayT]): +class ExtensionArray(Array): @property def storage(self) -> Any: ... @staticmethod def from_storage(typ: BaseExtensionType, - storage: _ArrayT) -> ExtensionArray[_ArrayT]: ... + storage: Array) -> ExtensionArray: ... -class JsonArray(ExtensionArray[_ArrayT]): +class JsonArray(ExtensionArray): ... -class UuidArray(ExtensionArray[_ArrayT]): +class UuidArray(ExtensionArray): ... -class FixedShapeTensorArray(ExtensionArray[_ArrayT]): +class FixedShapeTensorArray(ExtensionArray): def to_numpy_ndarray(self) -> np.ndarray: ... def to_tensor(self) -> Tensor: ... @@ -808,7 +723,7 @@ class FixedShapeTensorArray(ExtensionArray[_ArrayT]): ) -> Self: ... -class OpaqueArray(ExtensionArray[_ArrayT]): +class OpaqueArray(ExtensionArray): ... @@ -823,11 +738,11 @@ class Bool8Array(ExtensionArray): def from_numpy(cls, obj: np.ndarray) -> Self: ... -def concat_arrays(arrays: Iterable[_ArrayT], - memory_pool: MemoryPool | None = None) -> _ArrayT: ... +def concat_arrays(arrays: Iterable[Array], + memory_pool: MemoryPool | None = None) -> Array: ... -def _empty_array(type: _DataTypeT) -> Array[Scalar[_DataTypeT]]: ... +def _empty_array(type: DataType) -> Array: ... __all__ = [ @@ -891,5 +806,4 @@ __all__ = [ "FixedShapeTensorArray", "concat_arrays", "_empty_array", - "_CastAs", ] diff --git a/python/pyarrow-stubs/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi index 5e6a9f640591..e153f5ddc839 100644 --- a/python/pyarrow-stubs/pyarrow/table.pyi +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -27,7 +27,7 @@ else: from typing_extensions import TypeAlias from collections.abc import ( Collection, Generator, Iterable, Iterator, Sequence, Mapping) -from typing import Any, Generic, Literal, TypeVar +from typing import Any, Literal import builtins import numpy as np @@ -54,22 +54,18 @@ from pyarrow._stubs_typing import ( ) from pyarrow.compute import Expression from pyarrow.interchange.dataframe import _PyArrowDataFrame -from pyarrow.lib import Device, MemoryManager, MemoryPool, Schema -from pyarrow.lib import Field as _Field +from pyarrow.lib import Device, Field, MemoryManager, MemoryPool, Schema -from .array import Array, StructArray, _CastAs, _PandasConvertible +from .array import Array, StructArray, _PandasConvertible from .device import DeviceAllocationType from .io import Buffer from ._ipc import RecordBatchReader -from .scalar import BooleanScalar, Int64Scalar, Scalar, StructScalar +from .scalar import Int64Scalar, Scalar from .tensor import Tensor from ._stubs_typing import NullableCollection -from ._types import DataType, _AsPyType, _BasicDataType, _DataTypeT +from ._types import DataType -Field: TypeAlias = _Field[DataType] -_ScalarT = TypeVar("_ScalarT", bound=Scalar) -_Scalar_co = TypeVar("_Scalar_co", bound=Scalar, covariant=True) -ArrayOrChunkedArray: TypeAlias = Array[_Scalar_co] | ChunkedArray[_Scalar_co] +ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray _Aggregation: TypeAlias = Literal[ "all", @@ -127,14 +123,14 @@ NarySelector: TypeAlias = list[str] | tuple[str, ...] ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector -class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): +class ChunkedArray(_PandasConvertible): def as_py(self) -> list[Any]: ... @property def data(self) -> Self: ... @property - def type(self: ChunkedArray[Scalar[_DataTypeT]]) -> _DataTypeT: ... + def type(self) -> DataType: ... # Private attribute used internally for column names _name: str | None @@ -166,22 +162,22 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def __sizeof__(self) -> int: ... def __getitem__( - self, key: int | np.integer | builtins.slice) -> _Scalar_co | Self: ... + 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[BooleanScalar]: ... + def is_null(self, *, nan_is_null: bool = False) -> ChunkedArray: ... - def is_nan(self) -> ChunkedArray[BooleanScalar]: ... + def is_nan(self) -> ChunkedArray: ... - def is_valid(self) -> ChunkedArray[BooleanScalar]: ... + def is_valid(self) -> ChunkedArray: ... def cast( - self, target_type: _CastAs | str | None, safe: bool = True, + self, target_type: DataType | str | None, safe: bool = True, options: CastOptions | None = None, memory_pool: MemoryPool | None = None - ) -> Self | ChunkedArray[Scalar[_CastAs]]: ... + ) -> ChunkedArray: ... - def fill_null(self, fill_value: Scalar[_DataTypeT] | Any) -> Self: ... + def fill_null(self, fill_value: Any) -> Self: ... def equals(self, other: Self | Any) -> bool: ... @@ -193,12 +189,12 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def dictionary_encode(self, null_encoding: NullEncoding = "mask") -> Self: ... def flatten(self, memory_pool: MemoryPool | - None = None) -> list[ChunkedArray[Any]]: ... + None = None) -> list[ChunkedArray]: ... def combine_chunks(self, memory_pool: MemoryPool | - None = None) -> Array[_Scalar_co]: ... + None = None) -> Array: ... - def unique(self) -> ChunkedArray[_Scalar_co]: ... + def unique(self) -> ChunkedArray: ... def value_counts(self) -> StructArray: ... @@ -208,8 +204,8 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): null_selection_behavior: NullSelectionBehavior = "drop") -> Self: ... def index( - self: ChunkedArray[Scalar[_BasicDataType[_AsPyType]]], - value: Scalar[_DataTypeT] | _AsPyType, + self, + value: Any, start: int | None = None, end: int | None = None, *, @@ -227,22 +223,20 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): @property def num_chunks(self) -> int: ... - def chunk(self, i: int) -> Array[_Scalar_co]: ... + def chunk(self, i: int) -> Array: ... @property - def chunks(self) -> list[Array[_Scalar_co]]: ... + def chunks(self) -> list[Array]: ... - def iterchunks( - self: ArrayOrChunkedArray[_ScalarT], - ) -> Generator[Array, None, None]: ... + def iterchunks(self) -> Generator[Array, None, None]: ... - def __iter__(self) -> Iterator[_Scalar_co]: ... + def __iter__(self) -> Iterator[Scalar]: ... def to_pylist( - self: ChunkedArray[Scalar[_BasicDataType[_AsPyType]]], + self, *, maps_as_pydicts: Literal["lossy", "strict"] | None = None, - ) -> list[_AsPyType | None]: ... + ) -> list[Any]: ... def __arrow_c_stream__(self, requested_schema=None) -> Any: ... @@ -254,18 +248,15 @@ class ChunkedArray(_PandasConvertible[pd.Series], Generic[_Scalar_co]): def chunked_array( - arrays: Iterable[NullableCollection[Any]] + arrays: Iterable[NullableCollection] | Iterable[Iterable[Any] | SupportsArrowStream | SupportsArrowArray] - | Iterable[Array[_ScalarT]] | Array[_ScalarT] + | Iterable[Array] | Array | SupportsArrowArray | SupportsArrowStream, type: DataType | str | None = None, -) -> ChunkedArray[Scalar[Any]] | ChunkedArray[_ScalarT]: ... - - -_ColumnT = TypeVar("_ColumnT", bound=ArrayOrChunkedArray[Any]) +) -> ChunkedArray: ... -class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): +class _Tabular(_PandasConvertible): def __array__(self, dtype: np.dtype | None = None, copy: bool | None = None) -> np.ndarray: ... @@ -273,16 +264,16 @@ class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): self, nan_as_null: bool = False, allow_copy: bool = True ) -> _PyArrowDataFrame: ... - def __getitem__(self, key: int | str | slice) -> _ColumnT | Self: ... + def __getitem__(self, key: int | str | slice) -> ArrayOrChunkedArray | Self: ... def __len__(self) -> int: ... - def column(self, i: int | str) -> _ColumnT: ... + def column(self, i: int | str) -> ArrayOrChunkedArray: ... @property def column_names(self) -> list[str]: ... @property - def columns(self) -> list[_ColumnT]: ... + def columns(self) -> list[Array] | list[ChunkedArray]: ... def drop_null(self) -> Self: ... @@ -292,7 +283,7 @@ class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): def from_pydict( cls, mapping: - Mapping[Any, ArrayOrChunkedArray[Any] | list[Any] | np.ndarray | range], + Mapping[Any, ArrayOrChunkedArray | list[Any] | np.ndarray | range], schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> Self: ... @@ -305,7 +296,7 @@ class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> Self: ... - def itercolumns(self) -> Generator[_ColumnT, None, None]: ... + def itercolumns(self) -> Generator[ArrayOrChunkedArray, None, None]: ... @property def num_columns(self) -> int: ... @@ -342,14 +333,14 @@ class _Tabular(_PandasConvertible[pd.DataFrame], Generic[_ColumnT]): def drop_columns(self, columns: str | list[str]) -> Self: ... def add_column(self, i: int, field_: str | Field, - column: ArrayOrChunkedArray[Any] | list[list[Any]]) -> Self: ... + column: ArrayOrChunkedArray | list[list[Any]]) -> Self: ... def append_column( - self, field_: str | Field, column: ArrayOrChunkedArray[Any] | list[list[Any]] + self, field_: str | Field, column: ArrayOrChunkedArray | list[list[Any]] ) -> Self: ... -class RecordBatch(_Tabular[Array]): +class RecordBatch(_Tabular): def validate(self, *, full: bool = False) -> None: ... def replace_schema_metadata( @@ -378,7 +369,7 @@ class RecordBatch(_Tabular[Array]): def __sizeof__(self) -> int: ... def add_column( - self, i: int, field_: str | Field, column: ArrayOrChunkedArray[Any] | list + self, i: int, field_: str | Field, column: ArrayOrChunkedArray | list ) -> Self: ... def remove_column(self, i: int) -> Self: ... @@ -424,7 +415,7 @@ class RecordBatch(_Tabular[Array]): @classmethod def from_struct_array( - cls, struct_array: StructArray | ChunkedArray[StructScalar] + cls, struct_array: StructArray | ChunkedArray ) -> Self: ... def to_struct_array(self) -> StructArray: ... @@ -482,7 +473,7 @@ JoinType: TypeAlias = Literal[ ] -class Table(_Tabular[ChunkedArray[Any]]): +class Table(_Tabular): def validate(self, *, full: bool = False) -> None: ... def slice(self, offset: int = 0, length: int | None = None) -> Self: ... @@ -524,7 +515,7 @@ class Table(_Tabular[ChunkedArray[Any]]): def from_arrays( cls, arrays: - Collection[ArrayOrChunkedArray[Any] | Collection[NDArray[Any]] | list[Any]], + Collection[ArrayOrChunkedArray | Collection[NDArray[Any]] | list[Any]], names: list[str] | tuple[str, ...] | None = None, schema: Schema | None = None, metadata: Mapping[bytes, bytes] @@ -535,12 +526,12 @@ class Table(_Tabular[ChunkedArray[Any]]): @classmethod def from_struct_array( - cls, struct_array: StructArray | ChunkedArray[StructScalar] + cls, struct_array: StructArray | ChunkedArray ) -> Self: ... def to_struct_array( self, max_chunksize: int | None = None - ) -> ChunkedArray[StructScalar]: ... + ) -> ChunkedArray: ... @classmethod def from_batches(cls, batches: Iterable[RecordBatch], @@ -567,12 +558,12 @@ class Table(_Tabular[ChunkedArray[Any]]): def __sizeof__(self) -> int: ... def add_column(self, i: int, field_: str | Field, - column: ArrayOrChunkedArray[Any] | list[list[Any]]) -> Self: ... + column: ArrayOrChunkedArray | list[list[Any]]) -> Self: ... def remove_column(self, i: int) -> Self: ... def set_column(self, i: int, field_: str | Field, - column: ArrayOrChunkedArray[Any] | list[list[Any]]) -> Self: ... + column: ArrayOrChunkedArray | list[list[Any]]) -> Self: ... def rename_columns(self, names: list[str] | dict[str, str]) -> Self: ... @@ -610,8 +601,8 @@ class Table(_Tabular[ChunkedArray[Any]]): def record_batch( - data: Mapping[str, list[Any] | Array[Any]] - | Collection[Array[Any] | ChunkedArray[Any] | list[Any]] + data: Mapping[str, list[Any] | Array] + | Collection[Array | ChunkedArray | list[Any]] | pd.DataFrame | SupportsArrowArray | SupportsArrowDeviceArray, @@ -622,12 +613,12 @@ def record_batch( def table( - data: Collection[ArrayOrChunkedArray[Any] | list[Any] | range | str] + data: Collection[ArrayOrChunkedArray | list[Any] | range | str] | pd.DataFrame | SupportsArrowArray | SupportsArrowStream | SupportsArrowDeviceArray - | Mapping[str, list[Any] | Array[Any] | ChunkedArray[Any] | range] + | Mapping[str, list[Any] | Array | ChunkedArray | range] | Mapping[str, Any], names: list[str] | Schema | None = None, schema: Schema | None = None, diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index 99d1fbfe084d..c7e2fa6e45b5 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -81,26 +81,20 @@ _SignedInteger: TypeAlias = Int8Type | Int16Type | Int32Type | Int64Type _UnsignedInteger: TypeAlias = UInt8Type | UInt16Type | UInt32Type | UInt64Type _Integer: TypeAlias = _SignedInteger | _UnsignedInteger _Floating: TypeAlias = Float16Type | Float32Type | Float64Type -_Decimal: TypeAlias = ( - Decimal32Type[Any, Any] - | Decimal64Type[Any, Any] - | Decimal128Type[Any, Any] - | Decimal256Type[Any, Any] -) +_Decimal: TypeAlias = Decimal32Type | Decimal64Type | Decimal128Type | Decimal256Type _Date: TypeAlias = Date32Type | Date64Type -_Time: TypeAlias = Time32Type[Any] | Time64Type[Any] +_Time: TypeAlias = Time32Type | Time64Type _Interval: TypeAlias = MonthDayNanoIntervalType -_Temporal: TypeAlias = (TimestampType[Any, Any] - | DurationType[Any] | _Time | _Date | _Interval) +_Temporal: TypeAlias = TimestampType | DurationType | _Time | _Date | _Interval _Union: TypeAlias = SparseUnionType | DenseUnionType _Nested: TypeAlias = ( - ListType[Any] - | FixedSizeListType[Any, Any] - | LargeListType[Any] - | ListViewType[Any] - | LargeListViewType[Any] + ListType + | FixedSizeListType + | LargeListType + | ListViewType + | LargeListViewType | StructType - | MapType[Any, Any, Any] + | MapType | _Union ) @@ -122,21 +116,21 @@ 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[Any]]: ... -def is_large_list(t: DataType) -> TypeIs[LargeListType[Any]]: ... -def is_fixed_size_list(t: DataType) -> TypeIs[FixedSizeListType[Any, Any]]: ... -def is_list_view(t: DataType) -> TypeIs[ListViewType[Any]]: ... -def is_large_list_view(t: DataType) -> TypeIs[LargeListViewType[Any]]: ... +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[Any, Any]]: ... +def is_run_end_encoded(t: DataType) -> TypeIs[RunEndEncodedType]: ... def is_temporal(t: DataType) -> TypeIs[_Temporal]: ... -def is_timestamp(t: DataType) -> TypeIs[TimestampType[Any, Any]]: ... -def is_duration(t: DataType) -> TypeIs[DurationType[Any]]: ... +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[Any]]: ... -def is_time64(t: DataType) -> TypeIs[Time64Type[Any]]: ... +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]: ... @@ -149,13 +143,13 @@ 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[Any, Any, Any]]: ... +def is_map(t: DataType) -> TypeIs[MapType]: ... def is_decimal(t: DataType) -> TypeIs[_Decimal]: ... -def is_decimal32(t: DataType) -> TypeIs[Decimal32Type[Any, Any]]: ... -def is_decimal64(t: DataType) -> TypeIs[Decimal64Type[Any, Any]]: ... -def is_decimal128(t: DataType) -> TypeIs[Decimal128Type[Any, Any]]: ... -def is_decimal256(t: DataType) -> TypeIs[Decimal256Type[Any, Any]]: ... -def is_dictionary(t: DataType) -> TypeIs[DictionaryType[Any, Any, Any]]: ... +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: ... From bb8681e05072a1aad6fd7858bcb7f8e4a29688b2 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 12:17:17 +0100 Subject: [PATCH 25/93] Add tests for the type annotations of public APIs Add pyarrow-stubs/tests with pandas-stubs style checks: check(assert_type(lib.list_(lib.int32()), lib.ListType), lib.ListType) assert_type is verified by mypy, pyright and ty, which already include pyarrow-stubs/ in their runs, and check() verifies the runtime class when the tests run under pytest, so the stubs can't silently drift from what pyarrow returns. Usage that should be rejected goes under `if TYPE_CHECKING_INVALID_USAGE:` with a `# type: ignore[...]`; unused ignores are errors in these tests, so those fail if the usage starts type checking. To have something to test, lib.pyi now re-exports the stub modules that are annotated so far, still falling back to Any for everything else. The CI type annotation script runs the tests with pytest after the type checkers. --- ci/scripts/python_test_type_annotations.sh | 4 + docs/source/developers/python/development.rst | 41 +++++ python/pyarrow-stubs/pyarrow/lib.pyi | 17 +- python/pyarrow-stubs/tests/_check.py | 58 +++++++ python/pyarrow-stubs/tests/test_array.py | 88 ++++++++++ python/pyarrow-stubs/tests/test_table.py | 62 +++++++ python/pyarrow-stubs/tests/test_types.py | 161 ++++++++++++++++++ python/pyproject.toml | 11 ++ 8 files changed, 440 insertions(+), 2 deletions(-) create mode 100644 python/pyarrow-stubs/tests/_check.py create mode 100644 python/pyarrow-stubs/tests/test_array.py create mode 100644 python/pyarrow-stubs/tests/test_table.py create mode 100644 python/pyarrow-stubs/tests/test_types.py 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..70ba69d89949 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -169,6 +169,47 @@ 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 and +ty are configured in ``pyproject.toml``). + Doctest ======= diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi index 0c539af067ed..e1c64d028bcf 100644 --- a/python/pyarrow-stubs/pyarrow/lib.pyi +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -15,9 +15,22 @@ # specific language governing permissions and limitations # under the License. -# Placeholder stub for pyarrow.lib C extension module. -# Complete type annotations will be added in subsequent PRs. +# 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/tests/_check.py b/python/pyarrow-stubs/tests/_check.py new file mode 100644 index 000000000000..5207f72109d1 --- /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 the type checkers 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..ef461076a005 --- /dev/null +++ b/python/pyarrow-stubs/tests/test_array.py @@ -0,0 +1,88 @@ +# 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) + + # 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) + + +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) + check(assert_type(chunked.chunks, list[lib.Array]), list) + check(assert_type(chunked.combine_chunks(), 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..f9d0b63e7ad1 --- /dev/null +++ b/python/pyarrow-stubs/tests/test_table.py @@ -0,0 +1,62 @@ +# 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.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) + + +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..41041517a9e2 --- /dev/null +++ b/python/pyarrow-stubs/tests/test_types.py @@ -0,0 +1,161 @@ +# 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_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, + ) + + +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/pyproject.toml b/python/pyproject.toml index 53b6abb4e02a..affc194d11b1 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -136,6 +136,12 @@ 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 = "." }, +] # TODO: Enable type checking once stubs are merged [tool.ty.environment] @@ -149,3 +155,8 @@ exclude = [ "examples", "scripts", ] + +[[tool.ty.overrides]] +include = ["pyarrow-stubs/tests/**"] +[tool.ty.overrides.rules] +unused-type-ignore-comment = "error" From 376407ddc69cfd86626055fd5d365568e251cd98 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 15:46:29 +0100 Subject: [PATCH 26/93] Add placeholder stub for pyarrow._ipc table.pyi imports RecordBatchReader from ._ipc, which had no stub, so pyright and ty failed to resolve the import. Add a placeholder in the same style as io.pyi and scalar.pyi until _ipc gets fully annotated. --- python/pyarrow-stubs/pyarrow/_ipc.pyi | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) create mode 100644 python/pyarrow-stubs/pyarrow/_ipc.pyi diff --git a/python/pyarrow-stubs/pyarrow/_ipc.pyi b/python/pyarrow-stubs/pyarrow/_ipc.pyi new file mode 100644 index 000000000000..a5c111b40499 --- /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 + +class RecordBatchReader: ... + +def __getattr__(name: str) -> Any: ... From bc27123747ef07d179c02a23845b249928f600af Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 16:28:48 +0100 Subject: [PATCH 27/93] Skip type checking of modules that aren't annotated yet compute and interchange will be annotated in later PRs, so treat imports of them as Any instead of type checking their untyped sources (mypy) or failing to resolve them (ty). Recent ty versions also no longer honor `# type: ignore[]` comments, so disable the rules those comments suppress: method overrides in array.pyi, and invalid arguments in the tests, where mypy and pyright still check that the invalid usage is rejected. --- docs/source/developers/python/development.rst | 5 ++-- python/pyarrow-stubs/tests/_check.py | 2 +- python/pyproject.toml | 27 ++++++++++++++++++- 3 files changed, 30 insertions(+), 4 deletions(-) diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index 70ba69d89949..f2194a7aeb66 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -207,8 +207,9 @@ 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 and -ty are configured in ``pyproject.toml``). +``# 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/tests/_check.py b/python/pyarrow-stubs/tests/_check.py index 5207f72109d1..596129874188 100644 --- a/python/pyarrow-stubs/tests/_check.py +++ b/python/pyarrow-stubs/tests/_check.py @@ -29,7 +29,7 @@ Code that type checkers should *reject* goes under ``if TYPE_CHECKING_INVALID_USAGE:`` with a ``# type: ignore[]`` -comment. The block never runs, and the type checkers are configured to +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. """ diff --git a/python/pyproject.toml b/python/pyproject.toml index affc194d11b1..947effb6a9df 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" @@ -147,6 +153,16 @@ executionEnvironments = [ [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 = [ @@ -156,7 +172,16 @@ exclude = [ "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] -unused-type-ignore-comment = "error" +# mypy and pyright check that the invalid usage in the tests is rejected. +invalid-argument-type = "ignore" From ddd9cb4683b6528533e413bf4de7597738d48f65 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 16:40:04 +0100 Subject: [PATCH 28/93] Fix stub annotations that don't match the runtime - TypesEnum: list all members, with the runtime values. - is_nested() also matches run-end encoded types. - Arrays and chunked arrays have no as_py() method. - ChunkedArray.unique() returns an Array. - MapArray is a ListArray, so it has offsets and values. - value_lengths() of large list (view) arrays returns an Int64Array. - Array.from_pandas() arguments can be passed positionally. Add annotation tests for these. --- python/pyarrow-stubs/pyarrow/array.pyi | 19 +++++++--- python/pyarrow-stubs/pyarrow/table.pyi | 4 +- python/pyarrow-stubs/pyarrow/types.pyi | 48 ++++++++++++++++++++++-- python/pyarrow-stubs/tests/test_array.py | 19 ++++++++++ python/pyarrow-stubs/tests/test_types.py | 12 ++++++ 5 files changed, 90 insertions(+), 12 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi index 4c06b0c9c520..bca6cdbf37fd 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -142,8 +142,6 @@ class _PandasConvertible(_Weakrefable): class Array(_PandasConvertible): - def as_py(self) -> list[Any]: ... - def diff(self, other: Self) -> str: ... # Private attribute used internally (e.g., for column names in batches) @@ -172,7 +170,6 @@ class Array(_PandasConvertible): @staticmethod def from_pandas( obj: pd.Series | np.ndarray | ArrayLike, - *, mask: Mask | None = None, type: DataType | None = None, safe: bool = True, @@ -422,7 +419,7 @@ class BaseListArray(Array): def value_parent_indices(self) -> Int64Array: ... - def value_lengths(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array | Int64Array: ... class ListArray(BaseListArray): @@ -443,6 +440,8 @@ class ListArray(BaseListArray): @property def offsets(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + class LargeListArray(BaseListArray): @classmethod @@ -462,6 +461,8 @@ class LargeListArray(BaseListArray): @property def offsets(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + class ListViewArray(BaseListArray): @classmethod @@ -485,6 +486,8 @@ class ListViewArray(BaseListArray): @property def sizes(self) -> Int32Array: ... + def value_lengths(self) -> Int32Array: ... + class LargeListViewArray(BaseListArray): @classmethod @@ -508,6 +511,8 @@ class LargeListViewArray(BaseListArray): @property def sizes(self) -> Int64Array: ... + def value_lengths(self) -> Int64Array: ... + class FixedSizeListArray(BaseListArray): @classmethod @@ -523,10 +528,12 @@ class FixedSizeListArray(BaseListArray): @property def values(self) -> Array: ... + def value_lengths(self) -> Int32Array: ... + -class MapArray(BaseListArray): +class MapArray(ListArray): @classmethod - def from_arrays( + def from_arrays( # type: ignore[override] cls, offsets: Int32Array | list[int] | list[int | None], keys: Array | np.ndarray | list, diff --git a/python/pyarrow-stubs/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi index e153f5ddc839..e60218af9424 100644 --- a/python/pyarrow-stubs/pyarrow/table.pyi +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -125,8 +125,6 @@ ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector class ChunkedArray(_PandasConvertible): - def as_py(self) -> list[Any]: ... - @property def data(self) -> Self: ... @property @@ -194,7 +192,7 @@ class ChunkedArray(_PandasConvertible): def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Array: ... - def unique(self) -> ChunkedArray: ... + def unique(self) -> Array: ... def value_counts(self) -> StructArray: ... diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index c7e2fa6e45b5..78ceabf2242f 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -96,6 +96,7 @@ _Nested: TypeAlias = ( | StructType | MapType | _Union + | RunEndEncodedType ) @@ -216,6 +217,47 @@ __all__ = [ class TypesEnum(IntEnum): - INTERVAL_MONTHS = 0 - INTERVAL_DAY_TIME = 1 - INTERVAL_MONTH_DAY_NANO = 2 + 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/test_array.py b/python/pyarrow-stubs/tests/test_array.py index ef461076a005..d88fe2db0939 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -43,6 +43,24 @@ def test_array() -> None: 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: @@ -75,6 +93,7 @@ def test_chunked_array() -> None: check(assert_type(chunked.chunk(0), lib.Array), lib.Int64Array) 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) diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 41041517a9e2..93ec4d417ec4 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -113,6 +113,13 @@ def test_map_and_run_end_encoded_types() -> None: 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) @@ -152,6 +159,11 @@ def test_type_predicates_narrow() -> None: lib.Decimal64Type, ) + # 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: From 92250c550c029ab40058930eea7436f3ed6b1407 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:41 +0100 Subject: [PATCH 29/93] Fix test_timedelta_overflow for sub-unit truncation Converting timedelta.max to a millisecond or second duration truncates the sub-unit part, so it doesn't round trip. The comparisons used to be missing their assert, which hid this. --- python/pyarrow/tests/test_types.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/pyarrow/tests/test_types.py b/python/pyarrow/tests/test_types.py index 8d79674f7f46..04a842443989 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -1106,10 +1106,12 @@ def test_timedelta_overflow(): # microsecond resolution, not overflow 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]: - assert pa.scalar(d, type=pa.duration('ms')).as_py() == d - assert 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(): From 9636e5c769312ff33b192265cbb41c6d42ed4fdf Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:45 +0100 Subject: [PATCH 30/93] Keep RecordBatchReader as Any until it gets annotated An empty placeholder class rejects valid usage such as table.to_reader().read_all(), as it has no methods. --- python/pyarrow-stubs/pyarrow/_ipc.pyi | 4 ++-- python/pyarrow-stubs/tests/test_table.py | 5 +++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_ipc.pyi b/python/pyarrow-stubs/pyarrow/_ipc.pyi index a5c111b40499..c84964043e4a 100644 --- a/python/pyarrow-stubs/pyarrow/_ipc.pyi +++ b/python/pyarrow-stubs/pyarrow/_ipc.pyi @@ -18,8 +18,8 @@ # Placeholder stub - complete annotations in future PR. # At runtime, these symbols are provided by the pyarrow.lib C extension. -from typing import Any +from typing import Any, TypeAlias -class RecordBatchReader: ... +RecordBatchReader: TypeAlias = Any def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index f9d0b63e7ad1..908508eac533 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -42,6 +42,11 @@ def test_table() -> None: 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]}) From 4abbeccae36ce6ab4541b570dcfed63f60c54c9e Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:48 +0100 Subject: [PATCH 31/93] Use TypeGuard for is_fixed_size_binary() Decimal types subclass FixedSizeBinaryType, but is_fixed_size_binary() returns False for them. With TypeIs, type checkers would wrongly narrow decimal types away when the check fails. --- python/pyarrow-stubs/pyarrow/types.pyi | 5 +++-- python/pyarrow-stubs/tests/test_types.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index 78ceabf2242f..1bb66bd1209a 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -18,7 +18,7 @@ import sys from enum import IntEnum -from typing import Any +from typing import Any, TypeGuard if sys.version_info >= (3, 13): from typing import TypeIs @@ -138,7 +138,8 @@ 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]: ... -def is_fixed_size_binary(t: DataType) -> TypeIs[FixedSizeBinaryType]: ... +# 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]: ... diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 93ec4d417ec4..2133cb4de874 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -159,6 +159,20 @@ def test_type_predicates_narrow() -> None: 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) From 325ed1da02ff916856db065694706c5a8e0c137a Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:51 +0100 Subject: [PATCH 32/93] Test indexing of arrays, chunked arrays and tables --- python/pyarrow-stubs/tests/test_array.py | 6 ++++++ python/pyarrow-stubs/tests/test_table.py | 6 ++++++ 2 files changed, 12 insertions(+) diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index d88fe2db0939..d41826eed8dc 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -32,6 +32,9 @@ def test_array() -> None: 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) @@ -91,6 +94,9 @@ def test_chunked_array() -> None: 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) diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index 908508eac533..9544f44be794 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -35,6 +35,12 @@ def test_table() -> None: # 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) From 26fa0e0820b8f7cfe907b43f10826d4cff8b961d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:55 +0100 Subject: [PATCH 33/93] Add manual pre-commit hooks for linting and type checking the stubs Run ruff, mypy, pyright and ty on the stubs with pinned versions: pre-commit run --hook-stage manual python-stubs --all-files The hooks only run on request, so they don't affect the existing hooks or CI jobs. --- .pre-commit-config.yaml | 74 +++++++++++++++++++ docs/source/developers/python/development.rst | 12 +++ python/pyproject.toml | 8 ++ 3 files changed, 94 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index df3e8b799d77..3b6ea465b5de 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -156,6 +156,80 @@ repos: ( ?^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/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index f2194a7aeb66..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: diff --git a/python/pyproject.toml b/python/pyproject.toml index 947effb6a9df..fa471a5456aa 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -185,3 +185,11 @@ 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"] From 7d7a9fd65b1281ee9d010caa80966c9398e946de Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 27 Oct 2026 14:02:32 +0100 Subject: [PATCH 34/93] Add _compute.pyi and reportMissingModuleSource --- python/pyarrow-stubs/pyarrow/_compute.pyi | 26 +++++++++++++++++++++++ python/pyproject.toml | 2 ++ 2 files changed, 28 insertions(+) create mode 100644 python/pyarrow-stubs/pyarrow/_compute.pyi diff --git a/python/pyarrow-stubs/pyarrow/_compute.pyi b/python/pyarrow-stubs/pyarrow/_compute.pyi new file mode 100644 index 000000000000..e6e9cbd56bc8 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/_compute.pyi @@ -0,0 +1,26 @@ +# 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/pyproject.toml b/python/pyproject.toml index fa471a5456aa..2b713a39032a 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -148,6 +148,8 @@ 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] From 12ea05ff7bbdac61cee11c075c1cf557a0c19e00 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 27 Oct 2026 15:58:14 +0100 Subject: [PATCH 35/93] Fix linter error --- python/pyarrow-stubs/pyarrow/_compute.pyi | 1 - 1 file changed, 1 deletion(-) diff --git a/python/pyarrow-stubs/pyarrow/_compute.pyi b/python/pyarrow-stubs/pyarrow/_compute.pyi index e6e9cbd56bc8..3108cb063747 100644 --- a/python/pyarrow-stubs/pyarrow/_compute.pyi +++ b/python/pyarrow-stubs/pyarrow/_compute.pyi @@ -23,4 +23,3 @@ from typing import Any, TypeAlias Expression: TypeAlias = Any def __getattr__(name: str) -> Any: ... - From 25a866ba9da1dc47a48b33ed8f9b04a76ba77c24 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 27 Oct 2026 16:30:38 +0100 Subject: [PATCH 36/93] Stubs ruff-format --- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 5 +- python/pyarrow-stubs/pyarrow/_types.pyi | 4 +- python/pyarrow-stubs/pyarrow/array.pyi | 327 ++++-------------- python/pyarrow-stubs/pyarrow/builder.pyi | 9 - python/pyarrow-stubs/pyarrow/config.pyi | 7 - python/pyarrow-stubs/pyarrow/device.pyi | 17 +- python/pyarrow-stubs/pyarrow/memory.pyi | 38 +- python/pyarrow-stubs/pyarrow/table.pyi | 277 +++++---------- python/pyarrow-stubs/pyarrow/tensor.pyi | 96 ++--- python/pyarrow-stubs/pyarrow/types.pyi | 4 +- 10 files changed, 179 insertions(+), 605 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 89c36cba98a0..7a2d258f76b0 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -72,10 +72,7 @@ Mask: TypeAlias = ( Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray ) Indices: TypeAlias = ( - Sequence[int | None] - | NDArray[np.integer[Any]] - | lib.IntegerArray - | ChunkedArray + Sequence[int | None] | NDArray[np.integer[Any]] | lib.IntegerArray | ChunkedArray ) class SupportsDunderEQ(Protocol): diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/_types.pyi index a46f72394cd4..9bc9420bc6d3 100644 --- a/python/pyarrow-stubs/pyarrow/_types.pyi +++ b/python/pyarrow-stubs/pyarrow/_types.pyi @@ -526,9 +526,7 @@ 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: ... +def list_(value_type: DataType | Field, list_size: Literal[-1] = -1) -> ListType: ... @overload def list_( value_type: DataType | Field, list_size: int diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/array.pyi index bca6cdbf37fd..c86932cf2f7f 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/array.pyi @@ -55,10 +55,12 @@ 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, + values: NullableCollection + | Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray, type: Any | None = None, mask: Mask | pd.Series[bool] | None = None, size: int | None = None, @@ -66,59 +68,45 @@ def array( safe: bool = True, memory_pool: MemoryPool | None = None, ) -> ArrayLike: ... - - def asarray( - values: NullableCollection | Iterable[Any] | SupportsArrowArray + 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: ... - +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, @@ -140,7 +128,6 @@ class _PandasConvertible(_Weakrefable): coerce_temporal_nanoseconds: bool = False, ) -> pd.Series | pd.DataFrame: ... - class Array(_PandasConvertible): def diff(self, other: Self) -> str: ... @@ -154,19 +141,13 @@ class Array(_PandasConvertible): 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, @@ -175,7 +156,6 @@ class Array(_PandasConvertible): safe: bool = True, memory_pool: MemoryPool | None = None, ) -> Array: ... - @staticmethod def from_buffers( type: DataType, @@ -185,17 +165,13 @@ class Array(_PandasConvertible): 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, *, @@ -208,32 +184,21 @@ class Array(_PandasConvertible): 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, @@ -242,15 +207,13 @@ class Array(_PandasConvertible): *, 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 __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, *, @@ -259,49 +222,32 @@ class Array(_PandasConvertible): 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 NullArray(Array): ... class BooleanArray(Array): @property @@ -309,119 +255,38 @@ class BooleanArray(Array): @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 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( @@ -433,16 +298,12 @@ class ListArray(BaseListArray): 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( @@ -454,16 +315,12 @@ class LargeListArray(BaseListArray): 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( @@ -476,19 +333,14 @@ class ListViewArray(BaseListArray): 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( @@ -501,19 +353,14 @@ class LargeListViewArray(BaseListArray): 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( @@ -524,13 +371,10 @@ class FixedSizeListArray(BaseListArray): 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] @@ -543,26 +387,19 @@ class MapArray(ListArray): 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, @@ -571,7 +408,6 @@ class UnionArray(Array): field_names: list[str] | None = None, type_codes: Int8Array | list[int] | None = None, ) -> UnionArray: ... - @staticmethod def from_sparse( types: Int8Array, @@ -580,7 +416,6 @@ class UnionArray(Array): type_codes: Int8Array | list[int] | None = None, ) -> UnionArray: ... - class StringArray(Array): @staticmethod def from_buffers( # type: ignore[override] @@ -592,7 +427,6 @@ class StringArray(Array): offset: int | None = 0, ) -> StringArray: ... - class LargeStringArray(Array): @staticmethod def from_buffers( # type: ignore[override] @@ -604,34 +438,25 @@ class LargeStringArray(Array): offset: int | None = 0, ) -> LargeStringArray: ... - -class StringViewArray(Array): - ... - +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 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, @@ -641,7 +466,6 @@ class DictionaryArray(Array): null_count: int = -1, offset: int = 0, ) -> DictionaryArray: ... - @staticmethod def from_arrays( indices: Indices | Sequence[int | None], @@ -653,12 +477,9 @@ class DictionaryArray(Array): 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], @@ -668,18 +489,17 @@ class StructArray(Array): memory_pool: MemoryPool | None = None, type: StructType | None = None, ) -> StructArray: ... - - def sort(self, order: Order = "ascending", by: str | - None = None, **kwargs) -> 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, + values: Array | list[Any], + type: DataType | None = None, ) -> RunEndEncodedArray: ... - @staticmethod def from_buffers( # type: ignore[override] type: DataType, @@ -689,69 +509,46 @@ class RunEndEncodedArray(Array): 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): - ... + 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 + cls, obj: np.ndarray, dim_names: list[str] | tuple[str, ...] | None = None ) -> Self: ... - -class OpaqueArray(ExtensionArray): - ... - +class OpaqueArray(ExtensionArray): ... class Bool8Array(ExtensionArray): - def to_numpy(self, zero_copy_only: bool = ..., - writable: bool = ...) -> np.ndarray: ... - + 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 concat_arrays( + arrays: Iterable[Array], memory_pool: MemoryPool | None = None +) -> Array: ... def _empty_array(type: DataType) -> Array: ... - __all__ = [ "array", "asarray", diff --git a/python/pyarrow-stubs/pyarrow/builder.pyi b/python/pyarrow-stubs/pyarrow/builder.pyi index 9001d9835b6c..482bdc3bcd04 100644 --- a/python/pyarrow-stubs/pyarrow/builder.pyi +++ b/python/pyarrow-stubs/pyarrow/builder.pyi @@ -21,31 +21,22 @@ 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 index 069b70e553ac..5f2fbc5ee863 100644 --- a/python/pyarrow-stubs/pyarrow/config.pyi +++ b/python/pyarrow-stubs/pyarrow/config.pyi @@ -17,13 +17,11 @@ from typing import NamedTuple - class VersionInfo(NamedTuple): major: int minor: int patch: int - class CppBuildInfo(NamedTuple): version: str version_info: VersionInfo @@ -37,27 +35,22 @@ class CppBuildInfo(NamedTuple): 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", diff --git a/python/pyarrow-stubs/pyarrow/device.pyi b/python/pyarrow-stubs/pyarrow/device.pyi index 7787ac44deb8..c8e47ac7f82d 100644 --- a/python/pyarrow-stubs/pyarrow/device.pyi +++ b/python/pyarrow-stubs/pyarrow/device.pyi @@ -19,7 +19,6 @@ import enum from pyarrow.lib import _Weakrefable - class DeviceAllocationType(enum.Enum): CPU = enum.auto() CUDA = enum.auto() @@ -36,31 +35,27 @@ class DeviceAllocationType(enum.Enum): 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"] +__all__ = [ + "DeviceAllocationType", + "Device", + "MemoryManager", + "default_cpu_memory_manager", +] diff --git a/python/pyarrow-stubs/pyarrow/memory.pyi b/python/pyarrow-stubs/pyarrow/memory.pyi index f80e01ab21c0..2645f2fab067 100644 --- a/python/pyarrow-stubs/pyarrow/memory.pyi +++ b/python/pyarrow-stubs/pyarrow/memory.pyi @@ -17,65 +17,31 @@ 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): - ... - +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", diff --git a/python/pyarrow-stubs/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/table.pyi index e60218af9424..5bf6d94a686a 100644 --- a/python/pyarrow-stubs/pyarrow/table.pyi +++ b/python/pyarrow-stubs/pyarrow/table.pyi @@ -25,8 +25,7 @@ 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 collections.abc import Collection, Generator, Iterable, Iterator, Sequence, Mapping from typing import Any, Literal import builtins @@ -114,17 +113,20 @@ _AggregationPrefixed: TypeAlias = Literal[ "hash_variance", ] Aggregation: TypeAlias = _Aggregation | _AggregationPrefixed | str -AggregateOptions: TypeAlias = (ScalarAggregateOptions | CountOptions - | TDigestOptions | VarianceOptions | FunctionOptions) +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 @@ -148,59 +150,39 @@ class ChunkedArray(_PandasConvertible): 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, 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, + self, + target_type: DataType | str | None, + safe: bool = True, options: CastOptions | None = None, - memory_pool: MemoryPool | 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 __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 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 filter( + self, mask: Mask, null_selection_behavior: NullSelectionBehavior = "drop" + ) -> Self: ... def index( self, value: Any, @@ -209,83 +191,61 @@ class ChunkedArray(_PandasConvertible): *, 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, + | 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 __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], + 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, @@ -293,101 +253,80 @@ class _Tabular(_PandasConvertible): 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: ... - + 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 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 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 + | 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: ... - + 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, @@ -400,7 +339,6 @@ class RecordBatch(_Tabular): | Mapping[str, bytes] | None = None, ) -> Self: ... - @classmethod def from_pandas( cls, @@ -410,55 +348,36 @@ class RecordBatch(_Tabular): nthreads: int | None = None, columns: Sequence[str | int] | None = None, ) -> Self: ... - @classmethod - def from_struct_array( - cls, struct_array: StructArray | ChunkedArray - ) -> Self: ... - + 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", @@ -470,34 +389,30 @@ JoinType: TypeAlias = Literal[ "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 select( + self, columns: Iterable[str] | Iterable[int] | NDArray[np.str_] + ) -> Self: ... def replace_schema_metadata( - self, metadata: dict[str, str] + self, + metadata: dict[str, str] | dict[bytes, bytes] | dict[bytes, str] | dict[str, bytes] - | None = None + | 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: ... - + def cast( + self, + target_schema: Schema, + safe: bool | None = None, + options: CastOptions | None = None, + ) -> Self: ... @classmethod def from_pandas( cls, @@ -508,68 +423,49 @@ class Table(_Tabular): columns: Sequence[str | int] | None = None, safe: bool = True, ) -> Self: ... - @classmethod def from_arrays( cls, - arrays: - Collection[ArrayOrChunkedArray | Collection[NDArray[Any]] | list[Any]], + 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, + | 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: ... - + 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 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 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 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 group_by( + self, keys: str | list[str], use_threads: bool = True + ) -> TableGroupBy: ... def join( self, right_table: Self, @@ -581,7 +477,6 @@ class Table(_Tabular): coalesce_keys: bool = True, use_threads: bool = True, ) -> Self: ... - def join_asof( self, right_table: Self, @@ -591,13 +486,10 @@ class Table(_Tabular): 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]] @@ -608,8 +500,6 @@ def record_batch( schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> RecordBatch: ... - - def table( data: Collection[ArrayOrChunkedArray | list[Any] | range | str] | pd.DataFrame @@ -623,8 +513,6 @@ def table( metadata: Mapping[str | bytes, str | bytes] | None = None, nthreads: int | None = None, ) -> Table: ... - - def concat_tables( tables: Iterable[Table], memory_pool: MemoryPool | None = None, @@ -632,14 +520,12 @@ def concat_tables( **kwargs: Any, ) -> Table: ... - class TableGroupBy: - keys: str | list[str] - def __init__(self, table: Table, keys: str | - list[str], use_threads: bool = True): ... - + def __init__( + self, table: Table, keys: str | list[str], use_threads: bool = True + ): ... def aggregate( self, aggregations: Iterable[ @@ -647,17 +533,14 @@ class TableGroupBy: | 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", diff --git a/python/pyarrow-stubs/pyarrow/tensor.pyi b/python/pyarrow-stubs/pyarrow/tensor.pyi index b4bd52c832f0..6a4a06eb989f 100644 --- a/python/pyarrow-stubs/pyarrow/tensor.pyi +++ b/python/pyarrow-stubs/pyarrow/tensor.pyi @@ -30,48 +30,36 @@ 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 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: ... - + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: list[str] | None = None + ) -> Self: ... @classmethod def from_numpy( cls, @@ -80,28 +68,21 @@ class SparseCOOTensor(_Weakrefable): 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: ... - + 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: ... - + 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 @@ -111,7 +92,6 @@ class SparseCOOTensor(_Weakrefable): @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... - @property def dim_names(self) -> list[str]: ... @property @@ -121,12 +101,11 @@ class SparseCOOTensor(_Weakrefable): @property def type(self) -> DataType: ... - class SparseCSRMatrix(_Weakrefable): @classmethod - def from_dense_numpy(cls, obj: np.ndarray, - dim_names: list[str] | None = None) -> Self: ... - + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: list[str] | None = None + ) -> Self: ... @classmethod def from_numpy( cls, @@ -136,22 +115,16 @@ class SparseCSRMatrix(_Weakrefable): 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: ... - + 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 @@ -161,7 +134,6 @@ class SparseCSRMatrix(_Weakrefable): @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... - @property def dim_names(self) -> list[str]: ... @property @@ -169,12 +141,11 @@ class SparseCSRMatrix(_Weakrefable): @property def type(self) -> DataType: ... - class SparseCSCMatrix(_Weakrefable): @classmethod - def from_dense_numpy(cls, obj: np.ndarray, - dim_names: list[str] | None = None) -> Self: ... - + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: list[str] | None = None + ) -> Self: ... @classmethod def from_numpy( cls, @@ -184,22 +155,16 @@ class SparseCSCMatrix(_Weakrefable): 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: ... - + 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 @@ -209,7 +174,6 @@ class SparseCSCMatrix(_Weakrefable): @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... - @property def dim_names(self) -> list[str]: ... @property @@ -217,12 +181,11 @@ class SparseCSCMatrix(_Weakrefable): @property def type(self) -> DataType: ... - class SparseCSFTensor(_Weakrefable): @classmethod - def from_dense_numpy(cls, obj: np.ndarray, - dim_names: Sequence[str] | None = None) -> Self: ... - + def from_dense_numpy( + cls, obj: np.ndarray, dim_names: Sequence[str] | None = None + ) -> Self: ... @classmethod def from_numpy( cls, @@ -233,16 +196,11 @@ class SparseCSFTensor(_Weakrefable): 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 @@ -252,7 +210,6 @@ class SparseCSFTensor(_Weakrefable): @property def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... - @property def dim_names(self) -> list[str]: ... @property @@ -260,7 +217,6 @@ class SparseCSFTensor(_Weakrefable): @property def type(self) -> DataType: ... - __all__ = [ "Tensor", "SparseCOOTensor", diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index 1bb66bd1209a..1ba1e60965d2 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -99,7 +99,6 @@ _Nested: TypeAlias = ( | RunEndEncodedType ) - def is_null(t: DataType) -> TypeIs[NullType]: ... def is_boolean(t: DataType) -> TypeIs[BoolType]: ... def is_integer(t: DataType) -> TypeIs[_Integer]: ... @@ -138,6 +137,7 @@ 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]: ... @@ -158,7 +158,6 @@ 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", @@ -216,7 +215,6 @@ __all__ = [ "is_unsigned_integer", ] - class TypesEnum(IntEnum): NA = 0 BOOL = 1 From 94021e0c8ea96fc2c14113e3088a4791433c5ff1 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Wed, 28 Oct 2026 18:18:10 +0100 Subject: [PATCH 37/93] Move to __lib/ to fix Module errors for pa.array/table/scalar --- .../pyarrow-stubs/pyarrow/__lib/__init__.pyi | 18 +++++++++++++++++ .../pyarrow/{ => __lib}/array.pyi | 4 ++-- .../pyarrow/{ => __lib}/builder.pyi | 0 .../pyarrow/{ => __lib}/config.pyi | 0 .../pyarrow/{ => __lib}/device.pyi | 0 .../pyarrow/{ => __lib}/error.pyi | 0 .../pyarrow-stubs/pyarrow/{ => __lib}/io.pyi | 0 .../pyarrow/{ => __lib}/memory.pyi | 0 .../pyarrow/{ => __lib}/scalar.pyi | 0 .../pyarrow/{ => __lib}/table.pyi | 6 +++--- .../pyarrow/{ => __lib}/tensor.pyi | 0 python/pyarrow-stubs/pyarrow/lib.pyi | 20 +++++++++---------- python/pyproject.toml | 2 +- 13 files changed, 34 insertions(+), 16 deletions(-) create mode 100644 python/pyarrow-stubs/pyarrow/__lib/__init__.pyi rename python/pyarrow-stubs/pyarrow/{ => __lib}/array.pyi (99%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/builder.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/config.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/device.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/error.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/io.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/memory.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/scalar.pyi (100%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/table.pyi (99%) rename python/pyarrow-stubs/pyarrow/{ => __lib}/tensor.pyi (100%) diff --git a/python/pyarrow-stubs/pyarrow/__lib/__init__.pyi b/python/pyarrow-stubs/pyarrow/__lib/__init__.pyi new file mode 100644 index 000000000000..50d570015931 --- /dev/null +++ b/python/pyarrow-stubs/pyarrow/__lib/__init__.pyi @@ -0,0 +1,18 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Separate subpackage so that e.g. array.pyi doesn't shadow pyarrow.array() diff --git a/python/pyarrow-stubs/pyarrow/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi similarity index 99% rename from python/pyarrow-stubs/pyarrow/array.pyi rename to python/pyarrow-stubs/pyarrow/__lib/array.pyi index c86932cf2f7f..daf6df35f7c3 100644 --- a/python/pyarrow-stubs/pyarrow/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -52,8 +52,8 @@ import builtins from .scalar import Int64Scalar, Scalar from .device import DeviceAllocationType -from ._types import BaseExtensionType, DataType, Field, StructType -from ._stubs_typing import NullableCollection +from pyarrow._types import BaseExtensionType, DataType, Field, StructType +from pyarrow._stubs_typing import NullableCollection def array( values: NullableCollection diff --git a/python/pyarrow-stubs/pyarrow/builder.pyi b/python/pyarrow-stubs/pyarrow/__lib/builder.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/builder.pyi rename to python/pyarrow-stubs/pyarrow/__lib/builder.pyi diff --git a/python/pyarrow-stubs/pyarrow/config.pyi b/python/pyarrow-stubs/pyarrow/__lib/config.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/config.pyi rename to python/pyarrow-stubs/pyarrow/__lib/config.pyi diff --git a/python/pyarrow-stubs/pyarrow/device.pyi b/python/pyarrow-stubs/pyarrow/__lib/device.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/device.pyi rename to python/pyarrow-stubs/pyarrow/__lib/device.pyi diff --git a/python/pyarrow-stubs/pyarrow/error.pyi b/python/pyarrow-stubs/pyarrow/__lib/error.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/error.pyi rename to python/pyarrow-stubs/pyarrow/__lib/error.pyi diff --git a/python/pyarrow-stubs/pyarrow/io.pyi b/python/pyarrow-stubs/pyarrow/__lib/io.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/io.pyi rename to python/pyarrow-stubs/pyarrow/__lib/io.pyi diff --git a/python/pyarrow-stubs/pyarrow/memory.pyi b/python/pyarrow-stubs/pyarrow/__lib/memory.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/memory.pyi rename to python/pyarrow-stubs/pyarrow/__lib/memory.pyi diff --git a/python/pyarrow-stubs/pyarrow/scalar.pyi b/python/pyarrow-stubs/pyarrow/__lib/scalar.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/scalar.pyi rename to python/pyarrow-stubs/pyarrow/__lib/scalar.pyi diff --git a/python/pyarrow-stubs/pyarrow/table.pyi b/python/pyarrow-stubs/pyarrow/__lib/table.pyi similarity index 99% rename from python/pyarrow-stubs/pyarrow/table.pyi rename to python/pyarrow-stubs/pyarrow/__lib/table.pyi index 5bf6d94a686a..5bc11ccaffb7 100644 --- a/python/pyarrow-stubs/pyarrow/table.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/table.pyi @@ -58,11 +58,11 @@ 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 pyarrow._ipc import RecordBatchReader from .scalar import Int64Scalar, Scalar from .tensor import Tensor -from ._stubs_typing import NullableCollection -from ._types import DataType +from pyarrow._stubs_typing import NullableCollection +from pyarrow._types import DataType ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray diff --git a/python/pyarrow-stubs/pyarrow/tensor.pyi b/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/tensor.pyi rename to python/pyarrow-stubs/pyarrow/__lib/tensor.pyi diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi index e1c64d028bcf..873f69f08eb1 100644 --- a/python/pyarrow-stubs/pyarrow/lib.pyi +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -22,15 +22,15 @@ 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 +from pyarrow.__lib.array import * # noqa: F401,F403 +from pyarrow.__lib.builder import * # noqa: F401,F403 +from pyarrow.__lib.config import * # noqa: F401,F403 +from pyarrow.__lib.device import * # noqa: F401,F403 +from pyarrow.__lib.error import * # noqa: F401,F403 +from pyarrow.__lib.io import * # noqa: F401,F403 +from pyarrow.__lib.memory import * # noqa: F401,F403 +from pyarrow.__lib.scalar import * # noqa: F401,F403 +from pyarrow.__lib.table import * # noqa: F401,F403 +from pyarrow.__lib.tensor import * # noqa: F401,F403 def __getattr__(name: str) -> Any: ... diff --git a/python/pyproject.toml b/python/pyproject.toml index 2b713a39032a..e08254496ccd 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -177,7 +177,7 @@ exclude = [ # 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"] +include = ["pyarrow-stubs/pyarrow/__lib/array.pyi"] [tool.ty.overrides.rules] # Some array classes override methods with a different signature at runtime. invalid-method-override = "ignore" From ae3cad1a45a416792bd632e56c5534e0a4d61400 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Thu, 29 Oct 2026 20:36:01 +0100 Subject: [PATCH 38/93] Move remaining ipc/types.pyi to __lib/ --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 2 +- python/pyarrow-stubs/pyarrow/{_ipc.pyi => __lib/ipc.pyi} | 0 python/pyarrow-stubs/pyarrow/__lib/table.pyi | 4 ++-- python/pyarrow-stubs/pyarrow/__lib/tensor.pyi | 2 +- .../pyarrow-stubs/pyarrow/{_types.pyi => __lib/types.pyi} | 0 python/pyarrow-stubs/pyarrow/lib.pyi | 2 +- python/pyarrow-stubs/pyarrow/types.pyi | 2 +- python/scripts/update_stub_docstrings.py | 6 +----- 8 files changed, 7 insertions(+), 11 deletions(-) rename python/pyarrow-stubs/pyarrow/{_ipc.pyi => __lib/ipc.pyi} (100%) rename python/pyarrow-stubs/pyarrow/{_types.pyi => __lib/types.pyi} (100%) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index daf6df35f7c3..b64872dd76fb 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -52,7 +52,7 @@ import builtins from .scalar import Int64Scalar, Scalar from .device import DeviceAllocationType -from pyarrow._types import BaseExtensionType, DataType, Field, StructType +from .types import BaseExtensionType, DataType, Field, StructType from pyarrow._stubs_typing import NullableCollection def array( diff --git a/python/pyarrow-stubs/pyarrow/_ipc.pyi b/python/pyarrow-stubs/pyarrow/__lib/ipc.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/_ipc.pyi rename to python/pyarrow-stubs/pyarrow/__lib/ipc.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/table.pyi b/python/pyarrow-stubs/pyarrow/__lib/table.pyi index 5bc11ccaffb7..8461bd2690be 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/table.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/table.pyi @@ -58,11 +58,11 @@ from pyarrow.lib import Device, Field, MemoryManager, MemoryPool, Schema from .array import Array, StructArray, _PandasConvertible from .device import DeviceAllocationType from .io import Buffer -from pyarrow._ipc import RecordBatchReader +from .ipc import RecordBatchReader from .scalar import Int64Scalar, Scalar from .tensor import Tensor from pyarrow._stubs_typing import NullableCollection -from pyarrow._types import DataType +from .types import DataType ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray diff --git a/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi b/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi index 6a4a06eb989f..590ef0a9ac4b 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi @@ -26,7 +26,7 @@ from collections.abc import Sequence import numpy as np from pyarrow.lib import _Weakrefable -from pyarrow._types import DataType +from .types import DataType from scipy.sparse import coo_matrix, csc_matrix, csr_matrix from sparse import COO # type: ignore[import-untyped, import-not-found] diff --git a/python/pyarrow-stubs/pyarrow/_types.pyi b/python/pyarrow-stubs/pyarrow/__lib/types.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/_types.pyi rename to python/pyarrow-stubs/pyarrow/__lib/types.pyi diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi index 873f69f08eb1..8a1291188bec 100644 --- a/python/pyarrow-stubs/pyarrow/lib.pyi +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -21,7 +21,7 @@ from typing import Any -from pyarrow._types import * # noqa: F401,F403 +from pyarrow.__lib.types import * # noqa: F401,F403 from pyarrow.__lib.array import * # noqa: F401,F403 from pyarrow.__lib.builder import * # noqa: F401,F403 from pyarrow.__lib.config import * # noqa: F401,F403 diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/pyarrow/types.pyi index 1ba1e60965d2..609aafd7f7e4 100644 --- a/python/pyarrow-stubs/pyarrow/types.pyi +++ b/python/pyarrow-stubs/pyarrow/types.pyi @@ -31,7 +31,7 @@ else: import pyarrow.lib as lib -from pyarrow._types import ( +from pyarrow.__lib.types import ( BinaryType, BinaryViewType, BoolType, diff --git a/python/scripts/update_stub_docstrings.py b/python/scripts/update_stub_docstrings.py index 5ac424f138a0..982445bc3a56 100644 --- a/python/scripts/update_stub_docstrings.py +++ b/python/scripts/update_stub_docstrings.py @@ -170,10 +170,6 @@ def leave_FunctionDef(self, original_node, updated_node): return updated_node -LIB_MODULES = {"array", "builder", "compat", "config", "device", "error", "io", - "_ipc", "memory", "pandas_shim", "scalar", "table", "tensor", "_types"} - - def add_docstrings_to_stubs(stubs_dir): """Update all stub files in stubs_dir with docstrings from pyarrow runtime.""" stubs_dir = Path(stubs_dir) @@ -186,7 +182,7 @@ def add_docstrings_to_stubs(stubs_dir): continue module_name = stub_file.stem - if module_name in LIB_MODULES: + if stub_file.parent.name == "__lib": namespace = "lib" elif stub_file.parent.name in ("parquet", "interchange"): namespace = (stub_file.parent.name if module_name == "__init__" From ca5eb2e3d402a75c2a6ee7c59ee681116fa9f858 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Fri, 30 Oct 2026 17:51:57 +0100 Subject: [PATCH 39/93] Review type.pyi --- python/pyarrow-stubs/pyarrow/__lib/types.pyi | 54 +++++++++++--------- python/pyarrow-stubs/tests/test_types.py | 36 +++++++++++++ python/pyproject.toml | 7 +++ 3 files changed, 72 insertions(+), 25 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/types.pyi b/python/pyarrow-stubs/pyarrow/__lib/types.pyi index 9bc9420bc6d3..e7528788c1bc 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/types.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/types.pyi @@ -18,12 +18,12 @@ 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 +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] +from typing_extensions import CapsuleType, Self, deprecated from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit -from pyarrow.lib import ( # noqa: F401 +from pyarrow.lib import ( Array, Buffer, ChunkedArray, @@ -56,11 +56,11 @@ class DataType(_Weakrefable): ) -> 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: ... + @staticmethod + def _import_from_c(in_ptr: int) -> DataType: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> DataType: ... _PandasDtype: TypeAlias = ( type[np.generic] | np.dtype[Any] | pd.api.extensions.ExtensionDtype @@ -369,6 +369,10 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __iter__(self) -> Iterator[bytes]: ... def key(self, i: int) -> bytes: ... def value(self, i: int) -> bytes: ... + # The runtime yields generators instead of Mapping views + def keys(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def values(self) -> Iterator[bytes]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 + def items(self) -> Iterator[tuple[bytes, bytes]]: ... # type: ignore[override] # ty: ignore[invalid-method-override] # noqa: E501 def get_all(self, key: str | bytes) -> list[bytes]: ... def to_dict(self) -> dict[bytes, bytes]: ... @@ -386,18 +390,18 @@ class Field(_Weakrefable): def with_metadata( self, metadata: _KeyValueMetadataInput | KeyValueMetadata, - ) -> Self: ... - def remove_metadata(self) -> Self: ... + ) -> Field: ... + def remove_metadata(self) -> Field: ... def with_type(self, new_type: DataType) -> Field: ... - def with_name(self, name: str | bytes) -> Self: ... - def with_nullable(self, nullable: bool) -> Self: ... + def with_name(self, name: str | bytes) -> Field: ... + def with_nullable(self, nullable: bool) -> Field: ... def flatten(self) -> list[Field]: ... def _export_to_c(self, out_ptr: int) -> None: ... - @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: ... + @staticmethod + def _import_from_c(in_ptr: int) -> Field: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Field: ... _FieldTuple: TypeAlias = ( tuple[_FieldName, _DataTypeLike] @@ -462,11 +466,11 @@ class Schema(_Weakrefable): 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 _import_from_c(in_ptr: int) -> Schema: ... + def __arrow_c_schema__(self) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(schema: CapsuleType) -> Schema: ... def unify_schemas( schemas: Iterable[Schema], @@ -537,12 +541,12 @@ def large_list_view(value_type: DataType | Field) -> LargeListViewType: ... def map_( key_type: _DataTypeLike | Field, item_type: _DataTypeLike | Field, - keys_sorted: bool | None = None, + keys_sorted: bool = False, ) -> MapType: ... def dictionary( index_type: _IndexType | _DataTypeAliasInput, value_type: _DataTypeLike, - ordered: bool | None = None, + ordered: bool = False, ) -> DictionaryType: ... def struct( fields: _StructFieldsInput, @@ -577,7 +581,7 @@ def schema( fields: SupportsArrowSchema | _SchemaFieldsInput, metadata: _KeyValueMetadataInput | KeyValueMetadata = None, ) -> Schema: ... -def from_numpy_dtype(dtype: np.dtype[Any] | type | str) -> DataType: ... +def from_numpy_dtype(dtype: np.typing.DTypeLike) -> DataType: ... __all__ = [ "_Weakrefable", diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 2133cb4de874..6b74ba77b15a 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -16,10 +16,12 @@ # under the License. # mypy: warn-unused-ignores +from collections.abc import Iterator from typing import TYPE_CHECKING, Literal import pyarrow.lib as lib import pyarrow.types as types +from typing_extensions import CapsuleType from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -139,6 +141,15 @@ def test_schema() -> None: check(assert_type(schema[0], lib.Field), lib.Field) +def test_key_value_metadata() -> None: + # The runtime yields generators instead of Mapping views (duplicated keys) + metadata = lib.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) + check(assert_type(metadata.keys(), Iterator[bytes]), Iterator) + check(assert_type(metadata.values(), Iterator[bytes]), Iterator) + check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), + Iterator) + + def test_type_predicates_narrow() -> None: lst = _erase(lib.list_(lib.int32())) assert types.is_list(lst) @@ -179,6 +190,31 @@ def test_type_predicates_narrow() -> None: check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) +def test_import_from_c_capsule() -> None: + capsule = lib.int32().__arrow_c_schema__() + check(assert_type(capsule, CapsuleType), CapsuleType) + + # The import methods are static and return whatever type the C schema + # holds, not the class they are called on. + imported = lib.ListType._import_from_c_capsule(capsule) + check(assert_type(imported, lib.DataType), lib.DataType) + assert not isinstance(imported, lib.ListType) + + field = lib.field("a", lib.int32()) + field_capsule = field.__arrow_c_schema__() + check(assert_type(field_capsule, CapsuleType), CapsuleType) + + imported_field = lib.Field._import_from_c_capsule(field_capsule) + check(assert_type(imported_field, lib.Field), lib.Field) + + schema = lib.schema([field]) + schema_capsule = schema.__arrow_c_schema__() + check(assert_type(schema_capsule, CapsuleType), CapsuleType) + + imported_schema = lib.Schema._import_from_c_capsule(schema_capsule) + check(assert_type(imported_schema, lib.Schema), lib.Schema) + + def test_invalid_usage() -> None: if TYPE_CHECKING_INVALID_USAGE: # A Field needs a type unless it's created from an Arrow C schema. diff --git a/python/pyproject.toml b/python/pyproject.toml index e08254496ccd..19a5f429562f 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -174,6 +174,13 @@ exclude = [ "scripts", ] +# numpy, pandas and scipy are optional dependencies. A `ty: ignore` comment on +# their imports would be reported as unused whenever they are installed. +[[tool.ty.overrides]] +include = ["pyarrow-stubs/pyarrow/**"] +[tool.ty.overrides.rules] +unresolved-import = "ignore" + # ty doesn't honor `# type: ignore[]` comments, so disable the # rules that those comments suppress for mypy. [[tool.ty.overrides]] From d18e949a870511fcb515b200e7e810dc82929167 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Sat, 31 Oct 2026 18:52:52 +0100 Subject: [PATCH 40/93] Add typing_extensions to dependency-groups test --- python/pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 19a5f429562f..2adf0fa3f502 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -62,7 +62,8 @@ test = [ 'hypothesis', 'cffi', 'pytz', - 'pandas' + 'pandas', + 'typing_extensions>=4.12.0', ] [tool.numpydoc_validation] From cf32c2e41891a0eff7652bd37d08a1474bb38ae9 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Sat, 31 Oct 2026 22:44:25 +0100 Subject: [PATCH 41/93] Array staticmethod _import_from_c_capsule/device and test --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 16 ++++++++-------- python/pyarrow-stubs/tests/test_array.py | 6 ++++++ 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index b64872dd76fb..560ec10d8d39 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -227,17 +227,17 @@ class Array(_PandasConvertible): 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: ... + @staticmethod + def _import_from_c(in_ptr: int, type: int | DataType) -> Array: ... def __arrow_c_array__(self, requested_schema=None) -> Any: ... - @classmethod - def _import_from_c_capsule(cls, schema_capsule, array_capsule) -> Self: ... + @staticmethod + def _import_from_c_capsule(schema_capsule, array_capsule) -> Array: ... 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: ... + @staticmethod + def _import_from_c_device(in_ptr: int, type: DataType | int) -> Array: ... def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... - @classmethod - def _import_from_c_device_capsule(cls, schema_capsule, array_capsule) -> Self: ... + @staticmethod + def _import_from_c_device_capsule(schema_capsule, array_capsule) -> Array: ... def __dlpack__(self, stream: int | None = None) -> Any: ... def __dlpack_device__(self) -> tuple[int, int]: ... @property diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index d41826eed8dc..4a85666f4403 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -51,6 +51,12 @@ def test_array() -> None: lib.Int64Array) +def test_import_from_c_capsule() -> None: + capsules = lib.array([1, 2]).__arrow_c_array__() + imported = lib.BooleanArray._import_from_c_capsule(*capsules) + check(assert_type(imported, 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) From 516e8a9a4089b16f656b9f06045489bffc4a1e83 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Mon, 2 Nov 2026 09:38:46 +0100 Subject: [PATCH 42/93] Add hash stubs to ExtensionType and KeyValueMetadata --- python/pyarrow-stubs/pyarrow/__lib/types.pyi | 7 +++---- python/pyarrow-stubs/tests/test_types.py | 8 +++++++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/types.pyi b/python/pyarrow-stubs/pyarrow/__lib/types.pyi index e7528788c1bc..e40092f36c85 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/types.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/types.pyi @@ -16,7 +16,7 @@ # under the License. from collections.abc import Iterable, Iterator, Mapping, Sequence -from typing import Any, Literal, TypeAlias, overload +from typing import Any, ClassVar, Literal, TypeAlias, overload import numpy as np # type: ignore[import-not-found] import pandas as pd # type: ignore[import-not-found,import-untyped] @@ -50,7 +50,6 @@ class DataType(_Weakrefable): 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: ... @@ -259,6 +258,7 @@ class ExtensionType(BaseExtensionType): def __arrow_ext_deserialize__( cls, storage_type: DataType, serialized: bytes ) -> Self: ... + __hash__: ClassVar[None] # type: ignore[assignment] class FixedShapeTensorType(BaseExtensionType): @property @@ -362,6 +362,7 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, **kwargs: str | bytes, ) -> None: ... + __hash__: ClassVar[None] # type: ignore[assignment] def equals(self, other: KeyValueMetadata) -> bool: ... def __len__(self) -> int: ... def __contains__(self, /, __key: object) -> bool: ... @@ -378,7 +379,6 @@ class KeyValueMetadata(_Metadata, Mapping[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 @@ -424,7 +424,6 @@ class Schema(_Weakrefable): _field = __getitem__ def __iter__(self) -> Iterator[Field]: ... - def __hash__(self) -> int: ... def __sizeof__(self) -> int: ... @property def pandas_metadata(self) -> dict[str, Any] | None: ... diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 6b74ba77b15a..5e439e754c8d 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -16,7 +16,7 @@ # under the License. # mypy: warn-unused-ignores -from collections.abc import Iterator +from collections.abc import Hashable, Iterator from typing import TYPE_CHECKING, Literal import pyarrow.lib as lib @@ -122,6 +122,11 @@ def test_types_enum() -> None: == types.TypesEnum.INTERVAL_MONTH_DAY_NANO) +def test_data_type_hashable() -> None: + data_type: Hashable = lib.int32() + check(assert_type(hash(data_type), int), int) + + def test_field() -> None: field = lib.field("a", lib.int32()) check(assert_type(field, lib.Field), lib.Field) @@ -221,3 +226,4 @@ def test_invalid_usage() -> None: lib.field("a") # type: ignore[call-overload] lib.list_(1) # type: ignore[call-overload] lib.decimal128("10") # type: ignore[arg-type] + _: Hashable = lib.KeyValueMetadata({b"a": b"1"}) # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 From 401c3c9bbf9847af2d3b8a90288a13df3f6ce92c Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Mon, 2 Nov 2026 22:13:39 +0100 Subject: [PATCH 43/93] Simplified array inputs, ignore opt imports, fix DLPack signature and test --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 21 +++++++++++--------- python/pyarrow-stubs/tests/test_array.py | 11 ++++++++++ 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index 560ec10d8d39..adbf89b9174c 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -26,8 +26,8 @@ else: from typing import Any, Literal -import numpy as np -import pandas as pd +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] from pyarrow._compute import CastOptions from pyarrow._stubs_typing import ( @@ -56,8 +56,7 @@ from .types import BaseExtensionType, DataType, Field, StructType from pyarrow._stubs_typing import NullableCollection def array( - values: NullableCollection - | Iterable[Any] + values: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray | SupportsPyArrowArray, @@ -69,10 +68,7 @@ def array( memory_pool: MemoryPool | None = None, ) -> ArrayLike: ... def asarray( - values: NullableCollection - | Iterable[Any] - | SupportsArrowArray - | SupportsArrowDeviceArray, + values: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, type: Any | None = None, ) -> ArrayLike: ... def nulls( @@ -238,7 +234,14 @@ class Array(_PandasConvertible): def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... @staticmethod def _import_from_c_device_capsule(schema_capsule, array_capsule) -> Array: ... - def __dlpack__(self, stream: int | None = None) -> Any: ... + def __dlpack__( + self, + *, + stream: int | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> Any: ... def __dlpack_device__(self) -> tuple[int, int]: ... @property def device_type(self) -> DeviceAllocationType: ... diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index 4a85666f4403..d454d6431f27 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -19,6 +19,7 @@ from typing import Any, cast import pyarrow.lib as lib +from typing_extensions import CapsuleType from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -51,6 +52,16 @@ def test_array() -> None: lib.Int64Array) +def test_array_dlpack() -> None: + arr = cast(lib.Int64Array, lib.array([1, 2])) + capsule = arr.__dlpack__( + max_version=(1, 0), + dl_device=arr.__dlpack_device__(), + copy=False, + ) + check(assert_type(capsule, Any), CapsuleType) + + def test_import_from_c_capsule() -> None: capsules = lib.array([1, 2]).__arrow_c_array__() imported = lib.BooleanArray._import_from_c_capsule(*capsules) From 368a6dc22caa3c8621a6ee764951c67f74e04518 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 3 Nov 2026 18:54:24 +0100 Subject: [PATCH 44/93] Change Array.equals Any to Array and add tests --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 2 +- python/pyarrow-stubs/tests/test_array.py | 6 ++++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index adbf89b9174c..e9d48a74ea2c 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -179,7 +179,7 @@ class Array(_PandasConvertible): ) -> str: ... format = to_string - def equals(self, other: Array | Any) -> bool: ... + def equals(self, other: Array) -> bool: ... def __len__(self) -> int: ... def is_null(self, *, nan_is_null: bool = False) -> BooleanArray: ... def is_nan(self) -> BooleanArray: ... diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index d454d6431f27..6c41e5b19832 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -51,6 +51,10 @@ def test_array() -> None: lib.Array), lib.Int64Array) + other = check(cast(lib.StringArray, lib.array(["a", "b", None])), lib.StringArray) + check(assert_type(arr.equals(arr), bool), bool) + check(assert_type(arr.equals(other), bool), bool) + def test_array_dlpack() -> None: arr = cast(lib.Int64Array, lib.array([1, 2])) @@ -128,3 +132,5 @@ def test_invalid_usage() -> None: if TYPE_CHECKING_INVALID_USAGE: arr.cast(1) # type: ignore[arg-type] arr.slice("1") # type: ignore[arg-type] + arr.equals(object()) # type: ignore[arg-type] + arr.equals(None) # type: ignore[arg-type] From 91c3f0bd796783c7b439064b341cce6d4d9b639d Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Thu, 5 Nov 2026 13:36:05 +0100 Subject: [PATCH 45/93] Review table.pyi --- .pre-commit-config.yaml | 1 + python/pyarrow-stubs/pyarrow/__lib/array.pyi | 7 +- python/pyarrow-stubs/pyarrow/__lib/table.pyi | 288 ++++++++++-------- .../pyarrow-stubs/pyarrow/_stubs_typing.pyi | 2 + python/pyarrow-stubs/tests/test_table.py | 81 ++++- 5 files changed, 250 insertions(+), 129 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 3b6ea465b5de..bafa6097d574 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -202,6 +202,7 @@ repos: # Provides the modules that aren't annotated yet, e.g. pyarrow.compute. - pyarrow==26.0.0 - pyright==1.1.414 + - pytest==9.0.2 - scipy-stubs==1.17.1.5 - ty==0.0.92 - types-cffi==2.1.0.20260827 diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index e9d48a74ea2c..b0025637e512 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -34,6 +34,7 @@ from pyarrow._stubs_typing import ( ArrayLike, Indices, Mask, + MapsAsPydicts, Order, SupportsArrowArray, SupportsArrowDeviceArray, @@ -119,17 +120,13 @@ class _PandasConvertible(_Weakrefable): safe: bool = True, split_blocks: bool = False, self_destruct: bool = False, - maps_as_pydicts: Literal["None", "lossy", "strict"] | None = None, + maps_as_pydicts: MapsAsPydicts | None = None, types_mapper: Any = None, # Callable[[DataType], ExtensionDtype | None] | None coerce_temporal_nanoseconds: bool = False, ) -> pd.Series | pd.DataFrame: ... class Array(_PandasConvertible): 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, diff --git a/python/pyarrow-stubs/pyarrow/__lib/table.pyi b/python/pyarrow-stubs/pyarrow/__lib/table.pyi index 8461bd2690be..981f8991f273 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/table.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/table.pyi @@ -25,14 +25,15 @@ 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 +from collections.abc import Iterable, Iterator, Sequence, Mapping +from typing import Any, ClassVar, Literal +from typing_extensions import CapsuleType, deprecated import builtins -import numpy as np -import pandas as pd +import numpy as np # type: ignore[import-not-found] +import pandas as pd # type: ignore[import-not-found,import-untyped] -from numpy.typing import NDArray +from numpy.typing import NDArray # type: ignore[import-not-found] from pyarrow._compute import ( CastOptions, CountOptions, @@ -44,7 +45,9 @@ from pyarrow._compute import ( from pyarrow._stubs_typing import ( Indices, Mask, + MapsAsPydicts, NullEncoding, + NullPlacement, NullSelectionBehavior, Order, SupportsArrowArray, @@ -61,7 +64,6 @@ from .io import Buffer from .ipc import RecordBatchReader from .scalar import Int64Scalar, Scalar from .tensor import Tensor -from pyarrow._stubs_typing import NullableCollection from .types import DataType ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray @@ -76,6 +78,7 @@ _Aggregation: TypeAlias = Literal[ "distinct", "first", "first_last", + "kurtosis", "last", "list", "max", @@ -83,7 +86,9 @@ _Aggregation: TypeAlias = Literal[ "min", "min_max", "one", + "pivot_wider", "product", + "skew", "stddev", "sum", "tdigest", @@ -99,6 +104,7 @@ _AggregationPrefixed: TypeAlias = Literal[ "hash_distinct", "hash_first", "hash_first_last", + "hash_kurtosis", "hash_last", "hash_list", "hash_max", @@ -106,7 +112,9 @@ _AggregationPrefixed: TypeAlias = Literal[ "hash_min", "hash_min_max", "hash_one", + "hash_pivot_wider", "hash_product", + "hash_skew", "hash_stddev", "hash_sum", "hash_tdigest", @@ -121,6 +129,13 @@ AggregateOptions: TypeAlias = ( | FunctionOptions ) +ChunkedArrayConstructorInput: TypeAlias = ( + Iterable[Array | Iterable[Any] | SupportsArrowStream | SupportsArrowArray] + | Array + | SupportsArrowArray + | SupportsArrowStream +) + UnarySelector: TypeAlias = str NullarySelector: TypeAlias = tuple[()] NarySelector: TypeAlias = list[str] | tuple[str, ...] @@ -131,10 +146,6 @@ class ChunkedArray(_PandasConvertible): 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 @@ -146,9 +157,10 @@ class ChunkedArray(_PandasConvertible): window: int = 5, container_window: int = 2, skip_new_lines: bool = False, + element_size_limit: int = 100, ) -> str: ... - - format = to_string + @deprecated("Use 'to_string' instead") + def format(self, **kwargs) -> str: ... def validate(self, *, full: bool = False) -> None: ... @property def null_count(self) -> int: ... @@ -156,33 +168,36 @@ class ChunkedArray(_PandasConvertible): 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 __getitem__( + self, key: int | np.integer | builtins.slice + ) -> Scalar | ChunkedArray: ... + __hash__: ClassVar[None] # type: ignore[assignment] def is_null(self, *, nan_is_null: bool = False) -> ChunkedArray: ... def is_nan(self) -> ChunkedArray: ... def is_valid(self) -> ChunkedArray: ... def cast( self, - target_type: DataType | str | None, - safe: bool = True, + target_type: DataType | str | None = None, + safe: bool | None = None, 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 fill_null(self, fill_value: Any) -> ChunkedArray: ... + def equals(self, other: ChunkedArray | None) -> bool: ... def to_numpy(self, zero_copy_only: bool = False) -> np.ndarray: ... def __array__( self, dtype: np.dtype | None = None, copy: bool | None = None ) -> np.ndarray: ... - def dictionary_encode(self, null_encoding: NullEncoding = "mask") -> Self: ... + def dictionary_encode( + self, null_encoding: NullEncoding = "mask" + ) -> ChunkedArray: ... def flatten(self, memory_pool: MemoryPool | None = None) -> list[ChunkedArray]: ... def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Array: ... def unique(self) -> Array: ... def value_counts(self) -> StructArray: ... - def slice(self, offset: int = 0, length: int | None = None) -> Self: ... + def slice(self, offset: int = 0, length: int | None = None) -> ChunkedArray: ... def filter( self, mask: Mask, null_selection_behavior: NullSelectionBehavior = "drop" - ) -> Self: ... + ) -> ChunkedArray: ... def index( self, value: Any, @@ -191,35 +206,39 @@ class ChunkedArray(_PandasConvertible): *, 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: ... + def take(self, indices: Indices) -> ChunkedArray: ... + def drop_null(self) -> ChunkedArray: ... + def sort( + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + **kwargs, + ) -> ChunkedArray: ... + def unify_dictionaries( + self, memory_pool: MemoryPool | None = None + ) -> ChunkedArray: ... @property def num_chunks(self) -> int: ... def chunk(self, i: int) -> Array: ... @property def chunks(self) -> list[Array]: ... - def iterchunks(self) -> Generator[Array, None, None]: ... + def iterchunks(self) -> Iterator[Array]: ... def __iter__(self) -> Iterator[Scalar]: ... def to_pylist( self, *, - maps_as_pydicts: Literal["lossy", "strict"] | None = None, + maps_as_pydicts: MapsAsPydicts | None = None, ) -> list[Any]: ... - def __arrow_c_stream__(self, requested_schema=None) -> Any: ... - @classmethod - def _import_from_c_capsule(cls, stream) -> Self: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule(stream: CapsuleType) -> ChunkedArray: ... @property def is_cpu(self) -> bool: ... def chunked_array( - arrays: Iterable[NullableCollection] - | Iterable[Iterable[Any] | SupportsArrowStream | SupportsArrowArray] - | Iterable[Array] - | Array - | SupportsArrowArray - | SupportsArrowStream, + arrays: ChunkedArrayConstructorInput, type: DataType | str | None = None, ) -> ChunkedArray: ... @@ -242,7 +261,13 @@ class _Tabular(_PandasConvertible): @classmethod def from_pydict( cls, - mapping: Mapping[Any, ArrayOrChunkedArray | list[Any] | np.ndarray | range], + mapping: Mapping[ + Any, + ArrayOrChunkedArray + | Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray, + ], schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> Self: ... @@ -253,7 +278,7 @@ class _Tabular(_PandasConvertible): schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> Self: ... - def itercolumns(self) -> Generator[ArrayOrChunkedArray, None, None]: ... + def itercolumns(self) -> Iterator[ArrayOrChunkedArray]: ... @property def num_columns(self) -> int: ... @property @@ -262,9 +287,11 @@ class _Tabular(_PandasConvertible): 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 sort_by( + self, + sorting: str | Sequence[tuple[str, Order] | tuple[str, Order, NullPlacement]], + **kwargs, + ) -> Self: ... def take(self, indices: Indices) -> Self: ... def filter( self, @@ -272,10 +299,10 @@ class _Tabular(_PandasConvertible): null_selection_behavior: NullSelectionBehavior = "drop", ) -> Self: ... def to_pydict( - self, *, maps_as_pydicts: Literal["lossy", "strict"] | None = None + self, *, maps_as_pydicts: MapsAsPydicts | None = None ) -> dict[str, list[Any]]: ... def to_pylist( - self, *, maps_as_pydicts: Literal["lossy", "strict"] | None = None + self, *, maps_as_pydicts: MapsAsPydicts | None = None ) -> list[dict[str, Any]]: ... def to_string( self, *, show_metadata: bool = False, preview_cols: int = 0 @@ -298,7 +325,7 @@ class RecordBatch(_Tabular): | dict[bytes, str] | dict[str, bytes] | None = None, - ) -> Self: ... + ) -> RecordBatch: ... @property def num_columns(self) -> int: ... @property @@ -310,35 +337,36 @@ class RecordBatch(_Tabular): 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: ... + self, i: int, field_: str | Field, column: ArrayOrChunkedArray | list[Any] + ) -> RecordBatch: ... + def remove_column(self, i: int) -> RecordBatch: ... + def set_column( + self, i: int, field_: str | Field, column: Array | list + ) -> RecordBatch: ... + def rename_columns(self, names: list[str] | dict[str, str]) -> RecordBatch: ... def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... - def slice(self, offset: int = 0, length: int | None = None) -> Self: ... - def equals(self, other: Self | Any, check_metadata: bool = False) -> bool: ... + def slice(self, offset: int = 0, length: int | None = None) -> RecordBatch: ... + def equals(self, other: object, check_metadata: bool = False) -> bool: ... def select( self, columns: Iterable[str] | Iterable[int] | NDArray[np.str_] - ) -> Self: ... + ) -> RecordBatch: ... def cast( self, target_schema: Schema, safe: bool | None = None, options: CastOptions | None = None, - ) -> Self: ... - @classmethod + ) -> RecordBatch: ... + @staticmethod def from_arrays( - cls, - arrays: Iterable[Any], - names: list[str] | tuple[str, ...] | None = None, + arrays: list[Any], + names: Sequence[str] | None = None, schema: Schema | None = None, metadata: Mapping[bytes, bytes] | Mapping[str, str] | Mapping[bytes, str] | Mapping[str, bytes] | None = None, - ) -> Self: ... + ) -> RecordBatch: ... @classmethod def from_pandas( cls, @@ -346,10 +374,10 @@ class RecordBatch(_Tabular): 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: ... + columns: Iterable[str | int] | None = None, + ) -> RecordBatch: ... + @staticmethod + def from_struct_array(struct_array: StructArray) -> RecordBatch: ... def to_struct_array(self) -> StructArray: ... def to_tensor( self, @@ -358,25 +386,33 @@ class RecordBatch(_Tabular): 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: ... + @staticmethod + def _import_from_c(in_ptr: int, schema: Schema | int) -> RecordBatch: ... + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... + @staticmethod + def _import_from_c_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> RecordBatch: ... def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... - @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: ... + @staticmethod + def _import_from_c_device(in_ptr: int, schema: Schema | int) -> RecordBatch: ... + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs + ): ... + @staticmethod + def _import_from_c_device_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> RecordBatch: ... @property def device_type(self) -> DeviceAllocationType: ... @property def is_cpu(self) -> bool: ... - def copy_to(self, destination: MemoryManager | Device) -> Self: ... - -def table_to_blocks(options, table: Table, categories, extension_columns): ... + def copy_to(self, destination: MemoryManager | Device) -> RecordBatch: ... JoinType: TypeAlias = Literal[ "left semi", @@ -391,10 +427,8 @@ JoinType: TypeAlias = Literal[ 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 slice(self, offset: int = 0, length: int | None = None) -> Table: ... + def select(self, columns: Iterable[str | int]) -> Table: ... def replace_schema_metadata( self, metadata: dict[str, str] @@ -402,17 +436,17 @@ class Table(_Tabular): | 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: ... + ) -> Table: ... + def flatten(self, memory_pool: MemoryPool | None = None) -> Table: ... + def combine_chunks(self, memory_pool: MemoryPool | None = None) -> Table: ... + def unify_dictionaries(self, memory_pool: MemoryPool | None = None) -> Table: ... + def equals(self, other: Table | None, check_metadata: bool = False) -> bool: ... def cast( self, target_schema: Schema, safe: bool | None = None, options: CastOptions | None = None, - ) -> Self: ... + ) -> Table: ... @classmethod def from_pandas( cls, @@ -420,29 +454,34 @@ class Table(_Tabular): schema: Schema | None = None, preserve_index: bool | None = None, nthreads: int | None = None, - columns: Sequence[str | int] | None = None, + columns: Iterable[str | int] | None = None, safe: bool = True, - ) -> Self: ... - @classmethod + ) -> Table: ... + @staticmethod def from_arrays( - cls, - arrays: Collection[ArrayOrChunkedArray | Collection[NDArray[Any]] | list[Any]], - names: list[str] | tuple[str, ...] | None = None, + arrays: Sequence[ArrayOrChunkedArray | Iterable[Any]], + names: Sequence[str] | None = None, schema: Schema | None = None, metadata: Mapping[bytes, bytes] | Mapping[str, str] | Mapping[bytes, str] | Mapping[str, bytes] | None = None, - ) -> Self: ... - @classmethod - def from_struct_array(cls, struct_array: StructArray | ChunkedArray) -> Self: ... + ) -> Table: ... + @staticmethod + def from_struct_array(struct_array: StructArray | ChunkedArray) -> Table: ... def to_struct_array(self, max_chunksize: int | None = None) -> ChunkedArray: ... - @classmethod + @staticmethod def from_batches( - cls, batches: Iterable[RecordBatch], schema: Schema | None = None - ) -> Self: ... + batches: Iterable[RecordBatch], schema: Schema | None = None + ) -> Table: ... def to_batches(self, max_chunksize: int | None = None) -> list[RecordBatch]: ... + def to_tensor( + self, + null_to_nan: bool = False, + row_major: bool = True, + memory_pool: MemoryPool | None = None, + ) -> Tensor: ... def to_reader(self, max_chunksize: int | None = None) -> RecordBatchReader: ... @property def schema(self) -> Schema: ... @@ -455,20 +494,26 @@ class Table(_Tabular): 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: ... + self, + i: int, + field_: str | Field, + column: ChunkedArray | ChunkedArrayConstructorInput, + ) -> Table: ... + def remove_column(self, i: int) -> Table: ... 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: ... + self, + i: int, + field_: str | Field, + column: ChunkedArray | ChunkedArrayConstructorInput, + ) -> Table: ... + def rename_columns(self, names: list[str] | dict[str, str]) -> Table: ... + def drop(self, columns: str | list[str]) -> Table: ... def group_by( self, keys: str | list[str], use_threads: bool = True ) -> TableGroupBy: ... def join( self, - right_table: Self, + right_table: Table, keys: str | list[str], right_keys: str | list[str] | None = None, join_type: JoinType = "left outer", @@ -476,23 +521,26 @@ class Table(_Tabular): right_suffix: str | None = None, coalesce_keys: bool = True, use_threads: bool = True, - ) -> Self: ... + filter_expression: Expression | None = None, + ) -> Table: ... def join_asof( self, - right_table: Self, + right_table: Table, on: str, by: str | list[str], tolerance: int, right_on: str | list[str] | None = None, right_by: str | list[str] | None = None, - ) -> Self: ... - def __arrow_c_stream__(self, requested_schema=None): ... + ) -> Table: ... + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: ... @property def is_cpu(self) -> bool: ... def record_batch( - data: Mapping[str, list[Any] | Array] - | Collection[Array | ChunkedArray | list[Any]] + data: dict[str, Any] + | list[Any] | pd.DataFrame | SupportsArrowArray | SupportsArrowDeviceArray, @@ -501,13 +549,13 @@ def record_batch( metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> RecordBatch: ... def table( - data: Collection[ArrayOrChunkedArray | list[Any] | range | str] + data: dict[str, Any] + | list[Any] + | tuple[Any, ...] | pd.DataFrame | SupportsArrowArray | SupportsArrowStream - | SupportsArrowDeviceArray - | Mapping[str, list[Any] | Array | ChunkedArray | range] - | Mapping[str, Any], + | SupportsArrowDeviceArray, names: list[str] | Schema | None = None, schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, @@ -533,9 +581,6 @@ class TableGroupBy: | 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 @@ -546,13 +591,10 @@ __all__ = [ "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/_stubs_typing.pyi b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi index 7a2d258f76b0..60d7b9e6a5b4 100644 --- a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi +++ b/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi @@ -50,7 +50,9 @@ JoinType: TypeAlias = Literal[ Compression: TypeAlias = Literal[ "gzip", "bz2", "brotli", "lz4", "lz4_frame", "lz4_raw", "zstd", "snappy" ] +MapsAsPydicts: TypeAlias = Literal["lossy", "strict"] NullEncoding: TypeAlias = Literal["mask", "encode"] +NullPlacement: TypeAlias = Literal["at_start", "at_end"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index 9544f44be794..01069fa3d8ad 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -16,6 +16,10 @@ # under the License. # mypy: warn-unused-ignores +from types import MappingProxyType + +import pytest + import pyarrow.lib as lib from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -44,9 +48,29 @@ def test_table() -> None: 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.select(["a", 1]), lib.Table), lib.Table) + check(assert_type(table.equals(None), bool), bool) 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) + check(assert_type(lib.Table.from_arrays([(1, 2)], names=["a"]), lib.Table), + lib.Table) + + chunked = lib.chunked_array([[1, 2]]) + check(assert_type(table.add_column(0, "c", chunked), lib.Table), lib.Table) + check(assert_type(table.add_column(0, "c", lib.array([1, 2])), lib.Table), + lib.Table) + check(assert_type(table.add_column(0, "c", [[1], [2]]), lib.Table), + lib.Table) + check(assert_type(table.add_column(0, "c", [lib.array([1, 2])]), lib.Table), + lib.Table) + + table = lib.table(([1],), names=["a"]) + check(assert_type(table, lib.Table), lib.Table) + + chunked_struct = lib.chunked_array([[{"x": 1}]]) + check(assert_type(lib.Table.from_struct_array(chunked_struct), lib.Table), + lib.Table) # The reader isn't annotated yet, but must not reject valid usage. check(table.to_reader().read_all(), lib.Table) @@ -55,7 +79,7 @@ def test_table() -> None: def test_record_batch() -> None: - batch = lib.record_batch({"a": [1, 2]}) + batch = lib.record_batch({"a": [1, 2]}) # from a dict 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), @@ -64,10 +88,65 @@ def test_record_batch() -> None: check(assert_type(lib.Table.from_batches([batch]), lib.Table), lib.Table) check(assert_type(lib.concat_batches([batch, batch]), lib.RecordBatch), lib.RecordBatch) + check(assert_type(batch.add_column(0, "b", lib.array([3, 4])), + lib.RecordBatch), + lib.RecordBatch) + check(assert_type(batch.add_column(0, "b", [3, 4]), lib.RecordBatch), + lib.RecordBatch) + check(assert_type(batch.equals(batch), bool), bool) + # Any object can be compared, a non-batch returns False + check(assert_type(batch.equals(42), bool), bool) + + batch = lib.record_batch([lib.array([1])], names=["a"]) # from a list + check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + + struct_array = batch.to_struct_array() + check(assert_type(lib.RecordBatch.from_struct_array(struct_array), + lib.RecordBatch), + lib.RecordBatch) + + +def test_record_batch_from_pandas() -> None: + pd = pytest.importorskip("pandas") + df = pd.DataFrame({"a": [1], "b": [2]}) + # Columns only need to be iterable, e.g. a generator. + batch = lib.RecordBatch.from_pandas(df, columns=(c for c in ["a"])) + check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + + +def test_chunked_array() -> None: + chunked_array = lib.chunked_array([[1, 2], [None]]) + check(assert_type(chunked_array, lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(chunked_array.cast(lib.int32(), safe=None), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(chunked_array.fill_null(0), lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(chunked_array.sort(null_placement="at_start"), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(chunked_array.equals(None), bool), bool) + + arr = lib.array([1, 2]) + check(assert_type(lib.chunked_array(arr), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(lib.chunked_array([arr, [3]]), lib.ChunkedArray), + lib.ChunkedArray) def test_invalid_usage() -> None: table = lib.table({"a": [1, 2]}) + columns = {"a": lib.array([1])}.values() # dict as non-sequence example + if TYPE_CHECKING_INVALID_USAGE: table.column(1.5) # type: ignore[arg-type] table.slice("0") # type: ignore[arg-type] + lib.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] + + +def test_record_batch_invalid_usage() -> None: + arr = lib.array([1]) + chunked = lib.chunked_array([[{"a": 1}]]) + if TYPE_CHECKING_INVALID_USAGE: + lib.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] + lib.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] + lib.record_batch((arr,), names=["a"]) # type: ignore[arg-type] + # Only a `dict` is passed on to `from_pydict`, other mappings raise + lib.record_batch(MappingProxyType({"a": [1]})) # type: ignore[arg-type] From 1589c4cac641a063111685015fcafe43de2106dd Mon Sep 17 00:00:00 2001 From: Alenka Frim Date: Fri, 6 Nov 2026 13:49:07 +0100 Subject: [PATCH 46/93] Update ci/scripts/python_test_type_annotations.sh Co-authored-by: Joris Van den Bossche --- ci/scripts/python_test_type_annotations.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ci/scripts/python_test_type_annotations.sh b/ci/scripts/python_test_type_annotations.sh index 96fe9a3135de..ab7dddf66be2 100755 --- a/ci/scripts/python_test_type_annotations.sh +++ b/ci/scripts/python_test_type_annotations.sh @@ -38,5 +38,5 @@ 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. +# the expected annotated types against the runtime types. pytest -r s pyarrow-stubs/tests From f6892cda35417e447deb53a9ba0cda8fc5f05d79 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Fri, 6 Nov 2026 13:59:01 +0100 Subject: [PATCH 47/93] Remove sys.version_info >= (3, 11) check --- python/pyarrow-stubs/tests/_check.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/python/pyarrow-stubs/tests/_check.py b/python/pyarrow-stubs/tests/_check.py index 596129874188..7c02e1509c6e 100644 --- a/python/pyarrow-stubs/tests/_check.py +++ b/python/pyarrow-stubs/tests/_check.py @@ -34,13 +34,8 @@ 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 +from typing_extensions import assert_type __all__ = ["TYPE_CHECKING_INVALID_USAGE", "assert_type", "check"] From a1163e251e890f2c66c5c41bbeb446930c526407 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Fri, 6 Nov 2026 14:24:24 +0100 Subject: [PATCH 48/93] Assert struct array --- python/pyarrow-stubs/tests/test_table.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index 01069fa3d8ad..032181e25d4e 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -101,6 +101,7 @@ def test_record_batch() -> None: check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) struct_array = batch.to_struct_array() + check(assert_type(struct_array, lib.StructArray), lib.StructArray) check(assert_type(lib.RecordBatch.from_struct_array(struct_array), lib.RecordBatch), lib.RecordBatch) From e51b18e2bc6449e4ab6cabb0e554a27b8e0afaf2 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Fri, 6 Nov 2026 14:52:25 +0100 Subject: [PATCH 49/93] Add test for list_ --- python/pyarrow-stubs/tests/test_types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 5e439e754c8d..260dd86f8a74 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -84,7 +84,8 @@ def test_list_types() -> None: 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. + variable = lib.list_(lib.int32(), -1) + check(assert_type(variable, lib.ListType), lib.ListType) fixed = lib.list_(lib.int32(), 2) check(assert_type(fixed, lib.ListType | lib.FixedSizeListType), lib.FixedSizeListType) From d2973396a921aa98887b30e146a9ada12bed5acb Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Fri, 6 Nov 2026 17:31:51 +0100 Subject: [PATCH 50/93] Fix CI ty check failure latest v0.99 --- .pre-commit-config.yaml | 2 +- python/pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bafa6097d574..a086779210df 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -204,7 +204,7 @@ repos: - pyright==1.1.414 - pytest==9.0.2 - scipy-stubs==1.17.1.5 - - ty==0.0.92 + - ty==0.0.99 - types-cffi==2.1.0.20260827 - types-psutil==7.2.2.20260906 - types-python-dateutil==2.9.0.20261021 diff --git a/python/pyproject.toml b/python/pyproject.toml index 2adf0fa3f502..5d01b0b307ec 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -180,7 +180,7 @@ exclude = [ [[tool.ty.overrides]] include = ["pyarrow-stubs/pyarrow/**"] [tool.ty.overrides.rules] -unresolved-import = "ignore" +import-not-found = "ignore" # ty doesn't honor `# type: ignore[]` comments, so disable the # rules that those comments suppress for mypy. From 9dcdac66967aaa932ce292dd7946e63f9184cfd1 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Fri, 6 Nov 2026 19:13:50 +0100 Subject: [PATCH 51/93] More cases for test_table.py --- python/pyarrow-stubs/tests/test_table.py | 65 ++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index 032181e25d4e..14a8a309caff 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -17,14 +17,42 @@ # mypy: warn-unused-ignores from types import MappingProxyType +from typing import Any import pytest +from typing_extensions import CapsuleType import pyarrow.lib as lib from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check +class _PyArrowArrayProvider: + def __arrow_array__(self, type: lib.DataType | None = None) -> lib.Array: + return lib.array([1, 2], type=type) + + +class _ArrowArrayProvider: + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: + return lib.array([1, 2]).__arrow_c_array__(requested_schema) + + +class _ArrowDeviceArrayProvider: + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: + return lib.array([1, 2]).__arrow_c_device_array__(requested_schema, **kwargs) + + +class _ArrowStreamProvider: + def __arrow_c_stream__( + self, requested_schema: CapsuleType | None = None + ) -> CapsuleType: + return lib.chunked_array([[1, 2]]).__arrow_c_stream__(requested_schema) + + def test_table() -> None: table = lib.table({"a": [1, 2], "b": ["x", "y"]}) check(assert_type(table, lib.Table), lib.Table) @@ -46,6 +74,11 @@ def test_table() -> None: check(assert_type(table[:1], lib.Array | lib.ChunkedArray | lib.Table), lib.Table) + check(assert_type(table.set_column(0, "a", [_ArrowDeviceArrayProvider()]), + lib.Table), + lib.Table) + check(assert_type(table.append_column("c", ([3], [4])), 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.select(["a", 1]), lib.Table), lib.Table) @@ -78,6 +111,21 @@ def test_table() -> None: check(batch, lib.RecordBatch) +def test_table_from_arrays_protocols() -> None: + check(assert_type(lib.Table.from_arrays([_PyArrowArrayProvider()], names=["a"]), + lib.Table), + lib.Table) + + +def test_table_from_arrays_numpy() -> None: + pytest.importorskip("numpy") + import numpy as np + + columns = np.array([[1, 2], [3, 4]]) + check(assert_type(lib.Table.from_arrays(columns, names=["a", "b"]), lib.Table), + lib.Table) + + def test_record_batch() -> None: batch = lib.record_batch({"a": [1, 2]}) # from a dict check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) @@ -93,6 +141,8 @@ def test_record_batch() -> None: lib.RecordBatch) check(assert_type(batch.add_column(0, "b", [3, 4]), lib.RecordBatch), lib.RecordBatch) + check(assert_type(batch.append_column("b", [3, 4]), lib.RecordBatch), + lib.RecordBatch) check(assert_type(batch.equals(batch), bool), bool) # Any object can be compared, a non-batch returns False check(assert_type(batch.equals(42), bool), bool) @@ -130,16 +180,31 @@ def test_chunked_array() -> None: lib.ChunkedArray) check(assert_type(lib.chunked_array([arr, [3]]), lib.ChunkedArray), lib.ChunkedArray) + check(assert_type(lib.chunked_array([_PyArrowArrayProvider()]), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(lib.chunked_array([_ArrowDeviceArrayProvider()]), + lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(lib.chunked_array(_ArrowArrayProvider()), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(lib.chunked_array([_ArrowArrayProvider()]), lib.ChunkedArray), + lib.ChunkedArray) + check(assert_type(lib.chunked_array(_ArrowStreamProvider()), lib.ChunkedArray), + lib.ChunkedArray) def test_invalid_usage() -> None: table = lib.table({"a": [1, 2]}) + streams = [_ArrowStreamProvider()] columns = {"a": lib.array([1])}.values() # dict as non-sequence example if TYPE_CHECKING_INVALID_USAGE: table.column(1.5) # type: ignore[arg-type] table.slice("0") # type: ignore[arg-type] lib.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] + # Stream providers are accepted as the whole input, not individual chunks + lib.chunked_array(streams) # type: ignore[arg-type] + table.add_column(0, "b", streams) # type: ignore[arg-type] def test_record_batch_invalid_usage() -> None: From 2f49f15dbe0dd541439fbc11d5dfabad39726b8e Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Fri, 6 Nov 2026 20:16:20 +0100 Subject: [PATCH 52/93] table.pyi: Allow array protocols for inputs. Specify append_column --- python/pyarrow-stubs/pyarrow/__lib/table.pyi | 25 ++++++++++++++++++-- python/pyarrow-stubs/tests/test_table.py | 7 ++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/table.pyi b/python/pyarrow-stubs/pyarrow/__lib/table.pyi index 981f8991f273..c520a7e6e57a 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/table.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/table.pyi @@ -53,6 +53,7 @@ from pyarrow._stubs_typing import ( SupportsArrowArray, SupportsArrowDeviceArray, SupportsArrowStream, + SupportsPyArrowArray, ) from pyarrow.compute import Expression from pyarrow.interchange.dataframe import _PyArrowDataFrame @@ -130,7 +131,12 @@ AggregateOptions: TypeAlias = ( ) ChunkedArrayConstructorInput: TypeAlias = ( - Iterable[Array | Iterable[Any] | SupportsArrowStream | SupportsArrowArray] + Iterable[ + Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray + ] | Array | SupportsArrowArray | SupportsArrowStream @@ -339,6 +345,9 @@ class RecordBatch(_Tabular): def add_column( self, i: int, field_: str | Field, column: ArrayOrChunkedArray | list[Any] ) -> RecordBatch: ... + def append_column( + self, field_: str | Field, column: ArrayOrChunkedArray | list[Any] + ) -> RecordBatch: ... def remove_column(self, i: int) -> RecordBatch: ... def set_column( self, i: int, field_: str | Field, column: Array | list @@ -457,9 +466,16 @@ class Table(_Tabular): columns: Iterable[str | int] | None = None, safe: bool = True, ) -> Table: ... + # When no schema, arrays should also support indexing @staticmethod def from_arrays( - arrays: Sequence[ArrayOrChunkedArray | Iterable[Any]], + arrays: Sequence[ + Iterable[Any] + | SupportsArrowArray + | SupportsArrowDeviceArray + | SupportsPyArrowArray + ] + | NDArray[Any], names: Sequence[str] | None = None, schema: Schema | None = None, metadata: Mapping[bytes, bytes] @@ -499,6 +515,11 @@ class Table(_Tabular): field_: str | Field, column: ChunkedArray | ChunkedArrayConstructorInput, ) -> Table: ... + def append_column( + self, + field_: str | Field, + column: ChunkedArray | ChunkedArrayConstructorInput, + ) -> Table: ... def remove_column(self, i: int) -> Table: ... def set_column( self, diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs/tests/test_table.py index 14a8a309caff..e9d6391c073a 100644 --- a/python/pyarrow-stubs/tests/test_table.py +++ b/python/pyarrow-stubs/tests/test_table.py @@ -115,6 +115,12 @@ def test_table_from_arrays_protocols() -> None: check(assert_type(lib.Table.from_arrays([_PyArrowArrayProvider()], names=["a"]), lib.Table), lib.Table) + check(assert_type(lib.Table.from_arrays([_ArrowArrayProvider()], names=["a"]), + lib.Table), + lib.Table) + check(assert_type(lib.Table.from_arrays([_ArrowDeviceArrayProvider()], names=["a"]), + lib.Table), + lib.Table) def test_table_from_arrays_numpy() -> None: @@ -211,6 +217,7 @@ def test_record_batch_invalid_usage() -> None: arr = lib.array([1]) chunked = lib.chunked_array([[{"a": 1}]]) if TYPE_CHECKING_INVALID_USAGE: + # Tuples are failing at runtime. Recheck after the fix of GH-55418 lib.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] lib.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] lib.record_batch((arr,), names=["a"]) # type: ignore[arg-type] From 3301e775bbb79f53497ba53fecefc63d0cb7f767 Mon Sep 17 00:00:00 2001 From: Alenka Frim Date: Mon, 9 Nov 2026 09:33:08 +0100 Subject: [PATCH 53/93] Update python/pyarrow-stubs/tests/test_types.py Co-authored-by: Joris Van den Bossche --- python/pyarrow-stubs/tests/test_types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 260dd86f8a74..21297e275f31 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -124,7 +124,7 @@ def test_types_enum() -> None: def test_data_type_hashable() -> None: - data_type: Hashable = lib.int32() + data_type = lib.int32() check(assert_type(hash(data_type), int), int) From bf4a592f3d4418f993c59fbe443db65856736c35 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Mon, 9 Nov 2026 09:47:56 +0100 Subject: [PATCH 54/93] Add test for Field.metadata --- python/pyarrow-stubs/tests/test_types.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs/tests/test_types.py index 21297e275f31..64200bff56a6 100644 --- a/python/pyarrow-stubs/tests/test_types.py +++ b/python/pyarrow-stubs/tests/test_types.py @@ -132,11 +132,15 @@ 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.metadata, dict[bytes, bytes] | None), type(None)) 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) + field = lib.field("a", lib.int32(), metadata={"k": "v"}) + check(assert_type(field.metadata, dict[bytes, bytes] | None), dict) + def test_schema() -> None: schema = lib.schema([("a", lib.int32()), lib.field("b", lib.string())]) From a9cdd2061e7b1ded1e5979c87d1a97331ce6865e Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Mon, 9 Nov 2026 10:24:28 +0100 Subject: [PATCH 55/93] Remove sys.version_info >= (3, 11) on all files --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 9 +-------- python/pyarrow-stubs/pyarrow/__lib/table.pyi | 10 +--------- python/pyarrow-stubs/pyarrow/__lib/tensor.pyi | 7 +------ 3 files changed, 3 insertions(+), 23 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index b0025637e512..4c6d79534bb5 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -15,16 +15,9 @@ # 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 +from typing import Any, Literal, Self import numpy as np # type: ignore[import-not-found] import pandas as pd # type: ignore[import-not-found,import-untyped] diff --git a/python/pyarrow-stubs/pyarrow/__lib/table.pyi b/python/pyarrow-stubs/pyarrow/__lib/table.pyi index c520a7e6e57a..ddf7e9203066 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/table.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/table.pyi @@ -15,16 +15,8 @@ # specific language governing permissions and limitations # under the License. -import sys +from typing import Self, TypeAlias -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 Iterable, Iterator, Sequence, Mapping from typing import Any, ClassVar, Literal from typing_extensions import CapsuleType, deprecated diff --git a/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi b/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi index 590ef0a9ac4b..6d1e7cc05d41 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi @@ -15,12 +15,7 @@ # 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 typing import Self from collections.abc import Sequence import numpy as np From 905aa33f264219f982e1920c97e0b60f828c9b8d Mon Sep 17 00:00:00 2001 From: Alenka Frim Date: Mon, 9 Nov 2026 10:35:46 +0100 Subject: [PATCH 56/93] Update python/pyarrow-stubs/pyarrow/__lib/error.pyi Co-authored-by: Joris Van den Bossche --- python/pyarrow-stubs/pyarrow/__lib/error.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/error.pyi b/python/pyarrow-stubs/pyarrow/__lib/error.pyi index 67be73eadb6d..6ba1c8396075 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/error.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/error.pyi @@ -17,7 +17,7 @@ from types import TracebackType -from typing_extensions import Self +from typing import Self class ArrowException(Exception): ... class ArrowInvalid(ValueError, ArrowException): ... From a94a96ed6564e7e3d337a336bb2960a1364dc0f9 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Mon, 9 Nov 2026 16:44:48 +0100 Subject: [PATCH 57/93] More cases for test_array.py --- python/pyarrow-stubs/tests/test_array.py | 67 ++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index 6c41e5b19832..1b9a464f3dc3 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -18,7 +18,10 @@ from typing import Any, cast +import pytest + import pyarrow.lib as lib +from pyarrow._compute import CastOptions from typing_extensions import CapsuleType from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -45,6 +48,13 @@ def test_array() -> None: 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(arr.cast(lib.int32(), safe=None), lib.Array), lib.Int32Array) + # without target_type argument, as CastOptions already specifies it + check(assert_type(arr.cast(options=CastOptions.safe(lib.int32())), lib.Array), + lib.Int32Array) + check(assert_type(arr.view("uint64"), lib.Array), lib.UInt64Array) + check(assert_type(arr.to_string(element_size_limit=5), str), str) + check(assert_type(arr.to_tensor(allow_nulls=True), lib.Tensor), lib.Tensor) check(assert_type(lib.concat_arrays([arr, arr]), lib.Array), lib.Int64Array) check(assert_type(lib.Array.from_pandas([1, 2], None, lib.int64()), @@ -54,6 +64,24 @@ def test_array() -> None: other = check(cast(lib.StringArray, lib.array(["a", "b", None])), lib.StringArray) check(assert_type(arr.equals(arr), bool), bool) check(assert_type(arr.equals(other), bool), bool) + # Arrays of any type can diff() + check(assert_type(arr.diff(other), str), str) + + +def test_array_numpy() -> None: + pytest.importorskip("numpy") + import numpy as np + + arr = cast(lib.Int64Array, lib.array([1, 2])) + check(assert_type(arr[np.int64(0)], Any | lib.Int64Array), lib.Int64Scalar) + check(assert_type(lib.infer_type(np.array([1])), lib.DataType), lib.DataType) + offsets = np.array([0, 1, 2], dtype="int32") + check(assert_type(lib.ListArray.from_arrays(offsets, [1, 2]), lib.ListArray), + lib.ListArray) + if TYPE_CHECKING_INVALID_USAGE: + dtype = lib.dictionary(lib.int8(), lib.string()) + strs = np.array(["a"]) + lib.DictionaryArray.from_buffers(dtype, 1, [], strs) # type: ignore[arg-type] def test_array_dlpack() -> None: @@ -79,6 +107,23 @@ def test_list_arrays() -> None: large = lib.LargeListArray.from_arrays([0, 1, 3], lib.array([1, 2, 3])) check(assert_type(large.value_lengths(), lib.Int64Array), lib.Int64Array) + list_type = lib.list_(lib.int64()) + check(assert_type(lib.ListArray.from_arrays([0, 2], [1, 2], list_type), + lib.ListArray), + lib.ListArray) + check(assert_type(lib.LargeListArray.from_arrays([0, 2], [1, 2]), + lib.LargeListArray), + lib.LargeListArray) + check(assert_type(lib.ListViewArray.from_arrays([0], [2], [1, 2]), + lib.ListViewArray), + lib.ListViewArray) + check(assert_type(lib.LargeListViewArray.from_arrays([0], [2], [1, 2]), + lib.LargeListViewArray), + lib.LargeListViewArray) + check(assert_type(lib.FixedSizeListArray.from_arrays([1, 2], 2), + lib.FixedSizeListArray), + lib.FixedSizeListArray) + # 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) @@ -97,6 +142,12 @@ def test_dictionary_array() -> None: check(assert_type(dict_arr.dictionary_decode(), lib.Array), lib.StringArray) + # The validity buffer can be None + rebuilt = lib.DictionaryArray.from_buffers( + dict_arr.type, len(dict_arr), dict_arr.indices.buffers(), dict_arr.dictionary + ) + check(assert_type(rebuilt, lib.DictionaryArray), lib.DictionaryArray) + def test_run_end_encoded_array() -> None: ree = lib.RunEndEncodedArray.from_arrays([2, 3], ["a", "b"]) @@ -134,3 +185,19 @@ def test_invalid_usage() -> None: arr.slice("1") # type: ignore[arg-type] arr.equals(object()) # type: ignore[arg-type] arr.equals(None) # type: ignore[arg-type] + arr.format() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] # noqa: E501 + # Generators aren't accepted, only Sequences + lib.infer_type(x for x in [1]) # type: ignore[arg-type] + + +def test_buffer_constructors_invalid_usage() -> None: + types = cast(lib.Int8Array, lib.array([0], lib.int8())) + children: list[lib.Array] = [lib.array([1])] + strings = cast(lib.StringArray, lib.array(["a"])) + off, data = cast(list[lib.Buffer], strings.buffers()[1:]) + if TYPE_CHECKING_INVALID_USAGE: + # Only lists are accepted for children and type codes + lib.UnionArray.from_sparse(types, tuple(children)) # type: ignore[arg-type] + lib.UnionArray.from_sparse(types, children, None, + types) # type: ignore[arg-type] + lib.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] From 319b5c5b45d6f1e656b969bf708c1272c18193af Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Mon, 9 Nov 2026 18:43:13 +0100 Subject: [PATCH 58/93] Review array.pyi --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 70 +++++++++----------- 1 file changed, 33 insertions(+), 37 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index 4c6d79534bb5..54d5a4a29f1f 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -17,7 +17,7 @@ from collections.abc import Iterable, Iterator, Sequence -from typing import Any, Literal, Self +from typing import Any, Literal, Self, SupportsIndex import numpy as np # type: ignore[import-not-found] import pandas as pd # type: ignore[import-not-found,import-untyped] @@ -76,7 +76,9 @@ def repeat( memory_pool: MemoryPool | None = None, ) -> ArrayLike: ... def infer_type( - values: Iterable[Any], mask: Mask | None = None, from_pandas: bool = False + values: Sequence[Any] | np.ndarray, + mask: Mask | None = None, + from_pandas: bool = False, ) -> DataType: ... class ArrayStatistics(_Weakrefable): @@ -119,15 +121,15 @@ class _PandasConvertible(_Weakrefable): ) -> pd.Series | pd.DataFrame: ... class Array(_PandasConvertible): - def diff(self, other: Self) -> str: ... + def diff(self, other: Array) -> str: ... def cast( self, - target_type: DataType | str, - safe: bool = True, + target_type: DataType | str | None = None, + safe: bool | None = None, options: CastOptions | None = None, memory_pool: MemoryPool | None = None, ) -> Array: ... - def view(self, target_type: DataType) -> Array: ... + def view(self, target_type: DataType | str) -> Array: ... def sum(self, **kwargs) -> Scalar: ... @property def type(self) -> DataType: ... @@ -166,16 +168,15 @@ class Array(_PandasConvertible): window: int = 10, container_window: int = 2, skip_new_lines: bool = False, + element_size_limit: int = 100, ) -> str: ... - - format = to_string def equals(self, other: Array) -> bool: ... def __len__(self) -> int: ... def is_null(self, *, nan_is_null: bool = False) -> BooleanArray: ... def is_nan(self) -> BooleanArray: ... def is_valid(self) -> BooleanArray: ... def fill_null(self, fill_value: Any) -> Self: ... - def __getitem__(self, key: int | builtins.slice) -> Scalar | Self: ... + def __getitem__(self, key: SupportsIndex | builtins.slice) -> Scalar | Self: ... def slice(self, offset: int = 0, length: int | None = None) -> Self: ... def take(self, indices: Indices) -> Self: ... def drop_null(self) -> Self: ... @@ -207,6 +208,7 @@ class Array(_PandasConvertible): ) -> list[Any]: ... tolist = to_pylist + def to_tensor(self, *, allow_nulls: bool = False) -> Tensor: ... def validate(self, *, full: bool = False) -> None: ... @property def offset(self) -> int: ... @@ -284,9 +286,8 @@ class ListArray(BaseListArray): @classmethod def from_arrays( cls, - offsets: Int32Array | list[int] | list[int | None], + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, values: Array | list[Any], - *, type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, @@ -301,9 +302,8 @@ class LargeListArray(BaseListArray): @classmethod def from_arrays( cls, - offsets: Int64Array | list[int] | list[int | None], - values: Array, - *, + offsets: Int64Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, @@ -318,10 +318,9 @@ class ListViewArray(BaseListArray): @classmethod def from_arrays( cls, - offsets: Int32Array | list[int] | list[int | None], - sizes: Int32Array | list[int] | list[int | None], - values: Array, - *, + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, + sizes: Int32Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, @@ -338,10 +337,9 @@ class LargeListViewArray(BaseListArray): @classmethod def from_arrays( cls, - offsets: Int64Array | list[int] | list[int | None], - sizes: Int64Array | list[int] | list[int | None], - values: Array, - *, + offsets: Int64Array | list[int] | list[int | None] | np.ndarray, + sizes: Int64Array | list[int] | list[int | None] | np.ndarray, + values: Array | list[Any], type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, @@ -358,9 +356,8 @@ class FixedSizeListArray(BaseListArray): @classmethod def from_arrays( cls, - values: Array, + values: Array | list[Any], list_size: int | None = None, - *, type: DataType | None = None, mask: Mask | None = None, ) -> FixedSizeListArray: ... @@ -372,10 +369,9 @@ class MapArray(ListArray): @classmethod def from_arrays( # type: ignore[override] cls, - offsets: Int32Array | list[int] | list[int | None], + offsets: Int32Array | list[int] | list[int | None] | np.ndarray, keys: Array | np.ndarray | list, items: Array | np.ndarray | list, - *, type: DataType | None = None, pool: MemoryPool | None = None, mask: Mask | None = None, @@ -397,16 +393,16 @@ class UnionArray(Array): def from_dense( types: Int8Array, value_offsets: Int32Array, - children: NullableCollection, + children: list[Array], field_names: list[str] | None = None, - type_codes: Int8Array | list[int] | None = None, + type_codes: list[int] | None = None, ) -> UnionArray: ... @staticmethod def from_sparse( types: Int8Array, - children: NullableCollection, + children: list[Array], field_names: list[str] | None = None, - type_codes: Int8Array | list[int] | None = None, + type_codes: list[int] | None = None, ) -> UnionArray: ... class StringArray(Array): @@ -416,8 +412,8 @@ class StringArray(Array): value_offsets: Buffer, data: Buffer, null_bitmap: Buffer | None = None, - null_count: int | None = -1, - offset: int | None = 0, + null_count: int = -1, + offset: int = 0, ) -> StringArray: ... class LargeStringArray(Array): @@ -427,8 +423,8 @@ class LargeStringArray(Array): value_offsets: Buffer, data: Buffer, null_bitmap: Buffer | None = None, - null_count: int | None = -1, - offset: int | None = 0, + null_count: int = -1, + offset: int = 0, ) -> LargeStringArray: ... class StringViewArray(Array): ... @@ -454,8 +450,8 @@ class DictionaryArray(Array): def from_buffers( # type: ignore[override] type: DataType, length: int, - buffers: list[Buffer], - dictionary: Array | np.ndarray | pd.Series, + buffers: Sequence[Buffer | None], + dictionary: Array, null_count: int = -1, offset: int = 0, ) -> DictionaryArray: ... @@ -520,7 +516,6 @@ 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 @@ -548,6 +543,7 @@ __all__ = [ "nulls", "repeat", "infer_type", + "ArrayStatistics", "_PandasConvertible", "Array", "NullArray", From b642a3b322ffcebae26668de522bf15d6a4844dc Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Mon, 9 Nov 2026 23:32:07 +0100 Subject: [PATCH 59/93] Apply AlenkaF's suggestions and add tests --- python/pyarrow-stubs/pyarrow/__lib/array.pyi | 52 ++++++++++++++------ python/pyarrow-stubs/tests/test_array.py | 27 ++++++++-- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/pyarrow/__lib/array.pyi index 54d5a4a29f1f..7e7268fdc915 100644 --- a/python/pyarrow-stubs/pyarrow/__lib/array.pyi +++ b/python/pyarrow-stubs/pyarrow/__lib/array.pyi @@ -15,9 +15,9 @@ # specific language governing permissions and limitations # under the License. -from collections.abc import Iterable, Iterator, Sequence +from collections.abc import Collection, Iterable, Iterator, Sequence -from typing import Any, Literal, Self, SupportsIndex +from typing import Any, ClassVar, Self, SupportsIndex import numpy as np # type: ignore[import-not-found] import pandas as pd # type: ignore[import-not-found,import-untyped] @@ -28,6 +28,9 @@ from pyarrow._stubs_typing import ( Indices, Mask, MapsAsPydicts, + NullEncoding, + NullPlacement, + NullSelectionBehavior, Order, SupportsArrowArray, SupportsArrowDeviceArray, @@ -41,13 +44,12 @@ from pyarrow.lib import ( Tensor, _Weakrefable, ) -from typing_extensions import deprecated +from typing_extensions import CapsuleType, deprecated import builtins from .scalar import Int64Scalar, Scalar from .device import DeviceAllocationType from .types import BaseExtensionType, DataType, Field, StructType -from pyarrow._stubs_typing import NullableCollection def array( values: Iterable[Any] @@ -121,6 +123,7 @@ class _PandasConvertible(_Weakrefable): ) -> pd.Series | pd.DataFrame: ... class Array(_PandasConvertible): + __hash__: ClassVar[None] # type: ignore[assignment] def diff(self, other: Array) -> str: ... def cast( self, @@ -134,7 +137,9 @@ class Array(_PandasConvertible): @property def type(self) -> DataType: ... def unique(self) -> Self: ... - def dictionary_encode(self, null_encoding: str = "mask") -> DictionaryArray: ... + def dictionary_encode( + self, null_encoding: NullEncoding = "mask" + ) -> DictionaryArray: ... def value_counts(self) -> StructArray: ... @staticmethod def from_pandas( @@ -151,7 +156,7 @@ class Array(_PandasConvertible): buffers: Sequence[Buffer | None], null_count: int = -1, offset=0, - children: NullableCollection | None = None, + children: Collection[Array] | None = None, ) -> Array: ... @property def null_count(self) -> int: ... @@ -184,7 +189,7 @@ class Array(_PandasConvertible): self, mask: Mask, *, - null_selection_behavior: Literal["drop", "emit_null"] = "drop", + null_selection_behavior: NullSelectionBehavior = "drop", ) -> Self: ... def index( self, @@ -194,7 +199,12 @@ class Array(_PandasConvertible): *, memory_pool: MemoryPool | None = None, ) -> Int64Scalar: ... - def sort(self, order: Order = "ascending", **kwargs) -> Self: ... + def sort( + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + **kwargs, + ) -> Self: ... def __array__( self, dtype: np.dtype | None = None, copy: bool | None = None ) -> np.ndarray: ... @@ -204,7 +214,7 @@ class Array(_PandasConvertible): def to_pylist( self, *, - maps_as_pydicts: Literal["lossy", "strict"] | None = None, + maps_as_pydicts: MapsAsPydicts | None = None, ) -> list[Any]: ... tolist = to_pylist @@ -217,15 +227,23 @@ class Array(_PandasConvertible): def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... @staticmethod def _import_from_c(in_ptr: int, type: int | DataType) -> Array: ... - def __arrow_c_array__(self, requested_schema=None) -> Any: ... + def __arrow_c_array__( + self, requested_schema: CapsuleType | None = None + ) -> tuple[CapsuleType, CapsuleType]: ... @staticmethod - def _import_from_c_capsule(schema_capsule, array_capsule) -> Array: ... + def _import_from_c_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> Array: ... def _export_to_c_device(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... @staticmethod def _import_from_c_device(in_ptr: int, type: DataType | int) -> Array: ... - def __arrow_c_device_array__(self, requested_schema=None, **kwargs) -> Any: ... + def __arrow_c_device_array__( + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: ... @staticmethod - def _import_from_c_device_capsule(schema_capsule, array_capsule) -> Array: ... + def _import_from_c_device_capsule( + schema_capsule: CapsuleType, array_capsule: CapsuleType + ) -> Array: ... def __dlpack__( self, *, @@ -233,7 +251,7 @@ class Array(_PandasConvertible): max_version: tuple[int, int] | None = None, dl_device: tuple[int, int] | None = None, copy: bool | None = None, - ) -> Any: ... + ) -> CapsuleType: ... def __dlpack_device__(self) -> tuple[int, int]: ... @property def device_type(self) -> DeviceAllocationType: ... @@ -479,7 +497,11 @@ class StructArray(Array): type: StructType | None = None, ) -> StructArray: ... def sort( - self, order: Order = "ascending", by: str | None = None, **kwargs + self, + order: Order = "ascending", + null_placement: NullPlacement = "at_end", + by: str | None = None, + **kwargs, ) -> StructArray: ... class RunEndEncodedArray(Array): diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs/tests/test_array.py index 1b9a464f3dc3..1284720fb4e8 100644 --- a/python/pyarrow-stubs/tests/test_array.py +++ b/python/pyarrow-stubs/tests/test_array.py @@ -16,6 +16,7 @@ # under the License. # mypy: warn-unused-ignores +from collections.abc import Hashable from typing import Any, cast import pytest @@ -91,14 +92,29 @@ def test_array_dlpack() -> None: dl_device=arr.__dlpack_device__(), copy=False, ) - check(assert_type(capsule, Any), CapsuleType) + check(assert_type(capsule, CapsuleType), CapsuleType) def test_import_from_c_capsule() -> None: - capsules = lib.array([1, 2]).__arrow_c_array__() + arr = cast(lib.Int64Array, lib.array([1, 2])) + capsules = arr.__arrow_c_array__() + check(assert_type(capsules, tuple[CapsuleType, CapsuleType]), tuple) imported = lib.BooleanArray._import_from_c_capsule(*capsules) check(assert_type(imported, lib.Array), lib.Int64Array) + check(assert_type(arr.__arrow_c_device_array__(), + tuple[CapsuleType, CapsuleType]), + tuple) + + +def test_from_buffers() -> None: + values = cast(lib.Int64Array, lib.array([1, 2])) + offsets = lib.array([0, 2], lib.int32()).buffers()[1] + # Children can be any collection of arrays, like a tuple + list_arr = lib.Array.from_buffers(lib.list_(lib.int64()), 1, [None, offsets], + children=(values,)) + check(assert_type(list_arr, lib.Array), lib.ListArray) + def test_list_arrays() -> None: lst = lib.ListArray.from_arrays([0, 1, 3], [1, 2, 3]) @@ -142,7 +158,6 @@ def test_dictionary_array() -> None: check(assert_type(dict_arr.dictionary_decode(), lib.Array), lib.StringArray) - # The validity buffer can be None rebuilt = lib.DictionaryArray.from_buffers( dict_arr.type, len(dict_arr), dict_arr.indices.buffers(), dict_arr.dictionary ) @@ -186,6 +201,12 @@ def test_invalid_usage() -> None: arr.equals(object()) # type: ignore[arg-type] arr.equals(None) # type: ignore[arg-type] arr.format() # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] # noqa: E501 + # Arrays define __eq__ but not __hash__, hence unhashable at runtime + _: Hashable = arr # type: ignore[assignment] # ty: ignore[invalid-assignment] + arr.sort(null_placement="first") # type: ignore[arg-type] + arr.filter([True], null_selection_behavior="keep") # type: ignore[arg-type] + arr.dictionary_encode("drop") # type: ignore[arg-type] + arr.to_pylist(maps_as_pydicts="loose") # type: ignore[arg-type] # Generators aren't accepted, only Sequences lib.infer_type(x for x in [1]) # type: ignore[arg-type] From 4b7dc90fbefed1e669072078880266cc4842f8f2 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 00:12:28 +0100 Subject: [PATCH 60/93] Apply batched suggestions from code review Co-authored-by: Rok Mihevc --- python/pyarrow-stubs/pyarrow/lib.pyi | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/pyarrow/lib.pyi index 8a1291188bec..d17d05189c76 100644 --- a/python/pyarrow-stubs/pyarrow/lib.pyi +++ b/python/pyarrow-stubs/pyarrow/lib.pyi @@ -33,4 +33,14 @@ from pyarrow.__lib.scalar import * # noqa: F401,F403 from pyarrow.__lib.table import * # noqa: F401,F403 from pyarrow.__lib.tensor import * # noqa: F401,F403 +class MonthDayNano(tuple): + months: int + days: int + nanoseconds: int + + def __new__( + cls, + sequence: tuple[int, int, int] | list[int] = ..., + ) -> MonthDayNano: ... + def __getattr__(name: str) -> Any: ... From 41cc576f7428b9406af9256e593ba573dc23bfbc Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 00:14:15 +0100 Subject: [PATCH 61/93] Move stubs to a flat pyarrow-stubs package and tests to pyarrow-stubs-tests Lay out python/pyarrow-stubs/ as a PEP 561 stub-only package for pyarrow (python/pyarrow-stubs/pyarrow/* -> python/pyarrow-stubs/*), and move the type stub tests out of it to python/pyarrow-stubs-tests/, where they aren't read as pyarrow.tests. Files are moved unchanged; the configuration is updated in the next commit. --- python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/_check.py | 0 python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_array.py | 0 python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_table.py | 0 python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_types.py | 0 python/pyarrow-stubs/{pyarrow => }/__init__.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/__init__.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/array.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/builder.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/config.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/device.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/error.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/io.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/ipc.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/memory.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/scalar.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/table.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/tensor.pyi | 0 python/pyarrow-stubs/{pyarrow => }/__lib/types.pyi | 0 python/pyarrow-stubs/{pyarrow => }/_compute.pyi | 0 python/pyarrow-stubs/{pyarrow => }/_stubs_typing.pyi | 0 python/pyarrow-stubs/{pyarrow => }/lib.pyi | 0 python/pyarrow-stubs/{pyarrow => }/types.pyi | 0 22 files changed, 0 insertions(+), 0 deletions(-) rename python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/_check.py (100%) rename python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_array.py (100%) rename python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_table.py (100%) rename python/{pyarrow-stubs/tests => pyarrow-stubs-tests}/test_types.py (100%) rename python/pyarrow-stubs/{pyarrow => }/__init__.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/__init__.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/array.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/builder.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/config.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/device.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/error.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/io.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/ipc.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/memory.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/scalar.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/table.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/tensor.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/__lib/types.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/_compute.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/_stubs_typing.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/lib.pyi (100%) rename python/pyarrow-stubs/{pyarrow => }/types.pyi (100%) diff --git a/python/pyarrow-stubs/tests/_check.py b/python/pyarrow-stubs-tests/_check.py similarity index 100% rename from python/pyarrow-stubs/tests/_check.py rename to python/pyarrow-stubs-tests/_check.py diff --git a/python/pyarrow-stubs/tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py similarity index 100% rename from python/pyarrow-stubs/tests/test_array.py rename to python/pyarrow-stubs-tests/test_array.py diff --git a/python/pyarrow-stubs/tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py similarity index 100% rename from python/pyarrow-stubs/tests/test_table.py rename to python/pyarrow-stubs-tests/test_table.py diff --git a/python/pyarrow-stubs/tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py similarity index 100% rename from python/pyarrow-stubs/tests/test_types.py rename to python/pyarrow-stubs-tests/test_types.py diff --git a/python/pyarrow-stubs/pyarrow/__init__.pyi b/python/pyarrow-stubs/__init__.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__init__.pyi rename to python/pyarrow-stubs/__init__.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/__init__.pyi b/python/pyarrow-stubs/__lib/__init__.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/__init__.pyi rename to python/pyarrow-stubs/__lib/__init__.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/array.pyi rename to python/pyarrow-stubs/__lib/array.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/builder.pyi b/python/pyarrow-stubs/__lib/builder.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/builder.pyi rename to python/pyarrow-stubs/__lib/builder.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/config.pyi b/python/pyarrow-stubs/__lib/config.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/config.pyi rename to python/pyarrow-stubs/__lib/config.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/device.pyi b/python/pyarrow-stubs/__lib/device.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/device.pyi rename to python/pyarrow-stubs/__lib/device.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/error.pyi b/python/pyarrow-stubs/__lib/error.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/error.pyi rename to python/pyarrow-stubs/__lib/error.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/io.pyi b/python/pyarrow-stubs/__lib/io.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/io.pyi rename to python/pyarrow-stubs/__lib/io.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/ipc.pyi b/python/pyarrow-stubs/__lib/ipc.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/ipc.pyi rename to python/pyarrow-stubs/__lib/ipc.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/memory.pyi b/python/pyarrow-stubs/__lib/memory.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/memory.pyi rename to python/pyarrow-stubs/__lib/memory.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/scalar.pyi b/python/pyarrow-stubs/__lib/scalar.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/scalar.pyi rename to python/pyarrow-stubs/__lib/scalar.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/table.pyi rename to python/pyarrow-stubs/__lib/table.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/tensor.pyi b/python/pyarrow-stubs/__lib/tensor.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/tensor.pyi rename to python/pyarrow-stubs/__lib/tensor.pyi diff --git a/python/pyarrow-stubs/pyarrow/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/__lib/types.pyi rename to python/pyarrow-stubs/__lib/types.pyi diff --git a/python/pyarrow-stubs/pyarrow/_compute.pyi b/python/pyarrow-stubs/_compute.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/_compute.pyi rename to python/pyarrow-stubs/_compute.pyi diff --git a/python/pyarrow-stubs/pyarrow/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/_stubs_typing.pyi rename to python/pyarrow-stubs/_stubs_typing.pyi diff --git a/python/pyarrow-stubs/pyarrow/lib.pyi b/python/pyarrow-stubs/lib.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/lib.pyi rename to python/pyarrow-stubs/lib.pyi diff --git a/python/pyarrow-stubs/pyarrow/types.pyi b/python/pyarrow-stubs/types.pyi similarity index 100% rename from python/pyarrow-stubs/pyarrow/types.pyi rename to python/pyarrow-stubs/types.pyi From 82f4e5373f2ed708efb855934580e8d9ab415171 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 00:14:51 +0100 Subject: [PATCH 62/93] Address review comments on type checking setup and stubs - Update the configuration to the new layout: run mypy, pyright and ty on pyarrow-stubs/ and pyarrow-stubs-tests/ without excludes, and document how to run them. - Pin the type checkers and stub packages in python/requirements-stubs.txt, used by the CI script, and to the same versions in the stub pre-commit hooks. The Lint CI job runs these hooks. - Run stubtest in CI with scripts/run_stubtest.py and an allowlist of the known mismatches. - Move the data type classes that don't exist at runtime (Int32Type, ...) to pyarrow._stubs_typing and add an IntegerType alias there. A conftest.py provides the module at runtime, so the tests import them directly. - Make MapArray.from_arrays a static method that takes type positionally, like the runtime, and drop the type: ignore comments it required. - Drop flake8 on stubs, the numpy marks added to non-stub tests, and import ignores for packages that are now installed. - Enable mypy's warn_unused_ignores in pyproject.toml instead of per test file, and take assert_type from typing. --- .github/workflows/dev.yml | 3 + .pre-commit-config.yaml | 25 ++---- ci/scripts/python_test_type_annotations.sh | 14 +-- ci/scripts/python_wheel_validate_contents.py | 4 +- docs/source/developers/python/development.rst | 55 +++++++++--- python/CMakeLists.txt | 4 +- python/pyarrow-stubs-tests/_check.py | 3 +- python/pyarrow-stubs-tests/conftest.py | 37 ++++++++ .../stubtest_allowlist.txt | 67 ++++++++++++++ python/pyarrow-stubs-tests/test_array.py | 7 +- python/pyarrow-stubs-tests/test_table.py | 1 - python/pyarrow-stubs-tests/test_types.py | 55 ++++-------- python/pyarrow-stubs/__lib/array.pyi | 3 +- python/pyarrow-stubs/__lib/types.pyi | 88 ++++++------------- python/pyarrow-stubs/_stubs_typing.pyi | 38 ++++++++ python/pyarrow-stubs/lib.pyi | 3 +- python/pyarrow-stubs/types.pyi | 38 ++++---- python/pyarrow/fs.py | 2 +- python/pyarrow/tests/test_array.py | 16 ++-- python/pyarrow/tests/test_convert_builtin.py | 3 - python/pyarrow/tests/test_schema.py | 1 - python/pyarrow/tests/test_sparse_tensor.py | 5 +- python/pyproject.toml | 55 +++++------- python/requirements-stubs.txt | 17 ++++ python/scripts/run_stubtest.py | 72 +++++++++++++++ 25 files changed, 405 insertions(+), 211 deletions(-) create mode 100644 python/pyarrow-stubs-tests/conftest.py create mode 100644 python/pyarrow-stubs-tests/stubtest_allowlist.txt create mode 100644 python/requirements-stubs.txt create mode 100644 python/scripts/run_stubtest.py diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 7c3607bd1df3..5c24915a9426 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -67,6 +67,9 @@ jobs: - name: Run pre-commit run: | pre-commit run --all-files --color=always --show-diff-on-failure + - name: Run pre-commit for type stubs + run: | + pre-commit run --hook-stage manual python-stubs --all-files --color=always --show-diff-on-failure release: name: Source Release and Merge Script on ${{ matrix.runs-on }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a086779210df..bb7ceb65df98 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -136,29 +136,20 @@ 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: + # Linting, formatting and type checking of the type stubs. These hooks don't + # run on commit, only on request (the Lint CI job runs them too): # pre-commit run --hook-stage manual python-stubs --all-files + # Keep the versions of the type checkers and stub packages in sync with + # python/requirements-stubs.txt. - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.17.2 hooks: @@ -193,14 +184,14 @@ repos: entry: bash -c "cd python && exec mypy" pass_filenames: false files: >- - ^python/(pyarrow-stubs/|pyproject\.toml$) + ^python/(pyarrow-stubs(-tests)?/|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 + - pyarrow==26.0.1 - pyright==1.1.414 - pytest==9.0.2 - scipy-stubs==1.17.1.5 @@ -218,7 +209,7 @@ repos: entry: pyright -p python/pyproject.toml pass_filenames: false files: >- - ^python/(pyarrow-stubs/|pyproject\.toml$) + ^python/(pyarrow-stubs(-tests)?/|pyproject\.toml$) additional_dependencies: *python-type-check-dependencies - id: ty alias: python-stubs @@ -229,7 +220,7 @@ repos: entry: ty check --project python pass_filenames: false files: >- - ^python/(pyarrow-stubs/|pyproject\.toml$) + ^python/(pyarrow-stubs(-tests)?/|pyproject\.toml$) additional_dependencies: *python-type-check-dependencies - repo: https://github.com/MarcoGorelli/cython-lint rev: v0.21.1 diff --git a/ci/scripts/python_test_type_annotations.sh b/ci/scripts/python_test_type_annotations.sh index ab7dddf66be2..107c5ccf4f61 100755 --- a/ci/scripts/python_test_type_annotations.sh +++ b/ci/scripts/python_test_type_annotations.sh @@ -25,18 +25,20 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then . "${ARROW_PYTHON_VENV}/bin/activate" fi -# Install library stubs. Note some libraries contain their own type hints so they need to be installed. -pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil +cd "${pyarrow_dir}" -# Install type checkers -pip install mypy pyright ty +# Install type checkers and library stubs. Note some libraries contain their +# own type hints so they need to be installed. +pip install -r requirements-stubs.txt # Run type checkers -cd "${pyarrow_dir}" mypy pyright --stats ty check --verbose --output-format concise # The type annotation tests are type checked above; also run them to verify # the expected annotated types against the runtime types. -pytest -r s pyarrow-stubs/tests +pytest -r s pyarrow-stubs-tests + +# Check that the stubs match the signatures of the installed PyArrow. +python scripts/run_stubtest.py diff --git a/ci/scripts/python_wheel_validate_contents.py b/ci/scripts/python_wheel_validate_contents.py index 36956cf0b294..f11b24c11e0a 100644 --- a/ci/scripts/python_wheel_validate_contents.py +++ b/ci/scripts/python_wheel_validate_contents.py @@ -60,7 +60,7 @@ def validate_wheel(path): ), "pyarrow/py.typed is present in the wheel." source_root = Path(__file__).resolve().parents[2] - stubs_dir = source_root / "python" / "pyarrow-stubs" / "pyarrow" + stubs_dir = source_root / "python" / "pyarrow-stubs" assert stubs_dir.exists(), f"Stub source directory not found: {stubs_dir}" expected_stub_files = { @@ -75,7 +75,7 @@ def validate_wheel(path): } assert not (wheel_stub_files == expected_stub_files), ( - "Wheel .pyi files do not differ from python/pyarrow-stubs/pyarrow.\n" + "Wheel .pyi files do not differ from python/pyarrow-stubs.\n" f"Missing in wheel: {sorted(expected_stub_files - wheel_stub_files)}\n" f"Unexpected in wheel: {sorted(wheel_stub_files - expected_stub_files)}" ) diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index e50a78096ca8..af0c07c7d6fc 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -105,8 +105,10 @@ Type Checking ============= PyArrow provides type stubs (``*.pyi`` files) for static type checking. These -stubs are located in the ``pyarrow-stubs/`` directory and are automatically -included in the distributed wheel packages. +stubs are located in the ``pyarrow-stubs/`` directory, which is laid out as a +:pep:`561` stub-only package for ``pyarrow``, so type checkers and IDEs pick +it up from the ``arrow/python`` directory. The stubs are automatically included +in the distributed wheel packages. Running Type Checkers --------------------- @@ -122,9 +124,20 @@ also lints and formats the stubs with ruff: $ pre-commit run --hook-stage manual python-stubs --all-files -These hooks only run on request, not on commit. +These hooks only run on request, not on commit. The Lint CI job runs them as +well. CI also runs the type checkers, the type stub tests and stubtest (see +below) against the PyArrow it builds, with +``ci/scripts/python_test_type_annotations.sh``. -To run a type checker directly instead: +To run a type checker directly instead, install the pinned type checkers and +stub packages in your PyArrow development environment: + +.. code-block:: + + $ pip install -r requirements-stubs.txt + +Checking the PyArrow tests also needs their dependencies from +``requirements-test.txt``. **mypy** @@ -136,6 +149,12 @@ To run mypy on the PyArrow codebase: $ mypy The mypy configuration is in the ``[tool.mypy]`` section of ``pyproject.toml``. +mypy only uses the stubs when it checks them too, so pass ``pyarrow-stubs`` +along with any other files you check: + +.. code-block:: + + $ mypy pyarrow-stubs pyarrow/tests/test_array.py **pyright** @@ -157,6 +176,20 @@ To run ty (note: currently only partially configured): $ cd arrow/python $ ty check +**stubtest** + +mypy's stubtest compares the stubs with the signatures of the installed +PyArrow. Run it with: + +.. code-block:: + + $ cd arrow/python + $ python scripts/run_stubtest.py + +Stubs that aren't written yet are skipped. The known mismatches are listed in +``pyarrow-stubs-tests/stubtest_allowlist.txt``; remove the entries of the +mismatches you fix, as stubtest fails on unused entries. + Maintaining Type Stubs ----------------------- @@ -184,7 +217,7 @@ 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 +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: @@ -209,19 +242,13 @@ 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. + $ pytest pyarrow-stubs-tests 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. +fails once the usage is no longer rejected. ty doesn't honor mypy error codes, +so it doesn't check invalid usage. Doctest ======= diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 5bff8bcc6499..28c163328ff3 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -1077,7 +1077,7 @@ endif() # # TODO(GH-32609): Reintroduce type stubs into PyArrow wheels. # GH-49831: Temporarily do not install pyarrow-stubs into wheels. -# Stubs live in pyarrow-stubs/pyarrow/ during development but are installed +# Stubs live in pyarrow-stubs/ during development but are installed # alongside the package so type checkers can find them (PEP 561). # The stubs are currently incomplete, and some type checkers consume .pyi files # even without the py.typed marker. Re-enable this when the stubs are complete. @@ -1085,7 +1085,7 @@ if(PYARROW_REQUIRE_STUB_DOCSTRINGS) message(FATAL_ERROR "PYARROW_REQUIRE_STUB_DOCSTRINGS cannot be used while " "pyarrow-stubs are omitted from wheels (GH-49831).") endif() -# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs/pyarrow") +# set(PYARROW_STUBS_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/pyarrow-stubs") # if(EXISTS "${PYARROW_STUBS_SOURCE_DIR}") # install(DIRECTORY "${PYARROW_STUBS_SOURCE_DIR}/" # DESTINATION "." diff --git a/python/pyarrow-stubs-tests/_check.py b/python/pyarrow-stubs-tests/_check.py index 7c02e1509c6e..508f3c44a142 100644 --- a/python/pyarrow-stubs-tests/_check.py +++ b/python/pyarrow-stubs-tests/_check.py @@ -34,8 +34,7 @@ being an error. """ -from typing import Any, Final, TYPE_CHECKING, TypeVar -from typing_extensions import assert_type +from typing import Any, Final, TYPE_CHECKING, TypeVar, assert_type __all__ = ["TYPE_CHECKING_INVALID_USAGE", "assert_type", "check"] diff --git a/python/pyarrow-stubs-tests/conftest.py b/python/pyarrow-stubs-tests/conftest.py new file mode 100644 index 000000000000..7871f5c57118 --- /dev/null +++ b/python/pyarrow-stubs-tests/conftest.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import sys +import types + +import pyarrow.lib as lib + + +def _stub_only_type(name: str) -> type[lib.DataType]: + # The factories of the stub-only DataType subclasses (Int32Type, ...) + # return plain DataType instances at runtime. + if name.endswith("Type") and not name.startswith("_"): + return lib.DataType + raise AttributeError(name) + + +# pyarrow._stubs_typing only exists in the type stubs. Provide it at runtime, +# so that the tests can import the stub-only types without a TYPE_CHECKING +# block. +_stubs_typing = types.ModuleType("pyarrow._stubs_typing") +setattr(_stubs_typing, "__getattr__", _stub_only_type) +sys.modules["pyarrow._stubs_typing"] = _stubs_typing diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt new file mode 100644 index 000000000000..7cb289d267a8 --- /dev/null +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -0,0 +1,67 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Mismatches between the type stubs and the runtime that stubtest skips, one +# regex per line; see scripts/run_stubtest.py. Unused entries fail stubtest, +# so remove the entries of the mismatches you fix. Print the entries for the +# current mismatches with +# python scripts/run_stubtest.py --generate-allowlist + +# Stub-only modules. The pyarrow.__lib modules are checked against +# pyarrow.lib, which has no __all__. +pyarrow\._stubs_typing +pyarrow\.__lib\.[a-z]+\.__all__ + +# The __init__ of Cython extension types takes (*args, **kwargs) at runtime. +pyarrow\.__lib\.[a-z]+\.[A-Za-z_]+\.__init__ + +# Type aliases that are only defined in the stubs. +pyarrow\.__lib\.table\.AggregateOptions +pyarrow\.__lib\.table\.Aggregation +pyarrow\.__lib\.table\.ArrayOrChunkedArray +pyarrow\.__lib\.table\.ChunkedArrayConstructorInput +pyarrow\.__lib\.table\.ColumnSelector +pyarrow\.__lib\.table\.JoinType +pyarrow\.__lib\.table\.NarySelector +pyarrow\.__lib\.table\.NullarySelector +pyarrow\.__lib\.table\.UnarySelector + +# pyarrow.fs has no stub yet, so mypy reads its source. +pyarrow\.fs\.FSSpecHandler + +# Stubs to fix (GH-32609). +pyarrow\.__lib\.array\.Bool8Array\.from_numpy +pyarrow\.__lib\.array\.Bool8Array\.from_storage +pyarrow\.__lib\.array\.FixedShapeTensorArray\.from_numpy_ndarray +pyarrow\.__lib\.array\.FixedSizeListArray\.from_arrays +pyarrow\.__lib\.array\.LargeListArray\.from_arrays +pyarrow\.__lib\.array\.LargeListViewArray\.from_arrays +pyarrow\.__lib\.array\.ListArray\.from_arrays +pyarrow\.__lib\.array\.ListViewArray\.from_arrays +pyarrow\.__lib\.array\.UnionArray\.child +pyarrow\.__lib\.array\.array +pyarrow\.__lib\.device\.DeviceAllocationType\.[A-Z_]+ +pyarrow\.__lib\.device\.DeviceAllocationType\.__format__ +pyarrow\.__lib\.table\.ChunkedArray\.data +pyarrow\.__lib\.table\.ChunkedArray\.format +pyarrow\.__lib\.table\.concat_tables +pyarrow\.__lib\.tensor\.Sparse[A-Z]+(Tensor|Matrix)\.from_(numpy|pydata_sparse|scipy|tensor) +pyarrow\.__lib\.tensor\.Tensor\.from_numpy +pyarrow\.__lib\.types\.Schema\.__getitem__ +pyarrow\.__lib\.types\.Schema\.add_metadata +pyarrow\.__lib\.types\.Schema\.field_by_name +pyarrow\.types\.__all__ diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 1284720fb4e8..b1f9b9a3c04e 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -14,16 +14,15 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# mypy: warn-unused-ignores from collections.abc import Hashable from typing import Any, cast import pytest +from typing_extensions import CapsuleType import pyarrow.lib as lib from pyarrow._compute import CastOptions -from typing_extensions import CapsuleType from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -146,6 +145,10 @@ def test_list_arrays() -> None: 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) + map_type = lib.map_(lib.int64(), lib.string()) + check(assert_type(lib.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"], map_type), + lib.MapArray), + lib.ListArray) def test_dictionary_array() -> None: diff --git a/python/pyarrow-stubs-tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py index e9d6391c073a..0e36b73b4a1f 100644 --- a/python/pyarrow-stubs-tests/test_table.py +++ b/python/pyarrow-stubs-tests/test_table.py @@ -14,7 +14,6 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# mypy: warn-unused-ignores from types import MappingProxyType from typing import Any diff --git a/python/pyarrow-stubs-tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py index 64200bff56a6..27dcd0588937 100644 --- a/python/pyarrow-stubs-tests/test_types.py +++ b/python/pyarrow-stubs-tests/test_types.py @@ -14,39 +14,25 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# mypy: warn-unused-ignores from collections.abc import Hashable, Iterator -from typing import TYPE_CHECKING, Literal +from typing import Literal + +from typing_extensions import CapsuleType import pyarrow.lib as lib +from pyarrow._stubs_typing import ( + BoolType, + Int16Type, + Int32Type, + Int64Type, + IntegerType, + StringType, +) import pyarrow.types as types -from typing_extensions import CapsuleType 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. @@ -54,10 +40,10 @@ def _erase(t: lib.DataType) -> lib.DataType: 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) + 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: @@ -101,7 +87,7 @@ def test_dictionary_type() -> None: 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) + check(assert_type(dict_type.index_type, IntegerType), lib.DataType) def test_map_and_run_end_encoded_types() -> None: @@ -112,17 +98,10 @@ def test_map_and_run_end_encoded_types() -> None: 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"), + 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_data_type_hashable() -> None: data_type = lib.int32() check(assert_type(hash(data_type), int), int) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 7e7268fdc915..6f12a4325f2e 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -384,9 +384,8 @@ class FixedSizeListArray(BaseListArray): def value_lengths(self) -> Int32Array: ... class MapArray(ListArray): - @classmethod + @staticmethod def from_arrays( # type: ignore[override] - cls, offsets: Int32Array | list[int] | list[int | None] | np.ndarray, keys: Array | np.ndarray | list, items: Array | np.ndarray | list, diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index e40092f36c85..beb52cfeba0d 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -22,7 +22,33 @@ import numpy as np # type: ignore[import-not-found] import pandas as pd # type: ignore[import-not-found,import-untyped] from typing_extensions import CapsuleType, Self, deprecated -from pyarrow._stubs_typing import SupportsArrowSchema, TimeUnit +from pyarrow._stubs_typing import ( + BinaryType, + BinaryViewType, + BoolType, + Date32Type, + Date64Type, + Float16Type, + Float32Type, + Float64Type, + Int16Type, + Int32Type, + Int64Type, + Int8Type, + IntegerType, + LargeBinaryType, + LargeStringType, + MonthDayNanoIntervalType, + NullType, + StringType, + StringViewType, + SupportsArrowSchema, + TimeUnit, + UInt16Type, + UInt32Type, + UInt64Type, + UInt8Type, +) from pyarrow.lib import ( Array, Buffer, @@ -65,29 +91,6 @@ _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"] @@ -169,22 +172,11 @@ class FixedSizeListType(DataType): 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: ... + def index_type(self) -> IntegerType: ... @property def value_type(self) -> DataType: ... @@ -543,7 +535,7 @@ def map_( keys_sorted: bool = False, ) -> MapType: ... def dictionary( - index_type: _IndexType | _DataTypeAliasInput, + index_type: IntegerType | _DataTypeAliasInput, value_type: _DataTypeLike, ordered: bool = False, ) -> DictionaryType: ... @@ -586,28 +578,6 @@ __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", diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi index 60d7b9e6a5b4..d1b3c6840e2b 100644 --- a/python/pyarrow-stubs/_stubs_typing.pyi +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -114,3 +114,41 @@ class NullableCollection(Sized, Container[Any], Protocol): def __iter__(self) -> Iterator[Any]: ... def __len__(self) -> int: ... def __contains__(self, item: Any, /) -> bool: ... + +# Data types without a class of their own at runtime: their factories, e.g. +# pyarrow.int32(), return plain DataType instances. These classes only exist +# for type checking, where they distinguish the results of the factories and +# the narrowing of the pyarrow.types.is_* checks. +class NullType(lib.DataType): ... +class BoolType(lib.DataType): ... +class UInt8Type(lib.DataType): ... +class Int8Type(lib.DataType): ... +class UInt16Type(lib.DataType): ... +class Int16Type(lib.DataType): ... +class UInt32Type(lib.DataType): ... +class Int32Type(lib.DataType): ... +class UInt64Type(lib.DataType): ... +class Int64Type(lib.DataType): ... +class Float16Type(lib.DataType): ... +class Float32Type(lib.DataType): ... +class Float64Type(lib.DataType): ... +class Date32Type(lib.DataType): ... +class Date64Type(lib.DataType): ... +class MonthDayNanoIntervalType(lib.DataType): ... +class StringType(lib.DataType): ... +class LargeStringType(lib.DataType): ... +class StringViewType(lib.DataType): ... +class BinaryType(lib.DataType): ... +class LargeBinaryType(lib.DataType): ... +class BinaryViewType(lib.DataType): ... + +IntegerType: TypeAlias = ( + UInt8Type + | Int8Type + | UInt16Type + | Int16Type + | UInt32Type + | Int32Type + | UInt64Type + | Int64Type +) diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi index d17d05189c76..9ec39728d811 100644 --- a/python/pyarrow-stubs/lib.pyi +++ b/python/pyarrow-stubs/lib.pyi @@ -19,7 +19,7 @@ # 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 typing import Any, final from pyarrow.__lib.types import * # noqa: F401,F403 from pyarrow.__lib.array import * # noqa: F401,F403 @@ -33,6 +33,7 @@ from pyarrow.__lib.scalar import * # noqa: F401,F403 from pyarrow.__lib.table import * # noqa: F401,F403 from pyarrow.__lib.tensor import * # noqa: F401,F403 +@final class MonthDayNano(tuple): months: int days: int diff --git a/python/pyarrow-stubs/types.pyi b/python/pyarrow-stubs/types.pyi index 609aafd7f7e4..7fbbe0a4ec03 100644 --- a/python/pyarrow-stubs/types.pyi +++ b/python/pyarrow-stubs/types.pyi @@ -32,12 +32,7 @@ else: import pyarrow.lib as lib from pyarrow.__lib.types import ( - BinaryType, - BinaryViewType, - BoolType, DataType, - Date32Type, - Date64Type, Decimal32Type, Decimal64Type, Decimal128Type, @@ -47,6 +42,24 @@ from pyarrow.__lib.types import ( DurationType, FixedSizeBinaryType, FixedSizeListType, + LargeListType, + LargeListViewType, + ListType, + ListViewType, + MapType, + RunEndEncodedType, + SparseUnionType, + StructType, + Time32Type, + Time64Type, + TimestampType, +) +from pyarrow._stubs_typing import ( + BinaryType, + BinaryViewType, + BoolType, + Date32Type, + Date64Type, Float16Type, Float32Type, Float64Type, @@ -54,23 +67,13 @@ from pyarrow.__lib.types import ( Int16Type, Int32Type, Int64Type, + IntegerType, LargeBinaryType, - LargeListType, - LargeListViewType, LargeStringType, - ListType, - ListViewType, - MapType, MonthDayNanoIntervalType, NullType, - RunEndEncodedType, - SparseUnionType, StringType, StringViewType, - StructType, - Time32Type, - Time64Type, - TimestampType, UInt8Type, UInt16Type, UInt32Type, @@ -79,7 +82,6 @@ from pyarrow.__lib.types import ( _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 @@ -101,7 +103,7 @@ _Nested: TypeAlias = ( def is_null(t: DataType) -> TypeIs[NullType]: ... def is_boolean(t: DataType) -> TypeIs[BoolType]: ... -def is_integer(t: DataType) -> TypeIs[_Integer]: ... +def is_integer(t: DataType) -> TypeIs[IntegerType]: ... def is_signed_integer(t: DataType) -> TypeIs[_SignedInteger]: ... def is_unsigned_integer(t: DataType) -> TypeIs[_UnsignedInteger]: ... def is_int8(t: DataType) -> TypeIs[Int8Type]: ... diff --git a/python/pyarrow/fs.py b/python/pyarrow/fs.py index 482650699ea5..670ccaaf2455 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 # type: ignore[import-not-found,import-untyped] + import fsspec except ImportError: pass else: diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 71aa27a7cbb5..b597d12bb97f 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -18,11 +18,11 @@ from collections.abc import Iterable import datetime import decimal -import hypothesis as h # type: ignore[import-not-found] -import hypothesis.strategies as st # type: ignore[import-not-found] +import hypothesis as h +import hypothesis.strategies as st import itertools import operator -import pytest # type: ignore[import-not-found] +import pytest import struct import subprocess import sys @@ -1303,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) # type: ignore[arg-type] + result = pa.MapArray.from_arrays(offsets, keys, items) 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, # type: ignore[arg-type] + result = pa.MapArray.from_arrays(offsets, keys, items, 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() # type: ignore[arg-type] + offsets, keys, items, pa.string() ) with pytest.raises(pa.ArrowTypeError, match='Mismatching map items type'): - pa.MapArray.from_arrays(offsets, keys, items, pa.map_( # type: ignore[arg-type] + pa.MapArray.from_arrays(offsets, keys, items, pa.map_( keys.type, # Larger than the original i4 pa.int64() @@ -1358,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_( # type: ignore[arg-type] + pa.MapArray.from_arrays(offsets, keys, items, pa.map_( keys.type, items.type), mask=pa.array([False, True, False], type=pa.bool_()) diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index 19d33a6e2e05..f429172600a7 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -2457,7 +2457,6 @@ 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 @@ -2480,7 +2479,6 @@ 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(): @@ -2505,7 +2503,6 @@ 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_()), diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index 2dbfec5815f4..0cd004a5a97e 100644 --- a/python/pyarrow/tests/test_schema.py +++ b/python/pyarrow/tests/test_schema.py @@ -696,7 +696,6 @@ def test_empty_table(): assert table.schema == schema -@pytest.mark.numpy @pytest.mark.pandas def test_schema_from_pandas(): import pandas as pd diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index 2ce48b651b14..7cd5d853ab92 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -26,8 +26,7 @@ import pyarrow as pa try: - from scipy.sparse import ( # type: ignore[reportMissingModuleSource] - csr_array, coo_array, csr_matrix, coo_matrix) + from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix except ImportError: coo_matrix = None # type: ignore[assignment, misc] csr_matrix = None # type: ignore[assignment, misc] @@ -35,7 +34,7 @@ coo_array = None # type: ignore[assignment, misc] try: - import sparse # type: ignore[import-untyped, import-not-found] + import sparse # pyright: ignore[reportMissingImports] except ImportError: sparse = None diff --git a/python/pyproject.toml b/python/pyproject.toml index 5d01b0b307ec..2e8699b60aa9 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -114,14 +114,19 @@ fallback_version = '27.0.0a0' # TODO: Enable type checking once stubs are merged [tool.mypy] -files = ["pyarrow-stubs"] -mypy_path = "$MYPY_CONFIG_FILE_DIR/pyarrow-stubs" -exclude = [ - "^pyarrow/", - "^benchmarks/", - "^examples/", - "^scripts/", -] +# mypy can't find a "-stubs" package on mypy_path, but it maps pyarrow-stubs/ +# to the pyarrow module when checking it. Pass pyarrow-stubs as well when +# checking other files, e.g. `mypy pyarrow-stubs pyarrow/tests/test_array.py`. +files = ["pyarrow-stubs", "pyarrow-stubs-tests"] +# The type annotation tests use `# type: ignore[...]` to assert that some +# usage is rejected, so an ignore that is no longer needed is an error there. +warn_unused_ignores = true + +# The stubs ignore errors that depend on the installed packages, e.g. the +# imports of optional dependencies. +[[tool.mypy.overrides]] +module = ["pyarrow.*"] +warn_unused_ignores = false # Modules that will be annotated in later PRs (GH-32609). Treat them as Any # rather than type checking their sources. @@ -129,33 +134,27 @@ exclude = [ module = ["pyarrow._compute", "pyarrow.compute", "pyarrow.interchange.*"] follow_imports = "skip" +# Optional dependencies without type hints. +[[tool.mypy.overrides]] +module = ["fsspec", "sparse"] +ignore_missing_imports = true + # TODO: Enable type checking once stubs are merged [tool.pyright] pythonPlatform = "All" pythonVersion = "3.11" -include = ["pyarrow-stubs"] -exclude = [ - "pyarrow", - "benchmarks", - "examples", - "scripts", - "build", -] -stubPath = "pyarrow-stubs" +include = ["pyarrow-stubs", "pyarrow-stubs-tests"] typeCheckingMode = "basic" # The type annotation tests use `# type: ignore[...]` to assert that some # usage is rejected, so an ignore that is no longer needed is an error there. executionEnvironments = [ - { root = "pyarrow-stubs/tests", reportUnnecessaryTypeIgnoreComment = "error" }, + { root = "pyarrow-stubs-tests", extraPaths = ["."], 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. @@ -167,31 +166,25 @@ replace-imports-with-any = [ ] [tool.ty.src] -include = ["pyarrow-stubs"] -exclude = [ - "pyarrow", - "benchmarks", - "examples", - "scripts", -] +include = ["pyarrow-stubs", "pyarrow-stubs-tests"] # numpy, pandas and scipy are optional dependencies. A `ty: ignore` comment on # their imports would be reported as unused whenever they are installed. [[tool.ty.overrides]] -include = ["pyarrow-stubs/pyarrow/**"] +include = ["pyarrow-stubs/**"] [tool.ty.overrides.rules] import-not-found = "ignore" # ty doesn't honor `# type: ignore[]` comments, so disable the # rules that those comments suppress for mypy. [[tool.ty.overrides]] -include = ["pyarrow-stubs/pyarrow/__lib/array.pyi"] +include = ["pyarrow-stubs/__lib/array.pyi"] [tool.ty.overrides.rules] # Some array classes override methods with a different signature at runtime. invalid-method-override = "ignore" [[tool.ty.overrides]] -include = ["pyarrow-stubs/tests/**"] +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" diff --git a/python/requirements-stubs.txt b/python/requirements-stubs.txt new file mode 100644 index 000000000000..37d808571298 --- /dev/null +++ b/python/requirements-stubs.txt @@ -0,0 +1,17 @@ +# Type checkers and the type stubs of PyArrow's dependencies, pinned so that +# everyone gets the same results; see the "Type Checking" section of +# docs/source/developers/python/development.rst. Keep the versions in sync +# with the python-stubs hooks in .pre-commit-config.yaml. +fsspec==2026.9.0 +mypy==2.4.0 +pandas-stubs==3.0.5.260914 +pyright==1.1.414 +scipy-stubs==1.17.1.5 +ty==0.0.99 +types-cffi==2.1.0.20260827 +types-psutil==7.2.2.20260906 +types-python-dateutil==2.9.0.20261021 +types-requests==2.33.0.20261006 +# Annotated dependencies of the type stub tests, usually installed already. +numpy +pytest diff --git a/python/scripts/run_stubtest.py b/python/scripts/run_stubtest.py new file mode 100644 index 000000000000..3ec8a70b0fa9 --- /dev/null +++ b/python/scripts/run_stubtest.py @@ -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. + +""" +Check the type stubs in pyarrow-stubs/ against the installed PyArrow with +mypy's stubtest. + +Usage: + python scripts/run_stubtest.py [stubtest options] + +Stubs that aren't written yet are skipped (--ignore-missing-stub), and the +known mismatches are listed in pyarrow-stubs-tests/stubtest_allowlist.txt. +Pass --generate-allowlist to print the entries for the current mismatches. +""" + +import os +import sys +import tempfile +from pathlib import Path + +from mypy import stubtest + +import pyarrow.lib + +PYTHON_DIR = Path(__file__).resolve().parents[1] +STUBS_DIR = PYTHON_DIR / "pyarrow-stubs" +ALLOWLIST = PYTHON_DIR / "pyarrow-stubs-tests" / "stubtest_allowlist.txt" + + +def main(argv): + # The stubs of pyarrow.lib are split into the stub-only pyarrow.__lib + # package, one module per .pxi file. Alias those modules to pyarrow.lib, + # otherwise stubtest can't import them and skips everything they define. + for stub in (STUBS_DIR / "__lib").glob("*.pyi"): + name = "pyarrow.__lib" + if stub.stem != "__init__": + name += f".{stub.stem}" + sys.modules[name] = pyarrow.lib + + with tempfile.TemporaryDirectory() as tmpdir: + # mypy doesn't look for "-stubs" packages on MYPYPATH, so give it the + # stubs under the name of the package. + os.symlink(STUBS_DIR, Path(tmpdir) / "pyarrow") + os.environ["MYPYPATH"] = tmpdir + options = stubtest.parse_options([ + "pyarrow", + "--ignore-missing-stub", + # Every Cython extension type is a disjoint base (PEP 800). + "--ignore-disjoint-bases", + "--allowlist", str(ALLOWLIST), + "--mypy-config-file", str(PYTHON_DIR / "pyproject.toml"), + *argv, + ]) + return stubtest.test_stubs(options) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) From 79a8f1540095f5f5210d627c7c4a1b7c3b418935 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 00:51:18 +0100 Subject: [PATCH 63/93] Re-export annotated pyarrow.lib names in pyarrow and use pa.* in tests Re-export the names that pyarrow/__init__.py imports from pyarrow.lib and that have stubs already, so that pa.* resolves to their annotations. The module __getattr__ stays as the fallback for the names without stubs. Use the public pa.* names in the type stub tests instead of lib.*. --- docs/source/developers/python/development.rst | 10 +- python/pyarrow-stubs-tests/_check.py | 3 +- python/pyarrow-stubs-tests/conftest.py | 6 +- python/pyarrow-stubs-tests/test_array.py | 216 ++++++++--------- python/pyarrow-stubs-tests/test_table.py | 220 +++++++++--------- python/pyarrow-stubs-tests/test_types.py | 154 ++++++------ python/pyarrow-stubs/__init__.pyi | 206 ++++++++++++++++ 7 files changed, 510 insertions(+), 305 deletions(-) diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index af0c07c7d6fc..e792db89889c 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -223,17 +223,17 @@ verifies statically, with a runtime ``isinstance`` check: .. code-block:: python - import pyarrow.lib as lib + import pyarrow as pa 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) + lst = pa.list_(pa.int32()) + check(assert_type(lst, pa.ListType), pa.ListType) + check(assert_type(lst.value_type, pa.DataType), pa.DataType) if TYPE_CHECKING_INVALID_USAGE: - lib.list_(1) # type: ignore[call-overload] + pa.list_(1) # type: ignore[call-overload] The type checkers check these files as part of the regular runs above, and pytest runs them against the installed PyArrow to make sure the stubs agree diff --git a/python/pyarrow-stubs-tests/_check.py b/python/pyarrow-stubs-tests/_check.py index 508f3c44a142..d499c5f59949 100644 --- a/python/pyarrow-stubs-tests/_check.py +++ b/python/pyarrow-stubs-tests/_check.py @@ -19,8 +19,7 @@ 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) + check(assert_type(pa.list_(pa.int32()), pa.ListType), pa.ListType) ``assert_type`` is verified by every type checker we run in CI (mypy, pyright and ty) and is a no-op at runtime. ``check`` runs under pytest and diff --git a/python/pyarrow-stubs-tests/conftest.py b/python/pyarrow-stubs-tests/conftest.py index 7871f5c57118..fa4936a64ede 100644 --- a/python/pyarrow-stubs-tests/conftest.py +++ b/python/pyarrow-stubs-tests/conftest.py @@ -18,14 +18,14 @@ import sys import types -import pyarrow.lib as lib +import pyarrow as pa -def _stub_only_type(name: str) -> type[lib.DataType]: +def _stub_only_type(name: str) -> type[pa.DataType]: # The factories of the stub-only DataType subclasses (Int32Type, ...) # return plain DataType instances at runtime. if name.endswith("Type") and not name.startswith("_"): - return lib.DataType + return pa.DataType raise AttributeError(name) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index b1f9b9a3c04e..08fdb5df8de3 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -21,7 +21,7 @@ import pytest from typing_extensions import CapsuleType -import pyarrow.lib as lib +import pyarrow as pa from pyarrow._compute import CastOptions from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check @@ -30,38 +30,38 @@ 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) + arr = check(cast(pa.Int64Array, pa.array([1, 2, None])), pa.Int64Array) - check(assert_type(arr.type, lib.DataType), lib.DataType) + check(assert_type(arr.type, pa.DataType), pa.DataType) check(assert_type(arr.null_count, int), int) check(assert_type(len(arr), int), int) check(assert_type(arr.to_pylist(), list[Any]), list) # Scalars aren't annotated yet, so indexing returns Any | Int64Array. - check(assert_type(arr[0], Any | lib.Int64Array), lib.Int64Scalar) - check(assert_type(arr[1:], Any | lib.Int64Array), lib.Int64Array) + check(assert_type(arr[0], Any | pa.Int64Array), pa.Int64Scalar) + check(assert_type(arr[1:], Any | pa.Int64Array), pa.Int64Array) # Methods returning the same kind of array keep the subclass. - check(assert_type(arr.slice(1), 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(arr.cast(lib.int32(), safe=None), lib.Array), lib.Int32Array) + check(assert_type(arr.slice(1), pa.Int64Array), pa.Int64Array) + check(assert_type(arr.fill_null(0), pa.Int64Array), pa.Int64Array) + check(assert_type(arr.drop_null(), pa.Int64Array), pa.Int64Array) + + check(assert_type(arr.is_null(), pa.BooleanArray), pa.BooleanArray) + check(assert_type(arr.cast(pa.float64()), pa.Array), pa.DoubleArray) + check(assert_type(arr.cast("float64"), pa.Array), pa.DoubleArray) + check(assert_type(arr.cast(pa.int32(), safe=None), pa.Array), pa.Int32Array) # without target_type argument, as CastOptions already specifies it - check(assert_type(arr.cast(options=CastOptions.safe(lib.int32())), lib.Array), - lib.Int32Array) - check(assert_type(arr.view("uint64"), lib.Array), lib.UInt64Array) + check(assert_type(arr.cast(options=CastOptions.safe(pa.int32())), pa.Array), + pa.Int32Array) + check(assert_type(arr.view("uint64"), pa.Array), pa.UInt64Array) check(assert_type(arr.to_string(element_size_limit=5), str), str) - check(assert_type(arr.to_tensor(allow_nulls=True), lib.Tensor), lib.Tensor) - 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) - - other = check(cast(lib.StringArray, lib.array(["a", "b", None])), lib.StringArray) + check(assert_type(arr.to_tensor(allow_nulls=True), pa.Tensor), pa.Tensor) + check(assert_type(pa.concat_arrays([arr, arr]), pa.Array), + pa.Int64Array) + check(assert_type(pa.Array.from_pandas([1, 2], None, pa.int64()), + pa.Array), + pa.Int64Array) + + other = check(cast(pa.StringArray, pa.array(["a", "b", None])), pa.StringArray) check(assert_type(arr.equals(arr), bool), bool) check(assert_type(arr.equals(other), bool), bool) # Arrays of any type can diff() @@ -72,20 +72,20 @@ def test_array_numpy() -> None: pytest.importorskip("numpy") import numpy as np - arr = cast(lib.Int64Array, lib.array([1, 2])) - check(assert_type(arr[np.int64(0)], Any | lib.Int64Array), lib.Int64Scalar) - check(assert_type(lib.infer_type(np.array([1])), lib.DataType), lib.DataType) + arr = cast(pa.Int64Array, pa.array([1, 2])) + check(assert_type(arr[np.int64(0)], Any | pa.Int64Array), pa.Int64Scalar) + check(assert_type(pa.infer_type(np.array([1])), pa.DataType), pa.DataType) offsets = np.array([0, 1, 2], dtype="int32") - check(assert_type(lib.ListArray.from_arrays(offsets, [1, 2]), lib.ListArray), - lib.ListArray) + check(assert_type(pa.ListArray.from_arrays(offsets, [1, 2]), pa.ListArray), + pa.ListArray) if TYPE_CHECKING_INVALID_USAGE: - dtype = lib.dictionary(lib.int8(), lib.string()) + dtype = pa.dictionary(pa.int8(), pa.string()) strs = np.array(["a"]) - lib.DictionaryArray.from_buffers(dtype, 1, [], strs) # type: ignore[arg-type] + pa.DictionaryArray.from_buffers(dtype, 1, [], strs) # type: ignore[arg-type] def test_array_dlpack() -> None: - arr = cast(lib.Int64Array, lib.array([1, 2])) + arr = cast(pa.Int64Array, pa.array([1, 2])) capsule = arr.__dlpack__( max_version=(1, 0), dl_device=arr.__dlpack_device__(), @@ -95,11 +95,11 @@ def test_array_dlpack() -> None: def test_import_from_c_capsule() -> None: - arr = cast(lib.Int64Array, lib.array([1, 2])) + arr = cast(pa.Int64Array, pa.array([1, 2])) capsules = arr.__arrow_c_array__() check(assert_type(capsules, tuple[CapsuleType, CapsuleType]), tuple) - imported = lib.BooleanArray._import_from_c_capsule(*capsules) - check(assert_type(imported, lib.Array), lib.Int64Array) + imported = pa.BooleanArray._import_from_c_capsule(*capsules) + check(assert_type(imported, pa.Array), pa.Int64Array) check(assert_type(arr.__arrow_c_device_array__(), tuple[CapsuleType, CapsuleType]), @@ -107,97 +107,97 @@ def test_import_from_c_capsule() -> None: def test_from_buffers() -> None: - values = cast(lib.Int64Array, lib.array([1, 2])) - offsets = lib.array([0, 2], lib.int32()).buffers()[1] + values = cast(pa.Int64Array, pa.array([1, 2])) + offsets = pa.array([0, 2], pa.int32()).buffers()[1] # Children can be any collection of arrays, like a tuple - list_arr = lib.Array.from_buffers(lib.list_(lib.int64()), 1, [None, offsets], - children=(values,)) - check(assert_type(list_arr, lib.Array), lib.ListArray) + list_arr = pa.Array.from_buffers(pa.list_(pa.int64()), 1, [None, offsets], + children=(values,)) + check(assert_type(list_arr, pa.Array), pa.ListArray) 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) - - list_type = lib.list_(lib.int64()) - check(assert_type(lib.ListArray.from_arrays([0, 2], [1, 2], list_type), - lib.ListArray), - lib.ListArray) - check(assert_type(lib.LargeListArray.from_arrays([0, 2], [1, 2]), - lib.LargeListArray), - lib.LargeListArray) - check(assert_type(lib.ListViewArray.from_arrays([0], [2], [1, 2]), - lib.ListViewArray), - lib.ListViewArray) - check(assert_type(lib.LargeListViewArray.from_arrays([0], [2], [1, 2]), - lib.LargeListViewArray), - lib.LargeListViewArray) - check(assert_type(lib.FixedSizeListArray.from_arrays([1, 2], 2), - lib.FixedSizeListArray), - lib.FixedSizeListArray) + lst = pa.ListArray.from_arrays([0, 1, 3], [1, 2, 3]) + check(assert_type(lst.value_lengths(), pa.Int32Array), pa.Int32Array) + + large = pa.LargeListArray.from_arrays([0, 1, 3], pa.array([1, 2, 3])) + check(assert_type(large.value_lengths(), pa.Int64Array), pa.Int64Array) + + list_type = pa.list_(pa.int64()) + check(assert_type(pa.ListArray.from_arrays([0, 2], [1, 2], list_type), + pa.ListArray), + pa.ListArray) + check(assert_type(pa.LargeListArray.from_arrays([0, 2], [1, 2]), + pa.LargeListArray), + pa.LargeListArray) + check(assert_type(pa.ListViewArray.from_arrays([0], [2], [1, 2]), + pa.ListViewArray), + pa.ListViewArray) + check(assert_type(pa.LargeListViewArray.from_arrays([0], [2], [1, 2]), + pa.LargeListViewArray), + pa.LargeListViewArray) + check(assert_type(pa.FixedSizeListArray.from_arrays([1, 2], 2), + pa.FixedSizeListArray), + pa.FixedSizeListArray) # MapArray is a ListArray (of key-item structs) at runtime. - map_arr = 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) - map_type = lib.map_(lib.int64(), lib.string()) - check(assert_type(lib.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"], map_type), - lib.MapArray), - lib.ListArray) + map_arr = pa.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"]) + check(assert_type(map_arr, pa.MapArray), pa.ListArray) + check(assert_type(map_arr.offsets, pa.Int32Array), pa.Int32Array) + check(assert_type(map_arr.values, pa.Array), pa.StructArray) + check(assert_type(map_arr.keys, pa.Array), pa.Int64Array) + map_type = pa.map_(pa.int64(), pa.string()) + check(assert_type(pa.MapArray.from_arrays([0, 2], [1, 2], ["a", "b"], map_type), + pa.MapArray), + pa.ListArray) def test_dictionary_array() -> None: - arr = check(cast(lib.StringArray, lib.array(["a", "b", "a"])), - lib.StringArray) + arr = check(cast(pa.StringArray, pa.array(["a", "b", "a"])), + pa.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) + check(assert_type(dict_arr, pa.DictionaryArray), pa.DictionaryArray) + check(assert_type(dict_arr.indices, pa.IntegerArray), pa.Int32Array) + check(assert_type(dict_arr.dictionary, pa.Array), pa.StringArray) + check(assert_type(dict_arr.dictionary_decode(), pa.Array), + pa.StringArray) - rebuilt = lib.DictionaryArray.from_buffers( + rebuilt = pa.DictionaryArray.from_buffers( dict_arr.type, len(dict_arr), dict_arr.indices.buffers(), dict_arr.dictionary ) - check(assert_type(rebuilt, lib.DictionaryArray), lib.DictionaryArray) + check(assert_type(rebuilt, pa.DictionaryArray), pa.DictionaryArray) def test_run_end_encoded_array() -> None: - ree = lib.RunEndEncodedArray.from_arrays([2, 3], ["a", "b"]) - check(assert_type(ree, lib.RunEndEncodedArray), lib.RunEndEncodedArray) + ree = pa.RunEndEncodedArray.from_arrays([2, 3], ["a", "b"]) + check(assert_type(ree, pa.RunEndEncodedArray), pa.RunEndEncodedArray) check( assert_type(ree.run_ends, - lib.Int16Array | lib.Int32Array | lib.Int64Array), - lib.Int64Array, + pa.Int16Array | pa.Int32Array | pa.Int64Array), + pa.Int64Array, ) - check(assert_type(ree.values, lib.Array), lib.StringArray) + check(assert_type(ree.values, pa.Array), pa.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) + chunked = pa.chunked_array([[1, 2], [3]]) + check(assert_type(chunked, pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.type, pa.DataType), pa.DataType) check(assert_type(chunked.num_chunks, int), int) - check(assert_type(chunked.chunk(0), lib.Array), lib.Int64Array) + check(assert_type(chunked.chunk(0), pa.Array), pa.Int64Array) # Scalars aren't annotated yet, so indexing returns Any | ChunkedArray. - check(assert_type(chunked[0], Any | 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) + check(assert_type(chunked[0], Any | pa.ChunkedArray), pa.Int64Scalar) + check(assert_type(chunked[1:], Any | pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.chunks, list[pa.Array]), list) + check(assert_type(chunked.combine_chunks(), pa.Array), pa.Int64Array) + check(assert_type(chunked.unique(), pa.Array), pa.Int64Array) + check(assert_type(chunked.slice(1), pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked.cast(pa.float64()), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(chunked.is_null(), pa.ChunkedArray), pa.ChunkedArray) def test_invalid_usage() -> None: - arr = cast(lib.Int64Array, lib.array([1, 2])) + arr = cast(pa.Int64Array, pa.array([1, 2])) if TYPE_CHECKING_INVALID_USAGE: arr.cast(1) # type: ignore[arg-type] arr.slice("1") # type: ignore[arg-type] @@ -211,17 +211,17 @@ def test_invalid_usage() -> None: arr.dictionary_encode("drop") # type: ignore[arg-type] arr.to_pylist(maps_as_pydicts="loose") # type: ignore[arg-type] # Generators aren't accepted, only Sequences - lib.infer_type(x for x in [1]) # type: ignore[arg-type] + pa.infer_type(x for x in [1]) # type: ignore[arg-type] def test_buffer_constructors_invalid_usage() -> None: - types = cast(lib.Int8Array, lib.array([0], lib.int8())) - children: list[lib.Array] = [lib.array([1])] - strings = cast(lib.StringArray, lib.array(["a"])) - off, data = cast(list[lib.Buffer], strings.buffers()[1:]) + types = cast(pa.Int8Array, pa.array([0], pa.int8())) + children: list[pa.Array] = [pa.array([1])] + strings = cast(pa.StringArray, pa.array(["a"])) + off, data = cast(list[pa.Buffer], strings.buffers()[1:]) if TYPE_CHECKING_INVALID_USAGE: # Only lists are accepted for children and type codes - lib.UnionArray.from_sparse(types, tuple(children)) # type: ignore[arg-type] - lib.UnionArray.from_sparse(types, children, None, - types) # type: ignore[arg-type] - lib.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] + pa.UnionArray.from_sparse(types, tuple(children)) # type: ignore[arg-type] + pa.UnionArray.from_sparse(types, children, None, + types) # type: ignore[arg-type] + pa.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] diff --git a/python/pyarrow-stubs-tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py index 0e36b73b4a1f..1c9e80887472 100644 --- a/python/pyarrow-stubs-tests/test_table.py +++ b/python/pyarrow-stubs-tests/test_table.py @@ -21,105 +21,105 @@ import pytest from typing_extensions import CapsuleType -import pyarrow.lib as lib +import pyarrow as pa from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check class _PyArrowArrayProvider: - def __arrow_array__(self, type: lib.DataType | None = None) -> lib.Array: - return lib.array([1, 2], type=type) + def __arrow_array__(self, type: pa.DataType | None = None) -> pa.Array: + return pa.array([1, 2], type=type) class _ArrowArrayProvider: def __arrow_c_array__( self, requested_schema: CapsuleType | None = None ) -> tuple[CapsuleType, CapsuleType]: - return lib.array([1, 2]).__arrow_c_array__(requested_schema) + return pa.array([1, 2]).__arrow_c_array__(requested_schema) class _ArrowDeviceArrayProvider: def __arrow_c_device_array__( self, requested_schema: CapsuleType | None = None, **kwargs: Any ) -> tuple[CapsuleType, CapsuleType]: - return lib.array([1, 2]).__arrow_c_device_array__(requested_schema, **kwargs) + return pa.array([1, 2]).__arrow_c_device_array__(requested_schema, **kwargs) class _ArrowStreamProvider: def __arrow_c_stream__( self, requested_schema: CapsuleType | None = None ) -> CapsuleType: - return lib.chunked_array([[1, 2]]).__arrow_c_stream__(requested_schema) + return pa.chunked_array([[1, 2]]).__arrow_c_stream__(requested_schema) def test_table() -> None: - table = lib.table({"a": [1, 2], "b": ["x", "y"]}) - check(assert_type(table, lib.Table), lib.Table) + table = pa.table({"a": [1, 2], "b": ["x", "y"]}) + check(assert_type(table, pa.Table), pa.Table) check(assert_type(table.num_rows, int), int) check(assert_type(table.num_columns, int), int) check(assert_type(table.shape, tuple[int, int]), tuple) check(assert_type(table.column_names, list[str]), list) - check(assert_type(table.schema, lib.Schema), lib.Schema) - check(assert_type(table.field("a"), lib.Field), lib.Field) + check(assert_type(table.schema, pa.Schema), pa.Schema) + check(assert_type(table.field("a"), pa.Field), pa.Field) # Columns aren't typed per tabular class, so Table and RecordBatch both # return the union here. - check(assert_type(table.column("a"), 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.column("a"), pa.Array | pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(table["a"], pa.Array | pa.ChunkedArray | pa.Table), + pa.ChunkedArray) + check(assert_type(table[0], pa.Array | pa.ChunkedArray | pa.Table), + pa.ChunkedArray) + check(assert_type(table[:1], pa.Array | pa.ChunkedArray | pa.Table), + pa.Table) check(assert_type(table.set_column(0, "a", [_ArrowDeviceArrayProvider()]), - lib.Table), - lib.Table) - check(assert_type(table.append_column("c", ([3], [4])), lib.Table), lib.Table) + pa.Table), + pa.Table) + check(assert_type(table.append_column("c", ([3], [4])), pa.Table), pa.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.select(["a", 1]), lib.Table), lib.Table) + check(assert_type(table.slice(0, 1), pa.Table), pa.Table) + check(assert_type(table.select(["a"]), pa.Table), pa.Table) + check(assert_type(table.select(["a", 1]), pa.Table), pa.Table) check(assert_type(table.equals(None), bool), bool) - check(assert_type(table.combine_chunks(), lib.Table), lib.Table) - check(assert_type(table.to_batches(), list[lib.RecordBatch]), list) + check(assert_type(table.combine_chunks(), pa.Table), pa.Table) + check(assert_type(table.to_batches(), list[pa.RecordBatch]), list) check(assert_type(table.to_pydict(), dict[str, list]), dict) - check(assert_type(lib.Table.from_arrays([(1, 2)], names=["a"]), lib.Table), - lib.Table) + check(assert_type(pa.Table.from_arrays([(1, 2)], names=["a"]), pa.Table), + pa.Table) - chunked = lib.chunked_array([[1, 2]]) - check(assert_type(table.add_column(0, "c", chunked), lib.Table), lib.Table) - check(assert_type(table.add_column(0, "c", lib.array([1, 2])), lib.Table), - lib.Table) - check(assert_type(table.add_column(0, "c", [[1], [2]]), lib.Table), - lib.Table) - check(assert_type(table.add_column(0, "c", [lib.array([1, 2])]), lib.Table), - lib.Table) + chunked = pa.chunked_array([[1, 2]]) + check(assert_type(table.add_column(0, "c", chunked), pa.Table), pa.Table) + check(assert_type(table.add_column(0, "c", pa.array([1, 2])), pa.Table), + pa.Table) + check(assert_type(table.add_column(0, "c", [[1], [2]]), pa.Table), + pa.Table) + check(assert_type(table.add_column(0, "c", [pa.array([1, 2])]), pa.Table), + pa.Table) - table = lib.table(([1],), names=["a"]) - check(assert_type(table, lib.Table), lib.Table) + table = pa.table(([1],), names=["a"]) + check(assert_type(table, pa.Table), pa.Table) - chunked_struct = lib.chunked_array([[{"x": 1}]]) - check(assert_type(lib.Table.from_struct_array(chunked_struct), lib.Table), - lib.Table) + chunked_struct = pa.chunked_array([[{"x": 1}]]) + check(assert_type(pa.Table.from_struct_array(chunked_struct), pa.Table), + pa.Table) # The reader isn't annotated yet, but must not reject valid usage. - check(table.to_reader().read_all(), lib.Table) + check(table.to_reader().read_all(), pa.Table) for batch in table.to_reader(): - check(batch, lib.RecordBatch) + check(batch, pa.RecordBatch) def test_table_from_arrays_protocols() -> None: - check(assert_type(lib.Table.from_arrays([_PyArrowArrayProvider()], names=["a"]), - lib.Table), - lib.Table) - check(assert_type(lib.Table.from_arrays([_ArrowArrayProvider()], names=["a"]), - lib.Table), - lib.Table) - check(assert_type(lib.Table.from_arrays([_ArrowDeviceArrayProvider()], names=["a"]), - lib.Table), - lib.Table) + check(assert_type(pa.Table.from_arrays([_PyArrowArrayProvider()], names=["a"]), + pa.Table), + pa.Table) + check(assert_type(pa.Table.from_arrays([_ArrowArrayProvider()], names=["a"]), + pa.Table), + pa.Table) + check(assert_type(pa.Table.from_arrays([_ArrowDeviceArrayProvider()], names=["a"]), + pa.Table), + pa.Table) def test_table_from_arrays_numpy() -> None: @@ -127,98 +127,98 @@ def test_table_from_arrays_numpy() -> None: import numpy as np columns = np.array([[1, 2], [3, 4]]) - check(assert_type(lib.Table.from_arrays(columns, names=["a", "b"]), lib.Table), - lib.Table) + check(assert_type(pa.Table.from_arrays(columns, names=["a", "b"]), pa.Table), + pa.Table) def test_record_batch() -> None: - batch = lib.record_batch({"a": [1, 2]}) # from a dict - check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + batch = pa.record_batch({"a": [1, 2]}) # from a dict + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) check(assert_type(batch.num_rows, int), int) - check(assert_type(batch.column(0), 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) - check(assert_type(batch.add_column(0, "b", lib.array([3, 4])), - lib.RecordBatch), - lib.RecordBatch) - check(assert_type(batch.add_column(0, "b", [3, 4]), lib.RecordBatch), - lib.RecordBatch) - check(assert_type(batch.append_column("b", [3, 4]), lib.RecordBatch), - lib.RecordBatch) + check(assert_type(batch.column(0), pa.Array | pa.ChunkedArray), + pa.Array) + check(assert_type(batch.slice(1), pa.RecordBatch), pa.RecordBatch) + check(assert_type(pa.Table.from_batches([batch]), pa.Table), pa.Table) + check(assert_type(pa.concat_batches([batch, batch]), pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.add_column(0, "b", pa.array([3, 4])), + pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.add_column(0, "b", [3, 4]), pa.RecordBatch), + pa.RecordBatch) + check(assert_type(batch.append_column("b", [3, 4]), pa.RecordBatch), + pa.RecordBatch) check(assert_type(batch.equals(batch), bool), bool) # Any object can be compared, a non-batch returns False check(assert_type(batch.equals(42), bool), bool) - batch = lib.record_batch([lib.array([1])], names=["a"]) # from a list - check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + batch = pa.record_batch([pa.array([1])], names=["a"]) # from a list + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) struct_array = batch.to_struct_array() - check(assert_type(struct_array, lib.StructArray), lib.StructArray) - check(assert_type(lib.RecordBatch.from_struct_array(struct_array), - lib.RecordBatch), - lib.RecordBatch) + check(assert_type(struct_array, pa.StructArray), pa.StructArray) + check(assert_type(pa.RecordBatch.from_struct_array(struct_array), + pa.RecordBatch), + pa.RecordBatch) def test_record_batch_from_pandas() -> None: pd = pytest.importorskip("pandas") df = pd.DataFrame({"a": [1], "b": [2]}) # Columns only need to be iterable, e.g. a generator. - batch = lib.RecordBatch.from_pandas(df, columns=(c for c in ["a"])) - check(assert_type(batch, lib.RecordBatch), lib.RecordBatch) + batch = pa.RecordBatch.from_pandas(df, columns=(c for c in ["a"])) + check(assert_type(batch, pa.RecordBatch), pa.RecordBatch) def test_chunked_array() -> None: - chunked_array = lib.chunked_array([[1, 2], [None]]) - check(assert_type(chunked_array, lib.ChunkedArray), lib.ChunkedArray) - check(assert_type(chunked_array.cast(lib.int32(), safe=None), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(chunked_array.fill_null(0), lib.ChunkedArray), lib.ChunkedArray) - check(assert_type(chunked_array.sort(null_placement="at_start"), lib.ChunkedArray), - lib.ChunkedArray) + chunked_array = pa.chunked_array([[1, 2], [None]]) + check(assert_type(chunked_array, pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked_array.cast(pa.int32(), safe=None), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(chunked_array.fill_null(0), pa.ChunkedArray), pa.ChunkedArray) + check(assert_type(chunked_array.sort(null_placement="at_start"), pa.ChunkedArray), + pa.ChunkedArray) check(assert_type(chunked_array.equals(None), bool), bool) - arr = lib.array([1, 2]) - check(assert_type(lib.chunked_array(arr), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array([arr, [3]]), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array([_PyArrowArrayProvider()]), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array([_ArrowDeviceArrayProvider()]), - lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array(_ArrowArrayProvider()), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array([_ArrowArrayProvider()]), lib.ChunkedArray), - lib.ChunkedArray) - check(assert_type(lib.chunked_array(_ArrowStreamProvider()), lib.ChunkedArray), - lib.ChunkedArray) + arr = pa.array([1, 2]) + check(assert_type(pa.chunked_array(arr), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([arr, [3]]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_PyArrowArrayProvider()]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_ArrowDeviceArrayProvider()]), + pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array(_ArrowArrayProvider()), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array([_ArrowArrayProvider()]), pa.ChunkedArray), + pa.ChunkedArray) + check(assert_type(pa.chunked_array(_ArrowStreamProvider()), pa.ChunkedArray), + pa.ChunkedArray) def test_invalid_usage() -> None: - table = lib.table({"a": [1, 2]}) + table = pa.table({"a": [1, 2]}) streams = [_ArrowStreamProvider()] - columns = {"a": lib.array([1])}.values() # dict as non-sequence example + columns = {"a": pa.array([1])}.values() # dict as non-sequence example if TYPE_CHECKING_INVALID_USAGE: table.column(1.5) # type: ignore[arg-type] table.slice("0") # type: ignore[arg-type] - lib.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] + pa.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] # Stream providers are accepted as the whole input, not individual chunks - lib.chunked_array(streams) # type: ignore[arg-type] + pa.chunked_array(streams) # type: ignore[arg-type] table.add_column(0, "b", streams) # type: ignore[arg-type] def test_record_batch_invalid_usage() -> None: - arr = lib.array([1]) - chunked = lib.chunked_array([[{"a": 1}]]) + arr = pa.array([1]) + chunked = pa.chunked_array([[{"a": 1}]]) if TYPE_CHECKING_INVALID_USAGE: # Tuples are failing at runtime. Recheck after the fix of GH-55418 - lib.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] - lib.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] - lib.record_batch((arr,), names=["a"]) # type: ignore[arg-type] + pa.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] + pa.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] + pa.record_batch((arr,), names=["a"]) # type: ignore[arg-type] # Only a `dict` is passed on to `from_pydict`, other mappings raise - lib.record_batch(MappingProxyType({"a": [1]})) # type: ignore[arg-type] + pa.record_batch(MappingProxyType({"a": [1]})) # type: ignore[arg-type] diff --git a/python/pyarrow-stubs-tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py index 27dcd0588937..ab0219ef09e5 100644 --- a/python/pyarrow-stubs-tests/test_types.py +++ b/python/pyarrow-stubs-tests/test_types.py @@ -20,7 +20,7 @@ from typing_extensions import CapsuleType -import pyarrow.lib as lib +import pyarrow as pa from pyarrow._stubs_typing import ( BoolType, Int16Type, @@ -34,105 +34,105 @@ from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check -def _erase(t: lib.DataType) -> lib.DataType: +def _erase(t: pa.DataType) -> pa.DataType: # Hide the precise type, so we can test narrowing by the is_* functions. return t def test_primitive_factories() -> None: - check(assert_type(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) + check(assert_type(pa.bool_(), BoolType), pa.DataType) + check(assert_type(pa.int32(), Int32Type), pa.DataType) + check(assert_type(pa.string(), StringType), pa.DataType) + check(assert_type(pa.utf8(), StringType), pa.DataType) def test_temporal_types() -> None: - ts = lib.timestamp("ms", tz="UTC") - check(assert_type(ts, lib.TimestampType), lib.TimestampType) + ts = pa.timestamp("ms", tz="UTC") + check(assert_type(ts, pa.TimestampType), pa.TimestampType) check(assert_type(ts.unit, Literal["s", "ms", "us", "ns"]), str) check(assert_type(ts.tz, str | None), str) - check(assert_type(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) + check(assert_type(pa.time32("s").unit, Literal["s", "ms"]), str) + check(assert_type(pa.time64("ns").unit, Literal["us", "ns"]), str) + check(assert_type(pa.duration("us"), pa.DurationType), pa.DurationType) def test_decimal_types() -> None: - dec = lib.decimal128(10, 2) - check(assert_type(dec, lib.Decimal128Type), lib.Decimal128Type) + dec = pa.decimal128(10, 2) + check(assert_type(dec, pa.Decimal128Type), pa.Decimal128Type) check(assert_type(dec.precision, int), int) check(assert_type(dec.scale, int), int) def test_list_types() -> None: - lst = 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) + lst = pa.list_(pa.int32()) + check(assert_type(lst, pa.ListType), pa.ListType) + check(assert_type(lst.value_type, pa.DataType), pa.DataType) + check(assert_type(lst.value_field, pa.Field), pa.Field) - variable = lib.list_(lib.int32(), -1) - check(assert_type(variable, lib.ListType), lib.ListType) - fixed = lib.list_(lib.int32(), 2) - check(assert_type(fixed, lib.ListType | lib.FixedSizeListType), - lib.FixedSizeListType) + variable = pa.list_(pa.int32(), -1) + check(assert_type(variable, pa.ListType), pa.ListType) + fixed = pa.list_(pa.int32(), 2) + check(assert_type(fixed, pa.ListType | pa.FixedSizeListType), + pa.FixedSizeListType) - check(assert_type(lib.large_list(lib.int32()), lib.LargeListType), - lib.LargeListType) - check(assert_type(lib.list_view(lib.int32()), lib.ListViewType), - lib.ListViewType) + check(assert_type(pa.large_list(pa.int32()), pa.LargeListType), + pa.LargeListType) + check(assert_type(pa.list_view(pa.int32()), pa.ListViewType), + pa.ListViewType) def test_dictionary_type() -> None: - dict_type = lib.dictionary(lib.int8(), lib.string()) - check(assert_type(dict_type, lib.DictionaryType), lib.DictionaryType) + dict_type = pa.dictionary(pa.int8(), pa.string()) + check(assert_type(dict_type, pa.DictionaryType), pa.DictionaryType) check(assert_type(dict_type.ordered, bool), bool) - check(assert_type(dict_type.value_type, lib.DataType), lib.DataType) - check(assert_type(dict_type.index_type, IntegerType), lib.DataType) + check(assert_type(dict_type.value_type, pa.DataType), pa.DataType) + check(assert_type(dict_type.index_type, IntegerType), pa.DataType) def test_map_and_run_end_encoded_types() -> None: - map_type = 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) + map_type = pa.map_(pa.string(), pa.int64()) + check(assert_type(map_type, pa.MapType), pa.MapType) + check(assert_type(map_type.key_type, pa.DataType), pa.DataType) check(assert_type(map_type.keys_sorted, bool), bool) - ree = lib.run_end_encoded(lib.int16(), lib.string()) - check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) + ree = pa.run_end_encoded(pa.int16(), pa.string()) + check(assert_type(ree, pa.RunEndEncodedType), pa.RunEndEncodedType) check(assert_type(ree.run_end_type, Int16Type | Int32Type | Int64Type), - lib.DataType) + pa.DataType) def test_data_type_hashable() -> None: - data_type = lib.int32() + data_type = pa.int32() check(assert_type(hash(data_type), int), int) def test_field() -> None: - field = lib.field("a", lib.int32()) - check(assert_type(field, lib.Field), lib.Field) + field = pa.field("a", pa.int32()) + check(assert_type(field, pa.Field), pa.Field) check(assert_type(field.name, str), str) check(assert_type(field.metadata, dict[bytes, bytes] | None), type(None)) - check(assert_type(field.type, lib.DataType), lib.DataType) + check(assert_type(field.type, pa.DataType), pa.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) + check(assert_type(field.with_nullable(False), pa.Field), pa.Field) + check(assert_type(pa.field("b", "string"), pa.Field), pa.Field) - field = lib.field("a", lib.int32(), metadata={"k": "v"}) + field = pa.field("a", pa.int32(), metadata={"k": "v"}) check(assert_type(field.metadata, dict[bytes, bytes] | None), dict) def test_schema() -> None: - schema = lib.schema([("a", lib.int32()), lib.field("b", lib.string())]) - check(assert_type(schema, lib.Schema), lib.Schema) + schema = pa.schema([("a", pa.int32()), pa.field("b", pa.string())]) + check(assert_type(schema, pa.Schema), pa.Schema) check(assert_type(schema.names, list[str]), list) - check(assert_type(schema.types, list[lib.DataType]), list) - check(assert_type(schema.field("a"), lib.Field), lib.Field) - check(assert_type(schema[0], lib.Field), lib.Field) + check(assert_type(schema.types, list[pa.DataType]), list) + check(assert_type(schema.field("a"), pa.Field), pa.Field) + check(assert_type(schema[0], pa.Field), pa.Field) def test_key_value_metadata() -> None: # The runtime yields generators instead of Mapping views (duplicated keys) - metadata = lib.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) + metadata = pa.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) check(assert_type(metadata.keys(), Iterator[bytes]), Iterator) check(assert_type(metadata.values(), Iterator[bytes]), Iterator) check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), @@ -140,74 +140,74 @@ def test_key_value_metadata() -> None: def test_type_predicates_narrow() -> None: - lst = _erase(lib.list_(lib.int32())) + lst = _erase(pa.list_(pa.int32())) assert types.is_list(lst) - check(assert_type(lst, lib.ListType), lib.ListType) + check(assert_type(lst, pa.ListType), pa.ListType) - ts = _erase(lib.timestamp("s")) + ts = _erase(pa.timestamp("s")) assert types.is_timestamp(ts) - check(assert_type(ts, lib.TimestampType), lib.TimestampType) + check(assert_type(ts, pa.TimestampType), pa.TimestampType) - dec = _erase(lib.decimal64(5)) + dec = _erase(pa.decimal64(5)) assert types.is_decimal(dec) check( assert_type( dec, - lib.Decimal32Type | lib.Decimal64Type | lib.Decimal128Type - | lib.Decimal256Type, + pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type + | pa.Decimal256Type, ), - lib.Decimal64Type, + pa.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)) + dec128 = _erase(pa.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, + pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type + | pa.Decimal256Type, ), - lib.Decimal128Type, + pa.Decimal128Type, ) # Run-end encoded types are nested, so this doesn't narrow to Never. - ree = lib.run_end_encoded(lib.int32(), lib.string()) + ree = pa.run_end_encoded(pa.int32(), pa.string()) assert types.is_nested(ree) - check(assert_type(ree, lib.RunEndEncodedType), lib.RunEndEncodedType) + check(assert_type(ree, pa.RunEndEncodedType), pa.RunEndEncodedType) def test_import_from_c_capsule() -> None: - capsule = lib.int32().__arrow_c_schema__() + capsule = pa.int32().__arrow_c_schema__() check(assert_type(capsule, CapsuleType), CapsuleType) # The import methods are static and return whatever type the C schema # holds, not the class they are called on. - imported = lib.ListType._import_from_c_capsule(capsule) - check(assert_type(imported, lib.DataType), lib.DataType) - assert not isinstance(imported, lib.ListType) + imported = pa.ListType._import_from_c_capsule(capsule) + check(assert_type(imported, pa.DataType), pa.DataType) + assert not isinstance(imported, pa.ListType) - field = lib.field("a", lib.int32()) + field = pa.field("a", pa.int32()) field_capsule = field.__arrow_c_schema__() check(assert_type(field_capsule, CapsuleType), CapsuleType) - imported_field = lib.Field._import_from_c_capsule(field_capsule) - check(assert_type(imported_field, lib.Field), lib.Field) + imported_field = pa.Field._import_from_c_capsule(field_capsule) + check(assert_type(imported_field, pa.Field), pa.Field) - schema = lib.schema([field]) + schema = pa.schema([field]) schema_capsule = schema.__arrow_c_schema__() check(assert_type(schema_capsule, CapsuleType), CapsuleType) - imported_schema = lib.Schema._import_from_c_capsule(schema_capsule) - check(assert_type(imported_schema, lib.Schema), lib.Schema) + imported_schema = pa.Schema._import_from_c_capsule(schema_capsule) + check(assert_type(imported_schema, pa.Schema), pa.Schema) def test_invalid_usage() -> None: if TYPE_CHECKING_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] - _: Hashable = lib.KeyValueMetadata({b"a": b"1"}) # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 + pa.field("a") # type: ignore[call-overload] + pa.list_(1) # type: ignore[call-overload] + pa.decimal128("10") # type: ignore[arg-type] + _: Hashable = pa.KeyValueMetadata({b"a": b"1"}) # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs/__init__.pyi b/python/pyarrow-stubs/__init__.pyi index 4b2a8da206c2..68e7c457e65f 100644 --- a/python/pyarrow-stubs/__init__.pyi +++ b/python/pyarrow-stubs/__init__.pyi @@ -17,6 +17,212 @@ from typing import Any +# The names that pyarrow/__init__.py imports from pyarrow.lib and that have +# stubs already. Type checkers only fall back to __getattr__ below for the +# names that aren't defined here. +from pyarrow.lib import ( + Array as Array, + ArrowCancelled as ArrowCancelled, + ArrowCapacityError as ArrowCapacityError, + ArrowException as ArrowException, + ArrowIOError as ArrowIOError, + ArrowIndexError as ArrowIndexError, + ArrowInvalid as ArrowInvalid, + ArrowKeyError as ArrowKeyError, + ArrowMemoryError as ArrowMemoryError, + ArrowNotImplementedError as ArrowNotImplementedError, + ArrowSerializationError as ArrowSerializationError, + ArrowTypeError as ArrowTypeError, + BaseExtensionType as BaseExtensionType, + BinaryArray as BinaryArray, + BinaryViewArray as BinaryViewArray, + Bool8Array as Bool8Array, + Bool8Type as Bool8Type, + BooleanArray as BooleanArray, + BuildInfo as BuildInfo, + ChunkedArray as ChunkedArray, + CppBuildInfo as CppBuildInfo, + DataType as DataType, + Date32Array as Date32Array, + Date64Array as Date64Array, + Decimal128Array as Decimal128Array, + Decimal128Type as Decimal128Type, + Decimal256Array as Decimal256Array, + Decimal256Type as Decimal256Type, + Decimal32Array as Decimal32Array, + Decimal32Type as Decimal32Type, + Decimal64Array as Decimal64Array, + Decimal64Type as Decimal64Type, + DenseUnionType as DenseUnionType, + Device as Device, + DeviceAllocationType as DeviceAllocationType, + DictionaryArray as DictionaryArray, + DictionaryMemo as DictionaryMemo, + DictionaryType as DictionaryType, + DoubleArray as DoubleArray, + DurationArray as DurationArray, + DurationType as DurationType, + ExtensionArray as ExtensionArray, + ExtensionType as ExtensionType, + Field as Field, + FixedShapeTensorArray as FixedShapeTensorArray, + FixedShapeTensorType as FixedShapeTensorType, + FixedSizeBinaryArray as FixedSizeBinaryArray, + FixedSizeBinaryType as FixedSizeBinaryType, + FixedSizeListArray as FixedSizeListArray, + FixedSizeListType as FixedSizeListType, + FloatArray as FloatArray, + FloatingPointArray as FloatingPointArray, + HalfFloatArray as HalfFloatArray, + Int16Array as Int16Array, + Int32Array as Int32Array, + Int64Array as Int64Array, + Int8Array as Int8Array, + IntegerArray as IntegerArray, + JsonArray as JsonArray, + JsonType as JsonType, + KeyValueMetadata as KeyValueMetadata, + LargeBinaryArray as LargeBinaryArray, + LargeListArray as LargeListArray, + LargeListType as LargeListType, + LargeListViewArray as LargeListViewArray, + LargeListViewType as LargeListViewType, + LargeStringArray as LargeStringArray, + ListArray as ListArray, + ListType as ListType, + ListViewArray as ListViewArray, + ListViewType as ListViewType, + LoggingMemoryPool as LoggingMemoryPool, + MapArray as MapArray, + MapType as MapType, + MemoryManager as MemoryManager, + MemoryPool as MemoryPool, + MonthDayNano as MonthDayNano, + MonthDayNanoIntervalArray as MonthDayNanoIntervalArray, + NullArray as NullArray, + NumericArray as NumericArray, + OpaqueArray as OpaqueArray, + OpaqueType as OpaqueType, + ProxyMemoryPool as ProxyMemoryPool, + RecordBatch as RecordBatch, + RunEndEncodedArray as RunEndEncodedArray, + RunEndEncodedType as RunEndEncodedType, + RuntimeInfo as RuntimeInfo, + Schema as Schema, + SparseCOOTensor as SparseCOOTensor, + SparseCSCMatrix as SparseCSCMatrix, + SparseCSFTensor as SparseCSFTensor, + SparseCSRMatrix as SparseCSRMatrix, + SparseUnionType as SparseUnionType, + StringArray as StringArray, + StringViewArray as StringViewArray, + StructArray as StructArray, + StructType as StructType, + Table as Table, + TableGroupBy as TableGroupBy, + Tensor as Tensor, + Time32Array as Time32Array, + Time32Type as Time32Type, + Time64Array as Time64Array, + Time64Type as Time64Type, + TimestampArray as TimestampArray, + TimestampType as TimestampType, + UInt16Array as UInt16Array, + UInt32Array as UInt32Array, + UInt64Array as UInt64Array, + UInt8Array as UInt8Array, + UnionArray as UnionArray, + UnionType as UnionType, + UnknownExtensionType as UnknownExtensionType, + UuidArray as UuidArray, + UuidType as UuidType, + VersionInfo as VersionInfo, + array as array, + binary as binary, + binary_view as binary_view, + bool8 as bool8, + bool_ as bool_, + build_info as build_info, + chunked_array as chunked_array, + concat_arrays as concat_arrays, + concat_batches as concat_batches, + concat_tables as concat_tables, + cpp_build_info as cpp_build_info, + cpp_version as cpp_version, + cpp_version_info as cpp_version_info, + date32 as date32, + date64 as date64, + decimal128 as decimal128, + decimal256 as decimal256, + decimal32 as decimal32, + decimal64 as decimal64, + default_cpu_memory_manager as default_cpu_memory_manager, + default_memory_pool as default_memory_pool, + dense_union as dense_union, + dictionary as dictionary, + duration as duration, + enable_signal_handlers as enable_signal_handlers, + field as field, + fixed_shape_tensor as fixed_shape_tensor, + float16 as float16, + float32 as float32, + float64 as float64, + from_numpy_dtype as from_numpy_dtype, + infer_type as infer_type, + int16 as int16, + int32 as int32, + int64 as int64, + int8 as int8, + jemalloc_memory_pool as jemalloc_memory_pool, + jemalloc_set_decay_ms as jemalloc_set_decay_ms, + json_ as json_, + large_binary as large_binary, + large_list as large_list, + large_list_view as large_list_view, + large_string as large_string, + large_utf8 as large_utf8, + list_ as list_, + list_view as list_view, + log_memory_allocations as log_memory_allocations, + logging_memory_pool as logging_memory_pool, + map_ as map_, + mimalloc_memory_pool as mimalloc_memory_pool, + month_day_nano_interval as month_day_nano_interval, + null as null, + nulls as nulls, + opaque as opaque, + proxy_memory_pool as proxy_memory_pool, + record_batch as record_batch, + register_extension_type as register_extension_type, + repeat as repeat, + run_end_encoded as run_end_encoded, + runtime_info as runtime_info, + schema as schema, + set_memory_pool as set_memory_pool, + set_timezone_db_path as set_timezone_db_path, + sparse_union as sparse_union, + string as string, + string_view as string_view, + struct as struct, + supported_memory_backends as supported_memory_backends, + system_memory_pool as system_memory_pool, + table as table, + time32 as time32, + time64 as time64, + timestamp as timestamp, + total_allocated_bytes as total_allocated_bytes, + type_for_alias as type_for_alias, + uint16 as uint16, + uint32 as uint32, + uint64 as uint64, + uint8 as uint8, + unify_schemas as unify_schemas, + union as union, + unregister_extension_type as unregister_extension_type, + utf8 as utf8, + uuid as uuid, +) + # TODO(GH-48970): remove __getattr__ before release as this # will annotate non-existing attributes as Any. # https://github.com/apache/arrow/issues/48970 From 046f15f6eefdca728b07db4e807a5bc3330941bb Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 01:05:47 +0100 Subject: [PATCH 64/93] Run the ruff hooks for the type stubs on commit ruff is fast and installs in its own environment, so lint and format the stubs on commit too. Only the type checkers stay in the manual stage. --- .pre-commit-config.yaml | 8 ++------ docs/source/developers/python/development.rst | 4 ++-- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bb7ceb65df98..a5c4b871abab 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -145,8 +145,8 @@ repos: ( ?^python/pyarrow/vendored/| ) - # Linting, formatting and type checking of the type stubs. These hooks don't - # run on commit, only on request (the Lint CI job runs them too): + # Linting, formatting and type checking of the type stubs. The type checkers + # don't run on commit, only on request (the Lint CI job runs them too): # pre-commit run --hook-stage manual python-stubs --all-files # Keep the versions of the type checkers and stub packages in sync with # python/requirements-stubs.txt. @@ -156,8 +156,6 @@ repos: - id: ruff-check alias: python-stubs name: Python Stubs Lint (ruff) - stages: - - manual args: - "--fix" files: >- @@ -167,8 +165,6 @@ repos: - id: ruff-format alias: python-stubs name: Python Stubs Format (ruff) - stages: - - manual files: >- ^python/pyarrow-stubs/ types: diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index e792db89889c..17b55cefc991 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -124,8 +124,8 @@ also lints and formats the stubs with ruff: $ pre-commit run --hook-stage manual python-stubs --all-files -These hooks only run on request, not on commit. The Lint CI job runs them as -well. CI also runs the type checkers, the type stub tests and stubtest (see +The type checkers only run on request, while ruff also runs on commit. The +Lint CI job runs all of them as well. CI also runs the type checkers, the type stub tests and stubtest (see below) against the PyArrow it builds, with ``ci/scripts/python_test_type_annotations.sh``. From 5caf254858423d79e4b0e7ab5d77d017890dc506 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 01:28:26 +0100 Subject: [PATCH 65/93] Install the stub type checkers with uvx from requirements-stubs.txt Run mypy, pyright and ty in the stub pre-commit hooks with uvx, in the environment of python/requirements-stubs.txt, so that the versions are only pinned there. Mark the stubs as partial, so that type checkers read the modules without stubs yet from the source tree, and the hooks don't need to install PyArrow. --- .pre-commit-config.yaml | 41 ++++++++----------- docs/source/developers/python/development.rst | 9 ++-- python/pyarrow-stubs/lib.pyi | 5 +-- python/pyarrow-stubs/py.typed | 20 +++++++++ python/requirements-stubs.txt | 5 ++- 5 files changed, 47 insertions(+), 33 deletions(-) create mode 100644 python/pyarrow-stubs/py.typed diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a5c4b871abab..59facf7a2fa9 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -148,8 +148,10 @@ repos: # Linting, formatting and type checking of the type stubs. The type checkers # don't run on commit, only on request (the Lint CI job runs them too): # pre-commit run --hook-stage manual python-stubs --all-files - # Keep the versions of the type checkers and stub packages in sync with - # python/requirements-stubs.txt. + # The type checkers run with uvx, in the environment of + # python/requirements-stubs.txt. They don't need PyArrow installed: the stubs + # are marked partial, so the modules that aren't annotated yet, e.g. + # pyarrow.compute, are read from the source tree. - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.17.2 hooks: @@ -177,47 +179,36 @@ repos: stages: - manual language: python - entry: bash -c "cd python && exec mypy" + entry: bash -c "cd python && exec uvx --with-requirements requirements-stubs.txt mypy" pass_filenames: false files: >- - ^python/(pyarrow-stubs(-tests)?/|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.1 - - pyright==1.1.414 - - pytest==9.0.2 - - scipy-stubs==1.17.1.5 - - ty==0.0.99 - - types-cffi==2.1.0.20260827 - - types-psutil==7.2.2.20260906 - - types-python-dateutil==2.9.0.20261021 - - types-requests==2.33.0.20261006 + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: &uv + - uv==0.13.7 - id: pyright alias: python-stubs name: Python Stubs Type Check (pyright) stages: - manual language: python - entry: pyright -p python/pyproject.toml + entry: uvx --with-requirements python/requirements-stubs.txt pyright -p python/pyproject.toml pass_filenames: false files: >- - ^python/(pyarrow-stubs(-tests)?/|pyproject\.toml$) - additional_dependencies: *python-type-check-dependencies + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: *uv - id: ty alias: python-stubs name: Python Stubs Type Check (ty) stages: - manual language: python - entry: ty check --project python + # Unset the VIRTUAL_ENV of the hook, so that ty uses the environment + # of uvx. + entry: env -u VIRTUAL_ENV uvx --with-requirements python/requirements-stubs.txt ty check --project python pass_filenames: false files: >- - ^python/(pyarrow-stubs(-tests)?/|pyproject\.toml$) - additional_dependencies: *python-type-check-dependencies + ^python/(pyarrow-stubs(-tests)?/|(pyproject\.toml|requirements-stubs\.txt)$) + additional_dependencies: *uv - repo: https://github.com/MarcoGorelli/cython-lint rev: v0.21.1 hooks: diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index 17b55cefc991..2363b9dca81b 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -117,8 +117,8 @@ 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: +the type checkers and stub packages pinned in ``requirements-stubs.txt`` (with +uv). This also lints and formats the stubs with ruff: .. code-block:: @@ -212,7 +212,10 @@ process and will be included when users install PyArrow, enabling type checking 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`. +checkers that PyArrow supports type checking according to :pep:`561`. The +``py.typed`` file in ``pyarrow-stubs/`` marks the stubs as partial, so type +checkers use the PyArrow sources for the modules without stubs yet. Remove +``partial`` from it once the stubs are complete. Testing Type Stubs ------------------ diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi index 9ec39728d811..19e6918bba5e 100644 --- a/python/pyarrow-stubs/lib.pyi +++ b/python/pyarrow-stubs/lib.pyi @@ -19,7 +19,7 @@ # have been annotated so far; everything else falls back to Any via # __getattr__ until it gets annotated in subsequent PRs. -from typing import Any, final +from typing import Any, Self from pyarrow.__lib.types import * # noqa: F401,F403 from pyarrow.__lib.array import * # noqa: F401,F403 @@ -33,7 +33,6 @@ from pyarrow.__lib.scalar import * # noqa: F401,F403 from pyarrow.__lib.table import * # noqa: F401,F403 from pyarrow.__lib.tensor import * # noqa: F401,F403 -@final class MonthDayNano(tuple): months: int days: int @@ -42,6 +41,6 @@ class MonthDayNano(tuple): def __new__( cls, sequence: tuple[int, int, int] | list[int] = ..., - ) -> MonthDayNano: ... + ) -> Self: ... def __getattr__(name: str) -> Any: ... diff --git a/python/pyarrow-stubs/py.typed b/python/pyarrow-stubs/py.typed new file mode 100644 index 000000000000..e102ebf5076c --- /dev/null +++ b/python/pyarrow-stubs/py.typed @@ -0,0 +1,20 @@ +# 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. + +# The stubs are incomplete: type checkers use the PyArrow sources for the +# modules without stubs yet. Remove this once the stubs are complete. +partial diff --git a/python/requirements-stubs.txt b/python/requirements-stubs.txt index 37d808571298..25aefadfebea 100644 --- a/python/requirements-stubs.txt +++ b/python/requirements-stubs.txt @@ -1,7 +1,8 @@ # Type checkers and the type stubs of PyArrow's dependencies, pinned so that # everyone gets the same results; see the "Type Checking" section of -# docs/source/developers/python/development.rst. Keep the versions in sync -# with the python-stubs hooks in .pre-commit-config.yaml. +# docs/source/developers/python/development.rst. The python-stubs hooks in +# .pre-commit-config.yaml and ci/scripts/python_test_type_annotations.sh +# install them from this file. fsspec==2026.9.0 mypy==2.4.0 pandas-stubs==3.0.5.260914 From c093920bfe8d6d978fa90bbe8e29b679faac8f8d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 01:38:35 +0100 Subject: [PATCH 66/93] Mark MonthDayNano as final in the type stubs MonthDayNano can't be subclassed at runtime, which stubtest reports unless the stub is marked with @final. --- python/pyarrow-stubs/lib.pyi | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi index 19e6918bba5e..d09176780766 100644 --- a/python/pyarrow-stubs/lib.pyi +++ b/python/pyarrow-stubs/lib.pyi @@ -19,7 +19,7 @@ # have been annotated so far; everything else falls back to Any via # __getattr__ until it gets annotated in subsequent PRs. -from typing import Any, Self +from typing import Any, Self, final from pyarrow.__lib.types import * # noqa: F401,F403 from pyarrow.__lib.array import * # noqa: F401,F403 @@ -33,6 +33,7 @@ from pyarrow.__lib.scalar import * # noqa: F401,F403 from pyarrow.__lib.table import * # noqa: F401,F403 from pyarrow.__lib.tensor import * # noqa: F401,F403 +@final class MonthDayNano(tuple): months: int days: int From cce46489de9bbe2e7d01d4c9f91afa41ed977f91 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 01:45:11 +0100 Subject: [PATCH 67/93] Remove the partial marker so that pyright checks the type stubs pyright merges a partial stub package into the package it annotates, and then no longer finds the files in pyarrow-stubs, so it only checked the type stub tests. Remove the marker again, and disable reportMissingImports for the stubs instead, so that pyright doesn't need PyArrow installed to resolve the modules that aren't annotated yet. --- .pre-commit-config.yaml | 6 +++--- docs/source/developers/python/development.rst | 5 +---- python/pyarrow-stubs/py.typed | 20 ------------------- python/pyproject.toml | 5 ++++- 4 files changed, 8 insertions(+), 28 deletions(-) delete mode 100644 python/pyarrow-stubs/py.typed diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 59facf7a2fa9..f70cfbf73923 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -149,9 +149,9 @@ repos: # don't run on commit, only on request (the Lint CI job runs them too): # pre-commit run --hook-stage manual python-stubs --all-files # The type checkers run with uvx, in the environment of - # python/requirements-stubs.txt. They don't need PyArrow installed: the stubs - # are marked partial, so the modules that aren't annotated yet, e.g. - # pyarrow.compute, are read from the source tree. + # python/requirements-stubs.txt. They don't need PyArrow installed: the + # modules that aren't annotated yet, e.g. pyarrow.compute, are treated as Any + # (see python/pyproject.toml). - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.17.2 hooks: diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index 2363b9dca81b..c5bdaf6e51cf 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -212,10 +212,7 @@ process and will be included when users install PyArrow, enabling type checking 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`. The -``py.typed`` file in ``pyarrow-stubs/`` marks the stubs as partial, so type -checkers use the PyArrow sources for the modules without stubs yet. Remove -``partial`` from it once the stubs are complete. +checkers that PyArrow supports type checking according to :pep:`561`. Testing Type Stubs ------------------ diff --git a/python/pyarrow-stubs/py.typed b/python/pyarrow-stubs/py.typed deleted file mode 100644 index e102ebf5076c..000000000000 --- a/python/pyarrow-stubs/py.typed +++ /dev/null @@ -1,20 +0,0 @@ -# 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. - -# The stubs are incomplete: type checkers use the PyArrow sources for the -# modules without stubs yet. Remove this once the stubs are complete. -partial diff --git a/python/pyproject.toml b/python/pyproject.toml index 2e8699b60aa9..04e58b7deddc 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -149,7 +149,10 @@ typeCheckingMode = "basic" # usage is rejected, so an ignore that is no longer needed is an error there. executionEnvironments = [ { root = "pyarrow-stubs-tests", extraPaths = ["."], reportUnnecessaryTypeIgnoreComment = "error" }, - { root = "." }, + # The stubs import modules that will be annotated in later PRs (GH-32609), + # which only resolve when PyArrow is installed, and the optional sparse + # dependency. mypy reports the other imports that can't be resolved. + { root = ".", reportMissingImports = false }, ] # In editable install, pyright finds the stubs but compiled modules (e.g. pyarrow.lib) reportMissingModuleSource = false From 61c7c26739fe254de15dceb0fdc5357f1d661f01 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 01:45:14 +0100 Subject: [PATCH 68/93] Remove the import ignores from the type stubs numpy and the stubs of pandas and scipy are installed wherever the type checkers run, and the type checkers are configured to accept the missing sparse, so the type: ignore comments on these imports aren't needed. Remove the mypy and ty overrides that were only there for these comments too. --- python/pyarrow-stubs/__lib/array.pyi | 4 ++-- python/pyarrow-stubs/__lib/table.pyi | 6 +++--- python/pyarrow-stubs/__lib/tensor.pyi | 2 +- python/pyarrow-stubs/__lib/types.pyi | 4 ++-- python/pyproject.toml | 13 ------------- 5 files changed, 8 insertions(+), 21 deletions(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 6f12a4325f2e..39a792131b86 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -19,8 +19,8 @@ from collections.abc import Collection, Iterable, Iterator, Sequence from typing import Any, ClassVar, Self, SupportsIndex -import numpy as np # type: ignore[import-not-found] -import pandas as pd # type: ignore[import-not-found,import-untyped] +import numpy as np +import pandas as pd from pyarrow._compute import CastOptions from pyarrow._stubs_typing import ( diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi index ddf7e9203066..7367d4d849db 100644 --- a/python/pyarrow-stubs/__lib/table.pyi +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -22,10 +22,10 @@ from typing import Any, ClassVar, Literal from typing_extensions import CapsuleType, deprecated import builtins -import numpy as np # type: ignore[import-not-found] -import pandas as pd # type: ignore[import-not-found,import-untyped] +import numpy as np +import pandas as pd -from numpy.typing import NDArray # type: ignore[import-not-found] +from numpy.typing import NDArray from pyarrow._compute import ( CastOptions, CountOptions, diff --git a/python/pyarrow-stubs/__lib/tensor.pyi b/python/pyarrow-stubs/__lib/tensor.pyi index 6d1e7cc05d41..5f2a1a0869fe 100644 --- a/python/pyarrow-stubs/__lib/tensor.pyi +++ b/python/pyarrow-stubs/__lib/tensor.pyi @@ -23,7 +23,7 @@ import numpy as np from pyarrow.lib import _Weakrefable from .types import DataType from scipy.sparse import coo_matrix, csc_matrix, csr_matrix -from sparse import COO # type: ignore[import-untyped, import-not-found] +from sparse import COO class Tensor(_Weakrefable): @classmethod diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index beb52cfeba0d..155b1ab21c56 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -18,8 +18,8 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence from typing import Any, ClassVar, Literal, TypeAlias, overload -import numpy as np # type: ignore[import-not-found] -import pandas as pd # type: ignore[import-not-found,import-untyped] +import numpy as np +import pandas as pd from typing_extensions import CapsuleType, Self, deprecated from pyarrow._stubs_typing import ( diff --git a/python/pyproject.toml b/python/pyproject.toml index 04e58b7deddc..df33826eb830 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -122,12 +122,6 @@ files = ["pyarrow-stubs", "pyarrow-stubs-tests"] # usage is rejected, so an ignore that is no longer needed is an error there. warn_unused_ignores = true -# The stubs ignore errors that depend on the installed packages, e.g. the -# imports of optional dependencies. -[[tool.mypy.overrides]] -module = ["pyarrow.*"] -warn_unused_ignores = false - # Modules that will be annotated in later PRs (GH-32609). Treat them as Any # rather than type checking their sources. [[tool.mypy.overrides]] @@ -171,13 +165,6 @@ replace-imports-with-any = [ [tool.ty.src] include = ["pyarrow-stubs", "pyarrow-stubs-tests"] -# numpy, pandas and scipy are optional dependencies. A `ty: ignore` comment on -# their imports would be reported as unused whenever they are installed. -[[tool.ty.overrides]] -include = ["pyarrow-stubs/**"] -[tool.ty.overrides.rules] -import-not-found = "ignore" - # ty doesn't honor `# type: ignore[]` comments, so disable the # rules that those comments suppress for mypy. [[tool.ty.overrides]] From 6ed844e1ff15605754736b85bff7ac36c05c4ec9 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 09:30:21 +0100 Subject: [PATCH 69/93] Add tensor.pyi updates and tensor test file --- python/pyarrow-stubs-tests/test_tensor.py | 108 ++++++++++++++ python/pyarrow-stubs/__lib/tensor.pyi | 164 +++++++++++++--------- python/pyarrow/tensor.pxi | 4 +- 3 files changed, 205 insertions(+), 71 deletions(-) create mode 100644 python/pyarrow-stubs-tests/test_tensor.py diff --git a/python/pyarrow-stubs-tests/test_tensor.py b/python/pyarrow-stubs-tests/test_tensor.py new file mode 100644 index 000000000000..1d0adae97d06 --- /dev/null +++ b/python/pyarrow-stubs-tests/test_tensor.py @@ -0,0 +1,108 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# mypy: warn-unused-ignores + +from collections.abc import Hashable +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from typing_extensions import CapsuleType + +import pyarrow as pa + +from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + +if TYPE_CHECKING: + from scipy.sparse import coo_array, csc_array, csr_array # noqa: F401 + +DENSE = np.array([[1, 0, 2], [0, 3, 0]]) + + +def test_tensor() -> None: + tensor = pa.Tensor.from_numpy(DENSE, dim_names=("x", "y")) + check(assert_type(tensor, pa.Tensor), pa.Tensor) + check(assert_type(tensor.dim_names, list[str]), list) + + from_dlpack = pa.Tensor.from_dlpack(DENSE) + check(assert_type(from_dlpack, pa.Tensor), pa.Tensor) + check(assert_type(tensor.__dlpack__(max_version=(1, 0)), CapsuleType), + CapsuleType) + check(assert_type(tensor.__dlpack_device__(), tuple[int, int]), tuple) + # __dlpack_device__() output must be accepted back as `device` + round_trip = pa.Tensor.from_dlpack( + tensor, device=tensor.__dlpack_device__()) + check(assert_type(round_trip, pa.Tensor), pa.Tensor) + + +def test_sparse_from_numpy() -> None: + dim_names = ("x", "y") + coo = pa.SparseCOOTensor.from_numpy( + np.array([[1, 2, 3]]).T, np.array([[0, 0], [0, 2], [1, 1]]), + (2, 3), dim_names) + check(assert_type(coo, pa.SparseCOOTensor), pa.SparseCOOTensor) + check(assert_type(coo.dim_names, tuple[str, ...]), tuple) + + coo = pa.SparseCOOTensor.from_numpy( + np.array([[1, 2, 3]]).T, np.array([[0, 0], [0, 2], [1, 1]]), + [2, 3], dim_names) + check(assert_type(coo, pa.SparseCOOTensor), pa.SparseCOOTensor) + + csr = pa.SparseCSRMatrix.from_numpy( + np.array([[1, 2, 3]]).T, np.array([0, 2, 3]), np.array([0, 2, 1]), + (2, 3), dim_names) + check(assert_type(csr, pa.SparseCSRMatrix), pa.SparseCSRMatrix) + check(assert_type(csr.dim_names, tuple[str, ...]), tuple) + + csc = pa.SparseCSCMatrix.from_numpy( + np.array([[1, 3, 2]]).T, np.array([0, 1, 2, 3]), + np.array([0, 1, 0]), (2, 3), dim_names) + check(assert_type(csc, pa.SparseCSCMatrix), pa.SparseCSCMatrix) + check(assert_type(csc.dim_names, tuple[str, ...]), tuple) + + csf = pa.SparseCSFTensor.from_numpy( + np.array([[1, 2, 3]]).T, [np.array([0, 2, 3])], + [np.array([0, 1]), np.array([0, 2, 1])], (2, 3), + axis_order=(0, 1), dim_names=dim_names) + check(assert_type(csf, pa.SparseCSFTensor), pa.SparseCSFTensor) + check(assert_type(csf.dim_names, tuple[str, ...]), tuple) + + csf = pa.SparseCSFTensor.from_numpy( + np.array([[1, 2, 3]]).T, (np.array([0, 2, 3]),), + (np.array([0, 1]), np.array([0, 2, 1])), (2, 3)) + check(assert_type(csf, pa.SparseCSFTensor), pa.SparseCSFTensor) + + +def test_sparse_scipy() -> None: + sparse = pytest.importorskip("scipy.sparse") + + coo = pa.SparseCOOTensor.from_scipy(sparse.coo_array(DENSE)) + check(assert_type(coo.to_scipy(), "coo_array"), sparse.coo_array) + + csr = pa.SparseCSRMatrix.from_scipy(sparse.csr_array(DENSE)) + check(assert_type(csr.to_scipy(), "csr_array"), sparse.csr_array) + + csc = pa.SparseCSCMatrix.from_scipy(sparse.csc_array(DENSE)) + check(assert_type(csc.to_scipy(), "csc_array"), sparse.csc_array) + + +def test_invalid_usage() -> None: + tensor = pa.Tensor.from_numpy(DENSE) + if TYPE_CHECKING_INVALID_USAGE: + pa.Tensor.from_numpy(DENSE, dim_names=[1, 2]) # type: ignore[list-item] + _: Hashable = tensor # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs/__lib/tensor.pyi b/python/pyarrow-stubs/__lib/tensor.pyi index 5f2a1a0869fe..6cde448b34a1 100644 --- a/python/pyarrow-stubs/__lib/tensor.pyi +++ b/python/pyarrow-stubs/__lib/tensor.pyi @@ -15,21 +15,29 @@ # specific language governing permissions and limitations # under the License. -from typing import Self +from typing import ClassVar +from typing_extensions import CapsuleType -from collections.abc import Sequence +from collections.abc import Iterable, Sequence import numpy as np from pyarrow.lib import _Weakrefable from .types import DataType -from scipy.sparse import coo_matrix, csc_matrix, csr_matrix +from scipy.sparse import ( + coo_array, + coo_matrix, + csc_array, + csc_matrix, + csr_array, + csr_matrix, +) from sparse import COO class Tensor(_Weakrefable): - @classmethod + @staticmethod def from_numpy( - cls, obj: np.ndarray, dim_names: Sequence[str] | None = None - ) -> Self: ... + obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> Tensor: ... def to_numpy(self) -> np.ndarray: ... def equals(self, other: Tensor) -> bool: ... def dim_name(self, i: int) -> str: ... @@ -47,37 +55,54 @@ class Tensor(_Weakrefable): def shape(self) -> tuple[int, ...]: ... @property def strides(self) -> tuple[int, ...]: ... + @staticmethod + def from_dlpack( + x: object, + /, + *, + device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> Tensor: ... + def __dlpack__( + self, + *, + stream: int | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> CapsuleType: ... + def __dlpack_device__(self) -> tuple[int, int]: ... @property def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] class SparseCOOTensor(_Weakrefable): @classmethod def from_dense_numpy( - cls, obj: np.ndarray, dim_names: list[str] | None = None - ) -> Self: ... - @classmethod + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod def from_numpy( - cls, data: np.ndarray, coords: np.ndarray, - shape: Sequence[int], - dim_names: Sequence[str] | None = None, - ) -> Self: ... - @classmethod + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCOOTensor: ... + @staticmethod def from_scipy( - cls, obj: coo_matrix, dim_names: Sequence[str] | None = None - ) -> Self: ... - @classmethod + obj: coo_array | coo_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod def from_pydata_sparse( - cls, obj: COO, dim_names: Sequence[str] | None = None - ) -> Self: ... - @classmethod - def from_tensor(cls, obj: Tensor) -> Self: ... + obj: COO, dim_names: Iterable[str] | None = None + ) -> SparseCOOTensor: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCOOTensor: ... def to_numpy(self) -> tuple[np.ndarray, np.ndarray]: ... - def to_scipy(self) -> coo_matrix: ... + def to_scipy(self) -> coo_array: ... def to_pydata_sparse(self) -> COO: ... def to_tensor(self) -> Tensor: ... - def equals(self, other: Self) -> bool: ... + def equals(self, other: SparseCOOTensor) -> bool: ... @property def is_mutable(self) -> bool: ... @property @@ -88,38 +113,38 @@ class SparseCOOTensor(_Weakrefable): def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @property - def dim_names(self) -> list[str]: ... + def dim_names(self) -> tuple[str, ...]: ... @property def non_zero_length(self) -> int: ... @property def has_canonical_format(self) -> bool: ... @property def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] class SparseCSRMatrix(_Weakrefable): @classmethod def from_dense_numpy( - cls, obj: np.ndarray, dim_names: list[str] | None = None - ) -> Self: ... - @classmethod + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSRMatrix: ... + @staticmethod def from_numpy( - cls, data: np.ndarray, indptr: np.ndarray, indices: np.ndarray, - shape: Sequence[int], - dim_names: Sequence[str] | None = None, - ) -> Self: ... - @classmethod + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCSRMatrix: ... + @staticmethod def from_scipy( - cls, obj: csr_matrix, dim_names: Sequence[str] | None = None - ) -> Self: ... - @classmethod - def from_tensor(cls, obj: Tensor) -> Self: ... + obj: csr_array | csr_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCSRMatrix: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSRMatrix: ... def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... - def to_scipy(self) -> csr_matrix: ... + def to_scipy(self) -> csr_array: ... def to_tensor(self) -> Tensor: ... - def equals(self, other: Self) -> bool: ... + def equals(self, other: SparseCSRMatrix) -> bool: ... @property def is_mutable(self) -> bool: ... @property @@ -130,36 +155,36 @@ class SparseCSRMatrix(_Weakrefable): def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @property - def dim_names(self) -> list[str]: ... + def dim_names(self) -> tuple[str, ...]: ... @property def non_zero_length(self) -> int: ... @property def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] class SparseCSCMatrix(_Weakrefable): @classmethod def from_dense_numpy( - cls, obj: np.ndarray, dim_names: list[str] | None = None - ) -> Self: ... - @classmethod + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSCMatrix: ... + @staticmethod def from_numpy( - cls, data: np.ndarray, indptr: np.ndarray, indices: np.ndarray, - shape: Sequence[int], - dim_names: Sequence[str] | None = None, - ) -> Self: ... - @classmethod + shape: Iterable[int], + dim_names: Iterable[str] | None = None, + ) -> SparseCSCMatrix: ... + @staticmethod def from_scipy( - cls, obj: csc_matrix, dim_names: Sequence[str] | None = None - ) -> Self: ... - @classmethod - def from_tensor(cls, obj: Tensor) -> Self: ... + obj: csc_array | csc_matrix, dim_names: Iterable[str] | None = None + ) -> SparseCSCMatrix: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSCMatrix: ... def to_numpy(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: ... - def to_scipy(self) -> csc_matrix: ... + def to_scipy(self) -> csc_array: ... def to_tensor(self) -> Tensor: ... - def equals(self, other: Self) -> bool: ... + def equals(self, other: SparseCSCMatrix) -> bool: ... @property def is_mutable(self) -> bool: ... @property @@ -170,32 +195,32 @@ class SparseCSCMatrix(_Weakrefable): def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @property - def dim_names(self) -> list[str]: ... + def dim_names(self) -> tuple[str, ...]: ... @property def non_zero_length(self) -> int: ... @property def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] class SparseCSFTensor(_Weakrefable): @classmethod def from_dense_numpy( - cls, obj: np.ndarray, dim_names: Sequence[str] | None = None - ) -> Self: ... - @classmethod + cls, obj: np.ndarray, dim_names: Iterable[str] | None = None + ) -> SparseCSFTensor: ... + @staticmethod 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]: ... + indptr: list[np.ndarray] | tuple[np.ndarray, ...], + indices: list[np.ndarray] | tuple[np.ndarray, ...], + shape: Sequence[int], + axis_order: Iterable[int] | None = None, + dim_names: Iterable[str] | None = None, + ) -> SparseCSFTensor: ... + @staticmethod + def from_tensor(obj: Tensor) -> SparseCSFTensor: ... + def to_numpy(self) -> tuple[np.ndarray, list[np.ndarray], list[np.ndarray]]: ... def to_tensor(self) -> Tensor: ... - def equals(self, other: Self) -> bool: ... + def equals(self, other: SparseCSFTensor) -> bool: ... @property def is_mutable(self) -> bool: ... @property @@ -206,11 +231,12 @@ class SparseCSFTensor(_Weakrefable): def shape(self) -> tuple[int, ...]: ... def dim_name(self, i: int) -> str: ... @property - def dim_names(self) -> list[str]: ... + def dim_names(self) -> tuple[str, ...]: ... @property def non_zero_length(self) -> int: ... @property def type(self) -> DataType: ... + __hash__: ClassVar[None] # type: ignore[assignment] __all__ = [ "Tensor", diff --git a/python/pyarrow/tensor.pxi b/python/pyarrow/tensor.pxi index 7f006f294175..5ddcb02543b4 100644 --- a/python/pyarrow/tensor.pxi +++ b/python/pyarrow/tensor.pxi @@ -1376,14 +1376,14 @@ shape: {self.shape}""" ---------- data : numpy.ndarray Data used to populate the sparse tensor. - indptr : numpy.ndarray + indptr : list or tuple of numpy.ndarray The sparsity structure. Each two consecutive dimensions in a tensor correspond to a buffer in indices. A pair of consecutive values at `indptr[dim][i]` `indptr[dim][i + 1]` signify a range of nodes in `indices[dim + 1]` who are children of `indices[dim][i]` node. - indices : numpy.ndarray + indices : list or tuple of numpy.ndarray Stores values of nodes. Each tensor dimension corresponds to a buffer in indptr. shape : tuple From 6a4e6145741c22a17fac0bb3226415cc6b034b76 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 10:29:00 +0100 Subject: [PATCH 70/93] Update stubtest allowlist --- python/pyarrow-stubs-tests/stubtest_allowlist.txt | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt index 7cb289d267a8..bfae14baa66a 100644 --- a/python/pyarrow-stubs-tests/stubtest_allowlist.txt +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -59,8 +59,6 @@ pyarrow\.__lib\.device\.DeviceAllocationType\.__format__ pyarrow\.__lib\.table\.ChunkedArray\.data pyarrow\.__lib\.table\.ChunkedArray\.format pyarrow\.__lib\.table\.concat_tables -pyarrow\.__lib\.tensor\.Sparse[A-Z]+(Tensor|Matrix)\.from_(numpy|pydata_sparse|scipy|tensor) -pyarrow\.__lib\.tensor\.Tensor\.from_numpy pyarrow\.__lib\.types\.Schema\.__getitem__ pyarrow\.__lib\.types\.Schema\.add_metadata pyarrow\.__lib\.types\.Schema\.field_by_name From da3465afe0bb16265665051f7ced586f02f78483 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 11:00:23 +0100 Subject: [PATCH 71/93] Update _KeyValueMetadataInput --- .../stubtest_allowlist.txt | 1 - python/pyarrow-stubs-tests/test_types.py | 27 ++++++++++--------- python/pyarrow-stubs/__lib/types.pyi | 20 ++++++-------- 3 files changed, 22 insertions(+), 26 deletions(-) diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt index bfae14baa66a..2800b7f800c2 100644 --- a/python/pyarrow-stubs-tests/stubtest_allowlist.txt +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -60,6 +60,5 @@ pyarrow\.__lib\.table\.ChunkedArray\.data pyarrow\.__lib\.table\.ChunkedArray\.format pyarrow\.__lib\.table\.concat_tables pyarrow\.__lib\.types\.Schema\.__getitem__ -pyarrow\.__lib\.types\.Schema\.add_metadata pyarrow\.__lib\.types\.Schema\.field_by_name pyarrow\.types\.__all__ diff --git a/python/pyarrow-stubs-tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py index ab0219ef09e5..aef4f3477b79 100644 --- a/python/pyarrow-stubs-tests/test_types.py +++ b/python/pyarrow-stubs-tests/test_types.py @@ -73,13 +73,10 @@ def test_list_types() -> None: variable = pa.list_(pa.int32(), -1) check(assert_type(variable, pa.ListType), pa.ListType) fixed = pa.list_(pa.int32(), 2) - check(assert_type(fixed, pa.ListType | pa.FixedSizeListType), - pa.FixedSizeListType) + check(assert_type(fixed, pa.ListType | pa.FixedSizeListType), pa.FixedSizeListType) - check(assert_type(pa.large_list(pa.int32()), pa.LargeListType), - pa.LargeListType) - check(assert_type(pa.list_view(pa.int32()), pa.ListViewType), - pa.ListViewType) + check(assert_type(pa.large_list(pa.int32()), pa.LargeListType), pa.LargeListType) + check(assert_type(pa.list_view(pa.int32()), pa.ListViewType), pa.ListViewType) def test_dictionary_type() -> None: @@ -98,8 +95,7 @@ def test_map_and_run_end_encoded_types() -> None: ree = pa.run_end_encoded(pa.int16(), pa.string()) check(assert_type(ree, pa.RunEndEncodedType), pa.RunEndEncodedType) - check(assert_type(ree.run_end_type, Int16Type | Int32Type | Int64Type), - pa.DataType) + check(assert_type(ree.run_end_type, Int16Type | Int32Type | Int64Type), pa.DataType) def test_data_type_hashable() -> None: @@ -119,6 +115,11 @@ def test_field() -> None: field = pa.field("a", pa.int32(), metadata={"k": "v"}) check(assert_type(field.metadata, dict[bytes, bytes] | None), dict) + check(assert_type(field.with_metadata(None), pa.Field), pa.Field) + check( + assert_type(field.with_metadata(None).metadata, dict[bytes, bytes] | None), + dict # Empty dict + ) def test_schema() -> None: @@ -135,8 +136,7 @@ def test_key_value_metadata() -> None: metadata = pa.KeyValueMetadata([(b"a", b"1"), (b"a", b"2")]) check(assert_type(metadata.keys(), Iterator[bytes]), Iterator) check(assert_type(metadata.values(), Iterator[bytes]), Iterator) - check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), - Iterator) + check(assert_type(metadata.items(), Iterator[tuple[bytes, bytes]]), Iterator) def test_type_predicates_narrow() -> None: @@ -153,8 +153,7 @@ def test_type_predicates_narrow() -> None: check( assert_type( dec, - pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type - | pa.Decimal256Type, + pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type | pa.Decimal256Type, ), pa.Decimal64Type, ) @@ -167,7 +166,9 @@ def test_type_predicates_narrow() -> None: check( assert_type( dec128, - pa.Decimal32Type | pa.Decimal64Type | pa.Decimal128Type + pa.Decimal32Type + | pa.Decimal64Type + | pa.Decimal128Type | pa.Decimal256Type, ), pa.Decimal128Type, diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index 155b1ab21c56..8541a8fad514 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -287,7 +287,7 @@ _MetadataMapping: TypeAlias = ( | Mapping[str, bytes] ) _MetadataIterable: TypeAlias = Iterable[tuple[_StrOrBytes, _StrOrBytes]] -_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable | None +_KeyValueMetadataInput: TypeAlias = _MetadataMapping | _MetadataIterable _DataTypeAlias: TypeAlias = Literal[ "null", "bool", @@ -351,7 +351,7 @@ _DataTypeLike: TypeAlias = DataType | _DataTypeAliasInput class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): def __init__( self, - __arg0__: _KeyValueMetadataInput | KeyValueMetadata = None, + __arg0__: _KeyValueMetadataInput | KeyValueMetadata | None = None, **kwargs: str | bytes, ) -> None: ... __hash__: ClassVar[None] # type: ignore[assignment] @@ -381,7 +381,7 @@ class Field(_Weakrefable): def type(self) -> DataType: ... def with_metadata( self, - metadata: _KeyValueMetadataInput | KeyValueMetadata, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None, ) -> Field: ... def remove_metadata(self) -> Field: ... def with_type(self, new_type: DataType) -> Field: ... @@ -402,7 +402,7 @@ _FieldTuple: TypeAlias = ( _FieldName, _DataTypeLike, bool | None, - _KeyValueMetadataInput | KeyValueMetadata, + _KeyValueMetadataInput | KeyValueMetadata | None, ] ) _StructFieldsInput: TypeAlias = ( @@ -440,12 +440,8 @@ class Schema(_Weakrefable): 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 + self, metadata: _KeyValueMetadataInput | KeyValueMetadata | None ) -> Schema: ... def serialize(self, memory_pool: MemoryPool | None = None) -> Buffer: ... def remove_metadata(self) -> Schema: ... @@ -473,14 +469,14 @@ def field( name: SupportsArrowSchema, type: None = None, nullable: bool | None = None, - metadata: _KeyValueMetadataInput | KeyValueMetadata = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, ) -> Field: ... @overload def field( name: _FieldName, type: _DataTypeLike, nullable: bool | None = None, - metadata: _KeyValueMetadataInput | KeyValueMetadata = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, ) -> Field: ... def null() -> NullType: ... def bool_() -> BoolType: ... @@ -570,7 +566,7 @@ def opaque(storage_type: DataType, type_name: str, vendor_name: str) -> OpaqueTy def type_for_alias(name: str) -> DataType: ... def schema( fields: SupportsArrowSchema | _SchemaFieldsInput, - metadata: _KeyValueMetadataInput | KeyValueMetadata = None, + metadata: _KeyValueMetadataInput | KeyValueMetadata | None = None, ) -> Schema: ... def from_numpy_dtype(dtype: np.typing.DTypeLike) -> DataType: ... From 86fe42224a149e48b8e57edf072f0729b2dfa751 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 11:16:01 +0100 Subject: [PATCH 72/93] Update __contains__ and __getitem__ in KeyValueMetadata --- python/pyarrow-stubs/__lib/types.pyi | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index 8541a8fad514..7fc9c5d5fb73 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -357,8 +357,8 @@ class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]): __hash__: ClassVar[None] # type: ignore[assignment] def equals(self, other: KeyValueMetadata) -> bool: ... def __len__(self) -> int: ... - def __contains__(self, /, __key: object) -> bool: ... - def __getitem__(self, /, __key: str | bytes) -> bytes: ... + def __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: ... From 5ea9439c07f21313ffec89fadee611384e166cf9 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 12:02:42 +0100 Subject: [PATCH 73/93] test_array.py unused type-ignore --- python/pyarrow/tests/test_array.py | 22 +++++++++++----------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index b597d12bb97f..a2292150cf4d 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -77,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]) # type: ignore[reportCallIssue] + pa.Array([1, 2]) # pyright: ignore[reportCallIssue] def test_list_format(): @@ -787,7 +787,7 @@ def test_array_from_buffers(): with pytest.raises(TypeError): pa.Array.from_buffers( - pa.int16(), 3, ['', ''], offset=1 # type: ignore[reportArgumentType] + pa.int16(), 3, ['', ''], offset=1 # pyright: ignore[reportArgumentType] ) @@ -983,7 +983,7 @@ def test_struct_array_from_chunked(): with pytest.raises(TypeError, match="Expected Array"): pa.StructArray.from_arrays( - [chunked_arr], ["foo"]) # type: ignore[reportArgumentType] + [chunked_arr], ["foo"]) # pyright: ignore[reportArgumentType] @pytest.mark.parametrize("offset", (0, 1)) @@ -2875,7 +2875,7 @@ def test_interval_array_from_relativedelta(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, # type: ignore[arg-type] + pa.MonthDayNano([13, 8, (datetime.timedelta(seconds=1, microseconds=1, minutes=1, hours=1) // datetime.timedelta(microseconds=1)) * 1000])] @@ -2907,7 +2907,7 @@ def test_interval_array_from_tuple(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([1, 2, -3])] # type: ignore[arg-type] + pa.MonthDayNano([1, 2, -3])] expected = pa.array(expected_list) assert arr.equals(expected) assert arr.to_pylist() == expected_list @@ -2928,8 +2928,8 @@ def test_interval_array_from_dateoffset(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, 3661000001001]), # type: ignore[arg-type] - pa.MonthDayNano([0, 0, 0])] # type: ignore[arg-type] + pa.MonthDayNano([13, 8, 3661000001001]), + pa.MonthDayNano([0, 0, 0])] expected = pa.array(expected_list) assert arr.equals(expected) expected_from_pandas = [ @@ -3289,7 +3289,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])) # type: ignore[reportCallIssue] + repr(pa.Tensor([1])) # pyright: ignore[reportCallIssue] def test_invalid_tensor_construction(): @@ -3807,7 +3807,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]) # type: ignore[reportArgumentType] + mask=[1.0, 2.0, 3.0, 4.0]) # pyright: ignore[reportArgumentType] with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], @@ -4085,11 +4085,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]) # type: ignore[reportArgumentType] + pa.concat_arrays([None]) # pyright: ignore[reportArgumentType] arr = pa.chunked_array([[0, 1], [3, 4]]) with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays(arr) # type: ignore[reportArgumentType] + pa.concat_arrays(arr) # pyright: ignore[reportArgumentType] @pytest.mark.pandas From 42927bebff9e8b95a57fd742a27950643a6bb99f Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 12:15:54 +0100 Subject: [PATCH 74/93] Update mask arg in array() --- python/pyarrow-stubs-tests/test_array.py | 5 +++++ python/pyarrow-stubs/__lib/array.pyi | 3 ++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 08fdb5df8de3..1fc314f4ea10 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -67,6 +67,11 @@ def test_array() -> None: # Arrays of any type can diff() check(assert_type(arr.diff(other), str), str) + if TYPE_CHECKING_INVALID_USAGE: + pa.array([1, 2], mask=[True, None]) # type: ignore[list-item] + # TODO: when scalars get stubs + # pa.array([1, 2], mask=pa.chunked_array([[True, False]])) + def test_array_numpy() -> None: pytest.importorskip("numpy") diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 39a792131b86..1ba291a9a32c 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -22,6 +22,7 @@ from typing import Any, ClassVar, Self, SupportsIndex import numpy as np import pandas as pd +from numpy.typing import NDArray from pyarrow._compute import CastOptions from pyarrow._stubs_typing import ( ArrayLike, @@ -57,7 +58,7 @@ def array( | SupportsArrowDeviceArray | SupportsPyArrowArray, type: Any | None = None, - mask: Mask | pd.Series[bool] | None = None, + mask: BooleanArray | Iterable[bool] | NDArray[np.bool_] | None = None, size: int | None = None, from_pandas: bool | None = None, safe: bool = True, From 724d480466dae3c77ee1f1a92ede57c9e0f21d13 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 12:45:50 +0100 Subject: [PATCH 75/93] Update Mask alias --- python/pyarrow-stubs-tests/test_array.py | 7 +++++++ python/pyarrow-stubs/_stubs_typing.pyi | 6 ++---- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 1fc314f4ea10..4ba957a1356c 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -67,6 +67,9 @@ def test_array() -> None: # Arrays of any type can diff() check(assert_type(arr.diff(other), str), str) + bool_mask = cast(pa.BooleanArray, pa.array([True, False, True])) + check(assert_type(arr.filter(bool_mask), pa.Int64Array), pa.Int64Array) + if TYPE_CHECKING_INVALID_USAGE: pa.array([1, 2], mask=[True, None]) # type: ignore[list-item] # TODO: when scalars get stubs @@ -83,6 +86,10 @@ def test_array_numpy() -> None: offsets = np.array([0, 1, 2], dtype="int32") check(assert_type(pa.ListArray.from_arrays(offsets, [1, 2]), pa.ListArray), pa.ListArray) + + np_mask = np.array([True, False]) + check(assert_type(arr.filter(np_mask), pa.Int64Array), pa.Int64Array) + if TYPE_CHECKING_INVALID_USAGE: dtype = pa.dictionary(pa.int8(), pa.string()) strs = np.array(["a"]) diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi index d1b3c6840e2b..663fb129d254 100644 --- a/python/pyarrow-stubs/_stubs_typing.pyi +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -32,7 +32,7 @@ import numpy as np from numpy.typing import NDArray from pyarrow import lib -from pyarrow.lib import ChunkedArray +from pyarrow.lib import Array, ChunkedArray ArrayLike: TypeAlias = Any ScalarLike: TypeAlias = Any @@ -70,9 +70,7 @@ PyScalar: TypeAlias = ( ) NumpyScalar: TypeAlias = np.generic[Any] -Mask: TypeAlias = ( - Sequence[bool | None] | NDArray[np.bool_] | lib.BooleanArray | ChunkedArray -) +Mask: TypeAlias = Sequence[bool | None] | NDArray[np.bool_] | Array | ChunkedArray Indices: TypeAlias = ( Sequence[int | None] | NDArray[np.integer[Any]] | lib.IntegerArray | ChunkedArray ) From 133588549256abec90c993b49cc868b4dc42ec79 Mon Sep 17 00:00:00 2001 From: Alenka Frim Date: Tue, 10 Nov 2026 12:49:38 +0100 Subject: [PATCH 76/93] Update python/pyarrow-stubs/__lib/array.pyi Co-authored-by: Joris Van den Bossche --- python/pyarrow-stubs/__lib/array.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 1ba291a9a32c..becc4c17aae0 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -475,7 +475,7 @@ class DictionaryArray(Array): ) -> DictionaryArray: ... @staticmethod def from_arrays( - indices: Indices | Sequence[int | None], + indices: Indices, dictionary: Array | np.ndarray | pd.Series | list[Any], mask: np.ndarray | pd.Series | BooleanArray | None = None, ordered: bool = False, From bb913a8d9e806205920383d4054a81f6d926b8eb Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 13:09:12 +0100 Subject: [PATCH 77/93] pyarrow/tests/test_*.py remove/update ignores --- python/pyarrow/tests/test_convert_builtin.py | 10 +++--- python/pyarrow/tests/test_device.py | 4 +-- python/pyarrow/tests/test_schema.py | 8 ++--- python/pyarrow/tests/test_sparse_tensor.py | 2 +- python/pyarrow/tests/test_table.py | 35 ++++++++++---------- python/pyarrow/tests/test_tensor.py | 2 +- python/pyarrow/tests/test_types.py | 18 +++++----- 7 files changed, 40 insertions(+), 39 deletions(-) diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index f429172600a7..1552c03cb238 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -165,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()) # type: ignore[arg-type] + pa.array(ObjectWithOnlyGetitem()) # pyright: ignore[reportArgumentType] def _as_list(xs): @@ -2131,8 +2131,8 @@ def test_map_from_dicts(): assert arr.to_pylist() == expected # With omitted values - data[1] = None # type: ignore[call-overload] - expected[1] = None # type: ignore[call-overload] + data[1] = None # pyright: ignore[reportArgumentType, reportCallIssue] + expected[1] = None # pyright: ignore[reportArgumentType, reportCallIssue] arr = pa.array(expected, type=pa.map_(pa.binary(), pa.int32())) @@ -2540,9 +2540,9 @@ def test_array_from_pylist_offset_overflow(): pa.timestamp('us') ), ( - [pa.MonthDayNano([1, -1, -10100])], # type: ignore[call-arg, arg-type] + [pa.MonthDayNano([1, -1, -10100])], [pa.scalar( - pa.MonthDayNano([1, -1, -10100]) # type: ignore[call-arg, arg-type] + pa.MonthDayNano([1, -1, -10100]) )], pa.month_day_nano_interval() ), diff --git a/python/pyarrow/tests/test_device.py b/python/pyarrow/tests/test_device.py index 00f8bbf720da..67981b0a0822 100644 --- a/python/pyarrow/tests/test_device.py +++ b/python/pyarrow/tests/test_device.py @@ -67,7 +67,7 @@ def test_copy_to(): assert buffer.address != buffer_orig.address with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - arr.copy_to(mm.device.device_type) # type: ignore[arg-type] + arr.copy_to(mm.device.device_type) with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - batch.copy_to(mm.device.device_type) # type: ignore[arg-type] + batch.copy_to(mm.device.device_type) # pyright: ignore[reportArgumentType] diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index 0cd004a5a97e..b36a88a04511 100644 --- a/python/pyarrow/tests/test_schema.py +++ b/python/pyarrow/tests/test_schema.py @@ -253,7 +253,7 @@ def test_schema(): child 0, item: int8""" with pytest.raises(TypeError): - pa.schema([None]) # type: ignore[list-item] + pa.schema([None]) # pyright: ignore[reportArgumentType] def test_schema_weakref(): @@ -538,7 +538,7 @@ def test_schema_equals_invalid_type(): for val in [None, 'string', pa.array([1, 2])]: with pytest.raises(TypeError): - schema.equals(val) # type: ignore[invalid-argument-type] + schema.equals(val) def test_schema_equality_operators(): @@ -584,7 +584,7 @@ def test_schema_get_fields(): with pytest.raises(KeyError): schema.field('other') with pytest.raises(TypeError): - schema.field(0.0) # type: ignore[arg-type] + schema.field(0.0) # pyright: ignore[reportArgumentType] with pytest.raises(IndexError): schema.field(4) @@ -780,7 +780,7 @@ def test_schema_merge(): # raise proper error when passing a non-Schema value with pytest.raises(TypeError): - pa.unify_schemas([a, 1]) # type: ignore[list-item] + pa.unify_schemas([a, 1]) # pyright: ignore[reportArgumentType] def test_undecodable_metadata(): diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index 7cd5d853ab92..a224d1e4904f 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -483,7 +483,7 @@ def test_pydata_sparse_sparse_coo_tensor_roundtrip(dtype_str, arrow_type): shape = (4, 6) dim_names = ("x", "y") - sparse_array = sparse.COO( # type: ignore[reportOptionalMemberAccess] + sparse_array = sparse.COO( # pyright: ignore[reportOptionalMemberAccess] data=data, coords=coords, shape=shape) sparse_tensor = pa.SparseCOOTensor.from_pydata_sparse(sparse_array, dim_names=dim_names) diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index e28f00862035..2c2a6012537a 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -812,7 +812,7 @@ def test_recordbatch_get_field(): batch.field('d') with pytest.raises(TypeError): - batch.field(None) # type: ignore[arg-type] + batch.field(None) # pyright: ignore[reportArgumentType] with pytest.raises(IndexError): batch.field(4) @@ -833,7 +833,7 @@ def test_recordbatch_select_column(): batch.column('d') with pytest.raises(TypeError): - batch.column(None) # type: ignore[arg-type] + batch.column(None) # pyright: ignore[reportArgumentType] with pytest.raises(IndexError): batch.column(4) @@ -1227,7 +1227,7 @@ 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.float64, order="F") # type: ignore[no-matching-overload] + np.float64, order="F") expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1262,7 +1262,7 @@ 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") # type: ignore[no-matching-overload] + np.float32, order="F") expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1711,7 +1711,7 @@ def test_table_get_field(): table.field('d') with pytest.raises(TypeError): - table.field(None) # type: ignore[arg-type] + table.field(None) # pyright: ignore[reportArgumentType] with pytest.raises(IndexError): table.field(4) @@ -1732,7 +1732,7 @@ def test_table_select_column(): table.column('d') with pytest.raises(TypeError): - table.column(None) # type: ignore[arg-type] + table.column(None) # pyright: ignore[reportArgumentType] with pytest.raises(IndexError): table.column(4) @@ -2129,13 +2129,14 @@ 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") # type: ignore[arg-type] + pa.concat_tables( + [t, t], promote_options="invalid") # pyright: ignore[reportArgumentType] def test_concat_tables_none_table(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_tables([None]) # type: ignore[arg-type] + pa.concat_tables([None]) # pyright: ignore[reportArgumentType] @pytest.mark.pandas @@ -2267,7 +2268,7 @@ def test_concat_batches_different_schema(): def test_concat_batches_none_batches(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_batches([None]) # type: ignore[arg-type] + pa.concat_batches([None]) # pyright: ignore[reportArgumentType] @pytest.mark.parametrize( @@ -2418,7 +2419,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) # type: ignore[arg-type] + pa.Table.from_pydict(data, schema=schema) # pyright: ignore[reportArgumentType] # Cannot pass both schema and metadata / names with pytest.raises(ValueError): @@ -2523,7 +2524,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) # type: ignore[arg-type] + pa.Table.from_pydict(data, schema=schema) @pytest.mark.parametrize('data, klass', [ @@ -2540,7 +2541,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) # type: ignore[arg-type] + pa.Table.from_pydict(data, schema=schema) # data has columns not present in schema -> ignored schema = pa.schema([('strs', pa.utf8())]) @@ -2744,10 +2745,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") # type: ignore[arg-type] + pa.table("invalid input") # pyright: ignore[reportArgumentType] with pytest.raises(TypeError, match="Expected pandas DataFrame"): - pa.record_batch("invalid input") # type: ignore[arg-type] + pa.record_batch("invalid input") # pyright: ignore[reportArgumentType] def test_table_repr_to_string(): @@ -3814,7 +3815,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 # type: ignore[reportUnusedExpression] + _ = cuda_chunked_array == cuda_chunked_array # to_pandas() test with pytest.raises(NotImplementedError): @@ -4071,7 +4072,7 @@ def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch, # __dataframe__() test with pytest.raises(NotImplementedError): - from_dataframe(cuda_recordbatch.__dataframe__()) # type: ignore[misc] + from_dataframe(cuda_recordbatch.__dataframe__()) def verify_cuda_table(table, expected_schema): @@ -4270,7 +4271,7 @@ def test_table_non_cpu(cuda_context, cpu_table, cuda_table, # __dataframe__() test with pytest.raises(NotImplementedError): - from_dataframe(cuda_table.__dataframe__()) # type: ignore[misc] + from_dataframe(cuda_table.__dataframe__()) # __reduce__() test with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_tensor.py b/python/pyarrow/tests/test_tensor.py index c3726fdbbf41..9f94c2395d85 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) # type: ignore[reportArgumentType] + m = memoryview(tensor) # pyright: 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 04a842443989..cf490da2beb8 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -475,7 +475,7 @@ class BuggyTimezone2(datetime.tzinfo): def tzname(self, dt): return None - def utcoffset(self, dt): # type: ignore[override] + def utcoffset(self, dt): return "one hour" class BuggyTimezone3(datetime.tzinfo): @@ -483,7 +483,7 @@ class BuggyTimezone3(datetime.tzinfo): Wrong timezone name type """ - def tzname(self, dt): # type: ignore[override] + def tzname(self, dt): return 240 def utcoffset(self, dt): @@ -752,10 +752,10 @@ def test_struct_type(): # Neither integer nor string with pytest.raises(TypeError): - ty[None] # type: ignore[reportArgumentType] + ty[None] # pyright: ignore[reportArgumentType] with pytest.raises(TypeError): - ty.field(None) # type: ignore[reportArgumentType] + ty.field(None) # pyright: ignore[reportArgumentType] for a, b in zip(ty, fields): assert a == b @@ -882,7 +882,7 @@ def test_dictionary_type(): # invalid index type raises with pytest.raises(TypeError): - pa.dictionary(pa.string(), pa.int64()) # type: ignore[reportArgumentType] + pa.dictionary(pa.string(), pa.int64()) # pyright: ignore[reportArgumentType] def test_dictionary_ordered_equals(): @@ -971,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()) # type: ignore[reportArgumentType] + pa.run_end_encoded(pa.int8(), pa.utf8()) # pyright: ignore[reportArgumentType] @pytest.mark.parametrize('t,check_func', [ @@ -1149,11 +1149,11 @@ def test_key_value_metadata(): assert m1 != {'a': 'A', 'b': 'C'} with pytest.raises(TypeError): - pa.KeyValueMetadata({'a': 1}) # type: ignore[reportArgumentType] + pa.KeyValueMetadata({'a': 1}) # pyright: ignore[reportArgumentType] with pytest.raises(TypeError): - pa.KeyValueMetadata({1: 'a'}) # type: ignore[reportArgumentType] + pa.KeyValueMetadata({1: 'a'}) # pyright: ignore[reportArgumentType] with pytest.raises(TypeError): - pa.KeyValueMetadata(a=1) # type: ignore[reportArgumentType] + pa.KeyValueMetadata(a=1) # pyright: ignore[reportArgumentType] expected = [(b'a', b'A'), (b'b', b'B')] result = [(k, v) for k, v in m3.items()] From 17e3a08240be8c8f232e63f304b0d3ba70d06b08 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 02:18:58 +0100 Subject: [PATCH 78/93] Clean up the type stubs Remove the unused helper types from _stubs_typing, add the missing annotations, and import Self and TypeAlias from typing. --- python/pyarrow-stubs/__lib/array.pyi | 12 +++--- python/pyarrow-stubs/__lib/builder.pyi | 8 ++-- python/pyarrow-stubs/__lib/table.pyi | 12 +++--- python/pyarrow-stubs/__lib/types.pyi | 4 +- python/pyarrow-stubs/_stubs_typing.pyi | 60 +------------------------- python/pyarrow-stubs/types.pyi | 6 +-- 6 files changed, 20 insertions(+), 82 deletions(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index becc4c17aae0..329c4b8828bc 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -134,7 +134,7 @@ class Array(_PandasConvertible): memory_pool: MemoryPool | None = None, ) -> Array: ... def view(self, target_type: DataType | str) -> Array: ... - def sum(self, **kwargs) -> Scalar: ... + def sum(self, **kwargs: Any) -> Scalar: ... @property def type(self) -> DataType: ... def unique(self) -> Self: ... @@ -156,7 +156,7 @@ class Array(_PandasConvertible): length: int, buffers: Sequence[Buffer | None], null_count: int = -1, - offset=0, + offset: int = 0, children: Collection[Array] | None = None, ) -> Array: ... @property @@ -204,7 +204,7 @@ class Array(_PandasConvertible): self, order: Order = "ascending", null_placement: NullPlacement = "at_end", - **kwargs, + **kwargs: Any, ) -> Self: ... def __array__( self, dtype: np.dtype | None = None, copy: bool | None = None @@ -492,7 +492,7 @@ class StructArray(Array): arrays: Iterable[Array | np.ndarray | list], names: Sequence[str] | list[Field] | None = None, fields: list[Field] | None = None, - mask=None, + mask: BooleanArray | None = None, memory_pool: MemoryPool | None = None, type: StructType | None = None, ) -> StructArray: ... @@ -501,7 +501,7 @@ class StructArray(Array): order: Order = "ascending", null_placement: NullPlacement = "at_end", by: str | None = None, - **kwargs, + **kwargs: Any, ) -> StructArray: ... class RunEndEncodedArray(Array): @@ -517,7 +517,7 @@ class RunEndEncodedArray(Array): length: int, buffers: list[Buffer] | list[None], null_count: int = -1, - offset=0, + offset: int = 0, children: tuple[Array, Array] | list[list[int]] | None = None, ) -> RunEndEncodedArray: ... @property diff --git a/python/pyarrow-stubs/__lib/builder.pyi b/python/pyarrow-stubs/__lib/builder.pyi index 482bdc3bcd04..5d7c406f10e6 100644 --- a/python/pyarrow-stubs/__lib/builder.pyi +++ b/python/pyarrow-stubs/__lib/builder.pyi @@ -23,8 +23,8 @@ 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 append(self, value: str | bytes | float | None) -> None: ... + def append_values(self, values: Iterable[str | bytes | float | None]) -> None: ... def finish(self) -> StringArray: ... @property def null_count(self) -> int: ... @@ -32,8 +32,8 @@ class StringBuilder(_Weakrefable): 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 append(self, value: str | bytes | float | None) -> None: ... + def append_values(self, values: Iterable[str | bytes | float | None]) -> None: ... def finish(self) -> StringViewArray: ... @property def null_count(self) -> int: ... diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi index 7367d4d849db..a1d5b9412a4d 100644 --- a/python/pyarrow-stubs/__lib/table.pyi +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -158,7 +158,7 @@ class ChunkedArray(_PandasConvertible): element_size_limit: int = 100, ) -> str: ... @deprecated("Use 'to_string' instead") - def format(self, **kwargs) -> str: ... + def format(self, **kwargs: Any) -> str: ... def validate(self, *, full: bool = False) -> None: ... @property def null_count(self) -> int: ... @@ -210,7 +210,7 @@ class ChunkedArray(_PandasConvertible): self, order: Order = "ascending", null_placement: NullPlacement = "at_end", - **kwargs, + **kwargs: Any, ) -> ChunkedArray: ... def unify_dictionaries( self, memory_pool: MemoryPool | None = None @@ -288,7 +288,7 @@ class _Tabular(_PandasConvertible): def sort_by( self, sorting: str | Sequence[tuple[str, Order] | tuple[str, Order, NullPlacement]], - **kwargs, + **kwargs: Any, ) -> Self: ... def take(self, indices: Indices) -> Self: ... def filter( @@ -386,7 +386,7 @@ class RecordBatch(_Tabular): row_major: bool = True, memory_pool: MemoryPool | None = None, ) -> Tensor: ... - def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0): ... + def _export_to_c(self, out_ptr: int, out_schema_ptr: int = 0) -> None: ... @staticmethod def _import_from_c(in_ptr: int, schema: Schema | int) -> RecordBatch: ... def __arrow_c_array__( @@ -403,8 +403,8 @@ class RecordBatch(_Tabular): @staticmethod def _import_from_c_device(in_ptr: int, schema: Schema | int) -> RecordBatch: ... def __arrow_c_device_array__( - self, requested_schema: CapsuleType | None = None, **kwargs - ): ... + self, requested_schema: CapsuleType | None = None, **kwargs: Any + ) -> tuple[CapsuleType, CapsuleType]: ... @staticmethod def _import_from_c_device_capsule( schema_capsule: CapsuleType, array_capsule: CapsuleType diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index 7fc9c5d5fb73..45288c7947f9 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -16,11 +16,11 @@ # under the License. from collections.abc import Iterable, Iterator, Mapping, Sequence -from typing import Any, ClassVar, Literal, TypeAlias, overload +from typing import Any, ClassVar, Literal, Self, TypeAlias, overload import numpy as np import pandas as pd -from typing_extensions import CapsuleType, Self, deprecated +from typing_extensions import CapsuleType, deprecated from pyarrow._stubs_typing import ( BinaryType, diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi index 663fb129d254..a04f4b64dc7a 100644 --- a/python/pyarrow-stubs/_stubs_typing.pyi +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -15,16 +15,7 @@ # 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 collections.abc import Sequence from typing import Any, Literal, Protocol, TypeAlias import numpy as np @@ -35,62 +26,18 @@ from pyarrow import lib from pyarrow.lib import Array, 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" -] MapsAsPydicts: TypeAlias = Literal["lossy", "strict"] NullEncoding: TypeAlias = Literal["mask", "encode"] NullPlacement: TypeAlias = Literal["at_start", "at_end"] NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"] TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] -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_] | Array | 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: ... @@ -108,11 +55,6 @@ class SupportsArrowDeviceArray(Protocol): 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: ... - # Data types without a class of their own at runtime: their factories, e.g. # pyarrow.int32(), return plain DataType instances. These classes only exist # for type checking, where they distinguish the results of the factories and diff --git a/python/pyarrow-stubs/types.pyi b/python/pyarrow-stubs/types.pyi index 7fbbe0a4ec03..70cfc158cb08 100644 --- a/python/pyarrow-stubs/types.pyi +++ b/python/pyarrow-stubs/types.pyi @@ -18,16 +18,12 @@ import sys from enum import IntEnum -from typing import Any, TypeGuard +from typing import Any, TypeAlias, 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 e054197e0142d258b43bd9058d1e8d95a834043b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 02:19:08 +0100 Subject: [PATCH 79/93] Tidy the type checking setup Run stubtest independently of the working directory, skip the ruff hooks that the Lint job already ran in its type stubs step, and update the type checker comments in pyproject.toml. --- .github/workflows/dev.yml | 3 +++ python/pyarrow-stubs-tests/stubtest_allowlist.txt | 3 --- python/pyproject.toml | 7 +++---- python/scripts/run_stubtest.py | 11 +++++++---- 4 files changed, 13 insertions(+), 11 deletions(-) diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 5c24915a9426..8c6701a90a88 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -68,6 +68,9 @@ jobs: run: | pre-commit run --all-files --color=always --show-diff-on-failure - name: Run pre-commit for type stubs + env: + # The previous step already ran the ruff hooks. + SKIP: ruff-check,ruff-format run: | pre-commit run --hook-stage manual python-stubs --all-files --color=always --show-diff-on-failure diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt index 2800b7f800c2..f5004a3cc91e 100644 --- a/python/pyarrow-stubs-tests/stubtest_allowlist.txt +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -40,9 +40,6 @@ pyarrow\.__lib\.table\.NarySelector pyarrow\.__lib\.table\.NullarySelector pyarrow\.__lib\.table\.UnarySelector -# pyarrow.fs has no stub yet, so mypy reads its source. -pyarrow\.fs\.FSSpecHandler - # Stubs to fix (GH-32609). pyarrow\.__lib\.array\.Bool8Array\.from_numpy pyarrow\.__lib\.array\.Bool8Array\.from_storage diff --git a/python/pyproject.toml b/python/pyproject.toml index df33826eb830..7908fe1d7e67 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -112,7 +112,6 @@ version_scheme = 'guess-next-dev' git_describe_command = 'git describe --dirty --tags --long --match "apache-arrow-[0-9]*.*"' fallback_version = '27.0.0a0' -# TODO: Enable type checking once stubs are merged [tool.mypy] # mypy can't find a "-stubs" package on mypy_path, but it maps pyarrow-stubs/ # to the pyarrow module when checking it. Pass pyarrow-stubs as well when @@ -133,7 +132,6 @@ follow_imports = "skip" module = ["fsspec", "sparse"] ignore_missing_imports = true -# TODO: Enable type checking once stubs are merged [tool.pyright] pythonPlatform = "All" pythonVersion = "3.11" @@ -148,10 +146,11 @@ executionEnvironments = [ # dependency. mypy reports the other imports that can't be resolved. { root = ".", reportMissingImports = false }, ] -# In editable install, pyright finds the stubs but compiled modules (e.g. pyarrow.lib) +# Don't warn about imports that resolve to a stub without a Python source: +# pyarrow._stubs_typing only exists as a stub, and compiled modules such as +# pyarrow._compute have no source that pyright can read. reportMissingModuleSource = false -# TODO: Enable type checking once stubs are merged [tool.ty.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. diff --git a/python/scripts/run_stubtest.py b/python/scripts/run_stubtest.py index 3ec8a70b0fa9..e89e2a6e6db9 100644 --- a/python/scripts/run_stubtest.py +++ b/python/scripts/run_stubtest.py @@ -27,6 +27,7 @@ Pass --generate-allowlist to print the entries for the current mismatches. """ +import contextlib import os import sys import tempfile @@ -51,11 +52,13 @@ def main(argv): name += f".{stub.stem}" sys.modules[name] = pyarrow.lib - with tempfile.TemporaryDirectory() as tmpdir: - # mypy doesn't look for "-stubs" packages on MYPYPATH, so give it the - # stubs under the name of the package. + # mypy doesn't look for "-stubs" packages in its search path, so give it + # the stubs under the name of the package. Run it from the directory with + # the stubs: mypy searches the working directory too, so in python/ it + # would read the sources of the modules without stubs, e.g. pyarrow/fs.py, + # and stubtest would check them as if they were stubs. + with tempfile.TemporaryDirectory() as tmpdir, contextlib.chdir(tmpdir): os.symlink(STUBS_DIR, Path(tmpdir) / "pyarrow") - os.environ["MYPYPATH"] = tmpdir options = stubtest.parse_options([ "pyarrow", "--ignore-missing-stub", From 4be907a65f8ecaaa9ad8f955bea70977fd297ebb Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 12:54:06 +0100 Subject: [PATCH 80/93] Remove the type checker accommodations from the PyArrow tests The type checkers don't run on pyarrow/tests, so restore the import fallbacks, ignore comments, casts and isinstance asserts to what they were. Keep the assertions that were missing an assert. --- python/pyarrow/tests/test_array.py | 45 +++----- python/pyarrow/tests/test_convert_builtin.py | 37 +++--- python/pyarrow/tests/test_device.py | 10 +- python/pyarrow/tests/test_schema.py | 8 +- python/pyarrow/tests/test_sparse_tensor.py | 17 ++- python/pyarrow/tests/test_table.py | 115 ++++++------------- python/pyarrow/tests/test_tensor.py | 2 +- python/pyarrow/tests/test_types.py | 39 +++---- 8 files changed, 103 insertions(+), 170 deletions(-) diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index a2292150cf4d..1a1829f81a0d 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -27,15 +27,11 @@ import subprocess import sys import weakref -from typing import TYPE_CHECKING -if TYPE_CHECKING: +try: import numpy as np -else: - try: - import numpy as np - except ImportError: - np = None +except ImportError: + np = None import pyarrow as pa import pyarrow.tests.strategies as past @@ -77,7 +73,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]) # pyright: ignore[reportCallIssue] + pa.Array([1, 2]) def test_list_format(): @@ -770,8 +766,8 @@ def test_array_eq(): @pytest.mark.numpy def test_array_from_buffers(): - 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())) + values_buf = pa.py_buffer(np.int16([4, 5, 6, 7])) + nulls_buf = pa.py_buffer(np.uint8([0b00001101])) 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] @@ -786,9 +782,7 @@ 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 # pyright: ignore[reportArgumentType] - ) + pa.Array.from_buffers(pa.int16(), 3, ['', ''], offset=1) def test_string_binary_from_buffers(): @@ -982,8 +976,7 @@ 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"]) # pyright: ignore[reportArgumentType] + pa.StructArray.from_arrays([chunked_arr], ["foo"]) @pytest.mark.parametrize("offset", (0, 1)) @@ -1309,15 +1302,15 @@ def test_map_from_arrays(): 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, 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()) with pytest.raises(pa.ArrowTypeError, match='Mismatching map items type'): pa.MapArray.from_arrays(offsets, keys, items, pa.map_( @@ -3203,7 +3196,7 @@ def test_buffers_primitive(): assert struct.unpack('hhxxh', buffers[1].to_pybytes()) == (1, 2, 4) - a = pa.array(np.array([4, 5, 6], dtype=np.int8)) + a = pa.array(np.int8([4, 5, 6])) buffers = a.buffers() assert len(buffers) == 2 # No null bitmap from Numpy int array @@ -3289,7 +3282,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])) # pyright: ignore[reportCallIssue] + repr(pa.Tensor([1])) def test_invalid_tensor_construction(): @@ -3807,7 +3800,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]) # pyright: ignore[reportArgumentType] + mask=[1.0, 2.0, 3.0, 4.0]) with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], @@ -4085,11 +4078,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]) # pyright: ignore[reportArgumentType] + pa.concat_arrays([None]) arr = pa.chunked_array([[0, 1], [3, 4]]) with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays(arr) # pyright: ignore[reportArgumentType] + pa.concat_arrays(arr) @pytest.mark.pandas @@ -4618,7 +4611,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 1552c03cb238..3ac7be88eaa0 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -21,18 +21,13 @@ import itertools import math import re -from typing import TYPE_CHECKING, cast import hypothesis as h import pytest - -if TYPE_CHECKING: +try: import numpy as np -else: - try: - import numpy as np - except ImportError: - np = None +except ImportError: + np = None from pyarrow.pandas_compat import _pandas_api # noqa import pyarrow as pa @@ -71,7 +66,7 @@ def __int__(self): class MyBrokenInt: def __int__(self): - _ = 1/0 # MARKER + 1/0 # MARKER def test_iterable_types(): @@ -165,7 +160,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()) # pyright: ignore[reportArgumentType] + pa.array(ObjectWithOnlyGetitem()) def _as_list(xs): @@ -873,7 +868,7 @@ def test_large_binary_value(ty): assert isinstance(arr, pa.Array) assert arr.type == ty assert len(arr) == 4 - buf = cast(pa.FixedSizeBinaryScalar, arr[1]).as_buffer() + buf = arr[1].as_buffer() assert len(buf) == len(s) * nrepeats @@ -1119,11 +1114,11 @@ def expected_datetime_value(dt): ), ] utcdata = [ - pytz.utc.localize(cast(datetime.datetime, data[0])), + pytz.utc.localize(data[0]), data[1], None, - cast(datetime.datetime, data[3]).astimezone(pytz.utc), - cast(datetime.datetime, data[4]).astimezone(pytz.utc), + data[3].astimezone(pytz.utc), + data[4].astimezone(pytz.utc), ] ty = pa.timestamp(unit, tz=timezone) @@ -1251,9 +1246,9 @@ def test_sequence_timestamp_from_mixed_builtin_and_pandas_datetimes(): None, ] utcdata = [ - cast(datetime.datetime, data[0]).astimezone(pytz.utc), - pytz.utc.localize(cast(datetime.datetime, data[1])), - cast(datetime.datetime, data[2]).astimezone(pytz.utc), + data[0].astimezone(pytz.utc), + pytz.utc.localize(data[1]), + data[2].astimezone(pytz.utc), None, ] @@ -2131,8 +2126,8 @@ def test_map_from_dicts(): assert arr.to_pylist() == expected # With omitted values - data[1] = None # pyright: ignore[reportArgumentType, reportCallIssue] - expected[1] = None # pyright: ignore[reportArgumentType, reportCallIssue] + data[1] = None + expected[1] = None arr = pa.array(expected, type=pa.map_(pa.binary(), pa.int32())) @@ -2541,9 +2536,7 @@ def test_array_from_pylist_offset_overflow(): ), ( [pa.MonthDayNano([1, -1, -10100])], - [pa.scalar( - pa.MonthDayNano([1, -1, -10100]) - )], + [pa.scalar(pa.MonthDayNano([1, -1, -10100]))], 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 67981b0a0822..dc1a51e6d009 100644 --- a/python/pyarrow/tests/test_device.py +++ b/python/pyarrow/tests/test_device.py @@ -59,15 +59,11 @@ def test_copy_to(): batch_copied = batch.copy_to(dest) assert batch_copied.equals(batch) - 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 + assert batch_copied["col"].buffers()[1].device == mm.device + assert batch_copied["col"].buffers()[1].address != arr.buffers()[1].address with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): arr.copy_to(mm.device.device_type) with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - batch.copy_to(mm.device.device_type) # pyright: ignore[reportArgumentType] + batch.copy_to(mm.device.device_type) diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index b36a88a04511..01bf3c31a09b 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: - pass + np = None 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]) # pyright: ignore[reportArgumentType] + pa.schema([None]) def test_schema_weakref(): @@ -584,7 +584,7 @@ def test_schema_get_fields(): with pytest.raises(KeyError): schema.field('other') with pytest.raises(TypeError): - schema.field(0.0) # pyright: ignore[reportArgumentType] + schema.field(0.0) with pytest.raises(IndexError): schema.field(4) @@ -780,7 +780,7 @@ def test_schema_merge(): # raise proper error when passing a non-Schema value with pytest.raises(TypeError): - pa.unify_schemas([a, 1]) # pyright: ignore[reportArgumentType] + pa.unify_schemas([a, 1]) def test_undecodable_metadata(): diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index a224d1e4904f..eca8090d77a9 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -28,13 +28,13 @@ try: from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix except ImportError: - 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] + coo_matrix = None + csr_matrix = None + csr_array = None + coo_array = None try: - import sparse # pyright: ignore[reportMissingImports] + import sparse except ImportError: sparse = None @@ -401,7 +401,7 @@ def test_dense_to_sparse_tensor(dtype_str, arrow_type, sparse_tensor_type): assert np.array_equal(array, result_array) -@pytest.mark.skipif(coo_matrix is None, reason="requires scipy") +@pytest.mark.skipif(not coo_matrix, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (coo_array, coo_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, @@ -443,7 +443,7 @@ def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, assert out_scipy_matrix.has_canonical_format -@pytest.mark.skipif(csr_matrix is None, reason="requires scipy") +@pytest.mark.skipif(not csr_matrix, 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,8 +483,7 @@ def test_pydata_sparse_sparse_coo_tensor_roundtrip(dtype_str, arrow_type): shape = (4, 6) dim_names = ("x", "y") - sparse_array = sparse.COO( # pyright: ignore[reportOptionalMemberAccess] - data=data, coords=coords, shape=shape) + sparse_array = sparse.COO(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 2c2a6012537a..0fefabf02abe 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -18,13 +18,12 @@ from collections import OrderedDict from collections.abc import Iterable import sys -from typing import cast import weakref try: import numpy as np except ImportError: - pass + np = None import pytest import pyarrow as pa import pyarrow.compute as pc @@ -418,8 +417,7 @@ def test_to_pandas_empty_table(): table = pa.table(df) result = table.schema.empty_table().to_pandas() assert result.shape == (0, 2) - expected = df.iloc[:0] - tm.assert_frame_equal(result, expected) + tm.assert_frame_equal(result, df.iloc[:0]) @pytest.mark.pandas @@ -487,25 +485,12 @@ def test_chunked_array_unify_dictionaries(): pa.array(["foo", "bar", None, "foo"]).dictionary_encode(), pa.array(["quux", None, "foo"]).dictionary_encode(), ]) - 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"])) - + assert arr.chunk(0).dictionary.equals(pa.array(["foo", "bar"])) + assert arr.chunk(1).dictionary.equals(pa.array(["quux", "foo"])) arr = arr.unify_dictionaries() expected_dict = pa.array(["foo", "bar", "quux"]) - - 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.chunk(0).dictionary.equals(expected_dict) + assert arr.chunk(1).dictionary.equals(expected_dict) assert arr.to_pylist() == ["foo", "bar", None, "foo", "quux", None, "foo"] @@ -730,7 +715,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.column(0)._name == 'a0' + assert rb[0]._name == 'a0' def test_recordbatch_from_arrays_validate_schema(): @@ -812,7 +797,7 @@ def test_recordbatch_get_field(): batch.field('d') with pytest.raises(TypeError): - batch.field(None) # pyright: ignore[reportArgumentType] + batch.field(None) with pytest.raises(IndexError): batch.field(4) @@ -833,7 +818,7 @@ def test_recordbatch_select_column(): batch.column('d') with pytest.raises(TypeError): - batch.column(None) # pyright: ignore[reportArgumentType] + batch.column(None) with pytest.raises(IndexError): batch.column(4) @@ -969,10 +954,7 @@ def test_table_from_struct_array_chunked_array(): [[{"ints": 1}, {"floats": 1.0}]], type=pa.struct([("ints", pa.int32()), ("floats", pa.float32())]), ) - 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) + result = pa.Table.from_struct_array(chunked_struct_array) assert result.equals(pa.Table.from_arrays( [ pa.array([1, None], type=pa.int32()), @@ -1226,8 +1208,7 @@ 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") expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1261,8 +1242,7 @@ 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") expected = pa.Tensor.from_numpy(x) np.testing.assert_equal(result.to_numpy(), x) @@ -1484,7 +1464,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.column(0)._name == 'a0' + assert t[0]._name == 'a0' def test_table_equals(): @@ -1645,8 +1625,7 @@ 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])) - field0_metadata = table.field(0).metadata - assert field0_metadata is not None and b"a" in field0_metadata + assert b"a" in table.field(0).metadata assert table.field(1).nullable is False @@ -1711,7 +1690,7 @@ def test_table_get_field(): table.field('d') with pytest.raises(TypeError): - table.field(None) # pyright: ignore[reportArgumentType] + table.field(None) with pytest.raises(IndexError): table.field(4) @@ -1732,7 +1711,7 @@ def test_table_select_column(): table.column('d') with pytest.raises(TypeError): - table.column(None) # pyright: ignore[reportArgumentType] + table.column(None) with pytest.raises(IndexError): table.column(4) @@ -2025,41 +2004,22 @@ def test_table_unify_dictionaries(): table = pa.Table.from_batches([batch1, batch2]) table = table.replace_schema_metadata({b"key1": b"value1"}) - 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])) + 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])) table = table.unify_dictionaries(pa.default_memory_pool()) expected_dict_0 = pa.array(["foo", "bar", "quux"]) expected_dict_1 = pa.array([123, 456, 789]) - - 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.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) assert table.to_pydict() == { 'a': ["foo", "bar", None, "foo", "quux", "foo", None, "quux"], @@ -2129,14 +2089,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") # pyright: ignore[reportArgumentType] + pa.concat_tables([t, t], promote_options="invalid") def test_concat_tables_none_table(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_tables([None]) # pyright: ignore[reportArgumentType] + pa.concat_tables([None]) @pytest.mark.pandas @@ -2268,7 +2227,7 @@ def test_concat_batches_different_schema(): def test_concat_batches_none_batches(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_batches([None]) # pyright: ignore[reportArgumentType] + pa.concat_batches([None]) @pytest.mark.parametrize( @@ -2419,7 +2378,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) # pyright: ignore[reportArgumentType] + pa.Table.from_pydict(data, schema=schema) # Cannot pass both schema and metadata / names with pytest.raises(ValueError): @@ -2745,10 +2704,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") # pyright: ignore[reportArgumentType] + pa.table("invalid input") with pytest.raises(TypeError, match="Expected pandas DataFrame"): - pa.record_batch("invalid input") # pyright: ignore[reportArgumentType] + pa.record_batch("invalid input") def test_table_repr_to_string(): @@ -2882,8 +2841,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.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')) + 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')) def test_table_take_vanilla_functionality(): @@ -3815,7 +3774,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 # to_pandas() test with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_tensor.py b/python/pyarrow/tests/test_tensor.py index 9f94c2395d85..debb1066280c 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) # pyright: ignore[reportArgumentType] + m = memoryview(tensor) 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 cf490da2beb8..c4ac9e03cef9 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -25,22 +25,16 @@ import pytest import hypothesis as h import hypothesis.strategies as st -from typing import Any, TYPE_CHECKING +try: + import hypothesis.extra.pytz as tzst +except ImportError: + tzst = None import weakref -if TYPE_CHECKING: +try: import numpy as np - 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 - +except ImportError: + np = None import pyarrow as pa import pyarrow.types as types import pyarrow.tests.strategies as past @@ -421,7 +415,7 @@ def test_tzinfo_to_string_errors(): if tzst: timezones = tzst.timezones() else: - timezones = st.none() # type: ignore[assignment] + timezones = st.none() @h.given(timezones) @@ -752,10 +746,10 @@ def test_struct_type(): # Neither integer nor string with pytest.raises(TypeError): - ty[None] # pyright: ignore[reportArgumentType] + ty[None] with pytest.raises(TypeError): - ty.field(None) # pyright: ignore[reportArgumentType] + ty.field(None) for a, b in zip(ty, fields): assert a == b @@ -882,7 +876,7 @@ def test_dictionary_type(): # invalid index type raises with pytest.raises(TypeError): - pa.dictionary(pa.string(), pa.int64()) # pyright: ignore[reportArgumentType] + pa.dictionary(pa.string(), pa.int64()) def test_dictionary_ordered_equals(): @@ -971,7 +965,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()) # pyright: ignore[reportArgumentType] + pa.run_end_encoded(pa.int8(), pa.utf8()) @pytest.mark.parametrize('t,check_func', [ @@ -1149,11 +1143,11 @@ def test_key_value_metadata(): assert m1 != {'a': 'A', 'b': 'C'} with pytest.raises(TypeError): - pa.KeyValueMetadata({'a': 1}) # pyright: ignore[reportArgumentType] + pa.KeyValueMetadata({'a': 1}) with pytest.raises(TypeError): - pa.KeyValueMetadata({1: 'a'}) # pyright: ignore[reportArgumentType] + pa.KeyValueMetadata({1: 'a'}) with pytest.raises(TypeError): - pa.KeyValueMetadata(a=1) # pyright: ignore[reportArgumentType] + pa.KeyValueMetadata(a=1) expected = [(b'a', b'A'), (b'b', b'B')] result = [(k, v) for k, v in m3.items()] @@ -1280,7 +1274,6 @@ 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' @@ -1417,7 +1410,7 @@ def __arrow_c_schema__(self): return self.schema.__arrow_c_schema__() -class SchemaMapping(Mapping[Any, Any]): +class SchemaMapping(Mapping): def __init__(self, schema): self.schema = schema From 2d3f85c1a3b0b126571c04cde915cfe7bc39b204 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 13:10:12 +0100 Subject: [PATCH 81/93] Import scipy directly in the tensor type stub test Install scipy with the scipy-stubs extra, so the test doesn't need the TYPE_CHECKING import and importorskip. --- python/pyarrow-stubs-tests/test_tensor.py | 20 +++++++------------- python/requirements-stubs.txt | 3 ++- 2 files changed, 9 insertions(+), 14 deletions(-) diff --git a/python/pyarrow-stubs-tests/test_tensor.py b/python/pyarrow-stubs-tests/test_tensor.py index 1d0adae97d06..78c82582eb9a 100644 --- a/python/pyarrow-stubs-tests/test_tensor.py +++ b/python/pyarrow-stubs-tests/test_tensor.py @@ -17,20 +17,16 @@ # mypy: warn-unused-ignores from collections.abc import Hashable -from typing import TYPE_CHECKING import numpy as np -import pytest +from scipy.sparse import coo_array, csc_array, csr_array from typing_extensions import CapsuleType import pyarrow as pa from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check -if TYPE_CHECKING: - from scipy.sparse import coo_array, csc_array, csr_array # noqa: F401 - DENSE = np.array([[1, 0, 2], [0, 3, 0]]) @@ -89,16 +85,14 @@ def test_sparse_from_numpy() -> None: def test_sparse_scipy() -> None: - sparse = pytest.importorskip("scipy.sparse") - - coo = pa.SparseCOOTensor.from_scipy(sparse.coo_array(DENSE)) - check(assert_type(coo.to_scipy(), "coo_array"), sparse.coo_array) + coo = pa.SparseCOOTensor.from_scipy(coo_array(DENSE)) + check(assert_type(coo.to_scipy(), coo_array), coo_array) - csr = pa.SparseCSRMatrix.from_scipy(sparse.csr_array(DENSE)) - check(assert_type(csr.to_scipy(), "csr_array"), sparse.csr_array) + csr = pa.SparseCSRMatrix.from_scipy(csr_array(DENSE)) + check(assert_type(csr.to_scipy(), csr_array), csr_array) - csc = pa.SparseCSCMatrix.from_scipy(sparse.csc_array(DENSE)) - check(assert_type(csc.to_scipy(), "csc_array"), sparse.csc_array) + csc = pa.SparseCSCMatrix.from_scipy(csc_array(DENSE)) + check(assert_type(csc.to_scipy(), csc_array), csc_array) def test_invalid_usage() -> None: diff --git a/python/requirements-stubs.txt b/python/requirements-stubs.txt index 25aefadfebea..50ec8e495eac 100644 --- a/python/requirements-stubs.txt +++ b/python/requirements-stubs.txt @@ -7,7 +7,8 @@ fsspec==2026.9.0 mypy==2.4.0 pandas-stubs==3.0.5.260914 pyright==1.1.414 -scipy-stubs==1.17.1.5 +# The scipy extra installs the matching scipy, which the type stub tests import. +scipy-stubs[scipy]==1.17.1.5 ty==0.0.99 types-cffi==2.1.0.20260827 types-psutil==7.2.2.20260906 From 6ae817ed75db94b97c3dcab32f851e83ae64f2ec Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 13:19:13 +0100 Subject: [PATCH 82/93] Fix the type stub mismatches that stubtest skipped Match the runtime for the entries of the stubtest allowlist that were marked as stubs to fix: static methods for the from_arrays, from_storage, from_numpy and from_numpy_ndarray constructors, the obj parameter of array(), the integer values of DeviceAllocationType, the signatures of concat_tables and Schema.__getitem__, and no __all__ in pyarrow.types. Remove UnionArray.child, ChunkedArray.data, ChunkedArray.format and Schema.field_by_name, which no longer exist at runtime. Make the type aliases that only the table stubs define private, so that stubtest does not look for them at runtime. --- .../stubtest_allowlist.txt | 31 ---------- python/pyarrow-stubs/__lib/array.pyi | 35 +++++------ python/pyarrow-stubs/__lib/device.pyi | 30 +++++----- python/pyarrow-stubs/__lib/table.pyi | 58 +++++++++--------- python/pyarrow-stubs/__lib/types.pyi | 6 +- python/pyarrow-stubs/types.pyi | 59 ------------------- 6 files changed, 59 insertions(+), 160 deletions(-) diff --git a/python/pyarrow-stubs-tests/stubtest_allowlist.txt b/python/pyarrow-stubs-tests/stubtest_allowlist.txt index f5004a3cc91e..b361f3157c92 100644 --- a/python/pyarrow-stubs-tests/stubtest_allowlist.txt +++ b/python/pyarrow-stubs-tests/stubtest_allowlist.txt @@ -28,34 +28,3 @@ pyarrow\.__lib\.[a-z]+\.__all__ # The __init__ of Cython extension types takes (*args, **kwargs) at runtime. pyarrow\.__lib\.[a-z]+\.[A-Za-z_]+\.__init__ - -# Type aliases that are only defined in the stubs. -pyarrow\.__lib\.table\.AggregateOptions -pyarrow\.__lib\.table\.Aggregation -pyarrow\.__lib\.table\.ArrayOrChunkedArray -pyarrow\.__lib\.table\.ChunkedArrayConstructorInput -pyarrow\.__lib\.table\.ColumnSelector -pyarrow\.__lib\.table\.JoinType -pyarrow\.__lib\.table\.NarySelector -pyarrow\.__lib\.table\.NullarySelector -pyarrow\.__lib\.table\.UnarySelector - -# Stubs to fix (GH-32609). -pyarrow\.__lib\.array\.Bool8Array\.from_numpy -pyarrow\.__lib\.array\.Bool8Array\.from_storage -pyarrow\.__lib\.array\.FixedShapeTensorArray\.from_numpy_ndarray -pyarrow\.__lib\.array\.FixedSizeListArray\.from_arrays -pyarrow\.__lib\.array\.LargeListArray\.from_arrays -pyarrow\.__lib\.array\.LargeListViewArray\.from_arrays -pyarrow\.__lib\.array\.ListArray\.from_arrays -pyarrow\.__lib\.array\.ListViewArray\.from_arrays -pyarrow\.__lib\.array\.UnionArray\.child -pyarrow\.__lib\.array\.array -pyarrow\.__lib\.device\.DeviceAllocationType\.[A-Z_]+ -pyarrow\.__lib\.device\.DeviceAllocationType\.__format__ -pyarrow\.__lib\.table\.ChunkedArray\.data -pyarrow\.__lib\.table\.ChunkedArray\.format -pyarrow\.__lib\.table\.concat_tables -pyarrow\.__lib\.types\.Schema\.__getitem__ -pyarrow\.__lib\.types\.Schema\.field_by_name -pyarrow\.types\.__all__ diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 329c4b8828bc..1ca9ccc89cbb 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -45,7 +45,7 @@ from pyarrow.lib import ( Tensor, _Weakrefable, ) -from typing_extensions import CapsuleType, deprecated +from typing_extensions import CapsuleType import builtins from .scalar import Int64Scalar, Scalar @@ -53,7 +53,7 @@ from .device import DeviceAllocationType from .types import BaseExtensionType, DataType, Field, StructType def array( - values: Iterable[Any] + obj: Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray | SupportsPyArrowArray, @@ -302,9 +302,8 @@ class BaseListArray(Array): def value_lengths(self) -> Int32Array | Int64Array: ... class ListArray(BaseListArray): - @classmethod + @staticmethod def from_arrays( - cls, offsets: Int32Array | list[int] | list[int | None] | np.ndarray, values: Array | list[Any], type: DataType | None = None, @@ -318,9 +317,8 @@ class ListArray(BaseListArray): def value_lengths(self) -> Int32Array: ... class LargeListArray(BaseListArray): - @classmethod + @staticmethod def from_arrays( - cls, offsets: Int64Array | list[int] | list[int | None] | np.ndarray, values: Array | list[Any], type: DataType | None = None, @@ -334,9 +332,8 @@ class LargeListArray(BaseListArray): def value_lengths(self) -> Int64Array: ... class ListViewArray(BaseListArray): - @classmethod + @staticmethod def from_arrays( - cls, offsets: Int32Array | list[int] | list[int | None] | np.ndarray, sizes: Int32Array | list[int] | list[int | None] | np.ndarray, values: Array | list[Any], @@ -353,9 +350,8 @@ class ListViewArray(BaseListArray): def value_lengths(self) -> Int32Array: ... class LargeListViewArray(BaseListArray): - @classmethod + @staticmethod def from_arrays( - cls, offsets: Int64Array | list[int] | list[int | None] | np.ndarray, sizes: Int64Array | list[int] | list[int | None] | np.ndarray, values: Array | list[Any], @@ -372,9 +368,8 @@ class LargeListViewArray(BaseListArray): def value_lengths(self) -> Int64Array: ... class FixedSizeListArray(BaseListArray): - @classmethod + @staticmethod def from_arrays( - cls, values: Array | list[Any], list_size: int | None = None, type: DataType | None = None, @@ -400,8 +395,6 @@ class MapArray(ListArray): 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: ... @@ -538,10 +531,10 @@ class UuidArray(ExtensionArray): ... class FixedShapeTensorArray(ExtensionArray): def to_numpy_ndarray(self) -> np.ndarray: ... - @classmethod + @staticmethod def from_numpy_ndarray( - cls, obj: np.ndarray, dim_names: list[str] | tuple[str, ...] | None = None - ) -> Self: ... + obj: np.ndarray, dim_names: list[str] | tuple[str, ...] | None = None + ) -> FixedShapeTensorArray: ... class OpaqueArray(ExtensionArray): ... @@ -549,10 +542,10 @@ 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: ... + @staticmethod + def from_storage(storage: Int8Array) -> Bool8Array: ... # type: ignore[override] + @staticmethod + def from_numpy(obj: np.ndarray) -> Bool8Array: ... def concat_arrays( arrays: Iterable[Array], memory_pool: MemoryPool | None = None diff --git a/python/pyarrow-stubs/__lib/device.pyi b/python/pyarrow-stubs/__lib/device.pyi index c8e47ac7f82d..7bc20ffb7ec2 100644 --- a/python/pyarrow-stubs/__lib/device.pyi +++ b/python/pyarrow-stubs/__lib/device.pyi @@ -19,21 +19,21 @@ 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 DeviceAllocationType(enum.IntEnum): + CPU = 1 + CUDA = 2 + CUDA_HOST = 3 + OPENCL = 4 + VULKAN = 7 + METAL = 8 + VPI = 9 + ROCM = 10 + ROCM_HOST = 11 + EXT_DEV = 12 + CUDA_MANAGED = 13 + ONEAPI = 14 + WEBGPU = 15 + HEXAGON = 16 class Device(_Weakrefable): @property diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi index a1d5b9412a4d..28f2a1f8f3a3 100644 --- a/python/pyarrow-stubs/__lib/table.pyi +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -19,7 +19,7 @@ from typing import Self, TypeAlias from collections.abc import Iterable, Iterator, Sequence, Mapping from typing import Any, ClassVar, Literal -from typing_extensions import CapsuleType, deprecated +from typing_extensions import CapsuleType import builtins import numpy as np @@ -59,7 +59,7 @@ from .scalar import Int64Scalar, Scalar from .tensor import Tensor from .types import DataType -ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray +_ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray _Aggregation: TypeAlias = Literal[ "all", @@ -113,8 +113,8 @@ _AggregationPrefixed: TypeAlias = Literal[ "hash_tdigest", "hash_variance", ] -Aggregation: TypeAlias = _Aggregation | _AggregationPrefixed | str -AggregateOptions: TypeAlias = ( +_AggregationName: TypeAlias = _Aggregation | _AggregationPrefixed | str +_AggregateOptions: TypeAlias = ( ScalarAggregateOptions | CountOptions | TDigestOptions @@ -122,7 +122,7 @@ AggregateOptions: TypeAlias = ( | FunctionOptions ) -ChunkedArrayConstructorInput: TypeAlias = ( +_ChunkedArrayConstructorInput: TypeAlias = ( Iterable[ Iterable[Any] | SupportsArrowArray @@ -134,14 +134,12 @@ ChunkedArrayConstructorInput: TypeAlias = ( | SupportsArrowStream ) -UnarySelector: TypeAlias = str -NullarySelector: TypeAlias = tuple[()] -NarySelector: TypeAlias = list[str] | tuple[str, ...] -ColumnSelector: TypeAlias = UnarySelector | NullarySelector | NarySelector +_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: ... def length(self) -> int: ... @@ -157,8 +155,6 @@ class ChunkedArray(_PandasConvertible): skip_new_lines: bool = False, element_size_limit: int = 100, ) -> str: ... - @deprecated("Use 'to_string' instead") - def format(self, **kwargs: Any) -> str: ... def validate(self, *, full: bool = False) -> None: ... @property def null_count(self) -> int: ... @@ -236,7 +232,7 @@ class ChunkedArray(_PandasConvertible): def is_cpu(self) -> bool: ... def chunked_array( - arrays: ChunkedArrayConstructorInput, + arrays: _ChunkedArrayConstructorInput, type: DataType | str | None = None, ) -> ChunkedArray: ... @@ -247,9 +243,9 @@ class _Tabular(_PandasConvertible): def __dataframe__( self, nan_as_null: bool = False, allow_copy: bool = True ) -> _PyArrowDataFrame: ... - def __getitem__(self, key: int | str | slice) -> ArrayOrChunkedArray | Self: ... + def __getitem__(self, key: int | str | slice) -> _ArrayOrChunkedArray | Self: ... def __len__(self) -> int: ... - def column(self, i: int | str) -> ArrayOrChunkedArray: ... + def column(self, i: int | str) -> _ArrayOrChunkedArray: ... @property def column_names(self) -> list[str]: ... @property @@ -261,7 +257,7 @@ class _Tabular(_PandasConvertible): cls, mapping: Mapping[ Any, - ArrayOrChunkedArray + _ArrayOrChunkedArray | Iterable[Any] | SupportsArrowArray | SupportsArrowDeviceArray, @@ -276,7 +272,7 @@ class _Tabular(_PandasConvertible): schema: Schema | None = None, metadata: Mapping[str | bytes, str | bytes] | None = None, ) -> Self: ... - def itercolumns(self) -> Iterator[ArrayOrChunkedArray]: ... + def itercolumns(self) -> Iterator[_ArrayOrChunkedArray]: ... @property def num_columns(self) -> int: ... @property @@ -308,10 +304,13 @@ class _Tabular(_PandasConvertible): 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, + i: int, + field_: str | Field, + column: _ArrayOrChunkedArray | list[list[Any]], ) -> Self: ... def append_column( - self, field_: str | Field, column: ArrayOrChunkedArray | list[list[Any]] + self, field_: str | Field, column: _ArrayOrChunkedArray | list[list[Any]] ) -> Self: ... class RecordBatch(_Tabular): @@ -335,10 +334,10 @@ class RecordBatch(_Tabular): def get_total_buffer_size(self) -> int: ... def __sizeof__(self) -> int: ... def add_column( - self, i: int, field_: str | Field, column: ArrayOrChunkedArray | list[Any] + self, i: int, field_: str | Field, column: _ArrayOrChunkedArray | list[Any] ) -> RecordBatch: ... def append_column( - self, field_: str | Field, column: ArrayOrChunkedArray | list[Any] + self, field_: str | Field, column: _ArrayOrChunkedArray | list[Any] ) -> RecordBatch: ... def remove_column(self, i: int) -> RecordBatch: ... def set_column( @@ -415,7 +414,7 @@ class RecordBatch(_Tabular): def is_cpu(self) -> bool: ... def copy_to(self, destination: MemoryManager | Device) -> RecordBatch: ... -JoinType: TypeAlias = Literal[ +_JoinType: TypeAlias = Literal[ "left semi", "right semi", "left anti", @@ -505,19 +504,19 @@ class Table(_Tabular): self, i: int, field_: str | Field, - column: ChunkedArray | ChunkedArrayConstructorInput, + column: ChunkedArray | _ChunkedArrayConstructorInput, ) -> Table: ... def append_column( self, field_: str | Field, - column: ChunkedArray | ChunkedArrayConstructorInput, + column: ChunkedArray | _ChunkedArrayConstructorInput, ) -> Table: ... def remove_column(self, i: int) -> Table: ... def set_column( self, i: int, field_: str | Field, - column: ChunkedArray | ChunkedArrayConstructorInput, + column: ChunkedArray | _ChunkedArrayConstructorInput, ) -> Table: ... def rename_columns(self, names: list[str] | dict[str, str]) -> Table: ... def drop(self, columns: str | list[str]) -> Table: ... @@ -529,7 +528,7 @@ class Table(_Tabular): right_table: Table, keys: str | list[str], right_keys: str | list[str] | None = None, - join_type: JoinType = "left outer", + join_type: _JoinType = "left outer", left_suffix: str | None = None, right_suffix: str | None = None, coalesce_keys: bool = True, @@ -578,7 +577,6 @@ def concat_tables( tables: Iterable[Table], memory_pool: MemoryPool | None = None, promote_options: Literal["none", "default", "permissive"] = "none", - **kwargs: Any, ) -> Table: ... class TableGroupBy: @@ -590,8 +588,8 @@ class TableGroupBy: def aggregate( self, aggregations: Iterable[ - tuple[ColumnSelector, Aggregation] - | tuple[ColumnSelector, Aggregation, AggregateOptions | None] + tuple[_ColumnSelector, _AggregationName] + | tuple[_ColumnSelector, _AggregationName, _AggregateOptions | None] ], ) -> Table: ... diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index 45288c7947f9..82d231ed8766 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -20,7 +20,7 @@ from typing import Any, ClassVar, Literal, Self, TypeAlias, overload import numpy as np import pandas as pd -from typing_extensions import CapsuleType, deprecated +from typing_extensions import CapsuleType from pyarrow._stubs_typing import ( BinaryType, @@ -412,7 +412,7 @@ _SchemaFieldsInput: TypeAlias = _StructFieldsInput class Schema(_Weakrefable): def __len__(self) -> int: ... - def __getitem__(self, key: int) -> Field: ... + def __getitem__(self, key: int, /) -> Field: ... _field = __getitem__ def __iter__(self) -> Iterator[Field]: ... @@ -432,8 +432,6 @@ class Schema(_Weakrefable): 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: ... diff --git a/python/pyarrow-stubs/types.pyi b/python/pyarrow-stubs/types.pyi index 70cfc158cb08..06092148386a 100644 --- a/python/pyarrow-stubs/types.pyi +++ b/python/pyarrow-stubs/types.pyi @@ -25,8 +25,6 @@ if sys.version_info >= (3, 13): else: from typing_extensions import TypeIs -import pyarrow.lib as lib - from pyarrow.__lib.types import ( DataType, Decimal32Type, @@ -156,63 +154,6 @@ 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 From ff91bcfbc134829f6cb3fcd3c01a9293050259eb Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 14:35:49 +0100 Subject: [PATCH 83/93] Make the optional import fallbacks in the PyArrow tests pass mypy --- python/pyarrow/tests/test_array.py | 2 +- python/pyarrow/tests/test_convert_builtin.py | 2 +- python/pyarrow/tests/test_schema.py | 2 +- python/pyarrow/tests/test_sparse_tensor.py | 12 ++++++------ python/pyarrow/tests/test_table.py | 2 +- python/pyarrow/tests/test_types.py | 6 +++--- 6 files changed, 13 insertions(+), 13 deletions(-) diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 1a1829f81a0d..7619e65109b7 100644 --- a/python/pyarrow/tests/test_array.py +++ b/python/pyarrow/tests/test_array.py @@ -31,7 +31,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.tests.strategies as past diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index 3ac7be88eaa0..72b418715b4e 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -27,7 +27,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] from pyarrow.pandas_compat import _pandas_api # noqa import pyarrow as pa diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index 01bf3c31a09b..0ed3cec1408c 100644 --- a/python/pyarrow/tests/test_schema.py +++ b/python/pyarrow/tests/test_schema.py @@ -23,7 +23,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.tests.util as test_util diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index eca8090d77a9..a5b031ba4658 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -28,10 +28,10 @@ try: from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix except ImportError: - coo_matrix = None - csr_matrix = None - csr_array = None - coo_array = None + coo_matrix = None # type: ignore[assignment, misc] + csr_matrix = None # type: ignore[assignment, misc] + csr_array = None # type: ignore[assignment, misc] + coo_array = None # type: ignore[assignment, misc] try: import sparse @@ -401,7 +401,7 @@ def test_dense_to_sparse_tensor(dtype_str, arrow_type, sparse_tensor_type): assert np.array_equal(array, result_array) -@pytest.mark.skipif(not coo_matrix, reason="requires scipy") +@pytest.mark.skipif(coo_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (coo_array, coo_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, @@ -443,7 +443,7 @@ def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, assert out_scipy_matrix.has_canonical_format -@pytest.mark.skipif(not csr_matrix, reason="requires scipy") +@pytest.mark.skipif(csr_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (csr_array, csr_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_csr_matrix_scipy_roundtrip(dtype_str, arrow_type, diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index 0fefabf02abe..9a2a99b4e535 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -23,7 +23,7 @@ try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pytest import pyarrow as pa import pyarrow.compute as pc diff --git a/python/pyarrow/tests/test_types.py b/python/pyarrow/tests/test_types.py index c4ac9e03cef9..064eeb6e7a0e 100644 --- a/python/pyarrow/tests/test_types.py +++ b/python/pyarrow/tests/test_types.py @@ -28,13 +28,13 @@ try: import hypothesis.extra.pytz as tzst except ImportError: - tzst = None + tzst = None # type: ignore[assignment] import weakref try: import numpy as np except ImportError: - np = None + np = None # type: ignore[assignment] import pyarrow as pa import pyarrow.types as types import pyarrow.tests.strategies as past @@ -415,7 +415,7 @@ def test_tzinfo_to_string_errors(): if tzst: timezones = tzst.timezones() else: - timezones = st.none() + timezones = st.none() # type: ignore[assignment] @h.given(timezones) From c7d1d97759d4c68628ab42dbbe98b21bf3ed4aa0 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 15:30:43 +0100 Subject: [PATCH 84/93] Apply suggestion from @rok --- python/pyproject.toml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/pyproject.toml b/python/pyproject.toml index 7908fe1d7e67..36290f55c65b 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -178,8 +178,10 @@ include = ["pyarrow-stubs-tests/**"] # 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. +# Lints and formats the type stubs (*.pyi) only; see .pre-commit-config.yaml. +# The other Python files are checked with autopep8 and flake8. [tool.ruff] +include = ["pyarrow-stubs/**/*.pyi"] line-length = 88 [tool.ruff.lint] From d027724ec2e4f4314bc39409708aa69fccae02b2 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 16:25:19 +0100 Subject: [PATCH 85/93] Order imports and unify importing from pyarrow.lib --- python/pyarrow-stubs/__lib/array.pyi | 17 +++++----- python/pyarrow-stubs/__lib/builder.pyi | 9 +++-- python/pyarrow-stubs/__lib/error.pyi | 1 - python/pyarrow-stubs/__lib/table.pyi | 38 ++++++++++++--------- python/pyarrow-stubs/__lib/tensor.pyi | 8 ++--- python/pyarrow-stubs/__lib/types.pyi | 7 ++-- python/pyarrow-stubs/_stubs_typing.pyi | 3 +- python/pyarrow-stubs/lib.pyi | 2 +- python/pyarrow-stubs/types.pyi | 47 +++++++++++++------------- 9 files changed, 68 insertions(+), 64 deletions(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 1ca9ccc89cbb..757d277d93a9 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -15,20 +15,19 @@ # specific language governing permissions and limitations # under the License. +import builtins from collections.abc import Collection, Iterable, Iterator, Sequence - from typing import Any, ClassVar, Self, SupportsIndex import numpy as np import pandas as pd - from numpy.typing import NDArray from pyarrow._compute import CastOptions from pyarrow._stubs_typing import ( ArrayLike, Indices, - Mask, MapsAsPydicts, + Mask, NullEncoding, NullPlacement, NullSelectionBehavior, @@ -38,19 +37,21 @@ from pyarrow._stubs_typing import ( SupportsPyArrowArray, ) from pyarrow.lib import ( + BaseExtensionType, Buffer, + DataType, Device, + DeviceAllocationType, + Field, + Int64Scalar, MemoryManager, MemoryPool, + Scalar, + StructType, Tensor, _Weakrefable, ) from typing_extensions import CapsuleType -import builtins - -from .scalar import Int64Scalar, Scalar -from .device import DeviceAllocationType -from .types import BaseExtensionType, DataType, Field, StructType def array( obj: Iterable[Any] diff --git a/python/pyarrow-stubs/__lib/builder.pyi b/python/pyarrow-stubs/__lib/builder.pyi index 5d7c406f10e6..869e9c27727f 100644 --- a/python/pyarrow-stubs/__lib/builder.pyi +++ b/python/pyarrow-stubs/__lib/builder.pyi @@ -17,9 +17,12 @@ from collections.abc import Iterable -from pyarrow.lib import MemoryPool, _Weakrefable - -from .array import StringArray, StringViewArray +from pyarrow.lib import ( + MemoryPool, + StringArray, + StringViewArray, + _Weakrefable, +) class StringBuilder(_Weakrefable): def __init__(self, memory_pool: MemoryPool | None = None) -> None: ... diff --git a/python/pyarrow-stubs/__lib/error.pyi b/python/pyarrow-stubs/__lib/error.pyi index 6ba1c8396075..fa67de04535c 100644 --- a/python/pyarrow-stubs/__lib/error.pyi +++ b/python/pyarrow-stubs/__lib/error.pyi @@ -16,7 +16,6 @@ # under the License. from types import TracebackType - from typing import Self class ArrowException(Exception): ... diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi index 28f2a1f8f3a3..227419e03d41 100644 --- a/python/pyarrow-stubs/__lib/table.pyi +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -15,16 +15,12 @@ # specific language governing permissions and limitations # under the License. -from typing import Self, TypeAlias - -from collections.abc import Iterable, Iterator, Sequence, Mapping -from typing import Any, ClassVar, Literal -from typing_extensions import CapsuleType import builtins +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Any, ClassVar, Literal, Self, TypeAlias import numpy as np import pandas as pd - from numpy.typing import NDArray from pyarrow._compute import ( CastOptions, @@ -36,8 +32,8 @@ from pyarrow._compute import ( ) from pyarrow._stubs_typing import ( Indices, - Mask, MapsAsPydicts, + Mask, NullEncoding, NullPlacement, NullSelectionBehavior, @@ -47,17 +43,27 @@ from pyarrow._stubs_typing import ( SupportsArrowStream, SupportsPyArrowArray, ) +from pyarrow.lib import ( + Array, + Buffer, + DataType, + Device, + DeviceAllocationType, + Field, + Int64Scalar, + MemoryManager, + MemoryPool, + RecordBatchReader, + Scalar, + Schema, + StructArray, + Tensor, + _PandasConvertible, +) +from typing_extensions import CapsuleType + 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 .types import DataType _ArrayOrChunkedArray: TypeAlias = Array | ChunkedArray diff --git a/python/pyarrow-stubs/__lib/tensor.pyi b/python/pyarrow-stubs/__lib/tensor.pyi index 6cde448b34a1..281efe578595 100644 --- a/python/pyarrow-stubs/__lib/tensor.pyi +++ b/python/pyarrow-stubs/__lib/tensor.pyi @@ -15,14 +15,11 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Iterable, Sequence from typing import ClassVar -from typing_extensions import CapsuleType -from collections.abc import Iterable, Sequence import numpy as np - -from pyarrow.lib import _Weakrefable -from .types import DataType +from pyarrow.lib import DataType, _Weakrefable from scipy.sparse import ( coo_array, coo_matrix, @@ -32,6 +29,7 @@ from scipy.sparse import ( csr_matrix, ) from sparse import COO +from typing_extensions import CapsuleType class Tensor(_Weakrefable): @staticmethod diff --git a/python/pyarrow-stubs/__lib/types.pyi b/python/pyarrow-stubs/__lib/types.pyi index 82d231ed8766..d3f98a94c585 100644 --- a/python/pyarrow-stubs/__lib/types.pyi +++ b/python/pyarrow-stubs/__lib/types.pyi @@ -20,8 +20,6 @@ from typing import Any, ClassVar, Literal, Self, TypeAlias, overload import numpy as np import pandas as pd -from typing_extensions import CapsuleType - from pyarrow._stubs_typing import ( BinaryType, BinaryViewType, @@ -31,10 +29,10 @@ from pyarrow._stubs_typing import ( Float16Type, Float32Type, Float64Type, + Int8Type, Int16Type, Int32Type, Int64Type, - Int8Type, IntegerType, LargeBinaryType, LargeStringType, @@ -44,10 +42,10 @@ from pyarrow._stubs_typing import ( StringViewType, SupportsArrowSchema, TimeUnit, + UInt8Type, UInt16Type, UInt32Type, UInt64Type, - UInt8Type, ) from pyarrow.lib import ( Array, @@ -58,6 +56,7 @@ from pyarrow.lib import ( MemoryPool, Table, ) +from typing_extensions import CapsuleType class _Weakrefable: ... class _Metadata(_Weakrefable): ... diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi index a04f4b64dc7a..5d5751d42262 100644 --- a/python/pyarrow-stubs/_stubs_typing.pyi +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -19,11 +19,10 @@ from collections.abc import Sequence from typing import Any, Literal, Protocol, TypeAlias import numpy as np - from numpy.typing import NDArray +from pyarrow.lib import Array, ChunkedArray from pyarrow import lib -from pyarrow.lib import Array, ChunkedArray ArrayLike: TypeAlias = Any Order: TypeAlias = Literal["ascending", "descending"] diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi index d09176780766..046ceba7bcb5 100644 --- a/python/pyarrow-stubs/lib.pyi +++ b/python/pyarrow-stubs/lib.pyi @@ -21,7 +21,6 @@ from typing import Any, Self, final -from pyarrow.__lib.types import * # noqa: F401,F403 from pyarrow.__lib.array import * # noqa: F401,F403 from pyarrow.__lib.builder import * # noqa: F401,F403 from pyarrow.__lib.config import * # noqa: F401,F403 @@ -32,6 +31,7 @@ from pyarrow.__lib.memory import * # noqa: F401,F403 from pyarrow.__lib.scalar import * # noqa: F401,F403 from pyarrow.__lib.table import * # noqa: F401,F403 from pyarrow.__lib.tensor import * # noqa: F401,F403 +from pyarrow.__lib.types import * # noqa: F401,F403 @final class MonthDayNano(tuple): diff --git a/python/pyarrow-stubs/types.pyi b/python/pyarrow-stubs/types.pyi index 06092148386a..8f8773b08d56 100644 --- a/python/pyarrow-stubs/types.pyi +++ b/python/pyarrow-stubs/types.pyi @@ -17,7 +17,6 @@ import sys from enum import IntEnum - from typing import Any, TypeAlias, TypeGuard if sys.version_info >= (3, 13): @@ -25,29 +24,6 @@ if sys.version_info >= (3, 13): else: from typing_extensions import TypeIs -from pyarrow.__lib.types import ( - DataType, - Decimal32Type, - Decimal64Type, - Decimal128Type, - Decimal256Type, - DenseUnionType, - DictionaryType, - DurationType, - FixedSizeBinaryType, - FixedSizeListType, - LargeListType, - LargeListViewType, - ListType, - ListViewType, - MapType, - RunEndEncodedType, - SparseUnionType, - StructType, - Time32Type, - Time64Type, - TimestampType, -) from pyarrow._stubs_typing import ( BinaryType, BinaryViewType, @@ -73,6 +49,29 @@ from pyarrow._stubs_typing import ( UInt32Type, UInt64Type, ) +from pyarrow.lib import ( + DataType, + Decimal32Type, + Decimal64Type, + Decimal128Type, + Decimal256Type, + DenseUnionType, + DictionaryType, + DurationType, + FixedSizeBinaryType, + FixedSizeListType, + LargeListType, + LargeListViewType, + ListType, + ListViewType, + MapType, + RunEndEncodedType, + SparseUnionType, + StructType, + Time32Type, + Time64Type, + TimestampType, +) _SignedInteger: TypeAlias = Int8Type | Int16Type | Int32Type | Int64Type _UnsignedInteger: TypeAlias = UInt8Type | UInt16Type | UInt32Type | UInt64Type From 9fc46cdc1ed18712baeab8360901db116fb411f9 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 16:36:38 +0100 Subject: [PATCH 86/93] Fix the CI due to changes in imports --- python/pyarrow-stubs/__lib/array.pyi | 6 ++++-- python/pyarrow-stubs/__lib/table.pyi | 8 +++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 757d277d93a9..f8cb24812a18 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -22,6 +22,10 @@ from typing import Any, ClassVar, Self, SupportsIndex import numpy as np import pandas as pd from numpy.typing import NDArray + +# Placeholders (Any) until scalar.pyi is annotated; mypy rejects them as +# types when imported through pyarrow.lib. +from pyarrow.__lib.scalar import Int64Scalar, Scalar from pyarrow._compute import CastOptions from pyarrow._stubs_typing import ( ArrayLike, @@ -43,10 +47,8 @@ from pyarrow.lib import ( Device, DeviceAllocationType, Field, - Int64Scalar, MemoryManager, MemoryPool, - Scalar, StructType, Tensor, _Weakrefable, diff --git a/python/pyarrow-stubs/__lib/table.pyi b/python/pyarrow-stubs/__lib/table.pyi index 227419e03d41..3237d51d3e44 100644 --- a/python/pyarrow-stubs/__lib/table.pyi +++ b/python/pyarrow-stubs/__lib/table.pyi @@ -22,6 +22,11 @@ from typing import Any, ClassVar, Literal, Self, TypeAlias import numpy as np import pandas as pd from numpy.typing import NDArray + +# Placeholders (Any) until ipc.pyi and scalar.pyi are annotated; mypy +# rejects them as types when imported through pyarrow.lib. +from pyarrow.__lib.ipc import RecordBatchReader +from pyarrow.__lib.scalar import Int64Scalar, Scalar from pyarrow._compute import ( CastOptions, CountOptions, @@ -50,11 +55,8 @@ from pyarrow.lib import ( Device, DeviceAllocationType, Field, - Int64Scalar, MemoryManager, MemoryPool, - RecordBatchReader, - Scalar, Schema, StructArray, Tensor, From a8cc1aeda77611bcecb43a66b4d668536a04060d Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 10 Nov 2026 16:38:17 +0100 Subject: [PATCH 87/93] Update Indiced type alias --- python/pyarrow-stubs/_stubs_typing.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/_stubs_typing.pyi b/python/pyarrow-stubs/_stubs_typing.pyi index 5d5751d42262..e7b58eb63f75 100644 --- a/python/pyarrow-stubs/_stubs_typing.pyi +++ b/python/pyarrow-stubs/_stubs_typing.pyi @@ -34,7 +34,7 @@ TimeUnit: TypeAlias = Literal["s", "ms", "us", "ns"] Mask: TypeAlias = Sequence[bool | None] | NDArray[np.bool_] | Array | ChunkedArray Indices: TypeAlias = ( - Sequence[int | None] | NDArray[np.integer[Any]] | lib.IntegerArray | ChunkedArray + Sequence[int | None] | NDArray[np.integer[Any]] | Array | ChunkedArray ) class SupportsArrowStream(Protocol): From 96ba7b980b148005b90faae3b236ec47c7f036ec Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 16:48:15 +0100 Subject: [PATCH 88/93] array.pyi: from_pandas type to accept str --- python/pyarrow-stubs/__lib/array.pyi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index f8cb24812a18..67b1a330c8bb 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -149,7 +149,7 @@ class Array(_PandasConvertible): def from_pandas( obj: pd.Series | np.ndarray | ArrayLike, mask: Mask | None = None, - type: DataType | None = None, + type: DataType | str | None = None, safe: bool = True, memory_pool: MemoryPool | None = None, ) -> Array: ... From 0e446fcf19b55b3bf6fbdf0606eee1bc8a05ed4e Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 16:50:33 +0100 Subject: [PATCH 89/93] test_array.py: test from_pandas str and take Array indices --- python/pyarrow-stubs-tests/test_array.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 4ba957a1356c..c723d4fdf27a 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -44,6 +44,8 @@ def test_array() -> None: check(assert_type(arr.slice(1), pa.Int64Array), pa.Int64Array) check(assert_type(arr.fill_null(0), pa.Int64Array), pa.Int64Array) check(assert_type(arr.drop_null(), pa.Int64Array), pa.Int64Array) + indices: pa.Array = pa.array([1, 0]) + check(assert_type(arr.take(indices), pa.Int64Array), pa.Int64Array) check(assert_type(arr.is_null(), pa.BooleanArray), pa.BooleanArray) check(assert_type(arr.cast(pa.float64()), pa.Array), pa.DoubleArray) @@ -57,7 +59,7 @@ def test_array() -> None: check(assert_type(arr.to_tensor(allow_nulls=True), pa.Tensor), pa.Tensor) check(assert_type(pa.concat_arrays([arr, arr]), pa.Array), pa.Int64Array) - check(assert_type(pa.Array.from_pandas([1, 2], None, pa.int64()), + check(assert_type(pa.Array.from_pandas([1, 2], None, "int64"), pa.Array), pa.Int64Array) From faf04a92286181096d33ccc9895b325d4c1cb833 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 17:05:50 +0100 Subject: [PATCH 90/93] array.pyi: from_buffers accept Collection instead of Sequence --- python/pyarrow-stubs-tests/test_array.py | 5 +++-- python/pyarrow-stubs/__lib/array.pyi | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index c723d4fdf27a..7d7a22e44c34 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -123,8 +123,9 @@ def test_import_from_c_capsule() -> None: def test_from_buffers() -> None: values = cast(pa.Int64Array, pa.array([1, 2])) offsets = pa.array([0, 2], pa.int32()).buffers()[1] - # Children can be any collection of arrays, like a tuple - list_arr = pa.Array.from_buffers(pa.list_(pa.int64()), 1, [None, offsets], + # Buffers and children can be any collection, like dict values or a tuple + buffers = {"validity": None, "offsets": offsets} + list_arr = pa.Array.from_buffers(pa.list_(pa.int64()), 1, buffers.values(), children=(values,)) check(assert_type(list_arr, pa.Array), pa.ListArray) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index 67b1a330c8bb..a7dad314c3ab 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -157,7 +157,7 @@ class Array(_PandasConvertible): def from_buffers( type: DataType, length: int, - buffers: Sequence[Buffer | None], + buffers: Collection[Buffer | None], null_count: int = -1, offset: int = 0, children: Collection[Array] | None = None, @@ -464,7 +464,7 @@ class DictionaryArray(Array): def from_buffers( # type: ignore[override] type: DataType, length: int, - buffers: Sequence[Buffer | None], + buffers: Collection[Buffer | None], dictionary: Array, null_count: int = -1, offset: int = 0, From 4a899ebfffa7a84cab99f3ccc8c72c2153dcf9b9 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 18:25:01 +0100 Subject: [PATCH 91/93] Use TYPE_CHECKING for invalid usage in stub tests and fix review findings Replace the TYPE_CHECKING_INVALID_USAGE constant of the stub tests with typing.TYPE_CHECKING, in the tests and in the developer documentation. Require the argument of MonthDayNano, and test it together with RunEndEncodedArray.from_buffers, the names of StructArray.from_arrays and tables being unhashable. Move typing_extensions from the test dependency group to requirements-stubs.txt, and link the issue that tracks shipping the stubs in the documentation. --- docs/source/developers/python/development.rst | 25 +++++++++++-------- python/pyarrow-stubs-tests/_check.py | 15 +++++------ python/pyarrow-stubs-tests/test_array.py | 20 ++++++++++----- python/pyarrow-stubs-tests/test_table.py | 11 +++++--- python/pyarrow-stubs-tests/test_tensor.py | 5 ++-- python/pyarrow-stubs-tests/test_types.py | 7 +++--- python/pyarrow-stubs/lib.pyi | 5 +--- python/pyproject.toml | 3 +-- python/requirements-stubs.txt | 1 + 9 files changed, 51 insertions(+), 41 deletions(-) diff --git a/docs/source/developers/python/development.rst b/docs/source/developers/python/development.rst index c5bdaf6e51cf..ad46671b1cc4 100644 --- a/docs/source/developers/python/development.rst +++ b/docs/source/developers/python/development.rst @@ -107,8 +107,9 @@ Type Checking PyArrow provides type stubs (``*.pyi`` files) for static type checking. These stubs are located in the ``pyarrow-stubs/`` directory, which is laid out as a :pep:`561` stub-only package for ``pyarrow``, so type checkers and IDEs pick -it up from the ``arrow/python`` directory. The stubs are automatically included -in the distributed wheel packages. +it up from the ``arrow/python`` directory. The stubs aren't included in the +distributed wheel packages until they are complete (see +`GH-49831 `_). Running Type Checkers --------------------- @@ -125,8 +126,8 @@ uv). This also lints and formats the stubs with ruff: $ pre-commit run --hook-stage manual python-stubs --all-files The type checkers only run on request, while ruff also runs on commit. The -Lint CI job runs all of them as well. CI also runs the type checkers, the type stub tests and stubtest (see -below) against the PyArrow it builds, with +Lint CI job runs all of them as well. CI also runs the type checkers, the type +stub tests and stubtest (see below) against the PyArrow it builds, with ``ci/scripts/python_test_type_annotations.sh``. To run a type checker directly instead, install the pinned type checkers and @@ -223,16 +224,18 @@ verifies statically, with a runtime ``isinstance`` check: .. code-block:: python + from typing import TYPE_CHECKING + import pyarrow as pa - from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check + from _check import assert_type, check def test_list_type() -> None: lst = pa.list_(pa.int32()) check(assert_type(lst, pa.ListType), pa.ListType) check(assert_type(lst.value_type, pa.DataType), pa.DataType) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: pa.list_(1) # type: ignore[call-overload] The type checkers check these files as part of the regular runs above, and @@ -244,11 +247,11 @@ with what PyArrow actually returns: $ cd arrow/python $ pytest pyarrow-stubs-tests -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. ty doesn't honor mypy error codes, -so it doesn't check invalid usage. +Code under ``if TYPE_CHECKING:`` is never executed. It holds the usage that +should fail type checking, which is asserted with a +``# type: ignore[]`` comment: unused ignores are reported as errors +in these tests, so the test fails once the usage is no longer rejected. ty +doesn't honor mypy error codes, so it doesn't check invalid usage. Doctest ======= diff --git a/python/pyarrow-stubs-tests/_check.py b/python/pyarrow-stubs-tests/_check.py index d499c5f59949..4ebd47608b52 100644 --- a/python/pyarrow-stubs-tests/_check.py +++ b/python/pyarrow-stubs-tests/_check.py @@ -26,18 +26,15 @@ 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. +Code that type checkers should *reject* goes under ``if TYPE_CHECKING:`` +with a ``# type: ignore[]`` comment. The block never runs, and +mypy and pyright are configured to report unused ignores in these tests, so +the test fails if the usage stops being an error. """ -from typing import Any, Final, TYPE_CHECKING, TypeVar, assert_type +from typing import Any, TypeVar, assert_type -__all__ = ["TYPE_CHECKING_INVALID_USAGE", "assert_type", "check"] - -TYPE_CHECKING_INVALID_USAGE: Final = TYPE_CHECKING +__all__ = ["assert_type", "check"] _T = TypeVar("_T") diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 7d7a22e44c34..d804425af6e3 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -16,7 +16,7 @@ # under the License. from collections.abc import Hashable -from typing import Any, cast +from typing import Any, cast, TYPE_CHECKING import pytest from typing_extensions import CapsuleType @@ -24,7 +24,7 @@ import pyarrow as pa from pyarrow._compute import CastOptions -from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check +from _check import assert_type, check def test_array() -> None: @@ -72,7 +72,7 @@ def test_array() -> None: bool_mask = cast(pa.BooleanArray, pa.array([True, False, True])) check(assert_type(arr.filter(bool_mask), pa.Int64Array), pa.Int64Array) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: pa.array([1, 2], mask=[True, None]) # type: ignore[list-item] # TODO: when scalars get stubs # pa.array([1, 2], mask=pa.chunked_array([[True, False]])) @@ -92,7 +92,7 @@ def test_array_numpy() -> None: np_mask = np.array([True, False]) check(assert_type(arr.filter(np_mask), pa.Int64Array), pa.Int64Array) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: dtype = pa.dictionary(pa.int8(), pa.string()) strs = np.array(["a"]) pa.DictionaryArray.from_buffers(dtype, 1, [], strs) # type: ignore[arg-type] @@ -192,6 +192,11 @@ def test_run_end_encoded_array() -> None: ) check(assert_type(ree.values, pa.Array), pa.StringArray) + rebuilt = pa.RunEndEncodedArray.from_buffers( + ree.type, len(ree), [None], children=[ree.run_ends, ree.values] + ) + check(assert_type(rebuilt, pa.RunEndEncodedArray), pa.RunEndEncodedArray) + def test_chunked_array() -> None: chunked = pa.chunked_array([[1, 2], [3]]) @@ -213,7 +218,7 @@ def test_chunked_array() -> None: def test_invalid_usage() -> None: arr = cast(pa.Int64Array, pa.array([1, 2])) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: arr.cast(1) # type: ignore[arg-type] arr.slice("1") # type: ignore[arg-type] arr.equals(object()) # type: ignore[arg-type] @@ -234,9 +239,12 @@ def test_buffer_constructors_invalid_usage() -> None: children: list[pa.Array] = [pa.array([1])] strings = cast(pa.StringArray, pa.array(["a"])) off, data = cast(list[pa.Buffer], strings.buffers()[1:]) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: # Only lists are accepted for children and type codes pa.UnionArray.from_sparse(types, tuple(children)) # type: ignore[arg-type] pa.UnionArray.from_sparse(types, children, None, types) # type: ignore[arg-type] pa.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] + # Fields are passed as `fields`, `names` only takes strings + pa.StructArray.from_arrays( + children, names=[pa.field("a", pa.int64())]) # type: ignore[list-item] diff --git a/python/pyarrow-stubs-tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py index 1c9e80887472..4a1479190b23 100644 --- a/python/pyarrow-stubs-tests/test_table.py +++ b/python/pyarrow-stubs-tests/test_table.py @@ -15,15 +15,16 @@ # specific language governing permissions and limitations # under the License. +from collections.abc import Hashable from types import MappingProxyType -from typing import Any +from typing import Any, TYPE_CHECKING import pytest from typing_extensions import CapsuleType import pyarrow as pa -from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check +from _check import assert_type, check class _PyArrowArrayProvider: @@ -203,19 +204,21 @@ def test_invalid_usage() -> None: streams = [_ArrowStreamProvider()] columns = {"a": pa.array([1])}.values() # dict as non-sequence example - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: table.column(1.5) # type: ignore[arg-type] table.slice("0") # type: ignore[arg-type] pa.Table.from_arrays(columns, names=["a"]) # type: ignore[arg-type] # Stream providers are accepted as the whole input, not individual chunks pa.chunked_array(streams) # type: ignore[arg-type] table.add_column(0, "b", streams) # type: ignore[arg-type] + # Tables define __eq__ but not __hash__, hence unhashable at runtime + _: Hashable = table # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 def test_record_batch_invalid_usage() -> None: arr = pa.array([1]) chunked = pa.chunked_array([[{"a": 1}]]) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: # Tuples are failing at runtime. Recheck after the fix of GH-55418 pa.RecordBatch.from_arrays((arr,), names=["a"]) # type: ignore[arg-type] pa.RecordBatch.from_struct_array(chunked) # type: ignore[arg-type] diff --git a/python/pyarrow-stubs-tests/test_tensor.py b/python/pyarrow-stubs-tests/test_tensor.py index 78c82582eb9a..917736e7b3f0 100644 --- a/python/pyarrow-stubs-tests/test_tensor.py +++ b/python/pyarrow-stubs-tests/test_tensor.py @@ -17,6 +17,7 @@ # mypy: warn-unused-ignores from collections.abc import Hashable +from typing import TYPE_CHECKING import numpy as np @@ -25,7 +26,7 @@ import pyarrow as pa -from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check +from _check import assert_type, check DENSE = np.array([[1, 0, 2], [0, 3, 0]]) @@ -97,6 +98,6 @@ def test_sparse_scipy() -> None: def test_invalid_usage() -> None: tensor = pa.Tensor.from_numpy(DENSE) - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: pa.Tensor.from_numpy(DENSE, dim_names=[1, 2]) # type: ignore[list-item] _: Hashable = tensor # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs-tests/test_types.py b/python/pyarrow-stubs-tests/test_types.py index aef4f3477b79..f081cd62d8df 100644 --- a/python/pyarrow-stubs-tests/test_types.py +++ b/python/pyarrow-stubs-tests/test_types.py @@ -16,7 +16,7 @@ # under the License. from collections.abc import Hashable, Iterator -from typing import Literal +from typing import Literal, TYPE_CHECKING from typing_extensions import CapsuleType @@ -31,7 +31,7 @@ ) import pyarrow.types as types -from _check import TYPE_CHECKING_INVALID_USAGE, assert_type, check +from _check import assert_type, check def _erase(t: pa.DataType) -> pa.DataType: @@ -206,9 +206,10 @@ def test_import_from_c_capsule() -> None: def test_invalid_usage() -> None: - if TYPE_CHECKING_INVALID_USAGE: + if TYPE_CHECKING: # A Field needs a type unless it's created from an Arrow C schema. pa.field("a") # type: ignore[call-overload] pa.list_(1) # type: ignore[call-overload] pa.decimal128("10") # type: ignore[arg-type] + pa.MonthDayNano() # type: ignore[call-arg] # ty: ignore[missing-argument] _: Hashable = pa.KeyValueMetadata({b"a": b"1"}) # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 diff --git a/python/pyarrow-stubs/lib.pyi b/python/pyarrow-stubs/lib.pyi index 046ceba7bcb5..c585531eed4a 100644 --- a/python/pyarrow-stubs/lib.pyi +++ b/python/pyarrow-stubs/lib.pyi @@ -39,9 +39,6 @@ class MonthDayNano(tuple): days: int nanoseconds: int - def __new__( - cls, - sequence: tuple[int, int, int] | list[int] = ..., - ) -> Self: ... + def __new__(cls, sequence: tuple[int, int, int] | list[int]) -> Self: ... def __getattr__(name: str) -> Any: ... diff --git a/python/pyproject.toml b/python/pyproject.toml index 36290f55c65b..db8aba207f39 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -62,8 +62,7 @@ test = [ 'hypothesis', 'cffi', 'pytz', - 'pandas', - 'typing_extensions>=4.12.0', + 'pandas' ] [tool.numpydoc_validation] diff --git a/python/requirements-stubs.txt b/python/requirements-stubs.txt index 50ec8e495eac..ba6f48eec2fc 100644 --- a/python/requirements-stubs.txt +++ b/python/requirements-stubs.txt @@ -17,3 +17,4 @@ types-requests==2.33.0.20261006 # Annotated dependencies of the type stub tests, usually installed already. numpy pytest +typing_extensions>=4.12.0 From 91ae655fbdda19eb4b5f7fb3a5d4eac760a2aa57 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Nov 2026 19:00:10 +0100 Subject: [PATCH 92/93] Fix type checker errors in stub tests Pass the run-end encoded children as a tuple and remove two unused ignore comments reported by mypy, pyright and ty. --- python/pyarrow-stubs-tests/test_array.py | 4 ++-- python/pyarrow-stubs-tests/test_table.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index d804425af6e3..3ce5e76dcbb7 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -193,7 +193,7 @@ def test_run_end_encoded_array() -> None: check(assert_type(ree.values, pa.Array), pa.StringArray) rebuilt = pa.RunEndEncodedArray.from_buffers( - ree.type, len(ree), [None], children=[ree.run_ends, ree.values] + ree.type, len(ree), [None], children=(ree.run_ends, ree.values) ) check(assert_type(rebuilt, pa.RunEndEncodedArray), pa.RunEndEncodedArray) @@ -247,4 +247,4 @@ def test_buffer_constructors_invalid_usage() -> None: pa.StringArray.from_buffers(1, off, data, None, None) # type: ignore[arg-type] # Fields are passed as `fields`, `names` only takes strings pa.StructArray.from_arrays( - children, names=[pa.field("a", pa.int64())]) # type: ignore[list-item] + children, names=[pa.field("a", pa.int64())]) diff --git a/python/pyarrow-stubs-tests/test_table.py b/python/pyarrow-stubs-tests/test_table.py index 4a1479190b23..31732b991436 100644 --- a/python/pyarrow-stubs-tests/test_table.py +++ b/python/pyarrow-stubs-tests/test_table.py @@ -212,7 +212,7 @@ def test_invalid_usage() -> None: pa.chunked_array(streams) # type: ignore[arg-type] table.add_column(0, "b", streams) # type: ignore[arg-type] # Tables define __eq__ but not __hash__, hence unhashable at runtime - _: Hashable = table # type: ignore[assignment] # ty: ignore[invalid-assignment] # noqa: E501 + _: Hashable = table def test_record_batch_invalid_usage() -> None: From d2462bd5a344b73b2f71cdaaeebe98f134a0e4e8 Mon Sep 17 00:00:00 2001 From: Tadeja Kadunc Date: Tue, 10 Nov 2026 19:03:21 +0100 Subject: [PATCH 93/93] array.pyi: nulls/repeat return Array instead of ArrayLike --- python/pyarrow-stubs-tests/test_array.py | 2 ++ python/pyarrow-stubs/__lib/array.pyi | 4 ++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/python/pyarrow-stubs-tests/test_array.py b/python/pyarrow-stubs-tests/test_array.py index 3ce5e76dcbb7..57fabfa8f4e5 100644 --- a/python/pyarrow-stubs-tests/test_array.py +++ b/python/pyarrow-stubs-tests/test_array.py @@ -59,6 +59,8 @@ def test_array() -> None: check(assert_type(arr.to_tensor(allow_nulls=True), pa.Tensor), pa.Tensor) check(assert_type(pa.concat_arrays([arr, arr]), pa.Array), pa.Int64Array) + check(assert_type(pa.nulls(2, pa.int8()), pa.Array), pa.Int8Array) + check(assert_type(pa.repeat("a", 2), pa.Array), pa.StringArray) check(assert_type(pa.Array.from_pandas([1, 2], None, "int64"), pa.Array), pa.Int64Array) diff --git a/python/pyarrow-stubs/__lib/array.pyi b/python/pyarrow-stubs/__lib/array.pyi index a7dad314c3ab..21d8d732cd56 100644 --- a/python/pyarrow-stubs/__lib/array.pyi +++ b/python/pyarrow-stubs/__lib/array.pyi @@ -75,12 +75,12 @@ def nulls( size: int, type: Any | None = None, memory_pool: MemoryPool | None = None, -) -> ArrayLike: ... +) -> Array: ... def repeat( value: Any, size: int, memory_pool: MemoryPool | None = None, -) -> ArrayLike: ... +) -> Array: ... def infer_type( values: Sequence[Any] | np.ndarray, mask: Mask | None = None,