Improvements to iserInput handling

- Handle both direct values and references in userInput
- Don't convert direct values to references
- Accept any QueryChain in rank

5 files changed, 34 insertions, 31 deletions

diff --git a/client/src/main/java/ai/vespa/client/dsl/FixedQuery.java b/client/src/main/java/ai/vespa/client/dsl/FixedQuery.java
index 9976a405182..d957217a9c7 100644
--- a/client/src/main/java/ai/vespa/client/dsl/FixedQuery.java
+++ b/client/src/main/java/ai/vespa/client/dsl/FixedQuery.java
@@ -9,7 +9,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 /**
- * FixedQuery contains a 'Query' which is terminated by a 'semicolon'
+ * FixedQuery contains a 'Query'.
  * This object holds vespa or user defined parameters
  * https://docs.vespa.ai/en/reference/query-api-reference.html
  */
@@ -342,7 +342,6 @@ public class FixedQuery {
         if (queryMap != null) {
             return queryMap;
         }
-
         assignIndex();
 
         StringBuilder sb = new StringBuilder();
@@ -366,7 +365,7 @@ public class FixedQuery {
     }
 
     /**
-     * build the vespa query string join by '&'
+     * Builds the vespa query string joined by '&'
      *
      * @return the query string
      */
@@ -399,9 +398,7 @@ public class FixedQuery {
     private Map<String, String> getUserInputs(Query q) {
         Map<String, String> param = new HashMap<>();
         q.queries.forEach(qu -> {
-            if (qu instanceof UserInput) {
-                param.putAll(((UserInput) qu).getParam());
-            } else if (qu instanceof Query) {
+            if (qu instanceof Query) {
                 param.putAll(getUserInputs((Query) qu));
             }
         });
diff --git a/client/src/main/java/ai/vespa/client/dsl/Q.java b/client/src/main/java/ai/vespa/client/dsl/Q.java
index 70e0e644c07..2bb998cd3e5 100644
--- a/client/src/main/java/ai/vespa/client/dsl/Q.java
+++ b/client/src/main/java/ai/vespa/client/dsl/Q.java
@@ -69,7 +69,7 @@ public final class Q {
      * @param ranks the ranks
      * @return the rank query
      */
-    public static Rank rank(Query query, Query... ranks) {
+    public static Rank rank(Query query, QueryChain... ranks) {
         return new Rank(query, ranks);
     }
 
diff --git a/client/src/main/java/ai/vespa/client/dsl/Rank.java b/client/src/main/java/ai/vespa/client/dsl/Rank.java
index 86fe80b2909..aad70a29502 100644
--- a/client/src/main/java/ai/vespa/client/dsl/Rank.java
+++ b/client/src/main/java/ai/vespa/client/dsl/Rank.java
@@ -9,9 +9,9 @@ import java.util.stream.Stream;
 
 public class Rank extends QueryChain {
 
-    private final List<Query> queries = new ArrayList<>();
+    private final List<QueryChain> queries = new ArrayList<>();
 
-    Rank(Query query, Query... ranks) {
+    Rank(Query query, QueryChain... ranks) {
         this.query = query;
         this.nonEmpty = query.nonEmpty();
         queries.add(query);
diff --git a/client/src/main/java/ai/vespa/client/dsl/UserInput.java b/client/src/main/java/ai/vespa/client/dsl/UserInput.java
index f42b3d82b52..de048c57574 100644
--- a/client/src/main/java/ai/vespa/client/dsl/UserInput.java
+++ b/client/src/main/java/ai/vespa/client/dsl/UserInput.java
@@ -1,16 +1,14 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.client.dsl;
 
-import java.util.Collections;
-import java.util.Map;
 import java.util.UUID;
 
 public class UserInput extends QueryChain {
 
     private final Annotation annotation; // accept only defaultIndex annotation
     private final String value;
+    private final boolean valueIsReference;
     private final String indexField;
-    private String placeholder; // for generating unique param
     private boolean setDefaultIndex;
 
     UserInput(Sources sources, String value) {
@@ -21,6 +19,7 @@ public class UserInput extends QueryChain {
         this.sources = sources;
         this.annotation = annotation;
         this.value = value;
+        this.valueIsReference = value.startsWith("@");
         this.nonEmpty = true;
 
         if (annotation.contains("defaultIndex")) {
@@ -39,23 +38,21 @@ public class UserInput extends QueryChain {
         this(null, annotation, value);
     }
 
-    public void setIndex(int index) {
-        placeholder = setDefaultIndex
-                      ? "_" + index + "_" + indexField
-                      : "_" + index;
-    }
-
     @Override
     public String toString() {
-        //([{"defaultIndex": "shpdescfree"}](userInput(@_shpdescfree_1)))
-        return setDefaultIndex
-               ? Text.format("([%s]userInput(@%s))", annotation, placeholder)
-               : Text.format("userInput(@%s)", placeholder);
-    }
-
-
-    Map<String, String> getParam() {
-        return Collections.singletonMap(placeholder, value);
+        StringBuilder b = new StringBuilder();
+        if (setDefaultIndex)
+            b.append("(").append(annotation);
+        b.append("userInput(");
+        if ( ! valueIsReference)
+            b.append("\"");
+        b.append(value);
+        if ( ! valueIsReference)
+            b.append("\"");
+        b.append(")");
+        if (setDefaultIndex)
+            b.append(")");
+        return b.toString();
     }
 
     @Override
diff --git a/client/src/test/java/ai/vespa/client/dsl/QTest.java b/client/src/test/java/ai/vespa/client/dsl/QTest.java
index a249bc91411..114cbee0177 100644
--- a/client/src/test/java/ai/vespa/client/dsl/QTest.java
+++ b/client/src/test/java/ai/vespa/client/dsl/QTest.java
@@ -165,15 +165,24 @@ class QTest {
     }
 
     @Test
-    void userInput_with_and_with_out_defaultIndex() {
+    void userInput_with_and_without_defaultIndex() {
         String q = Q.select("*")
                 .from("sd1")
-                .where(Q.ui("value"))
+                .where(Q.ui("value1"))
                 .and(Q.ui("index", "value2"))
                 .semicolon()
                 .build();
+        assertEquals("yql=select * from sd1 where userInput(\"value1\") and ({\"defaultIndex\":\"index\"}userInput(\"value2\"));", q);
+    }
 
-        assertEquals(q, "yql=select * from sd1 where userInput(@_1) and ([{\"defaultIndex\":\"index\"}]userInput(@_2_index));&_2_index=value2&_1=value");
+    @Test
+    void userInput_with_rank() {
+        String q = Q.select("url")
+                    .from("site")
+                    .where(Q.rank(Q.p("docQ").nearestNeighbor("vectorQuery"),
+                                  Q.ui("@query")))
+                    .semicolon().build();
+        assertEquals("yql=select url from site where rank(nearestNeighbor(docQ, vectorQuery), userInput(@query));", q);
     }
 
     @Test
@@ -185,7 +194,7 @@ class QTest {
                 .semicolon()
                 .build();
 
-        assertEquals(q, "yql=select * from sd1 where dotProduct(f1, {\"a\":1,\"b\":2,\"c\":3}) and f2 contains \"1\";");
+        assertEquals("yql=select * from sd1 where dotProduct(f1, {\"a\":1,\"b\":2,\"c\":3}) and f2 contains \"1\";", q);
     }
 
     @Test