fixes

Signed-off-by: Robert Kruszewski <github@robertk.io>

Author: robert3005

AGENTS.md

@@ -72,6 +72,22 @@ Makefile target is used.
 
 If you touch documentation run doc tests via `cargo test --doc`.
 
+For Python binding changes under `vortex-python/`, run the narrow Python checks that match the
+files touched before broader test suites. Useful checks include:
+
+```bash
+python -m py_compile <changed-python-files>
+uv run --all-packages --reinstall-package vortex-data pytest <changed-python-tests>
+```
+
+If Python docstrings, `docs/api/python/`, or Sphinx configuration change, also run the docs checks
+from a clean Sphinx environment:
+
+```bash
+uv run --all-packages make -C docs clean html
+uv run --all-packages make -C docs clean doctest
+```
+
 ## Linting, Formatting, and Generated Files
 
 Run verification that matches the files changed. Do not run expensive Rust checks for changes that
@@ -80,6 +96,16 @@ with no Rust/API behavior impact. For docs/config-only changes, validate formatt
 or with a targeted doc/config command, and verify symlink or path changes with `ls`, `find`, and
 `git status`.
 
+For Python binding changes under `vortex-python/`, run the relevant Python lint and type checks:
+
+```bash
+uv run basedpyright vortex-python
+uv run ruff check <changed-python-files>
+```
+
+If PyO3 Rust files in `vortex-python/src/` change, include `cargo +nightly fmt --check -p
+vortex-python`. Always finish Python binding work with `git diff --check`.
+
 For Rust code, public API, feature flag, or generated-file changes, run these before stopping:
 
 ```bash

docs/api/python/index.rst

@@ -62,6 +62,7 @@ API Reference
    arrays
    expr
    compress
+   session
    io
    store
    dataset

docs/api/python/session.rst

@@ -0,0 +1,5 @@
+Session
+=======
+
+.. autoclass:: vortex.Session
+   :members:

vortex-python/python/vortex/__init__.py

@@ -79,7 +79,6 @@
 from .arrays import (
     Array,
     PyArray,
-    _unpickle_array,  # pyright: ignore[reportPrivateUsage]
     array,
 )
 from .file import VortexFile, open
@@ -193,7 +192,6 @@
     "dumps",
     "load",
     "loads",
-    "_unpickle_array",
     # File
     "VortexFile",
     "open",

vortex-python/python/vortex/_lib/runtime.pyi

@@ -0,0 +1,6 @@
+#  SPDX-License-Identifier: Apache-2.0
+#  SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+def set_worker_threads(n: int) -> None: ...
+def set_worker_threads_to_available_parallelism() -> None: ...
+def worker_count() -> int: ...

vortex-python/python/vortex/arrays.py

@@ -3,19 +3,15 @@
 from __future__ import annotations
 
 import abc
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 import pyarrow
 from typing_extensions import override
 
 import vortex._lib.arrays as _arrays  # pyright: ignore[reportMissingModuleSource]
 from vortex._lib.dtype import DType  # pyright: ignore[reportMissingModuleSource]
-from vortex._lib.serde import (  # pyright: ignore[reportMissingModuleSource]
-    ArrayContext,
-    SerializedArray,
-    decode_ipc_array_buffers,
-)
+from vortex._lib.serde import ArrayContext, SerializedArray  # pyright: ignore[reportMissingModuleSource]
 from vortex._lib.session import Session  # pyright: ignore[reportMissingModuleSource]
 
 try:
@@ -507,20 +503,3 @@ def decode(cls, parts: SerializedArray, ctx: ArrayContext, dtype: DType, len: in
         that represent the current array. Implementations of this function should validate this
         information, and then construct a new array.
         """
-
-
-def _unpickle_array(
-    array_buffers: Sequence[bytes | memoryview],
-    dtype_buffers: Sequence[bytes | memoryview],
-    *,
-    session: Session,
-) -> Array:  # pyright: ignore[reportUnusedFunction]
-    """Unpickle a Vortex array from IPC-encoded buffer lists.
-
-    This is an internal function used by the pickle module for both protocol 4 and 5.
-
-    For protocol 4, receives list[bytes] from __reduce__.
-    For protocol 5, receives list[PickleBuffer/memoryview] from __reduce_ex__.
-    Both use decode_ipc_array_buffers which concatenates the buffers during deserialization.
-    """
-    return decode_ipc_array_buffers(array_buffers, dtype_buffers, session=session)

vortex-python/python/vortex/pickle.py

@@ -6,7 +6,10 @@
 import io
 import pickle as _pickle
 from ast import literal_eval
-from typing import Any
+from collections.abc import Callable, Iterable, Sequence
+from typing import BinaryIO, TypeAlias, TypeGuard, cast
+
+from typing_extensions import override
 
 from vortex._lib.arrays import Array  # pyright: ignore[reportMissingModuleSource]
 from vortex._lib.serde import (  # pyright: ignore[reportMissingModuleSource]
@@ -18,27 +21,61 @@
 _ARRAY_PERSISTENT_ID = "vortex.array"
 _ARRAY_PERSISTENT_ID_VERSION = 1
 
+_BufferSequence: TypeAlias = Sequence[bytes | memoryview]
+_ArrayPersistentId: TypeAlias = tuple[str, int, _BufferSequence, _BufferSequence]
+_BufferCallback: TypeAlias = Callable[[_pickle.PickleBuffer], object | None]
+_OutOfBandBuffers: TypeAlias = Iterable[bytes | bytearray | memoryview | _pickle.PickleBuffer]
+
+
+def _is_buffer_sequence(obj: object) -> TypeGuard[_BufferSequence]:
+    return isinstance(obj, Sequence) and all(isinstance(buffer, bytes | memoryview) for buffer in obj)
+
+
+def _parse_array_persistent_id(pid: object) -> _ArrayPersistentId:
+    parsed_pid: object = pid
+    if isinstance(parsed_pid, str):
+        try:
+            parsed_pid = cast(object, literal_eval(parsed_pid))
+        except (SyntaxError, ValueError) as err:
+            raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}") from err
+
+    if not isinstance(parsed_pid, tuple):
+        raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
+
+    parsed_tuple = cast(tuple[object, ...], parsed_pid)
+    if len(parsed_tuple) != 4:
+        raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
+
+    tag, version, array_buffers, dtype_buffers = parsed_tuple
+    if tag != _ARRAY_PERSISTENT_ID or version != _ARRAY_PERSISTENT_ID_VERSION:
+        raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
+    if not _is_buffer_sequence(array_buffers) or not _is_buffer_sequence(dtype_buffers):
+        raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
+
+    return (_ARRAY_PERSISTENT_ID, _ARRAY_PERSISTENT_ID_VERSION, array_buffers, dtype_buffers)
+
 
 class Pickler(_pickle.Pickler):
     """Pickler that serializes Vortex arrays using an explicit session."""
 
     def __init__(
         self,
-        file: Any,  # pyright: ignore[reportExplicitAny]
+        file: BinaryIO,
         *,
         session: Session,
         protocol: int | None = None,
         fix_imports: bool = True,
-        buffer_callback: Any | None = None,  # pyright: ignore[reportExplicitAny]
+        buffer_callback: _BufferCallback | None = None,
     ) -> None:
         super().__init__(
             file,
             protocol=protocol,
             fix_imports=fix_imports,
             buffer_callback=buffer_callback,
         )
-        self._session = session
+        self._session: Session = session
 
+    @override
     def persistent_id(self, obj: object) -> object | None:
         if isinstance(obj, Array):
             array_buffers, dtype_buffers = encode_ipc_array_buffers(obj, session=self._session)
@@ -51,13 +88,13 @@ class Unpickler(_pickle.Unpickler):
 
     def __init__(
         self,
-        file: Any,  # pyright: ignore[reportExplicitAny]
+        file: BinaryIO,
         *,
         session: Session,
         fix_imports: bool = True,
         encoding: str = "ASCII",
         errors: str = "strict",
-        buffers: Any | None = None,  # pyright: ignore[reportExplicitAny]
+        buffers: _OutOfBandBuffers | None = None,
     ) -> None:
         super().__init__(
             file,
@@ -66,33 +103,22 @@ def __init__(
             errors=errors,
             buffers=buffers,
         )
-        self._session = session
+        self._session: Session = session
 
+    @override
     def persistent_load(self, pid: object) -> object:
-        if isinstance(pid, str):
-            try:
-                pid = literal_eval(pid)
-            except (SyntaxError, ValueError) as err:
-                raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}") from err
-
-        if not isinstance(pid, tuple) or len(pid) != 4:
-            raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
-
-        tag, version, array_buffers, dtype_buffers = pid
-        if tag != _ARRAY_PERSISTENT_ID or version != _ARRAY_PERSISTENT_ID_VERSION:
-            raise _pickle.UnpicklingError(f"unsupported persistent id: {pid!r}")
-
+        _, _, array_buffers, dtype_buffers = _parse_array_persistent_id(pid)
         return decode_ipc_array_buffers(array_buffers, dtype_buffers, session=self._session)
 
 
 def dump(
     obj: object,
-    file: Any,  # pyright: ignore[reportExplicitAny]
+    file: BinaryIO,
     *,
     session: Session,
     protocol: int | None = None,
     fix_imports: bool = True,
-    buffer_callback: Any | None = None,  # pyright: ignore[reportExplicitAny]
+    buffer_callback: _BufferCallback | None = None,
 ) -> None:
     Pickler(
         file,
@@ -109,7 +135,7 @@ def dumps(
     session: Session,
     protocol: int | None = None,
     fix_imports: bool = True,
-    buffer_callback: Any | None = None,  # pyright: ignore[reportExplicitAny]
+    buffer_callback: _BufferCallback | None = None,
 ) -> bytes:
     file = io.BytesIO()
     dump(
@@ -124,22 +150,25 @@ def dumps(
 
 
 def load(
-    file: Any,  # pyright: ignore[reportExplicitAny]
+    file: BinaryIO,
     *,
     session: Session,
     fix_imports: bool = True,
     encoding: str = "ASCII",
     errors: str = "strict",
-    buffers: Any | None = None,  # pyright: ignore[reportExplicitAny]
+    buffers: _OutOfBandBuffers | None = None,
 ) -> object:
-    return Unpickler(
-        file,
-        session=session,
-        fix_imports=fix_imports,
-        encoding=encoding,
-        errors=errors,
-        buffers=buffers,
-    ).load()
+    return cast(
+        object,
+        Unpickler(
+            file,
+            session=session,
+            fix_imports=fix_imports,
+            encoding=encoding,
+            errors=errors,
+            buffers=buffers,
+        ).load(),
+    )
 
 
 def loads(
@@ -149,7 +178,7 @@ def loads(
     fix_imports: bool = True,
     encoding: str = "ASCII",
     errors: str = "strict",
-    buffers: Any | None = None,  # pyright: ignore[reportExplicitAny]
+    buffers: _OutOfBandBuffers | None = None,
 ) -> object:
     return load(
         io.BytesIO(data),

vortex-python/src/session.rs

@@ -24,6 +24,10 @@ pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
     Ok(())
 }
 
+/// A Vortex execution session.
+///
+/// Sessions hold runtime state and registries used by Vortex operations. Create one session and
+/// pass it to APIs that read, write, convert, compress, or otherwise execute arrays.
 #[pyclass(name = "Session", module = "vortex", frozen)]
 pub struct PyVortexSession {
     inner: VortexSession,