Serialize argument type information when available

3 files changed, 13 insertions, 14 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/RankingExpression.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/RankingExpression.java
index 86d1c2f5066..722520fea08 100755
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/RankingExpression.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/RankingExpression.java
@@ -273,16 +273,6 @@ public class RankingExpression implements Serializable {
     }
 
     /**
-     * Returns the expression type rank property name for a given expression name.
-     *
-     * @param expressionName the expression name to return as an expression type property
-     * @return the property name
-     */
-    public static String propertyTypeName(String expressionName) {
-        return "rankingExpression(" + expressionName + ").type";
-    }
-
-    /**
      * Validates the type correctness of the given expression with the given context and
      * returns the type this expression will produce from the given type context
      *
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/rule/SerializationContext.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/rule/SerializationContext.java
index 58131cf42a3..94d663b4954 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/rule/SerializationContext.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/rule/SerializationContext.java
@@ -3,6 +3,8 @@ package com.yahoo.searchlib.rankingexpression.rule;
 
 import com.google.common.collect.ImmutableMap;
 import com.yahoo.searchlib.rankingexpression.ExpressionFunction;
+import com.yahoo.searchlib.rankingexpression.RankingExpression;
+import com.yahoo.tensor.TensorType;
 
 import java.util.Collection;
 import java.util.Collections;
@@ -80,9 +82,14 @@ public class SerializationContext extends FunctionReferenceContext {
         serializedFunctions.put(name, expressionString);
     }
 
-    /** Adds the serialization of the type of a function */
-    public void addFunctionTypeSerialization(String name, String expressionString) {
-        serializedFunctions.put(name, expressionString);
+    /** Adds the serialization of the an argument type to a function */
+    public void addArgumentTypeSerialization(String functionName, String argumentName, TensorType type) {
+        serializedFunctions.put("rankingExpression(" + functionName + ")." + argumentName + ".type", type.toString());
+    }
+
+    /** Adds the serialization of the return type of a function */
+    public void addFunctionTypeSerialization(String functionName, TensorType type) {
+        serializedFunctions.put("rankingExpression(" + functionName + ").type", type.toString());
     }
 
     @Override
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/aggregation/GroupingSerializationTest.java b/searchlib/src/test/java/com/yahoo/searchlib/aggregation/GroupingSerializationTest.java
index 118eba2cd96..969bc318391 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/aggregation/GroupingSerializationTest.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/aggregation/GroupingSerializationTest.java
@@ -15,6 +15,7 @@ import org.junit.Test;
 import java.io.*;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
 
 import static org.junit.Assert.fail;
@@ -40,7 +41,8 @@ public class GroupingSerializationTest {
             t.assertMatch(new FloatResultNode(7.3));
             t.assertMatch(new StringResultNode("7.3"));
             t.assertMatch(new StringResultNode(
-                    new String(new byte[]{(byte)0xe5, (byte)0xa6, (byte)0x82, (byte)0xe6, (byte)0x9e, (byte)0x9c})));
+                    new String(new byte[]{(byte)0xe5, (byte)0xa6, (byte)0x82, (byte)0xe6, (byte)0x9e, (byte)0x9c},
+                               StandardCharsets.UTF_8)));
             t.assertMatch(new RawResultNode(new byte[]{'7', '.', '4'}));
             t.assertMatch(new IntegerBucketResultNode());
             t.assertMatch(new FloatBucketResultNode());