fix: update arrow_cast function to accept string type for data_type parameter

kosiew · kosiew · commit 20cc78150e86 · 2024-11-26T17:47:10.000+08:00
diff --git a/python/datafusion/functions.py b/python/datafusion/functions.py
@@ -1100,9 +1100,9 @@ def arrow_typeof(arg: Expr) -> Expr:
     return Expr(f.arrow_typeof(arg.expr))
 
 
-def arrow_cast(expr: Expr, data_type: Expr) -> Expr:
+def arrow_cast(expr: Expr, data_type: str) -> Expr:
     """Casts an expression to a specified data type."""
-    return Expr(f.arrow_cast(expr.expr, data_type.expr))
+    return Expr(f.arrow_cast(expr.expr, data_type))
 
 
 def random() -> Expr:
diff --git a/python/tests/test_functions.py b/python/tests/test_functions.py
@@ -907,10 +907,8 @@ def test_temporal_functions(df):
 
 def test_arrow_cast(df):
     df = df.select(
-        f.arrow_cast(column("a"), literal("Float64").cast(pa.string())).alias(
-            "a_as_float"
-        ),
-        f.arrow_cast(column("a"), literal("Int32").cast(pa.string())).alias("a_as_int"),
+        f.arrow_cast(column("a"), "Float64").alias("a_as_float"),
+        f.arrow_cast(column("a"), "Int32").alias("a_as_int"),
     )
     result = df.collect()
     assert len(result) == 1
