Merge pull request #9003 from vespa-engine/bratseth/tensor-value-type

Bratseth/tensor value type

4 files changed, 17 insertions, 2 deletions

diff --git a/searchlib/abi-spec.json b/searchlib/abi-spec.json
index 79c633b9617..b8c51f4e33d 100644
--- a/searchlib/abi-spec.json
+++ b/searchlib/abi-spec.json
@@ -886,6 +886,7 @@
       "public final com.yahoo.searchlib.rankingexpression.rule.LambdaFunctionNode lambdaFunction()",
       "public final com.yahoo.tensor.functions.Reduce$Aggregator tensorReduceAggregator()",
       "public final com.yahoo.tensor.TensorType tensorTypeArgument()",
+      "public final com.yahoo.tensor.TensorType$Value optionalTensorValueTypeParameter()",
       "public final void tensorTypeDimension(com.yahoo.tensor.TensorType$Builder)",
       "public final java.lang.String tensorFunctionName()",
       "public final com.yahoo.searchlib.rankingexpression.rule.Function unaryFunctionName()",
diff --git a/searchlib/src/main/javacc/RankingExpressionParser.jj b/searchlib/src/main/javacc/RankingExpressionParser.jj
index 2f173ad0266..c83de4ced0a 100755
--- a/searchlib/src/main/javacc/RankingExpressionParser.jj
+++ b/searchlib/src/main/javacc/RankingExpressionParser.jj
@@ -598,9 +598,12 @@ Reduce.Aggregator tensorReduceAggregator() :
 
 TensorType tensorTypeArgument() :
 {
-    TensorType.Builder builder = new TensorType.Builder();
+    TensorType.Builder builder;
+    TensorType.Value valueType;
 }
 {
+    valueType = optionalTensorValueTypeParameter()
+    { builder = new TensorType.Builder(valueType); }
     <LBRACE>
     ( tensorTypeDimension(builder) ) ?
     ( <COMMA> tensorTypeDimension(builder) ) *
@@ -608,6 +611,15 @@ TensorType tensorTypeArgument() :
     { return builder.build(); }
 }
 
+TensorType.Value optionalTensorValueTypeParameter() :
+{
+    String valueType = "double";
+}
+{
+    ( <LT> valueType = identifier() <GT> )?
+    { return TensorTypeParser.toValueType(valueType); }
+}
+
 // NOTE: Only indexed bound dimensions are parsed currently, as that is what we need
 void tensorTypeDimension(TensorType.Builder builder) :
 {
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
index f2122bb5da9..f7e38862883 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTestCase.java
@@ -238,6 +238,8 @@ public class EvaluationTestCase {
                                "{{x:0}:1}", "{}", "{{y:0,z:0}:1}");
         tester.assertEvaluates("tensor(x{}):{}",
                                "tensor0 * tensor1", "{ {x:0}:3 }", "tensor(x{}):{ {x:1}:5 }");
+        tester.assertEvaluates("tensor<float>(x{}):{}",
+                               "tensor0 * tensor1", "{ {x:0}:3 }", "tensor<float>(x{}):{ {x:1}:5 }");
         tester.assertEvaluates("{ {x:0}:15 }",
                                "tensor0 * tensor1", "{ {x:0}:3 }", "{ {x:0}:5 }");
         tester.assertEvaluates("{ {x:0,y:0}:15 }",
diff --git a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTester.java b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTester.java
index ba0db4de5e1..488930a8eb9 100644
--- a/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTester.java
+++ b/searchlib/src/test/java/com/yahoo/searchlib/rankingexpression/evaluation/EvaluationTester.java
@@ -40,7 +40,7 @@ public class EvaluationTester {
         int argumentIndex = 0;
         for (String argumentString : tensorArgumentStrings) {
             Tensor argument;
-            if (argumentString.startsWith("tensor(")) // explicitly decided type
+            if (argumentString.startsWith("tensor")) // explicitly decided type
                 argument = Tensor.from(argumentString);
             else // use mappedTensors+dimensions in tensor to decide type
                 argument = Tensor.from(typeFrom(argumentString, mappedTensors), argumentString);