From 72976d3501f9c6ddef766ee6f40a92e3a4b7d61f Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Thu, 8 Feb 2024 12:56:47 -0400 Subject: [PATCH] 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; + } + +}