feat: make cast accept built-in Python types (#858)

Author: mesejo

python/datafusion/expr.py

@@ -28,7 +28,7 @@
     functions as functions_internal,
 )
 from datafusion.common import NullTreatment, RexType, DataTypeMap
-from typing import Any, Optional
+from typing import Any, Optional, Type
 import pyarrow as pa
 
 # The following are imported from the internal representation. We may choose to
@@ -372,8 +372,25 @@ def is_not_null(self) -> Expr:
         """Returns ``True`` if this expression is not null."""
         return Expr(self.expr.is_not_null())
 
-    def cast(self, to: pa.DataType[Any]) -> Expr:
+    _to_pyarrow_types = {
+        float: pa.float64(),
+        int: pa.int64(),
+        str: pa.string(),
+        bool: pa.bool_(),
+    }
+
+    def cast(
+        self, to: pa.DataType[Any] | Type[float] | Type[int] | Type[str] | Type[bool]
+    ) -> Expr:
         """Cast to a new data type."""
+        if not isinstance(to, pa.DataType):
+            try:
+                to = self._to_pyarrow_types[to]
+            except KeyError:
+                raise TypeError(
+                    "Expected instance of pyarrow.DataType or builtins.type"
+                )
+
         return Expr(self.expr.cast(to))
 
     def rex_type(self) -> RexType:

python/datafusion/tests/test_functions.py

@@ -44,8 +44,9 @@ def df():
                     datetime(2020, 7, 2),
                 ]
             ),
+            pa.array([False, True, True]),
         ],
-        names=["a", "b", "c", "d"],
+        names=["a", "b", "c", "d", "e"],
     )
     return ctx.create_dataframe([[batch]])
 
@@ -63,15 +64,14 @@ def test_named_struct(df):
     )
 
     expected = """DataFrame()
-+-------+---+---------+------------------------------+
-| a     | b | c       | d                            |
-+-------+---+---------+------------------------------+
-| Hello | 4 | hello   | {a: Hello, b: 4, c: hello }  |
-| World | 5 |  world  | {a: World, b: 5, c:  world } |
-| !     | 6 |  !      | {a: !, b: 6, c:  !}          |
-+-------+---+---------+------------------------------+
++-------+---+---------+------------------------------+-------+
+| a     | b | c       | d                            | e     |
++-------+---+---------+------------------------------+-------+
+| Hello | 4 | hello   | {a: Hello, b: 4, c: hello }  | false |
+| World | 5 |  world  | {a: World, b: 5, c:  world } | true  |
+| !     | 6 |  !      | {a: !, b: 6, c:  !}          | true  |
++-------+---+---------+------------------------------+-------+
 """.strip()
-
     assert str(df) == expected
 
 
@@ -978,3 +978,22 @@ def test_binary_string_functions(df):
     assert pa.array(result.column(1)).cast(pa.string()) == pa.array(
         ["Hello", "World", "!"]
     )
+
+
+@pytest.mark.parametrize(
+    "python_datatype, name, expected",
+    [
+        pytest.param(bool, "e", pa.bool_(), id="bool"),
+        pytest.param(int, "b", pa.int64(), id="int"),
+        pytest.param(float, "b", pa.float64(), id="float"),
+        pytest.param(str, "b", pa.string(), id="str"),
+    ],
+)
+def test_cast(df, python_datatype, name: str, expected):
+    df = df.select(
+        column(name).cast(python_datatype).alias("actual"),
+        column(name).cast(expected).alias("expected"),
+    )
+    result = df.collect()
+    result = result[0]
+    assert result.column(0) == result.column(1)