IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2025-02-25 18:55:28 -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> <dependency>
<groupId>com.google.protobuf</groupId> <groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java-util</artifactId> <artifactId>protobuf-java-util</artifactId>
<version>3.24.0</version> <version>3.25.3</version>
</dependency> </dependency>
<dependency> <dependency>
<groupId>io.milvus</groupId> <groupId>io.milvus</groupId>
<artifactId>milvus-sdk-java</artifactId> <artifactId>milvus-sdk-java</artifactId>
<version>2.3.4</version> <version>2.3.5</version>
</dependency> </dependency>
</dependencies> </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 * Create a new MilvusSpace Object which stores all stateful contextual information needed to interact
* with the Milvus/Zilliz database instance. * with the Milvus/Zilliz database instance.
* *
* @param name * @param name The name of this space
* The name of this space * @param cfg The configuration ({@link NBConfiguration}) for this nb run
* @param cfg
* The configuration ({@link NBConfiguration}) for this nb run
*/ */
public MilvusSpace(String name, NBConfiguration cfg) { public MilvusSpace(String name, NBConfiguration cfg) {
this.name = name; this.name = name;
@ -89,16 +87,13 @@ public class MilvusSpace implements AutoCloseable {
).orElseGet( ).orElseGet(
() -> cfg.getOptional("token") () -> cfg.getOptional("token")
.orElseThrow(() -> new RuntimeException("You must provide either a token_file or a token to " + .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); builder = builder.withToken(requiredToken);
ConnectParam connectParams = builder.build(); ConnectParam connectParams = builder.build();
logger.info("{}: Creating new Milvus/Zilliz Client with (masked) token [{}], uri/endpoint [{}]",
logger.info(this.name + ": Creating new Milvus/Zilliz Client with (masked) " + this.name, MilvusAdapterUtils.maskDigits(builder.getToken()), builder.getUri());
"token [" + MilvusAdapterUtils.maskDigits(builder.getToken()) + "], uri/endpoint [" + builder.getUri() + "]"
);
return new MilvusServiceClient(connectParams); return new MilvusServiceClient(connectParams);
} }
@ -117,8 +112,8 @@ public class MilvusSpace implements AutoCloseable {
.setDescription("the URI endpoint in which the database is running.") .setDescription("the URI endpoint in which the database is running.")
) )
.add( .add(
Param.optional(List.of("database_name","database")) Param.optional(List.of("database_name", "database"))
.setDescription("the name of the database to use. Defaults to 'baselines'") .setDescription("the name of the database to use.")
) )
.asReadOnly(); .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.client.MilvusServiceClient;
import io.milvus.common.clientenum.ConsistencyLevelEnum; import io.milvus.common.clientenum.ConsistencyLevelEnum;
import io.milvus.grpc.DataType; import io.milvus.grpc.DataType;
import io.milvus.param.collection.CollectionSchemaParam;
import io.milvus.param.collection.CreateCollectionParam; import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType; import io.milvus.param.collection.FieldType;
import io.nosqlbench.adapter.milvus.MilvusDriverAdapter; import io.nosqlbench.adapter.milvus.MilvusDriverAdapter;
@ -70,11 +71,10 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
CreateCollectionParam.Builder::withDatabaseName); CreateCollectionParam.Builder::withDatabaseName);
List<FieldType> fieldTypes = buildFieldTypesStruct( List<FieldType> fieldTypes = buildFieldTypesStruct(
op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey), op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey)
ebF
); );
final LongFunction<CreateCollectionParam.Builder> f = ebF; 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; final LongFunction<CreateCollectionParam.Builder> lastF = ebF;
return l -> lastF.apply(l).build(); 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}. * 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 fieldTypesData The static map of config data from the create collection request
* @param ebF
* @return a list of static field types * @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<>(); List<FieldType> fieldTypes = new ArrayList<>();
fieldTypesData.forEach((name, fieldspec) -> { fieldTypesData.forEach((name, fieldspec) -> {
FieldType.Builder builder = FieldType.newBuilder() FieldType.Builder builder = FieldType.newBuilder()