feat: support average='binary' in precision_score()

bigframes/ml/metrics/_metrics.py

@@ -15,9 +15,11 @@
 """Metrics functions for evaluating models. This module is styled after
 scikit-learn's metrics module: https://scikit-learn.org/stable/modules/metrics.html."""
 
+from __future__ import annotations
+
 import inspect
 import typing
-from typing import Tuple, Union
+from typing import Literal, overload, Tuple, Union
 
 import bigframes_vendored.constants as constants
 import bigframes_vendored.sklearn.metrics._classification as vendored_metrics_classification
@@ -259,31 +261,64 @@ def recall_score(
 recall_score.__doc__ = inspect.getdoc(vendored_metrics_classification.recall_score)
 
 
+@overload
 def precision_score(
-    y_true: Union[bpd.DataFrame, bpd.Series],
-    y_pred: Union[bpd.DataFrame, bpd.Series],
+    y_true: bpd.DataFrame | bpd.Series,
+    y_pred: bpd.DataFrame | bpd.Series,
     *,
-    average: typing.Optional[str] = "binary",
+    pos_label: int | float | bool | str = ...,
+    average: Literal["binary"] = ...,
+) -> float:
+    ...
+
+
+@overload
+def precision_score(
+    y_true: bpd.DataFrame | bpd.Series,
+    y_pred: bpd.DataFrame | bpd.Series,
+    *,
+    pos_label: int | float | bool | str = ...,
+    average: None = ...,
 ) -> pd.Series:
-    # TODO(ashleyxu): support more average type, default to "binary"
-    if average is not None:
-        raise NotImplementedError(
-            f"Only average=None is supported. {constants.FEEDBACK_LINK}"
-        )
+    ...
 
+
+def precision_score(
+    y_true: bpd.DataFrame | bpd.Series,
+    y_pred: bpd.DataFrame | bpd.Series,
+    *,
+    pos_label: int | float | bool | str = 1,
+    average: Literal["binary"] | None = "binary",
+) -> pd.Series | float:
     y_true_series, y_pred_series = utils.batch_convert_to_series(y_true, y_pred)
 
-    is_accurate = y_true_series == y_pred_series
+    if average is None:
+        return _precision_score_per_class(y_true_series, y_pred_series)
+
+    if average == "binary":
+        return _precision_score_binary_pos_only(y_true_series, y_pred_series, pos_label)
+
+    raise NotImplementedError(
+        f"Unsupported 'average' param value: {average}. {constants.FEEDBACK_LINK}"
+    )
+
+
+precision_score.__doc__ = inspect.getdoc(
+    vendored_metrics_classification.precision_score
+)
+
+
+def _precision_score_per_class(y_true: bpd.Series, y_pred: bpd.Series) -> pd.Series:
+    is_accurate = y_true == y_pred
     unique_labels = (
-        bpd.concat([y_true_series, y_pred_series], join="outer")
+        bpd.concat([y_true, y_pred], join="outer")
         .drop_duplicates()
         .sort_values(inplace=False)
     )
     index = unique_labels.to_list()
 
     precision = (
-        is_accurate.groupby(y_pred_series).sum()
-        / is_accurate.groupby(y_pred_series).count()
+        is_accurate.groupby(y_pred).sum() / is_accurate.groupby(y_pred).count()
     ).to_pandas()
 
     precision_score = pd.Series(0, index=index)
@@ -293,9 +328,32 @@ def precision_score(
     return precision_score
 
 
-precision_score.__doc__ = inspect.getdoc(
-    vendored_metrics_classification.precision_score
-)
+def _precision_score_binary_pos_only(
+    y_true: bpd.Series, y_pred: bpd.Series, pos_label: int | float | bool | str
+) -> float:
+    if y_true.drop_duplicates().count() != 2 or y_pred.drop_duplicates().count() != 2:
+        raise ValueError(
+            "Target is multiclass but average='binary'. Please choose another average setting."
+        )
+
+    total_labels = set(
+        y_true.drop_duplicates().to_list() + y_pred.drop_duplicates().to_list()
+    )
+
+    if len(total_labels) != 2:
+        raise ValueError(
+            "Target is multiclass but average='binary'. Please choose another average setting."
+        )
+
+    if pos_label not in total_labels:
+        raise ValueError(
+            f"pos_labe={pos_label} is not a valid label. It should be one of {list(total_labels)}"
+        )
+
+    target_elem_idx = y_pred == pos_label
+    is_accurate = y_pred[target_elem_idx] == y_true[target_elem_idx]
+
+    return is_accurate.sum() / is_accurate.count()
 
 
 def f1_score(

tests/system/small/ml/test_metrics.py

@@ -743,6 +743,57 @@ def test_precision_score_series(session):
     )
 
 
+@pytest.mark.parametrize(
+    ("pos_label", "expected_score"),
+    [
+        ("a", 1 / 3),
+        ("b", 0),
+    ],
+)
+def test_precision_score_binary(session, pos_label, expected_score):
+    pd_df = pd.DataFrame(
+        {
+            "y_true": ["a", "a", "a", "b", "b"],
+            "y_pred": ["b", "b", "a", "a", "a"],
+        }
+    )
+    df = session.read_pandas(pd_df)
+
+    precision_score = metrics.precision_score(
+        df["y_true"], df["y_pred"], average="binary", pos_label=pos_label
+    )
+
+    assert precision_score == pytest.approx(expected_score)
+
+
+@pytest.mark.parametrize(
+    ("y_true", "y_pred", "pos_label"),
+    [
+        pytest.param(
+            pd.Series([1, 2, 3]), pd.Series([1, 0]), 1, id="y_true-non-binary-label"
+        ),
+        pytest.param(
+            pd.Series([1, 0]), pd.Series([1, 2, 3]), 1, id="y_pred-non-binary-label"
+        ),
+        pytest.param(
+            pd.Series([1, 0]), pd.Series([1, 2]), 1, id="combined-non-binary-label"
+        ),
+        pytest.param(pd.Series([1, 0]), pd.Series([1, 0]), 2, id="invalid-pos_label"),
+    ],
+)
+def test_precision_score_binary_invalid_input_raise_error(
+    session, y_true, y_pred, pos_label
+):
+
+    bf_y_true = session.read_pandas(y_true)
+    bf_y_pred = session.read_pandas(y_pred)
+
+    with pytest.raises(ValueError):
+        metrics.precision_score(
+            bf_y_true, bf_y_pred, average="binary", pos_label=pos_label
+        )
+
+
 def test_f1_score(session):
     pd_df = pd.DataFrame(
         {

third_party/bigframes_vendored/sklearn/metrics/_classification.py

@@ -201,7 +201,7 @@ def precision_score(
                 default='binary'
             This parameter is required for multiclass/multilabel targets.
             Possible values are 'None', 'micro', 'macro', 'samples', 'weighted', 'binary'.
-            Only average=None is supported.
+            Only None and 'binary' is supported.
 
     Returns:
         precision: float (if average is not None) or Series of float of shape \