Initial version of count points impl

.run/qdrant_search_points_collection_glove_25.run.xml → .run/qdrant_search_points_glove_25.run.xml

@@ -1,5 +1,5 @@
 <component name="ProjectRunConfigurationManager">
-  <configuration default="false" name="qdrant_search_points_collection_glove_25" type="JarApplication" folderName="Qdrant">
+  <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" />

nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java

@@ -62,6 +62,7 @@ public OpDispenser<? extends QdrantBaseOp<?>> apply(ParsedOp op) {
                 new QdrantCreatePayloadIndexOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case search_points -> new QdrantSearchPointsOpDispenser(adapter, op, typeAndTarget.targetFunction);
             case upsert_points -> new QdrantUpsertPointsOpDispenser(adapter, op, typeAndTarget.targetFunction);
+            case count_points -> new QdrantCountPointsOpDispenser(adapter, op, typeAndTarget.targetFunction);
 //            default -> throw new RuntimeException("Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " +
 //                "mapping parsed op " + op);
         };

nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantCountPointsOpDispenser.java

@@ -0,0 +1,51 @@
+/*
+ * 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.ops.QdrantBaseOp;
+import io.nosqlbench.adapter.qdrant.ops.QdrantCountPointsOp;
+import io.nosqlbench.adapters.api.templating.ParsedOp;
+import io.qdrant.client.QdrantClient;
+import io.qdrant.client.grpc.Points.CountPoints;
+
+import java.util.function.LongFunction;
+
+public class QdrantCountPointsOpDispenser extends QdrantBaseOpDispenser<CountPoints> {
+    public QdrantCountPointsOpDispenser(QdrantDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
+        super(adapter, op, targetFunction);
+    }
+
+    @Override
+    public LongFunction<CountPoints> getParamFunc(
+        LongFunction<QdrantClient> clientF, ParsedOp op, LongFunction<String> targetF) {
+        LongFunction<CountPoints.Builder> ebF =
+            l -> CountPoints.newBuilder().setCollectionName(targetF.apply(l));
+
+        final LongFunction<CountPoints.Builder> lastF = ebF;
+        return l -> lastF.apply(l).build();
+    }
+
+    @Override
+    public LongFunction<QdrantBaseOp<CountPoints>> createOpFunc(
+        LongFunction<CountPoints> paramF,
+        LongFunction<QdrantClient> clientF,
+        ParsedOp op,
+        LongFunction<String> targetF) {
+        return l -> new QdrantCountPointsOp(clientF.apply(l), paramF.apply(l));
+    }
+}

nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantCountPointsOp.java

@@ -0,0 +1,46 @@
+/*
+ * 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.qdrant.client.QdrantClient;
+import io.qdrant.client.grpc.Points.CountPoints;
+
+import java.time.Duration;
+import java.util.concurrent.ExecutionException;
+
+public class QdrantCountPointsOp extends QdrantBaseOp<CountPoints> {
+    public QdrantCountPointsOp(QdrantClient client, CountPoints request) {
+        super(client, request);
+    }
+
+    @Override
+    public Object applyOp(long value) {
+        long result;
+        try {
+            result = client.countAsync(
+                request.getCollectionName(),
+                request.getFilter(),
+                request.getExact(),
+                Duration.ofMinutes(5) // opinionated default of 5 minutes
+            ).get();
+            logger.info("Total vector points counted: {}", result);
+        } catch (InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+        return result;
+    }
+}

nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/types/QdrantOpType.java

@@ -25,4 +25,7 @@ public enum QdrantOpType {
     // https://qdrant.tech/documentation/concepts/points/
     // https://qdrant.github.io/qdrant/redoc/index.html#tag/points/operation/upsert_points
     upsert_points,
+    // https://qdrant.github.io/qdrant/redoc/index.html#tag/points/operation/count_points
+    // https://qdrant.tech/documentation/concepts/points/#counting-points
+    count_points,
 }

nb-adapters/adapter-qdrant/src/main/resources/activities/qdrant_vectors_live.yaml

@@ -28,6 +28,11 @@ scenarios:
       errors===warn,counter
       cycles===TEMPLATE(train_cycles,TEMPLATE(trainsize,1000)) threads===TEMPLATE(train_threads,AUTO)
       uri=TEMPLATE(qdranthost) token_file=TEMPLATE(token_file)
+    count_vectors: >-
+      run tags==block:count_vectors
+      errors===stop
+      cycles===UNDEF threads===UNDEF
+      uri=TEMPLATE(qdranthost) token_file=TEMPLATE(token_file)
     search_points: >-
       run tags==block:search_points
       errors===warn,counter
@@ -149,3 +154,9 @@ blocks:
         # https://github.com/qdrant/qdrant/blob/v1.9.0/lib/api/src/grpc/proto/points.proto#L21-L25
         # 0 - All, 1 - Majority, 2 - Quorum
         read_consistency: 2
+
+  count_vectors:
+    ops:
+      count_points_op:
+        count_points: "TEMPLATE(collection)"
+        exact: true