Merge pull request #555 from weaviate/feat/query-profile

Author: bevzzz

.github/workflows/test.yaml

@@ -91,7 +91,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        WEAVIATE_VERSION: ["1.32.24", "1.33.11", "1.34.7", "1.35.2", "1.36.1"]
+        WEAVIATE_VERSION: ["1.32.24", "1.33.11", "1.34.7", "1.35.2", "1.36.9"]
     steps:
       - uses: actions/checkout@v4
 

src/it/java/io/weaviate/containers/Weaviate.java

@@ -44,7 +44,7 @@ public enum Version {
     V133(1, 33, 11),
     V134(1, 34, 7),
     V135(1, 35, 2),
-    V136(1, 36, 1);
+    V136(1, 36, 9);
 
     public final SemanticVersion semver;
 

src/it/java/io/weaviate/integration/SearchITest.java

@@ -36,7 +36,10 @@
 import io.weaviate.client6.v1.api.collections.query.Filter;
 import io.weaviate.client6.v1.api.collections.query.GroupBy;
 import io.weaviate.client6.v1.api.collections.query.Metadata;
+import io.weaviate.client6.v1.api.collections.query.QueryProfile;
+import io.weaviate.client6.v1.api.collections.query.QueryResponse;
 import io.weaviate.client6.v1.api.collections.query.QueryResponseGroup;
+import io.weaviate.client6.v1.api.collections.query.QueryResponseGrouped;
 import io.weaviate.client6.v1.api.collections.query.Rerank;
 import io.weaviate.client6.v1.api.collections.query.SortBy;
 import io.weaviate.client6.v1.api.collections.query.Target;
@@ -48,6 +51,7 @@
 import io.weaviate.containers.Img2VecNeural;
 import io.weaviate.containers.Model2Vec;
 import io.weaviate.containers.Weaviate;
+import io.weaviate.containers.Weaviate.Version;
 
 public class SearchITest extends ConcurrentTest {
   private static final ContainerGroup compose = Container.compose(
@@ -839,4 +843,45 @@ public void testVectorizerModel2VecPropeties() throws IOException {
           Assertions.assertThat(v.getSingle("title_vec")).isEqualTo(v.getSingle("author_vec"));
         });
   }
+
+  @Test
+  public void testQueryProfile() throws Exception {
+    Version.V136.orSkip();
+
+    var things = client.collections.use(COLLECTION);
+    var resp = things.query.nearVector(searchVector,
+        opt -> opt.distance(10f)
+            .returnMetadata(Metadata.QUERY_PROFILE));
+
+    Assertions.assertThat(resp)
+        .extracting(QueryResponse::queryProfile).isNotNull()
+        .extracting(QueryProfile::shards, InstanceOfAssertFactories.list(QueryProfile.ShardProfile.class))
+        .hasSize(1)
+        .allSatisfy(shardProfile -> Assertions.assertThat(shardProfile)
+            .extracting(QueryProfile.ShardProfile::searches,
+                InstanceOfAssertFactories.map(String.class, String.class))
+            .isNotEmpty());
+  }
+
+  @Test
+  @Ignore("v1.36.9 does not support query profiles for grouped queries")
+  public void testQueryProfile_groupBy() throws Exception {
+    Version.V136.orSkip();
+
+    var things = client.collections.use(COLLECTION);
+    var resp = things.query.nearVector(searchVector,
+        opt -> opt.distance(10f)
+            .returnMetadata(Metadata.QUERY_PROFILE),
+        GroupBy.property("category", 2, 5));
+
+    Assertions.assertThat(resp)
+        .extracting(QueryResponseGrouped::queryProfile).isNotNull()
+        .extracting(QueryProfile::shards,
+            InstanceOfAssertFactories.list(QueryProfile.ShardProfile.class))
+        .hasSize(1)
+        .allSatisfy(shardProfile -> Assertions.assertThat(shardProfile)
+            .extracting(QueryProfile.ShardProfile::searches,
+                InstanceOfAssertFactories.map(String.class, String.class))
+            .isNotEmpty());
+  }
 }

src/main/java/io/weaviate/client6/v1/api/collections/query/Metadata.java

@@ -11,7 +11,13 @@
 public interface Metadata {
   void appendTo(WeaviateProtoSearchGet.MetadataRequest.Builder metadata);
 
-  /** Include metadata in the metadata response. */
+  /**
+   * Include all metadata in the metadata response.
+   *
+   * <p>
+   * Collecting {@link #QUERY_PROFILE} involves significant overhead on the
+   * server side, so it is excluded from ALL and must be requested explicitly.
+   */
   public static final Metadata ALL = MetadataField.ALL;
   /** Include object creation time in the metadata response. */
   public static final Metadata CREATION_TIME_UNIX = MetadataField.CREATION_TIME_UNIX;
@@ -64,6 +70,13 @@ public interface Metadata {
    */
   public static final Metadata EXPLAIN_SCORE = MetadataField.EXPLAIN_SCORE;
 
+  /**
+   * Include a per-shard execution timing breakdowns for search queries.
+   *
+   * @see QueryProfile
+   */
+  public static final Metadata QUERY_PROFILE = MetadataField.QUERY_PROFILE;
+
   /**
    * MetadataField are collection properties that can be requested for any object.
    */
@@ -76,13 +89,15 @@ enum MetadataField implements Metadata {
     DISTANCE,
     CERTAINTY,
     SCORE,
-    EXPLAIN_SCORE;
+    EXPLAIN_SCORE,
+    QUERY_PROFILE;
 
     public void appendTo(WeaviateProtoSearchGet.MetadataRequest.Builder metadata) {
       switch (this) {
         case ALL:
           for (final var f : MetadataField.values()) {
-            if (f != ALL) {
+            // QUERY_PROFILE is expensive, require an explicit opt-in.
+            if (f != ALL && f != QUERY_PROFILE) {
               f.appendTo(metadata);
             }
           }
@@ -111,6 +126,8 @@ public void appendTo(WeaviateProtoSearchGet.MetadataRequest.Builder metadata) {
         case SCORE:
           metadata.setScore(true);
           break;
+        case QUERY_PROFILE:
+          metadata.setQueryProfile(true);
       }
     }
   }

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryProfile.java

@@ -0,0 +1,44 @@
+package io.weaviate.client6.v1.api.collections.query;
+
+import java.util.List;
+import java.util.Map;
+
+/** Per-shard execution timing breakdowns for search queries. */
+public record QueryProfile(List<ShardProfile> shards) {
+  /**
+   * Search profiling samples keyed by shard name.
+   * <p>
+   * Vector searches:
+   *
+   * <ul>
+   * <li>{@code vector_search_took}: time in the vector index
+   * <li>{@code hnsw_flat_search}: "true" / "false", whether flat search was used
+   * <li>{@code knn_search_layer_N_took}: per-layer HNSW traversal time
+   * <li>{@code knn_search_rescore_took}: PQ/BQ rescore time
+   * <li>{@code flat_search_iteration_took}, {@code flat_search_rescore_took}
+   * </ul>
+   *
+   * <p>
+   * Keyword (BM25) searches:
+   *
+   * <ul>
+   * <li>{@code kwd_method}, {@code kwd_time}, {@code kwd_filter_size}
+   * <li>{@code kwd_1_tok_time}: tokenization time
+   * <li>{@code kwd_3_term_time}: term lookup time
+   * <li>{@code kwd_4_wand_time / kwd_4_bmw_time}: scoring algorithm time
+   * <li>{@code kwd_5_objects_time}: object fetch time
+   * <li>{@code kwd_6_res_count}: result count
+   * </ul>
+   *
+   * <p>
+   * Filtered searches (any type):
+   *
+   * <ul>
+   * <li>{@code filters_build_allow_list_took}: filter evaluation time
+   * <li>{@code filters_ids_matched}: number of IDs matching the filter
+   * <li>{@code sort_took}, objects_took
+   * </ul>
+   */
+  public static record ShardProfile(Map<String, Map<String, String>> searches) {
+  }
+}

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryRequest.java

@@ -1,6 +1,7 @@
 package io.weaviate.client6.v1.api.collections.query;
 
 import io.weaviate.client6.v1.api.collections.CollectionHandleDefaults;
+import io.weaviate.client6.v1.internal.Debug;
 import io.weaviate.client6.v1.internal.grpc.Rpc;
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateGrpc.WeaviateBlockingStub;
 import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateGrpc.WeaviateFutureStub;

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponse.java

@@ -3,7 +3,9 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import io.weaviate.client6.v1.api.collections.GeoCoordinates;
@@ -19,7 +21,8 @@
 import io.weaviate.client6.v1.internal.orm.PropertiesBuilder;
 
 public record QueryResponse<PropertiesT>(
-    List<WeaviateObject<PropertiesT>> objects) {
+    List<WeaviateObject<PropertiesT>> objects,
+    QueryProfile queryProfile) {
 
   static <PropertiesT> QueryResponse<PropertiesT> unmarshal(WeaviateProtoSearchGet.SearchReply reply,
       CollectionDescriptor<PropertiesT> collection) {
@@ -29,7 +32,22 @@ static <PropertiesT> QueryResponse<PropertiesT> unmarshal(WeaviateProtoSearchGet
         .map(obj -> QueryResponse.unmarshalResultObject(
             obj.getProperties(), obj.getMetadata(), collection))
         .toList();
-    return new QueryResponse<>(objects);
+    return new QueryResponse<>(objects, unmarshalQueryProfile(reply.getQueryProfile()));
+  }
+
+  static QueryProfile unmarshalQueryProfile(WeaviateProtoSearchGet.QueryProfile queryProfile) {
+    return new QueryProfile(queryProfile
+        .getShardsList()
+        .stream()
+        .map(shard -> new QueryProfile.ShardProfile(
+            shard
+                .getSearchesMap()
+                .entrySet()
+                .stream()
+                .collect(Collectors.toMap(
+                    Map.Entry::getKey,
+                    entry -> entry.getValue().getDetailsMap()))))
+        .toList());
   }
 
   public static <PropertiesT> WeaviateObject<PropertiesT> unmarshalResultObject(

src/main/java/io/weaviate/client6/v1/api/collections/query/QueryResponseGrouped.java

@@ -14,7 +14,8 @@ public record QueryResponseGrouped<PropertiesT>(
     /** All objects retrieved in the query. */
     List<QueryObjectGrouped<PropertiesT>> objects,
     /** Grouped response objects. */
-    Map<String, QueryResponseGroup<PropertiesT>> groups) {
+    Map<String, QueryResponseGroup<PropertiesT>> groups,
+    QueryProfile queryProfile) {
 
   static <PropertiesT> QueryResponseGrouped<PropertiesT> unmarshal(
       WeaviateProtoSearchGet.SearchReply reply,
@@ -47,6 +48,8 @@ static <PropertiesT> QueryResponseGrouped<PropertiesT> unmarshal(
         // about the server's response.
         .collect(Collectors.toMap(QueryResponseGroup::name, Function.identity()));
 
-    return new QueryResponseGrouped<PropertiesT>(allObjects, groups);
+    return new QueryResponseGrouped<PropertiesT>(
+        allObjects, groups,
+        QueryResponse.unmarshalQueryProfile(reply.getQueryProfile()));
   }
 }