Browse Source

Simplify `listIDInSegment`

jianghua 4 years ago
parent
commit
52819633db

+ 2 - 4
src/main/java/io/milvus/client/MilvusClient.java

@@ -362,11 +362,9 @@ public interface MilvusClient {
    *
    * @param collectionName collection to get entity ids from
    * @param segmentId segment id in the collection
-   * @return <code>ListIDInSegmentResponse</code>
-   * @see ListIDInSegmentResponse
-   * @see Response
+   * @return a list of entity ids in the segment
    */
-  ListIDInSegmentResponse listIDInSegment(String collectionName, Long segmentId);
+  List<Long> listIDInSegment(String collectionName, Long segmentId);
 
   /**
    * Deletes data in a collection by a list of ids

+ 10 - 42
src/main/java/io/milvus/client/MilvusGrpcClient.java

@@ -483,48 +483,16 @@ abstract class AbstractMilvusGrpcClient implements MilvusClient {
   }
 
   @Override
-  public ListIDInSegmentResponse listIDInSegment(String collectionName, Long segmentId) {
-    if (!maybeAvailable()) {
-      logWarning("You are not connected to Milvus server");
-      return new ListIDInSegmentResponse(
-          new Response(Response.Status.CLIENT_NOT_CONNECTED), Collections.emptyList());
-    }
-
-    GetEntityIDsParam request =
-        GetEntityIDsParam.newBuilder()
-            .setCollectionName(collectionName)
-            .setSegmentId(segmentId)
-            .build();
-    EntityIds response;
-
-    try {
-      response = blockingStub().getEntityIDs(request);
-
-      if (response.getStatus().getErrorCode() == ErrorCode.SUCCESS) {
-
-        logInfo(
-            "listIDInSegment in collection `{}`, segment `{}` returned successfully!",
-            collectionName,
-            segmentId);
-        return new ListIDInSegmentResponse(
-            new Response(Response.Status.SUCCESS), response.getEntityIdArrayList());
-      } else {
-        logError(
-            "listIDInSegment in collection `{}`, segment `{}` failed:\n{}",
-            collectionName,
-            segmentId,
-            response.getStatus().toString());
-        return new ListIDInSegmentResponse(
-            new Response(
-                Response.Status.valueOf(response.getStatus().getErrorCodeValue()),
-                response.getStatus().getReason()),
-            Collections.emptyList());
-      }
-    } catch (StatusRuntimeException e) {
-      logError("listIDInSegment RPC failed:\n{}", e.getStatus().toString());
-      return new ListIDInSegmentResponse(
-          new Response(Response.Status.RPC_ERROR, e.toString()), Collections.emptyList());
-    }
+  public List<Long> listIDInSegment(String collectionName, Long segmentId) {
+    return translateExceptions(() -> {
+      GetEntityIDsParam request = GetEntityIDsParam.newBuilder()
+          .setCollectionName(collectionName)
+          .setSegmentId(segmentId)
+          .build();
+      EntityIds response = blockingStub().getEntityIDs(request);
+      checkResponseStatus(response.getStatus());
+      return response.getEntityIdArrayList();
+    });
   }
 
   @Override

+ 2 - 4
src/test/java/io/milvus/client/MilvusGrpcClientTest.java

@@ -741,10 +741,8 @@ class MilvusClientTest {
         .getJSONArray("segments")
         .getJSONObject(0);
 
-    ListIDInSegmentResponse listIDInSegmentResponse =
-        client.listIDInSegment(randomCollectionName, segmentInfo.getLong("id"));
-    assertTrue(listIDInSegmentResponse.ok());
-    assertFalse(listIDInSegmentResponse.getIds().isEmpty());
+    List<Long> entityIds = client.listIDInSegment(randomCollectionName, segmentInfo.getLong("id"));
+    assertFalse(entityIds.isEmpty());
   }
 
   @org.junit.jupiter.api.Test