IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2024-11-22 08:46:56 -06:00
Code Issues Packages Projects Releases Wiki Activity

Updates to Milvus adapter

Browse Source
This commit is contained in:
Madhavan Sridharan 2024-04-15 09:00:01 -04:00
parent af3928beed
commit 41911a0a02
3 changed files with 13 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nb-adapters/adapter-milvus/pom.xml
View File

@ -50,12 +50,12 @@
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java-util</artifactId>
<version>3.24.0</version>
<version>3.25.3</version>
</dependency>
<dependency>
<groupId>io.milvus</groupId>
<artifactId>milvus-sdk-java</artifactId>
<version>2.3.4</version>
<version>2.3.5</version>
</dependency>
</dependencies>

19
nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/MilvusSpace.java
View File

@ -51,10 +51,8 @@ public class MilvusSpace implements AutoCloseable {
* Create a new MilvusSpace Object which stores all stateful contextual information needed to interact
* with the Milvus/Zilliz database instance.
*
* @param name
* The name of this space
* @param cfg
* The configuration ({@link NBConfiguration}) for this nb run
* @param name The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
public MilvusSpace(String name, NBConfiguration cfg) {
this.name = name;
@ -89,16 +87,13 @@ public class MilvusSpace implements AutoCloseable {
).orElseGet(
() -> cfg.getOptional("token")
.orElseThrow(() -> new RuntimeException("You must provide either a token_file or a token to " +
"configure a Milvus client"))
"configure a Milvus/Zilliz client"))
);
builder = builder.withToken(requiredToken);
ConnectParam connectParams = builder.build();
logger.info(this.name + ": Creating new Milvus/Zilliz Client with (masked) " +
"token [" + MilvusAdapterUtils.maskDigits(builder.getToken()) + "], uri/endpoint [" + builder.getUri() + "]"
);
logger.info("{}: Creating new Milvus/Zilliz Client with (masked) token [{}], uri/endpoint [{}]",
this.name, MilvusAdapterUtils.maskDigits(builder.getToken()), builder.getUri());
return new MilvusServiceClient(connectParams);
}
@ -117,8 +112,8 @@ public class MilvusSpace implements AutoCloseable {
.setDescription("the URI endpoint in which the database is running.")
)
.add(
Param.optional(List.of("database_name","database"))
.setDescription("the name of the database to use. Defaults to 'baselines'")
Param.optional(List.of("database_name", "database"))
.setDescription("the name of the database to use.")
)
.asReadOnly();
}

9
nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusCreateCollectionOpDispenser.java
View File

@ -19,6 +19,7 @@ package io.nosqlbench.adapter.milvus.opdispensers;
import io.milvus.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType;
import io.milvus.param.collection.CollectionSchemaParam;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.nosqlbench.adapter.milvus.MilvusDriverAdapter;
@ -70,11 +71,10 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
CreateCollectionParam.Builder::withDatabaseName);
List<FieldType> fieldTypes = buildFieldTypesStruct(
op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey),
ebF
op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey)
);
final LongFunction<CreateCollectionParam.Builder> f = ebF;
ebF = l -> f.apply(l).withFieldTypes(fieldTypes);
ebF = l -> f.apply(l).withSchema(CollectionSchemaParam.newBuilder().withFieldTypes(fieldTypes).build());
final LongFunction<CreateCollectionParam.Builder> lastF = ebF;
return l -> lastF.apply(l).build();
@ -95,10 +95,9 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
* Function to build the {@link FieldType}s for the {@link CreateCollectionParam}.
*
* @param fieldTypesData The static map of config data from the create collection request
* @param ebF
* @return a list of static field types
*/
private List<FieldType> buildFieldTypesStruct(Map<String, ParsedOp> fieldTypesData, LongFunction<CreateCollectionParam.Builder> ebF) {
private List<FieldType> buildFieldTypesStruct(Map<String, ParsedOp> fieldTypesData) {
List<FieldType> fieldTypes = new ArrayList<>();
fieldTypesData.forEach((name, fieldspec) -> {
FieldType.Builder builder = FieldType.newBuilder()