distance matrix API

qdrant · Oct 4, 2024 · 403e07f · 403e07f
1 parent e70ea39
commit 403e07f
Show file tree

Hide file tree

Showing 3 changed files with 104 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -11,6 +11,7 @@ build/
 .idea/libraries/
 .idea/uiDesigner.xml
 .idea/codeStyles/codeStyleConfig.xml
+.idea/codeStyles/Project.xml
 *.iws
 *.iml
 *.ipr

diff --git a/src/main/java/io/qdrant/client/QdrantClient.java b/src/main/java/io/qdrant/client/QdrantClient.java
@@ -40,6 +40,7 @@
 import io.qdrant.client.grpc.Collections.VectorsConfig;
 import io.qdrant.client.grpc.CollectionsGrpc;
 import io.qdrant.client.grpc.JsonWithInt.Value;
+import io.qdrant.client.grpc.Points;
 import io.qdrant.client.grpc.Points.BatchResult;
 import io.qdrant.client.grpc.Points.ClearPayloadPoints;
 import io.qdrant.client.grpc.Points.CountPoints;
@@ -2804,6 +2805,72 @@ public ListenableFuture<List<PointGroup>> queryGroupsAsync(
         future, response -> response.getResult().getGroupsList(), MoreExecutors.directExecutor());
   }
 
+  // region distance matrix
+
+  /**
+   * Compute distance matrix for sampled points with a pair based output format.
+   *
+   * @param request the search matrix pairs request
+   * @return a new instance of {@link ListenableFuture}
+   */
+  public ListenableFuture<Points.SearchMatrixPairs> searchMatrixPairsAsync(
+      Points.SearchMatrixPoints request) {
+    return searchMatrixPairsAsync(request, null);
+  }
+
+  /**
+   * Compute distance matrix for sampled points with a pair based output format.
+   *
+   * @param request the search matrix pairs request
+   * @param timeout the timeout for the call.
+   * @return a new instance of {@link ListenableFuture}
+   */
+  public ListenableFuture<Points.SearchMatrixPairs> searchMatrixPairsAsync(
+      Points.SearchMatrixPoints request, @Nullable Duration timeout) {
+    Preconditions.checkArgument(
+        !request.getCollectionName().isEmpty(), "Collection name must not be empty");
+
+    logger.debug("Search matrix pairs on '{}'", request.getCollectionName());
+    ListenableFuture<Points.SearchMatrixPairsResponse> future =
+        getPoints(timeout).searchMatrixPairs(request);
+    addLogFailureCallback(future, "Search matrix pairs");
+    return Futures.transform(
+        future, Points.SearchMatrixPairsResponse::getResult, MoreExecutors.directExecutor());
+  }
+
+  /**
+   * Compute distance matrix for sampled points with an offset based output format
+   *
+   * @param request the search matrix pairs request
+   * @return a new instance of {@link ListenableFuture}
+   */
+  public ListenableFuture<Points.SearchMatrixOffsets> searchMatrixOffsetsAsync(
+      Points.SearchMatrixPoints request) {
+    return searchMatrixOffsetsAsync(request, null);
+  }
+
+  /**
+   * Compute distance matrix for sampled points with an offset based output format
+   *
+   * @param request the search matrix pairs request
+   * @param timeout the timeout for the call.
+   * @return a new instance of {@link ListenableFuture}
+   */
+  public ListenableFuture<Points.SearchMatrixOffsets> searchMatrixOffsetsAsync(
+      Points.SearchMatrixPoints request, @Nullable Duration timeout) {
+    Preconditions.checkArgument(
+        !request.getCollectionName().isEmpty(), "Collection name must not be empty");
+
+    logger.debug("Search matrix offsets on '{}'", request.getCollectionName());
+    ListenableFuture<Points.SearchMatrixOffsetsResponse> future =
+        getPoints(timeout).searchMatrixOffsets(request);
+    addLogFailureCallback(future, "Search matrix offsets");
+    return Futures.transform(
+        future, Points.SearchMatrixOffsetsResponse::getResult, MoreExecutors.directExecutor());
+  }
+
+  // endregion
+
   // region Snapshot Management
 
   /**

diff --git a/src/test/java/io/qdrant/client/PointsTest.java b/src/test/java/io/qdrant/client/PointsTest.java
@@ -813,6 +813,42 @@ public void queryGroups() throws ExecutionException, InterruptedException {
     assertEquals(1, groups.stream().filter(g -> g.getHitsCount() == 1).count());
   }
 
+  @Test
+  public void searchMatrixOffsets() throws ExecutionException, InterruptedException {
+    createAndSeedCollection(testName);
+
+    Points.SearchMatrixOffsets offsets =
+        client
+            .searchMatrixOffsetsAsync(
+                Points.SearchMatrixPoints.newBuilder()
+                    .setCollectionName(testName)
+                    .setSample(3)
+                    .setLimit(2)
+                    .build())
+            .get();
+
+    // Number of ids matches the limit
+    assertEquals(2, offsets.getIdsCount());
+  }
+
+  @Test
+  public void searchMatrixPairs() throws ExecutionException, InterruptedException {
+    createAndSeedCollection(testName);
+
+    Points.SearchMatrixPairs pairs =
+        client
+            .searchMatrixPairsAsync(
+                Points.SearchMatrixPoints.newBuilder()
+                    .setCollectionName(testName)
+                    .setSample(3)
+                    .setLimit(2)
+                    .build())
+            .get();
+
+    // Number of ids matches the limit
+    assertEquals(2, pairs.getPairsCount());
+  }
+
   private void createAndSeedCollection(String collectionName)
       throws ExecutionException, InterruptedException {
     CreateCollection request =