Skip to content

Commit cada0fc

Browse files
opensearch-trigger-bot[bot]github-actions[bot]vibrantvarun
authored
[Backport 2.x] Add Support for Hybrid Query Type (#857)
* Add Support for Hybrid Query Type (#850) * Add Support for Hybrid Query Type Signed-off-by: Varun Jain <[email protected]> * Add samples, guide and integ tests Signed-off-by: Varun Jain <[email protected]> * Removing wildcard imports Signed-off-by: Varun Jain <[email protected]> * Adding import Signed-off-by: Varun Jain <[email protected]> * Adding import Signed-off-by: Varun Jain <[email protected]> --------- Signed-off-by: Varun Jain <[email protected]> (cherry picked from commit 821dae6) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Vacha Shah <[email protected]> * Adding version check for hybrid query integ tests (#863) (#864) Signed-off-by: Varun Jain <[email protected]> Signed-off-by: Vacha Shah <[email protected]> --------- Signed-off-by: Varun Jain <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Vacha Shah <[email protected]> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Varun Jain <[email protected]>
1 parent 7acd729 commit cada0fc

File tree

9 files changed

+332
-30
lines changed

9 files changed

+332
-30
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+
- Add support for Hybrid query type ([#850](https://github.com/opensearch-project/opensearch-java/pull/850))
67

78
### Dependencies
89
- Bumps `org.ajoberstar.grgit:grgit-gradle` from 5.2.0 to 5.2.2

‎guides/search.md

+19
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,25 @@ for (int i = 0; i < searchResponse.hits().hits().size(); i++) {
8181
}
8282
```
8383

84+
### Search documents using a hybrid query
85+
```java
86+
Query searchQuery = Query.of(
87+
h -> h.hybrid(
88+
q -> q.queries(Arrays.asList(
89+
new MatchQuery.Builder().field("text").query(FieldValue.of("Text for document 2")).build().toQuery(),
90+
new TermQuery.Builder().field("passage_text").value(FieldValue.of("Foo bar")).build().toQuery(),
91+
new NeuralQuery.Builder().field("passage_embedding").queryText("Hi world").modelId("bQ1J8ooBpBj3wT4HVUsb").k(100).build().toQuery()
92+
)
93+
)
94+
)
95+
);
96+
SearchRequest searchRequest = new SearchRequest.Builder().query(searchQuery).build();
97+
SearchResponse<IndexData> searchResponse = client.search(searchRequest, IndexData.class);
98+
for (var hit : searchResponse.hits().hits()) {
99+
LOGGER.info("Found {} with score {}", hit.source(), hit.score());
100+
}
101+
```
102+
84103
### Search documents using suggesters
85104

86105
[AppData](../samples/src/main/java/org/opensearch/client/samples/util/AppData.java) refers to the sample data class used in the below samples.

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

+107
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
package org.opensearch.client.opensearch._types.query_dsl;
2+
3+
import jakarta.json.stream.JsonGenerator;
4+
import java.util.List;
5+
import java.util.function.Function;
6+
import org.opensearch.client.json.JsonpDeserializer;
7+
import org.opensearch.client.json.JsonpMapper;
8+
import org.opensearch.client.json.ObjectBuilderDeserializer;
9+
import org.opensearch.client.json.ObjectDeserializer;
10+
import org.opensearch.client.util.ApiTypeHelper;
11+
import org.opensearch.client.util.ObjectBuilder;
12+
13+
public class HybridQuery extends QueryBase implements QueryVariant {
14+
private final List<Query> queries;
15+
16+
private HybridQuery(HybridQuery.Builder builder) {
17+
super(builder);
18+
this.queries = ApiTypeHelper.unmodifiable(builder.queries);
19+
}
20+
21+
public static HybridQuery of(Function<HybridQuery.Builder, ObjectBuilder<HybridQuery>> fn) {
22+
return fn.apply(new HybridQuery.Builder()).build();
23+
}
24+
25+
/**
26+
* Required - list of search queries.
27+
*
28+
* @return list of queries provided under hybrid clause.
29+
*/
30+
public final List<Query> queries() {
31+
return this.queries;
32+
}
33+
34+
@Override
35+
protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) {
36+
super.serializeInternal(generator, mapper);
37+
generator.writeKey("queries");
38+
generator.writeStartArray();
39+
for (Query item0 : this.queries) {
40+
item0.serialize(generator, mapper);
41+
}
42+
generator.writeEnd();
43+
}
44+
45+
@Override
46+
public Query.Kind _queryKind() {
47+
return Query.Kind.Hybrid;
48+
}
49+
50+
public HybridQuery.Builder toBuilder() {
51+
return new HybridQuery.Builder().queries(queries);
52+
}
53+
54+
public static class Builder extends QueryBase.AbstractBuilder<HybridQuery.Builder> implements ObjectBuilder<HybridQuery> {
55+
private List<Query> queries;
56+
57+
/**
58+
* API name: {@code hybrid}
59+
* <p>
60+
* Adds all elements of <code>list</code> to <code>hybrid</code>.
61+
*/
62+
public final HybridQuery.Builder queries(List<Query> list) {
63+
this.queries = _listAddAll(this.queries, list);
64+
return this;
65+
}
66+
67+
/**
68+
* API name: {@code hybrid}
69+
* <p>
70+
* Adds one or more values to <code>hybrid</code>.
71+
*/
72+
public final HybridQuery.Builder queries(Query value, Query... values) {
73+
this.queries = _listAdd(this.queries, value, values);
74+
return this;
75+
}
76+
77+
/**
78+
* API name: {@code hybrid}
79+
* <p>
80+
* Adds a value to <code>hybrid</code> using a builder lambda.
81+
*/
82+
public final HybridQuery.Builder queries(Function<Query.Builder, ObjectBuilder<Query>> fn) {
83+
return queries(fn.apply(new Query.Builder()).build());
84+
}
85+
86+
@Override
87+
protected Builder self() {
88+
return this;
89+
}
90+
91+
@Override
92+
public HybridQuery build() {
93+
_checkSingleUse();
94+
return new HybridQuery(this);
95+
}
96+
}
97+
98+
public static final JsonpDeserializer<HybridQuery> _DESERIALIZER = ObjectBuilderDeserializer.lazy(
99+
HybridQuery.Builder::new,
100+
HybridQuery::setupHybridQueryDeserializer
101+
);
102+
103+
protected static void setupHybridQueryDeserializer(ObjectDeserializer<HybridQuery.Builder> op) {
104+
setupQueryBaseDeserializer(op);
105+
op.add(HybridQuery.Builder::queries, JsonpDeserializer.arrayDeserializer(Query._DESERIALIZER), "queries");
106+
}
107+
}

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

+30
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,8 @@ public enum Kind implements JsonEnum {
120120

121121
Neural("neural"),
122122

123+
Hybrid("hybrid"),
124+
123125
ParentId("parent_id"),
124126

125127
Percolate("percolate"),
@@ -725,6 +727,23 @@ public NeuralQuery neural() {
725727
return TaggedUnionUtils.get(this, Kind.Neural);
726728
}
727729

730+
/**
731+
* Is this variant instance of kind {@code hybrid}?
732+
*/
733+
public boolean isHybrid() {
734+
return _kind == Kind.Hybrid;
735+
}
736+
737+
/**
738+
* Get the {@code hybrid} variant value.
739+
*
740+
* @throws IllegalStateException
741+
* if the current variant is not of the {@code hybrid} kind.
742+
*/
743+
public HybridQuery hybrid() {
744+
return TaggedUnionUtils.get(this, Kind.Hybrid);
745+
}
746+
728747
/**
729748
* Is this variant instance of kind {@code parent_id}?
730749
*/
@@ -1510,6 +1529,16 @@ public ObjectBuilder<Query> neural(Function<NeuralQuery.Builder, ObjectBuilder<N
15101529
return this.neural(fn.apply(new NeuralQuery.Builder()).build());
15111530
}
15121531

1532+
public ObjectBuilder<Query> hybrid(HybridQuery v) {
1533+
this._kind = Kind.Hybrid;
1534+
this._value = v;
1535+
return this;
1536+
}
1537+
1538+
public ObjectBuilder<Query> hybrid(Function<HybridQuery.Builder, ObjectBuilder<HybridQuery>> fn) {
1539+
return this.hybrid(fn.apply(new HybridQuery.Builder()).build());
1540+
}
1541+
15131542
public ObjectBuilder<Query> parentId(ParentIdQuery v) {
15141543
this._kind = Kind.ParentId;
15151544
this._value = v;
@@ -1818,6 +1847,7 @@ protected static void setupQueryDeserializer(ObjectDeserializer<Builder> op) {
18181847
op.add(Builder::multiMatch, MultiMatchQuery._DESERIALIZER, "multi_match");
18191848
op.add(Builder::nested, NestedQuery._DESERIALIZER, "nested");
18201849
op.add(Builder::neural, NeuralQuery._DESERIALIZER, "neural");
1850+
op.add(Builder::hybrid, HybridQuery._DESERIALIZER, "hybrid");
18211851
op.add(Builder::parentId, ParentIdQuery._DESERIALIZER, "parent_id");
18221852
op.add(Builder::percolate, PercolateQuery._DESERIALIZER, "percolate");
18231853
op.add(Builder::pinned, PinnedQuery._DESERIALIZER, "pinned");

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

+7
Original file line numberDiff line numberDiff line change
@@ -261,6 +261,13 @@ public static NeuralQuery.Builder neural() {
261261
return new NeuralQuery.Builder();
262262
}
263263

264+
/**
265+
* Creates a builder for the {@link HybridQuery nested} {@code Query} variant.
266+
*/
267+
public static HybridQuery.Builder hybrid() {
268+
return new HybridQuery.Builder();
269+
}
270+
264271
/**
265272
* Creates a builder for the {@link ParentIdQuery parent_id} {@code Query}
266273
* variant.

‎java-client/src/test/java/org/opensearch/client/opensearch/_types/query_dsl/HybridQueryTest.java

+27
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
package org.opensearch.client.opensearch._types.query_dsl;
2+
3+
import java.util.Arrays;
4+
import org.junit.Test;
5+
import org.opensearch.client.opensearch._types.FieldValue;
6+
import org.opensearch.client.opensearch.model.ModelTestCase;
7+
8+
public class HybridQueryTest extends ModelTestCase {
9+
@Test
10+
public void toBuilder() {
11+
HybridQuery origin = new HybridQuery.Builder().queries(
12+
Arrays.asList(
13+
new TermQuery.Builder().field("passage_text").value(FieldValue.of("Foo bar")).build().toQuery(),
14+
new NeuralQuery.Builder().field("passage_embedding")
15+
.queryText("Hi world")
16+
.modelId("bQ1J8ooBpBj3wT4HVUsb")
17+
.k(100)
18+
.build()
19+
.toQuery(),
20+
new KnnQuery.Builder().field("passage_embedding").vector(new float[] { 0.01f, 0.02f }).k(2).build().toQuery()
21+
)
22+
).build();
23+
HybridQuery copied = origin.toBuilder().build();
24+
25+
assertEquals(toJson(copied), toJson(origin));
26+
}
27+
}

‎java-client/src/test/java/org/opensearch/client/opensearch/model/VariantsTest.java

+58
Original file line numberDiff line numberDiff line change
@@ -32,12 +32,17 @@
3232

3333
package org.opensearch.client.opensearch.model;
3434

35+
import java.util.Arrays;
3536
import org.junit.Test;
3637
import org.opensearch.client.json.JsonData;
38+
import org.opensearch.client.opensearch._types.FieldValue;
3739
import org.opensearch.client.opensearch._types.mapping.Property;
3840
import org.opensearch.client.opensearch._types.mapping.TypeMapping;
41+
import org.opensearch.client.opensearch._types.query_dsl.KnnQuery;
42+
import org.opensearch.client.opensearch._types.query_dsl.NeuralQuery;
3943
import org.opensearch.client.opensearch._types.query_dsl.Query;
4044
import org.opensearch.client.opensearch._types.query_dsl.QueryBuilders;
45+
import org.opensearch.client.opensearch._types.query_dsl.TermQuery;
4146
import org.opensearch.client.opensearch.core.SearchRequest;
4247
import org.opensearch.client.opensearch.indices.GetMappingResponse;
4348

@@ -243,4 +248,57 @@ public void testNeuralQueryFromJson() {
243248
assertEquals("bQ1J8ooBpBj3wT4HVUsb", searchRequest.query().neural().modelId());
244249
assertEquals(100, searchRequest.query().neural().k());
245250
}
251+
252+
@Test
253+
public void testHybridQuery() {
254+
255+
Query query = Query.of(
256+
h -> h.hybrid(
257+
q -> q.queries(
258+
Arrays.asList(
259+
new TermQuery.Builder().field("passage_text").value(FieldValue.of("Foo bar")).build().toQuery(),
260+
new NeuralQuery.Builder().field("passage_embedding")
261+
.queryText("Hi world")
262+
.modelId("bQ1J8ooBpBj3wT4HVUsb")
263+
.k(100)
264+
.build()
265+
.toQuery(),
266+
new KnnQuery.Builder().field("passage_embedding").vector(new float[] { 0.01f, 0.02f }).k(2).build().toQuery()
267+
)
268+
)
269+
)
270+
);
271+
SearchRequest searchRequest = SearchRequest.of(s -> s.query(query));
272+
assertEquals("passage_text", searchRequest.query().hybrid().queries().get(0).term().field());
273+
assertEquals("Foo bar", searchRequest.query().hybrid().queries().get(0).term().value().stringValue());
274+
assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(1).neural().field());
275+
assertEquals("Hi world", searchRequest.query().hybrid().queries().get(1).neural().queryText());
276+
assertEquals("bQ1J8ooBpBj3wT4HVUsb", searchRequest.query().hybrid().queries().get(1).neural().modelId());
277+
assertEquals(100, searchRequest.query().hybrid().queries().get(1).neural().k());
278+
assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(2).knn().field());
279+
assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().vector().length);
280+
assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().k());
281+
}
282+
283+
@Test
284+
public void testHybridQueryFromJson() {
285+
286+
String json = "{\"query\""
287+
+ ":{\"hybrid\":{\"queries\":[{\"term\":{\"passage_text\":\"Foo bar\"}},"
288+
+ "{\"neural\":{\"passage_embedding\":{\"query_text\":\"Hi world\",\"model_id\":\"bQ1J8ooBpBj3wT4HVUsb\",\"k\":100}}},"
289+
+ "{\"knn\":{\"passage_embedding\":{\"vector\":[0.01,0.02],\"k\":2}}}]}},\"size\":10"
290+
+ "}";
291+
292+
SearchRequest searchRequest = ModelTestCase.fromJson(json, SearchRequest.class, mapper);
293+
294+
assertEquals("passage_text", searchRequest.query().hybrid().queries().get(0).term().field());
295+
assertEquals("Foo bar", searchRequest.query().hybrid().queries().get(0).term().value().stringValue());
296+
assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(1).neural().field());
297+
assertEquals("Hi world", searchRequest.query().hybrid().queries().get(1).neural().queryText());
298+
assertEquals("bQ1J8ooBpBj3wT4HVUsb", searchRequest.query().hybrid().queries().get(1).neural().modelId());
299+
assertEquals(100, searchRequest.query().hybrid().queries().get(1).neural().k());
300+
assertEquals("passage_embedding", searchRequest.query().hybrid().queries().get(2).knn().field());
301+
assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().vector().length);
302+
assertEquals(2, searchRequest.query().hybrid().queries().get(2).knn().k());
303+
}
246304
}

0 commit comments

Comments
 (0)