feat: accept optional arguments upstream supports on trim, array_to_string, substr

- btrim, ltrim, rtrim, and trim take an optional characters argument
  naming the set to strip.
- array_to_string and its aliases take an optional null_string that is
  written in place of NULL elements.
- substr takes an optional length.

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

Author: timsaucer

crates/core/src/functions.rs

@@ -117,6 +117,20 @@ fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr
     .into()
 }
 
+#[pyfunction]
+#[pyo3(signature = (array, delimiter, null_string=None))]
+fn array_to_string(array: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr>) -> PyExpr {
+    let mut args = vec![array.into(), delimiter.into()];
+    if let Some(null_string) = null_string {
+        args.push(null_string.into());
+    }
+    Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
+        datafusion::functions_nested::string::array_to_string_udf(),
+        args,
+    ))
+    .into()
+}
+
 #[pyfunction]
 #[pyo3(signature = (start, stop, step=None))]
 fn gen_series(start: PyExpr, stop: PyExpr, step: Option<PyExpr>) -> PyExpr {
@@ -646,7 +660,6 @@ fn version() -> PyExpr {
 
 // Array Functions
 array_fn!(array_append, array element);
-array_fn!(array_to_string, array delimiter);
 array_fn!(array_dims, array);
 array_fn!(array_distinct, array);
 array_fn!(array_element, array element);

python/datafusion/functions/__init__.py

@@ -486,46 +486,71 @@ def decode(expr: Expr, encoding: Expr | str) -> Expr:
     return Expr(f.decode(expr.expr, encoding.expr))
 
 
-def array_to_string(expr: Expr, delimiter: Expr | str) -> Expr:
+def array_to_string(
+    expr: Expr, delimiter: Expr | str, null_string: Expr | str | None = None
+) -> Expr:
     """Converts each element to its text representation.
 
+    NULL elements are omitted unless ``null_string`` is given, in which case it
+    is written in their place.
+
     Examples:
         >>> ctx = dfn.SessionContext()
-        >>> df = ctx.from_pydict({"a": [[1, 2, 3]]})
+        >>> df = ctx.from_pydict({"a": [[1, None, 3]]})
         >>> result = df.select(
         ...     dfn.functions.array_to_string(dfn.col("a"), ",").alias("s"))
         >>> result.collect_column("s")[0].as_py()
-        '1,2,3'
+        '1,3'
+
+        >>> result = df.select(
+        ...     dfn.functions.array_to_string(
+        ...         dfn.col("a"), ",", null_string="*"
+        ...     ).alias("s"))
+        >>> result.collect_column("s")[0].as_py()
+        '1,*,3'
     """
     delimiter = coerce_to_expr(delimiter)
-    return Expr(f.array_to_string(expr.expr, delimiter.expr.cast(pa.string())))
+    null_string = coerce_to_expr_or_none(null_string)
+    return Expr(
+        f.array_to_string(
+            expr.expr,
+            delimiter.expr.cast(pa.string()),
+            null_string.expr if null_string is not None else None,
+        )
+    )
 
 
-def array_join(expr: Expr, delimiter: Expr | str) -> Expr:
+def array_join(
+    expr: Expr, delimiter: Expr | str, null_string: Expr | str | None = None
+) -> Expr:
     """Converts each element to its text representation.
 
     See Also:
         This is an alias for :py:func:`array_to_string`.
     """
-    return array_to_string(expr, delimiter)
+    return array_to_string(expr, delimiter, null_string=null_string)
 
 
-def list_to_string(expr: Expr, delimiter: Expr | str) -> Expr:
+def list_to_string(
+    expr: Expr, delimiter: Expr | str, null_string: Expr | str | None = None
+) -> Expr:
     """Converts each element to its text representation.
 
     See Also:
         This is an alias for :py:func:`array_to_string`.
     """
-    return array_to_string(expr, delimiter)
+    return array_to_string(expr, delimiter, null_string=null_string)
 
 
-def list_join(expr: Expr, delimiter: Expr | str) -> Expr:
+def list_join(
+    expr: Expr, delimiter: Expr | str, null_string: Expr | str | None = None
+) -> Expr:
     """Converts each element to its text representation.
 
     See Also:
         This is an alias for :py:func:`array_to_string`.
     """
-    return array_to_string(expr, delimiter)
+    return array_to_string(expr, delimiter, null_string=null_string)
 
 
 def lambda_var(name: str) -> Expr:
@@ -1128,17 +1153,29 @@ def bit_length(arg: Expr) -> Expr:
     return Expr(f.bit_length(arg.expr))
 
 
-def btrim(arg: Expr) -> Expr:
-    """Removes all characters, spaces by default, from both sides of a string.
+def btrim(arg: Expr, characters: Expr | str | None = None) -> Expr:
+    """Removes ``characters``, spaces by default, from both sides of a string.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [" a  "]})
         >>> trim_df = df.select(dfn.functions.btrim(dfn.col("a")).alias("trimmed"))
         >>> trim_df.collect_column("trimmed")[0].as_py()
         'a'
+
+        Trim a different set of characters:
+
+        >>> df = ctx.from_pydict({"a": ["xxaxx"]})
+        >>> trim_df = df.select(
+        ...     dfn.functions.btrim(dfn.col("a"), characters="x").alias("trimmed")
+        ... )
+        >>> trim_df.collect_column("trimmed")[0].as_py()
+        'a'
     """
-    return Expr(f.btrim(arg.expr))
+    args = [arg.expr]
+    if characters is not None:
+        args.append(coerce_to_expr(characters).expr)
+    return Expr(f.btrim(*args))
 
 
 def cbrt(arg: Expr) -> Expr:
@@ -1616,17 +1653,29 @@ def lpad(string: Expr, count: Expr | int, characters: Expr | str | None = None)
     return Expr(f.lpad(string.expr, count.expr, characters.expr))
 
 
-def ltrim(arg: Expr) -> Expr:
-    """Removes all characters, spaces by default, from the beginning of a string.
+def ltrim(arg: Expr, characters: Expr | str | None = None) -> Expr:
+    """Removes ``characters``, spaces by default, from the beginning of a string.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [" a  "]})
         >>> trim_df = df.select(dfn.functions.ltrim(dfn.col("a")).alias("trimmed"))
         >>> trim_df.collect_column("trimmed")[0].as_py()
         'a  '
+
+        Trim a different set of characters:
+
+        >>> df = ctx.from_pydict({"a": ["xxaxx"]})
+        >>> trim_df = df.select(
+        ...     dfn.functions.ltrim(dfn.col("a"), characters="x").alias("trimmed")
+        ... )
+        >>> trim_df.collect_column("trimmed")[0].as_py()
+        'axx'
     """
-    return Expr(f.ltrim(arg.expr))
+    args = [arg.expr]
+    if characters is not None:
+        args.append(coerce_to_expr(characters).expr)
+    return Expr(f.ltrim(*args))
 
 
 def md5(arg: Expr) -> Expr:
@@ -2138,17 +2187,29 @@ def rpad(string: Expr, count: Expr | int, characters: Expr | str | None = None)
     return Expr(f.rpad(string.expr, count.expr, characters.expr))
 
 
-def rtrim(arg: Expr) -> Expr:
-    """Removes all characters, spaces by default, from the end of a string.
+def rtrim(arg: Expr, characters: Expr | str | None = None) -> Expr:
+    """Removes ``characters``, spaces by default, from the end of a string.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [" a  "]})
         >>> trim_df = df.select(dfn.functions.rtrim(dfn.col("a")).alias("trimmed"))
         >>> trim_df.collect_column("trimmed")[0].as_py()
         ' a'
+
+        Trim a different set of characters:
+
+        >>> df = ctx.from_pydict({"a": ["xxaxx"]})
+        >>> trim_df = df.select(
+        ...     dfn.functions.rtrim(dfn.col("a"), characters="x").alias("trimmed")
+        ... )
+        >>> trim_df.collect_column("trimmed")[0].as_py()
+        'xxa'
     """
-    return Expr(f.rtrim(arg.expr))
+    args = [arg.expr]
+    if characters is not None:
+        args.append(coerce_to_expr(characters).expr)
+    return Expr(f.rtrim(*args))
 
 
 def sha224(arg: Expr) -> Expr:
@@ -2312,8 +2373,10 @@ def strpos(string: Expr, substring: Expr | str) -> Expr:
     return Expr(f.strpos(string.expr, substring.expr))
 
 
-def substr(string: Expr, position: Expr | int) -> Expr:
-    """Substring from the ``position`` to the end.
+def substr(
+    string: Expr, position: Expr | int, length: Expr | int | None = None
+) -> Expr:
+    """Substring from the ``position``, to the end or for ``length`` characters.
 
     Examples:
         >>> ctx = dfn.SessionContext()
@@ -2322,7 +2385,17 @@ def substr(string: Expr, position: Expr | int) -> Expr:
         ...     dfn.functions.substr(dfn.col("a"), 3).alias("s"))
         >>> result.collect_column("s")[0].as_py()
         'llo'
+
+        >>> result = df.select(
+        ...     dfn.functions.substr(dfn.col("a"), 2, length=3).alias("s"))
+        >>> result.collect_column("s")[0].as_py()
+        'ell'
+
+    See Also:
+        :py:func:`substring`.
     """
+    if length is not None:
+        return substring(string, position, length)
     position = coerce_to_expr(position)
     return Expr(f.substr(string.expr, position.expr))
 
@@ -2958,17 +3031,29 @@ def translate(string: Expr, from_val: Expr | str, to_val: Expr | str) -> Expr:
     return Expr(f.translate(string.expr, from_val.expr, to_val.expr))
 
 
-def trim(arg: Expr) -> Expr:
-    """Removes all characters, spaces by default, from both sides of a string.
+def trim(arg: Expr, characters: Expr | str | None = None) -> Expr:
+    """Removes ``characters``, spaces by default, from both sides of a string.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": ["  hello  "]})
         >>> result = df.select(dfn.functions.trim(dfn.col("a")).alias("t"))
         >>> result.collect_column("t")[0].as_py()
         'hello'
+
+        Trim a different set of characters:
+
+        >>> df = ctx.from_pydict({"a": ["xxhelloxx"]})
+        >>> result = df.select(
+        ...     dfn.functions.trim(dfn.col("a"), characters="x").alias("t")
+        ... )
+        >>> result.collect_column("t")[0].as_py()
+        'hello'
     """
-    return Expr(f.trim(arg.expr))
+    args = [arg.expr]
+    if characters is not None:
+        args.append(coerce_to_expr(characters).expr)
+    return Expr(f.trim(*args))
 
 
 def trunc(num: Expr, precision: Expr | int | None = None) -> Expr:

python/tests/test_functions.py

@@ -2547,3 +2547,46 @@ def test_backward_compat_with_lit(self):
             f.split_part(column("a"), literal(","), literal(2)).alias("s")
         ).collect()
         assert result[0].column(0)[0].as_py() == "b"
+
+
+@pytest.mark.parametrize(
+    ("fn", "expected"),
+    [
+        (f.btrim, "hi"),
+        (f.trim, "hi"),
+        (f.ltrim, "hixyx"),
+        (f.rtrim, "xyxhi"),
+    ],
+)
+def test_trim_characters(fn, expected):
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": ["xyxhixyx"]})
+    assert df.select(fn(column("a"), characters="xy").alias("r")).collect_column(
+        "r"
+    ).to_pylist() == [expected]
+    assert df.select(
+        fn(column("a"), characters=literal("xy")).alias("r")
+    ).collect_column("r").to_pylist() == [expected]
+
+
+@pytest.mark.parametrize(
+    "fn", [f.array_to_string, f.array_join, f.list_to_string, f.list_join]
+)
+def test_array_to_string_null_string(fn):
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": [[1, None, 3]]})
+    without = df.select(fn(column("a"), "-").alias("r")).collect_column("r")
+    with_null = df.select(fn(column("a"), "-", null_string="NA").alias("r"))
+    assert without.to_pylist() == ["1-3"]
+    assert with_null.collect_column("r").to_pylist() == ["1-NA-3"]
+
+
+def test_substr_length():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": ["hello"]})
+    r = df.select(
+        f.substr(column("a"), 2).alias("tail"),
+        f.substr(column("a"), 2, length=3).alias("mid"),
+        f.substr(column("a"), 2, length=literal(3)).alias("mid_expr"),
+    ).to_pydict()
+    assert r == {"tail": ["ello"], "mid": ["ell"], "mid_expr": ["ell"]}