From 72976d3501f9c6ddef766ee6f40a92e3a4b7d61f Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 12:56:47 -0400 Subject: [PATCH 1/6] KnnSearch op and dispenser first cut --- .../opensearch/OpenSearchOpMapper.java | 1 + .../adapter/opensearch/OpenSearchOpTypes.java | 3 +- .../dispensers/KnnSearchOpDispenser.java | 63 +++++++++++++++++++ .../adapter/opensearch/ops/KnnSearchOp.java | 41 ++++++++++++ 4 files changed, 107 insertions(+), 1 deletion(-) create mode 100644 adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java create mode 100644 adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/KnnSearchOp.java 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..ab9503e57 --- /dev/null +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java @@ -0,0 +1,63 @@ +/* + * 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.json.JsonData; +import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch._types.mapping.*; +import org.opensearch.client.opensearch._types.query_dsl.KnnQuery; +import org.opensearch.client.opensearch._types.query_dsl.MatchAllQuery; +import org.opensearch.client.opensearch._types.query_dsl.Query; +import org.opensearch.client.opensearch.core.IndexRequest; +import org.opensearch.client.opensearch.core.SearchRequest; +import org.opensearch.client.opensearch.indices.CreateIndexRequest; +import org.opensearch.client.opensearch.indices.GetIndexRequest; + +import java.util.Map; +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",float[].class, KnnQuery.Builder::vector); + //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(), Object.class); + } + + +} 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; + } + +} From 591bb27d2b1d1ccc93dc17505876a4689323219e Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 13:09:35 -0400 Subject: [PATCH 2/6] add yaml example of knn_search --- .../opensearch/dispensers/KnnSearchOpDispenser.java | 7 ------- .../src/main/resources/activities/osvectors.yaml | 4 ++++ 2 files changed, 4 insertions(+), 7 deletions(-) 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 index ab9503e57..bf0d524f9 100644 --- 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 @@ -19,18 +19,11 @@ 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.json.JsonData; import org.opensearch.client.opensearch.OpenSearchClient; -import org.opensearch.client.opensearch._types.mapping.*; import org.opensearch.client.opensearch._types.query_dsl.KnnQuery; -import org.opensearch.client.opensearch._types.query_dsl.MatchAllQuery; import org.opensearch.client.opensearch._types.query_dsl.Query; -import org.opensearch.client.opensearch.core.IndexRequest; import org.opensearch.client.opensearch.core.SearchRequest; -import org.opensearch.client.opensearch.indices.CreateIndexRequest; -import org.opensearch.client.opensearch.indices.GetIndexRequest; -import java.util.Map; import java.util.function.LongFunction; public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { diff --git a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml index 9add87b00..b06dc71bb 100644 --- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml +++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml @@ -14,3 +14,7 @@ blocks: create_index: TEMPLATE(indexname,vectors_index) mappings: m1: v1 + 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] From 139c1c3be6fffbd5608fc4dbbaf71cee4fd3d22f Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 13:22:41 -0400 Subject: [PATCH 3/6] add yaml example of knn_search --- .../src/main/resources/activities/osvectors.yaml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml index b06dc71bb..9e37f0462 100644 --- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml +++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml @@ -4,6 +4,7 @@ description: | 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 bindings: @@ -14,6 +15,8 @@ blocks: create_index: TEMPLATE(indexname,vectors_index) mappings: m1: v1 + query: + ops: search: knn_search: TEMPLATE(indexname,vectors_index) k: 100 From 0a4707879c231787447ecb9cf4f3013d99d3db5b Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 15:32:23 -0400 Subject: [PATCH 4/6] list to array conversion --- .../opensearch/dispensers/KnnSearchOpDispenser.java | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) 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 index bf0d524f9..37ff549f2 100644 --- 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 @@ -24,6 +24,7 @@ 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 { @@ -39,7 +40,8 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { 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",float[].class, KnnQuery.Builder::vector); + knnfunc = op.enhanceFuncOptionally(knnfunc, "vector", List.class, this::convertVector);//(b, v) -> b.vector((float[])v.toArray())); + //TODO: Implement the filter query builder here //knnfunc = op.enhanceFuncOptionally(knnfunc, "filter",Query.class, KnnQuery.Builder::filter); @@ -52,5 +54,13 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { return (long l) -> new KnnSearchOp(clientF.apply(l), bfunc.apply(l).build(), Object.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); + } + } From b7b790766ecf7f9012a3c86b44c332b2fc85426c Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 15:57:08 -0400 Subject: [PATCH 5/6] added mandatory field field --- .../adapter/opensearch/dispensers/KnnSearchOpDispenser.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index 37ff549f2..b1788c7b5 100644 --- 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 @@ -40,7 +40,8 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { 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);//(b, v) -> b.vector((float[])v.toArray())); + 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); From b16b04f4667f32ccd577bd83cd1699a937a8ff5b Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 17:21:22 -0400 Subject: [PATCH 6/6] adding Doc type for results --- .../dispensers/KnnSearchOpDispenser.java | 19 ++++++++++++++++++- .../main/resources/activities/osvectors.yaml | 1 + 2 files changed, 19 insertions(+), 1 deletion(-) 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 index b1788c7b5..15ae1d6ce 100644 --- 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 @@ -52,7 +52,7 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { .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(), Object.class); + return (long l) -> new KnnSearchOp(clientF.apply(l), bfunc.apply(l).build(), Doc.class); } private KnnQuery.Builder convertVector(KnnQuery.Builder builder, List list) { @@ -63,5 +63,22 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser { 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/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml index 9e37f0462..16dd0356c 100644 --- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml +++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml @@ -21,3 +21,4 @@ blocks: 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