diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..3303ec07542 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10460,6 +10460,7 @@ dependencies = [ "log", "object_store", "parking_lot", + "prost 0.14.4", "pyo3", "pyo3-bytes", "pyo3-log", diff --git a/docs/api/python/expr.rst b/docs/api/python/expr.rst index 6361b558347..9ee4a7d4450 100644 --- a/docs/api/python/expr.rst +++ b/docs/api/python/expr.rst @@ -11,25 +11,174 @@ the following expression represents the set of rows for which the `age` column l >>> age = vortex.expr.column("age") >>> (23 > age) & (age < 55) # doctest: +SKIP +Expressions are picklable, so a filter built in one process can be sent to another (for example to a +``multiprocessing`` worker or a Ray task). Pickling uses the same protobuf wire format exposed by +:meth:`vortex.expr.Expr.serialize` and :func:`vortex.expr.deserialize`. + .. autosummary:: :nosignatures: - ~vortex.expr.column ~vortex.expr.Expr + ~vortex.expr.root + ~vortex.expr.column + ~vortex.expr.literal + ~vortex.expr.get_item + ~vortex.expr.not_ + ~vortex.expr.and_ + ~vortex.expr.or_ + ~vortex.expr.and_collect + ~vortex.expr.or_collect + ~vortex.expr.eq + ~vortex.expr.not_eq + ~vortex.expr.gt + ~vortex.expr.gt_eq + ~vortex.expr.lt + ~vortex.expr.lt_eq + ~vortex.expr.add + ~vortex.expr.sub + ~vortex.expr.mul + ~vortex.expr.div + ~vortex.expr.between + ~vortex.expr.is_null + ~vortex.expr.is_not_null + ~vortex.expr.fill_null + ~vortex.expr.like + ~vortex.expr.ilike + ~vortex.expr.not_like + ~vortex.expr.not_ilike + ~vortex.expr.byte_length + ~vortex.expr.select + ~vortex.expr.select_exclude + ~vortex.expr.pack + ~vortex.expr.merge + ~vortex.expr.list_contains + ~vortex.expr.list_length + ~vortex.expr.list_sum + ~vortex.expr.case_when + ~vortex.expr.zip_ + ~vortex.expr.mask + ~vortex.expr.cast + ~vortex.expr.ext_storage + ~vortex.expr.variant_get + ~vortex.expr.deserialize .. raw:: html
+Leaves and scope +---------------- + +.. autofunction:: vortex.expr.root + .. autofunction:: vortex.expr.column +.. autofunction:: vortex.expr.literal + +.. autofunction:: vortex.expr.get_item + +Boolean logic +------------- + .. autofunction:: vortex.expr.not_ .. autofunction:: vortex.expr.and_ -.. autofunction:: vortex.expr.root +.. autofunction:: vortex.expr.or_ -.. autofunction:: vortex.expr.literal +.. autofunction:: vortex.expr.and_collect + +.. autofunction:: vortex.expr.or_collect + +Comparisons and arithmetic +-------------------------- + +.. autofunction:: vortex.expr.eq + +.. autofunction:: vortex.expr.not_eq + +.. autofunction:: vortex.expr.gt + +.. autofunction:: vortex.expr.gt_eq + +.. autofunction:: vortex.expr.lt + +.. autofunction:: vortex.expr.lt_eq + +.. autofunction:: vortex.expr.add + +.. autofunction:: vortex.expr.sub + +.. autofunction:: vortex.expr.mul + +.. autofunction:: vortex.expr.div + +.. autofunction:: vortex.expr.between + +Nullability +----------- + +.. autofunction:: vortex.expr.is_null + +.. autofunction:: vortex.expr.is_not_null + +.. autofunction:: vortex.expr.fill_null + +Strings +------- + +.. autofunction:: vortex.expr.like + +.. autofunction:: vortex.expr.ilike + +.. autofunction:: vortex.expr.not_like + +.. autofunction:: vortex.expr.not_ilike + +.. autofunction:: vortex.expr.byte_length + +Structs +------- + +.. autofunction:: vortex.expr.select + +.. autofunction:: vortex.expr.select_exclude + +.. autofunction:: vortex.expr.pack + +.. autofunction:: vortex.expr.merge + +Lists +----- + +.. autofunction:: vortex.expr.list_contains + +.. autofunction:: vortex.expr.list_length + +.. autofunction:: vortex.expr.list_sum + +Conditionals and conversions +---------------------------- + +.. autofunction:: vortex.expr.case_when + +.. autofunction:: vortex.expr.zip_ + +.. autofunction:: vortex.expr.mask + +.. autofunction:: vortex.expr.cast + +.. autofunction:: vortex.expr.ext_storage + +.. autofunction:: vortex.expr.variant_get + +Serialization +------------- + +.. autofunction:: vortex.expr.deserialize + +The expression class +-------------------- .. autoclass:: vortex.expr.Expr :members: diff --git a/vortex-array/src/expr/proto.rs b/vortex-array/src/expr/proto.rs index 4f544fecafb..1dee3963db6 100644 --- a/vortex-array/src/expr/proto.rs +++ b/vortex-array/src/expr/proto.rs @@ -98,6 +98,8 @@ pub fn deserialize_expr_proto( #[cfg(test)] mod tests { use prost::Message; + use rstest::rstest; + use vortex_error::VortexResult; use vortex_proto::expr as pb; use vortex_session::VortexSession; @@ -106,11 +108,14 @@ mod tests { use crate::expr::Expression; use crate::expr::and; use crate::expr::between; + use crate::expr::byte_length; use crate::expr::eq; use crate::expr::get_item; use crate::expr::lit; + use crate::expr::mask; use crate::expr::or; use crate::expr::root; + use crate::expr::zip_expr; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::session::ScalarFnSession; @@ -141,6 +146,20 @@ mod tests { assert_eq!(&deser_expr, &expr); } + /// `ByteLength`, `Mask` and `Zip` implement `serialize`/`deserialize` but were once missing from + /// `ScalarFnSession::default()`, so they serialized fine and then failed to deserialize with + /// "unknown expression id". + #[rstest] + #[case::byte_length(byte_length(root()))] + #[case::mask(mask(root(), lit(true)))] + #[case::zip(zip_expr(lit(true), root(), lit(0)))] + fn round_trips_through_proto(#[case] expr: Expression) -> VortexResult<()> { + let buf = expr.serialize_proto()?.encode_to_vec(); + let decoded = pb::Expr::decode(buf.as_slice())?; + assert_eq!(Expression::from_proto(&decoded, &array_session())?, expr); + Ok(()) + } + #[test] fn unknown_expression_id_allow_unknown() { let session = VortexSession::empty().with::(); diff --git a/vortex-array/src/scalar_fn/session.rs b/vortex-array/src/scalar_fn/session.rs index 612ded7a5c8..211227858ca 100644 --- a/vortex-array/src/scalar_fn/session.rs +++ b/vortex-array/src/scalar_fn/session.rs @@ -14,6 +14,7 @@ use crate::scalar_fn::ScalarFnPluginRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::byte_length::ByteLength; use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::ext_storage::ExtStorage; use crate::scalar_fn::fns::fill_null::FillNull; @@ -25,12 +26,14 @@ use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::list_length::ListLength; use crate::scalar_fn::fns::list_sum::ListSum; use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::fns::mask::Mask; use crate::scalar_fn::fns::merge::Merge; use crate::scalar_fn::fns::not::Not; use crate::scalar_fn::fns::pack::Pack; use crate::scalar_fn::fns::select::Select; use crate::scalar_fn::fns::stat::StatFn; use crate::scalar_fn::fns::variant_get::VariantGet; +use crate::scalar_fn::fns::zip::Zip; /// Registry of scalar function vtables. pub type ScalarFnRegistry = ArcSwapMap; @@ -62,6 +65,7 @@ impl Default for ScalarFnSession { // Register built-in expressions. this.register(Between); this.register(Binary); + this.register(ByteLength); this.register(Cast); this.register(ExtStorage); this.register(FillNull); @@ -73,12 +77,14 @@ impl Default for ScalarFnSession { this.register(ListLength); this.register(ListSum); this.register(Literal); + this.register(Mask); this.register(Merge); this.register(Not); this.register(Pack); this.register(Select); this.register(StatFn); this.register(VariantGet); + this.register(Zip); this } diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 2400046b9ac..b7a3f6e3588 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -45,6 +45,7 @@ object_store = { workspace = true, features = [ "http", ] } parking_lot = { workspace = true } +prost = { workspace = true } pyo3 = { workspace = true, features = ["abi3", "abi3-py311"] } pyo3-bytes = { workspace = true } pyo3-log = { workspace = true } diff --git a/vortex-python/python/vortex/_lib/expr.pyi b/vortex-python/python/vortex/_lib/expr.pyi index c69307266de..47172e3b87f 100644 --- a/vortex-python/python/vortex/_lib/expr.pyi +++ b/vortex-python/python/vortex/_lib/expr.pyi @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +from collections.abc import Iterable, Mapping, Sequence from datetime import date, datetime -from typing import TypeAlias, final +from typing import Literal, TypeAlias, final from typing_extensions import override from .dtype import DType from .scalar import ScalarPyType -IntoExpr: TypeAlias = Expr | int | str | date | datetime | None +IntoExpr: TypeAlias = Expr | bool | int | float | str | bytes | date | datetime | None +"""A value accepted anywhere an expression is expected. Non-``Expr`` values become literals.""" + +VariantPath: TypeAlias = str | int | Sequence[str | int] @final class Expr: @@ -22,17 +26,97 @@ class Expr: def __gt__(self, other: IntoExpr) -> Expr: ... def __ge__(self, other: IntoExpr) -> Expr: ... def __and__(self, other: IntoExpr) -> Expr: ... + def __rand__(self, other: IntoExpr) -> Expr: ... def __or__(self, other: IntoExpr) -> Expr: ... + def __ror__(self, other: IntoExpr) -> Expr: ... + def __invert__(self) -> Expr: ... def __add__(self, other: IntoExpr) -> Expr: ... + def __radd__(self, other: IntoExpr) -> Expr: ... def __sub__(self, other: IntoExpr) -> Expr: ... + def __rsub__(self, other: IntoExpr) -> Expr: ... def __mul__(self, other: IntoExpr) -> Expr: ... + def __rmul__(self, other: IntoExpr) -> Expr: ... def __truediv__(self, other: IntoExpr) -> Expr: ... + def __rtruediv__(self, other: IntoExpr) -> Expr: ... + def __getitem__(self, field: str) -> Expr: ... + def serialize(self) -> bytes: ... + @override + def __reduce__(self) -> tuple[object, tuple[bytes]]: ... -def column(name: str) -> Expr: ... +# Leaves and scope def root() -> Expr: ... +def column(name: str) -> Expr: ... def literal(dtype: DType, value: ScalarPyType) -> Expr: ... -def not_(child: Expr) -> Expr: ... -def and_(left: Expr, right: Expr) -> Expr: ... -def cast(child: Expr, dtype: DType) -> Expr: ... -def is_null(child: Expr) -> Expr: ... -def is_not_null(child: Expr) -> Expr: ... +def get_item(field: str, child: IntoExpr | None = None) -> Expr: ... + +# Boolean logic +def not_(child: IntoExpr) -> Expr: ... +def and_(left: IntoExpr, right: IntoExpr) -> Expr: ... +def or_(left: IntoExpr, right: IntoExpr) -> Expr: ... +def and_collect(exprs: Iterable[IntoExpr]) -> Expr | None: ... +def or_collect(exprs: Iterable[IntoExpr]) -> Expr | None: ... + +# Comparisons and arithmetic +def eq(left: IntoExpr, right: IntoExpr) -> Expr: ... +def not_eq(left: IntoExpr, right: IntoExpr) -> Expr: ... +def gt(left: IntoExpr, right: IntoExpr) -> Expr: ... +def gt_eq(left: IntoExpr, right: IntoExpr) -> Expr: ... +def lt(left: IntoExpr, right: IntoExpr) -> Expr: ... +def lt_eq(left: IntoExpr, right: IntoExpr) -> Expr: ... +def add(left: IntoExpr, right: IntoExpr) -> Expr: ... +def sub(left: IntoExpr, right: IntoExpr) -> Expr: ... +def mul(left: IntoExpr, right: IntoExpr) -> Expr: ... +def div(left: IntoExpr, right: IntoExpr) -> Expr: ... +def between( + child: IntoExpr, + lower: IntoExpr, + upper: IntoExpr, + *, + lower_strict: bool = False, + upper_strict: bool = False, +) -> Expr: ... + +# Nullability +def is_null(child: IntoExpr) -> Expr: ... +def is_not_null(child: IntoExpr) -> Expr: ... +def fill_null(child: IntoExpr, fill_value: IntoExpr) -> Expr: ... + +# Strings +def like(child: IntoExpr, pattern: IntoExpr) -> Expr: ... +def ilike(child: IntoExpr, pattern: IntoExpr) -> Expr: ... +def not_like(child: IntoExpr, pattern: IntoExpr) -> Expr: ... +def not_ilike(child: IntoExpr, pattern: IntoExpr) -> Expr: ... +def byte_length(child: IntoExpr) -> Expr: ... + +# Structs +def select(fields: str | Iterable[str], child: IntoExpr | None = None) -> Expr: ... +def select_exclude(fields: str | Iterable[str], child: IntoExpr | None = None) -> Expr: ... +def pack( + fields: Mapping[str, IntoExpr] | Iterable[tuple[str, IntoExpr]], + *, + nullable: bool = False, +) -> Expr: ... +def merge( + exprs: Iterable[IntoExpr], + *, + duplicate_handling: Literal["error", "rightmost"] = "error", +) -> Expr: ... + +# Lists +def list_contains(child: IntoExpr, value: IntoExpr) -> Expr: ... +def list_length(child: IntoExpr) -> Expr: ... +def list_sum(child: IntoExpr, *, skip_nans: bool = True) -> Expr: ... + +# Conditionals and misc +def case_when( + when_then: Iterable[tuple[IntoExpr, IntoExpr]], + else_value: IntoExpr | None = None, +) -> Expr: ... +def zip_(mask: IntoExpr, if_true: IntoExpr, if_false: IntoExpr) -> Expr: ... +def mask(child: IntoExpr, mask: IntoExpr) -> Expr: ... +def cast(child: IntoExpr, dtype: DType) -> Expr: ... +def ext_storage(child: IntoExpr) -> Expr: ... +def variant_get(child: IntoExpr, path: VariantPath, dtype: DType | None = None) -> Expr: ... + +# Serialization +def deserialize(data: bytes) -> Expr: ... diff --git a/vortex-python/python/vortex/expr.py b/vortex-python/python/vortex/expr.py index a14643a8463..a2ccefcdece 100644 --- a/vortex-python/python/vortex/expr.py +++ b/vortex-python/python/vortex/expr.py @@ -2,6 +2,94 @@ # SPDX-FileCopyrightText: Copyright the Vortex contributors -from ._lib.expr import Expr, and_, cast, column, literal, not_, root # pyright: ignore[reportMissingModuleSource] +from ._lib.expr import ( # pyright: ignore[reportMissingModuleSource] + Expr, + add, + and_, + and_collect, + between, + byte_length, + case_when, + cast, + column, + deserialize, + div, + eq, + ext_storage, + fill_null, + get_item, + gt, + gt_eq, + ilike, + is_not_null, + is_null, + like, + list_contains, + list_length, + list_sum, + literal, + lt, + lt_eq, + mask, + merge, + mul, + not_, + not_eq, + not_ilike, + not_like, + or_, + or_collect, + pack, + root, + select, + select_exclude, + sub, + variant_get, + zip_, +) -__all__ = ["Expr", "column", "literal", "root", "not_", "and_", "cast"] +__all__ = [ + "Expr", + "add", + "and_", + "and_collect", + "between", + "byte_length", + "case_when", + "cast", + "column", + "deserialize", + "div", + "eq", + "ext_storage", + "fill_null", + "get_item", + "gt", + "gt_eq", + "ilike", + "is_not_null", + "is_null", + "like", + "list_contains", + "list_length", + "list_sum", + "literal", + "lt", + "lt_eq", + "mask", + "merge", + "mul", + "not_", + "not_eq", + "not_ilike", + "not_like", + "or_", + "or_collect", + "pack", + "root", + "select", + "select_exclude", + "sub", + "variant_get", + "zip_", +] diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index 4e2fef96ef1..be3a0339143 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -3,39 +3,102 @@ use std::ops::Deref; +use prost::Message; +use pyo3::exceptions::PyTypeError; use pyo3::exceptions::PyValueError; +use pyo3::intern; use pyo3::prelude::*; use pyo3::types::*; +use vortex::aggregate_fn::NumericalAggregateOpts; use vortex::dtype::DType; +use vortex::dtype::FieldName; +use vortex::dtype::FieldNames; use vortex::dtype::Nullability; -use vortex::dtype::PType; use vortex::expr; use vortex::expr::Expression; -use vortex::expr::and; use vortex::expr::lit; -use vortex::expr::not; +use vortex::expr::proto::ExprSerializeProtoExt; +use vortex::proto::expr as pb; use vortex::scalar_fn::ScalarFnVTableExt; +use vortex::scalar_fn::fns::between::BetweenOptions; +use vortex::scalar_fn::fns::between::StrictComparison; use vortex::scalar_fn::fns::binary::Binary; -use vortex::scalar_fn::fns::get_item::GetItem; +use vortex::scalar_fn::fns::merge::DuplicateHandling; use vortex::scalar_fn::fns::operators::Operator; +use vortex::scalar_fn::fns::variant_get::VariantPath; +use vortex::scalar_fn::fns::variant_get::VariantPathElement; use crate::dtype::PyDType; +use crate::error::PyVortexResult; use crate::install_module; use crate::scalar::factory::scalar_helper; +use crate::session::session; pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { let m = PyModule::new(py, "expr")?; parent.add_submodule(&m)?; install_module("vortex._lib.expr", &m)?; - m.add_function(wrap_pyfunction!(column, &m)?)?; + // Leaves and scope m.add_function(wrap_pyfunction!(root, &m)?)?; + m.add_function(wrap_pyfunction!(column, &m)?)?; m.add_function(wrap_pyfunction!(literal, &m)?)?; + m.add_function(wrap_pyfunction!(get_item, &m)?)?; + + // Boolean logic m.add_function(wrap_pyfunction!(not_, &m)?)?; m.add_function(wrap_pyfunction!(and_, &m)?)?; - m.add_function(wrap_pyfunction!(cast, &m)?)?; + m.add_function(wrap_pyfunction!(or_, &m)?)?; + m.add_function(wrap_pyfunction!(and_collect, &m)?)?; + m.add_function(wrap_pyfunction!(or_collect, &m)?)?; + + // Comparisons and arithmetic + m.add_function(wrap_pyfunction!(eq, &m)?)?; + m.add_function(wrap_pyfunction!(not_eq, &m)?)?; + m.add_function(wrap_pyfunction!(gt, &m)?)?; + m.add_function(wrap_pyfunction!(gt_eq, &m)?)?; + m.add_function(wrap_pyfunction!(lt, &m)?)?; + m.add_function(wrap_pyfunction!(lt_eq, &m)?)?; + m.add_function(wrap_pyfunction!(add, &m)?)?; + m.add_function(wrap_pyfunction!(sub, &m)?)?; + m.add_function(wrap_pyfunction!(mul, &m)?)?; + m.add_function(wrap_pyfunction!(div, &m)?)?; + m.add_function(wrap_pyfunction!(between, &m)?)?; + + // Nullability m.add_function(wrap_pyfunction!(is_null, &m)?)?; m.add_function(wrap_pyfunction!(is_not_null, &m)?)?; + m.add_function(wrap_pyfunction!(fill_null, &m)?)?; + + // Strings + m.add_function(wrap_pyfunction!(like, &m)?)?; + m.add_function(wrap_pyfunction!(ilike, &m)?)?; + m.add_function(wrap_pyfunction!(not_like, &m)?)?; + m.add_function(wrap_pyfunction!(not_ilike, &m)?)?; + m.add_function(wrap_pyfunction!(byte_length, &m)?)?; + + // Structs + m.add_function(wrap_pyfunction!(select, &m)?)?; + m.add_function(wrap_pyfunction!(select_exclude, &m)?)?; + m.add_function(wrap_pyfunction!(pack, &m)?)?; + m.add_function(wrap_pyfunction!(merge, &m)?)?; + + // Lists + m.add_function(wrap_pyfunction!(list_contains, &m)?)?; + m.add_function(wrap_pyfunction!(list_length, &m)?)?; + m.add_function(wrap_pyfunction!(list_sum, &m)?)?; + + // Conditionals and misc + m.add_function(wrap_pyfunction!(case_when, &m)?)?; + m.add_function(wrap_pyfunction!(zip_, &m)?)?; + m.add_function(wrap_pyfunction!(mask, &m)?)?; + m.add_function(wrap_pyfunction!(cast, &m)?)?; + m.add_function(wrap_pyfunction!(ext_storage, &m)?)?; + m.add_function(wrap_pyfunction!(variant_get, &m)?)?; + + // Serialization + m.add_function(wrap_pyfunction!(deserialize, &m)?)?; + m.add_class::()?; Ok(()) @@ -75,37 +138,107 @@ impl PyExpr { } } +/// A Python value that can be coerced into an [`Expression`]. +/// +/// Accepts an existing [`PyExpr`], or any Python value convertible to a Vortex scalar (including +/// `None`, `bool`, `int`, `float`, `str`, `bytes`, `list`, `dict`, and `vortex.Scalar`), which is +/// wrapped in a literal expression. +pub struct PyIntoExpr(Expression); + +impl PyIntoExpr { + pub fn into_inner(self) -> Expression { + self.0 + } +} + +impl<'py> FromPyObject<'_, 'py> for PyIntoExpr { + type Error = PyErr; + + fn extract(ob: Borrowed<'_, 'py, PyAny>) -> Result { + coerce_expression(&ob).map(PyIntoExpr) + } +} + +/// Coerce an arbitrary Python object into an [`Expression`]. +fn coerce_expression(value: &Bound<'_, PyAny>) -> PyResult { + if let Ok(value) = value.cast::() { + return Ok(value.get().inner.clone()); + } + Ok(lit(scalar_helper(value, None).map_err(PyErr::from)?)) +} + fn py_binary_operator<'py>( left: PyRef<'py, PyExpr>, operator: Operator, - right: Bound<'py, PyExpr>, + right: &Bound<'py, PyAny>, ) -> PyResult> { + let right = coerce_expression(right)?; Bound::new( left.py(), PyExpr { - inner: Binary.new_expr(operator, [left.inner.clone(), right.borrow().inner.clone()]), + inner: Binary.new_expr(operator, [left.inner.clone(), right]), }, ) } -fn coerce_expr<'py>(value: &Bound<'py, PyAny>) -> PyResult> { - let nonnull = Nullability::NonNullable; - if let Ok(value) = value.cast::() { - Ok(value.clone()) - } else if let Ok(value) = value.cast::() { - scalar(DType::Null, value) - } else if let Ok(value) = value.cast::() { - scalar(DType::Primitive(PType::I64, nonnull), value) - } else if let Ok(value) = value.cast::() { - scalar(DType::Primitive(PType::F64, nonnull), value) - } else if let Ok(value) = value.cast::() { - scalar(DType::Utf8(nonnull), value) - } else if let Ok(value) = value.cast::() { - scalar(DType::Binary(nonnull), value) +/// The reflected form of [`py_binary_operator`], used for ` `. +fn py_reflected_operator<'py>( + right: PyRef<'py, PyExpr>, + operator: Operator, + left: &Bound<'py, PyAny>, +) -> PyResult> { + let left = coerce_expression(left)?; + Bound::new( + right.py(), + PyExpr { + inner: Binary.new_expr(operator, [left, right.inner.clone()]), + }, + ) +} + +fn field_names(fields: &Bound<'_, PyAny>) -> PyResult { + if let Ok(name) = fields.cast::() { + return Ok(FieldNames::from(vec![FieldName::from( + name.extract::()?, + )])); + } + Ok(fields + .try_iter()? + .map(|field| field?.extract::().map(FieldName::from)) + .collect::>>()? + .into()) +} + +/// Extract `(name, expression)` pairs from either a mapping or an iterable of 2-tuples. +fn named_exprs(fields: &Bound<'_, PyAny>) -> PyResult> { + let items: Vec> = if let Ok(dict) = fields.cast::() { + dict.items().iter().collect() + } else { + fields.try_iter()?.collect::>>()? + }; + + items + .into_iter() + .map(|item| { + let (name, value): (String, Bound<'_, PyAny>) = item.extract()?; + Ok((FieldName::from(name), coerce_expression(&value)?)) + }) + .collect() +} + +fn nullability(nullable: bool) -> Nullability { + if nullable { + Nullability::Nullable + } else { + Nullability::NonNullable + } +} + +fn strictness(strict: bool) -> StrictComparison { + if strict { + StrictComparison::Strict } else { - Err(PyValueError::new_err(format!( - "expected None, int, float, str, or bytes but found: {value}" - ))) + StrictComparison::NonStrict } } @@ -119,93 +252,217 @@ impl PyExpr { self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Eq, coerce_expr(right)?) + py_binary_operator(self_, Operator::Eq, right) } fn __ne__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::NotEq, coerce_expr(right)?) + py_binary_operator(self_, Operator::NotEq, right) } fn __gt__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Gt, coerce_expr(right)?) + py_binary_operator(self_, Operator::Gt, right) } fn __ge__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Gte, coerce_expr(right)?) + py_binary_operator(self_, Operator::Gte, right) } fn __lt__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Lt, coerce_expr(right)?) + py_binary_operator(self_, Operator::Lt, right) } fn __le__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Lte, coerce_expr(right)?) + py_binary_operator(self_, Operator::Lte, right) } fn __and__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::And, coerce_expr(right)?) + py_binary_operator(self_, Operator::And, right) + } + + fn __rand__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::And, left) } fn __or__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Or, coerce_expr(right)?) + py_binary_operator(self_, Operator::Or, right) + } + + fn __ror__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::Or, left) + } + + fn __invert__(self_: PyRef<'_, Self>) -> PyExpr { + PyExpr { + inner: expr::not(self_.inner.clone()), + } } fn __add__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Add, coerce_expr(right)?) + py_binary_operator(self_, Operator::Add, right) + } + + fn __radd__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::Add, left) } fn __sub__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Sub, coerce_expr(right)?) + py_binary_operator(self_, Operator::Sub, right) + } + + fn __rsub__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::Sub, left) } fn __mul__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Mul, coerce_expr(right)?) + py_binary_operator(self_, Operator::Mul, right) + } + + fn __rmul__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::Mul, left) } fn __truediv__<'py>( self_: PyRef<'py, Self>, right: &Bound<'py, PyAny>, ) -> PyResult> { - py_binary_operator(self_, Operator::Div, coerce_expr(right)?) + py_binary_operator(self_, Operator::Div, right) + } + + fn __rtruediv__<'py>( + self_: PyRef<'py, Self>, + left: &Bound<'py, PyAny>, + ) -> PyResult> { + py_reflected_operator(self_, Operator::Div, left) } // Special methods docstrings cannot be defined in Rust. Write a docstring in the corresponding // rST file. https://github.com/PyO3/pyo3/issues/4326 - fn __getitem__(self_: PyRef<'_, Self>, field: String) -> PyResult { - get_item(field, self_.clone()) + fn __getitem__(self_: PyRef<'_, Self>, field: String) -> PyExpr { + PyExpr { + inner: expr::get_item(field, self_.inner.clone()), + } + } + + /// Serialize this expression to its Vortex protobuf wire format. + /// + /// The result can be sent to another process or machine and rebuilt with + /// :func:`vortex.expr.deserialize`. + /// + /// Returns + /// ------- + /// :class:`.bytes` + /// + /// Raises + /// ------ + /// :class:`RuntimeError` + /// If the expression contains a scalar function that is not serializable. + /// + /// Examples + /// -------- + /// + /// ```python + /// >>> import vortex.expr as ve + /// >>> expr = ve.column("age") > 21 + /// >>> str(ve.deserialize(expr.serialize())) == str(expr) + /// True + /// ``` + fn serialize<'py>(self_: PyRef<'py, Self>) -> PyVortexResult> { + let proto = self_.inner.serialize_proto()?; + Ok(PyBytes::new(self_.py(), &proto.encode_to_vec())) + } + + /// Support for Python's pickle protocol, backed by the protobuf wire format. + /// + /// This lets expressions cross process boundaries, for example as filter pushdown in a + /// multiprocessing or Ray worker. + fn __reduce__<'py>( + self_: PyRef<'py, Self>, + ) -> PyVortexResult<(Bound<'py, PyAny>, (Bound<'py, PyBytes>,))> { + let py = self_.py(); + let proto = self_.inner.serialize_proto()?; + let bytes = PyBytes::new(py, &proto.encode_to_vec()); + + let module = PyModule::import(py, "vortex._lib.expr")?; + let deserialize_fn = module.getattr(intern!(py, "deserialize"))?; + + Ok((deserialize_fn, (bytes,))) } } +/// Rebuild an expression from its protobuf wire format. +/// +/// Parameters +/// ---------- +/// data : :class:`.bytes` +/// Bytes produced by :meth:`vortex.expr.Expr.serialize`. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> expr = ve.column("age") > 21 +/// >>> str(ve.deserialize(expr.serialize())) == str(expr) +/// True +/// ``` +#[pyfunction] +pub fn deserialize(data: &[u8]) -> PyVortexResult { + let proto = pb::Expr::decode(data) + .map_err(|err| PyValueError::new_err(format!("invalid Vortex expression bytes: {err}")))?; + Ok(PyExpr { + inner: Expression::from_proto(&proto, session())?, + }) +} + /// Create an expression that represents a literal value. /// /// Parameters @@ -294,6 +551,36 @@ pub fn column<'py>(name: &Bound<'py, PyString>) -> PyResult> ) } +/// Extract a named field from a struct expression. +/// +/// Parameters +/// ---------- +/// field : :class:`str` +/// The name of the field. +/// child : :class:`vortex.Expr`, optional +/// The struct expression to read from. Defaults to :func:`.root`. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.get_item("yy", ve.column("y")) +/// +/// ``` +#[pyfunction] +#[pyo3(signature = (field, child = None))] +pub fn get_item(field: String, child: Option) -> PyExpr { + let child = child.map_or_else(expr::root, PyIntoExpr::into_inner); + PyExpr { + inner: expr::get_item(field, child), + } +} + pub fn scalar<'py>(dtype: DType, value: &Bound<'py, PyAny>) -> PyResult> { let py = value.py(); Bound::new( @@ -304,12 +591,6 @@ pub fn scalar<'py>(dtype: DType, value: &Bound<'py, PyAny>) -> PyResult PyResult { - Ok(PyExpr { - inner: GetItem.new_expr(field.into(), [child.inner]), - }) -} - /// Negate a Boolean expression. /// /// Parameters @@ -331,10 +612,10 @@ pub fn get_item(field: String, child: PyExpr) -> PyResult { /// /// ``` #[pyfunction] -pub fn not_(child: PyExpr) -> PyResult { - Ok(PyExpr { - inner: not(child.inner), - }) +pub fn not_(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::not(child.into_inner()), + } } /// True if both arguments are true. @@ -361,12 +642,595 @@ pub fn not_(child: PyExpr) -> PyResult { /// /// ``` #[pyfunction] -pub fn and_(left: PyExpr, right: PyExpr) -> PyResult { +pub fn and_(left: PyIntoExpr, right: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::and(left.into_inner(), right.into_inner()), + } +} + +/// True if either argument is true. +/// +/// Parameters +/// ---------- +/// left : :class:`Expr` +/// A boolean expression. +/// +/// right : :class:`Expr` +/// A boolean expression. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> import vortex as vx +/// >>> ve.or_(ve.literal(vx.bool_(), True), ve.literal(vx.bool_(), False)) +/// +/// ``` +#[pyfunction] +pub fn or_(left: PyIntoExpr, right: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::or(left.into_inner(), right.into_inner()), + } +} + +/// Combine expressions with logical AND using a balanced tree. +/// +/// Parameters +/// ---------- +/// exprs : :class:`Iterable` +/// The boolean expressions to combine. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` or ``None`` +/// ``None`` if ``exprs`` is empty. +#[pyfunction] +pub fn and_collect(exprs: Vec) -> Option { + expr::and_collect(exprs.into_iter().map(PyIntoExpr::into_inner)).map(PyExpr::from) +} + +/// Combine expressions with logical OR using a balanced tree. +/// +/// Parameters +/// ---------- +/// exprs : :class:`Iterable` +/// The boolean expressions to combine. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` or ``None`` +/// ``None`` if ``exprs`` is empty. +#[pyfunction] +pub fn or_collect(exprs: Vec) -> Option { + expr::or_collect(exprs.into_iter().map(PyIntoExpr::into_inner)).map(PyExpr::from) +} + +macro_rules! binary_fn { + ($name:ident, $factory:path, $doc:literal) => { + #[doc = $doc] + /// Parameters + /// ---------- + /// left : :class:`Any` + /// right : :class:`Any` + /// + /// Returns + /// ------- + /// :class:`vortex.Expr` + #[pyfunction] + pub fn $name(left: PyIntoExpr, right: PyIntoExpr) -> PyExpr { + PyExpr { + inner: $factory(left.into_inner(), right.into_inner()), + } + } + }; +} + +binary_fn!(eq, expr::eq, "True where both arguments are equal."); +binary_fn!( + not_eq, + expr::not_eq, + "True where the arguments are not equal." +); +binary_fn!(gt, expr::gt, "True where `left` is greater than `right`."); +binary_fn!( + gt_eq, + expr::gt_eq, + "True where `left` is greater than or equal to `right`." +); +binary_fn!(lt, expr::lt, "True where `left` is less than `right`."); +binary_fn!( + lt_eq, + expr::lt_eq, + "True where `left` is less than or equal to `right`." +); +binary_fn!( + add, + expr::checked_add, + "The sum of the arguments, erroring on overflow." +); +binary_fn!(sub, sub_expr, "The difference between the arguments."); +binary_fn!(mul, mul_expr, "The product of the arguments."); +binary_fn!(div, div_expr, "`left` divided by `right`."); + +fn sub_expr(left: Expression, right: Expression) -> Expression { + Binary.new_expr(Operator::Sub, [left, right]) +} + +fn mul_expr(left: Expression, right: Expression) -> Expression { + Binary.new_expr(Operator::Mul, [left, right]) +} + +fn div_expr(left: Expression, right: Expression) -> Expression { + Binary.new_expr(Operator::Div, [left, right]) +} + +/// True where `child` lies between `lower` and `upper`. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// The expression to test. +/// lower : :class:`Any` +/// The lower bound. +/// upper : :class:`Any` +/// The upper bound. +/// lower_strict : :class:`bool` +/// If ``True``, compare the lower bound with ``<`` instead of ``<=``. +/// upper_strict : :class:`bool` +/// If ``True``, compare the upper bound with ``<`` instead of ``<=``. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.between(ve.column("age"), 23, 55) +/// +/// ``` +#[pyfunction] +#[pyo3(signature = (child, lower, upper, *, lower_strict = false, upper_strict = false))] +pub fn between( + child: PyIntoExpr, + lower: PyIntoExpr, + upper: PyIntoExpr, + lower_strict: bool, + upper_strict: bool, +) -> PyExpr { + PyExpr { + inner: expr::between( + child.into_inner(), + lower.into_inner(), + upper.into_inner(), + BetweenOptions { + lower_strict: strictness(lower_strict), + upper_strict: strictness(upper_strict), + }, + ), + } +} + +/// Checks which elements of its child are null. +/// +/// Parameters +/// ---------- +/// child : :class:`Expr` +/// Any expression. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn is_null(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::is_null(child.into_inner()), + } +} + +/// Creates an expression that checks for non-null values. +/// +/// Parameters +/// ---------- +/// child : :class:`vortex.Expr` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn is_not_null(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::is_not_null(child.into_inner()), + } +} + +/// Replace null values with a fill value. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// fill_value : :class:`Any` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.fill_null(ve.column("age"), 0) +/// +/// ``` +#[pyfunction] +pub fn fill_null(child: PyIntoExpr, fill_value: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::fill_null(child.into_inner(), fill_value.into_inner()), + } +} + +/// A SQL ``LIKE`` expression. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// The string expression to match. +/// pattern : :class:`Any` +/// The SQL ``LIKE`` pattern, where ``%`` matches any run of characters and ``_`` matches any +/// single character. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.like(ve.column("name"), "Ali%") +/// +/// ``` +#[pyfunction] +pub fn like(child: PyIntoExpr, pattern: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::like(child.into_inner(), pattern.into_inner()), + } +} + +/// A case-insensitive SQL ``ILIKE`` expression. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// pattern : :class:`Any` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn ilike(child: PyIntoExpr, pattern: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::ilike(child.into_inner(), pattern.into_inner()), + } +} + +/// A negated SQL ``NOT LIKE`` expression. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// pattern : :class:`Any` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn not_like(child: PyIntoExpr, pattern: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::not_like(child.into_inner(), pattern.into_inner()), + } +} + +/// A negated case-insensitive SQL ``NOT ILIKE`` expression. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// pattern : :class:`Any` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn not_ilike(child: PyIntoExpr, pattern: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::not_ilike(child.into_inner(), pattern.into_inner()), + } +} + +/// The byte length of each element, akin to SQL ``OCTET_LENGTH()``. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn byte_length(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::byte_length(child.into_inner()), + } +} + +/// Project only the named fields of a struct expression. +/// +/// Parameters +/// ---------- +/// fields : :class:`str` or :class:`Iterable` of :class:`str` +/// The field names to keep. +/// child : :class:`vortex.Expr`, optional +/// The struct expression to project. Defaults to :func:`.root`. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.select(["name", "age"]) +/// +/// ``` +#[pyfunction] +#[pyo3(signature = (fields, child = None))] +pub fn select(fields: &Bound<'_, PyAny>, child: Option) -> PyResult { + let child = child.map_or_else(expr::root, PyIntoExpr::into_inner); + Ok(PyExpr { + inner: expr::select(field_names(fields)?, child), + }) +} + +/// Project every field of a struct expression except the named ones. +/// +/// Parameters +/// ---------- +/// fields : :class:`str` or :class:`Iterable` of :class:`str` +/// The field names to drop. +/// child : :class:`vortex.Expr`, optional +/// The struct expression to project. Defaults to :func:`.root`. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +#[pyo3(signature = (fields, child = None))] +pub fn select_exclude(fields: &Bound<'_, PyAny>, child: Option) -> PyResult { + let child = child.map_or_else(expr::root, PyIntoExpr::into_inner); + Ok(PyExpr { + inner: expr::select_exclude(field_names(fields)?, child), + }) +} + +/// Pack expressions into a struct with named fields. +/// +/// Parameters +/// ---------- +/// fields : :class:`dict` or :class:`Iterable` of (:class:`str`, :class:`Any`) +/// The field names and their expressions. +/// nullable : :class:`bool` +/// Whether the resulting struct is nullable. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.pack({"id": ve.column("user_id"), "constant": 42}) +/// +/// ``` +#[pyfunction] +#[pyo3(signature = (fields, *, nullable = false))] +pub fn pack(fields: &Bound<'_, PyAny>, nullable: bool) -> PyResult { Ok(PyExpr { - inner: and(left.inner, right.inner), + inner: expr::pack(named_exprs(fields)?, nullability(nullable)), }) } +/// Merge struct expressions into a single struct. +/// +/// Parameters +/// ---------- +/// exprs : :class:`Iterable` of :class:`vortex.Expr` +/// The struct expressions to merge. +/// duplicate_handling : :class:`str` +/// Either ``"error"`` (the default) to reject duplicated field names, or ``"rightmost"`` to +/// take the value from the right-most expression. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +#[pyo3(signature = (exprs, *, duplicate_handling = "error"))] +pub fn merge(exprs: Vec, duplicate_handling: &str) -> PyResult { + let duplicate_handling = match duplicate_handling.to_ascii_lowercase().as_str() { + "error" => DuplicateHandling::Error, + "rightmost" | "right_most" => DuplicateHandling::RightMost, + other => { + return Err(PyValueError::new_err(format!( + "duplicate_handling must be 'error' or 'rightmost', but found: {other}" + ))); + } + }; + Ok(PyExpr { + inner: expr::merge_opts( + exprs.into_iter().map(PyIntoExpr::into_inner), + duplicate_handling, + ), + }) +} + +/// True where the list contains the given value. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// A list expression. +/// value : :class:`Any` +/// The value to search for. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn list_contains(child: PyIntoExpr, value: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::list_contains(child.into_inner(), value.into_inner()), + } +} + +/// The number of elements in each list, akin to SQL ``CARDINALITY()``. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// A list or fixed-size-list expression. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn list_length(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::list_length(child.into_inner()), + } +} + +/// The sum of the elements of each list. +/// +/// Follows SQL ``SUM`` semantics per list: null lists, empty lists, and lists whose elements are +/// all null yield null, and null elements are skipped. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// A list or fixed-size-list expression. +/// skip_nans : :class:`bool` +/// If ``True`` (the default), NaN float elements are skipped. Otherwise a single NaN poisons +/// the list's sum. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +#[pyo3(signature = (child, *, skip_nans = true))] +pub fn list_sum(child: PyIntoExpr, skip_nans: bool) -> PyExpr { + PyExpr { + inner: expr::list_sum_opts(child.into_inner(), NumericalAggregateOpts { skip_nans }), + } +} + +/// A ``CASE WHEN`` expression. +/// +/// Parameters +/// ---------- +/// when_then : :class:`Iterable` of (:class:`Any`, :class:`Any`) +/// One or more ``(condition, value)`` pairs, evaluated in order. +/// else_value : :class:`Any`, optional +/// The value to use when no condition matches. Defaults to null. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.case_when([(ve.column("age") > 21, "adult")], else_value="minor") +/// +/// ``` +#[pyfunction] +#[pyo3(signature = (when_then, else_value = None))] +pub fn case_when( + when_then: Vec<(PyIntoExpr, PyIntoExpr)>, + else_value: Option, +) -> PyResult { + if when_then.is_empty() { + return Err(PyValueError::new_err( + "case_when requires at least one (condition, value) pair", + )); + } + let when_then = when_then + .into_iter() + .map(|(condition, value)| (condition.into_inner(), value.into_inner())) + .collect(); + Ok(PyExpr { + inner: expr::nested_case_when(when_then, else_value.map(PyIntoExpr::into_inner)), + }) +} + +/// Select element-wise between two expressions based on a boolean mask. +/// +/// Parameters +/// ---------- +/// mask : :class:`Any` +/// A boolean expression. +/// if_true : :class:`Any` +/// The value used where ``mask`` is true. +/// if_false : :class:`Any` +/// The value used where ``mask`` is false. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction(name = "zip_")] +pub fn zip_(mask: PyIntoExpr, if_true: PyIntoExpr, if_false: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::zip_expr( + mask.into_inner(), + if_true.into_inner(), + if_false.into_inner(), + ), + } +} + +/// Null out the elements of an expression where the mask is true. +/// +/// Parameters +/// ---------- +/// child : :class:`Any` +/// The expression to mask. +/// mask : :class:`Any` +/// A boolean expression. +/// +/// Returns +/// ------- +/// :class:`vortex.Expr` +#[pyfunction] +pub fn mask(child: PyIntoExpr, mask: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::mask(child.into_inner(), mask.into_inner()), + } +} + /// Cast an expression to a compatible type. /// /// Parameters @@ -399,42 +1263,92 @@ pub fn and_(left: PyExpr, right: PyExpr) -> PyResult { /// /// ``` #[pyfunction] -pub fn cast(child: PyExpr, dtype: PyDType) -> PyResult { - Ok(PyExpr { +pub fn cast(child: PyIntoExpr, dtype: PyDType) -> PyExpr { + PyExpr { inner: expr::cast(child.into_inner(), dtype.into_inner()), - }) + } } -/// Checks which elements of its child are null. +/// Extract the storage values of an extension-typed expression. /// /// Parameters /// ---------- -/// child : :class:`Expr` -/// Any expression. +/// child : :class:`Any` /// /// Returns /// ------- /// :class:`vortex.Expr` -/// ``` #[pyfunction] -pub fn is_null(child: PyExpr) -> PyResult { - Ok(PyExpr { - inner: expr::is_null(child.into_inner()), - }) +pub fn ext_storage(child: PyIntoExpr) -> PyExpr { + PyExpr { + inner: expr::ext_storage(child.into_inner()), + } } -/// Creates an expression that checks for non-null values. +/// Extract a path from a Variant expression. +/// +/// Missing paths, traversal mismatches, and failed casts all return null. /// /// Parameters /// ---------- -/// child : :class:`vortex.Expr` +/// child : :class:`Any` +/// A Variant expression. +/// path : :class:`str`, :class:`int`, or :class:`Iterable` of :class:`str` or :class:`int` +/// The path to extract. Strings select object fields, integers select list elements. +/// dtype : :class:`vortex.DType`, optional +/// The requested output type. When omitted, the result is a nullable Variant. /// /// Returns /// ------- /// :class:`vortex.Expr` +/// +/// Examples +/// -------- +/// +/// ```python +/// >>> import vortex.expr as ve +/// >>> ve.variant_get(ve.column("payload"), ["user", "id"]) +/// +/// ``` #[pyfunction] -pub fn is_not_null(child: PyExpr) -> PyResult { +#[pyo3(signature = (child, path, dtype = None))] +pub fn variant_get( + child: PyIntoExpr, + path: &Bound<'_, PyAny>, + dtype: Option, +) -> PyResult { Ok(PyExpr { - inner: expr::is_not_null(child.into_inner()), + inner: expr::variant_get( + child.into_inner(), + variant_path(path)?, + dtype.map(PyDType::into_inner), + ), }) } + +fn variant_path(path: &Bound<'_, PyAny>) -> PyResult { + if let Ok(field) = path.cast::() { + return Ok(VariantPath::field(field.extract::()?)); + } + if let Ok(index) = path.cast::() { + return Ok(VariantPath::new([VariantPathElement::index( + index.extract::()?, + )])); + } + let elements = path + .try_iter()? + .map(|element| { + let element = element?; + if let Ok(field) = element.cast::() { + Ok(VariantPathElement::field(field.extract::()?)) + } else if let Ok(index) = element.cast::() { + Ok(VariantPathElement::index(index.extract::()?)) + } else { + Err(PyTypeError::new_err(format!( + "variant path elements must be str or int, but found: {element}" + ))) + } + }) + .collect::>>()?; + Ok(VariantPath::new(elements)) +} diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py new file mode 100644 index 00000000000..c70810ba25a --- /dev/null +++ b/vortex-python/test/test_expr.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""Coverage for the expression builders exposed by :mod:`vortex.expr`.""" + +import pickle +from collections.abc import Callable +from typing import cast + +import pyarrow as pa +import pytest + +import vortex as vx +import vortex.expr as ve +from vortex.expr import Expr + + +@pytest.fixture(scope="module") +def people(tmp_path_factory: pytest.TempPathFactory) -> vx.VortexFile: + path = tmp_path_factory.mktemp("expr") / "people.vortex" + array = vx.array( + pa.array( + [ + {"name": "Alice", "age": 30, "scores": [1, 2, 3], "nested": {"city": "Paris"}}, + {"name": "Bob", "age": 25, "scores": [4], "nested": {"city": "Berlin"}}, + {"name": "alfred", "age": None, "scores": [], "nested": {"city": "Paris"}}, + {"name": "Charlie", "age": 57, "scores": [5, 6], "nested": {"city": "Lima"}}, + ] + ) + ) + vx.io.write(array, str(path)) + return vx.open(str(path)) + + +def names(vxf: vx.VortexFile, expr: Expr) -> list[str | None]: + rows = cast(list[dict[str, object]], vxf.scan(["name"], expr=expr).read_all().to_arrow_table().to_pylist()) + return [cast("str | None", row["name"]) for row in rows] + + +def column_values(vxf: vx.VortexFile, projection: Expr) -> list[object]: + table = vxf.scan(ve.pack({"value": projection})).read_all().to_arrow_table() + return cast(list[object], table.column("value").to_pylist()) + + +# -------------------------------------------------------------------------------------- +# Every builder is reachable and produces an expression +# -------------------------------------------------------------------------------------- + +BUILDERS: dict[str, Callable[[], Expr]] = { + "root": ve.root, + "column": lambda: ve.column("age"), + "literal": lambda: ve.literal(vx.int_(), 42), + "get_item": lambda: ve.get_item("city", ve.column("nested")), + "get_item_default_child": lambda: ve.get_item("age"), + "not_": lambda: ve.not_(ve.is_null(ve.column("age"))), + "and_": lambda: ve.and_(ve.is_null(ve.column("age")), ve.is_not_null(ve.column("name"))), + "or_": lambda: ve.or_(ve.is_null(ve.column("age")), ve.is_not_null(ve.column("name"))), + "eq": lambda: ve.eq(ve.column("age"), 30), + "not_eq": lambda: ve.not_eq(ve.column("age"), 30), + "gt": lambda: ve.gt(ve.column("age"), 30), + "gt_eq": lambda: ve.gt_eq(ve.column("age"), 30), + "lt": lambda: ve.lt(ve.column("age"), 30), + "lt_eq": lambda: ve.lt_eq(ve.column("age"), 30), + "add": lambda: ve.add(ve.column("age"), 1), + "sub": lambda: ve.sub(ve.column("age"), 1), + "mul": lambda: ve.mul(ve.column("age"), 2), + "div": lambda: ve.div(ve.column("age"), 2), + "between": lambda: ve.between(ve.column("age"), 26, 40), + "between_strict": lambda: ve.between(ve.column("age"), 26, 40, lower_strict=True, upper_strict=True), + "is_null": lambda: ve.is_null(ve.column("age")), + "is_not_null": lambda: ve.is_not_null(ve.column("age")), + "fill_null": lambda: ve.fill_null(ve.column("age"), 0), + "like": lambda: ve.like(ve.column("name"), "A%"), + "ilike": lambda: ve.ilike(ve.column("name"), "a%"), + "not_like": lambda: ve.not_like(ve.column("name"), "A%"), + "not_ilike": lambda: ve.not_ilike(ve.column("name"), "a%"), + "byte_length": lambda: ve.byte_length(ve.column("name")), + "select": lambda: ve.select(["name", "age"]), + "select_str": lambda: ve.select("name"), + "select_exclude": lambda: ve.select_exclude(["scores"]), + "pack_dict": lambda: ve.pack({"n": ve.column("name"), "constant": 7}), + "pack_pairs": lambda: ve.pack([("n", ve.column("name"))]), + "pack_nullable": lambda: ve.pack({"n": ve.column("name")}, nullable=True), + "merge": lambda: ve.merge([ve.select(["name"]), ve.select(["age"])]), + "merge_rightmost": lambda: ve.merge([ve.select(["name"]), ve.select(["name"])], duplicate_handling="rightmost"), + "list_contains": lambda: ve.list_contains(ve.column("scores"), 5), + "list_length": lambda: ve.list_length(ve.column("scores")), + "list_sum": lambda: ve.list_sum(ve.column("scores")), + "list_sum_nans": lambda: ve.list_sum(ve.column("scores"), skip_nans=False), + "case_when": lambda: ve.case_when([(ve.gt(ve.column("age"), 30), "old")], else_value="young"), + "case_when_no_else": lambda: ve.case_when([(ve.gt(ve.column("age"), 30), "old")]), + "zip_": lambda: ve.zip_(ve.is_null(ve.column("age")), 0, ve.column("age")), + "mask": lambda: ve.mask(ve.column("age"), ve.is_null(ve.column("age"))), + "cast": lambda: ve.cast(ve.column("age"), vx.int_(32, nullable=True)), + "and_collect": lambda: cast(Expr, ve.and_collect([ve.is_null(ve.column("age")), ve.gt(ve.column("age"), 1)])), + "or_collect": lambda: cast(Expr, ve.or_collect([ve.is_null(ve.column("age")), ve.gt(ve.column("age"), 1)])), +} + + +@pytest.mark.parametrize("name", sorted(BUILDERS)) +def test_builder_returns_expr(name: str) -> None: + assert isinstance(BUILDERS[name](), Expr) + + +# `CaseWhen::serialize` is deliberately disabled in vortex-array pending a stable wire format, so +# `case_when` expressions cannot cross a process boundary yet. +NOT_SERIALIZABLE = {"case_when", "case_when_no_else"} + + +@pytest.mark.parametrize("name", sorted(set(BUILDERS) - NOT_SERIALIZABLE)) +def test_builder_round_trips_through_proto(name: str) -> None: + expr = BUILDERS[name]() + assert str(ve.deserialize(expr.serialize())) == str(expr) + + +@pytest.mark.parametrize("name", sorted(NOT_SERIALIZABLE)) +def test_unserializable_builder_reports_clearly(name: str) -> None: + with pytest.raises(RuntimeError, match="serial"): + _ = BUILDERS[name]().serialize() + + +def test_variant_get_paths() -> None: + # Variant columns are not exercised here; this only checks the path coercions are accepted. + assert isinstance(ve.variant_get(ve.column("payload"), "user"), Expr) + assert isinstance(ve.variant_get(ve.column("payload"), 0), Expr) + assert isinstance(ve.variant_get(ve.column("payload"), ["user", 1, "id"]), Expr) + assert isinstance(ve.variant_get(ve.column("payload"), "user", vx.int_(64, nullable=True)), Expr) + + +def test_variant_get_rejects_bad_path() -> None: + with pytest.raises(TypeError): + _ = ve.variant_get(ve.column("payload"), [1.5]) # pyright: ignore[reportArgumentType] + + +def test_ext_storage_returns_expr() -> None: + assert isinstance(ve.ext_storage(ve.column("ts")), Expr) + + +def test_merge_rejects_unknown_duplicate_handling() -> None: + with pytest.raises(ValueError): + _ = ve.merge([ve.select(["name"])], duplicate_handling="nonsense") # pyright: ignore[reportArgumentType] + + +def test_case_when_requires_a_pair() -> None: + with pytest.raises(ValueError): + _ = ve.case_when([]) + + +def test_collect_returns_none_when_empty() -> None: + assert ve.and_collect([]) is None + assert ve.or_collect([]) is None + + +# -------------------------------------------------------------------------------------- +# Semantics +# -------------------------------------------------------------------------------------- + + +def test_like_and_ilike(people: vx.VortexFile) -> None: + assert names(people, ve.like(ve.column("name"), "A%")) == ["Alice"] + assert names(people, ve.ilike(ve.column("name"), "a%")) == ["Alice", "alfred"] + assert names(people, ve.not_like(ve.column("name"), "A%")) == ["Bob", "alfred", "Charlie"] + assert names(people, ve.not_ilike(ve.column("name"), "a%")) == ["Bob", "Charlie"] + + +def test_between(people: vx.VortexFile) -> None: + assert names(people, ve.between(ve.column("age"), 25, 30)) == ["Alice", "Bob"] + assert names(people, ve.between(ve.column("age"), 25, 30, lower_strict=True)) == ["Alice"] + assert names(people, ve.between(ve.column("age"), 25, 30, upper_strict=True)) == ["Bob"] + + +def test_or_and_is_null(people: vx.VortexFile) -> None: + assert names(people, ve.is_null(ve.column("age"))) == ["alfred"] + assert names(people, ve.or_(ve.eq(ve.column("age"), 25), ve.eq(ve.column("age"), 57))) == ["Bob", "Charlie"] + + +def test_arithmetic_and_list_functions(people: vx.VortexFile) -> None: + assert column_values(people, ve.add(ve.column("age"), 1)) == [31, 26, None, 58] + assert column_values(people, ve.list_length(ve.column("scores"))) == [3, 1, 0, 2] + assert column_values(people, ve.list_sum(ve.column("scores"))) == [6, 4, None, 11] + assert column_values(people, ve.list_contains(ve.column("scores"), 5)) == [False, False, False, True] + assert column_values(people, ve.byte_length(ve.column("name"))) == [5, 3, 6, 7] + assert column_values(people, ve.fill_null(ve.column("age"), 0)) == [30, 25, 0, 57] + assert column_values(people, ve.get_item("city", ve.column("nested"))) == ["Paris", "Berlin", "Paris", "Lima"] + + +def test_case_when_semantics(people: vx.VortexFile) -> None: + expr = ve.case_when([(ve.gt(ve.column("age"), 40), "senior"), (ve.gt(ve.column("age"), 26), "mid")], "junior") + assert column_values(people, expr) == ["mid", "junior", "junior", "senior"] + + +def test_select_exclude(people: vx.VortexFile) -> None: + table = people.scan(ve.select_exclude(["scores", "nested"])).read_all().to_arrow_table() + assert table.column_names == ["name", "age"] + + +# -------------------------------------------------------------------------------------- +# Operators +# -------------------------------------------------------------------------------------- + + +def test_invert_operator(people: vx.VortexFile) -> None: + assert names(people, ~ve.is_null(ve.column("age"))) == ["Alice", "Bob", "Charlie"] + + +def test_reflected_operators(people: vx.VortexFile) -> None: + # `5 < expr` falls back to `expr.__gt__(5)`; `1 + expr` needs `__radd__`. + assert names(people, 26 < ve.column("age")) == ["Alice", "Charlie"] + assert column_values(people, 1 + ve.column("age")) == [31, 26, None, 58] + assert column_values(people, 100 - ve.column("age")) == [70, 75, None, 43] + assert column_values(people, 2 * ve.column("age")) == [60, 50, None, 114] + + +def test_bool_literals_are_boolean(people: vx.VortexFile) -> None: + # A Python bool is a subclass of int, so it must be checked before the int coercion. + assert names(people, ve.is_not_null(ve.column("age")) & True) == ["Alice", "Bob", "Charlie"] + + +# -------------------------------------------------------------------------------------- +# Serialization +# -------------------------------------------------------------------------------------- + + +def test_pickle_round_trip_preserves_filter(people: vx.VortexFile) -> None: + expr = (ve.column("age") > 26) & ve.ilike(ve.column("name"), "a%") + restored = cast(Expr, pickle.loads(pickle.dumps(expr))) + assert names(people, restored) == names(people, expr) == ["Alice"] + + +def test_deserialize_rejects_garbage() -> None: + with pytest.raises(ValueError): + _ = ve.deserialize(b"\xff\xff\xff\xff\xff\xff") + + +def test_serialize_is_stable() -> None: + expr = ve.column("age") > 21 + assert expr.serialize() == expr.serialize()