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.

Author: rok

python/pyarrow-stubs/pyarrow/array.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"]

python/pyarrow-stubs/pyarrow/builder.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",
+]

python/pyarrow-stubs/pyarrow/config.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"]

python/pyarrow-stubs/pyarrow/device.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",
+]