diff --git a/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java b/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java index aca93b77b..0788f43a8 100644 --- a/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java +++ b/src/main/java/org/opensearch/neuralsearch/search/query/HybridQueryPhaseSearcher.java @@ -14,6 +14,7 @@ import lombok.NoArgsConstructor; import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; +import org.apache.lucene.search.DisjunctionMaxQuery; import org.apache.lucene.search.Query; import org.opensearch.common.settings.Settings; import org.opensearch.index.mapper.MapperService; @@ -104,7 +105,7 @@ protected Query extractHybridQuery(final SearchContext searchContext, final Quer * } * ] * } - * TODO add similar validation for other compound type queries like dis_max, constant_score etc. + * TODO add similar validation for other compound type queries like constant_score, function_score etc. * * @param query query to validate */ @@ -114,6 +115,10 @@ private void validateQuery(final SearchContext searchContext, final Query query) for (BooleanClause booleanClause : booleanClauses) { validateNestedBooleanQuery(booleanClause.getQuery(), getMaxDepthLimit(searchContext)); } + } else if (query instanceof DisjunctionMaxQuery) { + for (Query disjunct : (DisjunctionMaxQuery) query) { + validateNestedDisJunctionQuery(disjunct, getMaxDepthLimit(searchContext)); + } } } @@ -135,6 +140,24 @@ private void validateNestedBooleanQuery(final Query query, final int level) { } } + private void validateNestedDisJunctionQuery(final Query query, final int level) { + if (query instanceof HybridQuery) { + throw new IllegalArgumentException("hybrid query must be a top level query and cannot be wrapped into other queries"); + } + if (level <= 0) { + // ideally we should throw an error here but this code is on the main search workflow path and that might block + // execution of some queries. Instead, we're silently exit and allow such query to execute and potentially produce incorrect + // results in case hybrid query is wrapped into such dis_max query + log.error("reached max nested query limit, cannot process dis_max query with that many nested clauses"); + return; + } + if (query instanceof DisjunctionMaxQuery) { + for (Query disjunct : (DisjunctionMaxQuery) query) { + validateNestedDisJunctionQuery(disjunct, level - 1); + } + } + } + private int getMaxDepthLimit(final SearchContext searchContext) { Settings indexSettings = searchContext.getQueryShardContext().getIndexSettings().getSettings(); return MapperService.INDEX_MAPPING_DEPTH_LIMIT_SETTING.get(indexSettings).intValue();