Additional APIs implemented

This commit is contained in:
Madhavan Sridharan 2024-05-13 22:10:09 -04:00
parent 6a20e8da07
commit ccafc97172
17 changed files with 476 additions and 9 deletions

View File

@ -35,16 +35,16 @@
<dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-protobuf</artifactId>
<version>1.63.0</version>
<!-- <version>1.63.0</version> -->
<!-- If Trying to match https://github.com/qdrant/java-client/blob/v1.9.1/build.gradle#L80 -->
<!-- <version>1.59.0</version> -->
<version>1.59.0</version>
</dependency>
<dependency>
<groupId>com.google.protobuf</groupId>
<artifactId>protobuf-java-util</artifactId>
<version>3.25.3</version>
<!-- <version>3.25.3</version> -->
<!-- If Trying to match https://github.com/qdrant/java-client/blob/v1.9.1/build.gradle#L81 -->
<!-- <version>3.24.0</version> -->
<version>3.24.0</version>
</dependency>
<dependency>
<groupId>com.google.guava</groupId>

View File

@ -63,6 +63,12 @@ public class QdrantOpMapper implements OpMapper<QdrantBaseOp<?>> {
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);
case list_collections -> new QdrantListCollectionsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case collection_info -> new QdrantCollectionInfoOpDispenser(adapter, op, typeAndTarget.targetFunction);
case collection_exists -> new QdrantCollectionExistsOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_collection_aliases ->
new QdrantListCollectionAliasesOpDispenser(adapter, op, typeAndTarget.targetFunction);
case list_snapshots -> new QdrantListSnapshotsOpDispenser(adapter, op, typeAndTarget.targetFunction);
// default -> throw new RuntimeException("Unrecognized op type '" + typeAndTarget.enumId.name() + "' while " +
// "mapping parsed op " + op);
};

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

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.opdispensers;
import io.nosqlbench.adapter.qdrant.QdrantDriverAdapter;
import io.nosqlbench.adapter.qdrant.ops.QdrantBaseOp;
import io.nosqlbench.adapter.qdrant.ops.QdrantCollectionInfoOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.qdrant.client.QdrantClient;
import java.util.function.LongFunction;
public class QdrantCollectionInfoOpDispenser extends QdrantBaseOpDispenser<String> {
public QdrantCollectionInfoOpDispenser(QdrantDriverAdapter adapter, ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
}
@Override
public LongFunction<String> getParamFunc(
LongFunction<QdrantClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> targetF.apply(l);
}
@Override
public LongFunction<QdrantBaseOp<String>> createOpFunc(
LongFunction<String> paramF,
LongFunction<QdrantClient> clientF,
ParsedOp op,
LongFunction<String> targetF) {
return l -> new QdrantCollectionInfoOp(clientF.apply(l), paramF.apply(l));
}
}

View File

@ -37,9 +37,13 @@ public class QdrantCountPointsOpDispenser extends QdrantBaseOpDispenser<CountPoi
LongFunction<CountPoints.Builder> ebF =
l -> CountPoints.newBuilder().setCollectionName(targetF.apply(l));
ebF = op.enhanceFuncOptionally(ebF, "exact", Boolean.class, CountPoints.Builder::setExact);
LongFunction<Filter.Builder> filterBuilder = getFilterFromOp(op);
final LongFunction<CountPoints.Builder> filterF = ebF;
ebF = l -> filterF.apply(l).setFilter(filterBuilder.apply(l));
if (filterBuilder != null) {
final LongFunction<CountPoints.Builder> filterF = ebF;
ebF = l -> filterF.apply(l).setFilter(filterBuilder.apply(l));
}
final LongFunction<CountPoints.Builder> lastF = ebF;
return l -> lastF.apply(l).build();

View File

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

View File

@ -0,0 +1,45 @@
/*
* 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.QdrantListCollectionsOp;
import io.nosqlbench.adapters.api.templating.ParsedOp;
import io.qdrant.client.QdrantClient;
import java.util.function.LongFunction;
public class QdrantListCollectionsOpDispenser extends QdrantBaseOpDispenser<Object> {
public QdrantListCollectionsOpDispenser(QdrantDriverAdapter adapter,
ParsedOp op, LongFunction<String> targetFunction) {
super(adapter, op, targetFunction);
}
@Override
public LongFunction<Object> getParamFunc(LongFunction<QdrantClient> clientF, ParsedOp op, LongFunction<String> targetF) {
return l -> "";
}
@Override
public LongFunction<QdrantBaseOp<Object>> createOpFunc(
LongFunction<Object> paramF,
LongFunction<QdrantClient> clientF,
ParsedOp op, LongFunction<String> targetF) {
return l -> new QdrantListCollectionsOp(clientF.apply(l), null);
}
}

View File

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

View File

@ -114,8 +114,10 @@ public class QdrantSearchPointsOpDispenser extends QdrantBaseOpDispenser<SearchP
}
LongFunction<Filter.Builder> filterBuilder = getFilterFromOp(op);
final LongFunction<SearchPoints.Builder> filterF = ebF;
ebF = l -> filterF.apply(l).setFilter(filterBuilder.apply(l));
if (filterBuilder != null) {
final LongFunction<SearchPoints.Builder> filterF = ebF;
ebF = l -> filterF.apply(l).setFilter(filterBuilder.apply(l));
}
final LongFunction<SearchPoints.Builder> lastF = ebF;
return l -> lastF.apply(l).build();

View File

@ -0,0 +1,40 @@
/*
* 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.Collections.CollectionExistsRequest;
import java.time.Duration;
public class QdrantCollectionExistsOp extends QdrantBaseOp<CollectionExistsRequest> {
public QdrantCollectionExistsOp(QdrantClient client, CollectionExistsRequest request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
Boolean response;
try {
response = client.collectionExistsAsync(request.getCollectionName(), Duration.ofSeconds(600)).get();
logger.info("Collection {} exists: {}", request.getCollectionName(), response);
} catch (Exception e) {
throw new RuntimeException(e);
}
return response;
}
}

View File

@ -0,0 +1,40 @@
/*
* 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.Collections.CollectionInfo;
import java.time.Duration;
public class QdrantCollectionInfoOp extends QdrantBaseOp<String> {
public QdrantCollectionInfoOp(QdrantClient client, String request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
CollectionInfo response;
try {
response = client.getCollectionInfoAsync(request, Duration.ofSeconds(600)).get();
logger.debug("Collection info: {}", response.toString());
} catch (Exception e) {
throw new RuntimeException(e);
}
return response;
}
}

View File

@ -31,9 +31,11 @@ public class QdrantCountPointsOp extends QdrantBaseOp<CountPoints> {
public Object applyOp(long value) {
long result;
try {
boolean hasFilters = request.getFilter() != null && (request.getFilter().getMustCount() > 0
|| request.getFilter().getMustNotCount() > 0 || request.getFilter().getShouldCount() > 0);
result = client.countAsync(
request.getCollectionName(),
request.getFilter(),
(hasFilters) ? request.getFilter() : null,
request.getExact(),
Duration.ofMinutes(5) // opinionated default of 5 minutes for timeout
).get();

View File

@ -0,0 +1,41 @@
/*
* 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.Collections.ListCollectionAliasesRequest;
import java.time.Duration;
import java.util.List;
public class QdrantListCollectionAliasesOp extends QdrantBaseOp<ListCollectionAliasesRequest> {
public QdrantListCollectionAliasesOp(QdrantClient client, ListCollectionAliasesRequest request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
List<String> response;
try {
response = client.listCollectionAliasesAsync(request.getCollectionName(), Duration.ofSeconds(600)).get();
logger.info("Collection aliases are {}", response);
} catch (Exception e) {
throw new RuntimeException(e);
}
return response;
}
}

View File

@ -0,0 +1,39 @@
/*
* 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 java.time.Duration;
import java.util.List;
public class QdrantListCollectionsOp extends QdrantBaseOp<Object> {
public QdrantListCollectionsOp(QdrantClient client, Object request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
List<String> response;
try {
response = client.listCollectionsAsync(Duration.ofSeconds(300)).get();
logger.info("[QdrantListCollectionsOp] Collections: {}", response);
} catch (Exception e) {
throw new RuntimeException(e);
}
return response;
}
}

View File

@ -0,0 +1,42 @@
/*
* 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.SnapshotsService.ListSnapshotsRequest;
import io.qdrant.client.grpc.SnapshotsService.SnapshotDescription;
import java.time.Duration;
import java.util.List;
public class QdrantListSnapshotsOp extends QdrantBaseOp<ListSnapshotsRequest> {
public QdrantListSnapshotsOp(QdrantClient client, ListSnapshotsRequest request) {
super(client, request);
}
@Override
public Object applyOp(long value) {
List<SnapshotDescription> response;
try {
response = client.listSnapshotAsync(request.getCollectionName(), Duration.ofSeconds(600)).get();
response.forEach(s -> logger.info("[QdrantListSnapshotsOp] Snapshot: " + s.toString()));
} catch(Exception e) {
throw new RuntimeException(e);
}
return response;
}
}

View File

@ -28,4 +28,9 @@ public enum QdrantOpType {
// https://qdrant.github.io/qdrant/redoc/index.html#tag/points/operation/count_points
// https://qdrant.tech/documentation/concepts/points/#counting-points
count_points,
list_collections,
collection_info,
collection_exists,
list_collection_aliases,
list_snapshots,
}

View File

@ -34,6 +34,11 @@ The following are a couple high level API operations.
* Drop Collection
* Search Points (vectors)
* Create Payload Index
* List Collections
* List Collection Aliases
* List Snapshots
* Collection Info
* Collection Exists
## Examples