feat(spark): add pyspark aliases and optional-length substr

- Add getbit, dateadd, datediff, datepart, sha, ceiling, printf,
  char_length, and character_length as aliases of their Spark primaries.
- Add substr, whose len argument is optional as in pyspark. It calls the
  UDF directly because upstream expr_fn::substring always takes a length.
- Rename the spark.last_day parameter from col to date to match pyspark,
  with an upgrade-guide note for keyword callers.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: timsaucer

crates/core/src/spark_functions.rs

@@ -254,6 +254,18 @@ spark_udf_vec!(format_string, udf::string::format_string);
 spark_expr_fn!(quote, arg1);
 spark_expr_fn!(space, arg1);
 spark_expr_fn!(substring, str pos length);
+/// `substr(str, pos, len=None)`. Upstream `expr_fn::substring` always takes a
+/// length, so call the UDF directly to allow the two-argument form.
+#[pyfunction]
+#[pyo3(signature = (str, pos, len=None))]
+fn substr(str: PyExpr, pos: PyExpr, len: Option<PyExpr>) -> PyExpr {
+    let args: Vec<Expr> = [Some(str), Some(pos), len]
+        .into_iter()
+        .flatten()
+        .map(Into::into)
+        .collect();
+    Expr::ScalarFunction(ScalarFunction::new_udf(udf::string::substring(), args)).into()
+}
 spark_expr_fn!(unbase64, str);
 spark_expr_fn!(soundex, str);
 spark_expr_fn!(is_valid_utf8, str);
@@ -372,6 +384,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(format_string))?;
     m.add_wrapped(wrap_pyfunction!(space))?;
     m.add_wrapped(wrap_pyfunction!(substring))?;
+    m.add_wrapped(wrap_pyfunction!(substr))?;
     m.add_wrapped(wrap_pyfunction!(unbase64))?;
     m.add_wrapped(wrap_pyfunction!(soundex))?;
     m.add_wrapped(wrap_pyfunction!(is_valid_utf8))?;

docs/source/user-guide/upgrade-guides.md

@@ -218,6 +218,17 @@ f.bit_and(column("a"), filter=my_filter)  # after
 Passing `filter` to `mean` previously raised a `TypeError`, whether passed
 positionally or by keyword; it now works.
 
+### `spark.last_day` renamed its parameter
+
+The parameter of {py:func}`datafusion.functions.spark.last_day` is now named
+`date`, matching `pyspark.sql.functions.last_day`. Positional calls are
+unaffected; update any call passing it by keyword.
+
+```python
+spark.last_day(col=d)  # before
+spark.last_day(date=d)  # after
+```
+
 ### Changes to the `datafusion-python-util` crate
 
 Extension libraries written in Rust usually depend on the

python/datafusion/functions/spark.py

@@ -356,6 +356,15 @@ def bit_get(col: Expr, pos: Expr | str) -> Expr:
     return Expr(_f.bit_get(col.expr, _to_raw_expr(pos)))
 
 
+def getbit(col: Expr, pos: Expr | str) -> Expr:
+    """Spark ``getbit``: returns the bit (0 or 1) at ``pos``.
+
+    See Also:
+        This is an alias for :py:func:`bit_get`.
+    """
+    return bit_get(col, pos)
+
+
 def bit_count(col: Expr) -> Expr:
     """Spark ``bit_count``: number of bits set in the integer's binary form.
 
@@ -539,6 +548,15 @@ def date_add(start: Expr, days: Expr | int) -> Expr:
     return Expr(_f.date_add(start.expr, _coerce_i32(days).expr))
 
 
+def dateadd(start: Expr, days: Expr | int) -> Expr:
+    """Spark ``dateadd``: date + N days.
+
+    See Also:
+        This is an alias for :py:func:`date_add`.
+    """
+    return date_add(start, days)
+
+
 def date_sub(start: Expr, days: Expr | int) -> Expr:
     """Spark ``date_sub``: date - N days.
 
@@ -627,7 +645,7 @@ def second(col: Expr) -> Expr:
     return Expr(_f.second(col.expr))
 
 
-def last_day(col: Expr) -> Expr:
+def last_day(date: Expr) -> Expr:
     """Spark ``last_day``: last day of the month containing the date.
 
     Examples:
@@ -640,7 +658,7 @@ def last_day(col: Expr) -> Expr:
         >>> r.collect_column("v")[0].as_py()
         datetime.date(2020, 1, 31)
     """
-    return Expr(_f.last_day(col.expr))
+    return Expr(_f.last_day(date.expr))
 
 
 def make_dt_interval(
@@ -754,6 +772,15 @@ def date_diff(end: Expr, start: Expr) -> Expr:
     return Expr(_f.date_diff(end.expr, start.expr))
 
 
+def datediff(end: Expr, start: Expr) -> Expr:
+    """Spark ``datediff``: number of days from ``start`` to ``end``.
+
+    See Also:
+        This is an alias for :py:func:`date_diff`.
+    """
+    return date_diff(end, start)
+
+
 def date_trunc(format: Expr | str, timestamp: Expr) -> Expr:
     """Spark ``date_trunc``: truncate timestamp to unit ``fmt``.
 
@@ -832,6 +859,15 @@ def date_part(field: Expr | str, source: Expr) -> Expr:
     return Expr(_f.date_part(coerce_to_expr(field).expr, source.expr))
 
 
+def datepart(field: Expr | str, source: Expr) -> Expr:
+    """Spark ``datepart``: extract ``field`` from a date/time/timestamp.
+
+    See Also:
+        This is an alias for :py:func:`date_part`.
+    """
+    return date_part(field, source)
+
+
 def from_utc_timestamp(timestamp: Expr, tz: Expr | str) -> Expr:
     """Spark ``from_utc_timestamp``: interpret ``ts`` as UTC, convert to ``tz``.
 
@@ -991,6 +1027,15 @@ def sha1(col: Expr) -> Expr:
     return Expr(_f.sha1(col.expr))
 
 
+def sha(col: Expr) -> Expr:
+    """Spark ``sha``: SHA-1 hash as a hex string.
+
+    See Also:
+        This is an alias for :py:func:`sha1`.
+    """
+    return sha1(col)
+
+
 def sha2(col: Expr, numBits: Expr | int) -> Expr:  # noqa: N803
     """Spark ``sha2``: SHA-2 family hash (224, 256, 384, 512). Bit length 0 = 256.
 
@@ -1175,6 +1220,15 @@ def ceil(col: Expr) -> Expr:
     return Expr(_f.ceil(col.expr))
 
 
+def ceiling(col: Expr) -> Expr:
+    """Spark ``ceiling``: smallest integer ≥ arg.
+
+    See Also:
+        This is an alias for :py:func:`ceil`.
+    """
+    return ceil(col)
+
+
 def expm1(col: Expr) -> Expr:
     """Spark ``expm1``: exp(arg) - 1.
 
@@ -1562,6 +1616,24 @@ def length(col: Expr) -> Expr:
     return Expr(_f.length(col.expr))
 
 
+def character_length(col: Expr) -> Expr:
+    """Spark ``character_length``: character length of a string, or bytes of binary.
+
+    See Also:
+        This is an alias for :py:func:`length`.
+    """
+    return length(col)
+
+
+def char_length(col: Expr) -> Expr:
+    """Spark ``char_length``: character length of a string, or bytes of binary.
+
+    See Also:
+        This is an alias for :py:func:`length`.
+    """
+    return length(col)
+
+
 def like(
     str: Expr,
     pattern: Expr | str,
@@ -1639,6 +1711,15 @@ def format_string(format: str | Expr, *cols: Expr) -> Expr:
     return Expr(_f.format_string(fmt_expr.expr, *[c.expr for c in cols]))
 
 
+def printf(format: str | Expr, *cols: Expr) -> Expr:
+    """Spark ``printf``: printf-style format string.
+
+    See Also:
+        This is an alias for :py:func:`format_string`.
+    """
+    return format_string(format, *cols)
+
+
 def space(col: Expr | int) -> Expr:
     """Spark ``space``: string of n spaces.
 
@@ -1673,6 +1754,28 @@ def substring(str: Expr, pos: Expr | int, len: Expr | int) -> Expr:
     )
 
 
+def substr(str: Expr, pos: Expr | int, len: Expr | int | None = None) -> Expr:
+    """Spark ``substr``: 1-indexed substring, to the end when ``len`` is omitted.
+
+    Same as :py:func:`substring` except that ``len`` is optional. ``pos`` and
+    ``len`` accept native ``int`` values or :class:`Expr`.
+
+    Examples:
+        >>> ctx = dfn.SessionContext()
+        >>> df = ctx.from_pydict({"x": [1]})
+        >>> r = df.select(dfn.functions.spark.substr(dfn.lit("hello"), 2).alias("v"))
+        >>> r.collect_column("v")[0].as_py()
+        'ello'
+
+        >>> r = df.select(
+        ...     dfn.functions.spark.substr(dfn.lit("hello"), 2, len=3).alias("v"))
+        >>> r.collect_column("v")[0].as_py()
+        'ell'
+    """
+    len_raw = coerce_to_expr(len).expr if len is not None else None
+    return Expr(_f.substr(str.expr, coerce_to_expr(pos).expr, len_raw))
+
+
 def unbase64(col: Expr) -> Expr:
     """Spark ``unbase64``: decode a base64 string to binary.
 
@@ -1867,7 +1970,10 @@ def url_encode(str: Expr) -> Expr:
     "bitmap_count",
     "bitwise_not",
     "ceil",
+    "ceiling",
     "char",
+    "char_length",
+    "character_length",
     "collect_list",
     "collect_set",
     "concat",
@@ -1880,12 +1986,16 @@ def url_encode(str: Expr) -> Expr:
     "date_part",
     "date_sub",
     "date_trunc",
+    "dateadd",
+    "datediff",
+    "datepart",
     "elt",
     "expm1",
     "factorial",
     "floor",
     "format_string",
     "from_utc_timestamp",
+    "getbit",
     "hex",
     "hour",
     "hypot",
@@ -1914,11 +2024,13 @@ def url_encode(str: Expr) -> Expr:
     "pmod",
     "pow",
     "power",
+    "printf",
     "quote",
     "rint",
     "round",
     "sec",
     "second",
+    "sha",
     "sha1",
     "sha2",
     "shiftleft",
@@ -1932,6 +2044,7 @@ def url_encode(str: Expr) -> Expr:
     "space",
     "spark_cast",
     "str_to_map",
+    "substr",
     "substring",
     "time_trunc",
     "to_utc_timestamp",

python/tests/test_spark_functions.py

@@ -500,3 +500,39 @@ def test_sql_concat_semantics_override():
         ctx2.sql("SELECT concat('a', NULL, 'b') AS c").collect_column("c")[0].as_py()
     )
     assert spark_out is None
+
+
+@pytest.mark.parametrize(
+    ("alias_fn", "primary_fn", "args"),
+    [
+        (spark.getbit, spark.bit_get, lambda: (lit(5), lit(0))),
+        (spark.dateadd, spark.date_add, lambda: (_ts().cast(pa.date32()), 3)),
+        (
+            spark.datediff,
+            spark.date_diff,
+            lambda: (_ts().cast(pa.date32()), lit("2020-01-01").cast(pa.date32())),
+        ),
+        (spark.datepart, spark.date_part, lambda: ("YEAR", _ts())),
+        (spark.sha, spark.sha1, lambda: (lit("abc"),)),
+        (spark.ceiling, spark.ceil, lambda: (lit(1.2),)),
+        (spark.printf, spark.format_string, lambda: ("%d-%s", lit(42), lit("hi"))),
+        (spark.char_length, spark.length, lambda: (lit("hello"),)),
+        (spark.character_length, spark.length, lambda: (lit("hello"),)),
+        (spark.power, spark.pow, lambda: (lit(2), lit(3))),
+        (spark.substr, spark.substring, lambda: (lit("hello"), 2, 3)),
+    ],
+)
+def test_aliases_match_primary(df, alias_fn, primary_fn, args):
+    assert _val(df, alias_fn(*args())) == _val(df, primary_fn(*args()))
+
+
+def test_substr_without_len(df):
+    assert _val(df, spark.substr(lit("hello"), 2)) == "ello"
+    assert _val(df, spark.substr(lit("hello"), -3)) == "llo"
+
+
+def test_last_day_date_keyword(df):
+    import datetime as dt
+
+    d = lit(pa.scalar(dt.date(2024, 2, 10), type=pa.date32()))
+    assert _val(df, spark.last_day(date=d)) == dt.date(2024, 2, 29)