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

fixing a couple problems, adding tests

Browse Source
This commit is contained in:
Mark Wolters 2023-05-12 21:39:20 +00:00 committed by Madhavan
parent 2d17ab0ab2
commit de7ef54cf2
5 changed files with 142 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeSpace.java
View File

@ -1,5 +1,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.pinecone.PineconeClient;
import io.pinecone.PineconeClientConfig;
@ -62,4 +64,8 @@ public class PineconeSpace {
return connection;
}
public static NBConfigModel getConfigModel() {
return ConfigModel.of(PineconeSpace.class);
}
}

13
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDeleteOpDispenser.java
View File

@ -76,9 +76,22 @@ public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
rFunc = l -> finalFunc.apply(l).setDeleteAll(af.apply(l));
}
//TODO: Add filters
LongFunction<DeleteRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
/* private LongFunction<Collection<AttributeDefinition>> resolveAttributeDefinitionFunction(ParsedOp cmd) {
LongFunction<? extends Map> attrsmap = cmd.getAsRequiredFunction("Attributes", Map.class);
return (long l) -> {
List<AttributeDefinition> defs = new ArrayList<>();
attrsmap.apply(l).forEach((k, v) -> {
defs.add(new AttributeDefinition(k.toString(), ScalarAttributeType.valueOf(v.toString())));
});
return defs;
};
}*/
}

5
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeDescribeIndexStatsOpDispenser.java
View File

@ -33,7 +33,10 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
}
private LongFunction<DescribeIndexStatsRequest> createDescribeIndexStatsRequestFunction(ParsedOp op) {
throw new NotSupportedException("Pinecone Describe Index Stats Op not yet supported");
LongFunction<DescribeIndexStatsRequest.Builder> rFunc = l -> DescribeIndexStatsRequest.newBuilder();
//TODO: Add filters
LongFunction<DescribeIndexStatsRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
@Override

47
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeQueryOpDispenser.java
View File

@ -31,7 +31,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
queryRequestFunc = createQueryRequestFunc(op, createQueryVectorFunc(op));
queryRequestFunc = createQueryRequestFunc(op);
}
/*
@ -57,7 +57,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
.build();
}
*/
private LongFunction<QueryRequest> createQueryRequestFunc(ParsedOp op, LongFunction<QueryVector> queryVectorFunc) {
private LongFunction<QueryRequest> createQueryRequestFunc(ParsedOp op) {
LongFunction<QueryRequest.Builder> rFunc = l -> QueryRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
@ -74,26 +74,24 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
rFunc = l -> finalFunc.apply(l).setTopK(af.apply(l));
}
Optional<LongFunction<Boolean>> mFunc = op.getAsOptionalFunction("includemetadata", Boolean.class);
Optional<LongFunction<Boolean>> mFunc = op.getAsOptionalFunction("include_metadata", Boolean.class);
if (mFunc.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<Boolean> af = mFunc.get();
rFunc = l -> finalFunc.apply(l).setIncludeMetadata(af.apply(l));
}
LongFunction<QueryRequest.Builder> returnFunc = rFunc;
rFunc = l -> returnFunc.apply(l).addQueries(queryVectorFunc.apply(l));
LongFunction<QueryRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
Optional<LongFunction<Boolean>> ivFunc = op.getAsOptionalFunction("include_values", Boolean.class);
if (ivFunc.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<Boolean> af = ivFunc.get();
rFunc = l -> finalFunc.apply(l).setIncludeValues(af.apply(l));
}
private LongFunction<QueryVector> createQueryVectorFunc(ParsedOp op) {
LongFunction<QueryVector.Builder> vFunc = l -> QueryVector.newBuilder();
Optional<LongFunction<String>> qFunc = op.getAsOptionalFunction("query", String.class);
if (qFunc.isPresent()) {
LongFunction<QueryVector.Builder> finalFunc = vFunc;
LongFunction<String> af = qFunc.get();
Optional<LongFunction<String>> vFunc = op.getAsOptionalFunction("vector", String.class);
if (vFunc.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = vFunc.get();
LongFunction<ArrayList<Float>> alf = l -> {
String[] vals = af.apply(l).split(",");
ArrayList<Float> fVals = new ArrayList<Float>();
@ -102,17 +100,18 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
}
return fVals;
};
vFunc = l -> finalFunc.apply(l).addAllValues(alf.apply(l));
rFunc = l -> finalFunc.apply(l).addAllVector(alf.apply(l));
}
//TODO: Add filters
//TODO: IF the above values are not populated we need to create and add the query vectors
//LongFunction<QueryRequest.Builder> returnFunc = rFunc;
//rFunc = l -> returnFunc.apply(l).addQueries(queryVectorFunc.apply(l));
LongFunction<QueryRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<QueryVector.Builder> finalFunc = vFunc;
LongFunction<String> af = nFunc.get();
vFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
//TODO: Add in optional filters
private LongFunction<QueryVector> createQueryVectorFunc(ParsedOp op) {
LongFunction<QueryVector.Builder> vFunc = l -> QueryVector.newBuilder();
LongFunction<QueryVector.Builder> finalVFunc = vFunc;
return l -> finalVFunc.apply(l).build();
}

96
adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java Normal file
View File

@ -0,0 +1,96 @@
package io.nosqlbench.adapter.pinecone;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.api.config.standard.NBConfiguration;
import io.nosqlbench.engine.api.activityconfig.OpsLoader;
import io.nosqlbench.engine.api.activityconfig.yaml.OpTemplate;
import io.nosqlbench.engine.api.activityconfig.yaml.OpTemplateFormat;
import io.nosqlbench.engine.api.activityconfig.yaml.OpsDocList;
import io.nosqlbench.engine.api.activityimpl.OpDispenser;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverSpaceCache;
import io.nosqlbench.engine.api.templating.ParsedOp;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import java.util.List;
import java.util.Map;
public class PineconeOpMapperTest {
private final static Logger logger = LogManager.getLogger(PineconeOpMapperTest.class);
static NBConfiguration cfg;
static PineconeDriverAdapter adapter;
static PineconeOpMapper mapper;
@BeforeAll
public static void initializeTestMapper() {
cfg = PineconeSpace.getConfigModel().apply(Map.of());
adapter = new PineconeDriverAdapter();
adapter.applyConfig(cfg);
DriverSpaceCache<? extends PineconeSpace> cache = adapter.getSpaceCache();
mapper = new PineconeOpMapper(adapter, cache, cfg);
}
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;
}
@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
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
//assertions go here...
}
@Test
public void testDeleteOpDispenser() {
}
@Test
public void testDescribeIndexStatsOpDispenser() {
}
@Test
public void testFetchOpDispenser() {
}
@Test
public void testUpdateOpDispenser() {
}
@Test
public void testUpsertOpDispenser() {
}
@Test
public void testQueryOpDispenserComplex() {
}
}