Add ntile window function

Author: timsaucer

python/datafusion/functions.py

@@ -1907,3 +1907,31 @@ def cume_dist() -> Expr:
     ref:`_window_functions` online documentation.
     """
     return Expr(f.cume_dist())
+
+
+def ntile(groups: int) -> Expr:
+    """Create a n-tile window function.
+
+    This window function orders the window frame into a give number of groups based on
+    the ordering criteria. It then returns which group the current row is assigned to.
+    Here is an example of a dataframe with a window ordered by descending ``points``
+    and the associated n-tile function.
+
+    ```
+    +--------+-------+
+    | points | ntile |
+    +--------+-------+
+    | 120    | 1     |
+    | 100    | 1     |
+    | 80     | 2     |
+    | 60     | 2     |
+    | 40     | 3     |
+    | 20     | 3     |
+    +--------+-------+
+    ```
+
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+    """
+    # Developer note: ntile only accepts literal values.
+    return Expr(f.ntile(Expr.literal(groups).expr))

python/datafusion/tests/test_dataframe.py

@@ -298,7 +298,7 @@ def test_distinct():
     ),
     (
         "ntile",
-        f.window("ntile", [literal(2)], order_by=[f.order_by(column("c"))]),
+        f.ntile(2).order_by(column("c").sort()).build(),
         [1, 1, 2],
     ),
     ("lead", f.lead(column("b")).order_by(column("b").sort()).build(), [5, 6, None]),

src/functions.rs

@@ -894,6 +894,11 @@ pub fn cume_dist() -> PyExpr {
     window_function::cume_dist().into()
 }
 
+#[pyfunction]
+pub fn ntile(arg: PyExpr) -> PyExpr {
+    window_function::ntile(arg.into()).into()
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;
     m.add_wrapped(wrap_pyfunction!(acos))?;
@@ -1078,13 +1083,15 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(array_slice))?;
     m.add_wrapped(wrap_pyfunction!(flatten))?;
 
+    // Window Functions
     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))?;
     m.add_wrapped(wrap_pyfunction!(dense_rank))?;
     m.add_wrapped(wrap_pyfunction!(percent_rank))?;
     m.add_wrapped(wrap_pyfunction!(cume_dist))?;
+    m.add_wrapped(wrap_pyfunction!(ntile))?;
 
     Ok(())
 }