IntenseFinance/nosqlbench
3
0
Fork 0
mirror of https://github.com/nosqlbench/nosqlbench.git synced 2024-11-21 16:27:51 -06:00
Code Issues Packages Projects Releases Wiki Activity

Initial draft skeleton for Qdrant driver adapter

Browse Source
This commit is contained in:
Madhavan Sridharan 2024-04-24 09:23:58 -04:00
parent 6545156ce3
commit 053dbec8f5
15 changed files with 965 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
nb-adapters/adapter-qdrant/pom.xml Normal file
View File

@ -0,0 +1,62 @@
<?xml version="1.0" encoding="UTF-8"?>
<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-qdrant</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 Qdrant database.
</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.grpc</groupId>
<artifactId>grpc-protobuf</artifactId>
<!-- <version>1.63.0</version> -->
<!-- Trying to match https://github.com/qdrant/java-client/blob/v1.9.0/build.gradle#L80 -->
<version>1.59.0</version>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java-util</artifactId>
<!--<version>3.25.3</version>-->
<!-- Trying to match https://github.com/qdrant/java-client/blob/master/build.gradle#L81 -->
<version>3.24.0</version>
</dependency>
<dependency>
<groupId>com.google.guava</groupId>
<artifactId>guava</artifactId>
<!--<version>33.1.0-jre</version>-->
<!-- Trying to match https://github.com/qdrant/java-client/blob/master/build.gradle#L93 -->
<version>30.1-jre</version>
</dependency>
<dependency>
<groupId>io.qdrant</groupId>
<artifactId>client</artifactId>
<version>1.9.0</version>
</dependency>
</dependencies>
</project>

73
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantAdapterUtils.java Normal file
View File

@ -0,0 +1,73 @@
/*
* Copyright (c) 2020-2024 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.qdrant;
import org.apache.commons.lang3.StringUtils;
import java.util.Arrays;
import java.util.List;
public class QdrantAdapterUtils {
public static final String QDRANT = "qdrant";
public static List<String> splitNames(String input) {
assert StringUtils.isNotBlank(input) && StringUtils.isNotEmpty(input);
return Arrays.stream(input.split("( +| *, *)"))
.filter(StringUtils::isNotBlank)
.toList();
}
public static List<Long> splitLongs(String input) {
assert StringUtils.isNotBlank(input) && StringUtils.isNotEmpty(input);
return Arrays.stream(input.split("( +| *, *)"))
.filter(StringUtils::isNotBlank)
.map(Long::parseLong)
.toList();
}
/**
* Mask the digits in the given string with '*'
*
* @param unmasked The string to mask
* @return The masked string
*/
protected static String maskDigits(String unmasked) {
assert StringUtils.isNotBlank(unmasked) && StringUtils.isNotEmpty(unmasked);
int inputLength = unmasked.length();
StringBuilder masked = new StringBuilder(inputLength);
for (char ch : unmasked.toCharArray()) {
if (Character.isDigit(ch)) {
masked.append("*");
} else {
masked.append(ch);
}
}
return masked.toString();
}
// public static int[] intArrayFromMilvusSearchResults(String fieldName, R<SearchResults> result) {
// SearchResultsWrapper wrapper = new SearchResultsWrapper(result.getData().getResults());
// List<String> fieldData = (List<String>) wrapper.getFieldData(fieldName, 0);
// int[] indices = new int[fieldData.size()];
// for (int i = 0; i < indices.length; i++) {
// indices[i] = Integer.parseInt(fieldData.get(i));
// }
// return indices;
// }
}

54
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapter.java Normal file
View File

@ -0,0 +1,54 @@
/*
* Copyright (c) 2020-2024 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.qdrant;
import io.nosqlbench.adapter.qdrant.ops.QdrantBaseOp;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.activityimpl.uniform.BaseDriverAdapter;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.components.core.NBComponent;
import io.nosqlbench.nb.api.config.standard.NBConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfiguration;
import io.nosqlbench.nb.api.labels.NBLabels;
import java.util.function.Function;
import static io.nosqlbench.adapter.qdrant.QdrantAdapterUtils.QDRANT;
@Service(value = DriverAdapter.class, selector = QDRANT)
public class QdrantDriverAdapter extends BaseDriverAdapter<QdrantBaseOp<?>, QdrantSpace> {
public QdrantDriverAdapter(NBComponent parentComponent, NBLabels labels) {
super(parentComponent, labels);
}
@Override
public OpMapper<QdrantBaseOp<?>> getOpMapper() {
return new QdrantOpMapper(this);
}
@Override
public Function<String, ? extends QdrantSpace> getSpaceInitializer(NBConfiguration cfg) {
return (s) -> new QdrantSpace(s, cfg);
}
@Override
public NBConfigModel getConfigModel() {
return super.getConfigModel().add(QdrantSpace.getConfigModel());
}
}

32
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapterLoader.java Normal file
View File

@ -0,0 +1,32 @@
/*
* Copyright (c) 2020-2024 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.qdrant;
import io.nosqlbench.adapter.diag.DriverAdapterLoader;
import io.nosqlbench.nb.annotations.Service;
import io.nosqlbench.nb.api.components.core.NBComponent;
import io.nosqlbench.nb.api.labels.NBLabels;
import static io.nosqlbench.adapter.qdrant.QdrantAdapterUtils.QDRANT;
@Service(value = DriverAdapterLoader.class, selector = QDRANT)
public class QdrantDriverAdapterLoader implements DriverAdapterLoader {
@Override
public QdrantDriverAdapter load(NBComponent parent, NBLabels childLabels) {
return new QdrantDriverAdapter(parent, childLabels);
}
}

66
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java Normal file
View File

@ -0,0 +1,66 @@
/*
* Copyright (c) 2020-2024 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.qdrant;
import io.nosqlbench.adapter.qdrant.opdispensers.QdrantBaseOpDispenser;
import io.nosqlbench.adapter.qdrant.opdispensers.QdrantCreateCollectionOpDispenser;
import io.nosqlbench.adapter.qdrant.ops.QdrantBaseOp;
import io.nosqlbench.adapter.qdrant.types.QdrantOpType;
import io.nosqlbench.adapters.api.activityimpl.OpDispenser;
import io.nosqlbench.adapters.api.activityimpl.OpMapper;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.nosqlbench.engine.api.templating.TypeAndTarget;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
public class QdrantOpMapper implements OpMapper<QdrantBaseOp<?>> {
private static final Logger logger = LogManager.getLogger(QdrantOpMapper.class);
private final QdrantDriverAdapter adapter;
/**
* Create a new QdrantOpMapper implementing the {@link OpMapper} interface.
*
* @param adapter The associated {@link QdrantDriverAdapter}
*/
public QdrantOpMapper(QdrantDriverAdapter adapter) {
this.adapter = adapter;
}
/**
* Given an instance of a {@link ParsedOp} returns the appropriate {@link QdrantBaseOpDispenser} subclass
*
* @param op The {@link ParsedOp} to be evaluated
* @return The correct {@link QdrantBaseOpDispenser} subclass based on the op type
*/
@Override
public OpDispenser<? extends QdrantBaseOp<?>> apply(ParsedOp op) {
TypeAndTarget<QdrantOpType, String> typeAndTarget = op.getTypeAndTarget(
QdrantOpType.class,
String.class,
"type",
"target"
);
logger.info(() -> "Using '" + typeAndTarget.enumId + "' op type for op template '" + op.getName() + "'");
return switch (typeAndTarget.enumId) {
// case drop_collection -> new QdrantDropCollectionOpDispenser(adapter, op, typeAndTarget.targetFunction);
case create_collection -> new QdrantCreateCollectionOpDispenser(adapter, op, typeAndTarget.targetFunction);
// default -> throw new RuntimeException("Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " +
// "mapping parsed op " + op);
};
}
}

137
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantSpace.java Normal file
View File

@ -0,0 +1,137 @@
/*
* Copyright (c) 2020-2024 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.qdrant;
import io.nosqlbench.nb.api.config.standard.ConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfigModel;
import io.nosqlbench.nb.api.config.standard.NBConfiguration;
import io.nosqlbench.nb.api.config.standard.Param;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.apache.commons.lang3.math.NumberUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.time.Duration;
/**
* The {@code QdrantSpace} class is a context object which stores all stateful contextual information needed to interact
* with the Qdrant database instance.
*
* @see <a href="https://qdrant.tech/documentation/cloud/quickstart-cloud/">Qdrant cloud quick start guide</a>
* @see <a href="https://qdrant.tech/documentation/quick-start/">Qdrant quick start guide</a>
* @see <a href="https://github.com/qdrant/java-client">Qdrant Java client</a>
*/
public class QdrantSpace implements AutoCloseable {
private final static Logger logger = LogManager.getLogger(QdrantSpace.class);
private final String name;
private final NBConfiguration cfg;
protected QdrantClient client;
// private final Map<String, ConnectParam> connections = new HashMap<>();
/**
* Create a new QdrantSpace Object which stores all stateful contextual information needed to interact
* with the Qdrant database instance.
*
* @param name The name of this space
* @param cfg The configuration ({@link NBConfiguration}) for this nb run
*/
public QdrantSpace(String name, NBConfiguration cfg) {
this.name = name;
this.cfg = cfg;
}
public synchronized QdrantClient getClient() {
if (client == null) {
client = createClient();
}
return client;
}
private QdrantClient createClient() {
String uri = cfg.get("uri");
int grpcPort = cfg.getOptional("grpc_port").map(Integer::parseInt).orElse(6334);
boolean useTls = cfg.getOptional("use_tls").map(Boolean::parseBoolean).orElse(true);
var builder = QdrantGrpcClient.newBuilder(uri, grpcPort, useTls);
var requiredToken = cfg.getOptional("token_file")
.map(Paths::get)
.map(
tokenFilePath -> {
try {
return Files.readAllLines(tokenFilePath).getFirst();
} catch (IOException e) {
String error = "Error while reading token from file:" + tokenFilePath;
logger.error(error, e);
throw new RuntimeException(e);
}
}
).orElseGet(
() -> cfg.getOptional("token")
.orElseThrow(() -> new RuntimeException("You must provide either a token_file or a token to " +
"configure a Qdrant client"))
);
builder = builder.withApiKey(requiredToken);
builder = builder.withTimeout(
Duration.ofMillis(NumberUtils.toInt(cfg.getOptional("timeout_ms").orElse("3000")))
);
logger.info("{}: Creating new Qdrant Client with (masked) token [{}], uri/endpoint [{}]",
this.name, QdrantAdapterUtils.maskDigits(requiredToken), cfg.get("uri").toString());
return new QdrantClient(builder.build());
}
public static NBConfigModel getConfigModel() {
return ConfigModel.of(QdrantSpace.class)
.add(
Param.optional("token_file", String.class, "the file to load the api token from")
)
.add(
Param.defaultTo("token", "qdrant")
.setDescription("the Qdrant api token to use to connect to the database")
)
.add(
Param.defaultTo("uri", "localhost")
.setDescription("the URI endpoint in which the database is running. Do not provide any suffix like https:// here.")
)
.add(
Param.defaultTo("use_tls", true)
.setDescription("whether to use TLS for the connection. Defaults to true.")
)
.add(
Param.defaultTo("timeout_ms", 3000)
.setDescription("sets the timeout in milliseconds for all requests. Defaults to 3000ms.")
)
.add(
Param.defaultTo("grpc_port", 6443)
.setDescription("the port to use for the gRPC connection. Defaults to 6334.")
)
.asReadOnly();
}
@Override
public void close() throws Exception {
if (client != null) {
client.close();
}
}
}

64
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantBaseOpDispenser.java Normal file
View File

@ -0,0 +1,64 @@
/*
* Copyright (c) 2020-2024 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.qdrant.opdispensers;
import io.nosqlbench.adapter.qdrant.QdrantDriverAdapter;
import io.nosqlbench.adapter.qdrant.QdrantSpace;
import io.nosqlbench.adapter.qdrant.ops.QdrantBaseOp;
import io.nosqlbench.adapters.api.activityimpl.BaseOpDispenser;
import io.nosqlbench.adapters.api.activityimpl.uniform.DriverAdapter;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.qdrant.client.QdrantClient;
import java.util.function.LongFunction;
public abstract class QdrantBaseOpDispenser<T> extends BaseOpDispenser<QdrantBaseOp<T>, QdrantSpace> {
protected final LongFunction<QdrantSpace> qdrantSpaceFunction;
protected final LongFunction<QdrantClient> clientFunction;
private final LongFunction<? extends QdrantBaseOp<T>> opF;
private final LongFunction<T> paramF;
protected QdrantBaseOpDispenser(QdrantDriverAdapter adapter, ParsedOp op, LongFunction<String> targetF) {
super((DriverAdapter)adapter, op);
this.qdrantSpaceFunction = adapter.getSpaceFunc(op);
this.clientFunction = (long l) -> this.qdrantSpaceFunction.apply(l).getClient();
this.paramF = getParamFunc(this.clientFunction,op,targetF);
this.opF = createOpFunc(paramF, this.clientFunction, op, targetF);
}
protected QdrantDriverAdapter getDriverAdapter() {
return (QdrantDriverAdapter) adapter;
}
public abstract LongFunction<T> getParamFunc(
LongFunction<QdrantClient> clientF,
ParsedOp op,
LongFunction<String> targetF
);
public abstract LongFunction<QdrantBaseOp<T>> createOpFunc(
LongFunction<T> paramF,
LongFunction<QdrantClient> clientF,
ParsedOp op,
LongFunction<String> targetF
);
@Override
public QdrantBaseOp<T> getOp(long value) {
return opF.apply(value);
}
}

153
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantCreateCollectionOpDispenser.java Normal file
View File

@ -0,0 +1,153 @@
/*
* Copyright (c) 2020-2020-2024 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.qdrant.opdispensers;
import io.nosqlbench.adapter.qdrant.QdrantDriverAdapter;
import io.nosqlbench.adapter.qdrant.ops.QdrantBaseOp;
import io.nosqlbench.adapter.qdrant.ops.QdrantCreateCollectionOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections.CreateCollection;
import io.qdrant.client.grpc.Collections.Distance;
import io.qdrant.client.grpc.Collections.VectorParams;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.HashMap;
import java.util.Map;
import java.util.function.LongFunction;
public class QdrantCreateCollectionOpDispenser extends QdrantBaseOpDispenser<CreateCollection> {
private static final Logger logger = LogManager.getLogger(QdrantCreateCollectionOpDispenser.class);
/**
* Create a new QdrantCreateCollectionOpDispenser subclassed from {@link QdrantBaseOpDispenser}.
*
* @param adapter The associated {@link QdrantDriverAdapter}
* @param op The {@link ParsedOp} encapsulating the activity for this cycle
* @param targetFunction A LongFunction that returns the specified Qdrant Index for this Op
* @see <a href="https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_collection">Qdrant Create Collection</a>.
*/
public QdrantCreateCollectionOpDispenser(QdrantDriverAdapter adapter,
ParsedOp op,
LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
}
@Override
public LongFunction<CreateCollection> getParamFunc(
LongFunction<QdrantClient> clientF,
ParsedOp op,
LongFunction<String> targetF
) {
LongFunction<CreateCollection.Builder> ebF =
l -> CreateCollection.newBuilder().setCollectionName(targetF.apply(l));
// LongFunction<VectorParams.Builder> ebF =
// l -> CreateCollectionParam.newBuilder().withCollectionName(targetF.apply(l));
Map<String, VectorParams> namedVectorsMap = buildNamedVectorsStruct(
op.getAsSubOps("vectors", ParsedOp.SubOpNaming.SubKey)
);
// ebF = op.enhanceFuncOptionally(ebF, "shards_num", Number.class,
// (VectorParams.Builder b, Number n) -> b.withShardsNum(n.intValue()));
// ebF = op.enhanceFuncOptionally(ebF, "partition_num", Number.class,
// (CreateCollectionParam.Builder b, Number n) -> b.withPartitionsNum(n.intValue()));
// ebF = op.enhanceFuncOptionally(ebF, "description", String.class,
// VectorParams.Builder::withDescription);
// ebF = op.enhanceEnumOptionally(ebF, "consistency_level",
// ConsistencyLevelEnum.class, CreateCollectionParam.Builder::withConsistencyLevel);
// ebF = op.enhanceFuncOptionally(ebF, "database_name", String.class,
// CreateCollectionParam.Builder::withDatabaseName);
// List<FieldType> fieldTypes = buildFieldTypesStruct(
// op.getAsSubOps("field_types", ParsedOp.SubOpNaming.SubKey)
// );
// TODO - HERE
// final LongFunction<VectorParams.Builder> f = ebF;
// ebF = l -> f.apply(l).withSchema(CollectionSchemaParam.newBuilder().withFieldTypes(fieldTypes).build());
//
// final LongFunction<VectorParams.Builder> lastF = ebF;
// return l -> lastF.apply(l).build();
return l -> ebF.apply(l).build();
}
private Map<String, VectorParams> buildNamedVectorsStruct(Map<String, ParsedOp> namedVectorsData) {
Map<String, VectorParams> namedVectors = new HashMap<>();
namedVectorsData.forEach((name, fieldspec) -> {
VectorParams.Builder builder = VectorParams.newBuilder();
// TODO - these are mandatory items; see how to achieve this.
fieldspec.getOptionalStaticConfig("distance", Distance.class)
.ifPresent(builder::setDistance);
fieldspec.getOptionalStaticConfig("size", Number.class)
.ifPresent((Number n) -> builder.setSize(n.intValue()));
namedVectors.put(name, builder.build());
});
return namedVectors;
}
// https://qdrant.tech/documentation/concepts/collections/#create-a-collection
@Override
public LongFunction<QdrantBaseOp<CreateCollection>> createOpFunc(
LongFunction<CreateCollection> paramF,
LongFunction<QdrantClient> clientF,
ParsedOp op,
LongFunction<String> targetF
) {
return l -> new QdrantCreateCollectionOp(clientF.apply(l), paramF.apply(l));
}
/**
* Function to build the {@link FieldType}s for the {@link VectorParams}.
*
* @param fieldTypesData The static map of config data from the create collection request
* @return a list of static field types
*/
// private List<FieldType> buildFieldTypesStruct(Map<String, ParsedOp> fieldTypesData) {
// List<FieldType> fieldTypes = new ArrayList<>();
// fieldTypesData.forEach((name, fieldspec) -> {
// FieldType.Builder builder = FieldType.newBuilder()
// .withName(name);
//
// fieldspec.getOptionalStaticValue("primary_key", Boolean.class)
// .ifPresent(builder::withPrimaryKey);
// fieldspec.getOptionalStaticValue("auto_id", Boolean.class)
// .ifPresent(builder::withAutoID);
// fieldspec.getOptionalStaticConfig("max_length", Number.class)
// .ifPresent((Number n) -> builder.withMaxLength(n.intValue()));
// fieldspec.getOptionalStaticConfig("max_capacity", Number.class)
// .ifPresent((Number n) -> builder.withMaxCapacity(n.intValue()));
// fieldspec.getOptionalStaticValue(List.of("partition_key", "partition"), Boolean.class)
// .ifPresent(builder::withPartitionKey);
// fieldspec.getOptionalStaticValue("dimension", Number.class)
// .ifPresent((Number n) -> builder.withDimension(n.intValue()));
// fieldspec.getOptionalStaticConfig("data_type", String.class)
// .map(DataType::valueOf)
// .ifPresent(builder::withDataType);
// fieldspec.getOptionalStaticConfig("type_params", Map.class)
// .ifPresent(builder::withTypeParams);
// fieldspec.getOptionalStaticConfig("element_type", String.class)
// .map(DataType::valueOf)
// .ifPresent(builder::withElementType);
//
// fieldTypes.add(builder.build());
// });
// return fieldTypes;
// }
}

75
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantBaseOp.java Normal file
View File

@ -0,0 +1,75 @@
/*
* Copyright (c) 2020-2024 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.qdrant.ops;
import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.CycleOp;
import io.qdrant.client.QdrantClient;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import java.util.function.LongFunction;
public abstract class QdrantBaseOp<T> implements CycleOp<Object> {
protected final static Logger logger = LogManager.getLogger(QdrantBaseOp.class);
protected final QdrantClient client;
protected final T request;
protected final LongFunction<Object> apiCall;
public QdrantBaseOp(QdrantClient client, T requestParam) {
this.client = client;
this.request = requestParam;
this.apiCall = this::applyOp;
}
public QdrantBaseOp(QdrantClient client, T requestParam, LongFunction<Object> call) {
this.client = client;
this.request = requestParam;
this.apiCall = call;
}
@Override
public final Object apply(long value) {
logger.trace("applying op: {}", this);
try {
Object result = applyOp(value);
// if (result instanceof R<?> r) {
// var error = r.getException();
// if (error != null) {
// throw error;
// }
// } else {
// logger.warn("Op '" + this.toString() + "' did not return a Result 'R' type." +
// " Exception handling will be bypassed"
// );
// }
return result;
} catch (Exception e) {
RuntimeException rte = (RuntimeException) e;
throw rte;
}
}
public abstract Object applyOp(long value);
@Override
public String toString() {
return "QdrantOp(" + this.request.getClass().getSimpleName() + ")";
}
}

53
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantCreateCollectionOp.java Normal file
View File

@ -0,0 +1,53 @@
/*
* Copyright (c) 2020-2024 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.qdrant.ops;
import com.google.common.util.concurrent.ListenableFuture;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections.CollectionOperationResponse;
import io.qdrant.client.grpc.Collections.CreateCollection;
public class QdrantCreateCollectionOp extends QdrantBaseOp<CreateCollection> {
/**
* Create a new {@link ParsedOp} encapsulating a call to the <b>Qdrant</b> create collection method.
*
* @param client The associated {@link QdrantClient} used to communicate with the database
* @param request The {@link CreateCollection} built for this operation
*/
public QdrantCreateCollectionOp(QdrantClient client, CreateCollection request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
// ListenableFuture<CollectionOperationResponse> response = client.createCollectionAsync(
// CreateCollection.newBuilder()
// .setCollectionName("test")
// .setVectorsConfig(VectorsConfig.newBuilder()
// .setParams(
// VectorParams.newBuilder()
// .setDistance(Distance.Cosine)
// .setSize(25)
// .build()
// ).build()
// ).build()
// );
ListenableFuture<CollectionOperationResponse> response = client.createCollectionAsync(request);
return response;
}
}

21
nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/types/QdrantOpType.java Normal file
View File

@ -0,0 +1,21 @@
/*
* Copyright (c) 2020-2024 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.qdrant.types;
public enum QdrantOpType {
create_collection,
}

53
nb-adapters/adapter-qdrant/src/main/resources/activities/qdrant_vectors_live.yaml Normal file
View File

@ -0,0 +1,53 @@
min_version: 5.21
description: |
This is a template for live vector search testing.
Template Variables:
schema: Install the schema required to run the test
rampup: Measure how long it takes to load a set of embeddings
search: Measure how the system responds to queries while it
is indexing recently ingested data.
search: Run vector search with a set of default (or overridden) parameters
In all of these phases, it is important to instance the metrics with distinct names.
Also, aggregates of recall should include total aggregate as well as a moving average.
scenarios:
qdrant_vectors:
schema_collection: >-
run tags==block:schema_collection
errors===stop
cycles===UNDEF threads===UNDEF
uri=TEMPLATE(qdranthost) token_file=TEMPLATE(token_file)
params:
driver: qdrant
instrument: true
bindings:
row_key: ToString()
row_key_batch: Mul(TEMPLATE(batch_size)L); ListSizedStepped(TEMPLATE(batch_size),long->ToString());
# filetype=hdf5 for TEMPLATE(filetype,hdf5)
test_floatlist_hdf5: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/test");
relevant_indices_hdf5: HdfFileToIntArray("testdata/TEMPLATE(dataset).hdf5", "/neighbors")
distance_floatlist_hdf5: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/distance")
train_floatlist_hdf5: HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/train");
train_floatlist_hdf5_batch: Mul(TEMPLATE(batch_size)L); ListSizedStepped(TEMPLATE(batch_size),HdfFileToFloatList("testdata/TEMPLATE(dataset).hdf5", "/train"));
# filetype=fvec for TEMPLATE(filetype,fvec)
test_floatlist_fvec: FVecReader("testdata/TEMPLATE(dataset)_TEMPLATE(trainsize)_query_vectors.fvec");
relevant_indices_fvec: IVecReader("testdata/TEMPLATE(dataset)_TEMPLATE(trainsize)_indices_query.ivec");
distance_floatlist_fvec: FVecReader("testdata/TEMPLATE(dataset)_TEMPLATE(testsize)_distances_count.fvec",TEMPLATE(dimensions),0);
train_floatlist_fvec: FVecReader("testdata/TEMPLATE(dataset)_TEMPLATE(trainsize)_base_vectors.fvec",TEMPLATE(dimensions),0);
train_floatlist_fvec_batch: Mul(TEMPLATE(batch_size,10)L); ListSizedStepped(TEMPLATE(batch_size),FVecReader("testdata/TEMPLATE(dataset)_TEMPLATE(trainsize)_base_vectors.fvec",TEMPLATE(dimensions),0));
blocks:
schema_collection:
ops:
# https://qdrant.github.io/qdrant/redoc/index.html#tag/collections/operation/create_collection
create_col_op:
create_collection: "TEMPLATE(collection)"
# description: "TEMPLATE(desc,a simple qdrant vector collection)"
# consistency_level: "TEMPLATE(write_cl,BOUNDED)"
vectors:
value:
size: TEMPLATE(dimensions,25)
distance: TEMPLATE(similarity_function,cosine)

97
nb-adapters/adapter-qdrant/src/main/resources/qdrant.md Normal file
View File

@ -0,0 +1,97 @@
# qdrant driver adapter
The qdrant driver adapter is a nb adapter for the qdrant driver, an open source Java driver for connecting to and
performing operations on an instance of a Qdrant Vector database. The driver is hosted on GitHub at
https://github.com/qdrant/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 [Qdrant database](https://qdrant.tech/documentation):
* `token` - In order to use the Qdrant database you must have an account. Once the account is created you can [request
an api key/token](https://qdrant.tech/documentation/cloud/authentication/). This key will need to be provided any
time a database connection is desired. Alternatively, the api key can be stored in a file securely and referenced via
the `token_file` config option pointing to the path of the file.
* `uri` - When a collection/index is created in the database the URI (aka endpoint) must be specified as well. The adapter will
use the default value of `localhost:6334` if none is provided at runtime. Remember to *not* provide the `https://`
suffix.
* `grpc_port` - the GRPC port used by the Qdrant database. Defaults to `6334`.
* `use_tls` - option to leverage TLS for the connection. Defaults to `true`.
* `timeout_ms` - sets the timeout in milliseconds for all requests. Defaults to `3000`ms.
## Op Templates
The Qdrant adapter supports [**all operations**](../java/io/nosqlbench/adapter/qdrant/ops) supported by the Java
driver published by Qdrant. The official Qdrant API reference can be found at
https://qdrant.github.io/java-client/io/qdrant/client/package-summary.html
The operations include a full-fledged support for all the APIs available in the Qdrant Java driver.
The following are a couple high level API operations.
# TODO - Below needs to be updated post driver development.
* Create Collection
* Create Index
* Drop Collection
* Drop Index
* Search (vectors)
## Examples
```yaml
ops:
example_create_collection:
create_collection: "example_collection"
description: "https://qdrant.io/api-reference/java/v2.3.x/Collection/createCollection().md"
collection_name: "example_collection"
shards_num: 10
consistency_level: BOUNDED # BOUNDED, SESSION, EVENTUAL
field_types:
field1:
primary_key: true # only for Int64 and VarChar types
description: "field description"
data_type: "Varchar"
# Bool, Int8, Int16, Int32, Int64,
# Float, Double, String, Varchar, BinaryVector, FloatVector
type_param:
example_param1: example_pvalue1
dimension: 1024 # >0
max_length: 1024 # for String only, >0
auto_id: false # Generate primary key?
partition_key: false # Primary key cannot be the partition key too
field2:
primary_key: false
description: "vector column/field"
data_type: "FloatVector"
dimension: 3
# https://qdrant.io/api-reference/java/v2.3.x/Index/dropIndex().md
example_drop_index:
drop_index: "exampe_collection_idx_name"
database_name: "my_database"
collection_name: "example_collection""
# https://qdrant.io/api-reference/java/v2.3.x/Collection/dropCollection().md
example_drop_collection:
drop_collection: "example_collection"
database_name: "my_database"
# https://qdrant.io/api-reference/java/v2.3.x/High-level%20API/insert().md
example_insert_op:
insert: "example_collection_name"
rows:
field1: "row_key"
field2: "[1.2, 3.4, 5.6]"
# https://qdrant.io/api-reference/java/v2.3.x/High-level%20API/search().md
# https://qdrant.io/api-reference/java/v2.3.x/Query%20and%20Search/search().md
example_search:
search: "example_collection"
vector: "[-0.4, 0.3, 0.99]"
metric_type: "COSINE"
out_fields:
- field1
- field2
vector_field_name: "field2"
top_k: 100
consistency_level: "EVENTUALLY"
```

14
nb-adapters/nb-adapters-included/pom.xml
View File

@ -238,6 +238,20 @@
</dependency>
</dependencies>
</profile>
<profile>
<id>adapter-qdrant-include</id>
<activation>
<activeByDefault>false</activeByDefault>
</activation>
<dependencies>
<dependency>
<groupId>io.nosqlbench</groupId>
<artifactId>adapter-qdrant</artifactId>
<version>${revision}</version>
</dependency>
</dependencies>
</profile>
</profiles>
</project>

12
nb-adapters/pom.xml
View File

@ -1,5 +1,5 @@
<!--
~ Copyright (c) 2022-2023 nosqlbench
~ Copyright (c) 2022-2024 nosqlbench
~
~ Licensed under the Apache License, Version 2.0 (the "License");
~ you may not use this file except in compliance with the License.
@ -174,5 +174,15 @@
</modules>
</profile>
<profile>
<id>adapter-qdrant-module</id>
<activation>
<activeByDefault>false</activeByDefault>
</activation>
<modules>
<module>adapter-qdrant</module>
</modules>
</profile>
</profiles>
</project>