diff --git a/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtils.java b/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtils.java similarity index 50% rename from adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtils.java rename to adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtils.java index 7e89f0323..14ebe6cd0 100644 --- a/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtils.java +++ b/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtils.java @@ -14,8 +14,9 @@ * limitations under the License. */ -package io.nosqlbench.engine.extensions.scriptingutils; +package io.nosqlbench.engine.extensions.vectormath; +import com.google.protobuf.Descriptors; import io.pinecone.proto.QueryResponse; import io.pinecone.proto.ScoredVector; @@ -29,4 +30,23 @@ public class PineconeScriptingUtils { return response.getMatchesList().stream().mapToInt(r -> Integer.parseInt(r.getId())).toArray(); } + public long[] responseIdsToLongArray(QueryResponse response) { + return response.getMatchesList().stream().mapToLong(r -> Long.parseLong(r.getId())).toArray(); + } + + public String[] responseFieldToStringArray(String fieldname, QueryResponse response) { + return response.getMatchesList().stream().map(r -> r.getMetadata().getFieldsMap() + .get(fieldname).getStringValue()).toArray(String[]::new); + } + + public int[] responseFieldToIntArray(String fieldname, QueryResponse response) { + return response.getMatchesList().stream().mapToInt(r -> Integer.parseInt(r.getMetadata() + .getFieldsMap().get(fieldname).getStringValue())).toArray(); + } + + public long[] responseFieldToLongArray(String fieldname, QueryResponse response) { + return response.getMatchesList().stream().mapToLong(r -> Long.parseLong(r.getMetadata() + .getFieldsMap().get(fieldname).getStringValue())).toArray(); + } + } diff --git a/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtilsPluginInfo.java b/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsPluginInfo.java similarity index 96% rename from adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtilsPluginInfo.java rename to adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsPluginInfo.java index 010d51b4e..b5fd1dd8b 100644 --- a/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/scriptingutils/PineconeScriptingUtilsPluginInfo.java +++ b/adapter-pinecone/src/main/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsPluginInfo.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package io.nosqlbench.engine.extensions.scriptingutils; +package io.nosqlbench.engine.extensions.vectormath; import com.codahale.metrics.MetricRegistry; import io.nosqlbench.api.config.LabeledScenarioContext; diff --git a/adapter-pinecone/src/test/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsTest.java b/adapter-pinecone/src/test/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsTest.java new file mode 100644 index 000000000..7bf75b3df --- /dev/null +++ b/adapter-pinecone/src/test/java/io/nosqlbench/engine/extensions/vectormath/PineconeScriptingUtilsTest.java @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2023 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.engine.extensions.vectormath; + +import com.google.protobuf.Struct; +import com.google.protobuf.Value; +import io.pinecone.proto.QueryResponse; +import io.pinecone.proto.ScoredVector; +import org.junit.jupiter.api.Test; + +import java.util.HashMap; +import java.util.Map; + +public class PineconeScriptingUtilsTest { + + + private QueryResponse generateQueryResponse() { + Value four = Value.newBuilder().setStringValue("4").build(); + Value five = Value.newBuilder().setStringValue("5").build(); + Value six = Value.newBuilder().setStringValue("6").build(); + Map metamap = new HashMap<>(); + metamap.put("a",four); + metamap.put("b",five); + metamap.put("c",six); + Struct metadata = Struct.newBuilder().putAllFields(metamap).build(); + return QueryResponse.newBuilder().addMatches( + ScoredVector.newBuilder().setId("1").setMetadata(metadata).build() + ).addMatches( + ScoredVector.newBuilder().setId("2").setMetadata(metadata).build() + ).addMatches( + ScoredVector.newBuilder().setId("3").setMetadata(metadata).build() + ).build(); + } + + @Test + void responseIdsToStringArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + String[] ids = utils.responseIdsToStringArray(response); + assert(ids.length == 3); + assert(ids[0].equals("1")); + assert(ids[1].equals("2")); + assert(ids[2].equals("3")); + } + + @Test + void responseIdsToIntArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + int[] ids = utils.responseIdsToIntArray(response); + assert(ids.length == 3); + assert(ids[0] == 1); + assert(ids[1] == 2); + assert(ids[2] == 3); + } + + @Test + void responseIdsToLongArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + long[] ids = utils.responseIdsToLongArray(response); + assert(ids.length == 3); + assert(ids[0] == 1L); + assert(ids[1] == 2L); + assert(ids[2] == 3L); + } + + @Test + void responseFieldToStringArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + String[] ids = utils.responseFieldToStringArray("a", response); + assert(ids.length == 3); + assert(ids[0].equals("4")); + assert(ids[1].equals("4")); + assert(ids[2].equals("4")); + } + + @Test + void responseFieldToIntArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + int[] ids = utils.responseFieldToIntArray("b", response); + assert(ids.length == 3); + assert(ids[0] == 5); + assert(ids[1] == 5); + assert(ids[2] == 5); + } + + @Test + void responseFieldToLongArrayTest() { + QueryResponse response = generateQueryResponse(); + PineconeScriptingUtils utils = new PineconeScriptingUtils(); + long[] ids = utils.responseFieldToLongArray("c", response); + assert(ids.length == 3); + assert(ids[0] == 6L); + assert(ids[1] == 6L); + assert(ids[2] == 6L); + } + + + + +}