Skip to content

Commit 3e6cf42

Browse files
yhmoyhmo
authored
GetPersistentSegmentInfo/GetQuerySegmentInfo enhancement (#1876)
Signed-off-by: yhmo <yihua.mo@zilliz.com>
1 parent a988668 commit 3e6cf42

7 files changed

Lines changed: 155 additions & 17 deletions

File tree

‎sdk-core/src/main/java/io/milvus/v2/service/utility/UtilityService.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,9 +377,11 @@ public GetPersistentSegmentInfoResp getPersistentSegmentInfo(MilvusServiceGrpc.M
377377
.segmentID(info.getSegmentID())
378378
.collectionID(info.getCollectionID())
379379
.partitionID(info.getPartitionID())
380+
.collectionName(collectionName)
380381
.numOfRows(info.getNumRows())
381382
.state(info.getState().name())
382383
.level(info.getLevel().name())
384+
.storageVersion(info.getStorageVersion())
383385
.isSorted(info.getIsSorted())
384386
.build());
385387
});
@@ -414,6 +416,7 @@ public GetQuerySegmentInfoResp getQuerySegmentInfo(MilvusServiceGrpc.MilvusServi
414416
.state(info.getState().name())
415417
.level(info.getLevel().name())
416418
.nodeIDs(info.getNodeIdsList())
419+
.storageVersion(info.getStorageVersion())
417420
.isSorted(info.getIsSorted())
418421
.build());
419422
});

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/GetPersistentSegmentInfoResp.java‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,22 @@ public static class PersistentSegmentInfo {
2727
private Long segmentID;
2828
private Long collectionID;
2929
private Long partitionID;
30+
private String collectionName;
3031
private Long numOfRows;
3132
private String state;
3233
private String level;
34+
private Long storageVersion;
3335
private Boolean isSorted;
3436

3537
private PersistentSegmentInfo(PersistentSegmentInfoBuilder builder) {
3638
this.segmentID = builder.segmentID;
3739
this.collectionID = builder.collectionID;
3840
this.partitionID = builder.partitionID;
41+
this.collectionName = builder.collectionName;
3942
this.numOfRows = builder.numOfRows;
4043
this.state = builder.state;
4144
this.level = builder.level;
45+
this.storageVersion = builder.storageVersion;
4246
this.isSorted = builder.isSorted;
4347
}
4448

@@ -70,6 +74,14 @@ public void setPartitionID(Long partitionID) {
7074
this.partitionID = partitionID;
7175
}
7276

77+
public String getCollectionName() {
78+
return collectionName;
79+
}
80+
81+
public void setCollectionName(String collectionName) {
82+
this.collectionName = collectionName;
83+
}
84+
7385
public Long getNumOfRows() {
7486
return numOfRows;
7587
}
@@ -94,6 +106,14 @@ public void setLevel(String level) {
94106
this.level = level;
95107
}
96108

109+
public Long getStorageVersion() {
110+
return storageVersion;
111+
}
112+
113+
public void setStorageVersion(Long storageVersion) {
114+
this.storageVersion = storageVersion;
115+
}
116+
97117
public Boolean getIsSorted() {
98118
return isSorted;
99119
}
@@ -108,9 +128,11 @@ public String toString() {
108128
"segmentID=" + segmentID +
109129
", collectionID=" + collectionID +
110130
", partitionID=" + partitionID +
131+
", collectionName='" + collectionName + '\'' +
111132
", numOfRows=" + numOfRows +
112133
", state='" + state + '\'' +
113134
", level='" + level + '\'' +
135+
", storageVersion=" + storageVersion +
114136
", isSorted=" + isSorted +
115137
'}';
116138
}
@@ -119,9 +141,11 @@ public static class PersistentSegmentInfoBuilder {
119141
private Long segmentID;
120142
private Long collectionID;
121143
private Long partitionID;
144+
private String collectionName;
122145
private Long numOfRows;
123146
private String state;
124147
private String level;
148+
private Long storageVersion;
125149
private Boolean isSorted;
126150

127151
public PersistentSegmentInfoBuilder segmentID(Long segmentID) {
@@ -139,6 +163,11 @@ public PersistentSegmentInfoBuilder partitionID(Long partitionID) {
139163
return this;
140164
}
141165

166+
public PersistentSegmentInfoBuilder collectionName(String collectionName) {
167+
this.collectionName = collectionName;
168+
return this;
169+
}
170+
142171
public PersistentSegmentInfoBuilder numOfRows(Long numOfRows) {
143172
this.numOfRows = numOfRows;
144173
return this;
@@ -154,6 +183,11 @@ public PersistentSegmentInfoBuilder level(String level) {
154183
return this;
155184
}
156185

186+
public PersistentSegmentInfoBuilder storageVersion(Long storageVersion) {
187+
this.storageVersion = storageVersion;
188+
return this;
189+
}
190+
157191
public PersistentSegmentInfoBuilder isSorted(Boolean isSorted) {
158192
this.isSorted = isSorted;
159193
return this;

‎sdk-core/src/main/java/io/milvus/v2/service/utility/response/GetQuerySegmentInfoResp.java‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ public static class QuerySegmentInfo {
3434
private String state;
3535
private String level;
3636
private List<Long> nodeIDs;
37+
private Long storageVersion;
3738
private Boolean isSorted;
3839

3940
private QuerySegmentInfo(QuerySegmentInfoBuilder builder) {
@@ -47,6 +48,7 @@ private QuerySegmentInfo(QuerySegmentInfoBuilder builder) {
4748
this.state = builder.state;
4849
this.level = builder.level;
4950
this.nodeIDs = builder.nodeIDs;
51+
this.storageVersion = builder.storageVersion;
5052
this.isSorted = builder.isSorted;
5153
}
5254

@@ -134,6 +136,14 @@ public void setNodeIDs(List<Long> nodeIDs) {
134136
this.nodeIDs = nodeIDs;
135137
}
136138

139+
public Long getStorageVersion() {
140+
return storageVersion;
141+
}
142+
143+
public void setStorageVersion(Long storageVersion) {
144+
this.storageVersion = storageVersion;
145+
}
146+
137147
public Boolean getIsSorted() {
138148
return isSorted;
139149
}
@@ -155,6 +165,7 @@ public String toString() {
155165
", state='" + state + '\'' +
156166
", level='" + level + '\'' +
157167
", nodeIDs=" + nodeIDs +
168+
", storageVersion=" + storageVersion +
158169
", isSorted=" + isSorted +
159170
'}';
160171
}
@@ -170,6 +181,7 @@ public static class QuerySegmentInfoBuilder {
170181
private String state;
171182
private String level;
172183
private List<Long> nodeIDs = new ArrayList<>();
184+
private Long storageVersion;
173185
private Boolean isSorted;
174186

175187
public QuerySegmentInfoBuilder segmentID(Long segmentID) {
@@ -222,6 +234,11 @@ public QuerySegmentInfoBuilder nodeIDs(List<Long> nodeIDs) {
222234
return this;
223235
}
224236

237+
public QuerySegmentInfoBuilder storageVersion(Long storageVersion) {
238+
this.storageVersion = storageVersion;
239+
return this;
240+
}
241+
225242
public QuerySegmentInfoBuilder isSorted(Boolean isSorted) {
226243
this.isSorted = isSorted;
227244
return this;

‎sdk-core/src/test/java/io/milvus/v2/BaseTest.java‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,36 @@ public void setUp() {
163163

164164
// utility api
165165
when(blockingStub.flush(any())).thenReturn(FlushResponse.newBuilder().setStatus(successStatus).build());
166+
when(blockingStub.getPersistentSegmentInfo(any())).thenReturn(GetPersistentSegmentInfoResponse.newBuilder()
167+
.setStatus(successStatus)
168+
.addInfos(PersistentSegmentInfo.newBuilder()
169+
.setSegmentID(1L)
170+
.setCollectionID(2L)
171+
.setPartitionID(3L)
172+
.setNumRows(4L)
173+
.setState(SegmentState.Flushed)
174+
.setLevel(SegmentLevel.L1)
175+
.setStorageVersion(5L)
176+
.setIsSorted(true)
177+
.build())
178+
.build());
179+
when(blockingStub.getQuerySegmentInfo(any())).thenReturn(GetQuerySegmentInfoResponse.newBuilder()
180+
.setStatus(successStatus)
181+
.addInfos(QuerySegmentInfo.newBuilder()
182+
.setSegmentID(6L)
183+
.setCollectionID(7L)
184+
.setPartitionID(8L)
185+
.setMemSize(9L)
186+
.setNumRows(10L)
187+
.setIndexName("test_index")
188+
.setIndexID(11L)
189+
.setState(SegmentState.Sealed)
190+
.setLevel(SegmentLevel.L1)
191+
.addNodeIds(12L)
192+
.setStorageVersion(13L)
193+
.setIsSorted(true)
194+
.build())
195+
.build());
166196
when(blockingStub.createAlias(any())).thenReturn(successStatus);
167197
when(blockingStub.dropAlias(any())).thenReturn(successStatus);
168198
when(blockingStub.alterAlias(any())).thenReturn(successStatus);

‎sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java‎

Lines changed: 26 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -366,25 +366,35 @@ void testFloatVectors() {
366366
.build());
367367

368368
// master branch, getPersistentSegmentInfo cannot ensure the segment is returned after flush()
369-
while (true) {
370-
// get persistent segment info
371-
GetPersistentSegmentInfoResp pSegInfo = client.getPersistentSegmentInfo(GetPersistentSegmentInfoReq.builder()
369+
long deadline = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(30);
370+
GetPersistentSegmentInfoResp pSegInfo = null;
371+
while (System.currentTimeMillis() < deadline) {
372+
pSegInfo = client.getPersistentSegmentInfo(GetPersistentSegmentInfoReq.builder()
372373
.collectionName(randomCollectionName)
373374
.build());
374-
if (pSegInfo.getSegmentInfos().size() == 0) {
375-
continue;
375+
if (!pSegInfo.getSegmentInfos().isEmpty()) {
376+
break;
377+
}
378+
try {
379+
TimeUnit.MILLISECONDS.sleep(200);
380+
} catch (InterruptedException e) {
381+
Thread.currentThread().interrupt();
382+
Assertions.fail("Interrupted while waiting for persistent segment info in collection: " + randomCollectionName);
376383
}
377-
Assertions.assertEquals(1, pSegInfo.getSegmentInfos().size());
378-
GetPersistentSegmentInfoResp.PersistentSegmentInfo pInfo = pSegInfo.getSegmentInfos().get(0);
379-
Assertions.assertTrue(pInfo.getSegmentID() > 0L);
380-
Assertions.assertTrue(pInfo.getCollectionID() > 0L);
381-
Assertions.assertTrue(pInfo.getPartitionID() > 0L);
382-
Assertions.assertEquals(count, pInfo.getNumOfRows());
383-
Assertions.assertEquals("Flushed", pInfo.getState());
384-
Assertions.assertEquals("L1", pInfo.getLevel());
385-
// Assertions.assertFalse(pInfo.getIsSorted());
386-
break;
387384
}
385+
Assertions.assertNotNull(pSegInfo, "Timed out waiting for persistent segment info response");
386+
Assertions.assertFalse(pSegInfo.getSegmentInfos().isEmpty(), "Timed out waiting for persistent segment info in collection: " + randomCollectionName);
387+
Assertions.assertEquals(1, pSegInfo.getSegmentInfos().size());
388+
GetPersistentSegmentInfoResp.PersistentSegmentInfo pInfo = pSegInfo.getSegmentInfos().get(0);
389+
Assertions.assertTrue(pInfo.getSegmentID() > 0L);
390+
Assertions.assertTrue(pInfo.getCollectionID() > 0L);
391+
Assertions.assertTrue(pInfo.getPartitionID() > 0L);
392+
Assertions.assertEquals(count, pInfo.getNumOfRows());
393+
Assertions.assertEquals(randomCollectionName, pInfo.getCollectionName());
394+
Assertions.assertEquals("Flushed", pInfo.getState());
395+
Assertions.assertEquals("L1", pInfo.getLevel());
396+
Assertions.assertNotNull(pInfo.getStorageVersion());
397+
// Assertions.assertFalse(pInfo.getIsSorted());
388398

389399
// compact
390400
CompactResp compactResp = client.compact(CompactReq.builder()
@@ -429,6 +439,7 @@ void testFloatVectors() {
429439
Assertions.assertEquals("L1", qInfo.getLevel());
430440
Assertions.assertEquals(1, qInfo.getNodeIDs().size());
431441
Assertions.assertTrue(qInfo.getNodeIDs().get(0) > 0L);
442+
Assertions.assertNotNull(qInfo.getStorageVersion());
432443
Assertions.assertTrue(qInfo.getIsSorted());
433444

434445
// create partition, upsert one row to the partition

‎sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2Test.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,7 @@ private static void randomCallMethod(Field field, Method method, Object callInst
274274
String obj = randomStr.generate(10);
275275
randomValues.put(field.getName(), obj);
276276
} else if (fieldType == Long.class || fieldType.getName().equals("long")) {
277-
Long obj = (long) random.nextInt(RANDOM_BOUND);
277+
Long obj = (long) random.nextInt(RANDOM_BOUND) + 1;
278278
randomValues.put(field.getName(), obj);
279279
} else if (fieldType == Integer.class || fieldType.getName().equals("int")) {
280280
Integer obj = random.nextInt(RANDOM_BOUND);

‎sdk-core/src/test/java/io/milvus/v2/service/utility/UtilityTest.java‎

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929

3030
import static org.junit.jupiter.api.Assertions.assertEquals;
3131
import static org.junit.jupiter.api.Assertions.assertNotNull;
32+
import static org.junit.jupiter.api.Assertions.assertTrue;
3233

3334
class UtilityTest extends BaseTest {
3435
Logger logger = LoggerFactory.getLogger(UtilityTest.class);
@@ -181,4 +182,46 @@ void testListFileResources() {
181182
assertEquals("test_resource_2", info2.getName());
182183
assertEquals("/data/test2.parquet", info2.getPath());
183184
}
184-
}
185+
186+
@Test
187+
void testGetPersistentSegmentInfo() {
188+
GetPersistentSegmentInfoResp resp = client_v2.getPersistentSegmentInfo(GetPersistentSegmentInfoReq.builder()
189+
.collectionName("test")
190+
.build());
191+
192+
assertEquals(1, resp.getSegmentInfos().size());
193+
GetPersistentSegmentInfoResp.PersistentSegmentInfo info = resp.getSegmentInfos().get(0);
194+
assertEquals(1L, info.getSegmentID());
195+
assertEquals(2L, info.getCollectionID());
196+
assertEquals(3L, info.getPartitionID());
197+
assertEquals("test", info.getCollectionName());
198+
assertEquals(4L, info.getNumOfRows());
199+
assertEquals("Flushed", info.getState());
200+
assertEquals("L1", info.getLevel());
201+
assertEquals(5L, info.getStorageVersion());
202+
assertTrue(info.getIsSorted());
203+
}
204+
205+
@Test
206+
void testGetQuerySegmentInfo() {
207+
GetQuerySegmentInfoResp resp = client_v2.getQuerySegmentInfo(GetQuerySegmentInfoReq.builder()
208+
.collectionName("test")
209+
.build());
210+
211+
assertEquals(1, resp.getSegmentInfos().size());
212+
GetQuerySegmentInfoResp.QuerySegmentInfo info = resp.getSegmentInfos().get(0);
213+
assertEquals(6L, info.getSegmentID());
214+
assertEquals(7L, info.getCollectionID());
215+
assertEquals(8L, info.getPartitionID());
216+
assertEquals(9L, info.getMemSize());
217+
assertEquals(10L, info.getNumOfRows());
218+
assertEquals("test_index", info.getIndexName());
219+
assertEquals(11L, info.getIndexID());
220+
assertEquals("Sealed", info.getState());
221+
assertEquals("L1", info.getLevel());
222+
assertEquals(1, info.getNodeIDs().size());
223+
assertEquals(12L, info.getNodeIDs().get(0));
224+
assertEquals(13L, info.getStorageVersion());
225+
assertTrue(info.getIsSorted());
226+
}
227+
}

0 commit comments

Comments
 (0)