forked from opensearch-project/opensearch-java
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support KnnQuery's
method_parameters and rescore
Signed-off-by: Thomas Farr <tsfarr@amazon.com>
- Loading branch information
Showing
4 changed files
with
258 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
118 changes: 118 additions & 0 deletions
118
...ient/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQueryRescore.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,118 @@ | ||
| /* | ||
| * 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 org.opensearch.client.json.JsonpDeserializable; | ||
| import org.opensearch.client.json.JsonpDeserializer; | ||
| import org.opensearch.client.json.JsonpMapper; | ||
| import org.opensearch.client.json.JsonpSerializable; | ||
| import org.opensearch.client.json.PlainJsonSerializable; | ||
| import org.opensearch.client.json.UnionDeserializer; | ||
| import org.opensearch.client.util.ApiTypeHelper; | ||
| import org.opensearch.client.util.ObjectBuilder; | ||
| import org.opensearch.client.util.ObjectBuilderBase; | ||
| import org.opensearch.client.util.TaggedUnion; | ||
| import org.opensearch.client.util.TaggedUnionUtils; | ||
|
|
||
| @JsonpDeserializable | ||
| public class KnnQueryRescore implements TaggedUnion<KnnQueryRescore.Kind, Object>, PlainJsonSerializable { | ||
| public enum Kind { | ||
| Enable, | ||
| Context | ||
| } | ||
|
|
||
| private final Kind _kind; | ||
| private final Object _value; | ||
|
|
||
| @Override | ||
| public Kind _kind() { | ||
| return _kind; | ||
| } | ||
|
|
||
| @Override | ||
| public Object _get() { | ||
| return _value; | ||
| } | ||
|
|
||
| private KnnQueryRescore(Kind kind, Object value) { | ||
| this._kind = kind; | ||
| this._value = value; | ||
| } | ||
|
|
||
| private KnnQueryRescore(Builder builder) { | ||
| this._kind = ApiTypeHelper.requireNonNull(builder._kind, builder, "<variant kind>"); | ||
| this._value = ApiTypeHelper.requireNonNull(builder._value, builder, "<variant value>"); | ||
| } | ||
|
|
||
| public static KnnQueryRescore of(Function<Builder, ObjectBuilder<KnnQueryRescore>> fn) { | ||
| return fn.apply(new Builder()).build(); | ||
| } | ||
|
|
||
| public boolean isEnable() { | ||
| return _kind == Kind.Enable; | ||
| } | ||
|
|
||
| public Boolean enable() { | ||
| return TaggedUnionUtils.get(this, Kind.Enable); | ||
| } | ||
|
|
||
| public boolean isContext() { | ||
| return _kind == Kind.Context; | ||
| } | ||
|
|
||
| public RescoreContext context() { | ||
| return TaggedUnionUtils.get(this, Kind.Context); | ||
| } | ||
|
|
||
| @Override | ||
| public void serialize(JsonGenerator generator, JsonpMapper mapper) { | ||
| if (_value instanceof JsonpSerializable) { | ||
| ((JsonpSerializable) _value).serialize(generator, mapper); | ||
| } else { | ||
| switch (_kind) { | ||
| case Enable: | ||
| generator.write((Boolean) _value); | ||
| break; | ||
| } | ||
| } | ||
| } | ||
|
|
||
| public static class Builder extends ObjectBuilderBase implements ObjectBuilder<KnnQueryRescore> { | ||
| private Kind _kind; | ||
| private Object _value; | ||
|
|
||
| public ObjectBuilder<KnnQueryRescore> enable(Boolean v) { | ||
| this._kind = Kind.Enable; | ||
| this._value = v; | ||
| return this; | ||
| } | ||
|
|
||
| public ObjectBuilder<KnnQueryRescore> context(RescoreContext v) { | ||
| this._kind = Kind.Context; | ||
| this._value = v; | ||
| return this; | ||
| } | ||
|
|
||
| @Override | ||
| public KnnQueryRescore build() { | ||
| _checkSingleUse(); | ||
| return new KnnQueryRescore(this); | ||
| } | ||
| } | ||
|
|
||
| private static JsonpDeserializer<KnnQueryRescore> buildKnnQueryRescoreDeserializer() { | ||
| return new UnionDeserializer.Builder<KnnQueryRescore, Kind, Object>(KnnQueryRescore::new, false) | ||
| .addMember(Kind.Enable, JsonpDeserializer.booleanDeserializer()) | ||
| .addMember(Kind.Context, RescoreContext._DESERIALIZER) | ||
| .build(); | ||
| } | ||
|
|
||
| public static final JsonpDeserializer<KnnQueryRescore> _DESERIALIZER = JsonpDeserializer.lazy(KnnQueryRescore::buildKnnQueryRescoreDeserializer); | ||
| } |
80 changes: 80 additions & 0 deletions
80
...lient/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/RescoreContext.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,80 @@ | ||
| /* | ||
| * 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 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.json.PlainJsonSerializable; | ||
| import org.opensearch.client.util.ObjectBuilder; | ||
| import org.opensearch.client.util.ObjectBuilderBase; | ||
|
|
||
| import javax.annotation.Nullable; | ||
| import java.util.function.Function; | ||
|
|
||
| @JsonpDeserializable | ||
| public class RescoreContext implements PlainJsonSerializable { | ||
| @Nullable | ||
| private final Float oversampleFactor; | ||
|
|
||
| private RescoreContext(Builder builder) { | ||
| this.oversampleFactor = builder.oversampleFactor; | ||
| } | ||
|
|
||
| public static RescoreContext of(Function<Builder, ObjectBuilder<RescoreContext>> fn) { | ||
| return fn.apply(new Builder()).build(); | ||
| } | ||
|
|
||
| @Nullable | ||
| public Float oversampleFactor() { | ||
| return this.oversampleFactor; | ||
| } | ||
|
|
||
| @Override | ||
| public void serialize(JsonGenerator generator, JsonpMapper mapper) { | ||
| generator.writeStartObject(); | ||
| serializeInternal(generator, mapper); | ||
| generator.writeEnd(); | ||
| } | ||
|
|
||
| protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) { | ||
| if (this.oversampleFactor != null) { | ||
| generator.writeKey("oversample_factor"); | ||
| generator.write(this.oversampleFactor); | ||
| } | ||
| } | ||
|
|
||
| public static class Builder extends ObjectBuilderBase implements ObjectBuilder<RescoreContext> { | ||
| @Nullable | ||
| private Float oversampleFactor; | ||
|
|
||
| public Builder oversampleFactor(@Nullable Float value) { | ||
| this.oversampleFactor = value; | ||
| return this; | ||
| } | ||
|
|
||
| @Override | ||
| public RescoreContext build() { | ||
| _checkSingleUse(); | ||
| return new RescoreContext(this); | ||
| } | ||
| } | ||
|
|
||
| public static final JsonpDeserializer<RescoreContext> _DESERIALIZER = ObjectBuilderDeserializer.lazy( | ||
| Builder::new, | ||
| RescoreContext::setupRescoreContextDeserializer | ||
| ); | ||
|
|
||
| protected static void setupRescoreContextDeserializer(ObjectDeserializer<Builder> op) { | ||
| op.add(Builder::oversampleFactor, JsonpDeserializer.floatDeserializer(), "oversample_factor"); | ||
| } | ||
| } |