Skip to content

Commit e019335

Browse files
XtansiaXtansia
committed
Support KnnQuery's method_parameters and rescore
Signed-off-by: Thomas Farr <[email protected]>
1 parent 31260b6 commit e019335

File tree

4 files changed

+265
-1
lines changed

4 files changed

+265
-1
lines changed

‎CHANGELOG.md

+1
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
33

44
## [Unreleased 2.x]
55
### Added
6+
- Added support for `KnnQuery`'s `method_parameters` and `rescore` properties ([#1407](https://github.com/opensearch-project/opensearch-java/pull/1407))
67

78
### Dependencies
89
- Bump `org.junit:junit-bom` from 5.11.3 to 5.11.4 ([#1367](https://github.com/opensearch-project/opensearch-java/pull/1367))

‎java-client/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQuery.java

+64-1
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,11 @@
99
package org.opensearch.client.opensearch._types.query_dsl;
1010

1111
import jakarta.json.stream.JsonGenerator;
12+
import java.util.Map;
1213
import java.util.function.Function;
14+
import javax.annotation.Nonnull;
1315
import javax.annotation.Nullable;
16+
import org.opensearch.client.json.JsonData;
1417
import org.opensearch.client.json.JsonpDeserializable;
1518
import org.opensearch.client.json.JsonpDeserializer;
1619
import org.opensearch.client.json.JsonpMapper;
@@ -31,6 +34,10 @@ public class KnnQuery extends QueryBase implements QueryVariant {
3134
private final Float maxDistance;
3235
@Nullable
3336
private final Query filter;
37+
@Nonnull
38+
private final Map<String, JsonData> methodParameters;
39+
@Nullable
40+
private final KnnQueryRescore rescore;
3441

3542
private KnnQuery(Builder builder) {
3643
super(builder);
@@ -41,6 +48,8 @@ private KnnQuery(Builder builder) {
4148
this.minScore = builder.minScore;
4249
this.maxDistance = builder.maxDistance;
4350
this.filter = builder.filter;
51+
this.methodParameters = ApiTypeHelper.unmodifiable(builder.methodParameters);
52+
this.rescore = builder.rescore;
4453
}
4554

4655
public static KnnQuery of(Function<Builder, ObjectBuilder<KnnQuery>> fn) {
@@ -108,6 +117,16 @@ public final Query filter() {
108117
return this.filter;
109118
}
110119

120+
@Nonnull
121+
public final Map<String, JsonData> methodParameters() {
122+
return this.methodParameters;
123+
}
124+
125+
@Nullable
126+
public final KnnQueryRescore rescore() {
127+
return this.rescore;
128+
}
129+
111130
@Override
112131
protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
113132
generator.writeStartObject(this.field);
@@ -138,11 +157,30 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
138157
this.filter.serialize(generator, mapper);
139158
}
140159

160+
if (ApiTypeHelper.isDefined(this.methodParameters)) {
161+
for (Map.Entry<String, JsonData> entry : this.methodParameters.entrySet()) {
162+
generator.writeKey(entry.getKey());
163+
entry.getValue().serialize(generator, mapper);
164+
}
165+
}
166+
167+
if (this.rescore != null) {
168+
generator.writeKey("rescore");
169+
this.rescore.serialize(generator, mapper);
170+
}
171+
141172
generator.writeEnd();
142173
}
143174

144175
public Builder toBuilder() {
145-
return toBuilder(new Builder()).field(field).vector(vector).k(k).minScore(minScore).maxDistance(maxDistance).filter(filter);
176+
return toBuilder(new Builder()).field(field)
177+
.vector(vector)
178+
.k(k)
179+
.minScore(minScore)
180+
.maxDistance(maxDistance)
181+
.filter(filter)
182+
.methodParameters(methodParameters)
183+
.rescore(rescore);
146184
}
147185

148186
/**
@@ -161,6 +199,10 @@ public static class Builder extends QueryBase.AbstractBuilder<Builder> implement
161199
private Float maxDistance;
162200
@Nullable
163201
private Query filter;
202+
@Nullable
203+
private Map<String, JsonData> methodParameters;
204+
@Nullable
205+
private KnnQueryRescore rescore;
164206

165207
/**
166208
* Required - The target field.
@@ -227,6 +269,25 @@ public Builder filter(@Nullable Query filter) {
227269
return this;
228270
}
229271

272+
public Builder methodParameters(@Nonnull Map<String, JsonData> value) {
273+
this.methodParameters = _mapPutAll(this.methodParameters, value);
274+
return this;
275+
}
276+
277+
public Builder methodParameters(String key, JsonData value) {
278+
this.methodParameters = _mapPut(this.methodParameters, key, value);
279+
return this;
280+
}
281+
282+
public Builder rescore(@Nullable KnnQueryRescore value) {
283+
this.rescore = value;
284+
return this;
285+
}
286+
287+
public Builder rescore(Function<KnnQueryRescore.Builder, ObjectBuilder<KnnQueryRescore>> fn) {
288+
return this.rescore(fn.apply(new KnnQueryRescore.Builder()).build());
289+
}
290+
230291
@Override
231292
protected Builder self() {
232293
return this;
@@ -264,6 +325,8 @@ protected static void setupKnnQueryDeserializer(ObjectDeserializer<Builder> op)
264325
op.add(Builder::minScore, JsonpDeserializer.floatDeserializer(), "min_score");
265326
op.add(Builder::maxDistance, JsonpDeserializer.floatDeserializer(), "max_distance");
266327
op.add(Builder::filter, Query._DESERIALIZER, "filter");
328+
op.add(Builder::methodParameters, JsonpDeserializer.stringMapDeserializer(JsonData._DESERIALIZER), "method_parameters");
329+
op.add(Builder::rescore, KnnQueryRescore._DESERIALIZER, "rescore");
267330

268331
op.setKey(Builder::field, JsonpDeserializer.stringDeserializer());
269332
}

‎java-client/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQueryRescore.java

+121
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.client.opensearch._types.query_dsl;
10+
11+
import jakarta.json.stream.JsonGenerator;
12+
import java.util.function.Function;
13+
import org.opensearch.client.json.JsonpDeserializable;
14+
import org.opensearch.client.json.JsonpDeserializer;
15+
import org.opensearch.client.json.JsonpMapper;
16+
import org.opensearch.client.json.JsonpSerializable;
17+
import org.opensearch.client.json.PlainJsonSerializable;
18+
import org.opensearch.client.json.UnionDeserializer;
19+
import org.opensearch.client.util.ApiTypeHelper;
20+
import org.opensearch.client.util.ObjectBuilder;
21+
import org.opensearch.client.util.ObjectBuilderBase;
22+
import org.opensearch.client.util.TaggedUnion;
23+
import org.opensearch.client.util.TaggedUnionUtils;
24+
25+
@JsonpDeserializable
26+
public class KnnQueryRescore implements TaggedUnion<KnnQueryRescore.Kind, Object>, PlainJsonSerializable {
27+
public enum Kind {
28+
Enable,
29+
Context
30+
}
31+
32+
private final Kind _kind;
33+
private final Object _value;
34+
35+
@Override
36+
public Kind _kind() {
37+
return _kind;
38+
}
39+
40+
@Override
41+
public Object _get() {
42+
return _value;
43+
}
44+
45+
private KnnQueryRescore(Kind kind, Object value) {
46+
this._kind = kind;
47+
this._value = value;
48+
}
49+
50+
private KnnQueryRescore(Builder builder) {
51+
this._kind = ApiTypeHelper.requireNonNull(builder._kind, builder, "<variant kind>");
52+
this._value = ApiTypeHelper.requireNonNull(builder._value, builder, "<variant value>");
53+
}
54+
55+
public static KnnQueryRescore of(Function<Builder, ObjectBuilder<KnnQueryRescore>> fn) {
56+
return fn.apply(new Builder()).build();
57+
}
58+
59+
public boolean isEnable() {
60+
return _kind == Kind.Enable;
61+
}
62+
63+
public Boolean enable() {
64+
return TaggedUnionUtils.get(this, Kind.Enable);
65+
}
66+
67+
public boolean isContext() {
68+
return _kind == Kind.Context;
69+
}
70+
71+
public RescoreContext context() {
72+
return TaggedUnionUtils.get(this, Kind.Context);
73+
}
74+
75+
@Override
76+
public void serialize(JsonGenerator generator, JsonpMapper mapper) {
77+
if (_value instanceof JsonpSerializable) {
78+
((JsonpSerializable) _value).serialize(generator, mapper);
79+
} else {
80+
switch (_kind) {
81+
case Enable:
82+
generator.write((Boolean) _value);
83+
break;
84+
}
85+
}
86+
}
87+
88+
public static class Builder extends ObjectBuilderBase implements ObjectBuilder<KnnQueryRescore> {
89+
private Kind _kind;
90+
private Object _value;
91+
92+
public ObjectBuilder<KnnQueryRescore> enable(Boolean v) {
93+
this._kind = Kind.Enable;
94+
this._value = v;
95+
return this;
96+
}
97+
98+
public ObjectBuilder<KnnQueryRescore> context(RescoreContext v) {
99+
this._kind = Kind.Context;
100+
this._value = v;
101+
return this;
102+
}
103+
104+
@Override
105+
public KnnQueryRescore build() {
106+
_checkSingleUse();
107+
return new KnnQueryRescore(this);
108+
}
109+
}
110+
111+
private static JsonpDeserializer<KnnQueryRescore> buildKnnQueryRescoreDeserializer() {
112+
return new UnionDeserializer.Builder<KnnQueryRescore, Kind, Object>(KnnQueryRescore::new, false).addMember(
113+
Kind.Enable,
114+
JsonpDeserializer.booleanDeserializer()
115+
).addMember(Kind.Context, RescoreContext._DESERIALIZER).build();
116+
}
117+
118+
public static final JsonpDeserializer<KnnQueryRescore> _DESERIALIZER = JsonpDeserializer.lazy(
119+
KnnQueryRescore::buildKnnQueryRescoreDeserializer
120+
);
121+
}

‎java-client/src/main/java/org/opensearch/client/opensearch/_types/query_dsl/RescoreContext.java

+79
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.client.opensearch._types.query_dsl;
10+
11+
import jakarta.json.stream.JsonGenerator;
12+
import java.util.function.Function;
13+
import javax.annotation.Nullable;
14+
import org.opensearch.client.json.JsonpDeserializable;
15+
import org.opensearch.client.json.JsonpDeserializer;
16+
import org.opensearch.client.json.JsonpMapper;
17+
import org.opensearch.client.json.ObjectBuilderDeserializer;
18+
import org.opensearch.client.json.ObjectDeserializer;
19+
import org.opensearch.client.json.PlainJsonSerializable;
20+
import org.opensearch.client.util.ObjectBuilder;
21+
import org.opensearch.client.util.ObjectBuilderBase;
22+
23+
@JsonpDeserializable
24+
public class RescoreContext implements PlainJsonSerializable {
25+
@Nullable
26+
private final Float oversampleFactor;
27+
28+
private RescoreContext(Builder builder) {
29+
this.oversampleFactor = builder.oversampleFactor;
30+
}
31+
32+
public static RescoreContext of(Function<Builder, ObjectBuilder<RescoreContext>> fn) {
33+
return fn.apply(new Builder()).build();
34+
}
35+
36+
@Nullable
37+
public Float oversampleFactor() {
38+
return this.oversampleFactor;
39+
}
40+
41+
@Override
42+
public void serialize(JsonGenerator generator, JsonpMapper mapper) {
43+
generator.writeStartObject();
44+
serializeInternal(generator, mapper);
45+
generator.writeEnd();
46+
}
47+
48+
protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
49+
if (this.oversampleFactor != null) {
50+
generator.writeKey("oversample_factor");
51+
generator.write(this.oversampleFactor);
52+
}
53+
}
54+
55+
public static class Builder extends ObjectBuilderBase implements ObjectBuilder<RescoreContext> {
56+
@Nullable
57+
private Float oversampleFactor;
58+
59+
public Builder oversampleFactor(@Nullable Float value) {
60+
this.oversampleFactor = value;
61+
return this;
62+
}
63+
64+
@Override
65+
public RescoreContext build() {
66+
_checkSingleUse();
67+
return new RescoreContext(this);
68+
}
69+
}
70+
71+
public static final JsonpDeserializer<RescoreContext> _DESERIALIZER = ObjectBuilderDeserializer.lazy(
72+
Builder::new,
73+
RescoreContext::setupRescoreContextDeserializer
74+
);
75+
76+
protected static void setupRescoreContextDeserializer(ObjectDeserializer<Builder> op) {
77+
op.add(Builder::oversampleFactor, JsonpDeserializer.floatDeserializer(), "oversample_factor");
78+
}
79+
}

0 commit comments

Comments
 (0)