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

kosiew · kosiew · commit 9e1ced7fb56c · 2024-11-27T15:58:56.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/src/functions.rs b/src/functions.rs
@@ -144,6 +144,12 @@ fn concat_ws(sep: String, args: Vec<PyExpr>) -> PyResult<PyExpr> {
     Ok(functions::string::expr_fn::concat_ws(lit(sep), args).into())
 }
 
+/// Casts an expression to a specified data type.
+#[pyfunction]
+fn arrow_cast(expr: PyExpr, data_type: &str) -> PyResult<PyExpr> {
+    Ok(functions::expr_fn::arrow_cast(expr.expr, data_type.to_string()).into())
+}
+
 #[pyfunction]
 #[pyo3(signature = (values, regex, flags=None))]
 fn regexp_like(values: PyExpr, regex: PyExpr, flags: Option<PyExpr>) -> PyResult<PyExpr> {
@@ -563,7 +569,6 @@ expr_fn_vec!(r#struct); // Use raw identifier since struct is a keyword
 expr_fn_vec!(named_struct);
 expr_fn!(from_unixtime, unixtime);
 expr_fn!(arrow_typeof, arg_1);
-expr_fn!(arrow_cast, arg_1 datatype);
 expr_fn!(random);
 
 // Array Functions
