Improve array function APIs: optional params, better naming, restore comment

- Make null_string optional in string_to_array/string_to_list
- Make step optional in gen_series/generate_series
- Rename second_array to element in array_contains/list_has/list_contains
- Restore # Window Functions section comment in __all__
- Add tests for optional parameter variants

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

Author: timsaucer

crates/core/src/functions.rs

@@ -109,6 +109,34 @@ fn arrays_zip(exprs: Vec<PyExpr>) -> PyExpr {
     datafusion::functions_nested::expr_fn::arrays_zip(exprs).into()
 }
 
+#[pyfunction]
+#[pyo3(signature = (string, delimiter, null_string=None))]
+fn string_to_array(string: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr>) -> PyExpr {
+    let mut args = vec![string.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::string_to_array_udf(),
+        args,
+    ))
+    .into()
+}
+
+#[pyfunction]
+#[pyo3(signature = (start, stop, step=None))]
+fn gen_series(start: PyExpr, stop: PyExpr, step: Option<PyExpr>) -> PyExpr {
+    let mut args = vec![start.into(), stop.into()];
+    if let Some(step) = step {
+        args.push(step.into());
+    }
+    Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
+        datafusion::functions_nested::range::gen_series_udf(),
+        args,
+    ))
+    .into()
+}
+
 #[pyfunction]
 #[pyo3(signature = (array, element, index=None))]
 fn array_position(array: PyExpr, element: PyExpr, index: Option<i64>) -> PyExpr {
@@ -687,8 +715,6 @@ array_fn!(array_any_value, array);
 array_fn!(array_max, array);
 array_fn!(array_min, array);
 array_fn!(array_reverse, array);
-array_fn!(string_to_array, string delimiter null_string);
-array_fn!(gen_series, start stop step);
 array_fn!(cardinality, array);
 array_fn!(flatten, array);
 array_fn!(range, start stop step);

python/datafusion/functions.py

@@ -317,6 +317,7 @@
     "var_samp",
     "var_sample",
     "when",
+    # Window Functions
     "window",
 ]
 
@@ -2848,22 +2849,22 @@ def array_has_any(first_array: Expr, second_array: Expr) -> Expr:
     return Expr(f.array_has_any(first_array.expr, second_array.expr))
 
 
-def array_contains(first_array: Expr, second_array: Expr) -> Expr:
-    """Returns true if the element appears in the first array, otherwise false.
+def array_contains(array: Expr, element: Expr) -> Expr:
+    """Returns true if the element appears in the array, otherwise false.
 
     See Also:
         This is an alias for :py:func:`array_has`.
     """
-    return array_has(first_array, second_array)
+    return array_has(array, element)
 
 
-def list_has(first_array: Expr, second_array: Expr) -> Expr:
-    """Returns true if the element appears in the first array, otherwise false.
+def list_has(array: Expr, element: Expr) -> Expr:
+    """Returns true if the element appears in the array, otherwise false.
 
     See Also:
         This is an alias for :py:func:`array_has`.
     """
-    return array_has(first_array, second_array)
+    return array_has(array, element)
 
 
 def list_has_all(first_array: Expr, second_array: Expr) -> Expr:
@@ -2884,13 +2885,13 @@ def list_has_any(first_array: Expr, second_array: Expr) -> Expr:
     return array_has_any(first_array, second_array)
 
 
-def list_contains(first_array: Expr, second_array: Expr) -> Expr:
-    """Returns true if the element appears in the first array, otherwise false.
+def list_contains(array: Expr, element: Expr) -> Expr:
+    """Returns true if the element appears in the array, otherwise false.
 
     See Also:
         This is an alias for :py:func:`array_has`.
     """
-    return array_has(first_array, second_array)
+    return array_has(array, element)
 
 
 def array_position(array: Expr, element: Expr, index: int | None = 1) -> Expr:
@@ -3590,25 +3591,30 @@ def list_zip(*arrays: Expr) -> Expr:
     return arrays_zip(*arrays)
 
 
-def string_to_array(string: Expr, delimiter: Expr, null_string: Expr) -> Expr:
+def string_to_array(
+    string: Expr, delimiter: Expr, null_string: Expr | None = None
+) -> Expr:
     """Splits a string based on a delimiter and returns an array of parts.
 
-    Any parts matching the ``null_string`` will be replaced with ``NULL``.
+    Any parts matching the optional ``null_string`` will be replaced with ``NULL``.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": ["hello,world"]})
         >>> result = df.select(
         ...     dfn.functions.string_to_array(
-        ...         dfn.col("a"), dfn.lit(","), dfn.lit(""),
+        ...         dfn.col("a"), dfn.lit(","),
         ...     ).alias("result"))
         >>> result.collect_column("result")[0].as_py()
         ['hello', 'world']
     """
-    return Expr(f.string_to_array(string.expr, delimiter.expr, null_string.expr))
+    null_expr = null_string.expr if null_string is not None else None
+    return Expr(f.string_to_array(string.expr, delimiter.expr, null_expr))
 
 
-def string_to_list(string: Expr, delimiter: Expr, null_string: Expr) -> Expr:
+def string_to_list(
+    string: Expr, delimiter: Expr, null_string: Expr | None = None
+) -> Expr:
     """Splits a string based on a delimiter and returns an array of parts.
 
     See Also:
@@ -3617,7 +3623,7 @@ def string_to_list(string: Expr, delimiter: Expr, null_string: Expr) -> Expr:
     return string_to_array(string, delimiter, null_string)
 
 
-def gen_series(start: Expr, stop: Expr, step: Expr) -> Expr:
+def gen_series(start: Expr, stop: Expr, step: Expr | None = None) -> Expr:
     """Creates a list of values in the range between start and stop.
 
     Unlike :py:func:`range`, this includes the upper bound.
@@ -3627,15 +3633,16 @@ def gen_series(start: Expr, stop: Expr, step: Expr) -> Expr:
         >>> df = ctx.from_pydict({"a": [0]})
         >>> result = df.select(
         ...     dfn.functions.gen_series(
-        ...         dfn.lit(1), dfn.lit(5), dfn.lit(1),
+        ...         dfn.lit(1), dfn.lit(5),
         ...     ).alias("result"))
         >>> result.collect_column("result")[0].as_py()
         [1, 2, 3, 4, 5]
     """
-    return Expr(f.gen_series(start.expr, stop.expr, step.expr))
+    step_expr = step.expr if step is not None else None
+    return Expr(f.gen_series(start.expr, stop.expr, step_expr))
 
 
-def generate_series(start: Expr, stop: Expr, step: Expr) -> Expr:
+def generate_series(start: Expr, stop: Expr, step: Expr | None = None) -> Expr:
     """Creates a list of values in the range between start and stop.
 
     Unlike :py:func:`range`, this includes the upper bound.

python/tests/test_functions.py

@@ -1567,35 +1567,48 @@ def test_string_to_array():
     ctx = SessionContext()
     df = ctx.from_pydict({"a": ["hello,world,foo"]})
     result = df.select(
-        f.string_to_array(column("a"), literal(","), literal("")).alias("v")
+        f.string_to_array(column("a"), literal(",")).alias("v")
     ).collect()
     assert result[0].column(0)[0].as_py() == ["hello", "world", "foo"]
 
 
-def test_string_to_list():
+def test_string_to_array_with_null_string():
     ctx = SessionContext()
-    df = ctx.from_pydict({"a": ["a-b-c"]})
+    df = ctx.from_pydict({"a": ["hello,NA,world"]})
     result = df.select(
-        f.string_to_list(column("a"), literal("-"), literal("")).alias("v")
+        f.string_to_array(column("a"), literal(","), literal("NA")).alias("v")
     ).collect()
+    values = result[0].column(0)[0].as_py()
+    assert values == ["hello", None, "world"]
+
+
+def test_string_to_list():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": ["a-b-c"]})
+    result = df.select(f.string_to_list(column("a"), literal("-")).alias("v")).collect()
     assert result[0].column(0)[0].as_py() == ["a", "b", "c"]
 
 
 def test_gen_series():
     ctx = SessionContext()
     df = ctx.from_pydict({"a": [0]})
-    result = df.select(
-        f.gen_series(literal(1), literal(5), literal(1)).alias("v")
-    ).collect()
+    result = df.select(f.gen_series(literal(1), literal(5)).alias("v")).collect()
     assert result[0].column(0)[0].as_py() == [1, 2, 3, 4, 5]
 
 
-def test_generate_series():
+def test_gen_series_with_step():
     ctx = SessionContext()
     df = ctx.from_pydict({"a": [0]})
     result = df.select(
-        f.generate_series(literal(1), literal(3), literal(1)).alias("v")
+        f.gen_series(literal(1), literal(10), literal(3)).alias("v")
     ).collect()
+    assert result[0].column(0)[0].as_py() == [1, 4, 7, 10]
+
+
+def test_generate_series():
+    ctx = SessionContext()
+    df = ctx.from_pydict({"a": [0]})
+    result = df.select(f.generate_series(literal(1), literal(3)).alias("v")).collect()
     assert result[0].column(0)[0].as_py() == [1, 2, 3]