[Raashish Aggarwal (raashish1601)]

Why are these changes needed?

sklearn_metric_loss_score handles ndcg@k per query group: it splits the rows with group_counts(groups) and averages the NDCG of each query. Plain ndcg, which is the default metric for task="rank", ignores groups and scores all validation rows as one ranked list. Documents from different queries are then compared with each other, so a model that ranks every query perfectly still gets a loss above 0:

import numpy as np
from flaml.automl.ml import sklearn_metric_loss_score

y_true = np.array([3, 2, 1, 0, 3, 2, 1, 0])
y_pred = np.array([10, 9, 8, 7, 4, 3, 2, 1])  # each query ranked perfectly
groups = np.array([0, 0, 0, 0, 1, 1, 1, 1])
sklearn_metric_loss_score("ndcg", y_pred, y_true, groups=groups)    # 0.0706
sklearn_metric_loss_score("ndcg@8", y_pred, y_true, groups=groups)  # 0.0

The other direction fails too: ndcg@k without groups raises ValueError: Computing NDCG is only meaningful when there is more than 1 document, because group_counts(None) returns a single group of size 1.

This change uses the same code for both: when groups is given, it averages the NDCG of each query (with k=None for plain ndcg); when it isn't, it scores one list (with k for ndcg@k). ndcg@k with groups gives the same numbers as before.

The Spark path in flaml/automl/spark/metrics.py also ignores groups for plain ndcg. I left it alone because I can't run the Spark tests locally.

Related issue number

None. Found while comparing the metric losses against sklearn.

Checks

• I've used pre-commit to lint the changes in this PR (note the same in integrated in our CI checks).
• I've included any doc changes needed for https://microsoft.github.io/FLAML/. No doc change needed.
• I've added tests (if relevant) corresponding to the changes introduced in this PR. test/automl/test_ndcg_groups.py fails on main (plain ndcg with groups, ndcg@3 without groups) and passes with this change. The rank test in test/automl/test_split.py still passes.
• I've made sure all auto checks have passed.

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

1. flaml/automl/ml.py:279-289: When a ranking query has one document, plain ndcg now calls sklearn.metrics.ndcg_score on a single row. That raises ValueError: Computing NDCG is only meaningful when there is more than 1 document; _eval_estimator converts it to infinite validation loss. The previous plain-ndcg path returned a finite score for the same dataset. Please handle singleton queries explicitly and cover one alongside a multi-document query in a regression test.
2. flaml/automl/spark/metrics.py:222-225: Spark predictions route through spark_metric_loss_score, where plain ndcg still ignores query groups. The default rank metric remains inaccurate on this supported path. Please make grouped scoring consistent with the sklearn path, with Spark coverage where possible.

The grouped ndcg@k behavior and new ungrouped ndcg@k support otherwise look sound.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

Thanks, both points addressed in c1aae97.

1. A query with one document now counts as NDCG 1 in sklearn_metric_loss_score (ndcg_score raises for it), for both ndcg and ndcg@k. test_ndcg_query_with_one_document adds a single-document query next to the two multi-document ones.
2. spark_metric_loss_score now uses the query groups for plain ndcg too, with the same structure as the sklearn path (per-group k is the group size when no @k is given). test_spark_ndcg_query_groups in test/spark/test_utils.py checks that the grouped loss averages the queries for ndcg and ndcg@3. On Spark a single row does not raise, so I did not special-case it there.

test/automl/test_ndcg_groups.py and the Spark metric tests pass locally (pyspark 3.5), and black/ruff at the pre-commit versions are clean.

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

The scikit-learn singleton-query fix resolves my first finding, but the new Spark path still does not measure query ranking correctly. At flaml/automl/spark/metrics.py:207-218, each document is wrapped as a separate one-element [prediction_score] / [relevance_label] pair by _process_df. Spark's RankingEvaluator interprets these arrays as ranked item IDs and relevant item IDs, not scores and relevance values; it compares their equality one row at a time rather than ranking documents within each query. For example, a query with y_true=[3,2] and y_pred=[10,9] is perfectly ordered (sklearn loss 0), but Spark receives [10] vs [3] and [9] vs [2], giving NDCG 0/loss 1. A singleton query with label 2 and prediction 5 similarly disagrees with the new sklearn handling. The new Spark test only gets its expected 0.5 grouped loss because its first query sets prediction equal to label on every row. Please compute ranking per query from predicted scores and relevance labels, and test a perfectly ordered query whose prediction values differ from the labels.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

You're right, RankingEvaluator treats those arrays as item ids, so the Spark path was not ranking anything. In bdcf7d8 the Spark ndcg branch now ranks the documents of each query by their predicted score and computes NDCG on the graded labels. It uses sklearn.metrics.ndcg_score per query, as the sklearn path does. Single-document queries count as 1. Without groups, all rows are one query, which is also what the sklearn path does. _process_df was only used here, so I removed it.

The Spark test now covers a perfectly ordered query whose predictions differ from the labels ([10, 9, -1] vs [3, 2, 0]) and a single-document query with label 2 and prediction 5. Both give loss 0. A mixed case and the no-groups case are checked against sklearn_metric_loss_score. The new test fails with the previous Spark code (loss 1.0 instead of 0).

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

The new Spark implementation fixes the item-ID versus score semantics I flagged, and its parity tests cover the previously wrong results. One scalability regression remains at flaml/automl/spark/metrics.py:191-204: both y_true.to_numpy() and y_predict.to_numpy() collect the entire distributed validation set into the driver's memory, including when there are no groups. The prior Spark path evaluated the dataset with Spark's distributed RankingEvaluator. Spark's own Series.to_numpy() documentation warns that it should only be used for small results because all data is loaded into driver memory. A large ranking validation set will now exhaust driver memory before NDCG can be scored. Please keep per-query NDCG computation distributed (or otherwise bound collected data) and include coverage beyond tiny local Spark inputs.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

Agreed, collecting both columns to the driver doesn't scale. I'd like to keep the per-query scoring distributed this way: attach the group id to the Spark frame, run groupBy(group).applyInPandas(...) so each executor computes the NDCG of its own queries with the same per-query rule as the sklearn path (singleton queries count as 1, k from @k), and then take mean over the groups in Spark. Only one scalar comes back to the driver. Without groups, the whole set is one query, so that case would still need a single partition. I'd use RankingEvaluator there only if it can take score/label arrays, and otherwise leave it as a documented limitation. Does that direction work for you?

[Raashish Aggarwal (raashish1601)]

The ubuntu 3.11 failure in test_spark_synapseml_rank came from this PR. In Spark cross-validation the training loss gets the groups of all rows, but y_train only has the fold's rows, so the positional slices ran past the end and ndcg_score got an empty query.

c18e322 fixes it and also follows the direction I described above for the driver memory point. The groups are joined to the predictions by index, each query is scored with groupBy(group).applyInPandas(...) on the executors, and only the mean comes back to the driver. Without groups all rows are still one query, scored in a single task. The Spark test now also covers groups that span more rows than the fold, and a set of 300 queries checked against sklearn_metric_loss_score.

I reproduced the failure on Linux with pyspark 3.5.1 (with a larger time budget so the first trial runs on my machine). With the change, test_spark_synapseml_rank, the Spark ndcg tests and test/automl/test_ndcg_groups.py pass.

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

The executor-side NDCG computation addresses my prior driver-memory finding. One scoring regression remains at flaml/automl/spark/metrics.py:196-210: when groups is a NumPy array, ps.Series(groups) creates a fresh 0-based index, but join() matches that index to the potentially nonzero or nonconsecutive indices of the Spark validation rows. GenericTask.validate_data can supply NumPy group arrays, and a validation fold can retain original row indices. For a four-row fold indexed [2,3,4,5] with positional groups [0,1,1,1], the join assigns [1,1,NULL,NULL], combining the wrong documents into queries. With labels [3,2,1,0] and predictions [0,10,9,8], the expected mean query NDCG loss is 0, but those joined groups yield about 0.0433. The new fold test supplies an already-indexed ps.Series, so it misses this path. Please align array-like groups with validation rows by position (or preserve their original indices) and add a regression test for a sliced fold with a NumPy groups argument.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

Good catch. In 8463368, array-like groups with one entry per row of y_true are now matched to the rows by position: they get the index of y_true before the join, so a fold indexed [2, 3, 4, 5] keeps its groups. Groups that cover more rows than y_true still use the 0-based index as before, and a ps.Series is still joined on its own index. Only the index of y_true goes to the driver, which is the same length as the groups array the driver already holds.

test_spark_ndcg_query_groups now has your example: a fold indexed [2, 3, 4, 5] with NumPy groups [0, 1, 1, 1], labels [3, 2, 1, 0] and predictions [0, 10, 9, 8]. It gives 0 for ndcg and ndcg@2, and 0.0433 without the change. I ran the Spark tests with pyspark 4.0.1.

The failing build jobs on c18e322 stop while loading test/conftest.py with ImportError: pyarrow requires NumPy 2.0 or newer, found 1.26.4. That comes from the CI environment (a newer pyarrow next to the pinned NumPy 1.26.4), not from this change.

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

The new equal-length NumPy-group handling resolves my previous finding, but the len(groups) != len(y_true) branch at flaml/automl/spark/metrics.py:198-204 still silently assumes an original 0-based row index. An array of groups has positional meaning, while ps.Series(groups, index=None) gives it indices 0..N-1; the following join() matches those indices to the original Spark row indices. A validation/training fold from a dataset indexed [10,11,12,13] can contain rows [10,11,12] with a full-length group array [0,1,1,1]. The join assigns null groups to all three rows instead of [0,1,1]; for labels [3,2,1] and predictions [1,10,9], the expected mean-query loss is 0 but grouping those rows together yields about 0.1325. The new test covers a fold from a default 0..N-1 index only. Please preserve the original full-data index when attaching longer positional arrays, or reject inputs that cannot be aligned instead of returning an incorrect score; test a nonzero/nonconsecutive original index with a full-length group array.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

Agreed, a longer positional array has no index to find the fold's rows in. In fd26ab3, array-like groups must now have one entry per row of y_true; otherwise spark_metric_loss_score raises a ValueError that asks for one group per row or a pandas-on-Spark Series indexed like the data. A ps.Series is still joined on its own index, which is what the Spark cross-validation passes (it joins groups to the data, so they are already a Series there).

The test now has your example: a dataset indexed [10, 11, 12, 13], the fold [10, 11, 12] and NumPy groups [0, 1, 1, 1] raise, and the same groups as a Series with that index give 0 for ndcg and ndcg@2. The Spark ndcg tests (pyspark 4.0.1, Linux) and test/automl/test_ndcg_groups.py pass.

[Li Jiang (thinkall)]

Overall review of the complete PR: changes requested.

The new guard resolves the incorrect score for full-length group arrays with custom indices. One scalability issue remains at flaml/automl/spark/metrics.py:189-223: when groups is None, F.lit(0) assigns every validation row to the same query, then groupBy(...).applyInPandas(query_ndcg, ...) materializes the entire query as one pandas DataFrame on one executor. Spark's GroupedData.applyInPandas documentation warns that all data of a group is loaded into memory, with OOM risk for a large group. An ungrouped ranking set of millions of rows can therefore still run out of memory, now on an executor rather than the driver; the new 300-query test does not exercise this case. Please compute the ungrouped NDCG without collecting all rows into a single pandas group, or explicitly bound and reject ungrouped inputs that cannot fit in memory rather than failing during execution.

The current build failures on Ubuntu 3.11/3.12 and Windows 3.12 are in test startup (pyarrow requires NumPy 2.0 or newer, found 1.26.4), before the changed tests run; they do not establish an NDCG test failure.

Posted by thinkall-agent-auto-reviewer

[Raashish Aggarwal (raashish1601)]

Without groups the rows are now ranked with Spark's sort (orderBy + zipWithIndex) and the DCG and ideal DCG are summed with Spark aggregations, so no single pandas group holds all rows. Tied predictions share their mean gain, as in sklearn's ndcg_score. Added test_spark_ndcg_without_groups, which blocks applyInPandas and checks 2000 rows with many ties against sklearn for ndcg, ndcg@10 and ndcg@301. In 596ddfb.

[Li Jiang (thinkall)]

Overall review of the complete PR: approved.

The default plain ndcg now averages per-query ranking scores, and ndcg@k also supports the ungrouped case. The subsequent fixes cover singleton queries, the Spark item-ID versus score mismatch, per-query executor-side scoring, and safe rejection of full-length group arrays that cannot be aligned to arbitrary fold indices. The ungrouped Spark path now computes position discounts with a distributed sort and handles prediction ties rather than loading the full query into one pandas group. The tests cover those paths against scikit-learn; no remaining blocking code issue found.

CI is not fully green: some build jobs have failed during test setup on the PyArrow/NumPy mismatch, and remaining jobs were still running at review time. Approval of the code is not a claim that all checks passed.

Posted by thinkall-agent-auto-reviewer

[Copilot]

🟡 Changes recommended

Singleton handling incorrectly rewards zero-relevance queries and still rejects ungrouped one-document queries in the local scorer.

2 open issues

• Handle ungrouped singleton queries in NDCG fallback · New
• Fix zero-relevance singleton NDCG shortcut handling · New

What changed in this PR

Aligns FLAML’s ranking losses so NDCG is averaged per query when groups are supplied and computed over one list otherwise.

Changes:

• Unifies local ndcg and ndcg@k scoring.
• Replaces Spark ranking evaluation with query-aware scoring and distributed ungrouped computation.
• Adds regression tests for grouping, alignment, ties, cutoffs, and singletons.

File	Description
test/​spark/​test_utils.py	Tests Spark scoring and parity with local metrics.
test/​automl/​test_ndcg_groups.py	Tests grouped, ungrouped, and singleton NDCG.
flaml/​automl/​spark/​metrics.py	Implements query-aware Spark NDCG.
flaml/​automl/​ml.py	Unifies local NDCG handling.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Singleton normalization and interleaved query handling still produce incorrect ranking losses.

5 open issues

• Support interleaved query groups before NDCG slicing · New
• Remove singleton shortcut bypassing zero-ideal normalization · New
• Handle zero-relevance singletons correctly in grouped Spark NDCG · New
• Fix zero-relevance singleton NDCG shortcut handling
• Handle ungrouped singleton queries in NDCG fallback

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The distributed Spark evaluator rewrite needs runtime validation and human review.

6 open issues

• Handle zero-relevance singletons correctly in grouped Spark NDCG
• Remove singleton shortcut bypassing zero-ideal normalization
• Support interleaved query groups before NDCG slicing
• Fix zero-relevance singleton NDCG shortcut handling
• Handle ungrouped singleton queries in NDCG fallback
• PR description omits Spark evaluation changes and validation · New

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Spark’s singleton shortcuts incorrectly award perfect NDCG to zero-relevance documents.

8 open issues

• Remove incorrect zero-relevance singleton NDCG shortcut · New
• Fix grouped singleton scoring for zero relevance · New
• Handle zero-relevance singletons correctly in grouped Spark NDCG
• Remove singleton shortcut bypassing zero-ideal normalization
• Support interleaved query groups before NDCG slicing
• Fix zero-relevance singleton NDCG shortcut handling
• Handle ungrouped singleton queries in NDCG fallback
• PR description omits Spark evaluation changes and validation

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

The substantial Spark scoring rewrite needs Spark-enabled validation and resolution of the retained correctness issue before approval.

9 open issues

• Remove singleton shortcut for zero relevance and k=0 · New
• Fix grouped singleton scoring for zero relevance
• Remove incorrect zero-relevance singleton NDCG shortcut
• Handle zero-relevance singletons correctly in grouped Spark NDCG
• Remove singleton shortcut bypassing zero-ideal normalization
• Support interleaved query groups before NDCG slicing
• Fix zero-relevance singleton NDCG shortcut handling
• Handle ungrouped singleton queries in NDCG fallback
• PR description omits Spark evaluation changes and validation

_{Review effort: Balanced}

[Li Jiang (thinkall)]

Re-reviewed commits 0db54e731b8fd5257361eaeaaa27608052a20bb9 -> a94a246710d0c87b716097ffd3c84d6b72c0e75c. The only delta is main's website/yarn.lock update of follow-redirects 1.16.0 -> 1.16.1. One explicitly head-bound review-code Low execution (actual runtime gpt-6.1-sol) found no new issue in this identical inherited patch. Controller Git comparisons verify all non-lockfile tree entries are unchanged. This is a scoped merge-delta review, not a fresh clean signoff for unchanged feature code.

Version/URL/caret/SRI shape checks pass. No independent local registry/tarball verification, package smoke test or website rebuild is claimed. Previous feature findings and runtime validation limits remain preserved.

At the 2026-10-26T12:13:44.943946+00:00 gate snapshot, actual current-head workflows were: docs completed/success; Build in_progress/no conclusion; Code formatting completed/success. Copilot has Balanced review without an Approval recommended verdict; 9 review thread(s) remain unresolved. Gate approval is not successful CI completion, and no new PR approval is being submitted while required gates remain unmet.

Posted by thinkall-agent-auto-reviewer