feat: cast expression (#5389)

Without this, it is impossible to compare two columns with different
bit-widths.

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

Author: danking

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

@@ -33,3 +33,4 @@ def root() -> 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: ...

vortex-python/python/vortex/expr.py

@@ -2,6 +2,6 @@
 # SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 
-from ._lib.expr import Expr, and_, column, literal, not_, root  # pyright: ignore[reportMissingModuleSource]
+from ._lib.expr import Expr, and_, cast, column, literal, not_, root  # pyright: ignore[reportMissingModuleSource]
 
-__all__ = ["Expr", "column", "literal", "root", "not_", "and_"]
+__all__ = ["Expr", "column", "literal", "root", "not_", "and_", "cast"]

vortex-python/src/expr/mod.rs

@@ -7,6 +7,7 @@ use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
 use pyo3::types::*;
 use vortex::dtype::{DType, Nullability, PType};
+use vortex::expr;
 use vortex::expr::{Binary, Expression, GetItem, Operator, VTableExt, and, lit, not};
 
 use crate::arrays::PyArrayRef;
@@ -25,6 +26,7 @@ pub(crate) fn init(py: Python, parent: &Bound<PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(literal, &m)?)?;
     m.add_function(wrap_pyfunction!(not_, &m)?)?;
     m.add_function(wrap_pyfunction!(and_, &m)?)?;
+    m.add_function(wrap_pyfunction!(cast, &m)?)?;
     m.add_class::<PyExpr>()?;
 
     Ok(())
@@ -256,7 +258,7 @@ pub fn literal<'py>(
 #[pyfunction]
 pub fn root() -> PyExpr {
     PyExpr {
-        inner: vortex::expr::root(),
+        inner: expr::root(),
     }
 }
 
@@ -290,7 +292,7 @@ pub fn column<'py>(name: &Bound<'py, PyString>) -> PyResult<Bound<'py, PyExpr>>
     Bound::new(
         py,
         PyExpr {
-            inner: vortex::expr::get_item(name, vortex::expr::root()),
+            inner: expr::get_item(name, expr::root()),
         },
     )
 }
@@ -342,7 +344,10 @@ pub fn not_(child: PyExpr) -> PyResult<PyExpr> {
 ///
 /// Parameters
 /// ----------
-/// child : :class:`Any`
+/// left : :class:`Expr`
+///     A boolean expression.
+///
+/// right : :class:`Expr`
 ///     A boolean expression.
 ///
 /// Returns
@@ -364,3 +369,41 @@ pub fn and_(left: PyExpr, right: PyExpr) -> PyResult<PyExpr> {
         inner: and(left.inner, right.inner),
     })
 }
+
+/// Cast an expression to a compatible type.
+///
+/// Parameters
+/// ----------
+/// child : :class:`Expr`
+///     The expression to cast.
+///
+/// Returns
+/// -------
+/// :class:`vortex.Expr`
+///
+/// Examples
+/// --------
+///
+/// Cast to a wider integer type:
+///
+/// ```python
+/// >>> import vortex.expr as ve
+/// >>> import vortex as vx
+/// >>> ve.cast(ve.literal(vx.int_(8), 1), vx.int_(16))
+/// <vortex.Expr object at ...>
+/// ```
+///
+/// Cast to a wider floating-point type:
+///
+/// ```python
+/// >>> import vortex.expr as ve
+/// >>> import vortex as vx
+/// >>> ve.cast(ve.literal(vx.float_(16), 3.145), vx.float_(64))
+/// <vortex.Expr object at ...>
+/// ```
+#[pyfunction]
+pub fn cast(child: PyExpr, dtype: PyDType) -> PyResult<PyExpr> {
+    Ok(PyExpr {
+        inner: expr::cast(child.into_inner(), dtype.into_inner()),
+    })
+}

vortex-python/test/test_scan.py

@@ -8,6 +8,7 @@
 import pytest
 
 import vortex as vx
+import vortex.expr as ve
 from vortex.scan import RepeatedScan
 
 
@@ -20,14 +21,19 @@ def record(x: int, columns: list[str] | set[str] | None = None) -> dict[str, int
 
 
 @pytest.fixture(scope="session")
-def vxscan(tmpdir_factory) -> vx.RepeatedScan:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
+def vxscan(vxfile: vx.VortexFile) -> vx.RepeatedScan:
+    return vxfile.to_repeated_scan()
+
+
+@pytest.fixture(scope="session")
+def vxfile(tmpdir_factory) -> vx.VortexFile:  # pyright: ignore[reportUnknownParameterType, reportMissingParameterType]
     fname = tmpdir_factory.mktemp("data") / "foo.vortex"  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
 
     if not os.path.exists(fname):  # pyright: ignore[reportUnknownArgumentType]
         a = pa.array([record(x) for x in range(1_000)])
         arr = vx.compress(vx.array(a))
         vx.io.write(arr, str(fname))  # pyright: ignore[reportUnknownArgumentType]
-    return vx.open(str(fname)).to_repeated_scan()  # pyright: ignore[reportUnknownArgumentType]
+    return vx.open(str(fname))  # pyright: ignore[reportUnknownArgumentType]
 
 
 def test_execute(vxscan: RepeatedScan):
@@ -50,3 +56,11 @@ def test_scalar_at(vxscan: RepeatedScan):
         "bool": True,
         "float": math.sqrt(10),
     }
+
+
+def test_scan_with_cast(vxfile: vx.VortexFile):
+    actual = vxfile.scan(expr=ve.cast(ve.column("index"), vx.int_(16)) == ve.literal(vx.int_(16), 1)).read_all()
+    expected = pa.array(
+        [{"index": 1, "string": pa.scalar("1", pa.string_view()), "bool": False, "float": math.sqrt(1)}]
+    )
+    assert str(actual.to_arrow_array()) == str(expected)