Add percent rank and dense rank

Author: timsaucer

python/datafusion/dataframe.py

@@ -123,11 +123,10 @@ def select(self, *exprs: Expr | str) -> DataFrame:
             df = df.select("a", col("b"), col("a").alias("alternate_a"))
 
         """
-        exprs = [
-            arg.expr if isinstance(arg, Expr) else Expr.column(arg).expr
-            for arg in exprs
+        exprs_internal = [
+            Expr.column(arg).expr if isinstance(arg, str) else arg.expr for arg in exprs
         ]
-        return DataFrame(self.df.select(*exprs))
+        return DataFrame(self.df.select(*exprs_internal))
 
     def filter(self, *predicates: Expr) -> DataFrame:
         """Return a DataFrame for which ``predicate`` evaluates to ``True``.

python/datafusion/functions.py

@@ -250,9 +250,14 @@
     "var",
     "var_pop",
     "var_samp",
+    # Window Functions
     "window",
     "lead",
     "lag",
+    "row_number",
+    "rank",
+    "dense_rank",
+    "percent_rank",
 ]
 
 
@@ -1810,13 +1815,64 @@ def rank() -> Expr:
     is an example of a dataframe with a window ordered by descending ``points`` and the
     associated rank.
 
+    You should set ``order_by`` to produce meaningful results.
+
     ```
     +--------+------+
     | points | rank |
     +--------+------+
     | 100    | 1    |
     | 100    | 1    |
     | 50     | 3    |
+    | 25     | 4    |
+    +--------+------+
+    ```
+
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+    """
+    return Expr(f.rank())
+
+
+def dense_rank() -> Expr:
+    """Create a dense_rank window function.
+
+    This window function is similar to :py:func:`rank` except that the returned values
+    will be consecutive. Here is an example of a dataframe with a window ordered by
+    descending ``points`` and the associated dense rank.
+
+    ```
+    +--------+------------+
+    | points | dense_rank |
+    +--------+------------+
+    | 100    | 1          |
+    | 100    | 1          |
+    | 50     | 2          |
+    | 25     | 3          |
+    +--------+------------+
+    ```
+
+    To set window function parameters use the window builder approach described in the
+    ref:`_window_functions` online documentation.
+    """
+    return Expr(f.rank())
+
+
+def percent_rank() -> Expr:
+    """Create a percent_rank window function.
+
+    This window function is similar to :py:func:`rank` except that the returned values
+    will be consecutive. Here is an example of a dataframe with a window ordered by
+    descending ``points`` and the associated dense rank.
+
+    ```
+    +--------+------+
+    | points | rank |
+    +--------+------+
+    | 100    | 1    |
+    | 100    | 1    |
+    | 50     | 2    |
+    | 25     | 3    |
     +--------+------+
     ```
 

python/datafusion/tests/test_dataframe.py

@@ -283,7 +283,7 @@ def test_distinct():
     ("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"))]),
+        f.dense_rank().order_by((column("c").sort())).build(),
         [2, 1, 2],
     ),
     (
@@ -839,15 +839,3 @@ def test_write_compressed_parquet_missing_compression_level(df, tmp_path, compre
 
     with pytest.raises(ValueError):
         df.write_parquet(str(path), compression=compression)
-
-
-# ctx = SessionContext()
-
-# # create a RecordBatch and a new DataFrame from it
-# batch = pa.RecordBatch.from_arrays(
-#     [pa.array([1, 2, 3]), pa.array([4, 5, 6]), pa.array([8, 5, 8])],
-#     names=["a", "b", "c"],
-# )
-
-# df = ctx.create_dataframe([[batch]])
-# test_execute_stream(df)

src/functions.rs

@@ -879,6 +879,16 @@ pub fn rank() -> PyExpr {
     window_function::rank().into()
 }
 
+#[pyfunction]
+pub fn dense_rank() -> PyExpr {
+    window_function::dense_rank().into()
+}
+
+#[pyfunction]
+pub fn percent_rank() -> PyExpr {
+    window_function::percent_rank().into()
+}
+
 pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(abs))?;
     m.add_wrapped(wrap_pyfunction!(acos))?;
@@ -1067,6 +1077,8 @@ pub(crate) fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     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))?;
 
     Ok(())
 }