Accept str for key parameter in arrow_metadata for consistency

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

Author: timsaucer

python/datafusion/functions.py

@@ -2550,19 +2550,24 @@ def arrow_cast(expr: Expr, data_type: Expr | str) -> Expr:
     return Expr(f.arrow_cast(expr.expr, data_type.expr))
 
 
-def arrow_metadata(*args: Expr) -> Expr:
+def arrow_metadata(expr: Expr, key: Expr | str | None = None) -> Expr:
     """Returns the metadata of the input expression.
 
     If called with one argument, returns a Map of all metadata key-value pairs.
     If called with two arguments, returns the value for the specified metadata key.
 
     Args:
-        args: An expression, optionally followed by a metadata key string.
+        expr: An expression whose metadata to retrieve.
+        key: Optional metadata key to look up. Can be a string or an Expr.
 
     Returns:
         A Map of metadata or a specific metadata value.
     """
-    return Expr(f.arrow_metadata(*[arg.expr for arg in args]))
+    if key is None:
+        return Expr(f.arrow_metadata(expr.expr))
+    if isinstance(key, str):
+        key = Expr.string_literal(key)
+    return Expr(f.arrow_metadata(expr.expr, key.expr))
 
 
 def get_field(expr: Expr, name: Expr | str) -> Expr:

python/tests/test_functions.py

@@ -20,7 +20,7 @@
 import numpy as np
 import pyarrow as pa
 import pytest
-from datafusion import SessionContext, column, literal, string_literal
+from datafusion import SessionContext, column, literal
 from datafusion import functions as f
 
 np.seterr(invalid="ignore")
@@ -1505,7 +1505,7 @@ def test_arrow_metadata():
 
     # 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"),
+        f.arrow_metadata(column("val"), "key1").alias("meta_val"),
     ).collect()[0]
     assert result.column(0)[0].as_py() == "value1"