IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2024-12-23 15:40:44 -06:00
Code Issues Packages Projects Releases Wiki Activity

KnnSearch op and dispenser first cut

Browse Source
This commit is contained in:
Mark Wolters 2024-02-08 12:56:47 -04:00
parent 99a1b287e2
commit 72976d3501
4 changed files with 107 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpMapper.java
View File

@ -44,6 +44,7 @@ public class OpenSearchOpMapper implements OpMapper<Op> {
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);
};

3
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchOpTypes.java
View File

@ -21,5 +21,6 @@ public enum OpenSearchOpTypes {
delete_index,
index,
update,
delete
delete,
knn_search
}

63
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java Normal file
View File

@ -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<String> targetF;
public KnnSearchOpDispenser(OpenSearchAdapter adapter, ParsedOp op, LongFunction<String> targetF) {
super(adapter, op);
this.targetF = targetF;
}
@Override
public LongFunction<KnnSearchOp> createOpFunc(LongFunction<OpenSearchClient> clientF, ParsedOp op) {
LongFunction<KnnQuery.Builder> 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<KnnQuery.Builder> finalKnnfunc = knnfunc;
LongFunction<SearchRequest.Builder> 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);
}
}

41
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/KnnSearchOp.java Normal file
View File

@ -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;
}
}