Rewrite large constant references from foo to constant(foo)

4 files changed, 19 insertions, 11 deletions

diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConstantTensorTransformer.java b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConstantTensorTransformer.java
index eb76446c045..caf5f0442eb 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConstantTensorTransformer.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/ConstantTensorTransformer.java
@@ -2,16 +2,15 @@
 package com.yahoo.searchdefinition.expressiontransforms;
 
 import com.yahoo.searchdefinition.FeatureNames;
+import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.TensorValue;
 import com.yahoo.searchlib.rankingexpression.evaluation.Value;
 import com.yahoo.searchlib.rankingexpression.rule.CompositeNode;
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
-import com.yahoo.searchlib.rankingexpression.rule.NameNode;
 import com.yahoo.searchlib.rankingexpression.rule.ReferenceNode;
 import com.yahoo.searchlib.rankingexpression.transform.ExpressionTransformer;
 
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -21,8 +20,6 @@ import java.util.List;
  */
 public class ConstantTensorTransformer extends ExpressionTransformer<RankProfileTransformContext> {
 
-    public static final String CONSTANT = "constant";
-
     @Override
     public ExpressionNode transform(ExpressionNode node, RankProfileTransformContext context) {
         if (node instanceof ReferenceNode) {
@@ -52,17 +49,18 @@ public class ConstantTensorTransformer extends ExpressionTransformer<RankProfile
     }
 
     private ExpressionNode transformConstantReference(ReferenceNode node, RankProfileTransformContext context) {
+        Reference constantReference = FeatureNames.asConstantFeature(node.getName());
         Value value = context.constants().get(node.getName());
         if (value == null || value.type().rank() == 0) {
+            if (context.rankProfile().rankingConstants().get(node.getName()) != null) // Large constants: Transform reference but don't add value
+                return new ReferenceNode(constantReference);
             return node;
         }
         TensorValue tensorValue = (TensorValue)value;
-        String featureName = CONSTANT + "(" + node.getName() + ")";
         String tensorType = tensorValue.asTensor().type().toString();
-        context.rankProperties().put(featureName + ".value", tensorValue.toString());
-        context.rankProperties().put(featureName + ".type", tensorType);
-        // TODO: This allows us to reference constant "a" as "a" instead of "constant(a)", but we shouldn't allow that
-        return new ReferenceNode(CONSTANT, Arrays.asList(new NameNode(node.getName())), null);
+        context.rankProperties().put(constantReference.toString() + ".value", tensorValue.toString());
+        context.rankProperties().put(constantReference.toString() + ".type", tensorType);
+        return new ReferenceNode(constantReference);
     }
 
 }
diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/RankProfileTransformContext.java b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/RankProfileTransformContext.java
index 434cae04747..2c0e1eaa56a 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/RankProfileTransformContext.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/RankProfileTransformContext.java
@@ -3,7 +3,9 @@ package com.yahoo.searchdefinition.expressiontransforms;
 
 import ai.vespa.rankingexpression.importer.configmodelview.ImportedMlModels;
 import com.yahoo.search.query.profile.QueryProfileRegistry;
+import com.yahoo.searchdefinition.MapEvaluationTypeContext;
 import com.yahoo.searchdefinition.RankProfile;
+import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.Value;
 import com.yahoo.searchlib.rankingexpression.transform.TransformContext;
 
@@ -22,6 +24,7 @@ public class RankProfileTransformContext extends TransformContext {
     private final ImportedMlModels importedModels;
     private final Map<String, RankProfile.RankingExpressionFunction> inlineFunctions;
     private final Map<String, String> rankProperties = new HashMap<>();
+    private final MapEvaluationTypeContext types;
 
     public RankProfileTransformContext(RankProfile rankProfile,
                                        QueryProfileRegistry queryProfiles,
@@ -33,6 +36,7 @@ public class RankProfileTransformContext extends TransformContext {
         this.queryProfiles = queryProfiles;
         this.importedModels = importedModels;
         this.inlineFunctions = inlineFunctions;
+        this.types = rankProfile.typeContext(queryProfiles);
     }
 
     public RankProfile rankProfile() { return rankProfile; }
@@ -41,4 +45,10 @@ public class RankProfileTransformContext extends TransformContext {
     public Map<String, RankProfile.RankingExpressionFunction> inlineFunctions() { return inlineFunctions; }
     public Map<String, String> rankProperties() { return rankProperties; }
 
+    /**
+     * Returns the types known in this context. We may have type information for references
+     * for which no value is available
+     */
+    public MapEvaluationTypeContext types() { return types; }
+
 }
diff --git a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/TensorTransformer.java b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/TensorTransformer.java
index fe232299363..5d03c323803 100644
--- a/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/TensorTransformer.java
+++ b/config-model/src/main/java/com/yahoo/searchdefinition/expressiontransforms/TensorTransformer.java
@@ -69,7 +69,7 @@ public class TensorTransformer extends ExpressionTransformer<RankProfileTransfor
         ExpressionNode arg1 = node.children().get(0);
         Optional<String> dimension = dimensionName(node.children().get(1));
         if (dimension.isPresent()) {
-            TensorType type = arg1.type(context.rankProfile().typeContext(context.queryProfiles()));
+            TensorType type = arg1.type(context.types());
             if (type.dimension(dimension.get()).isPresent()) {
                 return replaceMaxAndMinFunction(node);
             }
diff --git a/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java b/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
index a75699d2a1d..34b727f9f4e 100644
--- a/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
+++ b/config-model/src/test/java/com/yahoo/searchdefinition/processing/VespaMlModelTestCase.java
@@ -30,7 +30,7 @@ public class VespaMlModelTestCase {
             "rankingExpression(foo1).rankingScript : reduce(reduce(input1 * input2, sum, name) * constant(constant1), max, x) * 3.0\n" +
             "rankingExpression(foo1).input2.type : tensor(x[3])\n" +
             "rankingExpression(foo1).input1.type : tensor(name{},x[3])\n" +
-            "rankingExpression(foo2).rankingScript : max(reduce(input1 * input2, sum, name) * constant1asLarge,x) * 3.0\n" +
+            "rankingExpression(foo2).rankingScript : reduce(reduce(input1 * input2, sum, name) * constant(constant1asLarge), max, x) * 3.0\n" +
             "rankingExpression(foo2).input2.type : tensor(x[3])\n" +
             "rankingExpression(foo2).input1.type : tensor(name{},x[3])\n";