Update ranking model for java: 0.2.1

Reorder proposals only if they have recommedners scores

GitOrigin-RevId: 8ac442808264d1c4406821c54d2027d28fef3fb6

java/java-impl/intellij.java.impl.iml

@@ -59,13 +59,13 @@
     <orderEntry type="module" module-name="intellij.platform.diff.impl" />
     <orderEntry type="module-library">
       <library name="completion-ranking-java" type="repository">
-        <properties include-transitive-deps="false" maven-id="org.jetbrains.intellij.deps.completion:completion-ranking-java:0.1.7" />
+        <properties include-transitive-deps="false" maven-id="org.jetbrains.intellij.deps.completion:completion-ranking-java:0.2.1" />
         <CLASSES>
-          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/intellij/deps/completion/completion-ranking-java/0.1.7/completion-ranking-java-0.1.7.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/intellij/deps/completion/completion-ranking-java/0.2.1/completion-ranking-java-0.2.1.jar!/" />
         </CLASSES>
         <JAVADOC />
         <SOURCES>
-          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/intellij/deps/completion/completion-ranking-java/0.1.7/completion-ranking-java-0.1.7-sources.jar!/" />
+          <root url="jar://$MAVEN_REPOSITORY$/org/jetbrains/intellij/deps/completion/completion-ranking-java/0.2.1/completion-ranking-java-0.2.1-sources.jar!/" />
         </SOURCES>
       </library>
     </orderEntry>

plugins/stats-collector/src/com/intellij/completion/sorting/MLSorter.kt

@@ -9,6 +9,7 @@ import com.intellij.codeInsight.lookup.LookupManager
 import com.intellij.codeInsight.lookup.impl.LookupImpl
 import com.intellij.completion.ml.common.PrefixMatchingUtil
 import com.intellij.completion.settings.CompletionMLRankingSettings
+import com.intellij.lang.Language
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.util.Pair
 import com.intellij.openapi.util.registry.Registry
@@ -33,6 +34,7 @@ class MLSorter : CompletionFinalSorter() {
   }
 
   private val cachedScore: MutableMap<LookupElement, ItemRankInfo> = IdentityHashMap()
+  private lateinit var sortingRestrictions: SortingRestriction
   private val reorderOnlyTopItems: Boolean = Registry.`is`("completion.ml.reorder.only.top.items", true)
 
   override fun getRelevanceObjects(items: MutableIterable<LookupElement>): Map<LookupElement, List<Pair<String, Any>>> {
@@ -77,6 +79,10 @@ class MLSorter : CompletionFinalSorter() {
     val queryLength = lookup.queryLength()
     val prefix = lookup.prefix()
 
+    if (!this::sortingRestrictions.isInitialized) {
+      sortingRestrictions = SortingRestriction.forLanguage(lookupStorage.language)
+    }
+
     val element2score = mutableMapOf<LookupElement, Double?>()
     val elements = items.toList()
 
@@ -130,7 +136,10 @@ class MLSorter : CompletionFinalSorter() {
         PrefixMatchingUtil.calculateFeatures(element, prefix, additional)
       }
       val score = tracker.measure {
-        calculateElementScore(rankingModel, element, position, features.withElementFeatures(relevance, additional), queryLength)
+        val elementFeatures = features.withElementFeatures(relevance, additional)
+        val score = calculateElementScore(rankingModel, element, position, elementFeatures, queryLength)
+        sortingRestrictions.itemScored(elementFeatures)
+        return@measure score
       }
       element2score[element] = score
 
@@ -145,7 +154,7 @@ class MLSorter : CompletionFinalSorter() {
                              element2score: Map<LookupElement, Double?>,
                              positionsBefore: Map<LookupElement, Int>,
                              lookupStorage: MutableLookupStorage): Iterable<LookupElement> {
-    val mlScoresUsed = element2score.values.none { it == null }
+    val mlScoresUsed = element2score.values.none { it == null } && sortingRestrictions.shouldSort()
     if (LOG.isDebugEnabled) {
       LOG.debug("ML sorting in completion used=$mlScoresUsed for language=${lookupStorage.language.id}")
     }
@@ -250,6 +259,44 @@ class MLSorter : CompletionFinalSorter() {
       }
     }
   }
+
+  interface SortingRestriction {
+    companion object {
+      fun forLanguage(language: Language): SortingRestriction {
+        if (language.id.equals("Java", ignoreCase = true)) {
+          return SortOnlyWithRecommendersScore()
+        }
+        return SortAll()
+      }
+    }
+
+    fun itemScored(features: RankingFeatures)
+
+    fun shouldSort(): Boolean
+  }
+
+  private class SortAll : SortingRestriction {
+    override fun shouldSort(): Boolean = true
+    override fun itemScored(features: RankingFeatures) {}
+  }
+
+  private class SortOnlyWithRecommendersScore : SortingRestriction {
+    companion object {
+      private val REC_FEATURES_NAMES: List<String> = listOf("ml_rec-instances_probability", "ml_rec-statics2_probability")
+    }
+
+    private var recommendersScoreFound: Boolean = false
+
+    override fun itemScored(features: RankingFeatures) {
+      if (!recommendersScoreFound) {
+        recommendersScoreFound = REC_FEATURES_NAMES.any { features.hasFeature(it) }
+      }
+    }
+
+    override fun shouldSort(): Boolean {
+      return recommendersScoreFound
+    }
+  }
 }
 
 private data class ItemRankInfo(val positionBefore: Int, val mlRank: Double?, val prefixLength: Int)