Merge pull request #1280 from nosqlbench/pc-driver-adapter

Pc driver adapter
This commit is contained in:
Jonathan Shook
2023-05-30 00:46:11 -05:00
committed by GitHub
25 changed files with 1912 additions and 0 deletions
+38
View File
@@ -0,0 +1,38 @@
target/
!.mvn/wrapper/maven-wrapper.jar
!**/src/main/**/target/
!**/src/test/**/target/
### IntelliJ IDEA ###
.idea/modules.xml
.idea/jarRepositories.xml
.idea/compiler.xml
.idea/libraries/
*.iws
*.iml
*.ipr
### Eclipse ###
.apt_generated
.classpath
.factorypath
.project
.settings
.springBeans
.sts4-cache
### NetBeans ###
/nbproject/private/
/nbbuild/
/dist/
/nbdist/
/.nb-gradle/
build/
!**/src/main/**/build/
!**/src/test/**/build/
### VS Code ###
.vscode/
### Mac OS ###
.DS_Store
+58
View File
@@ -0,0 +1,58 @@
<!--
~ Copyright (c) 2022-2023 nosqlbench
~
~ Licensed under the Apache License, Version 2.0 (the "License");
~ you may not use this file except in compliance with the License.
~ You may obtain a copy of the License at
~
~ http://www.apache.org/licenses/LICENSE-2.0
~
~ Unless required by applicable law or agreed to in writing, software
~ distributed under the License is distributed on an "AS IS" BASIS,
~ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
~ See the License for the specific language governing permissions and
~ limitations under the License.
-->
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
<modelVersion>4.0.0</modelVersion>
<artifactId>adapter-pinecone</artifactId>
<packaging>jar</packaging>
<parent>
<artifactId>mvn-defaults</artifactId>
<groupId>io.nosqlbench</groupId>
<version>${revision}</version>
<relativePath>../mvn-defaults</relativePath>
</parent>
<name>${project.artifactId}</name>
<description>
An nosqlbench adapter driver module for the Pinecone Vector DB driver
</description>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>nb-annotations</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapters-api</artifactId>
<version>${revision}</version>
<scope>compile</scope>
</dependency>
<dependency>
<groupId>io.pinecone</groupId>
<artifactId>pinecone-client</artifactId>
<version>0.2.3</version>
</dependency>
</dependencies>
</project>
@@ -0,0 +1,44 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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;
import java.util.function.Function;
@Service(value = DriverAdapter.class, selector = "pinecone")
public class PineconeDriverAdapter extends BaseDriverAdapter<PineconeOp, PineconeSpace> {
@Override
public OpMapper<PineconeOp> getOpMapper() {
DriverSpaceCache<? extends PineconeSpace> spaceCache = getSpaceCache();
NBConfiguration adapterConfig = getConfiguration();
return new PineconeOpMapper(this, spaceCache, adapterConfig);
}
@Override
public Function<String, ? extends PineconeSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new PineconeSpace(s,cfg);
}
}
@@ -0,0 +1,85 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.api.config.standard.NBConfiguration;
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;
private final NBConfiguration cfg;
/**
* Create a new PineconeOpMapper implementing the {@link OpMapper} interface.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param spaceCache A cached context Object of thpe {@link PineconeSpace})
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
public PineconeOpMapper(PineconeDriverAdapter adapter,
DriverSpaceCache<? extends PineconeSpace> spaceCache,
NBConfiguration cfg) {
this.adapter = adapter;
this.spaceCache = spaceCache;
this.cfg = cfg;
}
/**
* Given an instance of a {@link ParsedOp} returns the appropriate {@link PineconeOpDispenser} subclass
*
* @param op The ParsedOp to be evaluated
* @return The correct PineconeOpDispenser subclass based on the op type
*/
@Override
public OpDispenser<? extends PineconeOp> apply(ParsedOp op) {
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", "index");
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);
};
}
}
@@ -0,0 +1,103 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.nosqlbench.api.config.standard.Param;
import io.pinecone.PineconeClient;
import io.pinecone.PineconeClientConfig;
import io.pinecone.PineconeConnection;
import io.pinecone.PineconeConnectionConfig;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.HashMap;
import java.util.Map;
public class PineconeSpace {
private final static Logger logger = LogManager.getLogger(PineconeSpace.class);
private final String apiKey;
private final String environment;
private final String projectName;
private final String name;
private final PineconeClient client;
private final PineconeClientConfig config;
private final Map<String,PineconeConnection> connections = new HashMap<>();
/**
* Create a new PineconeSpace Object which stores all stateful contextual information needed to interact
* with the Pinecone database instance.
*
* @param name The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
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()
.withApiKey(apiKey)
.withEnvironment(environment)
.withProjectName(projectName);
logger.info(this.name + ": Creating new Pinecone Client with api key " + apiKey + ", environment "
+ environment + " and project name " + projectName);
this.client = new PineconeClient(config);
}
/**
* 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.
*
* @param index The database index for which a connection is being requested
* @return The {@link PineconeConnection} for this database index
*/
public synchronized PineconeConnection getConnection(String index) {
PineconeConnection connection = connections.get(index);
if (connection == null) {
logger.info("Creating new Pinecone Connection to Index " + index);
PineconeConnectionConfig connectionConfig = new PineconeConnectionConfig().withIndexName(index);
connection = client.connect(connectionConfig);
connections.put(index, connection);
}
return connection;
}
public static NBConfigModel getConfigModel() {
return ConfigModel.of(PineconeSpace.class)
.add(
Param.required("apiKey",String.class)
.setDescription("the Pinecone API key to use to connect to the database")
)
.add(
Param.defaultTo("environment","us-east-1-aws")
.setDescription("the environment in which the desired index is running.")
)
.add(
Param.defaultTo("projectName","default")
.setDescription("the project name associated with the desired index")
)
.asReadOnly();
}
}
@@ -0,0 +1,116 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.PineconeDeleteOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.DeleteRequest;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.*;
import java.util.function.LongFunction;
public class PineconeDeleteOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeDeleteOpDispenser.class);
private final LongFunction<DeleteRequest> deleteRequestFunc;
/**
* Create a new PineconeDeleteOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeDeleteOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
deleteRequestFunc = createDeleteRequestFunction(op);
}
@Override
public PineconeOp apply(long value) {
return new PineconeDeleteOp(pcFunction.apply(value)
.getConnection(targetFunction.apply(value)),
deleteRequestFunc.apply(value));
}
/**
* @param op The ParsedOp used to build the Request
* @return A function that will take a long (the current cycle) and return a Pinecone DeleteRequest
*
* 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.
* For each method in the chain a function is created here and added to the chain of functions
* called at time of instantiation.
*/
private LongFunction<DeleteRequest> createDeleteRequestFunction(ParsedOp op) {
LongFunction<DeleteRequest.Builder> rFunc = l -> DeleteRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<DeleteRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = nFunc.get();
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
Optional<LongFunction<String>> iFunc = op.getAsOptionalFunction("ids", String.class);
if (iFunc.isPresent()) {
LongFunction<DeleteRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = iFunc.get();
LongFunction<List<String>> alf = l -> {
String[] vals = af.apply(l).split(",");
return Arrays.asList(vals);
};
rFunc = l -> finalFunc.apply(l).addAllIds(alf.apply(l));
}
Optional<LongFunction<Boolean>> aFunc = op.getAsOptionalFunction("deleteall", Boolean.class);
if (aFunc.isPresent()) {
LongFunction<DeleteRequest.Builder> finalFunc = rFunc;
LongFunction<Boolean> af = aFunc.get();
rFunc = l -> finalFunc.apply(l).setDeleteAll(af.apply(l));
}
Optional<LongFunction<String>> filterFunction = op.getAsOptionalFunction("filter", String.class);
if (filterFunction.isPresent()) {
LongFunction<DeleteRequest.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.parseInt(filterFields[2])).build()))
.build()).build();
};
rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l));
}
LongFunction<DeleteRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
}
@@ -0,0 +1,88 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.PineconeDescribeIndexStatsOp;
import io.nosqlbench.adapter.pinecone.ops.PineconeOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.DescribeIndexStatsRequest;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Optional;
import java.util.function.LongFunction;
public class PineconeDescribeIndexStatsOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeDescribeIndexStatsOpDispenser.class);
private final LongFunction<DescribeIndexStatsRequest> indexStatsRequestFunc;
/**
* Create a new PineconeDescribeIndexStatsOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeDescribeIndexStatsOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
indexStatsRequestFunc = createDescribeIndexStatsRequestFunction(op);
}
/**
* @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
*
* 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.
* 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
* 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.
*/
private LongFunction<DescribeIndexStatsRequest> createDescribeIndexStatsRequestFunction(ParsedOp op) {
LongFunction<DescribeIndexStatsRequest.Builder> rFunc = l -> DescribeIndexStatsRequest.newBuilder();
Optional<LongFunction<String>> filterFunction = op.getAsOptionalFunction("filter", String.class);
if (filterFunction.isPresent()) {
LongFunction<DescribeIndexStatsRequest.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.parseInt(filterFields[2])).build()))
.build()).build();
};
rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l));
}
LongFunction<DescribeIndexStatsRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
@Override
public PineconeOp apply(long value) {
return new PineconeDescribeIndexStatsOp(pcFunction.apply(value)
.getConnection(targetFunction.apply(value)),
indexStatsRequestFunc.apply(value));
}
}
@@ -0,0 +1,93 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.templating.ParsedOp;
import io.pinecone.proto.FetchRequest;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.function.LongFunction;
public class PineconeFetchOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeFetchOpDispenser.class);
private final LongFunction<FetchRequest> fetchRequestFunc;
/**
* Create a new PineconeFetchOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeFetchOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
fetchRequestFunc = createFetchRequestFunction(op);
}
/**
* @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
*
* 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.
* For each method in the chain a function is created here and added to the chain of functions
* called at time of instantiation.
*/
private LongFunction<FetchRequest> createFetchRequestFunction(ParsedOp op) {
LongFunction<FetchRequest.Builder> rFunc = l -> FetchRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<FetchRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = nFunc.get();
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
Optional<LongFunction<String>> iFunc = op.getAsOptionalFunction("ids", String.class);
if (iFunc.isPresent()) {
LongFunction<FetchRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = iFunc.get();
LongFunction<List<String>> alf = l -> {
String[] vals = af.apply(l).split(",");
return Arrays.asList(vals);
};
rFunc = l -> finalFunc.apply(l).addAllIds(alf.apply(l));
}
LongFunction<FetchRequest.Builder> finalRFunc = rFunc;
return l -> finalRFunc.apply(l).build();
}
@Override
public PineconeOp apply(long value) {
return new PineconeFetchOp(pcFunction.apply(value)
.getConnection(targetFunction.apply(value)),
fetchRequestFunc.apply(value));
}
}
@@ -0,0 +1,40 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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.BaseOpDispenser;
import io.nosqlbench.engine.api.templating.ParsedOp;
import java.util.function.LongFunction;
public abstract class PineconeOpDispenser extends BaseOpDispenser<PineconeOp, PineconeSpace> {
protected final LongFunction<PineconeSpace> pcFunction;
protected final LongFunction<String> targetFunction;
protected PineconeOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op);
this.pcFunction = pcFunction;
this.targetFunction = targetFunction;
}
}
@@ -0,0 +1,198 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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;
import io.nosqlbench.adapter.pinecone.ops.PineconeQueryOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.QueryRequest;
import io.pinecone.proto.QueryVector;
import io.pinecone.proto.SparseValues;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.*;
import java.util.function.LongFunction;
public class PineconeQueryOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeQueryOpDispenser.class);
private final LongFunction<QueryRequest.Builder> queryRequestFunc;
private final LongFunction<Collection<QueryVector>> queryVectorFunc;
/**
* Create a new PineconeQueryOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeQueryOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
queryRequestFunc = createQueryRequestFunc(op);
queryVectorFunc = createQueryVectorFunc(op);
}
/**
* @param op The ParsedOp used to build the Request
* @return A function that will take a long (the current cycle) and return a Pinecone QueryRequest Builder
*
* 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.
* For each method in the chain a function is created here and added to the chain of functions
* called at time of instantiation.
*
* The QueryVector objects used by the QueryRequest as sufficiently sophisticated in their own building process
* that it has 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.
*/
private LongFunction<QueryRequest.Builder> createQueryRequestFunc(ParsedOp op) {
LongFunction<QueryRequest.Builder> rFunc = l -> QueryRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = nFunc.get();
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
Optional<LongFunction<Integer>> tFunc = op.getAsOptionalFunction("topk", Integer.class);
if (tFunc.isPresent()) {
LongFunction<QueryRequest.Builder> finalFunc = rFunc;
LongFunction<Integer> af = tFunc.get();
rFunc = l -> finalFunc.apply(l).setTopK(af.apply(l));
}
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));
}
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));
}
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<>();
for (String val : vals) {
fVals.add(Float.valueOf(val));
}
return fVals;
};
rFunc = l -> finalFunc.apply(l).addAllVector(alf.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.parseInt(filterFields[2])).build()))
.build()).build();
};
rFunc = l -> finalFunc.apply(l).setFilter(builtFilter.apply(l));
}
return rFunc;
}
/**
* @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
*
* This method interrogates the subsection of the ParsedOp defined for QueryVector parameters and constructs
* a list of QueryVectors 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
* with a Builder pattern based on the values contained in the source ParsedOp.
*/
private LongFunction<Collection<QueryVector>> createQueryVectorFunc(ParsedOp op) {
Optional<LongFunction<List>> baseFunc =
op.getAsOptionalFunction("query_vectors", List.class);
return baseFunc.<LongFunction<Collection<QueryVector>>>map(listLongFunction -> l -> {
List<QueryVector> returnVectors = new ArrayList<>();
List<Map<String, Object>> vectors = listLongFunction.apply(l);
for (Map<String, Object> vector : vectors) {
QueryVector.Builder qvb = QueryVector.newBuilder();
String[] rawValues = ((String) vector.get("values")).split(",");
ArrayList<Float> floatValues = new ArrayList<>();
for (String val : rawValues) {
floatValues.add(Float.valueOf(val));
}
qvb.addAllValues(floatValues);
qvb.setNamespace((String) vector.get("namespace"));
if (vector.containsKey("top_k")) {
qvb.setTopK((Integer) vector.get("top_k"));
}
if (vector.containsKey("filter")) {
String[] rawVals = ((String)vector.get("filter")).split(" ");
qvb.setFilter(Struct.newBuilder().putFields(rawVals[0],
Value.newBuilder().setStructValue(Struct.newBuilder().putFields(rawVals[1],
Value.newBuilder().setNumberValue(Integer.parseInt(rawVals[2])).build()))
.build()).build());
}
if (vector.containsKey("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()
.addAllValues(floatValues)
.addAllIndices(intValues)
.build());
}
returnVectors.add(qvb.build());
}
return returnVectors;
}).orElse(null);
}
@Override
public PineconeOp apply(long value) {
QueryRequest.Builder qrb = queryRequestFunc.apply(value);
if (queryVectorFunc != null) {
qrb.addAllQueries(queryVectorFunc.apply(value));
}
return new PineconeQueryOp(pcFunction.apply(value).getConnection(targetFunction.apply(value)), qrb.build());
}
}
@@ -0,0 +1,173 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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;
import io.nosqlbench.adapter.pinecone.ops.PineconeUpdateOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.SparseValues;
import io.pinecone.proto.UpdateRequest;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.LongFunction;
public class PineconeUpdateOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeUpdateOpDispenser.class);
private final LongFunction<UpdateRequest.Builder> updateRequestFunc;
private final LongFunction<Struct> updateMetadataFunc;
private final LongFunction<SparseValues> sparseValuesFunc;
/**
* Create a new PineconeUpdateOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeUpdateOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
updateRequestFunc = createUpdateRequestFunction(op);
updateMetadataFunc = createUpdateMetadataFunction(op);
sparseValuesFunc = createSparseValuesFunction(op);
}
/**
* @param op the ParsedOp from which the SparseValues object will be built
* @return a SparseValues Object to be added to a Pinecone UpdateRequest
*
* 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
* base function returns either the SparseValues Object or null, while the interior function builds the SparseValues
* with a Builder pattern based on the values contained in the source ParsedOp.
*/
private LongFunction<SparseValues> createSparseValuesFunction(ParsedOp op) {
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("sparse_values", Map.class);
return mFunc.<LongFunction<SparseValues>>map(mapLongFunction -> 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()
.addAllValues(floatValues)
.addAllIndices(intValues)
.build();
}).orElse(null);
}
/**
* @param op the ParsedOp from which the Metadata objects will be built
* @return an Metadata Struct to be added to a Pinecone UpdateRequest
*
* 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
* base function returns either the Metadata Struct or null, while the interior function builds the Metadata
* with a Builder pattern based on the values contained in the source ParsedOp.
*/
private LongFunction<Struct> createUpdateMetadataFunction(ParsedOp op) {
Optional<LongFunction<Map>> mFunc = op.getAsOptionalFunction("metadata", Map.class);
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);
metadata_values_map.forEach(stringToValue);
return UpdateRequest.newBuilder().getSetMetadataBuilder().putAllFields(metadata_map).build();
}).orElse(null);
}
/**
* @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
*
* 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.
* For each method in the chain a function is created here and added to the chain of functions
* called at time of instantiation.
*
* 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
* and then added to the build chain by the builder returned by this method.
*/
private LongFunction<UpdateRequest.Builder> createUpdateRequestFunction(ParsedOp op) {
LongFunction<UpdateRequest.Builder> rFunc = l -> UpdateRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<UpdateRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = nFunc.get();
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
Optional<LongFunction<String>> iFunc = op.getAsOptionalFunction("id", String.class);
if (iFunc.isPresent()) {
LongFunction<UpdateRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = iFunc.get();
rFunc = l -> finalFunc.apply(l).setId(af.apply(l));
}
Optional<LongFunction<String>> vFunc = op.getAsOptionalFunction("values", String.class);
if (vFunc.isPresent()) {
LongFunction<UpdateRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = vFunc.get();
LongFunction<ArrayList<Float>> alf = l -> {
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));
}
return rFunc;
}
@Override
public PineconeOp apply(long value) {
UpdateRequest.Builder urb = updateRequestFunc.apply(value);
if (updateMetadataFunc != null) {
urb.setSetMetadata(updateMetadataFunc.apply(value));
}
if (sparseValuesFunc != null) {
urb.setSparseValues(sparseValuesFunc.apply(value));
}
return new PineconeUpdateOp(pcFunction.apply(value).getConnection(targetFunction.apply(value)), urb.build());
}
}
@@ -0,0 +1,150 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
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;
import io.nosqlbench.adapter.pinecone.ops.PineconeUpsertOp;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.SparseValues;
import io.pinecone.proto.UpsertRequest;
import io.pinecone.proto.Vector;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.LongFunction;
public class PineconeUpsertOpDispenser extends PineconeOpDispenser {
private static final Logger logger = LogManager.getLogger(PineconeUpsertOpDispenser.class);
private final LongFunction<UpsertRequest.Builder> upsertRequestFunc;
private final LongFunction<Collection<Vector>> upsertVectorFunc;
/**
* Create a new PineconeUpsertOpDispenser subclassed from {@link PineconeOpDispenser}.
*
* @param adapter The associated {@link PineconeDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param pcFunction A function to return the associated context of this activity (see {@link PineconeSpace})
* @param targetFunction A LongFunction that returns the specified Pinecone Index for this Op
*/
public PineconeUpsertOpDispenser(PineconeDriverAdapter adapter,
ParsedOp op,
LongFunction<PineconeSpace> pcFunction,
LongFunction<String> targetFunction) {
super(adapter, op, pcFunction, targetFunction);
upsertRequestFunc = createUpsertRequestFunc(op);
upsertVectorFunc = createUpsertRequestVectorsFunc(op);
}
/**
* @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
*
* 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
* base function returns either the List of vectors or null, while the interior function builds the vectors
* with a Builder pattern based on the values contained in the source ParsedOp.
*/
private LongFunction<Collection<Vector>> createUpsertRequestVectorsFunc(ParsedOp op) {
Optional<LongFunction<List>> baseFunc =
op.getAsOptionalFunction("upsert_vectors", List.class);
return baseFunc.<LongFunction<Collection<Vector>>>map(listLongFunction -> l -> {
List<Vector> returnVectors = new ArrayList<>();
List<Map<String, Object>> vectors = listLongFunction.apply(l);
for (Map<String, Object> vector : vectors) {
Vector.Builder vb = Vector.newBuilder();
String[] rawValues = ((String) vector.get("values")).split(",");
ArrayList<Float> floatValues = new ArrayList<>();
for (String val : rawValues) {
floatValues.add(Float.valueOf(val));
}
vb.addAllValues(floatValues);
if (vector.containsKey("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()
.addAllValues(floatValues)
.addAllIndices(intValues)
.build());
}
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");
metadata_values_map.forEach(stringToValue);
vb.setMetadata(Struct.newBuilder().putAllFields(metadata_map).build());
}
returnVectors.add(vb.build());
}
return returnVectors;
}).orElse(null);
}
/**
* @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
*
* 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.
* For each method in the chain a function is created here and added to the chain of functions
* called at time of instantiation.
*
* 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
* and then added to the build chain by the builder returned by this method.
*/
private LongFunction<UpsertRequest.Builder> createUpsertRequestFunc(ParsedOp op) {
LongFunction<UpsertRequest.Builder> rFunc = l -> UpsertRequest.newBuilder();
Optional<LongFunction<String>> nFunc = op.getAsOptionalFunction("namespace", String.class);
if (nFunc.isPresent()) {
LongFunction<UpsertRequest.Builder> finalFunc = rFunc;
LongFunction<String> af = nFunc.get();
rFunc = l -> finalFunc.apply(l).setNamespace(af.apply(l));
}
return rFunc;
}
@Override
public PineconeOp apply(long value) {
UpsertRequest.Builder urb = upsertRequestFunc.apply(value);
if (upsertVectorFunc != null) {
urb.addAllVectors(upsertVectorFunc.apply(value));
}
return new PineconeUpsertOp(pcFunction.apply(value).getConnection(targetFunction.apply(value)), urb.build());
}
}
@@ -0,0 +1,49 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.PineconeException;
import io.pinecone.proto.DeleteRequest;
import io.pinecone.proto.DeleteResponse;
import io.pinecone.PineconeConnection;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class PineconeDeleteOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeDeleteOp.class);
private final DeleteRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client delete method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link DeleteRequest} built for this operation
*/
public PineconeDeleteOp(PineconeConnection connection, DeleteRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
DeleteResponse response = connection.getBlockingStub().delete(request);
logger.debug("Pincecone delete request successful: " + response.toString());
}
}
@@ -0,0 +1,57 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.DescribeIndexStatsRequest;
import io.pinecone.proto.DescribeIndexStatsResponse;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.NamespaceSummary;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Map;
public class PineconeDescribeIndexStatsOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeDescribeIndexStatsOp.class);
private final DescribeIndexStatsRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client describeIndexStats method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link DescribeIndexStatsRequest} built for this operation
*/
public PineconeDescribeIndexStatsOp(PineconeConnection connection, DescribeIndexStatsRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
DescribeIndexStatsResponse response = connection.getBlockingStub().describeIndexStats(request);
if (logger.isDebugEnabled()) {
logger.debug("Vector counts:");
for (Map.Entry<String, NamespaceSummary> namespace : response.getNamespacesMap().entrySet()) {
logger.debug(namespace.getKey() + ": " + namespace.getValue().getVectorCount());
}
}
}
}
@@ -0,0 +1,55 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.FetchRequest;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.FetchResponse;
import io.pinecone.proto.Vector;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.Map;
public class PineconeFetchOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeFetchOp.class);
private final FetchRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client fetch method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link FetchRequest} built for this operation
*/
public PineconeFetchOp(PineconeConnection connection, FetchRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
FetchResponse response = connection.getBlockingStub().fetch(request);
if (logger.isDebugEnabled()) {
for (Map.Entry<String, Vector> vectors: response.getVectorsMap().entrySet()) {
logger.debug(vectors.getKey() + ": " + vectors.getValue().toString());
}
}
}
}
@@ -0,0 +1,29 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.activityimpl.uniform.flowtypes.RunnableOp;
import io.pinecone.PineconeConnection;
public abstract class PineconeOp implements RunnableOp {
protected final PineconeConnection connection;
public PineconeOp(PineconeConnection connection) {
this.connection = connection;
}
}
@@ -0,0 +1,26 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
public enum PineconeOpTypes {
query,
update,
upsert,
delete,
describeindexstats,
fetch
}
@@ -0,0 +1,56 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.QueryRequest;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.QueryResponse;
import io.pinecone.proto.ScoredVector;
import io.pinecone.proto.SingleQueryResults;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class PineconeQueryOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeQueryOp.class);
private final QueryRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client query method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link QueryRequest} built for this operation
*/
public PineconeQueryOp(PineconeConnection connection, QueryRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
QueryResponse response = connection.getBlockingStub().query(request);
if (logger.isDebugEnabled()) {
for (SingleQueryResults results : response.getResultsList()) {
for (ScoredVector scored : results.getMatchesList()) {
logger.debug(scored.getId() + ": " + scored.getScore());
}
}
}
}
}
@@ -0,0 +1,48 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.UpdateRequest;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.UpdateResponse;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class PineconeUpdateOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeUpdateOp.class);
private final UpdateRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client update method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link UpdateRequest} built for this operation
*/
public PineconeUpdateOp(PineconeConnection connection, UpdateRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
UpdateResponse response = connection.getBlockingStub().update(request);
logger.debug("UpdateResponse succesful: " + response.toString());
}
}
@@ -0,0 +1,48 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone.ops;
import io.nosqlbench.engine.api.templating.ParsedOp;
import io.pinecone.proto.UpsertRequest;
import io.pinecone.PineconeConnection;
import io.pinecone.proto.UpsertResponse;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class PineconeUpsertOp extends PineconeOp {
private static final Logger logger = LogManager.getLogger(PineconeUpsertOp.class);
private final UpsertRequest request;
/**
* Create a new {@link ParsedOp} encapsulating a call to the Pinecone client fetch method
*
* @param connection The associated {@link PineconeConnection} used to communicate with the database
* @param request The {@link UpsertRequest} built for this operation
*/
public PineconeUpsertOp(PineconeConnection connection, UpsertRequest request) {
super(connection);
this.request = request;
}
@Override
public void run() {
UpsertResponse response = connection.getBlockingStub().upsert(request);
logger.debug("Put " + response.getUpsertedCount() + " vectors into the index");
}
}
@@ -0,0 +1,19 @@
bindings:
vector: Need function to generate a list of float values here
comparator: Uniform(1,1000) -> int
blocks:
query:
ops:
query-op1:
query: "test-index"
vector: {vector}
namespace: "test-namespace"
top_k: 10
filters:
- filter_field: "value"
operator: "$lt"
comparator: {comparator}
include_values: true
include_metadata: true
@@ -0,0 +1,129 @@
# pinecone driver adapter
The pinecone driver adapter is a nb adapter for the pinecone driver, an open source Java driver for connecting to and
performing operations on an instance of a Pinecone Vector database. The driver is hosted on github at
https://github.com/pinecone-io/pinecone-java-client
## activity parameters
The following parameters must be supplied to the adapter at runtime in order to successfully connect to an
instance of the Pinecone database:
* api key - In order to use the pinecone database you must have an account. Once the account is created you can request
an api key. This key will need to be provided any time a database connection is desired.
* environment - When an Index is created in the database the environment must be specified as well. The adapter will
use the default value of us-east-1 if none is provided at runtime.
* project name - A project name must also be provided at time of index creation. This name then needs to be provided
to the adapter at runtime.
## Op Templates
The Pinecone adapter supports all operations supported by the Java driver published by Pinecone.
The official Pinecone API reference can be found at
https://docs.pinecone.io/reference/describe_index_stats_post
The operations include:
* Delete
* DescribeIndexStats
* Fetch
* Query
* Update
* Upsert
## Examples
```yaml
ops:
# A pinecone query op
query-example:
type: query
index: query_index
# The query vector. Use these fields if only querying a single vector. If querying multiple use the
# query_vectors structure below.
vector: my_array_of_floats
namespace: query_namespace
# 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/
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.
include_metadata: boolean
query_vectors:
- id: 1
values: csv_separated_floats
top_k: int_val
namespace: string_val
filter: <field operator compval>
sparse_values:
indices: list_of_ints
values: list_of_floats
- id: 2
values: csv_separated_floats
top_k: int_val
namespace: string_val
filter: <field operator compval>
sparse_values:
indices: list_of_ints
values: list_of_floats
# A delete op
# If specified, the metadata filter here will be used to select the vectors to delete. This is mutually exclusive
# with specifying ids to delete in the ids param or using delete_all=True. delete_all indicates that all vectors
# in the index namespace should be deleted.
delete-example:
type: delete
index: delete_index
namespace: delete_namespace
ids: csv_list_of_vectors_to_delete
deleteall: [true,false]
filter: <field operator compval>
# A describe index stats op. Specify metadata filters to narrow the range of indices described.
describe-index-stats-example:
type: describe-index-stats
index: describe_index
filter: <field operator compval>
# A pinecone fetch op
fetch-example:
fetch: fetch_index
namespace: fetch_namespace
ids: csv_list_of_vectors_to_fetch
# A pinecone update op
update-example:
type: update
index: update_index
id: string_id
values: list_of_floats
namespace: update_namespace
metadata:
key1: val1
key2: val2
key3: val3
sparse_values:
indices: list_of_ints
values: list_of_floats
# A pinecone upsert op
upsert-example:
type: upsert
index: upsert_index
namespace: upsert_namespace
upsert_vectors:
- id: 1
values: csv_separated_floats
sparse_values:
indices: list_of_ints
values: list_of_floats
metadata:
key1: val1
key2: val2
- id: 2
values: csv_separated_floats
sparse_values:
indices: list_of_ints
values: list_of_floats
```
@@ -0,0 +1,203 @@
/*
* Copyright (c) 2023 nosqlbench
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.nosqlbench.adapter.pinecone;
import io.nosqlbench.adapter.pinecone.opdispensers.*;
import io.nosqlbench.adapter.pinecone.ops.*;
import io.nosqlbench.api.config.NBLabeledElement;
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() {
Map<String,String> configMap = Map.of("apiKey","2f55b2f0-670f-4c51-9073-4d37142b761a",
"environment","us-east-1-aws",
"projectName","default");
cfg = PineconeSpace.getConfigModel().apply(configMap);
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);
NBLabeledElement parent = NBLabeledElement.EMPTY;
return new ParsedOp(opTemplate, cfg, List.of(adapter.getPreprocessor()), parent);
}
@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
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);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeQueryOp);
}
@Test
public void testDeleteOpDispenser() {
ParsedOp pop = parsedOpFor("""
ops:
op1:
type: "delete"
index: "test-index"
ids: "1.0,2.0,3.0"
namespace: "test-namespace"
deleteall: true
filter: "value $gt 10"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeDeleteOpDispenser);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeDeleteOp);
}
@Test
public void testDescribeIndexStatsOpDispenser() {
ParsedOp pop = parsedOpFor("""
ops:
op1:
type: "describeindexstats"
index: "test-index"
filter: "value $gt 10"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeDescribeIndexStatsOpDispenser);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeDescribeIndexStatsOp);
}
@Test
public void testFetchOpDispenser() {
ParsedOp pop = parsedOpFor("""
ops:
op1:
type: "fetch"
index: "test-index"
ids: "1.0,2.0,3.0"
namespace: "test-namespace"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeFetchOpDispenser);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeFetchOp);
}
@Test
public void testUpdateOpDispenser() {
ParsedOp pop = parsedOpFor("""
ops:
op1:
type: "update"
index: "test-index"
id: "id"
values: "1.0,2.0,3.0"
namespace: "test_namespace"
metadata:
key1: "val1"
key2: 2
key3: 3
sparse_values:
indices: "1,2,3"
values: "1.1,2.2,3.3"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeUpdateOpDispenser);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeUpdateOp);
}
@Test
public void testUpsertOpDispenser() {
ParsedOp pop = parsedOpFor("""
ops:
op1:
type: "upsert"
index: "test-index"
upsert_vectors:
- id: 1
values: "1.0,2.0,3.0"
sparse_values:
indices: "1,2,3"
values: "4.0,5.0,6.0"
metadata:
key1: "val1"
key2: 2
- id: 2
values: "7.0,8.0,9.0"
sparse_values:
indices: "4,5,6"
values: "1.1,2.2,3.3"
""");
OpDispenser<? extends PineconeOp> dispenser = mapper.apply(pop);
assert(dispenser instanceof PineconeUpsertOpDispenser);
PineconeOp op = dispenser.apply(0);
assert(op instanceof PineconeUpsertOp);
}
}
+6
View File
@@ -119,6 +119,12 @@
<version>${revision}</version>
</dependency>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapter-pinecone</artifactId>
<version>${revision}</version>
</dependency>
</dependencies>
<build>
+1
View File
@@ -111,6 +111,7 @@
<module>adapter-s4j</module>
<module>adapter-kafka</module>
<module>adapter-jdbc</module>
<module>adapter-pinecone</module>
<!-- VIRTDATA MODULES -->
<module>virtdata-api</module>