From 71a50a5841542fc3ede159a35f1b9da2d7efe2c5 Mon Sep 17 00:00:00 2001 From: Anush008 Date: Thu, 14 Dec 2023 14:51:28 +0530 Subject: [PATCH] feat: discover ops --- .../java/io/qdrant/client/QdrantClient.java | 75 ++++++++++++++++++- .../client/ShardKeySelectorFactory.java | 6 ++ 2 files changed, 77 insertions(+), 4 deletions(-) diff --git a/src/main/java/io/qdrant/client/QdrantClient.java b/src/main/java/io/qdrant/client/QdrantClient.java index a32c7da..91a50cb 100644 --- a/src/main/java/io/qdrant/client/QdrantClient.java +++ b/src/main/java/io/qdrant/client/QdrantClient.java @@ -10,6 +10,7 @@ import io.qdrant.client.grpc.CollectionsGrpc; import io.qdrant.client.grpc.PointsGrpc; import io.qdrant.client.grpc.SnapshotsGrpc; + import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -34,6 +35,10 @@ import static io.qdrant.client.grpc.Collections.DeleteCollection; import static io.qdrant.client.grpc.Collections.DeleteShardKeyRequest; import static io.qdrant.client.grpc.Collections.DeleteShardKeyResponse; +import static io.qdrant.client.grpc.Points.DiscoverBatchPoints; +import static io.qdrant.client.grpc.Points.DiscoverBatchResponse; +import static io.qdrant.client.grpc.Points.DiscoverPoints; +import static io.qdrant.client.grpc.Points.DiscoverResponse; import static io.qdrant.client.grpc.Collections.GetCollectionInfoRequest; import static io.qdrant.client.grpc.Collections.GetCollectionInfoResponse; import static io.qdrant.client.grpc.Collections.ListAliasesRequest; @@ -675,7 +680,7 @@ public ListenableFuture> listAliasesAsync(@Nullable Durat * Creates a shard key for a collection. * * @param createShardKey The request object for the operation. - * @return a new instance of {@link CreateShardKeyResponse} + * @return a new instance of {@link ListenableFuture} */ public ListenableFuture createShardKeyAsync(CreateShardKeyRequest createShardKey) { return createShardKeyAsync(createShardKey, null); @@ -686,7 +691,7 @@ public ListenableFuture createShardKeyAsync(CreateShardK * * @param createShardKey The request object for the operation. * @param timeout The timeout for the call. - * @return a new instance of {@link CreateShardKeyResponse} + * @return a new instance of {@link ListenableFuture} */ public ListenableFuture createShardKeyAsync(CreateShardKeyRequest createShardKey, @Nullable Duration timeout) { String collectionName = createShardKey.getCollectionName(); @@ -708,7 +713,7 @@ public ListenableFuture createShardKeyAsync(CreateShardK * Deletes a shard key for a collection. * * @param deleteShardKey The request object for the operation. - * @return a new instance of {@link DeleteShardKeyResponse} + * @return a new instance of {@link ListenableFuture} */ public ListenableFuture deleteShardKeyAsync(DeleteShardKeyRequest deleteShardKey) { return deleteShardKeyAsync(deleteShardKey, null); @@ -719,7 +724,7 @@ public ListenableFuture deleteShardKeyAsync(DeleteShardK * * @param deleteShardKey The request object for the operation. * @param timeout The timeout for the call. - * @return a new instance of {@link DeleteShardKeyResponse} + * @return a new instance of {@link ListenableFuture} */ public ListenableFuture deleteShardKeyAsync(DeleteShardKeyRequest deleteShardKey, @Nullable Duration timeout) { String collectionName = deleteShardKey.getCollectionName(); @@ -2227,6 +2232,68 @@ public ListenableFuture> recommendGroupsAsync(RecommendPointGro MoreExecutors.directExecutor()); } + /** + * Use the context and a target to find the most similar points to the target. + * Constraints by the context. + * + * @param request The discover points request + * @return a new instance of {@link ListenableFuture} + */ + public ListenableFuture> discoverAsync(DiscoverPoints request) { + return discoverAsync(request, null); + } + + /** + * Use the context and a target to find the most similar points to the target. + * Constraints by the context. + * + * @param request The discover points request + * @param timeout The timeout for the call. + * @return a new instance of {@link ListenableFuture} + */ + public ListenableFuture> discoverAsync(DiscoverPoints request, @Nullable Duration timeout) { + String collectionName = request.getCollectionName(); + Preconditions.checkArgument(!collectionName.isEmpty(), "Collection name must not be empty"); + logger.debug("Discover on '{}'", collectionName); + ListenableFuture future = getPoints(timeout).discover(request); + addLogFailureCallback(future, "Discover"); + return Futures.transform( + future, + response -> response.getResultList(), + MoreExecutors.directExecutor()); + } + + /** + * Use the context and a target to find the most similar points to the target in a batch. + * Constraints by the context. + * + * @param request The discover batch points request. + * @return a new instance of {@link ListenableFuture} + */ + public ListenableFuture> discoverBatchAsync(DiscoverBatchPoints request) { + return discoverBatchAsync(request, null); + } + + /** + * Use the context and a target to find the most similar points to the target in a batch. + * Constraints by the context. + * + * @param request The discover batch points request. + * @param timeout The timeout for the call. + * @return a new instance of {@link ListenableFuture} + */ + public ListenableFuture> discoverBatchAsync(DiscoverBatchPoints request, @Nullable Duration timeout) { + String collectionName = request.getCollectionName(); + Preconditions.checkArgument(!collectionName.isEmpty(), "Collection name must not be empty"); + logger.debug("Discover batch on '{}'", collectionName); + ListenableFuture future = getPoints(timeout).discoverBatch(request); + addLogFailureCallback(future, "Discover batch"); + return Futures.transform( + future, + response -> response.getResultList(), + MoreExecutors.directExecutor()); + } + /** * Count the points in a collection. The count is exact * diff --git a/src/main/java/io/qdrant/client/ShardKeySelectorFactory.java b/src/main/java/io/qdrant/client/ShardKeySelectorFactory.java index e477896..a171f11 100644 --- a/src/main/java/io/qdrant/client/ShardKeySelectorFactory.java +++ b/src/main/java/io/qdrant/client/ShardKeySelectorFactory.java @@ -7,7 +7,13 @@ import java.util.Arrays; +/** + * Convenience methods for constructing {@link ShardKeySelector} + */ public class ShardKeySelectorFactory { + private ShardKeySelectorFactory() { + } + /** * Creates a {@link ShardKeySelector} with the given shard keys. *