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

cleanup and fixing minor issues

Browse Source
This commit is contained in:
Mark Wolters 2023-05-17 19:42:20 +00:00 committed by Madhavan
parent 8b30d6251a
commit cc907de1b6
9 changed files with 80 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
adapter-pinecone/src/main/java/io/nosqlbench/adapter/pinecone/PineconeDriverAdapter.java
View File

@ -8,6 +8,8 @@ import io.nosqlbench.engine.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.engine.api.activityimpl.uniform.DriverSpaceCache;
import io.nosqlbench.nb.annotations.Service;
import java.util.function.Function;
@Service(value = DriverAdapter.class, selector = "pinecone")
public class PineconeDriverAdapter extends BaseDriverAdapter<PineconeOp, PineconeSpace> {
@ -18,4 +20,9 @@ public class PineconeDriverAdapter extends BaseDriverAdapter<PineconeOp, Pinecon
return new PineconeOpMapper(this, spaceCache, adapterConfig);
}
@Override
public Function<String, ? extends PineconeSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new PineconeSpace(s,cfg);
}
}

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

@ -25,7 +25,7 @@ public class PineconeSpace {
private final PineconeClient client;
private final PineconeClientConfig config;
private final Map<String,PineconeConnection> connections = new HashMap<String,PineconeConnection>();
private final Map<String,PineconeConnection> connections = new HashMap<>();
/**
* Create a new PineconeSpace Object which stores all stateful contextual information needed to interact

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

@ -76,7 +76,7 @@ public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
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()))
Value.newBuilder().setNumberValue(Integer.parseInt(filterFields[2])).build()))
.build()).build();
};
rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l));

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

@ -8,7 +8,6 @@ import io.nosqlbench.adapter.pinecone.ops.PineconeDescribeIndexStatsOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.DescribeIndexStatsRequest;
import jakarta.ws.rs.NotSupportedException;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@ -44,7 +43,7 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
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()))
Value.newBuilder().setNumberValue(Integer.parseInt(filterFields[2])).build()))
.build()).build();
};
rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l));

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

@ -5,7 +5,6 @@ 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.templating.ParsedOp;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.FetchRequest;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

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

@ -21,7 +21,6 @@ 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;
import java.util.function.LongFunction;

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

@ -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.PineconeOp;
@ -11,12 +13,15 @@ import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Optional;
import java.util.function.LongFunction;
public class PineconeQueryOpDispenser extends PineconeOpDispenser {
private static final Logger LOGGER = LogManager.getLogger(PineconeQueryOpDispenser.class);
private final LongFunction<QueryRequest> queryRequestFunc;
private final LongFunction<QueryRequest.Builder> queryRequestFunc;
private final LongFunction<Collection<QueryVector>> queryVectorFunc;
/**
* Create a new PineconeQueryOpDispenser subclassed from {@link PineconeOpDispenser}.
@ -32,32 +37,10 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
queryRequestFunc = createQueryRequestFunc(op);
queryVectorFunc = createQueryVectorFunc(op);
}
/*
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();
QueryRequest queryRequest = QueryRequest.newBuilder()
.addQueries(queryVector)
.setNamespace("default-namespace")
.setTopK(2)
.setIncludeMetadata(true)
.build();
}
*/
private LongFunction<QueryRequest> createQueryRequestFunc(ParsedOp op) {
private LongFunction<QueryRequest.Builder> createQueryRequestFunc(ParsedOp op) {
LongFunction<QueryRequest.Builder> rFunc = l -> QueryRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
@ -94,31 +77,49 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
LongFunction<String> af = vFunc.get();
LongFunction<ArrayList<Float>> alf = l -> {
String[] vals = af.apply(l).split(",");
ArrayList<Float> fVals = new ArrayList<Float>();
for (int i = 0; i < vals.length; i++) {
fVals.add(Float.valueOf(vals[i]));
ArrayList<Float> fVals = new ArrayList<>();
for (String val : vals) {
fVals.add(Float.valueOf(val));
}
return fVals;
};
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));
Optional<LongFunction<String>> filterFunction = op.getAsOptionalFunction("filter", String.class);
if (filterFunction.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<Struct> 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<QueryRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
return l -> finalRFunc.apply(l);
}
private LongFunction<QueryVector> createQueryVectorFunc(ParsedOp op) {
LongFunction<QueryVector.Builder> vFunc = l -> QueryVector.newBuilder();
LongFunction<QueryVector.Builder> finalVFunc = vFunc;
return l -> finalVFunc.apply(l).build();
private LongFunction<Collection<QueryVector>> createQueryVectorFunc(ParsedOp op) {
//Optional<LongFunction<Collection<Map<String,String>>>> baseFunc = op.getAsOptionalFunction("query_vectors", String.class);
// LongFunction<QueryVector.Builder> vFunc = l -> QueryVector.newBuilder();
// LongFunction<QueryVector.Builder> finalVFunc = vFunc;
//return l -> finalVFunc.apply(l).build();
return l -> null;
}
@Override
public PineconeOp apply(long value) {
return new PineconeQueryOp(pcFunction.apply(value).getConnection(targetFunction.apply(value)),
queryRequestFunc.apply(value));
QueryRequest.Builder qrb = queryRequestFunc.apply(value);
Collection<QueryVector> vectors = queryVectorFunc.apply(value);
if (vectors != null) {
qrb.addAllQueries(vectors);
}
return new PineconeQueryOp(pcFunction.apply(value).getConnection(targetFunction.apply(value)), qrb.build());
}
}

24
adapter-pinecone/src/main/resources/docs/pinecone.md
View File

@ -13,13 +13,7 @@ ops:
# The number of results to return for each query.
top_k: int_query_topk
# You can use vector metadata to limit your search. See https://www.pinecone.io/docs/metadata-filtering/
filters:
- filter_field: query_filter_field
operator: [$lt, $gt, $eq, ...]
comparator: query_compval
- filter_field: query_filter_field
operator: [$lt, $gt, $eq, ...]
comparator: query_compval
filter: <field operator compval>
# Indicates whether vector values are included in the response.
include_values: boolean
# Indicates whether metadata is included in the response as well as the ids.
@ -29,13 +23,7 @@ ops:
values: csv_separated_floats
top_k: int_val
namespace: string_val
filters:
- filter_field: query_vector_filter_field
operator: comparison_operator
comparator: comparator_val
- filter_field: query_filter_field
operator: [$lt, $gt, $eq, ...]
comparator: query_compval
filter: <field operator compval>
sparse_values:
indices: list_of_ints
values: list_of_floats
@ -43,13 +31,7 @@ ops:
values: csv_separated_floats
top_k: int_val
namespace: string_val
filters:
- filter_field: query_vector_filter_field
operator: comparison_operator
comparator: comparator_val
- filter_field: query_vector_filter_field
operator: comparison_operator
comparator: comparator_val
filter: <field operator compval>
sparse_values:
indices: list_of_ints
values: list_of_floats

43
adapter-pinecone/src/test/java/io/nosqlbench/adapter/pinecone/PineconeOpMapperTest.java
View File

@ -47,19 +47,31 @@ public class PineconeOpMapperTest {
@Test
public void testQueryOpDispenserSimple() {
ParsedOp pop = parsedOpFor("""
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
""");
ops:
op1:
type: "query"
index: "test-index"
vector: "1.0,2.0,3.0"
namespace: "test-namespace"
top_k: 10
filter: "value $lt 2"
include_values: true
include_metadata: true
query_vectors:
- id: 1
values: "1.0,2.0,3.0"
top_k: 8
namespace: "test-namespace"
filter: "value $lt 2"
sparse_values:
indices: "1,2,3"
values: "1.0,2.0,3.0"
- id: 2
values: "4.0,5.0,6.0"
top_k: 11
namespace: "test-namespace"
filter: "value $gt 10"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeQueryOpDispenser);
}
@ -74,12 +86,11 @@ public class PineconeOpMapperTest {
ids: "1.0,2.0,3.0"
namespace: "test-namespace"
deleteall: true
filters:
- "value $lt 2"
- "value $gt 10"
filter: "value $gt 10"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeDeleteOpDispenser);
PineconeOp op = dispenser.apply(0);
}
@Test