Accept str for field name and type parameters in scalar functions

Allow arrow_cast, get_field, and union_extract to accept plain str
arguments instead of requiring Expr wrappers. Also improve
arrow_metadata test coverage and fix parameter shadowing.

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

Author: timsaucer

python/datafusion/functions.py

@@ -2533,19 +2533,20 @@ 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: Expr | str) -> Expr:
     """Casts an expression to a specified data type.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [1]})
-        >>> data_type = dfn.string_literal("Float64")
         >>> result = df.select(
-        ...     dfn.functions.arrow_cast(dfn.col("a"), data_type).alias("c")
+        ...     dfn.functions.arrow_cast(dfn.col("a"), "Float64").alias("c")
         ... )
         >>> result.collect_column("c")[0].as_py()
         1.0
     """
+    if isinstance(data_type, str):
+        data_type = Expr.string_literal(data_type)
     return Expr(f.arrow_cast(expr.expr, data_type.expr))
 
 
@@ -2561,11 +2562,10 @@ def arrow_metadata(*args: Expr) -> Expr:
     Returns:
         A Map of metadata or a specific metadata value.
     """
-    args = [arg.expr for arg in args]
-    return Expr(f.arrow_metadata(*args))
+    return Expr(f.arrow_metadata(*[arg.expr for arg in args]))
 
 
-def get_field(expr: Expr, name: Expr) -> Expr:
+def get_field(expr: Expr, name: Expr | str) -> Expr:
     """Extracts a field from a struct or map by name.
 
     Args:
@@ -2575,10 +2575,12 @@ def get_field(expr: Expr, name: Expr) -> Expr:
     Returns:
         The value of the named field.
     """
+    if isinstance(name, str):
+        name = Expr.string_literal(name)
     return Expr(f.get_field(expr.expr, name.expr))
 
 
-def union_extract(union_expr: Expr, field_name: Expr) -> Expr:
+def union_extract(union_expr: Expr, field_name: Expr | str) -> Expr:
     """Extracts a value from a union type by field name.
 
     Returns the value of the named field if it is the currently selected
@@ -2591,6 +2593,8 @@ def union_extract(union_expr: Expr, field_name: Expr) -> Expr:
     Returns:
         The extracted value or NULL.
     """
+    if isinstance(field_name, str):
+        field_name = Expr.string_literal(field_name)
     return Expr(f.union_extract(union_expr.expr, field_name.expr))
 
 

python/tests/test_functions.py

@@ -1143,11 +1143,8 @@ def test_make_time(df):
 
 def test_arrow_cast(df):
     df = df.select(
-        # we use `string_literal` to return utf8 instead of `literal` which returns
-        # utf8view because datafusion.arrow_cast expects a utf8 instead of utf8view
-        # https://github.com/apache/datafusion/blob/86740bfd3d9831d6b7c1d0e1bf4a21d91598a0ac/datafusion/functions/src/core/arrow_cast.rs#L179
-        f.arrow_cast(column("b"), string_literal("Float64")).alias("b_as_float"),
-        f.arrow_cast(column("b"), string_literal("Int32")).alias("b_as_int"),
+        f.arrow_cast(column("b"), "Float64").alias("b_as_float"),
+        f.arrow_cast(column("b"), "Int32").alias("b_as_int"),
     )
     result = df.collect()
     assert len(result) == 1
@@ -1482,20 +1479,35 @@ def test_get_field(df):
         ),
     )
     result = df.select(
-        f.get_field(column("s"), string_literal("x")).alias("x_val"),
-        f.get_field(column("s"), string_literal("y")).alias("y_val"),
+        f.get_field(column("s"), "x").alias("x_val"),
+        f.get_field(column("s"), "y").alias("y_val"),
     ).collect()[0]
 
     assert result.column(0) == pa.array(["Hello", "World", "!"], type=pa.string_view())
     assert result.column(1) == pa.array([4, 5, 6])
 
 
-def test_arrow_metadata(df):
+def test_arrow_metadata():
+    ctx = SessionContext()
+    field = pa.field("val", pa.int64(), metadata={"key1": "value1", "key2": "value2"})
+    schema = pa.schema([field])
+    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], schema=schema)
+    df = ctx.create_dataframe([[batch]])
+
+    # One-argument form: returns a Map of all metadata key-value pairs
     result = df.select(
-        f.arrow_metadata(column("a")).alias("meta"),
+        f.arrow_metadata(column("val")).alias("meta"),
     ).collect()[0]
-    # The metadata column should be returned as a map type (possibly empty)
     assert result.column(0).type == pa.map_(pa.utf8(), pa.utf8())
+    meta = result.column(0)[0].as_py()
+    assert ("key1", "value1") in meta
+    assert ("key2", "value2") in meta
+
+    # Two-argument form: returns the value for a specific metadata key
+    result = df.select(
+        f.arrow_metadata(column("val"), string_literal("key1")).alias("meta_val"),
+    ).collect()[0]
+    assert result.column(0)[0].as_py() == "value1"
 
 
 def test_version():
@@ -1535,7 +1547,5 @@ def test_union_extract():
     arr = pa.UnionArray.from_dense(types, offsets, children, ["int", "str"], [0, 1])
     df = ctx.create_dataframe([[pa.RecordBatch.from_arrays([arr], names=["u"])]])
 
-    result = df.select(
-        f.union_extract(column("u"), string_literal("int")).alias("val")
-    ).collect()[0]
+    result = df.select(f.union_extract(column("u"), "int").alias("val")).collect()[0]
     assert result.column(0).to_pylist() == [1, None, 2]