Merge pull request #1942 from nosqlbench/ms/qdrant

Initial Qdrant driver implementation
nosqlbench · May 10, 2024 · 3db1504 · 3db1504
2 parents e22c44e + e253412
commit 3db1504
Show file tree

Hide file tree

Showing 41 changed files with 2,310 additions and 55 deletions.
diff --git a/.run/nb --list-drivers.run.xml b/.run/nb --list-drivers.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="nb --list-drivers" type="JarApplication" folderName="Common">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5-5.21.1-SNAPSHOT-jar-with-dependencies.jar" />
+    <option name="PROGRAM_PARAMETERS" value="--list-drivers --show-stacktraces" />
+    <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/local/" />
+    <option name="ALTERNATIVE_JRE_PATH" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/.run/nb --list-scenarios.run.xml b/.run/nb --list-scenarios.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="nb --list-scenarios" type="JarApplication" folderName="Common">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5-5.21.1-SNAPSHOT-jar-with-dependencies.jar" />
+    <option name="PROGRAM_PARAMETERS" value="--list-scenarios --show-stacktraces" />
+    <option name="WORKING_DIRECTORY" value="$PROJECT_DIR$/local/" />
+    <option name="ALTERNATIVE_JRE_PATH" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/.run/qdrant_delete_collection_glove_25.run.xml b/.run/qdrant_delete_collection_glove_25.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="qdrant_delete_collection_glove_25" type="JarApplication" folderName="Qdrant">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5.jar" />
+    <option name="PROGRAM_PARAMETERS" value="qdrant_vectors_live qdrant_vectors.delete_collection dimensions=25 testsize=10000 trainsize=1183514 dataset=glove-25-angular filetype=hdf5 collection=glove_25 similarity_function=1 qdranthost=ded78a51-8370-47d8-adb0-6147f0fcbba2.us-east4-0.gcp.cloud.qdrant.io token_file=./apikey grpc_port=6334 --progress console:1s -v --add-labels &quot;dimensions:25,dataset=glove-25&quot; --add-labels=&quot;target:qdrant,instance:vectors,vendor:qdrant_v191&quot; --report-prompush-to https://vector-perf.feat.apps.paas.datastax.com:8427/api/v1/import/prometheus/metrics/job/nosqlbench/instance/vectors --annotators &quot;[{'type':'log','level':'info'},{'type':'grafana','baseurl':'https://vector-perf.feat.apps.paas.datastax.com/'}]&quot; --report-interval 10 --show-stacktraces --logs-max 5" />
+    <option name="WORKING_DIRECTORY" value="$ProjectFileDir$/local/qdrant" />
+    <option name="ALTERNATIVE_JRE_PATH" value="jdk21" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/.run/qdrant_schema_collection_glove_25.run.xml b/.run/qdrant_schema_collection_glove_25.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="qdrant_schema_collection_glove_25" type="JarApplication" folderName="Qdrant">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5.jar" />
+    <option name="PROGRAM_PARAMETERS" value="qdrant_vectors_live qdrant_vectors.schema_collection dimensions=25 testsize=10000 trainsize=1183514 dataset=glove-25-angular filetype=hdf5 collection=glove_25 similarity_function=1 qdranthost=ded78a51-8370-47d8-adb0-6147f0fcbba2.us-east4-0.gcp.cloud.qdrant.io token_file=./apikey grpc_port=6334 --progress console:1s -v --add-labels &quot;dimensions:25,dataset=glove-25&quot; --add-labels=&quot;target:qdrant,instance:vectors,vendor:qdrant_v191&quot; --report-prompush-to https://vector-perf.feat.apps.paas.datastax.com:8427/api/v1/import/prometheus/metrics/job/nosqlbench/instance/vectors --annotators &quot;[{'type':'log','level':'info'},{'type':'grafana','baseurl':'https://vector-perf.feat.apps.paas.datastax.com/'}]&quot; --report-interval 10 --show-stacktraces --logs-max 5" />
+    <option name="WORKING_DIRECTORY" value="$ProjectFileDir$/local/qdrant" />
+    <option name="ALTERNATIVE_JRE_PATH" value="jdk21" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/.run/qdrant_search_points_glove_25.run.xml b/.run/qdrant_search_points_glove_25.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="qdrant_search_points_glove_25" type="JarApplication" folderName="Qdrant">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5.jar" />
+    <option name="PROGRAM_PARAMETERS" value="qdrant_vectors_live qdrant_vectors.search_points dimensions=25 testsize=10000 trainsize=1183514 dataset=glove-25-angular filetype=hdf5 collection=glove_25 similarity_function=1 qdranthost=ded78a51-8370-47d8-adb0-6147f0fcbba2.us-east4-0.gcp.cloud.qdrant.io token_file=./apikey grpc_port=6334 --progress console:1s -v --add-labels &quot;dimensions:25,dataset=glove-25&quot; --add-labels=&quot;target:qdrant,instance:vectors,vendor:qdrant_v191&quot; --report-prompush-to https://vector-perf.feat.apps.paas.datastax.com:8427/api/v1/import/prometheus/metrics/job/nosqlbench/instance/vectors --annotators &quot;[{'type':'log','level':'info'},{'type':'grafana','baseurl':'https://vector-perf.feat.apps.paas.datastax.com/'}]&quot; --report-interval 10 --report-csv-to metrics --show-stacktraces --logs-max 5" />
+    <option name="WORKING_DIRECTORY" value="$ProjectFileDir$/local/qdrant" />
+    <option name="ALTERNATIVE_JRE_PATH" value="jdk21" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/.run/qdrant_upsert_points_glove_25.run.xml b/.run/qdrant_upsert_points_glove_25.run.xml
@@ -0,0 +1,14 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="qdrant_upsert_points_glove_25" type="JarApplication" folderName="Qdrant">
+    <extension name="software.aws.toolkits.jetbrains.core.execution.JavaAwsConnectionExtension">
+      <option name="credential" />
+      <option name="region" />
+      <option name="useCurrentConnection" value="false" />
+    </extension>
+    <option name="JAR_PATH" value="$PROJECT_DIR$/nb5/target/nb5.jar" />
+    <option name="PROGRAM_PARAMETERS" value="qdrant_vectors_live qdrant_vectors.rampup dimensions=25 testsize=10000 trainsize=1183514 dataset=glove-25-angular filetype=hdf5 collection=glove_25 similarity_function=1 qdranthost=ded78a51-8370-47d8-adb0-6147f0fcbba2.us-east4-0.gcp.cloud.qdrant.io token_file=./apikey grpc_port=6334 --progress console:1s -v --add-labels &quot;dimensions:25,dataset=glove-25&quot; --add-labels=&quot;target:qdrant,instance:vectors,vendor:qdrant_v191&quot; --report-prompush-to https://vector-perf.feat.apps.paas.datastax.com:8427/api/v1/import/prometheus/metrics/job/nosqlbench/instance/vectors --annotators &quot;[{'type':'log','level':'info'},{'type':'grafana','baseurl':'https://vector-perf.feat.apps.paas.datastax.com/'}]&quot; --report-interval 10 --show-stacktraces --logs-max 5" />
+    <option name="WORKING_DIRECTORY" value="$ProjectFileDir$/local/qdrant" />
+    <option name="ALTERNATIVE_JRE_PATH" value="jdk21" />
+    <method v="2" />
+  </configuration>
+</component>
diff --git a/mvn-defaults/pom.xml b/mvn-defaults/pom.xml
@@ -41,6 +41,8 @@
 
         <PROG>nb5</PROG>
         <maven.plugin.validation>VERBOSE</maven.plugin.validation>
+
+        <jacoco.version>0.8.12</jacoco.version>
     </properties>
 
     <name>${project.artifactId}</name>
@@ -549,7 +551,7 @@
             <plugin>
                 <groupId>org.jacoco</groupId>
                 <artifactId>jacoco-maven-plugin</artifactId>
-                <version>0.8.10</version>
+                <version>${jacoco.version}</version>
                 <executions>
                     <execution>
                         <id>prepare-agent</id>
@@ -782,7 +784,7 @@
                 <plugin>
                     <groupId>org.jacoco</groupId>
                     <artifactId>jacoco-maven-plugin</artifactId>
-                    <version>0.8.10</version>
+                    <version>${jacoco.version}</version>
                 </plugin>
                 <plugin>
                     <groupId>org.apache.maven.plugins</groupId>

diff --git a/nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/MilvusSpace.java b/nb-adapters/adapter-milvus/src/main/java/io/nosqlbench/adapter/milvus/MilvusSpace.java
@@ -45,7 +45,7 @@ public class MilvusSpace implements AutoCloseable {
 
     protected MilvusServiceClient client;
 
-    private final Map<String, ConnectParam> connections = new HashMap<>();
+//    private final Map<String, ConnectParam> connections = new HashMap<>();
 
     /**
      * Create a new MilvusSpace Object which stores all stateful contextual information needed to interact

diff --git a/...apter-milvus/src/main/java/io/nosqlbench/adapter/milvus/ops/MilvusCreateCollectionOp.java b/...apter-milvus/src/main/java/io/nosqlbench/adapter/milvus/ops/MilvusCreateCollectionOp.java
@@ -24,7 +24,7 @@
 
 public class MilvusCreateCollectionOp extends MilvusBaseOp<CreateCollectionParam> {
     /**
-     * Create a new {@link ParsedOp} encapsulating a call to the Milvus/Zilliz client delete method
+     * Create a new {@link ParsedOp} encapsulating a call to the Milvus/Zilliz client create method.
      *
      * @param client    The associated {@link MilvusServiceClient} used to communicate with the database
      * @param request   The {@link CreateCollectionParam} built for this operation

diff --git a/nb-adapters/adapter-milvus/src/main/resources/milvus.md b/nb-adapters/adapter-milvus/src/main/resources/milvus.md
@@ -9,12 +9,15 @@ https://github.com/milvus-io/milvus-sdk-java.
 The following parameters must be supplied to the adapter at runtime in order to successfully connect to an
 instance of the Milvus/Zilliz database:
 
-* token - In order to use the pinecone database you must have an account. Once the account is created you can [request
+* `token` - In order to use the Milvus/Zilliz database you must have an account. Once the account is created you
+  can [request
   an api key/token](https://milvus.io/docs/users_and_roles.md#Users-and-Roles). This key will need to be provided any
-  time a
-  database connection is desired.
-* uri - When an Index is created in the database the uri must be specified as well. The adapter will
-  use the default value of localhost:19530 if none is provided at runtime.
+  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 an index is created in the database the URI/endpoint must be specified as well. The adapter will
+  use the default value of `localhost:19530` if none is provided at runtime.
+* `database_name` or `database` - the name of the database to use. For Zilliz, only `default` is supported.
 
 ## Op Templates
 

diff --git a/nb-adapters/adapter-qdrant/pom.xml b/nb-adapters/adapter-qdrant/pom.xml
@@ -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>
diff --git a/...dapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantAdapterUtils.java b/...dapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantAdapterUtils.java
@@ -0,0 +1,67 @@
+/*
+ * 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.qdrant.client.grpc.Points.ScoredPoint;
+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[] searchPointsResponseIdNumToIntArray(List<ScoredPoint> response) {
+        return response.stream().mapToInt(r -> ((Number) r.getId().getNum()).intValue()).toArray();
+    }
+}
diff --git a/...apters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapter.java b/...apters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapter.java
@@ -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());
+    }
+}
diff --git a/.../adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapterLoader.java b/.../adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantDriverAdapterLoader.java
@@ -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);
+    }
+}