From dc928b13ae7b29447ee48bf7e982c21c48321bf6 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Thu, 30 Jul 2026 17:07:56 +0200 Subject: [PATCH] feat(pyarrow): describe conversions in PyO3 introspection data PyO3's `experimental-inspect` feature records the Python type of every conversion so that tools like `maturin generate-stubs` can write a real `.pyi`. `PyArrowType` implements `FromPyObject` and `IntoPyObject` but leaves their `INPUT_TYPE`/`OUTPUT_TYPE` constants at the default, `_typeshed.Incomplete`. Since `PyArrowType` is how bindings exchange every Arrow value with Python, that default erases their whole public API. A binding whose functions take and return Arrow data generates stubs in which every signature is `Incomplete`, which is `Any`, so the stubs verify nothing. Add the hints to `FromPyArrow`, `ToPyArrow` and `IntoPyArrow` and forward them through `PyArrowType`, behind a new `experimental-inspect` feature that just turns on PyO3's. Nothing is compiled unless it is enabled, and the default stays `Incomplete` so an out-of-tree implementor of these traits is unaffected. Input and output hints are separate because they genuinely differ: `Vec` is built from anything iterable but handed back as a `list`. Input hints name the pyarrow class only. Every conversion here also accepts any object implementing the relevant PyCapsule interface method, but that is duck-typed and neither pyarrow nor typeshed defines a protocol to point at, so naming one would emit an import that does not resolve. This is documented on the trait; a binding that wants to advertise the wider protocol can declare its own `Protocol` and carry it on a newtype. `arrow` gets a matching feature that implies `pyarrow`. --- Cargo.lock | 25 ++++++ arrow-pyarrow/Cargo.toml | 5 ++ arrow-pyarrow/src/lib.rs | 177 +++++++++++++++++++++++++++++++++++++++ arrow/Cargo.toml | 3 + 4 files changed, 210 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index b53a7736fd27..1fbacedbb8cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2677,6 +2677,7 @@ dependencies = [ "portable-atomic", "pyo3-build-config", "pyo3-ffi", + "pyo3-macros", ] [[package]] @@ -2698,6 +2699,30 @@ dependencies = [ "pyo3-build-config", ] +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "quad-rand" version = "0.2.3" diff --git a/arrow-pyarrow/Cargo.toml b/arrow-pyarrow/Cargo.toml index 9f21a9dcb034..552c86088212 100644 --- a/arrow-pyarrow/Cargo.toml +++ b/arrow-pyarrow/Cargo.toml @@ -35,6 +35,11 @@ bench = false [package.metadata.docs.rs] all-features = true +[features] +# Emit the Python type of each conversion into PyO3's introspection data, so that tools like +# `maturin generate-stubs` can put a real type into a generated `.pyi` instead of `Incomplete`. +experimental-inspect = ["pyo3/experimental-inspect"] + [dependencies] arrow-array = { workspace = true, features = ["ffi"] } arrow-data = { workspace = true } diff --git a/arrow-pyarrow/src/lib.rs b/arrow-pyarrow/src/lib.rs index cfe51b807aef..312dbe0fbafc 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -51,6 +51,16 @@ //! as `pyarrow.RecordBatchReader`. (`Box` is typically //! easier to create.) //! +//! # Type stubs +//! +//! With the `experimental-inspect` feature enabled, each conversion records the pyarrow class it +//! maps to in PyO3's introspection data, so a generated `.pyi` says `pyarrow.Array` where it would +//! otherwise say `_typeshed.Incomplete`. The hints live on [`FromPyArrow::INPUT_TYPE`], +//! [`ToPyArrow::OUTPUT_TYPE`] and [`IntoPyArrow::OUTPUT_TYPE`], and [`PyArrowType`] forwards them. +//! +//! Input hints name the pyarrow class only, and are therefore narrower than what is accepted: the +//! PyCapsule interface is duck-typed and has no canonical Python type to name. +//! //! (2) Although arrow-rs offers [Table], a convenience wrapper for [pyarrow.Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table) //! that internally holds `Vec`, it is meant primarily for use cases where you already //! have `Vec` on the Rust side and want to export that in bulk as a `pyarrow.Table`. @@ -75,10 +85,14 @@ use arrow_data::ArrayData; use arrow_schema::{ArrowError, DataType, Field, Schema, SchemaRef}; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::ffi::Py_uintptr_t; +#[cfg(feature = "experimental-inspect")] +use pyo3::inspect::PyStaticExpr; use pyo3::prelude::*; use pyo3::sync::PyOnceLock; use pyo3::types::{PyCapsule, PyDict, PyList, PyString, PyType}; use pyo3::{CastError, import_exception, intern}; +#[cfg(feature = "experimental-inspect")] +use pyo3::{type_hint_identifier, type_hint_subscript}; import_exception!(pyarrow, ArrowException); /// Represents an exception raised by PyArrow. @@ -90,6 +104,20 @@ fn to_py_err(err: ArrowError) -> PyErr { /// Trait for converting Python objects to arrow-rs types. pub trait FromPyArrow: Sized { + /// The Python type this conversion accepts, as a type hint. + /// + /// Used by [`FromPyObject::INPUT_TYPE`] on [`PyArrowType`] so that a stub generator can write + /// `pyarrow.Array` where it would otherwise write `_typeshed.Incomplete`. + /// + /// This names the pyarrow class only. Every conversion here *also* accepts any object + /// implementing the relevant [PyCapsule interface](https://arrow.apache.org/docs/format/CDataInterface/PyCapsuleInterface.html) + /// method, which is duck-typed and has no canonical Python type to point at — neither pyarrow + /// nor typeshed defines one. The hint is therefore narrower than what is accepted at runtime. + /// A binding that wants to advertise the wider protocol can declare its own `Protocol` and + /// carry it on a newtype around the arrow-rs type. + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("_typeshed", "Incomplete"); + /// Convert a Python object to an arrow-rs type. /// /// Takes a GIL-bound value from Python and returns a result with the arrow-rs type. @@ -98,17 +126,33 @@ pub trait FromPyArrow: Sized { /// Create a new PyArrow object from a arrow-rs type. pub trait ToPyArrow { + /// The Python type this conversion produces, as a type hint. + /// + /// Unlike [`FromPyArrow::INPUT_TYPE`] this is exact: the conversion always constructs an + /// instance of the named pyarrow class. + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("_typeshed", "Incomplete"); + /// Convert the implemented type into a Python object without consuming it. fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult>; } /// Convert an arrow-rs type into a PyArrow object. pub trait IntoPyArrow { + /// The Python type this conversion produces, as a type hint. + /// + /// See [`ToPyArrow::OUTPUT_TYPE`]. + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("_typeshed", "Incomplete"); + /// Convert the implemented type into a Python object while consuming it. fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult>; } impl IntoPyArrow for T { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = ::OUTPUT_TYPE; + fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult> { self.to_pyarrow(py) } @@ -126,6 +170,9 @@ fn validate_class(expected: &Bound, value: &Bound) -> PyResult<() } impl FromPyArrow for DataType { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "DataType"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -153,6 +200,9 @@ impl FromPyArrow for DataType { } impl ToPyArrow for DataType { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "DataType"); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?; data_type_class(py)?.call_method1( @@ -163,6 +213,9 @@ impl ToPyArrow for DataType { } impl FromPyArrow for Field { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Field"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -190,6 +243,9 @@ impl FromPyArrow for Field { } impl ToPyArrow for Field { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Field"); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?; field_class(py)?.call_method1( @@ -200,6 +256,9 @@ impl ToPyArrow for Field { } impl FromPyArrow for Schema { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Schema"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -227,6 +286,9 @@ impl FromPyArrow for Schema { } impl ToPyArrow for Schema { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Schema"); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { let c_schema = FFI_ArrowSchema::try_from(self).map_err(to_py_err)?; schema_class(py)?.call_method1( @@ -237,6 +299,9 @@ impl ToPyArrow for Schema { } impl FromPyArrow for ArrayData { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Array"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -271,6 +336,9 @@ impl FromPyArrow for ArrayData { } impl ToPyArrow for ArrayData { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Array"); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { let array = FFI_ArrowArray::new(self); let schema = FFI_ArrowSchema::try_from(self.data_type()).map_err(to_py_err)?; @@ -285,6 +353,12 @@ impl ToPyArrow for ArrayData { } impl FromPyArrow for Vec { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_subscript!( + type_hint_identifier!("collections.abc", "Iterable"), + ::INPUT_TYPE + ); + fn from_pyarrow_bound(value: &Bound) -> PyResult { let mut v = Vec::with_capacity(value.len().unwrap_or(0)); for item in value.try_iter()? { @@ -295,6 +369,12 @@ impl FromPyArrow for Vec { } impl ToPyArrow for Vec { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_subscript!( + type_hint_identifier!("builtins", "list"), + ::OUTPUT_TYPE + ); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { self.iter() .map(|v| v.to_pyarrow(py)) @@ -304,6 +384,9 @@ impl ToPyArrow for Vec { } impl FromPyArrow for RecordBatch { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "RecordBatch"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -362,6 +445,9 @@ impl FromPyArrow for RecordBatch { } impl ToPyArrow for RecordBatch { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "RecordBatch"); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { // Workaround apache/arrow#37669 by returning RecordBatchIterator let reader = RecordBatchIterator::new(vec![Ok(self.clone())], self.schema()); @@ -373,6 +459,9 @@ impl ToPyArrow for RecordBatch { /// Supports conversion from `pyarrow.RecordBatchReader` to [ArrowArrayStreamReader]. impl FromPyArrow for ArrowArrayStreamReader { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "RecordBatchReader"); + fn from_pyarrow_bound(value: &Bound) -> PyResult { // Newer versions of PyArrow as well as other libraries with Arrow data implement this // method, so prefer it over _export_to_c. @@ -410,6 +499,9 @@ impl FromPyArrow for ArrowArrayStreamReader { /// Convert a [`RecordBatchReader`] into a `pyarrow.RecordBatchReader`. impl IntoPyArrow for Box { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "RecordBatchReader"); + // We can't implement `ToPyArrow` for `T: RecordBatchReader + Send` because // there is already a blanket implementation for `T: ToPyArrow`. fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult> { @@ -423,6 +515,9 @@ impl IntoPyArrow for Box { /// Convert a [`ArrowArrayStreamReader`] into a `pyarrow.RecordBatchReader`. impl IntoPyArrow for ArrowArrayStreamReader { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "RecordBatchReader"); + fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult> { let boxed: Box = Box::new(self); boxed.into_pyarrow(py) @@ -498,6 +593,9 @@ impl TryFrom> for Table { /// Convert a `pyarrow.Table` (or any other ArrowArrayStream compliant object) into [`Table`] impl FromPyArrow for Table { + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Table"); + fn from_pyarrow_bound(ob: &Bound) -> PyResult { let reader: Box = Box::new(ArrowArrayStreamReader::from_pyarrow_bound(ob)?); @@ -507,6 +605,9 @@ impl FromPyArrow for Table { /// Convert a [`Table`] into `pyarrow.Table`. impl IntoPyArrow for Table { + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("pyarrow", "Table"); + fn into_pyarrow(self, py: Python) -> PyResult> { let py_batches = PyList::new(py, self.record_batches.into_iter().map(PyArrowType))?; let py_schema = PyArrowType(Arc::unwrap_or_clone(self.schema)); @@ -564,6 +665,9 @@ pub struct PyArrowType(pub T); impl FromPyObject<'_, '_> for PyArrowType { type Error = PyErr; + #[cfg(feature = "experimental-inspect")] + const INPUT_TYPE: PyStaticExpr = ::INPUT_TYPE; + fn extract(value: Borrowed<'_, '_, PyAny>) -> PyResult { Ok(Self(T::from_pyarrow_bound(&value)?)) } @@ -576,6 +680,9 @@ impl<'py, T: IntoPyArrow> IntoPyObject<'py> for PyArrowType { type Error = PyErr; + #[cfg(feature = "experimental-inspect")] + const OUTPUT_TYPE: PyStaticExpr = ::OUTPUT_TYPE; + fn into_pyobject(self, py: Python<'py>) -> PyResult { self.0.into_pyarrow(py) } @@ -645,3 +752,73 @@ fn wrapping_type_error(py: Python<'_>, error: PyErr, message: String) -> PyErr { e.set_cause(py, Some(error)); e } + +#[cfg(all(test, feature = "experimental-inspect"))] +mod introspection_tests { + use super::*; + use pyo3::{FromPyObject, IntoPyObject}; + + /// The type hint a `PyArrowType` argument is described by. + fn input_type() -> String { + as FromPyObject<'_, '_>>::INPUT_TYPE.to_string() + } + + /// The type hint a `PyArrowType` return value is described by. + fn output_type() -> String + where + PyArrowType: for<'py> IntoPyObject<'py>, + { + as IntoPyObject<'_>>::OUTPUT_TYPE.to_string() + } + + #[test] + fn scalar_types_map_to_their_pyarrow_class() { + assert_eq!(input_type::(), "pyarrow.DataType"); + assert_eq!(output_type::(), "pyarrow.DataType"); + assert_eq!(input_type::(), "pyarrow.Field"); + assert_eq!(output_type::(), "pyarrow.Field"); + assert_eq!(input_type::(), "pyarrow.Schema"); + assert_eq!(output_type::(), "pyarrow.Schema"); + assert_eq!(input_type::(), "pyarrow.RecordBatch"); + assert_eq!(output_type::(), "pyarrow.RecordBatch"); + } + + /// `ArrayData` is the one case where the arrow-rs name and the pyarrow name differ. + #[test] + fn array_data_maps_to_pyarrow_array() { + assert_eq!(input_type::(), "pyarrow.Array"); + assert_eq!(output_type::(), "pyarrow.Array"); + } + + /// Asymmetric on purpose: `Vec` is built from anything iterable, but is handed back as a + /// list. + #[test] + fn vec_is_iterable_in_and_list_out() { + assert_eq!( + input_type::>(), + "collections.abc.Iterable[pyarrow.RecordBatch]" + ); + assert_eq!( + output_type::>(), + "list[pyarrow.RecordBatch]" + ); + } + + #[test] + fn readers_and_tables_map_to_their_pyarrow_class() { + assert_eq!( + input_type::(), + "pyarrow.RecordBatchReader" + ); + assert_eq!( + output_type::(), + "pyarrow.RecordBatchReader" + ); + assert_eq!( + output_type::>(), + "pyarrow.RecordBatchReader" + ); + assert_eq!(input_type::(), "pyarrow.Table"); + assert_eq!(output_type::
(), "pyarrow.Table"); + } +} diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index 2be4cfd1f104..373b5d45a5f4 100644 --- a/arrow/Cargo.toml +++ b/arrow/Cargo.toml @@ -74,6 +74,9 @@ prettyprint = ["arrow-cast/prettyprint"] # target without assuming an environment containing JavaScript. test_utils = ["dep:rand", "dep:half"] pyarrow = ["ffi", "dep:arrow-pyarrow"] +# Record the Python type of each pyarrow conversion in PyO3's introspection data, so that stub +# generators emit e.g. `pyarrow.Array` rather than `_typeshed.Incomplete`. Implies `pyarrow`. +experimental-inspect = ["pyarrow", "arrow-pyarrow/experimental-inspect"] # force_validate runs full data validation for all arrays that are created # this is not enabled by default as it is too computationally expensive # but is run as part of our CI checks