Skip to content

Commit 0644b73

Browse files
committed
Improve volume upload interaction
Signed-off-by: lentitude2tk <xushuang.hu@zilliz.com>
1 parent 83678b1 commit 0644b73

9 files changed

Lines changed: 1249 additions & 109 deletions

File tree

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriter.java

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121

2222
import com.google.common.collect.Lists;
2323
import com.google.gson.JsonObject;
24-
import io.milvus.bulkwriter.common.clientenum.ConnectType;
2524
import io.milvus.bulkwriter.model.UploadFilesResult;
2625
import io.milvus.bulkwriter.request.volume.UploadFilesRequest;
2726
import io.milvus.common.utils.ExceptionUtils;
@@ -64,7 +63,7 @@ public VolumeBulkWriter(VolumeBulkWriterParam bulkWriterParam) throws IOExceptio
6463
private VolumeFileManager initVolumeFileManagerParams(VolumeBulkWriterParam bulkWriterParam) throws IOException {
6564
VolumeFileManagerParam volumeFileManagerParam = VolumeFileManagerParam.newBuilder()
6665
.withCloudEndpoint(bulkWriterParam.getCloudEndpoint()).withApiKey(bulkWriterParam.getApiKey())
67-
.withVolumeName(bulkWriterParam.getVolumeName()).withConnectType(ConnectType.AUTO)
66+
.withVolumeName(bulkWriterParam.getVolumeName()).withConnectType(bulkWriterParam.getConnectType())
6867
.build();
6968
return new VolumeFileManager(volumeFileManagerParam);
7069
}
@@ -149,9 +148,13 @@ protected void callBack(List<String> fileList) {
149148

150149
@Override
151150
public void close() throws Exception {
152-
logger.info("execute remaining actions to prevent loss of memory data or residual empty directories.");
153-
exit();
154-
logger.info(String.format("RemoteBulkWriter done! output remote files: %s", getBatchFiles()));
151+
try {
152+
logger.info("execute remaining actions to prevent loss of memory data or residual empty directories.");
153+
exit();
154+
logger.info(String.format("RemoteBulkWriter done! output remote files: %s", getBatchFiles()));
155+
} finally {
156+
volumeFileManager.shutdownGracefully();
157+
}
155158
}
156159

157160
private void serialImportData(List<String> fileList) {
@@ -199,6 +202,6 @@ private static String getMinioFilePath(String remotePath, String relativeFilePat
199202
relativeFilePath = relativeFilePath.startsWith("/") ? relativeFilePath.substring(1) : relativeFilePath;
200203
Path relative = Paths.get(relativeFilePath);
201204
Path joinedPath = remote.resolve(relative);
202-
return joinedPath.toString();
205+
return joinedPath.toString().replace("\\", "/");
203206
}
204207
}

sdk-bulkwriter/src/main/java/io/milvus/bulkwriter/VolumeBulkWriterParam.java

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
package io.milvus.bulkwriter;
2121

2222
import io.milvus.bulkwriter.common.clientenum.BulkFileType;
23+
import io.milvus.bulkwriter.common.clientenum.ConnectType;
2324
import io.milvus.bulkwriter.common.utils.V2AdapterUtils;
2425
import io.milvus.exception.ParamException;
2526
import io.milvus.param.ParamUtils;
@@ -42,6 +43,7 @@ public class VolumeBulkWriterParam {
4243
private final String cloudEndpoint;
4344
private final String apiKey;
4445
private final String volumeName;
46+
private final ConnectType connectType;
4547

4648
private VolumeBulkWriterParam(Builder builder) {
4749
this.collectionSchema = builder.collectionSchema;
@@ -53,6 +55,7 @@ private VolumeBulkWriterParam(Builder builder) {
5355
this.cloudEndpoint = builder.cloudEndpoint;
5456
this.apiKey = builder.apiKey;
5557
this.volumeName = builder.volumeName;
58+
this.connectType = builder.connectType;
5659
}
5760

5861
public CreateCollectionReq.CollectionSchema getCollectionSchema() {
@@ -87,6 +90,10 @@ public String getVolumeName() {
8790
return volumeName;
8891
}
8992

93+
public ConnectType getConnectType() {
94+
return connectType;
95+
}
96+
9097
@Override
9198
public String toString() {
9299
return "VolumeBulkWriterParam{" +
@@ -96,6 +103,7 @@ public String toString() {
96103
", fileType=" + fileType +
97104
", cloudEndpoint='" + cloudEndpoint + '\'' +
98105
", volumeName='" + volumeName + '\'' +
106+
", connectType=" + connectType +
99107
'}';
100108
}
101109

@@ -118,6 +126,8 @@ public static final class Builder {
118126

119127
private String volumeName;
120128

129+
private ConnectType connectType = ConnectType.AUTO;
130+
121131
private Builder() {
122132
}
123133

@@ -184,6 +194,11 @@ public Builder withVolumeName(String volumeName) {
184194
return this;
185195
}
186196

197+
public Builder withConnectType(ConnectType connectType) {
198+
this.connectType = connectType;
199+
return this;
200+
}
201+
187202
/**
188203
* Verifies parameters and creates a new {@link VolumeBulkWriterParam} instance.
189204
*
@@ -195,6 +210,9 @@ public VolumeBulkWriterParam build() throws ParamException {
195210
if (collectionSchema == null) {
196211
throw new ParamException("collectionSchema cannot be null");
197212
}
213+
ParamUtils.CheckNullEmptyString(cloudEndpoint, "cloudEndpoint");
214+
ParamUtils.CheckNullEmptyString(apiKey, "apiKey");
215+
ParamUtils.CheckNullEmptyString(volumeName, "volumeName");
198216

199217
return new VolumeBulkWriterParam(this);
200218
}

0 commit comments

Comments
 (0)