feat: add is_pareto (#17)

Author: CangyuanLi

python/rapidstats/_polars/_numeric.py

@@ -169,3 +169,57 @@ def sum_horizontal(
         raise ValueError(
             f"Invalid `null_strategy` {null_method}, must be one of `kleene`, `ignore`, or `propagate`"
         )
+
+
+def is_pareto(*exprs: IntoExpr | Iterable[IntoExpr]) -> pl.Expr:
+    """Identifies whether each point lies on the Pareto frontier. A point is considered
+    Pareto-optimal (non-dominated) if there is no other point that is at least as large
+    in all dimensions and strictly larger in at least one dimension. All dimensions are
+    assumed to be maximized.
+
+    !!! warning
+
+        Currently, only 2 dimensions are supported.
+
+    Returns
+    -------
+    pl.Expr
+        A boolean expression indicating whether each row is pareto. Rows where there are
+        any nulls or NaNs are null.
+
+    Examples
+    --------
+    ``` py
+    import polars as pl
+    import rapidstats.polars as prs
+
+    df = pl.DataFrame({"x": [5, 1, 3, 2], "y": [1, 5, 3, 2]})
+    df.select(prs.is_pareto("x", "y"))
+    ```
+    ``` title="output"
+    shape: (4, 1)
+    ┌───────┐
+    │ x     │
+    │ ---   │
+    │ bool  │
+    ╞═══════╡
+    │ true  │
+    │ true  │
+    │ true  │
+    │ false │
+    └───────┘
+    ```
+
+    Added in version 0.4.0
+    ----------------------
+    """
+    parsed_exprs = _parse_into_list_of_exprs(*exprs)
+
+    if len(parsed_exprs) != 2:
+        raise NotImplementedError("Only 2 dimensions are currently supported")
+
+    return pl.plugins.register_plugin_function(
+        plugin_path=_PLUGIN_PATH,
+        function_name="pl_pareto_2d",
+        args=[e.cast(pl.Float64) for e in parsed_exprs],
+    )

python/rapidstats/polars.py

@@ -1,4 +1,4 @@
 # ruff: noqa: F401
 
 from ._polars._format import format
-from ._polars._numeric import auc, is_close, sum_horizontal
+from ._polars._numeric import auc, is_close, is_pareto, sum_horizontal

src/general.rs

@@ -1,5 +1,15 @@
+use std::f64;
+
 use polars::prelude::*;
-use pyo3_polars::derive::polars_expr;
+use pyo3_polars::{
+    derive::polars_expr,
+    export::{
+        polars_arrow::{
+            array::BooleanArray, bitmap::MutableBitmap, datatypes::ArrowDataType::Boolean,
+        },
+        polars_core::utils::Container,
+    },
+};
 
 pub fn trapezoidal_auc(x: &[f64], y: &[f64]) -> f64 {
     x.windows(2)
@@ -47,3 +57,53 @@ fn pl_auc(inputs: &[Series]) -> PolarsResult<Series> {
 
     Ok(Series::from_vec("auc".into(), vec![res]))
 }
+
+#[polars_expr(output_type=Boolean)]
+fn pl_pareto_2d(inputs: &[Series]) -> PolarsResult<Series> {
+    let x = &inputs[0];
+    let y = &inputs[1];
+
+    let df = df!("x" => x, "y" => y)?
+        .with_row_index("index".into(), Some(0))?
+        .sort(
+            ["x", "y"],
+            SortMultipleOptions::default().with_order_descending_multi([true, true]),
+        )?;
+
+    let index = df["index"].u32()?;
+    let x_sorted = df["x"].f64()?;
+    let y_sorted = df["y"].f64()?;
+
+    let mut res: Vec<bool> = vec![false; x.len()];
+    let mut validity = MutableBitmap::with_capacity(x.len());
+    validity.extend_constant(x.len(), true);
+
+    let mut best_y = -f64::INFINITY;
+    for ((i, x), y) in index.into_no_null_iter().zip(x_sorted).zip(y_sorted) {
+        let i_u = i as usize;
+
+        let (x, y) = match (x, y) {
+            (Some(x), Some(y)) => (x, y),
+            _ => {
+                validity.set(i_u, false);
+                continue;
+            }
+        };
+
+        if x.is_nan() || y.is_nan() {
+            validity.set(i_u, false);
+            continue;
+        }
+
+        if y > best_y {
+            res[i_u] = true;
+            best_y = y;
+        } else {
+            res[i_u] = false;
+        }
+    }
+
+    let arr = BooleanArray::new(Boolean, res.into(), validity.into());
+
+    Ok(BooleanChunked::with_chunk(PlSmallStr::EMPTY, arr).into_series())
+}

tests/test_polars.py

@@ -146,3 +146,17 @@ def test_sum_horizontal():
         ).to_series(),
         pl.Series("x", [3, None, 1]),
     )
+
+
+def test_is_pareto():
+    df = pl.DataFrame({"x": [5, 1, 3, 2], "y": [1, 5, 3, 2]}).with_columns(
+        prs.is_pareto("x", "y").alias("is_pareto")
+    )
+    assert df["is_pareto"].to_list() == [True, True, True, False]
+
+
+def test_is_pareto_with_nulls():
+    df = pl.DataFrame({"x": [5, 1, 3, 2], "y": [1, 5, None, 2]}).with_columns(
+        prs.is_pareto("x", "y").alias("is_pareto")
+    )
+    assert df["is_pareto"].to_list() == [True, True, None, True]