mirror of
https://github.com/nosqlbench/nosqlbench.git
synced 2026-07-29 23:58:23 -05:00
Merge pull request #1280 from nosqlbench/pc-driver-adapter
Pc driver adapter
This commit is contained in:
@@ -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
|
||||
@@ -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>
|
||||
|
||||
+44
@@ -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();
|
||||
}
|
||||
|
||||
}
|
||||
+116
@@ -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();
|
||||
}
|
||||
|
||||
|
||||
}
|
||||
+88
@@ -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));
|
||||
}
|
||||
}
|
||||
+93
@@ -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));
|
||||
}
|
||||
}
|
||||
+40
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
+198
@@ -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());
|
||||
}
|
||||
}
|
||||
+173
@@ -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());
|
||||
}
|
||||
}
|
||||
+150
@@ -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());
|
||||
}
|
||||
}
|
||||
+49
@@ -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());
|
||||
}
|
||||
}
|
||||
+57
@@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
+55
@@ -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;
|
||||
}
|
||||
|
||||
}
|
||||
+26
@@ -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
|
||||
}
|
||||
+56
@@ -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());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
+48
@@ -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());
|
||||
}
|
||||
}
|
||||
+48
@@ -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
|
||||
|
||||
|
||||
```
|
||||
+203
@@ -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);
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -119,6 +119,12 @@
|
||||
<version>${revision}</version>
|
||||
</dependency>
|
||||
|
||||
<dependency>
|
||||
<groupId>io.nosqlbench</groupId>
|
||||
<artifactId>adapter-pinecone</artifactId>
|
||||
<version>${revision}</version>
|
||||
</dependency>
|
||||
|
||||
</dependencies>
|
||||
|
||||
<build>
|
||||
|
||||
Reference in New Issue
Block a user