Skip to content

Commit 2694a5b

Browse files
nianliuunianliuu
authored
Merge pull request #30 from zilliztech/dev
dev
2 parents 461630d + fa1888c commit 2694a5b

File tree

7 files changed

+114
-79
lines changed

7 files changed

+114
-79
lines changed

‎pom.xml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141
<dependency>
4242
<groupId>io.milvus</groupId>
4343
<artifactId>milvus-sdk-java</artifactId>
44-
<version>2.4.1</version>
44+
<version>2.5.1</version>
4545
</dependency>
4646
<dependency>
4747
<groupId>junit</groupId>
@@ -68,7 +68,7 @@
6868
<dependency>
6969
<groupId>io.netty</groupId>
7070
<artifactId>netty-common</artifactId>
71-
<version>4.1.100.Final</version>
71+
<version>4.1.115.Final</version>
7272
</dependency>
7373
<dependency>
7474
<groupId>io.netty</groupId>

‎src/main/java/com/milvus/io/kafka/MilvusSinkConnector.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import java.util.List;
1313
import java.util.Map;
1414

15-
public class MilvusSinkConnector extends SinkConnector{
15+
public class MilvusSinkConnector extends SinkConnector {
1616

1717
private static final Logger log = LoggerFactory.getLogger(MilvusSinkConnector.class);
1818
private Map<String, String> configProperties;
@@ -24,7 +24,7 @@ public void start(Map<String, String> props) {
2424
configProperties = props;
2525
// validation
2626
new MilvusSinkConnectorConfig(props);
27-
}catch (ConfigException e){
27+
} catch (ConfigException e) {
2828
throw new ConfigException("Couldn't start MilvusSinkConnector due to configuration error", e);
2929
}
3030
}

‎src/main/java/com/milvus/io/kafka/MilvusSinkConnectorConfig.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,5 +34,7 @@ public Password getToken() {
3434
return getPassword(TOKEN);
3535
}
3636

37-
public String getCollectionName(){return getString(COLLECTION_NAME);}
37+
public String getCollectionName() {
38+
return getString(COLLECTION_NAME);
39+
}
3840
}

‎src/main/java/com/milvus/io/kafka/MilvusSinkTask.java‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
package com.milvus.io.kafka;
22

3-
import com.alibaba.fastjson.JSONObject;
3+
import com.google.gson.JsonObject;
4+
import static com.milvus.io.kafka.MilvusSinkConnectorConfig.TOKEN;
45
import com.milvus.io.kafka.helper.MilvusClientHelper;
56
import com.milvus.io.kafka.utils.DataConverter;
67
import com.milvus.io.kafka.utils.Utils;
78
import com.milvus.io.kafka.utils.VersionUtil;
89
import io.milvus.v2.client.MilvusClientV2;
9-
import io.milvus.v2.service.collection.request.CreateCollectionReq;
1010
import io.milvus.v2.service.collection.request.DescribeCollectionReq;
1111
import io.milvus.v2.service.collection.request.GetLoadStateReq;
1212
import io.milvus.v2.service.collection.request.HasCollectionReq;
@@ -18,9 +18,10 @@
1818
import org.slf4j.Logger;
1919
import org.slf4j.LoggerFactory;
2020

21-
import java.util.*;
22-
23-
import static com.milvus.io.kafka.MilvusSinkConnectorConfig.TOKEN;
21+
import java.util.ArrayList;
22+
import java.util.Collection;
23+
import java.util.List;
24+
import java.util.Map;
2425

2526
public class MilvusSinkTask extends SinkTask {
2627

@@ -58,7 +59,7 @@ private void preValidate() {
5859
throw new RuntimeException("Collection not exist" + config.getCollectionName());
5960
}
6061
// check if the collection is loaded
61-
if (!myMilvusClient.getLoadState(GetLoadStateReq.builder().collectionName(config.getCollectionName()).build())){
62+
if (!myMilvusClient.getLoadState(GetLoadStateReq.builder().collectionName(config.getCollectionName()).build())) {
6263
log.error("Collection not loaded");
6364
throw new RuntimeException("Collection not loaded" + config.getCollectionName());
6465
}
@@ -68,36 +69,36 @@ private void preValidate() {
6869
@Override
6970
public void put(Collection<SinkRecord> records) {
7071
log.info("Putting {} records to Milvus.", records.size());
71-
if(records.isEmpty()) {
72+
if (records.isEmpty()) {
7273
log.info("No records to put.");
7374
return;
7475
}
7576

7677
// not support dynamic schema for now, for dynamic schema, we need to put the data into a JSONObject
77-
List<JSONObject> datas = new ArrayList<>();
78+
List<JsonObject> datas = new ArrayList<>();
7879
for (SinkRecord record : records) {
7980
log.debug("Writing {} to Milvus.", record);
80-
if(record.value() == null) {
81+
if (record.value() == null) {
8182
log.warn("Skipping record with null value.");
8283
continue;
8384
}
8485
try {
85-
JSONObject data = converter.convertRecord(record, response.getCollectionSchema());
86+
JsonObject data = converter.convertRecord(record, response.getCollectionSchema());
8687
datas.add(data);
87-
}catch (Exception e){
88+
} catch (Exception e) {
8889
log.error("Failed to convert record to JSONObject, skip it", e);
8990
}
9091
}
9192

92-
if(!response.getAutoID()){
93+
if (!response.getAutoID()) {
9394
// default to use upsert
9495
UpsertReq upsertReq = UpsertReq.builder()
9596
.collectionName(config.getCollectionName())
9697
.data(datas)
9798
.build();
9899
log.info("Upserting data to collection: {} with datas: {}", config.getCollectionName(), datas);
99100
myMilvusClient.upsert(upsertReq);
100-
}else {
101+
} else {
101102
InsertReq insertReq = InsertReq.builder()
102103
.collectionName(config.getCollectionName())
103104
.data(datas)

‎src/main/java/com/milvus/io/kafka/helper/MilvusClientHelper.java‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,6 @@
22

33
import com.milvus.io.kafka.MilvusSinkConnectorConfig;
44
import com.milvus.io.kafka.utils.Utils;
5-
import io.milvus.client.MilvusServiceClient;
6-
import io.milvus.param.ConnectParam;
75
import io.milvus.v2.client.ConnectConfig;
86
import io.milvus.v2.client.MilvusClientV2;
97

‎src/main/java/com/milvus/io/kafka/utils/DataConverter.java‎

Lines changed: 45 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
package com.milvus.io.kafka.utils;
22

3-
import com.alibaba.fastjson.JSONObject;
43
import com.google.common.collect.Lists;
54
import com.google.gson.Gson;
5+
import com.google.gson.JsonObject;
6+
import com.google.gson.JsonParser;
67
import com.milvus.io.kafka.MilvusSinkConnectorConfig;
7-
import io.milvus.param.dml.InsertParam;
8+
import io.milvus.common.utils.JsonUtils;
89
import io.milvus.v2.common.DataType;
910
import io.milvus.v2.service.collection.request.CreateCollectionReq;
1011
import org.apache.kafka.connect.data.Struct;
@@ -13,63 +14,66 @@
1314
import org.slf4j.LoggerFactory;
1415

1516
import java.nio.ByteBuffer;
16-
import java.util.*;
17-
import java.util.stream.Collectors;
17+
import java.util.HashMap;
18+
import java.util.List;
1819

1920
public class DataConverter {
2021

21-
private final MilvusSinkConnectorConfig config;
22-
2322
private static final Logger log = LoggerFactory.getLogger(DataConverter.class);
23+
private final MilvusSinkConnectorConfig config;
2424

2525
public DataConverter(MilvusSinkConnectorConfig config) {
2626
this.config = config;
2727
}
28+
2829
/*
29-
* Convert SinkRecord to JSONObject
30+
* Convert SinkRecord to JsonObject
3031
*/
31-
public JSONObject convertRecord(SinkRecord sr, CreateCollectionReq.CollectionSchema collectionSchema) {
32-
// parse sinkRecord to get filed name and value
33-
if(sr.value() instanceof Struct) {
34-
return parseValue((Struct)sr.value(), collectionSchema);
35-
}else if (sr.value() instanceof HashMap) {
36-
return parseValue((HashMap<?, ?>)sr.value(), collectionSchema);
37-
}else {
38-
throw new RuntimeException("Unsupported SinkRecord data type" + sr.value());
32+
public JsonObject convertRecord(SinkRecord sr, CreateCollectionReq.CollectionSchema collectionSchema) {
33+
// parse sinkRecord to get field name and value
34+
if (sr.value() instanceof Struct) {
35+
return parseValue((Struct) sr.value(), collectionSchema);
36+
} else if (sr.value() instanceof HashMap) {
37+
return parseValue((HashMap<?, ?>) sr.value(), collectionSchema);
38+
} else {
39+
throw new RuntimeException("Unsupported SinkRecord data type: " + sr.value());
3940
}
4041
}
4142

42-
private JSONObject parseValue(HashMap<?, ?> mapValue, CreateCollectionReq.CollectionSchema collectionSchema) {
43-
JSONObject fields = new JSONObject();
43+
private JsonObject parseValue(HashMap<?, ?> mapValue, CreateCollectionReq.CollectionSchema collectionSchema) {
44+
JsonObject fields = new JsonObject();
45+
Gson gson = new Gson();
4446
mapValue.forEach((field, value) -> {
45-
if(collectionSchema.getField(field.toString())!=null){
47+
if (collectionSchema.getField(field.toString()) != null) {
4648
// if the key exists in the collection, store the value by collectionSchema DataType
47-
fields.put(field.toString(), castValueToType(value, collectionSchema.getField(field.toString()).getDataType()));
48-
}else {
49+
Object object = convertValueByMilvusType(value, collectionSchema.getField(field.toString()).getDataType());
50+
fields.add(field.toString(), gson.toJsonTree(object));
51+
} else {
4952
log.warn("Field {} not exists in collection", field);
5053
}
51-
5254
});
5355
return fields;
5456
}
5557

56-
private JSONObject parseValue(Struct structValue, CreateCollectionReq.CollectionSchema collectionSchema) {
57-
JSONObject fields = new JSONObject();
58-
58+
private JsonObject parseValue(Struct structValue, CreateCollectionReq.CollectionSchema collectionSchema) {
59+
JsonObject fields = new JsonObject();
60+
Gson gson = new Gson();
5961
structValue.schema().fields().forEach(field -> {
60-
if(collectionSchema.getField(field.name()) != null){
62+
if (collectionSchema.getField(field.name()) != null) {
6163
// if the key exists in the collection, store the value by collectionSchema DataType
62-
fields.put(field.toString(), castValueToType(structValue.get(field.name()), collectionSchema.getField(field.name()).getDataType()));
63-
}else {
64+
Object object = convertValueByMilvusType(structValue.get(field.name()), collectionSchema.getField(field.name()).getDataType());
65+
fields.add(field.name(), gson.toJsonTree(object));
66+
} else {
6467
log.warn("Field {} not exists in collection", field);
6568
}
6669
});
6770

6871
return fields;
6972
}
7073

71-
private Object castValueToType(Object value, DataType dataType) {
72-
switch (dataType){
74+
private Object convertValueByMilvusType(Object value, DataType dataType) {
75+
Gson gson = new Gson();
76+
switch (dataType) {
7377
case Bool:
7478
return Boolean.parseBoolean(value.toString());
7579
case Int8:
@@ -87,36 +91,37 @@ private Object castValueToType(Object value, DataType dataType) {
8791
case String:
8892
return value.toString();
8993
case JSON:
90-
Gson gson = new Gson();
9194
return gson.toJson(value);
9295
case BinaryVector:
9396
return parseBinaryVectorField(value.toString());
9497
case FloatVector:
9598
return parseFloatVectorField(value.toString());
99+
case SparseFloatVector:
100+
return gson.toJsonTree(value).getAsJsonObject();
96101
default:
97-
throw new RuntimeException("Unsupported data type" + dataType);
102+
throw new RuntimeException("Unsupported data type: " + dataType);
98103
}
99104
}
100105

101-
protected List<Float> parseFloatVectorField(String vectors){
106+
protected List<Float> parseFloatVectorField(String vectors) {
102107
try {
103108
log.debug("parse float vectors: {}", vectors);
104109

105110
String[] vectorArrays = vectors.replaceAll("\\[", "").replaceAll("\\]", "")
106-
.replaceAll(" ","").split(",");
111+
.replaceAll(" ", "").split(",");
107112

108113
List<Float> floatList = Lists.newLinkedList();
109114
for (String vector : vectorArrays) {
110115
floatList.add(Float.valueOf(vector));
111116
}
112117

113118
return floatList;
114-
}catch (Exception e){
115-
throw new RuntimeException("parse float vector field error: " + e.getMessage() + vectors);
119+
} catch (Exception e) {
120+
throw new RuntimeException("parse float vector field error: " + e.getMessage() + " " + vectors);
116121
}
117-
118122
}
119-
protected ByteBuffer parseBinaryVectorField(String vectors){
123+
124+
protected ByteBuffer parseBinaryVectorField(String vectors) {
120125
try {
121126
log.debug("parse binary vectors: {}", vectors);
122127

@@ -130,8 +135,8 @@ protected ByteBuffer parseBinaryVectorField(String vectors){
130135
}
131136

132137
return buffer;
133-
}catch (Exception e){
134-
throw new RuntimeException("parse binary vector field error: " + e.getMessage() + vectors);
138+
} catch (Exception e) {
139+
throw new RuntimeException("parse binary vector field error: " + e.getMessage() + " " + vectors);
135140
}
136141
}
137-
}
142+
}

0 commit comments

Comments
 (0)