Add rank window function

Author: timsaucer

python/datafusion/functions.py

@@ -1753,6 +1753,9 @@ def lead(arg: Expr, shift_offset: int = 1, default_value: Optional[Any] = None)
     return the 3rd following value in column ``b``. At the end of the partition, where
     no futher values can be returned it will return the default value of 5.
 
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+
     Args:
         arg: Value to return
         shift_offset: Number of rows following the current row.
@@ -1772,6 +1775,9 @@ def lag(arg: Expr, shift_offset: int = 1, default_value: Optional[Any] = None) -
     will return the 3rd previous value in column ``b``. At the beginnig of the
     partition, where no values can be returned it will return the default value of 5.
 
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+
     Args:
         arg: Value to return
         shift_offset: Number of rows before the current row.
@@ -1786,8 +1792,35 @@ def lag(arg: Expr, shift_offset: int = 1, default_value: Optional[Any] = None) -
 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.
+    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())
+
+
+def rank() -> Expr:
+    """Create a rank window function.
+
+    Returns the rank based upon the window order. Consecutive equal values will receive
+    the same rank, but the next different value will not be consecutive but rather the
+    number of rows that preceed it plus one. This is similar to Olympic medals. If two
+    people tie for gold, the next place is bronze. There would be no silver medal. Here
+    is an example of a dataframe with a window ordered by descending ``points`` and the
+    associated rank.
+
+    ```
+    +--------+------+
+    | points | rank |
+    +--------+------+
+    | 100    | 1    |
+    | 100    | 1    |
+    | 50     | 3    |
+    +--------+------+
+    ```
+
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+    """
+    return Expr(f.rank())

python/datafusion/tests/test_dataframe.py

@@ -280,7 +280,7 @@ def test_distinct():
 
 data_test_window_functions = [
     ("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]),
+    ("rank", f.rank().order_by(column("c").sort()).build(), [2, 1, 2]),
     (
         "dense_rank",
         f.window("dense_rank", [], order_by=[f.order_by(column("c"))]),

src/functions.rs

@@ -874,6 +874,11 @@ pub fn row_number() -> PyExpr {
     window_function::row_number().into()
 }
 
+#[pyfunction]
+pub fn rank() -> PyExpr {
+    window_function::rank().into()
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;
     m.add_wrapped(wrap_pyfunction!(acos))?;
@@ -1061,6 +1066,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))?;
+    m.add_wrapped(wrap_pyfunction!(rank))?;
 
     Ok(())
 }