Expose row_number window function

Author: timsaucer

python/datafusion/functions.py

@@ -1781,3 +1781,13 @@ def lag(arg: Expr, shift_offset: int = 1, default_value: Optional[Any] = None) -
         default_value = pa.scalar(default_value)
 
     return Expr(f.lag(arg.expr, shift_offset, default_value))
+
+
+def row_number() -> Expr:
+    """Create a row number window function.
+
+    Returns the row number of the window function. To set window function parameters
+    use the window builder approach described in the ref:`_window_functions` online
+    documentation.
+    """
+    return Expr(f.row_number())

python/datafusion/tests/test_dataframe.py

@@ -279,7 +279,7 @@ def test_distinct():
 
 
 data_test_window_functions = [
-    ("row", f.window("row_number", [], order_by=[f.order_by(column("c"))]), [2, 1, 3]),
+    ("row", f.row_number().order_by(column("c").sort()).build(), [2, 1, 3]),
     ("rank", f.window("rank", [], order_by=[f.order_by(column("c"))]), [2, 1, 2]),
     (
         "dense_rank",

src/functions.rs

@@ -869,6 +869,11 @@ pub fn lag(arg: PyExpr, shift_offset: i64, default_value: Option<ScalarValue>) -
     window_function::lag(arg.expr, Some(shift_offset), default_value).into()
 }
 
+#[pyfunction]
+pub fn row_number() -> PyExpr {
+    window_function::row_number().into()
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;
     m.add_wrapped(wrap_pyfunction!(acos))?;
@@ -1055,6 +1060,7 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
 
     m.add_wrapped(wrap_pyfunction!(lead))?;
     m.add_wrapped(wrap_pyfunction!(lag))?;
+    m.add_wrapped(wrap_pyfunction!(row_number))?;
 
     Ok(())
 }