mirror of
https://github.com/nosqlbench/nosqlbench.git
synced 2025-02-25 18:55:28 -06:00
merge fixups
This commit is contained in:
commit
5b48695069
@ -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);
|
||||
};
|
||||
|
@ -21,5 +21,6 @@ public enum OpenSearchOpTypes {
|
||||
delete_index,
|
||||
index,
|
||||
update,
|
||||
delete
|
||||
delete,
|
||||
knn_search
|
||||
}
|
||||
|
@ -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<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", 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<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(), 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 + "}";
|
||||
}
|
||||
}
|
||||
|
||||
}
|
@ -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;
|
||||
}
|
||||
|
||||
}
|
@ -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:
|
||||
|
Loading…
Reference in New Issue
Block a user