diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpMapper.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpMapper.java index ae218e10c..1d8e828bf 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpMapper.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpMapper.java @@ -44,6 +44,7 @@ public class OpenSearchOpMapper implements OpMapper { case index -> new IndexOpDispenser(adapter,op, typeAndTarget.targetFunction); case update -> new UpdateOpDispenser(adapter,op, typeAndTarget.targetFunction); case delete -> new DeleteOpDispenser(adapter,op, typeAndTarget.targetFunction); + case knn_search -> new KnnSearchOpDispenser(adapter,op, typeAndTarget.targetFunction); default -> throw new RuntimeException("Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " + "mapping parsed op " + op); }; diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpTypes.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpTypes.java index a926bedce..a733cc495 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpTypes.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpTypes.java @@ -21,5 +21,6 @@ public enum OpenSearchOpTypes { delete_index, index, update, - delete + delete, + knn_search } diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java new file mode 100644 index 000000000..15ae1d6ce --- /dev/null +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java @@ -0,0 +1,84 @@ +/* + * Copyright (c) 2024 nosqlbench + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.nosqlbench.adapter.opensearch.dispensers; + +import io.nosqlbench.adapter.opensearch.OpenSearchAdapter; +import io.nosqlbench.adapter.opensearch.ops.KnnSearchOp; +import io.nosqlbench.adapters.api.templating.ParsedOp; +import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch._types.query_dsl.KnnQuery; +import org.opensearch.client.opensearch._types.query_dsl.Query; +import org.opensearch.client.opensearch.core.SearchRequest; + +import java.util.List; +import java.util.function.LongFunction; + +public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { + + private final LongFunction targetF; + + public KnnSearchOpDispenser(OpenSearchAdapter adapter, ParsedOp op, LongFunction targetF) { + super(adapter, op); + this.targetF = targetF; + } + + @Override + public LongFunction createOpFunc(LongFunction clientF, ParsedOp op) { + LongFunction knnfunc = l -> new KnnQuery.Builder(); + knnfunc = op.enhanceFuncOptionally(knnfunc, "k",Integer.class, KnnQuery.Builder::k); + knnfunc = op.enhanceFuncOptionally(knnfunc, "vector", List.class, this::convertVector); + knnfunc = op.enhanceFuncOptionally(knnfunc, "field",String.class, KnnQuery.Builder::field); + + //TODO: Implement the filter query builder here + //knnfunc = op.enhanceFuncOptionally(knnfunc, "filter",Query.class, KnnQuery.Builder::filter); + + LongFunction finalKnnfunc = knnfunc; + LongFunction bfunc = + l -> new SearchRequest.Builder() + .index(targetF.apply(1)) + .query(new Query.Builder().knn(finalKnnfunc.apply(l).build()).build()); + + return (long l) -> new KnnSearchOp(clientF.apply(l), bfunc.apply(l).build(), Doc.class); + } + + private KnnQuery.Builder convertVector(KnnQuery.Builder builder, List list) { + float[] vector = new float[list.size()]; + for (int i = 0; i < list.size(); i++) { + vector[i] = (float) list.get(i); + } + return builder.vector(vector); + } + + public static class Doc { + private float[] values; + public Doc() {} + public Doc(float[] values) { + this.values = values; + } + public float[] getValues() { + return values; + } + public void setValues(float[] values) { + this.values = values; + } + @Override + public String toString() { + return "{" + "values=" + values + "}"; + } + } + +} diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/KnnSearchOp.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/KnnSearchOp.java new file mode 100644 index 000000000..16ca011e0 --- /dev/null +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/KnnSearchOp.java @@ -0,0 +1,41 @@ +/* + * Copyright (c) 2024 nosqlbench + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package io.nosqlbench.adapter.opensearch.ops; + +import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch.core.SearchRequest; +import org.opensearch.client.opensearch.core.SearchResponse; +import org.opensearch.client.opensearch.indices.GetIndexRequest; +import org.opensearch.client.opensearch.indices.GetIndexResponse; + +public class KnnSearchOp extends BaseOpenSearchOp { + private final SearchRequest rq; + private final Class doctype; + + public KnnSearchOp(OpenSearchClient client, SearchRequest rq, Class doctype) { + super(client); + this.rq = rq; + this.doctype = doctype; + } + + @Override + public Object applyOp(long value) throws Exception { + SearchResponse response = client.search(rq, doctype); + return response; + } + +} diff --git a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml index 011f8c035..bcb37126a 100644 --- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml +++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml @@ -10,6 +10,8 @@ params: instrument: true scenarios: + schema: run driver=opensearch tags=block:create threads===1 cycles===undef + query: run driver=opensearch tags=block:query threads=TEMPLATE(search_threads,10) cycles===undef # schema: run tags=block:schema threads===1 cycles===undef # rampup: run tags=block:rampup threads=auto cycles=TEMPLATE(rampup_cycles,10) @@ -44,6 +46,13 @@ blocks: create_index: TEMPLATE(indexname) mappings: m1: v1 + query: + ops: + search: + knn_search: TEMPLATE(indexname,vectors_index) + k: 100 + vector: [1f,2f,3f,4f,5f,6f,7f,8f,9f,10f,11f,12f,13f,14f,15f,16f,17f,18f,19f,20f,21f,22f,23f,24f,25f] + field: value bulkrampup: ops: bulk_index: