diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index df22097e5b9b..dd9a04c32467 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -230,6 +230,10 @@ jobs: - name: Run Rust tests run: | source venv/bin/activate + # arrow-pyarrow's own tests are feature-gated, so the workspace-wide `cargo test` in + # rust.yml (default features) does not build them. They need no interpreter at runtime, + # but pyo3 needs one to link against, hence the active venv. + cargo test -p arrow-pyarrow --all-features cd arrow-pyarrow-testing cargo test - name: Run Python tests diff --git a/Cargo.lock b/Cargo.lock index ee741323d02f..40077980e60a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2756,6 +2756,7 @@ dependencies = [ "portable-atomic", "pyo3-build-config", "pyo3-ffi", + "pyo3-macros", ] [[package]] @@ -2777,6 +2778,30 @@ dependencies = [ "pyo3-build-config", ] +[[package]] +name = "pyo3-macros" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91f9d455db760a9a0b0ddeaac25f1390b8a36ba73dfbda9f127cac6fc340d4d5" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e343bcec300ff262f5806a33a4e51b6d097a8a46435f512fcb83e95592581625" +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 0051de87436e..f7d6797c3f24 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..01d482e76de7 100644 --- a/arrow-pyarrow/src/lib.rs +++ b/arrow-pyarrow/src/lib.rs @@ -58,6 +58,16 @@ //! For example, a `pyarrow.Table` (or any other object that implements the ArrayStream PyCapsule //! interface) can be imported to Rust through `PyArrowType` instead of //! forcing eager reading into `Vec`. +//! +//! # 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 classes only, and are therefore narrower than what is accepted: the +//! PyCapsule interface is duck-typed and has no canonical Python type to name. use std::convert::{From, TryFrom}; use std::ffi::CStr; @@ -79,6 +89,20 @@ 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::{ + inspect::PyStaticExpr, type_hint_identifier, type_hint_subscript, type_hint_union, + type_object::PyTypeInfo, +}; + +/// Declares a `FromPyArrow::INPUT_TYPE` / `ToPyArrow::OUTPUT_TYPE` / `IntoPyArrow::OUTPUT_TYPE` +/// hint on an impl, expanding to nothing unless the `experimental-inspect` feature is enabled. +macro_rules! type_hint { + ($name:ident = $hint:expr) => { + #[cfg(feature = "experimental-inspect")] + const $name: PyStaticExpr = $hint; + }; +} import_exception!(pyarrow, ArrowException); /// Represents an exception raised by PyArrow. @@ -88,8 +112,35 @@ fn to_py_err(err: ArrowError) -> PyErr { PyArrowException::new_err(err.to_string()) } +/// The type hint shared by every conversion that imports through the ArrowArrayStream PyCapsule +/// interface, i.e. [`ArrowArrayStreamReader`] and [`Table`]. +/// +/// Both go through the same `__arrow_c_stream__` path and therefore accept exactly the same +/// objects, so naming only one of the two classes would make a stub generator reject usage this +/// crate's own documentation recommends — importing a `pyarrow.Table` as a +/// `PyArrowType`. +#[cfg(feature = "experimental-inspect")] +const ARRAY_STREAM_INPUT_TYPE: PyStaticExpr = type_hint_union!( + type_hint_identifier!("pyarrow", "RecordBatchReader"), + type_hint_identifier!("pyarrow", "Table") +); + /// 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 pyarrow classes 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 +149,32 @@ 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 { + type_hint!(OUTPUT_TYPE = ::OUTPUT_TYPE); + fn into_pyarrow<'py>(self, py: Python<'py>) -> PyResult> { self.to_pyarrow(py) } @@ -126,6 +192,8 @@ fn validate_class(expected: &Bound, value: &Bound) -> PyResult<() } impl FromPyArrow for DataType { + type_hint!(INPUT_TYPE = 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 +221,8 @@ impl FromPyArrow for DataType { } impl ToPyArrow for DataType { + type_hint!(OUTPUT_TYPE = 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 +233,8 @@ impl ToPyArrow for DataType { } impl FromPyArrow for Field { + type_hint!(INPUT_TYPE = 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 +262,8 @@ impl FromPyArrow for Field { } impl ToPyArrow for Field { + type_hint!(OUTPUT_TYPE = 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 +274,8 @@ impl ToPyArrow for Field { } impl FromPyArrow for Schema { + type_hint!(INPUT_TYPE = 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 +303,8 @@ impl FromPyArrow for Schema { } impl ToPyArrow for Schema { + type_hint!(OUTPUT_TYPE = 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 +315,8 @@ impl ToPyArrow for Schema { } impl FromPyArrow for ArrayData { + type_hint!(INPUT_TYPE = 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 +351,8 @@ impl FromPyArrow for ArrayData { } impl ToPyArrow for ArrayData { + type_hint!(OUTPUT_TYPE = 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 +367,13 @@ impl ToPyArrow for ArrayData { } impl FromPyArrow for Vec { + type_hint!( + INPUT_TYPE = 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 +384,10 @@ impl FromPyArrow for Vec { } impl ToPyArrow for Vec { + type_hint!( + OUTPUT_TYPE = type_hint_subscript!(PyList::TYPE_HINT, ::OUTPUT_TYPE) + ); + fn to_pyarrow<'py>(&self, py: Python<'py>) -> PyResult> { self.iter() .map(|v| v.to_pyarrow(py)) @@ -304,6 +397,8 @@ impl ToPyArrow for Vec { } impl FromPyArrow for RecordBatch { + type_hint!(INPUT_TYPE = 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 +457,8 @@ impl FromPyArrow for RecordBatch { } impl ToPyArrow for RecordBatch { + type_hint!(OUTPUT_TYPE = 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 +470,8 @@ impl ToPyArrow for RecordBatch { /// Supports conversion from `pyarrow.RecordBatchReader` to [ArrowArrayStreamReader]. impl FromPyArrow for ArrowArrayStreamReader { + type_hint!(INPUT_TYPE = ARRAY_STREAM_INPUT_TYPE); + 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 +509,8 @@ impl FromPyArrow for ArrowArrayStreamReader { /// Convert a [`RecordBatchReader`] into a `pyarrow.RecordBatchReader`. impl IntoPyArrow for Box { + type_hint!(OUTPUT_TYPE = 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 +524,8 @@ impl IntoPyArrow for Box { /// Convert a [`ArrowArrayStreamReader`] into a `pyarrow.RecordBatchReader`. impl IntoPyArrow for ArrowArrayStreamReader { + type_hint!(OUTPUT_TYPE = 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 +601,8 @@ impl TryFrom> for Table { /// Convert a `pyarrow.Table` (or any other ArrowArrayStream compliant object) into [`Table`] impl FromPyArrow for Table { + type_hint!(INPUT_TYPE = ARRAY_STREAM_INPUT_TYPE); + fn from_pyarrow_bound(ob: &Bound) -> PyResult { let reader: Box = Box::new(ArrowArrayStreamReader::from_pyarrow_bound(ob)?); @@ -507,6 +612,8 @@ impl FromPyArrow for Table { /// Convert a [`Table`] into `pyarrow.Table`. impl IntoPyArrow for Table { + type_hint!(OUTPUT_TYPE = 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 +671,8 @@ pub struct PyArrowType(pub T); impl FromPyObject<'_, '_> for PyArrowType { type Error = PyErr; + type_hint!(INPUT_TYPE = ::INPUT_TYPE); + fn extract(value: Borrowed<'_, '_, PyAny>) -> PyResult { Ok(Self(T::from_pyarrow_bound(&value)?)) } @@ -576,6 +685,8 @@ impl<'py, T: IntoPyArrow> IntoPyObject<'py> for PyArrowType { type Error = PyErr; + type_hint!(OUTPUT_TYPE = ::OUTPUT_TYPE); + fn into_pyobject(self, py: Python<'py>) -> PyResult { self.0.into_pyarrow(py) } @@ -645,3 +756,78 @@ 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::>(), + "builtins.list[pyarrow.RecordBatch]" + ); + } + + /// Outputs are exact, but both stream imports accept either class, because both go through + /// `__arrow_c_stream__`. + #[test] + fn readers_and_tables_map_to_their_pyarrow_class() { + assert_eq!( + input_type::(), + "pyarrow.RecordBatchReader | pyarrow.Table" + ); + assert_eq!( + output_type::(), + "pyarrow.RecordBatchReader" + ); + assert_eq!( + output_type::>(), + "pyarrow.RecordBatchReader" + ); + assert_eq!( + input_type::(), + "pyarrow.RecordBatchReader | pyarrow.Table" + ); + assert_eq!(output_type::
(), "pyarrow.Table"); + } +} diff --git a/arrow/Cargo.toml b/arrow/Cargo.toml index 9dbc59fe5d9b..c4bab6fc3f75 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`. +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 diff --git a/arrow/README.md b/arrow/README.md index fb5f6b9ab2bd..02c7a0cb72d6 100644 --- a/arrow/README.md +++ b/arrow/README.md @@ -64,6 +64,7 @@ The `arrow` crate provides the following features which may be enabled in your ` - `chrono-tz` - support of parsing timezone using [chrono-tz](https://docs.rs/chrono-tz/0.6.0/chrono_tz/) - `ffi` - bindings for the Arrow C [C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) - `pyarrow` - bindings for pyo3 to call arrow-rs from python +- `pyarrow-experimental-inspect` - record the pyarrow type of each conversion in PyO3's introspection data, so that stub generators emit e.g. `pyarrow.Array` rather than `_typeshed.Incomplete` (also enables `pyarrow`) - `canonical_extension_types` - definitions for [canonical extension types](https://arrow.apache.org/docs/format/CanonicalExtensions.html#format-canonical-extensions) - `async` - definitions for traits using `async`, intended to work with the async ecosystem