Support pyarrow DataType in arrow_cast

Allow arrow_cast to accept a pyarrow DataType in addition to str and
Expr. The DataType is converted to its string representation before
being passed to DataFusion. Adds test coverage for the new input type.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timsaucer

python/datafusion/functions.py

@@ -2634,10 +2634,10 @@ def arrow_typeof(arg: Expr) -> Expr:
     return Expr(f.arrow_typeof(arg.expr))
 
 
-def arrow_cast(expr: Expr, data_type: Expr | str) -> Expr:
+def arrow_cast(expr: Expr, data_type: Expr | str | pa.DataType) -> Expr:
     """Casts an expression to a specified data type.
 
-    The ``data_type`` can be a string or an ``Expr``.
+    The ``data_type`` can be a string, a ``pyarrow.DataType``, or an ``Expr``.
 
     Examples:
         >>> ctx = dfn.SessionContext()
@@ -2647,7 +2647,18 @@ def arrow_cast(expr: Expr, data_type: Expr | str) -> Expr:
         ... )
         >>> result.collect_column("c")[0].as_py()
         1.0
+
+        >>> import pyarrow as pa
+        >>> result = df.select(
+        ...     dfn.functions.arrow_cast(
+        ...         dfn.col("a"), data_type=pa.float64()
+        ...     ).alias("c")
+        ... )
+        >>> result.collect_column("c")[0].as_py()
+        1.0
     """
+    if isinstance(data_type, pa.DataType):
+        data_type = str(data_type)
     if isinstance(data_type, str):
         data_type = Expr.string_literal(data_type)
     return Expr(f.arrow_cast(expr.expr, data_type.expr))

python/tests/test_functions.py

@@ -1302,6 +1302,19 @@ def test_arrow_cast(df):
     assert result.column(1) == pa.array([4, 5, 6], type=pa.int32())
 
 
+def test_arrow_cast_with_pyarrow_type(df):
+    df = df.select(
+        f.arrow_cast(column("b"), pa.float64()).alias("b_as_float"),
+        f.arrow_cast(column("b"), pa.int32()).alias("b_as_int"),
+        f.arrow_cast(column("b"), pa.string()).alias("b_as_str"),
+    )
+    result = df.collect()[0]
+
+    assert result.column(0) == pa.array([4.0, 5.0, 6.0], type=pa.float64())
+    assert result.column(1) == pa.array([4, 5, 6], type=pa.int32())
+    assert result.column(2) == pa.array(["4", "5", "6"], type=pa.string())
+
+
 def test_case(df):
     df = df.select(
         f.case(column("b")).when(literal(4), literal(10)).otherwise(literal(8)),