Skip to content

Commit 51a87e9

Browse files
marko-bekhtamarko-bekhta
committed
HSEARCH-4950 Prevent adding knn in some cases for an Elasticsearch backend
1 parent d4c2da0 commit 51a87e9

File tree

9 files changed

+188
-27
lines changed

9 files changed

+188
-27
lines changed

‎backend/elasticsearch/src/main/java/org/hibernate/search/backend/elasticsearch/logging/impl/Log.java

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -864,19 +864,16 @@ SearchException vectorKnnMatchVectorTypeDiffersFromField(String absoluteFieldPat
864864
+ " Use the array of the same size as the vector field.")
865865
SearchException vectorKnnMatchVectorDimensionDiffersFromField(String absoluteFieldPath, int expected, int actual);
866866

867-
@Message(id = ID_OFFSET + 183, value = "A knn predicate can only be added as a should clause to the bool predicate.")
868-
SearchException knnPredicateCanOnlyBeShouldClause();
869-
870-
@Message(id = ID_OFFSET + 184, value = "A knn predicate cannot be added. "
867+
@Message(id = ID_OFFSET + 183, value = "A knn predicate cannot be added. "
871868
+ "Only adding knn predicate as a top-level predicate or as a should clause of a top-level bool predicate is allowed.")
872-
SearchException cannotBeNestedPredicate();
869+
SearchException cannotAddKnnClauseAtThisStep();
873870

874-
@Message(id = ID_OFFSET + 185,
871+
@Message(id = ID_OFFSET + 184,
875872
value = "An OpenSearch distribution does not allow specifying the `number of candidates` option. "
876873
+ "This option is only applicable to an Elasticsearch distribution of an Elasticsearch backend.")
877874
SearchException knnNumberOfCandidatesUnsupportedOption();
878875

879-
@Message(id = ID_OFFSET + 186,
876+
@Message(id = ID_OFFSET + 185,
880877
value = "An %1$s distribution version in use is not compatible with the Hibernate Search integration of vector search. "
881878
+ "Update your %1$s cluster to a %2$s series to get vector search integration enabled.")
882879
SearchException searchBackendVersionIncompatibleWithVectorIntegration(String distribution, String version);

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -34,8 +34,7 @@ public void checkNestableWithin(String expectedParentNestedPath) {
3434
}
3535

3636
@Override
37-
public JsonObject toJsonQuery(PredicateRequestContext context) {
38-
37+
public JsonObject buildJsonQuery(PredicateRequestContext context) {
3938
List<String> nestedPathHierarchy = getNestedPathHierarchy();
4039
String expectedNestedPath = nestedPathHierarchy.isEmpty()
4140
? null
@@ -44,7 +43,7 @@ public JsonObject toJsonQuery(PredicateRequestContext context) {
4443
if ( Objects.equals( context.getNestedPath(), expectedNestedPath ) ) {
4544
// Implicit nesting is not necessary
4645
checkNestableWithin( context.getNestedPath() );
47-
return super.toJsonQuery( context );
46+
return super.buildJsonQuery( context );
4847
}
4948

5049
// The context we expect this predicate to be built in.
@@ -54,7 +53,7 @@ public JsonObject toJsonQuery(PredicateRequestContext context) {
5453
context.withNestedPath( expectedNestedPath );
5554
checkNestableWithin( expectedNestedPath );
5655

57-
JsonObject result = super.toJsonQuery( contextAfterImplicitNesting );
56+
JsonObject result = super.buildJsonQuery( contextAfterImplicitNesting );
5857

5958
// traversing the nestedPathHierarchy in the inverted order
6059
int hierarchyLastIndex = nestedPathHierarchy.size() - 1;

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import com.google.gson.JsonObject;
1616

17-
public abstract class AbstractElasticsearchPredicate implements ElasticsearchSearchPredicate {
17+
public abstract class AbstractElasticsearchPredicate extends ElasticsearchSearchPredicate {
1818

1919
private static final JsonAccessor<Float> BOOST_ACCESSOR = JsonAccessor.root().property( "boost" ).asFloat();
2020

@@ -37,7 +37,7 @@ public Set<String> indexNames() {
3737
}
3838

3939
@Override
40-
public JsonObject toJsonQuery(PredicateRequestContext context) {
40+
public JsonObject buildJsonQuery(PredicateRequestContext context) {
4141
JsonObject outerObject = new JsonObject();
4242
JsonObject innerObject = new JsonObject();
4343

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -113,7 +113,7 @@ private void contributeClauses(PredicateRequestContext context, JsonObject inner
113113
JsonObject jsonQuery = clause.toJsonQuery( context );
114114
if ( jsonQuery == null ) {
115115
if ( !SHOULD_PROPERTY_NAME.equals( occurProperty ) ) {
116-
throw log.knnPredicateCanOnlyBeShouldClause();
116+
throw log.cannotAddKnnClauseAtThisStep();
117117
}
118118
// This is an exceptional case for a KNN search on Elasticsearch distribution.
119119
// A Knn predicate would contribute to a knn clause inside the request context itself

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -157,9 +157,9 @@ private ElasticsearchImpl(Builder<?> builder) {
157157
}
158158

159159
@Override
160-
public JsonObject toJsonQuery(PredicateRequestContext context) {
160+
public JsonObject buildJsonQuery(PredicateRequestContext context) {
161161
// we want the query to get created and passed to the request context
162-
context.contributeKnnClause( ( super.toJsonQuery( context ) ) );
162+
context.contributeKnnClause( ( super.buildJsonQuery( context ) ) );
163163
// but we don't want it to be an actual query so we return `null`:
164164
return null;
165165
}
@@ -184,7 +184,7 @@ protected JsonObject doToJsonQuery(PredicateRequestContext context, JsonObject o
184184
@Override
185185
public void checkNestableWithin(String expectedParentNestedPath) {
186186
if ( expectedParentNestedPath != null ) {
187-
throw log.cannotBeNestedPredicate();
187+
throw log.cannotAddKnnClauseAtThisStep();
188188
}
189189
}
190190

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

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,21 @@
1616

1717
import com.google.gson.JsonObject;
1818

19-
public interface ElasticsearchSearchPredicate extends SearchPredicate {
19+
public abstract class ElasticsearchSearchPredicate implements SearchPredicate {
2020

21-
Log log = LoggerFactory.make( Log.class, MethodHandles.lookup() );
21+
private static final Log log = LoggerFactory.make( Log.class, MethodHandles.lookup() );
2222

23-
Set<String> indexNames();
23+
public abstract Set<String> indexNames();
2424

25-
void checkNestableWithin(String expectedParentNestedPath);
25+
public abstract void checkNestableWithin(String expectedParentNestedPath);
2626

27-
JsonObject toJsonQuery(PredicateRequestContext context);
27+
public final JsonObject toJsonQuery(PredicateRequestContext context) {
28+
return buildJsonQuery( context.predicateContext( this ) );
29+
}
30+
31+
public abstract JsonObject buildJsonQuery(PredicateRequestContext context);
2832

29-
static ElasticsearchSearchPredicate from(ElasticsearchSearchIndexScope<?> scope, SearchPredicate predicate) {
33+
public static ElasticsearchSearchPredicate from(ElasticsearchSearchIndexScope<?> scope, SearchPredicate predicate) {
3034
if ( !( predicate instanceof ElasticsearchSearchPredicate ) ) {
3135
throw log.cannotMixElasticsearchSearchQueryWithOtherPredicates( predicate );
3236
}

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
import com.google.gson.JsonObject;
1414

15-
class ElasticsearchUserProvidedJsonPredicate implements ElasticsearchSearchPredicate {
15+
class ElasticsearchUserProvidedJsonPredicate extends ElasticsearchSearchPredicate {
1616

1717
private final Set<String> indexNames;
1818
private final JsonObject json;
@@ -34,7 +34,7 @@ public void checkNestableWithin(String expectedParentNestedPath) {
3434
}
3535

3636
@Override
37-
public JsonObject toJsonQuery(PredicateRequestContext context) {
37+
public JsonObject buildJsonQuery(PredicateRequestContext context) {
3838
return json;
3939
}
4040
}

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

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,14 @@
66
*/
77
package org.hibernate.search.backend.elasticsearch.search.predicate.impl;
88

9+
import java.lang.invoke.MethodHandles;
10+
911
import org.hibernate.search.backend.elasticsearch.gson.impl.JsonAccessor;
1012
import org.hibernate.search.backend.elasticsearch.gson.impl.JsonObjectAccessor;
13+
import org.hibernate.search.backend.elasticsearch.logging.impl.Log;
1114
import org.hibernate.search.backend.elasticsearch.lowlevel.query.impl.Queries;
1215
import org.hibernate.search.engine.backend.session.spi.BackendSessionContext;
16+
import org.hibernate.search.util.common.logging.impl.LoggerFactory;
1317

1418
import com.google.gson.JsonArray;
1519
import com.google.gson.JsonElement;
@@ -18,33 +22,52 @@
1822

1923
public class PredicateRequestContext {
2024

25+
private static final Log log = LoggerFactory.make( Log.class, MethodHandles.lookup() );
26+
27+
private final PredicateRequestContext parent;
28+
private final Class<? extends ElasticsearchSearchPredicate> predicateType;
2129
private final BackendSessionContext sessionContext;
2230
private final String nestedPath;
2331
private JsonElement jsonKnn = JsonNull.INSTANCE;
2432

2533
public PredicateRequestContext(BackendSessionContext sessionContext) {
2634
this.sessionContext = sessionContext;
2735
this.nestedPath = null;
36+
this.parent = null;
37+
this.predicateType = null;
2838
}
2939

30-
private PredicateRequestContext(BackendSessionContext sessionContext, String nestedPath) {
40+
private PredicateRequestContext(BackendSessionContext sessionContext, String nestedPath, PredicateRequestContext parent,
41+
Class<? extends ElasticsearchSearchPredicate> predicateType) {
3142
this.sessionContext = sessionContext;
3243
this.nestedPath = nestedPath;
44+
this.parent = parent;
45+
this.predicateType = predicateType;
3346
}
3447

3548
String getTenantId() {
3649
return sessionContext.tenantIdentifier();
3750
}
3851

3952
public PredicateRequestContext withNestedPath(String path) {
40-
return new PredicateRequestContext( sessionContext, path );
53+
return new PredicateRequestContext( sessionContext, path, this, predicateType );
54+
}
55+
56+
public PredicateRequestContext predicateContext(ElasticsearchSearchPredicate predicate) {
57+
return new PredicateRequestContext( sessionContext, nestedPath, this, predicate.getClass() );
4158
}
4259

4360
public String getNestedPath() {
4461
return nestedPath;
4562
}
4663

4764
public void contributeKnnClause(JsonObject knn) {
65+
if ( !canAcceptKnnClause() ) {
66+
throw log.cannotAddKnnClauseAtThisStep();
67+
}
68+
if ( parent != null ) {
69+
parent.contributeKnnClause( knn );
70+
}
4871
if ( jsonKnn.isJsonNull() ) {
4972
jsonKnn = knn;
5073
}
@@ -66,6 +89,24 @@ public JsonElement knnSearch(JsonArray filters) {
6689
return addFiltersToKnn( jsonKnn, filters );
6790
}
6891

92+
private boolean canAcceptKnnClause() {
93+
// to allow knn being added to a root context:
94+
return parent == null
95+
// we are adding a knn predicate directly as a root predicate
96+
|| parent.isRoot()
97+
// we are adding a knn should clause to a root boolean predicate.
98+
// the fact that it is a should clause is checked in the bool predicate itself.
99+
|| parent.isBooleanPredicateContext() && parent.parent != null && parent.parent.isRoot();
100+
}
101+
102+
private boolean isBooleanPredicateContext() {
103+
return ElasticsearchBooleanPredicate.class.equals( this.predicateType );
104+
}
105+
106+
private boolean isRoot() {
107+
return parent == null;
108+
}
109+
69110
private static JsonElement addFiltersToKnn(JsonElement jsonKnn, JsonArray filters) {
70111
if ( filters == null || filters.isEmpty() ) {
71112
return jsonKnn;
@@ -92,4 +133,13 @@ private static void addFiltersToKnn(JsonObject jsonKnn, JsonArray filters) {
92133
jsonKnn, Queries.boolFilter( filterAccessor.getOrCreate( jsonKnn, Queries::matchAll ), filters ) );
93134
}
94135
}
136+
137+
@Override
138+
public String toString() {
139+
return "PredicateRequestContext{" +
140+
"parent=" + parent +
141+
", predicateType=" + ( predicateType == null ? "root" : predicateType.getSimpleName() ) +
142+
", nestedPath='" + nestedPath + '\'' +
143+
'}';
144+
}
95145
}

0 commit comments

Comments
 (0)