diff --git a/docs/encoding.rst b/docs/encoding.rst index 8448777fba..ceeb92c406 100644 --- a/docs/encoding.rst +++ b/docs/encoding.rst @@ -4,4 +4,4 @@ Arrays .. automodule:: vortex.encoding :members: :imported-members: - :special-members: __len__ + :special-members: __len__, __lt__, __le__, __eq__, __ne__, __ge__, __gt__ diff --git a/pyvortex/src/array.rs b/pyvortex/src/array.rs index c74fe746fd..7dcb9dd751 100644 --- a/pyvortex/src/array.rs +++ b/pyvortex/src/array.rs @@ -5,7 +5,7 @@ use pyo3::prelude::*; use pyo3::types::{IntoPyDict, PyList}; use vortex::array::ChunkedArray; use vortex::compute::unary::fill_forward; -use vortex::compute::{slice, take}; +use vortex::compute::{compare, slice, take, Operator}; use vortex::{Array, ArrayDType, IntoCanonical}; use crate::dtype::PyDType; @@ -14,6 +14,68 @@ use crate::python_repr::PythonRepr; #[pyclass(name = "Array", module = "vortex", sequence, subclass)] /// An array of zero or more *rows* each with the same set of *columns*. +/// +/// Examples +/// -------- +/// +/// Arrays support all the standard comparison operations: +/// +/// >>> a = vortex.encoding.array(['dog', None, 'cat', 'mouse', 'fish']) +/// >>> b = vortex.encoding.array(['doug', 'jennifer', 'casper', 'mouse', 'faust']) +/// >>> (a < b).to_arrow_array() +/// +/// [ +/// true, +/// null, +/// false, +/// false, +/// false +/// ] +/// >>> (a <= b).to_arrow_array() +/// +/// [ +/// true, +/// null, +/// false, +/// true, +/// false +/// ] +/// >>> (a == b).to_arrow_array() +/// +/// [ +/// false, +/// null, +/// false, +/// true, +/// false +/// ] +/// >>> (a != b).to_arrow_array() +/// +/// [ +/// true, +/// null, +/// true, +/// false, +/// true +/// ] +/// >>> (a >= b).to_arrow_array() +/// +/// [ +/// false, +/// null, +/// true, +/// true, +/// true +/// ] +/// >>> (a > b).to_arrow_array() +/// +/// [ +/// false, +/// null, +/// true, +/// false, +/// true +/// ] pub struct PyArray { inner: Array, } @@ -139,6 +201,54 @@ impl PyArray { PyDType::wrap(self_.py(), self_.inner.dtype().clone()) } + // Rust docs are *not* copied into Python for __lt__: https://github.com/PyO3/pyo3/issues/4326 + fn __lt__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::Lt) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + + // Rust docs are *not* copied into Python for __le__: https://github.com/PyO3/pyo3/issues/4326 + fn __le__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::Lte) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + + // Rust docs are *not* copied into Python for __eq__: https://github.com/PyO3/pyo3/issues/4326 + fn __eq__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::Eq) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + + // Rust docs are *not* copied into Python for __ne__: https://github.com/PyO3/pyo3/issues/4326 + fn __ne__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::NotEq) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + + // Rust docs are *not* copied into Python for __ge__: https://github.com/PyO3/pyo3/issues/4326 + fn __ge__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::Gte) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + + // Rust docs are *not* copied into Python for __gt__: https://github.com/PyO3/pyo3/issues/4326 + fn __gt__(&self, other: &Bound) -> PyResult { + let other = other.borrow(); + compare(&self.inner, &other.inner, Operator::Gt) + .map(|arr| PyArray { inner: arr }) + .map_err(PyVortexError::map_err) + } + /// Filter an Array by another Boolean array. /// /// Parameters