From ee419b0da2c4c189d9c64c99175fdd5fee5f914d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 22 Dec 2025 00:44:32 +0100 Subject: [PATCH 01/36] 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 b09050f6ced40d12d4fdbfbb32ea98f658cc8d9b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Feb 2026 21:08:58 +0100 Subject: [PATCH 02/36] 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 28a44c7158368f68c7e1c7672b604054bc7dddf5 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Mar 2026 22:16:26 +0100 Subject: [PATCH 03/36] 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 3d7a71391ca74d49619bf9fc73183defd676080a Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:27:18 +0100 Subject: [PATCH 04/36] 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 fb50721dc984ea6b6ba43a1661efe64e4b39a468 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:39:04 +0100 Subject: [PATCH 05/36] 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 fc48948032f548bf54b4b60b3de49ab93aa96077 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Fri, 13 Mar 2026 16:39:47 +0100 Subject: [PATCH 06/36] 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 9b03d17cb60528900c324ad9b1a475013b4ec49b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 01:20:19 +0100 Subject: [PATCH 07/36] 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 3874966aebc432aaf8324c306c59d0b739dae2f2 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 01:50:36 +0100 Subject: [PATCH 08/36] 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 32e439de886ea46b2d463ce0b8f8b15822323def Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 16 Mar 2026 13:57:38 +0100 Subject: [PATCH 09/36] 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 93d807e17ad0baa3513ad557b52115c669a08252 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 02:10:18 +0100 Subject: [PATCH 10/36] 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 53f15ba4e7e8..2bbfadd89ae0 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 558321c8da582801f5c37a02fcc20c06b916301b Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 21:39:36 +0100 Subject: [PATCH 11/36] 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 6a1d410578ff956782b2c21f8ce4901fad01c325 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 17 Mar 2026 22:38:22 +0100 Subject: [PATCH 12/36] 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 de68c8920240acce5170078917940e888a949f11 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Wed, 6 May 2026 17:52:37 +0200 Subject: [PATCH 13/36] 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 2bbfadd89ae0..53f15ba4e7e8 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 e3105e34ac7bd6b3a796aefe4bab8a23f006e985 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 16:16:48 +0200 Subject: [PATCH 14/36] 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 de7a5b682943673ce8faf7da967d15c3d1e7f867 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 16:42:49 +0200 Subject: [PATCH 15/36] 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 5d920e908d09c0f0ae29c22a40675ef7d3c3fa4d Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:00:13 -0700 Subject: [PATCH 16/36] 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 f7573acb260336ad736beb8b0921eb98e8fc0d20 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:22:40 -0700 Subject: [PATCH 17/36] 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 a0c8c7a2d30e33ccb6c5903612f10440c20cce7a Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 09:44:05 -0700 Subject: [PATCH 18/36] 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 a05d894cf5ca8d409cbda36c33d81d776ec01c0c Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 21:15:15 +0200 Subject: [PATCH 19/36] 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 821e91438d1b3ed669b90426835f27cc46d2d9aa Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 12:30:43 -0700 Subject: [PATCH 20/36] 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 3003457c57764ae46bbe46698b0726932c72dc41 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 9 Feb 2026 20:54:52 +0100 Subject: [PATCH 21/36] 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 bf9a6d672280..fa3376fc8daf 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -3954,7 +3954,7 @@ cdef class FixedSizeListArray(BaseListArray): Or create from a values array, list size and matching type: >>> typ = pa.list_(pa.field("values", pa.int64()), 2) - >>> arr = pa.FixedSizeListArray.from_arrays(values,type=typ) + >>> arr = pa.FixedSizeListArray.from_arrays(values, type=typ) >>> arr [ diff --git a/python/pyarrow/scalar.pxi b/python/pyarrow/scalar.pxi index f96fb40572b9..ad562a0dbd5c 100644 --- a/python/pyarrow/scalar.pxi +++ b/python/pyarrow/scalar.pxi @@ -1108,7 +1108,7 @@ cdef class StructScalar(Scalar, Mapping): Parameters ---------- - index : Union[int, str] + key : Union[int, str] Index / position or name of the field. Returns diff --git a/python/pyarrow/tests/test_array.py b/python/pyarrow/tests/test_array.py index 74ff837e16e8..bcd7ed14edbd 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_()) @@ -2874,7 +2881,7 @@ def test_interval_array_from_relativedelta(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, + pa.MonthDayNano([13, 8, # type: ignore[arg-type] (datetime.timedelta(seconds=1, microseconds=1, minutes=1, hours=1) // datetime.timedelta(microseconds=1)) * 1000])] @@ -2906,7 +2913,7 @@ def test_interval_array_from_tuple(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([1, 2, -3])] + pa.MonthDayNano([1, 2, -3])] # type: ignore[arg-type] expected = pa.array(expected_list) assert arr.equals(expected) assert arr.to_pylist() == expected_list @@ -2927,8 +2934,8 @@ def test_interval_array_from_dateoffset(): assert arr.type == pa.month_day_nano_interval() expected_list = [ None, - pa.MonthDayNano([13, 8, 3661000001001]), - pa.MonthDayNano([0, 0, 0])] + pa.MonthDayNano([13, 8, 3661000001001]), # type: ignore[arg-type] + pa.MonthDayNano([0, 0, 0])] # type: ignore[arg-type] expected = pa.array(expected_list) assert arr.equals(expected) expected_from_pandas = [ @@ -3194,7 +3201,7 @@ def test_buffers_primitive(): # Slicing does not affect the buffers but the offset a_sliced = a[1:] buffers = a_sliced.buffers() - a_sliced.offset == 1 + assert a_sliced.offset == 1 assert len(buffers) == 2 null_bitmap = buffers[0].to_pybytes() assert 1 <= len(null_bitmap) <= 64 # XXX this is varying @@ -3202,7 +3209,7 @@ def test_buffers_primitive(): assert struct.unpack('hhxxh', buffers[1].to_pybytes()) == (1, 2, 4) - a = pa.array(np.int8([4, 5, 6])) + a = pa.array(np.array([4, 5, 6], dtype=np.int8)) buffers = a.buffers() assert len(buffers) == 2 # No null bitmap from Numpy int array @@ -3288,7 +3295,7 @@ def test_nbytes_size(): def test_invalid_tensor_constructor_repr(): # ARROW-2638: prevent calling extension class constructors directly with pytest.raises(TypeError): - repr(pa.Tensor([1])) + repr(pa.Tensor([1])) # type: ignore[reportCallIssue] def test_invalid_tensor_construction(): @@ -3806,7 +3813,7 @@ def test_array_supported_masks(): with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], - mask=[1.0, 2.0, 3.0, 4.0]) + mask=[1.0, 2.0, 3.0, 4.0]) # type: ignore[reportArgumentType] with pytest.raises(pa.ArrowTypeError): arr = pa.array([4, None, 4, 3], @@ -4084,11 +4091,11 @@ def test_concat_array_invalid_type(): # ARROW-9920 - do not segfault on non-array input with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays([None]) + pa.concat_arrays([None]) # type: ignore[reportArgumentType] arr = pa.chunked_array([[0, 1], [3, 4]]) with pytest.raises(TypeError, match="should contain Array objects"): - pa.concat_arrays(arr) + pa.concat_arrays(arr) # type: ignore[reportArgumentType] @pytest.mark.pandas @@ -4617,7 +4624,7 @@ def test_non_cpu_array(): with pytest.raises(NotImplementedError): [i for i in iter(arr)] with pytest.raises(NotImplementedError): - arr == arr2 + _ = arr == arr2 with pytest.raises(NotImplementedError): arr.is_null() with pytest.raises(NotImplementedError): diff --git a/python/pyarrow/tests/test_convert_builtin.py b/python/pyarrow/tests/test_convert_builtin.py index 3ac7be88eaa0..19d33a6e2e05 100644 --- a/python/pyarrow/tests/test_convert_builtin.py +++ b/python/pyarrow/tests/test_convert_builtin.py @@ -21,13 +21,18 @@ import itertools import math import re +from typing import TYPE_CHECKING, cast import hypothesis as h import pytest -try: + +if TYPE_CHECKING: import numpy as np -except ImportError: - np = None +else: + try: + import numpy as np + except ImportError: + np = None from pyarrow.pandas_compat import _pandas_api # noqa import pyarrow as pa @@ -66,7 +71,7 @@ def __int__(self): class MyBrokenInt: def __int__(self): - 1/0 # MARKER + _ = 1/0 # MARKER def test_iterable_types(): @@ -160,7 +165,7 @@ def test_object_with_getitem(): # https://github.com/apache/arrow/issues/34944 # considered as sequence because of __getitem__, but has no length with pytest.raises(TypeError, match="has no len()"): - pa.array(ObjectWithOnlyGetitem()) + pa.array(ObjectWithOnlyGetitem()) # type: ignore[arg-type] def _as_list(xs): @@ -868,7 +873,7 @@ def test_large_binary_value(ty): assert isinstance(arr, pa.Array) assert arr.type == ty assert len(arr) == 4 - buf = arr[1].as_buffer() + buf = cast(pa.FixedSizeBinaryScalar, arr[1]).as_buffer() assert len(buf) == len(s) * nrepeats @@ -1114,11 +1119,11 @@ def expected_datetime_value(dt): ), ] utcdata = [ - pytz.utc.localize(data[0]), + pytz.utc.localize(cast(datetime.datetime, data[0])), data[1], None, - data[3].astimezone(pytz.utc), - data[4].astimezone(pytz.utc), + cast(datetime.datetime, data[3]).astimezone(pytz.utc), + cast(datetime.datetime, data[4]).astimezone(pytz.utc), ] ty = pa.timestamp(unit, tz=timezone) @@ -1246,9 +1251,9 @@ def test_sequence_timestamp_from_mixed_builtin_and_pandas_datetimes(): None, ] utcdata = [ - data[0].astimezone(pytz.utc), - pytz.utc.localize(data[1]), - data[2].astimezone(pytz.utc), + cast(datetime.datetime, data[0]).astimezone(pytz.utc), + pytz.utc.localize(cast(datetime.datetime, data[1])), + cast(datetime.datetime, data[2]).astimezone(pytz.utc), None, ] @@ -2126,8 +2131,8 @@ def test_map_from_dicts(): assert arr.to_pylist() == expected # With omitted values - data[1] = None - expected[1] = None + data[1] = None # type: ignore[call-overload] + expected[1] = None # type: ignore[call-overload] arr = pa.array(expected, type=pa.map_(pa.binary(), pa.int32())) @@ -2452,6 +2457,7 @@ def test_nested_auto_chunking(ty, char): } +@pytest.mark.numpy @pytest.mark.large_memory def test_array_from_pylist_data_overflow(): # Regression test for ARROW-12983 @@ -2474,6 +2480,7 @@ def test_array_from_pylist_data_overflow(): assert len(arr.chunks) > 1 +@pytest.mark.numpy @pytest.mark.slow @pytest.mark.large_memory def test_array_from_pylist_offset_overflow(): @@ -2498,6 +2505,7 @@ def test_array_from_pylist_offset_overflow(): assert len(arr.chunks) > 1 +@pytest.mark.numpy @parametrize_with_collections_types @pytest.mark.parametrize(('data', 'scalar_data', 'value_type'), [ ([True, False, None], [pa.scalar(True), pa.scalar(False), None], pa.bool_()), @@ -2535,8 +2543,10 @@ def test_array_from_pylist_offset_overflow(): pa.timestamp('us') ), ( - [pa.MonthDayNano([1, -1, -10100])], - [pa.scalar(pa.MonthDayNano([1, -1, -10100]))], + [pa.MonthDayNano([1, -1, -10100])], # type: ignore[call-arg, arg-type] + [pa.scalar( + pa.MonthDayNano([1, -1, -10100]) # type: ignore[call-arg, arg-type] + )], pa.month_day_nano_interval() ), (["a", "b"], [pa.scalar("a"), pa.scalar("b")], pa.string()), diff --git a/python/pyarrow/tests/test_device.py b/python/pyarrow/tests/test_device.py index dc1a51e6d009..00f8bbf720da 100644 --- a/python/pyarrow/tests/test_device.py +++ b/python/pyarrow/tests/test_device.py @@ -59,11 +59,15 @@ def test_copy_to(): batch_copied = batch.copy_to(dest) assert batch_copied.equals(batch) - assert batch_copied["col"].buffers()[1].device == mm.device - assert batch_copied["col"].buffers()[1].address != arr.buffers()[1].address + buffer = batch_copied.column("col").buffers()[1] + assert buffer is not None + assert buffer.device == mm.device + buffer_orig = arr.buffers()[1] + assert buffer_orig is not None + assert buffer.address != buffer_orig.address with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - arr.copy_to(mm.device.device_type) + arr.copy_to(mm.device.device_type) # type: ignore[arg-type] with pytest.raises(TypeError, match="Argument 'destination' has incorrect type"): - batch.copy_to(mm.device.device_type) + batch.copy_to(mm.device.device_type) # type: ignore[arg-type] diff --git a/python/pyarrow/tests/test_schema.py b/python/pyarrow/tests/test_schema.py index b35b5350a814..15f61b1b4002 100644 --- a/python/pyarrow/tests/test_schema.py +++ b/python/pyarrow/tests/test_schema.py @@ -23,7 +23,7 @@ try: import numpy as np except ImportError: - np = None + pass import pyarrow as pa import pyarrow.tests.util as test_util @@ -253,7 +253,7 @@ def test_schema(): child 0, item: int8""" with pytest.raises(TypeError): - pa.schema([None]) + pa.schema([None]) # type: ignore[list-item] def test_schema_weakref(): @@ -542,7 +542,7 @@ def test_schema_equals_invalid_type(): for val in [None, 'string', pa.array([1, 2])]: with pytest.raises(TypeError): - schema.equals(val) + schema.equals(val) # type: ignore[invalid-argument-type] def test_schema_equality_operators(): @@ -588,7 +588,7 @@ def test_schema_get_fields(): with pytest.raises(KeyError): schema.field('other') with pytest.raises(TypeError): - schema.field(0.0) + schema.field(0.0) # type: ignore[arg-type] with pytest.raises(IndexError): schema.field(4) @@ -700,6 +700,7 @@ def test_empty_table(): assert table.schema == schema +@pytest.mark.numpy @pytest.mark.pandas def test_schema_from_pandas(): import pandas as pd @@ -784,7 +785,7 @@ def test_schema_merge(): # raise proper error when passing a non-Schema value with pytest.raises(TypeError): - pa.unify_schemas([a, 1]) + pa.unify_schemas([a, 1]) # type: ignore[list-item] def test_undecodable_metadata(): diff --git a/python/pyarrow/tests/test_sparse_tensor.py b/python/pyarrow/tests/test_sparse_tensor.py index eca8090d77a9..2ce48b651b14 100644 --- a/python/pyarrow/tests/test_sparse_tensor.py +++ b/python/pyarrow/tests/test_sparse_tensor.py @@ -26,15 +26,16 @@ import pyarrow as pa try: - from scipy.sparse import csr_array, coo_array, csr_matrix, coo_matrix + from scipy.sparse import ( # type: ignore[reportMissingModuleSource] + csr_array, coo_array, csr_matrix, coo_matrix) except ImportError: - coo_matrix = None - csr_matrix = None - csr_array = None - coo_array = None + coo_matrix = None # type: ignore[assignment, misc] + csr_matrix = None # type: ignore[assignment, misc] + csr_array = None # type: ignore[assignment, misc] + coo_array = None # type: ignore[assignment, misc] try: - import sparse + import sparse # type: ignore[import-untyped, import-not-found] except ImportError: sparse = None @@ -401,7 +402,7 @@ def test_dense_to_sparse_tensor(dtype_str, arrow_type, sparse_tensor_type): assert np.array_equal(array, result_array) -@pytest.mark.skipif(not coo_matrix, reason="requires scipy") +@pytest.mark.skipif(coo_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (coo_array, coo_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, @@ -443,7 +444,7 @@ def test_sparse_coo_tensor_scipy_roundtrip(dtype_str, arrow_type, assert out_scipy_matrix.has_canonical_format -@pytest.mark.skipif(not csr_matrix, reason="requires scipy") +@pytest.mark.skipif(csr_matrix is None, reason="requires scipy") @pytest.mark.parametrize('sparse_object', (csr_array, csr_matrix)) @pytest.mark.parametrize('dtype_str,arrow_type', scipy_type_pairs) def test_sparse_csr_matrix_scipy_roundtrip(dtype_str, arrow_type, @@ -483,7 +484,8 @@ def test_pydata_sparse_sparse_coo_tensor_roundtrip(dtype_str, arrow_type): shape = (4, 6) dim_names = ("x", "y") - sparse_array = sparse.COO(data=data, coords=coords, shape=shape) + sparse_array = sparse.COO( # type: ignore[reportOptionalMemberAccess] + data=data, coords=coords, shape=shape) sparse_tensor = pa.SparseCOOTensor.from_pydata_sparse(sparse_array, dim_names=dim_names) out_sparse_array = sparse_tensor.to_pydata_sparse() diff --git a/python/pyarrow/tests/test_table.py b/python/pyarrow/tests/test_table.py index eb02051693f2..2502292acf89 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -18,12 +18,13 @@ from collections import OrderedDict from collections.abc import Iterable import sys +from typing import cast import weakref try: import numpy as np except ImportError: - np = None + pass import pytest import pyarrow as pa import pyarrow.compute as pc @@ -417,7 +418,8 @@ def test_to_pandas_empty_table(): table = pa.table(df) result = table.schema.empty_table().to_pandas() assert result.shape == (0, 2) - tm.assert_frame_equal(result, df.iloc[:0]) + expected = df.iloc[:0] + tm.assert_frame_equal(result, expected) @pytest.mark.pandas @@ -485,12 +487,25 @@ def test_chunked_array_unify_dictionaries(): pa.array(["foo", "bar", None, "foo"]).dictionary_encode(), pa.array(["quux", None, "foo"]).dictionary_encode(), ]) - assert arr.chunk(0).dictionary.equals(pa.array(["foo", "bar"])) - assert arr.chunk(1).dictionary.equals(pa.array(["quux", "foo"])) + chunk_0 = arr.chunk(0) + assert isinstance(chunk_0, pa.DictionaryArray) + assert chunk_0.dictionary.equals(pa.array(["foo", "bar"])) + + chunk_1 = arr.chunk(1) + assert isinstance(chunk_1, pa.DictionaryArray) + assert chunk_1.dictionary.equals(pa.array(["quux", "foo"])) + arr = arr.unify_dictionaries() expected_dict = pa.array(["foo", "bar", "quux"]) - assert arr.chunk(0).dictionary.equals(expected_dict) - assert arr.chunk(1).dictionary.equals(expected_dict) + + chunk_0 = arr.chunk(0) + assert isinstance(chunk_0, pa.DictionaryArray) + assert chunk_0.dictionary.equals(expected_dict) + + chunk_1 = arr.chunk(1) + assert isinstance(chunk_1, pa.DictionaryArray) + assert chunk_1.dictionary.equals(expected_dict) + assert arr.to_pylist() == ["foo", "bar", None, "foo", "quux", None, "foo"] @@ -715,7 +730,7 @@ def test_recordbatch_take(): def test_recordbatch_column_sets_private_name(): # ARROW-6429 rb = pa.record_batch([pa.array([1, 2, 3, 4])], names=['a0']) - assert rb[0]._name == 'a0' + assert rb.column(0)._name == 'a0' def test_recordbatch_from_arrays_validate_schema(): @@ -797,7 +812,7 @@ def test_recordbatch_get_field(): batch.field('d') with pytest.raises(TypeError): - batch.field(None) + batch.field(None) # type: ignore[arg-type] with pytest.raises(IndexError): batch.field(4) @@ -818,7 +833,7 @@ def test_recordbatch_select_column(): batch.column('d') with pytest.raises(TypeError): - batch.column(None) + batch.column(None) # type: ignore[arg-type] with pytest.raises(IndexError): batch.column(4) @@ -937,7 +952,10 @@ def test_table_from_struct_array_chunked_array(): [[{"ints": 1}, {"floats": 1.0}]], type=pa.struct([("ints", pa.int32()), ("floats", pa.float32())]), ) - result = pa.Table.from_struct_array(chunked_struct_array) + assert isinstance(chunked_struct_array.type, pa.StructType) + # Cast to the proper type for type checker + struct_chunked_array = cast(pa.ChunkedArray, chunked_struct_array) + result = pa.Table.from_struct_array(struct_chunked_array) assert result.equals(pa.Table.from_arrays( [ pa.array([1, None], type=pa.int32()), @@ -1191,7 +1209,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) @@ -1225,7 +1243,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) @@ -1447,7 +1465,7 @@ def test_recordbatchlist_schema_equals(): def test_table_column_sets_private_name(): # ARROW-6429 t = pa.table([pa.array([1, 2, 3, 4])], names=['a0']) - assert t[0]._name == 'a0' + assert t.column(0)._name == 'a0' def test_table_equals(): @@ -1608,7 +1626,8 @@ def test_table_from_arrays_preserves_column_metadata(): field1 = pa.field('field2', pa.int64(), nullable=False) table = pa.Table.from_arrays([arr0, arr1], schema=pa.schema([field0, field1])) - assert b"a" in table.field(0).metadata + field0_metadata = table.field(0).metadata + assert field0_metadata is not None and b"a" in field0_metadata assert table.field(1).nullable is False @@ -1673,7 +1692,7 @@ def test_table_get_field(): table.field('d') with pytest.raises(TypeError): - table.field(None) + table.field(None) # type: ignore[arg-type] with pytest.raises(IndexError): table.field(4) @@ -1694,7 +1713,7 @@ def test_table_select_column(): table.column('d') with pytest.raises(TypeError): - table.column(None) + table.column(None) # type: ignore[arg-type] with pytest.raises(IndexError): table.column(4) @@ -1987,22 +2006,41 @@ def test_table_unify_dictionaries(): table = pa.Table.from_batches([batch1, batch2]) table = table.replace_schema_metadata({b"key1": b"value1"}) - assert table.column(0).chunk(0).dictionary.equals( - pa.array(["foo", "bar"])) - assert table.column(0).chunk(1).dictionary.equals( - pa.array(["quux", "foo"])) - assert table.column(1).chunk(0).dictionary.equals( - pa.array([123, 456, 789])) - assert table.column(1).chunk(1).dictionary.equals( - pa.array([456, 789])) + chunk_0_0 = table.column(0).chunk(0) + assert isinstance(chunk_0_0, pa.DictionaryArray) + assert chunk_0_0.dictionary.equals(pa.array(["foo", "bar"])) + + chunk_0_1 = table.column(0).chunk(1) + assert isinstance(chunk_0_1, pa.DictionaryArray) + assert chunk_0_1.dictionary.equals(pa.array(["quux", "foo"])) + + chunk_1_0 = table.column(1).chunk(0) + assert isinstance(chunk_1_0, pa.DictionaryArray) + assert chunk_1_0.dictionary.equals(pa.array([123, 456, 789])) + + chunk_1_1 = table.column(1).chunk(1) + assert isinstance(chunk_1_1, pa.DictionaryArray) + assert chunk_1_1.dictionary.equals(pa.array([456, 789])) table = table.unify_dictionaries(pa.default_memory_pool()) expected_dict_0 = pa.array(["foo", "bar", "quux"]) expected_dict_1 = pa.array([123, 456, 789]) - assert table.column(0).chunk(0).dictionary.equals(expected_dict_0) - assert table.column(0).chunk(1).dictionary.equals(expected_dict_0) - assert table.column(1).chunk(0).dictionary.equals(expected_dict_1) - assert table.column(1).chunk(1).dictionary.equals(expected_dict_1) + + chunk_0_0 = table.column(0).chunk(0) + assert isinstance(chunk_0_0, pa.DictionaryArray) + assert chunk_0_0.dictionary.equals(expected_dict_0) + + chunk_0_1 = table.column(0).chunk(1) + assert isinstance(chunk_0_1, pa.DictionaryArray) + assert chunk_0_1.dictionary.equals(expected_dict_0) + + chunk_1_0 = table.column(1).chunk(0) + assert isinstance(chunk_1_0, pa.DictionaryArray) + assert chunk_1_0.dictionary.equals(expected_dict_1) + + chunk_1_1 = table.column(1).chunk(1) + assert isinstance(chunk_1_1, pa.DictionaryArray) + assert chunk_1_1.dictionary.equals(expected_dict_1) assert table.to_pydict() == { 'a': ["foo", "bar", None, "foo", "quux", "foo", None, "quux"], @@ -2072,13 +2110,13 @@ def test_concat_tables_invalid_option(): t = pa.Table.from_arrays([list(range(10))], names=('a',)) with pytest.raises(ValueError, match="Invalid promote_options: invalid"): - pa.concat_tables([t, t], promote_options="invalid") + pa.concat_tables([t, t], promote_options="invalid") # type: ignore[arg-type] def test_concat_tables_none_table(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_tables([None]) + pa.concat_tables([None]) # type: ignore[arg-type] @pytest.mark.pandas @@ -2221,7 +2259,7 @@ def test_concat_batches_different_schema(): def test_concat_batches_none_batches(): # ARROW-11997 with pytest.raises(AttributeError): - pa.concat_batches([None]) + pa.concat_batches([None]) # type: ignore[arg-type] @pytest.mark.parametrize( @@ -2372,7 +2410,7 @@ def test_from_arrays_schema(data, klass): # with different and incompatible schema schema = pa.schema([('strs', pa.utf8()), ('floats', pa.timestamp('s'))]) with pytest.raises((NotImplementedError, TypeError)): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] # Cannot pass both schema and metadata / names with pytest.raises(ValueError): @@ -2477,7 +2515,7 @@ def test_table_from_pydict_arrow_arrays(data, klass): # with different and incompatible schema schema = pa.schema([('strs', pa.utf8()), ('floats', pa.timestamp('s'))]) with pytest.raises((NotImplementedError, TypeError)): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] @pytest.mark.parametrize('data, klass', [ @@ -2494,7 +2532,7 @@ def test_table_from_pydict_schema(data, klass): schema = pa.schema([('strs', pa.utf8()), ('floats', pa.float64()), ('ints', pa.int64())]) with pytest.raises(KeyError, match='ints'): - pa.Table.from_pydict(data, schema=schema) + pa.Table.from_pydict(data, schema=schema) # type: ignore[arg-type] # data has columns not present in schema -> ignored schema = pa.schema([('strs', pa.utf8())]) @@ -2698,10 +2736,10 @@ def test_table_factory_function_args_pandas(): def test_factory_functions_invalid_input(): with pytest.raises(TypeError, match="Expected pandas DataFrame, python"): - pa.table("invalid input") + pa.table("invalid input") # type: ignore[arg-type] with pytest.raises(TypeError, match="Expected pandas DataFrame"): - pa.record_batch("invalid input") + pa.record_batch("invalid input") # type: ignore[arg-type] def test_table_repr_to_string(): @@ -2835,8 +2873,8 @@ def test_table_function_unicode_schema(): schema = pa.schema([(col_a, pa.int32()), (col_b, pa.string())]) result = pa.table(d, schema=schema) - assert result[0].chunk(0).equals(pa.array([1, 2, 3], type='int32')) - assert result[1].chunk(0).equals(pa.array(['a', 'b', 'c'], type='string')) + assert result.column(0).chunk(0).equals(pa.array([1, 2, 3], type='int32')) + assert result.column(1).chunk(0).equals(pa.array(['a', 'b', 'c'], type='string')) def test_table_take_vanilla_functionality(): @@ -3768,7 +3806,7 @@ def test_chunked_array_non_cpu(cuda_context, cpu_chunked_array, cuda_chunked_arr # equals() test with pytest.raises(NotImplementedError): - cuda_chunked_array == cuda_chunked_array + cuda_chunked_array == cuda_chunked_array # type: ignore[reportUnusedExpression] # to_pandas() test with pytest.raises(NotImplementedError): @@ -4025,7 +4063,7 @@ def test_recordbatch_non_cpu(cuda_context, cpu_recordbatch, cuda_recordbatch, # __dataframe__() test with pytest.raises(NotImplementedError): - from_dataframe(cuda_recordbatch.__dataframe__()) + from_dataframe(cuda_recordbatch.__dataframe__()) # type: ignore[misc] def verify_cuda_table(table, expected_schema): @@ -4224,7 +4262,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 8886730413ddb14623d47317804a52f275c47db4 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 10 Mar 2026 01:07:51 +0100 Subject: [PATCH 22/36] 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 bcd7ed14edbd..cc42c81835e7 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 2502292acf89..60bb0e51902e 100644 --- a/python/pyarrow/tests/test_table.py +++ b/python/pyarrow/tests/test_table.py @@ -1209,7 +1209,8 @@ def test_recordbatch_to_tensor_null(): batch.to_tensor() result = batch.to_tensor(null_to_nan=True, row_major=False) - x = np.column_stack([arr1, arr2]).astype(np.float64, order="F") # 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) @@ -1243,7 +1244,8 @@ def test_recordbatch_to_tensor_null(): ) result = batch.to_tensor(null_to_nan=True, row_major=False) - x = np.column_stack([arr1, arr2]).astype(np.float32, order="F") # 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 5ee863c76b9bbf406c7edc23e30672c84aabdc22 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Sun, 14 Jun 2026 10:26:07 -0700 Subject: [PATCH 23/36] 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 4f6c101b36d35c2196d09e00d531cd2cff4d415c Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 11:43:30 +0100 Subject: [PATCH 24/36] 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 9a1185c3441c77f77076100cf026a71a22fa8cb1 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 12:17:17 +0100 Subject: [PATCH 25/36] 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 168c7c4ba40f397860aedf72a34d8a5a618d5825 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 15:46:29 +0100 Subject: [PATCH 26/36] 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 c347f7f2b456756528fe874d2dae724d4a2619ed Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 16:28:48 +0100 Subject: [PATCH 27/36] 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 0e55c813fa559e9905b6057347dd7c24b99493b0 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 16:40:04 +0100 Subject: [PATCH 28/36] 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 b096b3be6ded54a923500ec07def5f587f74f8fc Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:41 +0100 Subject: [PATCH 29/36] 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 147e02145ea6ea252fdd877171442274c87693c0 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:45 +0100 Subject: [PATCH 30/36] 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 67aeb3364debbda44ad6dc9b469caa7da4f06b50 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:48 +0100 Subject: [PATCH 31/36] 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 440fbeee2d713a54769011163d5bdfd88f0a5fa0 Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:51 +0100 Subject: [PATCH 32/36] 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 a72e35be411bbca7381828973559ef986878474f Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Mon, 26 Oct 2026 18:40:55 +0100 Subject: [PATCH 33/36] 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 120659e7215485262acbf598e319f189028da8bd Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 27 Oct 2026 14:02:32 +0100 Subject: [PATCH 34/36] 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 f93d912678ccfbff6af336d38ec52c744c160184 Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 27 Oct 2026 15:39:59 +0100 Subject: [PATCH 35/36] Fix Decima32Array typo in array.pxi --- python/pyarrow/array.pxi | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/pyarrow/array.pxi b/python/pyarrow/array.pxi index fa3376fc8daf..3b5dce1eda3a 100644 --- a/python/pyarrow/array.pxi +++ b/python/pyarrow/array.pxi @@ -2787,7 +2787,7 @@ cdef class FixedSizeBinaryArray(Array): Concrete class for Arrow arrays of a fixed-size binary data type. """ -cdef class Decima32Array(FixedSizeBinaryArray): +cdef class Decimal32Array(FixedSizeBinaryArray): """ Concrete class for Arrow arrays of decimal32 data type. """ From 1add31bc159de8ebde4dd1c77a7f2e689072e0bb Mon Sep 17 00:00:00 2001 From: AlenkaF Date: Tue, 27 Oct 2026 15:58:14 +0100 Subject: [PATCH 36/36] 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: ... -