Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support insert dynamic values by column-based #690

Merged
merged 1 commit into from
Oct 27, 2023
Merged

Support insert dynamic values by column-based #690

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/main/java/io/milvus/param/Constant.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ public class Constant {
public static final String DEFAULT_INDEX_NAME = "";
public final static String OFFSET = "offset";
public final static String LIMIT = "limit";
public final static String DYNAMIC_FIELD_NAME = "$meta";

// constant values for general
public static final String TTL_SECONDS = "collection.ttl.seconds";
Expand Down
49 changes: 37 additions & 12 deletions src/main/java/io/milvus/param/ParamUtils.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,13 +268,15 @@ private static void fillFieldsData(InsertParam requestParam, DescCollResponseWra
List<JSONObject> rowFields = requestParam.getRows();

if (CollectionUtils.isNotEmpty(columnFields)) {
checkAndSetColumnData(requestParam, wrapper.getFields(), insertBuilder, columnFields);
checkAndSetColumnData(wrapper, insertBuilder, columnFields);
} else {
checkAndSetRowData(wrapper, insertBuilder, rowFields);
}
}

private static void checkAndSetColumnData(InsertParam requestParam, List<FieldType> fieldTypes, InsertRequest.Builder insertBuilder, List<InsertParam.Field> fields) {
private static void checkAndSetColumnData(DescCollResponseWrapper wrapper, InsertRequest.Builder insertBuilder, List<InsertParam.Field> fields) {
List<FieldType> fieldTypes = wrapper.getFields();

// gen fieldData
// make sure the field order must be consisted with collection schema
for (FieldType fieldType : fieldTypes) {
Expand All @@ -288,7 +290,7 @@ private static void checkAndSetColumnData(InsertParam requestParam, List<FieldTy
checkFieldData(fieldType, field);

found = true;
insertBuilder.addFieldsData(genFieldData(field.getName(), fieldType.getDataType(), field.getValues()));
insertBuilder.addFieldsData(genFieldData(fieldType, field.getValues()));
break;
}

Expand All @@ -298,18 +300,40 @@ private static void checkAndSetColumnData(InsertParam requestParam, List<FieldTy
throw new ParamException(msg);
}
}

// deal with dynamicField
if (wrapper.getEnableDynamicField()) {
for (InsertParam.Field field : fields) {
if (field.getName().equals(Constant.DYNAMIC_FIELD_NAME)) {
FieldType dynamicType = FieldType.newBuilder()
.withName(Constant.DYNAMIC_FIELD_NAME)
.withDataType(DataType.JSON)
.withIsDynamic(true)
.build();
checkFieldData(dynamicType, field);
insertBuilder.addFieldsData(genFieldData(dynamicType, field.getValues(), true));
break;
}
}
}
}

private static void checkAndSetRowData(DescCollResponseWrapper wrapper, InsertRequest.Builder insertBuilder, List<JSONObject> rows) {
List<FieldType> fieldTypes = wrapper.getFields();

Map<String, InsertDataInfo> nameInsertInfo = new HashMap<>();
InsertDataInfo insertDynamicDataInfo = InsertDataInfo.builder().dataType(DataType.JSON).data(new LinkedList<>()).build();
InsertDataInfo insertDynamicDataInfo = InsertDataInfo.builder().fieldType(
FieldType.newBuilder()
.withName(Constant.DYNAMIC_FIELD_NAME)
.withDataType(DataType.JSON)
.withIsDynamic(true)
.build())
.data(new LinkedList<>()).build();
for (JSONObject row : rows) {
for (FieldType fieldType : fieldTypes) {
String fieldName = fieldType.getName();
InsertDataInfo insertDataInfo = nameInsertInfo.getOrDefault(fieldName, InsertDataInfo.builder()
.fieldName(fieldName).dataType(fieldType.getDataType()).data(new LinkedList<>()).build());
.fieldType(fieldType).data(new LinkedList<>()).build());

// check normalField
Object rowFieldData = row.get(fieldName);
Expand Down Expand Up @@ -345,10 +369,10 @@ private static void checkAndSetRowData(DescCollResponseWrapper wrapper, InsertRe

for (String fieldNameKey : nameInsertInfo.keySet()) {
InsertDataInfo insertDataInfo = nameInsertInfo.get(fieldNameKey);
insertBuilder.addFieldsData(genFieldData(insertDataInfo.getFieldName(), insertDataInfo.getDataType(), insertDataInfo.getData()));
insertBuilder.addFieldsData(genFieldData(insertDataInfo.getFieldType(), insertDataInfo.getData()));
}
if (wrapper.getEnableDynamicField()) {
insertBuilder.addFieldsData(genFieldData(insertDynamicDataInfo.getFieldName(), insertDynamicDataInfo.getDataType(), insertDynamicDataInfo.getData(), Boolean.TRUE));
insertBuilder.addFieldsData(genFieldData(insertDynamicDataInfo.getFieldType(), insertDynamicDataInfo.getData(), Boolean.TRUE));
}
}

Expand Down Expand Up @@ -544,15 +568,17 @@ private static long getGuaranteeTimestamp(ConsistencyLevelEnum consistencyLevel,
add(DataType.BinaryVector);
}};

private static FieldData genFieldData(String fieldName, DataType dataType, List<?> objects) {
return genFieldData(fieldName, dataType, objects, Boolean.FALSE);
private static FieldData genFieldData(FieldType fieldType, List<?> objects) {
return genFieldData(fieldType, objects, Boolean.FALSE);
}

@SuppressWarnings("unchecked")
private static FieldData genFieldData(String fieldName, DataType dataType, List<?> objects, boolean isDynamic) {
private static FieldData genFieldData(FieldType fieldType, List<?> objects, boolean isDynamic) {
if (objects == null) {
throw new ParamException("Cannot generate FieldData from null object");
}
DataType dataType = fieldType.getDataType();
String fieldName = fieldType.getName();
FieldData.Builder builder = FieldData.newBuilder();
if (vectorDataType.contains(dataType)) {
if (dataType == DataType.FloatVector) {
Expand Down Expand Up @@ -719,8 +745,7 @@ public static List<KeyValuePair> AssembleKvPair(Map<String, String> sourceMap) {
@Builder
@Getter
public static class InsertDataInfo {
private final String fieldName;
private final DataType dataType;
private final FieldType fieldType;
private final LinkedList<Object> data;
}
}
57 changes: 38 additions & 19 deletions src/test/java/io/milvus/client/MilvusClientDockerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1244,11 +1244,7 @@ void testFloatVectorIndex() {
indexTypes.put(IndexType.IVF_FLAT, "{\"nlist\":128}");
indexTypes.put(IndexType.IVF_SQ8, "{\"nlist\":128}");
indexTypes.put(IndexType.IVF_PQ, "{\"nlist\":128, \"m\":16, \"nbits\":8}");
indexTypes.put(IndexType.ANNOY, "{\"n_trees\":16}");
indexTypes.put(IndexType.HNSW, "{\"M\":16,\"efConstruction\":64}");
indexTypes.put(IndexType.RHNSW_FLAT, "{\"M\":16,\"efConstruction\":64}");
indexTypes.put(IndexType.RHNSW_PQ, "{\"M\":16,\"efConstruction\":64, \"PQM\":16}");
indexTypes.put(IndexType.RHNSW_SQ, "{\"M\":16,\"efConstruction\":64}");

List<MetricType> metricTypes = new ArrayList<>();
metricTypes.add(MetricType.L2);
Expand Down Expand Up @@ -1297,8 +1293,8 @@ void testBinaryVectorIndex() {

// test all supported indexes
List<MetricType> flatMetricTypes = new ArrayList<>();
flatMetricTypes.add(MetricType.SUBSTRUCTURE);
flatMetricTypes.add(MetricType.SUPERSTRUCTURE);
flatMetricTypes.add(MetricType.HAMMING);
flatMetricTypes.add(MetricType.JACCARD);

for (MetricType metric : flatMetricTypes) {
testIndex(randomCollectionName, field2Name, IndexType.BIN_FLAT, metric, "{}", Boolean.TRUE);
Expand All @@ -1308,7 +1304,6 @@ void testBinaryVectorIndex() {
List<MetricType> ivfMetricTypes = new ArrayList<>();
ivfMetricTypes.add(MetricType.HAMMING);
ivfMetricTypes.add(MetricType.JACCARD);
ivfMetricTypes.add(MetricType.TANIMOTO);

for (MetricType metric : ivfMetricTypes) {
testIndex(randomCollectionName, field2Name, IndexType.BIN_IVF_FLAT, metric, "{\"nlist\":128}", Boolean.TRUE);
Expand Down Expand Up @@ -1394,11 +1389,11 @@ void testDynamicField() {

// JSON field
JSONObject info = new JSONObject();
info.put("row-based-info", i);
info.put("row_based_info", i);
row.put(field3Name, info);

// extra meta is automatically stored in dynamic field
row.put("extra_meta", i % 3 == 0);
row.put("row_based_extra", i % 3 == 0);
row.put(generator.generate(5), 100);

rows.add(row);
Expand All @@ -1416,19 +1411,25 @@ void testDynamicField() {
// insert data by column-based
List<Long> ids = new ArrayList<>();
List<JSONObject> infos = new ArrayList<>();
List<JSONObject> dynamics = new ArrayList<>();
for (long i = 0L; i < rowCount; ++i) {
ids.add(rowCount + i);
JSONObject obj = new JSONObject();
obj.put("column-based-info", i);
obj.put("column_based_info", i);
obj.put(generator.generate(5), i);
infos.add(obj);

JSONObject dynamic = new JSONObject();
dynamic.put(String.format("column_based_extra_%d", i), i);
dynamics.add(dynamic);
}
List<List<Float>> vectors = generateFloatVectors(rowCount);

List<InsertParam.Field> fieldsInsert = new ArrayList<>();
fieldsInsert.add(new InsertParam.Field(field1Name, ids));
fieldsInsert.add(new InsertParam.Field(field2Name, vectors));
fieldsInsert.add(new InsertParam.Field(field3Name, infos));
fieldsInsert.add(new InsertParam.Field(Constant.DYNAMIC_FIELD_NAME, dynamics));

InsertParam insertColumnsParam = InsertParam.newBuilder()
.withCollectionName(randomCollectionName)
Expand All @@ -1451,8 +1452,8 @@ void testDynamicField() {
System.out.println("Collection row count: " + stat.getRowCount());

// retrieve rows
String expr = "extra_meta == true";
List<String> outputFields = Arrays.asList(field3Name, "extra_meta");
String expr = "row_based_extra == true";
List<String> outputFields = Arrays.asList(field3Name, "row_based_extra");
QueryParam queryParam = QueryParam.newBuilder()
.withCollectionName(randomCollectionName)
.withExpr(expr)
Expand All @@ -1464,13 +1465,11 @@ void testDynamicField() {

QueryResultsWrapper queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
List<QueryResultsWrapper.RowRecord> records = queryResultsWrapper.getRowRecords();
System.out.println("Query results:");
System.out.println("Query results with expr: " + expr);
for (QueryResultsWrapper.RowRecord record:records) {
System.out.println(record);
Object extraMeta = record.get("extra_meta");
if (extraMeta != null) {
System.out.println("'extra_meta' is from dynamic field, value: " + extraMeta);
}
Object extraMeta = record.get("row_based_extra");
System.out.println("'row_based_extra' is from dynamic field, value: " + extraMeta);
}

// search
Expand All @@ -1496,13 +1495,33 @@ void testDynamicField() {
System.out.println("The result of No." + i + " target vector:");
for (SearchResultsWrapper.IDScore score:scores) {
System.out.println(score);
Object extraMeta = score.get("extra_meta");
Object extraMeta = score.get("row_based_extra");
if (extraMeta != null) {
System.out.println("'extra_meta' is from dynamic field, value: " + extraMeta);
System.out.println("'row_based_extra' is from dynamic field, value: " + extraMeta);
}
}
}

// retrieve dynamic values inserted by column-based
expr = "column_based_extra_1 == 1";
queryParam = QueryParam.newBuilder()
.withCollectionName(randomCollectionName)
.withExpr("column_based_extra_1 == 1")
.withOutFields(Collections.singletonList("*"))
.build();

queryR = client.query(queryParam);
Assertions.assertEquals(R.Status.Success.getCode(), queryR.getStatus().intValue());

queryResultsWrapper = new QueryResultsWrapper(queryR.getData());
records = queryResultsWrapper.getRowRecords();
System.out.println("Query results with expr: " + expr);
for (QueryResultsWrapper.RowRecord record:records) {
System.out.println(record);
long id = (long)record.get(field1Name);
Assertions.assertEquals((long)rowCount+1L, id);
}

// drop collection
R<RpcStatus> dropR = client.dropCollection(DropCollectionParam.newBuilder()
.withCollectionName(randomCollectionName)
Expand Down
Loading