From 8b30d6251ad100ab40e3d668c5b3a5e8ae9762c9 Mon Sep 17 00:00:00 2001 From: Mark Wolters Date: Tue, 16 May 2023 22:43:16 +0000 Subject: [PATCH] unit tests, correct NBConfiguration creation and update to single filter only --- .../adapter/pinecone/PineconeSpace.java | 17 +++++- .../PineconeDeleteOpDispenser.java | 37 ++++++------- ...PineconeDescribeIndexStatsOpDispenser.java | 16 +++++- .../src/main/resources/docs/pinecone.md | 16 +----- .../pinecone/PineconeOpMapperTest.java | 52 ++++++++++++------- 5 files changed, 82 insertions(+), 56 deletions(-) diff --git a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeSpace.java b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeSpace.java index 6541cedd0..ef4fce75e 100644 --- a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeSpace.java +++ b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeSpace.java @@ -3,6 +3,7 @@ package io.nosqlbench.adapter.pinecone; import io.nosqlbench.api.config.standard.ConfigModel; import io.nosqlbench.api.config.standard.NBConfigModel; import io.nosqlbench.api.config.standard.NBConfiguration; +import io.nosqlbench.api.config.standard.Param; import io.pinecone.PineconeClient; import io.pinecone.PineconeClientConfig; import io.pinecone.PineconeConnection; @@ -65,7 +66,21 @@ public class PineconeSpace { } public static NBConfigModel getConfigModel() { - return ConfigModel.of(PineconeSpace.class); + + return ConfigModel.of(PineconeSpace.class) + .add( + Param.required("apiKey",String.class) + .setDescription("the Pinecone API key to use to connect to the database") + ) + .add( + Param.defaultTo("environment","us-east-1-aws") + .setDescription("the environment in which the desired index is running.") + ) + .add( + Param.defaultTo("projectName","default") + .setDescription("the project name associated with the desired index") + ) + .asReadOnly(); } } diff --git a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDeleteOpDispenser.java b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDeleteOpDispenser.java index ae3780aa3..b307e6279 100644 --- a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDeleteOpDispenser.java +++ b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDeleteOpDispenser.java @@ -1,5 +1,7 @@ package io.nosqlbench.adapter.pinecone.opdispensers; +import com.google.protobuf.Struct; +import com.google.protobuf.Value; import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter; import io.nosqlbench.adapter.pinecone.PineconeSpace; import io.nosqlbench.adapter.pinecone.ops.PineconeDeleteOp; @@ -9,9 +11,7 @@ import io.pinecone.proto.DeleteRequest; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import java.util.Arrays; -import java.util.List; -import java.util.Optional; +import java.util.*; import java.util.function.LongFunction; @@ -41,13 +41,6 @@ public class PineconeDeleteOpDispenser extends PineconeOpDispenser { deleteRequestFunc.apply(value)); } - /* - * return DeleteRequest.newBuilder() - * .setNamespace(namespace) - * .addAllIds(Arrays.asList(idsToDelete)) - * .setDeleteAll(false) - * .build(); - */ private LongFunction createDeleteRequestFunction(ParsedOp op) { LongFunction rFunc = l -> DeleteRequest.newBuilder(); @@ -76,22 +69,22 @@ public class PineconeDeleteOpDispenser extends PineconeOpDispenser { rFunc = l -> finalFunc.apply(l).setDeleteAll(af.apply(l)); } - //TODO: Add filters + Optional> filterFunction = op.getAsOptionalFunction("filter", String.class); + if (filterFunction.isPresent()) { + LongFunction finalFunc = rFunc; + LongFunction builtFilter = l -> { + String[] filterFields = filterFunction.get().apply(l).split(" "); + return Struct.newBuilder().putFields(filterFields[0], + Value.newBuilder().setStructValue(Struct.newBuilder().putFields(filterFields[1], + Value.newBuilder().setNumberValue(Integer.valueOf(filterFields[2])).build())) + .build()).build(); + }; + rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l)); + } LongFunction finalRFunc = rFunc; return l -> finalRFunc.apply(l).build(); } -/* private LongFunction> resolveAttributeDefinitionFunction(ParsedOp cmd) { - LongFunction attrsmap = cmd.getAsRequiredFunction("Attributes", Map.class); - return (long l) -> { - List defs = new ArrayList<>(); - attrsmap.apply(l).forEach((k, v) -> { - defs.add(new AttributeDefinition(k.toString(), ScalarAttributeType.valueOf(v.toString()))); - }); - return defs; - }; - }*/ - } diff --git a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDescribeIndexStatsOpDispenser.java b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDescribeIndexStatsOpDispenser.java index ae0aca9e1..9ce922456 100644 --- a/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDescribeIndexStatsOpDispenser.java +++ b/adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDescribeIndexStatsOpDispenser.java @@ -1,5 +1,7 @@ package io.nosqlbench.adapter.pinecone.opdispensers; +import com.google.protobuf.Struct; +import com.google.protobuf.Value; import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter; import io.nosqlbench.adapter.pinecone.PineconeSpace; import io.nosqlbench.adapter.pinecone.ops.PineconeDescribeIndexStatsOp; @@ -10,6 +12,7 @@ import jakarta.ws.rs.NotSupportedException; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import java.util.Optional; import java.util.function.LongFunction; public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser { @@ -34,7 +37,18 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser { private LongFunction createDescribeIndexStatsRequestFunction(ParsedOp op) { LongFunction rFunc = l -> DescribeIndexStatsRequest.newBuilder(); - //TODO: Add filters + Optional> filterFunction = op.getAsOptionalFunction("filter", String.class); + if (filterFunction.isPresent()) { + LongFunction finalFunc = rFunc; + LongFunction builtFilter = l -> { + String[] filterFields = filterFunction.get().apply(l).split(" "); + return Struct.newBuilder().putFields(filterFields[0], + Value.newBuilder().setStructValue(Struct.newBuilder().putFields(filterFields[1], + Value.newBuilder().setNumberValue(Integer.valueOf(filterFields[2])).build())) + .build()).build(); + }; + rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l)); + } LongFunction finalRFunc = rFunc; return l -> finalRFunc.apply(l).build(); } diff --git a/adapter-pinecone/src/main/resources/docs/pinecone.md b/adapter-pinecone/src/main/resources/docs/pinecone.md index c8308fae4..3f0bf320f 100644 --- a/adapter-pinecone/src/main/resources/docs/pinecone.md +++ b/adapter-pinecone/src/main/resources/docs/pinecone.md @@ -64,25 +64,13 @@ ops: namespace: delete_namespace ids: csv_list_of_vectors_to_delete deleteall: [true,false] - filters: - - filter_field: delete_filter_field - operator: [$lt, $gt, $eq, ...] - comparator: query_compval - - filter_field: delete_filter_field - operator: [$lt, $gt, $eq, ...] - comparator: query_compval + filter: # A describe index stats op. Specify metadata filters to narrow the range of indices described. describe-index-stats-example: type: describe-index-stats index: describe_index - filters: - - filter_field: delete_filter_field - operator: [$lt, $gt, $eq, ...] - comparator: query_compval - - filter_field: delete_filter_field - operator: [$lt, $gt, $eq, ...] - comparator: query_compval + filter: # A pinecone fetch op fetch-example: diff --git a/adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java b/adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java index 8c69e0e62..4ae232326 100644 --- a/adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java +++ b/adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java @@ -1,5 +1,7 @@ package io.nosqlbench.adapter.pinecone; +import io.nosqlbench.adapter.pinecone.opdispensers.PineconeDeleteOpDispenser; +import io.nosqlbench.adapter.pinecone.opdispensers.PineconeQueryOpDispenser; import io.nosqlbench.adapter.pinecone.ops.PineconeOp; import io.nosqlbench.api.config.standard.NBConfiguration; import io.nosqlbench.engine.api.activityconfig.OpsLoader; @@ -26,7 +28,10 @@ public class PineconeOpMapperTest { @BeforeAll public static void initializeTestMapper() { - cfg = PineconeSpace.getConfigModel().apply(Map.of()); + Map configMap = Map.of("apiKey","2f55b2f0-670f-4c51-9073-4d37142b761a", + "environment","us-east-1-aws", + "projectName","default"); + cfg = PineconeSpace.getConfigModel().apply(configMap); adapter = new PineconeDriverAdapter(); adapter.applyConfig(cfg); DriverSpaceCache cache = adapter.getSpaceCache(); @@ -36,34 +41,45 @@ public class PineconeOpMapperTest { private static ParsedOp parsedOpFor(String yaml) { OpsDocList docs = OpsLoader.loadString(yaml, OpTemplateFormat.yaml, Map.of(), null); OpTemplate opTemplate = docs.getOps().get(0); - ParsedOp parsedOp = new ParsedOp(opTemplate, cfg, List.of(adapter.getPreprocessor())); - return parsedOp; + return new ParsedOp(opTemplate, cfg, List.of(adapter.getPreprocessor())); } @Test public void testQueryOpDispenserSimple() { ParsedOp pop = parsedOpFor(""" - ops: - query-op1: - query: "test-index" - vector: "1.0,2.0,3.0" - namespace: "test-namespace" - top_k: 10 - filters: - - filter_field: "value" - operator: "$lt" - comparator: 2 - include_values: true - include_metadata: true - + ops: + op1: + type: "query" + index: "test-index" + vector: "1.0,2.0,3.0" + namespace: "test-namespace" + top_k: 10 + filters: + - "value $lt 2" + - "value $gt 10" + include_values: true + include_metadata: true """); OpDispenser dispenser = mapper.apply(pop); - //assertions go here... + assert(dispenser instanceof PineconeQueryOpDispenser); } @Test public void testDeleteOpDispenser() { - + ParsedOp pop = parsedOpFor(""" + ops: + op1: + type: "delete" + index: "test-index" + ids: "1.0,2.0,3.0" + namespace: "test-namespace" + deleteall: true + filters: + - "value $lt 2" + - "value $gt 10" + """); + OpDispenser dispenser = mapper.apply(pop); + assert(dispenser instanceof PineconeDeleteOpDispenser); } @Test