partial driver work on milvus

This commit is contained in:
Jonathan Shook 2024-03-17 19:51:34 -05:00
parent 94e1304fd2
commit 6f89033076
8 changed files with 106 additions and 113 deletions

View File

@ -28,7 +28,7 @@ import io.nosqlbench.nb.api.labels.NBLabels;
import java.util.function.Function;
@Service(value = DriverAdapter.class, selector = Utils.DRIVER_TYPE)
public class MilvusDriverAdapter extends BaseDriverAdapter<MilvusOp, MilvusSpace> {
public MilvusDriverAdapter(NBComponent parentComponent, NBLabels labels) {

View File

@ -23,7 +23,7 @@ import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.components.core.NBComponent;
import io.nosqlbench.nb.api.labels.NBLabels;
@Service(value = DriverAdapterLoader.class, selector = Utils.DRIVER_TYPE)
@Service(value = DriverAdapterLoader.class, selector = "milvus")
public class MilvusDriverAdapterLoader implements DriverAdapterLoader {
@Override
public MilvusDriverAdapter load(NBComponent parent, NBLabels childLabels) {

View File

@ -49,7 +49,12 @@ public class MilvusOpMapper implements OpMapper<MilvusOp> {
*/
@Override
public OpDispenser<? extends MilvusOp> apply(ParsedOp op) {
TypeAndTarget<MilvusOpTypes, String> typeAndTarget = op.getTypeAndTarget(MilvusOpTypes.class, String.class, "type", "index");
TypeAndTarget<MilvusOpTypes, String> typeAndTarget = op.getTypeAndTarget(
MilvusOpTypes.class,
String.class,
"type",
"target"
);
logger.info(() -> "Using " + typeAndTarget.enumId + " statement form for '" + op.getName());
return switch (typeAndTarget.enumId) {

View File

@ -17,7 +17,6 @@
package io.nosqlbench.adapter.milvus;
public class Utils {
public static final String DRIVER_TYPE = "milvus";
/**
* Mask the digits in the given string with '*'

View File

@ -17,7 +17,9 @@
package io.nosqlbench.adapter.milvus.opdispensers;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.nosqlbench.adapter.milvus.MilvusDriverAdapter;
import io.nosqlbench.adapter.milvus.ops.MilvusCreateCollectionOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
@ -45,9 +47,21 @@ public class MilvusCreateCollectionOpDispenser extends MilvusOpDispenser {
// https://milvus.io/docs/create_collection.md
@Override
public LongFunction<MilvusCreateCollectionOp> createOpFunc(LongFunction<MilvusServiceClient> clientF, ParsedOp op, LongFunction<String> targetF) {
CreateCollectionParam.Builder eb = CreateCollectionParam.newBuilder();
LongFunction<CreateCollectionParam.Builder> builderF = l ->CreateCollectionParam.newBuilder();
// builderF.apply(0).addFieldType(FieldType.newBuilder().withDataType(DataType.BinaryVector).build());
builderF = op.enhanceFuncOptionally(
builderF, "description",
String.class,
CreateCollectionParam.Builder::withDescription
);
LongFunction<CreateCollectionParam.Builder> f =
l -> CreateCollectionParam.newBuilder().withCollectionName(targetF.apply(l));
return l -> new MilvusCreateCollectionOp(clientF.apply(l), f.apply(1).build());
}
}

View File

@ -0,0 +1,58 @@
bindings:
scenarios:
default:
drop: run driver=milvus tags=block:drop threads==UNSET cycles==UNSET
schema: run driver=milvus tags=block:schema threads==UNSET cycles==UNSET
rampup: run driver=milvus tags=block:rampup threads=TEMPLATE(rampup_threads,auto) cycles=TEMPLATE(rampup_cycles,10k)
main: run driver=milvus tags=block:'main-.*' threads=TEMPLATE(rampup_threads,auto) cycles=TEMPLATE(rampup_cycles,10k)
blocks:
# drop:
# ops:
# drop_index:
schema:
ops:
create_collection_op:
create_collection: example_collection
description: "https://milvus.io/api-reference/java/v2.3.x/Collection/createCollection().md"
shards_num: 10
consistency_level: BOUNDED # BOUNDED, SESSION, EVENTUAL
partition_num: 1024 # number of partitions
field_types:
key:
data_type: Varchar
max_length: 1024 # for String only, >0
primary_key: true # only for Int64 and Varchar types
partition_key: true # how does this relate to primary_key?
description: field description
auto_id: false # Generate primary key?
value:
data_type: FloatVector
dimension: 1024 # >0
primary_key: false # only for Int64 and Varchar types
partition_key: false
description: A value within part_key
auto_id: false # Generate primary key?
create_index:
description: "https://milvus.io/api-reference/java/v2.3.x/Index/createIndex().md"
collection_name: "example_collection"
field_name: value
index_name: "example_idx"
index_type: DISKANN
metric_type: COSINE # L2 | IP | COSINE | HAMMING | JACCARD | TANIMOTO
# extra_param: "none"
sync_mode: true # whether to block till indexing is complete
sync_waiting_interval: 1000 # ms for polling interval on sync_mode: true
# rampup:
# ops:
# write:
#
# main-query:
# ops:
# query:

View File

@ -30,113 +30,24 @@ The operations include:
## Examples
```yaml
ops:
# A pinecone query op
query-example:
type: query
index: query_index
# The query vector. Use these fields if only querying a single vector. If querying multiple use the
# query_vectors structure below.
vector: my_array_of_floats
namespace: query_namespace
# The number of results to return for each query.
top_k: int_query_topk
# You can use vector metadata to limit your search. See https://www.pinecone.io/docs/metadata-filtering/
filter:
filterfield: metadata_field
operator: [$lt, $eq, $gt, ...]
comparator: value
# Indicates whether vector values are included in the response.
include_values: boolean
# Indicates whether metadata is included in the response as well as the ids.
include_metadata: boolean
query_vectors:
- id: 1
values: csv_separated_floats
top_k: int_val
namespace: string_val
filter:
filterfield: metadata_field
operator: [$lt, $eq, $gt, ...]
comparator: value
sparse_values:
indices: list_of_ints
values: list_of_floats
- id: 2
values: csv_separated_floats
top_k: int_val
namespace: string_val
filter:
filterfield: metadata_field
operator: [$lt, $eq, $gt, ...]
comparator: value
sparse_values:
indices: list_of_ints
values: list_of_floats
# A delete op
# If specified, the metadata filter here will be used to select the vectors to delete. This is mutually exclusive
# with specifying ids to delete in the ids param or using delete_all=True. delete_all indicates that all vectors
# in the index namespace should be deleted.
delete-example:
type: delete
index: delete_index
namespace: delete_namespace
ids: csv_list_of_vectors_to_delete
deleteall: [true,false]
filter:
filterfield: metadata_field
operator: [$lt, $eq, $gt, ...]
comparator: value
# A describe index stats op. Specify metadata filters to narrow the range of indices described.
describe-index-stats-example:
type: describe-index-stats
index: describe_index
filter:
filterfield: metadata_field
operator: [$lt, $eq, $gt, ...]
comparator: value
# A pinecone fetch op
fetch-example:
fetch: fetch_index
namespace: fetch_namespace
ids: csv_list_of_vectors_to_fetch
# A pinecone update op
update-example:
type: update
index: update_index
id: string_id
values: list_of_floats
namespace: update_namespace
metadata:
key1: val1
key2: val2
key3: val3
sparse_values:
indices: list_of_ints
values: list_of_floats
# A pinecone upsert op
upsert-example:
type: upsert
index: upsert_index
namespace: upsert_namespace
upsert_vectors:
- id: 1
values: csv_separated_floats
sparse_values:
indices: list_of_ints
values: list_of_floats
metadata:
key1: val1
key2: val2
- id: 2
values: csv_separated_floats
sparse_values:
indices: list_of_ints
values: list_of_floats
example_create_collection:
description: "https://milvus.io/api-reference/java/v2.3.x/Collection/createCollection().md"
collection_name: "example_collection""
shards_num: 10
field_types:
field1:
primary_key: true # only for Int64 and Varchar types
description: field description
data_type: Varchar
# Bool, Int8, Int16, Int32, Int64,
# Float, Double, String, Varchar, BinaryVector, FloatVector
type_param:
example_param1: example_pvalue1
dimension: 1024 # >0
max_length: 1024 # for String only, >0
auto_id: false # Generate primary key?
partition_key: true
consistency_level: BOUNDED # BOUNDED, SESSION, EVENTUAL
partition_num: 1024 # number of partitions
```

View File

@ -131,6 +131,12 @@
<version>${revision}</version>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapter-milvus</artifactId>
<version>${revision}</version>
</dependency>
</dependencies>
<build>