Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Improved rust docs and method visibility #119

Merged
merged 1 commit into from
Aug 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions arro3-core/python/arro3/core/_core.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ from .types import (

class Array:
"""An Arrow Array."""

def __init__(self, obj: Sequence[Any], /, type: ArrowSchemaExportable) -> None:
"""Create arro3.core.Array instance from a sequence of Python objects.

Expand Down Expand Up @@ -109,8 +110,8 @@ class ArrayReader:
This dunder method should not be called directly, but enables zero-copy data
transfer to other Python libraries that understand Arrow memory.

For example, you can call [`pyarrow.table()`][pyarrow.table] to convert this
ArrayReader to a pyarrow table, without copying memory.
For example, you can call [`pyarrow.chunked_array()`][pyarrow.chunked_array] to
convert this ArrayReader to a pyarrow ChunkedArray, without copying memory.
"""
def __iter__(self) -> ArrayReader: ...
def __next__(self) -> Array: ...
Expand Down
3 changes: 3 additions & 0 deletions mkdocs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,9 @@ plugins:
python:
paths: [arro3-compute/python, arro3-core/python, arro3-io/python]
options:
# We set allow_inspection: false to ensure that all docstrings come
# from the pyi files, not the Rust-facing doc comments.
allow_inspection: false
docstring_section_style: list
docstring_style: google
line_length: 80
Expand Down
34 changes: 16 additions & 18 deletions pyo3-arrow/src/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use arrow::datatypes::{
UInt64Type, UInt8Type,
};
use arrow_array::{
make_array, Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, LargeBinaryArray,
Array, ArrayRef, BinaryArray, BinaryViewArray, BooleanArray, LargeBinaryArray,
LargeStringArray, PrimitiveArray, StringArray, StringViewArray,
};
use arrow_schema::{ArrowError, DataType, Field, FieldRef};
Expand All @@ -26,6 +26,9 @@ use crate::interop::numpy::from_numpy::from_numpy;
use crate::interop::numpy::to_numpy::to_numpy;
use crate::{PyDataType, PyField};

/// A Python-facing Arrow array.
///
/// This is a wrapper around an [ArrayRef] and a [FieldRef].
#[pyclass(module = "arro3.core._core", name = "Array", subclass)]
pub struct PyArray {
array: ArrayRef,
Expand All @@ -52,11 +55,6 @@ impl PyArray {
Ok(Self { array, field })
}

pub fn from_array<A: Array>(array: A) -> Self {
let array = make_array(array.into_data());
Self::from_array_ref(array)
}

/// Create a new PyArray from an [ArrayRef], inferring its data type automatically.
pub fn from_array_ref(array: ArrayRef) -> Self {
let field = Field::new("", array.data_type().clone(), true);
Expand Down Expand Up @@ -132,7 +130,7 @@ impl Display for PyArray {
impl PyArray {
#[new]
#[pyo3(signature = (obj, /, r#type, *))]
pub fn init(py: Python, obj: PyObject, r#type: PyDataType) -> PyResult<Self> {
fn init(py: Python, obj: PyObject, r#type: PyDataType) -> PyResult<Self> {
macro_rules! impl_primitive {
($rust_type:ty, $arrow_type:ty) => {{
let values: Vec<$rust_type> = obj.extract(py)?;
Expand Down Expand Up @@ -194,7 +192,7 @@ impl PyArray {

#[pyo3(signature = (dtype=None, copy=None))]
#[allow(unused_variables)]
pub fn __array__(
fn __array__(
&self,
py: Python,
dtype: Option<PyObject>,
Expand All @@ -204,28 +202,28 @@ impl PyArray {
}

#[allow(unused_variables)]
pub fn __arrow_c_array__<'py>(
fn __arrow_c_array__<'py>(
&'py self,
py: Python<'py>,
requested_schema: Option<Bound<PyCapsule>>,
) -> PyArrowResult<Bound<PyTuple>> {
to_array_pycapsules(py, self.field.clone(), &self.array, requested_schema)
}

pub fn __eq__(&self, other: &PyArray) -> bool {
fn __eq__(&self, other: &PyArray) -> bool {
self.array.as_ref() == other.array.as_ref() && self.field == other.field
}

pub fn __len__(&self) -> usize {
fn __len__(&self) -> usize {
self.array.len()
}

pub fn __repr__(&self) -> String {
fn __repr__(&self) -> String {
self.to_string()
}

#[classmethod]
pub fn from_arrow(_cls: &Bound<PyType>, input: AnyArray) -> PyArrowResult<Self> {
fn from_arrow(_cls: &Bound<PyType>, input: AnyArray) -> PyArrowResult<Self> {
match input {
AnyArray::Array(array) => Ok(array),
AnyArray::Stream(stream) => {
Expand All @@ -239,7 +237,7 @@ impl PyArray {
}

#[classmethod]
pub fn from_arrow_pycapsule(
pub(crate) fn from_arrow_pycapsule(
_cls: &Bound<PyType>,
schema_capsule: &Bound<PyCapsule>,
array_capsule: &Bound<PyCapsule>,
Expand All @@ -249,7 +247,7 @@ impl PyArray {
}

#[classmethod]
pub fn from_numpy(
fn from_numpy(
_cls: &Bound<PyType>,
py: Python,
array: Bound<'_, PyAny>,
Expand Down Expand Up @@ -282,7 +280,7 @@ impl PyArray {
}

#[pyo3(signature = (offset=0, length=None))]
pub fn slice(&self, py: Python, offset: usize, length: Option<usize>) -> PyResult<PyObject> {
fn slice(&self, py: Python, offset: usize, length: Option<usize>) -> PyResult<PyObject> {
let length = length.unwrap_or_else(|| self.array.len() - offset);
let new_array = self.array.slice(offset, length);
PyArray::new(new_array, self.field().clone()).to_arro3(py)
Expand All @@ -293,12 +291,12 @@ impl PyArray {
Ok(PyArray::new(new_array, self.field.clone()).to_arro3(py)?)
}

pub fn to_numpy(&self, py: Python) -> PyResult<PyObject> {
fn to_numpy(&self, py: Python) -> PyResult<PyObject> {
self.__array__(py, None, None)
}

#[getter]
pub fn r#type(&self, py: Python) -> PyResult<PyObject> {
fn r#type(&self, py: Python) -> PyResult<PyObject> {
PyDataType::new(self.field.data_type().clone()).to_arro3(py)
}
}
28 changes: 16 additions & 12 deletions pyo3-arrow/src/array_reader.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,14 @@ use crate::ffi::{ArrayIterator, ArrayReader};
use crate::input::AnyArray;
use crate::{PyArray, PyChunkedArray, PyField};

/// A Python-facing Arrow array reader.
///
/// This is a wrapper around a [ArrayReader].
#[pyclass(module = "arro3.core._core", name = "ArrayReader", subclass)]
pub struct PyArrayReader(pub(crate) Option<Box<dyn ArrayReader + Send>>);

impl PyArrayReader {
/// Construct a new [PyArrayReader] from an existing [ArrayReader].
pub fn new(reader: Box<dyn ArrayReader + Send>) -> Self {
Self(Some(reader))
}
Expand All @@ -45,7 +49,7 @@ impl PyArrayReader {
for array in stream {
arrays.push(array?);
}
Ok(PyChunkedArray::new(arrays, field))
Ok(PyChunkedArray::try_new(arrays, field)?)
}

/// Access the [FieldRef] of this ArrayReader.
Expand Down Expand Up @@ -100,7 +104,7 @@ impl Display for PyArrayReader {
#[pymethods]
impl PyArrayReader {
#[allow(unused_variables)]
pub fn __arrow_c_stream__<'py>(
fn __arrow_c_stream__<'py>(
&'py mut self,
py: Python<'py>,
requested_schema: Option<Bound<PyCapsule>>,
Expand All @@ -122,23 +126,23 @@ impl PyArrayReader {
self.read_next_array(py)
}

pub fn __repr__(&self) -> String {
fn __repr__(&self) -> String {
self.to_string()
}

#[getter]
pub fn closed(&self) -> bool {
fn closed(&self) -> bool {
self.0.is_none()
}

#[classmethod]
pub fn from_arrow(_cls: &Bound<PyType>, input: AnyArray) -> PyArrowResult<Self> {
fn from_arrow(_cls: &Bound<PyType>, input: AnyArray) -> PyArrowResult<Self> {
let reader = input.into_reader()?;
Ok(Self::new(reader))
}

#[classmethod]
pub fn from_arrow_pycapsule(
pub(crate) fn from_arrow_pycapsule(
_cls: &Bound<PyType>,
capsule: &Bound<PyCapsule>,
) -> PyResult<Self> {
Expand All @@ -149,7 +153,7 @@ impl PyArrayReader {
}

#[classmethod]
pub fn from_arrays(_cls: &Bound<PyType>, field: PyField, arrays: Vec<PyArray>) -> Self {
fn from_arrays(_cls: &Bound<PyType>, field: PyField, arrays: Vec<PyArray>) -> Self {
let arrays = arrays
.into_iter()
.map(|array| {
Expand All @@ -164,16 +168,16 @@ impl PyArrayReader {
}

#[classmethod]
pub fn from_stream(_cls: &Bound<PyType>, data: &Bound<PyAny>) -> PyResult<Self> {
fn from_stream(_cls: &Bound<PyType>, data: &Bound<PyAny>) -> PyResult<Self> {
data.extract()
}

#[getter]
pub fn field(&self, py: Python) -> PyResult<PyObject> {
fn field(&self, py: Python) -> PyResult<PyObject> {
PyField::new(self.field_ref()?).to_arro3(py)
}

pub fn read_all(&mut self, py: Python) -> PyArrowResult<PyObject> {
fn read_all(&mut self, py: Python) -> PyArrowResult<PyObject> {
let stream = self
.0
.take()
Expand All @@ -183,10 +187,10 @@ impl PyArrayReader {
for array in stream {
arrays.push(array?);
}
Ok(PyChunkedArray::new(arrays, field).to_arro3(py)?)
Ok(PyChunkedArray::try_new(arrays, field)?.to_arro3(py)?)
}

pub fn read_next_array(&mut self, py: Python) -> PyArrowResult<PyObject> {
fn read_next_array(&mut self, py: Python) -> PyArrowResult<PyObject> {
let stream = self
.0
.as_mut()
Expand Down
Loading