refactor(spark): reshape varargs to match pyspark signatures

Replace generic ``*args`` with explicit pyspark-style signatures:
- json_tuple(col, *fields) — first positional is the JSON expr
- format_string(format, *cols) — `format` is the printf template; a plain
  ``str`` is auto-promoted to a literal
- parse_url(url, partToExtract, key=None) — `key` is optional and only
  meaningful with ``partToExtract='QUERY'``
- try_parse_url(url, partToExtract, key=None) — same shape
- url_decode(str), try_url_decode(str), url_encode(str) — single-argument
  forms (multi-arg calls were always semantically wrong)

Tests cover the three-arg parse_url path and the plain-str format_string
auto-promotion.

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

Author: timsaucer

python/datafusion/functions/spark.py

@@ -958,7 +958,7 @@ def xxhash64(*cols: Expr) -> Expr:
 # ---------------------------------------------------------------------------
 
 
-def json_tuple(*args: Expr) -> Expr:
+def json_tuple(col: Expr, *fields: Expr) -> Expr:
     """Spark ``json_tuple``: extract top-level fields from a JSON string.
 
     Examples:
@@ -972,7 +972,7 @@ def json_tuple(*args: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         {'c0': '1', 'c1': 'x'}
     """
-    return Expr(_f.json_tuple(*[a.expr for a in args]))
+    return Expr(_f.json_tuple(col.expr, *[f.expr for f in fields]))
 
 
 # ---------------------------------------------------------------------------
@@ -1439,23 +1439,25 @@ def luhn_check(col: Expr) -> Expr:
     return Expr(_f.luhn_check(col.expr))
 
 
-def format_string(*args: Expr) -> Expr:
+def format_string(format: str | Expr, *cols: Expr) -> Expr:
     """Spark ``format_string``: printf-style format string.
 
-    First arg is the format, remaining args are values to substitute.
+    ``format`` is the printf-style template (a plain ``str`` is auto-promoted
+    to a literal expression); remaining args are values to substitute.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"x": [1]})
         >>> r = df.select(
         ...     dfn.functions.spark.format_string(
-        ...         dfn.lit("%d-%s"), dfn.lit(42), dfn.lit("hi")
+        ...         "%d-%s", dfn.lit(42), dfn.lit("hi")
         ...     ).alias("v")
         ... )
         >>> r.collect_column("v")[0].as_py()
         '42-hi'
     """
-    return Expr(_f.format_string(*[a.expr for a in args]))
+    fmt_expr = format if isinstance(format, Expr) else Expr.literal(format)
+    return Expr(_f.format_string(fmt_expr.expr, *[c.expr for c in cols]))
 
 
 def space(col: Expr) -> Expr:
@@ -1555,9 +1557,17 @@ def make_valid_utf8(str: Expr) -> Expr:
 # ---------------------------------------------------------------------------
 
 
-def parse_url(*args: Expr) -> Expr:
+def parse_url(
+    url: Expr,
+    partToExtract: Expr,  # noqa: N803
+    key: Expr | None = None,
+) -> Expr:
     """Spark ``parse_url``: extract a part from a URL; errors on invalid URLs.
 
+    ``partToExtract`` is one of ``"HOST"``, ``"PATH"``, ``"QUERY"``,
+    ``"REF"``, ``"PROTOCOL"``, ``"FILE"``, ``"AUTHORITY"``, ``"USERINFO"``.
+    Pass ``key`` only with ``"QUERY"`` to extract a single parameter.
+
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"x": [1]})
@@ -1568,11 +1578,27 @@ def parse_url(*args: Expr) -> Expr:
         ... )
         >>> r.collect_column("v")[0].as_py()
         'example.com'
+
+        >>> r = df.select(
+        ...     dfn.functions.spark.parse_url(
+        ...         dfn.lit("http://example.com/path?q=1"),
+        ...         dfn.lit("QUERY"),
+        ...         key=dfn.lit("q"),
+        ...     ).alias("v")
+        ... )
+        >>> r.collect_column("v")[0].as_py()
+        '1'
     """
-    return Expr(_f.parse_url(*[a.expr for a in args]))
+    if key is None:
+        return Expr(_f.parse_url(url.expr, partToExtract.expr))
+    return Expr(_f.parse_url(url.expr, partToExtract.expr, key.expr))
 
 
-def try_parse_url(*args: Expr) -> Expr:
+def try_parse_url(
+    url: Expr,
+    partToExtract: Expr,  # noqa: N803
+    key: Expr | None = None,
+) -> Expr:
     """Spark ``try_parse_url``: like ``parse_url`` but returns NULL on invalid URLs.
 
     Examples:
@@ -1586,10 +1612,12 @@ def try_parse_url(*args: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         'example.com'
     """
-    return Expr(_f.try_parse_url(*[a.expr for a in args]))
+    if key is None:
+        return Expr(_f.try_parse_url(url.expr, partToExtract.expr))
+    return Expr(_f.try_parse_url(url.expr, partToExtract.expr, key.expr))
 
 
-def url_decode(*args: Expr) -> Expr:
+def url_decode(str: Expr) -> Expr:
     """Spark ``url_decode``: decode an application/x-www-form-urlencoded string.
 
     Examples:
@@ -1600,10 +1628,10 @@ def url_decode(*args: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         'a b'
     """
-    return Expr(_f.url_decode(*[a.expr for a in args]))
+    return Expr(_f.url_decode(str.expr))
 
 
-def try_url_decode(*args: Expr) -> Expr:
+def try_url_decode(str: Expr) -> Expr:
     """Spark ``try_url_decode``: like ``url_decode``; returns NULL on invalid input.
 
     Examples:
@@ -1614,10 +1642,10 @@ def try_url_decode(*args: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         'a b'
     """
-    return Expr(_f.try_url_decode(*[a.expr for a in args]))
+    return Expr(_f.try_url_decode(str.expr))
 
 
-def url_encode(*args: Expr) -> Expr:
+def url_encode(str: Expr) -> Expr:
     """Spark ``url_encode``: encode a string in application/x-www-form-urlencoded.
 
     Examples:
@@ -1628,7 +1656,7 @@ def url_encode(*args: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         'a+b'
     """
-    return Expr(_f.url_encode(*[a.expr for a in args]))
+    return Expr(_f.url_encode(str.expr))
 
 
 __all__ = [

python/tests/test_spark_functions.py

@@ -278,6 +278,22 @@ def test_like_escape_raises():
         spark.ilike(lit("a"), lit("a"), escapeChar="\\")
 
 
+def test_parse_url_three_arg():
+    """parse_url(url, partToExtract, key=...) extracts query parameters."""
+    ctx = SessionContext()
+    df = ctx.from_pydict({"x": [1]})
+    url = lit("http://example.com/p?q=hello&n=1")
+    assert _val(df, spark.parse_url(url, lit("QUERY"), key=lit("q"))) == "hello"
+    assert _val(df, spark.try_parse_url(url, lit("QUERY"), key=lit("n"))) == "1"
+
+
+def test_format_string_plain_str_format():
+    """format_string accepts a plain str format that is auto-promoted to lit."""
+    ctx = SessionContext()
+    df = ctx.from_pydict({"x": [1]})
+    assert _val(df, spark.format_string("%d-%s", lit(42), lit("hi"))) == "42-hi"
+
+
 def test_aggregate_positional_compat():
     """Pyspark-style positional calls still work after the rename to ``col``."""
     ctx = SessionContext()