Propagate type information

4 files changed, 15 insertions, 2 deletions

diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/Context.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/Context.java
index 4102d4078e6..4e046df11ca 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/Context.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/Context.java
@@ -5,6 +5,7 @@ import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.rule.Arguments;
 import com.yahoo.searchlib.rankingexpression.rule.ExpressionNode;
 import com.yahoo.tensor.Tensor;
+import com.yahoo.tensor.TensorType;
 import com.yahoo.tensor.evaluation.EvaluationContext;
 
 import java.util.Set;
@@ -25,6 +26,11 @@ public abstract class Context implements EvaluationContext<Reference> {
      */
     public abstract Value get(String name);
 
+    @Override
+    public TensorType getType(String reference) {
+        throw new UnsupportedOperationException("Not able to parse gereral references from string form");
+    }
+
     /** Returns a variable as a tensor */
     @Override
     public Tensor getTensor(String name) { return get(name).asTensor(); }
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/MapTypeContext.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/MapTypeContext.java
index 985878cfd66..2a42e2d92f7 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/MapTypeContext.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/MapTypeContext.java
@@ -23,6 +23,11 @@ public class MapTypeContext implements TypeContext<Reference> {
     }
 
     @Override
+    public TensorType getType(String reference) {
+        throw new UnsupportedOperationException("Not able to parse gereral references from string form");
+    }
+
+    @Override
     public TensorType getType(Reference reference) {
         return featureTypes.get(reference);
     }
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTForestNode.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTForestNode.java
index 8ee4cdbf297..649c70122f1 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTForestNode.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTForestNode.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.searchlib.rankingexpression.evaluation.gbdtoptimization;
 
+import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.DoubleValue;
 import com.yahoo.searchlib.rankingexpression.evaluation.Value;
@@ -26,7 +27,7 @@ public class GBDTForestNode extends ExpressionNode {
     }
 
     @Override
-    public final TensorType type(TypeContext context) { return TensorType.empty; }
+    public final TensorType type(TypeContext<Reference> context) { return TensorType.empty; }
 
     @Override
     public final Value evaluate(Context context) {
diff --git a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTNode.java b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTNode.java
index aac635b2545..53a286f09f6 100644
--- a/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTNode.java
+++ b/searchlib/src/main/java/com/yahoo/searchlib/rankingexpression/evaluation/gbdtoptimization/GBDTNode.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.searchlib.rankingexpression.evaluation.gbdtoptimization;
 
+import com.yahoo.searchlib.rankingexpression.Reference;
 import com.yahoo.searchlib.rankingexpression.evaluation.Context;
 import com.yahoo.searchlib.rankingexpression.evaluation.DoubleValue;
 import com.yahoo.searchlib.rankingexpression.evaluation.Value;
@@ -51,7 +52,7 @@ public final class GBDTNode extends ExpressionNode {
     public final double[] values() { return values; }
 
     @Override
-    public final TensorType type(TypeContext context) { return TensorType.empty; }
+    public final TensorType type(TypeContext<Reference> context) { return TensorType.empty; }
 
     @Override
     public final Value evaluate(Context context) {