com.aerospike.vector.client.internal.VectorSearchQuery Maven / Gradle / Ivy

Go to download
package com.aerospike.vector.client.internal;

import com.aerospike.vector.client.HnswSearchParams;
import com.aerospike.vector.client.IndexId;
import com.aerospike.vector.client.Vector;
import com.aerospike.vector.client.VectorSearchRequest;

/**
 * This class encapsulates all necessary parameters required for executing a vector search,
 * including namespace, index name, the vector for search, limit on results, search parameters,
 * and projection for result filtering.
 */
public class VectorSearchQuery {
    private final String namespace;
    private final String indexName;
    private final Vector vector;
    private final int limit;
    private final HnswSearchParams searchParams;
    private final Projection projection;

    /**
     * Constructs a VectorSearchQuery with specified search parameters.
     *
     * @param namespace The namespace of the Aerospike database where the search is to be performed.
     * @param indexName The name of the index within the namespace.
     * @param vector The vector against which other vectors will be searched.
     * @param limit The maximum number of search results to return.
     * @param searchParams Additional search parameters (e.g., algorithm-specific settings).
     * @param projection Projection settings to specify which fields should be included or excluded in the response.
     */
    private VectorSearchQuery(String namespace, String indexName, Vector vector, int limit, HnswSearchParams searchParams, Projection projection) {
        this.namespace = namespace;
        this.indexName = indexName;
        this.vector = vector;
        this.limit = limit;
        this.searchParams = searchParams;
        this.projection = projection;
    }

    public String getNamespace() {
        return namespace;
    }

    public String getIndexName() {
        return indexName;
    }

    public Vector getVector() {
        return vector;
    }

    public int getLimit() {
        return limit;
    }

    public HnswSearchParams getSearchParams() {
        return searchParams;
    }

    public Projection getProjection() {
        return projection;
    }

    /**
     * Builder class for constructing instances of VectorSearchQuery. This builder allows for the setting of
     * all parameters needed for a vector search query, including optional search parameters and result projection.
     */
    public static class Builder {
        private final String namespace;
        private final String indexName;
        private final Vector queryVector;
        private final int limit;

        private HnswSearchParams searchParams = null;
        private Projection projection = new Projection.Builder().withIncludeAll().build();

        public Builder(String namespace, String indexName, Vector queryVector, int limit) {
            this.namespace = namespace;
            this.indexName = indexName;
            this.queryVector = queryVector;
            this.limit = limit;
        }

        public Builder withHnswSearchParams(HnswSearchParams searchParams) {
            this.searchParams = searchParams;
            return this;
        }

        public Builder withProjection(Projection projection) {
            this.projection = projection;
            return this;
        }

        public VectorSearchQuery build() {
            return new VectorSearchQuery(namespace, indexName, queryVector, limit, searchParams, projection);
        }
    }

    /**
     * Converts this query object into a VectorSearchRequest
     *
     * @return a VectorSearchRequest object that contains all the information of this query.
     */
    public VectorSearchRequest toVectorSearchRequest() {
        VectorSearchRequest.Builder requestBuilder = VectorSearchRequest.newBuilder()
                .setIndex(IndexId.newBuilder().setNamespace(namespace).setName(indexName).build())
                .setQueryVector(vector)
                .setLimit(limit)
                .setProjection(projection.toProjectionSpec());

        if (searchParams != null) {
            requestBuilder.setHnswSearchParams(searchParams);
        }

        return requestBuilder.build();
    }
}