[ml-completion] don't cast ml features to string representation to reduce memory traffic

GitOrigin-RevId: 3aa4aad3edc49d49d7a3c30b76efc41a838a6711

plugins/completion-ml-ranking/src/com/intellij/completion/ml/features/ContextFeaturesStorage.kt

@@ -17,9 +17,9 @@ class ContextFeaturesStorage(private val featuresSnapshot: Map<String, MLFeature
   override fun categoricalValue(name: String): String? = findValue<MLFeatureValue.CategoricalValue>(name)?.value
 
   override fun classNameValue(name: String): String? = findValue<MLFeatureValue.ClassNameValue>(name)
-    ?.let { MLFeaturesUtil.valueAsString(it) }
+    ?.let { MLFeaturesUtil.getClassNameSafe(it) }
 
-  override fun asMap(): Map<String, String> = featuresSnapshot.mapValues { MLFeaturesUtil.valueAsString(it.value) }
+  override fun asMap(): Map<String, String> = featuresSnapshot.mapValues { MLFeaturesUtil.getRawValue(it.value).toString() }
 
   private inline fun <reified T : MLFeatureValue> findValue(name: String): T? {
     return featuresSnapshot[name] as? T

plugins/completion-ml-ranking/src/com/intellij/completion/ml/features/MLCompletionWeigher.kt

@@ -29,19 +29,17 @@ class MLCompletionWeigher : CompletionWeigher() {
     return if (result.isEmpty()) DummyComparable.EMPTY else DummyComparable(result)
   }
 
-  private class DummyComparable(values: Map<String, MLFeatureValue>) : Comparable<Any> {
-    val representation = calculateRepresentation(values)
+  internal class DummyComparable(values: Map<String, MLFeatureValue>) : Comparable<Any> {
+    val mlFeatures = values.mapValues { MLFeaturesUtil.getRawValue(it.value) }
 
     override fun compareTo(other: Any): Int = 0
 
-    override fun toString(): String = representation
+    override fun toString(): String {
+      return mlFeatures.entries.joinToString(",", "[", "]", transform = { "${it.key}=${it.value}" })
+    }
 
     companion object {
       val EMPTY = DummyComparable(emptyMap())
-
-      private fun calculateRepresentation(values: Map<String, MLFeatureValue>): String {
-        return values.entries.joinToString(",", "[", "]", transform = { "${it.key}=${MLFeaturesUtil.valueAsString(it.value)}" })
-      }
     }
   }
 }

plugins/completion-ml-ranking/src/com/intellij/completion/ml/features/MLFeaturesUtil.kt

@@ -6,10 +6,10 @@ import com.intellij.codeInsight.completion.ml.MLFeatureValue
 import com.intellij.internal.statistic.utils.getPluginInfo
 
 internal object MLFeaturesUtil {
-  fun valueAsString(featureValue: MLFeatureValue): String {
+  fun getRawValue(featureValue: MLFeatureValue): Any {
     return when (featureValue) {
-      is MLFeatureValue.BinaryValue -> if (featureValue.value) "1" else "0"
-      is MLFeatureValue.FloatValue -> featureValue.value.toString()
+      is MLFeatureValue.BinaryValue -> if (featureValue.value) 1 else 0
+      is MLFeatureValue.FloatValue -> featureValue.value
       is MLFeatureValue.CategoricalValue -> featureValue.value
       is MLFeatureValue.ClassNameValue -> getClassNameSafe(featureValue)
     }
@@ -25,7 +25,7 @@ internal object MLFeaturesUtil {
 
   private val CLASS_NAMES_CACHE = Caffeine.newBuilder().maximumSize(100).build<String, ClassNames>()
 
-  private fun getClassNameSafe(feature: MLFeatureValue.ClassNameValue): String {
+  fun getClassNameSafe(feature: MLFeatureValue.ClassNameValue): String {
     val clazz = feature.value
     val names = CLASS_NAMES_CACHE.get(clazz.name) { if (getPluginInfo(clazz).isSafeToReport()) clazz.getNames() else THIRD_PARTY_NAME }!!
     return if (feature.useSimpleName) names.simpleName else names.fullName

plugins/completion-ml-ranking/src/com/intellij/completion/ml/util/RelevanceUtil.kt

@@ -1,5 +1,6 @@
 package com.intellij.completion.ml.util
 
+import com.intellij.completion.ml.features.MLCompletionWeigher
 import com.intellij.completion.ml.sorting.FeatureUtils
 import com.intellij.internal.statistic.utils.PluginType
 import com.intellij.internal.statistic.utils.getPluginInfo
@@ -45,7 +46,7 @@ object RelevanceUtil {
           relevanceMap.addProximityValues("swift_prox", value)
         }
         "kotlin.callableWeight" -> relevanceMap.addDataClassValues("kotlin.callableWeight", value.toString())
-        "ml_weigh" -> additionalMap.addCompoundValues("ml", value.toString())
+        "ml_weigh" -> additionalMap.addMlFeatures("ml", value)
         else -> if (acceptValue(value) || name == FeatureUtils.ML_RANK) relevanceMap[name] = value
       }
     }
@@ -66,13 +67,15 @@ object RelevanceUtil {
     return substringBefore('@')
   }
 
-  /**
-   * Proximity features now came like [samePsiFile=true, openedInEditor=false], need to convert to proper map
-   */
-  private fun MutableMap<String, Any>.addCompoundValues(prefix: String, proximity: String) {
-    val items = proximity.replace("[", "").replace("]", "").split(",")
+  private fun MutableMap<String, Any>.addMlFeatures(prefix: String, comparable: Any) {
+    if (comparable !is MLCompletionWeigher.DummyComparable) {
+      LOG.error("Unexpected value type of `$prefix`: ${comparable.javaClass.simpleName}")
+      return
+    }
 
-    this.addProperties(prefix, items)
+    for ((name, value) in comparable.mlFeatures) {
+      this["${prefix}_$name"] = value
+    }
   }
 
   private fun MutableMap<String, Any>.addProximityValues(prefix: String, proximity: Any) {