fix: average ndcg over query groups

flaml/automl/ml.py

@@ -273,22 +273,28 @@ def sklearn_metric_loss_score(
     elif "ap" == metric_name:
         score = 1 - average_precision_score(y_true, y_predict, sample_weight=sample_weight)
     elif "ndcg" in metric_name:
-        if "@" in metric_name:
-            k = int(metric_name.split("@", 1)[-1])
+        k = int(metric_name.split("@", 1)[-1]) if "@" in metric_name else None
+        if groups is None:
+            score = 1 - ndcg_score([y_true], [y_predict], k=k)
+        else:
+            # average the NDCG of each query group
             counts = group_counts(groups)
             score = 0
             psum = 0
             for c in counts:
-                score -= ndcg_score(
-                    np.asarray([y_true[psum : psum + c]]),
-                    np.asarray([y_predict[psum : psum + c]]),
-                    k=k,
-                )
+                if c == 1:
+                    # a query with one document is always ranked perfectly, and
+                    # ndcg_score raises for it
+                    score -= 1
+                else:
+                    score -= ndcg_score(
+                        np.asarray([y_true[psum : psum + c]]),
+                        np.asarray([y_predict[psum : psum + c]]),
+                        k=k,
+                    )
                 psum += c
             score /= len(counts)
             score += 1
-        else:
-            score = 1 - ndcg_score([y_true], [y_predict])
     return score
 
 

flaml/automl/spark/metrics.py

@@ -189,40 +189,36 @@ def spark_metric_loss_score(
     elif "ndcg" in metric_name:
         # TODO: check if spark.ml ranker has the same format with
         # synapseML ranker, may need to adjust the format of df
-        if "@" in metric_name:
-            k = int(metric_name.split("@", 1)[-1])
-            if groups is None:
+        k = int(metric_name.split("@", 1)[-1]) if "@" in metric_name else None
+        if groups is None:
+            evaluator = RankingEvaluator(
+                metricName="ndcgAtK",
+                labelCol=label_col,
+                predictionCol=prediction_col,
+                **({} if k is None else {"k": k}),
+            )
+            df = _process_df(df, label_col, prediction_col)
+            score = 1 - evaluator.evaluate(df)
+        else:
+            # average the NDCG of each query group, as in sklearn_metric_loss_score
+            counts = ps_group_counts(groups)
+            score = 0
+            psum = 0
+            for c in counts:
+                y_true_ = y_true[psum : psum + c]
+                y_predict_ = y_predict[psum : psum + c]
+                df = y_true_.to_frame().join(y_predict_).to_spark()
+                df = _process_df(df, label_col, prediction_col)
                 evaluator = RankingEvaluator(
                     metricName="ndcgAtK",
                     labelCol=label_col,
                     predictionCol=prediction_col,
-                    k=k,
+                    k=c if k is None else k,
                 )
-                df = _process_df(df, label_col, prediction_col)
-                score = 1 - evaluator.evaluate(df)
-            else:
-                counts = ps_group_counts(groups)
-                score = 0
-                psum = 0
-                for c in counts:
-                    y_true_ = y_true[psum : psum + c]
-                    y_predict_ = y_predict[psum : psum + c]
-                    df = y_true_.to_frame().join(y_predict_).to_spark()
-                    df = _process_df(df, label_col, prediction_col)
-                    evaluator = RankingEvaluator(
-                        metricName="ndcgAtK",
-                        labelCol=label_col,
-                        predictionCol=prediction_col,
-                        k=k,
-                    )
-                    score -= evaluator.evaluate(df)
-                    psum += c
-                score /= len(counts)
-                score += 1
-        else:
-            evaluator = RankingEvaluator(metricName="ndcgAtK", labelCol=label_col, predictionCol=prediction_col)
-            df = _process_df(df, label_col, prediction_col)
-            score = 1 - evaluator.evaluate(df)
+                score -= evaluator.evaluate(df)
+                psum += c
+            score /= len(counts)
+            score += 1
         return score
     else:
         raise ValueError(f"Unknown metric name: {metric_name} for spark models.")

test/automl/test_ndcg_groups.py

@@ -0,0 +1,30 @@
+import numpy as np
+import pytest
+from sklearn.metrics import ndcg_score
+
+from flaml.automl.ml import sklearn_metric_loss_score
+
+# Two queries, each ranked perfectly, but the scores of the second query are all lower.
+Y_TRUE = np.array([3, 2, 1, 0, 3, 2, 1, 0])
+Y_PRED = np.array([10, 9, 8, 7, 4, 3, 2, 1])
+GROUPS = np.array([0, 0, 0, 0, 1, 1, 1, 1])
+
+
+@pytest.mark.parametrize("metric", ["ndcg", "ndcg@3", "ndcg@8"])
+def test_ndcg_averages_over_query_groups(metric):
+    assert sklearn_metric_loss_score(metric, Y_PRED, Y_TRUE, groups=GROUPS) == pytest.approx(0.0)
+
+
+@pytest.mark.parametrize("metric, k", [("ndcg", None), ("ndcg@3", 3)])
+def test_ndcg_without_groups_uses_one_query(metric, k):
+    expected = 1 - ndcg_score([Y_TRUE], [Y_PRED], k=k)
+    assert sklearn_metric_loss_score(metric, Y_PRED, Y_TRUE) == pytest.approx(expected)
+
+
+@pytest.mark.parametrize("metric", ["ndcg", "ndcg@3"])
+def test_ndcg_query_with_one_document(metric):
+    # a query with a single document is ranked perfectly and must not raise
+    y_true = np.append(Y_TRUE, 2)
+    y_pred = np.append(Y_PRED, 5)
+    groups = np.append(GROUPS, 2)
+    assert sklearn_metric_loss_score(metric, y_pred, y_true, groups=groups) == pytest.approx(0.0)

test/spark/test_utils.py

@@ -273,6 +273,22 @@ def test_iloc_pandas_on_spark():
     assert iloc_pandas_on_spark(psds, [0, 3]).tolist() == [1, 3]
 
 
+def test_spark_ndcg_query_groups():
+    spark = SparkSession.builder.getOrCreate()
+    # every row of the first query has prediction == label, none of the second one
+    dataset = spark.createDataFrame(
+        [(1.0, 1.0), (2.0, 2.0), (3.0, 3.0), (1.0, 2.0), (2.0, 1.0)], ["prediction", "label"]
+    )
+    dataset = to_pandas_on_spark(dataset)
+    groups = pd.Series([0, 0, 0, 1, 1])
+    # without groups, the rows are scored together; with groups, each query
+    # counts the same, for both ndcg and ndcg@k
+    assert spark_metric_loss_score("ndcg", dataset["prediction"], dataset["label"]) == pytest.approx(0.4)
+    for metric in ["ndcg", "ndcg@3"]:
+        loss = spark_metric_loss_score(metric, dataset["prediction"], dataset["label"], groups=groups)
+        assert loss == pytest.approx(0.5)
+
+
 def test_spark_metric_loss_score():
     spark = SparkSession.builder.getOrCreate()
     scoreAndLabels = map(