fixing index name identification

This commit is contained in:
Mark Wolters 2023-05-11 14:01:02 +00:00 committed by Madhavan
parent a38047f469
commit 3c9cae7f90
8 changed files with 57 additions and 47 deletions

View File

@ -1,15 +1,21 @@
package io.nosqlbench.adapter.pinecone;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.api.config.standard.NBConfiguration;
import io.nosqlbench.engine.api.activityimpl.OpMapper;
import io.nosqlbench.engine.api.activityimpl.uniform.BaseDriverAdapter;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverSpaceCache;
import io.nosqlbench.nb.annotations.Service;
@Service(value = DriverAdapter.class, selector = "pinecone")
public class PineconeDriverAdapter extends BaseDriverAdapter<PineconeOp, PineconeSpace> {
@Override
public OpMapper<PineconeOp> getOpMapper() {
return null;
DriverSpaceCache<? extends PineconeSpace> spaceCache = getSpaceCache();
NBConfiguration adapterConfig = getConfiguration();
return new PineconeOpMapper(this, spaceCache, adapterConfig);
}
}

View File

@ -23,7 +23,6 @@ import java.util.function.LongFunction;
*/
public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
private final String indexName;
private final LongFunction<DeleteRequest> deleteRequestFunc;
public PineconeDeleteOpDispenser(PineconeDriverAdapter adapter,
@ -32,13 +31,14 @@ public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
indexName = op.getAsRequiredFunction("query", String.class).apply(0);
indexNameFunc = op.getAsRequiredFunction("delete", String.class);
deleteRequestFunc = createDeleteRequestFunction(op);
}
@Override
public PineconeOp apply(long value) {
return new PineconeDeleteOp(pcFunction.apply(value).getConnection(indexName), deleteRequestFunc.apply(value));
return new PineconeDeleteOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
deleteRequestFunc.apply(value));
}
private LongFunction<DeleteRequest> createDeleteRequestFunction(ParsedOp op) {

View File

@ -11,7 +11,6 @@ import io.pinecone.proto.DescribeIndexStatsRequest;
import java.util.function.LongFunction;
public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
private final String indexName;
private final LongFunction<DescribeIndexStatsRequest> indexStatsRequestFunc;
public PineconeDescribeIndexStatsOpDispenser(PineconeDriverAdapter adapter,
@ -19,8 +18,7 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
indexName = op.getAsRequiredFunction("indexStatsRequest", String.class).apply(0);
indexNameFunc = op.getAsRequiredFunction("indexStatsRequest", String.class);
indexStatsRequestFunc = createDescribeIndexStatsRequestFunction(op);
}
@ -31,6 +29,7 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
@Override
public PineconeOp apply(long value) {
return new PineconeDescribeIndexStatsOp(pcFunction.apply(value).getConnection(indexName), indexStatsRequestFunc.apply(value));
return new PineconeDescribeIndexStatsOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
indexStatsRequestFunc.apply(value));
}
}

View File

@ -15,14 +15,13 @@ import java.util.function.LongFunction;
public class PineconeFetchOpDispenser extends PineconeOpDispenser {
private final LongFunction<FetchRequest> fetchRequestFunc;
private final String indexName;
public PineconeFetchOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
indexName = op.getAsRequiredFunction("fetch", String.class).apply(0);
indexNameFunc = op.getAsRequiredFunction("fetch", String.class);
fetchRequestFunc = createFetchRequestFunction(op);
}
@ -53,6 +52,7 @@ public class PineconeFetchOpDispenser extends PineconeOpDispenser {
@Override
public PineconeOp apply(long value) {
return new PineconeFetchOp(pcFunction.apply(value).getConnection(indexName), fetchRequestFunc.apply(value));
return new PineconeFetchOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
fetchRequestFunc.apply(value));
}
}

View File

@ -28,6 +28,8 @@ import java.util.function.LongFunction;
public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
protected LongFunction<PineconeSpace> pcFunction;
protected LongFunction<String> targetFunction;
protected LongFunction<String> indexNameFunc;
protected PineconeOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,

View File

@ -38,7 +38,6 @@ import java.util.function.LongFunction;
public class PineconeQueryOpDispenser extends PineconeOpDispenser {
private final LongFunction<QueryRequest> queryRequestFunc;
private final String indexName;
public PineconeQueryOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
@ -46,7 +45,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
indexName = op.getAsRequiredFunction("query", String.class).apply(0);
indexNameFunc = op.getAsRequiredFunction("query", String.class);
queryRequestFunc = createQueryRequestFunc(op, createQueryVectorFunc(op));
}
@ -112,6 +111,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
@Override
public PineconeOp apply(long value) {
return new PineconeQueryOp(pcFunction.apply(value).getConnection(indexName), queryRequestFunc.apply(value));
return new PineconeQueryOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
queryRequestFunc.apply(value));
}
}

View File

@ -11,8 +11,7 @@ import io.pinecone.proto.UpdateRequest;
import java.util.function.LongFunction;
public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
private UpdateRequest request;
private PineconeConnection connection;
private final LongFunction<UpdateRequest> updateRequestFunc;
public PineconeUpdateOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
@ -20,17 +19,18 @@ public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("update");
connection = pcFunction.apply(0).getConnection(indexName);
request = createUpdateRequest();
indexNameFunc = op.getAsRequiredFunction("update", String.class);
updateRequestFunc = createUpdateRequestFunction(op);
}
private UpdateRequest createUpdateRequest() {
return UpdateRequest.newBuilder().build();
private LongFunction<UpdateRequest> createUpdateRequestFunction(ParsedOp op) {
LongFunction<UpdateRequest.Builder> rFunc = l -> UpdateRequest.newBuilder();
return l -> rFunc.apply(l).build();
}
@Override
public PineconeOp apply(long value) {
return new PineconeUpdateOp(connection, request);
return new PineconeUpdateOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
updateRequestFunc.apply(value));
}
}

View File

@ -16,10 +16,30 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.LongFunction;
/**
* 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();
*/
public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
private UpsertRequest request;
private PineconeConnection connection;
private final LongFunction<UpsertRequest> upsertRequestFunc;
public PineconeUpsertOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
@ -27,34 +47,17 @@ public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
String indexName = op.getStaticValue("upsert");
connection = pcFunction.apply(0).getConnection(indexName);
request = createUpsertRequest();
indexNameFunc = op.getAsRequiredFunction("upsert", String.class);
upsertRequestFunc = createUpsertRequestFunc(op);
}
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();
private LongFunction<UpsertRequest> createUpsertRequestFunc(ParsedOp op) {
return null;
}
@Override
public PineconeOp apply(long value) {
return new PineconeUpsertOp(connection, request);
return new PineconeUpsertOp(pcFunction.apply(value).getConnection(indexNameFunc.apply(value)),
upsertRequestFunc.apply(value));
}
}