IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2025-02-20 11:38:28 -06:00
Code Issues Packages Projects Releases Wiki Activity

add opensearch converstion util

Browse Source
This commit is contained in:
Jonathan Shook 2024-02-08 17:58:47 -06:00
parent 9d83684565
commit da9e1e326a
3 changed files with 85 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/Utils.java Normal file
View File

@ -0,0 +1,33 @@
/*
* 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;
import io.nosqlbench.adapter.opensearch.pojos.Doc;
import org.opensearch.client.opensearch.core.SearchResponse;
import org.opensearch.client.opensearch.core.search.Hit;
public class Utils {
public static int[] DocHitsToIntIndicesArray(SearchResponse<Doc> response) {
int[] indices = response.hits().hits()
.stream()
.map(Hit::source)
.mapToInt(doc -> Integer.parseInt(doc.getKey()))
.toArray();
return indices;
}
}

30
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/KnnSearchOpDispenser.java
View File

@ -18,6 +18,7 @@ package io.nosqlbench.adapter.opensearch.dispensers;
import io.nosqlbench.adapter.opensearch.OpenSearchAdapter;
import io.nosqlbench.adapter.opensearch.ops.KnnSearchOp;
import io.nosqlbench.adapter.opensearch.pojos.Doc;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.opensearch._types.query_dsl.KnnQuery;
@ -63,33 +64,4 @@ public class KnnSearchOpDispenser extends BaseOpenSearchOpDispenser {
return builder.vector(vector);
}
public static class Doc {
private float[] value;
private String key;
public Doc() {}
public Doc(float[] value, String key) {
this.value = value;
this.key = key;
}
public float[] getValue() {
return value;
}
public void setValue(float[] value) {
this.value = value;
}
@Override
public String toString() {
return "{" + "values=" + value + "}";
}
public String getKey() {
return key;
}
public void setKey(String key) {
this.key = key;
}
}
}

51
adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/pojos/Doc.java Normal file
View File

@ -0,0 +1,51 @@
/*
* 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.pojos;
public class Doc {
private float[] value;
private String key;
public Doc() {
}
public Doc(float[] value, String key) {
this.value = value;
this.key = key;
}
public float[] getValue() {
return value;
}
public void setValue(float[] value) {
this.value = value;
}
@Override
public String toString() {
return "{" + "values=" + value + "}";
}
public String getKey() {
return key;
}
public void setKey(String key) {
this.key = key;
}
}