From a08cf2cee95c07a73c39cd984ba36395485a8726 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Wed, 7 Oct 2026 17:32:26 +0530
Subject: [PATCH 1/7] fix: average ndcg over query groups

With groups, ndcg@k averaged the NDCG of each query, but plain ndcg (the default metric for the rank task) scored all rows as one ranked list, so a model that ranks every query perfectly could still get a loss above 0. ndcg@k without groups also failed because group_counts(None) produced a single one-row group. Use the same per-group average for ndcg and ndcg@k, and a single list for both when there are no groups.
---
 flaml/automl/ml.py              |  9 +++++----
 test/automl/test_ndcg_groups.py | 21 +++++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)
 create mode 100644 test/automl/test_ndcg_groups.py

diff --git a/flaml/automl/ml.py b/flaml/automl/ml.py
index 3b09a4c9eb..57dce7f668 100644
--- a/flaml/automl/ml.py
+++ b/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
 
 
diff --git a/test/automl/test_ndcg_groups.py b/test/automl/test_ndcg_groups.py
new file mode 100644
index 0000000000..6ea26787f4
--- /dev/null
+++ b/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)

From c1aae975f7f00c7cf6f443d8e9d05c8efa676f6a Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Thu, 8 Oct 2026 01:16:42 +0530
Subject: [PATCH 2/7] fix: handle single-document queries and group plain ndcg
 on Spark

---
 flaml/automl/ml.py              | 15 ++++++---
 flaml/automl/spark/metrics.py   | 54 +++++++++++++++------------------
 test/automl/test_ndcg_groups.py |  9 ++++++
 test/spark/test_utils.py        | 16 ++++++++++
 4 files changed, 60 insertions(+), 34 deletions(-)

diff --git a/flaml/automl/ml.py b/flaml/automl/ml.py
index 57dce7f668..1a7f5d7cd8 100644
--- a/flaml/automl/ml.py
+++ b/flaml/automl/ml.py
@@ -282,11 +282,16 @@ def sklearn_metric_loss_score(
             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
diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index 7cb85c09d5..a56afd115b 100644
--- a/flaml/automl/spark/metrics.py
+++ b/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.")
diff --git a/test/automl/test_ndcg_groups.py b/test/automl/test_ndcg_groups.py
index 6ea26787f4..d157a12541 100644
--- a/test/automl/test_ndcg_groups.py
+++ b/test/automl/test_ndcg_groups.py
@@ -19,3 +19,12 @@ def test_ndcg_averages_over_query_groups(metric):
 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)
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index 8e787e9e83..05d2a42a36 100644
--- a/test/spark/test_utils.py
+++ b/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(

From bdcf7d8622273722f287c6c26ddabd3a2a8c0073 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Thu, 8 Oct 2026 05:16:00 +0530
Subject: [PATCH 3/7] fix: rank documents by score within each query for Spark
 ndcg

---
 flaml/automl/spark/metrics.py | 56 ++++++++++++-----------------------
 test/spark/test_utils.py      | 42 ++++++++++++++++++--------
 2 files changed, 49 insertions(+), 49 deletions(-)

diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index a56afd115b..e2e9fb77a8 100644
--- a/flaml/automl/spark/metrics.py
+++ b/flaml/automl/spark/metrics.py
@@ -22,12 +22,6 @@ def ps_group_counts(groups: Union[psSeries, np.ndarray]) -> np.ndarray:
     return c[np.argsort(i)].tolist()
 
 
-def _process_df(df, label_col, prediction_col):
-    df = df.withColumn(label_col, F.array([df[label_col]]))
-    df = df.withColumn(prediction_col, F.array([df[prediction_col]]))
-    return df
-
-
 def _compute_label_from_probability(df, probability_col, prediction_col):
     # array_max finds the maximum value in the 'probability' array
     # array_position finds the index of the maximum value in the 'probability' array
@@ -187,38 +181,26 @@ def spark_metric_loss_score(
             predictionCol=prediction_col,
         )
     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
+        # RankingEvaluator compares arrays of ranked and relevant item ids, not
+        # prediction scores and graded relevance labels, so rank the documents of
+        # each query by score here, the same way as sklearn_metric_loss_score
+        from sklearn.metrics import ndcg_score
+
         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=c if k is None else k,
-                )
-                score -= evaluator.evaluate(df)
-                psum += c
-            score /= len(counts)
-            score += 1
+        y_true = y_true.to_numpy()
+        y_predict = y_predict.to_numpy()
+        counts = [len(y_true)] if groups is None else ps_group_counts(groups)
+        score = 0
+        psum = 0
+        for c in counts:
+            if c == 1:
+                # a query with one document is always ranked perfectly
+                score -= 1
+            else:
+                score -= ndcg_score([y_true[psum : psum + c]], [y_predict[psum : psum + c]], k=k)
+            psum += c
+        score /= len(counts)
+        score += 1
         return score
     else:
         raise ValueError(f"Unknown metric name: {metric_name} for spark models.")
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index 05d2a42a36..4c89f9e68c 100644
--- a/test/spark/test_utils.py
+++ b/test/spark/test_utils.py
@@ -275,18 +275,36 @@ def test_iloc_pandas_on_spark():
 
 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 loss_pair(metric, predictions, labels, groups=None):
+        dataset = to_pandas_on_spark(spark.createDataFrame(list(zip(predictions, labels)), ["prediction", "label"]))
+        spark_loss = spark_metric_loss_score(metric, dataset["prediction"], dataset["label"], groups=groups)
+        sklearn_loss = sklearn_metric_loss_score(
+            metric, np.array(predictions), np.array(labels), groups=None if groups is None else groups.to_numpy()
+        )
+        return spark_loss, sklearn_loss
+
+    # a perfectly ordered query whose predictions differ from the labels, and a
+    # query with a single document, have no loss
+    for metric in ["ndcg", "ndcg@2"]:
+        spark_loss, sklearn_loss = loss_pair(
+            metric, [10.0, 9.0, -1.0, 5.0], [3.0, 2.0, 0.0, 2.0], pd.Series([0, 0, 0, 1])
+        )
+        assert spark_loss == pytest.approx(0)
+        assert sklearn_loss == pytest.approx(0)
+
+    # with a badly ordered query added, each query counts the same, as in sklearn
+    predictions = [10.0, 9.0, -1.0, 5.0, 1.0, 5.0]
+    labels = [3.0, 2.0, 0.0, 0.0, 2.0, 2.0]
+    for metric in ["ndcg", "ndcg@2"]:
+        spark_loss, sklearn_loss = loss_pair(metric, predictions, labels, pd.Series([0, 0, 0, 1, 1, 2]))
+        assert 0 < spark_loss < 1
+        assert spark_loss == pytest.approx(sklearn_loss)
+
+    # without groups, all rows are ranked as one query
+    spark_loss, sklearn_loss = loss_pair("ndcg", predictions, labels)
+    assert 0 < spark_loss < 1
+    assert spark_loss == pytest.approx(sklearn_loss)
 
 
 def test_spark_metric_loss_score():

From c18e322fc71f3cd945d998ffa2ec66aebed0ff6d Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Fri, 9 Oct 2026 21:07:31 +0530
Subject: [PATCH 4/7] fix: score Spark ndcg per query on the executors

Join the query groups to the predictions by index and compute the NDCG of each query with groupBy().applyInPandas(), so only the mean comes back to the driver. Matching by index also handles groups that cover more rows than the predictions, as for the training loss of a cross-validation fold, which failed in test_spark_synapseml_rank.
---
 flaml/automl/spark/metrics.py | 42 +++++++++++++++++++++--------------
 test/spark/test_utils.py      | 23 +++++++++++++++++++
 2 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index e2e9fb77a8..0bd309dd37 100644
--- a/flaml/automl/spark/metrics.py
+++ b/flaml/automl/spark/metrics.py
@@ -10,7 +10,7 @@
     RegressionEvaluator,
 )
 
-from flaml.automl.spark import F, T, psDataFrame, psSeries, sparkDataFrame
+from flaml.automl.spark import F, T, ps, psDataFrame, psSeries, sparkDataFrame
 
 
 def ps_group_counts(groups: Union[psSeries, np.ndarray]) -> np.ndarray:
@@ -183,25 +183,33 @@ def spark_metric_loss_score(
     elif "ndcg" in metric_name:
         # RankingEvaluator compares arrays of ranked and relevant item ids, not
         # prediction scores and graded relevance labels, so rank the documents of
-        # each query by score here, the same way as sklearn_metric_loss_score
-        from sklearn.metrics import ndcg_score
-
+        # each query by score here, the same way as sklearn_metric_loss_score.
+        # Each query is scored on the executors, only the mean reaches the driver.
         k = int(metric_name.split("@", 1)[-1]) if "@" in metric_name else None
-        y_true = y_true.to_numpy()
-        y_predict = y_predict.to_numpy()
-        counts = [len(y_true)] if groups is None else ps_group_counts(groups)
-        score = 0
-        psum = 0
-        for c in counts:
-            if c == 1:
+        group_col = "group"
+        if groups is None:
+            # all rows are one query
+            df = df.withColumn(group_col, F.lit(0))
+        else:
+            # match the groups to the rows by index, since they can cover more rows
+            # than y_true, e.g. the training loss of a cross-validation fold
+            if not isinstance(groups, psSeries):
+                groups = ps.Series(groups)
+            df = y_predict.to_frame().join(y_true).join(groups.rename(group_col)).to_spark()
+
+        def query_ndcg(pdf):
+            import pandas as pd
+            from sklearn.metrics import ndcg_score
+
+            if len(pdf) == 1:
                 # a query with one document is always ranked perfectly
-                score -= 1
+                score = 1.0
             else:
-                score -= ndcg_score([y_true[psum : psum + c]], [y_predict[psum : psum + c]], k=k)
-            psum += c
-        score /= len(counts)
-        score += 1
-        return score
+                score = ndcg_score([pdf[label_col].to_numpy()], [pdf[prediction_col].to_numpy()], k=k)
+            return pd.DataFrame({"ndcg": [score]})
+
+        ndcg = df.groupBy(group_col).applyInPandas(query_ndcg, schema="ndcg double")
+        return 1 - ndcg.agg(F.avg("ndcg")).first()[0]
     else:
         raise ValueError(f"Unknown metric name: {metric_name} for spark models.")
 
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index 4c89f9e68c..f0442a8c78 100644
--- a/test/spark/test_utils.py
+++ b/test/spark/test_utils.py
@@ -306,6 +306,29 @@ def loss_pair(metric, predictions, labels, groups=None):
     assert 0 < spark_loss < 1
     assert spark_loss == pytest.approx(sklearn_loss)
 
+    # the groups of all rows are matched to the rows of a fold by index, as for
+    # the training loss in cross-validation
+    dataset = to_pandas_on_spark(spark.createDataFrame(list(zip(predictions, labels)), ["prediction", "label"]))
+    rows = [2, 3, 4, 5]
+    fold = dataset.loc[rows]
+    groups = ps.Series([0, 0, 1, 1, 2, 2])
+    for metric in ["ndcg", "ndcg@2"]:
+        spark_loss = spark_metric_loss_score(metric, fold["prediction"], fold["label"], groups=groups)
+        sklearn_loss = sklearn_metric_loss_score(
+            metric, np.array(predictions)[rows], np.array(labels)[rows], groups=groups.to_numpy()[rows]
+        )
+        assert 0 < spark_loss < 1
+        assert spark_loss == pytest.approx(sklearn_loss)
+
+    # a larger set with many queries
+    rng = np.random.RandomState(0)
+    groups = np.repeat(np.arange(300), rng.randint(1, 10, size=300))
+    labels = rng.randint(0, 4, size=len(groups)).astype(float)
+    predictions = rng.rand(len(groups))
+    for metric in ["ndcg", "ndcg@3"]:
+        spark_loss, sklearn_loss = loss_pair(metric, predictions.tolist(), labels.tolist(), pd.Series(groups))
+        assert spark_loss == pytest.approx(sklearn_loss)
+
 
 def test_spark_metric_loss_score():
     spark = SparkSession.builder.getOrCreate()

From 84633683d2e9d844d09cf686ed24578e9fd910cb Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Fri, 9 Oct 2026 22:02:46 +0530
Subject: [PATCH 5/7] fix: match array-like groups to Spark rows by position

---
 flaml/automl/spark/metrics.py |  6 +++++-
 test/spark/test_utils.py      | 11 +++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index 0bd309dd37..894d172715 100644
--- a/flaml/automl/spark/metrics.py
+++ b/flaml/automl/spark/metrics.py
@@ -194,7 +194,11 @@ def spark_metric_loss_score(
             # match the groups to the rows by index, since they can cover more rows
             # than y_true, e.g. the training loss of a cross-validation fold
             if not isinstance(groups, psSeries):
-                groups = ps.Series(groups)
+                groups = np.asarray(groups)
+                # array-like groups for the rows of y_true follow them by position,
+                # whatever their index; longer ones use the default 0-based index
+                index = y_true.index.to_numpy() if len(groups) == len(y_true) else None
+                groups = ps.Series(groups, index=index)
             df = y_predict.to_frame().join(y_true).join(groups.rename(group_col)).to_spark()
 
         def query_ndcg(pdf):
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index f0442a8c78..fb28bbac05 100644
--- a/test/spark/test_utils.py
+++ b/test/spark/test_utils.py
@@ -320,6 +320,17 @@ def loss_pair(metric, predictions, labels, groups=None):
         assert 0 < spark_loss < 1
         assert spark_loss == pytest.approx(sklearn_loss)
 
+    # array-like groups of a fold follow its rows by position, not by index
+    dataset = to_pandas_on_spark(
+        spark.createDataFrame(
+            list(zip([1.0, 1.0, 0.0, 10.0, 9.0, 8.0], [0.0, 0.0, 3.0, 2.0, 1.0, 0.0])), ["prediction", "label"]
+        )
+    )
+    fold = dataset.loc[rows]
+    for metric in ["ndcg", "ndcg@2"]:
+        spark_loss = spark_metric_loss_score(metric, fold["prediction"], fold["label"], groups=np.array([0, 1, 1, 1]))
+        assert spark_loss == pytest.approx(0)
+
     # a larger set with many queries
     rng = np.random.RandomState(0)
     groups = np.repeat(np.arange(300), rng.randint(1, 10, size=300))

From fd26ab3eec142ffbbecf195227e46fef9a6312af Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 11:10:25 +0530
Subject: [PATCH 6/7] fix: reject array-like groups that do not match the Spark
 rows

---
 flaml/automl/spark/metrics.py | 12 ++++++++----
 test/spark/test_utils.py      | 13 +++++++++++++
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index 894d172715..1dc56628a6 100644
--- a/flaml/automl/spark/metrics.py
+++ b/flaml/automl/spark/metrics.py
@@ -195,10 +195,14 @@ def spark_metric_loss_score(
             # than y_true, e.g. the training loss of a cross-validation fold
             if not isinstance(groups, psSeries):
                 groups = np.asarray(groups)
-                # array-like groups for the rows of y_true follow them by position,
-                # whatever their index; longer ones use the default 0-based index
-                index = y_true.index.to_numpy() if len(groups) == len(y_true) else None
-                groups = ps.Series(groups, index=index)
+                # array-like groups follow the rows of y_true by position; longer
+                # ones carry no index to find the rows of y_true in them
+                if len(groups) != len(y_true):
+                    raise ValueError(
+                        f"groups has {len(groups)} entries but there are {len(y_true)} rows. "
+                        "Pass one group per row, or a pandas-on-Spark Series indexed like the data."
+                    )
+                groups = ps.Series(groups, index=y_true.index.to_numpy())
             df = y_predict.to_frame().join(y_true).join(groups.rename(group_col)).to_spark()
 
         def query_ndcg(pdf):
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index fb28bbac05..cd3679276e 100644
--- a/test/spark/test_utils.py
+++ b/test/spark/test_utils.py
@@ -331,6 +331,19 @@ def loss_pair(metric, predictions, labels, groups=None):
         spark_loss = spark_metric_loss_score(metric, fold["prediction"], fold["label"], groups=np.array([0, 1, 1, 1]))
         assert spark_loss == pytest.approx(0)
 
+    # array-like groups of all rows cannot be matched to the rows of a fold whose
+    # data has a nonzero index, so they are rejected; a Series with that index works
+    dataset = to_pandas_on_spark(
+        pd.DataFrame({"prediction": [1.0, 10.0, 9.0, 0.0], "label": [3.0, 2.0, 1.0, 0.0]}, index=[10, 11, 12, 13])
+    )
+    fold = dataset.loc[[10, 11, 12]]
+    with pytest.raises(ValueError, match="groups has 4 entries but there are 3 rows"):
+        spark_metric_loss_score("ndcg", fold["prediction"], fold["label"], groups=np.array([0, 1, 1, 1]))
+    groups = ps.Series([0, 1, 1, 1], index=[10, 11, 12, 13])
+    for metric in ["ndcg", "ndcg@2"]:
+        spark_loss = spark_metric_loss_score(metric, fold["prediction"], fold["label"], groups=groups)
+        assert spark_loss == pytest.approx(0)
+
     # a larger set with many queries
     rng = np.random.RandomState(0)
     groups = np.repeat(np.arange(300), rng.randint(1, 10, size=300))

From 596ddfb83c80c2ad851e2f113fe0957f805d56ce Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 16:21:29 +0530
Subject: [PATCH 7/7] fix: rank ungrouped Spark ndcg with Spark's sort

Without groups all rows are one query, and applyInPandas loaded that
whole query into one pandas DataFrame on one executor. Rank the rows with
orderBy and zipWithIndex instead and sum the discounted gains with Spark
aggregations. Rows with tied predictions share the mean of their gains,
the same as sklearn's ndcg_score.
---
 flaml/automl/spark/metrics.py | 34 ++++++++++++++++++++++++++++--
 test/spark/test_utils.py      | 39 +++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/flaml/automl/spark/metrics.py b/flaml/automl/spark/metrics.py
index 1dc56628a6..6b1f026b6f 100644
--- a/flaml/automl/spark/metrics.py
+++ b/flaml/automl/spark/metrics.py
@@ -31,6 +31,36 @@ def _compute_label_from_probability(df, probability_col, prediction_col):
     return df
 
 
+def _with_rank_discounts(df, sort_col, k=None):
+    """Rank the rows of df by sort_col, highest first, with Spark's distributed sort and
+    add the DCG discount of each row's position (0 from position k on)."""
+    ranked = df.orderBy(F.desc(sort_col)).rdd.zipWithIndex().map(lambda row: (*row[0], row[1]))
+    ranked = ranked.toDF(T.StructType(df.schema.fields + [T.StructField("position", T.LongType())]))
+    discount = 1 / F.log2(F.col("position") + 2)
+    if k is not None:
+        discount = F.when(F.col("position") < k, discount).otherwise(0.0)
+    return ranked.withColumn("discount", discount)
+
+
+def _ndcg_single_query(df, label_col, prediction_col, k=None):
+    """NDCG of all rows of df as one query, computed like sklearn's ndcg_score, where rows
+    with tied predictions share the mean of their gains, without loading the rows in one place."""
+    df = df.select(label_col, prediction_col)
+    if df.count() == 1:
+        # a query with one document is always ranked perfectly
+        return 1.0
+    dcg = (
+        _with_rank_discounts(df, prediction_col, k)
+        .groupBy(prediction_col)
+        .agg((F.avg(label_col) * F.sum("discount")).alias("dcg"))
+        .agg(F.sum("dcg"))
+        .first()[0]
+    )
+    ideal_dcg = _with_rank_discounts(df, label_col, k).agg(F.sum(F.col(label_col) * F.col("discount"))).first()[0]
+    # like sklearn, a query without relevant documents scores 0
+    return dcg / ideal_dcg if ideal_dcg else 0.0
+
+
 def string_to_array(s):
     try:
         return json.loads(s)
@@ -188,8 +218,8 @@ def spark_metric_loss_score(
         k = int(metric_name.split("@", 1)[-1]) if "@" in metric_name else None
         group_col = "group"
         if groups is None:
-            # all rows are one query
-            df = df.withColumn(group_col, F.lit(0))
+            # all rows are one query, which can be too large for one pandas group
+            return 1 - _ndcg_single_query(df, label_col, prediction_col, k)
         else:
             # match the groups to the rows by index, since they can cover more rows
             # than y_true, e.g. the training loss of a cross-validation fold
diff --git a/test/spark/test_utils.py b/test/spark/test_utils.py
index cd3679276e..18180b803f 100644
--- a/test/spark/test_utils.py
+++ b/test/spark/test_utils.py
@@ -354,6 +354,45 @@ def loss_pair(metric, predictions, labels, groups=None):
         assert spark_loss == pytest.approx(sklearn_loss)
 
 
+def test_spark_ndcg_without_groups(monkeypatch):
+    from pyspark.sql import GroupedData
+
+    spark = SparkSession.builder.getOrCreate()
+
+    # without groups, all rows are one query; they are ranked with Spark's sort
+    # instead of being loaded into a single pandas group
+    def fail(*args, **kwargs):
+        raise AssertionError("all rows were loaded into one pandas group")
+
+    monkeypatch.setattr(GroupedData, "applyInPandas", fail)
+
+    def loss_pair(metric, predictions, labels):
+        df = spark.createDataFrame(list(zip(predictions.tolist(), labels.tolist())), ["prediction", "label"])
+        dataset = to_pandas_on_spark(df)
+        spark_loss = spark_metric_loss_score(metric, dataset["prediction"], dataset["label"])
+        sklearn_loss = sklearn_metric_loss_score(metric, predictions, labels)
+        return spark_loss, sklearn_loss
+
+    # many tied predictions, which share their gains as in sklearn, and cutoffs
+    # inside groups of ties
+    rng = np.random.RandomState(0)
+    labels = rng.randint(0, 4, size=2000).astype(float)
+    predictions = rng.randint(0, 50, size=2000).astype(float)
+    for metric in ["ndcg", "ndcg@10", "ndcg@301"]:
+        spark_loss, sklearn_loss = loss_pair(metric, predictions, labels)
+        assert 0 < spark_loss < 1
+        assert spark_loss == pytest.approx(sklearn_loss)
+
+    # a perfectly ordered query whose predictions differ from the labels
+    spark_loss, _ = loss_pair("ndcg", np.array([10.0, 9.0, 5.0, -1.0]), np.array([3.0, 2.0, 2.0, 0.0]))
+    assert spark_loss == pytest.approx(0)
+
+    # a query without relevant documents scores 0, as in sklearn
+    spark_loss, sklearn_loss = loss_pair("ndcg", np.array([1.0, 2.0, 3.0]), np.zeros(3))
+    assert spark_loss == pytest.approx(1)
+    assert sklearn_loss == pytest.approx(1)
+
+
 def test_spark_metric_loss_score():
     spark = SparkSession.builder.getOrCreate()
     scoreAndLabels = map(