-
Notifications
You must be signed in to change notification settings - Fork 182
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add support for approximate k-NN queries #548
Changes from 8 commits
2f2cf7b
015fcd2
8985678
ad886ac
ff1b78a
adbbf94
75a1516
e280f83
41a79c1
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,202 @@ | ||
/* | ||
* SPDX-License-Identifier: Apache-2.0 | ||
* | ||
* The OpenSearch Contributors require contributions made to | ||
* this file be licensed under the Apache-2.0 license or a | ||
* compatible open source license. | ||
*/ | ||
|
||
package org.opensearch.client.opensearch._types.query_dsl; | ||
|
||
import jakarta.json.stream.JsonGenerator; | ||
import java.util.function.Function; | ||
import javax.annotation.Nullable; | ||
import org.opensearch.client.json.JsonpDeserializable; | ||
import org.opensearch.client.json.JsonpDeserializer; | ||
import org.opensearch.client.json.JsonpMapper; | ||
import org.opensearch.client.json.ObjectBuilderDeserializer; | ||
import org.opensearch.client.json.ObjectDeserializer; | ||
import org.opensearch.client.util.ApiTypeHelper; | ||
import org.opensearch.client.util.ObjectBuilder; | ||
|
||
@JsonpDeserializable | ||
public class KnnQuery extends QueryBase implements QueryVariant { | ||
private final String field; | ||
private final float[] vector; | ||
private final int k; | ||
@Nullable | ||
private final Query filter; | ||
|
||
private KnnQuery(Builder builder) { | ||
super(builder); | ||
|
||
this.field = ApiTypeHelper.requireNonNull(builder.field, this, "field"); | ||
this.vector = ApiTypeHelper.requireNonNull(builder.vector, this, "vector"); | ||
this.k = ApiTypeHelper.requireNonNull(builder.k, this, "k"); | ||
this.filter = builder.filter; | ||
} | ||
|
||
public static KnnQuery of(Function<Builder, ObjectBuilder<KnnQuery>> fn) { | ||
return fn.apply(new Builder()).build(); | ||
} | ||
|
||
/** | ||
* Query variant kind. | ||
* @return The query variant kind. | ||
*/ | ||
@Override | ||
public Query.Kind _queryKind() { | ||
return Query.Kind.Knn; | ||
} | ||
|
||
/** | ||
* Required - The target field. | ||
* @return The target field. | ||
*/ | ||
public final String field() { | ||
return this.field; | ||
} | ||
|
||
/** | ||
* Required - The vector to search for. | ||
* @return The vector to search for. | ||
*/ | ||
public final float[] vector() { | ||
return this.vector; | ||
} | ||
|
||
/** | ||
* Required - The number of neighbors the search of each graph will return. | ||
* @return The number of neighbors to return. | ||
*/ | ||
public final int k() { | ||
return this.k; | ||
} | ||
|
||
/** | ||
* Optional - A query to filter the results of the query. | ||
* @return The filter query. | ||
*/ | ||
@Nullable | ||
public final Query filter() { | ||
return this.filter; | ||
} | ||
|
||
@Override | ||
protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) { | ||
generator.writeStartObject(this.field); | ||
|
||
super.serializeInternal(generator, mapper); | ||
|
||
// TODO: Implement the rest of the serialization. | ||
|
||
generator.writeKey("vector"); | ||
generator.writeStartArray(); | ||
for (float value : this.vector) { | ||
generator.write(value); | ||
} | ||
generator.writeEnd(); | ||
|
||
generator.write("k", this.k); | ||
|
||
if (this.filter != null) { | ||
generator.writeKey("filter"); | ||
this.filter.serialize(generator, mapper); | ||
} | ||
|
||
generator.writeEnd(); | ||
} | ||
|
||
/** | ||
* Builder for {@link KnnQuery}. | ||
*/ | ||
public static class Builder extends QueryBase.AbstractBuilder<Builder> implements ObjectBuilder<KnnQuery> { | ||
@Nullable | ||
private String field; | ||
@Nullable | ||
private float[] vector; | ||
@Nullable | ||
private Integer k; | ||
@Nullable | ||
private Query filter; | ||
|
||
/** | ||
* Required - The target field. | ||
* @param field The target field. | ||
* @return This builder. | ||
*/ | ||
public Builder field(@Nullable String field) { | ||
this.field = field; | ||
return this; | ||
} | ||
|
||
/** | ||
* Required - The vector to search for. | ||
* | ||
* @param vector The vector to search for. | ||
* @return This builder. | ||
*/ | ||
public Builder vector(@Nullable float[] vector) { | ||
this.vector = vector; | ||
return this; | ||
} | ||
|
||
/** | ||
* Required - The number of neighbors the search of each graph will return. | ||
* | ||
* @param k The number of neighbors to return. | ||
* @return This builder. | ||
*/ | ||
public Builder k(@Nullable Integer k) { | ||
this.k = k; | ||
return this; | ||
} | ||
|
||
/** | ||
* Optional - A query to filter the results of the knn query. | ||
* | ||
* @param filter The filter query. | ||
* @return This builder. | ||
*/ | ||
public Builder filter(@Nullable Query filter) { | ||
this.filter = filter; | ||
return this; | ||
} | ||
|
||
@Override | ||
protected Builder self() { | ||
return this; | ||
} | ||
|
||
/** | ||
* Builds a {@link KnnQuery}. | ||
* | ||
* @return The built {@link KnnQuery}. | ||
*/ | ||
@Override | ||
public KnnQuery build() { | ||
_checkSingleUse(); | ||
|
||
return new KnnQuery(this); | ||
} | ||
} | ||
|
||
public static final JsonpDeserializer<KnnQuery> _DESERIALIZER = ObjectBuilderDeserializer | ||
.lazy(Builder::new, KnnQuery::setupKnnQueryDeserializer); | ||
|
||
protected static void setupKnnQueryDeserializer(ObjectDeserializer<Builder> op) { | ||
setupQueryBaseDeserializer(op); | ||
op.add((b, v) -> { | ||
float[] vector = new float[v.size()]; | ||
int i = 0; | ||
for (Float value : v) { | ||
vector[i++] = value; | ||
} | ||
b.vector(vector); | ||
}, JsonpDeserializer.arrayDeserializer(JsonpDeserializer.floatDeserializer()), "vector"); | ||
op.add(Builder::k, JsonpDeserializer.integerDeserializer(), "k"); | ||
op.add(Builder::filter, Query._DESERIALIZER, "filter"); | ||
|
||
op.setKey(Builder::field, JsonpDeserializer.stringDeserializer()); | ||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -1066,7 +1066,7 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) { | |
|
||
} | ||
if (this.knnAlgoParamEfSearch != null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Not related to this change, but it seems like a bunch of KNN index settings are missed:
https://github.com/opensearch-project/k-NN/blob/main/RFC.md?plain=1 |
||
generator.writeKey("knn.algo_param_ef_search"); | ||
generator.writeKey("knn.algo_param.ef_search"); | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 👍 |
||
generator.write(this.knnAlgoParamEfSearch); | ||
|
||
} | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
👍