From 05a53229d6f96ad363d70120b3dddd860326236d Mon Sep 17 00:00:00 2001
From: Madhavan Sridharan <madhavan.sridharan@datastax.com>
Date: Tue, 7 May 2024 11:57:51 -0400
Subject: [PATCH] Initial version of count points impl

---
 ... => qdrant_search_points_glove_25.run.xml} |  2 +-
 .../adapter/qdrant/QdrantOpMapper.java        |  1 +
 .../adapter/qdrant/QdrantSpace.java           |  2 -
 .../QdrantCountPointsOpDispenser.java         | 51 +++++++++++++++++++
 .../QdrantDeleteCollectionOpDispenser.java    |  3 --
 .../qdrant/ops/QdrantCountPointsOp.java       | 46 +++++++++++++++++
 .../adapter/qdrant/types/QdrantOpType.java    |  3 ++
 .../activities/qdrant_vectors_live.yaml       | 11 ++++
 8 files changed, 113 insertions(+), 6 deletions(-)
 rename .run/{qdrant_search_points_collection_glove_25.run.xml => qdrant_search_points_glove_25.run.xml} (88%)
 create mode 100644 nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantCountPointsOpDispenser.java
 create mode 100644 nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantCountPointsOp.java

diff --git a/.run/qdrant_search_points_collection_glove_25.run.xml b/.run/qdrant_search_points_glove_25.run.xml
similarity index 88%
rename from .run/qdrant_search_points_collection_glove_25.run.xml
rename to .run/qdrant_search_points_glove_25.run.xml
index 0c500df43..005845971 100644
--- a/.run/qdrant_search_points_collection_glove_25.run.xml
+++ b/.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" />
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java
index ca13372ce..3e83cdf58 100644
--- a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java
+++ b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantOpMapper.java
@@ -62,6 +62,7 @@ public class QdrantOpMapper implements OpMapper<QdrantBaseOp<?>> {
                 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);
         };
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantSpace.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantSpace.java
index 3cc19da29..32e2ce38f 100644
--- a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantSpace.java
+++ b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/QdrantSpace.java
@@ -46,8 +46,6 @@ public class QdrantSpace implements AutoCloseable {
 
     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.
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantCountPointsOpDispenser.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantCountPointsOpDispenser.java
new file mode 100644
index 000000000..f0f4edbed
--- /dev/null
+++ b/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));
+    }
+}
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantDeleteCollectionOpDispenser.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantDeleteCollectionOpDispenser.java
index 1a29c1f01..94fbd7d09 100644
--- a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantDeleteCollectionOpDispenser.java
+++ b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/opdispensers/QdrantDeleteCollectionOpDispenser.java
@@ -22,13 +22,10 @@ import io.nosqlbench.adapter.qdrant.ops.QdrantDeleteCollectionOp;
 import io.nosqlbench.adapters.api.templating.ParsedOp;
 import io.qdrant.client.QdrantClient;
 import io.qdrant.client.grpc.Collections.DeleteCollection;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
 
 import java.util.function.LongFunction;
 
 public class QdrantDeleteCollectionOpDispenser extends QdrantBaseOpDispenser<DeleteCollection> {
-    private static final Logger logger = LogManager.getLogger(QdrantDeleteCollectionOpDispenser.class);
 
     /**
      * Create a new {@link QdrantDeleteCollectionOpDispenser} subclassed from {@link QdrantBaseOpDispenser}.
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantCountPointsOp.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/ops/QdrantCountPointsOp.java
new file mode 100644
index 000000000..1cd82aca8
--- /dev/null
+++ b/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 for timeout
+            ).get();
+            logger.info("Total vector points counted: {}", result);
+        } catch (InterruptedException | ExecutionException e) {
+            throw new RuntimeException(e);
+        }
+        return result;
+    }
+}
diff --git a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/types/QdrantOpType.java b/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/types/QdrantOpType.java
index 61da130c2..49ab8e879 100644
--- a/nb-adapters/adapter-qdrant/src/main/java/io/nosqlbench/adapter/qdrant/types/QdrantOpType.java
+++ b/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,
 }
diff --git a/nb-adapters/adapter-qdrant/src/main/resources/activities/qdrant_vectors_live.yaml b/nb-adapters/adapter-qdrant/src/main/resources/activities/qdrant_vectors_live.yaml
index a1acfb003..d66434870 100644
--- a/nb-adapters/adapter-qdrant/src/main/resources/activities/qdrant_vectors_live.yaml
+++ b/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