Initial version of count points impl

This commit is contained in:
Madhavan Sridharan 2024-05-07 11:57:51 -04:00
parent 7aacbb0c50
commit 05a53229d6
8 changed files with 113 additions and 6 deletions

View File

@ -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" />

View File

@ -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);
};

View File

@ -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.

View File

@ -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));
}
}

View File

@ -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}.

View File

@ -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;
}
}

View File

@ -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,
}

View File

@ -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