fix: average ndcg over query groups

flaml/automl/ml.py

@@ -276,22 +276,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

@@ -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:
@@ -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
@@ -37,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)
@@ -187,43 +211,43 @@ 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
-        if "@" in metric_name:
-            k = int(metric_name.split("@", 1)[-1])
-            if groups is None:
-                evaluator = RankingEvaluator(
-                    metricName="ndcgAtK",
-                    labelCol=label_col,
-                    predictionCol=prediction_col,
-                    k=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
+        # 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.
+        # 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
+        group_col = "group"
+        if groups is None:
+            # 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:
-            evaluator = RankingEvaluator(metricName="ndcgAtK", labelCol=label_col, predictionCol=prediction_col)
-            df = _process_df(df, label_col, prediction_col)
-            score = 1 - evaluator.evaluate(df)
-        return 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 = np.asarray(groups)
+                # 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):
+            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.0
+            else:
+                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.")
 

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,126 @@ 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()
+
+    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)
+
+    # 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)
+
+    # 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)
+
+    # 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))
+    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_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(