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

fleshing out opMapper and OpDispensers

Browse Source
This commit is contained in:
Mark Wolters 2023-05-09 20:16:02 +00:00 committed by Madhavan
parent c6b550d53c
commit 15a13f7bcb
13 changed files with 265 additions and 39 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeOpMapper.java
View File

@ -1,13 +1,51 @@
package io.nosqlbench.adapter.pinecone;
import io.nosqlbench.adapter.pinecone.opdispensers.*;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOpTypes;
import io.nosqlbench.engine.api.activityimpl.OpDispenser;
import io.nosqlbench.engine.api.activityimpl.OpMapper;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverSpaceCache;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.nosqlbench.engine.api.templating.TypeAndTarget;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.function.LongFunction;
public class PineconeOpMapper implements OpMapper<PineconeOp> {
private static final Logger LOGGER = LogManager.getLogger(PineconeOpMapper.class);
private final PineconeDriverAdapter adapter;
private final DriverSpaceCache<? extends PineconeSpace> spaceCache;
public PineconeOpMapper(PineconeDriverAdapter adapter, DriverSpaceCache<? extends PineconeSpace> spaceCache) {
this.adapter = adapter;
this.spaceCache = spaceCache;
}
@Override
public OpDispenser<? extends PineconeOp> apply(ParsedOp op) {
return null;
LongFunction<String> spaceFunction = op.getAsFunctionOr("space", "default");
LongFunction<PineconeSpace> pcFunction = l -> spaceCache.get(spaceFunction.apply(l));
TypeAndTarget<PineconeOpTypes, String> opType = op.getTypeAndTarget(PineconeOpTypes.class, String.class, "type", "stmt");
LOGGER.info(() -> "Using " + opType.enumId + " statement form for '" + op.getName());
return switch (opType.enumId) {
case query ->
new PineconeQueryOpDispenser(adapter, op, pcFunction, opType.targetFunction);
case update ->
new PineconeUpdateOpDispenser(adapter, op, pcFunction, opType.targetFunction);
case upsert ->
new PineconeUpsertOpDispenser(adapter, op, pcFunction, opType.targetFunction);
case delete ->
new PineconeDeleteOpDispenser(adapter, op, pcFunction, opType.targetFunction);
case describeindexstats ->
new PineconeDescribeIndexStatsOpDispenser(adapter, op, pcFunction, opType.targetFunction);
case fetch ->
new PineconeFetchOpDispenser(adapter, op, pcFunction, opType.targetFunction);
};
}
}

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

@ -1,5 +1,6 @@
package io.nosqlbench.adapter.pinecone;
import io.nosqlbench.api.config.standard.NBConfiguration;
import io.pinecone.PineconeClient;
import io.pinecone.PineconeClientConfig;
import io.pinecone.PineconeConnection;
@ -27,10 +28,10 @@ public class PineconeSpace {
*/
private Map<String,PineconeConnection> connections = new HashMap<String,PineconeConnection>();
public PineconeSpace(String apiKey, String environment, String projectName, String name) {
this.apiKey = apiKey;
this.environment = environment;
this.projectName = projectName;
public PineconeSpace(String name, NBConfiguration cfg) {
this.apiKey = cfg.get("apiKey");
this.environment = cfg.get("environment");
this.projectName = cfg.get("projectName");
this.name = name;
config = new PineconeClientConfig()

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

@ -1,17 +1,47 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
import io.nosqlbench.adapter.pinecone.PineconeSpace;
import io.nosqlbench.adapter.pinecone.ops.PineconeDeleteOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.DeleteRequest;
import java.util.Arrays;
import java.util.function.LongFunction;
public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
protected PineconeDeleteOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private DeleteRequest request;
private PineconeConnection connection;
public PineconeDeleteOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("delete");
connection = pcFunction.apply(0).getConnection(indexName);
request = createDeleteRequest();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeDeleteOp(connection, request);
}
private DeleteRequest createDeleteRequest() {
// TODO: How do I pull these from the ParsedOp?
String[] idsToDelete = {"v2"};
String namespace = "ns";
return DeleteRequest.newBuilder()
.setNamespace(namespace)
.addAllIds(Arrays.asList(idsToDelete))
.setDeleteAll(false)
.build();
}
}

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

@ -1,17 +1,36 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
import io.nosqlbench.adapter.pinecone.PineconeSpace;
import io.nosqlbench.adapter.pinecone.ops.PineconeDescribeIndexStatsOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.DescribeIndexStatsRequest;
import java.util.function.LongFunction;
public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
protected PineconeDescribeIndexStatsOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private DescribeIndexStatsRequest request;
private PineconeConnection connection;
public PineconeDescribeIndexStatsOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("describeIndexStats");
connection = pcFunction.apply(0).getConnection(indexName);
request = createDescribeIndexStatsRequest();
}
private DescribeIndexStatsRequest createDescribeIndexStatsRequest() {
return DescribeIndexStatsRequest.newBuilder().build();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeDescribeIndexStatsOp(connection, request);
}
}

30
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeFetchOpDispenser.java
View File

@ -1,17 +1,39 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
import io.nosqlbench.adapter.pinecone.PineconeSpace;
import io.nosqlbench.adapter.pinecone.ops.PineconeFetchOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.FetchRequest;
import java.util.Arrays;
import java.util.List;
import java.util.function.LongFunction;
public class PineconeFetchOpDispenser extends PineconeOpDispenser {
protected PineconeFetchOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private FetchRequest request;
private PineconeConnection connection;
public PineconeFetchOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("fetch");
connection = pcFunction.apply(0).getConnection(indexName);
request = createFetchRequest();
}
private FetchRequest createFetchRequest() {
List<String> ids = Arrays.asList("v1","v2");
return FetchRequest.newBuilder().addAllIds(ids).setNamespace("default-namespace").build();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeFetchOp(connection, request);
}
}

19
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeOpDispenser.java
View File

@ -21,14 +21,21 @@ import io.nosqlbench.adapter.pinecone.PineconeSpace;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
public class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
protected PineconeOpDispenser(PineconeDriverAdapter adapter, ParsedOp op) {
import java.util.function.LongFunction;
public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
protected LongFunction<PineconeSpace> pcFunction;
protected LongFunction<String> targetFunction;
protected PineconeOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op);
this.pcFunction = pcFunction;
this.targetFunction = targetFunction;
}
@Override
public PineconeOp apply(long value) {
return null;
}
}

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

@ -1,17 +1,59 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import com.google.common.primitives.Floats;
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.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapter.pinecone.ops.PineconeQueryOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.QueryRequest;
import io.pinecone.proto.QueryVector;
import java.util.function.LongFunction;
public class PineconeQueryOpDispenser extends PineconeOpDispenser {
protected PineconeQueryOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private QueryRequest request;
private PineconeConnection connection;
public PineconeQueryOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("query");
connection = pcFunction.apply(0).getConnection(indexName);
request = createQueryRequest();
}
private QueryRequest createQueryRequest() {
float[] rawVector = {1.0F, 2.0F, 3.0F};
QueryVector queryVector = QueryVector.newBuilder()
.addAllValues(Floats.asList(rawVector))
.setFilter(Struct.newBuilder()
.putFields("some_field", Value.newBuilder()
.setStructValue(Struct.newBuilder()
.putFields("$lt", Value.newBuilder()
.setNumberValue(3)
.build()))
.build())
.build())
.setNamespace("default-namespace")
.build();
return QueryRequest.newBuilder()
.addQueries(queryVector)
.setNamespace("default-namespace")
.setTopK(2)
.setIncludeMetadata(true)
.build();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeQueryOp(connection, request);
}
}

27
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeUpdateOpDispenser.java
View File

@ -1,17 +1,36 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
import io.nosqlbench.adapter.pinecone.PineconeSpace;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapter.pinecone.ops.PineconeUpdateOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.UpdateRequest;
import java.util.function.LongFunction;
public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
protected PineconeUpdateOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private UpdateRequest request;
private PineconeConnection connection;
public PineconeUpdateOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("update");
connection = pcFunction.apply(0).getConnection(indexName);
request = createUpdateRequest();
}
private UpdateRequest createUpdateRequest() {
return UpdateRequest.newBuilder().build();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeUpdateOp(connection, request);
}
}

51
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/opdispensers/PineconeUpsertOpDispenser.java
View File

@ -1,17 +1,60 @@
package io.nosqlbench.adapter.pinecone.opdispensers;
import com.google.common.primitives.Floats;
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.PineconeOp;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapter.pinecone.ops.PineconeUpsertOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.UpsertRequest;
import io.pinecone.proto.Vector;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.LongFunction;
public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
protected PineconeUpsertOpDispenser(DriverAdapter<PineconeOp, PineconeSpace> adapter, ParsedOp op) {
super(adapter, op);
private UpsertRequest request;
private PineconeConnection connection;
public PineconeUpsertOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("upsert");
connection = pcFunction.apply(0).getConnection(indexName);
request = createUpsertRequest();
}
private UpsertRequest createUpsertRequest() {
float[][] upsertData = {{1.0F, 2.0F, 3.0F}, {4.0F, 5.0F, 6.0F}, {7.0F, 8.0F, 9.0F}};
List<String> upsertIds = Arrays.asList("v1", "v2", "v3");
List<Vector> upsertVectors = new ArrayList<>();
for (int i = 0; i < upsertData.length; i++) {
upsertVectors.add(Vector.newBuilder()
.addAllValues(Floats.asList(upsertData[i]))
.setMetadata(Struct.newBuilder()
.putFields("some_field", Value.newBuilder().setNumberValue(i).build())
.build())
.setId(upsertIds.get(i))
.build());
}
return UpsertRequest.newBuilder()
.addAllVectors(upsertVectors)
.setNamespace("default-namespace")
.build();
}
@Override
public PineconeOp apply(long value) {
return null;
return new PineconeUpsertOp(connection, request);
}
}

2
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/ops/PineconeDeleteOp.java
View File

@ -21,7 +21,7 @@ public class PineconeDeleteOp extends PineconeOp {
public void run() {
try {
DeleteResponse response = connection.getBlockingStub().delete(request);
// Do soemething with the response...
LOGGER.info(response.toString());
} catch (Exception e) {
LOGGER.error("Exception %s caught trying to do delete", e.getMessage());
LOGGER.error(e.getStackTrace());

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

@ -18,5 +18,9 @@ package io.nosqlbench.adapter.pinecone.ops;
public enum PineconeOpTypes {
query,
update
update,
upsert,
delete,
describeindexstats,
fetch
}

3
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/ops/PineconeQueryOp.java
View File

@ -21,7 +21,8 @@ public class PineconeQueryOp extends PineconeOp {
public void run() {
try {
QueryResponse response = connection.getBlockingStub().query(request);
// Do soemething with the response...
LOGGER.info("got query result ids: "
+ response.getResultsList().get(0).getMatchesList());
} catch (Exception e) {
LOGGER.error("Exception %s caught trying to do Query", e.getMessage());
LOGGER.error(e.getStackTrace());

2
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/ops/PineconeUpsertOp.java
View File

@ -21,7 +21,7 @@ public class PineconeUpsertOp extends PineconeOp {
public void run() {
try {
UpsertResponse response = connection.getBlockingStub().upsert(request);
// Do soemething with the response...
LOGGER.info("Put " + response.getUpsertedCount() + " vectors into the index");
} catch (Exception e) {
LOGGER.error("Exception %s caught trying to do upsert", e.getMessage());
LOGGER.error(e.getStackTrace());