fix: average ndcg over query groups

flaml/automl/ml.py

@@ -273,8 +273,11 @@ 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
@@ -287,8 +290,6 @@ def sklearn_metric_loss_score(
                 psum += c
             score /= len(counts)
             score += 1
-        else:
-            score = 1 - ndcg_score([y_true], [y_predict])
     return score
 
 

test/automl/test_ndcg_groups.py

@@ -0,0 +1,21 @@
+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)