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

Merge branch 'milvus_test' of github.com:nosqlbench/nosqlbench into milvus_test

Browse Source
This commit is contained in:
Jonathan Shook 2024-03-26 14:46:51 -05:00
commit b194f5ee02
10 changed files with 70 additions and 67 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusAlterCollectionOpDispenser.java
View File

@ -42,7 +42,8 @@ public class MilvusAlterCollectionOpDispenser extends MilvusBaseOpDispenser<Alte
LongFunction<AlterCollectionParam.Builder> ebF =
l -> AlterCollectionParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF,"ttl", Integer.class, AlterCollectionParam.Builder::withTTL);
ebF = op.enhanceFuncOptionally(ebF, "ttl", Number.class,
(AlterCollectionParam.Builder b, Number n) -> b.withTTL(n.intValue()));
final LongFunction<AlterCollectionParam.Builder> lastF = ebF;
final LongFunction<AlterCollectionParam> collectionParamF = l -> lastF.apply(l).build();
@ -56,6 +57,6 @@ public class MilvusAlterCollectionOpDispenser extends MilvusBaseOpDispenser<Alte
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusAlterCollectionOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusAlterCollectionOp(clientF.apply(l), paramF.apply(l));
}
}

36
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusCreateCollectionOpDispenser.java
View File

@ -39,12 +39,9 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
/**
* Create a new MilvusCreateCollectionOpDispenser subclassed from {@link MilvusBaseOpDispenser}.
*
* @param adapter
* The associated {@link MilvusDriverAdapter}
* @param op
* The {@link ParsedOp} encapsulating the activity for this cycle
* @param targetFunction
* A LongFunction that returns the specified Milvus Index for this Op
* @param adapter The associated {@link MilvusDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param targetFunction A LongFunction that returns the specified Milvus Index for this Op
*/
public MilvusCreateCollectionOpDispenser(MilvusDriverAdapter adapter,
ParsedOp op,
@ -61,10 +58,10 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
LongFunction<CreateCollectionParam.Builder> ebF =
l -> CreateCollectionParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF, "shards_num", Integer.class,
CreateCollectionParam.Builder::withShardsNum);
ebF = op.enhanceFuncOptionally(ebF, "partition_num", Integer.class,
CreateCollectionParam.Builder::withPartitionsNum);
ebF = op.enhanceFuncOptionally(ebF, "shards_num", Number.class,
(CreateCollectionParam.Builder b, Number n) -> b.withShardsNum(n.intValue()));
ebF = op.enhanceFuncOptionally(ebF, "partition_num", Number.class,
(CreateCollectionParam.Builder b, Number n) -> b.withPartitionsNum(n.intValue()));
ebF = op.enhanceFuncOptionally(ebF, "description", String.class,
CreateCollectionParam.Builder::withDescription);
ebF = op.enhanceEnumOptionally(ebF, "consistency_level",
@ -97,8 +94,7 @@ 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 fieldTypesData The static map of config data from the create collection request
* @param ebF
* @return a list of static field types
*/
@ -112,20 +108,20 @@ public class MilvusCreateCollectionOpDispenser extends MilvusBaseOpDispenser<Cre
.ifPresent(builder::withPrimaryKey);
fieldspec.getOptionalStaticValue("auto_id", Boolean.class)
.ifPresent(builder::withAutoID);
fieldspec.getOptionalStaticConfig("max_length", Integer.class)
.ifPresent(builder::withMaxLength);
fieldspec.getOptionalStaticConfig("max_capacity", Integer.class)
.ifPresent(builder::withMaxCapacity);
fieldspec.getOptionalStaticValue(List.of("partition_key","partition"), Boolean.class)
fieldspec.getOptionalStaticConfig("max_length", Number.class)
.ifPresent((Number n) -> builder.withMaxLength(n.intValue()));
fieldspec.getOptionalStaticConfig("max_capacity", Number.class)
.ifPresent((Number n) -> builder.withMaxCapacity(n.intValue()));
fieldspec.getOptionalStaticValue(List.of("partition_key", "partition"), Boolean.class)
.ifPresent(builder::withPartitionKey);
fieldspec.getOptionalStaticValue("dimension", Integer.class)
.ifPresent(builder::withDimension);
fieldspec.getOptionalStaticValue("dimension", Number.class)
.ifPresent((Number n) -> builder.withDimension(n.intValue()));
fieldspec.getOptionalStaticConfig("data_type", String.class)
.map(DataType::valueOf)
.ifPresent(builder::withDataType);
fieldspec.getOptionalStaticConfig("type_params", Map.class)
.ifPresent(builder::withTypeParams);
fieldspec.getOptionalStaticConfig("element_type",String.class)
fieldspec.getOptionalStaticConfig("element_type", String.class)
.map(DataType::valueOf)
.ifPresent(builder::withElementType);

5
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusCreateIndexOpDispenser.java
View File

@ -71,7 +71,10 @@ public class MilvusCreateIndexOpDispenser extends MilvusBaseOpDispenser<CreateIn
}
@Override
public LongFunction<MilvusBaseOp<CreateIndexParam>> createOpFunc(LongFunction<CreateIndexParam> paramF, LongFunction<MilvusServiceClient> clientF, ParsedOp op, LongFunction<String> targetF) {
public LongFunction<MilvusBaseOp<CreateIndexParam>> createOpFunc(
LongFunction<CreateIndexParam> paramF,
LongFunction<MilvusServiceClient> clientF,
ParsedOp op, LongFunction<String> targetF) {
return l -> new MilvusCreateIndexOp(clientF.apply(l), paramF.apply(l));
}
}

12
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusFlushOpDispenser.java
View File

@ -51,12 +51,12 @@ public class MilvusFlushOpDispenser extends MilvusBaseOpDispenser<FlushParam> {
};
LongFunction<FlushParam.Builder> finalEbF = ebF;
ebF = l -> finalEbF.apply(l).withCollectionNames(cnames.apply(l));
ebF = op.enhanceFuncOptionally(ebF, List.of("database_name","database"),String.class,
ebF = op.enhanceFuncOptionally(ebF, List.of("database_name", "database"), String.class,
FlushParam.Builder::withDatabaseName);
ebF = op.enhanceFuncOptionally(ebF, "sync_flush_waiting_interval",Long.class,
FlushParam.Builder::withSyncFlushWaitingInterval);
ebF = op.enhanceFuncOptionally(ebF, "sync_flush_waiting_timeout",Long.class,
FlushParam.Builder::withSyncFlushWaitingTimeout);
ebF = op.enhanceFuncOptionally(ebF, "sync_flush_waiting_interval", Number.class,
(FlushParam.Builder b, Number n) -> b.withSyncFlushWaitingInterval(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "sync_flush_waiting_timeout", Number.class,
(FlushParam.Builder b, Number n) -> b.withSyncFlushWaitingTimeout(n.longValue()));
final LongFunction<FlushParam.Builder> lastF = ebF;
final LongFunction<FlushParam> collectionParamF = l -> lastF.apply(l).build();
@ -70,6 +70,6 @@ public class MilvusFlushOpDispenser extends MilvusBaseOpDispenser<FlushParam> {
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusFlushOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusFlushOp(clientF.apply(l), paramF.apply(l));
}
}

7
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusGetFlushStateOpDispenser.java
View File

@ -52,9 +52,10 @@ public class MilvusGetFlushStateOpDispenser extends MilvusBaseOpDispenser<GetFlu
};
LongFunction<GetFlushStateParam.Builder> finalEbF = ebF;
ebF = l -> finalEbF.apply(l).withSegmentIDs(idsF.apply(l));
ebF = op.enhanceFuncOptionally(ebF,List.of("collection","collection_name"),String.class,
ebF = op.enhanceFuncOptionally(ebF, List.of("collection", "collection_name"), String.class,
GetFlushStateParam.Builder::withCollectionName);
ebF = op.enhanceFuncOptionally(ebF,"flush_ts",Long.class,GetFlushStateParam.Builder::withFlushTs);
ebF = op.enhanceFuncOptionally(ebF, "flush_ts", Number.class,
(GetFlushStateParam.Builder b, Number n) -> b.withFlushTs(n.longValue()));
final LongFunction<GetFlushStateParam.Builder> lastF = ebF;
final LongFunction<GetFlushStateParam> collectionParamF = l -> lastF.apply(l).build();
@ -68,6 +69,6 @@ public class MilvusGetFlushStateOpDispenser extends MilvusBaseOpDispenser<GetFlu
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusGetFlushStateOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusGetFlushStateOp(clientF.apply(l), paramF.apply(l));
}
}

5
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusListBulkInsertTasksOpDispenser.java
View File

@ -41,7 +41,8 @@ public class MilvusListBulkInsertTasksOpDispenser extends MilvusBaseOpDispenser<
) {
LongFunction<ListBulkInsertTasksParam.Builder> ebF =
l -> ListBulkInsertTasksParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF,"limit",Integer.class,ListBulkInsertTasksParam.Builder::withLimit);
ebF = op.enhanceFuncOptionally(ebF, "limit", Number.class,
(ListBulkInsertTasksParam.Builder b, Number n) -> b.withLimit(n.intValue()));
final LongFunction<ListBulkInsertTasksParam.Builder> lastF = ebF;
final LongFunction<ListBulkInsertTasksParam> collectionParamF = l -> lastF.apply(l).build();
@ -55,6 +56,6 @@ public class MilvusListBulkInsertTasksOpDispenser extends MilvusBaseOpDispenser<
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusListBulkInsertTasksOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusListBulkInsertTasksOp(clientF.apply(l), paramF.apply(l));
}
}

21
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusLoadCollectionOpDispenser.java
View File

@ -43,17 +43,18 @@ public class MilvusLoadCollectionOpDispenser extends MilvusBaseOpDispenser<LoadC
LongFunction<LoadCollectionParam.Builder> ebF =
l -> LoadCollectionParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF,List.of("database_name","database"),String.class,
ebF = op.enhanceFuncOptionally(ebF, List.of("database_name", "database"), String.class,
LoadCollectionParam.Builder::withDatabaseName);
ebF = op.enhanceFuncOptionally(ebF,"refresh",Boolean.class,LoadCollectionParam.Builder::withRefresh);
ebF = op.enhanceFuncOptionally(ebF,"sync_load",Boolean.class,LoadCollectionParam.Builder::withSyncLoad);
ebF = op.enhanceFuncOptionally(ebF,"replica_number",Integer.class,LoadCollectionParam.Builder::withReplicaNumber);
ebF = op.enhanceFuncOptionally(ebF,"resource_groups", List.class,LoadCollectionParam.Builder::withResourceGroups);
ebF = op.enhanceFuncOptionally(ebF,"sync_load_waiting_interval",Number.class,
(LoadCollectionParam.Builder b,Number n) -> b.withSyncLoadWaitingInterval(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF,"sync_load_waiting_timeout",Number.class,
(LoadCollectionParam.Builder b,Number n) -> b.withSyncLoadWaitingTimeout(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "refresh", Boolean.class, LoadCollectionParam.Builder::withRefresh);
ebF = op.enhanceFuncOptionally(ebF, "sync_load", Boolean.class, LoadCollectionParam.Builder::withSyncLoad);
ebF = op.enhanceFuncOptionally(ebF, "replica_number", Number.class,
(LoadCollectionParam.Builder b, Number n) -> b.withReplicaNumber(n.intValue()));
ebF = op.enhanceFuncOptionally(ebF, "resource_groups", List.class, LoadCollectionParam.Builder::withResourceGroups);
ebF = op.enhanceFuncOptionally(ebF, "sync_load_waiting_interval", Number.class,
(LoadCollectionParam.Builder b, Number n) -> b.withSyncLoadWaitingInterval(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "sync_load_waiting_timeout", Number.class,
(LoadCollectionParam.Builder b, Number n) -> b.withSyncLoadWaitingTimeout(n.longValue()));
final LongFunction<LoadCollectionParam.Builder> lastF = ebF;
final LongFunction<LoadCollectionParam> collectionParamF = l -> lastF.apply(l).build();
@ -67,6 +68,6 @@ public class MilvusLoadCollectionOpDispenser extends MilvusBaseOpDispenser<LoadC
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusLoadCollectionOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusLoadCollectionOp(clientF.apply(l), paramF.apply(l));
}
}

16
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusLoadPartitionsOpDispenser.java
View File

@ -50,22 +50,24 @@ public class MilvusLoadPartitionsOpDispenser extends MilvusBaseOpDispenser<LoadP
LongFunction<LoadPartitionsParam.Builder> ebF =
l -> LoadPartitionsParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFunc(ebF, List.of("partition_names","partitions"), List.class,
ebF = op.enhanceFunc(ebF, List.of("partition_names", "partitions"), List.class,
LoadPartitionsParam.Builder::withPartitionNames);
ebF = op.enhanceFuncOptionally(
ebF, "resource_groups", List.class,
LoadPartitionsParam.Builder::withResourceGroups
);
ebF = op.enhanceFuncOptionally(
ebF, List.of("database_name","database"), String.class,
ebF, List.of("database_name", "database"), String.class,
LoadPartitionsParam.Builder::withDatabaseName
);
ebF = op.enhanceFuncOptionally(ebF, "refresh", Boolean.class, LoadPartitionsParam.Builder::withRefresh);
ebF = op.enhanceFuncOptionally(ebF, "replica_number", Integer.class, LoadPartitionsParam.Builder::withReplicaNumber);
ebF = op.enhanceFuncOptionally(ebF,"sync_load",Boolean.class,LoadPartitionsParam.Builder::withSyncLoad);
ebF = op.enhanceFuncOptionally(ebF,"sync_load_waiting_interval",Long.class,LoadPartitionsParam.Builder::withSyncLoadWaitingInterval);
ebF = op.enhanceFuncOptionally(ebF,"sync_load_waiting_timeout",Long.class,
LoadPartitionsParam.Builder::withSyncLoadWaitingTimeout);
ebF = op.enhanceFuncOptionally(ebF, "replica_number", Number.class,
(LoadPartitionsParam.Builder b, Number n) -> b.withReplicaNumber(n.intValue()));
ebF = op.enhanceFuncOptionally(ebF, "sync_load", Boolean.class, LoadPartitionsParam.Builder::withSyncLoad);
ebF = op.enhanceFuncOptionally(ebF, "sync_load_waiting_interval", Number.class,
(LoadPartitionsParam.Builder b, Number n) -> b.withSyncLoadWaitingInterval(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "sync_load_waiting_timeout", Number.class,
(LoadPartitionsParam.Builder b, Number n) -> b.withSyncLoadWaitingTimeout(n.longValue()));
final LongFunction<LoadPartitionsParam.Builder> lastF = ebF;
final LongFunction<LoadPartitionsParam> collectionParamF = l -> lastF.apply(l).build();

16
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusQueryOpDispenser.java
View File

@ -44,14 +44,14 @@ public class MilvusQueryOpDispenser extends MilvusBaseOpDispenser<QueryParam> {
LongFunction<QueryParam.Builder> ebF =
l -> QueryParam.newBuilder().withCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF,List.of("partition_names","partitions"), List.class,
ebF = op.enhanceFuncOptionally(ebF, List.of("partition_names", "partitions"), List.class,
QueryParam.Builder::withPartitionNames);
ebF = op.enhanceEnumOptionally(ebF,"consistency_level", ConsistencyLevelEnum.class, QueryParam.Builder::withConsistencyLevel);
ebF = op.enhanceFuncOptionally(ebF,"expr",String.class,QueryParam.Builder::withExpr);
ebF = op.enhanceFuncOptionally(ebF,"limit",Long.class,QueryParam.Builder::withLimit);
ebF = op.enhanceFuncOptionally(ebF,"offset",Long.class,QueryParam.Builder::withOffset);
ebF = op.enhanceFuncOptionally(ebF,"ignore_growing",Boolean.class,QueryParam.Builder::withIgnoreGrowing);
ebF = op.enhanceFuncOptionally(ebF,"out_fields",List.class,QueryParam.Builder::withOutFields);
ebF = op.enhanceEnumOptionally(ebF, "consistency_level", ConsistencyLevelEnum.class, QueryParam.Builder::withConsistencyLevel);
ebF = op.enhanceFuncOptionally(ebF, "expr", String.class, QueryParam.Builder::withExpr);
ebF = op.enhanceFuncOptionally(ebF, "limit", Number.class, (QueryParam.Builder b, Number n) -> b.withLimit(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "offset", Number.class, (QueryParam.Builder b, Number n) -> b.withOffset(n.longValue()));
ebF = op.enhanceFuncOptionally(ebF, "ignore_growing", Boolean.class, QueryParam.Builder::withIgnoreGrowing);
ebF = op.enhanceFuncOptionally(ebF, "out_fields", List.class, QueryParam.Builder::withOutFields);
final LongFunction<QueryParam.Builder> lastF = ebF;
final LongFunction<QueryParam> collectionParamF = l -> lastF.apply(l).build();
@ -65,6 +65,6 @@ public class MilvusQueryOpDispenser extends MilvusBaseOpDispenser<QueryParam> {
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new MilvusQueryOp(clientF.apply(l),paramF.apply(l));
return l -> new MilvusQueryOp(clientF.apply(l), paramF.apply(l));
}
}

10
adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/opdispensers/MilvusSearchOpDispenser.java
View File

@ -24,12 +24,8 @@ import io.nosqlbench.adapter.milvus.MilvusDriverAdapter;
import io.nosqlbench.adapter.milvus.ops.MilvusBaseOp;
import io.nosqlbench.adapter.milvus.ops.MilvusSearchOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.LongFunction;
public class MilvusSearchOpDispenser extends MilvusBaseOpDispenser<SearchParam> {
@ -53,9 +49,11 @@ public class MilvusSearchOpDispenser extends MilvusBaseOpDispenser<SearchParam>
ebF = op.enhanceEnumOptionally(ebF, "consistency_level", ConsistencyLevelEnum.class, SearchParam.Builder::withConsistencyLevel);
ebF = op.enhanceFuncOptionally(ebF, "expr", String.class, SearchParam.Builder::withExpr);
ebF = op.enhanceDefaultFunc(ebF, "top_k", Integer.class, 100, SearchParam.Builder::withTopK);
ebF = op.enhanceDefaultFunc(ebF, "top_k", Number.class, 100,
(SearchParam.Builder b, Number n) -> b.withTopK(n.intValue()));
ebF = op.enhanceEnumOptionally(ebF, "metric_type", MetricType.class, SearchParam.Builder::withMetricType);
ebF = op.enhanceFuncOptionally(ebF, "round_decimal", Integer.class, SearchParam.Builder::withRoundDecimal);
ebF = op.enhanceFuncOptionally(ebF, "round_decimal", Number.class,
(SearchParam.Builder b, Number n) -> b.withRoundDecimal(n.intValue()));
ebF = op.enhanceFuncOptionally(ebF, "ignore_growing", Boolean.class, SearchParam.Builder::withIgnoreGrowing);
ebF = op.enhanceFuncOptionally(ebF, "params", String.class, SearchParam.Builder::withParams);
ebF = op.enhanceFunc(ebF, List.of("vector_field_name", "vector_field"), String.class,