feat: PyExpr::evaluate

This is primarily useful for diagnosing issues (e.g. #5385).

Signed-off-by: Daniel King <[email protected]>

Author: danking

vortex-python/python/vortex/_lib/expr.pyi

@@ -5,7 +5,9 @@ from datetime import date, datetime
 from typing import TypeAlias, final
 
 from typing_extensions import override
+from vortex.type_aliases import IntoArray
 
+from .arrays import Array
 from .dtype import DType
 from .scalar import ScalarPyType
 
@@ -23,6 +25,7 @@ class Expr:
     def __ge__(self, other: IntoExpr) -> Expr: ...
     def __and__(self, other: IntoExpr) -> Expr: ...
     def __or__(self, other: IntoExpr) -> Expr: ...
+    def evaluate(self, array: IntoArray) -> Array: ...
 
 def column(name: str) -> Expr: ...
 def root() -> Expr: ...

vortex-python/python/vortex/type_aliases.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright the Vortex contributors
-
-from typing import TypeAlias
+from typing import TypeAlias, Union
 
 import pyarrow as pa
 
@@ -12,6 +11,7 @@
 # TypeAliases do not support __doc__.
 IntoProjection: TypeAlias = Expr | list[str] | None
 IntoArrayIterator: TypeAlias = Array | ArrayIterator | pa.Table | pa.RecordBatchReader
+IntoArray: TypeAlias = Union[Array, "pa.Array[pa.Scalar[pa.DataType]]", pa.Table]
 
 # If you make an intersphinx reference to pyarrow.RecordBatchReader in the return type of a function
 # *and also* use the IntoProjection type alias in a parameter type, Sphinx thinks the type alias

vortex-python/src/arrays/into_array.rs

@@ -0,0 +1,62 @@
+use crate::PyVortex;
+use crate::arrays::py::PyPythonArray;
+use crate::arrays::{PyArrayRef, native::PyNativeArray};
+use crate::arrow::FromPyArrow;
+use arrow_array::ffi_stream::ArrowArrayStreamReader;
+use arrow_array::{RecordBatchReader as _, make_array};
+use arrow_data::ArrayData;
+use pyo3::{Bound, FromPyObject, PyAny, PyResult, exceptions::PyTypeError, types::PyAnyMethods};
+use vortex::dtype::DType;
+use vortex::dtype::arrow::FromArrowType as _;
+use vortex::error::VortexResult;
+use vortex::iter::{ArrayIteratorAdapter, ArrayIteratorExt};
+use vortex::{ArrayRef, arrow::FromArrowArray as _};
+
+/// Conversion type for converting Python objects into a [`vortex::Array`].
+pub struct PyIntoArray(PyArrayRef);
+
+impl PyIntoArray {
+    pub fn inner(&self) -> &ArrayRef {
+        self.0.inner()
+    }
+
+    #[allow(dead_code)]
+    pub fn into_inner(self) -> ArrayRef {
+        self.0.into_inner()
+    }
+}
+
+impl<'py> FromPyObject<'py> for PyIntoArray {
+    fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
+        if ob.is_instance_of::<PyNativeArray>() || ob.is_instance_of::<PyPythonArray>() {
+            return PyArrayRef::extract_bound(ob).map(PyIntoArray);
+        }
+
+        let py = ob.py();
+        let pa = py.import("pyarrow")?;
+
+        if ob.is_instance(&pa.getattr("Array")?)? {
+            let arrow_array_data = ArrayData::from_pyarrow_bound(ob)?;
+            return Ok(PyIntoArray(PyVortex(ArrayRef::from_arrow(
+                make_array(arrow_array_data).as_ref(),
+                false,
+            ))));
+        }
+
+        if ob.is_instance(&pa.getattr("Table")?)? {
+            let arrow_stream = ArrowArrayStreamReader::from_pyarrow_bound(ob)?;
+            let dtype = DType::from_arrow(arrow_stream.schema());
+            let vortex_iter = arrow_stream
+                .into_iter()
+                .map(|batch_result| -> VortexResult<_> {
+                    Ok(ArrayRef::from_arrow(batch_result?, false))
+                });
+            let array = ArrayIteratorAdapter::new(dtype, vortex_iter).read_all()?;
+            return Ok(PyIntoArray(PyVortex(array)));
+        }
+
+        Err(PyTypeError::new_err(
+            "Expected an object that can be converted to a Vortex ArrayRef (vortex.Array, pyarrow.Array, or pyarrow.Table)",
+        ))
+    }
+}

vortex-python/src/arrays/mod.rs

@@ -5,6 +5,7 @@ pub(crate) mod builtins;
 pub(crate) mod compressed;
 pub(crate) mod fastlanes;
 pub(crate) mod from_arrow;
+pub mod into_array;
 mod native;
 pub(crate) mod py;
 mod range_to_sequence;

vortex-python/src/expr/mod.rs

@@ -9,6 +9,8 @@ use pyo3::types::*;
 use vortex::dtype::{DType, Nullability, PType};
 use vortex::expr::{Binary, Expression, GetItem, Operator, VTableExt, and, lit, not};
 
+use crate::arrays::PyArrayRef;
+use crate::arrays::into_array::PyIntoArray;
 use crate::dtype::PyDType;
 use crate::install_module;
 use crate::scalar::factory::scalar_helper;
@@ -164,6 +166,45 @@ impl PyExpr {
     fn __getitem__(self_: PyRef<'_, Self>, field: String) -> PyResult<PyExpr> {
         get_item(field, self_.clone())
     }
+
+    /// Evaluate this expression on an in-memory array.
+    ///
+    /// Examples
+    /// --------
+    ///
+    /// Extract one column from a Vortex array:
+    ///
+    /// ```python
+    /// >>> import vortex.expr as ve
+    /// >>> import vortex as vx
+    /// >>> array = ve.column("a").evaluate(vx.array([{"a": 0, "b": "hello"}, {"a": 1, "b": "goodbye"}]))
+    /// >>> array.to_arrow_array()
+    /// <pyarrow.lib.Int64Array object at ...>
+    /// [
+    ///  0,
+    ///  1
+    /// ]
+    /// ```
+    ///
+    /// Evaluating an expression on an Arrow array or table implicitly converts it to a Vortex
+    /// array:
+    ///
+    /// >>> import pyarrow as pa
+    /// >>> array = ve.column("a").evaluate(pa.Table.from_arrays(
+    /// ...     [[0, 1, 2, 3]],
+    /// ...     names=['a'],
+    /// ... ))
+    /// >>> array
+    /// <vortex.PrimitiveArray object at ...>
+    ///
+    /// See also
+    /// --------
+    /// vortex.open : Open an on-disk Vortex array for scanning with an expression.
+    /// vortex.VortexFile : An on-disk Vortex array ready to scan with an expression.
+    /// vortex.VortexFile.scan : Scan an on-disk Vortex array with an expression.
+    fn evaluate(self_: PyRef<'_, Self>, array: PyIntoArray) -> PyResult<PyArrayRef> {
+        Ok(PyArrayRef::from(self_.evaluate(array.inner())?))
+    }
 }
 
 /// Create an expression that represents a literal value.