zhengruifeng commented on code in PR #36920:
URL: https://github.com/apache/spark/pull/36920#discussion_r903179811
##########
mllib/src/main/scala/org/apache/spark/mllib/evaluation/RankingMetrics.scala:
##########
@@ -154,54 +151,68 @@ class RankingMetrics[T: ClassTag] @Since("3.4.0") (
@Since("1.2.0")
def ndcgAt(k: Int): Double = {
require(k > 0, "ranking position k should be positive")
- predictionAndLabels.map { case (pred, lab, rel) =>
- val useBinary = rel.isEmpty
- val labSet = lab.toSet
- val relMap = lab.zip(rel).toMap
- if (useBinary && lab.size != rel.size) {
- logWarning(
- "# of ground truth set and # of relevance value set should be equal,
" +
- "check input data")
- }
-
- if (labSet.nonEmpty) {
- val labSetSize = labSet.size
- val n = math.min(math.max(pred.length, labSetSize), k)
- var maxDcg = 0.0
- var dcg = 0.0
- var i = 0
- while (i < n) {
- if (useBinary) {
- // Base of the log doesn't matter for calculating NDCG,
- // if the relevance value is binary.
- val gain = 1.0 / math.log(i + 2)
- if (i < pred.length && labSet.contains(pred(i))) {
- dcg += gain
- }
- if (i < labSetSize) {
- maxDcg += gain
+ predictionAndLabels
+ .map {
+ case (pred: Array[T], lab: Array[T]) =>
Review Comment:
I think `The previous one` maybe more concise.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
