Skip to content

Commit 4436e89

Browse files
committed
Fix KNNQuery Precision
Signed-off-by: luyuncheng <[email protected]>
1 parent 045c805 commit 4436e89

File tree

3 files changed

+32
-1
lines changed

3 files changed

+32
-1
lines changed

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

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

1111
import jakarta.json.stream.JsonGenerator;
12+
import java.math.BigDecimal;
1213
import java.util.function.Function;
1314
import javax.annotation.Nullable;
1415
import org.opensearch.client.json.JsonpDeserializable;
@@ -19,6 +20,8 @@
1920
import org.opensearch.client.util.ApiTypeHelper;
2021
import org.opensearch.client.util.ObjectBuilder;
2122

23+
import static java.math.RoundingMode.HALF_UP;
24+
2225
@JsonpDeserializable
2326
public class KnnQuery extends QueryBase implements QueryVariant {
2427
private final String field;
@@ -93,7 +96,9 @@ protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
9396
generator.writeKey("vector");
9497
generator.writeStartArray();
9598
for (float value : this.vector) {
96-
generator.write(value);
99+
BigDecimal b = new BigDecimal(value);
100+
double T = b.setScale(6, HALF_UP).doubleValue();
101+
generator.write(T);
97102
}
98103
generator.writeEnd();
99104

java-client/src/test/java/org/opensearch/client/opensearch/_types/query_dsl/KnnQueryTest.java

+7
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,11 @@ public void toBuilder() {
1111

1212
assertEquals(toJson(copied), toJson(origin));
1313
}
14+
15+
@Test
16+
public void toBuilderPrecision() {
17+
KnnQuery origin = new KnnQuery.Builder().field("field").vector(new float[] { 0.1f, 0.4f }).k(1).build();
18+
19+
assertEquals(toJson(origin), "{\"field\":{\"vector\":[0.1,0.4],\"k\":1}}");
20+
}
1421
}

java-client/src/test/java/org/opensearch/client/opensearch/model/RequestEncodingTest.java

+19
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
package org.opensearch.client.opensearch.model;
3434

3535
import org.junit.Test;
36+
import org.opensearch.client.json.jackson.JacksonJsonpMapper;
3637
import org.opensearch.client.json.jsonb.JsonbJsonpMapper;
3738
import org.opensearch.client.opensearch.core.SearchRequest;
3839

@@ -61,4 +62,22 @@ public void testParametersNotInJson() {
6162
assertNull(request.q());
6263

6364
}
65+
66+
@Test
67+
public void testKnnVectorPrecision() {
68+
69+
float[] vector = {0.4f, 0.3f};
70+
SearchRequest request = new SearchRequest.Builder().q("knn")
71+
.query(q -> q.knn(k -> k.field("values").vector(vector).k(1)))
72+
.build();
73+
74+
JacksonJsonpMapper mapper = new JacksonJsonpMapper();
75+
String str = toJson(request, mapper);
76+
assertEquals("{\"query\":{\"knn\":{\"values\":{\"vector\":[0.4,0.3],\"k\":1}}}}", str);
77+
78+
request = fromJson(str, SearchRequest.class, mapper);
79+
80+
assertTrue(request.query().isKnn());
81+
assertNull(request.q());
82+
}
6483
}

0 commit comments

Comments
 (0)