Skip to content

Commit db210b8

Browse files
marko-bekhtamarko-bekhta
committed
HSEARCH-5064 Integrate Elasticsearch's knn query introduced in 8.12
1 parent 5ddd54a commit db210b8

File tree

12 files changed

+174
-17
lines changed

12 files changed

+174
-17
lines changed

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/dialect/impl/ElasticsearchDialectFactory.java

+7-1
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
import org.hibernate.search.backend.elasticsearch.ElasticsearchDistributionName;
1313
import org.hibernate.search.backend.elasticsearch.ElasticsearchVersion;
1414
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch7ModelDialect;
15+
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch812ModelDialect;
1516
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch8ModelDialect;
1617
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.ElasticsearchModelDialect;
1718
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.OpenSearch1ModelDialect;
@@ -97,7 +98,12 @@ else if ( major == 7 ) {
9798
return new Elasticsearch7ModelDialect();
9899
}
99100
else {
100-
return new Elasticsearch8ModelDialect();
101+
// if there's no minor -- who knows which version it is, better stay safe
102+
// and assume that only 8 features are available, and nothing from 8.12+
103+
if ( major == 8 && ( minorOptional.isEmpty() || minorOptional.getAsInt() < 12 ) ) {
104+
return new Elasticsearch8ModelDialect();
105+
}
106+
return new Elasticsearch812ModelDialect();
101107
}
102108
}
103109

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/dialect/model/impl/Elasticsearch812ModelDialect.java

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
/*
2+
* Hibernate Search, full-text search for your domain model
3+
*
4+
* License: GNU Lesser General Public License (LGPL), version 2.1 or later
5+
* See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
6+
*/
7+
package org.hibernate.search.backend.elasticsearch.dialect.model.impl;
8+
9+
import org.hibernate.search.backend.elasticsearch.types.dsl.provider.impl.Elasticsearch812IndexFieldTypeFactoryProvider;
10+
import org.hibernate.search.backend.elasticsearch.types.dsl.provider.impl.ElasticsearchIndexFieldTypeFactoryProvider;
11+
import org.hibernate.search.backend.elasticsearch.validation.impl.Elasticsearch8PropertyMappingValidatorProvider;
12+
import org.hibernate.search.backend.elasticsearch.validation.impl.ElasticsearchPropertyMappingValidatorProvider;
13+
14+
import com.google.gson.Gson;
15+
16+
/**
17+
* The model dialect for Elasticsearch 8.12+.
18+
*/
19+
public class Elasticsearch812ModelDialect implements ElasticsearchModelDialect {
20+
21+
@Override
22+
public ElasticsearchIndexFieldTypeFactoryProvider createIndexTypeFieldFactoryProvider(Gson userFacingGson) {
23+
return new Elasticsearch812IndexFieldTypeFactoryProvider( userFacingGson );
24+
}
25+
26+
@Override
27+
public ElasticsearchPropertyMappingValidatorProvider createElasticsearchPropertyMappingValidatorProvider() {
28+
return new Elasticsearch8PropertyMappingValidatorProvider();
29+
}
30+
}

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/dialect/model/impl/Elasticsearch8ModelDialect.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
import com.google.gson.Gson;
1515

1616
/**
17-
* The model dialect for Elasticsearch 8.x.
17+
* The model dialect for Elasticsearch [8.0-8.11].
1818
*/
1919
public class Elasticsearch8ModelDialect implements ElasticsearchModelDialect {
2020

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/search/predicate/impl/ElasticsearchKnnPredicate.java

+65
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,19 @@ public KnnPredicateBuilder create(ElasticsearchSearchIndexScope<?> scope,
5858
}
5959
}
6060

61+
public static class Elasticsearch812Factory<F>
62+
extends AbstractElasticsearchCodecAwareSearchQueryElementFactory<KnnPredicateBuilder, F> {
63+
public Elasticsearch812Factory(ElasticsearchFieldCodec<F> codec) {
64+
super( codec );
65+
}
66+
67+
@Override
68+
public KnnPredicateBuilder create(ElasticsearchSearchIndexScope<?> scope,
69+
ElasticsearchSearchIndexValueFieldContext<F> field) {
70+
return new Elasticsearch812Impl.Builder<>( codec, scope, field );
71+
}
72+
}
73+
6174
public static class OpenSearchFactory<F>
6275
extends AbstractElasticsearchCodecAwareSearchQueryElementFactory<KnnPredicateBuilder, F> {
6376
public OpenSearchFactory(ElasticsearchFieldCodec<F> codec) {
@@ -215,6 +228,58 @@ public SearchPredicate build() {
215228
}
216229
}
217230

231+
private static class Elasticsearch812Impl extends ElasticsearchKnnPredicate {
232+
233+
private static final JsonObjectAccessor KNN_ACCESSOR = JsonAccessor.root().property( "knn" ).asObject();
234+
private static final JsonAccessor<String> FIELD_ACCESSOR = JsonAccessor.root().property( "field" ).asString();
235+
private static final JsonArrayAccessor VECTOR_ACCESSOR = JsonAccessor.root().property( "query_vector" ).asArray();
236+
private static final JsonObjectAccessor FILTER_ACCESSOR = JsonAccessor.root().property( "filter" ).asObject();
237+
private static final JsonAccessor<Integer> NUM_CANDIDATES_ACCESSOR =
238+
JsonAccessor.root().property( "num_candidates" ).asInteger();
239+
private static final JsonAccessor<Float> SIMILARITY_ACCESSOR = JsonAccessor.root().property( "similarity" ).asFloat();
240+
241+
242+
private Elasticsearch812Impl(Builder<?> builder) {
243+
super( builder );
244+
}
245+
246+
@Override
247+
protected JsonObject doToJsonQuery(PredicateRequestContext context, JsonObject outerObject, JsonObject innerObject) {
248+
KNN_ACCESSOR.set( outerObject, innerObject );
249+
250+
FIELD_ACCESSOR.set( innerObject, absoluteFieldPath );
251+
NUM_CANDIDATES_ACCESSOR.set( innerObject, k );
252+
VECTOR_ACCESSOR.set( innerObject, vector );
253+
254+
if ( filter != null ) {
255+
FILTER_ACCESSOR.set( innerObject, filter.toJsonQuery( context ) );
256+
}
257+
if ( similarity != null ) {
258+
SIMILARITY_ACCESSOR.set( innerObject, similarity );
259+
}
260+
261+
return outerObject;
262+
}
263+
264+
private static class Builder<F> extends AbstractKnnBuilder<F> {
265+
266+
private Builder(ElasticsearchFieldCodec<F> codec, ElasticsearchSearchIndexScope<?> scope,
267+
ElasticsearchSearchIndexValueFieldContext<F> field) {
268+
super( codec, scope, field );
269+
}
270+
271+
@Override
272+
public void requiredMinimumSimilarity(float similarity) {
273+
this.similarity = similarity;
274+
}
275+
276+
@Override
277+
public SearchPredicate build() {
278+
return new Elasticsearch812Impl( this );
279+
}
280+
}
281+
}
282+
218283
private static class OpenSearchImpl extends ElasticsearchKnnPredicate {
219284

220285
private static final JsonObjectAccessor KNN_ACCESSOR = JsonAccessor.root().property( "knn" ).asObject();

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/types/dsl/provider/impl/Elasticsearch812IndexFieldTypeFactoryProvider.java

+30
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
/*
2+
* Hibernate Search, full-text search for your domain model
3+
*
4+
* License: GNU Lesser General Public License (LGPL), version 2.1 or later
5+
* See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
6+
*/
7+
package org.hibernate.search.backend.elasticsearch.types.dsl.provider.impl;
8+
9+
import org.hibernate.search.backend.elasticsearch.types.mapping.impl.Elasticsearch812VectorFieldTypeMappingContributor;
10+
import org.hibernate.search.backend.elasticsearch.types.mapping.impl.ElasticsearchVectorFieldTypeMappingContributor;
11+
12+
import com.google.gson.Gson;
13+
14+
/**
15+
* The index field type factory provider for ES8.12+.
16+
*/
17+
public class Elasticsearch812IndexFieldTypeFactoryProvider extends AbstractIndexFieldTypeFactoryProvider {
18+
19+
private final Elasticsearch812VectorFieldTypeMappingContributor vectorFieldTypeMappingContributor =
20+
new Elasticsearch812VectorFieldTypeMappingContributor();
21+
22+
public Elasticsearch812IndexFieldTypeFactoryProvider(Gson userFacingGson) {
23+
super( userFacingGson );
24+
}
25+
26+
@Override
27+
protected ElasticsearchVectorFieldTypeMappingContributor vectorFieldTypeMappingContributor() {
28+
return vectorFieldTypeMappingContributor;
29+
}
30+
}

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/types/dsl/provider/impl/Elasticsearch8IndexFieldTypeFactoryProvider.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
import com.google.gson.Gson;
1313

1414
/**
15-
* The index field type factory provider for ES8.x.
15+
* The index field type factory provider for ES [8.0-8.11].
1616
*/
1717
public class Elasticsearch8IndexFieldTypeFactoryProvider extends AbstractIndexFieldTypeFactoryProvider {
1818

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/types/mapping/impl/Elasticsearch812VectorFieldTypeMappingContributor.java

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/*
2+
* Hibernate Search, full-text search for your domain model
3+
*
4+
* License: GNU Lesser General Public License (LGPL), version 2.1 or later
5+
* See the lgpl.txt file in the root directory or <http://www.gnu.org/licenses/lgpl-2.1.html>.
6+
*/
7+
package org.hibernate.search.backend.elasticsearch.types.mapping.impl;
8+
9+
import org.hibernate.search.backend.elasticsearch.search.predicate.impl.ElasticsearchKnnPredicate;
10+
import org.hibernate.search.backend.elasticsearch.types.impl.ElasticsearchIndexValueFieldType;
11+
import org.hibernate.search.engine.search.predicate.spi.PredicateTypeKeys;
12+
13+
public class Elasticsearch812VectorFieldTypeMappingContributor extends Elasticsearch8VectorFieldTypeMappingContributor {
14+
@Override
15+
public <F> void contribute(ElasticsearchIndexValueFieldType.Builder<F> builder, Context context) {
16+
builder.queryElementFactory( PredicateTypeKeys.KNN,
17+
new ElasticsearchKnnPredicate.Elasticsearch812Factory<>( builder.codec() ) );
18+
}
19+
}

‎backend/elasticsearch/src/test/java/org/hibernate/search/backend/elasticsearch/dialect/impl/ElasticsearchDialectFactoryTest.java

+6-5
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import org.hibernate.search.backend.elasticsearch.ElasticsearchDistributionName;
1616
import org.hibernate.search.backend.elasticsearch.ElasticsearchVersion;
1717
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch7ModelDialect;
18+
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch812ModelDialect;
1819
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.Elasticsearch8ModelDialect;
1920
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.ElasticsearchModelDialect;
2021
import org.hibernate.search.backend.elasticsearch.dialect.model.impl.OpenSearch1ModelDialect;
@@ -241,23 +242,23 @@ public static List<? extends Arguments> params() {
241242
),
242243
success(
243244
ElasticsearchDistributionName.ELASTIC, "8.12", "8.12.0",
244-
Elasticsearch8ModelDialect.class, Elasticsearch81ProtocolDialect.class
245+
Elasticsearch812ModelDialect.class, Elasticsearch81ProtocolDialect.class
245246
),
246247
success(
247248
ElasticsearchDistributionName.ELASTIC, "8.12.0", "8.12.0",
248-
Elasticsearch8ModelDialect.class, Elasticsearch81ProtocolDialect.class
249+
Elasticsearch812ModelDialect.class, Elasticsearch81ProtocolDialect.class
249250
),
250251
successWithWarning(
251252
ElasticsearchDistributionName.ELASTIC, "8.13", "8.13.0",
252-
Elasticsearch8ModelDialect.class, Elasticsearch81ProtocolDialect.class
253+
Elasticsearch812ModelDialect.class, Elasticsearch81ProtocolDialect.class
253254
),
254255
successWithWarning(
255256
ElasticsearchDistributionName.ELASTIC, "8.13.0", "8.13.0",
256-
Elasticsearch8ModelDialect.class, Elasticsearch81ProtocolDialect.class
257+
Elasticsearch812ModelDialect.class, Elasticsearch81ProtocolDialect.class
257258
),
258259
successWithWarning(
259260
ElasticsearchDistributionName.ELASTIC, "9.0.0", "9.0.0",
260-
Elasticsearch8ModelDialect.class, Elasticsearch81ProtocolDialect.class
261+
Elasticsearch812ModelDialect.class, Elasticsearch81ProtocolDialect.class
261262
),
262263
success(
263264
ElasticsearchDistributionName.OPENSEARCH, "1", "1.3.1",

‎documentation/src/main/asciidoc/public/reference/_search-dsl-predicate.adoc

+4-2
Original file line numberDiff line numberDiff line change
@@ -1578,7 +1578,7 @@ While when an <<backend-elasticsearch-compatibility-opensearch,OpenSearch>> dist
15781578
when the index is configured to have more than one shard.
15791579
See this section of the OpenSearch link:{openSearchDocUrl}/search-plugins/knn/approximate-knn/#get-started-with-approximate-k-nn[documentation] for more details.
15801580

1581-
With the Elasticsearch backend and when using <<backend-elasticsearch-compatibility-elasticsearch,Elasticsearch>>,
1581+
With the Elasticsearch backend and when using <<backend-elasticsearch-compatibility-elasticsearch,Elasticsearch>> version < 8.12,
15821582
a `knn` predicate can only be added as a top-level predicate,
15831583
i.e. a predicate directly passed to a where clause of a search query,
15841584
or as part of a top-level disjunction,
@@ -1588,7 +1588,9 @@ A `knn` predicate can be combined with other predicate types by adding them thro
15881588
Any other usages of a `knn` predicate, with this backend, would lead to an exception being thrown.
15891589
See this section of the Elasticsearch link:{elasticsearchDocUrl}/knn-search.html#_combine_approximate_knn_with_other_features[documentation]
15901590
in particular to learn more about how Elasticsearch combines regular queries with knn.
1591-
<<backend-elasticsearch-compatibility-opensearch,OpenSearch>> does not have these limitation.
1591+
For Elasticsearch 8.12+ a knn predicate is converted to a link:{elasticsearchDocUrl}/query-dsl-knn-query.html#query-dsl-knn-query[knn-query]
1592+
removing the limitations described for earlier versions.
1593+
<<backend-elasticsearch-compatibility-opensearch,OpenSearch>> also does not have these limitation.
15921594

15931595
.Multiple `knn` predicates added via `should` clauses
15941596
====

‎documentation/src/test/java/org/hibernate/search/documentation/search/predicate/PredicateDslIT.java

+1-1
Original file line numberDiff line numberDiff line change
@@ -1185,7 +1185,7 @@ void knn() {
11851185

11861186
if ( !BackendConfiguration.isElasticsearch()
11871187
|| ElasticsearchTestDialect.isActualVersion(
1188-
es -> false,
1188+
es -> !es.isLessThan( "8.12.0" ),
11891189
os -> !os.isLessThan( "2.0" ),
11901190
aoss -> true
11911191
) ) {

‎integrationtest/backend/elasticsearch/src/test/java/org/hibernate/search/integrationtest/backend/elasticsearch/testsupport/util/ElasticsearchTckBackendFeatures.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -250,16 +250,16 @@ public boolean supportsExplicitRefresh() {
250250
@Override
251251
public boolean supportsVectorSearch() {
252252
return isActualVersion(
253-
es -> !es.isLessThan( "8.0" ),
254-
os -> !os.isLessThan( "2.0" ),
253+
es -> !es.isLessThan( "8.0.0" ),
254+
os -> !os.isLessThan( "2.0.0" ),
255255
aoss -> true
256256
);
257257
}
258258

259259
@Override
260260
public boolean supportsVectorSearchInsideOtherPredicates() {
261261
return isActualVersion(
262-
es -> false,
262+
es -> !es.isLessThan( "8.12.0" ),
263263
os -> true,
264264
aoss -> true
265265
);
@@ -301,7 +301,7 @@ public boolean supportsSimilarity(VectorSimilarity vectorSimilarity) {
301301
@Override
302302
public boolean supportsVectorSearchKPerShard() {
303303
return isActualVersion(
304-
es -> false,
304+
es -> !es.isLessThan( "8.12.0" ),
305305
os -> true,
306306
aoss -> true
307307
);

‎integrationtest/backend/tck/src/main/java/org/hibernate/search/integrationtest/backend/tck/search/predicate/KnnPredicateSpecificsIT.java

+6-2
Original file line numberDiff line numberDiff line change
@@ -844,10 +844,14 @@ private static class NestedIndexBinding {
844844
nested = nestedField.toReference();
845845

846846
byteVector = nestedField.field(
847-
"byteVector", f -> f.asByteVector().dimension( 2 ).projectable( Projectable.YES ) )
847+
"byteVector",
848+
f -> f.asByteVector().dimension( 2 ).vectorSimilarity( VectorSimilarity.L2 )
849+
.projectable( Projectable.YES ) )
848850
.toReference();
849851
floatVector = nestedField
850-
.field( "floatVector", f -> f.asFloatVector().dimension( 2 ).projectable( Projectable.YES ) )
852+
.field( "floatVector",
853+
f -> f.asFloatVector().dimension( 2 ).vectorSimilarity( VectorSimilarity.L2 )
854+
.projectable( Projectable.YES ) )
851855
.toReference();
852856
}
853857
}

0 commit comments

Comments
 (0)