Merge pull request #25663 from vespa-engine/geirst/query-trace-profiling-parameters

Add query trace parameters for profiling backend query evaluation.

9 files changed, 303 insertions, 12 deletions

diff --git a/container-search/abi-spec.json b/container-search/abi-spec.json
index a71082ecba3..ae4d0e713f8 100644
--- a/container-search/abi-spec.json
+++ b/container-search/abi-spec.json
@@ -5698,6 +5698,7 @@
       "public void setTimestamps(boolean)",
       "public boolean getQuery()",
       "public void setQuery(boolean)",
+      "public com.yahoo.search.query.profiling.Profiling getProfiling()",
       "public void trace(java.lang.String, int)",
       "public void trace(java.lang.Object, int)",
       "public void trace(java.lang.String, boolean, int)",
@@ -5716,7 +5717,8 @@
       "public static final java.lang.String EXPLAIN_LEVEL",
       "public static final java.lang.String PROFILE_DEPTH",
       "public static final java.lang.String TIMESTAMPS",
-      "public static final java.lang.String QUERY"
+      "public static final java.lang.String QUERY",
+      "public static final java.lang.String PROFILING"
     ]
   },
   "com.yahoo.search.query.UniqueRequestId" : {
@@ -6679,6 +6681,52 @@
     ],
     "fields" : [ ]
   },
+  "com.yahoo.search.query.profiling.Profiling" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [
+      "java.lang.Cloneable"
+    ],
+    "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public void <init>()",
+      "public static com.yahoo.search.query.profile.types.QueryProfileType getArgumentType()",
+      "public com.yahoo.search.query.profiling.ProfilingParams getMatching()",
+      "public com.yahoo.search.query.profiling.ProfilingParams getFirstPhaseRanking()",
+      "public com.yahoo.search.query.profiling.ProfilingParams getSecondPhaseRanking()",
+      "public com.yahoo.search.query.profiling.ProfilingParams clone()",
+      "public boolean equals(java.lang.Object)",
+      "public int hashCode()",
+      "public bridge synthetic java.lang.Object clone()"
+    ],
+    "fields" : [
+      "public static final java.lang.String MATCHING",
+      "public static final java.lang.String FIRST_PHASE_RANKING",
+      "public static final java.lang.String SECOND_PHASE_RANKING"
+    ]
+  },
+  "com.yahoo.search.query.profiling.ProfilingParams" : {
+    "superClass" : "java.lang.Object",
+    "interfaces" : [ ],
+    "attributes" : [
+      "public"
+    ],
+    "methods" : [
+      "public void <init>()",
+      "public static com.yahoo.search.query.profile.types.QueryProfileType getArgumentType()",
+      "public int getDepth()",
+      "public void setDepth(int)",
+      "public com.yahoo.search.query.profiling.ProfilingParams clone()",
+      "public boolean equals(java.lang.Object)",
+      "public int hashCode()",
+      "public bridge synthetic java.lang.Object clone()"
+    ],
+    "fields" : [
+      "public static final java.lang.String PROFILING_PARAMS",
+      "public static final java.lang.String DEPTH"
+    ]
+  },
   "com.yahoo.search.query.properties.CloneHelper" : {
     "superClass" : "com.yahoo.processing.request.CloneHelper",
     "interfaces" : [ ],
diff --git a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
index 90e16e9dd44..0b6a6a97a35 100644
--- a/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
+++ b/container-search/src/main/java/com/yahoo/search/dispatch/rpc/ProtobufSerialization.java
@@ -24,6 +24,7 @@ import com.yahoo.search.query.QueryTree;
 import com.yahoo.search.query.Ranking;
 import com.yahoo.search.query.Sorting;
 import com.yahoo.search.query.Sorting.Order;
+import com.yahoo.search.query.profiling.Profiling;
 import com.yahoo.search.result.Coverage;
 import com.yahoo.search.result.ErrorMessage;
 import com.yahoo.searchlib.aggregation.Grouping;
@@ -90,8 +91,12 @@ public class ProtobufSerialization {
             builder.setCacheGrouping(true);
         }
 
-        builder.setTraceLevel(getTraceLevelForBackend(query));
+        int traceLevel = getTraceLevelForBackend(query);
+        builder.setTraceLevel(traceLevel);
         builder.setProfileDepth(query.getTrace().getProfileDepth());
+        if (traceLevel > 0) {
+            mergeToSearchRequestFromProfiling(query.getTrace().getProfiling(), builder);
+        }
 
         mergeToSearchRequestFromRanking(query.getRanking(), scratchPad, builder);
 
@@ -137,6 +142,22 @@ public class ProtobufSerialization {
         }
     }
 
+    private static void mergeToSearchRequestFromProfiling(Profiling prof, SearchProtocol.SearchRequest.Builder builder) {
+        var profBuilder = SearchProtocol.Profiling.newBuilder();
+        if (prof.getMatching().getDepth() != 0) {
+            profBuilder.setMatch(SearchProtocol.ProfilingParams.newBuilder().setDepth(prof.getMatching().getDepth()));
+        }
+        if (prof.getFirstPhaseRanking().getDepth() != 0) {
+            profBuilder.setFirstPhase(SearchProtocol.ProfilingParams.newBuilder().setDepth(prof.getFirstPhaseRanking().getDepth()));
+        }
+        if (prof.getSecondPhaseRanking().getDepth() != 0) {
+            profBuilder.setSecondPhase(SearchProtocol.ProfilingParams.newBuilder().setDepth(prof.getSecondPhaseRanking().getDepth()));
+        }
+        if (profBuilder.hasMatch() || profBuilder.hasFirstPhase() || profBuilder.hasSecondPhase()) {
+            builder.setProfiling(profBuilder);
+        }
+    }
+
     static SearchProtocol.DocsumRequest.Builder createDocsumRequestBuilder(Query query,
                                                                            String serverId,
                                                                            String summaryClass,
diff --git a/container-search/src/main/java/com/yahoo/search/query/Trace.java b/container-search/src/main/java/com/yahoo/search/query/Trace.java
index b8eacad67fc..62258390147 100644
--- a/container-search/src/main/java/com/yahoo/search/query/Trace.java
+++ b/container-search/src/main/java/com/yahoo/search/query/Trace.java
@@ -8,7 +8,9 @@ import com.yahoo.search.Query;
 import com.yahoo.search.query.profile.QueryProfileProperties;
 import com.yahoo.search.query.profile.compiled.CompiledQueryProfile;
 import com.yahoo.search.query.profile.types.FieldDescription;
+import com.yahoo.search.query.profile.types.QueryProfileFieldType;
 import com.yahoo.search.query.profile.types.QueryProfileType;
+import com.yahoo.search.query.profiling.Profiling;
 
 import java.util.HashSet;
 import java.util.Map;
@@ -37,6 +39,7 @@ public class Trace implements Cloneable {
     public static final String PROFILE_DEPTH = "profileDepth";
     public static final String TIMESTAMPS = "timestamps";
     public static final String QUERY = "query";
+    public static final String PROFILING = "profiling";
 
     static {
         argumentType = new QueryProfileType(TRACE);
@@ -47,6 +50,7 @@ public class Trace implements Cloneable {
         argumentType.addField(new FieldDescription(PROFILE_DEPTH, "integer"));
         argumentType.addField(new FieldDescription(TIMESTAMPS, "boolean"));
         argumentType.addField(new FieldDescription(QUERY, "boolean"));
+        argumentType.addField(new FieldDescription(PROFILING, new QueryProfileFieldType(Profiling.getArgumentType())));
         argumentType.freeze();
     }
 
@@ -59,6 +63,7 @@ public class Trace implements Cloneable {
     private int profileDepth = 0;
     private boolean timestamps = false;
     private boolean query = true;
+    private Profiling profiling = new Profiling();
 
     public Trace(Query parent) {
         this.parent = Objects.requireNonNull(parent);
@@ -74,7 +79,12 @@ public class Trace implements Cloneable {
     public int getExplainLevel() { return explainLevel; }
 
     /** Sets the profiling depth. Profiling enabled if non-zero. Higher numbers means increasingly more detail. */
-    public void setProfileDepth(int profileDepth) { this.profileDepth = profileDepth; }
+    public void setProfileDepth(int profileDepth) {
+        this.profileDepth = profileDepth;
+        profiling.getMatching().setDepth(profileDepth);
+        profiling.getFirstPhaseRanking().setDepth(profileDepth);
+        profiling.getSecondPhaseRanking().setDepth(profileDepth);
+    }
     public int getProfileDepth() { return profileDepth; }
 
     /** Returns whether trace entries should have a timestamp. Default is false. */
@@ -85,6 +95,10 @@ public class Trace implements Cloneable {
     public boolean getQuery() { return query; }
     public void setQuery(boolean query) { this.query = query; }
 
+    public Profiling getProfiling() {
+        return profiling;
+    }
+
     /**
      * Adds a context message to this query and to the info log,
      * if the context level of the query is sufficiently high.
@@ -218,17 +232,21 @@ public class Trace implements Cloneable {
 
     @Override
     public boolean equals(Object o) {
-        if (o == this ) return true;
-        if ( ! (o instanceof Trace other)) return false;
-        if (other.level != this.level) return false;
-        if (other.explainLevel != this.explainLevel) return false;
-        if (other.timestamps != this.timestamps) return false;
-        if (other.query != this.query) return false;
-        return true;
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        Trace trace = (Trace) o;
+        return level == trace.level &&
+                explainLevel == trace.explainLevel &&
+                profileDepth == trace.profileDepth &&
+                timestamps == trace.timestamps &&
+                query == trace.query &&
+                Objects.equals(profiling, trace.profiling);
     }
 
     @Override
-    public int hashCode() { return Objects.hash(level, explainLevel, timestamps, query); }
+    public int hashCode() {
+        return Objects.hash(level, explainLevel, profileDepth, timestamps, query, profiling);
+    }
 
     @Override
     public Trace clone() {
diff --git a/container-search/src/main/java/com/yahoo/search/query/profiling/Profiling.java b/container-search/src/main/java/com/yahoo/search/query/profiling/Profiling.java
new file mode 100644
index 00000000000..8541b9c3930
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/query/profiling/Profiling.java
@@ -0,0 +1,75 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.query.profiling;
+
+import com.yahoo.search.query.Trace;
+import com.yahoo.search.query.profile.types.FieldDescription;
+import com.yahoo.search.query.profile.types.QueryProfileFieldType;
+import com.yahoo.search.query.profile.types.QueryProfileType;
+
+import java.util.Objects;
+
+/**
+ * Contains settings for different parts of the backend query evaluation that should be profiled.
+ *
+ * @author geirst
+ */
+public class Profiling implements Cloneable {
+
+    public static final String MATCHING = "matching";
+    public static final String FIRST_PHASE_RANKING = "firstPhaseRanking";
+    public static final String SECOND_PHASE_RANKING = "secondPhaseRanking";
+
+    private static final QueryProfileType argumentType;
+
+    static {
+        argumentType = new QueryProfileType(Trace.PROFILING);
+        argumentType.setStrict(true);
+        argumentType.setBuiltin(true);
+        argumentType.addField(new FieldDescription(MATCHING, new QueryProfileFieldType(ProfilingParams.getArgumentType())));
+        argumentType.addField(new FieldDescription(FIRST_PHASE_RANKING, new QueryProfileFieldType(ProfilingParams.getArgumentType())));
+        argumentType.addField(new FieldDescription(SECOND_PHASE_RANKING, new QueryProfileFieldType(ProfilingParams.getArgumentType())));
+        argumentType.freeze();
+    }
+
+    public static QueryProfileType getArgumentType() { return argumentType; }
+
+    private ProfilingParams matching = new ProfilingParams();
+    private ProfilingParams firstPhaseRanking = new ProfilingParams();
+    private ProfilingParams secondPhaseRanking = new ProfilingParams();
+
+    public ProfilingParams getMatching() {
+        return matching;
+    }
+
+    public ProfilingParams getFirstPhaseRanking() {
+        return firstPhaseRanking;
+    }
+
+    public ProfilingParams getSecondPhaseRanking() {
+        return secondPhaseRanking;
+    }
+
+    @Override
+    public ProfilingParams clone() {
+        try {
+            return (ProfilingParams) super.clone();
+        } catch (CloneNotSupportedException e) {
+            throw new RuntimeException("Someone inserted a non-cloneable superclass", e);
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        Profiling profiling = (Profiling) o;
+        return Objects.equals(matching, profiling.matching) &&
+                Objects.equals(firstPhaseRanking, profiling.firstPhaseRanking) &&
+                Objects.equals(secondPhaseRanking, profiling.secondPhaseRanking);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(matching, firstPhaseRanking, secondPhaseRanking);
+    }
+}
diff --git a/container-search/src/main/java/com/yahoo/search/query/profiling/ProfilingParams.java b/container-search/src/main/java/com/yahoo/search/query/profiling/ProfilingParams.java
new file mode 100644
index 00000000000..87b4f6285a3
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/query/profiling/ProfilingParams.java
@@ -0,0 +1,62 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.search.query.profiling;
+
+import com.yahoo.search.query.profile.types.FieldDescription;
+import com.yahoo.search.query.profile.types.QueryProfileType;
+
+import java.util.Objects;
+
+/**
+ * Contains parameters for a part of the backend query evaluation that should be profiled.
+ *
+ * @author geirst
+ */
+public class ProfilingParams {
+
+    public static final String PROFILING_PARAMS = "profilingParams";
+    public static final String DEPTH = "depth";
+
+    private static final QueryProfileType argumentType;
+
+    static {
+        argumentType = new QueryProfileType(PROFILING_PARAMS);
+        argumentType.setStrict(true);
+        argumentType.setBuiltin(true);
+        argumentType.addField(new FieldDescription(DEPTH, "integer"));
+        argumentType.freeze();
+    }
+
+    public static QueryProfileType getArgumentType() { return argumentType; }
+
+    private int depth = 0;
+
+    public int getDepth() {
+        return depth;
+    }
+
+    public void setDepth(int value) {
+        depth = value;
+    }
+
+    @Override
+    public ProfilingParams clone() {
+        try {
+            return (ProfilingParams) super.clone();
+        } catch (CloneNotSupportedException e) {
+            throw new RuntimeException("Someone inserted a non-cloneable superclass", e);
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        ProfilingParams that = (ProfilingParams) o;
+        return Objects.equals(depth, that.depth);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(depth);
+    }
+}
diff --git a/container-search/src/main/java/com/yahoo/search/query/profiling/package-info.java b/container-search/src/main/java/com/yahoo/search/query/profiling/package-info.java
new file mode 100644
index 00000000000..827488684cc
--- /dev/null
+++ b/container-search/src/main/java/com/yahoo/search/query/profiling/package-info.java
@@ -0,0 +1,7 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+@ExportPackage
+@PublicApi
+package com.yahoo.search.query.profiling;
+
+import com.yahoo.api.annotations.PublicApi;
+import com.yahoo.osgi.annotation.ExportPackage;
diff --git a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
index e4a83972fae..da0051c527c 100644
--- a/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
+++ b/container-search/src/main/java/com/yahoo/search/query/properties/QueryProperties.java
@@ -16,6 +16,8 @@ import com.yahoo.search.query.profile.compiled.CompiledQueryProfileRegistry;
 import com.yahoo.search.query.profile.types.ConversionContext;
 import com.yahoo.search.query.profile.types.FieldDescription;
 import com.yahoo.search.query.profile.types.QueryProfileType;
+import com.yahoo.search.query.profiling.Profiling;
+import com.yahoo.search.query.profiling.ProfilingParams;
 import com.yahoo.search.query.ranking.Diversity;
 import com.yahoo.search.query.ranking.MatchPhase;
 import com.yahoo.search.query.ranking.Matching;
@@ -322,6 +324,15 @@ public class QueryProperties extends Properties {
                 if (key.last().equals(Trace.QUERY))
                     query.getTrace().setQuery(asBoolean(value, true));
             }
+            else if ((key.size() == 4) &&
+                    key.get(0).equals(Trace.TRACE) &&
+                    key.get(1).equals(Trace.PROFILING) &&
+                    key.get(3).equals(ProfilingParams.DEPTH)) {
+                var params = getProfilingParams(query.getTrace().getProfiling(), key.get(2));
+                if (params != null) {
+                    params.setDepth(asInteger(value, 0));
+                }
+            }
             else if (key.first().equals(Select.SELECT)) {
                 if (key.size() == 1) {
                     query.getSelect().setGroupingExpressionString(asString(value, ""));
@@ -364,6 +375,17 @@ public class QueryProperties extends Properties {
         }
     }
 
+    private static ProfilingParams getProfilingParams(Profiling prof, String name) {
+        if (name.equals(Profiling.MATCHING)) {
+            return prof.getMatching();
+        } else if (name.equals(Profiling.FIRST_PHASE_RANKING)) {
+            return prof.getFirstPhaseRanking();
+        } else if (name.equals(Profiling.SECOND_PHASE_RANKING)) {
+            return prof.getSecondPhaseRanking();
+        }
+        return null;
+    }
+
     @Override
     public Map<String, Object> listProperties(CompoundName prefix,
                                               Map<String,String> context,
diff --git a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/ProtobufSerializationTest.java b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/ProtobufSerializationTest.java
index 6f7a1f3bd8b..80f7375b6eb 100644
--- a/container-search/src/test/java/com/yahoo/search/dispatch/rpc/ProtobufSerializationTest.java
+++ b/container-search/src/test/java/com/yahoo/search/dispatch/rpc/ProtobufSerializationTest.java
@@ -44,6 +44,7 @@ public class ProtobufSerializationTest {
                 contentsOf(request1.getTensorFeatureOverrides(0).getValue()));
         assertEquals("\"\\006\\001\\001\\001x\\001?\\231\\231\\232\"",
                 contentsOf(request1.getTensorFeatureOverrides(1).getValue()));
+        assertFalse(request1.hasProfiling());
 
         query.prepare(); // calling prepare() moves "overrides" to "features" - content stays the same
         SearchProtocol.SearchRequest request2 = ProtobufSerialization.convertFromQuery(query, 9, "serverId", 0.5);
@@ -136,4 +137,26 @@ public class ProtobufSerializationTest {
         }
     }
 
+    @Test
+    void profiling_parameters_are_serialized_in_search_request() {
+        var q = new Query("?query=test&trace.level=1&" +
+                "trace.profiling.matching.depth=3&" +
+                "trace.profiling.firstPhaseRanking.depth=5&" +
+                "trace.profiling.secondPhaseRanking.depth=-7");
+        var req = ProtobufSerialization.convertFromQuery(q, 1, "serverId", 0.5);
+        assertEquals(3, req.getProfiling().getMatch().getDepth());
+        assertEquals(5, req.getProfiling().getFirstPhase().getDepth());
+        assertEquals(-7, req.getProfiling().getSecondPhase().getDepth());
+    }
+
+    @Test
+    void only_set_profiling_parameters_are_serialized_in_search_request() {
+        var q = new Query("?query=test&trace.level=1&" +
+                "trace.profiling.matching.depth=3");
+        var req = ProtobufSerialization.convertFromQuery(q, 1, "serverId", 0.5);
+        assertEquals(3, req.getProfiling().getMatch().getDepth());
+        assertFalse(req.getProfiling().hasFirstPhase());
+        assertFalse(req.getProfiling().hasSecondPhase());
+    }
+
 }
diff --git a/container-search/src/test/java/com/yahoo/search/test/QueryTestCase.java b/container-search/src/test/java/com/yahoo/search/test/QueryTestCase.java
index 9a75d994919..559b814a265 100644
--- a/container-search/src/test/java/com/yahoo/search/test/QueryTestCase.java
+++ b/container-search/src/test/java/com/yahoo/search/test/QueryTestCase.java
@@ -597,9 +597,24 @@ public class QueryTestCase {
     }
 
     @Test
-    void testProfilingDepth() {
+    void profile_depth_sets_default_profiling_parameters() {
         Query q = new Query("?query=foo&trace.profileDepth=2");
         assertEquals(2, q.getTrace().getProfileDepth());
+        assertEquals(2, q.getTrace().getProfiling().getMatching().getDepth());
+        assertEquals(2, q.getTrace().getProfiling().getFirstPhaseRanking().getDepth());
+        assertEquals(2, q.getTrace().getProfiling().getSecondPhaseRanking().getDepth());
+    }
+
+    @Test
+    void profiling_parameters_are_resolved() {
+        var q = new Query("?query=foo&" +
+                "trace.profiling.matching.depth=3&" +
+                "trace.profiling.firstPhaseRanking.depth=5&" +
+                "trace.profiling.secondPhaseRanking.depth=-7");
+        assertEquals(0, q.getTrace().getProfileDepth());
+        assertEquals(3, q.getTrace().getProfiling().getMatching().getDepth());
+        assertEquals(5, q.getTrace().getProfiling().getFirstPhaseRanking().getDepth());
+        assertEquals(-7, q.getTrace().getProfiling().getSecondPhaseRanking().getDepth());
     }
 
     @Test