mirror of
https://github.com/nosqlbench/nosqlbench.git
synced 2025-02-25 18:55:28 -06:00
Merge branch 'main' into nosqlbench-1295-vectorfuncs
This commit is contained in:
commit
525e1d2b66
@ -50,7 +50,7 @@
|
|||||||
<dependency>
|
<dependency>
|
||||||
<groupId>org.apache.tinkerpop</groupId>
|
<groupId>org.apache.tinkerpop</groupId>
|
||||||
<artifactId>gremlin-core</artifactId>
|
<artifactId>gremlin-core</artifactId>
|
||||||
<version>3.6.3</version>
|
<version>3.6.4</version>
|
||||||
</dependency>
|
</dependency>
|
||||||
<dependency>
|
<dependency>
|
||||||
<groupId>org.apache.tinkerpop</groupId>
|
<groupId>org.apache.tinkerpop</groupId>
|
||||||
|
@ -40,7 +40,7 @@ public class PineconeOpMapper implements OpMapper<PineconeOp> {
|
|||||||
* Create a new PineconeOpMapper implementing the {@link OpMapper} interface.
|
* Create a new PineconeOpMapper implementing the {@link OpMapper} interface.
|
||||||
*
|
*
|
||||||
* @param adapter The associated {@link PineconeDriverAdapter}
|
* @param adapter The associated {@link PineconeDriverAdapter}
|
||||||
* @param spaceCache A cached context Object of thpe {@link PineconeSpace})
|
* @param spaceCache A cached context Object of the {@link PineconeSpace})
|
||||||
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
|
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
|
||||||
*/
|
*/
|
||||||
public PineconeOpMapper(PineconeDriverAdapter adapter,
|
public PineconeOpMapper(PineconeDriverAdapter adapter,
|
||||||
|
@ -65,7 +65,7 @@ public class PineconeSpace {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* Connections are index-specific so we need to allow for multiple connection management across indices.
|
* Connections are index-specific, so we need to allow for multiple connection management across indices.
|
||||||
* However, note that a single connection object is thread safe and can be used by multiple clients.
|
* However, note that a single connection object is thread safe and can be used by multiple clients.
|
||||||
*
|
*
|
||||||
* @param index The database index for which a connection is being requested
|
* @param index The database index for which a connection is being requested
|
||||||
|
@ -16,9 +16,7 @@
|
|||||||
|
|
||||||
package io.nosqlbench.adapter.pinecone.opdispensers;
|
package io.nosqlbench.adapter.pinecone.opdispensers;
|
||||||
|
|
||||||
import com.google.protobuf.ListValue;
|
|
||||||
import com.google.protobuf.Struct;
|
import com.google.protobuf.Struct;
|
||||||
import com.google.protobuf.Value;
|
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
||||||
import io.nosqlbench.adapter.pinecone.ops.PineconeDeleteOp;
|
import io.nosqlbench.adapter.pinecone.ops.PineconeDeleteOp;
|
||||||
|
@ -17,7 +17,6 @@
|
|||||||
package io.nosqlbench.adapter.pinecone.opdispensers;
|
package io.nosqlbench.adapter.pinecone.opdispensers;
|
||||||
|
|
||||||
import com.google.protobuf.Struct;
|
import com.google.protobuf.Struct;
|
||||||
import com.google.protobuf.Value;
|
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
||||||
import io.nosqlbench.adapter.pinecone.ops.PineconeDescribeIndexStatsOp;
|
import io.nosqlbench.adapter.pinecone.ops.PineconeDescribeIndexStatsOp;
|
||||||
@ -54,11 +53,11 @@ public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
|
|||||||
/**
|
/**
|
||||||
* @param op The ParsedOp used to build the Request
|
* @param op The ParsedOp used to build the Request
|
||||||
* @return A function that will take a long (the current cycle) and return a Pinecone DescribeIndexStatsRequest
|
* @return A function that will take a long (the current cycle) and return a Pinecone DescribeIndexStatsRequest
|
||||||
*
|
* <p>
|
||||||
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
||||||
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
||||||
* For each method in the chain a function is created here and added to the chain of functions
|
* For each method in the chain a function is created here and added to the chain of functions
|
||||||
* called at time of instantiation. Additionally some of the arguments to the builder methods require
|
* called at time of instantiation. Additionally, some of the arguments to the builder methods require
|
||||||
* creation through their own builder process. In these cases the pattern adopted includes multiple layers of
|
* creation through their own builder process. In these cases the pattern adopted includes multiple layers of
|
||||||
* functions in order to build all objects in the correct manner and ordering.
|
* functions in order to build all objects in the correct manner and ordering.
|
||||||
*/
|
*/
|
||||||
|
@ -53,7 +53,7 @@ public class PineconeFetchOpDispenser extends PineconeOpDispenser {
|
|||||||
/**
|
/**
|
||||||
* @param op The ParsedOp used to build the Request
|
* @param op The ParsedOp used to build the Request
|
||||||
* @return A function that will take a long (the current cycle) and return a Pinecone FetchRequest
|
* @return A function that will take a long (the current cycle) and return a Pinecone FetchRequest
|
||||||
*
|
* <p>
|
||||||
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
||||||
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
||||||
* For each method in the chain a function is created here and added to the chain of functions
|
* For each method in the chain a function is created here and added to the chain of functions
|
||||||
|
@ -25,8 +25,11 @@ import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
|
|||||||
import io.nosqlbench.engine.api.activityimpl.BaseOpDispenser;
|
import io.nosqlbench.engine.api.activityimpl.BaseOpDispenser;
|
||||||
import io.nosqlbench.engine.api.templating.ParsedOp;
|
import io.nosqlbench.engine.api.templating.ParsedOp;
|
||||||
|
|
||||||
|
import java.util.ArrayList;
|
||||||
|
import java.util.HashMap;
|
||||||
import java.util.List;
|
import java.util.List;
|
||||||
import java.util.Map;
|
import java.util.Map;
|
||||||
|
import java.util.function.BiConsumer;
|
||||||
import java.util.function.LongFunction;
|
import java.util.function.LongFunction;
|
||||||
|
|
||||||
public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
|
public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
|
||||||
@ -73,6 +76,60 @@ public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, Pi
|
|||||||
return listValueBuilder.build();
|
return listValueBuilder.build();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
protected LongFunction<List<Float>> extractFloatVals(LongFunction<Object> af) {
|
||||||
|
return l -> this.getVectorValues(af.apply(l));
|
||||||
|
}
|
||||||
|
|
||||||
|
protected Map<String, Value> generateMetadataMap(Map<String, Object> metadata_values_map) {
|
||||||
|
Map<String, Value> metadata_map = new HashMap<>();
|
||||||
|
BiConsumer<String,Object> stringToValue = (key, val) -> {
|
||||||
|
Value targetval;
|
||||||
|
if (val instanceof String) targetval = Value.newBuilder().setStringValue((String)val).build();
|
||||||
|
else if (val instanceof Number) targetval = Value.newBuilder().setNumberValue((((Number) val).doubleValue())).build();
|
||||||
|
else if (val instanceof List) targetval = Value.newBuilder().setListValue(generateListValue((List) val)).build();
|
||||||
|
else if (val instanceof Boolean) targetval = Value.newBuilder().setBoolValue((Boolean) val).build();
|
||||||
|
else throw new RuntimeException("Unsupported metadata value type");
|
||||||
|
metadata_map.put(key, targetval);
|
||||||
|
};
|
||||||
|
metadata_values_map.forEach(stringToValue);
|
||||||
|
return metadata_map;
|
||||||
|
}
|
||||||
|
|
||||||
|
protected List<Float> getVectorValues(Object rawVectorValues) {
|
||||||
|
List<Float> floatValues;
|
||||||
|
if (rawVectorValues instanceof String) {
|
||||||
|
floatValues = new ArrayList<>();
|
||||||
|
String[] rawValues = (((String) rawVectorValues).split(","));
|
||||||
|
for (String val : rawValues) {
|
||||||
|
floatValues.add(Float.valueOf(val));
|
||||||
|
}
|
||||||
|
} else if (rawVectorValues instanceof List) {
|
||||||
|
floatValues = switch (((List<?>) rawVectorValues).get(0).getClass().getSimpleName()) {
|
||||||
|
case "Float" -> (List<Float>) rawVectorValues;
|
||||||
|
case "Double" -> ((List<Double>) rawVectorValues).stream().map(Double::floatValue).toList();
|
||||||
|
case "String" -> ((List<String>) rawVectorValues).stream().map(Float::parseFloat).toList();
|
||||||
|
default -> throw new RuntimeException("Invalid type specified for values");
|
||||||
|
};
|
||||||
|
} else {
|
||||||
|
throw new RuntimeException("Invalid type specified for values");
|
||||||
|
}
|
||||||
|
return floatValues;
|
||||||
|
}
|
||||||
|
|
||||||
|
protected List<Integer> getIndexValues(Object rawIndexValues) {
|
||||||
|
List<Integer> intValues;
|
||||||
|
if (rawIndexValues instanceof String) {
|
||||||
|
intValues = new ArrayList<>();
|
||||||
|
String[] rawValues = (((String) rawIndexValues).split(","));
|
||||||
|
for (String val : rawValues) {
|
||||||
|
intValues.add(Integer.valueOf(val));
|
||||||
|
}
|
||||||
|
} else if (rawIndexValues instanceof List) {
|
||||||
|
intValues = (List<Integer>) rawIndexValues;
|
||||||
|
}else {
|
||||||
|
throw new RuntimeException("Invalid type specified for Index values");
|
||||||
|
}
|
||||||
|
return intValues;
|
||||||
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
@ -77,7 +77,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
|
|||||||
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
|
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
|
||||||
}
|
}
|
||||||
|
|
||||||
Optional<LongFunction<Integer>> tFunc = op.getAsOptionalFunction("topk", Integer.class);
|
Optional<LongFunction<Integer>> tFunc = op.getAsOptionalFunction("top_k", Integer.class);
|
||||||
if (tFunc.isPresent()) {
|
if (tFunc.isPresent()) {
|
||||||
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
|
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
|
||||||
LongFunction<Integer> af = tFunc.get();
|
LongFunction<Integer> af = tFunc.get();
|
||||||
@ -98,18 +98,12 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
|
|||||||
rFunc = l -> finalFunc.apply(l).setIncludeValues(af.apply(l));
|
rFunc = l -> finalFunc.apply(l).setIncludeValues(af.apply(l));
|
||||||
}
|
}
|
||||||
|
|
||||||
Optional<LongFunction<String>> vFunc = op.getAsOptionalFunction("vector", String.class);
|
Optional<LongFunction<Object>> vFunc = op.getAsOptionalFunction("vector", Object.class);
|
||||||
if (vFunc.isPresent()) {
|
if (vFunc.isPresent()) {
|
||||||
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
|
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
|
||||||
LongFunction<String> af = vFunc.get();
|
LongFunction<Object> af = vFunc.get();
|
||||||
LongFunction<ArrayList<Float>> alf = l -> {
|
|
||||||
String[] vals = af.apply(l).split(",");
|
LongFunction<List<Float>> alf = extractFloatVals(af);
|
||||||
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));
|
rFunc = l -> finalFunc.apply(l).addAllVector(alf.apply(l));
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -123,6 +117,7 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
|
|||||||
return rFunc;
|
return rFunc;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* @param op the ParsedOp from which the Query Vector objects will be built
|
* @param op the ParsedOp from which the Query Vector objects will be built
|
||||||
* @return an Iterable Collection of QueryVector objects to be added to a Pinecone QueryRequest
|
* @return an Iterable Collection of QueryVector objects to be added to a Pinecone QueryRequest
|
||||||
@ -140,37 +135,20 @@ public class PineconeQueryOpDispenser extends PineconeOpDispenser {
|
|||||||
List<Map<String, Object>> vectors = listLongFunction.apply(l);
|
List<Map<String, Object>> vectors = listLongFunction.apply(l);
|
||||||
for (Map<String, Object> vector : vectors) {
|
for (Map<String, Object> vector : vectors) {
|
||||||
QueryVector.Builder qvb = QueryVector.newBuilder();
|
QueryVector.Builder qvb = QueryVector.newBuilder();
|
||||||
String[] rawValues = ((String) vector.get("values")).split(",");
|
qvb.addAllValues(getVectorValues(vector.get("values")));
|
||||||
ArrayList<Float> floatValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
floatValues.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
qvb.addAllValues(floatValues);
|
|
||||||
qvb.setNamespace((String) vector.get("namespace"));
|
qvb.setNamespace((String) vector.get("namespace"));
|
||||||
if (vector.containsKey("top_k")) {
|
if (vector.containsKey("top_k")) {
|
||||||
qvb.setTopK((Integer) vector.get("top_k"));
|
qvb.setTopK((Integer) vector.get("top_k"));
|
||||||
}
|
}
|
||||||
if (vector.containsKey("filter")) {
|
if (vector.containsKey("filter")) {
|
||||||
LongFunction<Struct> builtFilter = buildFilterStruct(l2 -> {
|
LongFunction<Struct> builtFilter = buildFilterStruct(l2 -> (Map) vector.get("filter"));
|
||||||
return (Map) vector.get("filter");
|
|
||||||
});
|
|
||||||
qvb.setFilter(builtFilter.apply(l));
|
qvb.setFilter(builtFilter.apply(l));
|
||||||
}
|
}
|
||||||
if (vector.containsKey("sparse_values")) {
|
if (vector.containsKey("sparse_values")) {
|
||||||
Map<String,String> sparse_values = (Map<String, String>) vector.get("sparse_values");
|
Map<String,String> sparse_values = (Map<String, String>) vector.get("sparse_values");
|
||||||
rawValues = ((String) sparse_values.get("values")).split(",");
|
|
||||||
floatValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
floatValues.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
rawValues = sparse_values.get("indices").split(",");
|
|
||||||
List<Integer> intValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
intValues.add(Integer.valueOf(val));
|
|
||||||
}
|
|
||||||
qvb.setSparseValues(SparseValues.newBuilder()
|
qvb.setSparseValues(SparseValues.newBuilder()
|
||||||
.addAllValues(floatValues)
|
.addAllValues(getVectorValues(sparse_values.get("values")))
|
||||||
.addAllIndices(intValues)
|
.addAllIndices(getIndexValues(sparse_values.get("indices")))
|
||||||
.build());
|
.build());
|
||||||
}
|
}
|
||||||
returnVectors.add(qvb.build());
|
returnVectors.add(qvb.build());
|
||||||
|
@ -17,7 +17,6 @@
|
|||||||
package io.nosqlbench.adapter.pinecone.opdispensers;
|
package io.nosqlbench.adapter.pinecone.opdispensers;
|
||||||
|
|
||||||
import com.google.protobuf.Struct;
|
import com.google.protobuf.Struct;
|
||||||
import com.google.protobuf.Value;
|
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
||||||
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
|
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
|
||||||
@ -28,8 +27,9 @@ import io.pinecone.proto.UpdateRequest;
|
|||||||
import org.apache.logging.log4j.LogManager;
|
import org.apache.logging.log4j.LogManager;
|
||||||
import org.apache.logging.log4j.Logger;
|
import org.apache.logging.log4j.Logger;
|
||||||
|
|
||||||
import java.util.*;
|
import java.util.List;
|
||||||
import java.util.function.BiConsumer;
|
import java.util.Map;
|
||||||
|
import java.util.Optional;
|
||||||
import java.util.function.LongFunction;
|
import java.util.function.LongFunction;
|
||||||
|
|
||||||
public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
||||||
@ -59,7 +59,7 @@ public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
|||||||
/**
|
/**
|
||||||
* @param op the ParsedOp from which the SparseValues object will be built
|
* @param op the ParsedOp from which the SparseValues object will be built
|
||||||
* @return a SparseValues Object to be added to a Pinecone UpdateRequest
|
* @return a SparseValues Object to be added to a Pinecone UpdateRequest
|
||||||
*
|
* <p>
|
||||||
* This method interrogates the subsection of the ParsedOp defined for SparseValues parameters and constructs
|
* This method interrogates the subsection of the ParsedOp defined for SparseValues parameters and constructs
|
||||||
* a SparseValues Object based on the included values, or returns null if this section is not populated. The
|
* a SparseValues Object based on the included values, or returns null if this section is not populated. The
|
||||||
* base function returns either the SparseValues Object or null, while the interior function builds the SparseValues
|
* base function returns either the SparseValues Object or null, while the interior function builds the SparseValues
|
||||||
@ -69,27 +69,17 @@ public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
|||||||
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("sparse_values", Map.class);
|
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("sparse_values", Map.class);
|
||||||
return mFunc.<LongFunction<SparseValues>>map(mapLongFunction -> l -> {
|
return mFunc.<LongFunction<SparseValues>>map(mapLongFunction -> l -> {
|
||||||
Map<String, String> sparse_values_map = mapLongFunction.apply(l);
|
Map<String, String> sparse_values_map = mapLongFunction.apply(l);
|
||||||
String[] rawValues = (sparse_values_map.get("values")).split(",");
|
|
||||||
ArrayList floatValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
floatValues.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
rawValues = sparse_values_map.get("indices").split(",");
|
|
||||||
List<Integer> intValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
intValues.add(Integer.valueOf(val));
|
|
||||||
}
|
|
||||||
return SparseValues.newBuilder()
|
return SparseValues.newBuilder()
|
||||||
.addAllValues(floatValues)
|
.addAllValues(getVectorValues(sparse_values_map.get("values")))
|
||||||
.addAllIndices(intValues)
|
.addAllIndices(getIndexValues(sparse_values_map.get("indices")))
|
||||||
.build();
|
.build();
|
||||||
}).orElse(null);
|
}).orElse(null);
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* @param op the ParsedOp from which the Metadata objects will be built
|
* @param op the ParsedOp from which the Metadata objects will be built
|
||||||
* @return an Metadata Struct to be added to a Pinecone UpdateRequest
|
* @return a Metadata Struct to be added to a Pinecone UpdateRequest
|
||||||
*
|
* <p>
|
||||||
* This method interrogates the subsection of the ParsedOp defined for metadata parameters and constructs
|
* This method interrogates the subsection of the ParsedOp defined for metadata parameters and constructs
|
||||||
* a Metadata Struct based on the included values, or returns null if this section is not populated. The
|
* a Metadata Struct based on the included values, or returns null if this section is not populated. The
|
||||||
* base function returns either the Metadata Struct or null, while the interior function builds the Metadata
|
* base function returns either the Metadata Struct or null, while the interior function builds the Metadata
|
||||||
@ -98,30 +88,23 @@ public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
|||||||
private LongFunction<Struct> createUpdateMetadataFunction(ParsedOp op) {
|
private LongFunction<Struct> createUpdateMetadataFunction(ParsedOp op) {
|
||||||
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("metadata", Map.class);
|
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("metadata", Map.class);
|
||||||
return mFunc.<LongFunction<Struct>>map(mapLongFunction -> l -> {
|
return mFunc.<LongFunction<Struct>>map(mapLongFunction -> l -> {
|
||||||
Map<String, Value> metadata_map = new HashMap<String,Value>();
|
|
||||||
BiConsumer<String,Object> stringToValue = (key, val) -> {
|
|
||||||
Value targetval = null;
|
|
||||||
if (val instanceof String) targetval = Value.newBuilder().setStringValue((String)val).build();
|
|
||||||
else if (val instanceof Number) targetval = Value.newBuilder().setNumberValue((((Number) val).doubleValue())).build();
|
|
||||||
metadata_map.put(key, targetval);
|
|
||||||
};
|
|
||||||
Map<String, Object> metadata_values_map = mapLongFunction.apply(l);
|
Map<String, Object> metadata_values_map = mapLongFunction.apply(l);
|
||||||
metadata_values_map.forEach(stringToValue);
|
return UpdateRequest.newBuilder().getSetMetadataBuilder()
|
||||||
return UpdateRequest.newBuilder().getSetMetadataBuilder().putAllFields(metadata_map).build();
|
.putAllFields(generateMetadataMap(metadata_values_map)).build();
|
||||||
}).orElse(null);
|
}).orElse(null);
|
||||||
}
|
}
|
||||||
|
|
||||||
/**
|
/**
|
||||||
* @param op The ParsedOp used to build the Request
|
* @param op The ParsedOp used to build the Request
|
||||||
* @return A function that will take a long (the current cycle) and return a Pinecone UpdateRequest Builder
|
* @return A function that will take a long (the current cycle) and return a Pinecone UpdateRequest Builder
|
||||||
*
|
* <p>
|
||||||
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
||||||
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
||||||
* For each method in the chain a function is created here and added to the chain of functions
|
* For each method in the chain a function is created here and added to the chain of functions
|
||||||
* called at time of instantiation.
|
* called at time of instantiation.
|
||||||
*
|
* <p>
|
||||||
* The Metadata and SparseValues objects used by the UpdateRequest are sufficiently sophisticated in their own
|
* The Metadata and SparseValues objects used by the UpdateRequest are sufficiently sophisticated in their own
|
||||||
* building process that they have been broken out into separate methods. At runtime they are built separately
|
* building process that they have been broken out into separate methods. At runtime, they are built separately
|
||||||
* and then added to the build chain by the builder returned by this method.
|
* and then added to the build chain by the builder returned by this method.
|
||||||
*/
|
*/
|
||||||
private LongFunction<UpdateRequest.Builder> createUpdateRequestFunction(ParsedOp op) {
|
private LongFunction<UpdateRequest.Builder> createUpdateRequestFunction(ParsedOp op) {
|
||||||
@ -141,18 +124,11 @@ public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
|
|||||||
rFunc = l -> finalFunc.apply(l).setId(af.apply(l));
|
rFunc = l -> finalFunc.apply(l).setId(af.apply(l));
|
||||||
}
|
}
|
||||||
|
|
||||||
Optional<LongFunction<String>> vFunc = op.getAsOptionalFunction("values", String.class);
|
Optional<LongFunction<Object>> vFunc = op.getAsOptionalFunction("values", Object.class);
|
||||||
if (vFunc.isPresent()) {
|
if (vFunc.isPresent()) {
|
||||||
LongFunction<UpdateRequest.Builder> finalFunc = rFunc;
|
LongFunction<UpdateRequest.Builder> finalFunc = rFunc;
|
||||||
LongFunction<String> af = vFunc.get();
|
LongFunction<Object> af = vFunc.get();
|
||||||
LongFunction<ArrayList<Float>> alf = l -> {
|
LongFunction<List<Float>> alf = extractFloatVals(af);
|
||||||
String[] vals = af.apply(l).split(",");
|
|
||||||
ArrayList<Float> fVals = new ArrayList<>();
|
|
||||||
for (String val : vals) {
|
|
||||||
fVals.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
return fVals;
|
|
||||||
};
|
|
||||||
rFunc = l -> finalFunc.apply(l).addAllValues(alf.apply(l));
|
rFunc = l -> finalFunc.apply(l).addAllValues(alf.apply(l));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -17,7 +17,6 @@
|
|||||||
package io.nosqlbench.adapter.pinecone.opdispensers;
|
package io.nosqlbench.adapter.pinecone.opdispensers;
|
||||||
|
|
||||||
import com.google.protobuf.Struct;
|
import com.google.protobuf.Struct;
|
||||||
import com.google.protobuf.Value;
|
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
import io.nosqlbench.adapter.pinecone.PineconeDriverAdapter;
|
||||||
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
import io.nosqlbench.adapter.pinecone.PineconeSpace;
|
||||||
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
|
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
|
||||||
@ -30,7 +29,6 @@ import org.apache.logging.log4j.LogManager;
|
|||||||
import org.apache.logging.log4j.Logger;
|
import org.apache.logging.log4j.Logger;
|
||||||
|
|
||||||
import java.util.*;
|
import java.util.*;
|
||||||
import java.util.function.BiConsumer;
|
|
||||||
import java.util.function.LongFunction;
|
import java.util.function.LongFunction;
|
||||||
|
|
||||||
public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
|
public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
|
||||||
@ -58,7 +56,7 @@ public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
|
|||||||
/**
|
/**
|
||||||
* @param op the ParsedOp from which the Vector objects will be built
|
* @param op the ParsedOp from which the Vector objects will be built
|
||||||
* @return an Iterable Collection of Vector objects to be added to a Pinecone UpsertRequest
|
* @return an Iterable Collection of Vector objects to be added to a Pinecone UpsertRequest
|
||||||
*
|
* <p>
|
||||||
* This method interrogates the subsection of the ParsedOp defined for Vector parameters and constructs
|
* This method interrogates the subsection of the ParsedOp defined for Vector parameters and constructs
|
||||||
* a list of Vectors based on the included values, or returns null if this section is not populated. The
|
* a list of Vectors based on the included values, or returns null if this section is not populated. The
|
||||||
* base function returns either the List of vectors or null, while the interior function builds the vectors
|
* base function returns either the List of vectors or null, while the interior function builds the vectors
|
||||||
@ -74,40 +72,17 @@ public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
|
|||||||
Vector.Builder vb = Vector.newBuilder();
|
Vector.Builder vb = Vector.newBuilder();
|
||||||
// No need to check for key, it is invalid if id is not there, let it throw an exception
|
// No need to check for key, it is invalid if id is not there, let it throw an exception
|
||||||
vb.setId(vector.get("id").toString());
|
vb.setId(vector.get("id").toString());
|
||||||
String[] rawValues = ((String) vector.get("values")).split(",");
|
vb.addAllValues(getVectorValues(vector.get("values")));
|
||||||
ArrayList<Float> floatValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
floatValues.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
vb.addAllValues(floatValues);
|
|
||||||
if (vector.containsKey("sparse_values")) {
|
if (vector.containsKey("sparse_values")) {
|
||||||
Map<String,String> sparse_values = (Map<String, String>) vector.get("sparse_values");
|
Map<String,String> sparse_values = (Map<String, String>) vector.get("sparse_values");
|
||||||
rawValues = ((String) sparse_values.get("values")).split(",");
|
|
||||||
floatValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
floatValues.add(Float.valueOf(val));
|
|
||||||
}
|
|
||||||
rawValues = sparse_values.get("indices").split(",");
|
|
||||||
List<Integer> intValues = new ArrayList<>();
|
|
||||||
for (String val : rawValues) {
|
|
||||||
intValues.add(Integer.valueOf(val));
|
|
||||||
}
|
|
||||||
vb.setSparseValues(SparseValues.newBuilder()
|
vb.setSparseValues(SparseValues.newBuilder()
|
||||||
.addAllValues(floatValues)
|
.addAllValues(getVectorValues(sparse_values.get("values")))
|
||||||
.addAllIndices(intValues)
|
.addAllIndices(getIndexValues(sparse_values.get("indices")))
|
||||||
.build());
|
.build());
|
||||||
}
|
}
|
||||||
if (vector.containsKey("metadata")) {
|
if (vector.containsKey("metadata")) {
|
||||||
Map<String, Value> metadata_map = new HashMap<String, Value>();
|
|
||||||
BiConsumer<String,Object> stringToValue = (key, val) -> {
|
|
||||||
Value targetval = null;
|
|
||||||
if (val instanceof String) targetval = Value.newBuilder().setStringValue((String)val).build();
|
|
||||||
else if (val instanceof Number) targetval = Value.newBuilder().setNumberValue((((Number) val).doubleValue())).build();
|
|
||||||
metadata_map.put(key, targetval);
|
|
||||||
};
|
|
||||||
Map<String, Object> metadata_values_map = (Map<String, Object>) vector.get("metadata");
|
Map<String, Object> metadata_values_map = (Map<String, Object>) vector.get("metadata");
|
||||||
metadata_values_map.forEach(stringToValue);
|
vb.setMetadata(Struct.newBuilder().putAllFields(generateMetadataMap(metadata_values_map)).build());
|
||||||
vb.setMetadata(Struct.newBuilder().putAllFields(metadata_map).build());
|
|
||||||
}
|
}
|
||||||
returnVectors.add(vb.build());
|
returnVectors.add(vb.build());
|
||||||
}
|
}
|
||||||
@ -118,14 +93,14 @@ public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
|
|||||||
/**
|
/**
|
||||||
* @param op The ParsedOp used to build the Request
|
* @param op The ParsedOp used to build the Request
|
||||||
* @return A function that will take a long (the current cycle) and return a Pinecone UpsertRequest Builder
|
* @return A function that will take a long (the current cycle) and return a Pinecone UpsertRequest Builder
|
||||||
*
|
* <p>
|
||||||
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
* The pattern used here is to accommodate the way Request types are constructed for Pinecone.
|
||||||
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
* Requests use a Builder pattern, so at time of instantiation the methods should be chained together.
|
||||||
* For each method in the chain a function is created here and added to the chain of functions
|
* For each method in the chain a function is created here and added to the chain of functions
|
||||||
* called at time of instantiation.
|
* called at time of instantiation.
|
||||||
*
|
* <p>
|
||||||
* The Vector objects used by the UpsertRequest are sufficiently sophisticated in their own
|
* The Vector objects used by the UpsertRequest are sufficiently sophisticated in their own
|
||||||
* building process that they have been broken out into a separate method. At runtime they are built separately
|
* building process that they have been broken out into a separate method. At runtime, they are built separately
|
||||||
* and then added to the build chain by the builder returned by this method.
|
* and then added to the build chain by the builder returned by this method.
|
||||||
*/
|
*/
|
||||||
private LongFunction<UpsertRequest.Builder> createUpsertRequestFunc(ParsedOp op) {
|
private LongFunction<UpsertRequest.Builder> createUpsertRequestFunc(ParsedOp op) {
|
||||||
|
@ -17,10 +17,9 @@
|
|||||||
package io.nosqlbench.adapter.pinecone.ops;
|
package io.nosqlbench.adapter.pinecone.ops;
|
||||||
|
|
||||||
import io.nosqlbench.engine.api.templating.ParsedOp;
|
import io.nosqlbench.engine.api.templating.ParsedOp;
|
||||||
import io.pinecone.PineconeException;
|
import io.pinecone.PineconeConnection;
|
||||||
import io.pinecone.proto.DeleteRequest;
|
import io.pinecone.proto.DeleteRequest;
|
||||||
import io.pinecone.proto.DeleteResponse;
|
import io.pinecone.proto.DeleteResponse;
|
||||||
import io.pinecone.PineconeConnection;
|
|
||||||
import org.apache.logging.log4j.LogManager;
|
import org.apache.logging.log4j.LogManager;
|
||||||
import org.apache.logging.log4j.Logger;
|
import org.apache.logging.log4j.Logger;
|
||||||
|
|
||||||
@ -44,6 +43,6 @@ public class PineconeDeleteOp extends PineconeOp {
|
|||||||
@Override
|
@Override
|
||||||
public void run() {
|
public void run() {
|
||||||
DeleteResponse response = connection.getBlockingStub().delete(request);
|
DeleteResponse response = connection.getBlockingStub().delete(request);
|
||||||
logger.debug("Pincecone delete request successful: " + response.toString());
|
logger.debug("Pinecone delete request successful: " + response.toString());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -43,6 +43,6 @@ public class PineconeUpdateOp extends PineconeOp {
|
|||||||
@Override
|
@Override
|
||||||
public void run() {
|
public void run() {
|
||||||
UpdateResponse response = connection.getBlockingStub().update(request);
|
UpdateResponse response = connection.getBlockingStub().update(request);
|
||||||
logger.debug("UpdateResponse succesful: " + response.toString());
|
logger.debug("UpdateResponse successful: " + response.toString());
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -0,0 +1,14 @@
|
|||||||
|
scenarios:
|
||||||
|
default:
|
||||||
|
query: run driver=pinecone tags==blocks:query cycles=10
|
||||||
|
|
||||||
|
blocks:
|
||||||
|
query:
|
||||||
|
ops:
|
||||||
|
op1:
|
||||||
|
query: "test-index"
|
||||||
|
vector: "0.8602578079921012,0.12103044768221516,0.7737329191858439,0.4521093269320254,0.29351661477669416,0.4261807015226558,0.14131665592103335,0.882370813029422,0.4412833140430886,0.9916525700115515"
|
||||||
|
namespace: "example_namespace"
|
||||||
|
top_k: 1000
|
||||||
|
include_values: true
|
||||||
|
include_metadata: true
|
@ -0,0 +1,37 @@
|
|||||||
|
scenarios:
|
||||||
|
default:
|
||||||
|
query: run driver=pinecone tags==blocks:query cycles=10
|
||||||
|
|
||||||
|
blocks:
|
||||||
|
query:
|
||||||
|
ops:
|
||||||
|
op1:
|
||||||
|
query: "test-index"
|
||||||
|
top_k: 3
|
||||||
|
query_vectors:
|
||||||
|
- id: "1"
|
||||||
|
values: "0.8602578079921012,0.12103044768221516,0.7737329191858439,0.4521093269320254,0.29351661477669416,0.4261807015226558,0.14131665592103335,0.882370813029422,0.4412833140430886,0.9916525700115515"
|
||||||
|
namespace: "example_namespace"
|
||||||
|
top_k: 3
|
||||||
|
filter:
|
||||||
|
filterfield: "category"
|
||||||
|
operator: "$eq"
|
||||||
|
comparator: "movies"
|
||||||
|
- id: "2"
|
||||||
|
values:
|
||||||
|
- 0.8602578079921012
|
||||||
|
- 0.12103044768221516
|
||||||
|
- 0.7737329191858439
|
||||||
|
- 0.4521093269320254
|
||||||
|
- 0.29351661477669416
|
||||||
|
- 0.4261807015226558
|
||||||
|
- 0.14131665592103335
|
||||||
|
- 0.882370813029422
|
||||||
|
- 0.4412833140430886
|
||||||
|
- 0.9916525700115515
|
||||||
|
namespace: "example_namespace"
|
||||||
|
top_k: 3
|
||||||
|
filter:
|
||||||
|
filterfield: "category"
|
||||||
|
operator: "$eq"
|
||||||
|
comparator: "sports"
|
@ -9,7 +9,7 @@ blocks:
|
|||||||
query: "test-index"
|
query: "test-index"
|
||||||
vector: "0.8602578079921012,0.12103044768221516,0.7737329191858439,0.4521093269320254,0.29351661477669416,0.4261807015226558,0.14131665592103335,0.882370813029422,0.4412833140430886,0.9916525700115515"
|
vector: "0.8602578079921012,0.12103044768221516,0.7737329191858439,0.4521093269320254,0.29351661477669416,0.4261807015226558,0.14131665592103335,0.882370813029422,0.4412833140430886,0.9916525700115515"
|
||||||
namespace: "example_namespace"
|
namespace: "example_namespace"
|
||||||
topk: 10
|
top_k: 10
|
||||||
include_values: true
|
include_values: true
|
||||||
include_metadata: true
|
include_metadata: true
|
||||||
filter:
|
filter:
|
||||||
|
@ -7,4 +7,10 @@ blocks:
|
|||||||
ops:
|
ops:
|
||||||
op1:
|
op1:
|
||||||
describeindexstats: "test-index"
|
describeindexstats: "test-index"
|
||||||
#filter: ""
|
filter:
|
||||||
|
filterfield: "colors"
|
||||||
|
operator: "$in"
|
||||||
|
comparator:
|
||||||
|
- "blue"
|
||||||
|
- "red"
|
||||||
|
- "green"
|
||||||
|
@ -0,0 +1,21 @@
|
|||||||
|
scenarios:
|
||||||
|
default:
|
||||||
|
upsert: run driver=pinecone tags==block:upsert cycles===TEMPLATE(upsert-cycles,10000000) threads=auto
|
||||||
|
query: run driver=pinecone tags==block:query cycles===TEMPLATE(query-cycles,10000000) threads=auto
|
||||||
|
|
||||||
|
bindings:
|
||||||
|
id: Mod(<<keycount:1000000000>>); ToString() -> String
|
||||||
|
vector_value: ListSizedHashed(<<dimensions:5>>,HashRange(0.0f,100.0f) -> float)
|
||||||
|
city: Cities()
|
||||||
|
|
||||||
|
blocks:
|
||||||
|
upsert:
|
||||||
|
ops:
|
||||||
|
op1:
|
||||||
|
upsert: "test-index"
|
||||||
|
namespace: "example_namespace"
|
||||||
|
upsert_vectors:
|
||||||
|
- id: "{id}"
|
||||||
|
values: "{vector_value}"
|
||||||
|
metadata:
|
||||||
|
city: "{city}"
|
@ -1,6 +1,6 @@
|
|||||||
scenarios:
|
scenarios:
|
||||||
default:
|
default:
|
||||||
query: run driver=pinecone tags==blocks:upsert cycles=1000
|
upsert: run driver=pinecone tags==blocks:upsert cycles=1000
|
||||||
|
|
||||||
blocks:
|
blocks:
|
||||||
upsert:
|
upsert:
|
||||||
@ -10,7 +10,17 @@ blocks:
|
|||||||
namespace: "example_namespace"
|
namespace: "example_namespace"
|
||||||
upsert_vectors:
|
upsert_vectors:
|
||||||
- id: "item_0"
|
- id: "item_0"
|
||||||
values: "0.8238042071878214,0.6579519242642494,0.2772098082162267,0.11434681368630942,0.7496610470760962,0.08914691629812001,0.13249049306247204,0.8875869003282576,0.21615924382746318,0.8454796243176312"
|
values:
|
||||||
|
- 0.8238042071878214
|
||||||
|
- 0.6579519242642494
|
||||||
|
- 0.2772098082162267
|
||||||
|
- 0.11434681368630942
|
||||||
|
- 0.7496610470760962
|
||||||
|
- 0.08914691629812001
|
||||||
|
- 0.13249049306247204
|
||||||
|
- 0.8875869003282576
|
||||||
|
- 0.21615924382746318
|
||||||
|
- 0.8454796243176312
|
||||||
metadata:
|
metadata:
|
||||||
category: "sports"
|
category: "sports"
|
||||||
colors: "blue,red,green"
|
colors: "blue,red,green"
|
||||||
|
@ -198,7 +198,10 @@ public class PineconeOpMapperTest {
|
|||||||
index: "test-index"
|
index: "test-index"
|
||||||
upsert_vectors:
|
upsert_vectors:
|
||||||
- id: 1
|
- id: 1
|
||||||
values: "1.0,2.0,3.0"
|
values:
|
||||||
|
- 1.0
|
||||||
|
- 2.0
|
||||||
|
- 3.0
|
||||||
sparse_values:
|
sparse_values:
|
||||||
indices: "1,2,3"
|
indices: "1,2,3"
|
||||||
values: "4.0,5.0,6.0"
|
values: "4.0,5.0,6.0"
|
||||||
@ -206,7 +209,10 @@ public class PineconeOpMapperTest {
|
|||||||
key1: "val1"
|
key1: "val1"
|
||||||
key2: 2
|
key2: 2
|
||||||
- id: 2
|
- id: 2
|
||||||
values: "7.0,8.0,9.0"
|
values:
|
||||||
|
- 7.0
|
||||||
|
- 8.0
|
||||||
|
- 9.0
|
||||||
sparse_values:
|
sparse_values:
|
||||||
indices: "4,5,6"
|
indices: "4,5,6"
|
||||||
values: "1.1,2.2,3.3"
|
values: "1.1,2.2,3.3"
|
||||||
|
Loading…
Reference in New Issue
Block a user