merge fixups

This commit is contained in:
Jonathan Shook 2024-02-08 16:29:40 -06:00
commit 5b48695069
5 changed files with 137 additions and 1 deletions

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

View File

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

View File

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

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

View File

@ -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: