handle more complex features and functions wrapped in rankingExpression(foo)

1 files changed, 8 insertions, 3 deletions

diff --git a/config-model/src/main/java/com/yahoo/schema/expressiontransforms/InputRecorder.java b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/InputRecorder.java
index b0f63ebb732..515745cb6bc 100644
--- a/config-model/src/main/java/com/yahoo/schema/expressiontransforms/InputRecorder.java
+++ b/config-model/src/main/java/com/yahoo/schema/expressiontransforms/InputRecorder.java
@@ -46,7 +46,12 @@ public class InputRecorder extends ExpressionTransformer<RankProfileTransformCon
         Reference ref = feature.reference();
         String name = ref.name();
         var args = ref.arguments();
-        if (args.size() == 0) {
+        boolean simpleFunctionOrIdentifier = (args.size() == 0) && (ref.output() == null);
+        if (ref.isSimpleRankingExpressionWrapper()) {
+            name = ref.simpleArgument().get();
+            simpleFunctionOrIdentifier = true;
+        }
+        if (simpleFunctionOrIdentifier) {
             var f = context.rankProfile().getFunctions().get(name);
             if (f != null && f.function().arguments().size() == 0) {
                 transform(f.function().getBody().getRoot(), context);
@@ -55,7 +60,7 @@ public class InputRecorder extends ExpressionTransformer<RankProfileTransformCon
             neededInputs.add(feature.toString());
             return;
         }
-        if (args.size() == 1) {
+        if (FeatureNames.isSimpleFeature(ref)) {
             if (FeatureNames.isAttributeFeature(ref)) {
                 neededInputs.add(feature.toString());
                 return;
@@ -96,6 +101,6 @@ public class InputRecorder extends ExpressionTransformer<RankProfileTransformCon
             }
             return;
         }
-        throw new IllegalArgumentException("cannot handle feature: " + feature);
+        neededInputs.add(feature.toString());
     }
 }