From 2f9458047048e8f7bdab2324877e38a841c49d59 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 14 Aug 2026 13:36:00 +0100 Subject: [PATCH] Expression constructors in python accept generators as an argument Signed-off-by: Robert Kruszewski --- vortex-python/src/expr/mod.rs | 43 ++++++++++++++++++++------------- vortex-python/test/test_expr.py | 25 +++++++++++++++++++ 2 files changed, 51 insertions(+), 17 deletions(-) diff --git a/vortex-python/src/expr/mod.rs b/vortex-python/src/expr/mod.rs index be3a0339143..8d76d20b051 100644 --- a/vortex-python/src/expr/mod.rs +++ b/vortex-python/src/expr/mod.rs @@ -209,6 +209,18 @@ fn field_names(fields: &Bound<'_, PyAny>) -> PyResult { .into()) } +/// Extract expressions from any iterable, so that a generator serves as well as a sequence. +/// +/// A `Vec` parameter would reject anything that is not a sequence, and these functions +/// are documented to take an iterable — which is what a caller building expressions in a +/// comprehension naturally has. +fn into_exprs(exprs: &Bound<'_, PyAny>) -> PyResult> { + exprs + .try_iter()? + .map(|expr| coerce_expression(&expr?)) + .collect() +} + /// 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::() { @@ -690,8 +702,8 @@ pub fn or_(left: PyIntoExpr, right: PyIntoExpr) -> PyExpr { /// :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) +pub fn and_collect(exprs: &Bound<'_, PyAny>) -> PyResult> { + Ok(expr::and_collect(into_exprs(exprs)?).map(PyExpr::from)) } /// Combine expressions with logical OR using a balanced tree. @@ -706,8 +718,8 @@ pub fn and_collect(exprs: Vec) -> Option { /// :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) +pub fn or_collect(exprs: &Bound<'_, PyAny>) -> PyResult> { + Ok(expr::or_collect(into_exprs(exprs)?).map(PyExpr::from)) } macro_rules! binary_fn { @@ -1068,7 +1080,7 @@ pub fn pack(fields: &Bound<'_, PyAny>, nullable: bool) -> PyResult { /// :class:`vortex.Expr` #[pyfunction] #[pyo3(signature = (exprs, *, duplicate_handling = "error"))] -pub fn merge(exprs: Vec, duplicate_handling: &str) -> PyResult { +pub fn merge(exprs: &Bound<'_, PyAny>, duplicate_handling: &str) -> PyResult { let duplicate_handling = match duplicate_handling.to_ascii_lowercase().as_str() { "error" => DuplicateHandling::Error, "rightmost" | "right_most" => DuplicateHandling::RightMost, @@ -1079,10 +1091,7 @@ pub fn merge(exprs: Vec, duplicate_handling: &str) -> PyResult PyExpr { /// ``` #[pyfunction] #[pyo3(signature = (when_then, else_value = None))] -pub fn case_when( - when_then: Vec<(PyIntoExpr, PyIntoExpr)>, - else_value: Option, -) -> PyResult { +pub fn case_when(when_then: &Bound<'_, PyAny>, else_value: Option) -> PyResult { + let when_then: Vec<(Expression, Expression)> = when_then + .try_iter()? + .map(|pair| { + let (condition, value): (Bound<'_, PyAny>, Bound<'_, PyAny>) = pair?.extract()?; + Ok((coerce_expression(&condition)?, coerce_expression(&value)?)) + }) + .collect::>()?; 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)), }) diff --git a/vortex-python/test/test_expr.py b/vortex-python/test/test_expr.py index c70810ba25a..f99ab9ffdad 100644 --- a/vortex-python/test/test_expr.py +++ b/vortex-python/test/test_expr.py @@ -149,6 +149,31 @@ def test_case_when_requires_a_pair() -> None: def test_collect_returns_none_when_empty() -> None: assert ve.and_collect([]) is None assert ve.or_collect([]) is None + assert ve.and_collect(expr for expr in ()) is None + assert ve.or_collect(expr for expr in ()) is None + + +def test_functions_taking_an_iterable_accept_a_generator() -> None: + # Building the expressions in a comprehension is the natural way to call these, so a generator + # has to work and not just a sequence. + columns = ["age", "id"] + assert str(ve.and_collect(ve.column(name) > 1 for name in columns)) == str( + ve.and_collect([ve.column("age") > 1, ve.column("id") > 1]) + ) + assert str(ve.or_collect(ve.column(name) > 1 for name in columns)) == str( + ve.or_collect([ve.column("age") > 1, ve.column("id") > 1]) + ) + assert str(ve.merge(ve.select([name]) for name in columns)) == str( + ve.merge([ve.select(["age"]), ve.select(["id"])]) + ) + assert str(ve.case_when((ve.column(name) > 1, name) for name in columns)) == str( + ve.case_when([(ve.column("age") > 1, "age"), (ve.column("id") > 1, "id")]) + ) + + +def test_case_when_rejects_an_empty_generator() -> None: + with pytest.raises(ValueError): + _ = ve.case_when(pair for pair in ()) # --------------------------------------------------------------------------------------