feat: accept native ints and a single argument in range and gen_series

range() required all three of start, stop, and step as Expr, although
upstream also accepts range(stop) and range(start, stop). Make stop and
step optional in the bindings for both range and gen_series, so a single
argument is the upper bound starting at 0, like Python's built-in range.

Accept plain ints for start, stop, and step, coerced to literals, so
callers no longer need lit(). Passing step without stop raises ValueError.

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

Author: timsaucer

crates/core/src/functions.rs

@@ -131,20 +131,46 @@ fn array_to_string(array: PyExpr, delimiter: PyExpr, null_string: Option<PyExpr>
     .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());
-    }
+/// Builds `range` or `gen_series` from its one, two, or three arguments.
+fn series_expr(
+    udf: std::sync::Arc<datafusion::logical_expr::ScalarUDF>,
+    start: PyExpr,
+    stop: Option<PyExpr>,
+    step: Option<PyExpr>,
+) -> PyExpr {
+    let args = std::iter::once(start)
+        .chain(stop)
+        .chain(step)
+        .map(Into::into)
+        .collect();
     Expr::ScalarFunction(datafusion::logical_expr::expr::ScalarFunction::new_udf(
-        datafusion::functions_nested::range::gen_series_udf(),
-        args,
+        udf, args,
     ))
     .into()
 }
 
+#[pyfunction]
+#[pyo3(signature = (start, stop=None, step=None))]
+fn range(start: PyExpr, stop: Option<PyExpr>, step: Option<PyExpr>) -> PyExpr {
+    series_expr(
+        datafusion::functions_nested::range::range_udf(),
+        start,
+        stop,
+        step,
+    )
+}
+
+#[pyfunction]
+#[pyo3(signature = (start, stop=None, step=None))]
+fn gen_series(start: PyExpr, stop: Option<PyExpr>, step: Option<PyExpr>) -> PyExpr {
+    series_expr(
+        datafusion::functions_nested::range::gen_series_udf(),
+        start,
+        stop,
+        step,
+    )
+}
+
 #[pyfunction]
 fn make_map(keys: Vec<PyExpr>, values: Vec<PyExpr>) -> PyExpr {
     let keys = keys.into_iter().map(|x| x.into()).collect();
@@ -701,7 +727,6 @@ array_fn!(array_min, array);
 array_fn!(array_reverse, array);
 array_fn!(cardinality, array);
 array_fn!(flatten, array);
-array_fn!(range, start stop step);
 
 // Map Functions
 array_fn!(map_keys, map);

python/datafusion/functions/__init__.py

@@ -3126,18 +3126,57 @@ def array(*args: Expr) -> Expr:
     return make_array(*args)
 
 
-def range(start: Expr, stop: Expr, step: Expr) -> Expr:
-    """Create a list of values in the range between start and stop.
+def _series(
+    fn: Callable[..., Any],
+    name: str,
+    start: Expr | int,
+    stop: Expr | int | None,
+    step: Expr | int | None,
+) -> Expr:
+    if stop is None and step is not None:
+        msg = f"{name}() requires stop when step is given"
+        raise ValueError(msg)
+    stop = coerce_to_expr_or_none(stop)
+    step = coerce_to_expr_or_none(step)
+    return Expr(
+        fn(
+            coerce_to_expr(start).expr,
+            stop.expr if stop is not None else None,
+            step.expr if step is not None else None,
+        )
+    )
+
+
+def range(
+    start: Expr | int,
+    stop: Expr | int | None = None,
+    step: Expr | int | None = None,
+) -> Expr:
+    """Create a list of values from ``start`` up to, but excluding, ``stop``.
+
+    With a single argument, it is the upper bound and the range starts at 0,
+    like Python's built-in :py:class:`range`.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [1]})
-        >>> result = df.select(
-        ...     dfn.functions.range(dfn.lit(0), dfn.lit(5), dfn.lit(2)).alias("r"))
+        >>> result = df.select(dfn.functions.range(5).alias("r"))
+        >>> result.collect_column("r")[0].as_py()
+        [0, 1, 2, 3, 4]
+
+        Specify a ``stop``:
+
+        >>> result = df.select(dfn.functions.range(1, stop=5).alias("r"))
+        >>> result.collect_column("r")[0].as_py()
+        [1, 2, 3, 4]
+
+        Specify a ``step``:
+
+        >>> result = df.select(dfn.functions.range(0, stop=5, step=2).alias("r"))
         >>> result.collect_column("r")[0].as_py()
         [0, 2, 4]
     """
-    return Expr(f.range(start.expr, stop.expr, step.expr))
+    return _series(f.range, "range", start, stop, step)
 
 
 def uuid() -> Expr:
@@ -5077,38 +5116,45 @@ def string_to_list(
     return string_to_array(string, delimiter, null_string)
 
 
-def gen_series(start: Expr, stop: Expr, step: Expr | None = None) -> Expr:
-    """Creates a list of values in the range between start and stop.
+def gen_series(
+    start: Expr | int,
+    stop: Expr | int | None = None,
+    step: Expr | int | None = None,
+) -> Expr:
+    """Creates a list of values from ``start`` up to and including ``stop``.
 
-    Unlike :py:func:`range`, this includes the upper bound.
+    Unlike :py:func:`range`, this includes the upper bound. With a single
+    argument, it is the upper bound and the series starts at 0.
 
     Examples:
         >>> ctx = dfn.SessionContext()
         >>> df = ctx.from_pydict({"a": [0]})
-        >>> result = df.select(
-        ...     dfn.functions.gen_series(
-        ...         dfn.lit(1), dfn.lit(5),
-        ...     ).alias("result"))
+        >>> result = df.select(dfn.functions.gen_series(3).alias("result"))
+        >>> result.collect_column("result")[0].as_py()
+        [0, 1, 2, 3]
+
+        Specify a ``stop``:
+
+        >>> result = df.select(dfn.functions.gen_series(1, stop=5).alias("result"))
         >>> result.collect_column("result")[0].as_py()
         [1, 2, 3, 4, 5]
 
-        Specify a custom ``step``:
+        Specify a ``step``:
 
         >>> result = df.select(
-        ...     dfn.functions.gen_series(
-        ...         dfn.lit(1), dfn.lit(10), step=dfn.lit(3),
-        ...     ).alias("result"))
+        ...     dfn.functions.gen_series(1, stop=10, step=3).alias("result"))
         >>> result.collect_column("result")[0].as_py()
         [1, 4, 7, 10]
     """
-    step_expr = step.expr if step is not None else None
-    return Expr(f.gen_series(start.expr, stop.expr, step_expr))
+    return _series(f.gen_series, "gen_series", start, stop, step)
 
 
-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.
+def generate_series(
+    start: Expr | int,
+    stop: Expr | int | None = None,
+    step: Expr | int | None = None,
+) -> Expr:
+    """Creates a list of values from ``start`` up to and including ``stop``.
 
     See Also:
         This is an alias for :py:func:`gen_series`.

python/tests/test_functions.py

@@ -2359,6 +2359,23 @@ def test_gen_series_with_step():
     assert result[0].column(0)[0].as_py() == [1, 4, 7, 10]
 
 
+@pytest.mark.parametrize(
+    ("func", "expected"),
+    [(f.range, [[0], [0, 1]]), (f.gen_series, [[0, 1], [0, 1, 2]])],
+)
+def test_series_single_arg_accepts_column(func, expected):
+    ctx = SessionContext()
+    df = ctx.from_pydict({"n": [1, 2]})
+    result = df.select(func(column("n")).alias("v"))
+    assert result.collect_column("v").to_pylist() == expected
+
+
+@pytest.mark.parametrize("func", [f.range, f.gen_series, f.generate_series])
+def test_series_step_requires_stop(func):
+    with pytest.raises(ValueError, match="requires stop"):
+        func(0, step=2)
+
+
 class TestPythonicNativeTypes:
     """Tests for accepting native Python types instead of requiring lit()."""